diff --git a/.clang-format b/.clang-format deleted file mode 100644 index c05c7633b..000000000 --- a/.clang-format +++ /dev/null @@ -1,30 +0,0 @@ ---- -BasedOnStyle: LLVM -IndentWidth: 4 # 缩进宽度，LLVM 默认值为 2，改为 4 -AccessModifierOffset: -4 # public/protected/private 访问控制符相对成员的偏移，与 IndentWidth 配合，LLVM 默认值为 -2 -AlignOperands: AlignAfterOperator # 双目运算符的行间对齐，LLVM 默认值为 Align，改为带符号一起换行 -BreakBeforeBinaryOperators: All # 在双目运算符之前换行，LLVM 默认值为 None，改为换行时总是把双目运算符放在行首，包括赋值（=） -ColumnLimit: 0 # 列宽限制，LLVM 默认值为 80，改为不限制 -AllowShortBlocksOnASingleLine: Always # 是否允许短块（单个语句的块）不换行，LLVM 默认值为 Never，改为允许 -AllowShortLoopsOnASingleLine: true # 是否允许短循环不换行，LLVM 默认值为 false，改为允许 -InsertBraces: true # 是否在 if/for/while/switch 等语句后插入大括号，LLVM 默认值为 false，改为允许 -BreakBeforeBraces: Custom # 大括号换行配置，LLVM 默认值为 LLVM，改为自定义以使 BraceWrapping 生效 -BraceWrapping: - AfterCaseLabel: false - AfterClass: false - AfterControlStatement: Never - AfterEnum: false - AfterFunction: false - AfterNamespace: false - AfterObjCDeclaration: false - AfterStruct: false - AfterUnion: false - AfterExternBlock: false - BeforeCatch: false - BeforeElse: false - BeforeLambdaBody: false - BeforeWhile: false - IndentBraces: false - SplitEmptyFunction: true - SplitEmptyRecord: true - SplitEmptyNamespace: true diff --git a/.github/ISSUE_TEMPLATE/bug.md b/.github/ISSUE_TEMPLATE/bug.md deleted file mode 100644 index e747a6674..000000000 --- a/.github/ISSUE_TEMPLATE/bug.md +++ /dev/null @@ -1,26 +0,0 @@ ---- -name: BUG反馈 -about: 描述Bug的详细信息 -title: "[BUG]" -labels: '' -assignees: '' - ---- - -**版本** -0.1.0 - -**问题描述** -简单描述问题表现，情包括问题发生的硬件平台和软件环境 - -**如何复现** -1. Go to '...' -2. Click on '....' -3. Scroll down to '....' -4. See error - -**预期结果** -简单描述预期得到的正确结果 - -**截图** -添加图片 diff --git a/.github/ISSUE_TEMPLATE/dev.md b/.github/ISSUE_TEMPLATE/dev.md deleted file mode 100644 index 54e80b5ee..000000000 --- a/.github/ISSUE_TEMPLATE/dev.md +++ /dev/null @@ -1,14 +0,0 @@ ---- -name: 功能开发 -about: 描述新功能的需求 -title: "[DEV]" -labels: '' -assignees: '' - ---- - -**目标版本** -main - -**功能描述** -简单描述新功能的需求 diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml deleted file mode 100644 index 620ef05c2..000000000 --- a/.github/workflows/build.yml +++ /dev/null @@ -1,63 +0,0 @@ -name: Build and test -on: - pull_request: - push: - paths-ignore: - - '**.md' - - 'LICENSE' - -jobs: - build: - name: Build - strategy: - fail-fast: false - matrix: - os: [windows-latest, ubuntu-latest] - type: [debug, release] - runs-on: ${{ matrix.os }} - steps: - - - name: checkout code - uses: actions/checkout@v4 - with: - submodules: recursive - fetch-depth: 0 - - - name: Check Format - if: matrix.os == 'ubuntu-latest' && matrix.type == 'debug' - run: | - sudo apt-get install -y clang-format-16 - pip install black ruff==0.15.20 colorama - if [ "${{ github.event_name }}" = "pull_request" ]; then - python3 scripts/format.py --ref "${{ github.event.pull_request.base.sha }}" --path src --check - elif [ "${{ github.event.before }}" != "0000000000000000000000000000000000000000" ]; then - python3 scripts/format.py --ref "${{ github.event.before }}" --path src --check - else - python3 scripts/format.py --path src --check - fi - - - name: install xmake - uses: xmake-io/github-action-setup-xmake@v1 - with: - xmake-version: latest - - - name: Build & Install - run: python scripts/install.py --omp=y -y - - - name: install python packages - run: | - pip install numpy - pip install torch - - - name: Python Test - run: python scripts/python_test.py --cpu - - - name: run infinirt-test --cpu on Linux - if: matrix.os == 'ubuntu-latest' - run: | - ./build/linux/x86_64/release/infinirt-test --cpu - - - name: run infinirt-test --cpu on Windows - if: matrix.os == 'windows-latest' - run: | - .\build\windows\x64\release\infinirt-test.exe --cpu diff --git a/.github/workflows/ruff.yml b/.github/workflows/ruff.yml deleted file mode 100644 index cfeb17909..000000000 --- a/.github/workflows/ruff.yml +++ /dev/null @@ -1,16 +0,0 @@ -name: Ruff -on: [push, pull_request] -jobs: - ruff: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - uses: chartboost/ruff-action@v1 - with: - src: './python/' - version: "0.15.20" - - uses: chartboost/ruff-action@v1 - with: - src: './python/' - args: format --check - version: "0.15.20" diff --git a/.gitignore b/.gitignore deleted file mode 100644 index d9479360b..000000000 --- a/.gitignore +++ /dev/null @@ -1,34 +0,0 @@ -# Xmake cache -.xmake/ -build/ -python/infinicore/lib/*.so - -# MacOS Cache -.DS_Store - -# Vscode -.vscode/ - -# Pycharm -.idea/ - -# Python -__pycache__/ -*.egg-info/ - -# Log -*.log - -# Cache -cache/ - -# JSON -*.json - -#GGUF -*.gguf - -# Compressed -*.gz -*.zip -*.tar diff --git a/.gitmodules b/.gitmodules index 37cafb336..4f1ec0ccf 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,34 +1,9 @@ -[submodule "third_party/spdlog"] - path = third_party/spdlog - url = https://github.com/gabime/spdlog.git -[submodule "third_party/nlohmann_json"] - path = third_party/nlohmann_json - url = https://github.com/nlohmann/json.git - branch = master -[submodule "third_party/flash-attention"] - path = third_party/flash-attention - url = https://github.com/Dao-AILab/flash-attention.git - ignore = untracked - update = none -[submodule "third_party/cutlass"] - path = third_party/cutlass - url = https://github.com/NVIDIA/cutlass.git - ignore = untracked - update = none [submodule "submodules/InfiniRT"] path = submodules/InfiniRT url = https://github.com/InfiniTensor/InfiniRT.git - branch = master [submodule "submodules/InfiniOps"] path = submodules/InfiniOps url = https://github.com/InfiniTensor/InfiniOps.git [submodule "submodules/InfiniCCL"] path = submodules/InfiniCCL url = https://github.com/InfiniTensor/InfiniCCL.git - branch = master -[submodule "third_party/mate"] - path = third_party/mate - url = https://github.com/MooreThreads/mate - branch = v0.1.3 - ignore = untracked - update = none diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 000000000..f718fae34 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,85 @@ +# Contributing Guide + +InfiniCore is a version manifest for the InfiniTensor core components. Changes +to component implementations belong in the corresponding component +repositories; changes here should be limited to component revisions and +manifest documentation. + +## General + +1. Keep changes minimal. Do not add what is not necessary for the manifest. +2. Prefer self-explanatory content over abundant comments. +3. Use Markdown syntax, such as backticks, when referring to identifiers. +4. Write comments and error messages in English and follow the conventions of + the file format or language. +5. End every file with a newline. +6. Review all changes before committing, especially AI-generated changes. + +## Submodule Updates + +1. Update a gitlink only to a reviewed, reachable component revision. +2. Keep independent component updates in separate commits when possible. When + multiple revisions must move together for compatibility, update them in one + commit and explain the relationship in the pull request. +3. Do not copy component source code, build logic, or runtime implementations + into InfiniCore. +4. Record the old and new revisions and link the corresponding upstream change + in the pull request. + +## Commits + +Commit messages must follow +[Conventional Commits](https://www.conventionalcommits.org/): + +```text +[optional scope][!]: +``` + +Use the type that best describes the change, such as `docs`, `fix`, `refactor`, +or `chore`. Mark breaking changes with `!` and describe the impact in the +commit body or a `BREAKING CHANGE` footer. + +Small changes should be represented by a single commit. Larger changes may use +multiple commits, but each commit must be meaningful, focused, and well-formed. + +## Branches + +Branch names use the format `/xxx-yyyy-zzzz`, where `` matches the +pull request title's Conventional Commits type and words are joined with +hyphens. For example: + +```text +docs/contribution-guidelines +refactor/component-manifest +``` + +## Pull Requests + +1. Use Conventional Commits format for the pull request title. Its type must + match the branch type. +2. Explain what changed, why it changed, which component revisions are + affected, and how the change was validated. +3. Keep small pull requests squashed. Large pull requests may retain multiple + commits when every commit is meaningful and well-formed. +4. Include validation evidence for every affected component and platform. Use + broader validation for high-risk changes, release preparation, shared build + changes, or cross-platform behavior changes. +5. State which checks were not run and why. + +## Validation + +InfiniCore has no build or test suite of its own. Validate the manifest itself +before opening a pull request: + +```shell +git diff --check +git ls-files --stage submodules +git submodule sync --recursive +git submodule update --init --recursive +git submodule status --recursive +``` + +Confirm that each component entry is a gitlink, each pinned revision can be +checked out, and the recursive submodule state is clean. For revision updates, +also include the relevant component or integration build and test evidence in +the pull request. diff --git a/DEV.md b/DEV.md deleted file mode 100644 index 77920c2bc..000000000 --- a/DEV.md +++ /dev/null @@ -1,207 +0,0 @@ -# InfiniCore 开发者手册 - -Dear 开发者，感谢你参与 InfiniCore 开源项目的开发！本文档将帮助你了解如何向 InfiniCore 项目贡献代码。 - -## 项目介绍 - -### 项目模块体系 - -- infinicore：统一计算框架。提供 Python 和 C++ 接口，支持多种硬件平台。 -- infinirt：统一底层运行时库，提供 C 语言接口，依赖 infini-utils。 -- infiniop：统一底层算子库，提供 C 语言接口，依赖 infinirt。除了 C++ 算子实现之外，也包括使用九齿（triton）的算子实现，这部分算子需要在编译之前使用脚本生成源文件。安装后可以运行位于 `test/infiniop` 中的单测脚本进行测试。 -- infiniccl：统一通信库，提供 C 语言接口，依赖 infinirt。 -- infini-utils：全模块通用工具代码。 -- utils-test：工具库测试代码，依赖 infini-utils。 -- infiniccl-test：通信库测试代码，使用前需要安装好 infiniccl。 - -### 文件目录结构 - -```bash -├── xmake.lua # 总体 xmake 编译配置，包含所有平台的编译选项和宏定义 -├── xmake/*.lua # 各平台 xmake 编译配置， 包含各平台特有的编译方式 -│ -├── include/ # 对外暴露的头文件目录，安装时会被复制到安装目录 -│ ├── infinicore/*.hpp # InfiniCore计算库头文件（C++） -│ ├── infiniop/*.h # InfiniOP算子库子头文件（C） -│ ├── *.h/.hpp # 模块核心头文件 -│ -├── src/ # 各模块源代码目录，包含源代码文件以及不对外暴露的头文件 -│ ├── infinicore/ # InfiniCore源代码目录 -│ │ ├── context/ # 张量运行时/硬件上下文管理源代码目录 -│ │ ├── nn/ # 机器学习模块源代码目录 -│ │ ├── ops/ # 张量算子源代码目录 -│ │ ├── pybind/ # pybind 接口源代码目录 -│ │ ├── tensor/ # 张量库源代码目录 -│ │ -│ ├── infiniop/ # InfiniOP算子库源代码目录 -│ │ ├── devices/ # 每个设备平台各自的通用代码目录 -│ │ ├── ops/ # 算子实现代码目录 -│ │ │ ├── [op]/ -│ │ │ │ ├── [device]/ # 各硬件平台算子实现代码目录 -│ │ │ │ ├── operator.cc # 算子C语言接口实现 -│ │ ├── reduce/ # 规约类算子通用代码目录 -│ │ ├── elementwise/ # 逐元素类算子通用代码目录 -│ │ ├── *.h # 核心结构体定义 -│ │ -│ ├── infinirt/ # InfiniRT运行时库源代码目录 -│ ├── infiniccl/ # InfiniCCL集合通信库源代码目录 -│ -├── test/ # 测试源代码目录 -│ ├── infinicore/ # InfiniCore测试目录 -│ │ ├── ops/*.py # 算子单测脚本（依赖各平台PyTorch） -│ ├── infiniop/ # InfiniOP算子库单元测试目录 -│ │ ├── *.py # 单测脚本（依赖各平台PyTorch） -│ -├── scripts/ # 脚本目录 -│ ├── install.py # 安装编译脚本 -│ ├── python_test.py # 运行所有单测脚本 -``` - -## 开发引导 - -### 代码提交流程 - -1. 在github仓库issue页面根据任务类型（开发或bug）创建 issue，所有commit必须有对应的 issue 编号。 -2. 外部人员需要通过 fork 代码仓库提交 PR。 -3. 根据 issue 编号建立分支，分支名字格式为 `issue/#` （# 为issue 编号）。如果出现重复，可在后面添加“-#”序号，或用“/”后增加说明。 -4. 所有 commit 信息必须以 `issue/#` 开头， -5. 分支推到远程后，建 Pull Request，标题需要以 `issue/#` 开头。在原issue页面上将该PR关联。 -6. PR必须添加至少两位审核员（模块负责人和项目管理员等），PR中需附上最后一次修改后测试通过的截图。 -7. PR通过审核，通过自动测试，无代码冲突后方可合并。合并后，关闭原 issue。 - -### 如何开发一个新算子 - -- 如果你想通过 C++ 以及硬件原生语言开发一个新的算子，请阅读 [infinicore::ops 开发指南](/src/infinicore/ops/README.md) - -### C++ 代码命名书写规范 - -1. 类型 - - 内部数据结构类型 `UpperCamelCase` - - ```c++ - // 尽量使用 Infinixx 开头 - struct InfiniopMatmulCudaDescriptor; - - template - class HashMap; - - using ValueMap = std::unordered_map; - ``` - - 对外暴露的指针类型和枚举类型 `infinixx[XxxXxx]_t` - - 常量使用 `INFINI_UPPER_SNAKE_CASE` - - ```c++ - typedef struct InfiniopMatmulCudaDescriptor *infiniopMatmulCudaDescriptor_t; - - typedef enum { - // INFINI... - INFINI_DTYPE_INVALID = 0, - } infiniDtype_t; - ``` - -2. 普通变量、形参、类数据成员，使用 `snake_case` - - 成员名前下划线特指private成员，其他情况应避免使用前下划线 - - ```c++ - int max_count; - - class Example { - public: - std::string getUserName(std::string user_id); - private: - // private数据成员名字前加下划线 - int _max_count; - std::string _user_name; - }; - - struct UrlTableProperties { - string name; - int num_entries; - static Pool* pool; - }; - ``` - - 当形参与函数内部变量或成员变量重名，可选择其中一个名字后加下划线。当函数内部临时变量和成员重名时，临时变量名字后加下划线。后下划线表示“临时” - - ```c++ - void do(int count_){ - int count = count_; - } - ``` - -3. 函数，使用 lowerCamelCase - - ```c++ - int getMaxValue() const; - ``` - - InfiniCore 中和 torch 对齐的接口，使用 `snake_case`。 - -4. const/volatile修饰符写在类型前面 - - ```c++ - const void *ptr; - const int num; - ``` - -### 代码格式化 - -本项目分别使用 `clang-format-16` 和 `black` 对 C/C++ 以及 Python 代码进行格式化。可以使用 [`scripts/format.py`](/scripts/format.py) 脚本实现代码格式化检查和操作。 - -使用 - -```shell -python scripts/format.py -h -``` - -查看脚本帮助信息： - -```plaintext -usage: format.py [-h] [--ref REF] [--path [PATH ...]] [--check] [--c C] [--py PY] - -options: - -h, --help show this help message and exit - --ref REF Git reference (commit hash) to compare against. - --path [PATH ...] Files to format or check. - --check Check files without modifying them. - --c C C formatter (default: clang-format-16) - --py PY Python formatter (default: black) -``` - -参数中： - -- `ref` 和 `path` 控制格式化的文件范围 - - 若 `ref` 和 `path` 都为空，格式化当前暂存（git added）的文件； - - 否则 - - 若 `ref` 非空，将比较指定 commit 和当前代码的差异，只格式化修改过的文件； - - 若 `path` 非空，可传入多个路径（`--path p0 p1 p2`），只格式化指定路径及其子目录中的文件； -- 若设置 `--check`，将检查代码是否需要修改格式，不修改文件内容； -- 通过 `--c` 指定 c/c++ 格式化器，默认为 `clang-format-16`； -- 通过 `--python` 指定 python 格式化器 `black`； - -### vscode 开发配置 - -基本配置见 [xmake 官方文档](https://xmake.io/#/zh-cn/plugin/more_plugins?id=%e9%85%8d%e7%bd%ae-intellsence)。 - -- TL;DR - - clangd - - 打开 *xmake.lua*，保存一次以触发编译命令生成，将在工作路径下自动生成 *.vscode/compile_commands.json* 文件。然后在这个文件夹下创建 *settings.json*，填入： - - > .vscode/settings.json - - ```json - { - "clangd.arguments": [ - "--compile-commands-dir=.vscode" - ], - "xmake.additionalConfigArguments": [ - // 在这里配置 XMAKE_CONFIG_FLAGS - "--nv-gpu=y" - ], - } - ``` diff --git a/MANIFEST.in b/MANIFEST.in deleted file mode 100644 index 462d79fd6..000000000 --- a/MANIFEST.in +++ /dev/null @@ -1 +0,0 @@ -global-include * diff --git a/README.md b/README.md index 21d490224..c0ddeec8a 100644 --- a/README.md +++ b/README.md @@ -1,303 +1,31 @@ # InfiniCore -[![Doc](https://img.shields.io/badge/Document-ready-blue)](https://github.com/InfiniTensor/InfiniCore-Documentation) -[![CI](https://github.com/InfiniTensor/InfiniCore/actions/workflows/build.yml/badge.svg?branch=main)](https://github.com/InfiniTensor/InfiniCore/actions) -[![license](https://img.shields.io/github/license/InfiniTensor/InfiniCore)](https://mit-license.org/) -![GitHub repo size](https://img.shields.io/github/repo-size/InfiniTensor/InfiniCore) -![GitHub code size in bytes](https://img.shields.io/github/languages/code-size/InfiniTensor/InfiniCore) +InfiniCore is the version manifest for the InfiniTensor core stack. It pins three independently developed components as Git submodules: -[![GitHub Issues](https://img.shields.io/github/issues/InfiniTensor/InfiniCore)](https://github.com/InfiniTensor/InfiniCore/issues) -[![GitHub Pull Requests](https://img.shields.io/github/issues-pr/InfiniTensor/InfiniCore)](https://github.com/InfiniTensor/InfiniCore/pulls) -![GitHub contributors](https://img.shields.io/github/contributors/InfiniTensor/InfiniCore) -![GitHub commit activity](https://img.shields.io/github/commit-activity/m/InfiniTensor/InfiniCore) +- [InfiniRT](https://github.com/InfiniTensor/InfiniRT) provides device and runtime services. +- [InfiniOps](https://github.com/InfiniTensor/InfiniOps) provides computational operators. +- [InfiniCCL](https://github.com/InfiniTensor/InfiniCCL) provides collective communication. -InfiniCore 是一个跨平台统一编程工具集，为不同芯片平台的功能（包括计算、运行时、通信等）提供统一 C 语言接口。目前支持的硬件和后端包括： +## Clone and update -- CPU； -- CUDA - - 英伟达 GPU； - - 摩尔线程 GPU； - - 天数智芯 GPU； - - 沐曦 GPU； - - 海光 DCU； - - 阿里 PPU； -- 华为昇腾 NPU； -- 寒武纪 MLU； -- 昆仑芯 XPU； - -API 定义以及使用方式详见 [`InfiniCore文档`](https://github.com/InfiniTensor/InfiniCore-Documentation)。 - -## 项目依赖 - -- [Xmake](https://xmake.io/)：跨平台自动构建工具，用于编译 InfiniCore 项目。 -- [gcc-11](https://gcc.gnu.org/) 以上或者 [clang-16](https://clang.llvm.org/)：基础编译器，需要支持 C++ 17 标准。 -- [Python>=3.10](https://www.python.org/) - - [PyTorch](https://pytorch.org/)：可选，用于对比测试。 -- 各个硬件平台的工具包：请参考各厂商官方文档（如英伟达平台需要安装 CUDA Toolkit）。 - -## 配置和使用 - -### 一、克隆项目 - -由于仓库中含有子模块，所以在克隆时请添加 `--recursive` 或 `--recurse-submodules`，如： +Clone the repository with its pinned submodules: ```shell -git clone --recursive https://github.com/InfiniTensor/InfiniCore.git +git clone --recurse-submodules https://github.com/InfiniTensor/InfiniCore.git ``` -或者在普通克隆后进行更新： +Initialize or update the submodules in an existing checkout: ```shell +git submodule sync --recursive git submodule update --init --recursive ``` -配置`INFINI_ROOT` 和 `LD_LIBRARY_PATH` 环境变量。 -默认`INFINI_ROOT`为`$HOME/.infini`，可以使用以下命令自动配置： - -```shell -source scripts/set_env_linux.sh -``` - -如果你需要在本地开发九齿算子（即需要对九齿算子库进行修改），推荐单独克隆[九齿算子库](https://github.com/InfiniTensor/ntops)，并从本地安装： - -```shell -git clone https://github.com/InfiniTensor/ntops.git -cd ntops -pip install -e . -``` - -### 二、编译安装 - -InfiniCore 项目主要包括： - -1. 底层 C 库（InfiniOP/InfiniRT/InfiniCCL）：[`一键安装`](#一键安装底层库)|[`手动安装`](#手动安装底层库)； -2. InfiniCore C++ 库：[`安装指令`](#2-安装-c-库) -3. InfiniCore Python 包（依赖[九齿算子库](https://github.com/InfiniTensor/ntops)）：[`安装指令`](#3-安装-python-包) - -三者需要按照顺序进行编译安装。 - -#### 1. 安装底层库 - -##### 一键安装底层库 - -在 `script/` 目录中提供了 `install.py` 安装脚本。使用方式如下： - -```shell -cd InfiniCore - -python scripts/install.py [XMAKE_CONFIG_FLAGS] -``` - -参数 `XMAKE_CONFIG_FLAGS` 是 xmake 构建配置，可配置下列可选项： - -| 选项 | 功能 | 默认值 -|--------------------------|-----------------------------------|:-: -| `--omp=[y\|n]` | 是否使用 OpenMP | y -| `--cpu=[y\|n]` | 是否编译 CPU 接口实现 | y -| `--nv-gpu=[y\|n]` | 是否编译英伟达 GPU 接口实现 | n -| `--ascend-npu=[y\|n]` | 是否编译昇腾 NPU 接口实现 | n -| `--cambricon-mlu=[y\|n]` | 是否编译寒武纪 MLU 接口实现 | n -| `--metax-gpu=[y\|n]` | 是否编译沐曦 GPU 接口实现 | n -| `--use-mc=[y\|n]` | 是否沐曦 GPU 接口实现使用maca SDK | n -| `--moore-gpu=[y\|n]` | 是否编译摩尔线程 GPU 接口实现 | n -| `--iluvatar-gpu=[y\|n]` | 是否编译天数 GPU 接口实现 | n -| `--qy-gpu=[y\|n]` | 是否编译QY GPU 接口实现 | n -| `--hygon-dcu=[y\|n]` | 是否编译海光 DCU 接口实现 | n -| `--kunlun-xpu=[y\|n]` | 是否编译昆仑 XPU 接口实现 | n -| `--ali-ppu=[y\|n]` | 是否编译阿里 PPU 接口实现 | n -| `--ninetoothed=[y\|n]` | 是否编译九齿实现 | n -| `--ccl=[y\|n]` | 是否编译 InfiniCCL 通信库接口实现 | n -| `--graph=[y\|n]` | 是否编译 cuda graph 接口实现 | n - -##### 手动安装底层库 - -0. 生成九齿算子（可选） - - - 克隆并安装[九齿算子库](https://github.com/InfiniTensor/ntops)。 - - - 在 `InfiniCore` 文件夹下运行以下命令 AOT 编译库中的九齿算子： - - ```shell - PYTHONPATH=${PYTHONPATH}:src python scripts/build_ntops.py - ``` - -1. 项目配置 - - windows系统上，建议使用`xmake v2.8.9`编译项目。 - - 查看当前配置 - - ```shell - xmake f -v - ``` - - - 配置 CPU（默认配置） - - ```shell - xmake f -cv - ``` - - - 配置加速卡 - - ```shell - # 英伟达 - # 可以指定 CUDA 路径， 一般环境变量为 `CUDA_HOME` 或者 `CUDA_ROOT` - # window系统：--cuda="%CUDA_HOME%" - # linux系统：--cuda=$CUDA_HOME - xmake f --nv-gpu=true --cuda=$CUDA_HOME -cv - - # QY - # 需要指定环境变量QY_ROOT来确认库所在位置，比如说export QY_ROOT=/usr/local/XX - xmake f --qy-gpu=true --cuda=$CUDA_HOME -cv - - # 寒武纪 - xmake f --cambricon-mlu=true -cv - - # 华为昇腾 - xmake f --ascend-npu=true -cv - ``` - -##### 试验功能 -- 使用英伟达平台 flash attention 库中的算子 - - ```shell - - # 该功能依赖 flash-attention 和 cutlass，默认不随仓库递归拉取。 - # 对应子模块固定为以下提交： - ## flash-attention commit: 10846960ca0793b993446f6dbaf696479c127a9d - ## cutlass commit: 087c84df83d254b5fb295a7a408f1a1d554085cf - - # 若需启用英伟达平台 flash attention 能力，请手动初始化对应子模块： - git -c submodule.third_party/flash-attention.update=checkout \ - -c submodule.third_party/cutlass.update=checkout \ - submodule update --init third_party/flash-attention third_party/cutlass - - # 上述命令只初始化这两个顶层子模块，并会切换到仓库记录的固定提交。 - - # 设置cutlass路径的环境变量CUTLASS_HOME(部分环境可选) - export CUTLASS_HOME=/InfiniCore/third_party/cutlass - - # xmake配置环节额外打开 --aten 开关，并设置 --flash-attn 库位置，例(cuda路径部分环境可使用默认)： - xmake f --nv-gpu=y --ccl=y --aten=y [--graph=y] [--cuda=$CUDA_HOME] --flash-attn=/InfiniCore/third_party/flash-attention -cv - - # 设置额外的环境变量 - export CPLUS_INCLUDE_PATH=$CUDA_HOME/include:$CPLUS_INCLUDE_PATH - - # flash attention库会伴随infinicore_cpp_api一同编译安装 - - ``` - -##### 试验功能 -- 使用摩尔线程开源 mate 提供的 flash attention 能力 - ```shell - #该功能依赖摩尔线程开源项目 mate（https://github.com/MooreThreads/mate） v0.1.3 版本，默认不随仓库递归拉取。 - - #若需启用摩尔线程开源项目 mate 提供的 flash attention 能力，请手动初始化对应子模块： - git -c submodule.third_party/mate.update=checkout submodule update --init --recursive third_party/mate - - #随后参考 mate v0.1.3 README 进行编译，之后在 xmake 配置环节额外打开 --aten 开关和 --flash-attn 使用 mate 提供的 flash attention 能力，可参考： - xmake f --moore-gpu=y --aten=y --flash-attn=y -cv - ``` - - -##### 试验功能 -- 使用海光 DCU 平台预编译 flash-attn 能力 - - ```shell - # 海光 DCU 不在 InfiniCore 内现场编译 flash-attn，而是链接 Python 环境中已经安装好的 flash-attn 运行库。 - # 因此 --flash-attn 需要指向 flash-attn 的 Python 安装根目录，通常是当前 Python 的 site-packages/dist-packages 目录。 - # 该目录下需要能找到以下两个文件： - # 1. flash_attn_2_cuda*.so - # 2. flash_attn/lib/libflash_attention.so - # 例如： - # /usr/local/lib/python3.10/dist-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so - # /usr/local/lib/python3.10/dist-packages/flash_attn/lib/libflash_attention.so - - # 若 flash_attn_2_cuda*.so 不在 --flash-attn 指定目录下，可通过 FLASH_ATTN_2_CUDA_SO 显式指定。 - export FLASH_ATTN_2_CUDA_SO=/usr/local/lib/python3.10/dist-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so - - # xmake 配置环节需要同时打开海光 DCU、ATen 和 flash-attn： - xmake f --hygon-dcu=y --aten=y --flash-attn=/usr/local/lib/python3.10/dist-packages -cv - - # 编译 Python/C++ 封装： - xmake build _infinicore - xmake install _infinicore - ``` - -##### 试验功能 -- 编译marlin相关算子 - - ```shell - - # 需要从github上克隆tvm_ffi仓库，克隆命令参考 - ## tvm-ffi commit: 35c99d0ac4cb784862115d0089f60c603acec8f9 - git clone https://github.com/apache/tvm-ffi.git --recursive - - # 设置TVM_ROOT - export TVM_ROOT=/tvm-ffi #用来搜索tvm相关头文件 - # 注意，编译gptq_marlin_gemm算子的时候除了指定TVM_ROOT以外，还需要指定cuda_arch - ``` - -2. 编译安装 - - 默认安装路径为 `$HOME/.infini`。 - - ```shell - xmake build && xmake install - ``` - -#### 2. 安装 C++ 库 - -```shell -xmake build _infinicore -xmake install _infinicore -``` - -#### 3. 安装 Python 包 - -```shell -pip install . -``` - -或 - -```shell -pip install -e . -``` - -注：开发时建议加入 `-e` 选项（即 `pip install -e .`），这样对 `python/infinicore` 做的更改将会实时得到反映，同时对 C++ 层所做的修改也只需要运行 `xmake build _infinicore && xmake install _infinicore` 便可以生效。 - -### 三、运行测试 - -#### 运行 InfiniCore Python算子接口测试 - -```bash -# 测试单算子 -python test/infinicore/ops/[operator].py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --Hygon | --ali] -# 测试全部算子 -python test/infinicore/run.py [--bench | --debug | --verbose] [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --ali] -``` - -使用 -h 查看更多参数。 +## Contributing -#### 运行 InfiniOP 算子测试 - -```shell -# 测试单算子 -python test/infiniop/[operator].py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --Hygon | --ali] -# 测试全部算子 -python scripts/python_test.py [--cpu | --nvidia | --cambricon | --ascend | --iluvatar | --metax | --moore | --kunlun | --Hygon | --ali] -``` - -#### 通信库（InfiniCCL）测试 - -编译（需要先安装底层库中的 InfiniCCL 库）： - -```shell -xmake build infiniccl-test -``` - -在英伟达平台运行测试（会自动使用所有可见的卡）： - -```shell -infiniccl-test --nvidia -``` +See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution and validation +guidelines. -## 如何开源贡献 +## License -见 [`InfiniCore开发者手册`](DEV.md)。 +InfiniCore is licensed under the MIT License. See [LICENSE](LICENSE). diff --git a/include/infiniccl.h b/include/infiniccl.h deleted file mode 100644 index b338d85f5..000000000 --- a/include/infiniccl.h +++ /dev/null @@ -1,75 +0,0 @@ -#ifndef __INFINICCL_API_H__ -#define __INFINICCL_API_H__ - -#include "infinirt.h" - -typedef enum { - INFINICCL_SUM = 0, - INFINICCL_PROD = 1, - INFINICCL_MAX = 2, - INFINICCL_MIN = 3, - INFINICCL_AVG = 4, -} infinicclReduceOp_t; - -struct InfinicclComm; - -typedef struct InfinicclComm *infinicclComm_t; - -__INFINI_C __export infiniStatus_t infinicclCommInitAll( - infiniDevice_t device_type, - infinicclComm_t *comms, - int ndevice, - const int *device_ids); - -__INFINI_C __export infiniStatus_t infinicclCommDestroy(infinicclComm_t comm); - -__INFINI_C __export infiniStatus_t infinicclGroupStart(infinicclComm_t comm); - -__INFINI_C __export infiniStatus_t infinicclGroupEnd(infinicclComm_t comm); - -__INFINI_C __export infiniStatus_t infinicclAllReduce( - void *sendbuf, - void *recvbuf, - size_t count, - infiniDtype_t dataype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream); - -__INFINI_C __export infiniStatus_t infinicclAllGather( - void *sendbuf, - void *recvbuf, - size_t send_count, - infiniDtype_t dataype, - infinicclComm_t comm, - infinirtStream_t stream); - -__INFINI_C __export infiniStatus_t infinicclAllGatherV( - void *sendbuf, - void *recvbuf, - const size_t *recv_counts, - int nranks, - infiniDtype_t dataype, - infinicclComm_t comm, - infinirtStream_t stream); - -__INFINI_C __export infiniStatus_t infinicclReduceScatter( - void *sendbuf, - void *recvbuf, - size_t recv_count, - infiniDtype_t dataype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream); - -__INFINI_C __export infiniStatus_t infinicclReduceScatterV( - void *sendbuf, - void *recvbuf, - const size_t *send_counts, - int nranks, - infiniDtype_t dataype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream); - -#endif diff --git a/include/infinicore.h b/include/infinicore.h deleted file mode 100644 index 849c62c14..000000000 --- a/include/infinicore.h +++ /dev/null @@ -1,77 +0,0 @@ -#ifndef __INFINICORE_API_H__ -#define __INFINICORE_API_H__ - -#if defined(_WIN32) -#define __export __declspec(dllexport) -#elif defined(__GNUC__) && ((__GNUC__ >= 4) || (__GNUC__ == 3 && __GNUC_MINOR__ >= 3)) -#define __export __attribute__((visibility("default"))) -#else -#define __export -#endif - -#ifdef __cplusplus -#define __INFINI_C extern "C" -#include -#else -#define __INFINI_C -#include -#endif - -typedef enum { - // Success - INFINI_STATUS_SUCCESS = 0, - // General Errors - INFINI_STATUS_INTERNAL_ERROR = 1, - INFINI_STATUS_NOT_IMPLEMENTED = 2, - INFINI_STATUS_BAD_PARAM = 3, - INFINI_STATUS_NULL_POINTER = 4, - INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED = 5, - INFINI_STATUS_DEVICE_NOT_FOUND = 6, - INFINI_STATUS_DEVICE_NOT_INITIALIZED = 7, - INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED = 8, - // Op Errors - INFINI_STATUS_BAD_TENSOR_DTYPE = 10, - INFINI_STATUS_BAD_TENSOR_SHAPE = 11, - INFINI_STATUS_BAD_TENSOR_STRIDES = 12, - INFINI_STATUS_INSUFFICIENT_WORKSPACE = 13, -} infiniStatus_t; - -typedef enum { - INFINI_DEVICE_CPU = 0, - INFINI_DEVICE_NVIDIA = 1, - INFINI_DEVICE_CAMBRICON = 2, - INFINI_DEVICE_ASCEND = 3, - INFINI_DEVICE_METAX = 4, - INFINI_DEVICE_MOORE = 5, - INFINI_DEVICE_ILUVATAR = 6, - INFINI_DEVICE_KUNLUN = 7, - INFINI_DEVICE_HYGON = 8, - INFINI_DEVICE_QY = 9, - INFINI_DEVICE_ALI = 10, - INFINI_DEVICE_TYPE_COUNT -} infiniDevice_t; - -typedef enum { - INFINI_DTYPE_INVALID = 0, - INFINI_DTYPE_BYTE = 1, - INFINI_DTYPE_BOOL = 2, - INFINI_DTYPE_I8 = 3, - INFINI_DTYPE_I16 = 4, - INFINI_DTYPE_I32 = 5, - INFINI_DTYPE_I64 = 6, - INFINI_DTYPE_U8 = 7, - INFINI_DTYPE_U16 = 8, - INFINI_DTYPE_U32 = 9, - INFINI_DTYPE_U64 = 10, - INFINI_DTYPE_F8 = 11, - INFINI_DTYPE_F16 = 12, - INFINI_DTYPE_F32 = 13, - INFINI_DTYPE_F64 = 14, - INFINI_DTYPE_C16 = 15, - INFINI_DTYPE_C32 = 16, - INFINI_DTYPE_C64 = 17, - INFINI_DTYPE_C128 = 18, - INFINI_DTYPE_BF16 = 19, -} infiniDtype_t; - -#endif // __INFINICORE_API_H__ diff --git a/include/infinicore.hpp b/include/infinicore.hpp deleted file mode 100644 index ddd2f3d75..000000000 --- a/include/infinicore.hpp +++ /dev/null @@ -1,8 +0,0 @@ -#pragma once - -#include "infinicore/device_event.hpp" -#include "infinicore/io.hpp" -#include "infinicore/nn.hpp" -#include "infinicore/ops.hpp" -#include "infinicore/quantization.hpp" -#include "infinicore/tensor.hpp" diff --git a/include/infinicore/adaptor/aten_adaptor.hpp b/include/infinicore/adaptor/aten_adaptor.hpp deleted file mode 100644 index 74b053dc5..000000000 --- a/include/infinicore/adaptor/aten_adaptor.hpp +++ /dev/null @@ -1,73 +0,0 @@ -#ifdef ENABLE_ATEN -#pragma once -#include "../context/context.hpp" -#include "../tensor.hpp" - -#include - -#if defined(ENABLE_HYGON_API) -#include -#include -#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) -#include -#include -#include -#endif - -#if defined(ENABLE_MOORE_API) -#include -#include -#include -#endif - -namespace infinicore::adaptor { -inline at::ScalarType to_at_dtype(DataType dtype) { - switch (dtype) { - case DataType::F32: - return at::kFloat; - case DataType::F16: - return at::kHalf; - case DataType::BF16: - return at::kBFloat16; - case DataType::I32: - return at::kInt; - case DataType::I64: - return at::kLong; - default: - throw std::runtime_error("Unsupported dtype for ATen"); - } -} - -inline at::Device to_at_device(const Device &device) { - // PyTorch ATen only exposes standard device types (e.g. kCPU/kCUDA). - // Treat MetaX/QY devices as CUDA devices for ATen tensor interoperability. - if (device.getType() == Device::Type::NVIDIA || device.getType() == Device::Type::METAX || device.getType() == Device::Type::QY || device.getType() == Device::Type::HYGON) { - return at::Device(at::kCUDA, device.getIndex()); - } else if (device.getType() == Device::Type::CPU) { - return at::Device(at::kCPU); - } -#if defined(ENABLE_MOORE_API) - else if (device.getType() == Device::Type::MOORE) { - return at::Device(at::DeviceType::PrivateUse1, device.getIndex()); - } -#endif - else { - throw std::runtime_error("Unsupported device type for ATen"); - } -} - -at::Tensor to_aten_tensor(const infinicore::Tensor &t); - -#if defined(ENABLE_HYGON_API) -c10::hip::HIPStream get_hip_stream(); -#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) -c10::cuda::CUDAStream get_cuda_stream(); -#endif - -#if defined(ENABLE_MOORE_API) -c10::musa::MUSAStream get_musa_stream(); -#endif - -} // namespace infinicore::adaptor - -#endif // ENABLE_ATEN diff --git a/include/infinicore/adaptor/flash_attention_adaptor.hpp b/include/infinicore/adaptor/flash_attention_adaptor.hpp deleted file mode 100644 index c5bf14858..000000000 --- a/include/infinicore/adaptor/flash_attention_adaptor.hpp +++ /dev/null @@ -1,139 +0,0 @@ -#ifdef ENABLE_FLASH_ATTN -#pragma once -#include "aten_adaptor.hpp" - -// NVIDIA flash-attn-nvidia.so uses namespace flash. The pip/MetaX flash_attn_2_cuda extension -// exports the same entry points at global scope (no namespace), matching FLASH_NAMESPACE builds -// where the namespace is empty. -#if !defined(ENABLE_METAX_API) -namespace flash { -#endif -std::vector -mha_fwd(at::Tensor &q, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) - const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) - const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8) - std::optional &out_, // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8) - std::optional &alibi_slopes_, // num_heads or batch_size x num_heads - const float p_dropout, - const float softmax_scale, - bool is_causal, - int window_size_left, - int window_size_right, - const float softcap, - const bool return_softmax, - std::optional gen_ -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. - , - std::optional &flash_attn_mars_ext_ -#endif -); - -std::vector -mha_varlen_fwd(at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i - const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. - const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. - std::optional &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i - const at::Tensor &cu_seqlens_q, // b+1 - const at::Tensor &cu_seqlens_k, // b+1 - std::optional &seqused_k, // b. If given, only this many elements of each batch element's keys are used. - std::optional &leftpad_k_, // batch_size - std::optional &block_table_, // batch_size x max_num_blocks_per_seq - std::optional &alibi_slopes_, // num_heads or b x num_heads - int max_seqlen_q, - const int max_seqlen_k, - const float p_dropout, - const float softmax_scale, - const bool zero_tensors, - bool is_causal, - int window_size_left, - int window_size_right, - const float softcap, - const bool return_softmax, - std::optional gen_ -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. - , - std::optional &flash_attn_mars_ext_ -#endif -); - -std::vector -mha_bwd(const at::Tensor &dout, // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8) - const at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size - const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x head_size - const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x head_size - const at::Tensor &out, // batch_size x seqlen_q x num_heads x head_size - const at::Tensor &softmax_lse, // b x h x seqlen_q - std::optional &dq_, // batch_size x seqlen_q x num_heads x head_size - std::optional &dk_, // batch_size x seqlen_k x num_heads_k x head_size - std::optional &dv_, // batch_size x seqlen_k x num_heads_k x head_size - std::optional &alibi_slopes_, // num_heads or batch_size x num_heads - const float p_dropout, // probability to drop - const float softmax_scale, - const bool is_causal, - int window_size_left, - int window_size_right, - const float softcap, - const bool deterministic, - std::optional gen_, - std::optional &rng_state); - -std::vector -mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads, x head_size - const at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i - const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i - const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i - const at::Tensor &out, // total_q x num_heads x head_size - const at::Tensor &softmax_lse, // h x total_q, softmax logsumexp - std::optional &dq_, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i - std::optional &dk_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i - std::optional &dv_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i - const at::Tensor &cu_seqlens_q, // b+1 - const at::Tensor &cu_seqlens_k, // b+1 - std::optional &alibi_slopes_, // num_heads or b x num_heads - const int max_seqlen_q, - const int max_seqlen_k, // max sequence length to choose the kernel - const float p_dropout, // probability to drop - const float softmax_scale, - const bool zero_tensors, - const bool is_causal, - int window_size_left, - int window_size_right, - const float softcap, - const bool deterministic, - std::optional gen_, - std::optional &rng_state); - -std::vector -mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size - const at::Tensor &kcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. - const at::Tensor &vcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. - std::optional &k_, // batch_size x seqlen_knew x num_heads_k x head_size - std::optional &v_, // batch_size x seqlen_knew x num_heads_k x head_size - std::optional &seqlens_k_, // batch_size - std::optional &rotary_cos_, // seqlen_ro x (rotary_dim / 2) - std::optional &rotary_sin_, // seqlen_ro x (rotary_dim / 2) - std::optional &cache_batch_idx_, // indices to index into the KV cache - std::optional &leftpad_k_, // batch_size - std::optional &block_table_, // batch_size x max_num_blocks_per_seq - std::optional &alibi_slopes_, // num_heads or batch_size x num_heads - std::optional &out_, // batch_size x seqlen_q x num_heads x head_size - const float softmax_scale, - bool is_causal, - int window_size_left, - int window_size_right, - const float softcap, - bool is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2 - int num_splits -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - // MetaX/Mars `flash_attn_2_cuda` (e.g. 2.6.x+mars) appends this argument vs upstream Dao-AILab flash-attn. - , - std::optional &flash_attn_mars_ext_ -#endif -); - -#if !defined(ENABLE_METAX_API) -} // namespace flash -#endif -#endif // ENABLE_FLASH_ATTN diff --git a/include/infinicore/analyzer.hpp b/include/infinicore/analyzer.hpp deleted file mode 100644 index 81e9ab1a9..000000000 --- a/include/infinicore/analyzer.hpp +++ /dev/null @@ -1,11 +0,0 @@ -#pragma once - -// Convenience header — includes all analyzer components. - -#include "analyzer/intent_generator.hpp" -#include "analyzer/mutual_awareness_analyzer.hpp" -#include "analyzer/op_trace.hpp" -#include "analyzer/op_type.hpp" -#include "analyzer/optimization_intent.hpp" -#include "analyzer/phase_detector.hpp" -#include "analyzer/resource_sensor.hpp" diff --git a/include/infinicore/analyzer/intent_generator.hpp b/include/infinicore/analyzer/intent_generator.hpp deleted file mode 100644 index b166e4c1e..000000000 --- a/include/infinicore/analyzer/intent_generator.hpp +++ /dev/null @@ -1,268 +0,0 @@ -#pragma once - -#include "op_trace.hpp" -#include "optimization_intent.hpp" - -#include -#include - -namespace infinicore::analyzer { - -/// IntentGenerator — the core "mutual awareness" logic. -/// -/// This is where task demand and resource supply are jointly -/// analyzed to produce an OptimizationIntent. It implements -/// the key insight: the same task phase has different optimization -/// needs under different resource conditions, and the same resource -/// state has different supply value under different task phases. -class IntentGenerator { -public: - IntentGenerator() = default; - - /// Generate the global semantic intent from phase detection - /// result and op trace window. - GlobalSemanticIntent generateGlobal( - PhaseType phase, - const std::vector &window, - const std::vector &device_intents) const { - - GlobalSemanticIntent intent; - intent.current_phase = phase; - intent.timestamp_ns = OpTraceEntry::now(); - - if (!window.empty()) { - intent.op_window_start = 0; - intent.op_window_end = static_cast(window.size()); - } - - // --- Compute intensity estimation --- - intent.compute_intensity = estimateComputeIntensity(phase, window); - - // --- Determine primary bottleneck (mutual awareness) --- - intent.primary_bottleneck = determineGlobalBottleneck(phase, device_intents); - - // --- Set optimization goal based on phase + bottleneck --- - intent.goal = determineGoal(phase, intent.primary_bottleneck); - - // --- Generate strategy hints --- - intent.strategy = generateStrategy(phase, intent.primary_bottleneck, device_intents); - - // --- Confidence --- - intent.confidence = computeConfidence(phase, window); - - return intent; - } - - /// Build the complete two-layer OptimizationIntent. - OptimizationIntent generate( - PhaseType phase, - const std::vector &window, - const std::vector &device_intents) const { - - OptimizationIntent result; - result.global = generateGlobal(phase, window, device_intents); - result.per_device = device_intents; - return result; - } - -private: - /// Estimate compute intensity (higher = more compute-heavy). - /// Uses a simple heuristic based on op type composition. - float estimateComputeIntensity( - PhaseType phase, - const std::vector &window) const { - - if (window.empty()) { - return 0.0f; - } - - size_t heavy_compute_ops = 0; - for (auto &e : window) { - if (isGemmMlpOp(e.op_type) || isAttentionOp(e.op_type)) { - heavy_compute_ops++; - } - } - return static_cast(heavy_compute_ops) / static_cast(window.size()); - } - - /// Determine global bottleneck by jointly considering phase and - /// per-device resource state (the core mutual awareness logic). - BottleneckType determineGlobalBottleneck( - PhaseType phase, - const std::vector &device_intents) const { - - bool any_memory_bound = false; - bool any_compute_bound = false; - bool any_bandwidth_bound = false; - bool any_communication_bound = false; - for (auto &d : device_intents) { - any_memory_bound = any_memory_bound || d.local_bottleneck == BottleneckType::MEMORY_BOUND; - any_compute_bound = any_compute_bound || d.local_bottleneck == BottleneckType::COMPUTE_BOUND; - any_bandwidth_bound = any_bandwidth_bound || d.local_bottleneck == BottleneckType::BANDWIDTH_BOUND; - any_communication_bound = any_communication_bound || d.local_bottleneck == BottleneckType::COMMUNICATION_BOUND; - } - - // --- Mutual awareness logic --- - // The same resource state has different "supply value" depending on phase: - - if (any_memory_bound) { - return BottleneckType::MEMORY_BOUND; - } - - if (phase == PhaseType::COMMUNICATION || any_communication_bound) { - return BottleneckType::COMMUNICATION_BOUND; - } - - switch (phase) { - case PhaseType::ATTENTION_DENSE: - case PhaseType::PREFILL: - // Attention/prefill is dominated by memory movement and KV access, - // so phase semantics should win unless memory/communication already - // forced an earlier return above. - if (any_bandwidth_bound) { - return BottleneckType::BANDWIDTH_BOUND; - } - return BottleneckType::BANDWIDTH_BOUND; - - case PhaseType::GEMM_MLP_DENSE: - if (any_compute_bound) { - return BottleneckType::COMPUTE_BOUND; - } - if (any_bandwidth_bound) { - return BottleneckType::BANDWIDTH_BOUND; - } - return BottleneckType::COMPUTE_BOUND; - - case PhaseType::DECODE: - if (any_bandwidth_bound) { - return BottleneckType::BANDWIDTH_BOUND; - } - if (any_compute_bound) { - return BottleneckType::COMPUTE_BOUND; - } - return BottleneckType::BANDWIDTH_BOUND; - - case PhaseType::KV_CACHE: - if (any_bandwidth_bound) { - return BottleneckType::BANDWIDTH_BOUND; - } - return BottleneckType::MEMORY_BOUND; - - default: - if (any_bandwidth_bound) { - return BottleneckType::BANDWIDTH_BOUND; - } - if (any_compute_bound) { - return BottleneckType::COMPUTE_BOUND; - } - return BottleneckType::BALANCED; - } - } - - /// Determine optimization goal based on phase and bottleneck. - OptimizationGoal determineGoal( - PhaseType phase, - BottleneckType bottleneck) const { - - // Under memory pressure, prioritize memory safety - if (bottleneck == BottleneckType::MEMORY_BOUND) { - return OptimizationGoal::MEMORY_SAFE; - } - - if (bottleneck == BottleneckType::COMMUNICATION_BOUND) { - return OptimizationGoal::STABILITY_FIRST; - } - - switch (phase) { - case PhaseType::DECODE: - // Decode latency is user-facing → latency first - return OptimizationGoal::LATENCY_FIRST; - - case PhaseType::PREFILL: - // Prefill processes a full prompt → throughput first - return OptimizationGoal::THROUGHPUT_FIRST; - - case PhaseType::ATTENTION_DENSE: - return OptimizationGoal::LATENCY_FIRST; - - case PhaseType::GEMM_MLP_DENSE: - return OptimizationGoal::THROUGHPUT_FIRST; - - default: - return OptimizationGoal::LATENCY_FIRST; - } - } - - /// Generate strategy hints from phase + bottleneck + resources. - StrategyHint generateStrategy( - PhaseType phase, - BottleneckType bottleneck, - const std::vector &device_intents) const { - - StrategyHint hint; - - // Fusion is beneficial for bandwidth-bound phases (reduce memory traffic) - hint.prefer_fused_ops = (bottleneck == BottleneckType::BANDWIDTH_BOUND) - || phase == PhaseType::DECODE; - - // In-place when memory is tight - hint.prefer_in_place = (bottleneck == BottleneckType::MEMORY_BOUND); - - // Recomputation (activation checkpointing) when memory is critical - bool extreme_memory = false; - for (auto &d : device_intents) { - if (d.memory_usage_ratio >= 0.95f) { - extreme_memory = true; - break; - } - } - hint.prefer_recomputation = extreme_memory; - - // Async comm overlap for multi-device and communication phases - hint.prefer_async_comm = (device_intents.size() > 1) - && (phase == PhaseType::GEMM_MLP_DENSE - || phase == PhaseType::COMMUNICATION); - - return hint; - } - - /// Compute confidence based on how clear the phase signal is. - float computeConfidence( - PhaseType phase, - const std::vector &window) const { - - if (window.empty() || phase == PhaseType::UNKNOWN) { - return 0.0f; - } - - // Count how many ops in the window match the detected phase - size_t matching = 0; - for (auto &e : window) { - bool match = false; - switch (phase) { - case PhaseType::ATTENTION_DENSE: - case PhaseType::PREFILL: - match = isAttentionOp(e.op_type); - break; - case PhaseType::GEMM_MLP_DENSE: - match = isGemmMlpOp(e.op_type) || isActivationOp(e.op_type); - break; - case PhaseType::KV_CACHE: - match = isKvCacheOp(e.op_type); - break; - case PhaseType::DECODE: - match = isAttentionOp(e.op_type) || isGemmMlpOp(e.op_type); - break; - default: - break; - } - if (match) { - matching++; - } - } - - return static_cast(matching) / static_cast(window.size()); - } -}; - -} // namespace infinicore::analyzer diff --git a/include/infinicore/analyzer/mutual_awareness_analyzer.hpp b/include/infinicore/analyzer/mutual_awareness_analyzer.hpp deleted file mode 100644 index d373da694..000000000 --- a/include/infinicore/analyzer/mutual_awareness_analyzer.hpp +++ /dev/null @@ -1,118 +0,0 @@ -#pragma once - -#include "intent_generator.hpp" -#include "op_trace.hpp" -#include "optimization_intent.hpp" -#include "phase_detector.hpp" -#include "resource_sensor.hpp" - -#include -#include - -namespace infinicore::analyzer { - -/// MutualAwarenessAnalyzer — the top-level facade for the -/// hardware-task mutual awareness requirements analysis module. -/// -/// This is the primary entry point exposed to external frameworks -/// (e.g., InfiniLM) via C++ function calls. It orchestrates: -/// 1. Op trace collection (via OpTraceRing) -/// 2. Phase detection (via PhaseDetector) -/// 3. Resource sensing (via ResourceSensor) -/// 4. Intent generation (via IntentGenerator) -/// -/// Usage: -/// auto& analyzer = MutualAwarenessAnalyzer::instance(); -/// // ... ops execute and get traced automatically ... -/// auto intent = analyzer.analyze(); // Produces OptimizationIntent -/// -/// Thread safety: analyze() is safe to call from any thread. -/// The analyzer reads a snapshot of the op trace ring. -class MutualAwarenessAnalyzer { -public: - /// Get the singleton instance. - static MutualAwarenessAnalyzer &instance(); - - // Non-copyable, non-movable - MutualAwarenessAnalyzer(const MutualAwarenessAnalyzer &) = delete; - MutualAwarenessAnalyzer &operator=(const MutualAwarenessAnalyzer &) = delete; - - /// Main analysis entry point. - /// Analyzes the current op trace window + resource state - /// and returns a complete OptimizationIntent. - /// - /// This is the function InfiniLM should call. - /// Latency: expected < 1ms for MVP rule-based analysis. - OptimizationIntent analyze(); - - /// Analyze with explicitly provided memory stats per device. - /// Use this when the caller can provide resource info directly. - OptimizationIntent analyze(const std::vector> &device_stats); - - /// Analyze with explicitly provided device resource snapshots. - /// This is the richer input path used by demand-analysis-oriented callers. - OptimizationIntent analyze(const std::vector &device_snapshots); - - /// Get the current phase without generating full intent. - /// Lightweight query for simple use cases. - PhaseType getCurrentPhase() const; - - /// Get the current optimization goal derived from the - /// latest analyzer result. - OptimizationGoal getCurrentOptimizationGoal() const; - - /// Get the most recent OptimizationIntent (cached from last analyze()). - const OptimizationIntent &lastIntent() const; - - /// Access the underlying components for configuration. - PhaseDetector &phaseDetector() { return phase_detector_; } - ResourceSensor &resourceSensor() { return resource_sensor_; } - OpTraceRing &opTrace() { return getGlobalOpTrace(); } - - /// Enable / disable the analyzer. - /// When disabled, analyze() returns a default intent and - /// op trace recording is skipped. - void setEnabled(bool enabled) { enabled_ = enabled; } - bool isEnabled() const { return enabled_; } - - /// Graph recording support: when graph recording stops, - /// analyze the recorded op sequence once and cache the result. - /// Subsequent calls return the cached intent without re-analysis. - void onGraphRecordingStop(); - void clearGraphCache(); - -private: - MutualAwarenessAnalyzer(); - - PhaseDetector phase_detector_; - ResourceSensor resource_sensor_; - IntentGenerator intent_generator_; - - OptimizationIntent last_intent_; - mutable std::mutex mutex_; - - bool enabled_ = true; - - // Graph recording cache - bool graph_intent_cached_ = false; - OptimizationIntent graph_cached_intent_; -}; - -// ============================================================ -// C-style API for external framework integration (e.g., InfiniLM) -// ============================================================ - -/// Analyze current state and return an OptimizationIntent. -/// This is the simplest API for external frameworks to call. -OptimizationIntent analyzeCurrentState(); - -/// Get the current detected phase. -PhaseType getCurrentPhase(); - -/// Get the current optimization goal. -OptimizationGoal getCurrentOptimizationGoal(); - -/// Enable / disable the mutual awareness analyzer. -void setAnalyzerEnabled(bool enabled); - -} // namespace infinicore::analyzer diff --git a/include/infinicore/analyzer/op_trace.hpp b/include/infinicore/analyzer/op_trace.hpp deleted file mode 100644 index 0b3018c28..000000000 --- a/include/infinicore/analyzer/op_trace.hpp +++ /dev/null @@ -1,160 +0,0 @@ -#pragma once - -#include "op_type.hpp" - -#include -#include -#include -#include -#include -#include - -namespace infinicore::analyzer { - -// ============================================================ -// OpTraceEntry — a single recorded op invocation -// ============================================================ - -/// Compact record of one operator invocation for phase detection. -/// Designed to be small (~80 bytes) and cheap to fill. -struct OpTraceEntry { - OpType op_type = OpType::UNKNOWN; - - // Tensor shape summary (up to 4 dims for the primary input tensor) - static constexpr size_t MAX_DIMS = 4; - uint32_t ndim = 0; - uint32_t shape[MAX_DIMS] = {}; - - // Data type of the primary input (cast from infiniDtype_t) - uint8_t dtype = 0; - - // Device info - uint8_t device_type = 0; - int8_t device_id = -1; - - // Timestamp (nanoseconds since epoch, from steady_clock) - uint64_t timestamp_ns = 0; - - /// Fill shape from a shape vector. - void setShape(const size_t *dims, size_t n) { - ndim = static_cast(n > MAX_DIMS ? MAX_DIMS : n); - for (uint32_t i = 0; i < ndim; ++i) { - shape[i] = static_cast(dims[i]); - } - } - - /// Get current steady_clock timestamp in nanoseconds. - static uint64_t now() { - return static_cast( - std::chrono::steady_clock::now().time_since_epoch().count()); - } -}; - -// ============================================================ -// OpTraceRing — lock-free ring buffer for op trace entries -// ============================================================ - -/// A fixed-capacity ring buffer for OpTraceEntry. -/// Single-producer (op execution thread) friendly. -/// Reader can safely read a snapshot via getRecentEntries(). -/// -/// Thread safety: -/// - write() is safe to call from the single producer thread -/// (typical in InfiniCore where ops are dispatched on one thread). -/// - getRecentEntries() takes a snapshot and is safe to call from -/// any thread (may see a partially written entry at the boundary, -/// which is acceptable for heuristic phase detection). -class OpTraceRing { -public: - static constexpr size_t DEFAULT_CAPACITY = 256; - - explicit OpTraceRing(size_t capacity = DEFAULT_CAPACITY) - : capacity_(capacity), - entries_(capacity), - write_pos_(0), - total_count_(0) { - } - - /// Record a new op trace entry. - void write(const OpTraceEntry &entry) { - size_t pos = write_pos_.load(std::memory_order_relaxed); - entries_[pos % capacity_] = entry; - write_pos_.store(pos + 1, std::memory_order_release); - total_count_.fetch_add(1, std::memory_order_relaxed); - } - - /// Get the most recent N entries (ordered oldest to newest). - /// Returns fewer entries if the ring hasn't filled up yet. - std::vector getRecentEntries(size_t n) const { - size_t wp = write_pos_.load(std::memory_order_acquire); - size_t available = wp < capacity_ ? wp : capacity_; - size_t count = n < available ? n : available; - - std::vector result; - result.reserve(count); - - // Read from (wp - count) to (wp - 1) - for (size_t i = wp - count; i < wp; ++i) { - result.push_back(entries_[i % capacity_]); - } - return result; - } - - /// Get all valid entries in the ring (ordered oldest to newest). - std::vector getAllEntries() const { - return getRecentEntries(capacity_); - } - - /// Total number of ops traced since creation. - size_t totalCount() const { - return total_count_.load(std::memory_order_relaxed); - } - - /// Current number of valid entries in the ring. - size_t size() const { - size_t wp = write_pos_.load(std::memory_order_relaxed); - return wp < capacity_ ? wp : capacity_; - } - - /// Ring capacity. - size_t capacity() const { return capacity_; } - - /// Clear all entries. - void clear() { - write_pos_.store(0, std::memory_order_relaxed); - total_count_.store(0, std::memory_order_relaxed); - } - -private: - size_t capacity_; - std::vector entries_; - std::atomic write_pos_; - std::atomic total_count_; -}; - -// ============================================================ -// Global OpTrace singleton access -// ============================================================ - -/// Get the global OpTraceRing instance. -/// This is the primary entry point for recording op traces. -OpTraceRing &getGlobalOpTrace(); - -/// Record an op invocation to the global trace ring. -/// This is the function called from the INFINICORE_GRAPH_OP_RECORD_OR_RUN -/// macro hook (when ENABLE_MUTUAL_AWARENESS is defined). -inline void traceOp(OpType op_type, - const size_t *shape, size_t ndim, - uint8_t dtype, - uint8_t device_type, int8_t device_id) { - OpTraceEntry entry; - entry.op_type = op_type; - entry.setShape(shape, ndim); - entry.dtype = dtype; - entry.device_type = device_type; - entry.device_id = device_id; - entry.timestamp_ns = OpTraceEntry::now(); - getGlobalOpTrace().write(entry); -} - -} // namespace infinicore::analyzer diff --git a/include/infinicore/analyzer/op_type.hpp b/include/infinicore/analyzer/op_type.hpp deleted file mode 100644 index cbb3ac102..000000000 --- a/include/infinicore/analyzer/op_type.hpp +++ /dev/null @@ -1,238 +0,0 @@ -#pragma once - -#include -#include - -namespace infinicore::analyzer { - -/// Op types recognized by the analyzer for phase detection. -/// Not an exhaustive list of all infiniop ops — only the ones -/// relevant for phase classification in LLM inference. -enum class OpType : uint8_t { - UNKNOWN = 0, - - // --- Attention-related --- - ATTENTION, - FLASH_ATTENTION, - CAUSAL_SOFTMAX, - PAGED_ATTENTION, - PAGED_ATTENTION_PREFILL, - MHA_KVCACHE, - MHA_VARLEN, - SOFTMAX, - - // --- GEMM / MLP --- - GEMM, - LINEAR, - MATMUL, - INT8_GEMM, - SCALED_MM_I8, - - // --- Activation --- - SILU, - SILU_AND_MUL, - GELU, - SWIGLU, - RELU, - SIGMOID, - - // --- Norm --- - RMS_NORM, - ADD_RMS_NORM, - LAYER_NORM, - - // --- Embedding / Positional --- - EMBEDDING, - ROPE, - - // --- KV Cache --- - KV_CACHING, - PAGED_CACHING, - - // --- Elementwise / Reduce --- - ADD, - MUL, - SUB, - SUM, - RECIPROCAL, - - // --- Quantization --- - PER_TENSOR_QUANT_I8, - PER_TENSOR_DEQUANT_I8, - PER_CHANNEL_QUANT_I8, - DEQUANTIZE_AWQ, - DEQUANTIZE_GPTQ, - - // --- Sampling --- - RANDOM_SAMPLE, - TOPK, - TOPK_ROUTER, - TOPK_SOFTMAX, - - // --- Communication (future) --- - ALLREDUCE, - - // --- Misc --- - REARRANGE, - ONES, - ZEROS, - TAKE, - - OP_TYPE_COUNT, -}; - -/// Convert OpType to human-readable string. -inline const char *opTypeToString(OpType type) { - switch (type) { - case OpType::ATTENTION: - return "attention"; - case OpType::FLASH_ATTENTION: - return "flash_attention"; - case OpType::CAUSAL_SOFTMAX: - return "causal_softmax"; - case OpType::PAGED_ATTENTION: - return "paged_attention"; - case OpType::PAGED_ATTENTION_PREFILL: - return "paged_attention_prefill"; - case OpType::MHA_KVCACHE: - return "mha_kvcache"; - case OpType::MHA_VARLEN: - return "mha_varlen"; - case OpType::SOFTMAX: - return "softmax"; - case OpType::GEMM: - return "gemm"; - case OpType::LINEAR: - return "linear"; - case OpType::MATMUL: - return "matmul"; - case OpType::INT8_GEMM: - return "int8_gemm"; - case OpType::SCALED_MM_I8: - return "scaled_mm_i8"; - case OpType::SILU: - return "silu"; - case OpType::SILU_AND_MUL: - return "silu_and_mul"; - case OpType::GELU: - return "gelu"; - case OpType::SWIGLU: - return "swiglu"; - case OpType::RELU: - return "relu"; - case OpType::SIGMOID: - return "sigmoid"; - case OpType::RMS_NORM: - return "rms_norm"; - case OpType::ADD_RMS_NORM: - return "add_rms_norm"; - case OpType::LAYER_NORM: - return "layer_norm"; - case OpType::EMBEDDING: - return "embedding"; - case OpType::ROPE: - return "rope"; - case OpType::KV_CACHING: - return "kv_caching"; - case OpType::PAGED_CACHING: - return "paged_caching"; - case OpType::ADD: - return "add"; - case OpType::MUL: - return "mul"; - case OpType::SUB: - return "sub"; - case OpType::SUM: - return "sum"; - case OpType::RECIPROCAL: - return "reciprocal"; - case OpType::PER_TENSOR_QUANT_I8: - return "per_tensor_quant_i8"; - case OpType::PER_TENSOR_DEQUANT_I8: - return "per_tensor_dequant_i8"; - case OpType::PER_CHANNEL_QUANT_I8: - return "per_channel_quant_i8"; - case OpType::DEQUANTIZE_AWQ: - return "dequantize_awq"; - case OpType::DEQUANTIZE_GPTQ: - return "dequantize_gptq"; - case OpType::RANDOM_SAMPLE: - return "random_sample"; - case OpType::TOPK: - return "topk"; - case OpType::TOPK_ROUTER: - return "topk_router"; - case OpType::TOPK_SOFTMAX: - return "topk_softmax"; - case OpType::ALLREDUCE: - return "allreduce"; - case OpType::REARRANGE: - return "rearrange"; - case OpType::ONES: - return "ones"; - case OpType::ZEROS: - return "zeros"; - case OpType::TAKE: - return "take"; - default: - return "unknown"; - } -} - -/// Check if an op type belongs to the attention family. -inline bool isAttentionOp(OpType type) { - switch (type) { - case OpType::ATTENTION: - case OpType::FLASH_ATTENTION: - case OpType::CAUSAL_SOFTMAX: - case OpType::PAGED_ATTENTION: - case OpType::PAGED_ATTENTION_PREFILL: - case OpType::MHA_KVCACHE: - case OpType::MHA_VARLEN: - return true; - default: - return false; - } -} - -/// Check if an op type belongs to the GEMM/MLP family. -inline bool isGemmMlpOp(OpType type) { - switch (type) { - case OpType::GEMM: - case OpType::LINEAR: - case OpType::MATMUL: - case OpType::INT8_GEMM: - case OpType::SCALED_MM_I8: - return true; - default: - return false; - } -} - -/// Check if an op type is an activation function. -inline bool isActivationOp(OpType type) { - switch (type) { - case OpType::SILU: - case OpType::SILU_AND_MUL: - case OpType::GELU: - case OpType::SWIGLU: - case OpType::RELU: - case OpType::SIGMOID: - return true; - default: - return false; - } -} - -/// Check if an op type is KV cache related. -inline bool isKvCacheOp(OpType type) { - switch (type) { - case OpType::KV_CACHING: - case OpType::PAGED_CACHING: - return true; - default: - return false; - } -} - -} // namespace infinicore::analyzer diff --git a/include/infinicore/analyzer/op_type_registry.hpp b/include/infinicore/analyzer/op_type_registry.hpp deleted file mode 100644 index 4445f3a51..000000000 --- a/include/infinicore/analyzer/op_type_registry.hpp +++ /dev/null @@ -1,49 +0,0 @@ -#pragma once - -#include "op_type.hpp" -#include -#include - -namespace infinicore::analyzer { - -/// Centralized class-name → OpType registry. -/// New ops only need one line added here — no changes to op headers. -inline OpType opTypeFromName(const char *name) { - static const std::unordered_map registry = { - // Attention - {"FlashAttention", OpType::FLASH_ATTENTION}, - {"CausalSoftmax", OpType::CAUSAL_SOFTMAX}, - {"PagedAttention", OpType::PAGED_ATTENTION}, - {"MhaKVCache", OpType::MHA_KVCACHE}, - {"MultiheadAttentionVarlen", OpType::MHA_VARLEN}, - // GEMM / MLP - {"Gemm", OpType::GEMM}, - {"I8Gemm", OpType::SCALED_MM_I8}, - // Activation - {"SiluAndMul", OpType::SILU_AND_MUL}, - {"SwiGLU", OpType::SWIGLU}, - // Norm - {"RMSNorm", OpType::RMS_NORM}, - {"AddRMSNorm", OpType::ADD_RMS_NORM}, - // Embedding / Positional - {"Embedding", OpType::EMBEDDING}, - {"RoPE", OpType::ROPE}, - // KV Cache - {"KVCaching", OpType::KV_CACHING}, - {"PagedCaching", OpType::PAGED_CACHING}, - // Elementwise - {"Add", OpType::ADD}, - {"Mul", OpType::MUL}, - // Quantization - {"PerTensorQuantI8", OpType::PER_TENSOR_QUANT_I8}, - {"PerTensorDequantI8", OpType::PER_TENSOR_DEQUANT_I8}, - {"PerChannelQuantI8", OpType::PER_CHANNEL_QUANT_I8}, - {"DequantizeAWQ", OpType::DEQUANTIZE_AWQ}, - // Misc - {"Rearrange", OpType::REARRANGE}, - }; - auto it = registry.find(name); - return it != registry.end() ? it->second : OpType::UNKNOWN; -} - -} // namespace infinicore::analyzer diff --git a/include/infinicore/analyzer/optimization_intent.hpp b/include/infinicore/analyzer/optimization_intent.hpp deleted file mode 100644 index 19d6c868c..000000000 --- a/include/infinicore/analyzer/optimization_intent.hpp +++ /dev/null @@ -1,187 +0,0 @@ -#pragma once - -#include "op_type.hpp" - -#include -#include -#include - -namespace infinicore::analyzer { - -// ============================================================ -// Phase & Bottleneck & Goal enumerations -// ============================================================ - -/// Recognized task phase types in LLM inference. -enum class PhaseType : uint8_t { - UNKNOWN = 0, - PREFILL, // Long-sequence prefill phase - DECODE, // Autoregressive decode phase - ATTENTION_DENSE, // Attention-dominated segment - GEMM_MLP_DENSE, // GEMM/MLP-dominated segment - MOE_ROUTING, // MoE routing segment (future) - KV_CACHE, // KV cache manipulation segment - COMMUNICATION, // Communication-dominated segment (future) -}; - -inline const char *phaseTypeToString(PhaseType phase) { - switch (phase) { - case PhaseType::PREFILL: - return "prefill"; - case PhaseType::DECODE: - return "decode"; - case PhaseType::ATTENTION_DENSE: - return "attention_dense"; - case PhaseType::GEMM_MLP_DENSE: - return "gemm_mlp_dense"; - case PhaseType::MOE_ROUTING: - return "moe_routing"; - case PhaseType::KV_CACHE: - return "kv_cache"; - case PhaseType::COMMUNICATION: - return "communication"; - default: - return "unknown"; - } -} - -/// Primary bottleneck type of the current phase. -enum class BottleneckType : uint8_t { - COMPUTE_BOUND = 0, - MEMORY_BOUND, - BANDWIDTH_BOUND, - COMMUNICATION_BOUND, - BALANCED, -}; - -inline const char *bottleneckTypeToString(BottleneckType bt) { - switch (bt) { - case BottleneckType::COMPUTE_BOUND: - return "compute_bound"; - case BottleneckType::MEMORY_BOUND: - return "memory_bound"; - case BottleneckType::BANDWIDTH_BOUND: - return "bandwidth_bound"; - case BottleneckType::COMMUNICATION_BOUND: - return "communication_bound"; - case BottleneckType::BALANCED: - return "balanced"; - default: - return "unknown"; - } -} - -/// Optimization goal / objective preference. -enum class OptimizationGoal : uint8_t { - LATENCY_FIRST = 0, - THROUGHPUT_FIRST, - MEMORY_SAFE, - STABILITY_FIRST, -}; - -inline const char *optimizationGoalToString(OptimizationGoal goal) { - switch (goal) { - case OptimizationGoal::LATENCY_FIRST: - return "latency_first"; - case OptimizationGoal::THROUGHPUT_FIRST: - return "throughput_first"; - case OptimizationGoal::MEMORY_SAFE: - return "memory_safe"; - case OptimizationGoal::STABILITY_FIRST: - return "stability_first"; - default: - return "unknown"; - } -} - -// ============================================================ -// Two-layer OptimizationIntent structure -// -// Design: Global Semantic Intent + Per-Device Local Refinement -// Rationale: Single-node multi-GPU scenario requires a unified -// task-level view PLUS per-device resource constraints. -// ============================================================ - -/// Strategy hints generated by the analyzer. -struct StrategyHint { - bool prefer_fused_ops = false; // Suggest op fusion? - bool prefer_in_place = false; // Suggest in-place computation? - bool prefer_recomputation = false; // Suggest recompute to save memory? - bool prefer_async_comm = false; // Suggest async comm overlap? -}; - -/// Global semantic-level intent — describes the task's overall -/// compute characteristics and optimization direction. -struct GlobalSemanticIntent { - // --- Current phase --- - PhaseType current_phase = PhaseType::UNKNOWN; - - // --- Primary bottleneck --- - BottleneckType primary_bottleneck = BottleneckType::BALANCED; - - // --- Optimization goal --- - OptimizationGoal goal = OptimizationGoal::LATENCY_FIRST; - - // --- Compute intensity estimate (ops/byte) --- - float compute_intensity = 0.0f; - - // --- Analysis confidence (0.0~1.0) --- - float confidence = 0.0f; - - // --- Strategy hints --- - StrategyHint strategy; - - // --- Production metadata --- - uint64_t timestamp_ns = 0; - uint32_t op_window_start = 0; // Start index of analyzed op window - uint32_t op_window_end = 0; // End index (exclusive) of analyzed op window -}; - -/// Per-device local intent — resource-aware refinement -/// for a specific device in multi-GPU scenario. -struct DeviceLocalIntent { - int device_id = -1; - - // Memory utilization (0.0 ~ 1.0) - float memory_usage_ratio = 0.0f; - - // Available memory in bytes - size_t memory_available_bytes = 0; - - // Device-local bottleneck (may differ from global) - BottleneckType local_bottleneck = BottleneckType::BALANCED; - - // Additional resource observations from the device snapshot - float compute_utilization = 0.0f; - float memory_bandwidth_utilization = 0.0f; - float communication_time_ratio = 0.0f; - - // Confidence of the resource-side classification (0.0 ~ 1.0) - float resource_confidence = 0.0f; -}; - -/// The complete OptimizationIntent: global task semantics -/// combined with per-device resource views. -/// -/// Produced by: MutualAwarenessAnalyzer -/// Consumed by (roadmap): -/// 1. MVP: query/record/validation (Python + C++ API) -/// 2. First: OpDispatcher / kernel selection -/// 3. Second: Allocator strategy (conservative vs aggressive) -/// 4. Third: Communication strategy (overlap, buffer, collective) -struct OptimizationIntent { - GlobalSemanticIntent global; - std::vector per_device; - - /// Get the local intent for a specific device, or nullptr if not found. - const DeviceLocalIntent *getDeviceIntent(int device_id) const { - for (auto &d : per_device) { - if (d.device_id == device_id) { - return &d; - } - } - return nullptr; - } -}; - -} // namespace infinicore::analyzer diff --git a/include/infinicore/analyzer/phase_detector.hpp b/include/infinicore/analyzer/phase_detector.hpp deleted file mode 100644 index 2f4828c6e..000000000 --- a/include/infinicore/analyzer/phase_detector.hpp +++ /dev/null @@ -1,126 +0,0 @@ -#pragma once - -#include "op_trace.hpp" -#include "optimization_intent.hpp" - -#include -#include - -namespace infinicore::analyzer { - -/// PhaseDetector — detects the current task phase from the -/// recent op trace window using rule-based pattern matching. -/// -/// Design choice (MVP): Fixed rule matching based on op type -/// composition in a sliding window. Will evolve to support -/// offline-generated phase templates in future iterations. -class PhaseDetector { -public: - /// Configuration for phase detection thresholds. - struct Config { - size_t window_size; - float attention_threshold; - float gemm_mlp_threshold; - float kv_cache_threshold; - uint32_t decode_seq_len_max; - uint32_t prefill_seq_len_min; - - Config() - : window_size(16), - attention_threshold(0.3f), - gemm_mlp_threshold(0.3f), - kv_cache_threshold(0.4f), - decode_seq_len_max(4), - prefill_seq_len_min(32) {} - }; - - explicit PhaseDetector(Config config = {}) : config_(config) {} - - /// Detect the current phase from a window of recent op traces. - PhaseType detect(const std::vector &window) const { - if (window.empty()) { - return PhaseType::UNKNOWN; - } - - // Count op categories in the window - size_t attention_count = 0; - size_t gemm_mlp_count = 0; - size_t kv_cache_count = 0; - size_t activation_count = 0; - size_t total = window.size(); - - // Track shape info for prefill/decode inference - uint32_t max_seq_len = 0; - uint32_t min_seq_len = UINT32_MAX; - bool has_attention_shape = false; - - for (auto &entry : window) { - if (isAttentionOp(entry.op_type)) { - attention_count++; - // For attention ops, shape[1] or shape[2] typically indicates seq_len - if (entry.ndim >= 2) { - // Heuristic: for attention-like ops, look at the sequence dimension - // Typically shape = [batch, seq_len, ...] or [batch, heads, seq_len, ...] - uint32_t seq_dim = (entry.ndim >= 3) ? entry.shape[2] : entry.shape[1]; - max_seq_len = std::max(max_seq_len, seq_dim); - min_seq_len = std::min(min_seq_len, seq_dim); - has_attention_shape = true; - } - } else if (isGemmMlpOp(entry.op_type)) { - gemm_mlp_count++; - } else if (isKvCacheOp(entry.op_type)) { - kv_cache_count++; - } else if (isActivationOp(entry.op_type)) { - activation_count++; - } - } - - float attention_ratio = static_cast(attention_count) / total; - float gemm_mlp_ratio = static_cast(gemm_mlp_count + activation_count) / total; - float kv_cache_ratio = static_cast(kv_cache_count) / total; - - // --- Phase classification --- - - // KV cache phase (high KV cache op ratio) - if (kv_cache_ratio >= config_.kv_cache_threshold) { - return PhaseType::KV_CACHE; - } - - // --- Prefill vs Decode inference from shape --- - // Self-inferred from sequence length, no external flags needed. - if (has_attention_shape) { - if (max_seq_len <= config_.decode_seq_len_max) { - return PhaseType::DECODE; - } - if (min_seq_len >= config_.prefill_seq_len_min) { - return PhaseType::PREFILL; - } - } - - // Attention-dense phase - if (attention_ratio >= config_.attention_threshold && attention_ratio >= gemm_mlp_ratio) { - return PhaseType::ATTENTION_DENSE; - } - - // GEMM/MLP-dense phase (include activation ops as co-indicators) - if (gemm_mlp_ratio >= config_.gemm_mlp_threshold && gemm_mlp_ratio >= attention_ratio) { - return PhaseType::GEMM_MLP_DENSE; - } - - return PhaseType::UNKNOWN; - } - - /// Convenience: detect from the global trace ring. - PhaseType detectFromTrace(const OpTraceRing &trace) const { - auto window = trace.getRecentEntries(config_.window_size); - return detect(window); - } - - const Config &config() const { return config_; } - void setConfig(Config config) { config_ = config; } - -private: - Config config_; -}; - -} // namespace infinicore::analyzer diff --git a/include/infinicore/analyzer/resource_sensor.hpp b/include/infinicore/analyzer/resource_sensor.hpp deleted file mode 100644 index 72c871080..000000000 --- a/include/infinicore/analyzer/resource_sensor.hpp +++ /dev/null @@ -1,181 +0,0 @@ -#pragma once - -#include "infinicore/device.hpp" -#include "optimization_intent.hpp" - -#include -#include -#include - -namespace infinicore::analyzer { - -/// Memory statistics from the allocator. -struct MemoryStats { - size_t allocated_bytes = 0; // Currently allocated bytes - size_t total_capacity = 0; // Total pool capacity in bytes - size_t peak_allocated = 0; // Peak allocation since last reset - size_t allocation_count = 0; // Number of active allocations - - float usageRatio() const { - return total_capacity > 0 - ? static_cast(allocated_bytes) / static_cast(total_capacity) - : 0.0f; - } -}; - -/// A normalized resource image consumed by the analyzer. -/// -/// This is intentionally vendor-neutral. Backend-specific runtimes -/// populate these fields through infinirt, and the analyzer consumes -/// only this normalized view. -struct DeviceResourceSnapshot { - int device_id = -1; - Device::Type device_type = Device::Type::CPU; - - bool has_memory_capacity = false; - bool has_compute_utilization = false; - bool has_memory_bandwidth_utilization = false; - bool has_kernel_time_ratio = false; - bool has_communication = false; - bool kernel_time_estimated = false; - - size_t free_bytes = 0; - size_t total_bytes = 0; - size_t used_bytes = 0; - size_t reserved_bytes = 0; - - float compute_utilization = 0.0f; - float memory_bandwidth_utilization = 0.0f; - - uint64_t bytes_read = 0; - uint64_t bytes_written = 0; - - float kernel_time_ratio = 0.0f; - float idle_time_ratio = 0.0f; - - float communication_time_ratio = 0.0f; - uint64_t communication_bytes = 0; - - float load_imbalance_score = 0.0f; - - float memoryUsageRatio() const { - if (total_bytes > 0) { - return static_cast(used_bytes) / static_cast(total_bytes); - } - if (reserved_bytes > 0 && used_bytes <= reserved_bytes) { - return static_cast(used_bytes) / static_cast(reserved_bytes); - } - return 0.0f; - } - - float resourceConfidence() const { - float confidence = 0.0f; - if (has_memory_capacity) { - confidence += 0.35f; - } - if (has_compute_utilization) { - confidence += 0.25f; - } - if (has_memory_bandwidth_utilization) { - confidence += 0.25f; - } - if (has_communication) { - confidence += 0.15f; - } - return std::min(confidence, 1.0f); - } -}; - -/// ResourceSensor — gathers current resource state from -/// the runtime and allocator subsystems. -/// -/// This is the "resource supply" side of the mutual-awareness -/// equation. It aggregates device type, memory stats, and -/// potentially timing info into a resource snapshot. -class ResourceSensor { -public: - ResourceSensor() = default; - - /// Build a DeviceLocalIntent from a normalized resource image. - DeviceLocalIntent sense(const DeviceResourceSnapshot &snapshot) const { - DeviceLocalIntent intent; - intent.device_id = snapshot.device_id; - intent.memory_usage_ratio = snapshot.memoryUsageRatio(); - intent.memory_available_bytes = snapshot.free_bytes > 0 - ? snapshot.free_bytes - : ((snapshot.total_bytes >= snapshot.used_bytes) ? (snapshot.total_bytes - snapshot.used_bytes) : 0); - intent.compute_utilization = snapshot.compute_utilization; - intent.memory_bandwidth_utilization = snapshot.memory_bandwidth_utilization; - intent.communication_time_ratio = snapshot.communication_time_ratio; - intent.resource_confidence = snapshot.resourceConfidence(); - - if (snapshot.has_communication && snapshot.communication_time_ratio > high_communication_threshold_) { - intent.local_bottleneck = BottleneckType::COMMUNICATION_BOUND; - return intent; - } - - if (intent.memory_usage_ratio > high_memory_threshold_) { - intent.local_bottleneck = BottleneckType::MEMORY_BOUND; - return intent; - } - - if (snapshot.has_memory_bandwidth_utilization - && snapshot.memory_bandwidth_utilization > high_bandwidth_threshold_ - && (!snapshot.has_compute_utilization - || snapshot.memory_bandwidth_utilization >= snapshot.compute_utilization + bandwidth_margin_)) { - intent.local_bottleneck = BottleneckType::BANDWIDTH_BOUND; - return intent; - } - - if (snapshot.has_compute_utilization && snapshot.compute_utilization > high_compute_threshold_) { - intent.local_bottleneck = BottleneckType::COMPUTE_BOUND; - return intent; - } - - if (intent.memory_usage_ratio > moderate_memory_threshold_) { - intent.local_bottleneck = BottleneckType::BALANCED; - return intent; - } - - // Preserve the MVP fallback: if all we know is that memory is low, - // compute is the most likely bottleneck. - intent.local_bottleneck = BottleneckType::COMPUTE_BOUND; - return intent; - } - - /// Build a DeviceLocalIntent from current resource state. - /// In MVP, this primarily queries allocator memory stats. - /// - /// @param device_id The device ID to query - /// @param stats Memory stats from the allocator - DeviceLocalIntent sense( - int device_id, - const MemoryStats &stats, - Device::Type device_type = Device::Type::CPU) const { - DeviceResourceSnapshot snapshot; - snapshot.device_id = device_id; - snapshot.device_type = device_type; - snapshot.has_memory_capacity = stats.total_capacity > 0; - snapshot.free_bytes = stats.total_capacity >= stats.allocated_bytes - ? (stats.total_capacity - stats.allocated_bytes) - : 0; - snapshot.total_bytes = stats.total_capacity; - snapshot.used_bytes = stats.allocated_bytes; - snapshot.reserved_bytes = stats.total_capacity; - return sense(snapshot); - } - - /// Thresholds for memory-based bottleneck classification. - void setHighMemoryThreshold(float t) { high_memory_threshold_ = t; } - void setModerateMemoryThreshold(float t) { moderate_memory_threshold_ = t; } - -private: - float high_memory_threshold_ = 0.85f; - float moderate_memory_threshold_ = 0.5f; - float high_compute_threshold_ = 0.75f; - float high_bandwidth_threshold_ = 0.75f; - float high_communication_threshold_ = 0.25f; - float bandwidth_margin_ = 0.05f; -}; - -} // namespace infinicore::analyzer diff --git a/include/infinicore/common/LRUCache.hpp b/include/infinicore/common/LRUCache.hpp deleted file mode 100644 index 0c13d1d09..000000000 --- a/include/infinicore/common/LRUCache.hpp +++ /dev/null @@ -1,137 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include -#include - -namespace infinicore::common { -template -class LRUCache { -public: - using KeyValuePair = std::pair; - using ListIt = typename std::list::iterator; - using Destructor = std::function; - - explicit LRUCache(size_t capacity = 100, Destructor destructor = nullptr) - : capacity_(capacity), destructor_(destructor) { - if (capacity == 0) { - capacity_ = UINT64_MAX; // effectively unbounded - } - } - - ~LRUCache() { - cleanup(); - } - - bool contains(const Key &key) const { - return map_.find(key) != map_.end(); - } - - void put(const Key &key, const Value &value) { - auto it = map_.find(key); - if (it != map_.end()) { - if (destructor_) { - destructor_(it->second->second); - } - it->second->second = value; - touch(it); - } else { - // insert new - if (list_.size() >= capacity_) { - evictLRU(); - } - list_.emplace_front(key, value); - map_[key] = list_.begin(); - } - } - - std::optional get(const Key &key) { - auto it = map_.find(key); - if (it == map_.end()) { - return std::nullopt; - } - touch(it); - return it->second->second; - } - - std::optional get(const Key &key) const { - auto it = map_.find(key); - if (it == map_.end()) { - return std::nullopt; - } - // Note: can't touch in const context - return it->second->second; - } - - void setDestructor(Destructor destructor) { - destructor_ = destructor; - } - - void setCapacity(size_t capacity) { - capacity_ = capacity; - while (list_.size() > capacity_) { - evictLRU(); - } - } - - void clear() { - if (destructor_) { - for (auto &item : list_) { - safeDestruct(item.second); - } - } - list_.clear(); - map_.clear(); - } - - const std::list &getAllItems() const { - return list_; - } - -protected: - std::list list_; // front = most recent, back = least - -private: - void touch(typename std::unordered_map::iterator it) { - // move this key to front (most recent) - list_.splice(list_.begin(), list_, it->second); - it->second = list_.begin(); - } - - void safeDestruct(Value &value) { - if (!destructor_) { - return; - } - - try { - destructor_(value); - } catch (const std::exception &e) { - // Built-in default error handling - std::cerr << "Cache destructor error (type: " << typeid(Value).name() - << "): " << e.what() << std::endl; - } - } - - void evictLRU() { - if (!list_.empty()) { - auto &kv = list_.back(); - safeDestruct(kv.second); - map_.erase(kv.first); - list_.pop_back(); - } - } - - void cleanup() { - clear(); - } - - size_t capacity_; - std::unordered_map map_; - Destructor destructor_; -}; - -} // namespace infinicore::common diff --git a/include/infinicore/common/hash.hpp b/include/infinicore/common/hash.hpp deleted file mode 100644 index 65178e4db..000000000 --- a/include/infinicore/common/hash.hpp +++ /dev/null @@ -1,72 +0,0 @@ -#pragma once - -#include "../tensor.hpp" - -#include -#include - -namespace infinicore { - -// Base hash_combine for arithmetic types -template -std::enable_if_t, void> -hash_combine(size_t &seed, const T &value) { - seed ^= std::hash{}(value) + 0x9e3779b9 + (seed << 6) + (seed >> 2); -} - -// Specialization for Tensor -inline void hash_combine(size_t &seed, Tensor tensor) { - if (!tensor) { - hash_combine(seed, static_cast(0)); - return; - } - - hash_combine(seed, static_cast(tensor->dtype())); - for (Size shape : tensor->shape()) { - hash_combine(seed, shape); - } - for (Stride stride : tensor->strides()) { - hash_combine(seed, static_cast(stride)); - } -} - -// Specialization for optional -template -inline void hash_combine(size_t &seed, const std::optional &opt) { - hash_combine(seed, opt.has_value()); - if (opt) { - hash_combine(seed, *opt); - } -} - -// Specialization for std::string -inline void hash_combine(size_t &seed, const std::string &str) { - hash_combine(seed, std::hash{}(str)); -} - -// Specialization for const char* -inline void hash_combine(size_t &seed, const char *str) { - hash_combine(seed, std::string(str)); -} - -// Variadic template for multiple arguments -template -void hash_combine(size_t &seed, const First &first, const Rest &...rest) { - hash_combine(seed, first); - hash_combine(seed, rest...); -} - -// Base case for variadic template -inline void hash_combine(size_t &seed) { - // Base case - do nothing -} - -// Convenience function to hash multiple values -template -size_t hash_combine(const Types &...values) { - size_t seed = 0; - hash_combine(seed, values...); - return seed; -} - -} // namespace infinicore diff --git a/include/infinicore/context/context.hpp b/include/infinicore/context/context.hpp deleted file mode 100644 index 82320ed52..000000000 --- a/include/infinicore/context/context.hpp +++ /dev/null @@ -1,57 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../memory.hpp" - -#include "../graph/graph.hpp" - -#include -#include - -#include - -namespace infinicore { - -namespace context { -void setDevice(Device device); -Device getDevice(); -size_t getDeviceCount(Device::Type type); - -infinirtStream_t getStream(); -infiniopHandle_t getInfiniopHandle(Device device); - -void syncStream(); -void syncDevice(); -void trimMemory(); - -std::shared_ptr allocateMemory(size_t size); -std::shared_ptr allocateHostMemory(size_t size); -std::shared_ptr allocatePinnedHostMemory(size_t size); - -void memcpyH2D(void *dst, const void *src, size_t size, bool async = true); -void memcpyD2H(void *dst, const void *src, size_t size); -void memcpyD2D(void *dst, const void *src, size_t size, bool async = true); -void memcpyH2H(void *dst, const void *src, size_t size); - -void setDeviceMemory(void *ptr, int value, size_t count); -void setDeviceMemoryAsync(void *ptr, int value, size_t count, infinirtStream_t stream); - -// Timing APIs for performance measurement -infinirtEvent_t createEvent(); -infinirtEvent_t createEventWithFlags(uint32_t flags); -void recordEvent(infinirtEvent_t event, infinirtStream_t stream = nullptr); -bool queryEvent(infinirtEvent_t event); -void synchronizeEvent(infinirtEvent_t event); -void destroyEvent(infinirtEvent_t event); -float elapsedTime(infinirtEvent_t start, infinirtEvent_t end); -void streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event); - -// Graph recording APIs -bool isGraphRecording(); -void startGraphRecording(); -void addGraphOperator(std::shared_ptr op); -std::shared_ptr stopGraphRecording(); - -} // namespace context - -} // namespace infinicore diff --git a/include/infinicore/device.hpp b/include/infinicore/device.hpp deleted file mode 100644 index 4bdab1088..000000000 --- a/include/infinicore/device.hpp +++ /dev/null @@ -1,53 +0,0 @@ -#pragma once - -#include -#include - -#include "infinicore.h" - -namespace infinicore { - -class Device { -public: - using Index = std::size_t; - - enum class Type { - CPU = INFINI_DEVICE_CPU, - NVIDIA = INFINI_DEVICE_NVIDIA, - CAMBRICON = INFINI_DEVICE_CAMBRICON, - ASCEND = INFINI_DEVICE_ASCEND, - METAX = INFINI_DEVICE_METAX, - MOORE = INFINI_DEVICE_MOORE, - ILUVATAR = INFINI_DEVICE_ILUVATAR, - KUNLUN = INFINI_DEVICE_KUNLUN, - HYGON = INFINI_DEVICE_HYGON, - QY = INFINI_DEVICE_QY, - ALI = INFINI_DEVICE_ALI, - COUNT = INFINI_DEVICE_TYPE_COUNT, - }; - - Device(const Type &type = Type::CPU, const Index &index = 0); - - const Type &getType() const; - - const Index &getIndex() const; - - std::string toString() const; - - static std::string toString(const Type &type); - - bool operator==(const Device &other) const; - - bool operator!=(const Device &other) const; - - inline static Device cpu() { - return Device(Type::CPU, 0); - } - -private: - Type type_; - - Index index_; -}; - -} // namespace infinicore diff --git a/include/infinicore/device_event.hpp b/include/infinicore/device_event.hpp deleted file mode 100644 index 8ce38d645..000000000 --- a/include/infinicore/device_event.hpp +++ /dev/null @@ -1,125 +0,0 @@ -#pragma once - -#include "device.hpp" -#include "infinirt.h" -#include -#include - -namespace infinicore { - -/** - * @brief A device event for timing operations and synchronization across devices. - * - * Similar to torch.cuda.Event, this class provides functionality to: - * - Record events on specific device streams - * - Synchronize with events - * - Measure elapsed time between events - * - Query event completion status - * - Make streams wait for events - */ -class DeviceEvent { -private: - infinirtEvent_t event_; // Underlying event handle - Device device_; // Device where this event was created - bool is_recorded_; // Whether the event has been recorded - -public: - /** - * @brief Construct a new DeviceEvent on the current device. - */ - DeviceEvent(); - - /** - * @brief Construct a new DeviceEvent on the current device with specific flags. - * @param flags Event creation flags (e.g., for timing, blocking sync) - */ - explicit DeviceEvent(uint32_t flags); - - /** - * @brief Construct a new DeviceEvent on a specific device. - * @param device Target device for this event - */ - explicit DeviceEvent(Device device); - - /** - * @brief Construct a new DeviceEvent on a specific device with flags. - * @param device Target device for this event - * @param flags Event creation flags - */ - DeviceEvent(Device device, uint32_t flags); - - // Disallow copying - DeviceEvent(const DeviceEvent &) = delete; - DeviceEvent &operator=(const DeviceEvent &) = delete; - - /** - * @brief Move constructor. - */ - DeviceEvent(DeviceEvent &&other) noexcept; - - /** - * @brief Move assignment operator. - */ - DeviceEvent &operator=(DeviceEvent &&other) noexcept; - - /** - * @brief Destroy the DeviceEvent and release underlying resources. - */ - ~DeviceEvent(); - - /** - * @brief Record the event on the current stream of its device. - */ - void record(); - - /** - * @brief Record the event on a specific stream. - * @param stream Stream to record the event on - */ - void record(infinirtStream_t stream); - - /** - * @brief Wait for the event to complete (blocking). - */ - void synchronize(); - - /** - * @brief Check if the event has been completed. - * @return true if completed, false otherwise - */ - bool query() const; - - /** - * @brief Calculate elapsed time between this event and another event (in milliseconds). - * @param other The other event to compare with - * @return Elapsed time in milliseconds - * @throws std::runtime_error if events are on different devices or not recorded - */ - float elapsed_time(const DeviceEvent &other) const; - - /** - * @brief Make a stream wait for this event to complete. - * @param stream Stream to make wait for this event (nullptr for current stream) - */ - void wait(infinirtStream_t stream = nullptr) const; - - /** - * @brief Get the device where this event was created. - * @return Device associated with this event - */ - Device device() const { return device_; } - - /** - * @brief Get the underlying event handle. - * @return Raw event handle - */ - infinirtEvent_t get() const { return event_; } - - /** - * @brief Check if the event has been recorded. - * @return true if recorded, false otherwise - */ - bool is_recorded() const { return is_recorded_; } -}; - -} // namespace infinicore diff --git a/include/infinicore/dtype.hpp b/include/infinicore/dtype.hpp deleted file mode 100644 index ea3f49286..000000000 --- a/include/infinicore/dtype.hpp +++ /dev/null @@ -1,33 +0,0 @@ -#pragma once - -#include -#include - -namespace infinicore { - -enum class DataType { - BYTE = INFINI_DTYPE_BYTE, - BOOL = INFINI_DTYPE_BOOL, - I8 = INFINI_DTYPE_I8, - I16 = INFINI_DTYPE_I16, - I32 = INFINI_DTYPE_I32, - I64 = INFINI_DTYPE_I64, - U8 = INFINI_DTYPE_U8, - U16 = INFINI_DTYPE_U16, - U32 = INFINI_DTYPE_U32, - U64 = INFINI_DTYPE_U64, - F8 = INFINI_DTYPE_F8, - F16 = INFINI_DTYPE_F16, - F32 = INFINI_DTYPE_F32, - F64 = INFINI_DTYPE_F64, - C16 = INFINI_DTYPE_C16, - C32 = INFINI_DTYPE_C32, - C64 = INFINI_DTYPE_C64, - C128 = INFINI_DTYPE_C128, - BF16 = INFINI_DTYPE_BF16, -}; - -std::string toString(const DataType &dtype); -size_t dsize(const DataType &dtype); - -} // namespace infinicore diff --git a/include/infinicore/graph/graph.hpp b/include/infinicore/graph/graph.hpp deleted file mode 100644 index be368f92d..000000000 --- a/include/infinicore/graph/graph.hpp +++ /dev/null @@ -1,133 +0,0 @@ -#pragma once - -#include -#include - -#include "../tensor.hpp" - -namespace infinicore::graph { -// Forward declarations -class GraphManager; - -class GraphTensor : public Tensor { -public: - GraphTensor(const Tensor &); -}; - -class GraphOperator { -public: - virtual void run() const = 0; - virtual ~GraphOperator() = default; -}; - -class DispatchableGraphOperator : public GraphOperator { -public: - void run() const override; - ~DispatchableGraphOperator() override; - -protected: - using run_schema = void (*)(void *); - using cleanup_schema = void (*)(void **); - void *planned_meta_; - run_schema runner_; - cleanup_schema deleter_; -}; - -class Graph { -public: - Graph(); - ~Graph(); - - void run() const; - -protected: - void add_operator(std::shared_ptr op); - void instantiate(); - std::vector> op_list_; - - friend class GraphManager; - -private: - struct DeviceGraph; - std::unique_ptr device_graph_; -}; -} // namespace infinicore::graph - -#define INFINICORE_GRAPH_OP_CLASS(__OP_NAME__, ...) \ - class __OP_NAME__ : public graph::DispatchableGraphOperator { \ - public: \ - using schema = void (*)(__VA_ARGS__); \ - using plan_schema = void *(*)(__VA_ARGS__); \ - static common::OpDispatcher &plan_dispatcher(); \ - static common::OpDispatcher &run_dispatcher(); \ - static common::OpDispatcher &cleanup_dispatcher(); \ - __OP_NAME__(__VA_ARGS__); \ - static void execute(__VA_ARGS__); \ - }; - -#define INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(__OP_NAME__) \ - common::OpDispatcher<__OP_NAME__::plan_schema> &__OP_NAME__::plan_dispatcher() { \ - static common::OpDispatcher<__OP_NAME__::plan_schema> dispatcher_; \ - return dispatcher_; \ - } \ - common::OpDispatcher<__OP_NAME__::run_schema> &__OP_NAME__::run_dispatcher() { \ - static common::OpDispatcher<__OP_NAME__::run_schema> dispatcher_; \ - return dispatcher_; \ - } \ - common::OpDispatcher<__OP_NAME__::cleanup_schema> &__OP_NAME__::cleanup_dispatcher() { \ - static common::OpDispatcher<__OP_NAME__::cleanup_schema> dispatcher_; \ - return dispatcher_; \ - } - -#define INFINICORE_GRAPH_OP_DISPATCH(__DEVICE_TYPE__, ...) \ - planned_meta_ = plan_dispatcher().lookup(__DEVICE_TYPE__)(__VA_ARGS__); \ - runner_ = run_dispatcher().lookup(__DEVICE_TYPE__); \ - deleter_ = cleanup_dispatcher().lookup(__DEVICE_TYPE__); - -#define INFINICORE_DETAIL_FIRST_ARG(__FIRST__, ...) __FIRST__ - -#ifdef ENABLE_MUTUAL_AWARENESS -#include "../analyzer/op_trace.hpp" -#include "../analyzer/op_type_registry.hpp" - -// Trace one op invocation into the global ring. Op type is resolved by -// stringified class name through `opTypeFromName`, so new graph ops are -// automatically discoverable without modifying the op header. -#define _INFINICORE_TRACE_OP(__OP_NAME__, __TRACE_TENSOR__) \ - do { \ - auto __op_type = ::infinicore::analyzer::opTypeFromName(#__OP_NAME__); \ - auto &&__trace_tensor = (__TRACE_TENSOR__); \ - if (__trace_tensor) { \ - const auto &__trace_shape = __trace_tensor->shape(); \ - const auto __trace_device = __trace_tensor->device(); \ - ::infinicore::analyzer::traceOp( \ - __op_type, \ - __trace_shape.data(), \ - __trace_shape.size(), \ - static_cast(__trace_tensor->dtype()), \ - static_cast(__trace_device.getType()), \ - static_cast(__trace_device.getIndex())); \ - } else { \ - ::infinicore::analyzer::traceOp(__op_type, nullptr, 0, 0, 0, -1); \ - } \ - } while (0) -#else -#define _INFINICORE_TRACE_OP(__OP_NAME__, __TRACE_TENSOR__) ((void)0) -#endif - -#define INFINICORE_GRAPH_OP_RECORD_OR_RUN(__OP_NAME__, ...) \ - auto ___op = std::make_shared<__OP_NAME__>(__VA_ARGS__); \ - if (context::isGraphRecording()) { \ - context::addGraphOperator(___op); \ - } else { \ - ___op->run(); \ - } \ - _INFINICORE_TRACE_OP(__OP_NAME__, INFINICORE_DETAIL_FIRST_ARG(__VA_ARGS__)); - -#define INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(__OP_NAME__, __PLAN_F__, __RUN_F__, __CLEANUP_F__) \ - static bool registered = []() { \ - __OP_NAME__::plan_dispatcher().registerAll(__PLAN_F__, false); \ - __OP_NAME__::run_dispatcher().registerAll(__RUN_F__, false); \ - __OP_NAME__::cleanup_dispatcher().registerAll(__CLEANUP_F__, false); \ - return true; \ - }(); diff --git a/include/infinicore/io.hpp b/include/infinicore/io.hpp deleted file mode 100644 index 8d11a88b2..000000000 --- a/include/infinicore/io.hpp +++ /dev/null @@ -1,86 +0,0 @@ -#pragma once - -#include "tensor.hpp" -#include - -namespace infinicore::print_options { - -/** - * @brief Sets the line width. After \a line_width chars, a new line is added. - * @param line_width The line width - */ -void set_line_width(int line_width); - -/** - * @brief Sets the threshold after which summarization is triggered (default: 1000). - * @param threshold The number of elements in the tensor that triggers summarization in the output - */ -void set_threshold(int threshold); - -/** - * @brief Sets the number of edge items. - * If the summarization is triggered, this value defines how many items of each dimension are printed. - * @param edge_items The number of edge items - */ -void set_edge_items(int edge_items); - -/** - * @brief Sets the precision for printing floating point values. - * @param precision The number of digits for floating point output - */ - -void set_precision(int precision); - -/** - * @brief Sets the sci mode of the floating point values when printing an Tensor. - * @param sci_mode The sci mode: -1 for auto decision, 0 to disable, 1 to enable - */ - -void set_sci_mode(int sci_mode); // -1: auto, 0: disable, 1: enable - -#define DEFINE_LOCAL_PRINT_OPTION(NAME) \ - class NAME { \ - public: \ - NAME(int value) : m_value(value) { id(); } \ - static int id() { \ - static int id = std::ios_base::xalloc(); \ - return id; \ - } \ - int value() const { return m_value; } \ - \ - private: \ - int m_value; \ - }; \ - \ - inline std::ostream &operator<<(std::ostream &out, const NAME &n) { \ - out.iword(NAME::id()) = n.value(); \ - return out; \ - } - -/** - * @class line_width - * io manipulator used to set the width of the lines when printing an Tensor. - * - * @code{.cpp} - * using po = infinicore::print_options; - * std::cout << po::line_width(100) << tensor << std::endl; - * @endcode - */ -DEFINE_LOCAL_PRINT_OPTION(line_width) - -/** - * io manipulator used to set the threshold after which summarization is triggered. - */ -DEFINE_LOCAL_PRINT_OPTION(threshold) - -/** - * io manipulator used to set the number of egde items if the summarization is triggered. - */ -DEFINE_LOCAL_PRINT_OPTION(edge_items) - -/** - * io manipulator used to set the precision of the floating point values when printing an Tensor. - */ -DEFINE_LOCAL_PRINT_OPTION(precision) - -} // namespace infinicore::print_options diff --git a/include/infinicore/memory.hpp b/include/infinicore/memory.hpp deleted file mode 100644 index 65d812d67..000000000 --- a/include/infinicore/memory.hpp +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once - -#include "device.hpp" - -#include -#include - -namespace infinicore { - -class Memory { -public: - using Deleter = std::function; - - Memory(std::byte *data, size_t size, Device device, Deleter deleter, bool pin_memory = false); - ~Memory(); - - std::byte *data(); - Device device() const; - size_t size() const; - bool is_pinned() const; - -private: - std::byte *data_; - size_t size_; - Device device_; - Deleter deleter_; - bool is_pinned_; -}; - -} // namespace infinicore diff --git a/include/infinicore/nn.hpp b/include/infinicore/nn.hpp deleted file mode 100644 index 16f463a72..000000000 --- a/include/infinicore/nn.hpp +++ /dev/null @@ -1,6 +0,0 @@ -#pragma once - -#include "nn/embedding.hpp" -#include "nn/linear.hpp" -#include "nn/rmsnorm.hpp" -#include "nn/rope.hpp" diff --git a/include/infinicore/nn/embedding.hpp b/include/infinicore/nn/embedding.hpp deleted file mode 100644 index 50a387325..000000000 --- a/include/infinicore/nn/embedding.hpp +++ /dev/null @@ -1,87 +0,0 @@ -#pragma once - -#include "../ops.hpp" -#include "module.hpp" -#include - -namespace infinicore::nn { - -/** - * @brief Embedding layer that maps indices to dense vectors - * - * A simple lookup table that stores embeddings of a fixed dictionary and size. - * This module is often used to store word embeddings and retrieve them using indices. - * The input to the module is a tensor of indices, and the output is the corresponding - * embedding vectors. - * - * Similar to PyTorch's nn.Embedding: - * https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html - * - * Example: - * @code - * // Create embedding: 10000 words, 300-dimensional embeddings - * auto embedding = Embedding(10000, 300); - * - * // Input: tensor of indices [batch_size, seq_len] - * auto indices = Tensor::from_data({2, 5}, {3, 5, 12, 8, 99, 0, 1, 45, 67, 23}); - * - * // Output: [batch_size, seq_len, embedding_dim] = [2, 5, 300] - * auto embeddings = embedding.forward(indices); - * @endcode - */ -class Embedding : public Module { -public: - /** - * @brief Construct an Embedding layer - * - * @param num_embeddings Size of the dictionary of embeddings (vocabulary size) - * @param embedding_dim The size of each embedding vector - * @param padding_idx If specified, the entries at padding_idx do not contribute to gradient - * and the embedding vector at padding_idx is not updated during training - * @param dtype Data type for the embedding weights (default: DataType::F32) - * @param device Device to create the embedding weight on - */ - Embedding(size_t num_embeddings, - size_t embedding_dim, - std::optional padding_idx = std::nullopt, - const DataType &dtype = DataType::F32, - const Device &device = Device()); - - /** - * @brief Forward pass: lookup embeddings for given indices - * - * @param indices Tensor containing indices into the embedding matrix. - * Can be any shape (*), typically [batch_size] or [batch_size, seq_len] - * @return Tensor containing the embedding vectors. - * Shape: (*, embedding_dim) where * matches the input shape - * - * Example: - * Input shape: [2, 3] -> Output shape: [2, 3, embedding_dim] - * Input shape: [10] -> Output shape: [10, embedding_dim] - */ - Tensor forward(const Tensor &indices) const; - - // Module information - size_t num_embeddings() const { return num_embeddings_; } - size_t embedding_dim() const { return embedding_dim_; } - std::optional padding_idx() const { return padding_idx_; } - DataType dtype() const { return dtype_; } - - // String representation - std::string extra_repr() const; - - // Accessors for parameters - Tensor weight() const { return weight_; } - -protected: - // Parameters - INFINICORE_NN_PARAMETER(weight); - -private: - size_t num_embeddings_; // Vocabulary size - size_t embedding_dim_; // Embedding dimension - std::optional padding_idx_; // Optional padding index - DataType dtype_; // Data type for embedding weights -}; - -} // namespace infinicore::nn diff --git a/include/infinicore/nn/layer_norm.hpp b/include/infinicore/nn/layer_norm.hpp deleted file mode 100644 index d69d0594f..000000000 --- a/include/infinicore/nn/layer_norm.hpp +++ /dev/null @@ -1,59 +0,0 @@ -#pragma once - -#include "../ops.hpp" -#include "module.hpp" - -namespace infinicore::nn { - -class LayerNorm : public Module { -public: - /** - * @brief Construct a LayerNorm layer - * - * @param normalized_shape Size of the feature dimension to normalize (typically hidden_size) - * @param eps Small constant for numerical stability (default: 1e-6) - * @param dtype Data type for the weight (default: DataType::F32) - * @param device Device to create the weight on - */ - LayerNorm(size_t normalized_shape, - double eps = 1e-6, - const DataType &dtype = DataType::F32, - const Device &device = Device()); - - /** - * @brief Forward pass: apply LayerNorm - * - * @param x Input tensor of shape (*, normalized_shape) where * is any number of dimensions - * @return Normalized tensor with same shape as input - * - * The normalization is applied over the last dimension. - * For example: - * Input: [batch, seq_len, hidden_size] -> normalize over hidden_size - * Input: [batch, hidden_size] -> normalize over hidden_size - */ - Tensor forward(const Tensor &x) const; - - // Module information - size_t normalized_shape() const { return normalized_shape_; } - double eps() const { return eps_; } - DataType dtype() const { return dtype_; } - - // String representation - std::string extra_repr() const; - - // Accessors for parameters - Tensor weight() const { return weight_; } - Tensor bias() const { return bias_; } - -protected: - // Parameters - INFINICORE_NN_PARAMETER(weight); - INFINICORE_NN_PARAMETER(bias); - -private: - size_t normalized_shape_; // Size of the feature dimension - double eps_; // Epsilon for numerical stability - DataType dtype_; // Data type for weight -}; - -} // namespace infinicore::nn diff --git a/include/infinicore/nn/linear.hpp b/include/infinicore/nn/linear.hpp deleted file mode 100644 index 00ea8a1a5..000000000 --- a/include/infinicore/nn/linear.hpp +++ /dev/null @@ -1,128 +0,0 @@ -#pragma once - -#include "../ops.hpp" -#include "../quantization.hpp" -#include "module.hpp" -#include -#include - -namespace infinicore::nn { - -class BaseLinear : public Module { -public: - BaseLinear(size_t in_features, size_t out_features, bool bias = true, - const DataType &dtype = DataType::F32, const Device &device = Device()); - - BaseLinear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias = true, - const DataType &dtype = DataType::F32, const Device &device = Device()); - - // Forward pass: output = input @ weight.T + bias - Tensor forward(Tensor &input) const; - - // Forward pass with residual connection (InfiniLM-style) - // output = input @ weight.T + bias + residual - Tensor forward(Tensor &input, Tensor &residual) const; - - // Module information - size_t in_features() const { return in_features_; } - size_t out_features() const { return out_features_; } - bool has_bias() const { return has_bias_; } - DataType dtype() const { return dtype_; } - float alpha() const { return alpha_; } - void set_alpha(float alpha) { alpha_ = alpha; } - - // Accessors for parameters - Tensor weight() const { return weight_; } - Tensor bias() const { return bias_; } - Tensor weight_scale() const { return weight_scale_; } - Tensor weight_zeros() const { return weight_zeros_; } - Tensor gidx() const { return gidx_; } - - std::shared_ptr get_quantization() const { return quantization_; } - void process_weights_after_loading() override; - -protected: - // Parameters - INFINICORE_NN_PARAMETER(weight); - INFINICORE_NN_PARAMETER(bias); - - INFINICORE_NN_PARAMETER(weight_scale); - INFINICORE_NN_PARAMETER(weight_zeros); - - INFINICORE_NN_PARAMETER(gidx); - -protected: - // Helper method for common forward computation - Tensor compute_linear(Tensor &input) const; - - size_t in_features_; - size_t out_features_; - bool has_bias_; - DataType dtype_; - float alpha_ = 1.0f; - std::shared_ptr quantization_ = std::make_shared(nullptr); -}; - -} // namespace infinicore::nn - -namespace infinicore::nn { - -class Linear : public BaseLinear { -public: - Linear(size_t in_features, size_t out_features, bool bias = true, - const DataType &dtype = DataType::F32, const Device &device = Device()); - - Linear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias = true, - const DataType &dtype = DataType::F32, const Device &device = Device()); - - // Forward pass: output = input @ weight.T + bias - Tensor forward(Tensor &input) const; - - // String representation - std::string extra_repr() const; -}; - -class ColumnParallelLinear : public BaseLinear { -public: - ColumnParallelLinear(size_t in_features, size_t out_features, bool bias = true, - const DataType &dtype = DataType::F32, const Device &device = Device(), - Size tp_rank = 0, Size tp_size = 1); - - ColumnParallelLinear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias = true, - const DataType &dtype = DataType::F32, const Device &device = Device(), - Size tp_rank = 0, Size tp_size = 1); - - // Forward pass: output = input @ weight.T + bias - Tensor forward(Tensor &input) const; - - // String representation - std::string extra_repr() const; - -protected: - Size tp_rank_ = 0; - Size tp_size_ = 1; -}; - -class RowParallelLinear : public BaseLinear { -public: - RowParallelLinear(size_t in_features, size_t out_features, bool bias = true, - const DataType &dtype = DataType::F32, const Device &device = Device(), - Size tp_rank = 0, Size tp_size = 1, infinicclComm_t communicator = nullptr); - - RowParallelLinear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias = true, - const DataType &dtype = DataType::F32, const Device &device = Device(), - Size tp_rank = 0, Size tp_size = 1, infinicclComm_t communicator = nullptr); - - // Forward pass: output = input @ weight.T + bias - Tensor forward(Tensor &input) const; - - // String representation - std::string extra_repr() const; - -protected: - Size tp_rank_ = 0; - Size tp_size_ = 1; - infinicclComm_t communicator_; -}; - -} // namespace infinicore::nn diff --git a/include/infinicore/nn/module.hpp b/include/infinicore/nn/module.hpp deleted file mode 100644 index 2b5747a16..000000000 --- a/include/infinicore/nn/module.hpp +++ /dev/null @@ -1,174 +0,0 @@ -#pragma once - -#include "../tensor.hpp" -#include "parameter.hpp" - -#include -#include -#include -#include -#include - -namespace infinicore::nn { -class Module { -public: - Module() = default; - - virtual ~Module() = default; - - std::unordered_map state_dict() const; - - std::vector state_dict_keys() const; - - void load_state_dict(const std::unordered_map &_state_dict); - - void load_parameter(const std::string &name, const Tensor ¶m); - - void load_parameters_no_sync(const std::unordered_map ¶ms, bool strict = true); - - void load_parameter_(const std::string &name, const Tensor ¶m); - - void load_parameter_from_blob(const std::string &name, const void *data); - - std::unordered_map modules_dict() const; - - virtual void process_weights_after_loading() {} - - virtual void reset_runtime_state() const {} - - const std::unordered_map> &children() const { - return submodules_; - } - -protected: - Tensor register_parameter(const std::string &name, Parameter param); - - Tensor register_buffer(const std::string &name, Parameter buffer); - - // Add an existing submodule to this module's hierarchy - // Template parameter M must be a type derived from Module - // Returns the submodule for convenience (allows method chaining) - template - std::shared_ptr add_module(const std::string &name, std::shared_ptr submodule) { - // Ensure M is derived from Module (compile-time check) - static_assert(std::is_base_of::value, - "Template parameter M must be derived from infinicore::nn::Module"); - - // Store in the submodules map (std::shared_ptr automatically converts to std::shared_ptr) - submodules_[name] = submodule; - - return submodule; - } - - // Create and register a new submodule by constructing it with the given arguments - // Template parameter M must be a type derived from Module - // Args are forwarded to M's constructor - template - std::shared_ptr register_module(const std::string &name, Args &&...args) { - // Ensure M is derived from Module (compile-time check) - static_assert(std::is_base_of::value, - "Template parameter M must be derived from infinicore::nn::Module"); - - // Construct the submodule - auto submodule = std::make_shared(std::forward(args)...); - - return add_module(name, submodule); - } - - // Create and register multiple submodules of the same type - // Each submodule is named as "name.0", "name.1", etc. - // Template parameter M must be a type derived from Module - template - std::vector> register_modules(size_t count, const std::string &name, Args &&...args) { - static_assert(std::is_base_of::value, - "Template parameter M must be derived from infinicore::nn::Module"); - - std::vector> modules; - modules.reserve(count); - for (size_t i = 0; i < count; i++) { - modules.push_back(register_module(name + "." + std::to_string(i), std::forward(args)...)); - } - return modules; - } - -protected: - Device device_; - std::unordered_map> submodules_; - std::unordered_map buffers_; - std::unordered_map parameters_; - -private: - void load_state_dict_recursively(const std::unordered_map &_state_dict, const std::string &prefix = ""); - void collect_all_parameters(std::unordered_map &all_params, const std::string &prefix = "") const; - void collect_all_parameter_names(std::vector &all_names, const std::string &prefix = "") const; - void collect_all_modules(std::unordered_map &out, const std::string &prefix) const; -}; - -// ============================================================================ -// PyTorch-like Macros for Convenient Module Registration -// ============================================================================ - -/** - * @brief Register submodules with automatic name inference from variable name - * - * Usage: - * @code - * class MyModel : public Module { - * protected: - * INFINICORE_NN_MODULE(Linear, layer1); - * INFINICORE_NN_MODULE(Linear, layer2); - * INFINICORE_NN_MODULE_VEC(Linear, layers); - * INFINICORE_NN_PARAMETER(scaling_factor); - * - * public: - * MyModel() { - * INFINICORE_NN_MODULE_INIT(layer1, 128, 64); - * INFINICORE_NN_MODULE_INIT(layer2, 64, 32); - * INFINICORE_NN_MODULE_VEC_INIT(layers, 3, Linear, 32, 16); - * INFINICORE_NN_PARAMETER_INIT(scaling_factor, ({1}, DataType::F32, Device())); - * } - * }; - * @endcode - */ - -// Declare a single module member variable -#define INFINICORE_NN_MODULE(ModuleType, name) \ - std::shared_ptr name##_ - -// Declare a vector of modules member variable -#define INFINICORE_NN_MODULE_VEC(ModuleType, name) \ - std::vector> name##_ - -// Initialize a module in constructor -#define INFINICORE_NN_MODULE_INIT(name, ...) \ - name##_ = this->register_module::type>(#name, ##__VA_ARGS__) - -// Initialize a vector of modules in constructor -// Usage: INFINICORE_NN_MODULE_VEC_INIT(layers, count, ModuleType, ctor_args...) -// Example: INFINICORE_NN_MODULE_VEC_INIT(layers, 3, Linear, 128, 64) -#define INFINICORE_NN_MODULE_VEC_INIT(name, count, ModuleType, ...) \ - name##_ = this->register_modules(count, #name, ##__VA_ARGS__) - -// Declare a parameter member variable -#define INFINICORE_NN_PARAMETER(name) \ - infinicore::nn::Parameter name##_ - -// Initialize a parameter in constructor -// Usage: INFINICORE_NN_PARAMETER_INIT(name, (shape, dtype, device)) -// Example: INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, DataType::F32, device)) -#define INFINICORE_NN_PARAMETER_INIT(name, args) \ - name##_ = infinicore::nn::Parameter args; \ - this->register_parameter(#name, name##_) - -// Declare a buffer member variable -#define INFINICORE_NN_BUFFER(name) \ - infinicore::nn::Parameter name##_ - -// Initialize a buffer in constructor -// Usage: INFINICORE_NN_BUFFER_INIT(name, (shape, dtype, device)) -// Example: INFINICORE_NN_BUFFER_INIT(cache, ({max_seq_len, head_dim}, DataType::F32, device)) -#define INFINICORE_NN_BUFFER_INIT(name, args) \ - name##_ = infinicore::nn::Parameter args; \ - this->register_buffer(#name, name##_) - -} // namespace infinicore::nn diff --git a/include/infinicore/nn/parameter.hpp b/include/infinicore/nn/parameter.hpp deleted file mode 100644 index 1602f58d6..000000000 --- a/include/infinicore/nn/parameter.hpp +++ /dev/null @@ -1,39 +0,0 @@ -#pragma once - -#include "../tensor.hpp" - -namespace infinicore::nn { -class Parameter : public Tensor { -public: - Parameter(); - - Parameter(const Tensor &tensor, - Size tp_dim = 0, - Size tp_rank = 0, - Size tp_size = 1, - Size num_shards = 0); - - Parameter(const Shape &shape, - const DataType &dtype, - const Device &device, - Size tp_dim = 0, - Size tp_rank = 0, - Size tp_size = 1, - Size num_shards = 0); - - Parameter(const Parameter &other); - - void load_blob(const void *data); - - void load(const Tensor &tensor); - - void load_no_sync(const Tensor &tensor); - -protected: - // Tensor parallel configs - Size tp_dim_; // dimension partitioned - Size tp_rank_; // rank of this partition among tp group - Size tp_size_; // total number of partitions - Size num_shards_ = 0; // number of logical shards, used when tp_size > num_kv_head -}; -} // namespace infinicore::nn diff --git a/include/infinicore/nn/rmsnorm.hpp b/include/infinicore/nn/rmsnorm.hpp deleted file mode 100644 index 5891819eb..000000000 --- a/include/infinicore/nn/rmsnorm.hpp +++ /dev/null @@ -1,96 +0,0 @@ -#pragma once - -#include "../ops.hpp" -#include "module.hpp" - -namespace infinicore::nn { - -/** - * @brief Root Mean Square Layer Normalization (RMSNorm) - * - * Applies Root Mean Square Layer Normalization over the last dimension. - * Unlike LayerNorm, RMSNorm doesn't subtract mean and doesn't use bias. - * - * Formula: y = (x / RMS(x)) * weight - * where RMS(x) = sqrt(mean(x^2) + eps) - * - * Used in LLaMA, Galactica, and other modern language models as a - * simpler and faster alternative to LayerNorm. - * - * Example: - * @code - * // Create RMSNorm for hidden size 4096 - * auto norm = RMSNorm(4096); - * - * // Input: [batch, seq_len, hidden_size] - * auto input = Tensor::randn({2, 10, 4096}); - * - * // Output: [batch, seq_len, hidden_size] - * auto output = norm.forward(input); - * @endcode - */ -class RMSNorm : public Module { -public: - /** - * @brief Construct a RMSNorm layer - * - * @param normalized_shape Size of the feature dimension to normalize (typically hidden_size) - * @param eps Small constant for numerical stability (default: 1e-6) - * @param dtype Data type for the weight (default: DataType::F32) - * @param device Device to create the weight on - */ - RMSNorm(size_t normalized_shape, - double eps = 1e-6, - const DataType &dtype = DataType::F32, - const Device &device = Device()); - - /** - * @brief Forward pass: apply RMSNorm - * - * @param x Input tensor of shape (*, normalized_shape) where * is any number of dimensions - * @return Normalized tensor with same shape as input - * - * The normalization is applied over the last dimension. - * For example: - * Input: [batch, seq_len, hidden_size] -> normalize over hidden_size - * Input: [batch, hidden_size] -> normalize over hidden_size - */ - Tensor forward(const Tensor &x) const; - - /** - * @brief Forward pass: apply RMSNorm in-place with residual - * - * @param x Input tensor of shape (*, normalized_shape) where * is any number of dimensions. - * Will be modified in-place to the normalized output. - * @param residual Residual tensor to add to input before normalization. - * Will be modified in-place to the sum of input and residual. - * - * The normalization is applied over the last dimension. - * For example: - * Input: [batch, seq_len, hidden_size] -> normalize over hidden_size - * Input: [batch, hidden_size] -> normalize over hidden_size - */ - void forward_inplace(Tensor &x, Tensor &residual) const; - - // Module information - size_t normalized_shape() const { return normalized_shape_; } - double eps() const { return eps_; } - DataType dtype() const { return dtype_; } - - // String representation - std::string extra_repr() const; - - // Accessors for parameters - Tensor weight() const { return weight_; } - -protected: - // Parameters - INFINICORE_NN_PARAMETER(weight); - -private: - size_t normalized_shape_; // Size of the feature dimension - double eps_; // Epsilon for numerical stability - DataType dtype_; // Data type for weight -}; - -} // namespace infinicore::nn diff --git a/include/infinicore/nn/rope.hpp b/include/infinicore/nn/rope.hpp deleted file mode 100644 index cbffb81c8..000000000 --- a/include/infinicore/nn/rope.hpp +++ /dev/null @@ -1,112 +0,0 @@ -#pragma once - -#include "../context/context.hpp" -#include "../tensor.hpp" -#include "module.hpp" -#include "rope_scaling_configs.hpp" -#include -#include -#include -#include -#include - -namespace infinicore::nn { - -class RoPE : public Module { -public: - /** - * @brief RoPE algorithm type - */ - enum class Algo { - GPT_J = 0, // GPT-J style RoPE algorithm (Interleave even and odd dimensions) - GPT_NEOX = 1, // GPT-NeoX style RoPE algorithm (First half dimensions for sin, second half for cos) - }; - - /** - * @brief Construct a RoPE layer - * - * @param head_dim Dimension of each attention head (must be even) - * @param rotary_dim Number of dimensions to apply rotation to (must be even). - * For full rotation models, this equals head_dim; - * for partial rotation models, this equals head_dim * partial_rotary_factor. - * @param max_seq_len Maximum sequence length for pre-computed cache - * @param theta Base frequency for rotary embeddings (default: 10000.0) - * @param algo RoPE algorithm type (default: Algo::GPT_J) - * @param dtype Data type for sin/cos cache (default: DataType::F32) - * @param device Device to create the cache on - * @param scaling RoPE scaling configuration (default: nullptr) - * @param mrope_section Optional MRoPE section sizes [t, h, w], whose sum must equal rotary_dim / 2. - * When set, pair forward overloads apply MRoPE to q/k using positions [3, num_tokens]. - * @param mrope_interleaved Whether to interleave MRoPE axes/frequency sections. - */ - RoPE(size_t head_dim, - size_t rotary_dim, - size_t max_seq_len, - double theta = 10000.0, - Algo algo = Algo::GPT_J, - const DataType &dtype = DataType::F32, - const Device &device = Device(), - std::shared_ptr scaling = nullptr, - std::optional> mrope_section = std::nullopt, - bool mrope_interleaved = false); - - /** - * @brief Forward pass: apply standard RoPE to a tensor - * - * @param x Input tensor of shape (..., rotary_dim) where ... is any number of dimensions - * @param pos Position IDs tensor of shape (*,) typically [seq_len] or [batch, seq_len] - * @param in_place If true, modify input tensor in place (default: false) - * @return Rotated tensor with same shape as input - */ - Tensor forward(const Tensor &x, const Tensor &pos, bool in_place = false) const; - - /** - * @brief Apply MRoPE to q and k. - * - * Requires construction with mrope_section. q/k may be either - * [num_tokens, num_heads * head_dim] or [num_tokens, num_heads, head_dim]. - * positions is [3, num_tokens] with axes ordered as t, h, w. - */ - std::pair forward(const Tensor &q, const Tensor &k, const Tensor &positions) const; - - /** - * @brief Apply MRoPE to q and k into caller-provided outputs. - */ - std::pair forward(const Tensor &q_out, - const Tensor &k_out, - const Tensor &q, - const Tensor &k, - const Tensor &positions) const; - - // Module information - size_t rotary_dim() const { return rotary_dim_; } - size_t head_dim() const { return head_dim_; } - size_t max_seq_len() const { return max_seq_len_; } - double theta() const { return theta_; } - Algo algo() const { return algo_; } - DataType dtype() const { return dtype_; } - const std::optional> &mrope_section() const { return mrope_section_; } - bool mrope_interleaved() const { return mrope_interleaved_; } - - // String representation - std::string extra_repr() const; - -protected: - // Buffers (sin and cos cache tables) - not exposed in state_dict - INFINICORE_NN_BUFFER(sin_cache); - INFINICORE_NN_BUFFER(cos_cache); - -private: - void initialize_cache(); - size_t rotary_dim_; // Number of dimensions to apply rotation to (must be even). - size_t head_dim_; // Dimension of each attention head - size_t max_seq_len_; // Maximum sequence length - double theta_; // Base frequency for rotary embeddings - Algo algo_; // RoPE algorithm type - DataType dtype_; // Data type for cache tables - std::shared_ptr scaling_; // RoPE scaling configuration - std::optional> mrope_section_; - bool mrope_interleaved_; -}; - -} // namespace infinicore::nn diff --git a/include/infinicore/nn/rope_scaling_configs.hpp b/include/infinicore/nn/rope_scaling_configs.hpp deleted file mode 100644 index 4129e6053..000000000 --- a/include/infinicore/nn/rope_scaling_configs.hpp +++ /dev/null @@ -1,131 +0,0 @@ -#pragma once -#include -#include - -namespace infinicore::nn { - -/** - * @brief Abstract base class for RoPE scaling strategies. - * Uses polymorphism to eliminate type checking (if-else) in the core RoPE loop. - */ -class RopeScalingConfig { -public: - virtual ~RopeScalingConfig() = default; - - /** - * @brief Calculate the frequency scaling factor for a specific position and dimension. - * - * @param pos Current sequence position - * @param dim_idx Current dimension index (0 to head_dim/2 - 1) - * @param base_inv_freq Pre-computed base inverse frequency for this dimension (1.0 / theta^(2j/head_dim)) - * @return Frequency scaling factor (default 1.0) - */ - virtual float get_freq_scale(size_t pos, size_t dim_idx, float base_inv_freq) const { - return 1.0f; - } - - /** - * @brief Calculate the magnitude scaling factor for a specific position and dimension. - * - * @param pos Current sequence position - * @param dim_idx Current dimension index (0 to head_dim/2 - 1) - * @param base_inv_freq Pre-computed base inverse frequency for this dimension - * @return Magnitude scaling factor (default 1.0) - */ - virtual float get_magnitude_scale(size_t pos, size_t dim_idx, float base_inv_freq) const { - return 1.0f; - } -}; - -/** - * @brief LongRoPE scaling configuration. - */ -class LongRopeScalingConfig : public RopeScalingConfig { -public: - LongRopeScalingConfig( - std::vector short_factor, - std::vector long_factor, - size_t original_max_position_embeddings, - float factor = 1.0f); - - float get_freq_scale(size_t pos, size_t dim_idx, float base_inv_freq) const override; - float get_magnitude_scale(size_t pos, size_t dim_idx, float base_inv_freq) const override; - - size_t original_max_position_embeddings() const { return original_max_position_embeddings_; } - const std::vector &short_factor() const { return short_factor_; } - const std::vector &long_factor() const { return long_factor_; } - float factor() const { return factor_; } - -private: - std::vector short_factor_; - std::vector long_factor_; - size_t original_max_position_embeddings_; - float factor_; -}; - -// TODO(rubik) implement in cpp -/** - * @brief Llama3 frequency-aware RoPE scaling configuration. - * Native support for Llama 3.1 RoPE scaling (smooth interpolation based on wavelength). - */ -class Llama3RopeScalingConfig : public RopeScalingConfig { -public: - Llama3RopeScalingConfig( - float factor, - float low_freq_factor, - float high_freq_factor, - size_t original_max_position_embeddings); - - float get_freq_scale(size_t pos, size_t dim_idx, float base_inv_freq) const override; - - // Llama3 does not use magnitude scaling, so it inherits the default get_magnitude_scale() returning 1.0f - -private: - float factor_; - float low_freq_factor_; - float high_freq_factor_; - size_t original_max_position_embeddings_; -}; - -/** - * @brief YaRN (Yet another RoPE extensioN) scaling configuration. - * - * rope_scaling fields: factor, original_max_position_embeddings, beta_fast, beta_slow, - * mscale, mscale_all_dim - * Model fields (must match RoPE): rotary_dim (e.g. qk_rope_head_dim), rope_theta - */ -class YarnRopeScalingConfig : public RopeScalingConfig { -public: - YarnRopeScalingConfig( - float factor, - size_t original_max_position_embeddings, - size_t rotary_dim, - float rope_theta, - int beta_fast = 32, - int beta_slow = 1, - float mscale = 1.0f, - float mscale_all_dim = 0.0f); - - float get_freq_scale(size_t pos, size_t dim_idx, float base_inv_freq) const override; - float get_magnitude_scale(size_t pos, size_t dim_idx, float base_inv_freq) const override; - - /** Recommended RoPE cache length: original_max_position_embeddings * factor. */ - static size_t max_seq_len(float factor, size_t original_max_position_embeddings) { - return static_cast( - static_cast(original_max_position_embeddings) * factor); - } - - float factor() const { return factor_; } - size_t original_max_position_embeddings() const { return original_max_position_embeddings_; } - -private: - float yarn_linear_ramp(size_t dim_idx) const; - - float factor_; - size_t original_max_position_embeddings_; - float magnitude_scale_; - float correction_low_; - float correction_high_; -}; - -} // namespace infinicore::nn diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp deleted file mode 100644 index b5c4ff18f..000000000 --- a/include/infinicore/ops.hpp +++ /dev/null @@ -1,76 +0,0 @@ -#pragma once - -#include "ops/adaptive_max_pool1d.hpp" -#include "ops/add.hpp" -#include "ops/add_rms_norm.hpp" -#include "ops/addcmul.hpp" -#include "ops/asin.hpp" -#include "ops/asinh.hpp" -#include "ops/asum.hpp" -#include "ops/atanh.hpp" -#include "ops/attention.hpp" -#include "ops/avg_pool1d.hpp" -#include "ops/axpy.hpp" -#include "ops/baddbmm.hpp" -#include "ops/bilinear.hpp" -#include "ops/binary_cross_entropy_with_logits.hpp" -#include "ops/blas_amax.hpp" -#include "ops/blas_amin.hpp" -#include "ops/blas_copy.hpp" -#include "ops/blas_dot.hpp" -#include "ops/causal_conv1d.hpp" -#include "ops/causal_softmax.hpp" -#include "ops/cdist.hpp" -#include "ops/chunk_gated_delta_rule.hpp" -#include "ops/conv2d.hpp" -#include "ops/cross_entropy.hpp" -#include "ops/deepseek_moe.hpp" -#include "ops/embedding.hpp" -#include "ops/flash_attention.hpp" -#include "ops/fmin.hpp" -#include "ops/fmod.hpp" -#include "ops/fused_gated_delta_net_gating.hpp" -#include "ops/gelu.hpp" -#include "ops/gelutanh.hpp" -#include "ops/hardswish.hpp" -#include "ops/hardtanh.hpp" -#include "ops/kv_caching.hpp" -#include "ops/layer_norm.hpp" -#include "ops/linear.hpp" -#include "ops/mamba_selective_scan.hpp" -#include "ops/matmul.hpp" -#include "ops/moe_align.hpp" -#include "ops/moe_fused_dense.hpp" -#include "ops/moe_fused_gate.hpp" -#include "ops/moe_sum.hpp" -#include "ops/moe_topk_sigmoid.hpp" -#include "ops/moe_topk_softmax.hpp" -#include "ops/nrm2.hpp" -#include "ops/ones.hpp" -#include "ops/paged_attention.hpp" -#include "ops/paged_attention_prefill.hpp" -#include "ops/paged_caching.hpp" -#include "ops/per_tensor_dequant_i8.hpp" -#include "ops/per_tensor_quant_i8.hpp" -#include "ops/prepare_moe_input.hpp" -#include "ops/quickgelu.hpp" -#include "ops/random_sample.hpp" -#include "ops/rearrange.hpp" -#include "ops/reciprocal.hpp" -#include "ops/recurrent_gated_delta_rule.hpp" -#include "ops/relu.hpp" -#include "ops/rms_norm.hpp" -#include "ops/rope.hpp" -#include "ops/rot.hpp" -#include "ops/rotg.hpp" -#include "ops/rotm.hpp" -#include "ops/rotmg.hpp" -#include "ops/rwkv5_wkv.hpp" -#include "ops/scal.hpp" -#include "ops/sigmoid.hpp" -#include "ops/silu.hpp" -#include "ops/silu_and_mul.hpp" -#include "ops/softmax.hpp" -#include "ops/swap.hpp" -#include "ops/swiglu.hpp" -#include "ops/topksoftmax.hpp" diff --git a/include/infinicore/ops/acos.hpp b/include/infinicore/ops/acos.hpp deleted file mode 100644 index 91aaee020..000000000 --- a/include/infinicore/ops/acos.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Acos { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor acos(Tensor input); -void acos_(Tensor output, Tensor input); -} // namespace infinicore::op diff --git a/include/infinicore/ops/adaptive_avg_pool1d.hpp b/include/infinicore/ops/adaptive_avg_pool1d.hpp deleted file mode 100644 index 3e8a95b8d..000000000 --- a/include/infinicore/ops/adaptive_avg_pool1d.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class AdaptiveAvgPool1d { -public: - // Schema: execute(Output, Input) - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor adaptive_avg_pool1d(Tensor input, int64_t output_size); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/adaptive_avg_pool3d.hpp b/include/infinicore/ops/adaptive_avg_pool3d.hpp deleted file mode 100644 index df6a4d2c3..000000000 --- a/include/infinicore/ops/adaptive_avg_pool3d.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class AdaptiveAvgPool3D { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor y, Tensor x); - static common::OpDispatcher &dispatcher(); -}; - -Tensor adaptive_avg_pool3d(Tensor x, std::vector output_size); -void adaptive_avg_pool3d_(Tensor y, Tensor x); -} // namespace infinicore::op diff --git a/include/infinicore/ops/adaptive_max_pool1d.hpp b/include/infinicore/ops/adaptive_max_pool1d.hpp deleted file mode 100644 index 51832cae0..000000000 --- a/include/infinicore/ops/adaptive_max_pool1d.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class AdaptiveMaxPool1d { -public: - using schema = void (*)(Tensor, Tensor, size_t); - static void execute(Tensor y, Tensor x, size_t output_size); - static common::OpDispatcher &dispatcher(); -}; - -Tensor adaptive_max_pool1d(Tensor x, size_t output_size); -void adaptive_max_pool1d_(Tensor y, Tensor x, size_t output_size); -} // namespace infinicore::op diff --git a/include/infinicore/ops/add.hpp b/include/infinicore/ops/add.hpp deleted file mode 100644 index 528cca18a..000000000 --- a/include/infinicore/ops/add.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Add, Tensor, const Tensor &, const Tensor &); - -Tensor add(const Tensor &a, const Tensor &b); -void add_(Tensor c, const Tensor &a, const Tensor &b); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/add_rms_norm.hpp b/include/infinicore/ops/add_rms_norm.hpp deleted file mode 100644 index 50064e0a4..000000000 --- a/include/infinicore/ops/add_rms_norm.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { -INFINICORE_GRAPH_OP_CLASS(AddRMSNorm, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &, float); - -// Fused Add and RMS Normalization -// Returns: (normalized_result, add_result) -// The add_result can be used as residual for subsequent layers -std::pair add_rms_norm(const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon = 1e-5f); -void add_rms_norm_(Tensor out, Tensor residual, const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon = 1e-5f); -// Fused Add and RMS Normalization (inplace) -// normalized_result wil be stored in input, add_result will be stored in residual -void add_rms_norm_inplace(Tensor input, Tensor residual, const Tensor &weight, float epsilon = 1e-5f); -} // namespace infinicore::op diff --git a/include/infinicore/ops/addbmm.hpp b/include/infinicore/ops/addbmm.hpp deleted file mode 100644 index 6c17a35d5..000000000 --- a/include/infinicore/ops/addbmm.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Addbmm { -public: - using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float, float); - static void execute(Tensor output, Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha); - - static common::OpDispatcher &dispatcher(); -}; -Tensor addbmm(Tensor input, Tensor batch1, Tensor batch2, float beta = 1.0f, float alpha = 1.0f); - -void addbmm_(Tensor output, Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/addcmul.hpp b/include/infinicore/ops/addcmul.hpp deleted file mode 100644 index 2806d1218..000000000 --- a/include/infinicore/ops/addcmul.hpp +++ /dev/null @@ -1,17 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Addcmul { -public: - // schema: out, input, t1, t2, value - using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float); - static void execute(Tensor out, Tensor input, Tensor t1, Tensor t2, float value); - static common::OpDispatcher &dispatcher(); -}; -Tensor addcmul(Tensor input, Tensor t1, Tensor t2, float value); -void addcmul_(Tensor out, Tensor input, Tensor t1, Tensor t2, float value); -} // namespace infinicore::op diff --git a/include/infinicore/ops/addr.hpp b/include/infinicore/ops/addr.hpp deleted file mode 100644 index cff78eade..000000000 --- a/include/infinicore/ops/addr.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Addr { -public: - using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float, float); - static void execute(Tensor out, Tensor input, Tensor vec1, Tensor vec2, float beta = 1.0f, float alpha = 1.0f); - static common::OpDispatcher &dispatcher(); -}; - -Tensor addr(Tensor input, Tensor vec1, Tensor vec2, float beta = 1.0f, float alpha = 1.0f); -void addr_(Tensor out, Tensor input, Tensor vec1, Tensor vec2, float beta = 1.0f, float alpha = 1.0f); -} // namespace infinicore::op diff --git a/include/infinicore/ops/affine_grid.hpp b/include/infinicore/ops/affine_grid.hpp deleted file mode 100644 index ea025c24d..000000000 --- a/include/infinicore/ops/affine_grid.hpp +++ /dev/null @@ -1,17 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -class AffineGrid { -public: - using schema = void (*)(Tensor, Tensor, bool); - static void execute(Tensor output, Tensor theta, bool align_corners); - static common::OpDispatcher &dispatcher(); -}; -Tensor affine_grid(Tensor theta, const std::vector &size, bool align_corners = false); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/all.hpp b/include/infinicore/ops/all.hpp deleted file mode 100644 index 50d76f2d7..000000000 --- a/include/infinicore/ops/all.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include -#include -namespace infinicore::op { -class All { -public: - using schema = void (*)(Tensor, Tensor, std::vector, bool); - static void execute(Tensor output, Tensor input, std::vector dim, bool keepdim = false); - static common::OpDispatcher &dispatcher(); -}; - -Tensor all(Tensor input, std::vector dim, bool keepdim = false); -void all_(Tensor output, Tensor input, std::vector dim, bool keepdim = false); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/argwhere.hpp b/include/infinicore/ops/argwhere.hpp deleted file mode 100644 index 24f4fa65e..000000000 --- a/include/infinicore/ops/argwhere.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Argwhere { -public: - using schema = void (*)(void **, size_t *, Tensor); - static void execute(void **, size_t *count, Tensor x); - static common::OpDispatcher &dispatcher(); -}; -Tensor argwhere(Tensor x); -} // namespace infinicore::op diff --git a/include/infinicore/ops/asin.hpp b/include/infinicore/ops/asin.hpp deleted file mode 100644 index bc5973978..000000000 --- a/include/infinicore/ops/asin.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Asin { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor asin(Tensor input); -void asin_(Tensor output, Tensor input); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/asinh.hpp b/include/infinicore/ops/asinh.hpp deleted file mode 100644 index 505eb97d9..000000000 --- a/include/infinicore/ops/asinh.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Asinh { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor y, Tensor x); - static common::OpDispatcher &dispatcher(); -}; - -Tensor asinh(Tensor x); -void asinh_(Tensor y, Tensor x); -} // namespace infinicore::op diff --git a/include/infinicore/ops/asum.hpp b/include/infinicore/ops/asum.hpp deleted file mode 100644 index 6471e6fe4..000000000 --- a/include/infinicore/ops/asum.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Asum, const Tensor &, Tensor); - -Tensor asum(const Tensor &x); -void asum_(const Tensor &x, Tensor result); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/atanh.hpp b/include/infinicore/ops/atanh.hpp deleted file mode 100644 index 5c3894e5c..000000000 --- a/include/infinicore/ops/atanh.hpp +++ /dev/null @@ -1,34 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Atanh { -public: - // schema 定义为：void(输出 Tensor, 输入 Tensor) - using schema = void (*)(Tensor, Tensor); - - // 执行函数 - static void execute(Tensor y, Tensor a); - - // 获取算子分发器，用于多后端（CPU/CUDA 等）匹配 - static common::OpDispatcher &dispatcher(); -}; - -/** - * @brief 计算输入 Tensor 的反双曲正切值 (out-of-place) - * @param a 输入 Tensor - * @return 包含结果的新 Tensor - */ -Tensor atanh(Tensor a); - -/** - * @brief 计算输入 Tensor 的反双曲正切值 (in-place / specified output) - * @param y 输出 Tensor - * @param a 输入 Tensor - */ -void atanh_(Tensor y, Tensor a); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/attention.hpp b/include/infinicore/ops/attention.hpp deleted file mode 100644 index 1bc447c77..000000000 --- a/include/infinicore/ops/attention.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Attention { -public: - using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, size_t); - static void execute(Tensor out, Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos); - static common::OpDispatcher &dispatcher(); -}; - -Tensor attention(Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos); -void attention_(Tensor out, Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos); -} // namespace infinicore::op diff --git a/include/infinicore/ops/avg_pool1d.hpp b/include/infinicore/ops/avg_pool1d.hpp deleted file mode 100644 index 4bf69bc2a..000000000 --- a/include/infinicore/ops/avg_pool1d.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class AvgPool1d { -public: - using schema = void (*)(Tensor, Tensor, size_t, size_t, size_t); - static void execute(Tensor output, Tensor input, size_t kernel_size, size_t stride, size_t padding); - static common::OpDispatcher &dispatcher(); -}; - -Tensor avg_pool1d(Tensor input, size_t kernel_size, size_t stride = 0, size_t padding = 0); -void avg_pool1d_(Tensor output, Tensor input, size_t kernel_size, size_t stride = 0, size_t padding = 0); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/awq_marlin_gemm.hpp b/include/infinicore/ops/awq_marlin_gemm.hpp deleted file mode 100644 index 164ee016b..000000000 --- a/include/infinicore/ops/awq_marlin_gemm.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(AwqMarlinGemm, Tensor, const Tensor &, const Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, int64_t, bool, bool, bool, bool); - -void awq_marlin_gemm_(Tensor c, const Tensor &a, const Tensor &b, Tensor &b_bias, Tensor &b_scales, Tensor &a_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/axpy.hpp b/include/infinicore/ops/axpy.hpp deleted file mode 100644 index 280d5ab60..000000000 --- a/include/infinicore/ops/axpy.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Axpy, const Tensor &, const Tensor &, Tensor); - -void axpy_(const Tensor &alpha, const Tensor &x, Tensor y); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/baddbmm.hpp b/include/infinicore/ops/baddbmm.hpp deleted file mode 100644 index 9988000bd..000000000 --- a/include/infinicore/ops/baddbmm.hpp +++ /dev/null @@ -1,15 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -Tensor baddbmm(Tensor input, Tensor batch1, Tensor batch2, - float beta = 1.0f, - float alpha = 1.0f); -void baddbmm_(Tensor out, Tensor input, Tensor batch1, Tensor batch2, - float beta = 1.0f, - float alpha = 1.0f); -} // namespace infinicore::op diff --git a/include/infinicore/ops/bilinear.hpp b/include/infinicore/ops/bilinear.hpp deleted file mode 100644 index 805fa1efc..000000000 --- a/include/infinicore/ops/bilinear.hpp +++ /dev/null @@ -1,12 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -Tensor bilinear(Tensor x1, Tensor x2, Tensor weight, std::optional bias); -void bilinear_(Tensor out, Tensor x1, Tensor x2, Tensor weight, std::optional bias); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/binary_cross_entropy_with_logits.hpp b/include/infinicore/ops/binary_cross_entropy_with_logits.hpp deleted file mode 100644 index 4cdead724..000000000 --- a/include/infinicore/ops/binary_cross_entropy_with_logits.hpp +++ /dev/null @@ -1,46 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -class BinaryCrossEntropyWithLogits { -public: - /** - * @brief BCEWithLogits 算子的函数原型 - * 参数顺序: out, logits, target, weight, pos_weight, reduction - */ - using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor, std::string); - - static void execute(Tensor out, - Tensor logits, - Tensor target, - Tensor weight, - Tensor pos_weight, - std::string reduction); - - static common::OpDispatcher &dispatcher(); -}; - -/** - * @brief 非原地操作接口 (Out-of-place) - */ -Tensor binary_cross_entropy_with_logits(Tensor logits, - Tensor target, - Tensor weight = {}, - Tensor pos_weight = {}, - std::string reduction = "mean"); - -/** - * @brief 显式指定输出张量的接口 - */ -void binary_cross_entropy_with_logits_(Tensor out, - Tensor logits, - Tensor target, - Tensor weight, - Tensor pos_weight, - std::string reduction); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/bitwise_right_shift.hpp b/include/infinicore/ops/bitwise_right_shift.hpp deleted file mode 100644 index db8d6a2e1..000000000 --- a/include/infinicore/ops/bitwise_right_shift.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "infinicore.h" - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(BitwiseRightShift, Tensor, const Tensor &, const Tensor &); - -__export Tensor bitwise_right_shift(const Tensor &input, const Tensor &other); -__export void bitwise_right_shift_(Tensor out, const Tensor &input, const Tensor &other); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/blas_amax.hpp b/include/infinicore/ops/blas_amax.hpp deleted file mode 100644 index a6a571f95..000000000 --- a/include/infinicore/ops/blas_amax.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(BlasAmax, const Tensor &, Tensor); - -Tensor blas_amax(const Tensor &x); -void blas_amax_(const Tensor &x, Tensor result); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/blas_amin.hpp b/include/infinicore/ops/blas_amin.hpp deleted file mode 100644 index a2ed21c7b..000000000 --- a/include/infinicore/ops/blas_amin.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(BlasAmin, const Tensor &, Tensor); - -Tensor blas_amin(const Tensor &x); -void blas_amin_(const Tensor &x, Tensor result); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/blas_copy.hpp b/include/infinicore/ops/blas_copy.hpp deleted file mode 100644 index dd32646a0..000000000 --- a/include/infinicore/ops/blas_copy.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(BlasCopy, const Tensor &, Tensor); - -void blas_copy_(const Tensor &x, Tensor y); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/blas_dot.hpp b/include/infinicore/ops/blas_dot.hpp deleted file mode 100644 index 157c167f6..000000000 --- a/include/infinicore/ops/blas_dot.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(BlasDot, const Tensor &, const Tensor &, Tensor); - -Tensor blas_dot(const Tensor &x, const Tensor &y); -void blas_dot_(const Tensor &x, const Tensor &y, Tensor result); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/block_diag.hpp b/include/infinicore/ops/block_diag.hpp deleted file mode 100644 index 89e227509..000000000 --- a/include/infinicore/ops/block_diag.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -#include - -namespace infinicore::op { - -class BlockDiag { -public: - using schema = void (*)(Tensor, const std::vector &); - static void execute(Tensor output, const std::vector &inputs); - static common::OpDispatcher &dispatcher(); -}; - -Tensor block_diag(const std::vector &inputs); -void block_diag_(Tensor output, const std::vector &inputs); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/broadcast_to.hpp b/include/infinicore/ops/broadcast_to.hpp deleted file mode 100644 index 94d98e065..000000000 --- a/include/infinicore/ops/broadcast_to.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(BroadcastTo, Tensor, Tensor); - -Tensor broadcast_to(Tensor x, const std::vector &shape); -void broadcast_to_(Tensor y, Tensor x); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/cat.hpp b/include/infinicore/ops/cat.hpp deleted file mode 100644 index 95be62103..000000000 --- a/include/infinicore/ops/cat.hpp +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include "common/op.hpp" - -namespace infinicore::op { - -Tensor cat(std::vector tensors, int dim); -void cat_(Tensor out, std::vector tensors, int dim); -} // namespace infinicore::op diff --git a/include/infinicore/ops/causal_conv1d.hpp b/include/infinicore/ops/causal_conv1d.hpp deleted file mode 100644 index ac20b779d..000000000 --- a/include/infinicore/ops/causal_conv1d.hpp +++ /dev/null @@ -1,41 +0,0 @@ -#pragma once - -#include "infinicore.h" - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(CausalConv1d, - Tensor, - Tensor, - std::optional, - const Tensor &, - const Tensor &, - std::optional, - std::optional, - std::optional, - std::optional); - -__export Tensor causal_conv1d(const Tensor &qkv, - Tensor conv_state, - const Tensor &weight, - std::optional bias = std::nullopt, - std::optional cu_seqlens = std::nullopt, - std::optional initial_state_indices = std::nullopt, - std::optional final_state_indices = std::nullopt); - -__export void causal_conv1d_(Tensor out, - Tensor conv_state, - std::optional final_conv_state, - const Tensor &qkv, - const Tensor &weight, - std::optional bias, - std::optional cu_seqlens, - std::optional initial_state_indices, - std::optional final_state_indices); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/causal_softmax.hpp b/include/infinicore/ops/causal_softmax.hpp deleted file mode 100644 index 2646852af..000000000 --- a/include/infinicore/ops/causal_softmax.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(CausalSoftmax, Tensor, const Tensor &); - -Tensor causal_softmax(const Tensor &input); -void causal_softmax_(Tensor output, const Tensor &input); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/cdist.hpp b/include/infinicore/ops/cdist.hpp deleted file mode 100644 index 7bc4902a0..000000000 --- a/include/infinicore/ops/cdist.hpp +++ /dev/null @@ -1,32 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Cdist { -public: - /** - * @brief 成对距离计算算子 (Pairwise distance) - * schema: out (M, N), x1 (M, D), x2 (N, D), p (norm degree) - */ - using schema = void (*)(Tensor, Tensor, Tensor, double); - - static void execute(Tensor out, Tensor x1, Tensor x2, double p); - - static common::OpDispatcher &dispatcher(); -}; - -/** - * @brief 非原地（Out-of-place）接口 - * @return 返回形状为 (M, N) 的新 Tensor - */ -Tensor cdist(Tensor x1, Tensor x2, double p = 2.0); - -/** - * @brief 显式指定输出接口 - */ -void cdist_(Tensor out, Tensor x1, Tensor x2, double p = 2.0); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/chunk_gated_delta_rule.hpp b/include/infinicore/ops/chunk_gated_delta_rule.hpp deleted file mode 100644 index 5102836e0..000000000 --- a/include/infinicore/ops/chunk_gated_delta_rule.hpp +++ /dev/null @@ -1,53 +0,0 @@ -#pragma once - -#include "infinicore.h" - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(ChunkGatedDeltaRule, - Tensor, - Tensor, - std::optional, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - std::optional, - std::optional, - std::optional, - bool, - size_t); - -__export Tensor chunk_gated_delta_rule(const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - Tensor initial_state, - std::optional cu_seqlens = std::nullopt, - std::optional initial_state_indices = std::nullopt, - std::optional final_state_indices = std::nullopt, - bool use_qk_l2norm = false, - size_t chunk_size = 64); - -__export void chunk_gated_delta_rule_(Tensor out, - Tensor initial_state, - std::optional final_state, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - std::optional cu_seqlens, - std::optional initial_state_indices, - std::optional final_state_indices, - bool use_qk_l2norm = false, - size_t chunk_size = 64); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/common/cache.hpp b/include/infinicore/ops/common/cache.hpp deleted file mode 100644 index ad3aa88b9..000000000 --- a/include/infinicore/ops/common/cache.hpp +++ /dev/null @@ -1,83 +0,0 @@ -#pragma once - -#include "../../common/LRUCache.hpp" -#include "../../context/context.hpp" -#include -#include -#include -#include - -namespace infinicore::op::common { - -template -class OpCache { -private: - using BaseCache = infinicore::common::LRUCache; - using Destructor = typename BaseCache::Destructor; - using CacheVector = std::vector; - -public: - explicit OpCache(size_t capacity = 100, Destructor destructor = nullptr) - : capacity_(capacity), destructor_(destructor) {} - - ~OpCache() { - clear(); - } - - BaseCache &getCache(Device::Type device_type, size_t device_index) { - auto &cache_vector = caches_[static_cast(device_type)]; - - if (cache_vector.size() <= device_index) { - cache_vector.resize(device_index + 1, BaseCache(capacity_, destructor_)); - } else { - cache_vector[device_index].setDestructor(destructor_); - } - - return cache_vector[device_index]; - } - - BaseCache &getCache(Device device) { - return getCache(device.getType(), device.getIndex()); - } - - void setCapacity(size_t capacity) { - capacity_ = capacity; - for (auto &vec : caches_) { - for (auto &cache : vec) { - cache.setCapacity(capacity); - } - } - } - - void clear() { - Device current_device = context::getDevice(); - - for (size_t type_idx = 0; type_idx < caches_.size(); ++type_idx) { - auto &vec = caches_[type_idx]; - for (size_t dev_idx = 0; dev_idx < vec.size(); ++dev_idx) { - Device target_device(static_cast(type_idx), dev_idx); - - if (current_device != target_device) { - context::setDevice(target_device); - } - - vec[dev_idx].clear(); - - if (current_device != target_device) { - context::setDevice(current_device); - } - } - vec.clear(); - } - - caches_ = {}; - } - -private: - size_t capacity_; - Destructor destructor_; - - std::array(Device::Type::COUNT)> caches_ = {}; -}; - -} // namespace infinicore::op::common diff --git a/include/infinicore/ops/common/dispatcher.hpp b/include/infinicore/ops/common/dispatcher.hpp deleted file mode 100644 index 9c5cac179..000000000 --- a/include/infinicore/ops/common/dispatcher.hpp +++ /dev/null @@ -1,80 +0,0 @@ -#pragma once - -#include "../../device.hpp" - -#ifdef ENABLE_MUTUAL_AWARENESS -#include "../../analyzer/optimization_intent.hpp" -#endif - -#include -#include - -namespace infinicore::op::common { -template -class OpDispatcher { -public: - void registerDevice(Device::Type device_type, Fn fn, bool override_existing = true) { - if (table_[(size_t)device_type] == nullptr || override_existing) { - table_[(size_t)device_type] = fn; - } - } - - void registerDevice(std::initializer_list device_types, Fn fn, bool override_existing = true) { - for (auto device_type : device_types) { - registerDevice(device_type, fn, override_existing); - } - } - - void registerAll(Fn fn, bool override_existing = true) { - for (size_t device_type = 0; device_type < static_cast(Device::Type::COUNT); ++device_type) { - registerDevice((Device::Type)device_type, fn, override_existing); - } - } - - Fn lookup(Device::Type device_type) const { - return table_.at((size_t)device_type); - } - -#ifdef ENABLE_MUTUAL_AWARENESS - // Goal-aware kernel registration. Backward compatible: callers that don't - // know about goals keep using the device-only overloads. Only kernels that - // want to specialize per OptimizationGoal need the goal-aware form. - static constexpr std::size_t kGoalCount = 4; - - void registerDevice(Device::Type device_type, - Fn fn, - analyzer::OptimizationGoal goal, - bool override_existing = true) { - std::size_t k = goalKey(device_type, goal); - if (goal_table_[k] == nullptr || override_existing) { - goal_table_[k] = fn; - } - } - - // Look up a kernel by (device, goal). If no goal-specific kernel is - // registered, fall back to the device-default kernel registered through - // the legacy lookup(device_type) path. - Fn lookup(Device::Type device_type, analyzer::OptimizationGoal goal) const { - std::size_t k = goalKey(device_type, goal); - Fn fn = goal_table_[k]; - if (fn != nullptr) { - return fn; - } - return lookup(device_type); - } - -private: - static std::size_t goalKey(Device::Type device_type, analyzer::OptimizationGoal goal) { - return static_cast(device_type) * kGoalCount - + static_cast(goal); - } - - std::array(Device::Type::COUNT) * kGoalCount> - goal_table_{}; -#endif - -private: - std::array(Device::Type::COUNT)> table_{}; -}; -} // namespace infinicore::op::common diff --git a/include/infinicore/ops/common/op.hpp b/include/infinicore/ops/common/op.hpp deleted file mode 100644 index b4f0434ef..000000000 --- a/include/infinicore/ops/common/op.hpp +++ /dev/null @@ -1,5 +0,0 @@ -#pragma once - -#include "../../context/context.hpp" -#include "../../tensor.hpp" -#include "dispatcher.hpp" diff --git a/include/infinicore/ops/conv2d.hpp b/include/infinicore/ops/conv2d.hpp deleted file mode 100644 index f1dda90ac..000000000 --- a/include/infinicore/ops/conv2d.hpp +++ /dev/null @@ -1,38 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -#include -#include - -namespace infinicore::op { -class Conv2d { -public: - using schema = void (*)(Tensor, Tensor, Tensor, Tensor, - const size_t *, const size_t *, const size_t *, size_t); - static void execute(Tensor output, - Tensor input, - Tensor weight, - Tensor bias, - const size_t *pads, - const size_t *strides, - const size_t *dilations, - size_t n); - static common::OpDispatcher &dispatcher(); -}; - -Tensor conv2d(Tensor input, - Tensor weight, - Tensor bias, - const std::vector &pads, - const std::vector &strides, - const std::vector &dilations); -void conv2d_(Tensor output, - Tensor input, - Tensor weight, - Tensor bias, - const std::vector &pads, - const std::vector &strides, - const std::vector &dilations); -} // namespace infinicore::op diff --git a/include/infinicore/ops/cross_entropy.hpp b/include/infinicore/ops/cross_entropy.hpp deleted file mode 100644 index 9a6d446d2..000000000 --- a/include/infinicore/ops/cross_entropy.hpp +++ /dev/null @@ -1,35 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class CrossEntropy { -public: - // Schema 定义：函数指针类型 - // CrossEntropy 需要接收三个 Tensor: Output (Loss), Input (Logits), Target (Labels) - using schema = void (*)(Tensor, Tensor, Tensor); - - // 执行入口 - static void execute(Tensor output, Tensor input, Tensor target); - - // 分发器访问接口 - static common::OpDispatcher &dispatcher(); -}; - -// ================================================================== -// 对外 Functional API -// ================================================================== - -// 1. Out-of-place 接口： -// 输入 Logits 和 Target，内部自动创建 Output Tensor 并返回 -Tensor cross_entropy(Tensor input, Tensor target); - -// 2. Explicit Output 接口 (类似于 In-place 风格)： -// 用户显式提供 Output Tensor 用于存储结果 -// 注意：虽然命名带有下划线 _，但通常 CrossEntropy 无法真正原地修改 input， -// 所以这里只是表示“写入指定的 output 内存” -void cross_entropy_(Tensor output, Tensor input, Tensor target); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/deepseek_moe.hpp b/include/infinicore/ops/deepseek_moe.hpp deleted file mode 100644 index e8e7ffef9..000000000 --- a/include/infinicore/ops/deepseek_moe.hpp +++ /dev/null @@ -1,41 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS( - DeepseekMoe, - Tensor, - const Tensor &, - const Tensor &, - const Tensor &, - const std::vector &, - const std::vector &, - const std::vector &, - size_t, - size_t); - -Tensor deepseek_moe(const Tensor &hidden, - const Tensor &topk_indices, - const Tensor &topk_weights, - const std::vector &gate_weights, - const std::vector &up_weights, - const std::vector &down_weights, - size_t intermediate_size, - size_t num_experts); - -void deepseek_moe_(Tensor out, - const Tensor &hidden, - const Tensor &topk_indices, - const Tensor &topk_weights, - const std::vector &gate_weights, - const std::vector &up_weights, - const std::vector &down_weights, - size_t intermediate_size, - size_t num_experts); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/dequantize_awq.hpp b/include/infinicore/ops/dequantize_awq.hpp deleted file mode 100644 index 50e4328f3..000000000 --- a/include/infinicore/ops/dequantize_awq.hpp +++ /dev/null @@ -1,10 +0,0 @@ -#pragma once -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { -INFINICORE_GRAPH_OP_CLASS(DequantizeAWQ, Tensor, const Tensor &, const Tensor &, const Tensor &); - -void dequantize_awq_(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zeros); -} // namespace infinicore::op diff --git a/include/infinicore/ops/diff.hpp b/include/infinicore/ops/diff.hpp deleted file mode 100644 index 1def7a61a..000000000 --- a/include/infinicore/ops/diff.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Diff, Tensor, const Tensor &, int, int); - -Tensor diff(const Tensor &x, int n = 1, int dim = -1); -void diff_(Tensor y, const Tensor &x, int n = 1, int dim = -1); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/digamma.hpp b/include/infinicore/ops/digamma.hpp deleted file mode 100644 index b61fc40f6..000000000 --- a/include/infinicore/ops/digamma.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Digamma, Tensor, const Tensor &); - -Tensor digamma(const Tensor &x); -void digamma_(Tensor y, const Tensor &x); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/dist.hpp b/include/infinicore/ops/dist.hpp deleted file mode 100644 index 459a52202..000000000 --- a/include/infinicore/ops/dist.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Dist, Tensor, const Tensor &, const Tensor &, double); - -Tensor dist(const Tensor &x1, const Tensor &x2, double p = 2.0); -void dist_(Tensor y, const Tensor &x1, const Tensor &x2, double p = 2.0); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/distributed/allgather.hpp b/include/infinicore/ops/distributed/allgather.hpp deleted file mode 100644 index f206cab32..000000000 --- a/include/infinicore/ops/distributed/allgather.hpp +++ /dev/null @@ -1,52 +0,0 @@ -#pragma once - -#include "../../device.hpp" -#include "../../graph/graph.hpp" -#include "../common/op.hpp" - -#include -#include - -namespace infinicore::op::distributed { - -class AllGather : public graph::GraphOperator { -public: - AllGather(Tensor output, const Tensor &input, infinicclComm_t communicator); - ~AllGather(); - void run() const override; - static void execute(Tensor output, const Tensor &input, infinicclComm_t communicator); - -private: - void *planned_meta_; -}; - -class AllGatherV : public graph::GraphOperator { -public: - AllGatherV(Tensor output, - const Tensor &input, - std::vector split_counts, - infinicclComm_t communicator); - ~AllGatherV(); - void run() const override; - static void execute(Tensor output, - const Tensor &input, - std::vector split_counts, - infinicclComm_t communicator); - -private: - void *planned_meta_; -}; - -Tensor allgather(const Tensor &input, size_t world_size, infinicclComm_t communicator); -void allgather_(Tensor output, const Tensor &input, infinicclComm_t communicator); -Tensor allgatherv(const Tensor &input, const std::vector &split_sizes, infinicclComm_t communicator); -void allgatherv_(Tensor output, const Tensor &input, const std::vector &split_sizes, infinicclComm_t communicator); -std::vector allgatherv_many(const std::vector &inputs, - const std::vector &split_sizes, - infinicclComm_t communicator); -void allgatherv_many_(const std::vector &outputs, - const std::vector &inputs, - const std::vector &split_sizes, - infinicclComm_t communicator); - -} // namespace infinicore::op::distributed diff --git a/include/infinicore/ops/distributed/allreduce.hpp b/include/infinicore/ops/distributed/allreduce.hpp deleted file mode 100644 index 39f74243a..000000000 --- a/include/infinicore/ops/distributed/allreduce.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include "../../device.hpp" -#include "../../graph/graph.hpp" -#include "../common/op.hpp" - -#include - -namespace infinicore::op::distributed { -class AllReduce : public graph::GraphOperator { -public: - AllReduce(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); - ~AllReduce(); - void run() const override; - static void execute(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); - -private: - void *planned_meta_; -}; - -Tensor allreduce(const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); -void allreduce_(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); - -} // namespace infinicore::op::distributed diff --git a/include/infinicore/ops/distributed/reduce_scatter.hpp b/include/infinicore/ops/distributed/reduce_scatter.hpp deleted file mode 100644 index ee5622872..000000000 --- a/include/infinicore/ops/distributed/reduce_scatter.hpp +++ /dev/null @@ -1,65 +0,0 @@ -#pragma once - -#include "../../device.hpp" -#include "../../graph/graph.hpp" -#include "../common/op.hpp" - -#include -#include - -namespace infinicore::op::distributed { - -class ReduceScatter : public graph::GraphOperator { -public: - ReduceScatter(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); - ~ReduceScatter(); - void run() const override; - static void execute(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); - -private: - void *planned_meta_; -}; - -class ReduceScatterV : public graph::GraphOperator { -public: - ReduceScatterV(Tensor output, - const Tensor &input, - std::vector split_counts, - infinicclReduceOp_t op, - infinicclComm_t communicator); - ~ReduceScatterV(); - void run() const override; - static void execute(Tensor output, - const Tensor &input, - std::vector split_counts, - infinicclReduceOp_t op, - infinicclComm_t communicator); - -private: - void *planned_meta_; -}; - -Tensor reduce_scatter(const Tensor &input, size_t world_size, infinicclReduceOp_t op, infinicclComm_t communicator); -void reduce_scatter_(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator); -Tensor reduce_scatterv(const Tensor &input, - const std::vector &split_sizes, - size_t rank, - infinicclReduceOp_t op, - infinicclComm_t communicator); -void reduce_scatterv_(Tensor output, - const Tensor &input, - const std::vector &split_sizes, - infinicclReduceOp_t op, - infinicclComm_t communicator); -std::vector reduce_scatterv_many(const std::vector &inputs, - const std::vector &split_sizes, - size_t rank, - infinicclReduceOp_t op, - infinicclComm_t communicator); -void reduce_scatterv_many_(const std::vector &outputs, - const std::vector &inputs, - const std::vector &split_sizes, - infinicclReduceOp_t op, - infinicclComm_t communicator); - -} // namespace infinicore::op::distributed diff --git a/include/infinicore/ops/embedding.hpp b/include/infinicore/ops/embedding.hpp deleted file mode 100644 index 43f18d090..000000000 --- a/include/infinicore/ops/embedding.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Embedding, Tensor, const Tensor &, const Tensor &); - -Tensor embedding(const Tensor &input, const Tensor &weight); -void embedding_(Tensor out, const Tensor &input, const Tensor &weight); -} // namespace infinicore::op diff --git a/include/infinicore/ops/equal.hpp b/include/infinicore/ops/equal.hpp deleted file mode 100644 index 1a158bf1e..000000000 --- a/include/infinicore/ops/equal.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Equal { -public: - using schema = void (*)(Tensor, Tensor, Tensor); - - static void execute(Tensor out, Tensor a, Tensor b); - static common::OpDispatcher &dispatcher(); -}; - -Tensor equal(Tensor a, Tensor b); -void equal_(Tensor out, Tensor a, Tensor b); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/flash_attention.hpp b/include/infinicore/ops/flash_attention.hpp deleted file mode 100644 index 24e33cfb6..000000000 --- a/include/infinicore/ops/flash_attention.hpp +++ /dev/null @@ -1,12 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(FlashAttention, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, float, bool); - -Tensor flash_attention(const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal); -void flash_attention_(Tensor out, const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal); -} // namespace infinicore::op diff --git a/include/infinicore/ops/flipud.hpp b/include/infinicore/ops/flipud.hpp deleted file mode 100644 index 7e449f9f5..000000000 --- a/include/infinicore/ops/flipud.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Flipud { -public: - // Schema signature: (Output, Input) - using schema = void (*)(Tensor, Tensor); - - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; -Tensor flipud(Tensor input); -void flipud_(Tensor output, Tensor input); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/float_power.hpp b/include/infinicore/ops/float_power.hpp deleted file mode 100644 index c461db2eb..000000000 --- a/include/infinicore/ops/float_power.hpp +++ /dev/null @@ -1,68 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class FloatPower { -public: - // ========================================================== - // Dispatcher Schemas - // ========================================================== - - // Output = Input ^ Scalar (scalar must be double!) - using schema_scalar = void (*)(Tensor output, - Tensor input, - double exponent); - - // Output = Input ^ Tensor - using schema_tensor = void (*)(Tensor output, - Tensor input, - Tensor exponent); - - // ========================================================== - // Execute Entry Points (called by functional interface) - // ========================================================== - - static void execute(Tensor output, - Tensor input, - double exponent); - - static void execute(Tensor output, - Tensor input, - Tensor exponent); - - // ========================================================== - // Dispatchers - // ========================================================== - - static common::OpDispatcher &dispatcher_scalar(); - static common::OpDispatcher &dispatcher_tensor(); -}; - -// ======================================================================= -// Functional Interface (Python-visible semantics) -// ======================================================================= - -// ------------------------------- -// 1. Scalar Exponent -// ------------------------------- - -// out-of-place: ALWAYS float64 -Tensor float_power(Tensor input, double exponent); - -// in-place -void float_power_(Tensor output, Tensor input, double exponent); - -// ------------------------------- -// 2. Tensor Exponent -// ------------------------------- - -// out-of-place: ALWAYS float64 -Tensor float_power(Tensor input, Tensor exponent); - -// in-place -void float_power_(Tensor output, Tensor input, Tensor exponent); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/floor.hpp b/include/infinicore/ops/floor.hpp deleted file mode 100644 index 11b52571b..000000000 --- a/include/infinicore/ops/floor.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Floor { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor floor(Tensor input); -void floor_(Tensor output, Tensor input); -} // namespace infinicore::op diff --git a/include/infinicore/ops/floor_divide.hpp b/include/infinicore/ops/floor_divide.hpp deleted file mode 100644 index 43267dce6..000000000 --- a/include/infinicore/ops/floor_divide.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class FloorDivide { -public: - using schema = void (*)(Tensor, Tensor, Tensor); - static void execute(Tensor c, Tensor a, Tensor b); - static common::OpDispatcher &dispatcher(); -}; - -Tensor floor_divide(Tensor a, Tensor b); -void floor_divide_(Tensor c, Tensor a, Tensor b); -} // namespace infinicore::op diff --git a/include/infinicore/ops/fmin.hpp b/include/infinicore/ops/fmin.hpp deleted file mode 100644 index 4ea00787f..000000000 --- a/include/infinicore/ops/fmin.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Fmin { -public: - using schema = void (*)(Tensor, Tensor, Tensor); - static void execute(Tensor c, Tensor a, Tensor b); - static common::OpDispatcher &dispatcher(); -}; - -Tensor fmin(Tensor a, Tensor b); -void fmin_(Tensor c, Tensor a, Tensor b); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/fmod.hpp b/include/infinicore/ops/fmod.hpp deleted file mode 100644 index 87b90d515..000000000 --- a/include/infinicore/ops/fmod.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Fmod { -public: - using schema = void (*)(Tensor, Tensor, Tensor); - static void execute(Tensor c, Tensor a, Tensor b); - static common::OpDispatcher &dispatcher(); -}; - -Tensor fmod(Tensor a, Tensor b); -void fmod_(Tensor c, Tensor a, Tensor b); -} // namespace infinicore::op diff --git a/include/infinicore/ops/fused_gated_delta_net_gating.hpp b/include/infinicore/ops/fused_gated_delta_net_gating.hpp deleted file mode 100644 index a61fd9937..000000000 --- a/include/infinicore/ops/fused_gated_delta_net_gating.hpp +++ /dev/null @@ -1,37 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(FusedGatedDeltaNetGating, - Tensor, - Tensor, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - float, - float); - -std::pair fused_gated_delta_net_gating(const Tensor &A_log, - const Tensor &a, - const Tensor &b, - const Tensor &dt_bias, - float beta = 1.0f, - float threshold = 20.0f); - -void fused_gated_delta_net_gating_(Tensor g, - Tensor beta_output, - const Tensor &A_log, - const Tensor &a, - const Tensor &b, - const Tensor &dt_bias, - float beta = 1.0f, - float threshold = 20.0f); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/gaussian_nll_loss.hpp b/include/infinicore/ops/gaussian_nll_loss.hpp deleted file mode 100644 index c40eb1e86..000000000 --- a/include/infinicore/ops/gaussian_nll_loss.hpp +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -#include "infinicore.h" - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(GaussianNllLoss, Tensor, const Tensor &, const Tensor &, const Tensor &, bool, double, int); - -__export Tensor gaussian_nll_loss(const Tensor &input, - const Tensor &target, - const Tensor &var, - bool full = false, - double eps = 1e-6, - int reduction = 1); - -__export void gaussian_nll_loss_(Tensor out, - const Tensor &input, - const Tensor &target, - const Tensor &var, - bool full = false, - double eps = 1e-6, - int reduction = 1); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/gelu.hpp b/include/infinicore/ops/gelu.hpp deleted file mode 100644 index 5e8c61347..000000000 --- a/include/infinicore/ops/gelu.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Gelu { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor gelu(Tensor input); -void gelu_(Tensor output, Tensor input); -} // namespace infinicore::op diff --git a/include/infinicore/ops/gelutanh.hpp b/include/infinicore/ops/gelutanh.hpp deleted file mode 100644 index c968c81fd..000000000 --- a/include/infinicore/ops/gelutanh.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class GeluTanh { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor gelu_tanh(Tensor input); -void gelu_tanh_(Tensor output, Tensor input); -} // namespace infinicore::op diff --git a/include/infinicore/ops/gemm.hpp b/include/infinicore/ops/gemm.hpp deleted file mode 100644 index 4f76cee26..000000000 --- a/include/infinicore/ops/gemm.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Gemm, Tensor, const Tensor &, const Tensor &, float, float); - -Tensor gemm(const Tensor &a, const Tensor &b, float alpha = 1.0f, float beta = 0.0f); -void gemm_(Tensor c, const Tensor &a, const Tensor &b, float alpha, float beta); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/gptq_marlin_gemm.hpp b/include/infinicore/ops/gptq_marlin_gemm.hpp deleted file mode 100644 index b16df0ad6..000000000 --- a/include/infinicore/ops/gptq_marlin_gemm.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(GptqMarlinGemm, Tensor, const Tensor &, const Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, int64_t, bool, bool, bool, bool); -INFINICORE_GRAPH_OP_CLASS(GptqMarlinGemmWithWorkspace, Tensor, Tensor, const Tensor &, const Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, Tensor &, int64_t, bool, bool, bool, bool); - -void gptq_marlin_gemm_(Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float); - -size_t gptq_marlin_gemm_workspace_size(Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm); -void gptq_marlin_gemm_with_workspace_(Tensor workspace, Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float); -void gptq_marlin_gemm_with_workspace_direct_(Tensor workspace, Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float); -} // namespace infinicore::op diff --git a/include/infinicore/ops/hardswish.hpp b/include/infinicore/ops/hardswish.hpp deleted file mode 100644 index 15313f461..000000000 --- a/include/infinicore/ops/hardswish.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Hardswish { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor hardswish(Tensor input); -void hardswish_(Tensor output, Tensor input); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/hardtanh.hpp b/include/infinicore/ops/hardtanh.hpp deleted file mode 100644 index 511408fee..000000000 --- a/include/infinicore/ops/hardtanh.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class HardTanh { -public: - using schema = void (*)(Tensor, Tensor, float, float); - static void execute(Tensor output, Tensor input, float min_val, float max_val); - static common::OpDispatcher &dispatcher(); -}; - -Tensor hardtanh(Tensor input, float min_val = -1.0f, float max_val = 1.0f); -void hardtanh_(Tensor output, Tensor input, float min_val = -1.0f, float max_val = 1.0f); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/hinge_embedding_loss.hpp b/include/infinicore/ops/hinge_embedding_loss.hpp deleted file mode 100644 index 9d0fcfbab..000000000 --- a/include/infinicore/ops/hinge_embedding_loss.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class HingeEmbeddingLoss { -public: - using schema = void (*)(Tensor, Tensor, Tensor, double, int); - static void execute(Tensor output, Tensor input, Tensor target, double margin, int reduction); - static common::OpDispatcher &dispatcher(); -}; - -Tensor hinge_embedding_loss(Tensor input, Tensor target, double margin = 1.0, int reduction = 1); -void hinge_embedding_loss_(Tensor output, Tensor input, Tensor target, double margin, int reduction); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/huber_loss.hpp b/include/infinicore/ops/huber_loss.hpp deleted file mode 100644 index c6212f0c1..000000000 --- a/include/infinicore/ops/huber_loss.hpp +++ /dev/null @@ -1,21 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class HuberLoss { -public: - // Schema: output, input, target, delta, reduction - using schema = void (*)(Tensor, Tensor, Tensor, float, int64_t); - - static void execute(Tensor output, Tensor input, Tensor target, float delta, int64_t reduction); - static common::OpDispatcher &dispatcher(); -}; - -// delta 默认为 1.0f，reduction 默认为 1 (MEAN) -Tensor huber_loss(Tensor input, Tensor target, float delta = 1.0f, int64_t reduction = 1); -void huber_loss_(Tensor output, Tensor input, Tensor target, float delta, int64_t reduction); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/hypot.hpp b/include/infinicore/ops/hypot.hpp deleted file mode 100644 index 24eebd44f..000000000 --- a/include/infinicore/ops/hypot.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Hypot { -public: - using schema = void (*)(Tensor, Tensor, Tensor); - - static void execute(Tensor output, Tensor input_a, Tensor input_b); - static common::OpDispatcher &dispatcher(); -}; - -Tensor hypot(Tensor input_a, Tensor input_b); - -void hypot_(Tensor output, Tensor input_a, Tensor input_b); -} // namespace infinicore::op diff --git a/include/infinicore/ops/index_add.hpp b/include/infinicore/ops/index_add.hpp deleted file mode 100644 index e880114e9..000000000 --- a/include/infinicore/ops/index_add.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class IndexAdd { -public: - using schema = void (*)(Tensor, Tensor, int64_t, Tensor, Tensor, float); - static void execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source, float alpha); - - static common::OpDispatcher &dispatcher(); -}; - -Tensor index_add(Tensor input, int64_t dim, Tensor index, Tensor source, float alpha = 1.0f); -void index_add_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source, float alpha); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/index_copy.hpp b/include/infinicore/ops/index_copy.hpp deleted file mode 100644 index b2437dead..000000000 --- a/include/infinicore/ops/index_copy.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class IndexCopy { -public: - using schema = void (*)(Tensor, Tensor, int64_t, Tensor, Tensor); - static void execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source); - - static common::OpDispatcher &dispatcher(); -}; -Tensor index_copy(Tensor input, int64_t dim, Tensor index, Tensor source); -void index_copy_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/inner.hpp b/include/infinicore/ops/inner.hpp deleted file mode 100644 index ae372c415..000000000 --- a/include/infinicore/ops/inner.hpp +++ /dev/null @@ -1,17 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Inner { -public: - using schema = void (*)(Tensor, Tensor, Tensor); - static void execute(Tensor out, Tensor input, Tensor other); - static common::OpDispatcher &dispatcher(); -}; - -Tensor inner(Tensor input, Tensor other); -void inner_(Tensor out, Tensor input, Tensor other); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/interpolate.hpp b/include/infinicore/ops/interpolate.hpp deleted file mode 100644 index 2a235ef3d..000000000 --- a/include/infinicore/ops/interpolate.hpp +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once - -#include "infinicore.h" - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -#include -#include -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Interpolate, Tensor, const Tensor &, std::string, std::vector, std::vector, int); - -__export Tensor interpolate(const Tensor &input, - std::string mode, - std::vector size, - std::vector scale_factor, - int align_corners); - -__export void interpolate_(Tensor out, - const Tensor &input, - std::string mode, - std::vector size, - std::vector scale_factor, - int align_corners); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/kron.hpp b/include/infinicore/ops/kron.hpp deleted file mode 100644 index 3a6834424..000000000 --- a/include/infinicore/ops/kron.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Kron { -public: - using schema = void (*)(Tensor, Tensor, Tensor); - static void execute(Tensor output, Tensor a, Tensor b); - static common::OpDispatcher &dispatcher(); -}; - -Tensor kron(Tensor a, Tensor b); -void kron_(Tensor output, Tensor a, Tensor b); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/kthvalue.hpp b/include/infinicore/ops/kthvalue.hpp deleted file mode 100644 index 1e9721d28..000000000 --- a/include/infinicore/ops/kthvalue.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -class Kthvalue { -public: - // Schema signature: values(out), indices(out), input, k, dim, keepdim - using schema = void (*)(Tensor, Tensor, Tensor, int64_t, int64_t, bool); - - static void execute(Tensor values, Tensor indices, Tensor input, int64_t k, int64_t dim, bool keepdim); - static common::OpDispatcher &dispatcher(); -}; - -// Functional API: Returns a tuple containing (values, indices) -std::tuple kthvalue(Tensor input, int64_t k, int64_t dim = -1, bool keepdim = false); - -// In-place/Output-provided API -void kthvalue_(Tensor values, Tensor indices, Tensor input, int64_t k, int64_t dim, bool keepdim); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/kv_caching.hpp b/include/infinicore/ops/kv_caching.hpp deleted file mode 100644 index 3a70c2824..000000000 --- a/include/infinicore/ops/kv_caching.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(KVCaching, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &); - -void kv_caching_(Tensor k_cache, - Tensor v_cache, - const Tensor &k, - const Tensor &v, - const Tensor &past_kv_lengths); -} // namespace infinicore::op diff --git a/include/infinicore/ops/layer_norm.hpp b/include/infinicore/ops/layer_norm.hpp deleted file mode 100644 index da6256b51..000000000 --- a/include/infinicore/ops/layer_norm.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(LayerNorm, Tensor, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &, float); - -Tensor layer_norm(const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon = 1e-5f); -void layer_norm_(Tensor y, Tensor standardization, Tensor std_deviation, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon = 1e-5f); -void layer_norm_(Tensor y, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon = 1e-5f); -void layer_norm_for_pybind(Tensor y, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon = 1e-5f); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/ldexp.hpp b/include/infinicore/ops/ldexp.hpp deleted file mode 100644 index bac7964c2..000000000 --- a/include/infinicore/ops/ldexp.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Ldexp { -public: - // Schema signature: output(out), input(x), other(exp) - using schema = void (*)(Tensor, Tensor, Tensor); - - static void execute(Tensor output, Tensor input, Tensor other); - static common::OpDispatcher &dispatcher(); -}; - -// Functional API: Returns a new Tensor containing input * (2^other) -Tensor ldexp(Tensor input, Tensor other); - -// In-place/Output-provided API -// Writes the result into 'output' -void ldexp_(Tensor output, Tensor input, Tensor other); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/lerp.hpp b/include/infinicore/ops/lerp.hpp deleted file mode 100644 index bf3e5cbec..000000000 --- a/include/infinicore/ops/lerp.hpp +++ /dev/null @@ -1,27 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Lerp { -public: - using schema_t = void (*)(Tensor, Tensor, Tensor, Tensor); - using schema_s = void (*)(Tensor, Tensor, Tensor, float); - - static void execute(Tensor output, Tensor start, Tensor end, Tensor weight); - static void execute(Tensor output, Tensor start, Tensor end, float weight); - - // 【核心修改】必须声明为模板函数，才能支持 dispatcher() 和 dispatcher() - template - static common::OpDispatcher &dispatcher(); -}; - -Tensor lerp(Tensor start, Tensor end, Tensor weight); -Tensor lerp(Tensor start, Tensor end, float weight); - -void lerp_(Tensor output, Tensor start, Tensor end, Tensor weight); -void lerp_(Tensor output, Tensor start, Tensor end, float weight); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/linear.hpp b/include/infinicore/ops/linear.hpp deleted file mode 100644 index 96ca733e4..000000000 --- a/include/infinicore/ops/linear.hpp +++ /dev/null @@ -1,12 +0,0 @@ -#pragma once - -#include "common/op.hpp" -#include - -namespace infinicore::op { - -Tensor linear(Tensor input, Tensor weight, std::optional bias, float alpha = 1.0f); - -void linear_(Tensor out, Tensor input, Tensor weight, std::optional bias, float alpha = 1.0f); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/linear_w4a16_awq.hpp b/include/infinicore/ops/linear_w4a16_awq.hpp deleted file mode 100644 index ebae7a685..000000000 --- a/include/infinicore/ops/linear_w4a16_awq.hpp +++ /dev/null @@ -1,12 +0,0 @@ -#pragma once - -#include "common/op.hpp" -#include - -namespace infinicore::op { - -Tensor linear_w4a16_awq(Tensor input, Tensor weight_packed, Tensor weight_scale, Tensor weight_zeros, std::optional bias); - -void linear_w4a16_awq_(Tensor out, Tensor input, Tensor weight_packed, Tensor weight_scale, Tensor weight_zeros, std::optional bias); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/linear_w4a16_gptq_qy.hpp b/include/infinicore/ops/linear_w4a16_gptq_qy.hpp deleted file mode 100644 index 72cc9eef4..000000000 --- a/include/infinicore/ops/linear_w4a16_gptq_qy.hpp +++ /dev/null @@ -1,12 +0,0 @@ -#pragma once - -#include "common/op.hpp" -#include - -namespace infinicore::op { - -Tensor linear_w4a16_gptq_qy(Tensor in, Tensor qweight, Tensor qzeros, Tensor scales, int64_t quant_type, int64_t bit); - -void linear_w4a16_gptq_qy_(Tensor out, Tensor in, Tensor qweights, Tensor scales, Tensor qzeros, int64_t quant_type, int64_t bit); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/linear_w8a8i8.hpp b/include/infinicore/ops/linear_w8a8i8.hpp deleted file mode 100644 index 08cadc111..000000000 --- a/include/infinicore/ops/linear_w8a8i8.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "../graph/graph.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -Tensor linear_w8a8i8(Tensor input, Tensor weight_packed, Tensor weight_scale, std::optional bias); - -void linear_w8a8i8_(Tensor out, Tensor input, Tensor weight_packed, Tensor weight_scale, std::optional bias); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/log_softmax.hpp b/include/infinicore/ops/log_softmax.hpp deleted file mode 100644 index 00e1c90c1..000000000 --- a/include/infinicore/ops/log_softmax.hpp +++ /dev/null @@ -1,23 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class LogSoftmax { -public: - // Schema signature: output(out), input, dim - using schema = void (*)(Tensor, Tensor, int64_t); - - static void execute(Tensor output, Tensor input, int64_t dim); - static common::OpDispatcher &dispatcher(); -}; - -// Functional API: Returns the result tensor -Tensor log_softmax(Tensor input, int64_t dim); - -// In-place/Output-provided API -void log_softmax_(Tensor output, Tensor input, int64_t dim); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/logaddexp.hpp b/include/infinicore/ops/logaddexp.hpp deleted file mode 100644 index fc84c10c2..000000000 --- a/include/infinicore/ops/logaddexp.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class LogAddExp { -public: - using schema = void (*)(Tensor, Tensor, Tensor); - static void execute(Tensor c, Tensor a, Tensor b); - static common::OpDispatcher &dispatcher(); -}; - -Tensor logaddexp(Tensor a, Tensor b); -void logaddexp_(Tensor c, Tensor a, Tensor b); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/logaddexp2.hpp b/include/infinicore/ops/logaddexp2.hpp deleted file mode 100644 index 7c7dbe96e..000000000 --- a/include/infinicore/ops/logaddexp2.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class LogAddExp2 { -public: - using schema = void (*)(Tensor, Tensor, Tensor); - static void execute(Tensor c, Tensor a, Tensor b); - static common::OpDispatcher &dispatcher(); -}; - -Tensor logaddexp2(Tensor a, Tensor b); -void logaddexp2_(Tensor c, Tensor a, Tensor b); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/logcumsumexp.hpp b/include/infinicore/ops/logcumsumexp.hpp deleted file mode 100644 index 212ff625d..000000000 --- a/include/infinicore/ops/logcumsumexp.hpp +++ /dev/null @@ -1,21 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class LogCumSumExp { -public: - using schema = void (*)(Tensor, Tensor, int, bool, bool); - - static void execute(Tensor y, Tensor x, int axis, bool exclusive, bool reverse); - - static common::OpDispatcher &dispatcher(); -}; - -Tensor logcumsumexp(Tensor x, int axis, bool exclusive = false, bool reverse = false); - -void logcumsumexp_(Tensor y, Tensor x, int axis, bool exclusive = false, bool reverse = false); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/logdet.hpp b/include/infinicore/ops/logdet.hpp deleted file mode 100644 index 69c1b6d0d..000000000 --- a/include/infinicore/ops/logdet.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Logdet, Tensor, const Tensor &); - -Tensor logdet(const Tensor &x); -void logdet_(Tensor y, const Tensor &x); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/logical_and.hpp b/include/infinicore/ops/logical_and.hpp deleted file mode 100644 index 41f45f1c6..000000000 --- a/include/infinicore/ops/logical_and.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class LogicalAnd { -public: - // LogicalAnd 是二元操作，schema 通常定义为 (Output, Input1, Input2) - using schema = void (*)(Tensor, Tensor, Tensor); - - static void execute(Tensor output, Tensor input1, Tensor input2); - static common::OpDispatcher &dispatcher(); -}; - -Tensor logical_and(Tensor input1, Tensor input2); -void logical_and_(Tensor output, Tensor input1, Tensor input2); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/logical_not.hpp b/include/infinicore/ops/logical_not.hpp deleted file mode 100644 index caf6bbb76..000000000 --- a/include/infinicore/ops/logical_not.hpp +++ /dev/null @@ -1,23 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class LogicalNot { -public: - // LogicalNot 是一元操作，schema 定义为 (Output, Input) - using schema = void (*)(Tensor, Tensor); - - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -// 构造新 Tensor 返回结果 -Tensor logical_not(Tensor input); - -// 将结果写入指定的 output Tensor -void logical_not_(Tensor output, Tensor input); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/mamba_selective_scan.hpp b/include/infinicore/ops/mamba_selective_scan.hpp deleted file mode 100644 index 3acdb1c4d..000000000 --- a/include/infinicore/ops/mamba_selective_scan.hpp +++ /dev/null @@ -1,43 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS( - MambaSelectiveScan, - Tensor, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - Tensor); - -Tensor mamba_selective_scan(const Tensor &x, - const Tensor &dt, - const Tensor &b, - const Tensor &c, - const Tensor &a_log, - const Tensor &d, - const Tensor &gate, - const Tensor &dt_bias, - Tensor state); - -void mamba_selective_scan_(Tensor out, - const Tensor &x, - const Tensor &dt, - const Tensor &b, - const Tensor &c, - const Tensor &a_log, - const Tensor &d, - const Tensor &gate, - const Tensor &dt_bias, - Tensor state); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/masked_select.hpp b/include/infinicore/ops/masked_select.hpp deleted file mode 100644 index d003e3bb8..000000000 --- a/include/infinicore/ops/masked_select.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class MaskedSelect { -public: - using schema = void (*)(Tensor, Tensor, void **, size_t *); - static void execute(Tensor input, Tensor mask, void **data_ptr, size_t *dlen_ptr); - static common::OpDispatcher &dispatcher(); -}; - -Tensor masked_select(Tensor input, Tensor mask); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/matmul.hpp b/include/infinicore/ops/matmul.hpp deleted file mode 100644 index 1fb41774e..000000000 --- a/include/infinicore/ops/matmul.hpp +++ /dev/null @@ -1,11 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -Tensor matmul(Tensor a, Tensor b, float alpha = 1.0f); -void matmul_(Tensor c, Tensor a, Tensor b, float alpha = 1.0f); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/mha.hpp b/include/infinicore/ops/mha.hpp deleted file mode 100644 index f62598c78..000000000 --- a/include/infinicore/ops/mha.hpp +++ /dev/null @@ -1,34 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS( - MultiheadAttention, - Tensor, - const Tensor &, - const Tensor &, - const Tensor &, - std::optional, - float, - bool); - -Tensor mha(const Tensor &q, - const Tensor &k, - const Tensor &v, - std::optional alibi_slopes, - float scale, - bool is_causal); - -void mha_(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - std::optional alibi_slopes, - float scale, - bool is_causal); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/mha_kvcache.hpp b/include/infinicore/ops/mha_kvcache.hpp deleted file mode 100644 index 69a11f61c..000000000 --- a/include/infinicore/ops/mha_kvcache.hpp +++ /dev/null @@ -1,51 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -// Flash Attention KV-cache decode op. -// -// Wraps flash::mha_fwd_kvcache for single-step (decode) attention over a -// paged KV cache. -// -// Tensor shapes: -// out : [batch_size, seqlen_q, num_heads, head_size] -// q : [batch_size, seqlen_q, num_heads, head_size] -// k_cache : [num_blocks, block_size, num_heads_k, head_size] (paged layout) -// v_cache : [num_blocks, block_size, num_heads_k, head_size] (paged layout) -// seqlens_k : [batch_size] int32 — total KV length per request -// block_table : [batch_size, max_num_blocks_per_seq] int32 - -INFINICORE_GRAPH_OP_CLASS( - MhaKVCache, - Tensor, // out - const Tensor &, // q - const Tensor &, // k_cache - const Tensor &, // v_cache - const Tensor &, // seqlens_k - const Tensor &, // block_table - std::optional, // alibi_slopes - float); // scale - -Tensor mha_kvcache(const Tensor &q, - const Tensor &k_cache, - const Tensor &v_cache, - const Tensor &seqlens_k, - const Tensor &block_table, - std::optional alibi_slopes, - float scale); - -void mha_kvcache_(Tensor out, - const Tensor &q, - const Tensor &k_cache, - const Tensor &v_cache, - const Tensor &seqlens_k, - const Tensor &block_table, - std::optional alibi_slopes, - float scale); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/mha_varlen.hpp b/include/infinicore/ops/mha_varlen.hpp deleted file mode 100644 index 4d226f383..000000000 --- a/include/infinicore/ops/mha_varlen.hpp +++ /dev/null @@ -1,46 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS( - MultiheadAttentionVarlen, - Tensor, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - std::optional, - int, - int, - std::optional, - float); - -Tensor mha_varlen(const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &cum_seqlens_q, - const Tensor &cum_seqlens_k, - std::optional block_table, - int max_seqlen_q, - int max_seqlen_k, - std::optional alibi_slopes, - float scale); - -void mha_varlen_(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &cum_seqlens_q, - const Tensor &cum_seqlens_k, - std::optional block_table, - int max_seqlen_q, - int max_seqlen_k, - std::optional alibi_slopes, - float scale); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/moe_align.hpp b/include/infinicore/ops/moe_align.hpp deleted file mode 100644 index 2ce7ca8f4..000000000 --- a/include/infinicore/ops/moe_align.hpp +++ /dev/null @@ -1,40 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "../tensor.hpp" -#include "common/op.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(MoeAlign, Tensor, Tensor, Tensor, const Tensor &, const size_t, const size_t, const bool); -INFINICORE_GRAPH_OP_CLASS(MoeAlignWithExpertMap, Tensor, Tensor, Tensor, const Tensor &, const Tensor &, const size_t, const size_t, const bool); - -std::tuple moe_align( - const Tensor &topk_ids, - size_t num_experts, - size_t block_size, - bool pad_sorted_token_ids = true); - -void moe_align_( - Tensor sorted_token_ids, - Tensor expert_ids, - Tensor num_tokens_post_padded, - const Tensor &topk_ids, - size_t num_experts, - size_t block_size, - bool pad_sorted_token_ids = true); - -void moe_align_with_expert_map_( - Tensor sorted_token_ids, - Tensor expert_ids, - Tensor num_tokens_post_padded, - const Tensor &topk_ids, - const Tensor &expert_map, - size_t num_experts, - size_t block_size, - bool pad_sorted_token_ids = true); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/moe_fused_dense.hpp b/include/infinicore/ops/moe_fused_dense.hpp deleted file mode 100644 index 0a34c734c..000000000 --- a/include/infinicore/ops/moe_fused_dense.hpp +++ /dev/null @@ -1,42 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "../tensor.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(MoeFusedDense, - Tensor, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &); - -Tensor moe_fused_dense( - const Tensor &hidden_states, - const Tensor &w13, - const Tensor &w2, - const Tensor &topk_weights, - const Tensor &topk_ids, - const Tensor &sorted_token_ids, - const Tensor &expert_ids, - const Tensor &num_tokens_post_padded); - -void moe_fused_dense_( - Tensor output, - const Tensor &hidden_states, - const Tensor &w13, - const Tensor &w2, - const Tensor &topk_weights, - const Tensor &topk_ids, - const Tensor &sorted_token_ids, - const Tensor &expert_ids, - const Tensor &num_tokens_post_padded); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/moe_fused_gate.hpp b/include/infinicore/ops/moe_fused_gate.hpp deleted file mode 100644 index e20d0a657..000000000 --- a/include/infinicore/ops/moe_fused_gate.hpp +++ /dev/null @@ -1,44 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "../tensor.hpp" -#include "common/op.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(MoeFusedGate, - Tensor, - Tensor, - const Tensor &, - const Tensor &, - const size_t, - const size_t, - const size_t, - const float, - const bool); - -std::tuple moe_fused_gate( - const Tensor &input, - const Tensor &bias, - size_t topk, - size_t num_expert_group, - size_t topk_group, - size_t num_fused_shared_experts, - float routed_scaling_factor, - bool apply_routed_scaling_factor_on_output); - -void moe_fused_gate_( - Tensor topk_weights, - Tensor topk_indices, - const Tensor &input, - const Tensor &bias, - size_t num_expert_group, - size_t topk_group, - size_t num_fused_shared_experts, - float routed_scaling_factor, - bool apply_routed_scaling_factor_on_output); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/moe_sum.hpp b/include/infinicore/ops/moe_sum.hpp deleted file mode 100644 index e9623e3b1..000000000 --- a/include/infinicore/ops/moe_sum.hpp +++ /dev/null @@ -1,15 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "../tensor.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(MoeSum, Tensor, const Tensor &); - -Tensor moe_sum(const Tensor &input); -void moe_sum_(Tensor output, const Tensor &input); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/moe_topk_sigmoid.hpp b/include/infinicore/ops/moe_topk_sigmoid.hpp deleted file mode 100644 index e796c1dc8..000000000 --- a/include/infinicore/ops/moe_topk_sigmoid.hpp +++ /dev/null @@ -1,32 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "../tensor.hpp" -#include "common/op.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(MoeTopkSigmoid, - Tensor, - Tensor, - const Tensor &, - const Tensor &, - const bool); - -std::tuple moe_topk_sigmoid( - const Tensor &gating_output, - size_t topk, - bool renormalize = false, - const Tensor &correction_bias = Tensor()); - -void moe_topk_sigmoid_( - Tensor topk_weights, - Tensor topk_indices, - const Tensor &gating_output, - const Tensor &correction_bias = Tensor(), - bool renormalize = false); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/moe_topk_softmax.hpp b/include/infinicore/ops/moe_topk_softmax.hpp deleted file mode 100644 index e6964aa88..000000000 --- a/include/infinicore/ops/moe_topk_softmax.hpp +++ /dev/null @@ -1,35 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "../tensor.hpp" -#include "common/op.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(MoeTopkSoftmax, - Tensor, - Tensor, - const Tensor &, - const Tensor &, - const bool, - const float); - -std::tuple moe_topk_softmax( - const Tensor &gating_output, - size_t topk, - bool renormalize = false, - float moe_softcapping = 0.0f, - const Tensor &correction_bias = Tensor()); - -void moe_topk_softmax_( - Tensor topk_weights, - Tensor topk_indices, - const Tensor &gating_output, - const Tensor &correction_bias = Tensor(), - bool renormalize = false, - float moe_softcapping = 0.0f); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/mrope.hpp b/include/infinicore/ops/mrope.hpp deleted file mode 100644 index 3ee2c1ef1..000000000 --- a/include/infinicore/ops/mrope.hpp +++ /dev/null @@ -1,51 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "../tensor.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(MRoPE, - Tensor, - Tensor, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - int, - int, - int, - int, - int, - bool); - -void mrope_(Tensor q_out, - Tensor k_out, - const Tensor &q, - const Tensor &k, - const Tensor &cos, - const Tensor &sin, - const Tensor &positions, - int head_size, - int rotary_dim, - int section_t, - int section_h, - int section_w, - bool interleaved); - -std::pair mrope(const Tensor &q, - const Tensor &k, - const Tensor &cos, - const Tensor &sin, - const Tensor &positions, - int head_size, - int rotary_dim, - int section_t, - int section_h, - int section_w, - bool interleaved); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/mul.hpp b/include/infinicore/ops/mul.hpp deleted file mode 100644 index 2eb480ddb..000000000 --- a/include/infinicore/ops/mul.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Mul, Tensor, const Tensor &, const Tensor &); - -Tensor mul(const Tensor &a, const Tensor &b); -void mul_(Tensor c, const Tensor &a, const Tensor &b); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/mul_scalar.hpp b/include/infinicore/ops/mul_scalar.hpp deleted file mode 100644 index 5ee5e0d0b..000000000 --- a/include/infinicore/ops/mul_scalar.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(MulScalar, Tensor, const Tensor &, double); - -Tensor mul_scalar(const Tensor &a, double alpha); -void mul_scalar_(Tensor c, const Tensor &a, double alpha); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/multi_margin_loss.hpp b/include/infinicore/ops/multi_margin_loss.hpp deleted file mode 100644 index e36cf469b..000000000 --- a/include/infinicore/ops/multi_margin_loss.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class MultiMarginLoss { -public: - using schema = void (*)(Tensor, Tensor, Tensor, Tensor, int64_t, float, int64_t); - - static void execute(Tensor output, Tensor input, Tensor target, Tensor weight, int64_t p, float margin, int64_t reduction); - static common::OpDispatcher &dispatcher(); -}; - -Tensor multi_margin_loss(Tensor input, Tensor target, Tensor weight = {}, int64_t p = 1, float margin = 1.0f, int64_t reduction = 1); -void multi_margin_loss_(Tensor output, Tensor input, Tensor target, Tensor weight, int64_t p, float margin, int64_t reduction); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/nrm2.hpp b/include/infinicore/ops/nrm2.hpp deleted file mode 100644 index c5552a4c1..000000000 --- a/include/infinicore/ops/nrm2.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Nrm2, const Tensor &, Tensor); - -Tensor nrm2(const Tensor &x); -void nrm2_(const Tensor &x, Tensor result); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/nsa_compress_paged_cache.hpp b/include/infinicore/ops/nsa_compress_paged_cache.hpp deleted file mode 100644 index bb5380632..000000000 --- a/include/infinicore/ops/nsa_compress_paged_cache.hpp +++ /dev/null @@ -1,15 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(NsaCompressPagedCache, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, int, bool); - -void nsa_compress_paged_cache_(Tensor k_cmp, Tensor v_cmp, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, int nsa_block_size, - bool update_last_only = false); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/nsa_paged_attention.hpp b/include/infinicore/ops/nsa_paged_attention.hpp deleted file mode 100644 index 00e9adc10..000000000 --- a/include/infinicore/ops/nsa_paged_attention.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(NsaPagedAttention, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, float, int, int, int); - -Tensor nsa_paged_attention(const Tensor &q, const Tensor &k_cmp, const Tensor &v_cmp, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, const Tensor &gates, - float scale, int nsa_block_size, int window_size, int select_blocks); - -void nsa_paged_attention_(Tensor out, const Tensor &q, const Tensor &k_cmp, const Tensor &v_cmp, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, const Tensor &gates, - float scale, int nsa_block_size, int window_size, int select_blocks); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/ones.hpp b/include/infinicore/ops/ones.hpp deleted file mode 100644 index 0ea747484..000000000 --- a/include/infinicore/ops/ones.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "common/op.hpp" - -namespace infinicore::op { -class Ones { - -public: - using schema = void (*)(Tensor); - static void execute(Tensor output); - static common::OpDispatcher &dispatcher(); -}; - -Tensor ones(); -void ones_(Tensor output); -} // namespace infinicore::op diff --git a/include/infinicore/ops/pad.hpp b/include/infinicore/ops/pad.hpp deleted file mode 100644 index e598a94b4..000000000 --- a/include/infinicore/ops/pad.hpp +++ /dev/null @@ -1,25 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -#include -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Pad, Tensor, const Tensor &, const std::vector &, const std::string &, double); - -Tensor pad(const Tensor &x, - const std::vector &pad, - const std::string &mode = "constant", - double value = 0.0); - -void pad_(Tensor y, - const Tensor &x, - const std::vector &pad, - const std::string &mode = "constant", - double value = 0.0); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/paged_attention.hpp b/include/infinicore/ops/paged_attention.hpp deleted file mode 100644 index 8c906c95e..000000000 --- a/include/infinicore/ops/paged_attention.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(PagedAttention, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, const Tensor &, std::optional, float); - -Tensor paged_attention(const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, - std::optional alibi_slopes, float scale); - -void paged_attention_(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, - std::optional alibi_slopes, float scale); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/paged_attention_prefill.hpp b/include/infinicore/ops/paged_attention_prefill.hpp deleted file mode 100644 index 952924528..000000000 --- a/include/infinicore/ops/paged_attention_prefill.hpp +++ /dev/null @@ -1,52 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -class PagedAttentionPrefill { -public: - /** - * @brief PagedAttentionPrefill operator signature - * * Argument order: - * 1. out: Output tensor (Packed format) - * 2. q: Current Query tensor (Packed format) - * 3. k_cache: Physical Key cache (Paged format) - * 4. v_cache: Physical Value cache (Paged format) - * 5. block_tables: Mapping table from logical blocks to physical blocks - * 6. total_kv_lens: lengths of Complete Key/Value for each request - * 7. cu_seqlens_q: Cumulative sequence lengths of Query (prefix sum for variable-length batch) - * 8. alibi_slopes: ALiBi bias slopes (optional) - * 9. scale: Scaling factor (typically 1/sqrt(head_size)) - */ - using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, std::optional, float); - - static void execute(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, - Tensor block_tables, Tensor total_kv_lens, Tensor cum_seqlens_q, - std::optional alibi_slopes, float scale); - - static common::OpDispatcher &dispatcher(); -}; - -Tensor paged_attention_prefill(Tensor q, - Tensor k_cache, - Tensor v_cache, - Tensor block_tables, - Tensor total_kv_lens, - Tensor cum_seqlens_q, - std::optional alibi_slopes, - float scale); - -void paged_attention_prefill_(Tensor out, - Tensor q, - Tensor k_cache, - Tensor v_cache, - Tensor block_tables, - Tensor total_kv_lens, - Tensor cum_seqlens_q, - std::optional alibi_slopes, - float scale); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/paged_caching.hpp b/include/infinicore/ops/paged_caching.hpp deleted file mode 100644 index 403b4b738..000000000 --- a/include/infinicore/ops/paged_caching.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(PagedCaching, Tensor, Tensor, const Tensor &, const Tensor &, const Tensor &); - -void paged_caching_(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/per_channel_quant_i8.hpp b/include/infinicore/ops/per_channel_quant_i8.hpp deleted file mode 100644 index 0b0296248..000000000 --- a/include/infinicore/ops/per_channel_quant_i8.hpp +++ /dev/null @@ -1,12 +0,0 @@ -#pragma once -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(PerChannelQuantI8, const Tensor &, Tensor, Tensor); - -void per_channel_quant_i8_(const Tensor &x, Tensor x_packed, Tensor x_scale); -} // namespace infinicore::op diff --git a/include/infinicore/ops/per_tensor_dequant_i8.hpp b/include/infinicore/ops/per_tensor_dequant_i8.hpp deleted file mode 100644 index a6c4b1381..000000000 --- a/include/infinicore/ops/per_tensor_dequant_i8.hpp +++ /dev/null @@ -1,11 +0,0 @@ -#pragma once -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(PerTensorDequantI8, Tensor, const Tensor &, const Tensor &, const Tensor &); - -void per_tensor_dequant_i8_(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zero); -} // namespace infinicore::op diff --git a/include/infinicore/ops/per_tensor_quant_i8.hpp b/include/infinicore/ops/per_tensor_quant_i8.hpp deleted file mode 100644 index 84ec99c37..000000000 --- a/include/infinicore/ops/per_tensor_quant_i8.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(PerTensorQuantI8, const Tensor &, Tensor, Tensor, Tensor, bool); - -void per_tensor_quant_i8_(const Tensor &x, Tensor x_packed, Tensor x_scale, Tensor x_zero, bool is_static); - -Tensor per_tensor_quant_i8(const Tensor &x, Tensor x_scale, Tensor x_zero, bool is_static); -} // namespace infinicore::op diff --git a/include/infinicore/ops/prelu.hpp b/include/infinicore/ops/prelu.hpp deleted file mode 100644 index e22c73d70..000000000 --- a/include/infinicore/ops/prelu.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "infinicore.h" - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Prelu, Tensor, const Tensor &, const Tensor &); - -__export Tensor prelu(const Tensor &input, const Tensor &weight); -__export void prelu_(Tensor out, const Tensor &input, const Tensor &weight); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/prepare_moe_input.hpp b/include/infinicore/ops/prepare_moe_input.hpp deleted file mode 100644 index c06951320..000000000 --- a/include/infinicore/ops/prepare_moe_input.hpp +++ /dev/null @@ -1,48 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "../tensor.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(PrepareMoeInput, - Tensor, - Tensor, - Tensor, - Tensor, - Tensor, - Tensor, - const Tensor &, - const size_t, - const size_t, - const size_t); - -struct PrepareMoeInputOutput { - Tensor expert_offsets; - Tensor problem_sizes1; - Tensor problem_sizes2; - Tensor input_permutation; - Tensor output_permutation; -}; - -PrepareMoeInputOutput prepare_moe_input( - const Tensor &topk_ids, - size_t num_experts, - size_t n, - size_t k); - -void prepare_moe_input_( - Tensor expert_offsets, - Tensor blockscale_offsets, - Tensor problem_sizes1, - Tensor problem_sizes2, - Tensor input_permutation, - Tensor output_permutation, - const Tensor &topk_ids, - size_t num_experts, - size_t n, - size_t k); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/quickgelu.hpp b/include/infinicore/ops/quickgelu.hpp deleted file mode 100644 index b67d7f71a..000000000 --- a/include/infinicore/ops/quickgelu.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class QuickGelu { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor quick_gelu(Tensor input); -void quick_gelu_(Tensor output, Tensor input); -} // namespace infinicore::op diff --git a/include/infinicore/ops/random_sample.hpp b/include/infinicore/ops/random_sample.hpp deleted file mode 100644 index a3a9d29f1..000000000 --- a/include/infinicore/ops/random_sample.hpp +++ /dev/null @@ -1,22 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -#include "infinicore/tensor.hpp" - -namespace infinicore::op { - -class RandomSample { -public: - using schema = void (*)(Tensor, Tensor, float, float, int, float); - static void execute(Tensor indices, Tensor logits, float random_val, float topp, int topk, float temperature); - static common::OpDispatcher &dispatcher(); -}; - -// Out-of-place API -Tensor random_sample(Tensor logits, float random_val, float topp, int topk, float temperature); -// In-place API -void random_sample_(Tensor indices, Tensor logits, float random_val, float topp, int topk, float temperature); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/rearrange.hpp b/include/infinicore/ops/rearrange.hpp deleted file mode 100644 index 5db983ef8..000000000 --- a/include/infinicore/ops/rearrange.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Rearrange, Tensor, const Tensor &); - -Tensor rearrange(const Tensor &x); -void rearrange_(Tensor y, const Tensor &x); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/reciprocal.hpp b/include/infinicore/ops/reciprocal.hpp deleted file mode 100644 index 172fcdf36..000000000 --- a/include/infinicore/ops/reciprocal.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Reciprocal { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor y, Tensor x); - static common::OpDispatcher &dispatcher(); -}; - -Tensor reciprocal(Tensor x); -void reciprocal_(Tensor y, Tensor x); -} // namespace infinicore::op diff --git a/include/infinicore/ops/recurrent_gated_delta_rule.hpp b/include/infinicore/ops/recurrent_gated_delta_rule.hpp deleted file mode 100644 index 7d837c67b..000000000 --- a/include/infinicore/ops/recurrent_gated_delta_rule.hpp +++ /dev/null @@ -1,55 +0,0 @@ -#pragma once - -#include "infinicore.h" - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(RecurrentGatedDeltaRule, - Tensor, - Tensor, - std::optional, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - std::optional, - std::optional, - bool); - -__export Tensor recurrent_gated_delta_rule(const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - const Tensor &initial_state, - bool use_qk_l2norm = false); - -__export Tensor recurrent_gated_delta_rule_indexed(const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - Tensor initial_state, - const Tensor &initial_state_indices, - const Tensor &final_state_indices, - bool use_qk_l2norm = false); - -__export void recurrent_gated_delta_rule_(Tensor out, - Tensor initial_state, - std::optional final_state, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - std::optional initial_state_indices, - std::optional final_state_indices, - bool use_qk_l2norm = false); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/relu.hpp b/include/infinicore/ops/relu.hpp deleted file mode 100644 index 59f2b26f5..000000000 --- a/include/infinicore/ops/relu.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Relu { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor relu(Tensor input); -void relu_(Tensor output, Tensor input); -} // namespace infinicore::op diff --git a/include/infinicore/ops/relu6.hpp b/include/infinicore/ops/relu6.hpp deleted file mode 100644 index 7dd5635f8..000000000 --- a/include/infinicore/ops/relu6.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "infinicore.h" - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Relu6, Tensor, const Tensor &); - -__export Tensor relu6(const Tensor &input); -__export void relu6_(Tensor out, const Tensor &input); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/rms_norm.hpp b/include/infinicore/ops/rms_norm.hpp deleted file mode 100644 index c7b2b2d72..000000000 --- a/include/infinicore/ops/rms_norm.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(RMSNorm, Tensor, const Tensor &, const Tensor &, float); - -Tensor rms_norm(const Tensor &x, const Tensor &weight, float epsilon = 1e-5f); -void rms_norm_(Tensor y, const Tensor &x, const Tensor &weight, float epsilon = 1e-5f); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/rope.hpp b/include/infinicore/ops/rope.hpp deleted file mode 100644 index 8fd630ce1..000000000 --- a/include/infinicore/ops/rope.hpp +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "../nn/rope.hpp" -#include "../tensor.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(RoPE, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, infinicore::nn::RoPE::Algo); - -// Internal -void rope_(Tensor x_out, - const Tensor &x, - const Tensor &pos, - const Tensor &sin_table, - const Tensor &cos_table, - infinicore::nn::RoPE::Algo algo); - -// Public API -Tensor rope(const Tensor &x, - const Tensor &pos, - const Tensor &sin_table, - const Tensor &cos_table, - infinicore::nn::RoPE::Algo algo); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/rot.hpp b/include/infinicore/ops/rot.hpp deleted file mode 100644 index ff473f33e..000000000 --- a/include/infinicore/ops/rot.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Rot, Tensor, Tensor, const Tensor &, const Tensor &); - -void rot_(Tensor x, Tensor y, const Tensor &c, const Tensor &s); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/rotg.hpp b/include/infinicore/ops/rotg.hpp deleted file mode 100644 index c65f211dd..000000000 --- a/include/infinicore/ops/rotg.hpp +++ /dev/null @@ -1,17 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Rotg { -public: - using schema = void (*)(Tensor, Tensor, Tensor, Tensor); - static void execute(Tensor x, Tensor y, Tensor c, Tensor s); - static common::OpDispatcher &dispatcher(); -}; - -void rotg_(Tensor x, Tensor y, Tensor c, Tensor s); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/rotm.hpp b/include/infinicore/ops/rotm.hpp deleted file mode 100644 index 813fc27ee..000000000 --- a/include/infinicore/ops/rotm.hpp +++ /dev/null @@ -1,17 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Rotm { -public: - using schema = void (*)(Tensor, Tensor, Tensor); - static void execute(Tensor x, Tensor y, Tensor param); - static common::OpDispatcher &dispatcher(); -}; - -void rotm_(Tensor x, Tensor y, Tensor param); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/rotmg.hpp b/include/infinicore/ops/rotmg.hpp deleted file mode 100644 index e245840a4..000000000 --- a/include/infinicore/ops/rotmg.hpp +++ /dev/null @@ -1,17 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Rotmg { -public: - using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor); - static void execute(Tensor d1, Tensor d2, Tensor x1, Tensor y1, Tensor param); - static common::OpDispatcher &dispatcher(); -}; - -void rotmg_(Tensor d1, Tensor d2, Tensor x1, Tensor y1, Tensor param); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/rwkv5_wkv.hpp b/include/infinicore/ops/rwkv5_wkv.hpp deleted file mode 100644 index f1e6accd1..000000000 --- a/include/infinicore/ops/rwkv5_wkv.hpp +++ /dev/null @@ -1,34 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS( - Rwkv5Wkv, - Tensor, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - const Tensor &, - Tensor); - -Tensor rwkv5_wkv(const Tensor &receptance, - const Tensor &key, - const Tensor &value, - const Tensor &time_decay, - const Tensor &time_faaaa, - Tensor state); - -void rwkv5_wkv_(Tensor out, - const Tensor &receptance, - const Tensor &key, - const Tensor &value, - const Tensor &time_decay, - const Tensor &time_faaaa, - Tensor state); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/scal.hpp b/include/infinicore/ops/scal.hpp deleted file mode 100644 index d6cb5ce8c..000000000 --- a/include/infinicore/ops/scal.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Scal, const Tensor &, Tensor); - -void scal_(const Tensor &alpha, Tensor x); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/scaled_mm_i8.hpp b/include/infinicore/ops/scaled_mm_i8.hpp deleted file mode 100644 index 427ed9d65..000000000 --- a/include/infinicore/ops/scaled_mm_i8.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(I8Gemm, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, std::optional); - -void scaled_mm_i8_(Tensor c, const Tensor &a_p, const Tensor &a_s, const Tensor &b_p, const Tensor &b_s, std::optional bias); -} // namespace infinicore::op diff --git a/include/infinicore/ops/scaled_mm_w4a16_gptq_qy.hpp b/include/infinicore/ops/scaled_mm_w4a16_gptq_qy.hpp deleted file mode 100644 index 33ee7c37b..000000000 --- a/include/infinicore/ops/scaled_mm_w4a16_gptq_qy.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(GptqQyblasGemm, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, int64_t, int64_t); - -void scaled_mm_w4a16_gptq_qy_(Tensor out, const Tensor &in, const Tensor &qweight, const Tensor &scales, const Tensor &qzeros, int64_t quant_type, int64_t bit); -} // namespace infinicore::op diff --git a/include/infinicore/ops/scatter.hpp b/include/infinicore/ops/scatter.hpp deleted file mode 100644 index 306bb1232..000000000 --- a/include/infinicore/ops/scatter.hpp +++ /dev/null @@ -1,21 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Scatter { -public: - using schema = void (*)(Tensor, Tensor, int64_t, Tensor, Tensor, int64_t); - - static void execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction); - static common::OpDispatcher &dispatcher(); -}; - -Tensor scatter(Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction = 0); - -// In-place / 指定 Output 接口 -void scatter_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/selu.hpp b/include/infinicore/ops/selu.hpp deleted file mode 100644 index 2aed070ed..000000000 --- a/include/infinicore/ops/selu.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Selu { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor selu(Tensor input); -void selu_(Tensor output, Tensor input); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/sigmoid.hpp b/include/infinicore/ops/sigmoid.hpp deleted file mode 100644 index e5346608f..000000000 --- a/include/infinicore/ops/sigmoid.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "infinicore.h" - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Sigmoid, Tensor, const Tensor &); - -__export Tensor sigmoid(const Tensor &input); -__export void sigmoid_(Tensor output, const Tensor &input); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/silu.hpp b/include/infinicore/ops/silu.hpp deleted file mode 100644 index d4cbac57d..000000000 --- a/include/infinicore/ops/silu.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Silu { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor silu(Tensor input); -void silu_(Tensor output, Tensor input); -} // namespace infinicore::op diff --git a/include/infinicore/ops/silu_and_mul.hpp b/include/infinicore/ops/silu_and_mul.hpp deleted file mode 100644 index 1b594111d..000000000 --- a/include/infinicore/ops/silu_and_mul.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(SiluAndMul, Tensor, const Tensor &); - -Tensor silu_and_mul(const Tensor &x); -void silu_and_mul_(Tensor out, const Tensor &x); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/sinh.hpp b/include/infinicore/ops/sinh.hpp deleted file mode 100644 index 31cb3c256..000000000 --- a/include/infinicore/ops/sinh.hpp +++ /dev/null @@ -1,18 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Sinh { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor sinh(Tensor input); -void sinh_(Tensor output, Tensor input); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/smooth_l1_loss.hpp b/include/infinicore/ops/smooth_l1_loss.hpp deleted file mode 100644 index 9a67c6109..000000000 --- a/include/infinicore/ops/smooth_l1_loss.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class SmoothL1Loss { -public: - using schema = void (*)(Tensor, Tensor, Tensor, float, int64_t); - - static void execute(Tensor output, Tensor input, Tensor target, float beta, int64_t reduction); - static common::OpDispatcher &dispatcher(); -}; - -Tensor smooth_l1_loss(Tensor input, Tensor target, float beta = 1.0f, int64_t reduction = 1); -void smooth_l1_loss_(Tensor output, Tensor input, Tensor target, float beta, int64_t reduction); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/softmax.hpp b/include/infinicore/ops/softmax.hpp deleted file mode 100644 index c96ab1810..000000000 --- a/include/infinicore/ops/softmax.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Softmax { -public: - using schema = void (*)(Tensor, Tensor, int); - static void execute(Tensor output, Tensor input, int axis); - static common::OpDispatcher &dispatcher(); -}; - -Tensor softmax(Tensor input, int axis = -1); -void softmax_(Tensor output, Tensor input, int axis = -1); -} // namespace infinicore::op diff --git a/include/infinicore/ops/softplus.hpp b/include/infinicore/ops/softplus.hpp deleted file mode 100644 index b6d665455..000000000 --- a/include/infinicore/ops/softplus.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Softplus { -public: - // 修改 1: Schema 增加 float beta, float threshold - using schema = void (*)(Tensor, Tensor, float, float); - static void execute(Tensor y, Tensor x, float beta, float threshold); - static common::OpDispatcher &dispatcher(); -}; -// default: beta = 1.0, threshold = 20.0 -Tensor softplus(Tensor x, float beta = 1.0f, float threshold = 20.0f); - -void softplus_(Tensor y, Tensor x, float beta = 1.0f, float threshold = 20.0f); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/softsign.hpp b/include/infinicore/ops/softsign.hpp deleted file mode 100644 index 05e8e54fa..000000000 --- a/include/infinicore/ops/softsign.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Softsign { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor y, Tensor x); - static common::OpDispatcher &dispatcher(); -}; -// 返回新 Tensor 的函数接口 -Tensor softsign(Tensor x); -void softsign_(Tensor y, Tensor x); -} // namespace infinicore::op diff --git a/include/infinicore/ops/sum.hpp b/include/infinicore/ops/sum.hpp deleted file mode 100644 index 0ead8de26..000000000 --- a/include/infinicore/ops/sum.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include -#include - -namespace infinicore::op { -class Sum { -public: - using schema = void (*)(Tensor, Tensor, std::vector, bool); - static void execute(Tensor output, Tensor input, std::vector dim, bool keepdim = false); - static common::OpDispatcher &dispatcher(); -}; - -Tensor sum(Tensor input, std::vector dim, bool keepdim = false); -void sum_(Tensor output, Tensor input, std::vector dim, bool keepdim = false); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/swap.hpp b/include/infinicore/ops/swap.hpp deleted file mode 100644 index aba3ad563..000000000 --- a/include/infinicore/ops/swap.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Swap, Tensor, Tensor); - -void swap_(Tensor x, Tensor y); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/swiglu.hpp b/include/infinicore/ops/swiglu.hpp deleted file mode 100644 index 7aa77e632..000000000 --- a/include/infinicore/ops/swiglu.hpp +++ /dev/null @@ -1,15 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "../tensor.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(SwiGLU, Tensor, const Tensor &, const Tensor &); - -Tensor swiglu(const Tensor &a, const Tensor &b); -void swiglu_(Tensor c, const Tensor &a, const Tensor &b); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/take.hpp b/include/infinicore/ops/take.hpp deleted file mode 100644 index 5672005ad..000000000 --- a/include/infinicore/ops/take.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Take { -public: - using schema = void (*)(Tensor, Tensor, Tensor); - - static void execute(Tensor output, Tensor input, Tensor indices); - static common::OpDispatcher &dispatcher(); -}; - -Tensor take(Tensor input, Tensor indices); - -void take_(Tensor output, Tensor input, Tensor indices); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/tan.hpp b/include/infinicore/ops/tan.hpp deleted file mode 100644 index 72410f866..000000000 --- a/include/infinicore/ops/tan.hpp +++ /dev/null @@ -1,17 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Tan { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor tan(Tensor input); -void tan_(Tensor output, Tensor input); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/tanhshrink.hpp b/include/infinicore/ops/tanhshrink.hpp deleted file mode 100644 index 83673b22c..000000000 --- a/include/infinicore/ops/tanhshrink.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { -class Tanhshrink { -public: - using schema = void (*)(Tensor, Tensor); - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -Tensor tanhshrink(Tensor input); -void tanhshrink_(Tensor output, Tensor input); -} // namespace infinicore::op diff --git a/include/infinicore/ops/topk.hpp b/include/infinicore/ops/topk.hpp deleted file mode 100644 index d8486112c..000000000 --- a/include/infinicore/ops/topk.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -namespace infinicore::op { -class TopK { -public: - using schema = void (*)(Tensor, Tensor, Tensor, size_t, size_t, bool, bool); - static void execute(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest = true, bool sorted = true); - static common::OpDispatcher &dispatcher(); -}; - -std::pair topk(Tensor input, size_t k, size_t dim, bool largest = true, bool sorted = true); -void topk_(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest = true, bool sorted = true); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/topksoftmax.hpp b/include/infinicore/ops/topksoftmax.hpp deleted file mode 100644 index 2cd3ac7c5..000000000 --- a/include/infinicore/ops/topksoftmax.hpp +++ /dev/null @@ -1,14 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "../graph/graph.hpp" -#include "../tensor.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_CLASS(Topksoftmax, Tensor, Tensor, const Tensor &, const size_t, const int); - -void topksoftmax(Tensor values, Tensor indices, const Tensor &x, const size_t topk, const int norm = 0); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/triplet_margin_loss.hpp b/include/infinicore/ops/triplet_margin_loss.hpp deleted file mode 100644 index 43b935eb8..000000000 --- a/include/infinicore/ops/triplet_margin_loss.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class TripletMarginLoss { -public: - // Schema signature: output, anchor, positive, negative, margin, p, eps, swap, reduction - using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float, int64_t, float, bool, int64_t); - - static void execute(Tensor output, Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction); - static common::OpDispatcher &dispatcher(); -}; - -// Functional API -// reduction: 0=None, 1=Mean, 2=Sum -Tensor triplet_margin_loss(Tensor anchor, Tensor positive, Tensor negative, float margin = 1.0f, int64_t p = 2, float eps = 1e-6f, bool swap = false, int64_t reduction = 1); - -// In-place / Explicit Output API -void triplet_margin_loss_(Tensor output, Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/triplet_margin_with_distance_loss.hpp b/include/infinicore/ops/triplet_margin_with_distance_loss.hpp deleted file mode 100644 index d6ea449c5..000000000 --- a/include/infinicore/ops/triplet_margin_with_distance_loss.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class TripletMarginWithDistanceLoss { -public: - // Schema signature: output(out), anchor, positive, negative, margin, swap, reduction - using schema = void (*)(Tensor, Tensor, Tensor, Tensor, double, bool, int64_t); - - static void execute(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction); - static common::OpDispatcher &dispatcher(); -}; - -// Functional API: Returns the result tensor -// margin default 1.0, swap default false, reduction default 1 (Mean) typically -Tensor triplet_margin_with_distance_loss(Tensor anchor, Tensor positive, Tensor negative, double margin = 1.0, bool swap = false, int64_t reduction = 1); - -// In-place/Output-provided API -void triplet_margin_with_distance_loss_(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/unfold.hpp b/include/infinicore/ops/unfold.hpp deleted file mode 100644 index f8c7a3b13..000000000 --- a/include/infinicore/ops/unfold.hpp +++ /dev/null @@ -1,35 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -class Unfold { -public: - // schema: output, input, kernel_sizes, dilations, paddings, strides - using schema = void (*)(Tensor, Tensor, const std::vector &, const std::vector &, const std::vector &, const std::vector &); - - static void execute(Tensor output, Tensor input, - const std::vector &kernel_sizes, - const std::vector &dilations, - const std::vector &paddings, - const std::vector &strides); - static common::OpDispatcher &dispatcher(); -}; - -// Functional API -Tensor unfold(Tensor input, - std::vector kernel_sizes, - std::vector dilations, - std::vector paddings, - std::vector strides); - -void unfold_(Tensor output, Tensor input, - std::vector kernel_sizes, - std::vector dilations, - std::vector paddings, - std::vector strides); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/upsample_bilinear.hpp b/include/infinicore/ops/upsample_bilinear.hpp deleted file mode 100644 index ba9aa2379..000000000 --- a/include/infinicore/ops/upsample_bilinear.hpp +++ /dev/null @@ -1,22 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -class UpsampleBilinear { -public: - // Schema signature: output, input, align_corners - using schema = void (*)(Tensor, Tensor, bool); - - static void execute(Tensor output, Tensor input, bool align_corners); - static common::OpDispatcher &dispatcher(); -}; - -// 需要传入 output_size (如 {H_out, W_out} 或 {N, C, H_out, W_out}) 来决定新 Tensor 的形状 -Tensor upsample_bilinear(Tensor input, std::vector output_size, bool align_corners = false); -void upsample_bilinear_(Tensor output, Tensor input, bool align_corners); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/upsample_nearest.hpp b/include/infinicore/ops/upsample_nearest.hpp deleted file mode 100644 index 188d1b923..000000000 --- a/include/infinicore/ops/upsample_nearest.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include - -namespace infinicore::op { - -class UpsampleNearest { -public: - // Schema signature: output(out), input - // Note: Scales are inferred from output.shape / input.shape - using schema = void (*)(Tensor, Tensor); - - static void execute(Tensor output, Tensor input); - static common::OpDispatcher &dispatcher(); -}; - -// Functional API: Returns the result tensor -// Requires output_size to calculate the shape of the result tensor -Tensor upsample_nearest(Tensor input, const std::vector &output_size); - -// In-place/Output-provided API -void upsample_nearest_(Tensor output, Tensor input); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/vander.hpp b/include/infinicore/ops/vander.hpp deleted file mode 100644 index 8198cfeb8..000000000 --- a/include/infinicore/ops/vander.hpp +++ /dev/null @@ -1,21 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Vander { -public: - // schema: output, input, N, increasing - using schema = void (*)(Tensor, Tensor, int64_t, bool); - - static void execute(Tensor output, Tensor input, int64_t N, bool increasing); - static common::OpDispatcher &dispatcher(); -}; - -// N defaults to 0 (implying N = input.size(0), i.e., a square matrix) -Tensor vander(Tensor input, int64_t N = 0, bool increasing = false); -void vander_(Tensor output, Tensor input, int64_t N, bool increasing); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/var.hpp b/include/infinicore/ops/var.hpp deleted file mode 100644 index d1e01e1bf..000000000 --- a/include/infinicore/ops/var.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include -#include -#include -namespace infinicore::op { -class Var { -public: - using schema = void (*)(Tensor, Tensor, std::vector, bool, bool); // var_output, input, dim, unbiased, keepdim - static void execute(Tensor var_output, Tensor input, std::vector dim, bool unbiased = true, bool keepdim = false); - static common::OpDispatcher &dispatcher(); -}; - -Tensor var(Tensor input, std::vector dim, bool unbiased = true, bool keepdim = false); -void var_(Tensor var_output, Tensor input, std::vector dim, bool unbiased = true, bool keepdim = false); - -} // namespace infinicore::op diff --git a/include/infinicore/ops/var_mean.hpp b/include/infinicore/ops/var_mean.hpp deleted file mode 100644 index a9679187c..000000000 --- a/include/infinicore/ops/var_mean.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" -#include -#include -#include -namespace infinicore::op { -class Var_Mean { -public: - using schema = void (*)(Tensor, Tensor, Tensor, std::vector, bool, bool); // var_output, mean_output, input, dim, unbiased, keepdim - static void execute(Tensor var_output, Tensor mean_output, Tensor input, std::vector dim, bool unbiased = true, bool keepdim = false); - static common::OpDispatcher &dispatcher(); -}; - -std::pair var_mean(Tensor input, std::vector dim, bool unbiased = true, bool keepdim = false); -void var_mean_(Tensor var_output, Tensor mean_output, Tensor input, std::vector dim, bool unbiased = true, bool keepdim = false); - -} // namespace infinicore::op diff --git a/include/infinicore/quantization.hpp b/include/infinicore/quantization.hpp deleted file mode 100644 index 58964c412..000000000 --- a/include/infinicore/quantization.hpp +++ /dev/null @@ -1,9 +0,0 @@ -#pragma once - -#include "quantization/awq.hpp" -#include "quantization/base_quantization.hpp" -#include "quantization/compressed_tensors.hpp" -#include "quantization/gptq.hpp" -#include "quantization/gptq_qy.hpp" -#include "quantization/none_quantizaiton.hpp" -#include "quantization/quantization_scheme.hpp" diff --git a/include/infinicore/quantization/awq.hpp b/include/infinicore/quantization/awq.hpp deleted file mode 100644 index d64b1315b..000000000 --- a/include/infinicore/quantization/awq.hpp +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once -#include "base_quantization.hpp" -namespace infinicore::quantization { - -class AWQ : public BaseQuantization { - // This is a temporary class that currently only returns AWQ_W4A16. - // Future enhancements should parse quant_config to extract detailed quantization - // information and support multiple quantization schemes. -public: - explicit AWQ(const nlohmann::json &quant_config) - : BaseQuantization(quant_config){}; - - infinicore::quantization::QuantScheme - get_quant_scheme() const override { - return infinicore::quantization::QuantScheme::AWQ_W4A16; - }; - - int get_packing_num() const { - // For AWQ, we pack 8 int4 weights into a single int32 value. - return 32 / this->get_or("bits", 4); // Default to 8 if not specified in config - } - - int get_group_size() const { - // For simplicity, we return a fixed group size here. In a more complete implementation, - // this could be extracted from quant_config_ to support different group sizes. - return this->get_or("group_size", 128); // Standard AWQ group size - } -}; - -} // namespace infinicore::quantization diff --git a/include/infinicore/quantization/base_quantization.hpp b/include/infinicore/quantization/base_quantization.hpp deleted file mode 100644 index 4ee4f7608..000000000 --- a/include/infinicore/quantization/base_quantization.hpp +++ /dev/null @@ -1,43 +0,0 @@ -#pragma once -#include "nlohmann/json.hpp" -#include "quantization_scheme.hpp" - -namespace infinicore::quantization { -class BaseQuantization { - // Base class for quantization schemes. Intended to be extended to support various quantization methods. -public: - explicit BaseQuantization(const nlohmann::json &quant_config) : quant_config_(quant_config){}; - virtual ~BaseQuantization() = default; - - const nlohmann::json &get_config() const { return quant_config_; } - - virtual infinicore::quantization::QuantScheme get_quant_scheme() const = 0; - template - T get(const std::string &key) const { - if (!quant_config_.contains(key)) { - throw std::out_of_range("Key '" + key + "' not found in config."); - } - try { - return quant_config_.at(key).get(); - } catch (const nlohmann::json::type_error &e) { - throw std::runtime_error("Type conversion failed for key '" + key + "': " + std::string(e.what())); - } - } - - template - T get_or(const std::string &key, const T &default_value) const { - if (!quant_config_.contains(key) || quant_config_.at(key).is_null()) { - return default_value; - } - try { - return quant_config_.at(key).get(); - } catch (const nlohmann::json::type_error &) { - // If type conversion fails, return default value - return default_value; - } - } - -protected: - nlohmann::json quant_config_; -}; -} // namespace infinicore::quantization diff --git a/include/infinicore/quantization/compressed_tensors.hpp b/include/infinicore/quantization/compressed_tensors.hpp deleted file mode 100644 index 0549c1f00..000000000 --- a/include/infinicore/quantization/compressed_tensors.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include "base_quantization.hpp" -namespace infinicore::quantization { - -class CompressedTensors : public BaseQuantization { - // This is a temporary class that currently only returns COMPRESSED_TENSOR_W8A8I8. - // Future enhancements should parse quant_config to extract detailed quantization - // information and support multiple quantization schemes. -public: - explicit CompressedTensors(const nlohmann::json &quant_config) - : BaseQuantization(quant_config){}; - - infinicore::quantization::QuantScheme - get_quant_scheme() const override { - return infinicore::quantization::QuantScheme::COMPRESSED_TENSOR_W8A8I8; - }; -}; - -} // namespace infinicore::quantization diff --git a/include/infinicore/quantization/gptq.hpp b/include/infinicore/quantization/gptq.hpp deleted file mode 100644 index c03416b07..000000000 --- a/include/infinicore/quantization/gptq.hpp +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once -#include "base_quantization.hpp" -namespace infinicore::quantization { - -class GPTQ : public BaseQuantization { - // This is a temporary class that currently only returns GPTQ W4A16. - // Future enhancements should parse quant_config to extract detailed quantization - // information and support multiple quantization schemes. -public: - explicit GPTQ(const nlohmann::json &quant_config) - : BaseQuantization(quant_config){}; - - infinicore::quantization::QuantScheme - get_quant_scheme() const override { - return infinicore::quantization::QuantScheme::GPTQ_W4A16; - }; - - int get_packing_num() const { - // For GPTQ, we pack 8 int4 weights into a single int32 value. - return 32 / this->get_or("bits", 4); // Default to 8 if not specified in config - } - - int get_group_size() const { - // For simplicity, we return a fixed group size here. In a more complete implementation, - // this could be extracted from quant_config_ to support different group sizes. - return this->get_or("group_size", 128); // Standard GPTQ group size - } -}; - -} // namespace infinicore::quantization diff --git a/include/infinicore/quantization/gptq_qy.hpp b/include/infinicore/quantization/gptq_qy.hpp deleted file mode 100644 index f3ca47c59..000000000 --- a/include/infinicore/quantization/gptq_qy.hpp +++ /dev/null @@ -1,320 +0,0 @@ -#pragma once - -#include "../tensor.hpp" -#include "base_quantization.hpp" -#include -#include -#include -#include -#include -#include -#include - -namespace { -#ifndef INFINICORE_FLOAT16_DEFINED -#define INFINICORE_FLOAT16_DEFINED -struct float16_raw { - uint16_t data; - float16_raw() : data(0) {} - explicit float16_raw(float f) : data(fp32_to_fp16_bits(f)) {} - - static uint16_t fp32_to_fp16_bits(float value) { - union { - float f; - uint32_t u; - } f2u; - f2u.f = value; - uint32_t x = f2u.u; - - uint32_t sign = (x >> 16) & 0x8000; - int32_t exp = ((x >> 23) & 0xFF) - 127; - uint32_t mantissa = x & 0x007FFFFF; - - if (exp == 128) { - if (mantissa == 0) { - return static_cast(sign | 0x7C00); - } else { - return static_cast(sign | 0x7C00 | (mantissa >> 13)); - } - } - if (exp > 15) { - return static_cast(sign | 0x7C00); - } - if (exp < -14) { - if (exp < -24) { - return static_cast(sign); - } - mantissa |= 0x00800000; - uint32_t shift = -exp - 14; - mantissa >>= shift; - if ((mantissa & 0x1000) && ((mantissa & 0x2FFF) != 0)) { - mantissa += 0x2000; - } - return static_cast(sign | (mantissa >> 13)); - } - - uint32_t exp16 = static_cast(exp + 15) << 10; - uint32_t mantissa16 = mantissa >> 13; - if ((mantissa & 0x1000) && ((mantissa & 0x2FFF) || (mantissa16 & 1))) { - mantissa16++; - if (mantissa16 == 0x400) { - exp16 += 0x400; - mantissa16 = 0; - } - } - return static_cast(sign | exp16 | mantissa16); - } -}; -using float16_t = float16_raw; -#endif - -inline std::vector float_to_fp16_bits(const std::vector &values) { - std::vector result; - result.reserve(values.size()); - for (float f : values) { -#ifdef INFINICORE_HAS_FLOAT16 - infinicore::float16_t h(f); - result.push_back(*reinterpret_cast(&h)); -#else - result.push_back(float16_raw::fp32_to_fp16_bits(f)); -#endif - } - return result; -} -} // anonymous namespace - -namespace infinicore::quantization { - -class GPTQ_QY : public BaseQuantization { -public: - explicit GPTQ_QY(const nlohmann::json &quant_config) - : BaseQuantization(quant_config) { - int bits = weight_bits(); - if (bits != 4) { - spdlog::warn("GPTQ_QY: bits={} not fully tested, expected 4", bits); - } - } - - QuantScheme get_quant_scheme() const override { - return QuantScheme::GPTQ_W4A16_QY; - } - - int get_packing_num() const { - return 32 / weight_bits(); - } - - int get_group_size() const { - return get_or("group_size", 128); - } - - void convert_from_gptq_w4a16(const Tensor &original_qweight, - const Tensor &original_qzeros, - const Tensor &original_scales, - const Tensor &g_idx, - const Device &target_device) { - if (converted_) { - spdlog::debug("GPTQ_QY: weights already converted, skipping"); - return; - } - - const int bits = weight_bits(); - const int values_per_int32 = 32 / bits; - - { - const auto &shape = original_qweight->shape(); - assert(shape.size() == 2); - size_t M = shape[0], N = shape[1]; - - auto weight_unpacked = unpack_int32_to_nibbles_3d_(original_qweight, bits); - auto weight_packed = combine_nibbles_last_dim_(weight_unpacked, M, values_per_int32, N); - - size_t dimY = N; - size_t total_bytes = M * values_per_int32 * (N / 2); - size_t dimX = total_bytes / dimY; - - assert(dimX * dimY == total_bytes && "Weight shape calculation mismatch"); - - converted_weight_ = make_tensor_from_host_( - weight_packed.data(), - total_bytes * sizeof(uint8_t), - {dimX, dimY}, - DataType::U8, - target_device); - } - - { - const auto &shape = original_qzeros->shape(); - assert(shape.size() == 2); - size_t P = shape[0], Q = shape[1]; - - auto zeros_fp32 = unpack_zeros_to_fp32_2d_(original_qzeros, bits); - auto zeros_fp16 = ::float_to_fp16_bits(zeros_fp32); - - converted_zeros_ = make_tensor_from_host_( - zeros_fp16.data(), - zeros_fp16.size() * sizeof(uint16_t), - {P, Q * static_cast(values_per_int32)}, - DataType::F16, - target_device); - } - - { - auto scales_cpu = original_scales->to(Device::Type::CPU); - size_t num_elements = scales_cpu->numel(); - const void *raw_data = scales_cpu->data(); - - std::vector scales_fp16(num_elements); - if (scales_cpu->dtype() == DataType::F16) { - std::memcpy(scales_fp16.data(), raw_data, num_elements * sizeof(uint16_t)); - } else if (scales_cpu->dtype() == DataType::F32) { - std::vector scales_fp32(num_elements); - std::memcpy(scales_fp32.data(), raw_data, num_elements * sizeof(float)); - scales_fp16 = ::float_to_fp16_bits(scales_fp32); - } else { - spdlog::error("Unsupported scales dtype, expected F16 or F32"); - assert(false && "Unsupported scales dtype"); - } - - converted_scales_ = make_tensor_from_host_( - scales_fp16.data(), - scales_fp16.size() * sizeof(uint16_t), - original_scales->shape(), - DataType::F16, - target_device); - } - - if (g_idx->numel() > 0) { - g_idx_ = g_idx->to(target_device); - } - - converted_ = true; - } - - void release_buffers() { - converted_weight_ = Tensor(); - converted_zeros_ = Tensor(); - converted_scales_ = Tensor(); - g_idx_ = Tensor(); - } - - void convert_and_take_ownership(Tensor &weight, Tensor &zeros, Tensor &scales, - const Tensor &g_idx, const Device &target_device) { - if (converted_) { - spdlog::warn("GPTQ_QY: Already converted, skipping"); - return; - } - - convert_from_gptq_w4a16(weight, zeros, scales, g_idx, target_device); - - weight = std::move(converted_weight_); - zeros = std::move(converted_zeros_); - scales = std::move(converted_scales_); - - converted_ = false; - spdlog::debug("GPTQ_QY: Ownership transferred, internal buffers cleared."); - } - - const Tensor &get_converted_weight() const { return std::move(converted_weight_); } - const Tensor &get_converted_zeros() const { return std::move(converted_zeros_); } - const Tensor &get_converted_scales() const { return std::move(converted_scales_); } - const Tensor &get_g_idx() const { return g_idx_; } - bool is_converted() const { return converted_; } - - int weight_bits() const { return get_or("bits", 4); } - bool desc_act() const { return get_or("desc_act", false); } - -private: - static inline std::vector unpack_int32_to_nibbles_3d_(const Tensor &packed, int bits) { - assert(bits == 4 || bits == 8); - const int values_per_int32 = 32 / bits; - - auto packed_cpu = packed->to(Device::Type::CPU); - const int32_t *packed_host = reinterpret_cast(packed_cpu->data()); - - const auto &shape = packed->shape(); - assert(shape.size() == 2); - size_t M = shape[0], N = shape[1]; - - std::vector unpacked(M * values_per_int32 * N); - - for (size_t i = 0; i < M; ++i) { - for (int k = 0; k < values_per_int32; ++k) { - for (size_t j = 0; j < N; ++j) { - int32_t val = packed_host[i * N + j]; - uint8_t extracted = static_cast((val >> (k * bits)) & ((1 << bits) - 1)); - size_t idx = i * (values_per_int32 * N) + k * N + j; - unpacked[idx] = extracted; - } - } - } - return unpacked; - } - - static inline std::vector combine_nibbles_last_dim_( - const std::vector &nibbles, size_t M, size_t K, size_t N) { - assert(N % 2 == 0 && "Last dimension must be even for nibble pairing"); - - std::vector combined(M * K * (N / 2)); - size_t out_idx = 0; - - for (size_t i = 0; i < M; ++i) { - for (size_t k = 0; k < K; ++k) { - size_t row_base = i * (K * N) + k * N; - for (size_t j = 0; j < N; j += 2) { - uint8_t low = nibbles[row_base + j] & 0x0F; - uint8_t high = nibbles[row_base + j + 1] & 0x0F; - combined[out_idx++] = static_cast((high << 4) | low); - } - } - } - return combined; - } - - static inline std::vector unpack_zeros_to_fp32_2d_(const Tensor &packed_zeros, int bits) { - assert(bits == 4 || bits == 8); - const int values_per_int32 = 32 / bits; - const int mask = (1 << bits) - 1; - - auto packed_cpu = packed_zeros->to(Device::Type::CPU); - const int32_t *packed_host = reinterpret_cast(packed_cpu->data()); - - const auto &shape = packed_zeros->shape(); - assert(shape.size() == 2); - size_t P = shape[0], Q = shape[1]; - - std::vector result(P * Q * values_per_int32); - size_t out_idx = 0; - - for (size_t p = 0; p < P; ++p) { - for (size_t q = 0; q < Q; ++q) { - int32_t val = packed_host[p * Q + q]; - for (int k = 0; k < values_per_int32; ++k) { - uint8_t extracted = static_cast((val >> (k * bits)) & mask); - int dequant_val = (static_cast(extracted) + 1) & mask; - result[out_idx++] = static_cast(dequant_val); - } - } - } - return result; - } - - static inline Tensor make_tensor_from_host_(const void *data, size_t bytes, - const std::vector &shape, - DataType dtype, const Device &device) { - auto tensor = Tensor::empty(shape, dtype, Device::Type::CPU); - std::memcpy(reinterpret_cast(tensor->data()), data, bytes); - - if (device != Device::Type::CPU) { - return tensor->to(device); - } - return tensor; - } - - Tensor converted_weight_; - Tensor converted_zeros_; - Tensor converted_scales_; - Tensor g_idx_; - bool converted_ = false; -}; - -} // namespace infinicore::quantization diff --git a/include/infinicore/quantization/none_quantizaiton.hpp b/include/infinicore/quantization/none_quantizaiton.hpp deleted file mode 100644 index 5009f0adc..000000000 --- a/include/infinicore/quantization/none_quantizaiton.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include "base_quantization.hpp" -namespace infinicore::quantization { - -class NoneQuantization : public BaseQuantization { - // This is a temporary class that currently only returns COMPRESSED_TENSOR_W8A8I8. - // Future enhancements should parse quant_config to extract detailed quantization - // information and support multiple quantization schemes. -public: - explicit NoneQuantization(const nlohmann::json &quant_config) - : BaseQuantization(quant_config){}; - - infinicore::quantization::QuantScheme - get_quant_scheme() const override { - return infinicore::quantization::QuantScheme::NONE; - }; -}; - -} // namespace infinicore::quantization diff --git a/include/infinicore/quantization/quantization_scheme.hpp b/include/infinicore/quantization/quantization_scheme.hpp deleted file mode 100644 index 3f2daff09..000000000 --- a/include/infinicore/quantization/quantization_scheme.hpp +++ /dev/null @@ -1,19 +0,0 @@ -// quant.hpp -#pragma once - -namespace infinicore::quantization { - -enum class QuantScheme { - NONE, - COMPRESSED_TENSOR_W8A8I8, - AWQ_W4A16, - GPTQ_W4A16_QY, - GPTQ_W4A16, -}; - -enum class KVQuantAlgo { - NONE, - INT8, -}; - -} // namespace infinicore::quantization diff --git a/include/infinicore/tensor.hpp b/include/infinicore/tensor.hpp deleted file mode 100644 index b1650417d..000000000 --- a/include/infinicore/tensor.hpp +++ /dev/null @@ -1,336 +0,0 @@ -#pragma once - -#include "device.hpp" -#include "dtype.hpp" -#include "memory.hpp" - -#include -#include -#include -#include -#include - -#include -namespace infinicore { - -using Size = std::size_t; -using Stride = std::ptrdiff_t; -using Shape = std::vector; -using Strides = std::vector; - -class TensorImpl; - -struct TensorMetaData { - Shape shape; - Strides strides; - DataType dtype; - infiniopTensorDescriptor_t desc; - - TensorMetaData(const Shape &shape, const Strides &strides, const DataType &dtype); - ~TensorMetaData(); -}; - -struct TensorData { - size_t offset; - std::shared_ptr memory; -}; - -struct TensorSliceParams { - size_t dim; - size_t start; - Size len; -}; - -class Tensor { -public: - static Tensor empty(const Shape &shape, - const DataType &dtype, - const Device &device, - bool pin_memory = false); - - static Tensor strided_empty(const Shape &shape, - const Strides &strides, - const DataType &dtype, - const Device &device, - bool pin_memory = false); - - static Tensor zeros(const Shape &shape, - const DataType &dtype, - const Device &device, - bool pin_memory = false); - - static Tensor ones(const Shape &shape, - const DataType &dtype, - const Device &device, - bool pin_memory = false); - - static Tensor from_blob(void *raw_ptr, - const Shape &shape, - const DataType &dtype, - const Device &device); - - static Tensor strided_from_blob(void *raw_ptr, - const Shape &shape, - const Strides &strides, - const DataType &dtype, - const Device &device); - - void reset() noexcept { - impl_.reset(); - } - - void reset(std::shared_ptr new_impl) noexcept { - impl_ = std::move(new_impl); - } - - bool empty() const noexcept { - return impl_ == nullptr; - } - - size_t use_count() const noexcept { - return impl_ ? impl_.use_count() : 0; - } - - Tensor() = default; - Tensor(const Tensor &) = default; - Tensor(Tensor &&) = default; - Tensor &operator=(const Tensor &) = default; - Tensor &operator=(Tensor &&) = default; - - TensorImpl *operator->(); - const TensorImpl *operator->() const; - - operator bool() const; - -protected: - Tensor(std::shared_ptr impl) : impl_(std::move(impl)) {} - std::shared_ptr impl_; - friend class TensorImpl; - friend std::ostream &operator<<(std::ostream &out, const Tensor &tensor); -}; - -class TensorImpl : public std::enable_shared_from_this { - -public: - TensorImpl(const Shape &shape, const DataType &dtype); - TensorImpl(const Shape &shape, const Strides &strides, const DataType &dtype); - - std::byte *data(); - const std::byte *data() const; - - const Shape &shape() const; - - const Strides &strides() const; - - bool is_contiguous() const; - - Size ndim() const; - - Size numel() const; - - Size size(size_t dim) const; - - size_t element_size() const; - - size_t nbytes() const; - - Stride stride(size_t dim) const; - - DataType dtype() const; - - Device device() const; - - infiniopTensorDescriptor_t desc() const; - - bool is_pinned() const; - - std::string info() const; - - void debug(const std::string &filename) const; - - void debug() const; - - /** - * Unsafe API that returns a new tensor with the same raw memory untracked by allocator - * This API is used for loosely tracking a piece of memory while allowing it to be reused, - * typically in a compute graph scenario. - */ - Tensor to_blob_() const; - - /** - * Unsafe API that returns a new tensor with the same memory and let allocator retracks the memory. - * Should only be used on the tensor returned by to_blob_(). - */ - Tensor resume_from_blob_() const; - - /// - /// Data Transfer APIs - /// - - /** - * Returns a new tensor with the same data on a different device. - * If the new device passed is same as the current device, the original tensor is returned. - * - * @param device The device of the new tensor - * - * @return A new tensor with the same data on the specified device - */ - Tensor to(Device device) const; - - /** - * Copy Data from another tensor to this tensor. - * - * @param src The source tensor to copy from - * - * @return A new tensor with the same data on the specified device - */ - void copy_from(Tensor src); - - /** - * Return a tensor with the same data in contiguous arrangement as current tensor. - * If this tensor is already contiguous, the original tensor is returned. - * - * @return A new tensor with the same data on the specified device - */ - Tensor contiguous() const; - - /// - /// View APIs - /// - - /** - * Returns a new tensor with a dimension of size one removed at the specified position. - * Throws runtime_error if the dimension to be removed is not of size 1. - * - * @param dim The dimension index to remove - * @return A new tensor with the removed dimension - * - * Example: - * // For a 3D tensor with shape [1, 3, 4], squeeze at dim 0 results in shape [3, 4] - * tensor->squeeze(0); - */ - Tensor squeeze(size_t dim) const; - - /** - * Returns a new tensor with a dimension of size one inserted at the specified position. - * The returned tensor shares the same underlying storage with the original tensor. - * - * @param dim The dimension index at which to insert the new dimension - * @return A new tensor with the added dimension - * - * Example: - * // For a 2D tensor with shape [3, 4], unsqueeze at dim 0 results in shape [1, 3, 4] - * // unsqueeze at dim 1 results in shape [3, 1, 4] - * // unsqueeze at dim 2 results in shape [3, 4, 1] - * tensor->unsqueeze(0); - */ - Tensor unsqueeze(size_t dim) const; - - /** - * Returns a new tensor that is a narrowed version of the current tensor. - * The returned tensor shares the same underlying storage with the original tensor. - * - * @param slices A vector of slice parameters specifying the dimension, start index, - * and length for each dimension to narrow - * @return A new tensor with narrowed dimensions - * - * Example: - * // Narrow dimension 0 from index 2 to 5 (length 3) - * // and dimension 1 from index 1 to 3 (length 2) - * tensor.narrow({{0, 2, 3}, {1, 1, 2}}); - */ - Tensor narrow(const std::vector &slices) const; - - /** - * Returns a new tensor with the dimensions permuted (reordered) according to the given order. - * The returned tensor shares the same underlying storage with the original tensor. - * - * @param order The desired ordering of dimensions - * @return A new tensor with permuted dimensions - * - * Example: - * // For a 3D tensor with shape [2, 3, 4], permute to [2, 0, 1] - * // This swaps the dimensions: dim0->dim2, dim1->dim0, dim2->dim1 - * tensor->permute({2, 0, 1}); - */ - Tensor permute(const Shape &order) const; - - /** - * Returns a new tensor with the same data but a different shape. - * The returned tensor shares the same underlying storage with the original tensor. - * The tensor is rearranged if the new shape is not compatible with the current shape. - * - * @param new_shape The desired new shape - * @return A new tensor with the specified shape - * - * Example: - * // Reshape a 2x3 tensor (6 elements) to a 3x2 tensor - * tensor->view({3, 2}); - */ - Tensor view(const Shape &new_shape) const; - - /** - * Insecurely returns a new tensor with the specified shape and strides. - * The returned tensor shares the same underlying storage with the original tensor. - * - * @param new_shape The desired new shape - * @param new_strides The desired new strides - * @return A new tensor with the specified shape and strides - * - * Example: - * // Create a non-contiguous view with custom strides - * tensor->as_strided({2, 3}, {6, 2}); // Stride of 6 for dim0, 2 for dim1 - */ - Tensor as_strided(const Shape &new_shape, const Strides &new_strides) const; - -protected: - static std::shared_ptr empty( - const Shape &shape, - const DataType &dtype, - const Device &device, - bool pin_memory = false); - - static std::shared_ptr strided_empty( - const Shape &shape, - const Strides &strides, - const DataType &dtype, - const Device &device, - bool pin_memory = false); - - static std::shared_ptr zeros( - const Shape &shape, - const DataType &dtype, - const Device &device, - bool pin_memory = false); - - static std::shared_ptr ones( - const Shape &shape, - const DataType &dtype, - const Device &device, - bool pin_memory = false); - - static std::shared_ptr from_blob( - void *raw_ptr, - const Shape &shape, - const DataType &dtype, - const Device &device); - - static std::shared_ptr strided_from_blob( - void *raw_ptr, - const Shape &shape, - const Strides &strides, - const DataType &dtype, - const Device &device); - - friend class Tensor; - -protected: - TensorMetaData meta_; - TensorData data_; - -private: - // Mark to indicate if the tensor is created from to_blob_() - bool to_blob_mark_ = false; -}; - -} // namespace infinicore diff --git a/include/infiniop.h b/include/infiniop.h deleted file mode 100644 index 95bf75a0d..000000000 --- a/include/infiniop.h +++ /dev/null @@ -1,159 +0,0 @@ -#ifndef __INFINIOP_API_H__ -#define __INFINIOP_API_H__ - -#include "infiniop/handle.h" -#include "infiniop/ops/acos.h" -#include "infiniop/ops/adaptive_avg_pool1d.h" -#include "infiniop/ops/adaptive_avg_pool3d.h" -#include "infiniop/ops/adaptive_max_pool1d.h" -#include "infiniop/ops/add.h" -#include "infiniop/ops/add_rms_norm.h" -#include "infiniop/ops/addbmm.h" -#include "infiniop/ops/addcmul.h" -#include "infiniop/ops/addr.h" -#include "infiniop/ops/affine_grid.h" -#include "infiniop/ops/all.h" -#include "infiniop/ops/asin.h" -#include "infiniop/ops/asinh.h" -#include "infiniop/ops/asum.h" -#include "infiniop/ops/atanh.h" -#include "infiniop/ops/attention.h" -#include "infiniop/ops/avg_pool1d.h" -#include "infiniop/ops/avg_pool3d.h" -#include "infiniop/ops/awq_marlin_gemm.h" -#include "infiniop/ops/axpy.h" -#include "infiniop/ops/binary_cross_entropy_with_logits.h" -#include "infiniop/ops/blas_amax.h" -#include "infiniop/ops/blas_amin.h" -#include "infiniop/ops/blas_copy.h" -#include "infiniop/ops/blas_dot.h" -#include "infiniop/ops/block_diag.h" -#include "infiniop/ops/broadcast_to.h" -#include "infiniop/ops/causal_conv1d.h" -#include "infiniop/ops/causal_softmax.h" -#include "infiniop/ops/cdist.h" -#include "infiniop/ops/chunk_gated_delta_rule.h" -#include "infiniop/ops/clip.h" -#include "infiniop/ops/conv.h" -#include "infiniop/ops/cross_entropy.h" -#include "infiniop/ops/deepseek_moe.h" -#include "infiniop/ops/dequant/per_tensor_dequant_int8.h" -#include "infiniop/ops/dequantize_awq.h" -#include "infiniop/ops/dequantize_gptq.h" -#include "infiniop/ops/diff.h" -#include "infiniop/ops/digamma.h" -#include "infiniop/ops/dist.h" -#include "infiniop/ops/dot.h" -#include "infiniop/ops/embedding.h" -#include "infiniop/ops/equal.h" -#include "infiniop/ops/erf.h" -#include "infiniop/ops/erfc.h" -#include "infiniop/ops/erfinv.h" -#include "infiniop/ops/flash_attention.h" -#include "infiniop/ops/flipud.h" -#include "infiniop/ops/float_power.h" -#include "infiniop/ops/floor.h" -#include "infiniop/ops/floor_divide.h" -#include "infiniop/ops/fmin.h" -#include "infiniop/ops/fmod.h" -#include "infiniop/ops/fused_gated_delta_net_gating.h" -#include "infiniop/ops/gelu.h" -#include "infiniop/ops/gelutanh.h" -#include "infiniop/ops/gemm.h" -#include "infiniop/ops/gptq_marlin_gemm.h" -#include "infiniop/ops/gptq_qyblas_gemm.h" -#include "infiniop/ops/hardswish.h" -#include "infiniop/ops/hardtanh.h" -#include "infiniop/ops/hinge_embedding_loss.h" -#include "infiniop/ops/histc.h" -#include "infiniop/ops/huber_loss.h" -#include "infiniop/ops/hypot.h" -#include "infiniop/ops/index_add.h" -#include "infiniop/ops/index_copy.h" -#include "infiniop/ops/inner.h" -#include "infiniop/ops/int8_gemm.h" -#include "infiniop/ops/kron.h" -#include "infiniop/ops/kthvalue.h" -#include "infiniop/ops/kv_caching.h" -#include "infiniop/ops/layer_norm.h" -#include "infiniop/ops/ldexp.h" -#include "infiniop/ops/lerp.h" -#include "infiniop/ops/log10.h" -#include "infiniop/ops/log1p.h" -#include "infiniop/ops/log_softmax.h" -#include "infiniop/ops/logaddexp.h" -#include "infiniop/ops/logaddexp2.h" -#include "infiniop/ops/logcumsumexp.h" -#include "infiniop/ops/logdet.h" -#include "infiniop/ops/lp_norm.h" -#include "infiniop/ops/mamba_selective_scan.h" -#include "infiniop/ops/masked_select.h" -#include "infiniop/ops/matrix_power.h" -#include "infiniop/ops/moe_align.h" -#include "infiniop/ops/moe_fused_dense.h" -#include "infiniop/ops/moe_fused_gate.h" -#include "infiniop/ops/moe_sum.h" -#include "infiniop/ops/moe_topk_sigmoid.h" -#include "infiniop/ops/moe_topk_softmax.h" -#include "infiniop/ops/mrope.h" -#include "infiniop/ops/mul.h" -#include "infiniop/ops/mul_scalar.h" -#include "infiniop/ops/multi_margin_loss.h" -#include "infiniop/ops/nrm2.h" -#include "infiniop/ops/nsa_compress_paged_cache.h" -#include "infiniop/ops/nsa_paged_attention.h" -#include "infiniop/ops/ones.h" -#include "infiniop/ops/pad.h" -#include "infiniop/ops/paged_attention.h" -#include "infiniop/ops/paged_attention_prefill.h" -#include "infiniop/ops/paged_caching.h" -#include "infiniop/ops/pixel_shuffle.h" -#include "infiniop/ops/prepare_moe_input.h" -#include "infiniop/ops/quant/per_channel_quant_int8.h" -#include "infiniop/ops/quant/per_tensor_quant_int8.h" -#include "infiniop/ops/quickgelu.h" -#include "infiniop/ops/random_sample.h" -#include "infiniop/ops/rearrange.h" -#include "infiniop/ops/reciprocal.h" -#include "infiniop/ops/recurrent_gated_delta_rule.h" -#include "infiniop/ops/relu.h" -#include "infiniop/ops/rms_norm.h" -#include "infiniop/ops/rope.h" -#include "infiniop/ops/rot.h" -#include "infiniop/ops/rotg.h" -#include "infiniop/ops/rotm.h" -#include "infiniop/ops/rotmg.h" -#include "infiniop/ops/rwkv5_wkv.h" -#include "infiniop/ops/scal.h" -#include "infiniop/ops/scatter.h" -#include "infiniop/ops/selu.h" -#include "infiniop/ops/sigmoid.h" -#include "infiniop/ops/silu.h" -#include "infiniop/ops/silu_and_mul.h" -#include "infiniop/ops/sinh.h" -#include "infiniop/ops/smooth_l1_loss.h" -#include "infiniop/ops/softmax.h" -#include "infiniop/ops/softplus.h" -#include "infiniop/ops/softsign.h" -#include "infiniop/ops/sub.h" -#include "infiniop/ops/sum.h" -#include "infiniop/ops/swap.h" -#include "infiniop/ops/swiglu.h" -#include "infiniop/ops/take.h" -#include "infiniop/ops/tan.h" -#include "infiniop/ops/tanh.h" -#include "infiniop/ops/tanhshrink.h" -#include "infiniop/ops/topk.h" -#include "infiniop/ops/topkrouter.h" -#include "infiniop/ops/topksoftmax.h" -#include "infiniop/ops/triplet_margin_loss.h" -#include "infiniop/ops/triplet_margin_with_distance_loss.h" -#include "infiniop/ops/unfold.h" -#include "infiniop/ops/upsample_bilinear.h" -#include "infiniop/ops/upsample_nearest.h" -#include "infiniop/ops/vander.h" -#include "infiniop/ops/var.h" -#include "infiniop/ops/var_mean.h" -#include "infiniop/ops/zeros.h" -#include "infiniop/tensor_descriptor.h" -#endif // __INFINIOP_API_H__ diff --git a/include/infiniop/handle.h b/include/infiniop/handle.h deleted file mode 100644 index 207ca2ce9..000000000 --- a/include/infiniop/handle.h +++ /dev/null @@ -1,16 +0,0 @@ -#ifndef __INFINIOP_HANDLE_API_H__ -#define __INFINIOP_HANDLE_API_H__ - -#include "../infinicore.h" - -struct InfiniopHandle; - -typedef struct InfiniopHandle *infiniopHandle_t; - -__INFINI_C __export infiniStatus_t infiniopSetRuntimeDevice(infiniDevice_t device, int device_id); - -__INFINI_C __export infiniStatus_t infiniopCreateHandle(infiniopHandle_t *handle_ptr); - -__INFINI_C __export infiniStatus_t infiniopDestroyHandle(infiniopHandle_t handle); - -#endif diff --git a/include/infiniop/operator_descriptor.h b/include/infiniop/operator_descriptor.h deleted file mode 100644 index e46935bc0..000000000 --- a/include/infiniop/operator_descriptor.h +++ /dev/null @@ -1,13 +0,0 @@ -#ifndef __INFINIOP_OPERATOR_DESCRIPTOR_API_H__ -#define __INFINIOP_OPERATOR_DESCRIPTOR_API_H__ - -#include "handle.h" -#include "tensor_descriptor.h" - -// Base descriptor for all operators -struct InfiniopDescriptor; - -__INFINI_C __export infiniStatus_t infiniopGetDescriptorDeviceType(const struct InfiniopDescriptor *desc_ptr, infiniDevice_t *device_type); -__INFINI_C __export infiniStatus_t infiniopGetDescriptorDeviceId(const struct InfiniopDescriptor *desc_ptr, int *device_id); - -#endif //__INFINIOP_OPERATOR_DESCRIPTOR_API_H__ diff --git a/include/infiniop/ops/acos.h b/include/infiniop/ops/acos.h deleted file mode 100644 index 011625d84..000000000 --- a/include/infiniop/ops/acos.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_ACOS_API_H__ -#define __INFINIOP_ACOS_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAcosDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAcosDescriptor(infiniopHandle_t handle, - infiniopAcosDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetAcosWorkspaceSize(infiniopAcosDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAcos(infiniopAcosDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAcosDescriptor(infiniopAcosDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/adaptive_avg_pool1d.h b/include/infiniop/ops/adaptive_avg_pool1d.h deleted file mode 100644 index 192dee03e..000000000 --- a/include/infiniop/ops/adaptive_avg_pool1d.h +++ /dev/null @@ -1,34 +0,0 @@ -#ifndef __INFINIOP_ADAPTIVE_AVG_POOL1D_API_H__ -#define __INFINIOP_ADAPTIVE_AVG_POOL1D_API_H__ - -#include "../operator_descriptor.h" - -// 定义算子描述符类型 -typedef struct InfiniopDescriptor *infiniopAdaptiveAvgPool1dDescriptor_t; - -// 1. 创建算子描述符 -__INFINI_C __export infiniStatus_t infiniopCreateAdaptiveAvgPool1dDescriptor( - infiniopHandle_t handle, - infiniopAdaptiveAvgPool1dDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc); - -// 2. 获取 Workspace 大小 -__INFINI_C __export infiniStatus_t infiniopGetAdaptiveAvgPool1dWorkspaceSize( - infiniopAdaptiveAvgPool1dDescriptor_t desc, - size_t *size); - -// 3. 执行计算 -__INFINI_C __export infiniStatus_t infiniopAdaptiveAvgPool1d( - infiniopAdaptiveAvgPool1dDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -// 4. 销毁描述符 -__INFINI_C __export infiniStatus_t infiniopDestroyAdaptiveAvgPool1dDescriptor( - infiniopAdaptiveAvgPool1dDescriptor_t desc); - -#endif // __INFINIOP_ADAPTIVE_AVG_POOL1D_API_H__ diff --git a/include/infiniop/ops/adaptive_avg_pool3d.h b/include/infiniop/ops/adaptive_avg_pool3d.h deleted file mode 100644 index bd23ebce8..000000000 --- a/include/infiniop/ops/adaptive_avg_pool3d.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef INFINIOP_ADAPTIVE_AVG_POOL3D_H_ -#define INFINIOP_ADAPTIVE_AVG_POOL3D_H_ - -#include "../operator_descriptor.h" -#include - -typedef struct InfiniopDescriptor *infiniopAdaptiveAvgPool3DDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAdaptiveAvgPool3DDescriptor( - infiniopHandle_t handle, - infiniopAdaptiveAvgPool3DDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - size_t *output_size); - -__INFINI_C __export infiniStatus_t infiniopGetAdaptiveAvgPool3DWorkspaceSize( - infiniopAdaptiveAvgPool3DDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAdaptiveAvgPool3D( - infiniopAdaptiveAvgPool3DDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAdaptiveAvgPool3DDescriptor(infiniopAdaptiveAvgPool3DDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/adaptive_max_pool1d.h b/include/infiniop/ops/adaptive_max_pool1d.h deleted file mode 100644 index 0b42844d1..000000000 --- a/include/infiniop/ops/adaptive_max_pool1d.h +++ /dev/null @@ -1,22 +0,0 @@ -#ifndef __INFINIOP_ADAPTIVE_MAX_POOL1D_H__ -#define __INFINIOP_ADAPTIVE_MAX_POOL1D_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAdaptiveMaxPool1dDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAdaptiveMaxPool1dDescriptor( - infiniopHandle_t handle, - infiniopAdaptiveMaxPool1dDescriptor_t *desc, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t output_size); - -__INFINI_C __export infiniStatus_t infiniopGetAdaptiveMaxPool1dWorkspaceSize(infiniopAdaptiveMaxPool1dDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAdaptiveMaxPool1d(infiniopAdaptiveMaxPool1dDescriptor_t desc, void *workspace, size_t workspace_size, - void *y, const void *x, void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAdaptiveMaxPool1dDescriptor(infiniopAdaptiveMaxPool1dDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/add.h b/include/infiniop/ops/add.h deleted file mode 100644 index 8a7592c1c..000000000 --- a/include/infiniop/ops/add.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_ADD_API_H__ -#define __INFINIOP_ADD_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAddDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAddDescriptor(infiniopHandle_t handle, - infiniopAddDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c, - infiniopTensorDescriptor_t a, - infiniopTensorDescriptor_t b); - -__INFINI_C __export infiniStatus_t infiniopGetAddWorkspaceSize(infiniopAddDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAdd(infiniopAddDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAddDescriptor(infiniopAddDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/add_rms_norm.h b/include/infiniop/ops/add_rms_norm.h deleted file mode 100644 index 2cde1db76..000000000 --- a/include/infiniop/ops/add_rms_norm.h +++ /dev/null @@ -1,32 +0,0 @@ -#ifndef __INFINIOP_ADD_RMS_NORM_API_H__ -#define __INFINIOP_ADD_RMS_NORM_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAddRMSNormDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAddRMSNormDescriptor( - infiniopHandle_t handle, - infiniopAddRMSNormDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t residual_out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t weight_desc, - float epsilon); - -__INFINI_C __export infiniStatus_t infiniopGetAddRMSNormWorkspaceSize(infiniopAddRMSNormDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAddRMSNorm(infiniopAddRMSNormDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - void *residual_out, - const void *a, - const void *b, - const void *weight, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAddRMSNormDescriptor(infiniopAddRMSNormDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/addbmm.h b/include/infiniop/ops/addbmm.h deleted file mode 100644 index 04f918f2d..000000000 --- a/include/infiniop/ops/addbmm.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef __INFINIOP_ADDBMM_API_H__ -#define __INFINIOP_ADDBMM_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAddbmmDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAddbmmDescriptor(infiniopHandle_t handle, - infiniopAddbmmDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t batch1_desc, - infiniopTensorDescriptor_t batch2_desc, - float alpha, - float beta); - -__INFINI_C __export infiniStatus_t infiniopGetAddbmmWorkspaceSize(infiniopAddbmmDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAddbmm(infiniopAddbmmDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *batch1, - const void *batch2, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAddbmmDescriptor(infiniopAddbmmDescriptor_t desc); - -#endif // __INFINIOP_ADDBMM_API_H__ diff --git a/include/infiniop/ops/addcmul.h b/include/infiniop/ops/addcmul.h deleted file mode 100644 index b858dccf0..000000000 --- a/include/infiniop/ops/addcmul.h +++ /dev/null @@ -1,57 +0,0 @@ -#ifndef __INFINIOP_ADDCMUL_API_H__ -#define __INFINIOP_ADDCMUL_API_H__ - -#include "../operator_descriptor.h" - -// 定义 addcmul 算子描述符类型 -typedef struct InfiniopDescriptor *infiniopAddcmulDescriptor_t; - -/** - * @brief 创建 Addcmul 算子描述符 - * @param handle 算子句柄 - * @param desc_ptr 指向返回的描述符指针 - * @param out 输出张量描述符 - * @param input 加项张量描述符 - * @param tensor1 乘项张量1描述符 - * @param tensor2 乘项张量2描述符 - * @param value 乘积的标量系数 - */ -__INFINI_C __export infiniStatus_t infiniopCreateAddcmulDescriptor(infiniopHandle_t handle, - infiniopAddcmulDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out, - infiniopTensorDescriptor_t input, - infiniopTensorDescriptor_t tensor1, - infiniopTensorDescriptor_t tensor2, - float value); - -/** - * @brief 获取 Addcmul 计算所需的临时空间大小 - */ -__INFINI_C __export infiniStatus_t infiniopGetAddcmulWorkspaceSize(infiniopAddcmulDescriptor_t desc, size_t *size); - -/** - * @brief 执行 Addcmul 计算 - * @param desc 算子描述符 - * @param workspace 临时空间指针 - * @param workspace_size 临时空间大小 - * @param out 输出数据指针 - * @param input 加项数据指针 - * @param tensor1 乘项1数据指针 - * @param tensor2 乘项2数据指针 - * @param stream 计算流 (CUDA stream 等) - */ -__INFINI_C __export infiniStatus_t infiniopAddcmul(infiniopAddcmulDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *input, - const void *tensor1, - const void *tensor2, - void *stream); - -/** - * @brief 销毁 Addcmul 算子描述符 - */ -__INFINI_C __export infiniStatus_t infiniopDestroyAddcmulDescriptor(infiniopAddcmulDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/addr.h b/include/infiniop/ops/addr.h deleted file mode 100644 index a2f26d330..000000000 --- a/include/infiniop/ops/addr.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef __INFINIOP_ADDR_API_H__ -#define __INFINIOP_ADDR_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAddrDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAddrDescriptor(infiniopHandle_t handle, - infiniopAddrDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out, - infiniopTensorDescriptor_t input, - infiniopTensorDescriptor_t vec1, - infiniopTensorDescriptor_t vec2, - float beta, - float alpha); - -__INFINI_C __export infiniStatus_t infiniopGetAddrWorkspaceSize(infiniopAddrDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAddr(infiniopAddrDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *input, - const void *vec1, - const void *vec2, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAddrDescriptor(infiniopAddrDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/affine_grid.h b/include/infiniop/ops/affine_grid.h deleted file mode 100644 index 390d4b85f..000000000 --- a/include/infiniop/ops/affine_grid.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __INFINIOP_AFFINE_GRID_API_H__ -#define __INFINIOP_AFFINE_GRID_API_H__ -#include "../operator_descriptor.h" -#include - -typedef struct InfiniopDescriptor *infiniopAffineGridDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAffineGridDescriptor(infiniopHandle_t handle, - infiniopAffineGridDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - uint8_t align_corners); - -__INFINI_C __export infiniStatus_t infiniopGetAffineGridWorkspaceSize(infiniopAffineGridDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAffineGrid(infiniopAffineGridDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAffineGridDescriptor(infiniopAffineGridDescriptor_t desc); - -#endif // __INFINIOP_AFFINE_GRID_API_H__ diff --git a/include/infiniop/ops/all.h b/include/infiniop/ops/all.h deleted file mode 100644 index 41d74cf9a..000000000 --- a/include/infiniop/ops/all.h +++ /dev/null @@ -1,31 +0,0 @@ -#ifndef __INFINIOP_ALL_API_H__ -#define __INFINIOP_ALL_API_H__ - -#include "../operator_descriptor.h" -#include -#include -typedef struct InfiniopDescriptor *infiniopAllDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAllDescriptor(infiniopHandle_t handle, - infiniopAllDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - size_t *dim, - size_t dim_size, - bool keepdim); - -__INFINI_C __export infiniStatus_t infiniopGetAllWorkspaceSize(infiniopAllDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAll(infiniopAllDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - size_t *dim, - size_t dim_size, - bool keepdim, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAllDescriptor(infiniopAllDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/argwhere.h b/include/infiniop/ops/argwhere.h deleted file mode 100644 index b1d2f9f99..000000000 --- a/include/infiniop/ops/argwhere.h +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef __INFINIOP_ARGWHERE_API_H__ -#define __INFINIOP_ARGWHERE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopArgwhereDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateArgwhereDescriptor( - infiniopHandle_t handle, - infiniopArgwhereDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t input_desc); - -__INFINI_C __export infiniStatus_t infiniopGetArgwhereWorkspaceSize( - infiniopArgwhereDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopArgwhere( - infiniopArgwhereDescriptor_t desc, - void *workspace, - size_t workspace_size, - void **output, - size_t *count, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyArgwhereDescriptor( - infiniopArgwhereDescriptor_t desc); - -#endif // __INFINIOP_ARGWHERE_API_H__ diff --git a/include/infiniop/ops/asin.h b/include/infiniop/ops/asin.h deleted file mode 100644 index b4d16d481..000000000 --- a/include/infiniop/ops/asin.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_ASIN_API_H__ -#define __INFINIOP_ASIN_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAsinDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAsinDescriptor(infiniopHandle_t handle, - infiniopAsinDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input); - -__INFINI_C __export infiniStatus_t infiniopGetAsinWorkspaceSize(infiniopAsinDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAsin(infiniopAsinDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAsinDescriptor(infiniopAsinDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/asinh.h b/include/infiniop/ops/asinh.h deleted file mode 100644 index 95c86bf79..000000000 --- a/include/infiniop/ops/asinh.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_ASINH_API_H__ -#define __INFINIOP_ASINH_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAsinhDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAsinhDescriptor(infiniopHandle_t handle, - infiniopAsinhDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetAsinhWorkspaceSize(infiniopAsinhDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAsinh(infiniopAsinhDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAsinhDescriptor(infiniopAsinhDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/asum.h b/include/infiniop/ops/asum.h deleted file mode 100644 index 89336bce7..000000000 --- a/include/infiniop/ops/asum.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_ASUM_API_H__ -#define __INFINIOP_ASUM_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAsumDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAsumDescriptor(infiniopHandle_t handle, - infiniopAsumDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t result); - -__INFINI_C __export infiniStatus_t infiniopGetAsumWorkspaceSize(infiniopAsumDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAsum(infiniopAsumDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAsumDescriptor(infiniopAsumDescriptor_t desc); - -#endif // __INFINIOP_ASUM_API_H__ diff --git a/include/infiniop/ops/atanh.h b/include/infiniop/ops/atanh.h deleted file mode 100644 index 0db82b164..000000000 --- a/include/infiniop/ops/atanh.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_Atanh_API_H__ -#define __INFINIOP_Atanh_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAtanhDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAtanhDescriptor(infiniopHandle_t handle, - infiniopAtanhDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t a); - -__INFINI_C __export infiniStatus_t infiniopGetAtanhWorkspaceSize(infiniopAtanhDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAtanh(infiniopAtanhDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *a, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAtanhDescriptor(infiniopAtanhDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/attention.h b/include/infiniop/ops/attention.h deleted file mode 100644 index 21a911489..000000000 --- a/include/infiniop/ops/attention.h +++ /dev/null @@ -1,34 +0,0 @@ -#ifndef __INFINIOP_ATTENTION_API_H__ -#define __INFINIOP_ATTENTION_API_H__ - -#include "../operator_descriptor.h" -#include "gemm.h" -#include "swiglu.h" - -typedef struct InfiniopDescriptor *infiniopAttentionDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAttentionDescriptor(infiniopHandle_t handle, - infiniopAttentionDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t k_cache_desc, - infiniopTensorDescriptor_t v_cache_desc, - size_t pos); - -__INFINI_C __export infiniStatus_t infiniopGetAttentionWorkspaceSize(infiniopAttentionDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAttention(infiniopAttentionDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *q, - const void *k, - const void *v, - void *k_cache, - void *v_cache, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAttentionDescriptor(infiniopAttentionDescriptor_t desc); -#endif diff --git a/include/infiniop/ops/avg_pool1d.h b/include/infiniop/ops/avg_pool1d.h deleted file mode 100644 index 81c489dd7..000000000 --- a/include/infiniop/ops/avg_pool1d.h +++ /dev/null @@ -1,32 +0,0 @@ -#ifndef __INFINIOP_AVG_POOL1D_API_H__ -#define __INFINIOP_AVG_POOL1D_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAvgPool1dDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAvgPool1dDescriptor( - infiniopHandle_t handle, - infiniopAvgPool1dDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - size_t kernel_size, - size_t stride, - size_t padding); - -__INFINI_C __export infiniStatus_t infiniopGetAvgPool1dWorkspaceSize( - infiniopAvgPool1dDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAvgPool1d( - infiniopAvgPool1dDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAvgPool1dDescriptor( - infiniopAvgPool1dDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/avg_pool3d.h b/include/infiniop/ops/avg_pool3d.h deleted file mode 100644 index a54fbb263..000000000 --- a/include/infiniop/ops/avg_pool3d.h +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef __INFINIOP_AVG_POOL3D_API_H__ -#define __INFINIOP_AVG_POOL3D_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAvgPool3dDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAvgPool3dDescriptor(infiniopHandle_t handle, - infiniopAvgPool3dDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - void *kernel_size, - void *stride, - void *padding); - -__INFINI_C __export infiniStatus_t infiniopGetAvgPool3dWorkspaceSize(infiniopAvgPool3dDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAvgPool3d(infiniopAvgPool3dDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAvgPool3dDescriptor(infiniopAvgPool3dDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/awq_marlin_gemm.h b/include/infiniop/ops/awq_marlin_gemm.h deleted file mode 100644 index 2d3fefcb5..000000000 --- a/include/infiniop/ops/awq_marlin_gemm.h +++ /dev/null @@ -1,46 +0,0 @@ -#ifndef __INFINIOP_AWQ_MARLIN_GEMM_API_H__ -#define __INFINIOP_AWQ_MARLIN_GEMM_API_H__ - -#include "../operator_descriptor.h" -#include - -typedef struct InfiniopDescriptor *infiniopAwqMarlinGemmDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAwqMarlinGemmDescriptor(infiniopHandle_t handle, - infiniopAwqMarlinGemmDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t b_bias_desc, - infiniopTensorDescriptor_t b_scales_desc, - infiniopTensorDescriptor_t a_scales_desc, - infiniopTensorDescriptor_t global_scales_desc, - infiniopTensorDescriptor_t b_zeros_desc, - infiniopTensorDescriptor_t g_idx_desc, - infiniopTensorDescriptor_t perm_desc); - -__INFINI_C __export infiniStatus_t infiniopGetAwqMarlinGemmWorkspaceSize(infiniopAwqMarlinGemmDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAwqMarlinGemm(infiniopAwqMarlinGemmDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *b_bias, - void *b_scales, - void *a_scales, - void *global_scales, - void *b_zeros, - void *g_idx, - void *perm, - int64_t b_q_type_id, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAwqMarlinGemmDescriptor(infiniopAwqMarlinGemmDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/awq_marlin_repack.h b/include/infiniop/ops/awq_marlin_repack.h deleted file mode 100644 index 017ff5568..000000000 --- a/include/infiniop/ops/awq_marlin_repack.h +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef __INFINIOP_AWQ_MARLIN_REPACK_API_H__ -#define __INFINIOP_AWQ_MARLIN_REPACK_API_H__ - -#include "../operator_descriptor.h" -#include - -typedef struct InfiniopDescriptor *infiniopAwqMarlinRepackDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAwqMarlinRepackDescriptor(infiniopHandle_t handle, - infiniopAwqMarlinRepackDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - int64_t num_bits, - bool is_a_8bit); - -__INFINI_C __export infiniStatus_t infiniopGetAwqMarlinRepackWorkspaceSize(infiniopAwqMarlinRepackDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAwqMarlinRepack(infiniopAwqMarlinRepackDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAwqMarlinRepackDescriptor(infiniopAwqMarlinRepackDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/axpy.h b/include/infiniop/ops/axpy.h deleted file mode 100644 index ce6b2a23b..000000000 --- a/include/infiniop/ops/axpy.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_AXPY_API_H__ -#define __INFINIOP_AXPY_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopAxpyDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateAxpyDescriptor(infiniopHandle_t handle, - infiniopAxpyDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t alpha, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t y); - -__INFINI_C __export infiniStatus_t infiniopGetAxpyWorkspaceSize(infiniopAxpyDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopAxpy(infiniopAxpyDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *alpha, - const void *x, - void *y, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyAxpyDescriptor(infiniopAxpyDescriptor_t desc); - -#endif // __INFINIOP_AXPY_API_H__ diff --git a/include/infiniop/ops/binary_cross_entropy_with_logits.h b/include/infiniop/ops/binary_cross_entropy_with_logits.h deleted file mode 100644 index 0bec6f9d5..000000000 --- a/include/infiniop/ops/binary_cross_entropy_with_logits.h +++ /dev/null @@ -1,73 +0,0 @@ -#ifndef __INFINIOP_BINARY_CROSS_ENTROPY_WITH_LOGITS_API_H__ -#define __INFINIOP_BINARY_CROSS_ENTROPY_WITH_LOGITS_API_H__ - -#include "../operator_descriptor.h" - -// 定义归约方式枚举 -typedef enum { - INFINIOP_REDUCTION_NONE = 0, - INFINIOP_REDUCTION_MEAN = 1, - INFINIOP_REDUCTION_SUM = 2 -} infiniopReduction_t; - -// 定义 BCEWithLogits 算子描述符类型 -typedef struct InfiniopDescriptor *infiniopBCEWithLogitsDescriptor_t; - -/** - * @brief 创建 BCEWithLogits 算子描述符 - * @param handle 算子句柄 - * @param desc_ptr 指向返回的描述符指针 - * @param out 输出张量描述符 (none时与input同形状，mean/sum时为标量) - * @param logits 输入 Logits 张量描述符 - * @param target 目标标签张量描述符 - * @param weight 样本权重描述符 (可选，不需要则传 NULL) - * @param pos_weight 正样本权重描述符 (可选，不需要则传 NULL) - * @param reduction 归约方式 (none, mean, sum) - */ -__INFINI_C __export infiniStatus_t infiniopCreateBCEWithLogitsDescriptor( - infiniopHandle_t handle, - infiniopBCEWithLogitsDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out, - infiniopTensorDescriptor_t logits, - infiniopTensorDescriptor_t target, - infiniopTensorDescriptor_t weight, - infiniopTensorDescriptor_t pos_weight, - infiniopReduction_t reduction); - -/** - * @brief 获取 BCEWithLogits 计算所需的临时空间大小 - */ -__INFINI_C __export infiniStatus_t infiniopGetBCEWithLogitsWorkspaceSize( - infiniopBCEWithLogitsDescriptor_t desc, - size_t *size); - -/** - * @brief 执行 BCEWithLogits 计算 - * @param desc 算子描述符 - * @param workspace 临时空间指针 - * @param workspace_size 临时空间大小 - * @param out 输出数据指针 - * @param logits Logits 数据指针 - * @param target Target 数据指针 - * @param weight 权重数据指针 (可选，传 NULL 表示权重全为 1) - * @param pos_weight 正样本权重数据指针 (可选，传 NULL 表示权重全为 1) - * @param stream 计算流 - */ -__INFINI_C __export infiniStatus_t infiniopBCEWithLogits( - infiniopBCEWithLogitsDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *logits, - const void *target, - const void *weight, - const void *pos_weight, - void *stream); - -/** - * @brief 销毁 BCEWithLogits 算子描述符 - */ -__INFINI_C __export infiniStatus_t infiniopDestroyBCEWithLogitsDescriptor( - infiniopBCEWithLogitsDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/bitwise_right_shift.h b/include/infiniop/ops/bitwise_right_shift.h deleted file mode 100644 index 64256fd8f..000000000 --- a/include/infiniop/ops/bitwise_right_shift.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_BITWISE_RIGHT_SHIFT_API_H__ -#define __INFINIOP_BITWISE_RIGHT_SHIFT_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopBitwiseRightShiftDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateBitwiseRightShiftDescriptor(infiniopHandle_t handle, - infiniopBitwiseRightShiftDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x1, - infiniopTensorDescriptor_t x2); - -__INFINI_C __export infiniStatus_t infiniopGetBitwiseRightShiftWorkspaceSize(infiniopBitwiseRightShiftDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopBitwiseRightShift(infiniopBitwiseRightShiftDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyBitwiseRightShiftDescriptor(infiniopBitwiseRightShiftDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/blas_amax.h b/include/infiniop/ops/blas_amax.h deleted file mode 100644 index 9981c7fb7..000000000 --- a/include/infiniop/ops/blas_amax.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_BLAS_AMAX_API_H__ -#define __INFINIOP_BLAS_AMAX_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopBlasAmaxDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateBlasAmaxDescriptor(infiniopHandle_t handle, - infiniopBlasAmaxDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t result); - -__INFINI_C __export infiniStatus_t infiniopGetBlasAmaxWorkspaceSize(infiniopBlasAmaxDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopBlasAmax(infiniopBlasAmaxDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyBlasAmaxDescriptor(infiniopBlasAmaxDescriptor_t desc); - -#endif // __INFINIOP_BLAS_AMAX_API_H__ diff --git a/include/infiniop/ops/blas_amin.h b/include/infiniop/ops/blas_amin.h deleted file mode 100644 index 6bc8680ba..000000000 --- a/include/infiniop/ops/blas_amin.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_BLAS_AMIN_API_H__ -#define __INFINIOP_BLAS_AMIN_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopBlasAminDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateBlasAminDescriptor(infiniopHandle_t handle, - infiniopBlasAminDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t result); - -__INFINI_C __export infiniStatus_t infiniopGetBlasAminWorkspaceSize(infiniopBlasAminDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopBlasAmin(infiniopBlasAminDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyBlasAminDescriptor(infiniopBlasAminDescriptor_t desc); - -#endif // __INFINIOP_BLAS_AMIN_API_H__ diff --git a/include/infiniop/ops/blas_copy.h b/include/infiniop/ops/blas_copy.h deleted file mode 100644 index 7c6f3611c..000000000 --- a/include/infiniop/ops/blas_copy.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_BLAS_COPY_API_H__ -#define __INFINIOP_BLAS_COPY_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopBlasCopyDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateBlasCopyDescriptor(infiniopHandle_t handle, - infiniopBlasCopyDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t y); - -__INFINI_C __export infiniStatus_t infiniopGetBlasCopyWorkspaceSize(infiniopBlasCopyDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopBlasCopy(infiniopBlasCopyDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *x, - void *y, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyBlasCopyDescriptor(infiniopBlasCopyDescriptor_t desc); - -#endif // __INFINIOP_BLAS_COPY_API_H__ diff --git a/include/infiniop/ops/blas_dot.h b/include/infiniop/ops/blas_dot.h deleted file mode 100644 index 9d03af2c0..000000000 --- a/include/infiniop/ops/blas_dot.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_BLAS_DOT_API_H__ -#define __INFINIOP_BLAS_DOT_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopBlasDotDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateBlasDotDescriptor(infiniopHandle_t handle, - infiniopBlasDotDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t result); - -__INFINI_C __export infiniStatus_t infiniopGetBlasDotWorkspaceSize(infiniopBlasDotDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopBlasDot(infiniopBlasDotDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *x, - const void *y, - void *result, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyBlasDotDescriptor(infiniopBlasDotDescriptor_t desc); - -#endif // __INFINIOP_BLAS_DOT_API_H__ diff --git a/include/infiniop/ops/block_diag.h b/include/infiniop/ops/block_diag.h deleted file mode 100644 index f1829d82a..000000000 --- a/include/infiniop/ops/block_diag.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __INFINIOP_BLOCK_DIAG_API_H__ -#define __INFINIOP_BLOCK_DIAG_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopBlockDiagDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateBlockDiagDescriptor(infiniopHandle_t handle, - infiniopBlockDiagDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t *x, - size_t num_inputs); - -__INFINI_C __export infiniStatus_t infiniopGetBlockDiagWorkspaceSize(infiniopBlockDiagDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopBlockDiag(infiniopBlockDiagDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void **x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyBlockDiagDescriptor(infiniopBlockDiagDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/broadcast_to.h b/include/infiniop/ops/broadcast_to.h deleted file mode 100644 index 2830ececb..000000000 --- a/include/infiniop/ops/broadcast_to.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_BROADCAST_TO_API_H__ -#define __INFINIOP_BROADCAST_TO_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopBroadcastToDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateBroadcastToDescriptor(infiniopHandle_t handle, - infiniopBroadcastToDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetBroadcastToWorkspaceSize(infiniopBroadcastToDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopBroadcastTo(infiniopBroadcastToDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyBroadcastToDescriptor(infiniopBroadcastToDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/causal_conv1d.h b/include/infiniop/ops/causal_conv1d.h deleted file mode 100644 index 431ff1858..000000000 --- a/include/infiniop/ops/causal_conv1d.h +++ /dev/null @@ -1,43 +0,0 @@ -#ifndef __INFINIOP_CAUSAL_CONV1D_API_H__ -#define __INFINIOP_CAUSAL_CONV1D_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopCausalConv1dDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateCausalConv1dDescriptor( - infiniopHandle_t handle, - infiniopCausalConv1dDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, // padded: [B, T, C]; varlen: [1, total_tokens, C] - infiniopTensorDescriptor_t conv_state_desc, // no-index: [B/num_requests, C, state_len]; pool: [pool_size, C, state_len] - infiniopTensorDescriptor_t final_conv_state_desc, // nullable when final_state_indices_desc is provided - infiniopTensorDescriptor_t qkv_desc, // padded: [B, T, C]; varlen: [1, total_tokens, C] - infiniopTensorDescriptor_t weight_desc, // [C, 1, K], depthwise; current backend supports K == 4 - infiniopTensorDescriptor_t bias_desc, // nullable; [C] - infiniopTensorDescriptor_t cu_seqlens_desc, // nullable; [num_requests + 1], int32/int64 - infiniopTensorDescriptor_t initial_state_indices_desc, // nullable; [num_requests], int32/int64 - infiniopTensorDescriptor_t final_state_indices_desc); // nullable; [num_requests], int32/int64; writes final state in-place - -__INFINI_C __export infiniStatus_t infiniopGetCausalConv1dWorkspaceSize( - infiniopCausalConv1dDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopCausalConv1d( - infiniopCausalConv1dDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - void *conv_state, - void *final_conv_state, - const void *qkv, - const void *weight, - const void *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyCausalConv1dDescriptor( - infiniopCausalConv1dDescriptor_t desc); - -#endif // __INFINIOP_CAUSAL_CONV1D_API_H__ diff --git a/include/infiniop/ops/causal_softmax.h b/include/infiniop/ops/causal_softmax.h deleted file mode 100644 index 13b325ffd..000000000 --- a/include/infiniop/ops/causal_softmax.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_CAUSAL_SOFTMAX_API_H__ -#define __INFINIOP_CAUSAL_SOFTMAX_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopCausalSoftmaxDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateCausalSoftmaxDescriptor( - infiniopHandle_t handle, - infiniopCausalSoftmaxDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc); - -__INFINI_C __export infiniStatus_t infiniopGetCausalSoftmaxWorkspaceSize(infiniopCausalSoftmaxDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopCausalSoftmax( - infiniopCausalSoftmaxDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyCausalSoftmaxDescriptor(infiniopCausalSoftmaxDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/cdist.h b/include/infiniop/ops/cdist.h deleted file mode 100644 index ab8d10877..000000000 --- a/include/infiniop/ops/cdist.h +++ /dev/null @@ -1,56 +0,0 @@ -#ifndef __INFINIOP_CDIST_API_H__ -#define __INFINIOP_CDIST_API_H__ - -#include "../operator_descriptor.h" - -// 定义 cdist 算子描述符类型 -typedef struct InfiniopDescriptor *infiniopCdistDescriptor_t; - -/** - * @brief 创建 Cdist 算子描述符 - * @param handle 算子句柄 - * @param desc_ptr 指向返回的描述符指针 - * @param y 输出张量描述符 (Shape: M x N) - * @param x1 输入张量1描述符 (Shape: M x D) - * @param x2 输入张量2描述符 (Shape: N x D) - * @param p 范数阶数 (L-p norm) - */ -__INFINI_C __export infiniStatus_t infiniopCreateCdistDescriptor( - infiniopHandle_t handle, - infiniopCdistDescriptor_t *desc_ptr, // 注意这里应该是具体类型的指针 - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p); - -/** - * @brief 获取 Cdist 计算所需的临时空间大小 - */ -__INFINI_C __export infiniStatus_t infiniopGetCdistWorkspaceSize(infiniopCdistDescriptor_t desc, - size_t *size); - -/** - * @brief 执行 Cdist 计算 - * @param desc 算子描述符 - * @param workspace 临时空间指针 - * @param workspace_size 临时空间大小 - * @param y 输出数据指针 - * @param x1 输入1数据指针 - * @param x2 输入2数据指针 - * @param stream 计算流 (CUDA stream 等) - */ -__INFINI_C __export infiniStatus_t infiniopCdist( - infiniopCdistDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream); - -/** - * @brief 销毁 Cdist 算子描述符 - */ -__INFINI_C __export infiniStatus_t infiniopDestroyCdistDescriptor(infiniopCdistDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/chunk_gated_delta_rule.h b/include/infiniop/ops/chunk_gated_delta_rule.h deleted file mode 100644 index 6caca6217..000000000 --- a/include/infiniop/ops/chunk_gated_delta_rule.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef __INFINIOP_CHUNK_GATED_DELTA_RULE_API_H__ -#define __INFINIOP_CHUNK_GATED_DELTA_RULE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopChunkGatedDeltaRuleDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateChunkGatedDeltaRuleDescriptor( - infiniopHandle_t handle, - infiniopChunkGatedDeltaRuleDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, // padded: [B, T, Hv, Dv]; varlen: [1, total_tokens, Hv, Dv] - infiniopTensorDescriptor_t initial_state_desc, // legacy: [B, Hv, Dv, Dk]; indexed pool: [pool_size, Hv, Dv, Dk] - infiniopTensorDescriptor_t final_state_desc, // null when final_state_indices_desc is provided - infiniopTensorDescriptor_t q_desc, // padded: [B, T, Hk, Dk]; varlen: [1, total_tokens, Hk, Dk] - infiniopTensorDescriptor_t k_desc, // same shape as q - infiniopTensorDescriptor_t v_desc, // padded: [B, T, Hv, Dv]; varlen: [1, total_tokens, Hv, Dv] - infiniopTensorDescriptor_t g_desc, // padded: [B, T, Hv]; varlen: [1, total_tokens, Hv] - infiniopTensorDescriptor_t beta_desc, // same shape/dtype as g - infiniopTensorDescriptor_t cu_seqlens_desc, // nullable; [B + 1], int32/int64 - infiniopTensorDescriptor_t initial_state_indices_desc, // nullable; [B], int32/int64; enables indexed state-pool reads - infiniopTensorDescriptor_t final_state_indices_desc, // nullable; [B], int32/int64; writes final state in-place to initial_state pool - bool use_qk_l2norm, - size_t chunk_size); - -__INFINI_C __export infiniStatus_t infiniopGetChunkGatedDeltaRuleWorkspaceSize( - infiniopChunkGatedDeltaRuleDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopChunkGatedDeltaRule( - infiniopChunkGatedDeltaRuleDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - void *initial_state, - void *final_state, - const void *q, - const void *k, - const void *v, - const void *g, - const void *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyChunkGatedDeltaRuleDescriptor( - infiniopChunkGatedDeltaRuleDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/clip.h b/include/infiniop/ops/clip.h deleted file mode 100644 index 965932ad8..000000000 --- a/include/infiniop/ops/clip.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_CLIP_API_H__ -#define __INFINIOP_CLIP_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopClipDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateClipDescriptor(infiniopHandle_t handle, - infiniopClipDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t min_val, - infiniopTensorDescriptor_t max_val); - -__INFINI_C __export infiniStatus_t infiniopGetClipWorkspaceSize(infiniopClipDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopClip(infiniopClipDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *min_val, - const void *max_val, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyClipDescriptor(infiniopClipDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/conv.h b/include/infiniop/ops/conv.h deleted file mode 100644 index eea64f5e9..000000000 --- a/include/infiniop/ops/conv.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __INFINIOP_CONV_API_H__ -#define __INFINIOP_CONV_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopConvDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateConvDescriptor(infiniopHandle_t handle, - infiniopConvDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t w_desc, - infiniopTensorDescriptor_t b_desc, - void *pads, - void *strides, - void *dilations, - size_t n); - -__INFINI_C __export infiniStatus_t infiniopGetConvWorkspaceSize(infiniopConvDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopConv(infiniopConvDescriptor_t desc, void *workspace, size_t workspace_size, void *y, const void *x, const void *w, const void *bias, void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyConvDescriptor(infiniopConvDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/cross_entropy.h b/include/infiniop/ops/cross_entropy.h deleted file mode 100644 index 2ebd4b168..000000000 --- a/include/infiniop/ops/cross_entropy.h +++ /dev/null @@ -1,31 +0,0 @@ -#ifndef __INFINIOP_CROSS_ENTROPY_API_H__ -#define __INFINIOP_CROSS_ENTROPY_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopCrossEntropyDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateCrossEntropyDescriptor( - infiniopHandle_t handle, - infiniopCrossEntropyDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t target_desc); - -__INFINI_C __export infiniStatus_t infiniopGetCrossEntropyWorkspaceSize( - infiniopCrossEntropyDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopCrossEntropy( - infiniopCrossEntropyDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *target, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyCrossEntropyDescriptor( - infiniopCrossEntropyDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/deepseek_moe.h b/include/infiniop/ops/deepseek_moe.h deleted file mode 100644 index ae55ae977..000000000 --- a/include/infiniop/ops/deepseek_moe.h +++ /dev/null @@ -1,57 +0,0 @@ -#ifndef __INFINIOP_DEEPSEEK_MOE_API_H__ -#define __INFINIOP_DEEPSEEK_MOE_API_H__ - -#include "../operator_descriptor.h" - -#ifdef __cplusplus -#include -#else -#include -#endif - -typedef struct InfiniopDescriptor *infiniopDeepseekMoeDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateDeepseekMoeDescriptor( - infiniopHandle_t handle, - infiniopDeepseekMoeDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t hidden_desc, - infiniopTensorDescriptor_t topk_indices_desc, - infiniopTensorDescriptor_t topk_weights_desc, - size_t intermediate_size, - size_t num_experts); - -__INFINI_C __export infiniStatus_t infiniopGetDeepseekMoeWorkspaceSize( - infiniopDeepseekMoeDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopDeepseekMoe( - infiniopDeepseekMoeDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *const *gate_weights, - const void *const *up_weights, - const void *const *down_weights, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDeepseekMoeWithDevicePtrs( - infiniopDeepseekMoeDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *gate_weight_ptrs, - const void *up_weight_ptrs, - const void *down_weight_ptrs, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyDeepseekMoeDescriptor( - infiniopDeepseekMoeDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/dequant/per_tensor_dequant_int8.h b/include/infiniop/ops/dequant/per_tensor_dequant_int8.h deleted file mode 100644 index 9614b4303..000000000 --- a/include/infiniop/ops/dequant/per_tensor_dequant_int8.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_PER_TENSOR_DEQUANT_INT8_API_H__ -#define __INFINIOP_PER_TENSOR_DEQUANT_INT8_API_H__ - -#include "../../operator_descriptor.h" - -typedef InfiniopDescriptor *infiniopPerTensorDequantI8Descriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreatePerTensorDequantI8Descriptor(infiniopHandle_t handle, - infiniopPerTensorDequantI8Descriptor_t *desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t x_packed_desc, - infiniopTensorDescriptor_t x_scale_desc, - infiniopTensorDescriptor_t x_zero_desc); - -__INFINI_C __export infiniStatus_t infiniopGetPerTensorDequantI8WorkspaceSize(infiniopPerTensorDequantI8Descriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopPerTensorDequantI8(infiniopPerTensorDequantI8Descriptor_t desc, - void *workspace, - size_t workspace_size, - void *x, - const void *x_packed, - const void *x_scale, - const void *x_zero, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyPerTensorDequantI8Descriptor(infiniopPerTensorDequantI8Descriptor_t desc); - -#endif diff --git a/include/infiniop/ops/dequantize_awq.h b/include/infiniop/ops/dequantize_awq.h deleted file mode 100644 index f4cdc4bbe..000000000 --- a/include/infiniop/ops/dequantize_awq.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_DEQUANTIZE_AWQ_API_H__ -#define __INFINIOP_DEQUANTIZE_AWQ_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopDequantizeAWQDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateDequantizeAWQDescriptor(infiniopHandle_t handle, - infiniopDequantizeAWQDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t qweight_desc, - infiniopTensorDescriptor_t scales_desc, - infiniopTensorDescriptor_t zeros_desc); - -__INFINI_C __export infiniStatus_t infiniopGetDequantizeAWQWorkspaceSize(infiniopDequantizeAWQDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopDequantizeAWQ(infiniopDequantizeAWQDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *qweight, - const void *scales, - const void *zeros, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyDequantizeAWQDescriptor(infiniopDequantizeAWQDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/dequantize_gptq.h b/include/infiniop/ops/dequantize_gptq.h deleted file mode 100644 index 78ff15839..000000000 --- a/include/infiniop/ops/dequantize_gptq.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef __INFINIOP_DEQUANTIZE_GPTQ_API_H__ -#define __INFINIOP_DEQUANTIZE_GPTQ_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopDequantizeGPTQDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateDequantizeGPTQDescriptor(infiniopHandle_t handle, - infiniopDequantizeGPTQDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t qweight_desc, - infiniopTensorDescriptor_t scales_desc, - infiniopTensorDescriptor_t zeros_desc, - infiniopTensorDescriptor_t g_idx_desc); // add g_idx - -__INFINI_C __export infiniStatus_t infiniopGetDequantizeGPTQWorkspaceSize(infiniopDequantizeGPTQDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopDequantizeGPTQ(infiniopDequantizeGPTQDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *qweight, - const void *scales, - const void *zeros, - const void *g_idx, // add g_idx - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyDequantizeGPTQDescriptor(infiniopDequantizeGPTQDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/diff.h b/include/infiniop/ops/diff.h deleted file mode 100644 index f1d8dbcf9..000000000 --- a/include/infiniop/ops/diff.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_DIFF_API_H__ -#define __INFINIOP_DIFF_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopDiffDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateDiffDescriptor(infiniopHandle_t handle, - infiniopDiffDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - int dim, - int n); - -__INFINI_C __export infiniStatus_t infiniopGetDiffWorkspaceSize(infiniopDiffDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopDiff(infiniopDiffDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyDiffDescriptor(infiniopDiffDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/digamma.h b/include/infiniop/ops/digamma.h deleted file mode 100644 index 2c53dc9a1..000000000 --- a/include/infiniop/ops/digamma.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_DIGAMMA_API_H__ -#define __INFINIOP_DIGAMMA_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopDigammaDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateDigammaDescriptor(infiniopHandle_t handle, - infiniopDigammaDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetDigammaWorkspaceSize(infiniopDigammaDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopDigamma(infiniopDigammaDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyDigammaDescriptor(infiniopDigammaDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/dist.h b/include/infiniop/ops/dist.h deleted file mode 100644 index 45e044aa8..000000000 --- a/include/infiniop/ops/dist.h +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef __INFINIOP_DIST_API_H__ -#define __INFINIOP_DIST_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopDistDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateDistDescriptor(infiniopHandle_t handle, - infiniopDistDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x1, - infiniopTensorDescriptor_t x2, - double p); - -__INFINI_C __export infiniStatus_t infiniopGetDistWorkspaceSize(infiniopDistDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopDist(infiniopDistDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyDistDescriptor(infiniopDistDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/dot.h b/include/infiniop/ops/dot.h deleted file mode 100644 index 1fb9197f9..000000000 --- a/include/infiniop/ops/dot.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_DOT_API_H__ -#define __INFINIOP_DOT_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopDotDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateDotDescriptor(infiniopHandle_t handle, - infiniopDotDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t a, - infiniopTensorDescriptor_t b); - -__INFINI_C __export infiniStatus_t infiniopGetDotWorkspaceSize(infiniopDotDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopDot(infiniopDotDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyDotDescriptor(infiniopDotDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/embedding.h b/include/infiniop/ops/embedding.h deleted file mode 100644 index 5528be131..000000000 --- a/include/infiniop/ops/embedding.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __INFINIOP_EMBEDDING_API_H__ -#define __INFINIOP_EMBEDDING_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopEmbeddingDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateEmbeddingDescriptor( - infiniopHandle_t handle, - infiniopEmbeddingDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t weight_desc); - -__INFINI_C __export infiniStatus_t infiniopEmbedding( - infiniopEmbeddingDescriptor_t desc, - void *output, - const void *input, - const void *weight, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyEmbeddingDescriptor( - infiniopEmbeddingDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/equal.h b/include/infiniop/ops/equal.h deleted file mode 100644 index 90c4f3386..000000000 --- a/include/infiniop/ops/equal.h +++ /dev/null @@ -1,31 +0,0 @@ -#ifndef __INFINIOP_EQUAL_API_H__ -#define __INFINIOP_EQUAL_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopEqualDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateEqualDescriptor( - infiniopHandle_t handle, - infiniopEqualDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c, - infiniopTensorDescriptor_t a, - infiniopTensorDescriptor_t b); - -__INFINI_C __export infiniStatus_t infiniopGetEqualWorkspaceSize( - infiniopEqualDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopEqual( - infiniopEqualDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyEqualDescriptor( - infiniopEqualDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/erf.h b/include/infiniop/ops/erf.h deleted file mode 100644 index 8786b31a3..000000000 --- a/include/infiniop/ops/erf.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_ERF_API_H__ -#define __INFINIOP_ERF_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopErfDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateErfDescriptor(infiniopHandle_t handle, - infiniopErfDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetErfWorkspaceSize(infiniopErfDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopErf(infiniopErfDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyErfDescriptor(infiniopErfDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/erfc.h b/include/infiniop/ops/erfc.h deleted file mode 100644 index 6bceebe35..000000000 --- a/include/infiniop/ops/erfc.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_ERFC_API_H__ -#define __INFINIOP_ERFC_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopErfcDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateErfcDescriptor(infiniopHandle_t handle, - infiniopErfcDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetErfcWorkspaceSize(infiniopErfcDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopErfc(infiniopErfcDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyErfcDescriptor(infiniopErfcDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/erfinv.h b/include/infiniop/ops/erfinv.h deleted file mode 100644 index b14975253..000000000 --- a/include/infiniop/ops/erfinv.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_ERFINV_API_H__ -#define __INFINIOP_ERFINV_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopErfinvDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateErfinvDescriptor(infiniopHandle_t handle, - infiniopErfinvDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetErfinvWorkspaceSize(infiniopErfinvDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopErfinv(infiniopErfinvDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyErfinvDescriptor(infiniopErfinvDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/flash_attention.h b/include/infiniop/ops/flash_attention.h deleted file mode 100644 index 06f1e14e1..000000000 --- a/include/infiniop/ops/flash_attention.h +++ /dev/null @@ -1,36 +0,0 @@ -#ifndef __INFINIOP_FLASH_ATTENTION_API_H__ -#define __INFINIOP_FLASH_ATTENTION_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopFlashAttentionDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateFlashAttentionDescriptor( - infiniopHandle_t handle, - infiniopFlashAttentionDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t total_kv_len, - float scale, - char is_causal); - -__INFINI_C __export infiniStatus_t infiniopGetFlashAttentionWorkspaceSize( - infiniopFlashAttentionDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopFlashAttention( - infiniopFlashAttentionDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *q, - const void *k, - const void *v, - const void *total_kv_len, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyFlashAttentionDescriptor( - infiniopFlashAttentionDescriptor_t desc); -#endif diff --git a/include/infiniop/ops/flipud.h b/include/infiniop/ops/flipud.h deleted file mode 100644 index 6f5a4768d..000000000 --- a/include/infiniop/ops/flipud.h +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef __INFINIOP_FLIPUD_API_H__ -#define __INFINIOP_FLIPUD_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopFlipudDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateFlipudDescriptor(infiniopHandle_t handle, - infiniopFlipudDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input); - -// 获取工作空间大小 -__INFINI_C __export infiniStatus_t infiniopGetFlipudWorkspaceSize(infiniopFlipudDescriptor_t desc, size_t *size); - -// 执行 Flipud 算子 -__INFINI_C __export infiniStatus_t infiniopFlipud(infiniopFlipudDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -// 销毁描述符 -__INFINI_C __export infiniStatus_t infiniopDestroyFlipudDescriptor(infiniopFlipudDescriptor_t desc); - -#endif // __INFINIOP_FLIPUD_API_H__ diff --git a/include/infiniop/ops/float_power.h b/include/infiniop/ops/float_power.h deleted file mode 100644 index 8a4c9a97e..000000000 --- a/include/infiniop/ops/float_power.h +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef __INFINIOP_FLOAT_POWER_API_H__ -#define __INFINIOP_FLOAT_POWER_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopFloatPowerDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateFloatPowerDescriptor(infiniopHandle_t handle, - infiniopFloatPowerDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t exponent, - float scalar_exponent); - -__INFINI_C __export infiniStatus_t infiniopGetFloatPowerWorkspaceSize(infiniopFloatPowerDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopFloatPower(infiniopFloatPowerDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *exponent, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyFloatPowerDescriptor(infiniopFloatPowerDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/floor.h b/include/infiniop/ops/floor.h deleted file mode 100644 index 037237ed7..000000000 --- a/include/infiniop/ops/floor.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_FLOOR_API_H__ -#define __INFINIOP_FLOOR_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopFloorDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateFloorDescriptor(infiniopHandle_t handle, - infiniopFloorDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t intput); - -__INFINI_C __export infiniStatus_t infiniopGetFloorWorkspaceSize(infiniopFloorDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopFloor(infiniopFloorDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *intput, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyFloorDescriptor(infiniopFloorDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/floor_divide.h b/include/infiniop/ops/floor_divide.h deleted file mode 100644 index 300290ede..000000000 --- a/include/infiniop/ops/floor_divide.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_FLOOR_DIVIDE_API_H__ -#define __INFINIOP_FLOOR_DIVIDE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopFloorDivideDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateFloorDivideDescriptor(infiniopHandle_t handle, - infiniopFloorDivideDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c, - infiniopTensorDescriptor_t a, - infiniopTensorDescriptor_t b); - -__INFINI_C __export infiniStatus_t infiniopGetFloorDivideWorkspaceSize(infiniopFloorDivideDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopFloorDivide(infiniopFloorDivideDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyFloorDivideDescriptor(infiniopFloorDivideDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/fmin.h b/include/infiniop/ops/fmin.h deleted file mode 100644 index c1f9f75bf..000000000 --- a/include/infiniop/ops/fmin.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_FMIN_H__ -#define __INFINIOP_FMIN_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopFminDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateFminDescriptor(infiniopHandle_t handle, - infiniopFminDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c, - infiniopTensorDescriptor_t a, - infiniopTensorDescriptor_t b); - -__INFINI_C __export infiniStatus_t infiniopGetFminWorkspaceSize(infiniopFminDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopFmin(infiniopFminDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyFminDescriptor(infiniopFminDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/fmod.h b/include/infiniop/ops/fmod.h deleted file mode 100644 index e51f1b3c8..000000000 --- a/include/infiniop/ops/fmod.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_FMOD_API_H__ -#define __INFINIOP_FMOD_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopFmodDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateFmodDescriptor(infiniopHandle_t handle, - infiniopFmodDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c, - infiniopTensorDescriptor_t a, - infiniopTensorDescriptor_t b); - -__INFINI_C __export infiniStatus_t infiniopGetFmodWorkspaceSize(infiniopFmodDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopFmod(infiniopFmodDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyFmodDescriptor(infiniopFmodDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/fused_ffn.h b/include/infiniop/ops/fused_ffn.h deleted file mode 100644 index 6becd77d5..000000000 --- a/include/infiniop/ops/fused_ffn.h +++ /dev/null @@ -1,37 +0,0 @@ -#ifndef __INFINIOP_FUSED_FFN_API_H__ -#define __INFINIOP_FUSED_FFN_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopFusedFFNDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateFusedFFNDescriptor( - infiniopHandle_t handle, - infiniopFusedFFNDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - infiniopTensorDescriptor_t residual_desc, - infiniopTensorDescriptor_t norm_weight_desc, - infiniopTensorDescriptor_t gate_up_weight_desc, - infiniopTensorDescriptor_t down_weight_desc, - float epsilon); - -__INFINI_C __export infiniStatus_t infiniopGetFusedFFNWorkspaceSize( - infiniopFusedFFNDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopFusedFFN( - infiniopFusedFFNDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *in, - const void *residual, - const void *norm_weight, - const void *gate_up_weight, - const void *down_weight, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyFusedFFNDescriptor( - infiniopFusedFFNDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/fused_gated_delta_net_gating.h b/include/infiniop/ops/fused_gated_delta_net_gating.h deleted file mode 100644 index dafea6710..000000000 --- a/include/infiniop/ops/fused_gated_delta_net_gating.h +++ /dev/null @@ -1,43 +0,0 @@ -#ifndef __INFINIOP_FUSED_GATED_DELTA_NET_GATING_API_H__ -#define __INFINIOP_FUSED_GATED_DELTA_NET_GATING_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopFusedGatedDeltaNetGatingDescriptor_t; - -__INFINI_C __export infiniStatus_t -infiniopCreateFusedGatedDeltaNetGatingDescriptor( - infiniopHandle_t handle, - infiniopFusedGatedDeltaNetGatingDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t g_desc, - infiniopTensorDescriptor_t beta_output_desc, - infiniopTensorDescriptor_t A_log_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t dt_bias_desc, - float beta, - float threshold); - -__INFINI_C __export infiniStatus_t -infiniopGetFusedGatedDeltaNetGatingWorkspaceSize( - infiniopFusedGatedDeltaNetGatingDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t -infiniopFusedGatedDeltaNetGating( - infiniopFusedGatedDeltaNetGatingDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *g, - void *beta_output, - const void *A_log, - const void *a, - const void *b, - const void *dt_bias, - void *stream); - -__INFINI_C __export infiniStatus_t -infiniopDestroyFusedGatedDeltaNetGatingDescriptor( - infiniopFusedGatedDeltaNetGatingDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/gaussian_nll_loss.h b/include/infiniop/ops/gaussian_nll_loss.h deleted file mode 100644 index cc09edee8..000000000 --- a/include/infiniop/ops/gaussian_nll_loss.h +++ /dev/null @@ -1,31 +0,0 @@ -#ifndef __INFINIOP_GAUSSIAN_NLL_LOSS_API_H__ -#define __INFINIOP_GAUSSIAN_NLL_LOSS_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopGaussianNllLossDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateGaussianNllLossDescriptor(infiniopHandle_t handle, - infiniopGaussianNllLossDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t input, - infiniopTensorDescriptor_t target, - infiniopTensorDescriptor_t var, - int full, - double eps, - int reduction); - -__INFINI_C __export infiniStatus_t infiniopGetGaussianNllLossWorkspaceSize(infiniopGaussianNllLossDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopGaussianNllLoss(infiniopGaussianNllLossDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - const void *var, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyGaussianNllLossDescriptor(infiniopGaussianNllLossDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/gelu.h b/include/infiniop/ops/gelu.h deleted file mode 100644 index 3ddc21df8..000000000 --- a/include/infiniop/ops/gelu.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_GELU_API_H__ -#define __INFINIOP_GELU_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopGeluDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateGeluDescriptor(infiniopHandle_t handle, - infiniopGeluDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t intput); - -__INFINI_C __export infiniStatus_t infiniopGetGeluWorkspaceSize(infiniopGeluDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopGelu(infiniopGeluDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *intput, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyGeluDescriptor(infiniopGeluDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/gelutanh.h b/include/infiniop/ops/gelutanh.h deleted file mode 100644 index 9e94514b8..000000000 --- a/include/infiniop/ops/gelutanh.h +++ /dev/null @@ -1,43 +0,0 @@ -#ifndef __INFINIOP_GELUTANH_API_H__ -#define __INFINIOP_GELUTANH_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopGeluTanhDescriptor_t; - -/** - * Create GELU-Tanh descriptor - * - * y = x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) - */ -__INFINI_C __export infiniStatus_t infiniopCreateGeluTanhDescriptor( - infiniopHandle_t handle, - infiniopGeluTanhDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -/** - * Query workspace size - */ -__INFINI_C __export infiniStatus_t infiniopGetGeluTanhWorkspaceSize( - infiniopGeluTanhDescriptor_t desc, - size_t *size); - -/** - * Launch GELU-Tanh operator - */ -__INFINI_C __export infiniStatus_t infiniopGeluTanh( - infiniopGeluTanhDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -/** - * Destroy descriptor - */ -__INFINI_C __export infiniStatus_t infiniopDestroyGeluTanhDescriptor( - infiniopGeluTanhDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/gemm.h b/include/infiniop/ops/gemm.h deleted file mode 100644 index 430e37003..000000000 --- a/include/infiniop/ops/gemm.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_GEMM_API_H__ -#define __INFINIOP_GEMM_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopGemmDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateGemmDescriptor(infiniopHandle_t handle, - infiniopGemmDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc); - -__INFINI_C __export infiniStatus_t infiniopGetGemmWorkspaceSize(infiniopGemmDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopGemm(infiniopGemmDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - void const *a, - void const *b, - float alpha, - float beta, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyGemmDescriptor(infiniopGemmDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/gptq_marlin_gemm.h b/include/infiniop/ops/gptq_marlin_gemm.h deleted file mode 100644 index 66b2e1bb5..000000000 --- a/include/infiniop/ops/gptq_marlin_gemm.h +++ /dev/null @@ -1,42 +0,0 @@ -#ifndef __INFINIOP_GPTQ_MARLIN_GEMM_API_H__ -#define __INFINIOP_GPTQ_MARLIN_GEMM_API_H__ - -#include "../operator_descriptor.h" -#include - -typedef struct InfiniopDescriptor *infiniopGptqMarlinGemmDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateGptqMarlinGemmDescriptor(infiniopHandle_t handle, - infiniopGptqMarlinGemmDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t b_scales_desc, - infiniopTensorDescriptor_t global_scales_desc, - infiniopTensorDescriptor_t b_zeros_desc, - infiniopTensorDescriptor_t g_idx_desc, - infiniopTensorDescriptor_t perm_desc); - -__INFINI_C __export infiniStatus_t infiniopGetGptqMarlinGemmWorkspaceSize(infiniopGptqMarlinGemmDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopGptqMarlinGemm(infiniopGptqMarlinGemmDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *a, - const void *b, - void *b_scales, - void *global_scales, - void *b_zeros, - void *g_idx, - void *perm, - int64_t b_q_type_id, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyGptqMarlinGemmDescriptor(infiniopGptqMarlinGemmDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/gptq_marlin_repack.h b/include/infiniop/ops/gptq_marlin_repack.h deleted file mode 100644 index c3b588fa5..000000000 --- a/include/infiniop/ops/gptq_marlin_repack.h +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef __INFINIOP_GPTQ_MARLIN_REPACK_API_H__ -#define __INFINIOP_GPTQ_MARLIN_REPACK_API_H__ - -#include "../operator_descriptor.h" -#include - -typedef struct InfiniopDescriptor *infiniopGptqMarlinRepackDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateGptqMarlinRepackDescriptor(infiniopHandle_t handle, - infiniopGptqMarlinRepackDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t perm_desc, - int64_t num_bits, - bool is_a_8bit); - -__INFINI_C __export infiniStatus_t infiniopGetGptqMarlinRepackWorkspaceSize(infiniopGptqMarlinRepackDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopGptqMarlinRepack(infiniopGptqMarlinRepackDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *perm, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyGptqMarlinRepackDescriptor(infiniopGptqMarlinRepackDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/gptq_qyblas_gemm.h b/include/infiniop/ops/gptq_qyblas_gemm.h deleted file mode 100644 index bb105132c..000000000 --- a/include/infiniop/ops/gptq_qyblas_gemm.h +++ /dev/null @@ -1,37 +0,0 @@ -#ifndef __INFINIOP_GPTQ_QYBLAS_GEMM_API_H__ -#define __INFINIOP_GPTQ_QYBLAS_GEMM_API_H__ - -#include "../operator_descriptor.h" -#include - -typedef struct InfiniopDescriptor *infiniopGptqQyblasGemmDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateGptqQyblasGemmDescriptor( - infiniopHandle_t handle, - infiniopGptqQyblasGemmDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t b_scales_desc, - infiniopTensorDescriptor_t b_zeros_desc); - -__INFINI_C __export infiniStatus_t infiniopGetGptqQyblasGemmWorkspaceSize( - infiniopGptqQyblasGemmDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopGptqQyblasGemm( - infiniopGptqQyblasGemmDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *a, - const void *b, - void *b_scale, - void *b_zero, - int64_t quant_type, - int64_t bit, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyGptqQyblasGemmDescriptor( - infiniopGptqQyblasGemmDescriptor_t desc); -#endif diff --git a/include/infiniop/ops/hardswish.h b/include/infiniop/ops/hardswish.h deleted file mode 100644 index 1cdeecf67..000000000 --- a/include/infiniop/ops/hardswish.h +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef __INFINIOP_HARDSWISH_API_H__ -#define __INFINIOP_HARDSWISH_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopHardSwishDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateHardSwishDescriptor( - infiniopHandle_t handle, - infiniopHardSwishDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input); - -__INFINI_C __export infiniStatus_t infiniopGetHardSwishWorkspaceSize( - infiniopHardSwishDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopHardSwish( - infiniopHardSwishDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyHardSwishDescriptor( - infiniopHardSwishDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/hardtanh.h b/include/infiniop/ops/hardtanh.h deleted file mode 100644 index d2f98cedd..000000000 --- a/include/infiniop/ops/hardtanh.h +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef __INFINIOP_HARDTANH_API_H__ -#define __INFINIOP_HARDTANH_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopHardTanhDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateHardTanhDescriptor(infiniopHandle_t handle, - infiniopHardTanhDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - float min_val, - float max_val); - -__INFINI_C __export infiniStatus_t infiniopGetHardTanhWorkspaceSize(infiniopHardTanhDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopHardTanh(infiniopHardTanhDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyHardTanhDescriptor(infiniopHardTanhDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/hinge_embedding_loss.h b/include/infiniop/ops/hinge_embedding_loss.h deleted file mode 100644 index 6058f96cd..000000000 --- a/include/infiniop/ops/hinge_embedding_loss.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_HINGE_EMBEDDING_LOSS_API_H__ -#define __INFINIOP_HINGE_EMBEDDING_LOSS_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopHingeEmbeddingLossDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateHingeEmbeddingLossDescriptor(infiniopHandle_t handle, - infiniopHingeEmbeddingLossDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t input, - infiniopTensorDescriptor_t target, - double margin, - int reduction); - -__INFINI_C __export infiniStatus_t infiniopGetHingeEmbeddingLossWorkspaceSize(infiniopHingeEmbeddingLossDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopHingeEmbeddingLoss(infiniopHingeEmbeddingLossDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyHingeEmbeddingLossDescriptor(infiniopHingeEmbeddingLossDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/histc.h b/include/infiniop/ops/histc.h deleted file mode 100644 index 3e7232259..000000000 --- a/include/infiniop/ops/histc.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_HISTC_API_H__ -#define __INFINIOP_HISTC_API_H__ - -#include "../operator_descriptor.h" -#include - -typedef struct InfiniopDescriptor *infiniopHistcDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateHistcDescriptor(infiniopHandle_t handle, - infiniopHistcDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - int64_t bins, - double min_val, - double max_val); - -__INFINI_C __export infiniStatus_t infiniopGetHistcWorkspaceSize(infiniopHistcDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopHistc(infiniopHistcDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyHistcDescriptor(infiniopHistcDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/huber_loss.h b/include/infiniop/ops/huber_loss.h deleted file mode 100644 index 4bf0d7716..000000000 --- a/include/infiniop/ops/huber_loss.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_HUBER_LOSS_API_H__ -#define __INFINIOP_HUBER_LOSS_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopHuberLossDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateHuberLossDescriptor(infiniopHandle_t handle, - infiniopHuberLossDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - infiniopTensorDescriptor_t target, - float delta, - int reduction); - -__INFINI_C __export infiniStatus_t infiniopGetHuberLossWorkspaceSize(infiniopHuberLossDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopHuberLoss(infiniopHuberLossDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *target, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyHuberLossDescriptor(infiniopHuberLossDescriptor_t desc); - -#endif // __INFINIOP_HUBER_LOSS_API_H__ diff --git a/include/infiniop/ops/hypot.h b/include/infiniop/ops/hypot.h deleted file mode 100644 index 791c18463..000000000 --- a/include/infiniop/ops/hypot.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_HYPOT_API_H__ -#define __INFINIOP_HYPOT_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopHypotDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateHypotDescriptor(infiniopHandle_t handle, - infiniopHypotDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input_a, - infiniopTensorDescriptor_t input_b); - -__INFINI_C __export infiniStatus_t infiniopGetHypotWorkspaceSize(infiniopHypotDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopHypot(infiniopHypotDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input_a, - const void *input_b, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyHypotDescriptor(infiniopHypotDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/index_add.h b/include/infiniop/ops/index_add.h deleted file mode 100644 index f8b23e268..000000000 --- a/include/infiniop/ops/index_add.h +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef __INFINIOP_INDEX_ADD_API_H__ -#define __INFINIOP_INDEX_ADD_API_H__ -#include "../operator_descriptor.h" -#include - -typedef struct InfiniopDescriptor *infiniopIndexAddDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateIndexAddDescriptor(infiniopHandle_t handle, - infiniopIndexAddDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - int64_t dim, - infiniopTensorDescriptor_t index, - infiniopTensorDescriptor_t source, - float alpha); - -__INFINI_C __export infiniStatus_t infiniopGetIndexAddWorkspaceSize(infiniopIndexAddDescriptor_t desc, size_t *size); -__INFINI_C __export infiniStatus_t infiniopIndexAdd(infiniopIndexAddDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *index, - const void *source, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyIndexAddDescriptor(infiniopIndexAddDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/index_copy.h b/include/infiniop/ops/index_copy.h deleted file mode 100644 index e44032fcd..000000000 --- a/include/infiniop/ops/index_copy.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_INDEX_COPY_API_H__ -#define __INFINIOP_INDEX_COPY_API_H__ -#include "../operator_descriptor.h" -#include - -typedef struct InfiniopDescriptor *infiniopIndexCopyDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateIndexCopyDescriptor(infiniopHandle_t handle, - infiniopIndexCopyDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - int64_t dim, - infiniopTensorDescriptor_t index, - infiniopTensorDescriptor_t source); - -__INFINI_C __export infiniStatus_t infiniopGetIndexCopyWorkspaceSize(infiniopIndexCopyDescriptor_t desc, size_t *size); -__INFINI_C __export infiniStatus_t infiniopIndexCopy(infiniopIndexCopyDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *index, - const void *source, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyIndexCopyDescriptor(infiniopIndexCopyDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/inner.h b/include/infiniop/ops/inner.h deleted file mode 100644 index 99fc8ab13..000000000 --- a/include/infiniop/ops/inner.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_INNER_API_H__ -#define __INFINIOP_INNER_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopInnerDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateInnerDescriptor( - infiniopHandle_t handle, - infiniopInnerDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t other_desc); - -__INFINI_C __export infiniStatus_t infiniopGetInnerWorkspaceSize(infiniopInnerDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopInner( - infiniopInnerDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *input, - const void *other, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyInnerDescriptor(infiniopInnerDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/int8_gemm.h b/include/infiniop/ops/int8_gemm.h deleted file mode 100644 index bc96c50a3..000000000 --- a/include/infiniop/ops/int8_gemm.h +++ /dev/null @@ -1,32 +0,0 @@ -#ifndef __INFINIOP_I8GEMM_API_H__ -#define __INFINIOP_I8GEMM_API_H__ - -#include "../operator_descriptor.h" - -typedef InfiniopDescriptor *infiniopI8GemmDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateI8GemmDescriptor(infiniopHandle_t handle, - infiniopI8GemmDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t bias_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t x_scale_desc, - infiniopTensorDescriptor_t weights_desc, - infiniopTensorDescriptor_t weights_scale_desc); - -__INFINI_C __export infiniStatus_t infiniopGetI8GemmWorkspaceSize(infiniopI8GemmDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopI8Gemm(infiniopI8GemmDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *bias, - const void *x, - const void *x_scale, - const void *weights, - const void *weights_scale, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyI8GemmDescriptor(infiniopI8GemmDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/interpolate.h b/include/infiniop/ops/interpolate.h deleted file mode 100644 index d52615267..000000000 --- a/include/infiniop/ops/interpolate.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_INTERPOLATE_API_H__ -#define __INFINIOP_INTERPOLATE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopInterpolateDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateInterpolateDescriptor(infiniopHandle_t handle, - infiniopInterpolateDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - const char *mode, - void *size, - void *scale_factor, - int align_corners); - -__INFINI_C __export infiniStatus_t infiniopGetInterpolateWorkspaceSize(infiniopInterpolateDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopInterpolate(infiniopInterpolateDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyInterpolateDescriptor(infiniopInterpolateDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/kron.h b/include/infiniop/ops/kron.h deleted file mode 100644 index 44608b8e2..000000000 --- a/include/infiniop/ops/kron.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_KRON_API_H__ -#define __INFINIOP_KRON_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopKronDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateKronDescriptor(infiniopHandle_t handle, - infiniopKronDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x1, - infiniopTensorDescriptor_t x2); - -__INFINI_C __export infiniStatus_t infiniopGetKronWorkspaceSize(infiniopKronDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopKron(infiniopKronDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyKronDescriptor(infiniopKronDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/kthvalue.h b/include/infiniop/ops/kthvalue.h deleted file mode 100644 index 1f16df01d..000000000 --- a/include/infiniop/ops/kthvalue.h +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef __INFINIOP_KTHVALUE_API_H__ -#define __INFINIOP_KTHVALUE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopKthvalueDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateKthvalueDescriptor(infiniopHandle_t handle, - infiniopKthvalueDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t values, - infiniopTensorDescriptor_t indices, - infiniopTensorDescriptor_t input, - int k, - int dim, - int keepdim); - -__INFINI_C __export infiniStatus_t infiniopGetKthvalueWorkspaceSize(infiniopKthvalueDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopKthvalue(infiniopKthvalueDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *values, - void *indices, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyKthvalueDescriptor(infiniopKthvalueDescriptor_t desc); - -#endif // __INFINIOP_KTHVALUE_API_H__ diff --git a/include/infiniop/ops/kv_caching.h b/include/infiniop/ops/kv_caching.h deleted file mode 100644 index 1e091cf37..000000000 --- a/include/infiniop/ops/kv_caching.h +++ /dev/null @@ -1,31 +0,0 @@ -#ifndef __INFINIOP_KV_CACHING_API_H__ -#define __INFINIOP_KV_CACHING_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopKVCachingDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateKVCachingDescriptor( - infiniopHandle_t handle, - infiniopKVCachingDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t k_cache, - infiniopTensorDescriptor_t v_cache, - infiniopTensorDescriptor_t k, - infiniopTensorDescriptor_t v, - infiniopTensorDescriptor_t past_kv_lengths); - -__INFINI_C __export infiniStatus_t infiniopGetKVCachingWorkspaceSize(infiniopKVCachingDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopKVCaching(infiniopKVCachingDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *k_cache, - void *v_cache, - const void *k, - const void *v, - const void *past_kv_lengths, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyKVCachingDescriptor(infiniopKVCachingDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/layer_norm.h b/include/infiniop/ops/layer_norm.h deleted file mode 100644 index 140d18f02..000000000 --- a/include/infiniop/ops/layer_norm.h +++ /dev/null @@ -1,34 +0,0 @@ -#ifndef __INFINIOP_LAYER_NORM_API_H__ -#define __INFINIOP_LAYER_NORM_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopLayerNormDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateLayerNormDescriptor( - infiniopHandle_t handle, - infiniopLayerNormDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_standardization_desc, - infiniopTensorDescriptor_t input_std_deviation_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t bias_desc, - float eps); - -__INFINI_C __export infiniStatus_t infiniopGetLayerNormWorkspaceSize(infiniopLayerNormDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopLayerNorm(infiniopLayerNormDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - void *input_standardization, - void *input_std_deviation, - const void *input, - const void *weight, - const void *bias, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyLayerNormDescriptor(infiniopLayerNormDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/ldexp.h b/include/infiniop/ops/ldexp.h deleted file mode 100644 index f01891212..000000000 --- a/include/infiniop/ops/ldexp.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_LDEXP_API_H__ -#define __INFINIOP_LDEXP_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopLdexpDescriptor_t; -__INFINI_C __export infiniStatus_t infiniopCreateLdexpDescriptor(infiniopHandle_t handle, - infiniopLdexpDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t exp); -__INFINI_C __export infiniStatus_t infiniopGetLdexpWorkspaceSize(infiniopLdexpDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopLdexp(infiniopLdexpDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *exp, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyLdexpDescriptor(infiniopLdexpDescriptor_t desc); - -#endif // __INFINIOP_LDEXP_API_H__ diff --git a/include/infiniop/ops/lerp.h b/include/infiniop/ops/lerp.h deleted file mode 100644 index 4b392b51c..000000000 --- a/include/infiniop/ops/lerp.h +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef __INFINIOP_LERP_API_H__ -#define __INFINIOP_LERP_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopLerpDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateLerpDescriptor(infiniopHandle_t handle, - infiniopLerpDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t start, - infiniopTensorDescriptor_t end, - infiniopTensorDescriptor_t weight, - float weight_scalar); - -__INFINI_C __export infiniStatus_t infiniopGetLerpWorkspaceSize(infiniopLerpDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopLerp(infiniopLerpDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *start, - const void *end, - const void *weight, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyLerpDescriptor(infiniopLerpDescriptor_t desc); - -#endif // __INFINIOP_LERP_API_H__ diff --git a/include/infiniop/ops/log10.h b/include/infiniop/ops/log10.h deleted file mode 100644 index 5ca8ccd8d..000000000 --- a/include/infiniop/ops/log10.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_LOG10_API_H__ -#define __INFINIOP_LOG10_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopLog10Descriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateLog10Descriptor(infiniopHandle_t handle, - infiniopLog10Descriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetLog10WorkspaceSize(infiniopLog10Descriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopLog10(infiniopLog10Descriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyLog10Descriptor(infiniopLog10Descriptor_t desc); - -#endif diff --git a/include/infiniop/ops/log1p.h b/include/infiniop/ops/log1p.h deleted file mode 100644 index 93f229c7d..000000000 --- a/include/infiniop/ops/log1p.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_LOG1P_API_H__ -#define __INFINIOP_LOG1P_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopLog1pDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateLog1pDescriptor(infiniopHandle_t handle, - infiniopLog1pDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetLog1pWorkspaceSize(infiniopLog1pDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopLog1p(infiniopLog1pDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyLog1pDescriptor(infiniopLog1pDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/log_softmax.h b/include/infiniop/ops/log_softmax.h deleted file mode 100644 index 8e2bef0f3..000000000 --- a/include/infiniop/ops/log_softmax.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __INFINIOP_LOG_SOFTMAX_API_H__ -#define __INFINIOP_LOG_SOFTMAX_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopLogSoftmaxDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateLogSoftmaxDescriptor(infiniopHandle_t handle, - infiniopLogSoftmaxDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - int dim); - -__INFINI_C __export infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopLogSoftmax(infiniopLogSoftmaxDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyLogSoftmaxDescriptor(infiniopLogSoftmaxDescriptor_t desc); - -#endif // __INFINIOP_LOG_SOFTMAX_API_H__ diff --git a/include/infiniop/ops/logaddexp.h b/include/infiniop/ops/logaddexp.h deleted file mode 100644 index df18a130f..000000000 --- a/include/infiniop/ops/logaddexp.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_LOGADDEXP_API_H__ -#define __INFINIOP_LOGADDEXP_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopLogAddExpDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateLogAddExpDescriptor(infiniopHandle_t handle, - infiniopLogAddExpDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c, - infiniopTensorDescriptor_t a, - infiniopTensorDescriptor_t b); - -__INFINI_C __export infiniStatus_t infiniopGetLogAddExpWorkspaceSize(infiniopLogAddExpDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopLogAddExp(infiniopLogAddExpDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyLogAddExpDescriptor(infiniopLogAddExpDescriptor_t desc); - -#endif // __INFINIOP_LOGADDEXP_API_H__ diff --git a/include/infiniop/ops/logaddexp2.h b/include/infiniop/ops/logaddexp2.h deleted file mode 100644 index 8572c235e..000000000 --- a/include/infiniop/ops/logaddexp2.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_LOGADDEXP2_API_H__ -#define __INFINIOP_LOGADDEXP2_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopLogAddExp2Descriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateLogAddExp2Descriptor(infiniopHandle_t handle, - infiniopLogAddExp2Descriptor_t *desc_ptr, - infiniopTensorDescriptor_t c, - infiniopTensorDescriptor_t a, - infiniopTensorDescriptor_t b); - -__INFINI_C __export infiniStatus_t infiniopGetLogAddExp2WorkspaceSize(infiniopLogAddExp2Descriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopLogAddExp2(infiniopLogAddExp2Descriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyLogAddExp2Descriptor(infiniopLogAddExp2Descriptor_t desc); - -#endif // __INFINIOP_LOGADDEXP2_API_H__ diff --git a/include/infiniop/ops/logcumsumexp.h b/include/infiniop/ops/logcumsumexp.h deleted file mode 100644 index 7e00727e3..000000000 --- a/include/infiniop/ops/logcumsumexp.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef __INFINIOP_LOGCUMSUMEXP_API_H__ -#define __INFINIOP_LOGCUMSUMEXP_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopLogCumSumExpDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateLogCumSumExpDescriptor(infiniopHandle_t handle, - infiniopLogCumSumExpDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - int axis, - int exclusive, - int reverse); - -/* 获取执行 LogCumSumExp 所需的临时空间大小 */ -__INFINI_C __export infiniStatus_t infiniopGetLogCumSumExpWorkspaceSize(infiniopLogCumSumExpDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopLogCumSumExp(infiniopLogCumSumExpDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -/* 销毁描述符 */ -__INFINI_C __export infiniStatus_t infiniopDestroyLogCumSumExpDescriptor(infiniopLogCumSumExpDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/logdet.h b/include/infiniop/ops/logdet.h deleted file mode 100644 index f92fb58f3..000000000 --- a/include/infiniop/ops/logdet.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_LOGDET_API_H__ -#define __INFINIOP_LOGDET_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopLogdetDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateLogdetDescriptor(infiniopHandle_t handle, - infiniopLogdetDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetLogdetWorkspaceSize(infiniopLogdetDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopLogdet(infiniopLogdetDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyLogdetDescriptor(infiniopLogdetDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/lp_norm.h b/include/infiniop/ops/lp_norm.h deleted file mode 100644 index 1022fdb85..000000000 --- a/include/infiniop/ops/lp_norm.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_LP_NORM_API_H__ -#define __INFINIOP_LP_NORM_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopLPNormDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateLPNormDescriptor( - infiniopHandle_t handle, - infiniopLPNormDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - int axis, - int p, - float eps); - -__INFINI_C __export infiniStatus_t infiniopGetLPNormWorkspaceSize(infiniopLPNormDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopLPNorm(infiniopLPNormDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyLPNormDescriptor(infiniopLPNormDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/mamba_selective_scan.h b/include/infiniop/ops/mamba_selective_scan.h deleted file mode 100644 index c4eb5f574..000000000 --- a/include/infiniop/ops/mamba_selective_scan.h +++ /dev/null @@ -1,50 +0,0 @@ -#ifndef __INFINIOP_MAMBA_SELECTIVE_SCAN_API_H__ -#define __INFINIOP_MAMBA_SELECTIVE_SCAN_API_H__ - -#include "../operator_descriptor.h" -#ifdef __cplusplus -#include -#else -#include -#endif - -typedef struct InfiniopDescriptor *infiniopMambaSelectiveScanDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateMambaSelectiveScanDescriptor( - infiniopHandle_t handle, - infiniopMambaSelectiveScanDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t dt_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_log_desc, - infiniopTensorDescriptor_t d_desc, - infiniopTensorDescriptor_t gate_desc, - infiniopTensorDescriptor_t dt_bias_desc, - infiniopTensorDescriptor_t state_desc); - -__INFINI_C __export infiniStatus_t infiniopGetMambaSelectiveScanWorkspaceSize( - infiniopMambaSelectiveScanDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMambaSelectiveScan( - infiniopMambaSelectiveScanDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *x, - const void *dt, - const void *b, - const void *c, - const void *a_log, - const void *d, - const void *gate, - const void *dt_bias, - void *state, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMambaSelectiveScanDescriptor( - infiniopMambaSelectiveScanDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/masked_select.h b/include/infiniop/ops/masked_select.h deleted file mode 100644 index d24ed11f7..000000000 --- a/include/infiniop/ops/masked_select.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_MASKED_SELECT_API_H__ -#define __INFINIOP_MASKED_SELECT_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopMaskedSelectDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateMaskedSelectDescriptor( - infiniopHandle_t handle, - infiniopMaskedSelectDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t mask_desc); - -__INFINI_C __export infiniStatus_t infiniopGetMaskedSelectWorkspaceSize(infiniopMaskedSelectDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMaskedSelect( - infiniopMaskedSelectDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *input, - const bool *mask, - void **data_ptr, - size_t *dlen, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMaskedSelectDescriptor(infiniopMaskedSelectDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/matrix_power.h b/include/infiniop/ops/matrix_power.h deleted file mode 100644 index 639d3cf02..000000000 --- a/include/infiniop/ops/matrix_power.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __INFINIOP_MATRIX_POWER_API_H__ -#define __INFINIOP_MATRIX_POWER_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopMatrixPowerDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateMatrixPowerDescriptor(infiniopHandle_t handle, - infiniopMatrixPowerDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - int n); - -__INFINI_C __export infiniStatus_t infiniopGetMatrixPowerWorkspaceSize(infiniopMatrixPowerDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMatrixPower(infiniopMatrixPowerDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMatrixPowerDescriptor(infiniopMatrixPowerDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/moe_align.h b/include/infiniop/ops/moe_align.h deleted file mode 100644 index 582bc781f..000000000 --- a/include/infiniop/ops/moe_align.h +++ /dev/null @@ -1,37 +0,0 @@ -#ifndef __INFINIOP_MOE_ALIGN_API_H__ -#define __INFINIOP_MOE_ALIGN_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopMoeAlignDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateMoeAlignDescriptor( - infiniopHandle_t handle, - infiniopMoeAlignDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t sorted_token_ids_desc, - infiniopTensorDescriptor_t expert_ids_desc, - infiniopTensorDescriptor_t num_tokens_post_padded_desc, - infiniopTensorDescriptor_t topk_ids_desc, - size_t num_experts, - size_t block_size); - -__INFINI_C __export infiniStatus_t infiniopGetMoeAlignWorkspaceSize( - infiniopMoeAlignDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMoeAlign( - infiniopMoeAlignDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *sorted_token_ids, - void *expert_ids, - void *num_tokens_post_padded, - const void *topk_ids, - const void *expert_map, - int pad_sorted_token_ids, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMoeAlignDescriptor( - infiniopMoeAlignDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/moe_fused_dense.h b/include/infiniop/ops/moe_fused_dense.h deleted file mode 100644 index 61812ce01..000000000 --- a/include/infiniop/ops/moe_fused_dense.h +++ /dev/null @@ -1,43 +0,0 @@ -#ifndef __INFINIOP_MOE_FUSED_DENSE_API_H__ -#define __INFINIOP_MOE_FUSED_DENSE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopMoeFusedDenseDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateMoeFusedDenseDescriptor( - infiniopHandle_t handle, - infiniopMoeFusedDenseDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t hidden_states_desc, - infiniopTensorDescriptor_t w13_desc, - infiniopTensorDescriptor_t w2_desc, - infiniopTensorDescriptor_t topk_weights_desc, - infiniopTensorDescriptor_t topk_ids_desc, - infiniopTensorDescriptor_t sorted_token_ids_desc, - infiniopTensorDescriptor_t expert_ids_desc, - infiniopTensorDescriptor_t num_tokens_post_padded_desc); - -__INFINI_C __export infiniStatus_t infiniopGetMoeFusedDenseWorkspaceSize( - infiniopMoeFusedDenseDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMoeFusedDense( - infiniopMoeFusedDenseDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *hidden_states, - const void *w13, - const void *w2, - const void *topk_weights, - const void *topk_ids, - const void *sorted_token_ids, - const void *expert_ids, - const void *num_tokens_post_padded, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMoeFusedDenseDescriptor( - infiniopMoeFusedDenseDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/moe_fused_gate.h b/include/infiniop/ops/moe_fused_gate.h deleted file mode 100644 index b64a901c1..000000000 --- a/include/infiniop/ops/moe_fused_gate.h +++ /dev/null @@ -1,38 +0,0 @@ -#ifndef __INFINIOP_MOE_FUSED_GATE_API_H__ -#define __INFINIOP_MOE_FUSED_GATE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopMoeFusedGateDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateMoeFusedGateDescriptor( - infiniopHandle_t handle, - infiniopMoeFusedGateDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t topk_weights_desc, - infiniopTensorDescriptor_t topk_indices_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t bias_desc, - size_t num_expert_group, - size_t topk_group, - size_t num_fused_shared_experts, - float routed_scaling_factor, - bool apply_routed_scaling_factor_on_output); - -__INFINI_C __export infiniStatus_t infiniopGetMoeFusedGateWorkspaceSize( - infiniopMoeFusedGateDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMoeFusedGate( - infiniopMoeFusedGateDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *topk_weights, - void *topk_indices, - const void *input, - const void *bias, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMoeFusedGateDescriptor( - infiniopMoeFusedGateDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/moe_sum.h b/include/infiniop/ops/moe_sum.h deleted file mode 100644 index 3b9a79f0d..000000000 --- a/include/infiniop/ops/moe_sum.h +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef __INFINIOP_MOE_SUM_API_H__ -#define __INFINIOP_MOE_SUM_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopMoeSumDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateMoeSumDescriptor( - infiniopHandle_t handle, - infiniopMoeSumDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc); - -__INFINI_C __export infiniStatus_t infiniopGetMoeSumWorkspaceSize( - infiniopMoeSumDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMoeSum( - infiniopMoeSumDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMoeSumDescriptor( - infiniopMoeSumDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/moe_topk_sigmoid.h b/include/infiniop/ops/moe_topk_sigmoid.h deleted file mode 100644 index 0afd40e13..000000000 --- a/include/infiniop/ops/moe_topk_sigmoid.h +++ /dev/null @@ -1,34 +0,0 @@ -#ifndef __INFINIOP_MOE_TOPK_SIGMOID_API_H__ -#define __INFINIOP_MOE_TOPK_SIGMOID_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopMoeTopkSigmoidDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateMoeTopkSigmoidDescriptor( - infiniopHandle_t handle, - infiniopMoeTopkSigmoidDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t topk_weights_desc, - infiniopTensorDescriptor_t topk_indices_desc, - infiniopTensorDescriptor_t gating_output_desc, - infiniopTensorDescriptor_t correction_bias_desc, - bool renormalize); - -__INFINI_C __export infiniStatus_t infiniopGetMoeTopkSigmoidWorkspaceSize( - infiniopMoeTopkSigmoidDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMoeTopkSigmoid( - infiniopMoeTopkSigmoidDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *topk_weights, - void *topk_indices, - const void *gating_output, - const void *correction_bias, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMoeTopkSigmoidDescriptor( - infiniopMoeTopkSigmoidDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/moe_topk_softmax.h b/include/infiniop/ops/moe_topk_softmax.h deleted file mode 100644 index a79425e25..000000000 --- a/include/infiniop/ops/moe_topk_softmax.h +++ /dev/null @@ -1,35 +0,0 @@ -#ifndef __INFINIOP_MOE_TOPK_SOFTMAX_API_H__ -#define __INFINIOP_MOE_TOPK_SOFTMAX_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopMoeTopkSoftmaxDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateMoeTopkSoftmaxDescriptor( - infiniopHandle_t handle, - infiniopMoeTopkSoftmaxDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t topk_weights_desc, - infiniopTensorDescriptor_t topk_indices_desc, - infiniopTensorDescriptor_t gating_output_desc, - infiniopTensorDescriptor_t correction_bias_desc, - bool renormalize, - float moe_softcapping); - -__INFINI_C __export infiniStatus_t infiniopGetMoeTopkSoftmaxWorkspaceSize( - infiniopMoeTopkSoftmaxDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMoeTopkSoftmax( - infiniopMoeTopkSoftmaxDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *topk_weights, - void *topk_indices, - const void *gating_output, - const void *correction_bias, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMoeTopkSoftmaxDescriptor( - infiniopMoeTopkSoftmaxDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/mrope.h b/include/infiniop/ops/mrope.h deleted file mode 100644 index be33bdcce..000000000 --- a/include/infiniop/ops/mrope.h +++ /dev/null @@ -1,42 +0,0 @@ -#ifndef __INFINIOP_MROPE_API_H__ -#define __INFINIOP_MROPE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopMRoPEDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateMRoPEDescriptor( - infiniopHandle_t handle, - infiniopMRoPEDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t q_out, - infiniopTensorDescriptor_t k_out, - infiniopTensorDescriptor_t q, - infiniopTensorDescriptor_t k, - infiniopTensorDescriptor_t cos, - infiniopTensorDescriptor_t sin, - infiniopTensorDescriptor_t positions, - int head_size, - int rotary_dim, - int section_t, - int section_h, - int section_w, - bool interleaved); - -__INFINI_C __export infiniStatus_t infiniopGetMRoPEWorkspaceSize(infiniopMRoPEDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMRoPE( - infiniopMRoPEDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *q_out, - void *k_out, - const void *q, - const void *k, - void const *cos, - void const *sin, - void const *positions, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMRoPEDescriptor(infiniopMRoPEDescriptor_t desc); - -#endif // __INFINIOP_MROPE_API_H__ diff --git a/include/infiniop/ops/mul.h b/include/infiniop/ops/mul.h deleted file mode 100644 index e3e48c39d..000000000 --- a/include/infiniop/ops/mul.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_MUL_API_H__ -#define __INFINIOP_MUL_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopMulDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateMulDescriptor(infiniopHandle_t handle, - infiniopMulDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c, - infiniopTensorDescriptor_t a, - infiniopTensorDescriptor_t b); - -__INFINI_C __export infiniStatus_t infiniopGetMulWorkspaceSize(infiniopMulDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMul(infiniopMulDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMulDescriptor(infiniopMulDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/mul_scalar.h b/include/infiniop/ops/mul_scalar.h deleted file mode 100644 index 9f90284ee..000000000 --- a/include/infiniop/ops/mul_scalar.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __INFINIOP_MUL_SCALAR_API_H__ -#define __INFINIOP_MUL_SCALAR_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopMulScalarDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateMulScalarDescriptor(infiniopHandle_t handle, - infiniopMulScalarDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input); - -__INFINI_C __export infiniStatus_t infiniopGetMulScalarWorkspaceSize(infiniopMulScalarDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMulScalar(infiniopMulScalarDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - double alpha, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMulScalarDescriptor(infiniopMulScalarDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/multi_margin_loss.h b/include/infiniop/ops/multi_margin_loss.h deleted file mode 100644 index 9789e2ae4..000000000 --- a/include/infiniop/ops/multi_margin_loss.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef __INFINIOP_MULTI_MARGIN_LOSS_API_H__ -#define __INFINIOP_MULTI_MARGIN_LOSS_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopMultiMarginLossDescriptor_t; -__INFINI_C __export infiniStatus_t infiniopCreateMultiMarginLossDescriptor(infiniopHandle_t handle, - infiniopMultiMarginLossDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - infiniopTensorDescriptor_t target, - infiniopTensorDescriptor_t weight, - int p, - float margin, - int reduction); - -__INFINI_C __export infiniStatus_t infiniopGetMultiMarginLossWorkspaceSize(infiniopMultiMarginLossDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopMultiMarginLoss(infiniopMultiMarginLossDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *target, - const void *weight, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyMultiMarginLossDescriptor(infiniopMultiMarginLossDescriptor_t desc); - -#endif // __INFINIOP_MULTI_MARGIN_LOSS_API_H__ diff --git a/include/infiniop/ops/nrm2.h b/include/infiniop/ops/nrm2.h deleted file mode 100644 index 2f1eed348..000000000 --- a/include/infiniop/ops/nrm2.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_NRM2_API_H__ -#define __INFINIOP_NRM2_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopNrm2Descriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateNrm2Descriptor(infiniopHandle_t handle, - infiniopNrm2Descriptor_t *desc_ptr, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t result); - -__INFINI_C __export infiniStatus_t infiniopGetNrm2WorkspaceSize(infiniopNrm2Descriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopNrm2(infiniopNrm2Descriptor_t desc, - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyNrm2Descriptor(infiniopNrm2Descriptor_t desc); - -#endif // __INFINIOP_NRM2_API_H__ diff --git a/include/infiniop/ops/nsa_compress_paged_cache.h b/include/infiniop/ops/nsa_compress_paged_cache.h deleted file mode 100644 index f69d3fb5d..000000000 --- a/include/infiniop/ops/nsa_compress_paged_cache.h +++ /dev/null @@ -1,39 +0,0 @@ -#ifndef __INFINIOP_NSA_COMPRESS_PAGED_CACHE_API_H__ -#define __INFINIOP_NSA_COMPRESS_PAGED_CACHE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopNsaCompressPagedCacheDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateNsaCompressPagedCacheDescriptor( - infiniopHandle_t handle, - infiniopNsaCompressPagedCacheDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t k_cmp_desc, - infiniopTensorDescriptor_t v_cmp_desc, - infiniopTensorDescriptor_t k_cache_desc, - infiniopTensorDescriptor_t v_cache_desc, - infiniopTensorDescriptor_t block_tables_desc, - infiniopTensorDescriptor_t seq_lens_desc, - int nsa_block_size, - int update_last_only); - -__INFINI_C __export infiniStatus_t infiniopGetNsaCompressPagedCacheWorkspaceSize( - infiniopNsaCompressPagedCacheDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopNsaCompressPagedCache( - infiniopNsaCompressPagedCacheDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *k_cmp, - void *v_cmp, - const void *k_cache, - const void *v_cache, - const void *block_tables, - const void *seq_lens, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyNsaCompressPagedCacheDescriptor( - infiniopNsaCompressPagedCacheDescriptor_t desc); - -#endif // __INFINIOP_NSA_COMPRESS_PAGED_CACHE_API_H__ diff --git a/include/infiniop/ops/nsa_paged_attention.h b/include/infiniop/ops/nsa_paged_attention.h deleted file mode 100644 index e92414f86..000000000 --- a/include/infiniop/ops/nsa_paged_attention.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef __INFINIOP_NSA_PAGED_ATTENTION_API_H__ -#define __INFINIOP_NSA_PAGED_ATTENTION_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopNsaPagedAttentionDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateNsaPagedAttentionDescriptor( - infiniopHandle_t handle, - infiniopNsaPagedAttentionDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_cmp_desc, - infiniopTensorDescriptor_t v_cmp_desc, - infiniopTensorDescriptor_t k_cache_desc, - infiniopTensorDescriptor_t v_cache_desc, - infiniopTensorDescriptor_t block_tables_desc, - infiniopTensorDescriptor_t seq_lens_desc, - infiniopTensorDescriptor_t gates_desc, - float scale, - int nsa_block_size, - int window_size, - int select_blocks); - -__INFINI_C __export infiniStatus_t infiniopGetNsaPagedAttentionWorkspaceSize( - infiniopNsaPagedAttentionDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopNsaPagedAttention( - infiniopNsaPagedAttentionDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *q, - const void *k_cmp, - const void *v_cmp, - const void *k_cache, - const void *v_cache, - const void *block_tables, - const void *seq_lens, - const void *gates, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyNsaPagedAttentionDescriptor( - infiniopNsaPagedAttentionDescriptor_t desc); - -#endif // __INFINIOP_NSA_PAGED_ATTENTION_API_H__ diff --git a/include/infiniop/ops/ones.h b/include/infiniop/ops/ones.h deleted file mode 100644 index e66b0f884..000000000 --- a/include/infiniop/ops/ones.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_ONES_API_H__ -#define __INFINIOP_ONES_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopOnesDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateOnesDescriptor(infiniopHandle_t handle, - infiniopOnesDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetOnesWorkspaceSize(infiniopOnesDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopOnes(infiniopOnesDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyOnesDescriptor(infiniopOnesDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/pad.h b/include/infiniop/ops/pad.h deleted file mode 100644 index af3af576f..000000000 --- a/include/infiniop/ops/pad.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_PAD_API_H__ -#define __INFINIOP_PAD_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopPadDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreatePadDescriptor(infiniopHandle_t handle, - infiniopPadDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - void *pad, - size_t pad_size, - const char *mode, - double value); - -__INFINI_C __export infiniStatus_t infiniopGetPadWorkspaceSize(infiniopPadDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopPad(infiniopPadDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyPadDescriptor(infiniopPadDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/paged_attention.h b/include/infiniop/ops/paged_attention.h deleted file mode 100644 index 7f1656ef3..000000000 --- a/include/infiniop/ops/paged_attention.h +++ /dev/null @@ -1,93 +0,0 @@ -#ifndef __INFINIOP_PAGED_ATTENTION_API_H__ -#define __INFINIOP_PAGED_ATTENTION_API_H__ - -#include "../operator_descriptor.h" - -// Define an opaque handle for the Paged Attention descriptor. -typedef struct InfiniopDescriptor *infiniopPagedAttentionDescriptor_t; - -/** - * @brief Creates a descriptor for the Paged Attention v1 operation. - * - * @param handle The library context handle. - * @param desc_ptr Pointer to the created descriptor. - * @param out_desc [Output] Shape: (num_seqs, num_heads, head_size). - * The output tensor for the attention mechanism. - * @param q_desc [Input] Shape: (num_seqs, num_heads, head_size). - * The query tensor. - * @param k_cache_desc [Input] Shape: (num_blocks, num_kv_heads, block_size, head_size). - * Paged key cache storing keys for all sequences. - * @param v_cache_desc [Input] Shape: (num_blocks, num_kv_heads, block_size, head_size). - * Paged value cache storing values for all sequences. - * @param block_tables_desc [Input] Shape: (num_seqs, max_num_blocks_per_seq). - * Maps each sequence to its physical block indices in the cache. - * Expected DType: int64_t (I64). - * @param seq_lens_desc [Input] Shape: (num_seqs,). - * The current logical length of each sequence. - * Expected DType: int64_t (I64). - * @param alibi_slopes_desc [Optional] Shape: (num_heads,). - * Slopes for ALiBi (Attention with Linear Biases). Can be NULL. - * @param scale The attention scaling factor (typically 1/sqrt(head_size)). - * @return infiniStatus_t Status code. - */ -__INFINI_C __export infiniStatus_t infiniopCreatePagedAttentionDescriptor( - infiniopHandle_t handle, - infiniopPagedAttentionDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_cache_desc, - infiniopTensorDescriptor_t v_cache_desc, - infiniopTensorDescriptor_t block_tables_desc, - infiniopTensorDescriptor_t seq_lens_desc, - infiniopTensorDescriptor_t alibi_slopes_desc, - float scale); - -/** - * @brief Retrieves the workspace size required for the Paged Attention operation. - * - * @param desc The Paged Attention descriptor. - * @param size A pointer to store the required workspace size in bytes. - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopGetPagedAttentionWorkspaceSize( - infiniopPagedAttentionDescriptor_t desc, size_t *size); - -/** - * @brief Executes the Paged Attention v1 operation. - * - * @param desc The Paged Attention descriptor. - * @param workspace Pointer to the workspace memory. - * @param workspace_size The size of the workspace. - * @param out Pointer to the output tensor data. - * @param q Pointer to the query tensor data. - * @param k_cache Pointer to the key cache data. - * @param v_cache Pointer to the value cache data. - * @param block_tables Pointer to the block tables data. - * @param seq_lens Pointer to the sequence lengths data. - * @param alibi_slopes Pointer to the ALiBi slopes data. Can be NULL. - * @param stream The CUDA stream for the operation. Can be NULL. - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopPagedAttention( - infiniopPagedAttentionDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *q, - const void *k_cache, - const void *v_cache, - const void *block_tables, - const void *seq_lens, - const void *alibi_slopes, - void *stream); - -/** - * @brief Destroys a Paged Attention descriptor. - * - * @param desc The descriptor to be destroyed. - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopDestroyPagedAttentionDescriptor( - infiniopPagedAttentionDescriptor_t desc); - -#endif // __INFINIOP_PAGED_ATTENTION_API_H__ diff --git a/include/infiniop/ops/paged_attention_prefill.h b/include/infiniop/ops/paged_attention_prefill.h deleted file mode 100644 index e2e93076b..000000000 --- a/include/infiniop/ops/paged_attention_prefill.h +++ /dev/null @@ -1,87 +0,0 @@ -#ifndef __INFINIOP_PAGED_ATTENTION_PREFILL_API_H__ -#define __INFINIOP_PAGED_ATTENTION_PREFILL_API_H__ - -#include "../operator_descriptor.h" - -// Define an opaque handle for the Paged Attention Prefill descriptor. -typedef struct InfiniopDescriptor *infiniopPagedAttentionPrefillDescriptor_t; - -/** - * @brief Creates a descriptor for the Paged Attention Prefill operation. - * @param handle The handle to the InfiniOP library context. - * @param desc_ptr A pointer to store the created descriptor. - * @param out_desc Descriptor for the output tensor. - * Shape: [total_q_tokens, num_heads, head_size] - * @param q_desc Descriptor for the query tensor (packed/flattened). - * Shape: [total_q_tokens, num_heads, head_size] - * @param k_cache_desc Descriptor for the global physical key cache. - * Shape: [max_num_blocks, num_kv_heads, block_size, head_size] - * @param v_cache_desc Descriptor for the global physical value cache. - * Shape: [max_num_blocks, num_kv_heads, block_size, head_size] - * @param block_tables_desc Descriptor for the block tables mapping logic to physical blocks. - * Shape: [batch_size, max_blocks_per_seq] - * @param seq_lens_desc Descriptor for the total KV lengths of each sequence. - * Shape: [batch_size] - * @param cum_seq_lens_q_desc Descriptor for the cumulative start position (prefix sum) of each Q sequence. - * Shape: [batch_size + 1] - * @param alibi_slopes_desc Optional descriptor for the ALiBi slopes tensor. Can be NULL. - * Shape: [num_heads] - * @param scale The attention scaling factor (typically 1.0 / sqrt(head_size)). - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopCreatePagedAttentionPrefillDescriptor( - infiniopHandle_t handle, - infiniopPagedAttentionPrefillDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_cache_desc, - infiniopTensorDescriptor_t v_cache_desc, - infiniopTensorDescriptor_t block_tables_desc, - infiniopTensorDescriptor_t seq_lens_desc, - infiniopTensorDescriptor_t cum_seq_lens_q_desc, - infiniopTensorDescriptor_t alibi_slopes_desc, - float scale); - -/** - * @brief Retrieves the workspace size required for the Paged Attention Prefill operation. - */ -__INFINI_C __export infiniStatus_t infiniopGetPagedAttentionPrefillWorkspaceSize( - infiniopPagedAttentionPrefillDescriptor_t desc, size_t *size); - -/** - * @brief Executes the Paged Attention Prefill operation. - * @param desc The Paged Attention Prefill descriptor. - * @param workspace Pointer to the workspace memory. - * @param workspace_size The size of the workspace. - * @param out Pointer to the output tensor data. - * @param q Pointer to the query tensor data (packed). - * @param k_cache Pointer to the global key cache data. - * @param v_cache Pointer to the global value cache data. - * @param block_tables Pointer to the block tables data. - * @param seq_lens Pointer to the KV lengths data. - * @param cum_seq_lens_q Pointer to the Q cumulative sequence lengths data (prefix sum). - * @param alibi_slopes Pointer to the ALiBi slopes data. Can be NULL. - * @param stream The device stream (e.g., cudaStream_t) for the operation. - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopPagedAttentionPrefill( - infiniopPagedAttentionPrefillDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *q, - const void *k_cache, - const void *v_cache, - const void *block_tables, - const void *seq_lens, - const void *cum_seq_lens_q, - const void *alibi_slopes, - void *stream); - -/** - * @brief Destroys a Paged Attention Prefill descriptor. - */ -__INFINI_C __export infiniStatus_t infiniopDestroyPagedAttentionPrefillDescriptor( - infiniopPagedAttentionPrefillDescriptor_t desc); - -#endif // __INFINIOP_PAGED_ATTENTION_PREFILL_API_H__ diff --git a/include/infiniop/ops/paged_caching.h b/include/infiniop/ops/paged_caching.h deleted file mode 100644 index d85125c30..000000000 --- a/include/infiniop/ops/paged_caching.h +++ /dev/null @@ -1,77 +0,0 @@ -#ifndef __INFINIOP_PAGED_CACHING_API_H__ -#define __INFINIOP_PAGED_CACHING_API_H__ - -#include "../operator_descriptor.h" - -// Define an opaque handle for the Paged Caching descriptor. -typedef struct InfiniopDescriptor *infiniopPagedCachingDescriptor_t; - -/** - * @brief Creates a descriptor for the Paged Caching operation. - * - * This function initializes a descriptor that holds all the metadata needed - * to copy key/value vectors into their respective cache pools. - * - * @param handle The handle to the InfiniOP library context. - * @param desc_ptr A pointer to store the created descriptor. - * @param k_cache_desc Descriptor for the key cache pool tensor. - * @param v_cache_desc Descriptor for the value cache pool tensor. - * @param k_desc Descriptor for the source key tensor. - * @param v_desc Descriptor for the source value tensor. - * @param slot_mapping_desc Descriptor for the slot mapping tensor. - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopCreatePagedCachingDescriptor( - infiniopHandle_t handle, - infiniopPagedCachingDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t k_cache_desc, - infiniopTensorDescriptor_t v_cache_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t slot_mapping_desc); - -/** - * @brief Retrieves the workspace size required for the Paged Caching operation. - * - * @param desc The Paged Caching descriptor. - * @param size A pointer to store the required workspace size in bytes (typically 0). - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopGetPagedCachingWorkspaceSize( - infiniopPagedCachingDescriptor_t desc, size_t *size); - -/** - * @brief Executes the Paged Caching operation. - * - * @param desc The Paged Caching descriptor. - * @param workspace Pointer to the workspace memory. - * @param workspace_size The size of the workspace. - * @param k_cache Pointer to the key cache pool data. - * @param v_cache Pointer to the value cache pool data. - * @param k Pointer to the source key tensor data. - * @param v Pointer to the source value tensor data. - * @param slot_mapping Pointer to the slot mapping data. - * @param stream The CUDA stream for the operation. Can be NULL. - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopPagedCaching( - infiniopPagedCachingDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *k_cache, - void *v_cache, - const void *k, - const void *v, - const void *slot_mapping, - void *stream); - -/** - * @brief Destroys a Paged Caching descriptor. - * - * @param desc The descriptor to be destroyed. - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopDestroyPagedCachingDescriptor( - infiniopPagedCachingDescriptor_t desc); - -#endif // __INFINIOP_PAGED_CACHING_API_H__ diff --git a/include/infiniop/ops/pixel_shuffle.h b/include/infiniop/ops/pixel_shuffle.h deleted file mode 100644 index fac8eff2c..000000000 --- a/include/infiniop/ops/pixel_shuffle.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __INFINIOP_PIXEL_SHUFFLE_API_H__ -#define __INFINIOP_PIXEL_SHUFFLE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopPixelShuffleDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreatePixelShuffleDescriptor(infiniopHandle_t handle, - infiniopPixelShuffleDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - int upscale_factor); - -__INFINI_C __export infiniStatus_t infiniopGetPixelShuffleWorkspaceSize(infiniopPixelShuffleDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopPixelShuffle(infiniopPixelShuffleDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyPixelShuffleDescriptor(infiniopPixelShuffleDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/prelu.h b/include/infiniop/ops/prelu.h deleted file mode 100644 index e90a571be..000000000 --- a/include/infiniop/ops/prelu.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_PRELU_API_H__ -#define __INFINIOP_PRELU_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopPreluDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreatePreluDescriptor(infiniopHandle_t handle, - infiniopPreluDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t weight); - -__INFINI_C __export infiniStatus_t infiniopGetPreluWorkspaceSize(infiniopPreluDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopPrelu(infiniopPreluDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *weight, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyPreluDescriptor(infiniopPreluDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/prepare_moe_input.h b/include/infiniop/ops/prepare_moe_input.h deleted file mode 100644 index 8621d044a..000000000 --- a/include/infiniop/ops/prepare_moe_input.h +++ /dev/null @@ -1,42 +0,0 @@ -#ifndef __INFINIOP_PREPARE_MOE_INPUT_API_H__ -#define __INFINIOP_PREPARE_MOE_INPUT_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopPrepareMoeInputDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreatePrepareMoeInputDescriptor( - infiniopHandle_t handle, - infiniopPrepareMoeInputDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t expert_offsets_desc, - infiniopTensorDescriptor_t blockscale_offsets_desc, - infiniopTensorDescriptor_t problem_sizes1_desc, - infiniopTensorDescriptor_t problem_sizes2_desc, - infiniopTensorDescriptor_t input_permutation_desc, - infiniopTensorDescriptor_t output_permutation_desc, - infiniopTensorDescriptor_t topk_ids_desc, - size_t num_experts, - size_t n, - size_t k); - -__INFINI_C __export infiniStatus_t infiniopGetPrepareMoeInputWorkspaceSize( - infiniopPrepareMoeInputDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopPrepareMoeInput( - infiniopPrepareMoeInputDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *expert_offsets, - void *blockscale_offsets, - void *problem_sizes1, - void *problem_sizes2, - void *input_permutation, - void *output_permutation, - const void *topk_ids, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyPrepareMoeInputDescriptor( - infiniopPrepareMoeInputDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/quant/per_channel_quant_int8.h b/include/infiniop/ops/quant/per_channel_quant_int8.h deleted file mode 100644 index faa5efd78..000000000 --- a/include/infiniop/ops/quant/per_channel_quant_int8.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_PER_CHANNEL_QUANT_INT8_API_H__ -#define __INFINIOP_PER_CHANNEL_QUANT_INT8_API_H__ - -#include "../../operator_descriptor.h" - -typedef InfiniopDescriptor *infiniopPerChannelQuantI8Descriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreatePerChannelQuantI8Descriptor(infiniopHandle_t handle, - infiniopPerChannelQuantI8Descriptor_t *desc_ptr, - infiniopTensorDescriptor_t x_packed_desc, - infiniopTensorDescriptor_t x_scale_desc, - infiniopTensorDescriptor_t x_zero_desc, - infiniopTensorDescriptor_t x_desc); - -__INFINI_C __export infiniStatus_t infiniopGetPerChannelQuantI8WorkspaceSize(infiniopPerChannelQuantI8Descriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopPerChannelQuantI8(infiniopPerChannelQuantI8Descriptor_t desc, - void *workspace, - size_t workspace_size, - void *x_packed, - void *x_scale, - void *x_zero, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyPerChannelQuantI8Descriptor(infiniopPerChannelQuantI8Descriptor_t desc); - -#endif diff --git a/include/infiniop/ops/quant/per_tensor_quant_int8.h b/include/infiniop/ops/quant/per_tensor_quant_int8.h deleted file mode 100644 index 16e1c2bc6..000000000 --- a/include/infiniop/ops/quant/per_tensor_quant_int8.h +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef __INFINIOP_PER_TENSOR_QUANT_INT8_API_H__ -#define __INFINIOP_PER_TENSOR_QUANT_INT8_API_H__ - -#include "../../operator_descriptor.h" - -typedef InfiniopDescriptor *infiniopPerTensorQuantI8Descriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreatePerTensorQuantI8Descriptor(infiniopHandle_t handle, - infiniopPerTensorQuantI8Descriptor_t *desc_ptr, - infiniopTensorDescriptor_t x_packed_desc, - infiniopTensorDescriptor_t x_scale_desc, - infiniopTensorDescriptor_t x_zero_desc, - infiniopTensorDescriptor_t x_desc); - -__INFINI_C __export infiniStatus_t infiniopGetPerTensorQuantI8WorkspaceSize(infiniopPerTensorQuantI8Descriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopPerTensorQuantI8(infiniopPerTensorQuantI8Descriptor_t desc, - void *workspace, - size_t workspace_size, - void *x_packed, - void *x_scale, - void *x_zero, - const void *x, - const bool is_static, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyPerTensorQuantI8Descriptor(infiniopPerTensorQuantI8Descriptor_t desc); - -#endif diff --git a/include/infiniop/ops/quickgelu.h b/include/infiniop/ops/quickgelu.h deleted file mode 100644 index be66f9495..000000000 --- a/include/infiniop/ops/quickgelu.h +++ /dev/null @@ -1,42 +0,0 @@ -#ifndef __INFINIOP_QUICKGELU_API_H__ -#define __INFINIOP_QUICKGELU_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopQuickGeluDescriptor_t; - -/** - * Create QuickGELU descriptor - * y = x * sigmoid(1.702 * x) - */ -__INFINI_C __export infiniStatus_t infiniopCreateQuickGeluDescriptor( - infiniopHandle_t handle, - infiniopQuickGeluDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -/** - * Query workspace size - */ -__INFINI_C __export infiniStatus_t infiniopGetQuickGeluWorkspaceSize( - infiniopQuickGeluDescriptor_t desc, - size_t *size); - -/** - * Launch QuickGELU operator - */ -__INFINI_C __export infiniStatus_t infiniopQuickGelu( - infiniopQuickGeluDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -/** - * Destroy descriptor - */ -__INFINI_C __export infiniStatus_t infiniopDestroyQuickGeluDescriptor( - infiniopQuickGeluDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/random_sample.h b/include/infiniop/ops/random_sample.h deleted file mode 100644 index c12dc0af3..000000000 --- a/include/infiniop/ops/random_sample.h +++ /dev/null @@ -1,52 +0,0 @@ -#ifndef __INFINIOP_RANDOM_SAMPLE_API_H__ -#define __INFINIOP_RANDOM_SAMPLE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopRandomSampleDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateRandomSampleDescriptor( - infiniopHandle_t handle, - infiniopRandomSampleDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t result, - infiniopTensorDescriptor_t probs); - -__INFINI_C __export infiniStatus_t infiniopGetRandomSampleWorkspaceSize( - infiniopRandomSampleDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopCreateRandomSampleBatchDescriptor( - infiniopHandle_t handle, - infiniopRandomSampleDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t result, - infiniopTensorDescriptor_t probs); - -__INFINI_C __export infiniStatus_t infiniopRandomSample( - infiniopRandomSampleDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *result, - const void *probs, - float random_val, - float topp, - int topk, - float temperature, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopRandomSampleBatch( - infiniopRandomSampleDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *result, - const void *probs, - const float *random_val, - const float *topp, - const int *topk, - const float *temperature, - int batch_size, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyRandomSampleDescriptor( - infiniopRandomSampleDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/rearrange.h b/include/infiniop/ops/rearrange.h deleted file mode 100644 index e28aeb97d..000000000 --- a/include/infiniop/ops/rearrange.h +++ /dev/null @@ -1,23 +0,0 @@ -#ifndef __INFINIOP_REARRANGE_API_H__ -#define __INFINIOP_REARRANGE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopRearrangeDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateRearrangeDescriptor( - infiniopHandle_t handle, - infiniopRearrangeDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t dst, - infiniopTensorDescriptor_t src); - -__INFINI_C __export infiniStatus_t infiniopRearrange( - infiniopRearrangeDescriptor_t desc, - void *dst, - const void *src, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyRearrangeDescriptor( - infiniopRearrangeDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/reciprocal.h b/include/infiniop/ops/reciprocal.h deleted file mode 100644 index bd1ea31c5..000000000 --- a/include/infiniop/ops/reciprocal.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_RECIPROCAL_API_H__ -#define __INFINIOP_RECIPROCAL_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopReciprocalDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateReciprocalDescriptor(infiniopHandle_t handle, - infiniopReciprocalDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetReciprocalWorkspaceSize(infiniopReciprocalDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopReciprocal(infiniopReciprocalDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyReciprocalDescriptor(infiniopReciprocalDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/recurrent_gated_delta_rule.h b/include/infiniop/ops/recurrent_gated_delta_rule.h deleted file mode 100644 index 6a1b29c65..000000000 --- a/include/infiniop/ops/recurrent_gated_delta_rule.h +++ /dev/null @@ -1,46 +0,0 @@ -#ifndef __INFINIOP_RECURRENT_GATED_DELTA_RULE_API_H__ -#define __INFINIOP_RECURRENT_GATED_DELTA_RULE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopRecurrentGatedDeltaRuleDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateRecurrentGatedDeltaRuleDescriptor( - infiniopHandle_t handle, - infiniopRecurrentGatedDeltaRuleDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, // [B, T, Hv, Dv], T must be 1; last dim contiguous - infiniopTensorDescriptor_t initial_state_desc, // legacy: [B, Hv, Dv, Dk]; indexed pool: [pool_size, Hv, Dv, Dk] - infiniopTensorDescriptor_t final_state_desc, // legacy/indexed out-of-place final state; null when final_state_indices_desc is provided - infiniopTensorDescriptor_t q_desc, // [B, T, Hk, Dk], T must be 1; last dim contiguous - infiniopTensorDescriptor_t k_desc, // [B, T, Hk, Dk], same shape as q; last dim contiguous - infiniopTensorDescriptor_t v_desc, // [B, T, Hv, Dv], Hv must be a multiple of Hk; last dim contiguous - infiniopTensorDescriptor_t g_desc, // [B, T, Hv]; may have a different fp dtype from q/k/v/out/state - infiniopTensorDescriptor_t beta_desc, // [B, T, Hv]; same dtype as g - infiniopTensorDescriptor_t initial_state_indices_desc, // nullable; [B], int32/int64; enables indexed pool mode - infiniopTensorDescriptor_t final_state_indices_desc, // nullable; [B], int32/int64; writes final state in-place to initial_state pool - bool use_qk_l2norm); - -__INFINI_C __export infiniStatus_t infiniopGetRecurrentGatedDeltaRuleWorkspaceSize( - infiniopRecurrentGatedDeltaRuleDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopRecurrentGatedDeltaRule( - infiniopRecurrentGatedDeltaRuleDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - void *initial_state, - void *final_state, - const void *q, - const void *k, - const void *v, - const void *g, - const void *beta, - const void *initial_state_indices, - const void *final_state_indices, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyRecurrentGatedDeltaRuleDescriptor( - infiniopRecurrentGatedDeltaRuleDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/relu.h b/include/infiniop/ops/relu.h deleted file mode 100644 index 8b129badd..000000000 --- a/include/infiniop/ops/relu.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_RELU_API_H__ -#define __INFINIOP_RELU_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopReluDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateReluDescriptor(infiniopHandle_t handle, - infiniopReluDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetReluWorkspaceSize(infiniopReluDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopRelu(infiniopReluDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyReluDescriptor(infiniopReluDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/relu6.h b/include/infiniop/ops/relu6.h deleted file mode 100644 index 2246bbc55..000000000 --- a/include/infiniop/ops/relu6.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_RELU6_API_H__ -#define __INFINIOP_RELU6_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopRelu6Descriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateRelu6Descriptor(infiniopHandle_t handle, - infiniopRelu6Descriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetRelu6WorkspaceSize(infiniopRelu6Descriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopRelu6(infiniopRelu6Descriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyRelu6Descriptor(infiniopRelu6Descriptor_t desc); - -#endif diff --git a/include/infiniop/ops/rms_norm.h b/include/infiniop/ops/rms_norm.h deleted file mode 100644 index f90da65e4..000000000 --- a/include/infiniop/ops/rms_norm.h +++ /dev/null @@ -1,23 +0,0 @@ -#ifndef __INFINIOP_RMS_NORM_API_H__ -#define __INFINIOP_RMS_NORM_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopRMSNormDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateRMSNormDescriptor( - infiniopHandle_t handle, - infiniopRMSNormDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t w_desc, - float epsilon); - -__INFINI_C __export infiniStatus_t infiniopGetRMSNormWorkspaceSize(infiniopRMSNormDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopRMSNorm(infiniopRMSNormDescriptor_t desc, void *workspace, size_t workspace_size, - void *y, const void *x, const void *w, void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyRMSNormDescriptor(infiniopRMSNormDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/rope.h b/include/infiniop/ops/rope.h deleted file mode 100644 index 7c91d51f9..000000000 --- a/include/infiniop/ops/rope.h +++ /dev/null @@ -1,40 +0,0 @@ -#ifndef __INFINIOP_ROPE_API_H__ -#define __INFINIOP_ROPE_API_H__ - -#include "../operator_descriptor.h" - -typedef enum { - INFINIOP_ROPE_ALGO_GPT_J = 0, // GPT-J style RoPE algorithm (Interleave even and odd dimensions) - INFINIOP_ROPE_ALGO_GPT_NEOX = 1, // GPT-NeoX style RoPE algorithm (First half dimensions for sin, second half for cos) - // Count - INFINIOP_ROPE_ALGO_COUNT = 2, -} infiniopRoPEAlgo_t; - -typedef struct InfiniopDescriptor *infiniopRoPEDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateRoPEDescriptor( - infiniopHandle_t handle, - infiniopRoPEDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t pos_ids, - infiniopTensorDescriptor_t sin_table, - infiniopTensorDescriptor_t cos_table, - infiniopRoPEAlgo_t algo); - -__INFINI_C __export infiniStatus_t infiniopGetRoPEWorkspaceSize(infiniopRoPEDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopRoPE( - infiniopRoPEDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void const *pos_ids, - void const *sin_table, - void const *cos_table, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyRoPEDescriptor(infiniopRoPEDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/rot.h b/include/infiniop/ops/rot.h deleted file mode 100644 index 0ecbae52d..000000000 --- a/include/infiniop/ops/rot.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_ROT_API_H__ -#define __INFINIOP_ROT_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopRotDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateRotDescriptor(infiniopHandle_t handle, - infiniopRotDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t c, - infiniopTensorDescriptor_t s); - -__INFINI_C __export infiniStatus_t infiniopGetRotWorkspaceSize(infiniopRotDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopRot(infiniopRotDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *x, - void *y, - const void *c, - const void *s, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyRotDescriptor(infiniopRotDescriptor_t desc); - -#endif // __INFINIOP_ROT_API_H__ diff --git a/include/infiniop/ops/rotg.h b/include/infiniop/ops/rotg.h deleted file mode 100644 index 63c2dad46..000000000 --- a/include/infiniop/ops/rotg.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_ROTG_API_H__ -#define __INFINIOP_ROTG_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopRotgDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateRotgDescriptor(infiniopHandle_t handle, - infiniopRotgDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t c, - infiniopTensorDescriptor_t s); - -__INFINI_C __export infiniStatus_t infiniopGetRotgWorkspaceSize(infiniopRotgDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopRotg(infiniopRotgDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *x, - void *y, - void *c, - void *s, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyRotgDescriptor(infiniopRotgDescriptor_t desc); - -#endif // __INFINIOP_ROTG_API_H__ diff --git a/include/infiniop/ops/rotm.h b/include/infiniop/ops/rotm.h deleted file mode 100644 index 6cc6a636c..000000000 --- a/include/infiniop/ops/rotm.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_ROTM_API_H__ -#define __INFINIOP_ROTM_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopRotmDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateRotmDescriptor(infiniopHandle_t handle, - infiniopRotmDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t param); - -__INFINI_C __export infiniStatus_t infiniopGetRotmWorkspaceSize(infiniopRotmDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopRotm(infiniopRotmDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *x, - void *y, - const void *param, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyRotmDescriptor(infiniopRotmDescriptor_t desc); - -#endif // __INFINIOP_ROTM_API_H__ diff --git a/include/infiniop/ops/rotmg.h b/include/infiniop/ops/rotmg.h deleted file mode 100644 index 0295339cb..000000000 --- a/include/infiniop/ops/rotmg.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef __INFINIOP_ROTMG_API_H__ -#define __INFINIOP_ROTMG_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopRotmgDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateRotmgDescriptor(infiniopHandle_t handle, - infiniopRotmgDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t d1, - infiniopTensorDescriptor_t d2, - infiniopTensorDescriptor_t x1, - infiniopTensorDescriptor_t y1, - infiniopTensorDescriptor_t param); - -__INFINI_C __export infiniStatus_t infiniopGetRotmgWorkspaceSize(infiniopRotmgDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopRotmg(infiniopRotmgDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *d1, - void *d2, - void *x1, - const void *y1, - void *param, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyRotmgDescriptor(infiniopRotmgDescriptor_t desc); - -#endif // __INFINIOP_ROTMG_API_H__ diff --git a/include/infiniop/ops/rwkv5_wkv.h b/include/infiniop/ops/rwkv5_wkv.h deleted file mode 100644 index 6f11c5497..000000000 --- a/include/infiniop/ops/rwkv5_wkv.h +++ /dev/null @@ -1,45 +0,0 @@ -#ifndef __INFINIOP_RWKV5_WKV_API_H__ -#define __INFINIOP_RWKV5_WKV_API_H__ - -#include "../operator_descriptor.h" - -#ifdef __cplusplus -#include -#else -#include -#endif - -typedef struct InfiniopDescriptor *infiniopRwkv5WkvDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateRwkv5WkvDescriptor( - infiniopHandle_t handle, - infiniopRwkv5WkvDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t receptance_desc, - infiniopTensorDescriptor_t key_desc, - infiniopTensorDescriptor_t value_desc, - infiniopTensorDescriptor_t time_decay_desc, - infiniopTensorDescriptor_t time_faaaa_desc, - infiniopTensorDescriptor_t state_desc); - -__INFINI_C __export infiniStatus_t infiniopGetRwkv5WkvWorkspaceSize( - infiniopRwkv5WkvDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopRwkv5Wkv( - infiniopRwkv5WkvDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *receptance, - const void *key, - const void *value, - const void *time_decay, - const void *time_faaaa, - void *state, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyRwkv5WkvDescriptor( - infiniopRwkv5WkvDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/scal.h b/include/infiniop/ops/scal.h deleted file mode 100644 index f7903de56..000000000 --- a/include/infiniop/ops/scal.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_SCAL_API_H__ -#define __INFINIOP_SCAL_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopScalDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateScalDescriptor(infiniopHandle_t handle, - infiniopScalDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t alpha, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetScalWorkspaceSize(infiniopScalDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopScal(infiniopScalDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *alpha, - void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyScalDescriptor(infiniopScalDescriptor_t desc); - -#endif // __INFINIOP_SCAL_API_H__ diff --git a/include/infiniop/ops/scatter.h b/include/infiniop/ops/scatter.h deleted file mode 100644 index 2fa22c534..000000000 --- a/include/infiniop/ops/scatter.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef __INFINIOP_SCATTER_API_H__ -#define __INFINIOP_SCATTER_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopScatterDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateScatterDescriptor(infiniopHandle_t handle, - infiniopScatterDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - infiniopTensorDescriptor_t indices, - infiniopTensorDescriptor_t updates, - int axis, - int reduction); - -__INFINI_C __export infiniStatus_t infiniopGetScatterWorkspaceSize(infiniopScatterDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopScatter(infiniopScatterDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *indices, - const void *updates, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyScatterDescriptor(infiniopScatterDescriptor_t desc); - -#endif // __INFINIOP_SCATTER_API_H__ diff --git a/include/infiniop/ops/selu.h b/include/infiniop/ops/selu.h deleted file mode 100644 index b1bf3dfcf..000000000 --- a/include/infiniop/ops/selu.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_SELU_API_H__ -#define __INFINIOP_SELU_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopSeluDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateSeluDescriptor(infiniopHandle_t handle, - infiniopSeluDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetSeluWorkspaceSize(infiniopSeluDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopSelu(infiniopSeluDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroySeluDescriptor(infiniopSeluDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/sigmoid.h b/include/infiniop/ops/sigmoid.h deleted file mode 100644 index 461fbdf7a..000000000 --- a/include/infiniop/ops/sigmoid.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_SIGMOID_API_H__ -#define __INFINIOP_SIGMOID_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopSigmoidDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateSigmoidDescriptor(infiniopHandle_t handle, - infiniopSigmoidDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetSigmoidWorkspaceSize(infiniopSigmoidDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopSigmoid(infiniopSigmoidDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroySigmoidDescriptor(infiniopSigmoidDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/silu.h b/include/infiniop/ops/silu.h deleted file mode 100644 index f864962c2..000000000 --- a/include/infiniop/ops/silu.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_SILU_API_H__ -#define __INFINIOP_SILU_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopSiluDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateSiluDescriptor(infiniopHandle_t handle, - infiniopSiluDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t intput); - -__INFINI_C __export infiniStatus_t infiniopGetSiluWorkspaceSize(infiniopSiluDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopSilu(infiniopSiluDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *intput, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroySiluDescriptor(infiniopSiluDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/silu_and_mul.h b/include/infiniop/ops/silu_and_mul.h deleted file mode 100644 index b1effcec8..000000000 --- a/include/infiniop/ops/silu_and_mul.h +++ /dev/null @@ -1,71 +0,0 @@ -#ifndef __INFINIOP_SILU_AND_MUL_API_H__ -#define __INFINIOP_SILU_AND_MUL_API_H__ - -#include "../operator_descriptor.h" - -/** - * @brief Opaque handle for the SiluAndMul descriptor. - */ -typedef struct InfiniopDescriptor *infiniopSiluAndMulDescriptor_t; - -/** - * @brief Creates a descriptor for the SiLU and Multiply (SiluAndMul) operation. - * - * Format: (input_shape, output_shape) - * Referencing vLLM kernel SiluAndMul interface: - * - input_shape is [..., 2*d] (last dimension is split into two halves for SiLU and multiplication) - * - output_shape is [..., d] (last dimension reduced to half) - * - * @param handle The handle to the InfiniOP library context. - * @param desc_ptr A pointer to store the created descriptor. - * @param output Descriptor for the output tensor. Shape [..., d]. - * @param input Descriptor for the input tensor. Shape [..., 2*d]. - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopCreateSiluAndMulDescriptor( - infiniopHandle_t handle, - infiniopSiluAndMulDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input); - -/** - * @brief Queries the workspace size required for SiluAndMul computation. - * @param desc The SiluAndMul descriptor. - * @param size Pointer to store the required workspace size in bytes. - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopGetSiluAndMulWorkspaceSize( - infiniopSiluAndMulDescriptor_t desc, - size_t *size); - -/** - * @brief Executes the SiluAndMul operation. - * - * Performs SiLU activation on the first half of the last dimension of `input`, - * multiplies element-wise with the second half, and stores the result in `output`. - * - * @param desc The SiluAndMul descriptor. - * @param workspace Pointer to workspace memory allocated according to GetWorkspaceSize(). - * @param workspace_size Size of the workspace in bytes. - * @param output Pointer to the output tensor memory. Shape [..., d]. - * @param input Pointer to the input tensor memory. Shape [..., 2*d]. - * @param stream Pointer to the execution stream (e.g., CUDA stream). Can be NULL for default stream. - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopSiluAndMul( - infiniopSiluAndMulDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -/** - * @brief Destroys a previously created SiluAndMul descriptor. - * @param desc The descriptor to destroy. - * @return infiniStatus_t Status code of the operation. - */ -__INFINI_C __export infiniStatus_t infiniopDestroySiluAndMulDescriptor( - infiniopSiluAndMulDescriptor_t desc); - -#endif // __INFINIOP_SILU_AND_MUL_API_H__ diff --git a/include/infiniop/ops/sinh.h b/include/infiniop/ops/sinh.h deleted file mode 100644 index 78d36b88b..000000000 --- a/include/infiniop/ops/sinh.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_SINH_API_H__ -#define __INFINIOP_SINH_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopSinhDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateSinhDescriptor(infiniopHandle_t handle, - infiniopSinhDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetSinhWorkspaceSize(infiniopSinhDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopSinh(infiniopSinhDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroySinhDescriptor(infiniopSinhDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/smooth_l1_loss.h b/include/infiniop/ops/smooth_l1_loss.h deleted file mode 100644 index 81678634f..000000000 --- a/include/infiniop/ops/smooth_l1_loss.h +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef __INFINIOP_SMOOTH_L1_LOSS_API_H__ -#define __INFINIOP_SMOOTH_L1_LOSS_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopSmoothL1LossDescriptor_t; -__INFINI_C __export infiniStatus_t infiniopCreateSmoothL1LossDescriptor(infiniopHandle_t handle, - infiniopSmoothL1LossDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - infiniopTensorDescriptor_t target, - float beta, - int reduction); - -__INFINI_C __export infiniStatus_t infiniopGetSmoothL1LossWorkspaceSize(infiniopSmoothL1LossDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopSmoothL1Loss(infiniopSmoothL1LossDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *target, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroySmoothL1LossDescriptor(infiniopSmoothL1LossDescriptor_t desc); - -#endif // __INFINIOP_SMOOTH_L1_LOSS_API_H__ diff --git a/include/infiniop/ops/softmax.h b/include/infiniop/ops/softmax.h deleted file mode 100644 index 8ab37b123..000000000 --- a/include/infiniop/ops/softmax.h +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef __INFINIOP_SOFTMAX_API_H__ -#define __INFINIOP_SOFTMAX_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopSoftmaxDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateSoftmaxDescriptor( - infiniopHandle_t handle, - infiniopSoftmaxDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int axis); - -__INFINI_C __export infiniStatus_t infiniopGetSoftmaxWorkspaceSize(infiniopSoftmaxDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopSoftmax( - infiniopSoftmaxDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroySoftmaxDescriptor(infiniopSoftmaxDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/softplus.h b/include/infiniop/ops/softplus.h deleted file mode 100644 index 7e0abc60b..000000000 --- a/include/infiniop/ops/softplus.h +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef __INFINIOP_OPS_SOFTPLUS_API_H__ -#define __INFINIOP_OPS_SOFTPLUS_API_H__ -#include "../tensor_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopSoftplusDescriptor_t; -__INFINI_C __export infiniStatus_t infiniopCreateSoftplusDescriptor( - infiniopHandle_t handle, - infiniopSoftplusDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - float beta, - float threshold); - -__INFINI_C __export infiniStatus_t infiniopGetSoftplusWorkspaceSize( - infiniopSoftplusDescriptor_t desc, - size_t *size); - -__INFINI_C __export infiniStatus_t infiniopSoftplus( - infiniopSoftplusDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroySoftplusDescriptor( - infiniopSoftplusDescriptor_t desc); - -#endif // __INFINIOP_OPS_SOFTPLUS_API_H__ diff --git a/include/infiniop/ops/softsign.h b/include/infiniop/ops/softsign.h deleted file mode 100644 index 62759dc54..000000000 --- a/include/infiniop/ops/softsign.h +++ /dev/null @@ -1,25 +0,0 @@ - -#ifndef __INFINIOP_SOFTSIGN_API_H__ -#define __INFINIOP_SOFTSIGN_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopSoftsignDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateSoftsignDescriptor(infiniopHandle_t handle, - infiniopSoftsignDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetSoftsignWorkspaceSize(infiniopSoftsignDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopSoftsign(infiniopSoftsignDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroySoftsignDescriptor(infiniopSoftsignDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/sub.h b/include/infiniop/ops/sub.h deleted file mode 100644 index 3a516c30c..000000000 --- a/include/infiniop/ops/sub.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_SUB_API_H__ -#define __INFINIOP_SUB_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopSubDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateSubDescriptor(infiniopHandle_t handle, - infiniopSubDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c, - infiniopTensorDescriptor_t a, - infiniopTensorDescriptor_t b); - -__INFINI_C __export infiniStatus_t infiniopGetSubWorkspaceSize(infiniopSubDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopSub(infiniopSubDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroySubDescriptor(infiniopSubDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/sum.h b/include/infiniop/ops/sum.h deleted file mode 100644 index c97104c90..000000000 --- a/include/infiniop/ops/sum.h +++ /dev/null @@ -1,31 +0,0 @@ -#ifndef __INFINIOP_SUM_API_H__ -#define __INFINIOP_SUM_API_H__ - -#include "../operator_descriptor.h" -#include -#include -typedef struct InfiniopDescriptor *infiniopSumDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateSumDescriptor(infiniopHandle_t handle, - infiniopSumDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - size_t *dim, - size_t dim_size, - bool keepdim); - -__INFINI_C __export infiniStatus_t infiniopGetSumWorkspaceSize(infiniopSumDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopSum(infiniopSumDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - size_t *dim, - size_t dim_size, - bool keepdim, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroySumDescriptor(infiniopSumDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/swap.h b/include/infiniop/ops/swap.h deleted file mode 100644 index 7eb14b57a..000000000 --- a/include/infiniop/ops/swap.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_SWAP_API_H__ -#define __INFINIOP_SWAP_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopSwapDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateSwapDescriptor(infiniopHandle_t handle, - infiniopSwapDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t y); - -__INFINI_C __export infiniStatus_t infiniopGetSwapWorkspaceSize(infiniopSwapDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopSwap(infiniopSwapDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *x, - void *y, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroySwapDescriptor(infiniopSwapDescriptor_t desc); - -#endif // __INFINIOP_SWAP_API_H__ diff --git a/include/infiniop/ops/swiglu.h b/include/infiniop/ops/swiglu.h deleted file mode 100644 index e11749c6b..000000000 --- a/include/infiniop/ops/swiglu.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_SWIGLU_API_H__ -#define __INFINIOP_SWIGLU_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopSwiGLUDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateSwiGLUDescriptor(infiniopHandle_t handle, - infiniopSwiGLUDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc); - -__INFINI_C __export infiniStatus_t infiniopGetSwiGLUWorkspaceSize(infiniopSwiGLUDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopSwiGLU(infiniopSwiGLUDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - void const *a, - void const *b, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroySwiGLUDescriptor(infiniopSwiGLUDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/take.h b/include/infiniop/ops/take.h deleted file mode 100644 index 58fb15311..000000000 --- a/include/infiniop/ops/take.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_TAKE_API_H__ -#define __INFINIOP_TAKE_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopTakeDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateTakeDescriptor(infiniopHandle_t handle, - infiniopTakeDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - infiniopTensorDescriptor_t indices); - -__INFINI_C __export infiniStatus_t infiniopGetTakeWorkspaceSize(infiniopTakeDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopTake(infiniopTakeDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *indices, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyTakeDescriptor(infiniopTakeDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/tan.h b/include/infiniop/ops/tan.h deleted file mode 100644 index 1b7bef23e..000000000 --- a/include/infiniop/ops/tan.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_TAN_API_H__ -#define __INFINIOP_TAN_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopTanDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateTanDescriptor( - infiniopHandle_t handle, - infiniopTanDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t out_desc); - -__INFINI_C __export infiniStatus_t infiniopGetTanWorkspaceSize(infiniopTanDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopTan( - infiniopTanDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyTanDescriptor(infiniopTanDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/tanh.h b/include/infiniop/ops/tanh.h deleted file mode 100644 index e623d854c..000000000 --- a/include/infiniop/ops/tanh.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_TANH_API_H__ -#define __INFINIOP_TANH_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopTanhDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateTanhDescriptor(infiniopHandle_t handle, - infiniopTanhDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input); - -__INFINI_C __export infiniStatus_t infiniopGetTanhWorkspaceSize(infiniopTanhDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopTanh(infiniopTanhDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyTanhDescriptor(infiniopTanhDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/tanhshrink.h b/include/infiniop/ops/tanhshrink.h deleted file mode 100644 index b2895ccab..000000000 --- a/include/infiniop/ops/tanhshrink.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_TANHSHRINK_API_H__ -#define __INFINIOP_TANHSHRINK_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopTanhshrinkDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateTanhshrinkDescriptor(infiniopHandle_t handle, - infiniopTanhshrinkDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t intput); - -__INFINI_C __export infiniStatus_t infiniopGetTanhshrinkWorkspaceSize(infiniopTanhshrinkDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopTanhshrink(infiniopTanhshrinkDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *intput, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyTanhshrinkDescriptor(infiniopTanhshrinkDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/topk.h b/include/infiniop/ops/topk.h deleted file mode 100644 index 3eaf94289..000000000 --- a/include/infiniop/ops/topk.h +++ /dev/null @@ -1,35 +0,0 @@ -#ifndef __INFINIOP_TOPK_API_H__ -#define __INFINIOP_TOPK_API_H__ - -#include "../operator_descriptor.h" -#include -#include -typedef struct InfiniopDescriptor *infiniopTopKDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateTopKDescriptor(infiniopHandle_t handle, - infiniopTopKDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t values_output_desc, - infiniopTensorDescriptor_t indices_output_desc, - infiniopTensorDescriptor_t input_desc, - size_t k, - size_t dim, - bool largest, - bool sorted); - -__INFINI_C __export infiniStatus_t infiniopGetTopKWorkspaceSize(infiniopTopKDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopTopK(infiniopTopKDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *values_output, - void *indices_output, - const void *input, - size_t k, - size_t dim, - bool largest, - bool sorted, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyTopKDescriptor(infiniopTopKDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/topkrouter.h b/include/infiniop/ops/topkrouter.h deleted file mode 100644 index 1d56a17cb..000000000 --- a/include/infiniop/ops/topkrouter.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_TOPKRouter_API_H__ -#define __INFINIOP_TOPKRouter_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopTopkrouterDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateTopkrouterDescriptor(infiniopHandle_t handle, - infiniopTopkrouterDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t correction_bias_desc); - -__INFINI_C __export infiniStatus_t infiniopGetTopkrouterWorkspaceSize(infiniopTopkrouterDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopTopkrouter(infiniopTopkrouterDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *values, - void *indices, - const void *x, - const void *correction_bias, - const float routed_scaling_factor, - const size_t topk, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyTopkrouterDescriptor(infiniopTopkrouterDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/topksoftmax.h b/include/infiniop/ops/topksoftmax.h deleted file mode 100644 index 74ea10d69..000000000 --- a/include/infiniop/ops/topksoftmax.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __INFINIOP_TOPKSOFTMAX_API_H__ -#define __INFINIOP_TOPKSOFTMAX_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopTopksoftmaxDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateTopksoftmaxDescriptor(infiniopHandle_t handle, - infiniopTopksoftmaxDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x_desc); - -__INFINI_C __export infiniStatus_t infiniopGetTopksoftmaxWorkspaceSize(infiniopTopksoftmaxDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopTopksoftmax(infiniopTopksoftmaxDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *values, - void *indices, - const void *x, - const size_t topk, - const int norm, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyTopksoftmaxDescriptor(infiniopTopksoftmaxDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/triplet_margin_loss.h b/include/infiniop/ops/triplet_margin_loss.h deleted file mode 100644 index a96dca22d..000000000 --- a/include/infiniop/ops/triplet_margin_loss.h +++ /dev/null @@ -1,33 +0,0 @@ -#ifndef __INFINIOP_TRIPLET_MARGIN_LOSS_API_H__ -#define __INFINIOP_TRIPLET_MARGIN_LOSS_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopTripletMarginLossDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateTripletMarginLossDescriptor(infiniopHandle_t handle, - infiniopTripletMarginLossDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t anchor, - infiniopTensorDescriptor_t positive, - infiniopTensorDescriptor_t negative, - float margin, - int p, - float eps, - int swap, // 0: False, 1: True - int reduction); // 0: None, 1: Mean, 2: Sum - -__INFINI_C __export infiniStatus_t infiniopGetTripletMarginLossWorkspaceSize(infiniopTripletMarginLossDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopTripletMarginLoss(infiniopTripletMarginLossDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *anchor, - const void *positive, - const void *negative, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyTripletMarginLossDescriptor(infiniopTripletMarginLossDescriptor_t desc); - -#endif // __INFINIOP_TRIPLET_MARGIN_LOSS_API_H__ diff --git a/include/infiniop/ops/triplet_margin_with_distance_loss.h b/include/infiniop/ops/triplet_margin_with_distance_loss.h deleted file mode 100644 index c834f32f5..000000000 --- a/include/infiniop/ops/triplet_margin_with_distance_loss.h +++ /dev/null @@ -1,32 +0,0 @@ -#ifndef __INFINIOP_TRIPLET_MARGIN_WITH_DISTANCE_LOSS_API_H__ -#define __INFINIOP_TRIPLET_MARGIN_WITH_DISTANCE_LOSS_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopTripletMarginWithDistanceLossDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateTripletMarginWithDistanceLossDescriptor( - infiniopHandle_t handle, - infiniopTripletMarginWithDistanceLossDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t anchor, - infiniopTensorDescriptor_t positive, - infiniopTensorDescriptor_t negative, - float margin, - int swap, - int reduction); -__INFINI_C __export infiniStatus_t infiniopGetTripletMarginWithDistanceLossWorkspaceSize( - infiniopTripletMarginWithDistanceLossDescriptor_t desc, - size_t *size); -__INFINI_C __export infiniStatus_t infiniopTripletMarginWithDistanceLoss(infiniopTripletMarginWithDistanceLossDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *anchor, - const void *positive, - const void *negative, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyTripletMarginWithDistanceLossDescriptor( - infiniopTripletMarginWithDistanceLossDescriptor_t desc); -#endif // __INFINIOP_TRIPLET_MARGIN_WITH_DISTANCE_LOSS_API_H__ diff --git a/include/infiniop/ops/unfold.h b/include/infiniop/ops/unfold.h deleted file mode 100644 index 0d1688978..000000000 --- a/include/infiniop/ops/unfold.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __INFINIOP_UNFOLD_API_H__ -#define __INFINIOP_UNFOLD_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopUnfoldDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateUnfoldDescriptor(infiniopHandle_t handle, - infiniopUnfoldDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - const int *kernel_sizes, - const int *strides, - const int *paddings, - const int *dilations); - -// 获取 Unfold 工作区大小 -__INFINI_C __export infiniStatus_t infiniopGetUnfoldWorkspaceSize(infiniopUnfoldDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopUnfold(infiniopUnfoldDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); -__INFINI_C __export infiniStatus_t infiniopDestroyUnfoldDescriptor(infiniopUnfoldDescriptor_t desc); - -#endif // __INFINIOP_UNFOLD_API_H__ diff --git a/include/infiniop/ops/upsample_bilinear.h b/include/infiniop/ops/upsample_bilinear.h deleted file mode 100644 index c53a5b24a..000000000 --- a/include/infiniop/ops/upsample_bilinear.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __INFINIOP_UPSAMPLE_BILINEAR_API_H__ -#define __INFINIOP_UPSAMPLE_BILINEAR_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopUpsampleBilinearDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateUpsampleBilinearDescriptor(infiniopHandle_t handle, - infiniopUpsampleBilinearDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - int align_corners); - -__INFINI_C __export infiniStatus_t infiniopGetUpsampleBilinearWorkspaceSize(infiniopUpsampleBilinearDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopUpsampleBilinear(infiniopUpsampleBilinearDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyUpsampleBilinearDescriptor(infiniopUpsampleBilinearDescriptor_t desc); - -#endif // __INFINIOP_UPSAMPLE_BILINEAR_API_H__ diff --git a/include/infiniop/ops/upsample_nearest.h b/include/infiniop/ops/upsample_nearest.h deleted file mode 100644 index de64badcc..000000000 --- a/include/infiniop/ops/upsample_nearest.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_UPSAMPLE_NEAREST_API_H__ -#define __INFINIOP_UPSAMPLE_NEAREST_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopUpsampleNearestDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateUpsampleNearestDescriptor(infiniopHandle_t handle, - infiniopUpsampleNearestDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input); - -__INFINI_C __export infiniStatus_t infiniopGetUpsampleNearestWorkspaceSize(infiniopUpsampleNearestDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopUpsampleNearest(infiniopUpsampleNearestDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyUpsampleNearestDescriptor(infiniopUpsampleNearestDescriptor_t desc); - -#endif // __INFINIOP_UPSAMPLE_NEAREST_API_H__ diff --git a/include/infiniop/ops/vander.h b/include/infiniop/ops/vander.h deleted file mode 100644 index 72773da5b..000000000 --- a/include/infiniop/ops/vander.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __INFINIOP_VANDER_API_H__ -#define __INFINIOP_VANDER_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopVanderDescriptor_t; -__INFINI_C __export infiniStatus_t infiniopCreateVanderDescriptor(infiniopHandle_t handle, - infiniopVanderDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - int N, - int increasing); - -__INFINI_C __export infiniStatus_t infiniopGetVanderWorkspaceSize(infiniopVanderDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopVander(infiniopVanderDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyVanderDescriptor(infiniopVanderDescriptor_t desc); - -#endif // __INFINIOP_VANDER_API_H__ diff --git a/include/infiniop/ops/var.h b/include/infiniop/ops/var.h deleted file mode 100644 index 7dc601a94..000000000 --- a/include/infiniop/ops/var.h +++ /dev/null @@ -1,33 +0,0 @@ -#ifndef __INFINIOP_VAR_API_H__ -#define __INFINIOP_VAR_API_H__ - -#include "../operator_descriptor.h" -#include -#include -typedef struct InfiniopDescriptor *infiniopVarDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateVarDescriptor(infiniopHandle_t handle, - infiniopVarDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t var_output_desc, - infiniopTensorDescriptor_t input_desc, - size_t *dim, - size_t dim_size, - bool unbiased, - bool keepdim); - -__INFINI_C __export infiniStatus_t infiniopGetVarWorkspaceSize(infiniopVarDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopVar(infiniopVarDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *var_output, - const void *input, - size_t *dim, - size_t dim_size, - bool unbiased, - bool keepdim, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyVarDescriptor(infiniopVarDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/var_mean.h b/include/infiniop/ops/var_mean.h deleted file mode 100644 index 358a55636..000000000 --- a/include/infiniop/ops/var_mean.h +++ /dev/null @@ -1,35 +0,0 @@ -#ifndef __INFINIOP_VAR_MEAN_API_H__ -#define __INFINIOP_VAR_MEAN_API_H__ - -#include "../operator_descriptor.h" -#include -#include -typedef struct InfiniopDescriptor *infiniopVarMeanDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateVarMeanDescriptor(infiniopHandle_t handle, - infiniopVarMeanDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t var_output_desc, - infiniopTensorDescriptor_t mean_output_desc, - infiniopTensorDescriptor_t input_desc, - size_t *dim, - size_t dim_size, - bool unbiased, - bool keepdim); - -__INFINI_C __export infiniStatus_t infiniopGetVarMeanWorkspaceSize(infiniopVarMeanDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopVarMean(infiniopVarMeanDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *var_output, - void *mean_output, - const void *input, - size_t *dim, - size_t dim_size, - bool unbiased, - bool keepdim, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyVarMeanDescriptor(infiniopVarMeanDescriptor_t desc); - -#endif diff --git a/include/infiniop/ops/zeros.h b/include/infiniop/ops/zeros.h deleted file mode 100644 index 5d7ce03d5..000000000 --- a/include/infiniop/ops/zeros.h +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef __INFINIOP_ZEROS_API_H__ -#define __INFINIOP_ZEROS_API_H__ - -#include "../operator_descriptor.h" - -typedef struct InfiniopDescriptor *infiniopZerosDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateZerosDescriptor(infiniopHandle_t handle, - infiniopZerosDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x); - -__INFINI_C __export infiniStatus_t infiniopGetZerosWorkspaceSize(infiniopZerosDescriptor_t desc, size_t *size); - -__INFINI_C __export infiniStatus_t infiniopZeros(infiniopZerosDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream); - -__INFINI_C __export infiniStatus_t infiniopDestroyZerosDescriptor(infiniopZerosDescriptor_t desc); - -#endif diff --git a/include/infiniop/tensor_descriptor.h b/include/infiniop/tensor_descriptor.h deleted file mode 100644 index 8390da81c..000000000 --- a/include/infiniop/tensor_descriptor.h +++ /dev/null @@ -1,14 +0,0 @@ -#ifndef __INFINIOP_TENSOR_DESCRIPTOR_API_H__ -#define __INFINIOP_TENSOR_DESCRIPTOR_API_H__ - -#include "../infinicore.h" - -struct InfiniopTensorDescriptor; - -typedef struct InfiniopTensorDescriptor *infiniopTensorDescriptor_t; - -__INFINI_C __export infiniStatus_t infiniopCreateTensorDescriptor(infiniopTensorDescriptor_t *desc_ptr, size_t ndim, const size_t *shape, const ptrdiff_t *strides, infiniDtype_t dtype); - -__INFINI_C __export infiniStatus_t infiniopDestroyTensorDescriptor(infiniopTensorDescriptor_t desc); - -#endif // __INFINIOP_TENSOR_DESCRIPTOR__ diff --git a/include/infinirt.h b/include/infinirt.h deleted file mode 100644 index 4ada848f5..000000000 --- a/include/infinirt.h +++ /dev/null @@ -1,134 +0,0 @@ -#ifndef __INFINIRT_API_H__ -#define __INFINIRT_API_H__ - -#include "infinicore.h" -#include - -typedef void *infinirtStream_t; -typedef void *infinirtEvent_t; -typedef void *infinirtGraph_t; -typedef void *infinirtGraphNode_t; -typedef void *infinirtGraphExec_t; - -// Bitmask describing which fields of an infinirtDeviceResourceSnapshot_t -// have been populated by the backend. Backends without a particular -// capability simply leave the flag clear. -typedef enum { - INFINIRT_RESOURCE_FIELD_NONE = 0, - INFINIRT_RESOURCE_FIELD_MEMORY_CAPACITY = 1ull << 0, - INFINIRT_RESOURCE_FIELD_COMPUTE_UTILIZATION = 1ull << 1, - INFINIRT_RESOURCE_FIELD_MEMORY_BANDWIDTH_UTILIZATION = 1ull << 2, - INFINIRT_RESOURCE_FIELD_KERNEL_TIME_RATIO = 1ull << 3, - INFINIRT_RESOURCE_FIELD_COMMUNICATION = 1ull << 4, -} infinirtResourceField_t; - -// Vendor-neutral resource snapshot consumed by the optional mutual -// awareness analyzer. Backends populate the fields they can observe and -// set the matching `valid_fields` bits; consumers must check the bits -// before relying on a value. -typedef struct { - infiniDevice_t device_type; - int device_id; - - uint64_t valid_fields; - uint64_t estimated_fields; - - size_t free_bytes; - size_t total_bytes; - size_t used_bytes; - size_t reserved_bytes; - - float compute_utilization; - float memory_bandwidth_utilization; - float kernel_time_ratio; - - float communication_time_ratio; - uint64_t communication_bytes; -} infinirtDeviceResourceSnapshot_t; - -__INFINI_C __export infiniStatus_t infinirtInit(); - -// Device -__INFINI_C __export infiniStatus_t infinirtGetAllDeviceCount(int *count_array); -__INFINI_C __export infiniStatus_t infinirtGetDeviceCount(infiniDevice_t device, int *count); -__INFINI_C __export infiniStatus_t infinirtSetDevice(infiniDevice_t device, int device_id); -__INFINI_C __export infiniStatus_t infinirtGetDevice(infiniDevice_t *device_ptr, int *device_id_ptr); -__INFINI_C __export infiniStatus_t infinirtGetMemInfo(infiniDevice_t device, int device_id, size_t *free_bytes, size_t *total_bytes); -__INFINI_C __export infiniStatus_t infinirtGetDeviceResourceSnapshot( - infiniDevice_t device, - int device_id, - infinirtDeviceResourceSnapshot_t *snapshot); -__INFINI_C __export infiniStatus_t infinirtDeviceSynchronize(); - -// Stream -__INFINI_C __export infiniStatus_t infinirtStreamCreate(infinirtStream_t *stream_ptr); -__INFINI_C __export infiniStatus_t infinirtStreamDestroy(infinirtStream_t stream); -__INFINI_C __export infiniStatus_t infinirtStreamSynchronize(infinirtStream_t stream); -__INFINI_C __export infiniStatus_t infinirtStreamWaitEvent(infinirtStream_t stream, infinirtEvent_t event); - -// Event -typedef enum { - INFINIRT_EVENT_COMPLETE = 0, - INFINIRT_EVENT_NOT_READY = 1, -} infinirtEventStatus_t; - -// Event flags for precise timing -typedef enum { - INFINIRT_EVENT_DEFAULT = 0x0, // Default event creation flags - INFINIRT_EVENT_DISABLE_TIMING = 0x1, // Event will not record timing data - INFINIRT_EVENT_BLOCKING_SYNC = 0x2, // Event uses blocking synchronization -} infinirtEventFlags_t; - -__INFINI_C __export infiniStatus_t infinirtEventCreate(infinirtEvent_t *event_ptr); -__INFINI_C __export infiniStatus_t infinirtEventCreateWithFlags(infinirtEvent_t *event_ptr, uint32_t flags); -__INFINI_C __export infiniStatus_t infinirtEventRecord(infinirtEvent_t event, infinirtStream_t stream); -__INFINI_C __export infiniStatus_t infinirtEventQuery(infinirtEvent_t event, infinirtEventStatus_t *status_ptr); -__INFINI_C __export infiniStatus_t infinirtEventSynchronize(infinirtEvent_t event); -__INFINI_C __export infiniStatus_t infinirtEventDestroy(infinirtEvent_t event); -__INFINI_C __export infiniStatus_t infinirtEventElapsedTime(float *ms_ptr, infinirtEvent_t start, infinirtEvent_t end); - -// Memory -typedef enum { - INFINIRT_MEMCPY_H2H = 0, - INFINIRT_MEMCPY_H2D = 1, - INFINIRT_MEMCPY_D2H = 2, - INFINIRT_MEMCPY_D2D = 3, -} infinirtMemcpyKind_t; - -__INFINI_C __export infiniStatus_t infinirtMalloc(void **p_ptr, size_t size); -__INFINI_C __export infiniStatus_t infinirtMallocHost(void **p_ptr, size_t size); -__INFINI_C __export infiniStatus_t infinirtFree(void *ptr); -__INFINI_C __export infiniStatus_t infinirtFreeHost(void *ptr); - -__INFINI_C __export infiniStatus_t infinirtMemcpy(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind); -__INFINI_C __export infiniStatus_t infinirtMemcpyAsync(void *dst, const void *src, size_t size, infinirtMemcpyKind_t kind, infinirtStream_t stream); - -// Stream-ordered memory -__INFINI_C __export infiniStatus_t infinirtMallocAsync(void **p_ptr, size_t size, infinirtStream_t stream); -__INFINI_C __export infiniStatus_t infinirtFreeAsync(void *ptr, infinirtStream_t stream); - -// Memset -__INFINI_C __export infiniStatus_t infinirtMemset(void *ptr, int value, size_t count); -__INFINI_C __export infiniStatus_t infinirtMemsetAsync(void *ptr, int value, size_t count, infinirtStream_t stream); - -// Graph -typedef enum { - INFINIRT_STREAM_CAPTURE_MODE_GLOBAL = 0, - INFINIRT_STREAM_CAPTURE_MODE_THREAD_LOCAL = 1, - INFINIRT_STREAM_CAPTURE_MODE_RELAXED = 2, - -} infinirtStreamCaptureMode_t; - -__INFINI_C __export infiniStatus_t infinirtStreamBeginCapture(infinirtStream_t stream, infinirtStreamCaptureMode_t mode); -__INFINI_C __export infiniStatus_t infinirtStreamEndCapture(infinirtStream_t stream, infinirtGraph_t *graph_ptr); -__INFINI_C __export infiniStatus_t infinirtGraphDestroy(infinirtGraph_t graph); -__INFINI_C __export infiniStatus_t infinirtGraphInstantiate( - infinirtGraphExec_t *graph_exec_ptr, - infinirtGraph_t graph, - infinirtGraphNode_t *node_ptr, - char *log_buffer, - size_t buffer_size); -__INFINI_C __export infiniStatus_t infinirtGraphExecDestroy(infinirtGraphExec_t graph_exec); -__INFINI_C __export infiniStatus_t infinirtGraphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stream); - -#endif // __INFINIRT_API_H__ diff --git a/pyproject.toml b/pyproject.toml deleted file mode 100644 index 4e422c568..000000000 --- a/pyproject.toml +++ /dev/null @@ -1,26 +0,0 @@ -[build-system] -requires = ["setuptools"] -build-backend = "setuptools.build_meta" - -[project] -name = "InfiniCore" -version = "0.1.0" -description = "InfiniCore 是一个跨平台统一编程工具集，为不同芯片平台的功能（包括计算、运行时、通信等）提供统一 C 语言接口。" -readme = "README.md" -dependencies = [] -requires-python = ">=3.8" -classifiers = [ - "Programming Language :: Python :: 3", - "License :: OSI Approved :: MIT License", - "Operating System :: OS Independent", -] - -[project.urls] -Homepage = "https://github.com/InfiniTensor/InfiniCore" -Issues = "https://github.com/InfiniTensor/InfiniCore/issues" - -[tool.ruff] -src = [".", "src"] - -[tool.ruff.lint] -select = ["E4", "E7", "E9", "F", "I"] diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py deleted file mode 100644 index 808995640..000000000 --- a/python/infinicore/__init__.py +++ /dev/null @@ -1,312 +0,0 @@ -import contextlib - -with contextlib.suppress(ImportError): - from ._preload import preload - - preload() - -import infinicore.context as context -import infinicore.nn as nn -from infinicore._tensor_str import printoptions, set_printoptions - -# Import context functions -from infinicore.context import ( - get_device, - get_device_count, - get_stream, - is_graph_recording, - set_device, - start_graph_recording, - stop_graph_recording, - sync_device, - sync_stream, -) -from infinicore.device import device -from infinicore.device_event import DeviceEvent -from infinicore.dtype import ( - bfloat16, - bool, - cdouble, - cfloat, - chalf, - complex32, - complex64, - complex128, - double, - dtype, - float, - float16, - float32, - float64, - half, - int, - int8, - int16, - int32, - int64, - long, - short, - uint8, -) -from infinicore.ops.acos import acos -from infinicore.ops.add import add -from infinicore.ops.add_rms_norm import add_rms_norm -from infinicore.ops.addbmm import addbmm -from infinicore.ops.addcmul import addcmul -from infinicore.ops.addr import addr -from infinicore.ops.all import all -from infinicore.ops.argwhere import argwhere -from infinicore.ops.asin import asin -from infinicore.ops.asinh import asinh -from infinicore.ops.asum import asum -from infinicore.ops.atanh import atanh -from infinicore.ops.attention import attention -from infinicore.ops.axpy import axpy -from infinicore.ops.baddbmm import baddbmm -from infinicore.ops.bilinear import bilinear -from infinicore.ops.binary_cross_entropy_with_logits import ( - binary_cross_entropy_with_logits, -) -from infinicore.ops.bitwise_right_shift import bitwise_right_shift -from infinicore.ops.blas_amax import blas_amax -from infinicore.ops.blas_amin import blas_amin -from infinicore.ops.blas_copy import blas_copy -from infinicore.ops.blas_dot import blas_dot -from infinicore.ops.block_diag import block_diag -from infinicore.ops.broadcast_to import broadcast_to -from infinicore.ops.cat import cat -from infinicore.ops.cdist import cdist -from infinicore.ops.conv2d import conv2d -from infinicore.ops.cross_entropy import cross_entropy -from infinicore.ops.diff import diff -from infinicore.ops.digamma import digamma -from infinicore.ops.dist import dist -from infinicore.ops.equal import equal -from infinicore.ops.flipud import flipud -from infinicore.ops.float_power import float_power -from infinicore.ops.floor import floor -from infinicore.ops.floor_divide import floor_divide -from infinicore.ops.fmin import fmin -from infinicore.ops.fmod import fmod -from infinicore.ops.hypot import hypot -from infinicore.ops.index_add import index_add -from infinicore.ops.index_copy import index_copy -from infinicore.ops.inner import inner -from infinicore.ops.kron import kron -from infinicore.ops.kthvalue import kthvalue -from infinicore.ops.kv_caching import kv_caching -from infinicore.ops.ldexp import ldexp -from infinicore.ops.lerp import lerp -from infinicore.ops.logaddexp import logaddexp -from infinicore.ops.logaddexp2 import logaddexp2 -from infinicore.ops.logcumsumexp import logcumsumexp -from infinicore.ops.logdet import logdet -from infinicore.ops.logical_and import logical_and -from infinicore.ops.logical_not import logical_not -from infinicore.ops.masked_select import masked_select -from infinicore.ops.matmul import matmul -from infinicore.ops.mha import mha -from infinicore.ops.mha_kvcache import mha_kvcache -from infinicore.ops.mha_varlen import mha_varlen -from infinicore.ops.moore_mate_flash_attn import ( - moore_mate_flash_attn_decode, - moore_mate_flash_attn_prefill, -) -from infinicore.ops.mrope import mrope -from infinicore.ops.mul import mul -from infinicore.ops.mul_scalar import mul_scalar -from infinicore.ops.narrow import narrow -from infinicore.ops.nrm2 import nrm2 -from infinicore.ops.paged_attention import paged_attention -from infinicore.ops.paged_attention_prefill import paged_attention_prefill -from infinicore.ops.paged_caching import paged_caching -from infinicore.ops.rearrange import rearrange -from infinicore.ops.reciprocal import reciprocal -from infinicore.ops.rot import rot -from infinicore.ops.rotg import rotg -from infinicore.ops.rotm import rotm -from infinicore.ops.rotmg import rotmg -from infinicore.ops.scal import scal -from infinicore.ops.scatter import scatter -from infinicore.ops.sinh import sinh -from infinicore.ops.squeeze import squeeze -from infinicore.ops.sum import sum -from infinicore.ops.swap import swap -from infinicore.ops.take import take -from infinicore.ops.tan import tan -from infinicore.ops.topk import topk -from infinicore.ops.unsqueeze import unsqueeze -from infinicore.ops.vander import vander -from infinicore.ops.var import var -from infinicore.ops.var_mean import var_mean -from infinicore.tensor import ( - Tensor, - empty, - empty_like, - from_blob, - from_list, - from_list_by_numpy, - from_numpy, - from_torch, - ones, - strided_empty, - strided_from_blob, - zeros, -) - -__all__ = [ - # Modules. - "context", - "nn", - # Classes. - "device", - "DeviceEvent", - "dtype", - "Tensor", - # Context functions. - "get_device", - "get_device_count", - "get_stream", - "set_device", - "sync_device", - "sync_stream", - "is_graph_recording", - "start_graph_recording", - "stop_graph_recording", - # Data Types. - "bfloat16", - "bool", - "cdouble", - "cfloat", - "chalf", - "complex32", - "complex64", - "complex128", - "double", - "float", - "float16", - "float32", - "float64", - "half", - "int", - "int8", - "int16", - "int32", - "int64", - "long", - "short", - "uint8", - # Operations. - "addcmul", - "atanh", - "binary_cross_entropy_with_logits", - "cdist", - "reciprocal", - "add", - "addr", - "add_rms_norm", - "argwhere", - "asin", - "asum", - "axpy", - "blas_amax", - "blas_amin", - "blas_copy", - "blas_dot", - "acos", - "addbmm", - "floor", - "attention", - "mrope", - "block_diag", - "kron", - "bitwise_right_shift", - "kv_caching", - "asinh", - "baddbmm", - "bilinear", - "fmod", - "cat", - "conv2d", - "inner", - "masked_select", - "logaddexp", - "logaddexp2", - "matmul", - "equal", - "mul", - "mul_scalar", - "diff", - "digamma", - "dist", - "logdet", - "narrow", - "nrm2", - "ldexp", - "lerp", - "kthvalue", - "squeeze", - "unsqueeze", - "rearrange", - "cross_entropy", - "tan", - "empty", - "empty_like", - "from_blob", - "from_list", - "from_list_by_numpy", - "from_numpy", - "from_torch", - "mha_kvcache", - "mha_varlen", - "mha", - "fmin", - "floor_divide", - "float_power", - "flipud", - "scatter", - "rot", - "rotg", - "rotm", - "rotmg", - "scal", - "logcumsumexp", - "logical_not", - "logical_and", - "vander", - "paged_caching", - "paged_attention", - "paged_attention_prefill", - "hypot", - "index_copy", - "index_add", - "take", - "sinh", - "swap", - "ones", - "broadcast_to", - "strided_empty", - "strided_from_blob", - "zeros", - "sum", - "var_mean", - "moore_mate_flash_attn_prefill", - "moore_mate_flash_attn_decode", - "var", - "topk", - "all", - "set_printoptions", - "printoptions", -] - -use_ntops = False - -with contextlib.suppress(ImportError, ModuleNotFoundError): - import sys - - import ntops - - for op_name in ntops.torch.__all__: - getattr(ntops.torch, op_name).__globals__["torch"] = sys.modules[__name__] - - use_ntops = True diff --git a/python/infinicore/_preload.py b/python/infinicore/_preload.py deleted file mode 100644 index 69121ff2f..000000000 --- a/python/infinicore/_preload.py +++ /dev/null @@ -1,224 +0,0 @@ -import ctypes -import glob -import importlib -import importlib.util -import os -import sys -from typing import Iterable, List - - -_PRELOADED_HANDLES: List[ctypes.CDLL] = [] - - -def _candidate_prefixes(path: str) -> List[str]: - """ - Return HPCC install prefixes to search for libs. - Prefer HPCC_PATH; if absent and explicitly opted-in, fall back to /opt/hpcc. - """ - prefixes: List[str] = [] - if path: - prefixes.append(path) - - seen = set() - unique: List[str] = [] - for p in prefixes: - if p and p not in seen: - seen.add(p) - unique.append(p) - return unique - - -def _try_load(paths: Iterable[str], name: str) -> bool: - """Try to load a shared library from given paths or system search path.""" - for path in paths: - full = os.path.join(path, "lib", name) - if os.path.exists(full): - try: - ctypes.CDLL(full, mode=ctypes.RTLD_GLOBAL) - return True - except OSError: - # Try next candidate - continue - # Last resort: rely on loader search path - try: - ctypes.CDLL(name, mode=ctypes.RTLD_GLOBAL) - return True - except OSError: - return False - - -def preload_hpcc() -> None: - """ - Best-effort preload of key HPCC runtime libs with RTLD_GLOBAL. - - This mirrors the behavior of torch's HPCC build that loads libtorch_global_deps.so, - but avoids introducing a hard torch dependency. All failures are swallowed. - """ - hpcc_path = os.getenv("HPCC_PATH") - if not hpcc_path: - return - - prefixes = _candidate_prefixes(hpcc_path) - libs = [ - "libhcruntime.so", - "libhcToolsExt.so", - "libruntime_cu.so", - "libhccompiler.so", - ] - - for lib in libs: - _try_load(prefixes, lib) - - -def preload_torch_hip() -> None: - """ - Best-effort import of the torch HIP runtime. - - Loading torch's shared libraries individually with ctypes bypasses - Python's extension initialization order and can corrupt HIP teardown. - """ - try: - importlib.import_module("torch") - except (ImportError, OSError): - pass - - -def _import_extension_global(module_name: str, path: str) -> None: - module = sys.modules.get(module_name) - if module is None: - spec = importlib.util.spec_from_file_location(module_name, path) - if spec is None or spec.loader is None: - raise ImportError(f"Cannot load extension module from {path}") - module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - try: - spec.loader.exec_module(module) - except Exception: - sys.modules.pop(module_name, None) - raise - - module_path = getattr(module, "__file__", None) or path - _PRELOADED_HANDLES.append( - ctypes.CDLL(module_path, mode=ctypes.RTLD_GLOBAL) - ) - - -def preload_flash_attn() -> None: - """ - Best-effort preload of flash_attn_2_cuda extension with RTLD_GLOBAL. - - InfiniCore hygon wrapper resolves C symbols like `mha_varlen_fwd` from the - flash-attn extension at runtime via dlsym(RTLD_DEFAULT, ...). The symbols - only need to be available when the operator is actually called, not at - library load time. So this preload is a convenience — if it fails, the - symbols will be resolved later when torch + flash_attn are imported by - the application (e.g. InfiniLM). - """ - candidates: List[str] = [] - from_env = os.getenv("FLASH_ATTN_PREBUILT") - if from_env: - if os.path.isfile(from_env): - candidates.append(from_env) - elif os.path.isdir(from_env): - candidates.extend( - glob.glob(os.path.join(from_env, "flash_attn_2_cuda*.so")) - ) - - # Try resolving via Python import metadata. - spec = importlib.util.find_spec("flash_attn_2_cuda") - if spec and spec.origin and os.path.exists(spec.origin): - candidates.append(spec.origin) - - # Fallback: scan python paths for extension module. - for p in sys.path: - if not p: - continue - candidates.extend(glob.glob(os.path.join(p, "flash_attn_2_cuda*.so"))) - - # Common installation locations. - candidates.extend( - glob.glob("/usr/local/lib/python*/dist-packages/flash_attn_2_cuda*.so") - ) - candidates.extend(glob.glob("/root/.infini/lib/flash_attn_2_cuda*.so")) - - seen = set() - for so_path in candidates: - if not so_path or so_path in seen: - continue - seen.add(so_path) - if not os.path.exists(so_path): - continue - try: - _import_extension_global("flash_attn_2_cuda", so_path) - return - except (ImportError, OSError): - continue - - -def _should_preload_device(device_type: str) -> bool: - """ - Check if preload is needed for a specific device type. - """ - device_env_map = { - "METAX": ["HPCC_PATH", "INFINICORE_PRELOAD_HPCC"], # HPCC/METAX - "HYGON": ["DTK_ROOT", "INFINICORE_PRELOAD_TORCH_HIP"], - # Add other device types here as needed: - # "ASCEND": ["ASCEND_PATH"], - # "CAMBRICON": ["NEUWARE_HOME"], - } - - env_vars = device_env_map.get(device_type, []) - for env_var in env_vars: - if os.getenv(env_var): - return True - if device_type == "HYGON": - dtk_root = os.getenv("DTK_ROOT") or "/opt/dtk" - if os.path.isdir(dtk_root): - return True - return False - - -def preload_device(device_type: str) -> None: - """ - Preload runtime libraries for a specific device type if needed. - - Args: - device_type: Device type name (e.g., "METAX", "ASCEND", etc.) - """ - if device_type == "METAX": - preload_hpcc() - elif device_type == "HYGON": - preload_torch_hip() - preload_flash_attn() - # Add other device preload functions here as needed: - # elif device_type == "ASCEND": - # preload_ascend() - # etc. - - -def preload() -> None: - """ - Universal preload function that loops through device types and preloads when required. - - This function detects available device types and preloads their runtime libraries - if the environment indicates they are needed. - """ - # Device types that may require preload. Keep Hygon-only preloads gated by - # Hygon environment markers so other CUDA-compatible platforms do not load - # unrelated torch/flash-attn libraries during package import. - device_types = [ - "METAX", # HPCC/METAX - "HYGON", - # Add other device types here as they are implemented: - # "ASCEND", - # "CAMBRICON", - # etc. - ] - - for device_type in device_types: - if _should_preload_device(device_type): - try: - preload_device(device_type) - except Exception: - # Swallow all errors - preload is best-effort - pass diff --git a/python/infinicore/_tensor_str.py b/python/infinicore/_tensor_str.py deleted file mode 100644 index 010964d91..000000000 --- a/python/infinicore/_tensor_str.py +++ /dev/null @@ -1,122 +0,0 @@ -import contextlib -import dataclasses -from typing import Any, Optional - -from infinicore.lib import _infinicore - - -@dataclasses.dataclass -class __PrinterOptions: - precision: int = 4 - threshold: float = 1000 - edgeitems: int = 3 - linewidth: int = 80 - sci_mode: Optional[bool] = None - - -PRINT_OPTS = __PrinterOptions() - - -def set_printoptions( - precision=None, - threshold=None, - edgeitems=None, - linewidth=None, - profile=None, - sci_mode=None, -): - r"""Set options for printing. - Args: - precision: Number of digits of precision for floating point output (default = 4). - threshold: Total number of array elements which trigger summarization rather than full `repr` (default = 1000). - edgeitems: Number of array items in summary at beginning and end of each dimension (default = 3). - linewidth: The number of characters per line (default = 80). - profile: Sane defaults for pretty printing. Can override with any of the above options. (any one of `default`, `short`, `full`) - sci_mode: Enable (True) or disable (False) scientific notation. - If None (default) is specified, the value is automatically chosen by the framework. - Example:: - >>> # Limit the precision of elements - >>> torch.set_printoptions(precision=2) - >>> torch.tensor([1.12345]) - tensor([1.12]) - """ - if profile is not None: - if profile == "default": - PRINT_OPTS.precision = 4 - PRINT_OPTS.threshold = 1000 - PRINT_OPTS.edgeitems = 3 - PRINT_OPTS.linewidth = 80 - elif profile == "short": - PRINT_OPTS.precision = 2 - PRINT_OPTS.threshold = 1000 - PRINT_OPTS.edgeitems = 2 - PRINT_OPTS.linewidth = 80 - elif profile == "full": - PRINT_OPTS.precision = 4 - PRINT_OPTS.threshold = 2147483647 # CPP_INT32_MAX - PRINT_OPTS.edgeitems = 3 - PRINT_OPTS.linewidth = 80 - else: - raise ValueError( - f"Invalid profile: {profile}. the profile must be one of 'default', 'short', 'full'" - ) - - if precision is not None: - PRINT_OPTS.precision = precision - if threshold is not None: - PRINT_OPTS.threshold = threshold - if edgeitems is not None: - PRINT_OPTS.edgeitems = edgeitems - if linewidth is not None: - PRINT_OPTS.linewidth = linewidth - PRINT_OPTS.sci_mode = sci_mode - - _infinicore.set_printoptions( - PRINT_OPTS.precision, - PRINT_OPTS.threshold, - PRINT_OPTS.edgeitems, - PRINT_OPTS.linewidth, - PRINT_OPTS.sci_mode, - ) - - -def get_printoptions() -> dict[str, Any]: - r"""Gets the current options for printing, as a dictionary that - can be passed as ``**kwargs`` to set_printoptions(). - """ - return dataclasses.asdict(PRINT_OPTS) - - -@contextlib.contextmanager -def printoptions( - precision=None, threshold=None, edgeitems=None, linewidth=None, sci_mode=None -): - r"""Context manager that temporarily changes the print options.""" - old_kwargs = get_printoptions() - set_printoptions( - precision=precision, - threshold=threshold, - edgeitems=edgeitems, - linewidth=linewidth, - sci_mode=sci_mode, - ) - try: - yield - finally: - set_printoptions(**old_kwargs) - - -def _str(self): - cpp_tensor_str = self._underlying.__str__() - py_dtype_str = ", dtype=" + self.dtype.__repr__() - - py_tensor_str = cpp_tensor_str.split(", d")[0] - if self.device.type != "cpu": - py_device_str = ", device='" + self.device.__str__() + "'" - py_tensor_str += py_device_str - py_tensor_str += py_dtype_str + ")\n" - - return py_tensor_str - - -set_printoptions() diff --git a/python/infinicore/context.py b/python/infinicore/context.py deleted file mode 100644 index 1181738c5..000000000 --- a/python/infinicore/context.py +++ /dev/null @@ -1,73 +0,0 @@ -import infinicore.device -from infinicore.graph import Graph -from infinicore.lib import _infinicore - - -def get_device(): - """Get the current active device. - - Returns: - device: The current active device object - """ - return _infinicore.get_device() - - -def get_device_count(device_type): - """Get the number of available devices of a specific type. - - Args: - device_type (str): The type of device to count (e.g., "cuda", "cpu", "npu") - - Returns: - int: The number of available devices of the specified type - """ - return _infinicore.get_device_count(infinicore.device(device_type)._underlying.type) - - -def set_device(device): - """Set the current active device. - - Args: - device: The device to set as active - """ - _infinicore.set_device(device._underlying) - - -def sync_stream(): - """Synchronize the current stream.""" - _infinicore.sync_stream() - - -def sync_device(): - """Synchronize the current device.""" - _infinicore.sync_device() - - -def get_stream(): - """Get the current stream. - - Returns: - stream: The current stream object - """ - return _infinicore.get_stream() - - -def is_graph_recording(): - """Check if the current graph is recording. - - Returns: - bool: True if the current graph is recording, False otherwise - """ - return _infinicore.is_graph_recording() - - -def start_graph_recording(device=None): - """Start recording the current graph.""" - if device is not None: - set_device(device) - _infinicore.start_graph_recording() - - -def stop_graph_recording(): - """Stop recording the current graph.""" - return Graph(_infinicore.stop_graph_recording()) diff --git a/python/infinicore/device.py b/python/infinicore/device.py deleted file mode 100644 index 8858a1235..000000000 --- a/python/infinicore/device.py +++ /dev/null @@ -1,138 +0,0 @@ -from infinicore.lib import _infinicore - -_infinicore_2_python_dict = {} -_python_2_infinicore_dict = {} - - -class device: - # Public attributes describing the device - type: str - index: int - _underlying: _infinicore.Device - - def __init__(self, type=None, index=None): - if isinstance(type, device): - self.type = type.type - self.index = type.index - return - - if type is None: - type = "cpu" - - if ":" in type: - if index is not None: - raise ValueError( - '`index` should not be provided when `type` contains `":"`.' - ) - - type, index = type.split(":") - index = int(index) - - self.type = type - self.index = index if index else 0 - - def __getattr__(self, name): - # Lazily construct and cache an attribute. - # such as, self._underlying . - if name == "_underlying": - setattr(self, name, device._to_infinicore_device(self.type, self.index)) - else: - raise AttributeError("{!r} object has no attribute {!r}".format(self, name)) - return getattr(self, name) - - def __repr__(self): - return f"device(type='{self.type}'{f', index={self.index}' if self.index is not None else ''})" - - def __str__(self): - return f"{self.type}{f':{self.index}' if self.index is not None else ''}" - - def __eq__(self, other): - """ - Compare two device objects for equality. - - Args: - other: The object to compare with - - Returns: - bool: True if both objects are device instances with the same type and index - """ - if not isinstance(other, device): - return False - return self.type == other.type and self.index == other.index - - @staticmethod - def _to_infinicore_device(type: str, index: int): - return _python_2_infinicore_dict[type][index] - - @staticmethod - def _from_infinicore_device(infinicore_device: _infinicore.Device): - return _infinicore_2_python_dict[infinicore_device.type][ - infinicore_device.index - ] - - -_TORCH_DEVICE_MAP = { - _infinicore.Device.Type.CPU: "cpu", - _infinicore.Device.Type.NVIDIA: "cuda", - _infinicore.Device.Type.CAMBRICON: "mlu", - _infinicore.Device.Type.ASCEND: "npu", - _infinicore.Device.Type.METAX: "cuda", - _infinicore.Device.Type.MOORE: "musa", - _infinicore.Device.Type.ILUVATAR: "cuda", - _infinicore.Device.Type.KUNLUN: "cuda", - _infinicore.Device.Type.HYGON: "cuda", - _infinicore.Device.Type.QY: "cuda", - _infinicore.Device.Type.ALI: "cuda", -} - - -def _initialize_device_relationship(all_device_types, all_device_count): - # python_device_type_set: {'cpu', 'musa', 'npu', 'cuda', 'mlu'} - python_device_type_set = set([type for type in _TORCH_DEVICE_MAP.values()]) - python_device_current_index = {type: 0 for type in python_device_type_set} - - infinicore_2_python_dict = {} - python_2_infinicore_dict = {} - for infinicore_device_type, infinicore_device_count in zip( - all_device_types, all_device_count - ): - if 0 == infinicore_device_count: - continue - - # Found one device - for infinicore_device_index in range(infinicore_device_count): - # Create instantiation objects for C++ devices - infinicore_instance = _infinicore.Device( - infinicore_device_type, infinicore_device_index - ) - - # Create instantiation objects for python devices - python_device_type = _TORCH_DEVICE_MAP[infinicore_device_type] - - python_device_index = python_device_current_index[python_device_type] - python_device_current_index[python_device_type] += 1 - - python_instance = device(python_device_type, python_device_index) - - # Cache corresponding relationship - if infinicore_2_python_dict.get(infinicore_device_type) is not None: - infinicore_2_python_dict[infinicore_device_type].append(python_instance) - else: - infinicore_2_python_dict[infinicore_device_type] = [python_instance] - - if python_2_infinicore_dict.get(python_device_type) is not None: - python_2_infinicore_dict[python_device_type].append(infinicore_instance) - else: - python_2_infinicore_dict[python_device_type] = [infinicore_instance] - - return infinicore_2_python_dict, python_2_infinicore_dict - - -_all_device_types = tuple(_infinicore.Device.Type.__members__.values())[:-1] -_all_device_count = tuple( - _infinicore.get_device_count(device) for device in _all_device_types -) - -_infinicore_2_python_dict, _python_2_infinicore_dict = _initialize_device_relationship( - _all_device_types, _all_device_count -) diff --git a/python/infinicore/device_event.py b/python/infinicore/device_event.py deleted file mode 100644 index b15d06bc4..000000000 --- a/python/infinicore/device_event.py +++ /dev/null @@ -1,136 +0,0 @@ -import infinicore.device -from infinicore.lib import _infinicore - - -class DeviceEvent: - """A device event for timing operations and synchronization across devices. - - Similar to torch.cuda.Event, this class provides functionality to: - - Record events on specific device streams - - Synchronize with events - - Measure elapsed time between events - - Query event completion status - - Make streams wait for events - - Args: - enable_timing: Whether the event should record timing data. Default: False. - blocking: Whether to use blocking synchronization. Default: False. - interprocess: Whether the event can be used for inter-process communication. Default: False. - external: Whether the event is an external event. Default: False. - device: Target device for this event. If None, uses current device. - """ - - def __init__(self, enable_timing=False, device=None): - # Build flags based on parameters - flags = 0 - if not enable_timing: - flags |= 0x2 # DISABLE_TIMING - # if blocking: - # flags |= 0x1 # BLOCKING_SYNC - - # Store parameters for reference - self._enable_timing = enable_timing - # self._blocking = blocking - # self._interprocess = interprocess - # self._external = external - - if device is None: - # Use current device - if flags == 0: - self._underlying = _infinicore.DeviceEvent() - else: - self._underlying = _infinicore.DeviceEvent(flags) - elif flags == 0: - # Construct with device only - self._underlying = _infinicore.DeviceEvent(device._underlying) - else: - # Construct with both device and flags - self._underlying = _infinicore.DeviceEvent(device._underlying, flags) - - def record(self, stream=None): - """Record the event. - - Args: - stream: Stream to record the event on. If None, uses current stream. - """ - if stream is None: - self._underlying.record() - else: - self._underlying.record(stream) - - def synchronize(self): - """Wait for the event to complete (blocking).""" - self._underlying.synchronize() - - def query(self): - """Check if the event has been completed. - - Returns: - bool: True if completed, False otherwise. - """ - return self._underlying.query() - - def elapsed_time(self, other): - """Calculate elapsed time between this event and another event. - - Args: - other: The other DeviceEvent to compare with - - Returns: - float: Elapsed time in milliseconds between this event and the other event - - Raises: - RuntimeError: If events are on different devices or not recorded, - or if timing is disabled on either event - """ - if not self._enable_timing or not other._enable_timing: - raise RuntimeError("Cannot measure elapsed time when timing is disabled") - return self._underlying.elapsed_time(other._underlying) - - def wait(self, stream=None): - """Make a stream wait for this event to complete. - - Args: - stream: Stream to make wait for this event. If None, uses current stream. - """ - self._underlying.wait(stream) - - @property - def device(self): - """Get the device where this event was created.""" - return infinicore.device._from_infinicore_device(self._underlying.device) - - @property - def is_recorded(self): - """Check if the event has been recorded.""" - return self._underlying.is_recorded - - @property - def enable_timing(self): - """Whether this event records timing data.""" - return self._enable_timing - - @property - def blocking(self): - """Whether this event uses blocking synchronization.""" - return self._blocking - - @property - def interprocess(self): - """Whether this event can be used for inter-process communication.""" - return self._interprocess - - def __repr__(self): - flags_str = [] - if not self._enable_timing: - flags_str.append("timing_disabled") - if self._blocking: - flags_str.append("blocking") - if self._interprocess: - flags_str.append("interprocess") - if self._external: - flags_str.append("external") - if not flags_str: - flags_str.append("default") - - return f"DeviceEvent(device={self.device}, flags={', '.join(flags_str)}, recorded={self.is_recorded})" diff --git a/python/infinicore/dtype.py b/python/infinicore/dtype.py deleted file mode 100644 index a323471c2..000000000 --- a/python/infinicore/dtype.py +++ /dev/null @@ -1,78 +0,0 @@ -from infinicore.lib import _infinicore - - -class dtype: - def __init__(self, data_type): - """An internal method. Please do not use this directly.""" - self._underlying = data_type - - def __repr__(self): - repr_map = { - _infinicore.DataType.BYTE: "uint8", - _infinicore.DataType.BOOL: "bool", - _infinicore.DataType.I8: "int8", - _infinicore.DataType.I16: "int16", - _infinicore.DataType.I32: "int32", - _infinicore.DataType.I64: "int64", - _infinicore.DataType.U8: "uint8", - _infinicore.DataType.U16: "uint16", - _infinicore.DataType.U32: "uint32", - _infinicore.DataType.U64: "uint64", - _infinicore.DataType.F8: "float8", - _infinicore.DataType.F16: "float16", - _infinicore.DataType.F32: "float32", - _infinicore.DataType.F64: "float64", - _infinicore.DataType.C16: "complex16", - _infinicore.DataType.C32: "complex32", - _infinicore.DataType.C64: "complex64", - _infinicore.DataType.C128: "complex128", - _infinicore.DataType.BF16: "bfloat16", - } - return f"infinicore.{repr_map[self._underlying]}" - - def __eq__(self, other): - """ - Compare two dtype objects for equality. - - Args: - other: The object to compare with - - Returns: - bool: True if both objects are dtype instances with the same underlying data type - """ - if not isinstance(other, dtype): - return False - return self._underlying == other._underlying - - def __hash__(self): - """ - Return a hash value for the dtype object. - - Returns: - int: Hash value based on the underlying data type - """ - return hash(self._underlying) - - -float32 = dtype(_infinicore.DataType.F32) -float = float32 -float64 = dtype(_infinicore.DataType.F64) -double = float64 -complex32 = dtype(_infinicore.DataType.C32) -chalf = complex32 -complex64 = dtype(_infinicore.DataType.C64) -cfloat = complex64 -complex128 = dtype(_infinicore.DataType.C128) -cdouble = complex128 -float16 = dtype(_infinicore.DataType.F16) -half = float16 -bfloat16 = dtype(_infinicore.DataType.BF16) -uint8 = dtype(_infinicore.DataType.U8) -int8 = dtype(_infinicore.DataType.I8) -int16 = dtype(_infinicore.DataType.I16) -short = int16 -int32 = dtype(_infinicore.DataType.I32) -int = int32 -int64 = dtype(_infinicore.DataType.I64) -long = int64 -bool = dtype(_infinicore.DataType.BOOL) diff --git a/python/infinicore/graph.py b/python/infinicore/graph.py deleted file mode 100644 index 7d7feb970..000000000 --- a/python/infinicore/graph.py +++ /dev/null @@ -1,18 +0,0 @@ -from infinicore.lib import _infinicore - - -class Graph: - """ - Python wrapper around a InfiniCore Graph instance. - """ - - def __init__(self, graph: _infinicore.Graph): - if not isinstance(graph, _infinicore.Graph): - raise TypeError("Expected _infinicore.Graph") - self._graph = graph - - def run(self): - return self._graph.run() - - def __repr__(self): - return f"" diff --git a/python/infinicore/nn/__init__.py b/python/infinicore/nn/__init__.py deleted file mode 100644 index 73c9f0aaa..000000000 --- a/python/infinicore/nn/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from infinicore.nn import functional -from infinicore.nn.modules import * # noqa: F403 -from infinicore.nn.parameter import InfiniCoreParameter as Parameter - -__all__ = ["functional", "Parameter"] diff --git a/python/infinicore/nn/functional/__init__.py b/python/infinicore/nn/functional/__init__.py deleted file mode 100644 index 0395c5626..000000000 --- a/python/infinicore/nn/functional/__init__.py +++ /dev/null @@ -1,103 +0,0 @@ -from infinicore.ops.conv2d import conv2d - -from .adaptive_avg_pool1d import adaptive_avg_pool1d -from .adaptive_avg_pool3d import adaptive_avg_pool3d -from .adaptive_max_pool1d import adaptive_max_pool1d -from .affine_grid import affine_grid -from .avg_pool1d import avg_pool1d -from .binary_cross_entropy_with_logits import binary_cross_entropy_with_logits -from .causal_conv1d import causal_conv1d -from .causal_softmax import causal_softmax -from .chunk_gated_delta_rule import chunk_gated_delta_rule -from .embedding import embedding -from .flash_attention import flash_attention -from .fused_gated_delta_net_gating import fused_gated_delta_net_gating -from .gaussian_nll_loss import gaussian_nll_loss -from .hardswish import hardswish -from .hardtanh import hardtanh -from .hinge_embedding_loss import hinge_embedding_loss -from .huber_loss import huber_loss -from .interpolate import interpolate -from .layer_norm import layer_norm -from .linear import linear -from .linear_w8a8i8 import linear_w8a8i8 -from .log_softmax import log_softmax -from .mamba_selective_scan import mamba_selective_scan -from .multi_margin_loss import multi_margin_loss -from .pad import pad -from .prelu import prelu -from .random_sample import random_sample -from .recurrent_gated_delta_rule import recurrent_gated_delta_rule -from .relu6 import relu6 -from .rms_norm import rms_norm -from .rope import RopeAlgo, rope -from .selu import selu -from .sigmoid import sigmoid -from .silu import silu -from .silu_and_mul import silu_and_mul -from .smooth_l1_loss import smooth_l1_loss -from .softplus import softplus -from .softsign import softsign -from .swiglu import swiglu -from .tanhshrink import tanhshrink -from .triplet_margin_loss import triplet_margin_loss -from .triplet_margin_with_distance_loss import triplet_margin_with_distance_loss -from .unfold import unfold -from .upsample_bilinear import upsample_bilinear - -__all__ = [ - "conv2d", - "adaptive_max_pool1d", - "causal_conv1d", - "causal_softmax", - "chunk_gated_delta_rule", - "embedding", - "flash_attention", - "fused_gated_delta_net_gating", - "gaussian_nll_loss", - "interpolate", - "linear", - "binary_cross_entropy_with_logits", - "random_sample", - "adaptive_avg_pool1d", - "affine_grid", - "prelu", - "relu6", - "rms_norm", - "recurrent_gated_delta_rule", - "sigmoid", - "silu", - "smooth_l1_loss", - "swiglu", - "interpolate", - "linear", - "triplet_margin_loss", - "upsample_bilinear", - "interpolate", - "log_softmax", - "mamba_selective_scan", - "upsample_nearest", - "triplet_margin_with_distance_loss", - "embedding", - "rope", - "unfold", - "RopeAlgo", - "rope", - "selu", - "hinge_embedding_loss", - "pad", - "silu", - "hardswish", - "hardtanh", - "avg_pool1d", - "swiglu", - "linear_w8a8i8", - "silu_and_mul", - "adaptive_avg_pool3d", - "tanhshrink", - "multi_margin_loss", - "softplus", - "softsign", - "huber_loss", - "layer_norm", -] diff --git a/python/infinicore/nn/functional/adaptive_avg_pool1d.py b/python/infinicore/nn/functional/adaptive_avg_pool1d.py deleted file mode 100644 index e06e08e3e..000000000 --- a/python/infinicore/nn/functional/adaptive_avg_pool1d.py +++ /dev/null @@ -1,7 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def adaptive_avg_pool1d(input: Tensor, output_size: int) -> Tensor: - r"""Apply a 1D adaptive average pooling.""" - return Tensor(_infinicore.adaptive_avg_pool1d(input._underlying, output_size)) diff --git a/python/infinicore/nn/functional/adaptive_avg_pool3d.py b/python/infinicore/nn/functional/adaptive_avg_pool3d.py deleted file mode 100644 index 797d8015d..000000000 --- a/python/infinicore/nn/functional/adaptive_avg_pool3d.py +++ /dev/null @@ -1,17 +0,0 @@ -from typing import List - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def adaptive_avg_pool3d(x: Tensor, output_size: List[int] = {1, 1, 1}) -> Tensor: - r"""Applies a 3D adaptive average pooling over an input signal composed of several input planes. - - Args: - x (Tensor): The input tensor of shape (N, C, D, H, W) - output_size (List[int]): The target output size of the form (d, h, w) - - Returns: - Tensor: The pooled output tensor - """ - return Tensor(_infinicore.adaptive_avg_pool3d(x._underlying, output_size)) diff --git a/python/infinicore/nn/functional/adaptive_max_pool1d.py b/python/infinicore/nn/functional/adaptive_max_pool1d.py deleted file mode 100644 index 8271a231f..000000000 --- a/python/infinicore/nn/functional/adaptive_max_pool1d.py +++ /dev/null @@ -1,37 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def adaptive_max_pool1d( - input: Tensor, - output_size: int, - *, - out=None, -) -> Tensor: - r"""Applies a 1D adaptive max pooling over an input signal composed of - several input planes. - - The output size is H_out. The algorithm used is fairly simple: - - .. math:: - \text{start} = \left\lfloor \frac{i \cdot L_{in}}{L_{out}} \right\rfloor - - \text{end} = \left\lceil \frac{(i + 1) \cdot L_{in}}{L_{out}} \right\rceil - - where :math:`L_{in}` is the size of the input dimension, and :math:`L_{out}` is the size of the output dimension. - - Args: - input (Tensor): Input tensor of shape (N, C, L_in) - output_size (int): The target output size (L_out) - out (Tensor, optional): Output tensor. - - Returns: - Tensor: The result of the adaptive max pooling operation. - """ - - if out is None: - return Tensor(_infinicore.adaptive_max_pool1d(input._underlying, output_size)) - - _infinicore.adaptive_max_pool1d_(out._underlying, input._underlying, output_size) - - return out diff --git a/python/infinicore/nn/functional/affine_grid.py b/python/infinicore/nn/functional/affine_grid.py deleted file mode 100644 index d01ac2a70..000000000 --- a/python/infinicore/nn/functional/affine_grid.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def affine_grid(theta: Tensor, size: list[int], align_corners: bool = False) -> Tensor: - r"""Generates a 2D flow field (sampling grid), given a batch of affine matrices theta.""" - - # 直接调用底层绑定 - # theta._underlying: 传递底层 C++ Tensor 对象 - # size: Python list[int] 自动转换为 C++ std::vector - return Tensor(_infinicore.affine_grid(theta._underlying, size, align_corners)) diff --git a/python/infinicore/nn/functional/avg_pool1d.py b/python/infinicore/nn/functional/avg_pool1d.py deleted file mode 100644 index 0cf4759ad..000000000 --- a/python/infinicore/nn/functional/avg_pool1d.py +++ /dev/null @@ -1,24 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def avg_pool1d( - input: Tensor, - kernel_size: int, - stride: int | None = None, - padding: int = 0, - *, - out=None, -) -> Tensor: - if stride is None: - stride = 0 - - if out is None: - return Tensor( - _infinicore.avg_pool1d(input._underlying, kernel_size, stride, padding) - ) - - _infinicore.avg_pool1d_( - out._underlying, input._underlying, kernel_size, stride, padding - ) - return out diff --git a/python/infinicore/nn/functional/binary_cross_entropy_with_logits.py b/python/infinicore/nn/functional/binary_cross_entropy_with_logits.py deleted file mode 100644 index 98d2e3ae9..000000000 --- a/python/infinicore/nn/functional/binary_cross_entropy_with_logits.py +++ /dev/null @@ -1,103 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def binary_cross_entropy_with_logits( - input: Tensor, - target: Tensor, - weight: Tensor | None = None, - pos_weight: Tensor | None = None, - reduction: str = "mean", - *, - out: Tensor | None = None, -) -> Tensor: - """Binary cross entropy loss with logits. - - This wraps the underlying C++/CUDA implementation exposed via `_infinicore`. - - The low-level binding treats missing ``weight`` / ``pos_weight`` via - default-constructed tensors. Here we avoid passing ``None`` down and - instead omit arguments when they are not provided, so pybind11 uses - its defaults. - """ - - # Out-of-place API - if out is None: - # Neither weight nor pos_weight - if weight is None and pos_weight is None: - return Tensor( - _infinicore.binary_cross_entropy_with_logits( - input._underlying, - target._underlying, - reduction=reduction, - ) - ) - - # weight provided only - if weight is not None and pos_weight is None: - return Tensor( - _infinicore.binary_cross_entropy_with_logits( - input._underlying, - target._underlying, - weight._underlying, - reduction=reduction, - ) - ) - - # pos_weight provided only - if weight is None and pos_weight is not None: - return Tensor( - _infinicore.binary_cross_entropy_with_logits( - input._underlying, - target._underlying, - pos_weight=pos_weight._underlying, - reduction=reduction, - ) - ) - - # both provided - return Tensor( - _infinicore.binary_cross_entropy_with_logits( - input._underlying, - target._underlying, - weight._underlying, - pos_weight._underlying, - reduction, - ) - ) - - # In-place-style API with explicit out - if weight is None and pos_weight is None: - _infinicore.binary_cross_entropy_with_logits_( - out._underlying, - input._underlying, - target._underlying, - reduction=reduction, - ) - elif weight is not None and pos_weight is None: - _infinicore.binary_cross_entropy_with_logits_( - out._underlying, - input._underlying, - target._underlying, - weight._underlying, - reduction=reduction, - ) - elif weight is None and pos_weight is not None: - _infinicore.binary_cross_entropy_with_logits_( - out._underlying, - input._underlying, - target._underlying, - pos_weight=pos_weight._underlying, - reduction=reduction, - ) - else: - _infinicore.binary_cross_entropy_with_logits_( - out._underlying, - input._underlying, - target._underlying, - weight._underlying, - pos_weight._underlying, - reduction, - ) - - return out diff --git a/python/infinicore/nn/functional/causal_conv1d.py b/python/infinicore/nn/functional/causal_conv1d.py deleted file mode 100644 index 04116a858..000000000 --- a/python/infinicore/nn/functional/causal_conv1d.py +++ /dev/null @@ -1,51 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def causal_conv1d( - qkv: Tensor, - conv_state: Tensor, - weight: Tensor, - bias: Tensor | None = None, - *, - cu_seqlens: Tensor | None = None, - initial_state_indices: Tensor | None = None, - final_state_indices: Tensor | None = None, -) -> Tensor: - """Run causal depthwise Conv1d and return only ``out``. - - Padded mode: - qkv/out: ``[B, T, C]`` - conv_state: ``[B, C, state_len]`` - - Continuous-batch mode: - Pass ``cu_seqlens`` with shape ``[num_requests + 1]``. - qkv/out: ``[1, total_tokens, C]`` - - Indexed state-pool mode: - conv_state is ``[pool_size, C, state_len]``. - ``initial_state_indices`` is optional and selects read slots. - If ``final_state_indices`` is provided, final states are written - in-place to ``conv_state[final_state_indices]``. Otherwise an internal - final-state tensor is allocated by the C++ wrapper and discarded. - - Notes: - The current infiniop backend supports ``K == 4`` only, where ``weight`` - has shape ``[C, 1, K]`` and ``conv_state`` has shape ``[*, C, K - 1]``. - No activation is applied; call ``silu`` separately when needed. - """ - return Tensor( - _infinicore.causal_conv1d( - qkv._underlying, - conv_state._underlying, - weight._underlying, - None if bias is None else bias._underlying, - None if cu_seqlens is None else cu_seqlens._underlying, - ( - None - if initial_state_indices is None - else initial_state_indices._underlying - ), - None if final_state_indices is None else final_state_indices._underlying, - ) - ) diff --git a/python/infinicore/nn/functional/causal_softmax.py b/python/infinicore/nn/functional/causal_softmax.py deleted file mode 100644 index 8a4e1ef22..000000000 --- a/python/infinicore/nn/functional/causal_softmax.py +++ /dev/null @@ -1,13 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def causal_softmax(input: Tensor, out=None) -> Tensor: - r"""Apply a causal softmax function.""" - - if out is None: - return Tensor(_infinicore.causal_softmax(input._underlying)) - - _infinicore.causal_softmax_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/nn/functional/chunk_gated_delta_rule.py b/python/infinicore/nn/functional/chunk_gated_delta_rule.py deleted file mode 100644 index dd87a9f80..000000000 --- a/python/infinicore/nn/functional/chunk_gated_delta_rule.py +++ /dev/null @@ -1,68 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def chunk_gated_delta_rule( - q: Tensor, - k: Tensor, - v: Tensor, - g: Tensor, - beta: Tensor, - initial_state: Tensor, - *, - cu_seqlens: Tensor | None = None, - initial_state_indices: Tensor | None = None, - final_state_indices: Tensor | None = None, - use_qk_l2norm: bool = False, - chunk_size: int = 64, -) -> Tensor: - """Run chunk gated delta rule and return only ``out``. - - Padded mode shapes: - q, k: ``[B, T, Hk, Dk]`` - v, out: ``[B, T, Hv, Dv]`` - g, beta: ``[B, T, Hv]`` - initial_state: ``[B, Hv, Dv, Dk]`` - - Continuous-batch mode shapes: - Pass ``cu_seqlens`` with shape ``[B + 1]`` and dtype int32/int64. - q, k: ``[1, total_tokens, Hk, Dk]`` - v, out: ``[1, total_tokens, Hv, Dv]`` - g, beta: ``[1, total_tokens, Hv]`` - - Indexed state-pool mode: - initial_state is ``[pool_size, Hv, Dv, Dk]``. - ``initial_state_indices`` and ``final_state_indices`` are both ``[B]`` - int32/int64 tensors. The final state is written in-place into - ``initial_state[final_state_indices]`` and no final state tensor is - returned. - - Notes: - ``Hv`` must be a multiple of ``Hk``. q/k/v/out may be strided in the - first three dimensions, but the last dimension must be contiguous. - g and beta may use a different floating dtype from q/k/v/state. - """ - if (initial_state_indices is None) != (final_state_indices is None): - raise ValueError( - "initial_state_indices and final_state_indices must be provided together" - ) - - return Tensor( - _infinicore.chunk_gated_delta_rule( - q._underlying, - k._underlying, - v._underlying, - g._underlying, - beta._underlying, - initial_state._underlying, - None if cu_seqlens is None else cu_seqlens._underlying, - ( - None - if initial_state_indices is None - else initial_state_indices._underlying - ), - None if final_state_indices is None else final_state_indices._underlying, - use_qk_l2norm, - chunk_size, - ) - ) diff --git a/python/infinicore/nn/functional/embedding.py b/python/infinicore/nn/functional/embedding.py deleted file mode 100644 index 592a12290..000000000 --- a/python/infinicore/nn/functional/embedding.py +++ /dev/null @@ -1,32 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def embedding( - input: Tensor, - weight: Tensor, - padding_idx=None, - max_norm=None, - norm_type=2.0, - scale_grad_by_freq=False, - sparse=False, - *, - out=None, -) -> Tensor: - r"""Generate a simple lookup table that looks up embeddings in a fixed dictionary and size.""" - - assert ( - (padding_idx is None) - and (max_norm is None) - and (scale_grad_by_freq is False) - and (sparse is False) - ), "Unsupported parameters." - - # Note: embedding now supports device-side input for graph recording - # The C++ implementation handles both CPU and device-side inputs - - if out is None: - return Tensor(_infinicore.embedding(input._underlying, weight._underlying)) - - _infinicore.embedding_(out._underlying, input._underlying, weight._underlying) - return out diff --git a/python/infinicore/nn/functional/flash_attention.py b/python/infinicore/nn/functional/flash_attention.py deleted file mode 100644 index 8f42e865f..000000000 --- a/python/infinicore/nn/functional/flash_attention.py +++ /dev/null @@ -1,34 +0,0 @@ -import math - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def flash_attention( - query, - key, - value, - total_kv_len, - attn_mask=None, - dropout_p=0, - is_causal=False, - scale=None, - enable_gqa=False, -): - assert attn_mask is None and dropout_p == 0 and not enable_gqa - - emb_dim = query.shape[-1] - - if scale is None: - scale = 1 / math.sqrt(emb_dim) - - return Tensor( - _infinicore.flash_attention( - query._underlying, - key._underlying, - value._underlying, - total_kv_len._underlying, - scale, - is_causal, - ) - ) diff --git a/python/infinicore/nn/functional/fused_gated_delta_net_gating.py b/python/infinicore/nn/functional/fused_gated_delta_net_gating.py deleted file mode 100644 index faa6a339f..000000000 --- a/python/infinicore/nn/functional/fused_gated_delta_net_gating.py +++ /dev/null @@ -1,37 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def fused_gated_delta_net_gating( - A_log: Tensor, - a: Tensor, - b: Tensor, - dt_bias: Tensor, - beta: float = 1.0, - threshold: float = 20.0, - *, - out: tuple[Tensor, Tensor] | None = None, -) -> tuple[Tensor, Tensor]: - if out is None: - g, beta_output = _infinicore.fused_gated_delta_net_gating( - A_log._underlying, - a._underlying, - b._underlying, - dt_bias._underlying, - beta, - threshold, - ) - return Tensor(g), Tensor(beta_output) - - g, beta_output = out - _infinicore.fused_gated_delta_net_gating_( - g._underlying, - beta_output._underlying, - A_log._underlying, - a._underlying, - b._underlying, - dt_bias._underlying, - beta, - threshold, - ) - return g, beta_output diff --git a/python/infinicore/nn/functional/gaussian_nll_loss.py b/python/infinicore/nn/functional/gaussian_nll_loss.py deleted file mode 100644 index 21f3bde92..000000000 --- a/python/infinicore/nn/functional/gaussian_nll_loss.py +++ /dev/null @@ -1,32 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - -_REDUCTION_TO_INT = { - "none": 0, - "mean": 1, - "sum": 2, -} - - -def gaussian_nll_loss( - input: Tensor, - target: Tensor, - var: Tensor, - full: bool = False, - eps: float = 1e-6, - reduction: str = "mean", -) -> Tensor: - reduction_i = _REDUCTION_TO_INT.get(reduction) - if reduction_i is None: - raise ValueError(f"Unsupported reduction: {reduction!r}") - - return Tensor( - _infinicore.gaussian_nll_loss( - input._underlying, - target._underlying, - var._underlying, - bool(full), - float(eps), - int(reduction_i), - ) - ) diff --git a/python/infinicore/nn/functional/hardswish.py b/python/infinicore/nn/functional/hardswish.py deleted file mode 100644 index b054b8978..000000000 --- a/python/infinicore/nn/functional/hardswish.py +++ /dev/null @@ -1,28 +0,0 @@ -import infinicore -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def hardswish(input: Tensor, inplace: bool = False, *, out=None) -> Tensor: - r"""Apply the Hardswish activation function element-wise.""" - - if ( - infinicore.use_ntops - and input.device.type in ("cuda", "musa") - and out is None - and hasattr(infinicore.ntops.torch, "hardswish") - ): - try: - return infinicore.ntops.torch.hardswish(input, inplace=inplace) - except AttributeError: - pass - - if inplace: - _infinicore.hardswish_(input._underlying, input._underlying) - return input - - if out is None: - return Tensor(_infinicore.hardswish(input._underlying)) - - _infinicore.hardswish_(out._underlying, input._underlying) - return out diff --git a/python/infinicore/nn/functional/hardtanh.py b/python/infinicore/nn/functional/hardtanh.py deleted file mode 100644 index 925de33d6..000000000 --- a/python/infinicore/nn/functional/hardtanh.py +++ /dev/null @@ -1,46 +0,0 @@ -import infinicore -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def hardtanh( - input: Tensor, - min_val: float = -1.0, - max_val: float = 1.0, - inplace: bool = False, - *, - out=None, -) -> Tensor: - """Clamp the input tensor to the range [min_val, max_val].""" - - if min_val > max_val: - raise ValueError("min_val must be less than or equal to max_val") - - if ( - infinicore.use_ntops - and input.device.type in ("cuda", "musa") - and out is None - and hasattr(infinicore.ntops.torch, "hardtanh") - ): - try: - return infinicore.ntops.torch.hardtanh( - input, min_val=min_val, max_val=max_val, inplace=inplace - ) - except AttributeError: - pass - - if inplace: - _infinicore.hardtanh_( - input._underlying, input._underlying, float(min_val), float(max_val) - ) - return input - - if out is None: - return Tensor( - _infinicore.hardtanh(input._underlying, float(min_val), float(max_val)) - ) - - _infinicore.hardtanh_( - out._underlying, input._underlying, float(min_val), float(max_val) - ) - return out diff --git a/python/infinicore/nn/functional/hinge_embedding_loss.py b/python/infinicore/nn/functional/hinge_embedding_loss.py deleted file mode 100644 index 824f4212a..000000000 --- a/python/infinicore/nn/functional/hinge_embedding_loss.py +++ /dev/null @@ -1,30 +0,0 @@ -from __future__ import annotations - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - -_REDUCTION_MAP: dict[str, int] = { - "none": 0, - "mean": 1, - "sum": 2, -} - - -def hinge_embedding_loss( - input: Tensor, - target: Tensor, - margin: float = 1.0, - reduction: str = "mean", -) -> Tensor: - reduction_enum = _REDUCTION_MAP.get(reduction) - if reduction_enum is None: - raise ValueError(f"Unsupported reduction: {reduction!r}") - - return Tensor( - _infinicore.hinge_embedding_loss( - input._underlying, - target._underlying, - float(margin), - reduction_enum, - ) - ) diff --git a/python/infinicore/nn/functional/huber_loss.py b/python/infinicore/nn/functional/huber_loss.py deleted file mode 100644 index 4f8e67472..000000000 --- a/python/infinicore/nn/functional/huber_loss.py +++ /dev/null @@ -1,45 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - -_REDUCTION_MODES = { - "none": 0, - "mean": 1, - "sum": 2, -} - - -def huber_loss( - input: Tensor, - target: Tensor, - delta: float = 1.0, - reduction: str = "mean", - *, - out: Optional[Tensor] = None, -) -> Tensor: - r"""Creates a criterion that uses a squared term if the absolute - element-wise error falls below delta and a delta-scaled L1 term otherwise. - """ - - if not input.is_contiguous(): - input = input.contiguous() - if not target.is_contiguous(): - target = target.contiguous() - - # 解析 reduction 参数 - if reduction not in _REDUCTION_MODES: - raise ValueError(f"{reduction} is not a valid value for reduction") - reduction_val = _REDUCTION_MODES[reduction] - - if out is not None: - _infinicore.huber_loss_( - out._underlying, input._underlying, target._underlying, delta, reduction_val - ) - return out - - return Tensor( - _infinicore.huber_loss( - input._underlying, target._underlying, delta, reduction_val - ) - ) diff --git a/python/infinicore/nn/functional/interpolate.py b/python/infinicore/nn/functional/interpolate.py deleted file mode 100644 index 739e4de36..000000000 --- a/python/infinicore/nn/functional/interpolate.py +++ /dev/null @@ -1,94 +0,0 @@ -from collections.abc import Iterable - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def _to_int64_list(value) -> list[int]: - if isinstance(value, int): - return [int(value)] - if isinstance(value, Iterable): - return [int(v) for v in value] - raise TypeError(f"Expected int or iterable of ints, got {type(value).__name__}") - - -def _to_double_list(value) -> list[float]: - if isinstance(value, (int, float)): - return [float(value)] - if isinstance(value, Iterable): - return [float(v) for v in value] - raise TypeError(f"Expected float or iterable of floats, got {type(value).__name__}") - - -# def interpolate( -# input: Tensor, -# size: Optional[Union[int, Sequence[int]]] = None, -# scale_factor: Optional[Union[float, Sequence[float]]] = None, -# mode: str = "nearest", -# align_corners: Optional[bool] = None, -# recompute_scale_factor: Optional[bool] = None, -# ) -> Tensor: -# if mode == "nearest": -# if align_corners is not None: -# raise ValueError( -# "align_corners option can only be set with the " -# "interpolating modes: linear | bilinear | bicubic | trilinear" -# ) -# return upsample_nearest(input, size, scale_factor) - -# if mode == "bilinear": -# if align_corners is None: -# align_corners = False -# return upsample_bilinear(input, size, scale_factor, align_corners) - -# raise NotImplementedError( -# f"Interpolation mode '{mode}' is not currently supported." -# ) - - -def interpolate( - input: Tensor, - size=None, - scale_factor=None, - mode: str = "nearest", - align_corners=None, -) -> Tensor: - size_list: list[int] = [] if size is None else _to_int64_list(size) - scale_list: list[float] = ( - [] if scale_factor is None else _to_double_list(scale_factor) - ) - - if bool(size_list) == bool(scale_list): - raise ValueError("Expected exactly one of size or scale_factor") - - spatial_ndim = input.ndim - 2 - if spatial_ndim < 1: - raise ValueError("interpolate expects input with at least 3 dimensions") - - if size_list: - if len(size_list) == 1 and spatial_ndim > 1: - size_list = size_list * spatial_ndim - if len(size_list) != spatial_ndim: - raise ValueError( - f"Expected size to have length {spatial_ndim}, got {len(size_list)}" - ) - - if scale_list: - if len(scale_list) == 1 and spatial_ndim > 1: - scale_list = scale_list * spatial_ndim - if len(scale_list) != spatial_ndim: - raise ValueError( - f"Expected scale_factor to have length {spatial_ndim}, got {len(scale_list)}" - ) - if any(v != scale_list[0] for v in scale_list[1:]): - raise ValueError( - "Per-dimension scale_factor is not supported; pass a scalar (or equal values)." - ) - - align_i = 0 if align_corners is None else int(bool(align_corners)) - - return Tensor( - _infinicore.interpolate( - input._underlying, str(mode), size_list, scale_list, align_i - ) - ) diff --git a/python/infinicore/nn/functional/layer_norm.py b/python/infinicore/nn/functional/layer_norm.py deleted file mode 100644 index b841fdc07..000000000 --- a/python/infinicore/nn/functional/layer_norm.py +++ /dev/null @@ -1,33 +0,0 @@ -from typing import List - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def layer_norm( - input: Tensor, - normalized_shape: List[int], - weight: Tensor, - bias: Tensor, - eps: float = 1e-5, - *, - out=None, -) -> Tensor: - r"""Apply Layer Normalization.""" - - assert normalized_shape == weight.shape, ( - "normalized_shape does not match weight.shape." - ) - - if out is None: - return Tensor( - _infinicore.layer_norm( - input._underlying, weight._underlying, bias._underlying, eps - ) - ) - - _infinicore.layer_norm_( - out._underlying, input._underlying, weight._underlying, bias._underlying, eps - ) - - return out diff --git a/python/infinicore/nn/functional/linear.py b/python/infinicore/nn/functional/linear.py deleted file mode 100644 index 8cc45e34c..000000000 --- a/python/infinicore/nn/functional/linear.py +++ /dev/null @@ -1,34 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - -__all__ = ["linear"] - - -def linear( - input: Tensor, - weight: Tensor, - bias=None, - *, - alpha: float = 1.0, - out=None, -) -> Tensor: - r"""Applies a linear transformation to the incoming data: y=alpha*xA^T+b.""" - - if out is None: - return Tensor( - _infinicore.linear( - input._underlying, - weight._underlying, - None if bias is None else bias._underlying, - alpha, - ) - ) - - _infinicore.linear_( - out._underlying, - input._underlying, - weight._underlying, - None if bias is None else bias._underlying, - alpha, - ) - return out diff --git a/python/infinicore/nn/functional/linear_w8a8i8.py b/python/infinicore/nn/functional/linear_w8a8i8.py deleted file mode 100644 index 33cb59b0e..000000000 --- a/python/infinicore/nn/functional/linear_w8a8i8.py +++ /dev/null @@ -1,31 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def linear_w8a8i8( - input: Tensor, - weight_packed: Tensor, - weight_scale: Tensor, - bias=None, - out=None, -) -> Tensor: - r"""Linear layer with weight quantized to int8 and input quantized to int8 with per-tensor scale.""" - - if out is None: - return Tensor( - _infinicore.linear_w8a8i8( - input._underlying, - weight_packed._underlying, - weight_scale._underlying, - None if bias is None else bias._underlying, - ) - ) - - _infinicore.linear_w8a8i8_( - out._underlying, - input._underlying, - weight_packed._underlying, - weight_scale._underlying, - None if bias is None else bias._underlying, - ) - return out diff --git a/python/infinicore/nn/functional/log_softmax.py b/python/infinicore/nn/functional/log_softmax.py deleted file mode 100644 index 1afce4f8c..000000000 --- a/python/infinicore/nn/functional/log_softmax.py +++ /dev/null @@ -1,26 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def log_softmax(input: Tensor, dim: int, *, out: Optional[Tensor] = None) -> Tensor: - r"""Applies a softmax followed by a logarithm. - While mathematically equivalent to log(softmax(x)), doing these two - operations separately is slower and numerically unstable. This function - uses an alternative formulation to compute the output and gradient correctly. - """ - - if not input.is_contiguous(): - input = input.contiguous() - - if out is not None: - if not isinstance(out, Tensor): - raise ValueError("out must be a Tensor") - - _infinicore.log_softmax_(out._underlying, input._underlying, dim) - return out - - ret = _infinicore.log_softmax(input._underlying, dim) - - return Tensor(ret) diff --git a/python/infinicore/nn/functional/mamba_selective_scan.py b/python/infinicore/nn/functional/mamba_selective_scan.py deleted file mode 100644 index 579e984e9..000000000 --- a/python/infinicore/nn/functional/mamba_selective_scan.py +++ /dev/null @@ -1,29 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def mamba_selective_scan( - x: Tensor, - dt: Tensor, - b: Tensor, - c: Tensor, - a_log: Tensor, - d: Tensor, - gate: Tensor, - dt_bias: Tensor, - state: Tensor, -) -> Tensor: - """Run Mamba selective scan and update ``state`` in-place.""" - return Tensor( - _infinicore.mamba_selective_scan( - x._underlying, - dt._underlying, - b._underlying, - c._underlying, - a_log._underlying, - d._underlying, - gate._underlying, - dt_bias._underlying, - state._underlying, - ) - ) diff --git a/python/infinicore/nn/functional/multi_margin_loss.py b/python/infinicore/nn/functional/multi_margin_loss.py deleted file mode 100644 index 1c1582581..000000000 --- a/python/infinicore/nn/functional/multi_margin_loss.py +++ /dev/null @@ -1,63 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - -_REDUCTION_MODES = { - "none": 0, - "mean": 1, - "sum": 2, -} - - -def multi_margin_loss( - input: Tensor, - target: Tensor, - p: int = 1, - margin: float = 1.0, - weight: Optional[Tensor] = None, - reduction: str = "mean", - *, - out: Optional[Tensor] = None, -) -> Tensor: - r"""Creates a criterion that optimizes a multi-class classification hinge - loss (margin-based loss) between input x and output y. - """ - - if not input.is_contiguous(): - input = input.contiguous() - if not target.is_contiguous(): - target = target.contiguous() - - weight_underlying = None - if weight is not None: - if not weight.is_contiguous(): - weight = weight.contiguous() - weight_underlying = weight._underlying - - # 解析 reduction 参数 - if reduction not in _REDUCTION_MODES: - raise ValueError(f"{reduction} is not a valid value for reduction") - reduction_val = _REDUCTION_MODES[reduction] - if out is not None: - _infinicore.multi_margin_loss_( - out._underlying, - input._underlying, - target._underlying, - weight_underlying, - p, - margin, - reduction_val, - ) - return out - - return Tensor( - _infinicore.multi_margin_loss( - input._underlying, - target._underlying, - weight_underlying, - p, - margin, - reduction_val, - ) - ) diff --git a/python/infinicore/nn/functional/pad.py b/python/infinicore/nn/functional/pad.py deleted file mode 100644 index 41e6bd955..000000000 --- a/python/infinicore/nn/functional/pad.py +++ /dev/null @@ -1,23 +0,0 @@ -from __future__ import annotations - -from collections.abc import Sequence -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def pad( - input: Tensor, - pad: Sequence[int], - mode: str = "constant", - value: float = 0.0, - *, - out: Optional[Tensor] = None, -) -> Tensor: - pad_list = list(pad) - if out is None: - return Tensor(_infinicore.pad(input._underlying, pad_list, mode, value)) - - _infinicore.pad_(out._underlying, input._underlying, pad_list, mode, value) - return out diff --git a/python/infinicore/nn/functional/prelu.py b/python/infinicore/nn/functional/prelu.py deleted file mode 100644 index d4b5aea47..000000000 --- a/python/infinicore/nn/functional/prelu.py +++ /dev/null @@ -1,6 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def prelu(input: Tensor, weight: Tensor) -> Tensor: - return Tensor(_infinicore.prelu(input._underlying, weight._underlying)) diff --git a/python/infinicore/nn/functional/random_sample.py b/python/infinicore/nn/functional/random_sample.py deleted file mode 100644 index 61410bb8d..000000000 --- a/python/infinicore/nn/functional/random_sample.py +++ /dev/null @@ -1,36 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def random_sample( - logits: Tensor, - random_val: float, - topp: float, - topk: int, - temperature: float, - *, - out=None, -) -> Tensor: - r"""Sample an index from logits with nucleus/top-k filtering.""" - - if out is None: - return Tensor( - _infinicore.random_sample( - logits._underlying, - random_val, - topp, - topk, - temperature, - ) - ) - - _infinicore.random_sample_( - out._underlying, - logits._underlying, - random_val, - topp, - topk, - temperature, - ) - - return out diff --git a/python/infinicore/nn/functional/recurrent_gated_delta_rule.py b/python/infinicore/nn/functional/recurrent_gated_delta_rule.py deleted file mode 100644 index 323701e93..000000000 --- a/python/infinicore/nn/functional/recurrent_gated_delta_rule.py +++ /dev/null @@ -1,47 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def recurrent_gated_delta_rule( - q: Tensor, - k: Tensor, - v: Tensor, - g: Tensor, - beta: Tensor, - initial_state: Tensor, - *, - initial_state_indices: Tensor | None = None, - final_state_indices: Tensor | None = None, - use_qk_l2norm: bool = False, -) -> Tensor: - if initial_state_indices is None and final_state_indices is None: - return Tensor( - _infinicore.recurrent_gated_delta_rule( - q._underlying, - k._underlying, - v._underlying, - g._underlying, - beta._underlying, - initial_state._underlying, - use_qk_l2norm, - ) - ) - - if initial_state_indices is None or final_state_indices is None: - raise ValueError( - "initial_state_indices and final_state_indices must be provided together" - ) - - return Tensor( - _infinicore.recurrent_gated_delta_rule_indexed( - q._underlying, - k._underlying, - v._underlying, - g._underlying, - beta._underlying, - initial_state._underlying, - initial_state_indices._underlying, - final_state_indices._underlying, - use_qk_l2norm, - ) - ) diff --git a/python/infinicore/nn/functional/relu6.py b/python/infinicore/nn/functional/relu6.py deleted file mode 100644 index 378c892f6..000000000 --- a/python/infinicore/nn/functional/relu6.py +++ /dev/null @@ -1,14 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def relu6(input: Tensor, inplace: bool = False, *, out: Tensor | None = None) -> Tensor: - if inplace: - _infinicore.relu6_(input._underlying, input._underlying) - return input - - if out is None: - return Tensor(_infinicore.relu6(input._underlying)) - - _infinicore.relu6_(out._underlying, input._underlying) - return out diff --git a/python/infinicore/nn/functional/rms_norm.py b/python/infinicore/nn/functional/rms_norm.py deleted file mode 100644 index 84cd82ae1..000000000 --- a/python/infinicore/nn/functional/rms_norm.py +++ /dev/null @@ -1,26 +0,0 @@ -from typing import List - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def rms_norm( - input: Tensor, - normalized_shape: List[int], - weight: Tensor, - eps: float = 1e-5, - *, - out=None, -) -> Tensor: - r"""Apply Root Mean Square Layer Normalization.""" - - assert normalized_shape == weight.shape, ( - "normalized_shape does not match weight.shape." - ) - - if out is None: - return Tensor(_infinicore.rms_norm(input._underlying, weight._underlying, eps)) - - _infinicore.rms_norm_(out._underlying, input._underlying, weight._underlying, eps) - - return out diff --git a/python/infinicore/nn/functional/rope.py b/python/infinicore/nn/functional/rope.py deleted file mode 100644 index ea0041452..000000000 --- a/python/infinicore/nn/functional/rope.py +++ /dev/null @@ -1,42 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -class RopeAlgo: - r"""Different types of RoPE algorithms.""" - - GPT_J = _infinicore.RoPEAlgo.GPT_J - GPT_NEOX = _infinicore.RoPEAlgo.GPT_NEOX - - -def rope( - x: Tensor, - pos_ids: Tensor, - sin_table: Tensor, - cos_table: Tensor, - algo: RopeAlgo = RopeAlgo.GPT_NEOX, - *, - out=None, -) -> Tensor: - r"""Rotary Position Embedding(RoPE).""" - - if out is None: - return Tensor( - _infinicore.rope( - x._underlying, - pos_ids._underlying, - sin_table._underlying, - cos_table._underlying, - algo, - ) - ) - - _infinicore.rope_( - out._underlying, - x._underlying, - pos_ids._underlying, - sin_table._underlying, - cos_table._underlying, - algo, - ) - return out diff --git a/python/infinicore/nn/functional/selu.py b/python/infinicore/nn/functional/selu.py deleted file mode 100644 index d3dafa124..000000000 --- a/python/infinicore/nn/functional/selu.py +++ /dev/null @@ -1,12 +0,0 @@ -from __future__ import annotations - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def selu(input: Tensor, *, out: Tensor | None = None) -> Tensor: - if out is None: - return Tensor(_infinicore.selu(input._underlying)) - - _infinicore.selu_(out._underlying, input._underlying) - return out diff --git a/python/infinicore/nn/functional/sigmoid.py b/python/infinicore/nn/functional/sigmoid.py deleted file mode 100644 index 59ebb464f..000000000 --- a/python/infinicore/nn/functional/sigmoid.py +++ /dev/null @@ -1,10 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def sigmoid(input: Tensor, *, out: Tensor | None = None) -> Tensor: - if out is None: - return Tensor(_infinicore.sigmoid(input._underlying)) - - _infinicore.sigmoid_(out._underlying, input._underlying) - return out diff --git a/python/infinicore/nn/functional/silu.py b/python/infinicore/nn/functional/silu.py deleted file mode 100644 index 36c7e4e3e..000000000 --- a/python/infinicore/nn/functional/silu.py +++ /dev/null @@ -1,21 +0,0 @@ -import infinicore -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def silu(input: Tensor, inplace: bool = False, *, out=None) -> Tensor: - r"""Apply the Sigmoid Linear Unit (SiLU) function, element-wise.""" - - if infinicore.use_ntops and input.device.type in ("cuda", "musa") and out is None: - return infinicore.ntops.torch.silu(input, inplace=inplace) - - if inplace: - _infinicore.silu_(input._underlying, input._underlying) - return input - - if out is None: - return Tensor(_infinicore.silu(input._underlying)) - - _infinicore.silu_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/nn/functional/silu_and_mul.py b/python/infinicore/nn/functional/silu_and_mul.py deleted file mode 100644 index 395754362..000000000 --- a/python/infinicore/nn/functional/silu_and_mul.py +++ /dev/null @@ -1,17 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def silu_and_mul(input: Tensor, out=None) -> Tensor: - r"""Apply the SiLU and Mul (SwiGLU) function. - - Formula: output = SiLU(input_gate) * input_up - Input shape: [..., 2*d], Output shape: [..., d] - """ - - if out is None: - return Tensor(_infinicore.silu_and_mul(input._underlying)) - - _infinicore.silu_and_mul_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/nn/functional/smooth_l1_loss.py b/python/infinicore/nn/functional/smooth_l1_loss.py deleted file mode 100644 index b8e8bf73e..000000000 --- a/python/infinicore/nn/functional/smooth_l1_loss.py +++ /dev/null @@ -1,55 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - -_REDUCTION_MODES = { - "none": 0, - "mean": 1, - "sum": 2, -} - - -def smooth_l1_loss( - input: Tensor, - target: Tensor, - beta: float = 1.0, - reduction: str = "mean", - *, - out: Optional[Tensor] = None, -) -> Tensor: - r"""Creates a criterion that uses a squared term if the absolute - element-wise error falls below beta and an L1 term otherwise. - - Args: - input (Tensor): the input tensor. - target (Tensor): the target tensor. - beta (float, optional): The threshold at which to change between L1 and L2 loss. - The value must be non-negative. Default: 1.0. - reduction (str, optional): Specifies the reduction to apply to the output: - 'none': no reduction will be applied, - 'mean': the sum of the output will be divided by the number of elements in the output, - 'sum': the output will be summed. Default: 'mean'. - out (Tensor, optional): the output tensor. - - Returns: - Tensor: The loss value. - """ - - if not input.is_contiguous(): - input = input.contiguous() - if not target.is_contiguous(): - target = target.contiguous() - if reduction not in _REDUCTION_MODES: - raise ValueError(f"{reduction} is not a valid value for reduction") - reduction_val = _REDUCTION_MODES[reduction] - if out is not None: - _infinicore.smooth_l1_loss_( - out._underlying, input._underlying, target._underlying, beta, reduction_val - ) - return out - return Tensor( - _infinicore.smooth_l1_loss( - input._underlying, target._underlying, beta, reduction_val - ) - ) diff --git a/python/infinicore/nn/functional/softplus.py b/python/infinicore/nn/functional/softplus.py deleted file mode 100644 index 76e96f978..000000000 --- a/python/infinicore/nn/functional/softplus.py +++ /dev/null @@ -1,13 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def softplus(input, beta=1, threshold=20, *, out=None): - if out is None: - # 修改：将 beta 和 threshold 传递给底层 C++ - return Tensor(_infinicore.softplus(input._underlying, beta, threshold)) - - # 修改：将 beta 和 threshold 传递给底层 C++ (In-place) - _infinicore.softplus_(out._underlying, input._underlying, beta, threshold) - - return out diff --git a/python/infinicore/nn/functional/softsign.py b/python/infinicore/nn/functional/softsign.py deleted file mode 100644 index 6a9074df6..000000000 --- a/python/infinicore/nn/functional/softsign.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def softsign(input, *, out=None): - if out is None: - return Tensor(_infinicore.softsign(input._underlying)) - - _infinicore.softsign_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/nn/functional/swiglu.py b/python/infinicore/nn/functional/swiglu.py deleted file mode 100644 index 050502683..000000000 --- a/python/infinicore/nn/functional/swiglu.py +++ /dev/null @@ -1,13 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def swiglu(input: Tensor, other: Tensor, *, out=None): - r"""Apply the Swish-Gated Linear Unit (SwiGLU) function, element-wise.""" - - if out is None: - return Tensor(_infinicore.swiglu(input._underlying, other._underlying)) - - _infinicore.swiglu_(out._underlying, input._underlying, other._underlying) - - return out diff --git a/python/infinicore/nn/functional/tanhshrink.py b/python/infinicore/nn/functional/tanhshrink.py deleted file mode 100644 index 10d17077d..000000000 --- a/python/infinicore/nn/functional/tanhshrink.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def tanhshrink(input, *, out=None): - if out is None: - return Tensor(_infinicore.tanhshrink(input._underlying)) - - _infinicore.tanhshrink_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/nn/functional/triplet_margin_loss.py b/python/infinicore/nn/functional/triplet_margin_loss.py deleted file mode 100644 index 673ff5e3b..000000000 --- a/python/infinicore/nn/functional/triplet_margin_loss.py +++ /dev/null @@ -1,65 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - -_REDUCTION_MODES = { - "none": 0, - "mean": 1, - "sum": 2, -} - - -def triplet_margin_loss( - anchor: Tensor, - positive: Tensor, - negative: Tensor, - margin: float = 1.0, - p: float = 2, - eps: float = 1e-6, - swap: bool = False, - reduction: str = "mean", - *, - out: Optional[Tensor] = None, -) -> Tensor: - r"""Creates a criterion that measures the triplet loss given an input - tensors x1, x2, x3 and a margin with a value greater than 0. - """ - - if not anchor.is_contiguous(): - anchor = anchor.contiguous() - if not positive.is_contiguous(): - positive = positive.contiguous() - if not negative.is_contiguous(): - negative = negative.contiguous() - - if reduction not in _REDUCTION_MODES: - raise ValueError(f"{reduction} is not a valid value for reduction") - reduction_val = _REDUCTION_MODES[reduction] - - if out is not None: - _infinicore.triplet_margin_loss_( - out._underlying, - anchor._underlying, - positive._underlying, - negative._underlying, - margin, - int(p), - eps, - swap, - reduction_val, - ) - return out - - return Tensor( - _infinicore.triplet_margin_loss( - anchor._underlying, - positive._underlying, - negative._underlying, - margin, - int(p), - eps, - swap, - reduction_val, - ) - ) diff --git a/python/infinicore/nn/functional/triplet_margin_with_distance_loss.py b/python/infinicore/nn/functional/triplet_margin_with_distance_loss.py deleted file mode 100644 index 723a29500..000000000 --- a/python/infinicore/nn/functional/triplet_margin_with_distance_loss.py +++ /dev/null @@ -1,58 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def triplet_margin_with_distance_loss( - anchor: Tensor, - positive: Tensor, - negative: Tensor, - *, - margin: float = 1.0, - swap: bool = False, - reduction: str = "mean", - out: Optional[Tensor] = None, -) -> Tensor: - r"""Calculates the triplet margin loss for a given triplet of tensors. - The loss is defined as: L(a, p, n) = max(d(a, p) - d(a, n) + margin, 0) - """ - - if not anchor.is_contiguous(): - anchor = anchor.contiguous() - if not positive.is_contiguous(): - positive = positive.contiguous() - if not negative.is_contiguous(): - negative = negative.contiguous() - - reduction_map = {"none": 0, "mean": 1, "sum": 2} - if reduction not in reduction_map: - raise ValueError(f"Invalid reduction mode: {reduction}") - - reduction_val = reduction_map[reduction] - - if out is not None: - if not isinstance(out, Tensor): - raise ValueError("out must be a Tensor") - - _infinicore.triplet_margin_with_distance_loss_( - out._underlying, - anchor._underlying, - positive._underlying, - negative._underlying, - margin, - swap, - reduction_val, - ) - return out - - ret = _infinicore.triplet_margin_with_distance_loss( - anchor._underlying, - positive._underlying, - negative._underlying, - margin, - swap, - reduction_val, - ) - - return Tensor(ret) diff --git a/python/infinicore/nn/functional/unfold.py b/python/infinicore/nn/functional/unfold.py deleted file mode 100644 index 223562cd5..000000000 --- a/python/infinicore/nn/functional/unfold.py +++ /dev/null @@ -1,47 +0,0 @@ -from typing import List, Optional, Tuple, Union - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def unfold( - input: Tensor, - kernel_size: Union[int, Tuple[int, ...], List[int]], - dilation: Union[int, Tuple[int, ...], List[int]] = 1, - padding: Union[int, Tuple[int, ...], List[int]] = 0, - stride: Union[int, Tuple[int, ...], List[int]] = 1, - *, - out: Optional[Tensor] = None, -) -> Tensor: - r"""Extracts sliding local blocks from a batched input tensor. - - Also known as im2col. The output tensor contains the flattened blocks. - - Args: - input (Tensor): The input tensor. - kernel_size (int or tuple): The size of the sliding blocks. - dilation (int or tuple, optional): The parameter that controls the stride of elements within the neighborhood. Default: 1. - padding (int or tuple, optional): Implicit zero padding to be added on both sides of input. Default: 0. - stride (int or tuple, optional): The stride of the sliding blocks. Default: 1. - out (Tensor, optional): The output tensor. - """ - - if not input.is_contiguous(): - input = input.contiguous() - - # Helper to ensure parameters are iterable (assuming 2D spatial dims for single int) - def _pair(x): - return (x, x) if isinstance(x, int) else x - - k_val = _pair(kernel_size) - d_val = _pair(dilation) - p_val = _pair(padding) - s_val = _pair(stride) - - if out is not None: - _infinicore.unfold_( - out._underlying, input._underlying, k_val, d_val, p_val, s_val - ) - return out - - return Tensor(_infinicore.unfold(input._underlying, k_val, d_val, p_val, s_val)) diff --git a/python/infinicore/nn/functional/upsample_bilinear.py b/python/infinicore/nn/functional/upsample_bilinear.py deleted file mode 100644 index 783b5846b..000000000 --- a/python/infinicore/nn/functional/upsample_bilinear.py +++ /dev/null @@ -1,74 +0,0 @@ -from typing import Optional, Sequence, Union - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def upsample_bilinear( - input: Tensor, - size: Optional[Union[int, Sequence[int]]] = None, - scale_factor: Optional[Union[float, Sequence[float]]] = None, - align_corners: bool = False, - *, - out: Optional[Tensor] = None, -) -> Tensor: - r""" - Applies bilinear interpolation upsampling to the input tensor. - """ - # 确保输入连续 - if not input.is_contiguous(): - input = input.contiguous() - - # 参数互斥检查 - if (size is None) == (scale_factor is None): - raise ValueError("Either size or scale_factor should be defined, but not both.") - - # 计算目标输出尺寸 (H, W) - output_size = [] - - if size is not None: - if isinstance(size, int): - # 如果是单个整数，应用于 H 和 W - output_size = [size, size] - elif isinstance(size, (list, tuple)): - if len(size) < 2: - raise ValueError( - "size sequence must contain at least 2 elements for bilinear upsampling" - ) - output_size = [size[0], size[1]] - else: - raise ValueError("size must be int or sequence of int") - else: - # 基于 scale_factor 计算 - if isinstance(scale_factor, float): - scale_h = scale_factor - scale_w = scale_factor - elif isinstance(scale_factor, (list, tuple)): - if len(scale_factor) < 2: - raise ValueError( - "scale_factor sequence must contain at least 2 elements" - ) - scale_h = scale_factor[0] - scale_w = scale_factor[1] - else: - raise ValueError("scale_factor must be float or sequence of float") - - # 假设输入是 (..., H, W)，取最后两维 - h_in = input.shape[-2] - w_in = input.shape[-1] - output_size = [int(h_in * scale_h), int(w_in * scale_w)] - - # 1. 显式输出 (In-place / Out parameter) - if out is not None: - if not out.is_contiguous(): - raise RuntimeError("out tensor must be contiguous") - - _infinicore.upsample_bilinear_( - out._underlying, input._underlying, align_corners - ) - return out - - # 2. 函数式调用 (Functional API) - return Tensor( - _infinicore.upsample_bilinear(input._underlying, output_size, align_corners) - ) diff --git a/python/infinicore/nn/functional/upsample_nearest.py b/python/infinicore/nn/functional/upsample_nearest.py deleted file mode 100644 index 7e64e7b5f..000000000 --- a/python/infinicore/nn/functional/upsample_nearest.py +++ /dev/null @@ -1,64 +0,0 @@ -from typing import Optional, Sequence, Union - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def upsample_nearest( - input: Tensor, - size: Optional[Union[int, Sequence[int]]] = None, - scale_factor: Optional[Union[float, Sequence[float]]] = None, - *, - out: Optional[Tensor] = None, -) -> Tensor: - if not input.is_contiguous(): - input = input.contiguous() - - if (size is None) == (scale_factor is None): - raise ValueError("Either size or scale_factor should be defined, but not both.") - - ndim = len(input.shape) - output_size = [] - - if size is not None: - if isinstance(size, int): - if ndim == 3: - output_size = [size] - else: - output_size = [size, size] - elif isinstance(size, (list, tuple)): - output_size = [int(s) for s in size] - else: - raise ValueError("size must be int or sequence of int") - else: - if isinstance(scale_factor, (float, int)): - scales = [float(scale_factor)] - elif isinstance(scale_factor, (list, tuple)): - scales = [float(s) for s in scale_factor] - else: - raise ValueError("scale_factor must be float or sequence of float") - - if ndim == 3: - w_in = input.shape[-1] - scale_w = scales[0] if len(scales) == 1 else scales[-1] - output_size = [int(w_in * scale_w)] - else: - if len(scales) == 1: - scale_h = scale_w = scales[0] - elif len(scales) >= 2: - scale_h, scale_w = scales[0], scales[1] - else: - raise ValueError("scale_factor sequence length mismatch") - - h_in = input.shape[-2] - w_in = input.shape[-1] - output_size = [int(h_in * scale_h), int(w_in * scale_w)] - - if out is not None: - if not out.is_contiguous(): - raise RuntimeError("out tensor must be contiguous") - - _infinicore.upsample_nearest_(out._underlying, input._underlying) - return out - - return Tensor(_infinicore.upsample_nearest(input._underlying, output_size)) diff --git a/python/infinicore/nn/modules/__init__.py b/python/infinicore/nn/modules/__init__.py deleted file mode 100644 index 1ec1349ad..000000000 --- a/python/infinicore/nn/modules/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from .container import InfiniCoreModuleList as ModuleList -from .linear import Linear -from .module import InfiniCoreModule as Module -from .normalization import RMSNorm -from .rope import RoPE -from .sparse import Embedding - -__all__ = ["Linear", "RMSNorm", "Embedding", "RoPE", "ModuleList", "Module"] diff --git a/python/infinicore/nn/modules/container.py b/python/infinicore/nn/modules/container.py deleted file mode 100644 index 2ec3ad426..000000000 --- a/python/infinicore/nn/modules/container.py +++ /dev/null @@ -1,188 +0,0 @@ -# ============================================ -# Copyright (c) 2025, InfiniCore -# -# This file implements InfiniCoreModuleList, which is similar to torch.nn.ModuleList -# but based on InfiniCoreModule for inference purposes. - -import operator -from collections import OrderedDict -from itertools import chain -from typing import Iterator, List, Optional, Sequence, TypeVar, Union - -from .module import InfiniCoreModule as Module - -# Define type variable for module compatibility (supports InfiniCoreModule) -ModuleType = TypeVar("ModuleType", bound=Union["Module"]) - - -class InfiniCoreModuleList(Module): - r"""Holds submodules in a list. - - InfiniCoreModuleList can be indexed like a regular Python list, but - modules it contains are properly registered, and will be visible by all - InfiniCoreModule methods. - - Args: - modules (iterable, optional): an iterable of modules to add - - Example:: - - >>> class MyModel(InfiniCoreModule): - ... def __init__(self): - ... super().__init__() - ... self.linears = InfiniCoreModuleList([ - ... torch.nn.Linear(10, 10) for i in range(10) - ... ]) - ... - ... def forward(self, x): - ... # ModuleList can act as an iterable, or be indexed using ints - ... for i, l in enumerate(self.linears): - ... x = self.linears[i // 2](x) + l(x) - ... return x - """ - - def __init__(self, modules: Optional[Sequence[ModuleType]] = None): - super().__init__() - if modules is not None: - self += modules - - def _get_abs_string_index(self, idx): - """Get the absolute index for the list of modules.""" - idx = operator.index(idx) - if not (-len(self) <= idx < len(self)): - raise IndexError(f"index {idx} is out of range") - if idx < 0: - idx += len(self) - return str(idx) - - def __getitem__( - self, idx: Union[int, slice] - ) -> Union[ModuleType, "InfiniCoreModuleList"]: - if isinstance(idx, slice): - return self.__class__(list(self._modules.values())[idx]) - else: - return self._modules[self._get_abs_string_index(idx)] - - def __setitem__(self, idx: int, module: ModuleType) -> None: - idx = self._get_abs_string_index(idx) - # Use add_module to register module - self.add_module(idx, module) - - def __delitem__(self, idx: Union[int, slice]) -> None: - if isinstance(idx, slice): - indices_to_delete = list(range(len(self._modules)))[idx] - for k in indices_to_delete: - if str(k) in self._modules: - del self._modules[str(k)] - else: - idx_str = self._get_abs_string_index(idx) - if idx_str in self._modules: - del self._modules[idx_str] - - # To preserve numbering, self._modules is being reconstructed with modules after deletion - if len(self._modules) > 0: - str_indices = [str(i) for i in range(len(self._modules))] - self._modules = OrderedDict(list(zip(str_indices, self._modules.values()))) - - def __len__(self) -> int: - return len(self._modules) - - def __iter__(self) -> Iterator[ModuleType]: - return iter(self._modules.values()) - - def __iadd__(self, modules: Sequence[ModuleType]) -> "InfiniCoreModuleList": - return self.extend(modules) - - def __add__( - self, other: Union[Sequence[ModuleType], "InfiniCoreModuleList"] - ) -> "InfiniCoreModuleList": - r"""Return a new InfiniCoreModuleList by concatenating with another iterable. - - Args: - other (iterable): iterable of modules to concatenate - """ - if not isinstance(other, (list, tuple, InfiniCoreModuleList)): - raise TypeError( - f"InfiniCoreModuleList can only be concatenated with list, tuple, or InfiniCoreModuleList, " - f"got {type(other).__name__}" - ) - - combined = InfiniCoreModuleList() - for i, module in enumerate(chain(self, other)): - combined.add_module(str(i), module) - return combined - - def append(self, module: ModuleType) -> "InfiniCoreModuleList": - r"""Append a given module to the end of the list. - - Args: - module (InfiniCoreModule): module to append - """ - self.add_module(str(len(self)), module) - return self - - def extend(self, modules: Sequence[ModuleType]) -> "InfiniCoreModuleList": - r"""Append modules from a Python iterable to the end of the list. - - Args: - modules (iterable): iterable of modules to append - """ - if not isinstance(modules, (list, tuple)): - try: - modules = list(modules) - except TypeError: - raise TypeError( - f"InfiniCoreModuleList.extend should be called with an " - f"iterable, but got {type(modules).__name__}" - ) - - offset = len(self) - for i, module in enumerate(modules): - self.add_module(str(offset + i), module) - return self - - def insert(self, index: int, module: ModuleType) -> None: - r"""Insert a given module before a given index in the list. - - Args: - index (int): index to insert. - module ( InfiniCoreModule): module to insert - """ - for i in range(len(self._modules), index, -1): - self._modules[str(i)] = self._modules[str(i - 1)] - self._modules[str(index)] = module - - def pop(self, idx: int = -1) -> ModuleType: - r"""Remove and return a module at the given index. - - Args: - idx (int): index of the module to pop. Default: -1 (last module) - - Returns: - Module: the module that was removed - """ - idx_str = self._get_abs_string_index(idx) - module = self._modules[idx_str] - # Use __delitem__ to ensure proper cleanup - self.__delitem__(int(idx_str)) - return module - - def __repr__(self) -> str: - """Return a string representation of the ModuleList.""" - if len(self) == 0: - return self.__class__.__name__ + "()" - - lines = [] - for i, module in enumerate(self): - lines.append(f"({i}): {repr(module)}") - - main_str = self.__class__.__name__ + "(\n " - main_str += "\n ".join(lines) + "\n)" - return main_str - - def __dir__(self) -> List[str]: - """Return a list of attribute names, excluding numeric keys.""" - keys = super().__dir__() - # Filter out numeric keys to avoid cluttering dir() output - keys = [key for key in keys if not key.isdigit()] - return keys diff --git a/python/infinicore/nn/modules/linear.py b/python/infinicore/nn/modules/linear.py deleted file mode 100644 index 2cd16e413..000000000 --- a/python/infinicore/nn/modules/linear.py +++ /dev/null @@ -1,70 +0,0 @@ -import infinicore -from infinicore.nn import functional as F - -from ...tensor import Tensor -from ..parameter import InfiniCoreParameter as Parameter -from .module import InfiniCoreModule as Module - - -class Linear(Module): - r"""Applies an affine linear transformation to the incoming data: :math:`y = xA^T + b`. - - Args: - in_features (int): size of each input sample - out_features (int): size of each output sample - bias (bool): If set to ``False``, the layer will not learn an additive bias. - Default: ``True`` - - Shape: - - Input: :(*, H_in) where :math:`*` means any number of dimensions, H_in = in_features. - - Output: :math:(*, H_out) where all but the last dimension are the same shape as the input and :math:H_out = out_features. - - Attributes: - weight(Tensor): the weights of the module of shape (out_features, in_features). - bias(Tensor): the bias of the module of shape (out_features). - """ - - __constants__ = ["in_features", "out_features"] - in_features: int - out_features: int - weight: Tensor - - def __init__( - self, - in_features: int, - out_features: int, - bias: bool = False, - device=None, - dtype=None, - ) -> None: - factory_kwargs = { - "device": infinicore.device("cpu", 0) if device is None else device, - "dtype": infinicore.float32 if dtype is None else dtype, - } - - super().__init__() - self.in_features = in_features - self.out_features = out_features - self._alpha = 1.0 - self.weight = Parameter( - infinicore.empty([out_features, in_features], **factory_kwargs) - ) - - if bias: - self.bias = Parameter(infinicore.empty([out_features], **factory_kwargs)) - else: - self.register_parameter("bias", None) - - def forward(self, input: Tensor) -> Tensor: - return F.linear(input, self.weight, self.bias, alpha=self._alpha) - - def extra_repr(self) -> str: - return f"in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}" - - @property - def alpha(self) -> float: - return self._alpha - - @alpha.setter - def alpha(self, value: float): - self._alpha = value diff --git a/python/infinicore/nn/modules/module.py b/python/infinicore/nn/modules/module.py deleted file mode 100644 index d21223903..000000000 --- a/python/infinicore/nn/modules/module.py +++ /dev/null @@ -1,857 +0,0 @@ -# Copyright (c) 2025, InfiniCore -# -# This file contains modified code derived from PyTorch's `torch.nn.Module` -# implementation, which is licensed under the BSD 3-Clause License. -# -# The modifications include adaptations for the InfiniCore framework, custom -# parameter/buffer registration mechanisms, and simplified state_dict handling. -# -# Original PyTorch source: -# https://github.com/pytorch/pytorch/blob/main/torch/nn/modules/module.py -# -# Referencing PyTorch v2.4.0 -# -# The use of this file is governed by the BSD 3-Clause License. - -import itertools -import warnings -from collections import OrderedDict, namedtuple -from typing import ( - Any, - Dict, - Iterator, - List, - Mapping, - Optional, - Set, - Tuple, - TypeVar, - Union, - overload, -) - -import infinicore - -from ...tensor import Tensor -from ..parameter import InfiniCoreParameter as Parameter - -_EXTRA_STATE_KEY_SUFFIX = "_extra_state" -T = TypeVar("T", bound="InfiniCoreModule") - - -class _IncompatibleKeys( - namedtuple("IncompatibleKeys", ["missing_keys", "unexpected_keys"]) -): - def __repr__(self): - if not self.missing_keys and not self.unexpected_keys: - return "" - return super().__repr__() - - __str__ = __repr__ - - -class InfiniCoreModule: - r"""Base class for InfiniCore neural network modules. - Your models should also subclass this class. - Modules can also contain other Modules, allowing to nest them in a tree structure. - """ - - _version: int = 1 - _parameters: Dict[str, Optional[Parameter]] - _buffers: Dict[str, Optional[Tensor]] - _non_persistent_buffers_set: Set[str] - _modules: Dict[str, Optional["InfiniCoreModule"]] - - def __init__(self): - super().__setattr__("_parameters", OrderedDict()) - super().__setattr__("_buffers", OrderedDict()) - super().__setattr__("_non_persistent_buffers_set", set()) - super().__setattr__("_modules", OrderedDict()) - - def __getattr__(self, name: str) -> Any: - if "_parameters" in self.__dict__: - _parameters = self.__dict__["_parameters"] - if name in _parameters: - return _parameters[name] - - if "_buffers" in self.__dict__: - _buffers = self.__dict__["_buffers"] - if name in _buffers: - return _buffers[name] - - if "_modules" in self.__dict__: - modules = self.__dict__["_modules"] - if name in modules: - return modules[name] - - raise AttributeError( - f"'{type(self).__name__}' object has no attribute '{name}'" - ) - - def __setattr__(self, name: str, value: Union[Tensor, "InfiniCoreModule"]) -> None: - def remove_from(*dicts_or_sets) -> None: - for d in dicts_or_sets: - if name in d: - if isinstance(d, dict): - del d[name] - else: - d.discard(name) - - params = self.__dict__.get("_parameters") - if params is None: - raise AttributeError( - "cannot assign parameters before Module.__init__() call" - ) - - if isinstance(value, Parameter): # the value is of type Parameter - remove_from( - self.__dict__, - self._buffers, - self._modules, - self._non_persistent_buffers_set, - ) - self.register_parameter(name, value) - elif name in params: # value will overwrite the name of params. - if not isinstance(value, Tensor): - raise TypeError( - f"cannot assign 'value' as parameter '{name}' (infinicore.nn.Parameter, Parameter or None expected)" - ) - self.register_parameter(name, value) - - else: - modules = self.__dict__.get("_modules") - if modules is None: - raise AttributeError( - "cannot assign module before Module.__init__() call" - ) - - if isinstance(value, InfiniCoreModule): - remove_from( - self.__dict__, - self._parameters, - self._buffers, - self._non_persistent_buffers_set, - ) - modules[name] = value - elif name in modules: # Do not overwrite this variable - raise TypeError( - f"cannot assign 'value' as child module '{name}' (infinicore.nn.Module or None expected)" - ) - else: - buffers = self.__dict__.get("_buffers") - if buffers is not None and name in buffers: - if value is not None and not isinstance(value, Tensor): - raise TypeError( - f"cannot assign 'value' as buffer '{name}' " - "(torch.Tensor or None expected)" - ) - buffers[name] = value - else: - super().__setattr__(name, value) - - def __call__(self, *input, **kwargs): - return self.forward(*input, **kwargs) - - def register_buffer( - self, name: str, tensor: Optional[Tensor], persistent: bool = True - ) -> None: - r"""Adds a buffer to the module. - - This is typically used to register a buffer that should not to be - considered a model parameter.Buffers, by default, are persistent - and will be saved alongside parameters. This behavior can be changed - by setting :attr:`persistent` to ``False``. The only difference between - a persistent buffer and a non-persistent buffer is that the latter - will not be a part of this module's :attr:`state_dict`. - - Buffers can be accessed as attributes using given names. - - Args: - name (str): name of the buffer. The buffer can be accessed - from this module using the given name - tensor (Tensor or None): buffer to be registered. If ``None``, then operations - that run on buffers, such as :attr:`cuda`, are ignored. If ``None``, - the buffer is **not** included in the module's :attr:`state_dict`. - persistent (bool): whether the buffer is part of this module's - :attr:`state_dict`. - - """ - if "_buffers" not in self.__dict__: - raise AttributeError("cannot assign buffer before Module.__init__() call") - elif not isinstance(name, str): - raise TypeError("buffer name should be a string. Got {}".format("name")) - elif "." in name: - raise KeyError('buffer name can\'t contain "."') - elif name == "": - raise KeyError('buffer name can\'t be empty string ""') - elif hasattr(self, name) and name not in self._buffers: - raise KeyError("attribute '{}' already exists".format(name)) - elif tensor is not None and not isinstance(tensor, Tensor): - raise TypeError( - "cannot assign '{}' object to buffer '{}' " - "(torch Tensor or None required)".format("tensor", name) - ) - else: - self._buffers[name] = tensor - if persistent: - self._non_persistent_buffers_set.discard(name) - else: - self._non_persistent_buffers_set.add(name) - - def add_module(self, name: str, module: Optional["InfiniCoreModule"]) -> None: - r"""Add a child module to the current module. - - The module can be accessed as an attribute using the given name. - - Args: - name (str): name of the child module. The child module can be - accessed from this module using the given name - module (Module or None): child module to be added to the module. If - ``None``, then operations that run on modules, such as :attr:`eval`, - are ignored. If ``None``, the module is **not** included in the - module's :attr:`children`. - """ - if not isinstance(name, str): - raise TypeError(f"module name should be a string. Got {name}") - elif "." in name: - raise KeyError(f'module name can\'t contain ".", got: {name}') - elif name == "": - raise KeyError('module name can\'t be empty string ""') - elif hasattr(self, name) and name not in self._modules: - raise KeyError(f"attribute '{name}' already exists") - - if module is not None and not isinstance(module, InfiniCoreModule): - raise TypeError(f"{module} is not a Module subclass") - - self._modules[name] = module - - def register_parameter(self, name: str, param: Parameter) -> None: - r"""Add a parameter to the module. - - The parameter can be accessed as an attribute using given name. - - Args: - name (str): name of the parameter. The parameter can be accessed - from this module using the given name - param (Parameter or None): parameter to be added to the module. If - ``None``, then operations that run on parameters, such as :attr:`cuda`, - are ignored. If ``None``, the parameter is **not** included in the - module's :attr:`state_dict`. - """ - - if "_parameters" not in self.__dict__: - raise AttributeError( - "cannot assign parameter before Module.__init__() call" - ) - elif not isinstance(name, str): - raise TypeError("parameter name should be a string.") - elif "." in name: - raise KeyError('parameter name can\'t contain "."') - elif name == "": - raise KeyError('parameter name can\'t be empty string ""') - elif hasattr(self, name) and name not in self._parameters: - raise KeyError(f"attribute '{name}' already exists") - - if param is None: - self._parameters[name] = None # 竟然可以是None - else: - if not isinstance(param, (Parameter, Tensor)): - raise TypeError( - f"cannot assign 'param' object to parameter '{name}' " - "(infinicore.nn.Parameter, Parameter or None required)" - ) - - self._parameters[name] = param - super().__setattr__(name, param) - - def get_extra_state(self) -> Any: - """Return any extra state to include in the module's state_dict. - - Implement this and a corresponding :func:`set_extra_state` for your module - if you need to store extra state. This function is called when building the - module's `state_dict()`. - - Note that extra state should be picklable to ensure working serialization - of the state_dict. We only provide provide backwards compatibility guarantees - for serializing Tensors; other objects may break backwards compatibility if - their serialized pickled form changes. - - Returns: - object: Any extra state to store in the module's state_dict - """ - raise RuntimeError( - "Reached a code path in Module.get_extra_state() that should never be called. " - ) - - def _save_to_state_dict(self, destination, prefix, keep_vars): - r"""Saves module state to `destination` dictionary, containing a state - of the module, but not its descendants. This is called on every - submodule in :meth:`~torch.nn.Module.state_dict`. - - In rare cases, subclasses can achieve class-specific behavior by - overriding this method with custom logic. - - Args: - destination (dict): a dict where state will be stored - prefix (str): the prefix for parameters and buffers used in this - module - """ - for name, param in self._parameters.items(): - if param is not None: - destination[prefix + name] = param if keep_vars else param - for name, buf in self._buffers.items(): - if buf is not None and name not in self._non_persistent_buffers_set: - destination[prefix + name] = buf if keep_vars else buf - extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX - if ( - getattr(self.__class__, "get_extra_state", InfiniCoreModule.get_extra_state) - is not InfiniCoreModule.get_extra_state - ): - destination[extra_state_key] = self.get_extra_state() - - # The user can pass an optional arbitrary mappable object to `state_dict`, in which case `state_dict` returns - # back that same object. But if they pass nothing, an `OrderedDict` is created and returned. - T_destination = TypeVar("T_destination", bound=Dict[str, Any]) - - @overload - def state_dict( - self, *, destination: T_destination, prefix: str = ..., keep_vars: bool = ... - ) -> T_destination: ... - - @overload - def state_dict( - self, *, prefix: str = ..., keep_vars: bool = ... - ) -> Dict[str, Any]: ... - - # TODO: Change `*args` to `*` and remove the copprespinding warning in docs when BC allows. - # Also remove the logic for arg parsing together. - def state_dict(self, *args, destination=None, prefix="", keep_vars=False): - r"""Returns a dictionary containing references to the whole state of the module. - - Both parameters and persistent buffers (e.g. running averages) are - included. Keys are corresponding parameter and buffer names. - Parameters and buffers set to ``None`` are not included. - - .. note:: - The returned object is a shallow copy. It contains references - to the module's parameters and buffers. - - .. warning:: - Currently ``state_dict()`` also accepts positional arguments for - ``destination``, ``prefix`` and ``keep_vars`` in order. However, - this is being deprecated and keyword arguments will be enforced in - future releases. - - .. warning:: - Please avoid the use of argument ``destination`` as it is not - designed for end-users. - - Args: - destination (dict, optional): If provided, the state of module will - be updated into the dict and the same object is returned. - Otherwise, an ``OrderedDict`` will be created and returned. - Default: ``None``. - prefix (str, optional): a prefix added to parameter and buffer - names to compose the keys in state_dict. Default: ``''``. - keep_vars (bool, optional): by default the :class:`~torch.Tensor` s - returned in the state dict are detached from autograd. If it's - set to ``True``, detaching will not be performed. - Default: ``False``. - - Returns: - dict: - a dictionary containing a whole state of the module - - Example:: - - >>> # xdoctest: +SKIP("undefined vars") - >>> module.state_dict().keys() - ['bias', 'weight'] - - """ - - # TODO: Remove `args` and the parsing logic when BC allows. - if len(args) > 0: - # DeprecationWarning is ignored by default - warnings.warn( - "Positional args are being deprecated, use kwargs instead. ", - FutureWarning, - stacklevel=2, - ) - if destination is None: - destination = args[0] - if len(args) > 1 and prefix == "": - prefix = args[1] - if len(args) > 2 and keep_vars is False: - keep_vars = args[2] - - if destination is None: - destination = OrderedDict() - destination._metadata = OrderedDict() - - local_metadata = dict(version=self._version) - if hasattr(destination, "_metadata"): - destination._metadata[prefix[:-1]] = local_metadata - - self._save_to_state_dict(destination, prefix, keep_vars) - for name, module in self._modules.items(): - if module is not None: - module.state_dict( - destination=destination, - prefix=prefix + name + ".", - keep_vars=keep_vars, - ) - return destination - - def set_extra_state(self, state: Any): - """ - This function is called from :func:`load_state_dict` to handle any extra state - found within the `state_dict`. Implement this function and a corresponding - :func:`get_extra_state` for your module if you need to store extra state within its - `state_dict`. - - Args: - state (dict): Extra state from the `state_dict` - """ - raise RuntimeError( - "Reached a code path in Module.set_extra_state() that should never be called. " - "Please file an issue at https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml " - "to report this bug." - ) - - def _load_from_state_dict( - self, - state_dict, - prefix, - local_metadata, - strict, - missing_keys, - unexpected_keys, - error_msgs, - ): - r"""Copies parameters and buffers from :attr:`state_dict` into only - this module, but not its descendants. This is called on every submodule - in :meth:`~torch.nn.Module.load_state_dict`. Metadata saved for this - module in input :attr:`state_dict` is provided as :attr:`local_metadata`. - For state dicts without metadata, :attr:`local_metadata` is empty. - Subclasses can achieve class-specific backward compatible loading using - the version number at `local_metadata.get("version", None)`. - - .. note:: - :attr:`state_dict` is not the same object as the input - :attr:`state_dict` to :meth:`~torch.nn.Module.load_state_dict`. So - it can be modified. - - Args: - state_dict (dict): a dict containing parameters and - persistent buffers. - prefix (str): the prefix for parameters and buffers used in this - module - local_metadata (dict): a dict containing the metadata for this module. - See - strict (bool): whether to strictly enforce that the keys in - :attr:`state_dict` with :attr:`prefix` match the names of - parameters and buffers in this module - missing_keys (list of str): if ``strict=True``, add missing keys to - this list - unexpected_keys (list of str): if ``strict=True``, add unexpected - keys to this list - error_msgs (list of str): error messages should be added to this - list, and will be reported together in - :meth:`~torch.nn.Module.load_state_dict` - """ - persistent_buffers = { - k: v - for k, v in self._buffers.items() - if k not in self._non_persistent_buffers_set - } - local_name_params = itertools.chain( - self._parameters.items(), persistent_buffers.items() - ) - local_state = {k: v for k, v in local_name_params if v is not None} - - for name, param in local_state.items(): - key = prefix + name - if key in state_dict: - input_param = state_dict[key] - - # input_param must be of type infinicore.Tensor - if not isinstance(input_param, Tensor): - raise TypeError( - f"While copying the parameter named {key}, expected Tensor from checkpoint but received {type(input_param)}" - ) - - if ( - (param.shape == input_param.shape) - and (param.dtype == input_param.dtype) - and (param.device == input_param.device) - ): - param.copy_(input_param) - else: - print(f"param '{name}' don't match input_param '{key}'") - setattr(self, name, input_param) - - elif strict: - missing_keys.append(key) - - extra_state_key = prefix + _EXTRA_STATE_KEY_SUFFIX - if ( - getattr(self.__class__, "set_extra_state", InfiniCoreModule.set_extra_state) - is not InfiniCoreModule.set_extra_state - ): - if extra_state_key in state_dict: - self.set_extra_state(state_dict[extra_state_key]) - elif strict: - missing_keys.append(extra_state_key) - elif strict and (extra_state_key in state_dict): - unexpected_keys.append(extra_state_key) - - if strict: - for key in state_dict.keys(): - if key.startswith(prefix): - input_name = key[len(prefix) :].split(".", 1) - # Must be Module if it have attributes - if len(input_name) > 1: - if input_name[0] not in self._modules: - unexpected_keys.append(key) - elif input_name[0] not in local_state: - unexpected_keys.append(key) - - def load_state_dict(self, state_dict: Mapping[str, Any], strict: bool = True): - r"""Copies parameters and buffers from :attr:`state_dict` into - this module and its descendants. If :attr:`strict` is ``True``, then - the keys of :attr:`state_dict` must exactly match the keys returned - by this module's :meth:`~torch.nn.Module.state_dict` function. - - Args: - state_dict (dict): a dict containing parameters and - persistent buffers. - strict (bool, optional): whether to strictly enforce that the keys - in :attr:`state_dict` match the keys returned by this module's - :meth:`~torch.nn.Module.state_dict` function. Default: ``True`` - - Returns: - ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields: - * **missing_keys** is a list of str containing the missing keys - * **unexpected_keys** is a list of str containing the unexpected keys - - Note: - If a parameter or buffer is registered as ``None`` and its corresponding key - exists in :attr:`state_dict`, :meth:`load_state_dict` will raise a - ``RuntimeError``. - """ - if not isinstance(state_dict, Mapping): - raise TypeError( - "Expected state_dict to be dict-like, got {}.".format(type(state_dict)) - ) - - missing_keys: List[str] = [] - unexpected_keys: List[str] = [] - error_msgs: List[str] = [] - - # copy state_dict so _load_from_state_dict can modify it - metadata = getattr(state_dict, "_metadata", None) - state_dict = OrderedDict(state_dict) - if metadata is not None: - state_dict._metadata = metadata # type: ignore[attr-defined] - - def load(module, local_state_dict, prefix=""): - local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {}) - module._load_from_state_dict( - local_state_dict, - prefix, - local_metadata, - True, - missing_keys, - unexpected_keys, - error_msgs, - ) - for name, child in module._modules.items(): - if child is not None: - child_prefix = prefix + name + "." - child_state_dict = { - k: v - for k, v in local_state_dict.items() - if k.startswith(child_prefix) - } - load(child, child_state_dict, child_prefix) # noqa: F821 - - load(self, state_dict) - del load - - if strict: - if len(unexpected_keys) > 0: - error_msgs.insert( - 0, - "Unexpected key(s) in state_dict: {}. ".format( - ", ".join('"{}"'.format(k) for k in unexpected_keys) - ), - ) - if len(missing_keys) > 0: - error_msgs.insert( - 0, - "Missing key(s) in state_dict: {}. ".format( - ", ".join('"{}"'.format(k) for k in missing_keys) - ), - ) - - if len(error_msgs) > 0: - raise RuntimeError( - "Error(s) in loading state_dict for {}:\n\t{}".format( - self.__class__.__name__, "\n\t".join(error_msgs) - ) - ) - return _IncompatibleKeys(missing_keys, unexpected_keys) - - def parameters(self, recurse: bool = True) -> Iterator["Parameter"]: - r"""Returns an iterator over module parameters. - - Args: - recurse (bool): if True, then yields parameters of this module - and all submodules. Otherwise, yields only parameters that - are direct members of this module. - - Yields: - Parameter: module parameter - - Example:: - - >>> # xdoctest: +SKIP("undefined vars") - >>> for param in model.parameters(): - ... print(type(param), param.size()) - - """ - for name, param in self.named_parameters(recurse=recurse): - yield param - - def named_parameters( - self, prefix: str = "", recurse: bool = True - ) -> Iterator[Tuple[str, "Parameter"]]: - r"""Returns an iterator over module parameters, yielding both the - name of the parameter as well as the parameter itself. - - Args: - prefix (str): prefix to prepend to all parameter names. - recurse (bool): if True, then yields parameters of this module - and all submodules. Otherwise, yields only parameters that - are direct members of this module. - - Yields: - (str, Parameter): Tuple containing the name and parameter - - Example:: - - >>> # xdoctest: +SKIP("undefined vars") - >>> for name, param in self.named_parameters(): - ... if name in ['bias']: - ... print(param.size()) - - """ - gen = self._named_members( - lambda module: module._parameters.items(), prefix=prefix, recurse=recurse - ) - for elem in gen: - yield elem - - def buffers(self, recurse: bool = True) -> Iterator[Tensor]: - r"""Returns an iterator over module buffers. - - Args: - recurse (bool): if True, then yields buffers of this module - and all submodules. Otherwise, yields only buffers that - are direct members of this module. - - Yields: - torch.Tensor: module buffer - - Example:: - - >>> # xdoctest: +SKIP("undefined vars") - >>> for buf in model.buffers(): - ... print(type(buf), buf.size()) - - """ - for name, buf in self.named_buffers(recurse=recurse): - yield buf - - def named_buffers( - self, prefix: str = "", recurse: bool = True - ) -> Iterator[Tuple[str, Tensor]]: - r"""Returns an iterator over module buffers, yielding both the - name of the buffer as well as the buffer itself. - - Args: - prefix (str): prefix to prepend to all buffer names. - recurse (bool): if True, then yields buffers of this module - and all submodules. Otherwise, yields only buffers that - are direct members of this module. - - Yields: - (str, torch.Tensor): Tuple containing the name and buffer - - Example:: - - >>> # xdoctest: +SKIP("undefined vars") - >>> for name, buf in self.named_buffers(): - ... if name in ['running_mean']: - ... print(buf.size()) - - """ - memo = set() - modules = self.named_modules(prefix=prefix) if recurse else [(prefix, self)] - for module_prefix, module in modules: - for k, v in module._buffers.items(): - if v is None or v in memo: - continue - if k in module._non_persistent_buffers_set: - continue - memo.add(v) - name = module_prefix + ("." if module_prefix else "") + k - yield (name, v) - - def _named_members(self, get_members_fn, prefix="", recurse=True): - r"""Helper method to yield members with their names.""" - memo = set() - modules = self.named_modules(prefix=prefix) if recurse else [(prefix, self)] - for module_prefix, module in modules: - members = get_members_fn(module) - for k, v in members: - if v is None or v in memo: - continue - memo.add(v) - name = module_prefix + ("." if module_prefix else "") + k - yield (name, v) - - def modules(self) -> Iterator["InfiniCoreModule"]: - r"""Returns an iterator over all modules in the network. - - Yields: - Module: a module in the network - - Note: - Duplicate modules are returned only once. In the following - example, ``l`` will be returned only once. - - Example:: - - >>> # xdoctest: +SKIP("undefined vars") - >>> l = nn.Linear(2, 2) - >>> net = nn.Sequential(l, l) - >>> for idx, m in enumerate(net.modules()): - ... print(idx, '->', m) - - 0 -> Sequential( - (0): Linear(in_features=2, out_features=2, bias=True) - (1): Linear(in_features=2, out_features=2, bias=True) - ) - 1 -> Linear(in_features=2, out_features=2, bias=True) - - """ - for name, module in self.named_modules(): - yield module - - def named_modules( - self, - memo: Optional[Set["InfiniCoreModule"]] = None, - prefix: str = "", - remove_duplicate: bool = True, - ): - r"""Returns an iterator over all modules in the network, yielding - both the name of the module as well as the module itself. - - Args: - memo: a memo to store the set of modules already added to the result - prefix: a prefix that will be added to the name of the module - remove_duplicate: whether to remove the duplicated module instances in the result - or not - - Yields: - (str, Module): Tuple of name and module - - Note: - Duplicate modules are returned only once. In the following - example, ``l`` will be returned only once. - - Example:: - - >>> # xdoctest: +SKIP("undefined vars") - >>> l = nn.Linear(2, 2) - >>> net = nn.Sequential(l, l) - >>> for idx, m in enumerate(net.named_modules()): - ... print(idx, '->', m) - - 0 -> ('', Sequential( - (0): Linear(in_features=2, out_features=2, bias=True) - (1): Linear(in_features=2, out_features=2, bias=True) - )) - 1 -> ('0', Linear(in_features=2, out_features=2, bias=True)) - - """ - if memo is None: - memo = set() - if remove_duplicate: - if self in memo: - return - memo.add(self) - yield prefix, self - for name, module in self._modules.items(): - if module is None: - continue - submodule_prefix = prefix + ("." if prefix else "") + name - # Handle both InfiniCoreModule and torch.nn.Module - if isinstance(module, InfiniCoreModule): - for m in module.named_modules(memo, submodule_prefix, remove_duplicate): - yield m - elif isinstance(module, infinicore.nn.Module): - # For torch.nn.Module, use its named_modules method - # torch.nn.Module.named_modules returns (name, module) tuples - for sub_name, sub_module in module.named_modules( - prefix=submodule_prefix, remove_duplicate=remove_duplicate - ): - yield (sub_name, sub_module) - - def children(self) -> Iterator["InfiniCoreModule"]: - r"""Returns an iterator over immediate children modules. - - Yields: - Module: a child module (can be InfiniCoreModule or torch.nn.Module) - """ - for name, module in self.named_children(): - yield module - - def named_children( - self, - ) -> Iterator[Tuple[str, "InfiniCoreModule"]]: - r"""Returns an iterator over immediate children modules, yielding both - the name of the module as well as the module itself. - - Yields: - (str, Module): Tuple containing a name and child module - - Example:: - - >>> # xdoctest: +SKIP("undefined vars") - >>> for name, module in model.named_children(): - >>> if name in ['conv4', 'conv5']: - >>> print(module) - - """ - memo = set() - for name, module in self._modules.items(): - if module is not None and module not in memo: - memo.add(module) - yield name, module - - def eval(self: T) -> T: - r"""Sets the module in evaluation mode. - - Returns: - Module: self - """ - pass - - def _apply(self, fn, recurse=True): - raise KeyError("not support") - - def to(self, *args, **kwargs): - raise KeyError("not support") diff --git a/python/infinicore/nn/modules/normalization.py b/python/infinicore/nn/modules/normalization.py deleted file mode 100644 index 4052b7ddb..000000000 --- a/python/infinicore/nn/modules/normalization.py +++ /dev/null @@ -1,59 +0,0 @@ -import numbers -from typing import Optional, Union - -import infinicore -from infinicore.nn import functional as F - -from ...tensor import Tensor -from ..parameter import InfiniCoreParameter as Parameter -from .module import InfiniCoreModule as Module - - -class RMSNorm(Module): - r"""Applies Root Mean Square Layer Normalization over a mini-batch of inputs. - The RMS is taken over the last dimensions, where normalized_shape is one-dimensional. - - Args: - normalized_shape (int or list): input shape from an expected input of size [* , normalized_shape[0]] - this module will normalize over the last dimension. - eps (float): a value added to the denominator for numerical stability. - - Shape: - - Input: (N, *) - - Output: (N, *) (same shape as input) - """ - - __constants__ = ["normalized_shape", "eps"] - normalized_shape: tuple[int] - eps: Optional[float] - - def __init__( - self, - normalized_shape: Union[int, list[int]], - eps=1e-6, - elementwise_affine=True, - device=None, - dtype=None, - ) -> None: - factory_kwargs = { - "dtype": infinicore.float32 if dtype is None else dtype, - "device": infinicore.device("cpu", 0) if device is None else device, - } - super().__init__() - - assert elementwise_affine, "elementwise_affine must be true." - - if isinstance(normalized_shape, numbers.Integral): - normalized_shape = [normalized_shape] - - self.normalized_shape = list(normalized_shape) - self.eps = eps - self.weight = Parameter( - infinicore.empty(self.normalized_shape, **factory_kwargs) - ) - - def forward(self, x: Tensor) -> Tensor: - return F.rms_norm(x, self.normalized_shape, self.weight, self.eps) - - def extra_repr(self) -> str: - return "{normalized_shape}, eps={eps}".format(**self.__dict__) diff --git a/python/infinicore/nn/modules/rope.py b/python/infinicore/nn/modules/rope.py deleted file mode 100644 index 1260b295c..000000000 --- a/python/infinicore/nn/modules/rope.py +++ /dev/null @@ -1,157 +0,0 @@ -import numpy as np - -import infinicore -from infinicore.nn import functional as F -from infinicore.ops.mrope import mrope - -from ...tensor import Tensor -from ..functional import RopeAlgo -from .module import InfiniCoreModule as Module - - -def create_sin_cos_table_numpy(max_position, rotary_dim, theta=10000.0): - assert rotary_dim % 2 == 0, "Embedding dimension must be even." - pos = np.arange(0, max_position) - freqs = 1.0 / ( - theta - ** (np.arange(0, rotary_dim, 2)[: (rotary_dim // 2)].astype(float) / rotary_dim) - ) - angles = np.outer(pos, freqs) - sin_table = np.sin(angles, dtype=np.float32) - cos_table = np.cos(angles, dtype=np.float32) - return sin_table, cos_table - - -def create_sin_cos_table( - max_position, - rotary_dim, - theta=10000.0, - device=None, - dtype=None, -): - sin_table_np, cos_table_np = create_sin_cos_table_numpy( - max_position, rotary_dim, theta - ) - - sin_table_infini = infinicore.from_numpy(sin_table_np, dtype=dtype, device=device) - cos_table_infini = infinicore.from_numpy(cos_table_np, dtype=dtype, device=device) - - return sin_table_infini, cos_table_infini - - -class RoPE(Module): - r"""Rotary Position Embedding(RoPE). - - Standard RoPE is used when ``mrope_section`` is None. MRoPE is enabled by passing - ``mrope_section=[t, h, w]`` and then calling the module as ``rope(q, k, positions)``. - """ - - __constants__ = [ - "max_position_embeddings", - "rope_theta", - "head_dim", - "rotary_dim", - "mrope_section", - "mrope_interleaved", - ] - max_position_embeddings: int - rope_theta: float - head_dim: int - rotary_dim: int - - def __init__( - self, - max_position_embeddings: int, - rope_theta: float, - head_dim: int, - device=None, - dtype=None, - rotary_dim: int | None = None, - mrope_section: list[int] | tuple[int, int, int] | None = None, - mrope_interleaved: bool = False, - ): - factory_kwargs = { - "device": infinicore.device("cpu", 0) if device is None else device, - "dtype": infinicore.float32 if dtype is None else dtype, - } - super().__init__() - - self.max_position_embeddings = max_position_embeddings - self.rope_theta = rope_theta - self.head_dim = head_dim - self.rotary_dim = head_dim if rotary_dim is None else rotary_dim - self.mrope_section = None if mrope_section is None else list(mrope_section) - self.mrope_interleaved = mrope_interleaved - - if ( - self.rotary_dim <= 0 - or self.rotary_dim > self.head_dim - or self.rotary_dim % 2 != 0 - ): - raise ValueError("rotary_dim must be positive, even, and <= head_dim") - if self.mrope_section is not None and ( - len(self.mrope_section) != 3 - or 2 * sum(self.mrope_section) != self.rotary_dim - ): - raise ValueError( - "mrope_section must contain 3 values and sum to rotary_dim / 2" - ) - - self._sin_table, self._cos_table = create_sin_cos_table( - self.max_position_embeddings, - rotary_dim=self.rotary_dim, - theta=self.rope_theta, - **factory_kwargs, - ) - - def forward( - self, - states: Tensor, - position_ids: Tensor, - *args, - algo=RopeAlgo.GPT_NEOX, - out=None, - ): - if args: - q = states - k = position_ids - positions = args[0] - if self.mrope_section is not None: - return mrope( - q, - k, - self._cos_table, - self._sin_table, - positions, - self.head_dim, - self.rotary_dim, - self.mrope_section[0], - self.mrope_section[1], - self.mrope_section[2], - self.mrope_interleaved, - out=out, - ) - - if out is None: - q_out = infinicore.empty(q.shape, dtype=q.dtype, device=q.device) - k_out = infinicore.empty(k.shape, dtype=k.dtype, device=k.device) - else: - q_out, k_out = out - F.rope(q, positions, self._sin_table, self._cos_table, algo=algo, out=q_out) - F.rope(k, positions, self._sin_table, self._cos_table, algo=algo, out=k_out) - return q_out, k_out - - if self.mrope_section is not None: - raise NotImplementedError( - "MRoPE single-tensor forward is not implemented; use fused forward(q, k, positions) instead" - ) - target = states if out is None else out - F.rope( - states, - position_ids, - self._sin_table, - self._cos_table, - algo=algo, - out=target, - ) - return target diff --git a/python/infinicore/nn/modules/sparse.py b/python/infinicore/nn/modules/sparse.py deleted file mode 100644 index 0fb8a4c69..000000000 --- a/python/infinicore/nn/modules/sparse.py +++ /dev/null @@ -1,77 +0,0 @@ -import infinicore -from infinicore.nn import functional as F - -from ...tensor import Tensor -from ..parameter import InfiniCoreParameter as Parameter -from .module import InfiniCoreModule as Module - - -class Embedding(Module): - r"""A simple lookup table that stores embeddings of a fixed dictionary and size. - - This module is often used to store word embeddings and retrieve them using indices. - The input to the module is a list of indices, and the output is the corresponding - word embeddings. - - Args: - num_embeddings (int): size of the dictionary of embeddings. - embedding_dim (int): the size of each embedding vector. - - Attributes: - weight (Tensor): the weights of the module of shape (num_embeddings, embedding_dim). - - Shape: - - Input: :(*), IntTensor or LongTensor of arbitrary shape containing the indices to extract. - - Output: (*, H), where `*` is the input shape and H=embedding_dim. - """ - - __constants__ = [ - "num_embeddings", - "embedding_dim", - ] - - num_embeddings: int - embedding_dim: int - weight: Tensor - - def __init__( - self, - num_embeddings: int, - embedding_dim: int, - padding_idx=None, - max_norm=None, - norm_type=2.0, - scale_grad_by_freq=False, - sparse=False, - _weight=None, - _freeze=False, - device=None, - dtype=None, - ) -> None: - factory_kwargs = { - "dtype": infinicore.float32 if dtype is None else dtype, - "device": infinicore.device("cpu", 0) if device is None else device, - } - super().__init__() - assert ( - (padding_idx is None) - and (max_norm is None) - and (scale_grad_by_freq is False) - and (sparse is False) - and (_weight is None) - and (_freeze is False) - ), "Unsupported parameters." - - self.num_embeddings = num_embeddings - self.embedding_dim = embedding_dim - - self.weight = Parameter( - infinicore.empty([num_embeddings, embedding_dim], **factory_kwargs) - ) - - def forward(self, input: Tensor) -> Tensor: - return F.embedding(input, self.weight) - - def extra_repr(self) -> str: - s = "{num_embeddings}, {embedding_dim}" - return s.format(**self.__dict__) diff --git a/python/infinicore/nn/parameter.py b/python/infinicore/nn/parameter.py deleted file mode 100644 index 598f90c1d..000000000 --- a/python/infinicore/nn/parameter.py +++ /dev/null @@ -1,34 +0,0 @@ -# Copyright (c) 2025, InfiniCore -# -# This file contains modified code derived from PyTorch's `torch.nn.Parameter` -# implementation, which is licensed under the BSD 3-Clause License. -# -# The modifications include adaptations for the InfiniCore framework. -# -# Original PyTorch source: -# https://github.com/pytorch/pytorch/blob/main/torch/nn/parameter.py -# -# Referencing PyTorch v2.4.0 -# -# The use of this file is governed by the BSD 3-Clause License. - - -from ..tensor import Tensor - - -class InfiniCoreParameter(Tensor): - r"""A kind of Tensor that is to be considered a module parameter.""" - - def __init__(self, data=None): - if not isinstance(data, Tensor): - raise ValueError("The `data` variable must be of type `infinicore.Tensor`.") - super().__init__(data._underlying) - - def __repr__(self): - return "Parameter containing:\n" + super().__repr__() - - def __deepcopy__(self, memo): - raise ValueError("not supported!") - - def __reduce_ex__(self, proto): - raise ValueError("not supported!") diff --git a/python/infinicore/ops/__init__.py b/python/infinicore/ops/__init__.py deleted file mode 100644 index 8b1378917..000000000 --- a/python/infinicore/ops/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/python/infinicore/ops/acos.py b/python/infinicore/ops/acos.py deleted file mode 100644 index 9027ac64f..000000000 --- a/python/infinicore/ops/acos.py +++ /dev/null @@ -1,10 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def acos(input, *, out=None): - if out is None: - return Tensor(_infinicore.acos(input._underlying)) - _infinicore.acos_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/ops/add.py b/python/infinicore/ops/add.py deleted file mode 100644 index ea3bc07e3..000000000 --- a/python/infinicore/ops/add.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def add(input, other, *, out=None): - if out is None: - return Tensor(_infinicore.add(input._underlying, other._underlying)) - - _infinicore.add_(out._underlying, input._underlying, other._underlying) - - return out diff --git a/python/infinicore/ops/add_rms_norm.py b/python/infinicore/ops/add_rms_norm.py deleted file mode 100644 index a5de7bd92..000000000 --- a/python/infinicore/ops/add_rms_norm.py +++ /dev/null @@ -1,34 +0,0 @@ -import infinicore.tensor as tensor -from infinicore.lib import _infinicore - - -def add_rms_norm(a, b, weight, epsilon=1e-5, *, out=None, residual=None): - """ - Fused Add and RMS Normalization. - - Args: - a: First input tensor - b: Second input tensor - weight: Scale weights - epsilon: Small constant for numerical stability, default is 1e-5 - out: Optional output tuple (y, residual_out) for in-place operation - - Returns: - Tuple of (normalized_result, add_result): (RMSNorm(a + b) * weight, a + b) - The add_result can be used as residual for subsequent layers. - """ - if out is None: - out = tensor.empty(a.shape, dtype=a.dtype, device=a.device) - if residual is None: - residual = tensor.empty(b.shape, dtype=b.dtype, device=b.device) - - _infinicore.add_rms_norm_( - out._underlying, - residual._underlying, - a._underlying, - b._underlying, - weight._underlying, - epsilon, - ) - - return out, residual diff --git a/python/infinicore/ops/addbmm.py b/python/infinicore/ops/addbmm.py deleted file mode 100644 index 7bb4e69f2..000000000 --- a/python/infinicore/ops/addbmm.py +++ /dev/null @@ -1,24 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def addbmm(input, batch1, batch2, *, beta=1.0, alpha=1.0, out=None): - # 1. Out-of-place 模式 (如果没有指定 out) - if out is None: - return Tensor( - _infinicore.addbmm( - input._underlying, batch1._underlying, batch2._underlying, beta, alpha - ) - ) - - # 2. In-place 模式 (指定了 out) - _infinicore.addbmm_( - out._underlying, - input._underlying, - batch1._underlying, - batch2._underlying, - beta, - alpha, - ) - - return out diff --git a/python/infinicore/ops/addcmul.py b/python/infinicore/ops/addcmul.py deleted file mode 100644 index 03cf06816..000000000 --- a/python/infinicore/ops/addcmul.py +++ /dev/null @@ -1,24 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def addcmul(input, tensor1, tensor2, value=1.0, *, out=None): - if out is None: - return Tensor( - _infinicore.addcmul( - input._underlying, - tensor1._underlying, - tensor2._underlying, - float(value), - ) - ) - - _infinicore.addcmul_( - out._underlying, - input._underlying, - tensor1._underlying, - tensor2._underlying, - float(value), - ) - - return out diff --git a/python/infinicore/ops/addr.py b/python/infinicore/ops/addr.py deleted file mode 100644 index f9a607961..000000000 --- a/python/infinicore/ops/addr.py +++ /dev/null @@ -1,28 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def addr( - input: Tensor, - vec1: Tensor, - vec2: Tensor, - beta: float = 1.0, - alpha: float = 1.0, - out=None, -) -> Tensor: - if out is None: - return Tensor( - _infinicore.addr( - input._underlying, vec1._underlying, vec2._underlying, beta, alpha - ) - ) - - _infinicore.addr_( - out._underlying, - input._underlying, - vec1._underlying, - vec2._underlying, - beta, - alpha, - ) - return out diff --git a/python/infinicore/ops/all.py b/python/infinicore/ops/all.py deleted file mode 100644 index 6aacd519d..000000000 --- a/python/infinicore/ops/all.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def all(input, dim=None, keepdim=False, out=None): - if out is None: - return Tensor(_infinicore.all(input._underlying, dim, keepdim)) - - _infinicore.all_(out._underlying, input._underlying, dim, keepdim) - - return out diff --git a/python/infinicore/ops/argwhere.py b/python/infinicore/ops/argwhere.py deleted file mode 100644 index 1e558c898..000000000 --- a/python/infinicore/ops/argwhere.py +++ /dev/null @@ -1,6 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def argwhere(x: Tensor) -> Tensor: - return Tensor(_infinicore.argwhere(x._underlying)) diff --git a/python/infinicore/ops/asin.py b/python/infinicore/ops/asin.py deleted file mode 100644 index fbb230737..000000000 --- a/python/infinicore/ops/asin.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def asin(input: Tensor, *, out=None): - """Arcsin activation function.""" - if out is None: - return Tensor(_infinicore.asin(input._underlying)) - - _infinicore.asin_(out._underlying, input._underlying) - return out diff --git a/python/infinicore/ops/asinh.py b/python/infinicore/ops/asinh.py deleted file mode 100644 index 05ec58779..000000000 --- a/python/infinicore/ops/asinh.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def asinh(input, *, out=None): - if out is None: - return Tensor(_infinicore.asinh(input._underlying)) - - _infinicore.asinh_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/ops/asum.py b/python/infinicore/ops/asum.py deleted file mode 100644 index 589a02129..000000000 --- a/python/infinicore/ops/asum.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def asum(x: Tensor, *, out=None): - if out is None: - return Tensor(_infinicore.asum(x._underlying)) - - _infinicore.asum_(x._underlying, out._underlying) - - return out diff --git a/python/infinicore/ops/atanh.py b/python/infinicore/ops/atanh.py deleted file mode 100644 index 4008e0268..000000000 --- a/python/infinicore/ops/atanh.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def atanh(input, *, out=None): - if out is None: - return Tensor(_infinicore.atanh(input._underlying)) - - _infinicore.atanh_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/ops/attention.py b/python/infinicore/ops/attention.py deleted file mode 100644 index c4b070484..000000000 --- a/python/infinicore/ops/attention.py +++ /dev/null @@ -1,28 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def attention(q, k, v, k_cache, v_cache, pos, *, out=None): - if out is None: - return Tensor( - _infinicore.attention( - q._underlying, - k._underlying, - v._underlying, - k_cache._underlying, - v_cache._underlying, - pos, - ) - ) - - _infinicore.attention_( - out._underlying, - q._underlying, - k._underlying, - v._underlying, - k_cache._underlying, - v_cache._underlying, - pos, - ) - - return out diff --git a/python/infinicore/ops/axpy.py b/python/infinicore/ops/axpy.py deleted file mode 100644 index 3457038fb..000000000 --- a/python/infinicore/ops/axpy.py +++ /dev/null @@ -1,7 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def axpy(alpha: Tensor, x: Tensor, y: Tensor): - _infinicore.axpy_(alpha._underlying, x._underlying, y._underlying) - return y diff --git a/python/infinicore/ops/baddbmm.py b/python/infinicore/ops/baddbmm.py deleted file mode 100644 index 4a34cbb64..000000000 --- a/python/infinicore/ops/baddbmm.py +++ /dev/null @@ -1,25 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def baddbmm(input, batch1, batch2, *, beta=1.0, alpha=1.0, out=None): - if out is None: - return Tensor( - _infinicore.baddbmm( - input._underlying, - batch1._underlying, - batch2._underlying, - float(beta), - float(alpha), - ) - ) - _infinicore.baddbmm_( - out._underlying, - input._underlying, - batch1._underlying, - batch2._underlying, - float(beta), - float(alpha), - ) - - return out diff --git a/python/infinicore/ops/bilinear.py b/python/infinicore/ops/bilinear.py deleted file mode 100644 index 4773dd825..000000000 --- a/python/infinicore/ops/bilinear.py +++ /dev/null @@ -1,23 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def bilinear(input1, input2, weight, bias=None, *, out=None): - if out is None: - return Tensor( - _infinicore.bilinear( - input1._underlying, - input2._underlying, - weight._underlying, - bias._underlying if bias is not None else None, - ) - ) - _infinicore.bilinear_( - out._underlying, - input1._underlying, - input2._underlying, - weight._underlying, - bias._underlying if bias is not None else None, - ) - - return out diff --git a/python/infinicore/ops/binary_cross_entropy_with_logits.py b/python/infinicore/ops/binary_cross_entropy_with_logits.py deleted file mode 100644 index b9b72b7ac..000000000 --- a/python/infinicore/ops/binary_cross_entropy_with_logits.py +++ /dev/null @@ -1,42 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def binary_cross_entropy_with_logits( - input, target, weight=None, pos_weight=None, reduction="mean", *, out=None -): - """ - input: Tensor (logits) - target: Tensor (labels) - weight: Tensor (optional, sample-wise weight) - pos_weight: Tensor (optional, class-wise weight) - reduction: str ('none', 'mean', 'sum') - """ - - # 提取底层 C++ 对象，处理可选 Tensor - weight_raw = weight._underlying if weight is not None else None - pos_weight_raw = pos_weight._underlying if pos_weight is not None else None - - if out is None: - # 调用非原地接口，返回新创建的 Tensor - return Tensor( - _infinicore.binary_cross_entropy_with_logits( - input._underlying, - target._underlying, - weight_raw, - pos_weight_raw, - str(reduction), - ) - ) - - # 调用显式指定输出的接口 (binary_cross_entropy_with_logits_) - _infinicore.binary_cross_entropy_with_logits_( - out._underlying, - input._underlying, - target._underlying, - weight_raw, - pos_weight_raw, - str(reduction), - ) - - return out diff --git a/python/infinicore/ops/bitwise_right_shift.py b/python/infinicore/ops/bitwise_right_shift.py deleted file mode 100644 index 714c6c777..000000000 --- a/python/infinicore/ops/bitwise_right_shift.py +++ /dev/null @@ -1,16 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def bitwise_right_shift( - input: Tensor, other: Tensor, *, out: Tensor | None = None -) -> Tensor: - if out is None: - return Tensor( - _infinicore.bitwise_right_shift(input._underlying, other._underlying) - ) - - _infinicore.bitwise_right_shift_( - out._underlying, input._underlying, other._underlying - ) - return out diff --git a/python/infinicore/ops/blas_amax.py b/python/infinicore/ops/blas_amax.py deleted file mode 100644 index 65279963d..000000000 --- a/python/infinicore/ops/blas_amax.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def blas_amax(x: Tensor, *, out=None): - if out is None: - return Tensor(_infinicore.blas_amax(x._underlying)) - - _infinicore.blas_amax_(x._underlying, out._underlying) - - return out diff --git a/python/infinicore/ops/blas_amin.py b/python/infinicore/ops/blas_amin.py deleted file mode 100644 index 472313faa..000000000 --- a/python/infinicore/ops/blas_amin.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def blas_amin(x: Tensor, *, out=None): - if out is None: - return Tensor(_infinicore.blas_amin(x._underlying)) - - _infinicore.blas_amin_(x._underlying, out._underlying) - - return out diff --git a/python/infinicore/ops/blas_copy.py b/python/infinicore/ops/blas_copy.py deleted file mode 100644 index 75d4abef3..000000000 --- a/python/infinicore/ops/blas_copy.py +++ /dev/null @@ -1,7 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def blas_copy(x: Tensor, y: Tensor): - _infinicore.blas_copy_(x._underlying, y._underlying) - return y diff --git a/python/infinicore/ops/blas_dot.py b/python/infinicore/ops/blas_dot.py deleted file mode 100644 index 60f0541b0..000000000 --- a/python/infinicore/ops/blas_dot.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def blas_dot(x: Tensor, y: Tensor, *, out=None): - if out is None: - return Tensor(_infinicore.blas_dot(x._underlying, y._underlying)) - - _infinicore.blas_dot_(x._underlying, y._underlying, out._underlying) - - return out diff --git a/python/infinicore/ops/block_diag.py b/python/infinicore/ops/block_diag.py deleted file mode 100644 index ac52c73fc..000000000 --- a/python/infinicore/ops/block_diag.py +++ /dev/null @@ -1,12 +0,0 @@ -from __future__ import annotations - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def block_diag(*tensors: Tensor) -> Tensor: - if not tensors: - raise ValueError("block_diag expects at least one input tensor") - - underlying = [t._underlying for t in tensors] - return Tensor(_infinicore.block_diag(underlying)) diff --git a/python/infinicore/ops/broadcast_to.py b/python/infinicore/ops/broadcast_to.py deleted file mode 100644 index e425bdd80..000000000 --- a/python/infinicore/ops/broadcast_to.py +++ /dev/null @@ -1,12 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -# 修改说明：将参数名 'shape' 改为 'size'，以匹配测试用例中的调用方式 kwargs={size=...} -def broadcast_to(input, size, *, out=None): - if out is None: - return Tensor(_infinicore.broadcast_to(input._underlying, size)) - - _infinicore.broadcast_to_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/ops/cat.py b/python/infinicore/ops/cat.py deleted file mode 100644 index e45e9294a..000000000 --- a/python/infinicore/ops/cat.py +++ /dev/null @@ -1,12 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def cat(tensors, dim: int = 0, *, out=None): - if out is None: - return Tensor(_infinicore.cat([tensor._underlying for tensor in tensors], dim)) - - _infinicore.cat_(out._underlying, [tensor._underlying for tensor in tensors], dim) - - # raise RuntimeError("breakpointer!!!") - return out diff --git a/python/infinicore/ops/cdist.py b/python/infinicore/ops/cdist.py deleted file mode 100644 index 111a81d30..000000000 --- a/python/infinicore/ops/cdist.py +++ /dev/null @@ -1,36 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def cdist(x1, x2, p=2.0, *, out=None): - """ - 计算两组向量集合中每一对向量之间的 p-norm 距离。 - - 参数: - x1 (Tensor): 形状为 (M, D) 的输入张量。 - x2 (Tensor): 形状为 (N, D) 的输入张量。 - p (float): p-norm 的阶数，默认为 2.0。 - out (Tensor, optional): 结果输出张量。 - - 返回: - Tensor: 形状为 (M, N) 的距离矩阵。 - """ - if out is None: - # 非原地操作：由底层 C++ 接口根据 x1, x2 推导形状并创建新 Tensor - return Tensor( - _infinicore.cdist( - x1._underlying, - x2._underlying, - float(p), - ) - ) - - # 原地/指定输出操作：结果写入用户提供的 out 张量 - _infinicore.cdist_( - out._underlying, - x1._underlying, - x2._underlying, - float(p), - ) - - return out diff --git a/python/infinicore/ops/conv2d.py b/python/infinicore/ops/conv2d.py deleted file mode 100644 index 5bc89651a..000000000 --- a/python/infinicore/ops/conv2d.py +++ /dev/null @@ -1,47 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor, zeros - - -def _pair(value): - if isinstance(value, int): - return [value, value] - return list(value) - - -def conv2d( - input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, *, out=None -): - if groups != 1: - raise NotImplementedError( - "infinicore.ops.conv2d currently supports groups=1 only" - ) - - stride = _pair(stride) - padding = _pair(padding) - dilation = _pair(dilation) - if bias is None: - bias = zeros((weight.shape[0],), dtype=weight.dtype, device=weight.device) - bias_raw = bias._underlying - - if out is None: - return Tensor( - _infinicore.conv2d( - input._underlying, - weight._underlying, - bias_raw, - padding, - stride, - dilation, - ) - ) - - _infinicore.conv2d_( - out._underlying, - input._underlying, - weight._underlying, - bias_raw, - padding, - stride, - dilation, - ) - return out diff --git a/python/infinicore/ops/cross_entropy.py b/python/infinicore/ops/cross_entropy.py deleted file mode 100644 index 5b47697b5..000000000 --- a/python/infinicore/ops/cross_entropy.py +++ /dev/null @@ -1,33 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def cross_entropy( - logits, - target, - weight=None, - *, - ignore_index=None, - reduction="none", - out=None, -): - """ - Token-wise cross entropy without reduction. The output tensor has the same - shape as target and uses the logits dtype. - """ - if weight is not None: - raise NotImplementedError("class weights are not supported yet.") - if ignore_index is not None: - raise NotImplementedError("ignore_index is not supported yet.") - if reduction not in (None, "none"): - raise NotImplementedError("Only reduction='none' is implemented.") - - if out is None: - return Tensor(_infinicore.cross_entropy(logits._underlying, target._underlying)) - - _infinicore.cross_entropy_( - out._underlying, - logits._underlying, - target._underlying, - ) - return out diff --git a/python/infinicore/ops/diff.py b/python/infinicore/ops/diff.py deleted file mode 100644 index 374176390..000000000 --- a/python/infinicore/ops/diff.py +++ /dev/null @@ -1,14 +0,0 @@ -from __future__ import annotations - -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def diff(input: Tensor, n: int = 1, dim: int = -1, *, out: Optional[Tensor] = None): - if out is None: - return Tensor(_infinicore.diff(input._underlying, n, dim)) - - _infinicore.diff_(out._underlying, input._underlying, n, dim) - return out diff --git a/python/infinicore/ops/digamma.py b/python/infinicore/ops/digamma.py deleted file mode 100644 index 6ab42da1c..000000000 --- a/python/infinicore/ops/digamma.py +++ /dev/null @@ -1,14 +0,0 @@ -from __future__ import annotations - -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def digamma(input: Tensor, *, out: Optional[Tensor] = None): - if out is None: - return Tensor(_infinicore.digamma(input._underlying)) - - _infinicore.digamma_(out._underlying, input._underlying) - return out diff --git a/python/infinicore/ops/dist.py b/python/infinicore/ops/dist.py deleted file mode 100644 index 5f3ae7e19..000000000 --- a/python/infinicore/ops/dist.py +++ /dev/null @@ -1,14 +0,0 @@ -from __future__ import annotations - -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def dist(input: Tensor, other: Tensor, p: float = 2.0, *, out: Optional[Tensor] = None): - if out is None: - return Tensor(_infinicore.dist(input._underlying, other._underlying, p)) - - _infinicore.dist_(out._underlying, input._underlying, other._underlying, p) - return out diff --git a/python/infinicore/ops/equal.py b/python/infinicore/ops/equal.py deleted file mode 100644 index 5a656ab30..000000000 --- a/python/infinicore/ops/equal.py +++ /dev/null @@ -1,10 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def equal(input, other, *, out=None): - if out is None: - return Tensor(_infinicore.equal(input._underlying, other._underlying)) - - _infinicore.equal_(out._underlying, input._underlying, other._underlying) - return out diff --git a/python/infinicore/ops/flipud.py b/python/infinicore/ops/flipud.py deleted file mode 100644 index 86cffbe0e..000000000 --- a/python/infinicore/ops/flipud.py +++ /dev/null @@ -1,24 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def flipud(input: Tensor, *, out: Optional[Tensor] = None) -> Tensor: - r"""Flip array in the up/down direction. - - Flips the entries in axis 0 (preserving the shape). - - Args: - input (Tensor): the input tensor. - out (Tensor, optional): the output tensor. - - Returns: - Tensor: The flipped tensor. - """ - if not input.is_contiguous(): - input = input.contiguous() - if out is not None: - _infinicore.flipud_(out._underlying, input._underlying) - return out - return Tensor(_infinicore.flipud(input._underlying)) diff --git a/python/infinicore/ops/float_power.py b/python/infinicore/ops/float_power.py deleted file mode 100644 index 4dec677ca..000000000 --- a/python/infinicore/ops/float_power.py +++ /dev/null @@ -1,38 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def float_power( - input: Tensor, exponent: float, *, out: Optional[Tensor] = None -) -> Tensor: - r"""Computes the power of each element in input with the given exponent. - - .. math:: - \text{out}_i = \text{input}_i^{\text{exponent}} - - Args: - input (Tensor): the input tensor. - exponent (float): the exponent value. - out (Tensor, optional): the output tensor. - - Returns: - Tensor: The result tensor. - """ - - # 1. 确保输入内存连续 (Contiguous check) - if not input.is_contiguous(): - input = input.contiguous() - - # 2. 分发计算 - # 如果用户提供了 output tensor，调用底层的 in-place/explicit 接口 - if out is not None: - if not out.is_contiguous(): - raise RuntimeError("Output tensor must be contiguous") - - _infinicore.float_power_(out._underlying, input._underlying, exponent) - return out - - # 否则调用底层的 functional 接口，返回新 Tensor - return Tensor(_infinicore.float_power(input._underlying, exponent)) diff --git a/python/infinicore/ops/floor.py b/python/infinicore/ops/floor.py deleted file mode 100644 index af673b82d..000000000 --- a/python/infinicore/ops/floor.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def floor(input, *, out=None): - if out is None: - return Tensor(_infinicore.floor(input._underlying)) - - _infinicore.floor_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/ops/floor_divide.py b/python/infinicore/ops/floor_divide.py deleted file mode 100644 index 1489fde08..000000000 --- a/python/infinicore/ops/floor_divide.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def floor_divide(input, other, *, out=None): - if out is None: - return Tensor(_infinicore.floor_divide(input._underlying, other._underlying)) - - _infinicore.floor_divide_(out._underlying, input._underlying, other._underlying) - - return out diff --git a/python/infinicore/ops/fmin.py b/python/infinicore/ops/fmin.py deleted file mode 100644 index 56e274a10..000000000 --- a/python/infinicore/ops/fmin.py +++ /dev/null @@ -1,10 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def fmin(input, other, *, out=None): - if out is None: - return Tensor(_infinicore.fmin(input._underlying, other._underlying)) - - _infinicore.fmin_(out._underlying, input._underlying, other._underlying) - return out diff --git a/python/infinicore/ops/fmod.py b/python/infinicore/ops/fmod.py deleted file mode 100644 index e52be82cb..000000000 --- a/python/infinicore/ops/fmod.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def fmod(input, other, *, out=None): - if out is None: - return Tensor(_infinicore.fmod(input._underlying, other._underlying)) - - _infinicore.fmod_(out._underlying, input._underlying, other._underlying) - - return out diff --git a/python/infinicore/ops/hypot.py b/python/infinicore/ops/hypot.py deleted file mode 100644 index 107010987..000000000 --- a/python/infinicore/ops/hypot.py +++ /dev/null @@ -1,10 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def hypot(input, other, *, out=None): - if out is None: - return Tensor(_infinicore.hypot(input._underlying, other._underlying)) - _infinicore.hypot_(out._underlying, input._underlying, other._underlying) - - return out diff --git a/python/infinicore/ops/index_add.py b/python/infinicore/ops/index_add.py deleted file mode 100644 index f02d94753..000000000 --- a/python/infinicore/ops/index_add.py +++ /dev/null @@ -1,21 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def index_add(input, dim, index, source, *, alpha=1.0, out=None): - if out is None: - return Tensor( - _infinicore.index_add( - input._underlying, dim, index._underlying, source._underlying, alpha - ) - ) - _infinicore.index_add_( - out._underlying, - input._underlying, - dim, - index._underlying, - source._underlying, - alpha, - ) - - return out diff --git a/python/infinicore/ops/index_copy.py b/python/infinicore/ops/index_copy.py deleted file mode 100644 index b99f3e101..000000000 --- a/python/infinicore/ops/index_copy.py +++ /dev/null @@ -1,17 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def index_copy(input, dim, index, source, *, out=None): - if out is None: - return Tensor( - _infinicore.index_copy( - input._underlying, dim, index._underlying, source._underlying - ) - ) - - _infinicore.index_copy_( - out._underlying, input._underlying, dim, index._underlying, source._underlying - ) - - return out diff --git a/python/infinicore/ops/inner.py b/python/infinicore/ops/inner.py deleted file mode 100644 index df87c5b34..000000000 --- a/python/infinicore/ops/inner.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def inner(input, other, *, out=None): - if out is None: - return Tensor(_infinicore.inner(input._underlying, other._underlying)) - - _infinicore.inner_(out._underlying, input._underlying, other._underlying) - - return out diff --git a/python/infinicore/ops/kron.py b/python/infinicore/ops/kron.py deleted file mode 100644 index 15c83c32b..000000000 --- a/python/infinicore/ops/kron.py +++ /dev/null @@ -1,8 +0,0 @@ -from __future__ import annotations - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def kron(a: Tensor, b: Tensor) -> Tensor: - return Tensor(_infinicore.kron(a._underlying, b._underlying)) diff --git a/python/infinicore/ops/kthvalue.py b/python/infinicore/ops/kthvalue.py deleted file mode 100644 index ded2e5194..000000000 --- a/python/infinicore/ops/kthvalue.py +++ /dev/null @@ -1,40 +0,0 @@ -from typing import Optional, Tuple - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def kthvalue( - input: Tensor, - k: int, - dim: int = -1, - keepdim: bool = False, - *, - out: Optional[Tuple[Tensor, Tensor]] = None, -) -> Tuple[Tensor, Tensor]: - r"""Returns a namedtuple (values, indices) where values is the k-th smallest - element of each row of the input tensor in the given dimension. - """ - - if not input.is_contiguous(): - input = input.contiguous() - - if out is not None: - if not isinstance(out, (tuple, list)) or len(out) != 2: - raise ValueError("out must be a tuple of two Tensors (values, indices)") - - out_values, out_indices = out - - _infinicore.kthvalue_( - out_values._underlying, - out_indices._underlying, - input._underlying, - k, - dim, - keepdim, - ) - return out - - ret = _infinicore.kthvalue(input._underlying, k, dim, keepdim) - - return (Tensor(ret[0]), Tensor(ret[1])) diff --git a/python/infinicore/ops/kv_caching.py b/python/infinicore/ops/kv_caching.py deleted file mode 100644 index b34f2346e..000000000 --- a/python/infinicore/ops/kv_caching.py +++ /dev/null @@ -1,13 +0,0 @@ -from infinicore.lib import _infinicore - - -def kv_caching(k_cache, v_cache, k, v, past_kv_lengths): - _infinicore.kv_caching_( - k_cache._underlying, - v_cache._underlying, - k._underlying, - v._underlying, - past_kv_lengths._underlying, - ) - - return k_cache, v_cache diff --git a/python/infinicore/ops/ldexp.py b/python/infinicore/ops/ldexp.py deleted file mode 100644 index c785f1489..000000000 --- a/python/infinicore/ops/ldexp.py +++ /dev/null @@ -1,33 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def ldexp(input: Tensor, other: Tensor, *, out: Optional[Tensor] = None) -> Tensor: - r"""Multiplies input by 2 raised to the power of other. - - Args: - input (Tensor): The input tensor (mantissa). - other (Tensor): The exponent tensor. - """ - - # 1. 确保输入内存连续 - if not input.is_contiguous(): - input = input.contiguous() - if not other.is_contiguous(): - other = other.contiguous() - - # 2. 处理 Explicit Output (out=...) - if out is not None: - if not isinstance(out, Tensor): - raise ValueError("out must be a Tensor") - - _infinicore.ldexp_(out._underlying, input._underlying, other._underlying) - return out - - # 3. 处理 Functional 调用 - ret = _infinicore.ldexp(input._underlying, other._underlying) - - # 4. 封装返回结果 - return Tensor(ret) diff --git a/python/infinicore/ops/lerp.py b/python/infinicore/ops/lerp.py deleted file mode 100644 index c9bf0dd2c..000000000 --- a/python/infinicore/ops/lerp.py +++ /dev/null @@ -1,44 +0,0 @@ -from typing import Optional, Union - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def lerp( - start: Tensor, - end: Tensor, - weight: Union[Tensor, float], - *, - out: Optional[Tensor] = None, -) -> Tensor: - r"""Does a linear interpolation of two tensors start and end based on a scalar or tensor weight. - - output = start + weight * (end - start) - """ - - # 检查输入 Tensor 的连续性 - if not start.is_contiguous(): - start = start.contiguous() - if not end.is_contiguous(): - end = end.contiguous() - - # 处理 weight 参数：可能是 Tensor 也可能是标量 - weight_arg = weight - if isinstance(weight, Tensor): - if not weight.is_contiguous(): - weight = weight.contiguous() - weight_arg = weight._underlying - elif isinstance(weight, (float, int)): - weight_arg = float(weight) - else: - raise TypeError(f"weight must be a Tensor or float, got {type(weight)}") - - # In-place / 输出到指定 Tensor - if out is not None: - _infinicore.lerp_( - out._underlying, start._underlying, end._underlying, weight_arg - ) - return out - - # 返回新 Tensor - return Tensor(_infinicore.lerp(start._underlying, end._underlying, weight_arg)) diff --git a/python/infinicore/ops/logaddexp.py b/python/infinicore/ops/logaddexp.py deleted file mode 100644 index f333540dd..000000000 --- a/python/infinicore/ops/logaddexp.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def logaddexp(input, other, *, out=None): - if out is None: - return Tensor(_infinicore.logaddexp(input._underlying, other._underlying)) - - _infinicore.logaddexp_(out._underlying, input._underlying, other._underlying) - - return out diff --git a/python/infinicore/ops/logaddexp2.py b/python/infinicore/ops/logaddexp2.py deleted file mode 100644 index cc62de74b..000000000 --- a/python/infinicore/ops/logaddexp2.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def logaddexp2(input, other, *, out=None): - if out is None: - return Tensor(_infinicore.logaddexp2(input._underlying, other._underlying)) - - _infinicore.logaddexp2_(out._underlying, input._underlying, other._underlying) - - return out diff --git a/python/infinicore/ops/logcumsumexp.py b/python/infinicore/ops/logcumsumexp.py deleted file mode 100644 index 1167f0875..000000000 --- a/python/infinicore/ops/logcumsumexp.py +++ /dev/null @@ -1,17 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def logcumsumexp(input, dim, exclusive=False, reverse=False, *, out=None): - if out is None: - # 调用 C++ 绑定的非原地操作，返回新创建的 _underlying 句柄 - return Tensor( - _infinicore.logcumsumexp(input._underlying, dim, exclusive, reverse) - ) - - # 调用 C++ 绑定的原地/指定输出操作 - _infinicore.logcumsumexp_( - out._underlying, input._underlying, dim, exclusive, reverse - ) - - return out diff --git a/python/infinicore/ops/logdet.py b/python/infinicore/ops/logdet.py deleted file mode 100644 index 5280ddf32..000000000 --- a/python/infinicore/ops/logdet.py +++ /dev/null @@ -1,14 +0,0 @@ -from __future__ import annotations - -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def logdet(input: Tensor, *, out: Optional[Tensor] = None): - if out is None: - return Tensor(_infinicore.logdet(input._underlying)) - - _infinicore.logdet_(out._underlying, input._underlying) - return out diff --git a/python/infinicore/ops/logical_and.py b/python/infinicore/ops/logical_and.py deleted file mode 100644 index 0aa17c1ef..000000000 --- a/python/infinicore/ops/logical_and.py +++ /dev/null @@ -1,16 +0,0 @@ -import infinicore -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def logical_and(input: Tensor, other: Tensor, *, out=None) -> Tensor: - r"""Computes the element-wise logical AND of the given input tensors.""" - if input.device.type not in ("cpu"): - assert infinicore.use_ntops - return infinicore.ntops.torch.logical_and(input, other, out=out) - - if out is None: - return Tensor(_infinicore.logical_and(input._underlying, other._underlying)) - - _infinicore.logical_and_(out._underlying, input._underlying, other._underlying) - return out diff --git a/python/infinicore/ops/logical_not.py b/python/infinicore/ops/logical_not.py deleted file mode 100644 index 04168fe0f..000000000 --- a/python/infinicore/ops/logical_not.py +++ /dev/null @@ -1,19 +0,0 @@ -import infinicore -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def logical_not(input: Tensor, *, out=None) -> Tensor: - r"""Computes the element-wise logical NOT of the given input tensors.""" - # 1. 非CPU平台调用 ntops 实现 - if input.device.type not in ("cpu"): - assert infinicore.use_ntops - return infinicore.ntops.torch.logical_not(input, out=out) - - # 2. 如果没有提供 out，创建一个新的 Tensor 并返回 - if out is None: - return Tensor(_infinicore.logical_not(input._underlying)) - - # 3. 如果提供了 out，进行原地操作 (In-place operation) - _infinicore.logical_not_(out._underlying, input._underlying) - return out diff --git a/python/infinicore/ops/masked_select.py b/python/infinicore/ops/masked_select.py deleted file mode 100644 index 14a3a1ff0..000000000 --- a/python/infinicore/ops/masked_select.py +++ /dev/null @@ -1,6 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def masked_select(input, mask): - return Tensor(_infinicore.masked_select(input._underlying, mask._underlying)) diff --git a/python/infinicore/ops/matmul.py b/python/infinicore/ops/matmul.py deleted file mode 100644 index 6d3821059..000000000 --- a/python/infinicore/ops/matmul.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def matmul(input, other, *, alpha=1.0, out=None): - if out is None: - return Tensor(_infinicore.matmul(input._underlying, other._underlying, alpha)) - - _infinicore.matmul_(out._underlying, input._underlying, other._underlying, alpha) - - return out diff --git a/python/infinicore/ops/mha.py b/python/infinicore/ops/mha.py deleted file mode 100644 index c879a5a4b..000000000 --- a/python/infinicore/ops/mha.py +++ /dev/null @@ -1,37 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def mha( - q: Tensor, - k: Tensor, - v: Tensor, - alibi_slopes: Tensor | None = None, - scale: float = 1.0, - is_causal: bool = False, - *, - out: Tensor | None = None, -): - if out is None: - return Tensor( - _infinicore.mha( - q._underlying, - k._underlying, - v._underlying, - alibi_slopes._underlying if alibi_slopes is not None else None, - scale, - is_causal, - ) - ) - - _infinicore.mha_( - out._underlying, - q._underlying, - k._underlying, - v._underlying, - alibi_slopes._underlying if alibi_slopes is not None else None, - scale, - is_causal, - ) - - return out diff --git a/python/infinicore/ops/mha_kvcache.py b/python/infinicore/ops/mha_kvcache.py deleted file mode 100644 index 8ced8df16..000000000 --- a/python/infinicore/ops/mha_kvcache.py +++ /dev/null @@ -1,67 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def mha_kvcache( - q: Tensor, - k_cache: Tensor, - v_cache: Tensor, - seqlens_k: Tensor, - block_table: Tensor, - alibi_slopes: Optional[Tensor] = None, - scale: float = 1.0, - *, - out: Optional[Tensor] = None, -) -> Tensor: - """ - Flash attention KV-cache decode for single-step attention over a paged KV cache. - - This function performs attention decoding using a paged KV cache layout, - which is efficient for inference with large sequence lengths. - - Args: - q: Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] - k_cache: Key cache tensor of shape [num_blocks, block_size, num_heads_k, head_size] (paged layout) - v_cache: Value cache tensor of shape [num_blocks, block_size, num_heads_k, head_size] (paged layout) - seqlens_k: Total KV length per request of shape [batch_size] (int32) - block_table: Block mapping table of shape [batch_size, max_num_blocks_per_seq] (int32) - alibi_slopes: Optional ALiBi slopes tensor, if None then ALiBi is disabled - scale: Scaling factor for attention scores (typically 1.0/sqrt(head_size)) - out: Optional output tensor. If provided, the operation will be performed in-place. - - Returns: - Output tensor of shape [batch_size, seqlen_q, num_heads, head_size] - - Note: - The KV cache uses a paged layout where: - - k_cache and v_cache are organized into fixed-size blocks - - block_table maps logical positions to physical blocks for each sequence - - seqlens_k indicates the current total length of each sequence in the cache - """ - if out is None: - return Tensor( - _infinicore.mha_kvcache( - q._underlying, - k_cache._underlying, - v_cache._underlying, - seqlens_k._underlying, - block_table._underlying, - alibi_slopes._underlying if alibi_slopes is not None else None, - scale, - ) - ) - - _infinicore.mha_kvcache_( - out._underlying, - q._underlying, - k_cache._underlying, - v_cache._underlying, - seqlens_k._underlying, - block_table._underlying, - alibi_slopes._underlying if alibi_slopes is not None else None, - scale, - ) - - return out diff --git a/python/infinicore/ops/mha_varlen.py b/python/infinicore/ops/mha_varlen.py deleted file mode 100644 index 9a65ed930..000000000 --- a/python/infinicore/ops/mha_varlen.py +++ /dev/null @@ -1,49 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def mha_varlen( - q: Tensor, - k: Tensor, - v: Tensor, - cum_seqlens_q: Tensor, - cum_seqlens_k: Tensor, - block_table: Tensor | None, - max_seqlen_q: int, - max_seqlen_k: int, - alibi_slopes: Tensor | None = None, - scale: float = 1.0, - *, - out: Tensor | None = None, -): - if out is None: - return Tensor( - _infinicore.mha_varlen( - q._underlying, - k._underlying, - v._underlying, - cum_seqlens_q._underlying, - cum_seqlens_k._underlying, - block_table._underlying if block_table is not None else None, - max_seqlen_q, - max_seqlen_k, - alibi_slopes._underlying if alibi_slopes is not None else None, - scale, - ) - ) - - _infinicore.mha_varlen_( - out._underlying, - q._underlying, - k._underlying, - v._underlying, - cum_seqlens_q._underlying, - cum_seqlens_k._underlying, - block_table._underlying if block_table is not None else None, - max_seqlen_q, - max_seqlen_k, - alibi_slopes._underlying if alibi_slopes is not None else None, - scale, - ) - - return out diff --git a/python/infinicore/ops/moore_mate_flash_attn.py b/python/infinicore/ops/moore_mate_flash_attn.py deleted file mode 100644 index e50d0f6a8..000000000 --- a/python/infinicore/ops/moore_mate_flash_attn.py +++ /dev/null @@ -1,304 +0,0 @@ -""" -Paged Flash-Attention wrapper backed by MooreThreads mate (flash_attn). - -Runtime requirements: - - torch (with MUSA) - - mate (repo : https://github.com/MooreThreads/mate) - -Provides three entry points: - - moore_mate_flash_attn_decode: decode with layout (num_blocks, block_size, num_kv_heads, head_size) - - moore_mate_flash_attn_prefill: variable-length prefill (used by mha_varlen) -""" - -import torch - -try: - from flash_attn import flash_attn_with_kvcache, get_scheduler_metadata - - _MATE_AVAILABLE = True -except ImportError: - _MATE_AVAILABLE = False - - -def is_available() -> bool: - """Return True if mate / flash_attn is installed and importable.""" - return _MATE_AVAILABLE - - -def _check_mate_available(): - """Raise a clear error if mate is not installed.""" - if not _MATE_AVAILABLE: - raise RuntimeError( - "flash_attn (mate) is not installed. " - "Please build and install MooreThreads/mate first." - ) - - -# ============================================================================= -# Decode kernels -# ============================================================================= - - -@torch.inference_mode() -def moore_mate_flash_attn_decode( - q: torch.Tensor, # (num_seqs, num_heads, head_size) - k_cache: torch.Tensor, # (num_blocks, block_size, num_kv_heads, head_size) - v_cache: torch.Tensor, # (num_blocks, block_size, num_kv_heads, head_size) - block_tables: torch.Tensor, # (num_seqs, max_blocks_per_seq) - seq_lens: torch.Tensor, # (num_seqs,) - scale: float, - block_size: int, - max_seq_len: int, -) -> torch.Tensor: - """ - Decode entry point with native flash_attn KV cache layout (B, P, H, D). - No layout conversion is performed. - """ - _check_mate_available() - - num_seqs, num_heads, head_size = q.shape - num_kv_heads = k_cache.shape[2] - device = q.device - - cache_seqlens = seq_lens.to(torch.int32) - page_table = block_tables.to(torch.int32) - cu_seqlens_q = torch.arange(0, num_seqs + 1, dtype=torch.int32, device=device) - pack_gqa = num_heads != num_kv_heads - - metadata = get_scheduler_metadata( - batch_size=num_seqs, - max_seqlen_q=1, - max_seqlen_k=max_seq_len, - num_heads_q=num_heads, - num_heads_kv=num_kv_heads, - headdim=head_size, - cache_seqlens=cache_seqlens, - qkv_dtype=q.dtype, - headdim_v=head_size, - cu_seqlens_q=cu_seqlens_q, - page_size=block_size, - causal=False, - window_size=(None, None), - pack_gqa=pack_gqa, - ) - - out, *_ = flash_attn_with_kvcache( - q=q, - k_cache=k_cache, - v_cache=v_cache, - cache_seqlens=cache_seqlens, - page_table=page_table, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=1, - softmax_scale=scale, - causal=False, - scheduler_metadata=metadata, - pack_gqa=pack_gqa, - return_softmax_lse=True, - ) - return out - - -# ============================================================================= -# Prefill kernel (variable-length) -# ============================================================================= - - -@torch.inference_mode() -def moore_mate_flash_attn_prefill( - q: torch.Tensor, # (total_q, num_heads, head_size) -- varlen unpad - k_cache: torch.Tensor, # (num_blocks, block_size, num_kv_heads, head_size) - v_cache: torch.Tensor, # (num_blocks, block_size, num_kv_heads, head_size) - cu_seqlens_q: torch.Tensor, # (batch+1,) int32 - cu_seqlens_k: torch.Tensor, # (batch+1,) int32 - block_tables: torch.Tensor, # (batch, max_blocks_per_seq) - scale: float, - max_seqlen_q: int, - max_seqlen_k: int, - block_size: int, - causal: bool = True, # prefill is typically causal -) -> torch.Tensor: - """ - Variable-length prefill entry point. Layout follows flash_attn (B, P, H, D). - Intended to be called from the C++ mha_varlen Moore branch. - """ - _check_mate_available() - - cu_seqlens_q = cu_seqlens_q.to(torch.int32) - cu_seqlens_k = cu_seqlens_k.to(torch.int32) - page_table = block_tables.to(torch.int32) - - # mate uses cache_seqlens (per-batch KV length), derived from cu_seqlens_k - cache_seqlens = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]).to(torch.int32).contiguous() - - batch_size = cache_seqlens.shape[0] - num_heads = q.shape[1] - head_size = q.shape[2] - num_kv_heads = k_cache.shape[2] - pack_gqa = num_heads != num_kv_heads - - metadata = get_scheduler_metadata( - batch_size=batch_size, - max_seqlen_q=int(max_seqlen_q), - max_seqlen_k=int(max_seqlen_k), - num_heads_q=num_heads, - num_heads_kv=num_kv_heads, - headdim=head_size, - cache_seqlens=cache_seqlens, - qkv_dtype=q.dtype, - headdim_v=head_size, - cu_seqlens_q=cu_seqlens_q, - page_size=block_size, - causal=causal, - window_size=(None, None), - pack_gqa=pack_gqa, - ) - - out, *_ = flash_attn_with_kvcache( - q=q, - k_cache=k_cache, - v_cache=v_cache, - cache_seqlens=cache_seqlens, - page_table=page_table, - cu_seqlens_q=cu_seqlens_q, - max_seqlen_q=int(max_seqlen_q), - softmax_scale=scale, - causal=causal, - scheduler_metadata=metadata, - pack_gqa=pack_gqa, - return_softmax_lse=True, - ) - return out - - -# ============================================================================= -# Self tests -# ============================================================================= - - -def _test_moore_mate_flash_attn_decode(): - """Test moore_mate_flash_attn_decode with flash_attn layout (B, P, H, D).""" - print("\n=== Test 1: moore_mate_flash_attn_decode (decode, flash_attn layout) ===") - device = torch.device("musa") - - num_seqs, num_heads, num_kv_heads = 2, 8, 2 - head_size, block_size, max_seq_len = 128, 16, 64 - num_blocks = 32 - - q = torch.randn(num_seqs, num_heads, head_size, dtype=torch.float16, device=device) - k_cache = torch.randn( - num_blocks, - block_size, - num_kv_heads, - head_size, - dtype=torch.float16, - device=device, - ) - v_cache = torch.randn( - num_blocks, - block_size, - num_kv_heads, - head_size, - dtype=torch.float16, - device=device, - ) - block_tables = torch.zeros(num_seqs, 4, dtype=torch.int32, device=device) - block_tables[0, 0] = 0 - block_tables[1, 0] = 1 - seq_lens = torch.tensor([32, 48], dtype=torch.int32, device=device) - - out = moore_mate_flash_attn_decode( - q=q, - k_cache=k_cache, - v_cache=v_cache, - block_tables=block_tables, - seq_lens=seq_lens, - scale=head_size**-0.5, - block_size=block_size, - max_seq_len=max_seq_len, - ) - torch.musa.synchronize() - print(f"output shape = {tuple(out.shape)}") - assert out.shape == q.shape - print("moore_mate_flash_attn_decode passed") - - -def _test_moore_mate_flash_attn_prefill(): - """Test moore_mate_flash_attn_prefill with variable-length input.""" - print("\n=== Test 2: moore_mate_flash_attn_prefill (varlen prefill) ===") - device = torch.device("musa") - torch.manual_seed(666) - torch.musa.manual_seed(666) - - batch_size = 2 - seqlens_q = [55, 222] - seqlens_kv = [55, 222] # prefill: q_len == k_len - num_heads, num_kv_heads = 8, 2 - head_size, block_size = 128, 16 - - total_q = sum(seqlens_q) - max_q = max(seqlens_q) - max_k = max(seqlens_kv) - num_blocks_per_seq = (max_k + block_size - 1) // block_size - num_blocks = batch_size * num_blocks_per_seq - - q_unpad = torch.randn( - total_q, num_heads, head_size, dtype=torch.float16, device=device - ) - k_cache = torch.randn( - num_blocks, - block_size, - num_kv_heads, - head_size, - dtype=torch.float16, - device=device, - ) - v_cache = torch.randn( - num_blocks, - block_size, - num_kv_heads, - head_size, - dtype=torch.float16, - device=device, - ) - - cu_seqlens_q = torch.tensor( - [0] + list(torch.cumsum(torch.tensor(seqlens_q), 0).tolist()), - dtype=torch.int32, - device=device, - ) - cu_seqlens_k = torch.tensor( - [0] + list(torch.cumsum(torch.tensor(seqlens_kv), 0).tolist()), - dtype=torch.int32, - device=device, - ) - block_tables = torch.arange(num_blocks, dtype=torch.int32, device=device).view( - batch_size, num_blocks_per_seq - ) - - out = moore_mate_flash_attn_prefill( - q=q_unpad, - k_cache=k_cache, - v_cache=v_cache, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_k=cu_seqlens_k, - block_tables=block_tables, - scale=head_size**-0.5, - max_seqlen_q=max_q, - max_seqlen_k=max_k, - block_size=block_size, - causal=True, - ) - torch.musa.synchronize() - print(f"output shape = {tuple(out.shape)}") - assert out.shape == q_unpad.shape - print("moore_mate_flash_attn_prefill passed") - - -if __name__ == "__main__": - if not is_available(): - raise SystemExit("mate / flash_attn not available, please build mate first.") - - _test_moore_mate_flash_attn_decode() - _test_moore_mate_flash_attn_prefill() - print("\nAll tests passed.") diff --git a/python/infinicore/ops/mrope.py b/python/infinicore/ops/mrope.py deleted file mode 100644 index 8111ce9d7..000000000 --- a/python/infinicore/ops/mrope.py +++ /dev/null @@ -1,52 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def mrope( - q: Tensor, - k: Tensor, - cos: Tensor, - sin: Tensor, - positions: Tensor, - head_size: int, - rotary_dim: int, - section_t: int, - section_h: int, - section_w: int, - interleaved: bool, - *, - out=None, -) -> tuple[Tensor, Tensor]: - if out is None: - q_out, k_out = _infinicore.mrope( - q._underlying, - k._underlying, - cos._underlying, - sin._underlying, - positions._underlying, - head_size, - rotary_dim, - section_t, - section_h, - section_w, - interleaved, - ) - return Tensor(q_out), Tensor(k_out) - - q_out, k_out = out - _infinicore.mrope_( - q_out._underlying, - k_out._underlying, - q._underlying, - k._underlying, - cos._underlying, - sin._underlying, - positions._underlying, - head_size, - rotary_dim, - section_t, - section_h, - section_w, - interleaved, - ) - return q_out, k_out diff --git a/python/infinicore/ops/mul.py b/python/infinicore/ops/mul.py deleted file mode 100644 index 8ce703917..000000000 --- a/python/infinicore/ops/mul.py +++ /dev/null @@ -1,17 +0,0 @@ -import numbers - -from infinicore.lib import _infinicore -from infinicore.ops.mul_scalar import mul_scalar -from infinicore.tensor import Tensor - - -def mul(input, other, *, out=None): - if isinstance(other, numbers.Real): - return mul_scalar(input, other, out=out) - - if out is None: - return Tensor(_infinicore.mul(input._underlying, other._underlying)) - - _infinicore.mul_(out._underlying, input._underlying, other._underlying) - - return out diff --git a/python/infinicore/ops/mul_scalar.py b/python/infinicore/ops/mul_scalar.py deleted file mode 100644 index e7508a570..000000000 --- a/python/infinicore/ops/mul_scalar.py +++ /dev/null @@ -1,12 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def mul_scalar(input, alpha, *, out=None): - alpha = float(alpha) - if out is None: - return Tensor(_infinicore.mul_scalar(input._underlying, alpha)) - - _infinicore.mul_scalar_(out._underlying, input._underlying, alpha) - - return out diff --git a/python/infinicore/ops/narrow.py b/python/infinicore/ops/narrow.py deleted file mode 100644 index 3d605a77b..000000000 --- a/python/infinicore/ops/narrow.py +++ /dev/null @@ -1,5 +0,0 @@ -from infinicore.tensor import Tensor - - -def narrow(input: Tensor, dim: int, start: int, length: int) -> Tensor: - return Tensor(input._underlying.narrow(dim, start, length)) diff --git a/python/infinicore/ops/nrm2.py b/python/infinicore/ops/nrm2.py deleted file mode 100644 index fcc3a0d3a..000000000 --- a/python/infinicore/ops/nrm2.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def nrm2(x: Tensor, *, out=None): - if out is None: - return Tensor(_infinicore.nrm2(x._underlying)) - - _infinicore.nrm2_(x._underlying, out._underlying) - - return out diff --git a/python/infinicore/ops/paged_attention.py b/python/infinicore/ops/paged_attention.py deleted file mode 100644 index dfefa6a76..000000000 --- a/python/infinicore/ops/paged_attention.py +++ /dev/null @@ -1,40 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def paged_attention( - q: Tensor, - k_cache: Tensor, - v_cache: Tensor, - block_tables: Tensor, - cache_lens: Tensor, - alibi_slopes: Tensor | None = None, - scale: float = 1.0, - *, - out: Tensor | None = None, -): - if out is None: - return Tensor( - _infinicore.paged_attention( - q._underlying, - k_cache._underlying, - v_cache._underlying, - block_tables._underlying, - cache_lens._underlying, - alibi_slopes._underlying if alibi_slopes is not None else None, - scale, - ) - ) - - _infinicore.paged_attention_( - out._underlying, - q._underlying, - k_cache._underlying, - v_cache._underlying, - block_tables._underlying, - cache_lens._underlying, - alibi_slopes._underlying if alibi_slopes is not None else None, - scale, - ) - - return out diff --git a/python/infinicore/ops/paged_attention_prefill.py b/python/infinicore/ops/paged_attention_prefill.py deleted file mode 100644 index 848f74abf..000000000 --- a/python/infinicore/ops/paged_attention_prefill.py +++ /dev/null @@ -1,45 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def paged_attention_prefill( - q: Tensor, - k_cache: Tensor, - v_cache: Tensor, - block_tables: Tensor, - history_lens: Tensor, - cu_seqlens_q: Tensor, - alibi_slopes: Tensor | None = None, - scale: float = 1.0, - *, - out: Tensor | None = None, -): - alibi_ptr = alibi_slopes._underlying if alibi_slopes is not None else None - - if out is None: - return Tensor( - _infinicore.paged_attention_prefill( - q._underlying, - k_cache._underlying, - v_cache._underlying, - block_tables._underlying, - history_lens._underlying, - cu_seqlens_q._underlying, - alibi_ptr, - scale, - ) - ) - - _infinicore.paged_attention_prefill_( - out._underlying, - q._underlying, - k_cache._underlying, - v_cache._underlying, - block_tables._underlying, - history_lens._underlying, - cu_seqlens_q._underlying, - alibi_ptr, - scale, - ) - - return out diff --git a/python/infinicore/ops/paged_caching.py b/python/infinicore/ops/paged_caching.py deleted file mode 100644 index e3b8d63fb..000000000 --- a/python/infinicore/ops/paged_caching.py +++ /dev/null @@ -1,21 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def paged_caching( - k_cache: Tensor, - v_cache: Tensor, - k: Tensor, - v: Tensor, - slot_mapping: Tensor, -): - Tensor( - _infinicore.paged_caching_( - k_cache._underlying, - v_cache._underlying, - k._underlying, - v._underlying, - slot_mapping._underlying, - ) - ) - return (k_cache, v_cache) diff --git a/python/infinicore/ops/rearrange.py b/python/infinicore/ops/rearrange.py deleted file mode 100644 index e876fba23..000000000 --- a/python/infinicore/ops/rearrange.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def rearrange(input, other, *, out=None): - if out is None: - return Tensor(_infinicore.rearrange(input._underlying)) - - _infinicore.rearrange_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/ops/reciprocal.py b/python/infinicore/ops/reciprocal.py deleted file mode 100644 index f5573ace8..000000000 --- a/python/infinicore/ops/reciprocal.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def reciprocal(input, *, out=None): - if out is None: - return Tensor(_infinicore.reciprocal(input._underlying)) - - _infinicore.reciprocal_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/ops/rot.py b/python/infinicore/ops/rot.py deleted file mode 100644 index 091775c4e..000000000 --- a/python/infinicore/ops/rot.py +++ /dev/null @@ -1,7 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def rot(x: Tensor, y: Tensor, c: Tensor, s: Tensor): - _infinicore.rot_(x._underlying, y._underlying, c._underlying, s._underlying) - return x, y diff --git a/python/infinicore/ops/rotg.py b/python/infinicore/ops/rotg.py deleted file mode 100644 index 8299e6205..000000000 --- a/python/infinicore/ops/rotg.py +++ /dev/null @@ -1,7 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def rotg(x: Tensor, y: Tensor, c: Tensor, s: Tensor): - _infinicore.rotg_(x._underlying, y._underlying, c._underlying, s._underlying) - return x, y, c, s diff --git a/python/infinicore/ops/rotm.py b/python/infinicore/ops/rotm.py deleted file mode 100644 index 68acfc231..000000000 --- a/python/infinicore/ops/rotm.py +++ /dev/null @@ -1,7 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def rotm(x: Tensor, y: Tensor, param: Tensor): - _infinicore.rotm_(x._underlying, y._underlying, param._underlying) - return x, y diff --git a/python/infinicore/ops/rotmg.py b/python/infinicore/ops/rotmg.py deleted file mode 100644 index 468dd687e..000000000 --- a/python/infinicore/ops/rotmg.py +++ /dev/null @@ -1,13 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def rotmg(d1: Tensor, d2: Tensor, x1: Tensor, y1: Tensor, param: Tensor): - _infinicore.rotmg_( - d1._underlying, - d2._underlying, - x1._underlying, - y1._underlying, - param._underlying, - ) - return d1, d2, x1, param diff --git a/python/infinicore/ops/scal.py b/python/infinicore/ops/scal.py deleted file mode 100644 index 8302e74a5..000000000 --- a/python/infinicore/ops/scal.py +++ /dev/null @@ -1,8 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def scal(x: Tensor, alpha: Tensor): - _infinicore.scal_(alpha._underlying, x._underlying) - - return x diff --git a/python/infinicore/ops/scatter.py b/python/infinicore/ops/scatter.py deleted file mode 100644 index 0711a8386..000000000 --- a/python/infinicore/ops/scatter.py +++ /dev/null @@ -1,58 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - -# Scatter 算子常用的 reduction 模式 -_SCATTER_REDUCTION_MODES = { - "none": 0, # 直接赋值/覆盖 - "add": 1, # 累加 - "multiply": 2, # 累乘 -} - - -def scatter( - input: Tensor, - index: Tensor, - src: Tensor, - dim: int, - reduction: str = "none", - *, - out: Optional[Tensor] = None, -) -> Tensor: - r"""Writes all values from the tensor src into input at the indices specified in the index tensor.""" - - if not input.is_contiguous(): - input = input.contiguous() - if not index.is_contiguous(): - index = index.contiguous() - if not src.is_contiguous(): - src = src.contiguous() - - # 解析 reduction 参数 - if reduction not in _SCATTER_REDUCTION_MODES: - raise ValueError(f"{reduction} is not a valid value for reduction") - reduction_val = _SCATTER_REDUCTION_MODES[reduction] - - # In-place 分支 (scatter_) - if out is not None: - _infinicore.scatter_( - out._underlying, - input._underlying, - index._underlying, # index (第3个) - src._underlying, # src (第4个) - dim, # dim (第5个) - reduction_val, - ) - return out - - # Out-of-place 分支 (scatter) - return Tensor( - _infinicore.scatter( - input._underlying, - index._underlying, # index (第2个) - src._underlying, # src (第3个) - dim, # dim (第4个) - reduction_val, - ) - ) diff --git a/python/infinicore/ops/sinh.py b/python/infinicore/ops/sinh.py deleted file mode 100644 index 34f3c589a..000000000 --- a/python/infinicore/ops/sinh.py +++ /dev/null @@ -1,12 +0,0 @@ -from __future__ import annotations - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def sinh(input: Tensor, *, out: Tensor | None = None) -> Tensor: - if out is None: - return Tensor(_infinicore.sinh(input._underlying)) - - _infinicore.sinh_(out._underlying, input._underlying) - return out diff --git a/python/infinicore/ops/squeeze.py b/python/infinicore/ops/squeeze.py deleted file mode 100644 index 8b78bf833..000000000 --- a/python/infinicore/ops/squeeze.py +++ /dev/null @@ -1,5 +0,0 @@ -from infinicore.tensor import Tensor - - -def squeeze(input: Tensor, dim: int) -> Tensor: - return Tensor(input._underlying.squeeze(dim)) diff --git a/python/infinicore/ops/sum.py b/python/infinicore/ops/sum.py deleted file mode 100644 index 5f264c24b..000000000 --- a/python/infinicore/ops/sum.py +++ /dev/null @@ -1,28 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def sum(input, dim=None, keepdim=False, out=None): - """ - Sum the elements of the input tensor along the given dimensions. - - Args: - input (Tensor): The input tensor. - out (Tensor, optional): The output tensor. - - Returns: - Tensor: The output tensor. - - Example: - >>> import infinicore - >>> input = infinicore.tensor([[1, 2, 3], [4, 5, 6]]) - >>> output = infinicore.sum(input) - >>> print(output) - tensor([15]) - """ - if out is None: - return Tensor(_infinicore.sum(input._underlying, dim, keepdim)) - - _infinicore.sum_(out._underlying, input._underlying, dim, keepdim) - - return out diff --git a/python/infinicore/ops/swap.py b/python/infinicore/ops/swap.py deleted file mode 100644 index f773a34eb..000000000 --- a/python/infinicore/ops/swap.py +++ /dev/null @@ -1,7 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def swap(x: Tensor, y: Tensor): - _infinicore.swap_(x._underlying, y._underlying) - return x, y diff --git a/python/infinicore/ops/take.py b/python/infinicore/ops/take.py deleted file mode 100644 index 4a0e36193..000000000 --- a/python/infinicore/ops/take.py +++ /dev/null @@ -1,27 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def take(input: Tensor, indices: Tensor, *, out: Optional[Tensor] = None) -> Tensor: - r"""Returns a new tensor with the elements of input at the given indices. - The input tensor is treated as if it were viewed as a 1-D tensor. - The result tensor has the same shape as the indices tensor. - - Args: - input (Tensor): the input tensor. - indices (Tensor): the indices into tensor, must be an Int or Long tensor. - out (Tensor, optional): the output tensor. - - Returns: - Tensor: A new tensor with the elements of input at the given indices. - """ - if not input.is_contiguous(): - input = input.contiguous() - - # 如果用户提供了 output tensor，调用底层的 in-place/explicit 接口 - if out is not None: - _infinicore.take_(out._underlying, input._underlying, indices._underlying) - return out - return Tensor(_infinicore.take(input._underlying, indices._underlying)) diff --git a/python/infinicore/ops/tan.py b/python/infinicore/ops/tan.py deleted file mode 100644 index da438e6d8..000000000 --- a/python/infinicore/ops/tan.py +++ /dev/null @@ -1,11 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def tan(input, *, out=None): - if out is None: - return Tensor(_infinicore.tan(input._underlying)) - - _infinicore.tan_(out._underlying, input._underlying) - - return out diff --git a/python/infinicore/ops/topk.py b/python/infinicore/ops/topk.py deleted file mode 100644 index 86eb32ee6..000000000 --- a/python/infinicore/ops/topk.py +++ /dev/null @@ -1,12 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def topk(input, k, dim, largest=True, sorted=True, out=None): - if out is None: - values, indices = _infinicore.topk(input._underlying, k, dim, largest, sorted) - return Tensor(values), Tensor(indices) - - _infinicore.topk_(out._underlying, input._underlying, k, dim, largest, sorted) - - return out diff --git a/python/infinicore/ops/unsqueeze.py b/python/infinicore/ops/unsqueeze.py deleted file mode 100644 index 0d421cb2b..000000000 --- a/python/infinicore/ops/unsqueeze.py +++ /dev/null @@ -1,5 +0,0 @@ -from infinicore.tensor import Tensor - - -def unsqueeze(input: Tensor, dim: int) -> Tensor: - return Tensor(input._underlying.unsqueeze(dim)) diff --git a/python/infinicore/ops/vander.py b/python/infinicore/ops/vander.py deleted file mode 100644 index d1c12975d..000000000 --- a/python/infinicore/ops/vander.py +++ /dev/null @@ -1,36 +0,0 @@ -from typing import Optional - -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def vander( - x: Tensor, - N: Optional[int] = None, - increasing: bool = False, - *, - out: Optional[Tensor] = None, -) -> Tensor: - r"""Generates a Vandermonde matrix. - - The columns of the output matrix are powers of the input vector. The - order of the powers is determined by the ``increasing`` boolean argument. - - Args: - x (Tensor): 1-D input tensor. - N (int, optional): Number of columns in the output. If None, defaults to the size of x. - increasing (bool, optional): Order of the powers. - If False (default), the powers are descending (x^(N-1), ..., x^0). - If True, the powers are ascending (x^0, ..., x^(N-1)). - out (Tensor, optional): The output tensor. - """ - - if not x.is_contiguous(): - x = x.contiguous() - N_val = N if N is not None else 0 - - if out is not None: - _infinicore.vander_(out._underlying, x._underlying, N_val, increasing) - return out - - return Tensor(_infinicore.vander(x._underlying, N_val, increasing)) diff --git a/python/infinicore/ops/var.py b/python/infinicore/ops/var.py deleted file mode 100644 index 71911ab10..000000000 --- a/python/infinicore/ops/var.py +++ /dev/null @@ -1,12 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def var(input, dim=None, unbiased=True, keepdim=False, out=None): - if out is None: - var_tensor = _infinicore.var(input._underlying, dim, unbiased, keepdim) - return Tensor(var_tensor) - var_output = out - _infinicore.var_(var_output._underlying, input._underlying, dim, unbiased, keepdim) - - return out diff --git a/python/infinicore/ops/var_mean.py b/python/infinicore/ops/var_mean.py deleted file mode 100644 index 0a9573938..000000000 --- a/python/infinicore/ops/var_mean.py +++ /dev/null @@ -1,21 +0,0 @@ -from infinicore.lib import _infinicore -from infinicore.tensor import Tensor - - -def var_mean(input, dim=None, unbiased=True, keepdim=False, out=None): - if out is None: - var_tensor, mean_tensor = _infinicore.var_mean( - input._underlying, dim, unbiased, keepdim - ) - return Tensor(var_tensor), Tensor(mean_tensor) - var_output, mean_output = out - _infinicore.var_mean_( - var_output._underlying, - mean_output._underlying, - input._underlying, - dim, - unbiased, - keepdim, - ) - - return out diff --git a/python/infinicore/tensor.py b/python/infinicore/tensor.py deleted file mode 100644 index bbe801f93..000000000 --- a/python/infinicore/tensor.py +++ /dev/null @@ -1,367 +0,0 @@ -import ctypes - -import numpy as np - -import infinicore.device -import infinicore.dtype -from infinicore.lib import _infinicore - -from ._tensor_str import _str -from .utils import ( - infinicore_to_numpy_dtype, - numpy_to_infinicore_dtype, - to_infinicore_dtype, -) - - -class Tensor: - # Public attributes describing the Tensor - _underlying: _infinicore.Tensor - _torch_ref: "torch.Tensor" # noqa: F821 - shape: list[int] - dtype: infinicore.dtype - device: infinicore.device - - def __init__(self, underlying, *, _torch_ref=None): - """An internal method. Please do not use this directly.""" - - self._underlying = underlying - self._torch_ref = _torch_ref - - def __getattr__(self, name): - # Lazily construct and cache an attribute. - # such as, self.shape, self.dtype, self.device . - if name == "shape": - setattr(self, name, getattr(self._underlying, name)) - elif name == "dtype": - setattr(self, name, infinicore.dtype(getattr(self._underlying, name))) - elif name == "device": - setattr( - self, - name, - infinicore.device._from_infinicore_device( - getattr(self._underlying, name) - ), - ) - else: - raise AttributeError( - "{!r} object has no attribute {!r}".format(__name__, name) - ) - - return getattr(self, name) - - @property - def ndim(self): - return self._underlying.ndim - - def data_ptr(self): - return self._underlying.data_ptr() - - def size(self, dim=None): - if dim is None: - return self.shape - - return self.shape[dim] - - def stride(self, dim=None): - if dim is None: - return self._underlying.strides - - return self._underlying.strides[dim] - - def numel(self): - return self._underlying.numel() - - def is_contiguous(self): - return self._underlying.is_contiguous() - - def is_pinned(self): - return self._underlying.is_pinned() - - def copy_(self, src): - self._underlying.copy_(src._underlying) - - def to(self, *args, **kwargs): - return Tensor( - self._underlying.to(*tuple(arg._underlying for arg in args), **kwargs) - ) - - def contiguous(self): - return Tensor(self._underlying.contiguous()) - - def as_strided(self, size, stride): - return Tensor(self._underlying.as_strided(size, stride)) - - def permute(self, dims): - return Tensor(self._underlying.permute(dims)) - - def view(self, shape): - return Tensor(self._underlying.view(shape)) - - def squeeze(self, dim): - return infinicore.squeeze(self, dim) - - def unsqueeze(self, dim): - return infinicore.unsqueeze(self, dim) - - def debug(self, filename=None): - """Print tensor data or save to file for debugging - - Args: - filename: Optional filename to save raw binary data. If None, prints to stdout. - """ - if filename is None: - self._underlying.debug() - else: - self._underlying.debug(filename) - - def __add__(self, other): - return infinicore.add(self, other) - - def __iadd__(self, other): - infinicore.add(self, other, out=self) - return self - - def __matmul__(self, other): - return infinicore.matmul(self, other) - - def __mul__(self, other): - return infinicore.mul(self, other) - - def __rmul__(self, other): - return infinicore.mul(self, other) - - def narrow(self, dim, start, length): - return infinicore.narrow(self, dim, start, length) - - def __repr__(self): - return _str(self) - - -def empty(size, *, dtype=None, device=None, pin_memory=False): - return Tensor( - _infinicore.empty(size, dtype._underlying, device._underlying, pin_memory) - ) - - -def empty_like(input, *, dtype=None, device=None): - if dtype is None: - dtype = input.dtype - - if device is None: - device = input.device - - return empty(input.size(), dtype=dtype, device=device) - - -def strided_empty(size, strides, *, dtype=None, device=None, pin_memory=False): - return Tensor( - _infinicore.strided_empty( - size, strides, dtype._underlying, device._underlying, pin_memory - ) - ) - - -def zeros(size, *, dtype=None, device=infinicore.device.device(), pin_memory=False): - return Tensor( - _infinicore.zeros(size, dtype._underlying, device._underlying, pin_memory) - ) - - -def ones(size, *, dtype=None, device=None, pin_memory=False): - return Tensor( - _infinicore.ones(size, dtype._underlying, device._underlying, pin_memory) - ) - - -def from_blob(data_ptr, size, *, dtype=None, device=None): - return Tensor( - _infinicore.from_blob(data_ptr, size, dtype._underlying, device._underlying) - ) - - -def strided_from_blob(data_ptr, size, strides, *, dtype=None, device=None): - return Tensor( - _infinicore.strided_from_blob( - data_ptr, size, strides, dtype._underlying, device._underlying - ) - ) - - -def from_torch(torch_tensor) -> Tensor: - infini_type = to_infinicore_dtype(torch_tensor.dtype) - infini_device = infinicore.device(torch_tensor.device.type, 0) - return Tensor( - _infinicore.from_blob( - torch_tensor.data_ptr(), - list(torch_tensor.shape), - dtype=infini_type._underlying, - device=infini_device._underlying, - ), - _torch_ref=torch_tensor, - ) - - -def from_numpy( - np_array, - *, - dtype: infinicore.dtype = None, - device: infinicore.device = None, -) -> Tensor: - """Convert a NumPy ndarray to an infinicore Tensor. - - Args: - np_array: NumPy ndarray to convert to tensor - dtype: Optional infinicore dtype. If None, inferred from numpy array - device: Optional infinicore device. If None, defaults to CPU device - - Returns: - Tensor: An infinicore tensor created from the numpy array - - Raises: - TypeError: If input data is not a numpy ndarray - ValueError: If input array is empty - - Note: - NumPy arrays can only be created on CPU. For CUDA devices, data is first - created on CPU, then copied to the target device. - """ - # Input validation - if not isinstance(np_array, np.ndarray): - raise TypeError( - f"Input data must be a np.ndarray, got {type(np_array).__name__}" - ) - - if np_array.size == 0: - raise ValueError("Input array cannot be empty") - - # Determine target numpy dtype - # If dtype is specified, convert it to numpy dtype first - if dtype is not None: - np_dtype = infinicore_to_numpy_dtype(dtype) - # Create a copy with the target dtype if dtype doesn't match - # Use copy=True to ensure we don't modify the original array - if np_dtype != np_array.dtype: - np_array = np_array.astype(np_dtype, copy=True) - # Ensure C-contiguous layout - elif not np_array.flags.c_contiguous: - np_array = np.ascontiguousarray(np_array) - else: - # Ensure C-contiguous layout - if not np_array.flags.c_contiguous: - np_array = np.ascontiguousarray(np_array) - - # Infer infinicore dtype if not provided - infini_type = ( - dtype if dtype is not None else numpy_to_infinicore_dtype(np_array.dtype) - ) - - # Default to CPU device if not provided - infini_device = device if device is not None else infinicore.device("cpu", 0) - cpu_device = infinicore.device("cpu", 0) - - # Create a temporary tensor on CPU using from_blob to reference numpy array - # This allows us to copy data without keeping numpy array reference - data_ptr = np_array.ctypes.data_as(ctypes.c_void_p).value - temp_tensor = Tensor( - _infinicore.from_blob( - data_ptr, - list(np_array.shape), - dtype=infini_type._underlying, - device=cpu_device._underlying, - ) - ) - - # Always create the result tensor on CPU first, then copy data - # This ensures we have a proper copy of the data - result = empty(list(np_array.shape), dtype=infini_type, device=cpu_device) - result.copy_(temp_tensor) - - # If target device is not CPU, move the tensor to the target device - # The temporary tensor and numpy array will be garbage collected - # since we don't keep references to them - if infini_device.type != "cpu": - result = result.to(infini_device) - - return result - - -def from_list(data, *, dtype) -> Tensor: - """Convert a 1D or 2D Python list to an infinicore Tensor on CPU. - - This follows a NumPy-like three-stage flow implemented in C++: - validate shape, allocate an empty tensor of the requested dtype, then - pack list elements directly into the destination buffer. - - Args: - data: A 1D list of scalars or a 2D list of scalars. - dtype: Required infinicore dtype. Must be specified explicitly; - ``from_list`` does not infer the dtype from list elements. - - Returns: - Tensor: A CPU infinicore tensor created from the list data. - - Raises: - TypeError: If ``dtype`` is missing, or if input is not a list/tuple - or contains invalid elements. - ValueError: If input is empty, irregular, or deeper than 2D. - - Note: - The result is always on CPU. Use ``tensor.to(device)`` to move it elsewhere. - """ - if dtype is None: - raise TypeError("from_list() requires `dtype` to be specified") - return Tensor(_infinicore.from_list(data, dtype._underlying)) - - -def from_list_by_numpy(data, *, dtype=None, device=None) -> Tensor: - """Convert a Python list to an infinicore Tensor. - - Args: - data: Python list or nested list to convert to tensor - dtype: Optional infinicore dtype. If None, inferred from numpy array - device: Optional infinicore device. If None, defaults to CPU device - - Returns: - Tensor: An infinicore tensor created from the list data - - Raises: - TypeError: If input data is not a list or tuple - ValueError: If input data is empty - - Note: - NumPy arrays can only be created on CPU. For CUDA devices, data is first - created on CPU, then copied to the target device. - This function internally converts the list to a numpy array and calls from_numpy. - """ - # Input validation - if not isinstance(data, (list, tuple)): - raise TypeError( - f"Input data must be a list or tuple, got {type(data).__name__}" - ) - - if not data: - raise ValueError("Input data cannot be empty") - - # Determine target numpy dtype - # If dtype is specified, convert it to numpy dtype first - # This ensures the numpy array has the correct dtype from the start - if dtype is not None: - np_dtype = infinicore_to_numpy_dtype(dtype) - else: - np_dtype = None # Let numpy infer - - # Convert Python list to numpy array with correct dtype - # NumPy arrays can only be created on CPU - # Use np.array(..., copy=True, order='C') to efficiently: - # - Convert data type (if dtype is specified) - # - Create a copy (ensuring data ownership) - # - Ensure C-contiguous memory layout - if np_dtype is not None: - np_array = np.array(data, dtype=np_dtype, copy=True, order="C") - else: - np_array = np.array(data, copy=True, order="C") - - # Reuse from_numpy to create the tensor - # This avoids code duplication and ensures consistent behavior - return from_numpy(np_array, dtype=dtype, device=device) diff --git a/python/infinicore/utils.py b/python/infinicore/utils.py deleted file mode 100644 index fd39db0e1..000000000 --- a/python/infinicore/utils.py +++ /dev/null @@ -1,118 +0,0 @@ -import numpy as np -import torch - -import infinicore - -try: - import ml_dtypes -except ModuleNotFoundError: - ml_dtypes = None - - -def to_torch_dtype(infini_dtype): - """Convert infinicore data type to PyTorch data type""" - if infini_dtype == infinicore.float16: - return torch.float16 - elif infini_dtype == infinicore.float32: - return torch.float32 - elif infini_dtype == infinicore.float64: - return torch.float64 - elif infini_dtype == infinicore.bfloat16: - return torch.bfloat16 - elif infini_dtype == infinicore.int8: - return torch.int8 - elif infini_dtype == infinicore.int16: - return torch.int16 - elif infini_dtype == infinicore.int32: - return torch.int32 - elif infini_dtype == infinicore.int64: - return torch.int64 - elif infini_dtype == infinicore.uint8: - return torch.uint8 - elif infini_dtype == infinicore.bool: - return torch.bool - else: - raise ValueError(f"Unsupported infinicore dtype: {infini_dtype}") - - -def to_infinicore_dtype(torch_dtype): - """Convert PyTorch data type to infinicore data type""" - if torch_dtype == torch.float32: - return infinicore.float32 - elif torch_dtype == torch.float16: - return infinicore.float16 - elif torch_dtype == torch.bfloat16: - return infinicore.bfloat16 - elif torch_dtype == torch.int8: - return infinicore.int8 - elif torch_dtype == torch.int16: - return infinicore.int16 - elif torch_dtype == torch.int32: - return infinicore.int32 - elif torch_dtype == torch.int64: - return infinicore.int64 - elif torch_dtype == torch.uint8: - return infinicore.uint8 - else: - raise ValueError(f"Unsupported torch dtype: {torch_dtype}") - - -def numpy_to_infinicore_dtype(numpy_dtype): - """Convert numpy data type to infinicore data type""" - if numpy_dtype == np.float32: - return infinicore.float32 - elif numpy_dtype == np.float64: - return infinicore.float64 - elif numpy_dtype == np.float16: - return infinicore.float16 - elif hasattr(np, "bfloat16") and numpy_dtype == np.bfloat16: - return infinicore.bfloat16 - elif ml_dtypes is not None and numpy_dtype == ml_dtypes.bfloat16: - return infinicore.bfloat16 - elif numpy_dtype == np.int8: - return infinicore.int8 - elif numpy_dtype == np.int16: - return infinicore.int16 - elif numpy_dtype == np.int32: - return infinicore.int32 - elif numpy_dtype == np.int64: - return infinicore.int64 - elif numpy_dtype == np.uint8: - return infinicore.uint8 - elif numpy_dtype == np.bool_: - return infinicore.bool - else: - raise ValueError(f"Unsupported numpy dtype: {numpy_dtype}") - - -def infinicore_to_numpy_dtype(infini_dtype): - """Convert infinicore data type to numpy data type""" - if infini_dtype == infinicore.float32: - return np.float32 - elif infini_dtype == infinicore.float64: - return np.float64 - elif infini_dtype == infinicore.float16: - return np.float16 - elif infini_dtype == infinicore.int8: - return np.int8 - elif infini_dtype == infinicore.int16: - return np.int16 - elif infini_dtype == infinicore.bfloat16: - if hasattr(np, "bfloat16"): - return np.bfloat16 - if ml_dtypes is None: - raise ModuleNotFoundError( - "ml_dtypes is required for bfloat16 numpy conversion. " - "Please install ml_dtypes." - ) - return ml_dtypes.bfloat16 - elif infini_dtype == infinicore.int32: - return np.int32 - elif infini_dtype == infinicore.int64: - return np.int64 - elif infini_dtype == infinicore.uint8: - return np.uint8 - elif infini_dtype == infinicore.bool: - return np.bool_ - else: - raise ValueError(f"Unsupported infinicore dtype: {infini_dtype}") diff --git a/scripts/build_ntops.py b/scripts/build_ntops.py deleted file mode 100644 index 601249615..000000000 --- a/scripts/build_ntops.py +++ /dev/null @@ -1,46 +0,0 @@ -import concurrent.futures -import importlib -import pathlib - -from infiniop.ninetoothed.build import BUILD_DIRECTORY_PATH - -CURRENT_FILE_PATH = pathlib.Path(__file__) - -SRC_DIR_PATH = CURRENT_FILE_PATH.parent.parent / "src" - - -def _find_and_build_ops(): - ops_path = SRC_DIR_PATH / "infiniop" / "ops" - - with concurrent.futures.ProcessPoolExecutor() as executor: - futures = [] - - for op_dir in ops_path.iterdir(): - ninetoothed_path = op_dir / "ninetoothed" - - if not ninetoothed_path.is_dir(): - continue - - build_file = ninetoothed_path / "build.py" - if not build_file.exists(): - continue - - futures.append(executor.submit(_build, ninetoothed_path)) - - for future in concurrent.futures.as_completed(futures): - future.result() - - -def _build(ninetoothed_path): - module_path = ninetoothed_path / "build" - relative_path = module_path.relative_to(SRC_DIR_PATH) - import_name = ".".join(relative_path.parts) - module = importlib.import_module(import_name) - - module.build() - - -if __name__ == "__main__": - BUILD_DIRECTORY_PATH.mkdir(parents=True, exist_ok=True) - - _find_and_build_ops() diff --git a/scripts/format.py b/scripts/format.py deleted file mode 100644 index 70124144c..000000000 --- a/scripts/format.py +++ /dev/null @@ -1,299 +0,0 @@ -import argparse -import os -import subprocess -from pathlib import Path - -from colorama import Fore, Style - -# Supported file types and their corresponding formatter categories -SUPPORTED_FILES = { - ".h": "c", - ".hh": "c", - ".hpp": "c", - ".c": "c", - ".cc": "c", - ".cpp": "c", - ".cxx": "c", - ".cu": "c", - ".cuh": "c", - ".mlu": "c", - ".cl": "c", - ".py": "py", -} - - -def format_file(file: Path, check: bool, formatter) -> bool: - formatter = formatter.get(SUPPORTED_FILES.get(file.suffix, None), None) - if not formatter: - return True # Unsupported file type, skip - - formatted = True - try: - cmd = [] - if formatter.startswith("clang-format"): - cmd = [formatter, "-style=file", "-i", file] - if check: - cmd.insert(2, "-dry-run") - process = subprocess.run( - cmd, - capture_output=True, - text=True, - check=True, - ) - if process.stderr: - print(f"{Fore.YELLOW}{file} is not formatted.{Style.RESET_ALL}") - print( - f"Use {Fore.CYAN}{formatter} -style=file -i {file}{Style.RESET_ALL} to format it." - ) - formatted = False - else: - subprocess.run( - cmd, - capture_output=True, - text=True, - check=True, - ) - print(f"{Fore.CYAN}Formatted: {file}{Style.RESET_ALL}") - elif formatter == "black": - cmd = [formatter, file] - if check: - cmd.insert(1, "--check") - process = subprocess.run( - cmd, - capture_output=True, - text=True, - check=False, - ) - if process.returncode != 0: - print(f"{Fore.YELLOW}{file} is not formatted.{Style.RESET_ALL}") - print( - f"Use {Fore.CYAN}{formatter} {file}{Style.RESET_ALL} to format it." - ) - formatted = False - else: - subprocess.run( - cmd, - capture_output=True, - text=True, - check=True, - ) - print(f"{Fore.CYAN}Formatted: {file}{Style.RESET_ALL}") - elif formatter == "ruff": - ruff_cmd = ["ruff", "check", file] - ruff_format_cmd = ["ruff", "format", file] - try: - if check: - process = subprocess.run( - ruff_cmd, - capture_output=True, - text=True, - check=False, - ) - if process.returncode != 0: - print(f"{Fore.YELLOW}{file} has ruff issues.{Style.RESET_ALL}") - print( - f"Use {Fore.CYAN}ruff check --fix {file}{Style.RESET_ALL} to fix it." - ) - formatted = False - - ruff_format_cmd.insert(2, "--check") - process = subprocess.run( - ruff_format_cmd, - capture_output=True, - text=True, - check=False, - ) - if process.returncode != 0: - print( - f"{Fore.YELLOW}{file} is not ruff-formatted.{Style.RESET_ALL}" - ) - print( - f"Use {Fore.CYAN}ruff format {file}{Style.RESET_ALL} to format it." - ) - formatted = False - else: - ruff_cmd.insert(2, "--fix") - process = subprocess.run( - ruff_cmd, - capture_output=True, - text=True, - check=False, - ) - if process.returncode != 0: - print(f"{Fore.RED}Ruff failed for {file}.{Style.RESET_ALL}") - if process.stdout: - print(process.stdout, end="") - if process.stderr: - print(process.stderr, end="") - formatted = False - else: - print(f"{Fore.CYAN}Ruff fixed: {file}{Style.RESET_ALL}") - - process = subprocess.run( - ruff_format_cmd, - capture_output=True, - text=True, - check=False, - ) - if process.returncode != 0: - print( - f"{Fore.RED}Ruff format failed for {file}.{Style.RESET_ALL}" - ) - if process.stdout: - print(process.stdout, end="") - if process.stderr: - print(process.stderr, end="") - formatted = False - else: - print(f"{Fore.CYAN}Ruff formatted: {file}{Style.RESET_ALL}") - except FileNotFoundError: - print( - f"{Fore.RED}Formatter ruff not found, {file} skipped.{Style.RESET_ALL}" - ) - formatted = False - except FileNotFoundError: - print( - f"{Fore.RED}Formatter {formatter} not found, {file} skipped.{Style.RESET_ALL}" - ) - formatted = False - except subprocess.CalledProcessError as e: - print(f"{Fore.RED}Formatter {formatter} failed: {e}{Style.RESET_ALL}") - formatted = False - - if not ensure_single_trailing_newline(file, check): - formatted = False - - return formatted - - -def ensure_single_trailing_newline(file: Path, check: bool) -> bool: - """Ensure the file ends with exactly one LF byte.""" - content = file.read_bytes() - fixed = content.rstrip(b"\r\n") + b"\n" - if content == fixed: - return True - - if check: - print( - f"{Fore.YELLOW}{file} does not end with exactly one newline.{Style.RESET_ALL}" - ) - return False - - file.write_bytes(fixed) - print(f"{Fore.CYAN}Fixed trailing newline: {file}{Style.RESET_ALL}") - return True - - -def git_added_files(): - """Get all staged files""" - try: - # Use git diff --cached --name-only to get all files added to staging area - result = subprocess.run( - ["git", "diff", "--cached", "--diff-filter=AMR", "--name-only"], - capture_output=True, - text=True, - check=True, - ) - for file in result.stdout.splitlines(): - yield Path(file.strip()) - except subprocess.CalledProcessError as e: - print(f"{Fore.RED}Git diff failed: {e}{Style.RESET_ALL}") - - -def git_modified_since_ref(ref): - """Get list of files modified from the specified Git reference to the current state""" - try: - result = subprocess.run( - ["git", "diff", f"{ref}..", "--diff-filter=AMR", "--name-only"], - capture_output=True, - text=True, - check=True, - ) - for file in result.stdout.splitlines(): - yield Path(file.strip()) - except subprocess.CalledProcessError as e: - print(f"{Fore.RED}Git diff failed: {e}{Style.RESET_ALL}") - - -def list_files(paths): - """Recursively get all files under the specified paths""" - for path in paths: - if path.is_file(): - yield path - elif path.is_dir(): - for dirpath, _, filenames in os.walk(path): - for name in filenames: - yield Path(dirpath) / name - else: - print( - f"{Fore.RED}Error: {path} is not a file or directory.{Style.RESET_ALL}" - ) - - -def filter_in_path(file: Path, path) -> bool: - """Check if file is within the specified paths""" - for p in path: - if file.is_relative_to(p): - return True - return False - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--ref", type=str, help="Git reference (commit hash) to compare against." - ) - parser.add_argument( - "--path", nargs="*", type=Path, help="Files to format or check." - ) - parser.add_argument( - "--check", action="store_true", help="Check files without modifying them." - ) - parser.add_argument( - "--c", default="clang-format-16", help="C formatter (default: clang-format-16)" - ) - parser.add_argument("--py", default="ruff", help="Python formatter (default: ruff)") - args = parser.parse_args() - - if args.ref is None and args.path is None: - # Last commit. - print(f"{Fore.GREEN}Formatting git staged files.{Style.RESET_ALL}") - files = git_added_files() - - else: - if args.ref is None: - print(f"{Fore.GREEN}Formatting files in {args.path}.{Style.RESET_ALL}") - files = list_files(args.path) - elif args.path is None: - print( - f"{Fore.GREEN}Formatting git modified files from {args.ref}.{Style.RESET_ALL}" - ) - files = git_modified_since_ref(args.ref) - else: - print( - f"{Fore.GREEN}Formatting git modified files from {args.ref} in {args.path}.{Style.RESET_ALL}" - ) - files = ( - file - for file in git_modified_since_ref(args.ref) - if filter_in_path(file, args.path) - ) - - formatted = True - for file in files: - if not format_file( - file, - args.check, - { - "c": args.c, - "py": args.py, - }, - ): - formatted = False - - if not formatted: - exit(1) - - -if __name__ == "__main__": - main() diff --git a/scripts/install.py b/scripts/install.py deleted file mode 100644 index 98a448254..000000000 --- a/scripts/install.py +++ /dev/null @@ -1,28 +0,0 @@ -import os -import subprocess -import platform -import sys -from set_env import ( - set_env, - set_env_by_config, -) - -PROJECT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) -os.chdir(PROJECT_DIR) - -def run_cmd(cmd): - subprocess.run(cmd, text=True, encoding="utf-8", check=True, shell=True) - - -def install(xmake_config_flags=""): - set_env_by_config(xmake_config_flags) - run_cmd(f"xmake f -y {xmake_config_flags} -cv") - run_cmd("xmake -y") - run_cmd("xmake install -y") - run_cmd("xmake build -y infiniop-test") - run_cmd("xmake install -y infiniop-test") - - -if __name__ == "__main__": - set_env() - install(" ".join(sys.argv[1:])) diff --git a/scripts/metax_env.py b/scripts/metax_env.py deleted file mode 100644 index df35c57a8..000000000 --- a/scripts/metax_env.py +++ /dev/null @@ -1,67 +0,0 @@ -import os - - -def _first_existing_dir(paths: list[str]) -> str: - for p in paths: - if p and os.path.isdir(p): - return p - return "" - - -def _metax_toolkit_root(use_mc: bool) -> str: - """Return toolkit root for MetaX builds (MACA when use-mc; otherwise HPCC).""" - if use_mc: - for key in ("MACA_PATH", "MACA_HOME", "MACA_ROOT"): - v = os.environ.get(key, "").strip() - if v: - return v - return _first_existing_dir(["/opt/maca"]) - return _first_existing_dir(["/opt/hpcc"]) - - -def _prepend_path_var(name: str, prefixes: list[str]) -> None: - """Prepend colon-separated *prefixes* to env var *name* (POSIX).""" - if not prefixes: - return - chunk = ":".join(prefixes) - cur = os.environ.get(name, "") - os.environ[name] = f"{chunk}:{cur}" if cur else chunk - - -def set_env_for_metax_gpu( - flags: str, - *, - parse_xmake_cli_flag_values, - truthy_flag_value, -) -> None: - """ - Prepend compiler include paths needed when building ATen-enabled C++ against torch headers. - - This chooses paths based on xmake backend flags (e.g. --metax-gpu) and toolkit selection - (e.g. MetaX HPCC vs MACA when --use-mc=y). - """ - d = parse_xmake_cli_flag_values(flags) - if not truthy_flag_value(d.get("aten", "n")): - return - - if truthy_flag_value(d.get("metax-gpu", "n")): - use_mc = truthy_flag_value(d.get("use-mc", "n")) - root = _metax_toolkit_root(use_mc=use_mc) - if not root: - return - dirs = [ - os.path.join(root, "tools", "cu-bridge", "include"), - os.path.join(root, "include", "hcr"), - # cu-bridge cuComplex.h includes "hcComplex.h" from HPCC include/common - os.path.join(root, "include", "common"), - # cu-bridge cusparse wrapper includes "hcsparse.h" under include/hcsparse - os.path.join(root, "include", "hcsparse"), - # cu-bridge cublasLt wrapper includes "hcblasLt.h" under include/hcblas - os.path.join(root, "include", "hcblas"), - # cu-bridge cusolver wrapper includes "hcsolver_common.h" under include/hcsolver - os.path.join(root, "include", "hcsolver"), - os.path.join(root, "include"), - ] - for var in ("CPATH", "CPLUS_INCLUDE_PATH", "C_INCLUDE_PATH"): - _prepend_path_var(var, dirs) - return diff --git a/scripts/python_test.py b/scripts/python_test.py deleted file mode 100644 index 4879da32a..000000000 --- a/scripts/python_test.py +++ /dev/null @@ -1,66 +0,0 @@ -import os -import subprocess -from set_env import set_env -import sys - -PROJECT_DIR = os.path.abspath( - os.path.join(os.path.dirname(__file__), "..", "test", "infiniop") -) -os.chdir(PROJECT_DIR) - - -def run_tests(args): - failed = [] - for test in [ - "add.py", - "attention.py", - "causal_softmax.py", - "clip.py", - "conv.py", - # "dequantize_awq.py", - # "dequantize_gptq.py", - "gelu.py", - "gemm.py", - # "layer_norm.py", - "logsoftmax.py", - # "lp_norm.py", - "mul.py", - "ones.py", - "random_sample.py", - "rearrange.py", - "relu.py", - "rms_norm.py", - "rope.py", - "sigmoid.py", - # "softmax.py", - "softplus.py", - "sub.py", - "swiglu.py", - "tanh.py", - "topkrouter.py", - "topksoftmax.py", - "zeros.py", - # "paged_attention.py", - # "paged_caching.py", - # "paged_attention_prefill.py" - "cross_entropy.py", - ]: - result = subprocess.run( - f"python {test} {args} --debug", text=True, encoding="utf-8", shell=True - ) - if result.returncode != 0: - failed.append(test) - - return failed - - -if __name__ == "__main__": - set_env() - failed = run_tests(" ".join(sys.argv[1:])) - if len(failed) == 0: - print("\033[92mAll tests passed!\033[0m") - else: - print("\033[91mThe following tests failed:\033[0m") - for test in failed: - print(f"\033[91m - {test}\033[0m") - exit(len(failed)) diff --git a/scripts/set_env.py b/scripts/set_env.py deleted file mode 100644 index f489c28ad..000000000 --- a/scripts/set_env.py +++ /dev/null @@ -1,65 +0,0 @@ -import os -import platform - -from metax_env import set_env_for_metax_gpu - - -def _parse_xmake_cli_flag_values(flags: str): - """Parse a string like '--metax-gpu=y --aten=y' into {key: value}.""" - parts = flags.replace("=", " ").split() - d = {} - i = 0 - n = len(parts) - while i < n: - p = parts[i] - if p.startswith("--") and len(p) > 2: - key = p[2:].lower() - i += 1 - if i < n and not parts[i].startswith("--"): - d[key] = parts[i].lower() - i += 1 - else: - d[key] = "y" - else: - i += 1 - return d - - -def _truthy_flag_value(v: str) -> bool: - return v in ("y", "yes", "true", "1", "on") - - -def set_env_by_config(flags: str) -> None: - """Set environment variables for InfiniCore builds with xmake config flags.""" - d = _parse_xmake_cli_flag_values(flags) - if _truthy_flag_value(d.get("metax-gpu", "n")): - set_env_for_metax_gpu( - flags, - parse_xmake_cli_flag_values=_parse_xmake_cli_flag_values, - truthy_flag_value=_truthy_flag_value, - ) - else: - pass - - -def set_env(): - if os.environ.get("INFINI_ROOT") == None: - os.environ["INFINI_ROOT"] = os.path.expanduser("~/.infini") - - if platform.system() == "Windows": - new_path = os.path.expanduser(os.environ.get("INFINI_ROOT") + "/bin") - if new_path not in os.environ.get("PATH", ""): - os.environ["PATH"] = f"{new_path};{os.environ.get('PATH', '')}" - - elif platform.system() == "Linux": - new_path = os.path.expanduser(os.environ.get("INFINI_ROOT") + "/bin") - if new_path not in os.environ.get("PATH", ""): - os.environ["PATH"] = f"{new_path}:{os.environ.get('PATH', '')}" - - new_lib_path = os.path.expanduser(os.environ.get("INFINI_ROOT") + "/lib") - if new_lib_path not in os.environ.get("LD_LIBRARY_PATH", ""): - os.environ["LD_LIBRARY_PATH"] = ( - f"{new_lib_path}:{os.environ.get('LD_LIBRARY_PATH', '')}" - ) - else: - raise RuntimeError("Unsupported platform.") diff --git a/scripts/set_env_linux.sh b/scripts/set_env_linux.sh deleted file mode 100644 index f97fce1c8..000000000 --- a/scripts/set_env_linux.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/bin/bash - -# Define the content to be added -content='export INFINI_ROOT="$HOME/.infini" -export LD_LIBRARY_PATH="$INFINI_ROOT/lib:$LD_LIBRARY_PATH"' - -# Check if bashrc file exists -bashrc_file="$HOME/.bashrc" - -if [ ! -f "$bashrc_file" ]; then - echo "Creating $bashrc_file file" - touch "$bashrc_file" -fi - -# Check if the content already exists -if grep -q "export INFINI_ROOT=" "$bashrc_file"; then - echo "INFINI_ROOT configuration already exists in $bashrc_file" -else - echo "Adding configuration to $bashrc_file" - echo "$content" >> "$bashrc_file" - echo "Configuration added successfully" -fi - -# Reload bashrc -echo "Reloading $bashrc_file" -source "$bashrc_file" - -echo "Done! INFINI_ROOT and LD_LIBRARY_PATH have been set" diff --git a/setup.py b/setup.py deleted file mode 100644 index e79e41f53..000000000 --- a/setup.py +++ /dev/null @@ -1,34 +0,0 @@ -import subprocess -from setuptools import setup, find_packages -from setuptools.command.build import build - -def run_xmake_build(): - print("Running xmake build...") - subprocess.run(["xmake", "build"], check=True) - subprocess.run(["xmake", "install"], check=True) - subprocess.run(["xmake", "build", "-y", "_infinicore"], check=True) - subprocess.run(["xmake", "install", "_infinicore"], check=True) - -class Build(build): - def run(self): - run_xmake_build() - super().run() - -setup( - # 1. Find main packages and manually add test/framework packages - packages=find_packages(where="python") + [ - "infinicore.test", - "infinicore.test.framework" - ], - - # 2. Directory mappings - package_dir={ - "": "python", # Root package is under python/ directory - "infinicore.test": "test/infinicore" # Intermediate package mapping - }, - - # 3. Register commands - cmdclass={ - "build": Build - } -) diff --git a/src/bridge/infini/rt.hpp b/src/bridge/infini/rt.hpp deleted file mode 100644 index c5bec7ecd..000000000 --- a/src/bridge/infini/rt.hpp +++ /dev/null @@ -1,72 +0,0 @@ -#pragma once - -#include "infinirt.h" - -#include - -namespace infinicore::bridge::infini::rt { - -inline infiniStatus_t translate(::infini::rt::runtime::Error error) { - switch (error) { - case ::infini::rt::runtime::kSuccess: - return INFINI_STATUS_SUCCESS; - default: - return INFINI_STATUS_INTERNAL_ERROR; - } -} - -inline constexpr ::infini::rt::Device::Type translate_to(infiniDevice_t device) { - switch (device) { - case INFINI_DEVICE_CPU: - return ::infini::rt::Device::Type::kCpu; - case INFINI_DEVICE_NVIDIA: - return ::infini::rt::Device::Type::kNvidia; - case INFINI_DEVICE_CAMBRICON: - return ::infini::rt::Device::Type::kCambricon; - case INFINI_DEVICE_ASCEND: - return ::infini::rt::Device::Type::kAscend; - case INFINI_DEVICE_METAX: - return ::infini::rt::Device::Type::kMetax; - case INFINI_DEVICE_MOORE: - return ::infini::rt::Device::Type::kMoore; - case INFINI_DEVICE_ILUVATAR: - return ::infini::rt::Device::Type::kIluvatar; - case INFINI_DEVICE_HYGON: - return ::infini::rt::Device::Type::kHygon; - default: - return ::infini::rt::Device::Type::kCount; - } -} - -inline constexpr infiniDevice_t translate_from(::infini::rt::Device::Type device) { - switch (device) { - case ::infini::rt::Device::Type::kCpu: - return INFINI_DEVICE_CPU; - case ::infini::rt::Device::Type::kNvidia: - return INFINI_DEVICE_NVIDIA; - case ::infini::rt::Device::Type::kCambricon: - return INFINI_DEVICE_CAMBRICON; - case ::infini::rt::Device::Type::kAscend: - return INFINI_DEVICE_ASCEND; - case ::infini::rt::Device::Type::kMetax: - return INFINI_DEVICE_METAX; - case ::infini::rt::Device::Type::kMoore: - return INFINI_DEVICE_MOORE; - case ::infini::rt::Device::Type::kIluvatar: - return INFINI_DEVICE_ILUVATAR; - case ::infini::rt::Device::Type::kHygon: - return INFINI_DEVICE_HYGON; - default: - return INFINI_DEVICE_TYPE_COUNT; - } -} - -inline ::infini::rt::runtime::Stream translate_to(infinirtStream_t stream) { - return reinterpret_cast<::infini::rt::runtime::Stream>(stream); -} - -inline infinirtStream_t translate_from(::infini::rt::runtime::Stream stream) { - return reinterpret_cast(stream); -} - -} // namespace infinicore::bridge::infini::rt diff --git a/src/infiniccl-test/infiniccl_test.cpp b/src/infiniccl-test/infiniccl_test.cpp deleted file mode 100644 index 0aa898484..000000000 --- a/src/infiniccl-test/infiniccl_test.cpp +++ /dev/null @@ -1,199 +0,0 @@ -#include "infiniccl_test.hpp" - -#include -#include -#include -#include -#include -#include - -#define TEST_INFINI(API__) CHECK_API_OR(API__, INFINI_STATUS_SUCCESS, return 1) -#define TEST_INFINI_THREAD(API__) CHECK_API_OR(API__, INFINI_STATUS_SUCCESS, return nullptr) - -const size_t MAX_COUNT = 8ULL * 1024 * 1024; -// const size_t MAX_COUNT = 512 * 1024; // for metax - -const size_t TEST_COUNTS[] = { - 128, - 1024, - 4 * 1024, - MAX_COUNT, -}; - -const infiniDtype_t TEST_DTYPES[] = {INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16}; - -const size_t WARM_UPS = 10; - -const size_t ITERATIONS = 100; - -struct ThreadArgs { - int rank; - int device_id; - infinicclComm_t comm; - infiniDevice_t device_type; - infiniDtype_t dtype; - size_t count; - const void *data; - const void *ans; - int *result; - double *time; -}; - -void setData(infiniDtype_t dtype, void *data, size_t count, float val) { - switch (dtype) { - case INFINI_DTYPE_F32: - for (size_t i = 0; i < count; i++) { - ((float *)data)[i] = val; - } - break; - - case INFINI_DTYPE_F16: - for (size_t i = 0; i < count; i++) { - ((fp16_t *)data)[i] = utils::cast(val); - } - break; - case INFINI_DTYPE_BF16: - for (size_t i = 0; i < count; i++) { - ((bf16_t *)data)[i] = utils::cast(val); - } - break; - default: - std::abort(); - break; - } -} - -template -int checkData(const T *actual_, const T *expected_, size_t count) { - int failed = 0; - for (size_t i = 0; i < count; i++) { - if constexpr (std::is_same::value) { - float actual = utils::cast(actual_[i]); - float expected = utils::cast(expected_[i]); - if (std::abs(actual - expected) > 1e-4) { - failed += 1; - } - } else if constexpr (std::is_same::value) { - float actual = utils::cast(actual_[i]); - float expected = utils::cast(expected_[i]); - if (std::abs(actual - expected) > 1e-4) { - failed += 1; - } - } else { - if (std::abs(actual_[i] - expected_[i]) > 1e-4) { - failed += 1; - } - } - } - return failed; -} - -int checkData(const void *actual, const void *expected, infiniDtype_t dtype, size_t count) { - switch (dtype) { - case INFINI_DTYPE_F32: - return checkData((const float *)actual, (const float *)expected, count); - case INFINI_DTYPE_F16: - return checkData((const fp16_t *)actual, (const fp16_t *)expected, count); - case INFINI_DTYPE_BF16: - return checkData((const bf16_t *)actual, (const bf16_t *)expected, count); - default: - std::abort(); - return 1; - } -} - -void *testAllReduceThread(void *arg) { - ThreadArgs *args = (ThreadArgs *)arg; - *(args->result) = 1; - TEST_INFINI_THREAD(infinirtSetDevice(args->device_type, args->device_id)); - infinirtStream_t stream; - TEST_INFINI_THREAD(infinirtStreamCreate(&stream)); - void *output = std::malloc(args->count * infiniSizeOf(args->dtype)); - std::memset(output, 0, args->count * infiniSizeOf(args->dtype)); - void *buf; - TEST_INFINI_THREAD(infinirtMalloc(&buf, args->count * infiniSizeOf(args->dtype))); - TEST_INFINI_THREAD(infinirtMemcpy(buf, args->data, args->count * infiniSizeOf(args->dtype), INFINIRT_MEMCPY_H2D)); - TEST_INFINI_THREAD(infinicclAllReduce(buf, buf, args->count, args->dtype, INFINICCL_SUM, args->comm, stream)); - TEST_INFINI_THREAD(infinirtStreamSynchronize(stream)); - TEST_INFINI_THREAD(infinirtMemcpy(output, buf, args->count * infiniSizeOf(args->dtype), INFINIRT_MEMCPY_D2H)); - - if (checkData(output, args->ans, args->dtype, args->count) != 0) { - std::free(output); - infinirtFree(buf); - infinirtStreamDestroy(stream); - return nullptr; - } - for (size_t i = 0; i < WARM_UPS; i++) { - TEST_INFINI_THREAD(infinicclAllReduce(buf, buf, args->count, args->dtype, INFINICCL_SUM, args->comm, stream)); - } - TEST_INFINI_THREAD(infinirtStreamSynchronize(stream)); - - // measure time - auto start = std::chrono::high_resolution_clock::now(); - for (size_t i = 0; i < ITERATIONS; i++) { - TEST_INFINI_THREAD(infinicclAllReduce(buf, buf, args->count, args->dtype, INFINICCL_SUM, args->comm, stream)); - } - TEST_INFINI_THREAD(infinirtStreamSynchronize(stream)); - auto end = std::chrono::high_resolution_clock::now(); - double elapsed_ms = std::chrono::duration(end - start).count(); - *args->time = elapsed_ms / ITERATIONS; - - *args->result = 0; - - std::free(output); - infinirtFree(buf); - infinirtStreamDestroy(stream); - return nullptr; -} - -int testAllReduce(infiniDevice_t device_type, int ndevice) { - std::vector thread_args(ndevice); - std::vector comms(ndevice); - std::vector threads(ndevice); - std::vector device_ids(ndevice); - std::vector results(ndevice); - std::vector times(ndevice); - void *data = std::malloc(MAX_COUNT * sizeof(float)); // Use float as max dtype size - void *ans = std::malloc(MAX_COUNT * sizeof(float)); - - for (int i = 0; i < ndevice; i++) { - device_ids[i] = i; - } - - for (infiniDtype_t dtype : TEST_DTYPES) { - setData(dtype, data, MAX_COUNT, 1.0f); - setData(dtype, ans, MAX_COUNT, 1.0f * ndevice); - for (size_t count : TEST_COUNTS) { - TEST_INFINI(infinicclCommInitAll(device_type, comms.data(), ndevice, device_ids.data())); - std::cout << "Testing AllReduce with " << count << " elements of " << infiniDtypeToString(dtype) << std::endl; - for (int rank = 0; rank < ndevice; rank++) { - thread_args[rank] = {rank, device_ids[rank], comms[rank], device_type, dtype, count, data, ans, &results[rank], ×[rank]}; - pthread_create(&threads[rank], NULL, testAllReduceThread, &thread_args[rank]); - } - for (int rank = 0; rank < ndevice; rank++) { - pthread_join(threads[rank], NULL); - } - int failed = std::accumulate(results.begin(), results.end(), 0); - for (int rank = 0; rank < ndevice; rank++) { - if (results[rank] != 0) { - std::cout << "Rank " << rank << ": incorrect results." << std::endl; - } else { - std::cout << "Rank " << rank << ": " << times[rank] << " ms." << std::endl; - } - } - - if (failed > 0) { - std::cout << "Failed with " << failed << " errors." << std::endl - << std::endl; - std::free(data); - std::free(ans); - return 1; - } - std::cout << std::endl; - } - } - - std::free(data); - std::free(ans); - return 0; -} diff --git a/src/infiniccl-test/infiniccl_test.hpp b/src/infiniccl-test/infiniccl_test.hpp deleted file mode 100644 index 0ea4bbed6..000000000 --- a/src/infiniccl-test/infiniccl_test.hpp +++ /dev/null @@ -1,10 +0,0 @@ -#ifndef INFINICCL_TEST_HPP -#define INFINICCL_TEST_HPP - -#include - -#include "../utils.h" - -int testAllReduce(infiniDevice_t device_type, int ndevice); - -#endif // INFINICCL_TEST_HPP diff --git a/src/infiniccl-test/main.cpp b/src/infiniccl-test/main.cpp deleted file mode 100644 index 8126aa15a..000000000 --- a/src/infiniccl-test/main.cpp +++ /dev/null @@ -1,80 +0,0 @@ -#include "infiniccl_test.hpp" - -#include - -struct ParsedArgs { - infiniDevice_t device_type; -}; - -void printUsage() { - std::cout << "Usage:" << std::endl - << std::endl; - std::cout << "infiniccl-test --" << std::endl - << std::endl; - std::cout << " --" << std::endl; - std::cout << " Specify the device type --(nvidia|cambricon|ascend|metax|moore|iluvatar|qy|kunlun|hygon|ali)." << std::endl - << std::endl; - std::cout << "The program will run tests on all visible devices of the specified device type." - << " Use Environmental Variables such as CUDA_VSIBLE_DEVICES to limit visible device IDs."; - exit(-1); -} - -#define PARSE_DEVICE(FLAG, DEVICE) \ - if (arg == FLAG) { \ - args.device_type = DEVICE; \ - } - -ParsedArgs parseArgs(int argc, char *argv[]) { - if (argc != 2) { - printUsage(); - } - - if (std::string(argv[1]) == "--help" || std::string(argv[1]) == "-h") { - printUsage(); - } - - ParsedArgs args; - try { - std::string arg = argv[1]; - // clang-format off - PARSE_DEVICE("--nvidia", INFINI_DEVICE_NVIDIA) - else PARSE_DEVICE("--cambricon", INFINI_DEVICE_CAMBRICON) - else PARSE_DEVICE("--ascend", INFINI_DEVICE_ASCEND) - else PARSE_DEVICE("--metax", INFINI_DEVICE_METAX) - else PARSE_DEVICE("--moore", INFINI_DEVICE_MOORE) - else PARSE_DEVICE("--iluvatar", INFINI_DEVICE_ILUVATAR) - else PARSE_DEVICE("--qy", INFINI_DEVICE_QY) - else PARSE_DEVICE("--kunlun", INFINI_DEVICE_KUNLUN) - else PARSE_DEVICE("--hygon", INFINI_DEVICE_HYGON) - else PARSE_DEVICE("--ali", INFINI_DEVICE_ALI) - else { - printUsage(); - } - // clang-format on - - } catch (const std::exception &) { - printUsage(); - } - - return args; -} - -int main(int argc, char *argv[]) { - ParsedArgs args = parseArgs(argc, argv); - int ndevice = 0; - infinirtInit(); - if (infinirtGetDeviceCount(args.device_type, &ndevice) != INFINI_STATUS_SUCCESS) { - std::cout << "Failed to get device count" << std::endl; - return -1; - } - if (ndevice == 0) { - std::cout << "No devices found. Tests skipped." << std::endl; - return 0; - } else { - std::cout << "Found " << ndevice << " devices. Running tests..." << std::endl; - } - - int failed = 0; - failed += testAllReduce(args.device_type, ndevice); - return failed; -} diff --git a/src/infiniccl/ascend/infiniccl_ascend.cc b/src/infiniccl/ascend/infiniccl_ascend.cc deleted file mode 100644 index 76f4c961b..000000000 --- a/src/infiniccl/ascend/infiniccl_ascend.cc +++ /dev/null @@ -1,100 +0,0 @@ -#include "infiniccl_ascend.h" - -#include "../../utils.h" - -#include -#include - -#include -#include - -#define CHECK_HCCL(API__) CHECK_INTERNAL(API__, HCCL_SUCCESS) - -inline aclrtStream getAscendStream(infinirtStream_t stream) { - if (stream == nullptr) { - return 0; - } - return static_cast(stream); -} - -inline HcclComm getHcclComm(infinicclComm_t comm) { - return static_cast(comm->comm); -} - -inline HcclDataType getAscendDtype(infiniDtype_t datatype) { - switch (datatype) { - case INFINI_DTYPE_F32: - return HCCL_DATA_TYPE_FP32; - case INFINI_DTYPE_F16: - return HCCL_DATA_TYPE_FP16; - case INFINI_DTYPE_BF16: - return HCCL_DATA_TYPE_BFP16; - default: - std::cerr << "Unsupported data type: " << datatype << std::endl; - std::abort(); - return HCCL_DATA_TYPE_FP16; - } -} - -inline HcclReduceOp getHcclRedOp(infinicclReduceOp_t op) { - switch (op) { - case INFINICCL_SUM: - return HCCL_REDUCE_SUM; - case INFINICCL_PROD: - return HCCL_REDUCE_PROD; - case INFINICCL_MAX: - return HCCL_REDUCE_MAX; - case INFINICCL_MIN: - return HCCL_REDUCE_MIN; - default: - std::abort(); - return HCCL_REDUCE_SUM; - } -} - -namespace infiniccl::ascend { - -infiniStatus_t commInitAll( - infinicclComm_t *comms, - int ndevice, - const int *device_ids) { - // Ascend requires all devices to be initialized before calling HcclCommInitAll. - for (int i = ndevice - 1; i >= 0; i--) { - aclrtSetDevice(device_ids[i]); - } - - std::vector hccl_comms(ndevice); - CHECK_HCCL(HcclCommInitAll(ndevice, (int32_t *)device_ids, hccl_comms.data())); - - for (int i = 0; i < ndevice; i++) { - comms[i] = new InfinicclComm{INFINI_DEVICE_ASCEND, device_ids[i], (void *)(hccl_comms[i])}; - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t commDestroy(infinicclComm_t comm) { - CHECK_HCCL(HcclCommDestroy(getHcclComm(comm))); - delete comm; - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t allReduce( - void *sendbuf, - void *recvbuf, - size_t count, - infiniDtype_t datatype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream) { - - CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16); - - CHECK_HCCL(HcclAllReduce(sendbuf, recvbuf, (uint64_t)count, - getAscendDtype(datatype), getHcclRedOp(op), - getHcclComm(comm), getAscendStream(stream))); - - return INFINI_STATUS_SUCCESS; -} -INFINICCL_DEVICE_API_UNSUPPORTED_EP_COLLECTIVES(); -} // namespace infiniccl::ascend diff --git a/src/infiniccl/ascend/infiniccl_ascend.h b/src/infiniccl/ascend/infiniccl_ascend.h deleted file mode 100644 index c15f99dc7..000000000 --- a/src/infiniccl/ascend/infiniccl_ascend.h +++ /dev/null @@ -1,12 +0,0 @@ -#ifndef INFINICCL_ASCEND_H_ -#define INFINICCL_ASCEND_H_ - -#include "../infiniccl_impl.h" - -#if defined(ENABLE_ASCEND_API) && defined(ENABLE_CCL) -INFINICCL_DEVICE_API_IMPL(ascend) -#else -INFINICCL_DEVICE_API_NOOP(ascend) -#endif - -#endif /* INFINICCL_ASCEND_H_ */ diff --git a/src/infiniccl/cambricon/infiniccl_cambricon.cc b/src/infiniccl/cambricon/infiniccl_cambricon.cc deleted file mode 100644 index 4bd1ca402..000000000 --- a/src/infiniccl/cambricon/infiniccl_cambricon.cc +++ /dev/null @@ -1,104 +0,0 @@ -#include "infiniccl_cambricon.h" - -#include "../../utils.h" -#include -#include -#include -#include - -#define CHECK_CNCL(API__) CHECK_INTERNAL(API__, CNCL_RET_SUCCESS) - -inline cnrtQueue_t getCambriconStream(infinirtStream_t stream) { - if (stream == nullptr) { - return (cnrtQueue_t)(0); - } - return static_cast(stream); -} - -inline cnclComm_t getCnclComm(infinicclComm_t comm) { - return static_cast(comm->comm); -} - -inline cnclDataType_t getCnclDtype(infiniDtype_t datatype) { - switch (datatype) { - case INFINI_DTYPE_F32: - return cnclFloat32; - case INFINI_DTYPE_F16: - return cnclFloat16; - case INFINI_DTYPE_BF16: - return cnclBfloat16; - default: - std::cerr << "Unsupported data type: " << datatype << std::endl; - std::abort(); - return cnclFloat16; - } -} - -inline cnclReduceOp_t getCnclRedOp(infinicclReduceOp_t op) { - switch (op) { - case INFINICCL_SUM: - return cnclSum; - case INFINICCL_PROD: - return cnclProd; - case INFINICCL_MAX: - return cnclMax; - case INFINICCL_MIN: - return cnclMin; - default: - std::abort(); - return cnclSum; - } -} - -namespace infiniccl::cambricon { - -infiniStatus_t commInitAll( - infinicclComm_t *comms, - int ndevice, - const int *device_ids) { - - std::vector cncl_comms(ndevice); - std::vector rank_list(ndevice); - - for (int i = 0; i < ndevice; i++) { - rank_list[i] = i; - CHECK_INTERNAL(cnrtSetDevice(device_ids[i]), cnrtSuccess); - } - - CHECK_CNCL(cnclInitComms(cncl_comms.data(), ndevice, - (int const *)device_ids, rank_list.data(), - ndevice, nullptr)); - - for (int i = 0; i < ndevice; i++) { - comms[i] = new InfinicclComm{INFINI_DEVICE_CAMBRICON, device_ids[i], (void *)(cncl_comms[i])}; - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t commDestroy(infinicclComm_t comm) { - CHECK_CNCL(cnclFreeComm(getCnclComm(comm))); - delete comm; - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t allReduce( - void *sendbuf, - void *recvbuf, - size_t count, - infiniDtype_t datatype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream) { - - CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16); - - CHECK_CNCL(cnclAllReduce(sendbuf, recvbuf, count, getCnclDtype(datatype), - getCnclRedOp(op), getCnclComm(comm), - getCambriconStream(stream))); - - return INFINI_STATUS_SUCCESS; -} - -INFINICCL_DEVICE_API_UNSUPPORTED_EP_COLLECTIVES(); -} // namespace infiniccl::cambricon diff --git a/src/infiniccl/cambricon/infiniccl_cambricon.h b/src/infiniccl/cambricon/infiniccl_cambricon.h deleted file mode 100644 index ad2a160fd..000000000 --- a/src/infiniccl/cambricon/infiniccl_cambricon.h +++ /dev/null @@ -1,12 +0,0 @@ -#ifndef INFINICCL_CAMBRICON_H_ -#define INFINICCL_CAMBRICON_H_ - -#include "../infiniccl_impl.h" - -#if defined(ENABLE_CAMBRICON_API) && defined(ENABLE_CCL) -INFINICCL_DEVICE_API_IMPL(cambricon) -#else -INFINICCL_DEVICE_API_NOOP(cambricon) -#endif - -#endif /* INFINICCL_CAMBRICON_H_ */ diff --git a/src/infiniccl/cuda/infiniccl_cuda.cu b/src/infiniccl/cuda/infiniccl_cuda.cu deleted file mode 100644 index 6a8442d21..000000000 --- a/src/infiniccl/cuda/infiniccl_cuda.cu +++ /dev/null @@ -1,220 +0,0 @@ -#include "infiniccl_cuda.h" - -#include -#include -#include -#include - -#include "../../utils.h" - -#define CHECK_NCCL(API__) CHECK_INTERNAL(API__, ncclSuccess) - -inline cudaStream_t getCudaStream(infinirtStream_t stream) { - if (stream == nullptr) { - return 0; - } - return static_cast(stream); -} - -inline ncclDataType_t getNcclDtype(infiniDtype_t datatype) { - switch (datatype) { - case INFINI_DTYPE_I32: - return ncclInt32; - case INFINI_DTYPE_I64: - return ncclInt64; - case INFINI_DTYPE_U32: - return ncclUint32; - case INFINI_DTYPE_U64: - return ncclUint64; - case INFINI_DTYPE_F32: - return ncclFloat; - case INFINI_DTYPE_F16: - return ncclHalf; - case INFINI_DTYPE_BF16: - return ncclBfloat16; - default: - std::abort(); - return ncclHalf; - } -} - -inline ncclRedOp_t getNcclRedOp(infinicclReduceOp_t op) { - switch (op) { - case INFINICCL_SUM: - return ncclSum; - case INFINICCL_PROD: - return ncclProd; - case INFINICCL_MAX: - return ncclMax; - case INFINICCL_MIN: - return ncclMin; - case INFINICCL_AVG: - return ncclAvg; - default: - std::abort(); - return ncclSum; - } -} - -inline ncclComm_t getNcclComm(infinicclComm_t comm) { - return static_cast(comm->comm); -} - -namespace infiniccl::cuda { - -infiniStatus_t commInitAll( - infinicclComm_t *comms, - int ndevice, - const int *device_ids) { - - std::vector nccl_comms(ndevice); - CHECK_NCCL(ncclCommInitAll(nccl_comms.data(), ndevice, (int const *)device_ids)); - - for (int i = 0; i < ndevice; i++) { - comms[i] = new InfinicclComm{INFINI_DEVICE_NVIDIA, device_ids[i], (void *)(nccl_comms[i]), i, ndevice}; - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t commDestroy(infinicclComm_t comm) { - CHECK_NCCL(ncclCommDestroy(getNcclComm(comm))); - delete comm; - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t groupStart(infinicclComm_t) { - CHECK_NCCL(ncclGroupStart()); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t groupEnd(infinicclComm_t) { - CHECK_NCCL(ncclGroupEnd()); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t allReduce( - void *sendbuf, - void *recvbuf, - size_t count, - infiniDtype_t datatype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream) { - - CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16); - - CHECK_NCCL(ncclAllReduce(sendbuf, recvbuf, count, getNcclDtype(datatype), - getNcclRedOp(op), getNcclComm(comm), getCudaStream(stream))); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t allGather( - void *sendbuf, - void *recvbuf, - size_t send_count, - infiniDtype_t datatype, - infinicclComm_t comm, - infinirtStream_t stream) { - - CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, - INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_U32, INFINI_DTYPE_U64); - - CHECK_NCCL(ncclAllGather(sendbuf, recvbuf, send_count, getNcclDtype(datatype), - getNcclComm(comm), getCudaStream(stream))); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t allGatherV( - void *sendbuf, - void *recvbuf, - const size_t *recv_counts, - int nranks, - infiniDtype_t datatype, - infinicclComm_t comm, - infinirtStream_t stream) { - - CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, - INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_U32, INFINI_DTYPE_U64); - CHECK_OR_DO(nranks == comm->world_size, return INFINI_STATUS_BAD_PARAM); - - auto cuda_stream = getCudaStream(stream); - ncclComm_t nccl_comm = getNcclComm(comm); - ncclDataType_t nccl_dtype = getNcclDtype(datatype); - size_t offset = 0; - - CHECK_NCCL(ncclGroupStart()); - for (int root = 0; root < nranks; ++root) { - CHECK_NCCL(ncclBroadcast( - sendbuf, - static_cast(recvbuf) + offset, - recv_counts[root], - nccl_dtype, - root, - nccl_comm, - cuda_stream)); - offset += recv_counts[root] * infiniSizeOf(datatype); - } - CHECK_NCCL(ncclGroupEnd()); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t reduceScatter( - void *sendbuf, - void *recvbuf, - size_t recv_count, - infiniDtype_t datatype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream) { - - CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, - INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_U32, INFINI_DTYPE_U64); - - CHECK_NCCL(ncclReduceScatter(sendbuf, recvbuf, recv_count, getNcclDtype(datatype), - getNcclRedOp(op), getNcclComm(comm), getCudaStream(stream))); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t reduceScatterV( - void *sendbuf, - void *recvbuf, - const size_t *send_counts, - int nranks, - infiniDtype_t datatype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream) { - - CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, - INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_U32, INFINI_DTYPE_U64); - CHECK_OR_DO(nranks == comm->world_size, return INFINI_STATUS_BAD_PARAM); - - auto cuda_stream = getCudaStream(stream); - ncclComm_t nccl_comm = getNcclComm(comm); - ncclDataType_t nccl_dtype = getNcclDtype(datatype); - ncclRedOp_t nccl_op = getNcclRedOp(op); - size_t offset = 0; - - CHECK_NCCL(ncclGroupStart()); - for (int root = 0; root < nranks; ++root) { - CHECK_NCCL(ncclReduce( - static_cast(sendbuf) + offset, - recvbuf, - send_counts[root], - nccl_dtype, - nccl_op, - root, - nccl_comm, - cuda_stream)); - offset += send_counts[root] * infiniSizeOf(datatype); - } - CHECK_NCCL(ncclGroupEnd()); - - return INFINI_STATUS_SUCCESS; -} -} // namespace infiniccl::cuda diff --git a/src/infiniccl/cuda/infiniccl_cuda.h b/src/infiniccl/cuda/infiniccl_cuda.h deleted file mode 100644 index 9bba02345..000000000 --- a/src/infiniccl/cuda/infiniccl_cuda.h +++ /dev/null @@ -1,13 +0,0 @@ -#ifndef INFINICCL_CUDA_H_ -#define INFINICCL_CUDA_H_ - -#include "../infiniccl_impl.h" - -// Windows does not support CUDA -#if (defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API)) && defined(ENABLE_CCL) && !defined(_WIN32) -INFINICCL_DEVICE_API_IMPL(cuda) -#else -INFINICCL_DEVICE_API_NOOP(cuda) -#endif - -#endif /* INFINICCL_CUDA_H_ */ diff --git a/src/infiniccl/infiniccl.cc b/src/infiniccl/infiniccl.cc deleted file mode 100644 index 48dbea425..000000000 --- a/src/infiniccl/infiniccl.cc +++ /dev/null @@ -1,296 +0,0 @@ -#include "infiniccl.h" - -#include "./ascend/infiniccl_ascend.h" -#include "./cambricon/infiniccl_cambricon.h" -#include "./cuda/infiniccl_cuda.h" -#include "./kunlun/infiniccl_kunlun.h" -#include "./metax/infiniccl_metax.h" -#include "./moore/infiniccl_moore.h" - -__INFINI_C infiniStatus_t infinicclCommInitAll( - infiniDevice_t device_type, - infinicclComm_t *comms, - int ndevice, - const int *device_ids) { - -#define COMM_INIT_ALL(CASE_, NAMESPACE_) \ - case CASE_: \ - return infiniccl::NAMESPACE_::commInitAll(comms, ndevice, device_ids) - - switch (device_type) { - COMM_INIT_ALL(INFINI_DEVICE_NVIDIA, cuda); - COMM_INIT_ALL(INFINI_DEVICE_ILUVATAR, cuda); - COMM_INIT_ALL(INFINI_DEVICE_QY, cuda); - COMM_INIT_ALL(INFINI_DEVICE_HYGON, cuda); - COMM_INIT_ALL(INFINI_DEVICE_ASCEND, ascend); - COMM_INIT_ALL(INFINI_DEVICE_CAMBRICON, cambricon); - COMM_INIT_ALL(INFINI_DEVICE_METAX, metax); - COMM_INIT_ALL(INFINI_DEVICE_MOORE, moore); - COMM_INIT_ALL(INFINI_DEVICE_KUNLUN, kunlun); - COMM_INIT_ALL(INFINI_DEVICE_ALI, cuda); - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef COMM_INIT_ALL -} - -__INFINI_C infiniStatus_t infinicclCommDestroy(infinicclComm_t comm) { - if (comm == nullptr) { - return INFINI_STATUS_SUCCESS; - } - -#define COMM_DESTROY(CASE_, NAMESPACE_) \ - case CASE_: \ - return infiniccl::NAMESPACE_::commDestroy(comm) - - switch (comm->device_type) { - COMM_DESTROY(INFINI_DEVICE_NVIDIA, cuda); - COMM_DESTROY(INFINI_DEVICE_ILUVATAR, cuda); - COMM_DESTROY(INFINI_DEVICE_QY, cuda); - COMM_DESTROY(INFINI_DEVICE_HYGON, cuda); - COMM_DESTROY(INFINI_DEVICE_ASCEND, ascend); - COMM_DESTROY(INFINI_DEVICE_CAMBRICON, cambricon); - COMM_DESTROY(INFINI_DEVICE_METAX, metax); - COMM_DESTROY(INFINI_DEVICE_MOORE, moore); - COMM_DESTROY(INFINI_DEVICE_KUNLUN, kunlun); - COMM_DESTROY(INFINI_DEVICE_ALI, cuda); - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef COMM_DESTROY -} - -__INFINI_C infiniStatus_t infinicclGroupStart(infinicclComm_t comm) { - if (comm == nullptr) { - return INFINI_STATUS_NULL_POINTER; - } - -#define GROUP_START(CASE_, NAMESPACE_) \ - case CASE_: \ - return infiniccl::NAMESPACE_::groupStart(comm) - - switch (comm->device_type) { - GROUP_START(INFINI_DEVICE_NVIDIA, cuda); - GROUP_START(INFINI_DEVICE_ILUVATAR, cuda); - GROUP_START(INFINI_DEVICE_QY, cuda); - GROUP_START(INFINI_DEVICE_HYGON, cuda); - GROUP_START(INFINI_DEVICE_ASCEND, ascend); - GROUP_START(INFINI_DEVICE_CAMBRICON, cambricon); - GROUP_START(INFINI_DEVICE_METAX, metax); - GROUP_START(INFINI_DEVICE_MOORE, moore); - GROUP_START(INFINI_DEVICE_KUNLUN, kunlun); - GROUP_START(INFINI_DEVICE_ALI, cuda); - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GROUP_START -} - -__INFINI_C infiniStatus_t infinicclGroupEnd(infinicclComm_t comm) { - if (comm == nullptr) { - return INFINI_STATUS_NULL_POINTER; - } - -#define GROUP_END(CASE_, NAMESPACE_) \ - case CASE_: \ - return infiniccl::NAMESPACE_::groupEnd(comm) - - switch (comm->device_type) { - GROUP_END(INFINI_DEVICE_NVIDIA, cuda); - GROUP_END(INFINI_DEVICE_ILUVATAR, cuda); - GROUP_END(INFINI_DEVICE_QY, cuda); - GROUP_END(INFINI_DEVICE_HYGON, cuda); - GROUP_END(INFINI_DEVICE_ASCEND, ascend); - GROUP_END(INFINI_DEVICE_CAMBRICON, cambricon); - GROUP_END(INFINI_DEVICE_METAX, metax); - GROUP_END(INFINI_DEVICE_MOORE, moore); - GROUP_END(INFINI_DEVICE_KUNLUN, kunlun); - GROUP_END(INFINI_DEVICE_ALI, cuda); - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GROUP_END -} - -__INFINI_C infiniStatus_t infinicclAllReduce( - void *sendbuf, - void *recvbuf, - size_t count, - infiniDtype_t dataype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream) { - - if (comm == nullptr) { - return INFINI_STATUS_NULL_POINTER; - } - -#define ALL_REDUCE(CASE_, NAMESPACE_) \ - case CASE_: \ - return infiniccl::NAMESPACE_::allReduce(sendbuf, recvbuf, count, dataype, op, comm, stream) - - switch (comm->device_type) { - ALL_REDUCE(INFINI_DEVICE_NVIDIA, cuda); - ALL_REDUCE(INFINI_DEVICE_ILUVATAR, cuda); - ALL_REDUCE(INFINI_DEVICE_QY, cuda); - ALL_REDUCE(INFINI_DEVICE_HYGON, cuda); - ALL_REDUCE(INFINI_DEVICE_ASCEND, ascend); - ALL_REDUCE(INFINI_DEVICE_CAMBRICON, cambricon); - ALL_REDUCE(INFINI_DEVICE_METAX, metax); - ALL_REDUCE(INFINI_DEVICE_MOORE, moore); - ALL_REDUCE(INFINI_DEVICE_KUNLUN, kunlun); - ALL_REDUCE(INFINI_DEVICE_ALI, cuda); - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef ALL_REDUCE -} - -__INFINI_C infiniStatus_t infinicclAllGather( - void *sendbuf, - void *recvbuf, - size_t send_count, - infiniDtype_t dataype, - infinicclComm_t comm, - infinirtStream_t stream) { - - if (comm == nullptr) { - return INFINI_STATUS_NULL_POINTER; - } - -#define ALL_GATHER(CASE_, NAMESPACE_) \ - case CASE_: \ - return infiniccl::NAMESPACE_::allGather(sendbuf, recvbuf, send_count, dataype, comm, stream) - - switch (comm->device_type) { - ALL_GATHER(INFINI_DEVICE_NVIDIA, cuda); - ALL_GATHER(INFINI_DEVICE_ILUVATAR, cuda); - ALL_GATHER(INFINI_DEVICE_QY, cuda); - ALL_GATHER(INFINI_DEVICE_HYGON, cuda); - ALL_GATHER(INFINI_DEVICE_ASCEND, ascend); - ALL_GATHER(INFINI_DEVICE_CAMBRICON, cambricon); - ALL_GATHER(INFINI_DEVICE_METAX, metax); - ALL_GATHER(INFINI_DEVICE_MOORE, moore); - ALL_GATHER(INFINI_DEVICE_KUNLUN, kunlun); - ALL_GATHER(INFINI_DEVICE_ALI, cuda); - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef ALL_GATHER -} - -__INFINI_C infiniStatus_t infinicclAllGatherV( - void *sendbuf, - void *recvbuf, - const size_t *recv_counts, - int nranks, - infiniDtype_t dataype, - infinicclComm_t comm, - infinirtStream_t stream) { - - if (comm == nullptr || recv_counts == nullptr) { - return INFINI_STATUS_NULL_POINTER; - } - -#define ALL_GATHER_V(CASE_, NAMESPACE_) \ - case CASE_: \ - return infiniccl::NAMESPACE_::allGatherV(sendbuf, recvbuf, recv_counts, nranks, dataype, comm, stream) - - switch (comm->device_type) { - ALL_GATHER_V(INFINI_DEVICE_NVIDIA, cuda); - ALL_GATHER_V(INFINI_DEVICE_ILUVATAR, cuda); - ALL_GATHER_V(INFINI_DEVICE_QY, cuda); - ALL_GATHER_V(INFINI_DEVICE_HYGON, cuda); - ALL_GATHER_V(INFINI_DEVICE_ASCEND, ascend); - ALL_GATHER_V(INFINI_DEVICE_CAMBRICON, cambricon); - ALL_GATHER_V(INFINI_DEVICE_METAX, metax); - ALL_GATHER_V(INFINI_DEVICE_MOORE, moore); - ALL_GATHER_V(INFINI_DEVICE_KUNLUN, kunlun); - ALL_GATHER_V(INFINI_DEVICE_ALI, cuda); - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef ALL_GATHER_V -} - -__INFINI_C infiniStatus_t infinicclReduceScatter( - void *sendbuf, - void *recvbuf, - size_t recv_count, - infiniDtype_t dataype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream) { - - if (comm == nullptr) { - return INFINI_STATUS_NULL_POINTER; - } - -#define REDUCE_SCATTER(CASE_, NAMESPACE_) \ - case CASE_: \ - return infiniccl::NAMESPACE_::reduceScatter(sendbuf, recvbuf, recv_count, dataype, op, comm, stream) - - switch (comm->device_type) { - REDUCE_SCATTER(INFINI_DEVICE_NVIDIA, cuda); - REDUCE_SCATTER(INFINI_DEVICE_ILUVATAR, cuda); - REDUCE_SCATTER(INFINI_DEVICE_QY, cuda); - REDUCE_SCATTER(INFINI_DEVICE_HYGON, cuda); - REDUCE_SCATTER(INFINI_DEVICE_ASCEND, ascend); - REDUCE_SCATTER(INFINI_DEVICE_CAMBRICON, cambricon); - REDUCE_SCATTER(INFINI_DEVICE_METAX, metax); - REDUCE_SCATTER(INFINI_DEVICE_MOORE, moore); - REDUCE_SCATTER(INFINI_DEVICE_KUNLUN, kunlun); - REDUCE_SCATTER(INFINI_DEVICE_ALI, cuda); - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef REDUCE_SCATTER -} - -__INFINI_C infiniStatus_t infinicclReduceScatterV( - void *sendbuf, - void *recvbuf, - const size_t *send_counts, - int nranks, - infiniDtype_t dataype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream) { - - if (comm == nullptr || send_counts == nullptr) { - return INFINI_STATUS_NULL_POINTER; - } - -#define REDUCE_SCATTER_V(CASE_, NAMESPACE_) \ - case CASE_: \ - return infiniccl::NAMESPACE_::reduceScatterV(sendbuf, recvbuf, send_counts, nranks, dataype, op, comm, stream) - - switch (comm->device_type) { - REDUCE_SCATTER_V(INFINI_DEVICE_NVIDIA, cuda); - REDUCE_SCATTER_V(INFINI_DEVICE_ILUVATAR, cuda); - REDUCE_SCATTER_V(INFINI_DEVICE_QY, cuda); - REDUCE_SCATTER_V(INFINI_DEVICE_HYGON, cuda); - REDUCE_SCATTER_V(INFINI_DEVICE_ASCEND, ascend); - REDUCE_SCATTER_V(INFINI_DEVICE_CAMBRICON, cambricon); - REDUCE_SCATTER_V(INFINI_DEVICE_METAX, metax); - REDUCE_SCATTER_V(INFINI_DEVICE_MOORE, moore); - REDUCE_SCATTER_V(INFINI_DEVICE_KUNLUN, kunlun); - REDUCE_SCATTER_V(INFINI_DEVICE_ALI, cuda); - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef REDUCE_SCATTER_V -} diff --git a/src/infiniccl/infiniccl_impl.h b/src/infiniccl/infiniccl_impl.h deleted file mode 100644 index 790355bd9..000000000 --- a/src/infiniccl/infiniccl_impl.h +++ /dev/null @@ -1,132 +0,0 @@ -#ifndef INFINICCL_IMPL_H -#define INFINICCL_IMPL_H - -#include "infiniccl.h" - -struct InfinicclComm { - infiniDevice_t device_type; - int device_id; // the actual device ID, not rank number - void *comm; // the actual communicator - int rank = 0; - int world_size = 1; -}; - -#define INFINICCL_DEVICE_API(NAMSPACE, IMPL) \ - namespace infiniccl::NAMSPACE { \ - infiniStatus_t commInitAll( \ - infinicclComm_t *comms, \ - int ndevice, \ - const int *device_ids) IMPL; \ - \ - infiniStatus_t commDestroy(infinicclComm_t comm) IMPL; \ - \ - infiniStatus_t groupStart(infinicclComm_t comm) IMPL; \ - \ - infiniStatus_t groupEnd(infinicclComm_t comm) IMPL; \ - \ - infiniStatus_t allReduce( \ - void *sendbuf, \ - void *recvbuf, \ - size_t count, \ - infiniDtype_t datatype, \ - infinicclReduceOp_t op, \ - infinicclComm_t comm, \ - infinirtStream_t stream) IMPL; \ - \ - infiniStatus_t allGather( \ - void *sendbuf, \ - void *recvbuf, \ - size_t send_count, \ - infiniDtype_t datatype, \ - infinicclComm_t comm, \ - infinirtStream_t stream) IMPL; \ - \ - infiniStatus_t allGatherV( \ - void *sendbuf, \ - void *recvbuf, \ - const size_t *recv_counts, \ - int nranks, \ - infiniDtype_t datatype, \ - infinicclComm_t comm, \ - infinirtStream_t stream) IMPL; \ - \ - infiniStatus_t reduceScatter( \ - void *sendbuf, \ - void *recvbuf, \ - size_t recv_count, \ - infiniDtype_t datatype, \ - infinicclReduceOp_t op, \ - infinicclComm_t comm, \ - infinirtStream_t stream) IMPL; \ - \ - infiniStatus_t reduceScatterV( \ - void *sendbuf, \ - void *recvbuf, \ - const size_t *send_counts, \ - int nranks, \ - infiniDtype_t datatype, \ - infinicclReduceOp_t op, \ - infinicclComm_t comm, \ - infinirtStream_t stream) IMPL; \ - }; - -#define INFINICCL_DEVICE_API_IMPL(NAMSPACE) \ - INFINICCL_DEVICE_API(NAMSPACE, ) - -#define INFINICCL_DEVICE_API_NOOP(NAMSPACE) \ - INFINICCL_DEVICE_API(NAMSPACE, { return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; }) - -#define INFINICCL_DEVICE_API_UNSUPPORTED_EP_COLLECTIVES() \ - infiniStatus_t groupStart(infinicclComm_t) { \ - return INFINI_STATUS_NOT_IMPLEMENTED; \ - } \ - \ - infiniStatus_t groupEnd(infinicclComm_t) { \ - return INFINI_STATUS_NOT_IMPLEMENTED; \ - } \ - \ - infiniStatus_t allGather( \ - void *, \ - void *, \ - size_t, \ - infiniDtype_t, \ - infinicclComm_t, \ - infinirtStream_t) { \ - return INFINI_STATUS_NOT_IMPLEMENTED; \ - } \ - \ - infiniStatus_t allGatherV( \ - void *, \ - void *, \ - const size_t *, \ - int, \ - infiniDtype_t, \ - infinicclComm_t, \ - infinirtStream_t) { \ - return INFINI_STATUS_NOT_IMPLEMENTED; \ - } \ - \ - infiniStatus_t reduceScatter( \ - void *, \ - void *, \ - size_t, \ - infiniDtype_t, \ - infinicclReduceOp_t, \ - infinicclComm_t, \ - infinirtStream_t) { \ - return INFINI_STATUS_NOT_IMPLEMENTED; \ - } \ - \ - infiniStatus_t reduceScatterV( \ - void *, \ - void *, \ - const size_t *, \ - int, \ - infiniDtype_t, \ - infinicclReduceOp_t, \ - infinicclComm_t, \ - infinirtStream_t) { \ - return INFINI_STATUS_NOT_IMPLEMENTED; \ - } - -#endif // INFINICCL_IMPL_H diff --git a/src/infiniccl/kunlun/infiniccl_kunlun.cc b/src/infiniccl/kunlun/infiniccl_kunlun.cc deleted file mode 100644 index 0fa19c4aa..000000000 --- a/src/infiniccl/kunlun/infiniccl_kunlun.cc +++ /dev/null @@ -1,101 +0,0 @@ -#include "infiniccl_kunlun.h" - -#include "../../utils.h" - -#include - -#include -#include - -#define CHECK_BKCL(API__) CHECK_INTERNAL(API__, BKCL_SUCCESS) - -typedef XPUStream kunlunStream_t; -typedef BKCLContext_t bkclComm_t; - -inline kunlunStream_t getKunlunStream(infinirtStream_t stream) { - if (stream == nullptr) { - return 0; - } - return reinterpret_cast(stream); -} - -inline bkclComm_t getBkclComm(infinicclComm_t comm) { - return reinterpret_cast(comm->comm); -} - -inline BKCLDataType getBkclDtype(infiniDtype_t datatype) { - switch (datatype) { - case INFINI_DTYPE_F32: - return BKCL_FLOAT; - case INFINI_DTYPE_F16: - return BKCL_FLOAT16; - case INFINI_DTYPE_BF16: - return BKCL_BFLOAT16; - default: - std::cerr << "Unsupported data type: " << datatype << std::endl; - std::abort(); - return BKCL_FLOAT16; - } -} - -inline BKCLOp getBkclRedOp(infinicclReduceOp_t op) { - switch (op) { - case INFINICCL_SUM: - return BKCL_ADD; - case INFINICCL_PROD: - return BKCL_PRODUCT; - case INFINICCL_MAX: - return BKCL_MAX; - case INFINICCL_MIN: - return BKCL_MIN; - default: - std::abort(); - return BKCL_ADD; - } -} - -namespace infiniccl::kunlun { - -infiniStatus_t commInitAll( - infinicclComm_t *comms, - int ndevice, - const int *device_ids) { - std::vector bkcl_comms(ndevice); - CHECK_BKCL(bkcl_comm_init_all(bkcl_comms.data(), ndevice, device_ids)); - - for (int i = 0; i < ndevice; i++) { - comms[i] = new InfinicclComm{INFINI_DEVICE_KUNLUN, device_ids[i], (void *)(bkcl_comms[i])}; - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t commDestroy(infinicclComm_t comm) { - CHECK_BKCL(bkcl_destroy_context(getBkclComm(comm))); - delete comm; - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t allReduce( - void *sendbuf, - void *recvbuf, - size_t count, - infiniDtype_t datatype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream) { - CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16); - CHECK_BKCL(bkcl_all_reduce( - getBkclComm(comm), - sendbuf, - recvbuf, - count, - getBkclDtype(datatype), - getBkclRedOp(op), - getKunlunStream(stream))); - - return INFINI_STATUS_SUCCESS; -} - -INFINICCL_DEVICE_API_UNSUPPORTED_EP_COLLECTIVES(); -} // namespace infiniccl::kunlun diff --git a/src/infiniccl/kunlun/infiniccl_kunlun.h b/src/infiniccl/kunlun/infiniccl_kunlun.h deleted file mode 100644 index 1855c8c5f..000000000 --- a/src/infiniccl/kunlun/infiniccl_kunlun.h +++ /dev/null @@ -1,12 +0,0 @@ -#ifndef INFINICCL_KUNLUN_H_ -#define INFINICCL_KUNLUN_H_ - -#include "../infiniccl_impl.h" - -#if defined(ENABLE_KUNLUN_API) && defined(ENABLE_CCL) -INFINICCL_DEVICE_API_IMPL(kunlun) -#else -INFINICCL_DEVICE_API_NOOP(kunlun) -#endif - -#endif /* INFINICCL_KUNLUN_H_ */ diff --git a/src/infiniccl/metax/infiniccl_metax.cc b/src/infiniccl/metax/infiniccl_metax.cc deleted file mode 100644 index 4b9f6303f..000000000 --- a/src/infiniccl/metax/infiniccl_metax.cc +++ /dev/null @@ -1,101 +0,0 @@ -#include "infiniccl_metax.h" - -#include "../../utils.h" - -#ifdef ENABLE_METAX_MC_API -#include -#include -#else -#include -#include -#endif - -#include -#include - -#define CHECK_HCCL(API__) CHECK_INTERNAL(API__, hcclSuccess) - -inline hcStream_t getMacaStream(infinirtStream_t stream) { - if (stream == nullptr) { - return 0; - } - return static_cast(stream); -} - -inline hcclDataType_t getHcclDtype(infiniDtype_t datatype) { - switch (datatype) { - case INFINI_DTYPE_F32: - return hcclFloat; - case INFINI_DTYPE_F16: - return hcclHalf; - case INFINI_DTYPE_BF16: - return hcclBfloat16; - default: - std::abort(); - return hcclHalf; - } -} - -inline hcclRedOp_t getHcclRedOp(infinicclReduceOp_t op) { - switch (op) { - case INFINICCL_SUM: - return hcclSum; - case INFINICCL_PROD: - return hcclProd; - case INFINICCL_MAX: - return hcclMax; - case INFINICCL_MIN: - return hcclMin; - case INFINICCL_AVG: - return hcclAvg; - default: - std::abort(); - return hcclSum; - } -} - -inline hcclComm_t getHcclComm(infinicclComm_t comm) { - return static_cast(comm->comm); -} - -namespace infiniccl::metax { - -infiniStatus_t commInitAll( - infinicclComm_t *comms, - int ndevice, - const int *device_ids) { - - std::vector hccl_comms(ndevice); - CHECK_HCCL(hcclCommInitAll(hccl_comms.data(), ndevice, (int const *)device_ids)); - - for (int i = 0; i < ndevice; i++) { - comms[i] = new InfinicclComm{INFINI_DEVICE_METAX, device_ids[i], (void *)(hccl_comms[i])}; - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t commDestroy(infinicclComm_t comm) { - CHECK_HCCL(hcclCommDestroy(getHcclComm(comm))); - delete comm; - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t allReduce( - void *sendbuf, - void *recvbuf, - size_t count, - infiniDtype_t datatype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream) { - - CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16); - - CHECK_HCCL(hcclAllReduce(sendbuf, recvbuf, count, getHcclDtype(datatype), - getHcclRedOp(op), getHcclComm(comm), getMacaStream(stream))); - - return INFINI_STATUS_SUCCESS; -} -INFINICCL_DEVICE_API_UNSUPPORTED_EP_COLLECTIVES(); -} // namespace infiniccl::metax diff --git a/src/infiniccl/metax/infiniccl_metax.h b/src/infiniccl/metax/infiniccl_metax.h deleted file mode 100644 index 7ea84afda..000000000 --- a/src/infiniccl/metax/infiniccl_metax.h +++ /dev/null @@ -1,13 +0,0 @@ -#ifndef INFINICCL_METAX_H_ -#define INFINICCL_METAX_H_ - -#include "../../infiniop/devices/metax/metax_ht2mc.h" -#include "../infiniccl_impl.h" - -#if defined(ENABLE_METAX_API) && defined(ENABLE_CCL) -INFINICCL_DEVICE_API_IMPL(metax) -#else -INFINICCL_DEVICE_API_NOOP(metax) -#endif - -#endif /* INFINICCL_METAX_H_ */ diff --git a/src/infiniccl/moore/infiniccl_moore.cc b/src/infiniccl/moore/infiniccl_moore.cc deleted file mode 100644 index 7c2faed65..000000000 --- a/src/infiniccl/moore/infiniccl_moore.cc +++ /dev/null @@ -1,109 +0,0 @@ -#include "infiniccl_moore.h" - -#include "../../utils.h" - -#include -#include - -#include -#include - -#define CHECK_MCCL(API__) CHECK_INTERNAL(API__, mcclSuccess) - -inline musaStream_t getMusaStream(infinirtStream_t stream) { - if (stream == nullptr) { - return 0; - } - return static_cast(stream); -} - -inline mcclDataType_t getMcclDtype(infiniDtype_t datatype) { - switch (datatype) { - case INFINI_DTYPE_F32: - return mcclFloat; - case INFINI_DTYPE_F16: - return mcclHalf; - -#if MARCH_TYPE == 310 - case INFINI_DTYPE_BF16: - return mcclBfloat16; -#endif - - default: - std::abort(); - return mcclHalf; - } -} - -inline mcclRedOp_t getMcclRedOp(infinicclReduceOp_t op) { - switch (op) { - case INFINICCL_SUM: - return mcclSum; - case INFINICCL_PROD: - return mcclProd; - case INFINICCL_MAX: - return mcclMax; - case INFINICCL_MIN: - return mcclMin; - case INFINICCL_AVG: - return mcclAvg; - default: - std::abort(); - return mcclSum; - } -} - -inline mcclComm_t getMcclComm(infinicclComm_t comm) { - return static_cast(comm->comm); -} - -namespace infiniccl::moore { - -infiniStatus_t commInitAll( - infinicclComm_t *comms, - int ndevice, - const int *device_ids) { - - std::vector mccl_comms(ndevice); - CHECK_MCCL(mcclCommInitAll(mccl_comms.data(), ndevice, (int const *)device_ids)); - - for (int i = 0; i < ndevice; i++) { - comms[i] = new InfinicclComm{INFINI_DEVICE_MOORE, device_ids[i], (void *)(mccl_comms[i])}; - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t commDestroy(infinicclComm_t comm) { - CHECK_MCCL(mcclCommDestroy(getMcclComm(comm))); - delete comm; - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t allReduce( - void *sendbuf, - void *recvbuf, - size_t count, - infiniDtype_t datatype, - infinicclReduceOp_t op, - infinicclComm_t comm, - infinirtStream_t stream) { - -#if MARCH_TYPE == 310 - CHECK_DTYPE(datatype, - INFINI_DTYPE_F32, - INFINI_DTYPE_F16, - INFINI_DTYPE_BF16); -#else - CHECK_DTYPE(datatype, - INFINI_DTYPE_F32, - INFINI_DTYPE_F16); -#endif - - CHECK_MCCL(mcclAllReduce(sendbuf, recvbuf, count, getMcclDtype(datatype), - getMcclRedOp(op), getMcclComm(comm), getMusaStream(stream))); - - return INFINI_STATUS_SUCCESS; -} -INFINICCL_DEVICE_API_UNSUPPORTED_EP_COLLECTIVES(); -} // namespace infiniccl::moore diff --git a/src/infiniccl/moore/infiniccl_moore.h b/src/infiniccl/moore/infiniccl_moore.h deleted file mode 100644 index 318fc468b..000000000 --- a/src/infiniccl/moore/infiniccl_moore.h +++ /dev/null @@ -1,12 +0,0 @@ -#ifndef INFINICCL_MOORE_H_ -#define INFINICCL_MOORE_H_ - -#include "../infiniccl_impl.h" - -#if defined(ENABLE_MOORE_API) && defined(ENABLE_CCL) -INFINICCL_DEVICE_API_IMPL(moore) -#else -INFINICCL_DEVICE_API_NOOP(moore) -#endif - -#endif /* INFINICCL_MOORE_H_ */ diff --git a/src/infinicore-test/README.md b/src/infinicore-test/README.md deleted file mode 100644 index a6349349d..000000000 --- a/src/infinicore-test/README.md +++ /dev/null @@ -1,279 +0,0 @@ -# InfiniCore Memory Management Test Suite - -This test suite provides comprehensive testing for the InfiniCore memory management system, focusing on the critical issues identified in the memory management architecture analysis. - -## Overview - -The test suite includes six main test categories: - -1. **Basic Memory Tests** - Basic allocation, deallocation, and memory operations -2. **Concurrency Tests** - Thread safety and concurrent access testing -3. **Exception Safety Tests** - Exception handling and safety testing -4. **Memory Leak Tests** - Memory leak detection and prevention -5. **Performance Tests** - Performance benchmarks and optimization validation -6. **Stress Tests** - High-load stress testing and edge cases - -## Building - -### Using XMake (if integrated with main build) -```bash -# From InfiniCore root directory -xmake build infinicore-test -``` - -## Running Tests - -### Run All Tests -```bash -./infinicore-test -``` - -### Run Specific Test Categories -```bash -# Basic memory tests -./infinicore-test --test basic - -# Concurrency tests -./infinicore-test --test concurrency - -# Exception safety tests -./infinicore-test --test exception - -# Memory leak tests -./infinicore-test --test leak - -# Performance tests -./infinicore-test --test performance - -# Stress tests -./infinicore-test --test stress -``` - -### Run with Specific Device -```bash -# Run on CPU -./infinicore-test --cpu - -# Run on NVIDIA GPU -./infinicore-test --nvidia - -# Run on other devices -./infinicore-test --cambricon -./infinicore-test --ascend -./infinicore-test --metax -./infinicore-test --moore -./infinicore-test --iluvatar -./infinicore-test --qy -./infinicore-test --kunlun -./infinicore-test --hygon -./infinicore-test --ali -``` - -### Customize Test Parameters -```bash -# Run with custom thread count -./infinicore-test --threads 8 - -# Run with custom iteration count -./infinicore-test --iterations 5000 - -# Combine options -./infinicore-test --nvidia --test concurrency --threads 16 --iterations 2000 -``` - -## Test Categories - -### 1. Basic Memory Tests -Tests fundamental memory operations: -- Memory allocation and deallocation -- Memory size and device properties -- Memory read/write operations -- Pinned memory allocation -- Memory data integrity - -### 2. Concurrency Tests -Tests thread safety and concurrent access: -- **Concurrent Allocations**: Multiple threads allocating memory simultaneously -- **Concurrent Device Switching**: Multiple threads switching device contexts -- **Memory Allocation Race**: Race condition testing for memory operations - -### 3. Exception Safety Tests -Tests exception handling and safety: -- **Allocation Failure**: Tests behavior when allocation fails -- **Deallocation Exception**: Tests exception safety during deallocation -- **Context Switch Exception**: Tests exception handling during device switching - -### 4. Memory Leak Tests -Tests memory leak detection and prevention: -- **Basic Leak Detection**: Basic memory leak detection -- **Cross-Device Leak Detection**: Memory leaks in cross-device scenarios -- **Exception Leak Detection**: Memory leaks during exception handling - -### 5. Performance Tests -Tests performance and benchmarks: -- **Allocation Performance**: Memory allocation speed benchmarks -- **Concurrent Performance**: Performance under concurrent load -- **Memory Copy Performance**: Memory copy bandwidth tests - -### 6. Stress Tests -Tests high-load scenarios and edge cases: -- **High Frequency Allocations**: Rapid allocation/deallocation cycles -- **Large Memory Allocations**: Large memory block allocation -- **Cross-Device Stress**: Stress testing across multiple devices - -## Expected Results - -### Critical Issues to Watch For - -The tests are designed to detect the critical issues identified in the memory management analysis: - -1. **Thread Safety Violations** - - Race conditions in concurrent allocations - - Inconsistent device context switching - - Global state corruption - -2. **Memory Leaks** - - Unfreed memory after deallocation - - Cross-device memory not properly cleaned up - - Exception-related memory leaks - -3. **Exception Safety Issues** - - Exceptions during allocation causing resource leaks - - Exceptions in destructors causing `std::terminate` - - Incomplete cleanup on exceptions - -4. **Performance Issues** - - Slow allocation/deallocation performance - - Poor concurrent performance - - Inefficient memory copy operations - -### Performance Thresholds - -The tests include performance thresholds: - -- **Allocation Performance**: < 100μs per allocation -- **Concurrent Performance**: < 200μs per allocation under load -- **Memory Bandwidth**: > 100 MB/s for memory copies - -## Test Output - -### Successful Test Run -``` -============================================== -InfiniCore Memory Management Test Suite -============================================== -Device: 0 -Threads: 4 -Iterations: 1000 -============================================== - -[SUITE] Running: BasicMemoryTest -[TEST] Starting: BasicMemoryTest -[TEST] PASSED: BasicMemoryTest (Duration: 1234μs) - -[SUITE] Running: ConcurrencyTest -[TEST] Starting: ConcurrencyTest -[TEST] PASSED: ConcurrencyTest (Duration: 5678μs) - -... - -============================================== -Test Summary -============================================== -Total Tests: 6 -Passed: 6 -Failed: 0 -Total Time: 12345μs -============================================== - -✅ All tests passed! -``` - -### Failed Test Run -``` -[TEST] FAILED: ConcurrencyTest - Concurrent allocation test failed: expected 8000 successes, got 7995 successes and 5 failures - -============================================== -Final Results -============================================== -Total Tests: 6 -Passed: 5 -Failed: 1 -============================================== - -❌ Some tests failed. Please review the output above. -``` - -## Debugging Failed Tests - -### Common Issues and Solutions - -1. **Thread Safety Failures** - - Check for race conditions in global state access - - Verify proper synchronization in allocators - - Review device context switching logic - -2. **Memory Leak Failures** - - Check deallocation logic in allocators - - Verify cross-device cleanup mechanisms - - Review exception safety in destructors - -3. **Performance Failures** - - Profile allocation/deallocation paths - - Check for unnecessary context switching - - Review memory copy implementations - -4. **Exception Safety Failures** - - Verify no-throw guarantees in destructors - - Check exception handling in allocation paths - - Review resource cleanup on exceptions - -## Integration with CI/CD - -### GitHub Actions Example -```yaml -- name: Run Memory Tests - run: | - cd src/infinicore-test - mkdir build && cd build - cmake .. - make - ./infinicore-test --test all -``` - -### Custom Test Targets -```bash -# Run specific test categories -make test-memory-basic -make test-memory-concurrency -make test-memory-exception -make test-memory-leak -make test-memory-performance -make test-memory-stress -make test-memory-all -``` - -## Contributing - -When adding new tests: - -1. Follow the existing test framework pattern -2. Add appropriate error messages and logging -3. Include performance thresholds where applicable -4. Test both success and failure scenarios -5. Update this README with new test descriptions - -## Dependencies - -- InfiniCore library (infinicore, infiniop, infinirt, infiniccl) -- C++17 compatible compiler -- Threading library (pthread on Linux) -- CMake 3.16+ (for CMake build) - -## Notes - -- Tests are designed to be deterministic where possible -- Some tests may have timing dependencies -- Performance tests may vary based on system load -- Memory leak detection is basic and may not catch all leaks -- Tests assume proper InfiniCore initialization diff --git a/src/infinicore-test/main.cc b/src/infinicore-test/main.cc deleted file mode 100644 index c2c32170d..000000000 --- a/src/infinicore-test/main.cc +++ /dev/null @@ -1,278 +0,0 @@ -#include "memory_test.h" -#include "test_mutual_awareness.h" -#include "test_nn_module.h" -#include "test_runner.h" -#include "test_tensor_destructor.h" -#include -#include -#include -#include - -struct ParsedArgs { - infiniDevice_t device_type = INFINI_DEVICE_CPU; - bool run_basic = true; - bool run_concurrency = true; - bool run_exception_safety = true; - bool run_memory_leak = true; - bool run_performance = true; - bool run_stress = true; - bool run_module = false; -#ifdef ENABLE_MUTUAL_AWARENESS - bool run_mutual_awareness = true; -#endif - int num_threads = 4; - int iterations = 1000; -}; - -void printUsage() { - std::cout << "Usage:" << std::endl - << " infinicore-test [--] [--test ] [--threads ] [--iterations ]" << std::endl - << std::endl - << "Options:" << std::endl - << " -- Specify the device type (default: cpu)" << std::endl - << " --test Run specific test (basic|concurrency|exception|leak|performance|stress|module|all)" << std::endl - << " --threads Number of threads for concurrency tests (default: 4)" << std::endl - << " --iterations Number of iterations for stress tests (default: 1000)" << std::endl - << " --help Show this help message" << std::endl - << std::endl - << "Available devices:" << std::endl - << " cpu - Default" << std::endl - << " nvidia" << std::endl - << " cambricon" << std::endl - << " ascend" << std::endl - << " metax" << std::endl - << " moore" << std::endl - << " iluvatar" << std::endl - << " qy" << std::endl - << " kunlun" << std::endl - << " hygon" << std::endl - << " ali" << std::endl - << std::endl - << "Available tests:" << std::endl - << " basic - Basic memory allocation and deallocation tests" << std::endl - << " concurrency - Thread safety and concurrent access tests" << std::endl - << " exception - Exception safety tests" << std::endl - << " leak - Memory leak detection tests" << std::endl - << " performance - Performance and benchmark tests" << std::endl - << " stress - Stress tests with high load" << std::endl - << " module - Neural network module tests" << std::endl - << " all - Run all tests (default)" << std::endl - << std::endl; - exit(EXIT_SUCCESS); -} - -ParsedArgs parseArgs(int argc, char *argv[]) { - ParsedArgs args; - - for (int i = 1; i < argc; ++i) { - std::string arg = argv[i]; - - if (arg == "--help" || arg == "-h") { - printUsage(); - } else if (arg == "--cpu") { - args.device_type = INFINI_DEVICE_CPU; - } else if (arg == "--nvidia") { - args.device_type = INFINI_DEVICE_NVIDIA; - } else if (arg == "--cambricon") { - args.device_type = INFINI_DEVICE_CAMBRICON; - } else if (arg == "--ascend") { - args.device_type = INFINI_DEVICE_ASCEND; - } else if (arg == "--metax") { - args.device_type = INFINI_DEVICE_METAX; - } else if (arg == "--moore") { - args.device_type = INFINI_DEVICE_MOORE; - } else if (arg == "--iluvatar") { - args.device_type = INFINI_DEVICE_ILUVATAR; - } else if (arg == "--qy") { - args.device_type = INFINI_DEVICE_QY; - } else if (arg == "--kunlun") { - args.device_type = INFINI_DEVICE_KUNLUN; - } else if (arg == "--hygon") { - args.device_type = INFINI_DEVICE_HYGON; - } else if (arg == "--ali") { - args.device_type = INFINI_DEVICE_ALI; - } else if (arg == "--test") { - if (i + 1 >= argc) { - std::cerr << "Error: --test requires a test name" << std::endl; - exit(EXIT_FAILURE); - } - - std::string test_name = argv[++i]; - args.run_basic = args.run_concurrency = args.run_exception_safety = args.run_memory_leak = args.run_performance = args.run_stress = args.run_module = false; -#ifdef ENABLE_MUTUAL_AWARENESS - args.run_mutual_awareness = false; -#endif - - if (test_name == "basic") { - args.run_basic = true; - } else if (test_name == "concurrency") { - args.run_concurrency = true; - } else if (test_name == "exception") { - args.run_exception_safety = true; - } else if (test_name == "leak") { - args.run_memory_leak = true; - } else if (test_name == "performance") { - args.run_performance = true; - } else if (test_name == "stress") { - args.run_stress = true; - } else if (test_name == "module") { - args.run_module = true; -#ifdef ENABLE_MUTUAL_AWARENESS - } else if (test_name == "mutual") { - args.run_mutual_awareness = true; -#endif - } else if (test_name == "all") { - args.run_basic = args.run_concurrency = args.run_exception_safety = args.run_memory_leak = args.run_performance = args.run_stress = args.run_module = true; -#ifdef ENABLE_MUTUAL_AWARENESS - args.run_mutual_awareness = true; -#endif - } else { - std::cerr << "Error: Unknown test name: " << test_name << std::endl; - exit(EXIT_FAILURE); - } - } else if (arg == "--threads") { - if (i + 1 >= argc) { - std::cerr << "Error: --threads requires a number" << std::endl; - exit(EXIT_FAILURE); - } - args.num_threads = std::stoi(argv[++i]); - if (args.num_threads <= 0) { - std::cerr << "Error: Number of threads must be positive" << std::endl; - exit(EXIT_FAILURE); - } - } else if (arg == "--iterations") { - if (i + 1 >= argc) { - std::cerr << "Error: --iterations requires a number" << std::endl; - exit(EXIT_FAILURE); - } - args.iterations = std::stoi(argv[++i]); - if (args.iterations <= 0) { - std::cerr << "Error: Number of iterations must be positive" << std::endl; - exit(EXIT_FAILURE); - } - } else { - std::cerr << "Error: Unknown argument: " << arg << std::endl; - exit(EXIT_FAILURE); - } - } - - return args; -} - -int main(int argc, char *argv[]) { - try { - ParsedArgs args = parseArgs(argc, argv); - spdlog::info("Arguments parsed successfully"); - - std::cout << "==============================================\n" - << "InfiniCore Memory Management Test Suite\n" - << "==============================================\n" - << "Device: " << static_cast(args.device_type) << "\n" - << "Threads: " << args.num_threads << "\n" - << "Iterations: " << args.iterations << "\n" - << "==============================================" << std::endl; - - spdlog::info("About to initialize InfiniCore context"); - // Initialize InfiniCore context - infinicore::context::setDevice(infinicore::Device(static_cast(args.device_type), 0)); - spdlog::info("InfiniCore context initialized successfully"); - - spdlog::info("Creating test runner"); - // Create test runner - infinicore::test::InfiniCoreTestRunner runner; - spdlog::info("Test runner created successfully"); - - // Add tests based on arguments - if (args.run_basic) { - runner.addTest(std::make_unique()); - - runner.addTest(std::make_unique()); - } - - if (args.run_module) { - runner.addTest(std::make_unique()); - } - - if (args.run_concurrency) { - runner.addTest(std::make_unique()); - } - - if (args.run_exception_safety) { - // runner.addTest(std::make_unique()); - } - - if (args.run_memory_leak) { - runner.addTest(std::make_unique()); - } - - if (args.run_performance) { - runner.addTest(std::make_unique()); - } - - if (args.run_stress) { - runner.addTest(std::make_unique()); - } - -#ifdef ENABLE_MUTUAL_AWARENESS - if (args.run_mutual_awareness) { - runner.addTest(std::make_unique()); - } -#endif - - spdlog::info("About to run all tests"); - // Run all tests - auto results = runner.runAllTests(); - spdlog::info("All tests completed"); - - // Count results and collect failed tests - size_t passed = 0, failed = 0; - std::vector failed_tests; - for (const auto &result : results) { - if (result.passed) { - passed++; - } else { - failed++; - failed_tests.push_back(result); - } - } - - // Print list of failed tests if any - if (!failed_tests.empty()) { - std::cout << "\n==============================================\n" - << "❌ FAILED TESTS\n" - << "==============================================" << std::endl; - for (const auto &test : failed_tests) { - std::cout << " • " << test.test_name; - if (!test.error_message.empty()) { - std::cout << "\n Error: " << test.error_message; - } - std::cout << "\n Duration: " << test.duration.count() << "μs" << std::endl; - } - } - - // Print final summary - std::cout << "\n==============================================\n" - << "Final Results\n" - << "==============================================\n" - << "Total Tests: " << results.size() << "\n" - << "Passed: " << passed << "\n" - << "Failed: " << failed << "\n" - << "==============================================" << std::endl; - - // Exit with appropriate code - if (failed > 0) { - std::cout << "\n❌ Some tests failed. Please review the failed tests list above." << std::endl; - return EXIT_FAILURE; - } else { - std::cout << "\n✅ All tests passed!" << std::endl; - return EXIT_SUCCESS; - } - - } catch (const std::exception &e) { - std::cerr << "Fatal error: " << e.what() << std::endl; - return EXIT_FAILURE; - } catch (...) { - std::cerr << "Fatal error: Unknown exception" << std::endl; - return EXIT_FAILURE; - } -} diff --git a/src/infinicore-test/memory_test.cc b/src/infinicore-test/memory_test.cc deleted file mode 100644 index 2029e0a81..000000000 --- a/src/infinicore-test/memory_test.cc +++ /dev/null @@ -1,900 +0,0 @@ -#include "memory_test.h" -#include -#include -#include - -namespace infinicore::test { - -// Basic Memory Test Implementation -TestResult BasicMemoryTest::run() { - return measureTime("BasicMemoryTest", [this]() -> bool { - try { - spdlog::debug("BasicMemoryTest: Starting test"); - // Test basic memory allocation - spdlog::debug("BasicMemoryTest: About to allocate memory"); - auto memory = context::allocateMemory(1024); - spdlog::debug("BasicMemoryTest: Memory allocated successfully"); - if (!memory) { - std::cerr << "Failed to allocate memory" << std::endl; - return false; - } - - spdlog::debug("BasicMemoryTest: Testing memory properties"); - // Test memory properties - if (memory->size() != 1024) { - std::cerr << "Memory size mismatch: expected 1024, got " << memory->size() << std::endl; - return false; - } - spdlog::debug("BasicMemoryTest: Memory size check passed"); - - spdlog::debug("BasicMemoryTest: Testing memory access"); - // Test memory access - std::byte *data = memory->data(); - spdlog::debug("BasicMemoryTest: Got memory data pointer: {}", static_cast(data)); - if (!data) { - std::cerr << "Memory data pointer is null" << std::endl; - return false; - } - spdlog::debug("BasicMemoryTest: Memory data pointer is valid"); - - // Check if this is GPU memory that can't be accessed directly - Device current_device = context::getDevice(); - spdlog::debug("BasicMemoryTest: Current device type: {}", static_cast(current_device.getType())); - spdlog::debug("BasicMemoryTest: Memory is pinned: {}", memory->is_pinned()); - - // For GPU memory, we shouldn't try to access it directly with memset - if (current_device.getType() != Device::Type::CPU) { - spdlog::debug("BasicMemoryTest: Skipping direct memory access for GPU device"); - spdlog::debug("BasicMemoryTest: GPU memory access test completed (skipped)"); - } else { - spdlog::debug("BasicMemoryTest: Testing memory write/read"); - // Test memory write/read - std::memset(data, 0xAB, 1024); - spdlog::debug("BasicMemoryTest: Memory memset completed"); - for (size_t i = 0; i < 1024; ++i) { - if (data[i] != static_cast(0xAB)) { - std::cerr << "Memory write/read test failed at index " << i << std::endl; - return false; - } - } - spdlog::debug("BasicMemoryTest: Memory write/read test completed"); - } - - spdlog::debug("BasicMemoryTest: Testing pinned memory allocation"); - // Test pinned memory allocation - auto pinned_memory = context::allocatePinnedHostMemory(512); - spdlog::debug("BasicMemoryTest: Pinned memory allocated"); - if (!pinned_memory) { - std::cerr << "Failed to allocate pinned memory" << std::endl; - return false; - } - - spdlog::debug("BasicMemoryTest: Checking pinned memory properties"); - // For CPU devices, pinned memory falls back to regular memory, so it may not be marked as pinned - Device pinned_device = context::getDevice(); - if (pinned_device.getType() != Device::Type::CPU && !pinned_memory->is_pinned()) { - std::cerr << "Pinned memory not marked as pinned" << std::endl; - return false; - } - spdlog::debug("BasicMemoryTest: Pinned memory test completed"); - - return true; - } catch (const std::exception &e) { - std::cerr << "BasicMemoryTest failed with exception: " << e.what() << std::endl; - return false; - } - }); -} - -// Concurrency Test Implementation -TestResult ConcurrencyTest::run() { - return measureTime("ConcurrencyTest", [this]() -> bool { - try { - // Run all concurrency subtests - auto result1 = testConcurrentAllocations(); - if (!result1.passed) { - std::cerr << "Concurrent allocations test failed: " << result1.error_message << std::endl; - return false; - } - - // auto result2 = testConcurrentDeviceSwitching(); - // if (!result2.passed) { - // std::cerr << "Concurrent device switching test failed: " << result2.error_message << std::endl; - // return false; - // } - - // auto result3 = testMemoryAllocationRace(); - // if (!result3.passed) { - // std::cerr << "Memory allocation race test failed: " << result3.error_message << std::endl; - // return false; - // } - - return true; - } catch (const std::exception &e) { - std::cerr << "ConcurrencyTest failed with exception: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult ConcurrencyTest::testConcurrentAllocations() { - return measureTime("ConcurrentAllocations", [this]() -> bool { - const int num_threads = 8; - const int allocations_per_thread = 100; - std::vector threads; - std::atomic success_count{0}; - std::atomic failure_count{0}; - - for (int i = 0; i < num_threads; ++i) { - threads.emplace_back([&, i]() { - try { - for (int j = 0; j < allocations_per_thread; ++j) { - // Allocate memory of random size - size_t size = 64 + (j % 1024); - auto memory = context::allocateMemory(size); - if (memory && memory->size() == size) { - success_count++; - } else { - failure_count++; - } - - // Small delay to increase chance of race conditions - std::this_thread::sleep_for(std::chrono::microseconds(1)); - } - } catch (const std::exception &e) { - failure_count++; - std::cerr << "Thread " << i << " failed: " << e.what() << std::endl; - } - }); - } - - for (auto &thread : threads) { - thread.join(); - } - - int total_expected = num_threads * allocations_per_thread; - if (success_count.load() != total_expected) { - std::cerr << "Concurrent allocation test failed: expected " << total_expected - << " successes, got " << success_count.load() - << " successes and " << failure_count.load() << " failures" << std::endl; - return false; - } - - return true; - }); -} - -TestResult ConcurrencyTest::testConcurrentDeviceSwitching() { - return measureTime("ConcurrentDeviceSwitching", [this]() -> bool { - const int num_threads = 4; - std::vector threads; - std::atomic success_count{0}; - std::atomic failure_count{0}; - - // Get available devices - std::vector devices; - for (int type = 0; type < static_cast(Device::Type::COUNT); ++type) { - size_t count = context::getDeviceCount(static_cast(type)); - for (size_t i = 0; i < count; ++i) { - devices.emplace_back(static_cast(type), i); - } - } - - if (devices.size() < 2) { - std::cout << "Skipping device switching test - need at least 2 devices" << std::endl; - return true; - } - - for (int i = 0; i < num_threads; ++i) { - threads.emplace_back([&, i, devices]() { - try { - for (int j = 0; j < 50; ++j) { - // Switch to random device - Device target_device = devices[j % devices.size()]; - context::setDevice(target_device); - - // Verify device was set correctly - Device current_device = context::getDevice(); - if (current_device == target_device) { - success_count++; - } else { - failure_count++; - std::cerr << "Device switching failed: expected " - << static_cast(target_device.getType()) - << ", got " << static_cast(current_device.getType()) << std::endl; - } - - // Allocate memory to test device context - auto memory = context::allocateMemory(256); - if (memory && memory->device() == target_device) { - success_count++; - } else { - failure_count++; - } - - std::this_thread::sleep_for(std::chrono::microseconds(10)); - } - } catch (const std::exception &e) { - failure_count++; - std::cerr << "Thread " << i << " failed: " << e.what() << std::endl; - } - }); - } - - for (auto &thread : threads) { - thread.join(); - } - - if (failure_count.load() > 0) { - std::cerr << "Concurrent device switching test failed: " - << failure_count.load() << " failures out of " - << (success_count.load() + failure_count.load()) << " operations" << std::endl; - return false; - } - - return true; - }); -} - -TestResult ConcurrencyTest::testMemoryAllocationRace() { - return measureTime("MemoryAllocationRace", [this]() -> bool { - const int num_threads = 16; - const int allocations_per_thread = 1000; - std::vector threads; - std::atomic success_count{0}; - std::atomic failure_count{0}; - std::vector> all_allocations; - std::mutex allocations_mutex; - - for (int i = 0; i < num_threads; ++i) { - threads.emplace_back([&, i]() { - std::vector> thread_allocations; - try { - for (int j = 0; j < allocations_per_thread; ++j) { - size_t size = 64 + (j % 1024); - auto memory = context::allocateMemory(size); - if (memory) { - thread_allocations.push_back(memory); - success_count++; - } else { - failure_count++; - } - - // Occasionally deallocate some memory to test concurrent alloc/dealloc - if (j % 10 == 0 && !thread_allocations.empty()) { - thread_allocations.pop_back(); - } - } - - // Store remaining allocations - std::lock_guard lock(allocations_mutex); - all_allocations.insert(all_allocations.end(), - thread_allocations.begin(), - thread_allocations.end()); - } catch (const std::exception &e) { - failure_count++; - std::cerr << "Thread " << i << " failed: " << e.what() << std::endl; - } - }); - } - - for (auto &thread : threads) { - thread.join(); - } - - // Verify all allocations are valid - for (const auto &memory : all_allocations) { - if (!memory || !memory->data()) { - std::cerr << "Invalid memory allocation found" << std::endl; - return false; - } - } - - int total_expected = num_threads * allocations_per_thread; - if (success_count.load() < total_expected * 0.9) { // Allow 10% failure rate - std::cerr << "Memory allocation race test failed: expected at least " - << total_expected * 0.9 << " successes, got " << success_count.load() << std::endl; - return false; - } - - return true; - }); -} - -// Exception Safety Test Implementation -TestResult ExceptionSafetyTest::run() { - return measureTime("ExceptionSafetyTest", [this]() -> bool { - try { - auto result1 = testAllocationFailure(); - if (!result1.passed) { - std::cerr << "Allocation failure test failed: " << result1.error_message << std::endl; - return false; - } - - auto result2 = testDeallocationException(); - if (!result2.passed) { - std::cerr << "Deallocation exception test failed: " << result2.error_message << std::endl; - return false; - } - - auto result3 = testContextSwitchException(); - if (!result3.passed) { - std::cerr << "Context switch exception test failed: " << result3.error_message << std::endl; - return false; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "ExceptionSafetyTest failed with exception: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult ExceptionSafetyTest::testAllocationFailure() { - return measureTime("AllocationFailure", [this]() -> bool { - try { - // Test allocation with extremely large size (should fail) - try { - auto memory = context::allocateMemory(SIZE_MAX); - std::cerr << "Expected allocation to fail with SIZE_MAX" << std::endl; - return false; - } catch (const std::exception &e) { - // Expected to fail - std::cout << "Allocation correctly failed with SIZE_MAX: " << e.what() << std::endl; - } - - // Test allocation with zero size - try { - auto memory = context::allocateMemory(0); - if (memory) { - std::cerr << "Zero-size allocation should return null or throw" << std::endl; - return false; - } - } catch (const std::exception &e) { - // Also acceptable - std::cout << "Zero-size allocation correctly failed: " << e.what() << std::endl; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "Allocation failure test failed with unexpected exception: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult ExceptionSafetyTest::testDeallocationException() { - return measureTime("DeallocationException", [this]() -> bool { - try { - // Test that deallocation doesn't throw exceptions - std::vector> memories; - - // Allocate some memory - for (int i = 0; i < 10; ++i) { - auto memory = context::allocateMemory(1024); - if (memory) { - memories.push_back(memory); - } - } - - // Test that destruction doesn't throw - try { - memories.clear(); // This should trigger deallocation - } catch (const std::exception &e) { - std::cerr << "Memory deallocation threw exception: " << e.what() << std::endl; - return false; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "Deallocation exception test failed: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult ExceptionSafetyTest::testContextSwitchException() { - return measureTime("ContextSwitchException", [this]() -> bool { - try { - // Test context switching with invalid device - Device original_device = context::getDevice(); - - try { - // Try to switch to a device that might not exist - Device invalid_device(Device::Type::COUNT, 999); - context::setDevice(invalid_device); - std::cerr << "Expected device switching to fail with invalid device" << std::endl; - return false; - } catch (const std::exception &e) { - // Expected to fail - std::cout << "Device switching correctly failed with invalid device: " << e.what() << std::endl; - } - - // Verify original device is still set - Device current_device = context::getDevice(); - if (current_device != original_device) { - std::cerr << "Device context not restored after exception" << std::endl; - return false; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "Context switch exception test failed: " << e.what() << std::endl; - return false; - } - }); -} - -// Memory Leak Test Implementation -TestResult MemoryLeakTest::run() { - return measureTime("MemoryLeakTest", [this]() -> bool { - try { - auto result1 = testBasicLeakDetection(); - if (!result1.passed) { - std::cerr << "Basic leak detection test failed: " << result1.error_message << std::endl; - return false; - } - - auto result2 = testCrossDeviceLeakDetection(); - if (!result2.passed) { - std::cerr << "Cross-device leak detection test failed: " << result2.error_message << std::endl; - return false; - } - - auto result3 = testExceptionLeakDetection(); - if (!result3.passed) { - std::cerr << "Exception leak detection test failed: " << result3.error_message << std::endl; - return false; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "MemoryLeakTest failed with exception: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult MemoryLeakTest::testBasicLeakDetection() { - return measureTime("BasicLeakDetection", [this]() -> bool { - try { - // Reset leak detector - MemoryLeakDetector::instance().reset(); - - // Allocate and deallocate memory - std::vector> memories; - for (int i = 0; i < 100; ++i) { - auto memory = context::allocateMemory(1024); - if (memory) { - memories.push_back(memory); - } - } - - // Clear memories to trigger deallocation - memories.clear(); - - // Force garbage collection if available - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - - // Check for leaks (this is a basic test - real leak detection would need more sophisticated tools) - size_t leaked_memory = MemoryLeakDetector::instance().getLeakedMemory(); - if (leaked_memory > 0) { - std::cerr << "Potential memory leak detected: " << leaked_memory << " bytes" << std::endl; - return false; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "Basic leak detection test failed: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult MemoryLeakTest::testCrossDeviceLeakDetection() { - return measureTime("CrossDeviceLeakDetection", [this]() -> bool { - try { - // Get available devices - std::vector devices; - for (int type = 0; type < static_cast(Device::Type::COUNT); ++type) { - size_t count = context::getDeviceCount(static_cast(type)); - for (size_t i = 0; i < count; ++i) { - devices.emplace_back(static_cast(type), i); - } - } - - if (devices.size() < 2) { - std::cout << "Skipping cross-device leak test - need at least 2 devices" << std::endl; - return true; - } - - // Allocate pinned memory on one device - context::setDevice(devices[0]); - auto pinned_memory = context::allocatePinnedHostMemory(1024); - - if (!pinned_memory) { - std::cerr << "Failed to allocate pinned memory" << std::endl; - return false; - } - - // Switch to another device and deallocate - context::setDevice(devices[1]); - pinned_memory.reset(); // This should trigger cross-device deallocation - - // Force garbage collection - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - - // Check for leaks - size_t leaked_memory = MemoryLeakDetector::instance().getLeakedMemory(); - if (leaked_memory > 0) { - std::cerr << "Potential cross-device memory leak detected: " << leaked_memory << " bytes" << std::endl; - return false; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "Cross-device leak detection test failed: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult MemoryLeakTest::testExceptionLeakDetection() { - return measureTime("ExceptionLeakDetection", [this]() -> bool { - try { - // Test that exceptions don't cause memory leaks - std::vector> memories; - - try { - // Allocate some memory - for (int i = 0; i < 10; ++i) { - auto memory = context::allocateMemory(1024); - if (memory) { - memories.push_back(memory); - } - } - - // Simulate an exception - throw std::runtime_error("Simulated exception"); - - } catch (const std::exception &e) { - // Memory should still be properly cleaned up - memories.clear(); - } - - // Force garbage collection - std::this_thread::sleep_for(std::chrono::milliseconds(100)); - - // Check for leaks - size_t leaked_memory = MemoryLeakDetector::instance().getLeakedMemory(); - if (leaked_memory > 0) { - std::cerr << "Potential exception-related memory leak detected: " << leaked_memory << " bytes" << std::endl; - return false; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "Exception leak detection test failed: " << e.what() << std::endl; - return false; - } - }); -} - -// Performance Test Implementation -TestResult PerformanceTest::run() { - return measureTime("PerformanceTest", [this]() -> bool { - try { - auto result1 = testAllocationPerformance(); - if (!result1.passed) { - std::cerr << "Allocation performance test failed: " << result1.error_message << std::endl; - return false; - } - - auto result2 = testConcurrentPerformance(); - if (!result2.passed) { - std::cerr << "Concurrent performance test failed: " << result2.error_message << std::endl; - return false; - } - - auto result3 = testMemoryCopyPerformance(); - if (!result3.passed) { - std::cerr << "Memory copy performance test failed: " << result3.error_message << std::endl; - return false; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "PerformanceTest failed with exception: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult PerformanceTest::testAllocationPerformance() { - return measureTime("AllocationPerformance", [this]() -> bool { - try { - const int num_allocations = 10000; - const size_t allocation_size = 1024; - - auto start = std::chrono::high_resolution_clock::now(); - - std::vector> memories; - for (int i = 0; i < num_allocations; ++i) { - auto memory = context::allocateMemory(allocation_size); - if (memory) { - memories.push_back(memory); - } - } - - auto end = std::chrono::high_resolution_clock::now(); - auto duration = std::chrono::duration_cast(end - start); - - double avg_time_per_allocation = static_cast(duration.count()) / num_allocations; - std::cout << "Average allocation time: " << avg_time_per_allocation << "μs" << std::endl; - - // Performance threshold: should be under 100μs per allocation - if (avg_time_per_allocation > 100.0) { - std::cerr << "Allocation performance too slow: " << avg_time_per_allocation << "μs per allocation" << std::endl; - return false; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "Allocation performance test failed: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult PerformanceTest::testConcurrentPerformance() { - return measureTime("ConcurrentPerformance", [this]() -> bool { - try { - const int num_threads = 4; - const int allocations_per_thread = 1000; - - auto start = std::chrono::high_resolution_clock::now(); - - std::vector threads; - std::atomic success_count{0}; - - for (int i = 0; i < num_threads; ++i) { - threads.emplace_back([&]() { - for (int j = 0; j < allocations_per_thread; ++j) { - auto memory = context::allocateMemory(512); - if (memory) { - success_count++; - } - } - }); - } - - for (auto &thread : threads) { - thread.join(); - } - - auto end = std::chrono::high_resolution_clock::now(); - auto duration = std::chrono::duration_cast(end - start); - - double total_allocations = num_threads * allocations_per_thread; - double avg_time_per_allocation = static_cast(duration.count()) / total_allocations; - std::cout << "Concurrent allocation time: " << avg_time_per_allocation << "μs per allocation" << std::endl; - - // Performance threshold: should be under 200μs per allocation under concurrent load - if (avg_time_per_allocation > 200.0) { - std::cerr << "Concurrent allocation performance too slow: " << avg_time_per_allocation << "μs per allocation" << std::endl; - return false; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "Concurrent performance test failed: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult PerformanceTest::testMemoryCopyPerformance() { - return measureTime("MemoryCopyPerformance", [this]() -> bool { - try { - const size_t data_size = 1024 * 1024; // 1MB - const int num_copies = 100; - - // Allocate source and destination memory - auto src_memory = context::allocateMemory(data_size); - auto dst_memory = context::allocateMemory(data_size); - - if (!src_memory || !dst_memory) { - std::cerr << "Failed to allocate memory for copy test" << std::endl; - return false; - } - - auto start = std::chrono::high_resolution_clock::now(); - - // Perform memory copies - for (int i = 0; i < num_copies; ++i) { - context::memcpyD2D(dst_memory->data(), src_memory->data(), data_size); - } - - // Synchronize to ensure copies complete - context::syncDevice(); - - auto end = std::chrono::high_resolution_clock::now(); - auto duration = std::chrono::duration_cast(end - start); - - double avg_time_per_copy = static_cast(duration.count()) / num_copies; - double bandwidth = (data_size * num_copies) / (duration.count() / 1e6) / (1024 * 1024); // MB/s - - std::cout << "Average copy time: " << avg_time_per_copy << "μs" << std::endl; - std::cout << "Memory bandwidth: " << bandwidth << " MB/s" << std::endl; - - // Performance threshold: should achieve at least 100 MB/s - if (bandwidth < 100.0) { - std::cerr << "Memory copy performance too slow: " << bandwidth << " MB/s" << std::endl; - return false; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "Memory copy performance test failed: " << e.what() << std::endl; - return false; - } - }); -} - -// Stress Test Implementation -TestResult StressTest::run() { - return measureTime("StressTest", [this]() -> bool { - try { - auto result1 = testHighFrequencyAllocations(); - if (!result1.passed) { - std::cerr << "High frequency allocations test failed: " << result1.error_message << std::endl; - return false; - } - - auto result2 = testLargeMemoryAllocations(); - if (!result2.passed) { - std::cerr << "Large memory allocations test failed: " << result2.error_message << std::endl; - return false; - } - - auto result3 = testCrossDeviceStress(); - if (!result3.passed) { - std::cerr << "Cross-device stress test failed: " << result3.error_message << std::endl; - return false; - } - - return true; - } catch (const std::exception &e) { - std::cerr << "StressTest failed with exception: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult StressTest::testHighFrequencyAllocations() { - return measureTime("HighFrequencyAllocations", [this]() -> bool { - try { - const int num_allocations = 100000; - std::vector> memories; - memories.reserve(num_allocations); - - auto start = std::chrono::high_resolution_clock::now(); - - for (int i = 0; i < num_allocations; ++i) { - size_t size = 64 + (i % 1024); - auto memory = context::allocateMemory(size); - if (memory) { - memories.push_back(memory); - } - - // Periodically deallocate some memory to test alloc/dealloc stress - if (i % 1000 == 0 && !memories.empty()) { - memories.erase(memories.begin(), memories.begin() + std::min(100, static_cast(memories.size()))); - } - } - - auto end = std::chrono::high_resolution_clock::now(); - auto duration = std::chrono::duration_cast(end - start); - - std::cout << "High frequency allocations completed: " << num_allocations - << " allocations in " << duration.count() << "ms" << std::endl; - - // Clear remaining memory - memories.clear(); - - return true; - } catch (const std::exception &e) { - std::cerr << "High frequency allocations test failed: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult StressTest::testLargeMemoryAllocations() { - return measureTime("LargeMemoryAllocations", [this]() -> bool { - try { - const size_t large_size = 100 * 1024 * 1024; // 100MB - const int num_allocations = 10; - - std::vector> memories; - - for (int i = 0; i < num_allocations; ++i) { - try { - auto memory = context::allocateMemory(large_size); - if (memory) { - memories.push_back(memory); - std::cout << "Allocated " << large_size / (1024 * 1024) << "MB memory block " << i + 1 << std::endl; - } - } catch (const std::exception &e) { - std::cout << "Large allocation " << i + 1 << " failed (expected): " << e.what() << std::endl; - break; // Expected to fail at some point due to memory limits - } - } - - std::cout << "Successfully allocated " << memories.size() << " large memory blocks" << std::endl; - - // Clear memory - memories.clear(); - - return true; - } catch (const std::exception &e) { - std::cerr << "Large memory allocations test failed: " << e.what() << std::endl; - return false; - } - }); -} - -TestResult StressTest::testCrossDeviceStress() { - return measureTime("CrossDeviceStress", [this]() -> bool { - try { - // Get available devices - std::vector devices; - for (int type = 0; type < static_cast(Device::Type::COUNT); ++type) { - size_t count = context::getDeviceCount(static_cast(type)); - for (size_t i = 0; i < count; ++i) { - devices.emplace_back(static_cast(type), i); - } - } - - if (devices.size() < 2) { - std::cout << "Skipping cross-device stress test - need at least 2 devices" << std::endl; - return true; - } - - const int num_operations = 1000; - std::vector> pinned_memories; - - for (int i = 0; i < num_operations; ++i) { - // Switch to random device - Device target_device = devices[i % devices.size()]; - context::setDevice(target_device); - - // Allocate pinned memory - auto pinned_memory = context::allocatePinnedHostMemory(1024); - if (pinned_memory) { - pinned_memories.push_back(pinned_memory); - } - - // Periodically deallocate some memory - if (i % 100 == 0 && !pinned_memories.empty()) { - pinned_memories.erase(pinned_memories.begin(), - pinned_memories.begin() + std::min(10, static_cast(pinned_memories.size()))); - } - } - - std::cout << "Cross-device stress test completed: " << num_operations - << " operations across " << devices.size() << " devices" << std::endl; - - // Clear remaining memory - pinned_memories.clear(); - - return true; - } catch (const std::exception &e) { - std::cerr << "Cross-device stress test failed: " << e.what() << std::endl; - return false; - } - }); -} - -} // namespace infinicore::test diff --git a/src/infinicore-test/memory_test.h b/src/infinicore-test/memory_test.h deleted file mode 100644 index 5fdffd518..000000000 --- a/src/infinicore-test/memory_test.h +++ /dev/null @@ -1,160 +0,0 @@ -#ifndef __INFINICORE_MEMORY_TEST_H__ -#define __INFINICORE_MEMORY_TEST_H__ - -#include "../infinicore/context/allocators/memory_allocator.hpp" -#include "test_runner.h" -#include -#include -#include -#include -#include -#include -#include - -namespace infinicore::test { - -// Mock allocator for testing exception safety -class MockAllocator : public infinicore::MemoryAllocator { -public: - MockAllocator(bool should_throw = false, size_t max_allocations = SIZE_MAX) - : should_throw_(should_throw), max_allocations_(max_allocations), - allocation_count_(0), total_allocated_(0) {} - - std::byte *allocate(size_t size) override { - if (should_throw_) { - throw std::runtime_error("Mock allocation failure"); - } - if (allocation_count_ >= max_allocations_) { - throw std::runtime_error("Mock allocation limit exceeded"); - } - allocation_count_++; - total_allocated_ += size; - return static_cast(std::malloc(size)); - } - - void deallocate(std::byte *ptr) override { - if (ptr) { - std::free(ptr); - } - } - - size_t getAllocationCount() const { return allocation_count_; } - size_t getTotalAllocated() const { return total_allocated_; } - -private: - bool should_throw_; - size_t max_allocations_; - std::atomic allocation_count_; - std::atomic total_allocated_; -}; - -// Memory leak detector -class MemoryLeakDetector { -public: - static MemoryLeakDetector &instance() { - static MemoryLeakDetector detector; - return detector; - } - - void recordAllocation(void *ptr, size_t size) { - std::lock_guard lock(mutex_); - allocations_[ptr] = size; - total_allocated_ += size; - } - - void recordDeallocation(void *ptr) { - std::lock_guard lock(mutex_); - auto it = allocations_.find(ptr); - if (it != allocations_.end()) { - total_allocated_ -= it->second; - allocations_.erase(it); - } - } - - size_t getLeakedMemory() const { - std::lock_guard lock(mutex_); - return total_allocated_; - } - - size_t getLeakCount() const { - std::lock_guard lock(mutex_); - return allocations_.size(); - } - - void reset() { - std::lock_guard lock(mutex_); - allocations_.clear(); - total_allocated_ = 0; - } - -private: - mutable std::mutex mutex_; - std::unordered_map allocations_; - size_t total_allocated_ = 0; -}; - -// Test categories -class BasicMemoryTest : public TestFramework { -public: - TestResult run() override; - std::string getName() const override { return "BasicMemoryTest"; } -}; - -class ConcurrencyTest : public TestFramework { -public: - TestResult run() override; - std::string getName() const override { return "ConcurrencyTest"; } - -private: - TestResult testConcurrentAllocations(); - TestResult testConcurrentDeviceSwitching(); - TestResult testMemoryAllocationRace(); -}; - -class ExceptionSafetyTest : public TestFramework { -public: - TestResult run() override; - std::string getName() const override { return "ExceptionSafetyTest"; } - -private: - TestResult testAllocationFailure(); - TestResult testDeallocationException(); - TestResult testContextSwitchException(); -}; - -class MemoryLeakTest : public TestFramework { -public: - TestResult run() override; - std::string getName() const override { return "MemoryLeakTest"; } - -private: - TestResult testBasicLeakDetection(); - TestResult testCrossDeviceLeakDetection(); - TestResult testExceptionLeakDetection(); -}; - -class PerformanceTest : public TestFramework { -public: - TestResult run() override; - std::string getName() const override { return "PerformanceTest"; } - -private: - TestResult testAllocationPerformance(); - TestResult testConcurrentPerformance(); - TestResult testMemoryCopyPerformance(); -}; - -class StressTest : public TestFramework { -public: - TestResult run() override; - std::string getName() const override { return "StressTest"; } - -private: - TestResult testHighFrequencyAllocations(); - TestResult testLargeMemoryAllocations(); - TestResult testCrossDeviceStress(); -}; - -} // namespace infinicore::test - -#endif // __INFINICORE_MEMORY_TEST_H__ diff --git a/src/infinicore-test/test_mutual_awareness.cc b/src/infinicore-test/test_mutual_awareness.cc deleted file mode 100644 index 506fb92fa..000000000 --- a/src/infinicore-test/test_mutual_awareness.cc +++ /dev/null @@ -1,261 +0,0 @@ -#include "test_mutual_awareness.h" - -#ifdef ENABLE_MUTUAL_AWARENESS - -#include "infinicore/analyzer.hpp" -#include "infinicore/ops/common/dispatcher.hpp" - -#include - -namespace infinicore::test { - -namespace { - -using namespace infinicore::analyzer; - -#define MA_ASSERT_TRUE(cond) \ - do { \ - if (!(cond)) { \ - spdlog::error("ASSERT_TRUE failed: {} ({}:{})", \ - #cond, __FILE__, __LINE__); \ - return false; \ - } \ - } while (0) - -#define MA_ASSERT_EQ(a, b) \ - do { \ - if (!((a) == (b))) { \ - spdlog::error("ASSERT_EQ failed: {} == {} ({}:{})", \ - #a, #b, __FILE__, __LINE__); \ - return false; \ - } \ - } while (0) - -std::vector makeWindow(const std::vector &types, - uint32_t seq_len = 128) { - std::vector window; - for (auto t : types) { - OpTraceEntry e; - e.op_type = t; - e.ndim = 3; - e.shape[0] = 1; - e.shape[1] = 32; - e.shape[2] = seq_len; - window.push_back(e); - } - return window; -} - -bool testOpTraceRing() { - OpTraceRing ring(4); - MA_ASSERT_EQ(ring.size(), 0u); - MA_ASSERT_EQ(ring.capacity(), 4u); - - OpTraceEntry first; - first.op_type = OpType::GEMM; - ring.write(first); - MA_ASSERT_EQ(ring.size(), 1u); - MA_ASSERT_EQ(ring.totalCount(), 1u); - - for (int i = 0; i < 6; ++i) { - OpTraceEntry x; - x.op_type = static_cast(i + 1); - ring.write(x); - } - MA_ASSERT_EQ(ring.size(), 4u); - MA_ASSERT_EQ(ring.totalCount(), 7u); - - auto recent = ring.getRecentEntries(2); - MA_ASSERT_EQ(recent.size(), 2u); - - ring.clear(); - MA_ASSERT_EQ(ring.size(), 0u); - return true; -} - -bool testOpTypeRegistry() { - MA_ASSERT_EQ(opTypeFromName("Add"), OpType::ADD); - MA_ASSERT_EQ(opTypeFromName("FlashAttention"), OpType::FLASH_ATTENTION); - MA_ASSERT_EQ(opTypeFromName("RMSNorm"), OpType::RMS_NORM); - MA_ASSERT_EQ(opTypeFromName("Rearrange"), OpType::REARRANGE); - MA_ASSERT_EQ(opTypeFromName("__nonexistent__"), OpType::UNKNOWN); - return true; -} - -bool testOpTypeClassification() { - MA_ASSERT_TRUE(isAttentionOp(OpType::FLASH_ATTENTION)); - MA_ASSERT_TRUE(isAttentionOp(OpType::PAGED_ATTENTION)); - MA_ASSERT_TRUE(!isAttentionOp(OpType::GEMM)); - MA_ASSERT_TRUE(isGemmMlpOp(OpType::GEMM)); - MA_ASSERT_TRUE(isGemmMlpOp(OpType::LINEAR)); - MA_ASSERT_TRUE(isActivationOp(OpType::SILU)); - MA_ASSERT_TRUE(isKvCacheOp(OpType::KV_CACHING)); - return true; -} - -bool testPhaseDetector() { - PhaseDetector detector; - - auto attention = makeWindow( - {OpType::ATTENTION, OpType::FLASH_ATTENTION, OpType::CAUSAL_SOFTMAX, - OpType::ATTENTION, OpType::RMS_NORM, OpType::ADD}, - /*seq_len=*/16); - MA_ASSERT_EQ(detector.detect(attention), PhaseType::ATTENTION_DENSE); - - auto gemm = makeWindow({OpType::LINEAR, OpType::SILU, OpType::LINEAR, - OpType::GEMM, OpType::GELU, OpType::LINEAR}); - MA_ASSERT_EQ(detector.detect(gemm), PhaseType::GEMM_MLP_DENSE); - - auto kv = makeWindow({OpType::KV_CACHING, OpType::KV_CACHING, - OpType::PAGED_CACHING, OpType::KV_CACHING, - OpType::ADD}); - MA_ASSERT_EQ(detector.detect(kv), PhaseType::KV_CACHE); - - auto decode = makeWindow( - {OpType::ATTENTION, OpType::CAUSAL_SOFTMAX, OpType::LINEAR, - OpType::SILU, OpType::GEMM, OpType::ADD}, - /*seq_len=*/1); - MA_ASSERT_EQ(detector.detect(decode), PhaseType::DECODE); - - auto prefill = makeWindow( - {OpType::ATTENTION, OpType::FLASH_ATTENTION, OpType::CAUSAL_SOFTMAX, - OpType::ATTENTION, OpType::RMS_NORM}, - /*seq_len=*/512); - MA_ASSERT_EQ(detector.detect(prefill), PhaseType::PREFILL); - - std::vector empty; - MA_ASSERT_EQ(detector.detect(empty), PhaseType::UNKNOWN); - return true; -} - -bool testResourceSensor() { - ResourceSensor sensor; - - MemoryStats high{900, 1000, 0, 0}; - auto high_intent = sensor.sense(0, high); - MA_ASSERT_EQ(high_intent.local_bottleneck, BottleneckType::MEMORY_BOUND); - - MemoryStats low{100, 1000, 0, 0}; - auto low_intent = sensor.sense(0, low); - MA_ASSERT_EQ(low_intent.local_bottleneck, BottleneckType::COMPUTE_BOUND); - - DeviceResourceSnapshot snap; - snap.device_id = 0; - snap.device_type = Device::Type::NVIDIA; - snap.has_memory_capacity = true; - snap.total_bytes = 1000; - snap.used_bytes = 300; - snap.free_bytes = 700; - snap.has_memory_bandwidth_utilization = true; - snap.memory_bandwidth_utilization = 0.9f; - auto bw_intent = sensor.sense(snap); - MA_ASSERT_EQ(bw_intent.local_bottleneck, BottleneckType::BANDWIDTH_BOUND); - return true; -} - -bool testIntentGenerator() { - IntentGenerator gen; - - DeviceLocalIntent pressure; - pressure.device_id = 0; - pressure.memory_usage_ratio = 0.95f; - pressure.local_bottleneck = BottleneckType::MEMORY_BOUND; - auto window = makeWindow({OpType::GEMM, OpType::LINEAR}); - auto intent = gen.generate(PhaseType::GEMM_MLP_DENSE, window, {pressure}); - MA_ASSERT_EQ(intent.global.primary_bottleneck, BottleneckType::MEMORY_BOUND); - MA_ASSERT_EQ(intent.global.goal, OptimizationGoal::MEMORY_SAFE); - MA_ASSERT_TRUE(intent.global.strategy.prefer_in_place); - MA_ASSERT_TRUE(intent.global.strategy.prefer_recomputation); - - DeviceLocalIntent fresh; - fresh.device_id = 0; - fresh.memory_usage_ratio = 0.3f; - fresh.local_bottleneck = BottleneckType::COMPUTE_BOUND; - auto attention = makeWindow({OpType::ATTENTION}, /*seq_len=*/1); - auto decode_intent = gen.generate(PhaseType::DECODE, attention, {fresh}); - MA_ASSERT_EQ(decode_intent.global.goal, OptimizationGoal::LATENCY_FIRST); - return true; -} - -bool testAnalyzerWithMemoryStats() { - auto &analyzer = MutualAwarenessAnalyzer::instance(); - auto &trace = getGlobalOpTrace(); - analyzer.setEnabled(true); - analyzer.clearGraphCache(); - trace.clear(); - - size_t dims[] = {1, 32, 512}; - traceOp(OpType::ATTENTION, dims, 3, 0, 0, 0); - traceOp(OpType::FLASH_ATTENTION, dims, 3, 0, 0, 0); - - MemoryStats stats{950 * 1024 * 1024, 1024 * 1024 * 1024UL, 0, 0}; - auto intent = analyzer.analyze({{0, stats}}); - MA_ASSERT_EQ(intent.global.goal, OptimizationGoal::MEMORY_SAFE); - - trace.clear(); - analyzer.clearGraphCache(); - return true; -} - -using ProbeFn = int (*)(int); - -int probeDefault(int v) { return v + 1; } -int probeThroughput(int v) { return v + 100; } - -bool testDispatcherGoalAware() { - op::common::OpDispatcher dispatcher; - - dispatcher.registerDevice(Device::Type::CPU, &probeDefault); - dispatcher.registerDevice(Device::Type::CPU, &probeThroughput, - OptimizationGoal::THROUGHPUT_FIRST); - - // Goal-specific lookup hits the registered kernel. - MA_ASSERT_EQ(dispatcher.lookup(Device::Type::CPU, - OptimizationGoal::THROUGHPUT_FIRST)(5), - 105); - - // Goal without a specific registration falls back to the device default. - MA_ASSERT_EQ(dispatcher.lookup(Device::Type::CPU, - OptimizationGoal::MEMORY_SAFE)(7), - 8); - - // Legacy device-only lookup still works. - MA_ASSERT_EQ(dispatcher.lookup(Device::Type::CPU)(9), 10); - return true; -} - -} // namespace - -TestResult MutualAwarenessTest::run() { - return measureTime("MutualAwarenessAllTests", [this]() { - bool ok = true; - struct Case { - const char *name; - bool (*fn)(); - }; - Case cases[] = { - {"OpTraceRing", &testOpTraceRing}, - {"OpTypeRegistry", &testOpTypeRegistry}, - {"OpTypeClassification", &testOpTypeClassification}, - {"PhaseDetector", &testPhaseDetector}, - {"ResourceSensor", &testResourceSensor}, - {"IntentGenerator", &testIntentGenerator}, - {"AnalyzerWithMemoryStats", &testAnalyzerWithMemoryStats}, - {"DispatcherGoalAware", &testDispatcherGoalAware}, - }; - for (auto &c : cases) { - spdlog::info("[mutual-awareness] running: {}", c.name); - if (!c.fn()) { - spdlog::error("[mutual-awareness] {} FAILED", c.name); - ok = false; - } else { - spdlog::info("[mutual-awareness] {} passed", c.name); - } - } - return ok; - }); -} - -} // namespace infinicore::test - -#endif // ENABLE_MUTUAL_AWARENESS diff --git a/src/infinicore-test/test_mutual_awareness.h b/src/infinicore-test/test_mutual_awareness.h deleted file mode 100644 index f76f335cd..000000000 --- a/src/infinicore-test/test_mutual_awareness.h +++ /dev/null @@ -1,19 +0,0 @@ -#ifndef __INFINICORE_TEST_MUTUAL_AWARENESS_H__ -#define __INFINICORE_TEST_MUTUAL_AWARENESS_H__ - -#ifdef ENABLE_MUTUAL_AWARENESS - -#include "test_runner.h" - -namespace infinicore::test { - -class MutualAwarenessTest : public TestFramework { -public: - TestResult run() override; - std::string getName() const override { return "MutualAwarenessTest"; } -}; - -} // namespace infinicore::test - -#endif // ENABLE_MUTUAL_AWARENESS -#endif // __INFINICORE_TEST_MUTUAL_AWARENESS_H__ diff --git a/src/infinicore-test/test_nn_module.cc b/src/infinicore-test/test_nn_module.cc deleted file mode 100644 index ca0b08ad7..000000000 --- a/src/infinicore-test/test_nn_module.cc +++ /dev/null @@ -1,2076 +0,0 @@ -#include "test_nn_module.h" -#include "infinicore/ops.hpp" - -namespace infinicore::test { - -// Helper function to format shape for logging -inline std::string formatShape(const std::vector &shape) { - std::ostringstream oss; - oss << "["; - for (size_t i = 0; i < shape.size(); ++i) { - if (i > 0) { - oss << ", "; - } - oss << shape[i]; - } - oss << "]"; - return oss.str(); -} - -// Test 1: Basic module operations (creation, parameters, state_dict, load_state_dict) -TestResult NNModuleTest::testBasicModuleCreation() { - return measureTime("BasicModuleOperations", [this]() { - try { - spdlog::info("=========================================="); - spdlog::info("Testing Basic Module Operations"); - spdlog::info("=========================================="); - - // Test 1a: Module creation and parameter registration - spdlog::info("Test 1a: Module creation and parameter registration"); - MockLinearModule module(8, 4, infinicore::Device()); - - // Verify the module was created successfully - auto state_dict = module.state_dict(); - if (state_dict.size() != 2) { - spdlog::error("Expected 2 parameters, got {}", state_dict.size()); - return false; - } - - // Test weight and bias parameters - const auto &weight = module.get_weight(); - const auto &bias = module.get_bias(); - - // Verify parameter shapes - if (weight->shape() != std::vector({4, 8})) { - spdlog::error("Weight shape mismatch. Expected {{4, 8}}"); - return false; - } - - if (bias->shape() != std::vector({4})) { - spdlog::error("Bias shape mismatch. Expected {{4}}"); - return false; - } - - spdlog::info("✓ Module creation and parameter registration passed"); - - // Test 1b: State dictionary functionality - spdlog::info("Test 1b: State dictionary functionality"); - - // Check if both parameters are in state dict - if (state_dict.find("weight") == state_dict.end()) { - spdlog::error("'weight' parameter not found in state dict"); - return false; - } - - if (state_dict.find("bias") == state_dict.end()) { - spdlog::error("'bias' parameter not found in state dict"); - return false; - } - - spdlog::debug("State dict contains {} parameters:", state_dict.size()); - for (const auto &[name, tensor] : state_dict) { - std::ostringstream shape_str; - shape_str << "["; - for (size_t i = 0; i < tensor->shape().size(); ++i) { - if (i > 0) { - shape_str << ", "; - } - shape_str << tensor->shape()[i]; - } - shape_str << "]"; - spdlog::debug(" - {} with shape: {}", name, shape_str.str()); - } - - spdlog::info("✓ State dict functionality passed"); - - // Test 1c: Load state dict functionality - spdlog::info("Test 1c: Load state dict functionality"); - - // Create new tensors to load - auto new_weight = infinicore::Tensor::ones({4, 8}, infinicore::DataType::F32, infinicore::Device()); - auto new_bias = infinicore::Tensor::zeros({4}, infinicore::DataType::F32, infinicore::Device()); - - // Load using load_parameter_ - module.load_parameter_("weight", new_weight); - module.load_parameter_("bias", new_bias); - - // Verify the parameters were updated - auto updated_state_dict = module.state_dict(); - if (!tensorsAllClose(updated_state_dict.at("weight"), new_weight, 1e-6, 1e-6)) { - spdlog::error("Weight parameter values do not match after load_parameter_"); - return false; - } - if (!tensorsAllClose(updated_state_dict.at("bias"), new_bias, 1e-6, 1e-6)) { - spdlog::error("Bias parameter values do not match after load_parameter_"); - return false; - } - - // Test load_state_dict - std::unordered_map new_state_dict; - new_state_dict.emplace("weight", infinicore::Tensor::ones({4, 8}, infinicore::DataType::F32, infinicore::Device())); - new_state_dict.emplace("bias", infinicore::Tensor::ones({4}, infinicore::DataType::F32, infinicore::Device())); - - module.load_state_dict(new_state_dict); - - auto final_state_dict = module.state_dict(); - if (final_state_dict.size() != 2) { - spdlog::error("State dict size mismatch after load_state_dict"); - return false; - } - - spdlog::info("✓ Load state dict functionality passed"); - - spdlog::info("=== All Basic Module Operations Passed ==="); - return true; - } catch (const std::exception &e) { - spdlog::error("Exception in testBasicModuleOperations: {}", e.what()); - return false; - } - }); -} - -TestResult NNModuleTest::testTensorParallelParameters() { - return measureTime("TensorParallelParameters", [this]() { - try { - spdlog::info("=========================================="); - spdlog::info("Testing Tensor Parallel Parameters"); - spdlog::info("=========================================="); - - auto device = infinicore::context::getDevice(); - - spdlog::info("Test Tensor Parallel Parameter"); - // Case 1: Partition along dimension 0 (row-wise partitioning) - infinicore::nn::Parameter param_dim0({8, 4}, infinicore::DataType::F32, device, 0, 0, 2); - if (param_dim0->shape() != std::vector({4, 4})) { - spdlog::error("TP dim0: Expected shape [4, 4], got [{}]", formatShape(param_dim0->shape())); - return false; - } - spdlog::info("✓ TP dim0 parameter created with correct partitioned shape"); - // Case 2: Partition along dimension 1 (column-wise partitioning) - infinicore::nn::Parameter param_dim1({8, 4}, infinicore::DataType::F32, device, 1, 0, 2); - if (param_dim1->shape() != std::vector({8, 2})) { - spdlog::error("TP dim1: Expected shape [8, 2], got [{}]", formatShape(param_dim1->shape())); - return false; - } - spdlog::info("✓ TP dim1 parameter created with correct partitioned shape"); - spdlog::info("✓ Parameter creation with tensor parallelism passed"); - - spdlog::info("Test Tensor Parallel Linear Module"); - auto w_data = std::vector(32 * 64); - auto b_data = std::vector(32); - for (size_t i = 0; i < 32; ++i) { - for (size_t j = 0; j < 64; ++j) { - w_data[i * 64 + j] = static_cast(j); - } - b_data[i] = static_cast(i); - } - { - spdlog::info("Test tp_size=4 tp_dim=0"); - Size tp_size = 4; - Size tp_dim = 0; - std::vector> tp_modules; - - for (Size tp_rank = 0; tp_rank < tp_size; ++tp_rank) { - auto module = std::make_unique(64, 32, device, tp_dim, tp_rank, tp_size); - tp_modules.push_back(std::move(module)); - } - - // Verify each partition has correct shape - for (size_t i = 0; i < tp_modules.size(); ++i) { - const auto &weight = tp_modules[i]->get_weight(); - const auto &bias = tp_modules[i]->get_bias(); - - // Weight should be partitioned along output dimension (dim 0) - if (weight->shape() != std::vector({8, 64})) { // 32/4 = 8 - spdlog::error("TP rank {}: Weight shape mismatch. Expected [8, 64], got [{}]", - i, formatShape(weight->shape())); - return false; - } - - // Bias should be partitioned along output dimension - if (bias->shape() != std::vector({8})) { // 32/4 = 8 - spdlog::error("TP rank {}: Bias shape mismatch. Expected [8], got [{}]", - i, formatShape(bias->shape())); - return false; - } - - spdlog::debug("TP rank {}: weight shape [{}], bias shape [{}]", - i, formatShape(weight->shape()), formatShape(bias->shape())); - - tp_modules[i]->load_parameter_from_blob("weight", w_data.data()); - tp_modules[i]->load_parameter_from_blob("bias", b_data.data()); - - auto weight_loaded = infinicore::Tensor::from_blob( - w_data.data(), - {32, 64}, - infinicore::DataType::F32, - infinicore::Device::cpu()) - ->narrow({{0, i * 8, 8}}) - ->to(device); // Narrow to get the partition - auto bias_loaded = infinicore::Tensor::from_blob( - b_data.data(), - {32}, - infinicore::DataType::F32, - infinicore::Device::cpu()) - ->narrow({{0, i * 8, 8}}) - ->to(device); // Narrow to get the partition - - if (!tensorsAllClose(tp_modules[i]->get_weight(), weight_loaded, 1e-6, 1e-6)) { - spdlog::error("TP rank {}: Weight values do not match after load_parameter_from_blob", i); - return false; - } - - if (!tensorsAllClose(tp_modules[i]->get_bias(), bias_loaded, 1e-6, 1e-6)) { - spdlog::error("TP rank {}: Bias values do not match after load_parameter_from_blob", i); - return false; - } - } - } - - { - spdlog::info("Test tp_size=4 tp_dim=1"); - Size tp_size = 4; - Size tp_dim = 1; - std::vector> tp_modules; - - for (Size tp_rank = 0; tp_rank < tp_size; ++tp_rank) { - auto module = std::make_unique(64, 32, device, tp_dim, tp_rank, tp_size); - tp_modules.push_back(std::move(module)); - } - - // Verify each partition has correct shape - for (size_t i = 0; i < tp_modules.size(); ++i) { - const auto &weight = tp_modules[i]->get_weight(); - const auto &bias = tp_modules[i]->get_bias(); - - // Weight should be partitioned along output dimension (dim 0) - if (weight->shape() != std::vector({32, 16})) { // 64/4 = 16 - spdlog::error("TP rank {}: Weight shape mismatch. Expected [32, 16], got [{}]", - i, formatShape(weight->shape())); - return false; - } - - // Bias should be partitioned along output dimension - if (bias->shape() != std::vector({32})) { // Bias not partitioned when tp_dim=1 - spdlog::error("TP rank {}: Bias shape mismatch. Expected [32], got [{}]", - i, formatShape(bias->shape())); - return false; - } - - spdlog::debug("TP rank {}: weight shape [{}], bias shape [{}]", - i, formatShape(weight->shape()), formatShape(bias->shape())); - ; - tp_modules[i]->load_parameter_from_blob("weight", w_data.data()); - tp_modules[i]->load_parameter_from_blob("bias", b_data.data()); - - auto weight_loaded = infinicore::Tensor::from_blob( - w_data.data(), - {32, 64}, - infinicore::DataType::F32, - infinicore::Device::cpu()) - ->narrow({{1, i * 16, 16}}) - ->to(device); // Narrow to get the partition - auto bias_loaded = infinicore::Tensor::from_blob( - b_data.data(), - {32}, - infinicore::DataType::F32, - infinicore::Device::cpu()) - ->to(device); // Narrow to get the partition - if (!tensorsAllClose(tp_modules[i]->get_weight(), weight_loaded, 1e-6, 1e-6)) { - spdlog::error("TP rank {}: Weight values do not match after load_parameter_from_blob", i); - return false; - } - if (!tensorsAllClose(tp_modules[i]->get_bias(), bias_loaded, 1e-6, 1e-6)) { - spdlog::error("TP rank {}: Bias values do not match after load_parameter_from_blob", i); - return false; - } - } - } - - spdlog::info("=== All Tensor Parallel Parameter Tests Passed ==="); - return true; - - } catch (const std::exception &e) { - spdlog::error("Exception in testTensorParallelParameters: {}", e.what()); - return false; - } - }); -} - -TestResult NNModuleTest::testParalleLinear() { - return measureTime("ParalleLinear", [this]() { - try { - spdlog::info("=========================================="); - spdlog::info(" Testing Tensor Parallel Linear "); - spdlog::info("=========================================="); - - auto device = infinicore::context::getDevice(); - - spdlog::info("Test Tensor Parallel Linear"); - - spdlog::info(device.toString()); - - auto w_data = std::vector(32 * 64); - auto b_data = std::vector(32); - for (size_t i = 0; i < 32; ++i) { - for (size_t j = 0; j < 64; ++j) { - w_data[i * 64 + j] = static_cast(j); - } - b_data[i] = static_cast(i); - } - { - spdlog::info("Test tp_size=4 tp_dim=0"); - Size tp_size = 4; - // Size tp_dim = 0; - std::vector> tp_modules; - - for (Size tp_rank = 0; tp_rank < tp_size; ++tp_rank) { - auto module = std::make_unique(64, 32, true, DataType::F32, device, tp_rank, tp_size); - tp_modules.push_back(std::move(module)); - } - - // Verify each partition has correct shape - for (size_t i = 0; i < tp_modules.size(); ++i) { - const auto &weight = tp_modules[i]->weight(); - const auto &bias = tp_modules[i]->bias(); - - // Weight should be partitioned along output dimension (dim 0) - if (weight->shape() != std::vector({8, 64})) { // 32/4 = 8 - spdlog::error("TP rank {}: Weight shape mismatch. Expected [8, 64], got [{}]", - i, formatShape(weight->shape())); - return false; - } - - // Bias should be partitioned along output dimension - if (bias->shape() != std::vector({8})) { // 32/4 = 8 - spdlog::error("TP rank {}: Bias shape mismatch. Expected [8], got [{}]", - i, formatShape(bias->shape())); - return false; - } - - spdlog::debug("TP rank {}: weight shape [{}], bias shape [{}]", - i, formatShape(weight->shape()), formatShape(bias->shape())); - - tp_modules[i]->load_parameter_from_blob("weight", w_data.data()); - tp_modules[i]->load_parameter_from_blob("bias", b_data.data()); - - auto weight_loaded = infinicore::Tensor::from_blob( - w_data.data(), - {32, 64}, - infinicore::DataType::F32, - infinicore::Device::cpu()) - ->narrow({{0, i * 8, 8}}) - ->to(device); // Narrow to get the partition - auto bias_loaded = infinicore::Tensor::from_blob( - b_data.data(), - {32}, - infinicore::DataType::F32, - infinicore::Device::cpu()) - ->narrow({{0, i * 8, 8}}) - ->to(device); // Narrow to get the partition - - if (!tensorsAllClose(tp_modules[i]->weight(), weight_loaded, 1e-6, 1e-6)) { - spdlog::error("TP rank {}: Weight values do not match after load_parameter_from_blob", i); - return false; - } - - if (!tensorsAllClose(tp_modules[i]->bias(), bias_loaded, 1e-6, 1e-6)) { - spdlog::error("TP rank {}: Bias values do not match after load_parameter_from_blob", i); - return false; - } - } - } - - { - spdlog::info("Test tp_size=4 tp_dim=1"); - Size tp_size = 4; - // Size tp_dim = 1; - std::vector> tp_modules; - - for (Size tp_rank = 0; tp_rank < tp_size; ++tp_rank) { - - auto module = std::make_unique(64, 32, true, DataType::F32, device, tp_rank, tp_size); - tp_modules.push_back(std::move(module)); - } - - // Verify each partition has correct shape - for (size_t i = 0; i < tp_modules.size(); ++i) { - const auto &weight = tp_modules[i]->weight(); - const auto &bias = tp_modules[i]->bias(); - - // Weight should be partitioned along output dimension (dim 0) - if (weight->shape() != std::vector({32, 16})) { // 64/4 = 16 - spdlog::error("TP rank {}: Weight shape mismatch. Expected [32, 16], got [{}]", - i, formatShape(weight->shape())); - return false; - } - - // Bias should be partitioned along output dimension - if (bias->shape() != std::vector({32})) { // Bias not partitioned when tp_dim=1 - spdlog::error("TP rank {}: Bias shape mismatch. Expected [32], got [{}]", - i, formatShape(bias->shape())); - return false; - } - - spdlog::debug("TP rank {}: weight shape [{}], bias shape [{}]", - i, formatShape(weight->shape()), formatShape(bias->shape())); - ; - tp_modules[i]->load_parameter_from_blob("weight", w_data.data()); - tp_modules[i]->load_parameter_from_blob("bias", b_data.data()); - - auto weight_loaded = infinicore::Tensor::from_blob( - w_data.data(), - {32, 64}, - infinicore::DataType::F32, - infinicore::Device::cpu()) - ->narrow({{1, i * 16, 16}}) - ->to(device); // Narrow to get the partition - auto bias_loaded = infinicore::Tensor::from_blob( - b_data.data(), - {32}, - infinicore::DataType::F32, - infinicore::Device::cpu()) - ->to(device); // Narrow to get the partition - if (!tensorsAllClose(tp_modules[i]->weight(), weight_loaded, 1e-6, 1e-6)) { - spdlog::error("TP rank {}: Weight values do not match after load_parameter_from_blob", i); - return false; - } - if (!tensorsAllClose(tp_modules[i]->bias(), bias_loaded, 1e-6, 1e-6)) { - spdlog::error("TP rank {}: Bias values do not match after load_parameter_from_blob", i); - return false; - } - } - } - - spdlog::info("=== All Tensor Parallel Linear Tests Passed ==="); - return true; - - } catch (const std::exception &e) { - spdlog::error("Exception in testTensorParallelParameters: {}", e.what()); - return false; - } - }); -} - -// Test 2: Advanced load state dict functionality (hierarchical modules) -TestResult NNModuleTest::testLoadStateDict() { - return measureTime("AdvancedLoadStateDict", [this]() { - try { - spdlog::info("=========================================="); - spdlog::info("Testing Advanced load_state_dict with Hierarchical Modules"); - spdlog::info("=========================================="); - - // Test: Deep nesting (2-level hierarchy) - spdlog::info("Test 4: Testing load_state_dict with 2-level deep nesting"); - - // Create parent -> child -> grandchild hierarchy using proper module definition - class DeepGrandchildModule : public infinicore::nn::Module { - protected: - INFINICORE_NN_MODULE(MockLinearModule, sublayer); - - public: - DeepGrandchildModule() { - INFINICORE_NN_MODULE_INIT(sublayer, 6, 4, infinicore::Device()); - } - }; - - class DeepChildModule : public infinicore::nn::Module { - protected: - INFINICORE_NN_MODULE(MockLinearModule, own_layer); - INFINICORE_NN_MODULE(DeepGrandchildModule, sublayer); - - public: - DeepChildModule() { - INFINICORE_NN_MODULE_INIT(own_layer, 8, 6, infinicore::Device()); - INFINICORE_NN_MODULE_INIT(sublayer); - } - }; - - class DeepParentModule : public infinicore::nn::Module { - protected: - INFINICORE_NN_MODULE(MockLinearModule, own_layer); - INFINICORE_NN_MODULE(DeepChildModule, layer1); - - public: - DeepParentModule() { - INFINICORE_NN_MODULE_INIT(own_layer, 10, 8, infinicore::Device()); - INFINICORE_NN_MODULE_INIT(layer1); - } - }; - - DeepParentModule deep_parent; - - // Verify initial state dict includes all 2-level hierarchical parameters - auto deep_initial_state = deep_parent.state_dict(); - spdlog::debug("Deep hierarchical state dict has {} parameters", deep_initial_state.size()); - - // Expected parameters: - // parent: own_layer.weight, own_layer.bias (2) - // layer1: layer1.own_layer.weight, layer1.own_layer.bias (2) - // sublayer: layer1.sublayer.sublayer.weight, layer1.sublayer.sublayer.bias (2) - // Total: 6 parameters - if (deep_initial_state.size() < 6) { - spdlog::error("Deep hierarchy state dict size mismatch. Expected at least 6, got {}", - deep_initial_state.size()); - return false; - } - - // Verify 2-level parameter names exist - bool has_sublayer_weight = deep_initial_state.find("layer1.sublayer.sublayer.weight") != deep_initial_state.end(); - bool has_sublayer_bias = deep_initial_state.find("layer1.sublayer.sublayer.bias") != deep_initial_state.end(); - - if (!has_sublayer_weight || !has_sublayer_bias) { - spdlog::error("2-level nested parameters missing from state dict"); - return false; - } - spdlog::debug("All 2-level hierarchical parameter names verified"); - - // Create state dict for 2-level hierarchy with all 1.0 values - std::unordered_map deep_state_dict; - deep_state_dict.emplace("own_layer.weight", infinicore::Tensor::ones({8, 10}, infinicore::DataType::F32, infinicore::Device())); - deep_state_dict.emplace("own_layer.bias", infinicore::Tensor::ones({8}, infinicore::DataType::F32, infinicore::Device())); - deep_state_dict.emplace("layer1.own_layer.weight", infinicore::Tensor::ones({6, 8}, infinicore::DataType::F32, infinicore::Device())); - deep_state_dict.emplace("layer1.own_layer.bias", infinicore::Tensor::ones({6}, infinicore::DataType::F32, infinicore::Device())); - deep_state_dict.emplace("layer1.sublayer.sublayer.weight", infinicore::Tensor::ones({4, 6}, infinicore::DataType::F32, infinicore::Device())); - deep_state_dict.emplace("layer1.sublayer.sublayer.bias", infinicore::Tensor::ones({4}, infinicore::DataType::F32, infinicore::Device())); - - // Load the deep hierarchical state dict - deep_parent.load_state_dict(deep_state_dict); - spdlog::debug("Successfully loaded 2-level deep hierarchical state dict"); - - // Verify all parameters were loaded correctly - auto deep_loaded_state = deep_parent.state_dict(); - - // Verify shapes at all levels - if (deep_loaded_state.at("own_layer.weight")->shape() != std::vector({8, 10})) { - spdlog::error("Deep parent weight shape mismatch"); - return false; - } - if (deep_loaded_state.at("layer1.own_layer.weight")->shape() != std::vector({6, 8})) { - spdlog::error("Deep layer1 weight shape mismatch"); - return false; - } - if (deep_loaded_state.at("layer1.sublayer.sublayer.weight")->shape() != std::vector({4, 6})) { - spdlog::error("Deep sublayer weight shape mismatch"); - return false; - } - spdlog::debug("All 2-level deep parameter shapes verified"); - - // Verify actual weight loading correctness by checking that loaded parameters - // match what we provided in the state dict (use the original tensors) - spdlog::info("Verifying weight loading correctness by direct comparison"); - - // Get the tensors we loaded from the state dict - auto loaded_parent_weight = deep_loaded_state.at("own_layer.weight"); - auto loaded_layer1_weight = deep_loaded_state.at("layer1.own_layer.weight"); - auto loaded_sublayer_weight = deep_loaded_state.at("layer1.sublayer.sublayer.weight"); - - // Compare with the original tensors we put in the state dict - if (!tensorsAllClose(loaded_parent_weight, deep_state_dict.at("own_layer.weight"), 1e-5, 1e-5)) { - spdlog::error("Deep parent weight not preserved after loading"); - return false; - } - if (!tensorsAllClose(loaded_layer1_weight, deep_state_dict.at("layer1.own_layer.weight"), 1e-5, 1e-5)) { - spdlog::error("Deep layer1 weight not preserved after loading"); - return false; - } - if (!tensorsAllClose(loaded_sublayer_weight, deep_state_dict.at("layer1.sublayer.sublayer.weight"), 1e-5, 1e-5)) { - spdlog::error("Deep sublayer weight not preserved after loading"); - return false; - } - - spdlog::info("✓ Weight loading correctness verified - loaded values match input state dict"); - spdlog::info("✓ 2-level deep hierarchy load_state_dict verification passed"); - - spdlog::info("=== All Advanced load_state_dict Tests Passed ==="); - return true; - } catch (const std::exception &e) { - spdlog::error("Exception in testLoadStateDict: {}", e.what()); - return false; - } - }); -} - -// Test 3: Module hierarchy (demonstrates proper hierarchical construction pattern) -TestResult NNModuleTest::testModuleHierarchy() { - return measureTime("ModuleHierarchy", [this]() { - try { - // Create a hierarchy using proper module definition: root -> layer1 -> layer2 - class Layer2Module : public infinicore::nn::Module { - protected: - INFINICORE_NN_MODULE(MockLinearModule, sublayer); - - public: - Layer2Module() { - INFINICORE_NN_MODULE_INIT(sublayer, 8, 4, infinicore::Device()); - } - }; - - class Layer1Module : public infinicore::nn::Module { - protected: - INFINICORE_NN_MODULE(MockLinearModule, sublayer); - INFINICORE_NN_MODULE(Layer2Module, layer2); - - public: - Layer1Module() { - INFINICORE_NN_MODULE_INIT(sublayer, 16, 8, infinicore::Device()); - INFINICORE_NN_MODULE_INIT(layer2); - } - }; - - class RootModule : public infinicore::nn::Module { - protected: - INFINICORE_NN_MODULE(MockLinearModule, root_layer); - INFINICORE_NN_MODULE(Layer1Module, layer1); - - public: - RootModule() { - INFINICORE_NN_MODULE_INIT(root_layer, 20, 16, infinicore::Device()); - INFINICORE_NN_MODULE_INIT(layer1); - } - }; - - RootModule root_module; - - // Check the complete state dict - auto root_state_dict = root_module.state_dict(); - - // Debug: Print all parameters - spdlog::debug("Found {} parameters:", root_state_dict.size()); - for (const auto &pair : root_state_dict) { - spdlog::debug(" - {}", pair.first); - } - - // Should have: root_layer.weight, root_layer.bias, - // layer1.sublayer.weight, layer1.sublayer.bias, - // layer1.layer2.sublayer.weight, layer1.layer2.sublayer.bias - if (root_state_dict.size() < 6) { - spdlog::error("Error: Expected at least 6 parameters in hierarchy, got {}", root_state_dict.size()); - return false; - } - - spdlog::info("Module hierarchy test passed. Root state dict has {} parameters", root_state_dict.size()); - - // Print the hierarchy - std::cout << "Module hierarchy:" << std::endl; - for (const auto &pair : root_state_dict) { - std::cout << " - " << pair.first << std::endl; - } - - // Additional: Test INFINICORE_NN_MODULE_VEC vector registration - spdlog::info("Testing INFINICORE_NN_MODULE_VEC (vector of submodules)"); - class VecModule : public infinicore::nn::Module { - protected: - INFINICORE_NN_MODULE_VEC(MockLinearModule, layers); - - public: - VecModule() { - INFINICORE_NN_MODULE_VEC_INIT(layers, 3, MockLinearModule, 16, 8, infinicore::Device()); - } - }; - - VecModule vec_mod; - auto vec_state = vec_mod.state_dict(); - - // Expect parameters for layers.0, layers.1, layers.2 (weight and bias for each) - std::vector expected_vec_params = { - "layers.0.weight", "layers.0.bias", - "layers.1.weight", "layers.1.bias", - "layers.2.weight", "layers.2.bias"}; - - for (const auto ¶m : expected_vec_params) { - if (vec_state.find(param) == vec_state.end()) { - spdlog::error("INFINICORE_NN_MODULE_VEC: missing '{}' in state_dict", param); - return false; - } - } - - spdlog::info("INFINICORE_NN_MODULE_VEC test passed - found all vector layer parameters"); - - return true; - } catch (const std::exception &e) { - spdlog::error("Exception in testModuleHierarchy: {}", e.what()); - return false; - } - }); -} - -// Test 4: Parameter loading from blob -TestResult NNModuleTest::testParameterLoading() { - return measureTime("ParameterLoading", [this]() { - try { - spdlog::info("=========================================="); - spdlog::info("Testing Parameter loading from blob"); - spdlog::info("=========================================="); - MockLinearModule module(3, 2, infinicore::Device()); - - // Create test data - std::vector weight_data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; - std::vector bias_data = {0.1f, 0.2f}; - - // Load parameters from blob data - module.load_parameter_from_blob("weight", weight_data.data()); - module.load_parameter_from_blob("bias", bias_data.data()); - - spdlog::info("Successfully loaded parameters from blob data"); - - // Verify parameters exist - auto state_dict = module.state_dict(); - if (state_dict.find("weight") == state_dict.end() || state_dict.find("bias") == state_dict.end()) { - spdlog::error("Error: Parameters not found after loading"); - return false; - } - - MockLinearModule module_row_parallel(3, 2, infinicore::Device(), 0, 1, 2); - - spdlog::info("Parameter loading test passed"); - return true; - } catch (const std::exception &e) { - spdlog::error("Exception in testParameterLoading: {}", e.what()); - return false; - } - }); -} - -// Test 5: Linear module implementation and behavior -TestResult NNModuleTest::testModuleLinear() { - return measureTime("ModuleLinear", [this]() { - try { - // Test with bias - spdlog::info("=========================================="); - spdlog::info("Testing Linear module with bias (8->4 features)"); - spdlog::info("=========================================="); - infinicore::nn::Linear m1(8, 4, true); - auto sd1 = m1.state_dict(); - if (sd1.find("weight") == sd1.end()) { - spdlog::error("weight missing"); - return false; - } - if (sd1.find("bias") == sd1.end()) { - spdlog::error("bias missing when bias=true"); - return false; - } - if (sd1.at("weight")->shape() != std::vector({4, 8})) { - spdlog::error("weight shape mismatch. Expected {{4, 8}}, got different shape"); - return false; - } - if (sd1.at("bias")->shape() != std::vector({4})) { - spdlog::error("bias shape mismatch. Expected {{4}}, got different shape"); - return false; - } - spdlog::debug("Parameter shapes verified: weight {{4, 8}}, bias {{4}}"); - - // Test module properties - if (m1.in_features() != 8) { - spdlog::error("in_features mismatch. Expected 8, got {}", m1.in_features()); - return false; - } - if (m1.out_features() != 4) { - spdlog::error("out_features mismatch. Expected 4, got {}", m1.out_features()); - return false; - } - if (!m1.has_bias()) { - spdlog::error("has_bias should be true"); - return false; - } - - // Test linear computation with bias - spdlog::info("Testing linear computation with bias"); - auto input1 = infinicore::Tensor::ones({2, 8}, infinicore::DataType::F32, infinicore::Device()); - auto output1 = m1.forward(input1); - if (output1->shape() != std::vector({2, 4})) { - spdlog::error("Linear output shape mismatch with bias. Expected {{2, 4}}, got different shape"); - return false; - } - spdlog::debug("Linear computation with bias passed. Input shape: {{2, 8}}, Output shape: {{2, 4}}"); - - // Test without bias - spdlog::info("Testing Linear module without bias (16->3 features)"); - infinicore::nn::Linear m2(16, 3, false); - auto sd2 = m2.state_dict(); - if (sd2.find("weight") == sd2.end()) { - spdlog::error("weight missing (no-bias)"); - return false; - } - if (sd2.find("bias") != sd2.end()) { - spdlog::error("bias should not exist when bias=false"); - return false; - } - if (sd2.at("weight")->shape() != std::vector({3, 16})) { - spdlog::error("weight shape mismatch (no-bias). Expected {{3, 16}}, got different shape"); - return false; - } - spdlog::debug("Parameter shapes verified: weight {{3, 16}}, no bias"); - - // Test module properties - if (m2.in_features() != 16) { - spdlog::error("in_features mismatch. Expected 16, got {}", m2.in_features()); - return false; - } - if (m2.out_features() != 3) { - spdlog::error("out_features mismatch. Expected 3, got {}", m2.out_features()); - return false; - } - if (m2.has_bias()) { - spdlog::error("has_bias should be false"); - return false; - } - - // Test linear computation without bias - spdlog::info("Testing linear computation without bias"); - auto input2 = infinicore::Tensor::ones({1, 16}, infinicore::DataType::F32, infinicore::Device()); - auto output2 = m2.forward(input2); - if (output2->shape() != std::vector({1, 3})) { - spdlog::error("Linear output shape mismatch without bias. Expected {{1, 3}}, got different shape"); - return false; - } - spdlog::debug("Linear computation without bias passed. Input shape: {{1, 16}}, Output shape: {{1, 3}}"); - - // Test load_state_dict for m2 (without bias) - spdlog::info("Testing load_state_dict on Linear without bias"); - auto m2_load_weight = infinicore::Tensor::ones({3, 16}, infinicore::DataType::F32, infinicore::Device()); - std::unordered_map m2_state_dict; - m2_state_dict.emplace("weight", m2_load_weight); - // Note: no bias parameter - m2.load_state_dict(m2_state_dict); - - // Verify via state_dict() and direct access - if (!tensorsAllClose(m2.state_dict().at("weight"), m2_load_weight, 1e-5, 1e-5)) { - spdlog::error("m2 weight not loaded correctly"); - return false; - } - if (!tensorsAllClose(m2.weight(), m2_load_weight, 1e-5, 1e-5)) { - spdlog::error("m2 weight field not synchronized"); - return false; - } - spdlog::debug("m2 load_state_dict verified - weight loaded correctly (no bias)"); - - // Test batch processing - spdlog::info("Testing batch linear computation (batch size 3)"); - auto input3 = infinicore::Tensor::ones({3, 8}, infinicore::DataType::F32, infinicore::Device()); - auto output3 = m1.forward(input3); - if (output3->shape() != std::vector({3, 4})) { - spdlog::error("Batch linear output shape mismatch. Expected {{3, 4}}, got different shape"); - return false; - } - spdlog::debug("Batch linear computation passed. Input shape: {{3, 8}}, Output shape: {{3, 4}}"); - - // Test parameter accessors - spdlog::info("Testing parameter accessors"); - auto weight_accessor = m1.weight(); - auto bias_accessor = m1.bias(); - if (weight_accessor->shape() != std::vector({4, 8})) { - spdlog::error("Weight accessor shape mismatch"); - return false; - } - if (bias_accessor->shape() != std::vector({4})) { - spdlog::error("Bias accessor shape mismatch"); - return false; - } - - // Test load_state_dict for m1 (with bias) - spdlog::info("Testing load_state_dict on Linear with bias"); - auto m1_load_weight = infinicore::Tensor::ones({4, 8}, infinicore::DataType::F32, infinicore::Device()); - auto m1_load_bias = infinicore::Tensor::ones({4}, infinicore::DataType::F32, infinicore::Device()); - std::unordered_map m1_state_dict; - m1_state_dict.emplace("weight", m1_load_weight); - m1_state_dict.emplace("bias", m1_load_bias); - m1.load_state_dict(m1_state_dict); - - // Verify via state_dict() and direct access - if (!tensorsAllClose(m1.state_dict().at("weight"), m1_load_weight, 1e-5, 1e-5)) { - spdlog::error("m1 weight not loaded correctly"); - return false; - } - if (!tensorsAllClose(m1.weight(), m1_load_weight, 1e-5, 1e-5)) { - spdlog::error("m1 weight field not synchronized"); - return false; - } - if (!tensorsAllClose(m1.bias(), m1_load_bias, 1e-5, 1e-5)) { - spdlog::error("m1 bias field not synchronized"); - return false; - } - spdlog::debug("m1 load_state_dict verified - parameters and fields synchronized"); - - // Test extra_repr - std::string repr = m1.extra_repr(); - spdlog::debug("Linear module representation: {}", repr); - - // Test forward with residual connection - spdlog::info("Testing Linear forward with residual connection"); - // auto residual = infinicore::Tensor::ones({2, 4}, infinicore::DataType::F32, infinicore::Device()); - auto output_with_residual = m1.forward(input1); - if (output_with_residual->shape() != std::vector({2, 4})) { - spdlog::error("Linear output with residual shape mismatch. Expected {{2, 4}}, got different shape"); - return false; - } - spdlog::debug("Linear forward with residual passed. Input shape: {{2, 8}}, Residual shape: {{2, 4}}, Output shape: {{2, 4}}"); - - // Test computation correctness: InfiniCore vs Naive implementation - spdlog::info("Testing computation correctness: InfiniCore vs Naive implementation"); - - // Create test data with known values for verification - auto test_input = infinicore::Tensor::ones({2, 8}, infinicore::DataType::F32, infinicore::Device()); - // auto test_residual = infinicore::Tensor::ones({2, 4}, infinicore::DataType::F32, infinicore::Device()); - - // Get InfiniCore result - auto infinicore_output = m1.forward(test_input); - - // Compute naive result: output = input @ weight.T + bias + residual - auto naive_output = infinicore::Tensor::empty({2, 4}, infinicore::DataType::F32, infinicore::Device()); - auto weight_naive = m1.weight(); - auto bias_naive = m1.bias(); - - // Naive computation step by step - auto weight_t = weight_naive->permute({1, 0}); // [4, 8] -> [8, 4] - auto matmul_result = infinicore::op::matmul(test_input, weight_t); // [2, 4] - - // Broadcast bias to [2, 4] - size_t ndim_diff = naive_output->ndim() - 1; - std::vector strides(ndim_diff, 0); - strides.push_back(bias_naive->stride(0)); - auto bias_view = bias_naive->as_strided(naive_output->shape(), strides); - - // Add bias to matmul result - infinicore::op::add_(naive_output, matmul_result, bias_view); - - // Add residual - // infinicore::op::add_(naive_output, naive_output, test_residual); - - // Compare results with actual value checking - if (infinicore_output->shape() != naive_output->shape()) { - spdlog::error("Shape mismatch between InfiniCore and naive implementation"); - return false; - } - - // Compare actual tensor values using local checker - if (!tensorsAllClose(infinicore_output, naive_output, 1e-5, 1e-5)) { - spdlog::error("Value mismatch between InfiniCore and naive implementation"); - return false; - } - spdlog::debug("Value comparison passed - InfiniCore and naive results match within tolerance"); - - spdlog::debug("Computation correctness test passed - both implementations produce identical results"); - spdlog::debug("InfiniCore output shape: {{2, 4}}, Naive output shape: {{2, 4}}"); - - // Test computation correctness without bias (using m2) - spdlog::info("Testing computation correctness without bias"); - auto test_input_no_bias = infinicore::Tensor::ones({1, 16}, infinicore::DataType::F32, infinicore::Device()); - // auto test_residual_no_bias = infinicore::Tensor::ones({1, 3}, infinicore::DataType::F32, infinicore::Device()); - - // Get InfiniCore result (no bias) - auto infinicore_output_no_bias = m2.forward(test_input_no_bias); - - // Compute naive result without bias: output = input @ weight.T + residual - auto naive_output_no_bias = infinicore::Tensor::empty({1, 3}, infinicore::DataType::F32, infinicore::Device()); - auto weight_no_bias_naive = m2.weight(); - - // Naive computation: just matmul + residual - auto weight_t_no_bias = weight_no_bias_naive->permute({1, 0}); // [3, 16] -> [16, 3] - auto matmul_result_no_bias = infinicore::op::matmul(test_input_no_bias, weight_t_no_bias); // [1, 3] - - // Add residual - // infinicore::op::add_(naive_output_no_bias, matmul_result_no_bias, test_residual_no_bias); - - // Compare results with actual value checking - if (infinicore_output_no_bias->shape() != naive_output_no_bias->shape()) { - spdlog::error("Shape mismatch between InfiniCore and naive implementation (no bias)"); - return false; - } - - // Compare actual tensor values for no-bias case - if (!tensorsAllClose(infinicore_output_no_bias, naive_output_no_bias, 1e-5, 1e-5)) { - spdlog::error("Value mismatch in no-bias computation"); - return false; - } - spdlog::debug("No-bias value comparison passed - results match within tolerance"); - - spdlog::debug("No-bias computation correctness test passed - both implementations produce identical results"); - spdlog::debug("InfiniCore no-bias output shape: {{1, 3}}, Naive no-bias output shape: {{1, 3}}"); - - // Test basic forward (no residual) vs naive - spdlog::info("Testing basic forward vs naive implementation"); - auto basic_infinicore = m1.forward(test_input); - auto basic_naive = infinicore::Tensor::empty({2, 4}, infinicore::DataType::F32, infinicore::Device()); - - // Naive basic computation: input @ weight.T + bias - auto basic_matmul = infinicore::op::matmul(test_input, weight_t); - infinicore::op::add_(basic_naive, basic_matmul, bias_view); - - if (basic_infinicore->shape() != basic_naive->shape()) { - spdlog::error("Shape mismatch in basic forward computation"); - return false; - } - - // Compare actual tensor values for basic forward - if (!tensorsAllClose(basic_infinicore, basic_naive, 1e-5, 1e-5)) { - spdlog::error("Value mismatch in basic forward computation"); - return false; - } - spdlog::debug("Basic forward value comparison passed - results match within tolerance"); - - spdlog::debug("Basic forward computation correctness test passed - both implementations produce identical results"); - spdlog::debug("Basic InfiniCore output shape: {{2, 4}}, Basic naive output shape: {{2, 4}}"); - - spdlog::info("All Linear module tests passed (with/without bias, load_state_dict, computation verification)"); - return true; - } catch (const std::exception &e) { - spdlog::error("Exception in testModuleLinear: {}", e.what()); - return false; - } - }); -} - -// Test 6: Embedding module implementation -TestResult NNModuleTest::testModuleEmbedding() { - return measureTime("ModuleEmbedding", [this]() { - try { - spdlog::info("=========================================="); - spdlog::info("Testing Embedding module implementation"); - spdlog::info("=========================================="); - - // Test 1: Basic embedding creation - spdlog::info("Test 1: Basic embedding creation (vocab=100, dim=64)"); - infinicore::nn::Embedding emb1(100, 64); - - auto state1 = emb1.state_dict(); - if (state1.find("weight") == state1.end()) { - spdlog::error("Embedding weight not found in state dict"); - return false; - } - - if (state1.at("weight")->shape() != std::vector({100, 64})) { - spdlog::error("Embedding weight shape mismatch. Expected {{100, 64}}"); - return false; - } - - if (emb1.num_embeddings() != 100) { - spdlog::error("num_embeddings mismatch. Expected 100, got {}", emb1.num_embeddings()); - return false; - } - - if (emb1.embedding_dim() != 64) { - spdlog::error("embedding_dim mismatch. Expected 64, got {}", emb1.embedding_dim()); - return false; - } - - spdlog::debug("Basic embedding creation passed"); - - // Test 2: Embedding with padding_idx - spdlog::info("Test 2: Embedding with padding_idx=0"); - infinicore::nn::Embedding emb2(50, 32, 0, infinicore::DataType::F32, infinicore::Device()); - - if (!emb2.padding_idx().has_value()) { - spdlog::error("padding_idx should have a value"); - return false; - } - - if (emb2.padding_idx().value() != 0) { - spdlog::error("padding_idx mismatch. Expected 0, got {}", emb2.padding_idx().value()); - return false; - } - - spdlog::debug("Embedding with padding_idx passed"); - - // Test 3: Forward pass - single index - spdlog::info("Test 3: Forward pass with single index"); - std::vector single_data = {5}; - auto indices_single = infinicore::Tensor::from_blob(single_data.data(), {1}, infinicore::DataType::I64, infinicore::Device()); - auto output_single = emb1.forward(indices_single); - - if (output_single->shape() != std::vector({1, 64})) { - spdlog::error("Single index output shape mismatch. Expected {{1, 64}}"); - return false; - } - - spdlog::debug("Single index forward pass passed. Output shape: {{1, 64}}"); - - // Test 4: Forward pass - batch of indices - spdlog::info("Test 4: Forward pass with batch of indices"); - std::vector batch_data = {0, 5, 10}; - auto indices_batch = infinicore::Tensor::from_blob(batch_data.data(), {3}, infinicore::DataType::I64, infinicore::Device()); - auto output_batch = emb1.forward(indices_batch); - - if (output_batch->shape() != std::vector({3, 64})) { - spdlog::error("Batch output shape mismatch. Expected {{3, 64}}"); - return false; - } - - spdlog::debug("Batch forward pass passed. Output shape: {{3, 64}}"); - - // Test 5: Forward pass - 2D indices (batch_size, seq_len) - spdlog::info("Test 5: Forward pass with 2D indices [batch, seq_len]"); - std::vector data_2d = {1, 2, 3, 4, 5, 6, 7, 8}; - auto indices_2d = infinicore::Tensor::from_blob(data_2d.data(), {2, 4}, - infinicore::DataType::I64, infinicore::Device()); - auto output_2d = emb1.forward(indices_2d); - - if (output_2d->shape() != std::vector({2, 4, 64})) { - spdlog::error("2D indices output shape mismatch. Expected {{2, 4, 64}}"); - return false; - } - - spdlog::debug("2D indices forward pass passed. Output shape: {{2, 4, 64}}"); - - // Test 6: Embedding lookup consistency - spdlog::info("Test 6: Testing embedding lookup consistency"); - std::vector idx_data = {7}; - auto idx1 = infinicore::Tensor::from_blob(idx_data.data(), {1}, infinicore::DataType::I64, infinicore::Device()); - auto idx2 = infinicore::Tensor::from_blob(idx_data.data(), {1}, infinicore::DataType::I64, infinicore::Device()); - - auto out1 = emb1.forward(idx1); - auto out2 = emb1.forward(idx2); - - // Same index should give same embedding - if (!tensorsAllClose(out1, out2, 1e-7, 1e-7)) { - spdlog::error("Same index should return identical embeddings"); - return false; - } - - spdlog::debug("Embedding lookup consistency passed"); - - // Test 7: load_state_dict - spdlog::info("Test 7: Testing load_state_dict for Embedding"); - auto new_weight = infinicore::Tensor::ones({100, 64}, infinicore::DataType::F32, infinicore::Device()); - std::unordered_map new_state; - new_state.emplace("weight", new_weight); - - emb1.load_state_dict(new_state); - - if (!tensorsAllClose(emb1.weight(), new_weight, 1e-7, 1e-7)) { - spdlog::error("Embedding weight not loaded correctly"); - return false; - } - - spdlog::debug("load_state_dict for Embedding passed"); - - // Test 8: extra_repr - spdlog::info("Test 8: Testing extra_repr"); - std::string repr1 = emb1.extra_repr(); - std::string repr2 = emb2.extra_repr(); - - spdlog::debug("Embedding repr (no padding): {}", repr1); - spdlog::debug("Embedding repr (with padding): {}", repr2); - - if (repr1.find("num_embeddings=100") == std::string::npos) { - spdlog::error("extra_repr should contain num_embeddings"); - return false; - } - - if (repr2.find("padding_idx=0") == std::string::npos) { - spdlog::error("extra_repr should contain padding_idx when specified"); - return false; - } - - spdlog::debug("extra_repr test passed"); - - spdlog::info("All Embedding module tests passed!"); - return true; - - } catch (const std::exception &e) { - spdlog::error("Exception in testModuleEmbedding: {}", e.what()); - return false; - } - }); -} - -// Test 7: RMSNorm module implementation -TestResult NNModuleTest::testModuleRMSNorm() { - return measureTime("ModuleRMSNorm", [this]() { - try { - spdlog::info("=========================================="); - spdlog::info("Testing RMSNorm module implementation"); - spdlog::info("=========================================="); - - // Test 1: Basic RMSNorm creation - spdlog::info("Test 1: Basic RMSNorm creation (hidden_size=768)"); - infinicore::nn::RMSNorm norm1(768); - - auto state1 = norm1.state_dict(); - if (state1.find("weight") == state1.end()) { - spdlog::error("RMSNorm weight not found in state dict"); - return false; - } - - if (state1.at("weight")->shape() != std::vector({768})) { - spdlog::error("RMSNorm weight shape mismatch. Expected {{768}}"); - return false; - } - - if (norm1.normalized_shape() != 768) { - spdlog::error("normalized_shape mismatch. Expected 768, got {}", norm1.normalized_shape()); - return false; - } - - spdlog::debug("Basic RMSNorm creation passed"); - - // Test 2: Forward pass - 2D input [batch, hidden] - spdlog::info("Test 2: Forward pass with 2D input [batch, hidden]"); - auto input_2d = infinicore::Tensor::ones({4, 768}, infinicore::DataType::F32, infinicore::Device()); - auto output_2d = norm1.forward(input_2d); - - if (output_2d->shape() != std::vector({4, 768})) { - spdlog::error("2D output shape mismatch. Expected {{4, 768}}"); - return false; - } - - spdlog::debug("2D forward pass passed. Output shape: {{4, 768}}"); - - // Test 3: Forward pass - 3D input [batch, seq_len, hidden] - spdlog::info("Test 3: Forward pass with 3D input [batch, seq_len, hidden]"); - auto input_3d = infinicore::Tensor::ones({2, 10, 768}, infinicore::DataType::F32, infinicore::Device()); - auto output_3d = norm1.forward(input_3d); - - if (output_3d->shape() != std::vector({2, 10, 768})) { - spdlog::error("3D output shape mismatch. Expected {{2, 10, 768}}"); - return false; - } - - spdlog::debug("3D forward pass passed. Output shape: {{2, 10, 768}}"); - - // Test 4: Test normalization properties - spdlog::info("Test 4: Testing RMSNorm properties"); - auto test_input = infinicore::Tensor::ones({1, 768}, infinicore::DataType::F32, infinicore::Device()); - auto test_output = norm1.forward(test_input); - - // Output should have same shape - if (test_output->shape() != test_input->shape()) { - spdlog::error("Output shape doesn't match input shape"); - return false; - } - - spdlog::debug("RMSNorm properties test passed"); - - // Test 5: load_state_dict - spdlog::info("Test 5: Testing load_state_dict for RMSNorm"); - auto new_weight = infinicore::Tensor::ones({768}, infinicore::DataType::F32, infinicore::Device()); - std::unordered_map new_state; - new_state.emplace("weight", new_weight); - - norm1.load_state_dict(new_state); - - if (!tensorsAllClose(norm1.weight(), new_weight, 1e-7, 1e-7)) { - spdlog::error("RMSNorm weight not loaded correctly"); - return false; - } - - spdlog::debug("load_state_dict for RMSNorm passed"); - - // Test 6: extra_repr - spdlog::info("Test 6: Testing extra_repr"); - std::string repr = norm1.extra_repr(); - spdlog::debug("RMSNorm repr: {}", repr); - - if (repr.find("normalized_shape=768") == std::string::npos) { - spdlog::error("extra_repr should contain normalized_shape"); - return false; - } - - if (repr.find("eps=") == std::string::npos) { - spdlog::error("extra_repr should contain eps"); - return false; - } - - spdlog::debug("extra_repr test passed"); - - // Test 7: Input validation - normalized_shape mismatch (op layer handles this) - spdlog::info("Test 7: Testing input validation - normalized_shape mismatch"); - auto input_wrong_shape = infinicore::Tensor::ones({4, 512}, infinicore::DataType::F32, infinicore::Device()); // normalized_shape=512, expected 768 - - try { - norm1.forward(input_wrong_shape); - spdlog::error("Should have thrown exception for normalized_shape mismatch"); - return false; - } catch (const std::exception &e) { - spdlog::debug("Correctly caught exception for normalized_shape mismatch (handled by op layer): {}", e.what()); - } catch (...) { - spdlog::error("Caught unexpected exception type"); - return false; - } - - spdlog::debug("Normalized_shape mismatch validation test passed"); - - // Test 8: Different hidden sizes - spdlog::info("Test 8: Testing different hidden sizes"); - infinicore::nn::RMSNorm norm_small(128, 1e-5); - infinicore::nn::RMSNorm norm_large(4096); - - auto input_small = infinicore::Tensor::ones({2, 128}, infinicore::DataType::F32, infinicore::Device()); - auto output_small = norm_small.forward(input_small); - - auto input_large = infinicore::Tensor::ones({2, 4096}, infinicore::DataType::F32, infinicore::Device()); - auto output_large = norm_large.forward(input_large); - - if (output_small->shape() != std::vector({2, 128})) { - spdlog::error("Small RMSNorm output shape mismatch"); - return false; - } - - if (output_large->shape() != std::vector({2, 4096})) { - spdlog::error("Large RMSNorm output shape mismatch"); - return false; - } - - spdlog::debug("Different hidden sizes test passed"); - - spdlog::info("All RMSNorm module tests passed!"); - return true; - - } catch (const std::exception &e) { - spdlog::error("Exception in testModuleRMSNorm: {}", e.what()); - return false; - } - }); -} - -// Test 7.5: RoPE module test -TestResult NNModuleTest::testModuleRoPE() { - return measureTime("ModuleRoPE", [this]() { - try { - spdlog::info("=========================================="); - spdlog::info("Testing RoPE module implementation"); - spdlog::info("=========================================="); - - // Test 1: Basic RoPE creation - spdlog::info("Test 1: Basic RoPE creation (head_dim=128, max_seq_len=2048)"); - infinicore::nn::RoPE rope1(128, 2048); - - auto state1 = rope1.state_dict(); - - if (rope1.head_dim() != 128) { - spdlog::error("head_dim mismatch. Expected 128, got {}", rope1.head_dim()); - return false; - } - - if (rope1.max_seq_len() != 2048) { - spdlog::error("max_seq_len mismatch. Expected 2048, got {}", rope1.max_seq_len()); - return false; - } - - spdlog::debug("Basic RoPE creation passed"); - - // Test 2: Forward pass - 3D input [seq_len, n_head, head_dim] - spdlog::info("Test 2: Forward pass with 3D input [seq_len, n_head, head_dim]"); - auto x_3d = infinicore::Tensor::ones({32, 32, 128}, infinicore::DataType::F32, infinicore::Device()); - - // Create position tensor [0, 1, 2, ..., 31] - std::vector pos_data(32); - for (size_t i = 0; i < 32; i++) { - pos_data[i] = static_cast(i); - } - auto pos = infinicore::Tensor::from_blob(pos_data.data(), {32}, infinicore::DataType::I32, infinicore::Device()); - - auto x_out = rope1.forward(x_3d, pos); - - if (x_out->shape() != std::vector({32, 32, 128})) { - spdlog::error("3D output shape mismatch. Expected {{32, 32, 128}}"); - return false; - } - - spdlog::debug("3D forward pass passed. Output shape: {{32, 32, 128}}"); - - // Test 3: Different algorithms - spdlog::info("Test 3: Testing different algorithms"); - infinicore::nn::RoPE rope_gptj(64, 1024, 10000.0, infinicore::nn::RoPE::Algo::GPT_J); - infinicore::nn::RoPE rope_gptneox(64, 1024, 10000.0, infinicore::nn::RoPE::Algo::GPT_NEOX); - - if (rope_gptj.algo() != infinicore::nn::RoPE::Algo::GPT_J) { - spdlog::error("GPT_J algorithm not set correctly"); - return false; - } - - if (rope_gptneox.algo() != infinicore::nn::RoPE::Algo::GPT_NEOX) { - spdlog::error("GPT_NEOX algorithm not set correctly"); - return false; - } - - auto x_test = infinicore::Tensor::ones({10, 32, 64}, infinicore::DataType::F32, infinicore::Device()); - - std::vector pos_test_data(10); - for (size_t i = 0; i < 10; i++) { - pos_test_data[i] = static_cast(i); - } - auto pos_test = infinicore::Tensor::from_blob(pos_test_data.data(), {10}, infinicore::DataType::I32, infinicore::Device()); - - auto x_gptj = rope_gptj.forward(x_test, pos_test); - auto x_neox = rope_gptneox.forward(x_test, pos_test); - - if (x_gptj->shape() != x_test->shape()) { - spdlog::error("GPT_J forward pass shape mismatch"); - return false; - } - - if (x_neox->shape() != x_test->shape()) { - spdlog::error("GPT_NEOX forward pass shape mismatch"); - return false; - } - - spdlog::debug("Different algorithms test passed"); - - // Test 4: Different theta values - spdlog::info("Test 4: Testing different theta values"); - infinicore::nn::RoPE rope_theta1(128, 2048, 1e5); - infinicore::nn::RoPE rope_theta2(128, 2048, 1e4); - - if (rope_theta1.theta() != 1e5) { - spdlog::error("theta mismatch. Expected 1e5, got {}", rope_theta1.theta()); - return false; - } - - if (rope_theta2.theta() != 1e4) { - spdlog::error("theta mismatch. Expected 1e4, got {}", rope_theta2.theta()); - return false; - } - - spdlog::debug("Different theta values test passed"); - - // Test 5: load_state_dict - std::unordered_map new_state; - rope1.load_state_dict(new_state); - spdlog::debug("load_state_dict for RoPE passed (no parameters to load)"); - - // Test 6: extra_repr - spdlog::info("Test 6: Testing extra_repr"); - std::string repr = rope1.extra_repr(); - spdlog::debug("RoPE repr: {}", repr); - - if (repr.find("head_dim=128") == std::string::npos) { - spdlog::error("extra_repr should contain head_dim"); - return false; - } - - if (repr.find("max_seq_len=2048") == std::string::npos) { - spdlog::error("extra_repr should contain max_seq_len"); - return false; - } - - if (repr.find("theta=") == std::string::npos) { - spdlog::error("extra_repr should contain theta"); - return false; - } - - spdlog::debug("extra_repr test passed"); - - // Test 7: Different head dimensions - spdlog::info("Test 7: Testing different head dimensions"); - infinicore::nn::RoPE rope_small(64, 1024); - infinicore::nn::RoPE rope_large(256, 4096); - - auto x_small = infinicore::Tensor::ones({10, 32, 64}, infinicore::DataType::F32, infinicore::Device()); - - std::vector pos_small_data(10); - for (size_t i = 0; i < 10; i++) { - pos_small_data[i] = static_cast(i); - } - auto pos_small = infinicore::Tensor::from_blob(pos_small_data.data(), {10}, infinicore::DataType::I32, infinicore::Device()); - - auto x_small_out = rope_small.forward(x_small, pos_small); - - if (x_small_out->shape() != std::vector({10, 32, 64})) { - spdlog::error("Small RoPE output shape mismatch"); - return false; - } - - auto x_large = infinicore::Tensor::ones({20, 32, 256}, infinicore::DataType::F32, infinicore::Device()); - - std::vector pos_large_data(20); - for (size_t i = 0; i < 20; i++) { - pos_large_data[i] = static_cast(i); - } - auto pos_large = infinicore::Tensor::from_blob(pos_large_data.data(), {20}, infinicore::DataType::I32, infinicore::Device()); - - auto x_large_out = rope_large.forward(x_large, pos_large); - - if (x_large_out->shape() != std::vector({20, 32, 256})) { - spdlog::error("Large RoPE output shape mismatch"); - return false; - } - - spdlog::debug("Different head dimensions test passed"); - - // Test 8: Invalid head_dim (odd number) - spdlog::info("Test 8: Testing invalid head_dim (odd number)"); - try { - infinicore::nn::RoPE rope_invalid(127, 2048); - spdlog::error("Should have thrown exception for odd head_dim"); - return false; - } catch (const std::invalid_argument &e) { - spdlog::debug("Correctly caught exception for odd head_dim: {}", e.what()); - } catch (...) { - spdlog::error("Caught unexpected exception type"); - return false; - } - - spdlog::debug("Invalid head_dim test passed"); - - // Test 9: Input validation - empty tensor (op layer handles this) - spdlog::info("Test 9: Testing input validation - empty tensor"); - auto x_empty = infinicore::Tensor::ones({}, infinicore::DataType::F32, infinicore::Device()); - std::vector pos_empty_data(1); - pos_empty_data[0] = 0; - auto pos_empty = infinicore::Tensor::from_blob(pos_empty_data.data(), {1}, infinicore::DataType::I32, infinicore::Device()); - - try { - rope1.forward(x_empty, pos_empty); - spdlog::error("Should have thrown exception for empty input tensor"); - return false; - } catch (const std::exception &e) { - spdlog::debug("Correctly caught exception for empty input (handled by op layer): {}", e.what()); - } catch (...) { - spdlog::error("Caught unexpected exception type"); - return false; - } - - spdlog::debug("Empty tensor validation test passed"); - - // Test 10: Input validation - head_dim mismatch (op layer handles this) - spdlog::info("Test 10: Testing input validation - head_dim mismatch"); - auto x_wrong_dim = infinicore::Tensor::ones({32, 32, 64}, infinicore::DataType::F32, infinicore::Device()); // head_dim=64, expected 128 - std::vector pos_wrong_data(32); - for (size_t i = 0; i < 32; i++) { - pos_wrong_data[i] = static_cast(i); - } - auto pos_wrong = infinicore::Tensor::from_blob(pos_wrong_data.data(), {32}, infinicore::DataType::I32, infinicore::Device()); - - try { - rope1.forward(x_wrong_dim, pos_wrong); - spdlog::error("Should have thrown exception for head_dim mismatch"); - return false; - } catch (const std::exception &e) { - spdlog::debug("Correctly caught exception for head_dim mismatch (handled by op layer): {}", e.what()); - } catch (...) { - spdlog::error("Caught unexpected exception type"); - return false; - } - - spdlog::debug("Head_dim mismatch validation test passed"); - - // Test 11: Different input shapes (from reference test cases) - spdlog::info("Test 11: Testing different input shapes"); - - // Test shape (1, 32, 128) - single sequence - auto x_single = infinicore::Tensor::ones({1, 32, 128}, infinicore::DataType::F32, infinicore::Device()); - std::vector pos_single_data(1); - pos_single_data[0] = 0; - auto pos_single = infinicore::Tensor::from_blob(pos_single_data.data(), {1}, infinicore::DataType::I32, infinicore::Device()); - auto x_single_out = rope1.forward(x_single, pos_single); - if (x_single_out->shape() != std::vector({1, 32, 128})) { - spdlog::error("Single sequence output shape mismatch"); - return false; - } - - // Test shape (10, 32, 64) - different head_dim - auto rope_64 = infinicore::nn::RoPE(64, 1024); - auto x_64 = infinicore::Tensor::ones({10, 32, 64}, infinicore::DataType::F32, infinicore::Device()); - std::vector pos_64_data(10); - for (size_t i = 0; i < 10; i++) { - pos_64_data[i] = static_cast(i); - } - auto pos_64 = infinicore::Tensor::from_blob(pos_64_data.data(), {10}, infinicore::DataType::I32, infinicore::Device()); - auto x_64_out = rope_64.forward(x_64, pos_64); - if (x_64_out->shape() != std::vector({10, 32, 64})) { - spdlog::error("Shape (10, 32, 64) output mismatch"); - return false; - } - - // Test shape (4, 1, 32) - single head - auto rope_32 = infinicore::nn::RoPE(32, 1024); - auto x_1head = infinicore::Tensor::ones({4, 1, 32}, infinicore::DataType::F32, infinicore::Device()); - std::vector pos_1head_data(4); - for (size_t i = 0; i < 4; i++) { - pos_1head_data[i] = static_cast(i); - } - auto pos_1head = infinicore::Tensor::from_blob(pos_1head_data.data(), {4}, infinicore::DataType::I32, infinicore::Device()); - auto x_1head_out = rope_32.forward(x_1head, pos_1head); - if (x_1head_out->shape() != std::vector({4, 1, 32})) { - spdlog::error("Shape (4, 1, 32) output mismatch"); - return false; - } - - spdlog::debug("Different input shapes test passed"); - - // Test 12: Position tensor validation - spdlog::info("Test 12: Testing position tensor edge cases"); - - // Test with seq_len less than max_seq_len - auto x_short = infinicore::Tensor::ones({10, 32, 128}, infinicore::DataType::F32, infinicore::Device()); - std::vector pos_short_data(10); - for (size_t i = 0; i < 10; i++) { - pos_short_data[i] = static_cast(i); - } - auto pos_short = infinicore::Tensor::from_blob(pos_short_data.data(), {10}, infinicore::DataType::I32, infinicore::Device()); - auto x_short_out = rope1.forward(x_short, pos_short); - if (x_short_out->shape() != std::vector({10, 32, 128})) { - spdlog::error("Short sequence output shape mismatch"); - return false; - } - - spdlog::debug("Position tensor edge cases test passed"); - - // Test 13: Test that outputs are on the same device as inputs - spdlog::info("Test 13: Testing device consistency"); - auto device = x_3d->device(); - if (x_out->device() != device) { - spdlog::error("Output tensor not on the same device as input"); - return false; - } - spdlog::debug("Device consistency test passed"); - - spdlog::info("All RoPE module tests passed!"); - return true; - - } catch (const std::exception &e) { - spdlog::error("Exception in testModuleRoPE: {}", e.what()); - return false; - } - }); -} - -// Test 8: Dtype assertion test -TestResult NNModuleTest::testDtypeAssertion() { - return measureTime("DtypeAssertionTest", [this]() { - try { - spdlog::info("=========================================="); - spdlog::info("Testing dtype assertions when loading parameters"); - spdlog::info("=========================================="); - - // Test 1: Successful load with matching dtype - spdlog::info("Test 1: Successful load with matching dtype (F32)"); - infinicore::nn::Linear linear1(8, 4, true); - auto matching_weight = infinicore::Tensor::ones({4, 8}, infinicore::DataType::F32, infinicore::Device()); - auto matching_bias = infinicore::Tensor::ones({4}, infinicore::DataType::F32, infinicore::Device()); - - std::unordered_map matching_state; - matching_state.emplace("weight", matching_weight); - matching_state.emplace("bias", matching_bias); - - // This should succeed without throwing - linear1.load_state_dict(matching_state); - spdlog::debug("✓ Matching dtype load succeeded"); - - // Test 2: Failed load with mismatched dtype (load_parameter_) - spdlog::info("Test 2: Failed load_parameter_ with mismatched dtype"); - infinicore::nn::Linear linear2(8, 4, true); - auto mismatched_weight = infinicore::Tensor::ones({4, 8}, infinicore::DataType::BF16, infinicore::Device()); - - bool exception_thrown = false; - try { - linear2.load_parameter_("weight", mismatched_weight); - } catch (const std::runtime_error &e) { - exception_thrown = true; - std::string error_msg = e.what(); - if (error_msg.find("dtype mismatch") == std::string::npos) { - spdlog::error("Exception message doesn't contain 'dtype mismatch'"); - return false; - } - spdlog::debug("✓ Mismatched dtype exception caught: {}", error_msg); - } - - if (!exception_thrown) { - spdlog::error("Expected exception for dtype mismatch in load_parameter_"); - return false; - } - - // Test 3: Failed load with mismatched dtype (load_state_dict) - spdlog::info("Test 3: Failed load_state_dict with mismatched dtype"); - infinicore::nn::Embedding embedding1(100, 64); - auto mismatched_embed_weight = infinicore::Tensor::ones({100, 64}, infinicore::DataType::BF16, infinicore::Device()); - - std::unordered_map mismatched_state; - mismatched_state.emplace("weight", mismatched_embed_weight); - - exception_thrown = false; - try { - embedding1.load_state_dict(mismatched_state); - } catch (const std::runtime_error &e) { - exception_thrown = true; - std::string error_msg = e.what(); - if (error_msg.find("dtype mismatch") == std::string::npos) { - spdlog::error("Exception message doesn't contain 'dtype mismatch'"); - return false; - } - if (error_msg.find("weight") == std::string::npos) { - spdlog::error("Exception message doesn't contain parameter name 'weight'"); - return false; - } - spdlog::debug("✓ Mismatched dtype exception caught: {}", error_msg); - } - - if (!exception_thrown) { - spdlog::error("Expected exception for dtype mismatch in load_state_dict"); - return false; - } - - // Test 4: Failed load with mismatched dtype (RMSNorm) - spdlog::info("Test 4: Failed load_state_dict with mismatched dtype (RMSNorm)"); - infinicore::nn::RMSNorm norm1(768); - auto mismatched_norm_weight = infinicore::Tensor::ones({768}, infinicore::DataType::BF16, infinicore::Device()); - - std::unordered_map mismatched_norm_state; - mismatched_norm_state.emplace("weight", mismatched_norm_weight); - - exception_thrown = false; - try { - norm1.load_state_dict(mismatched_norm_state); - } catch (const std::runtime_error &e) { - exception_thrown = true; - std::string error_msg = e.what(); - if (error_msg.find("dtype mismatch") == std::string::npos) { - spdlog::error("Exception message doesn't contain 'dtype mismatch'"); - return false; - } - spdlog::debug("✓ Mismatched dtype exception caught for RMSNorm: {}", error_msg); - } - - if (!exception_thrown) { - spdlog::error("Expected exception for dtype mismatch in RMSNorm load_state_dict"); - return false; - } - - // Test 5: Successful load with different module dtypes - spdlog::info("Test 5: Successful load with BF16 dtype (module created with BF16)"); - infinicore::nn::Linear linear3(8, 4, true, infinicore::DataType::BF16); - auto bf16_weight = infinicore::Tensor::ones({4, 8}, infinicore::DataType::BF16, infinicore::Device()); - auto bf16_bias = infinicore::Tensor::ones({4}, infinicore::DataType::BF16, infinicore::Device()); - - std::unordered_map bf16_state; - bf16_state.emplace("weight", bf16_weight); - bf16_state.emplace("bias", bf16_bias); - - // This should succeed - linear3.load_state_dict(bf16_state); - spdlog::debug("✓ BF16 dtype load succeeded"); - - spdlog::info("All dtype assertion tests passed!"); - return true; - - } catch (const std::exception &e) { - spdlog::error("Exception in testDtypeAssertion: {}", e.what()); - return false; - } - }); -} - -// Test 9: Comprehensive Tiny-Llama model test (construction + weight loading + validation) -TestResult NNModuleTest::testTinyLlamaConstruction() { - return measureTime("TinyLlamaModelTest", [this]() { - try { - spdlog::info("=========================================="); - spdlog::info("Testing Tiny-Llama Model Construction and Weight Loading"); - spdlog::info("=========================================="); - - // Tiny-Llama configuration (actual Tiny-Llama-1.1B-Chat-v1.0 specs) - struct TinyLlamaConfig { - size_t vocab_size = 32000; - size_t hidden_size = 2048; - size_t intermediate_size = 5632; - size_t num_hidden_layers = 22; - size_t num_attention_heads = 32; - size_t num_key_value_heads = 4; // GQA (Grouped Query Attention) - size_t max_position_embeddings = 2048; - double rms_norm_eps = 1e-5; - }; - - TinyLlamaConfig config; - - // ============================================ - // Phase 0: Use hard-coded TinyLlama configuration (CI-friendly) - // ============================================ - spdlog::info(""); - spdlog::info("Phase 0: Using hard-coded TinyLlama configuration (CI)"); - spdlog::info("------------------------------------------"); - - spdlog::info("Using Configuration:"); - spdlog::info(" vocab_size: {}", config.vocab_size); - spdlog::info(" hidden_size: {}", config.hidden_size); - spdlog::info(" intermediate_size: {}", config.intermediate_size); - spdlog::info(" num_layers: {}", config.num_hidden_layers); - spdlog::info(" num_attention_heads: {}", config.num_attention_heads); - spdlog::info(" num_key_value_heads: {} (GQA)", config.num_key_value_heads); - spdlog::info(" max_position_embeddings: {}", config.max_position_embeddings); - spdlog::info(" rms_norm_eps: {}", config.rms_norm_eps); - - // Create Tiny-Llama model skeleton closely matching HF/TinyLlama naming - class TinyLlamaModel : public infinicore::nn::Module { - protected: - // Inner modules to match naming like: layers.0.self_attn.q_proj.weight, layers.0.mlp.gate_proj.weight - class SelfAttn : public infinicore::nn::Module { - public: - INFINICORE_NN_MODULE(infinicore::nn::Linear, q_proj); - INFINICORE_NN_MODULE(infinicore::nn::Linear, k_proj); - INFINICORE_NN_MODULE(infinicore::nn::Linear, v_proj); - INFINICORE_NN_MODULE(infinicore::nn::Linear, o_proj); - - SelfAttn(size_t hidden_size, size_t kv_dim, const infinicore::Device &device) { - INFINICORE_NN_MODULE_INIT(q_proj, hidden_size, hidden_size, false, infinicore::DataType::F32, device); - INFINICORE_NN_MODULE_INIT(k_proj, hidden_size, kv_dim, false, infinicore::DataType::F32, device); - INFINICORE_NN_MODULE_INIT(v_proj, hidden_size, kv_dim, false, infinicore::DataType::F32, device); - INFINICORE_NN_MODULE_INIT(o_proj, hidden_size, hidden_size, false, infinicore::DataType::F32, device); - } - }; - - class MLP : public infinicore::nn::Module { - public: - INFINICORE_NN_MODULE(infinicore::nn::Linear, gate_proj); - INFINICORE_NN_MODULE(infinicore::nn::Linear, up_proj); - INFINICORE_NN_MODULE(infinicore::nn::Linear, down_proj); - - MLP(size_t hidden_size, size_t intermediate_size, const infinicore::Device &device) { - INFINICORE_NN_MODULE_INIT(gate_proj, hidden_size, intermediate_size, false, infinicore::DataType::F32, device); - INFINICORE_NN_MODULE_INIT(up_proj, hidden_size, intermediate_size, false, infinicore::DataType::F32, device); - INFINICORE_NN_MODULE_INIT(down_proj, intermediate_size, hidden_size, false, infinicore::DataType::F32, device); - } - }; - - class Block : public infinicore::nn::Module { - public: - INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, input_layernorm); - INFINICORE_NN_MODULE(SelfAttn, self_attn); - INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, post_attention_layernorm); - INFINICORE_NN_MODULE(MLP, mlp); - - Block(const TinyLlamaConfig &cfg, const infinicore::Device &device) { - size_t kv_dim = cfg.hidden_size * cfg.num_key_value_heads / cfg.num_attention_heads; - INFINICORE_NN_MODULE_INIT(input_layernorm, cfg.hidden_size, cfg.rms_norm_eps, infinicore::DataType::F32, device); - INFINICORE_NN_MODULE_INIT(self_attn, cfg.hidden_size, kv_dim, device); - INFINICORE_NN_MODULE_INIT(post_attention_layernorm, cfg.hidden_size, cfg.rms_norm_eps, infinicore::DataType::F32, device); - INFINICORE_NN_MODULE_INIT(mlp, cfg.hidden_size, cfg.intermediate_size, device); - } - }; - - public: - INFINICORE_NN_MODULE(infinicore::nn::Embedding, embed_tokens); - INFINICORE_NN_MODULE_VEC(Block, layers); - INFINICORE_NN_MODULE(infinicore::nn::RMSNorm, norm); - - TinyLlamaModel(const TinyLlamaConfig &config, const infinicore::Device &device) { - INFINICORE_NN_MODULE_INIT(embed_tokens, config.vocab_size, config.hidden_size, std::nullopt, infinicore::DataType::F32, device); - INFINICORE_NN_MODULE_VEC_INIT(layers, config.num_hidden_layers, Block, config, device); - INFINICORE_NN_MODULE_INIT(norm, config.hidden_size, config.rms_norm_eps, infinicore::DataType::F32, device); - } - }; - - // ============================================ - // Phase 1: Model Construction Verification - // ============================================ - spdlog::info(""); - spdlog::info("Phase 1: Model Construction Verification"); - spdlog::info("------------------------------------------"); - - // Construct the model - TinyLlamaModel model(config, infinicore::Device()); - - // Verify all components are created - auto state = model.state_dict(); - spdlog::info("✓ Model constructed with {} parameters", state.size()); - - // Parameter count expectation: - // embed_tokens.weight (1) + norm.weight (1) + per-layer (9 params) * num_layers - size_t expected_param_count = 1 + 1 + config.num_hidden_layers * 9; - if (state.size() != expected_param_count) { - spdlog::error("Parameter count mismatch. Got {}, expected {} (1 + {}*9 + 1)", - state.size(), expected_param_count, config.num_hidden_layers); - // Do not return false here to allow listing and detailed checks below - } - - // List all parameters for manual verification - spdlog::info("Listing all Tiny-Llama parameters (name -> shape):"); - for (const auto &kv : state) { - const auto &name = kv.first; - const auto &tensor = kv.second; - std::ostringstream shape_ss; - shape_ss << "["; - for (size_t i = 0; i < tensor->shape().size(); ++i) { - if (i) { - shape_ss << ", "; - } - shape_ss << tensor->shape()[i]; - } - shape_ss << "]"; - spdlog::info(" - {} -> {}", name, shape_ss.str()); - } - - // Automated verification: check all parameter shapes match hard-coded TinyLlama hierarchy - spdlog::info("Verifying listed parameters against hard-coded TinyLlama hierarchy..."); - - struct Expect { - std::string name; - std::vector shape; - }; - const size_t kv_dim = config.hidden_size * config.num_key_value_heads / config.num_attention_heads; - std::vector expected; - // embed and final norm - expected.push_back({"embed_tokens.weight", {config.vocab_size, config.hidden_size}}); - // per-layer expectations - for (size_t i = 0; i < config.num_hidden_layers; ++i) { - const std::string prefix = std::string("layers.") + std::to_string(i) + "."; - expected.push_back({prefix + "input_layernorm.weight", {config.hidden_size}}); - expected.push_back({prefix + "self_attn.q_proj.weight", {config.hidden_size, config.hidden_size}}); - expected.push_back({prefix + "self_attn.k_proj.weight", {kv_dim, config.hidden_size}}); - expected.push_back({prefix + "self_attn.v_proj.weight", {kv_dim, config.hidden_size}}); - expected.push_back({prefix + "self_attn.o_proj.weight", {config.hidden_size, config.hidden_size}}); - expected.push_back({prefix + "post_attention_layernorm.weight", {config.hidden_size}}); - expected.push_back({prefix + "mlp.gate_proj.weight", {config.intermediate_size, config.hidden_size}}); - expected.push_back({prefix + "mlp.up_proj.weight", {config.intermediate_size, config.hidden_size}}); - expected.push_back({prefix + "mlp.down_proj.weight", {config.hidden_size, config.intermediate_size}}); - } - expected.push_back({"norm.weight", {config.hidden_size}}); - - bool all_ok = true; - // Check expected ones (existence and shapes) - for (const auto &e : expected) { - auto it = state.find(e.name); - if (it == state.end()) { - spdlog::error("Missing expected parameter: {}", e.name); - all_ok = false; - continue; - } - auto got = it->second->shape(); - if (got != e.shape) { - std::ostringstream got_ss, exp_ss; - got_ss << "["; - for (size_t i = 0; i < got.size(); ++i) { - if (i) { - got_ss << ", "; - } - got_ss << got[i]; - } - got_ss << "]"; - exp_ss << "["; - for (size_t i = 0; i < e.shape.size(); ++i) { - if (i) { - exp_ss << ", "; - } - exp_ss << e.shape[i]; - } - exp_ss << "]"; - spdlog::error("Shape mismatch for '{}': got {}, expected {}", e.name, got_ss.str(), exp_ss.str()); - all_ok = false; - } - } - - // Check for unexpected extra parameters - for (const auto &kvp : state) { - const auto &name = kvp.first; - bool is_expected = false; - for (const auto &e : expected) { - if (e.name == name) { - is_expected = true; - break; - } - } - if (!is_expected) { - std::ostringstream got_ss; - auto got = kvp.second->shape(); - got_ss << "["; - for (size_t i = 0; i < got.size(); ++i) { - if (i) { - got_ss << ", "; - } - got_ss << got[i]; - } - got_ss << "]"; - spdlog::warn("Unexpected parameter present: {} with shape {}", name, got_ss.str()); - } - } - - if (!all_ok) { - spdlog::error("Tiny-Llama parameter verification: FAILED - see errors above"); - return false; - } - spdlog::info("Tiny-Llama parameter verification: PASSED"); - - // Create test weights - std::unordered_map test_state_dict; - for (const auto &[name, tensor] : state) { - // Create a test tensor with ones - test_state_dict.emplace(name, infinicore::Tensor::ones(tensor->shape(), - infinicore::DataType::F32, - infinicore::Device())); - } - - // Load the test weights - model.load_state_dict(test_state_dict); - - // Verify weights were loaded - auto loaded_state = model.state_dict(); - bool load_success = true; - for (const auto &[name, _] : test_state_dict) { - if (loaded_state.find(name) == loaded_state.end()) { - spdlog::error("Parameter '{}' not found after load_state_dict", name); - load_success = false; - } - } - - if (!load_success) { - spdlog::error("Weight loading verification failed"); - return false; - } - - spdlog::info("✓ State dict save/load mechanism verified"); - - // ============================================ - // Summary - // ============================================ - spdlog::info(""); - spdlog::info("=========================================="); - spdlog::info("✅ Tiny-Llama Model Test Summary"); - spdlog::info("=========================================="); - spdlog::info("✓ Metadata validation: PASSED (config matches actual model)"); - spdlog::info("✓ Model construction: PASSED"); - spdlog::info("✓ Parameter shapes: PASSED (11 parameters)"); - spdlog::info("✓ Forward passes: PASSED"); - spdlog::info("✓ Weight loading mechanism: PASSED"); - spdlog::info("✓ Architecture compatibility: Tiny-Llama-1.1B-Chat-v1.0"); - spdlog::info("✓ GQA support: num_key_value_heads={}", config.num_key_value_heads); - spdlog::info(""); - spdlog::info("Model is ready for:"); - spdlog::info(" - Full 22-layer implementation"); - spdlog::info(" - Safetensors/pickle weight loading"); - spdlog::info(" - Inference and fine-tuning"); - spdlog::info("=========================================="); - - return true; - - } catch (const std::exception &e) { - spdlog::error("Exception in testTinyLlamaConstruction: {}", e.what()); - return false; - } - }); -} - -// Main test runner -TestResult NNModuleTest::run() { - std::vector results; - - std::cout << "==============================================\n" - << "InfiniCore nn::Module Test Suite\n" - << "==============================================" << std::endl; - - results.push_back(testBasicModuleCreation()); // Merged: creation + parameters + state_dict + load - results.push_back(testTensorParallelParameters()); // Tensor-parallel parameters - results.push_back(testParalleLinear()); // ParalleLinear - results.push_back(testLoadStateDict()); // Advanced: hierarchical modules - results.push_back(testModuleHierarchy()); // Demonstrates hierarchical construction - results.push_back(testParameterLoading()); // Blob loading - results.push_back(testModuleLinear()); // Linear module comprehensive test - results.push_back(testModuleEmbedding()); // Embedding module test - results.push_back(testModuleRMSNorm()); // RMSNorm module test - results.push_back(testModuleRoPE()); // RoPE module test - results.push_back(testDtypeAssertion()); // Dtype assertion test - results.push_back(testTinyLlamaConstruction()); // Comprehensive: TinyLlama model test - - // Check if all tests passed - bool all_passed = true; - for (const auto &result : results) { - if (!result.passed) { - all_passed = false; - break; - } - } - - return TestResult("NNModuleTest", all_passed, - all_passed ? "" : "Some nn::module tests failed"); -} - -} // namespace infinicore::test diff --git a/src/infinicore-test/test_nn_module.h b/src/infinicore-test/test_nn_module.h deleted file mode 100644 index 7004c1f55..000000000 --- a/src/infinicore-test/test_nn_module.h +++ /dev/null @@ -1,107 +0,0 @@ -#ifndef __INFINICORE_TEST_NN_MODULE_H__ -#define __INFINICORE_TEST_NN_MODULE_H__ - -#include "infinicore/device.hpp" -#include "infinicore/nn/embedding.hpp" -#include "infinicore/nn/module.hpp" -#include "infinicore/nn/parameter.hpp" -#include "infinicore/nn/rmsnorm.hpp" -#include "infinicore/nn/rope.hpp" -#include "test_runner.h" -#include -#include -#include -#include -#include -#include -#include -#include - -namespace infinicore::test { - -// Simple test module that mimics torch.nn.Linear -class MockLinearModule : public infinicore::nn::Module { -public: - // Declare parameters using macros (torch-like style) - INFINICORE_NN_PARAMETER(weight); - INFINICORE_NN_PARAMETER(bias); - - MockLinearModule(int input_size, int output_size, const infinicore::Device &device, - Size tp_dim = 0, Size tp_rank = 0, Size tp_size = 1) - : input_size_(input_size), output_size_(output_size), device_(device), - tp_dim_(tp_dim), tp_rank_(tp_rank), tp_size_(tp_size) { - // Initialize parameters using macros - INFINICORE_NN_PARAMETER_INIT(weight, - ({static_cast(output_size), static_cast(input_size)}, - infinicore::DataType::F32, - device, - tp_dim_, - tp_rank_, - tp_size_)); - INFINICORE_NN_PARAMETER_INIT(bias, - ({static_cast(output_size)}, - infinicore::DataType::F32, - device, - 0, - tp_dim == 0 ? tp_rank_ : 0, - tp_dim == 0 ? tp_size_ : 1)); - } - - // Simple forward pass (conceptual - would need actual matrix operations) - infinicore::Tensor forward(const infinicore::Tensor &input) { - // This is a placeholder - in a real implementation, you'd do matrix multiplication - // For testing purposes, we'll just return the input - return input; - } - - infinicore::Tensor get_weight() const { - auto state_dict = this->state_dict(); - auto it = state_dict.find("weight"); - if (it != state_dict.end()) { - return it->second; - } - throw std::runtime_error("Weight parameter not found"); - } - - infinicore::Tensor get_bias() const { - auto state_dict = this->state_dict(); - auto it = state_dict.find("bias"); - if (it != state_dict.end()) { - return it->second; - } - throw std::runtime_error("Bias parameter not found"); - } - -private: - int input_size_; - int output_size_; - infinicore::Device device_; - - Size tp_dim_; - Size tp_rank_; - Size tp_size_; -}; - -class NNModuleTest : public TestFramework { -public: - TestResult run() override; - std::string getName() const override { return "NNModuleTest"; } - -private: - TestResult testBasicModuleCreation(); // Merged: creation, parameters, state_dict, load_state_dict - TestResult testTensorParallelParameters(); // Module with tensor parallel parameters - TestResult testParalleLinear(); // Module with ColumnParallelLinear, RowParallelLinear - TestResult testLoadStateDict(); // Advanced: hierarchical modules - TestResult testModuleHierarchy(); // Demonstrates proper hierarchical construction pattern - TestResult testParameterLoading(); // Test blob parameter loading - TestResult testModuleLinear(); // Comprehensive Linear module test - TestResult testModuleEmbedding(); // Embedding module test - TestResult testModuleRMSNorm(); // RMSNorm module test - TestResult testModuleRoPE(); // RoPE module test - TestResult testDtypeAssertion(); // Test dtype assertions when loading parameters - TestResult testTinyLlamaConstruction(); // Comprehensive: construction + weight loading + validation -}; - -} // namespace infinicore::test - -#endif // __INFINICORE_TEST_NN_MODULE_H__ diff --git a/src/infinicore-test/test_runner.h b/src/infinicore-test/test_runner.h deleted file mode 100644 index 8912a6a20..000000000 --- a/src/infinicore-test/test_runner.h +++ /dev/null @@ -1,260 +0,0 @@ -#ifndef __INFINICORE_TEST_RUNNER_H__ -#define __INFINICORE_TEST_RUNNER_H__ - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace infinicore::test { - -// ============================================================================ -// Common Test Utilities -// ============================================================================ - -/** - * @brief Compare two InfiniCore tensors elementwise with tolerance - * - * Compares two tensors for approximate equality, useful for testing numerical - * computations where exact equality is not expected due to floating-point arithmetic. - * - * @param actual The actual tensor result - * @param expected The expected tensor result - * @param rtol Relative tolerance (default: 1e-5) - * @param atol Absolute tolerance (default: 1e-5) - * @return true if tensors are approximately equal, false otherwise - * - * @note Currently only supports F32 dtype - * @note Tensors are automatically moved to CPU for comparison - * @note Reports up to 10 mismatches with detailed coordinates - */ -inline bool tensorsAllClose(const infinicore::Tensor &actual, - const infinicore::Tensor &expected, - double rtol = 1e-5, - double atol = 1e-5) { - if (actual->shape() != expected->shape()) { - spdlog::error("Shape mismatch: actual vs expected"); - return false; - } - - auto cpu = infinicore::Device(infinicore::Device::Type::CPU, 0); - auto a_cpu = actual->to(cpu); - a_cpu = a_cpu->contiguous(); - auto b_cpu = expected->to(cpu); - b_cpu = b_cpu->contiguous(); - - if (a_cpu->dtype() != b_cpu->dtype()) { - spdlog::error("DType mismatch"); - return false; - } - - // Only support F32 in this test - if (a_cpu->dtype() != infinicore::DataType::F32) { - spdlog::error("Unsupported dtype for comparison; only F32 supported in test"); - return false; - } - - size_t n = a_cpu->numel(); - const auto &shape = a_cpu->shape(); - - // Precompute strides for index -> coords mapping - std::vector stride(shape.size(), 1); - for (int i = static_cast(shape.size()) - 2; i >= 0; --i) { - stride[i] = stride[i + 1] * shape[i + 1]; - } - - const float *ap = reinterpret_cast(a_cpu->data()); - const float *bp = reinterpret_cast(b_cpu->data()); - size_t max_diff_index = 0; - float max_diff = 0.0f; - size_t num_fail_reported = 0; - - for (size_t i = 0; i < n; ++i) { - float av = ap[i]; - float bv = bp[i]; - float diff = std::fabs(av - bv); - if (diff > static_cast(atol + rtol * std::fabs(bv))) { - if (diff > max_diff) { - max_diff = diff; - max_diff_index = i; - } - if (num_fail_reported < 10) { - // Convert linear index to coordinates - std::vector coords(shape.size(), 0); - size_t t = i; - for (size_t d = 0; d < shape.size(); ++d) { - coords[d] = t / stride[d]; - t -= coords[d] * stride[d]; - } - std::stringstream ss; - ss << "["; - for (size_t d = 0; d < coords.size(); ++d) { - ss << coords[d] << (d + 1 < coords.size() ? "," : "]"); - } - double tol = atol + rtol * std::fabs(bv); - spdlog::error("Mismatch at index {} coords {}: actual={} expected={} diff={} tol={}", - i, ss.str(), av, bv, diff, tol); - num_fail_reported++; - } - } - } - - if (num_fail_reported > 0) { - // Report summary with max diff - std::vector coords(shape.size(), 0); - size_t t = max_diff_index; - for (size_t d = 0; d < shape.size(); ++d) { - coords[d] = t / stride[d]; - t -= coords[d] * stride[d]; - } - std::stringstream ss; - ss << "["; - for (size_t d = 0; d < coords.size(); ++d) { - ss << coords[d] << (d + 1 < coords.size() ? "," : "]"); - } - spdlog::error("Max diff {} at linear index {} coords {}", max_diff, max_diff_index, ss.str()); - return false; - } - - return true; -} - -// ============================================================================ -// Test Framework Classes -// ============================================================================ - -// Test result structure -struct TestResult { - std::string test_name; - bool passed; - std::string error_message; - std::chrono::microseconds duration; - - TestResult(const std::string &name, bool pass, const std::string &error = "", - std::chrono::microseconds dur = std::chrono::microseconds(0)) - : test_name(name), passed(pass), error_message(error), duration(dur) {} -}; - -// Test framework base class -class TestFramework { -public: - virtual ~TestFramework() = default; - virtual TestResult run() = 0; - virtual std::string getName() const = 0; - -protected: - void logTestStart(const std::string &test_name) { - std::cout << "[TEST] Starting: " << test_name << std::endl; - } - - void logTestResult(const TestResult &result) { - std::cout << "[TEST] " << (result.passed ? "PASSED" : "FAILED") - << ": " << result.test_name; - if (!result.passed && !result.error_message.empty()) { - std::cout << " - " << result.error_message; - } - std::cout << " (Duration: " << result.duration.count() << "μs)" << std::endl; - } - - template - TestResult measureTime(const std::string &test_name, Func &&func) { - auto start = std::chrono::high_resolution_clock::now(); - try { - bool result = func(); - auto end = std::chrono::high_resolution_clock::now(); - auto duration = std::chrono::duration_cast(end - start); - return TestResult(test_name, result, "", duration); - } catch (const std::exception &e) { - auto end = std::chrono::high_resolution_clock::now(); - auto duration = std::chrono::duration_cast(end - start); - return TestResult(test_name, false, e.what(), duration); - } - } -}; - -// Test runner -class InfiniCoreTestRunner { -public: - void addTest(std::unique_ptr test) { - tests_.push_back(std::move(test)); - } - - std::vector runAllTests() { - std::vector results; - - std::cout << "==============================================\n" - << "InfiniCore Test Suite\n" - << "==============================================" << std::endl; - - for (auto &test : tests_) { - logTestStart(test->getName()); - TestResult result = test->run(); - logTestResult(result); - results.push_back(result); - } - - printSummary(results); - return results; - } - -private: - std::vector> tests_; - - void logTestStart(const std::string &test_name) { - std::cout << "\n[SUITE] Running: " << test_name << std::endl; - } - - void logTestResult(const TestResult &result) { - std::cout << "[SUITE] " << (result.passed ? "PASSED" : "FAILED") - << ": " << result.test_name << std::endl; - } - - void printSummary(const std::vector &results) { - size_t passed = 0, failed = 0; - std::chrono::microseconds total_time(0); - std::vector failed_tests; - - for (const auto &result : results) { - if (result.passed) { - passed++; - } else { - failed++; - failed_tests.push_back(result); - } - total_time += result.duration; - } - - // Print list of failed tests if any - if (!failed_tests.empty()) { - std::cout << "\n==============================================\n" - << "❌ FAILED TESTS\n" - << "==============================================" << std::endl; - for (const auto &test : failed_tests) { - std::cout << " • " << test.test_name; - if (!test.error_message.empty()) { - std::cout << "\n Error: " << test.error_message; - } - std::cout << "\n Duration: " << test.duration.count() << "μs" << std::endl; - } - } - - std::cout << "\n==============================================\n" - << "Test Summary\n" - << "==============================================\n" - << "Total Tests: " << results.size() << "\n" - << "Passed: " << passed << "\n" - << "Failed: " << failed << "\n" - << "Total Time: " << total_time.count() << "μs\n" - << "==============================================" << std::endl; - } -}; - -} // namespace infinicore::test - -#endif // __INFINICORE_TEST_RUNNER_H__ diff --git a/src/infinicore-test/test_tensor_destructor.cc b/src/infinicore-test/test_tensor_destructor.cc deleted file mode 100644 index c75dd546b..000000000 --- a/src/infinicore-test/test_tensor_destructor.cc +++ /dev/null @@ -1,405 +0,0 @@ -#include "test_tensor_destructor.h" -#include "../src/utils/custom_types.h" -#include -namespace infinicore::test { - -// Test 1: Basic tensor creation and destruction -TestResult TensorDestructorTest::testBasicTensorDestruction() { - return measureTime("BasicTensorDestruction", [this]() { - { - // Create a tensor in a scope to test automatic destruction - auto tensor = Tensor::empty({2, 3}, DataType::F32, Device::Type::CPU); - - // Verify tensor was created successfully - if (!tensor.operator->()) { - return false; - } - if (tensor->shape().size() != 2) { - return false; - } - if (tensor->shape()[0] != 2) { - return false; - } - if (tensor->shape()[1] != 3) { - return false; - } - if (tensor->dtype() != DataType::F32) { - return false; - } - - std::cout << "Tensor created successfully with shape: "; - for (auto dim : tensor->shape()) { - std::cout << dim << " "; - } - std::cout << std::endl; - } - - // Tensor should be destroyed when it goes out of scope - // This should trigger the TensorMetaData destructor - std::cout << "Tensor destroyed successfully - destructor called" << std::endl; - return true; - }); -} - -// Test 2: Multiple tensor creation and destruction -TestResult TensorDestructorTest::testMultipleTensorDestruction() { - return measureTime("MultipleTensorDestruction", [this]() { - std::vector tensors; - - // Create multiple tensors with different shapes and types - tensors.push_back(Tensor::empty({1, 2, 3}, DataType::F32, Device::Type::CPU)); - tensors.push_back(Tensor::empty({4, 5}, DataType::F64, Device::Type::CPU)); - tensors.push_back(Tensor::zeros({2, 2, 2}, DataType::I32, Device::Type::CPU)); - tensors.push_back(Tensor::ones({3, 4}, DataType::F16, Device::Type::CPU)); - - // Verify all tensors were created - if (tensors.size() != 4) { - return false; - } - for (size_t i = 0; i < tensors.size(); ++i) { - if (!tensors[i].operator->()) { - return false; - } - std::cout << "Tensor " << i << " created with shape: "; - for (auto dim : tensors[i]->shape()) { - std::cout << dim << " "; - } - std::cout << std::endl; - } - - std::cout << "All " << tensors.size() << " tensors created successfully" << std::endl; - - // All tensors will be destroyed when the vector goes out of scope - // This should trigger all TensorMetaData destructors - return true; - }); -} - -// Test 3: Different data types -TestResult TensorDestructorTest::testDifferentDataTypes() { - return measureTime("DifferentDataTypes", [this]() { - std::vector> data_types = { - {DataType::F32, "F32"}, - {DataType::F64, "F64"}, - {DataType::F16, "F16"}, - {DataType::I32, "I32"}, - {DataType::I64, "I64"}, - {DataType::I8, "I8"}, - {DataType::U8, "U8"}, - {DataType::BOOL, "BOOL"}}; - - for (const auto &[dtype, name] : data_types) { - { - auto tensor = Tensor::empty({2, 2}, dtype, Device::Type::CPU); - if (!tensor.operator->()) { - return false; - } - if (tensor->dtype() != dtype) { - return false; - } - std::cout << "Created tensor with data type: " << name << std::endl; - } - std::cout << "Destroyed tensor with data type: " << name << std::endl; - } - - return true; - }); -} - -// Test 4: Different shapes -TestResult TensorDestructorTest::testDifferentShapes() { - return measureTime("DifferentShapes", [this]() { - std::vector shapes = { - {1}, // 1D - {2, 3}, // 2D - {4, 5, 6}, // 3D - {1, 2, 3, 4}, // 4D - {2, 3, 4, 5, 6}, // 5D - {1000}, // Large 1D - {100, 100}, // Large 2D - {10, 10, 10, 10} // Large 4D - }; - - for (const auto &shape : shapes) { - { - auto tensor = Tensor::empty(shape, DataType::F32, Device::Type::CPU); - if (!tensor.operator->()) { - return false; - } - if (tensor->shape() != shape) { - return false; - } - std::cout << "Created tensor with shape: "; - for (auto dim : shape) { - std::cout << dim << " "; - } - std::cout << std::endl; - } - std::cout << "Destroyed tensor with shape: "; - for (auto dim : shape) { - std::cout << dim << " "; - } - std::cout << std::endl; - } - - return true; - }); -} - -// Test 5: Tensor from blob -TestResult TensorDestructorTest::testTensorFromBlob() { - return measureTime("TensorFromBlob", [this]() { - // Create a blob of data - std::vector data = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}; - - { - // Create tensor from blob - auto tensor = Tensor::from_blob(data.data(), {2, 3}, DataType::F32, Device::Type::CPU); - if (!tensor.operator->()) { - return false; - } - if (tensor->shape() != Shape({2, 3})) { - return false; - } - if (tensor->dtype() != DataType::F32) { - return false; - } - - std::cout << "Created tensor from blob with shape: "; - for (auto dim : tensor->shape()) { - std::cout << dim << " "; - } - std::cout << std::endl; - } - - std::cout << "Destroyed tensor from blob successfully" << std::endl; - return true; - }); -} - -// Test 6: Strided tensor -TestResult TensorDestructorTest::testStridedTensor() { - return measureTime("StridedTensor", [this]() { - { - // Create a strided tensor - auto tensor = Tensor::empty({4, 4}, DataType::F32, Device::Type::CPU); - if (!tensor.operator->()) { - return false; - } - - // Create a narrowed view - std::vector slices = { - {0, 0, 2}, // dimension 0: start at 0, length 2 - {1, 0, 2} // dimension 1: start at 0, length 2 - }; - auto strided_tensor = tensor->narrow(slices); - if (!strided_tensor.operator->()) { - return false; - } - - std::cout << "Created strided tensor with shape: "; - for (auto dim : strided_tensor->shape()) { - std::cout << dim << " "; - } - std::cout << std::endl; - } - - std::cout << "Destroyed strided tensor successfully" << std::endl; - return true; - }); -} - -// Test 7: Memory leak detection -TestResult TensorDestructorTest::testMemoryLeakDetection() { - return measureTime("MemoryLeakDetection", [this]() { - // Reset memory leak detector - MemoryLeakDetector::instance().reset(); - - size_t initial_leaks = MemoryLeakDetector::instance().getLeakCount(); - - // Create and destroy many tensors - for (int i = 0; i < 100; ++i) { - { - auto tensor = Tensor::empty({10, 10}, DataType::F32, Device::Type::CPU); - if (!tensor.operator->()) { - return false; - } - } - } - - size_t final_leaks = MemoryLeakDetector::instance().getLeakCount(); - - std::cout << "Initial leaks: " << initial_leaks << std::endl; - std::cout << "Final leaks: " << final_leaks << std::endl; - - // Should not have more leaks than we started with - return final_leaks <= initial_leaks; - }); -} - -// Test 8: Tensor copy destruction -TestResult TensorDestructorTest::testTensorCopyDestruction() { - return measureTime("TensorCopyDestruction", [this]() { - { - auto original_tensor = Tensor::empty({3, 3}, DataType::F32, Device::Type::CPU); - if (!original_tensor.operator->()) { - return false; - } - - // Create a copy (using assignment operator) - auto copied_tensor = original_tensor; - if (!copied_tensor.operator->()) { - return false; - } - - std::cout << "Created original and copied tensors" << std::endl; - std::cout << "Original tensor shape: "; - for (auto dim : original_tensor->shape()) { - std::cout << dim << " "; - } - std::cout << std::endl; - std::cout << "Copied tensor shape: "; - for (auto dim : copied_tensor->shape()) { - std::cout << dim << " "; - } - std::cout << std::endl; - } - - std::cout << "Destroyed original and copied tensors successfully" << std::endl; - return true; - }); -} - -// Test 9: print options -TestResult TensorDestructorTest::testPrintOptions() { - return measureTime("PrintOptions", [this]() { - using namespace infinicore::print_options; - - // Prepare test data (stored outside loop to ensure lifetime) - bool bool_data[4] = {true, false, true, false}; - int8_t i8_data[4] = {-128, -64, 32, 127}; - int16_t i16_data[4] = {-32768, -16384, 8192, 32767}; - int32_t i32_data[4] = {-2147483648, -1073741824, 1073741824, 2147483647}; - int64_t i64_data[4] = {-1000000000000000000LL, -500000000000000000LL, - 500000000000000000LL, 1000000000000000000LL}; - uint8_t u8_data[4] = {0, 64, 192, 255}; - uint16_t u16_data[4] = {0, 16384, 49152, 65535}; - uint32_t u32_data[4] = {0, 1073741824, 3221225472, 4294967295}; - uint64_t u64_data[4] = {0, 1000000000000000000ULL, 5000000000000000000ULL, 10000000000000000000ULL}; - fp16_t f16_data[4] = {_f32_to_f16(1.234f), _f32_to_f16(2.345f), _f32_to_f16(4.567f), _f32_to_f16(5.678f)}; - bf16_t bf16_data[4] = {_f32_to_bf16(1.234f), _f32_to_bf16(2.345f), _f32_to_bf16(4.567f), _f32_to_bf16(5.678f)}; - float f32_data[6] = {1.23456789f, 2.3456789f, 3.456789f, 4.56789f, 5.6789f, 6.78901234f}; - double f64_data[4] = {1.23456789111, 2.3456789, 4.56789, 5.6789}; - - // Test cases list: (name, dtype, tensor) - using TestCaseTuple = std::tuple; - std::vector case_list = { - {"BOOL", Tensor::from_blob(bool_data, {2, 2}, DataType::BOOL, Device::Type::CPU)}, - {"I8", Tensor::from_blob(i8_data, {2, 2}, DataType::I8, Device::Type::CPU)}, - {"I16", Tensor::from_blob(i16_data, {2, 2}, DataType::I16, Device::Type::CPU)}, - {"I32", Tensor::from_blob(i32_data, {2, 2}, DataType::I32, Device::Type::CPU)}, - {"I64", Tensor::from_blob(i64_data, {2, 2}, DataType::I64, Device::Type::CPU)}, - {"U8", Tensor::from_blob(u8_data, {2, 2}, DataType::U8, Device::Type::CPU)}, - {"U16", Tensor::from_blob(u16_data, {2, 2}, DataType::U16, Device::Type::CPU)}, - {"U32", Tensor::from_blob(u32_data, {2, 2}, DataType::U32, Device::Type::CPU)}, - {"U64", Tensor::from_blob(u64_data, {2, 2}, DataType::U64, Device::Type::CPU)}, - {"F16", Tensor::from_blob(f16_data, {2, 2}, DataType::F16, Device::Type::CPU)}, - {"BF16", Tensor::from_blob(bf16_data, {2, 2}, DataType::BF16, Device::Type::CPU)}, - {"F32", Tensor::from_blob(f32_data, {2, 3}, DataType::F32, Device::Type::CPU)}, - {"F64", Tensor::from_blob(f64_data, {2, 2}, DataType::F64, Device::Type::CPU)}, - }; - - std::cout << "\n=== Testing Print Options for Different Data Types ===" << std::endl; - // Process each test case - for (const auto &test_case : case_list) { - const std::string &name = std::get<0>(test_case); - const Tensor &tensor = std::get<1>(test_case); - std::cout << "\n--- Testing " << name << " ---" << std::endl; - std::cout << tensor << std::endl; - } - - // Test print options with F32 tensor - std::cout << "\n=== Testing Print Options with F32 Tensor ===" << std::endl; - { - auto dataf32 = Tensor::from_blob(f32_data, {2, 3}, DataType::F32, Device::Type::CPU); - - // Test different precision values - std::cout << "\n--- Testing Precision Options ---" << std::endl; - std::cout << "Default (precision=-1, auto):" << std::endl; - std::cout << dataf32 << std::endl; - - std::cout << "\nWith precision=1:" << std::endl; - set_precision(1); - std::cout << dataf32 << std::endl; - - // Test different line_width values - std::cout << "\n--- Testing Line Width Options ---" << std::endl; - std::cout << "\nWith line_width=20:" << std::endl; - set_line_width(20); - std::cout << dataf32 << std::endl; - - // Test summarization - std::cout << "\nWith threshold=4:" << std::endl; - set_line_width(80); - set_edge_items(1); - set_threshold(4); - std::cout << dataf32 << std::endl; - - // Test sci_mode options - std::cout << "\nWith sci_mode=0 (normal notation):" << std::endl; - set_sci_mode(0); - std::cout << dataf32 << std::endl; - - std::cout << "\nWith sci_mode=1 (scientific notation):" << std::endl; - set_sci_mode(1); - std::cout << dataf32 << std::endl; - } - - // Test Temporary print options with F32 tensor - set_sci_mode(-1); - std::cout << "\n=== Testing Temporary Print Options with F32 Tensor ===" << std::endl; - { - auto dataf32 = Tensor::from_blob(f32_data, {2, 3}, DataType::F32, Device::Type::CPU); - std::cout << "\nWith precision=2 " << std::endl; - std::cout << precision(2) << dataf32 << std::endl; - std::cout << dataf32 << std::endl; - } - - std::cout << "\nPrint options test completed successfully" << std::endl; - return true; - }); -} - -// Main test runner -TestResult TensorDestructorTest::run() { - std::vector results; - - std::cout << "==============================================\n" - << "Tensor Destructor Test Suite\n" - << "==============================================" << std::endl; - - // Run all tests - results.push_back(testBasicTensorDestruction()); - results.push_back(testMultipleTensorDestruction()); - results.push_back(testDifferentDataTypes()); - results.push_back(testDifferentShapes()); - results.push_back(testTensorFromBlob()); - results.push_back(testStridedTensor()); - results.push_back(testMemoryLeakDetection()); - results.push_back(testTensorCopyDestruction()); - results.push_back(testPrintOptions()); - - // Check if all tests passed - bool all_passed = true; - for (const auto &result : results) { - if (!result.passed) { - all_passed = false; - break; - } - } - - return TestResult("TensorDestructorTest", all_passed, - all_passed ? "" : "Some tensor destructor tests failed"); -} - -} // namespace infinicore::test diff --git a/src/infinicore-test/test_tensor_destructor.h b/src/infinicore-test/test_tensor_destructor.h deleted file mode 100644 index b29505766..000000000 --- a/src/infinicore-test/test_tensor_destructor.h +++ /dev/null @@ -1,33 +0,0 @@ -#ifndef __INFINICORE_TEST_TENSOR_DESTRUCTOR_H__ -#define __INFINICORE_TEST_TENSOR_DESTRUCTOR_H__ - -#include "infinicore/context/context.hpp" -#include "infinicore/tensor.hpp" -#include "memory_test.h" -#include "test_runner.h" -#include -#include -#include - -namespace infinicore::test { - -class TensorDestructorTest : public TestFramework { -public: - TestResult run() override; - std::string getName() const override { return "TensorDestructorTest"; } - -private: - TestResult testBasicTensorDestruction(); - TestResult testMultipleTensorDestruction(); - TestResult testDifferentDataTypes(); - TestResult testDifferentShapes(); - TestResult testTensorFromBlob(); - TestResult testStridedTensor(); - TestResult testMemoryLeakDetection(); - TestResult testTensorCopyDestruction(); - TestResult testPrintOptions(); -}; - -} // namespace infinicore::test - -#endif // __INFINICORE_TEST_TENSOR_DESTRUCTOR_H__ diff --git a/src/infinicore/adaptor/aten_adaptor.cc b/src/infinicore/adaptor/aten_adaptor.cc deleted file mode 100644 index 73c01dfdb..000000000 --- a/src/infinicore/adaptor/aten_adaptor.cc +++ /dev/null @@ -1,56 +0,0 @@ -#ifdef ENABLE_ATEN -#include "infinicore/adaptor/aten_adaptor.hpp" - -namespace infinicore::adaptor { - -at::Tensor to_aten_tensor(const infinicore::Tensor &t) { - void *data_ptr = (void *)(t->data()); - - auto sizes = std::vector( - t->shape().begin(), - t->shape().end()); - - auto strides = t->strides(); - - auto dtype = to_at_dtype(t->dtype()); - auto device = to_at_device(t->device()); - - auto deleter_ = [](void * /*unused*/) mutable { - - }; - - at::TensorOptions options = at::TensorOptions() - .dtype(dtype) - .device(device) - .requires_grad(false); - - return at::from_blob( - data_ptr, - sizes, - strides, - deleter_, - options); -} - -#if defined(ENABLE_HYGON_API) -c10::hip::HIPStream get_hip_stream() { - return c10::hip::getStreamFromExternal( - hipStream_t(infinicore::context::getStream()), infinicore::context::getDevice().getIndex()); -} -#elif defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) -c10::cuda::CUDAStream get_cuda_stream() { - return c10::cuda::getStreamFromExternal( - cudaStream_t(infinicore::context::getStream()), infinicore::context::getDevice().getIndex()); -} -#endif - -#if defined(ENABLE_MOORE_API) -c10::musa::MUSAStream get_musa_stream() { - return c10::musa::getStreamFromExternal( - musaStream_t(infinicore::context::getStream()), infinicore::context::getDevice().getIndex()); -} -#endif - -} // namespace infinicore::adaptor - -#endif // ENABLE_ATEN diff --git a/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.cc b/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.cc deleted file mode 100644 index a26b400ba..000000000 --- a/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.cc +++ /dev/null @@ -1,173 +0,0 @@ -#if defined(ENABLE_HYGON_API) && defined(ENABLE_FLASH_ATTN) - -#include "flash_attn_hygon.hpp" - -#include -#include -#include -#include -#include -#include -#include - -using mha_fwd_kvcache_fn_t = std::vector (*)( - at::Tensor &q, - const at::Tensor &kcache, - const at::Tensor &vcache, - std::optional &k_, - std::optional &v_, - std::optional &seqlens_k_, - std::optional &rotary_cos_, - std::optional &rotary_sin_, - std::optional &cache_batch_idx_, - std::optional &leftpad_k_, - std::optional &block_table_, - std::optional &alibi_slopes_, - std::optional &out_, - float softmax_scale, - bool is_causal, - int window_size_left, - int window_size_right, - float softcap, - bool is_rotary_interleaved, - int num_splits, - const std::optional &s_aux_); - -using mha_varlen_fwd_fn_t = std::vector (*)( - at::Tensor &q, - const at::Tensor &k, - const at::Tensor &v, - std::optional &out_, - const at::Tensor &cu_seqlens_q, - const at::Tensor &cu_seqlens_k, - std::optional &seqused_k, - std::optional &leftpad_k_, - std::optional &block_table_, - std::optional &alibi_slopes_, - int max_seqlen_q, - int max_seqlen_k, - float p_dropout, - float softmax_scale, - bool zero_tensors, - bool is_causal, - int window_size_left, - int window_size_right, - float softcap, - bool return_softmax, - std::optional q_descale_, - std::optional k_descale_, - std::optional v_descale_, - std::optional gen_, - const std::optional &s_aux_); - -namespace { - -void *resolve_flash_symbol(const char *name) { - void *sym = dlsym(RTLD_DEFAULT, name); - if (!sym) { - throw std::runtime_error(std::string("flash_attn symbol not found: ") + name); - } - return sym; -} - -void *resolve_flash_extension_symbol(const char *name) { - static std::string last_error; - static void *handle = [&]() -> void * { - const char *candidates[] = { - "flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so", - "/usr/local/lib/python3.10/dist-packages/flash_attn_2_cuda.cpython-310-x86_64-linux-gnu.so", - }; - for (const char *candidate : candidates) { - dlerror(); - if (void *h = dlopen(candidate, RTLD_NOW | RTLD_GLOBAL)) { - return h; - } - if (const char *err = dlerror()) { - last_error += std::string(candidate) + ": " + err + "; "; - } - } - return nullptr; - }(); - if (handle) { - dlerror(); - if (void *sym = dlsym(handle, name)) { - return sym; - } - if (const char *err = dlerror()) { - throw std::runtime_error(std::string("flash_attn extension symbol not found: ") + name + "; " + err); - } - throw std::runtime_error(std::string("flash_attn extension symbol not found: ") + name); - } - throw std::runtime_error(std::string("flash_attn extension dlopen failed for symbol: ") + name + "; " + last_error); -} - -} // namespace - -namespace flash { - -std::vector -mha_fwd_kvcache(at::Tensor &q, - const at::Tensor &kcache, - const at::Tensor &vcache, - std::optional &k_, - std::optional &v_, - std::optional &seqlens_k_, - std::optional &rotary_cos_, - std::optional &rotary_sin_, - std::optional &cache_batch_idx_, - std::optional &leftpad_k_, - std::optional &block_table_, - std::optional &alibi_slopes_, - std::optional &out_, - float softmax_scale, - bool is_causal, - int window_size_left, - int window_size_right, - float softcap, - bool is_rotary_interleaved, - int num_splits) { - static auto fn = reinterpret_cast(resolve_flash_extension_symbol("mha_fwd_kvcache")); - std::optional s_aux = std::nullopt; - return fn(q, kcache, vcache, k_, v_, seqlens_k_, rotary_cos_, rotary_sin_, - cache_batch_idx_, leftpad_k_, block_table_, alibi_slopes_, out_, - softmax_scale, is_causal, window_size_left, window_size_right, - softcap, is_rotary_interleaved, num_splits, s_aux); -} - -std::vector -vllm_mha_varlen_fwd(at::Tensor &q, - const at::Tensor &k, - const at::Tensor &v, - std::optional &out_, - const at::Tensor &cu_seqlens_q, - const at::Tensor &cu_seqlens_k, - std::optional &seqused_k, - std::optional &leftpad_k_, - std::optional &block_table_, - std::optional &alibi_slopes_, - int max_seqlen_q, - int max_seqlen_k, - float p_dropout, - float softmax_scale, - bool zero_tensors, - bool is_causal, - int window_size_left, - int window_size_right, - float softcap, - bool return_softmax, - std::optional gen_) { - static auto fn = reinterpret_cast(resolve_flash_extension_symbol("mha_varlen_fwd")); - std::optional q_descale = std::nullopt; - std::optional k_descale = std::nullopt; - std::optional v_descale = std::nullopt; - std::optional s_aux = std::nullopt; - return fn(q, k, v, out_, cu_seqlens_q, cu_seqlens_k, seqused_k, leftpad_k_, - block_table_, alibi_slopes_, max_seqlen_q, max_seqlen_k, p_dropout, - softmax_scale, zero_tensors, is_causal, window_size_left, - window_size_right, softcap, return_softmax, q_descale, k_descale, - v_descale, gen_, s_aux); -} - -} // namespace flash - -#endif diff --git a/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.hpp b/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.hpp deleted file mode 100644 index 8600c8711..000000000 --- a/src/infinicore/adaptor/flash_attn/hygon/flash_attn_hygon.hpp +++ /dev/null @@ -1,60 +0,0 @@ -#ifndef INFINICORE_ADAPTOR_FLASH_ATTN_HYGON_FLASH_ATTN_HYGON_HPP -#define INFINICORE_ADAPTOR_FLASH_ATTN_HYGON_FLASH_ATTN_HYGON_HPP - -#if defined(ENABLE_FLASH_ATTN) && defined(ENABLE_HYGON_API) - -#include -#include -#include - -namespace flash { - -std::vector -mha_fwd_kvcache(at::Tensor &q, - const at::Tensor &kcache, - const at::Tensor &vcache, - std::optional &k_, - std::optional &v_, - std::optional &seqlens_k_, - std::optional &rotary_cos_, - std::optional &rotary_sin_, - std::optional &cache_batch_idx_, - std::optional &leftpad_k_, - std::optional &block_table_, - std::optional &alibi_slopes_, - std::optional &out_, - float softmax_scale, - bool is_causal, - int window_size_left, - int window_size_right, - float softcap, - bool is_rotary_interleaved, - int num_splits); - -std::vector -vllm_mha_varlen_fwd(at::Tensor &q, - const at::Tensor &k, - const at::Tensor &v, - std::optional &out_, - const at::Tensor &cu_seqlens_q, - const at::Tensor &cu_seqlens_k, - std::optional &seqused_k, - std::optional &leftpad_k_, - std::optional &block_table_, - std::optional &alibi_slopes_, - int max_seqlen_q, - int max_seqlen_k, - float p_dropout, - float softmax_scale, - bool zero_tensors, - bool is_causal, - int window_size_left, - int window_size_right, - float softcap, - bool return_softmax, - std::optional gen_); - -} // namespace flash - -#endif // ENABLE_FLASH_ATTN && ENABLE_HYGON_API -#endif // INFINICORE_ADAPTOR_FLASH_ATTN_HYGON_FLASH_ATTN_HYGON_HPP diff --git a/src/infinicore/analyzer/mutual_awareness_analyzer.cc b/src/infinicore/analyzer/mutual_awareness_analyzer.cc deleted file mode 100644 index 63b820ace..000000000 --- a/src/infinicore/analyzer/mutual_awareness_analyzer.cc +++ /dev/null @@ -1,272 +0,0 @@ -#include "infinicore/analyzer/mutual_awareness_analyzer.hpp" - -#include "infinirt.h" - -namespace infinicore::analyzer { - -namespace { - -DeviceResourceSnapshot buildSnapshotFromMemoryStats( - int device_id, - infinicore::Device::Type device_type, - const MemoryStats &stats) { - DeviceResourceSnapshot snapshot; - snapshot.device_id = device_id; - snapshot.device_type = device_type; - snapshot.has_memory_capacity = stats.total_capacity > 0; - snapshot.free_bytes = stats.total_capacity >= stats.allocated_bytes - ? (stats.total_capacity - stats.allocated_bytes) - : 0; - snapshot.total_bytes = stats.total_capacity; - snapshot.used_bytes = stats.allocated_bytes; - snapshot.reserved_bytes = stats.total_capacity; - return snapshot; -} - -DeviceResourceSnapshot buildSnapshotFromInfinirt( - infinicore::Device device, - const MemoryStats &allocator_stats) { - DeviceResourceSnapshot snapshot; - snapshot.device_id = static_cast(device.getIndex()); - snapshot.device_type = device.getType(); - - infinirtDeviceResourceSnapshot_t rt_snapshot{}; - auto status = infinirtGetDeviceResourceSnapshot( - static_cast(device.getType()), - static_cast(device.getIndex()), - &rt_snapshot); - - if (status == INFINI_STATUS_SUCCESS) { - snapshot.has_memory_capacity = (rt_snapshot.valid_fields & INFINIRT_RESOURCE_FIELD_MEMORY_CAPACITY) != 0; - snapshot.has_compute_utilization = (rt_snapshot.valid_fields & INFINIRT_RESOURCE_FIELD_COMPUTE_UTILIZATION) != 0; - snapshot.has_memory_bandwidth_utilization = (rt_snapshot.valid_fields & INFINIRT_RESOURCE_FIELD_MEMORY_BANDWIDTH_UTILIZATION) != 0; - snapshot.has_kernel_time_ratio = (rt_snapshot.valid_fields & INFINIRT_RESOURCE_FIELD_KERNEL_TIME_RATIO) != 0; - snapshot.has_communication = (rt_snapshot.valid_fields & INFINIRT_RESOURCE_FIELD_COMMUNICATION) != 0; - snapshot.kernel_time_estimated = (rt_snapshot.estimated_fields & INFINIRT_RESOURCE_FIELD_KERNEL_TIME_RATIO) != 0; - - snapshot.free_bytes = rt_snapshot.free_bytes; - snapshot.total_bytes = rt_snapshot.total_bytes; - snapshot.used_bytes = rt_snapshot.used_bytes; - snapshot.reserved_bytes = rt_snapshot.reserved_bytes; - snapshot.compute_utilization = rt_snapshot.compute_utilization; - snapshot.memory_bandwidth_utilization = rt_snapshot.memory_bandwidth_utilization; - snapshot.kernel_time_ratio = rt_snapshot.kernel_time_ratio; - snapshot.communication_time_ratio = rt_snapshot.communication_time_ratio; - snapshot.communication_bytes = rt_snapshot.communication_bytes; - } - - // Keep allocator pool information around even when vendor memory info exists. - if (snapshot.reserved_bytes == 0) { - snapshot.reserved_bytes = allocator_stats.total_capacity; - } - - // Fallback path for backends without a resource snapshot implementation. - if (!snapshot.has_memory_capacity || snapshot.total_bytes == 0) { - auto fallback = buildSnapshotFromMemoryStats( - static_cast(device.getIndex()), - device.getType(), - allocator_stats); - snapshot.has_memory_capacity = fallback.has_memory_capacity; - snapshot.free_bytes = fallback.free_bytes; - snapshot.total_bytes = fallback.total_bytes; - snapshot.used_bytes = fallback.used_bytes; - } - - return snapshot; -} - -std::vector collectRuntimeResourceSnapshots() { - std::vector device_snapshots; - // TODO: integrate with ContextImpl allocator stats when available. - // For now, return an empty snapshot list; the analyzer will use - // the fallback path that queries infinirt directly. - return device_snapshots; -} - -} // namespace - -// ============================================================ -// Singleton -// ============================================================ - -MutualAwarenessAnalyzer &MutualAwarenessAnalyzer::instance() { - static MutualAwarenessAnalyzer inst; - return inst; -} - -MutualAwarenessAnalyzer::MutualAwarenessAnalyzer() - : phase_detector_(), - resource_sensor_(), - intent_generator_(), - enabled_(true), - graph_intent_cached_(false) { -} - -// ============================================================ -// Main analysis entry points -// ============================================================ - -OptimizationIntent MutualAwarenessAnalyzer::analyze() { - if (!enabled_) { - return OptimizationIntent{}; - } - - // If we have a cached graph intent, return it - if (graph_intent_cached_) { - return graph_cached_intent_; - } - - // Get recent op trace window - auto &trace = getGlobalOpTrace(); - auto window = trace.getRecentEntries(phase_detector_.config().window_size); - - // Detect phase - PhaseType phase = phase_detector_.detect(window); - - auto device_snapshots = collectRuntimeResourceSnapshots(); - std::vector device_intents; - device_intents.reserve(device_snapshots.size()); - for (auto const &snapshot : device_snapshots) { - device_intents.push_back(resource_sensor_.sense(snapshot)); - } - - // Generate intent - auto intent = intent_generator_.generate(phase, window, device_intents); - - // Cache result - { - std::lock_guard lock(mutex_); - last_intent_ = intent; - } - - return intent; -} - -OptimizationIntent MutualAwarenessAnalyzer::analyze( - const std::vector> &device_stats) { - - if (!enabled_) { - return OptimizationIntent{}; - } - - // If we have a cached graph intent, return it - if (graph_intent_cached_) { - return graph_cached_intent_; - } - - // Get recent op trace window - auto &trace = getGlobalOpTrace(); - auto window = trace.getRecentEntries(phase_detector_.config().window_size); - - // Detect phase - PhaseType phase = phase_detector_.detect(window); - - // Build per-device intents from provided stats - std::vector device_intents; - device_intents.reserve(device_stats.size()); - for (auto &[dev_id, stats] : device_stats) { - device_intents.push_back(resource_sensor_.sense(dev_id, stats)); - } - - // Generate intent - auto intent = intent_generator_.generate(phase, window, device_intents); - - // Cache result - { - std::lock_guard lock(mutex_); - last_intent_ = intent; - } - - return intent; -} - -OptimizationIntent MutualAwarenessAnalyzer::analyze( - const std::vector &device_snapshots) { - - if (!enabled_) { - return OptimizationIntent{}; - } - - if (graph_intent_cached_) { - return graph_cached_intent_; - } - - auto &trace = getGlobalOpTrace(); - auto window = trace.getRecentEntries(phase_detector_.config().window_size); - - PhaseType phase = phase_detector_.detect(window); - - std::vector device_intents; - device_intents.reserve(device_snapshots.size()); - for (auto const &snapshot : device_snapshots) { - device_intents.push_back(resource_sensor_.sense(snapshot)); - } - - auto intent = intent_generator_.generate(phase, window, device_intents); - - { - std::lock_guard lock(mutex_); - last_intent_ = intent; - } - - return intent; -} - -PhaseType MutualAwarenessAnalyzer::getCurrentPhase() const { - if (!enabled_) { - return PhaseType::UNKNOWN; - } - - return phase_detector_.detectFromTrace(getGlobalOpTrace()); -} - -OptimizationGoal MutualAwarenessAnalyzer::getCurrentOptimizationGoal() const { - return const_cast(this)->analyze().global.goal; -} - -const OptimizationIntent &MutualAwarenessAnalyzer::lastIntent() const { - return last_intent_; -} - -// ============================================================ -// Graph recording support -// ============================================================ - -void MutualAwarenessAnalyzer::onGraphRecordingStop() { - if (!enabled_) { - return; - } - - // Analyze the op sequence recorded during graph capture - // and cache the result. Graph ops are static, so we only - // need to analyze once. - graph_cached_intent_ = analyze(); - graph_intent_cached_ = true; -} - -void MutualAwarenessAnalyzer::clearGraphCache() { - graph_intent_cached_ = false; - graph_cached_intent_ = OptimizationIntent{}; -} - -// ============================================================ -// C-style API for external framework integration -// ============================================================ - -OptimizationIntent analyzeCurrentState() { - return MutualAwarenessAnalyzer::instance().analyze(); -} - -PhaseType getCurrentPhase() { - return MutualAwarenessAnalyzer::instance().getCurrentPhase(); -} - -OptimizationGoal getCurrentOptimizationGoal() { - return MutualAwarenessAnalyzer::instance().getCurrentOptimizationGoal(); -} - -void setAnalyzerEnabled(bool enabled) { - MutualAwarenessAnalyzer::instance().setEnabled(enabled); -} - -} // namespace infinicore::analyzer diff --git a/src/infinicore/analyzer/op_trace.cc b/src/infinicore/analyzer/op_trace.cc deleted file mode 100644 index 566a83d09..000000000 --- a/src/infinicore/analyzer/op_trace.cc +++ /dev/null @@ -1,10 +0,0 @@ -#include "infinicore/analyzer/op_trace.hpp" - -namespace infinicore::analyzer { - -OpTraceRing &getGlobalOpTrace() { - static OpTraceRing instance(OpTraceRing::DEFAULT_CAPACITY); - return instance; -} - -} // namespace infinicore::analyzer diff --git a/src/infinicore/context/allocators/device_pinned_allocator.cc b/src/infinicore/context/allocators/device_pinned_allocator.cc deleted file mode 100644 index a8764d5f6..000000000 --- a/src/infinicore/context/allocators/device_pinned_allocator.cc +++ /dev/null @@ -1,43 +0,0 @@ -#include "device_pinned_allocator.hpp" - -#include "../../../bridge/infini/rt.hpp" - -#include "../../utils.hpp" - -namespace infinicore { -DevicePinnedHostAllocator::DevicePinnedHostAllocator(Device device) : MemoryAllocator(), owner_(device) {} - -DevicePinnedHostAllocator::~DevicePinnedHostAllocator() { - gc(); -} - -std::byte *DevicePinnedHostAllocator::allocate(size_t size) { - if (size == 0) { - return nullptr; - } - void *ptr; - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::MallocHost(&ptr, size))); - return (std::byte *)ptr; -} - -void DevicePinnedHostAllocator::deallocate(std::byte *ptr) { - if (ptr == nullptr) { - return; - } - if (owner_ == context::getDevice()) { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::FreeHost(ptr))); - gc(); - } else { - gc_queue_.push(ptr); - } -} - -void DevicePinnedHostAllocator::gc() { - while (gc_queue_.empty() == false) { - std::byte *p = gc_queue_.front(); - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::FreeHost(p))); - gc_queue_.pop(); - } -} - -} // namespace infinicore diff --git a/src/infinicore/context/allocators/device_pinned_allocator.hpp b/src/infinicore/context/allocators/device_pinned_allocator.hpp deleted file mode 100644 index 47a630b95..000000000 --- a/src/infinicore/context/allocators/device_pinned_allocator.hpp +++ /dev/null @@ -1,27 +0,0 @@ -#pragma once - -#include "memory_allocator.hpp" - -#include "../context_impl.hpp" - -#include - -namespace infinicore { -class DevicePinnedHostAllocator : public MemoryAllocator { -public: - explicit DevicePinnedHostAllocator(Device device); - ~DevicePinnedHostAllocator(); - - std::byte *allocate(size_t size) override; - void deallocate(std::byte *ptr) override; - - void gc(); - -private: - Device owner_; - - /// TODO: this is not thread-safe - std::queue gc_queue_; -}; - -} // namespace infinicore diff --git a/src/infinicore/context/allocators/host_allocator.cc b/src/infinicore/context/allocators/host_allocator.cc deleted file mode 100644 index 0f1cc7072..000000000 --- a/src/infinicore/context/allocators/host_allocator.cc +++ /dev/null @@ -1,18 +0,0 @@ -#include "host_allocator.hpp" - -namespace infinicore { -std::byte *HostAllocator::allocate(size_t size) { - if (size == 0) { - return nullptr; - } - return (std::byte *)std::malloc(size); -} - -void HostAllocator::deallocate(std::byte *ptr) { - if (ptr == nullptr) { - return; - } - std::free(ptr); -} - -} // namespace infinicore diff --git a/src/infinicore/context/allocators/host_allocator.hpp b/src/infinicore/context/allocators/host_allocator.hpp deleted file mode 100644 index 57cc00898..000000000 --- a/src/infinicore/context/allocators/host_allocator.hpp +++ /dev/null @@ -1,15 +0,0 @@ -#pragma once - -#include "memory_allocator.hpp" - -namespace infinicore { -class HostAllocator : public MemoryAllocator { -public: - HostAllocator() = default; - ~HostAllocator() = default; - - std::byte *allocate(size_t size) override; - void deallocate(std::byte *ptr) override; -}; - -} // namespace infinicore diff --git a/src/infinicore/context/allocators/memory_allocator.hpp b/src/infinicore/context/allocators/memory_allocator.hpp deleted file mode 100644 index 7fce31460..000000000 --- a/src/infinicore/context/allocators/memory_allocator.hpp +++ /dev/null @@ -1,15 +0,0 @@ -#pragma once - -#include "infinicore/memory.hpp" - -#include - -namespace infinicore { -class MemoryAllocator { -public: - virtual ~MemoryAllocator() = default; - - virtual std::byte *allocate(size_t size) = 0; - virtual void deallocate(std::byte *ptr) = 0; -}; -} // namespace infinicore diff --git a/src/infinicore/context/allocators/pinnable_block_allocator.cc b/src/infinicore/context/allocators/pinnable_block_allocator.cc deleted file mode 100644 index 3c70376f0..000000000 --- a/src/infinicore/context/allocators/pinnable_block_allocator.cc +++ /dev/null @@ -1,206 +0,0 @@ -#include "pinnable_block_allocator.hpp" - -#include "../context_impl.hpp" - -#include "../../utils.hpp" - -#include "../../../bridge/infini/rt.hpp" - -#include -#include - -namespace infinicore { - -// ------------------- Helper functions ------------------- - -// Round up size to nearest multiple of alignment -inline size_t align_up(size_t size, size_t alignment) { - return (size + alignment - 1) / alignment * alignment; -} - -// ------------------- Constructor ------------------- -PinnableBlockAllocator::PinnableBlockAllocator(Device device) - : device_(device) { - size_classes_ = { - {32 * 1024, {}}, // 32 KB - {256 * 1024, {}}, // 256 KB - {1 * 1024 * 1024, {}}, // 1 MB - {2 * 1024 * 1024, {}}, // 2 MB - {4 * 1024 * 1024, {}}, // 4 MB - {8 * 1024 * 1024, {}}, // 8 MB - {16 * 1024 * 1024, {}}, // 16 MB - {32 * 1024 * 1024, {}}, // 32 MB - {64 * 1024 * 1024, {}}, // 64 MB - {128 * 1024 * 1024, {}}, // 128 MB - {256 * 1024 * 1024, {}}, // 256 MB - }; -} - -// ------------------- allocate ------------------- -std::byte *PinnableBlockAllocator::allocate(size_t size) { - if (size == 0) { - return nullptr; - } - std::lock_guard lock(mutex_); - - // Align size to 256 bytes for GPU - size = align_up(size, 256); - - std::shared_ptr block; - - // 1. Try size-class allocation for small/medium - for (auto &cls : size_classes_) { - if (size <= cls.block_size) { - if (!cls.free_blocks.empty()) { - block = cls.free_blocks.back(); - while (block != nullptr && block->in_use) { - cls.free_blocks.pop_back(); - if (cls.free_blocks.empty()) { - block = nullptr; - break; - } - block = cls.free_blocks.back(); - } - if (block != nullptr) { - cls.free_blocks.pop_back(); - block->in_use = true; - block->use_count = 1; - return reinterpret_cast(block->ptr); - } - } - // Allocate a new block for this class - block = std::make_shared(); - block->size = cls.block_size; - block->frozen = pinned_mode_; - block->in_use = true; - block->use_count = 1; - - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Malloc(&block->ptr, block->size))); - - all_blocks_[block->ptr] = block; - return reinterpret_cast(block->ptr); - } - } - - // 2. Large block allocation - // Try to reuse a frozen or free large block - auto it = std::find_if(large_blocks_.begin(), large_blocks_.end(), - [size](const std::shared_ptr &b) { return b->size >= size && !b->in_use; }); - - if (it != large_blocks_.end()) { - block = *it; - block->in_use = true; - block->use_count = 1; - block->frozen = block->frozen || pinned_mode_; - return reinterpret_cast(block->ptr); - } - - // Allocate new large block - block = std::make_shared(); - block->size = size; - block->frozen = pinned_mode_; - block->in_use = true; - block->use_count = 1; - - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Malloc(&block->ptr, block->size))); - - large_blocks_.push_back(block); - all_blocks_[block->ptr] = block; - - return reinterpret_cast(block->ptr); -} - -// ------------------- deallocate ------------------- -void PinnableBlockAllocator::deallocate(std::byte *ptr) { - if (ptr == nullptr) { - return; - } - - std::lock_guard lock(mutex_); - - auto it = all_blocks_.find(reinterpret_cast(ptr)); - if (it == all_blocks_.end()) { - throw std::runtime_error("Pointer not allocated by this allocator"); - } - - auto block = it->second; - if (!block->in_use || block->use_count == 0) { - throw std::runtime_error("Double free detected in PinnableBlockAllocator"); - } - - --block->use_count; - if (block->use_count > 0) { - return; - } - - block->in_use = false; - for (auto &cls : size_classes_) { - if (block->size == cls.block_size) { - cls.free_blocks.push_back(block); - break; - } - } -} - -size_t PinnableBlockAllocator::mark_in_use_(void *ptr, bool in_use) { - std::lock_guard lock(mutex_); - - auto it = all_blocks_.find(reinterpret_cast(ptr)); - if (it == all_blocks_.end()) { - throw std::runtime_error("Pointer not allocated by this allocator"); - } - - auto block = it->second; - if (in_use) { - block->in_use = true; - ++block->use_count; - } else if (block->use_count > 0) { - --block->use_count; - block->in_use = block->use_count > 0; - } - return it->second->size; -} - -// ------------------- trim ------------------- -void PinnableBlockAllocator::trim() { - std::lock_guard lock(mutex_); - // Free non-frozen size-class blocks - for (auto &cls : size_classes_) { - for (auto it = cls.free_blocks.begin(); it != cls.free_blocks.end();) { - if (!(*it)->frozen) { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Free((*it)->ptr))); - all_blocks_.erase((*it)->ptr); - it = cls.free_blocks.erase(it); - } else { - ++it; - } - } - } - // Free non-frozen large blocks - for (auto it = large_blocks_.begin(); it != large_blocks_.end();) { - if (!(*it)->frozen && !(*it)->in_use) { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Free((*it)->ptr))); - all_blocks_.erase((*it)->ptr); - it = large_blocks_.erase(it); - } else { - ++it; - } - } -} - -// ------------------- Destructor ------------------- -PinnableBlockAllocator::~PinnableBlockAllocator() { - std::lock_guard lock(mutex_); - for (auto &p : all_blocks_) { - if (p.second->ptr) { - (void)infini::rt::runtime::Free(p.second->ptr); - } - } - all_blocks_.clear(); - large_blocks_.clear(); - for (auto &cls : size_classes_) { - cls.free_blocks.clear(); - } -} - -} // namespace infinicore diff --git a/src/infinicore/context/allocators/pinnable_block_allocator.hpp b/src/infinicore/context/allocators/pinnable_block_allocator.hpp deleted file mode 100644 index f506fcffd..000000000 --- a/src/infinicore/context/allocators/pinnable_block_allocator.hpp +++ /dev/null @@ -1,55 +0,0 @@ -#pragma once - -#include "memory_allocator.hpp" - -#include -#include -#include - -namespace infinicore { -class PinnableBlockAllocator : public MemoryAllocator { - // Represents a single memory block - struct Block { - void *ptr = nullptr; // Device pointer - size_t size = 0; // Block size in bytes - bool frozen = false; // True if used in pinned/graph mode - bool in_use = false; // Wether the block is currently in use - size_t use_count = 0; // Number of Memory owners for this block - }; - - // A simple size-class allocator for small/medium blocks - struct SizeClass { - size_t block_size; // Fixed size for this class - std::vector> free_blocks; - }; - -public: - PinnableBlockAllocator(Device device); - ~PinnableBlockAllocator(); - - std::byte *allocate(size_t size) override; - void deallocate(std::byte *ptr) override; - - // Switch pinned/graph mode - void set_pin_mode(bool pinned) { pinned_mode_ = pinned; } - - // internal use only, force set in_use flag for a mem block - // return the size of the block - size_t mark_in_use_(void *ptr, bool in_use); - - // trim cached blocks back to GPU (not pinned) - void trim(); - -private: - Device device_; - - bool pinned_mode_ = false; - - std::vector size_classes_; - std::vector> large_blocks_; - std::unordered_map> all_blocks_; - - std::mutex mutex_; // Thread safety -}; - -} // namespace infinicore diff --git a/src/infinicore/context/allocators/stream_ordered_allocator.cc b/src/infinicore/context/allocators/stream_ordered_allocator.cc deleted file mode 100644 index c6c54e47d..000000000 --- a/src/infinicore/context/allocators/stream_ordered_allocator.cc +++ /dev/null @@ -1,38 +0,0 @@ -#include "stream_ordered_allocator.hpp" - -#include "../../../bridge/infini/rt.hpp" - -#include "../../utils.hpp" - -namespace infinicore { -StreamOrderedAllocator::StreamOrderedAllocator(Device device) : MemoryAllocator(), device_(device) {} - -std::byte *StreamOrderedAllocator::allocate(size_t size) { - if (size == 0) { - return nullptr; - } - void *ptr = nullptr; - if (device_.getType() != Device::Type::CPU) { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::MallocAsync( - &ptr, - size, - bridge::infini::rt::translate_to(context::getStream())))); - } else { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Malloc(&ptr, size))); - } - return (std::byte *)ptr; -} - -void StreamOrderedAllocator::deallocate(std::byte *ptr) { - if (ptr == nullptr) { - return; - } - if (device_.getType() != Device::Type::CPU) { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::FreeAsync( - ptr, - bridge::infini::rt::translate_to(context::getStream())))); - } else { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Free(ptr))); - } -} -} // namespace infinicore diff --git a/src/infinicore/context/allocators/stream_ordered_allocator.hpp b/src/infinicore/context/allocators/stream_ordered_allocator.hpp deleted file mode 100644 index ed2fdd6e6..000000000 --- a/src/infinicore/context/allocators/stream_ordered_allocator.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include "memory_allocator.hpp" - -#include "../context_impl.hpp" - -namespace infinicore { -class StreamOrderedAllocator : public MemoryAllocator { -public: - explicit StreamOrderedAllocator(Device device); - ~StreamOrderedAllocator() = default; - - std::byte *allocate(size_t size) override; - void deallocate(std::byte *ptr) override; - -private: - Device device_; -}; - -} // namespace infinicore diff --git a/src/infinicore/context/context_impl.cc b/src/infinicore/context/context_impl.cc deleted file mode 100644 index 29a470679..000000000 --- a/src/infinicore/context/context_impl.cc +++ /dev/null @@ -1,233 +0,0 @@ -#include "context_impl.hpp" -#include "../../bridge/infini/rt.hpp" -#include "internal.hpp" - -#include "../utils.hpp" - -namespace infinicore { - -thread_local Runtime *ContextImpl::current_runtime_ = nullptr; - -Runtime *ContextImpl::getCurrentRuntime() { - if (current_runtime_ == nullptr) { - spdlog::debug("current_runtime_ is null, performing lazy initialization"); - // Lazy initialization: use the first available runtime - // Try to find the first non-CPU device, fallback to CPU - for (int i = int(Device::Type::COUNT) - 1; i > 0; i--) { - if (!runtime_table_[i].empty() && runtime_table_[i][0] != nullptr) { - current_runtime_ = runtime_table_[i][0].get()->activate(); - spdlog::debug("Lazy init: Set current_runtime_ to {} (ptr={})", current_runtime_->device().toString(), static_cast(current_runtime_)); - return current_runtime_; - } - } - // Fallback to CPU runtime - if (!runtime_table_[0].empty() && runtime_table_[0][0] != nullptr) { - current_runtime_ = runtime_table_[0][0].get()->activate(); - spdlog::debug("Lazy init: Set current_runtime_ to {} (ptr={})", current_runtime_->device().toString(), static_cast(current_runtime_)); - } - } else { - spdlog::debug("getCurrentRuntime() returning {} (ptr={})", current_runtime_->device().toString(), static_cast(current_runtime_)); - } - return current_runtime_; -} - -void ContextImpl::setDevice(Device device) { - if (device == getCurrentRuntime()->device()) { - // Do nothing if the device is already set. - return; - } - - thread_local bool warn_switch_runtime = false; - if (getCurrentRuntime()->isGraphRecording() && !warn_switch_runtime) { - spdlog::warn("Switching device runtime during graph recording may break the graph!"); - warn_switch_runtime = true; - } - - if (runtime_table_[int(device.getType())][device.getIndex()] == nullptr) { - // Lazy initialization of runtime if never set before. - runtime_table_[int(device.getType())][device.getIndex()] = std::unique_ptr(new Runtime(device)); - current_runtime_ = runtime_table_[int(device.getType())][device.getIndex()].get(); - } else { - current_runtime_ = runtime_table_[int(device.getType())][device.getIndex()].get()->activate(); - } -} - -size_t ContextImpl::getDeviceCount(Device::Type type) { - return runtime_table_[int(type)].size(); -} - -ContextImpl &ContextImpl::singleton() { - static ContextImpl instance; - return instance; -} - -template -void ContextImpl::initializeRuntime() { - constexpr auto runtime_device_type = bridge::infini::rt::translate_to( - static_cast(device_type)); - static_assert(runtime_device_type != infini::rt::Device::Type::kCount, - "InfiniCore device does not map to an InfiniRT runtime"); - - if constexpr (infini::rt::DeviceEnabled::value) { - constexpr auto device_index = static_cast(device_type); - infini::rt::set_runtime_device_type(runtime_device_type); - - int device_count = 0; - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate( - infini::rt::runtime::GetDeviceCount(&device_count))); - - auto &runtimes = runtime_table_[device_index]; - runtimes.resize(device_count); - if (device_count > 0) { - runtimes[0] = std::unique_ptr(new Runtime(Device(device_type, 0))); - if constexpr (device_type != Device::Type::CPU) { - current_runtime_ = runtimes[0].get(); - } - } - } -} - -ContextImpl::ContextImpl() { - initializeRuntime(); - initializeRuntime(); - initializeRuntime(); - initializeRuntime(); - initializeRuntime(); - - if (current_runtime_ == nullptr && !runtime_table_[static_cast(Device::Type::CPU)].empty()) { - current_runtime_ = runtime_table_[static_cast(Device::Type::CPU)][0].get(); - } -} - -namespace context { - -void setDevice(Device device) { - ContextImpl::singleton().setDevice(device); -} - -Device getDevice() { - return ContextImpl::singleton().getCurrentRuntime()->device(); -} - -size_t getDeviceCount(Device::Type type) { - return ContextImpl::singleton().getDeviceCount(type); -} - -infinirtStream_t getStream() { - return ContextImpl::singleton().getCurrentRuntime()->stream(); -} - -infiniopHandle_t getInfiniopHandle(Device device) { - if (device != getDevice()) { - setDevice(device); - } - return ContextImpl::singleton().getCurrentRuntime()->infiniopHandle(); -} - -void syncStream() { - return ContextImpl::singleton().getCurrentRuntime()->syncStream(); -} - -void syncDevice() { - return ContextImpl::singleton().getCurrentRuntime()->syncDevice(); -} - -void trimMemory() { - return ContextImpl::singleton().getCurrentRuntime()->trimMemory(); -} - -std::shared_ptr allocateMemory(size_t size) { - return ContextImpl::singleton().getCurrentRuntime()->allocateMemory(size); -} - -std::shared_ptr allocateHostMemory(size_t size) { - setDevice(Device::cpu()); - return allocateMemory(size); -} - -std::shared_ptr allocatePinnedHostMemory(size_t size) { - return ContextImpl::singleton().getCurrentRuntime()->allocatePinnedHostMemory(size); -} - -void memcpyH2D(void *dst, const void *src, size_t size, bool async) { - return ContextImpl::singleton().getCurrentRuntime()->memcpyH2D(dst, src, size, async); -} - -void memcpyD2H(void *dst, const void *src, size_t size) { - return ContextImpl::singleton().getCurrentRuntime()->memcpyD2H(dst, src, size); -} - -void memcpyD2D(void *dst, const void *src, size_t size, bool async) { - return ContextImpl::singleton().getCurrentRuntime()->memcpyD2D(dst, src, size, async); -} - -void memcpyH2H(void *dst, const void *src, size_t size) { - setDevice(Device::cpu()); - return ContextImpl::singleton().getCurrentRuntime()->memcpyD2D(dst, src, size); -} - -void setDeviceMemory(void *ptr, int value, size_t count) { - return ContextImpl::singleton().getCurrentRuntime()->setDeviceMemory(ptr, value, count); -} - -void setDeviceMemoryAsync(void *ptr, int value, size_t count, infinirtStream_t stream) { - return ContextImpl::singleton().getCurrentRuntime()->setDeviceMemoryAsync(ptr, value, count, stream); -} - -// Timing API implementations -infinirtEvent_t createEvent() { - return ContextImpl::singleton().getCurrentRuntime()->createEvent(); -} - -infinirtEvent_t createEventWithFlags(uint32_t flags) { - return ContextImpl::singleton().getCurrentRuntime()->createEventWithFlags(flags); -} - -void recordEvent(infinirtEvent_t event, infinirtStream_t stream) { - ContextImpl::singleton().getCurrentRuntime()->recordEvent(event, stream); -} - -bool queryEvent(infinirtEvent_t event) { - return ContextImpl::singleton().getCurrentRuntime()->queryEvent(event); -} - -void synchronizeEvent(infinirtEvent_t event) { - ContextImpl::singleton().getCurrentRuntime()->synchronizeEvent(event); -} - -void destroyEvent(infinirtEvent_t event) { - ContextImpl::singleton().getCurrentRuntime()->destroyEvent(event); -} - -float elapsedTime(infinirtEvent_t start, infinirtEvent_t end) { - return ContextImpl::singleton().getCurrentRuntime()->elapsedTime(start, end); -} - -void streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event) { - ContextImpl::singleton().getCurrentRuntime()->streamWaitEvent(stream, event); -} - -bool isGraphRecording() { - return ContextImpl::singleton().getCurrentRuntime()->isGraphRecording(); -} - -void startGraphRecording() { - ContextImpl::singleton().getCurrentRuntime()->startGraphRecording(); -} - -void addGraphOperator(std::shared_ptr op) { - ContextImpl::singleton().getCurrentRuntime()->addGraphOperator(op); -} - -std::shared_ptr stopGraphRecording() { - return ContextImpl::singleton().getCurrentRuntime()->stopGraphRecording(); -} - -std::shared_ptr reinstantiateBlob(std::shared_ptr blob) { - setDevice(blob->device()); - return ContextImpl::singleton().getCurrentRuntime()->reinstantiateBlob(blob); -} - -} // namespace context - -} // namespace infinicore diff --git a/src/infinicore/context/context_impl.hpp b/src/infinicore/context/context_impl.hpp deleted file mode 100644 index b05a56f5d..000000000 --- a/src/infinicore/context/context_impl.hpp +++ /dev/null @@ -1,33 +0,0 @@ -#pragma once -#include "infinicore/context/context.hpp" -#include "runtime/runtime.hpp" - -#include -#include - -namespace infinicore { -class ContextImpl { -private: - template - void initializeRuntime(); - - // Table of runtimes for every device (type and index) - std::array>, size_t(Device::Type::COUNT)> runtime_table_; - // Active runtime for current thread. Can use "static thread local" because context is a process singleton. - static thread_local Runtime *current_runtime_; - -protected: - ContextImpl(); - -public: - Runtime *getCurrentRuntime(); - - void setDevice(Device); - - size_t getDeviceCount(Device::Type type); - - static ContextImpl &singleton(); - - friend class Runtime; -}; -} // namespace infinicore diff --git a/src/infinicore/context/internal.hpp b/src/infinicore/context/internal.hpp deleted file mode 100644 index aeecaff51..000000000 --- a/src/infinicore/context/internal.hpp +++ /dev/null @@ -1,10 +0,0 @@ -#pragma once - -#include "infinicore/device.hpp" -#include "infinicore/memory.hpp" - -#include "infinicore/graph/graph.hpp" - -namespace infinicore::context { -std::shared_ptr reinstantiateBlob(std::shared_ptr blob); -}; diff --git a/src/infinicore/context/runtime/runtime.cc b/src/infinicore/context/runtime/runtime.cc deleted file mode 100644 index 0d5cd6de8..000000000 --- a/src/infinicore/context/runtime/runtime.cc +++ /dev/null @@ -1,212 +0,0 @@ -#include "runtime.hpp" -#include "../../../bridge/infini/rt.hpp" - -#include "../../utils.hpp" - -#include "../allocators/device_pinned_allocator.hpp" -#include "../allocators/host_allocator.hpp" -#include "../allocators/pinnable_block_allocator.hpp" -#include "../allocators/stream_ordered_allocator.hpp" - -namespace infinicore { -Runtime::Runtime(Device device) : device_(device), graph_manager_(std::make_unique()) { - activate(); - infini::rt::runtime::Stream stream = nullptr; - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::StreamCreate(&stream))); - stream_ = bridge::infini::rt::translate_from(stream); - INFINICORE_CHECK_ERROR(infiniopCreateHandle(&infiniop_handle_)); - if (device_.getType() == Device::Type::CPU) { - device_memory_allocator_ = std::make_unique(device); - } else { - device_memory_allocator_ = std::make_unique(device); - pinned_host_memory_allocator_ = std::make_unique(device); - } -} -Runtime::~Runtime() { - activate(); - if (pinned_host_memory_allocator_) { - pinned_host_memory_allocator_.reset(); - } - device_memory_allocator_.reset(); - infiniopDestroyHandle(infiniop_handle_); - (void)infini::rt::runtime::StreamDestroy(bridge::infini::rt::translate_to(stream_)); -} - -Runtime *Runtime::activate() { - auto rt_device = bridge::infini::rt::translate_to(static_cast(device_.getType())); - INFINICORE_ASSERT(rt_device != infini::rt::Device::Type::kCount); - infini::rt::set_runtime_device_type(rt_device); - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::SetDevice(static_cast(device_.getIndex())))); - return this; -} - -Device Runtime::device() const { - return device_; -} - -infinirtStream_t Runtime::stream() const { - return stream_; -} - -infiniopHandle_t Runtime::infiniopHandle() const { - return infiniop_handle_; -} - -void Runtime::syncStream() { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::StreamSynchronize(bridge::infini::rt::translate_to(stream_)))); -} - -void Runtime::syncDevice() { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::DeviceSynchronize())); -} - -void Runtime::trimMemory() { - device_memory_allocator_->trim(); -} - -std::shared_ptr Runtime::allocateMemory(size_t size) { - std::byte *data_ptr = device_memory_allocator_->allocate(size); - return std::make_shared( - data_ptr, size, device_, - [alloc = device_memory_allocator_.get()](std::byte *p) { - alloc->deallocate(p); - }); -} - -std::shared_ptr Runtime::allocatePinnedHostMemory(size_t size) { - if (!pinned_host_memory_allocator_) { - spdlog::warn("For CPU devices, pinned memory is not supported, falling back to regular host memory"); - return allocateMemory(size); - } - std::byte *data_ptr = pinned_host_memory_allocator_->allocate(size); - return std::make_shared( - data_ptr, size, device_, - [alloc = pinned_host_memory_allocator_.get()](std::byte *p) { - alloc->deallocate(p); - }, - true); -} - -std::shared_ptr Runtime::reinstantiateBlob(std::shared_ptr blob) { - std::lock_guard lock(reinstantiated_blob_mutex_); - - auto ptr = blob->data(); - auto it = reinstantiated_blobs_.find(ptr); - if (it != reinstantiated_blobs_.end()) { - if (auto memory = it->second.lock()) { - return memory; - } - } - - device_memory_allocator_.get()->mark_in_use_(ptr, true); - auto memory = std::make_shared( - blob->data(), blob->size(), device_, - [alloc = device_memory_allocator_.get()](std::byte *p) { - alloc->deallocate(p); - }); - reinstantiated_blobs_[ptr] = memory; - return memory; -} - -void Runtime::memcpyH2D(void *dst, const void *src, size_t size, bool async) { - if (async && device_.getType() != Device::Type::CPU) { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::MemcpyAsync(dst, src, size, infini::rt::runtime::kMemcpyHostToDevice, bridge::infini::rt::translate_to(stream_)))); - } else { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Memcpy(dst, src, size, infini::rt::runtime::kMemcpyHostToDevice))); - } -} - -void Runtime::memcpyD2H(void *dst, const void *src, size_t size) { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Memcpy(dst, src, size, infini::rt::runtime::kMemcpyDeviceToHost))); -} - -void Runtime::memcpyD2D(void *dst, const void *src, size_t size, bool async) { - if (async && device_.getType() != Device::Type::CPU) { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::MemcpyAsync(dst, src, size, infini::rt::runtime::kMemcpyDeviceToDevice, bridge::infini::rt::translate_to(stream_)))); - } else { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Memcpy(dst, src, size, infini::rt::runtime::kMemcpyDeviceToDevice))); - } -} - -void Runtime::setDeviceMemory(void *ptr, int value, size_t count) { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Memset(ptr, value, count))); -} - -void Runtime::setDeviceMemoryAsync(void *ptr, int value, size_t count, infinirtStream_t stream) { - if (device_.getType() != Device::Type::CPU) { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::MemsetAsync(ptr, value, count, bridge::infini::rt::translate_to(stream)))); - } else { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::Memset(ptr, value, count))); - } -} - -// Timing method implementations -infinirtEvent_t Runtime::createEvent() { - infinirtEvent_t event; - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::EventCreate(&event))); - return event; -} - -infinirtEvent_t Runtime::createEventWithFlags(uint32_t flags) { - infinirtEvent_t event; - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::EventCreateWithFlags(&event, flags))); - return event; -} - -void Runtime::recordEvent(infinirtEvent_t event, infinirtStream_t stream) { - if (stream == nullptr) { - stream = stream_; - } - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::EventRecord(event, bridge::infini::rt::translate_to(stream)))); -} - -bool Runtime::queryEvent(infinirtEvent_t event) { - return bridge::infini::rt::translate(infini::rt::runtime::EventQuery(event)) == INFINI_STATUS_SUCCESS; -} - -void Runtime::synchronizeEvent(infinirtEvent_t event) { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::EventSynchronize(event))); -} - -void Runtime::destroyEvent(infinirtEvent_t event) { - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::EventDestroy(event))); -} - -float Runtime::elapsedTime(infinirtEvent_t start, infinirtEvent_t end) { - float ms; - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::EventElapsedTime(&ms, start, end))); - return ms; -} - -void Runtime::streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event) { - // Use current stream if no specific stream is provided - if (stream == nullptr) { - stream = stream_; - } - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(infini::rt::runtime::StreamWaitEvent(bridge::infini::rt::translate_to(stream), event, 0))); -} - -bool Runtime::isGraphRecording() const { - return graph_manager_->is_recording(); -} - -void Runtime::startGraphRecording() { - device_memory_allocator_->set_pin_mode(true); - return graph_manager_->start_recording(); -} - -void Runtime::addGraphOperator(std::shared_ptr op) { - return graph_manager_->add_operator(op); -} - -std::shared_ptr Runtime::stopGraphRecording() { - auto graph = graph_manager_->stop_recording(); - device_memory_allocator_->set_pin_mode(false); - return graph; -} - -std::string Runtime::toString() const { - return fmt::format("Runtime({})", device_.toString()); -} - -} // namespace infinicore diff --git a/src/infinicore/context/runtime/runtime.hpp b/src/infinicore/context/runtime/runtime.hpp deleted file mode 100644 index 2ed382ab7..000000000 --- a/src/infinicore/context/runtime/runtime.hpp +++ /dev/null @@ -1,74 +0,0 @@ -#pragma once - -#include "../allocators/pinnable_block_allocator.hpp" - -#include "infinicore/context/context.hpp" - -#include "../../graph/graph_manager.hpp" - -#include -#include -#include -#include - -namespace infinicore { -class ContextImpl; -class Runtime { -private: - Device device_; - infinirtStream_t stream_; - infiniopHandle_t infiniop_handle_; - std::unique_ptr device_memory_allocator_; - std::unique_ptr pinned_host_memory_allocator_; - std::unique_ptr graph_manager_; - std::mutex reinstantiated_blob_mutex_; - std::unordered_map> reinstantiated_blobs_; - -protected: - Runtime(Device device); - -public: - ~Runtime(); - - Runtime *activate(); - - Device device() const; - infinirtStream_t stream() const; - infiniopHandle_t infiniopHandle() const; - - void syncStream(); - void syncDevice(); - void trimMemory(); - - std::shared_ptr allocateMemory(size_t size); - std::shared_ptr allocatePinnedHostMemory(size_t size); - std::shared_ptr reinstantiateBlob(std::shared_ptr blob); - - void memcpyH2D(void *dst, const void *src, size_t size, bool async = true); - void memcpyD2H(void *dst, const void *src, size_t size); - void memcpyD2D(void *dst, const void *src, size_t size, bool async = true); - - void setDeviceMemory(void *ptr, int value, size_t count); - void setDeviceMemoryAsync(void *ptr, int value, size_t count, infinirtStream_t stream); - - // Timing methods - infinirtEvent_t createEvent(); - infinirtEvent_t createEventWithFlags(uint32_t flags); - void recordEvent(infinirtEvent_t event, infinirtStream_t stream = nullptr); - bool queryEvent(infinirtEvent_t event); - void synchronizeEvent(infinirtEvent_t event); - void destroyEvent(infinirtEvent_t event); - float elapsedTime(infinirtEvent_t start, infinirtEvent_t end); - void streamWaitEvent(infinirtStream_t stream, infinirtEvent_t event); - - // Graph - bool isGraphRecording() const; - void startGraphRecording(); - void addGraphOperator(std::shared_ptr op); - std::shared_ptr stopGraphRecording(); - - std::string toString() const; - - friend class ContextImpl; -}; -} // namespace infinicore diff --git a/src/infinicore/device.cc b/src/infinicore/device.cc deleted file mode 100644 index 77c1f4986..000000000 --- a/src/infinicore/device.cc +++ /dev/null @@ -1,60 +0,0 @@ -#include -#include - -#include "infinicore.hpp" - -namespace infinicore { - -Device::Device(const Type &type, const Index &index) : type_{type}, index_{index} {} - -const Device::Type &Device::getType() const { - return type_; -} - -const Device::Index &Device::getIndex() const { - return index_; -} - -std::string Device::toString() const { - return toString(type_) + ":" + std::to_string(index_); -} - -std::string Device::toString(const Type &type) { - switch (type) { - case Type::CPU: - return "CPU"; - case Type::NVIDIA: - return "NVIDIA"; - case Type::CAMBRICON: - return "CAMBRICON"; - case Type::ASCEND: - return "ASCEND"; - case Type::METAX: - return "METAX"; - case Type::MOORE: - return "MOORE"; - case Type::ILUVATAR: - return "ILUVATAR"; - case Type::QY: - return "QY"; - case Type::KUNLUN: - return "KUNLUN"; - case Type::HYGON: - return "HYGON"; - case Type::ALI: - return "ALI"; - case Type::COUNT: - return "COUNT"; - default: - return "UNKNOWN"; - } -} - -bool Device::operator==(const Device &other) const { - return type_ == other.type_ && index_ == other.index_; -} - -bool Device::operator!=(const Device &other) const { - return type_ != other.type_ || index_ != other.index_; -} -} // namespace infinicore diff --git a/src/infinicore/device_event.cc b/src/infinicore/device_event.cc deleted file mode 100644 index 347e7effd..000000000 --- a/src/infinicore/device_event.cc +++ /dev/null @@ -1,180 +0,0 @@ -#include "infinicore.hpp" - -namespace infinicore { - -DeviceEvent::DeviceEvent() - : device_(context::getDevice()), is_recorded_(false) { - event_ = context::createEvent(); -} - -DeviceEvent::DeviceEvent(uint32_t flags) - : device_(context::getDevice()), is_recorded_(false) { - event_ = context::createEventWithFlags(flags); -} - -DeviceEvent::DeviceEvent(Device device) - : device_(device), is_recorded_(false) { - // Switch to target device for event creation - Device current_device = context::getDevice(); - context::setDevice(device_); - event_ = context::createEvent(); - // Restore original device - context::setDevice(current_device); -} - -DeviceEvent::DeviceEvent(Device device, uint32_t flags) - : device_(device), is_recorded_(false) { - // Switch to target device for event creation - Device current_device = context::getDevice(); - context::setDevice(device_); - event_ = context::createEventWithFlags(flags); - // Restore original device - context::setDevice(current_device); -} - -DeviceEvent::DeviceEvent(DeviceEvent &&other) noexcept - : event_(other.event_), device_(other.device_), is_recorded_(other.is_recorded_) { - other.event_ = nullptr; - other.is_recorded_ = false; -} - -DeviceEvent &DeviceEvent::operator=(DeviceEvent &&other) noexcept { - if (this != &other) { - // Clean up current resources - if (event_ != nullptr) { - context::destroyEvent(event_); - } - - // Transfer ownership - event_ = other.event_; - device_ = other.device_; - is_recorded_ = other.is_recorded_; - - // Reset source - other.event_ = nullptr; - other.is_recorded_ = false; - } - return *this; -} - -DeviceEvent::~DeviceEvent() { - if (event_ != nullptr) { - context::destroyEvent(event_); - } -} - -void DeviceEvent::record() { - Device current_device = context::getDevice(); - - // Ensure we're on the correct device - if (current_device != device_) { - context::setDevice(device_); - } - - context::recordEvent(event_); - is_recorded_ = true; - - // Restore original device if we changed it - if (current_device != device_) { - context::setDevice(current_device); - } -} - -void DeviceEvent::record(infinirtStream_t stream) { - Device current_device = context::getDevice(); - - // Ensure we're on the correct device - if (current_device != device_) { - context::setDevice(device_); - } - - context::recordEvent(event_, stream); - is_recorded_ = true; - - // Restore original device if we changed it - if (current_device != device_) { - context::setDevice(current_device); - } -} - -void DeviceEvent::synchronize() { - Device current_device = context::getDevice(); - - // Ensure we're on the correct device - if (current_device != device_) { - context::setDevice(device_); - } - - context::synchronizeEvent(event_); - - // Restore original device if we changed it - if (current_device != device_) { - context::setDevice(current_device); - } -} - -bool DeviceEvent::query() const { - Device current_device = context::getDevice(); - bool result = false; - - // Ensure we're on the correct device - if (current_device != device_) { - context::setDevice(device_); - } - - result = context::queryEvent(event_); - - // Restore original device if we changed it - if (current_device != device_) { - context::setDevice(current_device); - } - - return result; -} - -float DeviceEvent::elapsed_time(const DeviceEvent &other) const { - // Both events must be on the same device - if (device_ != other.device_) { - throw std::runtime_error("Cannot measure elapsed time between events on different devices"); - } - - // Both events must be recorded - if (!is_recorded_ || !other.is_recorded_) { - throw std::runtime_error("Both events must be recorded before measuring elapsed time"); - } - - Device current_device = context::getDevice(); - - // Switch to the device where events reside - if (current_device != device_) { - context::setDevice(device_); - } - - float elapsed_ms = context::elapsedTime(event_, other.event_); - - // Restore original device if we changed it - if (current_device != device_) { - context::setDevice(current_device); - } - - return elapsed_ms; -} - -void DeviceEvent::wait(infinirtStream_t stream) const { - Device current_device = context::getDevice(); - - // Ensure we're on the correct device - if (current_device != device_) { - context::setDevice(device_); - } - - // Make the stream wait for this event - context::streamWaitEvent(stream, event_); - - // Restore original device if we changed it - if (current_device != device_) { - context::setDevice(current_device); - } -} - -} // namespace infinicore diff --git a/src/infinicore/dtype.cc b/src/infinicore/dtype.cc deleted file mode 100644 index 9672a7ff7..000000000 --- a/src/infinicore/dtype.cc +++ /dev/null @@ -1,83 +0,0 @@ -#include - -namespace infinicore { - -std::string toString(const DataType &dtype) { - switch (dtype) { - case DataType::BYTE: - return "uint8"; - case DataType::BOOL: - return "bool"; - case DataType::I8: - return "int8"; - case DataType::I16: - return "int16"; - case DataType::I32: - return "int32"; - case DataType::I64: - return "int64"; - case DataType::U8: - return "uint8"; - case DataType::U16: - return "uint16"; - case DataType::U32: - return "uint32"; - case DataType::U64: - return "uint64"; - case DataType::F8: - return "float8"; - case DataType::F16: - return "float16"; - case DataType::F32: - return "float32"; - case DataType::F64: - return "float64"; - case DataType::C16: - return "complex16"; - case DataType::C32: - return "complex32"; - case DataType::C64: - return "complex64"; - case DataType::C128: - return "complex128"; - case DataType::BF16: - return "bfloat16"; - } - - // TODO: Add error handling. - return ""; -} - -size_t dsize(const DataType &dtype) { - switch (dtype) { - case DataType::BYTE: - case DataType::BOOL: - case DataType::F8: - case DataType::I8: - case DataType::U8: - return 1; - case DataType::I16: - case DataType::U16: - case DataType::F16: - case DataType::BF16: - case DataType::C16: - return 2; - case DataType::I32: - case DataType::U32: - case DataType::F32: - case DataType::C32: - return 4; - case DataType::I64: - case DataType::U64: - case DataType::F64: - case DataType::C64: - return 8; - case DataType::C128: - return 16; - } - - // TODO: Add error handling. - return 0; -} - -} // namespace infinicore diff --git a/src/infinicore/graph/graph.cc b/src/infinicore/graph/graph.cc deleted file mode 100644 index 9b28aae7f..000000000 --- a/src/infinicore/graph/graph.cc +++ /dev/null @@ -1,192 +0,0 @@ -#include "graph_manager.hpp" - -#include "../../bridge/infini/rt.hpp" -#include "../utils.hpp" -#include "infinicore/context/context.hpp" - -#ifdef USE_INFINIRT_GRAPH -#include -#endif - -namespace infinicore::graph { - -#ifdef USE_INFINIRT_GRAPH -namespace rt_runtime = ::infini::rt::runtime; -#endif - -/* ========================= - * GraphTensor - * ========================= */ - -GraphTensor::GraphTensor(const Tensor &tensor) : Tensor(tensor->to_blob_()) { -} - -/* ========================= - * GraphOperator - * ========================= */ - -void DispatchableGraphOperator::run() const { - runner_(planned_meta_); -} - -DispatchableGraphOperator::~DispatchableGraphOperator() { - if (deleter_) { - deleter_(&planned_meta_); - } -} - -/* ========================= - * Graph - * ========================= */ - -#ifdef USE_INFINIRT_GRAPH -struct Graph::DeviceGraph { - rt_runtime::Graph graph = nullptr; - rt_runtime::GraphExec exec = nullptr; - rt_runtime::Stream stream = nullptr; - ::infini::rt::Device::Type device_type = ::infini::rt::Device::Type::kCount; - int device_index = 0; - - ~DeviceGraph() { - if (exec) { - (void)rt_runtime::GraphExecDestroy(exec); - } - if (graph) { - (void)rt_runtime::GraphDestroy(graph); - } - } - - void launch() { - ::infini::rt::set_runtime_device_type(device_type); - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(rt_runtime::SetDevice(device_index))); - INFINICORE_CHECK_ERROR(bridge::infini::rt::translate(rt_runtime::GraphLaunch(exec, stream))); - } -}; -#else -struct Graph::DeviceGraph {}; -#endif - -Graph::Graph() { -} - -void Graph::run() const { -#ifdef USE_INFINIRT_GRAPH - if (device_graph_ != nullptr && device_graph_.get()->exec != nullptr) { - device_graph_.get()->launch(); - return; - } -#endif - for (auto &op : op_list_) { - op->run(); - } -} - -void Graph::add_operator(std::shared_ptr op) { - op_list_.push_back(op); -} - -void Graph::instantiate() { -#ifdef USE_INFINIRT_GRAPH - // Reset device graph - device_graph_ = std::make_unique(); - auto current_device = context::getDevice(); - device_graph_->device_type = bridge::infini::rt::translate_to(static_cast(current_device.getType())); - device_graph_->device_index = static_cast(current_device.getIndex()); - device_graph_->stream = bridge::infini::rt::translate_to(context::getStream()); - if (device_graph_->device_type == ::infini::rt::Device::Type::kCount) { - spdlog::warn("InfiniRT graph runtime does not support the current device. Falling back to eager execution."); - device_graph_.reset(); - return; - } - ::infini::rt::set_runtime_device_type(device_graph_->device_type); - auto set_device_status = bridge::infini::rt::translate(rt_runtime::SetDevice(device_graph_->device_index)); - if (set_device_status != INFINI_STATUS_SUCCESS) { - spdlog::warn("InfiniRT graph runtime failed to select the current device. Falling back to eager execution."); - device_graph_.reset(); - return; - } - - // warmup - for (size_t iter = 0; iter < 5; ++iter) { - this->run(); - } - infinicore::context::syncStream(); - - auto begin_status = bridge::infini::rt::translate(rt_runtime::StreamBeginCapture( - device_graph_->stream, - rt_runtime::StreamCaptureMode::kStreamCaptureModeRelaxed)); - if (begin_status != INFINI_STATUS_SUCCESS) { - spdlog::warn("Fail to begin device graph capture."); - device_graph_.reset(); - return; - } - - // Run and record - this->run(); - - auto end_status = bridge::infini::rt::translate(rt_runtime::StreamEndCapture( - device_graph_->stream, - &device_graph_->graph)); - if (end_status != INFINI_STATUS_SUCCESS) { - spdlog::warn("Fail to end device graph capture."); - device_graph_.reset(); - return; - } - - auto instantiate_status = bridge::infini::rt::translate(rt_runtime::GraphInstantiate( - &device_graph_->exec, - device_graph_->graph)); - if (instantiate_status != INFINI_STATUS_SUCCESS) { - static bool warned_once = false; - if (!warned_once) { - warned_once = true; - spdlog::warn("Fail to instantiate device graph."); - } - device_graph_.reset(); - return; - } - static bool logged_once = false; - if (!logged_once) { - logged_once = true; - spdlog::info("Using InfiniRT C++ graph runtime API for graph capture and replay."); - } -#endif -} - -Graph::~Graph() = default; - -/* ========================= - * GraphManager - * ========================= */ - -bool GraphManager::is_recording() const { - return recording_; -} - -void GraphManager::start_recording() { - if (is_recording()) { - spdlog::warn("Graph is already recording. Previous recording will be dropped."); - } - recording_ = true; - graph_ = std::make_shared(); -} - -void GraphManager::add_operator(std::shared_ptr op) { - INFINICORE_ASSERT(is_recording()); - - graph_->add_operator(op); -} - -std::shared_ptr GraphManager::stop_recording() { - if (!is_recording()) { - spdlog::warn("Graph is not recording. Please start recording first."); - return nullptr; - } - recording_ = false; -#ifdef USE_INFINIRT_GRAPH - graph_->instantiate(); -#endif - return std::exchange(graph_, nullptr); -} - -} // namespace infinicore::graph diff --git a/src/infinicore/graph/graph_manager.hpp b/src/infinicore/graph/graph_manager.hpp deleted file mode 100644 index 9cc040366..000000000 --- a/src/infinicore/graph/graph_manager.hpp +++ /dev/null @@ -1,25 +0,0 @@ -#pragma once - -#include "infinicore/graph/graph.hpp" - -#include -#include - -namespace infinicore::graph { - -class GraphManager { -public: - GraphManager() = default; - ~GraphManager() = default; - - bool is_recording() const; - void start_recording(); - void add_operator(std::shared_ptr op); - std::shared_ptr stop_recording(); - -private: - std::shared_ptr graph_; - bool recording_ = false; -}; - -} // namespace infinicore::graph diff --git a/src/infinicore/io.cc b/src/infinicore/io.cc deleted file mode 100644 index b9ed577c4..000000000 --- a/src/infinicore/io.cc +++ /dev/null @@ -1,673 +0,0 @@ -/* - * This file contains some code derived from xtensor (https://github.com/xtensor-stack/xtensor), - * which is licensed under the BSD-3-Clause license. - */ - -#include "infinicore/io.hpp" -#include "../utils/custom_types.h" -#include "infinicore/context/context.hpp" -#include "infinicore/dtype.hpp" -#include "infinicore/tensor.hpp" -#include -#include -#include -#include -#include -#include -#include - -namespace { - -using Tensor = infinicore::Tensor; -using TensorSliceParams = infinicore::TensorSliceParams; -using DataType = infinicore::DataType; -using Device = infinicore::Device; -using TensorImpl = infinicore::TensorImpl; -using Size = infinicore::Size; - -/** - * @brief This function extracts a scalar or sub-tensor from a tensor using a vector of indexes. - */ -inline Tensor at_impl(const Tensor &tensor, const std::vector &indexes) { - if (indexes.size() > tensor->ndim()) { - throw std::runtime_error("at_impl:: Number of indexes (" + std::to_string(indexes.size()) + ") exceeds tensor dimensions (" + std::to_string(tensor->ndim()) + ")"); - } - - for (size_t i = 0; i < indexes.size(); i++) { - if (indexes[i] >= tensor->shape()[i]) { - throw std::runtime_error("at_impl :: Index " + std::to_string(indexes[i]) + " is out of bounds for dimension " + std::to_string(i)); - } - } - - std::vector slices; - slices.reserve(indexes.size()); - for (size_t i = 0; i < indexes.size(); i++) { - slices.push_back({i, indexes[i], 1}); - } - - Tensor result = tensor->narrow(slices); - for (size_t i = 0; i < indexes.size(); i++) { - result = result->squeeze(0); - } - - return result; -} - -template -Tensor at(const Tensor &tensor, Args... args) { - std::vector indexes = {static_cast(args)...}; - return at_impl(tensor, indexes); -} - -[[maybe_unused]] Tensor at(const Tensor &tensor, std::initializer_list indexes) { - std::vector indexes_vec(indexes.begin(), indexes.end()); - return at_impl(tensor, indexes_vec); -} - -Tensor at(const Tensor &tensor, const std::vector &indexes) { - return at_impl(tensor, indexes); -} - -/** - * @brief read a value from raw data pointer based on DataType. - */ -template -T item_impl(const std::byte *data, DataType dtype) { - switch (dtype) { - case DataType::F16: { - const fp16_t *ptr = reinterpret_cast(data); - float f = _f16_to_f32(ptr[0]); - return static_cast(f); - } - case DataType::BF16: { - const bf16_t *ptr = reinterpret_cast(data); - float f = _bf16_to_f32(ptr[0]); - return static_cast(f); - } - case DataType::BOOL: { - const bool *ptr = reinterpret_cast(data); - return static_cast(ptr[0] ? 1 : 0); - } - default: - break; - } - - const T *ptr = reinterpret_cast(data); - return ptr[0]; -} - -/** - * @brief Extracts a scalar value from a single-element tensor. - * The tensor must have exactly one element and must be located on CPU device. - * - * @code{.cpp} - * float value = item(tensor, dtype); // Extract as float - * @endcode - */ -template -T item(const Tensor &tensor) { - if (tensor->numel() != 1) { - throw std::runtime_error("item() can only be called on a tensor with exactly one element, but got " + std::to_string(tensor->numel()) + " elements"); - } - - if (tensor->device().getType() != Device::Type::CPU) { - throw std::runtime_error("item() can only be called on a CPU tensor, but got device: " + tensor->device().toString()); - } - - const std::byte *data = tensor->data(); - DataType dtype = tensor->dtype(); - return item_impl(data, dtype); -} -} // namespace - -namespace infinicore { -namespace print_options { - -template -class fmtflags_guard { -public: - explicit fmtflags_guard(S &stream) - : m_stream(stream), m_flags(stream.flags()) {} - ~fmtflags_guard() { m_stream.flags(m_flags); } - -private: - S &m_stream; - std::ios_base::fmtflags m_flags; -}; - -struct PrintOptionsImpl { - int edge_items = 3; // default edge items: 3 means print 3 items of each dimension. - int line_width = 80; // default line width: 75 means print 75 chars per line. - int threshold = 1000; // default threshold: 1000 means print 1000 elements of the tensor. - int precision = 4; // default precision: -1 means no precision limit. - int sci_mode = -1; // default sci_mode: -1 means auto decision. -}; - -inline PrintOptionsImpl &print_options() { - static PrintOptionsImpl po; - return po; -} - -void set_line_width(int line_width) { - print_options().line_width = line_width; -} - -void set_threshold(int threshold) { - print_options().threshold = threshold; -} - -void set_edge_items(int edge_items) { - print_options().edge_items = edge_items; -} - -void set_precision(int precision) { - print_options().precision = precision; -} - -void set_sci_mode(int sci_mode) { - print_options().sci_mode = sci_mode; -} - -/** - * @brief read print options from the out stream and global settings. - */ -inline print_options::PrintOptionsImpl get_print_options(std::ostream &out) { - print_options::PrintOptionsImpl res; - -// Macro to read option from stream, apply default if not set, or reset stream value -#define PROCESS_PRINT_OPTION(OPTION) \ - res.OPTION = static_cast(out.iword(print_options::OPTION::id())); \ - if (res.OPTION > 0) { \ - out.iword(print_options::OPTION::id()) = long(-1); \ - } else { \ - res.OPTION = print_options::print_options().OPTION; \ - } - - // Process all print options - PROCESS_PRINT_OPTION(edge_items); - PROCESS_PRINT_OPTION(line_width); - PROCESS_PRINT_OPTION(threshold); - PROCESS_PRINT_OPTION(precision); - - res.sci_mode = print_options::print_options().sci_mode; - -#undef PROCESS_PRINT_OPTION - return res; -} - -template -struct Printer; - -/** - * @brief Printer specialization for floating-point types (float, double, long double). - */ -template -struct Printer::value>> { - using value_type = T; - using cache_type = std::vector; - using cache_iterator = typename cache_type::const_iterator; - - explicit Printer(std::streamsize precision, int sci_mode = 0) : m_precision(precision), m_sci_mode(sci_mode) {} - - void calculate() { - m_precision = m_precision > m_required_precision ? m_required_precision : m_precision; - m_it = m_cache.cbegin(); - - if (m_scientific) { - // 3 = sign, number and dot and 4 = "e+00" - m_width = m_precision + 7; - if (m_large_exponent) { - // = e+000 (additional number) - m_width += 1; - } - } else { - std::streamsize decimals = 1; // print a leading 0 - if (std::floor(m_max) != 0) { - decimals += std::streamsize(std::log10(std::floor(m_max))); - } - // 2 => sign and dot - m_width = 2 + decimals + m_precision; - } - if (!m_required_precision) { - --m_width; - } - } - - std::ostream &print_next(std::ostream &out) { - if ((1 == m_sci_mode) || ((-1 == m_sci_mode) && m_scientific)) { - if (!m_large_exponent) { - out << std::scientific; - out.width(m_width); - out << (*m_it); - } else { - std::stringstream buf; - buf.width(m_width); - buf << std::scientific; - buf.precision(m_precision); - buf << (*m_it); - std::string res = buf.str(); - - if (res[res.size() - 4] == 'e') { - res.erase(0, 1); - res.insert(res.size() - 2, "0"); - } - out << res; - } - } else { - std::stringstream buf; - buf.width(m_width); - buf << std::fixed; - buf.precision(m_precision); - buf << (*m_it); - if (!m_required_precision && !std::isinf(*m_it) && !std::isnan(*m_it)) { - buf << '.'; - } - std::string res = buf.str(); - auto sit = res.rbegin(); - while (*sit == '0') { - *sit = ' '; - ++sit; - } - out << res; - } - ++m_it; - return out; - } - - void update(const value_type &val) { - if (val != 0 && !std::isinf(val) && !std::isnan(val)) { - if (!m_scientific || !m_large_exponent) { - int exponent = 1 + int(std::log10(std::abs(val))); - if (exponent <= -5 || exponent > 7) { - m_scientific = true; - m_required_precision = m_precision; - if (exponent <= -100 || exponent >= 100) { - m_large_exponent = true; - } - } - } - - if (std::abs(val) > m_max) { - m_max = std::abs(val); - } - if (m_required_precision < m_precision) { - while (std::floor(val * std::pow(10, m_required_precision)) != val * std::pow(10, m_required_precision)) { - m_required_precision++; - } - } - } - m_cache.push_back(val); - } - - std::streamsize width() const { return m_width; } - -private: - bool m_large_exponent = false; - bool m_scientific = false; - - std::streamsize m_width = 9; - std::streamsize m_precision; - std::streamsize m_required_precision = 0; - value_type m_max = 0; - int m_sci_mode = -1; - cache_type m_cache; - cache_iterator m_it; -}; - -/** - * @brief Printer specialization for integer types (signed and unsigned integers). - */ -template -struct Printer< - T, std::enable_if_t::value && !std::is_same::value>> { - using value_type = T; - using cache_type = std::vector; - using cache_iterator = typename cache_type::const_iterator; - - explicit Printer(std::streamsize, int sci_mode = 0) {} - - void calculate() { - m_it = m_cache.cbegin(); - m_width = 1 + std::streamsize((m_max > 0) ? std::log10(m_max) : 0) + m_sign; - } - - std::ostream &print_next(std::ostream &out) { - // + enables printing of chars etc. as numbers - // TODO should chars be printed as numbers? - out.width(m_width); - out << +(*m_it); - ++m_it; - return out; - } - - void update(const value_type &val) { - // For unsigned types, abs is not needed (always non-negative) - // For signed types, we need to take absolute value - value_type abs_val; - if constexpr (std::is_signed::value) { - abs_val = (val < 0) ? -val : val; - } else { - abs_val = val; - } - - if (abs_val > m_max) { - m_max = abs_val; - } - - if (std::is_signed::value && val < 0) { - m_sign = true; - } - m_cache.push_back(val); - } - - std::streamsize width() { return m_width; } - -private: - std::streamsize m_width; - bool m_sign = false; - value_type m_max = 0; - - cache_type m_cache; - cache_iterator m_it; -}; - -/** - * @brief Printer specialization for bool type. - */ -template -struct Printer< - T, std::enable_if_t::value>> { - using value_type = bool; - using cache_type = std::vector; - using cache_iterator = typename cache_type::const_iterator; - - explicit Printer(std::streamsize, int sci_mode = 0) {} - - void calculate() { - m_it = m_cache.cbegin(); - } - - std::ostream &print_next(std::ostream &out) { - if (*m_it) { - out << " true"; - } else { - out << "false"; - } - // TODO: the following std::setw(5) isn't working correctly on OSX. - // out << std::boolalpha << std::setw(m_width) << (*m_it); - ++m_it; - return out; - } - - void update(const value_type &val) { m_cache.push_back(val); } - - std::streamsize width() { return m_width; } - -private: - std::streamsize m_width = 5; - cache_type m_cache; - cache_iterator m_it; -}; - -} // namespace print_options -} // namespace infinicore - -namespace infinicore { -namespace print_options { - -/** - * @brief Recursively traverses tensor dimensions to collect values for printing. - */ -template -void recurser_run(Printer &printer, - const Tensor &tensor, - std::vector indexes, - std::size_t lim = 0) { - - using size_type = Size; - const auto view = at(tensor, indexes); - if (view->ndim() == 0) { - T value = item(view); - printer.update(value); - } else { - size_type i = 0; - for (; i != static_cast(view->shape()[0] - 1); ++i) { - if (lim && size_type(view->shape()[0]) > (lim * 2) && i == lim) { - i = static_cast(view->shape()[0]) - lim; - if (lim <= 1) { - break; - } - } - indexes.push_back(static_cast(i)); - recurser_run(printer, tensor, indexes, lim); - indexes.pop_back(); - } - indexes.push_back(static_cast(i)); - recurser_run(printer, tensor, indexes, lim); - indexes.pop_back(); - } -} - -/** - * @brief Recursively prints tensor elements with proper formatting. - */ -template -std::ostream &xoutput(std::ostream &out, - const Tensor &tensor, - std::vector &indexes, - Printer &printer, - std::size_t blanks, - std::streamsize element_width, - std::size_t edge_items, - std::size_t line_width) { - - using size_type = Size; - const auto view = at(tensor, indexes); - if (view->ndim() == 0) { - printer.print_next(out); - } else { - std::string indents(blanks, ' '); - - size_type i = 0; - size_type elems_on_line = 0; - const size_type ewp2 = static_cast(element_width) + size_type(2); - const size_type line_lim = static_cast(std::floor(line_width / ewp2)); - - out << '['; - for (; i != size_type(view->shape()[0] - 1); ++i) { - - if (edge_items && size_type(view->shape()[0]) > (edge_items * 2) && i == edge_items) { - if (view->ndim() == 1 && line_lim != 0 && elems_on_line >= line_lim) { - out << " ...,"; - } else if (view->ndim() > 1) { - elems_on_line = 0; - out << "...," << std::endl - << indents; - } else { - out << "..., "; - } - i = size_type(view->shape()[0]) - edge_items; - if (edge_items <= 1) { - break; - } - } - if (view->ndim() == 1 && line_lim != 0 && elems_on_line >= line_lim) { - out << std::endl - << indents; - elems_on_line = 0; - } - - indexes.push_back(static_cast(i)); - xoutput(out, tensor, indexes, printer, blanks + 1, element_width, edge_items, - line_width) - << ','; - indexes.pop_back(); - elems_on_line++; - - if ((view->ndim() == 1) && !(line_lim != 0 && elems_on_line >= line_lim)) { - ; // out << ' '; - } else if (view->ndim() > 1) { - out << std::endl - << indents; - } - } - if (view->ndim() == 1 && line_lim != 0 && elems_on_line >= line_lim) { - out << std::endl - << indents; - } - - indexes.push_back(static_cast(i)); - xoutput(out, tensor, indexes, printer, blanks + 1, element_width, edge_items, - line_width) - << ']'; - indexes.pop_back(); - } - return out; -} - -template -std::ostream &pretty_print(const Tensor &original_tensor, - std::ostream &out = std::cout) { - Tensor tensor = original_tensor->to(Device::Type::CPU); - bool on_cpu = original_tensor->device() == Device::Type::CPU; - std::string device_str = original_tensor->device().toString(); - infinicore::context::syncDevice(); - - fmtflags_guard guard(out); - - std::size_t edge_items = 0; - Size sz = tensor->numel(); - auto po = get_print_options(out); - - if (sz > static_cast(po.threshold)) { - edge_items = static_cast(po.edge_items); - } - if (sz == 0) { - out << "[]"; - return out; - } - - auto temp_precision = out.precision(); - auto precision = temp_precision; - - if (po.precision != -1) { - out.precision(static_cast(po.precision)); - precision = static_cast(po.precision); - } - - Printer printer(precision, po.sci_mode); - std::vector indexes = {}; - - recurser_run(printer, tensor, indexes, edge_items); - - printer.calculate(); - indexes.clear(); - - auto element_width = printer.width(); - - out << "tensor("; - xoutput(out, - tensor, - indexes, - printer, - 1 + 7, - element_width, - edge_items, - static_cast(po.line_width)); - - if (!on_cpu) { - out << ", device=" << '\'' << device_str << '\''; - } - - out << ", dtype=infinicore." << toString(tensor->dtype()) << ")\n"; - out.precision(temp_precision); // restore precision - return out; -} - -} // namespace print_options -} // namespace infinicore - -namespace infinicore { -std::ostream &operator<<(std::ostream &out, const Tensor &tensor) { - if (!tensor) { - out << "tensor([])\n"; - return out; - } - - switch (tensor->dtype()) { - case DataType::BYTE: // 1 - { - throw std::runtime_error("cant not print INFINI.BYTE dtype tensor !!!"); - } - case DataType::BOOL: // 2 - { - return infinicore::print_options::pretty_print(tensor, out); - } - case DataType::I8: // 3 - { - return infinicore::print_options::pretty_print(tensor, out); - } - case DataType::I16: // 4 - { - return infinicore::print_options::pretty_print(tensor, out); - } - case DataType::I32: // 5 - { - return infinicore::print_options::pretty_print(tensor, out); - } - case DataType::I64: // 6 - { - return infinicore::print_options::pretty_print(tensor, out); - } - case DataType::U8: // 7 - { - return infinicore::print_options::pretty_print(tensor, out); - } - case DataType::U16: // 8 - { - return infinicore::print_options::pretty_print(tensor, out); - } - case DataType::U32: // 9 - { - return infinicore::print_options::pretty_print(tensor, out); - } - case DataType::U64: // 10 - { - return infinicore::print_options::pretty_print(tensor, out); - } - case DataType::F8: // 11 - { - throw std::runtime_error("cant not print INFINI.F8 dtype tensor !!!"); - } - case DataType::F16: // 12 - { - // Convert F16 to F32 for printing - return infinicore::print_options::pretty_print(tensor, out); - } - case DataType::F32: // 13 - { - return infinicore::print_options::pretty_print(tensor, out); - } - case DataType::F64: // 14 - { - return infinicore::print_options::pretty_print(tensor, out); - } - case DataType::C16: // 15 - case DataType::C32: // 16 - case DataType::C64: // 17 - case DataType::C128: // 18 - { - throw std::runtime_error("cant not print Complex dtype tensor !!!"); - } - case DataType::BF16: // 19 - { - // Convert BF16 to F32 for printing - return infinicore::print_options::pretty_print(tensor, out); - } - default: - throw std::runtime_error("cant not print unknown dtype tensor : " + toString(tensor->dtype())); - } - - return out; -} -} // namespace infinicore diff --git a/src/infinicore/memory.cc b/src/infinicore/memory.cc deleted file mode 100644 index 0a7cc2df3..000000000 --- a/src/infinicore/memory.cc +++ /dev/null @@ -1,33 +0,0 @@ -#include "infinicore/memory.hpp" - -namespace infinicore { - -Memory::Memory(std::byte *data, - size_t size, - Device device, - Memory::Deleter deleter, - bool pin_memory) - : data_{data}, size_{size}, device_{device}, deleter_{deleter}, is_pinned_(pin_memory) {} - -Memory::~Memory() { - if (deleter_) { - deleter_(data_); - } -} - -std::byte *Memory::data() { - return data_; -} - -Device Memory::device() const { - return device_; -} - -size_t Memory::size() const { - return size_; -} - -bool Memory::is_pinned() const { - return is_pinned_; -} -} // namespace infinicore diff --git a/src/infinicore/nn/embedding.cc b/src/infinicore/nn/embedding.cc deleted file mode 100644 index 2596b7f81..000000000 --- a/src/infinicore/nn/embedding.cc +++ /dev/null @@ -1,139 +0,0 @@ -#include "infinicore/nn/embedding.hpp" -#include "infinicore/context/context.hpp" -#include "infinicore/ops.hpp" -#include -#include -#include - -namespace infinicore::nn { - -Embedding::Embedding(size_t num_embeddings, - size_t embedding_dim, - std::optional padding_idx, - const DataType &dtype, - const Device &device) - : num_embeddings_(num_embeddings), - embedding_dim_(embedding_dim), - padding_idx_(padding_idx), - dtype_(dtype) { - - device_ = device; - - // Validate padding_idx - if (padding_idx_.has_value()) { - int64_t idx = padding_idx_.value(); - if (idx < 0 || idx >= static_cast(num_embeddings)) { - throw std::invalid_argument( - "padding_idx must be within num_embeddings range, got " + std::to_string(idx) + " for num_embeddings=" + std::to_string(num_embeddings)); - } - } - - // Initialize parameter using macro - INFINICORE_NN_PARAMETER_INIT(weight, ({num_embeddings, embedding_dim}, dtype_, device)); - - // If padding_idx is specified, initialize that row to zeros - if (padding_idx_.has_value()) { - // TODO: Set weight[padding_idx] to zeros - // This would require a slice operation - } - - SPDLOG_DEBUG("Created Embedding module: num_embeddings={}, embedding_dim={}, dtype={}, padding_idx={}", - num_embeddings, embedding_dim, static_cast(dtype_), - padding_idx_.has_value() ? std::to_string(padding_idx_.value()) : "None"); -} - -Tensor Embedding::forward(const Tensor &indices) const { - // TODO: Implement on-device embedding for all devices, then remove the condition and the classic approach - auto device_type = device_.getType(); - if (device_type != Device::Type::CPU) { - // Use op::embedding which supports device-side input and batch dimension - return op::embedding(indices->contiguous()->to(device_), weight_); - } - - // Get the shape of indices - auto indices_shape = indices->shape(); - - // Output shape: indices_shape + [embedding_dim] - std::vector output_shape = indices_shape; - output_shape.push_back(embedding_dim_); - - // Create output tensor on the same device as weight - auto out = Tensor::empty(output_shape, weight_->dtype(), weight_->device()); - - // Flatten indices for sequential row copies - auto cpu_device = Device(Device::Type::CPU, 0); - auto indices_cpu = indices->to(cpu_device)->contiguous(); - - // Calculate total number of lookups - size_t num_lookups = 1; - for (auto dim : indices_shape) { - num_lookups *= dim; - } - - const size_t row_bytes = embedding_dim_ * dsize(weight_->dtype()); - - // Helper lambda to read index based on dtype with bounds checking - auto read_index = [&](size_t i) -> int64_t { - auto dtype = indices_cpu->dtype(); - if (dtype == DataType::I32) { - const auto *data = reinterpret_cast(indices_cpu->data()); - return static_cast(data[i]); - } else if (dtype == DataType::I64) { - const auto *data = reinterpret_cast(indices_cpu->data()); - return data[i]; - } else if (dtype == DataType::U32) { - const auto *data = reinterpret_cast(indices_cpu->data()); - return static_cast(data[i]); - } else if (dtype == DataType::U64) { - const auto *data = reinterpret_cast(indices_cpu->data()); - uint64_t val = data[i]; - // Check if value can fit in int64_t - if (val > static_cast(std::numeric_limits::max())) { - throw std::out_of_range("Index value out of range for int64_t: " + std::to_string(val)); - } - return static_cast(val); - } else { - throw std::runtime_error("Embedding indices must be integer type, got dtype=" + std::to_string(static_cast(dtype))); - } - }; - - if (weight_->device().getType() == Device::Type::CPU) { - // CPU path: memcpy row by row - const auto *weight_base = reinterpret_cast(weight_->data()); - auto *out_base = reinterpret_cast(out->data()); - for (size_t i = 0; i < num_lookups; ++i) { - int64_t idx = read_index(i); - if (idx < 0 || idx >= static_cast(num_embeddings_)) { - throw std::out_of_range( - "Index out of range: " + std::to_string(idx) + " (num_embeddings=" + std::to_string(num_embeddings_) + ")"); - } - std::memcpy(out_base + i * row_bytes, weight_base + idx * row_bytes, row_bytes); - } - } else { - // Device fallback: copy rows through Tensor slices so device runtimes own stride/stream handling. - auto flat_out = out->view({num_lookups, embedding_dim_}); - for (size_t i = 0; i < num_lookups; ++i) { - int64_t idx = read_index(i); - if (idx < 0 || idx >= static_cast(num_embeddings_)) { - throw std::out_of_range( - "Index out of range: " + std::to_string(idx) + " (num_embeddings=" + std::to_string(num_embeddings_) + ")"); - } - auto dst = flat_out->narrow({{0, i, 1}}); - auto src = weight_->narrow({{0, static_cast(idx), 1}}); - dst->copy_from(src); - } - } - - return out; -} - -std::string Embedding::extra_repr() const { - std::string repr = "Embedding(num_embeddings=" + std::to_string(num_embeddings_) + ", embedding_dim=" + std::to_string(embedding_dim_) + ", dtype=" + std::to_string(static_cast(dtype_)); - if (padding_idx_.has_value()) { - repr += ", padding_idx=" + std::to_string(padding_idx_.value()); - } - repr += ")"; - return repr; -} - -} // namespace infinicore::nn diff --git a/src/infinicore/nn/layer_norm.cc b/src/infinicore/nn/layer_norm.cc deleted file mode 100644 index 45d3b9cb2..000000000 --- a/src/infinicore/nn/layer_norm.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/nn/layer_norm.hpp" -#include "infinicore/ops.hpp" -#include -#include - -namespace infinicore::nn { - -LayerNorm::LayerNorm(size_t normalized_shape, double eps, const DataType &dtype, const Device &device) - : normalized_shape_(normalized_shape), - eps_(eps), - dtype_(dtype) { - - device_ = device; - - INFINICORE_NN_PARAMETER_INIT(weight, ({normalized_shape}, dtype_, device)); - INFINICORE_NN_PARAMETER_INIT(bias, ({normalized_shape}, dtype_, device)); -} - -Tensor LayerNorm::forward(const Tensor &x) const { - return op::layer_norm(x, weight_, bias_, static_cast(eps_)); -} - -std::string LayerNorm::extra_repr() const { - return "LayerNorm(normalized_shape=" + std::to_string(normalized_shape_) + ", eps=" + std::to_string(eps_) + ", dtype=" + std::to_string(static_cast(dtype_)) + ")"; -} - -} // namespace infinicore::nn diff --git a/src/infinicore/nn/linear.cc b/src/infinicore/nn/linear.cc deleted file mode 100644 index 65d308c31..000000000 --- a/src/infinicore/nn/linear.cc +++ /dev/null @@ -1,564 +0,0 @@ -#include "infinicore/nn/linear.hpp" -#include "../utils.hpp" -#include "infinicore/ops.hpp" -#include "infinicore/ops/distributed/allreduce.hpp" -#include "infinicore/ops/linear.hpp" -#include "infinicore/ops/linear_w4a16_awq.hpp" -#include "infinicore/ops/linear_w4a16_gptq_qy.hpp" -#include "infinicore/ops/linear_w8a8i8.hpp" -#include -#include - -namespace infinicore::nn { - -BaseLinear::BaseLinear(size_t in_features, size_t out_features, bool bias, - const DataType &dtype, const Device &device) - : in_features_(in_features), - out_features_(out_features), - has_bias_(bias), - dtype_(dtype) { - - device_ = device; -} - -BaseLinear::BaseLinear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias, - const DataType &dtype, const Device &device) - : in_features_(in_features), - out_features_(out_features), - quantization_(quantization), - has_bias_(bias), - dtype_(dtype) { - - device_ = device; -} - -Tensor BaseLinear::compute_linear(Tensor &input) const { - switch (this->quantization_->get_quant_scheme()) { - case infinicore::quantization::QuantScheme::COMPRESSED_TENSOR_W8A8I8: { - Tensor input_contiguous = input->is_contiguous() ? input : input->contiguous(); - - Tensor weight_packed_tensor = static_cast(weight_); - Tensor weight_scale_tensor = static_cast(weight_scale_); - // weight_packed should be transposed and non-contiguous. - std::optional bias_opt = has_bias_ ? std::make_optional(static_cast(bias_)) : std::nullopt; - - auto output = infinicore::op::linear_w8a8i8(input_contiguous->contiguous(), weight_packed_tensor, weight_scale_tensor, bias_opt); - return output; - } - case infinicore::quantization::QuantScheme::AWQ_W4A16: { - Tensor input_contiguous = input->is_contiguous() ? input : input->contiguous(); - Tensor qweight = static_cast(weight_); - Tensor qzeros = static_cast(weight_zeros_); - Tensor scales = static_cast(weight_scale_); - std::optional bias_opt = has_bias_ ? std::make_optional(static_cast(bias_)) : std::nullopt; - auto output = infinicore::op::linear_w4a16_awq(input_contiguous->contiguous(), qweight, scales, qzeros, bias_opt); - return output; - } - case infinicore::quantization::QuantScheme::GPTQ_W4A16_QY: { - Tensor input_contiguous = input->is_contiguous() ? input : input->contiguous(); - Tensor qweight = static_cast(weight_); - Tensor qzeros = static_cast(weight_zeros_); - Tensor scales = static_cast(weight_scale_); - Tensor g_idx = static_cast(gidx_); - std::optional bias_opt = has_bias_ ? std::make_optional(static_cast(bias_)) : std::nullopt; - auto output = infinicore::op::linear_w4a16_gptq_qy(input_contiguous->contiguous(), qweight, qzeros, scales, 0, 4); - if (bias_opt.has_value()) { - infinicore::op::add_(output, output, bias_opt.value()->as_strided(output->shape(), {0, 0, 1})); - } - return output; - } - case infinicore::quantization::QuantScheme::GPTQ_W4A16: { - throw std::runtime_error("GPTQ_W4A16 quantization scheme is not yet supported in forward pass"); - } - default: { - // Ensure input is contiguous before creating views (required for matmul) - // This prevents hanging when input tensor has non-contiguous memory layout - Tensor input_contiguous = input->is_contiguous() ? input : input->contiguous(); - - // Use ops::linear_ directly to match Python backend's exact code path - // This ensures identical computation and numerical results - // Parameter inherits from Tensor, so we cast to Tensor explicitly - Tensor weight_tensor = static_cast(weight_); - std::optional bias_opt = has_bias_ ? std::make_optional(static_cast(bias_)) : std::nullopt; - - auto output = infinicore::op::linear(input_contiguous->contiguous(), weight_tensor->contiguous(), bias_opt, alpha_); - return output; - } - } - -} // namespace infinicore::nn - -Tensor BaseLinear::forward(Tensor &input) const { - return compute_linear(input); -} - -Tensor BaseLinear::forward(Tensor &input, Tensor &residual) const { - auto output = compute_linear(input); - - // Add residual: output = output + residual - infinicore::op::add_(output, output, residual); - - return output; -} - -void BaseLinear::process_weights_after_loading() { - if (quantization_->get_quant_scheme() == infinicore::quantization::QuantScheme::GPTQ_W4A16) { - - auto config = quantization_->get_config(); - auto gptq_qy = std::make_shared(config); - quantization_ = gptq_qy; - - { - auto orig_weight = weight_; - auto orig_zeros = weight_zeros_; - auto orig_scales = weight_scale_; - - gptq_qy->convert_from_gptq_w4a16( - orig_weight, orig_zeros, orig_scales, gidx_, device_); - } - - weight_.reset(); - weight_zeros_.reset(); - weight_scale_.reset(); - - weight_ = gptq_qy->get_converted_weight(); - weight_zeros_ = gptq_qy->get_converted_zeros(); - weight_scale_ = gptq_qy->get_converted_scales(); - gptq_qy->release_buffers(); - - assert(quantization_->get_quant_scheme() == infinicore::quantization::QuantScheme::GPTQ_W4A16_QY); - } -} -} // namespace infinicore::nn - -namespace infinicore::nn { - -Linear::Linear(size_t in_features, size_t out_features, bool bias, - const DataType &dtype, const Device &device) - : BaseLinear(in_features, out_features, bias, dtype, device_) { - - device_ = device; - - // Initialize parameters using macro - INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, dtype_, device)); - - // Register bias parameter if requested - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device)); - } else { - bias_ = Parameter(); // Default constructed empty parameter - } - - // SPDLOG_DEBUG("Created Linear module: in_features={}, out_features={}, bias={}, dtype={}", - // in_features, out_features, bias, static_cast(dtype_)); -} - -Linear::Linear(size_t in_features, size_t out_features, - std::shared_ptr quantization, bool bias, - const DataType &dtype, const Device &device) - : BaseLinear(in_features, out_features, quantization, bias, dtype, device_) { - - device_ = device; - - switch (this->quantization_->get_quant_scheme()) { - case infinicore::quantization::QuantScheme::COMPRESSED_TENSOR_W8A8I8: { - INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, infinicore::DataType::I8, device)); - INFINICORE_NN_PARAMETER_INIT(weight_scale, ({out_features, 1}, infinicore::DataType::F32, device)); - - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device)); - } else { - bias_ = Parameter(); - } - break; - } - case infinicore::quantization::QuantScheme::AWQ_W4A16: { - weight_ = infinicore::nn::Parameter({out_features, in_features}, infinicore::DataType::I32, device); - this->register_parameter("qweight", weight_); - weight_zeros_ = infinicore::nn::Parameter({out_features, in_features}, infinicore::DataType::I32, device); - this->register_parameter("qzeros", weight_zeros_); - weight_scale_ = infinicore::nn::Parameter({out_features, in_features}, dtype_, device); - this->register_parameter("scales", weight_scale_); - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device)); - } else { - bias_ = Parameter(); - } - break; - } - case infinicore::quantization::QuantScheme::GPTQ_W4A16_QY: { - weight_ = infinicore::nn::Parameter({in_features / 2, out_features}, infinicore::DataType::U8, device); - this->register_parameter("qweight", weight_); - weight_zeros_ = infinicore::nn::Parameter({in_features / 128, out_features}, dtype_, device); - this->register_parameter("qzeros", weight_zeros_); - weight_scale_ = infinicore::nn::Parameter({in_features / 128, out_features}, dtype_, device); - this->register_parameter("scales", weight_scale_); - - gidx_ = infinicore::nn::Parameter({in_features}, infinicore::DataType::I32, device); - this->register_parameter("g_idx", gidx_); - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device)); - } else { - bias_ = Parameter(); - } - break; - } - case infinicore::quantization::QuantScheme::GPTQ_W4A16: { - weight_ = infinicore::nn::Parameter({in_features / 8, out_features}, infinicore::DataType::I32, device); - this->register_parameter("qweight", weight_); - weight_zeros_ = infinicore::nn::Parameter({in_features / 128, out_features / 8}, infinicore::DataType::I32, device); - this->register_parameter("qzeros", weight_zeros_); - weight_scale_ = infinicore::nn::Parameter({in_features / 128, out_features}, dtype_, device); - this->register_parameter("scales", weight_scale_); - gidx_ = infinicore::nn::Parameter({in_features}, infinicore::DataType::I32, device); - this->register_parameter("g_idx", gidx_); - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device)); - } else { - bias_ = Parameter(); - } - break; - } - default: { - // Initialize parameters using macro - INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, dtype_, device)); - - // Register bias parameter if requested - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device)); - } else { - bias_ = Parameter(); // Default constructed empty parameter - } - - // SPDLOG_DEBUG("Created Linear module: in_features={}, out_features={}, bias={}, dtype={}", - // in_features, out_features, bias, static_cast(dtype_)); - break; - } - } -} - -Tensor Linear::forward(Tensor &input) const { - return BaseLinear::forward(input); -} - -std::string Linear::extra_repr() const { - return "Linear(in_features=" + std::to_string(in_features_) + ", out_features=" + std::to_string(out_features_) + ", bias=" + (has_bias_ ? "true" : "false") + ", dtype=" + std::to_string(static_cast(dtype_)) + ")"; -} - -} // namespace infinicore::nn - -namespace infinicore::nn { - -ColumnParallelLinear::ColumnParallelLinear(size_t in_features, size_t out_features, bool bias, - const DataType &dtype, const Device &device, - Size tp_rank, Size tp_size) - : BaseLinear(in_features, out_features, bias, dtype, device_), - tp_rank_(tp_rank), - tp_size_(tp_size) { - - device_ = device; - - // Initialize parameters using macro - INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, dtype_, device, - 0, tp_rank_, tp_size_)); - - // Register bias parameter if requested - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, - 0, tp_rank_, tp_size_)); - } else { - bias_ = Parameter(); // Default constructed empty parameter - } -} - -ColumnParallelLinear::ColumnParallelLinear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias, - const DataType &dtype, const Device &device, - Size tp_rank, Size tp_size) - : BaseLinear(in_features, out_features, quantization, bias, dtype, device_), - tp_rank_(tp_rank), - tp_size_(tp_size) { - - device_ = device; - - switch (this->quantization_->get_quant_scheme()) { - case infinicore::quantization::QuantScheme::COMPRESSED_TENSOR_W8A8I8: { - - INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, infinicore::DataType::I8, device, 0, tp_rank_, tp_size_)); - INFINICORE_NN_PARAMETER_INIT(weight_scale, ({out_features, 1}, infinicore::DataType::F32, device, 0, tp_rank_, tp_size_)); - - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, 0, 0, 1)); - } else { - bias_ = Parameter(); - } - break; - } - case infinicore::quantization::QuantScheme::AWQ_W4A16: { - auto awq_ptr = std::static_pointer_cast(this->quantization_); - int group_size = awq_ptr->get_group_size(); - int packing_num = awq_ptr->get_packing_num(); - - weight_ = infinicore::nn::Parameter({in_features, out_features / packing_num}, - infinicore::DataType::I32, - device, 1, tp_rank_, tp_size_); - this->register_parameter("qweight", weight_); - - // Weight scale: [out_features, in_features / group_size] - // One FP32 scale per group of weights (group_size=128) - - weight_scale_ = infinicore::nn::Parameter({in_features / group_size, out_features}, - dtype_, - device, 1, tp_rank_, tp_size_); - this->register_parameter("scales", weight_scale_); - - // Weight zeros (zero points): [out_features, in_features / group_size] - // AWQ implementations (e.g., AutoAWQ) typically store zero points as I32 - // for symmetric/asymmetric quantization support - weight_zeros_ = infinicore::nn::Parameter({in_features / group_size, out_features / packing_num}, - infinicore::DataType::I32, - device, 1, tp_rank_, tp_size_); - - this->register_parameter("qzeros", weight_zeros_); - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, 0, 0, 1)); - } else { - bias_ = Parameter(); - } - break; - } - case infinicore::quantization::QuantScheme::GPTQ_W4A16_QY: { - auto gptq_ptr = std::static_pointer_cast(this->quantization_); - int group_size = gptq_ptr->get_group_size(); - int packing_num = gptq_ptr->get_packing_num(); - weight_ = infinicore::nn::Parameter({in_features / 2, out_features}, infinicore::DataType::U8, device, 1, tp_rank_, tp_size_); - this->register_parameter("qweight", weight_); - weight_zeros_ = infinicore::nn::Parameter({in_features / group_size, out_features}, dtype_, device, 1, tp_rank_, tp_size_); - this->register_parameter("qzeros", weight_zeros_); - weight_scale_ = infinicore::nn::Parameter({in_features / group_size, out_features}, dtype_, device, 1, tp_rank_, tp_size_); - this->register_parameter("scales", weight_scale_); - gidx_ = infinicore::nn::Parameter({in_features}, - infinicore::DataType::I32, - device, 0, tp_rank_, tp_size_); - this->register_parameter("g_idx", gidx_); - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, 0, tp_rank_, tp_size_)); - } else { - bias_ = Parameter(); - } - break; - } - case infinicore::quantization::QuantScheme::GPTQ_W4A16: { - auto gptq_ptr = std::static_pointer_cast(this->quantization_); - int group_size = gptq_ptr->get_group_size(); - int packing_num = gptq_ptr->get_packing_num(); - weight_ = infinicore::nn::Parameter({in_features / 8, out_features}, infinicore::DataType::I32, device, 1, tp_rank_, tp_size_); - this->register_parameter("qweight", weight_); - weight_zeros_ = infinicore::nn::Parameter({in_features / group_size, out_features / 8}, infinicore::DataType::I32, device, 1, tp_rank_, tp_size_); - this->register_parameter("qzeros", weight_zeros_); - weight_scale_ = infinicore::nn::Parameter({in_features / group_size, out_features}, dtype_, device, 1, tp_rank_, tp_size_); - this->register_parameter("scales", weight_scale_); - gidx_ = infinicore::nn::Parameter({in_features}, - infinicore::DataType::I32, - device, 0, tp_rank_, tp_size_); - this->register_parameter("g_idx", gidx_); - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, 0, tp_rank_, tp_size_)); - } else { - bias_ = Parameter(); - } - break; - } - default: { - // Initialize parameters using macro - INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, dtype_, device, - 0, tp_rank_, tp_size_)); - - // Register bias parameter if requested - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, - 0, tp_rank_, tp_size_)); - } else { - bias_ = Parameter(); // Default constructed empty parameter - } - break; - } - } -} - -Tensor ColumnParallelLinear::forward(Tensor &input) const { - return BaseLinear::forward(input); -} - -std::string ColumnParallelLinear::extra_repr() const { - return "ColumnParallelLinear(in_features=" + std::to_string(in_features_) + ", out_features=" + std::to_string(out_features_) + ", bias=" + (has_bias_ ? "true" : "false") + ", dtype=" + std::to_string(static_cast(dtype_)) + ")"; -} - -} // namespace infinicore::nn - -namespace infinicore::nn { - -RowParallelLinear::RowParallelLinear(size_t in_features, size_t out_features, bool bias, - const DataType &dtype, const Device &device, - Size tp_rank, Size tp_size, infinicclComm_t communicator) - : BaseLinear(in_features, out_features, bias, dtype, device_), - tp_rank_(tp_rank), - tp_size_(tp_size), communicator_(communicator) { - - device_ = device; - - // Initialize parameters using macro - INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, dtype_, device, - 1, tp_rank_, tp_size_)); - - // Register bias parameter if requested - if (bias && (0 == tp_rank_)) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, 0, 0, 1)); - } else { - bias_ = Parameter(); // Default constructed empty parameter - } -} - -RowParallelLinear::RowParallelLinear(size_t in_features, size_t out_features, std::shared_ptr quantization, bool bias, - const DataType &dtype, const Device &device, - Size tp_rank, Size tp_size, infinicclComm_t communicator) - : BaseLinear(in_features, out_features, quantization, bias, dtype, device_), - tp_rank_(tp_rank), - tp_size_(tp_size), communicator_(communicator) { - - device_ = device; - - switch (this->quantization_->get_quant_scheme()) { - case infinicore::quantization::QuantScheme::COMPRESSED_TENSOR_W8A8I8: { - INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, infinicore::DataType::I8, device, 1, tp_rank_, tp_size_)); - INFINICORE_NN_PARAMETER_INIT(weight_scale, ({out_features, 1}, infinicore::DataType::F32, device, 0, 0, 1)); - - if (bias) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, 0, tp_rank_, tp_size_)); - } else { - bias_ = Parameter(); - } - break; - } - case infinicore::quantization::QuantScheme::AWQ_W4A16: { - // AWQ W4A16 for RowParallelLinear:切分维度为 in_features（权重矩阵的第1维） - // - Weight: packed int4 in I32 containers (8 int4 per I32) - // - Group-wise quantization with group_size=128 - // - Scale and zero points stored per group along in_features dimension - - auto awq_ptr = std::static_pointer_cast(this->quantization_); - int group_size = awq_ptr->get_group_size(); - int packing_num = awq_ptr->get_packing_num(); - - // Packed weight: [out_features, in_features / 8] - weight_ = infinicore::nn::Parameter({in_features, out_features / packing_num}, - infinicore::DataType::I32, - device, 0, tp_rank_, tp_size_); - this->register_parameter("qweight", weight_); - - // Weight scale: [out_features, in_features / group_size] - - weight_scale_ = infinicore::nn::Parameter({in_features / group_size, out_features}, - dtype_, - device, 0, tp_rank_, tp_size_); - this->register_parameter("scales", weight_scale_); - // Weight zeros (zero points): [out_features, in_features / group_size] - weight_zeros_ = infinicore::nn::Parameter({in_features / group_size, out_features / packing_num}, - infinicore::DataType::I32, - device, 0, tp_rank_, tp_size_); - this->register_parameter("qzeros", weight_zeros_); - - // Bias handling in RowParallelLinear: - // - Only rank 0 holds the full bias (after all-reduce on output) - // - Other ranks have empty bias parameter - if (bias && (0 == tp_rank_)) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, 0, 0, 1)); - } else { - bias_ = Parameter(); - } - break; - } - case infinicore::quantization::QuantScheme::GPTQ_W4A16_QY: { - // GPTQ W4A16 QY for RowParallelLinear:切分维度为 in_features（权重矩阵的第1维） - // - Weight: packed int4 in U8 containers (8 int4 per U8) - // - Group-wise quantization with group_size=128 - // - Scale and zero points stored per group along in_features dimension - - auto gptq_ptr = std::static_pointer_cast(this->quantization_); - int group_size = gptq_ptr->get_group_size(); - int packing_num = gptq_ptr->get_packing_num(); - - weight_ = infinicore::nn::Parameter({in_features / 2, out_features}, infinicore::DataType::U8, device, 0, tp_rank_, tp_size_); - this->register_parameter("qweight", weight_); - weight_zeros_ = infinicore::nn::Parameter({in_features / group_size, out_features}, dtype_, device, 0, tp_rank_, tp_size_); - this->register_parameter("qzeros", weight_zeros_); - weight_scale_ = infinicore::nn::Parameter({in_features / group_size, out_features}, dtype_, device, 0, tp_rank_, tp_size_); - this->register_parameter("scales", weight_scale_); - - gidx_ = infinicore::nn::Parameter({in_features}, - infinicore::DataType::I32, - device, 0, tp_rank_, tp_size_); - this->register_parameter("g_idx", gidx_); - if (bias && (0 == tp_rank_)) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, 0, 0, 1)); - } else { - bias_ = Parameter(); - } - break; - } - case infinicore::quantization::QuantScheme::GPTQ_W4A16: { - - auto gptq_ptr = std::static_pointer_cast(this->quantization_); - int group_size = gptq_ptr->get_group_size(); - int packing_num = gptq_ptr->get_packing_num(); - - weight_ = infinicore::nn::Parameter({in_features / 8, out_features}, infinicore::DataType::I32, device, 0, tp_rank_, tp_size_); - this->register_parameter("qweight", weight_); - weight_zeros_ = infinicore::nn::Parameter({in_features / group_size, out_features / 8}, infinicore::DataType::I32, device, 0, tp_rank_, tp_size_); - this->register_parameter("qzeros", weight_zeros_); - weight_scale_ = infinicore::nn::Parameter({in_features / group_size, out_features}, dtype_, device, 0, tp_rank_, tp_size_); - this->register_parameter("scales", weight_scale_); - - gidx_ = infinicore::nn::Parameter({in_features}, - infinicore::DataType::I32, - device, 0, tp_rank_, tp_size_); - this->register_parameter("g_idx", gidx_); - if (bias && (0 == tp_rank_)) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, 0, 0, 1)); - } else { - bias_ = Parameter(); - } - break; - } - default: { - // Initialize parameters using macro - INFINICORE_NN_PARAMETER_INIT(weight, ({out_features, in_features}, dtype_, device, - 1, tp_rank_, tp_size_)); - - // Register bias parameter if requested - if (bias && (0 == tp_rank_)) { - INFINICORE_NN_PARAMETER_INIT(bias, ({out_features}, dtype_, device, 0, 0, 1)); - } else { - bias_ = Parameter(); // Default constructed empty parameter - } - - // SPDLOG_DEBUG("Created RowParallelLinear module: in_features={}, out_features={}, bias={}, dtype={}", - // in_features, out_features, bias, static_cast(dtype_)); - break; - } - } -} - -Tensor RowParallelLinear::forward(Tensor &input) const { - auto output = BaseLinear::forward(input); - - if ((tp_size_ > 1) && (communicator_ != nullptr)) { - op::distributed::allreduce_(output, output, INFINICCL_SUM, communicator_); - } - return output; -} - -std::string RowParallelLinear::extra_repr() const { - return "RowParallelLinear(in_features=" + std::to_string(in_features_) + ", out_features=" + std::to_string(out_features_) + ", bias=" + (has_bias_ ? "true" : "false") + ", dtype=" + std::to_string(static_cast(dtype_)) + ")"; -} - -} // namespace infinicore::nn diff --git a/src/infinicore/nn/module.cc b/src/infinicore/nn/module.cc deleted file mode 100644 index 5ca255322..000000000 --- a/src/infinicore/nn/module.cc +++ /dev/null @@ -1,185 +0,0 @@ -#include "infinicore/nn/module.hpp" -#include -#include - -namespace infinicore::nn { -namespace { -Parameter get_state_dict_parameter( - const std::unordered_map &state_dict, - const std::string &name) { - auto it = state_dict.find(name); - if (it == state_dict.end()) { - throw std::runtime_error("Parameter '" + name + "' not found in module."); - } - return it->second; -} -} // namespace - -std::unordered_map Module::state_dict() const { - std::unordered_map result; - collect_all_parameters(result, ""); - return result; -} - -std::vector Module::state_dict_keys() const { - std::vector result; - collect_all_parameter_names(result, ""); - return result; -} - -void Module::load_state_dict(const std::unordered_map &_state_dict) { - load_state_dict_recursively(_state_dict, ""); -} - -void Module::load_parameter(const std::string &name, const Tensor ¶m) { - auto param_it = parameters_.find(name); - if (param_it != parameters_.end()) { - try { - param_it->second.load(param); - } catch (const std::exception &e) { - throw std::runtime_error("Error loading parameter '" + name + "'. \n" + e.what()); - } - return; - } - - std::shared_ptr matched_submodule; - std::string matched_prefix; - for (const auto &[sub_name, submodule] : submodules_) { - if (name.size() <= sub_name.size() || name.compare(0, sub_name.size(), sub_name) != 0 || name[sub_name.size()] != '.') { - continue; - } - if (sub_name.size() > matched_prefix.size()) { - matched_prefix = sub_name; - matched_submodule = submodule; - } - } - - if (matched_submodule) { - try { - matched_submodule->load_parameter(name.substr(matched_prefix.size() + 1), param); - } catch (const std::exception &e) { - throw std::runtime_error("Error loading parameter '" + name + "'. \n" + e.what()); - } - return; - } - - spdlog::debug("load_parameter: Parameter '{}' not found. Available direct params={}, submodules={}", - name, parameters_.size(), submodules_.size()); - throw std::runtime_error("Parameter '" + name + "' not found in module."); -} - -void Module::load_parameters_no_sync(const std::unordered_map ¶ms, bool strict) { - auto all_params = state_dict(); - for (const auto &[name, param] : params) { - auto it = all_params.find(name); - if (it == all_params.end()) { - if (strict) { - throw std::runtime_error("Parameter '" + name + "' not found in module."); - } - continue; - } - auto existing_param = it->second; - try { - existing_param.load_no_sync(param); - } catch (const std::exception &e) { - throw std::runtime_error("Error loading parameter '" + name + "'. \n" + e.what()); - } - } -} - -void Module::load_parameter_(const std::string &name, const Tensor ¶m) { - // This function only handles direct parameters (no hierarchical traversal) - auto it = parameters_.find(name); - if (it != parameters_.end()) { - auto existing_param = it->second; - try { - existing_param.load(param); - } catch (const std::exception &e) { - throw std::runtime_error("Error loading parameter '" + name + "'. \n" + e.what()); - } - return; - } - - // Parameter not found - spdlog::debug("load_parameter_: Parameter '{}' not found. Available: {} params", - name, parameters_.size()); - throw std::runtime_error("Parameter '" + name + "' not found in module."); -} - -void Module::load_parameter_from_blob(const std::string &name, const void *data) { - auto param = parameters_[name]; - param.load_blob(data); -} - -Tensor Module::register_parameter(const std::string &name, Parameter param) { - parameters_[name] = param; - return param; -} - -Tensor Module::register_buffer(const std::string &name, Parameter buffer) { - buffers_[name] = buffer; - return buffer; -} - -void Module::load_state_dict_recursively(const std::unordered_map &_state_dict, const std::string &prefix) { - // Load direct parameters with the given prefix - for (const auto &[param_name, param] : parameters_) { - std::string full_name = prefix.empty() ? param_name : prefix + "." + param_name; - auto it = _state_dict.find(full_name); - if (it != _state_dict.end()) { - load_parameter_(param_name, it->second); - } - } - - // Recursively load parameters from submodules with extended prefix - for (const auto &[sub_name, submodule] : submodules_) { - std::string sub_prefix = prefix.empty() ? sub_name : prefix + "." + sub_name; - submodule->load_state_dict_recursively(_state_dict, sub_prefix); - } -} - -void Module::collect_all_parameters(std::unordered_map &all_params, const std::string &prefix) const { - // Add direct parameters with the given prefix - for (const auto &[param_name, param] : parameters_) { - std::string full_name = prefix.empty() ? param_name : prefix + "." + param_name; - all_params[full_name] = param; - } - - // Recursively collect parameters from submodules with extended prefix - for (const auto &[sub_name, submodule] : submodules_) { - std::string sub_prefix = prefix.empty() ? sub_name : prefix + "." + sub_name; - submodule->collect_all_parameters(all_params, sub_prefix); - } -} - -void Module::collect_all_parameter_names(std::vector &all_names, const std::string &prefix) const { - for (const auto &[param_name, _] : parameters_) { - all_names.push_back(prefix.empty() ? param_name : prefix + "." + param_name); - } - - for (const auto &[sub_name, submodule] : submodules_) { - std::string sub_prefix = prefix.empty() ? sub_name : prefix + "." + sub_name; - submodule->collect_all_parameter_names(all_names, sub_prefix); - } -} - -std::unordered_map Module::modules_dict() const { - std::unordered_map result; - collect_all_modules(result, ""); - return result; -} - -void Module::collect_all_modules(std::unordered_map &out, const std::string &prefix) const { - // 记录当前模块（跳过根节点的空前缀，可按需改为 "root"） - if (!prefix.empty()) { - out[prefix] = const_cast(this); - } - - // 递归遍历子模块 - for (const auto &[name, sub] : submodules_) { - std::string sub_prefix = prefix.empty() ? name : prefix + "." + name; - sub->collect_all_modules(out, sub_prefix); - } -} - -} // namespace infinicore::nn diff --git a/src/infinicore/nn/parameter.cc b/src/infinicore/nn/parameter.cc deleted file mode 100644 index 1b8a82762..000000000 --- a/src/infinicore/nn/parameter.cc +++ /dev/null @@ -1,106 +0,0 @@ -#include "infinicore/nn/parameter.hpp" - -#include "infinicore/context/context.hpp" - -#include -#include - -namespace infinicore::nn { -Parameter::Parameter() - : Tensor() { -} - -inline Shape get_partition_shape_(const Shape &shape, Size tp_dim, Size tp_size, Size num_shards) { - if (tp_size <= 1) { - return shape; - } - Shape part_shape = shape; - if (tp_dim < shape.size()) { - Size partition_factor = (num_shards > 0) ? num_shards : tp_size; - if (shape[tp_dim] % partition_factor != 0) { - throw std::runtime_error("Tensor dimension " + std::to_string(tp_dim) + " with size " + std::to_string(shape[tp_dim]) + " is not divisible by " + (num_shards > 0 ? "num_shards " : "tp_size ") + std::to_string(partition_factor) + "."); - } - part_shape[tp_dim] = shape[tp_dim] / partition_factor; - } - return part_shape; -} - -Parameter::Parameter(const Tensor &tensor, Size tp_dim, Size tp_rank, Size tp_size, Size num_shards) : Tensor(tensor), tp_dim_(tp_dim), tp_rank_(tp_rank), tp_size_(tp_size), num_shards_(num_shards) { - if (tp_rank_ >= tp_size_) { - throw std::runtime_error("Tensor parallel rank " + std::to_string(tp_rank_) + " must be less than tensor parallel size " + std::to_string(tp_size_) + "."); - } -} - -Parameter::Parameter( - const Shape &shape, - const DataType &dtype, - const Device &device, - Size tp_dim, - Size tp_rank, - Size tp_size, - Size num_shards) - : Parameter(Tensor::empty(get_partition_shape_(shape, tp_dim, tp_size, num_shards), dtype, device, false), tp_dim, tp_rank, tp_size, num_shards) { -} - -Parameter::Parameter(const Parameter &other) - : Tensor(other), - tp_dim_(other.tp_dim_), - tp_rank_(other.tp_rank_), - tp_size_(other.tp_size_), - num_shards_(other.num_shards_) {} - -void Parameter::load_blob(const void *data) { - Shape expected_shape = Shape(impl_->shape()); - expected_shape[tp_dim_] *= tp_size_; - auto buffer = Tensor::empty(expected_shape, impl_->dtype(), Device(Device::Type::CPU, 0), true); - std::memcpy(buffer->data(), data, buffer->nbytes()); - this->load(buffer); -} - -void Parameter::load(const Tensor &tensor) { - load_no_sync(tensor); - infinicore::context::syncStream(); -} - -void Parameter::load_no_sync(const Tensor &tensor) { - if (impl_->dtype() != tensor->dtype()) { - throw std::runtime_error("Dtype mismatch when loading tensor into parameter. Weight: " + impl_->info() + ", Tensor: " + tensor->info() + "."); - } - - Shape expected_shape = Shape(impl_->shape()); - - if (num_shards_ == 0 || num_shards_ >= tp_size_) { - expected_shape[tp_dim_] *= tp_size_; - - if (expected_shape != tensor->shape()) { - throw std::runtime_error("Shape mismatch when loading tensor into parameter. Weight: " + impl_->info() + ", Tensor: " + tensor->info() + "."); - } - if (tp_size_ > 1) { - impl_->copy_from(tensor->narrow({{tp_dim_, tp_rank_ * impl_->size(tp_dim_), impl_->size(tp_dim_)}})); - } else { - impl_->copy_from(tensor); - } - } else { - if (num_shards_ == 0) { - throw std::runtime_error("num_shards_ is 0 but entered new logic branch!"); - } - - Size replica_size = tp_size_ / num_shards_; - if (replica_size == 0) { - throw std::runtime_error("replica_size is 0! tp_size_=" + std::to_string(tp_size_) + ", num_shards_=" + std::to_string(num_shards_)); - } - - Size shard_id = tp_rank_ / replica_size; - Size shard_size = impl_->size(tp_dim_); - Size offset = shard_id * shard_size; - - expected_shape[tp_dim_] *= num_shards_; - - if (offset + shard_size > tensor->shape()[tp_dim_]) { - throw std::runtime_error("Slice out of bounds! offset=" + std::to_string(offset) + ", shard_size=" + std::to_string(shard_size) + ", tensor_dim=" + std::to_string(tensor->shape()[tp_dim_])); - } - - impl_->copy_from(tensor->narrow({{tp_dim_, offset, shard_size}})); - } -} -} // namespace infinicore::nn diff --git a/src/infinicore/nn/rmsnorm.cc b/src/infinicore/nn/rmsnorm.cc deleted file mode 100644 index 567d88e49..000000000 --- a/src/infinicore/nn/rmsnorm.cc +++ /dev/null @@ -1,49 +0,0 @@ -#include "infinicore/nn/rmsnorm.hpp" -#include "infinicore/ops.hpp" -#include -#include - -namespace infinicore::nn { - -RMSNorm::RMSNorm(size_t normalized_shape, double eps, const DataType &dtype, const Device &device) - : normalized_shape_(normalized_shape), - eps_(eps), - dtype_(dtype) { - - device_ = device; - - INFINICORE_NN_PARAMETER_INIT(weight, ({normalized_shape}, dtype_, device)); -} - -Tensor RMSNorm::forward(const Tensor &x) const { - // Delegate to InfiniCore op (backed by InfiniRT/InfiniOP) - // Validation is handled by the op layer - return op::rms_norm(x, weight_, static_cast(eps_)); -} - -void RMSNorm::forward_inplace(Tensor &x, Tensor &residual) const { - if (!residual) { - residual = x; - x = op::rms_norm(x, weight_, static_cast(eps_)); - } else { - if (device_.getType() == Device::Type::CPU - || device_.getType() == Device::Type::NVIDIA - || device_.getType() == Device::Type::ILUVATAR - || device_.getType() == Device::Type::METAX - || device_.getType() == Device::Type::MOORE - || device_.getType() == Device::Type::ALI - || device_.getType() == Device::Type::CAMBRICON - || device_.getType() == Device::Type::HYGON) { - op::add_rms_norm_inplace(x, residual, weight_, static_cast(eps_)); - } else { - op::add_(residual, x, residual); - op::rms_norm_(x, residual, weight_, static_cast(eps_)); - } - } -} - -std::string RMSNorm::extra_repr() const { - return "RMSNorm(normalized_shape=" + std::to_string(normalized_shape_) + ", eps=" + std::to_string(eps_) + ", dtype=" + std::to_string(static_cast(dtype_)) + ")"; -} - -} // namespace infinicore::nn diff --git a/src/infinicore/nn/rope.cc b/src/infinicore/nn/rope.cc deleted file mode 100644 index 115e162e2..000000000 --- a/src/infinicore/nn/rope.cc +++ /dev/null @@ -1,266 +0,0 @@ -#include "infinicore/nn/rope.hpp" -#include "../../utils.h" -#include "../utils.hpp" -#include "infinicore/ops/mrope.hpp" -#include "infinicore/ops/rope.hpp" -#include -#include -#include -#include -#include -#include -#include - -namespace infinicore::nn { - -RoPE::RoPE(size_t head_dim, - size_t rotary_dim, - size_t max_seq_len, - double theta, - Algo algo, - const DataType &dtype, - const Device &device, - std::shared_ptr scaling, - std::optional> mrope_section, - bool mrope_interleaved) - : rotary_dim_(rotary_dim), - head_dim_(head_dim), - max_seq_len_(max_seq_len), - theta_(theta), - algo_(algo), - dtype_(dtype), - scaling_(scaling), - mrope_section_(mrope_section), - mrope_interleaved_(mrope_interleaved) { - if (rotary_dim % 2 != 0) { - throw std::invalid_argument("rotary_dim must be even for RoPE, got " + std::to_string(rotary_dim)); - } - assert((rotary_dim > 0) && (rotary_dim <= head_dim_)); - if (mrope_section_.has_value()) { - const auto §ion = mrope_section_.value(); - if (section.size() != 3 || section[0] <= 0 || section[1] <= 0 || section[2] <= 0) { - throw std::invalid_argument("mrope_section must contain 3 positive values"); - } - if (2 * static_cast(section[0] + section[1] + section[2]) != rotary_dim_) { - throw std::invalid_argument("MRoPE section sum must equal rotary_dim / 2"); - } - } - device_ = device; - - // Initialize cache tables - initialize_cache(); -} - -void RoPE::initialize_cache() { - size_t cache_dim = rotary_dim_ / 2; - - // Create sin and cos cache tables: [max_seq_len, cache_dim] - INFINICORE_NN_BUFFER_INIT(sin_cache, ({max_seq_len_, cache_dim}, dtype_, device_)); - INFINICORE_NN_BUFFER_INIT(cos_cache, ({max_seq_len_, cache_dim}, dtype_, device_)); - - // Pre-compute sin and cos values - // Frequency generation always uses GPT-J style (theta^(-2j/rotary_dim)). - // The rotation algorithm (algo_) controls how dimensions are paired in the kernel. - - // Compute on CPU first, then copy to device - auto cpu_device = Device(Device::Type::CPU, 0); - - // Allocate CPU buffers - std::vector sin_data(max_seq_len_ * cache_dim); - std::vector cos_data(max_seq_len_ * cache_dim); - - for (size_t pos = 0; pos < max_seq_len_; pos++) { - for (size_t dim_idx = 0; dim_idx < cache_dim; dim_idx++) { - // 1. Base inverse frequency (shared across all RoPE types) - float base_inv_freq = 1.0f / std::pow(static_cast(theta_), 2.0f * static_cast(dim_idx) / static_cast(rotary_dim_)); - - // 2. Polymorphic scaling resolution - // Passing pre-computed base_inv_freq avoids redundant pow() calculations in subclasses - float freq_scale = scaling_ ? scaling_->get_freq_scale(pos, dim_idx, base_inv_freq) : 1.0f; - float mag_scale = scaling_ ? scaling_->get_magnitude_scale(pos, dim_idx, base_inv_freq) : 1.0f; - - // 3. Compute final angle and sin/cos values - float angle = static_cast(pos) * base_inv_freq * freq_scale; - - sin_data[pos * cache_dim + dim_idx] = std::sin(angle) * mag_scale; - cos_data[pos * cache_dim + dim_idx] = std::cos(angle) * mag_scale; - } - } - - // Convert to target dtype on CPU (matching Python's numpy astype conversion pattern) - // Python: np_array.astype(ml_dtypes.bfloat16, copy=True) converts F32 -> BF16 - if (dtype_ == DataType::F32) { - // Direct use of F32 data - auto sin_f32_cpu = Tensor::from_blob(sin_data.data(), {max_seq_len_, cache_dim}, DataType::F32, cpu_device); - auto cos_f32_cpu = Tensor::from_blob(cos_data.data(), {max_seq_len_, cache_dim}, DataType::F32, cpu_device); - sin_cache_->copy_from(sin_f32_cpu); - cos_cache_->copy_from(cos_f32_cpu); - } else if (dtype_ == DataType::BF16) { - // Convert F32 to BF16 using the same conversion as Python's ml_dtypes.bfloat16 - // This uses round-to-nearest-even (matching _f32_to_bf16 implementation) - std::vector sin_bf16_data(max_seq_len_ * cache_dim); - std::vector cos_bf16_data(max_seq_len_ * cache_dim); - - for (size_t i = 0; i < sin_data.size(); i++) { - sin_bf16_data[i] = utils::cast(sin_data[i]); - cos_bf16_data[i] = utils::cast(cos_data[i]); - } - - auto sin_bf16_cpu = Tensor::from_blob(sin_bf16_data.data(), {max_seq_len_, cache_dim}, DataType::BF16, cpu_device); - auto cos_bf16_cpu = Tensor::from_blob(cos_bf16_data.data(), {max_seq_len_, cache_dim}, DataType::BF16, cpu_device); - - // copy_from handles cross-device copying to target device - sin_cache_->copy_from(sin_bf16_cpu); - cos_cache_->copy_from(cos_bf16_cpu); - } else if (dtype_ == DataType::F16) { - // Convert F32 to F16 - std::vector sin_f16_data(max_seq_len_ * cache_dim); - std::vector cos_f16_data(max_seq_len_ * cache_dim); - - for (size_t i = 0; i < sin_data.size(); i++) { - sin_f16_data[i] = utils::cast(sin_data[i]); - cos_f16_data[i] = utils::cast(cos_data[i]); - } - - auto sin_f16_cpu = Tensor::from_blob(sin_f16_data.data(), {max_seq_len_, cache_dim}, DataType::F16, cpu_device); - auto cos_f16_cpu = Tensor::from_blob(cos_f16_data.data(), {max_seq_len_, cache_dim}, DataType::F16, cpu_device); - - sin_cache_->copy_from(sin_f16_cpu); - cos_cache_->copy_from(cos_f16_cpu); - } else { - throw std::runtime_error( - "RoPE cache dtype conversion not yet supported for dtype: " - + std::to_string(static_cast(dtype_))); - } -} - -Tensor RoPE::forward(const Tensor &x, const Tensor &pos, bool in_place) const { - if (mrope_section_.has_value()) { - throw std::runtime_error("MRoPE single-tensor forward is not implemented; use fused forward(q, k, positions) instead"); - } - Tensor y; - if (in_place) { - y = Tensor(x); - } else { - y = Tensor::empty(x->shape(), x->dtype(), x->device()); - if (rotary_dim_ < head_dim_) { - y->copy_from(x); - } - } - - size_t ndim = x->ndim(); - op::rope_(y->narrow({{ndim - 1, 0, rotary_dim_}}), - x->narrow({{ndim - 1, 0, rotary_dim_}}), - pos, sin_cache_, cos_cache_, algo_); - return y; -} - -static Tensor mrope_flatten_input(const Tensor &x, size_t head_dim, const char *name) { - if (x->ndim() == 2) { - if (x->size(1) % head_dim != 0) { - throw std::runtime_error(std::string("MRoPE expects ") + name + " hidden size to be a multiple of head_dim"); - } - return x; - } - if (x->ndim() == 3 && x->size(2) == head_dim) { - return x->view({x->size(0), x->size(1) * head_dim}); - } - throw std::runtime_error(std::string("MRoPE expects ") + name + " with shape [num_tokens, num_heads * head_dim] or [num_tokens, num_heads, head_dim]"); -} - -static Tensor mrope_flatten_output(const Tensor &x, size_t head_dim, const char *name) { - if (x->ndim() == 2) { - if (x->size(1) % head_dim != 0) { - throw std::runtime_error(std::string("MRoPE expects ") + name + " hidden size to be a multiple of head_dim"); - } - return x->view({x->size(0), x->size(1)}); - } - if (x->ndim() == 3 && x->size(2) == head_dim) { - return x->view({x->size(0), x->size(1) * head_dim}); - } - throw std::runtime_error(std::string("MRoPE expects ") + name + " with shape [num_tokens, num_heads * head_dim] or [num_tokens, num_heads, head_dim]"); -} - -std::pair RoPE::forward(const Tensor &q, const Tensor &k, const Tensor &positions) const { - if (!mrope_section_.has_value()) { - auto q_out = Tensor::empty(q->shape(), q->dtype(), q->device()); - auto k_out = Tensor::empty(k->shape(), k->dtype(), k->device()); - return forward(q_out, k_out, q, k, positions); - } - auto q_flat = mrope_flatten_input(q, head_dim_, "q"); - auto k_flat = mrope_flatten_input(k, head_dim_, "k"); - auto q_out = Tensor::empty(q_flat->shape(), q_flat->dtype(), q_flat->device()); - auto k_out = Tensor::empty(k_flat->shape(), k_flat->dtype(), k_flat->device()); - const auto §ion = mrope_section_.value(); - op::mrope_(q_out, - k_out, - q_flat, - k_flat, - cos_cache_, - sin_cache_, - positions, - static_cast(head_dim_), - static_cast(rotary_dim_), - section[0], - section[1], - section[2], - mrope_interleaved_); - return {q_out->view(q->shape()), k_out->view(k->shape())}; -} - -std::pair RoPE::forward(const Tensor &q_out, - const Tensor &k_out, - const Tensor &q, - const Tensor &k, - const Tensor &positions) const { - if (!mrope_section_.has_value()) { - auto apply_standard = [this, &positions](Tensor out, const Tensor &in) { - if (rotary_dim_ < head_dim_) { - out->copy_from(in); - } - size_t ndim = in->ndim(); - op::rope_(out->narrow({{ndim - 1, 0, rotary_dim_}}), - in->narrow({{ndim - 1, 0, rotary_dim_}}), - positions, - sin_cache_, - cos_cache_, - algo_); - }; - apply_standard(q_out, q); - apply_standard(k_out, k); - return {q_out, k_out}; - } - auto q_flat = mrope_flatten_input(q, head_dim_, "q"); - auto k_flat = mrope_flatten_input(k, head_dim_, "k"); - auto q_out_flat = mrope_flatten_output(q_out, head_dim_, "q_out"); - auto k_out_flat = mrope_flatten_output(k_out, head_dim_, "k_out"); - const auto §ion = mrope_section_.value(); - op::mrope_(q_out_flat, - k_out_flat, - q_flat, - k_flat, - cos_cache_, - sin_cache_, - positions, - static_cast(head_dim_), - static_cast(rotary_dim_), - section[0], - section[1], - section[2], - mrope_interleaved_); - return {q_out, k_out}; -} - -std::string RoPE::extra_repr() const { - std::string algo_str = (algo_ == Algo::GPT_J) ? "GPT_J" : "GPT_NEOX"; - std::string repr = "RoPE(head_dim=" + std::to_string(head_dim_) + ", rotary_dim=" + std::to_string(rotary_dim_) + ", max_seq_len=" + std::to_string(max_seq_len_) + ", theta=" + std::to_string(theta_) + ", algo=" + algo_str + ", dtype=" + std::to_string(static_cast(dtype_)); - if (mrope_section_.has_value()) { - const auto §ion = mrope_section_.value(); - repr += ", mrope_section=[" + std::to_string(section[0]) + "," + std::to_string(section[1]) + "," + std::to_string(section[2]) + "]"; - repr += ", mrope_interleaved=" + std::string(mrope_interleaved_ ? "true" : "false"); - } - repr += ")"; - return repr; -} - -} // namespace infinicore::nn diff --git a/src/infinicore/nn/rope_scaling_configs.cc b/src/infinicore/nn/rope_scaling_configs.cc deleted file mode 100644 index d0901e7d3..000000000 --- a/src/infinicore/nn/rope_scaling_configs.cc +++ /dev/null @@ -1,193 +0,0 @@ -#include "infinicore/nn/rope_scaling_configs.hpp" -#include -#include -#include -#include -#include - -namespace infinicore::nn { - -namespace { -// Define a portable PI constant to avoid relying on the non-standard M_PI macro -// which is missing on MSVC (Windows) by default. -constexpr float kPi = 3.14159265358979323846f; -} // anonymous namespace - -// LongRopeScalingConfig Implementation -LongRopeScalingConfig::LongRopeScalingConfig( - std::vector short_factor, - std::vector long_factor, - size_t original_max_position_embeddings, - float factor) - : short_factor_(std::move(short_factor)), - long_factor_(std::move(long_factor)), - original_max_position_embeddings_(original_max_position_embeddings), - factor_(factor == 1.0f ? 1.0f : std::sqrt(1 + std::log(factor) / std::log(original_max_position_embeddings))) {} - -float LongRopeScalingConfig::get_freq_scale(size_t pos, size_t dim_idx, float base_inv_freq) const { - float _ext = (pos < original_max_position_embeddings_) ? short_factor_[dim_idx] : long_factor_[dim_idx]; - // The base inv_freq is multiplied by this scale. - // Original: inv_freq = 1.0f / (_ext * pow(theta, 2j/head_dim)) - // New: inv_freq = base_inv_freq * (1.0f / _ext) - return 1.0f / _ext; -} - -float LongRopeScalingConfig::get_magnitude_scale(size_t pos, size_t dim_idx, float base_inv_freq) const { - return factor_; -} - -// Llama3RopeScalingConfig Implementation -Llama3RopeScalingConfig::Llama3RopeScalingConfig( - float factor, - float low_freq_factor, - float high_freq_factor, - size_t original_max_position_embeddings) - : factor_(factor), - low_freq_factor_(low_freq_factor), - high_freq_factor_(high_freq_factor), - original_max_position_embeddings_(original_max_position_embeddings) {} - -float Llama3RopeScalingConfig::get_freq_scale(size_t pos, size_t dim_idx, float base_inv_freq) const { - // Calculate the wavelength corresponding to the current inverse frequency - float wavelen = 2.0f * static_cast(kPi) / base_inv_freq; - - // Compute the wavelength thresholds that separate high, mid, and low frequencies - float low_freq_wavelen = static_cast(original_max_position_embeddings_) / low_freq_factor_; - float high_freq_wavelen = static_cast(original_max_position_embeddings_) / high_freq_factor_; - - float scale = 1.0f; - - if (wavelen < low_freq_wavelen) { - // High-frequency band: short wavelengths retain the original scale - scale = 1.0f; - } else if (wavelen > high_freq_wavelen) { - // Low-frequency band: long wavelengths are directly scaled by the factor - scale = factor_; - } else { - // Mid-frequency band: apply smooth linear interpolation between 1.0 and factor_ - float smooth = (static_cast(original_max_position_embeddings_) / wavelen - low_freq_factor_) / (high_freq_factor_ - low_freq_factor_); - scale = 1.0f - smooth + smooth * factor_; - } - - // The framework applies the scale multiplicatively (inv_freq = base_inv_freq * return_value). - // Since the Llama3 logic divides the frequency (inv_freq = base_inv_freq / scale), - // we return the inverse of the computed scale. - return 1.0f / scale; -} - -namespace { - -float yarn_find_correction_dim( - int num_rotations, - size_t rotary_dim, - float base, - size_t original_max_position_embeddings) { - return (static_cast(rotary_dim) - * std::log(static_cast(original_max_position_embeddings) - / (static_cast(num_rotations) * 2.0f * kPi))) - / (2.0f * std::log(base)); -} - -std::pair yarn_find_correction_range( - int low_rot, - int high_rot, - size_t rotary_dim, - float base, - size_t original_max_position_embeddings, - bool truncate) { - float low = yarn_find_correction_dim( - low_rot, rotary_dim, base, original_max_position_embeddings); - float high = yarn_find_correction_dim( - high_rot, rotary_dim, base, original_max_position_embeddings); - if (truncate) { - low = std::floor(low); - high = std::ceil(high); - } - low = std::max(low, 0.0f); - high = std::min(high, static_cast(rotary_dim) - 1.0f); - return {low, high}; -} - -float yarn_get_mscale(float scale, float mscale_coeff = 1.0f) { - if (scale <= 1.0f) { - return 1.0f; - } - return 0.1f * mscale_coeff * std::log(scale) + 1.0f; -} - -} // anonymous namespace - -// YarnRopeScalingConfig Implementation -YarnRopeScalingConfig::YarnRopeScalingConfig( - float factor, - size_t original_max_position_embeddings, - size_t rotary_dim, - float rope_theta, - int beta_fast, - int beta_slow, - float mscale, - float mscale_all_dim) - : factor_(factor), - original_max_position_embeddings_(original_max_position_embeddings) { - if (factor <= 0.0f) { - throw std::invalid_argument( - "YarnRopeScalingConfig factor must be positive, got " - + std::to_string(factor)); - } - if (original_max_position_embeddings == 0) { - throw std::invalid_argument( - "YarnRopeScalingConfig original_max_position_embeddings must be positive"); - } - if (rope_theta <= 0.0f) { - throw std::invalid_argument( - "YarnRopeScalingConfig rope_theta must be positive, got " - + std::to_string(rope_theta)); - } - if (rotary_dim < 2 || rotary_dim % 2 != 0) { - throw std::invalid_argument( - "YarnRopeScalingConfig rotary_dim must be a positive even number, got " - + std::to_string(rotary_dim)); - } - - // vLLM: yarn_find_correction_range(beta_fast, beta_slow, ...) - auto [low, high] = yarn_find_correction_range( - beta_fast, - beta_slow, - rotary_dim, - rope_theta, - original_max_position_embeddings, - true); - correction_low_ = low; - correction_high_ = high; - - magnitude_scale_ = yarn_get_mscale(factor, mscale) / yarn_get_mscale(factor, mscale_all_dim); -} - -float YarnRopeScalingConfig::yarn_linear_ramp(size_t dim_idx) const { - float low = correction_low_; - float high = correction_high_; - if (low == high) { - high += 0.001f; // Prevent singularity (matches vLLM) - } - float linear = (static_cast(dim_idx) - low) / (high - low); - return std::clamp(linear, 0.0f, 1.0f); -} - -float YarnRopeScalingConfig::get_freq_scale( - size_t /*pos*/, - size_t dim_idx, - float /*base_inv_freq*/) const { - constexpr float kExtrapolationFactor = 1.0f; - float ramp = yarn_linear_ramp(dim_idx); - float inv_freq_mask = (1.0f - ramp) * kExtrapolationFactor; - return (1.0f - inv_freq_mask) / factor_ + inv_freq_mask; -} - -float YarnRopeScalingConfig::get_magnitude_scale( - size_t /*pos*/, - size_t /*dim_idx*/, - float /*base_inv_freq*/) const { - return magnitude_scale_; -} - -} // namespace infinicore::nn diff --git a/src/infinicore/ops/README.md b/src/infinicore/ops/README.md deleted file mode 100644 index 64e9fee37..000000000 --- a/src/infinicore/ops/README.md +++ /dev/null @@ -1,303 +0,0 @@ -# infinicore::ops 开发指南 - -infinicore::ops 模块包含了 InfiniCore 所有 C++ 算子的接口和实现。外部用户可以通过 `include/infinicore/ops/*OPNAME*/*OPNAME*.h` 中定义的 C++ 接口进行算子调用。部分算子会通过 pybind 暴露给 python 前端。 - -## 开发指南 - -### 1. 算子定义 - -创建 `include/infinicore/ops/*OPNAME*/*OPNAME*.h` 头文件，并根据算子名称定义算子的类以及外部计算接口（包括 in-place 和 out-of-place 两种模式），注意算子名称不能重复。 - -一个算子类主要包含以下部分： - -- schema 定义，用于描述算子的输入输出参数形式。 -- execute 函数，算子的计算逻辑。 -- dispatcher 分发器，用于注册算子在不同设备上的 kernel 实现。一个进程中，一种算子只有一个全局分发器，每种设备上只能同时注册一个 kernel 实现，可以多次注册对之前的实现进行覆盖。详细信息请参考 `include/infinicore/ops/common/dispatcher.hpp`。 - -示例 `Gemm` 算子的头文件如下： - -```c++ -#pragma once - -#include "../device.hpp" -#include "common/op.hpp" - -namespace infinicore::op { - -class Gemm { -public: - using schema = void (*)(Tensor, Tensor, Tensor, float, float); - static void execute(Tensor c, Tensor a, Tensor b, float alpha, float beta); - static common::OpDispatcher &dispatcher(); -}; - -Tensor gemm(Tensor a, Tensor b, float alpha = 1.0f, float beta = 0.0f); -void gemm_(Tensor c, Tensor a, Tensor b, float alpha, float beta); - -} -``` - -### 2. 算子实现 - -在 `src/infinicore/ops/*OPNAME*/*OPNAME*.cpp` 文件中实现算子的计算逻辑。 - -- execute 函数，使用算子的分发器，调用对应硬件上的核函数。可以通过 `context::setDevice` 来改变当前运行时的设备种类。 -- 计算接口，使用 execute 函数实现算子接口的计算逻辑，包括 in-place 和 out-of-place 两种模式，其中 in-place 模式的接口函数名以 `_` 结尾，将输出接口写入给定的参数中；out-of-place 模式的接口会为输出创建新的 Tensor。 - -示例 `Gemm` 算子的实现如下： - -```c++ -#include "infinicore/ops/gemm.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Gemm::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Gemm::execute(Tensor c, Tensor a, Tensor b, float alpha, float beta) { - // 检查张量设备是否一致 - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - // 将运行时设备设置为与张量一致。若设备为CPU时，该接口不会进行任何操作 - infinicore::context::setDevice(c->device()); - // 根据张量的设备种类选择 kernel，执行计算 - dispatcher().lookup(c->device().getType())(c, a, b, alpha, beta); -} - -Tensor gemm(Tensor a, Tensor b, float alpha, float beta) { - Shape shape = a->shape(); - Size size = a->ndim(); - shape[size - 1] = b->size(size - 1); - auto c = Tensor::empty(shape, a->dtype(), a->device()); - gemm_(c, a, b, alpha, beta); - return c; -} - -void gemm_(Tensor c, Tensor a, Tensor b, float alpha, float beta) { - Gemm::execute(c, a, b, alpha, beta); -} - -} -``` - -### 3. Kernel 注册 - -在 `src/infinicore/ops/*OPNAME*/` 目录中添加算子和函数实现，并在算子的分发器中进行注册。你可以选择为单个设备、多个设备、或全部平台注册 kernel 实现（函数指针），你还可以通过使用 `override_existing` 模式覆盖之前的实现。具体信息请参考 `include/infinicore/ops/common/dispatcher.hpp`： - -```c++ -// 为某个设备注册 kernel 实现 -void registerDevice(Device::Type device_type, Fn fn, bool override_existing = true); - -// 为多个设备注册 kernel 实现 -void registerDevice(std::initializer_list device_types, Fn fn, bool override_existing = true); - -// 为全部平台注册 kernel 实现 -void registerAll(Fn fn, bool override_existing = true); - -// 查找 kernel 实现 -Fn lookup(Device::Type device_type) const; -``` - -如果你为多个（或全部）设备注册了同一个 kernel 实现，那么你需要自行实现不同设备的分发机制。比如本框架中的 InfiniOP 算子库，其算子接口在不同平台都保持了一致，并根据当前设备类型自动分发，因此在注册时会为所有平台注册同一个计算函数。以 Gemm 算子为例： - -```c++ -namespace infinicore::op::matmul_impl::infiniop { - -// InfiniOP 算子缓存（线程级） -thread_local common::OpCache caches( - 100, - [](infiniopGemmDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyGemmDescriptor(desc)); - desc = nullptr; - } - }); - -// 计算函数 -void calculate(Tensor c, Tensor a, Tensor b, float alpha, float beta) { - // ... - INFINICORE_CHECK_ERROR(infiniopGemm( - desc, workspace->data(), workspace_size, - c->data(), a->data(), b->data(), alpha, beta, context::getStream())); -} - -// 在加载 InfiniCore 时为全平台注册 InfiniOP实现 -static bool registered = []() { - Gemm::dispatcher().registerAll(&calculate, false); - return true; -}(); -} -``` - -你可以仿照上面的例子单独为不同平台实现核函数并注册。请注意在 `xmake/*lua` 中添加对源文件的编译方式，并做好跨平台隔离工作以保证项目在别的平台上也可以正常编译。你可以选择像上面的例子一样，通过 `static bool registered = []() {...}` 方式在加载时注册核函数，但请注意避免加载时为同一个算子重复注册不同核函数的未定义行为。你也可以在程序运行时显式地注册算子。 - -如果你想通过 InfiniOP 库来实现算子，请参考 [`InfiniOP 开发者文档`](src/infiniop/README.md) 文件。 - -### 4. Python 接口 - -通过 pybind11 将 C++ 算子暴露给 Python 前端，需要在 `src/infinicore/pybind11/ops/*OPNAME*/` 目录中添加相应的头文件，并在 `src/infinicore/pybind11/ops.hpp` 中调用。之后你需要在 `python/infinicore/ops/` 目录中为算子添加一个 Python 文件，通过调用你刚才定义的 pybind 接口实现你的 Python 接口，并将 Python 接口通过 `python/infinicore/__init__.py` 暴露给外部。 - -### 5. Python 测试 - -在实现了 Python 接口后，你需要在 `/test/infinicore/ops/` 中添加相应的算子测试脚本，并确保测试通过。该目录下的测试使用了统一的测试框架，大部分测试功能已经实现，比如根据形状构建随机张量、自动测试算子的正确性和性能等。 - -你需要继承 `BaseOperatorTest` 类并实现以下方法： - -`get_test_cases()`: 返回测试用例列表 - -`torch_operator()`: PyTorch 参考实现 - -`infinicore_operator()`: InfiniCore 算子实现 - -```python -class OpTest(BaseOperatorTest): - """Add operator test with simplified implementation""" - - def __init__(self): - super().__init__("Add") - - def get_test_cases(self): - return parse_test_cases() - - def torch_operator(self, *args, **kwargs): - """PyTorch add implementation""" - return torch.add(*args, **kwargs) - - def infinicore_operator(self, *args, **kwargs): - """InfiniCore add implementation""" - return infinicore.add(*args, **kwargs) -``` - -在测试脚本中你需要为算子测试脚本添加测例。你可以使用统一的格式简洁地准备测试所需数据，如`_TEST_CASES_DATA`。 - -```python -# Test cases format: (shape, a_strides, b_strides, c_strides) -_TEST_CASES_DATA = [ - # Basic cases - ((13, 4), None, None, None), - ((13, 4), (10, 1), (10, 1), None), - # Strided cases - ((13, 4), None, None, (10, 1)), - ((13, 4), (10, 1), (10, 1), (10, 1)), - # 3D cases - ((13, 4, 4), None, None, None), - ((13, 4, 4), (20, 4, 1), (20, 4, 1), None), - # Broadcast cases - ((13, 4, 4), (4, 0, 1), (0, 4, 1), None), - # Large tensors - ((16, 5632), None, None, None), - ((16, 5632), (13312, 1), (13312, 1), None), -] -``` - - -对于支持多种精度的算子，你可以指定测试通过的误差范围。 - -```python -_TENSOR_DTYPES = [infinicore.float16, infinicore.bfloat16, infinicore.float32] - - -_TOLERANCE_MAP = { - infinicore.float16: {"atol": 1e-3, "rtol": 1e-3}, - infinicore.float32: {"atol": 1e-5, "rtol": 1e-5}, - infinicore.bfloat16: {"atol": 5e-3, "rtol": 1e-2}, -} -``` - -通过 `parse_test_cases` 函数来解析测例数据，构建具体的测例。OUT OF PLACE以及各种INPLACE操作，需要作为不同的测例传入`get_test_cases`函数。 - -```python -# Parse test cases -def parse_test_cases(): - """ - Parse test case data and return list of TestCase objects for all operation types. - Each test case contains all necessary information for execution and validation. - """ - test_cases = [] - - for data in _TEST_CASES_DATA: - shape = data[0] - a_strides = data[1] if len(data) > 1 else None - b_strides = data[2] if len(data) > 2 else None - c_strides = data[3] if len(data) > 3 else None - - # Check if tensors support in-place operations - a_supports_inplace = not is_broadcast(a_strides) - b_supports_inplace = not is_broadcast(b_strides) - c_supports_inplace = not is_broadcast(c_strides) - - # Generate test cases for all data types - for dtype in _TENSOR_DTYPES: - tolerance = _TOLERANCE_MAP.get(dtype, {"atol": 0, "rtol": 1e-3}) - - # Create typed tensor specs - a_spec = TensorSpec.from_tensor(shape, a_strides, dtype, name="a") - b_spec = TensorSpec.from_tensor(shape, b_strides, dtype, name="b") - c_spec = TensorSpec.from_tensor(shape, c_strides, dtype, name="c") - - # Test Case 1: Out-of-place (return value) - test_cases.append( - TestCase( - inputs=[a_spec, b_spec], - kwargs={}, - output_spec=None, - comparison_target=None, - tolerance=tolerance, - description=f"Add - OUT_OF_PLACE", - ) - ) - - # Test Case 2: In-place with explicit output tensor (add(a, b, out=c)) - if c_supports_inplace: - test_cases.append( - TestCase( - inputs=[a_spec, b_spec], - kwargs={}, - output_spec=c_spec, # Specify the output tensor spec - comparison_target="out", - tolerance=tolerance, - description=f"Add - INPLACE(out)", - ) - ) - - # Test Case 3: In-place on first input (add(a, b, out=a)) - if a_supports_inplace: - test_cases.append( - TestCase( - inputs=[a_spec, b_spec], - kwargs={"out": 0}, # Use index 0 for first input - output_spec=None, - comparison_target=0, # Compare first input - tolerance=tolerance, - description=f"Add - INPLACE(a)", - ) - ) - - # Test Case 4: In-place on second input (add(a, b, out=b)) - if b_supports_inplace: - test_cases.append( - TestCase( - inputs=[a_spec, b_spec], - kwargs={"out": 1}, # Use index 1 for second input - output_spec=None, - comparison_target=1, # Compare second input - tolerance=tolerance, - description=f"Add - INPLACE(b)", - ) - ) - - return test_cases - -``` - -运行测试指令检查算子的正确性和性能： - -```bash -python test/infinicore/ops/add --nvidia --verbose --bench --debug -``` -```bash -python test/infinicore/run.py --ops add matmul --nvidia --verbose --bench -``` diff --git a/src/infinicore/ops/acos/acos.cc b/src/infinicore/ops/acos/acos.cc deleted file mode 100644 index 743b477dd..000000000 --- a/src/infinicore/ops/acos/acos.cc +++ /dev/null @@ -1,24 +0,0 @@ -#include "infinicore/ops/acos.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Acos::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Acos::execute(Tensor output, Tensor input) { - dispatcher().lookup(context::getDevice().getType())(output, input); -} - -Tensor acos(Tensor input) { - auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); - acos_(output, input); - return output; -} - -void acos_(Tensor output, Tensor input) { - Acos::execute(output, input); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/acos/acos_infiniop.cc b/src/infinicore/ops/acos/acos_infiniop.cc deleted file mode 100644 index 37f3af066..000000000 --- a/src/infinicore/ops/acos/acos_infiniop.cc +++ /dev/null @@ -1,84 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/acos.hpp" -#include "infinicore/ops/common/cache.hpp" -#include -#include - -namespace infinicore::op::acos_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopAcosDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyAcosDescriptor(desc)); - desc = nullptr; - } - }); - -struct WorkspaceEntry { - size_t size = 0; - std::shared_ptr buf = nullptr; -}; - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - // 获取或创建 descriptor - auto desc_opt = cache.get(seed); - infiniopAcosDescriptor_t desc = nullptr; - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateAcosDescriptor( - context::getInfiniopHandle(output->device()), &desc, - output->desc(), input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - static thread_local std::unordered_map s_workspace_map; - auto it = s_workspace_map.find(desc); - - if (it == s_workspace_map.end()) { - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAcosWorkspaceSize(desc, &workspace_size)); - - WorkspaceEntry entry; - if (workspace_size > 0) { - entry.buf = context::allocateMemory(workspace_size); - entry.size = workspace_size; - } else { - entry.buf = nullptr; - entry.size = 0; - } - it = s_workspace_map.emplace(desc, std::move(entry)).first; - } else { - - size_t required_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAcosWorkspaceSize(desc, &required_size)); - if (required_size > it->second.size) { - it->second.buf = context::allocateMemory(required_size); - it->second.size = required_size; - } - } - void *workspace_ptr = (it != s_workspace_map.end() && it->second.buf) ? it->second.buf->data() : nullptr; - size_t workspace_size = (it != s_workspace_map.end()) ? it->second.size : 0; - INFINICORE_CHECK_ERROR(infiniopAcos( - desc, - workspace_ptr, - workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -static bool registered = []() { - Acos::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::acos_impl::infiniop diff --git a/src/infinicore/ops/adaptive_avg_pool1d/adaptive_avg_pool1d.cc b/src/infinicore/ops/adaptive_avg_pool1d/adaptive_avg_pool1d.cc deleted file mode 100644 index 06d267131..000000000 --- a/src/infinicore/ops/adaptive_avg_pool1d/adaptive_avg_pool1d.cc +++ /dev/null @@ -1,43 +0,0 @@ -#include "infinicore/ops/adaptive_avg_pool1d.hpp" -#include -#include - -namespace infinicore::op { - -common::OpDispatcher &AdaptiveAvgPool1d::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void AdaptiveAvgPool1d::execute(Tensor output, Tensor input) { - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No AdaptiveAvgPool1d implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor adaptive_avg_pool1d(Tensor input, int64_t output_size) { - size_t ndim = input->ndim(); - if (ndim != 2 && ndim != 3) { - throw std::runtime_error("AdaptiveAvgPool1d: Input tensor must be 2D or 3D."); - } - - if (output_size <= 0) { - throw std::runtime_error("AdaptiveAvgPool1d: output_size must be positive."); - } - - auto out_shape = input->shape(); - out_shape[ndim - 1] = output_size; - - auto output = Tensor::empty(out_shape, input->dtype(), input->device()); - - AdaptiveAvgPool1d::execute(output, input); - - return output; -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/adaptive_avg_pool1d/adaptive_avg_pool1d_infiniop.cc b/src/infinicore/ops/adaptive_avg_pool1d/adaptive_avg_pool1d_infiniop.cc deleted file mode 100644 index 5a4e286f8..000000000 --- a/src/infinicore/ops/adaptive_avg_pool1d/adaptive_avg_pool1d_infiniop.cc +++ /dev/null @@ -1,94 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/adaptive_avg_pool1d.hpp" -#include "infinicore/ops/common/cache.hpp" -#include -#include - -namespace infinicore::op::adaptive_avg_pool1d_impl::infiniop { - -// 1. 资源上下文 -struct AdaptiveAvgPool1dContext { - infiniopAdaptiveAvgPool1dDescriptor_t desc = nullptr; - std::shared_ptr workspace_buf = nullptr; - size_t workspace_size = 0; - - void *getWorkspacePtr() const { - return workspace_buf ? workspace_buf->data() : nullptr; - } -}; - -// 2. 缓存定义 -thread_local common::OpCache caches( - 256, - [](AdaptiveAvgPool1dContext &ctx) { - if (ctx.desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyAdaptiveAvgPool1dDescriptor(ctx.desc)); - ctx.desc = nullptr; - } - ctx.workspace_buf = nullptr; - }); - -// 3. 核心计算函数 -void calculate(Tensor output, Tensor input) { - size_t seed = reinterpret_cast(input.operator->()); - if (output->ndim() >= 3) { - seed ^= (output->shape()[2] << 1); - } - - static thread_local size_t last_seed = 0; - static thread_local bool last_ctx_valid = false; - static thread_local AdaptiveAvgPool1dContext last_ctx; - - AdaptiveAvgPool1dContext *active_ctx = nullptr; - - if (last_ctx_valid && seed == last_seed) { - active_ctx = &last_ctx; - } else { - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - auto &cache = caches.getCache(device_type, device_index); - - auto opt_ctx = cache.get(seed); - if (opt_ctx) { - last_ctx = *opt_ctx; - } else { - AdaptiveAvgPool1dContext new_ctx; - - INFINICORE_CHECK_ERROR(infiniopCreateAdaptiveAvgPool1dDescriptor( - context::getInfiniopHandle(output->device()), - &new_ctx.desc, - output->desc(), - input->desc())); - - INFINICORE_CHECK_ERROR(infiniopGetAdaptiveAvgPool1dWorkspaceSize(new_ctx.desc, &new_ctx.workspace_size)); - - if (new_ctx.workspace_size > 0) { - new_ctx.workspace_buf = context::allocateMemory(new_ctx.workspace_size); - } - - cache.put(seed, new_ctx); - last_ctx = new_ctx; - } - - last_seed = seed; - last_ctx_valid = true; - active_ctx = &last_ctx; - } - - INFINICORE_CHECK_ERROR(infiniopAdaptiveAvgPool1d( - active_ctx->desc, - active_ctx->getWorkspacePtr(), - active_ctx->workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -// 注册 -static bool registered = []() { - AdaptiveAvgPool1d::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::adaptive_avg_pool1d_impl::infiniop diff --git a/src/infinicore/ops/adaptive_avg_pool3d/adaptive_avg_pool3d.cc b/src/infinicore/ops/adaptive_avg_pool3d/adaptive_avg_pool3d.cc deleted file mode 100644 index 298a5d553..000000000 --- a/src/infinicore/ops/adaptive_avg_pool3d/adaptive_avg_pool3d.cc +++ /dev/null @@ -1,36 +0,0 @@ -#include "infinicore/ops/adaptive_avg_pool3d.hpp" -#include -#include -namespace infinicore::op { - -common::OpDispatcher &AdaptiveAvgPool3D::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void AdaptiveAvgPool3D::execute(Tensor y, Tensor x) { - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - if (func == nullptr) { - throw std::runtime_error("No AdaptiveAvgPool3D implementation found for device type: " + std::to_string(static_cast(device_type))); - } - func(y, x); -} - -Tensor adaptive_avg_pool3d(Tensor x, std::vector output_size) { - - // Create output tensor shap - Shape y_shape = x->shape(); - y_shape[2] = output_size[0]; // D dimension - y_shape[3] = output_size[1]; // H dimension - y_shape[4] = output_size[2]; // W dimension - - auto y = Tensor::empty(y_shape, x->dtype(), x->device()); - adaptive_avg_pool3d_(y, x); - return y; -} - -void adaptive_avg_pool3d_(Tensor y, Tensor x) { - AdaptiveAvgPool3D::execute(y, x); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/adaptive_avg_pool3d/adaptive_avg_pool3d_infiniop.cc b/src/infinicore/ops/adaptive_avg_pool3d/adaptive_avg_pool3d_infiniop.cc deleted file mode 100644 index 01341f9fd..000000000 --- a/src/infinicore/ops/adaptive_avg_pool3d/adaptive_avg_pool3d_infiniop.cc +++ /dev/null @@ -1,62 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/adaptive_avg_pool3d.hpp" -#include "infinicore/ops/common/cache.hpp" - -namespace infinicore::op::adaptive_avg_pool3d_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopAdaptiveAvgPool3DDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyAdaptiveAvgPool3DDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopAdaptiveAvgPool3DDescriptor_t desc = nullptr; - - if (!desc_opt) { - // Convert vector to array for output_size - std::vector output_size_vec = {output->size(2), output->size(3), output->size(4)}; - - INFINICORE_CHECK_ERROR(infiniopCreateAdaptiveAvgPool3DDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc(), - output_size_vec.data())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // Get workspace size and allocate if needed - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAdaptiveAvgPool3DWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR( - infiniopAdaptiveAvgPool3D( - desc, - workspace->data(), workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -static bool registered = []() { - AdaptiveAvgPool3D::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::adaptive_avg_pool3d_impl::infiniop diff --git a/src/infinicore/ops/adaptive_max_pool1d/adaptive_max_pool1d.cc b/src/infinicore/ops/adaptive_max_pool1d/adaptive_max_pool1d.cc deleted file mode 100644 index bd80b0771..000000000 --- a/src/infinicore/ops/adaptive_max_pool1d/adaptive_max_pool1d.cc +++ /dev/null @@ -1,30 +0,0 @@ -#include "infinicore/ops/adaptive_max_pool1d.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &AdaptiveMaxPool1d::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -void AdaptiveMaxPool1d::execute(Tensor y, Tensor x, size_t output_size) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - infinicore::context::setDevice(y->device()); - dispatcher().lookup(y->device().getType())(y, x, output_size); -} - -Tensor adaptive_max_pool1d(Tensor x, size_t output_size) { - infinicore::Shape y_shape = x->shape(); - y_shape.back() = output_size; - auto y = Tensor::empty(y_shape, x->dtype(), x->device()); - adaptive_max_pool1d_(y, x, output_size); - return y; -} - -void adaptive_max_pool1d_(Tensor y, Tensor x, size_t output_size) { - AdaptiveMaxPool1d::execute(y, x, output_size); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/adaptive_max_pool1d/adaptive_max_pool1d_infiniop.cc b/src/infinicore/ops/adaptive_max_pool1d/adaptive_max_pool1d_infiniop.cc deleted file mode 100644 index 451489e15..000000000 --- a/src/infinicore/ops/adaptive_max_pool1d/adaptive_max_pool1d_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/adaptive_max_pool1d.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::adaptive_max_pool1d_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopAdaptiveMaxPool1dDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyAdaptiveMaxPool1dDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor y, Tensor x, size_t out) { - size_t seed = hash_combine(y, x, out); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopAdaptiveMaxPool1dDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateAdaptiveMaxPool1dDescriptor( - context::getInfiniopHandle(y->device()), &desc, - y->desc(), x->desc(), out)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAdaptiveMaxPool1dWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopAdaptiveMaxPool1d( - desc, workspace->data(), workspace_size, - y->data(), x->data(), context::getStream())); -} - -static bool registered = []() { - AdaptiveMaxPool1d::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::adaptive_max_pool1d_impl::infiniop diff --git a/src/infinicore/ops/add/add.cc b/src/infinicore/ops/add/add.cc deleted file mode 100644 index 815a2de27..000000000 --- a/src/infinicore/ops/add/add.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/add.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Add); - -Add::Add(Tensor c, const Tensor &a, const Tensor &b) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - INFINICORE_GRAPH_OP_DISPATCH(c->device().getType(), c, a, b); -} - -void Add::execute(Tensor c, const Tensor &a, const Tensor &b) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Add, c, a, b); -} - -Tensor add(const Tensor &a, const Tensor &b) { - auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); - add_(c, a, b); - return c; -} - -void add_(Tensor c, const Tensor &a, const Tensor &b) { - Add::execute(c, a, b); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/add/add_infiniop.cc b/src/infinicore/ops/add/add_infiniop.cc deleted file mode 100644 index bb377d667..000000000 --- a/src/infinicore/ops/add/add_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "infinicore/ops/add.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::add_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Add, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, c, a, b; -}; - -void *plan(Tensor c, const Tensor &a, const Tensor &b) { - size_t seed = hash_combine(c, b, a); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Add, - seed, - c->desc(), a->desc(), b->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Add, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(c), - graph::GraphTensor(a), - graph::GraphTensor(b)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopAdd( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->c->data(), - planned->a->data(), - planned->b->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Add, &plan, &run, &cleanup); - -} // namespace infinicore::op::add_impl::infiniop diff --git a/src/infinicore/ops/add/add_infiniops.cc b/src/infinicore/ops/add/add_infiniops.cc deleted file mode 100644 index 4d505bcaf..000000000 --- a/src/infinicore/ops/add/add_infiniops.cc +++ /dev/null @@ -1,59 +0,0 @@ -#include "infinicore/ops/add.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -namespace infinicore::op::add_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -struct PlannedMeta { - TensorMeta c, a, b; - graph::GraphTensor c_tensor, a_tensor, b_tensor; -}; - -} // namespace - -void *plan(Tensor c, const Tensor &a, const Tensor &b) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(c->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - - return new PlannedMeta{ - TensorMeta(c), - TensorMeta(a), - TensorMeta(b), - graph::GraphTensor(c), - graph::GraphTensor(a), - graph::GraphTensor(b)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - infini::ops::Add::Call( - handle, - config, - planned->a.tensor(planned->a_tensor), - planned->b.tensor(planned->b_tensor), - planned->c.tensor(planned->c_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(Add::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(Add::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(Add::cleanup_dispatcher(), &cleanup); - return true; -}(); - -} // namespace infinicore::op::add_impl::infiniops -#endif diff --git a/src/infinicore/ops/add_rms_norm/add_rms_norm.cc b/src/infinicore/ops/add_rms_norm/add_rms_norm.cc deleted file mode 100644 index ccba62e21..000000000 --- a/src/infinicore/ops/add_rms_norm/add_rms_norm.cc +++ /dev/null @@ -1,33 +0,0 @@ -#include "infinicore/ops/add_rms_norm.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(AddRMSNorm); - -AddRMSNorm::AddRMSNorm(Tensor y, Tensor residual_out, const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, residual_out, a, b, weight); - INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, residual_out, a, b, weight, epsilon); -} - -void AddRMSNorm::execute(Tensor y, Tensor residual_out, const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(AddRMSNorm, y, residual_out, a, b, weight, epsilon); -} - -std::pair add_rms_norm(const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon) { - auto y = Tensor::empty(a->shape(), a->dtype(), a->device()); - auto residual_out = Tensor::empty(a->shape(), a->dtype(), a->device()); - add_rms_norm_(y, residual_out, a, b, weight, epsilon); - return std::make_pair(y, residual_out); -} - -void add_rms_norm_(Tensor out, Tensor residual, const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon) { - AddRMSNorm::execute(out, residual, a, b, weight, epsilon); -} - -void add_rms_norm_inplace(Tensor input, Tensor residual, const Tensor &weight, float epsilon) { - add_rms_norm_(input, residual, input, residual, weight, epsilon); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/add_rms_norm/add_rms_norm_infiniop.cc b/src/infinicore/ops/add_rms_norm/add_rms_norm_infiniop.cc deleted file mode 100644 index 53d30a2c7..000000000 --- a/src/infinicore/ops/add_rms_norm/add_rms_norm_infiniop.cc +++ /dev/null @@ -1,53 +0,0 @@ -#include "infinicore/ops/add_rms_norm.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::add_rms_norm_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, AddRMSNorm, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, residual, a, b, weight; - float epsilon; -}; - -void *plan(Tensor y, Tensor residual_out, const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon) { - size_t seed = hash_combine(y, residual_out, a, b, weight, epsilon); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, AddRMSNorm, - seed, y->desc(), residual_out->desc(), - a->desc(), b->desc(), weight->desc(), epsilon); - - INFINIOP_WORKSPACE_TENSOR(workspace, AddRMSNorm, descriptor); - - auto planned = new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(y), - graph::GraphTensor(residual_out), - graph::GraphTensor(a), - graph::GraphTensor(b), - graph::GraphTensor(weight), - epsilon}; - - return planned; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopAddRMSNorm( - planned->descriptor->desc, planned->workspace->data(), planned->workspace->numel(), - planned->out->data(), planned->residual->data(), planned->a->data(), planned->b->data(), planned->weight->data(), context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(AddRMSNorm, &plan, &run, &cleanup); - -} // namespace infinicore::op::add_rms_norm_impl::infiniop diff --git a/src/infinicore/ops/add_rms_norm/add_rms_norm_infiniops.cc b/src/infinicore/ops/add_rms_norm/add_rms_norm_infiniops.cc deleted file mode 100644 index f03cdec78..000000000 --- a/src/infinicore/ops/add_rms_norm/add_rms_norm_infiniops.cc +++ /dev/null @@ -1,72 +0,0 @@ -#include "infinicore/ops/add_rms_norm.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/add_rms_norm.h" - -#include - -namespace infinicore::op::add_rms_norm_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -struct PlannedMeta { - TensorMeta out, residual, a, b, weight; - graph::GraphTensor out_tensor, residual_tensor, a_tensor, b_tensor, weight_tensor; - float epsilon; -}; - -} // namespace - -void *plan(Tensor out, Tensor residual, const Tensor &a, const Tensor &b, const Tensor &weight, float epsilon) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(out->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, residual, a, b, weight); - - return new PlannedMeta{ - TensorMeta(out), - TensorMeta(residual), - TensorMeta(a), - TensorMeta(b), - TensorMeta(weight), - graph::GraphTensor(out), - graph::GraphTensor(residual), - graph::GraphTensor(a), - graph::GraphTensor(b), - graph::GraphTensor(weight), - epsilon}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - infini::ops::AddRmsNorm::Call( - handle, - config, - planned->a.tensor(planned->a_tensor), - planned->b.tensor(planned->b_tensor), - planned->weight.tensor(planned->weight_tensor), - std::optional{planned->epsilon}, - planned->out.tensor(planned->out_tensor), - planned->residual.tensor(planned->residual_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(AddRMSNorm::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(AddRMSNorm::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(AddRMSNorm::cleanup_dispatcher(), &cleanup); - return true; -}(); - -} // namespace infinicore::op::add_rms_norm_impl::infiniops -#endif diff --git a/src/infinicore/ops/addbmm/addbmm.cc b/src/infinicore/ops/addbmm/addbmm.cc deleted file mode 100644 index 88d51e06c..000000000 --- a/src/infinicore/ops/addbmm/addbmm.cc +++ /dev/null @@ -1,32 +0,0 @@ -#include "infinicore/ops/addbmm.hpp" -#include "../../utils.hpp" -#include "infinicore/ops/addbmm.hpp" - -namespace infinicore::op { - -// 1. 初始化 Dispatcher -common::OpDispatcher &Addbmm::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Addbmm::execute(Tensor output, Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha) { - - // 切换上下文 - infinicore::context::setDevice(output->device()); - - // 分发计算 - dispatcher().lookup(output->device().getType())(output, input, batch1, batch2, beta, alpha); -} - -Tensor addbmm(Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha) { - auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); - Addbmm::execute(output, input, batch1, batch2, beta, alpha); - return output; -} - -void addbmm_(Tensor output, Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha) { - Addbmm::execute(output, input, batch1, batch2, beta, alpha); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/addbmm/addbmm_infiniop.cc b/src/infinicore/ops/addbmm/addbmm_infiniop.cc deleted file mode 100644 index b0e2d14d7..000000000 --- a/src/infinicore/ops/addbmm/addbmm_infiniop.cc +++ /dev/null @@ -1,106 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/addbmm.hpp" -#include "infinicore/ops/common/cache.hpp" -#include -#include - -namespace infinicore::op::addbmm_impl::infiniop { - -struct AddbmmContext { - infiniopAddbmmDescriptor_t desc = nullptr; - std::shared_ptr workspace_buf = nullptr; - size_t workspace_size = 0; - - void *getWorkspacePtr() const { - return workspace_buf ? workspace_buf->data() : nullptr; - } -}; - -thread_local common::OpCache caches( - 256, - [](AddbmmContext &ctx) { - if (ctx.desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyAddbmmDescriptor(ctx.desc)); - ctx.desc = nullptr; - } - ctx.workspace_buf = nullptr; - }); - -inline size_t compute_key(const Tensor &output, const Tensor &input, - const Tensor &batch1, const Tensor &batch2, - float beta, float alpha) { - size_t seed = 0; - infinicore::hash_combine(seed, reinterpret_cast(output.operator->())); - infinicore::hash_combine(seed, reinterpret_cast(input.operator->())); - infinicore::hash_combine(seed, reinterpret_cast(batch1.operator->())); - infinicore::hash_combine(seed, reinterpret_cast(batch2.operator->())); - infinicore::hash_combine(seed, beta); - infinicore::hash_combine(seed, alpha); - return seed; -} - -void calculate(Tensor output, Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha) { - size_t seed = compute_key(output, input, batch1, batch2, beta, alpha); - - static thread_local size_t last_seed = 0; - static thread_local bool last_ctx_valid = false; - static thread_local AddbmmContext last_ctx; - - AddbmmContext *ctx_ptr = nullptr; - - if (last_ctx_valid && seed == last_seed) { - ctx_ptr = &last_ctx; - } else { - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - auto &cache = caches.getCache(device_type, device_index); - - auto opt_ctx = cache.get(seed); - if (opt_ctx) { - last_ctx = *opt_ctx; - } else { - AddbmmContext new_ctx; - - INFINICORE_CHECK_ERROR(infiniopCreateAddbmmDescriptor( - context::getInfiniopHandle(output->device()), - &new_ctx.desc, - output->desc(), - input->desc(), - batch1->desc(), - batch2->desc(), - alpha, - beta)); - - INFINICORE_CHECK_ERROR(infiniopGetAddbmmWorkspaceSize(new_ctx.desc, &new_ctx.workspace_size)); - - if (new_ctx.workspace_size > 0) { - new_ctx.workspace_buf = context::allocateMemory(new_ctx.workspace_size); - } - - cache.put(seed, new_ctx); - last_ctx = new_ctx; - } - - last_seed = seed; - last_ctx_valid = true; - ctx_ptr = &last_ctx; - } - - INFINICORE_CHECK_ERROR(infiniopAddbmm( - ctx_ptr->desc, - ctx_ptr->getWorkspacePtr(), - ctx_ptr->workspace_size, - output->data(), - input->data(), - batch1->data(), - batch2->data(), - context::getStream())); -} - -static bool registered = []() { - Addbmm::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::addbmm_impl::infiniop diff --git a/src/infinicore/ops/addcmul/addcmul.cc b/src/infinicore/ops/addcmul/addcmul.cc deleted file mode 100644 index 43c5b1377..000000000 --- a/src/infinicore/ops/addcmul/addcmul.cc +++ /dev/null @@ -1,29 +0,0 @@ -#include "infinicore/ops/addcmul.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Addcmul::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -// 执行核心逻辑：设备校验与后端分发 -void Addcmul::execute(Tensor out, Tensor input, Tensor t1, Tensor t2, float value) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, t1, t2); - infinicore::context::setDevice(out->device()); - dispatcher().lookup(out->device().getType())(out, input, t1, t2, value); -} - -// Out-of-place 接口：自动创建输出 Tensor -Tensor addcmul(Tensor input, Tensor t1, Tensor t2, float value) { - auto out = Tensor::empty(input->shape(), input->dtype(), input->device()); - addcmul_(out, input, t1, t2, value); - return out; -} - -void addcmul_(Tensor out, Tensor input, Tensor t1, Tensor t2, float value) { - Addcmul::execute(out, input, t1, t2, value); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/addcmul/addcmul_infiniop.cc b/src/infinicore/ops/addcmul/addcmul_infiniop.cc deleted file mode 100644 index 0558780b2..000000000 --- a/src/infinicore/ops/addcmul/addcmul_infiniop.cc +++ /dev/null @@ -1,51 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/addcmul.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::addcmul_impl::infiniop { - -// 定义线程局部的算子描述符缓存 -thread_local common::OpCache caches( - 100, - [](infiniopAddcmulDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyAddcmulDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor out, Tensor input, Tensor t1, Tensor t2, float value) { - size_t seed = hash_combine(out, input, t1, t2, value); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopAddcmulDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateAddcmulDescriptor( - context::getInfiniopHandle(device), &desc, - out->desc(), input->desc(), t1->desc(), t2->desc(), value)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAddcmulWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopAddcmul( - desc, workspace->data(), workspace_size, - out->data(), input->data(), t1->data(), t2->data(), context::getStream())); -} - -static bool registered = []() { - Addcmul::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::addcmul_impl::infiniop diff --git a/src/infinicore/ops/addr/addr.cc b/src/infinicore/ops/addr/addr.cc deleted file mode 100644 index 4691e13a8..000000000 --- a/src/infinicore/ops/addr/addr.cc +++ /dev/null @@ -1,32 +0,0 @@ -#include "infinicore/ops/addr.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Addr::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Addr::execute(Tensor out, Tensor input, Tensor vec1, Tensor vec2, float beta, float alpha) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, vec1, vec2); - infinicore::context::setDevice(out->device()); - dispatcher().lookup(out->device().getType())(out, input, vec1, vec2, beta, alpha); -} - -Tensor addr(Tensor input, Tensor vec1, Tensor vec2, float beta, float alpha) { - - size_t n = vec1->shape()[0]; - size_t m = vec2->shape()[0]; - - // Create output tensor - Tensor out = Tensor::empty({n, m}, input->dtype(), input->device()); - addr_(out, input, vec1, vec2, beta, alpha); - return out; -} - -void addr_(Tensor out, Tensor input, Tensor vec1, Tensor vec2, float beta, float alpha) { - Addr::execute(out, input, vec1, vec2, beta, alpha); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/addr/addr_infiniop.cc b/src/infinicore/ops/addr/addr_infiniop.cc deleted file mode 100644 index 1c0ca3733..000000000 --- a/src/infinicore/ops/addr/addr_infiniop.cc +++ /dev/null @@ -1,53 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/addr.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::addr_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopAddrDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyAddrDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor out, Tensor input, Tensor vec1, Tensor vec2, float beta, float alpha) { - // Hash the inputs including beta and alpha - size_t seed = hash_combine(out, input, vec1, vec2, beta, alpha); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopAddrDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateAddrDescriptor( - context::getInfiniopHandle(out->device()), &desc, - out->desc(), input->desc(), vec1->desc(), vec2->desc(), beta, alpha)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAddrWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopAddr( - desc, workspace->data(), workspace_size, - out->data(), input->data(), vec1->data(), vec2->data(), context::getStream())); -} - -static bool registered = []() { - Addr::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::addr_impl::infiniop diff --git a/src/infinicore/ops/affine_grid/affine_grid.cc b/src/infinicore/ops/affine_grid/affine_grid.cc deleted file mode 100644 index 2f22850b1..000000000 --- a/src/infinicore/ops/affine_grid/affine_grid.cc +++ /dev/null @@ -1,58 +0,0 @@ -#include "infinicore/ops/affine_grid.hpp" -#include -#include -#include - -namespace infinicore::op { - -common::OpDispatcher &AffineGrid::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void AffineGrid::execute(Tensor output, Tensor theta, bool align_corners) { - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No AffineGrid implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, theta, align_corners); -} - -Tensor affine_grid(Tensor theta, const std::vector &size, bool align_corners) { - if (theta->ndim() != 3) { - throw std::runtime_error("AffineGrid: Theta tensor must be 3D (N, 2, 3)."); - } - if (theta->shape()[1] != 2 || theta->shape()[2] != 3) { - throw std::runtime_error("AffineGrid: Theta tensor shape must be (N, 2, 3)."); - } - - if (size.size() != 4) { - throw std::runtime_error("AffineGrid: target size length must be 4 (N, C, H, W)."); - } - - if (static_cast(theta->shape()[0]) != size[0]) { - throw std::runtime_error("AffineGrid: Theta batch size does not match target size batch."); - } - - if (!theta->is_contiguous()) { - theta = theta->contiguous(); - } - - std::vector out_shape; - out_shape.reserve(4); - out_shape.push_back(static_cast(size[0])); - out_shape.push_back(static_cast(size[2])); - out_shape.push_back(static_cast(size[3])); - out_shape.push_back(2); - - auto output = Tensor::empty(out_shape, theta->dtype(), theta->device()); - - AffineGrid::execute(output, theta, align_corners); - - return output; -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/affine_grid/affine_grid_infiniop.cc b/src/infinicore/ops/affine_grid/affine_grid_infiniop.cc deleted file mode 100644 index fcbef4359..000000000 --- a/src/infinicore/ops/affine_grid/affine_grid_infiniop.cc +++ /dev/null @@ -1,62 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/affine_grid.hpp" // 引用算子定义 -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::affine_grid_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopAffineGridDescriptor_t &desc) { - if (desc != nullptr) { - // 销毁描述符 - INFINICORE_CHECK_ERROR(infiniopDestroyAffineGridDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor theta, bool align_corners) { - - size_t seed = hash_combine(output, theta, align_corners); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - // 获取当前设备对应的缓存 - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopAffineGridDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateAffineGridDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - theta->desc(), - align_corners)); // 传递 align_corners - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAffineGridWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopAffineGrid( - desc, - workspace->data(), - workspace_size, - output->data(), - theta->data(), - context::getStream())); -} - -static bool registered = []() { - AffineGrid::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::affine_grid_impl::infiniop diff --git a/src/infinicore/ops/all/all.cc b/src/infinicore/ops/all/all.cc deleted file mode 100644 index c695623b8..000000000 --- a/src/infinicore/ops/all/all.cc +++ /dev/null @@ -1,67 +0,0 @@ -#include "infinicore/ops/all.hpp" - -#include "../../utils.hpp" -#include -#include -#include -namespace infinicore::op { - -common::OpDispatcher &All::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; -void All::execute(Tensor output, Tensor input, std::vector dim, bool keepdim) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(input->device()); - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No All implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input, dim, keepdim); -} - -Tensor all(Tensor input, std::vector dim, bool keepdim) { - auto in_shape = input->shape(); - std::vector out_shape; - if (dim.empty()) { - for (size_t i = 0; i < in_shape.size(); i++) { - dim.push_back(i); - } - } - std::sort(dim.begin(), dim.end()); - if (dim.size() == in_shape.size() && !keepdim) { - out_shape = {}; - } else { - if (keepdim) { - size_t j = 0; - for (size_t i = 0; i < in_shape.size(); i++) { - if (j < dim.size() && dim[j] == i) { - out_shape.push_back(1); - j++; - } else { - out_shape.push_back(in_shape[i]); - } - } - } else { - size_t j = 0; - for (size_t i = 0; i < in_shape.size(); i++) { - if (j < dim.size() && dim[j] == i) { - j++; - } else { - out_shape.push_back(in_shape[i]); - } - } - } - } - auto output = Tensor::empty(out_shape, DataType::BOOL, input->device()); - all_(output, input, dim, keepdim); - return output; -} - -void all_(Tensor output, Tensor input, std::vector dim, bool keepdim) { - All::execute(output, input, dim, keepdim); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/all/all_infiniop.cc b/src/infinicore/ops/all/all_infiniop.cc deleted file mode 100644 index 3d24aca12..000000000 --- a/src/infinicore/ops/all/all_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/all.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::all_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopAllDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyAllDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, std::vector dim, bool keepdim) { - size_t seed = hash_combine(output, input, dim.size(), keepdim); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopAllDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateAllDescriptor( - context::getInfiniopHandle(output->device()), &desc, - output->desc(), input->desc(), dim.data(), dim.size(), keepdim)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAllWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopAll( - desc, workspace->data(), workspace_size, - output->data(), input->data(), dim.data(), dim.size(), keepdim, context::getStream())); -} - -static bool registered = []() { - All::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::all_impl::infiniop diff --git a/src/infinicore/ops/argwhere/argwhere.cc b/src/infinicore/ops/argwhere/argwhere.cc deleted file mode 100644 index 6e1f83615..000000000 --- a/src/infinicore/ops/argwhere/argwhere.cc +++ /dev/null @@ -1,33 +0,0 @@ -#include "infinicore/ops/argwhere.hpp" -#include "../../utils.hpp" -#include "infinicore.h" -#include "infinicore/dtype.hpp" -#include "infinicore/tensor.hpp" -#include - -namespace infinicore::op { - -common::OpDispatcher &Argwhere::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -void Argwhere::execute(void **y, size_t *count, Tensor x) { - auto device_type = context::getDevice().getType(); - - auto func = dispatcher().lookup(device_type); - if (func == nullptr) { - throw std::runtime_error("Argwhere op not implemented for device type " + std::to_string(static_cast(device_type))); - } - func(y, count, x); -} -Tensor argwhere(Tensor x) { - void *y = nullptr; - size_t count = 0; - Argwhere::execute(&y, &count, x); - auto result = Tensor::from_blob(y, Shape{count, x->ndim()}, DataType::I64, Device::cpu()); - result = result->to(x->device()); - return result; -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/argwhere/argwhere_infiniop.cc b/src/infinicore/ops/argwhere/argwhere_infiniop.cc deleted file mode 100644 index 7e04767bf..000000000 --- a/src/infinicore/ops/argwhere/argwhere_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/argwhere.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infiniop/ops/argwhere.h" - -namespace infinicore::op::argwhere_impl::infiniop { -thread_local common::OpCache caches( - 100, // capacity - [](infiniopArgwhereDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyArgwhereDescriptor(desc)); - desc = nullptr; - } - }); -void calculate(void **y, size_t *count, Tensor x) { - size_t seed = hash_combine(x); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopArgwhereDescriptor_t desc = nullptr; - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateArgwhereDescriptor( - context::getInfiniopHandle(x->device()), - &desc, - x->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetArgwhereWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopArgwhere( - desc, - workspace->data(), workspace_size, - y, - count, - x->data(), - context::getStream())); -} -static bool registered = []() { - Argwhere::dispatcher().registerAll(&calculate, false); - return true; -}(); -} // namespace infinicore::op::argwhere_impl::infiniop diff --git a/src/infinicore/ops/asin/asin.cc b/src/infinicore/ops/asin/asin.cc deleted file mode 100644 index f3a8efb7c..000000000 --- a/src/infinicore/ops/asin/asin.cc +++ /dev/null @@ -1,36 +0,0 @@ -#include "infinicore/ops/asin.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &Asin::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; -void Asin::execute(Tensor output, Tensor input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Asin implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor asin(Tensor input) { - Shape shape = input->shape(); - auto output = Tensor::empty(shape, input->dtype(), input->device()); - asin_(output, input); - return output; -} - -void asin_(Tensor output, Tensor input) { - Asin::execute(output, input); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/asin/asin_infiniop.cc b/src/infinicore/ops/asin/asin_infiniop.cc deleted file mode 100644 index a794a4682..000000000 --- a/src/infinicore/ops/asin/asin_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/asin.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::asin_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopAsinDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyAsinDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopAsinDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateAsinDescriptor( - context::getInfiniopHandle(output->device()), &desc, - output->desc(), input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAsinWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopAsin( - desc, workspace->data(), workspace_size, - output->data(), input->data(), context::getStream())); -} - -static bool registered = []() { - Asin::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::asin_impl::infiniop diff --git a/src/infinicore/ops/asinh/asinh.cc b/src/infinicore/ops/asinh/asinh.cc deleted file mode 100644 index fbf131d99..000000000 --- a/src/infinicore/ops/asinh/asinh.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/asinh.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Asinh::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Asinh::execute(Tensor y, Tensor x) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - infinicore::context::setDevice(y->device()); - dispatcher().lookup(y->device().getType())(y, x); -} - -Tensor asinh(Tensor x) { - auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); - asinh_(y, x); - return y; -} - -void asinh_(Tensor y, Tensor x) { - Asinh::execute(y, x); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/asinh/asinh_infiniop.cc b/src/infinicore/ops/asinh/asinh_infiniop.cc deleted file mode 100644 index ceed8d5a2..000000000 --- a/src/infinicore/ops/asinh/asinh_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/asinh.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::asinh_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopAsinhDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyAsinhDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor y, Tensor x) { - size_t seed = hash_combine(y, x); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopAsinhDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateAsinhDescriptor( - context::getInfiniopHandle(y->device()), &desc, - y->desc(), x->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAsinhWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopAsinh( - desc, workspace->data(), workspace_size, - y->data(), x->data(), context::getStream())); -} - -static bool registered = []() { - Asinh::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::asinh_impl::infiniop diff --git a/src/infinicore/ops/asum/asum.cc b/src/infinicore/ops/asum/asum.cc deleted file mode 100644 index c757574bb..000000000 --- a/src/infinicore/ops/asum/asum.cc +++ /dev/null @@ -1,28 +0,0 @@ -#include "infinicore/ops/asum.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Asum); - -Asum::Asum(const Tensor &x, Tensor result) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, result); - INFINICORE_GRAPH_OP_DISPATCH(result->device().getType(), x, result); -} - -void Asum::execute(const Tensor &x, Tensor result) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Asum, x, result); -} - -Tensor asum(const Tensor &x) { - auto result = Tensor::empty({}, x->dtype(), x->device()); - asum_(x, result); - return result; -} - -void asum_(const Tensor &x, Tensor result) { - Asum::execute(x, result); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/asum/asum_infiniop.cc b/src/infinicore/ops/asum/asum_infiniop.cc deleted file mode 100644 index 0cfd8b721..000000000 --- a/src/infinicore/ops/asum/asum_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/asum.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::asum_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Asum, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, x, result; -}; - -void *plan(const Tensor &x, Tensor result) { - size_t seed = hash_combine(x, result); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Asum, - seed, - x->desc(), result->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Asum, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x), - graph::GraphTensor(result)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopAsum( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->x->data(), - planned->result->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Asum, &plan, &run, &cleanup); - -} // namespace infinicore::op::asum_impl::infiniop diff --git a/src/infinicore/ops/atanh/atanh.cc b/src/infinicore/ops/atanh/atanh.cc deleted file mode 100644 index f104f14e5..000000000 --- a/src/infinicore/ops/atanh/atanh.cc +++ /dev/null @@ -1,37 +0,0 @@ -#include "infinicore/ops/atanh.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -// 获取单例分发器 -common::OpDispatcher &Atanh::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -// 执行入口：负责设备切换和后端查找 -void Atanh::execute(Tensor y, Tensor a) { - // 确保输入和输出在同一个设备上 - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, a); - - // 切换当前上下文到目标设备 - infinicore::context::setDevice(y->device()); - - // 根据设备类型（CPU/CUDA等）查找对应的实现并执行 - dispatcher().lookup(y->device().getType())(y, a); -} - -// Out-of-place 接口：自动创建结果 Tensor -Tensor atanh(Tensor a) { - // 创建一个与输入形状、类型、设备完全相同的空 Tensor - auto y = Tensor::empty(a->shape(), a->dtype(), a->device()); - atanh_(y, a); - return y; -} - -// In-place 或指定输出接口 -void atanh_(Tensor y, Tensor a) { - Atanh::execute(y, a); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/atanh/atanh_infiniop.cc b/src/infinicore/ops/atanh/atanh_infiniop.cc deleted file mode 100644 index 89a3da1c1..000000000 --- a/src/infinicore/ops/atanh/atanh_infiniop.cc +++ /dev/null @@ -1,58 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/atanh.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::atanh_impl::infiniop { - -// 定义线程局部的算子描述符缓存，避免重复创建 Descriptor 带来的开销 -thread_local common::OpCache caches( - 100, // 缓存容量 - [](infiniopAtanhDescriptor_t &desc) { - if (desc != nullptr) { - // 缓存释放时的回调：销毁 infiniop 算子描述符 - INFINICORE_CHECK_ERROR(infiniopDestroyAtanhDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor y, Tensor a) { - // 1. 根据 Tensor 的形状、步长、类型等信息生成唯一 Hash 值 - size_t seed = hash_combine(y, a); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - // 2. 尝试从缓存中获取已存在的描述符 - auto desc_opt = cache.get(seed); - infiniopAtanhDescriptor_t desc = nullptr; - - if (!desc_opt) { - // 如果缓存未命中，创建新的描述符 - INFINICORE_CHECK_ERROR(infiniopCreateAtanhDescriptor( - context::getInfiniopHandle(device), &desc, - y->desc(), a->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 3. 获取并分配必要的 Workspace 空间（如果有的话） - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAtanhWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - // 4. 执行底层计算 - INFINICORE_CHECK_ERROR(infiniopAtanh( - desc, workspace->data(), workspace_size, - y->data(), a->data(), context::getStream())); -} - -// 5. 自动注册逻辑：程序启动时将此实现注册到分发器中 -static bool registered = []() { - Atanh::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::atanh_impl::infiniop diff --git a/src/infinicore/ops/attention/attention.cc b/src/infinicore/ops/attention/attention.cc deleted file mode 100644 index 053798bb3..000000000 --- a/src/infinicore/ops/attention/attention.cc +++ /dev/null @@ -1,40 +0,0 @@ -#include "infinicore/ops/attention.hpp" -#include "../../utils.hpp" - -#ifdef ENABLE_MUTUAL_AWARENESS -#include "infinicore/analyzer/mutual_awareness_analyzer.hpp" -#endif - -namespace infinicore::op { - -common::OpDispatcher &Attention::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Attention::execute(Tensor out, Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k, v, k_cache, v_cache); - infinicore::context::setDevice(out->device()); -#ifdef ENABLE_MUTUAL_AWARENESS - auto goal = analyzer::MutualAwarenessAnalyzer::instance().getCurrentOptimizationGoal(); - dispatcher().lookup(out->device().getType(), goal)(out, q, k, v, k_cache, v_cache, pos); -#else - dispatcher().lookup(out->device().getType())(out, q, k, v, k_cache, v_cache, pos); -#endif -} - -Tensor attention(Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos) { - size_t n_q_head = q->shape()[0]; - size_t seq_len = q->shape()[1]; - size_t head_dim = q->shape()[2]; - Shape shape = {seq_len, n_q_head, head_dim}; - auto out = Tensor::empty(shape, q->dtype(), q->device()); - attention_(out, q, k, v, k_cache, v_cache, pos); - return out; -} - -void attention_(Tensor out, Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos) { - Attention::execute(out, q, k, v, k_cache, v_cache, pos); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/attention/attention_infiniop.cc b/src/infinicore/ops/attention/attention_infiniop.cc deleted file mode 100644 index 76b27bf46..000000000 --- a/src/infinicore/ops/attention/attention_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/attention.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::attention_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopAttentionDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyAttentionDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor out, Tensor q, Tensor k, Tensor v, Tensor k_cache, Tensor v_cache, size_t pos) { - size_t seed = hash_combine(out, q, k, v, k_cache, v_cache, pos); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopAttentionDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateAttentionDescriptor( - context::getInfiniopHandle(device), &desc, - out->desc(), q->desc(), k->desc(), v->desc(), - k_cache->desc(), v_cache->desc(), pos)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAttentionWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopAttention( - desc, workspace->data(), workspace_size, - out->data(), q->data(), k->data(), v->data(), - k_cache->data(), v_cache->data(), context::getStream())); -} - -static bool registered = []() { - Attention::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::attention_impl::infiniop diff --git a/src/infinicore/ops/avg_pool1d/avg_pool1d.cc b/src/infinicore/ops/avg_pool1d/avg_pool1d.cc deleted file mode 100644 index 907b25b00..000000000 --- a/src/infinicore/ops/avg_pool1d/avg_pool1d.cc +++ /dev/null @@ -1,68 +0,0 @@ -#include "infinicore/ops/avg_pool1d.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &AvgPool1d::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -void AvgPool1d::execute( - Tensor output, - Tensor input, - size_t kernel_size, - size_t stride, - size_t padding) { - - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - if (stride == 0) { - stride = kernel_size; - } - - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error( - "No AvgPool1d implementation for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input, kernel_size, stride, padding); -} - -Tensor avg_pool1d(Tensor input, size_t kernel_size, size_t stride, size_t padding) { - if (stride == 0) { - stride = kernel_size; - } - - const auto &shape = input->shape(); - if (shape.size() != 3) { - throw std::runtime_error("AvgPool1d expects tensors with shape [N, C, L]"); - } - - const size_t n = shape[0]; - const size_t c = shape[1]; - const size_t l_in = shape[2]; - - if (l_in + 2 * padding < kernel_size) { - throw std::runtime_error("AvgPool1d kernel_size is larger than padded length"); - } - - const size_t out_width = (l_in + 2 * padding - kernel_size) / stride + 1; - - Shape out_shape = {n, c, out_width}; - auto output = Tensor::empty(out_shape, input->dtype(), input->device()); - avg_pool1d_(output, input, kernel_size, stride, padding); - return output; -} - -void avg_pool1d_(Tensor output, Tensor input, size_t kernel_size, size_t stride, size_t padding) { - AvgPool1d::execute(output, input, kernel_size, stride, padding); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/avg_pool1d/avg_pool1d_infiniop.cc b/src/infinicore/ops/avg_pool1d/avg_pool1d_infiniop.cc deleted file mode 100644 index df7ebda8d..000000000 --- a/src/infinicore/ops/avg_pool1d/avg_pool1d_infiniop.cc +++ /dev/null @@ -1,69 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/avg_pool1d.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::avg_pool1d_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopAvgPool1dDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyAvgPool1dDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate( - Tensor output, - Tensor input, - size_t kernel_size, - size_t stride, - size_t padding) { - - if (stride == 0) { - stride = kernel_size; - } - - size_t seed = hash_combine(output, input, kernel_size, stride, padding); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopAvgPool1dDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateAvgPool1dDescriptor( - context::getInfiniopHandle(device), - &desc, - output->desc(), - input->desc(), - kernel_size, - stride, - padding)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAvgPool1dWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopAvgPool1d( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -static bool registered = []() { - AvgPool1d::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::avg_pool1d_impl::infiniop diff --git a/src/infinicore/ops/awq_marlin_gemm/awq_marlin_gemm.cc b/src/infinicore/ops/awq_marlin_gemm/awq_marlin_gemm.cc deleted file mode 100644 index eb275cf3c..000000000 --- a/src/infinicore/ops/awq_marlin_gemm/awq_marlin_gemm.cc +++ /dev/null @@ -1,21 +0,0 @@ -#include "infinicore/ops/awq_marlin_gemm.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(AwqMarlinGemm); - -AwqMarlinGemm::AwqMarlinGemm(Tensor c, const Tensor &a, const Tensor &b, Tensor &b_bias, Tensor &b_scales, Tensor &a_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b, b_bias, b_scales, a_scales, global_scales, b_zeros, g_idx, perm); - INFINICORE_GRAPH_OP_DISPATCH(c->device().getType(), c, a, b, b_bias, b_scales, a_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float); -} -void AwqMarlinGemm::execute(Tensor c, const Tensor &a, const Tensor &b, Tensor &b_bias, Tensor &b_scales, Tensor &a_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(AwqMarlinGemm, c, a, b, b_bias, b_scales, a_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float); -} - -void awq_marlin_gemm_(Tensor c, const Tensor &a, const Tensor &b, Tensor &b_bias, Tensor &b_scales, Tensor &a_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - AwqMarlinGemm::execute(c, a, b, b_bias, b_scales, a_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/awq_marlin_gemm/awq_marlin_gemm_infiniop.cc b/src/infinicore/ops/awq_marlin_gemm/awq_marlin_gemm_infiniop.cc deleted file mode 100644 index 345f88c3b..000000000 --- a/src/infinicore/ops/awq_marlin_gemm/awq_marlin_gemm_infiniop.cc +++ /dev/null @@ -1,87 +0,0 @@ -#include "../../utils.hpp" -#include "../infiniop_impl.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/awq_marlin_gemm.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::awq_marlin_gemm_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, AwqMarlinGemm, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, c, a, b, b_bias, b_scales, a_scales, global_scales, b_zeros, g_idx, perm; - int64_t b_q_type_id; - bool is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float; -}; - -void *plan(Tensor c, const Tensor &a, const Tensor &b, Tensor &b_bias, Tensor &b_scales, Tensor &a_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - size_t seed = hash_combine(c, a, b, b_bias, b_scales, a_scales, global_scales, b_zeros, g_idx, perm); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, AwqMarlinGemm, - seed, - c->desc(), a->desc(), - b->desc(), b_bias->desc(), b_scales->desc(), a_scales->desc(), global_scales->desc(), b_zeros->desc(), g_idx->desc(), perm->desc()); - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetAwqMarlinGemmWorkspaceSize(descriptor->desc, &workspace_size)); - Tensor workspace = Tensor::zeros({workspace_size}, DataType::U8, context::getDevice()); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(c), - graph::GraphTensor(a), - graph::GraphTensor(b), - graph::GraphTensor(b_bias), - graph::GraphTensor(b_scales), - graph::GraphTensor(a_scales), - graph::GraphTensor(global_scales), - graph::GraphTensor(b_zeros), - graph::GraphTensor(g_idx), - graph::GraphTensor(perm), - b_q_type_id, - is_k_full, - use_atomic_add, - use_fp32_reduce, - is_zp_float}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - auto optional_data = [](graph::GraphTensor &tensor) -> void * { - return tensor->numel() == 0 ? nullptr : tensor->data(); - }; - - INFINICORE_CHECK_ERROR(infiniopAwqMarlinGemm( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->c->data(), - planned->a->data(), - planned->b->data(), - optional_data(planned->b_bias), - planned->b_scales->data(), - optional_data(planned->a_scales), - optional_data(planned->global_scales), - optional_data(planned->b_zeros), - optional_data(planned->g_idx), - optional_data(planned->perm), - planned->b_q_type_id, - planned->is_k_full, - planned->use_atomic_add, - planned->use_fp32_reduce, - planned->is_zp_float, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(AwqMarlinGemm, &plan, &run, &cleanup); - -} // namespace infinicore::op::awq_marlin_gemm_impl::infiniop diff --git a/src/infinicore/ops/axpy/axpy.cc b/src/infinicore/ops/axpy/axpy.cc deleted file mode 100644 index d5d3241ec..000000000 --- a/src/infinicore/ops/axpy/axpy.cc +++ /dev/null @@ -1,22 +0,0 @@ -#include "infinicore/ops/axpy.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Axpy); - -Axpy::Axpy(const Tensor &alpha, const Tensor &x, Tensor y) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(alpha, x, y); - INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), alpha, x, y); -} - -void Axpy::execute(const Tensor &alpha, const Tensor &x, Tensor y) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Axpy, alpha, x, y); -} - -void axpy_(const Tensor &alpha, const Tensor &x, Tensor y) { - Axpy::execute(alpha, x, y); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/axpy/axpy_infiniop.cc b/src/infinicore/ops/axpy/axpy_infiniop.cc deleted file mode 100644 index 8d54c3823..000000000 --- a/src/infinicore/ops/axpy/axpy_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "infinicore/ops/axpy.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::axpy_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Axpy, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, alpha, x, y; -}; - -void *plan(const Tensor &alpha, const Tensor &x, Tensor y) { - size_t seed = hash_combine(y, alpha, x); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Axpy, - seed, - alpha->desc(), x->desc(), y->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Axpy, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(alpha), - graph::GraphTensor(x), - graph::GraphTensor(y)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopAxpy( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->alpha->data(), - planned->x->data(), - planned->y->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Axpy, &plan, &run, &cleanup); - -} // namespace infinicore::op::axpy_impl::infiniop diff --git a/src/infinicore/ops/baddbmm/baddbmm.cc b/src/infinicore/ops/baddbmm/baddbmm.cc deleted file mode 100644 index c4f8efaf5..000000000 --- a/src/infinicore/ops/baddbmm/baddbmm.cc +++ /dev/null @@ -1,100 +0,0 @@ -#include "infinicore/ops/baddbmm.hpp" -#include "infinicore/ops/gemm.hpp" -#include "infinicore/ops/rearrange.hpp" - -namespace infinicore::op { - -// 内联的 BLAS 兼容性检查，减少函数调用开销 -inline bool is_blas_compatible(const Tensor &t) { - const auto ndim = t->ndim(); - if (ndim == 2) { - const auto rs = t->stride(0); - const auto cs = t->stride(1); - if (rs != 1 && cs != 1) { - return false; - } - if (rs == 1 && cs == 1) { - return t->shape()[0] == 1 || t->shape()[1] == 1; - } - return true; - } else if (ndim == 3) { - const auto rs = t->stride(1); - const auto cs = t->stride(2); - if (t->shape()[0] > 1 && t->stride(0) == 0) { - return false; - } - if (rs != 1 && cs != 1) { - return false; - } - if (rs == 1 && cs == 1) { - return t->shape()[1] == 1 || t->shape()[2] == 1; - } - return true; - } - return false; -} - -inline void prepare_gemm_input(Tensor &output, Tensor &input, const size_t batch_size, const size_t m, const size_t n) { - const auto input_ndim = input->ndim(); - if (input_ndim == 2) { - rearrange_(output, input->as_strided( - {batch_size, m, n}, - {0, input->stride(0), input->stride(1)})); - } else if (input_ndim == 3 && input->shape()[0] == 1 && batch_size > 1) { - rearrange_(output, input->as_strided( - {batch_size, m, n}, - {0, input->stride(1), input->stride(2)})); - } else { - rearrange_(output, input); - } -} - -Tensor baddbmm(Tensor input, Tensor batch1, Tensor batch2, - float beta, - float alpha) { - const size_t batch_size = batch1->shape()[0]; - const size_t m = batch1->shape()[1]; - const size_t n = batch2->shape()[2]; - - const Tensor &a = is_blas_compatible(batch1) ? batch1 : rearrange(batch1); - const Tensor &b = is_blas_compatible(batch2) ? batch2 : rearrange(batch2); - - if (beta == 0.0f) { - return gemm(a, b, alpha, 0.0f); - } - - Tensor result = Tensor::empty({batch_size, m, n}, a->dtype(), a->device()); - - prepare_gemm_input(result, input, batch_size, m, n); - - gemm_(result, a, b, alpha, beta); - return result; -} - -void baddbmm_(Tensor out, Tensor input, Tensor batch1, Tensor batch2, - float beta, - float alpha) { - const size_t batch_size = batch1->shape()[0]; - const size_t m = batch1->shape()[1]; - const size_t n = batch2->shape()[2]; - - const Tensor &a = is_blas_compatible(batch1) ? batch1 : rearrange(batch1); - const Tensor &b = is_blas_compatible(batch2) ? batch2 : rearrange(batch2); - - const bool out_is_usable = out->is_contiguous() && out->ndim() == 3 && out->shape()[0] == batch_size && out->shape()[1] == m && out->shape()[2] == n; - - if (out_is_usable) { - if (beta != 0.0f && input->data() != out->data()) { - prepare_gemm_input(out, input, batch_size, m, n); - } - gemm_(out, a, b, alpha, beta); - } else { - Tensor result = Tensor::empty({batch_size, m, n}, a->dtype(), a->device()); - if (beta != 0.0f) { - prepare_gemm_input(result, input, batch_size, m, n); - } - gemm_(result, a, b, alpha, beta); - rearrange_(out, result); - } -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/bilinear/bilinear.cc b/src/infinicore/ops/bilinear/bilinear.cc deleted file mode 100644 index 930f9a31a..000000000 --- a/src/infinicore/ops/bilinear/bilinear.cc +++ /dev/null @@ -1,95 +0,0 @@ -#include "infinicore/ops/bilinear.hpp" -#include "infinicore/ops/add.hpp" -#include "infinicore/ops/matmul.hpp" -#include "infinicore/ops/rearrange.hpp" - -namespace infinicore::op { - -namespace { -inline bool is_gemm_compatible_3d(const Tensor &t) { - if (t->ndim() != 3) { - return false; - } - - const auto batch = t->shape()[0]; - const auto rows = t->shape()[1]; - const auto cols = t->shape()[2]; - const auto bs = t->stride(0); - const auto rs = t->stride(1); - const auto cs = t->stride(2); - - if (rs != 1 && cs != 1) { - return false; - } - - if (cs == 1) { - if (rs < static_cast(cols)) { - return false; - } - } else { - if (cs < static_cast(rows)) { - return false; - } - } - - if (batch > 1 && bs == 0) { - return false; - } - - return true; -} - -inline Tensor ensure_gemm_compatible(const Tensor &t) { - if (t->ndim() == 2) { - return t->is_contiguous() ? t : rearrange(t); - } else if (t->ndim() == 3) { - return is_gemm_compatible_3d(t) ? t : rearrange(t); - } - return t->is_contiguous() ? t : rearrange(t); -} - -} // anonymous namespace - -Tensor bilinear(Tensor x1, Tensor x2, Tensor weight, std::optional bias) { - const size_t batch_size = x1->shape()[0]; - const size_t in1_features = x1->shape()[1]; - const size_t in2_features = x2->shape()[1]; - const size_t out_features = weight->shape()[0]; - - Tensor x1_compat = ensure_gemm_compatible(x1); - Tensor x2_compat = ensure_gemm_compatible(x2); - Tensor weight_cont = weight->is_contiguous() ? weight : weight->contiguous(); - - Tensor weight_permuted = weight_cont->permute({1, 0, 2}); - Tensor weight_permuted_cont = weight_permuted->is_contiguous() - ? weight_permuted - : weight_permuted->contiguous(); - Tensor weight_matrix = weight_permuted_cont->view({in1_features, out_features * in2_features}); - - Tensor intermediate = matmul(x1_compat, weight_matrix, 1.0f); - - Tensor intermediate_3d = intermediate->view({batch_size, out_features, in2_features}); - Tensor intermediate_transposed = intermediate_3d->permute({0, 2, 1}); - Tensor intermediate_compat = ensure_gemm_compatible(intermediate_transposed); - - Tensor x2_row = x2_compat->view({batch_size, 1, in2_features}); - Tensor x2_row_compat = ensure_gemm_compatible(x2_row); - - Tensor out_3d = matmul(x2_row_compat, intermediate_compat, 1.0f); - Tensor out = out_3d->view({batch_size, out_features}); - - if (bias) { - Tensor bias_broadcast = (*bias)->as_strided( - {batch_size, out_features}, - {0, (*bias)->strides()[0]}); - out = add(out, bias_broadcast); - } - return out; -} - -void bilinear_(Tensor out, Tensor x1, Tensor x2, Tensor weight, std::optional bias) { - Tensor result = bilinear(x1, x2, weight, bias); - rearrange_(out, result); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits.cc b/src/infinicore/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits.cc deleted file mode 100644 index d1517c60f..000000000 --- a/src/infinicore/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits.cc +++ /dev/null @@ -1,67 +0,0 @@ -#include "infinicore/ops/binary_cross_entropy_with_logits.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -// 静态调度器实例化 -common::OpDispatcher &BinaryCrossEntropyWithLogits::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -/** - * 执行核心逻辑：设备校验、上下文设置与后端分发 - */ -void BinaryCrossEntropyWithLogits::execute(Tensor out, Tensor logits, Tensor target, Tensor weight, Tensor pos_weight, std::string reduction) { - // 1. 校验所有已定义的 Tensor 是否在同一设备上 - // 使用宏或循环校验 logits, target, out 以及可选的 weight/pos_weight - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, logits, target); - if (weight) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, weight); - } - if (pos_weight) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, pos_weight); - } - - // 2. 设置当前设备上下文 - infinicore::context::setDevice(out->device()); - - // 3. 根据设备类型查找并执行具体的后端实现（如 CUDA 或 CPU 实现） - dispatcher().lookup(out->device().getType())(out, logits, target, weight, pos_weight, reduction); -} - -/** - * Out-of-place 接口：根据 reduction 自动创建输出 Tensor - */ -Tensor binary_cross_entropy_with_logits(Tensor logits, Tensor target, Tensor weight, Tensor pos_weight, std::string reduction) { - std::vector out_shape; - - // 1. 根据归约方式确定输出形状 - if (reduction == "none") { - // 不归约，形状与输入 logits 一致 - auto in_shape = logits->shape(); - for (auto dim : in_shape) { - out_shape.push_back(static_cast(dim)); - } - } else { - // mean 或 sum 归约，输出为标量 (空 shape 向量表示 0-dim tensor) - out_shape = {}; - } - - // 2. 创建输出 Tensor - auto out = Tensor::empty(out_shape, logits->dtype(), logits->device()); - - // 3. 调用显式接口执行计算 - binary_cross_entropy_with_logits_(out, logits, target, weight, pos_weight, reduction); - - return out; -} - -/** - * 显式指定输出接口 - */ -void binary_cross_entropy_with_logits_(Tensor out, Tensor logits, Tensor target, Tensor weight, Tensor pos_weight, std::string reduction) { - BinaryCrossEntropyWithLogits::execute(out, logits, target, weight, pos_weight, reduction); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits_infiniop.cc b/src/infinicore/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits_infiniop.cc deleted file mode 100644 index bfb68a924..000000000 --- a/src/infinicore/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits_infiniop.cc +++ /dev/null @@ -1,99 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/binary_cross_entropy_with_logits.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::bce_logits_impl::infiniop { - -// 定义线程局部的 BCEWithLogits 算子描述符缓存 -thread_local common::OpCache caches( - 100, - [](infiniopBCEWithLogitsDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyBCEWithLogitsDescriptor(desc)); - desc = nullptr; - } - }); - -/** - * @brief 执行 BCEWithLogits 计算 - * @param out 输出 Tensor (根据 reduction 可能是标量或与 logits 同形状) - * @param logits 预测值 Tensor - * @param target 标签 Tensor - * @param weight 样本权重 Tensor (可选) - * @param pos_weight 正类权重 Tensor (可选) - * @param reduction_str 归约方式 ("none", "mean", "sum") - */ -void calculate(Tensor out, Tensor logits, Tensor target, Tensor weight, Tensor pos_weight, std::string reduction_str) { - // 1. 将字符串归约参数转换为底层 API 使用的枚举值 - infiniopReduction_t reduction; - if (reduction_str == "none") { - reduction = INFINIOP_REDUCTION_NONE; - } else if (reduction_str == "mean") { - reduction = INFINIOP_REDUCTION_MEAN; - } else if (reduction_str == "sum") { - reduction = INFINIOP_REDUCTION_SUM; - } else { - throw std::runtime_error("Unknown reduction mode: " + reduction_str); - } - - // 2. 生成唯一 Hash Seed 用于缓存查找 - // 包含所有输入 Tensor 的状态和 reduction 参数，确保缓存键的唯一性 - size_t seed = hash_combine(out, logits, target, weight, pos_weight, static_cast(reduction)); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopBCEWithLogitsDescriptor_t desc = nullptr; - - // 3. 如果缓存未命中，创建新的描述符并存入缓存 - if (!desc_opt) { - // 获取可选 Tensor 的描述符，若未定义则传 nullptr - auto weight_desc = weight ? weight->desc() : nullptr; - auto pos_weight_desc = pos_weight ? pos_weight->desc() : nullptr; - - INFINICORE_CHECK_ERROR(infiniopCreateBCEWithLogitsDescriptor( - context::getInfiniopHandle(device), - &desc, - out->desc(), - logits->desc(), - target->desc(), - weight_desc, - pos_weight_desc, - reduction)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 4. 动态获取并分配 Workspace 临时内存 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetBCEWithLogitsWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - // 5. 获取数据指针，处理可选 Tensor 的空指针逻辑 - const void *weight_ptr = weight ? weight->data() : nullptr; - const void *pos_weight_ptr = pos_weight ? pos_weight->data() : nullptr; - - // 6. 执行底层算子 - INFINICORE_CHECK_ERROR(infiniopBCEWithLogits( - desc, - workspace->data(), - workspace_size, - out->data(), - logits->data(), - target->data(), - weight_ptr, - pos_weight_ptr, - context::getStream())); -} - -// 7. 自动注册到调度器 (Dispatcher) -static bool registered = []() { - BinaryCrossEntropyWithLogits::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::bce_logits_impl::infiniop diff --git a/src/infinicore/ops/bitwise_right_shift/bitwise_right_shift.cc b/src/infinicore/ops/bitwise_right_shift/bitwise_right_shift.cc deleted file mode 100644 index caaad3193..000000000 --- a/src/infinicore/ops/bitwise_right_shift/bitwise_right_shift.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/bitwise_right_shift.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BitwiseRightShift); - -BitwiseRightShift::BitwiseRightShift(Tensor out, const Tensor &input, const Tensor &other) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, other); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, input, other); -} - -void BitwiseRightShift::execute(Tensor out, const Tensor &input, const Tensor &other) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(BitwiseRightShift, out, input, other); -} - -Tensor bitwise_right_shift(const Tensor &input, const Tensor &other) { - auto out = Tensor::empty(input->shape(), input->dtype(), input->device()); - bitwise_right_shift_(out, input, other); - return out; -} - -void bitwise_right_shift_(Tensor out, const Tensor &input, const Tensor &other) { - BitwiseRightShift::execute(out, input, other); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/bitwise_right_shift/bitwise_right_shift_infiniop.cc b/src/infinicore/ops/bitwise_right_shift/bitwise_right_shift_infiniop.cc deleted file mode 100644 index 3d3aa1e43..000000000 --- a/src/infinicore/ops/bitwise_right_shift/bitwise_right_shift_infiniop.cc +++ /dev/null @@ -1,53 +0,0 @@ -#include "infinicore/ops/bitwise_right_shift.hpp" - -#include "infiniop/ops/bitwise_right_shift.h" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::bitwise_right_shift_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, BitwiseRightShift, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, input, other; -}; - -void *plan(Tensor out, const Tensor &input, const Tensor &other) { - size_t seed = hash_combine(out, input, other); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, BitwiseRightShift, - seed, out->desc(), input->desc(), other->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, BitwiseRightShift, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(input), - graph::GraphTensor(other)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopBitwiseRightShift( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->input->data(), - planned->other->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(BitwiseRightShift, &plan, &run, &cleanup); - -} // namespace infinicore::op::bitwise_right_shift_impl::infiniop diff --git a/src/infinicore/ops/blas_amax/blas_amax.cc b/src/infinicore/ops/blas_amax/blas_amax.cc deleted file mode 100644 index 9579589ff..000000000 --- a/src/infinicore/ops/blas_amax/blas_amax.cc +++ /dev/null @@ -1,28 +0,0 @@ -#include "infinicore/ops/blas_amax.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BlasAmax); - -BlasAmax::BlasAmax(const Tensor &x, Tensor result) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, result); - INFINICORE_GRAPH_OP_DISPATCH(result->device().getType(), x, result); -} - -void BlasAmax::execute(const Tensor &x, Tensor result) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(BlasAmax, x, result); -} - -Tensor blas_amax(const Tensor &x) { - auto result = Tensor::empty({}, DataType::I32, x->device()); - blas_amax_(x, result); - return result; -} - -void blas_amax_(const Tensor &x, Tensor result) { - BlasAmax::execute(x, result); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/blas_amax/blas_amax_infiniop.cc b/src/infinicore/ops/blas_amax/blas_amax_infiniop.cc deleted file mode 100644 index 780fca744..000000000 --- a/src/infinicore/ops/blas_amax/blas_amax_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/blas_amax.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::blas_amax_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, BlasAmax, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, x, result; -}; - -void *plan(const Tensor &x, Tensor result) { - size_t seed = hash_combine(x, result); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, BlasAmax, - seed, - x->desc(), result->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, BlasAmax, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x), - graph::GraphTensor(result)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopBlasAmax( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->x->data(), - planned->result->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(BlasAmax, &plan, &run, &cleanup); - -} // namespace infinicore::op::blas_amax_impl::infiniop diff --git a/src/infinicore/ops/blas_amin/blas_amin.cc b/src/infinicore/ops/blas_amin/blas_amin.cc deleted file mode 100644 index e18e22739..000000000 --- a/src/infinicore/ops/blas_amin/blas_amin.cc +++ /dev/null @@ -1,28 +0,0 @@ -#include "infinicore/ops/blas_amin.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BlasAmin); - -BlasAmin::BlasAmin(const Tensor &x, Tensor result) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, result); - INFINICORE_GRAPH_OP_DISPATCH(result->device().getType(), x, result); -} - -void BlasAmin::execute(const Tensor &x, Tensor result) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(BlasAmin, x, result); -} - -Tensor blas_amin(const Tensor &x) { - auto result = Tensor::empty({}, DataType::I32, x->device()); - blas_amin_(x, result); - return result; -} - -void blas_amin_(const Tensor &x, Tensor result) { - BlasAmin::execute(x, result); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/blas_amin/blas_amin_infiniop.cc b/src/infinicore/ops/blas_amin/blas_amin_infiniop.cc deleted file mode 100644 index 00abf77e1..000000000 --- a/src/infinicore/ops/blas_amin/blas_amin_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/blas_amin.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::blas_amin_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, BlasAmin, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, x, result; -}; - -void *plan(const Tensor &x, Tensor result) { - size_t seed = hash_combine(x, result); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, BlasAmin, - seed, - x->desc(), result->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, BlasAmin, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x), - graph::GraphTensor(result)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopBlasAmin( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->x->data(), - planned->result->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(BlasAmin, &plan, &run, &cleanup); - -} // namespace infinicore::op::blas_amin_impl::infiniop diff --git a/src/infinicore/ops/blas_copy/blas_copy.cc b/src/infinicore/ops/blas_copy/blas_copy.cc deleted file mode 100644 index 77b2e7a5f..000000000 --- a/src/infinicore/ops/blas_copy/blas_copy.cc +++ /dev/null @@ -1,22 +0,0 @@ -#include "infinicore/ops/blas_copy.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BlasCopy); - -BlasCopy::BlasCopy(const Tensor &x, Tensor y) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, y); - INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), x, y); -} - -void BlasCopy::execute(const Tensor &x, Tensor y) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(BlasCopy, x, y); -} - -void blas_copy_(const Tensor &x, Tensor y) { - BlasCopy::execute(x, y); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/blas_copy/blas_copy_infiniop.cc b/src/infinicore/ops/blas_copy/blas_copy_infiniop.cc deleted file mode 100644 index 33a70523b..000000000 --- a/src/infinicore/ops/blas_copy/blas_copy_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/blas_copy.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::blas_copy_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, BlasCopy, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, x, y; -}; - -void *plan(const Tensor &x, Tensor y) { - size_t seed = hash_combine(x, y); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, BlasCopy, - seed, - x->desc(), y->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, BlasCopy, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x), - graph::GraphTensor(y)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopBlasCopy( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->x->data(), - planned->y->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(BlasCopy, &plan, &run, &cleanup); - -} // namespace infinicore::op::blas_copy_impl::infiniop diff --git a/src/infinicore/ops/blas_dot/blas_dot.cc b/src/infinicore/ops/blas_dot/blas_dot.cc deleted file mode 100644 index 8dba37acc..000000000 --- a/src/infinicore/ops/blas_dot/blas_dot.cc +++ /dev/null @@ -1,28 +0,0 @@ -#include "infinicore/ops/blas_dot.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BlasDot); - -BlasDot::BlasDot(const Tensor &x, const Tensor &y, Tensor result) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, y, result); - INFINICORE_GRAPH_OP_DISPATCH(result->device().getType(), x, y, result); -} - -void BlasDot::execute(const Tensor &x, const Tensor &y, Tensor result) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(BlasDot, x, y, result); -} - -Tensor blas_dot(const Tensor &x, const Tensor &y) { - auto result = Tensor::empty({}, x->dtype(), x->device()); - blas_dot_(x, y, result); - return result; -} - -void blas_dot_(const Tensor &x, const Tensor &y, Tensor result) { - BlasDot::execute(x, y, result); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/blas_dot/blas_dot_infiniop.cc b/src/infinicore/ops/blas_dot/blas_dot_infiniop.cc deleted file mode 100644 index 78559eea7..000000000 --- a/src/infinicore/ops/blas_dot/blas_dot_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "infinicore/ops/blas_dot.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::blas_dot_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, BlasDot, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, x, y, result; -}; - -void *plan(const Tensor &x, const Tensor &y, Tensor result) { - size_t seed = hash_combine(x, y, result); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, BlasDot, - seed, - x->desc(), y->desc(), result->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, BlasDot, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x), - graph::GraphTensor(y), - graph::GraphTensor(result)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopBlasDot( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->x->data(), - planned->y->data(), - planned->result->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(BlasDot, &plan, &run, &cleanup); - -} // namespace infinicore::op::blas_dot_impl::infiniop diff --git a/src/infinicore/ops/block_diag/block_diag.cc b/src/infinicore/ops/block_diag/block_diag.cc deleted file mode 100644 index 4ab1fa24a..000000000 --- a/src/infinicore/ops/block_diag/block_diag.cc +++ /dev/null @@ -1,63 +0,0 @@ -#include "infinicore/ops/block_diag.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &BlockDiag::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -void BlockDiag::execute(Tensor output, const std::vector &inputs) { - if (inputs.empty()) { - throw std::runtime_error("block_diag expects at least one input tensor"); - } - - for (const auto &x : inputs) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, x); - } - - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error( - "No BlockDiag implementation found for device type: " - + std::to_string(static_cast(device_type))); - } - - func(output, inputs); -} - -Tensor block_diag(const std::vector &inputs) { - if (inputs.empty()) { - throw std::runtime_error("block_diag expects at least one input tensor"); - } - - const auto &device = inputs.front()->device(); - const auto dtype = inputs.front()->dtype(); - - Size total_rows = 0; - Size total_cols = 0; - for (const auto &x : inputs) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(inputs.front(), x); - INFINICORE_ASSERT(x->dtype() == dtype); - INFINICORE_ASSERT(x->ndim() == 2); - total_rows += x->size(0); - total_cols += x->size(1); - } - - auto output = Tensor::empty({total_rows, total_cols}, dtype, device); - block_diag_(output, inputs); - return output; -} - -void block_diag_(Tensor output, const std::vector &inputs) { - BlockDiag::execute(output, inputs); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/block_diag/block_diag_infiniop.cc b/src/infinicore/ops/block_diag/block_diag_infiniop.cc deleted file mode 100644 index b2b8aea9c..000000000 --- a/src/infinicore/ops/block_diag/block_diag_infiniop.cc +++ /dev/null @@ -1,78 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/block_diag.hpp" -#include "infinicore/ops/common/cache.hpp" - -#include - -namespace infinicore::op::block_diag_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopBlockDiagDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyBlockDiagDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, const std::vector &inputs) { - if (inputs.empty()) { - throw std::runtime_error("block_diag expects at least one input tensor"); - } - - size_t seed = 0; - hash_combine(seed, output, static_cast(inputs.size())); - for (const auto &x : inputs) { - hash_combine(seed, x); - } - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopBlockDiagDescriptor_t desc = nullptr; - - std::vector input_descs; - input_descs.reserve(inputs.size()); - for (const auto &x : inputs) { - input_descs.push_back(x->desc()); - } - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateBlockDiagDescriptor( - context::getInfiniopHandle(device), - &desc, - output->desc(), - input_descs.data(), - input_descs.size())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetBlockDiagWorkspaceSize(desc, &workspace_size)); - auto workspace = context::allocateMemory(workspace_size); - - std::vector input_ptrs; - input_ptrs.reserve(inputs.size()); - for (const auto &x : inputs) { - input_ptrs.push_back(x->data()); - } - - INFINICORE_CHECK_ERROR(infiniopBlockDiag( - desc, - workspace->data(), - workspace_size, - output->data(), - input_ptrs.data(), - context::getStream())); -} - -static bool registered = []() { - BlockDiag::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::block_diag_impl::infiniop diff --git a/src/infinicore/ops/broadcast_to/broadcast_to.cc b/src/infinicore/ops/broadcast_to/broadcast_to.cc deleted file mode 100644 index b6f3079f5..000000000 --- a/src/infinicore/ops/broadcast_to/broadcast_to.cc +++ /dev/null @@ -1,28 +0,0 @@ -#include "infinicore/ops/broadcast_to.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(BroadcastTo); - -BroadcastTo::BroadcastTo(Tensor y, Tensor x) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, x); -} - -void BroadcastTo::execute(Tensor y, Tensor x) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(BroadcastTo, y, x); -} - -Tensor broadcast_to(Tensor x, const std::vector &shape) { - Shape target_shape(shape.begin(), shape.end()); - auto y = Tensor::empty(target_shape, x->dtype(), x->device()); - broadcast_to_(y, x); - return y; -} - -void broadcast_to_(Tensor y, Tensor x) { - BroadcastTo::execute(y, x); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/broadcast_to/broadcast_to_infiniop.cc b/src/infinicore/ops/broadcast_to/broadcast_to_infiniop.cc deleted file mode 100644 index 57c1ee901..000000000 --- a/src/infinicore/ops/broadcast_to/broadcast_to_infiniop.cc +++ /dev/null @@ -1,49 +0,0 @@ -#include "../infiniop_impl.hpp" -#include "infinicore/ops/broadcast_to.hpp" - -namespace infinicore::op::broadcast_to_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, BroadcastTo, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, y, x; -}; - -void *plan(Tensor y, Tensor x) { - size_t seed = hash_combine(y, x); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, BroadcastTo, - seed, - y->desc(), x->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, BroadcastTo, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(y), - graph::GraphTensor(x)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopBroadcastTo( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->y->data(), - planned->x->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(BroadcastTo, &plan, &run, &cleanup); - -} // namespace infinicore::op::broadcast_to_impl::infiniop diff --git a/src/infinicore/ops/cat/cat.cc b/src/infinicore/ops/cat/cat.cc deleted file mode 100644 index 3eafeae9e..000000000 --- a/src/infinicore/ops/cat/cat.cc +++ /dev/null @@ -1,200 +0,0 @@ -#include "infinicore/ops/cat.hpp" -#include "infinicore/context/context.hpp" -#include -#include -namespace infinicore::op { - -namespace { - -bool use_slice_copy_cat(Device::Type device_type, int dim, int ndim) { - // Keep this path limited to CUDA-like backends we have validated or can - // trace to the nvidia rearrange implementation. Other devices may be - // correct through copy_from, but their performance impact is unverified. - return dim == ndim - 1 - && (device_type == Device::Type::NVIDIA - || device_type == Device::Type::HYGON - || device_type == Device::Type::ILUVATAR - || device_type == Device::Type::ALI); -} - -class CatInfo { - - CatInfo() = default; - -public: - int dim; - int ndim; - size_t tensors_size; - - std::vector contiguous_dim; - std::vector copy_size; - - static CatInfo create(Tensor &out, std::vector &tensors, int dim) { - - int ndim = out->ndim(); - size_t tensors_size = tensors.size(); - - std::vector contiguous_dim(tensors_size, ndim); - std::vector copy_size(tensors_size, dsize(out->dtype())); - - for (int i = 0; i < tensors_size; i++) { - if (tensors[i]->ndim() == 1) { - continue; - } - if (tensors[i]->stride(ndim - 1) == 1) { - contiguous_dim[i] = ndim - 1; - for (int j = ndim - 2; j >= dim; j--) { - if (tensors[i]->stride(j) == tensors[i]->stride(j + 1) * tensors[i]->shape()[j + 1]) { - contiguous_dim[i] = j; - } - } - } - for (int j = contiguous_dim[i]; j < ndim; j++) { - copy_size[i] *= tensors[i]->shape()[j]; - } - } - - return CatInfo{dim, ndim, tensors_size, contiguous_dim, copy_size}; - } -}; - -void low_dim_copy( - CatInfo &info, Tensor &tensor, Tensor &out, - std::byte *tensor_ptr, std::byte *out_ptr, - int depth, int tensor_pos) { - - if (depth != info.contiguous_dim[tensor_pos]) { - std::byte *now_tensor_ptr = tensor_ptr; - std::byte *now_out_ptr = out_ptr; - - for (int i = 0; i < tensor->shape()[depth]; i++) { - - low_dim_copy(info, tensor, out, now_tensor_ptr, now_out_ptr, depth + 1, tensor_pos); - - now_tensor_ptr += tensor->stride(depth) * dsize(tensor->dtype()); - now_out_ptr += out->stride(depth) * dsize(out->dtype()); - } - } else { - if (out->device().getType() == Device::Type::CPU) { - - std::memcpy(out_ptr, tensor_ptr, info.copy_size[tensor_pos]); - } else { - - context::memcpyD2D(out_ptr, tensor_ptr, info.copy_size[tensor_pos]); - } - } -} - -void high_dim_split( - CatInfo &info, std::vector &tensors, Tensor &out, - std::vector tensors_ptr, std::byte *out_ptr, - int depth) { - - if (depth != info.dim) { - std::vector now_tensors_ptr = tensors_ptr; - std::byte *now_out_ptr = out_ptr; - - for (int i = 0; i < out->shape()[depth]; i++) { - - high_dim_split(info, tensors, out, now_tensors_ptr, now_out_ptr, depth + 1); - - for (int i = 0; i < info.tensors_size; i++) { - if (tensors[i]->ndim() == 1) { - continue; - } - now_tensors_ptr[i] += tensors[i]->stride(depth) * dsize(tensors[i]->dtype()); - } - now_out_ptr += out->stride(depth) * dsize(out->dtype()); - } - } else { - std::byte *now_out_ptr = out_ptr; - - for (int i = 0; i < info.tensors_size; i++) { - if (tensors[i]->ndim() == 1) { - continue; - } - - low_dim_copy(info, tensors[i], out, tensors_ptr[i], now_out_ptr, depth, i); - - now_out_ptr += tensors[i]->shape()[depth] * out->stride(depth) * dsize(out->dtype()); - } - } -} - -} // namespace - -Tensor cat(std::vector tensors, int dim) { - assert(tensors.size() >= 2); - int ndim = tensors[0]->ndim(); - assert(-ndim <= dim && dim < ndim); - dim = (dim + ndim) % ndim; - - Shape shape = tensors[0]->shape(); - for (int i = 1; i < tensors.size(); i++) { - assert(tensors[i]->ndim() == dim || tensors[i]->ndim() == 1); - if (tensors[i]->ndim() != ndim) { - continue; - } - shape[dim] += tensors[i]->shape()[dim]; - } - - auto out = Tensor::empty(shape, tensors[0]->dtype(), tensors[0]->device()); - cat_(out, tensors, dim); - return out; -} - -void cat_(Tensor out, std::vector tensors, int dim) { - // assert the parameter properties are correct. - assert(tensors.size() >= 2); - int ndim = out->ndim(); - assert(-ndim <= dim && dim < ndim); - dim = (dim + ndim) % ndim; - - size_t dim_shape = 0; - for (auto &tensor : tensors) { - assert(tensor->ndim() == ndim || tensors[i]->ndim() == 1); - if (tensor->ndim() == 1) { - assert(tensor->shape()[0] == 0); - continue; - } - for (int i = 0; i < ndim; i++) { - if (i != dim) { - assert(tensor->shape()[i] == out->shape()[i]); - } else { - dim_shape += tensor->shape()[i]; - } - } - } - assert(dim_shape == out->shape()[dim]); - - if (use_slice_copy_cat(out->device().getType(), dim, ndim)) { - // The generic path recursively issues one memcpy for every outer - // index. Concatenating MLA tensors on the last dimension can therefore - // enqueue hundreds of tiny D2D copy calls per layer. A strided output - // slice is semantically identical and copy_from lowers it to a single - // rearrange kernel per input tensor on CUDA-like backends. - size_t offset = 0; - for (auto &tensor : tensors) { - if (tensor->ndim() == 1) { - continue; - } - const size_t length = tensor->shape()[dim]; - auto output_slice = out->narrow({{static_cast(dim), offset, length}}); - output_slice->copy_from(tensor); - offset += length; - } - return; - } - - // Get info - CatInfo info = CatInfo::create(out, tensors, dim); - std::vector tensors_ptr(tensors.size()); - for (int i = 0; i < tensors.size(); i++) { - tensors_ptr[i] = tensors[i]->data(); - } - std::byte *out_ptr = out->data(); - - high_dim_split(info, tensors, out, tensors_ptr, out_ptr, 0); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/causal_conv1d/causal_conv1d.cc b/src/infinicore/ops/causal_conv1d/causal_conv1d.cc deleted file mode 100644 index 5645805e7..000000000 --- a/src/infinicore/ops/causal_conv1d/causal_conv1d.cc +++ /dev/null @@ -1,146 +0,0 @@ -#include "infinicore/ops/causal_conv1d.hpp" -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(CausalConv1d); - -CausalConv1d::CausalConv1d(Tensor out, - Tensor conv_state, - std::optional final_conv_state, - const Tensor &qkv, - const Tensor &weight, - std::optional bias, - std::optional cu_seqlens, - std::optional initial_state_indices, - std::optional final_state_indices) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, conv_state, qkv, weight); - if (final_conv_state.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_conv_state.value()); - } - if (bias.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, bias.value()); - } - if (cu_seqlens.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, cu_seqlens.value()); - } - if (initial_state_indices.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state_indices.value()); - } - if (final_state_indices.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_state_indices.value()); - } - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), - out, - conv_state, - final_conv_state, - qkv, - weight, - bias, - cu_seqlens, - initial_state_indices, - final_state_indices); -} - -void CausalConv1d::execute(Tensor out, - Tensor conv_state, - std::optional final_conv_state, - const Tensor &qkv, - const Tensor &weight, - std::optional bias, - std::optional cu_seqlens, - std::optional initial_state_indices, - std::optional final_state_indices) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(CausalConv1d, - out, - conv_state, - final_conv_state, - qkv, - weight, - bias, - cu_seqlens, - initial_state_indices, - final_state_indices); -} - -static void check_3d_tensor(const Tensor &x, const char *name, const char *shape_hint) { - if (x->shape().size() != 3) { - throw std::runtime_error(std::string("causal_conv1d expects ") + name + " with shape " + shape_hint); - } -} - -static size_t request_count_from_args(const Tensor &qkv, std::optional cu_seqlens) { - if (cu_seqlens.has_value()) { - const auto &cu_shape = cu_seqlens.value()->shape(); - if (cu_shape.size() != 1 || cu_shape[0] < 2) { - throw std::runtime_error("causal_conv1d expects cu_seqlens with shape [num_requests + 1]"); - } - return cu_shape[0] - 1; - } - return qkv->shape()[0]; -} - -static Shape final_state_shape(const Tensor &qkv, - const Tensor &conv_state, - const Tensor &weight, - std::optional cu_seqlens) { - const auto request_count = request_count_from_args(qkv, cu_seqlens); - return {request_count, conv_state->shape()[1], weight->shape()[2] - 1}; -} - -Tensor causal_conv1d(const Tensor &qkv, - Tensor conv_state, - const Tensor &weight, - std::optional bias, - std::optional cu_seqlens, - std::optional initial_state_indices, - std::optional final_state_indices) { - check_3d_tensor(qkv, "qkv", "[B, T, C] or [1, total_tokens, C]"); - check_3d_tensor(conv_state, "conv_state", "[B/num_requests, C, state_len] or [pool_size, C, state_len]"); - check_3d_tensor(weight, "weight", "[C, 1, state_len + 1]"); - auto out = Tensor::empty(qkv->shape(), qkv->dtype(), qkv->device()); - std::optional final_conv_state = std::nullopt; - if (!final_state_indices.has_value()) { - final_conv_state = Tensor::empty(final_state_shape(qkv, conv_state, weight, cu_seqlens), - conv_state->dtype(), - conv_state->device()); - } - causal_conv1d_(out, - conv_state, - final_conv_state, - qkv, - weight, - bias, - cu_seqlens, - initial_state_indices, - final_state_indices); - return out; -} - -void causal_conv1d_(Tensor out, - Tensor conv_state, - std::optional final_conv_state, - const Tensor &qkv, - const Tensor &weight, - std::optional bias, - std::optional cu_seqlens, - std::optional initial_state_indices, - std::optional final_state_indices) { - check_3d_tensor(out, "out", "[B, T, C] or [1, total_tokens, C]"); - check_3d_tensor(qkv, "qkv", "[B, T, C] or [1, total_tokens, C]"); - check_3d_tensor(conv_state, "conv_state", "[B/num_requests, C, state_len] or [pool_size, C, state_len]"); - check_3d_tensor(weight, "weight", "[C, 1, state_len + 1]"); - CausalConv1d::execute(out, - conv_state, - final_conv_state, - qkv, - weight, - bias, - cu_seqlens, - initial_state_indices, - final_state_indices); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/causal_conv1d/causal_conv1d_infiniop.cc b/src/infinicore/ops/causal_conv1d/causal_conv1d_infiniop.cc deleted file mode 100644 index 680c568b3..000000000 --- a/src/infinicore/ops/causal_conv1d/causal_conv1d_infiniop.cc +++ /dev/null @@ -1,93 +0,0 @@ -#include "infinicore/ops/causal_conv1d.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::causal_conv1d_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, CausalConv1d, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, conv_state, qkv, weight; - std::optional final_conv_state; - std::optional bias; - std::optional cu_seqlens; - std::optional initial_state_indices; - std::optional final_state_indices; -}; - -void *plan(Tensor out, - Tensor conv_state, - std::optional final_conv_state, - const Tensor &qkv, - const Tensor &weight, - std::optional bias, - std::optional cu_seqlens, - std::optional initial_state_indices, - std::optional final_state_indices) { - size_t seed = hash_combine(out, - conv_state, - final_conv_state, - qkv, - weight, - bias, - cu_seqlens, - initial_state_indices, - final_state_indices); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, CausalConv1d, - seed, - out->desc(), - conv_state->desc(), - final_conv_state.has_value() ? final_conv_state.value()->desc() : nullptr, - qkv->desc(), - weight->desc(), - bias.has_value() ? bias.value()->desc() : nullptr, - cu_seqlens.has_value() ? cu_seqlens.value()->desc() : nullptr, - initial_state_indices.has_value() ? initial_state_indices.value()->desc() : nullptr, - final_state_indices.has_value() ? final_state_indices.value()->desc() : nullptr); - - INFINIOP_WORKSPACE_TENSOR(workspace, CausalConv1d, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(conv_state), - graph::GraphTensor(qkv), - graph::GraphTensor(weight), - final_conv_state.has_value() ? std::optional(graph::GraphTensor(final_conv_state.value())) : std::nullopt, - bias.has_value() ? std::optional(graph::GraphTensor(bias.value())) : std::nullopt, - cu_seqlens.has_value() ? std::optional(graph::GraphTensor(cu_seqlens.value())) : std::nullopt, - initial_state_indices.has_value() ? std::optional(graph::GraphTensor(initial_state_indices.value())) : std::nullopt, - final_state_indices.has_value() ? std::optional(graph::GraphTensor(final_state_indices.value())) : std::nullopt}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopCausalConv1d( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->conv_state->data(), - planned->final_conv_state.has_value() ? planned->final_conv_state.value()->data() : nullptr, - planned->qkv->data(), - planned->weight->data(), - planned->bias.has_value() ? planned->bias.value()->data() : nullptr, - planned->cu_seqlens.has_value() ? planned->cu_seqlens.value()->data() : nullptr, - planned->initial_state_indices.has_value() ? planned->initial_state_indices.value()->data() : nullptr, - planned->final_state_indices.has_value() ? planned->final_state_indices.value()->data() : nullptr, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(CausalConv1d, &plan, &run, &cleanup); - -} // namespace infinicore::op::causal_conv1d_impl::infiniop diff --git a/src/infinicore/ops/causal_softmax/causal_softmax.cc b/src/infinicore/ops/causal_softmax/causal_softmax.cc deleted file mode 100644 index 328ff390e..000000000 --- a/src/infinicore/ops/causal_softmax/causal_softmax.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/causal_softmax.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(CausalSoftmax); - -CausalSoftmax::CausalSoftmax(Tensor output, const Tensor &input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - INFINICORE_GRAPH_OP_DISPATCH(output->device().getType(), output, input); -} - -void CausalSoftmax::execute(Tensor output, const Tensor &input) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(CausalSoftmax, output, input); -} - -Tensor causal_softmax(const Tensor &input) { - auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); - causal_softmax_(output, input); - return output; -} - -void causal_softmax_(Tensor output, const Tensor &input) { - CausalSoftmax::execute(output, input); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/causal_softmax/causal_softmax_infiniop.cc b/src/infinicore/ops/causal_softmax/causal_softmax_infiniop.cc deleted file mode 100644 index e0a0595fb..000000000 --- a/src/infinicore/ops/causal_softmax/causal_softmax_infiniop.cc +++ /dev/null @@ -1,49 +0,0 @@ -#include "infinicore/ops/causal_softmax.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::causal_softmax_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, CausalSoftmax, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, output, input; -}; - -void *plan(Tensor output, const Tensor &input) { - size_t seed = hash_combine(output, input); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, CausalSoftmax, - seed, output->desc(), input->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, CausalSoftmax, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(output), - graph::GraphTensor(input)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopCausalSoftmax( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->output->data(), - planned->input->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(CausalSoftmax, &plan, &run, &cleanup); - -} // namespace infinicore::op::causal_softmax_impl::infiniop diff --git a/src/infinicore/ops/causal_softmax/causal_softmax_infiniops.cc b/src/infinicore/ops/causal_softmax/causal_softmax_infiniops.cc deleted file mode 100644 index de35131d7..000000000 --- a/src/infinicore/ops/causal_softmax/causal_softmax_infiniops.cc +++ /dev/null @@ -1,58 +0,0 @@ -#include "infinicore/ops/causal_softmax.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/causal_softmax.h" - -namespace infinicore::op::causal_softmax_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -struct PlannedMeta { - TensorMeta output, input; - graph::GraphTensor output_tensor, input_tensor; -}; - -} // namespace - -void *plan(Tensor output, const Tensor &input) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - - return new PlannedMeta{ - TensorMeta(output), - TensorMeta(input), - graph::GraphTensor(output), - graph::GraphTensor(input)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - infini::ops::CausalSoftmax::Call( - handle, - config, - planned->input.tensor(planned->input_tensor), - planned->output.tensor(planned->output_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(CausalSoftmax::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(CausalSoftmax::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(CausalSoftmax::cleanup_dispatcher(), &cleanup); - return true; -}(); - -} // namespace infinicore::op::causal_softmax_impl::infiniops -#endif diff --git a/src/infinicore/ops/cdist/cdist.cc b/src/infinicore/ops/cdist/cdist.cc deleted file mode 100644 index f989488c9..000000000 --- a/src/infinicore/ops/cdist/cdist.cc +++ /dev/null @@ -1,56 +0,0 @@ -#include "infinicore/ops/cdist.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -// 静态调度器实例化 -common::OpDispatcher &Cdist::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -/** - * 执行核心逻辑：设备校验与后端分发 - */ -void Cdist::execute(Tensor out, Tensor x1, Tensor x2, double p) { - // 校验三个 Tensor 是否在同一设备上 - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, x1, x2); - - // 设置当前设备上下文 - infinicore::context::setDevice(out->device()); - - // 根据设备类型（CUDA/CPU/etc.）查找并执行注册的算子实现 - dispatcher().lookup(out->device().getType())(out, x1, x2, p); -} - -/** - * Out-of-place 接口：自动创建输出 Tensor - * x1: (M, D), x2: (N, D) -> out: (M, N) - */ -Tensor cdist(Tensor x1, Tensor x2, double p) { - // 1. 获取输入维度 - auto shape1 = x1->shape(); // 假设为 {M, D} - auto shape2 = x2->shape(); // 假设为 {N, D} - - // 将原来的 std::vector 修改为 std::vector - std::vector out_shape = { - static_cast(shape1[0]), - static_cast(shape2[0])}; - - // 或者使用更简洁的初始化列表方式，强制转换类型 - auto out = Tensor::empty({(uint64_t)shape1[0], (uint64_t)shape2[0]}, x1->dtype(), x1->device()); - - // 5. 调用执行接口 - cdist_(out, x1, x2, p); - - return out; -} - -/** - * 显式指定输出接口 - */ -void cdist_(Tensor out, Tensor x1, Tensor x2, double p) { - Cdist::execute(out, x1, x2, p); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/cdist/cdist_infiniop.cc b/src/infinicore/ops/cdist/cdist_infiniop.cc deleted file mode 100644 index 7eb6b0b4b..000000000 --- a/src/infinicore/ops/cdist/cdist_infiniop.cc +++ /dev/null @@ -1,66 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/cdist.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -namespace infinicore::op::cdist_impl::infiniop { - -// 定义线程局部的 cdist 算子描述符缓存 -// 缓存 key 为输入 Tensor 描述信息及参数 p 的哈希值 -thread_local common::OpCache caches( - 100, - [](infiniopCdistDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyCdistDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor out, Tensor x1, Tensor x2, double p) { - // 1. 生成唯一 Hash Seed 用于缓存查找 - size_t seed = hash_combine(out, x1, x2, p); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopCdistDescriptor_t desc = nullptr; - - // 2. 如果缓存未命中，创建新的描述符并存入缓存 - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateCdistDescriptor( - context::getInfiniopHandle(device), - &desc, - out->desc(), - x1->desc(), - x2->desc(), - p)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 3. 动态获取并分配 Workspace 临时内存 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetCdistWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - // 4. 执行底层算子 - INFINICORE_CHECK_ERROR(infiniopCdist( - desc, - workspace->data(), - workspace_size, - out->data(), - x1->data(), - x2->data(), - context::getStream())); -} - -// 5. 自动注册到调度器 (Dispatcher) -static bool registered = []() { - Cdist::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::cdist_impl::infiniop diff --git a/src/infinicore/ops/chunk_gated_delta_rule/chunk_gated_delta_rule.cc b/src/infinicore/ops/chunk_gated_delta_rule/chunk_gated_delta_rule.cc deleted file mode 100644 index 82e3d5b1d..000000000 --- a/src/infinicore/ops/chunk_gated_delta_rule/chunk_gated_delta_rule.cc +++ /dev/null @@ -1,167 +0,0 @@ -#include "infinicore/ops/chunk_gated_delta_rule.hpp" -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(ChunkGatedDeltaRule); - -ChunkGatedDeltaRule::ChunkGatedDeltaRule(Tensor out, - Tensor initial_state, - std::optional final_state, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - std::optional cu_seqlens, - std::optional initial_state_indices, - std::optional final_state_indices, - bool use_qk_l2norm, - size_t chunk_size) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state, q, k, v, g, beta); - if (final_state.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_state.value()); - } - if (cu_seqlens.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, cu_seqlens.value()); - } - if (initial_state_indices.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state_indices.value()); - } - if (final_state_indices.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_state_indices.value()); - } - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), - out, - initial_state, - final_state, - q, - k, - v, - g, - beta, - cu_seqlens, - initial_state_indices, - final_state_indices, - use_qk_l2norm, - chunk_size); -} - -void ChunkGatedDeltaRule::execute(Tensor out, - Tensor initial_state, - std::optional final_state, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - std::optional cu_seqlens, - std::optional initial_state_indices, - std::optional final_state_indices, - bool use_qk_l2norm, - size_t chunk_size) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(ChunkGatedDeltaRule, - out, - initial_state, - final_state, - q, - k, - v, - g, - beta, - cu_seqlens, - initial_state_indices, - final_state_indices, - use_qk_l2norm, - chunk_size); -} - -static void check_4d_sequence_tensor(const Tensor &x, const char *name) { - if (x->shape().size() != 4) { - throw std::runtime_error(std::string("chunk_gated_delta_rule expects ") + name + " with shape [B, T, H, D] or [1, total_tokens, H, D]"); - } -} - -static Shape chunk_final_state_shape(const Tensor &q, - const Tensor &v, - std::optional cu_seqlens) { - const auto &q_shape = q->shape(); - const auto &v_shape = v->shape(); - size_t B = cu_seqlens.has_value() ? cu_seqlens.value()->shape()[0] - 1 : v_shape[0]; - size_t Hv = v_shape[2]; - size_t Dk = q_shape[3]; - size_t Dv = v_shape[3]; - return {B, Hv, Dv, Dk}; -} - -Tensor chunk_gated_delta_rule(const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - Tensor initial_state, - std::optional cu_seqlens, - std::optional initial_state_indices, - std::optional final_state_indices, - bool use_qk_l2norm, - size_t chunk_size) { - check_4d_sequence_tensor(q, "q"); - check_4d_sequence_tensor(k, "k"); - check_4d_sequence_tensor(v, "v"); - auto out = Tensor::empty(v->shape(), v->dtype(), v->device()); - std::optional final_state = std::nullopt; - if (!final_state_indices.has_value()) { - final_state = Tensor::empty(chunk_final_state_shape(q, v, cu_seqlens), - initial_state->dtype(), - initial_state->device()); - } - chunk_gated_delta_rule_(out, - initial_state, - final_state, - q, - k, - v, - g, - beta, - cu_seqlens, - initial_state_indices, - final_state_indices, - use_qk_l2norm, - chunk_size); - return out; -} - -void chunk_gated_delta_rule_(Tensor out, - Tensor initial_state, - std::optional final_state, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - std::optional cu_seqlens, - std::optional initial_state_indices, - std::optional final_state_indices, - bool use_qk_l2norm, - size_t chunk_size) { - check_4d_sequence_tensor(q, "q"); - check_4d_sequence_tensor(k, "k"); - check_4d_sequence_tensor(v, "v"); - ChunkGatedDeltaRule::execute(out, - initial_state, - final_state, - q, - k, - v, - g, - beta, - cu_seqlens, - initial_state_indices, - final_state_indices, - use_qk_l2norm, - chunk_size); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_infiniop.cc b/src/infinicore/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_infiniop.cc deleted file mode 100644 index 96a2b3b44..000000000 --- a/src/infinicore/ops/chunk_gated_delta_rule/chunk_gated_delta_rule_infiniop.cc +++ /dev/null @@ -1,108 +0,0 @@ -#include "infinicore/ops/chunk_gated_delta_rule.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::chunk_gated_delta_rule_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, ChunkGatedDeltaRule, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, initial_state, q, k, v, g, beta; - std::optional final_state; - std::optional cu_seqlens; - std::optional initial_state_indices; - std::optional final_state_indices; -}; - -void *plan(Tensor out, - Tensor initial_state, - std::optional final_state, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - std::optional cu_seqlens, - std::optional initial_state_indices, - std::optional final_state_indices, - bool use_qk_l2norm, - size_t chunk_size) { - size_t seed = hash_combine(out, - initial_state, - final_state, - q, - k, - v, - g, - beta, - cu_seqlens, - initial_state_indices, - final_state_indices, - use_qk_l2norm, - chunk_size); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, ChunkGatedDeltaRule, - seed, - out->desc(), - initial_state->desc(), - final_state.has_value() ? final_state.value()->desc() : nullptr, - q->desc(), - k->desc(), - v->desc(), - g->desc(), - beta->desc(), - cu_seqlens.has_value() ? cu_seqlens.value()->desc() : nullptr, - initial_state_indices.has_value() ? initial_state_indices.value()->desc() : nullptr, - final_state_indices.has_value() ? final_state_indices.value()->desc() : nullptr, - use_qk_l2norm, - chunk_size); - - INFINIOP_WORKSPACE_TENSOR(workspace, ChunkGatedDeltaRule, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(initial_state), - graph::GraphTensor(q), - graph::GraphTensor(k), - graph::GraphTensor(v), - graph::GraphTensor(g), - graph::GraphTensor(beta), - final_state.has_value() ? std::optional(graph::GraphTensor(final_state.value())) : std::nullopt, - cu_seqlens.has_value() ? std::optional(graph::GraphTensor(cu_seqlens.value())) : std::nullopt, - initial_state_indices.has_value() ? std::optional(graph::GraphTensor(initial_state_indices.value())) : std::nullopt, - final_state_indices.has_value() ? std::optional(graph::GraphTensor(final_state_indices.value())) : std::nullopt}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopChunkGatedDeltaRule( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->initial_state->data(), - planned->final_state.has_value() ? planned->final_state.value()->data() : nullptr, - planned->q->data(), - planned->k->data(), - planned->v->data(), - planned->g->data(), - planned->beta->data(), - planned->cu_seqlens.has_value() ? planned->cu_seqlens.value()->data() : nullptr, - planned->initial_state_indices.has_value() ? planned->initial_state_indices.value()->data() : nullptr, - planned->final_state_indices.has_value() ? planned->final_state_indices.value()->data() : nullptr, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(ChunkGatedDeltaRule, &plan, &run, &cleanup); - -} // namespace infinicore::op::chunk_gated_delta_rule_impl::infiniop diff --git a/src/infinicore/ops/conv2d/conv2d.cc b/src/infinicore/ops/conv2d/conv2d.cc deleted file mode 100644 index 4aa040536..000000000 --- a/src/infinicore/ops/conv2d/conv2d.cc +++ /dev/null @@ -1,101 +0,0 @@ -#include "infinicore/ops/conv2d.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &Conv2d::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Conv2d::execute(Tensor output, - Tensor input, - Tensor weight, - Tensor bias, - const size_t *pads, - const size_t *strides, - const size_t *dilations, - size_t n) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, weight); - if (bias) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, bias); - } - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Conv2d implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input, weight, bias, pads, strides, dilations, n); -} - -Tensor conv2d(Tensor input, - Tensor weight, - Tensor bias, - const std::vector &pads, - const std::vector &strides, - const std::vector &dilations) { - const auto &in_shape = input->shape(); // [N, C_in, H_in, W_in] - const auto &w_shape = weight->shape(); // [C_out, C_in, kH, kW] - - // ------------------------------- - // Extract dimensions - // ------------------------------- - size_t N = in_shape[0]; - size_t C_in = in_shape[1]; - size_t H_in = in_shape[2]; - size_t W_in = in_shape[3]; - - size_t C_out = w_shape[0]; - size_t kH = w_shape[2]; - size_t kW = w_shape[3]; - - size_t pad_h = pads[0]; - size_t pad_w = pads[1]; - - size_t stride_h = strides[0]; - size_t stride_w = strides[1]; - - size_t dil_h = dilations[0]; - size_t dil_w = dilations[1]; - - auto calc_out = [](size_t in, size_t pad, size_t dilation, size_t kernel, size_t stride) { - return (in + 2 * pad - dilation * (kernel - 1) - 1) / stride + 1; - }; - size_t H_out = calc_out(H_in, pad_h, dil_h, kH, stride_h); - size_t W_out = calc_out(W_in, pad_w, dil_w, kW, stride_w); - if ((int64_t)H_out <= 0 || (int64_t)W_out <= 0) { - throw std::runtime_error("Invalid conv2d output shape (negative or zero)"); - } - Shape out_shape = {N, C_out, H_out, W_out}; - - auto output = Tensor::empty(out_shape, input->dtype(), input->device()); - conv2d_(output, input, weight, bias, pads, strides, dilations); - return output; -} - -void conv2d_(Tensor output, - Tensor input, - Tensor weight, - Tensor bias, - const std::vector &pads, - const std::vector &strides, - const std::vector &dilations) { - if (pads.size() != strides.size() || pads.size() != dilations.size()) { - throw std::runtime_error("conv2d_: pads/strides/dilations must have the same size"); - } - Conv2d::execute(output, - input, - weight, - bias, - pads.data(), - strides.data(), - dilations.data(), - pads.size()); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/conv2d/conv2d_infiniop.cc b/src/infinicore/ops/conv2d/conv2d_infiniop.cc deleted file mode 100644 index c15c14fdd..000000000 --- a/src/infinicore/ops/conv2d/conv2d_infiniop.cc +++ /dev/null @@ -1,69 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/conv2d.hpp" -#include - -namespace infinicore::op::conv2d_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopConvDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyConvDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, - Tensor input, - Tensor weight, - Tensor bias, - const size_t *pads, - const size_t *strides, - const size_t *dilations, - size_t n) { - size_t seed = hash_combine(output, input, weight, bias, n); - for (size_t i = 0; i < n; ++i) { - hash_combine(seed, pads[i], strides[i], dilations[i]); - } - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopConvDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateConvDescriptor( - context::getInfiniopHandle(device), &desc, - output->desc(), input->desc(), weight->desc(), - bias ? bias->desc() : nullptr, - const_cast(pads), - const_cast(strides), - const_cast(dilations), - n)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetConvWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopConv( - desc, workspace->data(), workspace_size, - output->data(), - input->data(), - weight->data(), - bias ? bias->data() : nullptr, - context::getStream())); -} - -static bool registered = []() { - Conv2d::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::conv2d_impl::infiniop diff --git a/src/infinicore/ops/conv2d/conv2d_infiniops.cc b/src/infinicore/ops/conv2d/conv2d_infiniops.cc deleted file mode 100644 index 3cdf0c63e..000000000 --- a/src/infinicore/ops/conv2d/conv2d_infiniops.cc +++ /dev/null @@ -1,72 +0,0 @@ -#include "infinicore/ops/conv2d.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/conv_infinilm.h" - -#include -#include - -namespace infinicore::op::conv2d_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -std::vector toInt64Vector(const size_t *values, size_t n) { - std::vector result; - result.reserve(n); - for (size_t i = 0; i < n; ++i) { - result.push_back(static_cast(values[i])); - } - return result; -} - -void calculate(Tensor output, - Tensor input, - Tensor weight, - Tensor bias, - const size_t *pads, - const size_t *strides, - const size_t *dilations, - size_t n) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, weight); - if (bias) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, bias); - } - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - TensorMeta output_meta(output); - TensorMeta input_meta(input); - TensorMeta weight_meta(weight); - std::optional bias_meta; - if (bias) { - bias_meta.emplace(bias); - } - - infini::ops::ConvInfinilm::Call( - handle, - config, - input_meta.tensor(input), - weight_meta.tensor(weight), - bias_meta ? std::optional{bias_meta->tensor(bias)} : std::nullopt, - toInt64Vector(pads, n), - toInt64Vector(strides, n), - toInt64Vector(dilations, n), - int64_t{1}, - output_meta.tensor(output)); -} - -} // namespace - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(Conv2d::dispatcher(), &calculate); - return true; -}(); - -} // namespace infinicore::op::conv2d_impl::infiniops -#endif diff --git a/src/infinicore/ops/cross_entropy/cross_entropy.cc b/src/infinicore/ops/cross_entropy/cross_entropy.cc deleted file mode 100644 index 84aebc1b1..000000000 --- a/src/infinicore/ops/cross_entropy/cross_entropy.cc +++ /dev/null @@ -1,45 +0,0 @@ -#include "infinicore/ops/cross_entropy.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &CrossEntropy::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void CrossEntropy::execute(Tensor output, Tensor input, Tensor target) { - - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(input, target); - - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No CrossEntropy implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input, target); -} - -Tensor cross_entropy(Tensor input, Tensor target) { - - Shape shape = target->shape(); - - auto output = Tensor::empty(shape, input->dtype(), input->device()); - - cross_entropy_(output, input, target); - return output; -} - -void cross_entropy_(Tensor output, Tensor input, Tensor target) { - CrossEntropy::execute(output, input, target); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/cross_entropy/cross_entropy_infiniop.cc b/src/infinicore/ops/cross_entropy/cross_entropy_infiniop.cc deleted file mode 100644 index 5fa7963d7..000000000 --- a/src/infinicore/ops/cross_entropy/cross_entropy_infiniop.cc +++ /dev/null @@ -1,64 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" - -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/cross_entropy.hpp" - -#include - -namespace infinicore::op::cross_entropy_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopCrossEntropyDescriptor_t &desc) { - if (desc != nullptr) { - - INFINICORE_CHECK_ERROR(infiniopDestroyCrossEntropyDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, Tensor target) { - - size_t seed = hash_combine(output, input, target); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopCrossEntropyDescriptor_t desc = nullptr; - - if (!desc_opt) { - - INFINICORE_CHECK_ERROR(infiniopCreateCrossEntropyDescriptor( - context::getInfiniopHandle(device), - &desc, - output->desc(), - input->desc(), - target->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetCrossEntropyWorkspaceSize(desc, &workspace_size)); - - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopCrossEntropy( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - target->data(), - context::getStream())); -} - -static bool registered = []() { - CrossEntropy::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::cross_entropy_impl::infiniop diff --git a/src/infinicore/ops/deepseek_moe/deepseek_moe.cc b/src/infinicore/ops/deepseek_moe/deepseek_moe.cc deleted file mode 100644 index 82e1c56dd..000000000 --- a/src/infinicore/ops/deepseek_moe/deepseek_moe.cc +++ /dev/null @@ -1,83 +0,0 @@ -#include "infinicore/ops/deepseek_moe.hpp" -#include "../../utils.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(DeepseekMoe); - -namespace { - -void check_weights(const std::vector &gate_weights, - const std::vector &up_weights, - const std::vector &down_weights, - size_t num_experts) { - if (gate_weights.size() != num_experts || up_weights.size() != num_experts || down_weights.size() != num_experts) { - throw std::runtime_error("DeepseekMoe: expert weight vector size mismatch"); - } -} - -} // namespace - -DeepseekMoe::DeepseekMoe(Tensor out, - const Tensor &hidden, - const Tensor &topk_indices, - const Tensor &topk_weights, - const std::vector &gate_weights, - const std::vector &up_weights, - const std::vector &down_weights, - size_t intermediate_size, - size_t num_experts) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, hidden, topk_indices, topk_weights); - check_weights(gate_weights, up_weights, down_weights, num_experts); - for (size_t i = 0; i < num_experts; ++i) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, gate_weights[i], up_weights[i], down_weights[i]); - } - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), - out, hidden, topk_indices, topk_weights, - gate_weights, up_weights, down_weights, - intermediate_size, num_experts); -} - -void DeepseekMoe::execute(Tensor out, - const Tensor &hidden, - const Tensor &topk_indices, - const Tensor &topk_weights, - const std::vector &gate_weights, - const std::vector &up_weights, - const std::vector &down_weights, - size_t intermediate_size, - size_t num_experts) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - DeepseekMoe, - out, hidden, topk_indices, topk_weights, - gate_weights, up_weights, down_weights, - intermediate_size, num_experts); -} - -void deepseek_moe_(Tensor out, - const Tensor &hidden, - const Tensor &topk_indices, - const Tensor &topk_weights, - const std::vector &gate_weights, - const std::vector &up_weights, - const std::vector &down_weights, - size_t intermediate_size, - size_t num_experts) { - DeepseekMoe::execute(out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, intermediate_size, num_experts); -} - -Tensor deepseek_moe(const Tensor &hidden, - const Tensor &topk_indices, - const Tensor &topk_weights, - const std::vector &gate_weights, - const std::vector &up_weights, - const std::vector &down_weights, - size_t intermediate_size, - size_t num_experts) { - auto out = Tensor::empty(hidden->shape(), hidden->dtype(), hidden->device()); - deepseek_moe_(out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, intermediate_size, num_experts); - return out; -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/deepseek_moe/deepseek_moe_infiniop.cc b/src/infinicore/ops/deepseek_moe/deepseek_moe_infiniop.cc deleted file mode 100644 index bc410946e..000000000 --- a/src/infinicore/ops/deepseek_moe/deepseek_moe_infiniop.cc +++ /dev/null @@ -1,113 +0,0 @@ -#include "infinicore/ops/deepseek_moe.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::deepseek_moe_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, DeepseekMoe, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - Tensor workspace_owner, out_owner, hidden_owner, topk_indices_owner, topk_weights_owner; - graph::GraphTensor workspace, out, hidden, topk_indices, topk_weights; - std::vector gate_weights, up_weights, down_weights; - std::vector gate_ptrs, up_ptrs, down_ptrs; - std::shared_ptr gate_ptrs_device, up_ptrs_device, down_ptrs_device; -}; - -static std::vector to_graph_tensors(const std::vector &tensors) { - std::vector result; - result.reserve(tensors.size()); - for (const auto &tensor : tensors) { - result.emplace_back(tensor); - } - return result; -} - -static std::vector data_ptrs(const std::vector &tensors) { - std::vector result; - result.reserve(tensors.size()); - for (const auto &tensor : tensors) { - result.push_back(tensor->data()); - } - return result; -} - -void *plan(Tensor out, - const Tensor &hidden, - const Tensor &topk_indices, - const Tensor &topk_weights, - const std::vector &gate_weights, - const std::vector &up_weights, - const std::vector &down_weights, - size_t intermediate_size, - size_t num_experts) { - size_t seed = hash_combine(out, hidden, topk_indices, topk_weights); - hash_combine(seed, intermediate_size, num_experts); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, DeepseekMoe, seed, - out->desc(), hidden->desc(), topk_indices->desc(), topk_weights->desc(), - intermediate_size, num_experts); - - INFINIOP_WORKSPACE_TENSOR(workspace, DeepseekMoe, descriptor); - - auto planned = new PlannedMeta{ - descriptor, - workspace, - out, - hidden, - topk_indices, - topk_weights, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(hidden), - graph::GraphTensor(topk_indices), - graph::GraphTensor(topk_weights), - to_graph_tensors(gate_weights), - to_graph_tensors(up_weights), - to_graph_tensors(down_weights), - {}, - {}, - {}, - nullptr, - nullptr, - nullptr}; - planned->gate_ptrs = data_ptrs(planned->gate_weights); - planned->up_ptrs = data_ptrs(planned->up_weights); - planned->down_ptrs = data_ptrs(planned->down_weights); - const size_t ptr_bytes = num_experts * sizeof(void *); - planned->gate_ptrs_device = context::allocateMemory(ptr_bytes); - planned->up_ptrs_device = context::allocateMemory(ptr_bytes); - planned->down_ptrs_device = context::allocateMemory(ptr_bytes); - context::memcpyH2D(planned->gate_ptrs_device->data(), planned->gate_ptrs.data(), ptr_bytes, false); - context::memcpyH2D(planned->up_ptrs_device->data(), planned->up_ptrs.data(), ptr_bytes, false); - context::memcpyH2D(planned->down_ptrs_device->data(), planned->down_ptrs.data(), ptr_bytes, false); - return planned; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopDeepseekMoeWithDevicePtrs( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->hidden->data(), - planned->topk_indices->data(), - planned->topk_weights->data(), - planned->gate_ptrs_device->data(), - planned->up_ptrs_device->data(), - planned->down_ptrs_device->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(DeepseekMoe, &plan, &run, &cleanup); - -} // namespace infinicore::op::deepseek_moe_impl::infiniop diff --git a/src/infinicore/ops/dequant/per_tensor_dequant_int8/per_tensor_quant_int8.cc b/src/infinicore/ops/dequant/per_tensor_dequant_int8/per_tensor_quant_int8.cc deleted file mode 100644 index dc403e18c..000000000 --- a/src/infinicore/ops/dequant/per_tensor_dequant_int8/per_tensor_quant_int8.cc +++ /dev/null @@ -1,20 +0,0 @@ -#include "../../../utils.hpp" -#include "infinicore/ops/per_tensor_dequant_i8.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PerTensorDequantI8); - -PerTensorDequantI8::PerTensorDequantI8(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zero) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, x_packed, x_scale, x_zero); - INFINICORE_GRAPH_OP_DISPATCH(x->device().getType(), x, x_packed, x_scale, x_zero); -} - -void PerTensorDequantI8::execute(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zero) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(PerTensorDequantI8, x, x_packed, x_scale, x_zero); -} - -void per_tensor_dequant_i8_(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zero) { - PerTensorDequantI8::execute(x, x_packed, x_scale, x_zero); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/dequant/per_tensor_dequant_int8/per_tensor_quant_int8_infiniop.cc b/src/infinicore/ops/dequant/per_tensor_dequant_int8/per_tensor_quant_int8_infiniop.cc deleted file mode 100644 index 353759ba8..000000000 --- a/src/infinicore/ops/dequant/per_tensor_dequant_int8/per_tensor_quant_int8_infiniop.cc +++ /dev/null @@ -1,53 +0,0 @@ -#include "../../infiniop_impl.hpp" -#include "infinicore/ops/per_tensor_dequant_i8.hpp" - -namespace infinicore::op::per_tensor_dequant_i8_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, PerTensorDequantI8, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, x, x_packed, x_scale, x_zero; -}; - -void *plan(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zero) { - size_t seed = hash_combine(x, x_packed, x_scale); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, PerTensorDequantI8, - seed, - x->desc(), x_packed->desc(), x_scale->desc(), (x_zero ? x_zero->desc() : nullptr)); - - INFINIOP_WORKSPACE_TENSOR(workspace, PerTensorDequantI8, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x), - graph::GraphTensor(x_packed), - graph::GraphTensor(x_scale), - graph::GraphTensor(x_zero)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopPerTensorDequantI8( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->x->data(), - planned->x_packed->data(), - planned->x_scale->data(), - nullptr, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(PerTensorDequantI8, &plan, &run, &cleanup); - -} // namespace infinicore::op::per_tensor_dequant_i8_impl::infiniop diff --git a/src/infinicore/ops/dequantize_awq/dequantize_awq.cc b/src/infinicore/ops/dequantize_awq/dequantize_awq.cc deleted file mode 100644 index dff92b6ec..000000000 --- a/src/infinicore/ops/dequantize_awq/dequantize_awq.cc +++ /dev/null @@ -1,20 +0,0 @@ -#include "infinicore/ops/dequantize_awq.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(DequantizeAWQ); - -DequantizeAWQ::DequantizeAWQ(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zeros) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, x_packed, x_scale, x_zeros); - INFINICORE_GRAPH_OP_DISPATCH(x->device().getType(), x, x_packed, x_scale, x_zeros); -} - -void DequantizeAWQ::execute(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zeros) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(DequantizeAWQ, x, x_packed, x_scale, x_zeros); -} - -void dequantize_awq_(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zeros) { - DequantizeAWQ::execute(x, x_packed, x_scale, x_zeros); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/dequantize_awq/dequantize_awq_infiniop.cc b/src/infinicore/ops/dequantize_awq/dequantize_awq_infiniop.cc deleted file mode 100644 index 3e643ee40..000000000 --- a/src/infinicore/ops/dequantize_awq/dequantize_awq_infiniop.cc +++ /dev/null @@ -1,56 +0,0 @@ -#include "../../utils.hpp" -#include "../infiniop_impl.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/dequantize_awq.hpp" -#include - -namespace infinicore::op::dequantize_awq_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, DequantizeAWQ, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, x, x_packed, x_scale, x_zeros; -}; - -void *plan(Tensor x, const Tensor &x_packed, const Tensor &x_scale, const Tensor &x_zeros) { - size_t seed = hash_combine(x, x_packed, x_scale, x_zeros); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, DequantizeAWQ, - seed, - x->desc(), x_packed->desc(), x_scale->desc(), x_zeros->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, DequantizeAWQ, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x), - graph::GraphTensor(x_packed), - graph::GraphTensor(x_scale), - graph::GraphTensor(x_zeros)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopDequantizeAWQ( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->x->data(), - planned->x_packed->data(), - planned->x_scale->data(), - planned->x_zeros->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(DequantizeAWQ, &plan, &run, &cleanup); -} // namespace infinicore::op::dequantize_awq_impl::infiniop diff --git a/src/infinicore/ops/diff/diff.cc b/src/infinicore/ops/diff/diff.cc deleted file mode 100644 index 66bc2c0a7..000000000 --- a/src/infinicore/ops/diff/diff.cc +++ /dev/null @@ -1,78 +0,0 @@ -#include "infinicore/ops/diff.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Diff); - -Diff::Diff(Tensor y, const Tensor &x, int dim, int n) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, x, dim, n); -} - -void Diff::execute(Tensor y, const Tensor &x, int dim, int n) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Diff, y, x, dim, n); -} - -static int normalize_dim(int dim, size_t ndim) { - if (ndim == 0) { - throw std::runtime_error("diff: input tensor must have at least one dimension."); - } - if (dim < 0) { - dim += static_cast(ndim); - } - if (dim < 0 || static_cast(dim) >= ndim) { - throw std::runtime_error("diff: dim out of range."); - } - return dim; -} - -Tensor diff(const Tensor &x, int n, int dim) { - if (n < 0) { - throw std::runtime_error("diff: n must be non-negative."); - } - Shape y_shape = x->shape(); - const int d = normalize_dim(dim, y_shape.size()); - const auto dim_size = y_shape[static_cast(d)]; - y_shape[static_cast(d)] = (dim_size >= static_cast(n)) ? (dim_size - static_cast(n)) : 0; - - auto y = Tensor::empty(y_shape, x->dtype(), x->device()); - if (n == 0) { - y->copy_from(x); - return y; - } - if (dim_size <= static_cast(n)) { - // Empty output by definition; nothing to compute. - return y; - } - - diff_(y, x, n, dim); - return y; -} - -void diff_(Tensor y, const Tensor &x, int n, int dim) { - if (n < 0) { - throw std::runtime_error("diff_: n must be non-negative."); - } - const int d = normalize_dim(dim, x->shape().size()); - Shape expected = x->shape(); - const auto dim_size = expected[static_cast(d)]; - expected[static_cast(d)] = (dim_size >= static_cast(n)) ? (dim_size - static_cast(n)) : 0; - if (y->shape() != expected) { - throw std::runtime_error("diff_: output tensor has incorrect shape."); - } - if (n == 0) { - y->copy_from(x); - return; - } - if (x->shape()[static_cast(d)] <= static_cast(n)) { - // Empty output by definition; nothing to compute. - return; - } - Diff::execute(y, x, d, n); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/diff/diff_infiniop.cc b/src/infinicore/ops/diff/diff_infiniop.cc deleted file mode 100644 index d6813782d..000000000 --- a/src/infinicore/ops/diff/diff_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/diff.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::diff_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Diff, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, y, x; -}; - -void *plan(Tensor y, const Tensor &x, int dim, int n) { - size_t seed = hash_combine(y, x, dim, n); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Diff, - seed, - y->desc(), x->desc(), dim, n); - - INFINIOP_WORKSPACE_TENSOR(workspace, Diff, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(y), - graph::GraphTensor(x)}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopDiff( - p->descriptor->desc, - p->workspace ? p->workspace->data() : nullptr, - p->workspace ? p->workspace->numel() : 0, - p->y->data(), - p->x->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Diff, &plan, &run, &cleanup); - -} // namespace infinicore::op::diff_impl::infiniop diff --git a/src/infinicore/ops/digamma/digamma.cc b/src/infinicore/ops/digamma/digamma.cc deleted file mode 100644 index 2c62f87ae..000000000 --- a/src/infinicore/ops/digamma/digamma.cc +++ /dev/null @@ -1,28 +0,0 @@ -#include "infinicore/ops/digamma.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Digamma); - -Digamma::Digamma(Tensor y, const Tensor &x) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, x); -} - -void Digamma::execute(Tensor y, const Tensor &x) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Digamma, y, x); -} - -Tensor digamma(const Tensor &x) { - auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); - digamma_(y, x); - return y; -} - -void digamma_(Tensor y, const Tensor &x) { - Digamma::execute(y, x); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/digamma/digamma_infiniop.cc b/src/infinicore/ops/digamma/digamma_infiniop.cc deleted file mode 100644 index c39fa3a0c..000000000 --- a/src/infinicore/ops/digamma/digamma_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/digamma.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::digamma_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Digamma, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, y, x; -}; - -void *plan(Tensor y, const Tensor &x) { - size_t seed = hash_combine(y, x); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Digamma, - seed, - y->desc(), x->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Digamma, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(y), - graph::GraphTensor(x)}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopDigamma( - p->descriptor->desc, - p->workspace ? p->workspace->data() : nullptr, - p->workspace ? p->workspace->numel() : 0, - p->y->data(), - p->x->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Digamma, &plan, &run, &cleanup); - -} // namespace infinicore::op::digamma_impl::infiniop diff --git a/src/infinicore/ops/dist/dist.cc b/src/infinicore/ops/dist/dist.cc deleted file mode 100644 index 794b192fc..000000000 --- a/src/infinicore/ops/dist/dist.cc +++ /dev/null @@ -1,30 +0,0 @@ -#include "infinicore/ops/dist.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Dist); - -Dist::Dist(Tensor y, const Tensor &x1, const Tensor &x2, double p) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x1, x2); - INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, x1, x2, p); -} - -void Dist::execute(Tensor y, const Tensor &x1, const Tensor &x2, double p) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Dist, y, x1, x2, p); -} - -Tensor dist(const Tensor &x1, const Tensor &x2, double p) { - auto y = Tensor::empty({}, x1->dtype(), x1->device()); - dist_(y, x1, x2, p); - return y; -} - -void dist_(Tensor y, const Tensor &x1, const Tensor &x2, double p) { - Dist::execute(y, x1, x2, p); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/dist/dist_infiniop.cc b/src/infinicore/ops/dist/dist_infiniop.cc deleted file mode 100644 index c42aae0d8..000000000 --- a/src/infinicore/ops/dist/dist_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "infinicore/ops/dist.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::dist_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Dist, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, y, x1, x2; -}; - -void *plan(Tensor y, const Tensor &x1, const Tensor &x2, double p) { - size_t seed = hash_combine(y, x1, x2, p); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Dist, - seed, - y->desc(), x1->desc(), x2->desc(), p); - - INFINIOP_WORKSPACE_TENSOR(workspace, Dist, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(y), - graph::GraphTensor(x1), - graph::GraphTensor(x2)}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopDist( - p->descriptor->desc, - p->workspace ? p->workspace->data() : nullptr, - p->workspace ? p->workspace->numel() : 0, - p->y->data(), - p->x1->data(), - p->x2->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Dist, &plan, &run, &cleanup); - -} // namespace infinicore::op::dist_impl::infiniop diff --git a/src/infinicore/ops/distributed/allgather.cc b/src/infinicore/ops/distributed/allgather.cc deleted file mode 100644 index fcf9ab810..000000000 --- a/src/infinicore/ops/distributed/allgather.cc +++ /dev/null @@ -1,201 +0,0 @@ -#include "infinicore/ops/distributed/allgather.hpp" -#include "../../utils.hpp" - -#include "infinicore/context/context.hpp" - -#include -#include - -namespace infinicore::op::distributed { -namespace { - -bool all_equal_to(const std::vector &values, size_t expected) { - return !values.empty() && std::all_of(values.begin(), values.end(), [&](size_t value) { - return value == expected; - }); -} - -bool can_use_grouped_allgather(const std::vector &inputs, - const std::vector &split_sizes) { - return !inputs.empty() && all_equal_to(split_sizes, inputs.front()->shape()[0]) && std::all_of(inputs.begin(), inputs.end(), [&](const Tensor &input) { - return input && input->ndim() > 0 && input->shape()[0] == split_sizes.front(); - }); -} - -} // namespace - -struct AllGatherPlannedMeta { - graph::GraphTensor output, input; - infinicclComm_t communicator; -}; - -struct AllGatherVPlannedMeta { - graph::GraphTensor output, input; - std::vector split_counts; - infinicclComm_t communicator; -}; - -AllGather::AllGather(Tensor output, const Tensor &input, infinicclComm_t communicator) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); - INFINICORE_ASSERT(input->numel() > 0); - INFINICORE_ASSERT(output->numel() % input->numel() == 0); - planned_meta_ = new AllGatherPlannedMeta{graph::GraphTensor(output), graph::GraphTensor(input), communicator}; -} - -AllGather::~AllGather() { - if (planned_meta_) { - auto *meta = reinterpret_cast(planned_meta_); - delete meta; - } -} - -void AllGather::run() const { - auto *meta = reinterpret_cast(planned_meta_); - INFINICORE_CHECK_ERROR(infinicclAllGather(meta->input->data(), - meta->output->data(), - meta->input->numel(), - static_cast(static_cast(meta->input->dtype())), - meta->communicator, - infinicore::context::getStream())); -} - -void AllGather::execute(Tensor output, const Tensor &input, infinicclComm_t communicator) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(AllGather, output, input, communicator); -} - -AllGatherV::AllGatherV(Tensor output, - const Tensor &input, - std::vector split_counts, - infinicclComm_t communicator) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); - INFINICORE_ASSERT(input->numel() > 0); - size_t total_count = 0; - for (auto count : split_counts) { - total_count += count; - } - INFINICORE_ASSERT(output->numel() == total_count); - planned_meta_ = new AllGatherVPlannedMeta{ - graph::GraphTensor(output), - graph::GraphTensor(input), - std::move(split_counts), - communicator, - }; -} - -AllGatherV::~AllGatherV() { - if (planned_meta_) { - auto *meta = reinterpret_cast(planned_meta_); - delete meta; - } -} - -void AllGatherV::run() const { - auto *meta = reinterpret_cast(planned_meta_); - INFINICORE_CHECK_ERROR(infinicclAllGatherV(meta->input->data(), - meta->output->data(), - meta->split_counts.data(), - static_cast(meta->split_counts.size()), - static_cast(static_cast(meta->input->dtype())), - meta->communicator, - infinicore::context::getStream())); -} - -void AllGatherV::execute(Tensor output, - const Tensor &input, - std::vector split_counts, - infinicclComm_t communicator) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(AllGatherV, output, input, std::move(split_counts), communicator); -} - -Tensor allgather(const Tensor &input, size_t world_size, infinicclComm_t communicator) { - INFINICORE_ASSERT(input->ndim() > 0); - auto shape = input->shape(); - shape[0] *= world_size; - auto output = Tensor::empty(shape, input->dtype(), input->device()); - allgather_(output, input, communicator); - return output; -} - -void allgather_(Tensor output, const Tensor &input, infinicclComm_t communicator) { - AllGather::execute(output, input, communicator); -} - -Tensor allgatherv(const Tensor &input, const std::vector &split_sizes, infinicclComm_t communicator) { - INFINICORE_ASSERT(input->ndim() > 0); - size_t total_dim0 = 0; - for (auto size : split_sizes) { - total_dim0 += size; - } - auto shape = input->shape(); - shape[0] = total_dim0; - auto output = Tensor::empty(shape, input->dtype(), input->device()); - allgatherv_(output, input, split_sizes, communicator); - return output; -} - -void allgatherv_(Tensor output, const Tensor &input, const std::vector &split_sizes, infinicclComm_t communicator) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); - INFINICORE_ASSERT(input->ndim() > 0); - INFINICORE_ASSERT(input->shape()[0] > 0); - if (all_equal_to(split_sizes, input->shape()[0])) { - INFINICORE_ASSERT(output->numel() == input->numel() * split_sizes.size()); - AllGather::execute(output, input, communicator); - return; - } - - const size_t inner = input->numel() / input->shape()[0]; - std::vector split_counts; - split_counts.reserve(split_sizes.size()); - size_t total_dim0 = 0; - for (auto size : split_sizes) { - split_counts.push_back(size * inner); - total_dim0 += size; - } - INFINICORE_ASSERT(output->numel() == total_dim0 * inner); - AllGatherV::execute(output, input, std::move(split_counts), communicator); -} - -std::vector allgatherv_many(const std::vector &inputs, - const std::vector &split_sizes, - infinicclComm_t communicator) { - std::vector outputs; - outputs.reserve(inputs.size()); - size_t total_dim0 = 0; - for (auto size : split_sizes) { - total_dim0 += size; - } - for (const auto &input : inputs) { - INFINICORE_ASSERT(input->ndim() > 0); - auto shape = input->shape(); - shape[0] = total_dim0; - outputs.push_back(Tensor::empty(shape, input->dtype(), input->device())); - } - allgatherv_many_(outputs, inputs, split_sizes, communicator); - return outputs; -} - -void allgatherv_many_(const std::vector &outputs, - const std::vector &inputs, - const std::vector &split_sizes, - infinicclComm_t communicator) { - INFINICORE_ASSERT(outputs.size() == inputs.size()); - if (inputs.empty()) { - return; - } - - const bool use_group = can_use_grouped_allgather(inputs, split_sizes) && !infinicore::context::isGraphRecording(); - if (use_group) { - INFINICORE_CHECK_ERROR(infinicclGroupStart(communicator)); - } - for (size_t i = 0; i < inputs.size(); ++i) { - allgatherv_(outputs[i], inputs[i], split_sizes, communicator); - } - if (use_group) { - INFINICORE_CHECK_ERROR(infinicclGroupEnd(communicator)); - } -} - -} // namespace infinicore::op::distributed diff --git a/src/infinicore/ops/distributed/allreduce.cc b/src/infinicore/ops/distributed/allreduce.cc deleted file mode 100644 index ddfc238c9..000000000 --- a/src/infinicore/ops/distributed/allreduce.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/distributed/allreduce.hpp" -#include "../../utils.hpp" - -namespace infinicore::op::distributed { - -struct PlannedMeta { - graph::GraphTensor output, input; - infinicclReduceOp_t op; - infinicclComm_t communicator; -}; - -AllReduce::AllReduce(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); - INFINICORE_ASSERT(output->numel() == input->numel()); - planned_meta_ = new PlannedMeta{graph::GraphTensor(output), graph::GraphTensor(input), op, communicator}; -} -AllReduce::~AllReduce() { - if (planned_meta_) { - PlannedMeta *meta = reinterpret_cast(planned_meta_); - delete meta; - } -} - -void AllReduce::run() const { - PlannedMeta *meta = reinterpret_cast(planned_meta_); - - INFINICORE_CHECK_ERROR(infinicclAllReduce(meta->input->data(), - meta->output->data(), - meta->input->numel(), - static_cast(static_cast(meta->input->dtype())), - meta->op, - meta->communicator, - infinicore::context::getStream())); -} - -void AllReduce::execute(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(AllReduce, output, input, op, communicator); -} - -Tensor allreduce(const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator) { - auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); - allreduce_(output, input, op, communicator); - return output; -} - -void allreduce_(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator) { - AllReduce::execute(output, input, op, communicator); -} -} // namespace infinicore::op::distributed diff --git a/src/infinicore/ops/distributed/reduce_scatter.cc b/src/infinicore/ops/distributed/reduce_scatter.cc deleted file mode 100644 index 60ec24624..000000000 --- a/src/infinicore/ops/distributed/reduce_scatter.cc +++ /dev/null @@ -1,214 +0,0 @@ -#include "infinicore/ops/distributed/reduce_scatter.hpp" -#include "../../utils.hpp" - -#include "infinicore/context/context.hpp" - -#include -#include - -namespace infinicore::op::distributed { -namespace { - -bool all_equal_to(const std::vector &values, size_t expected) { - return !values.empty() && std::all_of(values.begin(), values.end(), [&](size_t value) { - return value == expected; - }); -} - -bool can_use_grouped_reduce_scatter(const std::vector &outputs, - const std::vector &split_sizes) { - return !outputs.empty() && all_equal_to(split_sizes, outputs.front()->shape()[0]) && std::all_of(outputs.begin(), outputs.end(), [&](const Tensor &output) { - return output && output->ndim() > 0 && output->shape()[0] == split_sizes.front(); - }); -} - -} // namespace - -struct ReduceScatterPlannedMeta { - graph::GraphTensor output, input; - infinicclReduceOp_t op; - infinicclComm_t communicator; -}; - -struct ReduceScatterVPlannedMeta { - graph::GraphTensor output, input; - std::vector split_counts; - infinicclReduceOp_t op; - infinicclComm_t communicator; -}; - -ReduceScatter::ReduceScatter(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); - INFINICORE_ASSERT(output->numel() > 0); - INFINICORE_ASSERT(input->numel() % output->numel() == 0); - planned_meta_ = new ReduceScatterPlannedMeta{graph::GraphTensor(output), graph::GraphTensor(input), op, communicator}; -} - -ReduceScatter::~ReduceScatter() { - if (planned_meta_) { - auto *meta = reinterpret_cast(planned_meta_); - delete meta; - } -} - -void ReduceScatter::run() const { - auto *meta = reinterpret_cast(planned_meta_); - INFINICORE_CHECK_ERROR(infinicclReduceScatter(meta->input->data(), - meta->output->data(), - meta->output->numel(), - static_cast(static_cast(meta->input->dtype())), - meta->op, - meta->communicator, - infinicore::context::getStream())); -} - -void ReduceScatter::execute(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(ReduceScatter, output, input, op, communicator); -} - -ReduceScatterV::ReduceScatterV(Tensor output, - const Tensor &input, - std::vector split_counts, - infinicclReduceOp_t op, - infinicclComm_t communicator) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); - INFINICORE_ASSERT(output->numel() > 0); - size_t total_count = 0; - for (auto count : split_counts) { - total_count += count; - } - INFINICORE_ASSERT(input->numel() == total_count); - planned_meta_ = new ReduceScatterVPlannedMeta{ - graph::GraphTensor(output), - graph::GraphTensor(input), - std::move(split_counts), - op, - communicator, - }; -} - -ReduceScatterV::~ReduceScatterV() { - if (planned_meta_) { - auto *meta = reinterpret_cast(planned_meta_); - delete meta; - } -} - -void ReduceScatterV::run() const { - auto *meta = reinterpret_cast(planned_meta_); - INFINICORE_CHECK_ERROR(infinicclReduceScatterV(meta->input->data(), - meta->output->data(), - meta->split_counts.data(), - static_cast(meta->split_counts.size()), - static_cast(static_cast(meta->input->dtype())), - meta->op, - meta->communicator, - infinicore::context::getStream())); -} - -void ReduceScatterV::execute(Tensor output, - const Tensor &input, - std::vector split_counts, - infinicclReduceOp_t op, - infinicclComm_t communicator) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(ReduceScatterV, output, input, std::move(split_counts), op, communicator); -} - -Tensor reduce_scatter(const Tensor &input, size_t world_size, infinicclReduceOp_t op, infinicclComm_t communicator) { - INFINICORE_ASSERT(input->ndim() > 0); - INFINICORE_ASSERT(input->shape()[0] % world_size == 0); - auto shape = input->shape(); - shape[0] /= world_size; - auto output = Tensor::empty(shape, input->dtype(), input->device()); - reduce_scatter_(output, input, op, communicator); - return output; -} - -void reduce_scatter_(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator) { - ReduceScatter::execute(output, input, op, communicator); -} - -Tensor reduce_scatterv(const Tensor &input, - const std::vector &split_sizes, - size_t rank, - infinicclReduceOp_t op, - infinicclComm_t communicator) { - INFINICORE_ASSERT(input->ndim() > 0); - INFINICORE_ASSERT(rank < split_sizes.size()); - auto shape = input->shape(); - shape[0] = split_sizes[rank]; - auto output = Tensor::empty(shape, input->dtype(), input->device()); - reduce_scatterv_(output, input, split_sizes, op, communicator); - return output; -} - -void reduce_scatterv_(Tensor output, - const Tensor &input, - const std::vector &split_sizes, - infinicclReduceOp_t op, - infinicclComm_t communicator) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - INFINICORE_ASSERT(output->is_contiguous() && input->is_contiguous()); - INFINICORE_ASSERT(input->ndim() > 0); - INFINICORE_ASSERT(input->shape()[0] > 0); - if (all_equal_to(split_sizes, output->shape()[0])) { - INFINICORE_ASSERT(input->numel() == output->numel() * split_sizes.size()); - ReduceScatter::execute(output, input, op, communicator); - return; - } - - const size_t inner = input->numel() / input->shape()[0]; - std::vector split_counts; - split_counts.reserve(split_sizes.size()); - size_t total_dim0 = 0; - for (auto size : split_sizes) { - split_counts.push_back(size * inner); - total_dim0 += size; - } - INFINICORE_ASSERT(input->numel() == total_dim0 * inner); - ReduceScatterV::execute(output, input, std::move(split_counts), op, communicator); -} - -std::vector reduce_scatterv_many(const std::vector &inputs, - const std::vector &split_sizes, - size_t rank, - infinicclReduceOp_t op, - infinicclComm_t communicator) { - INFINICORE_ASSERT(rank < split_sizes.size()); - std::vector outputs; - outputs.reserve(inputs.size()); - for (const auto &input : inputs) { - INFINICORE_ASSERT(input->ndim() > 0); - auto shape = input->shape(); - shape[0] = split_sizes[rank]; - outputs.push_back(Tensor::empty(shape, input->dtype(), input->device())); - } - reduce_scatterv_many_(outputs, inputs, split_sizes, op, communicator); - return outputs; -} - -void reduce_scatterv_many_(const std::vector &outputs, - const std::vector &inputs, - const std::vector &split_sizes, - infinicclReduceOp_t op, - infinicclComm_t communicator) { - INFINICORE_ASSERT(outputs.size() == inputs.size()); - if (inputs.empty()) { - return; - } - - const bool use_group = can_use_grouped_reduce_scatter(outputs, split_sizes) && !infinicore::context::isGraphRecording(); - if (use_group) { - INFINICORE_CHECK_ERROR(infinicclGroupStart(communicator)); - } - for (size_t i = 0; i < inputs.size(); ++i) { - reduce_scatterv_(outputs[i], inputs[i], split_sizes, op, communicator); - } - if (use_group) { - INFINICORE_CHECK_ERROR(infinicclGroupEnd(communicator)); - } -} - -} // namespace infinicore::op::distributed diff --git a/src/infinicore/ops/embedding/embedding.cc b/src/infinicore/ops/embedding/embedding.cc deleted file mode 100644 index 4d4da708d..000000000 --- a/src/infinicore/ops/embedding/embedding.cc +++ /dev/null @@ -1,37 +0,0 @@ -#include "infinicore/ops/embedding.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Embedding); - -Embedding::Embedding(Tensor out, const Tensor &input, const Tensor &weight) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, input, weight); -} - -void Embedding::execute(Tensor out, const Tensor &input, const Tensor &weight) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Embedding, out, input, weight); -} - -Tensor embedding(const Tensor &input, // LongTensor of arbitrary shape containing the indices to extract - const Tensor &weight // Weight: Embedding matrix of floating point type with shape (V, embedding_dim), where V = maximum index + 1 -) { - auto input_shape = input->shape(); - auto weight_shape = weight->shape(); - auto embedding_dim = weight_shape[1]; - - // Assign memory to out variables - auto output_shape = input_shape; - output_shape.push_back(embedding_dim); - Tensor inputs_embeds = Tensor::empty(output_shape, weight->dtype(), weight->device()); - - embedding_(inputs_embeds, input, weight); - return inputs_embeds; -} - -void embedding_(Tensor out, const Tensor &input, const Tensor &weight) { - Embedding::execute(out, input, weight); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/embedding/embedding_infiniop.cc b/src/infinicore/ops/embedding/embedding_infiniop.cc deleted file mode 100644 index a9780d3ae..000000000 --- a/src/infinicore/ops/embedding/embedding_infiniop.cc +++ /dev/null @@ -1,44 +0,0 @@ -#include "../infiniop_impl.hpp" -#include "infinicore/ops/embedding.hpp" - -namespace infinicore::op::embedding_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Embedding, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor out, input, weight; -}; - -void *plan(Tensor out, const Tensor &input, const Tensor &weight) { - size_t seed = hash_combine(out, input, weight); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Embedding, - seed, out->desc(), input->desc(), weight->desc()); - - auto planned = new PlannedMeta{ - descriptor, - graph::GraphTensor(out), - graph::GraphTensor(input), - graph::GraphTensor(weight)}; - - return planned; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopEmbedding( - planned->descriptor->desc, - planned->out->data(), planned->input->data(), planned->weight->data(), context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Embedding, &plan, &run, cleanup); - -} // namespace infinicore::op::embedding_impl::infiniop diff --git a/src/infinicore/ops/embedding/embedding_infiniops.cc b/src/infinicore/ops/embedding/embedding_infiniops.cc deleted file mode 100644 index 330acb5fb..000000000 --- a/src/infinicore/ops/embedding/embedding_infiniops.cc +++ /dev/null @@ -1,61 +0,0 @@ -#include "infinicore/ops/embedding.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/embedding.h" - -namespace infinicore::op::embedding_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -struct PlannedMeta { - TensorMeta out, input, weight; - graph::GraphTensor out_tensor, input_tensor, weight_tensor; -}; - -} // namespace - -void *plan(Tensor out, const Tensor &input, const Tensor &weight) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(out->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight); - - return new PlannedMeta{ - TensorMeta(out), - TensorMeta(input), - TensorMeta(weight), - graph::GraphTensor(out), - graph::GraphTensor(input), - graph::GraphTensor(weight)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - infini::ops::Embedding::Call( - handle, - config, - planned->input.tensor(planned->input_tensor), - planned->weight.tensor(planned->weight_tensor), - planned->out.tensor(planned->out_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(Embedding::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(Embedding::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(Embedding::cleanup_dispatcher(), &cleanup); - return true; -}(); - -} // namespace infinicore::op::embedding_impl::infiniops -#endif diff --git a/src/infinicore/ops/equal/equal.cc b/src/infinicore/ops/equal/equal.cc deleted file mode 100644 index b6acc4d25..000000000 --- a/src/infinicore/ops/equal/equal.cc +++ /dev/null @@ -1,31 +0,0 @@ -#include "infinicore/ops/equal.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Equal::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Equal::execute(Tensor out, Tensor a, Tensor b) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, a, b); - infinicore::context::setDevice(out->device()); - dispatcher().lookup(out->device().getType())(out, a, b); -} - -Tensor equal(Tensor a, Tensor b) { - auto out = Tensor::empty(a->shape(), DataType::BOOL, a->device()); - equal_(out, a, b); - return out; -} - -void equal_(Tensor out, Tensor a, Tensor b) { - if (out->dtype() != DataType::BOOL) { - throw std::runtime_error("Equal expects bool output tensor."); - } - Equal::execute(out, a, b); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/equal/equal_infiniop.cc b/src/infinicore/ops/equal/equal_infiniop.cc deleted file mode 100644 index 1b4e4cffa..000000000 --- a/src/infinicore/ops/equal/equal_infiniop.cc +++ /dev/null @@ -1,57 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/equal.hpp" -#include - -namespace infinicore::op::equal_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopEqualDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyEqualDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor out, Tensor a, Tensor b) { - size_t seed = hash_combine(out, a, b); - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - infiniopEqualDescriptor_t desc = nullptr; - if (auto cached = cache.get(seed)) { - desc = *cached; - } else { - INFINICORE_CHECK_ERROR(infiniopCreateEqualDescriptor( - context::getInfiniopHandle(device), &desc, - out->desc(), a->desc(), b->desc())); - cache.put(seed, desc); - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetEqualWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace; - void *workspace_ptr = nullptr; - if (workspace_size != 0) { - workspace = context::allocateMemory(workspace_size); - workspace_ptr = workspace->data(); - } - - INFINICORE_CHECK_ERROR(infiniopEqual( - desc, - workspace_ptr, - workspace_size, - out->data(), - a->data(), - b->data(), - context::getStream())); -} - -static bool registered = []() { - Equal::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::equal_impl::infiniop diff --git a/src/infinicore/ops/flash_attention/flash_attention.cc b/src/infinicore/ops/flash_attention/flash_attention.cc deleted file mode 100644 index 21cd56010..000000000 --- a/src/infinicore/ops/flash_attention/flash_attention.cc +++ /dev/null @@ -1,31 +0,0 @@ -#include "infinicore/ops/flash_attention.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(FlashAttention); - -FlashAttention::FlashAttention(Tensor out, const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k, v); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), - out, q, k, v, total_kv_len, scale, is_causal); -} - -void FlashAttention::execute(Tensor out, const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(FlashAttention, out, q, k, v, total_kv_len, scale, is_causal); -} - -Tensor flash_attention(const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal) { - Shape shape = q->shape(); - int idx = shape.size() - 1; - shape[idx] = v->shape()[idx]; - auto out = Tensor::empty(shape, q->dtype(), q->device()); - flash_attention_(out, q, k, v, total_kv_len, scale, is_causal); - return out; -} - -void flash_attention_(Tensor out, const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal) { - FlashAttention::execute(out, q, k, v, total_kv_len, scale, is_causal); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/flash_attention/flash_attention_infiniop.cc b/src/infinicore/ops/flash_attention/flash_attention_infiniop.cc deleted file mode 100644 index f5207f0ee..000000000 --- a/src/infinicore/ops/flash_attention/flash_attention_infiniop.cc +++ /dev/null @@ -1,55 +0,0 @@ -#include "../../utils.hpp" -#include "../infiniop_impl.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/flash_attention.hpp" -#include - -namespace infinicore::op::flash_attention_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, FlashAttention, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, q, k, v, total_kv_len; - float scale; - bool is_causal; -}; - -void *plan(Tensor out, const Tensor &q, const Tensor &k, const Tensor &v, const Tensor &total_kv_len, float scale, bool is_causal) { - size_t seed = hash_combine(out, q, k, v, total_kv_len, scale, is_causal); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, FlashAttention, - seed, out->desc(), q->desc(), k->desc(), v->desc(), total_kv_len->desc(), scale, is_causal); - - INFINIOP_WORKSPACE_TENSOR(workspace, FlashAttention, descriptor); - - auto planned = new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(q), - graph::GraphTensor(k), - graph::GraphTensor(v), - graph::GraphTensor(total_kv_len), scale, is_causal}; - - return planned; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopFlashAttention( - planned->descriptor->desc, planned->workspace->data(), planned->workspace->numel(), - planned->out->data(), planned->q->data(), planned->k->data(), planned->v->data(), planned->total_kv_len->data(), context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(FlashAttention, &plan, &run, &cleanup); - -} // namespace infinicore::op::flash_attention_impl::infiniop diff --git a/src/infinicore/ops/flipud/flipud.cc b/src/infinicore/ops/flipud/flipud.cc deleted file mode 100644 index 8a4178641..000000000 --- a/src/infinicore/ops/flipud/flipud.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/flipud.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &Flipud::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -// 2. 静态执行函数 -void Flipud::execute(Tensor output, Tensor input) { - dispatcher().lookup(context::getDevice().getType())(output, input); -} -Tensor flipud(Tensor input) { - // Flipud 操作不改变张量的形状和数据类型 - // Output shape == Input shape - auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); - - flipud_(output, input); - return output; -} -void flipud_(Tensor output, Tensor input) { - Flipud::execute(output, input); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/flipud/flipud_infiniop.cc b/src/infinicore/ops/flipud/flipud_infiniop.cc deleted file mode 100644 index 9162fef6b..000000000 --- a/src/infinicore/ops/flipud/flipud_infiniop.cc +++ /dev/null @@ -1,60 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/flipud.hpp" -#include - -namespace infinicore::op::flipud_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopFlipudDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyFlipudDescriptor(desc)); - desc = nullptr; - } - }); - -// 执行函数 -void calculate(Tensor output, Tensor input) { - // 1. 计算缓存 Key - size_t seed = hash_combine(output, input); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopFlipudDescriptor_t desc = nullptr; - - // 2. 获取或创建描述符 - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateFlipudDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc())); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetFlipudWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - INFINICORE_CHECK_ERROR(infiniopFlipud( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -static bool registered = []() { - Flipud::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::flipud_impl::infiniop diff --git a/src/infinicore/ops/float_power/float_power.cc b/src/infinicore/ops/float_power/float_power.cc deleted file mode 100644 index c3bf5003b..000000000 --- a/src/infinicore/ops/float_power/float_power.cc +++ /dev/null @@ -1,71 +0,0 @@ -#include "infinicore/ops/float_power.hpp" -#include "infinicore/tensor.hpp" - -namespace infinicore::op { - -// ======================================================================= -// 1. Dispatcher 单例 -// ======================================================================= - -common::OpDispatcher &FloatPower::dispatcher_scalar() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -common::OpDispatcher &FloatPower::dispatcher_tensor() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -// ======================================================================= -// 2. Execute (执行入口) -// ======================================================================= - -void FloatPower::execute(Tensor output, Tensor input, double exponent) { - dispatcher_scalar() - .lookup(context::getDevice().getType())(output, input, exponent); -} - -void FloatPower::execute(Tensor output, Tensor input, Tensor exponent) { - dispatcher_tensor() - .lookup(context::getDevice().getType())(output, input, exponent); -} - -// ======================================================================= -// 3. Functional interface (out-of-place) -> 强制提升为 F64 -// ======================================================================= - -Tensor float_power(Tensor input, double exponent) { - auto output = Tensor::empty( - input->shape(), - infinicore::DataType::F64, - input->device()); - - float_power_(output, input, exponent); - return output; -} - -Tensor float_power(Tensor input, Tensor exponent) { - Shape output_shape = input->shape(); - auto output = Tensor::empty( - output_shape, - infinicore::DataType::F64, - input->device()); - - float_power_(output, input, exponent); - return output; -} - -// ======================================================================= -// 4. Explicit / in-place -// ======================================================================= - -void float_power_(Tensor output, Tensor input, double exponent) { - FloatPower::execute(output, input, exponent); -} - -void float_power_(Tensor output, Tensor input, Tensor exponent) { - FloatPower::execute(output, input, exponent); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/float_power/float_power_infiniop.cc b/src/infinicore/ops/float_power/float_power_infiniop.cc deleted file mode 100644 index e07c1fd02..000000000 --- a/src/infinicore/ops/float_power/float_power_infiniop.cc +++ /dev/null @@ -1,129 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/float_power.hpp" -#include - -namespace infinicore::op::float_power_impl::infiniop { - -// ======================================================================= -// Descriptor Cache -// ======================================================================= - -thread_local common::OpCache caches( - 100, - [](infiniopFloatPowerDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR( - infiniopDestroyFloatPowerDescriptor(desc)); - desc = nullptr; - } - }); - -// ======================================================================= -// 1. Scalar Exponent -// ======================================================================= - -void calculate_scalar(Tensor output, - Tensor input, - double exponent) { - // Hash: output / input meta + double exponent - size_t seed = hash_combine(output, input, exponent); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopFloatPowerDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR( - infiniopCreateFloatPowerDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc(), - nullptr, // exponent tensor descriptor = null - static_cast(exponent))); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR( - infiniopGetFloatPowerWorkspaceSize(desc, &workspace_size)); - - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR( - infiniopFloatPower( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - nullptr, // exponent data pointer = null - context::getStream())); -} - -// ======================================================================= -// 2. Tensor Exponent -// ======================================================================= - -void calculate_tensor(Tensor output, - Tensor input, - Tensor exponent) { - size_t seed = hash_combine(output, input, exponent); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopFloatPowerDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR( - infiniopCreateFloatPowerDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc(), - exponent->desc(), // tensor exponent - 0.0f // scalar ignored - )); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR( - infiniopGetFloatPowerWorkspaceSize(desc, &workspace_size)); - - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR( - infiniopFloatPower( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - exponent->data(), - context::getStream())); -} - -// ======================================================================= -// 3. Dispatcher Registration -// ======================================================================= - -static bool registered = []() { - FloatPower::dispatcher_scalar().registerAll(&calculate_scalar, false); - FloatPower::dispatcher_tensor().registerAll(&calculate_tensor, false); - return true; -}(); - -} // namespace infinicore::op::float_power_impl::infiniop diff --git a/src/infinicore/ops/floor/floor.cc b/src/infinicore/ops/floor/floor.cc deleted file mode 100644 index ac125d8f4..000000000 --- a/src/infinicore/ops/floor/floor.cc +++ /dev/null @@ -1,22 +0,0 @@ -#include "infinicore/ops/floor.hpp" - -namespace infinicore::op { -common::OpDispatcher &Floor::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; -void Floor::execute(Tensor output, Tensor input) { - dispatcher().lookup(context::getDevice().getType())(output, input); -} - -Tensor floor(Tensor input) { - - auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); - floor_(output, input); - return output; -} -void floor_(Tensor output, Tensor input) { - Floor::execute(output, input); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/floor/floor_infiniop.cc b/src/infinicore/ops/floor/floor_infiniop.cc deleted file mode 100644 index 2a20f6723..000000000 --- a/src/infinicore/ops/floor/floor_infiniop.cc +++ /dev/null @@ -1,54 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/floor.hpp" -#include - -namespace infinicore::op::floor_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopFloorDescriptor_t &desc) { - if (desc != nullptr) { - // 销毁 Floor 描述符 - INFINICORE_CHECK_ERROR(infiniopDestroyFloorDescriptor(desc)); - desc = nullptr; - } - }); - -// 计算函数实现 -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopFloorDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateFloorDescriptor( - context::getInfiniopHandle(output->device()), &desc, - output->desc(), input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetFloorWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - INFINICORE_CHECK_ERROR(infiniopFloor( - desc, - workspace->data(), workspace_size, - output->data(), input->data(), // 参数顺序通常是 Output, Input - context::getStream())); -} - -static bool registered = []() { - Floor::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::floor_impl::infiniop diff --git a/src/infinicore/ops/floor_divide/floor_divide.cc b/src/infinicore/ops/floor_divide/floor_divide.cc deleted file mode 100644 index 3dc87f3f2..000000000 --- a/src/infinicore/ops/floor_divide/floor_divide.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/floor_divide.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &FloorDivide::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void FloorDivide::execute(Tensor c, Tensor a, Tensor b) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - infinicore::context::setDevice(c->device()); - dispatcher().lookup(c->device().getType())(c, a, b); -} - -Tensor floor_divide(Tensor a, Tensor b) { - auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); - floor_divide_(c, a, b); - return c; -} - -void floor_divide_(Tensor c, Tensor a, Tensor b) { - FloorDivide::execute(c, a, b); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/floor_divide/floor_divide_infiniop.cc b/src/infinicore/ops/floor_divide/floor_divide_infiniop.cc deleted file mode 100644 index 332552b5f..000000000 --- a/src/infinicore/ops/floor_divide/floor_divide_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/floor_divide.hpp" -#include - -namespace infinicore::op::floor_divide_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopFloorDivideDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyFloorDivideDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor c, Tensor a, Tensor b) { - size_t seed = hash_combine(c, b, a); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopFloorDivideDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateFloorDivideDescriptor( - context::getInfiniopHandle(c->device()), &desc, - c->desc(), a->desc(), b->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetFloorDivideWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopFloorDivide( - desc, workspace->data(), workspace_size, - c->data(), a->data(), b->data(), context::getStream())); -} - -static bool registered = []() { - FloorDivide::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::floor_divide_impl::infiniop diff --git a/src/infinicore/ops/fmin/fmin.cc b/src/infinicore/ops/fmin/fmin.cc deleted file mode 100644 index ce6a92bb8..000000000 --- a/src/infinicore/ops/fmin/fmin.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/fmin.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Fmin::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -void Fmin::execute(Tensor c, Tensor a, Tensor b) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - infinicore::context::setDevice(c->device()); - dispatcher().lookup(c->device().getType())(c, a, b); -} - -Tensor fmin(Tensor a, Tensor b) { - auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); - fmin_(c, a, b); - return c; -} - -void fmin_(Tensor c, Tensor a, Tensor b) { - Fmin::execute(c, a, b); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/fmin/fmin_infiniop.cc b/src/infinicore/ops/fmin/fmin_infiniop.cc deleted file mode 100644 index c9560479a..000000000 --- a/src/infinicore/ops/fmin/fmin_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/fmin.hpp" -#include - -namespace infinicore::op::fmin_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopFminDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyFminDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor c, Tensor a, Tensor b) { - size_t seed = hash_combine(c, b, a); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopFminDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateFminDescriptor( - context::getInfiniopHandle(c->device()), &desc, - c->desc(), a->desc(), b->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetFminWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopFmin( - desc, workspace->data(), workspace_size, - c->data(), a->data(), b->data(), context::getStream())); -} - -static bool registered = []() { - Fmin::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::fmin_impl::infiniop diff --git a/src/infinicore/ops/fmod/fmod.cc b/src/infinicore/ops/fmod/fmod.cc deleted file mode 100644 index 30bee17d6..000000000 --- a/src/infinicore/ops/fmod/fmod.cc +++ /dev/null @@ -1,28 +0,0 @@ -#include "infinicore/ops/fmod.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Fmod::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Fmod::execute(Tensor c, Tensor a, Tensor b) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - infinicore::context::setDevice(c->device()); - dispatcher().lookup(c->device().getType())(c, a, b); -} - -Tensor fmod(Tensor a, Tensor b) { - auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); - fmod_(c, a, b); - return c; -} - -void fmod_(Tensor c, Tensor a, Tensor b) { - Fmod::execute(c, a, b); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/fmod/fmod_infiniop.cc b/src/infinicore/ops/fmod/fmod_infiniop.cc deleted file mode 100644 index c085ab695..000000000 --- a/src/infinicore/ops/fmod/fmod_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/fmod.hpp" -#include - -namespace infinicore::op::fmod_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopFmodDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyFmodDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor c, Tensor a, Tensor b) { - size_t seed = hash_combine(c, b, a); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopFmodDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateFmodDescriptor( - context::getInfiniopHandle(c->device()), &desc, - c->desc(), a->desc(), b->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetFmodWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopFmod( - desc, workspace->data(), workspace_size, - c->data(), a->data(), b->data(), context::getStream())); -} - -static bool registered = []() { - Fmod::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::fmod_impl::infiniop diff --git a/src/infinicore/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating.cc b/src/infinicore/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating.cc deleted file mode 100644 index 6ce94925b..000000000 --- a/src/infinicore/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating.cc +++ /dev/null @@ -1,85 +0,0 @@ -#include "infinicore/ops/fused_gated_delta_net_gating.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(FusedGatedDeltaNetGating); - -FusedGatedDeltaNetGating::FusedGatedDeltaNetGating(Tensor g, - Tensor beta_output, - const Tensor &A_log, - const Tensor &a, - const Tensor &b, - const Tensor &dt_bias, - float beta, - float threshold) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(g, beta_output, A_log, a, b, dt_bias); - INFINICORE_GRAPH_OP_DISPATCH(g->device().getType(), g, beta_output, A_log, a, b, dt_bias, beta, threshold); -} - -void FusedGatedDeltaNetGating::execute(Tensor g, - Tensor beta_output, - const Tensor &A_log, - const Tensor &a, - const Tensor &b, - const Tensor &dt_bias, - float beta, - float threshold) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(FusedGatedDeltaNetGating, g, beta_output, A_log, a, b, dt_bias, beta, threshold); -} - -static void validate_inputs(const Tensor &A_log, - const Tensor &a, - const Tensor &b, - const Tensor &dt_bias) { - if (a->shape().size() != 3 || b->shape().size() != 3) { - throw std::runtime_error("fused_gated_delta_net_gating expects a and b with shape [batch_size, seq_len, hidden]"); - } - if (a->shape() != b->shape()) { - throw std::runtime_error("fused_gated_delta_net_gating expects a and b to have the same shape"); - } - if (A_log->shape().size() != 1 || dt_bias->shape().size() != 1) { - throw std::runtime_error("fused_gated_delta_net_gating expects A_log and dt_bias with shape [hidden]"); - } - if (A_log->shape()[0] != a->shape()[2] || dt_bias->shape()[0] != a->shape()[2]) { - throw std::runtime_error("fused_gated_delta_net_gating hidden dimension mismatch"); - } -} - -std::pair fused_gated_delta_net_gating(const Tensor &A_log, - const Tensor &a, - const Tensor &b, - const Tensor &dt_bias, - float beta, - float threshold) { - validate_inputs(A_log, a, b, dt_bias); - - Tensor g = Tensor::empty(a->shape(), DataType::F32, a->device()); - Tensor beta_output = Tensor::empty(a->shape(), DataType::F32, a->device()); - fused_gated_delta_net_gating_(g, beta_output, A_log, a, b, dt_bias, beta, threshold); - return {g, beta_output}; -} - -void fused_gated_delta_net_gating_(Tensor g, - Tensor beta_output, - const Tensor &A_log, - const Tensor &a, - const Tensor &b, - const Tensor &dt_bias, - float beta, - float threshold) { - validate_inputs(A_log, a, b, dt_bias); - if (g->shape() != a->shape() || beta_output->shape() != a->shape()) { - throw std::runtime_error("fused_gated_delta_net_gating_ expects outputs with shape [batch_size, seq_len, hidden]"); - } - if (g->dtype() != DataType::F32 || beta_output->dtype() != DataType::F32) { - throw std::runtime_error("fused_gated_delta_net_gating_ expects float32 outputs"); - } - - FusedGatedDeltaNetGating::execute(g, beta_output, A_log, a, b, dt_bias, beta, threshold); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating_infiniop.cc b/src/infinicore/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating_infiniop.cc deleted file mode 100644 index 66aeeaa18..000000000 --- a/src/infinicore/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating_infiniop.cc +++ /dev/null @@ -1,71 +0,0 @@ -#include "infinicore/ops/fused_gated_delta_net_gating.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::fused_gated_delta_net_gating_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, FusedGatedDeltaNetGating, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, g, beta_output, A_log, a, b, dt_bias; -}; - -void *plan(Tensor g, - Tensor beta_output, - const Tensor &A_log, - const Tensor &a, - const Tensor &b, - const Tensor &dt_bias, - float beta, - float threshold) { - size_t seed = hash_combine(g, beta_output, A_log, a, b, dt_bias, beta, threshold); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, FusedGatedDeltaNetGating, seed, - g->desc(), - beta_output->desc(), - A_log->desc(), - a->desc(), - b->desc(), - dt_bias->desc(), - beta, - threshold); - - INFINIOP_WORKSPACE_TENSOR(workspace, FusedGatedDeltaNetGating, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(g), - graph::GraphTensor(beta_output), - graph::GraphTensor(A_log), - graph::GraphTensor(a), - graph::GraphTensor(b), - graph::GraphTensor(dt_bias)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopFusedGatedDeltaNetGating( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->g->data(), - planned->beta_output->data(), - planned->A_log->data(), - planned->a->data(), - planned->b->data(), - planned->dt_bias->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(FusedGatedDeltaNetGating, &plan, &run, &cleanup); - -} // namespace infinicore::op::fused_gated_delta_net_gating_impl::infiniop diff --git a/src/infinicore/ops/gaussian_nll_loss/gaussian_nll_loss.cc b/src/infinicore/ops/gaussian_nll_loss/gaussian_nll_loss.cc deleted file mode 100644 index 28ff2e686..000000000 --- a/src/infinicore/ops/gaussian_nll_loss/gaussian_nll_loss.cc +++ /dev/null @@ -1,51 +0,0 @@ -#include "infinicore/ops/gaussian_nll_loss.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(GaussianNllLoss); - -GaussianNllLoss::GaussianNllLoss(Tensor out, - const Tensor &input, - const Tensor &target, - const Tensor &var, - bool full, - double eps, - int reduction) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, target, var); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, input, target, var, full, eps, reduction); -} - -void GaussianNllLoss::execute(Tensor out, - const Tensor &input, - const Tensor &target, - const Tensor &var, - bool full, - double eps, - int reduction) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(GaussianNllLoss, out, input, target, var, full, eps, reduction); -} - -Tensor gaussian_nll_loss(const Tensor &input, - const Tensor &target, - const Tensor &var, - bool full, - double eps, - int reduction) { - std::vector out_shape = (reduction == 0) ? input->shape() : std::vector{}; - auto out = Tensor::empty(out_shape, input->dtype(), input->device()); - gaussian_nll_loss_(out, input, target, var, full, eps, reduction); - return out; -} - -void gaussian_nll_loss_(Tensor out, - const Tensor &input, - const Tensor &target, - const Tensor &var, - bool full, - double eps, - int reduction) { - GaussianNllLoss::execute(out, input, target, var, full, eps, reduction); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/gaussian_nll_loss/gaussian_nll_loss_infiniop.cc b/src/infinicore/ops/gaussian_nll_loss/gaussian_nll_loss_infiniop.cc deleted file mode 100644 index 1d2bc832d..000000000 --- a/src/infinicore/ops/gaussian_nll_loss/gaussian_nll_loss_infiniop.cc +++ /dev/null @@ -1,69 +0,0 @@ -#include "infinicore/ops/gaussian_nll_loss.hpp" - -#include "infiniop/ops/gaussian_nll_loss.h" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::gaussian_nll_loss_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, GaussianNllLoss, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, input, target, var; -}; - -void *plan(Tensor out, - const Tensor &input, - const Tensor &target, - const Tensor &var, - bool full, - double eps, - int reduction) { - const int full_i = full ? 1 : 0; - size_t seed = hash_combine(out, input, target, var, full_i, eps, reduction); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, GaussianNllLoss, - seed, - out->desc(), - input->desc(), - target->desc(), - var->desc(), - full_i, - eps, - reduction); - - INFINIOP_WORKSPACE_TENSOR(workspace, GaussianNllLoss, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(input), - graph::GraphTensor(target), - graph::GraphTensor(var)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopGaussianNllLoss( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->input->data(), - planned->target->data(), - planned->var->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(GaussianNllLoss, &plan, &run, &cleanup); - -} // namespace infinicore::op::gaussian_nll_loss_impl::infiniop diff --git a/src/infinicore/ops/gelu/gelu.cc b/src/infinicore/ops/gelu/gelu.cc deleted file mode 100644 index 612a2ecee..000000000 --- a/src/infinicore/ops/gelu/gelu.cc +++ /dev/null @@ -1,37 +0,0 @@ -#include "infinicore/ops/gelu.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &Gelu::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Gelu::execute(Tensor output, Tensor input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Gelu implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor gelu(Tensor input) { - Shape shape = input->shape(); - auto output = Tensor::empty(shape, input->dtype(), input->device()); - gelu_(output, input); - return output; -} - -void gelu_(Tensor output, Tensor input) { - Gelu::execute(output, input); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/gelu/gelu_infiniop.cc b/src/infinicore/ops/gelu/gelu_infiniop.cc deleted file mode 100644 index 6294a05c2..000000000 --- a/src/infinicore/ops/gelu/gelu_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/gelu.hpp" -#include - -namespace infinicore::op::gelu_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopGeluDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyGeluDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopGeluDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateGeluDescriptor( - context::getInfiniopHandle(device), &desc, - output->desc(), input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetGeluWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopGelu( - desc, workspace->data(), workspace_size, - output->data(), input->data(), context::getStream())); -} - -static bool registered = []() { - Gelu::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::gelu_impl::infiniop diff --git a/src/infinicore/ops/gelu/gelu_infiniops.cc b/src/infinicore/ops/gelu/gelu_infiniops.cc deleted file mode 100644 index de0b35fde..000000000 --- a/src/infinicore/ops/gelu/gelu_infiniops.cc +++ /dev/null @@ -1,41 +0,0 @@ -#include "infinicore/ops/gelu.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/gelu_infinilm.h" - -#include - -namespace infinicore::op::gelu_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -void calculate(Tensor output, Tensor input) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - TensorMeta output_meta(output); - TensorMeta input_meta(input); - infini::ops::GeluInfinilm::Call( - handle, - config, - input_meta.tensor(input), - std::string{"none"}, - output_meta.tensor(output)); -} - -} // namespace - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(Gelu::dispatcher(), &calculate); - return true; -}(); - -} // namespace infinicore::op::gelu_impl::infiniops -#endif diff --git a/src/infinicore/ops/gelutanh/gelutanh.cc b/src/infinicore/ops/gelutanh/gelutanh.cc deleted file mode 100644 index b6bae39dd..000000000 --- a/src/infinicore/ops/gelutanh/gelutanh.cc +++ /dev/null @@ -1,37 +0,0 @@ -#include "infinicore/ops/gelutanh.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &GeluTanh::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void GeluTanh::execute(Tensor output, Tensor input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No GeluTanh implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor gelu_tanh(Tensor input) { - Shape shape = input->shape(); - auto output = Tensor::empty(shape, input->dtype(), input->device()); - gelu_tanh_(output, input); - return output; -} - -void gelu_tanh_(Tensor output, Tensor input) { - GeluTanh::execute(output, input); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/gelutanh/gelutanh_infiniop.cc b/src/infinicore/ops/gelutanh/gelutanh_infiniop.cc deleted file mode 100644 index fb13ca98a..000000000 --- a/src/infinicore/ops/gelutanh/gelutanh_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/gelutanh.hpp" -#include - -namespace infinicore::op::gelutanh_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopGeluTanhDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyGeluTanhDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopGeluTanhDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateGeluTanhDescriptor( - context::getInfiniopHandle(device), &desc, - output->desc(), input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetGeluTanhWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopGeluTanh( - desc, workspace->data(), workspace_size, - output->data(), input->data(), context::getStream())); -} - -static bool registered = []() { - GeluTanh::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::gelutanh_impl::infiniop diff --git a/src/infinicore/ops/gelutanh/gelutanh_infiniops.cc b/src/infinicore/ops/gelutanh/gelutanh_infiniops.cc deleted file mode 100644 index e9faf4b35..000000000 --- a/src/infinicore/ops/gelutanh/gelutanh_infiniops.cc +++ /dev/null @@ -1,38 +0,0 @@ -#include "infinicore/ops/gelutanh.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/gelutanh_infinilm.h" - -namespace infinicore::op::gelutanh_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -void calculate(Tensor output, Tensor input) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - TensorMeta output_meta(output); - TensorMeta input_meta(input); - infini::ops::GelutanhInfinilm::Call( - handle, - config, - input_meta.tensor(input), - output_meta.tensor(output)); -} - -} // namespace - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(GeluTanh::dispatcher(), &calculate); - return true; -}(); - -} // namespace infinicore::op::gelutanh_impl::infiniops -#endif diff --git a/src/infinicore/ops/gemm/gemm.cc b/src/infinicore/ops/gemm/gemm.cc deleted file mode 100644 index 765bc869f..000000000 --- a/src/infinicore/ops/gemm/gemm.cc +++ /dev/null @@ -1,30 +0,0 @@ -#include "infinicore/ops/gemm.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Gemm); - -Gemm::Gemm(Tensor c, const Tensor &a, const Tensor &b, float alpha, float beta) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - INFINICORE_GRAPH_OP_DISPATCH(c->device().getType(), c, a, b, alpha, beta); -} - -void Gemm::execute(Tensor c, const Tensor &a, const Tensor &b, float alpha, float beta) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Gemm, c, a, b, alpha, beta); -} - -Tensor gemm(const Tensor &a, const Tensor &b, float alpha, float beta) { - Shape shape = a->shape(); - Size size = a->ndim(); - shape[size - 1] = b->size(size - 1); - auto c = Tensor::empty(shape, a->dtype(), a->device()); - gemm_(c, a, b, alpha, beta); - return c; -} - -void gemm_(Tensor c, const Tensor &a, const Tensor &b, float alpha, float beta) { - Gemm::execute(c, a, b, alpha, beta); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/gemm/gemm_infiniop.cc b/src/infinicore/ops/gemm/gemm_infiniop.cc deleted file mode 100644 index 33a7271c0..000000000 --- a/src/infinicore/ops/gemm/gemm_infiniop.cc +++ /dev/null @@ -1,49 +0,0 @@ -#include "../infiniop_impl.hpp" -#include "infinicore/ops/gemm.hpp" - -namespace infinicore::op::gemm_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Gemm, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, c, a, b; - float alpha, beta; -}; - -void *plan(Tensor c, const Tensor &a, const Tensor &b, float alpha, float beta) { - size_t seed = hash_combine(c, a, b); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Gemm, - seed, c->desc(), a->desc(), b->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Gemm, descriptor); - - auto planned = new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(c), - graph::GraphTensor(a), - graph::GraphTensor(b), - alpha, beta}; - - return planned; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopGemm( - planned->descriptor->desc, planned->workspace->data(), planned->workspace->numel(), - planned->c->data(), planned->a->data(), planned->b->data(), planned->alpha, planned->beta, context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Gemm, &plan, &run, &cleanup); - -} // namespace infinicore::op::gemm_impl::infiniop diff --git a/src/infinicore/ops/gemm/gemm_infiniops.cc b/src/infinicore/ops/gemm/gemm_infiniops.cc deleted file mode 100644 index f45dbec1f..000000000 --- a/src/infinicore/ops/gemm/gemm_infiniops.cc +++ /dev/null @@ -1,68 +0,0 @@ -#include "infinicore/ops/gemm.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include - -namespace infinicore::op::gemm_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -struct PlannedMeta { - TensorMeta c, a, b; - graph::GraphTensor c_tensor, a_tensor, b_tensor; - float alpha, beta; -}; - -} // namespace - -void *plan(Tensor c, const Tensor &a, const Tensor &b, float alpha, float beta) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(c->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - - return new PlannedMeta{ - TensorMeta(c), - TensorMeta(a), - TensorMeta(b), - graph::GraphTensor(c), - graph::GraphTensor(a), - graph::GraphTensor(b), - alpha, - beta}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - infini::ops::Gemm::Call( - handle, - config, - planned->a.tensor(planned->a_tensor), - planned->b.tensor(planned->b_tensor), - std::optional{planned->alpha}, - std::optional{planned->beta}, - std::optional{}, - std::optional{}, - planned->c.tensor(planned->c_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(Gemm::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(Gemm::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(Gemm::cleanup_dispatcher(), &cleanup); - return true; -}(); - -} // namespace infinicore::op::gemm_impl::infiniops -#endif diff --git a/src/infinicore/ops/gptq_marlin_gemm/gptq_marlin_gemm.cc b/src/infinicore/ops/gptq_marlin_gemm/gptq_marlin_gemm.cc deleted file mode 100644 index 56709b3ab..000000000 --- a/src/infinicore/ops/gptq_marlin_gemm/gptq_marlin_gemm.cc +++ /dev/null @@ -1,40 +0,0 @@ -#include "infinicore/ops/gptq_marlin_gemm.hpp" - -#include "../../utils.hpp" -#include "infinicore/context/context.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(GptqMarlinGemm); -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(GptqMarlinGemmWithWorkspace); - -GptqMarlinGemm::GptqMarlinGemm(Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, a, b, b_scales, global_scales, b_zeros, g_idx, perm); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, a, b, b_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float); -} -void GptqMarlinGemm::execute(Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(GptqMarlinGemm, out, a, b, b_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float); -} - -void gptq_marlin_gemm_(Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - GptqMarlinGemm::execute(out, a, b, b_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float); -} - -GptqMarlinGemmWithWorkspace::GptqMarlinGemmWithWorkspace(Tensor workspace, Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(workspace, out, a, b, b_scales, global_scales, b_zeros, g_idx, perm); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), workspace, out, a, b, b_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float); -} - -void GptqMarlinGemmWithWorkspace::execute(Tensor workspace, Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(GptqMarlinGemmWithWorkspace, workspace, out, a, b, b_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float); -} - -void gptq_marlin_gemm_with_workspace_(Tensor workspace, Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - if (context::isGraphRecording()) { - GptqMarlinGemmWithWorkspace::execute(workspace, out, a, b, b_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float); - } else { - gptq_marlin_gemm_with_workspace_direct_(workspace, out, a, b, b_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float); - } -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/gptq_marlin_gemm/gptq_marlin_gemm_infiniop.cc b/src/infinicore/ops/gptq_marlin_gemm/gptq_marlin_gemm_infiniop.cc deleted file mode 100644 index e0448e6b8..000000000 --- a/src/infinicore/ops/gptq_marlin_gemm/gptq_marlin_gemm_infiniop.cc +++ /dev/null @@ -1,220 +0,0 @@ -#include "../../utils.hpp" -#include "../infiniop_impl.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/gptq_marlin_gemm.hpp" -#include - -namespace infinicore::op::gptq_marlin_gemm_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, GptqMarlinGemm, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, a, b, b_scales, global_scales, b_zeros, g_idx, perm; - int64_t b_q_type_id; - bool is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float; -}; - -void *plan(Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - size_t seed = hash_combine(out, a, b, b_scales, global_scales, b_zeros, g_idx, perm); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, GptqMarlinGemm, - seed, - out->desc(), a->desc(), - b->desc(), b_scales->desc(), global_scales->desc(), b_zeros->desc(), g_idx->desc(), perm->desc()); - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetGptqMarlinGemmWorkspaceSize(descriptor->desc, &workspace_size)); - Tensor workspace = Tensor::zeros({workspace_size}, DataType::U8, context::getDevice()); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(a), - graph::GraphTensor(b), - graph::GraphTensor(b_scales), - graph::GraphTensor(global_scales), - graph::GraphTensor(b_zeros), - graph::GraphTensor(g_idx), - graph::GraphTensor(perm), - b_q_type_id, - is_k_full, - use_atomic_add, - use_fp32_reduce, - is_zp_float}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - auto optional_data = [](graph::GraphTensor &tensor) -> void * { - return tensor->numel() == 0 ? nullptr : tensor->data(); - }; - - INFINICORE_CHECK_ERROR(infiniopGptqMarlinGemm( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->a->data(), - planned->b->data(), - planned->b_scales->data(), - optional_data(planned->global_scales), - optional_data(planned->b_zeros), - optional_data(planned->g_idx), - optional_data(planned->perm), - planned->b_q_type_id, - planned->is_k_full, - planned->use_atomic_add, - planned->use_fp32_reduce, - planned->is_zp_float, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(GptqMarlinGemm, &plan, &run, &cleanup); - -} // namespace infinicore::op::gptq_marlin_gemm_impl::infiniop - -namespace infinicore::op::gptq_marlin_gemm_impl::infiniop { - -size_t workspace_size(Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm) { - size_t seed = hash_combine(out, a, b, b_scales, global_scales, b_zeros, g_idx, perm); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, GptqMarlinGemm, - seed, - out->desc(), a->desc(), - b->desc(), b_scales->desc(), global_scales->desc(), b_zeros->desc(), g_idx->desc(), perm->desc()); - - size_t size = 0; - INFINICORE_CHECK_ERROR(infiniopGetGptqMarlinGemmWorkspaceSize(descriptor->desc, &size)); - return size; -} - -struct PlannedWithWorkspaceMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, a, b, b_scales, global_scales, b_zeros, g_idx, perm; - int64_t b_q_type_id; - bool is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float; -}; - -void *plan_with_workspace(Tensor workspace, Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - size_t seed = hash_combine(out, a, b, b_scales, global_scales, b_zeros, g_idx, perm); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, GptqMarlinGemm, - seed, - out->desc(), a->desc(), - b->desc(), b_scales->desc(), global_scales->desc(), b_zeros->desc(), g_idx->desc(), perm->desc()); - - return new PlannedWithWorkspaceMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(a), - graph::GraphTensor(b), - graph::GraphTensor(b_scales), - graph::GraphTensor(global_scales), - graph::GraphTensor(b_zeros), - graph::GraphTensor(g_idx), - graph::GraphTensor(perm), - b_q_type_id, - is_k_full, - use_atomic_add, - use_fp32_reduce, - is_zp_float}; -} - -void run_with_workspace(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - auto optional_data = [](graph::GraphTensor &tensor) -> void * { - return tensor->numel() == 0 ? nullptr : tensor->data(); - }; - - context::setDeviceMemoryAsync(planned->workspace->data(), 0, planned->workspace->nbytes(), context::getStream()); - INFINICORE_CHECK_ERROR(infiniopGptqMarlinGemm( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->a->data(), - planned->b->data(), - planned->b_scales->data(), - optional_data(planned->global_scales), - optional_data(planned->b_zeros), - optional_data(planned->g_idx), - optional_data(planned->perm), - planned->b_q_type_id, - planned->is_k_full, - planned->use_atomic_add, - planned->use_fp32_reduce, - planned->is_zp_float, - context::getStream())); -} - -void cleanup_with_workspace(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered_with_workspace = []() { - GptqMarlinGemmWithWorkspace::plan_dispatcher().registerAll(&plan_with_workspace, false); - GptqMarlinGemmWithWorkspace::run_dispatcher().registerAll(&run_with_workspace, false); - GptqMarlinGemmWithWorkspace::cleanup_dispatcher().registerAll(&cleanup_with_workspace, false); - return true; -}(); - -void direct_with_workspace(Tensor workspace, Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - size_t seed = hash_combine(out, a, b, b_scales, global_scales, b_zeros, g_idx, perm); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, GptqMarlinGemm, - seed, - out->desc(), a->desc(), - b->desc(), b_scales->desc(), global_scales->desc(), b_zeros->desc(), g_idx->desc(), perm->desc()); - - auto optional_data = [](Tensor &tensor) -> void * { - return tensor->numel() == 0 ? nullptr : tensor->data(); - }; - - context::setDeviceMemoryAsync(workspace->data(), 0, workspace->nbytes(), context::getStream()); - INFINICORE_CHECK_ERROR(infiniopGptqMarlinGemm( - descriptor->desc, - workspace->data(), - workspace->numel(), - out->data(), - a->data(), - b->data(), - b_scales->data(), - optional_data(global_scales), - optional_data(b_zeros), - optional_data(g_idx), - optional_data(perm), - b_q_type_id, - is_k_full, - use_atomic_add, - use_fp32_reduce, - is_zp_float, - context::getStream())); -} - -} // namespace infinicore::op::gptq_marlin_gemm_impl::infiniop - -namespace infinicore::op { - -size_t gptq_marlin_gemm_workspace_size(Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm) { - return gptq_marlin_gemm_impl::infiniop::workspace_size(out, a, b, b_scales, global_scales, b_zeros, g_idx, perm); -} - -void gptq_marlin_gemm_with_workspace_direct_(Tensor workspace, Tensor out, const Tensor &a, const Tensor &b, Tensor &b_scales, Tensor &global_scales, Tensor &b_zeros, Tensor &g_idx, Tensor &perm, int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float) { - gptq_marlin_gemm_impl::infiniop::direct_with_workspace(workspace, out, a, b, b_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/hardswish/hardswish.cc b/src/infinicore/ops/hardswish/hardswish.cc deleted file mode 100644 index ec8db75ff..000000000 --- a/src/infinicore/ops/hardswish/hardswish.cc +++ /dev/null @@ -1,38 +0,0 @@ -#include "infinicore/ops/hardswish.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &Hardswish::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -void Hardswish::execute(Tensor output, Tensor input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error( - "No Hardswish implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor hardswish(Tensor input) { - auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); - hardswish_(output, input); - return output; -} - -void hardswish_(Tensor output, Tensor input) { - Hardswish::execute(output, input); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/hardswish/hardswish_infiniop.cc b/src/infinicore/ops/hardswish/hardswish_infiniop.cc deleted file mode 100644 index 44d4054e8..000000000 --- a/src/infinicore/ops/hardswish/hardswish_infiniop.cc +++ /dev/null @@ -1,61 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/hardswish.hpp" -#include - -namespace infinicore::op::hardswish_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopHardSwishDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyHardSwishDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopHardSwishDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateHardSwishDescriptor( - context::getInfiniopHandle(device), - &desc, - output->desc(), - input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetHardSwishWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace; - void *workspace_ptr = nullptr; - if (workspace_size != 0) { - workspace = context::allocateMemory(workspace_size); - workspace_ptr = workspace->data(); - } - - INFINICORE_CHECK_ERROR(infiniopHardSwish( - desc, - workspace_ptr, - workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -static bool registered = []() { - Hardswish::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::hardswish_impl::infiniop diff --git a/src/infinicore/ops/hardtanh/hardtanh.cc b/src/infinicore/ops/hardtanh/hardtanh.cc deleted file mode 100644 index 5a4df2142..000000000 --- a/src/infinicore/ops/hardtanh/hardtanh.cc +++ /dev/null @@ -1,38 +0,0 @@ -#include "infinicore/ops/hardtanh.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &HardTanh::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -void HardTanh::execute(Tensor output, Tensor input, float min_val, float max_val) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(output->device()); - - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - if (func == nullptr) { - throw std::runtime_error( - "No HardTanh implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input, min_val, max_val); -} - -Tensor hardtanh(Tensor input, float min_val, float max_val) { - auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); - hardtanh_(output, input, min_val, max_val); - return output; -} - -void hardtanh_(Tensor output, Tensor input, float min_val, float max_val) { - HardTanh::execute(output, input, min_val, max_val); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/hardtanh/hardtanh_infiniop.cc b/src/infinicore/ops/hardtanh/hardtanh_infiniop.cc deleted file mode 100644 index d8af439d8..000000000 --- a/src/infinicore/ops/hardtanh/hardtanh_infiniop.cc +++ /dev/null @@ -1,63 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/hardtanh.hpp" -#include - -namespace infinicore::op::hardtanh_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopHardTanhDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyHardTanhDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, float min_val, float max_val) { - size_t seed = hash_combine(output, input, min_val, max_val); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopHardTanhDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateHardTanhDescriptor( - context::getInfiniopHandle(device), - &desc, - output->desc(), - input->desc(), - min_val, - max_val)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetHardTanhWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace; - void *workspace_ptr = nullptr; - if (workspace_size != 0) { - workspace = context::allocateMemory(workspace_size); - workspace_ptr = workspace->data(); - } - - INFINICORE_CHECK_ERROR(infiniopHardTanh( - desc, - workspace_ptr, - workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -static bool registered = []() { - HardTanh::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::hardtanh_impl::infiniop diff --git a/src/infinicore/ops/hinge_embedding_loss/hinge_embedding_loss.cc b/src/infinicore/ops/hinge_embedding_loss/hinge_embedding_loss.cc deleted file mode 100644 index 5f46b98c1..000000000 --- a/src/infinicore/ops/hinge_embedding_loss/hinge_embedding_loss.cc +++ /dev/null @@ -1,57 +0,0 @@ -#include "infinicore/ops/hinge_embedding_loss.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &HingeEmbeddingLoss::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -void HingeEmbeddingLoss::execute( - Tensor output, - Tensor input, - Tensor target, - double margin, - int reduction) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input, target); - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error( - "No HingeEmbeddingLoss implementation found for device type: " - + std::to_string(static_cast(device_type))); - } - - func(output, input, target, margin, reduction); -} - -Tensor hinge_embedding_loss(Tensor input, Tensor target, double margin, int reduction) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(input, target); - INFINICORE_ASSERT(input->dtype() == target->dtype()); - - Shape output_shape; - if (reduction == 0) { - output_shape = input->shape(); - } - - auto output = Tensor::empty(output_shape, input->dtype(), input->device()); - hinge_embedding_loss_(output, input, target, margin, reduction); - return output; -} - -void hinge_embedding_loss_( - Tensor output, - Tensor input, - Tensor target, - double margin, - int reduction) { - HingeEmbeddingLoss::execute(output, input, target, margin, reduction); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/hinge_embedding_loss/hinge_embedding_loss_infiniop.cc b/src/infinicore/ops/hinge_embedding_loss/hinge_embedding_loss_infiniop.cc deleted file mode 100644 index a26615115..000000000 --- a/src/infinicore/ops/hinge_embedding_loss/hinge_embedding_loss_infiniop.cc +++ /dev/null @@ -1,61 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/hinge_embedding_loss.hpp" - -#include - -namespace infinicore::op::hinge_embedding_loss_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopHingeEmbeddingLossDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyHingeEmbeddingLossDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, Tensor target, double margin, int reduction) { - size_t seed = hash_combine(output, input, target, margin, reduction); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopHingeEmbeddingLossDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateHingeEmbeddingLossDescriptor( - context::getInfiniopHandle(device), - &desc, - output->desc(), - input->desc(), - target->desc(), - margin, - reduction)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetHingeEmbeddingLossWorkspaceSize(desc, &workspace_size)); - auto workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopHingeEmbeddingLoss( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - target->data(), - context::getStream())); -} - -static bool registered = []() { - HingeEmbeddingLoss::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::hinge_embedding_loss_impl::infiniop diff --git a/src/infinicore/ops/huber_loss/huber_loss.cc b/src/infinicore/ops/huber_loss/huber_loss.cc deleted file mode 100644 index 3542ca18c..000000000 --- a/src/infinicore/ops/huber_loss/huber_loss.cc +++ /dev/null @@ -1,35 +0,0 @@ -#include "infinicore/ops/huber_loss.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &HuberLoss::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void HuberLoss::execute(Tensor output, Tensor input, Tensor target, float delta, int64_t reduction) { - dispatcher().lookup(context::getDevice().getType())(output, input, target, delta, reduction); -} - -// 3. 函数式接口 -Tensor huber_loss(Tensor input, Tensor target, float delta, int64_t reduction) { - Shape output_shape; - if (reduction == 0) { // None - // HuberLoss 是 Element-wise 的，reduction='none' 时输出形状通常与输入一致 - output_shape = input->shape(); - } else { - output_shape = {}; // Scalar - } - - auto output = Tensor::empty(output_shape, input->dtype(), input->device()); - - huber_loss_(output, input, target, delta, reduction); - return output; -} - -void huber_loss_(Tensor output, Tensor input, Tensor target, float delta, int64_t reduction) { - HuberLoss::execute(output, input, target, delta, reduction); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/huber_loss/huber_loss_infiniop.cc b/src/infinicore/ops/huber_loss/huber_loss_infiniop.cc deleted file mode 100644 index 7a8a459ef..000000000 --- a/src/infinicore/ops/huber_loss/huber_loss_infiniop.cc +++ /dev/null @@ -1,66 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/huber_loss.hpp" -#include - -namespace infinicore::op::huber_loss_impl::infiniop { - -// 定义描述符缓存 -thread_local common::OpCache caches( - 100, // capacity - [](infiniopHuberLossDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyHuberLossDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, Tensor target, float delta, int64_t reduction) { - size_t seed = hash_combine(output, input, target, delta, reduction); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopHuberLossDescriptor_t desc = nullptr; - - if (!desc_opt) { - // 3. 创建描述符 - INFINICORE_CHECK_ERROR(infiniopCreateHuberLossDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc(), - target->desc(), - delta, - static_cast(reduction))); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 4. 获取 Workspace 并执行 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetHuberLossWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopHuberLoss( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - target->data(), - context::getStream())); -} - -static bool registered = []() { - HuberLoss::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::huber_loss_impl::infiniop diff --git a/src/infinicore/ops/hypot/hypot.cc b/src/infinicore/ops/hypot/hypot.cc deleted file mode 100644 index 27ea9f936..000000000 --- a/src/infinicore/ops/hypot/hypot.cc +++ /dev/null @@ -1,25 +0,0 @@ -#include "infinicore/ops/hypot.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &Hypot::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Hypot::execute(Tensor output, Tensor input_a, Tensor input_b) { - // lookup 需要传入设备类型，然后调用返回的函数指针 - dispatcher().lookup(context::getDevice().getType())(output, input_a, input_b); -} -Tensor hypot(Tensor input_a, Tensor input_b) { - auto output = Tensor::empty(input_a->shape(), input_a->dtype(), input_a->device()); - - hypot_(output, input_a, input_b); - return output; -} -void hypot_(Tensor output, Tensor input_a, Tensor input_b) { - Hypot::execute(output, input_a, input_b); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/hypot/hypot_infiniop.cc b/src/infinicore/ops/hypot/hypot_infiniop.cc deleted file mode 100644 index cc6eb4a99..000000000 --- a/src/infinicore/ops/hypot/hypot_infiniop.cc +++ /dev/null @@ -1,55 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/hypot.hpp" // 引入 Hypot 头文件 -#include - -namespace infinicore::op::hypot_impl::infiniop { -thread_local common::OpCache caches( - 100, // capacity - [](infiniopHypotDescriptor_t &desc) { - if (desc != nullptr) { - // 销毁 Hypot 描述符 - INFINICORE_CHECK_ERROR(infiniopDestroyHypotDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input_a, Tensor input_b) { - size_t seed = hash_combine(output, input_a, input_b); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopHypotDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateHypotDescriptor( - context::getInfiniopHandle(output->device()), &desc, - output->desc(), input_a->desc(), input_b->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetHypotWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopHypot( - desc, - workspace->data(), workspace_size, - output->data(), input_a->data(), input_b->data(), - context::getStream())); -} - -static bool registered = []() { - // 注册到 Hypot 的 dispatcher - Hypot::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::hypot_impl::infiniop diff --git a/src/infinicore/ops/index_add/index_add.cc b/src/infinicore/ops/index_add/index_add.cc deleted file mode 100644 index 4a4192b2c..000000000 --- a/src/infinicore/ops/index_add/index_add.cc +++ /dev/null @@ -1,114 +0,0 @@ -#include "infinicore/ops/index_add.hpp" -#include "infinicore/tensor.hpp" -#include -#include -#include - -namespace infinicore::op { - -common::OpDispatcher &IndexAdd::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void IndexAdd::execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source, float alpha) { - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No IndexAdd implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input, dim, index, source, alpha); -} - -static void check_index_add_args(const Tensor &input, int64_t &dim, const Tensor &index, const Tensor &source) { - int64_t ndim = static_cast(input->ndim()); - - if (dim < 0) { - dim += ndim; - } - if (dim < 0 || dim >= ndim) { - throw std::runtime_error("IndexAdd: Dimension out of range."); - } - - if (index->ndim() != 1) { - throw std::runtime_error("IndexAdd: Index tensor must be 1D."); - } - - // 使用 DataType::I64 和 I32 - if (index->dtype() != DataType::I64 && index->dtype() != DataType::I32) { - throw std::runtime_error("IndexAdd: Index tensor must be I32 or I64."); - } - - if (source->ndim() != input->ndim()) { - throw std::runtime_error("IndexAdd: Source tensor must have same number of dimensions as input tensor."); - } - - auto in_shape = input->shape(); - auto src_shape = source->shape(); - auto idx_len = index->shape()[0]; - - for (int64_t i = 0; i < ndim; ++i) { - if (i == dim) { - if (src_shape[i] != idx_len) { - throw std::runtime_error("IndexAdd: Source dimension mismatch."); - } - } else { - if (src_shape[i] != in_shape[i]) { - throw std::runtime_error("IndexAdd: Source non-index dimension mismatch."); - } - } - } -} - -// 1. Out-of-place 接口 -Tensor index_add(Tensor input, int64_t dim, Tensor index, Tensor source, float alpha) { - check_index_add_args(input, dim, index, source); - - Tensor output = Tensor::empty(input->shape(), input->dtype(), input->device()); - output->copy_from(input); - if (!index->is_contiguous()) { - index = index->contiguous(); - } - if (!source->is_contiguous()) { - source = source->contiguous(); - } - IndexAdd::execute(output, output, dim, index, source, alpha); - - return output; -} - -// 2. In-place 接口 -void index_add_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source, float alpha) { - check_index_add_args(input, dim, index, source); - - if (output->shape() != input->shape()) { - throw std::runtime_error("IndexAdd (In-place): Output shape must match Input shape."); - } - - if (output.operator->() != input.operator->()) { - output->copy_from(input); - } - - if (!index->is_contiguous()) { - index = index->contiguous(); - } - if (!source->is_contiguous()) { - source = source->contiguous(); - } - - if (!output->is_contiguous()) { - - Tensor contiguous_out = output->contiguous(); - - IndexAdd::execute(contiguous_out, contiguous_out, dim, index, source, alpha); - - output->copy_from(contiguous_out); - } else { - // 正常路径: Output 已经是连续的，直接原地执行 - IndexAdd::execute(output, input, dim, index, source, alpha); - } -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/index_add/index_add_infiniop.cc b/src/infinicore/ops/index_add/index_add_infiniop.cc deleted file mode 100644 index d9635cb5d..000000000 --- a/src/infinicore/ops/index_add/index_add_infiniop.cc +++ /dev/null @@ -1,67 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/index_add.hpp" // 引用算子定义 -#include - -namespace infinicore::op::index_add_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopIndexAddDescriptor_t &desc) { - if (desc != nullptr) { - // 销毁描述符 - INFINICORE_CHECK_ERROR(infiniopDestroyIndexAddDescriptor(desc)); - desc = nullptr; - } - }); - -// 计算函数实现 -void calculate(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source, float alpha) { - size_t seed = hash_combine(output, input, dim, index, source, alpha); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopIndexAddDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateIndexAddDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc(), - dim, // 传入 int64_t - index->desc(), - source->desc(), - alpha)); // 传入 float - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 3. 获取 Workspace 大小并分配 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetIndexAddWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - INFINICORE_CHECK_ERROR(infiniopIndexAdd( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - index->data(), - source->data(), - context::getStream())); -} - -// 5. 注册算子到 Dispatcher -static bool registered = []() { - IndexAdd::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::index_add_impl::infiniop diff --git a/src/infinicore/ops/index_copy/index_copy.cc b/src/infinicore/ops/index_copy/index_copy.cc deleted file mode 100644 index 866e8764b..000000000 --- a/src/infinicore/ops/index_copy/index_copy.cc +++ /dev/null @@ -1,114 +0,0 @@ -#include "infinicore/ops/index_copy.hpp" -#include "infinicore/tensor.hpp" -#include -#include -#include - -namespace infinicore::op { - -// ========================================================= -// Dispatcher & Execute -// ========================================================= - -common::OpDispatcher &IndexCopy::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; -void IndexCopy::execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source) { - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No IndexCopy implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input, dim, index, source); -} -static void check_index_copy_args(const Tensor &input, int64_t &dim, const Tensor &index, const Tensor &source) { - int64_t ndim = static_cast(input->ndim()); - - if (dim < 0) { - dim += ndim; - } - if (dim < 0 || dim >= ndim) { - throw std::runtime_error("IndexCopy: Dimension out of range."); - } - - if (index->ndim() != 1) { - throw std::runtime_error("IndexCopy: Index tensor must be 1D."); - } - - // 使用 DataType::I64 和 I32 - if (index->dtype() != DataType::I64 && index->dtype() != DataType::I32) { - throw std::runtime_error("IndexCopy: Index tensor must be I32 or I64."); - } - - if (source->ndim() != input->ndim()) { - throw std::runtime_error("IndexCopy: Source tensor must have same number of dimensions as input tensor."); - } - - auto in_shape = input->shape(); - auto src_shape = source->shape(); - auto idx_len = index->shape()[0]; - - for (int64_t i = 0; i < ndim; ++i) { - if (i == dim) { - if (src_shape[i] != idx_len) { - throw std::runtime_error("IndexCopy: Source dimension mismatch."); - } - } else { - if (src_shape[i] != in_shape[i]) { - throw std::runtime_error("IndexCopy: Source non-index dimension mismatch."); - } - } - } -} - -Tensor index_copy(Tensor input, int64_t dim, Tensor index, Tensor source) { - check_index_copy_args(input, dim, index, source); - Tensor output = Tensor::empty(input->shape(), input->dtype(), input->device()); - output->copy_from(input); - if (!index->is_contiguous()) { - index = index->contiguous(); - } - if (!source->is_contiguous()) { - source = source->contiguous(); - } - IndexCopy::execute(output, output, dim, index, source); - - return output; -} - -// 2. In-place 接口 -void index_copy_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source) { - check_index_copy_args(input, dim, index, source); - - if (output->shape() != input->shape()) { - throw std::runtime_error("IndexCopy (In-place): Output shape must match Input shape."); - } - - if (output.operator->() != input.operator->()) { - output->copy_from(input); - } - - if (!index->is_contiguous()) { - index = index->contiguous(); - } - if (!source->is_contiguous()) { - source = source->contiguous(); - } - - if (!output->is_contiguous()) { - // 策略: Copy -> Compute -> CopyBack - Tensor contiguous_out = output->contiguous(); - - IndexCopy::execute(contiguous_out, contiguous_out, dim, index, source); - - // 写回结果 - output->copy_from(contiguous_out); - } else { - IndexCopy::execute(output, input, dim, index, source); - } -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/index_copy/index_copy_infiniop.cc b/src/infinicore/ops/index_copy/index_copy_infiniop.cc deleted file mode 100644 index 556856245..000000000 --- a/src/infinicore/ops/index_copy/index_copy_infiniop.cc +++ /dev/null @@ -1,63 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/index_copy.hpp" -#include - -namespace infinicore::op::index_copy_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopIndexCopyDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyIndexCopyDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source) { - size_t seed = hash_combine(output, input, dim, index, source); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopIndexCopyDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateIndexCopyDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc(), - dim, - index->desc(), - source->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetIndexCopyWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopIndexCopy( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - index->data(), - source->data(), - context::getStream())); -} - -static bool registered = []() { - IndexCopy::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::index_copy_impl::infiniop diff --git a/src/infinicore/ops/infiniop_impl.hpp b/src/infinicore/ops/infiniop_impl.hpp deleted file mode 100644 index 67c09554c..000000000 --- a/src/infinicore/ops/infiniop_impl.hpp +++ /dev/null @@ -1,73 +0,0 @@ -#pragma once - -#include "../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include - -#define INFINIOP_CACHABLE_DESCRIPTOR(__DESC_TYPE__, __OP_NAME__, __SIZE__) \ - struct __DESC_TYPE__ { \ - infiniop##__OP_NAME__##Descriptor_t desc = nullptr; \ - \ - explicit __DESC_TYPE__(infiniop##__OP_NAME__##Descriptor_t d) \ - : desc(d) {} \ - \ - /* non-copyable */ \ - __DESC_TYPE__(const __DESC_TYPE__ &) = delete; \ - __DESC_TYPE__ &operator=(const __DESC_TYPE__ &) = delete; \ - \ - /* movable */ \ - __DESC_TYPE__(__DESC_TYPE__ &&other) noexcept \ - : desc(other.desc) { \ - other.desc = nullptr; \ - } \ - \ - __DESC_TYPE__ &operator=(__DESC_TYPE__ &&other) noexcept { \ - if (this != &other) { \ - if (desc != nullptr) { \ - infiniopDestroy##__OP_NAME__##Descriptor(desc); \ - } \ - desc = other.desc; \ - other.desc = nullptr; \ - } \ - return *this; \ - } \ - \ - ~__DESC_TYPE__() { \ - if (desc != nullptr) { \ - infiniopDestroy##__OP_NAME__##Descriptor(desc); \ - } \ - } \ - }; \ - \ - thread_local common::OpCache> \ - caches( \ - __SIZE__, \ - [](std::shared_ptr<__DESC_TYPE__> &desc) { \ - desc = nullptr; \ - }); - -#define INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(__DESC_TYPE__, __DESC_NAME__, __INFINIOP_NAME__, __HASH_KEY__, ...) \ - std::shared_ptr<__DESC_TYPE__> __DESC_NAME__; \ - { \ - auto device__ = context::getDevice(); \ - auto &cache__ = caches.getCache(device__); \ - __DESC_NAME__ = cache__.get(__HASH_KEY__).value_or(nullptr); \ - if (!__DESC_NAME__) { \ - __DESC_NAME__ = std::make_shared<__DESC_TYPE__>(nullptr); \ - INFINICORE_CHECK_ERROR(infiniopCreate##__INFINIOP_NAME__##Descriptor( \ - context::getInfiniopHandle(device__), \ - &__DESC_NAME__->desc, \ - __VA_ARGS__)); \ - cache__.put(__HASH_KEY__, __DESC_NAME__); \ - } \ - } - -#define INFINIOP_WORKSPACE_TENSOR(__TENSOR_NAME__, __INFINIOP_NAME__, __DESC_NAME__) \ - Tensor __TENSOR_NAME__; \ - { \ - auto device__ = context::getDevice(); \ - size_t workspace_size = 0; \ - INFINICORE_CHECK_ERROR(infiniopGet##__INFINIOP_NAME__##WorkspaceSize(__DESC_NAME__->desc, &workspace_size)); \ - __TENSOR_NAME__ = Tensor::empty({workspace_size}, DataType::U8, device__); \ - } diff --git a/src/infinicore/ops/infiniops_impl.hpp b/src/infinicore/ops/infiniops_impl.hpp deleted file mode 100644 index befb830cd..000000000 --- a/src/infinicore/ops/infiniops_impl.hpp +++ /dev/null @@ -1,104 +0,0 @@ -#pragma once - -#include "../utils.hpp" -#include "infinicore/tensor.hpp" - -#include - -#include "config.h" -#include "data_type.h" -#include "handle.h" -#include "infini/ops.h" -#include "tensor.h" - -namespace infinicore::op::infiniops { - -inline infini::ops::DataType toInfiniOpsDtype(DataType dtype) { - switch (dtype) { - case DataType::I8: - return infini::ops::DataType::kInt8; - case DataType::I16: - return infini::ops::DataType::kInt16; - case DataType::I32: - return infini::ops::DataType::kInt32; - case DataType::I64: - return infini::ops::DataType::kInt64; - case DataType::U8: - case DataType::BYTE: - return infini::ops::DataType::kUInt8; - case DataType::U16: - return infini::ops::DataType::kUInt16; - case DataType::U32: - return infini::ops::DataType::kUInt32; - case DataType::U64: - return infini::ops::DataType::kUInt64; - case DataType::F16: - return infini::ops::DataType::kFloat16; - case DataType::BF16: - return infini::ops::DataType::kBFloat16; - case DataType::F32: - return infini::ops::DataType::kFloat32; - case DataType::F64: - return infini::ops::DataType::kFloat64; - default: - throw std::runtime_error("InfiniOps backend does not support this tensor dtype."); - } -} - -inline infini::ops::Device toInfiniOpsDevice(const Device &device) { - switch (device.getType()) { - case Device::Type::NVIDIA: - return infini::ops::Device{infini::ops::Device::Type::kNvidia, static_cast(device.getIndex())}; - case Device::Type::METAX: - return infini::ops::Device{infini::ops::Device::Type::kMetax, static_cast(device.getIndex())}; - case Device::Type::MOORE: - return infini::ops::Device{infini::ops::Device::Type::kMoore, static_cast(device.getIndex())}; - case Device::Type::ILUVATAR: - return infini::ops::Device{infini::ops::Device::Type::kIluvatar, static_cast(device.getIndex())}; - default: - throw std::runtime_error("InfiniOps backend does not support this device type."); - } -} - -inline bool isSupportedDevice(Device::Type device_type) { - switch (device_type) { - case Device::Type::NVIDIA: - case Device::Type::METAX: - case Device::Type::MOORE: - case Device::Type::ILUVATAR: - return true; - default: - return false; - } -} - -template -void registerSupportedDevices(Dispatcher &dispatcher, Function function) { - dispatcher.registerDevice(Device::Type::NVIDIA, function); - dispatcher.registerDevice(Device::Type::METAX, function); - dispatcher.registerDevice(Device::Type::MOORE, function); - dispatcher.registerDevice(Device::Type::ILUVATAR, function); -} - -struct TensorMeta { - Shape shape; - Strides strides; - infini::ops::DataType dtype; - infini::ops::Device device; - - explicit TensorMeta(const Tensor &tensor) - : shape(tensor->shape()), - strides(tensor->strides()), - dtype(toInfiniOpsDtype(tensor->dtype())), - device(toInfiniOpsDevice(tensor->device())) {} - - infini::ops::Tensor tensor(const void *data) const { - return infini::ops::Tensor(const_cast(data), shape, dtype, device, strides); - } - - infini::ops::Tensor tensor(const Tensor &tensor) const { - return this->tensor(tensor->data()); - } -}; - -} // namespace infinicore::op::infiniops diff --git a/src/infinicore/ops/inner/inner.cc b/src/infinicore/ops/inner/inner.cc deleted file mode 100644 index c4a9dec8b..000000000 --- a/src/infinicore/ops/inner/inner.cc +++ /dev/null @@ -1,46 +0,0 @@ -#include -#include - -namespace infinicore::op { - -common::OpDispatcher &Inner::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Inner::execute(Tensor out, Tensor input, Tensor other) { - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Inner implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(out, input, other); -} - -Tensor inner(Tensor input, Tensor other) { - size_t input_ndim = input->ndim(); - size_t other_ndim = other->ndim(); - - assert(input->shape()[input_ndim - 1] == other->shape()[other_ndim - 1]); - - Shape out_shape; - for (int i = 0; i < input_ndim - 1; i++) { - out_shape.push_back(input->shape()[i]); - } - for (int i = 0; i < other_ndim - 1; i++) { - out_shape.push_back(other->shape()[i]); - } - auto out = Tensor::zeros(out_shape, input->dtype(), input->device()); - - inner_(out, input, other); - return out; -} - -void inner_(Tensor out, Tensor input, Tensor other) { - - Inner::execute(out, input, other); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/inner/inner_infiniop.cc b/src/infinicore/ops/inner/inner_infiniop.cc deleted file mode 100644 index c8670e996..000000000 --- a/src/infinicore/ops/inner/inner_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/inner.hpp" -#include - -namespace infinicore::op::inner_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopInnerDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyInnerDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor out, Tensor input, Tensor other) { - size_t seed = hash_combine(out, input, other); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopInnerDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateInnerDescriptor( - context::getInfiniopHandle(input->device()), &desc, - out->desc(), input->desc(), other->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetInnerWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopInner( - desc, workspace->data(), workspace_size, - out->data(), input->data(), other->data(), context::getStream())); -} - -static bool registered = []() { - Inner::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::inner_impl::infiniop diff --git a/src/infinicore/ops/interpolate/interpolate.cc b/src/infinicore/ops/interpolate/interpolate.cc deleted file mode 100644 index 4a4460478..000000000 --- a/src/infinicore/ops/interpolate/interpolate.cc +++ /dev/null @@ -1,145 +0,0 @@ -#include "infinicore/ops/interpolate.hpp" -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Interpolate); - -Interpolate::Interpolate(Tensor out, - const Tensor &input, - std::string mode, - std::vector size, - std::vector scale_factor, - int align_corners) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, input, std::move(mode), std::move(size), std::move(scale_factor), align_corners); -} - -void Interpolate::execute(Tensor out, - const Tensor &input, - std::string mode, - std::vector size, - std::vector scale_factor, - int align_corners) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Interpolate, out, input, std::move(mode), std::move(size), std::move(scale_factor), align_corners); -} - -static std::vector infer_interpolate_shape( - const std::vector &input_shape, - const std::vector &size, - const std::vector &scale_factor) { - if (input_shape.size() < 3) { - throw std::runtime_error("interpolate expects input with at least 3 dimensions"); - } - - const size_t spatial_ndim = input_shape.size() - 2; - std::vector out_shape = input_shape; - - const bool has_size = !size.empty(); - const bool has_scale = !scale_factor.empty(); - if (has_size == has_scale) { - throw std::runtime_error("interpolate expects exactly one of size or scale_factor"); - } - - if (has_size) { - if (size.size() != spatial_ndim) { - throw std::runtime_error("interpolate size dimensionality mismatch"); - } - for (size_t i = 0; i < spatial_ndim; ++i) { - if (size[i] < 0) { - throw std::runtime_error("interpolate size values must be non-negative"); - } - out_shape[i + 2] = static_cast(size[i]); - } - return out_shape; - } - - if (scale_factor.size() != spatial_ndim) { - throw std::runtime_error("interpolate scale_factor dimensionality mismatch"); - } - for (size_t i = 1; i < spatial_ndim; ++i) { - if (scale_factor[i] != scale_factor[0]) { - throw std::runtime_error("interpolate only supports scalar/uniform scale_factor"); - } - } - const double scale = scale_factor[0]; - if (!std::isfinite(scale) || scale < 0.0) { - throw std::runtime_error("interpolate scale_factor must be finite and non-negative"); - } - for (size_t i = 0; i < spatial_ndim; ++i) { - out_shape[i + 2] = static_cast(static_cast(input_shape[i + 2]) * scale); - } - return out_shape; -} - -static void normalize_interpolate_params( - const std::vector &input_shape, - std::vector &size, - std::vector &scale_factor) { - if (input_shape.size() < 3) { - throw std::runtime_error("interpolate expects input with at least 3 dimensions"); - } - - const size_t spatial_ndim = input_shape.size() - 2; - const bool has_size = !size.empty(); - const bool has_scale = !scale_factor.empty(); - if (has_size == has_scale) { - throw std::runtime_error("interpolate expects exactly one of size or scale_factor"); - } - - if (has_size) { - if (size.size() == 1 && spatial_ndim > 1) { - size.assign(spatial_ndim, size[0]); - } - if (size.size() != spatial_ndim) { - throw std::runtime_error("interpolate size dimensionality mismatch"); - } - for (size_t i = 0; i < spatial_ndim; ++i) { - if (size[i] < 0) { - throw std::runtime_error("interpolate size values must be non-negative"); - } - } - return; - } - - if (scale_factor.size() == 1 && spatial_ndim > 1) { - scale_factor.assign(spatial_ndim, scale_factor[0]); - } - if (scale_factor.size() != spatial_ndim) { - throw std::runtime_error("interpolate scale_factor dimensionality mismatch"); - } - for (size_t i = 1; i < spatial_ndim; ++i) { - if (scale_factor[i] != scale_factor[0]) { - throw std::runtime_error("interpolate only supports scalar/uniform scale_factor"); - } - } - if (!std::isfinite(scale_factor[0]) || scale_factor[0] < 0.0) { - throw std::runtime_error("interpolate scale_factor must be finite and non-negative"); - } -} - -Tensor interpolate(const Tensor &input, - std::string mode, - std::vector size, - std::vector scale_factor, - int align_corners) { - normalize_interpolate_params(input->shape(), size, scale_factor); - auto out_shape = infer_interpolate_shape(input->shape(), size, scale_factor); - auto out = Tensor::empty(out_shape, input->dtype(), input->device()); - interpolate_(out, input, std::move(mode), std::move(size), std::move(scale_factor), align_corners); - return out; -} - -void interpolate_(Tensor out, - const Tensor &input, - std::string mode, - std::vector size, - std::vector scale_factor, - int align_corners) { - normalize_interpolate_params(input->shape(), size, scale_factor); - Interpolate::execute(out, input, std::move(mode), std::move(size), std::move(scale_factor), align_corners); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/interpolate/interpolate_infiniop.cc b/src/infinicore/ops/interpolate/interpolate_infiniop.cc deleted file mode 100644 index 9f974332d..000000000 --- a/src/infinicore/ops/interpolate/interpolate_infiniop.cc +++ /dev/null @@ -1,84 +0,0 @@ -#include "infinicore/ops/interpolate.hpp" - -#include "infiniop/ops/interpolate.h" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::interpolate_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Interpolate, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, input; -}; - -static size_t hash_mode_and_params(const std::string &mode, - int align_corners, - const std::vector &size, - const std::vector &scale_factor) { - size_t seed = 0; - hash_combine(seed, mode); - hash_combine(seed, align_corners); - hash_combine(seed, size.size()); - for (auto v : size) { - hash_combine(seed, v); - } - hash_combine(seed, scale_factor.size()); - for (auto v : scale_factor) { - hash_combine(seed, v); - } - return seed; -} - -void *plan(Tensor out, - const Tensor &input, - std::string mode, - std::vector size, - std::vector scale_factor, - int align_corners) { - const size_t params_hash = hash_mode_and_params(mode, align_corners, size, scale_factor); - const size_t seed = hash_combine(out, input, params_hash); - - const void *size_ptr = size.empty() ? nullptr : static_cast(size.data()); - const void *scale_ptr = scale_factor.empty() ? nullptr : static_cast(scale_factor.data()); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Interpolate, - seed, - out->desc(), - input->desc(), - mode.c_str(), - const_cast(size_ptr), - const_cast(scale_ptr), - align_corners); - - INFINIOP_WORKSPACE_TENSOR(workspace, Interpolate, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(input)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopInterpolate( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->input->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Interpolate, &plan, &run, &cleanup); - -} // namespace infinicore::op::interpolate_impl::infiniop diff --git a/src/infinicore/ops/kron/kron.cc b/src/infinicore/ops/kron/kron.cc deleted file mode 100644 index f47be6a8d..000000000 --- a/src/infinicore/ops/kron/kron.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/kron.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &Kron::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -void Kron::execute(Tensor output, Tensor a, Tensor b) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, a, b); - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error( - "No Kron implementation found for device type: " - + std::to_string(static_cast(device_type))); - } - - func(output, a, b); -} - -Tensor kron(Tensor a, Tensor b) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(a, b); - INFINICORE_ASSERT(a->dtype() == b->dtype()); - INFINICORE_ASSERT(a->shape().size() == b->shape().size()); - - const auto &a_shape = a->shape(); - const auto &b_shape = b->shape(); - Shape y_shape(a_shape.size()); - for (size_t i = 0; i < a_shape.size(); ++i) { - y_shape[i] = a_shape[i] * b_shape[i]; - } - - auto output = Tensor::empty(y_shape, a->dtype(), a->device()); - kron_(output, a, b); - return output; -} - -void kron_(Tensor output, Tensor a, Tensor b) { - Kron::execute(output, a, b); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/kron/kron_infiniop.cc b/src/infinicore/ops/kron/kron_infiniop.cc deleted file mode 100644 index c596d31b3..000000000 --- a/src/infinicore/ops/kron/kron_infiniop.cc +++ /dev/null @@ -1,59 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/kron.hpp" - -#include - -namespace infinicore::op::kron_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopKronDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyKronDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor a, Tensor b) { - size_t seed = hash_combine(output, a, b); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopKronDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateKronDescriptor( - context::getInfiniopHandle(device), - &desc, - output->desc(), - a->desc(), - b->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetKronWorkspaceSize(desc, &workspace_size)); - auto workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopKron( - desc, - workspace->data(), - workspace_size, - output->data(), - a->data(), - b->data(), - context::getStream())); -} - -static bool registered = []() { - Kron::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::kron_impl::infiniop diff --git a/src/infinicore/ops/kthvalue/kthvalue.cc b/src/infinicore/ops/kthvalue/kthvalue.cc deleted file mode 100644 index c96cd3576..000000000 --- a/src/infinicore/ops/kthvalue/kthvalue.cc +++ /dev/null @@ -1,49 +0,0 @@ -#include "infinicore/ops/kthvalue.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &Kthvalue::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Kthvalue::execute(Tensor values, Tensor indices, Tensor input, int64_t k, int64_t dim, bool keepdim) { - dispatcher().lookup(context::getDevice().getType())(values, indices, input, k, dim, keepdim); -} - -// 3. 函数式接口 -std::tuple kthvalue(Tensor input, int64_t k, int64_t dim, bool keepdim) { - auto input_shape = input->shape(); - int64_t ndim = input_shape.size(); - - // 处理负数维度 - if (dim < 0) { - dim += ndim; - } - - Shape output_shape; - if (keepdim) { - output_shape = input_shape; - output_shape[dim] = 1; - } else { - output_shape.reserve(ndim - 1); - for (int64_t i = 0; i < ndim; ++i) { - if (i != dim) { - output_shape.push_back(input_shape[i]); - } - } - } - - // values 与 input 类型一致 - auto values = Tensor::empty(output_shape, input->dtype(), input->device()); - auto indices = Tensor::empty(output_shape, DataType::I64, input->device()); - kthvalue_(values, indices, input, k, dim, keepdim); - return std::make_tuple(values, indices); -} - -void kthvalue_(Tensor values, Tensor indices, Tensor input, int64_t k, int64_t dim, bool keepdim) { - Kthvalue::execute(values, indices, input, k, dim, keepdim); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/kthvalue/kthvalue_infiniop.cc b/src/infinicore/ops/kthvalue/kthvalue_infiniop.cc deleted file mode 100644 index 2952183f4..000000000 --- a/src/infinicore/ops/kthvalue/kthvalue_infiniop.cc +++ /dev/null @@ -1,67 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/kthvalue.hpp" -#include - -namespace infinicore::op::kthvalue_impl::infiniop { - -// 定义描述符缓存 -thread_local common::OpCache caches( - 100, // capacity - [](infiniopKthvalueDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyKthvalueDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor values, Tensor indices, Tensor input, int64_t k, int64_t dim, bool keepdim) { - size_t seed = hash_combine(values, indices, input, k, dim, keepdim); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopKthvalueDescriptor_t desc = nullptr; - - if (!desc_opt) { - // 3. 创建描述符 - INFINICORE_CHECK_ERROR(infiniopCreateKthvalueDescriptor( - context::getInfiniopHandle(input->device()), - &desc, - values->desc(), - indices->desc(), - input->desc(), - static_cast(k), - static_cast(dim), - static_cast(keepdim))); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 4. 获取 Workspace 并执行 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetKthvalueWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopKthvalue( - desc, - workspace->data(), - workspace_size, - values->data(), - indices->data(), - input->data(), - context::getStream())); -} - -static bool registered = []() { - Kthvalue::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::kthvalue_impl::infiniop diff --git a/src/infinicore/ops/kv_caching/kv_caching.cc b/src/infinicore/ops/kv_caching/kv_caching.cc deleted file mode 100644 index 0110f7973..000000000 --- a/src/infinicore/ops/kv_caching/kv_caching.cc +++ /dev/null @@ -1,42 +0,0 @@ -#include "infinicore/ops/kv_caching.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(KVCaching); - -KVCaching::KVCaching(Tensor k_cache, - Tensor v_cache, - const Tensor &k, - const Tensor &v, - const Tensor &past_kv_lengths) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, past_kv_lengths); - INFINICORE_GRAPH_OP_DISPATCH(k_cache->device().getType(), - k_cache, - v_cache, - k, - v, - past_kv_lengths); -} - -void KVCaching::execute(Tensor k_cache, - Tensor v_cache, - const Tensor &k, - const Tensor &v, - const Tensor &past_kv_lengths) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(KVCaching, - k_cache, - v_cache, - k, - v, - past_kv_lengths); -} - -void kv_caching_(Tensor k_cache, - Tensor v_cache, - const Tensor &k, - const Tensor &v, - const Tensor &past_kv_lengths) { - KVCaching::execute(k_cache, v_cache, k, v, past_kv_lengths); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/kv_caching/kv_caching_infiniop.cc b/src/infinicore/ops/kv_caching/kv_caching_infiniop.cc deleted file mode 100644 index 53ea5f0ae..000000000 --- a/src/infinicore/ops/kv_caching/kv_caching_infiniop.cc +++ /dev/null @@ -1,60 +0,0 @@ -#include "../infiniop_impl.hpp" -#include "infinicore/ops/kv_caching.hpp" - -namespace infinicore::op::kv_caching_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, KVCaching, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, k_cache, v_cache, k, v, past_kv_lengths; -}; - -void *plan(Tensor k_cache, - Tensor v_cache, - const Tensor &k, - const Tensor &v, - const Tensor &past_kv_lengths) { - size_t seed = hash_combine(k_cache, v_cache, k, v, past_kv_lengths); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, KVCaching, - seed, k_cache->desc(), v_cache->desc(), - k->desc(), v->desc(), past_kv_lengths->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, KVCaching, descriptor); - - auto planned = new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(k_cache), - graph::GraphTensor(v_cache), - graph::GraphTensor(k), - graph::GraphTensor(v), - graph::GraphTensor(past_kv_lengths)}; - - return planned; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopKVCaching( - planned->descriptor->desc, - nullptr, 0, - planned->k_cache->data(), - planned->v_cache->data(), - planned->k->data(), - planned->v->data(), - planned->past_kv_lengths->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(KVCaching, &plan, &run, cleanup); - -} // namespace infinicore::op::kv_caching_impl::infiniop diff --git a/src/infinicore/ops/kv_caching/kv_caching_infiniops.cc b/src/infinicore/ops/kv_caching/kv_caching_infiniops.cc deleted file mode 100644 index 5a2774f0d..000000000 --- a/src/infinicore/ops/kv_caching/kv_caching_infiniops.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/kv_caching.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/kv_caching_infinilm.h" - -namespace infinicore::op::kv_caching_impl::infiniops { -namespace { -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; -struct PlannedMeta { - TensorMeta k_cache, v_cache, k, v, past_kv_lengths; - graph::GraphTensor k_cache_tensor, v_cache_tensor, k_tensor, v_tensor, past_kv_lengths_tensor; -}; -} // namespace - -void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &past_kv_lengths) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(k_cache->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, past_kv_lengths); - return new PlannedMeta{TensorMeta(k_cache), TensorMeta(v_cache), TensorMeta(k), TensorMeta(v), TensorMeta(past_kv_lengths), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(past_kv_lengths)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - infini::ops::KvCachingInfinilm::Call( - handle, - config, - planned->k.tensor(planned->k_tensor), - planned->v.tensor(planned->v_tensor), - planned->past_kv_lengths.tensor(planned->past_kv_lengths_tensor), - planned->k_cache.tensor(planned->k_cache_tensor), - planned->v_cache.tensor(planned->v_cache_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(KVCaching::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(KVCaching::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(KVCaching::cleanup_dispatcher(), &cleanup); - return true; -}(); -} // namespace infinicore::op::kv_caching_impl::infiniops -#endif diff --git a/src/infinicore/ops/layer_norm/layer_norm.cc b/src/infinicore/ops/layer_norm/layer_norm.cc deleted file mode 100644 index 6749db68e..000000000 --- a/src/infinicore/ops/layer_norm/layer_norm.cc +++ /dev/null @@ -1,42 +0,0 @@ -#include "infinicore/ops/layer_norm.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(LayerNorm); - -LayerNorm::LayerNorm(Tensor y, Tensor standardization, Tensor std_deviation, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, standardization, std_deviation, x, weight); - INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, standardization, std_deviation, x, weight, bias, epsilon); -} - -void LayerNorm::execute(Tensor y, Tensor standardization, Tensor std_deviation, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(LayerNorm, y, standardization, std_deviation, x, weight, bias, epsilon); -} - -Tensor layer_norm(const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { - auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); - auto reduced_shape = x->shape(); - reduced_shape.pop_back(); - auto standardization = Tensor::empty(x->shape(), x->dtype(), x->device()); - auto std_deviation = Tensor::empty(reduced_shape, x->dtype(), x->device()); - layer_norm_(y, standardization, std_deviation, x, weight, bias, epsilon); - return y; -} - -void layer_norm_(Tensor y, Tensor standardization, Tensor std_deviation, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { - LayerNorm::execute(y, standardization, std_deviation, x, weight, bias, epsilon); -} - -void layer_norm_(Tensor y, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { - auto reduced_shape = x->shape(); - reduced_shape.pop_back(); - auto standardization = Tensor::empty(x->shape(), x->dtype(), x->device()); - auto std_deviation = Tensor::empty(reduced_shape, x->dtype(), x->device()); - LayerNorm::execute(y, standardization, std_deviation, x, weight, bias, epsilon); -} - -void layer_norm_for_pybind(Tensor y, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { - layer_norm_(y, x, weight, bias, epsilon); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/layer_norm/layer_norm_infiniop.cc b/src/infinicore/ops/layer_norm/layer_norm_infiniop.cc deleted file mode 100644 index 2b4a5f335..000000000 --- a/src/infinicore/ops/layer_norm/layer_norm_infiniop.cc +++ /dev/null @@ -1,65 +0,0 @@ -#include "infinicore/ops/layer_norm.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::layer_norm_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, LayerNorm, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, y, standardization, std_deviation, x, weight, bias; -}; - -void *plan(Tensor y, Tensor standardization, Tensor std_deviation, const Tensor &x, const Tensor &weight, const Tensor &bias, float epsilon) { - size_t seed = hash_combine(y, standardization, std_deviation, x, weight, bias, epsilon); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, LayerNorm, - seed, - y->desc(), - standardization->desc(), - std_deviation->desc(), - x->desc(), - weight->desc(), - bias->desc(), - epsilon); - - INFINIOP_WORKSPACE_TENSOR(workspace, LayerNorm, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(y), - graph::GraphTensor(standardization), - graph::GraphTensor(std_deviation), - graph::GraphTensor(x), - graph::GraphTensor(weight), - graph::GraphTensor(bias)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR( - infiniopLayerNorm( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->y->data(), - planned->standardization->data(), - planned->std_deviation->data(), - planned->x->data(), - planned->weight->data(), - planned->bias->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(LayerNorm, &plan, &run, &cleanup); - -} // namespace infinicore::op::layer_norm_impl::infiniop diff --git a/src/infinicore/ops/ldexp/ldexp.cc b/src/infinicore/ops/ldexp/ldexp.cc deleted file mode 100644 index 31d237f37..000000000 --- a/src/infinicore/ops/ldexp/ldexp.cc +++ /dev/null @@ -1,51 +0,0 @@ -#include "infinicore/ops/ldexp.hpp" -#include // for std::max - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &Ldexp::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Ldexp::execute(Tensor output, Tensor input, Tensor other) { - dispatcher().lookup(context::getDevice().getType())(output, input, other); -} - -// 2. 函数式接口 -Tensor ldexp(Tensor input, Tensor other) { - // 计算广播后的输出形状 (Broadcasting Logic) - const auto &shape_a = input->shape(); - const auto &shape_b = other->shape(); - - size_t ndim_a = shape_a.size(); - size_t ndim_b = shape_b.size(); - size_t ndim_out = std::max(ndim_a, ndim_b); - - Shape output_shape(ndim_out); - - // 从后往前对齐维度进行广播检查 - for (size_t i = 0; i < ndim_out; ++i) { - // 获取对应的维度大小，若越界则视为 1 (右对齐) - int64_t dim_a = (i >= ndim_out - ndim_a) ? shape_a[i - (ndim_out - ndim_a)] : 1; - int64_t dim_b = (i >= ndim_out - ndim_b) ? shape_b[i - (ndim_out - ndim_b)] : 1; - output_shape[i] = std::max(dim_a, dim_b); - } - - // 分配输出 Tensor - // ldexp 的输出类型通常跟随 input (尾数)，设备跟随 input - auto output = Tensor::empty(output_shape, input->dtype(), input->device()); - - // 调用 Explicit output 接口 - ldexp_(output, input, other); - - return output; -} - -// 3. Explicit Output 接口 -void ldexp_(Tensor output, Tensor input, Tensor other) { - Ldexp::execute(output, input, other); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/ldexp/ldexp_infiniop.cc b/src/infinicore/ops/ldexp/ldexp_infiniop.cc deleted file mode 100644 index bb59c17a6..000000000 --- a/src/infinicore/ops/ldexp/ldexp_infiniop.cc +++ /dev/null @@ -1,84 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/ldexp.hpp" -#include - -namespace infinicore::op::ldexp_impl::infiniop { - -// 定义描述符缓存 -thread_local common::OpCache caches( - 100, // capacity - [](infiniopLdexpDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyLdexpDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, Tensor other) { - // 1. 计算哈希值 - // 注意：必须手动哈希 strides 以区分 Broadcasting 和 Inplace 情况 - size_t seed = 0; - auto combine_tensor_meta = [&](Tensor t) { - infinicore::hash_combine(seed, static_cast(t->dtype())); - for (auto s : t->shape()) { - infinicore::hash_combine(seed, s); - } - for (auto str : t->strides()) { - infinicore::hash_combine(seed, str); - } - }; - - combine_tensor_meta(output); - combine_tensor_meta(input); - combine_tensor_meta(other); - - // 2. 获取缓存对象 - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopLdexpDescriptor_t desc = nullptr; - - if (!desc_opt) { - // 3. 创建描述符 - // 这里后端 (CPU/GPU) 会根据 input/other 类型决定是否需要 workspace - INFINICORE_CHECK_ERROR(infiniopCreateLdexpDescriptor( - context::getInfiniopHandle(input->device()), - &desc, - output->desc(), - input->desc(), - other->desc())); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 4. 获取 Workspace 并执行 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetLdexpWorkspaceSize(desc, &workspace_size)); - - // 如果后端检测到需要转换 Int32 -> Float，workspace_size 会 > 0 - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopLdexp( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - other->data(), - context::getStream())); -} - -// 5. 注册算子 -static bool registered = []() { - // 注册为普通算子 (dispatcher 会自动处理 inplace 逻辑) - Ldexp::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::ldexp_impl::infiniop diff --git a/src/infinicore/ops/lerp/lerp.cc b/src/infinicore/ops/lerp/lerp.cc deleted file mode 100644 index be05b68d7..000000000 --- a/src/infinicore/ops/lerp/lerp.cc +++ /dev/null @@ -1,132 +0,0 @@ -#include "infinicore/ops/lerp.hpp" -#include // for std::max -#include // for std::runtime_error -#include - -namespace infinicore::op { - -// ======================================================================== -// 0. 内部辅助函数：手动实现形状广播推导 -// ======================================================================== -namespace { - -Shape compute_broadcast_shape(const std::vector &shapes) { - if (shapes.empty()) { - return {}; - } - - // 1. 找出最大的维度数 (Max Rank) - size_t max_ndim = 0; - for (const auto &shape : shapes) { - max_ndim = std::max(max_ndim, shape.size()); - } - - Shape out_shape(max_ndim); - - // 2. 从右向左遍历每一个维度 (Standard Broadcasting Rule) - for (size_t i = 0; i < max_ndim; ++i) { - size_t current_dim_val = 1; - bool set = false; - - for (const auto &shape : shapes) { - // 计算当前 shape 对应的维度索引 (从右对齐) - // 比如 max_ndim=4, 当前 shape_ndim=2, i=0 (最右边) - // shape index = 2 - 1 - 0 = 1 - if (i < shape.size()) { - size_t dim = shape[shape.size() - 1 - i]; - - if (dim == 1) { - continue; // 1 可以被广播，忽略 - } - - if (!set) { - current_dim_val = dim; - set = true; - } else if (current_dim_val != dim) { - // 维度不相等，且都不为 1，无法广播 - throw std::runtime_error( - "Lerp: Shapes are not broadcastable. Mismatch at dimension offset " + std::to_string(i)); - } - } - } - // 填充输出形状 (从右向左填，或者填好后由 vector 自动管理) - out_shape[max_ndim - 1 - i] = current_dim_val; - } - - return out_shape; -} - -} // namespace - -// ======================================================================== -// 1. 定义 Dispatcher 单例 -// ======================================================================== - -template <> -common::OpDispatcher &Lerp::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -template <> -common::OpDispatcher &Lerp::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -// ======================================================================== -// 2. Execute 静态方法实现 -// ======================================================================== - -void Lerp::execute(Tensor output, Tensor start, Tensor end, Tensor weight) { - dispatcher().lookup(context::getDevice().getType())(output, start, end, weight); -} - -void Lerp::execute(Tensor output, Tensor start, Tensor end, float weight) { - dispatcher().lookup(context::getDevice().getType())(output, start, end, weight); -} - -// ======================================================================== -// 3. 函数式接口 (Functional API) - 集成形状推导 -// ======================================================================== - -Tensor lerp(Tensor start, Tensor end, Tensor weight) { - // 1. 调用本地实现的推导函数，计算 start, end, weight 三者的广播形状 - Shape output_shape = compute_broadcast_shape({start->shape(), - end->shape(), - weight->shape()}); - - // 2. 分配输出内存 - auto output = Tensor::empty(output_shape, start->dtype(), start->device()); - - // 3. 执行计算 - lerp_(output, start, end, weight); - return output; -} - -Tensor lerp(Tensor start, Tensor end, float weight) { - // 1. 计算 start, end 两者的广播形状 (标量 weight 不参与形状计算) - Shape output_shape = compute_broadcast_shape({start->shape(), - end->shape()}); - - // 2. 分配输出内存 - auto output = Tensor::empty(output_shape, start->dtype(), start->device()); - - // 3. 执行计算 - lerp_(output, start, end, weight); - return output; -} - -// ======================================================================== -// 4. In-place / Output-buffer 接口 -// ======================================================================== - -void lerp_(Tensor output, Tensor start, Tensor end, Tensor weight) { - Lerp::execute(output, start, end, weight); -} - -void lerp_(Tensor output, Tensor start, Tensor end, float weight) { - Lerp::execute(output, start, end, weight); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/lerp/lerp_infiniop.cc b/src/infinicore/ops/lerp/lerp_infiniop.cc deleted file mode 100644 index 7d58af02b..000000000 --- a/src/infinicore/ops/lerp/lerp_infiniop.cc +++ /dev/null @@ -1,114 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/lerp.hpp" -#include - -namespace infinicore::op::lerp_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopLerpDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyLerpDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor start, Tensor end, Tensor weight) { - size_t seed = hash_combine(output, start, end, weight); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - auto desc_opt = cache.get(seed); - - infiniopLerpDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateLerpDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - start->desc(), - end->desc(), - weight->desc(), - 0.0f)); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetLerpWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopLerp( - desc, - workspace->data(), - workspace_size, - output->data(), - start->data(), - end->data(), - weight->data(), - context::getStream())); -} - -void calculate(Tensor output, Tensor start, Tensor end, float weight) { - size_t seed = hash_combine(output, start, end, weight); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - auto desc_opt = cache.get(seed); - - infiniopLerpDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateLerpDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - start->desc(), - end->desc(), - nullptr, - weight)); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetLerpWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopLerp( - desc, - workspace->data(), - workspace_size, - output->data(), - start->data(), - end->data(), - nullptr, - context::getStream())); -} - -static bool registered = []() { - using SchemaTensor = void (*)(Tensor, Tensor, Tensor, Tensor); - Lerp::dispatcher().registerAll( - static_cast(&calculate), - false); - - using SchemaScalar = void (*)(Tensor, Tensor, Tensor, float); - Lerp::dispatcher().registerAll( - static_cast(&calculate), - false); - - return true; -}(); - -} // namespace infinicore::op::lerp_impl::infiniop diff --git a/src/infinicore/ops/linear/linear.cc b/src/infinicore/ops/linear/linear.cc deleted file mode 100644 index 976bf0bf9..000000000 --- a/src/infinicore/ops/linear/linear.cc +++ /dev/null @@ -1,60 +0,0 @@ -#include "infinicore/ops/linear.hpp" -#include "infinicore/ops/gemm.hpp" -#include "infinicore/ops/rearrange.hpp" - -namespace infinicore::op { - -Tensor linear(Tensor input, - Tensor weight, - std::optional bias, - float alpha) { - - Size ndim = input->ndim(); - Size out_features = weight->shape()[0]; - - // Assign memory to out variables - auto output_shape = input->shape(); - output_shape[ndim - 1] = out_features; - auto out = Tensor::empty(output_shape, input->dtype(), input->device()); - - // Inplace Calculate - linear_(out, input, weight, bias, alpha); - return out; -} - -void linear_(Tensor out, - Tensor input, - Tensor weight, - std::optional bias, - float alpha) { - - auto weight_shape = weight->shape(); - Size out_features = weight_shape[0]; - Size in_features = weight_shape[1]; - - Size ndim = input->ndim(); - assert(out->ndim() == ndim); - - // Calculate the number of features - Size N = 1; - auto input_shape = input->shape(); - for (size_t i = 0; i < ndim - 1; ++i) { - N *= input_shape[i]; - } - - // linear transformation - Tensor out_view = out->view({N, out_features}); - // Add bias - float beta = 0.0f; - if (bias.has_value()) { - rearrange_(out_view, - bias.value()->as_strided({N, out_features}, {0, 1})); - beta = 1.0f; - } - - gemm_(out_view, - input->view({N, in_features}), - weight->permute({1, 0}), alpha, beta); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/linear_w4a16_awq/linear_w4a16_awq.cc b/src/infinicore/ops/linear_w4a16_awq/linear_w4a16_awq.cc deleted file mode 100644 index b602574e0..000000000 --- a/src/infinicore/ops/linear_w4a16_awq/linear_w4a16_awq.cc +++ /dev/null @@ -1,64 +0,0 @@ -#include "infinicore/ops/linear_w4a16_awq.hpp" -#include "infinicore/ops/dequantize_awq.hpp" -#include "infinicore/ops/gemm.hpp" -#include "infinicore/ops/rearrange.hpp" -namespace infinicore::op { - -Tensor linear_w4a16_awq(Tensor input, - Tensor weight_packed, - Tensor weight_scale, - Tensor weight_zeros, - std::optional bias) { - - // Input is of shape [M, K], Weight_packed is of shape [N, K],stirdes is [N, 1] - Size ndim = input->ndim(); - Size element_size = weight_packed->element_size(); - Size out_features = weight_packed->shape()[1] * element_size * 2; - - // Assign memory to out variables - auto output_shape = input->shape(); - output_shape[ndim - 1] = out_features; - auto out = Tensor::empty(output_shape, input->dtype(), input->device()); - - // Inplace Calculate - linear_w4a16_awq_(out, input, weight_packed, weight_scale, weight_zeros, bias); - return out; -} - -void linear_w4a16_awq_(Tensor out, - Tensor input, - Tensor weight_packed, - Tensor weight_scale, - Tensor weight_zeros, - std::optional bias) { - - auto weight_packed_shape = weight_packed->shape(); - Size out_features = weight_packed_shape[0]; - Size in_features = weight_packed_shape[1] * 8; - - Size ndim = input->ndim(); - assert(out->ndim() == ndim); - - Size N = 1; - auto input_shape = input->shape(); - for (size_t i = 0; i < ndim - 1; ++i) { - N *= input_shape[i]; - } - auto weight = Tensor::empty( - {out_features, in_features}, - out->dtype(), - weight_packed->device()); - float alpha = 1.0f; - float beta = 0.0f; - op::dequantize_awq_(weight, weight_packed, weight_scale, weight_zeros); - if (bias.has_value()) { - rearrange_(out, - bias.value()->as_strided({N, in_features}, {0, 1})); - beta = 1.0f; - } - gemm_(out, - input->view({N, out_features}), - weight, alpha, beta); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/linear_w4a16_gptq_qy/linear_w4a16_gptq_qy.cc b/src/infinicore/ops/linear_w4a16_gptq_qy/linear_w4a16_gptq_qy.cc deleted file mode 100644 index bb088d709..000000000 --- a/src/infinicore/ops/linear_w4a16_gptq_qy/linear_w4a16_gptq_qy.cc +++ /dev/null @@ -1,54 +0,0 @@ -#include "infinicore/ops/linear_w4a16_gptq_qy.hpp" -#include "infinicore/ops/scaled_mm_w4a16_gptq_qy.hpp" -#include -namespace infinicore::op { - -Tensor linear_w4a16_gptq_qy(Tensor input, Tensor qweight, Tensor qzeros, Tensor scales, int64_t quant_type, int64_t bit) { - - Size ndim = input->ndim(); - - Size out_features = qweight->shape()[1]; - - // 2. 计算输出形状 [..., out_features] - auto output_shape = input->shape(); - output_shape[ndim - 1] = out_features; - - // 3. 分配输出显存 - auto out = Tensor::zeros(output_shape, input->dtype(), input->device()); - - // 4. 执行计算 - linear_w4a16_gptq_qy_(out, input, qweight, scales, qzeros, quant_type, bit); - - return out; -} - -void linear_w4a16_gptq_qy_(Tensor out, Tensor in, Tensor qweights, Tensor scales, Tensor qzeros, int64_t quant_type, int64_t bit) { - - Size in_features = qweights->shape()[0] * 2; // ✅ 修正：第 0 维是 in/2 - Size out_features = qweights->shape()[1]; // ✅ 修正：第 1 维是 out - - // 检查输入输出维度 - Size ndim = in->ndim(); - - // ======================================================================== - // 合并 Batch 维度 - // ======================================================================== - Size N = 1; - auto input_shape = in->shape(); - for (size_t i = 0; i < ndim - 1; ++i) { - N *= input_shape[i]; - } - - op::scaled_mm_w4a16_gptq_qy_( - out->view({N, out_features}), // Output: [N, out] - in->view({N, in_features}), // Input: [N, in] - qweights, // Weight: [in/2, out] - scales, // Scales: [in/group, out] - qzeros, // QZeros: [in/group, out] - quant_type, // Quantization type - bit // Bit width - ); - // out->debug(); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/linear_w8a8i8/linear_w8a8i8.cc b/src/infinicore/ops/linear_w8a8i8/linear_w8a8i8.cc deleted file mode 100644 index d69e0e7a2..000000000 --- a/src/infinicore/ops/linear_w8a8i8/linear_w8a8i8.cc +++ /dev/null @@ -1,66 +0,0 @@ -#include "infinicore/ops/linear_w8a8i8.hpp" -#include "infinicore/ops/per_channel_quant_i8.hpp" -#include "infinicore/ops/scaled_mm_i8.hpp" - -namespace infinicore::op { - -Tensor linear_w8a8i8(Tensor input, - Tensor weight_packed, - Tensor weight_scale, - std::optional bias) { - - // Input is of shape [M, K], Weight_packed is of shape [N, K],stirdes is [N, 1] - Size ndim = input->ndim(); - Size out_features = weight_packed->shape()[0]; - - // Assign memory to out variables - auto output_shape = input->shape(); - output_shape[ndim - 1] = out_features; - auto out = Tensor::empty(output_shape, input->dtype(), input->device()); - - // Inplace Calculate - linear_w8a8i8_(out, input, weight_packed, weight_scale, bias); - return out; -} - -void linear_w8a8i8_(Tensor out, - Tensor input, - Tensor weight_packed, - Tensor weight_scale, - std::optional bias) { - - auto weight_packed_shape = weight_packed->shape(); - Size out_features = weight_packed_shape[0]; - Size in_features = weight_packed_shape[1]; - - Size ndim = input->ndim(); - assert(out->ndim() == ndim); - - Size N = 1; - auto input_shape = input->shape(); - for (size_t i = 0; i < ndim - 1; ++i) { - N *= input_shape[i]; - } - - auto input_packed = Tensor::empty( - {N, input_shape[ndim - 1]}, - DataType::I8, - input->device()); - auto input_scale = Tensor::empty( - {N, 1}, - DataType::F32, - input->device()); - op::per_channel_quant_i8_(input->view({N, in_features}), input_packed, input_scale); - if (bias.has_value()) { - bias = std::make_optional(bias.value()->as_strided({N, out_features}, {0, 1})); - } - op::scaled_mm_i8_( - out->view({N, out_features}), - input_packed, - input_scale, - weight_packed->permute({1, 0}), - weight_scale, - bias); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/log_softmax/log_softmax.cc b/src/infinicore/ops/log_softmax/log_softmax.cc deleted file mode 100644 index 08ea8550a..000000000 --- a/src/infinicore/ops/log_softmax/log_softmax.cc +++ /dev/null @@ -1,34 +0,0 @@ -#include "infinicore/ops/log_softmax.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &LogSoftmax::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void LogSoftmax::execute(Tensor output, Tensor input, int64_t dim) { - dispatcher().lookup(context::getDevice().getType())(output, input, dim); -} - -// 3. 函数式接口 -Tensor log_softmax(Tensor input, int64_t dim) { - int64_t ndim = input->shape().size(); - - // 处理负数维度 - if (dim < 0) { - dim += ndim; - } - - // LogSoftmax 输出形状与输入一致，dtype 与 input 一致 - auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); - log_softmax_(output, input, dim); - return output; -} - -void log_softmax_(Tensor output, Tensor input, int64_t dim) { - LogSoftmax::execute(output, input, dim); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/log_softmax/log_softmax_infiniop.cc b/src/infinicore/ops/log_softmax/log_softmax_infiniop.cc deleted file mode 100644 index c613a22a2..000000000 --- a/src/infinicore/ops/log_softmax/log_softmax_infiniop.cc +++ /dev/null @@ -1,63 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/log_softmax.hpp" -#include - -namespace infinicore::op::log_softmax_impl::infiniop { - -// 定义描述符缓存 -thread_local common::OpCache caches( - 100, // capacity - [](infiniopLogSoftmaxDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyLogSoftmaxDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, int64_t dim) { - size_t seed = hash_combine(output, input, dim); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopLogSoftmaxDescriptor_t desc = nullptr; - - if (!desc_opt) { - // 3. 创建描述符 - INFINICORE_CHECK_ERROR(infiniopCreateLogSoftmaxDescriptor( - context::getInfiniopHandle(input->device()), - &desc, - output->desc(), - input->desc(), - static_cast(dim))); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 4. 获取 Workspace 并执行 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetLogSoftmaxWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopLogSoftmax( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -static bool registered = []() { - LogSoftmax::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::log_softmax_impl::infiniop diff --git a/src/infinicore/ops/logaddexp/logaddexp.cc b/src/infinicore/ops/logaddexp/logaddexp.cc deleted file mode 100644 index f9b2c749c..000000000 --- a/src/infinicore/ops/logaddexp/logaddexp.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/logaddexp.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &LogAddExp::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void LogAddExp::execute(Tensor c, Tensor a, Tensor b) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - infinicore::context::setDevice(c->device()); - dispatcher().lookup(c->device().getType())(c, a, b); -} - -Tensor logaddexp(Tensor a, Tensor b) { - auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); - logaddexp_(c, a, b); - return c; -} - -void logaddexp_(Tensor c, Tensor a, Tensor b) { - LogAddExp::execute(c, a, b); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/logaddexp/logaddexp_infiniop.cc b/src/infinicore/ops/logaddexp/logaddexp_infiniop.cc deleted file mode 100644 index 776573761..000000000 --- a/src/infinicore/ops/logaddexp/logaddexp_infiniop.cc +++ /dev/null @@ -1,48 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/logaddexp.hpp" -#include - -namespace infinicore::op::logaddexp_impl::infiniop { -thread_local common::OpCache caches( - 100, // capacity - [](infiniopLogAddExpDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyLogAddExpDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor c, Tensor a, Tensor b) { - size_t seed = hash_combine(c, a, b); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopLogAddExpDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateLogAddExpDescriptor( - context::getInfiniopHandle(device), &desc, - c->desc(), a->desc(), b->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetLogAddExpWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopLogAddExp( - desc, workspace->data(), workspace_size, - c->data(), a->data(), b->data(), context::getStream())); -} - -static bool registered = []() { - LogAddExp::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::logaddexp_impl::infiniop diff --git a/src/infinicore/ops/logaddexp2/logaddxep2.cc b/src/infinicore/ops/logaddexp2/logaddxep2.cc deleted file mode 100644 index d196545b7..000000000 --- a/src/infinicore/ops/logaddexp2/logaddxep2.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/ops/logaddexp2.hpp" - -namespace infinicore::op { - -common::OpDispatcher &LogAddExp2::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void LogAddExp2::execute(Tensor c, Tensor a, Tensor b) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - infinicore::context::setDevice(c->device()); - dispatcher().lookup(c->device().getType())(c, a, b); -} - -Tensor logaddexp2(Tensor a, Tensor b) { - auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); - logaddexp2_(c, a, b); - return c; -} - -void logaddexp2_(Tensor c, Tensor a, Tensor b) { - LogAddExp2::execute(c, a, b); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/logaddexp2/logaddxep2_infiniop.cc b/src/infinicore/ops/logaddexp2/logaddxep2_infiniop.cc deleted file mode 100644 index ffe09ceac..000000000 --- a/src/infinicore/ops/logaddexp2/logaddxep2_infiniop.cc +++ /dev/null @@ -1,48 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/logaddexp2.hpp" -#include - -namespace infinicore::op::logaddexp2_impl::infiniop { -thread_local common::OpCache caches( - 100, // capacity - [](infiniopLogAddExp2Descriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyLogAddExp2Descriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor c, Tensor a, Tensor b) { - size_t seed = hash_combine(c, a, b); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopLogAddExp2Descriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateLogAddExp2Descriptor( - context::getInfiniopHandle(device), &desc, - c->desc(), a->desc(), b->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetLogAddExp2WorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopLogAddExp2( - desc, workspace->data(), workspace_size, - c->data(), a->data(), b->data(), context::getStream())); -} - -static bool registered = []() { - LogAddExp2::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::logaddexp2_impl::infiniop diff --git a/src/infinicore/ops/logcumsumexp/logcumsumexp.cc b/src/infinicore/ops/logcumsumexp/logcumsumexp.cc deleted file mode 100644 index 60289fa24..000000000 --- a/src/infinicore/ops/logcumsumexp/logcumsumexp.cc +++ /dev/null @@ -1,35 +0,0 @@ -#include "infinicore/ops/logcumsumexp.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -// 初始化 Dispatcher 单例 -common::OpDispatcher &LogCumSumExp::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -// 算子执行逻辑：校验设备并分发任务 -void LogCumSumExp::execute(Tensor y, Tensor x, int axis, bool exclusive, bool reverse) { - // 确保输入输出张量在同一设备上 - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - - // 切换到目标设备的上下文 - infinicore::context::setDevice(y->device()); - dispatcher().lookup(y->device().getType())(y, x, axis, exclusive, reverse); -} - -// 函数式接口：自动创建输出张量并返回 -Tensor logcumsumexp(Tensor x, int axis, bool exclusive, bool reverse) { - // 创建一个与输入 x 形状、类型和设备相同的空张量作为输出 - auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); - logcumsumexp_(y, x, axis, exclusive, reverse); - return y; -} - -// 原地/指定输出接口 -void logcumsumexp_(Tensor y, Tensor x, int axis, bool exclusive, bool reverse) { - LogCumSumExp::execute(y, x, axis, exclusive, reverse); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/logcumsumexp/logcumsumexp_infiniop.cc b/src/infinicore/ops/logcumsumexp/logcumsumexp_infiniop.cc deleted file mode 100644 index 187786f87..000000000 --- a/src/infinicore/ops/logcumsumexp/logcumsumexp_infiniop.cc +++ /dev/null @@ -1,69 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/logcumsumexp.hpp" -#include - -namespace infinicore::op::logcumsumexp_impl::infiniop { - -// 定义缓存：用于管理 infiniop 算子描述符的生命周期 -thread_local common::OpCache caches( - 100, // 缓存容量 - [](infiniopLogCumSumExpDescriptor_t &desc) { - if (desc != nullptr) { - // 使用 API 定义中的销毁函数 - INFINICORE_CHECK_ERROR(infiniopDestroyLogCumSumExpDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor y, Tensor x, int axis, bool exclusive, bool reverse) { - // 1. 生成 Hash Key：必须包含张量信息和算子特有参数 (axis, exclusive, reverse) - size_t seed = hash_combine(y, x, axis, exclusive, reverse); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - // 2. 获取当前设备对应的缓存实例 - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopLogCumSumExpDescriptor_t desc = nullptr; - - // 3. 如果缓存未命中，创建新的描述符 - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateLogCumSumExpDescriptor( - context::getInfiniopHandle(y->device()), - &desc, - y->desc(), // 输出张量描述符 - x->desc(), // 输入张量描述符 - axis, - exclusive, - reverse)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 4. 获取算子执行所需的 Workspace 大小并分配内存 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetLogCumSumExpWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - // 5. 执行算子 - INFINICORE_CHECK_ERROR(infiniopLogCumSumExp( - desc, - workspace->data(), - workspace_size, - y->data(), - x->data(), - context::getStream())); -} - -// 6. 自动注册：将此实现注册到 LogCumSumExp 的调度器中 -static bool registered = []() { - LogCumSumExp::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::logcumsumexp_impl::infiniop diff --git a/src/infinicore/ops/logdet/logdet.cc b/src/infinicore/ops/logdet/logdet.cc deleted file mode 100644 index 60d29a774..000000000 --- a/src/infinicore/ops/logdet/logdet.cc +++ /dev/null @@ -1,28 +0,0 @@ -#include "infinicore/ops/logdet.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Logdet); - -Logdet::Logdet(Tensor y, const Tensor &x) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, x); -} - -void Logdet::execute(Tensor y, const Tensor &x) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Logdet, y, x); -} - -Tensor logdet(const Tensor &x) { - auto y = Tensor::empty({}, x->dtype(), x->device()); - logdet_(y, x); - return y; -} - -void logdet_(Tensor y, const Tensor &x) { - Logdet::execute(y, x); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/logdet/logdet_infiniop.cc b/src/infinicore/ops/logdet/logdet_infiniop.cc deleted file mode 100644 index bb478ed19..000000000 --- a/src/infinicore/ops/logdet/logdet_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/logdet.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::logdet_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Logdet, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, y, x; -}; - -void *plan(Tensor y, const Tensor &x) { - size_t seed = hash_combine(y, x); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Logdet, - seed, - y->desc(), x->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Logdet, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(y), - graph::GraphTensor(x)}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopLogdet( - p->descriptor->desc, - p->workspace ? p->workspace->data() : nullptr, - p->workspace ? p->workspace->numel() : 0, - p->y->data(), - p->x->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Logdet, &plan, &run, &cleanup); - -} // namespace infinicore::op::logdet_impl::infiniop diff --git a/src/infinicore/ops/logical_and/logical_and.cc b/src/infinicore/ops/logical_and/logical_and.cc deleted file mode 100644 index 7b3196771..000000000 --- a/src/infinicore/ops/logical_and/logical_and.cc +++ /dev/null @@ -1,37 +0,0 @@ -#include "infinicore/ops/logical_and.hpp" -#include - -namespace infinicore::op { - -common::OpDispatcher &LogicalAnd::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void LogicalAnd::execute(Tensor output, Tensor input1, Tensor input2) { - // --- 修正点：去掉第二个参数 true --- - infinicore::context::setDevice(input1->device()); - - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No LogicalAnd implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input1, input2); -} - -Tensor logical_and(Tensor input1, Tensor input2) { - Shape shape = input1->shape(); - auto output = Tensor::empty(shape, input1->dtype(), input1->device()); - - logical_and_(output, input1, input2); - return output; -} - -void logical_and_(Tensor output, Tensor input1, Tensor input2) { - LogicalAnd::execute(output, input1, input2); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/logical_and/logical_and_cpu.cc b/src/infinicore/ops/logical_and/logical_and_cpu.cc deleted file mode 100644 index 19c1503b9..000000000 --- a/src/infinicore/ops/logical_and/logical_and_cpu.cc +++ /dev/null @@ -1,128 +0,0 @@ -#include "../../../utils.h" -#include "infinicore/device.hpp" -#include "infinicore/ops/logical_and.hpp" -#include "infinicore/tensor.hpp" -#include -#include -#include - -namespace infinicore::op::logical_and_impl::cpu { - -void calculate(Tensor output, Tensor input1, Tensor input2) { - auto ndim = output->ndim(); - auto numel = output->numel(); - auto shapes = output->shape(); - - auto strides1 = input1->strides(); - auto strides2 = input2->strides(); - auto out_strides = output->strides(); - - auto dtype = input1->dtype(); - auto dtype_size = input1->element_size(); - - auto out_dtype = output->dtype(); - auto out_dtype_size = output->element_size(); - - // 假设 Tensor::data() 返回的是支持字节加法的指针 (char* 或 uint8_t*) - // 如果是 void*，建议显式强转为 uint8_t* 或 char* - auto input1_base = reinterpret_cast(input1->data()); - auto input2_base = reinterpret_cast(input2->data()); - auto output_base = reinterpret_cast(output->data()); - - std::vector indices(ndim, 0); - - for (size_t idx = 0; idx < numel; ++idx) { - size_t offset1 = 0; - size_t offset2 = 0; - size_t out_offset = 0; - - for (size_t dim = 0; dim < ndim; ++dim) { - offset1 += indices[dim] * strides1[dim]; - offset2 += indices[dim] * strides2[dim]; - out_offset += indices[dim] * out_strides[dim]; - } - - bool result = false; - - // ========================================== - // INPUT TYPE DISPATCH (输入类型分发) - // ========================================== - - // 1. 浮点型 - if (dtype == DataType::F32) { - auto *p1 = reinterpret_cast(input1_base + offset1 * dtype_size); - auto *p2 = reinterpret_cast(input2_base + offset2 * dtype_size); - result = (*p1 != 0.0f) && (*p2 != 0.0f); - } else if (dtype == DataType::F64) { - auto *p1 = reinterpret_cast(input1_base + offset1 * dtype_size); - auto *p2 = reinterpret_cast(input2_base + offset2 * dtype_size); - result = (*p1 != 0.0) && (*p2 != 0.0); - } else if (dtype == DataType::F16) { - auto *p1 = reinterpret_cast(input1_base + offset1 * dtype_size); - auto *p2 = reinterpret_cast(input2_base + offset2 * dtype_size); - float v1 = utils::cast(*p1); - float v2 = utils::cast(*p2); - result = (v1 != 0.0f) && (v2 != 0.0f); - - // 2. 布尔与8位整型 - } else if (dtype == DataType::BOOL || dtype == DataType::U8) { - auto *p1 = reinterpret_cast(input1_base + offset1 * dtype_size); - auto *p2 = reinterpret_cast(input2_base + offset2 * dtype_size); - result = (*p1 != 0) && (*p2 != 0); - - // 3. 【新增】32位整型 (修复 int32 测试失败的关键！) - } else if (dtype == DataType::I32 || dtype == DataType::U32) { - // 无论是 I32 还是 U32，做非零判断逻辑是一样的，直接强转成 int32_t 读取即可 - auto *p1 = reinterpret_cast(input1_base + offset1 * dtype_size); - auto *p2 = reinterpret_cast(input2_base + offset2 * dtype_size); - result = (*p1 != 0) && (*p2 != 0); - - // 4. 【新增】64位整型 (增强健壮性) - } else if (dtype == DataType::I64 || dtype == DataType::U64) { - auto *p1 = reinterpret_cast(input1_base + offset1 * dtype_size); - auto *p2 = reinterpret_cast(input2_base + offset2 * dtype_size); - result = (*p1 != 0) && (*p2 != 0); - - // 5. 【新增】16位整型 (增强健壮性) - } else if (dtype == DataType::I16 || dtype == DataType::U16) { - auto *p1 = reinterpret_cast(input1_base + offset1 * dtype_size); - auto *p2 = reinterpret_cast(input2_base + offset2 * dtype_size); - result = (*p1 != 0) && (*p2 != 0); - - } else { - // 如果遇到 I8，可以合并到 U8 处理；如果没有，这里抛出异常是正确的 - throw std::runtime_error("Unsupported data type for logical_and operation."); - } - - // ========================================== - // OUTPUT TYPE DISPATCH (输出类型分发) - // ========================================== - if (out_dtype == DataType::BOOL || out_dtype == DataType::U8) { - auto *output_ptr = reinterpret_cast(output_base + out_offset * out_dtype_size); - *output_ptr = result ? 1 : 0; - } else if (out_dtype == DataType::F32) { - *reinterpret_cast(output_base + out_offset * out_dtype_size) = result ? 1.0f : 0.0f; - } else if (out_dtype == DataType::I32) { // 预防性增加对 int32 输出的支持 - *reinterpret_cast(output_base + out_offset * out_dtype_size) = result ? 1 : 0; - } else { - // 也可以选择在这里 throw，或者默认不做处理 - } - - // --- 维度索引递增逻辑 --- - for (ptrdiff_t dim = ndim - 1; dim >= 0; --dim) { - indices[dim]++; - if (indices[dim] < shapes[dim]) { - break; - } else { - indices[dim] = 0; - } - } - } -} - -static bool registered = []() { - LogicalAnd::dispatcher().registerDevice(Device::Type::CPU, &calculate); - return true; -}(); - -} // namespace infinicore::op::logical_and_impl::cpu diff --git a/src/infinicore/ops/logical_not/logical_not.cc b/src/infinicore/ops/logical_not/logical_not.cc deleted file mode 100644 index 161331a64..000000000 --- a/src/infinicore/ops/logical_not/logical_not.cc +++ /dev/null @@ -1,35 +0,0 @@ -#include "infinicore/ops/logical_not.hpp" -#include - -namespace infinicore::op { - -common::OpDispatcher & -LogicalNot::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -void LogicalNot::execute(Tensor output, Tensor input) { - infinicore::context::setDevice(input->device()); - - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - if (func == nullptr) { - throw std::runtime_error( - "No LogicalNot implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor logical_not(Tensor input) { - auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); - logical_not_(output, input); - return output; -} - -void logical_not_(Tensor output, Tensor input) { - LogicalNot::execute(output, input); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/logical_not/logical_not_cpu.cc b/src/infinicore/ops/logical_not/logical_not_cpu.cc deleted file mode 100644 index 713e36c37..000000000 --- a/src/infinicore/ops/logical_not/logical_not_cpu.cc +++ /dev/null @@ -1,171 +0,0 @@ -#include "../../../utils.h" -#include "infinicore/device.hpp" -#include "infinicore/ops/logical_not.hpp" -#include "infinicore/tensor.hpp" - -#include -#include // memcpy -#include -#include - -namespace infinicore::op::logical_not_impl::cpu { - -// ---------- safe load/store (avoid strict-aliasing & alignment UB) ---------- -template -static inline T load_scalar(const uint8_t *p) { - T v; - std::memcpy(&v, p, sizeof(T)); - return v; -} - -template -static inline void store_scalar(uint8_t *p, T v) { - std::memcpy(p, &v, sizeof(T)); -} - -// ---------- read "truthiness": treat nonzero as true ---------- -static inline bool read_truth(const uint8_t *base, ptrdiff_t elem_off, DataType dtype, size_t elem_size) { - const uint8_t *p = base + elem_off * static_cast(elem_size); - - switch (dtype) { - case DataType::BOOL: - case DataType::U8: - return load_scalar(p) != 0; - - case DataType::I16: - return load_scalar(p) != 0; - case DataType::U16: - return load_scalar(p) != 0; - - case DataType::I32: - return load_scalar(p) != 0; - case DataType::U32: - return load_scalar(p) != 0; - - case DataType::I64: - return load_scalar(p) != 0; - case DataType::U64: - return load_scalar(p) != 0; - - case DataType::F32: - return load_scalar(p) != 0.0f; - case DataType::F64: - return load_scalar(p) != 0.0; - - case DataType::F16: { - fp16_t h = load_scalar(p); - float fv = utils::cast(h); - return fv != 0.0f; - } - - default: - throw std::runtime_error("logical_not(cpu): unsupported input dtype."); - } -} - -// ---------- write result (bool -> out_dtype as 0/1) ---------- -static inline void write_bool(uint8_t *base, ptrdiff_t elem_off, DataType out_dtype, size_t out_elem_size, bool b) { - uint8_t *p = base + elem_off * static_cast(out_elem_size); - - switch (out_dtype) { - case DataType::BOOL: - case DataType::U8: - store_scalar(p, b ? 1 : 0); - return; - - case DataType::I16: - store_scalar(p, b ? 1 : 0); - return; - case DataType::U16: - store_scalar(p, b ? 1 : 0); - return; - - case DataType::I32: - store_scalar(p, b ? 1 : 0); - return; - case DataType::U32: - store_scalar(p, b ? 1u : 0u); - return; - - case DataType::I64: - store_scalar(p, b ? 1 : 0); - return; - case DataType::U64: - store_scalar(p, b ? 1ull : 0ull); - return; - - case DataType::F32: - store_scalar(p, b ? 1.0f : 0.0f); - return; - case DataType::F64: - store_scalar(p, b ? 1.0 : 0.0); - return; - - case DataType::F16: { - fp16_t h = utils::cast(b ? 1.0f : 0.0f); - store_scalar(p, h); - return; - } - - default: - throw std::runtime_error("logical_not(cpu): unsupported output dtype."); - } -} - -void calculate(Tensor output, Tensor input) { - const auto ndim = output->ndim(); - if (input->ndim() != ndim) { - throw std::runtime_error("logical_not(cpu): input/output ndim mismatch."); - } - - const auto numel = output->numel(); - const auto shape = output->shape(); - - // IMPORTANT: strides in TensorImpl are "element strides" (see tensor.cc calculate_contiguous_strides) - const auto in_strides = input->strides(); - const auto out_strides = output->strides(); - - const auto in_dtype = input->dtype(); - const auto out_dtype = output->dtype(); - - const auto in_es = input->element_size(); - const auto out_es = output->element_size(); - - const uint8_t *in_base = reinterpret_cast(input->data()); - uint8_t *out_base = reinterpret_cast(output->data()); - - std::vector idx(ndim, 0); - - for (size_t linear = 0; linear < numel; ++linear) { - // use ptrdiff_t all the way (must NOT cast strides to size_t) - ptrdiff_t in_elem_off = 0; - ptrdiff_t out_elem_off = 0; - - for (size_t d = 0; d < ndim; ++d) { - in_elem_off += static_cast(idx[d]) * in_strides[d]; - out_elem_off += static_cast(idx[d]) * out_strides[d]; - } - - // logical_not: result == (input == 0) == !(truthiness) - const bool truth = read_truth(in_base, in_elem_off, in_dtype, in_es); - const bool result = !truth; - - write_bool(out_base, out_elem_off, out_dtype, out_es, result); - - // increment multi-d index - for (ptrdiff_t d = static_cast(ndim) - 1; d >= 0; --d) { - idx[static_cast(d)]++; - if (idx[static_cast(d)] < shape[static_cast(d)]) { - break; - } - idx[static_cast(d)] = 0; - } - } -} - -static bool registered = []() { - LogicalNot::dispatcher().registerDevice(Device::Type::CPU, &calculate); - return true; -}(); - -} // namespace infinicore::op::logical_not_impl::cpu diff --git a/src/infinicore/ops/mamba_selective_scan/mamba_selective_scan.cc b/src/infinicore/ops/mamba_selective_scan/mamba_selective_scan.cc deleted file mode 100644 index b0df1815c..000000000 --- a/src/infinicore/ops/mamba_selective_scan/mamba_selective_scan.cc +++ /dev/null @@ -1,28 +0,0 @@ -#include "infinicore/ops/mamba_selective_scan.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MambaSelectiveScan); - -MambaSelectiveScan::MambaSelectiveScan(Tensor out, const Tensor &x, const Tensor &dt, - const Tensor &b, const Tensor &c, const Tensor &a_log, - const Tensor &d, const Tensor &gate, const Tensor &dt_bias, - Tensor state) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, x, dt, b, c, a_log, d, gate, dt_bias, state); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, x, dt, b, c, a_log, d, gate, dt_bias, state); -} -void MambaSelectiveScan::execute(Tensor out, const Tensor &x, const Tensor &dt, - const Tensor &b, const Tensor &c, const Tensor &a_log, - const Tensor &d, const Tensor &gate, const Tensor &dt_bias, - Tensor state) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(MambaSelectiveScan, out, x, dt, b, c, a_log, d, gate, dt_bias, state); -} -Tensor mamba_selective_scan(const Tensor &x, const Tensor &dt, const Tensor &b, const Tensor &c, const Tensor &a_log, const Tensor &d, const Tensor &gate, const Tensor &dt_bias, Tensor state) { - auto output = Tensor::empty(x->shape(), x->dtype(), x->device()); - mamba_selective_scan_(output, x, dt, b, c, a_log, d, gate, dt_bias, state); - return output; -} -void mamba_selective_scan_(Tensor out, const Tensor &x, const Tensor &dt, const Tensor &b, const Tensor &c, const Tensor &a_log, const Tensor &d, const Tensor &gate, const Tensor &dt_bias, Tensor state) { - MambaSelectiveScan::execute(out, x, dt, b, c, a_log, d, gate, dt_bias, state); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/mamba_selective_scan/mamba_selective_scan_infiniop.cc b/src/infinicore/ops/mamba_selective_scan/mamba_selective_scan_infiniop.cc deleted file mode 100644 index a64ae6c51..000000000 --- a/src/infinicore/ops/mamba_selective_scan/mamba_selective_scan_infiniop.cc +++ /dev/null @@ -1,25 +0,0 @@ -#include "../infiniop_impl.hpp" -#include "infinicore/ops/mamba_selective_scan.hpp" - -namespace infinicore::op::mamba_selective_scan_impl::infiniop { -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, MambaSelectiveScan, 100); -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, x, dt, b, c, a_log, d, gate, dt_bias, state; -}; -void *plan(Tensor out, const Tensor &x, const Tensor &dt, const Tensor &b, const Tensor &c, const Tensor &a_log, const Tensor &d, const Tensor &gate, const Tensor &dt_bias, Tensor state) { - size_t seed = hash_combine(out, x, dt, b, c, a_log, d, gate, dt_bias, state); - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE(Descriptor, descriptor, MambaSelectiveScan, seed, out->desc(), x->desc(), dt->desc(), b->desc(), c->desc(), a_log->desc(), d->desc(), gate->desc(), dt_bias->desc(), state->desc()); - INFINIOP_WORKSPACE_TENSOR(workspace, MambaSelectiveScan, descriptor); - return new PlannedMeta{descriptor, graph::GraphTensor(workspace), graph::GraphTensor(out), graph::GraphTensor(x), graph::GraphTensor(dt), graph::GraphTensor(b), graph::GraphTensor(c), graph::GraphTensor(a_log), graph::GraphTensor(d), graph::GraphTensor(gate), graph::GraphTensor(dt_bias), graph::GraphTensor(state)}; -} -void run(void *planned_meta) { - auto p = reinterpret_cast(planned_meta); - INFINICORE_CHECK_ERROR(infiniopMambaSelectiveScan(p->descriptor->desc, p->workspace->data(), p->workspace->numel(), p->out->data(), p->x->data(), p->dt->data(), p->b->data(), p->c->data(), p->a_log->data(), p->d->data(), p->gate->data(), p->dt_bias->data(), p->state->data(), context::getStream())); -} -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MambaSelectiveScan, &plan, &run, &cleanup); -} // namespace infinicore::op::mamba_selective_scan_impl::infiniop diff --git a/src/infinicore/ops/masked_select/masked_select.cc b/src/infinicore/ops/masked_select/masked_select.cc deleted file mode 100644 index 1ade25a09..000000000 --- a/src/infinicore/ops/masked_select/masked_select.cc +++ /dev/null @@ -1,33 +0,0 @@ -#include -#include - -namespace infinicore::op { - -common::OpDispatcher &MaskedSelect::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void MaskedSelect::execute(Tensor input, Tensor mask, void **data_ptr, size_t *dlen_ptr) { - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No MaskedSelect implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(input, mask, data_ptr, dlen_ptr); -} - -Tensor masked_select(Tensor input, Tensor mask) { - - std::byte *data; - size_t dlen; - MaskedSelect::execute(input, mask, (void **)&data, &dlen); - - auto out = Tensor::from_blob(data, {dlen}, input->dtype(), input->device()); - - return out; -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/masked_select/masked_select_infiniop.cc b/src/infinicore/ops/masked_select/masked_select_infiniop.cc deleted file mode 100644 index befef2d9e..000000000 --- a/src/infinicore/ops/masked_select/masked_select_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/masked_select.hpp" -#include - -namespace infinicore::op::masked_select_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopMaskedSelectDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyMaskedSelectDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor input, Tensor mask, void **data_ptr, size_t *dlen_ptr) { - size_t seed = hash_combine(input, mask, (std::uintptr_t)data_ptr, (std::uintptr_t)dlen_ptr); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopMaskedSelectDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateMaskedSelectDescriptor( - context::getInfiniopHandle(input->device()), &desc, - input->desc(), mask->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetMaskedSelectWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopMaskedSelect( - desc, workspace->data(), workspace_size, - input->data(), (const bool *)mask->data(), data_ptr, dlen_ptr, context::getStream())); -} - -static bool registered = []() { - MaskedSelect::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::masked_select_impl::infiniop diff --git a/src/infinicore/ops/matmul/matmul.cc b/src/infinicore/ops/matmul/matmul.cc deleted file mode 100644 index 58119bc10..000000000 --- a/src/infinicore/ops/matmul/matmul.cc +++ /dev/null @@ -1,13 +0,0 @@ -#include "infinicore/ops/matmul.hpp" -#include "infinicore/ops/gemm.hpp" - -namespace infinicore::op { - -Tensor matmul(Tensor a, Tensor b, float alpha) { - return gemm(a, b, alpha, 0.0f); -} - -void matmul_(Tensor c, Tensor a, Tensor b, float alpha) { - Gemm::execute(c, a, b, alpha, 0.0f); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc b/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc deleted file mode 100644 index 9d2cdc2b5..000000000 --- a/src/infinicore/ops/mha_kvcache/hygon/mha_kvcache_flashattn_hygon.cc +++ /dev/null @@ -1,167 +0,0 @@ -#if defined(ENABLE_HYGON_API) && defined(ENABLE_FLASH_ATTN) -#include "infinicore/ops/mha_kvcache.hpp" - -#include "../../../adaptor/flash_attn/hygon/flash_attn_hygon.hpp" -#include "infinicore/adaptor/aten_adaptor.hpp" - -#include -#include -#include -#include - -#include -#include - -namespace infinicore::op::mha_kvcache_impl::flashattn { - -struct PlannedMeta { - graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; - std::optional alibi_slopes; - float scale; -}; - -void *plan(Tensor out, - const Tensor &q, - const Tensor &k_cache, - const Tensor &v_cache, - const Tensor &seqlens_k, - const Tensor &block_table, - std::optional alibi_slopes, - float scale) { - return new PlannedMeta{ - graph::GraphTensor(out), - graph::GraphTensor(q), - graph::GraphTensor(k_cache), - graph::GraphTensor(v_cache), - graph::GraphTensor(seqlens_k), - graph::GraphTensor(block_table), - alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, - scale}; -} - -void run(void *planned_meta) { - c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); - auto *p = reinterpret_cast(planned_meta); - - // Paged KV caches must be contiguous for flash-attn; avoid extra copies for q/metadata when already dense. - const bool out_need_copy_back = !p->out->is_contiguous(); - Tensor out_work = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); - auto out_tensor = infinicore::adaptor::to_aten_tensor(out_work); - auto q = infinicore::adaptor::to_aten_tensor(p->q); - auto k_cache = infinicore::adaptor::to_aten_tensor(p->k_cache); - auto v_cache = infinicore::adaptor::to_aten_tensor(p->v_cache); - auto seqlens_k = std::optional(infinicore::adaptor::to_aten_tensor(p->seqlens_k)); - auto block_table = std::optional(infinicore::adaptor::to_aten_tensor(p->block_table)); - auto alibi_slopes = p->alibi_slopes - ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) - : std::nullopt; - - if (std::getenv("INFINICORE_HYGON_ATEN_FALLBACK")) { - namespace idx = at::indexing; - auto seqlens_t = infinicore::adaptor::to_aten_tensor(p->seqlens_k); - auto block_table_t = infinicore::adaptor::to_aten_tensor(p->block_table); - auto seqlens_cpu = seqlens_t.to(at::kCPU); - auto block_table_cpu = block_table_t.to(at::kCPU); - - auto result = at::empty_like(out_tensor); - const int64_t batch_size = q.size(0); - const int64_t seqlen_q = q.size(1); - const int64_t num_heads = q.size(2); - const int64_t block_size = k_cache.size(1); - const int64_t num_kv_heads = k_cache.size(2); - const int64_t group_size = num_heads / num_kv_heads; - - for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { - const int64_t seq_len = seqlens_cpu.index({batch_idx}).item(); - std::vector keys; - std::vector values; - keys.reserve(seq_len); - values.reserve(seq_len); - for (int64_t logical_pos = 0; logical_pos < seq_len; ++logical_pos) { - const int64_t block_id = block_table_cpu.index({batch_idx, logical_pos / block_size}).item(); - const int64_t off = logical_pos % block_size; - keys.push_back(k_cache.index({block_id, off, idx::Slice(), idx::Slice()})); - values.push_back(v_cache.index({block_id, off, idx::Slice(), idx::Slice()})); - } - auto K = at::stack(keys, 0); - auto V = at::stack(values, 0); - if (group_size > 1) { - K = K.repeat_interleave(group_size, 1); - V = V.repeat_interleave(group_size, 1); - } - auto cur_q = q.index({batch_idx}); - auto scores = at::matmul(cur_q.permute({1, 0, 2}).to(at::kFloat), K.permute({1, 2, 0}).to(at::kFloat)) * p->scale; - auto mask = at::full({seqlen_q, seq_len}, -std::numeric_limits::infinity(), q.options().dtype(at::kFloat)); - const int64_t prefix_len = seq_len - seqlen_q; - for (int64_t query_pos = 0; query_pos < seqlen_q; ++query_pos) { - mask.index_put_({query_pos, idx::Slice(0, prefix_len + query_pos + 1)}, 0.0); - } - auto attn = at::softmax(scores + mask.unsqueeze(0), -1).to(q.dtype()); - auto cur_out = at::matmul(attn, V.permute({1, 0, 2})).permute({1, 0, 2}); - result.index_put_({batch_idx}, cur_out); - } - - out_tensor.copy_(result); - if (out_need_copy_back) { - p->out->copy_from(out_work); - } - return; - } - - std::optional k_new = std::nullopt; - std::optional v_new = std::nullopt; - std::optional rotary_cos = std::nullopt; - std::optional rotary_sin = std::nullopt; - std::optional cache_batch_idx = std::nullopt; - std::optional leftpad_k = std::nullopt; - const bool use_dynamic_out = q.dim() == 4 && k_cache.dim() == 4 - && q.size(1) == 1 && q.size(2) > k_cache.size(2) - && q.size(3) % 8 == 0 && !alibi_slopes.has_value(); - - auto out = use_dynamic_out ? std::optional(std::nullopt) - : std::optional(out_tensor); - - auto result = flash::mha_fwd_kvcache( - q, - k_cache, - v_cache, - k_new, - v_new, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - p->scale, - true, - -1, - -1, - 0.0f, - false, - 0); - - if (!result.empty() && result[0].defined()) { - out_tensor.copy_(result[0]); - } - if (out_need_copy_back) { - p->out->copy_from(out_work); - } -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - MhaKVCache::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); - MhaKVCache::run_dispatcher().registerDevice(Device::Type::HYGON, &run); - MhaKVCache::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); - return true; -}(); - -} // namespace infinicore::op::mha_kvcache_impl::flashattn -#endif // ENABLE_HYGON_API && ENABLE_FLASH_ATTN diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache.cc deleted file mode 100644 index 0c5b3ae8c..000000000 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache.cc +++ /dev/null @@ -1,58 +0,0 @@ -#include "infinicore/ops/mha_kvcache.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MhaKVCache); - -MhaKVCache::MhaKVCache(Tensor out, - const Tensor &q, - const Tensor &k_cache, - const Tensor &v_cache, - const Tensor &seqlens_k, - const Tensor &block_table, - std::optional alibi_slopes, - float scale) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, seqlens_k, block_table); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), - out, q, k_cache, v_cache, seqlens_k, block_table, alibi_slopes, scale); -} - -void MhaKVCache::execute(Tensor out, - const Tensor &q, - const Tensor &k_cache, - const Tensor &v_cache, - const Tensor &seqlens_k, - const Tensor &block_table, - std::optional alibi_slopes, - float scale) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - MhaKVCache, - out, q, k_cache, v_cache, seqlens_k, block_table, alibi_slopes, scale); -} - -void mha_kvcache_(Tensor out, - const Tensor &q, - const Tensor &k_cache, - const Tensor &v_cache, - const Tensor &seqlens_k, - const Tensor &block_table, - std::optional alibi_slopes, - float scale) { - MhaKVCache::execute(out, q, k_cache, v_cache, seqlens_k, block_table, alibi_slopes, scale); -} - -Tensor mha_kvcache(const Tensor &q, - const Tensor &k_cache, - const Tensor &v_cache, - const Tensor &seqlens_k, - const Tensor &block_table, - std::optional alibi_slopes, - float scale) { - // Output shape matches q: [batch_size, seqlen_q, num_heads, head_size] - auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); - mha_kvcache_(out, q, k_cache, v_cache, seqlens_k, block_table, alibi_slopes, scale); - return out; -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc deleted file mode 100644 index 0167c17df..000000000 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn.cc +++ /dev/null @@ -1,136 +0,0 @@ -#include "infinicore/ops/mha_kvcache.hpp" - -#include "infinicore/adaptor/flash_attention_adaptor.hpp" - -#include - -#ifdef ENABLE_FLASH_ATTN -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) -#include -#endif -#endif - -#if defined(ENABLE_METAX_API) -#define INFINICORE_FLASH_OP(name) ::name -#else -#define INFINICORE_FLASH_OP(name) flash::name -#endif - -namespace infinicore::op::mha_kvcache_impl::flashattn { - -struct PlannedMeta { - graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; - std::optional alibi_slopes; - float scale; -}; - -void *plan(Tensor out, - const Tensor &q, - const Tensor &k_cache, - const Tensor &v_cache, - const Tensor &seqlens_k, - const Tensor &block_table, - std::optional alibi_slopes, - float scale) { - return new PlannedMeta{ - graph::GraphTensor(out), - graph::GraphTensor(q), - graph::GraphTensor(k_cache), - graph::GraphTensor(v_cache), - graph::GraphTensor(seqlens_k), - graph::GraphTensor(block_table), - alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, - scale}; -} - -void run(void *planned_meta) { -#ifdef ENABLE_FLASH_ATTN -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) - c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream()); -#endif - auto *p = reinterpret_cast(planned_meta); - - // Paged KV caches must be contiguous for flash-attn; avoid extra copies for q/metadata when already dense. - const bool out_need_copy_back = !p->out->is_contiguous(); - Tensor out_work = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); - auto out_tensor = infinicore::adaptor::to_aten_tensor(out_work); - auto q = infinicore::adaptor::to_aten_tensor(p->q); -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) - auto k_cache = infinicore::adaptor::to_aten_tensor(p->k_cache); - auto v_cache = infinicore::adaptor::to_aten_tensor(p->v_cache); -#elif defined(ENABLE_QY_API) - Tensor k_cache_work = p->k_cache->contiguous(); - Tensor v_cache_work = p->v_cache->contiguous(); - auto k_cache = infinicore::adaptor::to_aten_tensor(k_cache_work); - auto v_cache = infinicore::adaptor::to_aten_tensor(v_cache_work); -#endif - auto seqlens_k = std::optional(infinicore::adaptor::to_aten_tensor(p->seqlens_k)); - auto block_table = std::optional(infinicore::adaptor::to_aten_tensor(p->block_table)); - auto alibi_slopes = p->alibi_slopes - ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) - : std::nullopt; - - std::optional k_new = std::nullopt; - std::optional v_new = std::nullopt; - std::optional rotary_cos = std::nullopt; - std::optional rotary_sin = std::nullopt; - std::optional cache_batch_idx = std::nullopt; - std::optional leftpad_k = std::nullopt; - - const bool use_dynamic_out = q.dim() == 4 && k_cache.dim() == 4 - && q.size(1) == 1 && q.size(2) > k_cache.size(2) - && q.size(3) % 8 == 0 && !alibi_slopes.has_value(); - - auto out = use_dynamic_out ? std::optional(std::nullopt) - : std::optional(out_tensor); - -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - std::optional flash_attn_mars_ext = std::nullopt; -#endif - - auto result = INFINICORE_FLASH_OP(mha_fwd_kvcache)( - q, - k_cache, - v_cache, - k_new, - v_new, - seqlens_k, - rotary_cos, - rotary_sin, - cache_batch_idx, - leftpad_k, - block_table, - alibi_slopes, - out, - p->scale, - true, - -1, - -1, - 0.0f, - false, - 0 -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - , - flash_attn_mars_ext -#endif - ); - - if (use_dynamic_out) { - out_tensor.copy_(result[0]); - } - if (out_need_copy_back) { - p->out->copy_from(out_work); - } -#else - throw std::runtime_error("FlashAttention is not enabled in this build"); -#endif -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MhaKVCache, &plan, &run, &cleanup); - -} // namespace infinicore::op::mha_kvcache_impl::flashattn diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_moore.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_moore.cc deleted file mode 100644 index ddaabb937..000000000 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache_flashattn_moore.cc +++ /dev/null @@ -1,127 +0,0 @@ -#if defined(ENABLE_MOORE_MATE_FLASH_ATTN) - -#include "infinicore/ops/mha_kvcache.hpp" - -#include "infinicore/adaptor/aten_adaptor.hpp" - -#include -#include -#include -#include - -namespace infinicore::op::mha_kvcache_impl::flashattn_moore { - -namespace py = pybind11; - -// Lightweight RAII: Binds MUSA streams, -// avoiding the need to include -namespace { -class LocalMUSAStreamGuard { -public: - explicit LocalMUSAStreamGuard(const c10::musa::MUSAStream &s) - : prev_(c10::musa::getCurrentMUSAStream(s.device_index())) { - c10::musa::setCurrentMUSAStream(s); - } - ~LocalMUSAStreamGuard() { - c10::musa::setCurrentMUSAStream(prev_); - } - LocalMUSAStreamGuard(const LocalMUSAStreamGuard &) = delete; - LocalMUSAStreamGuard &operator=(const LocalMUSAStreamGuard &) = delete; - -private: - c10::musa::MUSAStream prev_; -}; -} // namespace - -struct PlannedMeta { - graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; - std::optional alibi_slopes; - float scale; -}; - -void *plan(Tensor out, - const Tensor &q, - const Tensor &k_cache, - const Tensor &v_cache, - const Tensor &seqlens_k, - const Tensor &block_table, - std::optional alibi_slopes, - float scale) { - return new PlannedMeta{ - graph::GraphTensor(out), - graph::GraphTensor(q), - graph::GraphTensor(k_cache), - graph::GraphTensor(v_cache), - graph::GraphTensor(seqlens_k), - graph::GraphTensor(block_table), - alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, - scale}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - if (p->alibi_slopes.has_value()) { - throw std::runtime_error( - "[mha_kvcache/moore] ALiBi not supported by mate flash_attn_with_kvcache"); - } - - LocalMUSAStreamGuard guard(infinicore::adaptor::get_musa_stream()); - - auto out_tensor = infinicore::adaptor::to_aten_tensor(p->out); - auto q_4d = infinicore::adaptor::to_aten_tensor(p->q); - auto k_cache = infinicore::adaptor::to_aten_tensor(p->k_cache); - auto v_cache = infinicore::adaptor::to_aten_tensor(p->v_cache); - auto seqlens_k = infinicore::adaptor::to_aten_tensor(p->seqlens_k); - auto block_table = infinicore::adaptor::to_aten_tensor(p->block_table); - - auto q_3d = q_4d.squeeze(1); - - const int64_t block_size = k_cache.size(1); - const int64_t max_seq_len = block_table.size(1) * block_size; - - try { - py::gil_scoped_acquire gil; - py::module_ wrapper = py::module_::import("infinicore.ops.moore_mate_flash_attn"); - - py::object py_q = py::cast(q_3d); - py::object py_k_cache = py::cast(k_cache); - py::object py_v_cache = py::cast(v_cache); - py::object py_seqlens_k = py::cast(seqlens_k); - py::object py_blk_tbl = py::cast(block_table); - - py::object result = wrapper.attr("moore_mate_flash_attn_decode")( - py_q, - py_k_cache, - py_v_cache, - py_blk_tbl, - py_seqlens_k, - p->scale, - block_size, - max_seq_len); - - at::Tensor result_t = result.cast(); - out_tensor.copy_(result_t.unsqueeze(1)); - - result = py::none(); - py_q = py_k_cache = py_v_cache = py_seqlens_k = py_blk_tbl = py::none(); - } catch (const py::error_already_set &e) { - throw std::runtime_error( - std::string("[mha_kvcache/moore] Python error: ") + e.what()); - } -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - MhaKVCache::plan_dispatcher().registerDevice(Device::Type::MOORE, &plan); - MhaKVCache::run_dispatcher().registerDevice(Device::Type::MOORE, &run); - MhaKVCache::cleanup_dispatcher().registerDevice(Device::Type::MOORE, &cleanup); - return true; -}(); - -} // namespace infinicore::op::mha_kvcache_impl::flashattn_moore - -#endif // ENABLE_MOORE_MATE_FLASH_ATTN diff --git a/src/infinicore/ops/moe_align/moe_align.cc b/src/infinicore/ops/moe_align/moe_align.cc deleted file mode 100644 index 0874a2f63..000000000 --- a/src/infinicore/ops/moe_align/moe_align.cc +++ /dev/null @@ -1,155 +0,0 @@ -#include "infinicore/ops/moe_align.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeAlign); -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeAlignWithExpertMap); - -MoeAlign::MoeAlign(Tensor sorted_token_ids, - Tensor expert_ids, - Tensor num_tokens_post_padded, - const Tensor &topk_ids, - const size_t num_experts, - const size_t block_size, - const bool pad_sorted_token_ids) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(sorted_token_ids, expert_ids, num_tokens_post_padded, topk_ids); - INFINICORE_GRAPH_OP_DISPATCH( - sorted_token_ids->device().getType(), - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - topk_ids, - num_experts, - block_size, - pad_sorted_token_ids); -} - -void MoeAlign::execute(Tensor sorted_token_ids, - Tensor expert_ids, - Tensor num_tokens_post_padded, - const Tensor &topk_ids, - const size_t num_experts, - const size_t block_size, - const bool pad_sorted_token_ids) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - MoeAlign, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - topk_ids, - num_experts, - block_size, - pad_sorted_token_ids); -} - -MoeAlignWithExpertMap::MoeAlignWithExpertMap(Tensor sorted_token_ids, - Tensor expert_ids, - Tensor num_tokens_post_padded, - const Tensor &topk_ids, - const Tensor &expert_map, - const size_t num_experts, - const size_t block_size, - const bool pad_sorted_token_ids) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(sorted_token_ids, expert_ids, num_tokens_post_padded, topk_ids, expert_map); - INFINICORE_GRAPH_OP_DISPATCH( - sorted_token_ids->device().getType(), - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - topk_ids, - expert_map, - num_experts, - block_size, - pad_sorted_token_ids); -} - -void MoeAlignWithExpertMap::execute(Tensor sorted_token_ids, - Tensor expert_ids, - Tensor num_tokens_post_padded, - const Tensor &topk_ids, - const Tensor &expert_map, - const size_t num_experts, - const size_t block_size, - const bool pad_sorted_token_ids) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - MoeAlignWithExpertMap, - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - topk_ids, - expert_map, - num_experts, - block_size, - pad_sorted_token_ids); -} - -std::tuple moe_align( - const Tensor &topk_ids, - size_t num_experts, - size_t block_size, - bool pad_sorted_token_ids) { - auto shape = topk_ids->shape(); - INFINICORE_ASSERT(shape.size() == 2); - const size_t numel = shape[0] * shape[1]; - const size_t align_num_experts = num_experts + 1; - const size_t max_num_tokens_padded = numel < align_num_experts - ? numel * block_size - : numel + align_num_experts * (block_size - 1); - const size_t sorted_token_ids_capacity = ((max_num_tokens_padded + 3) / 4) * 4; - const size_t max_num_blocks = (max_num_tokens_padded + block_size - 1) / block_size; - - auto sorted_token_ids = Tensor::empty({sorted_token_ids_capacity}, DataType::I32, topk_ids->device()); - auto expert_ids = Tensor::empty({max_num_blocks}, DataType::I32, topk_ids->device()); - auto num_tokens_post_padded = Tensor::empty({1}, DataType::I32, topk_ids->device()); - - moe_align_( - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - topk_ids, - num_experts, - block_size, - pad_sorted_token_ids); - - return {sorted_token_ids, expert_ids, num_tokens_post_padded}; -} - -void moe_align_(Tensor sorted_token_ids, - Tensor expert_ids, - Tensor num_tokens_post_padded, - const Tensor &topk_ids, - size_t num_experts, - size_t block_size, - bool pad_sorted_token_ids) { - MoeAlign::execute( - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - topk_ids, - num_experts, - block_size, - pad_sorted_token_ids); -} - -void moe_align_with_expert_map_(Tensor sorted_token_ids, - Tensor expert_ids, - Tensor num_tokens_post_padded, - const Tensor &topk_ids, - const Tensor &expert_map, - size_t num_experts, - size_t block_size, - bool pad_sorted_token_ids) { - MoeAlignWithExpertMap::execute( - sorted_token_ids, - expert_ids, - num_tokens_post_padded, - topk_ids, - expert_map, - num_experts, - block_size, - pad_sorted_token_ids); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/moe_align/moe_align_infiniop.cc b/src/infinicore/ops/moe_align/moe_align_infiniop.cc deleted file mode 100644 index 812f64bda..000000000 --- a/src/infinicore/ops/moe_align/moe_align_infiniop.cc +++ /dev/null @@ -1,174 +0,0 @@ -#include "infinicore/ops/moe_align.hpp" - -#include - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::moe_align_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, MoeAlign, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace; - graph::GraphTensor sorted_token_ids; - graph::GraphTensor expert_ids; - graph::GraphTensor num_tokens_post_padded; - graph::GraphTensor topk_ids; - bool pad_sorted_token_ids; -}; - -void *plan(Tensor sorted_token_ids, - Tensor expert_ids, - Tensor num_tokens_post_padded, - const Tensor &topk_ids, - const size_t num_experts, - const size_t block_size, - const bool pad_sorted_token_ids) { - size_t seed = hash_combine(sorted_token_ids, expert_ids, num_tokens_post_padded, topk_ids, num_experts, block_size); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, MoeAlign, - seed, - sorted_token_ids->desc(), - expert_ids->desc(), - num_tokens_post_padded->desc(), - topk_ids->desc(), - num_experts, - block_size); - - INFINIOP_WORKSPACE_TENSOR(workspace, MoeAlign, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(sorted_token_ids), - graph::GraphTensor(expert_ids), - graph::GraphTensor(num_tokens_post_padded), - graph::GraphTensor(topk_ids), - pad_sorted_token_ids}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopMoeAlign( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->sorted_token_ids->data(), - planned->expert_ids->data(), - planned->num_tokens_post_padded->data(), - planned->topk_ids->data(), - nullptr, - static_cast(planned->pad_sorted_token_ids), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MoeAlign, &plan, &run, cleanup); - -} // namespace infinicore::op::moe_align_impl::infiniop - -namespace infinicore::op::moe_align_with_expert_map_impl::infiniop { - -struct Descriptor { - infiniopMoeAlignDescriptor_t desc = nullptr; - - explicit Descriptor(infiniopMoeAlignDescriptor_t d) - : desc(d) {} - - Descriptor(const Descriptor &) = delete; - Descriptor &operator=(const Descriptor &) = delete; - - ~Descriptor() { - if (desc != nullptr) { - infiniopDestroyMoeAlignDescriptor(desc); - } - } -}; - -thread_local common::OpCache> caches( - 100, - [](std::shared_ptr &desc) { - desc = nullptr; - }); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace; - graph::GraphTensor sorted_token_ids; - graph::GraphTensor expert_ids; - graph::GraphTensor num_tokens_post_padded; - graph::GraphTensor topk_ids; - graph::GraphTensor expert_map; - bool pad_sorted_token_ids; -}; - -void *plan(Tensor sorted_token_ids, - Tensor expert_ids, - Tensor num_tokens_post_padded, - const Tensor &topk_ids, - const Tensor &expert_map, - const size_t num_experts, - const size_t block_size, - const bool pad_sorted_token_ids) { - size_t seed = hash_combine(sorted_token_ids, expert_ids, num_tokens_post_padded, topk_ids, expert_map, num_experts, block_size); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, MoeAlign, - seed, - sorted_token_ids->desc(), - expert_ids->desc(), - num_tokens_post_padded->desc(), - topk_ids->desc(), - num_experts, - block_size); - - Tensor workspace; - { - auto device = context::getDevice(); - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetMoeAlignWorkspaceSize(descriptor->desc, &workspace_size)); - workspace = Tensor::empty({workspace_size}, DataType::U8, device); - } - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(sorted_token_ids), - graph::GraphTensor(expert_ids), - graph::GraphTensor(num_tokens_post_padded), - graph::GraphTensor(topk_ids), - graph::GraphTensor(expert_map), - pad_sorted_token_ids}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopMoeAlign( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->sorted_token_ids->data(), - planned->expert_ids->data(), - planned->num_tokens_post_padded->data(), - planned->topk_ids->data(), - planned->expert_map->data(), - static_cast(planned->pad_sorted_token_ids), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MoeAlignWithExpertMap, &plan, &run, cleanup); - -} // namespace infinicore::op::moe_align_with_expert_map_impl::infiniop diff --git a/src/infinicore/ops/moe_fused_dense/moe_fused_dense.cc b/src/infinicore/ops/moe_fused_dense/moe_fused_dense.cc deleted file mode 100644 index 81a6983e6..000000000 --- a/src/infinicore/ops/moe_fused_dense/moe_fused_dense.cc +++ /dev/null @@ -1,73 +0,0 @@ -#include "infinicore/ops/moe_fused_dense.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeFusedDense); - -MoeFusedDense::MoeFusedDense(Tensor output, - const Tensor &hidden_states, - const Tensor &w13, - const Tensor &w2, - const Tensor &topk_weights, - const Tensor &topk_ids, - const Tensor &sorted_token_ids, - const Tensor &expert_ids, - const Tensor &num_tokens_post_padded) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE( - output, hidden_states, w13, w2, topk_weights, topk_ids, - sorted_token_ids, expert_ids, num_tokens_post_padded); - INFINICORE_GRAPH_OP_DISPATCH( - output->device().getType(), output, hidden_states, w13, w2, topk_weights, topk_ids, - sorted_token_ids, expert_ids, num_tokens_post_padded); -} - -void MoeFusedDense::execute(Tensor output, - const Tensor &hidden_states, - const Tensor &w13, - const Tensor &w2, - const Tensor &topk_weights, - const Tensor &topk_ids, - const Tensor &sorted_token_ids, - const Tensor &expert_ids, - const Tensor &num_tokens_post_padded) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - MoeFusedDense, output, hidden_states, w13, w2, topk_weights, topk_ids, - sorted_token_ids, expert_ids, num_tokens_post_padded); -} - -Tensor moe_fused_dense( - const Tensor &hidden_states, - const Tensor &w13, - const Tensor &w2, - const Tensor &topk_weights, - const Tensor &topk_ids, - const Tensor &sorted_token_ids, - const Tensor &expert_ids, - const Tensor &num_tokens_post_padded) { - auto shape = hidden_states->shape(); - INFINICORE_ASSERT(shape.size() == 2); - auto output = Tensor::empty(shape, hidden_states->dtype(), hidden_states->device()); - moe_fused_dense_( - output, hidden_states, w13, w2, topk_weights, topk_ids, - sorted_token_ids, expert_ids, num_tokens_post_padded); - return output; -} - -void moe_fused_dense_( - Tensor output, - const Tensor &hidden_states, - const Tensor &w13, - const Tensor &w2, - const Tensor &topk_weights, - const Tensor &topk_ids, - const Tensor &sorted_token_ids, - const Tensor &expert_ids, - const Tensor &num_tokens_post_padded) { - MoeFusedDense::execute( - output, hidden_states, w13, w2, topk_weights, topk_ids, - sorted_token_ids, expert_ids, num_tokens_post_padded); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/moe_fused_dense/moe_fused_dense_infiniop.cc b/src/infinicore/ops/moe_fused_dense/moe_fused_dense_infiniop.cc deleted file mode 100644 index fa6fcdfa7..000000000 --- a/src/infinicore/ops/moe_fused_dense/moe_fused_dense_infiniop.cc +++ /dev/null @@ -1,90 +0,0 @@ -#include "infinicore/ops/moe_fused_dense.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::moe_fused_dense_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, MoeFusedDense, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace; - graph::GraphTensor output; - graph::GraphTensor hidden_states; - graph::GraphTensor w13; - graph::GraphTensor w2; - graph::GraphTensor topk_weights; - graph::GraphTensor topk_ids; - graph::GraphTensor sorted_token_ids; - graph::GraphTensor expert_ids; - graph::GraphTensor num_tokens_post_padded; -}; - -void *plan(Tensor output, - const Tensor &hidden_states, - const Tensor &w13, - const Tensor &w2, - const Tensor &topk_weights, - const Tensor &topk_ids, - const Tensor &sorted_token_ids, - const Tensor &expert_ids, - const Tensor &num_tokens_post_padded) { - size_t seed = hash_combine( - output, hidden_states, w13, w2, topk_weights, topk_ids, - sorted_token_ids, expert_ids, num_tokens_post_padded); - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, - descriptor, - MoeFusedDense, - seed, - output->desc(), - hidden_states->desc(), - w13->desc(), - w2->desc(), - topk_weights->desc(), - topk_ids->desc(), - sorted_token_ids->desc(), - expert_ids->desc(), - num_tokens_post_padded->desc()); - INFINIOP_WORKSPACE_TENSOR(workspace, MoeFusedDense, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(output), - graph::GraphTensor(hidden_states), - graph::GraphTensor(w13), - graph::GraphTensor(w2), - graph::GraphTensor(topk_weights), - graph::GraphTensor(topk_ids), - graph::GraphTensor(sorted_token_ids), - graph::GraphTensor(expert_ids), - graph::GraphTensor(num_tokens_post_padded)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - INFINICORE_CHECK_ERROR(infiniopMoeFusedDense( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->output->data(), - planned->hidden_states->data(), - planned->w13->data(), - planned->w2->data(), - planned->topk_weights->data(), - planned->topk_ids->data(), - planned->sorted_token_ids->data(), - planned->expert_ids->data(), - planned->num_tokens_post_padded->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MoeFusedDense, &plan, &run, cleanup); - -} // namespace infinicore::op::moe_fused_dense_impl::infiniop diff --git a/src/infinicore/ops/moe_fused_gate/moe_fused_gate.cc b/src/infinicore/ops/moe_fused_gate/moe_fused_gate.cc deleted file mode 100644 index a006a2cad..000000000 --- a/src/infinicore/ops/moe_fused_gate/moe_fused_gate.cc +++ /dev/null @@ -1,101 +0,0 @@ -#include "infinicore/ops/moe_fused_gate.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeFusedGate); - -MoeFusedGate::MoeFusedGate(Tensor topk_weights, - Tensor topk_indices, - const Tensor &input, - const Tensor &bias, - const size_t num_expert_group, - const size_t topk_group, - const size_t num_fused_shared_experts, - const float routed_scaling_factor, - const bool apply_routed_scaling_factor_on_output) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_indices, input, bias); - INFINICORE_GRAPH_OP_DISPATCH( - topk_weights->device().getType(), - topk_weights, - topk_indices, - input, - bias, - num_expert_group, - topk_group, - num_fused_shared_experts, - routed_scaling_factor, - apply_routed_scaling_factor_on_output); -} - -void MoeFusedGate::execute(Tensor topk_weights, - Tensor topk_indices, - const Tensor &input, - const Tensor &bias, - const size_t num_expert_group, - const size_t topk_group, - const size_t num_fused_shared_experts, - const float routed_scaling_factor, - const bool apply_routed_scaling_factor_on_output) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - MoeFusedGate, - topk_weights, - topk_indices, - input, - bias, - num_expert_group, - topk_group, - num_fused_shared_experts, - routed_scaling_factor, - apply_routed_scaling_factor_on_output); -} - -std::tuple moe_fused_gate( - const Tensor &input, - const Tensor &bias, - size_t topk, - size_t num_expert_group, - size_t topk_group, - size_t num_fused_shared_experts, - float routed_scaling_factor, - bool apply_routed_scaling_factor_on_output) { - auto shape = input->shape(); - INFINICORE_ASSERT(shape.size() == 2); - auto topk_weights = Tensor::empty({shape[0], topk}, DataType::F32, input->device()); - auto topk_indices = Tensor::empty({shape[0], topk}, DataType::I32, input->device()); - moe_fused_gate_( - topk_weights, - topk_indices, - input, - bias, - num_expert_group, - topk_group, - num_fused_shared_experts, - routed_scaling_factor, - apply_routed_scaling_factor_on_output); - return {topk_weights, topk_indices}; -} - -void moe_fused_gate_(Tensor topk_weights, - Tensor topk_indices, - const Tensor &input, - const Tensor &bias, - size_t num_expert_group, - size_t topk_group, - size_t num_fused_shared_experts, - float routed_scaling_factor, - bool apply_routed_scaling_factor_on_output) { - MoeFusedGate::execute( - topk_weights, - topk_indices, - input, - bias, - num_expert_group, - topk_group, - num_fused_shared_experts, - routed_scaling_factor, - apply_routed_scaling_factor_on_output); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/moe_fused_gate/moe_fused_gate_infiniop.cc b/src/infinicore/ops/moe_fused_gate/moe_fused_gate_infiniop.cc deleted file mode 100644 index 99be1aa6d..000000000 --- a/src/infinicore/ops/moe_fused_gate/moe_fused_gate_infiniop.cc +++ /dev/null @@ -1,82 +0,0 @@ -#include "infinicore/ops/moe_fused_gate.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::moe_fused_gate_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, MoeFusedGate, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace; - graph::GraphTensor topk_weights; - graph::GraphTensor topk_indices; - graph::GraphTensor input; - graph::GraphTensor bias; -}; - -void *plan(Tensor topk_weights, - Tensor topk_indices, - const Tensor &input, - const Tensor &bias, - const size_t num_expert_group, - const size_t topk_group, - const size_t num_fused_shared_experts, - const float routed_scaling_factor, - const bool apply_routed_scaling_factor_on_output) { - size_t seed = hash_combine( - topk_weights, - topk_indices, - input, - bias, - num_expert_group, - topk_group, - num_fused_shared_experts, - routed_scaling_factor, - apply_routed_scaling_factor_on_output); - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, - descriptor, - MoeFusedGate, - seed, - topk_weights->desc(), - topk_indices->desc(), - input->desc(), - bias->desc(), - num_expert_group, - topk_group, - num_fused_shared_experts, - routed_scaling_factor, - apply_routed_scaling_factor_on_output); - INFINIOP_WORKSPACE_TENSOR(workspace, MoeFusedGate, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(topk_weights), - graph::GraphTensor(topk_indices), - graph::GraphTensor(input), - graph::GraphTensor(bias)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - INFINICORE_CHECK_ERROR(infiniopMoeFusedGate( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->topk_weights->data(), - planned->topk_indices->data(), - planned->input->data(), - planned->bias->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MoeFusedGate, &plan, &run, cleanup); - -} // namespace infinicore::op::moe_fused_gate_impl::infiniop diff --git a/src/infinicore/ops/moe_sum/moe_sum.cc b/src/infinicore/ops/moe_sum/moe_sum.cc deleted file mode 100644 index 3a5dbc852..000000000 --- a/src/infinicore/ops/moe_sum/moe_sum.cc +++ /dev/null @@ -1,30 +0,0 @@ -#include "infinicore/ops/moe_sum.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeSum); - -MoeSum::MoeSum(Tensor output, const Tensor &input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - INFINICORE_GRAPH_OP_DISPATCH(output->device().getType(), output, input); -} - -void MoeSum::execute(Tensor output, const Tensor &input) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(MoeSum, output, input); -} - -Tensor moe_sum(const Tensor &input) { - auto shape = input->shape(); - INFINICORE_ASSERT(shape.size() == 3); - auto output = Tensor::empty({shape[0], shape[2]}, input->dtype(), input->device()); - moe_sum_(output, input); - return output; -} - -void moe_sum_(Tensor output, const Tensor &input) { - MoeSum::execute(output, input); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/moe_sum/moe_sum_infiniop.cc b/src/infinicore/ops/moe_sum/moe_sum_infiniop.cc deleted file mode 100644 index 71a887e5f..000000000 --- a/src/infinicore/ops/moe_sum/moe_sum_infiniop.cc +++ /dev/null @@ -1,51 +0,0 @@ -#include "infinicore/ops/moe_sum.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::moe_sum_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, MoeSum, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace; - graph::GraphTensor output; - graph::GraphTensor input; -}; - -void *plan(Tensor output, const Tensor &input) { - size_t seed = hash_combine(output, input); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, MoeSum, - seed, output->desc(), input->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, MoeSum, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(output), - graph::GraphTensor(input)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopMoeSum( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->output->data(), - planned->input->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MoeSum, &plan, &run, cleanup); - -} // namespace infinicore::op::moe_sum_impl::infiniop diff --git a/src/infinicore/ops/moe_topk_sigmoid/moe_topk_sigmoid.cc b/src/infinicore/ops/moe_topk_sigmoid/moe_topk_sigmoid.cc deleted file mode 100644 index a612a0565..000000000 --- a/src/infinicore/ops/moe_topk_sigmoid/moe_topk_sigmoid.cc +++ /dev/null @@ -1,63 +0,0 @@ -#include "infinicore/ops/moe_topk_sigmoid.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeTopkSigmoid); - -MoeTopkSigmoid::MoeTopkSigmoid(Tensor topk_weights, - Tensor topk_indices, - const Tensor &gating_output, - const Tensor &correction_bias, - const bool renormalize) { - if (correction_bias) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_indices, gating_output, correction_bias); - } else { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_indices, gating_output); - } - INFINICORE_GRAPH_OP_DISPATCH( - topk_weights->device().getType(), - topk_weights, - topk_indices, - gating_output, - correction_bias, - renormalize); -} - -void MoeTopkSigmoid::execute(Tensor topk_weights, - Tensor topk_indices, - const Tensor &gating_output, - const Tensor &correction_bias, - const bool renormalize) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - MoeTopkSigmoid, - topk_weights, - topk_indices, - gating_output, - correction_bias, - renormalize); -} - -std::tuple moe_topk_sigmoid( - const Tensor &gating_output, - size_t topk, - bool renormalize, - const Tensor &correction_bias) { - auto shape = gating_output->shape(); - INFINICORE_ASSERT(shape.size() == 2); - auto topk_weights = Tensor::empty({shape[0], topk}, DataType::F32, gating_output->device()); - auto topk_indices = Tensor::empty({shape[0], topk}, DataType::I32, gating_output->device()); - moe_topk_sigmoid_(topk_weights, topk_indices, gating_output, correction_bias, renormalize); - return {topk_weights, topk_indices}; -} - -void moe_topk_sigmoid_(Tensor topk_weights, - Tensor topk_indices, - const Tensor &gating_output, - const Tensor &correction_bias, - bool renormalize) { - MoeTopkSigmoid::execute(topk_weights, topk_indices, gating_output, correction_bias, renormalize); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/moe_topk_sigmoid/moe_topk_sigmoid_infiniop.cc b/src/infinicore/ops/moe_topk_sigmoid/moe_topk_sigmoid_infiniop.cc deleted file mode 100644 index 192134a79..000000000 --- a/src/infinicore/ops/moe_topk_sigmoid/moe_topk_sigmoid_infiniop.cc +++ /dev/null @@ -1,72 +0,0 @@ -#include "infinicore/ops/moe_topk_sigmoid.hpp" - -#include "../infiniop_impl.hpp" - -#include - -namespace infinicore::op::moe_topk_sigmoid_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, MoeTopkSigmoid, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace; - graph::GraphTensor topk_weights; - graph::GraphTensor topk_indices; - graph::GraphTensor gating_output; - std::optional correction_bias; -}; - -void *plan(Tensor topk_weights, - Tensor topk_indices, - const Tensor &gating_output, - const Tensor &correction_bias, - const bool renormalize) { - size_t seed = hash_combine(topk_weights, topk_indices, gating_output, correction_bias, renormalize); - infiniopTensorDescriptor_t correction_bias_desc = correction_bias ? correction_bias->desc() : nullptr; - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, - descriptor, - MoeTopkSigmoid, - seed, - topk_weights->desc(), - topk_indices->desc(), - gating_output->desc(), - correction_bias_desc, - renormalize); - INFINIOP_WORKSPACE_TENSOR(workspace, MoeTopkSigmoid, descriptor); - - std::optional correction_bias_graph; - if (correction_bias) { - correction_bias_graph.emplace(correction_bias); - } - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(topk_weights), - graph::GraphTensor(topk_indices), - graph::GraphTensor(gating_output), - correction_bias_graph}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - INFINICORE_CHECK_ERROR(infiniopMoeTopkSigmoid( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->topk_weights->data(), - planned->topk_indices->data(), - planned->gating_output->data(), - planned->correction_bias ? (*planned->correction_bias)->data() : nullptr, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MoeTopkSigmoid, &plan, &run, cleanup); - -} // namespace infinicore::op::moe_topk_sigmoid_impl::infiniop diff --git a/src/infinicore/ops/moe_topk_softmax/moe_topk_softmax.cc b/src/infinicore/ops/moe_topk_softmax/moe_topk_softmax.cc deleted file mode 100644 index 84dee3b49..000000000 --- a/src/infinicore/ops/moe_topk_softmax/moe_topk_softmax.cc +++ /dev/null @@ -1,69 +0,0 @@ -#include "infinicore/ops/moe_topk_softmax.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MoeTopkSoftmax); - -MoeTopkSoftmax::MoeTopkSoftmax(Tensor topk_weights, - Tensor topk_indices, - const Tensor &gating_output, - const Tensor &correction_bias, - const bool renormalize, - const float moe_softcapping) { - if (correction_bias) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_indices, gating_output, correction_bias); - } else { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(topk_weights, topk_indices, gating_output); - } - INFINICORE_GRAPH_OP_DISPATCH( - topk_weights->device().getType(), - topk_weights, - topk_indices, - gating_output, - correction_bias, - renormalize, - moe_softcapping); -} - -void MoeTopkSoftmax::execute(Tensor topk_weights, - Tensor topk_indices, - const Tensor &gating_output, - const Tensor &correction_bias, - const bool renormalize, - const float moe_softcapping) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - MoeTopkSoftmax, - topk_weights, - topk_indices, - gating_output, - correction_bias, - renormalize, - moe_softcapping); -} - -std::tuple moe_topk_softmax( - const Tensor &gating_output, - size_t topk, - bool renormalize, - float moe_softcapping, - const Tensor &correction_bias) { - auto shape = gating_output->shape(); - INFINICORE_ASSERT(shape.size() == 2); - auto topk_weights = Tensor::empty({shape[0], topk}, DataType::F32, gating_output->device()); - auto topk_indices = Tensor::empty({shape[0], topk}, DataType::I32, gating_output->device()); - moe_topk_softmax_(topk_weights, topk_indices, gating_output, correction_bias, renormalize, moe_softcapping); - return {topk_weights, topk_indices}; -} - -void moe_topk_softmax_(Tensor topk_weights, - Tensor topk_indices, - const Tensor &gating_output, - const Tensor &correction_bias, - bool renormalize, - float moe_softcapping) { - MoeTopkSoftmax::execute(topk_weights, topk_indices, gating_output, correction_bias, renormalize, moe_softcapping); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/moe_topk_softmax/moe_topk_softmax_infiniop.cc b/src/infinicore/ops/moe_topk_softmax/moe_topk_softmax_infiniop.cc deleted file mode 100644 index cf6bee84b..000000000 --- a/src/infinicore/ops/moe_topk_softmax/moe_topk_softmax_infiniop.cc +++ /dev/null @@ -1,74 +0,0 @@ -#include "infinicore/ops/moe_topk_softmax.hpp" - -#include "../infiniop_impl.hpp" - -#include - -namespace infinicore::op::moe_topk_softmax_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, MoeTopkSoftmax, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace; - graph::GraphTensor topk_weights; - graph::GraphTensor topk_indices; - graph::GraphTensor gating_output; - std::optional correction_bias; -}; - -void *plan(Tensor topk_weights, - Tensor topk_indices, - const Tensor &gating_output, - const Tensor &correction_bias, - const bool renormalize, - const float moe_softcapping) { - size_t seed = hash_combine(topk_weights, topk_indices, gating_output, correction_bias, renormalize, moe_softcapping); - infiniopTensorDescriptor_t correction_bias_desc = correction_bias ? correction_bias->desc() : nullptr; - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, - descriptor, - MoeTopkSoftmax, - seed, - topk_weights->desc(), - topk_indices->desc(), - gating_output->desc(), - correction_bias_desc, - renormalize, - moe_softcapping); - INFINIOP_WORKSPACE_TENSOR(workspace, MoeTopkSoftmax, descriptor); - - std::optional correction_bias_graph; - if (correction_bias) { - correction_bias_graph.emplace(correction_bias); - } - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(topk_weights), - graph::GraphTensor(topk_indices), - graph::GraphTensor(gating_output), - correction_bias_graph}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - INFINICORE_CHECK_ERROR(infiniopMoeTopkSoftmax( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->topk_weights->data(), - planned->topk_indices->data(), - planned->gating_output->data(), - planned->correction_bias ? (*planned->correction_bias)->data() : nullptr, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MoeTopkSoftmax, &plan, &run, cleanup); - -} // namespace infinicore::op::moe_topk_softmax_impl::infiniop diff --git a/src/infinicore/ops/mrope/mrope.cc b/src/infinicore/ops/mrope/mrope.cc deleted file mode 100644 index 6dcd8b5fe..000000000 --- a/src/infinicore/ops/mrope/mrope.cc +++ /dev/null @@ -1,100 +0,0 @@ -#include "infinicore/ops/mrope.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MRoPE); - -MRoPE::MRoPE(Tensor q_out, - Tensor k_out, - const Tensor &q, - const Tensor &k, - const Tensor &cos, - const Tensor &sin, - const Tensor &positions, - int head_size, - int rotary_dim, - int section_t, - int section_h, - int section_w, - bool interleaved) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(q_out, k_out, q, k, cos, sin, positions); - INFINICORE_GRAPH_OP_DISPATCH(q->device().getType(), - q_out, - k_out, - q, - k, - cos, - sin, - positions, - head_size, - rotary_dim, - section_t, - section_h, - section_w, - interleaved); -} - -void MRoPE::execute(Tensor q_out, - Tensor k_out, - const Tensor &q, - const Tensor &k, - const Tensor &cos, - const Tensor &sin, - const Tensor &positions, - int head_size, - int rotary_dim, - int section_t, - int section_h, - int section_w, - bool interleaved) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(MRoPE, - q_out, - k_out, - q, - k, - cos, - sin, - positions, - head_size, - rotary_dim, - section_t, - section_h, - section_w, - interleaved); -} - -void mrope_(Tensor q_out, - Tensor k_out, - const Tensor &q, - const Tensor &k, - const Tensor &cos, - const Tensor &sin, - const Tensor &positions, - int head_size, - int rotary_dim, - int section_t, - int section_h, - int section_w, - bool interleaved) { - MRoPE::execute(q_out, k_out, q, k, cos, sin, positions, head_size, rotary_dim, section_t, section_h, section_w, interleaved); -} - -std::pair mrope(const Tensor &q, - const Tensor &k, - const Tensor &cos, - const Tensor &sin, - const Tensor &positions, - int head_size, - int rotary_dim, - int section_t, - int section_h, - int section_w, - bool interleaved) { - auto q_out = Tensor::empty(q->shape(), q->dtype(), q->device()); - auto k_out = Tensor::empty(k->shape(), k->dtype(), k->device()); - mrope_(q_out, k_out, q, k, cos, sin, positions, head_size, rotary_dim, section_t, section_h, section_w, interleaved); - return {q_out, k_out}; -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/mrope/mrope_infiniop.cc b/src/infinicore/ops/mrope/mrope_infiniop.cc deleted file mode 100644 index 13280319f..000000000 --- a/src/infinicore/ops/mrope/mrope_infiniop.cc +++ /dev/null @@ -1,105 +0,0 @@ -#include "infinicore/ops/mrope.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::mrope_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, MRoPE, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace; - graph::GraphTensor q_out; - graph::GraphTensor k_out; - graph::GraphTensor q; - graph::GraphTensor k; - graph::GraphTensor cos; - graph::GraphTensor sin; - graph::GraphTensor positions; -}; - -void *plan(Tensor q_out, - Tensor k_out, - const Tensor &q, - const Tensor &k, - const Tensor &cos, - const Tensor &sin, - const Tensor &positions, - int head_size, - int rotary_dim, - int section_t, - int section_h, - int section_w, - bool interleaved) { - size_t key = hash_combine(q_out, - k_out, - q, - k, - cos, - sin, - positions, - head_size, - rotary_dim, - section_t, - section_h, - section_w, - interleaved); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, - descriptor, - MRoPE, - key, - q_out->desc(), - k_out->desc(), - q->desc(), - k->desc(), - cos->desc(), - sin->desc(), - positions->desc(), - head_size, - rotary_dim, - section_t, - section_h, - section_w, - interleaved); - - INFINIOP_WORKSPACE_TENSOR(workspace, MRoPE, descriptor); - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(q_out), - graph::GraphTensor(k_out), - graph::GraphTensor(q), - graph::GraphTensor(k), - graph::GraphTensor(cos), - graph::GraphTensor(sin), - graph::GraphTensor(positions)}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR( - infiniopMRoPE( - p->descriptor->desc, - p->workspace->data(), - p->workspace->numel(), - p->q_out->data(), - p->k_out->data(), - p->q->data(), - p->k->data(), - p->cos->data(), - p->sin->data(), - p->positions->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MRoPE, &plan, &run, &cleanup); - -} // namespace infinicore::op::mrope_impl::infiniop diff --git a/src/infinicore/ops/mul/mul.cc b/src/infinicore/ops/mul/mul.cc deleted file mode 100644 index 6923fed9c..000000000 --- a/src/infinicore/ops/mul/mul.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/mul.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Mul); - -Mul::Mul(Tensor c, const Tensor &a, const Tensor &b) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - INFINICORE_GRAPH_OP_DISPATCH(c->device().getType(), c, a, b); -} - -void Mul::execute(Tensor c, const Tensor &a, const Tensor &b) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Mul, c, a, b); -} - -Tensor mul(const Tensor &a, const Tensor &b) { - auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); - mul_(c, a, b); - return c; -} - -void mul_(Tensor c, const Tensor &a, const Tensor &b) { - Mul::execute(c, a, b); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/mul/mul_infiniop.cc b/src/infinicore/ops/mul/mul_infiniop.cc deleted file mode 100644 index 39a7bd87d..000000000 --- a/src/infinicore/ops/mul/mul_infiniop.cc +++ /dev/null @@ -1,51 +0,0 @@ -#include "infinicore/ops/mul.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::mul_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Mul, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, c, a, b; -}; - -void *plan(Tensor c, const Tensor &a, const Tensor &b) { - size_t seed = hash_combine(c, b, a); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Mul, - seed, c->desc(), a->desc(), b->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Mul, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(c), - graph::GraphTensor(a), - graph::GraphTensor(b)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopMul( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->c->data(), - planned->a->data(), - planned->b->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Mul, &plan, &run, &cleanup); - -} // namespace infinicore::op::mul_impl::infiniop diff --git a/src/infinicore/ops/mul_scalar/mul_scalar.cc b/src/infinicore/ops/mul_scalar/mul_scalar.cc deleted file mode 100644 index b5be25e7e..000000000 --- a/src/infinicore/ops/mul_scalar/mul_scalar.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/mul_scalar.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MulScalar); - -MulScalar::MulScalar(Tensor c, const Tensor &a, double alpha) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a); - INFINICORE_GRAPH_OP_DISPATCH(c->device().getType(), c, a, alpha); -} - -void MulScalar::execute(Tensor c, const Tensor &a, double alpha) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(MulScalar, c, a, alpha); -} - -Tensor mul_scalar(const Tensor &a, double alpha) { - auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); - mul_scalar_(c, a, alpha); - return c; -} - -void mul_scalar_(Tensor c, const Tensor &a, double alpha) { - MulScalar::execute(c, a, alpha); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/mul_scalar/mul_scalar_infiniop.cc b/src/infinicore/ops/mul_scalar/mul_scalar_infiniop.cc deleted file mode 100644 index 13f3c615c..000000000 --- a/src/infinicore/ops/mul_scalar/mul_scalar_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "infinicore/ops/mul_scalar.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::mul_scalar_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, MulScalar, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, c, a; - double alpha; -}; - -void *plan(Tensor c, const Tensor &a, double alpha) { - size_t seed = hash_combine(c, a, alpha); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, MulScalar, - seed, c->desc(), a->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, MulScalar, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(c), - graph::GraphTensor(a), - alpha}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopMulScalar( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->c->data(), - planned->a->data(), - planned->alpha, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MulScalar, &plan, &run, &cleanup); - -} // namespace infinicore::op::mul_scalar_impl::infiniop diff --git a/src/infinicore/ops/multi_head_attention/hygon/mha_flashattn_hygon.cc b/src/infinicore/ops/multi_head_attention/hygon/mha_flashattn_hygon.cc deleted file mode 100644 index 247507c0c..000000000 --- a/src/infinicore/ops/multi_head_attention/hygon/mha_flashattn_hygon.cc +++ /dev/null @@ -1,102 +0,0 @@ -#if defined(ENABLE_HYGON_API) && defined(ENABLE_FLASH_ATTN) -#include "infinicore/ops/mha.hpp" - -#include "../../../adaptor/flash_attn/hygon/flash_attn_hygon.hpp" -#include "infinicore/adaptor/aten_adaptor.hpp" - -#include -#include - -#include - -namespace infinicore::op::mha_impl::flashattn { -struct PlannedMeta { - graph::GraphTensor out, q, k, v; - std::optional alibi_slopes; - float scale; - bool is_causal; -}; - -void *plan(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - std::optional alibi_slopes, - float scale, - bool is_causal) { - - return new PlannedMeta{ - graph::GraphTensor(out), - graph::GraphTensor(q), - graph::GraphTensor(k), - graph::GraphTensor(v), - alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, - scale, - is_causal}; -} - -void run(void *planned_meta) { - c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); - auto *p = reinterpret_cast(planned_meta); - - auto q = infinicore::adaptor::to_aten_tensor(p->q); - auto k = infinicore::adaptor::to_aten_tensor(p->k); - auto v = infinicore::adaptor::to_aten_tensor(p->v); - - const bool out_need_copy_back = !p->out->is_contiguous(); - Tensor out_work_ic = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); - auto out_work = infinicore::adaptor::to_aten_tensor(out_work_ic); - auto scale = p->scale; - auto is_causal = p->is_causal; - - const int64_t num_q_heads = q.size(2); - const int64_t num_kv_heads = k.size(2); - if (q.dim() == 4 && k.dim() == 4 && v.dim() == 4 && q.size(1) == 1 && k.size(1) == 1 && v.size(1) == 1) { - at::Tensor value = v; - if (num_kv_heads != num_q_heads) { - const int64_t repeat = num_q_heads / num_kv_heads; - value = v.unsqueeze(3).repeat({1, 1, 1, repeat, 1}).reshape(q.sizes()); - } - out_work.copy_(value); - if (out_need_copy_back) { - p->out->copy_from(out_work_ic); - } - return; - } - - at::Tensor k_attn = k; - at::Tensor v_attn = v; - if (num_kv_heads != num_q_heads) { - const int64_t repeat = num_q_heads / num_kv_heads; - k_attn = k.unsqueeze(3).repeat({1, 1, 1, repeat, 1}).reshape(q.sizes()); - v_attn = v.unsqueeze(3).repeat({1, 1, 1, repeat, 1}).reshape(q.sizes()); - } - - auto attn_weight = at::matmul(q.permute({0, 2, 1, 3}), k_attn.permute({0, 2, 3, 1})) * scale; - if (is_causal) { - auto mask = at::tril(at::ones_like(attn_weight), -1).flip({-2, -1}); - auto neg_inf = at::full_like(attn_weight, -std::numeric_limits::infinity()); - attn_weight = at::where(mask == 1, neg_inf, attn_weight); - } - auto attn = at::softmax(attn_weight, -1); - auto result = at::matmul(attn, v_attn.permute({0, 2, 1, 3})).permute({0, 2, 1, 3}); - out_work.copy_(result); - if (out_need_copy_back) { - p->out->copy_from(out_work_ic); - } -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - MultiheadAttention::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); - MultiheadAttention::run_dispatcher().registerDevice(Device::Type::HYGON, &run); - MultiheadAttention::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); - return true; -}(); - -} // namespace infinicore::op::mha_impl::flashattn -#endif // ENABLE_HYGON_API && ENABLE_FLASH_ATTN diff --git a/src/infinicore/ops/multi_head_attention/mha.cc b/src/infinicore/ops/multi_head_attention/mha.cc deleted file mode 100644 index 46d6ffffb..000000000 --- a/src/infinicore/ops/multi_head_attention/mha.cc +++ /dev/null @@ -1,54 +0,0 @@ -#include "infinicore/ops/mha.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MultiheadAttention); - -MultiheadAttention::MultiheadAttention(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - std::optional alibi_slopes, - float scale, - bool is_causal) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k, v); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), - out, q, k, v, alibi_slopes, scale, is_causal); -} - -void MultiheadAttention::execute(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - std::optional alibi_slopes, - float scale, - bool is_causal) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - MultiheadAttention, - out, q, k, v, alibi_slopes, scale, is_causal); -} - -Tensor mha( - const Tensor &q, - const Tensor &k, - const Tensor &v, - std::optional alibi_slopes, - float scale, - bool is_causal) { - auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); - mha_(out, q, k, v, alibi_slopes, scale, is_causal); - return out; -} - -void mha_(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - std::optional alibi_slopes, - float scale, - bool is_causal) { - MultiheadAttention::execute(out, q, k, v, alibi_slopes, scale, is_causal); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/multi_head_attention/mha_flashattn.cc b/src/infinicore/ops/multi_head_attention/mha_flashattn.cc deleted file mode 100644 index 13c96b94d..000000000 --- a/src/infinicore/ops/multi_head_attention/mha_flashattn.cc +++ /dev/null @@ -1,115 +0,0 @@ -#include "infinicore/ops/mha.hpp" - -#include "infinicore/adaptor/flash_attention_adaptor.hpp" - -#include - -#ifdef ENABLE_FLASH_ATTN -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) -#include -#endif -#endif - -namespace infinicore::op::mha_impl::flashattn { - -struct PlannedMeta { - graph::GraphTensor out, q, k, v; - std::optional alibi_slopes; - float scale; - bool is_causal; -}; - -void *plan(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - std::optional alibi_slopes, - float scale, - bool is_causal) { - - return new PlannedMeta{ - graph::GraphTensor(out), - graph::GraphTensor(q), - graph::GraphTensor(k), - graph::GraphTensor(v), - alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, - scale, - is_causal}; -} - -namespace { - -// Only support nv for now -#if defined(ENABLE_FLASH_ATTN) && defined(ENABLE_NVIDIA_API) -// MetaX/hpcc pip `flash_attn_2_cuda` exports `mha_fwd` at global scope (no namespace), -// while NVIDIA `flash-attn-nvidia.so` uses `flash::mha_fwd`. -#if defined(ENABLE_METAX_API) -#define INFINICORE_FLASH_OP(name) ::name -#else -#define INFINICORE_FLASH_OP(name) flash::name -#endif - -#endif // ENABLE_FLASH_ATTN -} // namespace - -void run(void *planned_meta) { -// Only support nv for now -#if defined(ENABLE_FLASH_ATTN) && defined(ENABLE_NVIDIA_API) - c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream()); - auto *p = reinterpret_cast(planned_meta); - - auto q = infinicore::adaptor::to_aten_tensor(p->q); - auto k = infinicore::adaptor::to_aten_tensor(p->k); - auto v = infinicore::adaptor::to_aten_tensor(p->v); - - const bool out_need_copy_back = !p->out->is_contiguous(); - Tensor out_work_ic = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); - auto out_work = infinicore::adaptor::to_aten_tensor(out_work_ic); - auto out = std::optional(out_work); - - auto alibi_slopes = p->alibi_slopes ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; - auto scale = p->scale; - auto is_causal = p->is_causal; - -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - std::optional flash_attn_mars_ext = std::nullopt; -#endif - - INFINICORE_FLASH_OP(mha_fwd) - ( - q, - k, - v, - out, - alibi_slopes, - 0.0, - scale, - is_causal, - -1, - -1, - 0.0, - false, - std::nullopt -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - , - flash_attn_mars_ext -#endif - ); - - if (out_need_copy_back) { - p->out->copy_from(out_work_ic); - } - -#else - throw std::runtime_error("FlashAttention is not enabled in this build"); -#endif -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MultiheadAttention, &plan, &run, &cleanup); - -} // namespace infinicore::op::mha_impl::flashattn diff --git a/src/infinicore/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc b/src/infinicore/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc deleted file mode 100644 index b9ffe1570..000000000 --- a/src/infinicore/ops/multi_head_attention_varlen/hygon/mha_varlen_flashattn_hygon.cc +++ /dev/null @@ -1,266 +0,0 @@ -#if defined(ENABLE_HYGON_API) && defined(ENABLE_FLASH_ATTN) -#include "infinicore/ops/mha_varlen.hpp" - -#ifdef ENABLE_ATEN -#include "infinicore/adaptor/aten_adaptor.hpp" -#include -#include -#include -#endif - -#include "../../../adaptor/flash_attn/hygon/flash_attn_hygon.hpp" - -#include -#include -#include -#include - -namespace infinicore::op::mha_varlen_impl::flashattn { - -struct PlannedMeta { - graph::GraphTensor out, q, k, v, cum_seqlens_q, cum_seqlens_k; - std::optional block_table; - int max_seqlen_q, max_seqlen_k; - std::optional alibi_slopes; - float scale; -}; - -void *plan(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &cum_seqlens_q, - const Tensor &cum_seqlens_k, - std::optional block_table, - int max_seqlen_q, - int max_seqlen_k, - std::optional alibi_slopes, - float scale) { - - return new PlannedMeta{ - graph::GraphTensor(out), - graph::GraphTensor(q), - graph::GraphTensor(k), - graph::GraphTensor(v), - graph::GraphTensor(cum_seqlens_q), - graph::GraphTensor(cum_seqlens_k), - block_table ? std::optional(graph::GraphTensor(*block_table)) : std::nullopt, - max_seqlen_q, - max_seqlen_k, - alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, - scale}; -} - -void run(void *planned_meta) { -#ifndef ENABLE_ATEN - (void)planned_meta; - throw std::runtime_error("ATen is not enabled in this build"); -#else - c10::hip::HIPStreamGuard guard(infinicore::adaptor::get_hip_stream()); - auto *p = reinterpret_cast(planned_meta); - - auto q = infinicore::adaptor::to_aten_tensor(p->q); - auto k = infinicore::adaptor::to_aten_tensor(p->k); - auto v = infinicore::adaptor::to_aten_tensor(p->v); - - const bool out_need_copy_back = !p->out->is_contiguous(); - Tensor out_work_ic = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); - auto out_work = infinicore::adaptor::to_aten_tensor(out_work_ic); - - auto cu_seqlens_q = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_q); - auto cu_seqlens_kv = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_k); - - const bool dense_sdpa = !p->block_table.has_value() - && !p->alibi_slopes.has_value() - && q.dim() == 3 && k.dim() == 3 && v.dim() == 3 - && p->max_seqlen_q > 0 && p->max_seqlen_k > 0 - && p->max_seqlen_q == p->max_seqlen_k - && cu_seqlens_q.dim() == 1 - && cu_seqlens_q.size(0) == cu_seqlens_kv.size(0) - && q.size(0) == (cu_seqlens_q.size(0) - 1) * p->max_seqlen_q - && k.size(0) == (cu_seqlens_kv.size(0) - 1) * p->max_seqlen_k - && ((q.size(2) > 256) || (v.size(2) != q.size(2))); - if (dense_sdpa) { - const int64_t batch_size = cu_seqlens_q.size(0) - 1; - const int64_t seqlen = p->max_seqlen_q; - const int64_t num_heads = q.size(1); - const int64_t num_kv_heads = k.size(1); - const int64_t head_dim = q.size(2); - const int64_t value_dim = v.size(2); - auto q_4d = q.reshape({batch_size, seqlen, num_heads, head_dim}).permute({0, 2, 1, 3}); - auto k_4d = k.reshape({batch_size, seqlen, num_kv_heads, head_dim}).permute({0, 2, 1, 3}); - auto v_4d = v.reshape({batch_size, seqlen, num_kv_heads, value_dim}).permute({0, 2, 1, 3}); - if (num_heads != num_kv_heads) { - if (num_heads % num_kv_heads != 0) { - throw std::runtime_error("mha_varlen dense SDPA fallback requires num_heads to be divisible by num_kv_heads"); - } - const int64_t groups = num_heads / num_kv_heads; - k_4d = k_4d.unsqueeze(2).expand({batch_size, num_kv_heads, groups, seqlen, head_dim}).reshape({batch_size, num_heads, seqlen, head_dim}); - v_4d = v_4d.unsqueeze(2).expand({batch_size, num_kv_heads, groups, seqlen, value_dim}).reshape({batch_size, num_heads, seqlen, value_dim}); - } - auto result = at::scaled_dot_product_attention( - q_4d, - k_4d, - v_4d, - std::nullopt, - 0.0, - true, - std::optional(static_cast(p->scale))); - out_work.copy_(result.permute({0, 2, 1, 3}).reshape({q.size(0), num_heads, value_dim})); - if (out_need_copy_back) { - p->out->copy_from(out_work_ic); - } - return; - } - - if (std::getenv("INFINICORE_HYGON_ATEN_FALLBACK")) { - namespace idx = at::indexing; - auto cu_q_cpu = cu_seqlens_q.to(at::kCPU); - auto cu_k_cpu = cu_seqlens_kv.to(at::kCPU); - const int64_t num_seqs = cu_q_cpu.size(0) - 1; - auto result = at::zeros_like(out_work); - - if (!p->block_table.has_value()) { - for (int64_t i = 0; i < num_seqs; ++i) { - const int64_t q_start = cu_q_cpu.index({i}).item(); - const int64_t q_end = cu_q_cpu.index({i + 1}).item(); - const int64_t k_start = cu_k_cpu.index({i}).item(); - const int64_t k_end = cu_k_cpu.index({i + 1}).item(); - auto cur_q = q.index({idx::Slice(q_start, q_end)}).unsqueeze(0).transpose(1, 2); - auto cur_k = k.index({idx::Slice(k_start, k_end)}).unsqueeze(0).transpose(1, 2); - auto cur_v = v.index({idx::Slice(k_start, k_end)}).unsqueeze(0).transpose(1, 2); - auto cur_out = at::scaled_dot_product_attention( - cur_q, cur_k, cur_v, std::nullopt, 0.0, true, std::optional(static_cast(p->scale))); - result.index_put_({idx::Slice(q_start, q_end)}, cur_out.transpose(1, 2).squeeze(0)); - } - } else { - auto block_table_t = infinicore::adaptor::to_aten_tensor(*p->block_table); - auto block_table_cpu = block_table_t.to(at::kCPU); - const int64_t block_size = k.size(1); - for (int64_t i = 0; i < num_seqs; ++i) { - const int64_t q_start = cu_q_cpu.index({i}).item(); - const int64_t q_end = cu_q_cpu.index({i + 1}).item(); - const int64_t q_len = q_end - q_start; - const int64_t h_len = (cu_k_cpu.index({i + 1}).item() - cu_k_cpu.index({i}).item()) - q_len; - const int64_t total_len = h_len + q_len; - auto cur_q = q.index({idx::Slice(q_start, q_end)}); - std::vector keys; - std::vector values; - keys.reserve(total_len); - values.reserve(total_len); - for (int64_t j = 0; j < total_len; ++j) { - const int64_t b_id = block_table_cpu.index({i, j / block_size}).item(); - const int64_t off = j % block_size; - keys.push_back(k.index({b_id, off, idx::Slice(), idx::Slice()})); - values.push_back(v.index({b_id, off, idx::Slice(), idx::Slice()})); - } - auto K = at::stack(keys, 0); - auto V = at::stack(values, 0); - const int64_t q_heads = cur_q.size(1); - const int64_t kv_heads = K.size(1); - if (q_heads != kv_heads) { - const int64_t repeat = q_heads / kv_heads; - K = K.repeat_interleave(repeat, 1); - V = V.repeat_interleave(repeat, 1); - } - auto scores = at::matmul(cur_q.permute({1, 0, 2}).to(at::kFloat), K.permute({1, 2, 0}).to(at::kFloat)) * p->scale; - auto mask = at::full({q_len, total_len}, -std::numeric_limits::infinity(), q.options().dtype(at::kFloat)); - for (int64_t t = 0; t < q_len; ++t) { - mask.index_put_({t, idx::Slice(0, h_len + t + 1)}, 0.0); - } - auto attn = at::softmax(scores + mask.unsqueeze(0), -1).to(q.dtype()); - auto cur_out = at::matmul(attn, V.permute({1, 0, 2})).permute({1, 0, 2}); - result.index_put_({idx::Slice(q_start, q_end)}, cur_out); - } - } - - out_work.copy_(result); - if (out_need_copy_back) { - p->out->copy_from(out_work_ic); - } - return; - } - - auto out = std::optional(out_work); - std::optional seqused_k = std::nullopt; - std::optional leftpad_k = std::nullopt; - auto block_table = p->block_table ? std::optional(infinicore::adaptor::to_aten_tensor(*p->block_table)) : std::nullopt; - auto max_seqlen_q = p->max_seqlen_q; - auto max_seqlen_k = p->max_seqlen_k; - auto alibi_slopes = p->alibi_slopes ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; - auto scale = p->scale; - - if (alibi_slopes.has_value()) { - throw std::runtime_error("[mha_varlen/hygon] ALiBi is not supported by the direct libflash_attention varlen_fwd ABI"); - } - auto q_work = q.contiguous(); - auto k_work = k.contiguous(); - auto v_work = v.contiguous(); - if (block_table.has_value() && k.dim() == 4 && v.dim() == 4) { - const int64_t num_blocks = k.size(0); - const int64_t block_size = k.size(1); - const int64_t num_kv_heads = k.size(2); - const int64_t head_dim = k.size(3); - if (block_size % 64 != 0) { - throw std::runtime_error("[mha_varlen/hygon] flash-attn requires paged KV block size to be divisible by 64"); - } - const int64_t pages_per_block = block_size / 64; - k_work = k_work.reshape({num_blocks, pages_per_block, 64, num_kv_heads, head_dim}) - .reshape({num_blocks * pages_per_block, 64, num_kv_heads, head_dim}) - .contiguous(); - v_work = v_work.reshape({num_blocks, pages_per_block, 64, num_kv_heads, head_dim}) - .reshape({num_blocks * pages_per_block, 64, num_kv_heads, head_dim}) - .contiguous(); - if (pages_per_block != 1) { - auto offsets = at::arange(pages_per_block, block_table->options()).view({1, 1, pages_per_block}); - block_table = ((*block_table).unsqueeze(-1) * pages_per_block + offsets) - .reshape({block_table->size(0), block_table->size(1) * pages_per_block}) - .contiguous(); - } - } - auto result = flash::vllm_mha_varlen_fwd( - q_work, - k_work, - v_work, - out, - cu_seqlens_q, - cu_seqlens_kv, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - 0.0, - scale, - false, - true, - -1, - -1, - 0.0, - false, - std::nullopt); - if (!result.empty() && result[0].defined()) { - out_work.copy_(result[0]); - } - if (out_need_copy_back) { - p->out->copy_from(out_work_ic); - } - -#endif -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - MultiheadAttentionVarlen::plan_dispatcher().registerDevice(Device::Type::HYGON, &plan); - MultiheadAttentionVarlen::run_dispatcher().registerDevice(Device::Type::HYGON, &run); - MultiheadAttentionVarlen::cleanup_dispatcher().registerDevice(Device::Type::HYGON, &cleanup); - return true; -}(); - -} // namespace infinicore::op::mha_varlen_impl::flashattn -#endif // ENABLE_HYGON_API && ENABLE_FLASH_ATTN diff --git a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen.cc b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen.cc deleted file mode 100644 index 50f5407ff..000000000 --- a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen.cc +++ /dev/null @@ -1,75 +0,0 @@ -#include "infinicore/ops/mha_varlen.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(MultiheadAttentionVarlen); - -MultiheadAttentionVarlen::MultiheadAttentionVarlen(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &cum_seqlens_q, - const Tensor &cum_seqlens_kv, - std::optional block_table, - int max_seqlen_q, - int max_seqlen_k, - std::optional alibi_slopes, - float scale) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k, v, cum_seqlens_q, cum_seqlens_kv); - if (block_table.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, block_table.value()); - } - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), - out, q, k, v, cum_seqlens_q, cum_seqlens_kv, block_table, max_seqlen_q, max_seqlen_k, alibi_slopes, scale); -} - -void MultiheadAttentionVarlen::execute(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &cum_seqlens_q, - const Tensor &cum_seqlens_kv, - std::optional block_table, - int max_seqlen_q, - int max_seqlen_k, - std::optional alibi_slopes, - float scale) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - MultiheadAttentionVarlen, - out, q, k, v, cum_seqlens_q, cum_seqlens_kv, block_table, max_seqlen_q, max_seqlen_k, alibi_slopes, scale); -} - -Tensor mha_varlen( - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &cum_seqlens_q, - const Tensor &cum_seqlens_kv, - std::optional block_table, - int max_seqlen_q, - int max_seqlen_k, - std::optional alibi_slopes, - float scale) { - auto out_shape = q->shape(); - out_shape.back() = v->shape().back(); - auto out = Tensor::empty(out_shape, q->dtype(), q->device()); - mha_varlen_(out, q, k, v, cum_seqlens_q, cum_seqlens_kv, block_table, max_seqlen_q, max_seqlen_k, alibi_slopes, scale); - return out; -} - -void mha_varlen_(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &cum_seqlens_q, - const Tensor &cum_seqlens_kv, - std::optional block_table, - int max_seqlen_q, - int max_seqlen_k, - std::optional alibi_slopes, - float scale) { - MultiheadAttentionVarlen::execute(out, q, k, v, cum_seqlens_q, cum_seqlens_kv, block_table, max_seqlen_q, max_seqlen_k, alibi_slopes, scale); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc deleted file mode 100644 index 4a233bf5b..000000000 --- a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn.cc +++ /dev/null @@ -1,191 +0,0 @@ -#include "infinicore/ops/mha_varlen.hpp" - -#ifdef ENABLE_ATEN -#include "infinicore/adaptor/aten_adaptor.hpp" -#include -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) -#include -#endif -#endif - -#ifdef ENABLE_FLASH_ATTN -#include "infinicore/adaptor/flash_attention_adaptor.hpp" -#endif - -#include - -namespace infinicore::op::mha_varlen_impl::flashattn { - -struct PlannedMeta { - graph::GraphTensor out, q, k, v, cum_seqlens_q, cum_seqlens_k; - std::optional block_table; - int max_seqlen_q, max_seqlen_k; - std::optional alibi_slopes; - float scale; -}; - -void *plan(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &cum_seqlens_q, - const Tensor &cum_seqlens_k, - std::optional block_table, - int max_seqlen_q, - int max_seqlen_k, - std::optional alibi_slopes, - float scale) { - - return new PlannedMeta{ - graph::GraphTensor(out), - graph::GraphTensor(q), - graph::GraphTensor(k), - graph::GraphTensor(v), - graph::GraphTensor(cum_seqlens_q), - graph::GraphTensor(cum_seqlens_k), - block_table ? std::optional(graph::GraphTensor(*block_table)) : std::nullopt, - max_seqlen_q, - max_seqlen_k, - alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, - scale}; -} - -namespace { - -#ifdef ENABLE_FLASH_ATTN -// MetaX/hpcc pip `flash_attn_2_cuda` exports `mha_varlen_fwd` at global scope (no namespace), -// while NVIDIA `flash-attn-nvidia.so` uses `flash::mha_varlen_fwd`. -#if defined(ENABLE_METAX_API) -#define INFINICORE_FLASH_OP(name) ::name -#else -#define INFINICORE_FLASH_OP(name) flash::name -#endif - -#endif // ENABLE_FLASH_ATTN -} // namespace - -void run(void *planned_meta) { -#ifndef ENABLE_ATEN - (void)planned_meta; - throw std::runtime_error("ATen is not enabled in this build"); -#else -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_METAX_API) || defined(ENABLE_QY_API) - c10::cuda::CUDAStreamGuard guard(infinicore::adaptor::get_cuda_stream()); -#endif - auto *p = reinterpret_cast(planned_meta); - - auto q = infinicore::adaptor::to_aten_tensor(p->q); - auto k = infinicore::adaptor::to_aten_tensor(p->k); - auto v = infinicore::adaptor::to_aten_tensor(p->v); - - const bool out_need_copy_back = !p->out->is_contiguous(); - Tensor out_work_ic = out_need_copy_back ? p->out->contiguous() : Tensor(p->out); - auto out_work = infinicore::adaptor::to_aten_tensor(out_work_ic); - - auto cu_seqlens_q = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_q); - auto cu_seqlens_kv = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_k); - - const bool dense_sdpa = !p->block_table.has_value() - && !p->alibi_slopes.has_value() - && q.dim() == 3 && k.dim() == 3 && v.dim() == 3 - && p->max_seqlen_q > 0 && p->max_seqlen_k > 0 - && p->max_seqlen_q == p->max_seqlen_k - && cu_seqlens_q.dim() == 1 - && cu_seqlens_q.size(0) == cu_seqlens_kv.size(0) - && q.size(0) == (cu_seqlens_q.size(0) - 1) * p->max_seqlen_q - && k.size(0) == (cu_seqlens_kv.size(0) - 1) * p->max_seqlen_k - && ((q.size(2) > 256) || (v.size(2) != q.size(2))); - if (dense_sdpa) { - const int64_t batch_size = cu_seqlens_q.size(0) - 1; - const int64_t seqlen = p->max_seqlen_q; - const int64_t num_heads = q.size(1); - const int64_t num_kv_heads = k.size(1); - const int64_t head_dim = q.size(2); - const int64_t value_dim = v.size(2); - auto q_4d = q.reshape({batch_size, seqlen, num_heads, head_dim}).permute({0, 2, 1, 3}); - auto k_4d = k.reshape({batch_size, seqlen, num_kv_heads, head_dim}).permute({0, 2, 1, 3}); - auto v_4d = v.reshape({batch_size, seqlen, num_kv_heads, value_dim}).permute({0, 2, 1, 3}); - if (num_heads != num_kv_heads) { - if (num_heads % num_kv_heads != 0) { - throw std::runtime_error("mha_varlen dense SDPA fallback requires num_heads to be divisible by num_kv_heads"); - } - const int64_t groups = num_heads / num_kv_heads; - k_4d = k_4d.unsqueeze(2).expand({batch_size, num_kv_heads, groups, seqlen, head_dim}).reshape({batch_size, num_heads, seqlen, head_dim}); - v_4d = v_4d.unsqueeze(2).expand({batch_size, num_kv_heads, groups, seqlen, value_dim}).reshape({batch_size, num_heads, seqlen, value_dim}); - } - auto result = at::scaled_dot_product_attention( - q_4d, - k_4d, - v_4d, - std::nullopt, - 0.0, - true, - std::optional(static_cast(p->scale))); - out_work.copy_(result.permute({0, 2, 1, 3}).reshape({q.size(0), num_heads, value_dim})); - if (out_need_copy_back) { - p->out->copy_from(out_work_ic); - } - return; - } - -#ifdef ENABLE_FLASH_ATTN - auto out = std::optional(out_work); - std::optional seqused_k = std::nullopt; - std::optional leftpad_k = std::nullopt; - auto block_table = p->block_table ? std::optional(infinicore::adaptor::to_aten_tensor(*p->block_table)) : std::nullopt; - auto max_seqlen_q = p->max_seqlen_q; - auto max_seqlen_k = p->max_seqlen_k; - auto alibi_slopes = p->alibi_slopes ? std::optional(infinicore::adaptor::to_aten_tensor(*p->alibi_slopes)) : std::nullopt; - auto scale = p->scale; - -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - std::optional flash_attn_mars_ext = std::nullopt; -#endif - - INFINICORE_FLASH_OP(mha_varlen_fwd) - ( - q, - k, - v, - out, - cu_seqlens_q, - cu_seqlens_kv, - seqused_k, - leftpad_k, - block_table, - alibi_slopes, - max_seqlen_q, - max_seqlen_k, - 0.0, - scale, - false, - true, - -1, - -1, - 0.0, - false, - std::nullopt -#if defined(ENABLE_METAX_API) && defined(INFINICORE_HPCC_VERSION_MAJOR) && (INFINICORE_HPCC_VERSION_MAJOR >= 3) - , - flash_attn_mars_ext -#endif - ); - - if (out_need_copy_back) { - p->out->copy_from(out_work_ic); - } - -#else - throw std::runtime_error("FlashAttention is not enabled in this build and dense SDPA fallback is not applicable"); -#endif -#endif -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(MultiheadAttentionVarlen, &plan, &run, &cleanup); - -} // namespace infinicore::op::mha_varlen_impl::flashattn diff --git a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn_moore.cc b/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn_moore.cc deleted file mode 100644 index 83f17866b..000000000 --- a/src/infinicore/ops/multi_head_attention_varlen/mha_varlen_flashattn_moore.cc +++ /dev/null @@ -1,139 +0,0 @@ -#if defined(ENABLE_MOORE_MATE_FLASH_ATTN) - -#include "infinicore/ops/mha_varlen.hpp" - -#include "infinicore/adaptor/aten_adaptor.hpp" - -#include -#include -#include -#include - -namespace infinicore::op::mha_varlen_impl::flashattn_moore { - -namespace py = pybind11; - -namespace { -class LocalMUSAStreamGuard { -public: - explicit LocalMUSAStreamGuard(const c10::musa::MUSAStream &s) - : prev_(c10::musa::getCurrentMUSAStream(s.device_index())) { - c10::musa::setCurrentMUSAStream(s); - } - ~LocalMUSAStreamGuard() { - c10::musa::setCurrentMUSAStream(prev_); - } - LocalMUSAStreamGuard(const LocalMUSAStreamGuard &) = delete; - LocalMUSAStreamGuard &operator=(const LocalMUSAStreamGuard &) = delete; - -private: - c10::musa::MUSAStream prev_; -}; -} // namespace - -struct PlannedMeta { - graph::GraphTensor out, q, k, v, cum_seqlens_q, cum_seqlens_k; - std::optional block_table; - int max_seqlen_q, max_seqlen_k; - std::optional alibi_slopes; - float scale; -}; - -void *plan(Tensor out, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &cum_seqlens_q, - const Tensor &cum_seqlens_k, - std::optional block_table, - int max_seqlen_q, - int max_seqlen_k, - std::optional alibi_slopes, - float scale) { - return new PlannedMeta{ - graph::GraphTensor(out), - graph::GraphTensor(q), - graph::GraphTensor(k), - graph::GraphTensor(v), - graph::GraphTensor(cum_seqlens_q), - graph::GraphTensor(cum_seqlens_k), - block_table ? std::optional(graph::GraphTensor(*block_table)) : std::nullopt, - max_seqlen_q, - max_seqlen_k, - alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, - scale}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - if (p->alibi_slopes.has_value()) { - throw std::runtime_error( - "[mha_varlen/moore] ALiBi not supported by mate flash_attn_varlen"); - } - if (!p->block_table.has_value()) { - throw std::runtime_error( - "[mha_varlen/moore] dense KV is not supported by mate flash_attn_varlen"); - } - - LocalMUSAStreamGuard guard(infinicore::adaptor::get_musa_stream()); - - auto out_tensor = infinicore::adaptor::to_aten_tensor(p->out); - auto q_tensor = infinicore::adaptor::to_aten_tensor(p->q); - auto k_cache = infinicore::adaptor::to_aten_tensor(p->k); - auto v_cache = infinicore::adaptor::to_aten_tensor(p->v); - auto cu_seqlens_q = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_q); - auto cu_seqlens_k = infinicore::adaptor::to_aten_tensor(p->cum_seqlens_k); - auto block_table = infinicore::adaptor::to_aten_tensor(*p->block_table); - - const int64_t block_size = k_cache.size(1); - - int max_seqlen_q_bound = static_cast(q_tensor.size(0)); - int max_seqlen_k_bound = static_cast(q_tensor.size(0)); - - try { - py::gil_scoped_acquire gil; - py::module_ wrapper = py::module_::import("infinicore.ops.moore_mate_flash_attn"); - - py::object py_q = py::cast(q_tensor); - py::object py_k = py::cast(k_cache); - py::object py_v = py::cast(v_cache); - py::object py_cuq = py::cast(cu_seqlens_q); - py::object py_cuk = py::cast(cu_seqlens_k); - py::object py_blk = py::cast(block_table); - - py::object result = wrapper.attr("moore_mate_flash_attn_prefill")( - py_q, - py_k, - py_v, - py_cuq, - py_cuk, - py_blk, - p->scale, - max_seqlen_q_bound, - max_seqlen_k_bound, - block_size, - true); - - at::Tensor result_t = result.cast(); - out_tensor.copy_(result_t); - } catch (const py::error_already_set &e) { - throw std::runtime_error( - std::string("[mha_varlen/moore] Python error: ") + e.what()); - } -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - MultiheadAttentionVarlen::plan_dispatcher().registerDevice(Device::Type::MOORE, &plan); - MultiheadAttentionVarlen::run_dispatcher().registerDevice(Device::Type::MOORE, &run); - MultiheadAttentionVarlen::cleanup_dispatcher().registerDevice(Device::Type::MOORE, &cleanup); - return true; -}(); - -} // namespace infinicore::op::mha_varlen_impl::flashattn_moore - -#endif // ENABLE_MOORE_MATE_FLASH_ATTN diff --git a/src/infinicore/ops/multi_margin_loss/multi_margin_loss.cc b/src/infinicore/ops/multi_margin_loss/multi_margin_loss.cc deleted file mode 100644 index c7d2c21ee..000000000 --- a/src/infinicore/ops/multi_margin_loss/multi_margin_loss.cc +++ /dev/null @@ -1,34 +0,0 @@ -#include "infinicore/ops/multi_margin_loss.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &MultiMarginLoss::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void MultiMarginLoss::execute(Tensor output, Tensor input, Tensor target, Tensor weight, int64_t p, float margin, int64_t reduction) { - dispatcher().lookup(context::getDevice().getType())(output, input, target, weight, p, margin, reduction); -} - -// 3. 函数式接口 -Tensor multi_margin_loss(Tensor input, Tensor target, Tensor weight, int64_t p, float margin, int64_t reduction) { - Shape output_shape; - if (reduction == 0) { // None - output_shape = {input->shape()[0]}; - } else { - output_shape = {}; // Scalar - } - - auto output = Tensor::empty(output_shape, input->dtype(), input->device()); - - multi_margin_loss_(output, input, target, weight, p, margin, reduction); - return output; -} - -void multi_margin_loss_(Tensor output, Tensor input, Tensor target, Tensor weight, int64_t p, float margin, int64_t reduction) { - MultiMarginLoss::execute(output, input, target, weight, p, margin, reduction); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/multi_margin_loss/multi_margin_loss_infiniop.cc b/src/infinicore/ops/multi_margin_loss/multi_margin_loss_infiniop.cc deleted file mode 100644 index 4214bfda0..000000000 --- a/src/infinicore/ops/multi_margin_loss/multi_margin_loss_infiniop.cc +++ /dev/null @@ -1,82 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/multi_margin_loss.hpp" -#include - -namespace infinicore::op::multi_margin_loss_impl::infiniop { - -// 定义描述符缓存 -thread_local common::OpCache caches( - 100, // capacity - [](infiniopMultiMarginLossDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyMultiMarginLossDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, Tensor target, Tensor weight, int64_t p, float margin, int64_t reduction) { - bool has_weight = static_cast(weight); - size_t seed; - if (has_weight) { - seed = hash_combine(output, input, target, weight, p, margin, reduction); - } else { - seed = hash_combine(output, input, target, size_t(0), p, margin, reduction); - } - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopMultiMarginLossDescriptor_t desc = nullptr; - infiniopTensorDescriptor_t weight_desc = nullptr; - const void *weight_data = nullptr; - - if (has_weight) { - weight_desc = weight->desc(); - weight_data = weight->data(); - } - - if (!desc_opt) { - // 创建描述符 - INFINICORE_CHECK_ERROR(infiniopCreateMultiMarginLossDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc(), - target->desc(), - weight_desc, - static_cast(p), - margin, - static_cast(reduction))); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 获取 Workspace 并执行 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetMultiMarginLossWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopMultiMarginLoss( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - target->data(), - weight_data, - context::getStream())); -} - -static bool registered = []() { - MultiMarginLoss::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::multi_margin_loss_impl::infiniop diff --git a/src/infinicore/ops/nrm2/nrm2.cc b/src/infinicore/ops/nrm2/nrm2.cc deleted file mode 100644 index 276e1015e..000000000 --- a/src/infinicore/ops/nrm2/nrm2.cc +++ /dev/null @@ -1,28 +0,0 @@ -#include "infinicore/ops/nrm2.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Nrm2); - -Nrm2::Nrm2(const Tensor &x, Tensor result) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, result); - INFINICORE_GRAPH_OP_DISPATCH(result->device().getType(), x, result); -} - -void Nrm2::execute(const Tensor &x, Tensor result) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Nrm2, x, result); -} - -Tensor nrm2(const Tensor &x) { - auto result = Tensor::empty({}, x->dtype(), x->device()); - nrm2_(x, result); - return result; -} - -void nrm2_(const Tensor &x, Tensor result) { - Nrm2::execute(x, result); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/nrm2/nrm2_infiniop.cc b/src/infinicore/ops/nrm2/nrm2_infiniop.cc deleted file mode 100644 index 3f3ca8c74..000000000 --- a/src/infinicore/ops/nrm2/nrm2_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/nrm2.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::nrm2_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Nrm2, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, x, result; -}; - -void *plan(const Tensor &x, Tensor result) { - size_t seed = hash_combine(x, result); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Nrm2, - seed, - x->desc(), result->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Nrm2, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x), - graph::GraphTensor(result)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopNrm2( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->x->data(), - planned->result->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Nrm2, &plan, &run, &cleanup); - -} // namespace infinicore::op::nrm2_impl::infiniop diff --git a/src/infinicore/ops/nsa_compress_paged_cache/nsa_compress_paged_cache.cc b/src/infinicore/ops/nsa_compress_paged_cache/nsa_compress_paged_cache.cc deleted file mode 100644 index 5d0269ff6..000000000 --- a/src/infinicore/ops/nsa_compress_paged_cache/nsa_compress_paged_cache.cc +++ /dev/null @@ -1,29 +0,0 @@ -#include "infinicore/ops/nsa_compress_paged_cache.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(NsaCompressPagedCache); - -NsaCompressPagedCache::NsaCompressPagedCache(Tensor k_cmp, Tensor v_cmp, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, int nsa_block_size, - bool update_last_only) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens); - INFINICORE_GRAPH_OP_DISPATCH(k_cmp->device().getType(), k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, nsa_block_size, update_last_only); -} - -void NsaCompressPagedCache::execute(Tensor k_cmp, Tensor v_cmp, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, int nsa_block_size, - bool update_last_only) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - NsaCompressPagedCache, - k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, nsa_block_size, update_last_only); -} - -void nsa_compress_paged_cache_(Tensor k_cmp, Tensor v_cmp, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, int nsa_block_size, - bool update_last_only) { - NsaCompressPagedCache::execute(k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, nsa_block_size, update_last_only); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/nsa_compress_paged_cache/nsa_compress_paged_cache_infiniop.cc b/src/infinicore/ops/nsa_compress_paged_cache/nsa_compress_paged_cache_infiniop.cc deleted file mode 100644 index 97f28d59b..000000000 --- a/src/infinicore/ops/nsa_compress_paged_cache/nsa_compress_paged_cache_infiniop.cc +++ /dev/null @@ -1,59 +0,0 @@ -#include "infinicore/ops/nsa_compress_paged_cache.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::nsa_compress_paged_cache_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, NsaCompressPagedCache, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens; -}; - -void *plan(Tensor k_cmp, Tensor v_cmp, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &cache_lens, int nsa_block_size, bool update_last_only) { - size_t seed = hash_combine(k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, nsa_block_size, update_last_only); - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, NsaCompressPagedCache, - seed, - k_cmp->desc(), v_cmp->desc(), k_cache->desc(), v_cache->desc(), - block_tables->desc(), cache_lens->desc(), nsa_block_size, static_cast(update_last_only)); - - INFINIOP_WORKSPACE_TENSOR(workspace, NsaCompressPagedCache, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(k_cmp), - graph::GraphTensor(v_cmp), - graph::GraphTensor(k_cache), - graph::GraphTensor(v_cache), - graph::GraphTensor(block_tables), - graph::GraphTensor(cache_lens)}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - INFINICORE_CHECK_ERROR( - infiniopNsaCompressPagedCache( - p->descriptor->desc, - p->workspace->data(), - p->workspace->numel(), - p->k_cmp->data(), - p->v_cmp->data(), - p->k_cache->data(), - p->v_cache->data(), - p->block_tables->data(), - p->cache_lens->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(NsaCompressPagedCache, &plan, &run, &cleanup); - -} // namespace infinicore::op::nsa_compress_paged_cache_impl::infiniop diff --git a/src/infinicore/ops/nsa_paged_attention/nsa_paged_attention.cc b/src/infinicore/ops/nsa_paged_attention/nsa_paged_attention.cc deleted file mode 100644 index 8d74a2e64..000000000 --- a/src/infinicore/ops/nsa_paged_attention/nsa_paged_attention.cc +++ /dev/null @@ -1,40 +0,0 @@ -#include "infinicore/ops/nsa_paged_attention.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(NsaPagedAttention); - -NsaPagedAttention::NsaPagedAttention(Tensor out, const Tensor &q, const Tensor &k_cmp, const Tensor &v_cmp, - const Tensor &k_cache, const Tensor &v_cache, const Tensor &block_tables, - const Tensor &kv_lens, const Tensor &gates, float scale, int nsa_block_size, - int window_size, int select_blocks) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, gates); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), - out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, gates, scale, nsa_block_size, window_size, select_blocks); -} - -void NsaPagedAttention::execute(Tensor out, const Tensor &q, const Tensor &k_cmp, const Tensor &v_cmp, - const Tensor &k_cache, const Tensor &v_cache, const Tensor &block_tables, - const Tensor &kv_lens, const Tensor &gates, float scale, int nsa_block_size, - int window_size, int select_blocks) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - NsaPagedAttention, - out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, gates, scale, nsa_block_size, window_size, select_blocks); -} - -Tensor nsa_paged_attention(const Tensor &q, const Tensor &k_cmp, const Tensor &v_cmp, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, const Tensor &gates, - float scale, int nsa_block_size, int window_size, int select_blocks) { - auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); - nsa_paged_attention_(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, gates, scale, nsa_block_size, window_size, select_blocks); - return out; -} - -void nsa_paged_attention_(Tensor out, const Tensor &q, const Tensor &k_cmp, const Tensor &v_cmp, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, const Tensor &gates, - float scale, int nsa_block_size, int window_size, int select_blocks) { - NsaPagedAttention::execute(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, kv_lens, gates, scale, nsa_block_size, window_size, select_blocks); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/nsa_paged_attention/nsa_paged_attention_infiniop.cc b/src/infinicore/ops/nsa_paged_attention/nsa_paged_attention_infiniop.cc deleted file mode 100644 index d2f26af06..000000000 --- a/src/infinicore/ops/nsa_paged_attention/nsa_paged_attention_infiniop.cc +++ /dev/null @@ -1,67 +0,0 @@ -#include "infinicore/ops/nsa_paged_attention.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::nsa_paged_attention_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, NsaPagedAttention, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, gates; -}; - -void *plan(Tensor out, const Tensor &q, const Tensor &k_cmp, const Tensor &v_cmp, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &cache_lens, const Tensor &gates, - float scale, int nsa_block_size, int window_size, int select_blocks) { - size_t seed = hash_combine(out, q, k_cmp, v_cmp, k_cache, v_cache, block_tables, cache_lens, gates, nsa_block_size, window_size, select_blocks); - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, NsaPagedAttention, - seed, - out->desc(), q->desc(), k_cmp->desc(), v_cmp->desc(), k_cache->desc(), v_cache->desc(), - block_tables->desc(), cache_lens->desc(), gates->desc(), - scale, nsa_block_size, window_size, select_blocks); - - INFINIOP_WORKSPACE_TENSOR(workspace, NsaPagedAttention, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(q), - graph::GraphTensor(k_cmp), - graph::GraphTensor(v_cmp), - graph::GraphTensor(k_cache), - graph::GraphTensor(v_cache), - graph::GraphTensor(block_tables), - graph::GraphTensor(cache_lens), - graph::GraphTensor(gates)}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - INFINICORE_CHECK_ERROR( - infiniopNsaPagedAttention( - p->descriptor->desc, - p->workspace->data(), - p->workspace->numel(), - p->out->data(), - p->q->data(), - p->k_cmp->data(), - p->v_cmp->data(), - p->k_cache->data(), - p->v_cache->data(), - p->block_tables->data(), - p->cache_lens->data(), - p->gates->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(NsaPagedAttention, &plan, &run, &cleanup); - -} // namespace infinicore::op::nsa_paged_attention_impl::infiniop diff --git a/src/infinicore/ops/ones/ones.cc b/src/infinicore/ops/ones/ones.cc deleted file mode 100644 index c28403eaf..000000000 --- a/src/infinicore/ops/ones/ones.cc +++ /dev/null @@ -1,13 +0,0 @@ -#include "infinicore/ops/ones.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Ones::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Ones::execute(Tensor output) { -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/pad/pad.cc b/src/infinicore/ops/pad/pad.cc deleted file mode 100644 index f1229b654..000000000 --- a/src/infinicore/ops/pad/pad.cc +++ /dev/null @@ -1,56 +0,0 @@ -#include "infinicore/ops/pad.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Pad); - -Pad::Pad(Tensor y, const Tensor &x, const std::vector &pad, const std::string &mode, double value) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, x, pad, mode, value); -} - -void Pad::execute(Tensor y, const Tensor &x, const std::vector &pad, const std::string &mode, double value) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Pad, y, x, pad, mode, value); -} - -static Shape infer_padded_shape(const Shape &in_shape, const std::vector &pad) { - if (pad.empty() || (pad.size() % 2) != 0) { - throw std::runtime_error("pad: pad must have even length."); - } - - Shape out_shape = in_shape; - const size_t ndim = out_shape.size(); - const size_t dims_padded = pad.size() / 2; - if (dims_padded > ndim) { - throw std::runtime_error("pad: pad has more dimensions than input."); - } - - for (size_t j = 0; j < dims_padded; ++j) { - const int left = pad[2 * j]; - const int right = pad[2 * j + 1]; - if (left < 0 || right < 0) { - throw std::runtime_error("pad: negative pad is not supported."); - } - const size_t dim = ndim - 1 - j; - out_shape[dim] += static_cast(left + right); - } - - return out_shape; -} - -Tensor pad(const Tensor &x, const std::vector &pad, const std::string &mode, double value) { - auto y_shape = infer_padded_shape(x->shape(), pad); - auto y = Tensor::empty(y_shape, x->dtype(), x->device()); - pad_(y, x, pad, mode, value); - return y; -} - -void pad_(Tensor y, const Tensor &x, const std::vector &pad, const std::string &mode, double value) { - Pad::execute(y, x, pad, mode, value); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/pad/pad_infiniop.cc b/src/infinicore/ops/pad/pad_infiniop.cc deleted file mode 100644 index 5b258ee36..000000000 --- a/src/infinicore/ops/pad/pad_infiniop.cc +++ /dev/null @@ -1,62 +0,0 @@ -#include "infinicore/ops/pad.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::pad_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Pad, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, y, x; -}; - -void *plan(Tensor y, - const Tensor &x, - const std::vector &pad, - const std::string &mode, - double value) { - size_t seed = hash_combine(y, x, mode, value, static_cast(pad.size())); - for (int v : pad) { - hash_combine(seed, v); - } - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Pad, - seed, - y->desc(), - x->desc(), - const_cast(pad.data()), - pad.size() * sizeof(int), - mode.c_str(), - value); - - INFINIOP_WORKSPACE_TENSOR(workspace, Pad, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(y), - graph::GraphTensor(x)}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopPad( - p->descriptor->desc, - p->workspace ? p->workspace->data() : nullptr, - p->workspace ? p->workspace->numel() : 0, - p->y->data(), - p->x->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Pad, &plan, &run, &cleanup); - -} // namespace infinicore::op::pad_impl::infiniop diff --git a/src/infinicore/ops/paged_attention/paged_attention.cc b/src/infinicore/ops/paged_attention/paged_attention.cc deleted file mode 100644 index 60de2ae66..000000000 --- a/src/infinicore/ops/paged_attention/paged_attention.cc +++ /dev/null @@ -1,38 +0,0 @@ -#include "infinicore/ops/paged_attention.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PagedAttention); - -PagedAttention::PagedAttention(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, - std::optional alibi_slopes, float scale) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, kv_lens); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), - out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale); -} - -void PagedAttention::execute(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, - std::optional alibi_slopes, float scale) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - PagedAttention, - out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale); -} - -Tensor paged_attention(const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, - std::optional alibi_slopes, float scale) { - auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); - paged_attention_(out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale); - return out; -} - -void paged_attention_(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &kv_lens, - std::optional alibi_slopes, float scale) { - PagedAttention::execute(out, q, k_cache, v_cache, block_tables, kv_lens, alibi_slopes, scale); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/paged_attention/paged_attention_infiniop.cc b/src/infinicore/ops/paged_attention/paged_attention_infiniop.cc deleted file mode 100644 index 733733a6b..000000000 --- a/src/infinicore/ops/paged_attention/paged_attention_infiniop.cc +++ /dev/null @@ -1,68 +0,0 @@ -#include "infinicore/ops/paged_attention.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::paged_attention_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, PagedAttention, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, q, k_cache, v_cache, block_tables, cache_lens; - std::optional alibi_slopes; - float scale; -}; - -void *plan(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, - const Tensor &block_tables, const Tensor &cache_lens, - std::optional alibi_slopes, float scale) { - size_t seed = hash_combine(out, q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes); - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, PagedAttention, - seed, - out->desc(), q->desc(), k_cache->desc(), v_cache->desc(), - block_tables->desc(), cache_lens->desc(), - alibi_slopes ? alibi_slopes.value()->desc() : nullptr, - scale); - - INFINIOP_WORKSPACE_TENSOR(workspace, PagedAttention, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(q), - graph::GraphTensor(k_cache), - graph::GraphTensor(v_cache), - graph::GraphTensor(block_tables), - graph::GraphTensor(cache_lens), - alibi_slopes ? std::optional(graph::GraphTensor(*alibi_slopes)) : std::nullopt, - scale}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR( - infiniopPagedAttention( - p->descriptor->desc, - p->workspace->data(), - p->workspace->numel(), - p->out->data(), - p->q->data(), - p->k_cache->data(), - p->v_cache->data(), - p->block_tables->data(), - p->cache_lens->data(), - p->alibi_slopes.has_value() ? p->alibi_slopes.value()->data() : nullptr, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(PagedAttention, &plan, &run, &cleanup); - -} // namespace infinicore::op::paged_attention_impl::infiniop diff --git a/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc b/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc deleted file mode 100644 index 1f4419129..000000000 --- a/src/infinicore/ops/paged_attention/paged_attention_infiniops.cc +++ /dev/null @@ -1,88 +0,0 @@ -#include "infinicore/ops/paged_attention.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/paged_attention_infinilm.h" - -#include -#include - -namespace infinicore::op::paged_attention_impl::infiniops { -namespace { -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -constexpr std::size_t kMaxPagedAttentionSplits = 8; - -std::size_t WorkspaceSizeInBytes(const Tensor &q) { - return kMaxPagedAttentionSplits - * static_cast(q->size(0)) - * static_cast(q->size(1)) - * static_cast(q->size(2) + 2) - * sizeof(float); -} - -struct PlannedMeta { - TensorMeta out, q, k_cache, v_cache, block_tables, cache_lens; - std::optional alibi_slopes; - graph::GraphTensor workspace, out_tensor, q_tensor, k_cache_tensor, v_cache_tensor, block_tables_tensor, cache_lens_tensor; - std::optional alibi_slopes_tensor; - float scale; -}; -} // namespace - -void *plan(Tensor out, - const Tensor &q, - const Tensor &k_cache, - const Tensor &v_cache, - const Tensor &block_tables, - const Tensor &cache_lens, - std::optional alibi_slopes, - float scale) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(out->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, cache_lens); - if (alibi_slopes) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, *alibi_slopes); - } - return new PlannedMeta{ - TensorMeta(out), TensorMeta(q), TensorMeta(k_cache), TensorMeta(v_cache), TensorMeta(block_tables), TensorMeta(cache_lens), - alibi_slopes ? std::optional{TensorMeta(*alibi_slopes)} : std::nullopt, - graph::GraphTensor(Tensor::empty({WorkspaceSizeInBytes(q)}, DataType::U8, out->device())), - graph::GraphTensor(out), graph::GraphTensor(q), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(block_tables), graph::GraphTensor(cache_lens), - alibi_slopes ? std::optional{graph::GraphTensor(*alibi_slopes)} : std::nullopt, - scale}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - handle.set_workspace(planned->workspace->data()); - handle.set_workspace_size_in_bytes(planned->workspace->numel()); - infini::ops::Config config; - infini::ops::PagedAttentionInfinilm::Call( - handle, - config, - planned->q.tensor(planned->q_tensor), - planned->k_cache.tensor(planned->k_cache_tensor), - planned->v_cache.tensor(planned->v_cache_tensor), - planned->block_tables.tensor(planned->block_tables_tensor), - planned->cache_lens.tensor(planned->cache_lens_tensor), - planned->alibi_slopes ? std::optional{planned->alibi_slopes->tensor(planned->alibi_slopes_tensor.value()->data())} : std::nullopt, - planned->scale, - planned->out.tensor(planned->out_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(PagedAttention::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(PagedAttention::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(PagedAttention::cleanup_dispatcher(), &cleanup); - return true; -}(); -} // namespace infinicore::op::paged_attention_impl::infiniops -#endif diff --git a/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill.cc b/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill.cc deleted file mode 100644 index d223198b4..000000000 --- a/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill.cc +++ /dev/null @@ -1,39 +0,0 @@ -#include "infinicore/ops/paged_attention_prefill.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &PagedAttentionPrefill::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void PagedAttentionPrefill::execute(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, - Tensor block_tables, Tensor kv_lens, Tensor cum_seqlens_q, - std::optional alibi_slopes, float scale) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q); - - infinicore::context::setDevice(out->device()); - - dispatcher().lookup(out->device().getType())(out, q, k_cache, v_cache, block_tables, - kv_lens, cum_seqlens_q, alibi_slopes, scale); -} - -Tensor paged_attention_prefill(Tensor q, Tensor k_cache, Tensor v_cache, - Tensor block_tables, Tensor kv_lens, Tensor cum_seqlens_q, - std::optional alibi_slopes, float scale) { - - auto out = Tensor::empty(q->shape(), q->dtype(), q->device()); - paged_attention_prefill_(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q, alibi_slopes, scale); - return out; -} - -void paged_attention_prefill_(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, - Tensor block_tables, Tensor kv_lens, Tensor cum_seqlens_q, - std::optional alibi_slopes, float scale) { - - PagedAttentionPrefill::execute(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q, alibi_slopes, scale); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniop.cc b/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniop.cc deleted file mode 100644 index 05de90cb8..000000000 --- a/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniop.cc +++ /dev/null @@ -1,70 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/paged_attention_prefill.hpp" -#include - -namespace infinicore::op::paged_attention_prefill_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopPagedAttentionPrefillDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyPagedAttentionPrefillDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, - Tensor block_tables, Tensor kv_lens, Tensor cum_seqlens_q, - std::optional alibi_slopes, float scale) { - size_t seed = hash_combine(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q, alibi_slopes, scale); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopPagedAttentionPrefillDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreatePagedAttentionPrefillDescriptor( - context::getInfiniopHandle(device), &desc, - out->desc(), - q->desc(), - k_cache->desc(), - v_cache->desc(), - block_tables->desc(), - kv_lens->desc(), - cum_seqlens_q->desc(), - alibi_slopes.has_value() ? alibi_slopes.value()->desc() : nullptr, - scale)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetPagedAttentionPrefillWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopPagedAttentionPrefill( - desc, - workspace->data(), - workspace_size, - out->data(), - q->data(), - k_cache->data(), - v_cache->data(), - block_tables->data(), - kv_lens->data(), - cum_seqlens_q->data(), - alibi_slopes.has_value() ? alibi_slopes.value()->data() : nullptr, - context::getStream())); -} - -static bool registered = []() { - PagedAttentionPrefill::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::paged_attention_prefill_impl::infiniop diff --git a/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniops.cc b/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniops.cc deleted file mode 100644 index bb842a6c6..000000000 --- a/src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniops.cc +++ /dev/null @@ -1,68 +0,0 @@ -#include "infinicore/ops/paged_attention_prefill.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/paged_attention_prefill_infinilm.h" - -#include - -namespace infinicore::op::paged_attention_prefill_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -void calculate(Tensor out, - Tensor q, - Tensor k_cache, - Tensor v_cache, - Tensor block_tables, - Tensor kv_lens, - Tensor cum_seqlens_q, - std::optional alibi_slopes, - float scale) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(out->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, kv_lens, cum_seqlens_q); - if (alibi_slopes) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, *alibi_slopes); - } - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - TensorMeta out_meta(out); - TensorMeta q_meta(q); - TensorMeta k_cache_meta(k_cache); - TensorMeta v_cache_meta(v_cache); - TensorMeta block_tables_meta(block_tables); - TensorMeta kv_lens_meta(kv_lens); - TensorMeta cum_seqlens_q_meta(cum_seqlens_q); - std::optional alibi_slopes_meta; - if (alibi_slopes) { - alibi_slopes_meta.emplace(*alibi_slopes); - } - - infini::ops::PagedAttentionPrefillInfinilm::Call( - handle, - config, - q_meta.tensor(q), - k_cache_meta.tensor(k_cache), - v_cache_meta.tensor(v_cache), - block_tables_meta.tensor(block_tables), - kv_lens_meta.tensor(kv_lens), - cum_seqlens_q_meta.tensor(cum_seqlens_q), - alibi_slopes_meta ? std::optional{alibi_slopes_meta->tensor(*alibi_slopes)} : std::nullopt, - scale, - out_meta.tensor(out)); -} - -} // namespace - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(PagedAttentionPrefill::dispatcher(), &calculate); - return true; -}(); - -} // namespace infinicore::op::paged_attention_prefill_impl::infiniops -#endif diff --git a/src/infinicore/ops/paged_caching/paged_caching.cc b/src/infinicore/ops/paged_caching/paged_caching.cc deleted file mode 100644 index afc8bf0c6..000000000 --- a/src/infinicore/ops/paged_caching/paged_caching.cc +++ /dev/null @@ -1,21 +0,0 @@ -#include "infinicore/ops/paged_caching.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PagedCaching); - -PagedCaching::PagedCaching(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, slot_mapping); - INFINICORE_GRAPH_OP_DISPATCH(k->device().getType(), k_cache, v_cache, k, v, slot_mapping); -} - -void PagedCaching::execute(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(PagedCaching, k_cache, v_cache, k, v, slot_mapping); -} - -void paged_caching_(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { - PagedCaching::execute(k_cache, v_cache, k, v, slot_mapping); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/paged_caching/paged_caching_infiniop.cc b/src/infinicore/ops/paged_caching/paged_caching_infiniop.cc deleted file mode 100644 index 5e8be049a..000000000 --- a/src/infinicore/ops/paged_caching/paged_caching_infiniop.cc +++ /dev/null @@ -1,57 +0,0 @@ -#include "infinicore/ops/paged_caching.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::paged_caching_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, PagedCaching, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - - graph::GraphTensor workspace, k_cache, v_cache, k, v, slot_mapping; -}; - -void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { - size_t key = hash_combine(k_cache, v_cache, k, v, slot_mapping); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, PagedCaching, - key, k_cache->desc(), v_cache->desc(), k->desc(), v->desc(), slot_mapping->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, PagedCaching, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(k_cache), - graph::GraphTensor(v_cache), - graph::GraphTensor(k), - graph::GraphTensor(v), - graph::GraphTensor(slot_mapping)}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR( - infiniopPagedCaching( - p->descriptor->desc, - p->workspace->data(), - p->workspace->numel(), - p->k_cache->data(), - p->v_cache->data(), - p->k->data(), - p->v->data(), - p->slot_mapping->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(PagedCaching, &plan, &run, &cleanup); - -} // namespace infinicore::op::paged_caching_impl::infiniop diff --git a/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc b/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc deleted file mode 100644 index 664472904..000000000 --- a/src/infinicore/ops/paged_caching/paged_caching_infiniops.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/paged_caching.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/paged_caching_infinilm.h" - -namespace infinicore::op::paged_caching_impl::infiniops { -namespace { -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; -struct PlannedMeta { - TensorMeta k_cache, v_cache, k, v, slot_mapping; - graph::GraphTensor k_cache_tensor, v_cache_tensor, k_tensor, v_tensor, slot_mapping_tensor; -}; -} // namespace - -void *plan(Tensor k_cache, Tensor v_cache, const Tensor &k, const Tensor &v, const Tensor &slot_mapping) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(k_cache->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(k_cache, v_cache, k, v, slot_mapping); - return new PlannedMeta{TensorMeta(k_cache), TensorMeta(v_cache), TensorMeta(k), TensorMeta(v), TensorMeta(slot_mapping), graph::GraphTensor(k_cache), graph::GraphTensor(v_cache), graph::GraphTensor(k), graph::GraphTensor(v), graph::GraphTensor(slot_mapping)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - infini::ops::PagedCachingInfinilm::Call( - handle, - config, - planned->k.tensor(planned->k_tensor), - planned->v.tensor(planned->v_tensor), - planned->slot_mapping.tensor(planned->slot_mapping_tensor), - planned->k_cache.tensor(planned->k_cache_tensor), - planned->v_cache.tensor(planned->v_cache_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(PagedCaching::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(PagedCaching::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(PagedCaching::cleanup_dispatcher(), &cleanup); - return true; -}(); -} // namespace infinicore::op::paged_caching_impl::infiniops -#endif diff --git a/src/infinicore/ops/per_channel_quant_i8/per_channel_quant_i8.cc b/src/infinicore/ops/per_channel_quant_i8/per_channel_quant_i8.cc deleted file mode 100644 index 40ddefbfe..000000000 --- a/src/infinicore/ops/per_channel_quant_i8/per_channel_quant_i8.cc +++ /dev/null @@ -1,20 +0,0 @@ -#include "infinicore/ops/per_channel_quant_i8.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PerChannelQuantI8); - -PerChannelQuantI8::PerChannelQuantI8(const Tensor &x, Tensor x_packed, Tensor x_scale) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, x_packed, x_scale); - INFINICORE_GRAPH_OP_DISPATCH(x->device().getType(), x, x_packed, x_scale); -} - -void PerChannelQuantI8::execute(const Tensor &x, Tensor x_packed, Tensor x_scale) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(PerChannelQuantI8, x, x_packed, x_scale); -} - -void per_channel_quant_i8_(const Tensor &x, Tensor x_packed, Tensor x_scale) { - PerChannelQuantI8::execute(x, x_packed, x_scale); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/per_channel_quant_i8/per_channel_quant_i8_infiniop.cc b/src/infinicore/ops/per_channel_quant_i8/per_channel_quant_i8_infiniop.cc deleted file mode 100644 index 569c9fdae..000000000 --- a/src/infinicore/ops/per_channel_quant_i8/per_channel_quant_i8_infiniop.cc +++ /dev/null @@ -1,56 +0,0 @@ -#include "../../utils.hpp" -#include "../infiniop_impl.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/per_channel_quant_i8.hpp" -#include - -namespace infinicore::op::per_channel_quant_i8_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, PerChannelQuantI8, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, x, x_packed, x_scale; -}; - -void *plan(const Tensor &x, Tensor x_packed, Tensor x_scale) { - size_t seed = hash_combine(x, x_packed, x_scale); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, PerChannelQuantI8, - seed, - x_packed->desc(), x_scale->desc(), nullptr, x->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, PerChannelQuantI8, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x), - graph::GraphTensor(x_packed), - graph::GraphTensor(x_scale)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopPerChannelQuantI8( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->x_packed->data(), - planned->x_scale->data(), - nullptr, - planned->x->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(PerChannelQuantI8, &plan, &run, &cleanup); - -} // namespace infinicore::op::per_channel_quant_i8_impl::infiniop diff --git a/src/infinicore/ops/prelu/prelu.cc b/src/infinicore/ops/prelu/prelu.cc deleted file mode 100644 index 42441f99b..000000000 --- a/src/infinicore/ops/prelu/prelu.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/prelu.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Prelu); - -Prelu::Prelu(Tensor out, const Tensor &input, const Tensor &weight) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input, weight); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, input, weight); -} - -void Prelu::execute(Tensor out, const Tensor &input, const Tensor &weight) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Prelu, out, input, weight); -} - -Tensor prelu(const Tensor &input, const Tensor &weight) { - auto out = Tensor::empty(input->shape(), input->dtype(), input->device()); - prelu_(out, input, weight); - return out; -} - -void prelu_(Tensor out, const Tensor &input, const Tensor &weight) { - Prelu::execute(out, input, weight); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/prelu/prelu_infiniop.cc b/src/infinicore/ops/prelu/prelu_infiniop.cc deleted file mode 100644 index b4d66a4fe..000000000 --- a/src/infinicore/ops/prelu/prelu_infiniop.cc +++ /dev/null @@ -1,56 +0,0 @@ -#include "infinicore/ops/prelu.hpp" - -#include "infiniop/ops/prelu.h" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::prelu_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Prelu, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, input, weight; -}; - -void *plan(Tensor out, const Tensor &input, const Tensor &weight) { - size_t seed = hash_combine(out, input, weight); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Prelu, - seed, - out->desc(), - input->desc(), - weight->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Prelu, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(input), - graph::GraphTensor(weight)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopPrelu( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->input->data(), - planned->weight->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Prelu, &plan, &run, &cleanup); - -} // namespace infinicore::op::prelu_impl::infiniop diff --git a/src/infinicore/ops/prepare_moe_input/prepare_moe_input.cc b/src/infinicore/ops/prepare_moe_input/prepare_moe_input.cc deleted file mode 100644 index 7a92adbf2..000000000 --- a/src/infinicore/ops/prepare_moe_input/prepare_moe_input.cc +++ /dev/null @@ -1,124 +0,0 @@ -#include "infinicore/ops/prepare_moe_input.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PrepareMoeInput); - -PrepareMoeInput::PrepareMoeInput(Tensor expert_offsets, - Tensor blockscale_offsets, - Tensor problem_sizes1, - Tensor problem_sizes2, - Tensor input_permutation, - Tensor output_permutation, - const Tensor &topk_ids, - const size_t num_experts, - const size_t n, - const size_t k) { - if (blockscale_offsets) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE( - expert_offsets, blockscale_offsets, problem_sizes1, problem_sizes2, input_permutation, output_permutation, topk_ids); - } else { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE( - expert_offsets, problem_sizes1, problem_sizes2, input_permutation, output_permutation, topk_ids); - } - INFINICORE_GRAPH_OP_DISPATCH( - expert_offsets->device().getType(), - expert_offsets, - blockscale_offsets, - problem_sizes1, - problem_sizes2, - input_permutation, - output_permutation, - topk_ids, - num_experts, - n, - k); -} - -void PrepareMoeInput::execute(Tensor expert_offsets, - Tensor blockscale_offsets, - Tensor problem_sizes1, - Tensor problem_sizes2, - Tensor input_permutation, - Tensor output_permutation, - const Tensor &topk_ids, - const size_t num_experts, - const size_t n, - const size_t k) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN( - PrepareMoeInput, - expert_offsets, - blockscale_offsets, - problem_sizes1, - problem_sizes2, - input_permutation, - output_permutation, - topk_ids, - num_experts, - n, - k); -} - -PrepareMoeInputOutput prepare_moe_input( - const Tensor &topk_ids, - size_t num_experts, - size_t n, - size_t k) { - auto shape = topk_ids->shape(); - INFINICORE_ASSERT(shape.size() == 2); - const size_t topk_length = shape[0] * shape[1]; - - auto expert_offsets = Tensor::empty({num_experts + 1}, DataType::I32, topk_ids->device()); - auto problem_sizes1 = Tensor::empty({num_experts, 3}, DataType::I32, topk_ids->device()); - auto problem_sizes2 = Tensor::empty({num_experts, 3}, DataType::I32, topk_ids->device()); - auto input_permutation = Tensor::empty({topk_length}, DataType::I32, topk_ids->device()); - auto output_permutation = Tensor::empty({topk_length}, DataType::I32, topk_ids->device()); - - prepare_moe_input_( - expert_offsets, - Tensor(), - problem_sizes1, - problem_sizes2, - input_permutation, - output_permutation, - topk_ids, - num_experts, - n, - k); - - return PrepareMoeInputOutput{ - expert_offsets, - problem_sizes1, - problem_sizes2, - input_permutation, - output_permutation, - }; -} - -void prepare_moe_input_( - Tensor expert_offsets, - Tensor blockscale_offsets, - Tensor problem_sizes1, - Tensor problem_sizes2, - Tensor input_permutation, - Tensor output_permutation, - const Tensor &topk_ids, - size_t num_experts, - size_t n, - size_t k) { - PrepareMoeInput::execute( - expert_offsets, - blockscale_offsets, - problem_sizes1, - problem_sizes2, - input_permutation, - output_permutation, - topk_ids, - num_experts, - n, - k); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/prepare_moe_input/prepare_moe_input_infiniop.cc b/src/infinicore/ops/prepare_moe_input/prepare_moe_input_infiniop.cc deleted file mode 100644 index 681b7f8cf..000000000 --- a/src/infinicore/ops/prepare_moe_input/prepare_moe_input_infiniop.cc +++ /dev/null @@ -1,105 +0,0 @@ -#include "infinicore/ops/prepare_moe_input.hpp" - -#include "../infiniop_impl.hpp" - -#include - -namespace infinicore::op::prepare_moe_input_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, PrepareMoeInput, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace; - graph::GraphTensor expert_offsets; - std::optional blockscale_offsets; - graph::GraphTensor problem_sizes1; - graph::GraphTensor problem_sizes2; - graph::GraphTensor input_permutation; - graph::GraphTensor output_permutation; - graph::GraphTensor topk_ids; -}; - -void *plan(Tensor expert_offsets, - Tensor blockscale_offsets, - Tensor problem_sizes1, - Tensor problem_sizes2, - Tensor input_permutation, - Tensor output_permutation, - const Tensor &topk_ids, - const size_t num_experts, - const size_t n, - const size_t k) { - size_t seed = hash_combine( - expert_offsets, - blockscale_offsets, - problem_sizes1, - problem_sizes2, - input_permutation, - output_permutation, - topk_ids, - num_experts, - n, - k); - - infiniopTensorDescriptor_t blockscale_desc = blockscale_offsets ? blockscale_offsets->desc() : nullptr; - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, - descriptor, - PrepareMoeInput, - seed, - expert_offsets->desc(), - blockscale_desc, - problem_sizes1->desc(), - problem_sizes2->desc(), - input_permutation->desc(), - output_permutation->desc(), - topk_ids->desc(), - num_experts, - n, - k); - - INFINIOP_WORKSPACE_TENSOR(workspace, PrepareMoeInput, descriptor); - - std::optional blockscale_offsets_graph; - if (blockscale_offsets) { - blockscale_offsets_graph.emplace(blockscale_offsets); - } - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(expert_offsets), - blockscale_offsets_graph, - graph::GraphTensor(problem_sizes1), - graph::GraphTensor(problem_sizes2), - graph::GraphTensor(input_permutation), - graph::GraphTensor(output_permutation), - graph::GraphTensor(topk_ids)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopPrepareMoeInput( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->expert_offsets->data(), - planned->blockscale_offsets ? (*planned->blockscale_offsets)->data() : nullptr, - planned->problem_sizes1->data(), - planned->problem_sizes2->data(), - planned->input_permutation->data(), - planned->output_permutation->data(), - planned->topk_ids->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(PrepareMoeInput, &plan, &run, cleanup); - -} // namespace infinicore::op::prepare_moe_input_impl::infiniop diff --git a/src/infinicore/ops/quant/per_tensor_quant_int8/per_tensor_quant_int8.cc b/src/infinicore/ops/quant/per_tensor_quant_int8/per_tensor_quant_int8.cc deleted file mode 100644 index 498af541e..000000000 --- a/src/infinicore/ops/quant/per_tensor_quant_int8/per_tensor_quant_int8.cc +++ /dev/null @@ -1,26 +0,0 @@ -#include "../../../utils.hpp" -#include "infinicore/ops/per_tensor_quant_i8.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(PerTensorQuantI8); - -PerTensorQuantI8::PerTensorQuantI8(const Tensor &x, Tensor x_packed, Tensor x_scale, Tensor x_zero, bool is_static) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, x_packed, x_scale, x_zero); - INFINICORE_GRAPH_OP_DISPATCH(x->device().getType(), x, x_packed, x_scale, x_zero, is_static); -} - -void PerTensorQuantI8::execute(const Tensor &x, Tensor x_packed, Tensor x_scale, Tensor x_zero, bool is_static) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(PerTensorQuantI8, x, x_packed, x_scale, x_zero, is_static); -} - -void per_tensor_quant_i8_(const Tensor &x, Tensor x_packed, Tensor x_scale, Tensor x_zero, bool is_static) { - PerTensorQuantI8::execute(x, x_packed, x_scale, x_zero, is_static); -} - -Tensor per_tensor_quant_i8(const Tensor &x, Tensor x_scale, Tensor x_zero, bool is_static) { - auto x_packed = Tensor::strided_empty(x->shape(), x->strides(), infinicore::DataType::I8, x->device()); - PerTensorQuantI8::execute(x, x_packed, x_scale, x_zero, is_static); - return x_packed; -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/quant/per_tensor_quant_int8/per_tensor_quant_int8_infiniop.cc b/src/infinicore/ops/quant/per_tensor_quant_int8/per_tensor_quant_int8_infiniop.cc deleted file mode 100644 index 59c9cc243..000000000 --- a/src/infinicore/ops/quant/per_tensor_quant_int8/per_tensor_quant_int8_infiniop.cc +++ /dev/null @@ -1,56 +0,0 @@ -#include "../../infiniop_impl.hpp" -#include "infinicore/ops/per_tensor_quant_i8.hpp" - -namespace infinicore::op::per_tensor_quant_i8_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, PerTensorQuantI8, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, x, x_packed, x_scale, x_zero; - const bool is_static; -}; - -void *plan(const Tensor &x, Tensor x_packed, Tensor x_scale, Tensor x_zero, bool is_static) { - size_t seed = hash_combine(x, x_packed, x_scale); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, PerTensorQuantI8, - seed, - x_packed->desc(), x_scale->desc(), (x_zero ? x_zero->desc() : nullptr), x->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, PerTensorQuantI8, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x), - graph::GraphTensor(x_packed), - graph::GraphTensor(x_scale), - graph::GraphTensor(x_zero), - is_static}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - const bool is_static = planned->is_static; - INFINICORE_CHECK_ERROR(infiniopPerTensorQuantI8( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->x_packed->data(), - planned->x_scale->data(), - nullptr, - planned->x->data(), - is_static, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(PerTensorQuantI8, &plan, &run, &cleanup); - -} // namespace infinicore::op::per_tensor_quant_i8_impl::infiniop diff --git a/src/infinicore/ops/quickgelu/quickgelu.cc b/src/infinicore/ops/quickgelu/quickgelu.cc deleted file mode 100644 index 4ab46c49a..000000000 --- a/src/infinicore/ops/quickgelu/quickgelu.cc +++ /dev/null @@ -1,37 +0,0 @@ -#include "infinicore/ops/quickgelu.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &QuickGelu::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void QuickGelu::execute(Tensor output, Tensor input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No QuickGelu implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor quick_gelu(Tensor input) { - Shape shape = input->shape(); - auto output = Tensor::empty(shape, input->dtype(), input->device()); - quick_gelu_(output, input); - return output; -} - -void quick_gelu_(Tensor output, Tensor input) { - QuickGelu::execute(output, input); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/quickgelu/quickgelu_infiniop.cc b/src/infinicore/ops/quickgelu/quickgelu_infiniop.cc deleted file mode 100644 index 3b18c315a..000000000 --- a/src/infinicore/ops/quickgelu/quickgelu_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/quickgelu.hpp" -#include - -namespace infinicore::op::quickgelu_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopQuickGeluDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyQuickGeluDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopQuickGeluDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateQuickGeluDescriptor( - context::getInfiniopHandle(device), &desc, - output->desc(), input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetQuickGeluWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopQuickGelu( - desc, workspace->data(), workspace_size, - output->data(), input->data(), context::getStream())); -} - -static bool registered = []() { - QuickGelu::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::quickgelu_impl::infiniop diff --git a/src/infinicore/ops/random_sample/random_sample.cc b/src/infinicore/ops/random_sample/random_sample.cc deleted file mode 100644 index bb0b41ce8..000000000 --- a/src/infinicore/ops/random_sample/random_sample.cc +++ /dev/null @@ -1,42 +0,0 @@ -#include "infinicore/ops/random_sample.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &RandomSample::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void RandomSample::execute( - Tensor indices, Tensor logits, - float random_val, float topp, int topk, float temperature) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(indices, logits); - infinicore::context::setDevice(logits->device()); - dispatcher().lookup(logits->device().getType())( - indices, logits, random_val, topp, topk, temperature); -} - -Tensor random_sample( - Tensor logits, - float random_val, - float topp, - int topk, - float temperature) { - auto indices = Tensor::empty({}, DataType::I32, logits->device()); - random_sample_(indices, logits, random_val, topp, topk, temperature); - return indices; -} - -void random_sample_( - Tensor indices, - Tensor logits, - float random_val, - float topp, - int topk, - float temperature) { - RandomSample::execute(indices, logits, random_val, topp, topk, temperature); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/random_sample/random_sample_infiniop.cc b/src/infinicore/ops/random_sample/random_sample_infiniop.cc deleted file mode 100644 index af6ac0fcc..000000000 --- a/src/infinicore/ops/random_sample/random_sample_infiniop.cc +++ /dev/null @@ -1,64 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/random_sample.hpp" -#include - -namespace infinicore::op::random_sample_impl::infiniop_backend { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopRandomSampleDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyRandomSampleDescriptor(desc)); - desc = nullptr; - } - }); - -static void calculate( - Tensor indices, - Tensor logits, - float random_val, - float topp, - int topk, - float temperature) { - // cache per (result desc + logits desc) on device - size_t seed = hash_combine(indices, logits); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopRandomSampleDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateRandomSampleDescriptor( - context::getInfiniopHandle(device), &desc, - indices->desc(), logits->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetRandomSampleWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopRandomSample( - desc, - workspace->data(), workspace_size, - indices->data(), logits->data(), - random_val, topp, topk, temperature, - context::getStream())); -} - -} // namespace infinicore::op::random_sample_impl::infiniop_backend - -namespace infinicore::op { - -static bool registered = []() { - RandomSample::dispatcher().registerAll(&random_sample_impl::infiniop_backend::calculate, false); - return true; -}(); - -} // namespace infinicore::op diff --git a/src/infinicore/ops/random_sample/random_sample_infiniops.cc b/src/infinicore/ops/random_sample/random_sample_infiniops.cc deleted file mode 100644 index 5146b1443..000000000 --- a/src/infinicore/ops/random_sample/random_sample_infiniops.cc +++ /dev/null @@ -1,42 +0,0 @@ -#include "infinicore/ops/random_sample.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/random_sample_infinilm.h" - -namespace infinicore::op::random_sample_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -void calculate(Tensor indices, Tensor logits, float random_val, float topp, int topk, float temperature) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(indices->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(indices, logits); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - TensorMeta indices_meta(indices); - TensorMeta logits_meta(logits); - infini::ops::RandomSampleInfinilm::Call( - handle, - config, - logits_meta.tensor(logits), - random_val, - topp, - static_cast(topk), - temperature, - indices_meta.tensor(indices)); -} - -} // namespace - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(RandomSample::dispatcher(), &calculate); - return true; -}(); - -} // namespace infinicore::op::random_sample_impl::infiniops -#endif diff --git a/src/infinicore/ops/rearrange/rearrange.cc b/src/infinicore/ops/rearrange/rearrange.cc deleted file mode 100644 index 191d0871f..000000000 --- a/src/infinicore/ops/rearrange/rearrange.cc +++ /dev/null @@ -1,32 +0,0 @@ -#include "infinicore/ops/rearrange.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Rearrange); - -Rearrange::Rearrange(Tensor y, const Tensor &x) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, x); -} - -void Rearrange::execute(Tensor y, const Tensor &x) { - auto op = std::make_shared(y, x); - if (context::isGraphRecording()) { - context::addGraphOperator(op); - } else { - op->run(); - } -} - -Tensor rearrange(const Tensor &x) { - auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); - rearrange_(y, x); - return y; -} - -void rearrange_(Tensor y, const Tensor &x) { - Rearrange::execute(y, x); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/rearrange/rearrange_infiniop.cc b/src/infinicore/ops/rearrange/rearrange_infiniop.cc deleted file mode 100644 index f30b09e79..000000000 --- a/src/infinicore/ops/rearrange/rearrange_infiniop.cc +++ /dev/null @@ -1,46 +0,0 @@ -#include "infinicore/ops/rearrange.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::rearrange_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Rearrange, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor y, x; -}; - -void *plan(Tensor y, const Tensor &x) { - size_t seed = hash_combine(y, x); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Rearrange, - seed, y->desc(), - x->desc()); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(y), - graph::GraphTensor(x)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR( - infiniopRearrange( - planned->descriptor->desc, - planned->y->data(), - planned->x->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Rearrange, &plan, &run, &cleanup); - -} // namespace infinicore::op::rearrange_impl::infiniop diff --git a/src/infinicore/ops/rearrange/rearrange_infiniops.cc b/src/infinicore/ops/rearrange/rearrange_infiniops.cc deleted file mode 100644 index 24b275763..000000000 --- a/src/infinicore/ops/rearrange/rearrange_infiniops.cc +++ /dev/null @@ -1,41 +0,0 @@ -#include "infinicore/ops/rearrange.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/rearrange_infinilm.h" - -namespace infinicore::op::rearrange_impl::infiniops { -namespace { -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; -struct PlannedMeta { - TensorMeta y, x; - graph::GraphTensor y_tensor, x_tensor; -}; -} // namespace - -void *plan(Tensor y, const Tensor &x) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(y->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - return new PlannedMeta{TensorMeta(y), TensorMeta(x), graph::GraphTensor(y), graph::GraphTensor(x)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - infini::ops::RearrangeInfinilm::Call( - handle, config, planned->x.tensor(planned->x_tensor), planned->y.tensor(planned->y_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -// ReArrange is used by Tensor::copy_from during test result conversion. Keep -// the InfiniCore/InfiniOp implementation active until this adapter gets the same -// CUDA shim treatment as Gemm. -} // namespace infinicore::op::rearrange_impl::infiniops -#endif diff --git a/src/infinicore/ops/reciprocal/reciprocal.cc b/src/infinicore/ops/reciprocal/reciprocal.cc deleted file mode 100644 index e24b7aa50..000000000 --- a/src/infinicore/ops/reciprocal/reciprocal.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/reciprocal.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Reciprocal::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Reciprocal::execute(Tensor y, Tensor x) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - infinicore::context::setDevice(y->device()); - dispatcher().lookup(y->device().getType())(y, x); -} - -Tensor reciprocal(Tensor x) { - auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); - reciprocal_(y, x); - return y; -} - -void reciprocal_(Tensor y, Tensor x) { - Reciprocal::execute(y, x); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/reciprocal/reciprocal_infiniop.cc b/src/infinicore/ops/reciprocal/reciprocal_infiniop.cc deleted file mode 100644 index fa323ed15..000000000 --- a/src/infinicore/ops/reciprocal/reciprocal_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/reciprocal.hpp" -#include - -namespace infinicore::op::reciprocal_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopReciprocalDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyReciprocalDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor y, Tensor x) { - size_t seed = hash_combine(y, x); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopReciprocalDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateReciprocalDescriptor( - context::getInfiniopHandle(device), &desc, - y->desc(), x->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetReciprocalWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopReciprocal( - desc, workspace->data(), workspace_size, - y->data(), x->data(), context::getStream())); -} - -static bool registered = []() { - Reciprocal::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::reciprocal_impl::infiniop diff --git a/src/infinicore/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule.cc b/src/infinicore/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule.cc deleted file mode 100644 index 06edd9955..000000000 --- a/src/infinicore/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule.cc +++ /dev/null @@ -1,166 +0,0 @@ -#include "infinicore/ops/recurrent_gated_delta_rule.hpp" -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(RecurrentGatedDeltaRule); - -RecurrentGatedDeltaRule::RecurrentGatedDeltaRule(Tensor out, - Tensor initial_state, - std::optional final_state, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - std::optional initial_state_indices, - std::optional final_state_indices, - bool use_qk_l2norm) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state, q, k, v, g, beta); - if (final_state.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_state.value()); - } - if (initial_state_indices.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, initial_state_indices.value()); - } - if (final_state_indices.has_value()) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, final_state_indices.value()); - } - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), - out, - initial_state, - final_state, - q, - k, - v, - g, - beta, - initial_state_indices, - final_state_indices, - use_qk_l2norm); -} - -void RecurrentGatedDeltaRule::execute(Tensor out, - Tensor initial_state, - std::optional final_state, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - std::optional initial_state_indices, - std::optional final_state_indices, - bool use_qk_l2norm) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(RecurrentGatedDeltaRule, - out, - initial_state, - final_state, - q, - k, - v, - g, - beta, - initial_state_indices, - final_state_indices, - use_qk_l2norm); -} - -static Tensor ensure_4d_sequence_tensor(const Tensor &x, const char *name) { - if (x->shape().size() == 4) { - return x; - } - if (x->shape().size() == 3) { - return x->unsqueeze(1); - } - throw std::runtime_error(std::string("recurrent_gated_delta_rule expects ") + name + " with shape [B, T, H, D] or [B, H, D]"); -} - -static Shape recurrent_output_shape(const Tensor &v) { - const auto &shape = v->shape(); - return {shape[0], shape[1], shape[2], shape[3]}; -} - -Tensor recurrent_gated_delta_rule(const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - const Tensor &initial_state, - bool use_qk_l2norm) { - Tensor q4 = ensure_4d_sequence_tensor(q, "q"); - Tensor k4 = ensure_4d_sequence_tensor(k, "k"); - Tensor v4 = ensure_4d_sequence_tensor(v, "v"); - auto out = Tensor::empty(recurrent_output_shape(v4), v4->dtype(), v4->device()); - Shape final_state_shape = {v4->shape()[0], v4->shape()[2], v4->shape()[3], q4->shape()[3]}; - auto final_state = Tensor::empty(final_state_shape, initial_state->dtype(), initial_state->device()); - recurrent_gated_delta_rule_(out, - initial_state, - final_state, - q4, - k4, - v4, - g, - beta, - std::nullopt, - std::nullopt, - use_qk_l2norm); - return out; -} - -Tensor recurrent_gated_delta_rule_indexed(const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - Tensor initial_state, - const Tensor &initial_state_indices, - const Tensor &final_state_indices, - bool use_qk_l2norm) { - Tensor q4 = ensure_4d_sequence_tensor(q, "q"); - Tensor k4 = ensure_4d_sequence_tensor(k, "k"); - Tensor v4 = ensure_4d_sequence_tensor(v, "v"); - auto out = Tensor::empty(recurrent_output_shape(v4), v4->dtype(), v4->device()); - recurrent_gated_delta_rule_(out, - initial_state, - std::nullopt, - q4, - k4, - v4, - g, - beta, - initial_state_indices, - final_state_indices, - use_qk_l2norm); - return out; -} - -void recurrent_gated_delta_rule_(Tensor out, - Tensor initial_state, - std::optional final_state, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - std::optional initial_state_indices, - std::optional final_state_indices, - bool use_qk_l2norm) { - Tensor q4 = ensure_4d_sequence_tensor(q, "q"); - Tensor k4 = ensure_4d_sequence_tensor(k, "k"); - Tensor v4 = ensure_4d_sequence_tensor(v, "v"); - RecurrentGatedDeltaRule::execute(out, - initial_state, - final_state, - q4, - k4, - v4, - g, - beta, - initial_state_indices, - final_state_indices, - use_qk_l2norm); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule_infiniop.cc b/src/infinicore/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule_infiniop.cc deleted file mode 100644 index ba52b2c40..000000000 --- a/src/infinicore/ops/recurrent_gated_delta_rule/recurrent_gated_delta_rule_infiniop.cc +++ /dev/null @@ -1,99 +0,0 @@ -#include "infinicore/ops/recurrent_gated_delta_rule.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::recurrent_gated_delta_rule_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, RecurrentGatedDeltaRule, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, initial_state, q, k, v, g, beta; - std::optional final_state; - std::optional initial_state_indices; - std::optional final_state_indices; -}; - -void *plan(Tensor out, - Tensor initial_state, - std::optional final_state, - const Tensor &q, - const Tensor &k, - const Tensor &v, - const Tensor &g, - const Tensor &beta, - std::optional initial_state_indices, - std::optional final_state_indices, - bool use_qk_l2norm) { - size_t seed = hash_combine(out, - initial_state, - final_state, - q, - k, - v, - g, - beta, - initial_state_indices, - final_state_indices, - use_qk_l2norm); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, RecurrentGatedDeltaRule, - seed, - out->desc(), - initial_state->desc(), - final_state.has_value() ? final_state.value()->desc() : nullptr, - q->desc(), - k->desc(), - v->desc(), - g->desc(), - beta->desc(), - initial_state_indices.has_value() ? initial_state_indices.value()->desc() : nullptr, - final_state_indices.has_value() ? final_state_indices.value()->desc() : nullptr, - use_qk_l2norm); - - INFINIOP_WORKSPACE_TENSOR(workspace, RecurrentGatedDeltaRule, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(initial_state), - graph::GraphTensor(q), - graph::GraphTensor(k), - graph::GraphTensor(v), - graph::GraphTensor(g), - graph::GraphTensor(beta), - final_state.has_value() ? std::optional(graph::GraphTensor(final_state.value())) : std::nullopt, - initial_state_indices.has_value() ? std::optional(graph::GraphTensor(initial_state_indices.value())) : std::nullopt, - final_state_indices.has_value() ? std::optional(graph::GraphTensor(final_state_indices.value())) : std::nullopt}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopRecurrentGatedDeltaRule( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->initial_state->data(), - planned->final_state.has_value() ? planned->final_state.value()->data() : nullptr, - planned->q->data(), - planned->k->data(), - planned->v->data(), - planned->g->data(), - planned->beta->data(), - planned->initial_state_indices.has_value() ? planned->initial_state_indices.value()->data() : nullptr, - planned->final_state_indices.has_value() ? planned->final_state_indices.value()->data() : nullptr, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(RecurrentGatedDeltaRule, &plan, &run, &cleanup); - -} // namespace infinicore::op::recurrent_gated_delta_rule_impl::infiniop diff --git a/src/infinicore/ops/relu/relu.cc b/src/infinicore/ops/relu/relu.cc deleted file mode 100644 index dceb618b1..000000000 --- a/src/infinicore/ops/relu/relu.cc +++ /dev/null @@ -1,37 +0,0 @@ -#include "infinicore/ops/relu.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &Relu::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Relu::execute(Tensor output, Tensor input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Relu implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor relu(Tensor input) { - Shape shape = input->shape(); - auto output = Tensor::empty(shape, input->dtype(), input->device()); - relu_(output, input); - return output; -} - -void relu_(Tensor output, Tensor input) { - Relu::execute(output, input); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/relu/relu_infiniop.cc b/src/infinicore/ops/relu/relu_infiniop.cc deleted file mode 100644 index dc80535a8..000000000 --- a/src/infinicore/ops/relu/relu_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/relu.hpp" -#include - -namespace infinicore::op::relu_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopReluDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyReluDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopReluDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateReluDescriptor( - context::getInfiniopHandle(device), &desc, - output->desc(), input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetReluWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopRelu( - desc, workspace->data(), workspace_size, - output->data(), input->data(), context::getStream())); -} - -static bool registered = []() { - Relu::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::relu_impl::infiniop diff --git a/src/infinicore/ops/relu/relu_infiniops.cc b/src/infinicore/ops/relu/relu_infiniops.cc deleted file mode 100644 index 7890765fe..000000000 --- a/src/infinicore/ops/relu/relu_infiniops.cc +++ /dev/null @@ -1,38 +0,0 @@ -#include "infinicore/ops/relu.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/relu_infinilm.h" - -namespace infinicore::op::relu_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -void calculate(Tensor output, Tensor input) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - TensorMeta output_meta(output); - TensorMeta input_meta(input); - infini::ops::ReluInfinilm::Call( - handle, - config, - input_meta.tensor(input), - output_meta.tensor(output)); -} - -} // namespace - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(Relu::dispatcher(), &calculate); - return true; -}(); - -} // namespace infinicore::op::relu_impl::infiniops -#endif diff --git a/src/infinicore/ops/relu6/relu6.cc b/src/infinicore/ops/relu6/relu6.cc deleted file mode 100644 index ebdb91b43..000000000 --- a/src/infinicore/ops/relu6/relu6.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/relu6.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Relu6); - -Relu6::Relu6(Tensor out, const Tensor &input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, input); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, input); -} - -void Relu6::execute(Tensor out, const Tensor &input) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Relu6, out, input); -} - -Tensor relu6(const Tensor &input) { - auto out = Tensor::empty(input->shape(), input->dtype(), input->device()); - relu6_(out, input); - return out; -} - -void relu6_(Tensor out, const Tensor &input) { - Relu6::execute(out, input); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/relu6/relu6_infiniop.cc b/src/infinicore/ops/relu6/relu6_infiniop.cc deleted file mode 100644 index ddb7379de..000000000 --- a/src/infinicore/ops/relu6/relu6_infiniop.cc +++ /dev/null @@ -1,53 +0,0 @@ -#include "infinicore/ops/relu6.hpp" - -#include "infiniop/ops/relu6.h" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::relu6_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Relu6, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, input; -}; - -void *plan(Tensor out, const Tensor &input) { - size_t seed = hash_combine(out, input); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Relu6, - seed, - out->desc(), - input->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Relu6, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(input)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopRelu6( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->input->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Relu6, &plan, &run, &cleanup); - -} // namespace infinicore::op::relu6_impl::infiniop diff --git a/src/infinicore/ops/rms_norm/rms_norm.cc b/src/infinicore/ops/rms_norm/rms_norm.cc deleted file mode 100644 index 0f8f2e57b..000000000 --- a/src/infinicore/ops/rms_norm/rms_norm.cc +++ /dev/null @@ -1,26 +0,0 @@ -#include "infinicore/ops/rms_norm.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(RMSNorm); - -RMSNorm::RMSNorm(Tensor y, const Tensor &x, const Tensor &weight, float epsilon) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x, weight); - INFINICORE_GRAPH_OP_DISPATCH(y->device().getType(), y, x, weight, epsilon); -} - -void RMSNorm::execute(Tensor y, const Tensor &x, const Tensor &weight, float epsilon) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(RMSNorm, y, x, weight, epsilon); -} - -Tensor rms_norm(const Tensor &x, const Tensor &weight, float epsilon) { - auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); - rms_norm_(y, x, weight, epsilon); - return y; -} - -void rms_norm_(Tensor y, const Tensor &x, const Tensor &weight, float epsilon) { - RMSNorm::execute(y, x, weight, epsilon); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/rms_norm/rms_norm_infiniop.cc b/src/infinicore/ops/rms_norm/rms_norm_infiniop.cc deleted file mode 100644 index 9e4622a28..000000000 --- a/src/infinicore/ops/rms_norm/rms_norm_infiniop.cc +++ /dev/null @@ -1,55 +0,0 @@ -#include "infinicore/ops/rms_norm.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::rms_norm_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, RMSNorm, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, y, x, weight; -}; - -void *plan(Tensor y, const Tensor &x, const Tensor &weight, float epsilon) { - size_t seed = hash_combine(y, x, weight, epsilon); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, RMSNorm, - seed, y->desc(), - x->desc(), - weight->desc(), - epsilon); - - INFINIOP_WORKSPACE_TENSOR(workspace, RMSNorm, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(y), - graph::GraphTensor(x), - graph::GraphTensor(weight)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR( - infiniopRMSNorm( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->y->data(), - planned->x->data(), - planned->weight->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(RMSNorm, &plan, &run, &cleanup); - -} // namespace infinicore::op::rms_norm_impl::infiniop diff --git a/src/infinicore/ops/rms_norm/rms_norm_infiniops.cc b/src/infinicore/ops/rms_norm/rms_norm_infiniops.cc deleted file mode 100644 index d01d989f7..000000000 --- a/src/infinicore/ops/rms_norm/rms_norm_infiniops.cc +++ /dev/null @@ -1,64 +0,0 @@ -#include "infinicore/ops/rms_norm.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/rms_norm.h" - -namespace infinicore::op::rms_norm_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -struct PlannedMeta { - TensorMeta y, x, weight; - graph::GraphTensor y_tensor, x_tensor, weight_tensor; - float epsilon; -}; - -} // namespace - -void *plan(Tensor y, const Tensor &x, const Tensor &weight, float epsilon) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(y->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x, weight); - - return new PlannedMeta{ - TensorMeta(y), - TensorMeta(x), - TensorMeta(weight), - graph::GraphTensor(y), - graph::GraphTensor(x), - graph::GraphTensor(weight), - epsilon}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - infini::ops::RmsNorm::Call( - handle, - config, - planned->x.tensor(planned->x_tensor), - planned->weight.tensor(planned->weight_tensor), - planned->epsilon, - planned->y.tensor(planned->y_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(RMSNorm::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(RMSNorm::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(RMSNorm::cleanup_dispatcher(), &cleanup); - return true; -}(); - -} // namespace infinicore::op::rms_norm_impl::infiniops -#endif diff --git a/src/infinicore/ops/rope/rope.cc b/src/infinicore/ops/rope/rope.cc deleted file mode 100644 index d28951b7d..000000000 --- a/src/infinicore/ops/rope/rope.cc +++ /dev/null @@ -1,46 +0,0 @@ -#include "infinicore/ops/rope.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(RoPE); - -RoPE::RoPE(Tensor x_out, - const Tensor &x, - const Tensor &pos, - const Tensor &sin_table, - const Tensor &cos_table, - infinicore::nn::RoPE::Algo algo) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x_out, x, pos, sin_table, cos_table); - INFINICORE_GRAPH_OP_DISPATCH(x_out->device().getType(), x_out, x, pos, sin_table, cos_table, algo); -} - -void RoPE::execute(Tensor x_out, - const Tensor &x, - const Tensor &pos, - const Tensor &sin_table, - const Tensor &cos_table, - infinicore::nn::RoPE::Algo algo) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(RoPE, x_out, x, pos, sin_table, cos_table, algo); -} - -void rope_(Tensor x_out, - const Tensor &x, - const Tensor &pos, - const Tensor &sin_table, - const Tensor &cos_table, - infinicore::nn::RoPE::Algo algo) { - RoPE::execute(x_out, x, pos, sin_table, cos_table, algo); -} - -Tensor rope(const Tensor &x, - const Tensor &pos, - const Tensor &sin_table, - const Tensor &cos_table, - infinicore::nn::RoPE::Algo algo) { - auto x_out = Tensor::empty(x->shape(), x->dtype(), x->device()); - rope_(x_out, x, pos, sin_table, cos_table, algo); - return x_out; -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/rope/rope_infiniop.cc b/src/infinicore/ops/rope/rope_infiniop.cc deleted file mode 100644 index 850c2d0a2..000000000 --- a/src/infinicore/ops/rope/rope_infiniop.cc +++ /dev/null @@ -1,81 +0,0 @@ -#include "infinicore/ops/rope.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::rope_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, RoPE, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace; - graph::GraphTensor x_out; - graph::GraphTensor x; - graph::GraphTensor pos; - graph::GraphTensor sin; - graph::GraphTensor cos; -}; - -static infiniopRoPEAlgo_t to_infiniop_algo(infinicore::nn::RoPE::Algo algo) { - switch (algo) { - case infinicore::nn::RoPE::Algo::GPT_J: - return INFINIOP_ROPE_ALGO_GPT_J; - case infinicore::nn::RoPE::Algo::GPT_NEOX: - return INFINIOP_ROPE_ALGO_GPT_NEOX; - default: - throw std::runtime_error("Unsupported RoPE algorithm"); - } -} - -void *plan(Tensor x_out, - const Tensor &x, - const Tensor &pos, - const Tensor &sin, - const Tensor &cos, - infinicore::nn::RoPE::Algo algo) { - auto infiniop_algo = to_infiniop_algo(algo); - size_t key = hash_combine(x_out, x, pos, sin, cos, static_cast(infiniop_algo)); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, RoPE, key, x_out->desc(), - x->desc(), - pos->desc(), - sin->desc(), - cos->desc(), - infiniop_algo); - - INFINIOP_WORKSPACE_TENSOR(workspace, RoPE, descriptor); - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x_out), - graph::GraphTensor(x), - graph::GraphTensor(pos), - graph::GraphTensor(sin), - graph::GraphTensor(cos)}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR( - infiniopRoPE( - p->descriptor->desc, - p->workspace->data(), - p->workspace->numel(), - p->x_out->data(), - p->x->data(), - p->pos->data(), - p->sin->data(), - p->cos->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(RoPE, &plan, &run, &cleanup); - -} // namespace infinicore::op::rope_impl::infiniop diff --git a/src/infinicore/ops/rope/rope_infiniops.cc b/src/infinicore/ops/rope/rope_infiniops.cc deleted file mode 100644 index 045a5aca3..000000000 --- a/src/infinicore/ops/rope/rope_infiniops.cc +++ /dev/null @@ -1,73 +0,0 @@ -#include "infinicore/ops/rope.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/rotary_embedding_infinilm.h" - -#include - -namespace infinicore::op::rope_impl::infiniops { -namespace { -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; -struct PlannedMeta { - TensorMeta x_out, x, pos, sin, cos; - graph::GraphTensor x_out_tensor, x_tensor, pos_tensor, sin_tensor, cos_tensor; - bool is_neox; -}; - -bool toInfiniOpsIsNeox(infinicore::nn::RoPE::Algo algo) { - switch (algo) { - case infinicore::nn::RoPE::Algo::GPT_J: - return true; - case infinicore::nn::RoPE::Algo::GPT_NEOX: - return false; - default: - throw std::runtime_error("Unsupported RoPE algorithm"); - } -} -} // namespace - -void *plan(Tensor x_out, - const Tensor &x, - const Tensor &pos, - const Tensor &sin, - const Tensor &cos, - infinicore::nn::RoPE::Algo algo) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(x_out->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x_out, x, pos, sin, cos); - return new PlannedMeta{ - TensorMeta(x_out), TensorMeta(x), TensorMeta(pos), TensorMeta(sin), TensorMeta(cos), - graph::GraphTensor(x_out), graph::GraphTensor(x), graph::GraphTensor(pos), graph::GraphTensor(sin), graph::GraphTensor(cos), - toInfiniOpsIsNeox(algo)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - infini::ops::RotaryEmbeddingInfinilm::Call( - handle, - config, - planned->x.tensor(planned->x_tensor), - planned->pos.tensor(planned->pos_tensor), - planned->sin.tensor(planned->sin_tensor), - planned->cos.tensor(planned->cos_tensor), - planned->is_neox, - planned->x_out.tensor(planned->x_out_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(RoPE::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(RoPE::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(RoPE::cleanup_dispatcher(), &cleanup); - return true; -}(); -} // namespace infinicore::op::rope_impl::infiniops -#endif diff --git a/src/infinicore/ops/rot/rot.cc b/src/infinicore/ops/rot/rot.cc deleted file mode 100644 index 262cce001..000000000 --- a/src/infinicore/ops/rot/rot.cc +++ /dev/null @@ -1,22 +0,0 @@ -#include "infinicore/ops/rot.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Rot); - -Rot::Rot(Tensor x, Tensor y, const Tensor &c, const Tensor &s) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, y, c, s); - INFINICORE_GRAPH_OP_DISPATCH(x->device().getType(), x, y, c, s); -} - -void Rot::execute(Tensor x, Tensor y, const Tensor &c, const Tensor &s) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Rot, x, y, c, s); -} - -void rot_(Tensor x, Tensor y, const Tensor &c, const Tensor &s) { - Rot::execute(x, y, c, s); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/rot/rot_infiniop.cc b/src/infinicore/ops/rot/rot_infiniop.cc deleted file mode 100644 index a114bf110..000000000 --- a/src/infinicore/ops/rot/rot_infiniop.cc +++ /dev/null @@ -1,54 +0,0 @@ -#include "infinicore/ops/rot.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::rot_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Rot, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, x, y, c, s; -}; - -void *plan(Tensor x, Tensor y, const Tensor &c, const Tensor &s) { - size_t seed = hash_combine(x, y, c, s); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Rot, - seed, - x->desc(), y->desc(), c->desc(), s->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Rot, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x), - graph::GraphTensor(y), - graph::GraphTensor(c), - graph::GraphTensor(s)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopRot( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->x->data(), - planned->y->data(), - planned->c->data(), - planned->s->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Rot, &plan, &run, &cleanup); - -} // namespace infinicore::op::rot_impl::infiniop diff --git a/src/infinicore/ops/rotg/rotg.cc b/src/infinicore/ops/rotg/rotg.cc deleted file mode 100644 index 233177b68..000000000 --- a/src/infinicore/ops/rotg/rotg.cc +++ /dev/null @@ -1,22 +0,0 @@ -#include "infinicore/ops/rotg.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Rotg::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Rotg::execute(Tensor x, Tensor y, Tensor c, Tensor s) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, y, c, s); - infinicore::context::setDevice(x->device()); - dispatcher().lookup(x->device().getType())(x, y, c, s); -} - -void rotg_(Tensor x, Tensor y, Tensor c, Tensor s) { - Rotg::execute(x, y, c, s); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/rotg/rotg_infiniop.cc b/src/infinicore/ops/rotg/rotg_infiniop.cc deleted file mode 100644 index d5e7a0e99..000000000 --- a/src/infinicore/ops/rotg/rotg_infiniop.cc +++ /dev/null @@ -1,56 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/rotg.hpp" -#include - -namespace infinicore::op::rotg_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopRotgDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyRotgDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor x, Tensor y, Tensor c, Tensor s) { - size_t seed = hash_combine(x, y, c, s); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopRotgDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateRotgDescriptor( - context::getInfiniopHandle(x->device()), &desc, - x->desc(), y->desc(), c->desc(), s->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetRotgWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopRotg( - desc, workspace->data(), workspace_size, - x->data(), y->data(), c->data(), s->data(), context::getStream())); -} - -static bool registered = []() { - Rotg::dispatcher().registerDevice({Device::Type::CPU, - Device::Type::CAMBRICON, - Device::Type::METAX}, - &calculate, - false); - return true; -}(); - -} // namespace infinicore::op::rotg_impl::infiniop diff --git a/src/infinicore/ops/rotm/rotm.cc b/src/infinicore/ops/rotm/rotm.cc deleted file mode 100644 index 2b1921f8b..000000000 --- a/src/infinicore/ops/rotm/rotm.cc +++ /dev/null @@ -1,22 +0,0 @@ -#include "infinicore/ops/rotm.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Rotm::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Rotm::execute(Tensor x, Tensor y, Tensor param) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, y, param); - infinicore::context::setDevice(x->device()); - dispatcher().lookup(x->device().getType())(x, y, param); -} - -void rotm_(Tensor x, Tensor y, Tensor param) { - Rotm::execute(x, y, param); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/rotm/rotm_infiniop.cc b/src/infinicore/ops/rotm/rotm_infiniop.cc deleted file mode 100644 index 9df5c9ab1..000000000 --- a/src/infinicore/ops/rotm/rotm_infiniop.cc +++ /dev/null @@ -1,56 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/rotm.hpp" -#include - -namespace infinicore::op::rotm_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopRotmDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyRotmDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor x, Tensor y, Tensor param) { - size_t seed = hash_combine(x, y, param); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopRotmDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateRotmDescriptor( - context::getInfiniopHandle(x->device()), &desc, - x->desc(), y->desc(), param->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetRotmWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopRotm( - desc, workspace->data(), workspace_size, - x->data(), y->data(), param->data(), context::getStream())); -} - -static bool registered = []() { - Rotm::dispatcher().registerDevice({Device::Type::CPU, - Device::Type::CAMBRICON, - Device::Type::METAX}, - &calculate, - false); - return true; -}(); - -} // namespace infinicore::op::rotm_impl::infiniop diff --git a/src/infinicore/ops/rotmg/rotmg.cc b/src/infinicore/ops/rotmg/rotmg.cc deleted file mode 100644 index 8a2860d63..000000000 --- a/src/infinicore/ops/rotmg/rotmg.cc +++ /dev/null @@ -1,22 +0,0 @@ -#include "infinicore/ops/rotmg.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Rotmg::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Rotmg::execute(Tensor d1, Tensor d2, Tensor x1, Tensor y1, Tensor param) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(d1, d2, x1, y1, param); - infinicore::context::setDevice(d1->device()); - dispatcher().lookup(d1->device().getType())(d1, d2, x1, y1, param); -} - -void rotmg_(Tensor d1, Tensor d2, Tensor x1, Tensor y1, Tensor param) { - Rotmg::execute(d1, d2, x1, y1, param); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/rotmg/rotmg_infiniop.cc b/src/infinicore/ops/rotmg/rotmg_infiniop.cc deleted file mode 100644 index ade0bd7aa..000000000 --- a/src/infinicore/ops/rotmg/rotmg_infiniop.cc +++ /dev/null @@ -1,56 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/rotmg.hpp" -#include - -namespace infinicore::op::rotmg_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopRotmgDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyRotmgDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor d1, Tensor d2, Tensor x1, Tensor y1, Tensor param) { - size_t seed = hash_combine(d1, d2, x1, y1, param); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopRotmgDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateRotmgDescriptor( - context::getInfiniopHandle(d1->device()), &desc, - d1->desc(), d2->desc(), x1->desc(), y1->desc(), param->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetRotmgWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopRotmg( - desc, workspace->data(), workspace_size, - d1->data(), d2->data(), x1->data(), y1->data(), param->data(), context::getStream())); -} - -static bool registered = []() { - Rotmg::dispatcher().registerDevice({Device::Type::CPU, - Device::Type::CAMBRICON, - Device::Type::METAX}, - &calculate, - false); - return true; -}(); - -} // namespace infinicore::op::rotmg_impl::infiniop diff --git a/src/infinicore/ops/rwkv5_wkv/rwkv5_wkv.cc b/src/infinicore/ops/rwkv5_wkv/rwkv5_wkv.cc deleted file mode 100644 index b33ff0131..000000000 --- a/src/infinicore/ops/rwkv5_wkv/rwkv5_wkv.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/rwkv5_wkv.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Rwkv5Wkv); - -Rwkv5Wkv::Rwkv5Wkv(Tensor out, - const Tensor &receptance, - const Tensor &key, - const Tensor &value, - const Tensor &time_decay, - const Tensor &time_faaaa, - Tensor state) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, receptance, key, value, time_decay, time_faaaa, state); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, receptance, key, value, time_decay, time_faaaa, state); -} - -void Rwkv5Wkv::execute(Tensor out, - const Tensor &receptance, - const Tensor &key, - const Tensor &value, - const Tensor &time_decay, - const Tensor &time_faaaa, - Tensor state) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Rwkv5Wkv, out, receptance, key, value, time_decay, time_faaaa, state); -} - -Tensor rwkv5_wkv(const Tensor &receptance, - const Tensor &key, - const Tensor &value, - const Tensor &time_decay, - const Tensor &time_faaaa, - Tensor state) { - auto output = Tensor::empty(receptance->shape(), receptance->dtype(), receptance->device()); - rwkv5_wkv_(output, receptance, key, value, time_decay, time_faaaa, state); - return output; -} - -void rwkv5_wkv_(Tensor out, - const Tensor &receptance, - const Tensor &key, - const Tensor &value, - const Tensor &time_decay, - const Tensor &time_faaaa, - Tensor state) { - Rwkv5Wkv::execute(out, receptance, key, value, time_decay, time_faaaa, state); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/rwkv5_wkv/rwkv5_wkv_infiniop.cc b/src/infinicore/ops/rwkv5_wkv/rwkv5_wkv_infiniop.cc deleted file mode 100644 index db071f971..000000000 --- a/src/infinicore/ops/rwkv5_wkv/rwkv5_wkv_infiniop.cc +++ /dev/null @@ -1,72 +0,0 @@ -#include "infinicore/ops/rwkv5_wkv.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::rwkv5_wkv_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Rwkv5Wkv, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, receptance, key, value, time_decay, time_faaaa, state; -}; - -void *plan(Tensor out, - const Tensor &receptance, - const Tensor &key, - const Tensor &value, - const Tensor &time_decay, - const Tensor &time_faaaa, - Tensor state) { - size_t seed = hash_combine(out, receptance, key, value, time_decay, time_faaaa, state); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Rwkv5Wkv, - seed, - out->desc(), - receptance->desc(), - key->desc(), - value->desc(), - time_decay->desc(), - time_faaaa->desc(), - state->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Rwkv5Wkv, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(receptance), - graph::GraphTensor(key), - graph::GraphTensor(value), - graph::GraphTensor(time_decay), - graph::GraphTensor(time_faaaa), - graph::GraphTensor(state)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopRwkv5Wkv( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->receptance->data(), - planned->key->data(), - planned->value->data(), - planned->time_decay->data(), - planned->time_faaaa->data(), - planned->state->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Rwkv5Wkv, &plan, &run, &cleanup); - -} // namespace infinicore::op::rwkv5_wkv_impl::infiniop diff --git a/src/infinicore/ops/scal/scal.cc b/src/infinicore/ops/scal/scal.cc deleted file mode 100644 index 21258af2a..000000000 --- a/src/infinicore/ops/scal/scal.cc +++ /dev/null @@ -1,22 +0,0 @@ -#include "infinicore/ops/scal.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Scal); - -Scal::Scal(const Tensor &alpha, Tensor x) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(alpha, x); - INFINICORE_GRAPH_OP_DISPATCH(x->device().getType(), alpha, x); -} - -void Scal::execute(const Tensor &alpha, Tensor x) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Scal, alpha, x); -} - -void scal_(const Tensor &alpha, Tensor x) { - Scal::execute(alpha, x); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/scal/scal_infiniop.cc b/src/infinicore/ops/scal/scal_infiniop.cc deleted file mode 100644 index 25521f4ee..000000000 --- a/src/infinicore/ops/scal/scal_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/scal.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::scal_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Scal, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, alpha, x; -}; - -void *plan(const Tensor &alpha, Tensor x) { - size_t seed = hash_combine(alpha, x); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Scal, - seed, - alpha->desc(), x->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Scal, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(alpha), - graph::GraphTensor(x)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopScal( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->alpha->data(), - planned->x->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Scal, &plan, &run, &cleanup); - -} // namespace infinicore::op::scal_impl::infiniop diff --git a/src/infinicore/ops/scaled_mm_i8/scaled_mm_i8.cc b/src/infinicore/ops/scaled_mm_i8/scaled_mm_i8.cc deleted file mode 100644 index 6d6a5f8ff..000000000 --- a/src/infinicore/ops/scaled_mm_i8/scaled_mm_i8.cc +++ /dev/null @@ -1,21 +0,0 @@ -#include "infinicore/ops/scaled_mm_i8.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(I8Gemm); - -I8Gemm::I8Gemm(Tensor c, const Tensor &a_p, const Tensor &a_s, const Tensor &b_p, const Tensor &b_s, std::optional bias) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a_p, a_s, b_p, b_s); - INFINICORE_GRAPH_OP_DISPATCH(c->device().getType(), c, a_p, a_s, b_p, b_s, bias); -} -void I8Gemm::execute(Tensor c, const Tensor &a_p, const Tensor &a_s, const Tensor &b_p, const Tensor &b_s, std::optional bias) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(I8Gemm, c, a_p, a_s, b_p, b_s, bias); -} - -void scaled_mm_i8_(Tensor c, const Tensor &a_p, const Tensor &a_s, const Tensor &b_p, const Tensor &b_s, std::optional bias) { - I8Gemm::execute(c, a_p, a_s, b_p, b_s, bias); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/scaled_mm_i8/scaled_mm_i8_infiniop.cc b/src/infinicore/ops/scaled_mm_i8/scaled_mm_i8_infiniop.cc deleted file mode 100644 index 952b570cc..000000000 --- a/src/infinicore/ops/scaled_mm_i8/scaled_mm_i8_infiniop.cc +++ /dev/null @@ -1,65 +0,0 @@ -#include "../../utils.hpp" -#include "../infiniop_impl.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/scaled_mm_i8.hpp" -#include - -namespace infinicore::op::scaled_mm_i8_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, I8Gemm, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, c, a_p, a_s, b_p, b_s; - std::optional bias; -}; - -void *plan(Tensor c, const Tensor &a_p, const Tensor &a_s, const Tensor &b_p, const Tensor &b_s, std::optional bias) { - size_t seed = hash_combine(c, a_p, a_s, b_p, b_s); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, I8Gemm, - seed, - c->desc(), bias.has_value() ? bias.value()->desc() : nullptr, - a_p->desc(), a_s->desc(), b_p->desc(), b_s->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, I8Gemm, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(c), - graph::GraphTensor(a_p), - graph::GraphTensor(a_s), - graph::GraphTensor(b_p), - graph::GraphTensor(b_s), - // bias.has_value() ? bias.value()->desc() : nullptr}; - bias ? std::optional(graph::GraphTensor(*bias)) : std::nullopt}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopI8Gemm( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->c->data(), - // planned->bias->data(), - planned->bias.has_value() ? planned->bias.value()->data() : nullptr, - planned->a_p->data(), - planned->a_s->data(), - planned->b_p->data(), - planned->b_s->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(I8Gemm, &plan, &run, &cleanup); - -} // namespace infinicore::op::scaled_mm_i8_impl::infiniop diff --git a/src/infinicore/ops/scaled_mm_w4a16_gptq_qy/scaled_mm_w4a16_gptq_qy.cc b/src/infinicore/ops/scaled_mm_w4a16_gptq_qy/scaled_mm_w4a16_gptq_qy.cc deleted file mode 100644 index 3afb715f0..000000000 --- a/src/infinicore/ops/scaled_mm_w4a16_gptq_qy/scaled_mm_w4a16_gptq_qy.cc +++ /dev/null @@ -1,23 +0,0 @@ -#include "infinicore/ops/scaled_mm_w4a16_gptq_qy.hpp" -#include "../../utils.hpp" -#include - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(GptqQyblasGemm); - -GptqQyblasGemm::GptqQyblasGemm(Tensor out, const Tensor &in, const Tensor &qweight, const Tensor &scales, const Tensor &qzeros, int64_t quant_type, int64_t bit) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, in, qweight, scales, qzeros); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, in, qweight, scales, qzeros, quant_type, bit); -} - -void GptqQyblasGemm::execute(Tensor out, const Tensor &in, const Tensor &qweight, const Tensor &scales, const Tensor &qzeros, int64_t quant_type, int64_t bit) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(GptqQyblasGemm, out, in, qweight, scales, qzeros, quant_type, bit); -} - -void scaled_mm_w4a16_gptq_qy_(Tensor out, const Tensor &in, const Tensor &qweight, const Tensor &scales, const Tensor &qzeros, int64_t quant_type, int64_t bit) { - - GptqQyblasGemm::execute(out, in, qweight, scales, qzeros, quant_type, bit); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/scaled_mm_w4a16_gptq_qy/scaled_mm_w4a16_gptq_qy_infiniop.cc b/src/infinicore/ops/scaled_mm_w4a16_gptq_qy/scaled_mm_w4a16_gptq_qy_infiniop.cc deleted file mode 100644 index 5ead44772..000000000 --- a/src/infinicore/ops/scaled_mm_w4a16_gptq_qy/scaled_mm_w4a16_gptq_qy_infiniop.cc +++ /dev/null @@ -1,59 +0,0 @@ -#include "../../utils.hpp" -#include "../infiniop_impl.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/scaled_mm_w4a16_gptq_qy.hpp" -#include - -namespace infinicore::op::scaled_mm_w4a16_gptq_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, GptqQyblasGemm, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, out, in, qweight, scales, qzeros; - int64_t quant_type, bit; -}; - -void *plan(Tensor out, const Tensor &in, const Tensor &qweight, const Tensor &scales, const Tensor &qzeros, int64_t quant_type, int64_t bit) { - size_t seed = hash_combine(out, in, qweight, scales, qzeros); - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, GptqQyblasGemm, - seed, - out->desc(), in->desc(), qweight->desc(), scales->desc(), qzeros->desc()); - INFINIOP_WORKSPACE_TENSOR(workspace, GptqQyblasGemm, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(out), - graph::GraphTensor(in), - graph::GraphTensor(qweight), - graph::GraphTensor(scales), - graph::GraphTensor(qzeros), - quant_type, bit}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - INFINICORE_CHECK_ERROR(infiniopGptqQyblasGemm( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->out->data(), - planned->in->data(), - planned->qweight->data(), - planned->scales->data(), - planned->qzeros->data(), - planned->quant_type, planned->bit, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(GptqQyblasGemm, &plan, &run, &cleanup); - -} // namespace infinicore::op::scaled_mm_w4a16_gptq_impl::infiniop diff --git a/src/infinicore/ops/scatter/scatter.cc b/src/infinicore/ops/scatter/scatter.cc deleted file mode 100644 index 402baf223..000000000 --- a/src/infinicore/ops/scatter/scatter.cc +++ /dev/null @@ -1,26 +0,0 @@ -#include "infinicore/ops/scatter.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Scatter::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Scatter::execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction) { - dispatcher().lookup(context::getDevice().getType())(output, input, dim, index, src, reduction); -} - -Tensor scatter(Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction) { - // 创建与 input 形状、数据类型、设备一致的 Output Tensor - auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); - scatter_(output, input, dim, index, src, reduction); - - return output; -} - -void scatter_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction) { - Scatter::execute(output, input, dim, index, src, reduction); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/scatter/scatter_infiniop.cc b/src/infinicore/ops/scatter/scatter_infiniop.cc deleted file mode 100644 index 3f080b6d7..000000000 --- a/src/infinicore/ops/scatter/scatter_infiniop.cc +++ /dev/null @@ -1,71 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/scatter.hpp" -#include - -namespace infinicore::op::scatter_impl::infiniop { - -// 定义描述符缓存 -thread_local common::OpCache caches( - 100, // capacity - [](infiniopScatterDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyScatterDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor src, int64_t reduction) { - // Scatter 算子输入 input, index, src 均为必须存在的 Tensor，直接参与 hash - size_t seed = hash_combine(output, input, dim, index, src, reduction); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopScatterDescriptor_t desc = nullptr; - - if (!desc_opt) { - // 3. 创建描述符 - // C++ Op 参数: output, input, dim, index, src, reduction - // C API 参数: output, input, indices, updates, axis, reduction - INFINICORE_CHECK_ERROR(infiniopCreateScatterDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc(), - index->desc(), // 对应 C API indices - src->desc(), // 对应 C API updates - static_cast(dim), - static_cast(reduction))); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 4. 获取 Workspace 并执行 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetScatterWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopScatter( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - index->data(), - src->data(), - context::getStream())); -} - -static bool registered = []() { - Scatter::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::scatter_impl::infiniop diff --git a/src/infinicore/ops/selu/selu.cc b/src/infinicore/ops/selu/selu.cc deleted file mode 100644 index 880b83207..000000000 --- a/src/infinicore/ops/selu/selu.cc +++ /dev/null @@ -1,40 +0,0 @@ -#include "infinicore/ops/selu.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &Selu::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -void Selu::execute(Tensor output, Tensor input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error( - "No Selu implementation found for device type: " - + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor selu(Tensor input) { - Shape shape = input->shape(); - auto output = Tensor::empty(shape, input->dtype(), input->device()); - selu_(output, input); - return output; -} - -void selu_(Tensor output, Tensor input) { - Selu::execute(output, input); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/selu/selu_infiniop.cc b/src/infinicore/ops/selu/selu_infiniop.cc deleted file mode 100644 index 9e517e868..000000000 --- a/src/infinicore/ops/selu/selu_infiniop.cc +++ /dev/null @@ -1,57 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/selu.hpp" - -#include - -namespace infinicore::op::selu_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopSeluDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroySeluDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopSeluDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateSeluDescriptor( - context::getInfiniopHandle(device), - &desc, - output->desc(), - input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetSeluWorkspaceSize(desc, &workspace_size)); - auto workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopSelu( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -static bool registered = []() { - Selu::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::selu_impl::infiniop diff --git a/src/infinicore/ops/sigmoid/sigmoid.cc b/src/infinicore/ops/sigmoid/sigmoid.cc deleted file mode 100644 index d2dd48d72..000000000 --- a/src/infinicore/ops/sigmoid/sigmoid.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/sigmoid.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Sigmoid); - -Sigmoid::Sigmoid(Tensor output, const Tensor &input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - INFINICORE_GRAPH_OP_DISPATCH(output->device().getType(), output, input); -} - -void Sigmoid::execute(Tensor output, const Tensor &input) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Sigmoid, output, input); -} - -Tensor sigmoid(const Tensor &input) { - auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); - sigmoid_(output, input); - return output; -} - -void sigmoid_(Tensor output, const Tensor &input) { - Sigmoid::execute(output, input); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/sigmoid/sigmoid_infiniop.cc b/src/infinicore/ops/sigmoid/sigmoid_infiniop.cc deleted file mode 100644 index 7e3086d98..000000000 --- a/src/infinicore/ops/sigmoid/sigmoid_infiniop.cc +++ /dev/null @@ -1,51 +0,0 @@ -#include "infinicore/ops/sigmoid.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::sigmoid_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Sigmoid, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, output, input; -}; - -void *plan(Tensor output, const Tensor &input) { - size_t seed = hash_combine(output, input); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Sigmoid, - seed, - output->desc(), - input->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Sigmoid, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(output), - graph::GraphTensor(input)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopSigmoid( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->output->data(), - planned->input->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Sigmoid, &plan, &run, &cleanup); - -} // namespace infinicore::op::sigmoid_impl::infiniop diff --git a/src/infinicore/ops/sigmoid/sigmoid_infiniops.cc b/src/infinicore/ops/sigmoid/sigmoid_infiniops.cc deleted file mode 100644 index 0327e6c7c..000000000 --- a/src/infinicore/ops/sigmoid/sigmoid_infiniops.cc +++ /dev/null @@ -1,58 +0,0 @@ -#include "infinicore/ops/sigmoid.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/sigmoid_infinilm.h" - -namespace infinicore::op::sigmoid_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -struct PlannedMeta { - TensorMeta output, input; - graph::GraphTensor output_tensor, input_tensor; -}; - -} // namespace - -void *plan(Tensor output, const Tensor &input) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - - return new PlannedMeta{ - TensorMeta(output), - TensorMeta(input), - graph::GraphTensor(output), - graph::GraphTensor(input)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - infini::ops::SigmoidInfinilm::Call( - handle, - config, - planned->input.tensor(planned->input_tensor), - planned->output.tensor(planned->output_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(Sigmoid::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(Sigmoid::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(Sigmoid::cleanup_dispatcher(), &cleanup); - return true; -}(); - -} // namespace infinicore::op::sigmoid_impl::infiniops -#endif diff --git a/src/infinicore/ops/silu/silu.cc b/src/infinicore/ops/silu/silu.cc deleted file mode 100644 index afb976a89..000000000 --- a/src/infinicore/ops/silu/silu.cc +++ /dev/null @@ -1,37 +0,0 @@ -#include "infinicore/ops/silu.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &Silu::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Silu::execute(Tensor output, Tensor input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Silu implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor silu(Tensor input) { - Shape shape = input->shape(); - auto output = Tensor::empty(shape, input->dtype(), input->device()); - silu_(output, input); - return output; -} - -void silu_(Tensor output, Tensor input) { - Silu::execute(output, input); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/silu/silu_infiniop.cc b/src/infinicore/ops/silu/silu_infiniop.cc deleted file mode 100644 index 5edad2038..000000000 --- a/src/infinicore/ops/silu/silu_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/silu.hpp" -#include - -namespace infinicore::op::silu_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopSiluDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroySiluDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopSiluDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateSiluDescriptor( - context::getInfiniopHandle(device), &desc, - output->desc(), input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetSiluWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopSilu( - desc, workspace->data(), workspace_size, - output->data(), input->data(), context::getStream())); -} - -static bool registered = []() { - Silu::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::silu_impl::infiniop diff --git a/src/infinicore/ops/silu/silu_infiniops.cc b/src/infinicore/ops/silu/silu_infiniops.cc deleted file mode 100644 index b1009079b..000000000 --- a/src/infinicore/ops/silu/silu_infiniops.cc +++ /dev/null @@ -1,38 +0,0 @@ -#include "infinicore/ops/silu.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/silu.h" - -namespace infinicore::op::silu_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -void calculate(Tensor output, Tensor input) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - TensorMeta output_meta(output); - TensorMeta input_meta(input); - infini::ops::Silu::Call( - handle, - config, - input_meta.tensor(input), - output_meta.tensor(output)); -} - -} // namespace - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(Silu::dispatcher(), &calculate); - return true; -}(); - -} // namespace infinicore::op::silu_impl::infiniops -#endif diff --git a/src/infinicore/ops/silu_and_mul/silu_and_mul.cc b/src/infinicore/ops/silu_and_mul/silu_and_mul.cc deleted file mode 100644 index 33fe4ff8a..000000000 --- a/src/infinicore/ops/silu_and_mul/silu_and_mul.cc +++ /dev/null @@ -1,35 +0,0 @@ -#include "infinicore/ops/silu_and_mul.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(SiluAndMul); - -SiluAndMul::SiluAndMul(Tensor out, const Tensor &x) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, x); - INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, x); -} - -void SiluAndMul::execute(Tensor out, const Tensor &x) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(SiluAndMul, out, x); -} - -Tensor silu_and_mul(const Tensor &x) { - Shape shape = x->shape(); - size_t ndim = x->ndim(); - - if (shape[ndim - 1] % 2 != 0) { - throw std::runtime_error("SiluAndMul input last dim must be even."); - } - shape[ndim - 1] /= 2; - - auto out = Tensor::empty(shape, x->dtype(), x->device()); - silu_and_mul_(out, x); - return out; -} - -void silu_and_mul_(Tensor out, const Tensor &x) { - SiluAndMul::execute(out, x); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/silu_and_mul/silu_and_mul_infiniop.cc b/src/infinicore/ops/silu_and_mul/silu_and_mul_infiniop.cc deleted file mode 100644 index d5d243d13..000000000 --- a/src/infinicore/ops/silu_and_mul/silu_and_mul_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "../infiniop_impl.hpp" -#include "infinicore/ops/silu_and_mul.hpp" - -namespace infinicore::op::silu_and_mul_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, SiluAndMul, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, output, input; -}; - -void *plan(Tensor output, const Tensor &input) { - size_t seed = hash_combine(output, input); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, SiluAndMul, - seed, output->desc(), input->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, SiluAndMul, descriptor); - - auto planned = new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(output), - graph::GraphTensor(input)}; - - return planned; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopSiluAndMul( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->output->data(), - planned->input->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(SiluAndMul, &plan, &run, &cleanup); - -} // namespace infinicore::op::silu_and_mul_impl::infiniop diff --git a/src/infinicore/ops/silu_and_mul/silu_and_mul_infiniops.cc b/src/infinicore/ops/silu_and_mul/silu_and_mul_infiniops.cc deleted file mode 100644 index aa1bf3413..000000000 --- a/src/infinicore/ops/silu_and_mul/silu_and_mul_infiniops.cc +++ /dev/null @@ -1,44 +0,0 @@ -#include "infinicore/ops/silu_and_mul.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/silu_and_mul_infinilm.h" - -namespace infinicore::op::silu_and_mul_impl::infiniops { -namespace { -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; -struct PlannedMeta { - TensorMeta output, input; - graph::GraphTensor output_tensor, input_tensor; -}; -} // namespace - -void *plan(Tensor output, const Tensor &input) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - return new PlannedMeta{TensorMeta(output), TensorMeta(input), graph::GraphTensor(output), graph::GraphTensor(input)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - infini::ops::SiluAndMulInfinilm::Call( - handle, config, planned->input.tensor(planned->input_tensor), planned->output.tensor(planned->output_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(SiluAndMul::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(SiluAndMul::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(SiluAndMul::cleanup_dispatcher(), &cleanup); - return true; -}(); -} // namespace infinicore::op::silu_and_mul_impl::infiniops -#endif diff --git a/src/infinicore/ops/sinh/sinh.cc b/src/infinicore/ops/sinh/sinh.cc deleted file mode 100644 index 7a4ce770c..000000000 --- a/src/infinicore/ops/sinh/sinh.cc +++ /dev/null @@ -1,40 +0,0 @@ -#include "infinicore/ops/sinh.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &Sinh::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -} - -void Sinh::execute(Tensor output, Tensor input) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error( - "No Sinh implementation found for device type: " - + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor sinh(Tensor input) { - Shape shape = input->shape(); - auto output = Tensor::empty(shape, input->dtype(), input->device()); - sinh_(output, input); - return output; -} - -void sinh_(Tensor output, Tensor input) { - Sinh::execute(output, input); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/sinh/sinh_infiniop.cc b/src/infinicore/ops/sinh/sinh_infiniop.cc deleted file mode 100644 index 2f0e7d860..000000000 --- a/src/infinicore/ops/sinh/sinh_infiniop.cc +++ /dev/null @@ -1,57 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/sinh.hpp" - -#include - -namespace infinicore::op::sinh_impl::infiniop { - -thread_local common::OpCache caches( - 100, - [](infiniopSinhDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroySinhDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopSinhDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateSinhDescriptor( - context::getInfiniopHandle(device), - &desc, - output->desc(), - input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetSinhWorkspaceSize(desc, &workspace_size)); - auto workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopSinh( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -static bool registered = []() { - Sinh::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::sinh_impl::infiniop diff --git a/src/infinicore/ops/smooth_l1_loss/smooth_l1_loss.cc b/src/infinicore/ops/smooth_l1_loss/smooth_l1_loss.cc deleted file mode 100644 index 7649cc825..000000000 --- a/src/infinicore/ops/smooth_l1_loss/smooth_l1_loss.cc +++ /dev/null @@ -1,34 +0,0 @@ -#include "infinicore/ops/smooth_l1_loss.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &SmoothL1Loss::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void SmoothL1Loss::execute(Tensor output, Tensor input, Tensor target, float beta, int64_t reduction) { - dispatcher().lookup(context::getDevice().getType())(output, input, target, beta, reduction); -} - -Tensor smooth_l1_loss(Tensor input, Tensor target, float beta, int64_t reduction) { - Shape output_shape; - if (reduction == 0) { - // Reduction::None -> 输出形状与输入一致 - output_shape = input->shape(); - } else { - output_shape = {}; - } - - auto output = Tensor::empty(output_shape, input->dtype(), input->device()); - - smooth_l1_loss_(output, input, target, beta, reduction); - return output; -} - -void smooth_l1_loss_(Tensor output, Tensor input, Tensor target, float beta, int64_t reduction) { - SmoothL1Loss::execute(output, input, target, beta, reduction); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/smooth_l1_loss/smooth_l1_loss_infiniop.cc b/src/infinicore/ops/smooth_l1_loss/smooth_l1_loss_infiniop.cc deleted file mode 100644 index 02e8a9444..000000000 --- a/src/infinicore/ops/smooth_l1_loss/smooth_l1_loss_infiniop.cc +++ /dev/null @@ -1,63 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/smooth_l1_loss.hpp" -#include - -namespace infinicore::op::smooth_l1_loss_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopSmoothL1LossDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroySmoothL1LossDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, Tensor target, float beta, int64_t reduction) { - size_t seed = hash_combine(output, input, target, beta, reduction); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopSmoothL1LossDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateSmoothL1LossDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc(), - target->desc(), - beta, - static_cast(reduction))); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetSmoothL1LossWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopSmoothL1Loss( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - target->data(), - context::getStream())); -} - -static bool registered = []() { - SmoothL1Loss::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::smooth_l1_loss_impl::infiniop diff --git a/src/infinicore/ops/softmax/softmax.cc b/src/infinicore/ops/softmax/softmax.cc deleted file mode 100644 index 1856f61de..000000000 --- a/src/infinicore/ops/softmax/softmax.cc +++ /dev/null @@ -1,37 +0,0 @@ -#include "infinicore/ops/softmax.hpp" - -#include "../../utils.hpp" - -#include - -namespace infinicore::op { - -common::OpDispatcher &Softmax::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Softmax::execute(Tensor output, Tensor input, int axis) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(output->device()); - auto device_type = output->device().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Softmax implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input, axis); -} - -Tensor softmax(Tensor input, int axis) { - Shape shape = input->shape(); - auto output = Tensor::empty(shape, input->dtype(), input->device()); - softmax_(output, input, axis); - return output; -} - -void softmax_(Tensor output, Tensor input, int axis) { - Softmax::execute(output, input, axis); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/softmax/softmax_infiniop.cc b/src/infinicore/ops/softmax/softmax_infiniop.cc deleted file mode 100644 index 2c1dfd6e5..000000000 --- a/src/infinicore/ops/softmax/softmax_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/softmax.hpp" -#include - -namespace infinicore::op::softmax_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopSoftmaxDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroySoftmaxDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, int axis) { - size_t seed = hash_combine(output, input, axis); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopSoftmaxDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateSoftmaxDescriptor( - context::getInfiniopHandle(device), &desc, - output->desc(), input->desc(), axis)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetSoftmaxWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopSoftmax( - desc, workspace->data(), workspace_size, - output->data(), input->data(), context::getStream())); -} - -static bool registered = []() { - Softmax::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::softmax_impl::infiniop diff --git a/src/infinicore/ops/softmax/softmax_infiniops.cc b/src/infinicore/ops/softmax/softmax_infiniops.cc deleted file mode 100644 index c0087e1f8..000000000 --- a/src/infinicore/ops/softmax/softmax_infiniops.cc +++ /dev/null @@ -1,42 +0,0 @@ -#include "infinicore/ops/softmax.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/softmax_infinilm.h" - -#include - -namespace infinicore::op::softmax_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -void calculate(Tensor output, Tensor input, int axis) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(output->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - TensorMeta output_meta(output); - TensorMeta input_meta(input); - infini::ops::SoftmaxInfinilm::Call( - handle, - config, - input_meta.tensor(input), - static_cast(axis), - std::optional{}, - output_meta.tensor(output)); -} - -} // namespace - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(Softmax::dispatcher(), &calculate); - return true; -}(); - -} // namespace infinicore::op::softmax_impl::infiniops -#endif diff --git a/src/infinicore/ops/softplus/softplus.cc b/src/infinicore/ops/softplus/softplus.cc deleted file mode 100644 index 750d887d2..000000000 --- a/src/infinicore/ops/softplus/softplus.cc +++ /dev/null @@ -1,34 +0,0 @@ -#include "infinicore/ops/softplus.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Softplus::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -// 修改：增加 beta 和 threshold 参数 -void Softplus::execute(Tensor y, Tensor x, float beta, float threshold) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - infinicore::context::setDevice(y->device()); - - // 修改：将 beta 和 threshold 传递给底层的实现 (infiniop wrapper) - dispatcher().lookup(y->device().getType())(y, x, beta, threshold); -} - -// 修改：增加 beta 和 threshold 参数 -Tensor softplus(Tensor x, float beta, float threshold) { - auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); - // 传递参数 - softplus_(y, x, beta, threshold); - return y; -} - -// 修改：增加 beta 和 threshold 参数 -void softplus_(Tensor y, Tensor x, float beta, float threshold) { - // 传递参数 - Softplus::execute(y, x, beta, threshold); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/softplus/softplus_infiniop.cc b/src/infinicore/ops/softplus/softplus_infiniop.cc deleted file mode 100644 index f7ccc7e20..000000000 --- a/src/infinicore/ops/softplus/softplus_infiniop.cc +++ /dev/null @@ -1,49 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/softplus.hpp" -#include - -namespace infinicore::op::softplus_impl::infiniop { -thread_local common::OpCache caches( - 100, // capacity - [](infiniopSoftplusDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroySoftplusDescriptor(desc)); - desc = nullptr; - } - }); -void calculate(Tensor y, Tensor x, float beta, float threshold) { - size_t seed = hash_combine(y, x, beta, threshold); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopSoftplusDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateSoftplusDescriptor( - context::getInfiniopHandle(device), &desc, - y->desc(), x->desc(), beta, threshold)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 获取并分配 Workspace - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetSoftplusWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - // 执行计算 - INFINICORE_CHECK_ERROR(infiniopSoftplus( - desc, workspace->data(), workspace_size, - y->data(), x->data(), context::getStream())); -} -static bool registered = []() { - Softplus::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::softplus_impl::infiniop diff --git a/src/infinicore/ops/softsign/softsign.cc b/src/infinicore/ops/softsign/softsign.cc deleted file mode 100644 index 662d5bd6a..000000000 --- a/src/infinicore/ops/softsign/softsign.cc +++ /dev/null @@ -1,27 +0,0 @@ -#include "infinicore/ops/softsign.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { - -common::OpDispatcher &Softsign::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Softsign::execute(Tensor y, Tensor x) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(y, x); - infinicore::context::setDevice(y->device()); - dispatcher().lookup(y->device().getType())(y, x); -} - -Tensor softsign(Tensor x) { - auto y = Tensor::empty(x->shape(), x->dtype(), x->device()); - softsign_(y, x); - return y; -} - -void softsign_(Tensor y, Tensor x) { - Softsign::execute(y, x); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/softsign/softsign_infiniop.cc b/src/infinicore/ops/softsign/softsign_infiniop.cc deleted file mode 100644 index 0fa0642de..000000000 --- a/src/infinicore/ops/softsign/softsign_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/softsign.hpp" -#include - -namespace infinicore::op::softsign_impl::infiniop { -thread_local common::OpCache caches( - 100, // capacity - [](infiniopSoftsignDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroySoftsignDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor y, Tensor x) { - size_t seed = hash_combine(y, x); - - auto device = context::getDevice(); - auto &cache = caches.getCache(device); - - auto desc_opt = cache.get(seed); - infiniopSoftsignDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateSoftsignDescriptor( - context::getInfiniopHandle(device), &desc, - y->desc(), x->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 获取工作空间大小 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetSoftsignWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - // 执行计算 - INFINICORE_CHECK_ERROR(infiniopSoftsign( - desc, workspace->data(), workspace_size, - y->data(), x->data(), context::getStream())); -} -static bool registered = []() { - Softsign::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::softsign_impl::infiniop diff --git a/src/infinicore/ops/sum/sum.cc b/src/infinicore/ops/sum/sum.cc deleted file mode 100644 index 5fcecda5e..000000000 --- a/src/infinicore/ops/sum/sum.cc +++ /dev/null @@ -1,67 +0,0 @@ -#include "infinicore/ops/sum.hpp" - -#include "../../utils.hpp" -#include -#include - -namespace infinicore::op { - -common::OpDispatcher &Sum::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; -void Sum::execute(Tensor output, Tensor input, std::vector dim, bool keepdim) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(output, input); - infinicore::context::setDevice(input->device()); - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Sum implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input, dim, keepdim); -} - -Tensor sum(Tensor input, std::vector dim, bool keepdim) { - auto in_shape = input->shape(); - std::vector out_shape; - if (dim.empty()) { - for (size_t i = 0; i < in_shape.size(); i++) { - dim.push_back(i); - } - } - std::sort(dim.begin(), dim.end()); - if (dim.size() == in_shape.size() && !keepdim) { - out_shape = {}; - } else { - if (keepdim) { - size_t j = 0; - for (size_t i = 0; i < in_shape.size(); i++) { - if (j < dim.size() && dim[j] == i) { - out_shape.push_back(1); - j++; - } else { - out_shape.push_back(in_shape[i]); - } - } - } else { - size_t j = 0; - for (size_t i = 0; i < in_shape.size(); i++) { - if (j < dim.size() && dim[j] == i) { - j++; - } else { - out_shape.push_back(in_shape[i]); - } - } - } - } - auto output = Tensor::empty(out_shape, input->dtype(), input->device()); - sum_(output, input, dim, keepdim); - return output; -} - -void sum_(Tensor output, Tensor input, std::vector dim, bool keepdim) { - Sum::execute(output, input, dim, keepdim); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/sum/sum_infiniop.cc b/src/infinicore/ops/sum/sum_infiniop.cc deleted file mode 100644 index 700ab04fa..000000000 --- a/src/infinicore/ops/sum/sum_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/sum.hpp" -#include - -namespace infinicore::op::sum_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopSumDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroySumDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, std::vector dim, bool keepdim) { - size_t seed = hash_combine(output, input, dim.size(), keepdim); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopSumDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateSumDescriptor( - context::getInfiniopHandle(output->device()), &desc, - output->desc(), input->desc(), dim.data(), dim.size(), keepdim)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetSumWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopSum( - desc, workspace->data(), workspace_size, - output->data(), input->data(), dim.data(), dim.size(), keepdim, context::getStream())); -} - -static bool registered = []() { - Sum::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::sum_impl::infiniop diff --git a/src/infinicore/ops/swap/swap.cc b/src/infinicore/ops/swap/swap.cc deleted file mode 100644 index 6ce4b86b9..000000000 --- a/src/infinicore/ops/swap/swap.cc +++ /dev/null @@ -1,22 +0,0 @@ -#include "infinicore/ops/swap.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Swap); - -Swap::Swap(Tensor x, Tensor y) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, y); - INFINICORE_GRAPH_OP_DISPATCH(x->device().getType(), x, y); -} - -void Swap::execute(Tensor x, Tensor y) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Swap, x, y); -} - -void swap_(Tensor x, Tensor y) { - Swap::execute(x, y); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/swap/swap_infiniop.cc b/src/infinicore/ops/swap/swap_infiniop.cc deleted file mode 100644 index 32a163f65..000000000 --- a/src/infinicore/ops/swap/swap_infiniop.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "infinicore/ops/swap.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::swap_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Swap, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, x, y; -}; - -void *plan(Tensor x, Tensor y) { - size_t seed = hash_combine(x, y); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Swap, - seed, - x->desc(), y->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Swap, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(x), - graph::GraphTensor(y)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopSwap( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->x->data(), - planned->y->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Swap, &plan, &run, &cleanup); - -} // namespace infinicore::op::swap_impl::infiniop diff --git a/src/infinicore/ops/swiglu/swiglu.cc b/src/infinicore/ops/swiglu/swiglu.cc deleted file mode 100644 index 8ee0682ad..000000000 --- a/src/infinicore/ops/swiglu/swiglu.cc +++ /dev/null @@ -1,26 +0,0 @@ -#include "infinicore/ops/swiglu.hpp" -#include "../../utils.hpp" - -namespace infinicore::op { -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(SwiGLU); - -SwiGLU::SwiGLU(Tensor c, const Tensor &a, const Tensor &b) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - INFINICORE_GRAPH_OP_DISPATCH(c->device().getType(), c, a, b); -} - -void SwiGLU::execute(Tensor c, const Tensor &a, const Tensor &b) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(SwiGLU, c, a, b); -} - -Tensor swiglu(const Tensor &a, const Tensor &b) { - auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); - swiglu_(c, a, b); - return c; -} - -void swiglu_(Tensor c, const Tensor &a, const Tensor &b) { - SwiGLU::execute(c, a, b); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/swiglu/swiglu_infiniop.cc b/src/infinicore/ops/swiglu/swiglu_infiniop.cc deleted file mode 100644 index fbb76b570..000000000 --- a/src/infinicore/ops/swiglu/swiglu_infiniop.cc +++ /dev/null @@ -1,55 +0,0 @@ -#include "infinicore/ops/swiglu.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::swiglu_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, SwiGLU, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace; - graph::GraphTensor c; - graph::GraphTensor a; - graph::GraphTensor b; -}; - -void *plan(Tensor c, const Tensor &a, const Tensor &b) { - size_t key = hash_combine(c, a, b); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, SwiGLU, - key, c->desc(), a->desc(), b->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, SwiGLU, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(c), - graph::GraphTensor(a), - graph::GraphTensor(b)}; -} - -void run(void *planned_meta) { - auto *p = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR( - infiniopSwiGLU( - p->descriptor->desc, - p->workspace->data(), - p->workspace->numel(), - p->c->data(), - p->a->data(), - p->b->data(), - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(SwiGLU, &plan, &run, &cleanup); - -} // namespace infinicore::op::swiglu_impl::infiniop diff --git a/src/infinicore/ops/swiglu/swiglu_infiniops.cc b/src/infinicore/ops/swiglu/swiglu_infiniops.cc deleted file mode 100644 index 5b99a8eef..000000000 --- a/src/infinicore/ops/swiglu/swiglu_infiniops.cc +++ /dev/null @@ -1,61 +0,0 @@ -#include "infinicore/ops/swiglu.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/swiglu.h" - -namespace infinicore::op::swiglu_impl::infiniops { -namespace { - -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; - -struct PlannedMeta { - TensorMeta c, a, b; - graph::GraphTensor c_tensor, a_tensor, b_tensor; -}; - -} // namespace - -void *plan(Tensor c, const Tensor &a, const Tensor &b) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(c->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); - - return new PlannedMeta{ - TensorMeta(c), - TensorMeta(a), - TensorMeta(b), - graph::GraphTensor(c), - graph::GraphTensor(a), - graph::GraphTensor(b)}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - - infini::ops::Swiglu::Call( - handle, - config, - planned->a.tensor(planned->a_tensor), - planned->b.tensor(planned->b_tensor), - planned->c.tensor(planned->c_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(SwiGLU::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(SwiGLU::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(SwiGLU::cleanup_dispatcher(), &cleanup); - return true; -}(); - -} // namespace infinicore::op::swiglu_impl::infiniops -#endif diff --git a/src/infinicore/ops/take/take.cc b/src/infinicore/ops/take/take.cc deleted file mode 100644 index fb3abc8d6..000000000 --- a/src/infinicore/ops/take/take.cc +++ /dev/null @@ -1,28 +0,0 @@ -#include "infinicore/ops/take.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &Take::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -// 2. Execute 实现：查找对应设备的核函数并执行 -void Take::execute(Tensor output, Tensor input, Tensor indices) { - dispatcher().lookup(context::getDevice().getType())(output, input, indices); -} - -Tensor take(Tensor input, Tensor indices) { - // 【关键区别】Take 的输出形状取决于 indices 的形状，但数据类型取决于 input - auto output = Tensor::empty(indices->shape(), input->dtype(), input->device()); - - take_(output, input, indices); - return output; -} - -void take_(Tensor output, Tensor input, Tensor indices) { - Take::execute(output, input, indices); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/take/take_infiniop.cc b/src/infinicore/ops/take/take_infiniop.cc deleted file mode 100644 index c39653886..000000000 --- a/src/infinicore/ops/take/take_infiniop.cc +++ /dev/null @@ -1,58 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/take.hpp" -#include - -namespace infinicore::op::take_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopTakeDescriptor_t &desc) { - if (desc != nullptr) { - // 销毁 Take 描述符 - INFINICORE_CHECK_ERROR(infiniopDestroyTakeDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, Tensor indices) { - size_t seed = hash_combine(output, input, indices); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopTakeDescriptor_t desc = nullptr; - - if (!desc_opt) { - // 2. 创建描述符：传入 input 和 indices 的 descriptor - INFINICORE_CHECK_ERROR(infiniopCreateTakeDescriptor( - context::getInfiniopHandle(output->device()), &desc, - output->desc(), input->desc(), indices->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetTakeWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - // 3. 执行计算：传入 input 和 indices 的数据指针 - INFINICORE_CHECK_ERROR(infiniopTake( - desc, - workspace->data(), workspace_size, - output->data(), input->data(), indices->data(), - context::getStream())); -} - -static bool registered = []() { - // 注册到 Take 的 dispatcher - Take::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::take_impl::infiniop diff --git a/src/infinicore/ops/tan/tan.cc b/src/infinicore/ops/tan/tan.cc deleted file mode 100644 index 20bcab4b5..000000000 --- a/src/infinicore/ops/tan/tan.cc +++ /dev/null @@ -1,32 +0,0 @@ -#include "infinicore/ops/tan.hpp" -#include - -namespace infinicore::op { - -common::OpDispatcher &Tan::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Tan::execute(Tensor output, Tensor input) { - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Tan implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor tan(Tensor input) { - Shape shape = input->shape(); - auto output = Tensor::empty(shape, input->dtype(), input->device()); - tan_(output, input); - return output; -} - -void tan_(Tensor output, Tensor input) { - Tan::execute(output, input); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/tan/tan_infiniop.cc b/src/infinicore/ops/tan/tan_infiniop.cc deleted file mode 100644 index cf3fed96c..000000000 --- a/src/infinicore/ops/tan/tan_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/tan.hpp" -#include - -namespace infinicore::op::tan_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopTanDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyTanDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopTanDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateTanDescriptor( - context::getInfiniopHandle(output->device()), &desc, - output->desc(), input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetTanWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopTan( - desc, workspace->data(), workspace_size, - output->data(), input->data(), context::getStream())); -} - -static bool registered = []() { - Tan::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::tan_impl::infiniop diff --git a/src/infinicore/ops/tanhshrink/tanhshrink.cc b/src/infinicore/ops/tanhshrink/tanhshrink.cc deleted file mode 100644 index ea96e46f4..000000000 --- a/src/infinicore/ops/tanhshrink/tanhshrink.cc +++ /dev/null @@ -1,32 +0,0 @@ -#include "infinicore/ops/tanhshrink.hpp" -#include - -namespace infinicore::op { - -common::OpDispatcher &Tanhshrink::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Tanhshrink::execute(Tensor output, Tensor input) { - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Tanhshrink implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(output, input); -} - -Tensor tanhshrink(Tensor input) { - Shape shape = input->shape(); - auto output = Tensor::empty(shape, input->dtype(), input->device()); - tanhshrink_(output, input); - return output; -} - -void tanhshrink_(Tensor output, Tensor input) { - Tanhshrink::execute(output, input); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/tanhshrink/tanhshrink_infiniop.cc b/src/infinicore/ops/tanhshrink/tanhshrink_infiniop.cc deleted file mode 100644 index 1a8885386..000000000 --- a/src/infinicore/ops/tanhshrink/tanhshrink_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/tanhshrink.hpp" -#include - -namespace infinicore::op::tanhshrink_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopTanhshrinkDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyTanhshrinkDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopTanhshrinkDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateTanhshrinkDescriptor( - context::getInfiniopHandle(output->device()), &desc, - output->desc(), input->desc())); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetTanhshrinkWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopTanhshrink( - desc, workspace->data(), workspace_size, - output->data(), input->data(), context::getStream())); -} - -static bool registered = []() { - Tanhshrink::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::tanhshrink_impl::infiniop diff --git a/src/infinicore/ops/topk/topk.cc b/src/infinicore/ops/topk/topk.cc deleted file mode 100644 index a5b52fccf..000000000 --- a/src/infinicore/ops/topk/topk.cc +++ /dev/null @@ -1,40 +0,0 @@ -#include "infinicore/ops/topk.hpp" - -#include "../../utils.hpp" -#include -#include - -namespace infinicore::op { - -common::OpDispatcher &TopK::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; -void TopK::execute(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest, bool sorted) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(values_output, input); - infinicore::context::setDevice(input->device()); - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Topk implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(values_output, indices_output, input, k, dim, largest, sorted); -} - -std::pair topk(Tensor input, size_t k, size_t dim, bool largest, bool sorted) { - auto in_shape = input->shape(); - std::vector out_shape = in_shape; - out_shape[dim] = k; - - auto values_output = Tensor::empty(out_shape, input->dtype(), input->device()); - auto indices_output = Tensor::empty(out_shape, DataType::I32, input->device()); - topk_(values_output, indices_output, input, k, dim, largest, sorted); - return {values_output, indices_output}; -} - -void topk_(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest, bool sorted) { - TopK::execute(values_output, indices_output, input, k, dim, largest, sorted); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/topk/topk_infiniop.cc b/src/infinicore/ops/topk/topk_infiniop.cc deleted file mode 100644 index db7e0471f..000000000 --- a/src/infinicore/ops/topk/topk_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/topk.hpp" -#include - -namespace infinicore::op::topk_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopTopKDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyTopKDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest, bool sorted) { - size_t seed = hash_combine(values_output, indices_output, input, k, dim, largest, sorted); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopTopKDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateTopKDescriptor( - context::getInfiniopHandle(values_output->device()), &desc, - values_output->desc(), indices_output->desc(), input->desc(), k, dim, largest, sorted)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetTopKWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopTopK( - desc, workspace->data(), workspace_size, - values_output->data(), indices_output->data(), input->data(), k, dim, largest, sorted, context::getStream())); -} - -static bool registered = []() { - TopK::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::topk_impl::infiniop diff --git a/src/infinicore/ops/topksoftmax/topksoftmax.cc b/src/infinicore/ops/topksoftmax/topksoftmax.cc deleted file mode 100644 index c0572ca83..000000000 --- a/src/infinicore/ops/topksoftmax/topksoftmax.cc +++ /dev/null @@ -1,30 +0,0 @@ -#include "infinicore/ops/topksoftmax.hpp" - -#include "../../utils.hpp" - -namespace infinicore::op { - -INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(Topksoftmax); - -Topksoftmax::Topksoftmax(Tensor values, - Tensor indices, - const Tensor &x, - const size_t topk, - const int norm) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(values, indices, x); - INFINICORE_GRAPH_OP_DISPATCH(values->device().getType(), values, indices, x, topk, norm); -} - -void Topksoftmax::execute(Tensor values, - Tensor indices, - const Tensor &x, - const size_t topk, - const int norm) { - INFINICORE_GRAPH_OP_RECORD_OR_RUN(Topksoftmax, values, indices, x, topk, norm); -} - -void topksoftmax(Tensor values, Tensor indices, const Tensor &x, const size_t topk, const int norm) { - Topksoftmax::execute(values, indices, x, topk, norm); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/topksoftmax/topksoftmax_infiniop.cc b/src/infinicore/ops/topksoftmax/topksoftmax_infiniop.cc deleted file mode 100644 index 1c016538c..000000000 --- a/src/infinicore/ops/topksoftmax/topksoftmax_infiniop.cc +++ /dev/null @@ -1,56 +0,0 @@ -#include "infinicore/ops/topksoftmax.hpp" - -#include "../infiniop_impl.hpp" - -namespace infinicore::op::topksoftmax_impl::infiniop { - -INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, Topksoftmax, 100); - -struct PlannedMeta { - std::shared_ptr descriptor; - graph::GraphTensor workspace, values, indices, x; - size_t topk; - int norm; -}; - -void *plan(Tensor values, Tensor indices, const Tensor &x, const size_t topk, const int norm) { - size_t seed = hash_combine(values, indices, x); - - INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( - Descriptor, descriptor, Topksoftmax, seed, x->desc()); - - INFINIOP_WORKSPACE_TENSOR(workspace, Topksoftmax, descriptor); - - return new PlannedMeta{ - descriptor, - graph::GraphTensor(workspace), - graph::GraphTensor(values), - graph::GraphTensor(indices), - graph::GraphTensor(x), - topk, - norm}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - - INFINICORE_CHECK_ERROR(infiniopTopksoftmax( - planned->descriptor->desc, - planned->workspace->data(), - planned->workspace->numel(), - planned->values->data(), - planned->indices->data(), - planned->x->data(), - planned->topk, - planned->norm, - context::getStream())); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(Topksoftmax, &plan, &run, cleanup); - -} // namespace infinicore::op::topksoftmax_impl::infiniop diff --git a/src/infinicore/ops/topksoftmax/topksoftmax_infiniops.cc b/src/infinicore/ops/topksoftmax/topksoftmax_infiniops.cc deleted file mode 100644 index 4a4679417..000000000 --- a/src/infinicore/ops/topksoftmax/topksoftmax_infiniops.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "infinicore/ops/topksoftmax.hpp" - -#ifdef ENABLE_INFINIOPS_API -#include "../infiniops_impl.hpp" - -#include "base/topksoftmax_infinilm.h" - -namespace infinicore::op::topksoftmax_impl::infiniops { -namespace { -using TensorMeta = ::infinicore::op::infiniops::TensorMeta; -struct PlannedMeta { - TensorMeta values, indices, x; - graph::GraphTensor values_tensor, indices_tensor, x_tensor; - size_t topk; - int norm; -}; -} // namespace - -void *plan(Tensor values, Tensor indices, const Tensor &x, const size_t topk, const int norm) { - INFINICORE_ASSERT(::infinicore::op::infiniops::isSupportedDevice(values->device().getType())); - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(values, indices, x); - return new PlannedMeta{TensorMeta(values), TensorMeta(indices), TensorMeta(x), graph::GraphTensor(values), graph::GraphTensor(indices), graph::GraphTensor(x), topk, norm}; -} - -void run(void *planned_meta) { - auto planned = reinterpret_cast(planned_meta); - infini::ops::Handle handle; - handle.set_stream(context::getStream()); - infini::ops::Config config; - infini::ops::TopksoftmaxInfinilm::Call( - handle, - config, - planned->x.tensor(planned->x_tensor), - static_cast(planned->topk), - planned->norm != 0, - planned->values.tensor(planned->values_tensor), - planned->indices.tensor(planned->indices_tensor)); -} - -void cleanup(void **planned_meta_ptr) { - delete *reinterpret_cast(planned_meta_ptr); - *planned_meta_ptr = nullptr; -} - -static bool registered = []() { - ::infinicore::op::infiniops::registerSupportedDevices(Topksoftmax::plan_dispatcher(), &plan); - ::infinicore::op::infiniops::registerSupportedDevices(Topksoftmax::run_dispatcher(), &run); - ::infinicore::op::infiniops::registerSupportedDevices(Topksoftmax::cleanup_dispatcher(), &cleanup); - return true; -}(); -} // namespace infinicore::op::topksoftmax_impl::infiniops -#endif diff --git a/src/infinicore/ops/triplet_margin_loss/triplet_margin_loss.cc b/src/infinicore/ops/triplet_margin_loss/triplet_margin_loss.cc deleted file mode 100644 index 5c46eec37..000000000 --- a/src/infinicore/ops/triplet_margin_loss/triplet_margin_loss.cc +++ /dev/null @@ -1,37 +0,0 @@ -#include "infinicore/ops/triplet_margin_loss.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &TripletMarginLoss::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void TripletMarginLoss::execute(Tensor output, Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction) { - dispatcher().lookup(context::getDevice().getType())(output, anchor, positive, negative, margin, p, eps, swap, reduction); -} - -// 3. 函数式接口 -Tensor triplet_margin_loss(Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction) { - Shape output_shape; - if (reduction == 0) { // None - // TripletMarginLoss 输入通常为 (N, D)，reduction='none' 时输出为 (N) - // 取第 0 维作为 Batch Size - output_shape = {anchor->shape()[0]}; - } else { - output_shape = {}; // Scalar - } - - // 使用 anchor 的属性创建输出 Tensor - auto output = Tensor::empty(output_shape, anchor->dtype(), anchor->device()); - - triplet_margin_loss_(output, anchor, positive, negative, margin, p, eps, swap, reduction); - return output; -} - -void triplet_margin_loss_(Tensor output, Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction) { - TripletMarginLoss::execute(output, anchor, positive, negative, margin, p, eps, swap, reduction); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/triplet_margin_loss/triplet_margin_loss_infiniop.cc b/src/infinicore/ops/triplet_margin_loss/triplet_margin_loss_infiniop.cc deleted file mode 100644 index 6ccca3f5d..000000000 --- a/src/infinicore/ops/triplet_margin_loss/triplet_margin_loss_infiniop.cc +++ /dev/null @@ -1,72 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/triplet_margin_loss.hpp" -#include - -namespace infinicore::op::triplet_margin_loss_impl::infiniop { - -// 定义描述符缓存 -thread_local common::OpCache caches( - 100, // capacity - [](infiniopTripletMarginLossDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyTripletMarginLossDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction) { - // 1. 计算 Hash Seed 作为 Cache Key - size_t seed = hash_combine(output, anchor, positive, negative, margin, p, eps, swap, reduction); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopTripletMarginLossDescriptor_t desc = nullptr; - - if (!desc_opt) { - // 2. 创建描述符 - INFINICORE_CHECK_ERROR(infiniopCreateTripletMarginLossDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - anchor->desc(), - positive->desc(), - negative->desc(), - margin, - static_cast(p), - eps, - static_cast(swap), // bool -> int - static_cast(reduction))); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 3. 获取 Workspace 并执行 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetTripletMarginLossWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopTripletMarginLoss( - desc, - workspace->data(), - workspace_size, - output->data(), - anchor->data(), - positive->data(), - negative->data(), - context::getStream())); -} - -static bool registered = []() { - TripletMarginLoss::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::triplet_margin_loss_impl::infiniop diff --git a/src/infinicore/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.cc b/src/infinicore/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.cc deleted file mode 100644 index cec52e806..000000000 --- a/src/infinicore/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.cc +++ /dev/null @@ -1,38 +0,0 @@ -#include "infinicore/ops/triplet_margin_with_distance_loss.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &TripletMarginWithDistanceLoss::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void TripletMarginWithDistanceLoss::execute(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction) { - dispatcher().lookup(context::getDevice().getType())(output, anchor, positive, negative, margin, swap, reduction); -} - -// 3. 函数式接口 -Tensor triplet_margin_with_distance_loss(Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction) { - Shape out_shape; - - // 推断输出形状 - if (reduction == 0) { - // Reduction::None -> 输出形状取决于输入的广播结果 - out_shape = anchor->shape(); - } else { - // Reduction::Mean 或 Reduction::Sum -> 输出为标量 - out_shape = {}; - } - - auto output = Tensor::empty(out_shape, anchor->dtype(), anchor->device()); - - triplet_margin_with_distance_loss_(output, anchor, positive, negative, margin, swap, reduction); - return output; -} - -void triplet_margin_with_distance_loss_(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction) { - TripletMarginWithDistanceLoss::execute(output, anchor, positive, negative, margin, swap, reduction); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss_infiniop.cc b/src/infinicore/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss_infiniop.cc deleted file mode 100644 index 5d84f4ec9..000000000 --- a/src/infinicore/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss_infiniop.cc +++ /dev/null @@ -1,66 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/triplet_margin_with_distance_loss.hpp" -#include - -namespace infinicore::op::triplet_margin_with_distance_loss_impl::infiniop { - -// 定义描述符缓存 -thread_local common::OpCache caches( - 100, // capacity - [](infiniopTripletMarginWithDistanceLossDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyTripletMarginWithDistanceLossDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction) { - size_t seed = hash_combine(output, anchor, positive, negative, margin, swap, reduction); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopTripletMarginWithDistanceLossDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateTripletMarginWithDistanceLossDescriptor( - context::getInfiniopHandle(anchor->device()), - &desc, - output->desc(), - anchor->desc(), - positive->desc(), - negative->desc(), - static_cast(margin), - static_cast(swap), - static_cast(reduction))); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetTripletMarginWithDistanceLossWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopTripletMarginWithDistanceLoss( - desc, - workspace->data(), - workspace_size, - output->data(), - anchor->data(), - positive->data(), - negative->data(), - context::getStream())); -} - -static bool registered = []() { - TripletMarginWithDistanceLoss::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::triplet_margin_with_distance_loss_impl::infiniop diff --git a/src/infinicore/ops/unfold/unfold.cc b/src/infinicore/ops/unfold/unfold.cc deleted file mode 100644 index 1d96902ee..000000000 --- a/src/infinicore/ops/unfold/unfold.cc +++ /dev/null @@ -1,72 +0,0 @@ -#include "infinicore/ops/unfold.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &Unfold::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -// 2. Execute 方法实现 -void Unfold::execute(Tensor output, Tensor input, - const std::vector &kernel_sizes, - const std::vector &dilations, - const std::vector &paddings, - const std::vector &strides) { - dispatcher().lookup(context::getDevice().getType())(output, input, kernel_sizes, dilations, paddings, strides); -} - -// 3. 函数式接口 -Tensor unfold(Tensor input, - std::vector kernel_sizes, - std::vector dilations, - std::vector paddings, - std::vector strides) { - - // 基础维度校验与获取 - const auto &input_shape = input->shape(); - int64_t n_dim = input->ndim(); - int64_t spatial_dims = n_dim - 2; // N, C, D1, D2... -> spatial starts at 2 - int64_t N = input_shape[0]; - int64_t C = input_shape[1]; - - // 计算 dim 1: C * kernel_sizes[0] * kernel_sizes[1] ... - int64_t output_dim1 = C; - for (auto k : kernel_sizes) { - output_dim1 *= k; - } - - int64_t L = 1; - for (int i = 0; i < spatial_dims; ++i) { - int64_t input_dim = input_shape[i + 2]; - int64_t k = kernel_sizes[i]; - int64_t p = paddings[i]; - int64_t d = dilations[i]; - int64_t s = strides[i]; - - // 公式: out = floor((in + 2*p - d*(k-1) - 1) / s + 1) - int64_t output_spatial = (input_dim + 2 * p - d * (k - 1) - 1) / s + 1; - L *= output_spatial; - } - Shape output_shape = { - static_cast(N), - static_cast(output_dim1), - static_cast(L)}; - - auto output = Tensor::empty(output_shape, input->dtype(), input->device()); - - unfold_(output, input, kernel_sizes, dilations, paddings, strides); - return output; -} - -// 4. In-place / 显式输出接口 -void unfold_(Tensor output, Tensor input, - std::vector kernel_sizes, - std::vector dilations, - std::vector paddings, - std::vector strides) { - Unfold::execute(output, input, kernel_sizes, dilations, paddings, strides); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/unfold/unfold_infiniop.cc b/src/infinicore/ops/unfold/unfold_infiniop.cc deleted file mode 100644 index 5bb995353..000000000 --- a/src/infinicore/ops/unfold/unfold_infiniop.cc +++ /dev/null @@ -1,107 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/unfold.hpp" -#include -#include -#include - -namespace infinicore::op::unfold_impl::infiniop { - -// 定义描述符缓存 -thread_local common::OpCache caches( - 100, // capacity - [](infiniopUnfoldDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyUnfoldDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, - const std::vector &kernel_sizes, - const std::vector &dilations, - const std::vector &paddings, - const std::vector &strides) { - - // 1. 计算 Hash Key (修复点：手动拆解，避开 hash.hpp 的递归 bug 和 vector 不支持问题) - size_t seed = 0; - - // 基础 Tensor 支持 (hash.hpp 中有 Tensor 重载) - hash_combine(seed, output); - hash_combine(seed, input); - - // Vector 类型必须手动遍历 (hash.hpp 不支持 vector 直接 hash) - for (auto v : kernel_sizes) { - hash_combine(seed, v); - } - for (auto v : dilations) { - hash_combine(seed, v); - } - for (auto v : paddings) { - hash_combine(seed, v); - } - for (auto v : strides) { - hash_combine(seed, v); - } - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopUnfoldDescriptor_t desc = nullptr; - - if (!desc_opt) { - // 2. 创建描述符 - - // 辅助函数：将 int64_t vector 转换为 int vector 以匹配 C API 的 int* 签名 - auto to_int_vec = [](const std::vector &src) { - std::vector dst(src.size()); - std::transform(src.begin(), src.end(), dst.begin(), - [](int64_t val) { return static_cast(val); }); - return dst; - }; - - std::vector k_int = to_int_vec(kernel_sizes); - std::vector s_int = to_int_vec(strides); - std::vector p_int = to_int_vec(paddings); - std::vector d_int = to_int_vec(dilations); - - INFINICORE_CHECK_ERROR(infiniopCreateUnfoldDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc(), - k_int.data(), - s_int.data(), - p_int.data(), - d_int.data())); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 3. 获取 Workspace 并执行 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetUnfoldWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopUnfold( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -// 4. 注册算子实现 -static bool registered = []() { - Unfold::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::unfold_impl::infiniop diff --git a/src/infinicore/ops/upsample_bilinear/upsample_bilinear.cc b/src/infinicore/ops/upsample_bilinear/upsample_bilinear.cc deleted file mode 100644 index bd6cf1c8b..000000000 --- a/src/infinicore/ops/upsample_bilinear/upsample_bilinear.cc +++ /dev/null @@ -1,41 +0,0 @@ -#include "infinicore/ops/upsample_bilinear.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &UpsampleBilinear::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void UpsampleBilinear::execute(Tensor output, Tensor input, bool align_corners) { - dispatcher().lookup(context::getDevice().getType())(output, input, align_corners); -} - -// 3. 函数式接口 -Tensor upsample_bilinear(Tensor input, std::vector output_size, bool align_corners) { - // 构造输出 Shape - // 假设 input 是 (N, C, H_in, W_in) 或 (C, H_in, W_in) - // output_size 通常只包含 (H_out, W_out) - Shape input_shape = input->shape(); - size_t ndim = input_shape.size(); - - Shape output_shape = input_shape; - - // 更新最后两个维度为 output_size 指定的大小 - if (output_size.size() == 2 && ndim >= 2) { - output_shape[ndim - 2] = output_size[0]; - output_shape[ndim - 1] = output_size[1]; - } - - auto output = Tensor::empty(output_shape, input->dtype(), input->device()); - - upsample_bilinear_(output, input, align_corners); - return output; -} - -void upsample_bilinear_(Tensor output, Tensor input, bool align_corners) { - UpsampleBilinear::execute(output, input, align_corners); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/upsample_bilinear/upsample_bilinear_infiniopo.cc b/src/infinicore/ops/upsample_bilinear/upsample_bilinear_infiniopo.cc deleted file mode 100644 index a599836b6..000000000 --- a/src/infinicore/ops/upsample_bilinear/upsample_bilinear_infiniopo.cc +++ /dev/null @@ -1,66 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/upsample_bilinear.hpp" -#include - -namespace infinicore::op::upsample_bilinear_impl::infiniop { - -// 定义描述符缓存 -thread_local common::OpCache caches( - 100, // capacity - [](infiniopUpsampleBilinearDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyUpsampleBilinearDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, bool align_corners) { - // 1. 计算 Hash Seed - // align_corners 是 bool，可以直接参与 hash - size_t seed = hash_combine(output, input, align_corners); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopUpsampleBilinearDescriptor_t desc = nullptr; - - if (!desc_opt) { - // 3. 创建描述符 - // 注意：C 接口中 align_corners 通常用 int 传递 - INFINICORE_CHECK_ERROR(infiniopCreateUpsampleBilinearDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc(), - static_cast(align_corners))); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 4. 获取 Workspace 并执行 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetUpsampleBilinearWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopUpsampleBilinear( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -static bool registered = []() { - UpsampleBilinear::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::upsample_bilinear_impl::infiniop diff --git a/src/infinicore/ops/upsample_nearest/upsample_nearest.cc b/src/infinicore/ops/upsample_nearest/upsample_nearest.cc deleted file mode 100644 index 6acbb8966..000000000 --- a/src/infinicore/ops/upsample_nearest/upsample_nearest.cc +++ /dev/null @@ -1,61 +0,0 @@ -#include "infinicore/ops/upsample_nearest.hpp" -#include -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &UpsampleNearest::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void UpsampleNearest::execute(Tensor output, Tensor input) { - dispatcher().lookup(context::getDevice().getType())(output, input); -} - -// 3. 函数式接口 -Tensor upsample_nearest(Tensor input, const std::vector &output_size) { - Shape input_shape = input->shape(); - size_t ndim = input_shape.size(); - - // 校验 - if (ndim < 3 || ndim > 4) { - if (ndim != 3 && ndim != 4) { - throw std::runtime_error("upsample_nearest: Only supports 3D (N,C,W) or 4D (N,C,H,W) input"); - } - } - - Shape output_shape = input_shape; - - if (ndim == 3) { - // [N, C, W] - // output_size 可能是 [W_out] (size=1) 或者 [1, W_out] (size=2) - int64_t target_w = 0; - if (output_size.size() == 1) { - target_w = output_size[0]; - } else if (output_size.size() == 2) { - target_w = output_size[1]; - } else { - throw std::runtime_error("upsample_nearest: output_size for 3D input must be [w] or [1, w]"); - } - output_shape[2] = target_w; - - } else if (ndim == 4) { - // [N, C, H, W] - if (output_size.size() != 2) { - throw std::runtime_error("upsample_nearest: output_size for 4D input must be [h, w]"); - } - output_shape[2] = output_size[0]; - output_shape[3] = output_size[1]; - } - - auto output = Tensor::empty(output_shape, input->dtype(), input->device()); - - upsample_nearest_(output, input); - return output; -} - -void upsample_nearest_(Tensor output, Tensor input) { - UpsampleNearest::execute(output, input); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/upsample_nearest/upsample_nearest_infiniop.cc b/src/infinicore/ops/upsample_nearest/upsample_nearest_infiniop.cc deleted file mode 100644 index 388b01b17..000000000 --- a/src/infinicore/ops/upsample_nearest/upsample_nearest_infiniop.cc +++ /dev/null @@ -1,59 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/upsample_nearest.hpp" -#include - -namespace infinicore::op::upsample_nearest_impl::infiniop { - -// 定义描述符缓存 -thread_local common::OpCache caches( - 100, // capacity - [](infiniopUpsampleNearestDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyUpsampleNearestDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input) { - size_t seed = hash_combine(output, input); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopUpsampleNearestDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateUpsampleNearestDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc())); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetUpsampleNearestWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopUpsampleNearest( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -static bool registered = []() { - UpsampleNearest::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::upsample_nearest_impl::infiniop diff --git a/src/infinicore/ops/vander/vander.cc b/src/infinicore/ops/vander/vander.cc deleted file mode 100644 index a78c3c215..000000000 --- a/src/infinicore/ops/vander/vander.cc +++ /dev/null @@ -1,35 +0,0 @@ -#include "infinicore/ops/vander.hpp" - -namespace infinicore::op { - -// 1. 定义 Dispatcher 单例 -common::OpDispatcher &Vander::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -// 2. Execute 方法实现 -void Vander::execute(Tensor output, Tensor input, int64_t N, bool increasing) { - dispatcher().lookup(context::getDevice().getType())(output, input, N, increasing); -} - -// 3. 函数式接口 -Tensor vander(Tensor input, int64_t N, bool increasing) { - int64_t input_size = input->shape()[0]; - int64_t cols = (N > 0) ? N : input_size; - Shape output_shape = { - static_cast(input_size), - static_cast(cols)}; - - auto output = Tensor::empty(output_shape, input->dtype(), input->device()); - - vander_(output, input, N, increasing); - return output; -} - -// 4. In-place / 显式输出接口 -void vander_(Tensor output, Tensor input, int64_t N, bool increasing) { - Vander::execute(output, input, N, increasing); -} - -} // namespace infinicore::op diff --git a/src/infinicore/ops/vander/vander_infiniop.cc b/src/infinicore/ops/vander/vander_infiniop.cc deleted file mode 100644 index 74a7839c6..000000000 --- a/src/infinicore/ops/vander/vander_infiniop.cc +++ /dev/null @@ -1,68 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/vander.hpp" -#include - -namespace infinicore::op::vander_impl::infiniop { - -// 定义描述符缓存 -thread_local common::OpCache caches( - 100, // capacity - [](infiniopVanderDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyVanderDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor output, Tensor input, int64_t N, bool increasing) { - // 1. 计算 Hash Key - // 直接组合 output, input 以及标量参数 N, increasing - size_t seed = hash_combine(output, input, N, increasing); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopVanderDescriptor_t desc = nullptr; - - if (!desc_opt) { - // 2. 创建描述符 - // 注意：将 int64_t N 转换为 int，bool increasing 转换为 int 以匹配 C API 签名 - INFINICORE_CHECK_ERROR(infiniopCreateVanderDescriptor( - context::getInfiniopHandle(output->device()), - &desc, - output->desc(), - input->desc(), - static_cast(N), - static_cast(increasing))); - - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - // 3. 获取 Workspace 并执行 - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetVanderWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopVander( - desc, - workspace->data(), - workspace_size, - output->data(), - input->data(), - context::getStream())); -} - -// 4. 注册算子实现 -static bool registered = []() { - Vander::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::vander_impl::infiniop diff --git a/src/infinicore/ops/var/var.cc b/src/infinicore/ops/var/var.cc deleted file mode 100644 index bc0849e64..000000000 --- a/src/infinicore/ops/var/var.cc +++ /dev/null @@ -1,68 +0,0 @@ -#include "infinicore/ops/var.hpp" - -#include "../../utils.hpp" -#include -#include - -namespace infinicore::op { - -common::OpDispatcher &Var::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Var::execute(Tensor var_output, Tensor input, std::vector dim, bool unbiased, bool keepdim) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(var_output, input); - infinicore::context::setDevice(input->device()); - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Var implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(var_output, input, dim, unbiased, keepdim); -} - -Tensor var(Tensor input, std::vector dim, bool unbiased, bool keepdim) { - auto in_shape = input->shape(); - std::vector out_shape; - if (dim.empty()) { - for (size_t i = 0; i < in_shape.size(); i++) { - dim.push_back(i); - } - } - std::sort(dim.begin(), dim.end()); - if (dim.size() == in_shape.size() && !keepdim) { - out_shape = {}; - } else { - if (keepdim) { - size_t j = 0; - for (size_t i = 0; i < in_shape.size(); i++) { - if (j < dim.size() && dim[j] == i) { - out_shape.push_back(1); - j++; - } else { - out_shape.push_back(in_shape[i]); - } - } - } else { - size_t j = 0; - for (size_t i = 0; i < in_shape.size(); i++) { - if (j < dim.size() && dim[j] == i) { - j++; - } else { - out_shape.push_back(in_shape[i]); - } - } - } - } - auto var_output = Tensor::empty(out_shape, input->dtype(), input->device()); - var_(var_output, input, dim, unbiased, keepdim); - return var_output; -} - -void var_(Tensor var_output, Tensor input, std::vector dim, bool unbiased, bool keepdim) { - Var::execute(var_output, input, dim, unbiased, keepdim); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/var/var_infiniop.cc b/src/infinicore/ops/var/var_infiniop.cc deleted file mode 100644 index 9c162cec7..000000000 --- a/src/infinicore/ops/var/var_infiniop.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/var.hpp" -#include - -namespace infinicore::op::var_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopVarDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyVarDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor var_output, Tensor input, std::vector dim, bool unbiased, bool keepdim) { - size_t seed = hash_combine(var_output, input, dim.size(), unbiased, keepdim); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopVarDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateVarDescriptor( - context::getInfiniopHandle(var_output->device()), &desc, - var_output->desc(), input->desc(), dim.data(), dim.size(), unbiased, keepdim)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetVarWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopVar( - desc, workspace->data(), workspace_size, - var_output->data(), input->data(), dim.data(), dim.size(), unbiased, keepdim, context::getStream())); -} - -static bool registered = []() { - Var::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::var_impl::infiniop diff --git a/src/infinicore/ops/var_mean/var_mean.cc b/src/infinicore/ops/var_mean/var_mean.cc deleted file mode 100644 index 817be7bcf..000000000 --- a/src/infinicore/ops/var_mean/var_mean.cc +++ /dev/null @@ -1,69 +0,0 @@ -#include "infinicore/ops/var_mean.hpp" - -#include "../../utils.hpp" -#include -#include - -namespace infinicore::op { - -common::OpDispatcher &Var_Mean::dispatcher() { - static common::OpDispatcher dispatcher_; - return dispatcher_; -}; - -void Var_Mean::execute(Tensor var_output, Tensor mean_output, Tensor input, std::vector dim, bool unbiased, bool keepdim) { - INFINICORE_ASSERT_TENSORS_SAME_DEVICE(var_output, mean_output, input); - infinicore::context::setDevice(input->device()); - auto device_type = context::getDevice().getType(); - auto func = dispatcher().lookup(device_type); - - if (func == nullptr) { - throw std::runtime_error("No Var_Mean implementation found for device type: " + std::to_string(static_cast(device_type))); - } - - func(var_output, mean_output, input, dim, unbiased, keepdim); -} - -std::pair var_mean(Tensor input, std::vector dim, bool unbiased, bool keepdim) { - auto in_shape = input->shape(); - std::vector out_shape; - if (dim.empty()) { - for (size_t i = 0; i < in_shape.size(); i++) { - dim.push_back(i); - } - } - std::sort(dim.begin(), dim.end()); - if (dim.size() == in_shape.size() && !keepdim) { - out_shape = {}; - } else { - if (keepdim) { - size_t j = 0; - for (size_t i = 0; i < in_shape.size(); i++) { - if (j < dim.size() && dim[j] == i) { - out_shape.push_back(1); - j++; - } else { - out_shape.push_back(in_shape[i]); - } - } - } else { - size_t j = 0; - for (size_t i = 0; i < in_shape.size(); i++) { - if (j < dim.size() && dim[j] == i) { - j++; - } else { - out_shape.push_back(in_shape[i]); - } - } - } - } - auto var_output = Tensor::empty(out_shape, input->dtype(), input->device()); - auto mean_output = Tensor::empty(out_shape, input->dtype(), input->device()); - var_mean_(var_output, mean_output, input, dim, unbiased, keepdim); - return {var_output, mean_output}; -} - -void var_mean_(Tensor var_output, Tensor mean_output, Tensor input, std::vector dim, bool unbiased, bool keepdim) { - Var_Mean::execute(var_output, mean_output, input, dim, unbiased, keepdim); -} -} // namespace infinicore::op diff --git a/src/infinicore/ops/var_mean/var_mean_infiniop.cc b/src/infinicore/ops/var_mean/var_mean_infiniop.cc deleted file mode 100644 index 5f74a3dbf..000000000 --- a/src/infinicore/ops/var_mean/var_mean_infiniop.cc +++ /dev/null @@ -1,54 +0,0 @@ -#include "../../utils.hpp" -#include "infinicore/common/hash.hpp" -#include "infinicore/ops/common/cache.hpp" -#include "infinicore/ops/var_mean.hpp" -#include - -// todo 实现需要修改calculate函数 - -namespace infinicore::op::var_mean_impl::infiniop { - -thread_local common::OpCache caches( - 100, // capacity - [](infiniopVarMeanDescriptor_t &desc) { - if (desc != nullptr) { - INFINICORE_CHECK_ERROR(infiniopDestroyVarMeanDescriptor(desc)); - desc = nullptr; - } - }); - -void calculate(Tensor var_output, Tensor mean_output, Tensor input, std::vector dim, bool unbiased, bool keepdim) { - size_t seed = hash_combine(var_output, mean_output, input, dim.size(), unbiased, keepdim); - - auto device_type = context::getDevice().getType(); - auto device_index = context::getDevice().getIndex(); - - auto &cache = caches.getCache(device_type, device_index); - - auto desc_opt = cache.get(seed); - infiniopVarMeanDescriptor_t desc = nullptr; - - if (!desc_opt) { - INFINICORE_CHECK_ERROR(infiniopCreateVarMeanDescriptor( - context::getInfiniopHandle(var_output->device()), &desc, - var_output->desc(), mean_output->desc(), input->desc(), dim.data(), dim.size(), unbiased, keepdim)); - cache.put(seed, desc); - } else { - desc = *desc_opt; - } - - size_t workspace_size = 0; - INFINICORE_CHECK_ERROR(infiniopGetVarMeanWorkspaceSize(desc, &workspace_size)); - std::shared_ptr workspace = context::allocateMemory(workspace_size); - - INFINICORE_CHECK_ERROR(infiniopVarMean( - desc, workspace->data(), workspace_size, - var_output->data(), mean_output->data(), input->data(), dim.data(), dim.size(), unbiased, keepdim, context::getStream())); -} - -static bool registered = []() { - Var_Mean::dispatcher().registerAll(&calculate, false); - return true; -}(); - -} // namespace infinicore::op::var_mean_impl::infiniop diff --git a/src/infinicore/pybind11/analyzer.hpp b/src/infinicore/pybind11/analyzer.hpp deleted file mode 100644 index 1d799991a..000000000 --- a/src/infinicore/pybind11/analyzer.hpp +++ /dev/null @@ -1,168 +0,0 @@ -#pragma once - -#include -#include - -#include "infinicore/analyzer/mutual_awareness_analyzer.hpp" - -namespace py = pybind11; - -namespace infinicore::analyzer::pybind { - -inline void bind(py::module &m) { - auto analyzer_mod = m.def_submodule("analyzer", - "Hardware-Task Mutual Awareness Analysis Module"); - - // --- Enums --- - py::enum_(analyzer_mod, "PhaseType") - .value("UNKNOWN", PhaseType::UNKNOWN) - .value("PREFILL", PhaseType::PREFILL) - .value("DECODE", PhaseType::DECODE) - .value("ATTENTION_DENSE", PhaseType::ATTENTION_DENSE) - .value("GEMM_MLP_DENSE", PhaseType::GEMM_MLP_DENSE) - .value("MOE_ROUTING", PhaseType::MOE_ROUTING) - .value("KV_CACHE", PhaseType::KV_CACHE) - .value("COMMUNICATION", PhaseType::COMMUNICATION) - .export_values(); - - py::enum_(analyzer_mod, "BottleneckType") - .value("COMPUTE_BOUND", BottleneckType::COMPUTE_BOUND) - .value("MEMORY_BOUND", BottleneckType::MEMORY_BOUND) - .value("BANDWIDTH_BOUND", BottleneckType::BANDWIDTH_BOUND) - .value("COMMUNICATION_BOUND", BottleneckType::COMMUNICATION_BOUND) - .value("BALANCED", BottleneckType::BALANCED) - .export_values(); - - py::enum_(analyzer_mod, "OptimizationGoal") - .value("LATENCY_FIRST", OptimizationGoal::LATENCY_FIRST) - .value("THROUGHPUT_FIRST", OptimizationGoal::THROUGHPUT_FIRST) - .value("MEMORY_SAFE", OptimizationGoal::MEMORY_SAFE) - .value("STABILITY_FIRST", OptimizationGoal::STABILITY_FIRST) - .export_values(); - - py::enum_(analyzer_mod, "OpType") - .value("UNKNOWN", OpType::UNKNOWN) - .value("ATTENTION", OpType::ATTENTION) - .value("FLASH_ATTENTION", OpType::FLASH_ATTENTION) - .value("GEMM", OpType::GEMM) - .value("LINEAR", OpType::LINEAR) - .value("MATMUL", OpType::MATMUL) - .value("SILU", OpType::SILU) - .value("GELU", OpType::GELU) - .value("RMS_NORM", OpType::RMS_NORM) - .value("KV_CACHING", OpType::KV_CACHING) - .value("PAGED_CACHING", OpType::PAGED_CACHING) - .value("EMBEDDING", OpType::EMBEDDING) - .value("ROPE", OpType::ROPE) - .export_values(); - - // --- StrategyHint --- - py::class_(analyzer_mod, "StrategyHint") - .def(py::init<>()) - .def_readwrite("prefer_fused_ops", &StrategyHint::prefer_fused_ops) - .def_readwrite("prefer_in_place", &StrategyHint::prefer_in_place) - .def_readwrite("prefer_recomputation", &StrategyHint::prefer_recomputation) - .def_readwrite("prefer_async_comm", &StrategyHint::prefer_async_comm) - .def("__repr__", [](const StrategyHint &s) { - return ""; - }); - - // --- GlobalSemanticIntent --- - py::class_(analyzer_mod, "GlobalSemanticIntent") - .def(py::init<>()) - .def_readwrite("current_phase", &GlobalSemanticIntent::current_phase) - .def_readwrite("primary_bottleneck", &GlobalSemanticIntent::primary_bottleneck) - .def_readwrite("goal", &GlobalSemanticIntent::goal) - .def_readwrite("compute_intensity", &GlobalSemanticIntent::compute_intensity) - .def_readwrite("confidence", &GlobalSemanticIntent::confidence) - .def_readwrite("strategy", &GlobalSemanticIntent::strategy) - .def_readwrite("timestamp_ns", &GlobalSemanticIntent::timestamp_ns) - .def("__repr__", [](const GlobalSemanticIntent &i) { - return ""; - }); - - // --- DeviceLocalIntent --- - py::class_(analyzer_mod, "DeviceLocalIntent") - .def(py::init<>()) - .def_readwrite("device_id", &DeviceLocalIntent::device_id) - .def_readwrite("memory_usage_ratio", &DeviceLocalIntent::memory_usage_ratio) - .def_readwrite("memory_available_bytes", &DeviceLocalIntent::memory_available_bytes) - .def_readwrite("local_bottleneck", &DeviceLocalIntent::local_bottleneck) - .def_readwrite("compute_utilization", &DeviceLocalIntent::compute_utilization) - .def_readwrite("memory_bandwidth_utilization", &DeviceLocalIntent::memory_bandwidth_utilization) - .def_readwrite("communication_time_ratio", &DeviceLocalIntent::communication_time_ratio) - .def_readwrite("resource_confidence", &DeviceLocalIntent::resource_confidence) - .def("__repr__", [](const DeviceLocalIntent &d) { - return ""; - }); - - // --- OptimizationIntent --- - py::class_(analyzer_mod, "OptimizationIntent") - .def(py::init<>()) - .def_readwrite("global_intent", &OptimizationIntent::global) - .def_readwrite("per_device", &OptimizationIntent::per_device) - .def("get_device_intent", &OptimizationIntent::getDeviceIntent, - py::return_value_policy::reference, - py::arg("device_id")) - .def("__repr__", [](const OptimizationIntent &i) { - return ""; - }); - - // --- Top-level functions --- - analyzer_mod.def("analyze", &analyzeCurrentState, - "Analyze current state and return an OptimizationIntent"); - analyzer_mod.def("get_current_phase", &getCurrentPhase, - "Get the current detected task phase"); - analyzer_mod.def("set_enabled", &setAnalyzerEnabled, - "Enable/disable the mutual awareness analyzer", - py::arg("enabled")); - analyzer_mod.def( - "trace_op_for_test", - [](OpType op_type, - const std::vector &shape, - uint8_t dtype, - uint8_t device_type, - int device_id) { - traceOp(op_type, shape.data(), shape.size(), dtype, device_type, static_cast(device_id)); - }, - "Inject an OpTrace entry for testing", - py::arg("op_type"), - py::arg("shape"), - py::arg("dtype") = 0, - py::arg("device_type") = 0, - py::arg("device_id") = 0); - analyzer_mod.def( - "clear_trace", []() { - getGlobalOpTrace().clear(); - MutualAwarenessAnalyzer::instance().clearGraphCache(); - }, - "Clear the global OpTrace ring and analyzer graph cache"); - - // --- Access to analyzer instance for advanced usage --- - analyzer_mod.def("get_analyzer", &MutualAwarenessAnalyzer::instance, - py::return_value_policy::reference, - "Get the MutualAwarenessAnalyzer singleton instance"); - - py::class_(analyzer_mod, "MutualAwarenessAnalyzer") - .def("analyze", py::overload_cast<>(&MutualAwarenessAnalyzer::analyze)) - .def("get_current_phase", &MutualAwarenessAnalyzer::getCurrentPhase) - .def("last_intent", &MutualAwarenessAnalyzer::lastIntent, - py::return_value_policy::reference) - .def("set_enabled", &MutualAwarenessAnalyzer::setEnabled) - .def("is_enabled", &MutualAwarenessAnalyzer::isEnabled) - .def("on_graph_recording_stop", &MutualAwarenessAnalyzer::onGraphRecordingStop) - .def("clear_graph_cache", &MutualAwarenessAnalyzer::clearGraphCache); -} - -} // namespace infinicore::analyzer::pybind diff --git a/src/infinicore/pybind11/context.hpp b/src/infinicore/pybind11/context.hpp deleted file mode 100644 index 9c24a322e..000000000 --- a/src/infinicore/pybind11/context.hpp +++ /dev/null @@ -1,34 +0,0 @@ -#pragma once - -#include - -#include "infinicore.hpp" - -namespace py = pybind11; - -namespace infinicore::context { - -inline void bind(py::module &m) { - // Device management - m.def("get_device", &getDevice, "Get the current active device"); - m.def("get_device_count", &getDeviceCount, - "Get the number of available devices of a specific type", - py::arg("device_type")); - m.def("set_device", &setDevice, - "Set the current active device", - py::arg("device")); - - // Stream and handle management - m.def("get_stream", &getStream, "Get the current stream"); - - // Synchronization - m.def("sync_stream", &syncStream, "Synchronize the current stream"); - m.def("sync_device", &syncDevice, "Synchronize the current device"); - - // Graph - m.def("is_graph_recording", &isGraphRecording, "Check if graph recording is turned on"); - m.def("start_graph_recording", &startGraphRecording, "Start graph recording"); - m.def("stop_graph_recording", &stopGraphRecording, "Stop graph recording and return the graph"); -} - -} // namespace infinicore::context diff --git a/src/infinicore/pybind11/device.hpp b/src/infinicore/pybind11/device.hpp deleted file mode 100644 index 74cb0a23b..000000000 --- a/src/infinicore/pybind11/device.hpp +++ /dev/null @@ -1,36 +0,0 @@ -#pragma once - -#include - -#include "infinicore.hpp" - -namespace py = pybind11; - -namespace infinicore::device { - -inline void bind(py::module &m) { - py::class_ device(m, "Device"); - - py::enum_(device, "Type") - .value("CPU", Device::Type::CPU) - .value("NVIDIA", Device::Type::NVIDIA) - .value("CAMBRICON", Device::Type::CAMBRICON) - .value("ASCEND", Device::Type::ASCEND) - .value("METAX", Device::Type::METAX) - .value("MOORE", Device::Type::MOORE) - .value("ILUVATAR", Device::Type::ILUVATAR) - .value("QY", Device::Type::QY) - .value("KUNLUN", Device::Type::KUNLUN) - .value("HYGON", Device::Type::HYGON) - .value("ALI", Device::Type::ALI) - .value("COUNT", Device::Type::COUNT); - - device - .def(py::init(), - py::arg("type") = Device::Type::CPU, py::arg("index") = 0) - .def_property_readonly("type", &Device::getType) - .def_property_readonly("index", &Device::getIndex) - .def("__str__", static_cast(&Device::toString)); -} - -} // namespace infinicore::device diff --git a/src/infinicore/pybind11/device_event.hpp b/src/infinicore/pybind11/device_event.hpp deleted file mode 100644 index f482422bb..000000000 --- a/src/infinicore/pybind11/device_event.hpp +++ /dev/null @@ -1,42 +0,0 @@ -#pragma once - -#include "infinicore.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::device_event { - -inline void bind(py::module &m) { - py::class_(m, "DeviceEvent") - .def(py::init<>(), "Construct a DeviceEvent on the current device") - .def(py::init(), "Construct a DeviceEvent with specific flags", py::arg("flags")) - .def(py::init(), "Construct a DeviceEvent on a specific device", py::arg("device")) - .def(py::init(), "Construct a DeviceEvent on a specific device with flags", - py::arg("device"), py::arg("flags")) - - .def("record", py::overload_cast<>(&DeviceEvent::record), - "Record the event on the current stream of its device") - .def("record", py::overload_cast(&DeviceEvent::record), - "Record the event on a specific stream", py::arg("stream")) - - .def("synchronize", &DeviceEvent::synchronize, - "Wait for the event to complete (blocking)") - .def("query", &DeviceEvent::query, - "Check if the event has been completed") - - .def("elapsed_time", &DeviceEvent::elapsed_time, - "Calculate elapsed time between this event and another event (in milliseconds)", - py::arg("other")) - - .def("wait", &DeviceEvent::wait, - "Make a stream wait for this event to complete", - py::arg("stream") = nullptr) - - .def_property_readonly("device", &DeviceEvent::device, - "Get the device where this event was created") - .def_property_readonly("is_recorded", &DeviceEvent::is_recorded, - "Check if the event has been recorded"); -} - -} // namespace infinicore::device_event diff --git a/src/infinicore/pybind11/dtype.hpp b/src/infinicore/pybind11/dtype.hpp deleted file mode 100644 index cee339023..000000000 --- a/src/infinicore/pybind11/dtype.hpp +++ /dev/null @@ -1,34 +0,0 @@ -#pragma once - -#include - -#include "infinicore.hpp" - -namespace py = pybind11; - -namespace infinicore::dtype { - -inline void bind(py::module &m) { - py::enum_(m, "DataType") - .value("BYTE", DataType::BYTE) - .value("BOOL", DataType::BOOL) - .value("I8", DataType::I8) - .value("I16", DataType::I16) - .value("I32", DataType::I32) - .value("I64", DataType::I64) - .value("U8", DataType::U8) - .value("U16", DataType::U16) - .value("U32", DataType::U32) - .value("U64", DataType::U64) - .value("F8", DataType::F8) - .value("F16", DataType::F16) - .value("F32", DataType::F32) - .value("F64", DataType::F64) - .value("C16", DataType::C16) - .value("C32", DataType::C32) - .value("C64", DataType::C64) - .value("C128", DataType::C128) - .value("BF16", DataType::BF16); -} - -} // namespace infinicore::dtype diff --git a/src/infinicore/pybind11/from_list.cc b/src/infinicore/pybind11/from_list.cc deleted file mode 100644 index 47b9eb05e..000000000 --- a/src/infinicore/pybind11/from_list.cc +++ /dev/null @@ -1,356 +0,0 @@ -#include "../../utils/custom_types.h" -#include "infinicore.hpp" -#include -#include -#include -#include -#include -#include - -namespace py = pybind11; - -namespace infinicore::tensor { - -namespace { - -constexpr const char *kElementTypeError = "List elements must be bool, int, or float"; - -// Zero-copy view over list / tuple / generic sequence; items() exposes the -// underlying PyObject** for hot loop indexing. - -class SequenceView { -public: - enum class Kind : uint8_t { List, - Tuple, - Fast }; - - explicit SequenceView(py::handle obj, const char *err = "Could not convert object to sequence") - : SequenceView(obj.ptr(), err) {} - - explicit SequenceView(PyObject *obj, const char *err = "Could not convert object to sequence") { - if (PyList_Check(obj)) { - kind_ = Kind::List; - borrowed_ = obj; - fast_ = nullptr; - return; - } - if (PyTuple_Check(obj)) { - kind_ = Kind::Tuple; - borrowed_ = obj; - fast_ = nullptr; - return; - } - kind_ = Kind::Fast; - borrowed_ = nullptr; - fast_ = PySequence_Fast(obj, err); - if (!fast_) { - throw py::error_already_set(); - } - } - - ~SequenceView() { - if (fast_ != nullptr) { - Py_DECREF(fast_); - } - } - - SequenceView(const SequenceView &) = delete; - SequenceView &operator=(const SequenceView &) = delete; - - bool is_list() const { - return kind_ == Kind::List; - } - - Py_ssize_t size() const { - return size_of(ptr()); - } - - PyObject *item(Py_ssize_t index) const { - return item_at(ptr(), index); - } - - // Direct pointer to ob_item; valid for list / tuple / PySequence_Fast result. - PyObject **items() const { - return PySequence_Fast_ITEMS(ptr()); - } - - static Py_ssize_t size_of(PyObject *obj) { - if (PyList_Check(obj)) { - return PyList_GET_SIZE(obj); - } - if (PyTuple_Check(obj)) { - return PyTuple_GET_SIZE(obj); - } - return PySequence_Size(obj); - } - - static PyObject *item_at(PyObject *obj, Py_ssize_t index) { - if (PyList_Check(obj)) { - return PyList_GET_ITEM(obj, index); - } - if (PyTuple_Check(obj)) { - return PyTuple_GET_ITEM(obj, index); - } - return PySequence_GetItem(obj, index); - } - -private: - PyObject *ptr() const { - return borrowed_ != nullptr ? borrowed_ : fast_; - } - - Kind kind_ = Kind::Fast; - PyObject *borrowed_ = nullptr; - PyObject *fast_ = nullptr; -}; - -// Shape discovery + uniform per-scalar traversal. - -inline bool is_python_scalar(PyObject *obj) { - return PyBool_Check(obj) || PyFloat_Check(obj) || PyLong_Check(obj); -} - -inline bool is_python_sequence(PyObject *obj) { - return PyList_Check(obj) || PyTuple_Check(obj); -} - -struct ListLayout { - SequenceView seq; - int ndim = 0; - Shape shape; - bool list_of_lists = false; - - explicit ListLayout(py::handle obj) - : seq(obj) { - if (!is_python_sequence(obj.ptr())) { - throw py::type_error("Input data must be a list or tuple"); - } - if (seq.size() == 0) { - throw py::value_error("Input data cannot be empty"); - } - - PyObject *first = seq.item(0); - if (is_python_scalar(first)) { - ndim = 1; - shape = Shape{static_cast(seq.size())}; - return; - } - - if (!is_python_sequence(first)) { - throw py::type_error("List elements must be scalars or nested lists"); - } - if (SequenceView::size_of(first) == 0) { - throw py::value_error("Nested list cannot be empty"); - } - if (!is_python_scalar(SequenceView::item_at(first, 0))) { - throw py::value_error("Only 1D and 2D lists are supported"); - } - - ndim = 2; - const Size rows = static_cast(seq.size()); - const Size cols = static_cast(SequenceView::size_of(first)); - list_of_lists = seq.is_list() && PyList_Check(first); - - for (Py_ssize_t i = 1; i < seq.size(); ++i) { - PyObject *row_obj = seq.item(i); - if (!is_python_sequence(row_obj)) { - throw py::value_error("Input must be a regular 2D list with equal row lengths"); - } - if (static_cast(SequenceView::size_of(row_obj)) != cols) { - throw py::value_error("Input must be a regular 2D list with equal row lengths"); - } - if (list_of_lists && !PyList_Check(row_obj)) { - list_of_lists = false; - } - } - - shape = Shape{rows, cols}; - } -}; - -template -void for_each_scalar(const ListLayout &layout, Fn &&fn) { - const Py_ssize_t rows = layout.seq.size(); - PyObject **rows_items = layout.seq.items(); - - if (layout.ndim == 1) { - for (Py_ssize_t i = 0; i < rows; ++i) { - fn(rows_items[i]); - } - return; - } - - const Py_ssize_t cols = static_cast(layout.shape[1]); - if (layout.list_of_lists) { - for (Py_ssize_t i = 0; i < rows; ++i) { - PyObject **row_items = PySequence_Fast_ITEMS(rows_items[i]); - for (Py_ssize_t j = 0; j < cols; ++j) { - fn(row_items[j]); - } - } - return; - } - - for (Py_ssize_t i = 0; i < rows; ++i) { - PyObject *row = rows_items[i]; - for (Py_ssize_t j = 0; j < cols; ++j) { - fn(SequenceView::item_at(row, j)); - } - } -} - -// read_pylong_fast inlines compact int reads on Python >= 3.12 via -// PyUnstable_Long_*; falls back to PyLong_AsLongLong otherwise. -inline int64_t read_pylong_fast(PyObject *obj) { -#if PY_VERSION_HEX >= 0x030C0000 - PyLongObject *lo = reinterpret_cast(obj); - if (PyUnstable_Long_IsCompact(lo)) { - return static_cast(PyUnstable_Long_CompactValue(lo)); - } -#endif - const int64_t value = PyLong_AsLongLong(obj); - if (value == -1 && PyErr_Occurred()) { - throw py::error_already_set(); - } - return value; -} - -// Convert int64 -> Dst with range check for narrow ints (i8/i16/i32, -// u8/u16/u32) and uint64 (forbids negatives); bool/i64/float/double truncate. -template -inline Dst narrow_int_to(int64_t v) { - if constexpr (std::is_floating_point_v) { - return static_cast(v); - } else if constexpr (std::is_same_v) { - return v != 0; - } else if constexpr (std::is_same_v) { - return v; - } else if constexpr (std::is_same_v) { - if (v < 0) { - throw std::overflow_error("Integer value out of range for target dtype"); - } - return static_cast(v); - } else if constexpr (std::is_unsigned_v) { - if (v < 0 || static_cast(v) > std::numeric_limits::max()) { - throw std::overflow_error("Integer value out of range for target dtype"); - } - return static_cast(v); - } else { - if (v < std::numeric_limits::min() || v > std::numeric_limits::max()) { - throw std::overflow_error("Integer value out of range for target dtype"); - } - return static_cast(v); - } -} - -// Write path: branch order is PyLong/PyFloat exact first (most inputs), then -// bool (int subclass), then long/float subclass fallback. F16/BF16 share -// write_dtype_half via float pivot; other dtypes share write_dtype_native. -template -void write_dtype_native(const ListLayout &layout, std::byte *dest) { - Dst *out = reinterpret_cast(dest); - for_each_scalar(layout, [&](PyObject *obj) { - if (PyLong_CheckExact(obj)) { - *out++ = narrow_int_to(read_pylong_fast(obj)); - } else if (PyFloat_CheckExact(obj)) { - *out++ = static_cast(PyFloat_AS_DOUBLE(obj)); - } else if (PyBool_Check(obj)) { - *out++ = static_cast(obj == Py_True); - } else if (PyLong_Check(obj)) { - *out++ = narrow_int_to(read_pylong_fast(obj)); - } else if (PyFloat_Check(obj)) { - *out++ = static_cast(PyFloat_AS_DOUBLE(obj)); - } else { - throw py::type_error(kElementTypeError); - } - }); -} - -template -void write_dtype_half(const ListLayout &layout, std::byte *dest) { - Dst *out = reinterpret_cast(dest); - for_each_scalar(layout, [&](PyObject *obj) { - float f; - if (PyLong_CheckExact(obj)) { - f = static_cast(read_pylong_fast(obj)); - } else if (PyFloat_CheckExact(obj)) { - f = static_cast(PyFloat_AS_DOUBLE(obj)); - } else if (PyBool_Check(obj)) { - f = (obj == Py_True) ? 1.0f : 0.0f; - } else if (PyLong_Check(obj)) { - f = static_cast(read_pylong_fast(obj)); - } else if (PyFloat_Check(obj)) { - f = static_cast(PyFloat_AS_DOUBLE(obj)); - } else { - throw py::type_error(kElementTypeError); - } - *out++ = utils::cast(f); - }); -} - -void write_with_dtype(const ListLayout &layout, DataType dtype, std::byte *dest) { - // Case order must match dtype.hpp DataType enum exactly: - // BYTE, BOOL, I8, I16, I32, I64, U8, U16, U32, U64, - // F8, F16, F32, F64, C16, C32, C64, C128, BF16 - switch (dtype) { - case DataType::BYTE: - throw py::type_error( - std::string("Unsupported dtype for from_list: ") + toString(dtype)); - case DataType::BOOL: - return write_dtype_native(layout, dest); - case DataType::I8: - return write_dtype_native(layout, dest); - case DataType::I16: - return write_dtype_native(layout, dest); - case DataType::I32: - return write_dtype_native(layout, dest); - case DataType::I64: - return write_dtype_native(layout, dest); - case DataType::U8: - return write_dtype_native(layout, dest); - case DataType::U16: - return write_dtype_native(layout, dest); - case DataType::U32: - return write_dtype_native(layout, dest); - case DataType::U64: - return write_dtype_native(layout, dest); - case DataType::F8: - throw py::type_error( - std::string("Unsupported dtype for from_list: ") + toString(dtype)); - case DataType::F16: - return write_dtype_half(layout, dest); - case DataType::F32: - return write_dtype_native(layout, dest); - case DataType::F64: - return write_dtype_native(layout, dest); - case DataType::C16: - throw py::type_error( - std::string("Unsupported dtype for from_list: ") + toString(dtype)); - case DataType::C32: - throw py::type_error( - std::string("Unsupported dtype for from_list: ") + toString(dtype)); - case DataType::C64: - throw py::type_error( - std::string("Unsupported dtype for from_list: ") + toString(dtype)); - case DataType::C128: - throw py::type_error( - std::string("Unsupported dtype for from_list: ") + toString(dtype)); - case DataType::BF16: - return write_dtype_half(layout, dest); - default: - throw py::type_error( - std::string("Unsupported dtype for from_list: ") + toString(dtype)); - } -} - -} // namespace - -// Entry (exported for pybind registration in tensor.hpp). -Tensor from_list_py(py::handle data, DataType dtype) { - const ListLayout layout(data); - auto tensor = Tensor::empty(layout.shape, dtype, Device(Device::Type::CPU, 0)); - write_with_dtype(layout, dtype, tensor->data()); - return tensor; -} - -} // namespace infinicore::tensor diff --git a/src/infinicore/pybind11/graph.hpp b/src/infinicore/pybind11/graph.hpp deleted file mode 100644 index d45c9b32c..000000000 --- a/src/infinicore/pybind11/graph.hpp +++ /dev/null @@ -1,17 +0,0 @@ -#pragma once - -#include -#include - -#include "infinicore.hpp" - -namespace py = pybind11; - -namespace infinicore::graph { -inline void bind(py::module_ &m) { - py::class_>(m, "Graph") - .def(py::init<>()) // allow construction - .def("run", &infinicore::graph::Graph::run); -} -} // namespace infinicore::graph diff --git a/src/infinicore/pybind11/infinicore.cc b/src/infinicore/pybind11/infinicore.cc deleted file mode 100644 index 9820eab1b..000000000 --- a/src/infinicore/pybind11/infinicore.cc +++ /dev/null @@ -1,35 +0,0 @@ -#include -#include - -#include "../utils.hpp" -#include "context.hpp" -#include "device.hpp" -#include "device_event.hpp" -#include "dtype.hpp" -#include "graph.hpp" -#include "io.hpp" -#include "ops.hpp" -#include "tensor.hpp" - -#ifdef ENABLE_MUTUAL_AWARENESS -#include "analyzer.hpp" -#endif - -namespace infinicore { - -PYBIND11_MODULE(_infinicore, m) { - context::bind(m); - device::bind(m); - device_event::bind(m); - dtype::bind(m); - ops::bind(m); - tensor::bind(m); - io::bind(m); - graph::bind(m); - -#ifdef ENABLE_MUTUAL_AWARENESS - analyzer::pybind::bind(m); -#endif -} - -} // namespace infinicore diff --git a/src/infinicore/pybind11/io.hpp b/src/infinicore/pybind11/io.hpp deleted file mode 100644 index fc87049ef..000000000 --- a/src/infinicore/pybind11/io.hpp +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -#include "infinicore.hpp" -#include -#include - -namespace py = pybind11; - -namespace infinicore::io { - -inline void bind(py::module &m) { - m.def( - "set_printoptions", [](int precision, int threshold, int edge_items, int line_width, py::object sci_mode) { - infinicore::print_options::set_precision(precision); - infinicore::print_options::set_threshold(threshold); - infinicore::print_options::set_edge_items(edge_items); - infinicore::print_options::set_line_width(line_width); - - // Handle sci_mode: None -> -1 (auto), True -> 1 (enable), False -> 0 (disable) - int sci_mode_value = -1; // default: auto - if (!sci_mode.is_none()) { - sci_mode_value = static_cast(py::cast(sci_mode)); // True -> 1, False -> 0 - } - - infinicore::print_options::set_sci_mode(sci_mode_value); }, pybind11::arg("precision"), pybind11::arg("threshold"), pybind11::arg("edge_items"), pybind11::arg("line_width"), pybind11::arg("sci_mode")); -} - -} // namespace infinicore::io diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp deleted file mode 100644 index 087261382..000000000 --- a/src/infinicore/pybind11/ops.hpp +++ /dev/null @@ -1,275 +0,0 @@ -#pragma once - -#include - -#include "ops/acos.hpp" -#include "ops/adaptive_avg_pool1d.hpp" -#include "ops/adaptive_avg_pool3d.hpp" -#include "ops/adaptive_max_pool1d.hpp" -#include "ops/add.hpp" -#include "ops/add_rms_norm.hpp" -#include "ops/addbmm.hpp" -#include "ops/addcmul.hpp" -#include "ops/addr.hpp" -#include "ops/affine_grid.hpp" -#include "ops/all.hpp" -#include "ops/argwhere.hpp" -#include "ops/asin.hpp" -#include "ops/asinh.hpp" -#include "ops/asum.hpp" -#include "ops/atanh.hpp" -#include "ops/attention.hpp" -#include "ops/avg_pool1d.hpp" -#include "ops/axpy.hpp" -#include "ops/baddbmm.hpp" -#include "ops/bilinear.hpp" -#include "ops/binary_cross_entropy_with_logits.hpp" -#include "ops/bitwise_right_shift.hpp" -#include "ops/blas_amax.hpp" -#include "ops/blas_amin.hpp" -#include "ops/blas_copy.hpp" -#include "ops/blas_dot.hpp" -#include "ops/block_diag.hpp" -#include "ops/broadcast_to.hpp" -#include "ops/cat.hpp" -#include "ops/causal_conv1d.hpp" -#include "ops/causal_softmax.hpp" -#include "ops/cdist.hpp" -#include "ops/chunk_gated_delta_rule.hpp" -#include "ops/conv2d.hpp" -#include "ops/cross_entropy.hpp" -#include "ops/diff.hpp" -#include "ops/digamma.hpp" -#include "ops/dist.hpp" -#include "ops/embedding.hpp" -#include "ops/equal.hpp" -#include "ops/flash_attention.hpp" -#include "ops/flipud.hpp" -#include "ops/float_power.hpp" -#include "ops/floor.hpp" -#include "ops/floor_divide.hpp" -#include "ops/fmin.hpp" -#include "ops/fmod.hpp" -#include "ops/fused_gated_delta_net_gating.hpp" -#include "ops/gaussian_nll_loss.hpp" -#include "ops/hardswish.hpp" -#include "ops/hardtanh.hpp" -#include "ops/hinge_embedding_loss.hpp" -#include "ops/huber_loss.hpp" -#include "ops/hypot.hpp" -#include "ops/index_add.hpp" -#include "ops/index_copy.hpp" -#include "ops/inner.hpp" -#include "ops/interpolate.hpp" -#include "ops/kron.hpp" -#include "ops/kthvalue.hpp" -#include "ops/kv_caching.hpp" -#include "ops/layer_norm.hpp" -#include "ops/ldexp.hpp" -#include "ops/lerp.hpp" -#include "ops/linear.hpp" -#include "ops/linear_w8a8i8.hpp" -#include "ops/log_softmax.hpp" -#include "ops/logaddexp.hpp" -#include "ops/logaddexp2.hpp" -#include "ops/logcumsumexp.hpp" -#include "ops/logdet.hpp" -#include "ops/logical_and.hpp" -#include "ops/logical_not.hpp" -#include "ops/mamba_selective_scan.hpp" -#include "ops/masked_select.hpp" -#include "ops/matmul.hpp" -#include "ops/mha.hpp" -#include "ops/mha_kvcache.hpp" -#include "ops/mha_varlen.hpp" -#include "ops/moe_topk_softmax.hpp" -#include "ops/mrope.hpp" -#include "ops/mul.hpp" -#include "ops/mul_scalar.hpp" -#include "ops/multi_margin_loss.hpp" -#include "ops/nrm2.hpp" -#include "ops/pad.hpp" -#include "ops/paged_attention.hpp" -#include "ops/paged_attention_prefill.hpp" -#include "ops/paged_caching.hpp" -#include "ops/prelu.hpp" -#include "ops/random_sample.hpp" -#include "ops/rearrange.hpp" -#include "ops/reciprocal.hpp" -#include "ops/recurrent_gated_delta_rule.hpp" -#include "ops/relu6.hpp" -#include "ops/rms_norm.hpp" -#include "ops/rope.hpp" -#include "ops/rot.hpp" -#include "ops/rotg.hpp" -#include "ops/rotm.hpp" -#include "ops/rotmg.hpp" -#include "ops/rwkv5_wkv.hpp" -#include "ops/scal.hpp" -#include "ops/scatter.hpp" -#include "ops/selu.hpp" -#include "ops/sigmoid.hpp" -#include "ops/silu.hpp" -#include "ops/silu_and_mul.hpp" -#include "ops/sinh.hpp" -#include "ops/smooth_l1_loss.hpp" -#include "ops/softplus.hpp" -#include "ops/softsign.hpp" -#include "ops/sum.hpp" -#include "ops/swap.hpp" -#include "ops/swiglu.hpp" -#include "ops/take.hpp" -#include "ops/tan.hpp" -#include "ops/tanhshrink.hpp" -#include "ops/topk.hpp" -#include "ops/topksoftmax.hpp" -#include "ops/triplet_margin_loss.hpp" -#include "ops/triplet_margin_with_distance_loss.hpp" -#include "ops/unfold.hpp" -#include "ops/upsample_bilinear.hpp" -#include "ops/upsample_nearest.hpp" -#include "ops/vander.hpp" -#include "ops/var.hpp" -#include "ops/var_mean.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind(py::module &m) { - bind_adaptive_max_pool1d(m); - bind_add(m); - bind_adaptive_avg_pool3d(m); - bind_argwhere(m); - bind_addr(m); - bind_asin(m); - bind_add_rms_norm(m); - bind_add(m); - bind_addbmm(m); - bind_acos(m); - bind_affine_grid(m); - bind_floor(m); - bind_adaptive_avg_pool1d(m); - bind_attention(m); - bind_asinh(m); - bind_asum(m); - bind_axpy(m); - bind_baddbmm(m); - bind_bilinear(m); - bind_blas_amax(m); - bind_blas_amin(m); - bind_blas_copy(m); - bind_blas_dot(m); - bind_block_diag(m); - bind_bitwise_right_shift(m); - bind_causal_conv1d(m); - bind_causal_softmax(m); - bind_diff(m); - bind_digamma(m); - bind_dist(m); - bind_flash_attention(m); - bind_hinge_embedding_loss(m); - bind_kv_caching(m); - bind_fmod(m); - bind_fused_gated_delta_net_gating(m); - bind_fmin(m); - bind_cat(m); - bind_causal_softmax(m); - bind_inner(m); - bind_random_sample(m); - bind_masked_select(m); - bind_log_softmax(m); - bind_logaddexp(m); - bind_logaddexp2(m); - bind_linear(m); - bind_logdet(m); - bind_matmul(m); - bind_mamba_selective_scan(m); - bind_kron(m); - bind_mul(m); - bind_mul_scalar(m); - bind_nrm2(m); - bind_mha_kvcache(m); - bind_mha_varlen(m); - bind_mha(m); - bind_moe_topk_softmax(m); - bind_mrope(m); - bind_hardswish(m); - bind_hardtanh(m); - bind_gaussian_nll_loss(m); - bind_interpolate(m); - bind_paged_attention(m); - bind_paged_attention_prefill(m); - bind_paged_caching(m); - bind_pad(m); - bind_prelu(m); - bind_random_sample(m); - bind_cross_entropy(m); - bind_conv2d(m); - bind_hypot(m); - bind_take(m); - bind_index_copy(m); - bind_index_add(m); - bind_smooth_l1_loss(m); - bind_rearrange(m); - bind_relu6(m); - bind_rms_norm(m); - bind_avg_pool1d(m); - bind_sigmoid(m); - bind_silu(m); - bind_swiglu(m); - bind_tan(m); - bind_tanhshrink(m); - bind_logcumsumexp(m); - bind_logical_and(m); - bind_logical_not(m); - bind_vander(m); - bind_unfold(m); - bind_rope(m); - bind_rot(m); - bind_rotg(m); - bind_rotm(m); - bind_rotmg(m); - bind_rwkv5_wkv(m); - bind_floor_divide(m); - bind_float_power(m); - bind_flipud(m); - bind_multi_margin_loss(m); - bind_scatter(m); - bind_scal(m); - bind_broadcast_to(m); - bind_softplus(m); - bind_softsign(m); - bind_linear(m); - bind_huber_loss(m); - bind_triplet_margin_with_distance_loss(m); - bind_upsample_nearest(m); - bind_embedding(m); - bind_linear_w8a8i8(m); - bind_silu_and_mul(m); - bind_sum(m); - bind_var_mean(m); - bind_var(m); - bind_topk(m); - bind_all(m); - bind_equal(m); - bind_atanh(m); - bind_addcmul(m); - bind_cdist(m); - bind_binary_cross_entropy_with_logits(m); - bind_chunk_gated_delta_rule(m); - bind_reciprocal(m); - bind_recurrent_gated_delta_rule(m); - bind_upsample_bilinear(m); - bind_kthvalue(m); - bind_ldexp(m); - bind_lerp(m); - bind_triplet_margin_loss(m); - bind_selu(m); - bind_swap(m); - bind_sinh(m); - bind_layer_norm(m); - bind_topksoftmax(m); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/acos.hpp b/src/infinicore/pybind11/ops/acos.hpp deleted file mode 100644 index 4c4f21db3..000000000 --- a/src/infinicore/pybind11/ops/acos.hpp +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -#include "infinicore/ops/acos.hpp" // 引用核心算子头文件 -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_acos(py::module &m) { - // 绑定 out-of-place 接口: output = acos(input) - m.def("acos", - &op::acos, - py::arg("input"), - R"doc(Computes the inverse cosine (arccosine) of each element of input. - -Returns a new tensor with the arccosine of the elements of input. -The range of the result is [0, pi].)doc"); - - // 绑定 in-place 接口: acos_(output, input) - m.def("acos_", - &op::acos_, - py::arg("output"), - py::arg("input"), - R"doc(In-place acos operation. Writes result into output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/adaptive_avg_pool1d.hpp b/src/infinicore/pybind11/ops/adaptive_avg_pool1d.hpp deleted file mode 100644 index 8d58232f6..000000000 --- a/src/infinicore/pybind11/ops/adaptive_avg_pool1d.hpp +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/adaptive_avg_pool1d.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_adaptive_avg_pool1d(py::module &m) { - // 绑定函数接口: output = adaptive_avg_pool1d(input, output_size) - m.def("adaptive_avg_pool1d", - &op::adaptive_avg_pool1d, - py::arg("input"), - py::arg("output_size"), - R"doc(Applies a 1D adaptive average pooling over an input signal composed of several input planes. - -Args: - input (Tensor): Input tensor of shape (C, L) or (N, C, L). - output_size (int): The target output size. - -Returns: - Tensor: Output tensor of shape (C, output_size) or (N, C, output_size). -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/adaptive_avg_pool3d.hpp b/src/infinicore/pybind11/ops/adaptive_avg_pool3d.hpp deleted file mode 100644 index 758baf232..000000000 --- a/src/infinicore/pybind11/ops/adaptive_avg_pool3d.hpp +++ /dev/null @@ -1,22 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/adaptive_avg_pool3d.hpp" - -namespace py = pybind11; -namespace infinicore::ops { -inline void bind_adaptive_avg_pool3d(py::module &m) { - m.def("adaptive_avg_pool3d", - &op::adaptive_avg_pool3d, - py::arg("x"), - py::arg("output_size"), - R"doc( Adaptive Average Pooling 3D.)doc"); - - m.def("adaptive_avg_pool3d_", - &op::adaptive_avg_pool3d_, - py::arg("y"), - py::arg("x"), - R"doc(In-place, Adaptive Average Pooling 3D.)doc"); -} -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/adaptive_max_pool1d.hpp b/src/infinicore/pybind11/ops/adaptive_max_pool1d.hpp deleted file mode 100644 index 99dc3bbc5..000000000 --- a/src/infinicore/pybind11/ops/adaptive_max_pool1d.hpp +++ /dev/null @@ -1,39 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/adaptive_max_pool1d.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_adaptive_max_pool1d(py::module &m) { - m.def("adaptive_max_pool1d", - &op::adaptive_max_pool1d, - py::arg("x"), - py::arg("output_size"), - R"doc(1D Adaptive Max Pooling. - -Args: - x: Input tensor of shape (N, C, L_in) or (N, L_in) - output_size: Target output size L_out -Returns: - Output tensor of shape (N, C, L_out) or (N, L_out) -)doc"); - - m.def("adaptive_max_pool1d_", - &op::adaptive_max_pool1d_, - py::arg("y"), - py::arg("x"), - py::arg("output_size"), - R"doc(In-place 1D Adaptive Max Pooling. - -Args: - y: Output tensor of shape (N, C, L_out) or (N, L_out) - x: Input tensor of shape (N, C, L_in) or (N, L_in) - output_size: Target output size L_out -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/add.hpp b/src/infinicore/pybind11/ops/add.hpp deleted file mode 100644 index fe7ac852f..000000000 --- a/src/infinicore/pybind11/ops/add.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/add.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_add(py::module &m) { - m.def("add", - &op::add, - py::arg("a"), - py::arg("b"), - R"doc(Addition of two tensors.)doc"); - - m.def("add_", - &op::add_, - py::arg("c"), - py::arg("a"), - py::arg("b"), - R"doc(In-place tensor addition.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/add_rms_norm.hpp b/src/infinicore/pybind11/ops/add_rms_norm.hpp deleted file mode 100644 index 5f9b243e5..000000000 --- a/src/infinicore/pybind11/ops/add_rms_norm.hpp +++ /dev/null @@ -1,51 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/add_rms_norm.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_add_rms_norm(py::module &m) { - m.def("add_rms_norm", - &op::add_rms_norm, - py::arg("a"), - py::arg("b"), - py::arg("weight"), - py::arg("epsilon") = 1e-5f, - R"doc(Fused Add and RMS Normalization. - -Args: - a: First input tensor - b: Second input tensor - weight: Scale weights - epsilon: Small constant for numerical stability, default is 1e-5 - -Returns: - Tuple of (normalized_result, add_result): (RMSNorm(a + b) * weight, a + b) - The add_result can be used as residual for subsequent layers. -)doc"); - - m.def("add_rms_norm_", - &op::add_rms_norm_, - py::arg("y"), - py::arg("residual_out"), - py::arg("a"), - py::arg("b"), - py::arg("weight"), - py::arg("epsilon") = 1e-5f, - R"doc(In-place Fused Add and RMS Normalization. - -Args: - y: Output tensor for normalized result - residual_out: Output tensor for add result (a + b) before normalization - a: First input tensor - b: Second input tensor - weight: Scale weights - epsilon: Small constant for numerical stability, default is 1e-5 -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/addbmm.hpp b/src/infinicore/pybind11/ops/addbmm.hpp deleted file mode 100644 index 3756d9e59..000000000 --- a/src/infinicore/pybind11/ops/addbmm.hpp +++ /dev/null @@ -1,52 +0,0 @@ -#pragma once - -#include "infinicore/ops/addbmm.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_addbmm(py::module &m) { - // ----------------------------------------------------------- - // 1. Out-of-place 接口: output = addbmm(...) - // ----------------------------------------------------------- - m.def("addbmm", - &op::addbmm, - py::arg("input"), - py::arg("batch1"), - py::arg("batch2"), - py::arg("beta") = 1.0f, - py::arg("alpha") = 1.0f, - R"doc(Performs a batch matrix-matrix product of matrices stored in batch1 and batch2, -with a reduced add step (summing over all matrices in the batch). - -.. math:: - \text{out} = \beta \times \text{input} + \alpha \times \sum_{i=0}^{b-1} (\text{batch1}_i \mathbin{@} \text{batch2}_i) - -Args: - input (Tensor): Matrix to be added. Shape (n, p). - batch1 (Tensor): The first batch of matrices to be multiplied. Shape (b, n, m). - batch2 (Tensor): The second batch of matrices to be multiplied. Shape (b, m, p). - beta (float, optional): Multiplier for input. Default: 1.0. - alpha (float, optional): Multiplier for batch1 @ batch2. Default: 1.0. - -Returns: - Tensor: Output tensor of shape (n, p). -)doc"); - - // ----------------------------------------------------------- - // 2. [新增] In-place 接口: addbmm_(out, ...) - // ----------------------------------------------------------- - m.def("addbmm_", - &op::addbmm_, // 绑定到 C++ 的 void addbmm_(...) - py::arg("out"), // 第一个参数通常是输出 Tensor - py::arg("input"), - py::arg("batch1"), - py::arg("batch2"), - py::arg("beta") = 1.0f, - py::arg("alpha") = 1.0f, - "In-place version of addbmm"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/addcmul.hpp b/src/infinicore/pybind11/ops/addcmul.hpp deleted file mode 100644 index 7301ce27a..000000000 --- a/src/infinicore/pybind11/ops/addcmul.hpp +++ /dev/null @@ -1,51 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/addcmul.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_addcmul(py::module &m) { - // 绑定 out-of-place 接口: out = addcmul(input, t1, t2, value) - m.def("addcmul", - &op::addcmul, - py::arg("input"), - py::arg("tensor1"), - py::arg("tensor2"), - py::arg("value") = 1.0f, - R"doc(Performs the element-wise multiplication of tensor1 by tensor2, -multiplies the result by value and adds it to input. - -Args: - input: Tensor to be added - tensor1: First tensor for multiplication - tensor2: Second tensor for multiplication - value: Scalar multiplier for tensor1 * tensor2 (default: 1.0) - -Returns: - The output tensor -)doc"); - - // 绑定 in-place / specified output 接口: addcmul_(out, input, t1, t2, value) - m.def("addcmul_", - &op::addcmul_, - py::arg("out"), - py::arg("input"), - py::arg("tensor1"), - py::arg("tensor2"), - py::arg("value") = 1.0f, - R"doc(In-place version of addcmul. - -Args: - out: The destination tensor to store the result - input: Tensor to be added - tensor1: First tensor for multiplication - tensor2: Second tensor for multiplication - value: Scalar multiplier for tensor1 * tensor2 (default: 1.0) -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/addr.hpp b/src/infinicore/pybind11/ops/addr.hpp deleted file mode 100644 index 912b2b5da..000000000 --- a/src/infinicore/pybind11/ops/addr.hpp +++ /dev/null @@ -1,29 +0,0 @@ -#pragma once - -#include "infinicore/ops/addr.hpp" -#include - -namespace py = pybind11; -namespace infinicore::ops { -inline void bind_addr(py::module_ &m) { - m.def( - "addr", - &op::addr, - py::arg("input"), - py::arg("vec1"), - py::arg("vec2"), - py::arg("alpha"), - py::arg("beta"), - R"doc(Addr.)doc"); - m.def( - "addr_", - &op::addr_, - py::arg("out"), - py::arg("input"), - py::arg("vec1"), - py::arg("vec2"), - py::arg("beta"), - py::arg("alpha"), - R"doc(Addr.)doc"); -} -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/affine_grid.hpp b/src/infinicore/pybind11/ops/affine_grid.hpp deleted file mode 100644 index c12279bcc..000000000 --- a/src/infinicore/pybind11/ops/affine_grid.hpp +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once - -#include "infinicore/ops/affine_grid.hpp" -#include -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_affine_grid(py::module &m) { - // 绑定函数接口: grid = affine_grid(theta, size, align_corners) - m.def("affine_grid", - &op::affine_grid, - py::arg("theta"), - py::arg("size"), - py::arg("align_corners") = false, // 设置默认值 - R"doc(Generates a 2D or 3D flow field (sampling grid), given a batch of affine matrices theta. - -Args: - theta (Tensor): Input affine matrices of shape (N, 2, 3) for 2D or (N, 3, 4) for 3D. - size (List[int]): The target output image size. Usually (N, C, H, W) for 2D or (N, C, D, H, W) for 3D. - align_corners (bool, optional): Geometrically, we consider the pixels of the input as squares rather than points. If set to True, the extrema (-1 and 1) are considered as referring to the center points of the input's corner pixels. If set to False, they are instead considered as referring to the corner points of the input's corner pixels, making the sampling more resolution agnostic. Defaults to False. - -Returns: - Tensor: Output tensor of shape (N, H, W, 2). -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/all.hpp b/src/infinicore/pybind11/ops/all.hpp deleted file mode 100644 index 4ccac685b..000000000 --- a/src/infinicore/pybind11/ops/all.hpp +++ /dev/null @@ -1,60 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/all.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_all(Tensor input, py::object dim, bool keepdim) { - if (dim.is_none()) { - std::vector dim_vec; - for (int i = 0; i < input->shape().size(); i++) { - dim_vec.push_back(i); - } - return op::all(input, dim_vec, keepdim); - } else if (py::isinstance(dim) || py::isinstance(dim)) { - return op::all(input, dim.cast>(), keepdim); - } else if (py::isinstance(dim)) { - return op::all(input, std::vector(1, dim.cast()), keepdim); - } else { - throw std::invalid_argument("dim must be a tuple or an integer"); - } -} - -void py_all_(Tensor output, Tensor input, py::object dim, bool keepdim) { - if (dim.is_none()) { - std::vector dim_vec; - for (int i = 0; i < input->shape().size(); i++) { - dim_vec.push_back(i); - } - op::all_(output, input, dim_vec, keepdim); - } else if (py::isinstance(dim) || py::isinstance(dim)) { - op::all_(output, input, dim.cast>(), keepdim); - } else if (py::isinstance(dim)) { - op::all_(output, input, std::vector(1, dim.cast()), keepdim); - } else { - throw std::invalid_argument("dim must be a tuple or an integer"); - } -} - -inline void bind_all(py::module &m) { - m.def("all", - &py_all, - py::arg("input"), - py::arg("dim"), - py::arg("keepdim"), - R"doc(All of input tensor along the given dimensions.)doc"); - - m.def("all_", - &py_all_, - py::arg("output"), - py::arg("input"), - py::arg("dim"), - py::arg("keepdim"), - R"doc(In-place tensor all.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/argwhere.hpp b/src/infinicore/pybind11/ops/argwhere.hpp deleted file mode 100644 index cb91142c8..000000000 --- a/src/infinicore/pybind11/ops/argwhere.hpp +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "infinicore/ops/argwhere.hpp" -#include -namespace py = pybind11; -namespace infinicore::ops { -inline void bind_argwhere(py::module &m) { - m.def("argwhere", - &op::argwhere, - py::arg("x"), - R"doc(Argwhere.)doc"); -} -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/asin.hpp b/src/infinicore/pybind11/ops/asin.hpp deleted file mode 100644 index ebd14adbc..000000000 --- a/src/infinicore/pybind11/ops/asin.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/asin.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_asin(py::module &m) { - m.def("asin", - &op::asin, - py::arg("input"), - R"doc(Arcsin activation function.)doc"); - - m.def("asin_", - &op::asin_, - py::arg("output"), - py::arg("input"), - R"doc(In-place arcsin activation function.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/asinh.hpp b/src/infinicore/pybind11/ops/asinh.hpp deleted file mode 100644 index bf1fcca23..000000000 --- a/src/infinicore/pybind11/ops/asinh.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/asinh.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_asinh(py::module &m) { - m.def("asinh", - &op::asinh, - py::arg("x"), - R"doc(Element-wise inverse hyperbolic sine function.)doc"); - - m.def("asinh_", - &op::asinh_, - py::arg("y"), - py::arg("x"), - R"doc(In-place element-wise inverse hyperbolic sine function.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/asum.hpp b/src/infinicore/pybind11/ops/asum.hpp deleted file mode 100644 index b094d12f5..000000000 --- a/src/infinicore/pybind11/ops/asum.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/asum.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_asum(py::module &m) { - m.def("asum", - &op::asum, - py::arg("x"), - R"doc(BLAS level-1 asum.)doc"); - - m.def("asum_", - &op::asum_, - py::arg("x"), - py::arg("result"), - R"doc(In-place BLAS level-1 asum.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/atanh.hpp b/src/infinicore/pybind11/ops/atanh.hpp deleted file mode 100644 index 44fd74169..000000000 --- a/src/infinicore/pybind11/ops/atanh.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/atanh.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_atanh(py::module &m) { - m.def("atanh", - &op::atanh, - py::arg("a"), - R"doc(Inverse hyperbolic tangent of a tensor.)doc"); - - m.def("atanh_", - &op::atanh_, - py::arg("y"), - py::arg("a"), - R"doc(Compute inverse hyperbolic tangent and store in the provided output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/attention.hpp b/src/infinicore/pybind11/ops/attention.hpp deleted file mode 100644 index 4af2d5f74..000000000 --- a/src/infinicore/pybind11/ops/attention.hpp +++ /dev/null @@ -1,56 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/attention.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_attention(py::module &m) { - m.def("attention", - &op::attention, - py::arg("q"), - py::arg("k"), - py::arg("v"), - py::arg("k_cache"), - py::arg("v_cache"), - py::arg("pos"), - R"doc(Attention mechanism with KV caching. - -Args: - q: Query tensor - k: Key tensor - v: Value tensor - k_cache: Key cache tensor - v_cache: Value cache tensor - pos: Current position in the sequence - -Returns: - Output tensor from attention computation -)doc"); - - m.def("attention_", - &op::Attention::execute, - py::arg("out"), - py::arg("q"), - py::arg("k"), - py::arg("v"), - py::arg("k_cache"), - py::arg("v_cache"), - py::arg("pos"), - R"doc(In-place attention mechanism with KV caching. - -Args: - out: Output tensor - q: Query tensor - k: Key tensor - v: Value tensor - k_cache: Key cache tensor - v_cache: Value cache tensor - pos: Current position in the sequence -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/avg_pool1d.hpp b/src/infinicore/pybind11/ops/avg_pool1d.hpp deleted file mode 100644 index 32394552a..000000000 --- a/src/infinicore/pybind11/ops/avg_pool1d.hpp +++ /dev/null @@ -1,37 +0,0 @@ -#pragma once - -#include -#include - -#include "infinicore/ops/avg_pool1d.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_avg_pool1d(py::module &m) { - m.def( - "avg_pool1d", - [](::infinicore::Tensor input, size_t kernel_size, std::optional stride, size_t padding) { - return op::avg_pool1d(input, kernel_size, stride.value_or(0), padding); - }, - py::arg("input"), - py::arg("kernel_size"), - py::arg("stride") = py::none(), - py::arg("padding") = 0, - R"doc(AvgPool1d out-of-place.)doc"); - - m.def( - "avg_pool1d_", - [](::infinicore::Tensor output, ::infinicore::Tensor input, size_t kernel_size, std::optional stride, size_t padding) { - op::avg_pool1d_(output, input, kernel_size, stride.value_or(0), padding); - }, - py::arg("output"), - py::arg("input"), - py::arg("kernel_size"), - py::arg("stride") = py::none(), - py::arg("padding") = 0, - R"doc(AvgPool1d in-place variant writing to provided output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/axpy.hpp b/src/infinicore/pybind11/ops/axpy.hpp deleted file mode 100644 index fd0aff633..000000000 --- a/src/infinicore/pybind11/ops/axpy.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/axpy.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_axpy(py::module &m) { - m.def("axpy_", - &op::axpy_, - py::arg("alpha"), - py::arg("x"), - py::arg("y"), - R"doc(In-place BLAS level-1 axpy, updating y.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/baddbmm.hpp b/src/infinicore/pybind11/ops/baddbmm.hpp deleted file mode 100644 index bb5709afc..000000000 --- a/src/infinicore/pybind11/ops/baddbmm.hpp +++ /dev/null @@ -1,56 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/baddbmm.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_baddbmm(Tensor input, Tensor batch1, Tensor batch2, float beta = 1.0f, float alpha = 1.0f) { - return op::baddbmm(input, batch1, batch2, beta, alpha); -} - -void py_baddbmm_(Tensor out, Tensor input, Tensor batch1, Tensor batch2, float beta = 1.0f, float alpha = 1.0f) { - op::baddbmm_(out, input, batch1, batch2, beta, alpha); -} - -inline void bind_baddbmm(py::module &m) { - m.def("baddbmm", - &py_baddbmm, - py::arg("input"), - py::arg("batch1"), - py::arg("batch2"), - py::arg("beta") = 1.0f, - py::arg("alpha") = 1.0f, - R"doc(Batched matrix-matrix product with addition. -Args: - input: Input tensor - batch1: First batch of matrices - batch2: Second batch of matrices - beta: Scaling factor for input tensor - alpha: Scaling factor for the product of batch1 and batch2 -Returns: - Output tensor after baddbmm operation -)doc"); - m.def("baddbmm_", - &py_baddbmm_, - py::arg("out"), - py::arg("input"), - py::arg("batch1"), - py::arg("batch2"), - py::arg("beta") = 1.0f, - py::arg("alpha") = 1.0f, - R"doc(In-place batched matrix-matrix product with addition. -Args: - out: Output tensor - input: Input tensor - batch1: First batch of matrices - batch2: Second batch of matrices - beta: Scaling factor for input tensor - alpha: Scaling factor for the product of batch1 and batch2 -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/bilinear.hpp b/src/infinicore/pybind11/ops/bilinear.hpp deleted file mode 100644 index 2288f3602..000000000 --- a/src/infinicore/pybind11/ops/bilinear.hpp +++ /dev/null @@ -1,61 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/bilinear.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_bilinear(Tensor x1, Tensor x2, Tensor weight, pybind11::object bias) { - std::optional bias_tensor = std::nullopt; - if (!bias.is_none()) { - bias_tensor = bias.cast(); - } - return op::bilinear(x1, x2, weight, bias_tensor); -} - -void py_bilinear_(Tensor out, Tensor x1, Tensor x2, Tensor weight, pybind11::object bias) { - std::optional bias_tensor = std::nullopt; - if (!bias.is_none()) { - bias_tensor = bias.cast(); - } - op::bilinear_(out, x1, x2, weight, bias_tensor); -} - -inline void bind_bilinear(py::module &m) { - m.def("bilinear", - &py_bilinear, - py::arg("x1"), - py::arg("x2"), - py::arg("weight"), - py::arg("bias"), - R"doc(Bilinear transformation of two input tensors. -Args: - x1: First input tensor - x2: Second input tensor - weight: Weight tensor - bias: Bias tensor (optional) -Returns: - Output tensor after bilinear transformation -)doc"); - - m.def("bilinear_", - &py_bilinear_, - py::arg("out"), - py::arg("x1"), - py::arg("x2"), - py::arg("weight"), - py::arg("bias"), - R"doc(In-place bilinear transformation of two input tensors. -Args: - out: Output tensor - x1: First input tensor - x2: Second input tensor - weight: Weight tensor - bias: Bias tensor (optional) -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/binary_cross_entropy_with_logits.hpp b/src/infinicore/pybind11/ops/binary_cross_entropy_with_logits.hpp deleted file mode 100644 index 48943804f..000000000 --- a/src/infinicore/pybind11/ops/binary_cross_entropy_with_logits.hpp +++ /dev/null @@ -1,77 +0,0 @@ -#pragma once - -#include "infinicore/ops/binary_cross_entropy_with_logits.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_binary_cross_entropy_with_logits(py::module &m) { - // 1. 绑定 out-of-place 接口: out = binary_cross_entropy_with_logits(...) - m.def( - "binary_cross_entropy_with_logits", - [](Tensor logits, - Tensor target, - py::object weight, - py::object pos_weight, - std::string reduction) { - Tensor w = weight.is_none() ? Tensor() : weight.cast(); - Tensor pw = pos_weight.is_none() ? Tensor() : pos_weight.cast(); - - return op::binary_cross_entropy_with_logits( - logits, target, w, pw, reduction); - }, - py::arg("input"), - py::arg("target"), - py::arg("weight") = py::none(), - py::arg("pos_weight") = py::none(), - py::arg("reduction") = "mean", - R"doc(Measures Binary Cross Entropy between target and output logits. - - -Args: - input: Tensor of arbitrary shape as unnormalized scores (logits). - target: Tensor of the same shape as input with values between 0 and 1. - weight: Optional rescaling weight for each loss component. - pos_weight: Optional weight for positive examples (must be broadcastable). - reduction: Specfies the reduction to apply: 'none' | 'mean' | 'sum'. - -Returns: - A tensor representing the loss. -)doc"); - - // 2. 绑定指定输出接口: binary_cross_entropy_with_logits_(out, ...) - m.def( - "binary_cross_entropy_with_logits_", - [](Tensor output, - Tensor logits, - Tensor target, - py::object weight, - py::object pos_weight, - std::string reduction) { - Tensor w = weight.is_none() ? Tensor() : weight.cast(); - Tensor pw = pos_weight.is_none() ? Tensor() : pos_weight.cast(); - - return op::binary_cross_entropy_with_logits_( - output, logits, target, w, pw, reduction); - }, - py::arg("out"), - py::arg("input"), - py::arg("target"), - py::arg("weight") = py::none(), - py::arg("pos_weight") = py::none(), - py::arg("reduction") = "mean", - R"doc(Specified output version of binary_cross_entropy_with_logits. - -Args: - out: The destination tensor to store the loss. - input: Logits tensor. - target: Target tensor. - weight: Optional sample weight. - pos_weight: Optional positive class weight. - reduction: Specfies the reduction to apply. -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/bitwise_right_shift.hpp b/src/infinicore/pybind11/ops/bitwise_right_shift.hpp deleted file mode 100644 index a540a33bc..000000000 --- a/src/infinicore/pybind11/ops/bitwise_right_shift.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/bitwise_right_shift.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_bitwise_right_shift(py::module &m) { - m.def("bitwise_right_shift", - &op::bitwise_right_shift, - py::arg("input"), - py::arg("other"), - R"doc(Element-wise bitwise right shift.)doc"); - - m.def("bitwise_right_shift_", - &op::bitwise_right_shift_, - py::arg("out"), - py::arg("input"), - py::arg("other"), - R"doc(In-place element-wise bitwise right shift.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/blas_amax.hpp b/src/infinicore/pybind11/ops/blas_amax.hpp deleted file mode 100644 index 51e8cfe4d..000000000 --- a/src/infinicore/pybind11/ops/blas_amax.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/blas_amax.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_blas_amax(py::module &m) { - m.def("blas_amax", - &op::blas_amax, - py::arg("x"), - R"doc(BLAS level-1 amax.)doc"); - - m.def("blas_amax_", - &op::blas_amax_, - py::arg("x"), - py::arg("result"), - R"doc(In-place BLAS level-1 amax.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/blas_amin.hpp b/src/infinicore/pybind11/ops/blas_amin.hpp deleted file mode 100644 index 8961a9363..000000000 --- a/src/infinicore/pybind11/ops/blas_amin.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/blas_amin.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_blas_amin(py::module &m) { - m.def("blas_amin", - &op::blas_amin, - py::arg("x"), - R"doc(BLAS level-1 amin.)doc"); - - m.def("blas_amin_", - &op::blas_amin_, - py::arg("x"), - py::arg("result"), - R"doc(In-place BLAS level-1 amin.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/blas_copy.hpp b/src/infinicore/pybind11/ops/blas_copy.hpp deleted file mode 100644 index c348ac38b..000000000 --- a/src/infinicore/pybind11/ops/blas_copy.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/blas_copy.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_blas_copy(py::module &m) { - m.def("blas_copy_", - &op::blas_copy_, - py::arg("x"), - py::arg("y"), - R"doc(In-place BLAS level-1 copy from x to y.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/blas_dot.hpp b/src/infinicore/pybind11/ops/blas_dot.hpp deleted file mode 100644 index 73b4f0bc9..000000000 --- a/src/infinicore/pybind11/ops/blas_dot.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/blas_dot.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_blas_dot(py::module &m) { - m.def("blas_dot", - &op::blas_dot, - py::arg("x"), - py::arg("y"), - R"doc(BLAS level-1 dot.)doc"); - - m.def("blas_dot_", - &op::blas_dot_, - py::arg("x"), - py::arg("y"), - py::arg("result"), - R"doc(In-place BLAS level-1 dot.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/block_diag.hpp b/src/infinicore/pybind11/ops/block_diag.hpp deleted file mode 100644 index b8f9100fe..000000000 --- a/src/infinicore/pybind11/ops/block_diag.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include -#include - -#include "infinicore/ops/block_diag.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_block_diag(py::module &m) { - m.def( - "block_diag", - &op::block_diag, - py::arg("tensors"), - R"doc(Construct a block diagonal matrix from a list of 2D tensors.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/broadcast_to.hpp b/src/infinicore/pybind11/ops/broadcast_to.hpp deleted file mode 100644 index 6ae1e6482..000000000 --- a/src/infinicore/pybind11/ops/broadcast_to.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include -#include - -#include "infinicore/ops/broadcast_to.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_broadcast_to(py::module &m) { - m.def("broadcast_to", - &op::broadcast_to, - py::arg("x"), - py::arg("shape"), - R"doc(Broadcast tensor to target shape.)doc"); - - m.def("broadcast_to_", - &op::broadcast_to_, - py::arg("y"), - py::arg("x"), - R"doc(In-place/Out broadcast tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/cat.hpp b/src/infinicore/pybind11/ops/cat.hpp deleted file mode 100644 index 093b15ff2..000000000 --- a/src/infinicore/pybind11/ops/cat.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include "infinicore/ops/cat.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_cat(py::module &m) { - - m.def("cat", - &op::cat, - py::arg("tensors"), - py::arg("dim") = 0, - R"doc(opertor: torch.cat, out-of-place mode)doc"); - - m.def("cat_", - &op::cat_, - py::arg("out"), - py::arg("tensors"), - py::arg("dim") = 0, - R"doc(opertor: torch.cat, in-place mode)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/causal_conv1d.hpp b/src/infinicore/pybind11/ops/causal_conv1d.hpp deleted file mode 100644 index 4fbeb3c69..000000000 --- a/src/infinicore/pybind11/ops/causal_conv1d.hpp +++ /dev/null @@ -1,38 +0,0 @@ -#pragma once - -#include -#include - -#include "infinicore/ops/causal_conv1d.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_causal_conv1d(py::module &m) { - m.def("causal_conv1d", - &op::causal_conv1d, - py::arg("qkv"), - py::arg("conv_state"), - py::arg("weight"), - py::arg("bias") = std::nullopt, - py::arg("cu_seqlens") = std::nullopt, - py::arg("initial_state_indices") = std::nullopt, - py::arg("final_state_indices") = std::nullopt, - R"doc(Causal depthwise Conv1d. Returns out only. - -Padded mode: - qkv/out: [B, T, C], conv_state: [B, C, state_len]. - -Continuous-batch mode: - pass cu_seqlens [num_requests + 1]; qkv/out: [1, total_tokens, C]. - -Indexed pool mode: - conv_state is [pool_size, C, state_len]. Provide initial_state_indices [num_requests] - to read states. Provide final_state_indices [num_requests] to write final states - in-place to conv_state. The current backend supports K == 4, where - weight is [C, 1, K] and conv_state is [*, C, K - 1]. -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/causal_softmax.hpp b/src/infinicore/pybind11/ops/causal_softmax.hpp deleted file mode 100644 index 926a96d90..000000000 --- a/src/infinicore/pybind11/ops/causal_softmax.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/causal_softmax.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_causal_softmax(py::module &m) { - m.def("causal_softmax", - &op::causal_softmax, - py::arg("input"), - R"doc(Causal softmax activation function.)doc"); - - m.def("causal_softmax_", - &op::causal_softmax_, - py::arg("output"), - py::arg("input"), - R"doc(In-place causal softmax activation function.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/cdist.hpp b/src/infinicore/pybind11/ops/cdist.hpp deleted file mode 100644 index a52c24af3..000000000 --- a/src/infinicore/pybind11/ops/cdist.hpp +++ /dev/null @@ -1,46 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/cdist.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_cdist(py::module &m) { - // 1. 绑定 out-of-place 接口: out = cdist(x1, x2, p) - m.def("cdist", - &op::cdist, - py::arg("x1"), - py::arg("x2"), - py::arg("p") = 2.0, - R"doc(Computes batched pairwise distance between vectors in x1 and x2 using p-norm. - -Args: - x1: First set of vectors, shape (M, D) - x2: Second set of vectors, shape (N, D) - p: The p-norm to apply (default: 2.0) - -Returns: - A matrix containing pairwise distances, shape (M, N) -)doc"); - - // 2. 绑定 in-place / specified output 接口: cdist_(out, x1, x2, p) - m.def("cdist_", - &op::cdist_, - py::arg("out"), - py::arg("x1"), - py::arg("x2"), - py::arg("p") = 2.0, - R"doc(In-place version of cdist. Stores the results in the 'out' tensor. - -Args: - out: The destination tensor, shape (M, N) - x1: First set of vectors, shape (M, D) - x2: Second set of vectors, shape (N, D) - p: The p-norm to apply (default: 2.0) -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/chunk_gated_delta_rule.hpp b/src/infinicore/pybind11/ops/chunk_gated_delta_rule.hpp deleted file mode 100644 index c39174f86..000000000 --- a/src/infinicore/pybind11/ops/chunk_gated_delta_rule.hpp +++ /dev/null @@ -1,42 +0,0 @@ -#pragma once - -#include -#include - -#include "infinicore/ops/chunk_gated_delta_rule.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_chunk_gated_delta_rule(py::module &m) { - m.def("chunk_gated_delta_rule", - &op::chunk_gated_delta_rule, - py::arg("q"), - py::arg("k"), - py::arg("v"), - py::arg("g"), - py::arg("beta"), - py::arg("initial_state"), - py::arg("cu_seqlens") = std::nullopt, - py::arg("initial_state_indices") = std::nullopt, - py::arg("final_state_indices") = std::nullopt, - py::arg("use_qk_l2norm") = false, - py::arg("chunk_size") = 64, - R"doc(Chunk gated delta rule. Returns out only. - -Padded mode: - q/k: [B, T, Hk, Dk], v/out: [B, T, Hv, Dv], g/beta: [B, T, Hv], - initial_state: [B, Hv, Dv, Dk]. - -Continuous-batch mode: - pass cu_seqlens [B + 1]; q/k: [1, total_tokens, Hk, Dk], - v/out: [1, total_tokens, Hv, Dv], g/beta: [1, total_tokens, Hv]. - -Indexed pool mode: - initial_state is [pool_size, Hv, Dv, Dk]. Provide both initial_state_indices - and final_state_indices [B]; final state is written in-place to initial_state. -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/conv2d.hpp b/src/infinicore/pybind11/ops/conv2d.hpp deleted file mode 100644 index f282840e8..000000000 --- a/src/infinicore/pybind11/ops/conv2d.hpp +++ /dev/null @@ -1,61 +0,0 @@ -#pragma once - -#include "infinicore/ops/conv2d.hpp" - -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_conv2d(Tensor input, - Tensor weight, - pybind11::object bias, - const std::vector &padding, - const std::vector &stride, - const std::vector &dilation) { - Tensor bias_tensor; - if (!bias.is_none()) { - bias_tensor = bias.cast(); - } - return op::conv2d(input, weight, bias_tensor, padding, stride, dilation); -} - -void py_conv2d_(Tensor out, - Tensor input, - Tensor weight, - pybind11::object bias, - const std::vector &padding, - const std::vector &stride, - const std::vector &dilation) { - Tensor bias_tensor; - if (!bias.is_none()) { - bias_tensor = bias.cast(); - } - op::conv2d_(out, input, weight, bias_tensor, padding, stride, dilation); -} - -inline void bind_conv2d(py::module &m) { - m.def("conv2d", - &ops::py_conv2d, - py::arg("input"), - py::arg("weight"), - py::arg("bias") = py::none(), - py::arg("padding") = std::vector{0, 0}, - py::arg("stride") = std::vector{1, 1}, - py::arg("dilation") = std::vector{1, 1}, - R"doc(Applies a 2D convolution over an input tensor.)doc"); - - m.def("conv2d_", - &ops::py_conv2d_, - py::arg("out"), - py::arg("input"), - py::arg("weight"), - py::arg("bias") = py::none(), - py::arg("padding") = std::vector{0, 0}, - py::arg("stride") = std::vector{1, 1}, - py::arg("dilation") = std::vector{1, 1}, - R"doc(In-place 2D convolution.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/cross_entropy.hpp b/src/infinicore/pybind11/ops/cross_entropy.hpp deleted file mode 100644 index 8105642a6..000000000 --- a/src/infinicore/pybind11/ops/cross_entropy.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/cross_entropy.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_cross_entropy(py::module &m) { - m.def("cross_entropy", - &op::cross_entropy, - py::arg("logits"), - py::arg("target"), - R"doc(Token-wise cross entropy loss without reduction.)doc"); - - m.def("cross_entropy_", - &op::cross_entropy_, - py::arg("loss"), - py::arg("logits"), - py::arg("target"), - R"doc(Write cross entropy loss into a provided tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/diff.hpp b/src/infinicore/pybind11/ops/diff.hpp deleted file mode 100644 index 4a83d9607..000000000 --- a/src/infinicore/pybind11/ops/diff.hpp +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/diff.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_diff(py::module &m) { - m.def("diff", - &op::diff, - py::arg("x"), - py::arg("n") = 1, - py::arg("dim") = -1, - R"doc(Difference of adjacent elements along a dimension.)doc"); - - m.def("diff_", - &op::diff_, - py::arg("y"), - py::arg("x"), - py::arg("n") = 1, - py::arg("dim") = -1, - R"doc(Out variant of diff.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/digamma.hpp b/src/infinicore/pybind11/ops/digamma.hpp deleted file mode 100644 index 3110c7a07..000000000 --- a/src/infinicore/pybind11/ops/digamma.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/digamma.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_digamma(py::module &m) { - m.def("digamma", - &op::digamma, - py::arg("x"), - R"doc(Digamma function.)doc"); - - m.def("digamma_", - &op::digamma_, - py::arg("y"), - py::arg("x"), - R"doc(Out variant of digamma.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/dist.hpp b/src/infinicore/pybind11/ops/dist.hpp deleted file mode 100644 index 9aec0cca6..000000000 --- a/src/infinicore/pybind11/ops/dist.hpp +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/dist.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_dist(py::module &m) { - m.def("dist", - &op::dist, - py::arg("x1"), - py::arg("x2"), - py::arg("p") = 2.0, - R"doc(p-norm distance between two tensors.)doc"); - - m.def("dist_", - &op::dist_, - py::arg("y"), - py::arg("x1"), - py::arg("x2"), - py::arg("p") = 2.0, - R"doc(Out variant of dist.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/embedding.hpp b/src/infinicore/pybind11/ops/embedding.hpp deleted file mode 100644 index 44e14b61c..000000000 --- a/src/infinicore/pybind11/ops/embedding.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include "infinicore/ops/embedding.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_embedding(py::module &m) { - - m.def("embedding", - &op::embedding, - py::arg("input"), - py::arg("weight"), - R"doc(Generate a simple lookup table that looks up embeddings in a fixed dictionary and size..)doc"); - - m.def("embedding_", - &op::embedding_, - py::arg("out"), - py::arg("input"), - py::arg("weight"), - R"doc(In-place, Generate a simple lookup table that looks up embeddings in a fixed dictionary and size..)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/equal.hpp b/src/infinicore/pybind11/ops/equal.hpp deleted file mode 100644 index d14a6b61d..000000000 --- a/src/infinicore/pybind11/ops/equal.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/equal.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_equal(py::module &m) { - m.def("equal", - &op::equal, - py::arg("a"), - py::arg("b"), - R"doc(Elementwise equality returning a bool tensor.)doc"); - - m.def("equal_", - &op::equal_, - py::arg("out"), - py::arg("a"), - py::arg("b"), - R"doc(In-place elementwise equality writing into `out`.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/flash_attention.hpp b/src/infinicore/pybind11/ops/flash_attention.hpp deleted file mode 100644 index 6e3766796..000000000 --- a/src/infinicore/pybind11/ops/flash_attention.hpp +++ /dev/null @@ -1,22 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/flash_attention.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_flash_attention(py::module &m) { - m.def("flash_attention", - &op::flash_attention, - py::arg("q"), - py::arg("k"), - py::arg("v"), - py::arg("total_kv_len"), - py::arg("scale"), - py::arg("is_causal")); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/flipud.hpp b/src/infinicore/pybind11/ops/flipud.hpp deleted file mode 100644 index 97c5641d6..000000000 --- a/src/infinicore/pybind11/ops/flipud.hpp +++ /dev/null @@ -1,31 +0,0 @@ -#pragma once - -#include "infinicore/ops/flipud.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_flipud(py::module &m) { - // 1. 绑定 out-of-place 接口: output = flipud(input) - m.def("flipud", - &op::flipud, - py::arg("input"), - R"doc(Flip array in the up/down direction. - - Flips the entries in axis 0 (preserving the shape). - - Args: - input (Tensor): The input tensor. - )doc"); - - // 2. 绑定 explicit output 接口: flipud_(output, input) - m.def("flipud_", - &op::flipud_, - py::arg("output"), - py::arg("input"), - R"doc(Explicit output FlipUD operation. Writes the result into the output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/float_power.hpp b/src/infinicore/pybind11/ops/float_power.hpp deleted file mode 100644 index 8ee3a5d77..000000000 --- a/src/infinicore/pybind11/ops/float_power.hpp +++ /dev/null @@ -1,59 +0,0 @@ -#include "../tensor.hpp" -#include "infinicore/ops/float_power.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -using infinicore::Tensor; -using infinicore::op::float_power; -using infinicore::op::float_power_; - -inline Tensor unwrap(py::handle obj) { - try { - return obj.cast(); - } catch (...) {} - - if (py::hasattr(obj, "_underlying")) { - return obj.attr("_underlying").cast(); - } - - throw py::type_error("Expected infinicore.Tensor, but got " + py::repr(obj.get_type()).cast()); -} - -void bind_float_power(py::module &m) { - - // --- Out-of-place: float_power(input, exponent) --- - m.def( - "float_power", [](py::object input_obj, py::object exp_obj) -> Tensor { - Tensor input = unwrap(input_obj); - - // 处理标量指数的情况 (float 或 int) - if (py::isinstance(exp_obj) || py::isinstance(exp_obj)) { - return float_power(input, exp_obj.cast()); - } - - // 处理张量指数的情况 - Tensor exponent = unwrap(exp_obj); - return float_power(input, exponent); - }, - py::arg("input"), py::arg("exponent")); - - // --- In-place: float_power_(out, input, exponent) --- - m.def( - "float_power_", [](py::object out_obj, py::object input_obj, py::object exp_obj) { - Tensor out = unwrap(out_obj); - Tensor input = unwrap(input_obj); - - if (py::isinstance(exp_obj) || py::isinstance(exp_obj)) { - float_power_(out, input, exp_obj.cast()); - } else { - Tensor exponent = unwrap(exp_obj); - float_power_(out, input, exponent); - } - }, - py::arg("out"), py::arg("input"), py::arg("exponent")); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/floor.hpp b/src/infinicore/pybind11/ops/floor.hpp deleted file mode 100644 index a209ea5bc..000000000 --- a/src/infinicore/pybind11/ops/floor.hpp +++ /dev/null @@ -1,25 +0,0 @@ -#pragma once - -#include "infinicore/ops/floor.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_floor(py::module &m) { - // 绑定 out-of-place 接口: output = floor(input) - m.def("floor", - &op::floor, - py::arg("input"), - R"doc(Computes the floor of each element of input.)doc"); - - // 绑定 in-place 接口: floor_(output, input) - m.def("floor_", - &op::floor_, - py::arg("output"), - py::arg("input"), - R"doc(In-place floor operation. Writes result into output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/floor_divide.hpp b/src/infinicore/pybind11/ops/floor_divide.hpp deleted file mode 100644 index 626767218..000000000 --- a/src/infinicore/pybind11/ops/floor_divide.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/floor_divide.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_floor_divide(py::module &m) { - m.def("floor_divide", - &op::floor_divide, - py::arg("a"), - py::arg("b"), - R"doc(Floor division of two tensors.)doc"); - - m.def("floor_divide_", - &op::floor_divide_, - py::arg("c"), - py::arg("a"), - py::arg("b"), - R"doc(In-place tensor floor division.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/fmin.hpp b/src/infinicore/pybind11/ops/fmin.hpp deleted file mode 100644 index afdbe4b33..000000000 --- a/src/infinicore/pybind11/ops/fmin.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/fmin.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_fmin(py::module &m) { - m.def("fmin", - &op::fmin, - py::arg("a"), - py::arg("b"), - R"doc(fmin of two tensors.)doc"); - - m.def("fmin_", - &op::fmin_, - py::arg("c"), - py::arg("a"), - py::arg("b"), - R"doc(In-place tensor fmin.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/fmod.hpp b/src/infinicore/pybind11/ops/fmod.hpp deleted file mode 100644 index 97af57da2..000000000 --- a/src/infinicore/pybind11/ops/fmod.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/fmod.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_fmod(py::module &m) { - m.def("fmod", - &op::fmod, - py::arg("a"), - py::arg("b"), - R"doc(Element-wise floating point remainder of division of two tensors.)doc"); - - m.def("fmod_", - &op::fmod_, - py::arg("c"), - py::arg("a"), - py::arg("b"), - R"doc(In-place element-wise floating point remainder of division of two tensors.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/fused_gated_delta_net_gating.hpp b/src/infinicore/pybind11/ops/fused_gated_delta_net_gating.hpp deleted file mode 100644 index e43566977..000000000 --- a/src/infinicore/pybind11/ops/fused_gated_delta_net_gating.hpp +++ /dev/null @@ -1,44 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/fused_gated_delta_net_gating.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_fused_gated_delta_net_gating(py::module &m) { - m.def( - "fused_gated_delta_net_gating", - [](const Tensor &A_log, - const Tensor &a, - const Tensor &b, - const Tensor &dt_bias, - float beta, - float threshold) { - auto result = op::fused_gated_delta_net_gating(A_log, a, b, dt_bias, beta, threshold); - return py::make_tuple(result.first, result.second); - }, - py::arg("A_log"), - py::arg("a"), - py::arg("b"), - py::arg("dt_bias"), - py::arg("beta") = 1.0f, - py::arg("threshold") = 20.0f, - R"doc(Fused GatedDeltaNet gating out-of-place.)doc"); - - m.def("fused_gated_delta_net_gating_", - &op::fused_gated_delta_net_gating_, - py::arg("g"), - py::arg("beta_output"), - py::arg("A_log"), - py::arg("a"), - py::arg("b"), - py::arg("dt_bias"), - py::arg("beta") = 1.0f, - py::arg("threshold") = 20.0f, - R"doc(Fused GatedDeltaNet gating writing to provided outputs.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/gaussian_nll_loss.hpp b/src/infinicore/pybind11/ops/gaussian_nll_loss.hpp deleted file mode 100644 index 974f8e29b..000000000 --- a/src/infinicore/pybind11/ops/gaussian_nll_loss.hpp +++ /dev/null @@ -1,34 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/gaussian_nll_loss.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_gaussian_nll_loss(py::module &m) { - m.def("gaussian_nll_loss", - &op::gaussian_nll_loss, - py::arg("input"), - py::arg("target"), - py::arg("var"), - py::arg("full") = false, - py::arg("eps") = 1e-6, - py::arg("reduction") = 1, - R"doc(Gaussian negative log-likelihood loss.)doc"); - - m.def("gaussian_nll_loss_", - &op::gaussian_nll_loss_, - py::arg("out"), - py::arg("input"), - py::arg("target"), - py::arg("var"), - py::arg("full") = false, - py::arg("eps") = 1e-6, - py::arg("reduction") = 1, - R"doc(In-place Gaussian negative log-likelihood loss.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/hardswish.hpp b/src/infinicore/pybind11/ops/hardswish.hpp deleted file mode 100644 index daaccec62..000000000 --- a/src/infinicore/pybind11/ops/hardswish.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/hardswish.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_hardswish(py::module &m) { - m.def("hardswish", - &op::hardswish, - py::arg("input"), - R"doc(Out-of-place Hardswish activation.)doc"); - - m.def("hardswish_", - &op::hardswish_, - py::arg("output"), - py::arg("input"), - R"doc(In-place Hardswish activation.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/hardtanh.hpp b/src/infinicore/pybind11/ops/hardtanh.hpp deleted file mode 100644 index ff9abb872..000000000 --- a/src/infinicore/pybind11/ops/hardtanh.hpp +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/hardtanh.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_hardtanh(py::module &m) { - m.def("hardtanh", - &op::hardtanh, - py::arg("input"), - py::arg("min_val") = -1.0f, - py::arg("max_val") = 1.0f, - R"doc(Apply the HardTanh activation.)doc"); - - m.def("hardtanh_", - &op::hardtanh_, - py::arg("output"), - py::arg("input"), - py::arg("min_val") = -1.0f, - py::arg("max_val") = 1.0f, - R"doc(In-place HardTanh activation.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/hinge_embedding_loss.hpp b/src/infinicore/pybind11/ops/hinge_embedding_loss.hpp deleted file mode 100644 index fe481d853..000000000 --- a/src/infinicore/pybind11/ops/hinge_embedding_loss.hpp +++ /dev/null @@ -1,22 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/hinge_embedding_loss.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_hinge_embedding_loss(py::module &m) { - m.def( - "hinge_embedding_loss", - &op::hinge_embedding_loss, - py::arg("input"), - py::arg("target"), - py::arg("margin") = 1.0, - py::arg("reduction") = 1, - R"doc(Hinge embedding loss.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/huber_loss.hpp b/src/infinicore/pybind11/ops/huber_loss.hpp deleted file mode 100644 index 41c5f4e50..000000000 --- a/src/infinicore/pybind11/ops/huber_loss.hpp +++ /dev/null @@ -1,46 +0,0 @@ -#pragma once - -#include "infinicore/ops/huber_loss.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_huber_loss(py::module &m) { - m.def( - "huber_loss", - [](const Tensor &input, const Tensor &target, float delta, int reduction) { - return op::huber_loss(input, target, delta, reduction); - }, - py::arg("input"), - py::arg("target"), - py::arg("delta") = 1.0f, - py::arg("reduction") = 1, - R"doc(Computes the Huber Loss between input and target. - - Args: - input (Tensor): Input tensor of arbitrary shape. - target (Tensor): Ground truth labels, same shape as input. - delta (float, optional): The threshold at which to change between delta-scaled L1 and L2 loss. Default: 1.0. - reduction (int, optional): Specifies the reduction to apply to the output: 0=None, 1=Mean, 2=Sum. Default: 1. - )doc"); - - // ------------------------------------------------------------------------- - // 2. 绑定 in-place 接口 (huber_loss_) - // ------------------------------------------------------------------------- - m.def( - "huber_loss_", - [](Tensor &output, const Tensor &input, const Tensor &target, float delta, int reduction) { - // 调用底层 - op::huber_loss_(output, input, target, delta, reduction); - }, - py::arg("output"), - py::arg("input"), - py::arg("target"), - py::arg("delta") = 1.0f, - py::arg("reduction") = 1, - R"doc(Explicit output Huber Loss operation. Writes the result into the output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/hypot.hpp b/src/infinicore/pybind11/ops/hypot.hpp deleted file mode 100644 index 70d08ef19..000000000 --- a/src/infinicore/pybind11/ops/hypot.hpp +++ /dev/null @@ -1,27 +0,0 @@ -#pragma once - -#include "infinicore/ops/hypot.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_hypot(py::module &m) { - // 绑定 out-of-place 接口: output = hypot(input, other) - m.def("hypot", - &op::hypot, - py::arg("input"), - py::arg("other"), - R"doc(Computes the hypotenuse of input and other arguments, i.e. sqrt(input^2 + other^2).)doc"); - - // 绑定 in-place / 指定输出接口: hypot_(output, input, other) - m.def("hypot_", - &op::hypot_, - py::arg("output"), - py::arg("input"), - py::arg("other"), - R"doc(In-place hypot operation. Writes result into output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/index_add.hpp b/src/infinicore/pybind11/ops/index_add.hpp deleted file mode 100644 index 4fadef4d1..000000000 --- a/src/infinicore/pybind11/ops/index_add.hpp +++ /dev/null @@ -1,31 +0,0 @@ -#pragma once - -#include "infinicore/ops/index_add.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_index_add(py::module &m) { - m.def("index_add", - &op::index_add, - py::arg("input"), - py::arg("dim"), - py::arg("index"), - py::arg("source"), - py::arg("alpha") = 1.0f, - R"doc(Accumulate elements of source into input by adding to the indices in the order given in index. - Formula: output[index[i]] = input[index[i]] + alpha * source[i])doc"); - m.def("index_add_", - &op::index_add_, - py::arg("output"), - py::arg("input"), - py::arg("dim"), - py::arg("index"), - py::arg("source"), - py::arg("alpha") = 1.0f, - R"doc(In-place version of index_add. Writes result into output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/index_copy.hpp b/src/infinicore/pybind11/ops/index_copy.hpp deleted file mode 100644 index c631be95c..000000000 --- a/src/infinicore/pybind11/ops/index_copy.hpp +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once - -#include "infinicore/ops/index_copy.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_index_copy(py::module &m) { - // 1. Out-of-place version (returns new tensor) - m.def("index_copy", - &op::index_copy, - py::arg("input"), - py::arg("dim"), - py::arg("index"), - py::arg("source"), - R"doc(Copies elements of source into input at the indices given in index. - Formula: output[index[i]] = source[i])doc"); - m.def("index_copy_", - &op::index_copy_, - py::arg("output"), - py::arg("input"), - py::arg("dim"), - py::arg("index"), - py::arg("source"), - R"doc(In-place version of index_copy. Writes result into output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/inner.hpp b/src/infinicore/pybind11/ops/inner.hpp deleted file mode 100644 index b54d6117b..000000000 --- a/src/infinicore/pybind11/ops/inner.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include "infinicore/ops/inner.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_inner(py::module &m) { - - m.def("inner", - &op::inner, - py::arg("input"), - py::arg("other"), - R"doc(opertor: torch.inner, out-of-place mode)doc"); - - m.def("inner_", - &op::inner_, - py::arg("out"), - py::arg("input"), - py::arg("other"), - R"doc(opertor: torch.inner, in-place mode)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/interpolate.hpp b/src/infinicore/pybind11/ops/interpolate.hpp deleted file mode 100644 index 26fb3b812..000000000 --- a/src/infinicore/pybind11/ops/interpolate.hpp +++ /dev/null @@ -1,33 +0,0 @@ -#pragma once - -#include -#include - -#include "infinicore/ops/interpolate.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_interpolate(py::module &m) { - m.def("interpolate", - &op::interpolate, - py::arg("input"), - py::arg("mode"), - py::arg("size"), - py::arg("scale_factor"), - py::arg("align_corners"), - R"doc(Interpolate (upsample/downsample) a tensor.)doc"); - - m.def("interpolate_", - &op::interpolate_, - py::arg("out"), - py::arg("input"), - py::arg("mode"), - py::arg("size"), - py::arg("scale_factor"), - py::arg("align_corners"), - R"doc(In-place interpolate (writes to out).)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/kron.hpp b/src/infinicore/pybind11/ops/kron.hpp deleted file mode 100644 index 751ef879f..000000000 --- a/src/infinicore/pybind11/ops/kron.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/kron.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_kron(py::module &m) { - m.def( - "kron", - &op::kron, - py::arg("a"), - py::arg("b"), - R"doc(Kronecker product.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/kthvalue.hpp b/src/infinicore/pybind11/ops/kthvalue.hpp deleted file mode 100644 index a0f67ae2a..000000000 --- a/src/infinicore/pybind11/ops/kthvalue.hpp +++ /dev/null @@ -1,39 +0,0 @@ -#pragma once - -#include "infinicore/ops/kthvalue.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_kthvalue(py::module &m) { - // 1. 绑定 functional 接口: (values, indices) = kthvalue(input, k, dim, keepdim) - m.def("kthvalue", - &op::kthvalue, - py::arg("input"), - py::arg("k"), - py::arg("dim") = -1, - py::arg("keepdim") = false, - R"doc(Returns the k-th smallest element of each row of the input tensor in the given dimension. - - Args: - input (Tensor): The input tensor. - k (int): The k value. - dim (int): The dimension to find the k-th value along. - keepdim (bool): Whether to keep the output dimension. - )doc"); - - // 2. 绑定 explicit output 接口: kthvalue_(values, indices, input, k, dim, keepdim) - m.def("kthvalue_", - &op::kthvalue_, - py::arg("values"), - py::arg("indices"), - py::arg("input"), - py::arg("k"), - py::arg("dim") = -1, - py::arg("keepdim") = false, - R"doc(Explicit output Kthvalue operation. Writes results into values and indices tensors.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/kv_caching.hpp b/src/infinicore/pybind11/ops/kv_caching.hpp deleted file mode 100644 index 2864312b2..000000000 --- a/src/infinicore/pybind11/ops/kv_caching.hpp +++ /dev/null @@ -1,32 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/kv_caching.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_kv_caching(py::module &m) { - m.def("kv_caching_", - &op::kv_caching_, - py::arg("k_cache"), - py::arg("v_cache"), - py::arg("k"), - py::arg("v"), - py::arg("past_kv_lengths"), - R"doc(In-place Key-Value Caching. - -Updates the KV cache in-place with new key and value tensors. - -Args: - k_cache: Key cache tensor to update in-place - v_cache: Value cache tensor to update in-place - k: New key tensor to append - v: New value tensor to append - past_kv_lengths: Tensor containing current sequence lengths for each batch -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/layer_norm.hpp b/src/infinicore/pybind11/ops/layer_norm.hpp deleted file mode 100644 index 5ca6a87d3..000000000 --- a/src/infinicore/pybind11/ops/layer_norm.hpp +++ /dev/null @@ -1,48 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/layer_norm.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_layer_norm(py::module &m) { - m.def("layer_norm", - &op::layer_norm, - py::arg("x"), - py::arg("weight"), - py::arg("bias"), - py::arg("epsilon") = 1e-5f, - R"doc(Layer Normalization. - -Args: - x: Input tensor - weight: Scale weights - bias: Bias weights - epsilon: Small constant for numerical stability, default is 1e-5 - -Returns: - Normalized tensor with same shape as input -)doc"); - - m.def("layer_norm_", - &op::layer_norm_for_pybind, - py::arg("y"), - py::arg("x"), - py::arg("weight"), - py::arg("bias"), - py::arg("epsilon") = 1e-5f, - R"doc(In-place Layer Normalization. - -Args: - y: Output tensor - x: Input tensor - weight: Scale weights - bias: Bias weights - epsilon: Small constant for numerical stability, default is 1e-5 -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/ldexp.hpp b/src/infinicore/pybind11/ops/ldexp.hpp deleted file mode 100644 index 948c03d68..000000000 --- a/src/infinicore/pybind11/ops/ldexp.hpp +++ /dev/null @@ -1,40 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/ldexp.hpp" - -namespace py = pybind11; - -#pragma once - -#include "infinicore/ops/ldexp.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_ldexp(py::module &m) { - // 1. 绑定 functional 接口: output = ldexp(input, other) - m.def("ldexp", - &op::ldexp, - py::arg("input"), - py::arg("other"), - R"doc(Multiplies input by 2 raised to the power of other. - - Args: - input (Tensor): The input tensor (mantissa). - other (Tensor): The exponent tensor. - )doc"); - - // 2. 绑定 explicit output 接口: ldexp_(output, input, other) - m.def("ldexp_", - &op::ldexp_, - py::arg("output"), - py::arg("input"), - py::arg("other"), - R"doc(Explicit output Ldexp operation. Writes result into output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/lerp.hpp b/src/infinicore/pybind11/ops/lerp.hpp deleted file mode 100644 index d1100b8f2..000000000 --- a/src/infinicore/pybind11/ops/lerp.hpp +++ /dev/null @@ -1,64 +0,0 @@ -#pragma once - -#include "infinicore/ops/lerp.hpp" -#include -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_lerp(py::module &m) { - // 定义函数指针别名，用于区分重载 - using LerpTensorFunc = Tensor (*)(Tensor, Tensor, Tensor); - using LerpScalarFunc = Tensor (*)(Tensor, Tensor, float); - using LerpTensorInplaceFunc = void (*)(Tensor, Tensor, Tensor, Tensor); - using LerpScalarInplaceFunc = void (*)(Tensor, Tensor, Tensor, float); - - // ======================================================================== - // 1. 绑定 functional 接口 - // ======================================================================== - - // 重载 1: weight 为 Tensor - m.def("lerp", - static_cast(&op::lerp), - py::arg("start"), - py::arg("end"), - py::arg("weight"), - R"doc(Does a linear interpolation of two tensors start and end based on a tensor weight. - - output = start + weight * (end - start) - )doc"); - - // 重载 2: weight 为 float - m.def("lerp", - static_cast(&op::lerp), - py::arg("start"), - py::arg("end"), - py::arg("weight"), - R"doc(Does a linear interpolation of two tensors start and end based on a scalar weight.)doc"); - - // ======================================================================== - // 2. 绑定 explicit output 接口 (In-place) - // ======================================================================== - - // 重载 1: weight 为 Tensor - m.def("lerp_", - static_cast(&op::lerp_), - py::arg("output"), - py::arg("start"), - py::arg("end"), - py::arg("weight"), - R"doc(Explicit output Lerp operation with tensor weight. Writes the result into the output tensor.)doc"); - - // 重载 2: weight 为 float - m.def("lerp_", - static_cast(&op::lerp_), - py::arg("output"), - py::arg("start"), - py::arg("end"), - py::arg("weight"), - R"doc(Explicit output Lerp operation with scalar weight. Writes the result into the output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/linear.hpp b/src/infinicore/pybind11/ops/linear.hpp deleted file mode 100644 index e8f9a90a2..000000000 --- a/src/infinicore/pybind11/ops/linear.hpp +++ /dev/null @@ -1,56 +0,0 @@ -#pragma once - -#include "infinicore/ops/linear.hpp" - -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_linear(Tensor input, - Tensor weight, - pybind11::object bias, - float alpha = 1.0f) { - std::optional bias_tensor = std::nullopt; - if (!bias.is_none()) { - bias_tensor = bias.cast(); - } - return op::linear(input, weight, bias_tensor, alpha); -} - -void py_linear_(Tensor out, - Tensor input, - Tensor weight, - pybind11::object bias, - float alpha = 1.0f) { - - std::optional bias_tensor = std::nullopt; - if (!bias.is_none()) { - bias_tensor = bias.cast(); - } - - op::linear_(out, input, weight, bias_tensor, alpha); -} - -inline void bind_linear(py::module &m) { - - m.def("linear", - &ops::py_linear, - py::arg("input"), - py::arg("weight"), - py::arg("bias") = py::none(), - py::arg("alpha") = 1.0f, - R"doc(Applies a linear transformation to the incoming data: y=alpha*xA^T+b.)doc"); - - m.def("linear_", - &ops::py_linear_, - py::arg("out"), - py::arg("input"), - py::arg("weight"), - py::arg("bias") = py::none(), - py::arg("alpha") = 1.0f, - R"doc(In-place, applies a linear transformation to the incoming data: y=alpha*xA^T+b.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/linear_w8a8i8.hpp b/src/infinicore/pybind11/ops/linear_w8a8i8.hpp deleted file mode 100644 index 926d554b1..000000000 --- a/src/infinicore/pybind11/ops/linear_w8a8i8.hpp +++ /dev/null @@ -1,54 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/linear_w8a8i8.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_linear_w8a8i8(Tensor input, - Tensor weight_packed, - Tensor weight_scale, - pybind11::object bias) { - std::optional bias_tensor = std::nullopt; - if (!bias.is_none()) { - bias_tensor = bias.cast(); - } - return op::linear_w8a8i8(input, weight_packed, weight_scale, bias_tensor); -} - -void py_linear_w8a8i8_(Tensor out, - Tensor input, - Tensor weight_packed, - Tensor weight_scale, - pybind11::object bias) { - - std::optional bias_tensor = std::nullopt; - if (!bias.is_none()) { - bias_tensor = bias.cast(); - } - - op::linear_w8a8i8_(out, input, weight_packed, weight_scale, bias_tensor); -} - -inline void bind_linear_w8a8i8(py::module &m) { - m.def("linear_w8a8i8", - &ops::py_linear_w8a8i8, - py::arg("input"), - py::arg("weight_packed"), - py::arg("weight_scale"), - py::arg("bias") = py::none(), - R"doc(linear_w8a8i8.)doc"); - m.def("linear_w8a8i8_", - &ops::py_linear_w8a8i8_, - py::arg("out"), - py::arg("input"), - py::arg("weight_packed"), - py::arg("weight_scale"), - py::arg("bias") = py::none(), - R"doc(linear_w8a8i8_.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/log_softmax.hpp b/src/infinicore/pybind11/ops/log_softmax.hpp deleted file mode 100644 index b2eb8f47d..000000000 --- a/src/infinicore/pybind11/ops/log_softmax.hpp +++ /dev/null @@ -1,32 +0,0 @@ -#pragma once - -#include "infinicore/ops/log_softmax.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_log_softmax(py::module &m) { - // 1. 绑定 functional 接口: output = log_softmax(input, dim) - m.def("log_softmax", - &op::log_softmax, - py::arg("input"), - py::arg("dim"), - R"doc(Applies a softmax followed by a logarithm. - - Args: - input (Tensor): The input tensor. - dim (int): A dimension along which log_softmax will be computed. - )doc"); - - // 2. 绑定 explicit output 接口: log_softmax_(output, input, dim) - m.def("log_softmax_", - &op::log_softmax_, - py::arg("output"), - py::arg("input"), - py::arg("dim"), - R"doc(Explicit output LogSoftmax operation. Writes results into output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/logaddexp.hpp b/src/infinicore/pybind11/ops/logaddexp.hpp deleted file mode 100644 index 8f1bc6b18..000000000 --- a/src/infinicore/pybind11/ops/logaddexp.hpp +++ /dev/null @@ -1,25 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/logaddexp.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_logaddexp(py::module &m) { - m.def("logaddexp", - &op::logaddexp, - py::arg("a"), - py::arg("b"), - R"doc(Logarithm of the sum of exponentiations of the inputs.)doc"); - m.def("logaddexp_", - &op::logaddexp_, - py::arg("c"), - py::arg("a"), - py::arg("b"), - R"doc(In-place logaddexp operation. Writes results into c tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/logaddexp2.hpp b/src/infinicore/pybind11/ops/logaddexp2.hpp deleted file mode 100644 index 77ee154d7..000000000 --- a/src/infinicore/pybind11/ops/logaddexp2.hpp +++ /dev/null @@ -1,25 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/logaddexp2.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_logaddexp2(py::module &m) { - m.def("logaddexp2", - &op::logaddexp2, - py::arg("a"), - py::arg("b"), - R"doc(Logarithm of the sum of exponentiations of the inputs in base-2.)doc"); - m.def("logaddexp2_", - &op::logaddexp2_, - py::arg("c"), - py::arg("a"), - py::arg("b"), - R"doc(In-place logaddexp2 operation. Writes results into c tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/logcumsumexp.hpp b/src/infinicore/pybind11/ops/logcumsumexp.hpp deleted file mode 100644 index aace6cc07..000000000 --- a/src/infinicore/pybind11/ops/logcumsumexp.hpp +++ /dev/null @@ -1,31 +0,0 @@ -#pragma once - -#include "infinicore/ops/logcumsumexp.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_logcumsumexp(py::module &m) { - // 绑定非原地操作接口 (返回新 Tensor) - m.def("logcumsumexp", - &op::logcumsumexp, - py::arg("input"), - py::arg("dim"), - py::arg("exclusive") = false, - py::arg("reverse") = false, - R"doc(Computes the logarithm of the cumulative summation of the exponentiation of elements.)doc"); - - // 绑定原地/指定输出接口 - m.def("logcumsumexp_", - &op::logcumsumexp_, - py::arg("out"), - py::arg("input"), - py::arg("dim"), - py::arg("exclusive") = false, - py::arg("reverse") = false, - R"doc(In-place version of logcumsumexp.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/logdet.hpp b/src/infinicore/pybind11/ops/logdet.hpp deleted file mode 100644 index e2a8f2d28..000000000 --- a/src/infinicore/pybind11/ops/logdet.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/logdet.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_logdet(py::module &m) { - m.def("logdet", - &op::logdet, - py::arg("x"), - R"doc(Log determinant of a square matrix (NaN if determinant is negative).)doc"); - - m.def("logdet_", - &op::logdet_, - py::arg("y"), - py::arg("x"), - R"doc(Out variant of logdet.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/logical_and.hpp b/src/infinicore/pybind11/ops/logical_and.hpp deleted file mode 100644 index 1ca34fdc0..000000000 --- a/src/infinicore/pybind11/ops/logical_and.hpp +++ /dev/null @@ -1,27 +0,0 @@ -#pragma once -#include - -#include "infinicore/ops/logical_and.hpp" - -namespace py = pybind11; -namespace infinicore::ops { - -inline void bind_logical_and(py::module &m) { - // 绑定常规函数: logical_and(input, other) -> Tensor - m.def("logical_and", - &op::logical_and, - py::arg("input"), - py::arg("other"), - R"doc(Computes the element-wise logical AND of the given input tensors.)doc"); - - // 绑定底层输出指定函数: logical_and_(output, input, other) - // 对应 Python 调用: _infinicore.logical_and_(out, input, other) - m.def("logical_and_", - &op::logical_and_, - py::arg("output"), - py::arg("input"), - py::arg("other"), - R"doc(Explicit output logical AND computation.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/logical_not.hpp b/src/infinicore/pybind11/ops/logical_not.hpp deleted file mode 100644 index 94edd2388..000000000 --- a/src/infinicore/pybind11/ops/logical_not.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include "infinicore/ops/logical_not.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_logical_not(py::module &m) { - // Out-of-place: output = logical_not(input) - m.def("logical_not", - &op::logical_not, - py::arg("input"), - R"doc(Logical NOT of the tensor.)doc"); - - // In-place / Explicit Output: logical_not_(output, input) - // 对应 C++: void logical_not_(Tensor output, Tensor input) - m.def("logical_not_", - &op::logical_not_, - py::arg("output"), - py::arg("input"), - R"doc(In-place logical NOT computation.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/mamba_selective_scan.hpp b/src/infinicore/pybind11/ops/mamba_selective_scan.hpp deleted file mode 100644 index 7bf6e9508..000000000 --- a/src/infinicore/pybind11/ops/mamba_selective_scan.hpp +++ /dev/null @@ -1,34 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/mamba_selective_scan.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_mamba_selective_scan(py::module &m) { - m.def("mamba_selective_scan", - &op::mamba_selective_scan, - py::arg("x"), - py::arg("dt"), - py::arg("b"), - py::arg("c"), - py::arg("a_log"), - py::arg("d"), - py::arg("gate"), - py::arg("dt_bias"), - py::arg("state"), - R"doc(Mamba selective scan. Returns out and updates state in-place. - -Shapes: - x, dt, gate, out: [batch, seq_len, intermediate] - b, c: [batch, seq_len, state_size] - a_log: [intermediate, state_size] - d, dt_bias: [intermediate] - state: [batch, intermediate, state_size], float32 -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/masked_select.hpp b/src/infinicore/pybind11/ops/masked_select.hpp deleted file mode 100644 index 949566b6b..000000000 --- a/src/infinicore/pybind11/ops/masked_select.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include "infinicore/ops/masked_select.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_masked_select(py::module &m) { - - m.def("masked_select", - &op::masked_select, - py::arg("input"), - py::arg("mask"), - R"doc(opertor: torch.masked_select, out-of-place mode)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/matmul.hpp b/src/infinicore/pybind11/ops/matmul.hpp deleted file mode 100644 index c00bd2566..000000000 --- a/src/infinicore/pybind11/ops/matmul.hpp +++ /dev/null @@ -1,28 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/matmul.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_matmul(py::module &m) { - m.def("matmul", - &op::matmul, - py::arg("a"), - py::arg("b"), - py::arg("alpha") = 1.0f, - R"doc(Matrix multiplication of two tensors.)doc"); - - m.def("matmul_", - &op::matmul_, - py::arg("c"), - py::arg("a"), - py::arg("b"), - py::arg("alpha") = 1.0f, - R"doc(In-place matrix multiplication.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/mha.hpp b/src/infinicore/pybind11/ops/mha.hpp deleted file mode 100644 index 58dd21734..000000000 --- a/src/infinicore/pybind11/ops/mha.hpp +++ /dev/null @@ -1,78 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/mha.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_mha(Tensor q, - Tensor k, - Tensor v, - pybind11::object alibi_slopes, - float scale, - bool is_causal) { - std::optional alibi_slopes_tensor = std::nullopt; - if (!alibi_slopes.is_none()) { - alibi_slopes_tensor = alibi_slopes.cast(); - } - - return op::mha( - q, - k, - v, - alibi_slopes_tensor, - scale, - is_causal); -} - -void py_mha_(Tensor out, - Tensor q, - Tensor k, - Tensor v, - pybind11::object alibi_slopes, - float scale, - bool is_causal) { - std::optional alibi_slopes_tensor = std::nullopt; - if (!alibi_slopes.is_none()) { - alibi_slopes_tensor = alibi_slopes.cast(); - } - - op::mha_( - out, - q, - k, - v, - alibi_slopes_tensor, - scale, - is_causal); -} - -inline void bind_mha(py::module &m) { - m.def( - "mha", - &ops::py_mha, - py::arg("q"), - py::arg("k"), - py::arg("v"), - py::arg("alibi_slopes"), - py::arg("scale"), - py::arg("is_causal"), - R"doc(Variable-length multi-head attention.)doc"); - - m.def( - "mha_", - &ops::py_mha_, - py::arg("out"), - py::arg("q"), - py::arg("k"), - py::arg("v"), - py::arg("alibi_slopes"), - py::arg("scale"), - py::arg("is_causal"), - R"doc(In-place variable-length multi-head attention.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/mha_kvcache.hpp b/src/infinicore/pybind11/ops/mha_kvcache.hpp deleted file mode 100644 index 38934233e..000000000 --- a/src/infinicore/pybind11/ops/mha_kvcache.hpp +++ /dev/null @@ -1,127 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/mha_kvcache.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_mha_kvcache(Tensor q, - Tensor k_cache, - Tensor v_cache, - Tensor seqlens_k, - Tensor block_table, - pybind11::object alibi_slopes, - float scale) { - std::optional alibi_slopes_tensor = std::nullopt; - if (!alibi_slopes.is_none()) { - alibi_slopes_tensor = alibi_slopes.cast(); - } - - return op::mha_kvcache( - q, - k_cache, - v_cache, - seqlens_k, - block_table, - alibi_slopes_tensor, - scale); -} - -void py_mha_kvcache_(Tensor out, - Tensor q, - Tensor k_cache, - Tensor v_cache, - Tensor seqlens_k, - Tensor block_table, - pybind11::object alibi_slopes, - float scale) { - std::optional alibi_slopes_tensor = std::nullopt; - if (!alibi_slopes.is_none()) { - alibi_slopes_tensor = alibi_slopes.cast(); - } - - op::mha_kvcache_( - out, - q, - k_cache, - v_cache, - seqlens_k, - block_table, - alibi_slopes_tensor, - scale); -} - -inline void bind_mha_kvcache(py::module &m) { - m.def( - "mha_kvcache", - &ops::py_mha_kvcache, - py::arg("q"), - py::arg("k_cache"), - py::arg("v_cache"), - py::arg("seqlens_k"), - py::arg("block_table"), - py::arg("alibi_slopes"), - py::arg("scale"), - R"doc(Flash attention KV-cache decode for single-step attention over a paged KV cache. - -Parameters ----------- -q : Tensor - Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] -k_cache : Tensor - Key cache tensor of shape [num_blocks, block_size, num_heads_k, head_size] (paged layout) -v_cache : Tensor - Value cache tensor of shape [num_blocks, block_size, num_heads_k, head_size] (paged layout) -seqlens_k : Tensor - Total KV length per request of shape [batch_size] (int32) -block_table : Tensor - Block mapping table of shape [batch_size, max_num_blocks_per_seq] (int32) -alibi_slopes : Optional[Tensor] - ALiBi slopes tensor, if None then ALiBi is disabled -scale : float - Scaling factor for attention scores (typically 1.0/sqrt(head_size)) - -Returns -------- -Tensor - Output tensor of shape [batch_size, seqlen_q, num_heads, head_size] -)doc"); - - m.def( - "mha_kvcache_", - &ops::py_mha_kvcache_, - py::arg("out"), - py::arg("q"), - py::arg("k_cache"), - py::arg("v_cache"), - py::arg("seqlens_k"), - py::arg("block_table"), - py::arg("alibi_slopes"), - py::arg("scale"), - R"doc(In-place flash attention KV-cache decode. - -Parameters ----------- -out : Tensor - Output tensor of shape [batch_size, seqlen_q, num_heads, head_size] -q : Tensor - Query tensor of shape [batch_size, seqlen_q, num_heads, head_size] -k_cache : Tensor - Key cache tensor of shape [num_blocks, block_size, num_heads_k, head_size] (paged layout) -v_cache : Tensor - Value cache tensor of shape [num_blocks, block_size, num_heads_k, head_size] (paged layout) -seqlens_k : Tensor - Total KV length per request of shape [batch_size] (int32) -block_table : Tensor - Block mapping table of shape [batch_size, max_num_blocks_per_seq] (int32) -alibi_slopes : Optional[Tensor] - ALiBi slopes tensor, if None then ALiBi is disabled -scale : float - Scaling factor for attention scores (typically 1.0/sqrt(head_size)) -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/mha_varlen.hpp b/src/infinicore/pybind11/ops/mha_varlen.hpp deleted file mode 100644 index b1b463a93..000000000 --- a/src/infinicore/pybind11/ops/mha_varlen.hpp +++ /dev/null @@ -1,110 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/mha_varlen.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_mha_varlen(Tensor q, - Tensor k, - Tensor v, - Tensor cum_seqlens_q, - Tensor cum_seqlens_k, - pybind11::object block_table, - int max_seqlen_q, - int max_seqlen_k, - pybind11::object alibi_slopes, - float scale) { - std::optional block_table_tensor = std::nullopt; - if (!block_table.is_none()) { - block_table_tensor = block_table.cast(); - } - std::optional alibi_slopes_tensor = std::nullopt; - if (!alibi_slopes.is_none()) { - alibi_slopes_tensor = alibi_slopes.cast(); - } - - return op::mha_varlen( - q, - k, - v, - cum_seqlens_q, - cum_seqlens_k, - block_table_tensor, - max_seqlen_q, - max_seqlen_k, - alibi_slopes_tensor, - scale); -} - -void py_mha_varlen_(Tensor out, - Tensor q, - Tensor k, - Tensor v, - Tensor cum_seqlens_q, - Tensor cum_seqlens_k, - pybind11::object block_table, - int max_seqlen_q, - int max_seqlen_k, - pybind11::object alibi_slopes, - float scale) { - std::optional block_table_tensor = std::nullopt; - if (!block_table.is_none()) { - block_table_tensor = block_table.cast(); - } - std::optional alibi_slopes_tensor = std::nullopt; - if (!alibi_slopes.is_none()) { - alibi_slopes_tensor = alibi_slopes.cast(); - } - - op::mha_varlen_( - out, - q, - k, - v, - cum_seqlens_q, - cum_seqlens_k, - block_table_tensor, - max_seqlen_q, - max_seqlen_k, - alibi_slopes_tensor, - scale); -} - -inline void bind_mha_varlen(py::module &m) { - m.def( - "mha_varlen", - &ops::py_mha_varlen, - py::arg("q"), - py::arg("k"), - py::arg("v"), - py::arg("cum_seqlens_q"), - py::arg("cum_seqlens_k"), - py::arg("block_table"), - py::arg("max_seqlen_q"), - py::arg("max_seqlen_k"), - py::arg("alibi_slopes"), - py::arg("scale"), - R"doc(Variable-length multi-head attention.)doc"); - - m.def( - "mha_varlen_", - &ops::py_mha_varlen_, - py::arg("out"), - py::arg("q"), - py::arg("k"), - py::arg("v"), - py::arg("cum_seqlens_q"), - py::arg("cum_seqlens_k"), - py::arg("block_table"), - py::arg("max_seqlen_q"), - py::arg("max_seqlen_k"), - py::arg("alibi_slopes"), - py::arg("scale"), - R"doc(In-place variable-length multi-head attention.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/moe_topk_softmax.hpp b/src/infinicore/pybind11/ops/moe_topk_softmax.hpp deleted file mode 100644 index 5466fcd79..000000000 --- a/src/infinicore/pybind11/ops/moe_topk_softmax.hpp +++ /dev/null @@ -1,57 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/moe_topk_softmax.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -std::tuple py_moe_topk_softmax(Tensor gating_output, - size_t topk, - bool renormalize, - float moe_softcapping, - py::object correction_bias) { - Tensor bias; - if (!correction_bias.is_none()) { - bias = correction_bias.cast(); - } - return op::moe_topk_softmax(gating_output, topk, renormalize, moe_softcapping, bias); -} - -void py_moe_topk_softmax_(Tensor topk_weights, - Tensor topk_indices, - Tensor gating_output, - py::object correction_bias, - bool renormalize, - float moe_softcapping) { - Tensor bias; - if (!correction_bias.is_none()) { - bias = correction_bias.cast(); - } - op::moe_topk_softmax_(topk_weights, topk_indices, gating_output, bias, renormalize, moe_softcapping); -} - -inline void bind_moe_topk_softmax(py::module &m) { - m.def("moe_topk_softmax", - &py_moe_topk_softmax, - py::arg("gating_output"), - py::arg("topk"), - py::arg("renormalize") = false, - py::arg("moe_softcapping") = 0.0f, - py::arg("correction_bias") = py::none(), - R"doc(MoE top-k softmax.)doc"); - - m.def("moe_topk_softmax_", - &py_moe_topk_softmax_, - py::arg("topk_weights"), - py::arg("topk_indices"), - py::arg("gating_output"), - py::arg("correction_bias") = py::none(), - py::arg("renormalize") = false, - py::arg("moe_softcapping") = 0.0f, - R"doc(In-place MoE top-k softmax.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/mrope.hpp b/src/infinicore/pybind11/ops/mrope.hpp deleted file mode 100644 index 765584880..000000000 --- a/src/infinicore/pybind11/ops/mrope.hpp +++ /dev/null @@ -1,46 +0,0 @@ -#pragma once - -#include -#include - -#include "infinicore/ops/mrope.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_mrope(py::module &m) { - m.def("mrope", - &op::mrope, - py::arg("q"), - py::arg("k"), - py::arg("cos"), - py::arg("sin"), - py::arg("positions"), - py::arg("head_size"), - py::arg("rotary_dim"), - py::arg("section_t"), - py::arg("section_h"), - py::arg("section_w"), - py::arg("interleaved"), - R"doc(Multimodal rotary position embedding for q and k.)doc"); - - m.def("mrope_", - &op::mrope_, - py::arg("q_out"), - py::arg("k_out"), - py::arg("q"), - py::arg("k"), - py::arg("cos"), - py::arg("sin"), - py::arg("positions"), - py::arg("head_size"), - py::arg("rotary_dim"), - py::arg("section_t"), - py::arg("section_h"), - py::arg("section_w"), - py::arg("interleaved"), - R"doc(In-place multimodal rotary position embedding for q and k.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/mul.hpp b/src/infinicore/pybind11/ops/mul.hpp deleted file mode 100644 index fb8e4144b..000000000 --- a/src/infinicore/pybind11/ops/mul.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/mul.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_mul(py::module &m) { - m.def("mul", - &op::mul, - py::arg("a"), - py::arg("b"), - R"doc(Element-wise multiplication of two tensors.)doc"); - - m.def("mul_", - &op::mul_, - py::arg("c"), - py::arg("a"), - py::arg("b"), - R"doc(In-place element-wise tensor multiplication.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/mul_scalar.hpp b/src/infinicore/pybind11/ops/mul_scalar.hpp deleted file mode 100644 index 2bce862bc..000000000 --- a/src/infinicore/pybind11/ops/mul_scalar.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/mul_scalar.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_mul_scalar(py::module &m) { - m.def("mul_scalar", - &op::mul_scalar, - py::arg("a"), - py::arg("alpha"), - R"doc(Multiply a tensor by a host scalar.)doc"); - - m.def("mul_scalar_", - &op::mul_scalar_, - py::arg("c"), - py::arg("a"), - py::arg("alpha"), - R"doc(Out-of-place tensor-scalar multiplication into c.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/multi_margin_loss.hpp b/src/infinicore/pybind11/ops/multi_margin_loss.hpp deleted file mode 100644 index f59fcec8e..000000000 --- a/src/infinicore/pybind11/ops/multi_margin_loss.hpp +++ /dev/null @@ -1,57 +0,0 @@ -#pragma once - -#include "infinicore/ops/multi_margin_loss.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_multi_margin_loss(py::module &m) { - m.def( - "multi_margin_loss", - [](const Tensor &input, const Tensor &target, py::object weight, int p, float margin, int reduction) { - Tensor weight_tensor; - if (!weight.is_none()) { - weight_tensor = weight.cast(); - } - return op::multi_margin_loss(input, target, weight_tensor, p, margin, reduction); - }, - py::arg("input"), - py::arg("target"), - py::arg("weight") = py::none(), // Python 端看到默认值是 None - py::arg("p") = 1, - py::arg("margin") = 1.0f, - py::arg("reduction") = 1, - R"doc(Computes the Multi Margin Loss between input and target. - - Args: - input (Tensor): Input tensor of shape (N, C). - target (Tensor): Ground truth labels of shape (N,). - weight (Tensor, optional): Manual rescaling weight given to each class. If given, has to be a Tensor of size C. - p (int, optional): The norm degree for pairwise distance. p=1 or p=2. Default: 1. - margin (float, optional): Margin value. Default: 1.0. - reduction (int, optional): Specifies the reduction to apply to the output: 0=None, 1=Mean, 2=Sum. Default: 1. - )doc"); - - m.def( - "multi_margin_loss_", - [](Tensor &output, const Tensor &input, const Tensor &target, py::object weight, int p, float margin, int reduction) { - Tensor weight_tensor; - if (!weight.is_none()) { - weight_tensor = weight.cast(); - } - // 调用底层 - op::multi_margin_loss_(output, input, target, weight_tensor, p, margin, reduction); - }, - py::arg("output"), - py::arg("input"), - py::arg("target"), - py::arg("weight") = py::none(), - py::arg("p") = 1, - py::arg("margin") = 1.0f, - py::arg("reduction") = 1, - R"doc(Explicit output Multi Margin Loss operation. Writes the result into the output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/nrm2.hpp b/src/infinicore/pybind11/ops/nrm2.hpp deleted file mode 100644 index 02b21f53b..000000000 --- a/src/infinicore/pybind11/ops/nrm2.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/nrm2.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_nrm2(py::module &m) { - m.def("nrm2", - &op::nrm2, - py::arg("x"), - R"doc(BLAS level-1 nrm2.)doc"); - - m.def("nrm2_", - &op::nrm2_, - py::arg("x"), - py::arg("result"), - R"doc(In-place BLAS level-1 nrm2.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/pad.hpp b/src/infinicore/pybind11/ops/pad.hpp deleted file mode 100644 index 19c090308..000000000 --- a/src/infinicore/pybind11/ops/pad.hpp +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/pad.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_pad(py::module &m) { - m.def("pad", - &op::pad, - py::arg("x"), - py::arg("pad"), - py::arg("mode") = std::string("constant"), - py::arg("value") = 0.0, - R"doc(Pad a tensor (PyTorch padding order).)doc"); - - m.def("pad_", - &op::pad_, - py::arg("y"), - py::arg("x"), - py::arg("pad"), - py::arg("mode") = std::string("constant"), - py::arg("value") = 0.0, - R"doc(Out variant of pad.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/paged_attention.hpp b/src/infinicore/pybind11/ops/paged_attention.hpp deleted file mode 100644 index ab77c87a4..000000000 --- a/src/infinicore/pybind11/ops/paged_attention.hpp +++ /dev/null @@ -1,53 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/paged_attention.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_paged_attention(Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, pybind11::object alibi_slopes, float scale) { - std::optional alibi_slopes_tensor = std::nullopt; - if (!alibi_slopes.is_none()) { - alibi_slopes_tensor = alibi_slopes.cast(); - } - return op::paged_attention(q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes_tensor, scale); -} - -void py_paged_attention_(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, pybind11::object alibi_slopes, float scale) { - std::optional alibi_slopes_tensor = std::nullopt; - if (!alibi_slopes.is_none()) { - alibi_slopes_tensor = alibi_slopes.cast(); - } - - op::paged_attention_(out, q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes_tensor, scale); -} - -inline void bind_paged_attention(py::module &m) { - m.def("paged_attention", - &ops::py_paged_attention, - py::arg("q"), - py::arg("k_cache"), - py::arg("v_cache"), - py::arg("block_tables"), - py::arg("cache_lens"), - py::arg("alibi_slopes"), - py::arg("scale"), - R"doc(Paged attention of query and key cache tensors.)doc"); - - m.def("paged_attention_", - &ops::py_paged_attention_, - py::arg("out"), - py::arg("q"), - py::arg("k_cache"), - py::arg("v_cache"), - py::arg("block_tables"), - py::arg("cache_lens"), - py::arg("alibi_slopes"), - py::arg("scale"), - R"doc(In-place paged attention of query and key cache tensors.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/paged_attention_prefill.hpp b/src/infinicore/pybind11/ops/paged_attention_prefill.hpp deleted file mode 100644 index 13b2c2683..000000000 --- a/src/infinicore/pybind11/ops/paged_attention_prefill.hpp +++ /dev/null @@ -1,69 +0,0 @@ -#pragma once - -#include "infinicore/ops/paged_attention_prefill.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_paged_attention_prefill(Tensor q, - Tensor k_cache, - Tensor v_cache, - Tensor block_tables, - Tensor history_lens, - Tensor cu_seqlens_q, - py::object alibi_slopes, - float scale) { - std::optional alibi_slopes_tensor = std::nullopt; - if (!alibi_slopes.is_none()) { - alibi_slopes_tensor = alibi_slopes.cast(); - } - return op::paged_attention_prefill( - q, k_cache, v_cache, block_tables, history_lens, cu_seqlens_q, alibi_slopes_tensor, scale); -} - -void py_paged_attention_prefill_(Tensor out, - Tensor q, - Tensor k_cache, - Tensor v_cache, - Tensor block_tables, - Tensor history_lens, - Tensor cu_seqlens_q, - py::object alibi_slopes, - float scale) { - std::optional alibi_slopes_tensor = std::nullopt; - if (!alibi_slopes.is_none()) { - alibi_slopes_tensor = alibi_slopes.cast(); - } - op::paged_attention_prefill_(out, q, k_cache, v_cache, block_tables, history_lens, cu_seqlens_q, alibi_slopes_tensor, scale); -} - -inline void bind_paged_attention_prefill(py::module &m) { - m.def("paged_attention_prefill", - &ops::py_paged_attention_prefill, - py::arg("q"), - py::arg("k_cache"), - py::arg("v_cache"), - py::arg("block_tables"), - py::arg("history_lens"), - py::arg("cu_seqlens_q"), - py::arg("alibi_slopes") = py::none(), - py::arg("scale") = 1.0, - R"doc(Paged attention prefill for packed variable-length queries.)doc"); - - m.def("paged_attention_prefill_", - &ops::py_paged_attention_prefill_, - py::arg("out"), - py::arg("q"), - py::arg("k_cache"), - py::arg("v_cache"), - py::arg("block_tables"), - py::arg("history_lens"), - py::arg("cu_seqlens_q"), - py::arg("alibi_slopes") = py::none(), - py::arg("scale") = 1.0, - R"doc(In-place paged attention prefill for packed variable-length queries.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/paged_caching.hpp b/src/infinicore/pybind11/ops/paged_caching.hpp deleted file mode 100644 index 4320b4eef..000000000 --- a/src/infinicore/pybind11/ops/paged_caching.hpp +++ /dev/null @@ -1,22 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/paged_caching.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_paged_caching(py::module &m) { - m.def("paged_caching_", - &op::paged_caching_, - py::arg("k_cache"), - py::arg("v_cache"), - py::arg("k"), - py::arg("v"), - py::arg("slot_mapping"), - R"doc(Paged caching of key and value tensors.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/per_channel_quant_i8.hpp b/src/infinicore/pybind11/ops/per_channel_quant_i8.hpp deleted file mode 100644 index da6f9f592..000000000 --- a/src/infinicore/pybind11/ops/per_channel_quant_i8.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/per_channel_quant_i8.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_per_channel_quant_i8(py::module &m) { - m.def("per_channel_quant_i8_", - &op::per_channel_quant_i8_, - py::arg("x"), - py::arg("x_packed"), - py::arg("x_scale"), - R"doc(Per-channel quantization of a tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/prelu.hpp b/src/infinicore/pybind11/ops/prelu.hpp deleted file mode 100644 index ee195c52f..000000000 --- a/src/infinicore/pybind11/ops/prelu.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/prelu.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_prelu(py::module &m) { - m.def("prelu", - &op::prelu, - py::arg("input"), - py::arg("weight"), - R"doc(Parametric ReLU.)doc"); - - m.def("prelu_", - &op::prelu_, - py::arg("out"), - py::arg("input"), - py::arg("weight"), - R"doc(In-place Parametric ReLU (writes to out).)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/random_sample.hpp b/src/infinicore/pybind11/ops/random_sample.hpp deleted file mode 100644 index e5beb9e4e..000000000 --- a/src/infinicore/pybind11/ops/random_sample.hpp +++ /dev/null @@ -1,32 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/random_sample.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_random_sample(py::module &m) { - m.def("random_sample", - &op::random_sample, - py::arg("logits"), - py::arg("random_val"), - py::arg("topp"), - py::arg("topk"), - py::arg("temperature"), - R"doc(Random sampling: returns an int32 scalar index.)doc"); - - m.def("random_sample_", - &op::random_sample_, - py::arg("indices"), - py::arg("logits"), - py::arg("random_val"), - py::arg("topp"), - py::arg("topk"), - py::arg("temperature"), - R"doc(In-place random sampling into provided int32 scalar tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/rearrange.hpp b/src/infinicore/pybind11/ops/rearrange.hpp deleted file mode 100644 index 816b00079..000000000 --- a/src/infinicore/pybind11/ops/rearrange.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/rearrange.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_rearrange(py::module &m) { - m.def("rearrange", - &op::rearrange, - py::arg("x"), - R"doc(Matrix rearrangement of a tensor.)doc"); - - m.def("rearrange_", - &op::rearrange_, - py::arg("y"), - py::arg("x"), - R"doc(In-place tensor rearrangement.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/reciprocal.hpp b/src/infinicore/pybind11/ops/reciprocal.hpp deleted file mode 100644 index 3e49205f6..000000000 --- a/src/infinicore/pybind11/ops/reciprocal.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/reciprocal.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_reciprocal(py::module &m) { - m.def("reciprocal", - &op::reciprocal, - py::arg("x"), - R"doc(Computes the reciprocal of the input tensor.)doc"); - - m.def("reciprocal_", - &op::reciprocal_, - py::arg("y"), - py::arg("x"), - R"doc(Computes the reciprocal of the input tensor and stores in the output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/recurrent_gated_delta_rule.hpp b/src/infinicore/pybind11/ops/recurrent_gated_delta_rule.hpp deleted file mode 100644 index 127b62042..000000000 --- a/src/infinicore/pybind11/ops/recurrent_gated_delta_rule.hpp +++ /dev/null @@ -1,37 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/recurrent_gated_delta_rule.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_recurrent_gated_delta_rule(py::module &m) { - m.def("recurrent_gated_delta_rule", - &op::recurrent_gated_delta_rule, - py::arg("q"), - py::arg("k"), - py::arg("v"), - py::arg("g"), - py::arg("beta"), - py::arg("initial_state"), - py::arg("use_qk_l2norm") = false, - R"doc(Recurrent gated delta rule. Returns out only.)doc"); - - m.def("recurrent_gated_delta_rule_indexed", - &op::recurrent_gated_delta_rule_indexed, - py::arg("q"), - py::arg("k"), - py::arg("v"), - py::arg("g"), - py::arg("beta"), - py::arg("initial_state"), - py::arg("initial_state_indices"), - py::arg("final_state_indices"), - py::arg("use_qk_l2norm") = false, - R"doc(Recurrent gated delta rule with indexed in-place state pool. Returns out only.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/relu6.hpp b/src/infinicore/pybind11/ops/relu6.hpp deleted file mode 100644 index e69e3568a..000000000 --- a/src/infinicore/pybind11/ops/relu6.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/relu6.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_relu6(py::module &m) { - m.def("relu6", - &op::relu6, - py::arg("input"), - R"doc(ReLU6 activation.)doc"); - - m.def("relu6_", - &op::relu6_, - py::arg("out"), - py::arg("input"), - R"doc(In-place ReLU6 activation (writes to out).)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/rms_norm.hpp b/src/infinicore/pybind11/ops/rms_norm.hpp deleted file mode 100644 index 1fd899c44..000000000 --- a/src/infinicore/pybind11/ops/rms_norm.hpp +++ /dev/null @@ -1,44 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/rms_norm.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_rms_norm(py::module &m) { - m.def("rms_norm", - &op::rms_norm, - py::arg("x"), - py::arg("weight"), - py::arg("epsilon") = 1e-5f, - R"doc(Root Mean Square Normalization. - -Args: - x: Input tensor - weight: Scale weights - epsilon: Small constant for numerical stability, default is 1e-5 - -Returns: - Normalized tensor with same shape as input -)doc"); - - m.def("rms_norm_", - &op::rms_norm_, - py::arg("y"), - py::arg("x"), - py::arg("weight"), - py::arg("epsilon") = 1e-5f, - R"doc(In-place Root Mean Square Normalization. - -Args: - y: Output tensor - x: Input tensor - weight: Scale weights - epsilon: Small constant for numerical stability, default is 1e-5 -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/rope.hpp b/src/infinicore/pybind11/ops/rope.hpp deleted file mode 100644 index a1836955a..000000000 --- a/src/infinicore/pybind11/ops/rope.hpp +++ /dev/null @@ -1,37 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/rope.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_rope(py::module &m) { - - py::enum_(m, "RoPEAlgo") - .value("GPT_J", infinicore::nn::RoPE::Algo::GPT_J) - .value("GPT_NEOX", infinicore::nn::RoPE::Algo::GPT_NEOX); - - m.def("rope", - &op::rope, - py::arg("x"), - py::arg("pos"), - py::arg("sin_table"), - py::arg("cos_table"), - py::arg("algo"), - R"doc( Rotary Position Embedding(RoPE).)doc"); - - m.def("rope_", - &op::rope_, - py::arg("x_out"), - py::arg("x"), - py::arg("pos"), - py::arg("sin_table"), - py::arg("cos_table"), - py::arg("algo"), - R"doc(In-place, Rotary Position Embedding(RoPE).)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/rot.hpp b/src/infinicore/pybind11/ops/rot.hpp deleted file mode 100644 index 359cb9745..000000000 --- a/src/infinicore/pybind11/ops/rot.hpp +++ /dev/null @@ -1,21 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/rot.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_rot(py::module &m) { - m.def("rot_", - &op::rot_, - py::arg("x"), - py::arg("y"), - py::arg("c"), - py::arg("s"), - R"doc(In-place BLAS level-1 rot, updating x and y.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/rotg.hpp b/src/infinicore/pybind11/ops/rotg.hpp deleted file mode 100644 index a37e79336..000000000 --- a/src/infinicore/pybind11/ops/rotg.hpp +++ /dev/null @@ -1,21 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/rotg.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_rotg(py::module &m) { - m.def("rotg_", - py::overload_cast(&op::rotg_), - py::arg("x"), - py::arg("y"), - py::arg("c"), - py::arg("s"), - R"doc(In-place BLAS level-1 rotg, updating x, y, c, and s.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/rotm.hpp b/src/infinicore/pybind11/ops/rotm.hpp deleted file mode 100644 index a88db38ba..000000000 --- a/src/infinicore/pybind11/ops/rotm.hpp +++ /dev/null @@ -1,20 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/rotm.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_rotm(py::module &m) { - m.def("rotm_", - py::overload_cast(&op::rotm_), - py::arg("x"), - py::arg("y"), - py::arg("param"), - R"doc(In-place BLAS level-1 rotm, updating x and y.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/rotmg.hpp b/src/infinicore/pybind11/ops/rotmg.hpp deleted file mode 100644 index 72b816751..000000000 --- a/src/infinicore/pybind11/ops/rotmg.hpp +++ /dev/null @@ -1,22 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/rotmg.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_rotmg(py::module &m) { - m.def("rotmg_", - py::overload_cast(&op::rotmg_), - py::arg("d1"), - py::arg("d2"), - py::arg("x1"), - py::arg("y1"), - py::arg("param"), - R"doc(In-place BLAS level-1 rotmg, updating d1, d2, x1, and param.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/rwkv5_wkv.hpp b/src/infinicore/pybind11/ops/rwkv5_wkv.hpp deleted file mode 100644 index ec59316d0..000000000 --- a/src/infinicore/pybind11/ops/rwkv5_wkv.hpp +++ /dev/null @@ -1,34 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/rwkv5_wkv.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_rwkv5_wkv(py::module &m) { - m.def("rwkv5_wkv", - &op::rwkv5_wkv, - py::arg("receptance"), - py::arg("key"), - py::arg("value"), - py::arg("time_decay"), - py::arg("time_faaaa"), - py::arg("state"), - R"doc(RWKV5 weighted key-value recurrence. Updates state in-place and returns output.)doc"); - - m.def("rwkv5_wkv_", - &op::rwkv5_wkv_, - py::arg("out"), - py::arg("receptance"), - py::arg("key"), - py::arg("value"), - py::arg("time_decay"), - py::arg("time_faaaa"), - py::arg("state"), - R"doc(Explicit-output RWKV5 weighted key-value recurrence. Updates out and state in-place.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/scal.hpp b/src/infinicore/pybind11/ops/scal.hpp deleted file mode 100644 index 75a2c5ca8..000000000 --- a/src/infinicore/pybind11/ops/scal.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/scal.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_scal(py::module &m) { - m.def("scal_", - &op::scal_, - py::arg("alpha"), - py::arg("x"), - R"doc(In-place BLAS level-1 scal, updating x.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/scaled_mm_i8.hpp b/src/infinicore/pybind11/ops/scaled_mm_i8.hpp deleted file mode 100644 index c3d46d9df..000000000 --- a/src/infinicore/pybind11/ops/scaled_mm_i8.hpp +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/scaled_mm_i8.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_scaled_mm_i8(py::module &m) { - m.def("scaled_mm_i8", - &op::scaled_mm_i8, - py::arg("a_p"), - py::arg("a_s"), - py::arg("b_p"), - py::arg("b_s"), - py::arg("bias"), - R"doc(Scaled matrix multiplication of two tensors.)doc"); - - m.def("scaled_mm_i8_", - &op::scaled_mm_i8_, - py::arg("a"), - py::arg("b"), - py::arg("a_scale"), - py::arg("b_scale"), - R"doc(In-place Scaled matrix multiplication of two tensors.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/scatter.hpp b/src/infinicore/pybind11/ops/scatter.hpp deleted file mode 100644 index 950ed49e8..000000000 --- a/src/infinicore/pybind11/ops/scatter.hpp +++ /dev/null @@ -1,56 +0,0 @@ -#pragma once - -#include "infinicore/ops/scatter.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_scatter(py::module &m) { - // ========================================================================= - // 1. 绑定 out-of-place 接口: scatter - // ========================================================================= - // 为了匹配测试脚本的行为（将所有 Tensor 作为位置参数传入，属性作为 kwargs 传入）， - // 我们将参数顺序调整为: input, index, src, dim, reduction - // ========================================================================= - m.def( - "scatter", - [](const Tensor &input, const Tensor &index, const Tensor &src, int64_t dim, int64_t reduction) { - // 调用底层 C++ 实现时，必须恢复正确的参数顺序: (input, dim, index, src, reduction) - return op::scatter(input, dim, index, src, reduction); - }, - py::arg("input"), - py::arg("index"), - py::arg("src"), - py::arg("dim"), // 关键修改：将 dim 移到 Tensor 参数之后 - py::arg("reduction") = 0, - R"doc( - Scatter operator. - Note: Parameter order in this binding is adapted for the test runner: (input, index, src, dim, reduction). - )doc"); - - // ========================================================================= - // 2. 绑定 in-place 接口: scatter_ - // ========================================================================= - // 参数顺序调整为: output, input, index, src, dim, reduction - // ========================================================================= - m.def( - "scatter_", - [](Tensor &output, const Tensor &input, const Tensor &index, const Tensor &src, int64_t dim, int64_t reduction) { - // 调用底层 C++ 实现 - op::scatter_(output, input, dim, index, src, reduction); - }, - py::arg("output"), - py::arg("input"), - py::arg("index"), - py::arg("src"), - py::arg("dim"), // 关键修改：将 dim 移到 Tensor 参数之后 - py::arg("reduction") = 0, - R"doc( - In-place Scatter operator. - Writes result into output. - )doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/selu.hpp b/src/infinicore/pybind11/ops/selu.hpp deleted file mode 100644 index ea7ffe553..000000000 --- a/src/infinicore/pybind11/ops/selu.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/selu.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_selu(py::module &m) { - m.def( - "selu", - &op::selu, - py::arg("input"), - R"doc(SELU activation function.)doc"); - - m.def( - "selu_", - &op::selu_, - py::arg("output"), - py::arg("input"), - R"doc(In-place SELU activation function.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/sigmoid.hpp b/src/infinicore/pybind11/ops/sigmoid.hpp deleted file mode 100644 index 8f24c0c8e..000000000 --- a/src/infinicore/pybind11/ops/sigmoid.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/sigmoid.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_sigmoid(py::module &m) { - m.def("sigmoid", - &op::sigmoid, - py::arg("input"), - R"doc(Out-of-place Sigmoid activation.)doc"); - - m.def("sigmoid_", - &op::sigmoid_, - py::arg("output"), - py::arg("input"), - R"doc(In-place Sigmoid activation.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/silu.hpp b/src/infinicore/pybind11/ops/silu.hpp deleted file mode 100644 index 52aa7c3a8..000000000 --- a/src/infinicore/pybind11/ops/silu.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/silu.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_silu(py::module &m) { - m.def("silu", - &op::silu, - py::arg("input"), - R"doc(SiLU (Swish) activation function.)doc"); - - m.def("silu_", - &op::silu_, - py::arg("output"), - py::arg("input"), - R"doc(In-place SiLU (Swish) activation function.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/silu_and_mul.hpp b/src/infinicore/pybind11/ops/silu_and_mul.hpp deleted file mode 100644 index 009c3b533..000000000 --- a/src/infinicore/pybind11/ops/silu_and_mul.hpp +++ /dev/null @@ -1,29 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/silu_and_mul.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_silu_and_mul(py::module &m) { - m.def("silu_and_mul", - &op::silu_and_mul, - py::arg("input"), - R"doc( - SiLU and Mul (SwiGLU) activation function. - Input should be [..., 2*d], output will be [..., d]. - )doc"); - - m.def("silu_and_mul_", - &op::silu_and_mul_, - py::arg("output"), - py::arg("input"), - R"doc( - In-place or destination-specified SiLU and Mul (SwiGLU) activation function. - )doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/sinh.hpp b/src/infinicore/pybind11/ops/sinh.hpp deleted file mode 100644 index 69ff34e99..000000000 --- a/src/infinicore/pybind11/ops/sinh.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/sinh.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_sinh(py::module &m) { - m.def( - "sinh", - &op::sinh, - py::arg("input"), - R"doc(Sinh activation function.)doc"); - - m.def( - "sinh_", - &op::sinh_, - py::arg("output"), - py::arg("input"), - R"doc(In-place Sinh activation function.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/smooth_l1_loss.hpp b/src/infinicore/pybind11/ops/smooth_l1_loss.hpp deleted file mode 100644 index 4402b1b4b..000000000 --- a/src/infinicore/pybind11/ops/smooth_l1_loss.hpp +++ /dev/null @@ -1,36 +0,0 @@ -#pragma once - -#include "infinicore/ops/smooth_l1_loss.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_smooth_l1_loss(py::module &m) { - // 1. 绑定 out-of-place 接口: output = smooth_l1_loss(input, target, beta, reduction) - m.def("smooth_l1_loss", - &op::smooth_l1_loss, - py::arg("input"), - py::arg("target"), - py::arg("beta") = 1.0f, - py::arg("reduction") = 1, - R"doc(Computes the Smooth L1 Loss between input and target. - - Args: - input (Tensor): Predicted values. - target (Tensor): Ground truth values. - beta (float, optional): The threshold at which to change between L1 and L2 loss. Default: 1.0. - reduction (int, optional): Specifies the reduction to apply to the output: 0=None, 1=Mean, 2=Sum. Default: 1. - )doc"); - m.def("smooth_l1_loss_", - &op::smooth_l1_loss_, - py::arg("output"), - py::arg("input"), - py::arg("target"), - py::arg("beta") = 1.0f, - py::arg("reduction") = 1, - R"doc(Explicit output Smooth L1 Loss operation. Writes the result into the output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/softplus.hpp b/src/infinicore/pybind11/ops/softplus.hpp deleted file mode 100644 index 768f99f38..000000000 --- a/src/infinicore/pybind11/ops/softplus.hpp +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/softplus.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_softplus(py::module &m) { - // Functional interface: returns a new Tensor - m.def("softplus", - &op::softplus, - py::arg("x"), - py::arg("beta") = 1.0f, - py::arg("threshold") = 20.0f, - R"doc(Computes the softplus function element-wise: y = 1/beta * log(1 + exp(beta * x)).)doc"); - - // In-place/Out-variant interface: writes to provided output Tensor - m.def("softplus_", - &op::softplus_, - py::arg("y"), - py::arg("x"), - py::arg("beta") = 1.0f, - py::arg("threshold") = 20.0f, - R"doc(In-place softplus activation. Writes result into y.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/softsign.hpp b/src/infinicore/pybind11/ops/softsign.hpp deleted file mode 100644 index 1191abe02..000000000 --- a/src/infinicore/pybind11/ops/softsign.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/softsign.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_softsign(py::module &m) { - m.def("softsign", - &op::softsign, - py::arg("x"), - R"doc(Softsign activation function.)doc"); - - m.def("softsign_", - &op::softsign_, - py::arg("y"), - py::arg("x"), - R"doc(In-place softsign activation.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/sum.hpp b/src/infinicore/pybind11/ops/sum.hpp deleted file mode 100644 index 50fef7539..000000000 --- a/src/infinicore/pybind11/ops/sum.hpp +++ /dev/null @@ -1,60 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/sum.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_sum(Tensor input, py::object dim, bool keepdim) { - if (dim.is_none()) { - std::vector dim_vec; - for (int i = 0; i < input->shape().size(); i++) { - dim_vec.push_back(i); - } - return op::sum(input, dim_vec, keepdim); - } else if (py::isinstance(dim) || py::isinstance(dim)) { - return op::sum(input, dim.cast>(), keepdim); - } else if (py::isinstance(dim)) { - return op::sum(input, std::vector(1, dim.cast()), keepdim); - } else { - throw std::invalid_argument("dim must be a tuple or an integer"); - } -} - -void py_sum_(Tensor output, Tensor input, py::object dim, bool keepdim) { - if (dim.is_none()) { - std::vector dim_vec; - for (int i = 0; i < input->shape().size(); i++) { - dim_vec.push_back(i); - } - op::sum_(output, input, dim_vec, keepdim); - } else if (py::isinstance(dim) || py::isinstance(dim)) { - op::sum_(output, input, dim.cast>(), keepdim); - } else if (py::isinstance(dim)) { - op::sum_(output, input, std::vector(1, dim.cast()), keepdim); - } else { - throw std::invalid_argument("dim must be a tuple or an integer"); - } -} - -inline void bind_sum(py::module &m) { - m.def("sum", - &py_sum, - py::arg("input"), - py::arg("dim"), - py::arg("keepdim"), - R"doc(Sum of input tensor along the given dimensions.)doc"); - - m.def("sum_", - &py_sum_, - py::arg("output"), - py::arg("input"), - py::arg("dim"), - py::arg("keepdim"), - R"doc(In-place tensor sum.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/swap.hpp b/src/infinicore/pybind11/ops/swap.hpp deleted file mode 100644 index 0b8f2ae3e..000000000 --- a/src/infinicore/pybind11/ops/swap.hpp +++ /dev/null @@ -1,19 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/swap.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_swap(py::module &m) { - m.def("swap_", - &op::swap_, - py::arg("x"), - py::arg("y"), - R"doc(In-place BLAS level-1 swap.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/swiglu.hpp b/src/infinicore/pybind11/ops/swiglu.hpp deleted file mode 100644 index a5c98609f..000000000 --- a/src/infinicore/pybind11/ops/swiglu.hpp +++ /dev/null @@ -1,26 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/swiglu.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_swiglu(py::module &m) { - m.def("swiglu", - &op::swiglu, - py::arg("a"), - py::arg("b"), - R"doc(SwiGLU activation function.)doc"); - - m.def("swiglu_", - &op::swiglu_, - py::arg("c"), - py::arg("a"), - py::arg("b"), - R"doc(In-place SwiGLU activation function.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/take.hpp b/src/infinicore/pybind11/ops/take.hpp deleted file mode 100644 index c6660cf22..000000000 --- a/src/infinicore/pybind11/ops/take.hpp +++ /dev/null @@ -1,25 +0,0 @@ -#pragma once - -#include "infinicore/ops/take.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_take(py::module &m) { - m.def("take", - &op::take, - py::arg("input"), - py::arg("indices"), - R"doc(Extracts elements from the input tensor along the given indices. -The input tensor is treated as a flattened 1D array.)doc"); - m.def("take_", - &op::take_, - py::arg("output"), - py::arg("input"), - py::arg("indices"), - R"doc(Explicit output take operation. Writes the result into the output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/tan.hpp b/src/infinicore/pybind11/ops/tan.hpp deleted file mode 100644 index acd62e048..000000000 --- a/src/infinicore/pybind11/ops/tan.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/tan.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_tan(py::module &m) { - m.def("tan", - &op::tan, - py::arg("input"), - R"doc(opertor: torch.tan, out-of-place mode)doc"); - - m.def("tan_", - &op::tan_, - py::arg("output"), - py::arg("input"), - R"doc(opertor: torch.tan, in-place mode)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/tanhshrink.hpp b/src/infinicore/pybind11/ops/tanhshrink.hpp deleted file mode 100644 index 3780555f4..000000000 --- a/src/infinicore/pybind11/ops/tanhshrink.hpp +++ /dev/null @@ -1,24 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/tanhshrink.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_tanhshrink(py::module &m) { - m.def("tanhshrink", - &op::tanhshrink, - py::arg("input"), - R"doc(opertor: torch.tanhshrink, out-of-place mode)doc"); - - m.def("tanhshrink_", - &op::tanhshrink_, - py::arg("output"), - py::arg("input"), - R"doc(opertor: torch.tanhshrink, in-place mode)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/topk.hpp b/src/infinicore/pybind11/ops/topk.hpp deleted file mode 100644 index 1341f39fa..000000000 --- a/src/infinicore/pybind11/ops/topk.hpp +++ /dev/null @@ -1,54 +0,0 @@ -#pragma once - -#include -#include // 添加这行 - -#include "infinicore/ops/topk.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -std::pair py_topk(Tensor input, size_t k, int dim, bool largest, bool sorted) { - if (dim == -1) { - return op::topk(input, k, input->ndim() - 1, largest, sorted); - } else if (dim >= 0) { - return op::topk(input, k, static_cast(dim), largest, sorted); - } else { - throw std::invalid_argument("invalid argument: dim"); - } -} - -void py_topk_(Tensor values_output, Tensor indices_output, Tensor input, size_t k, int dim, bool largest, bool sorted) { - if (dim == -1) { - op::topk_(values_output, indices_output, input, k, input->ndim() - 1, largest, sorted); - } else if (dim >= 0) { - op::topk_(values_output, indices_output, input, k, static_cast(dim), largest, sorted); - } else { - throw std::invalid_argument("invalid argument: dim"); - } -} - -inline void bind_topk(py::module &m) { - m.def("topk", - &py_topk, - py::arg("input"), - py::arg("k"), - py::arg("dim"), - py::arg("largest"), - py::arg("sorted"), - R"doc(topk of input tensor along the given dimensions.)doc"); - - m.def("topk_", - &py_topk_, - py::arg("values_output"), - py::arg("indices_output"), - py::arg("input"), - py::arg("k"), - py::arg("dim"), - py::arg("largest"), - py::arg("sorted"), - R"doc(In-place tensor topk_.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/topksoftmax.hpp b/src/infinicore/pybind11/ops/topksoftmax.hpp deleted file mode 100644 index ca93edff0..000000000 --- a/src/infinicore/pybind11/ops/topksoftmax.hpp +++ /dev/null @@ -1,32 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/topksoftmax.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_topksoftmax(py::module &m) { - m.def("topksoftmax", - &op::topksoftmax, - py::arg("values"), - py::arg("indices"), - py::arg("x"), - py::arg("topk"), - py::arg("norm") = 0, - R"doc(In-place Top-k Softmax. - -Writes results to pre-allocated values and indices tensors. - -Args: - values: Output tensor for softmax weights [N, topk] - indices: Output tensor for selected indices [N, topk], int32 - x: Input tensor [N, width], router logits - topk: Number of top values to select - norm: Whether to re-normalize top-k probabilities (1=yes, 0=no), default 0 -)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/triplet_margin_loss.hpp b/src/infinicore/pybind11/ops/triplet_margin_loss.hpp deleted file mode 100644 index f003bac94..000000000 --- a/src/infinicore/pybind11/ops/triplet_margin_loss.hpp +++ /dev/null @@ -1,50 +0,0 @@ -#pragma once - -#include "infinicore/ops/triplet_margin_loss.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_triplet_margin_loss(py::module &m) { - // 1. 绑定 functional 接口: output = triplet_margin_loss(anchor, positive, negative, ...) - m.def("triplet_margin_loss", - &op::triplet_margin_loss, - py::arg("anchor"), - py::arg("positive"), - py::arg("negative"), - py::arg("margin") = 1.0f, - py::arg("p") = 2, - py::arg("eps") = 1e-6f, - py::arg("swap") = false, - py::arg("reduction") = 1, - R"doc(Computes the triplet margin loss. - - Args: - anchor (Tensor): The anchor tensor. - positive (Tensor): The positive tensor. - negative (Tensor): The negative tensor. - margin (float): Default: 1.0. - p (int): The norm degree for pairwise distance. Default: 2. - eps (float): Small constant for numerical stability. Default: 1e-6. - swap (bool): The distance swap is described in the paper Learning shallow convolutional feature descriptors with triplet losses. Default: False. - reduction (int): Specifies the reduction to apply to the output: 0 (none), 1 (mean), 2 (sum). Default: 1. - )doc"); - - // 2. 绑定 explicit output 接口: triplet_margin_loss_(output, anchor, positive, negative, ...) - m.def("triplet_margin_loss_", - &op::triplet_margin_loss_, - py::arg("output"), - py::arg("anchor"), - py::arg("positive"), - py::arg("negative"), - py::arg("margin") = 1.0f, - py::arg("p") = 2, - py::arg("eps") = 1e-6f, - py::arg("swap") = false, - py::arg("reduction") = 1, - R"doc(Explicit output TripletMarginLoss operation. Writes the result into the output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/triplet_margin_with_distance_loss.hpp b/src/infinicore/pybind11/ops/triplet_margin_with_distance_loss.hpp deleted file mode 100644 index e4b20208a..000000000 --- a/src/infinicore/pybind11/ops/triplet_margin_with_distance_loss.hpp +++ /dev/null @@ -1,41 +0,0 @@ -#pragma once - -#include "infinicore/ops/triplet_margin_with_distance_loss.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_triplet_margin_with_distance_loss(py::module &m) { - m.def("triplet_margin_with_distance_loss", - &op::triplet_margin_with_distance_loss, - py::arg("anchor"), - py::arg("positive"), - py::arg("negative"), - py::arg("margin") = 1.0, - py::arg("swap") = false, - py::arg("reduction") = 1, - R"doc(Computes the triplet margin loss with distance. - - Args: - anchor (Tensor): The anchor input tensor. - positive (Tensor): The positive input tensor. - negative (Tensor): The negative input tensor. - margin (float, optional): Default: 1.0. - swap (bool, optional): The distance swap is described in the paper Learning shallow convolutional feature descriptors with triplet losses. Default: False. - reduction (int, optional): Specifies the reduction to apply to the output: 0 (None), 1 (Mean), 2 (Sum). Default: 1. - )doc"); - m.def("triplet_margin_with_distance_loss_", - &op::triplet_margin_with_distance_loss_, - py::arg("output"), - py::arg("anchor"), - py::arg("positive"), - py::arg("negative"), - py::arg("margin"), - py::arg("swap"), - py::arg("reduction"), - R"doc(Explicit output TripletMarginWithDistanceLoss operation. Writes results into output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/unfold.hpp b/src/infinicore/pybind11/ops/unfold.hpp deleted file mode 100644 index b09d2ff70..000000000 --- a/src/infinicore/pybind11/ops/unfold.hpp +++ /dev/null @@ -1,61 +0,0 @@ -#pragma once - -#include "infinicore/ops/unfold.hpp" -#include -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_unfold(py::module &m) { - // ------------------------------------------------------------------------- - // 1. 绑定函数式接口 (unfold) - // ------------------------------------------------------------------------- - m.def( - "unfold", - [](const Tensor &input, - std::vector kernel_sizes, - std::vector dilations, - std::vector paddings, - std::vector strides) { - return op::unfold(input, kernel_sizes, dilations, paddings, strides); - }, - py::arg("input"), - py::arg("kernel_sizes"), - py::arg("dilations"), - py::arg("paddings"), - py::arg("strides"), - R"doc(Extracts sliding local blocks from a batched input tensor. - - Args: - input (Tensor): The input tensor. - kernel_sizes (List[int]): The size of the sliding blocks. - dilations (List[int]): The parameter that controls the stride of elements within the neighborhood. - paddings (List[int]): Implicit zero padding to be added on both sides of input. - strides (List[int]): The stride of the sliding blocks. - )doc"); - - // ------------------------------------------------------------------------- - // 2. 绑定 in-place 接口 (unfold_) - // ------------------------------------------------------------------------- - m.def( - "unfold_", - [](Tensor &output, - const Tensor &input, - std::vector kernel_sizes, - std::vector dilations, - std::vector paddings, - std::vector strides) { - op::unfold_(output, input, kernel_sizes, dilations, paddings, strides); - }, - py::arg("output"), - py::arg("input"), - py::arg("kernel_sizes"), - py::arg("dilations"), - py::arg("paddings"), - py::arg("strides"), - R"doc(Explicit output Unfold operation. Writes the result into the output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/upsample_bilinear.hpp b/src/infinicore/pybind11/ops/upsample_bilinear.hpp deleted file mode 100644 index 4a7619ffa..000000000 --- a/src/infinicore/pybind11/ops/upsample_bilinear.hpp +++ /dev/null @@ -1,35 +0,0 @@ -#pragma once - -#include "infinicore/ops/upsample_bilinear.hpp" -#include -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_upsample_bilinear(py::module &m) { - // 1. 绑定 functional 接口: output = upsample_bilinear(input, output_size, align_corners) - m.def("upsample_bilinear", - &op::upsample_bilinear, - py::arg("input"), - py::arg("output_size"), - py::arg("align_corners") = false, - R"doc(Upsample the input using bilinear interpolation. - - Args: - input (Tensor): The input tensor. - output_size (List[int]): The output spatial size (e.g. [H_out, W_out]). - align_corners (bool): If True, the corner pixels of the input and output tensors are aligned. - )doc"); - - // 2. 绑定 explicit output 接口: upsample_bilinear_(output, input, align_corners) - m.def("upsample_bilinear_", - &op::upsample_bilinear_, - py::arg("output"), - py::arg("input"), - py::arg("align_corners") = false, - R"doc(Explicit output UpsampleBilinear operation. Writes the result into the output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/upsample_nearest.hpp b/src/infinicore/pybind11/ops/upsample_nearest.hpp deleted file mode 100644 index 7e4154e9b..000000000 --- a/src/infinicore/pybind11/ops/upsample_nearest.hpp +++ /dev/null @@ -1,32 +0,0 @@ -#pragma once - -#include "infinicore/ops/upsample_nearest.hpp" -#include -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_upsample_nearest(py::module &m) { - // 1. 绑定 functional 接口: output = upsample_nearest(input, output_size) - m.def("upsample_nearest", - &op::upsample_nearest, - py::arg("input"), - py::arg("output_size"), - R"doc(Upsample the input using nearest neighbor interpolation. - - Args: - input (Tensor): The input tensor. - output_size (List[int]): The output spatial size (e.g. [H_out, W_out]). - )doc"); - - // 2. 绑定 explicit output 接口: upsample_nearest_(output, input) - m.def("upsample_nearest_", - &op::upsample_nearest_, - py::arg("output"), - py::arg("input"), - R"doc(Explicit output UpsampleNearest operation. Writes the result into the output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/vander.hpp b/src/infinicore/pybind11/ops/vander.hpp deleted file mode 100644 index 89a851f08..000000000 --- a/src/infinicore/pybind11/ops/vander.hpp +++ /dev/null @@ -1,42 +0,0 @@ -#pragma once - -#include "infinicore/ops/vander.hpp" -#include - -namespace py = pybind11; - -namespace infinicore::ops { - -inline void bind_vander(py::module &m) { - m.def( - "vander", - [](const Tensor &input, int64_t N, bool increasing) { - return op::vander(input, N, increasing); - }, - py::arg("input"), - py::arg("N") = 0, - py::arg("increasing") = false, - R"doc(Generates a Vandermonde matrix. - - Args: - input (Tensor): 1-D input tensor. - N (int, optional): Number of columns in the output. If 0, defaults to input size (square matrix). Default: 0. - increasing (bool, optional): Order of the powers. If True, powers increase (x^0, x^1...). Default: False. - )doc"); - - // ------------------------------------------------------------------------- - // 2. 绑定 in-place 接口 (vander_) - // ------------------------------------------------------------------------- - m.def( - "vander_", - [](Tensor &output, const Tensor &input, int64_t N, bool increasing) { - op::vander_(output, input, N, increasing); - }, - py::arg("output"), - py::arg("input"), - py::arg("N") = 0, - py::arg("increasing") = false, - R"doc(Explicit output Vander operation. Writes the result into the output tensor.)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/var.hpp b/src/infinicore/pybind11/ops/var.hpp deleted file mode 100644 index 9668fef5f..000000000 --- a/src/infinicore/pybind11/ops/var.hpp +++ /dev/null @@ -1,62 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/var.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -Tensor py_var(Tensor input, py::object dim, bool unbiased, bool keepdim) { - if (dim.is_none()) { - std::vector dim_vec; - for (int i = 0; i < input->shape().size(); i++) { - dim_vec.push_back(i); - } - return op::var(input, dim_vec, unbiased, keepdim); - } else if (py::isinstance(dim) || py::isinstance(dim)) { - return op::var(input, dim.cast>(), unbiased, keepdim); - } else if (py::isinstance(dim)) { - return op::var(input, std::vector(1, dim.cast()), unbiased, keepdim); - } else { - throw std::invalid_argument("dim must be a tuple or an integer"); - } -} - -void py_var_(Tensor var_output, Tensor input, py::object dim, bool unbiased, bool keepdim) { - if (dim.is_none()) { - std::vector dim_vec; - for (int i = 0; i < input->shape().size(); i++) { - dim_vec.push_back(i); - } - op::var_(var_output, input, dim_vec, unbiased, keepdim); - } else if (py::isinstance(dim) || py::isinstance(dim)) { - op::var_(var_output, input, dim.cast>(), unbiased, keepdim); - } else if (py::isinstance(dim)) { - op::var_(var_output, input, std::vector(1, dim.cast()), unbiased, keepdim); - } else { - throw std::invalid_argument("dim must be a list/tuple or an integer"); - } -} - -inline void bind_var(py::module &m) { - m.def("var", - &py_var, - py::arg("input"), - py::arg("dim"), - py::arg("unbiased"), - py::arg("keepdim"), - R"doc(Var of input tensor along the given dimensions.)doc"); - - m.def("var_", - &py_var_, - py::arg("var_output"), - py::arg("input"), - py::arg("dim"), - py::arg("unbiased"), - py::arg("keepdim"), - R"doc(In-place tensor Var .)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/ops/var_mean.hpp b/src/infinicore/pybind11/ops/var_mean.hpp deleted file mode 100644 index 986ec49f7..000000000 --- a/src/infinicore/pybind11/ops/var_mean.hpp +++ /dev/null @@ -1,63 +0,0 @@ -#pragma once - -#include - -#include "infinicore/ops/var_mean.hpp" - -namespace py = pybind11; - -namespace infinicore::ops { - -std::pair py_var_mean(Tensor input, py::object dim, bool unbiased, bool keepdim) { - if (dim.is_none()) { - std::vector dim_vec; - for (int i = 0; i < input->shape().size(); i++) { - dim_vec.push_back(i); - } - return op::var_mean(input, dim_vec, unbiased, keepdim); - } else if (py::isinstance(dim) || py::isinstance(dim)) { - return op::var_mean(input, dim.cast>(), unbiased, keepdim); - } else if (py::isinstance(dim)) { - return op::var_mean(input, std::vector(1, dim.cast()), unbiased, keepdim); - } else { - throw std::invalid_argument("dim must be a tuple or an integer"); - } -} - -void py_var_mean_(Tensor var_output, Tensor mean_output, Tensor input, py::object dim, bool unbiased, bool keepdim) { - if (dim.is_none()) { - std::vector dim_vec; - for (int i = 0; i < input->shape().size(); i++) { - dim_vec.push_back(i); - } - op::var_mean_(var_output, mean_output, input, dim_vec, unbiased, keepdim); - } else if (py::isinstance(dim) || py::isinstance(dim)) { - op::var_mean_(var_output, mean_output, input, dim.cast>(), unbiased, keepdim); - } else if (py::isinstance(dim)) { - op::var_mean_(var_output, mean_output, input, std::vector(1, dim.cast()), unbiased, keepdim); - } else { - throw std::invalid_argument("dim must be a list/tuple or an integer"); - } -} - -inline void bind_var_mean(py::module &m) { - m.def("var_mean", - &py_var_mean, - py::arg("input"), - py::arg("dim"), - py::arg("unbiased"), - py::arg("keepdim"), - R"doc(Var & Mean of input tensor along the given dimensions.)doc"); - - m.def("var_mean_", - &py_var_mean_, - py::arg("var_output"), - py::arg("mean_output"), - py::arg("input"), - py::arg("dim"), - py::arg("unbiased"), - py::arg("keepdim"), - R"doc(In-place tensor Var & Mean .)doc"); -} - -} // namespace infinicore::ops diff --git a/src/infinicore/pybind11/tensor.hpp b/src/infinicore/pybind11/tensor.hpp deleted file mode 100644 index 8f117591b..000000000 --- a/src/infinicore/pybind11/tensor.hpp +++ /dev/null @@ -1,100 +0,0 @@ -#pragma once - -#include "infinicore.hpp" -#include -#include - -namespace py = pybind11; - -namespace infinicore::tensor { - -Tensor from_list_py(py::handle data, DataType dtype); - -inline void bind(py::module &m) { - py::class_(m, "Tensor") - .def_property_readonly("shape", [](const Tensor &tensor) { return tensor->shape(); }) - .def_property_readonly("strides", [](const Tensor &tensor) { return tensor->strides(); }) - .def_property_readonly("ndim", [](const Tensor &tensor) { return tensor->ndim(); }) - .def_property_readonly("dtype", [](const Tensor &tensor) { return tensor->dtype(); }) - .def_property_readonly("device", [](const Tensor &tensor) { return tensor->device(); }) - .def("data_ptr", [](const Tensor &tensor) { return reinterpret_cast(tensor->data()); }) - .def("size", [](const Tensor &tensor, std::size_t dim) { return tensor->size(dim); }) - .def("stride", [](const Tensor &tensor, std::size_t dim) { return tensor->stride(dim); }) - .def("numel", [](const Tensor &tensor) { return tensor->numel(); }) - .def("is_contiguous", [](const Tensor &tensor) { return tensor->is_contiguous(); }) - .def("is_pinned", [](const Tensor &tensor) { return tensor->is_pinned(); }) - .def("info", [](const Tensor &tensor) { return tensor->info(); }) - - .def("debug", [](const Tensor &tensor) { return tensor->debug(); }) - .def("debug", [](const Tensor &tensor, const std::string &filename) { return tensor->debug(filename); }) - - .def("copy_", [](Tensor &tensor, const Tensor &other) { tensor->copy_from(other); }) - .def("to", [](const Tensor &tensor, const Device &device) { return tensor->to(device); }) - .def("contiguous", [](const Tensor &tensor) { return tensor->contiguous(); }) - - .def("as_strided", [](const Tensor &tensor, const Shape &shape, const Strides &strides) { return tensor->as_strided(shape, strides); }) - .def("narrow", [](const Tensor &tensor, std::size_t dim, std::size_t start, std::size_t length) { return tensor->narrow({{dim, start, length}}); }) - .def("permute", [](const Tensor &tensor, const Shape &dims) { return tensor->permute(dims); }) - .def("view", [](const Tensor &tensor, const Shape &shape) { return tensor->view(shape); }) - .def("unsqueeze", [](const Tensor &tensor, std::size_t dim) { return tensor->unsqueeze(dim); }) - .def("squeeze", [](const Tensor &tensor, std::size_t dim) { return tensor->squeeze(dim); }) - .def("reset", static_cast(&Tensor::reset)) - .def("use_count", &Tensor::use_count) - .def("__str__", [](const Tensor &tensor) { - std::ostringstream oss; - oss << tensor; - return oss.str(); - }) - .def("__repr__", [](const Tensor &tensor) { - std::ostringstream oss; - oss << tensor; - return oss.str(); - }) - .def("__bool__", [](const Tensor &tensor) { - return bool(tensor); - }); - - using EmptyFuncType = Tensor (*)(const Shape &, const DataType &, const Device &, bool); - using StridedEmptyFuncType = Tensor (*)(const Shape &, const Strides &, const DataType &, const Device &, bool); - - m.def("empty", static_cast(&Tensor::empty), - py::arg("shape"), - py::arg("dtype"), - py::arg("device"), - py::arg("pin_memory") = false); - - m.def("strided_empty", static_cast(&Tensor::strided_empty), - py::arg("shape"), - py::arg("strides"), - py::arg("dtype"), - py::arg("device"), - py::arg("pin_memory") = false); - - m.def("zeros", static_cast(&Tensor::zeros), - py::arg("shape"), - py::arg("dtype"), - py::arg("device"), - py::arg("pin_memory") = false); - - m.def("ones", static_cast(&Tensor::ones), - py::arg("shape"), - py::arg("dtype"), - py::arg("device"), - py::arg("pin_memory") = false); - - m.def( - "from_blob", [](uintptr_t raw_ptr, Shape &shape, const DataType &dtype, const Device &device) { - return Tensor{infinicore::Tensor::from_blob(reinterpret_cast(raw_ptr), shape, dtype, device)}; - }, - pybind11::arg("raw_ptr"), pybind11::arg("shape"), pybind11::arg("dtype"), pybind11::arg("device")); - - m.def( - "strided_from_blob", [](uintptr_t raw_ptr, Shape &shape, Strides &strides, const DataType &dtype, const Device &device) { - return Tensor{infinicore::Tensor::strided_from_blob(reinterpret_cast(raw_ptr), shape, strides, dtype, device)}; - }, - pybind11::arg("raw_ptr"), pybind11::arg("shape"), pybind11::arg("strides"), pybind11::arg("dtype"), pybind11::arg("device")); - - m.def("from_list", &from_list_py, py::arg("data"), py::arg("dtype")); -} - -} // namespace infinicore::tensor diff --git a/src/infinicore/tensor/copy.cc b/src/infinicore/tensor/copy.cc deleted file mode 100644 index 9db8bebe8..000000000 --- a/src/infinicore/tensor/copy.cc +++ /dev/null @@ -1,67 +0,0 @@ -#include "infinicore/context/context.hpp" -#include "infinicore/dtype.hpp" -#include "infinicore/ops.hpp" -#include "infinicore/tensor.hpp" - -#include -#include -#include -namespace infinicore { -Tensor TensorImpl::to(Device device) const { - if (device == data_.memory->device()) { - return Tensor(const_cast(this)->shared_from_this()); - } else { - std::shared_ptr _t = empty(meta_.shape, meta_.dtype, device); - _t->copy_from(Tensor(const_cast(this)->shared_from_this())); - return Tensor(_t); - } -} - -void TensorImpl::copy_from(Tensor src) { - if (src->shape() != this->shape()) { - throw std::runtime_error( - "Cannot copy from tensor with different shape. Src: " + src->info() + " Dst: " + this->info()); - } - if (this->device() == src->device()) { - op::rearrange_(Tensor(const_cast(this)->shared_from_this()), src); - } else { - if (!src->is_contiguous()) { - src = src->contiguous(); - } - - // Use nbytes() to get the actual tensor size, not the full memory size - size_t copy_size = std::min(this->nbytes(), src->nbytes()); - if (this->device().getType() == Device::Type::CPU) { - if (this->is_contiguous()) { - context::setDevice(src->device()); - context::memcpyD2H(this->data(), src->data(), copy_size); - } else { - auto local_src = Tensor::empty(this->shape(), this->dtype(), this->device()); - context::setDevice(src->device()); - context::memcpyD2H(local_src->data(), src->data(), copy_size); - op::rearrange_(Tensor(const_cast(this)->shared_from_this()), local_src); - } - } else if (src->device().getType() == Device::Type::CPU) { - context::setDevice(this->device()); - // Keep Ascend H2D synchronous because copy_from does not retain the host source. - const bool async = this->device().getType() != Device::Type::ASCEND; - if (this->is_contiguous()) { - context::memcpyH2D(this->data(), src->data(), copy_size, async); - } else { - auto local_src = Tensor::empty(this->shape(), this->dtype(), this->device()); - context::memcpyH2D(local_src->data(), src->data(), copy_size, async); - op::rearrange_(Tensor(const_cast(this)->shared_from_this()), local_src); - } - } - } -} - -Tensor TensorImpl::contiguous() const { - if (is_contiguous()) { - return Tensor(const_cast(this)->shared_from_this()); - } else { - return op::rearrange(Tensor(const_cast(this)->shared_from_this())); - } -} - -} // namespace infinicore diff --git a/src/infinicore/tensor/debug.cc b/src/infinicore/tensor/debug.cc deleted file mode 100644 index 9b4e97777..000000000 --- a/src/infinicore/tensor/debug.cc +++ /dev/null @@ -1,233 +0,0 @@ -#include "infinicore/context/context.hpp" -#include "infinicore/dtype.hpp" -#include "infinicore/tensor.hpp" - -#include -#include -#include -#include -#include -#include - -namespace infinicore { - -inline float f16_to_f32(uint16_t h) { - uint32_t sign = (h & 0x8000) << 16; - int32_t exponent = (h >> 10) & 0x1F; - uint32_t mantissa = h & 0x3FF; - - uint32_t f32; - if (exponent == 31) { - if (mantissa != 0) { - f32 = sign | 0x7F800000 | (mantissa << 13); - } else { - f32 = sign | 0x7F800000; - } - } else if (exponent == 0) { - if (mantissa == 0) { - f32 = sign; - } else { - exponent = -14; - while ((mantissa & 0x400) == 0) { - mantissa <<= 1; - exponent--; - } - mantissa &= 0x3FF; - f32 = sign | ((exponent + 127) << 23) | (mantissa << 13); - } - } else { - f32 = sign | ((exponent + 127 - 15) << 23) | (mantissa << 13); - } - - float result; - std::memcpy(&result, &f32, sizeof(result)); - return result; -} - -inline float bf16_to_f32(uint16_t val) { - uint32_t bits32 = static_cast(val) << 16; - float out; - std::memcpy(&out, &bits32, sizeof(out)); - return out; -} - -// Template function for printing data recursively -template -void print_data(const T *data, const Shape &shape, const Strides &strides, size_t dim) { - if (dim == shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - std::cout << data[i * strides[dim]] << " "; - } - std::cout << std::endl; - } else if (dim < shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - print_data(data + i * strides[dim], shape, strides, dim + 1); - } - } -} - -// Specialization for F16 (uint16_t) -template <> -void print_data(const uint16_t *data, const Shape &shape, const Strides &strides, size_t dim) { - if (dim == shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - std::cout << f16_to_f32(data[i * strides[dim]]) << " "; - } - std::cout << std::endl; - } else if (dim < shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - print_data(data + i * strides[dim], shape, strides, dim + 1); - } - } -} - -// Function for printing BF16 data -void print_data_bf16(const uint16_t *data, const Shape &shape, const Strides &strides, size_t dim) { - if (dim == shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - std::cout << bf16_to_f32(data[i * strides[dim]]) << " "; - } - std::cout << std::endl; - } else if (dim < shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - print_data_bf16(data + i * strides[dim], shape, strides, dim + 1); - } - } -} - -// Function for printing I8 data -void print_data_i8(const int8_t *data, const Shape &shape, const Strides &strides, size_t dim) { - if (dim == shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - std::cout << static_cast(data[i * strides[dim]]) << " "; - } - std::cout << std::endl; - } else if (dim < shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - print_data_i8(data + i * strides[dim], shape, strides, dim + 1); - } - } -} - -// Function for printing U8 data -void print_data_u8(const uint8_t *data, const Shape &shape, const Strides &strides, size_t dim) { - if (dim == shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - std::cout << static_cast(data[i * strides[dim]]) << " "; - } - std::cout << std::endl; - } else if (dim < shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - print_data_u8(data + i * strides[dim], shape, strides, dim + 1); - } - } -} - -// Template function for writing data recursively to binary file (handles non-contiguous tensors) -template -void write_binary_data(std::ofstream &out, const T *data, const Shape &shape, const Strides &strides, size_t dim) { - if (dim == shape.size() - 1) { - // Write the innermost dimension - for (size_t i = 0; i < shape[dim]; i++) { - out.write(reinterpret_cast(&data[i * strides[dim]]), sizeof(T)); - } - } else { - // Recursively process higher dimensions - for (size_t i = 0; i < shape[dim]; i++) { - write_binary_data(out, data + i * strides[dim], shape, strides, dim + 1); - } - } -} - -void TensorImpl::debug(const std::string &filename) const { - // Synchronize device if needed - context::syncDevice(); - std::cout << info() << std::endl; - std::unique_ptr allocated_memory; // RAII: 自动管理内存 - auto cpu_tensor = this->contiguous()->to(Device::cpu()); - const std::byte *cpu_data = cpu_tensor->data(); - // If filename is provided, save to binary file - if (!filename.empty()) { - std::ofstream outFile(filename, std::ios::binary); - if (!outFile) { - std::cerr << "Error opening file for writing: " << filename << "\n"; - return; // allocated_memory 会自动释放（RAII） - } - // Fast path: contiguous tensor, write in one go - size_t mem_size = cpu_tensor->numel() * dsize(cpu_tensor->dtype()); - outFile.write(reinterpret_cast(cpu_data), mem_size); - // 显式关闭文件并检查是否成功 - outFile.close(); - if (!outFile) { - std::cerr << "Error: Failed to write data to file: " << filename << "\n"; - return; - } - std::cout << "Data written to binary file: " << filename; - std::cout << "\n"; - return; - } - // Print data based on dtype - switch (cpu_tensor->dtype()) { - case DataType::F16: - print_data(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - case DataType::F32: - print_data(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - case DataType::F64: - print_data(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - case DataType::U64: - print_data(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - case DataType::I64: - print_data(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - case DataType::U32: - print_data(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - case DataType::I32: - print_data(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - case DataType::U16: - print_data(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - case DataType::I16: - print_data(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - case DataType::U8: - print_data_u8(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - case DataType::I8: - print_data_i8(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - case DataType::BF16: - print_data_bf16(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - case DataType::BOOL: - print_data(reinterpret_cast(cpu_data), - cpu_tensor->shape(), cpu_tensor->strides(), 0); - break; - default: - std::cout << "Unsupported data type for debug" << std::endl; - break; - } -} - -void TensorImpl::debug() const { - this->debug(""); -} - -} // namespace infinicore diff --git a/src/infinicore/tensor/tensor.cc b/src/infinicore/tensor/tensor.cc deleted file mode 100644 index e7e54aa59..000000000 --- a/src/infinicore/tensor/tensor.cc +++ /dev/null @@ -1,302 +0,0 @@ -#include "infinicore/tensor.hpp" -#include "../context/internal.hpp" -#include "../utils.hpp" -#include "infinicore/context/context.hpp" -#include "infinicore/dtype.hpp" - -#include - -namespace { -// Helper function to calculate contiguous strides -inline infinicore::Strides calculate_contiguous_strides(const infinicore::Shape &shape) { - infinicore::Strides strides(shape.size()); - infinicore::Stride stride = 1; - for (int i = shape.size() - 1; i >= 0; --i) { - strides[i] = stride; - stride *= shape[i]; - } - return strides; -} -} // namespace - -namespace infinicore { -TensorImpl *Tensor::operator->() { return impl_.get(); } - -const TensorImpl *Tensor::operator->() const { return impl_.get(); } - -Tensor Tensor::empty(const Shape &shape, - const DataType &dtype, - const Device &device, - bool pin_memory) { - return Tensor{TensorImpl::empty(shape, dtype, device, pin_memory)}; -} - -Tensor Tensor::strided_empty(const Shape &shape, - const Strides &strides, - const DataType &dtype, - const Device &device, - bool pin_memory) { - return Tensor{TensorImpl::strided_empty(shape, strides, dtype, device, pin_memory)}; -} - -Tensor Tensor::zeros(const Shape &shape, - const DataType &dtype, - const Device &device, - bool pin_memory) { - return Tensor{TensorImpl::zeros(shape, dtype, device, pin_memory)}; -} - -Tensor Tensor::ones(const Shape &shape, - const DataType &dtype, - const Device &device, - bool pin_memory) { - return Tensor{TensorImpl::ones(shape, dtype, device, pin_memory)}; -} - -Tensor Tensor::from_blob(void *raw_ptr, const Shape &shape, const DataType &dtype, const Device &device) { - return Tensor{TensorImpl::from_blob(raw_ptr, shape, dtype, device)}; -} - -Tensor Tensor::strided_from_blob(void *raw_ptr, const Shape &shape, const Strides &strides, const DataType &dtype, const Device &device) { - return Tensor{TensorImpl::strided_from_blob(raw_ptr, shape, strides, dtype, device)}; -} - -Tensor::operator bool() const { - return impl_ != nullptr; -} - -TensorMetaData::TensorMetaData(const Shape &_shape, const Strides &_strides, const DataType &_dtype) - : shape(_shape), strides(_strides), dtype(_dtype) { - INFINICORE_CHECK_ERROR(infiniopCreateTensorDescriptor(&desc, shape.size(), shape.data(), strides.data(), (infiniDtype_t)dtype)); -} - -TensorMetaData::~TensorMetaData() { - if (desc) { - infiniopDestroyTensorDescriptor(desc); - desc = nullptr; - } -} - -TensorImpl::TensorImpl(const Shape &shape, const DataType &dtype) - : meta_(TensorMetaData(shape, calculate_contiguous_strides(shape), dtype)) {} - -TensorImpl::TensorImpl(const Shape &shape, const Strides &strides, const DataType &dtype) - : meta_(TensorMetaData(shape, strides, dtype)) {} - -std::byte *TensorImpl::data() { - return data_.memory->data() + data_.offset; -} - -const std::byte *TensorImpl::data() const { - return data_.memory->data() + data_.offset; -} - -const Shape &TensorImpl::shape() const { - return meta_.shape; -} - -const Strides &TensorImpl::strides() const { - return meta_.strides; -} - -Size TensorImpl::ndim() const { - return meta_.shape.size(); -} - -bool TensorImpl::is_contiguous() const { - Stride expected_stride = 1; - for (int i = meta_.shape.size() - 1; i >= 0; --i) { - if (meta_.strides[i] != expected_stride) { - return false; - } - expected_stride *= meta_.shape[i]; - } - return true; -} - -Size TensorImpl::numel() const { - Size total = 1; - for (const auto &dim : meta_.shape) { - total *= dim; - } - return total; -} - -size_t TensorImpl::element_size() const { - return dsize(dtype()); -} - -size_t TensorImpl::nbytes() const { - return numel() * element_size(); -} - -Size TensorImpl::size(size_t dim) const { - return meta_.shape[dim]; -} - -Stride TensorImpl::stride(size_t dim) const { - return meta_.strides[dim]; -} - -DataType TensorImpl::dtype() const { - return meta_.dtype; -} - -Device TensorImpl::device() const { - return data_.memory->device(); -} - -infiniopTensorDescriptor_t TensorImpl::desc() const { - return meta_.desc; -} - -bool TensorImpl::is_pinned() const { - return data_.memory->is_pinned(); -} - -std::string TensorImpl::info() const { - std::stringstream ss; - - ss << "Tensor: " - << "shape[ "; - for (auto s : this->shape()) { - ss << s << " "; - } - ss << "] strides[ "; - for (auto s : this->strides()) { - ss << s << " "; - } - ss << "] dtype=" << toString(this->dtype()); - ss << " device=" << this->device().toString(); - - return ss.str(); -} - -std::shared_ptr TensorImpl::empty(const Shape &shape, - const DataType &dtype, - const Device &device, - bool pin_memory) { - auto t = std::shared_ptr(new TensorImpl(shape, dtype)); - t->data_.offset = 0; - - if (device == Device::Type::CPU) { - if (pin_memory) { - if (context::getDevice() == Device::Type::CPU) { - spdlog::warn("Tensor memory is not pinned by any device with CPU runtime."); - t->data_.memory = context::allocateHostMemory(t->numel() * dsize(dtype)); - } else { - t->data_.memory = context::allocatePinnedHostMemory(t->numel() * dsize(dtype)); - } - } else { - t->data_.memory = context::allocateHostMemory(t->numel() * dsize(dtype)); - } - } else { - context::setDevice(device); - t->data_.memory = context::allocateMemory(t->numel() * dsize(dtype)); - } - - return t; -} - -std::shared_ptr TensorImpl::strided_empty( - const Shape &shape, - const Strides &strides, - const DataType &dtype, - const Device &device, - bool pin_memory) { - - auto impl = std::shared_ptr(new TensorImpl(shape, strides, dtype)); - impl->data_.offset = 0; - - size_t max_offset = 0; - - for (size_t i = 0; i < shape.size(); ++i) { - if (shape[i] > 0) { - max_offset += (shape[i] - 1) * strides[i]; - } - } - - size_t required_elements = max_offset + 1; - size_t required_bytes = required_elements * dsize(dtype); - - if (device == Device::Type::CPU) { - if (pin_memory) { - if (context::getDevice() == Device::Type::CPU) { - spdlog::warn("Tensor memory is not pinned by any device with CPU runtime."); - impl->data_.memory = context::allocateHostMemory(required_bytes); - } else { - impl->data_.memory = context::allocatePinnedHostMemory(required_bytes); - } - } else { - impl->data_.memory = context::allocateHostMemory(required_bytes); - } - } else { - context::setDevice(device); - impl->data_.memory = context::allocateMemory(required_bytes); - } - - return impl; -} - -std::shared_ptr TensorImpl::zeros(const Shape &shape, - const DataType &dtype, - const Device &device, - bool pin_memory) { - - auto result = empty(shape, dtype, device, pin_memory); - context::setDeviceMemoryAsync(result->data(), 0, result->nbytes(), context::getStream()); - return result; -} - -std::shared_ptr TensorImpl::ones(const Shape &shape, - const DataType &dtype, - const Device &device, - bool pin_memory) { - // TODO: Implement this. - return empty(shape, dtype, device, pin_memory); -} - -std::shared_ptr TensorImpl::from_blob( - void *raw_ptr, - const Shape &shape, - const DataType &dtype, - const Device &device) { - auto t = std::shared_ptr(new TensorImpl(shape, dtype)); - t->data_.offset = 0; - t->data_.memory = std::make_shared((std::byte *)raw_ptr, t->numel() * dsize(dtype), device, nullptr); - return t; -} - -std::shared_ptr TensorImpl::strided_from_blob( - void *raw_ptr, - const Shape &shape, - const Strides &strides, - const DataType &dtype, - const Device &device) { - auto t = std::shared_ptr(new TensorImpl(shape, strides, dtype)); - t->data_.offset = 0; - t->data_.memory = std::make_shared((std::byte *)raw_ptr, t->numel() * dsize(dtype), device, nullptr); - return t; -} - -Tensor TensorImpl::to_blob_() const { - auto t = std::shared_ptr(new TensorImpl(shape(), strides(), dtype())); - t->data_.offset = this->data_.offset; - t->data_.memory = std::make_shared(this->data_.memory->data(), this->data_.memory->size(), this->data_.memory->device(), nullptr); - t->to_blob_mark_ = true; - return Tensor{t}; -} - -Tensor TensorImpl::resume_from_blob_() const { - auto t = std::shared_ptr(new TensorImpl(shape(), strides(), dtype())); - t->data_.offset = this->data_.offset; - if (to_blob_mark_) { - t->data_.memory = context::reinstantiateBlob(this->data_.memory); - } else { - t->data_.memory = this->data_.memory; - } - - return Tensor{t}; -} - -} // namespace infinicore diff --git a/src/infinicore/tensor/view.cc b/src/infinicore/tensor/view.cc deleted file mode 100644 index 051ee42c0..000000000 --- a/src/infinicore/tensor/view.cc +++ /dev/null @@ -1,159 +0,0 @@ -#include "infinicore/context/context.hpp" -#include "infinicore/dtype.hpp" -#include "infinicore/tensor.hpp" - -#include "../utils.hpp" - -#include -#include - -namespace infinicore { -Tensor TensorImpl::squeeze(size_t dim) const { - // Create new shape with dimension of size one removed at dim - if (meta_.shape[dim] != 1) { - spdlog::error("Dimension {} is not of size 1 for squeeze operation on {}.", dim, this->info()); - throw std::runtime_error("Invalid squeeze operation on tensor."); - } - Shape new_shape = meta_.shape; - new_shape.erase(new_shape.begin() + dim); - Strides new_strides = meta_.strides; - new_strides.erase(new_strides.begin() + dim); - - auto tensor_impl = std::make_shared(new_shape, new_strides, meta_.dtype); - tensor_impl->data_ = data_; - - return Tensor(tensor_impl); -} - -Tensor TensorImpl::unsqueeze(size_t dim) const { - // Create new shape with dimension of size one inserted at dim - Shape new_shape = meta_.shape; - new_shape.insert(new_shape.begin() + dim, 1); - - // Create new strides with stride of zero for the new dimension - Strides new_strides = meta_.strides; - new_strides.insert(new_strides.begin() + dim, 0); - - auto tensor_impl = std::make_shared(new_shape, new_strides, meta_.dtype); - tensor_impl->data_ = data_; - - return Tensor(tensor_impl); -} - -Tensor TensorImpl::narrow(const std::vector &slices) const { - // Create new shape and calculate offset - Shape new_shape = meta_.shape; - size_t offset = data_.offset; - - for (const auto &slice : slices) { - if (meta_.shape[slice.dim] < slice.start + slice.len) { - spdlog::error("Invalid slice [dim={}, start={}, len={}] on {}.", slice.dim, slice.start, slice.len, this->info()); - throw std::runtime_error("Invalid slice on tensor."); - } - new_shape[slice.dim] = slice.len; - offset += slice.start * meta_.strides[slice.dim] * dsize(meta_.dtype); - } - - // Create new tensor with the same strides but narrowed shape - auto tensor_impl = std::make_shared(new_shape, meta_.strides, meta_.dtype); - tensor_impl->data_.offset = offset; - tensor_impl->data_.memory = data_.memory; - - return Tensor(tensor_impl); -} - -Tensor TensorImpl::permute(const Shape &order) const { - // Validate input - INFINICORE_ASSERT(meta_.shape.size() == order.size()); - - // Check that order contains all indices from 0 to n-1 exactly once - for (size_t i = 0; i < order.size(); i++) { - INFINICORE_ASSERT(std::find(order.begin(), order.end(), i) != order.end()); - } - - // Permute shape and strides - Shape new_shape(order.size()); - Strides new_strides(order.size()); - - for (size_t i = 0; i < order.size(); i++) { - new_shape[i] = meta_.shape[order[i]]; - new_strides[i] = meta_.strides[order[i]]; - } - - auto tensor_impl = std::make_shared(new_shape, new_strides, meta_.dtype); - tensor_impl->data_ = data_; - - return Tensor(tensor_impl); -} - -Tensor TensorImpl::view(const Shape &new_shape) const { - // Step 1: Validate total size - Size numel = 1; - for (Size dim : meta_.shape) { - numel *= dim; - } - - Size new_numel = 1; - for (Size dim : new_shape) { - new_numel *= dim; - } - - assert(numel == new_numel); - - // Step 2: Get current shape and strides - const Shape &old_shape = meta_.shape; - const Strides &old_strides = meta_.strides; - - // Step 3: Create merged shape and strides - Shape merged_shape; - Strides merged_strides; - - if (!old_shape.empty()) { - merged_shape.push_back(old_shape[0]); - merged_strides.push_back(old_strides[0]); - - for (size_t i = 1; i < old_shape.size(); ++i) { - if (old_strides[i] * static_cast(old_shape[i]) == merged_strides.back()) { - merged_shape.back() *= old_shape[i]; - merged_strides.back() = old_strides[i]; - } else { - merged_shape.push_back(old_shape[i]); - merged_strides.push_back(old_strides[i]); - } - } - } - - // Step 4: Compute new strides by splitting merged dimensions - Strides new_strides(new_shape.size()); - size_t merged_idx = 0; - Stride current_stride = merged_strides[0]; - Size remaining_size = merged_shape[0]; - - for (size_t i = 0; i < new_shape.size(); ++i) { - // Find which merged dimension contains this new dimension - while (new_shape[i] > remaining_size) { - if (++merged_idx >= merged_shape.size()) { - throw std::runtime_error("Incompatible shape for view operation."); - } - current_stride = merged_strides[merged_idx]; - remaining_size = merged_shape[merged_idx]; - } - - if (remaining_size % new_shape[i] != 0) { - throw std::runtime_error("Incompatible shape for view operation."); - }; - - new_strides[i] = current_stride * (remaining_size / new_shape[i]); - remaining_size /= new_shape[i]; - } - - return this->as_strided(new_shape, new_strides); -} - -Tensor TensorImpl::as_strided(const Shape &new_shape, const Strides &new_strides) const { - auto tensor_impl = std::make_shared(new_shape, new_strides, meta_.dtype); - tensor_impl->data_ = data_; - - return Tensor(tensor_impl); -} -} // namespace infinicore diff --git a/src/infinicore/utils.hpp b/src/infinicore/utils.hpp deleted file mode 100644 index cf8e69789..000000000 --- a/src/infinicore/utils.hpp +++ /dev/null @@ -1,63 +0,0 @@ -#pragma once - -#include "../utils/infini_status_string.h" - -#include -#include -#include - -inline struct SpdlogInitializer { - SpdlogInitializer() { - if (!std::getenv("INFINICORE_LOG_LEVEL")) { - spdlog::set_level(spdlog::level::info); - } else { - spdlog::cfg::load_env_levels("INFINICORE_LOG_LEVEL"); - } - // Set pattern for logging - // Using SPDLOG_* macros enables source location support (%s and %#) - // Format: [timestamp] [level] [file:line] message - spdlog::set_pattern("[%Y-%m-%d %H:%M:%S.%e] [%^%l%$] [%s:%#] %v"); - } -} spdlog_initializer; - -#define STRINGIZE_(x) #x -#define STRINGIZE(x) STRINGIZE_(x) - -#define INFINICORE_CHECK_ERROR(call) \ - do { \ - SPDLOG_DEBUG("Entering `" #call "` at `" __FILE__ ":" STRINGIZE(__LINE__) "`."); \ - infiniStatus_t ret = (call); \ - SPDLOG_DEBUG("Exiting `" #call "` at `" __FILE__ ":" STRINGIZE(__LINE__) "`."); \ - if (ret != INFINI_STATUS_SUCCESS) { \ - throw std::runtime_error("`" #call "` failed with error: " + std::string(infini_status_string(ret)) \ - + " from " + std::string(__func__) \ - + " at " + std::string(__FILE__) \ - + ":" + std::to_string(__LINE__) + "."); \ - } \ - } while (false) - -#define INFINICORE_ASSERT_TENSORS_SAME_DEVICE(FIRST___, ...) \ - do { \ - const auto &first_device___ = (FIRST___)->device(); \ - for (const auto &tensor___ : {__VA_ARGS__}) { \ - if (first_device___ != (tensor___)->device()) { \ - throw std::runtime_error("Tensor devices mismatch " \ - + first_device___.toString() + " vs " \ - + (tensor___)->device().toString() \ - + " from " + std::string(__func__) \ - + " at " + std::string(__FILE__) \ - + ":" + std::to_string(__LINE__) + "."); \ - } \ - } \ - } while (0) - -#define INFINICORE_ASSERT(CONDITION__) \ - do { \ - if (!(CONDITION__)) { \ - SPDLOG_ERROR( \ - "Assertion `{}` failed from {} at {}:{}", \ - #CONDITION__, __func__, __FILE__, __LINE__); \ - throw std::runtime_error( \ - std::string("Assertion `") + #CONDITION__ + "` failed from " + __func__ + " at " + __FILE__ + ":" + std::to_string(__LINE__)); \ - } \ - } while (0) diff --git a/src/infiniop-test/include/file_mapping.hpp b/src/infiniop-test/include/file_mapping.hpp deleted file mode 100644 index ebbcca6fe..000000000 --- a/src/infiniop-test/include/file_mapping.hpp +++ /dev/null @@ -1,31 +0,0 @@ -#ifndef __INFINIOPTEST_FILE_MAPPING_HPP__ -#define __INFINIOPTEST_FILE_MAPPING_HPP__ - -#ifdef _WIN32 // windows -#include -#else // linux -#include -#include -#include -#include -#endif - -#include -#include -#include - -class FileMapping { -private: - void *_ptr; - size_t _size; -#ifdef _WIN32 - HANDLE _file_handle = NULL; - HANDLE _file_mapping = NULL; -#endif -public: - FileMapping(const std::string &filepath); - ~FileMapping(); - void *ptr() const; - size_t size() const; -}; -#endif // __INFINIOPTEST_FILE_MAPPING_HPP__ diff --git a/src/infiniop-test/include/gguf.hpp b/src/infiniop-test/include/gguf.hpp deleted file mode 100644 index 0c0414616..000000000 --- a/src/infiniop-test/include/gguf.hpp +++ /dev/null @@ -1,250 +0,0 @@ -#ifndef __INFINIOPTEST_GGUF_HPP__ -#define __INFINIOPTEST_GGUF_HPP__ -#include "file_mapping.hpp" -#include -#include -#include -#include -#include -#include - -#ifdef _WIN32 -#include -#endif - -typedef enum { - GGUF_TYPE_UINT8 = 0, - GGUF_TYPE_INT8 = 1, - GGUF_TYPE_UINT16 = 2, - GGUF_TYPE_INT16 = 3, - GGUF_TYPE_UINT32 = 4, - GGUF_TYPE_INT32 = 5, - GGUF_TYPE_FLOAT32 = 6, - GGUF_TYPE_BOOL = 7, - GGUF_TYPE_STRING = 8, - GGUF_TYPE_ARRAY = 9, - GGUF_TYPE_UINT64 = 10, - GGUF_TYPE_INT64 = 11, - GGUF_TYPE_FLOAT64 = 12, - GGUF_TYPE_COUNT, // marks the end of the enum -} GGUF_TYPE; - -constexpr const char *GGUF_TYPE_NAME[GGUF_TYPE_COUNT] = { - "GGUF_TYPE_UINT8", - "GGUF_TYPE_INT8", - "GGUF_TYPE_UINT16", - "GGUF_TYPE_INT16", - "GGUF_TYPE_UINT32", - "GGUF_TYPE_INT32", - "GGUF_TYPE_FLOAT32", - "GGUF_TYPE_BOOL", - "GGUF_TYPE_STRING", - "GGUF_TYPE_ARRAY", - "GGUF_TYPE_UINT64", - "GGUF_TYPE_INT64", - "GGUF_TYPE_FLOAT64", -}; - -struct gguf_str { - uint64_t n; - char *data; -}; - -static const size_t GGUF_TYPE_SIZE[GGUF_TYPE_COUNT] = { - sizeof(uint8_t), // GGUF_TYPE_UINT8 - sizeof(int8_t), // GGUF_TYPE_INT8 - sizeof(uint16_t), // GGUF_TYPE_UINT16 - sizeof(int16_t), // GGUF_TYPE_INT16 - sizeof(uint32_t), // GGUF_TYPE_UINT32 - sizeof(int32_t), // GGUF_TYPE_INT32 - sizeof(float), // GGUF_TYPE_FLOAT32 - sizeof(bool), // GGUF_TYPE_BOOL - sizeof(gguf_str), // GGUF_TYPE_STRING - 0, // GGUF_TYPE_ARRAY (undefined) - sizeof(uint64_t), // GGUF_TYPE_UINT64 - sizeof(int64_t), // GGUF_TYPE_INT64 - sizeof(double), // GGUF_TYPE_FLOAT64 -}; - -inline std::string ggufDataToString(const uint8_t *data, GGUF_TYPE gguf_type) { - switch (gguf_type) { - -#define RETURN_GGUF_DATA(CASE, CTYPE) \ - case CASE: \ - return std::to_string(*reinterpret_cast(data)); - - RETURN_GGUF_DATA(GGUF_TYPE_UINT8, uint8_t) - RETURN_GGUF_DATA(GGUF_TYPE_INT8, int8_t) - RETURN_GGUF_DATA(GGUF_TYPE_UINT16, uint16_t) - RETURN_GGUF_DATA(GGUF_TYPE_INT16, int16_t) - RETURN_GGUF_DATA(GGUF_TYPE_UINT32, uint32_t) - RETURN_GGUF_DATA(GGUF_TYPE_INT32, int32_t) - RETURN_GGUF_DATA(GGUF_TYPE_FLOAT32, float) - RETURN_GGUF_DATA(GGUF_TYPE_BOOL, bool) - RETURN_GGUF_DATA(GGUF_TYPE_UINT64, uint64_t) - RETURN_GGUF_DATA(GGUF_TYPE_INT64, int64_t) - RETURN_GGUF_DATA(GGUF_TYPE_FLOAT64, double) - RETURN_GGUF_DATA(GGUF_TYPE_STRING, char) - - case GGUF_TYPE_ARRAY: - throw std::runtime_error("GGUF_TYPE_ARRAY should be processed element by element"); - - default: - return "GGUF_TYPE_UNKNOWN"; - } - -#undef RETURN_GGUF_DATA -} - -struct GGUFKeyValue { - std::string key; - GGUF_TYPE gguf_type; // gguf_type - std::vector value; - - std::string toString() const; -}; - -typedef enum { - GGML_TYPE_F32 = 0, - GGML_TYPE_F16 = 1, - GGML_TYPE_Q4_0 = 2, - GGML_TYPE_Q4_1 = 3, - GGML_TYPE_Q5_0 = 6, - GGML_TYPE_Q5_1 = 7, - GGML_TYPE_Q8_0 = 8, - GGML_TYPE_Q8_1 = 9, - GGML_TYPE_Q2_K = 10, - GGML_TYPE_Q3_K = 11, - GGML_TYPE_Q4_K = 12, - GGML_TYPE_Q5_K = 13, - GGML_TYPE_Q6_K = 14, - GGML_TYPE_Q8_K = 15, - GGML_TYPE_IQ2_XXS = 16, - GGML_TYPE_IQ2_XS = 17, - GGML_TYPE_IQ3_XXS = 18, - GGML_TYPE_IQ1_S = 19, - GGML_TYPE_IQ4_NL = 20, - GGML_TYPE_IQ3_S = 21, - GGML_TYPE_IQ2_S = 22, - GGML_TYPE_IQ4_XS = 23, - GGML_TYPE_I8 = 24, - GGML_TYPE_I16 = 25, - GGML_TYPE_I32 = 26, - GGML_TYPE_I64 = 27, - GGML_TYPE_F64 = 28, - GGML_TYPE_IQ1_M = 29, - GGML_TYPE_BF16 = 30, - GGML_TYPE_TQ1_0 = 34, - GGML_TYPE_TQ2_0 = 35, - GGML_TYPE_COUNT = 36, -} GGML_TYPE; - -inline size_t ggmlTypeSize(GGML_TYPE ggml_type) { - switch (ggml_type) { - case GGML_TYPE_Q8_K: - return 1; - case GGML_TYPE_I8: - return 1; - case GGML_TYPE_I16: - return 2; - case GGML_TYPE_I32: - return 4; - case GGML_TYPE_I64: - return 8; - case GGML_TYPE_BF16: - return 2; - case GGML_TYPE_F16: - return 2; - case GGML_TYPE_F32: - return 4; - case GGML_TYPE_F64: - return 8; - default: - throw std::runtime_error("GGML_TYPE_SIZE: Unsupported GGML_TYPE"); - } - return 0; -} - -constexpr const char *GGML_TYPE_NAME[GGML_TYPE_COUNT] = { - "F32", - "F16", - "Q4_0", - "Q4_1", - nullptr, // 4 (gap) - nullptr, // 5 (gap) - "Q5_0", - "Q5_1", - "Q8_0", - "Q8_1", - "Q2_K", - "Q3_K", - "Q4_K", - "Q5_K", - "Q6_K", - "Q8_K", - "IQ2_XXS", - "IQ2_XS", - "IQ3_XXS", - "IQ1_S", - "IQ4_NL", - "IQ3_S", - "IQ2_S", - "IQ4_XS", - "I8", - "I16", - "I32", - "I64", - "F64", - "IQ1_M", - "BF16", - nullptr, // 31 (gap) - nullptr, // 32 (gap) - nullptr, // 33 (gap) - "TQ1_0", - "TQ2_0", -}; - -struct GGUFTensorInfo { - std::string name; - uint32_t ndim; - std::vector shape; - GGML_TYPE ggml_type; - uint64_t data_offset; - - std::string toString() const; -}; - -class GGUFFileReader { -public: - GGUFFileReader(const std::string &filepath); - ~GGUFFileReader() = default; - std::string toString() const; - - const std::unordered_map> &getAttributeMap() const; - const std::unordered_map> &getTensorInfoMap() const; - - std::shared_ptr getFileMapping() const { return _file; } - void *getGgmlStart() const { return _cursor; } - -private: - void readHeader(); - void readMetaKVs(); - void readTensorInfos(); - std::string readString(); - template - T read(); - - std::shared_ptr _file; - void *_data = nullptr; - uint8_t *_cursor = nullptr; - uint32_t _version; - int64_t _num_tensors; - int64_t _num_meta_kvs; - std::vector> _meta_kvs; - std::vector> _tensor_infos; - - std::unordered_map> _attributes_map; - std::unordered_map> _tensors_info_map; -}; - -#endif diff --git a/src/infiniop-test/include/ops.hpp b/src/infiniop-test/include/ops.hpp deleted file mode 100644 index 2490a8291..000000000 --- a/src/infiniop-test/include/ops.hpp +++ /dev/null @@ -1,84 +0,0 @@ -#ifndef __INFINIOPTEST_OPS_HPP__ -#define __INFINIOPTEST_OPS_HPP__ -#include "test.hpp" - -/* - * Declare all the tests here - */ -DECLARE_INFINIOP_TEST(gemm) -DECLARE_INFINIOP_TEST(random_sample) -DECLARE_INFINIOP_TEST(rms_norm) -DECLARE_INFINIOP_TEST(mul) -DECLARE_INFINIOP_TEST(rope) -DECLARE_INFINIOP_TEST(clip) -DECLARE_INFINIOP_TEST(swiglu) -DECLARE_INFINIOP_TEST(add) -DECLARE_INFINIOP_TEST(atanh) -DECLARE_INFINIOP_TEST(addcmul) -DECLARE_INFINIOP_TEST(cdist) -DECLARE_INFINIOP_TEST(binary_cross_entropy_with_logits) -DECLARE_INFINIOP_TEST(reciprocal) -DECLARE_INFINIOP_TEST(causal_softmax) -DECLARE_INFINIOP_TEST(rearrange) -DECLARE_INFINIOP_TEST(silu) -DECLARE_INFINIOP_TEST(sub) -DECLARE_INFINIOP_TEST(zeros) -DECLARE_INFINIOP_TEST(ones) - -DECLARE_INFINIOP_TEST(sigmoid) -DECLARE_INFINIOP_TEST(topkrouter) -DECLARE_INFINIOP_TEST(topksoftmax) -#define REGISTER_INFINIOP_TEST(name) \ - { \ - #name, \ - { \ - infiniop_test::name::Test::build, \ - infiniop_test::name::Test::attribute_names(), \ - infiniop_test::name::Test::tensor_names(), \ - infiniop_test::name::Test::output_names(), \ - }}, - -/* - * Register all the tests here - */ -#define TEST_BUILDER_MAPPINGS \ - { \ - REGISTER_INFINIOP_TEST(gemm) \ - REGISTER_INFINIOP_TEST(random_sample) \ - REGISTER_INFINIOP_TEST(add) \ - REGISTER_INFINIOP_TEST(mul) \ - REGISTER_INFINIOP_TEST(clip) \ - REGISTER_INFINIOP_TEST(swiglu) \ - REGISTER_INFINIOP_TEST(rope) \ - REGISTER_INFINIOP_TEST(rms_norm) \ - REGISTER_INFINIOP_TEST(causal_softmax) \ - REGISTER_INFINIOP_TEST(rearrange) \ - REGISTER_INFINIOP_TEST(sub) \ - REGISTER_INFINIOP_TEST(zeros) \ - REGISTER_INFINIOP_TEST(ones) \ - REGISTER_INFINIOP_TEST(sigmoid) \ - REGISTER_INFINIOP_TEST(topkrouter) \ - REGISTER_INFINIOP_TEST(topksoftmax) \ - REGISTER_INFINIOP_TEST(silu) \ - } - -namespace infiniop_test { - -// Global variable for {op_name: builder} mappings -extern std::unordered_map TEST_BUILDERS; - -template -bool check_names( - const std::unordered_map &map, - const std::vector &names) { - for (auto const &name : names) { - if (map.find(name) == map.end()) { - return false; - } - } - return true; -} - -} // namespace infiniop_test - -#endif diff --git a/src/infiniop-test/include/tensor.hpp b/src/infiniop-test/include/tensor.hpp deleted file mode 100644 index e6a7c5b69..000000000 --- a/src/infiniop-test/include/tensor.hpp +++ /dev/null @@ -1,82 +0,0 @@ -#ifndef __INFINIOPTEST_TENSOR_HPP__ -#define __INFINIOPTEST_TENSOR_HPP__ -#include "file_mapping.hpp" -#include "gguf.hpp" -#include - -inline infiniDtype_t ggmlTypeToInfiniType(GGML_TYPE type) { - switch (type) { - case GGML_TYPE_Q8_K: - return INFINI_DTYPE_BOOL; - case GGML_TYPE_I8: - return INFINI_DTYPE_I8; - case GGML_TYPE_I16: - return INFINI_DTYPE_I16; - case GGML_TYPE_I32: - return INFINI_DTYPE_I32; - case GGML_TYPE_I64: - return INFINI_DTYPE_I64; - case GGML_TYPE_BF16: - return INFINI_DTYPE_BF16; - case GGML_TYPE_F16: - return INFINI_DTYPE_F16; - case GGML_TYPE_F32: - return INFINI_DTYPE_F32; - case GGML_TYPE_F64: - return INFINI_DTYPE_F64; - default: - throw std::runtime_error("Unsupported GGML type"); - } -} - -namespace infiniop_test { -class Memory { -private: - void *_ptr; - size_t _size; - infiniDevice_t _device; - int _device_id; - std::shared_ptr _file_mapping; - -public: - Memory(size_t size, infiniDevice_t device, int device_id); - Memory(const std::shared_ptr &file_mapping, void *ptr, size_t size); - ~Memory(); - void *ptr() const { return _ptr; } - size_t size() const { return _size; } - infiniDevice_t device() const { return _device; } - int device_id() const { return _device_id; } -}; - -class Tensor { -private: - infiniopTensorDescriptor_t _desc; - std::shared_ptr _memory; - std::vector _shape; - std::vector _strides; - size_t _offset; - GGML_TYPE _ggml_type; - -public: - Tensor(const GGUFTensorInfo *info, - const void *ggml_ptr, - const GGUFKeyValue *shape_meta = nullptr, - const GGUFKeyValue *strides_meta = nullptr, - bool isOutput = false); - Tensor(std::shared_ptr memory, size_t offset, - const std::vector &shape, - const std::vector &strides, - GGML_TYPE dtype); - ~Tensor(); - infiniopTensorDescriptor_t desc() const { return _desc; } - std::vector shape() const { return std::vector(_shape); } - std::vector strides() const { return std::vector(_strides); } - GGML_TYPE ggml_type() const { return _ggml_type; } - void *data() const; - std::shared_ptr to(infiniDevice_t device, int device_id = 0) const; - std::string info() const; - void debug() const; -}; -} // namespace infiniop_test - -#endif diff --git a/src/infiniop-test/include/test.hpp b/src/infiniop-test/include/test.hpp deleted file mode 100644 index e2dd45f9f..000000000 --- a/src/infiniop-test/include/test.hpp +++ /dev/null @@ -1,132 +0,0 @@ -#ifndef __INFINIOPTEST_HPP__ -#define __INFINIOPTEST_HPP__ - -#include "gguf.hpp" -#include "tensor.hpp" -#include -#include -#include -#include - -#define RESET "\033[0m" -#define GREEN "\033[32m" -#define RED "\033[31m" -#define YELLOW "\033[33m" - -namespace infiniop_test { -enum class TestStatus { - PASS, - TEST_INIT_FAILED, - OP_CREATION_FAILED, - OP_EXECUTION_FAILED, - RESULT_INCORRECT, -}; - -// Result of a testcase -class Result { -private: - TestStatus _status; - double _time = 0.; - std::string _description; - std::string _error_message; - -public: - Result(TestStatus status_, double time_, const std::string &description_, const std::string &error_message_) - : _status(status_), _time(time_), _description(description_), _error_message(error_message_) {} - bool isPassed() const { return _status == TestStatus::PASS; } - std::string toString() const; -}; - -// Quick macro for creating a test result -#define TEST_PASSED(delay) std::make_shared(infiniop_test::TestStatus::PASS, delay, toString(), "") -#define TEST_FAILED(reason, msg) std::make_shared(infiniop_test::TestStatus::reason, 0., toString(), msg) -#define TEST_INIT_FAILED(op_name) std::make_shared(infiniop_test::TestStatus::TEST_INIT_FAILED, 0., "Invalid " + std::string(op_name), "") - -// Run all tests read from a GGUF file -std::vector> runAllTests( - const GGUFFileReader &, - infiniDevice_t device, int device_id, - size_t warm_ups, size_t iterations, - double rtol, double atol); - -// Run a single test read from a GGUF file -std::shared_ptr runTest( - const GGUFFileReader &, - infiniDevice_t device, int device_id, - size_t warm_ups, size_t iterations, - double rtol, double atol, - size_t test_id); - -// Check if two tensors are close within given tolerance -void allClose(std::shared_ptr actual, std::shared_ptr expected, double rtol = 1e-3, double atol = 1e-3); - -// Check if two tensors are equal -void allEqual(std::shared_ptr actual, std::shared_ptr expected); - -// Helper function for benchmarking a function -double benchmark(std::function func, size_t warmups, size_t iterations); -} // namespace infiniop_test - -namespace infiniop_test::base { -// Base class for a testcase, each operator test should inherit from this class -class Test { -public: - virtual std::shared_ptr run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, - size_t warm_ups, size_t iterations) - = 0; - virtual std::string toString() const = 0; -}; - -} // namespace infiniop_test::base - -// Quick macro for declaring a new testcase -#define DECLARE_INFINIOP_TEST(name) \ - namespace infiniop_test::name { \ - class Test : public infiniop_test::base::Test { \ - double _rtol, _atol; \ - \ - public: \ - static std::string op_name() { return #name; } \ - static std::shared_ptr build( \ - std::unordered_map> attributes, \ - std::unordered_map> tensors, \ - double, double); \ - \ - static std::vector attribute_names(); \ - static std::vector tensor_names(); \ - static std::vector output_names(); \ - \ - std::shared_ptr run( \ - infiniopHandle_t handle, infiniDevice_t device, int device_id, \ - size_t warm_ups, size_t iterations) override; \ - \ - std::string toString() const override; \ - \ - ~Test(); \ - \ - private: \ - struct Attributes; \ - Attributes *_attributes; \ - Test() = delete; \ - Test(double rtol, double atol) : _rtol(rtol), _atol(atol) {} \ - }; \ - } - -namespace infiniop_test { -using BuilderFunc = std::function( - std::unordered_map>, - std::unordered_map>, - double, double)>; - -// Testcase Registry -// Each testcase should provid a formatted builder, attribute names, and tensor names -struct TestBuilder { - BuilderFunc build; - std::vector attribute_names; - std::vector tensor_names; - std::vector output_names; -}; -} // namespace infiniop_test - -#endif diff --git a/src/infiniop-test/include/utils.hpp b/src/infiniop-test/include/utils.hpp deleted file mode 100644 index cdb85ce5d..000000000 --- a/src/infiniop-test/include/utils.hpp +++ /dev/null @@ -1,60 +0,0 @@ -#ifndef __INFINIOPTEST_UTILS_HPP__ -#define __INFINIOPTEST_UTILS_HPP__ -#include "../../utils.h" -#include "gguf.hpp" -#include -#include - -#define CHECK_OR(cmd, action) CHECK_API_OR(cmd, INFINI_STATUS_SUCCESS, action) - -inline double getVal(void *ptr, GGML_TYPE ggml_type) { - switch (ggml_type) { - case GGML_TYPE_BF16: - return utils::cast(*(bf16_t *)ptr); - case GGML_TYPE_F16: - return utils::cast(*(fp16_t *)ptr); - case GGML_TYPE_F32: - return *(float *)ptr; - case GGML_TYPE_F64: - return *(double *)ptr; - case GGML_TYPE_Q8_K: - return *(bool *)ptr; - case GGML_TYPE_I8: - return *(int8_t *)ptr; - case GGML_TYPE_I16: - return *(int16_t *)ptr; - case GGML_TYPE_I32: - return *(int32_t *)ptr; - case GGML_TYPE_I64: - return (double)(*(int64_t *)ptr); - default: - throw std::runtime_error("Unsupported data type"); - } -} - -inline size_t ggmlSizeOf(GGML_TYPE ggml_type) { - switch (ggml_type) { - case GGML_TYPE_BF16: - return sizeof(bf16_t); - case GGML_TYPE_F16: - return sizeof(fp16_t); - case GGML_TYPE_F32: - return sizeof(float); - case GGML_TYPE_F64: - return sizeof(double); - case GGML_TYPE_Q8_K: - return sizeof(bool); - case GGML_TYPE_I8: - return sizeof(int8_t); - case GGML_TYPE_I16: - return sizeof(int16_t); - case GGML_TYPE_I32: - return sizeof(int32_t); - case GGML_TYPE_I64: - return sizeof(int64_t); - default: - throw std::runtime_error("Unsupported data type"); - } -} - -#endif diff --git a/src/infiniop-test/src/file_mapping.cpp b/src/infiniop-test/src/file_mapping.cpp deleted file mode 100644 index 74acb15c7..000000000 --- a/src/infiniop-test/src/file_mapping.cpp +++ /dev/null @@ -1,65 +0,0 @@ -#include "file_mapping.hpp" -#include - -FileMapping::FileMapping(const std::string &filepath) { -#ifdef _WIN32 - _file_handle = CreateFile(filepath.c_str(), GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL); - if (_file_handle == INVALID_HANDLE_VALUE) { - throw std::runtime_error("Failed to open GGUF file"); - } - _file_mapping = CreateFileMapping(_file_handle, NULL, PAGE_READONLY, 0, 0, NULL); - if (!_file_mapping) { - CloseHandle(_file_handle); - throw std::runtime_error("Failed to create file mapping"); - } - _ptr = MapViewOfFile(_file_mapping, FILE_MAP_READ, 0, 0, 0); - if (!_ptr) { - CloseHandle(_file_mapping); - CloseHandle(_file_handle); - throw std::runtime_error("Failed to map view of file"); - } - _size = GetFileSize(_file_handle, NULL); -#else - int fd = open(filepath.c_str(), O_RDONLY); - if (fd == -1) { - throw std::runtime_error("Failed to open GGUF file"); - } - struct stat sb; - if (fstat(fd, &sb) == -1) { - close(fd); - throw std::runtime_error("Failed to get file size"); - } - _size = sb.st_size; - _ptr = mmap(NULL, _size, PROT_READ, MAP_PRIVATE, fd, 0); - close(fd); - if (_ptr == MAP_FAILED) { - throw std::runtime_error("Failed to mmap file"); - } -#endif -} - -FileMapping::~FileMapping() { -#ifdef _WIN32 - if (_ptr) { - UnmapViewOfFile(_ptr); - } - if (_file_mapping) { - CloseHandle(_file_mapping); - } - if (_file_handle) { - CloseHandle(_file_handle); - } -#else - if (_ptr) { - munmap(_ptr, _size); - } -#endif -} - -void *FileMapping::ptr() const { - return _ptr; -} - -size_t FileMapping::size() const { - return _size; -} diff --git a/src/infiniop-test/src/gguf.cpp b/src/infiniop-test/src/gguf.cpp deleted file mode 100644 index a4b200033..000000000 --- a/src/infiniop-test/src/gguf.cpp +++ /dev/null @@ -1,167 +0,0 @@ -#include "gguf.hpp" -#include -#include -#include -#include -#include - -#ifdef _WIN32 -#include -#else -#include -#include -#include - -#endif - -std::string GGUFKeyValue::toString() const { - std::ostringstream oss; - oss << "Key: " << key << ", Type: " << GGUF_TYPE_NAME[gguf_type] << ", Value: "; - if (gguf_type == GGUF_TYPE_STRING) { - std::string str(value.begin(), value.end()); - oss << str; - } else if (value.size() > GGUF_TYPE_SIZE[gguf_type]) { - oss << "["; - for (size_t i = 0; i < value.size() / GGUF_TYPE_SIZE[gguf_type]; ++i) { - oss << ggufDataToString(value.data() + i * GGUF_TYPE_SIZE[gguf_type], gguf_type); - if (i < value.size() / GGUF_TYPE_SIZE[gguf_type] - 1) { - oss << ", "; - } - } - oss << "]"; - } else { - oss << ggufDataToString(value.data(), gguf_type); - } - - return oss.str(); -} - -std::string GGUFTensorInfo::toString() const { - std::ostringstream oss; - oss << "Name: " << name << ", NDims: " << ndim << ", Shape: ["; - for (size_t i = 0; i < shape.size(); ++i) { - oss << shape[i]; - if (i < shape.size() - 1) { - oss << ", "; - } - } - oss << "], DataType: " << GGML_TYPE_NAME[ggml_type] << ", DataOffset: " << data_offset; - return oss.str(); -} - -GGUFFileReader::GGUFFileReader(const std::string &filepath) { - try { - _file = std::make_shared(filepath); - } catch (const std::exception &e) { - throw e; - } - _data = _file->ptr(); - _cursor = reinterpret_cast(_data); - readHeader(); - readMetaKVs(); - readTensorInfos(); - size_t padding = (size_t)(32 - ((char *)_cursor - (char *)_data) % 32) % 32; - _cursor += padding; -} - -const std::unordered_map> & -GGUFFileReader::getAttributeMap() const { - return _attributes_map; -} - -const std::unordered_map> & -GGUFFileReader::getTensorInfoMap() const { - return _tensors_info_map; -} - -void GGUFFileReader::readHeader() { - if (std::memcmp(_cursor, "GGUF", 4) != 0) { - throw std::runtime_error("Invalid GGUF magic"); - } - _cursor += 4; - - _version = read(); - _num_tensors = read(); - _num_meta_kvs = read(); - _attributes_map = std::unordered_map>(); - _tensors_info_map = std::unordered_map>(); -} - -void GGUFFileReader::readMetaKVs() { - for (int64_t i = 0; i < _num_meta_kvs; ++i) { - auto kv = std::make_shared(); - kv->key = readString(); - kv->gguf_type = read(); - - if (kv->gguf_type == GGUF_TYPE_ARRAY) { - GGUF_TYPE array_type = read(); - uint64_t array_size = read(); - kv->value.resize(array_size * GGUF_TYPE_SIZE[array_type]); - kv->gguf_type = array_type; - std::memcpy(kv->value.data(), _cursor, kv->value.size()); - _cursor += kv->value.size(); - } else if (kv->gguf_type == GGUF_TYPE_STRING) { - uint64_t str_size = read(); - kv->value.resize(str_size); - std::memcpy(kv->value.data(), _cursor, str_size); - _cursor += str_size; - } else { - kv->value.resize(GGUF_TYPE_SIZE[kv->gguf_type]); - std::memcpy(kv->value.data(), _cursor, kv->value.size()); - _cursor += kv->value.size(); - } - - _meta_kvs.push_back(kv); - _attributes_map.emplace(kv->key, kv); - } -} - -void GGUFFileReader::readTensorInfos() { - for (int64_t i = 0; i < _num_tensors; ++i) { - auto tensor_info = std::make_shared(); - tensor_info->name = readString(); - tensor_info->ndim = read(); - tensor_info->shape.resize(tensor_info->ndim); - for (size_t j = 0; j < tensor_info->ndim; ++j) { - tensor_info->shape[j] = read(); - } - tensor_info->ggml_type = read(); - tensor_info->data_offset = read(); - _tensor_infos.push_back(tensor_info); - _tensors_info_map.emplace(tensor_info->name, tensor_info); - } -} - -std::string GGUFFileReader::readString() { - uint64_t length = read(); - std::string str(reinterpret_cast(_cursor), length); - _cursor += length; - return str; -} - -template -T GGUFFileReader::read() { - T value; - std::memcpy(&value, _cursor, sizeof(T)); - _cursor += sizeof(T); - return value; -} - -std::string GGUFFileReader::toString() const { - std::ostringstream oss; - oss << "GGUF File Contents: " << std::endl; - oss << "Version: " << _version << std::endl; - oss << "Number of Meta KVs: " << _num_meta_kvs << std::endl; - oss << "Number of Tensors: " << _num_tensors << std::endl - << std::endl; - oss << "Meta KVs: " << std::endl; - for (const auto &kv : _meta_kvs) { - oss << kv->toString() << std::endl; - } - oss << std::endl; - oss << "Tensor INFOs: " << std::endl; - for (const auto &info : _tensor_infos) { - oss << info->toString() << std::endl; - } - return oss.str(); -} diff --git a/src/infiniop-test/src/main.cpp b/src/infiniop-test/src/main.cpp deleted file mode 100644 index 00e16cc99..000000000 --- a/src/infiniop-test/src/main.cpp +++ /dev/null @@ -1,143 +0,0 @@ -#include "gguf.hpp" -#include "test.hpp" -#include -#include - -struct ParsedArgs { - std::string file_path; // Mandatory argument: test.gguf file path - infiniDevice_t device_type = INFINI_DEVICE_CPU; // Default to CPU - int device_id = 0; // CUDA device ID (if specified) - int warmups = 0; // Default to 0 if not given - int iterations = 0; // Default to 0 if not given - double atol = 0.0015; // Default absolute tolerance - double rtol = 0.001; // Default relative tolerance -}; - -void printUsage() { - std::cout << "Usage:" << std::endl - << std::endl; - std::cout << "infiniop-test [--[:id]] [--warmup ] [--run ] [--atol ] [--rtol ]" << std::endl - << std::endl; - std::cout << " >" << std::endl; - std::cout << " Path to the test gguf file" << std::endl - << std::endl; - std::cout << " --[:id]" << std::endl; - std::cout << " (Optional) Specify the device type --(cpu|nvidia|cambricon|ascend|metax|moore|iluvatar|qy|kunlun|hygon|ali) and device ID (optional). CPU by default." << std::endl - << std::endl; - std::cout << " --warmup " << std::endl; - std::cout << " (Optional) Number of warmups to perform before timing. Default to 0." << std::endl - << std::endl; - std::cout << " --run " << std::endl; - std::cout << " (Optional) Number of iterations to perform for timing. Default to 0." << std::endl - << std::endl; - std::cout << " --atol " << std::endl; - std::cout << " (Optional) Absolute tolerance for correctness check. Default to 0.001" << std::endl - << std::endl; - std::cout << " --rtol " << std::endl; - std::cout << " (Optional) Relative tolerance for correctness check. Default to 0.001" << std::endl - << std::endl; - exit(-1); -} - -#define PARSE_DEVICE(FLAG, DEVICE) \ - else if (arg.find(FLAG) == 0) { \ - size_t colon_pos = arg.find(':'); \ - args.device_type = DEVICE; \ - if (colon_pos != std::string::npos) { \ - args.device_id = std::stoi(arg.substr(colon_pos + 1)); \ - } else { \ - args.device_id = 0; \ - } \ - } - -ParsedArgs parseArgs(int argc, char *argv[]) { - if (argc < 2) { - printUsage(); - } - - if (std::string(argv[1]) == "--help" || std::string(argv[1]) == "-h") { - printUsage(); - } - - ParsedArgs args; - args.file_path = argv[1]; // First argument is always the test.gguf file - - std::unordered_map options; - - try { - for (int i = 2; i < argc; ++i) { - std::string arg = argv[i]; - - if (arg.find("--cpu") == 0) { - args.device_id = 0; - } - PARSE_DEVICE("--nvidia", INFINI_DEVICE_NVIDIA) - PARSE_DEVICE("--cambricon", INFINI_DEVICE_CAMBRICON) - PARSE_DEVICE("--ascend", INFINI_DEVICE_ASCEND) - PARSE_DEVICE("--metax", INFINI_DEVICE_METAX) - PARSE_DEVICE("--moore", INFINI_DEVICE_MOORE) - PARSE_DEVICE("--iluvatar", INFINI_DEVICE_ILUVATAR) - PARSE_DEVICE("--qy", INFINI_DEVICE_QY) - PARSE_DEVICE("--kunlun", INFINI_DEVICE_KUNLUN) - PARSE_DEVICE("--hygon", INFINI_DEVICE_HYGON) - PARSE_DEVICE("--ali", INFINI_DEVICE_ALI) - else if (arg == "--warmup" && i + 1 < argc) { - args.warmups = std::stoi(argv[++i]); - } - else if (arg == "--run" && i + 1 < argc) { - args.iterations = std::stoi(argv[++i]); - } - else if (arg == "--atol" && i + 1 < argc) { - args.atol = std::stod(argv[++i]); - } - else if (arg == "--rtol" && i + 1 < argc) { - args.rtol = std::stod(argv[++i]); - } - else { - printUsage(); - } - } - } catch (const std::exception &) { - printUsage(); - } - - return args; -} - -int main(int argc, char *argv[]) { - ParsedArgs args = parseArgs(argc, argv); - int failed = 0; - try { - std::cout << args.file_path << std::endl; - GGUFFileReader reader = GGUFFileReader(args.file_path); - // std::cout << reader.toString() << std::endl; - - if (infinirtInit() != INFINI_STATUS_SUCCESS) { - std::cerr << "Error: Failed to initialize InfiniRT" << std::endl; - return -1; - } - - auto results = infiniop_test::runAllTests( - reader, - (infiniDevice_t)args.device_type, args.device_id, - args.warmups, args.iterations, - args.rtol, args.atol); - - std::cout << "=====================================" << std::endl; - for (auto result : results) { - if (!result->isPassed()) { - failed++; - } - std::cout << result->toString() << std::endl; - std::cout << "=====================================" << std::endl; - } - if (failed == 0) { - std::cout << GREEN << "All tests passed" << RESET << std::endl; - } else { - std::cout << RED << failed << " of " << results.size() << " tests failed" << RESET << std::endl; - } - } catch (const std::exception &e) { - std::cerr << "Error: " << e.what() << std::endl; - } - return failed; -} diff --git a/src/infiniop-test/src/ops/add.cpp b/src/infiniop-test/src/ops/add.cpp deleted file mode 100644 index 27f69d687..000000000 --- a/src/infiniop-test/src/ops/add.cpp +++ /dev/null @@ -1,109 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::add { -struct Test::Attributes { - std::shared_ptr a; - std::shared_ptr b; - std::shared_ptr c; - std::shared_ptr ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (tensors.find("a") == tensors.end() - || tensors.find("b") == tensors.end() - || tensors.find("c") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->a = tensors["a"]; - test->_attributes->b = tensors["b"]; - test->_attributes->c = tensors["c"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - infiniopAddDescriptor_t op_desc; - auto a = _attributes->a->to(device, device_id); - auto b = _attributes->b->to(device, device_id); - auto c = _attributes->c->to(device, device_id); - CHECK_OR(infiniopCreateAddDescriptor(handle, &op_desc, - c->desc(), - a->desc(), - b->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - size_t workspace_size; - CHECK_OR(infiniopGetAddWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - CHECK_OR(infiniopAdd(op_desc, workspace, workspace_size, - c->data(), - a->data(), - b->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(c, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopAdd( - op_desc, workspace, workspace_size, - c->data(), - a->data(), - b->data(), - nullptr); - }, - warm_ups, iterations); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} - -std::vector Test::tensor_names() { - return {"a", "b", "c", "ans"}; -} - -std::vector Test::output_names() { - return {"c"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- a: " << _attributes->a->info() << std::endl; - oss << "- b: " << _attributes->b->info() << std::endl; - oss << "- c: " << _attributes->c->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::add diff --git a/src/infiniop-test/src/ops/addcmul.cpp b/src/infiniop-test/src/ops/addcmul.cpp deleted file mode 100644 index c1dc11679..000000000 --- a/src/infiniop-test/src/ops/addcmul.cpp +++ /dev/null @@ -1,139 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::addcmul { - -struct Test::Attributes { - std::shared_ptr input; - std::shared_ptr t1; - std::shared_ptr t2; - std::shared_ptr out; - std::shared_ptr ans; - float value; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - - // 校验张量是否存在 - if (tensors.find("input") == tensors.end() || tensors.find("t1") == tensors.end() || tensors.find("t2") == tensors.end() || tensors.find("out") == tensors.end() || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Addcmul Test: Missing tensors"); - } - - // 获取标量属性 value - test->_attributes->value = 1.0f; // 默认值 - if (attributes.find("value") != attributes.end()) { - test->_attributes->value = *reinterpret_cast(attributes["value"].data()); - } - - test->_attributes->input = tensors["input"]; - test->_attributes->t1 = tensors["t1"]; - test->_attributes->t2 = tensors["t2"]; - test->_attributes->out = tensors["out"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - - infiniopAddcmulDescriptor_t op_desc; - - // 数据迁移至指定设备 - auto input = _attributes->input->to(device, device_id); - auto t1 = _attributes->t1->to(device, device_id); - auto t2 = _attributes->t2->to(device, device_id); - auto out = _attributes->out->to(device, device_id); - - // 创建算子描述符 - CHECK_OR(infiniopCreateAddcmulDescriptor(handle, &op_desc, - out->desc(), - input->desc(), - t1->desc(), - t2->desc(), - _attributes->value), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create addcmul descriptor.")); - - // Workspace 处理 - size_t workspace_size; - CHECK_OR(infiniopGetAddcmulWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - - // 执行计算 - CHECK_OR(infiniopAddcmul(op_desc, workspace, workspace_size, - out->data(), - input->data(), - t1->data(), - t2->data(), - nullptr), // stream - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - // 结果验证 - try { - allClose(out, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - // 性能测试 - double elapsed_time = benchmark( - [=]() { - infiniopAddcmul(op_desc, workspace, workspace_size, - out->data(), - input->data(), - t1->data(), - t2->data(), - nullptr); - }, - warm_ups, iterations); - - // 资源清理 - infinirtFree(workspace); - infiniopDestroyAddcmulDescriptor(op_desc); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {"value"}; -} - -std::vector Test::tensor_names() { - return {"input", "t1", "t2", "out", "ans"}; -} - -std::vector Test::output_names() { - return {"out"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- value: " << _attributes->value << std::endl; - oss << "- input: " << _attributes->input->info() << std::endl; - oss << "- t1: " << _attributes->t1->info() << std::endl; - oss << "- t2: " << _attributes->t2->info() << std::endl; - oss << "- out: " << _attributes->out->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::addcmul diff --git a/src/infiniop-test/src/ops/atanh.cpp b/src/infiniop-test/src/ops/atanh.cpp deleted file mode 100644 index b24018bb9..000000000 --- a/src/infiniop-test/src/ops/atanh.cpp +++ /dev/null @@ -1,116 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::atanh { -struct Test::Attributes { - std::shared_ptr a; // 输入 - std::shared_ptr y; // 输出 - std::shared_ptr ans; // 参考结果 -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - - // atanh 只需要 a (input), y (output) 和 ans (reference) - if (tensors.find("a") == tensors.end() - || tensors.find("y") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Atanh Test: Missing tensors."); - } - - test->_attributes->a = tensors["a"]; - test->_attributes->y = tensors["y"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - - infiniopAtanhDescriptor_t op_desc; - auto a = _attributes->a->to(device, device_id); - auto y = _attributes->y->to(device, device_id); - - // 调用修正后的 4 参数版本接口 (handle, desc, y, a) - CHECK_OR(infiniopCreateAtanhDescriptor(handle, &op_desc, - y->desc(), - a->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create atanh descriptor.")); - - size_t workspace_size; - CHECK_OR(infiniopGetAtanhWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - - // 执行计算 (移除 b 相关的参数) - CHECK_OR(infiniopAtanh(op_desc, workspace, workspace_size, - y->data(), - a->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during atanh execution.")); - - // 验证结果 - try { - allClose(y, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - // 性能测试 (Benchmark) - double elapsed_time = 0.; - elapsed_time = benchmark( - [=]() { - infiniopAtanh( - op_desc, workspace, workspace_size, - y->data(), - a->data(), - nullptr); - }, - warm_ups, iterations); - - // 释放资源 (可选：根据框架决定是否在此释放 op_desc) - // infiniopDestroyAtanhDescriptor(op_desc); - // infinirtFree(workspace); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} - -std::vector Test::tensor_names() { - return {"a", "y", "ans"}; -} - -std::vector Test::output_names() { - return {"y"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- a: " << _attributes->a->info() << std::endl; - oss << "- y: " << _attributes->y->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::atanh diff --git a/src/infiniop-test/src/ops/binary_cross_entropy_with_logits.cpp b/src/infiniop-test/src/ops/binary_cross_entropy_with_logits.cpp deleted file mode 100644 index 88c7e7123..000000000 --- a/src/infiniop-test/src/ops/binary_cross_entropy_with_logits.cpp +++ /dev/null @@ -1,153 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::binary_cross_entropy_with_logits { - -struct Test::Attributes { - std::shared_ptr logits; - std::shared_ptr target; - std::shared_ptr weight; // 可选 - std::shared_ptr pos_weight; // 可选 - std::shared_ptr out; - std::shared_ptr ans; - int reduction; // 0: none, 1: mean, 2: sum -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - - // 1. 校验必要张量是否存在 - if (tensors.find("logits") == tensors.end() || tensors.find("target") == tensors.end() || tensors.find("out") == tensors.end() || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid BCE Test: Missing mandatory tensors"); - } - - // 2. 获取 reduction 属性 (默认为 1: mean) - test->_attributes->reduction = 1; - if (attributes.find("reduction") != attributes.end()) { - test->_attributes->reduction = *reinterpret_cast(attributes["reduction"].data()); - } - - // 3. 填充张量（处理可选张量） - test->_attributes->logits = tensors["logits"]; - test->_attributes->target = tensors["target"]; - test->_attributes->out = tensors["out"]; - test->_attributes->ans = tensors["ans"]; - - // 如果 tensors 中存在则赋值，否则为 nullptr - test->_attributes->weight = tensors.count("weight") ? tensors["weight"] : nullptr; - test->_attributes->pos_weight = tensors.count("pos_weight") ? tensors["pos_weight"] : nullptr; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - - infiniopBCEWithLogitsDescriptor_t op_desc; - - // 4. 数据迁移 - auto logits = _attributes->logits->to(device, device_id); - auto target = _attributes->target->to(device, device_id); - auto out = _attributes->out->to(device, device_id); - - // 处理可选张量迁移 - std::shared_ptr weight = (_attributes->weight) ? _attributes->weight->to(device, device_id) : nullptr; - std::shared_ptr pos_weight = (_attributes->pos_weight) ? _attributes->pos_weight->to(device, device_id) : nullptr; - - // 5. 创建描述符 (注意处理 NULL 描述符) - auto w_desc = weight ? weight->desc() : nullptr; - auto pw_desc = pos_weight ? pos_weight->desc() : nullptr; - - CHECK_OR(infiniopCreateBCEWithLogitsDescriptor(handle, &op_desc, - out->desc(), - logits->desc(), - target->desc(), - w_desc, - pw_desc, - static_cast(_attributes->reduction)), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create BCE descriptor.")); - - // 6. Workspace 管理 - size_t workspace_size; - CHECK_OR(infiniopGetBCEWithLogitsWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - - // 7. 执行计算 - auto w_data = weight ? weight->data() : nullptr; - auto pw_data = pos_weight ? pos_weight->data() : nullptr; - - CHECK_OR(infiniopBCEWithLogits(op_desc, workspace, workspace_size, - out->data(), - logits->data(), - target->data(), - w_data, - pw_data, - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - // 8. 结果验证 - try { - allClose(out, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - // 9. 性能 Benchmark - double elapsed_time = benchmark( - [=]() { - infiniopBCEWithLogits(op_desc, workspace, workspace_size, - out->data(), logits->data(), target->data(), - w_data, pw_data, nullptr); - }, - warm_ups, iterations); - - // 10. 资源清理 - infinirtFree(workspace); - infiniopDestroyBCEWithLogitsDescriptor(op_desc); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {"reduction"}; -} - -std::vector Test::tensor_names() { - return {"logits", "target", "weight", "pos_weight", "out", "ans"}; -} - -std::vector Test::output_names() { - return {"out"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- reduction: " << _attributes->reduction << std::endl; - oss << "- logits: " << _attributes->logits->info() << std::endl; - if (_attributes->weight) { - oss << "- weight: " << _attributes->weight->info() << std::endl; - } - oss << "- out: " << _attributes->out->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::binary_cross_entropy_with_logits diff --git a/src/infiniop-test/src/ops/causal_softmax.cpp b/src/infiniop-test/src/ops/causal_softmax.cpp deleted file mode 100644 index 29612960a..000000000 --- a/src/infiniop-test/src/ops/causal_softmax.cpp +++ /dev/null @@ -1,103 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::causal_softmax { -struct Test::Attributes { - std::shared_ptr x; - std::shared_ptr y; - std::shared_ptr ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (tensors.find("x") == tensors.end() - || tensors.find("y") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->x = tensors["x"]; - test->_attributes->y = tensors["y"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - infiniopCausalSoftmaxDescriptor_t op_desc; - auto y = _attributes->y->to(device, device_id); - auto x = _attributes->x->to(device, device_id); - auto ans = _attributes->ans->to(device, device_id); - CHECK_OR(infiniopCreateCausalSoftmaxDescriptor(handle, &op_desc, - y->desc(), - x->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - size_t workspace_size; - CHECK_OR(infiniopGetCausalSoftmaxWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - CHECK_OR(infiniopCausalSoftmax(op_desc, workspace, workspace_size, - y->data(), - x->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(y, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopCausalSoftmax( - op_desc, workspace, workspace_size, - y->data(), - x->data(), - nullptr); - }, - warm_ups, iterations); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} - -std::vector Test::tensor_names() { - return {"x", "y", "ans"}; -} - -std::vector Test::output_names() { - return {"y"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- x: " << _attributes->x->info() << std::endl; - oss << "- y: " << _attributes->y->info() << std::endl; - oss << "- ans: " << _attributes->ans->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::causal_softmax diff --git a/src/infiniop-test/src/ops/cdist.cpp b/src/infiniop-test/src/ops/cdist.cpp deleted file mode 100644 index b370683c2..000000000 --- a/src/infiniop-test/src/ops/cdist.cpp +++ /dev/null @@ -1,132 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::cdist { - -struct Test::Attributes { - std::shared_ptr x1; - std::shared_ptr x2; - std::shared_ptr out; - std::shared_ptr ans; - double p; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - - // 1. 校验张量是否存在 (x1, x2, out, ans) - if (tensors.find("x1") == tensors.end() || tensors.find("x2") == tensors.end() || tensors.find("out") == tensors.end() || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Cdist Test: Missing tensors"); - } - - // 2. 获取标量属性 p (注意 cdist 通常用 double) - test->_attributes->p = 2.0; // 默认值 - if (attributes.find("p") != attributes.end()) { - test->_attributes->p = *reinterpret_cast(attributes["p"].data()); - } - - test->_attributes->x1 = tensors["x1"]; - test->_attributes->x2 = tensors["x2"]; - test->_attributes->out = tensors["out"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - - infiniopCdistDescriptor_t op_desc; - - // 3. 数据迁移至指定设备 (M x D, N x D) - auto x1 = _attributes->x1->to(device, device_id); - auto x2 = _attributes->x2->to(device, device_id); - auto out = _attributes->out->to(device, device_id); - - // 4. 创建算子描述符 - CHECK_OR(infiniopCreateCdistDescriptor(handle, &op_desc, - out->desc(), - x1->desc(), - x2->desc(), - _attributes->p), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create cdist descriptor.")); - - // 5. Workspace 动态内存分配 - size_t workspace_size; - CHECK_OR(infiniopGetCdistWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - - // 6. 执行计算 (计算 M x N 距离矩阵) - CHECK_OR(infiniopCdist(op_desc, workspace, workspace_size, - out->data(), - x1->data(), - x2->data(), - nullptr), // stream - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - // 7. 结果数值验证 - try { - allClose(out, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - // 8. 性能 Benchmark - double elapsed_time = benchmark( - [=]() { - infiniopCdist(op_desc, workspace, workspace_size, - out->data(), - x1->data(), - x2->data(), - nullptr); - }, - warm_ups, iterations); - - // 9. 资源清理 - infinirtFree(workspace); - infiniopDestroyCdistDescriptor(op_desc); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {"p"}; -} - -std::vector Test::tensor_names() { - return {"x1", "x2", "out", "ans"}; -} - -std::vector Test::output_names() { - return {"out"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- p: " << _attributes->p << std::endl; - oss << "- x1: " << _attributes->x1->info() << std::endl; - oss << "- x2: " << _attributes->x2->info() << std::endl; - oss << "- out: " << _attributes->out->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::cdist diff --git a/src/infiniop-test/src/ops/clip.cpp b/src/infiniop-test/src/ops/clip.cpp deleted file mode 100644 index 82a0e9b10..000000000 --- a/src/infiniop-test/src/ops/clip.cpp +++ /dev/null @@ -1,120 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::clip { -struct Test::Attributes { - std::shared_ptr x; - std::shared_ptr min_val; - std::shared_ptr max_val; - std::shared_ptr y; - std::shared_ptr ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (tensors.find("x") == tensors.end() - || tensors.find("min_val") == tensors.end() - || tensors.find("max_val") == tensors.end() - || tensors.find("y") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->x = tensors["x"]; - test->_attributes->min_val = tensors["min_val"]; - test->_attributes->max_val = tensors["max_val"]; - test->_attributes->y = tensors["y"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - infiniopClipDescriptor_t op_desc; - auto x = _attributes->x->to(device, device_id); - auto min_val = _attributes->min_val->to(device, device_id); - auto max_val = _attributes->max_val->to(device, device_id); - auto y = _attributes->y->to(device, device_id); - CHECK_OR(infiniopCreateClipDescriptor(handle, &op_desc, - y->desc(), - x->desc(), - min_val->desc(), - max_val->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create clip descriptor.")); - size_t workspace_size; - CHECK_OR(infiniopGetClipWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - CHECK_OR(infiniopClip(op_desc, workspace, workspace_size, - y->data(), - x->data(), - min_val->data(), - max_val->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(y, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopClip( - op_desc, workspace, workspace_size, - y->data(), - x->data(), - min_val->data(), - max_val->data(), - nullptr); - }, - warm_ups, iterations); - - infiniopDestroyClipDescriptor(op_desc); - infinirtFree(workspace); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} - -std::vector Test::tensor_names() { - return {"x", "min_val", "max_val", "y", "ans"}; -} - -std::vector Test::output_names() { - return {"y"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- x: " << _attributes->x->info() << std::endl; - oss << "- min_val: " << _attributes->min_val->info() << std::endl; - oss << "- max_val: " << _attributes->max_val->info() << std::endl; - oss << "- y: " << _attributes->y->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::clip diff --git a/src/infiniop-test/src/ops/gemm.cpp b/src/infiniop-test/src/ops/gemm.cpp deleted file mode 100644 index 37c8ed6fe..000000000 --- a/src/infiniop-test/src/ops/gemm.cpp +++ /dev/null @@ -1,131 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::gemm { -struct Test::Attributes { - float alpha; - float beta; - - std::shared_ptr a; - std::shared_ptr b; - std::shared_ptr c; - std::shared_ptr ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (!check_names(attributes, Test::attribute_names()) || !check_names(tensors, Test::tensor_names())) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->alpha = *reinterpret_cast(attributes["alpha"].data()); - test->_attributes->beta = *reinterpret_cast(attributes["beta"].data()); - - test->_attributes->a = tensors["a"]; - test->_attributes->b = tensors["b"]; - test->_attributes->c = tensors["c"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - infiniopGemmDescriptor_t op_desc; - auto alpha = _attributes->alpha; - auto beta = _attributes->beta; - auto a = _attributes->a->to(device, device_id); - auto b = _attributes->b->to(device, device_id); - auto c = _attributes->c->to(device, device_id); - CHECK_OR(infiniopCreateGemmDescriptor(handle, &op_desc, - c->desc(), - a->desc(), - b->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - size_t workspace_size; - CHECK_OR(infiniopGetGemmWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - CHECK_OR(infiniopGemm(op_desc, workspace, workspace_size, - c->data(), - a->data(), - b->data(), - alpha, - beta, - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(c, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - // add and subtract to avoid overflow - float beta_ = beta == .0f ? .0f : 1.f / beta; - float alpha_ = beta == .0f ? alpha : -beta_; - - elapsed_time = benchmark( - [=]() { - infiniopGemm( - op_desc, workspace, workspace_size, - c->data(), - a->data(), - b->data(), - alpha, - beta, - nullptr); - infiniopGemm( - op_desc, workspace, workspace_size, - c->data(), - a->data(), - b->data(), - alpha_, - beta_, - nullptr); - }, - (warm_ups + 1) / 2, (iterations + 1) / 2); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {"alpha", "beta"}; -} - -std::vector Test::tensor_names() { - return {"a", "b", "c", "ans"}; -} - -std::vector Test::output_names() { - return {}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- alpha=" << _attributes->alpha << ", beta=" << _attributes->beta << std::endl; - oss << "- a: " << _attributes->a->info() << std::endl; - oss << "- b: " << _attributes->b->info() << std::endl; - oss << "- c: " << _attributes->c->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::gemm diff --git a/src/infiniop-test/src/ops/mul.cpp b/src/infiniop-test/src/ops/mul.cpp deleted file mode 100644 index 8ebfc426b..000000000 --- a/src/infiniop-test/src/ops/mul.cpp +++ /dev/null @@ -1,109 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::mul { -struct Test::Attributes { - std::shared_ptr a; - std::shared_ptr b; - std::shared_ptr c; - std::shared_ptr ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (tensors.find("a") == tensors.end() - || tensors.find("b") == tensors.end() - || tensors.find("c") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->a = tensors["a"]; - test->_attributes->b = tensors["b"]; - test->_attributes->c = tensors["c"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - infiniopMulDescriptor_t op_desc; - auto a = _attributes->a->to(device, device_id); - auto b = _attributes->b->to(device, device_id); - auto c = _attributes->c->to(device, device_id); - CHECK_OR(infiniopCreateMulDescriptor(handle, &op_desc, - c->desc(), - a->desc(), - b->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - size_t workspace_size; - CHECK_OR(infiniopGetMulWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - CHECK_OR(infiniopMul(op_desc, workspace, workspace_size, - c->data(), - a->data(), - b->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(c, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopMul( - op_desc, workspace, workspace_size, - c->data(), - a->data(), - b->data(), - nullptr); - }, - warm_ups, iterations); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} - -std::vector Test::tensor_names() { - return {"a", "b", "c", "ans"}; -} - -std::vector Test::output_names() { - return {"c"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- a: " << _attributes->a->info() << std::endl; - oss << "- b: " << _attributes->b->info() << std::endl; - oss << "- c: " << _attributes->c->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::mul diff --git a/src/infiniop-test/src/ops/ones.cpp b/src/infiniop-test/src/ops/ones.cpp deleted file mode 100644 index f8d992bba..000000000 --- a/src/infiniop-test/src/ops/ones.cpp +++ /dev/null @@ -1,103 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::ones { -struct Test::Attributes { - std::shared_ptr x; - std::shared_ptr y; - std::shared_ptr ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (tensors.find("x") == tensors.end() - || tensors.find("y") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->x = tensors["x"]; - test->_attributes->y = tensors["y"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - - infiniopOnesDescriptor_t op_desc; - auto x = _attributes->x->to(device, device_id); - auto y = _attributes->y->to(device, device_id); - CHECK_OR(infiniopCreateOnesDescriptor(handle, &op_desc, - y->desc(), - x->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - size_t workspace_size; - CHECK_OR(infiniopGetOnesWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - CHECK_OR(infiniopOnes(op_desc, workspace, workspace_size, - y->data(), - x->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(y, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopOnes( - op_desc, workspace, workspace_size, - y->data(), - x->data(), - nullptr); - }, - warm_ups, iterations); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} - -std::vector Test::tensor_names() { - return {"x", "y", "ans"}; -} - -std::vector Test::output_names() { - return {"y"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- x: " << _attributes->x->info() << std::endl; - oss << "- y: " << _attributes->y->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::ones diff --git a/src/infiniop-test/src/ops/random_sample.cpp b/src/infiniop-test/src/ops/random_sample.cpp deleted file mode 100644 index a11e0f446..000000000 --- a/src/infiniop-test/src/ops/random_sample.cpp +++ /dev/null @@ -1,128 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::random_sample { -struct Test::Attributes { - float random_val; - float topp; - int topk; - int voc; - float temperature; - - std::shared_ptr data; - std::shared_ptr ans; - std::shared_ptr result; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (!check_names(attributes, Test::attribute_names()) || !check_names(tensors, Test::tensor_names())) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->random_val = *reinterpret_cast(attributes["random_val"].data()); - test->_attributes->topp = *reinterpret_cast(attributes["topp"].data()); - test->_attributes->topk = *reinterpret_cast(attributes["topk"].data()); - test->_attributes->voc = *reinterpret_cast(attributes["voc"].data()); - test->_attributes->temperature = *reinterpret_cast(attributes["temperature"].data()); - - test->_attributes->data = tensors["data"]; - test->_attributes->ans = tensors["ans"]; - test->_attributes->result = tensors["result"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - infiniopRandomSampleDescriptor_t op_desc; - auto random_val = _attributes->random_val; - auto topp = _attributes->topp; - auto topk = _attributes->topk; - auto temperature = _attributes->temperature; - auto data = _attributes->data->to(device, device_id); - auto result = _attributes->result->to(device, device_id); - CHECK_OR(infiniopCreateRandomSampleDescriptor(handle, &op_desc, - result->desc(), - data->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - size_t workspace_size; - CHECK_OR(infiniopGetRandomSampleWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - CHECK_OR(infiniopRandomSample(op_desc, workspace, workspace_size, - result->data(), - data->data(), - random_val, - topp, - topk, - temperature, - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(result, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopRandomSample( - op_desc, workspace, workspace_size, - result->data(), - data->data(), - random_val, - topp, - topk, - temperature, - nullptr); - }, - warm_ups, iterations); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {"random_val", "topp", "topk", "voc", "temperature"}; -} - -std::vector Test::tensor_names() { - return {"data", "ans", "result"}; -} - -std::vector Test::output_names() { - return {}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- random_val=" << _attributes->random_val - << ", topp=" << _attributes->topp << std::endl - << ", topk=" << _attributes->topk << std::endl - << ", voc=" << _attributes->voc << std::endl - << ", temperature=" << _attributes->temperature << std::endl; - oss << "- data: " << _attributes->data->info() << std::endl; - oss << "- result: " << _attributes->result->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::random_sample diff --git a/src/infiniop-test/src/ops/rearrange.cpp b/src/infiniop-test/src/ops/rearrange.cpp deleted file mode 100644 index 9fbf6f2cb..000000000 --- a/src/infiniop-test/src/ops/rearrange.cpp +++ /dev/null @@ -1,97 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::rearrange { -struct Test::Attributes { - std::shared_ptr dst, src, ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (!check_names(attributes, Test::attribute_names()) || !check_names(tensors, Test::tensor_names())) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->dst = tensors["dst"]; - test->_attributes->src = tensors["src"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, - infiniDevice_t device, - int device_id, - size_t warm_ups, - size_t iterations) { - - infiniopRearrangeDescriptor_t op_desc; - auto dst = _attributes->dst->to(device, device_id); - auto src = _attributes->src->to(device, device_id); - CHECK_OR(infiniopCreateRearrangeDescriptor( - handle, &op_desc, - dst->desc(), - src->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - CHECK_OR(infiniopRearrange( - op_desc, - dst->data(), - src->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allEqual(dst, _attributes->ans); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopRearrange( - op_desc, - dst->data(), - src->data(), - nullptr); - }, - warm_ups, iterations); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} - -std::vector Test::tensor_names() { - return {"dst", "src", "ans"}; -} - -std::vector Test::output_names() { - return {"dst"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl - << "- dst: " << _attributes->dst->info() << std::endl - << "- src: " << _attributes->src->info() << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::rearrange diff --git a/src/infiniop-test/src/ops/reciprocal.cpp b/src/infiniop-test/src/ops/reciprocal.cpp deleted file mode 100644 index ba6e1e622..000000000 --- a/src/infiniop-test/src/ops/reciprocal.cpp +++ /dev/null @@ -1,104 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::reciprocal { -struct Test::Attributes { - std::shared_ptr x; - std::shared_ptr y; - std::shared_ptr ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (tensors.find("x") == tensors.end() - || tensors.find("y") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->x = tensors["x"]; - test->_attributes->y = tensors["y"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - infiniopReciprocalDescriptor_t op_desc; - auto x = _attributes->x->to(device, device_id); - auto y = _attributes->y->to(device, device_id); - CHECK_OR(infiniopCreateReciprocalDescriptor(handle, &op_desc, - y->desc(), - x->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - - size_t workspace_size; - CHECK_OR(infiniopGetReciprocalWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - - CHECK_OR(infiniopReciprocal(op_desc, workspace, workspace_size, - y->data(), - x->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(y, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopReciprocal( - op_desc, workspace, workspace_size, - y->data(), - x->data(), - nullptr); - }, - warm_ups, iterations); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} - -std::vector Test::tensor_names() { - return {"x", "y", "ans"}; -} - -std::vector Test::output_names() { - return {"y"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- x: " << _attributes->x->info() << std::endl; - oss << "- y: " << _attributes->y->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::reciprocal diff --git a/src/infiniop-test/src/ops/rms_norm.cpp b/src/infiniop-test/src/ops/rms_norm.cpp deleted file mode 100644 index 8359a4536..000000000 --- a/src/infiniop-test/src/ops/rms_norm.cpp +++ /dev/null @@ -1,128 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::rms_norm { -struct Test::Attributes { - float epsilon; - std::shared_ptr x; - std::shared_ptr w; - std::shared_ptr ans; - std::shared_ptr y; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - - if (attributes.find("epsilon") == attributes.end() - || tensors.find("x") == tensors.end() - || tensors.find("w") == tensors.end() - || tensors.find("ans") == tensors.end() - || tensors.find("y") == tensors.end()) { - throw std::runtime_error("Invalid Test: Missing attributes or tensors"); - } - - test->_attributes->epsilon = *reinterpret_cast(attributes["epsilon"].data()); - - test->_attributes->x = tensors["x"]; - test->_attributes->w = tensors["w"]; - test->_attributes->ans = tensors["ans"]; - test->_attributes->y = tensors["y"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, - size_t warm_ups, size_t iterations) { - - infiniopRMSNormDescriptor_t op_desc; - CHECK_OR(infiniopCreateRMSNormDescriptor(handle, &op_desc, - _attributes->y->desc(), - _attributes->x->desc(), - _attributes->w->desc(), - _attributes->epsilon), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create RMSNorm descriptor")); - - auto x = _attributes->x->to(device, device_id); - auto w = _attributes->w->to(device, device_id); - auto y = _attributes->y->to(device, device_id); - - size_t workspace_size; - CHECK_OR(infiniopGetRMSNormWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size")); - void *workspace = nullptr; - if (workspace_size > 0) { - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace")); - } - - CHECK_OR(infiniopRMSNorm(op_desc, - workspace, workspace_size, - y->data(), - x->data(), - w->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "RMSNorm execution failed")); - - try { - allClose(y, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopRMSNorm(op_desc, - workspace, workspace_size, - y->data(), - x->data(), - w->data(), - nullptr); - }, - warm_ups, iterations); - - if (workspace != nullptr) { - infinirtFree(workspace); - } - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {"epsilon"}; -} - -std::vector Test::tensor_names() { - return {"x", "w", "ans", "y"}; -} - -std::vector Test::output_names() { - return {"y"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- epsilon=" << _attributes->epsilon << std::endl; - oss << "- x: " << _attributes->x->info() << std::endl; - oss << "- w: " << _attributes->w->info() << std::endl; - oss << "- y: " << _attributes->y->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::rms_norm diff --git a/src/infiniop-test/src/ops/rope.cpp b/src/infiniop-test/src/ops/rope.cpp deleted file mode 100644 index 510406234..000000000 --- a/src/infiniop-test/src/ops/rope.cpp +++ /dev/null @@ -1,140 +0,0 @@ -#include "infiniop/ops/rope.h" -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::rope { -struct Test::Attributes { - infiniopRoPEAlgo_t algo; - - std::shared_ptr y; - std::shared_ptr x; - std::shared_ptr pos_ids; - std::shared_ptr sin_table; - std::shared_ptr cos_table; - std::shared_ptr ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - - if (!check_names(attributes, Test::attribute_names()) || tensors.find("y") == tensors.end() - || tensors.find("x") == tensors.end() - || tensors.find("pos_ids") == tensors.end() - || tensors.find("sin_table") == tensors.end() - || tensors.find("cos_table") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->algo = *reinterpret_cast(attributes["algo"].data()); - - test->_attributes->y = tensors["y"]; - test->_attributes->x = tensors["x"]; - test->_attributes->pos_ids = tensors["pos_ids"]; - test->_attributes->sin_table = tensors["sin_table"]; - test->_attributes->cos_table = tensors["cos_table"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - infiniopRoPEDescriptor_t op_desc; - infiniopRoPEAlgo_t algo = _attributes->algo; - auto y = _attributes->y->to(device, device_id); - auto x = _attributes->x->to(device, device_id); - auto pos_ids = _attributes->pos_ids->to(device, device_id); - auto sin_table = _attributes->sin_table->to(device, device_id); - auto cos_table = _attributes->cos_table->to(device, device_id); - - CHECK_OR(infiniopCreateRoPEDescriptor(handle, &op_desc, - y->desc(), - x->desc(), - pos_ids->desc(), - sin_table->desc(), - cos_table->desc(), - algo), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - - size_t workspace_size; - - CHECK_OR(infiniopGetRoPEWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - - void *workspace; - - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - - CHECK_OR(infiniopRoPE(op_desc, workspace, workspace_size, - y->data(), - x->data(), - pos_ids->data(), - sin_table->data(), - cos_table->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(y, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopRoPE( - op_desc, workspace, workspace_size, - y->data(), - x->data(), - pos_ids->data(), - sin_table->data(), - cos_table->data(), - nullptr); - }, - warm_ups, iterations); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {"algo"}; -} - -std::vector Test::tensor_names() { - return {"y", "x", "pos_ids", "sin_table", "cos_table", "ans"}; -} - -std::vector Test::output_names() { - return {"y"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- y: " << _attributes->y->info() << std::endl; - oss << "- x: " << _attributes->x->info() << std::endl; - oss << "- pos_ids: " << _attributes->pos_ids->info() << std::endl; - oss << "- sin_table: " << _attributes->sin_table->info() << std::endl; - oss << "- cos_table: " << _attributes->cos_table->info() << std::endl; - oss << "- algo: " << _attributes->algo << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::rope diff --git a/src/infiniop-test/src/ops/sigmoid.cpp b/src/infiniop-test/src/ops/sigmoid.cpp deleted file mode 100644 index bb3a0f70a..000000000 --- a/src/infiniop-test/src/ops/sigmoid.cpp +++ /dev/null @@ -1,103 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::sigmoid { -struct Test::Attributes { - std::shared_ptr x; - std::shared_ptr y; - std::shared_ptr ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (tensors.find("x") == tensors.end() - || tensors.find("y") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->x = tensors["x"]; - test->_attributes->y = tensors["y"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - infiniopSigmoidDescriptor_t op_desc; - auto x = _attributes->x->to(device, device_id); - auto y = _attributes->y->to(device, device_id); - CHECK_OR(infiniopCreateSigmoidDescriptor(handle, &op_desc, - y->desc(), - x->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - size_t workspace_size; - CHECK_OR(infiniopGetSigmoidWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - CHECK_OR(infiniopSigmoid(op_desc, workspace, workspace_size, - y->data(), - x->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(y, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopSigmoid( - op_desc, workspace, workspace_size, - y->data(), - x->data(), - nullptr); - }, - warm_ups, iterations); - - infiniopDestroySigmoidDescriptor(op_desc); - infinirtFree(workspace); - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} - -std::vector Test::tensor_names() { - return {"x", "y", "ans"}; -} - -std::vector Test::output_names() { - return {"y"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- x: " << _attributes->x->info() << std::endl; - oss << "- y: " << _attributes->y->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::sigmoid diff --git a/src/infiniop-test/src/ops/silu.cpp b/src/infiniop-test/src/ops/silu.cpp deleted file mode 100644 index 75684503c..000000000 --- a/src/infiniop-test/src/ops/silu.cpp +++ /dev/null @@ -1,101 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::silu { -struct Test::Attributes { - std::shared_ptr input; - std::shared_ptr output; - std::shared_ptr ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (tensors.find("input") == tensors.end() - || tensors.find("output") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->input = tensors["input"]; - test->_attributes->output = tensors["output"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - infiniopSiluDescriptor_t op_desc; - auto input = _attributes->input->to(device, device_id); - auto output = _attributes->output->to(device, device_id); - CHECK_OR(infiniopCreateSiluDescriptor(handle, &op_desc, - output->desc(), - input->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - size_t workspace_size; - CHECK_OR(infiniopGetSiluWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - CHECK_OR(infiniopSilu(op_desc, workspace, workspace_size, - output->data(), - input->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(output, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopSilu( - op_desc, workspace, workspace_size, - output->data(), - input->data(), - nullptr); - }, - warm_ups, iterations); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} - -std::vector Test::tensor_names() { - return {"input", "output", "ans"}; -} - -std::vector Test::output_names() { - return {"output"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- input: " << _attributes->input->info() << std::endl; - oss << "- output: " << _attributes->output->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::silu diff --git a/src/infiniop-test/src/ops/sub.cpp b/src/infiniop-test/src/ops/sub.cpp deleted file mode 100644 index 6bb1fd1eb..000000000 --- a/src/infiniop-test/src/ops/sub.cpp +++ /dev/null @@ -1,109 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::sub { -struct Test::Attributes { - std::shared_ptr a; - std::shared_ptr b; - std::shared_ptr c; - std::shared_ptr ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (tensors.find("a") == tensors.end() - || tensors.find("b") == tensors.end() - || tensors.find("c") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->a = tensors["a"]; - test->_attributes->b = tensors["b"]; - test->_attributes->c = tensors["c"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - infiniopSubDescriptor_t op_desc; - auto a = _attributes->a->to(device, device_id); - auto b = _attributes->b->to(device, device_id); - auto c = _attributes->c->to(device, device_id); - CHECK_OR(infiniopCreateSubDescriptor(handle, &op_desc, - c->desc(), - a->desc(), - b->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - size_t workspace_size; - CHECK_OR(infiniopGetSubWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - CHECK_OR(infiniopSub(op_desc, workspace, workspace_size, - c->data(), - a->data(), - b->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(c, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopSub( - op_desc, workspace, workspace_size, - c->data(), - a->data(), - b->data(), - nullptr); - }, - warm_ups, iterations); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} - -std::vector Test::tensor_names() { - return {"a", "b", "c", "ans"}; -} - -std::vector Test::output_names() { - return {"c"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- a: " << _attributes->a->info() << std::endl; - oss << "- b: " << _attributes->b->info() << std::endl; - oss << "- c: " << _attributes->c->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::sub diff --git a/src/infiniop-test/src/ops/swiglu.cpp b/src/infiniop-test/src/ops/swiglu.cpp deleted file mode 100644 index 96b75efc5..000000000 --- a/src/infiniop-test/src/ops/swiglu.cpp +++ /dev/null @@ -1,104 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::swiglu { -struct Test::Attributes { - std::shared_ptr a; - std::shared_ptr b; - std::shared_ptr ans; - std::shared_ptr c; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - - if (tensors.find("a") == tensors.end() - || tensors.find("b") == tensors.end() - || tensors.find("c") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->a = tensors["a"]; - test->_attributes->b = tensors["b"]; - test->_attributes->c = tensors["c"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - infiniopSwiGLUDescriptor_t op_desc; - auto a = _attributes->a->to(device, device_id); - auto b = _attributes->b->to(device, device_id); - auto c = _attributes->c->to(device, device_id); - CHECK_OR(infiniopCreateSwiGLUDescriptor(handle, &op_desc, - c->desc(), - a->desc(), - b->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - size_t workspace_size; - CHECK_OR(infiniopGetSwiGLUWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - CHECK_OR(infiniopSwiGLU(op_desc, workspace, workspace_size, c->data(), a->data(), b->data(), nullptr), - return TEST_FAILED(OP_CREATION_FAILED, "Failed during execution.")); - try { - allClose(c, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopSwiGLU( - op_desc, - workspace, - workspace_size, - c->data(), - a->data(), - b->data(), - nullptr); - }, - warm_ups, iterations); - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} -std::vector Test::tensor_names() { - return {"a", "b", "c", "ans"}; -} - -std::vector Test::output_names() { - return {"c"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- a: " << _attributes->a->info() << std::endl; - oss << "- b: " << _attributes->b->info() << std::endl; - oss << "- c: " << _attributes->c->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::swiglu diff --git a/src/infiniop-test/src/ops/topkrouter.cpp b/src/infiniop-test/src/ops/topkrouter.cpp deleted file mode 100644 index 367087fd6..000000000 --- a/src/infiniop-test/src/ops/topkrouter.cpp +++ /dev/null @@ -1,130 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::topkrouter { -struct Test::Attributes { - std::shared_ptr values; - std::shared_ptr indices; - std::shared_ptr x; - std::shared_ptr correction_bias; - float routed_scaling_factor; - int topk; - std::shared_ptr lable_values; - std::shared_ptr lable_indices; -}; - -std::shared_ptr Test::build(std::unordered_map> attributes, - std::unordered_map> tensors, double rtol, - double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - - if (attributes.find("routed_scaling_factor") == attributes.end() || attributes.find("topk") == attributes.end() || tensors.find("values") == tensors.end() || tensors.find("indices") == tensors.end() || tensors.find("x") == tensors.end() || tensors.find("correction_bias") == tensors.end() || tensors.find("lable_values") == tensors.end() || tensors.find("lable_indices") == tensors.end()) { - throw std::runtime_error("Invalid Test: Missing attributes or tensors"); - } - - test->_attributes->values = tensors["values"]; - test->_attributes->indices = tensors["indices"]; - test->_attributes->x = tensors["x"]; - test->_attributes->correction_bias = tensors["correction_bias"]; - - test->_attributes->routed_scaling_factor = *reinterpret_cast(attributes["routed_scaling_factor"].data()); - test->_attributes->topk = *reinterpret_cast(attributes["topk"].data()); - - test->_attributes->lable_values = tensors["lable_values"]; - test->_attributes->lable_indices = tensors["lable_indices"]; - - return test; -} - -std::shared_ptr Test::run(infiniopHandle_t handle, infiniDevice_t device, int device_id, - size_t warm_ups, size_t iterations) { - infiniopTopkrouterDescriptor_t op_desc; - CHECK_OR(infiniopCreateTopkrouterDescriptor(handle, &op_desc, _attributes->x->desc(), - _attributes->correction_bias->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create Topkrouter descriptor")); - - // - auto values = _attributes->values->to(device, device_id); - auto indices = _attributes->indices->to(device, device_id); - auto x = _attributes->x->to(device, device_id); - auto correction_bias = _attributes->correction_bias->to(device, device_id); - - float routed_scaling_factor = _attributes->routed_scaling_factor; - int topk = _attributes->topk; - - size_t workspace_size; - CHECK_OR(infiniopGetTopkrouterWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size")); - void *workspace = nullptr; - if (workspace_size > 0) { - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace")); - } - - CHECK_OR(infiniopTopkrouter(op_desc, workspace, workspace_size, values->data(), indices->data(), x->data(), - correction_bias->data(), routed_scaling_factor, topk, nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Topkrouter execution failed")); - - try { - allClose(values, _attributes->lable_values, _rtol, _atol); - allClose(indices, _attributes->lable_indices, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopTopkrouter(op_desc, workspace, workspace_size, values->data(), indices->data(), x->data(), - correction_bias->data(), routed_scaling_factor, topk, nullptr); - }, - warm_ups, iterations); - - if (workspace != nullptr) { - infinirtFree(workspace); - } - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {"routed_scaling_factor", "topk"}; -} - -std::vector Test::tensor_names() { - return {"values", "indices", "x", "correction_bias", "lable_values", "lable_indices"}; -} - -std::vector Test::output_names() { - return {"values", "indices"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- routed_scaling_factor=" << _attributes->routed_scaling_factor << std::endl; - oss << "- topk=" << _attributes->topk << std::endl; - - oss << "- values: " << _attributes->values->info() << std::endl; - oss << "- indices: " << _attributes->indices->info() << std::endl; - oss << "- x: " << _attributes->x->info() << std::endl; - oss << "- correction_bias: " << _attributes->correction_bias->info() << std::endl; - - oss << "- lable_values: " << _attributes->lable_values->info() << std::endl; - oss << "- lable_indices: " << _attributes->lable_indices->info() << std::endl; - - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::topkrouter diff --git a/src/infiniop-test/src/ops/topksoftmax.cpp b/src/infiniop-test/src/ops/topksoftmax.cpp deleted file mode 100644 index f6b0af7a4..000000000 --- a/src/infiniop-test/src/ops/topksoftmax.cpp +++ /dev/null @@ -1,122 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::topksoftmax { -struct Test::Attributes { - std::shared_ptr values; - std::shared_ptr indices; - std::shared_ptr x; - int topk; - bool norm; - std::shared_ptr lable_values; - std::shared_ptr lable_indices; -}; - -std::shared_ptr Test::build(std::unordered_map> attributes, - std::unordered_map> tensors, double rtol, - double atol) { - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - - if (attributes.find("topk") == attributes.end() || attributes.find("norm") == attributes.end() || tensors.find("values") == tensors.end() || tensors.find("indices") == tensors.end() || tensors.find("x") == tensors.end() || tensors.find("lable_values") == tensors.end() || tensors.find("lable_indices") == tensors.end()) { - throw std::runtime_error("Invalid Test: Missing attributes or tensors"); - } - - test->_attributes->values = tensors["values"]; - test->_attributes->indices = tensors["indices"]; - test->_attributes->x = tensors["x"]; - - test->_attributes->topk = *reinterpret_cast(attributes["topk"].data()); - test->_attributes->norm = *reinterpret_cast(attributes["norm"].data()); - test->_attributes->lable_values = tensors["lable_values"]; - test->_attributes->lable_indices = tensors["lable_indices"]; - - return test; -} - -std::shared_ptr Test::run(infiniopHandle_t handle, infiniDevice_t device, int device_id, - size_t warm_ups, size_t iterations) { - infiniopTopksoftmaxDescriptor_t op_desc; - CHECK_OR(infiniopCreateTopksoftmaxDescriptor(handle, &op_desc, _attributes->x->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create Topksoftmax descriptor")); - - // - auto values = _attributes->values->to(device, device_id); - auto indices = _attributes->indices->to(device, device_id); - auto x = _attributes->x->to(device, device_id); - - int topk = _attributes->topk; - bool norm = _attributes->norm; - - size_t workspace_size; - CHECK_OR(infiniopGetTopksoftmaxWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size")); - void *workspace = nullptr; - if (workspace_size > 0) { - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace")); - } - - CHECK_OR(infiniopTopksoftmax(op_desc, workspace, workspace_size, values->data(), indices->data(), x->data(), topk, norm, nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Topksoftmax execution failed")); - - try { - allClose(values, _attributes->lable_values, _rtol, _atol); - allClose(indices, _attributes->lable_indices, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - elapsed_time = benchmark( - [=]() { - infiniopTopksoftmax(op_desc, workspace, workspace_size, values->data(), indices->data(), x->data(), topk, norm, nullptr); - }, - warm_ups, iterations); - - if (workspace != nullptr) { - infinirtFree(workspace); - } - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {"topk", "norm"}; -} - -std::vector Test::tensor_names() { - return {"values", "indices", "x", "lable_values", "lable_indices"}; -} - -std::vector Test::output_names() { - return {"values", "indices"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - - oss << "- topk=" << _attributes->topk << std::endl; - oss << "- norm=" << _attributes->norm << std::endl; - - oss << "- values: " << _attributes->values->info() << std::endl; - oss << "- indices: " << _attributes->indices->info() << std::endl; - oss << "- x: " << _attributes->x->info() << std::endl; - - oss << "- lable_values: " << _attributes->lable_values->info() << std::endl; - oss << "- lable_indices: " << _attributes->lable_indices->info() << std::endl; - - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::topksoftmax diff --git a/src/infiniop-test/src/ops/zeros.cpp b/src/infiniop-test/src/ops/zeros.cpp deleted file mode 100644 index c4d37fcf4..000000000 --- a/src/infiniop-test/src/ops/zeros.cpp +++ /dev/null @@ -1,103 +0,0 @@ -#include "ops.hpp" -#include "utils.hpp" -#include -#include -#include - -namespace infiniop_test::zeros { -struct Test::Attributes { - std::shared_ptr x; - std::shared_ptr y; - std::shared_ptr ans; -}; - -std::shared_ptr Test::build( - std::unordered_map> attributes, - std::unordered_map> tensors, - double rtol, double atol) { - - auto test = std::shared_ptr(new Test(rtol, atol)); - test->_attributes = new Attributes(); - if (tensors.find("x") == tensors.end() - || tensors.find("y") == tensors.end() - || tensors.find("ans") == tensors.end()) { - throw std::runtime_error("Invalid Test"); - } - - test->_attributes->x = tensors["x"]; - test->_attributes->y = tensors["y"]; - test->_attributes->ans = tensors["ans"]; - - return test; -} - -std::shared_ptr Test::run( - infiniopHandle_t handle, infiniDevice_t device, int device_id, size_t warm_ups, size_t iterations) { - - infiniopZerosDescriptor_t op_desc; - auto x = _attributes->x->to(device, device_id); - auto y = _attributes->y->to(device, device_id); - CHECK_OR(infiniopCreateZerosDescriptor(handle, &op_desc, - y->desc(), - x->desc()), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to create op descriptor.")); - size_t workspace_size; - CHECK_OR(infiniopGetZerosWorkspaceSize(op_desc, &workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to get workspace size.")); - void *workspace; - CHECK_OR(infinirtMalloc(&workspace, workspace_size), - return TEST_FAILED(OP_CREATION_FAILED, "Failed to allocate workspace.")); - CHECK_OR(infiniopZeros(op_desc, workspace, workspace_size, - y->data(), - x->data(), - nullptr), - return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution.")); - - try { - allClose(y, _attributes->ans, _rtol, _atol); - } catch (const std::exception &e) { - return TEST_FAILED(RESULT_INCORRECT, e.what()); - } - - double elapsed_time = 0.; - - elapsed_time = benchmark( - [=]() { - infiniopZeros( - op_desc, workspace, workspace_size, - y->data(), - x->data(), - nullptr); - }, - warm_ups, iterations); - - return TEST_PASSED(elapsed_time); -} - -std::vector Test::attribute_names() { - return {}; -} - -std::vector Test::tensor_names() { - return {"x", "y", "ans"}; -} - -std::vector Test::output_names() { - return {"y"}; -} - -std::string Test::toString() const { - std::ostringstream oss; - oss << op_name() << std::endl; - oss << "- x: " << _attributes->x->info() << std::endl; - oss << "- y: " << _attributes->y->info() << std::endl; - oss << std::scientific << std::setprecision(2); - oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl; - return oss.str(); -} - -Test::~Test() { - delete _attributes; -} - -} // namespace infiniop_test::zeros diff --git a/src/infiniop-test/src/tensor.cpp b/src/infiniop-test/src/tensor.cpp deleted file mode 100644 index 2f1c3629a..000000000 --- a/src/infiniop-test/src/tensor.cpp +++ /dev/null @@ -1,324 +0,0 @@ -#include "tensor.hpp" -#include "gguf.hpp" -#include "utils.hpp" -#include -#include -#include - -template -void printData(const T *data, const std::vector &shape, const std::vector &strides, size_t dim) { - if (dim == shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - std::cout << *(data + i * strides[dim]) << " "; - } - std::cout << std::endl; - } else if (dim < shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - printData(data + i * strides[dim], shape, strides, dim + 1); - std::cout << std::endl; - } - } -} - -// The type int8_t is represented by signed char, with a range of –128 to 127. -// It may contain non-printable characters and thus cannot be printed directly. -template <> -void printData(const int8_t *data, const std::vector &shape, - const std::vector &strides, size_t dim) { - if (dim == shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - std::cout << static_cast(*(data + i * strides[dim])) << " "; - } - std::cout << std::endl; - } else if (dim < shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - printData(data + i * strides[dim], shape, strides, dim + 1); - std::cout << std::endl; - } - } -} - -template <> -void printData(const bf16_t *data, const std::vector &shape, - const std::vector &strides, size_t dim) { - if (dim == shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - std::cout << utils::cast(*(data + i * strides[dim])) << " "; - } - std::cout << std::endl; - } else if (dim < shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - printData(data + i * strides[dim], shape, strides, dim + 1); - std::cout << std::endl; - } - } -} - -template <> -void printData(const fp16_t *data, const std::vector &shape, - const std::vector &strides, size_t dim) { - if (dim == shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - std::cout << utils::cast(*(data + i * strides[dim])) << " "; - } - std::cout << std::endl; - } else if (dim < shape.size() - 1) { - for (size_t i = 0; i < shape[dim]; i++) { - printData(data + i * strides[dim], shape, strides, dim + 1); - std::cout << std::endl; - } - } -} - -// Calculate memory size & offset given shape & strides -inline void calculateTensorMemory(size_t &size, size_t &offset, - std::vector shape, - std::vector strides, - size_t data_size) { - size_t ndim = shape.size(); - offset = 0; - size = 0; - for (size_t i = 0; i < ndim; i++) { - if (shape[i] == 0) { - offset = 0; - size = 0; - return; - } - if (strides[i] > 0) { - size += (shape[i] - 1) * strides[i] * data_size; - } else if (strides[i] < 0) { - offset += (shape[i] - 1) * (size_t)(-strides[i]) * data_size; - } - } - size = offset + size + data_size; -} - -namespace infiniop_test { - -Memory::Memory(size_t size, infiniDevice_t device, int device_id) { - _file_mapping = nullptr; - _device = device; - _device_id = device_id; - _size = size; - if (device == INFINI_DEVICE_CPU) { - _ptr = std::malloc(size); - } else { - CHECK_OR(infinirtSetDevice(_device, _device_id), throw std::runtime_error("Error Creating Memory: set device")); - CHECK_OR(infinirtMalloc(&_ptr, _size), throw std::runtime_error("Error Creating Memory: malloc")); - } -} - -Memory::Memory(const std::shared_ptr &file_mapping, void *ptr, size_t size) { - _device = INFINI_DEVICE_CPU; - _device_id = 0; - _size = size; - _ptr = ptr; - _file_mapping = file_mapping; -} - -Memory::~Memory() { - // if memory does not map to a file, free it manually - if (_file_mapping == nullptr) { - if (_device == INFINI_DEVICE_CPU) { - std::free(_ptr); - } else { - infinirtSetDevice(_device, _device_id); - infinirtFree(_ptr); - } - } -} - -void *Tensor::data() const { - return (char *)(_memory->ptr()) + _offset; -} - -Tensor::Tensor(const GGUFTensorInfo *info, - const void *ggml_ptr, - const GGUFKeyValue *shape_meta, - const GGUFKeyValue *strides_meta, - bool isOutput) { - - _ggml_type = info->ggml_type; - _offset = 0; - size_t ndim = static_cast(info->ndim); - // `_shape`存储真实的tensor形状（来自shape_meta），`temp_shape`存储用于rearrange和计算内存的tensor形状 - _shape = std::vector(ndim); - std::vector temp_shape(ndim); - _strides = std::vector(ndim); - std::vector contiguous_strides(ndim); - for (size_t i = 0; i < ndim; i++) { - temp_shape[i] = static_cast(info->shape[ndim - 1 - i]); - if (i == 0) { - contiguous_strides[ndim - 1] = (ptrdiff_t)1; - } else { - contiguous_strides[ndim - 1 - i] = (ptrdiff_t)info->shape[i - 1] * contiguous_strides[ndim - i]; - } - if (isOutput) { - contiguous_strides[i] = (ptrdiff_t)0; - } - } - - if (strides_meta == nullptr) { - for (size_t i = 0; i < ndim; i++) { - _strides[i] = contiguous_strides[i]; - } - } else { - for (size_t i = 0; i < ndim; i++) { - if (strides_meta->gguf_type == GGUF_TYPE_INT64) { - _strides[i] = (ptrdiff_t)(reinterpret_cast( - strides_meta->value.data())[ndim - 1 - i]); - } else if (strides_meta->gguf_type == GGUF_TYPE_INT32) { - _strides[i] = (ptrdiff_t)(reinterpret_cast( - strides_meta->value.data())[ndim - 1 - i]); - } else { - throw std::runtime_error("Error Creating Tensor: Unsupported strides type"); - } - } - } - - if (isOutput) { - if (shape_meta == nullptr) { - throw std::runtime_error("Error Creating Tensor: shape_meta cannot be null for output tensor"); - } - for (size_t i = 0; i < ndim; i++) { - if (shape_meta->gguf_type == GGUF_TYPE_INT64) { - int64_t val = reinterpret_cast(shape_meta->value.data())[i]; - if (val < 0) { - throw std::runtime_error("Shape must be non-negative"); - } - temp_shape[i] = static_cast(val); - } else if (shape_meta->gguf_type == GGUF_TYPE_INT32) { - int32_t val = reinterpret_cast(shape_meta->value.data())[i]; - if (val < 0) { - throw std::runtime_error("Shape must be non-negative"); - } - temp_shape[i] = static_cast(val); - } else { - throw std::runtime_error("Error Creating Tensor: Unsupported shape type"); - } - } - } - infiniopCreateTensorDescriptor(&_desc, ndim, temp_shape.data(), _strides.data(), ggmlTypeToInfiniType(_ggml_type)); - size_t size; - calculateTensorMemory(size, _offset, temp_shape, _strides, ggmlTypeSize(_ggml_type)); - _memory = std::make_shared(size, INFINI_DEVICE_CPU, 0); - utils::rearrange( - (char *)_memory->ptr() + _offset, - (char *)ggml_ptr + info->data_offset, - temp_shape.data(), - _strides.data(), - contiguous_strides.data(), - ndim, - ggmlTypeSize(_ggml_type)); - - if (shape_meta == nullptr) { - _shape = temp_shape; - } else { - for (size_t i = 0; i < ndim; i++) { - if (shape_meta->gguf_type == GGUF_TYPE_INT64) { - int64_t val = reinterpret_cast(shape_meta->value.data())[i]; - if (val < 0) { - throw std::runtime_error("Shape must be non-negative"); - } - _shape[i] = static_cast(val); - } else if (shape_meta->gguf_type == GGUF_TYPE_INT32) { - int32_t val = reinterpret_cast(shape_meta->value.data())[i]; - if (val < 0) { - throw std::runtime_error("Shape must be non-negative"); - } - _shape[i] = static_cast(val); - } else { - throw std::runtime_error("Error Creating Tensor: Unsupported shape type"); - } - } - } -} - -Tensor::Tensor(std::shared_ptr memory, size_t offset, - const std::vector &shape, - const std::vector &strides, - GGML_TYPE dtype) : _memory(memory), _shape(shape), _strides(strides), _offset(offset), _ggml_type(dtype) { - infiniopCreateTensorDescriptor(&_desc, shape.size(), shape.data(), strides.data(), ggmlTypeToInfiniType(dtype)); -} - -std::shared_ptr Tensor::to(infiniDevice_t device, int device_id) const { - if (device == _memory->device() && (device_id == _memory->device_id() || device == INFINI_DEVICE_CPU)) { - return std::make_shared(_memory, _offset, _shape, _strides, _ggml_type); - } - std::shared_ptr memory; - if (device == INFINI_DEVICE_CPU) { - memory = std::make_shared(_memory->size(), INFINI_DEVICE_CPU, 0); - CHECK_OR(infinirtSetDevice(_memory->device(), _memory->device_id()), throw std::runtime_error("Error Tensor::to: set device")); - CHECK_OR(infinirtMemcpy(memory->ptr(), _memory->ptr(), _memory->size(), INFINIRT_MEMCPY_D2H), throw std::runtime_error("Error Tensor::to: cpy")); - } else if (_memory->device() == INFINI_DEVICE_CPU) { - memory = std::make_shared(_memory->size(), device, device_id); - CHECK_OR(infinirtMemcpy(memory->ptr(), _memory->ptr(), _memory->size(), INFINIRT_MEMCPY_H2D), throw std::runtime_error("Error Tensor::to: cpy")); - } else { - return to(INFINI_DEVICE_CPU, 0)->to(device, device_id); - } - return std::make_shared(memory, _offset, _shape, _strides, _ggml_type); -} - -void Tensor::debug() const { - auto tensor = to(INFINI_DEVICE_CPU, 0); - std::cout << "Tensor: " << tensor->info() << std::endl; - switch (_ggml_type) { - case GGML_TYPE_BF16: - printData((bf16_t *)(tensor->data()), _shape, _strides, 0); - case GGML_TYPE_F16: - printData((fp16_t *)(tensor->data()), _shape, _strides, 0); - break; - case GGML_TYPE_F32: - printData((float *)(tensor->data()), _shape, _strides, 0); - break; - case GGML_TYPE_F64: - printData((double *)(tensor->data()), _shape, _strides, 0); - break; - case GGML_TYPE_Q8_K: - printData((bool *)(tensor->data()), _shape, _strides, 0); - break; - case GGML_TYPE_I8: - printData((int8_t *)(tensor->data()), _shape, _strides, 0); - break; - case GGML_TYPE_I16: - printData((int16_t *)(tensor->data()), _shape, _strides, 0); - break; - case GGML_TYPE_I32: - printData((int32_t *)(tensor->data()), _shape, _strides, 0); - break; - case GGML_TYPE_I64: - printData((int64_t *)(tensor->data()), _shape, _strides, 0); - break; - default: - std::cout << "Unsupported GGML type" << std::endl; - break; - } -} - -std::string Tensor::info() const { - std::ostringstream oss; - oss << "Shape: ["; - for (size_t i = 0; i < _shape.size(); ++i) { - oss << _shape[i]; - if (i != _shape.size() - 1) { - oss << ", "; - } - } - oss << "]"; - oss << ", Strides: ["; - for (size_t i = 0; i < _strides.size(); ++i) { - oss << _strides[i]; - if (i != _strides.size() - 1) { - oss << ", "; - } - } - oss << "]"; - oss << ", Type: " << GGML_TYPE_NAME[_ggml_type]; - - return oss.str(); -} - -Tensor::~Tensor() { - infiniopDestroyTensorDescriptor(_desc); -} -} // namespace infiniop_test diff --git a/src/infiniop-test/src/test.cpp b/src/infiniop-test/src/test.cpp deleted file mode 100644 index e312ac5f5..000000000 --- a/src/infiniop-test/src/test.cpp +++ /dev/null @@ -1,233 +0,0 @@ -#include "ops.hpp" -#include "tensor.hpp" -#include "utils.hpp" -#include -#include -#include -#include -#include - -namespace infiniop_test { -std::unordered_map TEST_BUILDERS = TEST_BUILDER_MAPPINGS; - -std::string Result::toString() const { - std::ostringstream oss; - oss << "Status: "; - switch (_status) { - case TestStatus::PASS: - oss << GREEN << "PASS" << RESET; - break; - case TestStatus::TEST_INIT_FAILED: - oss << RED << "INVALID TEST" << RESET; - break; - case TestStatus::OP_CREATION_FAILED: - oss << RED << "OP CREATION FAILED" << RESET; - break; - case TestStatus::OP_EXECUTION_FAILED: - oss << RED << "EXECUTION FAILED" << RESET; - break; - case TestStatus::RESULT_INCORRECT: - oss << RED << "WRONG ANSWER" << RESET; - break; - default: - oss << YELLOW << "SKIPPED" << RESET; - break; - } - oss << std::endl; - oss << "Description: " << _description << std::endl; - if (_time > 0.) { - oss << "Time: " << _time << " us" << std::endl; - } else { - oss << "Time: N/A" << std::endl; - } - if (_error_message.size() > 0) { - oss << "Error: " << _error_message << std::endl; - } - return oss.str(); -} - -std::vector> runAllTests(const GGUFFileReader &gguf_reader, - infiniDevice_t device, int device_id, - size_t warm_ups, size_t iterations, - double rtol, double atol) { - auto meta = gguf_reader.getAttributeMap(); - auto count_meta = meta.find("test_count"); - if (count_meta == meta.end()) { - throw std::runtime_error("Invalid GGUF file: missing test_count attribute"); - } - size_t count = *(size_t *)(count_meta->second->value.data()); - std::cout << "Found " << count << " tests" << std::endl; - auto results = std::vector>(count); - try { - for (size_t i = 0; i < count; i++) { - results[i] = runTest(gguf_reader, device, device_id, warm_ups, iterations, rtol, atol, i); - } - } catch (const std::exception &e) { - std::cerr << "Error: " << e.what() << std::endl; - } - - return results; -} - -std::shared_ptr runTest(const GGUFFileReader &gguf_reader, - infiniDevice_t device, int device_id, - size_t warm_ups, size_t iterations, - double rtol, double atol, size_t test_id) { - auto meta = gguf_reader.getAttributeMap(); - auto tensor_info = gguf_reader.getTensorInfoMap(); - auto name_meta = meta.find("test." + std::to_string(test_id) + ".op_name"); - if (name_meta != meta.end()) { - std::string op_name(name_meta->second->value.begin(), name_meta->second->value.end()); - auto builder = TEST_BUILDERS.find(op_name)->second; - auto attrs = std::unordered_map>(); - auto tensors = std::unordered_map>(); - infiniopHandle_t handle; - CHECK_OR(infinirtSetDevice(device, device_id), throw std::runtime_error("Failed to set device")); - CHECK_OR(infiniopCreateHandle(&handle), throw std::runtime_error("Failed to create handle")); - for (auto attr_name : builder.attribute_names) { - auto attr = meta.find("test." + std::to_string(test_id) + "." + attr_name); - if (attr != meta.end()) { - attrs[attr_name] = attr->second->value; - } - } - - for (auto tensor_name : builder.tensor_names) { - auto info = tensor_info.find("test." + std::to_string(test_id) + "." + tensor_name); - if (info != tensor_info.end()) { - auto shape = meta.find("test." + std::to_string(test_id) + "." + tensor_name + ".shape"); - auto strides = meta.find("test." + std::to_string(test_id) + "." + tensor_name + ".strides"); - bool is_output = std::find(builder.output_names.begin(), builder.output_names.end(), tensor_name) != builder.output_names.end(); - tensors[tensor_name] = std::make_shared( - info->second.get(), - gguf_reader.getGgmlStart(), - shape != meta.end() ? shape->second.get() : nullptr, - strides != meta.end() ? strides->second.get() : nullptr, - is_output); - } - } - std::shared_ptr test; - try { - test = builder.build(attrs, tensors, rtol, atol); - } catch (const std::exception &e) { - return TEST_INIT_FAILED(op_name + "/n" + e.what()); - } - - std::shared_ptr result; - try { - result = test->run(handle, device, device_id, warm_ups, iterations); - } catch (const std::exception &e) { - return TEST_INIT_FAILED(op_name + "/n" + e.what()); - } - - CHECK_OR(infiniopDestroyHandle(handle), throw std::runtime_error("Failed to destroy handle")); - return result; - } - return TEST_INIT_FAILED(""); -} - -void incrementOffset(ptrdiff_t &offset_1, const std::vector &strides_1, size_t data_size_1, - ptrdiff_t &offset_2, const std::vector &strides_2, size_t data_size_2, - std::vector &counter, const std::vector &shape) { - for (ptrdiff_t d = shape.size() - 1; d >= 0; d--) { - counter[d] += 1; - offset_1 += strides_1[d] * data_size_1; - offset_2 += strides_2[d] * data_size_2; - if (counter[d] < shape[d]) { - break; - } - counter[d] = 0; - offset_1 -= shape[d] * strides_1[d] * data_size_1; - offset_2 -= shape[d] * strides_2[d] * data_size_2; - } -} - -void allClose(std::shared_ptr actual_, std::shared_ptr expected_, double rtol, double atol) { - auto actual = actual_->to(INFINI_DEVICE_CPU); - auto expected = expected_->to(INFINI_DEVICE_CPU); - auto shape = actual->shape(); - if (shape != expected->shape()) { - throw std::runtime_error("Shape mismatch."); - } - auto ndim = shape.size(); - size_t total = std::accumulate(shape.begin(), shape.end(), (size_t)1, std::multiplies()); - auto counter = std::vector(ndim, 0); - ptrdiff_t actual_offset = 0, - expected_offset = 0; - size_t num_failed = 0; - std::string first_failed_msg; - for (size_t i = 0; i < total; i++) { - double a_ = getVal((char *)actual->data() + actual_offset, actual->ggml_type()); - double e_ = getVal((char *)expected->data() + expected_offset, expected->ggml_type()); - if (std::fabs(a_ - e_) > atol && std::fabs(a_ - e_) > rtol * std::fmax(std::fabs(a_), std::fabs(e_))) { - if (num_failed == 0) { - first_failed_msg = "First failed at index " + std::to_string(i) + " with value " + std::to_string(a_) + " but should be " + std::to_string(e_) + "."; - } - num_failed++; - } - incrementOffset(actual_offset, actual->strides(), ggmlTypeSize(actual->ggml_type()), - expected_offset, expected->strides(), ggmlTypeSize(expected->ggml_type()), - counter, shape); - } - if (num_failed > 0) { - throw std::runtime_error(std::to_string(num_failed) + " out of " + std::to_string(total) + " values failed. " + first_failed_msg); - } -} - -void allEqual(std::shared_ptr actual_, std::shared_ptr expected_) { - auto actual = actual_->to(INFINI_DEVICE_CPU); - auto expected = expected_->to(INFINI_DEVICE_CPU); - auto ggml_type = actual->ggml_type(); - auto shape = actual->shape(); - - if (ggml_type != expected->ggml_type()) { - throw std::runtime_error("Data type mismatch."); - } - if (shape != expected->shape()) { - throw std::runtime_error("Shape mismatch."); - } - - auto ndim = shape.size(); - size_t total = std::accumulate(shape.begin(), shape.end(), (size_t)1, std::multiplies()); - auto counter = std::vector(ndim, 0); - ptrdiff_t actual_offset = 0, - expected_offset = 0; - size_t num_failed = 0; - std::string first_failed_msg; - for (size_t i = 0; i < total; i++) { - char *a_ = (char *)actual->data() + actual_offset, - *e_ = (char *)expected->data() + expected_offset; - if (std::memcmp(a_, e_, ggmlSizeOf(ggml_type))) { - if (num_failed == 0) { - first_failed_msg = "First failed at index " + std::to_string(i); - } - num_failed++; - } - incrementOffset(actual_offset, actual->strides(), ggmlTypeSize(actual->ggml_type()), - expected_offset, expected->strides(), ggmlTypeSize(expected->ggml_type()), - counter, shape); - } - if (num_failed > 0) { - throw std::runtime_error(std::to_string(num_failed) + " out of " + std::to_string(total) + " values failed. " + first_failed_msg); - } -} - -double benchmark(std::function func, size_t warmups, size_t iterations) { - if (iterations == 0) { - return 0.0; - } - for (size_t i = 0; i < warmups; ++i) { - func(); - } - infinirtDeviceSynchronize(); - auto start = std::chrono::high_resolution_clock::now(); - for (size_t i = 0; i < iterations; ++i) { - func(); - } - infinirtDeviceSynchronize(); - auto end = std::chrono::high_resolution_clock::now(); - auto duration = std::chrono::duration_cast(end - start); - double average_time = duration.count() / iterations / 1e3; // average in us - - return average_time; -} -} // namespace infiniop_test diff --git a/src/infiniop/README.md b/src/infiniop/README.md deleted file mode 100644 index 0268eb1b4..000000000 --- a/src/infiniop/README.md +++ /dev/null @@ -1,48 +0,0 @@ -# InfiniOP 开发者文档 - -InfiniOP 是 InfiniCore 下属的统一底层算子框架，为相同算子在不同平台提供统一的 C 语言多段式接口。 - -## 开发流程 - -1. 根据算子定义设计算子接口，在 [`InfiniCore文档`](https://github.com/InfiniTensor/InfiniCore-Documentation) 中添加算子文档。提交文档 PR 。 - -2. 在 `include/infiniop/` 中添加算子头文件，并 include 到 `include/infiniop.h` 中。每个算子暴露的接口包括：创建算子描述、获取工作空间大小、执行算子、销毁算子描述。比如： - - ```c - #ifndef __INFINIOP_ADD_API_H__ - #define __INFINIOP_ADD_API_H__ - - #include "../operator_descriptor.h" - - typedef struct InfiniopDescriptor *infiniopAddDescriptor_t; - - __INFINI_C __export infiniStatus_t infiniopCreateAddDescriptor(infiniopHandle_t handle, - infiniopAddDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c, - infiniopTensorDescriptor_t a, - infiniopTensorDescriptor_t b); - - __INFINI_C __export infiniStatus_t infiniopGetAddWorkspaceSize(infiniopAddDescriptor_t desc, size_t *size); - - __INFINI_C __export infiniStatus_t infiniopAdd(infiniopAddDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream); - - __INFINI_C __export infiniStatus_t infiniopDestroyAddDescriptor(infiniopAddDescriptor_t desc); - - #endif - ``` - - 在任何平台都不需要工作空间的算子也可以不提供获取工作空间大小接口。 - -3. 在 `src/infiniop/ops/` 中添加算子实现目录，并在目录中创建 `operator.cc` 文件实现头文件中的接口，并根据硬件环境分发至不同平台的核函数。你还可以在目录中创建该算子在全平台通用的代码，比如 `causal_softmax/info.h` 中就包含了对 Causal Softmax 算子在创建算子描述时的一些通用的信息获取和输入输出检查。像逐元素类的算子除了计算内核以外大部分逻辑都是一样的，你可以使用 `src/infiniop/elementwise/` 中的通用代码快速适配算子。 - -4. 在 `src/infiniop/ops/[op]/[device]/` 中添加平台算子实现。注意复用平台公共代码，比如规约计算（`src/infiniop/reduce/`），开发过程中把未来可复用的代码写在相应公用代码目录里。 - - 一些 CUDA kernel 可以被多个支持 CUDA 的平台公用，可以考虑在头文件中实现，并在多个源文件中使用。 比如 `mul/cuda/kernel.cuh` 中只有 device 测代码，会被多个支持 CUDA 的平台源代码引用。 - -5. 算子实现可以成功编译安装后，在 `test/infiniop/` 中添加单测脚本，与 PyTorch 实现进行正确性和性能比较。你可以仿照已有的测试脚本进行开发，以使用各种通用的测试功能。测例应覆盖算子常用类型和形状。测试成功之后可以将测例添加至 `scripts/python_test.py` 一键测试脚本中（这样 Github 自动测试也会包含该算子）。 diff --git a/src/infiniop/binary/binary.h b/src/infiniop/binary/binary.h deleted file mode 100644 index 73d4a4e94..000000000 --- a/src/infiniop/binary/binary.h +++ /dev/null @@ -1,98 +0,0 @@ -#ifndef __INFINIOP_BINARY_H__ -#define __INFINIOP_BINARY_H__ - -#include "../operator.h" -#include "../tensor.h" -#include - -/** - * 该类的设计基于 matmul.h 中 YdrMaster 设计的 DESCRIPTOR 宏。 - */ - -#define BINARY_DESCRIPTOR(OP, NAMESPACE) \ - \ - namespace op::OP::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - infiniDtype_t _dtype; \ - op::binary::BinaryInfo _info; \ - \ - Descriptor( \ - infiniDtype_t dtype, \ - op::binary::BinaryInfo info, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _dtype(dtype), \ - _info(info) {} \ - \ - public: \ - ~Descriptor(); \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t c_desc, \ - infiniopTensorDescriptor_t a_desc, \ - infiniopTensorDescriptor_t b_desc); \ - \ - infiniStatus_t calculate( \ - void *c, \ - const void *a, \ - const void *b, \ - void *stream) const; \ - }; \ - } - -namespace op::binary { - -// Stores metadata for binary operations on CPU -struct BinaryInfo { - size_t c_data_size; - size_t ndim; - bool contiguous; - bool broadcasted; - std::vector c_shape; - std::vector a_shape; - std::vector b_shape; - std::vector c_strides; - std::vector a_strides; - std::vector b_strides; -}; - -inline infiniStatus_t createBinaryInfo(BinaryInfo &info, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - - if (!c_desc || !a_desc || !b_desc) { - return INFINI_STATUS_BAD_PARAM; - } - - info.c_data_size = c_desc->numel(); - info.ndim = c_desc->ndim(); - info.contiguous = c_desc->isContiguous() && a_desc->isContiguous() && b_desc->isContiguous(); - - // Destination cannot have broadcast setup - if (c_desc->hasBroadcastDim()) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - const bool ndim_match = (c_desc->ndim() == a_desc->ndim()) && (c_desc->ndim() == b_desc->ndim()); - info.broadcasted = !info.contiguous && (!ndim_match || a_desc->hasBroadcastDim() || b_desc->hasBroadcastDim()); - - info.c_shape = std::move(c_desc->shape()); - info.a_shape = std::move(a_desc->shape()); - info.b_shape = std::move(b_desc->shape()); - info.c_strides = std::move(c_desc->strides()); - info.a_strides = std::move(a_desc->strides()); - info.b_strides = std::move(b_desc->strides()); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::binary - -#endif // __INFINIOP_BINARY_H__ diff --git a/src/infiniop/binary/cpu/binary_cpu.h b/src/infiniop/binary/cpu/binary_cpu.h deleted file mode 100644 index 208733729..000000000 --- a/src/infiniop/binary/cpu/binary_cpu.h +++ /dev/null @@ -1,57 +0,0 @@ -#ifndef __INFINIOP_BINARY_CPU_H__ -#define __INFINIOP_BINARY_CPU_H__ - -#include "../../devices/cpu/common_cpu.h" -#include "../binary.h" -#include - -namespace op::common_cpu { - -namespace binary_op { - -// Perform binary computation when inputs and the output can have different dtypes -template -void calculate(op::binary::BinaryInfo info, void *c, const void *a, const void *b, Args &&...args) { - auto a_ = reinterpret_cast(a); - auto b_ = reinterpret_cast(b); - auto c_ = reinterpret_cast(c); - ptrdiff_t data_size = info.c_data_size; - -#pragma omp parallel for - for (ptrdiff_t i = 0; i < data_size; ++i) { - size_t a_index = info.contiguous ? i : op::common_cpu::indexToOffset(i, info.ndim, info.a_shape.data(), info.a_strides.data()); - size_t b_index = info.contiguous ? i : op::common_cpu::indexToOffset(i, info.ndim, info.b_shape.data(), info.b_strides.data()); - size_t c_index = info.contiguous ? i : (op::common_cpu::indexToOffset(i, info.ndim, info.c_shape.data(), info.c_strides.data())); - - c_[c_index] = BinaryOp{}(a_[a_index], b_[b_index], std::forward(args)...); - } -} - -// Perform binary computation when all inputs and the output share the same dtype -template -void calculate(op::binary::BinaryInfo info, void *c, const void *a, const void *b, Args &&...args) { - auto a_ = reinterpret_cast(a); - auto b_ = reinterpret_cast(b); - auto c_ = reinterpret_cast(c); - ptrdiff_t data_size = info.c_data_size; - -#pragma omp parallel for - for (ptrdiff_t i = 0; i < data_size; ++i) { - size_t a_index = info.contiguous ? i : op::common_cpu::indexToOffset(i, info.ndim, info.a_shape.data(), info.a_strides.data()); - size_t b_index = info.contiguous ? i : op::common_cpu::indexToOffset(i, info.ndim, info.b_shape.data(), info.b_strides.data()); - size_t c_index = info.contiguous ? i : (op::common_cpu::indexToOffset(i, info.ndim, info.c_shape.data(), info.c_strides.data())); - - if constexpr (std::is_same_v) { - float a_val = utils::cast(a_[a_index]); - float b_val = utils::cast(b_[b_index]); - c_[c_index] = utils::cast(BinaryOp{}(a_val, b_val, std::forward(args)...)); - } else { - c_[c_index] = BinaryOp{}(a_[a_index], b_[b_index], std::forward(args)...); - } - } -} - -} // namespace binary_op -} // namespace op::common_cpu - -#endif // __INFINIOP_BINARY_CPU_H__ diff --git a/src/infiniop/devices/ascend/CMakeLists.txt b/src/infiniop/devices/ascend/CMakeLists.txt deleted file mode 100644 index 5229f770c..000000000 --- a/src/infiniop/devices/ascend/CMakeLists.txt +++ /dev/null @@ -1,36 +0,0 @@ -cmake_minimum_required(VERSION 3.16.0) - -# project information -project(Ascend_C) -set(SOC_VERSION "Ascend910B3" CACHE STRING "system on chip type") -set(ASCEND_CANN_PACKAGE_PATH $ENV{ASCEND_TOOLKIT_HOME} CACHE PATH "ASCEND CANN package installation directory") -set(RUN_MODE "npu" CACHE STRING "run mode: npu") -set(CMAKE_BUILD_TYPE "Release" CACHE STRING "Build type Release/Debug (default Debug)" FORCE) -set(CMAKE_INSTALL_PREFIX "${CMAKE_CURRENT_LIST_DIR}/out" CACHE STRING "path for install()" FORCE) - -if(EXISTS ${ASCEND_CANN_PACKAGE_PATH}/tools/tikcpp/ascendc_kernel_cmake) - set(ASCENDC_CMAKE_DIR ${ASCEND_CANN_PACKAGE_PATH}/tools/tikcpp/ascendc_kernel_cmake) -elseif(EXISTS ${ASCEND_CANN_PACKAGE_PATH}/compiler/tikcpp/ascendc_kernel_cmake) - set(ASCENDC_CMAKE_DIR ${ASCEND_CANN_PACKAGE_PATH}/compiler/tikcpp/ascendc_kernel_cmake) -elseif(EXISTS ${ASCEND_CANN_PACKAGE_PATH}/ascendc_devkit/tikcpp/samples/cmake) - set(ASCENDC_CMAKE_DIR ${ASCEND_CANN_PACKAGE_PATH}/ascendc_devkit/tikcpp/samples/cmake) -else() - message(FATAL_ERROR "ascendc_kernel_cmake does not exist, please check whether the cann package is installed.") -endif() - -include(${ASCENDC_CMAKE_DIR}/ascendc.cmake) -include_directories( - ${CMAKE_SOURCE_DIR}/../../../../include/infiniop/ -) - -ascendc_library(ascend_kernels STATIC - ../../ops/swiglu/ascend/swiglu_ascend_kernel.cpp - ../../ops/rope/ascend/rope_ascend_kernel.cpp - ../../ops/random_sample/ascend/random_sample_kernel.cpp - ../../ops/rms_norm/ascend/cast_kernel.cpp - ../../ops/paged_caching/ascend/paged_caching_ascend_kernel.cpp - ../../ops/paged_attention/ascend/paged_attention_ascend_kernel.cpp - ../../ops/paged_attention_prefill/ascend/paged_attention_prefill_ascend_kernel.cpp -) - -target_include_directories(ascend_kernels PRIVATE ../../../../include) diff --git a/src/infiniop/devices/ascend/Makefile b/src/infiniop/devices/ascend/Makefile deleted file mode 100644 index 7af26076d..000000000 --- a/src/infiniop/devices/ascend/Makefile +++ /dev/null @@ -1,10 +0,0 @@ -.PHONY: build clean - -MKFILE_PATH := $(abspath $(lastword $(MAKEFILE_LIST))) -MKFILE_DIR := $(dir $(MKFILE_PATH)) - -build: - mkdir -p build && cd build && cmake .. && make -j8 - -clean: - rm -rf build diff --git a/src/infiniop/devices/ascend/ascend_handle.cc b/src/infiniop/devices/ascend/ascend_handle.cc deleted file mode 100644 index df1cf5513..000000000 --- a/src/infiniop/devices/ascend/ascend_handle.cc +++ /dev/null @@ -1,14 +0,0 @@ -#include "ascend_handle.h" - -namespace device::ascend { - -Handle::Handle(int device_id) - : InfiniopHandle{INFINI_DEVICE_ASCEND, device_id} {} - -infiniStatus_t Handle::create(InfiniopHandle **Handle_ptr, int device_id) { - *Handle_ptr = new Handle(device_id); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace device::ascend diff --git a/src/infiniop/devices/ascend/ascend_handle.h b/src/infiniop/devices/ascend/ascend_handle.h deleted file mode 100644 index 4cf41d599..000000000 --- a/src/infiniop/devices/ascend/ascend_handle.h +++ /dev/null @@ -1,17 +0,0 @@ -#ifndef __INFINIOP_ASCEND_HANDLE_H__ -#define __INFINIOP_ASCEND_HANDLE_H__ - -#include "../../handle.h" - -namespace device::ascend { - -class Handle : public InfiniopHandle { - - Handle(int device_id); - -public: - static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id); -}; -} // namespace device::ascend - -#endif diff --git a/src/infiniop/devices/ascend/ascend_kernel_common.h b/src/infiniop/devices/ascend/ascend_kernel_common.h deleted file mode 100644 index b67dd9eb3..000000000 --- a/src/infiniop/devices/ascend/ascend_kernel_common.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef __INFINIOP_ASCEND_KERNEL_COMMON_H__ -#define __INFINIOP_ASCEND_KERNEL_COMMON_H__ - -#include "../../../../include/infinicore.h" -#include "kernel_operator.h" - -constexpr size_t BLOCK_NUM = 8; -constexpr size_t BUFFER_NUM = 2; -constexpr size_t BYTE_ALIGN = 32; -constexpr size_t BLOCK_LEN = 256; - -template -__aicore__ inline size_t alignTileLen(size_t tile_len, size_t byte_align) { - size_t bytes = tile_len * sizeof(T); - size_t aligned_bytes = (bytes % byte_align == 0) - ? bytes - : (bytes + (byte_align - bytes % byte_align)); - return aligned_bytes / sizeof(T); -} - -#endif diff --git a/src/infiniop/devices/ascend/common_ascend.cc b/src/infiniop/devices/ascend/common_ascend.cc deleted file mode 100644 index d4f35d728..000000000 --- a/src/infiniop/devices/ascend/common_ascend.cc +++ /dev/null @@ -1,259 +0,0 @@ -#include "common_ascend.h" - -std::vector inferStorageShape(std::vector shape, std::vector strides) { - if (shape.size() != strides.size()) { - throw std::invalid_argument("Shape and strides must have the same length."); - } - - int64_t max_offset = 0; - for (size_t i = 0; i < shape.size(); ++i) { - max_offset += (shape[i] - 1) * strides[i]; - } - - // storage shape is 1D buffer that must cover all accessed elements - return {max_offset + 1}; -} - -size_t aclnnTensorDescriptor::numel() const { - return std::accumulate(shape.begin(), shape.end(), (size_t)1, std::multiplies()); -} - -aclnnTensorDescriptor::aclnnTensorDescriptor(infiniopTensorDescriptor_t desc, void *data) { - this->ndim = desc->ndim(); - this->shape = std::vector(ndim); - this->strides = std::vector(ndim); - for (uint64_t i = 0; i < ndim; ++i) { - this->shape[i] = static_cast(desc->dim(i)); - this->strides[i] = static_cast(desc->stride(i)); - } - this->storageShape = inferStorageShape(this->shape, this->strides); - this->dataType = toAclDataType(desc->dtype()); - // TODO: support other formats - this->format = aclFormat::ACL_FORMAT_ND; - this->tensor = aclCreateTensor(this->shape.data(), - this->ndim, - this->dataType, - this->strides.data(), - this->offset, - this->format, - this->storageShape.data(), - this->storageNdim, - data); -} - -aclnnTensorDescriptor::aclnnTensorDescriptor(aclDataType dtype, const std::vector &shape, const std::vector &strides, void *data) { - this->ndim = shape.size(); - this->shape = shape; - this->strides = strides; - this->dataType = dtype; - this->format = aclFormat::ACL_FORMAT_ND; - if (this->ndim != 0) { - this->storageShape = inferStorageShape(this->shape, this->strides); - } else { - this->storageShape = shape; - this->storageNdim = 0; - } - this->tensor = aclCreateTensor(this->shape.data(), - this->ndim, - this->dataType, - this->strides.data(), - this->offset, - this->format, - this->storageShape.data(), - this->storageNdim, - data); -} - -aclnnTensorDescriptor::~aclnnTensorDescriptor() { - if (this->tensor) { - aclDestroyTensor(this->tensor); - this->tensor = nullptr; - } -} - -aclDataType toAclDataType(infiniDtype_t dt) { - if (dt == INFINI_DTYPE_I8) { - return aclDataType::ACL_INT8; - } else if (dt == INFINI_DTYPE_I16) { - return aclDataType::ACL_INT16; - } else if (dt == INFINI_DTYPE_I32) { - return aclDataType::ACL_INT32; - } else if (dt == INFINI_DTYPE_I64) { - return aclDataType::ACL_INT64; - } else if (dt == INFINI_DTYPE_U8) { - return aclDataType::ACL_UINT8; - } else if (dt == INFINI_DTYPE_U16) { - return aclDataType::ACL_UINT16; - } else if (dt == INFINI_DTYPE_U32) { - return aclDataType::ACL_UINT32; - } else if (dt == INFINI_DTYPE_U64) { - return aclDataType::ACL_UINT64; - } else if (dt == INFINI_DTYPE_F16) { - return aclDataType::ACL_FLOAT16; - } else if (dt == INFINI_DTYPE_BF16) { - return aclDataType::ACL_BF16; - } else if (dt == INFINI_DTYPE_F32) { - return aclDataType::ACL_FLOAT; - } else if (dt == INFINI_DTYPE_F64) { - return aclDataType::ACL_DOUBLE; - } else { - return aclDataType::ACL_DT_UNDEFINED; - } -} - -const char *dataTypeToString(aclDataType dtype) { - switch (dtype) { - case ACL_DT_UNDEFINED: - return "ACL_DT_UNDEFINED"; - case ACL_FLOAT: - return "ACL_FLOAT"; - case ACL_FLOAT16: - return "ACL_FLOAT16"; - case ACL_INT8: - return "ACL_INT8"; - case ACL_INT32: - return "ACL_INT32"; - case ACL_UINT8: - return "ACL_UINT8"; - case ACL_INT16: - return "ACL_INT16"; - case ACL_UINT16: - return "ACL_UINT16"; - case ACL_UINT32: - return "ACL_UINT32"; - case ACL_INT64: - return "ACL_INT64"; - case ACL_UINT64: - return "ACL_UINT64"; - case ACL_DOUBLE: - return "ACL_DOUBLE"; - case ACL_BOOL: - return "ACL_BOOL"; - case ACL_STRING: - return "ACL_STRING"; - case ACL_COMPLEX64: - return "ACL_COMPLEX64"; - case ACL_COMPLEX128: - return "ACL_COMPLEX128"; - case ACL_BF16: - return "ACL_BF16"; - case ACL_INT4: - return "ACL_INT4"; - case ACL_UINT1: - return "ACL_UINT1"; - case ACL_COMPLEX32: - return "ACL_COMPLEX32"; - default: - return "UNKNOWN"; - } -} - -const char *formatToString(aclFormat format) { - switch (format) { - case ACL_FORMAT_UNDEFINED: - return "ACL_FORMAT_UNDEFINED"; - case ACL_FORMAT_NCHW: - return "ACL_FORMAT_NCHW"; - case ACL_FORMAT_NHWC: - return "ACL_FORMAT_NHWC"; - case ACL_FORMAT_ND: - return "ACL_FORMAT_ND"; - case ACL_FORMAT_NC1HWC0: - return "ACL_FORMAT_NC1HWC0"; - case ACL_FORMAT_FRACTAL_Z: - return "ACL_FORMAT_FRACTAL_Z"; - case ACL_FORMAT_NC1HWC0_C04: - return "ACL_FORMAT_NC1HWC0_C04"; - case ACL_FORMAT_HWCN: - return "ACL_FORMAT_HWCN"; - case ACL_FORMAT_NDHWC: - return "ACL_FORMAT_NDHWC"; - case ACL_FORMAT_FRACTAL_NZ: - return "ACL_FORMAT_FRACTAL_NZ"; - case ACL_FORMAT_NCDHW: - return "ACL_FORMAT_NCDHW"; - case ACL_FORMAT_NDC1HWC0: - return "ACL_FORMAT_NDC1HWC0"; - case ACL_FRACTAL_Z_3D: - return "ACL_FRACTAL_Z_3D"; - case ACL_FORMAT_NC: - return "ACL_FORMAT_NC"; - case ACL_FORMAT_NCL: - return "ACL_FORMAT_NCL"; - default: - return "UNKNOWN"; - } -} - -std::string aclnnTensorDescriptor::toString() { - std::ostringstream oss; - - // 写入 ndim - oss << "ndim: " << this->ndim << "\n"; - - // 写入 shape - oss << "shape: ["; - for (uint64_t i = 0; i < this->ndim; ++i) { - oss << this->shape[i]; - if (i < this->ndim - 1) { - oss << ", "; - } - } - oss << "]\n"; - - // 写入 stride - oss << "stride: ["; - for (uint64_t i = 0; i < this->ndim; ++i) { - oss << this->strides[i]; - if (i < this->ndim - 1) { - oss << ", "; - } - } - oss << "]\n"; - - // 写入 offset - oss << "offset: " << this->offset << "\n"; - - // 写入 dataType - oss << "dataType: " << dataTypeToString(this->dataType) << "\n"; - - // 写入 format - oss << "format: " << formatToString(this->format) << "\n"; - - // 写入 storageShape - oss << "storageShape: ["; - for (int64_t i = 0; i < this->storageNdim; ++i) { - oss << this->storageShape[i]; - if (i < this->storageNdim - 1) { - oss << ", "; - } - } - oss << "]\n"; - - // 写入 storageNdim - oss << "storageNdim: " << this->storageNdim << "\n"; - - // 返回构建的字符串 - return oss.str(); -} - -aclnnScalarDescriptor::aclnnScalarDescriptor(aclDataType dtype, const void *value, size_t size) { - this->dataType = dtype; - this->size = size; - this->value = const_cast(value); - this->scalar = aclCreateScalar(this->value, this->dataType); -} - -aclnnScalarDescriptor::aclnnScalarDescriptor(infiniDtype_t dtype, const void *value, size_t size) { - this->dataType = toAclDataType(dtype); - this->size = size; - this->value = const_cast(value); - this->scalar = aclCreateScalar(this->value, this->dataType); -} - -aclnnScalarDescriptor::~aclnnScalarDescriptor() { - if (this->scalar) { - aclDestroyScalar(this->scalar); - this->scalar = nullptr; - } -} diff --git a/src/infiniop/devices/ascend/common_ascend.h b/src/infiniop/devices/ascend/common_ascend.h deleted file mode 100644 index 4c2e42c9b..000000000 --- a/src/infiniop/devices/ascend/common_ascend.h +++ /dev/null @@ -1,68 +0,0 @@ -#ifndef __INFINIOP_COMMON_ASCEND_H__ -#define __INFINIOP_COMMON_ASCEND_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include "ascend_handle.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif -#define CHECK_ACL(API) CHECK_INTERNAL(API, ACL_SUCCESS) -#ifdef __cplusplus -}; -#endif - -struct aclnnTensorDescriptor { - uint64_t ndim; - std::vector shape; - std::vector strides; - int64_t offset = 0; - aclDataType dataType; - aclFormat format; - std::vector storageShape; - int64_t storageNdim = 1; - aclTensor *tensor; - - aclnnTensorDescriptor(aclDataType dtype, const std::vector &shape, const std::vector &strides, void *data = nullptr); - aclnnTensorDescriptor(infiniopTensorDescriptor_t y_desc, void *data = nullptr); - ~aclnnTensorDescriptor(); - size_t numel() const; - - std::string toString(); -}; -typedef aclnnTensorDescriptor *aclnnTensorDescriptor_t; - -struct aclnnScalarDescriptor { - aclDataType dataType; - void *value; - size_t size; - aclScalar *scalar; - - aclnnScalarDescriptor(aclDataType dtype, const void *value, size_t size); - aclnnScalarDescriptor(infiniDtype_t dtype, const void *value, size_t size); - ~aclnnScalarDescriptor(); -}; -typedef aclnnScalarDescriptor *aclnnScalarDescriptor_t; - -aclDataType toAclDataType(infiniDtype_t dt); - -#define GetRecentErrMsg() \ - { \ - auto tmp_err_msg = aclGetRecentErrMsg(); \ - if (tmp_err_msg != NULL) { \ - printf(" ERROR Message : %s \n ", tmp_err_msg); \ - } \ - } - -#endif diff --git a/src/infiniop/devices/bang/bang_handle.cc b/src/infiniop/devices/bang/bang_handle.cc deleted file mode 100644 index 59c16bac4..000000000 --- a/src/infiniop/devices/bang/bang_handle.cc +++ /dev/null @@ -1,96 +0,0 @@ -#include "../../tensor.h" -#include "common_bang.h" -#include "infiniop/tensor_descriptor.h" -#include -#include - -namespace device::bang { - -Handle::Handle(infiniDevice_t device, int device_id) - : InfiniopHandle{device, device_id}, - _internal(std::make_shared(device_id)) {} - -auto Handle::internal() const -> const std::shared_ptr & { - return _internal; -} - -Handle::Internal::Internal(int device_id) { - cnrtDeviceGetAttribute(&_cluster_count, cnrtAttrClusterCount, device_id); - cnrtDeviceGetAttribute(&_core_per_cluster, cnrtAttrMcorePerCluster, device_id); -} - -infiniStatus_t Handle::Internal::useCnnl(cnrtQueue_t queue, const Fn &f) const { - auto handle = cnnl_handles.pop(); - if (!handle) { - CHECK_BANG(cnnlCreate(&(*handle))); - } - CHECK_BANG(cnnlSetQueue(*handle, queue)); - CHECK_STATUS(f(*handle)); - cnnl_handles.push(std::move(*handle)); - return INFINI_STATUS_SUCCESS; -} - -int Handle::Internal::getCorePerCluster() const { return _core_per_cluster; } -int Handle::Internal::getClusterCount() const { return _cluster_count; } - -cnnlDataType_t getCnnlDtype(infiniDtype_t dt) { - switch (dt) { - case INFINI_DTYPE_F32: - return CNNL_DTYPE_FLOAT; - case INFINI_DTYPE_F64: - return CNNL_DTYPE_DOUBLE; - case INFINI_DTYPE_F16: - return CNNL_DTYPE_HALF; - case INFINI_DTYPE_I8: - return CNNL_DTYPE_INT8; - case INFINI_DTYPE_I32: - return CNNL_DTYPE_INT32; - case INFINI_DTYPE_U8: - return CNNL_DTYPE_UINT8; - case INFINI_DTYPE_BF16: - return CNNL_DTYPE_BFLOAT16; - case INFINI_DTYPE_I64: - return CNNL_DTYPE_INT64; - default: - return CNNL_DTYPE_INVALID; - } -} - -infiniStatus_t setCnnlTensor(cnnlTensorDescriptor_t desc, - const InfiniopTensorDescriptor *layout) { - std::vector dims(layout->ndim()); - for (size_t i = 0; i < layout->ndim(); i++) { - dims[i] = static_cast(layout->shape()[i]); - } - CHECK_BANG(cnnlSetTensorDescriptor(desc, CNNL_LAYOUT_ARRAY, - getCnnlDtype(layout->dtype()), dims.size(), - dims.data())); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t setCnnlTensorEx(cnnlTensorDescriptor_t desc, - const InfiniopTensorDescriptor *layout) { - std::vector dim_size(layout->ndim()), dim_stride(layout->ndim()); - for (size_t i = 0; i < layout->ndim(); i++) { - dim_size[i] = static_cast(layout->shape()[i]); - dim_stride[i] = static_cast(layout->strides()[i]); - } - CHECK_BANG(cnnlSetTensorDescriptorEx( - desc, CNNL_LAYOUT_ARRAY, getCnnlDtype(layout->dtype()), - dim_size.size(), dim_size.data(), dim_stride.data())); - return INFINI_STATUS_SUCCESS; -} - -namespace cambricon { - -Handle::Handle(int device_id) - : bang::Handle(INFINI_DEVICE_CAMBRICON, device_id) {} - -infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int device_id) { - *handle_ptr = new Handle(device_id); - return INFINI_STATUS_SUCCESS; -} - -} // namespace cambricon - -} // namespace device::bang diff --git a/src/infiniop/devices/bang/bang_handle.h b/src/infiniop/devices/bang/bang_handle.h deleted file mode 100644 index 68a542703..000000000 --- a/src/infiniop/devices/bang/bang_handle.h +++ /dev/null @@ -1,33 +0,0 @@ -#ifndef __INFINIOP_BANG_HANDLE_H__ -#define __INFINIOP_BANG_HANDLE_H__ - -#include "../../handle.h" -#include - -namespace device::bang { - -struct Handle : public InfiniopHandle { - class Internal; - auto internal() const -> const std::shared_ptr &; - -protected: - Handle(infiniDevice_t device, int device_id); - -private: - std::shared_ptr _internal; -}; - -namespace cambricon { - -class Handle : public bang::Handle { - Handle(int device_id); - -public: - static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id); -}; - -} // namespace cambricon - -} // namespace device::bang - -#endif // __INFINIOP_BANG_HANDLE_H__ diff --git a/src/infiniop/devices/bang/bang_kernel_common.h b/src/infiniop/devices/bang/bang_kernel_common.h deleted file mode 100644 index fc6f8b51b..000000000 --- a/src/infiniop/devices/bang/bang_kernel_common.h +++ /dev/null @@ -1,211 +0,0 @@ -#ifndef __INFINIOP_BANG_KERNEL_COMMON_H__ -#define __INFINIOP_BANG_KERNEL_COMMON_H__ - -// Include Cambricon CNNL and CNRT headers for MLU (Machine Learning Unit) specific functions -#include "cnnl.h" -#include "cnrt.h" - -namespace device::bang::kernel { - -template -__mlu_device__ float to_float(const T &v) { - return static_cast(v); -} - -template -__mlu_device__ bfloat16_t to_bfloat16(const T &v) { - return static_cast(v); -} - -template -__mlu_device__ half to_half(const T &v) { - return static_cast(v); -} - -/** - * @brief Converts a flattened index to a memory offset considering tensor striding. - * - * This is the general case for non-contiguous tensors where elements are not - * stored sequentially in memory. - * - * @param flat_index The flattened index in the tensor - * @param ndim Number of dimensions - * @param shape Tensor shape - * @param strides Tensor strides (in elements) - * @return size_t Offset in the tensor's memory - */ -inline __mlu_device__ size_t indexToOffset( - size_t flat_index, - size_t ndim, - const size_t *shape, - const ptrdiff_t *strides) { - - size_t res = 0; - // Process dimensions from highest to lowest - for (size_t i = ndim; i-- > 0;) { - // Add contribution from this dimension - res += (flat_index % shape[i]) * strides[i]; - // Remove the contribution from this dimension - flat_index /= shape[i]; - } - return res; -} - -/** - * @brief Helper struct for computing input tensor indices considering broadcasting and striding. - * - * This is particularly useful for operations where inputs may be broadcasted - * to match the output shape, or may have non-contiguous memory layouts. - */ -struct InputIndexer { - size_t idx; // Base index for this task - size_t ndim; // Number of dimensions - const bool *input_contiguous; // Array indicating which inputs are contiguous - const bool *input_broadcasted; // Array indicating which inputs are broadcasted - const size_t *input_shapes; // Array of input shapes (concatenated) - const ptrdiff_t *input_strides; // Array of input strides (concatenated) - const ptrdiff_t *output_strides; // Output tensor strides - - /** - * @brief Computes memory offset for input tensor element. - * - * @param input_id Input tensor ID. - * @param element_idx Element index in output tensor. - * @return size_t Memory offset in input tensor. - */ - __mlu_device__ size_t operator()(size_t input_id, size_t element_idx) const { - size_t global_idx = idx + element_idx; - return input_contiguous[input_id] - ? global_idx // Simple case: contiguous memory - : indexToOffset(global_idx, ndim, input_shapes + input_id * ndim, input_strides + input_id * ndim); - } -}; - -/** - * @brief Computes output tensor index considering striding. - * - * @param idx Linear index. - * @param is_contiguous Whether output is contiguous. - * @param ndim Number of dimensions. - * @param shape Output tensor shape. - * @param strides Output tensor strides. - * @return size_t Memory offset in output tensor. - */ -inline __mlu_device__ size_t -getOutputIndex(size_t idx, - bool is_contiguous, - size_t ndim, - const size_t *shape, - const ptrdiff_t *strides) { - return is_contiguous ? idx : indexToOffset(idx, ndim, shape, strides); -} /** - * @brief Calculates optimal chunk size for memory operations based on tensor contiguity. - * - * This function doesn't handle tensors with non-standard strides, which - * require more general optimizations not specific to Cambricon. - * - * @param global_idx_ Starting global index. - * @param ndim Number of dimensions. - * @param shape Tensor shape. - * @param strides Tensor strides. - * @param max_len Maximum allowed chunk size. - * @return size_t Optimal chunk size for memory operations. - */ -__mlu_device__ size_t calculateChunkSize( - size_t global_idx_, - size_t ndim, - const size_t *shape, - const ptrdiff_t *strides, - size_t max_len) { - // Find the last dimension that is contiguous - int last_contiguous_dim = -1; - ptrdiff_t expected_stride = 1; - - for (int i = (int)ndim - 1; i >= 0; --i) { - if (strides[i] != expected_stride) { - break; - } - last_contiguous_dim = i; - if (i > 0) { - expected_stride *= shape[i]; - } - } - - if (last_contiguous_dim < 0) { - return 1; - } - - // Calculate position in the contiguous block - size_t global_idx = global_idx_; - size_t pos_in_block = 0; - size_t block_size = 1; - - for (int i = (int)ndim - 1; i >= last_contiguous_dim; --i) { - size_t dim_idx = global_idx % shape[i]; - pos_in_block += dim_idx * block_size; - block_size *= shape[i]; - global_idx /= shape[i]; - } - - size_t remaining_in_block = block_size - pos_in_block; - return std::min(max_len, remaining_in_block); -} - -/** - * @brief Helper function for non-contiguous memory copy - * - * @param dst Destination buffer - * @param src Source buffer - * @param direction Memory copy direction (GDRAM2NRAM or NRAM2GDRAM) - * @param indexer Input indexer helper (for input copies) - * @param input_idx Input tensor index (for input copies) - * @param processed Number of elements already processed - * @param curr_batch Current batch size - * @param start_idx Starting index for this task - * @param ndim Number of dimensions - * @param shape Tensor shape - * @param strides Tensor strides - * @param is_input_copy Whether this is an input copy operation - */ -template -__mlu_device__ void nonContiguousMemcpy( - Tdata *dst, - Tdata *src, - mluMemcpyDirection_t direction, - InputIndexer &indexer, - size_t input_idx, - size_t processed, - size_t curr_batch, - size_t start_idx, - size_t ndim, - const size_t *shape, - const ptrdiff_t *strides, - bool is_input_copy) { - - size_t remaining = curr_batch; - size_t current_pos = 0; - - while (remaining > 0) { - size_t element_offset = is_input_copy ? indexer(input_idx, processed + current_pos) : getOutputIndex(start_idx + processed + current_pos, - false, // output_contiguous is false for non-contiguous - ndim, shape, strides); - - size_t chunk_size = calculateChunkSize(start_idx + processed + current_pos, - ndim, - shape, - strides, - remaining); - - __memcpy_async(dst + (is_input_copy ? current_pos : element_offset), - src + (is_input_copy ? element_offset : current_pos), - chunk_size * sizeof(Tdata), - direction); - - current_pos += chunk_size; - remaining -= chunk_size; - } -} - -} // namespace device::bang::kernel - -#endif diff --git a/src/infiniop/devices/bang/common_bang.h b/src/infiniop/devices/bang/common_bang.h deleted file mode 100644 index e1fe98c2b..000000000 --- a/src/infiniop/devices/bang/common_bang.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef __COMMON_BANG_H__ -#define __COMMON_BANG_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include "../pool.h" -#include "bang_handle.h" -#include "cnnl.h" -#include "cnrt.h" -#include - -#define CHECK_BANG(API) CHECK_INTERNAL(API, CNNL_STATUS_SUCCESS) - -#define NRAM_MAX_SIZE 1024 * 240 -constexpr size_t ALIGN_SIZE = 128; - -namespace device::bang { - -class Handle::Internal { - Pool cnnl_handles; - - int _core_per_cluster; - int _cluster_count; - - template - using Fn = std::function; - -public: - Internal(int); - - infiniStatus_t useCnnl(cnrtQueue_t queue, const Fn &f) const; - - int getCorePerCluster() const; - int getClusterCount() const; -}; - -cnnlDataType_t getCnnlDtype(infiniDtype_t dt); - -// set cnnl tensor descriptor without strides -infiniStatus_t setCnnlTensor(cnnlTensorDescriptor_t desc, - const InfiniopTensorDescriptor *layout); - -// set cnnl tensor descriptor with strides -infiniStatus_t setCnnlTensorEx(cnnlTensorDescriptor_t desc, - const InfiniopTensorDescriptor *layout); - -} // namespace device::bang - -#endif // __COMMON_BANG_H__ diff --git a/src/infiniop/devices/cpu/common_cpu.cc b/src/infiniop/devices/cpu/common_cpu.cc deleted file mode 100644 index 6032fa03f..000000000 --- a/src/infiniop/devices/cpu/common_cpu.cc +++ /dev/null @@ -1,41 +0,0 @@ -#include "common_cpu.h" - -namespace op::common_cpu { - -size_t indexToOffset( - size_t flat_index, - size_t ndim, - const size_t *shape, - const ptrdiff_t *strides) { - size_t res = 0; - for (size_t i = ndim; i-- > 0;) { - res += (flat_index % shape[i]) * strides[i]; - flat_index /= shape[i]; - } - return res; -} - -size_t getPaddedSize( - size_t ndim, - size_t *shape, - const size_t *pads) { - size_t total_size = 1; - for (size_t i = 0; i < ndim; ++i) { - total_size *= shape[i] + (i < 2 ? 0 : 2 * pads[i - 2]); - } - return total_size; -} - -std::vector getPaddedShape( - size_t ndim, - const size_t *shape, - const size_t *pads) { - std::vector padded_shape(ndim); - memcpy(padded_shape.data(), shape, ndim * sizeof(size_t)); - for (size_t i = 2; i < ndim; ++i) { - padded_shape[i] += 2 * pads[i - 2]; - } - return padded_shape; -} - -} // namespace op::common_cpu diff --git a/src/infiniop/devices/cpu/common_cpu.h b/src/infiniop/devices/cpu/common_cpu.h deleted file mode 100644 index 1ae16ed83..000000000 --- a/src/infiniop/devices/cpu/common_cpu.h +++ /dev/null @@ -1,32 +0,0 @@ -#ifndef __INFINIOP_COMMON_CPU_H__ -#define __INFINIOP_COMMON_CPU_H__ - -#include "../../../utils.h" -#include "cpu_handle.h" -#include -#include -#include -#include -#include - -#ifdef ENABLE_OMP -#include -#endif - -namespace op::common_cpu { - -// return the memory offset a tensor given flattened index -size_t indexToOffset(size_t flat_index, size_t ndim, const size_t *shape, const ptrdiff_t *strides); - -/** - * get the total array size (element count) after applying padding for a - * ndim-ary tensor with the given shape - */ -size_t getPaddedSize(size_t ndim, size_t *shape, const size_t *pads); - -// calculate the padded shape and store the result in padded_shape -std::vector getPaddedShape(size_t ndim, const size_t *shape, const size_t *pads); - -} // namespace op::common_cpu - -#endif // __INFINIOP__COMMON_CPU_H__ diff --git a/src/infiniop/devices/cpu/cpu_handle.cc b/src/infiniop/devices/cpu/cpu_handle.cc deleted file mode 100644 index 2fb3b5854..000000000 --- a/src/infiniop/devices/cpu/cpu_handle.cc +++ /dev/null @@ -1,12 +0,0 @@ -#include "cpu_handle.h" - -namespace device::cpu { - -Handle::Handle() : InfiniopHandle{INFINI_DEVICE_CPU, 0} {} - -infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int) { - *handle_ptr = new Handle{}; - return INFINI_STATUS_SUCCESS; -} - -} // namespace device::cpu diff --git a/src/infiniop/devices/cpu/cpu_handle.h b/src/infiniop/devices/cpu/cpu_handle.h deleted file mode 100644 index 0600d5179..000000000 --- a/src/infiniop/devices/cpu/cpu_handle.h +++ /dev/null @@ -1,17 +0,0 @@ -#ifndef __INFINIOP_CPU_HANDLE_H__ -#define __INFINIOP_CPU_HANDLE_H__ - -#include "../../handle.h" - -namespace device::cpu { - -class Handle : public InfiniopHandle { - Handle(); - -public: - static infiniStatus_t create(InfiniopHandle **handle_ptr, int); -}; - -} // namespace device::cpu - -#endif diff --git a/src/infiniop/devices/handle.cc b/src/infiniop/devices/handle.cc deleted file mode 100644 index 9057f7203..000000000 --- a/src/infiniop/devices/handle.cc +++ /dev/null @@ -1,143 +0,0 @@ -#include "infiniop/handle.h" -#include "../../bridge/infini/rt.hpp" -#include "../../utils.h" - -#ifdef ENABLE_CPU_API -#include "cpu/cpu_handle.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) -#include "nvidia/nvidia_handle.h" -#endif -#ifdef ENABLE_CAMBRICON_API -#include "bang/bang_handle.h" -#endif -#ifdef ENABLE_ASCEND_API -#include "ascend/ascend_handle.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/moore_handle.h" -#endif -#ifdef ENABLE_KUNLUN_API -#include "kunlun/kunlun_handle.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/metax_handle.h" -#endif - -__INFINI_C infiniStatus_t infiniopSetRuntimeDevice(infiniDevice_t device, int device_id) { - auto rt_device = infinicore::bridge::infini::rt::translate_to(device); - if (rt_device == infini::rt::Device::Type::kCount) { - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - - infini::rt::set_runtime_device_type(rt_device); - CHECK_STATUS(infinicore::bridge::infini::rt::translate(infini::rt::runtime::SetDevice(device_id))); - return INFINI_STATUS_SUCCESS; -} - -__INFINI_C infiniStatus_t infiniopCreateHandle(infiniopHandle_t *handle_ptr) { - if (handle_ptr == nullptr) { - return INFINI_STATUS_NULL_POINTER; - } - - infiniDevice_t device = infinicore::bridge::infini::rt::translate_from(infini::rt::runtime_device_type()); - if (device == INFINI_DEVICE_TYPE_COUNT) { - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - - int device_id = 0; - CHECK_STATUS(infinicore::bridge::infini::rt::translate(infini::rt::runtime::GetDevice(&device_id))); - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return device::NAMESPACE::Handle::create(handle_ptr, device_id) - - switch (device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, iluvatar); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, ali); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, qy); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang::cambricon); -#endif -#ifdef ENABLE_ASCEND_API - CREATE(INFINI_DEVICE_ASCEND, ascend); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_KUNLUN_API - CREATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, hygon); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopDestroyHandle(infiniopHandle_t handle) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(handle); \ - return INFINI_STATUS_SUCCESS - - switch (handle->device) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, iluvatar); -#endif -#ifdef ENABLE_ALI_API - DELETE(INFINI_DEVICE_ALI, ali); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, qy); -#endif -#ifdef ENABLE_CAMBRICON_API - DELETE(INFINI_DEVICE_CAMBRICON, bang::cambricon); -#endif -#ifdef ENABLE_ASCEND_API - DELETE(INFINI_DEVICE_ASCEND, ascend); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_KUNLUN_API - DELETE(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, hygon); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/devices/kunlun/kunlun_common.h b/src/infiniop/devices/kunlun/kunlun_common.h deleted file mode 100644 index 3e1bd8d1d..000000000 --- a/src/infiniop/devices/kunlun/kunlun_common.h +++ /dev/null @@ -1,17 +0,0 @@ -#ifndef __KUNLUN_COMMON_H__ -#define __KUNLUN_COMMON_H__ - -#include "../../../utils.h" -#include -#include -#include - -namespace xdnn = baidu::xpu::api; - -typedef XPUStream kunlunStream_t; -typedef XPUEvent kunlunEvent_t; -typedef xdnn::Context *xdnnHandle_t; - -#define CHECK_KUNLUN(API) CHECK_INTERNAL(API, XPU_SUCCESS) - -#endif diff --git a/src/infiniop/devices/kunlun/kunlun_handle.cc b/src/infiniop/devices/kunlun/kunlun_handle.cc deleted file mode 100644 index 3086b8045..000000000 --- a/src/infiniop/devices/kunlun/kunlun_handle.cc +++ /dev/null @@ -1,29 +0,0 @@ -#include "kunlun_handle.h" - -namespace device::kunlun { - -Handle::Handle(int device_id) - : InfiniopHandle{INFINI_DEVICE_KUNLUN, device_id}, - _internal(std::make_shared()) {} - -auto Handle::internal() const -> const std::shared_ptr & { - return _internal; -} - -infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int device_id) { - *handle_ptr = new Handle(device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Handle::Internal::useXdnn(kunlunStream_t stream, const Fn &f) const { - auto handle = dnn_handles.pop(); - if (!handle) { - *handle = xdnn::create_context(); - } - (*handle)->set_stream(stream); - CHECK_STATUS(f(*handle)); - dnn_handles.push(std::move(*handle)); - return INFINI_STATUS_SUCCESS; -} - -} // namespace device::kunlun diff --git a/src/infiniop/devices/kunlun/kunlun_handle.h b/src/infiniop/devices/kunlun/kunlun_handle.h deleted file mode 100644 index 1e0a7c804..000000000 --- a/src/infiniop/devices/kunlun/kunlun_handle.h +++ /dev/null @@ -1,35 +0,0 @@ -#ifndef __INFINIOP_KUNLUN_HANDLE_H__ -#define __INFINIOP_KUNLUN_HANDLE_H__ - -#include "../../handle.h" -#include "../pool.h" -#include "kunlun_common.h" -#include - -namespace device::kunlun { - -struct Handle : public InfiniopHandle { - class Internal; - auto internal() const -> const std::shared_ptr &; - - Handle(int device_id); - -private: - std::shared_ptr _internal; - -public: - static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id); -}; - -class Handle::Internal { - Pool dnn_handles; - template - using Fn = std::function; - -public: - infiniStatus_t useXdnn(kunlunStream_t stream, const Fn &f) const; -}; - -} // namespace device::kunlun - -#endif // __INFINIOP_KUNLUN_HANDLE_H__ diff --git a/src/infiniop/devices/kunlun/kunlun_kernel_common.h b/src/infiniop/devices/kunlun/kunlun_kernel_common.h deleted file mode 100644 index 7ff3374fd..000000000 --- a/src/infiniop/devices/kunlun/kunlun_kernel_common.h +++ /dev/null @@ -1,218 +0,0 @@ -#ifndef __INFINIOP_KUNLUN_KERNEL_COMMON_H__ -#define __INFINIOP_KUNLUN_KERNEL_COMMON_H__ - -// This header file will only be include by .xpu file -#include "xpu/runtime.h" -#include -#include -#include -#include -#include -#include - -namespace device::kunlun::kernel { - -#define SM_SIZE 40960 - -/** - * @brief Define ptrdiff_t and size_t for kunlun xpu - * ptrdiff_t is 32 bit, size_t is 32 bit in xpu kernel - * We padding it into 64 bit for convience of DATACOPY - */ -typedef struct _ptrdiff_t { - int32_t value; // 32 bit - int32_t padding; // 32 bit -} _ptrdiff_t; -// same as ptrdiff -typedef struct _size_t { - uint32_t value; - uint32_t padding; -} _size_t; - -// Get mask for kunlun xpu 512bit register calculation -// if data is not enough to 512bit, padding zero and use -// mask to identify real data -// 0 - i bit 1, others 0 -inline __device__ float lowerBitMask(int i) { - return (1 << (i + 1)) - 1; -} - -// Load len data from shared memory -template -__device__ inline void loadsm(__shared_ptr__ const T *p, T *v, int len) { - __builtin_memcpy(v, p, len * sizeof(T)); -} - -template -__device__ inline T loadShared(__shared_ptr__ const T *ptr) { - T value; - __builtin_memcpy(&value, ptr, sizeof(T)); - return value; -} - -template -__device__ inline void storeShared(__shared_ptr__ T *ptr, const T &value) { - __builtin_memcpy(ptr, &value, sizeof(T)); -} - -/** - * @brief atomicAdd for kunlun xpu - * @param ptr: pointer to shared memory - * @param value: value to add - */ -template -inline __device__ T atomicAdd(__shared_ptr__ T *ptr, T value) { - T x = atomicadd(ptr, value); - return x; -} -// Specialize atomicAdd for half -template <> -inline __device__ half atomicAdd(__shared_ptr__ half *ptr, half value) { - ticket_lock_mix(); - half old = loadShared(ptr); - - float of = __half2float(old); - float vf = __half2float(value); - float sumf = of + vf; - half sum = __float2half_rn(sumf); - storeShared(ptr, sum); - - mfence_sm(); - ticket_unlock_mix(); - return old; -} -// Specialize atomicAdd for bfloat16_t -template <> -inline __device__ bfloat16_t atomicAdd(__shared_ptr__ bfloat16_t *ptr, bfloat16_t value) { - ticket_lock_mix(); - bfloat16_t old = loadShared(ptr); - - float of = __bfloat162float(old); - float vf = __bfloat162float(value); - float sumf = of + vf; - bfloat16_t sum = __float2bfloat16_rn(sumf); - storeShared(ptr, sum); - - mfence_sm(); - ticket_unlock_mix(); - return old; -} - -/** - * @brief atomicMax for kunlun xpu - * @param ptr: pointer to shared memory - * @param value: value to compare - */ -template -inline __device__ T atomicMax(__shared_ptr__ T *ptr, T value) { - ticket_lock_mix(); - T old = loadShared(ptr); - - if constexpr (std::is_same::value) { - float of = __bfloat162float(old); - float vf = __bfloat162float(value); - float maxf = fmax(of, vf); - bfloat16_t max = __float2bfloat16_rn(maxf); - storeShared(ptr, max); - } else { - const T max = fmax(old, value); - storeShared(ptr, max); - } - mfence_sm(); - ticket_unlock_mix(); - return old; -} - -/** - * @brief Get real offset of input index - * flat_index: flatten index input - * ndim: dim of input tensor - * shape: shape of input tensor - * strides: strides of input tensor - */ -inline __device__ int indexToOffset( - int flat_index, - int ndim, - const _size_t *shape, - const _ptrdiff_t *strides) { - - int res = 0; - for (int i = ndim; i-- > 0;) { - res += (flat_index % shape[i].value) * strides[i].value; - flat_index /= shape[i].value; - } - return res; -} - -/** - * @brief Get max of a array of local mem - * @param data: pointer to local memory - * @param len: length of array - * @return max value - */ -template -__inline__ __device__ T max(const T *data_ptr, size_t len) { - T max_val = data_ptr[0]; - for (size_t i = 0; i < len; ++i) { - max_val = fmax(max_val, data_ptr[i]); - } - return max_val; -} - -// Use simd vector instruction to calculate max of a half array -template <> -__inline__ __device__ half max(const half *data_ptr, size_t len) { - int remain = len % 32; - int offset_last = len - remain; - half res = data_ptr[0]; - for (int i = offset_last; i < len; i++) { - res = fmax(res, *(data_ptr + i)); - } - mfence(); - if (offset_last != 0) { - __local__ half acc_buf[32]; - float16x32_t v_mv = vload_lm_float16x32_mz(data_ptr); - // for every 16 float data - for (int i = 32; i < offset_last; i += 32) { - float16x32_t v_0 = vload_lm_float16x32_mz(data_ptr + i); - v_mv = vvmax_float16x32_mz(v_mv, v_0); - } - vstore_lm_float16x32_mz(acc_buf, v_mv); - mfence(); - for (int i = 0; i < 32; i++) { - res = fmax(res, acc_buf[i]); - } - } - return res; -} - -// Use simd vector instruction to calculate max of a half array -template <> -__inline__ __device__ float max(const float *data_ptr, size_t len) { - int remain = len % 16; - int offset_last = len - remain; - float res = data_ptr[0]; - for (int i = offset_last; i < len; i++) { - res = fmax(res, *(data_ptr + i)); - } - mfence(); - if (offset_last != 0) { - __local__ float acc_buf[16]; - float32x16_t v_mv = vload_lm_float32x16_mz(data_ptr); - // for every 16 float data - for (int i = 16; i < offset_last; i += 16) { - float32x16_t v_0 = vload_lm_float32x16_mz(data_ptr + i); - v_mv = vvmax_float32x16_mz(v_mv, v_0); - } - vstore_lm_float32x16_mz(acc_buf, v_mv); - mfence(); - for (int i = 0; i < 16; i++) { - res = fmax(res, acc_buf[i]); - } - } - return res; -} - -} // namespace device::kunlun::kernel - -#endif // __INFINIOP_KUNLUN_KERNEL_COMMON_H__ diff --git a/src/infiniop/devices/kunlun/kunlun_xblas.cc b/src/infiniop/devices/kunlun/kunlun_xblas.cc deleted file mode 100644 index 8f071e6a6..000000000 --- a/src/infiniop/devices/kunlun/kunlun_xblas.cc +++ /dev/null @@ -1,30 +0,0 @@ -#include "kunlun_xblas.h" - -namespace device::kunlun::blas { - -Handle::Handle(int device_id) - : InfiniopHandle{INFINI_DEVICE_KUNLUN, device_id}, - _internal(std::make_shared()) {} - -auto Handle::internal() const -> const std::shared_ptr & { - return _internal; -} - -infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int device_id) { - *handle_ptr = new Handle(device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Handle::Internal::useCublas(cudaStream_t stream, const Fn &f) const { - - auto handle = blas_handles.pop(); - if (!handle) { - CHECK_CUBLAS(cublasCreate(&(*handle))); - } - CHECK_CUBLAS(cublasSetStream(*handle, stream)); - CHECK_STATUS(f(*handle)); - blas_handles.push(std::move(*handle)); - return INFINI_STATUS_SUCCESS; -} - -} // namespace device::kunlun::blas diff --git a/src/infiniop/devices/kunlun/kunlun_xblas.h b/src/infiniop/devices/kunlun/kunlun_xblas.h deleted file mode 100644 index 3e6d2a71c..000000000 --- a/src/infiniop/devices/kunlun/kunlun_xblas.h +++ /dev/null @@ -1,38 +0,0 @@ -#ifndef __KUNLUN_XBLAS_H__ -#define __KUNLUN_XBLAS_H__ - -#include "../../handle.h" -#include "../pool.h" -#include "kunlun_common.h" -#include -#include - -#define CHECK_CUBLAS(API) CHECK_INTERNAL(API, CUBLAS_STATUS_SUCCESS) - -namespace device::kunlun::blas { - -struct Handle : public InfiniopHandle { - class Internal; - auto internal() const -> const std::shared_ptr &; - - Handle(int device_id); - -private: - std::shared_ptr _internal; - -public: - static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id); -}; - -class Handle::Internal { - Pool blas_handles; - template - using Fn = std::function; - -public: - infiniStatus_t useCublas(cudaStream_t stream, const Fn &f) const; -}; - -} // namespace device::kunlun::blas - -#endif // __KUNLUN_XBLAS_H__ diff --git a/src/infiniop/devices/metax/metax_common.h b/src/infiniop/devices/metax/metax_common.h deleted file mode 100644 index 2aef6c786..000000000 --- a/src/infiniop/devices/metax/metax_common.h +++ /dev/null @@ -1,53 +0,0 @@ -#pragma once -#include "../../../utils.h" -#include "../pool.h" -#include "metax_handle.h" -#include "metax_ht2mc.h" - -#ifdef ENABLE_METAX_MC_API -#include -#include -#include -#else -#include -#include -#include -#endif -#include -#include - -#define CHECK_MCBLAS(API) CHECK_INTERNAL(API, HCBLAS_STATUS_SUCCESS) -#define CHECK_MCDNN(API) CHECK_INTERNAL(API, HCDNN_STATUS_SUCCESS) - -namespace device::metax { - -class Handle::Internal { - Pool mcblas_handles; - Pool mcdnn_handles; - - template - using Fn = std::function; - - int _warp_size, - _max_threads_per_block, - _block_size[3], - _grid_size[3]; - -public: - Internal(int); - infiniStatus_t useMcblas(hcStream_t stream, const Fn &f) const; - infiniStatus_t useMcdnn(hcStream_t stream, const Fn &f) const; - - int warpSize() const; - int maxThreadsPerBlock() const; - int blockSizeX() const; - int blockSizeY() const; - int blockSizeZ() const; - int gridSizeX() const; - int gridSizeY() const; - int gridSizeZ() const; -}; - -hcdnnDataType_t getHcdnnDtype(infiniDtype_t dt); - -} // namespace device::metax diff --git a/src/infiniop/devices/metax/metax_handle.cc b/src/infiniop/devices/metax/metax_handle.cc deleted file mode 100644 index d56561c02..000000000 --- a/src/infiniop/devices/metax/metax_handle.cc +++ /dev/null @@ -1,86 +0,0 @@ -#include "metax_common.h" - -namespace device::metax { -Handle::Handle(infiniDevice_t device, int device_id) - : InfiniopHandle{device, device_id}, - _internal(std::make_shared(device_id)) {} - -Handle::Handle(int device_id) : Handle(INFINI_DEVICE_METAX, device_id) {} - -auto Handle::internal() const -> const std::shared_ptr & { - return _internal; -} - -Handle::Internal::Internal(int device_id) { - hcDeviceProp_t prop; - hcGetDeviceProperties(&prop, device_id); - _warp_size = prop.warpSize; - _max_threads_per_block = prop.maxThreadsPerBlock; - _block_size[0] = prop.maxThreadsDim[0]; - _block_size[1] = prop.maxThreadsDim[1]; - _block_size[2] = prop.maxThreadsDim[2]; - _grid_size[0] = prop.maxGridSize[0]; - _grid_size[1] = prop.maxGridSize[1]; - _grid_size[2] = prop.maxGridSize[2]; -} - -infiniStatus_t Handle::Internal::useMcblas(hcStream_t stream, const Fn &f) const { - auto handle = mcblas_handles.pop(); - if (!handle) { - CHECK_MCBLAS(hcblasCreate(&(*handle))); - } - CHECK_MCBLAS(hcblasSetStream(*handle, stream)); - CHECK_STATUS(f(*handle)); - mcblas_handles.push(std::move(*handle)); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Handle::Internal::useMcdnn(hcStream_t stream, const Fn &f) const { - auto handle = mcdnn_handles.pop(); - if (!handle) { - CHECK_MCDNN(hcdnnCreate(&(*handle))); - } - CHECK_MCDNN(hcdnnSetStream(*handle, stream)); - CHECK_STATUS(f(*handle)); - mcdnn_handles.push(std::move(*handle)); - return INFINI_STATUS_SUCCESS; -} - -int Handle::Internal::warpSize() const { return _warp_size; } -int Handle::Internal::maxThreadsPerBlock() const { return _max_threads_per_block; } -int Handle::Internal::blockSizeX() const { return _block_size[0]; } -int Handle::Internal::blockSizeY() const { return _block_size[1]; } -int Handle::Internal::blockSizeZ() const { return _block_size[2]; } -int Handle::Internal::gridSizeX() const { return _grid_size[0]; } -int Handle::Internal::gridSizeY() const { return _grid_size[1]; } -int Handle::Internal::gridSizeZ() const { return _grid_size[2]; } - -hcdnnDataType_t getHcdnnDtype(infiniDtype_t dt) { - switch (dt) { - case INFINI_DTYPE_F16: - return HCDNN_DATA_HALF; - case INFINI_DTYPE_F32: - return HCDNN_DATA_FLOAT; - case INFINI_DTYPE_F64: - return HCDNN_DATA_DOUBLE; - case INFINI_DTYPE_BF16: - return HCDNN_DATA_BFLOAT16; - case INFINI_DTYPE_I8: - return HCDNN_DATA_INT8; - case INFINI_DTYPE_I32: - return HCDNN_DATA_INT32; - case INFINI_DTYPE_I64: - return HCDNN_DATA_INT64; - case INFINI_DTYPE_U8: - return HCDNN_DATA_UINT8; - default: - return HCDNN_DATA_FLOAT; - } -} - -infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int device_id) { - *handle_ptr = new Handle(INFINI_DEVICE_METAX, device_id); - return INFINI_STATUS_SUCCESS; -} - -} // namespace device::metax diff --git a/src/infiniop/devices/metax/metax_handle.h b/src/infiniop/devices/metax/metax_handle.h deleted file mode 100644 index ff01150f9..000000000 --- a/src/infiniop/devices/metax/metax_handle.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __INFINIOP_METAX_HANDLE_H__ -#define __INFINIOP_METAX_HANDLE_H__ - -#include "../../handle.h" -#include - -namespace device::metax { -struct Handle : public InfiniopHandle { - Handle(int device_id); - class Internal; - auto internal() const -> const std::shared_ptr &; - -public: - static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id); - -protected: - Handle(infiniDevice_t device, int device_id); - -private: - std::shared_ptr _internal; -}; - -} // namespace device::metax - -#endif // __INFINIOP_METAX_HANDLE_H__ diff --git a/src/infiniop/devices/metax/metax_ht2mc.h b/src/infiniop/devices/metax/metax_ht2mc.h deleted file mode 100644 index 9e0da85ef..000000000 --- a/src/infiniop/devices/metax/metax_ht2mc.h +++ /dev/null @@ -1,156 +0,0 @@ -#ifdef ENABLE_METAX_MC_API -#define hpccDataType macaDataType -#define HPCC_R_32F MACA_R_32F -#define HPCC_R_64F MACA_R_64F -#define HPCC_R_16F MACA_R_16F -#define HPCC_R_16BF MACA_R_16BF -#define hpcc_bfloat162 maca_bfloat162 -#define hpcc_bfloat16 maca_bfloat16 -#define hcSuccess mcSuccess -#define hcStreamWaitEvent mcStreamWaitEvent -#define hcStreamSynchronize mcStreamSynchronize -#define hcStreamEndCapture mcStreamEndCapture -#define hcStreamDestroy mcStreamDestroy -#define hcStreamCreate mcStreamCreate -#define hcStreamCaptureModeThreadLocal mcStreamCaptureModeThreadLocal -#define hcStreamCaptureModeRelaxed mcStreamCaptureModeRelaxed -#define hcStreamCaptureModeGlobal mcStreamCaptureModeGlobal -#define hcStreamCaptureMode mcStreamCaptureMode -#define hcStreamBeginCapture mcStreamBeginCapture -#define hcStream_t mcStream_t -#define hcSetDevice mcSetDevice -#define hcMemset mcMemset -#define hcMemsetAsync mcMemsetAsync -#define hcMemcpyKind mcMemcpyKind -#define hcMemcpyHostToHost mcMemcpyHostToHost -#define hcMemcpyHostToDevice mcMemcpyHostToDevice -#define hcMemcpyDeviceToHost mcMemcpyDeviceToHost -#define hcMemcpyDeviceToDevice mcMemcpyDeviceToDevice -#define hcMemcpyDefault mcMemcpyDefault -#define hcMemcpyAsync mcMemcpyAsync -#define hcMemcpy mcMemcpy -#define hcMemGetInfo mcMemGetInfo -#define hcMallocHost mcMallocHost -#define hcMallocAsync mcMallocAsync -#define hcMalloc mcMalloc -#define hcLaunchKernel mcLaunchKernel -#define hcGraphNode_t mcGraphNode_t -#define hcGraphLaunch mcGraphLaunch -#define hcGraphInstantiate mcGraphInstantiate -#define hcGraphExecDestroy mcGraphExecDestroy -#define hcGraphExec_t mcGraphExec_t -#define hcGraphDestroy mcGraphDestroy -#define hcGraph_t mcGraph_t -#define hcGetLastError mcGetLastError -#define hcGetDeviceProperties mcGetDeviceProperties -#define hcGetDeviceCount mcGetDeviceCount -#define hcGetDevice mcGetDevice -#define hcFreeHost mcFreeHost -#define hcFreeAsync mcFreeAsync -#define hcFree mcFree -#define hcEventSynchronize mcEventSynchronize -#define hcEventRecord mcEventRecord -#define hcEventQuery mcEventQuery -#define hcEventElapsedTime mcEventElapsedTime -#define hcEventDisableTiming mcEventDisableTiming -#define hcEventDestroy mcEventDestroy -#define hcEventDefault mcEventDefault -#define hcEventCreateWithFlags mcEventCreateWithFlags -#define hcEventCreate mcEventCreate -#define hcEventBlockingSync mcEventBlockingSync -#define hcEvent_t mcEvent_t -#define hcError_t mcError_t -#define hcdnnTensorFormat_t mcdnnTensorFormat_t -#define hcdnnTensorDescriptor_t mcdnnTensorDescriptor_t -#define hcdnnStatus_t mcdnnStatus_t -#define hcdnnSetTensor4dDescriptor mcdnnSetTensor4dDescriptor -#define hcdnnSetTensorNdDescriptor mcdnnSetTensorNdDescriptor -#define hcdnnSetStream mcdnnSetStream -#define hcdnnSetPooling2dDescriptor mcdnnSetPooling2dDescriptor -#define hcdnnSetPoolingNdDescriptor mcdnnSetPoolingNdDescriptor -#define hcdnnPoolingMode_t mcdnnPoolingMode_t -#define hcdnnPoolingForward mcdnnPoolingForward -#define hcdnnPoolingDescriptor_t mcdnnPoolingDescriptor_t -#define hcdnnHandle_t mcdnnHandle_t -#define hcdnnDestroyTensorDescriptor mcdnnDestroyTensorDescriptor -#define hcdnnDestroyPoolingDescriptor mcdnnDestroyPoolingDescriptor -#define hcdnnDestroyPoolingDescriptor mcdnnDestroyPoolingDescriptor -#define hcdnnDataType_t mcdnnDataType_t -#define hcdnnCreateTensorDescriptor mcdnnCreateTensorDescriptor -#define hcdnnCreateTensorDescriptor mcdnnCreateTensorDescriptor -#define hcdnnCreatePoolingDescriptor mcdnnCreatePoolingDescriptor -#define hcdnnCreate mcdnnCreate -#define HCDNN_TENSOR_NCHW MCDNN_TENSOR_NCHW -#define HCDNN_STATUS_SUCCESS MCDNN_STATUS_SUCCESS -#define HCDNN_STATUS_SUCCESS MCDNN_STATUS_SUCCESS -#define HCDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING MCDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING -#define HCDNN_NOT_PROPAGATE_NAN MCDNN_NOT_PROPAGATE_NAN -#define HCDNN_DATA_UINT8 MCDNN_DATA_UINT8 -#define HCDNN_DATA_INT8 MCDNN_DATA_INT8 -#define HCDNN_DATA_INT64 MCDNN_DATA_INT64 -#define HCDNN_DATA_INT32 MCDNN_DATA_INT32 -#define HCDNN_DATA_HALF MCDNN_DATA_HALF -#define HCDNN_DATA_HALF MCDNN_DATA_HALF -#define HCDNN_DATA_FLOAT MCDNN_DATA_FLOAT -#define HCDNN_DATA_FLOAT MCDNN_DATA_FLOAT -#define HCDNN_DATA_DOUBLE MCDNN_DATA_DOUBLE -#define HCDNN_DATA_BFLOAT16 MCDNN_DATA_BFLOAT16 -#define hcDeviceSynchronize mcDeviceSynchronize -#define hcDeviceProp_t mcDeviceProp_t -#define hcDeviceGetAttribute mcDeviceGetAttribute -#define hcDeviceAttributeMultiProcessorCount mcDeviceAttributeMultiProcessorCount -#define hcclSum mcclSum -#define hcclSuccess mcclSuccess -#define hcclRedOp_t mcclRedOp_t -#define hcclProd mcclProd -#define hcclMin mcclMin -#define hcclMax mcclMax -#define hcclHalf mcclHalf -#define hcclFloat mcclFloat -#define hcclDataType_t mcclDataType_t -#define hcclCommInitAll mcclCommInitAll -#define hcclCommDestroy mcclCommDestroy -#define hcclComm_t mcclComm_t -#define hcclBfloat16 mcclBfloat16 -#define hcclAvg mcclAvg -#define hcclAllReduce mcclAllReduce -#define hcblasSetStream mcblasSetStream -#define hcblasHandle_t mcblasHandle_t -#define hcblasGemmStridedBatchedEx mcblasGemmStridedBatchedEx -#define hcblasGemmEx mcblasGemmEx -#define hcblasCreate mcblasCreate -#define hcblasComputeType_t mcblasComputeType_t -#define hcblasSetPointerMode mcblasSetPointerMode -#define hcblasIsamax mcblasIsamax -#define hcblasIdamax mcblasIdamax -#define hcblasIsamin mcblasIsamin -#define hcblasIdamin mcblasIdamin -#define hcblasSasum mcblasSasum -#define hcblasDasum mcblasDasum -#define hcblasAxpyEx mcblasAxpyEx -#define hcblasScopy mcblasScopy -#define hcblasDcopy mcblasDcopy -#define hcblasDotEx mcblasDotEx -#define hcblasNrm2Ex mcblasNrm2Ex -#define hcblasRotEx mcblasRotEx -#define hcblasSrotg mcblasSrotg -#define hcblasDrotg mcblasDrotg -#define hcblasSrotm mcblasSrotm -#define hcblasDrotm mcblasDrotm -#define hcblasSrotmg mcblasSrotmg -#define hcblasDrotmg mcblasDrotmg -#define hcblasScalEx mcblasScalEx -#define hcblasSswap mcblasSswap -#define hcblasDswap mcblasDswap -#define HCBLAS_STATUS_SUCCESS MCBLAS_STATUS_SUCCESS -#define HCBLAS_OP_T MCBLAS_OP_T -#define HCBLAS_OP_N MCBLAS_OP_N -#define HCBLAS_GEMM_DEFAULT_TENSOR_OP MCBLAS_GEMM_DEFAULT_TENSOR_OP -#define HCBLAS_GEMM_DEFAULT MCBLAS_GEMM_DEFAULT -#define HCBLAS_COMPUTE_32F_FAST_TF32 MCBLAS_COMPUTE_32F_FAST_TF32 -#define HCBLAS_COMPUTE_32F MCBLAS_COMPUTE_32F -#define HCBLAS_POINTER_MODE_DEVICE MCBLAS_POINTER_MODE_DEVICE -#define HCBLAS_POINTER_MODE_HOST MCBLAS_POINTER_MODE_HOST -#define __hpcc_fp8_e4m3 __maca_fp8_e4m3 -#define __hpcc_bfloat16 __maca_bfloat16 -#endif diff --git a/src/infiniop/devices/metax/metax_kernel_common.h b/src/infiniop/devices/metax/metax_kernel_common.h deleted file mode 100644 index 7066a6ebd..000000000 --- a/src/infiniop/devices/metax/metax_kernel_common.h +++ /dev/null @@ -1,78 +0,0 @@ -#pragma once -#define INFINIOP_METAX_KERNEL __global__ void - -#ifdef ENABLE_METAX_MC_API -#include -#include -#include -#include -#else -#include -#include -#include -#include -#endif - -// Posible maximum number of threads per block for METAX architectures -// Used for picking correct kernel launch configuration -#define METAX_BLOCK_SIZE_512 512 -#define METAX_BLOCK_SIZE_1024 1024 -#define METAX_BLOCK_SIZE_2048 2048 -#define METAX_BLOCK_SIZE_4096 4096 - -#define CHECK_METAX(API) CHECK_INTERNAL(API, hcSuccess) - -using cuda_bfloat16 = hpcc_bfloat16; -using cuda_bfloat162 = hpcc_bfloat162; -using cuda_fp8_e4m3 = __hpcc_fp8_e4m3; - -#ifdef ENABLE_METAX_MC_API -using __nv_bfloat16 = __maca_bfloat16; -#else -using __nv_bfloat16 = __hpcc_bfloat16; -#endif - -namespace device::metax { - -// get the memory offset of the given element in a tensor given its flat index -__forceinline__ __device__ __host__ size_t -indexToOffset( - size_t flat_index, - size_t ndim, - const size_t *shape, - const ptrdiff_t *strides) { - size_t res = 0; - for (size_t i = ndim; i-- > 0;) { - res += (flat_index % shape[i]) * strides[i]; - flat_index /= shape[i]; - } - return res; -} -} // namespace device::metax - -using device::metax::indexToOffset; - -__forceinline__ __device__ float -exp_(const float val) { - return expf(val); -} - -__forceinline__ __device__ long double -exp_(const long double val) { - return exp(val); -} - -__forceinline__ __device__ double -exp_(const double val) { - return exp(val); -} - -__forceinline__ __device__ __half -exp_(const __half x) { - return hexp(x); -} - -__forceinline__ __device__ __hpcc_bfloat16 -exp_(const __hpcc_bfloat16 x) { - return hexp(x); -} diff --git a/src/infiniop/devices/moore/moore_common.h b/src/infiniop/devices/moore/moore_common.h deleted file mode 100644 index c60677127..000000000 --- a/src/infiniop/devices/moore/moore_common.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef __INFINIOP_MOORE_COMMON_H__ -#define __INFINIOP_MOORE_COMMON_H__ -#include "../../../utils.h" -#include "../pool.h" -#include "moore_handle.h" -#include -#include -#include -#include -#include - -#define CHECK_MUBLAS(API) CHECK_INTERNAL(API, MUBLAS_STATUS_SUCCESS) -#define CHECK_MUDNN(API) CHECK_INTERNAL((int)API, (int)::musa::dnn::Status::SUCCESS) - -namespace device::moore { - -class Handle::Internal { - Pool> mublas_handles; - Pool> mudnn_handles; - - int _warp_size, - _max_threads_per_block, - _block_size[3], - _grid_size[3]; - - int _device_id; - template - using Fn = std::function; - -public: - Internal(int device_id); - - infiniStatus_t useMublas(musaStream_t stream, const Fn &f) const; - infiniStatus_t useMudnn(musaStream_t stream, const Fn<::musa::dnn::Handle &> &f) const; - - int warpSize() const; - int maxThreadsPerBlock() const; - int blockSizeX() const; - int blockSizeY() const; - int blockSizeZ() const; - int gridSizeX() const; - int gridSizeY() const; - int gridSizeZ() const; -}; - -} // namespace device::moore -#endif diff --git a/src/infiniop/devices/moore/moore_handle.cc b/src/infiniop/devices/moore/moore_handle.cc deleted file mode 100644 index 5404cb9a2..000000000 --- a/src/infiniop/devices/moore/moore_handle.cc +++ /dev/null @@ -1,71 +0,0 @@ -#include "moore_common.h" - -namespace device::moore { -Handle::Handle(infiniDevice_t device, int device_id) - : InfiniopHandle{device, device_id}, - _internal(std::make_shared(device_id)) {} - -Handle::Handle(int device_id) : Handle(INFINI_DEVICE_MOORE, device_id) {} - -auto Handle::internal() const -> const std::shared_ptr & { - return _internal; -} - -Handle::Internal::Internal(int device_id) - : _device_id(device_id) { - musaDeviceProp prop; - musaGetDeviceProperties(&prop, device_id); - _warp_size = prop.warpSize; - _max_threads_per_block = prop.maxThreadsPerBlock; - _block_size[0] = prop.maxThreadsDim[0]; - _block_size[1] = prop.maxThreadsDim[1]; - _block_size[2] = prop.maxThreadsDim[2]; - _grid_size[0] = prop.maxGridSize[0]; - _grid_size[1] = prop.maxGridSize[1]; - _grid_size[2] = prop.maxGridSize[2]; -} - -infiniStatus_t Handle::Internal::useMublas(musaStream_t stream, const Fn &f) const { - std::unique_ptr handle; - auto opt_handle = mublas_handles.pop(); - if (opt_handle.has_value()) { - handle = std::move(*opt_handle); - } else { - handle = std::make_unique(); - CHECK_MUBLAS(mublasCreate(&(*handle))); - } - CHECK_MUBLAS(mublasSetStream(*handle, stream)); - CHECK_STATUS(f(*handle)); - mublas_handles.push(std::move(handle)); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Handle::Internal::useMudnn(musaStream_t stream, const Fn<::musa::dnn::Handle &> &f) const { - std::unique_ptr<::musa::dnn::Handle> handle; - auto opt_handle = mudnn_handles.pop(); - if (opt_handle.has_value()) { - handle = std::move(*opt_handle); - } else { - handle = std::make_unique<::musa::dnn::Handle>(_device_id); - } - CHECK_MUDNN(handle->SetStream(stream)); - CHECK_STATUS(f(*handle)); - mudnn_handles.push(std::move(handle)); - return INFINI_STATUS_SUCCESS; -} - -int Handle::Internal::warpSize() const { return _warp_size; } -int Handle::Internal::maxThreadsPerBlock() const { return _max_threads_per_block; } -int Handle::Internal::blockSizeX() const { return _block_size[0]; } -int Handle::Internal::blockSizeY() const { return _block_size[1]; } -int Handle::Internal::blockSizeZ() const { return _block_size[2]; } -int Handle::Internal::gridSizeX() const { return _grid_size[0]; } -int Handle::Internal::gridSizeY() const { return _grid_size[1]; } -int Handle::Internal::gridSizeZ() const { return _grid_size[2]; } - -infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int device_id) { - *handle_ptr = new Handle(INFINI_DEVICE_MOORE, device_id); - return INFINI_STATUS_SUCCESS; -} - -} // namespace device::moore diff --git a/src/infiniop/devices/moore/moore_handle.h b/src/infiniop/devices/moore/moore_handle.h deleted file mode 100644 index 52157dd84..000000000 --- a/src/infiniop/devices/moore/moore_handle.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef __INFINIOP_MOORE_HANDLE_H__ -#define __INFINIOP_MOORE_HANDLE_H__ - -#include "../../handle.h" -#include - -namespace device::moore { -struct Handle : public InfiniopHandle { - Handle(int device_id); - class Internal; - auto internal() const -> const std::shared_ptr &; - -public: - static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id); - -protected: - Handle(infiniDevice_t device, int device_id); - -private: - std::shared_ptr _internal; -}; - -} // namespace device::moore - -#endif // __INFINIOP_MOORE_HANDLE_H__ diff --git a/src/infiniop/devices/moore/moore_kernel_common.h b/src/infiniop/devices/moore/moore_kernel_common.h deleted file mode 100644 index e121fddf1..000000000 --- a/src/infiniop/devices/moore/moore_kernel_common.h +++ /dev/null @@ -1,77 +0,0 @@ -#ifndef __INFINIOP_MOORE_KERNEL_COMMON_H__ -#define __INFINIOP_MOORE_KERNEL_COMMON_H__ -#define INFINIOP_MOORE_KERNEL __global__ void - -#include -#include -#include - -// Posible maximum number of threads per block for MUSA architectures -// Used for picking correct kernel launch configuration -#define MOORE_BLOCK_SIZE_4096 4096 -#define MOORE_BLOCK_SIZE_2048 2048 -#define MOORE_BLOCK_SIZE_1024 1024 -#define MOORE_BLOCK_SIZE_512 512 - -#define CHECK_MOORE(API) CHECK_INTERNAL(API, musaSuccess) - -using cuda_bfloat16 = mt_bfloat16; -using cuda_bfloat162 = mt_bfloat162; -using cuda_fp8_e4m3 = __mt_fp8_e4m3; - -using __nv_bfloat16 = __mt_bfloat16; - -namespace device::moore { - -// get the memory offset of the given element in a tensor given its flat index -__forceinline__ __device__ __host__ size_t -indexToOffset( - size_t flat_index, - size_t ndim, - const size_t *shape, - const ptrdiff_t *strides) { - size_t res = 0; - for (size_t i = ndim; i-- > 0;) { - res += (flat_index % shape[i]) * strides[i]; - flat_index /= shape[i]; - } - return res; -} -} // namespace device::moore - -using device::moore::indexToOffset; - -__forceinline__ __device__ float -exp_(const float val) { - return expf(val); -} - -// Computes exp for long double on Moore GPU, -// casts to double to resolve ambiguous exp call, -// due to conflicting double/float definitions in MUSA math libraries. -__forceinline__ __device__ long double -exp_(const long double val) { - return static_cast(exp(static_cast(val))); -} - -__forceinline__ __device__ double -exp_(const double val) { - return exp(val); -} - -// may not support hexp -__forceinline__ __device__ __half -exp_(const __half x) { - float f_val = __half2float(x); - float f_result = expf(f_val); - return __float2half(f_result); -} - -// may not support hexp -__forceinline__ __device__ __mt_bfloat16 -exp_(const __mt_bfloat16 x) { - float f_val = __bfloat162float(x); - float f_result = expf(f_val); - return __float2bfloat16(f_result); -} -#endif diff --git a/src/infiniop/devices/nvidia/nvidia_common.cu b/src/infiniop/devices/nvidia/nvidia_common.cu deleted file mode 100644 index eca122dce..000000000 --- a/src/infiniop/devices/nvidia/nvidia_common.cu +++ /dev/null @@ -1,148 +0,0 @@ -#include "nvidia_handle.cuh" - -namespace device { - -namespace nvidia { - -Handle::Handle(infiniDevice_t device, int device_id) - : InfiniopHandle{device, device_id}, - _internal(std::make_shared(device_id)) {} - -auto Handle::internal() const -> const std::shared_ptr & { - return _internal; -} - -Handle::Internal::Internal(int device_id) { - cudaDeviceProp prop; - cudaGetDeviceProperties(&prop, device_id); - _warp_size = prop.warpSize; - _max_threads_per_block = prop.maxThreadsPerBlock; - _block_size[0] = prop.maxThreadsDim[0]; - _block_size[1] = prop.maxThreadsDim[1]; - _block_size[2] = prop.maxThreadsDim[2]; - _grid_size[0] = prop.maxGridSize[0]; - _grid_size[1] = prop.maxGridSize[1]; - _grid_size[2] = prop.maxGridSize[2]; - this->useCublas(nullptr, [](cublasHandle_t handle) { return INFINI_STATUS_SUCCESS; }); -#ifdef ENABLE_CUDNN_API - this->useCudnn(nullptr, [](cudnnHandle_t handle) { return INFINI_STATUS_SUCCESS; }); -#endif -} - -infiniStatus_t Handle::Internal::useCublas(cudaStream_t stream, const Fn &f) const { - auto handle = blas_handles.pop(); - if (!handle) { - CHECK_CUBLAS(cublasCreate(&(*handle))); - } - CHECK_CUBLAS(cublasSetStream(*handle, stream)); - CHECK_STATUS(f(*handle)); - blas_handles.push(std::move(*handle)); - return INFINI_STATUS_SUCCESS; -} - -#ifdef ENABLE_CUDNN_API -infiniStatus_t Handle::Internal::useCudnn(cudaStream_t stream, const Fn &f) const { - auto handle = dnn_handles.pop(); - if (!handle) { - CHECK_CUDNN(cudnnCreate(&(*handle))); - } - CHECK_CUDNN(cudnnSetStream(*handle, stream)); - CHECK_STATUS(f(*handle)); - dnn_handles.push(std::move(*handle)); - return INFINI_STATUS_SUCCESS; -} -#endif - -int Handle::Internal::warpSize() const { return _warp_size; } -int Handle::Internal::maxThreadsPerBlock() const { return _max_threads_per_block; } -int Handle::Internal::blockSizeX() const { return _block_size[0]; } -int Handle::Internal::blockSizeY() const { return _block_size[1]; } -int Handle::Internal::blockSizeZ() const { return _block_size[2]; } -int Handle::Internal::gridSizeX() const { return _grid_size[0]; } -int Handle::Internal::gridSizeY() const { return _grid_size[1]; } -int Handle::Internal::gridSizeZ() const { return _grid_size[2]; } - -#ifdef ENABLE_CUDNN_API -cudnnDataType_t getCudnnDtype(infiniDtype_t dt) { - switch (dt) { - case INFINI_DTYPE_F16: - return CUDNN_DATA_HALF; - case INFINI_DTYPE_F32: - return CUDNN_DATA_FLOAT; -#ifndef ENABLE_QY_API - case INFINI_DTYPE_F64: - return CUDNN_DATA_DOUBLE; -#endif - case INFINI_DTYPE_BF16: - return CUDNN_DATA_BFLOAT16; - case INFINI_DTYPE_I8: - return CUDNN_DATA_INT8; - case INFINI_DTYPE_I32: - return CUDNN_DATA_INT32; -#ifndef ENABLE_ILUVATAR_API - case INFINI_DTYPE_I64: - return CUDNN_DATA_INT64; -#endif - case INFINI_DTYPE_U8: - return CUDNN_DATA_UINT8; - default: - return CUDNN_DATA_FLOAT; - } -} -#endif - -infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int device_id) { - *handle_ptr = new Handle(INFINI_DEVICE_NVIDIA, device_id); - return INFINI_STATUS_SUCCESS; -} - -} // namespace nvidia - -namespace iluvatar { - -Handle::Handle(int device_id) - : nvidia::Handle(INFINI_DEVICE_ILUVATAR, device_id) {} - -infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int device_id) { - *handle_ptr = new Handle(device_id); - return INFINI_STATUS_SUCCESS; -} - -} // namespace iluvatar - -namespace ali { - -Handle::Handle(int device_id) - : nvidia::Handle(INFINI_DEVICE_ALI, device_id) {} - -infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int device_id) { - *handle_ptr = new Handle(device_id); - return INFINI_STATUS_SUCCESS; -} - -} // namespace ali - -namespace qy { - -Handle::Handle(int device_id) - : nvidia::Handle(INFINI_DEVICE_QY, device_id) {} - -infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int device_id) { - *handle_ptr = new Handle(device_id); - return INFINI_STATUS_SUCCESS; -} -} // namespace qy - -namespace hygon { - -Handle::Handle(int device_id) - : nvidia::Handle(INFINI_DEVICE_HYGON, device_id) {} - -infiniStatus_t Handle::create(InfiniopHandle **handle_ptr, int device_id) { - *handle_ptr = new Handle(device_id); - return INFINI_STATUS_SUCCESS; -} - -} // namespace hygon - -} // namespace device diff --git a/src/infiniop/devices/nvidia/nvidia_common.cuh b/src/infiniop/devices/nvidia/nvidia_common.cuh deleted file mode 100644 index 8f5055c0b..000000000 --- a/src/infiniop/devices/nvidia/nvidia_common.cuh +++ /dev/null @@ -1,15 +0,0 @@ -#ifndef __INFINIOP_CUDA_COMMON_CUH__ -#define __INFINIOP_CUDA_COMMON_CUH__ - -#include "infinicore.h" -#include "nvidia_handle.cuh" - -namespace device::nvidia { - -#ifdef ENABLE_CUDNN_API -cudnnDataType_t getCudnnDtype(infiniDtype_t dt); -#endif - -} // namespace device::nvidia - -#endif // __INFINIOP_CUDA_COMMON_CUH__ diff --git a/src/infiniop/devices/nvidia/nvidia_handle.cuh b/src/infiniop/devices/nvidia/nvidia_handle.cuh deleted file mode 100644 index 1dcb4521e..000000000 --- a/src/infiniop/devices/nvidia/nvidia_handle.cuh +++ /dev/null @@ -1,53 +0,0 @@ -#ifndef __INFINIOP_CUDA_HANDLE_CUH__ -#define __INFINIOP_CUDA_HANDLE_CUH__ - -#include "../../../utils.h" -#include "../pool.h" -#include "nvidia_handle.h" -#include -#include - -#ifdef ENABLE_CUDNN_API -#include -#endif - -#define CHECK_CUBLAS(API) CHECK_INTERNAL(API, CUBLAS_STATUS_SUCCESS) -#define CHECK_CUDNN(API) CHECK_INTERNAL(API, CUDNN_STATUS_SUCCESS) - -namespace device::nvidia { - -class Handle::Internal { - Pool blas_handles; -#ifdef ENABLE_CUDNN_API - Pool dnn_handles; -#endif - - int _warp_size, - _max_threads_per_block, - _block_size[3], - _grid_size[3]; - - template - using Fn = std::function; - -public: - Internal(int); - - infiniStatus_t useCublas(cudaStream_t stream, const Fn &f) const; -#ifdef ENABLE_CUDNN_API - infiniStatus_t useCudnn(cudaStream_t stream, const Fn &f) const; -#endif - - int warpSize() const; - int maxThreadsPerBlock() const; - int blockSizeX() const; - int blockSizeY() const; - int blockSizeZ() const; - int gridSizeX() const; - int gridSizeY() const; - int gridSizeZ() const; -}; - -} // namespace device::nvidia - -#endif // __INFINIOP_CUDA_HANDLE_CUH__ diff --git a/src/infiniop/devices/nvidia/nvidia_handle.h b/src/infiniop/devices/nvidia/nvidia_handle.h deleted file mode 100644 index b644266e8..000000000 --- a/src/infiniop/devices/nvidia/nvidia_handle.h +++ /dev/null @@ -1,73 +0,0 @@ -#ifndef __INFINIOP_CUDA_HANDLE_H__ -#define __INFINIOP_CUDA_HANDLE_H__ - -#include "../../handle.h" -#include - -namespace device { - -namespace nvidia { - -struct Handle : public InfiniopHandle { - class Internal; - auto internal() const -> const std::shared_ptr &; - -protected: - Handle(infiniDevice_t device, int device_id); - -public: - static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id); - -private: - std::shared_ptr _internal; -}; - -} // namespace nvidia - -namespace iluvatar { - -struct Handle : public nvidia::Handle { - Handle(int device_id); - -public: - static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id); -}; - -} // namespace iluvatar - -namespace ali { - -struct Handle : public nvidia::Handle { - Handle(int device_id); - -public: - static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id); -}; - -} // namespace ali - -namespace qy { - -struct Handle : public nvidia::Handle { - Handle(int device_id); - -public: - static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id); -}; - -} // namespace qy - -namespace hygon { - -struct Handle : public nvidia::Handle { - Handle(int device_id); - -public: - static infiniStatus_t create(InfiniopHandle **handle_ptr, int device_id); -}; - -} // namespace hygon - -} // namespace device - -#endif // __INFINIOP_CUDA_HANDLE_H__ diff --git a/src/infiniop/devices/nvidia/nvidia_kernel_common.cuh b/src/infiniop/devices/nvidia/nvidia_kernel_common.cuh deleted file mode 100644 index 2d1b50148..000000000 --- a/src/infiniop/devices/nvidia/nvidia_kernel_common.cuh +++ /dev/null @@ -1,82 +0,0 @@ -#ifndef __INFINIOP_CUDA_KERNEL_COMMON_CUH__ -#define __INFINIOP_CUDA_KERNEL_COMMON_CUH__ - -#if defined(ENABLE_HYGON_API) -#define INFINIOP_CUDA_KERNEL __launch_bounds__(1024) __global__ void -#else -#define INFINIOP_CUDA_KERNEL __global__ void -#endif - -#include -#include -#ifndef ENABLE_HYGON_API -#include -#endif - -// Posible maximum number of threads per block for CUDA architectures -// Used for picking correct kernel launch configuration -#define CUDA_BLOCK_SIZE_4096 4096 -#define CUDA_BLOCK_SIZE_2048 2048 -#define CUDA_BLOCK_SIZE_1024 1024 -#define CUDA_BLOCK_SIZE_512 512 - -#define CHECK_CUDA(API) CHECK_INTERNAL(API, cudaSuccess) - -#ifdef ENABLE_HYGON_API -// Hygon DCU uses different bfloat16 type definitions -using cuda_bfloat16 = __nv_bfloat16; -using cuda_bfloat162 = __nv_bfloat162; -#else -using cuda_bfloat16 = nv_bfloat16; -using cuda_bfloat162 = nv_bfloat162; -using cuda_fp8_e4m3 = __nv_fp8_e4m3; -#endif - -namespace device::nvidia { - -// get the memory offset of the given element in a tensor given its flat index -__forceinline__ __device__ __host__ size_t -indexToOffset( - size_t flat_index, - size_t ndim, - const size_t *shape, - const ptrdiff_t *strides) { - size_t res = 0; - for (size_t i = ndim; i-- > 0;) { - res += (flat_index % shape[i]) * strides[i]; - flat_index /= shape[i]; - } - return res; -} -} // namespace device::nvidia - -using device::nvidia::indexToOffset; - -__forceinline__ __device__ float -exp_(const float val) { - return expf(val); -} - -#if !defined(ENABLE_ILUVATAR_API) && !defined(ENABLE_QY_API) && !defined(ENABLE_HYGON_API) && !defined(ENABLE_ALI_API) -__forceinline__ __device__ long double -exp_(const long double val) { - return expl(val); -} -#endif - -__forceinline__ __device__ double -exp_(const double val) { - return exp(val); -} - -__forceinline__ __device__ __half -exp_(const __half x) { - return hexp(x); -} - -__forceinline__ __device__ __nv_bfloat16 -exp_(const __nv_bfloat16 x) { - return hexp(x); -} - -#endif // __INFINIOP_CUDA_KERNEL_COMMON_CUH__ diff --git a/src/infiniop/devices/pool.h b/src/infiniop/devices/pool.h deleted file mode 100644 index ed15a075b..000000000 --- a/src/infiniop/devices/pool.h +++ /dev/null @@ -1,50 +0,0 @@ -#ifndef __POOL_H__ -#define __POOL_H__ - -#include -#include -#include - -template -class Pool { -public: - Pool() : _head(nullptr) {} - - Pool(const Pool &) = delete; - - Pool(Pool &&pool) noexcept : _head(pool._head.exchange(nullptr)) {} - - ~Pool() { - while (this->pop()) {} - } - - void push(T &&val) const { - Node *new_node = new Node(std::move(val)); - new_node->next = _head.load(); - while (!_head.compare_exchange_weak(new_node->next, new_node)) {} - } - - std::optional pop() const { - Node *top = _head.load(); - Node *new_head = nullptr; - do { - if (!top) { - return std::nullopt; - } - new_head = top->next; - } while (!_head.compare_exchange_weak(top, new_head)); - return {std::move(top->data)}; - } - -private: - template - struct Node { - U data; - Node *next; - Node(U &&data) : data(std::move(data)), next(nullptr) {} - }; - - mutable std::atomic *> _head; -}; - -#endif // __POOL_H__ diff --git a/src/infiniop/elementwise/bang/elementwise_bang.h b/src/infiniop/elementwise/bang/elementwise_bang.h deleted file mode 100644 index cf5c0a63b..000000000 --- a/src/infiniop/elementwise/bang/elementwise_bang.h +++ /dev/null @@ -1,212 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_BANG_H__ -#define __INFINIOP_ELEMENTWISE_BANG_H__ - -#include "../../../utils.h" -#include "../../devices/bang/common_bang.h" -#include "elementwise_bang_api.h" - -namespace op::elementwise::bang { - -/** - * @brief Opaque implementation structure for BANG device operations. - * - * Contains device-specific resources and implementation methods. - */ -struct DeviceImpl::Opaque { - std::shared_ptr internal; - - /** - * @brief Constructs an Opaque instance with device handle internals. - * - * @param internal_ Shared pointer to BANG device handle internals. - */ - Opaque(const std::shared_ptr &internal_) - : internal(internal_) {} - - /** - * @brief Implements elementwise calculation for BANG device. - * - * @tparam N Number of input tensors. - * @tparam Op Operator functor type. - * @tparam Tdata Data type for inputs and output. - * @tparam Args Additional arguments for the operator. - * - * @param info Elementwise operation metadata (shapes, strides, etc.). - * @param workspace Device workspace memory. - * @param output Output tensor buffer. - * @param inputs Vector of input tensor pointers. - * @param queue BANG queue for asynchronous execution. - * @param args Additional arguments for the operator. - * @return infiniStatus_t Status indicating success or failure. - */ - template - infiniStatus_t calculateImpl(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - cnrtQueue_t queue, - Args &&...args) { - auto output_size = info.getOutputSize(); - if (output_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - // Device pointers for metadata - const void **d_inputs_arr = nullptr; - const bool *d_input_contiguous = nullptr; - const bool *d_input_broadcasted = nullptr; - const size_t *d_output_shape = nullptr; - const ptrdiff_t *d_output_strides = nullptr; - const size_t *d_input_shapes = nullptr; - const ptrdiff_t *d_input_strides = nullptr; - - // Copy metadata to device and setup pointers - CHECK_STATUS(infoToDevice(info, workspace, inputs.data(), d_inputs_arr, - d_input_contiguous, d_input_broadcasted, - d_output_shape, d_output_strides, - d_input_shapes, d_input_strides)); - - // Launch the elementwise kernel - Op::template launch( - output_size, - info.getNdim(), - info.isOutputContiguous(), - reinterpret_cast(d_input_contiguous), - reinterpret_cast(d_input_broadcasted), - reinterpret_cast(d_output_shape), - reinterpret_cast(d_input_shapes), - reinterpret_cast(d_output_strides), - reinterpret_cast(d_input_strides), - output, - reinterpret_cast(d_inputs_arr), - queue, - internal, - args...); - - // Synchronize queue to ensure completion - CNRT_CHECK(cnrtQueueSync(queue)); - - return INFINI_STATUS_SUCCESS; - } - -private: - /** - * @brief Transfers elementwise operation metadata to device memory. - * - * @tparam N Number of input tensors. - * - * @param info Elementwise operation metadata. - * @param workspace Device workspace memory. - * @param h_inputs_arr Host array of input pointers. - * @param d_inputs_arr Output reference to device input pointers. - * @param d_input_contiguous Output reference to contiguous flags. - * @param d_input_broadcasted Output reference to broadcast flags. - * @param d_output_shape Output reference to output shape. - * @param d_output_strides Output reference to output strides. - * @param d_input_shapes Output reference to input shapes. - * @param d_input_strides Output reference to input strides. - * @return infiniStatus_t Status indicating success or failure. - */ - template - infiniStatus_t infoToDevice( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - const void *const *h_inputs_arr, - const void **&d_inputs_arr, - const bool *&d_input_contiguous, - const bool *&d_input_broadcasted, - const size_t *&d_output_shape, - const ptrdiff_t *&d_output_strides, - const size_t *&d_input_shapes, - const ptrdiff_t *&d_input_strides) const { - - constexpr auto input_size = N; - const auto ndim = info.getNdim(); - constexpr auto input_arr_size = N * sizeof(*h_inputs_arr); - const int8_t *info_meta_start = info.getMetaStart(); - const int8_t *d_meta_start = reinterpret_cast(workspace) + input_arr_size; - - // Copy input pointer array and metadata to device - CNRT_CHECK(cnrtMemcpy(workspace, (void *)h_inputs_arr, input_arr_size, cnrtMemcpyHostToDev)); - CNRT_CHECK(cnrtMemcpy((void *)d_meta_start, (void *)info_meta_start, info.getMetaMemSize(), cnrtMemcpyHostToDev)); - - // Setup pointers to device memory regions - d_inputs_arr = reinterpret_cast(workspace); - d_output_shape = reinterpret_cast(d_meta_start); - d_output_strides = reinterpret_cast(d_output_shape + ndim); - d_input_shapes = reinterpret_cast(d_output_strides + ndim); - d_input_strides = reinterpret_cast(d_input_shapes + input_size * ndim); - d_input_contiguous = reinterpret_cast(d_input_strides + input_size * ndim); - d_input_broadcasted = reinterpret_cast(d_input_contiguous + input_size); - - return INFINI_STATUS_SUCCESS; - } -}; - -/** - * @brief Creates a DeviceImpl instance for BANG device. - * - * @tparam Args Argument types for Opaque construction. - * @param args Arguments forwarded to Opaque constructor. - * @return utils::Result Result containing new DeviceImpl instance. - */ -template -utils::Result DeviceImpl::create(Args &&...args) { - auto opaque = std::make_shared(std::forward(args)...); - return utils::Result(new DeviceImpl(opaque)); -} - -/** - * @brief Calculates elementwise operation for BANG device. - * - * @tparam Op Operator functor type. - * @tparam Tdata Data type for inputs and output. - * @tparam Args Additional arguments for the operator. - * - * @param info Elementwise operation metadata. - * @param workspace Device workspace memory. - * @param output Output tensor buffer. - * @param inputs Vector of input tensor pointers. - * @param queue BANG queue (as void*). - * @param args Additional arguments for the operator. - * @return infiniStatus_t Status indicating success or failure. - */ -template -infiniStatus_t DeviceImpl::calculate(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *queue, - Args &&...args) { - constexpr size_t N = Op::num_inputs; - return _opaque->calculateImpl( - info, workspace, output, inputs, - reinterpret_cast(queue), - std::forward(args)...); -} -} // namespace op::elementwise::bang - -/** - * @brief Macro for declaring BANG kernel interface. - * - * @param OpName Name of the elementwise operation. - */ -#define LAUNCH_ELEMENTWISE_KERNEL(OpName) \ - template \ - void launch##OpName##Kernel( \ - size_t output_size, \ - size_t ndim, \ - bool output_contiguous, \ - const void *input_contiguous, \ - const void *input_broadcasted, \ - const void *output_shape, \ - const void *input_shapes, \ - const void *output_strides, \ - const void *input_strides, \ - void *output, \ - const void *const *inputs, \ - cnrtQueue_t queue, \ - const std::shared_ptr &internal, \ - Args... args); - -#endif // __INFINIOP_ELEMENTWISE_BANG_H__ diff --git a/src/infiniop/elementwise/bang/elementwise_bang_api.h b/src/infiniop/elementwise/bang/elementwise_bang_api.h deleted file mode 100644 index 65e8ad170..000000000 --- a/src/infiniop/elementwise/bang/elementwise_bang_api.h +++ /dev/null @@ -1,84 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_BANG_API_H__ -#define __INFINIOP_ELEMENTWISE_BANG_API_H__ - -#include "../elementwise.h" - -namespace op::elementwise::bang { - -/** - * @brief BANG device implementation for elementwise operations. - * - * Provides interface for creating and executing elementwise operations on BANG devices. - */ -class DeviceImpl final { - struct Opaque; - std::shared_ptr _opaque; - - DeviceImpl(std::shared_ptr opaque) : _opaque(std::move(opaque)) {} - -public: - ~DeviceImpl() = default; - - /** - * @brief Creates a DeviceImpl instance. - * - * @tparam Args Argument types for construction. - * @param args Arguments forwarded to implementation. - * @return utils::Result Result containing new instance. - */ - template - static utils::Result create(Args &&...args); - - /** - * @brief Executes elementwise operation on BANG device. - * - * @tparam Op Operator functor type. - * @tparam Tdata Data type for inputs and output. - * @tparam Args Additional arguments for the operator. - * - * @param info Elementwise operation metadata. - * @param workspace Device workspace memory. - * @param output Output tensor buffer. - * @param inputs Vector of input tensor pointers. - * @param queue BANG queue (as void*). - * @param args Additional arguments for the operator. - * @return infiniStatus_t Status indicating success or failure. - */ - template - infiniStatus_t calculate( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *queue, - Args &&...args); -}; -} // namespace op::elementwise::bang - -/** - * @brief Macro for creating BANG elementwise operation descriptor. - * - * @param HANDLE Device handle. - * @param DTYPE Output data type. - * @param OUT_DESC Output tensor descriptor. - * @param INPUT_DESC_VEC Vector of input tensor descriptors. - */ -#define CREATE_ELEMENTWISE_BANG_DESCRIPTOR(HANDLE, DTYPE, OUT_DESC, INPUT_DESC_VEC) \ - \ - auto info_result = op::elementwise::ElementwiseInfo::create(OUT_DESC, INPUT_DESC_VEC); \ - CHECK_RESULT(info_result); \ - auto info = info_result.take(); \ - auto workspace_size = info.getMetaMemSize() + info.getInputSize() * sizeof(void *); \ - \ - auto device_impl_result = op::elementwise::bang::DeviceImpl::create(HANDLE->internal()); \ - CHECK_RESULT(device_impl_result); \ - \ - *desc_ptr = new Descriptor( \ - DTYPE, \ - std::move(info), \ - std::move(device_impl_result.take()), \ - workspace_size, \ - HANDLE->device, \ - HANDLE->device_id); - -#endif // __INFINIOP_ELEMENTWISE_BANG_API_H__ diff --git a/src/infiniop/elementwise/bang/elementwise_bang_kernel.h b/src/infiniop/elementwise/bang/elementwise_bang_kernel.h deleted file mode 100644 index aeb1bdefe..000000000 --- a/src/infiniop/elementwise/bang/elementwise_bang_kernel.h +++ /dev/null @@ -1,335 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_BANG_KERNEL_MLU__ -#define __INFINIOP_ELEMENTWISE_BANG_KERNEL_MLU__ - -#include "../../devices/bang/bang_kernel_common.h" -#include "../../devices/bang/common_bang.h" - -using namespace device::bang::kernel; - -/** - * @brief Core elementwise operation implementation for BANG device. - * - * @tparam N Number of input tensors. - * @tparam Op Operator functor type. - * @tparam Tdata Data type for inputs and output. - * @tparam Args Additional arguments for operator. - * - * @param typed_inputs Array of typed input pointers. - * @param output Output tensor pointer. - * @param nram_buf NRAM buffer for temporary storage. - * @param input_indexes Precomputed input indexes. - * @param output_index Starting output index. - * @param num_elements Number of elements to process. - * @param output_contiguous Whether output is contiguous. - * @param input_contiguous Array indicating input contiguity. - * @param ndim Number of dimensions. - * @param input_shape Input shape in global memory. - * @param input_strides Input strides in global memory. - * @param output_shape Output shape in global memory. - * @param output_strides Output strides in global memory. - * @param indexer Input indexer helper. - * @param start_idx Starting index for this task. - * @param args Additional arguments for operator. - */ -template -__mlu_device__ void launchOp( - Tdata **typed_inputs, - Tdata *output, - Tdata *nram_buf, - size_t *input_indexes, - size_t output_index, - size_t num_elements, - bool output_contiguous, - const bool *input_contiguous, - const bool *input_broadcasted, - size_t ndim, - const size_t *input_shapes, - const ptrdiff_t *input_strides, - const size_t *output_shape, - const ptrdiff_t *output_strides, - InputIndexer indexer, - size_t start_idx, - Args... args) { - - static_assert(N == Op::num_inputs, "template N is not equal to Op::num_inputs!"); - - // NRAM memory planning - const size_t nram_usable = NRAM_MAX_SIZE - (ALIGN_SIZE * (N + 1)); - const size_t max_batch = nram_usable / ((N + 1) * sizeof(Tdata)); - - size_t processed = 0; - while (processed < num_elements) { - size_t curr_batch = std::min(max_batch, num_elements - processed); - - // Align memory address - Tdata *aligned_buf = reinterpret_cast( - (reinterpret_cast(nram_buf) + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1)); - - // 1. Copy input data to NRAM - Tdata *input_buffers[N]; - for (size_t i = 0; i < N; ++i) { - input_buffers[i] = aligned_buf + i * max_batch; - - if (input_contiguous[i]) { - // Contiguous case - bulk copy - __memcpy_async(input_buffers[i], - typed_inputs[i] + input_indexes[i] + processed, - curr_batch * sizeof(Tdata), - GDRAM2NRAM); - } else { - // Non-contiguous case - copy in contiguous chunks - nonContiguousMemcpy( - input_buffers[i], - typed_inputs[i], - GDRAM2NRAM, - indexer, - i, - processed, - curr_batch, - start_idx, - ndim, - input_shapes + i * ndim, - input_strides + i * ndim, - true); - } - } - __sync_io(); - - // 2. Execute operation - Tdata *output_buffer = aligned_buf + N * max_batch; - Op op; - if constexpr (N == 1) { - op(output_buffer, input_buffers[0], input_buffers[0], curr_batch, args...); - } else { - op(output_buffer, input_buffers[0], input_buffers[1], curr_batch, args...); - } - __sync_compute(); - - // 3. Write back results - if (output_contiguous) { - // Contiguous output - bulk copy - __memcpy_async(output + output_index + processed, - output_buffer, - curr_batch * sizeof(Tdata), - NRAM2GDRAM); - } else { - // Non-contiguous output - copy in contiguous chunks - nonContiguousMemcpy( - output, - output_buffer, - NRAM2GDRAM, - indexer, - 0, // unused for output - processed, - curr_batch, - start_idx, - ndim, - output_shape, - output_strides, - false); - } - - processed += curr_batch; - } -} - -/** - * @brief BANG kernel for elementwise operations. - * - * @tparam N Number of input tensors. - * @tparam Op Operator functor type. - * @tparam Tdata Data type for inputs and output. - * @tparam Args Additional arguments for operator. - * - * @param output_size Total output elements. - * @param ndim Number of dimensions. - * @param output_contiguous Whether output is contiguous. - * @param input_contiguous Input contiguity flags in global memory. - * @param input_broadcasted Input broadcast flags in global memory. - * @param output_shape Output shape in global memory. - * @param input_shapes Input shapes in global memory. - * @param output_strides Output strides in global memory. - * @param input_strides Input strides in global memory. - * @param output Output tensor pointer. - * @param inputs Array of input pointers. - * @param args Additional arguments for operator. - */ -template -__mlu_global__ void elementwiseKernel( - size_t output_size, - size_t ndim, - bool output_contiguous, - const bool *input_contiguous, - const bool *input_broadcasted, - const size_t *output_shape, - const size_t *input_shapes, - const ptrdiff_t *output_strides, - const ptrdiff_t *input_strides, - Tdata *output, - const void *const *inputs, - Args... args) { - - // Cast input pointers to the correct type - Tdata *typed_inputs[N]; - for (size_t i = 0; i < N; ++i) { - typed_inputs[i] = reinterpret_cast(const_cast(inputs[i])); - } - - // Calculate workload per task - size_t elements_per_task = (output_size + taskDim - 1) / taskDim; - size_t start_idx = taskId * elements_per_task; - size_t end_idx = std::min(start_idx + elements_per_task, output_size); - size_t num_elements = end_idx > start_idx ? end_idx - start_idx : 0; - - if (num_elements == 0) { - return; - } - - // Allocate NRAM buffer (shared by all inputs and output) - __nram__ Tdata nram_buf[NRAM_MAX_SIZE / sizeof(Tdata)]; - - // Get output index - size_t output_index = getOutputIndex(start_idx, output_contiguous, - ndim, output_shape, output_strides); - - // Create input indexer - InputIndexer indexer{ - static_cast(start_idx), - ndim, - input_contiguous, - input_broadcasted, - input_shapes, - input_strides, - output_strides}; - - // Get index offsets for each operand - size_t input_indexes[N]; - for (size_t i = 0; i < N; ++i) { - input_indexes[i] = indexer(i, 0); - } - - // Launch the operation with all required parameters - launchOp(typed_inputs, output, nram_buf, input_indexes, - output_index, num_elements, output_contiguous, - input_contiguous, input_broadcasted, ndim, - input_shapes, input_strides, output_shape, - output_strides, indexer, start_idx, args...); -} - -/** - * @brief Intermediate layer that determines optimal launch configuration before calling elementwiseKernel. - * - * @tparam N Number of input tensors. - * @tparam Op Operator functor type. - * @tparam Tdata Data type for inputs and output. - * @tparam Args Additional arguments for operator. - */ -template -void launchElementwiseKernelWrapper( - size_t output_size, - size_t ndim, - bool output_contiguous, - const bool *input_contiguous, - const bool *input_broadcasted, - const size_t *output_shape, - const size_t *input_shapes, - const ptrdiff_t *output_strides, - const ptrdiff_t *input_strides, - Tdata *output, - const void *const *inputs, - cnrtQueue_t queue, - const std::shared_ptr &internal, - Args... args) { - - // Get hardware information from internal handle - int core_per_cluster = internal->getCorePerCluster(); - int cluster_count = internal->getClusterCount(); - - // Set kernel launch dimensions - cnrtDim3_t dim; - dim.x = core_per_cluster; - dim.y = cluster_count; - dim.z = 1; - - // Choose kernel type based on problem characteristics - cnrtFunctionType_t func_type = cnrtFuncTypeBlock; - if (output_size > 1024 * 1024 && output_contiguous) { - // For large contiguous operations, use UNION type - func_type = cnrtFuncTypeUnion1; - } - - // Launch the kernel with optimal configuration - elementwiseKernel<<>>( - output_size, ndim, output_contiguous, - input_contiguous, input_broadcasted, - output_shape, input_shapes, - output_strides, input_strides, - output, inputs, args...); -} - -/** - * @brief Macro for implementing elementwise kernel launch. - * - * @param OpName Name of the operation. - * @param Op Operator functor type. - */ -#define LAUNCH_ELEMENTWISE_KERNEL_IMPL(OpName, Op) \ - template \ - void launch##OpName##Kernel( \ - size_t output_size, \ - size_t ndim, \ - bool output_contiguous, \ - const void *input_contiguous, \ - const void *input_broadcasted, \ - const void *output_shape, \ - const void *input_shapes, \ - const void *output_strides, \ - const void *input_strides, \ - void *output, \ - const void *const *inputs, \ - cnrtQueue_t queue, \ - const std::shared_ptr &internal, \ - Args... args) { \ - launchElementwiseKernelWrapper( \ - output_size, ndim, output_contiguous, \ - reinterpret_cast(input_contiguous), \ - reinterpret_cast(input_broadcasted), \ - reinterpret_cast(output_shape), \ - reinterpret_cast(input_shapes), \ - reinterpret_cast(output_strides), \ - reinterpret_cast(input_strides), \ - reinterpret_cast(output), inputs, queue, internal, args...); \ - } - -/** - * @brief Macro for instantiating elementwise kernel for specific types. - * - * @param OpName Name of the operation. - * @param T Data type. - * @param ... Additional template arguments. - */ -/** - * @brief Macro for instantiating elementwise kernel for specific types. - * - * @param OpName Name of the operation. - * @param T Data type. - * @param ... Additional template arguments. - */ -#define LAUNCH_ELEMENTWISE_KERNEL_INSTANTIATE(OpName, T, ...) \ - template void launch##OpName##Kernel( \ - size_t output_size, \ - size_t ndim, \ - bool output_contiguous, \ - const void *input_contiguous, \ - const void *input_broadcasted, \ - const void *output_shape, \ - const void *input_shapes, \ - const void *output_strides, \ - const void *input_strides, \ - void *output, \ - const void *const *inputs, \ - cnrtQueue_t queue, \ - const std::shared_ptr &internal, \ - ##__VA_ARGS__); - -#endif diff --git a/src/infiniop/elementwise/cpu/elementwise_cpu.h b/src/infiniop/elementwise/cpu/elementwise_cpu.h deleted file mode 100644 index 487cb5bdb..000000000 --- a/src/infiniop/elementwise/cpu/elementwise_cpu.h +++ /dev/null @@ -1,197 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_CPU_H__ -#define __INFINIOP_ELEMENTWISE_CPU_H__ - -#include "../../devices/cpu/common_cpu.h" -#include "../elementwise.h" -#include - -/** - * @brief Define the process for initializing a Descriptor of an elementwise operation - * for its CPU implementation - * - * @param HANDLE The device handle. - * @param DTYPE The output dtype. - * @param OUT_DESC The output tensor descriptor. - * @param INPUT_DESC_VEC A vector containing input tensor descriptors. - */ -#define CREATE_ELEMENTWISE_CPU_DESCRIPTOR(HANDLE, DTYPE, OUT_DESC, INPUT_DESC_VEC) \ - \ - auto info_result = op::elementwise::ElementwiseInfo::create(OUT_DESC, INPUT_DESC_VEC); \ - CHECK_RESULT(info_result); \ - \ - *desc_ptr = new Descriptor( \ - DTYPE, \ - info_result.take(), \ - nullptr, \ - 0, \ - HANDLE->device, \ - HANDLE->device_id); - -namespace op::elementwise::cpu { - -/** - * @brief CPU-specific device implementation for resource management and - * calculation implementations. - * - * This class encapsulates device-specific behavior and execution logic. - * Use the static create() method to instantiate a DeviceImpl. - */ -class DeviceImpl final { - struct Opaque; - std::shared_ptr _opaque; - - DeviceImpl(std::shared_ptr opaque) : _opaque(std::move(opaque)) {} - -public: - ~DeviceImpl() = default; - - template - static utils::Result create(Args &&...args); - - /** - * @brief Dispatches an elementwise operation with uniform input types. - * - * @tparam Op The elementwise operation to perform. - * @tparam Tdata The common data type of all inputs and output. - * @tparam Args Additional backend-specific arguments. - * @param info Precomputed tensor metadata (shapes, strides, etc.). - * @param output Pointer to the output tensor buffer. - * @param inputs Vector of input tensor data pointers. - * @param stream Device execution stream. - * @param args Additional backend-specific arguments. - * @return infiniStatus_t Status indicating success or failure. - */ - template - infiniStatus_t calculate( - const op::elementwise::ElementwiseInfo &info, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args); - - /** - * @brief Dispatches an elementwise operation with heterogeneous input types. - * - * Supports operations where each input may have a different type, as defined by Op. - * The number of input types must match the operation's expected input count. - * - * @tparam Op The elementwise operation to perform. - * @tparam Tout Output data type. - * @tparam Tin Variadic input data types. - * @tparam Args Additional backend-specific arguments. - * @param info Precomputed tensor metadata (shapes, strides, etc.). - * @param output Pointer to the output tensor buffer. - * @param inputs Vector of input tensor data pointers. - * @param stream Device execution stream. - * @param args Additional backend-specific arguments. - * @return infiniStatus_t Status indicating success or failure. - */ - template = 0> - infiniStatus_t calculate( - const op::elementwise::ElementwiseInfo &info, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args); -}; - -// Define the Opaque struct for CPU, which is empty -struct DeviceImpl::Opaque {}; - -template -utils::Result DeviceImpl::create(Args &&...args) { - return INFINI_STATUS_NOT_IMPLEMENTED; -} - -// Perform elementwise operation for different input types -template = 0> -void calculate_impl(const op::elementwise::ElementwiseInfo &info, - void *output, - const std::vector &inputs, - std::index_sequence, - Args &&...args) { - - Tout *out = reinterpret_cast(output); - std::tuple input_ptrs = {reinterpret_cast(inputs[Is])...}; - ptrdiff_t output_size = info.getOutputSize(); - -#pragma omp parallel for - for (ptrdiff_t i = 0; i < output_size; ++i) { - size_t out_idx = info.isOutputContiguous() - ? i - : op::common_cpu::indexToOffset(i, info.getNdim(), info.getOutputShape(), info.getOutputStrides()); - - auto get_input_idx = [&](size_t input_id) { - return info.getInputContiguous()[input_id] - ? i - : op::common_cpu::indexToOffset(i, info.getNdim(), info.getInputShape(input_id), info.getInputStrides(input_id)); - }; - - out[out_idx] = utils::cast( - Op{}.template operator()(std::get(input_ptrs)[get_input_idx(Is)]..., std::forward(args)...)); - } -} - -// Invoke elementwise operation for different input types -template > -infiniStatus_t DeviceImpl::calculate(const op::elementwise::ElementwiseInfo &info, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args) { - - static_assert(sizeof...(Tin) == Op::num_inputs, "Input type count mismatch"); - calculate_impl(info, output, inputs, std::make_index_sequence{}, std::forward(args)...); - return INFINI_STATUS_SUCCESS; -} - -// Perform elementwise operation when all inputs have the same type -template -void calculate_impl(const op::elementwise::ElementwiseInfo &info, - void *output, - const std::vector &inputs, - std::index_sequence, - Args &&...args) { - - Tdata *out = reinterpret_cast(output); - std::array ins = {reinterpret_cast(inputs[Is])...}; - const ptrdiff_t output_size = info.getOutputSize(); - -#pragma omp parallel for if (output_size > 1024) - for (ptrdiff_t i = 0; i < output_size; ++i) { - size_t out_idx = info.isOutputContiguous() - ? i - : op::common_cpu::indexToOffset(i, info.getNdim(), info.getOutputShape(), info.getOutputStrides()); - - auto get_input_idx = [&](size_t input_id) { - return info.getInputContiguous()[input_id] - ? i - : op::common_cpu::indexToOffset(i, info.getNdim(), info.getInputShape(input_id), info.getInputStrides(input_id)); - }; - - if constexpr (std::is_same_v || std::is_same_v) { - out[out_idx] = utils::cast(Op{}(utils::cast(ins[Is][get_input_idx(Is)])..., std::forward(args)...)); - } else { - out[out_idx] = Op{}(ins[Is][get_input_idx(Is)]..., std::forward(args)...); - } - } -} - -// Invoke elementwise operation when all inputs have the same type -template -infiniStatus_t DeviceImpl::calculate(const op::elementwise::ElementwiseInfo &info, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args) { - constexpr size_t N = Op::num_inputs; - calculate_impl(info, output, inputs, std::make_index_sequence{}, std::forward(args)...); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::elementwise::cpu - -#endif // __INFINIOP_ELEMENTWISE_CPU_H__ diff --git a/src/infiniop/elementwise/elementwise.h b/src/infiniop/elementwise/elementwise.h deleted file mode 100644 index a6a5477f4..000000000 --- a/src/infiniop/elementwise/elementwise.h +++ /dev/null @@ -1,206 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_H__ -#define __INFINIOP_ELEMENTWISE_H__ - -#include "../../utils.h" -#include "../operator.h" -#include "../tensor.h" -#include -#include -#include -#include -#include -#include -#include - -#define ELEMENTWISE_DESCRIPTOR(OP, NAMESPACE) \ - \ - namespace op::OP::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - infiniDtype_t _dtype; \ - op::elementwise::ElementwiseInfo _info; \ - std::unique_ptr _device_info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - infiniDtype_t dtype, \ - op::elementwise::ElementwiseInfo info, \ - op::elementwise::NAMESPACE::DeviceImpl *device_info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _dtype(dtype), \ - _info(std::move(info)), \ - _device_info(std::move(device_info)), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t output_desc, \ - std::vector input_descs); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *output, \ - std::vector inputs, \ - void *stream) const; \ - }; \ - } - -namespace op::elementwise { - -/** - * @brief Stores the metadata required for performing an elementwise operation. - * - * This struct encapsulates shape, stride, and layout information for both - * output and multiple input tensors involved in an elementwise operation. - * - * Memory is manually managed and freed in the destructor. - * Supports move construction but disallows copy construction and copy/move assignment. - * - * Use ElementwiseInfo::create(...) to safely construct an instance from tensor descriptors. - */ -struct ElementwiseInfo { -private: - std::vector _meta; - size_t _output_size; - size_t _input_size; - size_t _ndim; - bool _output_contiguous; - - ElementwiseInfo(std::vector meta, - size_t output_size, - size_t input_size, - size_t ndim, - bool output_contiguous) - : _meta(std::move(meta)), _output_size(output_size), - _input_size(input_size), _ndim(ndim), - _output_contiguous(output_contiguous) {} - -public: - // Get the Memory size of the meta data in bytes - inline size_t getMetaMemSize() const { - return _meta.size() * sizeof(size_t); - } - inline const int8_t *getMetaStart() const { - return reinterpret_cast(_meta.data()); - } - inline size_t getOutputSize() const { - return _output_size; - } - inline size_t getInputSize() const { - return _input_size; - } - inline size_t getNdim() const { - return _ndim; - } - inline bool isOutputContiguous() const { - return _output_contiguous; - } - inline const size_t *getOutputShape() const { - return reinterpret_cast(_meta.data()); - } - inline const ptrdiff_t *getOutputStrides() const { - return reinterpret_cast(getOutputShape() + _ndim); - } - inline const size_t *getAllInputShapes() const { - return reinterpret_cast(getOutputStrides() + _ndim); - } - inline const size_t *getInputShape(const size_t &index) const { - if (index < _input_size) { - return reinterpret_cast(getAllInputShapes() + index * _ndim); - } - return nullptr; - } - inline const ptrdiff_t *getAllInputStrides() const { - return reinterpret_cast(getAllInputShapes() + _input_size * _ndim); - } - inline const ptrdiff_t *getInputStrides(const size_t &index) const { - if (index < _input_size) { - return reinterpret_cast(getAllInputStrides() + index * _ndim); - } - return nullptr; - } - inline const bool *getInputContiguous() const { - return reinterpret_cast(getAllInputStrides() + _input_size * _ndim); - } - inline const bool *getInputBroadcasted() const { - return reinterpret_cast(getInputContiguous() + _input_size); - } - - using ResultType = utils::Result; - - /** - * @brief Construct ElementwiseInfo from output and input tensor descriptors. - * @param output_desc Descriptor of the output tensor. - * @param input_descs Descriptors of the input tensors. - * @return Result with the successfully constructed ElementwiseInfo, - * or the status code. - */ - static ResultType create( - infiniopTensorDescriptor_t output_desc, - std::vector input_descs) { - - if (!output_desc || input_descs.empty()) { - return INFINI_STATUS_BAD_PARAM; - } - - // Destination cannot have broadcast setup - if (output_desc->hasBroadcastDim()) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - auto input_size = input_descs.size(); - auto ndim = output_desc->ndim(); - auto output_size = output_desc->numel(); - auto output_contiguous = output_desc->isContiguous(); - - // Allocate memory for meta - auto shape_unit = output_desc->dim(0); - auto stride_unit = output_desc->stride(0); - size_t meta_mem_size = ndim * (sizeof(shape_unit) + sizeof(stride_unit)) - + input_size * ndim * sizeof(shape_unit) - + input_size * ndim * sizeof(stride_unit) - + 2 * input_size * sizeof(bool); - std::vector meta(CEIL_DIV(meta_mem_size, sizeof(size_t))); - int8_t *meta_ptr = reinterpret_cast(meta.data()); - - const auto output_shape = output_desc->shape(); - const auto output_strides = output_desc->strides(); - - // Pointers to the sections within _meta - size_t *output_shape_p = reinterpret_cast(meta_ptr); - ptrdiff_t *output_strides_p = reinterpret_cast(output_shape_p + ndim); - size_t *input_shapes = reinterpret_cast(output_strides_p + ndim); - ptrdiff_t *input_strides = reinterpret_cast(input_shapes + input_size * ndim); - bool *input_contiguous = reinterpret_cast(input_strides + input_size * ndim); - bool *input_broadcasted = input_contiguous + input_size; - - // Copy output shape and strides - std::memcpy(output_shape_p, output_shape.data(), ndim * sizeof(*output_shape_p)); - std::memcpy(output_strides_p, output_strides.data(), ndim * sizeof(*output_strides_p)); - - // Copy input shapes, strides, contiguous, and broadcasted flags - for (size_t i = 0; i < input_size; ++i) { - auto &desc = input_descs[i]; - const auto in_shape = desc->shape(); - const auto in_strides = desc->strides(); - std::memcpy(input_shapes + i * ndim, in_shape.data(), ndim * sizeof(*input_shapes)); - std::memcpy(input_strides + i * ndim, in_strides.data(), ndim * sizeof(*input_strides)); - input_contiguous[i] = desc->isContiguous(); - input_broadcasted[i] = !input_contiguous[i] && (desc->ndim() != ndim || desc->hasBroadcastDim()); - } - - ElementwiseInfo info(std::move(meta), output_size, input_size, ndim, output_contiguous); - return ResultType(std::move(info)); - } -}; -} // namespace op::elementwise - -#endif // __INFINIOP_ELEMENTWISE_H__ diff --git a/src/infiniop/elementwise/kunlun/elementwise_kunlun.h b/src/infiniop/elementwise/kunlun/elementwise_kunlun.h deleted file mode 100644 index a54534265..000000000 --- a/src/infiniop/elementwise/kunlun/elementwise_kunlun.h +++ /dev/null @@ -1,317 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_KUNLUN_XPU__ -#define __INFINIOP_ELEMENTWISE_KUNLUN_XPU__ - -#include "../../../utils.h" -#include "../../devices/kunlun/kunlun_common.h" -#include "../../devices/kunlun/kunlun_handle.h" -#include "../../devices/kunlun/kunlun_kernel_common.h" -#include "elementwise_kunlun_api.h" - -namespace op::elementwise::kunlun { - -using namespace device::kunlun::kernel; - -template -__device__ const T *typedInputPtr(const void *ptr) { - return reinterpret_cast(ptr); -} - -/** - * @brief Computes input tile offset - */ -struct InputIndexer { - int idx; - int ndim; - const bool *input_contiguous; - const bool *input_broadcasted; - const _size_t *input_shapes; - const _ptrdiff_t *input_strides; - const _ptrdiff_t *output_strides; - - inline __device__ int operator()(int input_id) const { - return input_contiguous[input_id] - ? idx - : indexToOffset(idx, ndim, input_shapes + input_id * ndim, input_strides + input_id * ndim); - } -}; - -/** - * @brief Computes the output index in memory, accounting for strides if non-contiguous. - * - * @param idx Linear index. - * @param is_contiguous Whether the output tensor is contiguous. - * @param ndim Number of dimensions. - * @param shape Shape of the output tensor. - * @param strides Strides of the output tensor. - * @return Memory offset index. - */ -inline __device__ int -getOutputIndex(int idx, - bool is_contiguous, - int ndim, - const _size_t *shape, - const _ptrdiff_t *strides) { - return is_contiguous ? idx : indexToOffset(idx, ndim, shape, strides); -} - -/** - * @brief Computes elements of input indexes - */ -template -__device__ void launchOp( - __global_ptr__ Tdata **typed_inputs, // gm pointer - __global_ptr__ Tdata *output, // gm pointer output - Tdata *inputs_buf, // local mem buffer - int *input_indexes, - int output_index, - Args... args) { - - static_assert(N == Op::num_inputs, "template N is not equal to Op::num_inputs!\n"); - -#pragma unroll - // Copy inputs to buf - for (int i = 0; i < N; i++) { - auto gm = typed_inputs[i] + input_indexes[i]; - auto lm = inputs_buf + i; - GM2LM_ASYNC(gm, lm, 1 * sizeof(Tdata)); - } - mfence(); - - // Calculate elementwise - // Inputs save all operands - Tdata out = Op{}(inputs_buf, args...); - - // Copy out to gm - LM2GM_ASYNC(&out, output + output_index, 1 * sizeof(Tdata)); - mfence(); -} - -template -__global__ void elementwiseKernel( - int output_size, - int ndim, - bool output_contiguous, - const bool *input_contiguous_gm, - const bool *input_broadcasted_gm, - const void *output_shape_gm, - const void *input_shapes_gm, - const void *output_strides_gm, - const void *input_strides_gm, - Tdata *output, - const void *const *inputs, - Args... args) { - - int cid = core_id(); - int ncores = core_num(); - if (cid >= ncores) { - return; - } - int thread_id = ncores * cluster_id() + cid; - int nthreads = ncores * cluster_num(); - - // Cast input gm pointer type - auto typed_inputs = reinterpret_cast(inputs); - - const int BUFF_SIZE = 64; - // Input data cache - __local__ Tdata inputs_buf[N]; - // Input contiguous/broadcasted flags - __local__ bool input_contiguous[N]; - __local__ bool input_broadcasted[N]; - // Input shape/strides - __local__ _size_t input_shapes[N * ndim]; - __local__ _ptrdiff_t input_strides[N * ndim]; - // Output shape/strides - __local__ _size_t output_shape[ndim]; - __local__ _ptrdiff_t output_strides[ndim]; - // Inputs gm ptr buf - __local__ __global_ptr__ Tdata *typed_inputs_ptr[N]; - - // Load from gm - GM2LM_ASYNC(input_contiguous_gm, input_contiguous, N * sizeof(bool)); - GM2LM_ASYNC(input_broadcasted_gm, input_broadcasted, N * sizeof(bool)); - GM2LM_ASYNC(input_shapes_gm, input_shapes, N * ndim * sizeof(_size_t)); - GM2LM_ASYNC(input_strides_gm, input_strides, N * ndim * sizeof(_ptrdiff_t)); - GM2LM_ASYNC(output_shape_gm, output_shape, ndim * sizeof(_size_t)); - GM2LM_ASYNC(output_strides_gm, output_strides, ndim * sizeof(_ptrdiff_t)); - GM2LM_ASYNC(typed_inputs, typed_inputs_ptr, N * sizeof(__global_ptr__ Tdata *)); - mfence(); - - int len_per_loop = min(BUFF_SIZE, roundup_div(output_size, nthreads)); - - for (int start = thread_id * len_per_loop; start < output_size; start += nthreads * len_per_loop) { - int read_len = min(len_per_loop, output_size - start); - for (int idx = start; idx < start + read_len; ++idx) { - int out_idx = getOutputIndex(idx, output_contiguous, - ndim, output_shape, output_strides); - InputIndexer indexer{idx, ndim, input_contiguous, input_broadcasted, - input_shapes, input_strides, output_strides}; - // Get index offset for every operand - int indexes[N]; - for (int i = 0; i < N; i++) { - indexes[i] = indexer(i); - } - // Launch operater - launchOp(&typed_inputs_ptr[0], output, inputs_buf, indexes, out_idx, args...); - } - } - sync_cluster(); -} - -struct DeviceImpl::Opaque { - std::shared_ptr internal; - - Opaque(const std::shared_ptr &internal_) - : internal(internal_) {} - - template - infiniStatus_t calculateImpl(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - kunlunStream_t stream, - Args &&...args) { - return launchElementwiseKernel( - info, - workspace, - reinterpret_cast(output), - inputs, - elementwiseKernel, - stream, - std::forward(args)...); - } - -private: - /** - * @brief Transfers elementwise operation metadata and input pointers from host to device memory. - * - * @tparam N Number of input tensors. - * - * @param info Elementwise operation metadata (shapes, strides, flags, etc.). - * @param workspace Pointer to device workspace memory for storing metadata and input pointers. - * @param h_inputs_arr Host array of input tensor pointers. - * @param d_inputs_arr Input reference to device array of input tensor pointers. - * @param d_input_contiguous Input reference to device array indicating whether each input is contiguous. - * @param d_input_broadcasted Input reference to device array indicating whether each input is broadcasted. - * @param d_output_shape Output reference to device array holding the output tensor shape. - * @param d_output_strides Output reference to device array holding output tensor strides. - * @param d_input_shapes Output reference to flattened input tensor shapes (N * ndim). - * @param d_input_strides Output reference to flattened input tensor strides (N * ndim). - * @param stream KUNLUN stream used for asynchronous memory transfer. - * @return infiniStatus_t Status indicating success or failure of the memory transfer and setup. - */ - template - infiniStatus_t infoToDevice( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - const void *const *h_inputs_arr, - __global_ptr__ const void **&d_inputs_arr, - __global_ptr__ const bool *&d_input_contiguous, - __global_ptr__ const bool *&d_input_broadcasted, - __global_ptr__ const size_t *&d_output_shape, - __global_ptr__ const ptrdiff_t *&d_output_strides, - __global_ptr__ const size_t *&d_input_shapes, - __global_ptr__ const ptrdiff_t *&d_input_strides, - kunlunStream_t stream) const { - - constexpr auto input_size = N; - const auto ndim = info.getNdim(); - constexpr auto input_arr_size = N * sizeof(*h_inputs_arr); - auto info_meta_start = info.getMetaStart(); // host meta pointer - - auto d_meta_start = reinterpret_cast<__global_ptr__ int8_t *>(workspace) - + input_arr_size; // device meta pointer - - // copy the input pointer array and meta to device - CHECK_KUNLUN(xpu_memcpy_async(workspace, h_inputs_arr, input_arr_size, XPU_HOST_TO_DEVICE, stream)); - CHECK_KUNLUN(xpu_memcpy_async((void *)d_meta_start, info_meta_start, info.getMetaMemSize(), XPU_HOST_TO_DEVICE, stream)); - - // offset/assign the pointers - d_inputs_arr = reinterpret_cast<__global_ptr__ const void **>(workspace); - d_output_shape = reinterpret_cast<__global_ptr__ const size_t *>(d_meta_start); - d_output_strides = reinterpret_cast<__global_ptr__ const ptrdiff_t *>(d_output_shape + ndim); - d_input_shapes = reinterpret_cast<__global_ptr__ const size_t *>(d_output_strides + ndim); - d_input_strides = reinterpret_cast<__global_ptr__ const ptrdiff_t *>(d_input_shapes + input_size * ndim); - d_input_contiguous = reinterpret_cast(d_input_strides + input_size * ndim); - d_input_broadcasted = reinterpret_cast(d_input_contiguous + input_size); - - // contiguous / broadcast 信息 - const bool *contiguous = info.getInputContiguous(); - const bool *broadcasted = info.getInputBroadcasted(); - - return INFINI_STATUS_SUCCESS; - } - - /** - * @brief Launch elementwise kernel - */ - template - infiniStatus_t launchElementwiseKernel( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - Tout *output, - const std::vector &inputs, - KernelFunc kernel_func, - kunlunStream_t stream, - Args &&...args) { - - auto output_size = info.getOutputSize(); - if (output_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - // Device pointers - __global_ptr__ const void **d_inputs_arr = nullptr; - __global_ptr__ const bool *d_input_contiguous = nullptr; - __global_ptr__ const bool *d_input_broadcasted = nullptr; - __global_ptr__ const size_t *d_output_shape = nullptr; - __global_ptr__ const ptrdiff_t *d_output_strides = nullptr; - __global_ptr__ const size_t *d_input_shapes = nullptr; - __global_ptr__ const ptrdiff_t *d_input_strides = nullptr; - - CHECK_STATUS(infoToDevice(info, workspace, inputs.data(), d_inputs_arr, - d_input_contiguous, d_input_broadcasted, - d_output_shape, d_output_strides, - d_input_shapes, d_input_strides, stream)); - - kernel_func<<>>( - output_size, - info.getNdim(), - info.isOutputContiguous(), - d_input_contiguous, - d_input_broadcasted, - reinterpret_cast<__global_ptr__ const void *>(d_output_shape), - reinterpret_cast<__global_ptr__ const void *>(d_input_shapes), - reinterpret_cast<__global_ptr__ const void *>(d_output_strides), - reinterpret_cast<__global_ptr__ const void *>(d_input_strides), - output, - reinterpret_cast<__global_ptr__ const void **>(d_inputs_arr), - args...); - - return INFINI_STATUS_SUCCESS; - } -}; - -template -utils::Result DeviceImpl::create(Args &&...args) { - auto opaque = std::make_shared(std::forward(args)...); - return utils::Result(new DeviceImpl(opaque)); -} - -template -infiniStatus_t DeviceImpl::calculate(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args) { - constexpr int N = Op::num_inputs; - return _opaque->calculateImpl( - info, workspace, output, inputs, - reinterpret_cast(stream), - std::forward(args)...); -} - -} // namespace op::elementwise::kunlun - -#endif diff --git a/src/infiniop/elementwise/kunlun/elementwise_kunlun_api.h b/src/infiniop/elementwise/kunlun/elementwise_kunlun_api.h deleted file mode 100644 index cd4f8032e..000000000 --- a/src/infiniop/elementwise/kunlun/elementwise_kunlun_api.h +++ /dev/null @@ -1,67 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_KUNLUN_API_H__ -#define __INFINIOP_ELEMENTWISE_KUNLUN_API_H__ - -#include "../elementwise.h" - -namespace op::elementwise::kunlun { - -class DeviceImpl final { - struct Opaque; - std::shared_ptr _opaque; - - DeviceImpl(std::shared_ptr opaque) : _opaque(std::move(opaque)) {} - -public: - ~DeviceImpl() = default; - - template - static utils::Result create(Args &&...args); - - /** - * @brief Launches elementwise operation. Operands have same types. - */ - template - infiniStatus_t calculate( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args); - - // /** - // * @brief Launches elementwise operation where operands' types differ - // */ - // template = 0> - // infiniStatus_t calculate( - // const op::elementwise::ElementwiseInfo &info, - // void *workspace, - // void *output, - // const std::vector &inputs, - // void *stream, - // Args &&...args); -}; - -} // namespace op::elementwise::kunlun - -#define CREATE_ELEMENTWISE_KUNLUN_DESCRIPTOR(HANDLE, DTYPE, OUT_DESC, INPUT_DESC_VEC) \ - \ - auto info_result = op::elementwise::ElementwiseInfo::create(OUT_DESC, INPUT_DESC_VEC); \ - CHECK_RESULT(info_result); \ - auto info = info_result.take(); \ - auto workspace_size = info.getMetaMemSize() + info.getInputSize() * sizeof(void *); \ - \ - auto device_impl_result = op::elementwise::kunlun::DeviceImpl::create(HANDLE->internal()); \ - CHECK_RESULT(device_impl_result); \ - \ - *desc_ptr = new Descriptor( \ - DTYPE, \ - std::move(info), \ - std::move(device_impl_result.take()), \ - workspace_size, \ - HANDLE->device, \ - HANDLE->device_id); - -#endif diff --git a/src/infiniop/elementwise/metax/elementwise_metax.h b/src/infiniop/elementwise/metax/elementwise_metax.h deleted file mode 100644 index 084677ea7..000000000 --- a/src/infiniop/elementwise/metax/elementwise_metax.h +++ /dev/null @@ -1,262 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_METAX_H__ -#define __INFINIOP_ELEMENTWISE_METAX_H__ - -#include "../../../utils.h" -#include "../../devices/metax/metax_common.h" -#include "../../devices/metax/metax_kernel_common.h" -#include "elementwise_metax_api.h" - -namespace op::elementwise::metax { -template -__device__ __forceinline__ const T *typedInputPtr(const void *ptr) { - return reinterpret_cast(ptr); -} - -__device__ __forceinline__ size_t getOutputIndex(size_t idx, bool is_contiguous, size_t ndim, - const size_t *shape, const ptrdiff_t *strides) { - return is_contiguous ? idx : device::metax::indexToOffset(idx, ndim, shape, strides); -} - -struct InputIndexer { - size_t idx; - size_t ndim; - const bool *input_contiguous; - const bool *input_broadcasted; - const size_t *input_shapes; - const ptrdiff_t *input_strides; - const ptrdiff_t *output_strides; - - __device__ __forceinline__ size_t operator()(size_t input_id) const { - return input_contiguous[input_id] - ? idx - : device::metax::indexToOffset(idx, ndim, input_shapes + input_id * ndim, input_strides + input_id * ndim); - } -}; - -template -__device__ __forceinline__ void unpackInputsAndApply(F &&f, std::index_sequence) { - f(std::integral_constant{}...); -} - -template -INFINIOP_METAX_KERNEL elementwiseKernel( - size_t output_size, - size_t ndim, - bool output_contiguous, - const bool *__restrict__ input_contiguous, - const bool *__restrict__ input_broadcasted, - const size_t *__restrict__ output_shape, - const size_t *__restrict__ input_shapes, - const ptrdiff_t *__restrict__ output_strides, - const ptrdiff_t *__restrict__ input_strides, - Tdata *output, - const void *const *inputs, - size_t offset, - Args... args) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; - - if (idx < output_size) { - const Tdata *const *typed_inputs = reinterpret_cast(inputs); - size_t out_idx = getOutputIndex(idx, output_contiguous, ndim, output_shape, output_strides); - InputIndexer indexer{idx, ndim, input_contiguous, input_broadcasted, input_shapes, input_strides, output_strides}; - - unpackInputsAndApply( - [&](auto... Is) { - output[out_idx] = Op{}(typed_inputs[Is.value][indexer(Is.value)]..., std::forward(args)...); - }, - std::make_index_sequence{}); - } -} - -template -INFINIOP_METAX_KERNEL elementwiseKernel( - size_t output_size, - size_t ndim, - bool output_contiguous, - const bool *__restrict__ input_contiguous, - const bool *__restrict__ input_broadcasted, - const size_t *__restrict__ output_shape, - const size_t *__restrict__ input_shapes, - const ptrdiff_t *__restrict__ output_strides, - const ptrdiff_t *__restrict__ input_strides, - Tout *output, - const void *const *__restrict__ inputs, - size_t offset) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; - - if (idx < output_size) { - size_t out_idx = getOutputIndex(idx, output_contiguous, ndim, output_shape, output_strides); - InputIndexer indexer{idx, ndim, input_contiguous, input_broadcasted, input_shapes, input_strides, output_strides}; - - unpackInputsAndApply( - [&](auto... Is) { - output[out_idx] = Op{}.template operator()( - (typedInputPtr(inputs[Is.value])[indexer(Is.value)])...); - }, - std::index_sequence_for{}); - } -} - -struct DeviceImpl::Opaque { - std::shared_ptr internal; - - Opaque(const std::shared_ptr &internal) - : internal(internal) {} - - template - infiniStatus_t calculateImpl(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - hcStream_t stream, - Args &&...args) { - return launchElementwiseKernel( - info, workspace, - reinterpret_cast(output), inputs, - elementwiseKernel, - stream, - std::forward(args)...); - } - - template = 0> - infiniStatus_t calculateImpl(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - hcStream_t stream, - Args &&...args) { - return launchElementwiseKernel( - info, workspace, - reinterpret_cast(output), inputs, - elementwiseKernel, - stream); - } - -private: - template - infiniStatus_t infoToDevice( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - const void *const *h_inputs_arr, - const void **&d_inputs_arr, - const bool *&d_input_contiguous, - const bool *&d_input_broadcasted, - const size_t *&d_output_shape, - const ptrdiff_t *&d_output_strides, - const size_t *&d_input_shapes, - const ptrdiff_t *&d_input_strides, - hcStream_t stream) const { - - constexpr auto input_size = N; - const auto ndim = info.getNdim(); - constexpr auto input_arr_size = N * sizeof(*h_inputs_arr); - const int8_t *info_meta_start = info.getMetaStart(); - const int8_t *d_meta_start = reinterpret_cast(workspace) + input_arr_size; - - // copy the input pointer array and meta to device - CHECK_METAX(hcMemcpyAsync(workspace, h_inputs_arr, input_arr_size, hcMemcpyHostToDevice, stream)); - CHECK_METAX(hcMemcpyAsync((void *)d_meta_start, info_meta_start, info.getMetaMemSize(), hcMemcpyHostToDevice, stream)); - - // offset/assign the pointers - d_inputs_arr = reinterpret_cast(workspace); - d_output_shape = reinterpret_cast(d_meta_start); - d_output_strides = reinterpret_cast(d_output_shape + ndim); - d_input_shapes = reinterpret_cast(d_output_strides + ndim); - d_input_strides = reinterpret_cast(d_input_shapes + input_size * ndim); - d_input_contiguous = reinterpret_cast(d_input_strides + input_size * ndim); - d_input_broadcasted = reinterpret_cast(d_input_contiguous + input_size); - - return INFINI_STATUS_SUCCESS; - } - - template - infiniStatus_t launchElementwiseKernel( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - Tout *output, - const std::vector &inputs, - KernelFunc kernel_func, - hcStream_t stream, - Args &&...args) { - - auto output_size = info.getOutputSize(); - if (output_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - // Device pointers - const void **d_inputs_arr = nullptr; - const bool *d_input_contiguous = nullptr; - const bool *d_input_broadcasted = nullptr; - const size_t *d_output_shape = nullptr; - const ptrdiff_t *d_output_strides = nullptr; - const size_t *d_input_shapes = nullptr; - const ptrdiff_t *d_input_strides = nullptr; - - CHECK_STATUS(infoToDevice(info, workspace, inputs.data(), d_inputs_arr, - d_input_contiguous, d_input_broadcasted, - d_output_shape, d_output_strides, - d_input_shapes, d_input_strides, stream)); - - dim3 blockDims(std::min(BLOCK_SIZE, static_cast(internal->maxThreadsPerBlock()))); - dim3 gridDims(std::min(uint32_t(CEIL_DIV(output_size, blockDims.x)), static_cast(internal->gridSizeX()))); - size_t step = gridDims.x * blockDims.x; - - for (size_t i = 0; i < output_size; i += step) { - kernel_func<<>>( - output_size, info.getNdim(), info.isOutputContiguous(), - d_input_contiguous, d_input_broadcasted, - d_output_shape, d_input_shapes, - d_output_strides, d_input_strides, - output, reinterpret_cast(d_inputs_arr), - i, std::forward(args)...); - } - - return INFINI_STATUS_SUCCESS; - } -}; - -template -utils::Result DeviceImpl::create(Args &&...args) { - auto opaque = std::make_shared(std::forward(args)...); - return utils::Result(new DeviceImpl(opaque)); -} - -/* Invoke elementwise operation for different input types */ -template > -infiniStatus_t DeviceImpl::calculate(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args) { - constexpr size_t N = Op::num_inputs; - static_assert(sizeof...(Tin) == N, "Input type count mismatch"); - return _opaque->calculateImpl( - info, workspace, output, inputs, - reinterpret_cast(stream), - std::forward(args)...); -} - -/* Invoke elementwise operation when all inputs have the same dtype */ -template -infiniStatus_t DeviceImpl::calculate(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args) { - constexpr size_t N = Op::num_inputs; - return _opaque->calculateImpl( - info, workspace, output, inputs, - reinterpret_cast(stream), - std::forward(args)...); -} - -} // namespace op::elementwise::metax - -#endif diff --git a/src/infiniop/elementwise/metax/elementwise_metax_api.h b/src/infiniop/elementwise/metax/elementwise_metax_api.h deleted file mode 100644 index b59c14da5..000000000 --- a/src/infiniop/elementwise/metax/elementwise_metax_api.h +++ /dev/null @@ -1,59 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_METAX_API_H__ -#define __INFINIOP_ELEMENTWISE_METAX_API_H__ - -#include "../elementwise.h" - -namespace op::elementwise::metax { - -class DeviceImpl final { - struct Opaque; - std::shared_ptr _opaque; - - DeviceImpl(std::shared_ptr opaque) : _opaque(std::move(opaque)) {} - -public: - ~DeviceImpl() = default; - - template - static utils::Result create(Args &&...args); - - template - infiniStatus_t calculate( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args); - - template = 0> - infiniStatus_t calculate( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args); -}; -} // namespace op::elementwise::metax -#define CREATE_ELEMENTWISE_METAX_DESCRIPTOR(HANDLE, DTYPE, OUT_DESC, INPUT_DESC_VEC) \ - \ - auto info_result = op::elementwise::ElementwiseInfo::create(OUT_DESC, INPUT_DESC_VEC); \ - CHECK_RESULT(info_result); \ - auto info = info_result.take(); \ - auto workspace_size = info.getMetaMemSize() + info.getInputSize() * sizeof(void *); \ - \ - auto device_impl_result = op::elementwise::metax::DeviceImpl::create(HANDLE->internal()); \ - CHECK_RESULT(device_impl_result); \ - \ - *desc_ptr = new Descriptor( \ - DTYPE, \ - std::move(info), \ - std::move(device_impl_result.take()), \ - workspace_size, \ - HANDLE->device, \ - HANDLE->device_id); - -#endif // __INFINIOP_ELEMENTWISE_METAX_API_H__ diff --git a/src/infiniop/elementwise/moore/elementwise_moore.h b/src/infiniop/elementwise/moore/elementwise_moore.h deleted file mode 100644 index 088f76b6a..000000000 --- a/src/infiniop/elementwise/moore/elementwise_moore.h +++ /dev/null @@ -1,262 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_MOORE_H__ -#define __INFINIOP_ELEMENTWISE_MOORE_H__ - -#include "../../../utils.h" -#include "../../devices/moore/moore_common.h" -#include "../../devices/moore/moore_kernel_common.h" -#include "elementwise_moore_api.h" - -namespace op::elementwise::moore { -template -__device__ __forceinline__ const T *typedInputPtr(const void *ptr) { - return reinterpret_cast(ptr); -} - -__device__ __forceinline__ size_t getOutputIndex(size_t idx, bool is_contiguous, size_t ndim, - const size_t *shape, const ptrdiff_t *strides) { - return is_contiguous ? idx : device::moore::indexToOffset(idx, ndim, shape, strides); -} - -struct InputIndexer { - size_t idx; - size_t ndim; - const bool *input_contiguous; - const bool *input_broadcasted; - const size_t *input_shapes; - const ptrdiff_t *input_strides; - const ptrdiff_t *output_strides; - - __device__ __forceinline__ size_t operator()(size_t input_id) const { - return input_contiguous[input_id] - ? idx - : device::moore::indexToOffset(idx, ndim, input_shapes + input_id * ndim, input_strides + input_id * ndim); - } -}; - -template -__device__ __forceinline__ void unpackInputsAndApply(F &&f, std::index_sequence) { - f(std::integral_constant{}...); -} - -template -INFINIOP_MOORE_KERNEL elementwiseKernel( - size_t output_size, - size_t ndim, - bool output_contiguous, - const bool *__restrict__ input_contiguous, - const bool *__restrict__ input_broadcasted, - const size_t *__restrict__ output_shape, - const size_t *__restrict__ input_shapes, - const ptrdiff_t *__restrict__ output_strides, - const ptrdiff_t *__restrict__ input_strides, - Tdata *output, - const void *const *inputs, - size_t offset, - Args... args) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; - - if (idx < output_size) { - const Tdata *const *typed_inputs = reinterpret_cast(inputs); - size_t out_idx = getOutputIndex(idx, output_contiguous, ndim, output_shape, output_strides); - InputIndexer indexer{idx, ndim, input_contiguous, input_broadcasted, input_shapes, input_strides, output_strides}; - - unpackInputsAndApply( - [&](auto... Is) { - output[out_idx] = Op{}(typed_inputs[Is.value][indexer(Is.value)]..., std::forward(args)...); - }, - std::make_index_sequence{}); - } -} - -template -INFINIOP_MOORE_KERNEL elementwiseKernel( - size_t output_size, - size_t ndim, - bool output_contiguous, - const bool *__restrict__ input_contiguous, - const bool *__restrict__ input_broadcasted, - const size_t *__restrict__ output_shape, - const size_t *__restrict__ input_shapes, - const ptrdiff_t *__restrict__ output_strides, - const ptrdiff_t *__restrict__ input_strides, - Tout *output, - const void *const *__restrict__ inputs, - size_t offset) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; - - if (idx < output_size) { - size_t out_idx = getOutputIndex(idx, output_contiguous, ndim, output_shape, output_strides); - InputIndexer indexer{idx, ndim, input_contiguous, input_broadcasted, input_shapes, input_strides, output_strides}; - - unpackInputsAndApply( - [&](auto... Is) { - output[out_idx] = Op{}.template operator()( - (typedInputPtr(inputs[Is.value])[indexer(Is.value)])...); - }, - std::index_sequence_for{}); - } -} - -struct DeviceImpl::Opaque { - std::shared_ptr internal; - - Opaque(const std::shared_ptr &internal) - : internal(internal) {} - - template - infiniStatus_t calculateImpl(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - musaStream_t stream, - Args &&...args) { - return launchElementwiseKernel( - info, workspace, - reinterpret_cast(output), inputs, - elementwiseKernel, - stream, - std::forward(args)...); - } - - template = 0> - infiniStatus_t calculateImpl(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - musaStream_t stream, - Args &&...args) { - return launchElementwiseKernel( - info, workspace, - reinterpret_cast(output), inputs, - elementwiseKernel, - stream); - } - -private: - template - infiniStatus_t infoToDevice( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - const void *const *h_inputs_arr, - const void **&d_inputs_arr, - const bool *&d_input_contiguous, - const bool *&d_input_broadcasted, - const size_t *&d_output_shape, - const ptrdiff_t *&d_output_strides, - const size_t *&d_input_shapes, - const ptrdiff_t *&d_input_strides, - musaStream_t stream) const { - - constexpr auto input_size = N; - const auto ndim = info.getNdim(); - constexpr auto input_arr_size = N * sizeof(*h_inputs_arr); - const int8_t *info_meta_start = info.getMetaStart(); - const int8_t *d_meta_start = reinterpret_cast(workspace) + input_arr_size; - - // copy the input pointer array and meta to device - CHECK_MOORE(musaMemcpyAsync(workspace, h_inputs_arr, input_arr_size, musaMemcpyHostToDevice, stream)); - CHECK_MOORE(musaMemcpyAsync((void *)d_meta_start, info_meta_start, info.getMetaMemSize(), musaMemcpyHostToDevice, stream)); - - // offset/assign the pointers - d_inputs_arr = reinterpret_cast(workspace); - d_output_shape = reinterpret_cast(d_meta_start); - d_output_strides = reinterpret_cast(d_output_shape + ndim); - d_input_shapes = reinterpret_cast(d_output_strides + ndim); - d_input_strides = reinterpret_cast(d_input_shapes + input_size * ndim); - d_input_contiguous = reinterpret_cast(d_input_strides + input_size * ndim); - d_input_broadcasted = reinterpret_cast(d_input_contiguous + input_size); - - return INFINI_STATUS_SUCCESS; - } - - template - infiniStatus_t launchElementwiseKernel( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - Tout *output, - const std::vector &inputs, - KernelFunc kernel_func, - musaStream_t stream, - Args &&...args) { - - auto output_size = info.getOutputSize(); - if (output_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - // Device pointers - const void **d_inputs_arr = nullptr; - const bool *d_input_contiguous = nullptr; - const bool *d_input_broadcasted = nullptr; - const size_t *d_output_shape = nullptr; - const ptrdiff_t *d_output_strides = nullptr; - const size_t *d_input_shapes = nullptr; - const ptrdiff_t *d_input_strides = nullptr; - - CHECK_STATUS(infoToDevice(info, workspace, inputs.data(), d_inputs_arr, - d_input_contiguous, d_input_broadcasted, - d_output_shape, d_output_strides, - d_input_shapes, d_input_strides, stream)); - - dim3 blockDims(std::min(BLOCK_SIZE, static_cast(internal->maxThreadsPerBlock()))); - dim3 gridDims(std::min(uint32_t(CEIL_DIV(output_size, blockDims.x)), static_cast(internal->gridSizeX()))); - size_t step = gridDims.x * blockDims.x; - - for (size_t i = 0; i < output_size; i += step) { - kernel_func<<>>( - output_size, info.getNdim(), info.isOutputContiguous(), - d_input_contiguous, d_input_broadcasted, - d_output_shape, d_input_shapes, - d_output_strides, d_input_strides, - output, reinterpret_cast(d_inputs_arr), - i, std::forward(args)...); - } - - return INFINI_STATUS_SUCCESS; - } -}; - -template -utils::Result DeviceImpl::create(Args &&...args) { - auto opaque = std::make_shared(std::forward(args)...); - return utils::Result(new DeviceImpl(opaque)); -} - -/* Invoke elementwise operation for different input types */ -template > -infiniStatus_t DeviceImpl::calculate(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args) { - constexpr size_t N = Op::num_inputs; - static_assert(sizeof...(Tin) == N, "Input type count mismatch"); - return _opaque->calculateImpl( - info, workspace, output, inputs, - reinterpret_cast(stream), - std::forward(args)...); -} - -/* Invoke elementwise operation when all inputs have the same dtype */ -template -infiniStatus_t DeviceImpl::calculate(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args) { - constexpr size_t N = Op::num_inputs; - return _opaque->calculateImpl( - info, workspace, output, inputs, - reinterpret_cast(stream), - std::forward(args)...); -} - -} // namespace op::elementwise::moore - -#endif diff --git a/src/infiniop/elementwise/moore/elementwise_moore_api.h b/src/infiniop/elementwise/moore/elementwise_moore_api.h deleted file mode 100644 index a62ae4f49..000000000 --- a/src/infiniop/elementwise/moore/elementwise_moore_api.h +++ /dev/null @@ -1,59 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_MOORE_API_H__ -#define __INFINIOP_ELEMENTWISE_MOORE_API_H__ - -#include "../elementwise.h" - -namespace op::elementwise::moore { - -class DeviceImpl final { - struct Opaque; - std::shared_ptr _opaque; - - DeviceImpl(std::shared_ptr opaque) : _opaque(std::move(opaque)) {} - -public: - ~DeviceImpl() = default; - - template - static utils::Result create(Args &&...args); - - template - infiniStatus_t calculate( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args); - - template = 0> - infiniStatus_t calculate( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args); -}; -} // namespace op::elementwise::moore -#define CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(HANDLE, DTYPE, OUT_DESC, INPUT_DESC_VEC) \ - \ - auto info_result = op::elementwise::ElementwiseInfo::create(OUT_DESC, INPUT_DESC_VEC); \ - CHECK_RESULT(info_result); \ - auto info = info_result.take(); \ - auto workspace_size = info.getMetaMemSize() + info.getInputSize() * sizeof(void *); \ - \ - auto device_impl_result = op::elementwise::moore::DeviceImpl::create(HANDLE->internal()); \ - CHECK_RESULT(device_impl_result); \ - \ - *desc_ptr = new Descriptor( \ - DTYPE, \ - std::move(info), \ - std::move(device_impl_result.take()), \ - workspace_size, \ - HANDLE->device, \ - HANDLE->device_id); - -#endif // __INFINIOP_ELEMENTWISE_MOORE_API_H__ diff --git a/src/infiniop/elementwise/nvidia/elementwise_nvidia.cuh b/src/infiniop/elementwise/nvidia/elementwise_nvidia.cuh deleted file mode 100644 index adebf758f..000000000 --- a/src/infiniop/elementwise/nvidia/elementwise_nvidia.cuh +++ /dev/null @@ -1,421 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_CUDA_H__ -#define __INFINIOP_ELEMENTWISE_CUDA_H__ - -#include "../../../utils.h" -#include "../../devices/nvidia/nvidia_common.cuh" -#include "../../devices/nvidia/nvidia_kernel_common.cuh" -#include "elementwise_nvidia_api.cuh" - -namespace op::elementwise::nvidia { - -/** - * @brief Casts an untyped device pointer to a typed pointer of type T. - * - * @tparam T Desired pointer type. - * - * @param ptr Untyped pointer. - * @return Pointer of type const T*. - */ -template -__device__ __forceinline__ const T *typedInputPtr(const void *ptr) { - return reinterpret_cast(ptr); -} - -/** - * @brief Computes the output index in memory, accounting for strides if non-contiguous. - * - * @param idx Linear index. - * @param is_contiguous Whether the output tensor is contiguous. - * @param ndim Number of dimensions. - * @param shape Shape of the output tensor. - * @param strides Strides of the output tensor. - * @return Memory offset index. - */ -__device__ __forceinline__ size_t getOutputIndex(size_t idx, bool is_contiguous, size_t ndim, - const size_t *shape, const ptrdiff_t *strides) { - return is_contiguous ? idx : device::nvidia::indexToOffset(idx, ndim, shape, strides); -} - -/** - * @brief Computes input element offset for broadcasting and strided access. - * - * Used to map a linear output index to the corresponding index in an input tensor, - * considering contiguity and broadcasting. - */ -struct InputIndexer { - size_t idx; - size_t ndim; - const bool *input_contiguous; - const bool *input_broadcasted; - const size_t *input_shapes; - const ptrdiff_t *input_strides; - const ptrdiff_t *output_strides; - - /** - * @brief Computes the memory offset for a given input tensor at current index. - * - * @param input_id ID of the input tensor. - * @return Offset into the input tensor. - */ - __device__ __forceinline__ size_t operator()(size_t input_id) const { - return input_contiguous[input_id] - ? idx - : device::nvidia::indexToOffset(idx, ndim, input_shapes + input_id * ndim, input_strides + input_id * ndim); - } -}; - -/** - * @brief Invokes a callable with compile-time index constants. - * - * Used to unpack index sequence for variadic template processing of inputs. - * - * @tparam F Callable type. - * @tparam Is Compile-time index sequence. - * - * @param f Callable to invoke with index constants. - */ -template -__device__ __forceinline__ void unpackInputsAndApply(F &&f, std::index_sequence) { - f(std::integral_constant{}...); -} - -/** - * @brief CUDA kernel for performing elementwise operations on tensors where all inputs share the same data type. - * - * @tparam N Number of input tensors. - * @tparam Op Operator type implementing operator()(Tdata...). - * @tparam Tdata Common data type for inputs and output. - * @tparam Args Additional arguments to pass to the operator. - * - * @param output_size Total number of output elements. - * @param ndim Number of dimensions in tensors. - * @param output_contiguous Whether the output tensor is contiguous in memory. - * @param input_contiguous Array indicating if each input tensor is contiguous. - * @param input_broadcasted Array indicating if each input tensor is broadcasted. - * @param output_shape Shape of the output tensor. - * @param input_shapes Shapes of the input tensors. - * @param output_strides Strides for the output tensor. - * @param input_strides Strides for each input tensor. - * @param output Output buffer. - * @param inputs Array of input pointers, all of type Tdata. - * @param offset Linear offset to support partitioned execution. - * @param args Additional arguments passed to the operator. - */ -template -INFINIOP_CUDA_KERNEL elementwiseKernel( - size_t output_size, - size_t ndim, - bool output_contiguous, - const bool *__restrict__ input_contiguous, - const bool *__restrict__ input_broadcasted, - const size_t *__restrict__ output_shape, - const size_t *__restrict__ input_shapes, - const ptrdiff_t *__restrict__ output_strides, - const ptrdiff_t *__restrict__ input_strides, - Tdata *output, - const void *const *inputs, - size_t offset, - Args... args) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; - - if (idx < output_size) { - const Tdata *const *typed_inputs = reinterpret_cast(inputs); - size_t out_idx = getOutputIndex(idx, output_contiguous, ndim, output_shape, output_strides); - InputIndexer indexer{idx, ndim, input_contiguous, input_broadcasted, input_shapes, input_strides, output_strides}; - - unpackInputsAndApply( - [&](auto... Is) { -#if defined(ENABLE_HYGON_API) - output[out_idx] = Op{}(typed_inputs[Is.value][indexer(Is.value)]..., args...); -#else - output[out_idx] = Op{}(typed_inputs[Is.value][indexer(Is.value)]..., std::forward(args)...); -#endif - }, - std::make_index_sequence{}); - } -} - -/** - * @brief CUDA kernel for performing an elementwise operation on tensors with support - * for broadcasting and mixed data types. - * - * @tparam Op Operator type implementing a templated operator() for (Tout, Tin...). - * @tparam Tout Output data type. - * @tparam Tin Variadic input data types. - * - * @param output_size Total number of output elements. - * @param ndim Number of dimensions in the tensors. - * @param output_contiguous Whether the output tensor is contiguous. - * @param input_contiguous Array indicating whether each input is contiguous. - * @param input_broadcasted Array indicating whether each input is broadcasted. - * @param output_shape Shape of the output tensor. - * @param input_shapes Shapes of the input tensors. - * @param output_strides Strides of the output tensor. - * @param input_strides Strides of the input tensors. - * @param output Pointer to the output buffer. - * @param inputs Array of untyped input pointers. - * @param offset Linear offset into the output for partitioned execution. - */ -template -INFINIOP_CUDA_KERNEL elementwiseKernel( - size_t output_size, - size_t ndim, - bool output_contiguous, - const bool *__restrict__ input_contiguous, - const bool *__restrict__ input_broadcasted, - const size_t *__restrict__ output_shape, - const size_t *__restrict__ input_shapes, - const ptrdiff_t *__restrict__ output_strides, - const ptrdiff_t *__restrict__ input_strides, - Tout *output, - const void *const *__restrict__ inputs, - size_t offset) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; - - if (idx < output_size) { - size_t out_idx = getOutputIndex(idx, output_contiguous, ndim, output_shape, output_strides); - InputIndexer indexer{idx, ndim, input_contiguous, input_broadcasted, input_shapes, input_strides, output_strides}; - - unpackInputsAndApply( - [&](auto... Is) { - output[out_idx] = Op{}.template operator()( - (typedInputPtr(inputs[Is.value])[indexer(Is.value)])...); - }, - std::index_sequence_for{}); - } -} - -struct DeviceImpl::Opaque { - std::shared_ptr internal; - - Opaque(const std::shared_ptr &internal) - : internal(internal) {} - - /** - * @brief Executes an elementwise operation where all inputs and the output share the same data type. - * - * @tparam BLOCK_SIZE CUDA block size used for kernel launch. - * @tparam N Number of input tensors. - * @tparam Op Functor representing the elementwise operation. - * @tparam Tdata Data type of both input and output tensors. - * @tparam Args Optional additional arguments passed to the operation. - * - * @param info Metadata about the operation including shape, size, and dimensionality. - * @param workspace Temporary workspace used for storing metadata on device. - * @param output Pointer to the output buffer. - * @param inputs Vector of pointers to input buffers. - * @param stream CUDA stream for asynchronous execution. - * @param args Additional arguments forwarded to the operation. - * @return infiniStatus_t Returns success or failure status. - */ - template - infiniStatus_t calculateImpl(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - cudaStream_t stream, - Args &&...args) { - return launchElementwiseKernel( - info, workspace, - reinterpret_cast(output), inputs, - elementwiseKernel, - stream, - std::forward(args)...); - } - - /** - * @brief Executes an elementwise operation with mixed input and output data types. - * - * @tparam BLOCK_SIZE CUDA block size used for kernel launch. - * @tparam N Number of input tensors. - * @tparam Op Functor representing the elementwise operation. - * @tparam Tout Data type of the output tensor. - * @tparam Tin... Data types of the input tensors. - * @tparam Args Optional additional arguments passed to the operation.(UNUSED) - * - * @param info Metadata about the operation including shape, size, and dimensionality. - * @param workspace Temporary workspace used for storing metadata on device. - * @param output Pointer to the output buffer. - * @param inputs Vector of pointers to input buffers. - * @param stream CUDA stream for asynchronous execution. - * @param args Additional arguments forwarded to the operation. - * @return infiniStatus_t Returns success or failure status. - */ - template = 0> - infiniStatus_t calculateImpl(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - cudaStream_t stream, - Args &&...args) { - return launchElementwiseKernel( - info, workspace, - reinterpret_cast(output), inputs, - elementwiseKernel, - stream); - } - -private: - /** - * @brief Transfers elementwise operation metadata and input pointers from host to device memory. - * - * @tparam N Number of input tensors. - * - * @param info Elementwise operation metadata (shapes, strides, flags, etc.). - * @param workspace Pointer to device workspace memory for storing metadata and input pointers. - * @param h_inputs_arr Host array of input tensor pointers. - * @param d_inputs_arr Input reference to device array of input tensor pointers. - * @param d_input_contiguous Input reference to device array indicating whether each input is contiguous. - * @param d_input_broadcasted Input reference to device array indicating whether each input is broadcasted. - * @param d_output_shape Output reference to device array holding the output tensor shape. - * @param d_output_strides Output reference to device array holding output tensor strides. - * @param d_input_shapes Output reference to flattened input tensor shapes (N * ndim). - * @param d_input_strides Output reference to flattened input tensor strides (N * ndim). - * @param stream CUDA stream used for asynchronous memory transfer. - * @return infiniStatus_t Status indicating success or failure of the memory transfer and setup. - */ - template - infiniStatus_t infoToDevice( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - const void *const *h_inputs_arr, - const void **&d_inputs_arr, - const bool *&d_input_contiguous, - const bool *&d_input_broadcasted, - const size_t *&d_output_shape, - const ptrdiff_t *&d_output_strides, - const size_t *&d_input_shapes, - const ptrdiff_t *&d_input_strides, - cudaStream_t stream) const { - - constexpr auto input_size = N; - const auto ndim = info.getNdim(); - constexpr auto input_arr_size = N * sizeof(*h_inputs_arr); - const int8_t *info_meta_start = info.getMetaStart(); - const int8_t *d_meta_start = reinterpret_cast(workspace) + input_arr_size; - - // copy the input pointer array and meta to device - CHECK_CUDA(cudaMemcpyAsync(workspace, h_inputs_arr, input_arr_size, cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemcpyAsync((void *)d_meta_start, info_meta_start, info.getMetaMemSize(), cudaMemcpyHostToDevice, stream)); - - // offset/assign the pointers - d_inputs_arr = reinterpret_cast(workspace); - d_output_shape = reinterpret_cast(d_meta_start); - d_output_strides = reinterpret_cast(d_output_shape + ndim); - d_input_shapes = reinterpret_cast(d_output_strides + ndim); - d_input_strides = reinterpret_cast(d_input_shapes + input_size * ndim); - d_input_contiguous = reinterpret_cast(d_input_strides + input_size * ndim); - d_input_broadcasted = reinterpret_cast(d_input_contiguous + input_size); - - return INFINI_STATUS_SUCCESS; - } - - /** - * @brief Launches the elementwise kernel for the specified operation. - * - * @tparam BLOCK_SIZE Number of threads per block. - * @tparam N Number of input tensors. - * @tparam KernelFunc Type of the kernel function pointer. - * @tparam Tout Output data type. - * @tparam Args Additional arguments to be forwarded to the kernel. - * - * @param info Metadata about the elementwise operation (shapes, strides, etc.). - * @param workspace CUDA memory used for storing metadata. - * @param output Pointer to output buffer on device. - * @param inputs Vector of device pointers to input tensors. - * @param kernel_func Kernel function to launch. - * @param stream CUDA stream for asynchronous execution. - * @param args Additional arguments passed to the kernel. - * @return infiniStatus_t Status code indicating success or failure. - */ - template - infiniStatus_t launchElementwiseKernel( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - Tout *output, - const std::vector &inputs, - KernelFunc kernel_func, - cudaStream_t stream, - Args &&...args) { - - auto output_size = info.getOutputSize(); - if (output_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - // Device pointers - const void **d_inputs_arr = nullptr; - const bool *d_input_contiguous = nullptr; - const bool *d_input_broadcasted = nullptr; - const size_t *d_output_shape = nullptr; - const ptrdiff_t *d_output_strides = nullptr; - const size_t *d_input_shapes = nullptr; - const ptrdiff_t *d_input_strides = nullptr; - - CHECK_STATUS(infoToDevice(info, workspace, inputs.data(), d_inputs_arr, - d_input_contiguous, d_input_broadcasted, - d_output_shape, d_output_strides, - d_input_shapes, d_input_strides, stream)); - - dim3 blockDims(std::min(BLOCK_SIZE, static_cast(internal->maxThreadsPerBlock()))); - dim3 gridDims(std::min(uint32_t(CEIL_DIV(output_size, blockDims.x)), static_cast(internal->gridSizeX()))); - size_t step = gridDims.x * blockDims.x; - - for (size_t i = 0; i < output_size; i += step) { - kernel_func<<>>( - output_size, info.getNdim(), info.isOutputContiguous(), - d_input_contiguous, d_input_broadcasted, - d_output_shape, d_input_shapes, - d_output_strides, d_input_strides, - output, reinterpret_cast(d_inputs_arr), - i, std::forward(args)...); - } - - return INFINI_STATUS_SUCCESS; - } -}; - -template -utils::Result DeviceImpl::create(Args &&...args) { - auto opaque = std::make_shared(std::forward(args)...); - return utils::Result(new DeviceImpl(opaque)); -} - -/* Invoke elementwise operation for different input types */ -template > -infiniStatus_t DeviceImpl::calculate(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args) { - constexpr size_t N = Op::num_inputs; - static_assert(sizeof...(Tin) == N, "Input type count mismatch"); - return _opaque->calculateImpl( - info, workspace, output, inputs, - reinterpret_cast(stream), - std::forward(args)...); -} - -/* Invoke elementwise operation when all inputs have the same dtype */ -template -infiniStatus_t DeviceImpl::calculate(const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args) { - constexpr size_t N = Op::num_inputs; - return _opaque->calculateImpl( - info, workspace, output, inputs, - reinterpret_cast(stream), - std::forward(args)...); -} - -} // namespace op::elementwise::nvidia - -#endif // __INFINIOP_ELEMENTWISE_CUDA_H__ diff --git a/src/infiniop/elementwise/nvidia/elementwise_nvidia_api.cuh b/src/infiniop/elementwise/nvidia/elementwise_nvidia_api.cuh deleted file mode 100644 index 7e7442210..000000000 --- a/src/infiniop/elementwise/nvidia/elementwise_nvidia_api.cuh +++ /dev/null @@ -1,109 +0,0 @@ -#ifndef __INFINIOP_ELEMENTWISE_CUDA_API_H__ -#define __INFINIOP_ELEMENTWISE_CUDA_API_H__ - -#include "../elementwise.h" - -namespace op::elementwise::nvidia { - -/** - * @brief Define the methods and info needed by CUDA to perform elementwise operation - */ -class DeviceImpl final { - struct Opaque; - std::shared_ptr _opaque; - - DeviceImpl(std::shared_ptr opaque) : _opaque(std::move(opaque)) {} - -public: - ~DeviceImpl() = default; - - template - static utils::Result create(Args &&...args); - - /** - * @brief Launches elementwise operation where all input types are the same. - * - * Calls the corresponding templated `calculateImpl` with a unified input type. - * - * @tparam BLOCK_SIZE Number of threads per block. - * @tparam Op Operation functor defining the computation. - * @tparam Tdata Data type for both input and output tensors. - * @tparam Args... Additional arguments passed to the operation. - * - * @param info Metadata describing tensor shapes, strides, etc. - * @param workspace Pointer to workspace buffer on device. - * @param output Pointer to output buffer on device. - * @param inputs Vector of input pointers (device memory). - * @param stream CUDA stream (opaque void*). - * @param args Additional operation-specific arguments. - * @return infiniStatus_t Status indicating success or failure. - */ - template - infiniStatus_t calculate( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args); - - /** - * @brief Launches elementwise operation where input types may differ. - * - * Dispatches to templated `calculateImpl` using specified output and input types. - * - * @tparam BLOCK_SIZE Number of threads per block. - * @tparam Op Operation functor defining the computation. - * @tparam Tout Output data type. - * @tparam Tin... Input data types (must match Op::num_inputs). - * @tparam Args... Additional arguments passed to the operation. - * - * @param info Metadata describing tensor shapes, strides, etc. - * @param workspace Pointer to workspace buffer on device. - * @param output Pointer to output buffer on device. - * @param inputs Vector of input pointers (device memory). - * @param stream CUDA stream (opaque void*). - * @param args (UNUSED) Additional operation-specific arguments. - * @return infiniStatus_t Status indicating success or failure. - */ - template = 0> - infiniStatus_t calculate( - const op::elementwise::ElementwiseInfo &info, - void *workspace, - void *output, - const std::vector &inputs, - void *stream, - Args &&...args); -}; -} // namespace op::elementwise::nvidia - -/** - * @brief Define the process for initializing a Descriptor of an elementwise operation - * for its CUDA implementation - * - * @param HANDLE The device handle. - * @param DTYPE The output dtype. - * @param OUT_DESC The output tensor descriptor. - * @param INPUT_DESC_VEC A vector containing input tensor descriptors. - */ -#define CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(HANDLE, DTYPE, OUT_DESC, INPUT_DESC_VEC) \ - \ - auto info_result = op::elementwise::ElementwiseInfo::create(OUT_DESC, INPUT_DESC_VEC); \ - CHECK_RESULT(info_result); \ - auto info = info_result.take(); \ - auto workspace_size = info.getMetaMemSize() + info.getInputSize() * sizeof(void *); \ - \ - auto device_impl_result = op::elementwise::nvidia::DeviceImpl::create(HANDLE->internal()); \ - CHECK_RESULT(device_impl_result); \ - \ - *desc_ptr = new Descriptor( \ - DTYPE, \ - std::move(info), \ - std::move(device_impl_result.take()), \ - workspace_size, \ - HANDLE->device, \ - HANDLE->device_id); - -#endif // __INFINIOP_ELEMENTWISE_CUDA_API_H__ diff --git a/src/infiniop/handle.h b/src/infiniop/handle.h deleted file mode 100644 index d790be947..000000000 --- a/src/infiniop/handle.h +++ /dev/null @@ -1,11 +0,0 @@ -#ifndef __INFINIOP_HANDLE_H__ -#define __INFINIOP_HANDLE_H__ - -#include "infiniop/handle.h" - -struct InfiniopHandle { - infiniDevice_t device; - int device_id; -}; - -#endif //__INFINIOP_HANDLE_H__ diff --git a/src/infiniop/ninetoothed/build.py b/src/infiniop/ninetoothed/build.py deleted file mode 100644 index 153e6b9f5..000000000 --- a/src/infiniop/ninetoothed/build.py +++ /dev/null @@ -1,131 +0,0 @@ -import concurrent.futures -import functools -import inspect -import itertools -import pathlib - -import ninetoothed -from ninetoothed.aot import _HEADER_PATH - -CURRENT_FILE_PATH = pathlib.Path(__file__) - -BUILD_DIRECTORY_PATH = ( - CURRENT_FILE_PATH.parent.parent.parent.parent / "build" / "ninetoothed" -) - - -def build(premake, constexpr_param_grid, caller, op_name, output_dir): - headers = [] - all_param_names = [] - combinations = [] - launches = [] - - with concurrent.futures.ProcessPoolExecutor() as executor: - futures = [] - - for combination in tuple( - _generate_param_value_combinations(constexpr_param_grid) - ): - future = executor.submit( - _make, premake, combination, caller, op_name, output_dir - ) - - futures.append(future) - - for future in concurrent.futures.as_completed(futures): - header, param_names, combination, launch = future.result() - - headers.append(header) - all_param_names.append(param_names) - combinations.append(combination) - launches.append(launch) - - includes = "\n".join(f'#include "{header}"' for header in headers) - - param_names = list( - functools.reduce( - lambda x, y: dict.fromkeys(x) | dict.fromkeys(y), - sorted(all_param_names, key=len, reverse=True), - {}, - ) - ) - param_types = [ - "NineToothedStream", - ] + ["NineToothedTensor" for _ in range(len(param_names) - 1)] - - for param_name in functools.reduce(lambda x, y: x | y, combinations, {}): - param_names.append(param_name) - param_types.append("int") - - param_decls = ", ".join( - f"{type} {param}" for param, type in zip(param_names, param_types) - ) - - source_file_name = f"{op_name}.c" - header_file_name = f"{op_name}.h" - - func_sig = f"NineToothedResult launch_{op_name}({param_decls})" - - joined_launches = "\n".join(launches) - - op_decl = f'#ifdef __cplusplus\nextern "C" {func_sig};\n#else\n{func_sig};\n#endif' - op_def = f"""{func_sig} {{ -{joined_launches} - return INFINI_STATUS_NOT_IMPLEMENTED; -}}""" - - source_content = f"""#include "{header_file_name}" - -#include "infinicore.h" - -{includes}\n\n{op_def}\n""" - header_content = f"""#include "{_HEADER_PATH}" -\n{op_decl}\n""" - - (BUILD_DIRECTORY_PATH / source_file_name).write_text(source_content) - (BUILD_DIRECTORY_PATH / header_file_name).write_text(header_content) - - -def _make(premake, combination, caller, op_name, output_dir): - arrangement, application, tensors = premake(**combination) - - for param_name, param_value in combination.items(): - if isinstance(param_value, str): - combination[param_name] = ( - f"INFINI_DTYPE_{combination[param_name].replace('fp', 'F').upper()}" - ) - - combination = {f"{name}_": value for name, value in combination.items()} - - kernel_name = f"{op_name}_{_generate_suffix(combination.values())}" - - ninetoothed.make( - arrangement, - application, - tensors, - caller=caller, - kernel_name=kernel_name, - output_dir=output_dir, - ) - - header = output_dir / f"{kernel_name}.h" - param_names = ("stream",) + tuple(inspect.signature(application).parameters.keys()) - launch = f""" if ({_generate_condition(combination)}) - return launch_{kernel_name}({", ".join(param_names)});""" - - return header, param_names, combination, launch - - -def _generate_condition(combination): - return " && ".join(f"{param} == {value}" for param, value in combination.items()) - - -def _generate_suffix(values): - return "_".join(f"{value}" for value in values) - - -def _generate_param_value_combinations(param_grid): - keys = list(param_grid.keys()) - value_combinations = itertools.product(*param_grid.values()) - - return tuple(dict(zip(keys, combination)) for combination in value_combinations) diff --git a/src/infiniop/ninetoothed/utils.h b/src/infiniop/ninetoothed/utils.h deleted file mode 100644 index 1b7d1fe3a..000000000 --- a/src/infiniop/ninetoothed/utils.h +++ /dev/null @@ -1,75 +0,0 @@ -#ifndef __NINETOOTHED_UTILS__ -#define __NINETOOTHED_UTILS__ - -#include -#include -#include -#include - -namespace ninetoothed { - -template -class Tensor { -public: - using Data = decltype(NineToothedTensor::data); - - using Size = std::remove_pointer_t; - - using Stride = std::remove_pointer_t; - - template - Tensor(const void *data, Shape shape, Strides strides) : data_{data}, shape_{shape}, strides_{strides}, ndim_{shape_.size()} {} - - Tensor(const void *data, std::initializer_list shape, std::initializer_list strides) : Tensor{data, decltype(shape_){shape}, decltype(strides_){strides}} {} - - Tensor(const void *data, const Size *shape, const Stride *strides, Size ndim) : data_{data}, shape_{shape, shape + ndim}, strides_{strides, strides + ndim}, ndim_{shape_.size()} {} - - Tensor(const T value) : value_{value}, data_{&value_}, ndim_{0} {} - - operator NineToothedTensor() { return {const_cast(data_), shape_.data(), strides_.data()}; } - - template - Tensor expand(const Shape &sizes) const { - auto new_ndim{sizes.size()}; - - decltype(shape_) shape(new_ndim, 1); - decltype(strides_) strides(new_ndim, 0); - - auto num_new_dims{new_ndim - ndim_}; - - for (auto dim{decltype(ndim_){0}}; dim < ndim_; ++dim) { - shape[dim + num_new_dims] = shape_[dim]; - strides[dim + num_new_dims] = strides_[dim]; - } - - for (auto dim{decltype(new_ndim){0}}; dim < new_ndim; ++dim) { - if (sizes[dim] == std::numeric_limits>::max() || shape[dim] != 1) { - continue; - } - - shape[dim] = sizes[dim]; - strides[dim] = 0; - } - - return {data_, shape, strides}; - } - - Tensor expand_as(const Tensor &other) const { - return expand(other.shape_); - } - -private: - const void *data_{nullptr}; - - std::vector shape_; - - std::vector strides_; - - Size ndim_{0}; - - T value_{0}; -}; - -} // namespace ninetoothed - -#endif diff --git a/src/infiniop/operator.h b/src/infiniop/operator.h deleted file mode 100644 index efad3d5ba..000000000 --- a/src/infiniop/operator.h +++ /dev/null @@ -1,11 +0,0 @@ -#ifndef __INFINIOP_OPERATOR_H__ -#define __INFINIOP_OPERATOR_H__ - -#include "infiniop/operator_descriptor.h" - -struct InfiniopDescriptor { - infiniDevice_t device_type; - int device_id; -}; - -#endif diff --git a/src/infiniop/operator_descriptor.cc b/src/infiniop/operator_descriptor.cc deleted file mode 100644 index dac8cf008..000000000 --- a/src/infiniop/operator_descriptor.cc +++ /dev/null @@ -1,15 +0,0 @@ -#include "operator.h" - -infiniStatus_t infiniopGetDescriptorDeviceType( - const InfiniopDescriptor *desc_ptr, - infiniDevice_t *device_type) { - *device_type = desc_ptr->device_type; - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t infiniopGetDescriptorDeviceId( - const InfiniopDescriptor *desc_ptr, - int *device_id) { - *device_id = desc_ptr->device_id; - return INFINI_STATUS_SUCCESS; -} diff --git a/src/infiniop/ops/acos/cpu/acos_cpu.cc b/src/infiniop/ops/acos/cpu/acos_cpu.cc deleted file mode 100644 index f829134ca..000000000 --- a/src/infiniop/ops/acos/cpu/acos_cpu.cc +++ /dev/null @@ -1,59 +0,0 @@ -#include "acos_cpu.h" - -namespace op::acos::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - CHECK_DTYPE(dtype, - INFINI_DTYPE_BF16, - INFINI_DTYPE_F16, - INFINI_DTYPE_F32, - INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - // create CPU elementwise descriptor - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - // === 浮点类型 === - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - - // 移除了所有整数 Case - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::acos::cpu diff --git a/src/infiniop/ops/acos/cpu/acos_cpu.h b/src/infiniop/ops/acos/cpu/acos_cpu.h deleted file mode 100644 index 37d957bd0..000000000 --- a/src/infiniop/ops/acos/cpu/acos_cpu.h +++ /dev/null @@ -1,32 +0,0 @@ -#ifndef __ACOS_CPU_H__ -#define __ACOS_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" - -// 使用宏声明 Descriptor 类 -ELEMENTWISE_DESCRIPTOR(acos, cpu) - -#include -#include - -namespace op::acos::cpu { - -typedef struct AcosOp { -public: - static constexpr size_t num_inputs = 1; - - template - T operator()(const T &x) const { - if constexpr (std::is_integral_v) { - return static_cast(std::acos(static_cast(x))); - } else if constexpr (std::is_same_v || std::is_same_v) { - return std::acos(x); - } else { - return static_cast(std::acos(static_cast(x))); - } - } -} AcosOp; - -} // namespace op::acos::cpu - -#endif // __ACOS_CPU_H__ diff --git a/src/infiniop/ops/acos/cuda/kernel.cuh b/src/infiniop/ops/acos/cuda/kernel.cuh deleted file mode 100644 index fa0bf40e0..000000000 --- a/src/infiniop/ops/acos/cuda/kernel.cuh +++ /dev/null @@ -1,95 +0,0 @@ -#ifndef __ACOS_CUDA_H__ -#define __ACOS_CUDA_H__ - -#include -#include -#include - -namespace op::acos::cuda { - -// ---------------------- -// Fast acos approximation -// ---------------------- -__device__ __forceinline__ float fast_acosf(float x) { - // 高性能多项式近似 acos(x) - float ax = fabsf(x); - float t = sqrtf(1.0f - ax); - float r = ((-0.0187293f * ax + 0.0742610f) * ax - 0.2121144f) * ax + 1.5707288f; - return (x >= 0.0f ? t * r : 3.14159265358979323846f - t * r); -} - -// ---------------------- -// float kernel (F32) -// ---------------------- -template -__device__ __forceinline__ T acos_impl(T val); - -template <> -__device__ __forceinline__ float acos_impl(float val) { - return fast_acosf(val); -} - -// ---------------------- -// half kernel (F16) -// ---------------------- -template <> -__device__ __forceinline__ half acos_impl(half val) { -#if (__CUDA_ARCH__ >= 530) - float f = __half2float(val); - return __float2half(fast_acosf(f)); -#else - float f = __half2float(val); - return __float2half(fast_acosf(f)); -#endif -} - -// ---------------------- -// half2 kernel (F16x2 vectorized) -// ---------------------- -template <> -__device__ __forceinline__ half2 acos_impl(half2 val) { -#if (__CUDA_ARCH__ >= 530) - float2 f = __half22float2(val); - f.x = fast_acosf(f.x); - f.y = fast_acosf(f.y); - return __float22half2_rn(f); -#else - float2 f = __half22float2(val); - f.x = fast_acosf(f.x); - f.y = fast_acosf(f.y); - return __float22half2_rn(f); -#endif -} - -// ---------------------- -// bfloat16 kernel (BF16) -// ---------------------- -template <> -__device__ __forceinline__ cuda_bfloat16 acos_impl(cuda_bfloat16 val) { - float f = __bfloat162float(val); - return __float2bfloat16(fast_acosf(f)); -} - -// ---------------------- -// Fallback kernel -// ---------------------- -template -__device__ __forceinline__ T acos_impl(T val) { - return static_cast(fast_acosf(static_cast(val))); -} - -// ---------------------- -// AcosOp struct -// ---------------------- -struct AcosOp { - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(const T &a) const { - return acos_impl(a); - } -}; - -} // namespace op::acos::cuda - -#endif // __ACOS_CUDA_H__ diff --git a/src/infiniop/ops/acos/metax/acos_metax.h b/src/infiniop/ops/acos/metax/acos_metax.h deleted file mode 100644 index 9ec3d40c4..000000000 --- a/src/infiniop/ops/acos/metax/acos_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ACOS_METAX_API_H__ -#define __ACOS_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(acos, metax) - -#endif // __ACOS_METAX_API_H__ diff --git a/src/infiniop/ops/acos/metax/acos_metax.maca b/src/infiniop/ops/acos/metax/acos_metax.maca deleted file mode 100644 index aec018534..000000000 --- a/src/infiniop/ops/acos/metax/acos_metax.maca +++ /dev/null @@ -1,58 +0,0 @@ -#include "acos_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::acos::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - const auto &a_desc = input_desc_vec.at(0); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(c_shape, a_shape); - - // create CUDA elementwise descriptor - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::AcosOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::AcosOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::AcosOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::AcosOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::acos::metax diff --git a/src/infiniop/ops/acos/moore/acos_moore.h b/src/infiniop/ops/acos/moore/acos_moore.h deleted file mode 100644 index 50089ecda..000000000 --- a/src/infiniop/ops/acos/moore/acos_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ACOS_MOORE_API_H__ -#define __ACOS_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(acos, moore) - -#endif // __ACOS_MOORE_API_H__ diff --git a/src/infiniop/ops/acos/moore/acos_moore.mu b/src/infiniop/ops/acos/moore/acos_moore.mu deleted file mode 100644 index a11720b78..000000000 --- a/src/infiniop/ops/acos/moore/acos_moore.mu +++ /dev/null @@ -1,68 +0,0 @@ -#include "acos_moore.h" - -// 引入 Moore 平台的通用 Elementwise 描述符宏 -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "acos_moore_kernel.h" - -namespace op::acos::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - // Acos is a unary operator (y = acos(x)) - const auto &in_desc = input_desc_vec.at(0); - const auto &out_shape = out_desc->shape(); - const auto &in_shape = in_desc->shape(); - - // Acos supports floating point types. - // Unlike floor, acos generally doesn't support integer outputs directly. - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - // Check if output shape matches input shape - CHECK_SAME_SHAPE(out_shape, in_shape); - - // create MOORE elementwise descriptor - // 这里的宏会自动生成描述符初始化的通用代码 - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // Use moore::AcosOp template defined in acos_moore_kernel.h - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, moore::AcosOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, moore::AcosOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, moore::AcosOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, moore::AcosOp, double>(_info, workspace, output, inputs, stream); - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::acos::moore diff --git a/src/infiniop/ops/acos/moore/acos_moore_kernel.h b/src/infiniop/ops/acos/moore/acos_moore_kernel.h deleted file mode 100644 index bcce169a8..000000000 --- a/src/infiniop/ops/acos/moore/acos_moore_kernel.h +++ /dev/null @@ -1,54 +0,0 @@ -#ifndef __ACOS_MOORE_KERNEL_H__ -#define __ACOS_MOORE_KERNEL_H__ - -namespace op::acos::moore { - -typedef struct AcosOp { -public: - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(const T &input) const { - // ----------------------------------------------------------------- - // 1. Half2 (FP16x2) - // ----------------------------------------------------------------- - if constexpr (std::is_same_v) { - - float f1 = __low2float(input); - float f2 = __high2float(input); - return __floats2half2_rn(::acosf(f1), ::acosf(f2)); - } - // ----------------------------------------------------------------- - // 2. Half (FP16) - // ----------------------------------------------------------------- - else if constexpr (std::is_same_v) { - // Half fallback to float - float val_f = __half2float(input); - return __float2half(::acosf(val_f)); - } - // ----------------------------------------------------------------- - // 3. Bfloat16 - // ----------------------------------------------------------------- - else if constexpr (std::is_same_v) { - // BF16 fallback to float - float val_f = __bfloat162float(input); - return __float2bfloat16(::acosf(val_f)); - } - // ----------------------------------------------------------------- - // 4. Float32 - // ----------------------------------------------------------------- - else if constexpr (std::is_same_v) { - - return ::acosf(input); - } - // ----------------------------------------------------------------- - // 5. Double / Other - // ----------------------------------------------------------------- - else { - return ::acos(input); - } - } -} AcosOp; -} // namespace op::acos::moore - -#endif // __ACOS_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/acos/nvidia/acos_nvidia.cu b/src/infiniop/ops/acos/nvidia/acos_nvidia.cu deleted file mode 100644 index 8015e5031..000000000 --- a/src/infiniop/ops/acos/nvidia/acos_nvidia.cu +++ /dev/null @@ -1,67 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "acos_nvidia.cuh" - -namespace op::acos::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, - INFINI_DTYPE_BF16, - INFINI_DTYPE_F16, - INFINI_DTYPE_F32, - INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - // 使用通用的 Elementwise 描述符创建宏 - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // ----------------------------------------------------------- - // 算子分发：使用 cuda::AcosOp - // ----------------------------------------------------------- - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::AcosOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::AcosOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::AcosOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::AcosOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::acos::nvidia diff --git a/src/infiniop/ops/acos/nvidia/acos_nvidia.cuh b/src/infiniop/ops/acos/nvidia/acos_nvidia.cuh deleted file mode 100644 index d599fb96b..000000000 --- a/src/infiniop/ops/acos/nvidia/acos_nvidia.cuh +++ /dev/null @@ -1,6 +0,0 @@ -#ifndef __ACOS_NVIDIA_CUH__ -#define __ACOS_NVIDIA_CUH__ -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" -ELEMENTWISE_DESCRIPTOR(acos, nvidia) - -#endif // __FLOOR_NVIDIA_CUH__ diff --git a/src/infiniop/ops/acos/operator.cc b/src/infiniop/ops/acos/operator.cc deleted file mode 100644 index 16f2649ae..000000000 --- a/src/infiniop/ops/acos/operator.cc +++ /dev/null @@ -1,187 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/acos.h" - -// --- 后端实现头文件 --- -#ifdef ENABLE_CPU_API -#include "cpu/acos_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/acos_nvidia.cuh" -#endif - -#ifdef ENABLE_METAX_API -#include "metax/acos_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/acos_moore.h" -#endif - -extern "C" { - -// ======================================================================= -// 1. 创建算子描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopCreateAcosDescriptor( - infiniopHandle_t handle, - infiniopAcosDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::acos::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output, \ - {input}) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ======================================================================= -// 2. 获取 Workspace 大小 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopGetAcosWorkspaceSize(infiniopAcosDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -// ======================================================================= -// 3. 执行计算 (Calculate) -// ======================================================================= -__INFINI_C infiniStatus_t infiniopAcos( - infiniopAcosDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, {input}, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -// ======================================================================= -// 4. 销毁描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopDestroyAcosDescriptor(infiniopAcosDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} - -} // extern "C" diff --git a/src/infiniop/ops/adaptive_avg_pool1d/adaptive_avg_pool1d.h b/src/infiniop/ops/adaptive_avg_pool1d/adaptive_avg_pool1d.h deleted file mode 100644 index 4534d0b93..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/adaptive_avg_pool1d.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef ADAPTIVE_AVG_POOL1D_H -#define ADAPTIVE_AVG_POOL1D_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::adaptive_avg_pool1d::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AdaptiveAvgPool1dInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - AdaptiveAvgPool1dInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t in_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *output, \ - const void *input, \ - void *stream) const; \ - }; \ - } - -#endif // ADAPTIVE_AVG_POOL1D_H diff --git a/src/infiniop/ops/adaptive_avg_pool1d/cpu/adaptive_avg_pool1d_cpu.cc b/src/infiniop/ops/adaptive_avg_pool1d/cpu/adaptive_avg_pool1d_cpu.cc deleted file mode 100644 index b6f412f6f..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/cpu/adaptive_avg_pool1d_cpu.cc +++ /dev/null @@ -1,112 +0,0 @@ -#include "adaptive_avg_pool1d_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include -#include - -namespace op::adaptive_avg_pool1d::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); - auto result = AdaptiveAvgPool1dInfo::create(out_desc, in_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - nullptr, // Opaque* - result.take(), // Info - 0, // Workspace Size - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -void calculate( - const AdaptiveAvgPool1dInfo &info, - void *output, - const void *input) { - - size_t num_channels = info.num_channels(); - size_t isize = info.input_size(); - size_t osize = info.output_size(); - - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - -#pragma omp parallel for - for (ptrdiff_t c = 0; c < (ptrdiff_t)num_channels; ++c) { - - const Tdata *in_c = in_ptr + c * isize; - Tdata *out_c = out_ptr + c * osize; - - for (size_t i = 0; i < osize; ++i) { - - size_t istart = (i * isize) / osize; - size_t iend = ((i + 1) * isize + osize - 1) / osize; - - size_t klen = iend - istart; - - float sum = 0.0f; - - for (size_t j = istart; j < iend; ++j) { - if constexpr (std::is_same_v || std::is_same_v) { - sum += utils::cast(in_c[j]); - } else { - sum += (float)in_c[j]; - } - } - - float avg = (klen > 0) ? (sum / (float)klen) : 0.0f; - - if constexpr (std::is_same_v || std::is_same_v) { - out_c[i] = utils::cast(avg); - } else { - out_c[i] = (Tdata)avg; - } - } - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) const { - - // 从 Info 中获取 dtype - auto dtype = _info.dtype(); - - switch (dtype) { - case INFINI_DTYPE_F16: - cpu::calculate(_info, output, input); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_BF16: - cpu::calculate(_info, output, input); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_F32: - cpu::calculate(_info, output, input); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_F64: - cpu::calculate(_info, output, input); - return INFINI_STATUS_SUCCESS; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::adaptive_avg_pool1d::cpu diff --git a/src/infiniop/ops/adaptive_avg_pool1d/cpu/adaptive_avg_pool1d_cpu.h b/src/infiniop/ops/adaptive_avg_pool1d/cpu/adaptive_avg_pool1d_cpu.h deleted file mode 100644 index 3a76ae1bb..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/cpu/adaptive_avg_pool1d_cpu.h +++ /dev/null @@ -1,6 +0,0 @@ -#ifndef __ADAPTIVE_AVG_POOL1D_CPU_H__ -#define __ADAPTIVE_AVG_POOL1D_CPU_H__ - -#include "../adaptive_avg_pool1d.h" -DESCRIPTOR(cpu) -#endif // __ADAPTIVE_AVG_POOL1D_CPU_H__ diff --git a/src/infiniop/ops/adaptive_avg_pool1d/cuda/kernel.cuh b/src/infiniop/ops/adaptive_avg_pool1d/cuda/kernel.cuh deleted file mode 100644 index 073b13290..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/cuda/kernel.cuh +++ /dev/null @@ -1,199 +0,0 @@ -#ifndef __ADAPTIVE_AVG_POOL1D_CUDA_H__ -#define __ADAPTIVE_AVG_POOL1D_CUDA_H__ - -#include -#include - -namespace op::adaptive_avg_pool1d::cuda { - -// ------------------------------------------- -// 工具：Warp 级归约求和 -// ------------------------------------------- -template -__device__ __forceinline__ T warp_reduce_sum(T val) { -#pragma unroll - for (int offset = 16; offset > 0; offset /= 2) { - val += __shfl_down_sync(0xffffffff, val, offset); - } - return val; -} - -// ------------------------------------------- -// 工具：数值转换 -// ------------------------------------------- -template -__device__ __forceinline__ float to_float(const T &x) { - if constexpr (std::is_same_v) { - return __half2float(x); - } -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - else if constexpr (std::is_same_v) { - return __bfloat162float(x); - } -#endif - else { - return static_cast(x); - } -} - -template -__device__ __forceinline__ T from_float(float x) { - if constexpr (std::is_same_v) { - return __float2half(x); - } -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - else if constexpr (std::is_same_v) { - return __float2bfloat16(x); - } -#endif - else { - return static_cast(x); - } -} - -// ------------------------------------------- -// Optimization 1: Global Average Pool 特化 Kernel (osize == 1) -// 策略：1 Block 处理 1 个 Channel，Block 内归约 -// ------------------------------------------- -template -__global__ void global_avg_pool1d_kernel( - T *output, - const T *input, - size_t total_channels, // batch * channels - size_t isize) { - // 每一个 Block 处理一个 (Batch, Channel) 任务 - size_t channel_idx = blockIdx.x; - if (channel_idx >= total_channels) { - return; - } - - const T *channel_input = input + channel_idx * isize; - float sum = 0.0f; - - // Grid-Stride Loop within the channel (handle isize > blockDim.x) - for (size_t i = threadIdx.x; i < isize; i += blockDim.x) { - sum += to_float(channel_input[i]); - } - - // Block 内归约 - // 1. Warp Reduce - sum = warp_reduce_sum(sum); - - // 2. Shared Memory Reduce (跨 Warp) - static __shared__ float shared_sum[32]; // Max 1024 threads / 32 = 32 warps - int lane = threadIdx.x % 32; - int wid = threadIdx.x / 32; - - if (lane == 0) { - shared_sum[wid] = sum; - } - __syncthreads(); - - // 3. 让第一个 Warp 把 shared memory 里的结果加起来 - if (wid == 0) { - float val = (threadIdx.x < (blockDim.x + 31) / 32) ? shared_sum[lane] : 0.0f; - val = warp_reduce_sum(val); - if (threadIdx.x == 0) { - output[channel_idx] = from_float(val / static_cast(isize)); - } - } -} - -// ------------------------------------------- -// Optimization 2: 通用 Kernel -// 策略：使用浮点数计算索引，避免整数除法；移除不安全的向量化 -// ------------------------------------------- -template -__global__ void adaptive_avg_pool1d_general_kernel( - T *output, - const T *input, - size_t batch_channels, - size_t isize, - size_t osize) { - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - size_t stride = gridDim.x * blockDim.x; - size_t total_elements = batch_channels * osize; - - // 预计算缩放因子，避免循环内除法 - float stride_factor = static_cast(isize) / static_cast(osize); - - for (; idx < total_elements; idx += stride) { - size_t bc_idx = idx / osize; - size_t out_idx = idx % osize; - - const T *in_ptr = input + bc_idx * isize; - - // 使用浮点数计算起止点，替代 (i * isize) / osize - // 注意：PyTorch 官方实现使用 float 进行索引计算 - int istart = static_cast(floorf(out_idx * stride_factor)); - int iend = static_cast(ceilf((out_idx + 1) * stride_factor)); - - // 边界保护 - istart = max(0, istart); - iend = min(static_cast(isize), iend); - - float sum = 0.0f; - int klen = iend - istart; - - // 标量循环：利用 L1 Cache 和编译器优化 - // 移除手动向量化，因为 istart 不保证对齐 - for (int i = istart; i < iend; ++i) { - sum += to_float(in_ptr[i]); - } - - output[idx] = (klen > 0) ? from_float(sum / klen) : from_float(0.0f); - } -} - -// ------------------------------------------- -// Launcher -// ------------------------------------------- -template -void launch_adaptive_avg_pool1d( - T *output, - const T *input, - size_t batch_channels, - size_t isize, - size_t osize, - cudaStream_t stream) { - // 策略分发 - if (osize == 1) { - // Case 1: Global Average Pooling (Gap) - // 每个 Block 处理一个 Channel - int threads = 256; - // 如果 isize 很小，减少线程数 - if (isize < 256) { - threads = 128; - } - if (isize < 128) { - threads = 64; - } - if (isize < 64) { - threads = 32; - } - - dim3 block(threads); - dim3 grid(batch_channels); - - global_avg_pool1d_kernel<<>>( - output, input, batch_channels, isize); - } else { - // Case 2: General Case - // 这里的并行度基于输出元素个数 - size_t total_output = batch_channels * osize; - int threads = 256; - int blocks = (total_output + threads - 1) / threads; - - // 限制最大 Grid 大小，防止超限 - if (blocks > 65535) { - blocks = 65535; - } - - adaptive_avg_pool1d_general_kernel<<>>( - output, input, batch_channels, isize, osize); - } -} - -} // namespace op::adaptive_avg_pool1d::cuda - -#endif // __ADAPTIVE_AVG_POOL1D_CUDA_H__ diff --git a/src/infiniop/ops/adaptive_avg_pool1d/info.h b/src/infiniop/ops/adaptive_avg_pool1d/info.h deleted file mode 100644 index 610780d14..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/info.h +++ /dev/null @@ -1,62 +0,0 @@ -#ifndef __ADAPTIVE_AVG_POOL1D_INFO_H__ -#define __ADAPTIVE_AVG_POOL1D_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::adaptive_avg_pool1d { - -class AdaptiveAvgPool1dInfo { - AdaptiveAvgPool1dInfo() = default; - -public: - size_t _input_size; - size_t _output_size; - size_t _num_channels; - int _dtype; - - size_t input_size() const { return _input_size; } - size_t output_size() const { return _output_size; } - size_t num_channels() const { return _num_channels; } - int dtype() const { return _dtype; } - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc) { - - // 1. 检查数据类型一致性 - if (out_desc->dtype() != in_desc->dtype()) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // 2. 检查维度 (至少 2 维: C, L 或 N, C, L) - size_t ndim = in_desc->ndim(); - if (ndim < 2 || out_desc->ndim() != ndim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t num_channels = 1; - for (size_t i = 0; i < ndim - 1; ++i) { - if (in_desc->shape()[i] != out_desc->shape()[i]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - num_channels *= in_desc->shape()[i]; - } - - // 4. 获取输入和输出的长度 (L) - size_t input_size = in_desc->shape()[ndim - 1]; - size_t output_size = out_desc->shape()[ndim - 1]; - int dtype = in_desc->dtype(); - - return utils::Result(AdaptiveAvgPool1dInfo{ - input_size, - output_size, - num_channels, - dtype}); - } -}; - -} // namespace op::adaptive_avg_pool1d - -#endif // __ADAPTIVE_AVG_POOL1D_INFO_H__ diff --git a/src/infiniop/ops/adaptive_avg_pool1d/metax/adaptive_avg_pool1d_metax.h b/src/infiniop/ops/adaptive_avg_pool1d/metax/adaptive_avg_pool1d_metax.h deleted file mode 100644 index 8162610e9..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/metax/adaptive_avg_pool1d_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADAPTIVE_AVG_POOL1D_METAX_H__ -#define __ADAPTIVE_AVG_POOL1D_METAX_H__ - -#include "../adaptive_avg_pool1d.h" - -DESCRIPTOR(metax) - -#endif // __ADAPTIVE_AVG_POOL1D_METAX_H__ diff --git a/src/infiniop/ops/adaptive_avg_pool1d/metax/adaptive_avg_pool1d_metax.maca b/src/infiniop/ops/adaptive_avg_pool1d/metax/adaptive_avg_pool1d_metax.maca deleted file mode 100644 index 7bc130e70..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/metax/adaptive_avg_pool1d_metax.maca +++ /dev/null @@ -1,188 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "adaptive_avg_pool1d_metax.h" - -namespace op::adaptive_avg_pool1d::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc) { - - auto handle = reinterpret_cast(handle_); - - // 检查数据类型 - auto dtype = out_desc->dtype(); - if (dtype != INFINI_DTYPE_F16 && dtype != INFINI_DTYPE_F32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // 调用 Info::create (只传两个参数) - auto result = AdaptiveAvgPool1dInfo::create(out_desc, in_desc); - - // 简单的错误检查，Info::create 返回 Result 类型 - if (!result) { - return INFINI_STATUS_BAD_PARAM; - } - - *desc_ptr = new Descriptor( - new Opaque{handle->internal()}, - result.take(), - 0, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) const { - - // 准备 HCDNN 参数 - hcdnnDataType_t data_type; - hcdnnPoolingMode_t pooling_mode = HCDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING; - hcdnnTensorFormat_t tensor_format = HCDNN_TENSOR_NCHW; - - float alpha = 1.0f; - float beta = 0.0f; - - switch (_info.dtype()) { - case INFINI_DTYPE_F16: - data_type = HCDNN_DATA_HALF; - break; - case INFINI_DTYPE_F32: - data_type = HCDNN_DATA_FLOAT; - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // 维度映射: - // Info 类已经将 N 和 C 展平为 num_channels。 - // 为了适配 hcdnn (NCHW)，我们将 num_channels 视为 N (Batch)，C=1, H=1, W=Length。 - // 这样做对于 1D Pooling 是数学上等价且高效的。 - int n = static_cast(_info.num_channels()); - int c = 1; - int h_in = 1; - int w_in = static_cast(_info.input_size()); - - int h_out = 1; - int w_out = static_cast(_info.output_size()); - - // 动态计算 Kernel 和 Stride (Adaptive -> Fixed转换) - // 1. Stride = Input / Output - // 2. Kernel = Input - (Output - 1) * Stride - // 注意: 如果 Input 不能被 Output 整除，这种固定 Stride/Kernel 的方式只是近似， - // 但在使用标准 Pooling API 实现 Adaptive Pooling 时是常用做法。 - int stride_w = w_in / w_out; - int kernel_w = w_in - (w_out - 1) * stride_w; - int pad_w = 0; - - int windowHeight = 1; - int windowWidth = kernel_w; - int verticalPadding = 0; - int horizontalPadding = pad_w; - int verticalStride = 1; - int horizontalStride = stride_w; - - return _opaque->internal->useMcdnn( - (hcStream_t)stream, - [&](hcdnnHandle_t handle) { - hcdnnStatus_t status; - // 1. 输入张量描述符 - hcdnnTensorDescriptor_t input_desc; - status = hcdnnCreateTensorDescriptor(&input_desc); - if (status != HCDNN_STATUS_SUCCESS) { - return INFINI_STATUS_INTERNAL_ERROR; - } - - status = hcdnnSetTensor4dDescriptor( - input_desc, - tensor_format, - data_type, - n, c, h_in, w_in); - if (status != HCDNN_STATUS_SUCCESS) { - hcdnnDestroyTensorDescriptor(input_desc); - return INFINI_STATUS_INTERNAL_ERROR; - } - - // 2. 输出张量描述符 - hcdnnTensorDescriptor_t output_desc; - status = hcdnnCreateTensorDescriptor(&output_desc); - if (status != HCDNN_STATUS_SUCCESS) { - hcdnnDestroyTensorDescriptor(input_desc); - return INFINI_STATUS_INTERNAL_ERROR; - } - - status = hcdnnSetTensor4dDescriptor( - output_desc, - tensor_format, - data_type, - n, c, h_out, w_out); - if (status != HCDNN_STATUS_SUCCESS) { - hcdnnDestroyTensorDescriptor(input_desc); - hcdnnDestroyTensorDescriptor(output_desc); - return INFINI_STATUS_INTERNAL_ERROR; - } - - // 3. Pooling 描述符 - hcdnnPoolingDescriptor_t pool_desc; - status = hcdnnCreatePoolingDescriptor(&pool_desc); - if (status != HCDNN_STATUS_SUCCESS) { - hcdnnDestroyTensorDescriptor(input_desc); - hcdnnDestroyTensorDescriptor(output_desc); - return INFINI_STATUS_INTERNAL_ERROR; - } - - status = hcdnnSetPooling2dDescriptor( - pool_desc, - pooling_mode, - HCDNN_NOT_PROPAGATE_NAN, - windowHeight, windowWidth, - verticalPadding, horizontalPadding, - verticalStride, horizontalStride); - - if (status != HCDNN_STATUS_SUCCESS) { - hcdnnDestroyTensorDescriptor(input_desc); - hcdnnDestroyTensorDescriptor(output_desc); - hcdnnDestroyPoolingDescriptor(pool_desc); - return INFINI_STATUS_INTERNAL_ERROR; - } - - // 4. 执行 - status = hcdnnPoolingForward( - handle, - pool_desc, - &alpha, - input_desc, - input, - &beta, - output_desc, - output); - - // 清理 - hcdnnDestroyTensorDescriptor(input_desc); - hcdnnDestroyTensorDescriptor(output_desc); - hcdnnDestroyPoolingDescriptor(pool_desc); - - if (status != HCDNN_STATUS_SUCCESS) { - return INFINI_STATUS_INTERNAL_ERROR; - } - - return INFINI_STATUS_SUCCESS; - }); -} - -} // namespace op::adaptive_avg_pool1d::metax diff --git a/src/infiniop/ops/adaptive_avg_pool1d/moore/adaptive_avg_pool1d_moore.h b/src/infiniop/ops/adaptive_avg_pool1d/moore/adaptive_avg_pool1d_moore.h deleted file mode 100644 index 6f59d26e7..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/moore/adaptive_avg_pool1d_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADAPTIVE_AVG_POOL1D_MOORE_API_H__ -#define __ADAPTIVE_AVG_POOL1D_MOORE_API_H__ - -// 引入上层定义的 Descriptor 宏和基础类 -#include "../adaptive_avg_pool1d.h" -DESCRIPTOR(moore) - -#endif // __ADAPTIVE_AVG_POOL1D_MOORE_API_H__ diff --git a/src/infiniop/ops/adaptive_avg_pool1d/moore/adaptive_avg_pool1d_moore.mu b/src/infiniop/ops/adaptive_avg_pool1d/moore/adaptive_avg_pool1d_moore.mu deleted file mode 100644 index 251a57ce3..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/moore/adaptive_avg_pool1d_moore.mu +++ /dev/null @@ -1,175 +0,0 @@ -#include "adaptive_avg_pool1d_moore.h" -#include "adaptive_avg_pool1d_moore_kernel.h" - -#include -#include -#include -#include - -#include "../../../devices/moore/moore_handle.h" - -namespace op::adaptive_avg_pool1d::moore { - -// ================================================================== -// 1. Kernel Implementation -// ================================================================== - -template -__global__ void adaptive_avg_pool1d_kernel( - const int total_elements, - const int input_size, - const int output_size, - const T *input, - T *output) { - - int idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < total_elements) { - int w_out = idx % output_size; - int bc = idx / output_size; - - int start = (w_out * input_size) / output_size; - int end = ((w_out + 1) * input_size + output_size - 1) / output_size; - - start = (start < 0) ? 0 : start; - end = (end > input_size) ? input_size : end; - - int kernel_size = end - start; - if (kernel_size < 1) { - kernel_size = 1; - } - - const T *in_ptr = input + bc * input_size; - - float sum = 0.0f; - for (int i = start; i < end; ++i) { - T val = in_ptr[i]; - if constexpr (std::is_same_v) { - sum += __half2float(val); - } else if constexpr (std::is_same_v) { - sum += __bfloat162float(val); - } else { - sum += static_cast(val); - } - } - - float avg = sum / static_cast(kernel_size); - - if constexpr (std::is_same_v) { - output[idx] = __float2half(avg); - } else if constexpr (std::is_same_v) { - output[idx] = __float2bfloat16(avg); - } else { - output[idx] = static_cast(avg); - } - } -} - -// ================================================================== -// 2. Launcher Implementation -// ================================================================== - -template -void adaptive_avg_pool1d_moore_launch( - const AdaptiveAvgPool1dInfo &info, - T *output, - const T *input, - void *stream) { - - int input_size = info.input_size(); - int output_size = info.output_size(); - - size_t total_elements = info.num_channels() * output_size; - - int threads = 256; - int blocks = (total_elements + threads - 1) / threads; - - adaptive_avg_pool1d_kernel<<>>( - total_elements, - input_size, - output_size, - input, - output); -} - -// ================================================================== -// 3. Descriptor Implementation -// ================================================================== - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc) { - - auto handle = reinterpret_cast(handle_); - - auto info_result = AdaptiveAvgPool1dInfo::create(out_desc, in_desc); - if (!info_result) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - *desc_ptr = new Descriptor( - nullptr, - *info_result, - 0, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_info.dtype()) { - case INFINI_DTYPE_F16: - adaptive_avg_pool1d_moore_launch( - _info, - static_cast(output), - static_cast(input), - stream); - break; - - case INFINI_DTYPE_BF16: - adaptive_avg_pool1d_moore_launch<__mt_bfloat16>( - _info, - static_cast<__mt_bfloat16 *>(output), - static_cast(input), - stream); - break; - - case INFINI_DTYPE_F32: - adaptive_avg_pool1d_moore_launch( - _info, - static_cast(output), - static_cast(input), - stream); - break; - - case INFINI_DTYPE_F64: - adaptive_avg_pool1d_moore_launch( - _info, - static_cast(output), - static_cast(input), - stream); - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::adaptive_avg_pool1d::moore diff --git a/src/infiniop/ops/adaptive_avg_pool1d/moore/adaptive_avg_pool1d_moore_kernel.h b/src/infiniop/ops/adaptive_avg_pool1d/moore/adaptive_avg_pool1d_moore_kernel.h deleted file mode 100644 index 8021eae32..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/moore/adaptive_avg_pool1d_moore_kernel.h +++ /dev/null @@ -1,60 +0,0 @@ -#ifndef __ADAPTIVE_AVG_POOL1D_MOORE_KERNEL_H__ -#define __ADAPTIVE_AVG_POOL1D_MOORE_KERNEL_H__ - -#include -#include -#include - -#include - -namespace op::adaptive_avg_pool1d::moore { - -typedef struct AdaptiveAvgPool1dOp { -public: - template - __device__ __forceinline__ void operator()( - const int w_out, - const int input_size, - const int output_size, - const T *input_base, - T *output_ptr) const { - - int start = (w_out * input_size) / output_size; - int end = ((w_out + 1) * input_size + output_size - 1) / output_size; - - start = (start < 0) ? 0 : start; - end = (end > input_size) ? input_size : end; - - int kernel_size = end - start; - kernel_size = (kernel_size < 1) ? 1 : kernel_size; - - float sum = 0.0f; - - for (int i = start; i < end; ++i) { - T val = input_base[i]; - - if constexpr (std::is_same_v) { - sum += __half2float(val); - } else if constexpr (std::is_same_v) { - sum += __bfloat162float(val); - } else { - sum += static_cast(val); - } - } - - float avg = sum / static_cast(kernel_size); - - if constexpr (std::is_same_v) { - *output_ptr = __float2half(avg); - } else if constexpr (std::is_same_v) { - *output_ptr = __float2bfloat16(avg); - } else { - *output_ptr = static_cast(avg); - } - } - -} AdaptiveAvgPool1dOp; - -} // namespace op::adaptive_avg_pool1d::moore - -#endif // __ADAPTIVE_AVG_POOL1D_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/adaptive_avg_pool1d/nvidia/adaptive_avg_pool1d_nvidia.cu b/src/infiniop/ops/adaptive_avg_pool1d/nvidia/adaptive_avg_pool1d_nvidia.cu deleted file mode 100644 index dd1d382de..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/nvidia/adaptive_avg_pool1d_nvidia.cu +++ /dev/null @@ -1,100 +0,0 @@ - -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" - -#include "../cuda/kernel.cuh" -#include "adaptive_avg_pool1d_nvidia.cuh" - -namespace op::adaptive_avg_pool1d::nvidia { - -template -void launch_kernel( - void *output, - const void *input, - size_t num_channels, // 这里实际上是 total_channels (Batch * C) - size_t isize, - size_t osize, - void *stream) { - - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - auto cuda_stream = reinterpret_cast(stream); - - cuda::launch_adaptive_avg_pool1d( - out_ptr, - in_ptr, - num_channels, - isize, - osize, - cuda_stream); -} - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc) { - - // 1. 使用 Info 类解析参数 - auto info_result = AdaptiveAvgPool1dInfo::create(out_desc, in_desc); - if (!info_result) { - return info_result.status(); - } - auto info = info_result.take(); - - // 2. 创建 Descriptor - *desc_ptr = new Descriptor( - new Opaque(), // Opaque 指针 - info, // Info 对象 - 0, // Workspace size - handle->device, // Device Type - handle->device_id // Device ID - ); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) const { - - auto dtype = _info.dtype(); - auto num_channels = _info.num_channels(); // 这里通常是 Batch * Channel - auto input_size = _info.input_size(); - auto output_size = _info.output_size(); - - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel(output, input, num_channels, input_size, output_size, stream); - break; -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - case INFINI_DTYPE_BF16: - launch_kernel(output, input, num_channels, input_size, output_size, stream); - break; -#endif - case INFINI_DTYPE_F32: - launch_kernel(output, input, num_channels, input_size, output_size, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, num_channels, input_size, output_size, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::adaptive_avg_pool1d::nvidia diff --git a/src/infiniop/ops/adaptive_avg_pool1d/nvidia/adaptive_avg_pool1d_nvidia.cuh b/src/infiniop/ops/adaptive_avg_pool1d/nvidia/adaptive_avg_pool1d_nvidia.cuh deleted file mode 100644 index b95e1088f..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/nvidia/adaptive_avg_pool1d_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADAPTIVE_AVG_POOL1D_CUH__ -#define __ADAPTIVE_AVG_POOL1D_CUH__ - -#include "../adaptive_avg_pool1d.h" - -DESCRIPTOR(nvidia) - -#endif // __GEMM_CUDA_CUH__ diff --git a/src/infiniop/ops/adaptive_avg_pool1d/operator.cc b/src/infiniop/ops/adaptive_avg_pool1d/operator.cc deleted file mode 100644 index 3c92f6d96..000000000 --- a/src/infiniop/ops/adaptive_avg_pool1d/operator.cc +++ /dev/null @@ -1,195 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/adaptive_avg_pool1d.h" - -// --- 后端实现头文件 --- -#ifdef ENABLE_CPU_API -#include "cpu/adaptive_avg_pool1d_cpu.h" -#endif - -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/adaptive_avg_pool1d_nvidia.cuh" -#endif - -// [Metax Support] -#ifdef ENABLE_METAX_API -#include "metax/adaptive_avg_pool1d_metax.h" -#endif - -// [Moore Threads Support] -#ifdef ENABLE_MOORE_API -#include "moore/adaptive_avg_pool1d_moore.h" -#endif - -extern "C" { - -// ======================================================================= -// 1. 创建算子描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopCreateAdaptiveAvgPool1dDescriptor( - infiniopHandle_t handle, - infiniopAdaptiveAvgPool1dDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::adaptive_avg_pool1d::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output_desc, \ - input_desc) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ======================================================================= -// 2. 获取 Workspace 大小 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopGetAdaptiveAvgPool1dWorkspaceSize(infiniopAdaptiveAvgPool1dDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (reinterpret_cast(desc)->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -// ======================================================================= -// 3. 执行计算 (Calculate) -// ======================================================================= -__INFINI_C infiniStatus_t infiniopAdaptiveAvgPool1d( - infiniopAdaptiveAvgPool1dDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, input, stream) - - switch (reinterpret_cast(desc)->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -// ======================================================================= -// 4. 销毁描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopDestroyAdaptiveAvgPool1dDescriptor(infiniopAdaptiveAvgPool1dDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (reinterpret_cast(desc)->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} - -} // extern "C" diff --git a/src/infiniop/ops/adaptive_avg_pool3d/adaptive_avg_pool3d.h b/src/infiniop/ops/adaptive_avg_pool3d/adaptive_avg_pool3d.h deleted file mode 100644 index 33280aad9..000000000 --- a/src/infiniop/ops/adaptive_avg_pool3d/adaptive_avg_pool3d.h +++ /dev/null @@ -1,90 +0,0 @@ -#ifndef INFINIOP_ADAPTIVE_AVG_POOL3D_DESCRIPTOR_H_ -#define INFINIOP_ADAPTIVE_AVG_POOL3D_DESCRIPTOR_H_ -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" -#include "infiniop/ops/adaptive_avg_pool3d.h" -#include - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::adaptive_avg_pool3d::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AdaptiveAvgPool3DInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - AdaptiveAvgPool3DInfo info, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t y_desc, \ - infiniopTensorDescriptor_t x_desc, \ - size_t *output_size); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *y, \ - const void *x, \ - void *stream) const; \ - }; \ - } - -class AdaptiveAvgPool3DInfo { -private: - AdaptiveAvgPool3DInfo() = default; - -public: - infiniDtype_t dtype; - size_t x_d, x_h, x_w; - size_t y_d, y_h, y_w; - size_t N, C; - std::vector x_strides; - std::vector y_strides; - - static utils::Result - create( - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t *output_size) { - CHECK_OR_RETURN(x_desc != nullptr && output_size != nullptr, - INFINI_STATUS_NULL_POINTER); - - const infiniDtype_t data_type = x_desc->dtype(); - CHECK_DTYPE(data_type, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - const size_t ndim = x_desc->ndim(); - CHECK_OR_RETURN(ndim == 5, INFINI_STATUS_BAD_TENSOR_SHAPE); - - return utils::Result(AdaptiveAvgPool3DInfo{ - data_type, - x_desc->dim(2), - x_desc->dim(3), - x_desc->dim(4), - output_size[0], - output_size[1], - output_size[2], - x_desc->dim(0), - x_desc->dim(1), - x_desc->strides(), - y_desc->strides()}); - } -}; - -#endif diff --git a/src/infiniop/ops/adaptive_avg_pool3d/cpu/adaptive_avg_pool3d_cpu.cc b/src/infiniop/ops/adaptive_avg_pool3d/cpu/adaptive_avg_pool3d_cpu.cc deleted file mode 100644 index 6006e8ea6..000000000 --- a/src/infiniop/ops/adaptive_avg_pool3d/cpu/adaptive_avg_pool3d_cpu.cc +++ /dev/null @@ -1,136 +0,0 @@ -#include "adaptive_avg_pool3d_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include -namespace op::adaptive_avg_pool3d::cpu { -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t *output_size) { - - auto handle = reinterpret_cast(handle_); - - auto info = AdaptiveAvgPool3DInfo::create(y_desc, x_desc, output_size); - CHECK_RESULT(info); - - *desc_ptr = new Descriptor( - info.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateAdaptiveAvgPool3D( - const AdaptiveAvgPool3DInfo &info, - Tdata *y, - const Tdata *x) { - - std::array y_strides; - y_strides[4] = 1; - y_strides[3] = info.y_w * y_strides[4]; - y_strides[2] = info.y_h * y_strides[3]; - y_strides[1] = info.y_d * y_strides[2]; - y_strides[0] = info.C * y_strides[1]; - - const size_t total = info.N * info.C * info.y_d * info.y_h * info.y_w; - -#pragma omp parallel for - for (ptrdiff_t idx = 0; idx < (ptrdiff_t)total; idx++) { - - // 🔹 Decode flattened index → (n, c, od, oh, ow) - ptrdiff_t tmp = idx; - - ptrdiff_t ow = tmp % info.y_w; - tmp /= info.y_w; - ptrdiff_t oh = tmp % info.y_h; - tmp /= info.y_h; - ptrdiff_t od = tmp % info.y_d; - tmp /= info.y_d; - ptrdiff_t c = tmp % info.C; - tmp /= info.C; - ptrdiff_t n = tmp; - - // 🔹 Compute pooling region - size_t x_start_d = od * info.x_d / info.y_d; - size_t x_end_d = ((od + 1) * info.x_d + info.y_d - 1) / info.y_d; - - size_t x_start_h = oh * info.x_h / info.y_h; - size_t x_end_h = ((oh + 1) * info.x_h + info.y_h - 1) / info.y_h; - - size_t x_start_w = ow * info.x_w / info.y_w; - size_t x_end_w = ((ow + 1) * info.x_w + info.y_w - 1) / info.y_w; - - size_t count = (x_end_d - x_start_d) * (x_end_h - x_start_h) * (x_end_w - x_start_w); - - size_t y_offset = n * y_strides[0] + c * y_strides[1] + od * y_strides[2] + oh * y_strides[3] + ow * y_strides[4]; - - // 🔹 Accumulate - if constexpr (std::is_same::value || std::is_same::value) { - - float sum = 0.0f; - - for (size_t id = x_start_d; id < x_end_d; id++) { - for (size_t ih = x_start_h; ih < x_end_h; ih++) { - for (size_t iw = x_start_w; iw < x_end_w; iw++) { - - size_t x_offset = n * info.x_strides[0] + c * info.x_strides[1] + id * info.x_strides[2] + ih * info.x_strides[3] + iw * info.x_strides[4]; - - sum += utils::cast(x[x_offset]); - } - } - } - - y[y_offset] = utils::cast(sum / (float)count); - - } else { - - Tdata sum = (Tdata)0; - - for (size_t id = x_start_d; id < x_end_d; id++) { - for (size_t ih = x_start_h; ih < x_end_h; ih++) { - for (size_t iw = x_start_w; iw < x_end_w; iw++) { - - size_t x_offset = n * info.x_strides[0] + c * info.x_strides[1] + id * info.x_strides[2] + ih * info.x_strides[3] + iw * info.x_strides[4]; - - sum += x[x_offset]; - } - } - } - - y[y_offset] = sum / (Tdata)count; - } - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - switch (_info.dtype) { - case INFINI_DTYPE_F16: - return calculateAdaptiveAvgPool3D(_info, reinterpret_cast(y), reinterpret_cast(x)); - case INFINI_DTYPE_F32: - return calculateAdaptiveAvgPool3D(_info, reinterpret_cast(y), reinterpret_cast(x)); - case INFINI_DTYPE_F64: - return calculateAdaptiveAvgPool3D(_info, reinterpret_cast(y), reinterpret_cast(x)); - case INFINI_DTYPE_BF16: - return calculateAdaptiveAvgPool3D(_info, reinterpret_cast(y), reinterpret_cast(x)); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::adaptive_avg_pool3d::cpu diff --git a/src/infiniop/ops/adaptive_avg_pool3d/cpu/adaptive_avg_pool3d_cpu.h b/src/infiniop/ops/adaptive_avg_pool3d/cpu/adaptive_avg_pool3d_cpu.h deleted file mode 100644 index f20cfcc77..000000000 --- a/src/infiniop/ops/adaptive_avg_pool3d/cpu/adaptive_avg_pool3d_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __INFINIOP_ADAPTIVE_AVG_POOL3D_CPU_H__ -#define __INFINIOP_ADAPTIVE_AVG_POOL3D_CPU_H__ - -#include "../adaptive_avg_pool3d.h" - -DESCRIPTOR(cpu) - -#endif diff --git a/src/infiniop/ops/adaptive_avg_pool3d/cuda/kernel.cuh b/src/infiniop/ops/adaptive_avg_pool3d/cuda/kernel.cuh deleted file mode 100644 index 6493ea786..000000000 --- a/src/infiniop/ops/adaptive_avg_pool3d/cuda/kernel.cuh +++ /dev/null @@ -1,44 +0,0 @@ -#ifndef __ADAPTIVE_AVG_POOL3D_CUDA_H__ -#define __ADAPTIVE_AVG_POOL3D_CUDA_H__ - -#include - -template -__device__ void adaptiveAvgPool3DKernel(T *y, const T *x, size_t N, size_t C, - size_t x_d, size_t x_h, size_t x_w, - size_t y_d, size_t y_h, size_t y_w, - const ptrdiff_t *x_strides, - const ptrdiff_t *y_strides) { - - size_t index = blockIdx.x * blockDim.x + threadIdx.x; - size_t total = N * C * y_d * y_h * y_w; - size_t n, c, od, oh, ow; - n = index / (C * y_d * y_h * y_w); - c = (index / (y_d * y_h * y_w)) % C; - od = (index / (y_h * y_w)) % y_d; - oh = (index / y_w) % y_h; - ow = index % y_w; - - if (index < total) { - size_t x_d_start = (od * x_d) / y_d; - size_t x_d_end = ((od + 1) * x_d + y_d - 1) / y_d; - size_t x_h_start = (oh * x_h) / y_h; - size_t x_h_end = ((oh + 1) * x_h + y_h - 1) / y_h; - size_t x_w_start = (ow * x_w) / y_w; - size_t x_w_end = ((ow + 1) * x_w + y_w - 1) / y_w; - - T sum = static_cast(0); - size_t count = (x_d_end - x_d_start) * (x_h_end - x_h_start) * (x_w_end - x_w_start); - for (size_t id = x_d_start; id < x_d_end; ++id) { - for (size_t ih = x_h_start; ih < x_h_end; ++ih) { - for (size_t iw = x_w_start; iw < x_w_end; ++iw) { - size_t x_index = n * x_strides[0] + c * x_strides[1] + id * x_strides[2] + ih * x_strides[3] + iw * x_strides[4]; - sum += x[x_index]; - } - } - } - size_t y_index = n * y_strides[0] + c * y_strides[1] + od * y_strides[2] + oh * y_strides[3] + ow * y_strides[4]; - y[y_index] = sum / static_cast(static_cast(count)); - } -} -#endif diff --git a/src/infiniop/ops/adaptive_avg_pool3d/metax/adaptive_avg_pool3d_metax.h b/src/infiniop/ops/adaptive_avg_pool3d/metax/adaptive_avg_pool3d_metax.h deleted file mode 100644 index 9a1c04384..000000000 --- a/src/infiniop/ops/adaptive_avg_pool3d/metax/adaptive_avg_pool3d_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef INFINIOP_OPS_ADAPTIVE_AVG_POOL3D_METAX_ADAPTIVE_AVG_POOL3D_METAX_H_ -#define INFINIOP_OPS_ADAPTIVE_AVG_POOL3D_METAX_ADAPTIVE_AVG_POOL3D_METAX_H_ - -#include "../adaptive_avg_pool3d.h" - -DESCRIPTOR(metax) - -#endif diff --git a/src/infiniop/ops/adaptive_avg_pool3d/metax/adaptive_avg_pool3d_metax.maca b/src/infiniop/ops/adaptive_avg_pool3d/metax/adaptive_avg_pool3d_metax.maca deleted file mode 100644 index 5e49830a2..000000000 --- a/src/infiniop/ops/adaptive_avg_pool3d/metax/adaptive_avg_pool3d_metax.maca +++ /dev/null @@ -1,88 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../cuda/kernel.cuh" -#include "adaptive_avg_pool3d_metax.h" - -template -INFINIOP_METAX_KERNEL -adaptiveAvgPool3D(T *y, const T *x, size_t N, size_t C, size_t x_d, size_t x_h, - size_t x_w, size_t y_d, size_t y_h, size_t y_w, - const ptrdiff_t *x_strides, const ptrdiff_t *y_strides) { - adaptiveAvgPool3DKernel(y, x, N, C, x_d, x_h, x_w, y_d, y_h, - y_w, x_strides, y_strides); -} - -namespace op::adaptive_avg_pool3d::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create(infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t *output_size) { - auto info = AdaptiveAvgPool3DInfo::create(y_desc, x_desc, output_size); - CHECK_RESULT(info); - size_t workspace_size = 10 * sizeof(ptrdiff_t); // for x_strides and y_strides - *desc_ptr = new Descriptor( - info.take(), workspace_size, - new Opaque{reinterpret_cast(handle)->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(void *y, const void *x, infiniDtype_t dtype, - size_t N, size_t C, size_t x_d, size_t x_h, - size_t x_w, size_t y_d, size_t y_h, size_t y_w, - const ptrdiff_t *x_strides, - const ptrdiff_t *y_strides, hcStream_t stream) { - size_t num_blocks = (N * C * y_d * y_h * y_w + BLOCK_SIZE - 1) / BLOCK_SIZE; - if (dtype == INFINI_DTYPE_F16) { - adaptiveAvgPool3D<<>>( - (half *)y, (const half *)x, N, C, x_d, x_h, x_w, y_d, y_h, y_w, - x_strides, y_strides); - } else if (dtype == INFINI_DTYPE_F32) { - adaptiveAvgPool3D<<>>( - (float *)y, (const float *)x, N, C, x_d, x_h, x_w, y_d, y_h, y_w, - x_strides, y_strides); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *y, const void *x, - void *stream_) const { - hcStream_t stream = (hcStream_t)stream_; - - // Prepare strides - ptrdiff_t *x_strides = (ptrdiff_t *)workspace; - ptrdiff_t *y_strides = (ptrdiff_t *)workspace + _info.x_strides.size(); - hcMemcpyAsync(x_strides, _info.x_strides.data(), - _info.x_strides.size() * sizeof(ptrdiff_t), - hcMemcpyHostToDevice, stream); - hcMemcpyAsync(y_strides, _info.y_strides.data(), - _info.y_strides.size() * sizeof(ptrdiff_t), - hcMemcpyHostToDevice, stream); - - if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.N, _info.C, _info.x_d, _info.x_h, _info.x_w, - _info.y_d, _info.y_h, _info.y_w, x_strides, y_strides, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.N, _info.C, _info.x_d, _info.x_h, _info.x_w, - _info.y_d, _info.y_h, _info.y_w, x_strides, y_strides, stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::adaptive_avg_pool3d::metax diff --git a/src/infiniop/ops/adaptive_avg_pool3d/moore/adaptive_avg_pool3d_kernel.h b/src/infiniop/ops/adaptive_avg_pool3d/moore/adaptive_avg_pool3d_kernel.h deleted file mode 100644 index f4eccade2..000000000 --- a/src/infiniop/ops/adaptive_avg_pool3d/moore/adaptive_avg_pool3d_kernel.h +++ /dev/null @@ -1,42 +0,0 @@ -#ifndef ADAPTIVE_AVG_POOL3D_KERNEL_H_ -#define ADAPTIVE_AVG_POOL3D_KERNEL_H_ - -template -__device__ void adaptiveAvgPool3DKernel(T *y, const T *x, size_t N, size_t C, - size_t x_d, size_t x_h, size_t x_w, - size_t y_d, size_t y_h, size_t y_w, - const ptrdiff_t *x_strides, - const ptrdiff_t *y_strides) { - - size_t index = blockIdx.x * blockDim.x + threadIdx.x; - size_t total = N * C * y_d * y_h * y_w; - size_t n, c, od, oh, ow; - n = index / (C * y_d * y_h * y_w); - c = (index / (y_d * y_h * y_w)) % C; - od = (index / (y_h * y_w)) % y_d; - oh = (index / y_w) % y_h; - ow = index % y_w; - - if (index < total) { - size_t x_d_start = (od * x_d) / y_d; - size_t x_d_end = ((od + 1) * x_d + y_d - 1) / y_d; - size_t x_h_start = (oh * x_h) / y_h; - size_t x_h_end = ((oh + 1) * x_h + y_h - 1) / y_h; - size_t x_w_start = (ow * x_w) / y_w; - size_t x_w_end = ((ow + 1) * x_w + y_w - 1) / y_w; - - T sum = static_cast(0); - size_t count = (x_d_end - x_d_start) * (x_h_end - x_h_start) * (x_w_end - x_w_start); - for (size_t id = x_d_start; id < x_d_end; ++id) { - for (size_t ih = x_h_start; ih < x_h_end; ++ih) { - for (size_t iw = x_w_start; iw < x_w_end; ++iw) { - size_t x_index = n * x_strides[0] + c * x_strides[1] + id * x_strides[2] + ih * x_strides[3] + iw * x_strides[4]; - sum += x[x_index]; - } - } - } - size_t y_index = n * y_strides[0] + c * y_strides[1] + od * y_strides[2] + oh * y_strides[3] + ow * y_strides[4]; - y[y_index] = sum / static_cast(count); - } -} -#endif diff --git a/src/infiniop/ops/adaptive_avg_pool3d/moore/adaptive_avg_pool3d_moore.h b/src/infiniop/ops/adaptive_avg_pool3d/moore/adaptive_avg_pool3d_moore.h deleted file mode 100644 index 37967bf16..000000000 --- a/src/infiniop/ops/adaptive_avg_pool3d/moore/adaptive_avg_pool3d_moore.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef ADAPTIVE_AVG_POOL3D_MOORE_H_ -#define ADAPTIVE_AVG_POOL3D_MOORE_H_ -#include "../adaptive_avg_pool3d.h" - -DESCRIPTOR(moore) - -#endif // ADAPTIVE_AVG_POOL3D_MOORE_H_ diff --git a/src/infiniop/ops/adaptive_avg_pool3d/moore/adaptive_avg_pool3d_moore.mu b/src/infiniop/ops/adaptive_avg_pool3d/moore/adaptive_avg_pool3d_moore.mu deleted file mode 100644 index 23f7f0dfa..000000000 --- a/src/infiniop/ops/adaptive_avg_pool3d/moore/adaptive_avg_pool3d_moore.mu +++ /dev/null @@ -1,89 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "adaptive_avg_pool3d_kernel.h" -#include "adaptive_avg_pool3d_moore.h" - -template -INFINIOP_MOORE_KERNEL -adaptiveAvgPool3D(T *y, const T *x, size_t N, size_t C, size_t x_d, size_t x_h, - size_t x_w, size_t y_d, size_t y_h, size_t y_w, - const ptrdiff_t *x_strides, const ptrdiff_t *y_strides) { - adaptiveAvgPool3DKernel(y, x, N, C, x_d, x_h, x_w, y_d, y_h, - y_w, x_strides, y_strides); -} -namespace op::adaptive_avg_pool3d::moore { -struct Descriptor::Opaque { - std::shared_ptr internal; -}; -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create(infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t *output_size) { - auto info = AdaptiveAvgPool3DInfo::create(y_desc, x_desc, output_size); - CHECK_RESULT(info); - size_t workspace_size = 10 * sizeof(ptrdiff_t); // for x_strides and y_strides - *desc_ptr = new Descriptor( - info.take(), workspace_size, - new Opaque{ - reinterpret_cast(handle)->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(void *y, const void *x, infiniDtype_t dtype, - size_t N, size_t C, size_t x_d, size_t x_h, - size_t x_w, size_t y_d, size_t y_h, size_t y_w, - const ptrdiff_t *x_strides, - const ptrdiff_t *y_strides, musaStream_t stream) { - size_t num_blocks = (N * C * y_d * y_h * y_w + BLOCK_SIZE - 1) / BLOCK_SIZE; - if (dtype == INFINI_DTYPE_F16) { - adaptiveAvgPool3D<<>>( - (half *)y, (const half *)x, N, C, x_d, x_h, x_w, y_d, y_h, y_w, - x_strides, y_strides); - } else if (dtype == INFINI_DTYPE_F32) { - adaptiveAvgPool3D<<>>( - (float *)y, (const float *)x, N, C, x_d, x_h, x_w, y_d, y_h, y_w, - x_strides, y_strides); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *y, const void *x, - void *stream_) const { - musaStream_t stream = (musaStream_t)stream_; - - // Prepare strides - ptrdiff_t *x_strides = (ptrdiff_t *)workspace; - ptrdiff_t *y_strides = (ptrdiff_t *)workspace + _info.x_strides.size(); - musaMemcpyAsync(x_strides, _info.x_strides.data(), - _info.x_strides.size() * sizeof(ptrdiff_t), - musaMemcpyHostToDevice, stream); - musaMemcpyAsync(y_strides, _info.y_strides.data(), - _info.y_strides.size() * sizeof(ptrdiff_t), - musaMemcpyHostToDevice, stream); - - if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.N, _info.C, _info.x_d, _info.x_h, _info.x_w, - _info.y_d, _info.y_h, _info.y_w, x_strides, y_strides, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.N, _info.C, _info.x_d, _info.x_h, _info.x_w, - _info.y_d, _info.y_h, _info.y_w, x_strides, y_strides, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_2048) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.N, _info.C, _info.x_d, _info.x_h, _info.x_w, - _info.y_d, _info.y_h, _info.y_w, x_strides, y_strides, stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} -} // namespace op::adaptive_avg_pool3d::moore diff --git a/src/infiniop/ops/adaptive_avg_pool3d/nvidia/adaptive_avg_pool3d_nvidia.cu b/src/infiniop/ops/adaptive_avg_pool3d/nvidia/adaptive_avg_pool3d_nvidia.cu deleted file mode 100644 index b6f3738b1..000000000 --- a/src/infiniop/ops/adaptive_avg_pool3d/nvidia/adaptive_avg_pool3d_nvidia.cu +++ /dev/null @@ -1,97 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../cuda/kernel.cuh" -#include "adaptive_avg_pool3d_nvidia.cuh" -#include - -template -INFINIOP_CUDA_KERNEL -adaptiveAvgPool3D(T *y, const T *x, size_t N, size_t C, size_t x_d, size_t x_h, - size_t x_w, size_t y_d, size_t y_h, size_t y_w, - const ptrdiff_t *x_strides, const ptrdiff_t *y_strides) { - adaptiveAvgPool3DKernel(y, x, N, C, x_d, x_h, x_w, y_d, y_h, - y_w, x_strides, y_strides); -} - -namespace op::adaptive_avg_pool3d::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create(infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t *output_size) { - auto info = AdaptiveAvgPool3DInfo::create(y_desc, x_desc, output_size); - CHECK_RESULT(info); - size_t workspace_size = 10 * sizeof(ptrdiff_t); // for x_strides and y_strides - *desc_ptr = new Descriptor( - info.take(), workspace_size, - new Opaque{ - reinterpret_cast(handle)->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(void *y, const void *x, infiniDtype_t dtype, - size_t N, size_t C, size_t x_d, size_t x_h, - size_t x_w, size_t y_d, size_t y_h, size_t y_w, - const ptrdiff_t *x_strides, - const ptrdiff_t *y_strides, cudaStream_t stream) { - size_t num_blocks = (N * C * y_d * y_h * y_w + BLOCK_SIZE - 1) / BLOCK_SIZE; - if (dtype == INFINI_DTYPE_F16) { - adaptiveAvgPool3D<<>>( - (half *)y, (const half *)x, N, C, x_d, x_h, x_w, y_d, y_h, y_w, - x_strides, y_strides); - } else if (dtype == INFINI_DTYPE_F32) { - adaptiveAvgPool3D<<>>( - (float *)y, (const float *)x, N, C, x_d, x_h, x_w, y_d, y_h, y_w, - x_strides, y_strides); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *y, const void *x, - void *stream_) const { - cudaStream_t stream = (cudaStream_t)stream_; - - // Prepare strides - ptrdiff_t *x_strides = (ptrdiff_t *)workspace; - ptrdiff_t *y_strides = (ptrdiff_t *)workspace + _info.x_strides.size(); - cudaMemcpyAsync(x_strides, _info.x_strides.data(), - _info.x_strides.size() * sizeof(ptrdiff_t), - cudaMemcpyHostToDevice, stream); - cudaMemcpyAsync(y_strides, _info.y_strides.data(), - _info.y_strides.size() * sizeof(ptrdiff_t), - cudaMemcpyHostToDevice, stream); - - if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.N, _info.C, _info.x_d, _info.x_h, _info.x_w, - _info.y_d, _info.y_h, _info.y_w, x_strides, y_strides, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.N, _info.C, _info.x_d, _info.x_h, _info.x_w, - _info.y_d, _info.y_h, _info.y_w, x_strides, y_strides, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_2048) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.N, _info.C, _info.x_d, _info.x_h, _info.x_w, - _info.y_d, _info.y_h, _info.y_w, x_strides, y_strides, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.N, _info.C, _info.x_d, _info.x_h, _info.x_w, - _info.y_d, _info.y_h, _info.y_w, x_strides, y_strides, stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} -} // namespace op::adaptive_avg_pool3d::nvidia diff --git a/src/infiniop/ops/adaptive_avg_pool3d/nvidia/adaptive_avg_pool3d_nvidia.cuh b/src/infiniop/ops/adaptive_avg_pool3d/nvidia/adaptive_avg_pool3d_nvidia.cuh deleted file mode 100644 index 09d2f5cc8..000000000 --- a/src/infiniop/ops/adaptive_avg_pool3d/nvidia/adaptive_avg_pool3d_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADAPTIVE_AVG_POOL3D_NVIDIA_H__ -#define __ADAPTIVE_AVG_POOL3D_NVIDIA_H__ - -#include "../adaptive_avg_pool3d.h" - -DESCRIPTOR(nvidia) - -#endif diff --git a/src/infiniop/ops/adaptive_avg_pool3d/operator.cc b/src/infiniop/ops/adaptive_avg_pool3d/operator.cc deleted file mode 100644 index 0b2036a6a..000000000 --- a/src/infiniop/ops/adaptive_avg_pool3d/operator.cc +++ /dev/null @@ -1,179 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infinicore.h" -#include "infiniop/ops/adaptive_avg_pool3d.h" - -#ifdef ENABLE_CPU_API -#include "cpu/adaptive_avg_pool3d_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/adaptive_avg_pool3d_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/adaptive_avg_pool3d_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/adaptive_avg_pool3d_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAdaptiveAvgPool3DDescriptor( - infiniopHandle_t handle, - infiniopAdaptiveAvgPool3DDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - size_t *output_size) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::adaptive_avg_pool3d::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y, \ - x, \ - output_size) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopGetAdaptiveAvgPool3DWorkspaceSize( - infiniopAdaptiveAvgPool3DDescriptor_t desc, - size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopAdaptiveAvgPool3D( - infiniopAdaptiveAvgPool3DDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, x, stream); - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopDestroyAdaptiveAvgPool3DDescriptor(infiniopAdaptiveAvgPool3DDescriptor_t desc) { -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} diff --git a/src/infiniop/ops/adaptive_max_pool1d/adaptive_max_pool1d.h b/src/infiniop/ops/adaptive_max_pool1d/adaptive_max_pool1d.h deleted file mode 100644 index 575c79894..000000000 --- a/src/infiniop/ops/adaptive_max_pool1d/adaptive_max_pool1d.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef ADAPTIVE_MAX_POOL1D_H -#define ADAPTIVE_MAX_POOL1D_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::adaptive_max_pool1d::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AdaptiveMaxPool1dInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - AdaptiveMaxPool1dInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t y_desc, \ - infiniopTensorDescriptor_t x_desc, \ - size_t output_size); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *y, \ - const void *x, \ - void *stream) const; \ - }; \ - } - -#endif // ADAPTIVE_MAX_POOL1D_H diff --git a/src/infiniop/ops/adaptive_max_pool1d/cpu/adaptive_max_pool1d_cpu.cc b/src/infiniop/ops/adaptive_max_pool1d/cpu/adaptive_max_pool1d_cpu.cc deleted file mode 100644 index 526d90415..000000000 --- a/src/infiniop/ops/adaptive_max_pool1d/cpu/adaptive_max_pool1d_cpu.cc +++ /dev/null @@ -1,94 +0,0 @@ -#include "adaptive_max_pool1d_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include "../../../reduce/cpu/reduce.h" -#include -#include - -namespace op::adaptive_max_pool1d::cpu { - -Descriptor::~Descriptor() {} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t output_size) { - auto result = AdaptiveMaxPool1dInfo::create(y_desc, x_desc, output_size); - CHECK_RESULT(result); - *desc_ptr = new Descriptor(nullptr, result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t adaptiveMaxPool1d(const AdaptiveMaxPool1dInfo *info, T *y, const T *x) { - - const size_t ndim = info->ndim(); - const size_t batch_size = info->shape[0]; - const size_t channels = ndim > 2 ? info->shape[1] : 1; - - const size_t input_length = info->input_length(); - const size_t output_length = info->output_length(); - - // 计算总的任务块数 (Batch * Channels) - const ptrdiff_t total_blocks = static_cast(batch_size * channels); - - const ptrdiff_t x_stride_last = info->x_strides.back(); - -#pragma omp parallel for - for (ptrdiff_t block_idx = 0; block_idx < total_blocks; ++block_idx) { - const size_t i = block_idx / channels; // batch index - const size_t j = block_idx % channels; // channel index - - const T *x_ptr_base; - T *y_ptr_base; - - if (ndim > 2) { // (N, C, L) - x_ptr_base = x + i * info->x_strides[0] + j * info->x_strides[1]; - y_ptr_base = y + i * info->y_strides[0] + j * info->y_strides[1]; - } else { // (N, L) - x_ptr_base = x + i * info->x_strides[0]; - y_ptr_base = y + i * info->y_strides[0]; - } - - for (size_t out_idx = 0; out_idx < output_length; ++out_idx) { - size_t start_index = (out_idx * input_length) / output_length; - size_t end_index = ((out_idx + 1) * input_length + output_length - 1) / output_length; - - start_index = std::max(start_index, size_t(0)); - end_index = std::min(end_index, input_length); - size_t window_len = end_index - start_index; - - if (window_len <= 0) { - continue; - } - - const T *window_ptr = x_ptr_base + start_index * x_stride_last; - - auto max_val = op::common_cpu::reduce_op::max(window_ptr, window_len, x_stride_last); - y_ptr_base[out_idx] = utils::cast(max_val); - } - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *y, const void *x, - void *stream) const { - - if (_info.atype == INFINI_DTYPE_F32) { - return adaptiveMaxPool1d(&_info, (float *)y, (const float *)x); - } else if (_info.atype == INFINI_DTYPE_F16) { - return adaptiveMaxPool1d(&_info, (fp16_t *)y, (const fp16_t *)x); - } else if (_info.atype == INFINI_DTYPE_BF16) { - return adaptiveMaxPool1d(&_info, (bf16_t *)y, (const bf16_t *)x); - } else if (_info.atype == INFINI_DTYPE_F64) { - return adaptiveMaxPool1d(&_info, (double *)y, (const double *)x); - } - - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -} // namespace op::adaptive_max_pool1d::cpu diff --git a/src/infiniop/ops/adaptive_max_pool1d/cpu/adaptive_max_pool1d_cpu.h b/src/infiniop/ops/adaptive_max_pool1d/cpu/adaptive_max_pool1d_cpu.h deleted file mode 100644 index e36c18100..000000000 --- a/src/infiniop/ops/adaptive_max_pool1d/cpu/adaptive_max_pool1d_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADAPTIVE_MAX_POOL1D_CPU_H__ -#define __ADAPTIVE_MAX_POOL1D_CPU_H__ - -#include "../adaptive_max_pool1d.h" - -DESCRIPTOR(cpu) - -#endif diff --git a/src/infiniop/ops/adaptive_max_pool1d/cuda/kernel.cuh b/src/infiniop/ops/adaptive_max_pool1d/cuda/kernel.cuh deleted file mode 100644 index 5112cf9ed..000000000 --- a/src/infiniop/ops/adaptive_max_pool1d/cuda/kernel.cuh +++ /dev/null @@ -1,54 +0,0 @@ -#ifndef __ADAPTIVE_MAX_POOL1D_CUDA_KERNEL_H__ -#define __ADAPTIVE_MAX_POOL1D_CUDA_KERNEL_H__ - -#include -#include - -template -__device__ void adaptiveMaxPool1dBlock( - Tdata *__restrict__ y, - ptrdiff_t stride_y_batch, - ptrdiff_t stride_y_channel, - const Tdata *__restrict__ x, - ptrdiff_t stride_x_batch, - ptrdiff_t stride_x_channel, - ptrdiff_t stride_x_length, - size_t channels, - size_t input_length, - size_t output_length, - size_t ndim) { - - size_t block_idx = blockIdx.x; - size_t batch_idx = block_idx / channels; - size_t channel_idx = block_idx % channels; - - const Tdata *x_ptr; - Tdata *y_ptr; - - if (ndim > 2) { - x_ptr = x + batch_idx * stride_x_batch + channel_idx * stride_x_channel; - y_ptr = y + batch_idx * stride_y_batch + channel_idx * stride_y_channel; - } else { - x_ptr = x + batch_idx * stride_x_batch; - y_ptr = y + batch_idx * stride_y_batch; - } - - for (size_t out_idx = threadIdx.x; out_idx < output_length; out_idx += BLOCK_SIZE) { - int start_index = static_cast(floorf((float)out_idx * input_length / output_length)); - int end_index = static_cast(ceilf((float)(out_idx + 1) * input_length / output_length)); - - if (end_index <= start_index) { - continue; - } - - Tcompute max_val = Tcompute(x_ptr[start_index * stride_x_length]); - for (int i = start_index + 1; i < end_index; ++i) { - Tcompute val = Tcompute(x_ptr[i * stride_x_length]); - max_val = max(max_val, val); - } - - y_ptr[out_idx] = Tdata(max_val); - } -} - -#endif diff --git a/src/infiniop/ops/adaptive_max_pool1d/info.h b/src/infiniop/ops/adaptive_max_pool1d/info.h deleted file mode 100644 index d8cd6a629..000000000 --- a/src/infiniop/ops/adaptive_max_pool1d/info.h +++ /dev/null @@ -1,65 +0,0 @@ -#ifndef __ADAPATIVE_MAX_POOL1D_H__ -#define __ADAPATIVE_MAX_POOL1D_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::adaptive_max_pool1d { - -class AdaptiveMaxPool1dInfo { - AdaptiveMaxPool1dInfo() = default; - -public: - infiniDtype_t atype; - std::vector shape; - std::vector y_strides; - std::vector x_strides; - size_t input_size; - size_t output_size; - size_t ndim() const { return shape.size(); } - size_t input_length() const { return input_size; } - size_t output_length() const { return output_size; } - - static utils::Result create( - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t output_size) { - - auto atype = y_desc->dtype(); - if (x_desc->dtype() != atype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - if (atype != INFINI_DTYPE_F16 && atype != INFINI_DTYPE_BF16 && atype != INFINI_DTYPE_F32 && atype != INFINI_DTYPE_F64) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - const size_t y_ndim = y_desc->ndim(); - const size_t x_ndim = x_desc->ndim(); - - if (y_ndim != x_ndim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - for (size_t i = 0; i < y_ndim - 1; ++i) { - if (x_desc->dim(i) != y_desc->dim(i)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - if (y_desc->dim(y_ndim - 1) != output_size) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - return utils::Result(AdaptiveMaxPool1dInfo{ - atype, - y_desc->shape(), - y_desc->strides(), - x_desc->strides(), - x_desc->dim(x_ndim - 1), - output_size}); - } -}; -} // namespace op::adaptive_max_pool1d - -#endif // __ADAPATIVE_MAX_POOL1D_H__ diff --git a/src/infiniop/ops/adaptive_max_pool1d/metax/adaptive_max_pool1d_metax.cuh b/src/infiniop/ops/adaptive_max_pool1d/metax/adaptive_max_pool1d_metax.cuh deleted file mode 100644 index 1a2305d19..000000000 --- a/src/infiniop/ops/adaptive_max_pool1d/metax/adaptive_max_pool1d_metax.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADAPTIVE_MAX_POOL1D_METAX_CUH__ -#define __ADAPTIVE_MAX_POOL1D_METAX_CUH__ - -#include "../adaptive_max_pool1d.h" - -DESCRIPTOR(metax) - -#endif diff --git a/src/infiniop/ops/adaptive_max_pool1d/metax/adaptive_max_pool1d_metax.maca b/src/infiniop/ops/adaptive_max_pool1d/metax/adaptive_max_pool1d_metax.maca deleted file mode 100644 index 7327b67aa..000000000 --- a/src/infiniop/ops/adaptive_max_pool1d/metax/adaptive_max_pool1d_metax.maca +++ /dev/null @@ -1,130 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "adaptive_max_pool1d_metax.cuh" - -#include "../../../devices/metax/metax_kernel_common.h" - -#include "../cuda/kernel.cuh" - -template -INFINIOP_METAX_KERNEL adaptiveMaxPool1dKernel( - Tdata *__restrict__ y, - ptrdiff_t stride_y_batch, - ptrdiff_t stride_y_channel, - const Tdata *__restrict__ x, - ptrdiff_t stride_x_batch, - ptrdiff_t stride_x_channel, - ptrdiff_t stride_x_length, - size_t channels, - size_t input_length, - size_t output_length, - size_t ndim) { - - adaptiveMaxPool1dBlock( - y, stride_y_batch, stride_y_channel, - x, stride_x_batch, stride_x_channel, stride_x_length, - channels, input_length, output_length, ndim); -} - -namespace op::adaptive_max_pool1d::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t output_size) { - - auto result = AdaptiveMaxPool1dInfo::create(y_desc, x_desc, output_size); - CHECK_RESULT(result); - auto info = result.take(); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - std::move(info), - 0, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel( - uint32_t numblock, - void *y, infiniDtype_t dtype, - ptrdiff_t stride_y_batch, ptrdiff_t stride_y_channel, - const void *x, - ptrdiff_t stride_x_batch, ptrdiff_t stride_x_channel, ptrdiff_t stride_x_length, - size_t channels, size_t input_length, size_t output_length, size_t ndim, - hcStream_t stream) { - -#define LAUNCH_KERNEL(Tdata, Tcompute) \ - adaptiveMaxPool1dKernel<<>>( \ - reinterpret_cast(y), \ - stride_y_batch, stride_y_channel, \ - reinterpret_cast(x), \ - stride_x_batch, stride_x_channel, stride_x_length, \ - channels, input_length, output_length, ndim) - - if (dtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(half, float); - } else if (dtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(__hpcc_bfloat16, float); - } else if (dtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(float, float); - } else if (dtype == INFINI_DTYPE_F64) { - LAUNCH_KERNEL(double, double); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -#undef LAUNCH_KERNEL - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *y, const void *x, - void *stream_) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - const size_t ndim = _info.ndim(); - const size_t batch_size = _info.shape[0]; - const size_t channels = ndim > 2 ? _info.shape[1] : 1; - const size_t input_length = _info.input_length(); - const size_t output_length = _info.output_length(); - - ptrdiff_t stride_x_batch = _info.x_strides[0]; - ptrdiff_t stride_x_channel = ndim > 2 ? _info.x_strides[1] : 0; - ptrdiff_t stride_x_length = _info.x_strides.back(); - - ptrdiff_t stride_y_batch = _info.y_strides[0]; - ptrdiff_t stride_y_channel = ndim > 2 ? _info.y_strides[1] : 0; - - uint32_t num_blocks = static_cast(batch_size * channels); - auto stream = reinterpret_cast(stream_); - - if (_opaque->internal->maxThreadsPerBlock() >= METAX_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - num_blocks, y, _info.atype, - stride_y_batch, stride_y_channel, - x, stride_x_batch, stride_x_channel, stride_x_length, - channels, input_length, output_length, ndim, - stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::adaptive_max_pool1d::metax diff --git a/src/infiniop/ops/adaptive_max_pool1d/moore/adaptive_max_pool1d_moore.h b/src/infiniop/ops/adaptive_max_pool1d/moore/adaptive_max_pool1d_moore.h deleted file mode 100644 index c56ad6fd4..000000000 --- a/src/infiniop/ops/adaptive_max_pool1d/moore/adaptive_max_pool1d_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADAPTIVE_MAX_POOL1D_MOOORE_H__ -#define __ADAPTIVE_MAX_POOL1D_MOOORE_H__ - -#include "../adaptive_max_pool1d.h" - -DESCRIPTOR(moore) - -#endif diff --git a/src/infiniop/ops/adaptive_max_pool1d/moore/adaptive_max_pool1d_moore.mu b/src/infiniop/ops/adaptive_max_pool1d/moore/adaptive_max_pool1d_moore.mu deleted file mode 100644 index ba2416a71..000000000 --- a/src/infiniop/ops/adaptive_max_pool1d/moore/adaptive_max_pool1d_moore.mu +++ /dev/null @@ -1,144 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "adaptive_max_pool1d_moore.h" - -#include "../../../devices/moore/moore_kernel_common.h" - -#include "../cuda/kernel.cuh" - -template -INFINIOP_MOORE_KERNEL adaptiveMaxPool1dKernel( - Tdata *__restrict__ y, - ptrdiff_t stride_y_batch, - ptrdiff_t stride_y_channel, - const Tdata *__restrict__ x, - ptrdiff_t stride_x_batch, - ptrdiff_t stride_x_channel, - ptrdiff_t stride_x_length, - size_t channels, - size_t input_length, - size_t output_length, - size_t ndim) { - - adaptiveMaxPool1dBlock( - y, stride_y_batch, stride_y_channel, - x, stride_x_batch, stride_x_channel, stride_x_length, - channels, input_length, output_length, ndim); -} - -namespace op::adaptive_max_pool1d::moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t output_size) { - auto result = AdaptiveMaxPool1dInfo::create(y_desc, x_desc, output_size); - CHECK_RESULT(result); - auto info = result.take(); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - std::move(info), - 0, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel( - uint32_t num_blocks, - void *y, infiniDtype_t dtype, - ptrdiff_t stride_y_batch, ptrdiff_t stride_y_channel, - const void *x, - ptrdiff_t stride_x_batch, ptrdiff_t stride_x_channel, ptrdiff_t stride_x_length, - size_t channels, size_t input_length, size_t output_length, size_t ndim, - musaStream_t musa_stream) { - -#define LAUNCH_KERNEL(Tdata, Tcompute) \ - adaptiveMaxPool1dKernel<<>>( \ - reinterpret_cast(y), \ - stride_y_batch, stride_y_channel, \ - reinterpret_cast(x), \ - stride_x_batch, stride_x_channel, stride_x_length, \ - channels, input_length, output_length, ndim) - - if (dtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(half, float); - } else if (dtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(__mt_bfloat16, float); - } else if (dtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(float, float); - } else if (dtype == INFINI_DTYPE_F64) { - LAUNCH_KERNEL(double, double); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_KERNEL - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *y, const void *x, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - const size_t ndim = _info.ndim(); - const size_t batch_size = _info.shape[0]; - const size_t channels = ndim > 2 ? _info.shape[1] : 1; - const size_t input_length = _info.input_length(); - const size_t output_length = _info.output_length(); - - ptrdiff_t stride_x_batch = _info.x_strides[0]; - ptrdiff_t stride_x_channel = ndim > 2 ? _info.x_strides[1] : 0; - ptrdiff_t stride_x_length = _info.x_strides.back(); - - ptrdiff_t stride_y_batch = _info.y_strides[0]; - ptrdiff_t stride_y_channel = ndim > 2 ? _info.y_strides[1] : 0; - - uint32_t num_blocks = static_cast(batch_size * channels); - auto musa_stream = reinterpret_cast(stream); - - if (_opaque->internal->maxThreadsPerBlock() >= MOORE_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - num_blocks, y, _info.atype, - stride_y_batch, stride_y_channel, - x, stride_x_batch, stride_x_channel, stride_x_length, - channels, input_length, output_length, ndim, - musa_stream)); - } else if (_opaque->internal->maxThreadsPerBlock() >= MOORE_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - num_blocks, y, _info.atype, - stride_y_batch, stride_y_channel, - x, stride_x_batch, stride_x_channel, stride_x_length, - channels, input_length, output_length, ndim, - musa_stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_2048) { - CHECK_STATUS(launchKernel( - num_blocks, y, _info.atype, - stride_y_batch, stride_y_channel, - x, stride_x_batch, stride_x_channel, stride_x_length, - channels, input_length, output_length, ndim, - musa_stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::adaptive_max_pool1d::moore diff --git a/src/infiniop/ops/adaptive_max_pool1d/nvidia/adaptive_max_pool1d_nvidia.cu b/src/infiniop/ops/adaptive_max_pool1d/nvidia/adaptive_max_pool1d_nvidia.cu deleted file mode 100644 index dc453a2d1..000000000 --- a/src/infiniop/ops/adaptive_max_pool1d/nvidia/adaptive_max_pool1d_nvidia.cu +++ /dev/null @@ -1,144 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "adaptive_max_pool1d_nvidia.cuh" - -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" - -#include "../cuda/kernel.cuh" - -template -INFINIOP_CUDA_KERNEL adaptiveMaxPool1dKernel( - Tdata *__restrict__ y, - ptrdiff_t stride_y_batch, - ptrdiff_t stride_y_channel, - const Tdata *__restrict__ x, - ptrdiff_t stride_x_batch, - ptrdiff_t stride_x_channel, - ptrdiff_t stride_x_length, - size_t channels, - size_t input_length, - size_t output_length, - size_t ndim) { - - adaptiveMaxPool1dBlock( - y, stride_y_batch, stride_y_channel, - x, stride_x_batch, stride_x_channel, stride_x_length, - channels, input_length, output_length, ndim); -} - -namespace op::adaptive_max_pool1d::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t output_size) { - auto result = AdaptiveMaxPool1dInfo::create(y_desc, x_desc, output_size); - CHECK_RESULT(result); - auto info = result.take(); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - std::move(info), - 0, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel( - uint32_t num_blocks, - void *y, infiniDtype_t dtype, - ptrdiff_t stride_y_batch, ptrdiff_t stride_y_channel, - const void *x, - ptrdiff_t stride_x_batch, ptrdiff_t stride_x_channel, ptrdiff_t stride_x_length, - size_t channels, size_t input_length, size_t output_length, size_t ndim, - cudaStream_t cuda_stream) { - -#define LAUNCH_KERNEL(Tdata, Tcompute) \ - adaptiveMaxPool1dKernel<<>>( \ - reinterpret_cast(y), \ - stride_y_batch, stride_y_channel, \ - reinterpret_cast(x), \ - stride_x_batch, stride_x_channel, stride_x_length, \ - channels, input_length, output_length, ndim) - - if (dtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(half, float); - } else if (dtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(__nv_bfloat16, float); - } else if (dtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(float, float); - } else if (dtype == INFINI_DTYPE_F64) { - LAUNCH_KERNEL(double, double); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_KERNEL - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *y, const void *x, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - const size_t ndim = _info.ndim(); - const size_t batch_size = _info.shape[0]; - const size_t channels = ndim > 2 ? _info.shape[1] : 1; - const size_t input_length = _info.input_length(); - const size_t output_length = _info.output_length(); - - ptrdiff_t stride_x_batch = _info.x_strides[0]; - ptrdiff_t stride_x_channel = ndim > 2 ? _info.x_strides[1] : 0; - ptrdiff_t stride_x_length = _info.x_strides.back(); - - ptrdiff_t stride_y_batch = _info.y_strides[0]; - ptrdiff_t stride_y_channel = ndim > 2 ? _info.y_strides[1] : 0; - - uint32_t num_blocks = static_cast(batch_size * channels); - auto cuda_stream = reinterpret_cast(stream); - - if (_opaque->internal->maxThreadsPerBlock() >= CUDA_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - num_blocks, y, _info.atype, - stride_y_batch, stride_y_channel, - x, stride_x_batch, stride_x_channel, stride_x_length, - channels, input_length, output_length, ndim, - cuda_stream)); - } else if (_opaque->internal->maxThreadsPerBlock() >= CUDA_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - num_blocks, y, _info.atype, - stride_y_batch, stride_y_channel, - x, stride_x_batch, stride_x_channel, stride_x_length, - channels, input_length, output_length, ndim, - cuda_stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) { - CHECK_STATUS(launchKernel( - num_blocks, y, _info.atype, - stride_y_batch, stride_y_channel, - x, stride_x_batch, stride_x_channel, stride_x_length, - channels, input_length, output_length, ndim, - cuda_stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::adaptive_max_pool1d::nvidia diff --git a/src/infiniop/ops/adaptive_max_pool1d/nvidia/adaptive_max_pool1d_nvidia.cuh b/src/infiniop/ops/adaptive_max_pool1d/nvidia/adaptive_max_pool1d_nvidia.cuh deleted file mode 100644 index 3675fabb6..000000000 --- a/src/infiniop/ops/adaptive_max_pool1d/nvidia/adaptive_max_pool1d_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADAPTIVE_MAX_POOL1D_CUDA_H__ -#define __ADAPTIVE_MAX_POOL1D_CUDA_H__ - -#include "../adaptive_max_pool1d.h" - -DESCRIPTOR(nvidia) - -#endif diff --git a/src/infiniop/ops/adaptive_max_pool1d/operator.cc b/src/infiniop/ops/adaptive_max_pool1d/operator.cc deleted file mode 100644 index 52d9e805f..000000000 --- a/src/infiniop/ops/adaptive_max_pool1d/operator.cc +++ /dev/null @@ -1,159 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/adaptive_max_pool1d.h" - -#ifdef ENABLE_CPU_API -#include "cpu/adaptive_max_pool1d_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) -#include "nvidia/adaptive_max_pool1d_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/adaptive_max_pool1d_metax.cuh" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/adaptive_max_pool1d_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAdaptiveMaxPool1dDescriptor( - infiniopHandle_t handle, - infiniopAdaptiveMaxPool1dDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t output_size) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::adaptive_max_pool1d::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - x_desc, \ - output_size) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - } -#undef CREATE - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopGetAdaptiveMaxPool1dWorkspaceSize( - infiniopAdaptiveMaxPool1dDescriptor_t desc, - size_t *size) { -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopAdaptiveMaxPool1d( - infiniopAdaptiveMaxPool1dDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc)->calculate( \ - workspace, workspace_size, y, x, stream); - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - } -#undef CALCULATE - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopDestroyAdaptiveMaxPool1dDescriptor( - infiniopAdaptiveMaxPool1dDescriptor_t desc) { -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DESTROY(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - DESTROY(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DESTROY(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DESTROY(INFINI_DEVICE_HYGON, nvidia); -#endif - } -#undef DESTROY - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} diff --git a/src/infiniop/ops/add/add.h b/src/infiniop/ops/add/add.h deleted file mode 100644 index c7caa6b57..000000000 --- a/src/infiniop/ops/add/add.h +++ /dev/null @@ -1,46 +0,0 @@ -#ifndef ADD_H -#define ADD_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::add::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AddInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - AddInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t c_desc, \ - std::vector input_descs); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *c, \ - std::vector inputs, \ - void *stream) const; \ - }; \ - } - -#endif // ADD_H diff --git a/src/infiniop/ops/add/ascend/add_ascend.cc b/src/infiniop/ops/add/ascend/add_ascend.cc deleted file mode 100644 index 1a9465f5f..000000000 --- a/src/infiniop/ops/add/ascend/add_ascend.cc +++ /dev/null @@ -1,143 +0,0 @@ -#include "add_ascend.h" -#include "../../../devices/ascend/common_ascend.h" -#include -#include -#include - -namespace op::add::ascend { - -// Opaque structure must be defined AFTER the class declaration (which is in add.h via DESCRIPTOR macro) -struct Descriptor::Opaque { - aclnnTensorDescriptor_t a; - aclnnTensorDescriptor_t b; - aclnnTensorDescriptor_t c; - aclnnScalarDescriptor_t alpha; - void *alpha_value; - size_t workspaceSize; - aclOpExecutor *executor; - - Opaque(aclnnTensorDescriptor_t a_, aclnnTensorDescriptor_t b_, aclnnTensorDescriptor_t c_, - aclnnScalarDescriptor_t alpha_, void *alpha_value_, size_t ws, aclOpExecutor *exec) - : a(a_), b(b_), c(c_), alpha(alpha_), alpha_value(alpha_value_), workspaceSize(ws), executor(exec) {} - - ~Opaque() { - delete a; - delete b; - delete c; - delete alpha; - std::free(alpha_value); - aclDestroyAclOpExecutor(executor); - } -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t c_desc, - std::vector input_descs) { - - if (input_descs.size() != 2) { - return INFINI_STATUS_BAD_PARAM; - } - - auto a_desc = input_descs[0]; - auto b_desc = input_descs[1]; - - // Create AddInfo first - auto result = AddInfo::create(c_desc, a_desc, b_desc); - CHECK_RESULT(result); - - auto handle_ascend = reinterpret_cast(handle); - - aclnnTensorDescriptor_t a = new aclnnTensorDescriptor(a_desc); - aclnnTensorDescriptor_t b = new aclnnTensorDescriptor(b_desc); - aclnnTensorDescriptor_t c = new aclnnTensorDescriptor(c_desc); - - void *alpha_value = nullptr; - size_t alpha_value_size = 0; - infiniDtype_t alpha_dtype = INFINI_DTYPE_F32; - -#define SET_ALPHA(TYPE, DTYPE, VALUE) \ - do { \ - alpha_value_size = sizeof(TYPE); \ - alpha_value = std::malloc(alpha_value_size); \ - if (alpha_value == nullptr) { \ - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; \ - } \ - *static_cast(alpha_value) = (TYPE)(VALUE); \ - alpha_dtype = DTYPE; \ - } while (0) - - switch (c_desc->dtype()) { - case INFINI_DTYPE_I32: - SET_ALPHA(int32_t, INFINI_DTYPE_I32, 1); - break; - case INFINI_DTYPE_I64: - SET_ALPHA(int64_t, INFINI_DTYPE_I64, 1); - break; - default: - SET_ALPHA(float, INFINI_DTYPE_F32, 1.0f); - break; - } - -#undef SET_ALPHA - - aclnnScalarDescriptor_t alpha = new aclnnScalarDescriptor( - alpha_dtype, alpha_value, alpha_value_size); - - size_t workspace_size = 0; - aclOpExecutor *executor = nullptr; - - CHECK_ACL(aclnnAddGetWorkspaceSize( - a->tensor, - b->tensor, - alpha->scalar, - c->tensor, - &workspace_size, - &executor)); - - aclSetAclOpExecutorRepeatable(executor); - - *desc_ptr = new Descriptor( - new Opaque{a, b, c, alpha, alpha_value, workspace_size, executor}, - result.take(), - workspace_size, - handle_ascend->device, - handle_ascend->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *c, std::vector inputs, - void *stream) const { - - if (inputs.size() != 2) { - return INFINI_STATUS_BAD_PARAM; - } - - if (workspace_size < workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // Set input/output tensor addresses in the executor - // Parameters: executor, tensor_index, tensor_descriptor, data_pointer - AclSetTensorAddr(_opaque->executor, 0, _opaque->a->tensor, const_cast(inputs[0])); - AclSetTensorAddr(_opaque->executor, 1, _opaque->b->tensor, const_cast(inputs[1])); - AclSetTensorAddr(_opaque->executor, 2, _opaque->c->tensor, c); - - CHECK_ACL(aclnnAdd( - workspace, - workspace_size, - _opaque->executor, - stream)); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::add::ascend diff --git a/src/infiniop/ops/add/ascend/add_ascend.h b/src/infiniop/ops/add/ascend/add_ascend.h deleted file mode 100644 index f63831498..000000000 --- a/src/infiniop/ops/add/ascend/add_ascend.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADD_ASCEND_H__ -#define __ADD_ASCEND_H__ - -#include "../add.h" - -DESCRIPTOR(ascend) - -#endif // __ADD_ASCEND_H__ diff --git a/src/infiniop/ops/add/bang/add_bang.h b/src/infiniop/ops/add/bang/add_bang.h deleted file mode 100644 index 054d8961e..000000000 --- a/src/infiniop/ops/add/bang/add_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADD_BANG_API_H__ -#define __ADD_BANG_API_H__ - -#include "../../../elementwise/bang/elementwise_bang.h" - -ELEMENTWISE_DESCRIPTOR(add, bang) - -#endif // __ADD_BANG_API_H__ diff --git a/src/infiniop/ops/add/bang/add_bang.mlu b/src/infiniop/ops/add/bang/add_bang.mlu deleted file mode 100644 index 4925a604c..000000000 --- a/src/infiniop/ops/add/bang/add_bang.mlu +++ /dev/null @@ -1,72 +0,0 @@ -#include "add_bang.h" - -// Operator Interface Declaration -LAUNCH_ELEMENTWISE_KERNEL(Add) - -namespace op::add::bang { - -typedef struct AddOp { - static constexpr size_t num_inputs = 2; - template - static infiniStatus_t launch(Args... args) { - launchAddKernel(args...); - return INFINI_STATUS_SUCCESS; - } -} AddOp; - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create Bang elementwise descriptor - CREATE_ELEMENTWISE_BANG_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *queue) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, workspace, output, inputs, queue); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, workspace, output, inputs, queue); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, workspace, output, inputs, queue); - case INFINI_DTYPE_I32: - return _device_info->calculate(_info, workspace, output, inputs, queue); - case INFINI_DTYPE_I64: - return _device_info->calculate(_info, workspace, output, inputs, queue); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::add::bang diff --git a/src/infiniop/ops/add/bang/add_bang_internal.mlu b/src/infiniop/ops/add/bang/add_bang_internal.mlu deleted file mode 100644 index a6d2d24cb..000000000 --- a/src/infiniop/ops/add/bang/add_bang_internal.mlu +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef __ADD_BANG_INTERNAL_H__ -#define __ADD_BANG_INTERNAL_H__ - -#include "../../../elementwise/bang/elementwise_bang_kernel.h" - -typedef struct AddOp { -public: - static constexpr size_t num_inputs = 2; - template - __mlu_device__ void operator()(T *out, const T *a, const T *b, size_t num_elements) const { - if constexpr (std::is_same_v || std::is_same_v || std::is_same_v || std::is_same_v || std::is_same_v) { - __bang_add(out, a, b, num_elements); - } else { - out = a + b; - } - } -} AddOp; - -LAUNCH_ELEMENTWISE_KERNEL_IMPL(Add, AddOp) - -LAUNCH_ELEMENTWISE_KERNEL_INSTANTIATE(Add, half) -LAUNCH_ELEMENTWISE_KERNEL_INSTANTIATE(Add, bfloat16_t) -LAUNCH_ELEMENTWISE_KERNEL_INSTANTIATE(Add, float) -LAUNCH_ELEMENTWISE_KERNEL_INSTANTIATE(Add, int32_t) -LAUNCH_ELEMENTWISE_KERNEL_INSTANTIATE(Add, int64_t) - -#endif // __ADD_BANG_INTERNAL_H__ diff --git a/src/infiniop/ops/add/cpu/add_cpu.cc b/src/infiniop/ops/add/cpu/add_cpu.cc deleted file mode 100644 index 3c03c6d7f..000000000 --- a/src/infiniop/ops/add/cpu/add_cpu.cc +++ /dev/null @@ -1,58 +0,0 @@ -#include "add_cpu.h" - -namespace op::add::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create CPU elementwise descriptor - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::add::cpu diff --git a/src/infiniop/ops/add/cpu/add_cpu.h b/src/infiniop/ops/add/cpu/add_cpu.h deleted file mode 100644 index 736468acf..000000000 --- a/src/infiniop/ops/add/cpu/add_cpu.h +++ /dev/null @@ -1,19 +0,0 @@ -#ifndef __ADD_CPU_H__ -#define __ADD_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" - -ELEMENTWISE_DESCRIPTOR(add, cpu) - -namespace op::add::cpu { -typedef struct AddOp { -public: - static constexpr size_t num_inputs = 2; - template - T operator()(const T &a, const T &b) const { - return a + b; - } -} AddOp; -} // namespace op::add::cpu - -#endif // __ADD_CPU_H__ diff --git a/src/infiniop/ops/add/cuda/kernel.cuh b/src/infiniop/ops/add/cuda/kernel.cuh deleted file mode 100644 index 14abc2e51..000000000 --- a/src/infiniop/ops/add/cuda/kernel.cuh +++ /dev/null @@ -1,23 +0,0 @@ -#ifndef __ADD_CUDA_H__ -#define __ADD_CUDA_H__ - -namespace op::add::cuda { -typedef struct AddOp { -public: - static constexpr size_t num_inputs = 2; - template - __device__ __forceinline__ T operator()(const T &a, const T &b) const { - if constexpr (std::is_same_v) { - return __hadd2(a, b); - } else if constexpr (std::is_same_v || std::is_same_v) { - return __hadd(a, b); - } else if constexpr (std::is_same_v) { - return __fadd_rd(a, b); - } else { - return a + b; - } - } -} AddOp; -} // namespace op::add::cuda - -#endif // __ADD_CUDA_H__ diff --git a/src/infiniop/ops/add/info.h b/src/infiniop/ops/add/info.h deleted file mode 100644 index ca2736222..000000000 --- a/src/infiniop/ops/add/info.h +++ /dev/null @@ -1,59 +0,0 @@ -#ifndef __ADD_INFO_H__ -#define __ADD_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::add { - -class AddInfo { - AddInfo() = default; - -public: - infiniDtype_t dtype; - std::vector shape; - size_t numel; - - static utils::Result create( - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - - auto dtype = c_desc->dtype(); - - // Check dtype compatibility - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - // Check shape compatibility (broadcast) - auto c_shape = c_desc->shape(); - auto a_shape = a_desc->shape(); - auto b_shape = b_desc->shape(); - - auto c_ndim = c_desc->ndim(); - - // Require same ndim and shape for now - if (c_ndim != a_desc->ndim() || c_ndim != b_desc->ndim()) { - CHECK_STATUS(INFINI_STATUS_BAD_TENSOR_SHAPE); - } - - CHECK_SAME_SHAPE(c_shape, a_shape); - CHECK_SAME_SHAPE(c_shape, b_shape); - - size_t numel = 1; - std::vector shape; - for (size_t i = 0; i < c_ndim; i++) { - shape.push_back(c_shape[i]); - numel *= c_shape[i]; - } - - return utils::Result(AddInfo{ - dtype, - shape, - numel}); - } -}; - -} // namespace op::add - -#endif // __ADD_INFO_H__ diff --git a/src/infiniop/ops/add/kunlun/add_kunlun.h b/src/infiniop/ops/add/kunlun/add_kunlun.h deleted file mode 100644 index 877365199..000000000 --- a/src/infiniop/ops/add/kunlun/add_kunlun.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADD_KUNLUN_API_H__ -#define __ADD_KUNLUN_API_H__ - -#include "../../../elementwise/kunlun/elementwise_kunlun_api.h" - -ELEMENTWISE_DESCRIPTOR(add, kunlun) - -#endif // __ADD_KUNLUN_API_H__ diff --git a/src/infiniop/ops/add/kunlun/add_kunlun.xpu b/src/infiniop/ops/add/kunlun/add_kunlun.xpu deleted file mode 100644 index a27f065ed..000000000 --- a/src/infiniop/ops/add/kunlun/add_kunlun.xpu +++ /dev/null @@ -1,62 +0,0 @@ -#include "../../../elementwise/kunlun/elementwise_kunlun.h" -#include "add_kunlun.h" -#include "kernel.h" - -namespace op::add::kunlun { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create KUNLUN elementwise descriptor - CREATE_ELEMENTWISE_KUNLUN_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<8, AddOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<8, AddOp, bfloat16_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<8, AddOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<8, AddOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<8, AddOp, int64_t>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::add::kunlun diff --git a/src/infiniop/ops/add/kunlun/kernel.h b/src/infiniop/ops/add/kunlun/kernel.h deleted file mode 100644 index cabe37c2e..000000000 --- a/src/infiniop/ops/add/kunlun/kernel.h +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef __ADD_KUNLUN_KERNEL_H__ -#define __ADD_KUNLUN_KERNEL_H__ - -namespace op::add::kunlun { - -typedef struct AddOp { -public: - static constexpr int num_inputs = 2; - template - inline __device__ T operator()(const T *inputs) const { - T a = inputs[0]; - T b = inputs[1]; - return a + b; - } - // bfloat16 - cast to flloat - inline __device__ bfloat16_t operator()(const bfloat16_t *inputs) const { - float a_f = __bfloat162float(inputs[0]); - float b_f = __bfloat162float(inputs[1]); - return __float2bfloat16(a_f + b_f); - } - inline __device__ int64_t operator()(const int64_t *inputs) const { - return inputs[0] + inputs[1]; - } -} AddOp; -} // namespace op::add::kunlun - -#endif // __ADD_KUNLUN_KERNEL_H__ diff --git a/src/infiniop/ops/add/metax/add_metax.h b/src/infiniop/ops/add/metax/add_metax.h deleted file mode 100644 index 5a0ad5fd9..000000000 --- a/src/infiniop/ops/add/metax/add_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADD_METAX_API_H__ -#define __ADD_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(add, metax) - -#endif // __ADD_METAX_API_H__ diff --git a/src/infiniop/ops/add/metax/add_metax.maca b/src/infiniop/ops/add/metax/add_metax.maca deleted file mode 100644 index 2cc7de418..000000000 --- a/src/infiniop/ops/add/metax/add_metax.maca +++ /dev/null @@ -1,66 +0,0 @@ -#include "add_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::add::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create CUDA elementwise descriptor - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::AddOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::AddOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::AddOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::AddOp, double>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, cuda::AddOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, cuda::AddOp, int64_t>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::add::metax diff --git a/src/infiniop/ops/add/moore/add_moore.h b/src/infiniop/ops/add/moore/add_moore.h deleted file mode 100644 index db774c252..000000000 --- a/src/infiniop/ops/add/moore/add_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADD_MOORE_API_H__ -#define __ADD_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(add, moore) - -#endif // __ADD_MOORE_API_H__ diff --git a/src/infiniop/ops/add/moore/add_moore.mu b/src/infiniop/ops/add/moore/add_moore.mu deleted file mode 100644 index 84df6bcb8..000000000 --- a/src/infiniop/ops/add/moore/add_moore.mu +++ /dev/null @@ -1,66 +0,0 @@ -#include "add_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "add_moore_kernel.h" - -namespace op::add::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create MOORE elementwise descriptor - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, moore::AddOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, moore::AddOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, moore::AddOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, moore::AddOp, double>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, moore::AddOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, moore::AddOp, int64_t>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::add::moore diff --git a/src/infiniop/ops/add/moore/add_moore_kernel.h b/src/infiniop/ops/add/moore/add_moore_kernel.h deleted file mode 100644 index 9957e5d03..000000000 --- a/src/infiniop/ops/add/moore/add_moore_kernel.h +++ /dev/null @@ -1,38 +0,0 @@ -#ifndef __ADD_MOORE_KERNEL_H__ -#define __ADD_MOORE_KERNEL_H__ - -/* - * This file contains the Add operation implementation for the MUSA backend. - * - * It uses the 'op::add::cuda' namespace to maintain a consistent code structure - * and interface with the CUDA implementation, ensuring code alignment across different - * hardware platforms. - */ - -namespace op::add::moore { -typedef struct AddOp { -public: - static constexpr size_t num_inputs = 2; - template - __device__ __forceinline__ T operator()(const T &a, const T &b) const { - if constexpr (std::is_same_v) { - return __hadd2(a, b); - } else if constexpr (std::is_same_v) { - return __hadd(a, b); - } else if constexpr (std::is_same_v) { - // On MUSA platform, convert to float, add, then convert back to avoid ambiguous conversion - // from int (returned by __hadd) to __mt_bfloat16 - float a_f = __bfloat162float(a); - float b_f = __bfloat162float(b); - return __float2bfloat16_rn(a_f + b_f); - } else if constexpr (std::is_same_v) { - // Use __fadd_rn instead of __fadd_rd for moore platform compatibility - return __fadd_rn(a, b); - } else { - return a + b; - } - } -} AddOp; -} // namespace op::add::moore - -#endif // __ADD_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/add/nvidia/add_nvidia.cu b/src/infiniop/ops/add/nvidia/add_nvidia.cu deleted file mode 100644 index 787f45a0e..000000000 --- a/src/infiniop/ops/add/nvidia/add_nvidia.cu +++ /dev/null @@ -1,156 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "add_nvidia.cuh" - -namespace op::add::nvidia { -namespace { - -template -INFINIOP_CUDA_KERNEL addKernel( - size_t output_size, - size_t ndim, - bool output_contiguous, - const bool *__restrict__ input_contiguous, - const bool *__restrict__ input_broadcasted, - const size_t *__restrict__ output_shape, - const size_t *__restrict__ input_shapes, - const ptrdiff_t *__restrict__ output_strides, - const ptrdiff_t *__restrict__ input_strides, - T *output, - const T *__restrict__ a, - const T *__restrict__ b, - size_t offset) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; - if (idx >= output_size) { - return; - } - - size_t out_idx = op::elementwise::nvidia::getOutputIndex(idx, output_contiguous, ndim, output_shape, output_strides); - op::elementwise::nvidia::InputIndexer indexer{ - idx, ndim, input_contiguous, input_broadcasted, input_shapes, input_strides, output_strides}; - output[out_idx] = cuda::AddOp{}(a[indexer(0)], b[indexer(1)]); -} - -template -infiniStatus_t launchAddKernel( - const op::elementwise::ElementwiseInfo &info, - const std::shared_ptr &internal, - void *workspace, - void *output, - const void *a, - const void *b, - cudaStream_t stream) { - - auto output_size = info.getOutputSize(); - if (output_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - auto ndim = info.getNdim(); - auto *d_meta_start = reinterpret_cast(workspace); - CHECK_CUDA(cudaMemcpyAsync(d_meta_start, info.getMetaStart(), info.getMetaMemSize(), cudaMemcpyHostToDevice, stream)); - - auto *d_output_shape = reinterpret_cast(d_meta_start); - auto *d_output_strides = reinterpret_cast(d_output_shape + ndim); - auto *d_input_shapes = reinterpret_cast(d_output_strides + ndim); - auto *d_input_strides = reinterpret_cast(d_input_shapes + info.getInputSize() * ndim); - auto *d_input_contiguous = reinterpret_cast(d_input_strides + info.getInputSize() * ndim); - auto *d_input_broadcasted = reinterpret_cast(d_input_contiguous + info.getInputSize()); - - dim3 block_dims(std::min(256U, static_cast(internal->maxThreadsPerBlock()))); - dim3 grid_dims(std::min(uint32_t(CEIL_DIV(output_size, block_dims.x)), static_cast(internal->gridSizeX()))); - size_t step = grid_dims.x * block_dims.x; - - for (size_t i = 0; i < output_size; i += step) { - addKernel<<>>( - output_size, - ndim, - info.isOutputContiguous(), - d_input_contiguous, - d_input_broadcasted, - d_output_shape, - d_input_shapes, - d_output_strides, - d_input_strides, - reinterpret_cast(output), - reinterpret_cast(a), - reinterpret_cast(b), - i); - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - auto info_result = op::elementwise::ElementwiseInfo::create(out_desc, input_desc_vec); - CHECK_RESULT(info_result); - auto info = info_result.take(); - auto workspace_size = info.getMetaMemSize(); - - *desc_ptr = new Descriptor( - dtype, - std::move(info), - handle->internal(), - workspace_size, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *a, - const void *b, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return launchAddKernel(_info, _internal, workspace, output, a, b, reinterpret_cast(stream)); - case INFINI_DTYPE_BF16: - return launchAddKernel(_info, _internal, workspace, output, a, b, reinterpret_cast(stream)); - case INFINI_DTYPE_F32: - return launchAddKernel(_info, _internal, workspace, output, a, b, reinterpret_cast(stream)); - case INFINI_DTYPE_I32: - return launchAddKernel(_info, _internal, workspace, output, a, b, reinterpret_cast(stream)); - case INFINI_DTYPE_I64: - return launchAddKernel(_info, _internal, workspace, output, a, b, reinterpret_cast(stream)); - case INFINI_DTYPE_F64: - return launchAddKernel(_info, _internal, workspace, output, a, b, reinterpret_cast(stream)); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::add::nvidia diff --git a/src/infiniop/ops/add/nvidia/add_nvidia.cuh b/src/infiniop/ops/add/nvidia/add_nvidia.cuh deleted file mode 100644 index 1abc71b4f..000000000 --- a/src/infiniop/ops/add/nvidia/add_nvidia.cuh +++ /dev/null @@ -1,48 +0,0 @@ -#ifndef __ADD_CUDA_API_H__ -#define __ADD_CUDA_API_H__ - -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -namespace op::add::nvidia { -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - op::elementwise::ElementwiseInfo _info; - std::shared_ptr _internal; - size_t _workspace_size; - - Descriptor( - infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - std::shared_ptr internal, - size_t workspace_size, - infiniDevice_t device_type, - int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)), - _internal(std::move(internal)), - _workspace_size(workspace_size) {} - -public: - ~Descriptor(); - - size_t workspaceSize() const { return _workspace_size; } - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - std::vector input_descs); - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *a, - const void *b, - void *stream) const; -}; -} // namespace op::add::nvidia - -#endif // __ADD_CUDA_API_H__ diff --git a/src/infiniop/ops/add/operator.cc b/src/infiniop/ops/add/operator.cc deleted file mode 100644 index 9546a7577..000000000 --- a/src/infiniop/ops/add/operator.cc +++ /dev/null @@ -1,280 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/add.h" - -#ifdef ENABLE_CPU_API -#include "cpu/add_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) -#include "nvidia/add_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/add_metax.h" -#endif -#ifdef ENABLE_KUNLUN_API -#include "kunlun/add_kunlun.h" -#endif -#ifdef ENABLE_CAMBRICON_API -#include "bang/add_bang.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/add_moore.h" -#endif -#ifdef ENABLE_ASCEND_API -#include "ascend/add_ascend.h" -#endif -#include - -__INFINI_C infiniStatus_t infiniopCreateAddDescriptor( - infiniopHandle_t handle, - infiniopAddDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::add::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - c_desc, \ - {a_desc, \ - b_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_KUNLUN_API - CREATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ASCEND_API - CREATE(INFINI_DEVICE_ASCEND, ascend); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAddWorkspaceSize(infiniopAddDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - GET(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_KUNLUN_API - GET(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_CAMBRICON_API - GET(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ASCEND_API - GET(INFINI_DEVICE_ASCEND, ascend); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -namespace { - -template -infiniStatus_t calculateAdd( - const Descriptor *desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream) { - const std::vector inputs{a, b}; - return desc->calculate(workspace, workspace_size, c, inputs, stream); -} - -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) -infiniStatus_t calculateAdd( - const op::add::nvidia::Descriptor *desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream) { - return desc->calculate(workspace, workspace_size, c, a, b, stream); -} -#endif - -} // namespace - -__INFINI_C infiniStatus_t infiniopAdd( - infiniopAddDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream) { - -// NVIDIA Add keeps explicit a/b pointers because the generic elementwise -// input-vector path copies inputs.data() from host to device workspace before -// launching the kernel. During CUDA graph capture, that H2D node records the -// host source address; if infiniopAdd used a temporary vector such as {a, b}, -// graph replay could later read from an invalid host address and copy bad input -// pointers into device workspace. Other backends keep their original vector -// interface through calculateAdd's default forwarding path. -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return calculateAdd( \ - reinterpret_cast(desc), \ - workspace, workspace_size, c, a, b, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_KUNLUN_API - CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ASCEND_API - CALCULATE(INFINI_DEVICE_ASCEND, ascend); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyAddDescriptor(infiniopAddDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - DELETE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_KUNLUN_API - DELETE(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_CAMBRICON_API - DELETE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ASCEND_API - DELETE(INFINI_DEVICE_ASCEND, ascend); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/add_rms_norm/add_rms_norm.h b/src/infiniop/ops/add_rms_norm/add_rms_norm.h deleted file mode 100644 index 76451e982..000000000 --- a/src/infiniop/ops/add_rms_norm/add_rms_norm.h +++ /dev/null @@ -1,53 +0,0 @@ -#ifndef ADD_RMS_NORM_H -#define ADD_RMS_NORM_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::add_rms_norm::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AddRMSNormInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - AddRMSNormInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t y_desc, \ - infiniopTensorDescriptor_t residual_out_desc, \ - infiniopTensorDescriptor_t a_desc, \ - infiniopTensorDescriptor_t b_desc, \ - infiniopTensorDescriptor_t weight_desc, \ - float epsilon); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *y, \ - void *residual_out, \ - const void *a, \ - const void *b, \ - const void *weight, \ - void *stream) const; \ - }; \ - } - -#endif // ADD_RMS_NORM_H diff --git a/src/infiniop/ops/add_rms_norm/bang/add_rms_norm_bang.h b/src/infiniop/ops/add_rms_norm/bang/add_rms_norm_bang.h deleted file mode 100644 index 0fdc89718..000000000 --- a/src/infiniop/ops/add_rms_norm/bang/add_rms_norm_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADD_RMS_NORM_BANG_H__ -#define __ADD_RMS_NORM_BANG_H__ - -#include "../add_rms_norm.h" - -DESCRIPTOR(bang) - -#endif // __ADD_RMS_NORM_BANG_H__ diff --git a/src/infiniop/ops/add_rms_norm/bang/add_rms_norm_bang.mlu b/src/infiniop/ops/add_rms_norm/bang/add_rms_norm_bang.mlu deleted file mode 100644 index aa36a29bf..000000000 --- a/src/infiniop/ops/add_rms_norm/bang/add_rms_norm_bang.mlu +++ /dev/null @@ -1,480 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "../../../reduce/bang/reduce_bang.h" -#include "add_rms_norm_bang.h" - -#include -#include - -__nram__ char nram_buffer[NRAM_MAX_SIZE]; - -static size_t getPartsPerRow(size_t rows, size_t dim, int task_dim) { - if (rows == 0 || task_dim <= 0) { - return 1; - } - size_t max_parts_by_tasks = std::max(1, static_cast(task_dim) / rows); - size_t parts_by_dim = std::max(1, (dim + 1023) / 1024); - return std::max(1, std::min(max_parts_by_tasks, parts_by_dim)); -} - -template -__mlu_device__ void loadToFloat(float *dst, T *cache, const T *src, size_t n) { - __memcpy(cache, src, n * sizeof(T), GDRAM2NRAM); - if constexpr (std::is_same::value) { - __bang_half2float(dst, cache, n); - } else if constexpr (std::is_same::value) { - __bang_bfloat162float(dst, cache, n); - } else { - __memcpy(dst, cache, n * sizeof(float), NRAM2NRAM); - } -} - -template -__mlu_device__ void storeFromFloat(T *dst, T *cache, float *src, size_t n) { - if constexpr (std::is_same::value) { - __bang_float2half(cache, src, n); - } else if constexpr (std::is_same::value) { - __bang_float2bfloat16(cache, src, n); - } else { - __memcpy(cache, src, n * sizeof(float), NRAM2NRAM); - } - __memcpy(dst, cache, n * sizeof(T), NRAM2GDRAM); -} - -__mlu_device__ int getMaxBatchSize(size_t dim, size_t data_size, size_t weight_size) { - constexpr int reduce_buffer_size = 128 / sizeof(float); - int max_batch_size = (NRAM_MAX_SIZE - 256 - reduce_buffer_size * sizeof(float)) / (3 * data_size + weight_size + 3 * sizeof(float)); - max_batch_size = std::min(max_batch_size, static_cast(dim)); - if (max_batch_size > 64) { - max_batch_size = (max_batch_size / 64) * 64; - } - return std::max(max_batch_size, 1); -} - -template -__mlu_device__ void bindNramBuffers( - int max_batch_size, - float **reduction_buffer, - Tdata **a_cache, - Tdata **b_cache, - Tdata **out_cache, - Tweight **weight_cache, - float **a_float, - float **b_float, - float **weight_float) { - constexpr int reduce_buffer_size = 128 / sizeof(float); - *reduction_buffer = reinterpret_cast(nram_buffer); - *a_cache = reinterpret_cast(*reduction_buffer + reduce_buffer_size); - *b_cache = *a_cache + max_batch_size; - *out_cache = *b_cache + max_batch_size; - *weight_cache = reinterpret_cast(*out_cache + max_batch_size); - *a_float = reinterpret_cast(*weight_cache + max_batch_size); - *b_float = *a_float + max_batch_size; - *weight_float = *b_float + max_batch_size; -} - -template -__mlu_global__ void addRMSNormKernel( - Tdata *__restrict__ y, - Tdata *__restrict__ residual_out, - const Tdata *__restrict__ a, - const Tdata *__restrict__ b, - const Tweight *__restrict__ weight, - size_t batch_size, - size_t nhead, - size_t dim, - ptrdiff_t stride_y_batch, - ptrdiff_t stride_y_nhead, - ptrdiff_t stride_residual_batch, - ptrdiff_t stride_residual_nhead, - ptrdiff_t stride_a_batch, - ptrdiff_t stride_a_nhead, - ptrdiff_t stride_b_batch, - ptrdiff_t stride_b_nhead, - float epsilon) { - - int max_batch_size = getMaxBatchSize(dim, sizeof(Tdata), sizeof(Tweight)); - float *reduction_buffer; - Tdata *a_cache; - Tdata *b_cache; - Tdata *out_cache; - Tweight *weight_cache; - float *a_float; - float *b_float; - float *weight_float; - bindNramBuffers(max_batch_size, &reduction_buffer, &a_cache, &b_cache, &out_cache, &weight_cache, &a_float, &b_float, &weight_float); - - size_t rows = batch_size * nhead; - for (size_t row = taskId; row < rows; row += taskDim) { - size_t batch_idx = row / nhead; - size_t head_idx = row - batch_idx * nhead; - - Tdata *y_ptr = y + batch_idx * stride_y_batch + head_idx * stride_y_nhead; - Tdata *residual_ptr = residual_out + batch_idx * stride_residual_batch + head_idx * stride_residual_nhead; - const Tdata *a_ptr = a + batch_idx * stride_a_batch + head_idx * stride_a_nhead; - const Tdata *b_ptr = b + batch_idx * stride_b_batch + head_idx * stride_b_nhead; - - float sum_squared = 0.0f; - size_t processed = 0; - while (processed < dim) { - size_t current = std::min(static_cast(max_batch_size), dim - processed); - - loadToFloat(a_float, a_cache, a_ptr + processed, current); - loadToFloat(b_float, b_cache, b_ptr + processed, current); - __bang_add(a_float, a_float, b_float, current); - storeFromFloat(residual_ptr + processed, out_cache, a_float, current); - - __bang_mul(a_float, a_float, a_float, current); - if (current >= 128) { - op::common_bang::reduce_op::sumInternal(reduction_buffer, a_float, current); - sum_squared += reduction_buffer[0]; - } else { - for (size_t i = 0; i < current; ++i) { - sum_squared += a_float[i]; - } - } - processed += current; - } - - float inv_rms = 1.0f / sqrtf(sum_squared / static_cast(dim) + epsilon); - - processed = 0; - while (processed < dim) { - size_t current = std::min(static_cast(max_batch_size), dim - processed); - - loadToFloat(a_float, a_cache, residual_ptr + processed, current); - loadToFloat(weight_float, weight_cache, weight + processed, current); - __bang_mul(a_float, a_float, weight_float, current); - __bang_mul_scalar(a_float, a_float, inv_rms, current); - storeFromFloat(y_ptr + processed, out_cache, a_float, current); - - processed += current; - } - } -} - -template -__mlu_global__ void addRMSNormPartialSumKernel( - Tdata *__restrict__ residual_out, - float *__restrict__ partial_sums, - const Tdata *__restrict__ a, - const Tdata *__restrict__ b, - size_t batch_size, - size_t nhead, - size_t dim, - size_t parts_per_row, - ptrdiff_t stride_residual_batch, - ptrdiff_t stride_residual_nhead, - ptrdiff_t stride_a_batch, - ptrdiff_t stride_a_nhead, - ptrdiff_t stride_b_batch, - ptrdiff_t stride_b_nhead) { - - int max_batch_size = getMaxBatchSize(dim, sizeof(Tdata), sizeof(Tweight)); - float *reduction_buffer; - Tdata *a_cache; - Tdata *b_cache; - Tdata *out_cache; - Tweight *weight_cache; - float *a_float; - float *b_float; - float *weight_float; - bindNramBuffers(max_batch_size, &reduction_buffer, &a_cache, &b_cache, &out_cache, &weight_cache, &a_float, &b_float, &weight_float); - - size_t rows = batch_size * nhead; - size_t logical_tasks = rows * parts_per_row; - size_t part_size = (dim + parts_per_row - 1) / parts_per_row; - - for (size_t logical = taskId; logical < logical_tasks; logical += taskDim) { - size_t row = logical / parts_per_row; - size_t part = logical - row * parts_per_row; - size_t start = part * part_size; - size_t end = std::min(dim, start + part_size); - if (start >= end) { - partial_sums[logical] = 0.0f; - continue; - } - - size_t batch_idx = row / nhead; - size_t head_idx = row - batch_idx * nhead; - Tdata *residual_ptr = residual_out + batch_idx * stride_residual_batch + head_idx * stride_residual_nhead; - const Tdata *a_ptr = a + batch_idx * stride_a_batch + head_idx * stride_a_nhead; - const Tdata *b_ptr = b + batch_idx * stride_b_batch + head_idx * stride_b_nhead; - - float sum_squared = 0.0f; - size_t processed = start; - while (processed < end) { - size_t current = std::min(static_cast(max_batch_size), end - processed); - - loadToFloat(a_float, a_cache, a_ptr + processed, current); - loadToFloat(b_float, b_cache, b_ptr + processed, current); - __bang_add(a_float, a_float, b_float, current); - storeFromFloat(residual_ptr + processed, out_cache, a_float, current); - - __bang_mul(a_float, a_float, a_float, current); - if (current >= 128) { - op::common_bang::reduce_op::sumInternal(reduction_buffer, a_float, current); - sum_squared += reduction_buffer[0]; - } else { - for (size_t i = 0; i < current; ++i) { - sum_squared += a_float[i]; - } - } - processed += current; - } - partial_sums[logical] = sum_squared; - } -} - -template -__mlu_global__ void addRMSNormApplyKernel( - Tdata *__restrict__ y, - const Tdata *__restrict__ residual_out, - const Tweight *__restrict__ weight, - const float *__restrict__ partial_sums, - size_t batch_size, - size_t nhead, - size_t dim, - size_t parts_per_row, - ptrdiff_t stride_y_batch, - ptrdiff_t stride_y_nhead, - ptrdiff_t stride_residual_batch, - ptrdiff_t stride_residual_nhead, - float epsilon) { - - int max_batch_size = getMaxBatchSize(dim, sizeof(Tdata), sizeof(Tweight)); - float *reduction_buffer; - Tdata *a_cache; - Tdata *b_cache; - Tdata *out_cache; - Tweight *weight_cache; - float *a_float; - float *b_float; - float *weight_float; - bindNramBuffers(max_batch_size, &reduction_buffer, &a_cache, &b_cache, &out_cache, &weight_cache, &a_float, &b_float, &weight_float); - - size_t rows = batch_size * nhead; - size_t logical_tasks = rows * parts_per_row; - size_t part_size = (dim + parts_per_row - 1) / parts_per_row; - - for (size_t logical = taskId; logical < logical_tasks; logical += taskDim) { - size_t row = logical / parts_per_row; - size_t part = logical - row * parts_per_row; - size_t start = part * part_size; - size_t end = std::min(dim, start + part_size); - if (start >= end) { - continue; - } - - float sum_squared = 0.0f; - const float *row_sums = partial_sums + row * parts_per_row; - for (size_t i = 0; i < parts_per_row; ++i) { - sum_squared += row_sums[i]; - } - float inv_rms = 1.0f / sqrtf(sum_squared / static_cast(dim) + epsilon); - - size_t batch_idx = row / nhead; - size_t head_idx = row - batch_idx * nhead; - Tdata *y_ptr = y + batch_idx * stride_y_batch + head_idx * stride_y_nhead; - const Tdata *residual_ptr = residual_out + batch_idx * stride_residual_batch + head_idx * stride_residual_nhead; - - size_t processed = start; - while (processed < end) { - size_t current = std::min(static_cast(max_batch_size), end - processed); - - loadToFloat(a_float, a_cache, residual_ptr + processed, current); - loadToFloat(weight_float, weight_cache, weight + processed, current); - __bang_mul(a_float, a_float, weight_float, current); - __bang_mul_scalar(a_float, a_float, inv_rms, current); - storeFromFloat(y_ptr + processed, out_cache, a_float, current); - - processed += current; - } - } -} - -template -static infiniStatus_t launchAddRMSNorm( - void *workspace, - int core_per_cluster, - int cluster_count, - cnrtQueue_t queue, - size_t batch_size, - size_t nhead, - size_t dim, - size_t parts_per_row, - void *y, - void *residual_out, - const void *a, - const void *b, - const void *weight, - ptrdiff_t stride_y_batch, - ptrdiff_t stride_y_nhead, - ptrdiff_t stride_residual_batch, - ptrdiff_t stride_residual_nhead, - ptrdiff_t stride_a_batch, - ptrdiff_t stride_a_nhead, - ptrdiff_t stride_b_batch, - ptrdiff_t stride_b_nhead, - float epsilon) { - - cnrtDim3_t kernel_dim; - kernel_dim.x = core_per_cluster; - kernel_dim.y = cluster_count; - kernel_dim.z = 1; - - if (parts_per_row > 1) { - auto partial_sums = reinterpret_cast(workspace); - addRMSNormPartialSumKernel<<>>( - reinterpret_cast(residual_out), - partial_sums, - reinterpret_cast(a), - reinterpret_cast(b), - batch_size, - nhead, - dim, - parts_per_row, - stride_residual_batch, - stride_residual_nhead, - stride_a_batch, - stride_a_nhead, - stride_b_batch, - stride_b_nhead); - addRMSNormApplyKernel<<>>( - reinterpret_cast(y), - reinterpret_cast(residual_out), - reinterpret_cast(weight), - partial_sums, - batch_size, - nhead, - dim, - parts_per_row, - stride_y_batch, - stride_y_nhead, - stride_residual_batch, - stride_residual_nhead, - epsilon); - } else { - addRMSNormKernel<<>>( - reinterpret_cast(y), - reinterpret_cast(residual_out), - reinterpret_cast(a), - reinterpret_cast(b), - reinterpret_cast(weight), - batch_size, - nhead, - dim, - stride_y_batch, - stride_y_nhead, - stride_residual_batch, - stride_residual_nhead, - stride_a_batch, - stride_a_nhead, - stride_b_batch, - stride_b_nhead, - epsilon); - } - cnrtQueueSync(queue); - return INFINI_STATUS_SUCCESS; -} - -namespace op::add_rms_norm::bang { - -struct Descriptor::Opaque { - std::shared_ptr internal; - size_t parts_per_row; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t residual_out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t weight_desc, - float epsilon) { - auto result = AddRMSNormInfo::create(y_desc, residual_out_desc, a_desc, b_desc, weight_desc, epsilon); - CHECK_RESULT(result); - auto info = result.take(); - - auto internal = reinterpret_cast(handle)->internal(); - size_t rows = info.shape[0] * (info.shape.size() > 2 ? info.shape[1] : 1); - int task_dim = internal->getCorePerCluster() * internal->getClusterCount(); - size_t parts_per_row = getPartsPerRow(rows, info.dim(), task_dim); - size_t workspace_size = parts_per_row > 1 ? rows * parts_per_row * sizeof(float) : 0; - - *desc_ptr = new Descriptor( - new Opaque{internal, parts_per_row}, - std::move(info), - workspace_size, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *y, void *residual_out, const void *a, const void *b, const void *weight, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - auto queue = reinterpret_cast(stream); - int core_per_cluster = _opaque->internal->getCorePerCluster(); - int cluster_count = _opaque->internal->getClusterCount(); - - size_t batch_size = _info.shape[0]; - size_t nhead = _info.shape.size() > 2 ? _info.shape[1] : 1; - size_t dim = _info.dim(); - - ptrdiff_t stride_y_batch = _info.y_strides[0]; - ptrdiff_t stride_y_nhead = _info.shape.size() > 2 ? _info.y_strides[1] : 0; - ptrdiff_t stride_residual_batch = _info.residual_out_strides[0]; - ptrdiff_t stride_residual_nhead = _info.shape.size() > 2 ? _info.residual_out_strides[1] : 0; - ptrdiff_t stride_a_batch = _info.a_strides[0]; - ptrdiff_t stride_a_nhead = _info.shape.size() > 2 ? _info.a_strides[1] : 0; - ptrdiff_t stride_b_batch = _info.b_strides[0]; - ptrdiff_t stride_b_nhead = _info.shape.size() > 2 ? _info.b_strides[1] : 0; - -#define DISPATCH(Tdata, Tweight) \ - return launchAddRMSNorm( \ - workspace, core_per_cluster, cluster_count, queue, batch_size, nhead, dim, \ - _opaque->parts_per_row, y, residual_out, a, b, weight, \ - stride_y_batch, stride_y_nhead, \ - stride_residual_batch, stride_residual_nhead, \ - stride_a_batch, stride_a_nhead, \ - stride_b_batch, stride_b_nhead, \ - _info.epsilon) - - if (_info.atype == INFINI_DTYPE_F16) { - if (_info.wtype == INFINI_DTYPE_F16) { - DISPATCH(half, half); - } else if (_info.wtype == INFINI_DTYPE_BF16) { - DISPATCH(half, bfloat16_t); - } else if (_info.wtype == INFINI_DTYPE_F32) { - DISPATCH(half, float); - } - } else if (_info.atype == INFINI_DTYPE_BF16) { - if (_info.wtype == INFINI_DTYPE_BF16) { - DISPATCH(bfloat16_t, bfloat16_t); - } else if (_info.wtype == INFINI_DTYPE_F16) { - DISPATCH(bfloat16_t, half); - } else if (_info.wtype == INFINI_DTYPE_F32) { - DISPATCH(bfloat16_t, float); - } - } else if (_info.atype == INFINI_DTYPE_F32 && _info.wtype == INFINI_DTYPE_F32) { - DISPATCH(float, float); - } - -#undef DISPATCH - - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -} // namespace op::add_rms_norm::bang diff --git a/src/infiniop/ops/add_rms_norm/cpu/add_rms_norm_cpu.cc b/src/infiniop/ops/add_rms_norm/cpu/add_rms_norm_cpu.cc deleted file mode 100644 index a3099c5c4..000000000 --- a/src/infiniop/ops/add_rms_norm/cpu/add_rms_norm_cpu.cc +++ /dev/null @@ -1,147 +0,0 @@ -#include "add_rms_norm_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include "../../../reduce/cpu/reduce.h" - -namespace op::add_rms_norm::cpu { - -Descriptor::~Descriptor() {} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t residual_out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t weight_desc, - float epsilon) { - auto result = AddRMSNormInfo::create(y_desc, residual_out_desc, a_desc, b_desc, weight_desc, epsilon); - CHECK_RESULT(result); - *desc_ptr = new Descriptor(nullptr, result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t add_rmsnorm(const AddRMSNormInfo *info, T *y, T *residual_out, const T *a, const T *b, const T *w) { - const size_t batch_size = info->shape[0]; - const size_t nhead = info->ndim() > 2 ? info->shape[1] : 1; - const size_t dim = info->dim(); - const ptrdiff_t total_blocks = static_cast(batch_size * nhead); - -#pragma omp parallel for - for (ptrdiff_t block_idx = 0; block_idx < total_blocks; ++block_idx) { - const size_t i = block_idx / nhead; // batch index - const size_t j = block_idx % nhead; // head index - - const T *a_ptr = a + i * info->a_strides[0] + j * info->a_strides[1]; - const T *b_ptr = b + i * info->b_strides[0] + j * info->b_strides[1]; - T *y_ptr = y + i * info->y_strides[0] + j * info->y_strides[1]; - T *residual_out_ptr = residual_out + i * info->residual_out_strides[0] + j * info->residual_out_strides[1]; - - // Compute add(a, b) once and store it - T sum_squared = (T)0; - for (size_t k = 0; k < dim; k++) { - T sum_val = a_ptr[k] + b_ptr[k]; - residual_out_ptr[k] = sum_val; // Store add result - sum_squared += sum_val * sum_val; - } - - // Compute RMS: 1 / (sqrt(mean(sum^2) + eps)) - // Note: mean = sum_squared / dim - T rms = (T)1 / std::sqrt(sum_squared / (T)(dim) + (T)(info->epsilon)); - - // Apply normalization: y = (a + b) * w * rms - // Reuse stored values from residual_out - for (size_t k = 0; k < dim; k++) { - y_ptr[k] = residual_out_ptr[k] * w[k] * rms; - } - } - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t add_rmsnormHalfPrecision(const AddRMSNormInfo *info, T *y, T *residual_out, const T *a, const T *b, const Tw *w) { - static_assert(std::is_same::value || std::is_same::value, - "T must be fp16_t or bf16_t"); - - const size_t batch_size = info->shape[0]; - const size_t nhead = info->ndim() > 2 ? info->shape[1] : 1; - const size_t dim = info->dim(); - const ptrdiff_t total_blocks = static_cast(batch_size * nhead); - -#pragma omp parallel for - for (ptrdiff_t block_idx = 0; block_idx < total_blocks; ++block_idx) { - const size_t i = block_idx / nhead; // batch index - const size_t j = block_idx % nhead; // head index - - const T *a_ptr = a + i * info->a_strides[0] + j * info->a_strides[1]; - const T *b_ptr = b + i * info->b_strides[0] + j * info->b_strides[1]; - T *y_ptr = y + i * info->y_strides[0] + j * info->y_strides[1]; - T *residual_out_ptr = residual_out + i * info->residual_out_strides[0] + j * info->residual_out_strides[1]; - - // Compute sum of squares for RMS normalization and store add result - float sum_squared = 0.0f; - for (size_t k = 0; k < dim; k++) { - float sum_val = utils::cast(a_ptr[k]) + utils::cast(b_ptr[k]); - residual_out_ptr[k] = utils::cast(sum_val); // Store add result - sum_squared += sum_val * sum_val; - } - - // Compute RMS: 1 / (sqrt(sum/dim + eps)) - float rms = 1.f / std::sqrt(sum_squared / (float)(dim) + info->epsilon); - - // Apply normalization: y = (a + b) * w * rms - // Reuse stored values from residual_out - for (size_t k = 0; k < dim; k++) { - float sum_val = utils::cast(residual_out_ptr[k]); - float val; - if constexpr (std::is_same::value) { - val = sum_val * w[k] * rms; - } else if constexpr (std::is_same::value || std::is_same_v || std::is_same_v) { - val = sum_val * utils::cast(w[k]) * rms; - } else { - std::abort(); - } - y_ptr[k] = utils::cast(val); - } - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *y, void *residual_out, const void *a, const void *b, const void *weight, - void *stream) const { - if (_info.atype == INFINI_DTYPE_F16) { - if (_info.wtype == INFINI_DTYPE_F16) { - CHECK_STATUS(add_rmsnormHalfPrecision(&_info, (fp16_t *)y, (fp16_t *)residual_out, (const fp16_t *)a, (const fp16_t *)b, (const fp16_t *)weight)); - } else if (_info.wtype == INFINI_DTYPE_F32) { - CHECK_STATUS(add_rmsnormHalfPrecision(&_info, (fp16_t *)y, (fp16_t *)residual_out, (const fp16_t *)a, (const fp16_t *)b, (const float *)weight)); - } else if (_info.wtype == INFINI_DTYPE_BF16) { - CHECK_STATUS(add_rmsnormHalfPrecision(&_info, (fp16_t *)y, (fp16_t *)residual_out, (const fp16_t *)a, (const fp16_t *)b, (const bf16_t *)weight)); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (_info.atype == INFINI_DTYPE_BF16) { - if (_info.wtype == INFINI_DTYPE_BF16) { - CHECK_STATUS(add_rmsnormHalfPrecision(&_info, (bf16_t *)y, (bf16_t *)residual_out, (const bf16_t *)a, (const bf16_t *)b, (const bf16_t *)weight)); - } else if (_info.wtype == INFINI_DTYPE_F32) { - CHECK_STATUS(add_rmsnormHalfPrecision(&_info, (bf16_t *)y, (bf16_t *)residual_out, (const bf16_t *)a, (const bf16_t *)b, (const float *)weight)); - } else if (_info.wtype == INFINI_DTYPE_F16) { - CHECK_STATUS(add_rmsnormHalfPrecision(&_info, (bf16_t *)y, (bf16_t *)residual_out, (const bf16_t *)a, (const bf16_t *)b, (const fp16_t *)weight)); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (_info.atype == INFINI_DTYPE_F32) { - CHECK_STATUS(add_rmsnorm(&_info, (float *)y, (float *)residual_out, (const float *)a, (const float *)b, (const float *)weight)); - } else if (_info.atype == INFINI_DTYPE_F64) { - CHECK_STATUS(add_rmsnorm(&_info, (double *)y, (double *)residual_out, (const double *)a, (const double *)b, (const double *)weight)); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::add_rms_norm::cpu diff --git a/src/infiniop/ops/add_rms_norm/cpu/add_rms_norm_cpu.h b/src/infiniop/ops/add_rms_norm/cpu/add_rms_norm_cpu.h deleted file mode 100644 index f4173adbf..000000000 --- a/src/infiniop/ops/add_rms_norm/cpu/add_rms_norm_cpu.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __ADD_RMS_NORM_CPU_H__ -#define __ADD_RMS_NORM_CPU_H__ -#include "../add_rms_norm.h" - -DESCRIPTOR(cpu) - -#endif diff --git a/src/infiniop/ops/add_rms_norm/cuda/kernel.cuh b/src/infiniop/ops/add_rms_norm/cuda/kernel.cuh deleted file mode 100644 index 97f254a09..000000000 --- a/src/infiniop/ops/add_rms_norm/cuda/kernel.cuh +++ /dev/null @@ -1,63 +0,0 @@ -#ifndef __ADD_RMS_NORM_CUDA_KERNEL_H__ -#define __ADD_RMS_NORM_CUDA_KERNEL_H__ - -#include - -template -__device__ void add_rmsnormBlock( - Tdata *__restrict__ y, - Tdata *__restrict__ residual_out, - ptrdiff_t stride_y_batch, - ptrdiff_t stride_y_nhead, - ptrdiff_t stride_residual_out_batch, - ptrdiff_t stride_residual_out_nhead, - const Tdata *__restrict__ a, - ptrdiff_t stride_a_batch, - ptrdiff_t stride_a_nhead, - const Tdata *__restrict__ b, - ptrdiff_t stride_b_batch, - ptrdiff_t stride_b_nhead, - const Tweight *__restrict__ w, - size_t nhead, - size_t dim, - float epsilon) { - // Each block takes care of one head in one batch - // Each thread deals with every block_size element in the row - size_t batch_idx = blockIdx.x / nhead; - size_t head_idx = blockIdx.x % nhead; - - auto y_ptr = y + batch_idx * stride_y_batch + head_idx * stride_y_nhead; - auto a_ptr = a + batch_idx * stride_a_batch + head_idx * stride_a_nhead; - auto b_ptr = b + batch_idx * stride_b_batch + head_idx * stride_b_nhead; - auto w_ptr = w; - Tdata *residual_out_ptr = residual_out + batch_idx * stride_residual_out_batch + head_idx * stride_residual_out_nhead; - - // Compute add(a, b) and sum of squares in one pass - Tcompute sum_squared = 0; - for (size_t i = threadIdx.x; i < dim; i += BLOCK_SIZE) { - Tcompute sum_val = Tcompute(a_ptr[i]) + Tcompute(b_ptr[i]); - residual_out_ptr[i] = Tdata(sum_val); // Store add result - sum_squared += sum_val * sum_val; - } - - // Block-reduce sum of squares - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - sum_squared = BlockReduce(temp_storage).Sum(sum_squared); - - // Thread_0 computes RMS=1/sqrt(ss/dim+epsilon) and stores in shared memory - __shared__ Tcompute rms; - if (threadIdx.x == 0) { - rms = Tcompute(rsqrtf(sum_squared / Tcompute(dim) + epsilon)); - } - __syncthreads(); - - // Apply normalization: y = (a + b) * w * rms - // Reuse stored values from residual_out - for (size_t i = threadIdx.x; i < dim; i += BLOCK_SIZE) { - Tcompute sum_val = Tcompute(residual_out_ptr[i]); // Reuse stored value - y_ptr[i] = Tdata(sum_val * Tcompute(w_ptr[i]) * rms); - } -} - -#endif diff --git a/src/infiniop/ops/add_rms_norm/info.h b/src/infiniop/ops/add_rms_norm/info.h deleted file mode 100644 index 883aed343..000000000 --- a/src/infiniop/ops/add_rms_norm/info.h +++ /dev/null @@ -1,132 +0,0 @@ -#ifndef __ADD_RMS_NORM_INFO_H__ -#define __ADD_RMS_NORM_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::add_rms_norm { - -class AddRMSNormInfo { - AddRMSNormInfo() = default; - -public: - infiniDtype_t wtype; - infiniDtype_t atype; - float epsilon; - std::vector shape; - std::vector y_strides; - std::vector residual_out_strides; - std::vector a_strides; - std::vector b_strides; - bool has_residual_out; - - size_t ndim() const { return shape.size(); } - size_t dim() const { return shape[ndim() - 1]; } - - static utils::Result create( - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t residual_out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t weight_desc, - float epsilon) { - - auto atype = y_desc->dtype(); - auto wtype = weight_desc->dtype(); - - // Check that all input tensors have the same dtype - if (a_desc->dtype() != atype || b_desc->dtype() != atype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - if (atype == INFINI_DTYPE_F16 || atype == INFINI_DTYPE_BF16) { - // For half-precision types (FP16/BF16), weights can be the same half-precision type or FP32 - if (wtype != atype && wtype != INFINI_DTYPE_F32 && wtype != INFINI_DTYPE_BF16 && wtype != INFINI_DTYPE_F16) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (atype == INFINI_DTYPE_F32 || atype == INFINI_DTYPE_F64) { - // For FP32/FP64, activations and weights must be of the same type - if (atype != wtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - const size_t y_ndim = y_desc->ndim(); - const size_t a_ndim = a_desc->ndim(); - const size_t b_ndim = b_desc->ndim(); - const size_t w_ndim = weight_desc->ndim(); - - if (y_ndim != a_ndim || y_ndim != b_ndim || w_ndim != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t batch = 1; - size_t nhead = 1; - size_t dim = 0; - - if (y_ndim == 2) { - batch = y_desc->dim(0); - dim = y_desc->dim(1); - - if (a_desc->dim(0) != batch || a_desc->dim(1) != dim || b_desc->dim(0) != batch || b_desc->dim(1) != dim || weight_desc->dim(0) != dim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } else if (y_ndim == 3) { - batch = y_desc->dim(0); - nhead = y_desc->dim(1); - dim = y_desc->dim(2); - - if (a_desc->dim(0) != batch || a_desc->dim(1) != nhead || a_desc->dim(2) != dim || b_desc->dim(0) != batch || b_desc->dim(1) != nhead || b_desc->dim(2) != dim || weight_desc->dim(0) != dim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } else { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // Check contiguity of the last dimension - if (y_desc->stride(y_ndim - 1) != 1 || a_desc->stride(a_ndim - 1) != 1 || b_desc->stride(b_ndim - 1) != 1 || weight_desc->stride(w_ndim - 1) != 1) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - // residual_out_desc is required (always needed for fused operator) - if (residual_out_desc == nullptr) { - return INFINI_STATUS_BAD_PARAM; - } - - const size_t residual_out_ndim = residual_out_desc->ndim(); - if (residual_out_ndim != y_ndim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (residual_out_desc->dtype() != atype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - // Check shape matches - for (size_t i = 0; i < y_ndim; i++) { - if (residual_out_desc->dim(i) != y_desc->dim(i)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - if (residual_out_desc->stride(residual_out_ndim - 1) != 1) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - AddRMSNormInfo info; - info.wtype = wtype; - info.atype = atype; - info.epsilon = epsilon; - info.shape = y_desc->shape(); - info.y_strides = y_desc->strides(); - info.a_strides = a_desc->strides(); - info.b_strides = b_desc->strides(); - info.has_residual_out = true; // Always true now - info.residual_out_strides = residual_out_desc->strides(); - return utils::Result(info); - } -}; - -} // namespace op::add_rms_norm - -#endif // __ADD_RMS_NORM_INFO_H__ diff --git a/src/infiniop/ops/add_rms_norm/metax/add_rms_norm_metax.cuh b/src/infiniop/ops/add_rms_norm/metax/add_rms_norm_metax.cuh deleted file mode 100644 index 3d6b13b53..000000000 --- a/src/infiniop/ops/add_rms_norm/metax/add_rms_norm_metax.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADD_RMS_NORM_METAX_CUH__ -#define __ADD_RMS_NORM_METAX_CUH__ - -#include "../add_rms_norm.h" - -DESCRIPTOR(metax) - -#endif diff --git a/src/infiniop/ops/add_rms_norm/metax/add_rms_norm_metax.maca b/src/infiniop/ops/add_rms_norm/metax/add_rms_norm_metax.maca deleted file mode 100644 index 04355e927..000000000 --- a/src/infiniop/ops/add_rms_norm/metax/add_rms_norm_metax.maca +++ /dev/null @@ -1,191 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "add_rms_norm_metax.cuh" - -#include "../../../devices/metax/metax_kernel_common.h" -#include - -#include "../../../reduce/cuda/reduce.cuh" - -#include "../cuda/kernel.cuh" - -// Kernel function template for add_rms_norm on Metax platform -template -INFINIOP_METAX_KERNEL add_rmsnormKernel( - Tdata *__restrict__ y, - Tdata *__restrict__ residual_out, - ptrdiff_t stride_y_batch, - ptrdiff_t stride_y_nhead, - ptrdiff_t stride_residual_out_batch, - ptrdiff_t stride_residual_out_nhead, - const Tdata *__restrict__ a, - ptrdiff_t stride_a_batch, - ptrdiff_t stride_a_nhead, - const Tdata *__restrict__ b, - ptrdiff_t stride_b_batch, - ptrdiff_t stride_b_nhead, - const Tweight *__restrict__ w, - size_t nhead, - size_t dim, - float epsilon) { - add_rmsnormBlock( - y, residual_out, - stride_y_batch, stride_y_nhead, - stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - w, nhead, dim, epsilon); -} - -namespace op::add_rms_norm::metax { - -// Internal opaque structure for Metax device handle -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -// Destructor -Descriptor::~Descriptor() { - delete _opaque; -} - -// Create descriptor for add_rms_norm operator -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t residual_out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t weight_desc, - float epsilon) { - auto result = AddRMSNormInfo::create(y_desc, residual_out_desc, a_desc, b_desc, weight_desc, epsilon); - CHECK_RESULT(result); - auto info = result.take(); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - std::move(info), - 0, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -// Launch kernel with different data types -template -infiniStatus_t launchKernel( - uint32_t batch_size, size_t nhead, size_t dim, - void *y, infiniDtype_t atype, ptrdiff_t stride_y_batch, ptrdiff_t stride_y_nhead, - void *residual_out, ptrdiff_t stride_residual_out_batch, ptrdiff_t stride_residual_out_nhead, - const void *a, ptrdiff_t stride_a_batch, ptrdiff_t stride_a_nhead, - const void *b, ptrdiff_t stride_b_batch, ptrdiff_t stride_b_nhead, - const void *w, infiniDtype_t wtype, - float epsilon, - hcStream_t stream) { - -#define LAUNCH_KERNEL(Tdata, Tweight, Tcompute) \ - add_rmsnormKernel<<>>( \ - reinterpret_cast(y), \ - reinterpret_cast(residual_out), \ - stride_y_batch, \ - stride_y_nhead, \ - stride_residual_out_batch, \ - stride_residual_out_nhead, \ - reinterpret_cast(a), \ - stride_a_batch, \ - stride_a_nhead, \ - reinterpret_cast(b), \ - stride_b_batch, \ - stride_b_nhead, \ - reinterpret_cast(w), \ - nhead, \ - dim, \ - epsilon) - - // Handle different data type combinations following Metax pattern - if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(half, half, float); - } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(half, __hpcc_bfloat16, float); - } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(half, float, float); - } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(__hpcc_bfloat16, __hpcc_bfloat16, float); - } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(__hpcc_bfloat16, half, float); - } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(__hpcc_bfloat16, float, float); - } else if (atype == INFINI_DTYPE_F32 && wtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(float, float, float); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_KERNEL - - return INFINI_STATUS_SUCCESS; -} - -// Main calculation function -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *y, void *residual_out, const void *a, const void *b, const void *weight, - void *stream) const { - - // Check workspace size - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // Extract tensor strides and dimensions - auto stride_a_batch = _info.a_strides[0]; - auto stride_a_nhead = _info.a_strides[1]; - auto stride_b_batch = _info.b_strides[0]; - auto stride_b_nhead = _info.b_strides[1]; - auto stride_y_batch = _info.y_strides[0]; - auto stride_y_nhead = _info.y_strides[1]; - auto stride_residual_out_batch = _info.residual_out_strides[0]; - auto stride_residual_out_nhead = _info.residual_out_strides[1]; - auto dim = _info.dim(); - uint32_t batch_size = static_cast(_info.shape[0]); - size_t nhead = _info.shape.size() > 2 ? _info.shape[1] : 1; - auto stream_ = reinterpret_cast(stream); - - // Launch kernel with different block sizes - if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, - residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - weight, _info.wtype, _info.epsilon, stream_)); - } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, - residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - weight, _info.wtype, _info.epsilon, stream_)); - } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_2048) { - CHECK_STATUS(launchKernel( - batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, - residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - weight, _info.wtype, _info.epsilon, stream_)); - } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_4096) { - CHECK_STATUS(launchKernel( - batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, - residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - weight, _info.wtype, _info.epsilon, stream_)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} -} // namespace op::add_rms_norm::metax diff --git a/src/infiniop/ops/add_rms_norm/moore/add_rms_norm_moore.h b/src/infiniop/ops/add_rms_norm/moore/add_rms_norm_moore.h deleted file mode 100644 index 9d3f810f2..000000000 --- a/src/infiniop/ops/add_rms_norm/moore/add_rms_norm_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADD_RMS_NORM_MOORE_H__ -#define __ADD_RMS_NORM_MOORE_H__ - -#include "../add_rms_norm.h" - -DESCRIPTOR(moore) - -#endif diff --git a/src/infiniop/ops/add_rms_norm/moore/add_rms_norm_moore.mu b/src/infiniop/ops/add_rms_norm/moore/add_rms_norm_moore.mu deleted file mode 100644 index 90c027ead..000000000 --- a/src/infiniop/ops/add_rms_norm/moore/add_rms_norm_moore.mu +++ /dev/null @@ -1,183 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "add_rms_norm_moore.h" - -#include "../../../devices/moore/moore_kernel_common.h" -#include - -#include "../../../reduce/cuda/reduce.cuh" - -#include "../cuda/kernel.cuh" - -// Kernel function template for add_rms_norm on Moore platform -template -INFINIOP_MOORE_KERNEL add_rmsnormKernel( - Tdata *__restrict__ y, - Tdata *__restrict__ residual_out, - ptrdiff_t stride_y_batch, - ptrdiff_t stride_y_nhead, - ptrdiff_t stride_residual_out_batch, - ptrdiff_t stride_residual_out_nhead, - const Tdata *__restrict__ a, - ptrdiff_t stride_a_batch, - ptrdiff_t stride_a_nhead, - const Tdata *__restrict__ b, - ptrdiff_t stride_b_batch, - ptrdiff_t stride_b_nhead, - const Tweight *__restrict__ w, - size_t nhead, - size_t dim, - float epsilon) { - add_rmsnormBlock( - y, residual_out, - stride_y_batch, stride_y_nhead, - stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - w, nhead, dim, epsilon); -} - -namespace op::add_rms_norm::moore { - -// Internal opaque structure for Moore device handle -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -// Destructor -Descriptor::~Descriptor() { - delete _opaque; -} - -// Create descriptor for add_rms_norm operator -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t residual_out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t weight_desc, - float epsilon) { - auto result = AddRMSNormInfo::create(y_desc, residual_out_desc, a_desc, b_desc, weight_desc, epsilon); - CHECK_RESULT(result); - auto info = result.take(); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - std::move(info), - 0, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -// Launch kernel with different data types -template -infiniStatus_t launchKernel( - uint32_t batch_size, size_t nhead, size_t dim, - void *y, infiniDtype_t atype, ptrdiff_t stride_y_batch, ptrdiff_t stride_y_nhead, - void *residual_out, ptrdiff_t stride_residual_out_batch, ptrdiff_t stride_residual_out_nhead, - const void *a, ptrdiff_t stride_a_batch, ptrdiff_t stride_a_nhead, - const void *b, ptrdiff_t stride_b_batch, ptrdiff_t stride_b_nhead, - const void *w, infiniDtype_t wtype, - float epsilon, - musaStream_t musa_stream) { - -#define LAUNCH_KERNEL(Tdata, Tweight, Tcompute) \ - add_rmsnormKernel<<>>( \ - reinterpret_cast(y), \ - reinterpret_cast(residual_out), \ - stride_y_batch, \ - stride_y_nhead, \ - stride_residual_out_batch, \ - stride_residual_out_nhead, \ - reinterpret_cast(a), \ - stride_a_batch, \ - stride_a_nhead, \ - reinterpret_cast(b), \ - stride_b_batch, \ - stride_b_nhead, \ - reinterpret_cast(w), \ - nhead, \ - dim, \ - epsilon) - - // Handle different data type combinations - if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(half, half, float); - } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(half, __mt_bfloat16, float); - } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(half, float, float); - } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(__mt_bfloat16, __mt_bfloat16, float); - } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(__mt_bfloat16, half, float); - } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(__mt_bfloat16, float, float); - } else if (atype == INFINI_DTYPE_F32 && wtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(float, float, float); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_KERNEL - - return INFINI_STATUS_SUCCESS; -} - -// Main calculation function -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *y, void *residual_out, const void *a, const void *b, const void *weight, - void *stream) const { - - // Check workspace size - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // Extract tensor strides and dimensions - auto stride_a_batch = _info.a_strides[0]; - auto stride_a_nhead = _info.a_strides[1]; - auto stride_b_batch = _info.b_strides[0]; - auto stride_b_nhead = _info.b_strides[1]; - auto stride_y_batch = _info.y_strides[0]; - auto stride_y_nhead = _info.y_strides[1]; - auto stride_residual_out_batch = _info.residual_out_strides[0]; - auto stride_residual_out_nhead = _info.residual_out_strides[1]; - auto dim = _info.dim(); - uint32_t batch_size = static_cast(_info.shape[0]); - size_t nhead = _info.shape.size() > 2 ? _info.shape[1] : 1; - auto musa_stream = reinterpret_cast(stream); - - // Launch kernel with appropriate block size based on device capability - if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, - residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - weight, _info.wtype, _info.epsilon, musa_stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, - residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - weight, _info.wtype, _info.epsilon, musa_stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_2048) { - CHECK_STATUS(launchKernel( - batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, - residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - weight, _info.wtype, _info.epsilon, musa_stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} -} // namespace op::add_rms_norm::moore diff --git a/src/infiniop/ops/add_rms_norm/nvidia/add_rms_norm_nvidia.cu b/src/infiniop/ops/add_rms_norm/nvidia/add_rms_norm_nvidia.cu deleted file mode 100644 index 652f8adf3..000000000 --- a/src/infiniop/ops/add_rms_norm/nvidia/add_rms_norm_nvidia.cu +++ /dev/null @@ -1,183 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "add_rms_norm_nvidia.cuh" - -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include - -#include "../../../reduce/cuda/reduce.cuh" - -#include "../cuda/kernel.cuh" - -template -INFINIOP_CUDA_KERNEL add_rmsnormKernel( - Tdata *__restrict__ y, - Tdata *__restrict__ residual_out, - ptrdiff_t stride_y_batch, - ptrdiff_t stride_y_nhead, - ptrdiff_t stride_residual_out_batch, - ptrdiff_t stride_residual_out_nhead, - const Tdata *__restrict__ a, - ptrdiff_t stride_a_batch, - ptrdiff_t stride_a_nhead, - const Tdata *__restrict__ b, - ptrdiff_t stride_b_batch, - ptrdiff_t stride_b_nhead, - const Tweight *__restrict__ w, - size_t nhead, - size_t dim, - float epsilon) { - add_rmsnormBlock( - y, residual_out, - stride_y_batch, stride_y_nhead, - stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - w, nhead, dim, epsilon); -} - -namespace op::add_rms_norm::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t residual_out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t weight_desc, - float epsilon) { - auto result = AddRMSNormInfo::create(y_desc, residual_out_desc, a_desc, b_desc, weight_desc, epsilon); - CHECK_RESULT(result); - auto info = result.take(); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - std::move(info), - 0, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -// launch kernel with different data types -template -infiniStatus_t launchKernel( - uint32_t batch_size, size_t nhead, size_t dim, - void *y, infiniDtype_t atype, ptrdiff_t stride_y_batch, ptrdiff_t stride_y_nhead, - void *residual_out, ptrdiff_t stride_residual_out_batch, ptrdiff_t stride_residual_out_nhead, - const void *a, ptrdiff_t stride_a_batch, ptrdiff_t stride_a_nhead, - const void *b, ptrdiff_t stride_b_batch, ptrdiff_t stride_b_nhead, - const void *w, infiniDtype_t wtype, - float epsilon, - cudaStream_t cuda_stream) { - -#define LAUNCH_KERNEL(Tdata, Tweight, Tcompute) \ - add_rmsnormKernel<<>>( \ - reinterpret_cast(y), \ - reinterpret_cast(residual_out), \ - stride_y_batch, \ - stride_y_nhead, \ - stride_residual_out_batch, \ - stride_residual_out_nhead, \ - reinterpret_cast(a), \ - stride_a_batch, \ - stride_a_nhead, \ - reinterpret_cast(b), \ - stride_b_batch, \ - stride_b_nhead, \ - reinterpret_cast(w), \ - nhead, \ - dim, \ - epsilon) - - if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(half, half, float); - } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(half, __nv_bfloat16, float); - } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(half, float, float); - } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(__nv_bfloat16, __nv_bfloat16, float); - } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(__nv_bfloat16, half, float); - } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(__nv_bfloat16, float, float); - } else if (atype == INFINI_DTYPE_F32 && wtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(float, float, float); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_KERNEL - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *y, void *residual_out, const void *a, const void *b, const void *weight, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - auto stride_a_batch = _info.a_strides[0]; - auto stride_a_nhead = _info.a_strides[1]; - auto stride_b_batch = _info.b_strides[0]; - auto stride_b_nhead = _info.b_strides[1]; - auto stride_y_batch = _info.y_strides[0]; - auto stride_y_nhead = _info.y_strides[1]; - auto stride_residual_out_batch = _info.residual_out_strides[0]; - auto stride_residual_out_nhead = _info.residual_out_strides[1]; - auto dim = _info.dim(); - uint32_t batch_size = static_cast(_info.shape[0]); - size_t nhead = _info.shape.size() > 2 ? _info.shape[1] : 1; - auto cuda_stream = reinterpret_cast(stream); - - // launch kernel with different block sizes - if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, - residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - weight, _info.wtype, _info.epsilon, cuda_stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, - residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - weight, _info.wtype, _info.epsilon, cuda_stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_2048) { - CHECK_STATUS(launchKernel( - batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, - residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - weight, _info.wtype, _info.epsilon, cuda_stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) { - CHECK_STATUS(launchKernel( - batch_size, nhead, dim, - y, _info.atype, stride_y_batch, stride_y_nhead, - residual_out, stride_residual_out_batch, stride_residual_out_nhead, - a, stride_a_batch, stride_a_nhead, - b, stride_b_batch, stride_b_nhead, - weight, _info.wtype, _info.epsilon, cuda_stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} -} // namespace op::add_rms_norm::nvidia diff --git a/src/infiniop/ops/add_rms_norm/nvidia/add_rms_norm_nvidia.cuh b/src/infiniop/ops/add_rms_norm/nvidia/add_rms_norm_nvidia.cuh deleted file mode 100644 index 548be83e1..000000000 --- a/src/infiniop/ops/add_rms_norm/nvidia/add_rms_norm_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADD_RMS_NORM_NVIDIA_CUDA_H__ -#define __ADD_RMS_NORM_NVIDIA_CUDA_H__ - -#include "../add_rms_norm.h" - -DESCRIPTOR(nvidia) - -#endif diff --git a/src/infiniop/ops/add_rms_norm/operator.cc b/src/infiniop/ops/add_rms_norm/operator.cc deleted file mode 100644 index 7d502770d..000000000 --- a/src/infiniop/ops/add_rms_norm/operator.cc +++ /dev/null @@ -1,234 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/add_rms_norm.h" - -#ifdef ENABLE_CPU_API -#include "cpu/add_rms_norm_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) -#include "nvidia/add_rms_norm_nvidia.cuh" -#endif -#ifdef ENABLE_ASCEND_API -// TODO: Add Ascend implementation -// #include "ascend/add_rms_norm_aclnn.h" -#endif -#ifdef ENABLE_CAMBRICON_API -#include "bang/add_rms_norm_bang.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/add_rms_norm_metax.cuh" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/add_rms_norm_moore.h" -#endif -#ifdef ENABLE_KUNLUN_API -// TODO: Add Kunlun implementation -// #include "kunlun/add_rms_norm_kunlun.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAddRMSNormDescriptor( - infiniopHandle_t handle, - infiniopAddRMSNormDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t residual_out_desc, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t weight_desc, - float epsilon) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::add_rms_norm::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - residual_out_desc, \ - a_desc, \ - b_desc, \ - weight_desc, \ - epsilon) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_KUNLUN_API - // CREATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAddRMSNormWorkspaceSize(infiniopAddRMSNormDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - GET(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - GET(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_KUNLUN_API - // GET(INFINI_DEVICE_KUNLUN, kunlun); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopAddRMSNorm( - infiniopAddRMSNormDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - void *residual_out, - const void *a, - const void *b, - const void *weight, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, residual_out, a, b, weight, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_KUNLUN_API - // CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyAddRMSNormDescriptor(infiniopAddRMSNormDescriptor_t desc) { - if (desc == nullptr) { - return INFINI_STATUS_SUCCESS; - } - -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DESTROY(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - DESTROY(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_MOORE_API - DESTROY(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - DESTROY(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - DESTROY(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_QY_API - DESTROY(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DESTROY(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_KUNLUN_API - // DESTROY(INFINI_DEVICE_KUNLUN, kunlun); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DESTROY -} diff --git a/src/infiniop/ops/addbmm/addbmm.h b/src/infiniop/ops/addbmm/addbmm.h deleted file mode 100644 index 8e54a500d..000000000 --- a/src/infiniop/ops/addbmm/addbmm.h +++ /dev/null @@ -1,52 +0,0 @@ -#ifndef ADDBMM_H -#define ADDBMM_H - -#include "../../operator.h" -#include "info.h" // 对应 addbmm_info.h - -// 宏定义：用于生成不同命名空间下的 Descriptor 类 -// 注意：addbmm 需要处理 alpha, beta 和多个输入张量 -#define DESCRIPTOR(NAMESPACE) \ - namespace op::addbmm::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AddbmmInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - AddbmmInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - /* create 函数接收 Tensor 描述符列表和标量系数 */ \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - std::vector input_desc_vec, \ - float alpha, \ - float beta); \ - \ - /* calculate 函数接收数据指针列表 */ \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *output, \ - std::vector inputs, \ - void *stream) const; \ - }; \ - } - -#endif // ADDBMM_H diff --git a/src/infiniop/ops/addbmm/cpu/addbmm_cpu.cc b/src/infiniop/ops/addbmm/cpu/addbmm_cpu.cc deleted file mode 100644 index 7b430ac60..000000000 --- a/src/infiniop/ops/addbmm/cpu/addbmm_cpu.cc +++ /dev/null @@ -1,185 +0,0 @@ -#include "addbmm_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include "../../../handle.h" -#include -#include -#include -#include -#include -#include - -namespace op::addbmm::cpu { - -Descriptor::~Descriptor() = default; - -// ================================================================== -// 辅助函数：通用 stride 寻址 -// ================================================================== - -// 计算 2D 张量的偏移量 -inline size_t offset_2d(size_t r, size_t c, const int64_t *strides) { - return r * strides[0] + c * strides[1]; -} - -// 计算 3D 张量的偏移量 -inline size_t offset_3d(size_t b, size_t r, size_t c, const int64_t *strides) { - return b * strides[0] + r * strides[1] + c * strides[2]; -} - -// ================================================================== -// 核心 Kernel 实现 -// ================================================================== - -/** - * @brief Addbmm 核心 CPU 计算函数 (支持任意 Stride) - */ -template -void calculate_impl( - const AddbmmInfo &info, - void *output, - const void *input, - const void *batch1, - const void *batch2) { - - // [变更 1] 使用 Getter 获取维度 - size_t b_dim = info.b(); - size_t n = info.n(); - size_t m = info.m(); - size_t p = info.p(); - - float alpha = info.alpha(); - float beta = info.beta(); - - // 指针转换 - Tdata *out_ptr = reinterpret_cast(output); - const Tdata *inp_ptr = reinterpret_cast(input); - const Tdata *b1_ptr = reinterpret_cast(batch1); - const Tdata *b2_ptr = reinterpret_cast(batch2); - - const int64_t *out_strides = info.out_strides().data(); - const int64_t *in_strides = info.in_strides().data(); - const int64_t *b1_strides = info.b1_strides().data(); - const int64_t *b2_strides = info.b2_strides().data(); - - // 1. 初始化 output = beta * input - for (size_t i = 0; i < n; ++i) { - for (size_t k = 0; k < p; ++k) { - size_t out_idx = offset_2d(i, k, out_strides); - size_t in_idx = offset_2d(i, k, in_strides); - - float val_in = (beta != 0.0f) ? utils::cast(inp_ptr[in_idx]) : 0.0f; - - if (beta == 0.0f && alpha == 0.0f) { - out_ptr[out_idx] = utils::cast(0.0f); - } else { - out_ptr[out_idx] = utils::cast(val_in * beta); - } - } - } - - // 2. 累加矩阵乘法: out += alpha * sum(b1 @ b2) - for (size_t b = 0; b < b_dim; ++b) { // Batch - for (size_t i = 0; i < n; ++i) { // Row - for (size_t k = 0; k < p; ++k) { // Col - - float dot_product = 0.0f; - - // 内部点积 (Inner dimension m) - for (size_t j = 0; j < m; ++j) { - size_t b1_idx = offset_3d(b, i, j, b1_strides); - size_t b2_idx = offset_3d(b, j, k, b2_strides); - - float v1 = utils::cast(b1_ptr[b1_idx]); - float v2 = utils::cast(b2_ptr[b2_idx]); - - dot_product += v1 * v2; - } - - // 累加到 Output - size_t out_idx = offset_2d(i, k, out_strides); - float current_val = utils::cast(out_ptr[out_idx]); - out_ptr[out_idx] = utils::cast(current_val + alpha * dot_product); - } - } - } -} - -// ================================================================== -// Descriptor 接口实现 -// ================================================================== - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float alpha, - float beta) { - - if (input_desc_vec.size() != 3) { - return INFINI_STATUS_BAD_PARAM; - } - - infiniopTensorDescriptor_t in_desc = input_desc_vec[0]; - infiniopTensorDescriptor_t batch1_desc = input_desc_vec[1]; - infiniopTensorDescriptor_t batch2_desc = input_desc_vec[2]; - - auto dtype = out_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); - - // 创建 Info 对象 - auto result = AddbmmInfo::create(out_desc, in_desc, batch1_desc, batch2_desc, alpha, beta); - CHECK_RESULT(result); - - auto handle = reinterpret_cast(handle_); - - *desc_ptr = new Descriptor( - nullptr, - result.take(), - 0, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (inputs.size() != 3) { - return INFINI_STATUS_BAD_PARAM; - } - - const void *input = inputs[0]; - const void *batch1 = inputs[1]; - const void *batch2 = inputs[2]; - - auto dtype = _info.dtype(); - - switch (dtype) { - case INFINI_DTYPE_F16: - calculate_impl(_info, output, input, batch1, batch2); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_BF16: - calculate_impl(_info, output, input, batch1, batch2); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_F32: - calculate_impl(_info, output, input, batch1, batch2); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_F64: - calculate_impl(_info, output, input, batch1, batch2); - return INFINI_STATUS_SUCCESS; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::addbmm::cpu diff --git a/src/infiniop/ops/addbmm/cpu/addbmm_cpu.h b/src/infiniop/ops/addbmm/cpu/addbmm_cpu.h deleted file mode 100644 index 2b726d7bb..000000000 --- a/src/infiniop/ops/addbmm/cpu/addbmm_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADDBMM_CPU_H__ -#define __ADDBMM_CPU_H__ - -#include "../addbmm.h" - -DESCRIPTOR(cpu) - -#endif //_GRID_CPU_H_ diff --git a/src/infiniop/ops/addbmm/cuda/kernel.cuh b/src/infiniop/ops/addbmm/cuda/kernel.cuh deleted file mode 100644 index 4a7c057b5..000000000 --- a/src/infiniop/ops/addbmm/cuda/kernel.cuh +++ /dev/null @@ -1,274 +0,0 @@ -#ifndef __ADDBMM_NVIDIA_CUH__ -#define __ADDBMM_NVIDIA_CUH__ - -#include - -namespace op::addbmm::nvidia { - -// --- 常量定义 --- -constexpr int BLOCK_SIZE = 16; // 16x16 线程块，处理 FP32/FP16 比较通用 - -template -__device__ __forceinline__ float to_float_acc(const T &x) { - if constexpr (std::is_same_v) { - return __half2float(x); - } -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - else if constexpr (std::is_same_v) { - return __bfloat162float(x); - } -#endif - else { - return static_cast(x); - } -} - -template -__device__ __forceinline__ T from_float_res(float x) { - if constexpr (std::is_same_v) { - return __float2half(x); - } -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - else if constexpr (std::is_same_v) { - return __float2bfloat16(x); - } -#endif - else { - return static_cast(x); - } -} - -template -__global__ void addbmm_tiled_kernel( - T *output, - const T *input, - const T *batch1, - const T *batch2, - size_t b, size_t n, size_t m, size_t p, - float alpha, float beta, - // Strides - ptrdiff_t out_s0, ptrdiff_t out_s1, - ptrdiff_t inp_s0, ptrdiff_t inp_s1, - ptrdiff_t b1_s0, ptrdiff_t b1_s1, ptrdiff_t b1_s2, - ptrdiff_t b2_s0, ptrdiff_t b2_s1, ptrdiff_t b2_s2) { - // Block 行列索引 (Output 的坐标) - int row = blockIdx.y * BLOCK_SIZE + threadIdx.y; - int col = blockIdx.x * BLOCK_SIZE + threadIdx.x; - - float acc = 0.0f; - - // Shared Memory 缓存 - // 大小: 2个矩阵 * BLOCK_SIZE * BLOCK_SIZE - __shared__ float s_b1[BLOCK_SIZE][BLOCK_SIZE]; - __shared__ float s_b2[BLOCK_SIZE][BLOCK_SIZE]; - - // 遍历每一个 Batch - for (int batch_idx = 0; batch_idx < b; ++batch_idx) { - - // 遍历 K 维度 (即 m 维度)，步长为 BLOCK_SIZE - for (int k = 0; k < m; k += BLOCK_SIZE) { - - if (row < n && (k + threadIdx.x) < m) { - // b1: [batch_idx, row, k + tx] - size_t idx = batch_idx * b1_s0 + row * b1_s1 + (k + threadIdx.x) * b1_s2; - s_b1[threadIdx.y][threadIdx.x] = to_float_acc(batch1[idx]); - } else { - s_b1[threadIdx.y][threadIdx.x] = 0.0f; - } - - if ((k + threadIdx.y) < m && col < p) { - // b2: [batch_idx, k + ty, col] - size_t idx = batch_idx * b2_s0 + (k + threadIdx.y) * b2_s1 + col * b2_s2; - s_b2[threadIdx.y][threadIdx.x] = to_float_acc(batch2[idx]); - } else { - s_b2[threadIdx.y][threadIdx.x] = 0.0f; - } - - // 等待所有线程加载完毕 - __syncthreads(); - -// 2. 计算子块乘积 (Partial Accumulation) -// 循环展开以提高指令吞吐量 -#pragma unroll - for (int e = 0; e < BLOCK_SIZE; ++e) { - acc += s_b1[threadIdx.y][e] * s_b2[e][threadIdx.x]; - } - - // 等待计算完毕，准备加载下一个 tile - __syncthreads(); - } - } - - // 3. 写入结果 (Scale + Add Input) - if (row < n && col < p) { - float res = 0.0f; - - // Input 部分: beta * input - if (beta != 0.0f) { - size_t inp_idx = row * inp_s0 + col * inp_s1; - res = beta * to_float_acc(input[inp_idx]); - } - - // Matmul 部分: alpha * sum - res += alpha * acc; - - size_t out_idx = row * out_s0 + col * out_s1; - output[out_idx] = from_float_res(res); - } -} - -// ================================================================== -// 2. Launcher -// ================================================================== -template -void launch_kernel( - void *output, const void *input, const void *batch1, const void *batch2, - size_t b, size_t n, size_t m, size_t p, - float alpha, float beta, - // Strides - ptrdiff_t out_s0, ptrdiff_t out_s1, - ptrdiff_t inp_s0, ptrdiff_t inp_s1, - ptrdiff_t b1_s0, ptrdiff_t b1_s1, ptrdiff_t b1_s2, - ptrdiff_t b2_s0, ptrdiff_t b2_s1, ptrdiff_t b2_s2, - void *stream) { - - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - auto b1_ptr = reinterpret_cast(batch1); - auto b2_ptr = reinterpret_cast(batch2); - - // 2D Grid 配置 - dim3 block(BLOCK_SIZE, BLOCK_SIZE); // 16x16 threads - dim3 grid( - (p + BLOCK_SIZE - 1) / BLOCK_SIZE, // x轴覆盖 col (p) - (n + BLOCK_SIZE - 1) / BLOCK_SIZE // y轴覆盖 row (n) - ); - - auto cuda_stream = reinterpret_cast(stream); - - addbmm_tiled_kernel<<>>( - out_ptr, in_ptr, b1_ptr, b2_ptr, - b, n, m, p, - alpha, beta, - out_s0, out_s1, inp_s0, inp_s1, - b1_s0, b1_s1, b1_s2, b2_s0, b2_s1, b2_s2); -} - -// // tiled version fails in some cases -template -__global__ void addbmm_kernel( - const size_t B, const size_t N, const size_t M, const size_t P, - const float alpha, const float beta, - T *output, - const T *input, - const T *batch1, - const T *batch2, - const int64_t out_s0, const int64_t out_s1, - const int64_t in_s0, const int64_t in_s1, - const int64_t b1_s0, const int64_t b1_s1, const int64_t b1_s2, - const int64_t b2_s0, const int64_t b2_s1, const int64_t b2_s2) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - size_t total_elements = N * P; - - if (idx < total_elements) { - size_t n = idx / P; - size_t p = idx % P; - - float matmul_sum = 0.0f; - - int64_t b1_n_offset = n * b1_s1; - int64_t b2_p_offset = p * b2_s2; - - for (size_t b = 0; b < B; ++b) { - int64_t b1_b_offset = b * b1_s0; - int64_t b2_b_offset = b * b2_s0; - - for (size_t m = 0; m < M; ++m) { - // 直接计算偏移：Batch1[b, n, m] - int64_t offset1 = b1_b_offset + b1_n_offset + m * b1_s2; - // 直接计算偏移：Batch2[b, m, p] - int64_t offset2 = b2_b_offset + m * b2_s1 + b2_p_offset; - - T val1 = batch1[offset1]; - T val2 = batch2[offset2]; - - float v1_f, v2_f; - if constexpr (std::is_same_v) { - v1_f = __half2float(val1); - v2_f = __half2float(val2); - } else if constexpr (std::is_same_v) { - v1_f = __bfloat162float(val1); - v2_f = __bfloat162float(val2); - } else { - v1_f = static_cast(val1); - v2_f = static_cast(val2); - } - matmul_sum += v1_f * v2_f; - } - } - - // 直接计算偏移：Input[n, p] - int64_t in_offset = n * in_s0 + p * in_s1; - T in_val = input[in_offset]; - - float in_val_f; - if constexpr (std::is_same_v) { - in_val_f = __half2float(in_val); - } else if constexpr (std::is_same_v) { - in_val_f = __bfloat162float(in_val); - } else { - in_val_f = static_cast(in_val); - } - - float result = beta * in_val_f + alpha * matmul_sum; - - // 直接计算偏移：Output[n, p] - int64_t out_offset = n * out_s0 + p * out_s1; - - if constexpr (std::is_same_v) { - output[out_offset] = __float2half(result); - } else if constexpr (std::is_same_v) { - output[out_offset] = __float2bfloat16(result); - } else { - output[out_offset] = static_cast(result); - } - } -} - -// ================================================================== -// 2. Launcher Implementation -// ================================================================== - -template -void launch_addbmm_naive( - void *output, const void *input, - const void *batch1, const void *batch2, - size_t b, size_t n, size_t m, size_t p, - float alpha, float beta, - ptrdiff_t out_s0, ptrdiff_t out_s1, - ptrdiff_t inp_s0, ptrdiff_t inp_s1, - ptrdiff_t b1_s0, ptrdiff_t b1_s1, ptrdiff_t b1_s2, - ptrdiff_t b2_s0, ptrdiff_t b2_s1, ptrdiff_t b2_s2, - void *stream) { - - size_t total_elements = n * p; - int threads = 256; - int blocks = (total_elements + threads - 1) / threads; - - addbmm_kernel<<>>( - b, n, m, p, - alpha, beta, - (T *)output, - (const T *)input, - (const T *)batch1, - (const T *)batch2, - out_s0, out_s1, - inp_s0, inp_s1, - b1_s0, b1_s1, b1_s2, - b2_s0, b2_s1, b2_s2); -} - -} // namespace op::addbmm::nvidia - -#endif // __ADDBMM_NVIDIA_CUH__ diff --git a/src/infiniop/ops/addbmm/info.h b/src/infiniop/ops/addbmm/info.h deleted file mode 100644 index 690d2be13..000000000 --- a/src/infiniop/ops/addbmm/info.h +++ /dev/null @@ -1,124 +0,0 @@ -#ifndef __ADDBMM_INFO_H__ -#define __ADDBMM_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::addbmm { - -class AddbmmInfo { - AddbmmInfo() = default; - -public: - // 矩阵运算维度的定义: - // batch1: (b, n, m) - // batch2: (b, m, p) - // input/output: (n, p) - size_t _b; - size_t _n; - size_t _m; - size_t _p; - - // 【新增】步长信息 (Strides) - // 用于处理非连续内存布局 - std::vector _out_strides; // [stride_n, stride_p] - std::vector _in_strides; // [stride_n, stride_p] - std::vector _b1_strides; // [stride_b, stride_n, stride_m] - std::vector _b2_strides; // [stride_b, stride_m, stride_p] - - // 标量系数 - float _alpha; - float _beta; - - // 数据类型 - int _dtype; - - // Getters - size_t b() const { return _b; } - size_t n() const { return _n; } - size_t m() const { return _m; } - size_t p() const { return _p; } - - // 【新增】Strides Getters - const std::vector &out_strides() const { return _out_strides; } - const std::vector &in_strides() const { return _in_strides; } - const std::vector &b1_strides() const { return _b1_strides; } - const std::vector &b2_strides() const { return _b2_strides; } - - float alpha() const { return _alpha; } - float beta() const { return _beta; } - int dtype() const { return _dtype; } - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - infiniopTensorDescriptor_t batch1_desc, - infiniopTensorDescriptor_t batch2_desc, - float alpha, - float beta) { - - // 1. 检查数据类型一致性 - int dtype = out_desc->dtype(); - if (in_desc->dtype() != dtype || batch1_desc->dtype() != dtype || batch2_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // 2. 检查维度数量 (ndim) - if (batch1_desc->ndim() != 3 || batch2_desc->ndim() != 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (in_desc->ndim() != 2 || out_desc->ndim() != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const auto &b1_shape = batch1_desc->shape(); - const auto &b2_shape = batch2_desc->shape(); - const auto &in_shape = in_desc->shape(); - const auto &out_shape = out_desc->shape(); - - // 3. 解析并校验维度 - size_t b = b1_shape[0]; - size_t n = b1_shape[1]; - size_t m = b1_shape[2]; - - if (b2_shape[0] != b || b2_shape[1] != m) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - size_t p = b2_shape[2]; - - if (out_shape[0] != n || out_shape[1] != p) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (in_shape[0] != n || in_shape[1] != p) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // 4. 【新增】提取 Strides - auto out_strides = out_desc->strides(); - auto in_strides = in_desc->strides(); - auto b1_strides = batch1_desc->strides(); - auto b2_strides = batch2_desc->strides(); - - // 5. 返回 Info 对象 - AddbmmInfo info; - info._b = b; - info._n = n; - info._m = m; - info._p = p; - info._out_strides = out_strides; - info._in_strides = in_strides; - info._b1_strides = b1_strides; - info._b2_strides = b2_strides; - info._alpha = alpha; - info._beta = beta; - info._dtype = dtype; - - return utils::Result(info); - } -}; - -} // namespace op::addbmm - -#endif // __ADDBMM_INFO_H__ diff --git a/src/infiniop/ops/addbmm/metax/addbmm_metax.h b/src/infiniop/ops/addbmm/metax/addbmm_metax.h deleted file mode 100644 index b8a007173..000000000 --- a/src/infiniop/ops/addbmm/metax/addbmm_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADDBMM_METAX_H__ -#define __ADDBMM_METAX_H__ - -#include "../addbmm.h" - -DESCRIPTOR(metax) - -#endif // __ADDBMM_METAX_H__ diff --git a/src/infiniop/ops/addbmm/metax/addbmm_metax.maca b/src/infiniop/ops/addbmm/metax/addbmm_metax.maca deleted file mode 100644 index 1093a962c..000000000 --- a/src/infiniop/ops/addbmm/metax/addbmm_metax.maca +++ /dev/null @@ -1,385 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "addbmm_metax.h" -#include -#include -#include - -namespace op::addbmm::metax { - -// ================================================================== -// 辅助 Kernel -// ================================================================== -template -__global__ void copy_tensor_kernel( - T *dst, - const T *src, - int batch, int rows, int cols, - int64_t stride_b, int64_t stride_h, int64_t stride_w) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total = batch * rows * cols; - if (idx >= total) { - return; - } - - int c = idx % cols; - int r = (idx / cols) % rows; - int b = idx / (rows * cols); - - int64_t src_offset = b * stride_b + r * stride_h + c * stride_w; - dst[idx] = src[src_offset]; -} - -template -__global__ void copy_back_kernel( - T *dst, - const T *src, - int rows, int cols, - int64_t stride_h, int64_t stride_w) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total = rows * cols; - if (idx >= total) { - return; - } - - int c = idx % cols; - int r = idx / cols; - - int64_t dst_offset = r * stride_h + c * stride_w; - dst[dst_offset] = src[idx]; -} - -bool needs_copy(const std::vector &strides, size_t rows, size_t cols) { - if (strides.empty()) { - return false; - } - - int64_t stride_row = strides[strides.size() - 2]; - int64_t stride_col = strides[strides.size() - 1]; - - if (stride_col != 1) { - return true; - } - if (stride_row != static_cast(cols)) { - return true; - } - return false; -} - -std::pair get_rc_strides(const std::vector &s) { - if (s.size() < 2) { - return {0, 1}; // Should not happen - } - return {s[s.size() - 2], s[s.size() - 1]}; -} - -// ================================================================== -// Descriptor 实现 -// ================================================================== - -struct Descriptor::Opaque { - std::shared_ptr internal; - bool copy_b1 = false; - bool copy_b2 = false; - bool copy_out = false; - size_t size_b1 = 0; - size_t size_b2 = 0; - size_t size_out = 0; - size_t offset_b1 = 0; - size_t offset_b2 = 0; - size_t offset_out = 0; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float alpha, - float beta) { - - auto handle = reinterpret_cast(handle_); - - if (input_desc_vec.size() != 3) { - return INFINI_STATUS_BAD_PARAM; - } - - auto in_desc = input_desc_vec[0]; - auto b1_desc = input_desc_vec[1]; - auto b2_desc = input_desc_vec[2]; - - auto dtype = out_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = AddbmmInfo::create(out_desc, in_desc, b1_desc, b2_desc, alpha, beta); - CHECK_RESULT(result); - auto info = result.take(); - - auto opaque = new Opaque{handle->internal()}; - - size_t dtype_size = (dtype == INFINI_DTYPE_F32) ? 4 : 2; - size_t total_workspace = 0; - - // B1: (b, n, m) - if (needs_copy(info.b1_strides(), info.n(), info.m())) { - opaque->copy_b1 = true; - opaque->size_b1 = info.b() * info.n() * info.m() * dtype_size; - opaque->offset_b1 = total_workspace; - total_workspace += opaque->size_b1; - total_workspace = (total_workspace + 255) / 256 * 256; - } - - // B2: (b, m, p) - if (needs_copy(info.b2_strides(), info.m(), info.p())) { - opaque->copy_b2 = true; - opaque->size_b2 = info.b() * info.m() * info.p() * dtype_size; - opaque->offset_b2 = total_workspace; - total_workspace += opaque->size_b2; - total_workspace = (total_workspace + 255) / 256 * 256; - } - - // Out: (n, p) - if (needs_copy(info.out_strides(), info.n(), info.p())) { - opaque->copy_out = true; - opaque->size_out = info.n() * info.p() * dtype_size; - opaque->offset_out = total_workspace; - total_workspace += opaque->size_out; - } - - *desc_ptr = new Descriptor(opaque, info, total_workspace, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (inputs.size() != 3) { - return INFINI_STATUS_BAD_PARAM; - } - - const void *input_ptr = inputs[0]; - const void *b1_ptr = inputs[1]; - const void *b2_ptr = inputs[2]; - auto hc_stream = (hcStream_t)stream; - - decltype(HPCC_R_32F) a_type, b_type, c_type; - hcblasComputeType_t compute_type; - - switch (_info.dtype()) { - case INFINI_DTYPE_F16: - a_type = b_type = c_type = HPCC_R_16F; - compute_type = HCBLAS_COMPUTE_32F; - break; - case INFINI_DTYPE_BF16: - a_type = b_type = c_type = HPCC_R_16BF; - compute_type = HCBLAS_COMPUTE_32F; - break; - case INFINI_DTYPE_F32: - a_type = b_type = c_type = HPCC_R_32F; - compute_type = HCBLAS_COMPUTE_32F; - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - char *ws_base = (char *)workspace; - - // --- 准备 B1 --- - const void *b1_active = b1_ptr; - - long long strideb_blas = _info.b1_strides()[0]; - int ldb_blas = static_cast(_info.b1_strides()[1]); - - if (_opaque->copy_b1) { - void *b1_ws = ws_base + _opaque->offset_b1; - size_t total = _info.b() * _info.n() * _info.m(); - size_t block = 256; - size_t grid = (total + block - 1) / block; - - // 提取 stride - auto &s = _info.b1_strides(); - int64_t sb = s[0], sh = s[1], sw = s[2]; - - if (_info.dtype() == INFINI_DTYPE_F32) { - copy_tensor_kernel<<>>( - (float *)b1_ws, (const float *)b1_ptr, - _info.b(), _info.n(), _info.m(), sb, sh, sw); - } else if (_info.dtype() == INFINI_DTYPE_F16) { - copy_tensor_kernel<<>>( - (half *)b1_ws, (const half *)b1_ptr, - _info.b(), _info.n(), _info.m(), sb, sh, sw); - } else { - copy_tensor_kernel<<>>( - (cuda_bfloat16 *)b1_ws, (const cuda_bfloat16 *)b1_ptr, - _info.b(), _info.n(), _info.m(), sb, sh, sw); - } - b1_active = b1_ws; - strideb_blas = _info.n() * _info.m(); - ldb_blas = _info.m(); - } - - // --- 准备 B2 --- - const void *b2_active = b2_ptr; - long long stridea_blas = _info.b2_strides()[0]; - int lda_blas = static_cast(_info.b2_strides()[1]); - - if (_opaque->copy_b2) { - void *b2_ws = ws_base + _opaque->offset_b2; - size_t total = _info.b() * _info.m() * _info.p(); - size_t block = 256; - size_t grid = (total + block - 1) / block; - - auto &s = _info.b2_strides(); - int64_t sb = s[0], sh = s[1], sw = s[2]; - - if (_info.dtype() == INFINI_DTYPE_F32) { - copy_tensor_kernel<<>>( - (float *)b2_ws, (const float *)b2_ptr, - _info.b(), _info.m(), _info.p(), sb, sh, sw); - } else if (_info.dtype() == INFINI_DTYPE_F16) { - copy_tensor_kernel<<>>( - (half *)b2_ws, (const half *)b2_ptr, - _info.b(), _info.m(), _info.p(), sb, sh, sw); - } else { - copy_tensor_kernel<<>>( - (cuda_bfloat16 *)b2_ws, (const cuda_bfloat16 *)b2_ptr, - _info.b(), _info.m(), _info.p(), sb, sh, sw); - } - b2_active = b2_ws; - stridea_blas = _info.m() * _info.p(); - lda_blas = _info.p(); - } - - // --- 准备 Output / Input(Beta) --- - void *out_active = output; - - auto out_rc = get_rc_strides(_info.out_strides()); - int ldc_blas = static_cast(out_rc.first); // Row Stride - - if (_opaque->copy_out) { - out_active = ws_base + _opaque->offset_out; - ldc_blas = _info.p(); - - size_t total = _info.n() * _info.p(); - size_t block = 256; - size_t grid = (total + block - 1) / block; - - auto in_rc = get_rc_strides(_info.in_strides()); - - if (_info.dtype() == INFINI_DTYPE_F32) { - copy_tensor_kernel<<>>( - (float *)out_active, (const float *)input_ptr, - 1, _info.n(), _info.p(), - 0, in_rc.first, in_rc.second); - } else if (_info.dtype() == INFINI_DTYPE_F16) { - copy_tensor_kernel<<>>( - (half *)out_active, (const half *)input_ptr, - 1, _info.n(), _info.p(), - 0, in_rc.first, in_rc.second); - } else { - copy_tensor_kernel<<>>( - (cuda_bfloat16 *)out_active, (const cuda_bfloat16 *)input_ptr, - 1, _info.n(), _info.p(), - 0, in_rc.first, in_rc.second); - } - } else { - if (output != input_ptr) { - size_t total = _info.n() * _info.p(); - size_t block = 256; - size_t grid = (total + block - 1) / block; - - // [修复] 强制使用 kernel copy，并安全获取 stride - auto in_rc = get_rc_strides(_info.in_strides()); - - if (_info.dtype() == INFINI_DTYPE_F32) { - copy_tensor_kernel<<>>( - (float *)output, (const float *)input_ptr, - 1, _info.n(), _info.p(), - 0, in_rc.first, in_rc.second); - } else if (_info.dtype() == INFINI_DTYPE_F16) { - copy_tensor_kernel<<>>( - (half *)output, (const half *)input_ptr, - 1, _info.n(), _info.p(), - 0, in_rc.first, in_rc.second); - } else { - copy_tensor_kernel<<>>( - (cuda_bfloat16 *)output, (const cuda_bfloat16 *)input_ptr, - 1, _info.n(), _info.p(), - 0, in_rc.first, in_rc.second); - } - } - } - - // --- 执行 BLAS 计算 --- - int m_blas = static_cast(_info.p()); - int n_blas = static_cast(_info.n()); - int k_blas = static_cast(_info.m()); - float alpha = _info.alpha(); - float beta_user = _info.beta(); - size_t batch_size = _info.b(); - int data_width = (_info.dtype() == INFINI_DTYPE_F32 ? 4 : 2); - - CHECK_STATUS(_opaque->internal->useMcblas( - (hcStream_t)stream, - [&](hcblasHandle_t handle) { - for (size_t i = 0; i < batch_size; ++i) { - float current_beta = (i == 0) ? beta_user : 1.0f; - - const void *curr_b2 = static_cast(b2_active) + i * stridea_blas * data_width; - const void *curr_b1 = static_cast(b1_active) + i * strideb_blas * data_width; - - hcblasGemmEx( - handle, - HCBLAS_OP_N, HCBLAS_OP_N, - m_blas, n_blas, k_blas, - &alpha, - curr_b2, a_type, lda_blas, - curr_b1, b_type, ldb_blas, - ¤t_beta, - out_active, c_type, ldc_blas, - compute_type, - HCBLAS_GEMM_DEFAULT); - } - return INFINI_STATUS_SUCCESS; - })); - - // --- 拷回结果 --- - if (_opaque->copy_out) { - size_t total = _info.n() * _info.p(); - size_t block = 256; - size_t grid = (total + block - 1) / block; - - // [修复] 安全获取 output stride - auto out_rc = get_rc_strides(_info.out_strides()); - - if (_info.dtype() == INFINI_DTYPE_F32) { - copy_back_kernel<<>>( - (float *)output, (const float *)out_active, - _info.n(), _info.p(), - out_rc.first, out_rc.second); - } else if (_info.dtype() == INFINI_DTYPE_F16) { - copy_back_kernel<<>>( - (half *)output, (const half *)out_active, - _info.n(), _info.p(), - out_rc.first, out_rc.second); - } else { - copy_back_kernel<<>>( - (cuda_bfloat16 *)output, (const cuda_bfloat16 *)out_active, - _info.n(), _info.p(), - out_rc.first, out_rc.second); - } - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::addbmm::metax diff --git a/src/infiniop/ops/addbmm/moore/addbmm_moore.h b/src/infiniop/ops/addbmm/moore/addbmm_moore.h deleted file mode 100644 index 86ef1957d..000000000 --- a/src/infiniop/ops/addbmm/moore/addbmm_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADDBMM_MOORE_API_H__ -#define __ADDBMM_MOORE_API_H__ - -#include "../addbmm.h" - -DESCRIPTOR(moore) - -#endif // __ADDBMM_MOORE_API_H__ diff --git a/src/infiniop/ops/addbmm/moore/addbmm_moore.mu b/src/infiniop/ops/addbmm/moore/addbmm_moore.mu deleted file mode 100644 index 3d605b1f8..000000000 --- a/src/infiniop/ops/addbmm/moore/addbmm_moore.mu +++ /dev/null @@ -1,237 +0,0 @@ -#include "addbmm_moore.h" - -#include "../../../devices/moore/moore_handle.h" -#include -#include -#include -#include - -namespace op::addbmm::moore { - -// ================================================================== -// 1. Kernel Implementation (Removed Macros) -// ================================================================== - -template -__global__ void addbmm_kernel( - const size_t B, const size_t N, const size_t M, const size_t P, - const float alpha, const float beta, - T *output, - const T *input, - const T *batch1, - const T *batch2, - const int64_t out_s0, const int64_t out_s1, - const int64_t in_s0, const int64_t in_s1, - const int64_t b1_s0, const int64_t b1_s1, const int64_t b1_s2, - const int64_t b2_s0, const int64_t b2_s1, const int64_t b2_s2) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - size_t total_elements = N * P; - - if (idx < total_elements) { - size_t n = idx / P; - size_t p = idx % P; - - float matmul_sum = 0.0f; - - // 预先计算与 b 无关的偏移量部分，略微优化性能 - int64_t b1_n_offset = n * b1_s1; - int64_t b2_p_offset = p * b2_s2; - - for (size_t b = 0; b < B; ++b) { - int64_t b1_b_offset = b * b1_s0; - int64_t b2_b_offset = b * b2_s0; - - for (size_t m = 0; m < M; ++m) { - // 直接计算偏移：Batch1[b, n, m] - int64_t offset1 = b1_b_offset + b1_n_offset + m * b1_s2; - // 直接计算偏移：Batch2[b, m, p] - int64_t offset2 = b2_b_offset + m * b2_s1 + b2_p_offset; - - T val1 = batch1[offset1]; - T val2 = batch2[offset2]; - - float v1_f, v2_f; - if constexpr (std::is_same_v) { - v1_f = __half2float(val1); - v2_f = __half2float(val2); - } else if constexpr (std::is_same_v) { - v1_f = __bfloat162float(val1); - v2_f = __bfloat162float(val2); - } else { - v1_f = static_cast(val1); - v2_f = static_cast(val2); - } - matmul_sum += v1_f * v2_f; - } - } - - // 直接计算偏移：Input[n, p] - int64_t in_offset = n * in_s0 + p * in_s1; - T in_val = input[in_offset]; - - float in_val_f; - if constexpr (std::is_same_v) { - in_val_f = __half2float(in_val); - } else if constexpr (std::is_same_v) { - in_val_f = __bfloat162float(in_val); - } else { - in_val_f = static_cast(in_val); - } - - float result = beta * in_val_f + alpha * matmul_sum; - - // 直接计算偏移：Output[n, p] - int64_t out_offset = n * out_s0 + p * out_s1; - - if constexpr (std::is_same_v) { - output[out_offset] = __float2half(result); - } else if constexpr (std::is_same_v) { - output[out_offset] = __float2bfloat16(result); - } else { - output[out_offset] = static_cast(result); - } - } -} - -// ================================================================== -// 2. Launcher Implementation -// ================================================================== - -template -void addbmm_moore_launch( - const AddbmmInfo &info, - T *output, - const T *input, - const T *batch1, - const T *batch2, - void *stream) { - - size_t total_elements = info.n() * info.p(); - int threads = 256; - int blocks = (total_elements + threads - 1) / threads; - - const auto &out_strides = info.out_strides(); - const auto &in_strides = info.in_strides(); - const auto &b1_strides = info.b1_strides(); - const auto &b2_strides = info.b2_strides(); - - addbmm_kernel<<>>( - info.b(), info.n(), info.m(), info.p(), - info.alpha(), info.beta(), - output, input, batch1, batch2, - out_strides[0], out_strides[1], - in_strides[0], in_strides[1], - b1_strides[0], b1_strides[1], b1_strides[2], - b2_strides[0], b2_strides[1], b2_strides[2]); -} - -// ================================================================== -// 3. Descriptor Implementation -// ================================================================== - -Descriptor::~Descriptor() = default; - -// 匹配 std::vector 接口 -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float alpha, - float beta) { - - if (input_desc_vec.size() != 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - infiniopTensorDescriptor_t in_desc = input_desc_vec[0]; - infiniopTensorDescriptor_t batch1_desc = input_desc_vec[1]; - infiniopTensorDescriptor_t batch2_desc = input_desc_vec[2]; - - auto handle = reinterpret_cast(handle_); - auto info_result = AddbmmInfo::create(out_desc, in_desc, batch1_desc, batch2_desc, alpha, beta); - - if (!info_result) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - *desc_ptr = new Descriptor( - nullptr, - *info_result, - 0, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -// 匹配 std::vector 接口 -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (inputs.size() != 3) { - return INFINI_STATUS_BAD_PARAM; - } - - const void *input = inputs[0]; - const void *batch1 = inputs[1]; - const void *batch2 = inputs[2]; - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_info.dtype()) { - case INFINI_DTYPE_F16: - addbmm_moore_launch( - _info, - static_cast(output), - static_cast(input), - static_cast(batch1), - static_cast(batch2), - stream); - break; - - case INFINI_DTYPE_BF16: - addbmm_moore_launch<__mt_bfloat16>( - _info, - static_cast<__mt_bfloat16 *>(output), - static_cast(input), - static_cast(batch1), - static_cast(batch2), - stream); - break; - - case INFINI_DTYPE_F32: - addbmm_moore_launch( - _info, - static_cast(output), - static_cast(input), - static_cast(batch1), - static_cast(batch2), - stream); - break; - - case INFINI_DTYPE_F64: - addbmm_moore_launch( - _info, - static_cast(output), - static_cast(input), - static_cast(batch1), - static_cast(batch2), - stream); - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::addbmm::moore diff --git a/src/infiniop/ops/addbmm/nvidia/addbmm_nvidia.cu b/src/infiniop/ops/addbmm/nvidia/addbmm_nvidia.cu deleted file mode 100644 index 95623b074..000000000 --- a/src/infiniop/ops/addbmm/nvidia/addbmm_nvidia.cu +++ /dev/null @@ -1,162 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" // Handle 定义 - -#include "../cuda/kernel.cuh" // Descriptor 基类定义 (关键！) -#include "addbmm_nvidia.cuh" // Tiled Kernel Launcher 定义 -#include "infinicore/ops/addbmm.hpp" // Descriptor 声明 -#include - -// ================================================================== -// 匿名命名空间：辅助函数 Wrapper -// ================================================================== -namespace { - -// 引用 Info 类 -using AddbmmInfo = ::op::addbmm::AddbmmInfo; - -// 泛型 Wrapper：负责从 Info 提取参数并调用底层 Launcher -template -void launch_kernel_wrapper( - void *output, - const void *input, - const void *batch1, - const void *batch2, - const AddbmmInfo &info, // 接收 Info 对象 - void *stream) { - - // 1. 提取维度 - size_t b = info.b(); - size_t n = info.n(); - size_t m = info.m(); - size_t p = info.p(); - float alpha = info.alpha(); - float beta = info.beta(); - - // 2. 提取 Strides - const auto &os = info.out_strides(); - const auto &is = info.in_strides(); - const auto &b1s = info.b1_strides(); - const auto &b2s = info.b2_strides(); - - // 3. 调用 .cuh 中的优化版 Launcher - // 【关键修复】不再使用 addbmm_kernel<<<...>>> - // 而是调用 op::addbmm::nvidia::launch_kernel - ::op::addbmm::nvidia::launch_addbmm_naive( - output, input, batch1, batch2, - b, n, m, p, - alpha, beta, - // 显式转换为 ptrdiff_t，匹配 .cuh 签名 - static_cast(os[0]), static_cast(os[1]), - static_cast(is[0]), static_cast(is[1]), - static_cast(b1s[0]), static_cast(b1s[1]), static_cast(b1s[2]), - static_cast(b2s[0]), static_cast(b2s[1]), static_cast(b2s[2]), - stream); -} - -} // anonymous namespace - -// ================================================================== -// Descriptor 成员函数实现 -// ================================================================== -namespace op::addbmm::nvidia { - -// Opaque 结构体定义 -struct Descriptor::Opaque {}; - -// 析构函数 -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -// Create 函数实现 -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, // 接收 Vector - float alpha, - float beta) { - - // 1. 参数校验 - if (input_desc_vec.size() != 3) { - return INFINI_STATUS_BAD_PARAM; - } - - // 2. 调用 Info::create 解析参数 - auto info_result = ::op::addbmm::AddbmmInfo::create( - out_desc, - input_desc_vec[0], // input - input_desc_vec[1], // batch1 - input_desc_vec[2], // batch2 - alpha, - beta); - - if (!info_result) { - return info_result.status(); - } - - // 3. 创建 Descriptor 实例 - *desc_ptr = new Descriptor( - new Opaque(), - info_result.take(), - 0, // Tiled Kernel 不需要 workspace - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -// Calculate 函数实现 -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - // 1. 参数校验 - if (inputs.size() != 3) { - return INFINI_STATUS_BAD_PARAM; - } - - const void *input_ptr = inputs[0]; - const void *batch1_ptr = inputs[1]; - const void *batch2_ptr = inputs[2]; - - // 2. 提取参数 - auto dtype = _info.dtype(); - - // 3. 分发 Kernel - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel_wrapper( - output, input_ptr, batch1_ptr, batch2_ptr, _info, stream); - break; - - case INFINI_DTYPE_BF16: - launch_kernel_wrapper( - output, input_ptr, batch1_ptr, batch2_ptr, _info, stream); - break; - - case INFINI_DTYPE_F32: - launch_kernel_wrapper( - output, input_ptr, batch1_ptr, batch2_ptr, _info, stream); - break; - - case INFINI_DTYPE_F64: - // 假设 double 也使用 Tiled Kernel (如果 .cuh 支持) - launch_kernel_wrapper( - output, input_ptr, batch1_ptr, batch2_ptr, _info, stream); - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::addbmm::nvidia diff --git a/src/infiniop/ops/addbmm/nvidia/addbmm_nvidia.cuh b/src/infiniop/ops/addbmm/nvidia/addbmm_nvidia.cuh deleted file mode 100644 index 3cedb1741..000000000 --- a/src/infiniop/ops/addbmm/nvidia/addbmm_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ADDBMM_CUH__ -#define __ADDBMM_CUH__ - -#include "../addbmm.h" - -DESCRIPTOR(nvidia) - -#endif // __GEMM_CUDA_CUH__ diff --git a/src/infiniop/ops/addbmm/operator.cc b/src/infiniop/ops/addbmm/operator.cc deleted file mode 100644 index 53ca28e2e..000000000 --- a/src/infiniop/ops/addbmm/operator.cc +++ /dev/null @@ -1,206 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/addbmm.h" - -// --- 后端实现头文件 --- -#ifdef ENABLE_CPU_API -#include "cpu/addbmm_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/addbmm_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/addbmm_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/addbmm_moore.h" -#endif - -// ======================================================================= -// [修复] 定义结构体 -// ======================================================================= -struct infiniopAddbmmDescriptor { - int device_type; -}; - -extern "C" { - -// ======================================================================= -// 1. 创建算子描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopCreateAddbmmDescriptor( - infiniopHandle_t handle, - infiniopAddbmmDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - infiniopTensorDescriptor_t batch1, - infiniopTensorDescriptor_t batch2, - float alpha, - float beta) { - -// 宏：根据不同后端调用对应的 C++ create 方法 -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::addbmm::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output, \ - {input, batch1, batch2}, \ - alpha, \ - beta) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ======================================================================= -// 2. 获取 Workspace 大小 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopGetAddbmmWorkspaceSize(infiniopAddbmmDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -// ======================================================================= -// 3. 执行计算 (Calculate) -// ======================================================================= -__INFINI_C infiniStatus_t infiniopAddbmm( - infiniopAddbmmDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *batch1, - const void *batch2, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, {input, batch1, batch2}, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -// ======================================================================= -// 4. 销毁描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopDestroyAddbmmDescriptor(infiniopAddbmmDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} - -} // extern "C" diff --git a/src/infiniop/ops/addcmul/cpu/addcmul_cpu.cc b/src/infiniop/ops/addcmul/cpu/addcmul_cpu.cc deleted file mode 100644 index 1e7cb479b..000000000 --- a/src/infiniop/ops/addcmul/cpu/addcmul_cpu.cc +++ /dev/null @@ -1,61 +0,0 @@ -#include "addcmul_cpu.h" - -namespace op::addcmul::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float value) { // 额外接收 value 参数 - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - // 1. 类型检查 - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - // 2. 形状检查 (仿照 atanh，这里至少检查第一个输入) - const auto &y_shape = out_desc->shape(); - for (const auto &in_desc : input_desc_vec) { - CHECK_SAME_SHAPE(y_shape, in_desc->shape()); - } - - // 3. 使用通用的 Elementwise 宏创建描述符 - // 该宏会实例化 Descriptor 并将其赋值给 *desc_ptr - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - // 4. 将标量属性 value 存入 Descriptor 内部 - (*desc_ptr)->_value = value; - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - // 仿照 atanh，使用 switch 分发不同数据类型 - // 这里的模板参数是 AddcmulOp，它在 addcmul_cpu.h 中定义 - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream, _value); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream, _value); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream, _value); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream, _value); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::addcmul::cpu diff --git a/src/infiniop/ops/addcmul/cpu/addcmul_cpu.h b/src/infiniop/ops/addcmul/cpu/addcmul_cpu.h deleted file mode 100644 index cd491ecf7..000000000 --- a/src/infiniop/ops/addcmul/cpu/addcmul_cpu.h +++ /dev/null @@ -1,78 +0,0 @@ -#ifndef __ADDCMUL_CPU_H__ -#define __ADDCMUL_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" -#include -#include - -namespace op::addcmul::cpu { - -struct AddcmulOp { -public: - // addcmul 是三元算子: out = input + value * t1 * t2 - static constexpr size_t num_inputs = 3; - - template - T operator()(const T &input, const T &t1, const T &t2, Scalar value) const { - // 对于 float, double 等原生浮点类型 - if constexpr (std::is_floating_point_v) { - return input + static_cast(value) * t1 * t2; - } else { - // 对于 fp16, bf16 等类型，提升至 float 计算以保证精度并处理标量乘法 - float f_input = static_cast(input); - float f_t1 = static_cast(t1); - float f_t2 = static_cast(t2); - float v = static_cast(value); - return static_cast(f_input + v * f_t1 * f_t2); - } - } -}; - -// 为 addcmul 在 CPU 端自定义 Descriptor，支持额外的标量参数 value -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - op::elementwise::ElementwiseInfo _info; - std::unique_ptr _device_info; - size_t _workspace_size; - float _value; // 标量系数 value - - Descriptor( - infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - op::elementwise::cpu::DeviceImpl *device_info, - size_t workspace_size, - infiniDevice_t device_type, - int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)), - _device_info(device_info), - _workspace_size(workspace_size), - _value(0.0f) {} - -public: - ~Descriptor(); - - size_t workspaceSize() const { return _workspace_size; } - - // 额外接收 value 参数 - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - std::vector input_descs, - float value); - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const; - - float getValue() const { return _value; } -}; - -} // namespace op::addcmul::cpu - -#endif // __ADDCMUL_CPU_H__ diff --git a/src/infiniop/ops/addcmul/cuda/kernel.cuh b/src/infiniop/ops/addcmul/cuda/kernel.cuh deleted file mode 100644 index 2ec83c26c..000000000 --- a/src/infiniop/ops/addcmul/cuda/kernel.cuh +++ /dev/null @@ -1,44 +0,0 @@ -#ifndef __ADDCMUL_CUDA_CUH__ -#define __ADDCMUL_CUDA_CUH__ - -#include - -namespace op::addcmul::cuda { - -struct AddcmulOp { -public: - // addcmul 是三元算子：out = input + value * t1 * t2 - static constexpr size_t num_inputs = 3; - - template - __device__ __host__ __forceinline__ T operator()(const T &input, const T &t1, const T &t2, float value) const { - float v = value; - if constexpr (std::is_same_v) { - // 提升至 float 计算以保证精度并简化标量乘法 - float f_input = __half2float(input); - float f_t1 = __half2float(t1); - float f_t2 = __half2float(t2); - return __float2half(f_input + v * f_t1 * f_t2); - - } else if constexpr (std::is_same_v) { - float f_input = __bfloat162float(input); - float f_t1 = __bfloat162float(t1); - float f_t2 = __bfloat162float(t2); - return __float2bfloat16(f_input + v * f_t1 * f_t2); - - } else if constexpr (std::is_same_v) { - return input + v * t1 * t2; - - } else if constexpr (std::is_same_v) { - return input + static_cast(v) * t1 * t2; - - } else { - // 兜底逻辑 - return static_cast(static_cast(input) + v * static_cast(t1) * static_cast(t2)); - } - } -}; - -} // namespace op::addcmul::cuda - -#endif // __ADDCMUL_CUDA_CUH__ diff --git a/src/infiniop/ops/addcmul/metax/addcmul_metax.h b/src/infiniop/ops/addcmul/metax/addcmul_metax.h deleted file mode 100644 index 556375fa1..000000000 --- a/src/infiniop/ops/addcmul/metax/addcmul_metax.h +++ /dev/null @@ -1,72 +0,0 @@ -#ifndef __ADDCMUL_METAX_H__ -#define __ADDCMUL_METAX_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -namespace op::addcmul::metax { - -// 为 addcmul 在 METAX 端自定义 Descriptor，支持额外的标量参数 value -class Descriptor final : public InfiniopDescriptor { - // 为保持与通用 Elementwise 框架的兼容，仍然保留这些成员 - infiniDtype_t _dtype; - op::elementwise::ElementwiseInfo _info; - std::unique_ptr _device_info; - size_t _workspace_size; - float _value; // 标量系数 value - -public: - // 为自定义 CUDA kernel 记录张量元信息 - static constexpr int MAX_NDIM = 8; - - struct TensorMeta { - int ndim; - size_t shape[MAX_NDIM]; - ptrdiff_t strides[MAX_NDIM]; - }; - - TensorMeta _out_meta{}; - TensorMeta _input_meta{}; - TensorMeta _t1_meta{}; - TensorMeta _t2_meta{}; - size_t _output_size{0}; - - Descriptor( - infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - op::elementwise::metax::DeviceImpl *device_info, - size_t workspace_size, - infiniDevice_t device_type, - int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)), - _device_info(device_info), - _workspace_size(workspace_size), - _value(0.0f) {} - -public: - ~Descriptor(); - - size_t workspaceSize() const { return _workspace_size; } - - // 额外接收 value 参数 - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - std::vector input_descs, - float value); - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const; - - float getValue() const { return _value; } -}; - -} // namespace op::addcmul::metax - -#endif // __ADDCMUL_METAX_H__ diff --git a/src/infiniop/ops/addcmul/metax/addcmul_metax.maca b/src/infiniop/ops/addcmul/metax/addcmul_metax.maca deleted file mode 100644 index 6d851217b..000000000 --- a/src/infiniop/ops/addcmul/metax/addcmul_metax.maca +++ /dev/null @@ -1,180 +0,0 @@ -#include "../../../devices/metax/metax_handle.h" -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "addcmul_metax.h" -#include "addcmul_metax_kernel.h" - -namespace op::addcmul::metax { - -Descriptor::~Descriptor() = default; - -// 将 TensorDescriptor 中的 shape/strides 填充到 TensorMeta 结构中 -static inline infiniStatus_t fill_tensor_meta( - infiniopTensorDescriptor_t desc, - Descriptor::TensorMeta &meta) { - - auto ndim = desc->ndim(); - if (ndim > Descriptor::MAX_NDIM) { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - - meta.ndim = static_cast(ndim); - const auto &shape = desc->shape(); - const auto &strides = desc->strides(); - for (int i = 0; i < meta.ndim; ++i) { - meta.shape[i] = shape[i]; - meta.strides[i] = strides[i]; - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float value) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - // 1. 类型检查 - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); - - // 2. 形状检查：要求输出与三个输入形状一致（若不支持广播） - const auto &out_shape = out_desc->shape(); - const auto &input_desc = input_desc_vec.at(0); - const auto &t1_desc = input_desc_vec.at(1); - const auto &t2_desc = input_desc_vec.at(2); - CHECK_SAME_SHAPE(out_shape, input_desc->shape()); - CHECK_SAME_SHAPE(out_shape, t1_desc->shape()); - CHECK_SAME_SHAPE(out_shape, t2_desc->shape()); - - // 3. 创建底层的 Elementwise METAX 描述符 - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - // 4. 记录张量元信息和输出元素个数，供自定义 METAX kernel 使用 - auto *desc = *desc_ptr; - desc->_output_size = out_desc->numel(); - - CHECK_STATUS(fill_tensor_meta(out_desc, desc->_out_meta)); - CHECK_STATUS(fill_tensor_meta(input_desc, desc->_input_meta)); - CHECK_STATUS(fill_tensor_meta(t1_desc, desc->_t1_meta)); - CHECK_STATUS(fill_tensor_meta(t2_desc, desc->_t2_meta)); - - // 5. 将标量属性 value 存入 Descriptor 内部 - desc->_value = value; - - return INFINI_STATUS_SUCCESS; -} - -// 自定义 addcmul METAX kernel：使用 Descriptor 中的 TensorMeta 做通用 strided 访问 -template -INFINIOP_METAX_KERNEL addcmul_kernel( - size_t output_size, - Descriptor::TensorMeta out_meta, - Descriptor::TensorMeta in_meta, - Descriptor::TensorMeta t1_meta, - Descriptor::TensorMeta t2_meta, - T *out, - const T *input, - const T *t1, - const T *t2, - float value) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= output_size) { - return; - } - - // 根据输出 shape/stride 计算各个张量的偏移 - ptrdiff_t out_offset = 0; - ptrdiff_t in_offset = 0; - ptrdiff_t t1_offset = 0; - ptrdiff_t t2_offset = 0; - - size_t linear = idx; - for (int dim = out_meta.ndim - 1; dim >= 0; --dim) { - size_t dim_size = out_meta.shape[dim]; - size_t coord = linear % dim_size; - linear /= dim_size; - - out_offset += static_cast(coord) * out_meta.strides[dim]; - in_offset += static_cast(coord) * in_meta.strides[dim]; - t1_offset += static_cast(coord) * t1_meta.strides[dim]; - t2_offset += static_cast(coord) * t2_meta.strides[dim]; - } - - T in_val = input[in_offset]; - T t1_val = t1[t1_offset]; - T t2_val = t2[t2_offset]; - - out[out_offset] = op::addcmul::metax::AddcmulOp{}(in_val, t1_val, t2_val, value); -} - -template -static inline infiniStatus_t launch_addcmul_kernel( - const Descriptor *desc, - void *output, - const std::vector &inputs, - void *stream) { - - size_t output_size = desc->_output_size; - if (output_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - auto *out_ptr = reinterpret_cast(output); - auto *in_ptr = reinterpret_cast(inputs.at(0)); - auto *t1_ptr = reinterpret_cast(inputs.at(1)); - auto *t2_ptr = reinterpret_cast(inputs.at(2)); - - hcStream_t metax_stream = reinterpret_cast(stream); - - constexpr uint32_t BLOCK_SIZE = 256; - uint32_t grid = static_cast((output_size + BLOCK_SIZE - 1) / BLOCK_SIZE); - - addcmul_kernel<<>>( - output_size, - desc->_out_meta, - desc->_input_meta, - desc->_t1_meta, - desc->_t2_meta, - out_ptr, - in_ptr, - t1_ptr, - t2_ptr, - desc->getValue()); - - CHECK_METAX(hcGetLastError()); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - // 目前不依赖 workspace 内容，只检查大小是否足够以保持与其他算子一致的接口语义 - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // 直接调用自定义 METAX kernel，避免通过通用 elementwise 框架 - switch (_dtype) { - case INFINI_DTYPE_F16: - return launch_addcmul_kernel(this, output, inputs, stream); - case INFINI_DTYPE_BF16: - return launch_addcmul_kernel(this, output, inputs, stream); - case INFINI_DTYPE_F32: - return launch_addcmul_kernel(this, output, inputs, stream); - case INFINI_DTYPE_F64: - return launch_addcmul_kernel(this, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} -} // namespace op::addcmul::metax diff --git a/src/infiniop/ops/addcmul/metax/addcmul_metax_kernel.h b/src/infiniop/ops/addcmul/metax/addcmul_metax_kernel.h deleted file mode 100644 index fa0c92472..000000000 --- a/src/infiniop/ops/addcmul/metax/addcmul_metax_kernel.h +++ /dev/null @@ -1,44 +0,0 @@ -#ifndef __ADDCMUL_METAX_KERNEL_H__ -#define __ADDCMUL_METAX_KERNEL_H__ - -/* - * This file contains the Addcmul operation implementation for the MUSA backend. - * Formula: out = input + value * tensor1 * tensor2 - */ - -namespace op::addcmul::metax { - -typedef struct AddcmulOp { -public: - // 三元算子，输入为 input, tensor1, tensor2 - static constexpr size_t num_inputs = 3; - - template - __device__ __forceinline__ T operator()(const T &in, const T &t1, const T &t2, float value) const { - if constexpr (std::is_same_v) { - // F32 直接使用乘加指令 - return in + value * t1 * t2; - } else if constexpr (std::is_same_v) { - // F16 提升到 float 计算以防止中间乘法溢出 - float f_in = __half2float(in); - float f_t1 = __half2float(t1); - float f_t2 = __half2float(t2); - return __float2half(f_in + value * f_t1 * f_t2); - } else if constexpr (std::is_same_v) { - // BF16 同样提升到 float 计算 - float f_in = __bfloat162float(in); - float f_t1 = __bfloat162float(t1); - float f_t2 = __bfloat162float(t2); - return __float2bfloat16_rn(f_in + value * f_t1 * f_t2); - } else if constexpr (std::is_same_v) { - return in + (double)value * t1 * t2; - } else { - // 整数类型或其他类型 - return in + static_cast(value) * t1 * t2; - } - } -} AddcmulOp; - -} // namespace op::addcmul::metax - -#endif // __ADDCMUL_METAX_KERNEL_H__ diff --git a/src/infiniop/ops/addcmul/moore/addcmul_moore.h b/src/infiniop/ops/addcmul/moore/addcmul_moore.h deleted file mode 100644 index 8e1ac68ec..000000000 --- a/src/infiniop/ops/addcmul/moore/addcmul_moore.h +++ /dev/null @@ -1,76 +0,0 @@ -#ifndef __ADDCMUL_MOORE_H__ -#define __ADDCMUL_MOORE_H__ - -// 1. 切换到 Moore 平台的 Elementwise API -#include "../../../elementwise/moore/elementwise_moore_api.h" - -namespace op::addcmul::moore { - -/** - * 为 addcmul 在 Moore 端自定义 Descriptor - * 保持与 NVIDIA 版本一致的结构，以便于跨平台对齐 - */ -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - op::elementwise::ElementwiseInfo _info; - // 2. 切换到 Moore 设备的实现指针 - std::unique_ptr _device_info; - size_t _workspace_size; - float _value; // 标量系数 value - -public: - // 摩尔线程 MUSA 同样支持 stride 访问，记录张量元信息 - static constexpr int MAX_NDIM = 8; - - struct TensorMeta { - int ndim; - size_t shape[MAX_NDIM]; - ptrdiff_t strides[MAX_NDIM]; - }; - - TensorMeta _out_meta{}; - TensorMeta _input_meta{}; - TensorMeta _t1_meta{}; - TensorMeta _t2_meta{}; - size_t _output_size{0}; - - Descriptor( - infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - op::elementwise::moore::DeviceImpl *device_info, // 3. 修改构造函数参数类型 - size_t workspace_size, - infiniDevice_t device_type, - int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)), - _device_info(device_info), - _workspace_size(workspace_size), - _value(0.0f) {} - -public: - ~Descriptor(); - - size_t workspaceSize() const { return _workspace_size; } - - // 4. 保持相同的接口，接收 value 参数 - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - std::vector input_descs, - float value); - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const; - - float getValue() const { return _value; } -}; - -} // namespace op::addcmul::moore - -#endif // __ADDCMUL_MOORE_H__ diff --git a/src/infiniop/ops/addcmul/moore/addcmul_moore.mu b/src/infiniop/ops/addcmul/moore/addcmul_moore.mu deleted file mode 100644 index 0c4d39099..000000000 --- a/src/infiniop/ops/addcmul/moore/addcmul_moore.mu +++ /dev/null @@ -1,163 +0,0 @@ -#include "../../../elementwise/moore/elementwise_moore.h" -#include "addcmul_moore.h" -#include "addcmul_moore_kernel.h" -#include - -namespace op::addcmul::moore { - -Descriptor::~Descriptor() = default; - -// 1. 填充 TensorMeta，逻辑与 NVIDIA 一致，用于 MUSA Kernel 中的 Strided 寻址 -static inline infiniStatus_t fill_tensor_meta( - infiniopTensorDescriptor_t desc, - Descriptor::TensorMeta &meta) { - - auto ndim = desc->ndim(); - if (ndim > Descriptor::MAX_NDIM) { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - - meta.ndim = static_cast(ndim); - const auto &shape = desc->shape(); - const auto &strides = desc->strides(); - for (int i = 0; i < meta.ndim; ++i) { - meta.shape[i] = shape[i]; - meta.strides[i] = strides[i]; - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float value) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - // 类型检查 - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); - - // 形状检查 (A, T1, T2 需一致) - const auto &out_shape = out_desc->shape(); - const auto &input_desc = input_desc_vec.at(0); - const auto &t1_desc = input_desc_vec.at(1); - const auto &t2_desc = input_desc_vec.at(2); - CHECK_SAME_SHAPE(out_shape, input_desc->shape()); - CHECK_SAME_SHAPE(out_shape, t1_desc->shape()); - CHECK_SAME_SHAPE(out_shape, t2_desc->shape()); - - // 2. 调用 Moore 平台的描述符创建宏 - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - auto *desc = *desc_ptr; - desc->_output_size = out_desc->numel(); - - // 填充元数据 - CHECK_STATUS(fill_tensor_meta(out_desc, desc->_out_meta)); - CHECK_STATUS(fill_tensor_meta(input_desc, desc->_input_meta)); - CHECK_STATUS(fill_tensor_meta(t1_desc, desc->_t1_meta)); - CHECK_STATUS(fill_tensor_meta(t2_desc, desc->_t2_meta)); - - desc->_value = value; - - return INFINI_STATUS_SUCCESS; -} - -// 3. MUSA Kernel 实现：逻辑保持一致 -template -__global__ void addcmul_kernel( - size_t output_size, - Descriptor::TensorMeta out_meta, - Descriptor::TensorMeta in_meta, - Descriptor::TensorMeta t1_meta, - Descriptor::TensorMeta t2_meta, - T *out, - const T *input, - const T *t1, - const T *t2, - float value) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= output_size) { - return; - } - - ptrdiff_t out_offset = 0, in_offset = 0, t1_offset = 0, t2_offset = 0; - size_t linear = idx; - - // 通用多维索引转偏移逻辑 - for (int dim = out_meta.ndim - 1; dim >= 0; --dim) { - size_t dim_size = out_meta.shape[dim]; - size_t coord = linear % dim_size; - linear /= dim_size; - - out_offset += static_cast(coord) * out_meta.strides[dim]; - in_offset += static_cast(coord) * in_meta.strides[dim]; - t1_offset += static_cast(coord) * t1_meta.strides[dim]; - t2_offset += static_cast(coord) * t2_meta.strides[dim]; - } - - // 调用 Moore 平台定义的 AddcmulOp - out[out_offset] = op::addcmul::moore::AddcmulOp{}(input[in_offset], t1[t1_offset], t2[t2_offset], value); -} - -// 4. 内核启动封装 -template -static inline infiniStatus_t launch_addcmul_kernel( - const Descriptor *desc, - void *output, - const std::vector &inputs, - void *stream) { - - size_t output_size = desc->_output_size; - if (output_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - auto *out_ptr = reinterpret_cast(output); - auto *in_ptr = reinterpret_cast(inputs.at(0)); - auto *t1_ptr = reinterpret_cast(inputs.at(1)); - auto *t2_ptr = reinterpret_cast(inputs.at(2)); - - musaStream_t musa_stream = reinterpret_cast(stream); - - constexpr uint32_t BLOCK_SIZE = 256; - uint32_t grid = static_cast((output_size + BLOCK_SIZE - 1) / BLOCK_SIZE); - - addcmul_kernel<<>>( - output_size, desc->_out_meta, desc->_input_meta, desc->_t1_meta, desc->_t2_meta, - out_ptr, in_ptr, t1_ptr, t2_ptr, desc->getValue()); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return launch_addcmul_kernel(this, output, inputs, stream); - case INFINI_DTYPE_BF16: - // 使用 Moore 平台对应的 bf16 类型 - return launch_addcmul_kernel(this, output, inputs, stream); - case INFINI_DTYPE_F32: - return launch_addcmul_kernel(this, output, inputs, stream); - case INFINI_DTYPE_F64: - return launch_addcmul_kernel(this, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} -} // namespace op::addcmul::moore diff --git a/src/infiniop/ops/addcmul/moore/addcmul_moore_kernel.h b/src/infiniop/ops/addcmul/moore/addcmul_moore_kernel.h deleted file mode 100644 index dc1788014..000000000 --- a/src/infiniop/ops/addcmul/moore/addcmul_moore_kernel.h +++ /dev/null @@ -1,44 +0,0 @@ -#ifndef __ADDCMUL_MOORE_KERNEL_H__ -#define __ADDCMUL_MOORE_KERNEL_H__ - -/* - * This file contains the Addcmul operation implementation for the MUSA backend. - * Formula: out = input + value * tensor1 * tensor2 - */ - -namespace op::addcmul::moore { - -typedef struct AddcmulOp { -public: - // 三元算子，输入为 input, tensor1, tensor2 - static constexpr size_t num_inputs = 3; - - template - __device__ __forceinline__ T operator()(const T &in, const T &t1, const T &t2, float value) const { - if constexpr (std::is_same_v) { - // F32 直接使用乘加指令 - return in + value * t1 * t2; - } else if constexpr (std::is_same_v) { - // F16 提升到 float 计算以防止中间乘法溢出 - float f_in = __half2float(in); - float f_t1 = __half2float(t1); - float f_t2 = __half2float(t2); - return __float2half(f_in + value * f_t1 * f_t2); - } else if constexpr (std::is_same_v) { - // BF16 同样提升到 float 计算 - float f_in = __bfloat162float(in); - float f_t1 = __bfloat162float(t1); - float f_t2 = __bfloat162float(t2); - return __float2bfloat16_rn(f_in + value * f_t1 * f_t2); - } else if constexpr (std::is_same_v) { - return in + (double)value * t1 * t2; - } else { - // 整数类型或其他类型 - return in + static_cast(value) * t1 * t2; - } - } -} AddcmulOp; - -} // namespace op::addcmul::moore - -#endif // __ADDCMUL_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/addcmul/nvidia/addcmul_nvidia.cu b/src/infiniop/ops/addcmul/nvidia/addcmul_nvidia.cu deleted file mode 100644 index ff83ca016..000000000 --- a/src/infiniop/ops/addcmul/nvidia/addcmul_nvidia.cu +++ /dev/null @@ -1,178 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" -#include "../cuda/kernel.cuh" -#include "addcmul_nvidia.cuh" - -namespace op::addcmul::nvidia { - -Descriptor::~Descriptor() = default; - -// 将 TensorDescriptor 中的 shape/strides 填充到 TensorMeta 结构中 -static inline infiniStatus_t fill_tensor_meta( - infiniopTensorDescriptor_t desc, - Descriptor::TensorMeta &meta) { - - auto ndim = desc->ndim(); - if (ndim > Descriptor::MAX_NDIM) { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - - meta.ndim = static_cast(ndim); - const auto &shape = desc->shape(); - const auto &strides = desc->strides(); - for (int i = 0; i < meta.ndim; ++i) { - meta.shape[i] = shape[i]; - meta.strides[i] = strides[i]; - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float value) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - // 1. 类型检查 - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); - - // 2. 形状检查：要求输出与三个输入形状一致（若不支持广播） - const auto &out_shape = out_desc->shape(); - const auto &input_desc = input_desc_vec.at(0); - const auto &t1_desc = input_desc_vec.at(1); - const auto &t2_desc = input_desc_vec.at(2); - CHECK_SAME_SHAPE(out_shape, input_desc->shape()); - CHECK_SAME_SHAPE(out_shape, t1_desc->shape()); - CHECK_SAME_SHAPE(out_shape, t2_desc->shape()); - - // 3. 创建底层的 Elementwise CUDA 描述符 - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - // 4. 记录张量元信息和输出元素个数，供自定义 CUDA kernel 使用 - auto *desc = *desc_ptr; - desc->_output_size = out_desc->numel(); - - CHECK_STATUS(fill_tensor_meta(out_desc, desc->_out_meta)); - CHECK_STATUS(fill_tensor_meta(input_desc, desc->_input_meta)); - CHECK_STATUS(fill_tensor_meta(t1_desc, desc->_t1_meta)); - CHECK_STATUS(fill_tensor_meta(t2_desc, desc->_t2_meta)); - - // 5. 将标量属性 value 存入 Descriptor 内部 - desc->_value = value; - - return INFINI_STATUS_SUCCESS; -} - -// 自定义 addcmul CUDA kernel：使用 Descriptor 中的 TensorMeta 做通用 strided 访问 -template -INFINIOP_CUDA_KERNEL addcmul_kernel( - size_t output_size, - Descriptor::TensorMeta out_meta, - Descriptor::TensorMeta in_meta, - Descriptor::TensorMeta t1_meta, - Descriptor::TensorMeta t2_meta, - T *out, - const T *input, - const T *t1, - const T *t2, - float value) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= output_size) { - return; - } - - // 根据输出 shape/stride 计算各个张量的偏移 - ptrdiff_t out_offset = 0; - ptrdiff_t in_offset = 0; - ptrdiff_t t1_offset = 0; - ptrdiff_t t2_offset = 0; - - size_t linear = idx; - for (int dim = out_meta.ndim - 1; dim >= 0; --dim) { - size_t dim_size = out_meta.shape[dim]; - size_t coord = linear % dim_size; - linear /= dim_size; - - out_offset += static_cast(coord) * out_meta.strides[dim]; - in_offset += static_cast(coord) * in_meta.strides[dim]; - t1_offset += static_cast(coord) * t1_meta.strides[dim]; - t2_offset += static_cast(coord) * t2_meta.strides[dim]; - } - - T in_val = input[in_offset]; - T t1_val = t1[t1_offset]; - T t2_val = t2[t2_offset]; - - out[out_offset] = op::addcmul::cuda::AddcmulOp{}(in_val, t1_val, t2_val, value); -} - -template -static inline infiniStatus_t launch_addcmul_kernel( - const Descriptor *desc, - void *output, - const std::vector &inputs, - void *stream) { - - size_t output_size = desc->_output_size; - if (output_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - auto *out_ptr = reinterpret_cast(output); - auto *in_ptr = reinterpret_cast(inputs.at(0)); - auto *t1_ptr = reinterpret_cast(inputs.at(1)); - auto *t2_ptr = reinterpret_cast(inputs.at(2)); - - cudaStream_t cuda_stream = reinterpret_cast(stream); - - constexpr uint32_t BLOCK_SIZE = 256; - uint32_t grid = static_cast((output_size + BLOCK_SIZE - 1) / BLOCK_SIZE); - - addcmul_kernel<<>>( - output_size, - desc->_out_meta, - desc->_input_meta, - desc->_t1_meta, - desc->_t2_meta, - out_ptr, - in_ptr, - t1_ptr, - t2_ptr, - desc->getValue()); - - CHECK_CUDA(cudaGetLastError()); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - // 目前不依赖 workspace 内容，只检查大小是否足够以保持与其他算子一致的接口语义 - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // 直接调用自定义 CUDA kernel，避免通过通用 elementwise 框架 - switch (_dtype) { - case INFINI_DTYPE_F16: - return launch_addcmul_kernel(this, output, inputs, stream); - case INFINI_DTYPE_BF16: - return launch_addcmul_kernel(this, output, inputs, stream); - case INFINI_DTYPE_F32: - return launch_addcmul_kernel(this, output, inputs, stream); - case INFINI_DTYPE_F64: - return launch_addcmul_kernel(this, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} -} // namespace op::addcmul::nvidia diff --git a/src/infiniop/ops/addcmul/nvidia/addcmul_nvidia.cuh b/src/infiniop/ops/addcmul/nvidia/addcmul_nvidia.cuh deleted file mode 100644 index 0e52a0a4f..000000000 --- a/src/infiniop/ops/addcmul/nvidia/addcmul_nvidia.cuh +++ /dev/null @@ -1,72 +0,0 @@ -#ifndef __ADDCMUL_NVIDIA_H__ -#define __ADDCMUL_NVIDIA_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -namespace op::addcmul::nvidia { - -// 为 addcmul 在 NVIDIA 端自定义 Descriptor，支持额外的标量参数 value -class Descriptor final : public InfiniopDescriptor { - // 为保持与通用 Elementwise 框架的兼容，仍然保留这些成员 - infiniDtype_t _dtype; - op::elementwise::ElementwiseInfo _info; - std::unique_ptr _device_info; - size_t _workspace_size; - float _value; // 标量系数 value - -public: - // 为自定义 CUDA kernel 记录张量元信息 - static constexpr int MAX_NDIM = 8; - - struct TensorMeta { - int ndim; - size_t shape[MAX_NDIM]; - ptrdiff_t strides[MAX_NDIM]; - }; - - TensorMeta _out_meta{}; - TensorMeta _input_meta{}; - TensorMeta _t1_meta{}; - TensorMeta _t2_meta{}; - size_t _output_size{0}; - - Descriptor( - infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - op::elementwise::nvidia::DeviceImpl *device_info, - size_t workspace_size, - infiniDevice_t device_type, - int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)), - _device_info(device_info), - _workspace_size(workspace_size), - _value(0.0f) {} - -public: - ~Descriptor(); - - size_t workspaceSize() const { return _workspace_size; } - - // 额外接收 value 参数 - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - std::vector input_descs, - float value); - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const; - - float getValue() const { return _value; } -}; - -} // namespace op::addcmul::nvidia - -#endif // __ADDCMUL_NVIDIA_H__ diff --git a/src/infiniop/ops/addcmul/operator.cc b/src/infiniop/ops/addcmul/operator.cc deleted file mode 100644 index ca071adb5..000000000 --- a/src/infiniop/ops/addcmul/operator.cc +++ /dev/null @@ -1,175 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/addcmul.h" - -#ifdef ENABLE_CPU_API -#include "cpu/addcmul_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/addcmul_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/addcmul_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/addcmul_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAddcmulDescriptor( - infiniopHandle_t handle, - infiniopAddcmulDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t t1_desc, - infiniopTensorDescriptor_t t2_desc, - float value) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::addcmul::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - out_desc, \ - {input_desc, t1_desc, t2_desc}, \ - value) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAddcmulWorkspaceSize(infiniopAddcmulDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopAddcmul( - infiniopAddcmulDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *input, - const void *t1, - const void *t2, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, out, {input, t1, t2}, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyAddcmulDescriptor(infiniopAddcmulDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} diff --git a/src/infiniop/ops/addr/addr.h b/src/infiniop/ops/addr/addr.h deleted file mode 100644 index 426bdd6b3..000000000 --- a/src/infiniop/ops/addr/addr.h +++ /dev/null @@ -1,95 +0,0 @@ -#ifndef INFINIOP_ADDR_DESCRIPTOR_H_ -#define INFINIOP_ADDR_DESCRIPTOR_H_ -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" -#include "infiniop/ops/addr.h" -#include - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::addr::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AddrInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - AddrInfo info, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t input_desc, \ - infiniopTensorDescriptor_t vec1_desc, \ - infiniopTensorDescriptor_t vec2_desc, \ - float beta, \ - float alpha); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *out, \ - const void *input, \ - const void *vec1, \ - const void *vec2, \ - void *stream) const; \ - }; \ - } - -struct AddrInfo { - infiniDtype_t dtype; - size_t vec1_size; - size_t vec2_size; - float beta; - float alpha; - ptrdiff_t input_stride0, input_stride1; - ptrdiff_t output_stride0, output_stride1; - ptrdiff_t vec1_stride; - ptrdiff_t vec2_stride; - - static utils::Result - create(infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t vec1_desc, - infiniopTensorDescriptor_t vec2_desc, - float beta = 1.0f, float alpha = 1.0f) { - CHECK_OR_RETURN(input_desc->ndim() == 2, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(vec1_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(vec2_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(input_desc->dim(0) == vec1_desc->dim(0) && input_desc->dim(1) == vec2_desc->dim(0), INFINI_STATUS_BAD_TENSOR_SHAPE); - const infiniDtype_t data_type = input_desc->dtype(); - CHECK_DTYPE(data_type, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - return utils::Result(AddrInfo{ - data_type, - vec1_desc->dim(0), - vec2_desc->dim(0), - beta, - alpha, - input_desc->stride(0), - input_desc->stride(1), - output_desc->stride(0), - output_desc->stride(1), - vec1_desc->stride(0), - vec2_desc->stride(0), - }); - } -}; - -#endif diff --git a/src/infiniop/ops/addr/cpu/addr_cpu.cc b/src/infiniop/ops/addr/cpu/addr_cpu.cc deleted file mode 100644 index f1069aef1..000000000 --- a/src/infiniop/ops/addr/cpu/addr_cpu.cc +++ /dev/null @@ -1,111 +0,0 @@ -#include "addr_cpu.h" -#include "../../../devices/cpu/common_cpu.h" - -namespace op::addr::cpu { -Descriptor::~Descriptor() = default; - -// Template function to handle different data types -template -infiniStatus_t addr_impl(Tdata *out, - const Tdata *input, - const Tdata *vec1, - const Tdata *vec2, - const AddrInfo &info, - void *workspace, - size_t workspace_size) { - size_t n = info.vec1_size; - size_t m = info.vec2_size; - float beta = info.beta; - float alpha = info.alpha; - - size_t total = n * m; - -#pragma omp parallel for - for (ptrdiff_t idx = 0; idx < (ptrdiff_t)total; ++idx) { - - // 🔹 Decode (i, j) - size_t i = idx / m; - size_t j = idx % m; - - size_t v1_idx = i * info.vec1_stride; - size_t v2_idx = j * info.vec2_stride; - size_t in_idx = i * info.input_stride0 + j * info.input_stride1; - size_t out_idx = i * info.output_stride0 + j * info.output_stride1; - - if constexpr (std::is_same::value || std::is_same::value) { - - float a = utils::cast(vec1[v1_idx]); - float b = utils::cast(vec2[v2_idx]); - float c = utils::cast(input[in_idx]); - - out[out_idx] = utils::cast(alpha * a * b + beta * c); - - } else { - - float a = (float)vec1[v1_idx]; - float b = (float)vec2[v2_idx]; - float c = (float)input[in_idx]; - - out[out_idx] = utils::cast(alpha * a * b + beta * c); - } - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t vec1_desc, - infiniopTensorDescriptor_t vec2_desc, - float beta, - float alpha) { - - auto handle = reinterpret_cast(handle_); - auto info = AddrInfo::create(input_desc, out_desc, vec1_desc, vec2_desc, beta, alpha); - CHECK_RESULT(info); - - *desc_ptr = new Descriptor(info.take(), 0, nullptr, - INFINI_DEVICE_CPU, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *input, - const void *vec1, - const void *vec2, - void *stream) const { - - switch (_info.dtype) { - case INFINI_DTYPE_F32: - return addr_impl(reinterpret_cast(out), - reinterpret_cast(input), - reinterpret_cast(vec1), - reinterpret_cast(vec2), - _info, workspace, workspace_size); - break; - case INFINI_DTYPE_F16: - return addr_impl(reinterpret_cast(out), - reinterpret_cast(input), - reinterpret_cast(vec1), - reinterpret_cast(vec2), - _info, workspace, workspace_size); - case INFINI_DTYPE_BF16: - return addr_impl(reinterpret_cast(out), - reinterpret_cast(input), - reinterpret_cast(vec1), - reinterpret_cast(vec2), - _info, workspace, workspace_size); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::addr::cpu diff --git a/src/infiniop/ops/addr/cpu/addr_cpu.h b/src/infiniop/ops/addr/cpu/addr_cpu.h deleted file mode 100644 index e2da1026a..000000000 --- a/src/infiniop/ops/addr/cpu/addr_cpu.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef INFINIOP_ADDR_CPU_H_ -#define INFINIOP_ADDR_CPU_H_ -#include "../addr.h" - -DESCRIPTOR(cpu) - -#endif diff --git a/src/infiniop/ops/addr/cuda/kernel.cu b/src/infiniop/ops/addr/cuda/kernel.cu deleted file mode 100644 index e920c372f..000000000 --- a/src/infiniop/ops/addr/cuda/kernel.cu +++ /dev/null @@ -1,38 +0,0 @@ -#ifndef __ADDR_CUDA_H__ -#define __ADDR_CUDA_H__ -#include -#include - -template -__device__ void addr_kernel(T *out, const T *input, const T *vec1, - const T *vec2, size_t n, size_t m, float beta, - float alpha, ptrdiff_t stride1, ptrdiff_t stride2, - ptrdiff_t out_stride_0, ptrdiff_t out_stride_1, - ptrdiff_t in_stride_0, ptrdiff_t in_stride_1) { - size_t i = blockIdx.x * blockDim.x + threadIdx.x; - size_t j = blockIdx.y * blockDim.y + threadIdx.y; - if (i >= n || j >= m) { - return; - } - size_t out_idx = i * out_stride_0 + j * out_stride_1; - size_t in_idx = i * in_stride_0 + j * in_stride_1; - size_t vec1_idx = i * stride1; - size_t vec2_idx = j * stride2; - T in_val = input[in_idx]; - T vec1_val = vec1[vec1_idx]; - T vec2_val = vec2[vec2_idx]; - T out_val; - if constexpr (std::is_same_v) { - out_val = __hadd(__hmul(__hmul(vec1_val, vec2_val), __float2half(alpha)), - __hmul(__float2half(beta), in_val)); - } else if constexpr (std::is_same_v) { - out_val = __hadd(__hmul(__hmul(vec1_val, vec2_val), __float2bfloat16(alpha)), - __hmul(__float2bfloat16(beta), in_val)); - } else { - out_val = beta * in_val + alpha * vec1_val * vec2_val; - } - - out[out_idx] = out_val; - __syncthreads(); -} -#endif diff --git a/src/infiniop/ops/addr/metax/addr_metax.h b/src/infiniop/ops/addr/metax/addr_metax.h deleted file mode 100644 index 3c7dad51b..000000000 --- a/src/infiniop/ops/addr/metax/addr_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef INFINIOP_ADDR_METAX_H_ -#define INFINIOP_ADDR_METAX_H_ - -#include "../addr.h" - -DESCRIPTOR(metax) - -#endif diff --git a/src/infiniop/ops/addr/metax/addr_nvidia.maca b/src/infiniop/ops/addr/metax/addr_nvidia.maca deleted file mode 100644 index 1187c81e2..000000000 --- a/src/infiniop/ops/addr/metax/addr_nvidia.maca +++ /dev/null @@ -1,107 +0,0 @@ - -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../cuda/kernel.cu" -#include "addr_metax.h" -#include "infinicore.h" -#include - -template -INFINIOP_METAX_KERNEL Addr(Tdata *output, const Tdata *input, const Tdata *vec1, - const Tdata *vec2, size_t n, size_t m, float beta, - float alpha, ptrdiff_t stride1, ptrdiff_t stride2, - ptrdiff_t out_stride_0, ptrdiff_t out_stride_1, - ptrdiff_t in_stride_0, ptrdiff_t in_stride_1) { - addr_kernel(output, input, vec1, vec2, n, m, beta, alpha, stride1, - stride2, out_stride_0, out_stride_1, in_stride_0, - in_stride_1); -} -namespace op::addr::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create(infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t vec1_desc, - infiniopTensorDescriptor_t vec2_desc, - float beta, float alpha) { - - auto info = AddrInfo::create(input_desc, out_desc, vec1_desc, vec2_desc, beta, alpha); - CHECK_RESULT(info); - size_t workspace_size = 0; - - *desc_ptr = new Descriptor( - info.take(), workspace_size, - new Opaque{reinterpret_cast(handle)->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t -launchKernel(infiniDtype_t dtype, void *output, const void *input, - const void *vec1, const void *vec2, size_t n, size_t m, float beta, - float alpha, ptrdiff_t stride1, ptrdiff_t stride2, - ptrdiff_t out_stride_0, ptrdiff_t out_stride_1, - ptrdiff_t in_stride_0, ptrdiff_t in_stride_1, hcStream_t stream) { - unsigned int dn = std::min((unsigned int)n, BLOCK_SIZE), - dm = std::min((unsigned int)m, BLOCK_SIZE); - dim3 grid = {(unsigned int)((n + dn - 1) / dn), - (unsigned int)((m + dm - 1) / dm)}; - dim3 block = {dn, dm}; - switch (dtype) { - case INFINI_DTYPE_F32: - Addr<<>>( - (float *)output, (const float *)input, (const float *)vec1, - (const float *)vec2, n, m, beta, alpha, stride1, stride2, out_stride_0, - out_stride_1, in_stride_0, in_stride_1); - break; - case INFINI_DTYPE_F16: - Addr<__half><<>>( - (__half *)output, (__half *)input, (__half *)vec1, (__half *)vec2, n, m, - beta, alpha, stride1, stride2, out_stride_0, out_stride_1, in_stride_0, - in_stride_1); - break; - case INFINI_DTYPE_BF16: - Addr<__hpcc_bfloat16><<>>( - (__hpcc_bfloat16 *)output, (__hpcc_bfloat16 *)input, - (__hpcc_bfloat16 *)vec1, (__hpcc_bfloat16 *)vec2, n, m, beta, alpha, - stride1, stride2, out_stride_0, out_stride_1, in_stride_0, in_stride_1); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *out, const void *input, - const void *vec1, const void *vec2, - void *stream_) const { - hcStream_t stream = (hcStream_t)stream_; - - if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - _info.dtype, out, input, vec1, vec2, _info.vec1_size, _info.vec2_size, - _info.beta, _info.alpha, _info.vec1_stride, _info.vec2_stride, - _info.output_stride0, _info.output_stride1, _info.input_stride0, - _info.input_stride1, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - _info.dtype, out, input, vec1, vec2, _info.vec1_size, _info.vec2_size, - _info.beta, _info.alpha, _info.vec1_stride, _info.vec2_stride, - _info.output_stride0, _info.output_stride1, _info.input_stride0, - _info.input_stride1, stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::addr::metax diff --git a/src/infiniop/ops/addr/moore/addr_kernel.h b/src/infiniop/ops/addr/moore/addr_kernel.h deleted file mode 100644 index dcc2f7d49..000000000 --- a/src/infiniop/ops/addr/moore/addr_kernel.h +++ /dev/null @@ -1,36 +0,0 @@ -#ifndef __ADDR_MOORE_H__ -#define __ADDR_MOORE_H__ - -template -__device__ void addr_kernel(T *out, const T *input, const T *vec1, - const T *vec2, size_t n, size_t m, float beta, - float alpha, ptrdiff_t stride1, ptrdiff_t stride2, - ptrdiff_t out_stride_0, ptrdiff_t out_stride_1, - ptrdiff_t in_stride_0, ptrdiff_t in_stride_1) { - size_t i = blockIdx.x * blockDim.x + threadIdx.x; - size_t j = blockIdx.y * blockDim.y + threadIdx.y; - if (i >= n || j >= m) { - return; - } - size_t out_idx = i * out_stride_0 + j * out_stride_1; - size_t in_idx = i * in_stride_0 + j * in_stride_1; - size_t vec1_idx = i * stride1; - size_t vec2_idx = j * stride2; - T in_val = input[in_idx]; - T vec1_val = vec1[vec1_idx]; - T vec2_val = vec2[vec2_idx]; - T out_val; - if constexpr (std::is_same_v) { - out_val = __hadd(__hmul(__hmul(vec1_val, vec2_val), __float2half(alpha)), - __hmul(__float2half(beta), in_val)); - } else if constexpr (std::is_same_v) { - float a = __bfloat162float(vec1_val), b = __bfloat162float(vec2_val), in = __bfloat162float(in_val); - out_val = __float2bfloat16_rn(a * b * alpha + in * beta); - } else { - out_val = beta * in_val + alpha * vec1_val * vec2_val; - } - - out[out_idx] = out_val; - __syncthreads(); -} -#endif diff --git a/src/infiniop/ops/addr/moore/addr_moore.h b/src/infiniop/ops/addr/moore/addr_moore.h deleted file mode 100644 index e090e6182..000000000 --- a/src/infiniop/ops/addr/moore/addr_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef INFINIOP_ADDR_MOORE_H_ -#define INFINIOP_ADDR_MOORE_H_ - -#include "../addr.h" - -DESCRIPTOR(moore) - -#endif diff --git a/src/infiniop/ops/addr/moore/addr_moore.mu b/src/infiniop/ops/addr/moore/addr_moore.mu deleted file mode 100644 index b00037033..000000000 --- a/src/infiniop/ops/addr/moore/addr_moore.mu +++ /dev/null @@ -1,107 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "addr_kernel.h" -#include "addr_moore.h" -#include "infinicore.h" -#include - -template -INFINIOP_MOORE_KERNEL Addr(Tdata *output, const Tdata *input, const Tdata *vec1, - const Tdata *vec2, size_t n, size_t m, float beta, - float alpha, ptrdiff_t stride1, ptrdiff_t stride2, - ptrdiff_t out_stride_0, ptrdiff_t out_stride_1, - ptrdiff_t in_stride_0, ptrdiff_t in_stride_1) { - addr_kernel(output, input, vec1, vec2, n, m, beta, alpha, stride1, - stride2, out_stride_0, out_stride_1, in_stride_0, - in_stride_1); -} -namespace op::addr::moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create(infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t vec1_desc, - infiniopTensorDescriptor_t vec2_desc, - float beta, float alpha) { - - auto info = AddrInfo::create(input_desc, out_desc, vec1_desc, vec2_desc, beta, alpha); - CHECK_RESULT(info); - size_t workspace_size = 0; - - *desc_ptr = new Descriptor( - info.take(), workspace_size, - new Opaque{ - reinterpret_cast(handle)->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(infiniDtype_t dtype, void *output, - const void *input, const void *vec1, - const void *vec2, size_t n, size_t m, float beta, - float alpha, ptrdiff_t stride1, ptrdiff_t stride2, - ptrdiff_t out_stride_0, ptrdiff_t out_stride_1, - ptrdiff_t in_stride_0, ptrdiff_t in_stride_1, - musaStream_t stream) { - unsigned int dn = std::min((unsigned int)n, BLOCK_SIZE), - dm = std::min((unsigned int)m, BLOCK_SIZE); - dim3 grid = {(unsigned int)((n + dn - 1) / dn), - (unsigned int)((m + dm - 1) / dm)}; - dim3 block = {dn, dm}; - switch (dtype) { - case INFINI_DTYPE_F32: - Addr<<>>( - (float *)output, (const float *)input, (const float *)vec1, - (const float *)vec2, n, m, beta, alpha, stride1, stride2, out_stride_0, - out_stride_1, in_stride_0, in_stride_1); - break; - case INFINI_DTYPE_F16: - Addr<__half><<>>( - (__half *)output, (__half *)input, (__half *)vec1, (__half *)vec2, n, m, - beta, alpha, stride1, stride2, out_stride_0, out_stride_1, in_stride_0, - in_stride_1); - break; - case INFINI_DTYPE_BF16: - Addr<<>>( - (cuda_bfloat16 *)output, (cuda_bfloat16 *)input, (cuda_bfloat16 *)vec1, - (cuda_bfloat16 *)vec2, n, m, beta, alpha, stride1, stride2, - out_stride_0, out_stride_1, in_stride_0, in_stride_1); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *out, const void *input, - const void *vec1, const void *vec2, - void *stream_) const { - musaStream_t stream = (musaStream_t)stream_; - - if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - _info.dtype, out, input, vec1, vec2, _info.vec1_size, _info.vec2_size, - _info.beta, _info.alpha, _info.vec1_stride, _info.vec2_stride, - _info.output_stride0, _info.output_stride1, _info.input_stride0, - _info.input_stride1, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - _info.dtype, out, input, vec1, vec2, _info.vec1_size, _info.vec2_size, - _info.beta, _info.alpha, _info.vec1_stride, _info.vec2_stride, - _info.output_stride0, _info.output_stride1, _info.input_stride0, - _info.input_stride1, stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::addr::moore diff --git a/src/infiniop/ops/addr/nvidia/addr_nvidia.cu b/src/infiniop/ops/addr/nvidia/addr_nvidia.cu deleted file mode 100644 index 00deeab48..000000000 --- a/src/infiniop/ops/addr/nvidia/addr_nvidia.cu +++ /dev/null @@ -1,120 +0,0 @@ - -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../cuda/kernel.cu" -#include "addr_nvidia.cuh" -#include "infinicore.h" -#include - -template -INFINIOP_CUDA_KERNEL Addr(Tdata *output, const Tdata *input, const Tdata *vec1, - const Tdata *vec2, size_t n, size_t m, float beta, - float alpha, ptrdiff_t stride1, ptrdiff_t stride2, - ptrdiff_t out_stride_0, ptrdiff_t out_stride_1, - ptrdiff_t in_stride_0, ptrdiff_t in_stride_1) { - addr_kernel(output, input, vec1, vec2, n, m, beta, alpha, stride1, - stride2, out_stride_0, out_stride_1, in_stride_0, - in_stride_1); -} -namespace op::addr::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create(infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t vec1_desc, - infiniopTensorDescriptor_t vec2_desc, - float beta, float alpha) { - - auto info = AddrInfo::create(input_desc, out_desc, vec1_desc, vec2_desc, beta, alpha); - CHECK_RESULT(info); - size_t workspace_size = 0; - - *desc_ptr = new Descriptor( - info.take(), workspace_size, - new Opaque{ - reinterpret_cast(handle)->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(infiniDtype_t dtype, void *output, - const void *input, const void *vec1, - const void *vec2, size_t n, size_t m, float beta, - float alpha, ptrdiff_t stride1, ptrdiff_t stride2, - ptrdiff_t out_stride_0, ptrdiff_t out_stride_1, - ptrdiff_t in_stride_0, ptrdiff_t in_stride_1, - cudaStream_t stream) { - unsigned int dn = std::min((unsigned int)n, BLOCK_SIZE), - dm = std::min((unsigned int)m, BLOCK_SIZE); - dim3 grid = {(unsigned int)((n + dn - 1) / dn), - (unsigned int)((m + dm - 1) / dm)}; - dim3 block = {dn, dm}; - switch (dtype) { - case INFINI_DTYPE_F32: - Addr<<>>( - (float *)output, (const float *)input, (const float *)vec1, - (const float *)vec2, n, m, beta, alpha, stride1, stride2, out_stride_0, - out_stride_1, in_stride_0, in_stride_1); - break; - case INFINI_DTYPE_F16: - Addr<__half><<>>( - (__half *)output, (__half *)input, (__half *)vec1, (__half *)vec2, n, m, - beta, alpha, stride1, stride2, out_stride_0, out_stride_1, in_stride_0, - in_stride_1); - break; - case INFINI_DTYPE_BF16: - Addr<__nv_bfloat16><<>>( - (__nv_bfloat16 *)output, (__nv_bfloat16 *)input, (__nv_bfloat16 *)vec1, - (__nv_bfloat16 *)vec2, n, m, beta, alpha, stride1, stride2, - out_stride_0, out_stride_1, in_stride_0, in_stride_1); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *out, const void *input, - const void *vec1, const void *vec2, - void *stream_) const { - cudaStream_t stream = (cudaStream_t)stream_; - - if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - _info.dtype, out, input, vec1, vec2, _info.vec1_size, _info.vec2_size, - _info.beta, _info.alpha, _info.vec1_stride, _info.vec2_stride, - _info.output_stride0, _info.output_stride1, _info.input_stride0, - _info.input_stride1, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - _info.dtype, out, input, vec1, vec2, _info.vec1_size, _info.vec2_size, - _info.beta, _info.alpha, _info.vec1_stride, _info.vec2_stride, - _info.output_stride0, _info.output_stride1, _info.input_stride0, - _info.input_stride1, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_2048) { - CHECK_STATUS(launchKernel( - _info.dtype, out, input, vec1, vec2, _info.vec1_size, _info.vec2_size, - _info.beta, _info.alpha, _info.vec1_stride, _info.vec2_stride, - _info.output_stride0, _info.output_stride1, _info.input_stride0, - _info.input_stride1, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) { - CHECK_STATUS(launchKernel( - _info.dtype, out, input, vec1, vec2, _info.vec1_size, _info.vec2_size, - _info.beta, _info.alpha, _info.vec1_stride, _info.vec2_stride, - _info.output_stride0, _info.output_stride1, _info.input_stride0, - _info.input_stride1, stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::addr::nvidia diff --git a/src/infiniop/ops/addr/nvidia/addr_nvidia.cuh b/src/infiniop/ops/addr/nvidia/addr_nvidia.cuh deleted file mode 100644 index 151da410b..000000000 --- a/src/infiniop/ops/addr/nvidia/addr_nvidia.cuh +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef INFINIOP_ADDR_NVIDIA_H_ -#define INFINIOP_ADDR_NVIDIA_H_ -#include "../addr.h" - -DESCRIPTOR(nvidia) - -#endif diff --git a/src/infiniop/ops/addr/operator.cc b/src/infiniop/ops/addr/operator.cc deleted file mode 100644 index 1b5101352..000000000 --- a/src/infiniop/ops/addr/operator.cc +++ /dev/null @@ -1,194 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infinicore.h" -#include "infiniop/ops/addr.h" - -#ifdef ENABLE_CPU_API -#include "cpu/addr_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/addr_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/addr_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/addr_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAddrDescriptor( - infiniopHandle_t handle, - infiniopAddrDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out, - infiniopTensorDescriptor_t input, - infiniopTensorDescriptor_t vec1, - infiniopTensorDescriptor_t vec2, - float beta, - float alpha) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::addr::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - out, \ - input, \ - vec1, \ - vec2, \ - beta, \ - alpha) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopGetAddrWorkspaceSize( - infiniopAddrDescriptor_t desc, - size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopAddr( - infiniopAddrDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *input, - const void *vec1, - const void *vec2, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, out, input, vec1, vec2, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopDestroyAddrDescriptor(infiniopAddrDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} diff --git a/src/infiniop/ops/affine_grid/affine_grid.h b/src/infiniop/ops/affine_grid/affine_grid.h deleted file mode 100644 index 723f2e743..000000000 --- a/src/infiniop/ops/affine_grid/affine_grid.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef AFFINE_GRID_H -#define AFFINE_GRID_H - -#include "../../operator.h" -#include "info.h" - -// 宏定义：用于生成不同命名空间下的 Descriptor 类 -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::affine_grid::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AffineGridInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - AffineGridInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t in_desc, \ - bool align_corners); /* 增加 align_corners 参数 */ \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *output, \ - const void *input, \ - void *stream) const; \ - }; \ - } - -#endif // AFFINE_GRID_H diff --git a/src/infiniop/ops/affine_grid/cpu/affine_grid_cpu.cc b/src/infiniop/ops/affine_grid/cpu/affine_grid_cpu.cc deleted file mode 100644 index e0296a4b6..000000000 --- a/src/infiniop/ops/affine_grid/cpu/affine_grid_cpu.cc +++ /dev/null @@ -1,143 +0,0 @@ -#include "affine_grid_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include -#include -#include - -namespace op::affine_grid::cpu { - -template -inline float to_float(T val) { - if constexpr (std::is_same::value || std::is_same::value) { - return utils::cast(val); - } else { - return static_cast(val); - } -} - -template -inline T from_float(float val) { - if constexpr (std::is_same::value || std::is_same::value) { - return utils::cast(val); - } else { - return static_cast(val); - } -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - bool align_corners) { // 接收 align_corners - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - // 1. 检查数据类型 - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); - - // 2. 创建 Info 对象 (传递 align_corners) - auto result = AffineGridInfo::create(out_desc, in_desc, align_corners); - CHECK_RESULT(result); - - // 3. 创建 Descriptor - *desc_ptr = new Descriptor( - nullptr, // Opaque* - result.take(), // Info - 0, // Workspace Size (AffineGrid CPU 不需要额外 workspace) - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -void calculate_cpu_impl( - const AffineGridInfo &info, - void *output, - const void *input) { - - size_t batch = info.batch(); - size_t H = info.height(); - size_t W = info.width(); - bool align_corners = info.align_corners(); - - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - - // 并行化处理 Batch -#pragma omp parallel for - for (ptrdiff_t n = 0; n < (ptrdiff_t)batch; ++n) { - - const Tdata *theta_n = in_ptr + n * 6; - - // 提取仿射矩阵参数并转为 float - float r00 = to_float(theta_n[0]); - float r01 = to_float(theta_n[1]); - float tx = to_float(theta_n[2]); - float r10 = to_float(theta_n[3]); - float r11 = to_float(theta_n[4]); - float ty = to_float(theta_n[5]); - - // 遍历空间维度 - for (size_t h = 0; h < H; ++h) { - for (size_t w = 0; w < W; ++w) { - // 1. 计算归一化坐标 (-1 到 1) - float x_norm, y_norm; - - if (align_corners) { - x_norm = (W > 1) ? (2.0f * w) / (W - 1.0f) - 1.0f : 0.0f; - y_norm = (H > 1) ? (2.0f * h) / (H - 1.0f) - 1.0f : 0.0f; - } else { - x_norm = (2.0f * w + 1.0f) / W - 1.0f; - y_norm = (2.0f * h + 1.0f) / H - 1.0f; - } - - // 2. 应用仿射变换 - float grid_x = r00 * x_norm + r01 * y_norm + tx; - float grid_y = r10 * x_norm + r11 * y_norm + ty; - size_t offset = (n * H * W + h * W + w) * 2; - - out_ptr[offset + 0] = from_float(grid_x); - out_ptr[offset + 1] = from_float(grid_y); - } - } - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) const { - - // 从 Info 中获取 dtype - auto dtype = _info.dtype(); - - switch (dtype) { - case INFINI_DTYPE_F16: - cpu::calculate_cpu_impl(_info, output, input); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_BF16: - cpu::calculate_cpu_impl(_info, output, input); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_F32: - cpu::calculate_cpu_impl(_info, output, input); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_F64: - cpu::calculate_cpu_impl(_info, output, input); - return INFINI_STATUS_SUCCESS; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::affine_grid::cpu diff --git a/src/infiniop/ops/affine_grid/cpu/affine_grid_cpu.h b/src/infiniop/ops/affine_grid/cpu/affine_grid_cpu.h deleted file mode 100644 index 1e2ee9580..000000000 --- a/src/infiniop/ops/affine_grid/cpu/affine_grid_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __AFFINE_GRID_CPU_H__ -#define __AFFINE_GRID_CPU_H__ - -#include "../affine_grid.h" - -DESCRIPTOR(cpu) - -#endif //_GRID_CPU_H_ diff --git a/src/infiniop/ops/affine_grid/cuda/kernel.cuh b/src/infiniop/ops/affine_grid/cuda/kernel.cuh deleted file mode 100644 index 8e0597266..000000000 --- a/src/infiniop/ops/affine_grid/cuda/kernel.cuh +++ /dev/null @@ -1,103 +0,0 @@ -#ifndef __AFFINE_GRID_CUDA_H__ -#define __AFFINE_GRID_CUDA_H__ - -#include -#include - -namespace op::affine_grid::cuda { - -template -__device__ __forceinline__ float to_float_acc(const T &x) { - if constexpr (std::is_same_v) { - return __half2float(x); - } else if constexpr (std::is_same_v) { - return __bfloat162float(x); - } else { - return static_cast(x); - } -} - -template -__global__ void affine_grid_kernel( - T *__restrict__ output, // [优化1] 使用 __restrict__ - const T *__restrict__ theta, // [优化1] 使用 __restrict__ - size_t N, - size_t H, - size_t W, - bool align_corners) { - // 扁平化索引 - size_t total_elements = N * H * W; - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx >= total_elements) { - return; - } - - float w_scale, h_scale, w_bias, h_bias; - - if (align_corners) { - // align_corners = True: formula is (2*i)/(size-1) - 1 - // => i * (2/(size-1)) - 1 - w_scale = (W > 1) ? 2.0f / (W - 1.0f) : 0.0f; - h_scale = (H > 1) ? 2.0f / (H - 1.0f) : 0.0f; - w_bias = -1.0f; - h_bias = -1.0f; - } else { - // align_corners = False: formula is (2*i + 1)/size - 1 - // => i * (2/size) + (1/size - 1) - w_scale = 2.0f / W; - h_scale = 2.0f / H; - w_bias = 1.0f / W - 1.0f; - h_bias = 1.0f / H - 1.0f; - } - - size_t w = idx % W; - size_t temp = idx / W; - size_t h = temp % H; - size_t n = temp / H; // 此时 temp = n * H + h - - // 2. 计算归一化坐标 (使用乘法代替除法) - float x_norm = (float)w * w_scale + w_bias; - float y_norm = (float)h * h_scale + h_bias; - - // 如果 align_corners=True 且 size=1，特判修正 - if (align_corners) { - if (W <= 1) { - x_norm = 0.0f; - } - if (H <= 1) { - y_norm = 0.0f; - } - } - - const T *theta_ptr = theta + n * 6; - - float r00 = to_float_acc(theta_ptr[0]); - float r01 = to_float_acc(theta_ptr[1]); - float tx = to_float_acc(theta_ptr[2]); - float r10 = to_float_acc(theta_ptr[3]); - float r11 = to_float_acc(theta_ptr[4]); - float ty = to_float_acc(theta_ptr[5]); - - float grid_x = r00 * x_norm + r01 * y_norm + tx; - float grid_y = r10 * x_norm + r11 * y_norm + ty; - - // 5. 向量化写入 (Vectorized Store) - if constexpr (std::is_same_v) { - float2 *out_vec = reinterpret_cast(output); - out_vec[idx] = make_float2(grid_x, grid_y); - } else if constexpr (std::is_same_v) { - half2 *out_vec = reinterpret_cast(output); - out_vec[idx] = __floats2half2_rn(grid_x, grid_y); - } else if constexpr (std::is_same_v) { - cuda_bfloat162 *out_vec = reinterpret_cast(output); - out_vec[idx] = __floats2bfloat162_rn(grid_x, grid_y); - } else { - output[idx * 2 + 0] = static_cast(grid_x); - output[idx * 2 + 1] = static_cast(grid_y); - } -} - -} // namespace op::affine_grid::cuda - -#endif // __AFFINE_GRID_CUDA_H__ diff --git a/src/infiniop/ops/affine_grid/info.h b/src/infiniop/ops/affine_grid/info.h deleted file mode 100644 index bf9c9b39e..000000000 --- a/src/infiniop/ops/affine_grid/info.h +++ /dev/null @@ -1,78 +0,0 @@ -#ifndef __AFFINE_GRID_INFO_H__ -#define __AFFINE_GRID_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::affine_grid { - -class AffineGridInfo { - AffineGridInfo() = default; - -public: - size_t _batch; - size_t _height; - size_t _width; - bool _align_corners; - int _dtype; - - size_t batch() const { return _batch; } - size_t height() const { return _height; } - size_t width() const { return _width; } - bool align_corners() const { return _align_corners; } - int dtype() const { return _dtype; } - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - bool align_corners) { - - // 1. 检查数据类型一致性 - if (out_desc->dtype() != in_desc->dtype()) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // 2. 检查输入 Theta 的形状 - // 标准 2D Affine Grid 输入必须是 (N, 2, 3) - if (in_desc->ndim() != 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (in_desc->shape()[1] != 2 || in_desc->shape()[2] != 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // 3. 检查输出 Grid 的形状 - // 标准 2D Affine Grid 输出必须是 (N, H, W, 2) - if (out_desc->ndim() != 4) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - // 最后一维必须是 2 (代表 x, y 坐标) - if (out_desc->shape()[3] != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // 4. 检查 Batch Size 是否匹配 - if (in_desc->shape()[0] != out_desc->shape()[0]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // 5. 提取维度信息 - size_t batch = out_desc->shape()[0]; - size_t height = out_desc->shape()[1]; - size_t width = out_desc->shape()[2]; - int dtype = in_desc->dtype(); - - // 6. 返回 Info 对象 - return utils::Result(AffineGridInfo{ - batch, - height, - width, - align_corners, - dtype}); - } -}; - -} // namespace op::affine_grid - -#endif // __AFFINE_GRID_INFO_H__ diff --git a/src/infiniop/ops/affine_grid/metax/affine_grid_metax.h b/src/infiniop/ops/affine_grid/metax/affine_grid_metax.h deleted file mode 100644 index c3aa24b6b..000000000 --- a/src/infiniop/ops/affine_grid/metax/affine_grid_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __AFFINE_GRID_METAX_H__ -#define __AFFINE_GRID_METAX_H__ - -#include "../affine_grid.h" - -DESCRIPTOR(metax) - -#endif // __AFFINE_GRID_METAX_H__ diff --git a/src/infiniop/ops/affine_grid/metax/affine_grid_metax.maca b/src/infiniop/ops/affine_grid/metax/affine_grid_metax.maca deleted file mode 100644 index 8f5e64d5b..000000000 --- a/src/infiniop/ops/affine_grid/metax/affine_grid_metax.maca +++ /dev/null @@ -1,192 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "affine_grid_metax.h" -#include -#include -#include - -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" - -namespace op::affine_grid::metax { - -// ================================================================== -// Kernel: Index-Based Double Precision (Maximum Accuracy) -// ================================================================== -__global__ void affine_grid_kernel_f32_double_index( - float *__restrict__ output, - const float *__restrict__ theta, - int N, int H, int W, int align_corners_int, - int in_s_n, int in_s_h, int in_s_w, - int out_s_n, int out_s_h, int out_s_w, int out_s_c) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total_elements = N * H * W; - - if (idx >= total_elements) { - return; - } - - // 1. 索引分解 - int w_idx = idx % W; - int temp = idx / W; - int h_idx = temp % H; - int n_idx = temp / H; - - // 2. 坐标生成 (Double Precision + Index Formula) - double x, y; - - if (align_corners_int == 1) { - // align_corners = True - // Formula: (i * 2) / (N - 1) - 1 - if (W > 1) { - x = ((double)w_idx * 2.0) / (double)(W - 1) - 1.0; - } else { - x = 0.0; - } - - if (H > 1) { - y = ((double)h_idx * 2.0) / (double)(H - 1) - 1.0; - } else { - y = 0.0; - } - } else { - // align_corners = False - // Formula: (2 * i + 1) / N - 1 - x = ((double)(2 * w_idx + 1)) / (double)W - 1.0; - y = ((double)(2 * h_idx + 1)) / (double)H - 1.0; - } - - // 3. 读取 Theta (Convert to Double) - int theta_base = n_idx * in_s_n; - double r00 = (double)theta[theta_base + 0 * in_s_h + 0 * in_s_w]; - double r01 = (double)theta[theta_base + 0 * in_s_h + 1 * in_s_w]; - double tx = (double)theta[theta_base + 0 * in_s_h + 2 * in_s_w]; - - double r10 = (double)theta[theta_base + 1 * in_s_h + 0 * in_s_w]; - double r11 = (double)theta[theta_base + 1 * in_s_h + 1 * in_s_w]; - double ty = (double)theta[theta_base + 1 * in_s_h + 2 * in_s_w]; - - double grid_x = r00 * x + r01 * y + tx; - double grid_y = r10 * x + r11 * y + ty; - - // 5. 结果写入 (Cast back to float) - int out_base = n_idx * out_s_n + h_idx * out_s_h + w_idx * out_s_w; - output[out_base + 0 * out_s_c] = (float)grid_x; - output[out_base + 1 * out_s_c] = (float)grid_y; -} - -// ================================================================== -// Launch Helper & Descriptor -// ================================================================== -template -void launch_kernel( - void *output, const void *input, size_t batch, size_t height, size_t width, bool align_corners, - void *stream, const int64_t *in_strides = nullptr, const int64_t *out_strides = nullptr) { - auto hc_stream = reinterpret_cast(stream); - size_t total_elements = batch * height * width; - size_t block_size = 256; - size_t grid_size = (total_elements + block_size - 1) / block_size; - - if constexpr (std::is_same_v) { - - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - int align_corners_i32 = align_corners ? 1 : 0; - - int is_n = 6, is_h = 3, is_w = 1; - if (in_strides) { - is_n = (int)in_strides[0]; - is_h = (int)in_strides[1]; - is_w = (int)in_strides[2]; - } - - int os_n = height * width * 2, os_h = width * 2, os_w = 2, os_c = 1; - if (out_strides) { - os_n = (int)out_strides[0]; - os_h = (int)out_strides[1]; - os_w = (int)out_strides[2]; - os_c = (int)out_strides[3]; - } - - affine_grid_kernel_f32_double_index<<>>( - out_ptr, in_ptr, (int)batch, (int)height, (int)width, align_corners_i32, - is_n, is_h, is_w, os_n, os_h, os_w, os_c); - } else { - // [FP16 / BF16 路径] 使用通用模板 - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - op::affine_grid::cuda::affine_grid_kernel<<>>( - out_ptr, in_ptr, batch, height, width, align_corners); - } -} - -// ================================================================== -// Descriptor Implementation -// ================================================================== -struct Descriptor::Opaque { - std::shared_ptr internal; - std::vector in_strides; - std::vector out_strides; -}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, Descriptor **desc_ptr, infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t in_desc, bool align_corners) { - auto handle = reinterpret_cast(handle_); - auto info_result = AffineGridInfo::create(out_desc, in_desc, align_corners); - if (!info_result) { - return info_result.status(); - } - auto opaque = new Opaque{handle->internal()}; - - // 获取 Input Strides - auto idesc = reinterpret_cast(in_desc); - if (idesc && idesc->ndim() == 3) { - auto s = idesc->strides(); - for (auto val : s) { - opaque->in_strides.push_back((int64_t)val); - } - } - // 获取 Output Strides - auto odesc = reinterpret_cast(out_desc); - if (odesc && odesc->ndim() == 4) { - auto s = odesc->strides(); - for (auto val : s) { - opaque->out_strides.push_back((int64_t)val); - } - } - *desc_ptr = new Descriptor(opaque, info_result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, void *output, const void *input, void *stream) const { - auto dtype = _info.dtype(); - const int64_t *in_strides_ptr = _opaque->in_strides.empty() ? nullptr : _opaque->in_strides.data(); - const int64_t *out_strides_ptr = _opaque->out_strides.empty() ? nullptr : _opaque->out_strides.data(); - - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel<__half>(output, input, _info.batch(), _info.height(), _info.width(), _info.align_corners(), stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, _info.batch(), _info.height(), _info.width(), _info.align_corners(), stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(output, input, _info.batch(), _info.height(), _info.width(), _info.align_corners(), stream, in_strides_ptr, out_strides_ptr); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, _info.batch(), _info.height(), _info.width(), _info.align_corners(), stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} -} // namespace op::affine_grid::metax diff --git a/src/infiniop/ops/affine_grid/moore/affine_grid_moore.h b/src/infiniop/ops/affine_grid/moore/affine_grid_moore.h deleted file mode 100644 index d96763b24..000000000 --- a/src/infiniop/ops/affine_grid/moore/affine_grid_moore.h +++ /dev/null @@ -1,9 +0,0 @@ -#ifndef __AFFINE_GRID_MOORE_API_H__ -#define __AFFINE_GRID_MOORE_API_H__ - -#include "../affine_grid.h" - -// 使用 affine_grid.h 中定义的宏生成 op::affine_grid::moore::Descriptor 类定义 -DESCRIPTOR(moore) - -#endif // __AFFINE_GRID_MOORE_API_H__ diff --git a/src/infiniop/ops/affine_grid/moore/affine_grid_moore.mu b/src/infiniop/ops/affine_grid/moore/affine_grid_moore.mu deleted file mode 100644 index 71cee7cdc..000000000 --- a/src/infiniop/ops/affine_grid/moore/affine_grid_moore.mu +++ /dev/null @@ -1,143 +0,0 @@ -#include "affine_grid_moore.h" -#include "affine_grid_moore_kernel.h" -#include - -// 引用 Handle 路径 -#include "../../../devices/moore/moore_handle.h" - -namespace op::affine_grid::moore { - -template -__global__ void affine_grid_kernel( - const int N, const int H, const int W, - const bool align_corners, - const T *theta, - T *output) { - - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total_pixels = N * H * W; - - if (idx < total_pixels) { - int w = idx % W; - int h = (idx / W) % H; - int n = idx / (W * H); - - const T *current_theta = theta + n * 6; - T *out_ptr = output + idx * 2; - - AffineGridOp op; - op(w, h, W, H, current_theta, align_corners, &out_ptr[0], &out_ptr[1]); - } -} - -// ================================================================== -// 2. Launcher Implementation -// ================================================================== - -template -void affine_grid_moore_launch( - const AffineGridInfo &info, - T *output, - const T *input, - void *stream) { - - size_t num_pixels = info.batch() * info.height() * info.width(); - - int threads = 256; - int blocks = (num_pixels + threads - 1) / threads; - - affine_grid_kernel<<>>( - info.batch(), - info.height(), - info.width(), - info.align_corners(), - input, - output); -} - -// ================================================================== -// 3. Descriptor Implementation -// ================================================================== - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - bool align_corners) { - - auto handle = reinterpret_cast(handle_); - - auto info_result = AffineGridInfo::create(out_desc, in_desc, align_corners); - - if (!info_result) { - - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - *desc_ptr = new Descriptor( - nullptr, - *info_result, - 0, - handle->device, // 原: handle->device_type() - handle->device_id // 原: handle->device_id() - ); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_info.dtype()) { - case INFINI_DTYPE_F16: - affine_grid_moore_launch( - _info, - static_cast(output), - static_cast(input), - stream); - break; - - case INFINI_DTYPE_BF16: - - affine_grid_moore_launch<__mt_bfloat16>( - _info, - static_cast<__mt_bfloat16 *>(output), - static_cast(input), - stream); - break; - - case INFINI_DTYPE_F32: - affine_grid_moore_launch( - _info, - static_cast(output), - static_cast(input), - stream); - break; - - case INFINI_DTYPE_F64: - affine_grid_moore_launch( - _info, - static_cast(output), - static_cast(input), - stream); - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::affine_grid::moore diff --git a/src/infiniop/ops/affine_grid/moore/affine_grid_moore_kernel.h b/src/infiniop/ops/affine_grid/moore/affine_grid_moore_kernel.h deleted file mode 100644 index fd6c307b8..000000000 --- a/src/infiniop/ops/affine_grid/moore/affine_grid_moore_kernel.h +++ /dev/null @@ -1,80 +0,0 @@ -#ifndef __AFFINE_GRID_MOORE_KERNEL_H__ -#define __AFFINE_GRID_MOORE_KERNEL_H__ - -#include // 包含 __mt_bfloat16 定义 -#include -#include - -namespace op::affine_grid::moore { -typedef struct AffineGridOp { -public: - static constexpr size_t num_dimensions = 2; - - template - __device__ __forceinline__ void operator()( - const int w_idx, const int h_idx, - const int W, const int H, - const T *theta, - const bool align_corners, - T *out_x, T *out_y) const { - // 1. 归一化坐标计算 - float x_norm, y_norm; - if (align_corners) { - x_norm = (float)(w_idx * 2 - (W - 1)) / (float)(MAX(W - 1, 1)); - y_norm = (float)(h_idx * 2 - (H - 1)) / (float)(MAX(H - 1, 1)); - } else { - x_norm = (float)(w_idx * 2 + 1) / (float)W - 1.0f; - y_norm = (float)(h_idx * 2 + 1) / (float)H - 1.0f; - } - - // 2. 仿射变换逻辑 - if constexpr (std::is_same_v) { - float t00 = __half2float(theta[0]); - float t01 = __half2float(theta[1]); - float t02 = __half2float(theta[2]); - float t10 = __half2float(theta[3]); - float t11 = __half2float(theta[4]); - float t12 = __half2float(theta[5]); - - float res_x = t00 * x_norm + t01 * y_norm + t02; - float res_y = t10 * x_norm + t11 * y_norm + t12; - - *out_x = __float2half(res_x); - *out_y = __float2half(res_y); - - } else if constexpr (std::is_same_v) { // 【修改】使用 __mt_bfloat16 - // 显式转换 __mt_bfloat16 -> float - float t00 = __bfloat162float(theta[0]); - float t01 = __bfloat162float(theta[1]); - float t02 = __bfloat162float(theta[2]); - float t10 = __bfloat162float(theta[3]); - float t11 = __bfloat162float(theta[4]); - float t12 = __bfloat162float(theta[5]); - - float res_x = t00 * x_norm + t01 * y_norm + t02; - float res_y = t10 * x_norm + t11 * y_norm + t12; - - // 转换回 __mt_bfloat16 - *out_x = __float2bfloat16(res_x); - *out_y = __float2bfloat16(res_y); - - } else if constexpr (std::is_same_v) { - float res_x = __fadd_rn(__fmul_rn(theta[0], x_norm), __fadd_rn(__fmul_rn(theta[1], y_norm), theta[2])); - float res_y = __fadd_rn(__fmul_rn(theta[3], x_norm), __fadd_rn(__fmul_rn(theta[4], y_norm), theta[5])); - *out_x = res_x; - *out_y = res_y; - } else { - *out_x = theta[0] * static_cast(x_norm) + theta[1] * static_cast(y_norm) + theta[2]; - *out_y = theta[3] * static_cast(x_norm) + theta[4] * static_cast(y_norm) + theta[5]; - } - } - -private: - __device__ __forceinline__ int MAX(int a, int b) const { - return a > b ? a : b; - } - -} AffineGridOp; -} // namespace op::affine_grid::moore - -#endif // __AFFINE_GRID_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/affine_grid/nvidia/affine_grid_nvidia.cu b/src/infiniop/ops/affine_grid/nvidia/affine_grid_nvidia.cu deleted file mode 100644 index fd8955c28..000000000 --- a/src/infiniop/ops/affine_grid/nvidia/affine_grid_nvidia.cu +++ /dev/null @@ -1,124 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" - -#include "../cuda/kernel.cuh" -#include "affine_grid_nvidia.cuh" - -namespace op::affine_grid::nvidia { - -// ================================================================== -// Kernel Launch Helper -// ================================================================== -template -void launch_kernel( - void *output, - const void *input, // 这里对应 Theta - size_t batch, - size_t height, - size_t width, - bool align_corners, - void *stream) { - - // 指针强转 - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - - // 计算总线程数: N * H * W - // 每个线程负责生成一个 (x, y) 坐标对 - size_t total_elements = batch * height * width; - - size_t block_size = 256; - size_t grid_size = (total_elements + block_size - 1) / block_size; - - auto cuda_stream = reinterpret_cast(stream); - - cuda::affine_grid_kernel<<>>( - out_ptr, - in_ptr, - batch, - height, - width, - align_corners); -} - -// ================================================================== -// Descriptor Implementation -// ================================================================== - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -// 创建算子描述符 -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - bool align_corners) { - - // 1. 使用 Info 类解析并校验参数 - // create 方法内部会检查 Tensor 维度和类型一致性 - auto info_result = AffineGridInfo::create(out_desc, in_desc, align_corners); - if (!info_result) { - return info_result.status(); - } - auto info = info_result.take(); - - // 2. 创建 Descriptor - *desc_ptr = new Descriptor( - new Opaque(), // Opaque 指针 - info, // Info 对象 (包含 N, H, W, align_corners) - 0, // Workspace size (AffineGrid 不需要额外 workspace) - handle->device, // Device Type - handle->device_id // Device ID - ); - - return INFINI_STATUS_SUCCESS; -} - -// 执行计算 -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) const { - auto dtype = _info.dtype(); - auto batch = _info.batch(); - auto height = _info.height(); - auto width = _info.width(); - auto align_corners = _info.align_corners(); - - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel(output, input, batch, height, width, align_corners, stream); - break; - - case INFINI_DTYPE_BF16: - - launch_kernel(output, input, batch, height, width, align_corners, stream); - break; - - case INFINI_DTYPE_F32: - launch_kernel(output, input, batch, height, width, align_corners, stream); - break; - - case INFINI_DTYPE_F64: - - launch_kernel(output, input, batch, height, width, align_corners, stream); - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::affine_grid::nvidia diff --git a/src/infiniop/ops/affine_grid/nvidia/affine_grid_nvidia.cuh b/src/infiniop/ops/affine_grid/nvidia/affine_grid_nvidia.cuh deleted file mode 100644 index 57c65b10f..000000000 --- a/src/infiniop/ops/affine_grid/nvidia/affine_grid_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __AFFINE_GRID_CUH__ -#define __AFFINE_GRID_CUH__ - -#include "../affine_grid.h" - -DESCRIPTOR(nvidia) - -#endif // __GEMM_CUDA_CUH__ diff --git a/src/infiniop/ops/affine_grid/operator.cc b/src/infiniop/ops/affine_grid/operator.cc deleted file mode 100644 index d5212f798..000000000 --- a/src/infiniop/ops/affine_grid/operator.cc +++ /dev/null @@ -1,182 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/affine_grid.h" - -// --- 后端实现头文件 --- -#ifdef ENABLE_CPU_API -#include "cpu/affine_grid_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/affine_grid_nvidia.cuh" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/affine_grid_moore.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/affine_grid_metax.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAffineGridDescriptor( - infiniopHandle_t handle, - infiniopAffineGridDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - uint8_t align_corners) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::affine_grid::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output_desc, \ - input_desc, \ - align_corners) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAffineGridWorkspaceSize(infiniopAffineGridDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (reinterpret_cast(desc)->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopAffineGrid( - infiniopAffineGridDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, input, stream) - - switch (reinterpret_cast(desc)->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyAffineGridDescriptor(infiniopAffineGridDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (reinterpret_cast(desc)->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/all/all_desc.h b/src/infiniop/ops/all/all_desc.h deleted file mode 100644 index 9b7a1e0d6..000000000 --- a/src/infiniop/ops/all/all_desc.h +++ /dev/null @@ -1,53 +0,0 @@ -#ifndef INFINIOP_ALL_DESCRIPTOR_H_ -#define INFINIOP_ALL_DESCRIPTOR_H_ -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" - -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::all::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AllInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - AllInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t output_desc, \ - infiniopTensorDescriptor_t input_desc, \ - size_t *dim, \ - size_t dim_size, \ - bool keepdim); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *output, \ - const void *input, \ - size_t *dim, \ - size_t dim_size, \ - bool keepdim, \ - void *stream) const; \ - }; \ - } - -#endif diff --git a/src/infiniop/ops/all/cpu/all_cpu.cc b/src/infiniop/ops/all/cpu/all_cpu.cc deleted file mode 100644 index dbe03fc3b..000000000 --- a/src/infiniop/ops/all/cpu/all_cpu.cc +++ /dev/null @@ -1,77 +0,0 @@ -#include "all_cpu.h" -#include "../../../../utils.h" -#include "../../../devices/cpu/common_cpu.h" -#include -namespace op::all::cpu { - -Descriptor::~Descriptor() {} -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - size_t *dim, - size_t dim_size, - bool keepdim) { - auto result = AllInfo::create(output_desc, input_desc, dim, dim_size, keepdim); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor(nullptr, result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { -template -infiniStatus_t calculateAll( - const AllInfo &info, - bool *output, - const Tdata *input, - size_t *dim, - size_t dim_size, - bool keepdim) { - if (info.reduce_dim_size == info.ndim) { - bool result = true; - for (size_t index = 0; index < info.input_size; index++) { - size_t input_offset = op::common_cpu::indexToOffset(index, info.ndim, info.permuted_input_shape.data(), info.permuted_input_strides.data()); - result = result && input[input_offset]; - } - output[0] = result; - return INFINI_STATUS_SUCCESS; - } else { - for (size_t i = info.output_size; i-- > 0;) { - size_t output_offset = op::common_cpu::indexToOffset(i, info.output_shape.size(), info.output_shape.data(), info.output_strides.data()); - bool result = true; - for (size_t j = 0; j < info.reduce_num; j++) { - size_t input_flat = j + i * info.reduce_num; - size_t input_offset = op::common_cpu::indexToOffset(input_flat, info.ndim, info.permuted_input_shape.data(), info.permuted_input_strides.data()); - Tdata input_val = input[input_offset]; - bool bool_val = static_cast(input_val); - result = result && bool_val; - } - output[output_offset] = result; - } - return INFINI_STATUS_SUCCESS; - } -} -} // namespace -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - size_t *dim, - size_t dim_size, - bool keepdim, - void *stream) const { - switch (_info.dtype) { - case INFINI_DTYPE_BOOL: - return calculateAll(_info, reinterpret_cast(output), reinterpret_cast(input), dim, dim_size, keepdim); - case INFINI_DTYPE_U8: - return calculateAll(_info, reinterpret_cast(output), reinterpret_cast(input), dim, dim_size, keepdim); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::all::cpu diff --git a/src/infiniop/ops/all/cpu/all_cpu.h b/src/infiniop/ops/all/cpu/all_cpu.h deleted file mode 100644 index 71fd83689..000000000 --- a/src/infiniop/ops/all/cpu/all_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __INFINIOP_ALL_CPU_H__ -#define __INFINIOP_ALL_CPU_H__ - -#include "../all_desc.h" - -DESCRIPTOR(cpu); - -#endif // __INFINIOP_ALL_CPU_H__ diff --git a/src/infiniop/ops/all/cuda/kernel.cuh b/src/infiniop/ops/all/cuda/kernel.cuh deleted file mode 100644 index 8c7233174..000000000 --- a/src/infiniop/ops/all/cuda/kernel.cuh +++ /dev/null @@ -1,84 +0,0 @@ -#ifndef __ALL_CUDA_H__ -#define __ALL_CUDA_H__ - -template -__global__ void allReduceTempKernel( - bool *temp_output, - const Tdata *input, - size_t input_size, - size_t permuted_input_shape_size, - size_t *permuted_input_shape, - ptrdiff_t *permuted_input_strides) { - __shared__ bool s_data[BLOCK_SIZE]; - size_t tid = threadIdx.x; - size_t idx = tid + blockIdx.x * blockDim.x; - if (idx < input_size) { - size_t input_offset = indexToOffset(idx, permuted_input_shape_size, permuted_input_shape, permuted_input_strides); - s_data[tid] = static_cast(input[input_offset]); - } else { - s_data[tid] = true; - } - __syncthreads(); - for (size_t s = blockDim.x / 2; s > 0; s >>= 1) { - if (tid < s) { - s_data[tid] = s_data[tid] && s_data[tid + s]; - } - __syncthreads(); - } - if (tid == 0) { - temp_output[blockIdx.x] = s_data[0]; - } -} - -template -__global__ void finalAllReduceKernel( - bool *output, - const bool *block_results, - size_t num_blocks) { - __shared__ bool s_data[BLOCK_SIZE]; - size_t tid = threadIdx.x; - bool thread_val = true; - for (size_t i = tid; i < num_blocks; i += blockDim.x) { - thread_val = thread_val && block_results[i]; - } - s_data[tid] = thread_val; - __syncthreads(); - for (size_t s = BLOCK_SIZE / 2; s > 0; s >>= 1) { - if (tid < s) { - s_data[tid] = s_data[tid] && s_data[tid + s]; - } - __syncthreads(); - } - - if (tid == 0) { - *output = s_data[0]; - } -} - -template -__global__ void allKernel( - bool *output, - const Tdata *input, - size_t permuted_input_shape_size, - size_t output_shape_size, - size_t output_size, - size_t reduce_num, - size_t *permuted_input_shape, - size_t *output_shape, - ptrdiff_t *permuted_input_strides, - ptrdiff_t *output_strides) { - size_t tid = threadIdx.x; - size_t idx = tid + blockIdx.x * blockDim.x; - if (idx >= output_size) { - return; - } - size_t output_index = indexToOffset(idx, output_shape_size, output_shape, output_strides); - bool tempRes = true; - for (size_t i = 0; i < reduce_num; i++) { - size_t input_offset = indexToOffset(i + idx * reduce_num, permuted_input_shape_size, permuted_input_shape, permuted_input_strides); - tempRes = tempRes && static_cast(input[input_offset]); - } - output[output_index] = tempRes; -} - -#endif // __ALL_CUDA_H__ diff --git a/src/infiniop/ops/all/info.h b/src/infiniop/ops/all/info.h deleted file mode 100644 index f3f333fc8..000000000 --- a/src/infiniop/ops/all/info.h +++ /dev/null @@ -1,66 +0,0 @@ -#ifndef __ALL_INFO_H__ -#define __ALL_INFO_H__ -#include "../../../utils.h" -#include "../../tensor.h" -#include -#include -#include - -namespace op::all { -class AllInfo { - AllInfo() = default; - -public: - infiniDtype_t dtype; - std::vector permuted_input_shape; // need to permute - std::vector output_shape; - std::vector permuted_input_strides; // need to permute - std::vector output_strides; - size_t reduce_dim_size; // reduce dim size - size_t reduce_num; // number of elements to reduce for each output element - size_t input_size; // total number of input elements - size_t output_size; // total number of output elements - size_t ndim; // number of dimensions - static utils::Result create( - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - size_t *dim, - size_t dim_size, - bool keepdim) { - auto input_shape = input_desc->shape(); - auto input_strides = input_desc->strides(); - size_t input_ndim = input_desc->ndim(); - size_t reduce_num = 1; - for (size_t i = 0; i < dim_size; i++) { - reduce_num *= input_shape[dim[i]]; - } - std::vector permute_order; - for (size_t i = 0; i < input_ndim; i++) { - if (std::find(dim, dim + dim_size, i) == dim + dim_size) { - permute_order.push_back(i); - } - } - for (size_t i = 0; i < dim_size; i++) { - permute_order.push_back(dim[i]); - } - std::vector permuted_input_shape; - std::vector permuted_input_strides; - for (size_t i = 0; i < permute_order.size(); i++) { - permuted_input_shape.push_back(input_shape[permute_order[i]]); - permuted_input_strides.push_back(input_strides[permute_order[i]]); - } - return utils::Result(AllInfo{input_desc->dtype(), - permuted_input_shape, - output_desc->shape(), - permuted_input_strides, - output_desc->strides(), - dim_size, - reduce_num, - input_desc->numel(), - output_desc->numel(), - input_ndim}); - } -}; -} // namespace op::all - -#endif diff --git a/src/infiniop/ops/all/metax/all_metax.h b/src/infiniop/ops/all/metax/all_metax.h deleted file mode 100644 index 0f0ecc742..000000000 --- a/src/infiniop/ops/all/metax/all_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ALL_METAX_H__ -#define __ALL_METAX_H__ - -#include "../all_desc.h" - -DESCRIPTOR(metax); - -#endif diff --git a/src/infiniop/ops/all/metax/all_metax.maca b/src/infiniop/ops/all/metax/all_metax.maca deleted file mode 100644 index b95936585..000000000 --- a/src/infiniop/ops/all/metax/all_metax.maca +++ /dev/null @@ -1,117 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../cuda/kernel.cuh" -#include "all_metax.h" - -namespace op::all::metax { -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - size_t *dim, - size_t dim_size, - bool keepdim) { - auto result = AllInfo::create(output_desc, input_desc, dim, dim_size, keepdim); - CHECK_RESULT(result); - auto info = result.take(); - size_t workspace_size = 0; - workspace_size += (input_desc->ndim() + output_desc->ndim()) * (sizeof(size_t) + sizeof(ptrdiff_t)); - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info, workspace_size, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -infiniStatus_t launchKernel( - const AllInfo &info, - bool *output, const Tdata *input, - hcStream_t stream, void *workspace, size_t workspace_size) { - size_t input_ndim = info.permuted_input_shape.size(); - size_t output_ndim = info.output_shape.size(); - size_t input_size = info.input_size; - size_t output_size = info.output_size; - size_t reduce_num = info.reduce_num; - unsigned char *workspace_ptr = reinterpret_cast(workspace); - size_t workspace_offset = 0; - size_t *permuted_input_shape_hc = reinterpret_cast(workspace_ptr + workspace_offset); - size_t *output_shape_hc = permuted_input_shape_hc + input_ndim; - workspace_offset += (input_ndim + output_ndim) * sizeof(size_t); - - ptrdiff_t *permuted_input_strides_hc = reinterpret_cast(workspace_ptr + workspace_offset); - ptrdiff_t *output_strides_hc = permuted_input_strides_hc + input_ndim; - workspace_offset += (input_ndim + output_ndim) * sizeof(ptrdiff_t); - - CHECK_METAX(hcMemcpyAsync(permuted_input_shape_hc, info.permuted_input_shape.data(), input_ndim * sizeof(size_t), hcMemcpyHostToDevice, stream)); - CHECK_METAX(hcMemcpyAsync(output_shape_hc, info.output_shape.data(), output_ndim * sizeof(size_t), hcMemcpyHostToDevice, stream)); - CHECK_METAX(hcMemcpyAsync(permuted_input_strides_hc, info.permuted_input_strides.data(), input_ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, stream)); - CHECK_METAX(hcMemcpyAsync(output_strides_hc, info.output_strides.data(), output_ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, stream)); - - if (info.reduce_num == input_size) { - size_t grid_size = (input_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - bool *temp_output; - CHECK_METAX(hcMalloc(&temp_output, grid_size * sizeof(bool))); - allReduceTempKernel<<>>( - temp_output, input, input_size, input_ndim, permuted_input_shape_hc, permuted_input_strides_hc); - finalAllReduceKernel<<<1, BLOCK_SIZE>>>(output, temp_output, grid_size); - CHECK_METAX(hcFree(temp_output)); - } else { - size_t grid_size = (info.output_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - allKernel<<>>( - output, input, input_ndim, output_ndim, output_size, reduce_num, - permuted_input_shape_hc, output_shape_hc, permuted_input_strides_hc, output_strides_hc); - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - size_t *dim, - size_t dim_size, - bool keepdim, - void *stream_) const { - - hcStream_t stream = (hcStream_t)stream_; - -#define CALCULATE_ALL(BLOCK_SIZE, Tdata) \ - launchKernel( \ - _info, \ - (bool *)output, (const Tdata *)input, \ - stream, workspace, workspace_size) - -#define CALCULATE_ALL_WITH_BLOCK_SIZE(BLOCK_SIZE) \ - { \ - if (_info.dtype == INFINI_DTYPE_BOOL) \ - return CALCULATE_ALL(BLOCK_SIZE, bool); \ - else if (_info.dtype == INFINI_DTYPE_U8) \ - return CALCULATE_ALL(BLOCK_SIZE, uint8_t); \ - else \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } - - if (_opaque->internal->maxThreadsPerBlock() >= 256) { - CALCULATE_ALL_WITH_BLOCK_SIZE(256) - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::all::metax diff --git a/src/infiniop/ops/all/moore/all_moore.h b/src/infiniop/ops/all/moore/all_moore.h deleted file mode 100644 index d7dab5396..000000000 --- a/src/infiniop/ops/all/moore/all_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ALL_MOORE_H__ -#define __ALL_MOORE_H__ - -#include "../all_desc.h" - -DESCRIPTOR(moore); - -#endif diff --git a/src/infiniop/ops/all/moore/all_moore.mu b/src/infiniop/ops/all/moore/all_moore.mu deleted file mode 100644 index 624d47391..000000000 --- a/src/infiniop/ops/all/moore/all_moore.mu +++ /dev/null @@ -1,117 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "../cuda/kernel.cuh" -#include "all_moore.h" - -namespace op::all::moore { -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - size_t *dim, - size_t dim_size, - bool keepdim) { - auto result = AllInfo::create(output_desc, input_desc, dim, dim_size, keepdim); - CHECK_RESULT(result); - auto info = result.take(); - size_t workspace_size = 0; - workspace_size += (input_desc->ndim() + output_desc->ndim()) * (sizeof(size_t) + sizeof(ptrdiff_t)); - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info, workspace_size, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -infiniStatus_t launchKernel( - const AllInfo &info, - bool *output, const Tdata *input, - musaStream_t stream, void *workspace, size_t workspace_size) { - size_t input_ndim = info.permuted_input_shape.size(); - size_t output_ndim = info.output_shape.size(); - size_t input_size = info.input_size; - size_t output_size = info.output_size; - size_t reduce_num = info.reduce_num; - unsigned char *workspace_ptr = reinterpret_cast(workspace); - size_t workspace_offset = 0; - size_t *permuted_input_shape_musa = reinterpret_cast(workspace_ptr + workspace_offset); - size_t *output_shape_musa = permuted_input_shape_musa + input_ndim; - workspace_offset += (input_ndim + output_ndim) * sizeof(size_t); - - ptrdiff_t *permuted_input_strides_musa = reinterpret_cast(workspace_ptr + workspace_offset); - ptrdiff_t *output_strides_musa = permuted_input_strides_musa + input_ndim; - workspace_offset += (input_ndim + output_ndim) * sizeof(ptrdiff_t); - - CHECK_MOORE(musaMemcpyAsync(permuted_input_shape_musa, info.permuted_input_shape.data(), input_ndim * sizeof(size_t), musaMemcpyHostToDevice, stream)); - CHECK_MOORE(musaMemcpyAsync(output_shape_musa, info.output_shape.data(), output_ndim * sizeof(size_t), musaMemcpyHostToDevice, stream)); - CHECK_MOORE(musaMemcpyAsync(permuted_input_strides_musa, info.permuted_input_strides.data(), input_ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, stream)); - CHECK_MOORE(musaMemcpyAsync(output_strides_musa, info.output_strides.data(), output_ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, stream)); - - if (info.reduce_num == input_size) { - size_t grid_size = (input_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - bool *temp_output; - CHECK_MOORE(musaMalloc(&temp_output, grid_size * sizeof(bool))); - allReduceTempKernel<<>>( - temp_output, input, input_size, input_ndim, permuted_input_shape_musa, permuted_input_strides_musa); - finalAllReduceKernel<<<1, BLOCK_SIZE>>>(output, temp_output, grid_size); - CHECK_MOORE(musaFree(temp_output)); - } else { - size_t grid_size = (info.output_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - allKernel<<>>( - output, input, input_ndim, output_ndim, output_size, reduce_num, - permuted_input_shape_musa, output_shape_musa, permuted_input_strides_musa, output_strides_musa); - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - size_t *dim, - size_t dim_size, - bool keepdim, - void *stream_) const { - - musaStream_t stream = (musaStream_t)stream_; - -#define CALCULATE_ALL(BLOCK_SIZE, Tdata) \ - launchKernel( \ - _info, \ - (bool *)output, (const Tdata *)input, \ - stream, workspace, workspace_size) - -#define CALCULATE_ALL_WITH_BLOCK_SIZE(BLOCK_SIZE) \ - { \ - if (_info.dtype == INFINI_DTYPE_BOOL) \ - return CALCULATE_ALL(BLOCK_SIZE, bool); \ - else if (_info.dtype == INFINI_DTYPE_U8) \ - return CALCULATE_ALL(BLOCK_SIZE, uint8_t); \ - else \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } - - if (_opaque->internal->maxThreadsPerBlock() >= 256) { - CALCULATE_ALL_WITH_BLOCK_SIZE(256) - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::all::moore diff --git a/src/infiniop/ops/all/nvidia/all_nvidia.cu b/src/infiniop/ops/all/nvidia/all_nvidia.cu deleted file mode 100644 index f0858d2f7..000000000 --- a/src/infiniop/ops/all/nvidia/all_nvidia.cu +++ /dev/null @@ -1,117 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../cuda/kernel.cuh" -#include "all_nvidia.cuh" - -namespace op::all::nvidia { -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - size_t *dim, - size_t dim_size, - bool keepdim) { - auto result = AllInfo::create(output_desc, input_desc, dim, dim_size, keepdim); - CHECK_RESULT(result); - auto info = result.take(); - size_t workspace_size = 0; - workspace_size += (input_desc->ndim() + output_desc->ndim()) * (sizeof(size_t) + sizeof(ptrdiff_t)); - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info, workspace_size, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -infiniStatus_t launchKernel( - const AllInfo &info, - bool *output, const Tdata *input, - cudaStream_t stream, void *workspace, size_t workspace_size) { - size_t input_ndim = info.permuted_input_shape.size(); - size_t output_ndim = info.output_shape.size(); - size_t input_size = info.input_size; - size_t output_size = info.output_size; - size_t reduce_num = info.reduce_num; - unsigned char *workspace_ptr = reinterpret_cast(workspace); - size_t workspace_offset = 0; - size_t *permuted_input_shape_cuda = reinterpret_cast(workspace_ptr + workspace_offset); - size_t *output_shape_cuda = permuted_input_shape_cuda + input_ndim; - workspace_offset += (input_ndim + output_ndim) * sizeof(size_t); - - ptrdiff_t *permuted_input_strides_cuda = reinterpret_cast(workspace_ptr + workspace_offset); - ptrdiff_t *output_strides_cuda = permuted_input_strides_cuda + input_ndim; - workspace_offset += (input_ndim + output_ndim) * sizeof(ptrdiff_t); - - CHECK_CUDA(cudaMemcpyAsync(permuted_input_shape_cuda, info.permuted_input_shape.data(), input_ndim * sizeof(size_t), cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemcpyAsync(output_shape_cuda, info.output_shape.data(), output_ndim * sizeof(size_t), cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemcpyAsync(permuted_input_strides_cuda, info.permuted_input_strides.data(), input_ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemcpyAsync(output_strides_cuda, info.output_strides.data(), output_ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, stream)); - - if (info.reduce_num == input_size) { - size_t grid_size = (input_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - bool *temp_output; - CHECK_CUDA(cudaMalloc(&temp_output, grid_size * sizeof(bool))); - allReduceTempKernel<<>>( - temp_output, input, input_size, input_ndim, permuted_input_shape_cuda, permuted_input_strides_cuda); - finalAllReduceKernel<<<1, BLOCK_SIZE>>>(output, temp_output, grid_size); - CHECK_CUDA(cudaFree(temp_output)); - } else { - size_t grid_size = (info.output_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - allKernel<<>>( - output, input, input_ndim, output_ndim, output_size, reduce_num, - permuted_input_shape_cuda, output_shape_cuda, permuted_input_strides_cuda, output_strides_cuda); - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - size_t *dim, - size_t dim_size, - bool keepdim, - void *stream_) const { - - cudaStream_t stream = (cudaStream_t)stream_; - -#define CALCULATE_ALL(BLOCK_SIZE, Tdata) \ - launchKernel( \ - _info, \ - (bool *)output, (const Tdata *)input, \ - stream, workspace, workspace_size) - -#define CALCULATE_ALL_WITH_BLOCK_SIZE(BLOCK_SIZE) \ - { \ - if (_info.dtype == INFINI_DTYPE_BOOL) \ - return CALCULATE_ALL(BLOCK_SIZE, bool); \ - else if (_info.dtype == INFINI_DTYPE_U8) \ - return CALCULATE_ALL(BLOCK_SIZE, uint8_t); \ - else \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } - - if (_opaque->internal->maxThreadsPerBlock() >= 256) { - CALCULATE_ALL_WITH_BLOCK_SIZE(256) - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::all::nvidia diff --git a/src/infiniop/ops/all/nvidia/all_nvidia.cuh b/src/infiniop/ops/all/nvidia/all_nvidia.cuh deleted file mode 100644 index 111e0816f..000000000 --- a/src/infiniop/ops/all/nvidia/all_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ALL_NVIDIA_H__ -#define __ALL_NVIDIA_H__ - -#include "../all_desc.h" - -DESCRIPTOR(nvidia); - -#endif // __ALL_CUDA_API_H__ diff --git a/src/infiniop/ops/all/operator.cc b/src/infiniop/ops/all/operator.cc deleted file mode 100644 index 8ffb34b63..000000000 --- a/src/infiniop/ops/all/operator.cc +++ /dev/null @@ -1,191 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/all.h" -#include - -#ifdef ENABLE_CPU_API -#include "cpu/all_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/all_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/all_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/all_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAllDescriptor( - infiniopHandle_t handle, - infiniopAllDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - size_t *dim, - size_t dim_size, - bool keepdim) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::all::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output_desc, \ - input_desc, \ - dim, \ - dim_size, \ - keepdim) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAllWorkspaceSize(infiniopAllDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopAll( - infiniopAllDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - size_t *dim, - size_t dim_size, - bool keepdim, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, input, dim, dim_size, keepdim, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyAllDescriptor(infiniopAllDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/argwhere/argwhere.h b/src/infiniop/ops/argwhere/argwhere.h deleted file mode 100644 index 09fc872f4..000000000 --- a/src/infiniop/ops/argwhere/argwhere.h +++ /dev/null @@ -1,75 +0,0 @@ -#ifndef __INFINIOP_ARGWHERE_H__ -#define __INFINIOP_ARGWHERE_H__ -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" -#include "infinicore.h" -#include - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::argwhere::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - ArgwhereInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - ArgwhereInfo info, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t x_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void **y, \ - size_t *count, \ - const void *x, \ - void *stream) const; \ - }; \ - } - -class ArgwhereInfo { -private: - ArgwhereInfo() = default; - -public: - infiniDtype_t dtype; - std::vector strides; - std::vector shapes; - size_t num_elements; - - static utils::Result - create( - infiniopTensorDescriptor_t x_desc) { - CHECK_OR_RETURN(x_desc != nullptr, - INFINI_STATUS_NULL_POINTER); - - const infiniDtype_t data_type = x_desc->dtype(); - CHECK_DTYPE(data_type, INFINI_DTYPE_F32); - - return utils::Result(ArgwhereInfo{ - data_type, - x_desc->strides(), - x_desc->shape(), - x_desc->numel()}); - } -}; -#endif diff --git a/src/infiniop/ops/argwhere/cpu/argwhere_cpu.cc b/src/infiniop/ops/argwhere/cpu/argwhere_cpu.cc deleted file mode 100644 index 98782f958..000000000 --- a/src/infiniop/ops/argwhere/cpu/argwhere_cpu.cc +++ /dev/null @@ -1,89 +0,0 @@ -#include "argwhere_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include "infinicore.h" -#include -#include -#include -namespace op::argwhere::cpu { -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc) { - - auto handle = reinterpret_cast(handle_); - - auto info = ArgwhereInfo::create(x_desc); - CHECK_RESULT(info); - - *desc_ptr = new Descriptor( - info.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateArgWhere( - const ArgwhereInfo &info, - void *workspace, - size_t workspace_size, - void **y, - size_t *count, - const void *x) { - - const Tdata *x_data = reinterpret_cast(x); - - std::vector positions; - const size_t ndim = info.shapes.size(); - - for (size_t i = 0; i < info.num_elements; i++) { - size_t pos = 0; - size_t tmp = i; - - std::vector coord(ndim); - - // unravel index - for (size_t j = ndim; j-- > 0;) { - coord[j] = tmp % info.shapes[j]; - tmp /= info.shapes[j]; - pos += coord[j] * info.strides[j]; - } - - // PyTorch semantics: != 0 - if (x_data[pos] != Tdata(0)) { - for (size_t j = 0; j < ndim; j++) { - positions.push_back(coord[j]); - } - } - } - - *count = positions.size() / ndim; - - *y = new int64_t[positions.size()]; - memcpy(*y, positions.data(), positions.size() * sizeof(int64_t)); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void **y, - size_t *count, - const void *x, - void *stream) const { - switch (_info.dtype) { - case INFINI_DTYPE_F32: - return calculateArgWhere(_info, workspace, workspace_size, y, count, x); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::argwhere::cpu diff --git a/src/infiniop/ops/argwhere/cpu/argwhere_cpu.h b/src/infiniop/ops/argwhere/cpu/argwhere_cpu.h deleted file mode 100644 index 46d828fde..000000000 --- a/src/infiniop/ops/argwhere/cpu/argwhere_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __INFINIOP_ARGWHERE_CPU_API_H__ -#define __INFINIOP_ARGWHERE_CPU_API_H__ - -#include "../argwhere.h" - -DESCRIPTOR(cpu) - -#endif diff --git a/src/infiniop/ops/argwhere/cuda/kernel.cuh b/src/infiniop/ops/argwhere/cuda/kernel.cuh deleted file mode 100644 index e562cf88e..000000000 --- a/src/infiniop/ops/argwhere/cuda/kernel.cuh +++ /dev/null @@ -1,219 +0,0 @@ -#ifndef ARGWHERE_NVIDIA_KERNEL_H -#define ARGWHERE_NVIDIA_KERNEL_H -#include -#include - -__device__ void Index2Pos(size_t index, size_t ndim, const size_t *shapes, - size_t *pos) { - for (int i = ndim - 1; i >= 0; i--) { - pos[i] = index % shapes[i]; - index /= shapes[i]; - } -} -__device__ size_t pos2dest(size_t *pos, size_t ndim, const ptrdiff_t *strides) { - size_t dest = 0; - for (size_t i = 0; i < ndim; i++) { - dest += pos[i] * strides[i]; - } - return dest; -} - -template -__global__ void parallel_block_argwhere_kernel(T *data, int64_t *results, - size_t N, const size_t *shapes, const ptrdiff_t *strides, size_t ndim, size_t *count) { - extern __shared__ char smem[]; - size_t *tmp = reinterpret_cast(smem); - - size_t pos1[5], pos2[5]; // 两个数的在tensor中的索引 - bool is_zero1 = false, is_zero2 = false; - int tid = threadIdx.x; - int leaf_num = blockDim.x * 2; // equals to length of tmp - - if (tid * 2 < N) { - Index2Pos(tid * 2, ndim, shapes, pos1); - is_zero1 = fabs(data[pos2dest(pos1, ndim, strides)]) <= 1e-5; - tmp[tid * 2] = !is_zero1; - } - if (tid * 2 + 1 < N) { - Index2Pos(tid * 2 + 1, ndim, shapes, pos2); - is_zero2 = fabs(data[pos2dest(pos2, ndim, strides)]) <= 1e-5; - tmp[tid * 2 + 1] = !is_zero2; - } - - __syncthreads(); - - int offset = 1; - for (int d = leaf_num >> 1; d > 0; d >>= 1) { - if (tid < d) { - int ai = offset * (2 * tid + 1) - 1; - int bi = offset * (2 * tid + 2) - 1; - tmp[bi] += tmp[ai]; - } - offset *= 2; - __syncthreads(); - } - - if (tid == 0) { - tmp[leaf_num - 1] = 0; - } - __syncthreads(); - - for (int d = 1; d < leaf_num; d *= 2) { - offset >>= 1; - if (tid < d) { - int ai = offset * (2 * tid + 1) - 1; - int bi = offset * (2 * tid + 2) - 1; - - int v = tmp[ai]; - tmp[ai] = tmp[bi]; - tmp[bi] += v; - } - __syncthreads(); - } - - // 写入最终结果 - if (!is_zero1 && tid * 2 < N) { - for (int i = 0; i < ndim; i++) { - results[tmp[2 * tid] * ndim + i] = pos1[i]; - } - } - if (!is_zero2 && tid * 2 + 1 < N) { - for (int i = 0; i < ndim; i++) { - results[tmp[2 * tid + 1] * ndim + i] = pos2[i]; - } - } - if (tid == blockDim.x - 1) { - // printf("blockIdxDim = %d\n", blockDim.x); - *count = tmp[N - 1] + (N == blockDim.x * 2 ? 1 : tmp[N] != tmp[N - 1]); - // printf("finally: count = %d\n", tmp[leaf_num - 1]); - } -} - -template -__global__ void parallel_block_scan_kernel(size_t N, int64_t *pre_sum) { - // single block scan - extern __shared__ char smem[]; - int64_t *tmp = reinterpret_cast(smem); - - int tid = threadIdx.x; - int leaf_num = blockDim.x * 2; // equals to length of tmp - - tmp[tid * 2] = tid * 2 < N - 1 ? pre_sum[tid * 2] : 0; - tmp[tid * 2 + 1] = tid * 2 + 1 < N - 1 ? pre_sum[tid * 2 + 1] : 0; - __syncthreads(); - - int offset = 1; - for (int d = leaf_num >> 1; d > 0; d >>= 1) { - if (tid < d) { - int ai = offset * (2 * tid + 1) - 1; - int bi = offset * (2 * tid + 2) - 1; - tmp[bi] += tmp[ai]; - } - offset *= 2; - __syncthreads(); - } - if (tid == 0) { - tmp[leaf_num - 1] = 0; - } - __syncthreads(); - for (int d = 1; d < leaf_num; d <<= 1) { - offset >>= 1; - if (tid < d) { - int ai = offset * (2 * tid + 1) - 1; - int bi = offset * (2 * tid + 2) - 1; - - int v = tmp[ai]; - tmp[ai] = tmp[bi]; - tmp[bi] += v; - } - __syncthreads(); - } - if (tid * 2 < N) { - pre_sum[tid * 2] = tmp[tid * 2]; - } - if (tid * 2 + 1 < N) { - pre_sum[tid * 2 + 1] = tmp[tid * 2 + 1]; - } -} - -template -__global__ void -parallel_large_argwhere_kernel(const T *data, int64_t *block_sum, - int64_t *results, size_t N, const size_t *shapes, - const ptrdiff_t *strides, size_t ndim) { - // To be implemented for large N - extern __shared__ int64_t tmp_argwhere[]; - int bid = blockIdx.x, tid = threadIdx.x; - size_t pos1[5], pos2[5]; // 两个数的在tensor中的索引 - bool is_zero1 = false, is_zero2 = false; - int block_offset = bid * blockDim.x * 2, leaf_num = blockDim.x * 2; - tmp_argwhere[2 * tid] = tmp_argwhere[2 * tid + 1] = 0; - if (block_offset + tid * 2 < N) { - Index2Pos(block_offset + tid * 2, ndim, shapes, pos1); - is_zero1 = fabs(data[pos2dest(pos1, ndim, strides)]) <= 1e-9; - tmp_argwhere[2 * tid] = !is_zero1; - } - if (block_offset + tid * 2 + 1 < N) { - Index2Pos(block_offset + tid * 2 + 1, ndim, shapes, pos2); - is_zero2 = fabs(data[pos2dest(pos2, ndim, strides)]) <= 1e-9; - tmp_argwhere[2 * tid + 1] = !is_zero2; - } - __syncthreads(); - int offset = 1; - for (int d = leaf_num >> 1; d > 0; d >>= 1) { - if (tid < d) { - int ai = offset * (2 * tid + 1) - 1; - int bi = offset * (2 * tid + 2) - 1; - tmp_argwhere[bi] += tmp_argwhere[ai]; - } - offset *= 2; - __syncthreads(); - } - if (tid == 0) { - block_sum[bid] = tmp_argwhere[leaf_num - 1]; - // printf("tmp_argwhere[%d] = %lld\n", leaf_num - 1, - // tmp_argwhere[leaf_num - 1]); - // printf("block_sum[%d] = %lld\n", bid, block_sum[bid]); - tmp_argwhere[leaf_num - 1] = 0; - } - __syncthreads(); - for (int d = 1; d < leaf_num; d <<= 1) { - offset >>= 1; - if (tid < d) { - int ai = offset * (2 * tid + 1) - 1; - int bi = offset * (2 * tid + 2) - 1; - int v = tmp_argwhere[ai]; - tmp_argwhere[ai] = tmp_argwhere[bi]; - tmp_argwhere[bi] += v; - } - __syncthreads(); - } - - if (!is_zero1 && block_offset + tid * 2 < N) { - for (int i = 0; i < ndim; i++) { - results[(tmp_argwhere[2 * tid] + block_offset) * ndim + i] = pos1[i]; - } - } - if (!is_zero2 && block_offset + tid * 2 + 1 < N) { - for (int i = 0; i < ndim; i++) { - results[(tmp_argwhere[2 * tid + 1] + block_offset) * ndim + i] = pos2[i]; - } - } -} - -__global__ void add_block_offset_kernel(int64_t *results, int64_t *tmp, - int64_t *block_sums, size_t ndim) { - int bid = blockIdx.x, tid = threadIdx.x; - size_t block_offset = block_sums[bid], origin_offset = bid * blockDim.x * 2; - if (2 * tid < block_sums[bid + 1] - block_sums[bid]) { - for (int i = 0; i < ndim; i++) { - results[(block_offset + 2 * tid) * ndim + i] = tmp[(origin_offset + tid * 2) * ndim + i]; - } - } - if (2 * tid + 1 < block_sums[bid + 1] - block_sums[bid]) { - for (int i = 0; i < ndim; i++) { - results[(block_offset + 2 * tid + 1) * ndim + i] = tmp[(origin_offset + tid * 2 + 1) * ndim + i]; - } - } -} -#endif diff --git a/src/infiniop/ops/argwhere/metax/argwhere_metax.h b/src/infiniop/ops/argwhere/metax/argwhere_metax.h deleted file mode 100644 index 8335f413b..000000000 --- a/src/infiniop/ops/argwhere/metax/argwhere_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef ARGWHERE_METAX_H -#define ARGWHERE_METAX_H - -#include "../argwhere.h" - -DESCRIPTOR(metax) - -#endif diff --git a/src/infiniop/ops/argwhere/metax/argwhere_metax.maca b/src/infiniop/ops/argwhere/metax/argwhere_metax.maca deleted file mode 100644 index 41e621c43..000000000 --- a/src/infiniop/ops/argwhere/metax/argwhere_metax.maca +++ /dev/null @@ -1,94 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../cuda/kernel.cuh" -#include "argwhere_metax.h" -#include "infinicore.h" -#include -#include - -infiniStatus_t launchKernel(const void *data, int64_t *results, size_t N, - size_t M, const size_t *shapes, - const ptrdiff_t *strides, size_t ndim, - infiniDtype_t dtype, size_t *count) { - - if (dtype == INFINI_DTYPE_F32) { - parallel_block_argwhere_kernel<<<1, M / 2, M * sizeof(size_t)>>>( - (float *)data, results, N, shapes, strides, ndim, count); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -namespace op::argwhere::metax { -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create(infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc) { - auto info = ArgwhereInfo::create(x_desc); - CHECK_RESULT(info); - size_t workspace_size = x_desc->ndim() * sizeof(size_t) * 2 + x_desc->ndim() * sizeof(int64_t) * x_desc->numel() + sizeof(size_t); - *desc_ptr = new Descriptor( - info.take(), workspace_size, - new Opaque{reinterpret_cast(handle)->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -int nextPowerOfTwo(int x) { - int power = 1; - while (power < x) { - power *= 2; - } - return power; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void **y, size_t *count, const void *x, - void *stream) const { - hcStream_t hc_stream = static_cast(stream); - size_t ndim = _info.strides.size(); - ptrdiff_t *strides = static_cast(workspace); - size_t *shapes = reinterpret_cast(strides + ndim); - int64_t *result = reinterpret_cast(shapes + ndim); - size_t *count_cuda = reinterpret_cast(result + _info.num_elements * ndim); - - hcMemcpyAsync(shapes, _info.shapes.data(), - _info.shapes.size() * sizeof(size_t), hcMemcpyHostToDevice, - hc_stream); - hcMemcpyAsync(strides, _info.strides.data(), - _info.strides.size() * sizeof(ptrdiff_t), hcMemcpyHostToDevice, - hc_stream); - // hcStreamSynchronize(hc_stream); - size_t M = nextPowerOfTwo(_info.num_elements); - CHECK_STATUS(launchKernel(x, result, _info.num_elements, M, shapes, strides, - ndim, INFINI_DTYPE_F32, count_cuda)); - // 从设备内存中读取 count_cuda 的值 - hcMemcpyAsync(count, count_cuda, sizeof(size_t), hcMemcpyDeviceToHost); - hcStreamSynchronize(hc_stream); - - // 写回结果 - *y = new int64_t[(*count) * ndim]; - // cudaStreamSynchronize(cuda_stream); - - // spdlog::debug("count_cuda:{}", *count_cuda); - // spdlog::debug("N:{}", N); - - hcMemcpyAsync(*y, result, sizeof(int64_t) * (*count) * ndim, - hcMemcpyDeviceToHost, hc_stream); - - // hcStreamSynchronize(hc_stream); - // for (size_t i = 0; i < (*count) * ndim; i++) { - // spdlog::debug("(*y)[{}]:{}", i, static_cast(*y)[i]); - // } - // hcFreeAsync(result, hc_stream); - // hcFreeAsync(count_cuda, hc_stream); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::argwhere::metax diff --git a/src/infiniop/ops/argwhere/moore/argwhere_kernel.h b/src/infiniop/ops/argwhere/moore/argwhere_kernel.h deleted file mode 100644 index e32727e9a..000000000 --- a/src/infiniop/ops/argwhere/moore/argwhere_kernel.h +++ /dev/null @@ -1,91 +0,0 @@ -#ifndef ARGWHERE_MOORE_KERNEL_H -#define ARGWHERE_MOORE_KERNEL_H - -__device__ void Index2Pos(size_t index, size_t ndim, const size_t *shapes, - size_t *pos) { - for (int i = ndim - 1; i >= 0; i--) { - pos[i] = index % shapes[i]; - index /= shapes[i]; - } -} -__device__ size_t pos2dest(size_t *pos, size_t ndim, const size_t *shapes, - const ptrdiff_t *strides) { - size_t dest = 0; - for (size_t i = 0; i < ndim; i++) { - dest += pos[i] * strides[i]; - } - return dest; -} -template -__global__ void parallel_block_argwhere_kernel(T *data, int64_t *results, - size_t N, const size_t *shapes, - const ptrdiff_t *strides, - size_t ndim, size_t *count) { - extern __shared__ size_t tmp[]; - size_t pos1[5], pos2[5]; // 两个数的在tensor中的索引 - bool is_zero1 = false, is_zero2 = false; - int tid = threadIdx.x; - int leaf_num = blockDim.x * 2; - - if (tid * 2 < N) { - Index2Pos(tid * 2, ndim, shapes, pos1); - is_zero1 = fabs(data[pos2dest(pos1, ndim, shapes, strides)]) <= 1e-5; - tmp[tid * 2] = !is_zero1; - } else { - tmp[tid * 2] = 0; - } - - if (tid * 2 + 1 < N) { - Index2Pos(tid * 2 + 1, ndim, shapes, pos2); - is_zero2 = fabs(data[pos2dest(pos2, ndim, shapes, strides)]) <= 1e-5; - tmp[tid * 2 + 1] = !is_zero2; - } else { - tmp[tid * 2 + 1] = 0; - } - - __syncthreads(); - - int offset = 1; - for (int d = leaf_num >> 1; d > 0; d >>= 1) { - if (tid < d) { - int ai = offset * (2 * tid + 1) - 1; - int bi = offset * (2 * tid + 2) - 1; - tmp[bi] += tmp[ai]; - } - offset *= 2; - __syncthreads(); - } - - if (tid == 0) { - tmp[leaf_num - 1] = 0; - } - __syncthreads(); - - for (int d = 1; d < leaf_num; d *= 2) { - offset >>= 1; - if (tid < d) { - int ai = offset * (2 * tid + 1) - 1; - int bi = offset * (2 * tid + 2) - 1; - - int v = tmp[ai]; - tmp[ai] = tmp[bi]; - tmp[bi] += v; - } - __syncthreads(); - } - if (!is_zero1 && tid * 2 < N) { - for (int i = 0; i < ndim; i++) { - results[tmp[2 * tid] * ndim + i] = pos1[i]; - } - } - if (!is_zero2 && tid * 2 + 1 < N) { - for (int i = 0; i < ndim; i++) { - results[tmp[2 * tid + 1] * ndim + i] = pos2[i]; - } - } - if (tid == blockDim.x - 1) { - *count = tmp[N - 1] + (N == blockDim.x * 2 ? 1 : tmp[N] != tmp[N - 1]); - } -} - -#endif diff --git a/src/infiniop/ops/argwhere/moore/argwhere_moore.h b/src/infiniop/ops/argwhere/moore/argwhere_moore.h deleted file mode 100644 index 903915893..000000000 --- a/src/infiniop/ops/argwhere/moore/argwhere_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef ARGWHERE_MOORE_H -#define ARGWHERE_MOORE_H - -#include "../argwhere.h" - -DESCRIPTOR(moore) - -#endif diff --git a/src/infiniop/ops/argwhere/moore/argwhere_moore.mu b/src/infiniop/ops/argwhere/moore/argwhere_moore.mu deleted file mode 100644 index ac95d78e8..000000000 --- a/src/infiniop/ops/argwhere/moore/argwhere_moore.mu +++ /dev/null @@ -1,86 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "argwhere_kernel.h" -#include "argwhere_moore.h" -#include "infinicore.h" - -infiniStatus_t launchKernel(const void *data, int64_t *results, size_t N, - size_t M, const size_t *shapes, - const ptrdiff_t *strides, size_t ndim, - infiniDtype_t dtype, size_t *count) { - - if (dtype == INFINI_DTYPE_F32) { - parallel_block_argwhere_kernel<<<1, M / 2, M * sizeof(size_t)>>>( - (float *)data, results, N, shapes, strides, ndim, count); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -namespace op::argwhere::moore { -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create(infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc) { - auto info = ArgwhereInfo::create(x_desc); - CHECK_RESULT(info); - size_t workspace_size = x_desc->ndim() * sizeof(size_t) * 2 + x_desc->ndim() * sizeof(int64_t) * x_desc->numel() + sizeof(size_t); - *desc_ptr = new Descriptor( - info.take(), workspace_size, - new Opaque{ - reinterpret_cast(handle)->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -int nextPowerOfTwo(int x) { - int power = 1; - while (power < x) { - power *= 2; - } - return power; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void **y, size_t *count, const void *x, - void *stream) const { - musaStream_t moore_stream = static_cast(stream); - size_t ndim = _info.strides.size(); - ptrdiff_t *strides = static_cast(workspace); - size_t *shapes = reinterpret_cast(strides + ndim); - int64_t *result = reinterpret_cast(shapes + ndim); - size_t *count_cuda = reinterpret_cast(result + _info.num_elements * ndim); - - musaMemcpyAsync(shapes, _info.shapes.data(), - _info.shapes.size() * sizeof(size_t), musaMemcpyHostToDevice, - moore_stream); - musaMemcpyAsync(strides, _info.strides.data(), - _info.strides.size() * sizeof(ptrdiff_t), - musaMemcpyHostToDevice, moore_stream); - // musaStreamSynchronize(moore_stream); - size_t M = nextPowerOfTwo(_info.num_elements); - // musaStreamSynchronize(moore_stream); - CHECK_STATUS(launchKernel(x, result, _info.num_elements, M, shapes, strides, - ndim, INFINI_DTYPE_F32, count_cuda)); - // 从设备内存中读取 count_cuda 的值 - musaMemcpyAsync(count, count_cuda, sizeof(size_t), musaMemcpyDeviceToHost, - moore_stream); - musaStreamSynchronize(moore_stream); - - // 写回结果 - *y = new int64_t[(*count) * ndim]; - // cudaStreamSynchronize(cuda_stream); - - musaMemcpyAsync(*y, result, sizeof(int64_t) * (*count) * ndim, - musaMemcpyDeviceToHost, moore_stream); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::argwhere::moore diff --git a/src/infiniop/ops/argwhere/nvidia/argwhere_nvidia.cu b/src/infiniop/ops/argwhere/nvidia/argwhere_nvidia.cu deleted file mode 100644 index 20e46f7bc..000000000 --- a/src/infiniop/ops/argwhere/nvidia/argwhere_nvidia.cu +++ /dev/null @@ -1,107 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../cuda/kernel.cuh" -#include "argwhere_nvidia.cuh" -#include "infinicore.h" -#include -#include - -int nextPowerOfTwo(int x) { - int power = 1; - while (power < x) { - power *= 2; - } - return power; -} - -infiniStatus_t launchKernel(const void *data, int64_t *tmp, int64_t *results, - int64_t *block_sum, size_t N, const size_t *shapes, - const ptrdiff_t *strides, size_t ndim, - infiniDtype_t dtype, size_t *count, - int maxThreadsPerBlock, cudaStream_t stream) { - - if (dtype == INFINI_DTYPE_F32) { - int block_size = maxThreadsPerBlock * 2; - int num_blocks = (N + block_size - 1) / block_size; - parallel_large_argwhere_kernel - <<>>( - static_cast(data), block_sum, tmp, N, shapes, - strides, ndim); - if (num_blocks > 1) { - // 计算前缀和 - parallel_block_scan_kernel - <<<1, maxThreadsPerBlock, sizeof(int64_t) * block_size>>>( - num_blocks + 1, block_sum); - // 重新整理结果 - add_block_offset_kernel<<>>( - results, tmp, block_sum, ndim); - cudaMemcpyAsync(count, &block_sum[num_blocks], sizeof(size_t), - cudaMemcpyDeviceToHost, stream); - } else { - cudaMemcpyAsync(count, &block_sum[0], sizeof(size_t), - cudaMemcpyDeviceToHost, stream); - cudaMemcpyAsync(results, tmp, sizeof(int64_t) * (*count) * ndim, - cudaMemcpyDeviceToDevice, stream); - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -namespace op::argwhere::nvidia { -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create(infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc) { - auto info = ArgwhereInfo::create(x_desc); - CHECK_RESULT(info); - Opaque *opaque = new Opaque{ - reinterpret_cast(handle)->internal()}; - size_t workspace_size = x_desc->ndim() * sizeof(size_t) * 2 + x_desc->ndim() * sizeof(int64_t) * x_desc->numel() * 2 + sizeof(int64_t) * opaque->internal->maxThreadsPerBlock(); - *desc_ptr = new Descriptor(info.take(), workspace_size, opaque, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void **y, size_t *count, const void *x, - void *stream) const { - cudaStream_t cuda_stream = static_cast(stream); - size_t ndim = _info.strides.size(); - ptrdiff_t *strides = static_cast(workspace); - size_t *shapes = reinterpret_cast(strides + ndim); - int64_t *tmp = reinterpret_cast(shapes + ndim); - int64_t *result = reinterpret_cast(tmp + _info.num_elements * ndim); - int64_t *block_sum = reinterpret_cast(result + _info.num_elements * ndim); - - cudaMemcpyAsync(shapes, _info.shapes.data(), - _info.shapes.size() * sizeof(size_t), cudaMemcpyHostToDevice, - cuda_stream); - cudaMemcpyAsync(strides, _info.strides.data(), - _info.strides.size() * sizeof(ptrdiff_t), - cudaMemcpyHostToDevice, cuda_stream); - - CHECK_STATUS(launchKernel(x, tmp, result, block_sum, _info.num_elements, - shapes, strides, ndim, INFINI_DTYPE_F32, count, - _opaque->internal->maxThreadsPerBlock(), - cuda_stream)); - // 从设备内存中读取 count_cuda 的值 - - cudaStreamSynchronize(cuda_stream); - - // 写回结果 - *y = new int64_t[(*count) * ndim]; - - cudaMemcpyAsync(*y, result, sizeof(int64_t) * (*count) * ndim, - cudaMemcpyDeviceToHost, cuda_stream); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::argwhere::nvidia diff --git a/src/infiniop/ops/argwhere/nvidia/argwhere_nvidia.cuh b/src/infiniop/ops/argwhere/nvidia/argwhere_nvidia.cuh deleted file mode 100644 index b4e369ad2..000000000 --- a/src/infiniop/ops/argwhere/nvidia/argwhere_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef ARGWHERE_NVIDIA_H -#define ARGWHERE_NVIDIA_H - -#include "../argwhere.h" - -DESCRIPTOR(nvidia) - -#endif diff --git a/src/infiniop/ops/argwhere/operator.cc b/src/infiniop/ops/argwhere/operator.cc deleted file mode 100644 index e1eaccfd3..000000000 --- a/src/infiniop/ops/argwhere/operator.cc +++ /dev/null @@ -1,182 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/argwhere.h" - -#ifdef ENABLE_CPU_API -#include "cpu/argwhere_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/argwhere_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/argwhere_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/argwhere_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateArgwhereDescriptor( - infiniopHandle_t handle, - infiniopArgwhereDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t input_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::argwhere::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - input_desc) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetArgwhereWorkspaceSize( - infiniopArgwhereDescriptor_t desc, - size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopArgwhere( - infiniopArgwhereDescriptor_t desc, - void *workspace, - size_t workspace_size, - void **output, - size_t *count, - const void *input, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, count, input, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyArgwhereDescriptor(infiniopArgwhereDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/asin/cpu/asin_cpu.cc b/src/infiniop/ops/asin/cpu/asin_cpu.cc deleted file mode 100644 index 48a05d76c..000000000 --- a/src/infiniop/ops/asin/cpu/asin_cpu.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "asin_cpu.h" - -namespace op::asin::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - // create CPU elementwise descriptor - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::asin::cpu diff --git a/src/infiniop/ops/asin/cpu/asin_cpu.h b/src/infiniop/ops/asin/cpu/asin_cpu.h deleted file mode 100644 index 409f15779..000000000 --- a/src/infiniop/ops/asin/cpu/asin_cpu.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef __ASIN_CPU_H__ -#define __ASIN_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" -#include - -ELEMENTWISE_DESCRIPTOR(asin, cpu) - -namespace op::asin::cpu { -typedef struct AsinOp { -public: - static constexpr size_t num_inputs = 1; - - template - T operator()(const T &input) const { - return std::asin(input); - } -} AsinOp; -} // namespace op::asin::cpu - -#endif // __ASIN_CPU_H__ diff --git a/src/infiniop/ops/asin/cuda/kernel.cuh b/src/infiniop/ops/asin/cuda/kernel.cuh deleted file mode 100644 index 00802692b..000000000 --- a/src/infiniop/ops/asin/cuda/kernel.cuh +++ /dev/null @@ -1,42 +0,0 @@ -#ifndef __ASIN_CUDA_H__ -#define __ASIN_CUDA_H__ - -namespace op::asin::cuda { -typedef struct AsinOp { - static constexpr size_t num_inputs = 1; - - __device__ __forceinline__ float asin_f32_func(float x) const { - return asinf(x); - } - template - __device__ __forceinline__ T operator()(const T &input) const { - if constexpr (std::is_same_v) { - float2 vf = __half22float2(input); - float2 vr = make_float2(asin_f32_func(vf.x), asin_f32_func(vf.y)); - return __float22half2_rn(vr); - } else if constexpr (std::is_same_v) { - float xf = __half2float(input); - float yf = asin_f32_func(xf); - return __float2half_rn(yf); - } else if constexpr (std::is_same_v) { - float f0 = __bfloat162float(__low2bfloat16(input)); - float f1 = __bfloat162float(__high2bfloat16(input)); - float r0 = asin_f32_func(f0); - float r1 = asin_f32_func(f1); - return __floats2bfloat162_rn(r0, r1); - } else if constexpr (std::is_same_v) { - float xf = __bfloat162float(input); - float rf = asin_f32_func(xf); - return __float2bfloat16_rn(rf); - } else if constexpr (std::is_same_v) { - return asin_f32_func(input); - } else if constexpr (std::is_same_v) { - return std::asin(input); - } else { - return std::asin(input); - } - } -} AsinOp; -} // namespace op::asin::cuda - -#endif // __ASIN_CUDA_H__ diff --git a/src/infiniop/ops/asin/metax/asin_metax.h b/src/infiniop/ops/asin/metax/asin_metax.h deleted file mode 100644 index e1805508d..000000000 --- a/src/infiniop/ops/asin/metax/asin_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ASIN_METAX_API_H__ -#define __ASIN_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(asin, metax) - -#endif // __ASIN_METAX_API_H__ diff --git a/src/infiniop/ops/asin/metax/asin_metax.maca b/src/infiniop/ops/asin/metax/asin_metax.maca deleted file mode 100644 index b09ebd936..000000000 --- a/src/infiniop/ops/asin/metax/asin_metax.maca +++ /dev/null @@ -1,61 +0,0 @@ -#include "../../../elementwise/metax/elementwise_metax.h" -#include "asin_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::asin::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t -Descriptor::create(infiniopHandle_t handle_, Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, - INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - // create CUDA elementwise descriptor - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::AsinOp, half>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::AsinOp, cuda_bfloat16>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::AsinOp, float>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::AsinOp, double>( - _info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::asin::metax diff --git a/src/infiniop/ops/asin/moore/asin_moore.h b/src/infiniop/ops/asin/moore/asin_moore.h deleted file mode 100644 index 49394ea9d..000000000 --- a/src/infiniop/ops/asin/moore/asin_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ASIN_MOORE_API_H__ -#define __ASIN_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(asin, moore) - -#endif // __ASIN_MOORE_API_H__ diff --git a/src/infiniop/ops/asin/moore/asin_moore.mu b/src/infiniop/ops/asin/moore/asin_moore.mu deleted file mode 100644 index 7e2150e6c..000000000 --- a/src/infiniop/ops/asin/moore/asin_moore.mu +++ /dev/null @@ -1,60 +0,0 @@ -#include "asin_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "asin_moore_kernel.h" - -namespace op::asin::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - // create MOORE elementwise descriptor - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, moore::AsinOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, moore::AsinOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, moore::AsinOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, moore::AsinOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::asin::moore diff --git a/src/infiniop/ops/asin/moore/asin_moore_kernel.h b/src/infiniop/ops/asin/moore/asin_moore_kernel.h deleted file mode 100644 index 0401417f2..000000000 --- a/src/infiniop/ops/asin/moore/asin_moore_kernel.h +++ /dev/null @@ -1,42 +0,0 @@ -#ifndef __ASIN_MOORE_KERNEL_H__ -#define __ASIN_MOORE_KERNEL_H__ - -namespace op::asin::moore { -typedef struct AsinOp { - static constexpr size_t num_inputs = 1; - - __device__ __forceinline__ float asin_f32_func(float x) const { - return asinf(x); - } - template - __device__ __forceinline__ T operator()(const T &input) const { - if constexpr (std::is_same_v) { - float2 vf = __half22float2(input); - float2 vr = make_float2(asin_f32_func(vf.x), asin_f32_func(vf.y)); - return __float22half2_rn(vr); - } else if constexpr (std::is_same_v) { - float xf = __half2float(input); - float yf = asin_f32_func(xf); - return __float2half_rn(yf); - } else if constexpr (std::is_same_v) { - float f0 = __bfloat162float(__low2bfloat16(input)); - float f1 = __bfloat162float(__high2bfloat16(input)); - float r0 = asin_f32_func(f0); - float r1 = asin_f32_func(f1); - return __floats2bfloat162_rn(r0, r1); - } else if constexpr (std::is_same_v) { - float xf = __bfloat162float(input); - float rf = asin_f32_func(xf); - return __float2bfloat16_rn(rf); - } else if constexpr (std::is_same_v) { - return asin_f32_func(input); - } else if constexpr (std::is_same_v) { - return std::asin(input); - } else { - return std::asin(input); - } - } -} AsinOp; -} // namespace op::asin::moore - -#endif // __ASIN_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/asin/nvidia/asin_nvidia.cu b/src/infiniop/ops/asin/nvidia/asin_nvidia.cu deleted file mode 100644 index 6def60963..000000000 --- a/src/infiniop/ops/asin/nvidia/asin_nvidia.cu +++ /dev/null @@ -1,61 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "asin_nvidia.cuh" - -namespace op::asin::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t -Descriptor::create(infiniopHandle_t handle_, Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, - INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - // create CUDA elementwise descriptor - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::AsinOp, half>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::AsinOp, cuda_bfloat16>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::AsinOp, float>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::AsinOp, double>( - _info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::asin::nvidia diff --git a/src/infiniop/ops/asin/nvidia/asin_nvidia.cuh b/src/infiniop/ops/asin/nvidia/asin_nvidia.cuh deleted file mode 100644 index 624676428..000000000 --- a/src/infiniop/ops/asin/nvidia/asin_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ASIN_CUDA_API_H__ -#define __ASIN_CUDA_API_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(asin, nvidia) - -#endif // __ASIN_CUDA_API_H__ diff --git a/src/infiniop/ops/asin/operator.cc b/src/infiniop/ops/asin/operator.cc deleted file mode 100644 index e372286b3..000000000 --- a/src/infiniop/ops/asin/operator.cc +++ /dev/null @@ -1,169 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/asin.h" - -#ifdef ENABLE_CPU_API -#include "cpu/asin_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) -#include "nvidia/asin_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/asin_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/asin_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAsinDescriptor( - infiniopHandle_t handle, - infiniopAsinDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::asin::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output_desc, \ - {input_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAsinWorkspaceSize(infiniopAsinDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopAsin( - infiniopAsinDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, {input}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyAsinDescriptor(infiniopAsinDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/asinh/cpu/asinh_cpu.cc b/src/infiniop/ops/asinh/cpu/asinh_cpu.cc deleted file mode 100644 index 8b2a1047c..000000000 --- a/src/infiniop/ops/asinh/cpu/asinh_cpu.cc +++ /dev/null @@ -1,50 +0,0 @@ -#include "asinh_cpu.h" - -namespace op::asinh::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &x_desc = input_desc_vec.at(0); - const auto &y_shape = out_desc->shape(); - const auto &x_shape = x_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(y_shape, x_shape); - - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::asinh::cpu diff --git a/src/infiniop/ops/asinh/cpu/asinh_cpu.h b/src/infiniop/ops/asinh/cpu/asinh_cpu.h deleted file mode 100644 index 0a999b63b..000000000 --- a/src/infiniop/ops/asinh/cpu/asinh_cpu.h +++ /dev/null @@ -1,22 +0,0 @@ -#ifndef __ASINH_CPU_H__ -#define __ASINH_CPU_H__ - -#include - -#include "../../../elementwise/cpu/elementwise_cpu.h" - -ELEMENTWISE_DESCRIPTOR(asinh, cpu) - -namespace op::asinh::cpu { -typedef struct AsinhOp { -public: - static constexpr size_t num_inputs = 1; - - template - T operator()(const T &x) const { - return std::asinh(x); - } -} AsinhOp; -} // namespace op::asinh::cpu - -#endif // __ASINH_CPU_H__ diff --git a/src/infiniop/ops/asinh/cuda/kernel.cuh b/src/infiniop/ops/asinh/cuda/kernel.cuh deleted file mode 100644 index 70d51fb0b..000000000 --- a/src/infiniop/ops/asinh/cuda/kernel.cuh +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef __ASINH_CUDA_KERNEL_H__ -#define __ASINH_CUDA_KERNEL_H__ - -namespace op::asinh::cuda { - -typedef struct AsinhOp { -public: - static constexpr size_t num_inputs = 1; - template - __device__ __forceinline__ T operator()(const T &x) const { - - if constexpr (std::is_same_v) { - float x_f = __half2float(x); - return __float2half(asinhf(x_f)); - } else if constexpr (std::is_same_v) { - float x_f = __bfloat162float(x); - return __float2bfloat16(asinhf(x_f)); - } else if constexpr (std::is_same_v) { - return asinhf(x); - } else { - return ::asinh(x); - } - } - -} AsinhOp; - -} // namespace op::asinh::cuda - -#endif // __ASINH_CUDA_KERNEL_H__ diff --git a/src/infiniop/ops/asinh/metax/asinh.maca b/src/infiniop/ops/asinh/metax/asinh.maca deleted file mode 100644 index b178cfa8e..000000000 --- a/src/infiniop/ops/asinh/metax/asinh.maca +++ /dev/null @@ -1,58 +0,0 @@ -#include "../../../elementwise/metax/elementwise_metax.h" -#include "asinh_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::asinh::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &x_desc = input_desc_vec.at(0); - const auto &y_shape = out_desc->shape(); - const auto &x_shape = x_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(y_shape, x_shape); - - // create CUDA elementwise descriptor - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::AsinhOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::AsinhOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::AsinhOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::AsinhOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::asinh::metax diff --git a/src/infiniop/ops/asinh/metax/asinh_metax.h b/src/infiniop/ops/asinh/metax/asinh_metax.h deleted file mode 100644 index a6753d52a..000000000 --- a/src/infiniop/ops/asinh/metax/asinh_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ASINH_METAX_API_H__ -#define __ASINH_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(asinh, metax) - -#endif // __ASINH_METAX_API_H__ diff --git a/src/infiniop/ops/asinh/moore/asinh_moore.h b/src/infiniop/ops/asinh/moore/asinh_moore.h deleted file mode 100644 index 7b9372579..000000000 --- a/src/infiniop/ops/asinh/moore/asinh_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ASINH_MOORE_API_H__ -#define __ASINH_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(asinh, moore) - -#endif // __ASINH_MOORE_API_H__ diff --git a/src/infiniop/ops/asinh/moore/asinh_moore.mu b/src/infiniop/ops/asinh/moore/asinh_moore.mu deleted file mode 100644 index 6ee10c4b3..000000000 --- a/src/infiniop/ops/asinh/moore/asinh_moore.mu +++ /dev/null @@ -1,59 +0,0 @@ -#include "asinh_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "../cuda/kernel.cuh" - -namespace op::asinh::moore { -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &x_desc = input_desc_vec.at(0); - const auto &y_shape = out_desc->shape(); - const auto &x_shape = x_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(y_shape, x_shape); - - // create MOORE elementwise descriptor - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::AsinhOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::AsinhOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::AsinhOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::AsinhOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::asinh::moore diff --git a/src/infiniop/ops/asinh/nvidia/asinh_nvidia.cu b/src/infiniop/ops/asinh/nvidia/asinh_nvidia.cu deleted file mode 100644 index 3e9c2412d..000000000 --- a/src/infiniop/ops/asinh/nvidia/asinh_nvidia.cu +++ /dev/null @@ -1,56 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "asinh_nvidia.cuh" - -namespace op::asinh::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &x_desc = input_desc_vec.at(0); - const auto &y_shape = out_desc->shape(); - const auto &x_shape = x_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(y_shape, x_shape); - - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::AsinhOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::AsinhOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::AsinhOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::AsinhOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::asinh::nvidia diff --git a/src/infiniop/ops/asinh/nvidia/asinh_nvidia.cuh b/src/infiniop/ops/asinh/nvidia/asinh_nvidia.cuh deleted file mode 100644 index 4df770f7b..000000000 --- a/src/infiniop/ops/asinh/nvidia/asinh_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ASINH_NVIDIA_API_H__ -#define __ASINH_NVIDIA_API_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(asinh, nvidia) - -#endif // __ASINH_NVIDIA_API_H diff --git a/src/infiniop/ops/asinh/operator.cc b/src/infiniop/ops/asinh/operator.cc deleted file mode 100644 index 021ddb6b0..000000000 --- a/src/infiniop/ops/asinh/operator.cc +++ /dev/null @@ -1,153 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/asinh.h" - -#ifdef ENABLE_CPU_API -#include "cpu/asinh_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) -#include "nvidia/asinh_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/asinh_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/asinh_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAsinhDescriptor( - infiniopHandle_t handle, - infiniopAsinhDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::asinh::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - {x_desc}) - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAsinhWorkspaceSize(infiniopAsinhDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopAsinh(infiniopAsinhDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, {x}, stream); - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyAsinhDescriptor(infiniopAsinhDescriptor_t desc) { -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DESTROY(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - DESTROY(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DESTROY(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DESTROY(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DESTROY -} diff --git a/src/infiniop/ops/asum/asum.h b/src/infiniop/ops/asum/asum.h deleted file mode 100644 index dd7aef22e..000000000 --- a/src/infiniop/ops/asum/asum.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef __ASUM_H__ -#define __ASUM_H__ - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::asum::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AsumInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - AsumInfo info, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(std::move(info)), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t x_desc, \ - infiniopTensorDescriptor_t result_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - const void *x, \ - void *result, \ - void *stream) const; \ - }; \ - } - -#endif // __ASUM_H__ diff --git a/src/infiniop/ops/asum/bang/asum_bang.h b/src/infiniop/ops/asum/bang/asum_bang.h deleted file mode 100644 index bf388c744..000000000 --- a/src/infiniop/ops/asum/bang/asum_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ASUM_BANG_H__ -#define __ASUM_BANG_H__ - -#include "../asum.h" - -DESCRIPTOR(bang) - -#endif // __ASUM_BANG_H__ diff --git a/src/infiniop/ops/asum/bang/asum_bang.mlu b/src/infiniop/ops/asum/bang/asum_bang.mlu deleted file mode 100644 index 1c8cba0c3..000000000 --- a/src/infiniop/ops/asum/bang/asum_bang.mlu +++ /dev/null @@ -1,95 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "asum_bang.h" -#include "asum_bang_kernel.mlu" - -namespace op::asum::bang { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = AsumInfo::createAsumInfo(x_desc, result_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateAsum( - const AsumInfo &info, - const Tdata *x, - Tdata *result, - cnrtQueue_t queue) { - - const int n = utils::cast(info.n); - const int incx = utils::cast(info.incx); - - cnrtDim3_t k_dim; - cnrtFunctionType_t k_type; - - k_dim.x = 4; - k_dim.y = 1; - k_dim.z = 1; - k_type = cnrtFuncTypeUnion1; - - if (incx == 1) { - asumKernelContiguous<<>>( - n, - x, - result); - } else { - asumKernelStrided<<>>( - n, - x, - incx, - result); - } - - cnrtQueueSync(queue); - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_ASUM(TDATA) \ - calculateAsum(_info, \ - (const TDATA *)x, \ - (TDATA *)result, \ - (cnrtQueue_t)stream) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - switch (_info.data_type) { - case INFINI_DTYPE_F16: - return CALCULATE_ASUM(half); - case INFINI_DTYPE_F32: - return CALCULATE_ASUM(float); - case INFINI_DTYPE_BF16: - return CALCULATE_ASUM(bfloat16_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE_ASUM - -} // namespace op::asum::bang diff --git a/src/infiniop/ops/asum/bang/asum_bang_kernel.mlu b/src/infiniop/ops/asum/bang/asum_bang_kernel.mlu deleted file mode 100644 index 633ad767d..000000000 --- a/src/infiniop/ops/asum/bang/asum_bang_kernel.mlu +++ /dev/null @@ -1,155 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "asum_bang.h" - -#include - -__nram__ char nram_buffer[NRAM_MAX_SIZE]; - -template -__mlu_device__ void asumToCompute(float *dst, const Tdata *src, int size) { - if constexpr (std::is_same_v) { - __bang_half2float(dst, src, size); - } else if constexpr (std::is_same_v) { - __bang_bfloat162float(dst, src, size); - } else { - __memcpy(dst, src, size * sizeof(float), NRAM2NRAM); - } -} - -template -__mlu_device__ float asumToCompute(Tdata value) { - if constexpr (std::is_same_v) { - return __half2float(value); - } else if constexpr (std::is_same_v) { - return __bfloat162float(value); - } else { - return static_cast(value); - } -} - -template -__mlu_device__ void asumStoreResult(Tdata *result, Tdata *nram_result, float *nram_compute, float value) { - nram_compute[0] = value; - if constexpr (std::is_same_v) { - __bang_float2half(nram_result, nram_compute, 1); - result[0] = nram_result[0]; - } else if constexpr (std::is_same_v) { - __bang_float2bfloat16(nram_result, nram_compute, 1); - result[0] = nram_result[0]; - } else { - result[0] = nram_compute[0]; - } -} - -template -__mlu_global__ void asumKernelContiguous( - int n, - const Tdata *x, - Tdata *result) { - - __mlu_shared__ float shared_partial_sum[4]; - - char *nram_aligned = (char *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1)); - - size_t nram_usable = NRAM_MAX_SIZE - (nram_aligned - nram_buffer); - size_t max_chunk_elements = nram_usable / (sizeof(Tdata) + sizeof(float)); - - size_t align_elements = ALIGN_SIZE / sizeof(Tdata); - if (align_elements == 0) { - align_elements = 1; - } - int chunk_size = (int)((max_chunk_elements / align_elements) * align_elements); - - Tdata *nram_x = (Tdata *)nram_aligned; - float *nram_compute = (float *)(nram_x + chunk_size); - - int elements_per_core = n / taskDim; - int remain = n % taskDim; - int core_elements = elements_per_core + (taskId < remain ? 1 : 0); - int core_offset = taskId < remain ? taskId * core_elements : taskId * elements_per_core + remain; - - int chunks = core_elements / chunk_size; - int chunk_rem = core_elements % chunk_size; - - float partial_sum = 0.0f; - - for (int c = 0; c < chunks; c++) { - int current_offset = core_offset + c * chunk_size; - - __memcpy(nram_x, x + current_offset, chunk_size * sizeof(Tdata), GDRAM2NRAM); - - asumToCompute(nram_compute, nram_x, chunk_size); - __bang_abs(nram_compute, nram_compute, chunk_size); - - partial_sum += __bang_sum(nram_compute, chunk_size); - } - - if (chunk_rem > 0) { - int current_offset = core_offset + chunks * chunk_size; - - __memcpy(nram_x, x + current_offset, chunk_rem * sizeof(Tdata), GDRAM2NRAM); - - asumToCompute(nram_compute, nram_x, chunk_rem); - __bang_abs(nram_compute, nram_compute, chunk_rem); - - partial_sum += __bang_sum(nram_compute, chunk_rem); - } - - shared_partial_sum[coreId] = partial_sum; - - __sync_cluster(); - - if (coreId == 0) { - float cluster_sum = 0.0f; - - for (int i = 0; i < coreDim; i++) { - cluster_sum += shared_partial_sum[i]; - } - - asumStoreResult(result, nram_x, nram_compute, cluster_sum); - } -} - -template -__mlu_global__ void asumKernelStrided( - int n, - const Tdata *x, - int incx, - Tdata *result) { - - __mlu_shared__ float shared_partial_sum[4]; - - char *nram_aligned = (char *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1)); - - float *nram_compute = (float *)nram_aligned; - Tdata *nram_result = (Tdata *)(nram_compute + 1); - - int elements_per_core = n / taskDim; - int remain = n % taskDim; - int actual_tasks = elements_per_core + (taskId < remain ? 1 : 0); - int start_idx = taskId < remain ? taskId * actual_tasks : taskId * elements_per_core + remain; - - float partial_sum = 0.0f; - - for (int i = start_idx; i < start_idx + actual_tasks; ++i) { - int offset = i * incx; - float x_val = asumToCompute(x[offset]); - float abs_val = x_val > 0.0f ? x_val : -x_val; - - partial_sum += abs_val; - } - - shared_partial_sum[coreId] = partial_sum; - - __sync_cluster(); - - if (coreId == 0) { - float cluster_sum = 0.0f; - - for (int i = 0; i < coreDim; i++) { - cluster_sum += shared_partial_sum[i]; - } - - asumStoreResult(result, nram_result, nram_compute, cluster_sum); - } -} diff --git a/src/infiniop/ops/asum/cpu/asum_cpu.cc b/src/infiniop/ops/asum/cpu/asum_cpu.cc deleted file mode 100644 index c1178a74c..000000000 --- a/src/infiniop/ops/asum/cpu/asum_cpu.cc +++ /dev/null @@ -1,91 +0,0 @@ -#include "asum_cpu.h" -#include "../../../devices/cpu/common_cpu.h" - -namespace op::asum::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = AsumInfo::createAsumInfo(x_desc, result_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateAsum( - const AsumInfo &info, - const Tdata *x, - Tdata *result) { - - const size_t n = info.n; - const ptrdiff_t incx = info.incx; - - if constexpr (std::is_same::value || std::is_same::value) { - float total_sum = 0.0; - - for (size_t i = 0; i < n; ++i) { - const ptrdiff_t idx = utils::cast(i) * incx; - total_sum += std::abs(utils::cast(x[idx])); - } - - result[0] = utils::cast(total_sum); - } else { - Tdata total_sum = 0.0; - - for (size_t i = 0; i < n; ++i) { - const ptrdiff_t idx = utils::cast(i) * incx; - total_sum += std::abs(x[idx]); - } - - result[0] = total_sum; - } - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_ASUM(TDATA) \ - calculateAsum(_info, \ - (const TDATA *)x, \ - (TDATA *)result) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - switch (_info.data_type) { - case INFINI_DTYPE_F16: - return CALCULATE_ASUM(fp16_t); - case INFINI_DTYPE_BF16: - return CALCULATE_ASUM(bf16_t); - case INFINI_DTYPE_F32: - return CALCULATE_ASUM(float); - case INFINI_DTYPE_F64: - return CALCULATE_ASUM(double); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE_ASUM - -} // namespace op::asum::cpu diff --git a/src/infiniop/ops/asum/cpu/asum_cpu.h b/src/infiniop/ops/asum/cpu/asum_cpu.h deleted file mode 100644 index 84b7572d8..000000000 --- a/src/infiniop/ops/asum/cpu/asum_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ASUM_CPU_H__ -#define __ASUM_CPU_H__ - -#include "../asum.h" - -DESCRIPTOR(cpu) - -#endif // __ASUM_CPU_H__ diff --git a/src/infiniop/ops/asum/info.h b/src/infiniop/ops/asum/info.h deleted file mode 100644 index 3efcc0e39..000000000 --- a/src/infiniop/ops/asum/info.h +++ /dev/null @@ -1,41 +0,0 @@ -#ifndef __ASUM_INFO_H__ -#define __ASUM_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" - -class AsumInfo { -private: - AsumInfo() = default; - -public: - size_t n; - ptrdiff_t incx; - infiniDtype_t data_type; - - static utils::Result createAsumInfo( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - - CHECK_OR_RETURN(x_desc != nullptr, INFINI_STATUS_NULL_POINTER); - CHECK_OR_RETURN(result_desc != nullptr, INFINI_STATUS_NULL_POINTER); - - auto data_type = x_desc->dtype(); - - CHECK_OR_RETURN(result_desc->dtype() == data_type, INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_DTYPE(data_type, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_OR_RETURN(x_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(result_desc->numel() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - - auto n = x_desc->numel(); - auto incx = x_desc->stride(0); - - return utils::Result(AsumInfo{ - n, - incx, - data_type}); - } -}; - -#endif // __ASUM_INFO_H__ diff --git a/src/infiniop/ops/asum/metax/asum_metax.cc b/src/infiniop/ops/asum/metax/asum_metax.cc deleted file mode 100644 index 4a084fbf9..000000000 --- a/src/infiniop/ops/asum/metax/asum_metax.cc +++ /dev/null @@ -1,71 +0,0 @@ -#include "asum_metax.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" - -namespace op::asum::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = AsumInfo::createAsumInfo(x_desc, result_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - new Opaque{handle->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - const int n = utils::cast(_info.n); - const int incx = utils::cast(_info.incx); - const infiniDtype_t data_type = _info.data_type; - - CHECK_STATUS(_opaque->internal->useMcblas( - (hcStream_t)stream, - [&](hcblasHandle_t handle) { - CHECK_MCBLAS(hcblasSetPointerMode(handle, HCBLAS_POINTER_MODE_DEVICE)); - - switch (data_type) { - case INFINI_DTYPE_F32: - CHECK_MCBLAS(hcblasSasum(handle, n, (const float *)x, incx, (float *)result)); - break; - case INFINI_DTYPE_F64: - CHECK_MCBLAS(hcblasDasum(handle, n, (const double *)x, incx, (double *)result)); - break; - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - - return INFINI_STATUS_SUCCESS; - })); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::asum::metax diff --git a/src/infiniop/ops/asum/metax/asum_metax.h b/src/infiniop/ops/asum/metax/asum_metax.h deleted file mode 100644 index f5ade8e58..000000000 --- a/src/infiniop/ops/asum/metax/asum_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ASUM_METAX_H__ -#define __ASUM_METAX_H__ - -#include "../asum.h" - -DESCRIPTOR(metax) - -#endif // __ASUM_METAX_H__ diff --git a/src/infiniop/ops/asum/operator.cc b/src/infiniop/ops/asum/operator.cc deleted file mode 100644 index 51769f0d3..000000000 --- a/src/infiniop/ops/asum/operator.cc +++ /dev/null @@ -1,124 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/asum.h" - -#ifdef ENABLE_CPU_API -#include "cpu/asum_cpu.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/asum_metax.h" -#endif -#ifdef ENABLE_CAMBRICON_API -#include "bang/asum_bang.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAsumDescriptor( - infiniopHandle_t handle, - infiniopAsumDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::asum::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - x_desc, \ - result_desc) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAsumWorkspaceSize(infiniopAsumDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - GET(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopAsum( - infiniopAsumDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, x, result, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyAsumDescriptor(infiniopAsumDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - DELETE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/atanh/cpu/atanh_cpu.cc b/src/infiniop/ops/atanh/cpu/atanh_cpu.cc deleted file mode 100644 index 3aeecfe29..000000000 --- a/src/infiniop/ops/atanh/cpu/atanh_cpu.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "atanh_cpu.h" - -namespace op::atanh::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &y_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(y_shape, a_shape); - - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - // 分发到对应的数据类型进行计算，模板参数为我们在 atanh_cpu.h 中定义的 AtanhOp - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::atanh::cpu diff --git a/src/infiniop/ops/atanh/cpu/atanh_cpu.h b/src/infiniop/ops/atanh/cpu/atanh_cpu.h deleted file mode 100644 index 543d82d57..000000000 --- a/src/infiniop/ops/atanh/cpu/atanh_cpu.h +++ /dev/null @@ -1,31 +0,0 @@ -#ifndef __ATANH_CPU_H__ -#define __ATANH_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" -#include -#include - -// 注册 atanh 算子在 cpu 后端的 descriptor -ELEMENTWISE_DESCRIPTOR(atanh, cpu) - -namespace op::atanh::cpu { -typedef struct AtanhOp { -public: - // atanh 是一元算子 - static constexpr size_t num_inputs = 1; - - template - T operator()(const T &a) const { - // 对于 float, double 等原生支持的类型直接调用 std::atanh - if constexpr (std::is_floating_point_v) { - return std::atanh(a); - } else { - // 对于 half, bfloat16 等自定义类型，先转为 float 计算再转回 - // 假设这些类型支持 static_cast 到 float - return static_cast(std::atanhf(static_cast(a))); - } - } -} AtanhOp; -} // namespace op::atanh::cpu - -#endif // __ATANH_CPU_H__ diff --git a/src/infiniop/ops/atanh/cuda/kernel.cuh b/src/infiniop/ops/atanh/cuda/kernel.cuh deleted file mode 100644 index f984c0686..000000000 --- a/src/infiniop/ops/atanh/cuda/kernel.cuh +++ /dev/null @@ -1,37 +0,0 @@ -#ifndef __ATANH_CUDA_H__ -#define __ATANH_CUDA_H__ - -namespace op::atanh::cuda { -typedef struct AtanhOp { -public: - // atanh 是一元算子，只需要一个输入 - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(const T &a) const { - if constexpr (std::is_same_v) { - // 对 half2 的两个部分分别求 atanh - float2 f = __half22float2(a); - f.x = atanhf(f.x); - f.y = atanhf(f.y); - return __float22half2_rn(f); - } else if constexpr (std::is_same_v) { - // half 类型先转为 float 计算再转回 - return __float2half(atanhf(__half2float(a))); - } else if constexpr (std::is_same_v) { - // bfloat16 类型处理同上 - return __float2bfloat16(atanhf(__bfloat162float(a))); - } else if constexpr (std::is_same_v) { - // float 直接调用标准数学库函数 - return atanhf(a); - } else if constexpr (std::is_same_v) { - return ::atanh(a); - } else { - // 其他整数类型或不支持类型理论上不应进入，此处做简单 fallback - return static_cast(atanhf(static_cast(a))); - } - } -} AtanhOp; -} // namespace op::atanh::cuda - -#endif // __ATANH_CUDA_H__ diff --git a/src/infiniop/ops/atanh/metax/atanh_metax.h b/src/infiniop/ops/atanh/metax/atanh_metax.h deleted file mode 100644 index cc9b1e0de..000000000 --- a/src/infiniop/ops/atanh/metax/atanh_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ATANH_METAX_API_H__ -#define __ATANH_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(atanh, metax) - -#endif // __ATANH_METAX_API_H__ diff --git a/src/infiniop/ops/atanh/metax/atanh_metax.maca b/src/infiniop/ops/atanh/metax/atanh_metax.maca deleted file mode 100644 index 24428b594..000000000 --- a/src/infiniop/ops/atanh/metax/atanh_metax.maca +++ /dev/null @@ -1,58 +0,0 @@ -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "atanh_metax.h" -#include "atanh_metax_kernel.h" - -namespace op::atanh::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &y_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(y_shape, a_shape); - - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, metax::AtanhOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, metax::AtanhOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, metax::AtanhOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, metax::AtanhOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::atanh::metax diff --git a/src/infiniop/ops/atanh/metax/atanh_metax_kernel.h b/src/infiniop/ops/atanh/metax/atanh_metax_kernel.h deleted file mode 100644 index 83deb1237..000000000 --- a/src/infiniop/ops/atanh/metax/atanh_metax_kernel.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef __ATANH_METAX_KERNEL_H__ -#define __ATANH_METAX_KERNEL_H__ - -/* - * This file contains the Atanh operation implementation for the MUSA backend. - * - * It follows the consistent code structure to ensure alignment across different - * hardware platforms within the Moore Threads (MUSA) ecosystem. - */ -namespace op::atanh::metax { - -typedef struct AtanhOp { -public: - // 一元算子，输入数量为 1 - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(const T &a) const { - if constexpr (std::is_same_v) { - // 针对 half2 进行并行计算 - float2 f2 = __half22float2(a); - f2.x = atanhf(f2.x); - f2.y = atanhf(f2.y); - return __float22half2_rn(f2); - } else if constexpr (std::is_same_v) { - // 转为 float 计算以保证精度并匹配 MUSA 数学库 - return __float2half(atanhf(__half2float(a))); - } else if constexpr (std::is_same_v) { - // BF16 同样提升到 float 计算，避免转换歧义 - float a_f = __bfloat162float(a); - return __float2bfloat16_rn(atanhf(a_f)); - } else if constexpr (std::is_same_v) { - // 调用 MUSA 内置的单精度反双曲正切函数 - return atanhf(a); - } else if constexpr (std::is_same_v) { - // 调用双精度版本 - return ::atanh(a); - } else { - // 兜底实现（如果是整数类型，通常会隐式转为 float） - return static_cast(atanhf(static_cast(a))); - } - } -} AtanhOp; - -} // namespace op::atanh::metax - -#endif // __ATANH_METAX_KERNEL_H__ diff --git a/src/infiniop/ops/atanh/moore/atanh_moore.h b/src/infiniop/ops/atanh/moore/atanh_moore.h deleted file mode 100644 index fb256bbc1..000000000 --- a/src/infiniop/ops/atanh/moore/atanh_moore.h +++ /dev/null @@ -1,11 +0,0 @@ -#ifndef __ATANH_MOORE_API_H__ -#define __ATANH_MOORE_API_H__ - -// 1. 修改包含路径，指向 moore 平台的 elementwise API 定义 -#include "../../../elementwise/moore/elementwise_moore_api.h" - -// 2. 使用 ELEMENTWISE_DESCRIPTOR 宏，平台参数改为 moore -// 这将自动生成 op::atanh::moore::Descriptor 类的声明 -ELEMENTWISE_DESCRIPTOR(atanh, moore) - -#endif // __ATANH_MOORE_API_H__ diff --git a/src/infiniop/ops/atanh/moore/atanh_moore.mu b/src/infiniop/ops/atanh/moore/atanh_moore.mu deleted file mode 100644 index bcb1bd2df..000000000 --- a/src/infiniop/ops/atanh/moore/atanh_moore.mu +++ /dev/null @@ -1,64 +0,0 @@ -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "atanh_moore.h" -#include "atanh_moore_kernel.h" - -namespace op::atanh::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - // 1. 转换 Handle 为 Moore 平台类型 - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &y_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - - // 2. 检查数据类型支持情况 - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); - - // 3. 校验 Shape 一致性 - CHECK_SAME_SHAPE(y_shape, a_shape); - - // 4. 创建 Moore 平台的 Elementwise 描述符 - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // 5. 根据数据类型分发到具体的 MUSA Kernel 逻辑 - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, moore::AtanhOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - // 注意：这里将 nv_bfloat16 替换为 Moore 环境下的 bfloat16 类型名 - return _device_info->calculate<256, moore::AtanhOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, moore::AtanhOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, moore::AtanhOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::atanh::moore diff --git a/src/infiniop/ops/atanh/moore/atanh_moore_kernel.h b/src/infiniop/ops/atanh/moore/atanh_moore_kernel.h deleted file mode 100644 index 490a7ffe2..000000000 --- a/src/infiniop/ops/atanh/moore/atanh_moore_kernel.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef __ATANH_MOORE_KERNEL_H__ -#define __ATANH_MOORE_KERNEL_H__ - -/* - * This file contains the Atanh operation implementation for the MUSA backend. - * - * It follows the consistent code structure to ensure alignment across different - * hardware platforms within the Moore Threads (MUSA) ecosystem. - */ -namespace op::atanh::moore { - -typedef struct AtanhOp { -public: - // 一元算子，输入数量为 1 - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(const T &a) const { - if constexpr (std::is_same_v) { - // 针对 half2 进行并行计算 - float2 f2 = __half22float2(a); - f2.x = atanhf(f2.x); - f2.y = atanhf(f2.y); - return __float22half2_rn(f2); - } else if constexpr (std::is_same_v) { - // 转为 float 计算以保证精度并匹配 MUSA 数学库 - return __float2half(atanhf(__half2float(a))); - } else if constexpr (std::is_same_v) { - // BF16 同样提升到 float 计算，避免转换歧义 - float a_f = __bfloat162float(a); - return __float2bfloat16_rn(atanhf(a_f)); - } else if constexpr (std::is_same_v) { - // 调用 MUSA 内置的单精度反双曲正切函数 - return atanhf(a); - } else if constexpr (std::is_same_v) { - // 调用双精度版本 - return ::atanh(a); - } else { - // 兜底实现（如果是整数类型，通常会隐式转为 float） - return static_cast(atanhf(static_cast(a))); - } - } -} AtanhOp; - -} // namespace op::atanh::moore - -#endif // __ATANH_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/atanh/nvidia/atanh_nvidia.cu b/src/infiniop/ops/atanh/nvidia/atanh_nvidia.cu deleted file mode 100644 index 1143cd589..000000000 --- a/src/infiniop/ops/atanh/nvidia/atanh_nvidia.cu +++ /dev/null @@ -1,58 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "atanh_nvidia.cuh" - -namespace op::atanh::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &y_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(y_shape, a_shape); - - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::AtanhOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::AtanhOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::AtanhOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::AtanhOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::atanh::nvidia diff --git a/src/infiniop/ops/atanh/nvidia/atanh_nvidia.cuh b/src/infiniop/ops/atanh/nvidia/atanh_nvidia.cuh deleted file mode 100644 index d954a390d..000000000 --- a/src/infiniop/ops/atanh/nvidia/atanh_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ATANH_CUDA_API_H__ -#define __ATANH_CUDA_API_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(atanh, nvidia) - -#endif // __ATANH_CUDA_API_H__ diff --git a/src/infiniop/ops/atanh/operator.cc b/src/infiniop/ops/atanh/operator.cc deleted file mode 100644 index 0472d6d43..000000000 --- a/src/infiniop/ops/atanh/operator.cc +++ /dev/null @@ -1,179 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/atanh.h" - -#ifdef ENABLE_CPU_API -#include "cpu/atanh_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/atanh_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/atanh_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/atanh_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAtanhDescriptor( - infiniopHandle_t handle, - infiniopAtanhDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::atanh::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - {a_desc}) // 一元算子只传入 a_desc - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAtanhWorkspaceSize(infiniopAtanhDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopAtanh( - infiniopAtanhDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *a, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, {a}, stream) // 一元算子只传入 a - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyAtanhDescriptor(infiniopAtanhDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/attention/attention.h b/src/infiniop/ops/attention/attention.h deleted file mode 100644 index d4740d25e..000000000 --- a/src/infiniop/ops/attention/attention.h +++ /dev/null @@ -1,37 +0,0 @@ -#ifndef ATTENTION_H -#define ATTENTION_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::attention::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t y_desc, \ - infiniopTensorDescriptor_t x_desc); \ - }; \ - } - -#endif // ATTENTION_H diff --git a/src/infiniop/ops/attention/operator.cc b/src/infiniop/ops/attention/operator.cc deleted file mode 100644 index 14279591b..000000000 --- a/src/infiniop/ops/attention/operator.cc +++ /dev/null @@ -1,291 +0,0 @@ -#include "../../operator.h" -#include "../../../utils.h" -#include "../../../utils/check.h" -#include "../../handle.h" -#include "../../tensor.h" -#include "infiniop/ops/attention.h" -#include "infiniop/ops/causal_softmax.h" -#include "infiniop/ops/gemm.h" -#include "infiniop/ops/rearrange.h" - -#include -#include - -struct InfiniopAttentionDescriptor { - InfiniopDescriptor _super; - infiniopRearrangeDescriptor_t rearrange_desc_k; - infiniopRearrangeDescriptor_t rearrange_desc_v; - infiniopRearrangeDescriptor_t rearrange_desc_q; - infiniopRearrangeDescriptor_t rearrange_desc_out; - infiniopGemmDescriptor_t matmul_desc1; - infiniopGemmDescriptor_t matmul_desc2; - infiniopCausalSoftmaxDescriptor_t softmax_desc; - size_t workspace_size; - size_t op_workspace_offset; - size_t op_workspace_size; - size_t q_cont_offset; - size_t att_score_offset; - size_t att_val_offset; - size_t k_cache_offset; - size_t v_cache_offset; - float qk_alpha; -}; - -__INFINI_C __export infiniStatus_t infiniopCreateAttentionDescriptor(infiniopHandle_t handle, - infiniopAttentionDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t k_cache_desc, - infiniopTensorDescriptor_t v_cache_desc, - size_t pos) { - if (out_desc->ndim() != 3 || q_desc->ndim() != 3 || k_desc->ndim() != 3 || v_desc->ndim() != 3 || k_cache_desc->ndim() != 3 || v_cache_desc->ndim() != 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (!out_desc->isContiguous()) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - if (q_desc->strides()[2] != 1 || k_desc->strides()[2] != 1 || v_desc->strides()[2] != 1 || k_cache_desc->strides()[2] != 1 || v_cache_desc->strides()[2] != 1) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - size_t n_q_head = q_desc->shape()[0]; - size_t seq_len = q_desc->shape()[1]; - size_t head_dim = q_desc->shape()[2]; - size_t hidden_size = n_q_head * head_dim; - size_t n_kv_head = k_desc->shape()[0]; - size_t total_seq_len = seq_len + pos; - size_t n_group = n_q_head / n_kv_head; - size_t alignment = 256; - - if (out_desc->shape()[0] != seq_len || out_desc->shape()[1] != n_q_head || out_desc->shape()[2] != head_dim) { - return INFINI_STATUS_BAD_PARAM; - } - - // k: [n_kv_head, seq_len, head_dim] - if (k_desc->shape()[0] != n_kv_head || k_desc->shape()[1] != seq_len || k_desc->shape()[2] != head_dim) { - return INFINI_STATUS_BAD_PARAM; - } - - // v: [n_kv_head, seq_len, head_dim] - if (v_desc->shape()[0] != n_kv_head || v_desc->shape()[1] != seq_len || v_desc->shape()[2] != head_dim) { - return INFINI_STATUS_BAD_PARAM; - } - - // k_cache: [n_kv_head, _, head_dim] - if (k_cache_desc->shape()[0] != n_kv_head || k_cache_desc->shape()[1] < total_seq_len || k_cache_desc->shape()[2] != head_dim) { - return INFINI_STATUS_BAD_PARAM; - } - - // v_cache: [n_kv_head, _, head_dim] - if (v_cache_desc->shape()[0] != n_kv_head || v_cache_desc->shape()[1] < total_seq_len || v_cache_desc->shape()[2] != head_dim) { - return INFINI_STATUS_BAD_PARAM; - } - - // Rearrange k into k_cache - infiniopTensorDescriptor_t dst_k_desc; - CHECK_STATUS(infiniopCreateTensorDescriptor(&dst_k_desc, 3, k_desc->shape().data(), k_cache_desc->strides().data(), k_cache_desc->dtype())); - infiniopRearrangeDescriptor_t rearrange_desc_k; - CHECK_STATUS(infiniopCreateRearrangeDescriptor(handle, &rearrange_desc_k, dst_k_desc, k_desc)); - - // Rearrange v into v_cache - infiniopTensorDescriptor_t dst_v_desc; - CHECK_STATUS(infiniopCreateTensorDescriptor(&dst_v_desc, 3, v_desc->shape().data(), v_cache_desc->strides().data(), v_cache_desc->dtype())); - infiniopRearrangeDescriptor_t rearrange_desc_v; - CHECK_STATUS(infiniopCreateRearrangeDescriptor(handle, &rearrange_desc_v, dst_v_desc, v_desc)); - - infiniopRearrangeDescriptor_t rearrange_desc_q = nullptr; - size_t q_cont_size = 0; - infiniopTensorDescriptor_t rearranged_q_desc; - // Rearrange q into contiguous - if (!q_desc->isContiguous(0, 1)) { - CHECK_STATUS(infiniopCreateTensorDescriptor(&rearranged_q_desc, 3, q_desc->shape().data(), nullptr, q_desc->dtype())); - q_cont_size = utils::align(rearranged_q_desc->numel() * infiniSizeOf(rearranged_q_desc->dtype()), alignment); - rearrange_desc_q = new InfiniopDescriptor; - CHECK_STATUS(infiniopCreateRearrangeDescriptor(handle, &rearrange_desc_q, rearranged_q_desc, q_desc)); - } - - // Matmul1: q * full_k - // q: [n_q_head, seq_len, head_dim] -> [n_kv_head, n_group *seq_len, head_dim] - infiniopTensorDescriptor_t reshaped_q_desc; - CHECK_STATUS(infiniopCreateTensorDescriptor(&reshaped_q_desc, 3, q_desc->shape().data(), nullptr, q_desc->dtype())); - TRANSFORM_TENSOR_DESC(reshaped_q_desc, dimSplit(0, {n_kv_head, n_group})); - TRANSFORM_TENSOR_DESC(reshaped_q_desc, dimMerge(1, 2)); - // full_k: [n_kv_head, head_dim, total_seq_len] - infiniopTensorDescriptor_t full_k_desc; - size_t full_k_shape[3] = {n_kv_head, total_seq_len, head_dim}; - CHECK_STATUS(infiniopCreateTensorDescriptor(&full_k_desc, 3, full_k_shape, k_cache_desc->strides().data(), k_cache_desc->dtype())); - TRANSFORM_TENSOR_DESC(full_k_desc, dimPermute({0, 2, 1})); - // qk: [n_kv_head, n_group * seq_len, total_seq_len] - infiniopTensorDescriptor_t qk_desc; - size_t qk_shape[3] = {n_kv_head, n_group * seq_len, total_seq_len}; - CHECK_STATUS(infiniopCreateTensorDescriptor(&qk_desc, 3, qk_shape, nullptr, q_desc->dtype())); - // matmul1_desc - // qk_alpha - float qk_alpha = 1 / sqrt(head_dim); - infiniopGemmDescriptor_t matmul1_desc; - CHECK_STATUS(infiniopCreateGemmDescriptor(handle, &matmul1_desc, qk_desc, reshaped_q_desc, full_k_desc)); - // matmul1 workspace size - size_t matmul1_workspace_size; - CHECK_STATUS(infiniopGetGemmWorkspaceSize(matmul1_desc, &matmul1_workspace_size)); - matmul1_workspace_size = utils::align(matmul1_workspace_size, alignment); - // attention score tensor size - size_t attn_score_size = utils::align(qk_desc->numel() * infiniSizeOf(qk_desc->dtype()), alignment); - - // CausalSoftmax: softmax(qk) - // qk: [n_kv_head, n_group * seq_len, total_seq_len] -> [n_q_head, seq_len, total_seq_len] - TRANSFORM_TENSOR_DESC(qk_desc, dimSplit(1, {n_group, seq_len})); - TRANSFORM_TENSOR_DESC(qk_desc, dimMerge(0, 1)); - infiniopCausalSoftmaxDescriptor_t softmax_desc; - CHECK_STATUS(infiniopCreateCausalSoftmaxDescriptor(handle, &softmax_desc, qk_desc, qk_desc)); - // softmax workspace size - size_t softmax_workspace_size; - CHECK_STATUS(infiniopGetCausalSoftmaxWorkspaceSize(softmax_desc, &softmax_workspace_size)); - softmax_workspace_size = utils::align(softmax_workspace_size, alignment); - - // Matmul2: softmax(qk) * full_v - // softmax(qk): [n_q_head, seq_len, total_seq_len] -> [n_kv_head, n_group * seq_len, total_seq_len] - // full_v: [n_kv_head, total_seq_len, head_dim] - TRANSFORM_TENSOR_DESC(qk_desc, dimSplit(0, {n_kv_head, n_group})); - TRANSFORM_TENSOR_DESC(qk_desc, dimMerge(1, 2)); - infiniopTensorDescriptor_t full_v_desc; - size_t full_v_shape[3] = {n_kv_head, total_seq_len, head_dim}; - CHECK_STATUS(infiniopCreateTensorDescriptor(&full_v_desc, 3, full_v_shape, v_cache_desc->strides().data(), v_cache_desc->dtype())); - // temp_out: [n_kv_head, n_group * seq_len, head_dim] - infiniopTensorDescriptor_t att_val_desc; - size_t temp_out_shape[3] = {n_kv_head, n_group * seq_len, head_dim}; - CHECK_STATUS(infiniopCreateTensorDescriptor(&att_val_desc, 3, temp_out_shape, nullptr, q_desc->dtype())); - // matmul2_desc - infiniopGemmDescriptor_t matmul2_desc; - CHECK_STATUS(infiniopCreateGemmDescriptor(handle, &matmul2_desc, att_val_desc, qk_desc, full_v_desc)); - // matmul2 workspace size - size_t matmul2_workspace_size; - CHECK_STATUS(infiniopGetGemmWorkspaceSize(matmul2_desc, &matmul2_workspace_size)); - matmul2_workspace_size = utils::align(matmul2_workspace_size, alignment); - // attention value tensor size - size_t att_val_size = utils::align(att_val_desc->numel() * infiniSizeOf(att_val_desc->dtype()), alignment); - - // Rearrange temp_out into out - // out: [seq_len, n_q_head, head_dim] - // temp_out: [n_kv_head, n_group * seq_len, head_dim] -> [n_q_head, seq_len, head_dim] -> [seq_len, n_q_head, head_dim] - TRANSFORM_TENSOR_DESC(att_val_desc, dimSplit(1, {n_group, seq_len})); - TRANSFORM_TENSOR_DESC(att_val_desc, dimMerge(0, 1)); - TRANSFORM_TENSOR_DESC(att_val_desc, dimPermute({1, 0, 2})); - infiniopRearrangeDescriptor_t rearrange_desc_out; - CHECK_STATUS(infiniopCreateRearrangeDescriptor(handle, &rearrange_desc_out, out_desc, att_val_desc)); - - // workspace size - size_t op_workspace_size = utils::align(std::max(std::max(matmul1_workspace_size, matmul2_workspace_size), softmax_workspace_size), alignment); - size_t temp_tensors_size = attn_score_size + std::max(q_cont_size, att_val_size); - size_t workspace_size = temp_tensors_size + op_workspace_size; - - // k_cache_offset - size_t k_cache_offset = 0; - if (pos > 0) { - k_cache_offset = pos * k_cache_desc->getByteStrides()[1]; - } - - // v_cache_offset - size_t v_cache_offset = 0; - if (pos > 0) { - v_cache_offset = pos * v_cache_desc->getByteStrides()[1]; - } - - // create attention descriptor - *(InfiniopAttentionDescriptor **)desc_ptr = new InfiniopAttentionDescriptor{ - {handle->device, handle->device_id}, - rearrange_desc_k, - rearrange_desc_v, - rearrange_desc_q, - rearrange_desc_out, - matmul1_desc, - matmul2_desc, - softmax_desc, - workspace_size, - temp_tensors_size, - op_workspace_size, - attn_score_size, - 0, - attn_score_size, - k_cache_offset, - v_cache_offset, - 1.f / std::sqrt(float(head_dim)), - }; - - return INFINI_STATUS_SUCCESS; -} - -__INFINI_C __export infiniStatus_t infiniopGetAttentionWorkspaceSize(infiniopAttentionDescriptor_t desc, size_t *size) { - *size = ((InfiniopAttentionDescriptor *)desc)->workspace_size; - return INFINI_STATUS_SUCCESS; -} - -__INFINI_C __export infiniStatus_t infiniopAttention(infiniopAttentionDescriptor_t desc_, - void *workspace_, - size_t workspace_size_, - void *out, - void const *q, - void const *k, - void const *v, - void *k_cache, - void *v_cache, - void *stream) { - auto desc = (InfiniopAttentionDescriptor *)desc_; - if (workspace_size_ < desc->workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; // STATUS_MEMORY_NOT_ALLOCATED - } - void *workspace = (char *)workspace_ + desc->op_workspace_offset; - size_t workspace_size = desc->op_workspace_size; - void *att_score = (char *)workspace_ + desc->att_score_offset; - void *att_val = (char *)workspace_ + desc->att_val_offset; - void const *q_ = q; - // concat k and v to k_cache and v_cache - CHECK_STATUS(infiniopRearrange(desc->rearrange_desc_k, - (char *)k_cache + desc->k_cache_offset, k, stream)); - - CHECK_STATUS(infiniopRearrange(desc->rearrange_desc_v, - (char *)v_cache + desc->v_cache_offset, v, stream)); - - // rearrange q into contiguous - if (desc->rearrange_desc_q) { - void *q_cont = (char *)workspace_ + desc->q_cont_offset; - CHECK_STATUS(infiniopRearrange(desc->rearrange_desc_q, q_cont, q, stream)); - q_ = q_cont; - } - - // matmul1: q * full_k - CHECK_STATUS(infiniopGemm(desc->matmul_desc1, - workspace, workspace_size, - att_score, q_, k_cache, desc->qk_alpha, 0.0, stream)); - // softmax(qk) - CHECK_STATUS(infiniopCausalSoftmax(desc->softmax_desc, - workspace, workspace_size, - att_score, att_score, stream)); - // matmul2: softmax(qk) * full_v - CHECK_STATUS(infiniopGemm(desc->matmul_desc2, - workspace, workspace_size, - att_val, att_score, v_cache, 1.0, 0.0, stream)); - // rearrange out - CHECK_STATUS(infiniopRearrange(desc->rearrange_desc_out, out, att_val, stream)); - - return INFINI_STATUS_SUCCESS; -} - -__INFINI_C __export infiniStatus_t infiniopDestroyAttentionDescriptor(infiniopAttentionDescriptor_t desc_) { - auto desc = (InfiniopAttentionDescriptor *)desc_; - if (desc->rearrange_desc_q) { - CHECK_STATUS(infiniopDestroyRearrangeDescriptor(desc->rearrange_desc_q)); - } - CHECK_STATUS(infiniopDestroyRearrangeDescriptor(desc->rearrange_desc_k)); - CHECK_STATUS(infiniopDestroyRearrangeDescriptor(desc->rearrange_desc_v)); - CHECK_STATUS(infiniopDestroyRearrangeDescriptor(desc->rearrange_desc_out)); - CHECK_STATUS(infiniopDestroyGemmDescriptor(desc->matmul_desc1)); - CHECK_STATUS(infiniopDestroyGemmDescriptor(desc->matmul_desc2)); - CHECK_STATUS(infiniopDestroyCausalSoftmaxDescriptor(desc->softmax_desc)); - delete desc; - - return INFINI_STATUS_SUCCESS; -} diff --git a/src/infiniop/ops/avg_pool1d/avg_pool1d.h b/src/infiniop/ops/avg_pool1d/avg_pool1d.h deleted file mode 100644 index a81f46464..000000000 --- a/src/infiniop/ops/avg_pool1d/avg_pool1d.h +++ /dev/null @@ -1,103 +0,0 @@ -#ifndef __AVG_POOL1D_H__ -#define __AVG_POOL1D_H__ - -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" -#include "infiniop/ops/avg_pool1d.h" - -#define DESCRIPTOR(NAMESPACE) \ - namespace op::avg_pool1d::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AvgPool1dInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - AvgPool1dInfo info, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t y_desc, \ - infiniopTensorDescriptor_t x_desc, \ - size_t kernel_size, \ - size_t stride, \ - size_t padding); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *y, \ - const void *x, \ - void *stream) const; \ - }; \ - } - -class AvgPool1dInfo { -private: - AvgPool1dInfo() = default; - -public: - infiniDtype_t dtype; - size_t batch, channels, in_width, out_width; - size_t kernel_size, stride, padding; - - ptrdiff_t y_stride_batch, y_stride_channel, y_stride_width; - ptrdiff_t x_stride_batch, x_stride_channel, x_stride_width; - - static utils::Result createAvgPool1dInfo( - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t kernel_size, - size_t stride, - size_t padding) { - - CHECK_OR_RETURN(y_desc != nullptr && x_desc != nullptr, INFINI_STATUS_NULL_POINTER); - - const infiniDtype_t dtype = y_desc->dtype(); - CHECK_OR_RETURN(dtype == x_desc->dtype(), INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_OR_RETURN(y_desc->ndim() == 3 && x_desc->ndim() == 3, INFINI_STATUS_BAD_TENSOR_SHAPE); - - size_t batch = x_desc->dim(0); - size_t channels = x_desc->dim(1); - size_t in_width = x_desc->dim(2); - - CHECK_OR_RETURN(y_desc->dim(0) == batch, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(y_desc->dim(1) == channels, INFINI_STATUS_BAD_TENSOR_SHAPE); - - size_t padded_len = in_width + 2 * padding; - - CHECK_OR_RETURN(padded_len >= kernel_size, INFINI_STATUS_BAD_TENSOR_SHAPE); - - size_t expected_out_width = (padded_len - kernel_size) / stride + 1; - CHECK_OR_RETURN(y_desc->dim(2) == expected_out_width, INFINI_STATUS_BAD_TENSOR_SHAPE); - - size_t out_width = expected_out_width; - - return utils::Result(AvgPool1dInfo{ - dtype, - batch, channels, in_width, out_width, - kernel_size, stride, padding, - y_desc->stride(0), y_desc->stride(1), y_desc->stride(2), - x_desc->stride(0), x_desc->stride(1), x_desc->stride(2)}); - } -}; - -#endif diff --git a/src/infiniop/ops/avg_pool1d/cpu/avg_pool1d_cpu.cc b/src/infiniop/ops/avg_pool1d/cpu/avg_pool1d_cpu.cc deleted file mode 100644 index 67e5b6623..000000000 --- a/src/infiniop/ops/avg_pool1d/cpu/avg_pool1d_cpu.cc +++ /dev/null @@ -1,96 +0,0 @@ -#include "avg_pool1d_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include - -namespace op::avg_pool1d::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t kernel_size, - size_t stride, - size_t padding) { - - auto handle = reinterpret_cast(handle_); - - auto info = AvgPool1dInfo::createAvgPool1dInfo(y_desc, x_desc, kernel_size, stride, padding); - CHECK_RESULT(info); - - *desc_ptr = new Descriptor( - info.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateAvgPool1d(const AvgPool1dInfo &info, - T *y, - const T *x) { - const float inv_kernel = 1.0f / static_cast(info.kernel_size); - -#pragma omp parallel for - for (ptrdiff_t bc = 0; bc < ptrdiff_t(info.batch * info.channels); ++bc) { - - ptrdiff_t b = bc / info.channels; - ptrdiff_t c = bc % info.channels; - - size_t y_base = b * info.y_stride_batch + c * info.y_stride_channel; - size_t x_base = b * info.x_stride_batch + c * info.x_stride_channel; - - for (size_t ow = 0; ow < info.out_width; ++ow) { - size_t y_offset = y_base + ow * info.y_stride_width; - - long long start_w = static_cast(ow * info.stride) - info.padding; - long long end_w = start_w + info.kernel_size; - - long long valid_start = std::max(0LL, start_w); - long long valid_end = std::min(static_cast(info.in_width), end_w); - - float sum = 0.0f; - for (long long iw = valid_start; iw < valid_end; ++iw) { - size_t x_offset = x_base + iw * info.x_stride_width; - sum += utils::cast(x[x_offset]); - } - - const float avg = sum * inv_kernel; - y[y_offset] = utils::cast(avg); - } - } - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE(TDATA) calculateAvgPool1d(_info, (TDATA *)y, (const TDATA *)x) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - switch (_info.dtype) { - case INFINI_DTYPE_F16: - return CALCULATE(fp16_t); - case INFINI_DTYPE_BF16: - return CALCULATE(bf16_t); - case INFINI_DTYPE_F32: - return CALCULATE(float); - case INFINI_DTYPE_F64: - return CALCULATE(double); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE - -} // namespace op::avg_pool1d::cpu diff --git a/src/infiniop/ops/avg_pool1d/cpu/avg_pool1d_cpu.h b/src/infiniop/ops/avg_pool1d/cpu/avg_pool1d_cpu.h deleted file mode 100644 index 0b9f6c666..000000000 --- a/src/infiniop/ops/avg_pool1d/cpu/avg_pool1d_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __INFINIOP_AVG_POOL1D_CPU_H__ -#define __INFINIOP_AVG_POOL1D_CPU_H__ - -#include "../avg_pool1d.h" - -DESCRIPTOR(cpu) - -#endif diff --git a/src/infiniop/ops/avg_pool1d/cuda/kernel.cuh b/src/infiniop/ops/avg_pool1d/cuda/kernel.cuh deleted file mode 100644 index d052ef97f..000000000 --- a/src/infiniop/ops/avg_pool1d/cuda/kernel.cuh +++ /dev/null @@ -1,58 +0,0 @@ -#ifndef __INFINIOP_AVG_POOL1D_CUDA_KERNEL_CUH__ -#define __INFINIOP_AVG_POOL1D_CUDA_KERNEL_CUH__ - -template -__device__ void avgPool1dKernel( - T *y, - const T *x, - size_t batch, - size_t channels, - size_t in_width, - size_t out_width, - size_t kernel_size, - size_t stride, - size_t padding, - - ptrdiff_t y_stride_batch, - ptrdiff_t y_stride_channel, - ptrdiff_t y_stride_width, - ptrdiff_t x_stride_batch, - ptrdiff_t x_stride_channel, - ptrdiff_t x_stride_width) { - - size_t total_elements = batch * channels * out_width; - - for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - idx < total_elements; - idx += gridDim.x * blockDim.x) { - - size_t ow = idx % out_width; - size_t temp = idx / out_width; - size_t c = temp % channels; - size_t b = temp / channels; - - size_t y_offset = b * y_stride_batch + c * y_stride_channel + ow * y_stride_width; - - long long start_w = static_cast(ow * stride) - padding; - - T sum = 0; - - for (size_t k = 0; k < kernel_size; ++k) { - long long iw = start_w + k; - - if (iw >= 0 && iw < static_cast(in_width)) { - size_t x_offset = b * x_stride_batch + c * x_stride_channel + iw * x_stride_width; - sum += x[x_offset]; - } - } - -#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_QY_API) - // Iluvatar __half doesn't accept size_t directly. - y[y_offset] = sum / static_cast(static_cast(kernel_size)); -#else - y[y_offset] = sum / static_cast(kernel_size); -#endif - } -} - -#endif diff --git a/src/infiniop/ops/avg_pool1d/metax/avg_pool1d_metax.h b/src/infiniop/ops/avg_pool1d/metax/avg_pool1d_metax.h deleted file mode 100644 index 576da66de..000000000 --- a/src/infiniop/ops/avg_pool1d/metax/avg_pool1d_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __INFINIOP_AVG_POOL1D_METAX_H__ -#define __INFINIOP_AVG_POOL1D_METAX_H__ - -#include "../avg_pool1d.h" - -DESCRIPTOR(metax) - -#endif // __INFINIOP_AVG_POOL1D_METAX_H__ diff --git a/src/infiniop/ops/avg_pool1d/metax/avg_pool1d_metax.maca b/src/infiniop/ops/avg_pool1d/metax/avg_pool1d_metax.maca deleted file mode 100644 index 9b3f15b9a..000000000 --- a/src/infiniop/ops/avg_pool1d/metax/avg_pool1d_metax.maca +++ /dev/null @@ -1,170 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "avg_pool1d_metax.h" -#include "../../../devices/metax/metax_kernel_common.h" - -#include - -namespace op::avg_pool1d::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t kernel_size, - size_t stride, - size_t padding) { - - auto handle = reinterpret_cast(handle_); - - auto info = AvgPool1dInfo::createAvgPool1dInfo(y_desc, x_desc, kernel_size, stride, padding); - CHECK_RESULT(info); - - *desc_ptr = new Descriptor( - info.take(), - 0, - new Opaque{handle->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -__device__ __forceinline__ Tdata castToOutput(Tcompute val) { - if constexpr (std::is_same_v) { - return __float2half(static_cast(val)); - } else if constexpr (std::is_same_v) { - return __float2bfloat16(static_cast(val)); - } else { - return static_cast(val); - } -} - -template -INFINIOP_METAX_KERNEL avgPool1dGlobalKernel( - Tdata *y, - const Tdata *x, - size_t batch, - size_t channels, - size_t in_width, - size_t out_width, - size_t kernel_size, - size_t stride, - size_t padding, - ptrdiff_t y_stride_batch, - ptrdiff_t y_stride_channel, - ptrdiff_t y_stride_width, - ptrdiff_t x_stride_batch, - ptrdiff_t x_stride_channel, - ptrdiff_t x_stride_width) { - - size_t total_elements = batch * channels * out_width; - Tcompute inv_kernel = Tcompute(1) / static_cast(kernel_size); - - for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - idx < total_elements; - idx += gridDim.x * blockDim.x) { - - size_t ow = idx % out_width; - size_t temp = idx / out_width; - size_t c = temp % channels; - size_t b = temp / channels; - - size_t y_offset = b * y_stride_batch + c * y_stride_channel + ow * y_stride_width; - size_t x_base = b * x_stride_batch + c * x_stride_channel; - - long long start_w = static_cast(ow * stride) - static_cast(padding); - long long end_w = start_w + static_cast(kernel_size); - long long iw_start = start_w < 0 ? 0 : start_w; - long long iw_end = end_w > static_cast(in_width) ? static_cast(in_width) : end_w; - - Tcompute sum = Tcompute(0); - if (iw_start < iw_end) { - size_t x_offset = x_base + static_cast(iw_start) * x_stride_width; - for (long long iw = iw_start; iw < iw_end; ++iw) { - sum += static_cast(x[x_offset]); - x_offset += x_stride_width; - } - } - - y[y_offset] = castToOutput(sum * inv_kernel); - } -} - -template -infiniStatus_t calculateAvgPool1d( - const AvgPool1dInfo &info, - int max_threads_per_block, - Tdata *y, - const Tdata *x, - hcStream_t stream) { - - size_t total_elements = info.batch * info.channels * info.out_width; - - int block_size = 256; - if (max_threads_per_block > 0 && max_threads_per_block < block_size) { - block_size = max_threads_per_block; - } - - size_t grid_size = (total_elements + block_size - 1) / block_size; - if (grid_size > 65535) { - grid_size = 65535; - } - - avgPool1dGlobalKernel<<>>( - y, x, - info.batch, info.channels, info.in_width, info.out_width, - info.kernel_size, info.stride, info.padding, - info.y_stride_batch, info.y_stride_channel, info.y_stride_width, - info.x_stride_batch, info.x_stride_channel, info.x_stride_width); - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE(TDATA, TCOMPUTE) \ - calculateAvgPool1d( \ - _info, \ - _opaque->internal->maxThreadsPerBlock(), \ - (TDATA *)y, \ - (const TDATA *)x, \ - (hcStream_t)stream) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - (void)workspace; - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_info.dtype) { - case INFINI_DTYPE_F16: - return CALCULATE(half, float); - case INFINI_DTYPE_BF16: - return CALCULATE(cuda_bfloat16, float); - case INFINI_DTYPE_F32: - return CALCULATE(float, float); - case INFINI_DTYPE_F64: - return CALCULATE(double, double); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE - -} // namespace op::avg_pool1d::metax diff --git a/src/infiniop/ops/avg_pool1d/moore/avg_pool1d_kernel.h b/src/infiniop/ops/avg_pool1d/moore/avg_pool1d_kernel.h deleted file mode 100644 index 9034d7358..000000000 --- a/src/infiniop/ops/avg_pool1d/moore/avg_pool1d_kernel.h +++ /dev/null @@ -1,72 +0,0 @@ -#ifndef __INFINIOP_AVG_POOL1D_MOORE_KERNEL_H__ -#define __INFINIOP_AVG_POOL1D_MOORE_KERNEL_H__ - -#include - -namespace op::avg_pool1d::moore { - -template -__device__ __forceinline__ Tdata castToOutput(Tcompute val) { - if constexpr (std::is_same_v) { - return __float2half(static_cast(val)); - } else if constexpr (std::is_same_v) { - return __float2bfloat16_rn(static_cast(val)); - } else { - return static_cast(val); - } -} - -template -__device__ void avgPool1dKernel( - Tdata *y, - const Tdata *x, - size_t batch, - size_t channels, - size_t in_width, - size_t out_width, - size_t kernel_size, - size_t stride, - size_t padding, - ptrdiff_t y_stride_batch, - ptrdiff_t y_stride_channel, - ptrdiff_t y_stride_width, - ptrdiff_t x_stride_batch, - ptrdiff_t x_stride_channel, - ptrdiff_t x_stride_width) { - - size_t total_elements = batch * channels * out_width; - Tcompute inv_kernel = Tcompute(1) / static_cast(kernel_size); - - for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - idx < total_elements; - idx += gridDim.x * blockDim.x) { - - size_t ow = idx % out_width; - size_t temp = idx / out_width; - size_t c = temp % channels; - size_t b = temp / channels; - - size_t y_offset = b * y_stride_batch + c * y_stride_channel + ow * y_stride_width; - size_t x_base = b * x_stride_batch + c * x_stride_channel; - - long long start_w = static_cast(ow * stride) - static_cast(padding); - long long end_w = start_w + static_cast(kernel_size); - long long iw_start = start_w < 0 ? 0 : start_w; - long long iw_end = end_w > static_cast(in_width) ? static_cast(in_width) : end_w; - - Tcompute sum = Tcompute(0); - if (iw_start < iw_end) { - size_t x_offset = x_base + static_cast(iw_start) * x_stride_width; - for (long long iw = iw_start; iw < iw_end; ++iw) { - sum += static_cast(x[x_offset]); - x_offset += x_stride_width; - } - } - - y[y_offset] = castToOutput(sum * inv_kernel); - } -} - -} // namespace op::avg_pool1d::moore - -#endif // __INFINIOP_AVG_POOL1D_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/avg_pool1d/moore/avg_pool1d_moore.h b/src/infiniop/ops/avg_pool1d/moore/avg_pool1d_moore.h deleted file mode 100644 index 604d06012..000000000 --- a/src/infiniop/ops/avg_pool1d/moore/avg_pool1d_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __INFINIOP_AVG_POOL1D_MOORE_H__ -#define __INFINIOP_AVG_POOL1D_MOORE_H__ - -#include "../avg_pool1d.h" - -DESCRIPTOR(moore) - -#endif // __INFINIOP_AVG_POOL1D_MOORE_H__ diff --git a/src/infiniop/ops/avg_pool1d/moore/avg_pool1d_moore.mu b/src/infiniop/ops/avg_pool1d/moore/avg_pool1d_moore.mu deleted file mode 100644 index 518d249b9..000000000 --- a/src/infiniop/ops/avg_pool1d/moore/avg_pool1d_moore.mu +++ /dev/null @@ -1,135 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "avg_pool1d_moore.h" - -#include "../../../devices/moore/moore_kernel_common.h" - -#include "avg_pool1d_kernel.h" - -namespace op::avg_pool1d::moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t kernel_size, - size_t stride, - size_t padding) { - - auto handle = reinterpret_cast(handle_); - - auto info = AvgPool1dInfo::createAvgPool1dInfo(y_desc, x_desc, kernel_size, stride, padding); - CHECK_RESULT(info); - - *desc_ptr = new Descriptor( - info.take(), - 0, - new Opaque{handle->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -INFINIOP_MOORE_KERNEL avgPool1dGlobalKernel( - Tdata *y, - const Tdata *x, - size_t batch, - size_t channels, - size_t in_width, - size_t out_width, - size_t kernel_size, - size_t stride, - size_t padding, - ptrdiff_t y_stride_batch, - ptrdiff_t y_stride_channel, - ptrdiff_t y_stride_width, - ptrdiff_t x_stride_batch, - ptrdiff_t x_stride_channel, - ptrdiff_t x_stride_width) { - - avgPool1dKernel( - y, x, - batch, channels, in_width, out_width, - kernel_size, stride, padding, - y_stride_batch, y_stride_channel, y_stride_width, - x_stride_batch, x_stride_channel, x_stride_width); -} - -template -infiniStatus_t calculateAvgPool1d( - const AvgPool1dInfo &info, - int max_threads_per_block, - Tdata *y, - const Tdata *x, - musaStream_t stream) { - - size_t total_elements = info.batch * info.channels * info.out_width; - - int block_size = 256; - if (max_threads_per_block > 0 && max_threads_per_block < block_size) { - block_size = max_threads_per_block; - } - - size_t grid_size = (total_elements + block_size - 1) / block_size; - if (grid_size > 65535) { - grid_size = 65535; - } - - avgPool1dGlobalKernel<<>>( - y, x, - info.batch, info.channels, info.in_width, info.out_width, - info.kernel_size, info.stride, info.padding, - info.y_stride_batch, info.y_stride_channel, info.y_stride_width, - info.x_stride_batch, info.x_stride_channel, info.x_stride_width); - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE(TDATA, TCOMPUTE) \ - calculateAvgPool1d(\ - _info,\ - _opaque->internal->maxThreadsPerBlock(),\ - (TDATA *)y,\ - (const TDATA *)x,\ - (musaStream_t)stream) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - (void)workspace; - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_info.dtype) { - case INFINI_DTYPE_F16: - return CALCULATE(half, float); - case INFINI_DTYPE_BF16: - return CALCULATE(cuda_bfloat16, float); - case INFINI_DTYPE_F32: - return CALCULATE(float, float); - case INFINI_DTYPE_F64: - return CALCULATE(double, double); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE - -} // namespace op::avg_pool1d::moore diff --git a/src/infiniop/ops/avg_pool1d/nvidia/avg_pool1d_nvidia.cu b/src/infiniop/ops/avg_pool1d/nvidia/avg_pool1d_nvidia.cu deleted file mode 100644 index 634ce9018..000000000 --- a/src/infiniop/ops/avg_pool1d/nvidia/avg_pool1d_nvidia.cu +++ /dev/null @@ -1,126 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../cuda/kernel.cuh" -#include "avg_pool1d_nvidia.cuh" - -template -__global__ void avgPool1dGlobalKernel( - T *y, - const T *x, - size_t batch, - size_t channels, - size_t in_width, - size_t out_width, - size_t kernel_size, - size_t stride, - size_t padding, - ptrdiff_t y_stride_batch, - ptrdiff_t y_stride_channel, - ptrdiff_t y_stride_width, - ptrdiff_t x_stride_batch, - ptrdiff_t x_stride_channel, - ptrdiff_t x_stride_width) { - - avgPool1dKernel( - y, x, - batch, channels, in_width, out_width, - kernel_size, stride, padding, - y_stride_batch, y_stride_channel, y_stride_width, - x_stride_batch, x_stride_channel, x_stride_width); -} - -namespace op::avg_pool1d::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - size_t kernel_size, - size_t stride, - size_t padding) { - - auto handle = reinterpret_cast(handle_); - - auto info = AvgPool1dInfo::createAvgPool1dInfo(y_desc, x_desc, kernel_size, stride, padding); - CHECK_RESULT(info); - - *desc_ptr = new Descriptor( - info.take(), - 0, - new Opaque{reinterpret_cast(handle)->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateAvgPool1d( - const AvgPool1dInfo &info, - int max_threads_per_block, - T *y, - const T *x, - cudaStream_t stream) { - - size_t total_elements = info.batch * info.channels * info.out_width; - - int block_size = 256; - if (max_threads_per_block > 0 && max_threads_per_block < 256) { - block_size = max_threads_per_block; - } - - size_t grid_size = (total_elements + block_size - 1) / block_size; - if (grid_size > 65535) { - grid_size = 65535; - } - - avgPool1dGlobalKernel<<>>( - y, x, - info.batch, info.channels, info.in_width, info.out_width, - info.kernel_size, info.stride, info.padding, - info.y_stride_batch, info.y_stride_channel, info.y_stride_width, - info.x_stride_batch, info.x_stride_channel, info.x_stride_width); - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE(TDATA) \ - calculateAvgPool1d(_info, \ - _opaque->internal->maxThreadsPerBlock(), \ - (TDATA *)y, \ - (const TDATA *)x, \ - (cudaStream_t)stream) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - switch (_info.dtype) { - case INFINI_DTYPE_F16: - return CALCULATE(half); - case INFINI_DTYPE_BF16: - return CALCULATE(cuda_bfloat16); - case INFINI_DTYPE_F32: - return CALCULATE(float); - case INFINI_DTYPE_F64: - return CALCULATE(double); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE - -} // namespace op::avg_pool1d::nvidia diff --git a/src/infiniop/ops/avg_pool1d/nvidia/avg_pool1d_nvidia.cuh b/src/infiniop/ops/avg_pool1d/nvidia/avg_pool1d_nvidia.cuh deleted file mode 100644 index 1019ce354..000000000 --- a/src/infiniop/ops/avg_pool1d/nvidia/avg_pool1d_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __INFINIOP_AVG_POOL1D_CUDA_H__ -#define __INFINIOP_AVG_POOL1D_CUDA_H__ - -#include "../avg_pool1d.h" - -DESCRIPTOR(nvidia) - -#endif diff --git a/src/infiniop/ops/avg_pool1d/operator.cc b/src/infiniop/ops/avg_pool1d/operator.cc deleted file mode 100644 index 54ea110fd..000000000 --- a/src/infiniop/ops/avg_pool1d/operator.cc +++ /dev/null @@ -1,180 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/avg_pool1d.h" - -#ifdef ENABLE_CPU_API -#include "cpu/avg_pool1d_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/avg_pool1d_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/avg_pool1d_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/avg_pool1d_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAvgPool1dDescriptor( - infiniopHandle_t handle, - infiniopAvgPool1dDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - size_t kernel_size, - size_t stride, - size_t padding) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::avg_pool1d::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y, \ - x, \ - kernel_size, \ - stride, \ - padding) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAvgPool1dWorkspaceSize(infiniopAvgPool1dDescriptor_t desc, - size_t *size) { -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET -} - -__INFINI_C infiniStatus_t infiniopAvgPool1d( - infiniopAvgPool1dDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, x, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyAvgPool1dDescriptor(infiniopAvgPool1dDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/avg_pool3d/cpu/avg_pool3d_cpu.cc b/src/infiniop/ops/avg_pool3d/cpu/avg_pool3d_cpu.cc deleted file mode 100644 index 9553c17f5..000000000 --- a/src/infiniop/ops/avg_pool3d/cpu/avg_pool3d_cpu.cc +++ /dev/null @@ -1,243 +0,0 @@ -#include "avg_pool3d_cpu.h" -#include "../../../../utils.h" -#include -#include -#include -#include - -namespace op::avg_pool3d::cpu { - -utils::Result AvgPool3dInfo::create( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - void *kernel_size, - void *stride, - void *padding) { - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x_shape.size() != 5 || y_shape.size() != 5) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t batch = x_shape[0]; - size_t channels = x_shape[1]; - size_t input_d = x_shape[2]; - size_t input_h = x_shape[3]; - size_t input_w = x_shape[4]; - - // Parse kernel_size - size_t kernel_d, kernel_h, kernel_w; - if (kernel_size) { - size_t *ks = reinterpret_cast(kernel_size); - if (ks[0] == 0 || ks[1] == 0 || ks[2] == 0) { - return INFINI_STATUS_BAD_PARAM; - } - kernel_d = ks[0]; - kernel_h = ks[1]; - kernel_w = ks[2]; - } else { - return INFINI_STATUS_BAD_PARAM; - } - - // Parse stride (default to kernel_size if not provided) - size_t stride_d, stride_h, stride_w; - if (stride) { - size_t *s = reinterpret_cast(stride); - if (s[0] == 0 || s[1] == 0 || s[2] == 0) { - return INFINI_STATUS_BAD_PARAM; - } - stride_d = s[0]; - stride_h = s[1]; - stride_w = s[2]; - } else { - stride_d = kernel_d; - stride_h = kernel_h; - stride_w = kernel_w; - } - - // Parse padding - size_t pad_d, pad_h, pad_w; - if (padding) { - size_t *p = reinterpret_cast(padding); - // Assume it's always a tuple of 3 values for 3D pooling - pad_d = p[0]; - pad_h = p[1]; - pad_w = p[2]; - } else { - pad_d = pad_h = pad_w = 0; - } - - // Calculate output dimensions. Guard against unsigned underflow when kernel > input + 2*pad. - if (pad_d > (std::numeric_limits::max() - input_d) / 2 || pad_h > (std::numeric_limits::max() - input_h) / 2 || pad_w > (std::numeric_limits::max() - input_w) / 2) { - return INFINI_STATUS_BAD_PARAM; - } - size_t effective_d = input_d + 2 * pad_d; - size_t effective_h = input_h + 2 * pad_h; - size_t effective_w = input_w + 2 * pad_w; - if (kernel_d > effective_d || kernel_h > effective_h || kernel_w > effective_w) { - return INFINI_STATUS_BAD_PARAM; - } - - size_t output_d = (effective_d - kernel_d) / stride_d + 1; - size_t output_h = (effective_h - kernel_h) / stride_h + 1; - size_t output_w = (effective_w - kernel_w) / stride_w + 1; - - // Verify output shape - if (y_shape[0] != batch || y_shape[1] != channels || y_shape[2] != output_d || y_shape[3] != output_h || y_shape[4] != output_w) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - AvgPool3dInfo info; - info.batch = batch; - info.channels = channels; - info.input_d = input_d; - info.input_h = input_h; - info.input_w = input_w; - info.output_d = output_d; - info.output_h = output_h; - info.output_w = output_w; - info.kernel_d = kernel_d; - info.kernel_h = kernel_h; - info.kernel_w = kernel_w; - info.stride_d = stride_d; - info.stride_h = stride_h; - info.stride_w = stride_w; - info.pad_d = pad_d; - info.pad_h = pad_h; - info.pad_w = pad_w; - info.input_strides = x_desc->strides(); - info.output_strides = y_desc->strides(); - - if (info.input_strides.size() != x_shape.size() || info.output_strides.size() != y_shape.size()) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - // Reject broadcasted (0-stride) or negative strides for dimensions that are actually indexed. - // The kernel computes indices using size_t, so negative strides would underflow and go OOB. - for (size_t i = 0; i < x_shape.size(); ++i) { - if (x_shape[i] > 1 && info.input_strides[i] <= 0) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - } - for (size_t i = 0; i < y_shape.size(); ++i) { - if (y_shape[i] > 1 && info.output_strides[i] <= 0) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - } - - return utils::Result(std::move(info)); -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - void *kernel_size, - void *stride, - void *padding) { - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - if (y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto info_result = AvgPool3dInfo::create(x_desc, y_desc, kernel_size, stride, padding); - CHECK_RESULT(info_result); - - *desc_ptr = new Descriptor(dtype, info_result.take(), handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -void avg_pool3d_impl( - const AvgPool3dInfo &info, - T *y, - const T *x) { - - const size_t kernel_size = info.kernel_d * info.kernel_h * info.kernel_w; - using Tacc = std::conditional_t, double, float>; - const Tacc inv_kernel_size = Tacc(1) / static_cast(kernel_size); - - size_t total = info.batch * info.channels * info.output_d * info.output_h * info.output_w; - -#pragma omp parallel for - for (ptrdiff_t idx = 0; idx < (ptrdiff_t)total; ++idx) { - - size_t tmp = idx; - - size_t ow = tmp % info.output_w; - tmp /= info.output_w; - size_t oh = tmp % info.output_h; - tmp /= info.output_h; - size_t od = tmp % info.output_d; - tmp /= info.output_d; - size_t c = tmp % info.channels; - tmp /= info.channels; - size_t b = tmp; - - using Tacc = std::conditional_t, double, float>; - Tacc sum = Tacc(0); - - ptrdiff_t id_start = (ptrdiff_t)od * (ptrdiff_t)info.stride_d - (ptrdiff_t)info.pad_d; - ptrdiff_t ih_start = (ptrdiff_t)oh * (ptrdiff_t)info.stride_h - (ptrdiff_t)info.pad_h; - ptrdiff_t iw_start = (ptrdiff_t)ow * (ptrdiff_t)info.stride_w - (ptrdiff_t)info.pad_w; - - for (size_t kd = 0; kd < info.kernel_d; ++kd) { - for (size_t kh = 0; kh < info.kernel_h; ++kh) { - for (size_t kw = 0; kw < info.kernel_w; ++kw) { - - ptrdiff_t id = id_start + kd; - ptrdiff_t ih = ih_start + kh; - ptrdiff_t iw = iw_start + kw; - - if (id >= 0 && id < (ptrdiff_t)info.input_d && ih >= 0 && ih < (ptrdiff_t)info.input_h && iw >= 0 && iw < (ptrdiff_t)info.input_w) { - - size_t x_idx = b * info.input_strides[0] + c * info.input_strides[1] + (size_t)id * info.input_strides[2] + (size_t)ih * info.input_strides[3] + (size_t)iw * info.input_strides[4]; - - sum += utils::cast(x[x_idx]); - } - } - } - } - - size_t y_idx = b * info.output_strides[0] + c * info.output_strides[1] + od * info.output_strides[2] + oh * info.output_strides[3] + ow * info.output_strides[4]; - - y[y_idx] = utils::cast(sum * inv_kernel_size); - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - avg_pool3d_impl(_info, reinterpret_cast(y), reinterpret_cast(x)); - break; - case INFINI_DTYPE_BF16: - avg_pool3d_impl(_info, reinterpret_cast(y), reinterpret_cast(x)); - break; - case INFINI_DTYPE_F32: - avg_pool3d_impl(_info, reinterpret_cast(y), reinterpret_cast(x)); - break; - case INFINI_DTYPE_F64: - avg_pool3d_impl(_info, reinterpret_cast(y), reinterpret_cast(x)); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::avg_pool3d::cpu diff --git a/src/infiniop/ops/avg_pool3d/cpu/avg_pool3d_cpu.h b/src/infiniop/ops/avg_pool3d/cpu/avg_pool3d_cpu.h deleted file mode 100644 index 57be33f9e..000000000 --- a/src/infiniop/ops/avg_pool3d/cpu/avg_pool3d_cpu.h +++ /dev/null @@ -1,64 +0,0 @@ -#ifndef __AVG_POOL3D_CPU_H__ -#define __AVG_POOL3D_CPU_H__ - -#include "../../../devices/cpu/common_cpu.h" -#include "../../../operator.h" -#include "../../../tensor.h" -#include - -namespace op::avg_pool3d::cpu { - -struct AvgPool3dInfo { - size_t batch; - size_t channels; - size_t input_d, input_h, input_w; - size_t output_d, output_h, output_w; - size_t kernel_d, kernel_h, kernel_w; - size_t stride_d, stride_h, stride_w; - size_t pad_d, pad_h, pad_w; - std::vector input_strides; - std::vector output_strides; - - static utils::Result create( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - void *kernel_size, - void *stride, - void *padding); -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - AvgPool3dInfo _info; - - Descriptor(infiniDtype_t dtype, AvgPool3dInfo info, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - void *kernel_size, - void *stride, - void *padding); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::avg_pool3d::cpu - -#endif // __AVG_POOL3D_CPU_H__ diff --git a/src/infiniop/ops/avg_pool3d/metax/avg_pool3d_metax.h b/src/infiniop/ops/avg_pool3d/metax/avg_pool3d_metax.h deleted file mode 100644 index fdc838868..000000000 --- a/src/infiniop/ops/avg_pool3d/metax/avg_pool3d_metax.h +++ /dev/null @@ -1,42 +0,0 @@ -#ifndef __AVG_POOL3D_METAX_H__ -#define __AVG_POOL3D_METAX_H__ - -#include "../../../operator.h" -#include "../../../tensor.h" -#include - -namespace op::avg_pool3d::metax { - -class Descriptor final : public InfiniopDescriptor { - struct Opaque; - std::unique_ptr _opaque; - infiniDtype_t _dtype; - - Descriptor(infiniDtype_t dtype, std::unique_ptr opaque, - infiniDevice_t device_type, int device_id); - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - void *kernel_size, - void *stride, - void *padding); - - size_t workspaceSize() const; - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::avg_pool3d::metax - -#endif // __AVG_POOL3D_METAX_H__ diff --git a/src/infiniop/ops/avg_pool3d/metax/avg_pool3d_metax.maca b/src/infiniop/ops/avg_pool3d/metax/avg_pool3d_metax.maca deleted file mode 100644 index fdd2f7d15..000000000 --- a/src/infiniop/ops/avg_pool3d/metax/avg_pool3d_metax.maca +++ /dev/null @@ -1,179 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/metax/metax_common.h" -#include "avg_pool3d_metax.h" - -namespace op::avg_pool3d::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; - hcdnnTensorDescriptor_t x_desc = nullptr; - hcdnnTensorDescriptor_t y_desc = nullptr; - hcdnnPoolingDescriptor_t pool_desc = nullptr; - - size_t workspace_size = 0; - - Opaque(std::shared_ptr internal_ptr) - : internal(internal_ptr) {} - - ~Opaque() { - if (x_desc) { - hcdnnDestroyTensorDescriptor(x_desc); - } - if (y_desc) { - hcdnnDestroyTensorDescriptor(y_desc); - } - if (pool_desc) { - hcdnnDestroyPoolingDescriptor(pool_desc); - } - } -}; - -Descriptor::Descriptor(infiniDtype_t dtype, std::unique_ptr opaque, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _opaque(std::move(opaque)), - _dtype(dtype) {} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - void *kernel_size, - void *stride, - void *padding) { - - auto metax_handle = reinterpret_cast(handle); - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x_shape.size() != 5 || y_shape.size() != 5) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t *ks = reinterpret_cast(kernel_size); - size_t kernel_d = ks[0], kernel_h = ks[1], kernel_w = ks[2]; - - size_t stride_d, stride_h, stride_w; - if (stride) { - size_t *s = reinterpret_cast(stride); - stride_d = s[0]; - stride_h = s[1]; - stride_w = s[2]; - } else { - stride_d = kernel_d; - stride_h = kernel_h; - stride_w = kernel_w; - } - - size_t pad_d, pad_h, pad_w; - if (padding) { - size_t *p = reinterpret_cast(padding); - pad_d = p[0]; - pad_h = p[1]; - pad_w = p[2]; - } else { - pad_d = pad_h = pad_w = 0; - } - - auto opaque = std::make_unique(metax_handle->internal()); - - // Create hcdnn descriptors - CHECK_MCDNN(hcdnnCreateTensorDescriptor(&opaque->x_desc)); - CHECK_MCDNN(hcdnnCreateTensorDescriptor(&opaque->y_desc)); - CHECK_MCDNN(hcdnnCreatePoolingDescriptor(&opaque->pool_desc)); - - // Set tensor descriptors - int n = static_cast(x_shape[0]); - int c = static_cast(x_shape[1]); - int d = static_cast(x_shape[2]); - int h = static_cast(x_shape[3]); - int w = static_cast(x_shape[4]); - int out_d = static_cast(y_shape[2]); - int out_h = static_cast(y_shape[3]); - int out_w = static_cast(y_shape[4]); - - int input_dims[5] = {n, c, d, h, w}; - int input_strides[5] = { - static_cast(c * d * h * w), - static_cast(d * h * w), - static_cast(h * w), - static_cast(w), - 1}; - - int output_dims[5] = {n, c, out_d, out_h, out_w}; - int output_strides[5] = { - static_cast(c * out_d * out_h * out_w), - static_cast(out_d * out_h * out_w), - static_cast(out_h * out_w), - static_cast(out_w), - 1}; - - hcdnnDataType_t hcdnn_dtype = device::metax::getHcdnnDtype(dtype); - CHECK_MCDNN(hcdnnSetTensorNdDescriptor( - opaque->x_desc, hcdnn_dtype, 5, input_dims, input_strides)); - CHECK_MCDNN(hcdnnSetTensorNdDescriptor( - opaque->y_desc, hcdnn_dtype, 5, output_dims, output_strides)); - - // Set pooling descriptor - int window_dims[3] = {static_cast(kernel_d), static_cast(kernel_h), static_cast(kernel_w)}; - int padding_dims[3] = {static_cast(pad_d), static_cast(pad_h), static_cast(pad_w)}; - int stride_dims[3] = {static_cast(stride_d), static_cast(stride_h), static_cast(stride_w)}; - - CHECK_MCDNN(hcdnnSetPoolingNdDescriptor( - opaque->pool_desc, - HCDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING, - HCDNN_NOT_PROPAGATE_NAN, - 3, window_dims, padding_dims, stride_dims)); - - *desc_ptr = new Descriptor(dtype, std::move(opaque), handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -size_t Descriptor::workspaceSize() const { - return _opaque->workspace_size; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - auto hc_stream = reinterpret_cast(stream); - return _opaque->internal->useMcdnn(hc_stream, [&](hcdnnHandle_t hcdnn_handle) { - const void *alpha = nullptr; - const void *beta = nullptr; - if (_dtype == INFINI_DTYPE_F32) { - static const float alpha_val = 1.0f, beta_val = 0.0f; - alpha = &alpha_val; - beta = &beta_val; - } else if (_dtype == INFINI_DTYPE_F64) { - static const double alpha_val = 1.0, beta_val = 0.0; - alpha = &alpha_val; - beta = &beta_val; - } else { - static const float alpha_val = 1.0f, beta_val = 0.0f; - alpha = &alpha_val; - beta = &beta_val; - } - - CHECK_MCDNN(hcdnnPoolingForward( - hcdnn_handle, - _opaque->pool_desc, - alpha, - _opaque->x_desc, x, - beta, - _opaque->y_desc, y)); - - return INFINI_STATUS_SUCCESS; - }); -} - -} // namespace op::avg_pool3d::metax diff --git a/src/infiniop/ops/avg_pool3d/moore/avg_pool3d_moore.h b/src/infiniop/ops/avg_pool3d/moore/avg_pool3d_moore.h deleted file mode 100644 index 4ee472b28..000000000 --- a/src/infiniop/ops/avg_pool3d/moore/avg_pool3d_moore.h +++ /dev/null @@ -1,42 +0,0 @@ -#ifndef __AVG_POOL3D_MOORE_H__ -#define __AVG_POOL3D_MOORE_H__ - -#include "../../../operator.h" -#include "../../../tensor.h" -#include - -namespace op::avg_pool3d::moore { - -class Descriptor final : public InfiniopDescriptor { - struct Opaque; - std::unique_ptr _opaque; - infiniDtype_t _dtype; - - Descriptor(infiniDtype_t dtype, std::unique_ptr opaque, - infiniDevice_t device_type, int device_id); - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - void *kernel_size, - void *stride, - void *padding); - - size_t workspaceSize() const; - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::avg_pool3d::moore - -#endif // __AVG_POOL3D_MOORE_H__ diff --git a/src/infiniop/ops/avg_pool3d/moore/avg_pool3d_moore.mu b/src/infiniop/ops/avg_pool3d/moore/avg_pool3d_moore.mu deleted file mode 100644 index a538bb92c..000000000 --- a/src/infiniop/ops/avg_pool3d/moore/avg_pool3d_moore.mu +++ /dev/null @@ -1,51 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/moore/moore_common.h" -#include "avg_pool3d_moore.h" - -namespace op::avg_pool3d::moore { - -// MOORE platform uses musa API which is CUDA-compatible -// We can reuse the NVIDIA implementation structure -// For now, return NOT_IMPLEMENTED as a placeholder -// Full implementation would require adapting NVIDIA code to use musaStream_t - -struct Descriptor::Opaque {}; - -Descriptor::Descriptor(infiniDtype_t dtype, std::unique_ptr opaque, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _opaque(std::move(opaque)), - _dtype(dtype) {} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - void *kernel_size, - void *stride, - void *padding) { - - // MOORE implementation would be similar to NVIDIA but using musa API - // For now, delegate to a CPU fallback or implement custom kernel - // This is a simplified placeholder - full implementation needed - return INFINI_STATUS_NOT_IMPLEMENTED; -} - -size_t Descriptor::workspaceSize() const { - return 0; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - // Placeholder - full implementation needed - return INFINI_STATUS_NOT_IMPLEMENTED; -} - -} // namespace op::avg_pool3d::moore diff --git a/src/infiniop/ops/avg_pool3d/nvidia/avg_pool3d_nvidia.cu b/src/infiniop/ops/avg_pool3d/nvidia/avg_pool3d_nvidia.cu deleted file mode 100644 index 8bee6c14a..000000000 --- a/src/infiniop/ops/avg_pool3d/nvidia/avg_pool3d_nvidia.cu +++ /dev/null @@ -1,200 +0,0 @@ -#include "../../../../utils.h" -#include "avg_pool3d_nvidia.cuh" -#include -#include - -namespace op::avg_pool3d::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; - cudnnTensorDescriptor_t x_desc = nullptr; - cudnnTensorDescriptor_t y_desc = nullptr; - cudnnPoolingDescriptor_t pool_desc = nullptr; - size_t workspace_size = 0; - - Opaque(std::shared_ptr internal_ptr) - : internal(internal_ptr) {} - - ~Opaque() { - if (x_desc) { - cudnnDestroyTensorDescriptor(x_desc); - } - if (y_desc) { - cudnnDestroyTensorDescriptor(y_desc); - } - if (pool_desc) { - cudnnDestroyPoolingDescriptor(pool_desc); - } - } -}; - -Descriptor::Descriptor(infiniDtype_t dtype, std::unique_ptr opaque, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _opaque(std::move(opaque)), - _dtype(dtype) {} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - void *kernel_size, - void *stride, - void *padding) { - - auto nvidia_handle = reinterpret_cast(handle); - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - if (y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x_shape.size() != 5 || y_shape.size() != 5) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (!kernel_size) { - return INFINI_STATUS_BAD_PARAM; - } - size_t *ks = reinterpret_cast(kernel_size); - if (ks[0] == 0 || ks[1] == 0 || ks[2] == 0) { - return INFINI_STATUS_BAD_PARAM; - } - size_t kernel_d = ks[0], kernel_h = ks[1], kernel_w = ks[2]; - - size_t stride_d, stride_h, stride_w; - if (stride) { - size_t *s = reinterpret_cast(stride); - stride_d = s[0]; - stride_h = s[1]; - stride_w = s[2]; - } else { - stride_d = kernel_d; - stride_h = kernel_h; - stride_w = kernel_w; - } - - size_t pad_d, pad_h, pad_w; - if (padding) { - size_t *p = reinterpret_cast(padding); - pad_d = p[0]; - pad_h = p[1]; - pad_w = p[2]; - } else { - pad_d = pad_h = pad_w = 0; - } - - auto opaque = std::make_unique(nvidia_handle->internal()); - - // Create cuDNN descriptors - CHECK_CUDNN(cudnnCreateTensorDescriptor(&opaque->x_desc)); - CHECK_CUDNN(cudnnCreateTensorDescriptor(&opaque->y_desc)); - CHECK_CUDNN(cudnnCreatePoolingDescriptor(&opaque->pool_desc)); - - // Set tensor descriptors - int n = static_cast(x_shape[0]); - int c = static_cast(x_shape[1]); - int d = static_cast(x_shape[2]); - int h = static_cast(x_shape[3]); - int w = static_cast(x_shape[4]); - int out_d = static_cast(y_shape[2]); - int out_h = static_cast(y_shape[3]); - int out_w = static_cast(y_shape[4]); - - int input_dims[5] = {n, c, d, h, w}; - auto x_strides = x_desc->strides(); - if (x_strides.size() != 5) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - int input_strides[5] = {}; - for (size_t i = 0; i < 5; ++i) { - if (x_strides[i] <= 0 || x_strides[i] > std::numeric_limits::max()) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - input_strides[i] = static_cast(x_strides[i]); - } - - int output_dims[5] = {n, c, out_d, out_h, out_w}; - auto y_strides = y_desc->strides(); - if (y_strides.size() != 5) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - int output_strides[5] = {}; - for (size_t i = 0; i < 5; ++i) { - if (y_strides[i] <= 0 || y_strides[i] > std::numeric_limits::max()) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - output_strides[i] = static_cast(y_strides[i]); - } - - cudnnDataType_t cudnn_dtype = device::nvidia::getCudnnDtype(dtype); - CHECK_CUDNN(cudnnSetTensorNdDescriptor( - opaque->x_desc, cudnn_dtype, 5, input_dims, input_strides)); - CHECK_CUDNN(cudnnSetTensorNdDescriptor( - opaque->y_desc, cudnn_dtype, 5, output_dims, output_strides)); - - // Set pooling descriptor - int window_dims[3] = {static_cast(kernel_d), static_cast(kernel_h), static_cast(kernel_w)}; - int padding_dims[3] = {static_cast(pad_d), static_cast(pad_h), static_cast(pad_w)}; - int stride_dims[3] = {static_cast(stride_d), static_cast(stride_h), static_cast(stride_w)}; - - CHECK_CUDNN(cudnnSetPoolingNdDescriptor( - opaque->pool_desc, - CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING, - CUDNN_NOT_PROPAGATE_NAN, - 3, window_dims, padding_dims, stride_dims)); - - *desc_ptr = new Descriptor(dtype, std::move(opaque), handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -size_t Descriptor::workspaceSize() const { - return _opaque->workspace_size; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - return _opaque->internal->useCudnn(cuda_stream, [&](cudnnHandle_t cudnn_handle) { - const void *alpha = nullptr; - const void *beta = nullptr; - if (_dtype == INFINI_DTYPE_F32) { - static const float alpha_val = 1.0f, beta_val = 0.0f; - alpha = &alpha_val; - beta = &beta_val; - } else if (_dtype == INFINI_DTYPE_F64) { - static const double alpha_val = 1.0, beta_val = 0.0; - alpha = &alpha_val; - beta = &beta_val; - } else { - // For F16/BF16, use float alpha/beta - static const float alpha_val = 1.0f, beta_val = 0.0f; - alpha = &alpha_val; - beta = &beta_val; - } - - CHECK_CUDNN(cudnnPoolingForward( - cudnn_handle, - _opaque->pool_desc, - alpha, - _opaque->x_desc, x, - beta, - _opaque->y_desc, y)); - - return INFINI_STATUS_SUCCESS; - }); -} - -} // namespace op::avg_pool3d::nvidia diff --git a/src/infiniop/ops/avg_pool3d/nvidia/avg_pool3d_nvidia.cuh b/src/infiniop/ops/avg_pool3d/nvidia/avg_pool3d_nvidia.cuh deleted file mode 100644 index b2a1801cb..000000000 --- a/src/infiniop/ops/avg_pool3d/nvidia/avg_pool3d_nvidia.cuh +++ /dev/null @@ -1,46 +0,0 @@ -#ifndef __AVG_POOL3D_NVIDIA_H__ -#define __AVG_POOL3D_NVIDIA_H__ - -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "../../../operator.h" -#include "../../../tensor.h" -#include - -namespace op::avg_pool3d::nvidia { - -class Descriptor final : public InfiniopDescriptor { - struct Opaque; - std::unique_ptr _opaque; - infiniDtype_t _dtype; - - // Defined out-of-line (in .cu) to avoid requiring Opaque to be complete - // in every translation unit that includes this header. - Descriptor(infiniDtype_t dtype, std::unique_ptr opaque, - infiniDevice_t device_type, int device_id); - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - void *kernel_size, - void *stride, - void *padding); - - size_t workspaceSize() const; - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::avg_pool3d::nvidia - -#endif // __AVG_POOL3D_NVIDIA_H__ diff --git a/src/infiniop/ops/avg_pool3d/operator.cc b/src/infiniop/ops/avg_pool3d/operator.cc deleted file mode 100644 index df58e1282..000000000 --- a/src/infiniop/ops/avg_pool3d/operator.cc +++ /dev/null @@ -1,163 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/avg_pool3d.h" - -#ifdef ENABLE_CPU_API -#include "cpu/avg_pool3d_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) -#include "nvidia/avg_pool3d_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/avg_pool3d_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/avg_pool3d_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAvgPool3dDescriptor( - infiniopHandle_t handle, - infiniopAvgPool3dDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - void *kernel_size, - void *stride, - void *padding) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::avg_pool3d::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - x_desc, \ - kernel_size, \ - stride, \ - padding) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAvgPool3dWorkspaceSize(infiniopAvgPool3dDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopAvgPool3d( - infiniopAvgPool3dDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, x, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyAvgPool3dDescriptor(infiniopAvgPool3dDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/awq_marlin_gemm/awq_marlin_gemm.h b/src/infiniop/ops/awq_marlin_gemm/awq_marlin_gemm.h deleted file mode 100644 index f1df9e736..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/awq_marlin_gemm.h +++ /dev/null @@ -1,57 +0,0 @@ -#ifndef AWQ_MARLIN_GEMM_H -#define AWQ_MARLIN_GEMM_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::awq_marlin_gemm::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AwqMarlinGemmInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - AwqMarlinGemmInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t a_desc, \ - infiniopTensorDescriptor_t b_desc, \ - infiniopTensorDescriptor_t b_bias_desc, \ - infiniopTensorDescriptor_t b_scales_desc, \ - infiniopTensorDescriptor_t a_scales_desc, \ - infiniopTensorDescriptor_t global_scales_desc, \ - infiniopTensorDescriptor_t b_zeros_desc, \ - infiniopTensorDescriptor_t g_idx_desc, \ - infiniopTensorDescriptor_t perm_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *c, \ - const void *a, const void *b, \ - void *b_bias, void *b_scales, void *a_scales, void *global_scales, \ - void *b_zeros, void *g_idx, void *perm, \ - int64_t b_q_type_id, bool is_k_full, bool use_atomic_add, bool use_fp32_reduce, bool is_zp_float, \ - void *stream) const; \ - }; \ - } - -#endif // AWQ_MARLIN_GEMM_H diff --git a/src/infiniop/ops/awq_marlin_gemm/core/scalar_type.hpp b/src/infiniop/ops/awq_marlin_gemm/core/scalar_type.hpp deleted file mode 100644 index 7cc0aa09d..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/core/scalar_type.hpp +++ /dev/null @@ -1,236 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include -#include - -namespace vllm { - -class ScalarType { -public: - enum NanRepr : uint8_t { - NAN_NONE = 0, - NAN_IEEE_754 = 1, - NAN_EXTD_RANGE_MAX_MIN = 2, - - NAN_REPR_ID_MAX - }; - - constexpr ScalarType(uint8_t exponent, uint8_t mantissa, bool signed_, - int32_t bias, bool finite_values_only = false, - NanRepr nan_repr = NAN_IEEE_754) - : exponent(exponent), - mantissa(mantissa), - signed_(signed_), - bias(bias), - finite_values_only(finite_values_only), - nan_repr(nan_repr) {} - - static constexpr ScalarType int_(uint8_t size_bits, int32_t bias = 0) { - return ScalarType(0, size_bits - 1, true, bias); - } - - static constexpr ScalarType uint(uint8_t size_bits, int32_t bias = 0) { - return ScalarType(0, size_bits, false, bias); - } - - static constexpr ScalarType float_IEEE754(uint8_t exponent, - uint8_t mantissa) { - return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754); - } - - static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa, - bool finite_values_only, - NanRepr nan_repr) { - return ScalarType(exponent, mantissa, true, 0, - finite_values_only, nan_repr); - } - - static inline ScalarType float_checked(uint8_t exponent, - uint8_t mantissa, - bool finite_values_only, - NanRepr nan_repr) { - if (!(nan_repr < NAN_REPR_ID_MAX)) { - throw std::runtime_error("Invalid NanRepr"); - } - - if (!(mantissa > 0 && exponent > 0)) { - throw std::runtime_error("mantissa/exponent must > 0"); - } - - if (nan_repr == NAN_IEEE_754) { - throw std::runtime_error("use float_IEEE754"); - } - - return float_(exponent, mantissa, finite_values_only, nan_repr); - } - - uint8_t const exponent; - uint8_t const mantissa; - bool const signed_; - int32_t const bias; - - bool const finite_values_only; - NanRepr const nan_repr; - - using Id = int64_t; - -private: - template - static constexpr size_t member_id_field_width() { - using T = std::decay_t; - return std::is_same::value ? 1 : sizeof(T) * 8; - } - - template - static constexpr auto reduce_members_helper(Fn f, Init val, Member member, - Rest... rest) { - auto new_val = f(val, member); - if constexpr (sizeof...(rest) > 0) { - return reduce_members_helper(f, new_val, rest...); - } else { - return new_val; - } - } - - template - constexpr auto reduce_members(Fn f, Init init) const { - return reduce_members_helper(f, init, exponent, mantissa, signed_, bias, - finite_values_only, nan_repr); - } - - template - static constexpr auto reduce_member_types(Fn f, Init init) { - constexpr auto dummy = ScalarType(0, 0, false, 0, false, NAN_NONE); - return dummy.reduce_members(f, init); - } - - static constexpr auto id_size_bits() { - return reduce_member_types( - [](int acc, auto member) -> int { - return acc + member_id_field_width(); - }, - 0); - } - -public: - constexpr Id id() const { - static_assert(id_size_bits() <= sizeof(Id) * 8, - "ScalarType id too large"); - - auto fn = [](std::pair result, auto member) { - auto [id, offset] = result; - constexpr auto bits = member_id_field_width(); - return std::pair{ - id | (int64_t(member) & ((uint64_t(1) << bits) - 1)) << offset, - offset + bits}; - }; - - return reduce_members(fn, std::pair{}).first; - } - - static constexpr ScalarType from_id(Id id) { - auto fn = [id](auto result, auto member) { - using T = decltype(member); - auto [tuple, offset] = result; - constexpr auto bits = member_id_field_width(); - auto val = static_cast((id >> offset) & ((uint64_t(1) << bits) - 1)); - return std::pair{std::tuple_cat(tuple, std::make_tuple(val)), offset + bits}; - }; - - auto [args, _] = reduce_member_types(fn, std::pair, int>{}); - - return std::apply([](auto... xs) { return ScalarType(xs...); }, args); - } - - constexpr int64_t size_bits() const { - return mantissa + exponent + (signed_ ? 1 : 0); - } - - constexpr bool is_signed() const { return signed_; } - constexpr bool is_integer() const { return exponent == 0; } - constexpr bool is_floating_point() const { return exponent > 0; } - - constexpr bool is_ieee_754() const { - return is_floating_point() && !finite_values_only && nan_repr == NAN_IEEE_754; - } - - constexpr bool has_nans() const { - return is_floating_point() && nan_repr != NAN_NONE; - } - - constexpr bool has_infs() const { - return is_floating_point() && !finite_values_only; - } - - constexpr bool has_bias() const { return bias != 0; } - - std::string str() const { - if (is_floating_point()) { - auto ret = "float" + std::to_string(size_bits()) + "_e" + std::to_string(exponent) + "m" + std::to_string(mantissa); - - if (!is_ieee_754()) { - if (finite_values_only) { - ret += "f"; - } - if (nan_repr != NAN_NONE) { - ret += "n"; - } - } - return ret; - } else { - auto ret = (signed_ ? "int" : "uint") + std::to_string(size_bits()); - if (has_bias()) { - ret += "b" + std::to_string(bias); - } - return ret; - } - } - - constexpr bool operator==(ScalarType const &other) const { - return mantissa == other.mantissa && exponent == other.exponent && bias == other.bias && signed_ == other.signed_ && finite_values_only == other.finite_values_only && nan_repr == other.nan_repr; - } -}; - -using ScalarTypeId = ScalarType::Id; - -static inline constexpr auto kS4 = ScalarType::int_(4); -static inline constexpr auto kU4 = ScalarType::uint(4); -static inline constexpr auto kU4B8 = ScalarType::uint(4, 8); -static inline constexpr auto kS8 = ScalarType::int_(8); -static inline constexpr auto kU8 = ScalarType::uint(8); -static inline constexpr auto kU8B128 = ScalarType::uint(8, 128); - -static inline constexpr auto kFE2M1f = ScalarType::float_(2, 1, true, ScalarType::NAN_NONE); -static inline constexpr auto kFE3M2f = ScalarType::float_(3, 2, true, ScalarType::NAN_NONE); -static inline constexpr auto kFE4M3fn = ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); -static inline constexpr auto kFE8M0fnu = ScalarType(8, 0, false, 0, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); -static inline constexpr auto kFE5M2 = ScalarType::float_IEEE754(5, 2); -static inline constexpr auto kFE8M7 = ScalarType::float_IEEE754(8, 7); -static inline constexpr auto kFE5M10 = ScalarType::float_IEEE754(5, 10); - -static inline constexpr auto kInt4 = kS4; -static inline constexpr auto kUint4 = kU4; -static inline constexpr auto kUint4b8 = kU4B8; -static inline constexpr auto kInt8 = kS8; -static inline constexpr auto kUint8 = kU8; -static inline constexpr auto kUint8b128 = kU8B128; - -static inline constexpr auto kFloat4_e2m1f = kFE2M1f; -static inline constexpr auto kFloat6_e3m2f = kFE3M2f; -static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn; -static inline constexpr auto kFloat8_e5m2 = kFE5M2; -static inline constexpr auto kFloat16_e8m7 = kFE8M7; -static inline constexpr auto kFloat16_e5m10 = kFE5M10; - -static inline constexpr auto kHalf = kFE5M10; -static inline constexpr auto kFloat16 = kHalf; -static inline constexpr auto kBFloat16 = kFE8M7; - -static inline constexpr auto kFloat16Id = kFloat16.id(); - -} // namespace vllm diff --git a/src/infiniop/ops/awq_marlin_gemm/core/source_location.h b/src/infiniop/ops/awq_marlin_gemm/core/source_location.h deleted file mode 100644 index 9a06fb380..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/core/source_location.h +++ /dev/null @@ -1,40 +0,0 @@ -/// \file source_location.h -/// \brief Portable `source_location` wrapper. -/// -/// Uses `std::source_location` when available (C++20), otherwise falls -/// back to a minimal stub that returns empty/zero values. - -#pragma once -#include - -/// NOTE: fallback to a minimal source_location implementation -#if defined(__cpp_lib_source_location) -#include - -using source_location_t = std::source_location; - -#else - -struct source_location_fallback { -public: - static constexpr source_location_fallback current() noexcept { - return source_location_fallback{}; - } - constexpr source_location_fallback() noexcept = default; - constexpr unsigned line() const noexcept { - return 0; - } - constexpr unsigned column() const noexcept { - return 0; - } - constexpr const char *file_name() const noexcept { - return ""; - } - constexpr const char *function_name() const noexcept { - return ""; - } -}; - -using source_location_t = source_location_fallback; - -#endif diff --git a/src/infiniop/ops/awq_marlin_gemm/core/utils.h b/src/infiniop/ops/awq_marlin_gemm/core/utils.h deleted file mode 100644 index c07956d09..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/core/utils.h +++ /dev/null @@ -1,212 +0,0 @@ -/// \file utils.h -/// \brief Host-side C++ utilities used by JIT kernel wrappers. -/// -/// Provides: -/// - `DebugInfo` - wraps `std::source_location` for error reporting. -/// - `RuntimeCheck` - runtime assertion with formatted error messages. -/// - `Panic` - unconditional abort with formatted error messages. -/// - `pointer::offset` - safe void-pointer arithmetic (host side). -/// - `div_ceil` - integer ceiling division. -/// - `dtype_bytes` - byte width of a `DLDataType`. -/// - `irange` - Python-style integer range for range-for loops. - -#pragma once - -// ref: https://forums.developer.nvidia.com/t/c-20s-source-location-compilation-error-when-using-nvcc-12-1/258026/3 -#ifdef __CUDACC__ -#include -#if CUDA_VERSION <= 12010 - -#pragma push_macro("__cpp_consteval") -#pragma push_macro("_NODISCARD") -#pragma push_macro("__builtin_LINE") - -#pragma clang diagnostic push -#pragma clang diagnostic ignored "-Wbuiltin-macro-redefined" -#define __cpp_consteval 201811L -#pragma clang diagnostic pop - -#ifdef _NODISCARD -#undef _NODISCARD -#define _NODISCARD -#endif - -#define consteval constexpr - -#include "source_location.h" - -#undef consteval -#pragma pop_macro("__cpp_consteval") -#pragma pop_macro("_NODISCARD") -#else // __CUDACC__ && CUDA_VERSION > 12010 -#include "source_location.h" -#endif -#else // no __CUDACC__ -#include "source_location.h" -#endif - -#include "../dlpack/dlpack.h" -#include -// #include -#include -#include -// #include -#include -#include - -namespace host { - -template -inline constexpr bool dependent_false_v = false; - -/// \brief Source-location wrapper for debug/error messages. -struct DebugInfo : public source_location_t { - DebugInfo(source_location_t loc = source_location_t::current()) : source_location_t(loc) {} -}; - -/// \brief Exception type thrown by `RuntimeCheck` and `Panic`. -struct PanicError : public std::runtime_error { -public: - explicit PanicError(std::string msg) : runtime_error(msg), m_message(std::move(msg)) {} - auto root_cause() const -> std::string_view { - const auto str = std::string_view{m_message}; - const auto pos = str.find(": "); - return pos == std::string_view::npos ? str : str.substr(pos + 2); - } - -private: - std::string m_message; -}; - -/// \brief Unconditionally abort with a formatted error message. -template -[[noreturn]] inline auto panic(DebugInfo location, Args &&...args) -> void { - std::ostringstream os; - os << "Runtime check failed at " << location.file_name() << ":" << location.line(); - if constexpr (sizeof...(args) > 0) { - os << ": "; - (os << ... << std::forward(args)); - } else { - os << " in " << location.function_name(); - } - throw PanicError(std::move(os).str()); -} - -/** - * \brief Runtime assertion: panics with a formatted message when `condition` - * is false. Extra `args` are streamed to the error message. - * - * Example: - * \code - * RuntimeCheck(n > 0, "n must be positive, got ", n); - * \endcode - */ -template -struct RuntimeCheck { - template - explicit RuntimeCheck(Cond &&condition, Args &&...args, DebugInfo location = {}) { - if (condition) { - return; - } - [[unlikely]] ::host::panic(location, std::forward(args)...); - } - template - explicit RuntimeCheck(DebugInfo location, Cond &&condition, Args &&...args) { - if (condition) { - return; - } - [[unlikely]] ::host::panic(location, std::forward(args)...); - } -}; - -template -struct Panic { - explicit Panic(Args &&...args, DebugInfo location = {}) { - ::host::panic(location, std::forward(args)...); - } - explicit Panic(DebugInfo location, Args &&...args) { - ::host::panic(location, std::forward(args)...); - } - [[noreturn]] ~Panic() { - std::terminate(); - } -}; - -template -explicit RuntimeCheck(Cond &&, Args &&...) -> RuntimeCheck; - -template -explicit RuntimeCheck(DebugInfo, Cond &&, Args &&...) -> RuntimeCheck; - -template -explicit Panic(Args &&...) -> Panic; - -template -explicit Panic(DebugInfo, Args &&...) -> Panic; - -namespace pointer { - -// we only allow void * pointer arithmetic for safety - -template ::value && ...)>> -inline auto offset(void *ptr, U... offset) -> void * { - return static_cast(ptr) + (... + offset); -} - -template ::value && ...)>> -inline auto offset(const void *ptr, U... offset) -> const void * { - return static_cast(ptr) + (... + offset); -} - -} // namespace pointer - -/// \brief Integer ceiling division: ceil(a / b). -template -inline constexpr auto div_ceil(T a, U b) { - static_assert(std::is_integral::value, "T must be integral"); - static_assert(std::is_integral::value, "U must be integral"); - return (a + b - 1) / b; -} - -/// \brief Returns the byte width of a DLPack data type. -inline auto dtype_bytes(DLDataType dtype) -> std::size_t { - return static_cast(dtype.bits / 8); -} - -template -struct IntegerRange { - T start_; - T end_; - - struct Iterator { - T value; - - T operator*() const { return value; } - Iterator &operator++() { - ++value; - return *this; - } - bool operator!=(const Iterator &other) const { - return value != other.value; - } - }; - - Iterator begin() const { return {start_}; } - Iterator end() const { return {end_}; } -}; - -/// Python-style integer range: irange(n) -> [0, n) -template -IntegerRange irange(T end) { - return {0, end}; -} - -/// Python-style integer range: irange(start, end) -> [start, end) -template -IntegerRange irange(T start, T end) { - return {start, end}; -} - -} // namespace host diff --git a/src/infiniop/ops/awq_marlin_gemm/dlpack/dlpack.h b/src/infiniop/ops/awq_marlin_gemm/dlpack/dlpack.h deleted file mode 100644 index 3f4ed201a..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/dlpack/dlpack.h +++ /dev/null @@ -1,647 +0,0 @@ -/*! - * Copyright (c) 2017 - by Contributors - * \file dlpack.h - * \brief The common header of DLPack. - */ -#ifndef DLPACK_DLPACK_H_ -#define DLPACK_DLPACK_H_ - -/** - * \brief Compatibility with C++ - */ -#ifdef __cplusplus -#define DLPACK_EXTERN_C extern "C" -#else -#define DLPACK_EXTERN_C -#endif - -/*! \brief The current major version of dlpack */ -#define DLPACK_MAJOR_VERSION 1 - -/*! \brief The current minor version of dlpack */ -#define DLPACK_MINOR_VERSION 3 - -/*! \brief DLPACK_DLL prefix for windows */ -#ifdef _WIN32 -#ifdef DLPACK_EXPORTS -#define DLPACK_DLL __declspec(dllexport) -#else -#define DLPACK_DLL __declspec(dllimport) -#endif -#else -#define DLPACK_DLL -#endif - -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -/*! - * \brief The DLPack version. - * - * A change in major version indicates that we have changed the - * data layout of the ABI - DLManagedTensorVersioned. - * - * A change in minor version indicates that we have added new - * code, such as a new device type, but the ABI is kept the same. - * - * If an obtained DLPack tensor has a major version that disagrees - * with the version number specified in this header file - * (i.e. major != DLPACK_MAJOR_VERSION), the consumer must call the deleter - * (and it is safe to do so). It is not safe to access any other fields - * as the memory layout will have changed. - * - * In the case of a minor version mismatch, the tensor can be safely used as - * long as the consumer knows how to interpret all fields. Minor version - * updates indicate the addition of enumeration values. - */ -typedef struct { - /*! \brief DLPack major version. */ - uint32_t major; - /*! \brief DLPack minor version. */ - uint32_t minor; -} DLPackVersion; - -/*! - * \brief The device type in DLDevice. - */ -#ifdef __cplusplus -typedef enum : int32_t { -#else -typedef enum { -#endif - /*! \brief CPU device */ - kDLCPU = 1, - /*! \brief CUDA GPU device */ - kDLCUDA = 2, - /*! - * \brief Pinned CUDA CPU memory by cudaMallocHost - */ - kDLCUDAHost = 3, - /*! \brief OpenCL devices. */ - kDLOpenCL = 4, - /*! \brief Vulkan buffer for next generation graphics. */ - kDLVulkan = 7, - /*! \brief Metal for Apple GPU. */ - kDLMetal = 8, - /*! \brief Verilog simulator buffer */ - kDLVPI = 9, - /*! \brief ROCm GPUs for AMD GPUs */ - kDLROCM = 10, - /*! - * \brief Pinned ROCm CPU memory allocated by hipMallocHost - */ - kDLROCMHost = 11, - /*! - * \brief Reserved extension device type, - * used for quickly test extension device - * The semantics can differ depending on the implementation. - */ - kDLExtDev = 12, - /*! - * \brief CUDA managed/unified memory allocated by cudaMallocManaged - */ - kDLCUDAManaged = 13, - /*! - * \brief Unified shared memory allocated on a oneAPI non-partititioned - * device. Call to oneAPI runtime is required to determine the device - * type, the USM allocation type and the sycl context it is bound to. - * - */ - kDLOneAPI = 14, - /*! \brief GPU support for next generation WebGPU standard. */ - kDLWebGPU = 15, - /*! \brief Qualcomm Hexagon DSP */ - kDLHexagon = 16, - /*! \brief Microsoft MAIA devices */ - kDLMAIA = 17, - /*! \brief AWS Trainium */ - kDLTrn = 18, -} DLDeviceType; - -/*! - * \brief A Device for Tensor and operator. - */ -typedef struct { - /*! \brief The device type used in the device. */ - DLDeviceType device_type; - /*! - * \brief The device index. - * For vanilla CPU memory, pinned memory, or managed memory, this is set to 0. - */ - int32_t device_id; -} DLDevice; - -/*! - * \brief The type code options DLDataType. - */ -typedef enum { - /*! \brief signed integer */ - kDLInt = 0U, - /*! \brief unsigned integer */ - kDLUInt = 1U, - /*! \brief IEEE floating point */ - kDLFloat = 2U, - /*! - * \brief Opaque handle type, reserved for testing purposes. - * Frameworks need to agree on the handle data type for the exchange to be well-defined. - */ - kDLOpaqueHandle = 3U, - /*! \brief bfloat16 */ - kDLBfloat = 4U, - /*! - * \brief complex number - * (C/C++/Python layout: compact struct per complex number) - */ - kDLComplex = 5U, - /*! \brief boolean */ - kDLBool = 6U, - /*! \brief FP8 data types */ - kDLFloat8_e3m4 = 7U, - kDLFloat8_e4m3 = 8U, - kDLFloat8_e4m3b11fnuz = 9U, - kDLFloat8_e4m3fn = 10U, - kDLFloat8_e4m3fnuz = 11U, - kDLFloat8_e5m2 = 12U, - kDLFloat8_e5m2fnuz = 13U, - kDLFloat8_e8m0fnu = 14U, - /*! \brief FP6 data types - * Setting bits != 6 is currently unspecified, and the producer must ensure it is set - * while the consumer must stop importing if the value is unexpected. - */ - kDLFloat6_e2m3fn = 15U, - kDLFloat6_e3m2fn = 16U, - /*! \brief FP4 data types - * Setting bits != 4 is currently unspecified, and the producer must ensure it is set - * while the consumer must stop importing if the value is unexpected. - */ - kDLFloat4_e2m1fn = 17U, -} DLDataTypeCode; - -/*! - * \brief The data type the tensor can hold. The data type is assumed to follow the - * native endian-ness. An explicit error message should be raised when attempting to - * export an array with non-native endianness - * - * Examples - * - float: type_code = 2, bits = 32, lanes = 1 - * - float4(vectorized 4 float): type_code = 2, bits = 32, lanes = 4 - * - int8: type_code = 0, bits = 8, lanes = 1 - * - std::complex: type_code = 5, bits = 64, lanes = 1 - * - bool: type_code = 6, bits = 8, lanes = 1 (as per common array library convention, the underlying storage size of bool is 8 bits) - * - float8_e4m3: type_code = 8, bits = 8, lanes = 1 (packed in memory) - * - float6_e3m2fn: type_code = 16, bits = 6, lanes = 1 (packed in memory) - * - float4_e2m1fn: type_code = 17, bits = 4, lanes = 1 (packed in memory) - * - * When a sub-byte type is packed, DLPack requires the data to be in little bit-endian, i.e., - * for a packed data set D ((D >> (i * bits)) && bit_mask) stores the i-th element. - */ -typedef struct { - /*! - * \brief Type code of base types. - * We keep it uint8_t instead of DLDataTypeCode for minimal memory - * footprint, but the value should be one of DLDataTypeCode enum values. - * */ - uint8_t code; - /*! - * \brief Number of bits, common choices are 8, 16, 32. - */ - uint8_t bits; - /*! \brief Number of lanes in the type, used for vector types. */ - uint16_t lanes; -} DLDataType; - -/*! - * \brief Plain C Tensor object, does not manage memory. - */ -typedef struct { - /*! - * \brief The data pointer points to the allocated data. This will be CUDA - * device pointer or cl_mem handle in OpenCL. It may be opaque on some device - * types. This pointer is always aligned to 256 bytes as in CUDA. The - * `byte_offset` field should be used to point to the beginning of the data. - * - * Note that as of Nov 2021, multiple libraries (CuPy, PyTorch, TensorFlow, - * TVM, perhaps others) do not adhere to this 256 byte alignment requirement - * on CPU/CUDA/ROCm, and always use `byte_offset=0`. This must be fixed - * (after which this note will be updated); at the moment it is recommended - * to not rely on the data pointer being correctly aligned. - * - * For given DLTensor, the size of memory required to store the contents of - * data is calculated as follows: - * - * \code{.c} - * static inline size_t GetDataSize(const DLTensor* t) { - * size_t size = 1; - * for (tvm_index_t i = 0; i < t->ndim; ++i) { - * size *= t->shape[i]; - * } - * size *= (t->dtype.bits * t->dtype.lanes + 7) / 8; - * return size; - * } - * \endcode - * - * Note that if the tensor is of size zero, then the data pointer should be - * set to `NULL`. - */ - void *data; - /*! \brief The device of the tensor */ - DLDevice device; - /*! \brief Number of dimensions */ - int32_t ndim; - /*! \brief The data type of the pointer*/ - DLDataType dtype; - /*! - * \brief The shape of the tensor - * - * When ndim == 0, shape can be set to NULL. - */ - int64_t *shape; - /*! - * \brief strides of the tensor (in number of elements, not bytes), - * can not be NULL if ndim != 0, must points to - * an array of ndim elements that specifies the strides, - * so consumer can always rely on strides[dim] being valid for 0 <= dim < ndim. - * - * When ndim == 0, strides can be set to NULL. - * - * \note Before DLPack v1.2, strides can be NULL to indicate contiguous data. - * This is not allowed in DLPack v1.2 and later. The rationale - * is to simplify the consumer handling. - */ - int64_t *strides; - /*! \brief The offset in bytes to the beginning pointer to data */ - uint64_t byte_offset; -} DLTensor; - -/*! - * \brief C Tensor object, manage memory of DLTensor. This data structure is - * intended to facilitate the borrowing of DLTensor by another framework. It is - * not meant to transfer the tensor. When the borrowing framework doesn't need - * the tensor, it should call the deleter to notify the host that the resource - * is no longer needed. - * - * \note This data structure is used as Legacy DLManagedTensor - * in DLPack exchange and is deprecated after DLPack v0.8 - * Use DLManagedTensorVersioned instead. - * This data structure may get renamed or deleted in future versions. - * - * \sa DLManagedTensorVersioned - */ -typedef struct DLManagedTensor { - /*! \brief DLTensor which is being memory managed */ - DLTensor dl_tensor; - /*! \brief the context of the original host framework of DLManagedTensor in - * which DLManagedTensor is used in the framework. It can also be NULL. - */ - void *manager_ctx; - /*! - * \brief Destructor - this should be called - * to destruct the manager_ctx which backs the DLManagedTensor. It can be - * NULL if there is no way for the caller to provide a reasonable destructor. - * The destructor deletes the argument self as well. - */ - void (*deleter)(struct DLManagedTensor *self); -} DLManagedTensor; - -// bit masks used in the DLManagedTensorVersioned - -/*! \brief bit mask to indicate that the tensor is read only. */ -#define DLPACK_FLAG_BITMASK_READ_ONLY (1UL << 0UL) - -/*! - * \brief bit mask to indicate that the tensor is a copy made by the producer. - * - * If set, the tensor is considered solely owned throughout its lifetime by the - * consumer, until the producer-provided deleter is invoked. - */ -#define DLPACK_FLAG_BITMASK_IS_COPIED (1UL << 1UL) - -/*! - * \brief bit mask to indicate that whether a sub-byte type is packed or padded. - * - * The default for sub-byte types (ex: fp4/fp6) is assumed packed. This flag can - * be set by the producer to signal that a tensor of sub-byte type is padded. - */ -#define DLPACK_FLAG_BITMASK_IS_SUBBYTE_TYPE_PADDED (1UL << 2UL) - -/*! - * \brief A versioned and managed C Tensor object, manage memory of DLTensor. - * - * This data structure is intended to facilitate the borrowing of DLTensor by - * another framework. It is not meant to transfer the tensor. When the borrowing - * framework doesn't need the tensor, it should call the deleter to notify the - * host that the resource is no longer needed. - * - * \note This is the current standard DLPack exchange data structure. - */ -typedef struct DLManagedTensorVersioned { - /*! - * \brief The API and ABI version of the current managed Tensor - */ - DLPackVersion version; - /*! - * \brief the context of the original host framework. - * - * Stores DLManagedTensorVersioned is used in the - * framework. It can also be NULL. - */ - void *manager_ctx; - /*! - * \brief Destructor. - * - * This should be called to destruct manager_ctx which holds the DLManagedTensorVersioned. - * It can be NULL if there is no way for the caller to provide a reasonable - * destructor. The destructor deletes the argument self as well. - */ - void (*deleter)(struct DLManagedTensorVersioned *self); - /*! - * \brief Additional bitmask flags information about the tensor. - * - * By default the flags should be set to 0. - * - * \note Future ABI changes should keep everything until this field - * stable, to ensure that deleter can be correctly called. - * - * \sa DLPACK_FLAG_BITMASK_READ_ONLY - * \sa DLPACK_FLAG_BITMASK_IS_COPIED - */ - uint64_t flags; - /*! \brief DLTensor which is being memory managed */ - DLTensor dl_tensor; -} DLManagedTensorVersioned; - -//---------------------------------------------------------------------- -// DLPack `__dlpack_c_exchange_api__` fast exchange protocol definitions -//---------------------------------------------------------------------- -/*! - * \brief Request a producer library to create a new tensor. - * - * Create a new `DLManagedTensorVersioned` within the context of the producer - * library. The allocation is defined via the prototype DLTensor. - * - * This function is exposed by the framework through the DLPackExchangeAPI. - * - * \param prototype The prototype DLTensor. Only the dtype, ndim, shape, - * and device fields are used. - * \param out The output DLManagedTensorVersioned. - * \param error_ctx Context for `SetError`. - * \param SetError The function to set the error. - * \return The owning DLManagedTensorVersioned* or NULL on failure. - * SetError is called exactly when NULL is returned (the implementer - * must ensure this). - * \note - As a C function, must not thrown C++ exceptions. - * - Error propagation via SetError to avoid any direct need - * of Python API. Due to this `SetError` may have to ensure the GIL is - * held since it will presumably set a Python error. - * - * \sa DLPackExchangeAPI - */ -typedef int (*DLPackManagedTensorAllocator)( // - DLTensor *prototype, DLManagedTensorVersioned **out, void *error_ctx, // - void (*SetError)(void *error_ctx, const char *kind, const char *message) // -); - -/*! - * \brief Exports a PyObject* Tensor/NDArray to a DLManagedTensorVersioned. - * - * This function does not perform any stream synchronization. The consumer should query - * DLPackCurrentWorkStream to get the current work stream and launch kernels on it. - * - * This function is exposed by the framework through the DLPackExchangeAPI. - * - * \param py_object The Python object to convert. Must have the same type - * as the one the `DLPackExchangeAPI` was discovered from. - * \param out The output DLManagedTensorVersioned. - * \return The owning DLManagedTensorVersioned* or NULL on failure with a - * Python exception set. If the data cannot be described using DLPack - * this should be a BufferError if possible. - * \note - As a C function, must not thrown C++ exceptions. - * - * \sa DLPackExchangeAPI, DLPackCurrentWorkStream - */ -typedef int (*DLPackManagedTensorFromPyObjectNoSync)( // - void *py_object, // - DLManagedTensorVersioned **out // -); - -/*! - * \brief Exports a PyObject* Tensor/NDArray to a provided DLTensor. - * - * This function provides a faster interface for temporary, non-owning, exchange. - * The producer (implementer) still owns the memory of data, strides, shape. - * The liveness of the DLTensor and the data it views is only guaranteed until - * control is returned. - * - * This function currently assumes that the producer (implementer) can fill - * in the DLTensor shape and strides without the need for temporary allocations. - * - * This function does not perform any stream synchronization. The consumer should query - * DLPackCurrentWorkStream to get the current work stream and launch kernels on it. - * - * This function is exposed by the framework through the DLPackExchangeAPI. - * - * \param py_object The Python object to convert. Must have the same type - * as the one the `DLPackExchangeAPI` was discovered from. - * \param out The output DLTensor, whose space is pre-allocated on stack. - * \return 0 on success, -1 on failure with a Python exception set. - * \note - As a C function, must not thrown C++ exceptions. - * - * \sa DLPackExchangeAPI, DLPackCurrentWorkStream - */ -typedef int (*DLPackDLTensorFromPyObjectNoSync)( // - void *py_object, // - DLTensor *out // -); - -/*! - * \brief Obtain the current work stream of a device. - * - * Obtain the current work stream of a device from the producer framework. - * For example, it should map to torch.cuda.current_stream in PyTorch. - * - * When device_type is kDLCPU, the consumer do not have to query the stream - * and the producer can simply return NULL when queried. - * The consumer do not have to do anything on stream sync or setting. - * So CPU only framework can just provide a dummy implementation that - * always set out_current_stream[0] to NULL. - * - * \param device_type The device type. - * \param device_id The device id. - * \param out_current_stream The output current work stream. - * - * \return 0 on success, -1 on failure with a Python exception set. - * \note - As a C function, must not thrown C++ exceptions. - * - * \sa DLPackExchangeAPI - */ -typedef int (*DLPackCurrentWorkStream)( // - DLDeviceType device_type, // - int32_t device_id, // - void **out_current_stream // -); - -/*! - * \brief Imports a DLManagedTensorVersioned to a PyObject* Tensor/NDArray. - * - * Convert an owning DLManagedTensorVersioned* to the Python tensor of the - * producer (implementer) library with the correct type. - * - * This function does not perform any stream synchronization. - * - * This function is exposed by the framework through the DLPackExchangeAPI. - * - * \param tensor The DLManagedTensorVersioned to convert the ownership of the - * tensor is stolen. - * \param out_py_object The output Python object. - * \return 0 on success, -1 on failure with a Python exception set. - * - * \sa DLPackExchangeAPI - */ -typedef int (*DLPackManagedTensorToPyObjectNoSync)( // - DLManagedTensorVersioned *tensor, // - void **out_py_object // -); - -/*! - * \brief DLPackExchangeAPI stable header. - * \sa DLPackExchangeAPI - */ -typedef struct DLPackExchangeAPIHeader { - /*! - * \brief The provided DLPack version the consumer must check major version - * compatibility before using this struct. - */ - DLPackVersion version; - /*! - * \brief Optional pointer to an older DLPackExchangeAPI in the chain. - * - * It must be NULL if the framework does not support older versions. - * If the current major version is larger than the one supported by the - * consumer, the consumer may walk this to find an earlier supported version. - * - * \sa DLPackExchangeAPI - */ - struct DLPackExchangeAPIHeader *prev_api; -} DLPackExchangeAPIHeader; - -/*! - * \brief Framework-specific function pointers table for DLPack exchange. - * - * Additionally to `__dlpack__()` we define a C function table sharable by - * - * Python implementations via `__dlpack_c_exchange_api__`. - * This attribute must be set on the type as a Python PyCapsule - * with name "dlpack_exchange_api". - * - * A consumer library may use a pattern such as: - * - * \code - * - * PyObject *api_capsule = PyObject_GetAttrString( - * (PyObject *)Py_TYPE(tensor_obj), "__dlpack_c_exchange_api__") - * ); - * if (api_capsule == NULL) { goto handle_error; } - * MyDLPackExchangeAPI *api = (MyDLPackExchangeAPI *)PyCapsule_GetPointer( - * api_capsule, "dlpack_exchange_api" - * ); - * Py_DECREF(api_capsule); - * if (api == NULL) { goto handle_error; } - * - * \endcode - * - * Note that this must be defined on the type. The consumer should look up the - * attribute on the type and may cache the result for each unique type. - * - * The precise API table is given by: - * \code - * struct MyDLPackExchangeAPI : public DLPackExchangeAPI { - * MyDLPackExchangeAPI() { - * header.version.major = DLPACK_MAJOR_VERSION; - * header.version.minor = DLPACK_MINOR_VERSION; - * header.prev_version_api = nullptr; - * - * managed_tensor_allocator = MyDLPackManagedTensorAllocator; - * managed_tensor_from_py_object_no_sync = MyDLPackManagedTensorFromPyObjectNoSync; - * managed_tensor_to_py_object_no_sync = MyDLPackManagedTensorToPyObjectNoSync; - * dltensor_from_py_object_no_sync = MyDLPackDLTensorFromPyObjectNoSync; - * current_work_stream = MyDLPackCurrentWorkStream; - * } - * - * static const DLPackExchangeAPI* Global() { - * static MyDLPackExchangeAPI inst; - * return &inst; - * } - * }; - * \endcode - * - * Guidelines for leveraging DLPackExchangeAPI: - * - * There are generally two kinds of consumer needs for DLPack exchange: - * - N0: library support, where consumer.kernel(x, y, z) would like to run a kernel - * with the data from x, y, z. The consumer is also expected to run the kernel with the same - * stream context as the producer. For example, when x, y, z is torch.Tensor, - * consumer should query exchange_api->current_work_stream to get the - * current stream and launch the kernel with the same stream. - * This setup is necessary for no synchronization in kernel launch and maximum compatibility - * with CUDA graph capture in the producer. - * This is the desirable behavior for library extension support for frameworks like PyTorch. - * - N1: data ingestion and retention - * - * Note that obj.__dlpack__() API should provide useful ways for N1. - * The primary focus of the current DLPackExchangeAPI is to enable faster exchange N0 - * with the support of the function pointer current_work_stream. - * - * Array/Tensor libraries should statically create and initialize this structure - * then return a pointer to DLPackExchangeAPI as an int value in Tensor/Array. - * The DLPackExchangeAPI* must stay alive throughout the lifetime of the process. - * - * One simple way to do so is to create a static instance of DLPackExchangeAPI - * within the framework and return a pointer to it. The following code - * shows an example to do so in C++. It should also be reasonably easy - * to do so in other languages. - */ -typedef struct DLPackExchangeAPI { - /*! - * \brief The header that remains stable across versions. - */ - DLPackExchangeAPIHeader header; - /*! - * \brief Producer function pointer for DLPackManagedTensorAllocator - * This function must not be NULL. - * \sa DLPackManagedTensorAllocator - */ - DLPackManagedTensorAllocator managed_tensor_allocator; - /*! - * \brief Producer function pointer for DLPackManagedTensorFromPyObject - * This function must be not NULL. - * \sa DLPackManagedTensorFromPyObject - */ - DLPackManagedTensorFromPyObjectNoSync managed_tensor_from_py_object_no_sync; - /*! - * \brief Producer function pointer for DLPackManagedTensorToPyObject - * This function must be not NULL. - * \sa DLPackManagedTensorToPyObjectNoSync - */ - DLPackManagedTensorToPyObjectNoSync managed_tensor_to_py_object_no_sync; - /*! - * \brief Producer function pointer for DLPackDLTensorFromPyObject - * This function can be NULL when the producer does not support this function. - * \sa DLPackDLTensorFromPyObjectNoSync - */ - DLPackDLTensorFromPyObjectNoSync dltensor_from_py_object_no_sync; - /*! - * \brief Producer function pointer for DLPackCurrentWorkStream - * This function must be not NULL. - * \sa DLPackCurrentWorkStream - */ - DLPackCurrentWorkStream current_work_stream; -} DLPackExchangeAPI; - -#ifdef __cplusplus -} // DLPACK_EXTERN_C -#endif -#endif // DLPACK_DLPACK_H_ diff --git a/src/infiniop/ops/awq_marlin_gemm/info.h b/src/infiniop/ops/awq_marlin_gemm/info.h deleted file mode 100644 index 9c74b413f..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/info.h +++ /dev/null @@ -1,61 +0,0 @@ -#ifndef __AWQ_MARLIN_GEMM_INFO_H__ -#define __AWQ_MARLIN_GEMM_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include "marlin/marlin.cuh" -#include - -#include - -namespace op::awq_marlin_gemm { - -class AwqMarlinGemmInfo { - AwqMarlinGemmInfo() = default; - -public: - infiniDtype_t a_dtype, b_dtype, c_dtype, s_dtype; - size_t size_m, size_k, size_n; - int num_groups; - size_t b_q_size_0, b_q_size_1, b_zeros_size_1; - ptrdiff_t a_stride_0; - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t b_bias_desc, - infiniopTensorDescriptor_t b_scales_desc, - infiniopTensorDescriptor_t a_scales_desc, - infiniopTensorDescriptor_t global_scales_desc, - infiniopTensorDescriptor_t b_zeros_desc, - infiniopTensorDescriptor_t g_idx_desc, - infiniopTensorDescriptor_t perm_desc) { - CHECK_OR_RETURN( - out_desc != nullptr && a_desc != nullptr && b_desc != nullptr && b_scales_desc != nullptr, - INFINI_STATUS_NULL_POINTER); - const infiniDtype_t a_dtype = a_desc->dtype(); - const infiniDtype_t b_dtype = b_desc->dtype(); - const infiniDtype_t c_dtype = out_desc->dtype(); - const infiniDtype_t s_dtype = b_scales_desc->dtype(); - - size_t size_m = a_desc->dim(0); - size_t size_k = a_desc->dim(1); - size_t size_n = out_desc->dim(1); - - int num_groups = static_cast(b_scales_desc->dim(0)); - size_t b_q_size_0 = b_desc->dim(0); - size_t b_q_size_1 = b_desc->dim(1); - size_t b_zeros_size_1 = b_zeros_desc->dim(1); - ptrdiff_t a_stride_0 = a_desc->strides()[0]; - - return utils::Result( - AwqMarlinGemmInfo{a_dtype, b_dtype, c_dtype, s_dtype, - size_m, size_k, size_n, - num_groups, b_q_size_0, b_q_size_1, b_zeros_size_1, a_stride_0}); - } -}; - -} // namespace op::awq_marlin_gemm - -#endif // __AWQ_MARLIN_GEMM_INFO_H__ diff --git a/src/infiniop/ops/awq_marlin_gemm/marlin/dequant.h b/src/infiniop/ops/awq_marlin_gemm/marlin/dequant.h deleted file mode 100644 index 1340b6bbe..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/marlin/dequant.h +++ /dev/null @@ -1,601 +0,0 @@ -/* -Fast Dequantization (Converting INT4/INT8/FP4/FP8 to FP16/BF16) - -The process of fast dequantization can be summarized as a combination -of bitwise operations and floating-point computations: - -weight =>(bit_op / bitwise operations)=> -f16_value =>(flop / floating-point computation)=> -dequantized_weight - -Since the dequantized weights typically require subtracting the zero point and -applying a scale factor, the floating-point computation step can be fused with -the zero-point subtraction and scaling operations. - -The following are the parts that need to be modified for the fused operation -of zero-point subtraction and scaling. - -## INT4 => FP16/BF16 or INT8 => FP16 - -The floating-point computation is `__hsub2` - -If has zero points: - - flop(bit_op(weight)) - flop(bit_op(zp)) - = sub(bit_op(weight), bias) - sub(bit_op(zp), bias) - = bit_op(weight) - bit_op(zp) - -so we don't need additional modification. - -If has float zero points: - - flop(bit_op(weight)) - fzp - = sub(bit_op(weight), bias) - fzp - = bit_op(weight) - (fzp + bias) - -where the `fzp + bias` can be computed at weight loading. But this -may have accuracy issue, so we should not use this in most cases. - -If has not zero points: - - scale(flop(bit_op(weight))) - = scale(sub(bit_op(weight), bias)) - = scale(bit_op(weight)) - scale(bias) - = fma(bit_op(weight), scale_factor, scale(bias)) - -where the `scale(bias)` can be cached. But this may have accuracy issue, -so we should not use this in most cases. - - -## INT8 => BF16 - -INT8 => BF16 is a special case, it use byte_perm instead of flop. -We cannot fused byte_perm with scaling. - - -## FP4/FP8 => FP16/BF16 - - scale(flop(bit_op(weight))) - = scale(mul(bit_op(weight), multiplier)) - = mul(bit_op(weight), scale_factor * multiplier) - -where `scale_factor * multiplier` can be computed at weight loading. - -*/ - -#include "marlin_dtypes.cuh" - -namespace MARLIN_NAMESPACE_NAME { - -#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 750 -// Lookup-table based 3-input logical operation; explicitly used for -// dequantization as the compiler does not seem to automatically recognize it in -// all cases. -template -__device__ inline int lop3(int a, int b, int c) { - int res; - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(res) - : "r"(a), "r"(b), "r"(c), "n"(lut)); - return res; -} - -// Constructs destination register by taking bytes from 2 sources (based on -// mask) -template -__device__ inline uint32_t prmt(uint32_t a) { - uint32_t res; - asm volatile("prmt.b32 %0, %1, %2, %3;\n" - : "=r"(res) - : "r"(a), "n"(start_byte), "n"(mask)); - return res; -} - -template -__device__ inline void dequant(int q, scalar_t2 *frag_b); - -// -// Efficiently dequantize 4bit values packed in an int32 value into a full -// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below, -// with some small changes: -// - FP16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287 -// - BF16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385 -// -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - const int MASK = 0x000f000f; - const int EX = 0x64006400; - // Guarantee that the `(a & b) | c` operations are LOP3s. - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - q >>= 4; - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - - frag_b[0] = *reinterpret_cast(&lo); - frag_b[1] = *reinterpret_cast(&hi); -} - -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - const int LO = 0x000f000f; - const int HI = 0x00f000f0; - const int EX = 0x64006400; - // Guarantee that the `(a & b) | c` operations are LOP3s. - // clang-format off - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); - // clang-format on - // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point - // directly into `SUB` and `ADD`. - const int SUB = 0x64086408; - const int MUL = 0x2c002c00; - const int ADD = 0xd480d480; - frag_b[0] = __hsub2(*reinterpret_cast(&lo), - *reinterpret_cast(&SUB)); - frag_b[1] = __hfma2(*reinterpret_cast(&hi), - *reinterpret_cast(&MUL), - *reinterpret_cast(&ADD)); -} - -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - dequant(q, frag_b); -} - -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - const int LO = 0x000f000f; - const int HI = 0x00f000f0; - const int EX = 0x64006400; - // Guarantee that the `(a & b) | c` operations are LOP3s. - // clang-format off - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); - // clang-format on - // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point - // directly into `SUB` and `ADD`. - const int SUB = 0x64006400; - const int MUL = 0x2c002c00; - const int ADD = 0xd400d400; - frag_b[0] = __hsub2(*reinterpret_cast(&lo), - *reinterpret_cast(&SUB)); - frag_b[1] = __hfma2(*reinterpret_cast(&hi), - *reinterpret_cast(&MUL), - *reinterpret_cast(&ADD)); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - static constexpr uint32_t MASK = 0x000f000f; - static constexpr uint32_t EX = 0x43004300; - - // Guarantee that the `(a & b) | c` operations are LOP3s. - // clang-format off - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - q >>= 4; - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - // clang-format on - - frag_b[0] = *reinterpret_cast(&lo); - frag_b[1] = *reinterpret_cast(&hi); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); - - static constexpr uint32_t SUB = 0x43084308; - - frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&SUB)); - frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&SUB)); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); - - static constexpr uint32_t SUB = 0x43004300; - - frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&SUB)); - frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&SUB)); -} - -// -// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or -// bf16 Reference: -// - FP16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85 -// - BF16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175 -// -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - static constexpr uint32_t mask_for_elt_01 = 0x5250; - static constexpr uint32_t mask_for_elt_23 = 0x5351; - static constexpr uint32_t start_byte_for_fp16 = 0x64646464; - - uint32_t lo = prmt(q); - uint32_t hi = prmt(q); - - frag_b[0] = *reinterpret_cast(&lo); - frag_b[1] = *reinterpret_cast(&hi); -} - -template <> -__device__ inline void dequant( - int q, half2 *frag_b) { - dequant(q, frag_b); - - static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480; - frag_b[0] = __hsub2(frag_b[0], - *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); - frag_b[1] = __hsub2(frag_b[1], - *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); -} - -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - dequant(q, frag_b); -} - -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - dequant(q, frag_b); - - static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400; - frag_b[0] = __hsub2(frag_b[0], - *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); - frag_b[1] = __hsub2(frag_b[1], - *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - float fp32_intermediates[4]; - uint32_t *fp32_intermediates_casted = reinterpret_cast(fp32_intermediates); - - static constexpr uint32_t fp32_base = 0x4B000000; - fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); - fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); - fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); - fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); - - fp32_intermediates[0] -= 8388736.f; - fp32_intermediates[1] -= 8388736.f; - fp32_intermediates[2] -= 8388736.f; - fp32_intermediates[3] -= 8388736.f; - - uint32_t *bf16_result_ptr = reinterpret_cast(frag_b); - bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], - fp32_intermediates_casted[1], 0x7632); - bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], - fp32_intermediates_casted[3], 0x7632); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - float fp32_intermediates[4]; - uint32_t *fp32_intermediates_casted = reinterpret_cast(fp32_intermediates); - - static constexpr uint32_t fp32_base = 0x4B000000; - fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); - fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); - fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); - fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); - - fp32_intermediates[0] -= 8388608.f; - fp32_intermediates[1] -= 8388608.f; - fp32_intermediates[2] -= 8388608.f; - fp32_intermediates[3] -= 8388608.f; - - uint32_t *bf16_result_ptr = reinterpret_cast(frag_b); - bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], - fp32_intermediates_casted[1], 0x7632); - bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], - fp32_intermediates_casted[3], 0x7632); -} - -template <> -__device__ inline void dequant( - int q, half2 *frag_b) { - // Constants for FP8 (E4M3) and FP16 formats - constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5; - constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT; - constexpr int MASK = 0x7F007F00; - - // Extract and shift FP8 values to FP16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 8; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant( - int q, half2 *frag_b) { - dequant(q, frag_b); - - // Constants for FP8 (E4M3) and FP16 formats - constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5; - - // Construct and apply exponent bias - constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); - const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); - - // Convert to half2 and apply bias - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - // Constants for FP8 (E4M3) and BF16 formats - constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; - constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; - - constexpr int MASK = 0x7F007F00; - - // Extract and shift FP8 values to BF16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 8; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); - - // Constants for FP8 (E4M3) and BF16 formats - constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; - - // Construct and apply exponent bias - constexpr int BIAS_OFFSET = (1 << (BF16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); - // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent - // position - constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; - const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); - - // Convert to bfloat162 and apply bias - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); -} - -template <> -__device__ inline void dequant(int q, - half2 *frag_b) { - // Constants for FP4 (E2M1) and FP16 formats - constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5; - constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP4_EXPONENT; - constexpr int MASK = 0x70007000; - - // Extract and shift FP4 values to FP16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 4; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant( - int q, half2 *frag_b) { - dequant(q, frag_b); - - // Constants for FP4 (E2M1) and FP16 formats - constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5; - - // Construct and apply exponent bias - constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1)); - const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); - - // Convert to half2 and apply bias - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - // Constants for FP4 (E2M1) and FP16 formats - constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8; - constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP4_EXPONENT; - constexpr int MASK = 0x70007000; - - // Extract and shift FP4 values to FP16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 4; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant( - int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); - - // Constants for FP4 (E2M1) and BF16 formats - constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8; - - // Construct and apply exponent bias - constexpr int BIAS_OFFSET = (1 << (BF16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1)); - // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent - // position - constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; - const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); - - // Convert to half2 and apply bias - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); -} - -template <> -__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kFE2M1f.id(), true>( - int q, __nv_fp8x4_e4m3 *frag_b) { - // Constants for FP4 (E2M1) and FP16 formats - constexpr int FP4_EXPONENT = 2, FP8_EXPONENT = 4; - constexpr int RIGHT_SHIFT = FP8_EXPONENT - FP4_EXPONENT; - constexpr int MASK = 0x70707070; - - // Extract and shift FP4 values to FP16 format - int Out1 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 4; - int Out2 = (q & 0x80808080) | ((q & MASK) >> RIGHT_SHIFT); - - // Note1: reverse indexing is intentional because weights are permuted - // Note2: when dequant to 8bit type, we write to `frag_b[2]` instead of - // `frag_b[1]` to fit the layout of tensorcore - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant( - int q, int32_t *frag_b) { - constexpr int repeated_zp = 0x08080808; - constexpr int MASK = 0x80808080; - - frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK; - q >>= 4; - frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK; -} - -template <> -__device__ inline void dequant<__nv_fp8x4_e4m3, vllm::kU4B8.id(), true>( - int q, __nv_fp8x4_e4m3 *frag_b) { - int s = q & 0x08080808; - int Out1 = ((q & 0x07070707) | (s << 4)) + (s >> 3); - q >>= 4; - s = q & 0x08080808; - int Out2 = ((q & 0x07070707) | (s << 4)) + (s >> 3); - - frag_b[0] = *reinterpret_cast(&Out1); - frag_b[1] = *reinterpret_cast(&Out2); -} - -template -__device__ inline void dequant_fp8_scales(int q, scalar_t2 *frag_b); - -template <> -__device__ inline void dequant_fp8_scales( - int q, half2 *frag_b) { - int Out1 = (q & 0xFF00FF00) >> 1; - ; - q <<= 8; - int Out2 = (q & 0xFF00FF00) >> 1; - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -}; - -template <> -__device__ inline void dequant_fp8_scales( - int q, nv_bfloat162 *frag_b) { - constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; - constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; - constexpr int MASK = 0x7F007F00; - - // Extract and shift FP8 values to BF16 format - int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 8; - int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant_fp8_scales( - int q, nv_bfloat162 *frag_b) { - // In this conversion, 2 ** -127 in FP8E8M0 would become 0 in BF16, - // but we assume that such a extreme value would not occur in real models. - int Out1 = (q & 0xFF00FF00) >> 1; - q <<= 7; - int Out2 = q & 0x7F807F80; - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -}; - -// subtract zero point in quanted format and then dequant -template -__device__ inline void sub_zp_and_dequant(int q, scalar_t2 *frag_b, int zp); - -template <> -__device__ inline void sub_zp_and_dequant( - int q, int32_t *frag_b, int zp) { - // INT4 with zp -> INT8 - // see https://github.com/vllm-project/vllm/pull/24722 - int repeated_zp = 0x01010101 * zp; - int MASK = 0x80808080; - - frag_b[0] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK; - q >>= 4; - frag_b[1] = ((q & 0x0F0F0F0F | MASK) - repeated_zp) ^ MASK; -} - -template <> -__device__ inline void sub_zp_and_dequant<__nv_fp8x4_e4m3, vllm::kU4.id(), - true>(int q, __nv_fp8x4_e4m3 *frag_b, - int zp) { - // INT4 with zp -> FP8 - // see https://github.com/vllm-project/vllm/pull/24722 - uint32_t u_q = *reinterpret_cast(&q); - uint32_t u_zp = *reinterpret_cast(&zp); - uint32_t u_zp1 = u_zp + 1; - uint32_t repeated_zp = 0x01010101 * u_zp; - - uint32_t q0, s; - q0 = (u_q & 0x0F0F0F0F) | 0x70707070; - s = (q0 + repeated_zp) & 0x80808080; - uint32_t Out1 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s; - - u_q >>= 4; - q0 = (u_q & 0x0F0F0F0F) | 0x70707070; - s = (q0 + repeated_zp) & 0x80808080; - uint32_t Out2 = (q0 + (s >> 7) * u_zp1) & 0x0F0F0F0F | s; - - frag_b[0] = *reinterpret_cast(&Out1); - frag_b[1] = *reinterpret_cast(&Out2); -} - -#endif - -} // namespace MARLIN_NAMESPACE_NAME diff --git a/src/infiniop/ops/awq_marlin_gemm/marlin/kernel.h b/src/infiniop/ops/awq_marlin_gemm/marlin/kernel.h deleted file mode 100644 index b1fabe323..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/marlin/kernel.h +++ /dev/null @@ -1,43 +0,0 @@ - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -#include "../core/scalar_type.hpp" -#include "marlin.cuh" -#include "marlin_dtypes.cuh" - -#define MARLIN_KERNEL_PARAMS \ - const int4 *__restrict__ A, const int4 *__restrict__ B, \ - int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ - const int4 *__restrict__ b_bias_ptr, \ - const float *__restrict__ a_scales_ptr, \ - const int4 *__restrict__ scales_ptr, \ - const float *__restrict__ global_scale_ptr, \ - const int4 *__restrict__ zp_ptr, const int *__restrict__ g_idx, \ - int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \ - bool has_bias, bool use_atomic_add, bool use_fp32_reduce, \ - int max_shared_mem - -namespace MARLIN_NAMESPACE_NAME { -template shared - // fetch pipeline - const int group_blocks, // number of consecutive 16x16 blocks - // with a separate quantization scale - const bool is_zp_float // is zero point of float16 type? - > -__global__ void Marlin(MARLIN_KERNEL_PARAMS); - -} diff --git a/src/infiniop/ops/awq_marlin_gemm/marlin/marlin.cuh b/src/infiniop/ops/awq_marlin_gemm/marlin/marlin.cuh deleted file mode 100644 index f3d897d27..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/marlin/marlin.cuh +++ /dev/null @@ -1,178 +0,0 @@ -#pragma once - -#ifndef _marlin_cuh -#define _marlin_cuh - -#include -#include -#include -#include - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -template -__device__ __forceinline__ uint32_t __cvta_generic_to_shared(T *ptr) { - size_t smem_addr; - asm volatile( - "cvta.to.shared.u64 %0, %1;" - : "=l"(smem_addr) - : "l"(ptr)); - return static_cast(smem_addr); -} - -namespace MARLIN_NAMESPACE_NAME { - -// Marlin params - -// 8 warps are a good choice since every SM has 4 schedulers and having more -// than 1 warp per schedule allows some more latency hiding. At the same time, -// we want relatively few warps to have many registers per warp and small tiles. -static constexpr int default_threads = 256; - -static constexpr int pipe_stages = 4; // 4 pipeline stages fit into shared memory - -static constexpr int min_thread_n = 64; -static constexpr int min_thread_k = 64; -static constexpr int max_thread_n = 256; - -static constexpr int tile_size = 16; -static constexpr int max_par = 16; - -// Repack params -static constexpr int repack_stages = 8; - -static constexpr int repack_threads = 256; - -static constexpr int tile_k_size = tile_size; -static constexpr int tile_n_size = tile_k_size * 4; - -// Helpers -template -struct Vec { - T elems[n]; - __device__ T &operator[](int i) { return elems[i]; } -}; - -using I4 = Vec; - -constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; } - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - -__device__ inline void cp_async1_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async2_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4(void *smem_ptr, const void *glob_ptr) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; -} - -__device__ inline void cp_async_fence() {} - -template -__device__ inline void cp_async_wait() {} - -#else - -__device__ inline void cp_async1_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 4; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async2_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 8; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.cg.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4(void *smem_ptr, const void *glob_ptr) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " cp.async.cg.shared.global [%0], [%1], %2;\n" - "}\n" ::"r"(smem), - "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async_fence() { - asm volatile("cp.async.commit_group;\n" ::); -} - -template -__device__ inline void cp_async_wait() { - asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); -} - -#endif - -} // namespace MARLIN_NAMESPACE_NAME - -#endif diff --git a/src/infiniop/ops/awq_marlin_gemm/marlin/marlin_dtypes.cuh b/src/infiniop/ops/awq_marlin_gemm/marlin/marlin_dtypes.cuh deleted file mode 100644 index 3ee06faef..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/marlin/marlin_dtypes.cuh +++ /dev/null @@ -1,155 +0,0 @@ - -#ifndef _data_types_cuh -#define _data_types_cuh -#include "../core/scalar_type.hpp" -#include "marlin.cuh" -#include -#include -#include - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -namespace MARLIN_NAMESPACE_NAME { - -template -class MarlinScalarType { -}; - -template <> -class MarlinScalarType { -public: - using scalar_t = half; - using scalar_t2 = half2; - using scalar_t4 = half2; - using scalar_32bit_t = half2; - - // Matrix fragments for tensor core instructions; their precise layout is - // documented here: - // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type - using FragA = Vec; - using FragB = Vec; - using FragC = Vec; - using FragS = Vec; - using FragS0 = Vec<__nv_fp8x2_e4m3, 1>; - using FragZP = Vec; - - static __device__ float inline num2float(const half x) { - return __half2float(x); - } - - static __device__ half2 inline num2num2(const half x) { - return __half2half2(x); - } - - static __device__ half2 inline nums2num2(const half x1, const half x2) { - return __halves2half2(x1, x2); - } - - static __host__ __device__ half inline float2num(const float x) { - return __float2half(x); - } - - static __host__ __device__ float2 inline num22float2(const half2 x) { - return __half22float2(x); - } -}; - -template <> -class MarlinScalarType { -public: - using scalar_t = nv_bfloat16; - using scalar_t2 = nv_bfloat162; - using scalar_t4 = nv_bfloat162; - using scalar_32bit_t = nv_bfloat162; - - using FragA = Vec; - using FragB = Vec; - using FragC = Vec; - using FragS = Vec; - using FragS0 = Vec<__nv_fp8x2_e4m3, 1>; - using FragZP = Vec; - -#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 - static __device__ float inline num2float(const nv_bfloat16 x) { - return __bfloat162float(x); - } - - static __device__ nv_bfloat162 inline num2num2(const nv_bfloat16 x) { - return __bfloat162bfloat162(x); - } - - static __device__ nv_bfloat162 inline nums2num2(const nv_bfloat16 x1, - const nv_bfloat16 x2) { - return __halves2bfloat162(x1, x2); - } - - static __host__ __device__ nv_bfloat16 inline float2num(const float x) { - return __float2bfloat16(x); - } - - static __host__ __device__ float2 inline num22float2(const nv_bfloat162 x) { - return __bfloat1622float2(x); - } -#endif -}; - -template <> -class MarlinScalarType { -public: - using scalar_t = __nv_fp8_e4m3; - using scalar_t2 = __nv_fp8x2_e4m3; - using scalar_t4 = __nv_fp8x4_e4m3; - using scalar_32bit_t = __nv_fp8x4_e4m3; - - using FragA = Vec<__nv_fp8x4_e4m3, 4>; - using FragB = Vec<__nv_fp8x4_e4m3, 2>; - using FragC = Vec; - using FragZP = Vec<__nv_fp8x2_e4m3, 4>; - - static __host__ __device__ - float2 inline num22float2(const __nv_fp8x2_e4m3 x) { - return (float2)x; - } -}; - -template <> -class MarlinScalarType { -public: - using scalar_t = int8_t; - using scalar_t2 = int16_t; - using scalar_t4 = int32_t; - using scalar_32bit_t = int32_t; - - using FragA = Vec; - using FragB = Vec; - using FragC = Vec; - using FragZP = Vec; -}; - -template -class MarlinScalarType2 { -}; - -template <> -class MarlinScalarType2 : public MarlinScalarType { -}; - -template <> -class MarlinScalarType2 - : public MarlinScalarType { -}; - -template <> -class MarlinScalarType2<__nv_fp8_e4m3> - : public MarlinScalarType { -}; - -template <> -class MarlinScalarType2 : public MarlinScalarType { -}; - -} // namespace MARLIN_NAMESPACE_NAME - -#endif diff --git a/src/infiniop/ops/awq_marlin_gemm/marlin/marlin_mma.h b/src/infiniop/ops/awq_marlin_gemm/marlin/marlin_mma.h deleted file mode 100644 index d8fb2946e..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/marlin/marlin_mma.h +++ /dev/null @@ -1,267 +0,0 @@ - -#include "marlin_dtypes.cuh" - -namespace MARLIN_NAMESPACE_NAME { - -// m16n8k16 tensor core mma instruction with fp16 inputs and fp32 -// output/accumulation. -template -__device__ inline void mma( - const typename MarlinScalarType::FragA &a_frag, - const typename MarlinScalarType::FragB &frag_b, - typename MarlinScalarType::FragC &frag_c, int idx = 0) { - const uint32_t *a = reinterpret_cast(&a_frag); - const uint32_t *b = reinterpret_cast(&frag_b); - using scalar_t = typename MarlinScalarType::scalar_t; - if constexpr (!std::is_same::value || k_size != 16) { - static_assert(!use_fp16_accum); - } - - if constexpr (k_size == 16) { - if constexpr (std::is_same::value && !use_fp16_accum) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(b[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]), - "f"(c[3])); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[2]), "r"(a[3]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), - "f"(c[3])); -#else - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); -#endif - } else if constexpr (std::is_same::value && use_fp16_accum) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - uint32_t *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " - "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(a[0]), "r"(a[1]), "r"(b[0]), "r"(c[0]), "r"(c[1])); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " - "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(a[2]), "r"(a[3]), "r"(b[1]), "r"(c[0]), "r"(c[1])); -#else - uint32_t *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " - "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), - "r"(c[0]), "r"(c[1])); -#endif - } else if constexpr (std::is_same::value) { - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); - } else if constexpr (std::is_same::value) { - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "f"(c[0]), - "f"(c[1]), "f"(c[2]), "f"(c[3])); - } else if constexpr (std::is_same::value) { - int32_t *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) - : "r"(a[idx * 2]), "r"(a[idx * 2 + 1]), "r"(b[idx]), "r"(c[0]), - "r"(c[1]), "r"(c[2]), "r"(c[3])); - } - } else if (k_size == 32) { - if constexpr (std::is_same::value) { - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); - } else if constexpr (std::is_same::value) { - int32_t *c = reinterpret_cast(&frag_c); -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(a[0]), "r"(b[0]), "r"(c[0]), "r"(c[1])); - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[2]), "=r"(c[3]) - : "r"(a[1]), "r"(b[0]), "r"(c[2]), "r"(c[3])); - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(a[2]), "r"(b[1]), "r"(c[0]), "r"(c[1])); - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[2]), "=r"(c[3]) - : "r"(a[3]), "r"(b[1]), "r"(c[2]), "r"(c[3])); -#else - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), - "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3])); -#endif - } - } -} - -template -__device__ inline void mma_trans( - const typename MarlinScalarType::FragA &a_frag, - const typename MarlinScalarType::FragB &frag_b, - const typename MarlinScalarType::FragB &frag_b2, - typename MarlinScalarType::FragC &frag_c) { - const uint32_t *a = reinterpret_cast(&a_frag); - const uint32_t *b = reinterpret_cast(&frag_b); - const uint32_t *b2 = reinterpret_cast(&frag_b2); - float *c = reinterpret_cast(&frag_c); - using scalar_t = typename MarlinScalarType::scalar_t; - if constexpr (!std::is_same::value || k_size != 16) { - static_assert(!use_fp16_accum); - } - - if constexpr (k_size == 16) { - if constexpr (std::is_same::value && !use_fp16_accum) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]), - "f"(c[3])); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[1]), "r"(b2[1]), "r"(a[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), - "f"(c[3])); -#else - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); -#endif - } else if constexpr (std::is_same::value && use_fp16_accum) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - uint32_t *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " - "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1])); - asm volatile( - "mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 " - "{%0,%1}, {%2,%3}, {%4}, {%5,%6};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(b[1]), "r"(b2[1]), "r"(a[1]), "r"(c[0]), "r"(c[1])); -#else - uint32_t *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " - "{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), - "r"(c[0]), "r"(c[1])); -#endif - } else if constexpr (std::is_same::value) { - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); - } else if constexpr (std::is_same::value) { - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.e4m3.e4m3.f32 " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "f"(c[0]), "f"(c[1]), "f"(c[2]), - "f"(c[3])); - } else if constexpr (std::is_same::value) { - int32_t *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1,%2,%3}, {%4,%5}, {%6}, {%7,%8,%9,%10};\n" - : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(a[0]), "r"(c[0]), "r"(c[1]), "r"(c[2]), - "r"(c[3])); - } - } else { - if constexpr (std::is_same::value) { - float *c = reinterpret_cast(&frag_c); - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); - } else if constexpr (std::is_same::value) { - int32_t *c = reinterpret_cast(&frag_c); -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(b[0]), "r"(a[0]), "r"(c[0]), "r"(c[1])); - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[2]), "=r"(c[3]) - : "r"(b2[1]), "r"(a[0]), "r"(c[2]), "r"(c[3])); - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[0]), "=r"(c[1]) - : "r"(b[0]), "r"(a[1]), "r"(c[0]), "r"(c[1])); - asm volatile( - "mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1}, {%2}, {%3}, {%4,%5};\n" - : "=r"(c[2]), "=r"(c[3]) - : "r"(b2[1]), "r"(a[1]), "r"(c[2]), "r"(c[3])); -#else - asm volatile( - "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32.satfinite " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) - : "r"(b[0]), "r"(b2[0]), "r"(b[1]), "r"(b2[1]), "r"(a[0]), "r"(a[1]), - "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3])); -#endif - } - } -} - -} // namespace MARLIN_NAMESPACE_NAME diff --git a/src/infiniop/ops/awq_marlin_gemm/marlin/marlin_template.h b/src/infiniop/ops/awq_marlin_gemm/marlin/marlin_template.h deleted file mode 100644 index d705af8fb..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/marlin/marlin_template.h +++ /dev/null @@ -1,1988 +0,0 @@ -/* - * Modified by Neural Magic - * Copyright (C) Marlin.2024 Elias Frantar - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * Adapted from https://github.com/IST-DASLab/marlin - */ - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -#include "../core/scalar_type.hpp" -#include "dequant.h" -#include "marlin.cuh" -#include "marlin_dtypes.cuh" -#include "marlin_mma.h" - -#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ - static_assert(std::is_same::value || std::is_same::value, \ - "only float16 and bfloat16 is supported"); - -namespace MARLIN_NAMESPACE_NAME { - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750 - -template shared - // fetch pipeline - const bool has_act_order, // whether act_order is enabled - const int group_blocks, // number of consecutive 16x16 blocks - // with a separate quantization scale - const bool is_zp_float // is zero point of float16 type? - > -__global__ void Marlin( - const int4 *__restrict__ A, // fp16 input matrix of shape mxk - const int4 *__restrict__ B, // 4bit quantized weight matrix of shape kxn - int4 *__restrict__ C, // fp16 output buffer of shape mxn - int4 *__restrict__ C_tmp, // fp32 tmp output buffer (for reduce) - const int4 *__restrict__ scales_ptr, // fp16 quantization scales of shape - // (k/groupsize)xn - const int *__restrict__ g_idx, // int32 group indices of shape k - int num_groups, // number of scale groups per output channel - int prob_m, // batch dimension m - int prob_n, // output dimension n - int prob_k, // reduction dimension k - int *locks, // extra global storage for barrier synchronization - bool use_fp32_reduce // whether to use fp32 global reduce -) { -} - -} // namespace marlin - -#else - -// Instruction for loading a full 16x16 matrix fragment of operand A from shared -// memory, directly in tensor core layout. -template -__device__ inline void ldsm(typename MarlinScalarType::FragA &frag_a, - const void *smem_ptr) { - uint32_t *a = reinterpret_cast(&frag_a); - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - if constexpr (count == 4) { - asm volatile( - "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" - : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) - : "r"(smem)); - } else if constexpr (count == 2) { - asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" - : "=r"(a[0]), "=r"(a[1]) - : "r"(smem)); - } else if constexpr (count == 1) { - asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" - : "=r"(a[0]) - : "r"(smem)); - } else { - static_assert(count == 1 || count == 2 || count == 4, "invalid count"); - } -} - -// Multiply dequantized values by the corresponding quantization scale; used -// only for grouped quantization. -template -__device__ inline void scale(typename MarlinScalarType::FragB &frag_b, - typename MarlinScalarType::FragS &frag_s, - int i) { - using scalar_t = typename MarlinScalarType::scalar_t; - using scalar_t2 = typename MarlinScalarType::scalar_t2; - scalar_t2 s = MarlinScalarType::num2num2( - reinterpret_cast(&frag_s)[i]); - frag_b[0] = __hmul2(frag_b[0], s); - frag_b[1] = __hmul2(frag_b[1], s); -} - -template -__device__ inline void scale_and_sub( - typename MarlinScalarType::FragB &frag_b, - typename MarlinScalarType::scalar_t s, - typename MarlinScalarType::scalar_t zp) { - using scalar_t = typename MarlinScalarType::scalar_t; - using scalar_t2 = typename MarlinScalarType::scalar_t2; - scalar_t2 s2 = MarlinScalarType::num2num2(s); - scalar_t2 zp2 = MarlinScalarType::num2num2(zp); - frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2)); - frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2)); -} - -template -__device__ inline void sub_zp( - typename MarlinScalarType::FragB &frag_b, - typename MarlinScalarType::scalar_t2 &frag_zp, int i) { - using scalar_t = typename MarlinScalarType::scalar_t; - using scalar_t2 = typename MarlinScalarType::scalar_t2; - scalar_t2 zp = MarlinScalarType::num2num2( - reinterpret_cast(&frag_zp)[i]); - frag_b[0] = __hsub2(frag_b[0], zp); - frag_b[1] = __hsub2(frag_b[1], zp); -} - -// Same as above, but for act_order (each K is multiplied individually) -template -__device__ inline void scale4( - typename MarlinScalarType::FragB &frag_b, - typename MarlinScalarType::FragS &frag_s_1, - typename MarlinScalarType::FragS &frag_s_2, - typename MarlinScalarType::FragS &frag_s_3, - typename MarlinScalarType::FragS &frag_s_4, int i) { - using scalar_t = typename MarlinScalarType::scalar_t; - using scalar_t2 = typename MarlinScalarType::scalar_t2; - - scalar_t2 s_val_1_2; - s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; - s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; - - scalar_t2 s_val_3_4; - s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; - s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; - - frag_b[0] = __hmul2(frag_b[0], s_val_1_2); - frag_b[1] = __hmul2(frag_b[1], s_val_3_4); -} - -// Given 2 floats multiply by 2 scales (halves) -template -__device__ inline void scale_float( - float *c, typename MarlinScalarType::FragS &s) { - using scalar_t = typename MarlinScalarType::scalar_t; - scalar_t *s_ptr = reinterpret_cast(&s); - c[0] = __fmul_rn(c[0], MarlinScalarType::num2float(s_ptr[0])); - c[1] = __fmul_rn(c[1], MarlinScalarType::num2float(s_ptr[1])); -} - -// Wait until barrier reaches `count`, then lock for current threadblock. -__device__ inline void barrier_acquire(int *lock, int count) { - if (threadIdx.x == 0) { - int state = -1; - do { - // Guarantee that subsequent writes by this threadblock will be visible - // globally. - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" - : "=r"(state) - : "l"(lock)); - } while (state != count); - } - __syncthreads(); -} - -// Release barrier and increment visitation count. -__device__ inline void barrier_release(int *lock, bool reset = false) { - __syncthreads(); - if (threadIdx.x == 0) { - if (reset) { - lock[0] = 0; - return; - } - int val = 1; - // Make sure that all writes since acquiring this barrier are visible - // globally, while releasing the barrier. - asm volatile("fence.acq_rel.gpu;\n"); - asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" - : - : "l"(lock), "r"(val)); - } -} - -// Wait until value of lock to be negative, and then add 1 -__device__ inline void wait_negative_and_add(int *lock) { - if (threadIdx.x == 0) { - int state = 0; - do { - // Guarantee that subsequent writes by this threadblock will be visible - // globally. - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" - : "=r"(state) - : "l"(lock)); - } while (state >= 0); - atomicAdd(lock, 1); - } - __syncthreads(); -} - -template shared - // fetch pipeline - const int group_blocks, // number of consecutive 16x16 blocks - // with a separate quantization scale - const bool is_zp_float // is zero point of float16 type? - > -__global__ void Marlin( - const int4 *__restrict__ A0, // fp16 input matrix of shape mxk - const int4 *__restrict__ B, // 4bit quantized weight matrix of shape kxn - int4 *__restrict__ C0, // fp16 output buffer of shape mxn - int4 *__restrict__ C_tmp, // fp32 tmp output buffer (for reduce) - const int4 *__restrict__ b_bias_ptr, - // float scales of input matrix, only used when is_a_8bit == true. - // shape (m,) - const float *__restrict__ a_scales_ptr, - // fp16 quantization scales. shape (k/groupsize, n) - const int4 *__restrict__ scales_ptr, - // float global scale (for nvfp4// only) - const float *__restrict__ global_scale_ptr, - // 4bit packed zero-points of shape - // (k/groupsize, n/pack_factor) - const int4 *__restrict__ zp_ptr, - // int32 group indices of shape k - const int *__restrict__ g_idx, - int num_groups, // number of scale groups per output channel - int prob_m, // batch dimension m - int prob_n, // output dimension n - int prob_k, // reduction dimension k - int lda, // A.stride(0), equal to prob_k is A is contiguous - int *locks, // extra global storage for barrier synchronization - bool has_bias, - bool use_atomic_add, // whether to use atomic add to reduce - bool use_fp32_reduce, // whether to use fp32 global reduce - int max_shared_mem) { - // Each threadblock processes one "stripe" of the B matrix with (roughly) the - // same size, which might involve multiple column "slices" (of width 16 * - // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM - // example: - // 0 1 3 - // 0 2 3 - // 1 2 4 - // While this kind of partitioning makes things somewhat more complicated, it - // ensures good utilization of all SMs for many kinds of shape and GPU - // configurations, while requiring as few slow global cross-threadblock - // reductions as possible. - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 890 - // FP8 computation is only supported for Ada Lovelace or newer architectures. - if constexpr (a_type_id == vllm::kFE4M3fn.id()) { - return; - } -#endif - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - // Turing TensorCore only supports fp16 and int8 - if constexpr (a_type_id != vllm::kFloat16.id() && a_type_id != vllm::kS8.id()) { - return; - } -#endif - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 - constexpr auto num_bits = vllm::ScalarType::from_id(b_type_id).size_bits(); - // Disable use_fp16_accum for NVFP4 and cases when group_size == -1 && - // num_bits == 4 - constexpr bool use_fp16_accum = a_type_id == vllm::kFloat16.id() && (!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) && !(group_blocks == -1 && num_bits == 4)); -#else - constexpr bool use_fp16_accum = false; -#endif - using Adtype = MarlinScalarType; - using Cdtype = MarlinScalarType; - const int4 *A = A0; - int4 *C = C0; - - using scalar_t = typename MarlinScalarType::scalar_t; - using scalar_t2 = typename MarlinScalarType::scalar_t2; - using scalar_32bit_t = typename MarlinScalarType::scalar_32bit_t; - - using c_scalar_t = typename MarlinScalarType::scalar_t; - using c_scalar_t2 = typename MarlinScalarType::scalar_t2; - - using FragA = typename MarlinScalarType::FragA; - using FragB = typename MarlinScalarType::FragB; - using FragC = typename MarlinScalarType::FragC; - using FragS = typename MarlinScalarType::FragS; - using FragZP = typename MarlinScalarType::FragZP; - - static constexpr auto a_type = vllm::ScalarType::from_id(a_type_id); - static constexpr auto b_type = vllm::ScalarType::from_id(b_type_id); - static constexpr auto c_type = vllm::ScalarType::from_id(c_type_id); - static constexpr auto s_type = vllm::ScalarType::from_id(s_type_id); - if constexpr (b_type == vllm::kFE2M1f) { - static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 || s_type == vllm::kFE8M0fnu && group_blocks == 2); - } else if constexpr (std::is_same::value) { - static_assert(s_type == vllm::kBFloat16); - } else if constexpr (std::is_same::value) { - static_assert(s_type == vllm::kFloat16); - } - - constexpr bool is_a_8bit = a_type.size_bits() == 8; - if constexpr (!is_a_8bit) { - static_assert(std::is_same::value); - } - constexpr bool has_zp = b_type == vllm::kU4 || b_type == vllm::kU8; - constexpr bool is_int_type = b_type == vllm::kU4 || b_type == vllm::kU8 || b_type == vllm::kS4 || b_type == vllm::kS8 || b_type == vllm::kU4B8 || b_type == vllm::kU8B128; - // see comments of dequant.h for more details - constexpr bool dequant_skip_flop = is_a_8bit || b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn || has_zp && !is_zp_float && !std::is_same::value || has_zp && !is_zp_float && !(b_type == vllm::kU8); - - float global_scale_f32 = 1.0f; - - if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) { - global_scale_f32 = global_scale_ptr[0]; - } - - constexpr bool has_act_order = group_blocks == 0; - constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); - - extern __shared__ int4 sh[]; - float *sh_a_s = reinterpret_cast(sh); - int4 *sh_new = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0); - constexpr int pack_factor = 32 / b_type.size_bits(); - static_assert(thread_m_blocks == 1 || !m_block_size_8); - - // For larger GEMMs we run multiple batchsize 64 versions in parallel for a - // better partitioning with less reductions - int parallel = 1; - if (prob_m > m_block_size) { - parallel = prob_m / m_block_size; - prob_m = m_block_size; - } - - int k_tiles = prob_k / 16 / thread_k_blocks; - int n_tiles = prob_n / 16 / thread_n_blocks; - - int global_mn_tiles = parallel * n_tiles; - int part2_mn_tiles = global_mn_tiles; - int part1_mn_iters = 0; - bool in_part2 = false; - - if (global_mn_tiles > gridDim.x) { - part2_mn_tiles = global_mn_tiles % gridDim.x; - if (part2_mn_tiles * 3 <= gridDim.x) { - part2_mn_tiles += gridDim.x; - } - part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x; - } - - int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x); - - if constexpr (!has_act_order && group_blocks != -1) { - if (group_blocks >= thread_k_blocks) { - // Ensure that the number of tiles in each stripe is a multiple of the - // groupsize; this avoids an annoying special case where a stripe starts - // in the middle of group. - iters = (group_blocks / thread_k_blocks) * div_ceil(iters, (group_blocks / thread_k_blocks)); - } - } - - int slice_row = 0; - int slice_col_par = blockIdx.x; - int slice_col; - int slice_iters = k_tiles; // number of threadblock tiles in the current slice - // total number of active threadblocks in the current slice - int slice_count = 1; - // index of threadblock in current slice; numbered bottom to top - int slice_idx = 0; - - int par_id = 0; - int locks_off = 0; - - if (part2_mn_tiles >= gridDim.x) { - // when part2_mn_tiles >= sms - // then there are at most $sms$ conflict tile blocks - locks_off = blockIdx.x; - } else { - locks_off = (iters * blockIdx.x) / k_tiles - 1; - } - - // Compute all information about the current slice which is required for - // synchronization. - bool first_init = true; - auto init_part2_slice = [&]() { - slice_iters = iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); - if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) { - slice_iters = 0; - } - if (slice_iters == 0) { - return; - } - if (slice_row + slice_iters > k_tiles) { - slice_iters = k_tiles - slice_row; - } - slice_count = 1; - slice_idx = 0; - int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); - if (col_first <= k_tiles * (slice_col_par + 1)) { - int col_off = col_first - k_tiles * slice_col_par; - slice_count = div_ceil(k_tiles - col_off, iters); - if (col_off > 0) { - slice_count++; - } - int delta_first = iters * blockIdx.x - col_first; - if (delta_first < 0 || (col_off == 0 && delta_first == 0)) { - slice_idx = slice_count - 1; - } else { - slice_idx = slice_count - 1 - delta_first / iters; - if (col_off > 0) { - slice_idx--; - } - } - } - if (part2_mn_tiles >= gridDim.x) { - if (slice_count > 1 && slice_idx == slice_count - 1) { - locks_off++; - } - } else { - locks_off++; - } - - if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) { - constexpr int threads_per_m = 16 * thread_n_blocks / 8; - int m_per_thread = div_ceil(thread_m_blocks * 16, threads / threads_per_m); - if (m_block_size_8) { - m_per_thread = div_ceil(8, threads / threads_per_m); - } - for (int i = 0; i < m_per_thread; i++) { - int row = threads / threads_per_m * i + threadIdx.x / threads_per_m; - if (row < prob_m) { - int col = slice_col * 16 * thread_n_blocks / 8 + threadIdx.x % threads_per_m; - C[row * prob_n / 8 + col] = {0, 0, 0, 0}; - } - } - // After write zero to output, write a negative value to lock. - // Every SM that processes the same slice would wait for - // the negative value, and then atomicAdd 1 to it. - // After all SMs are processed, the lock value would back to 0 again. - __syncthreads(); - if (threadIdx.x == 0) { - locks[locks_off] = 1 - slice_count; - } - } - - if (slice_col == n_tiles) { - A += 16 * thread_m_blocks * lda / (is_a_8bit ? 16 : 8); - C += 16 * thread_m_blocks * prob_n / 8; - slice_col = 0; - par_id++; - } - if (is_a_8bit && (first_init || slice_col == 0)) { - __syncthreads(); - int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x; - cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd], - threadIdx.x < prob_m); - } - }; - - auto init_part1_slice = [&]() { - if (part1_mn_iters) { - part1_mn_iters--; - par_id = slice_col_par / n_tiles; - slice_col = slice_col_par % n_tiles; - slice_iters = k_tiles; - A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda; - C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n; - if (is_a_8bit) { - __syncthreads(); - int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x; - cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd], - threadIdx.x < prob_m); - } - } - }; - - auto init_slice = [&]() { - if (!in_part2 && !part1_mn_iters) { - in_part2 = true; - slice_col_par = (iters * blockIdx.x) / k_tiles; - slice_row = (iters * blockIdx.x) % k_tiles; - slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles; - par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles; - A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda; - C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n; - } - if (!in_part2) { - init_part1_slice(); - } else { - init_part2_slice(); - first_init = false; - } - }; - - init_slice(); - - // A sizes/strides - - // stride of the A matrix in global memory - int a_gl_stride = lda / (is_a_8bit ? 16 : 8); - // stride of an A matrix tile in shared memory - constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); - // delta between subsequent A tiles in global memory - constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); - // between subsequent accesses within a tile - int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); - // between shared memory writes - constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); - // within a shared memory tile - constexpr int a_sh_rd_delta_i = a_sh_stride * 16; - // overall size of a tile - constexpr int a_sh_stage = a_sh_stride * m_block_size; - // number of shared write iterations for a tile - constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); - - // B sizes/strides - int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4)); - constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4); - constexpr int b_thread_vecs = b_type.size_bits() == 4 ? 1 : 2; - constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; - - int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); - constexpr int b_sh_wr_delta = threads * b_thread_vecs; - constexpr int b_sh_stage = b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); - constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; - - // Scale sizes/strides without act_order - int s_gl_stride = prob_n / (b_type == vllm::kFE2M1f ? 16 : 8); - constexpr int s_sh_stride = 16 * thread_n_blocks / (b_type == vllm::kFE2M1f ? 16 : 8); - constexpr int s_tb_groups = !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks - ? thread_k_blocks / group_blocks - : 1; - constexpr int s_sh_stage = s_tb_groups * s_sh_stride; - int s_gl_rd_delta = s_gl_stride; - - // Scale size/strides with act_order - constexpr int tb_k = 16 * thread_k_blocks; - constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; - // constexpr int act_s_row_stride = 1; - // int act_s_col_stride = act_s_row_stride * num_groups; - constexpr int act_s_max_num_groups = 32; - int act_s_col_stride = 1; - int act_s_col_warp_stride = act_s_col_stride * 8; - - constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4); - int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; - - // Zero-points sizes/strides - int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4; - constexpr int zp_sh_stride = is_zp_float - ? 16 * thread_n_blocks / 8 - : ((16 * thread_n_blocks) / pack_factor) / 4; - constexpr int zp_tb_groups = s_tb_groups; - constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0; - int zp_gl_rd_delta = zp_gl_stride; - - // Global A read index of current thread. - int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - a_gl_rd += a_gl_rd_delta_o * slice_row; - // Shared write index of current thread. - int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - // Shared read index. - int a_sh_rd = a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) + (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1)); - a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters; - - int b_gl_rd; - if (threads <= b_sh_stride) { - b_gl_rd = threadIdx.x; - } else { - b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); - } - - b_gl_rd += b_sh_stride * slice_col; - b_gl_rd += b_gl_rd_delta_o * slice_row; - auto b_sh_rd = threadIdx.x * b_thread_vecs; - b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1)); - - // For act_order - int slice_k_start = tb_k * slice_row; - int slice_k_finish = slice_k_start + tb_k * slice_iters; - int slice_k_start_shared_fetch = slice_k_start; - int slice_n_offset = act_s_col_tb_stride * slice_col; - - // No act_order - int s_gl_rd; - if constexpr (!has_act_order) { - if constexpr (group_blocks == -1) { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - } else if constexpr (group_blocks >= thread_k_blocks) { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col + threadIdx.x; - } else { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / s_sh_stride) + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; - } - } - auto s_sh_wr = threadIdx.x; - bool s_sh_wr_pred = threadIdx.x < s_sh_stage; - - // Zero-points - int zp_gl_rd; - if constexpr (has_zp) { - if constexpr (group_blocks == -1) { - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; - } else if constexpr (group_blocks >= thread_k_blocks) { - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + zp_sh_stride * slice_col + threadIdx.x; - } else { - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / zp_sh_stride) + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; - } - } - auto zp_sh_wr = threadIdx.x; - bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage; - - // We use a different scale layout for grouped and column-wise quantization as - // we scale a `half2` tile in column-major layout in the former and in - // row-major in the latter case. - int s_sh_rd; - if constexpr (is_a_8bit) { - s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4); - } else if constexpr (group_blocks != -1) { - s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4; - } else if constexpr (group_blocks == -1 && (m_block_size_8 || (has_zp && !dequant_skip_flop))) { - s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; - } else { - s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; - } - - int bias_sh_rd; - if constexpr (m_block_size_8) { - bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; - } else { - bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; - } - - int bias_sh_wr = threadIdx.x; - int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; - - // Zero-points have the same read layout as the scales - // (without column-wise case) - constexpr int num_col_threads = 8; - constexpr int num_row_threads = 4; - constexpr int num_ints_per_thread = 8 / pack_factor; - int zp_sh_rd; - if constexpr (has_zp) { - if constexpr (is_zp_float) { - if constexpr (group_blocks != -1) { - zp_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4; - } - } else if (is_a_8bit) { - zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % tb_n_warps / 2) + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); - } else { - zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % tb_n_warps) + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); - } - } - - // Precompute which thread should not read memory in which iterations; this is - // needed if there are more threads than required for a certain tilesize or - // when the batchsize is not a multiple of 16. - bool a_sh_wr_pred[a_sh_wr_iters]; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) { - a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m; - } - - // To ensure that writing and reading A tiles to/from shared memory, the - // latter in fragment format, is fully bank conflict free, we need to use a - // rather fancy XOR-based layout. The key here is that neither reads nor - // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the - // same shared memory banks. Further, it seems (based on NSight-Compute) that - // each warp must also write a consecutive memory segment? - auto transform_a = [&](int i) { - int row = i / a_gl_rd_delta_o; - return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8); - }; - // Since the computation of this remapping is non-trivial and, due to our main - // loop unrolls, all shared memory accesses are static, we simply precompute - // both transformed reads and writes. - int a_sh_wr_trans[a_sh_wr_iters]; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) { - a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); - } - int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; -#pragma unroll - for (int i = 0; i < b_sh_wr_iters; i++) { -#pragma unroll - for (int j = 0; j < thread_m_blocks; j++) { - a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd); - } - } - - // Since B-accesses have non-constant stride they have to be computed at - // runtime; we break dependencies between subsequent accesses with a tile by - // maintining multiple pointers (we have enough registers), a tiny - // optimization. - - // Shared memory storage for global fetch pipelines. - constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks; - constexpr int sh_b_size = stages * b_sh_stage; - int4 *sh_b = sh_new; - int4 *sh_red = sh_new; - constexpr int sh_size_b_red_min = (sh_red_size < sh_b_size ? sh_red_size : sh_b_size); - constexpr int sh_size_b_red_max = (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); - constexpr int sh_bias_size = (thread_n_blocks * 16 / 8); - constexpr int sh_b_red_bias_size = sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size) - ? sh_size_b_red_max - : (sh_size_b_red_min + sh_bias_size); - - int4 *sh_bias = sh_new + sh_size_b_red_min; - int4 *sh_g_idx = sh_new + sh_b_red_bias_size; - int4 *sh_zp = sh_g_idx + (stages * g_idx_stage); - constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride) - : (stages * s_sh_stage); - int4 *sh_s = sh_zp + (stages * zp_sh_stage); - int4 *sh_a = sh_s + sh_s_size; - - // Register storage for double buffer of shared memory reads. - FragA frag_a[2][thread_m_blocks]; - I4 frag_b_quant[2][b_thread_vecs]; - FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2]; - FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2]; - FragS frag_s[2][4]; // No act-order - FragS frag_bias[2][4]; - FragS act_frag_s[2][4][4]; // For act-order - int frag_qzp[2][num_ints_per_thread]; // Zero-points - FragZP frag_zp; // Zero-points in fp16 - FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ - - if constexpr (is_a_8bit) { -#pragma unroll - for (int j = 0; j < 2; j++) { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int g = 0; g < 4; g++) { - frag_c_tmp[i][j][0][g] = 0.0f; - } - -#pragma unroll - for (int g = 0; g < 4; g++) { - frag_c_tmp[i][j][1][g] = 0.0f; - } - } - } - } - - // Zero accumulators. - auto zero_accums = [&]() { -#pragma unroll - for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) - reinterpret_cast(frag_c)[i] = 0; - }; - - int sh_first_group_id = -1; - int sh_num_groups = -1; - - auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id, - int last_group_id) { - sh_first_group_id = first_group_id; - sh_num_groups = last_group_id - first_group_id + 1; - - if (sh_num_groups > act_s_max_num_groups) { - sh_num_groups = act_s_max_num_groups; - } - - if (sh_first_group_id + sh_num_groups > num_groups) { - sh_num_groups = num_groups - sh_first_group_id; - } - - int row_offset = first_group_id * s_gl_stride; - - if (is_async) { - for (int i = 0; i < sh_num_groups; i++) { - if (threadIdx.x < s_sh_stride) { - cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x], - &scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]); - } - } - } else { - for (int i = 0; i < sh_num_groups; i++) { - if (threadIdx.x < s_sh_stride) { - sh_s[(i * s_sh_stride) + threadIdx.x] = scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]; - } - } - } - }; - // Asynchronously fetch the next A, B and s tile from global to the next - // shared memory pipeline location. - auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { - if (pred) { - int4 *sh_a_stage = sh_a + a_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) { - cp_async4_pred( - &sh_a_stage[a_sh_wr_trans[i]], - &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off], - a_sh_wr_pred[i]); - } - int4 *sh_b_stage = sh_b + b_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < (b_sh_wr_iters * b_thread_vecs); i++) { - constexpr int count = div_ceil(b_sh_stride, threads); - int b_gl_idx = b_gl_rd + (i % count) * threads + b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride); - - cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]); - } - - b_gl_rd += b_gl_rd_delta_o; - - if constexpr (has_act_order) { - // Fetch g_idx thread-block portion - int full_pipe = a_off; - int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; - if (cur_k < prob_k && cur_k < slice_k_finish) { - int4 *sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - - int4 const *cur_g_idx_stage_ptr = reinterpret_cast(&g_idx[cur_k]); - - if (threadIdx.x < g_idx_stage) { - cp_async4_pred(&sh_g_idx_stage[threadIdx.x], - &cur_g_idx_stage_ptr[threadIdx.x]); - } - } - } else { - if constexpr (group_blocks != -1) { - int4 *sh_s_stage = sh_s + s_sh_stage * pipe; - - // Only fetch scales if this tile starts a new group - if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { - if (s_sh_wr_pred) { - cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); - } - s_gl_rd += s_gl_rd_delta * s_tb_groups; - } - } - - if constexpr (has_zp && group_blocks != -1) { - int4 *sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - // Only fetch zero points if this tile starts a new group - if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { - if (zp_sh_wr_pred) { - cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; - } - } - } - } - // Insert a fence even when we are winding down the pipeline to ensure that - // waiting is also correct at this point. - cp_async_fence(); - }; - - auto fetch_col_zp_to_shared = [&]() { - if (zp_sh_wr_pred) { - cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - }; - - auto fetch_col_scale_to_shared = [&]() { - if (s_sh_wr_pred) { - cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); - } - }; - - // Wait until the next thread tile has been loaded to shared memory. - auto wait_for_stage = [&]() { - // We only have `stages - 2` active fetches since we are double buffering - // and can only issue the next fetch when it is guaranteed that the previous - // shared memory load is fully complete (as it may otherwise be - // overwritten). - cp_async_wait(); - __syncthreads(); - }; - - // Load the next sub-tile from the current location in the shared memory pipe - // into the current register buffer. - auto fetch_to_registers = [&](int k, int pipe) { - int4 *sh_a_stage = sh_a + a_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - ldsm( - frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); - int4 *sh_b_stage = sh_b + b_sh_stage * pipe; - -#pragma unroll - for (int i = 0; i < b_thread_vecs; i++) { - frag_b_quant[k % 2][i] = *reinterpret_cast( - &sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]); - } - }; - - bool is_same_group[stages]; - int same_group_id[stages]; - - auto init_same_group = [&](int pipe) { - if constexpr (!has_act_order) { - return; - } - - int4 *sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int *sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - int group_id_1 = sh_g_idx_int_ptr[0]; - int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; - - is_same_group[pipe] = group_id_1 == group_id_2; - same_group_id[pipe] = group_id_1; - }; - - auto fetch_scales_to_registers = [&](int k, int full_pipe) { - int pipe = full_pipe % stages; - using IT1 = typename std::conditional_t; - using IT0 = typename std::conditional_t; - constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1); - - if constexpr (!has_act_order) { - // No act-order case - if constexpr (group_blocks == -1) { - // load only when starting a new slice - if (k == 0 && full_pipe == 0 && dequant_skip_flop) { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; - reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; - } - } else if constexpr (group_blocks != -1) { - if constexpr (group_blocks >= thread_k_blocks) { - constexpr int g = group_blocks / thread_k_blocks; - if (pipe % g == 0) { - if (k % b_sh_wr_iters == 0) { - int4 *sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); - reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; - } else { - reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; - } - } - } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) { - auto warp_id = threadIdx.x / 32; - int warp_row = warp_id / tb_n_warps; - - int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; - int cur_group_id = k_blocks / group_blocks2; - - int4 *sh_s_stage = sh_s + s_sh_stage * pipe; - - if constexpr (b_type_id != vllm::kFE2M1f.id()) { - reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; - } else { - reinterpret_cast(&frag_s[k % 2])[0] = reinterpret_cast( - sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; - } - } else if (group_blocks >= b_sh_wr_iters) { - if constexpr (b_type_id != vllm::kFE2M1f.id()) { - reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; - } else { - reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; - } - } - } - - return; - } - - // Act-order case - - // Determine K of the "current" thread-block - int cur_k = slice_k_start + tb_k * full_pipe; - if (cur_k >= prob_k || cur_k >= slice_k_finish) { - return; - } - - // Reset (to current thread-block) since we read g_idx portion from the - // shared memory - cur_k = 0; - - // Progress to current iteration - cur_k += k % b_sh_wr_iters; - - // Determine "position" inside the thread-block (based on warp and - // thread-id) - auto warp_id = threadIdx.x / 32; - int warp_row = warp_id / tb_n_warps; - int warp_col = warp_id % tb_n_warps; - - cur_k += warp_row * 16 * b_sh_wr_iters; - - auto th_id = threadIdx.x % 32; - cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix - - int s_col_shift = - /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + (th_id / 4) * act_s_col_stride; - - if (is_same_group[pipe]) { - if (k % 2 == 0) { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + s_col_shift]; - } else { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); - } - - for (int i = 1; i < 4; i++) { - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); - } - return; - } - - int4 *sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int *sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - constexpr int k_frag_offsets[4] = {0, 1, 8, - 9}; // Tensor core offsets per thread - -#pragma unroll - for (int i = 0; i < 4; i++) { - int actual_k = cur_k + k_frag_offsets[i]; - - int group_id = sh_g_idx_int_ptr[actual_k]; - int rel_group_id = group_id - sh_first_group_id; - - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = sh_s[rel_group_id * s_sh_stride + s_col_shift]; - } - }; - - auto fetch_zp_to_registers = [&](int k, int full_pipe) { - // This code does not handle group_blocks == 0, - // which signifies act_order. - // has_zp implies AWQ, which doesn't have act_order, - static_assert(!has_zp || group_blocks != 0); - - if constexpr (has_zp && !is_zp_float) { - int pipe = full_pipe % stages; - - if constexpr (group_blocks == -1) { - // load only when starting a new slice - if (k == 0 && full_pipe == 0 || is_a_8bit) { -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp))[zp_sh_rd + i]; - } - } - } else if constexpr (group_blocks >= thread_k_blocks) { - constexpr int g = group_blocks / thread_k_blocks; - if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) { - int4 *sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; - } - } - } else { - auto warp_id = threadIdx.x / 32; - - int warp_row = warp_id / tb_n_warps; - - int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; - int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1); - - int4 *sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - sh_zp_stage += cur_group_id * zp_sh_stride; - -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; - } - } - } - - else if constexpr (has_zp && is_zp_float) { - int pipe = full_pipe % stages; - - if constexpr (group_blocks != -1) { - if constexpr (group_blocks >= thread_k_blocks) { - constexpr int g = group_blocks / thread_k_blocks; - if (pipe % g == 0 && k % b_sh_wr_iters == 0) { - int4 *sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); - reinterpret_cast(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd]; - } - } else if (group_blocks < b_sh_wr_iters || k % b_sh_wr_iters == 0) { - auto warp_id = threadIdx.x / 32; - - int warp_row = warp_id / tb_n_warps; - int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; - int cur_group_id = k_blocks / group_blocks; - - int4 *sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - reinterpret_cast(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride]; - } - } - } - }; - - auto dequant_data = [&](int q, scalar_32bit_t *frag_b_ptr, int zp = 0) { - if constexpr (a_type.size_bits() != b_type.size_bits()) { - if constexpr (is_a_8bit && has_zp) { - sub_zp_and_dequant( - q, frag_b_ptr, zp); - } else { - dequant(q, frag_b_ptr); - } - } - }; - - // Execute the actual tensor core matmul of a sub-tile. - bool is_first_matmul_in_slice = true; - auto matmul = [&](int k, int pipe) { - if (is_a_8bit) { - return; - } - int k2 = k % 2; - constexpr int g = group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1; - const bool is_new_zp = (group_blocks == 0) || ((group_blocks > 0) && (group_blocks < b_sh_wr_iters || k == 0)) && (pipe % g == 0) || (group_blocks == -1 && is_first_matmul_in_slice); - if constexpr (has_zp && !is_zp_float) { - if (is_new_zp) { - if constexpr (group_blocks == -1) { - is_first_matmul_in_slice = false; - } - int zp_quant_0, zp_quant_1; - - if constexpr (b_type.size_bits() == 4) { - zp_quant_0 = frag_qzp[k2][0]; - zp_quant_1 = zp_quant_0 >> 8; - } else { - static_assert(b_type.size_bits() == 8); - zp_quant_0 = frag_qzp[k2][0]; - zp_quant_1 = frag_qzp[k2][1]; - } - - dequant_data(zp_quant_0, reinterpret_cast(&frag_zp)); - dequant_data(zp_quant_1, - reinterpret_cast(&frag_zp) + 2); - } - } - if constexpr (!dequant_skip_flop && has_zp && is_zp_float) { - if (is_new_zp) { - reinterpret_cast(&frag_zp)[0] = reinterpret_cast(&frag_zpf[k2])[0]; - } - } - - if constexpr (b_type == vllm::kFE2M1f) { - int s_quant_0 = reinterpret_cast(frag_s[k2])[0]; - int s_quant_1 = reinterpret_cast(frag_s[k2])[1]; - - dequant_fp8_scales( - s_quant_0, reinterpret_cast(&frag_s[k2])); - dequant_fp8_scales( - s_quant_1, reinterpret_cast(&frag_s[k2]) + 2); - } - -// We have the m dimension as the inner loop in order to encourage overlapping -// dequantization and matmul operations. -#pragma unroll - for (int j = 0; j < 4; j++) { - FragB frag_b0; - FragB frag_b1; - int b_quant_0, b_quant_1; - - if constexpr (b_type_id == vllm::kFE2M1f.id()) { - b_quant_1 = frag_b_quant[k2][0][j]; - b_quant_0 = b_quant_1 << 8; - } else if constexpr (b_type.size_bits() == 4) { - b_quant_0 = frag_b_quant[k2][0][j]; - b_quant_1 = b_quant_0 >> 8; - } else { - static_assert(b_type.size_bits() == 8); - int *frag_b_quant_ptr = reinterpret_cast(frag_b_quant[k2]); - b_quant_0 = frag_b_quant_ptr[j * 2 + 0]; - b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; - } - - dequant_data(b_quant_0, reinterpret_cast(&frag_b0)); - dequant_data(b_quant_1, reinterpret_cast(&frag_b1)); - - if constexpr (dequant_skip_flop && has_zp && !is_zp_float && !is_a_8bit) { - sub_zp(frag_b0, frag_zp[j], 0); - sub_zp(frag_b1, frag_zp[j], 1); - } - - // Apply scale to frag_b0 - if constexpr (has_act_order && !is_a_8bit) { - static_assert(group_blocks != -1); - scale4(frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j], - act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0); - scale4(frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], - act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1); - } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && group_blocks == -1 && !is_a_8bit) { - int idx = (threadIdx.x / 4) % 2; - scalar_t2 s2 = Adtype::nums2num2( - reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], - reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 1])[idx]); - if (is_new_zp) - frag_zp[j] = __hmul2(frag_zp[j], s2); - scale_and_sub(frag_b0, s2.x, frag_zp[j].x); - scale_and_sub(frag_b1, s2.y, frag_zp[j].y); - } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1 && !is_a_8bit) { - if (is_new_zp) - frag_zp[j] = __hmul2(frag_zp[j], - *reinterpret_cast(&frag_s[k2][j])); - scale_and_sub(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x); - scale_and_sub(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y); - } else if constexpr (group_blocks != -1 && !is_a_8bit) { - scale(frag_b0, frag_s[k2][j], 0); - scale(frag_b1, frag_s[k2][j], 1); - } - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { - if constexpr (m_block_size_8) { - mma_trans(frag_a[k2][i], frag_b0, frag_b1, - frag_c[i][j][0]); - } else { - mma(frag_a[k2][i], frag_b0, - frag_c[i][j][0]); - mma(frag_a[k2][i], frag_b1, - frag_c[i][j][1]); - } - } - } - }; - - auto matmul_a8 = [&](int k) { - int k2 = k % 2; -#pragma unroll - for (int j = 0; j < 2; j++) { - FragB frag_b[2]; - - if (is_a_8bit && b_type.size_bits() == 4 && !has_zp) { - dequant_data(frag_b_quant[k2][0][j * 2], - reinterpret_cast(&frag_b)); - dequant_data(frag_b_quant[k2][0][j * 2 + 1], - reinterpret_cast(&frag_b) + 2); - } else if (is_a_8bit && b_type.size_bits() == 4 && has_zp) { - int off = (threadIdx.x / 32) % 2 * 2 + j; - int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF; - dequant_data(frag_b_quant[k2][0][j * 2], - reinterpret_cast(&frag_b), zp); - zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF; - dequant_data(frag_b_quant[k2][0][j * 2 + 1], - reinterpret_cast(&frag_b) + 2, zp); - } else { - reinterpret_cast(&frag_b)[0] = reinterpret_cast(&frag_b_quant[k2][j])[0]; - reinterpret_cast(&frag_b)[1] = reinterpret_cast(&frag_b_quant[k2][j])[1]; - } - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { - mma( - frag_a[k2][i], frag_b[0], - (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]); - mma( - frag_a[k2][i], frag_b[1], - (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]); - } - - if constexpr (group_blocks != -1) { - if (group_blocks == 2 || k == 1) { - if constexpr (a_type == vllm::kS8) { - int2 s_vals[2]; - s_vals[0] = { - (int)reinterpret_cast(&frag_s[k2][j * 2][0])[0], - (int)reinterpret_cast(&frag_s[k2][j * 2][0])[1]}; - s_vals[1] = { - (int)reinterpret_cast(&frag_s[k2][j * 2 + 1][0])[0], - (int)reinterpret_cast(&frag_s[k2][j * 2 + 1][0])[1]}; - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int g = 0; g < 4; g++) { - int scale = reinterpret_cast(&s_vals[0])[g % 2]; - *reinterpret_cast(&frag_c[i][j][0][g]) += *reinterpret_cast(&frag_c_tmp[i][j][0][g]) * scale; - frag_c_tmp[i][j][0][g] = 0.0f; - } - -#pragma unroll - for (int g = 0; g < 4; g++) { - int scale = reinterpret_cast(&s_vals[1])[g % 2]; - *reinterpret_cast(&frag_c[i][j][1][g]) += *reinterpret_cast(&frag_c_tmp[i][j][1][g]) * scale; - frag_c_tmp[i][j][1][g] = 0.0f; - } - } - } else { - float2 s_vals[2]; - if constexpr (s_type_id != vllm::kFE8M0fnu.id()) { - static_assert(a_type.size_bits() == 16 || s_type.size_bits() == 16); - s_vals[0] = Cdtype::num22float2(frag_s[k2][j * 2][0]); - s_vals[1] = Cdtype::num22float2(frag_s[k2][j * 2 + 1][0]); - } else { - int32_t *s_vals_int = reinterpret_cast(&s_vals[0]); - int32_t s_vals_e8m0 = *reinterpret_cast(&frag_s[k2][j][0]); - - s_vals_int[0] = (s_vals_e8m0 & 0xFF) << 23; - s_vals_int[1] = (s_vals_e8m0 & 0xFF00) << 15; - s_vals_int[2] = (s_vals_e8m0 & 0xFF0000) << 7; - s_vals_int[3] = (s_vals_e8m0 & 0xFF000000) >> 1; - } - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int g = 0; g < 4; g++) { - float scale = reinterpret_cast(&s_vals[0])[g % 2]; - frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale; - frag_c_tmp[i][j][0][g] = 0.0f; - } - -#pragma unroll - for (int g = 0; g < 4; g++) { - float scale = reinterpret_cast(&s_vals[1])[g % 2]; - frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale; - frag_c_tmp[i][j][1][g] = 0.0f; - } - } - } - } - } - } - }; - - // Since we slice across the k dimension of a tile in order to increase the - // number of warps while keeping the n dimension of a tile reasonable, we have - // multiple warps that accumulate their partial sums of the same output - // location; which we have to reduce over in the end. We do in shared memory. - auto thread_block_reduce = [&]() { - constexpr int red_off = threads / b_sh_stride_threads / 2; - if (red_off >= 1) { - auto red_idx = threadIdx.x / b_sh_stride_threads; - constexpr int red_sh_stride = b_sh_stride_threads * (is_a_8bit ? 2 : 4) * 2; - constexpr int red_sh_delta = b_sh_stride_threads; - int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + (threadIdx.x % b_sh_stride_threads); - - // Parallel logarithmic shared memory reduction. We make sure to avoid any - // unnecessary read or write iterations, e.g., for two warps we write only - // once by warp 1 and read only once by warp 0. - -#pragma unroll - for (int m_block = 0; m_block < thread_m_blocks; m_block++) { -#pragma unroll - for (int i = red_off; i > 0; i /= 2) { - if (i <= red_idx && red_idx < 2 * i) { -#pragma unroll - for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2; - j += (m_block_size_8 ? 2 : 1)) { - int red_sh_wr = red_sh_delta * j + (red_sh_rd - red_sh_stride * i); - if (i < red_off) { - float *c_rd = reinterpret_cast( - &sh_red[red_sh_delta * j + red_sh_rd]); - float *c_wr = reinterpret_cast(&sh_red[red_sh_wr]); -#pragma unroll - for (int k = 0; k < 4; k++) - reinterpret_cast( - frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] - += c_rd[k] + c_wr[k]; - } - sh_red[red_sh_wr] = reinterpret_cast( - &frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j]; - } - } - __syncthreads(); - } - if (red_idx == 0) { -#pragma unroll - for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2; - i += (m_block_size_8 ? 2 : 1)) { - float *c_rd = reinterpret_cast(&sh_red[red_sh_delta * i + red_sh_rd]); -#pragma unroll - for (int j = 0; j < 4; j++) - reinterpret_cast( - frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] - += c_rd[j]; - } - } - __syncthreads(); - } - } - }; - - // Since multiple threadblocks may process parts of the same column slice, we - // finally have to globally reduce over the results. As the striped - // partitioning minimizes the number of such reductions and our outputs are - // usually rather small, we perform this reduction serially in L2 cache. - auto global_reduce_fp16 = [&](bool first = false, bool last = false) { - // We are very careful here to reduce directly in the output buffer to - // maximize L2 cache utilization in this step. To do this, we write out - // results in FP16 (but still reduce with FP32 compute). - constexpr int active_threads = 32 * tb_n_warps; - if (threadIdx.x < active_threads) { - int c_gl_stride = prob_n / 8; - int c_gl_wr_delta_o = 8 * c_gl_stride * (is_a_8bit ? 2 : 1); - int c_gl_wr_delta_i = 4 * (active_threads / 32); - int c_gl_wr; - if constexpr (m_block_size_8) { - c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8; - c_gl_wr += (2 * thread_n_blocks) * slice_col; - } else { - c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) * (is_a_8bit ? 2 : 1) + 4 * (threadIdx.x / 32) + threadIdx.x % 4; - c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1); - } - constexpr int c_sh_wr_delta = active_threads; - auto c_sh_wr = threadIdx.x; - - int row = (threadIdx.x % 32) / 4; - - if (!first) { -// Interestingly, doing direct global accesses here really seems to mess up -// the compiler and lead to slowdowns, hence we also use async-copies even -// though these fetches are not actually asynchronous. -#pragma unroll - for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { - if constexpr (m_block_size_8) { - cp_async4_pred(&sh_red[c_sh_wr + c_sh_wr_delta * i], - &C[c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i], - (threadIdx.x % 4) * 2 + i < prob_m); - } else if constexpr (is_a_8bit) { - int2 *sh_red_int2 = reinterpret_cast(sh_red); - int2 *c_int2 = reinterpret_cast(C); - cp_async2_ca_pred( - &sh_red_int2[c_sh_wr + c_sh_wr_delta * i], - &c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)], - i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); - } else { - cp_async4_pred( - &sh_red[c_sh_wr + c_sh_wr_delta * i], - &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)], - i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); - } - } - cp_async_fence(); - cp_async_wait<0>(); - } - -#pragma unroll - for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { - bool mask = (!m_block_size_8) && (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) || (m_block_size_8) && ((threadIdx.x % 4) * 2 + i < prob_m); - if (mask) { - if (!first) { - c_scalar_t *c_red_f16; - if constexpr (is_a_8bit) { - int2 tmp = reinterpret_cast(sh_red)[c_sh_wr + i * c_sh_wr_delta]; - c_red_f16 = reinterpret_cast(&tmp); - } else { - int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta]; - c_red_f16 = reinterpret_cast(&tmp); - } -#pragma unroll - for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) { - int delta = 0; - if constexpr (m_block_size_8) { - delta = j % 2 == 1 ? -2 : 0; - } - reinterpret_cast( - &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] - += Cdtype::num2float(c_red_f16[j]); - } - } - if (!last) { - c_scalar_t c_f16[is_a_8bit ? 4 : 8]; -#pragma unroll - for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) { - int delta = 0; - if constexpr (m_block_size_8) { - delta = j % 2 == 1 ? -2 : 0; - } - c_f16[j] = Cdtype::float2num(reinterpret_cast( - &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]); - } - if constexpr (m_block_size_8) { - C[c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i] = *reinterpret_cast(c_f16); - } else if constexpr (is_a_8bit) { - int2 *c_int2 = reinterpret_cast(C); - c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] = *reinterpret_cast(c_f16); - } else { - C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] = *reinterpret_cast(c_f16); - } - } - } - } - } - }; - - // Globally reduce over threadblocks that compute the same column block. - // We use a tmp C buffer to reduce in full fp32 precision. - auto global_reduce_fp32 = [&](bool first = false, bool last = false) { - constexpr int tb_m = thread_m_blocks * 16; - constexpr int tb_n = thread_n_blocks * 16; - - constexpr int c_size = tb_m * tb_n * sizeof(float) / 16; - - constexpr int active_threads = 32 * tb_n_warps; - bool is_th_active = threadIdx.x < active_threads; - - constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4; - constexpr int th_size = num_floats * sizeof(float) / 16; - - int c_cur_offset = locks_off * c_size; - - if (!is_th_active) { - return; - } - - if (!first) { - float *frag_c_ptr = reinterpret_cast(&frag_c); -#pragma unroll - for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) { - sh_red[threadIdx.x] = C_tmp[c_cur_offset + active_threads * k + threadIdx.x]; - - float *sh_c_ptr = reinterpret_cast(&sh_red[threadIdx.x]); -#pragma unroll - for (int f = 0; f < 4; f++) { - frag_c_ptr[k * 4 + f] += sh_c_ptr[f]; - } - } - } - - if (!last) { - int4 *frag_c_ptr = reinterpret_cast(&frag_c); -#pragma unroll - for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) { - C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k]; - } - } - }; - - // Write out the reduce final result in the correct layout. We only actually - // reshuffle matrix fragments in this step, the reduction above is performed - // in fragment layout. - auto write_result = [&](bool last) { - int c_gl_stride = prob_n / 8; - constexpr int c_sh_stride = 2 * thread_n_blocks + 1; - int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); - constexpr int c_sh_rd_delta = c_sh_stride * (threads / (2 * thread_n_blocks)); - - int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); - c_gl_wr += (2 * thread_n_blocks) * slice_col; - int c_sh_wr; - if constexpr (m_block_size_8) { - c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) + (threadIdx.x % 32) / 4; - c_sh_wr += 64 * (threadIdx.x / 32); - } else { - c_sh_wr = (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; - c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32); - } - - int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); - - int c_gl_wr_end = c_gl_stride * prob_m; - // We first reorder in shared memory to guarantee the most efficient final - // global write patterns - auto write = [&](int idx, float c0, float c1, FragS &s, FragS &b_bias) { - if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) { - c0 *= global_scale_f32; - c1 *= global_scale_f32; - } - c_scalar_t2 res = Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1)); - - // For per-column quantization we finally apply the scale here (only for - // 4-bit) - if constexpr (!has_act_order && group_blocks == -1 && !is_a_8bit && b_type.size_bits() == 4 && (has_zp && dequant_skip_flop || !has_zp)) { - c_scalar_t2 tmp_scale = s[0]; - if constexpr (m_block_size_8) { - tmp_scale = Cdtype::num2num2( - reinterpret_cast(&s[0])[(threadIdx.x % 8) / 4]); - } - res = __hmul2(res, tmp_scale); - } - if (has_bias && last) { - c_scalar_t2 tmp_bias = b_bias[0]; - if constexpr (m_block_size_8) { - tmp_bias = Cdtype::num2num2( - reinterpret_cast(&b_bias[0])[(threadIdx.x % 8) / 4]); - } - res = __hadd2(res, tmp_bias); - } - - if constexpr (m_block_size_8) { - ((c_scalar_t *)sh_red)[idx] = res.x; - ((c_scalar_t *)sh_red)[idx + 8 * c_sh_stride] = res.y; - } else { - ((c_scalar_t2 *)sh_red)[idx] = res; - } - }; - - if (threadIdx.x / 32 < tb_n_warps) { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int j = 0; j < (is_a_8bit ? 2 : 4); j++) { - if constexpr (m_block_size_8) { - int wr = c_sh_wr + 16 * j; - write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1], - frag_s[j / 2][2 * (j % 2) + 0], - frag_bias[j / 2][2 * (j % 2) + 0]); - write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3], - frag_s[j / 2][2 * (j % 2) + 1], - frag_bias[j / 2][2 * (j % 2) + 1]); - } else { - int wr = c_sh_wr + 8 * j; - write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], - frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0], - frag_bias[j / 2][2 * (j % 2) + 0]); - write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], - frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0], - frag_bias[j / 2][2 * (j % 2) + 0]); - write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], - frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1], - frag_bias[j / 2][2 * (j % 2) + 1]); - write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], - frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1], - frag_bias[j / 2][2 * (j % 2) + 1]); - } - } - c_sh_wr += 16 * (4 * c_sh_stride); - } - } - __syncthreads(); - -#pragma unroll - for (int i = 0; - i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); - i++) { - if (c_gl_wr < c_gl_wr_end) { - if (use_atomic_add && slice_count > 1) { - c_scalar_t2 *C_half2 = reinterpret_cast(&C[c_gl_wr]); - c_scalar_t2 *sh_red_half2 = reinterpret_cast(&sh_red[c_sh_rd]); -#pragma unroll - for (int a = 0; a < 4; a++) { - atomicAdd(&C_half2[a], sh_red_half2[a]); - } - } else { - C[c_gl_wr] = sh_red[c_sh_rd]; - } - c_gl_wr += c_gl_wr_delta; - c_sh_rd += c_sh_rd_delta; - } - } - __syncthreads(); - }; - - // Start global fetch and register load pipelines. - auto start_pipes = [&]() { - -#pragma unroll - for (int i = 0; i < stages - 1; i++) { - if (has_act_order && i == 0) { - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) { - last_g_idx = prob_k - 1; - } - fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], - g_idx[last_g_idx]); - } - - if constexpr (has_zp && !is_zp_float && group_blocks == -1) { - if (i == 0) { - fetch_col_zp_to_shared(); - if constexpr (!dequant_skip_flop) { - fetch_col_scale_to_shared(); - } - } - } - fetch_to_shared(i, i, i < slice_iters); - } - - zero_accums(); - wait_for_stage(); - init_same_group(0); - fetch_to_registers(0, 0); - fetch_scales_to_registers(0, 0); - fetch_zp_to_registers(0, 0); - a_gl_rd += a_gl_rd_delta_o * (stages - 1); - if constexpr (has_act_order) { - slice_k_start_shared_fetch += tb_k * (stages - 1); - } - }; - if (slice_iters) { - start_pipes(); - } - - // Main loop. - while (slice_iters) { - // We unroll over both the global fetch and the register load pipeline to - // ensure all shared memory accesses are static. Note that both pipelines - // have even length meaning that the next iteration will always start at - // index 0. - -#pragma unroll - for (int pipe = 0; pipe < stages;) { -#pragma unroll - for (int k = 0; k < b_sh_wr_iters; k++) { - fetch_to_registers(k + 1, pipe % stages); - fetch_scales_to_registers(k + 1, pipe); - fetch_zp_to_registers(k + 1, pipe); - if (k == b_sh_wr_iters - 2) { - fetch_to_shared((pipe + stages - 1) % stages, pipe, - slice_iters >= stages); - pipe++; - wait_for_stage(); - init_same_group(pipe % stages); - } - - if constexpr (!is_a_8bit) { - matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0)); - } else { - static_assert(group_blocks != 0 && group_blocks != 1); - matmul_a8(k); - } - } - slice_iters--; - if (slice_iters == 0) { - break; - } - } - - a_gl_rd += a_gl_rd_delta_o * stages; - - if constexpr (has_act_order) { - slice_k_start += tb_k * stages; - - if (slice_k_start < prob_k) { - slice_k_start_shared_fetch += tb_k * stages; - int first_group_id = g_idx[slice_k_start]; - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) { - last_g_idx = prob_k - 1; - } - int last_group_id = g_idx[last_g_idx]; - if (last_group_id >= sh_first_group_id + sh_num_groups) { - fetch_act_order_scales_to_shared(false, first_group_id, - last_group_id); - __syncthreads(); - } - } - } - - // Process results and, if necessary, proceed to the next column slice. - // While this pattern may not be the most readable, other ways of writing - // the loop seemed to noticeably worse performance after compilation. - if (slice_iters == 0) { - // convert fp16 accum to fp32 for reduction - if constexpr (use_fp16_accum) { -#pragma unroll - for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) { - float *frag_c_part_float = reinterpret_cast(frag_c) + i * 4; - scalar_t *frag_c_part_half = reinterpret_cast(frag_c_part_float); - -#pragma unroll - for (int i = 3; i >= 0; i--) { - frag_c_part_float[i] = Cdtype::num2float(frag_c_part_half[i]); - } - } - } - - if constexpr (is_a_8bit) { - float frag_a_s[2 * thread_m_blocks]; - - for (int i = 0; i < 2 * thread_m_blocks; i++) { - frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4]; - } - -#pragma unroll - for (int j = 0; j < 2; j++) { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int g = 0; g < 4; g++) { - float c_val = frag_c[i][j][0][g]; - - if constexpr (a_type == vllm::kS8) { - c_val = __int2float_rn(*reinterpret_cast(&c_val)); - } - float s_val = frag_a_s[i * 2 + g / 2]; - frag_c[i][j][0][g] = c_val * s_val; - } -#pragma unroll - for (int g = 0; g < 4; g++) { - float c_val = frag_c[i][j][1][g]; - - if constexpr (a_type == vllm::kS8) { - c_val = __int2float_rn(*reinterpret_cast(&c_val)); - } - float s_val = frag_a_s[i * 2 + g / 2]; - frag_c[i][j][1][g] = c_val * s_val; - } - } - } - } - - cp_async_wait<0>(); - bool last = slice_idx == slice_count - 1; - // For per-column scales, we only fetch them here in the final step before - // write-out - if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) { - if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) { - if (s_sh_wr_pred) { - cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); - } - cp_async_fence(); - } - } - - thread_block_reduce(); - - if (has_bias && last) { - __syncthreads(); - cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd], - threadIdx.x < 16 * thread_n_blocks / 8); - cp_async_fence(); - } - - if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) { - if constexpr (is_a_8bit) { - cp_async_wait<0>(); - __syncthreads(); - if (threadIdx.x / 32 < tb_n_warps) { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; - } - } else if (b_type.size_bits() == 8 || (last || use_atomic_add)) { - cp_async_wait<0>(); - __syncthreads(); - if (threadIdx.x / 32 < tb_n_warps) { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; - reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; - if constexpr (m_block_size_8) { - int idx = (threadIdx.x / 4) % 2; - c_scalar_t2 *frag_s_half2 = reinterpret_cast(frag_s); -#pragma unroll - for (int i = 0; i < 8; i++) { - frag_s_half2[i] = Cdtype::num2num2( - reinterpret_cast(&frag_s_half2[i])[idx]); - } - } - } - } - } - - // For 8-bit channelwise, we apply the scale before the global reduction - // that converts the fp32 results to fp16 (so that we avoid possible - // overflow in fp16) - if constexpr (!has_act_order && group_blocks == -1 && is_a_8bit) { -#pragma unroll - for (int j = 0; j < 2; j++) { - float2 aa[2]; - aa[0] = Cdtype::num22float2(frag_s[0][j * 2][0]); - aa[1] = Cdtype::num22float2(frag_s[0][j * 2 + 1][0]); - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int g = 0; g < 4; g++) { - float scale = reinterpret_cast(&aa[0])[g % 2]; - frag_c[i][j][0][g] *= scale; - } - -#pragma unroll - for (int g = 0; g < 4; g++) { - float scale = reinterpret_cast(&aa[1])[g % 2]; - frag_c[i][j][1][g] *= scale; - } - } - } - } else if (!has_act_order && group_blocks == -1 && b_type.size_bits() == 8 && (has_zp && dequant_skip_flop || !has_zp)) { - if (threadIdx.x / 32 < tb_n_warps) { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int j = 0; j < 4; j++) { - scale_float( - reinterpret_cast(&frag_c[i][j][0][0]), - frag_s[j / 2][2 * (j % 2) + 0]); - scale_float( - reinterpret_cast(&frag_c[i][j][0][2]), - frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]); - - if constexpr (!m_block_size_8) { - scale_float( - reinterpret_cast(&frag_c[i][j][1][0]), - frag_s[j / 2][2 * (j % 2) + 1]); - scale_float( - reinterpret_cast(&frag_c[i][j][1][2]), - frag_s[j / 2][2 * (j % 2) + 1]); - } - } - } - } - } - - if (slice_count > 1 && !use_atomic_add) { - // only globally reduce if there is more than one block in a slice - barrier_acquire(&locks[locks_off], slice_idx); - if (use_fp32_reduce) { - global_reduce_fp32(slice_idx == 0, last); - } else { - global_reduce_fp16(slice_idx == 0, last); - } - barrier_release(&locks[locks_off], last); - } - - if (has_bias && last) { - cp_async_wait<0>(); - __syncthreads(); - reinterpret_cast(&frag_bias)[0] = sh_bias[bias_sh_rd]; - if constexpr (!is_a_8bit) { - reinterpret_cast(&frag_bias)[1] = sh_bias[bias_sh_rd + 4]; - } - __syncthreads(); - } - - if (use_atomic_add && slice_count > 1 && slice_idx != 0) { - wait_negative_and_add(&locks[locks_off]); - } - if (last || use_atomic_add) { - // only the last block in a slice actually writes the result - write_result(last); - } - slice_row = 0; - if (!in_part2) { - slice_col_par += gridDim.x; - } else { - slice_col_par++; - slice_col++; - } - is_first_matmul_in_slice = true; - init_slice(); - - if (slice_iters) { - a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - a_gl_rd += a_gl_rd_delta_o * slice_row; - b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); - b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row; - - bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; - // Update slice k/n for scales loading - if constexpr (has_act_order) { - slice_k_start = tb_k * slice_row; - slice_k_finish = slice_k_start + tb_k * slice_iters; - slice_k_start_shared_fetch = slice_k_start; - slice_n_offset = act_s_col_tb_stride * slice_col; - } else { - if constexpr (group_blocks == -1) { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; - } else if constexpr (group_blocks >= thread_k_blocks) { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + s_sh_stride * slice_col + threadIdx.x; - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + zp_sh_stride * slice_col + threadIdx.x; - } else { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / s_sh_stride) + s_sh_stride * slice_col + threadIdx.x % s_sh_stride; - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + threadIdx.x / zp_sh_stride) + zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; - } - } - start_pipes(); - } - } - } -} - -} // namespace MARLIN_NAMESPACE_NAME - -#endif diff --git a/src/infiniop/ops/awq_marlin_gemm/nvidia/.gitignore b/src/infiniop/ops/awq_marlin_gemm/nvidia/.gitignore deleted file mode 100644 index 7dc482a89..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/nvidia/.gitignore +++ /dev/null @@ -1,3 +0,0 @@ -sm*_kernel_*.cu -kernel_selector.h -kernel_*.cu diff --git a/src/infiniop/ops/awq_marlin_gemm/nvidia/awq_marlin_gemm_nvidia.cu b/src/infiniop/ops/awq_marlin_gemm/nvidia/awq_marlin_gemm_nvidia.cu deleted file mode 100644 index 9dbc855fc..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/nvidia/awq_marlin_gemm_nvidia.cu +++ /dev/null @@ -1,363 +0,0 @@ -#if defined(ENABLE_NVIDIA_API) -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../core/utils.h" -#include "awq_marlin_gemm_nvidia.cuh" -#include "kernel.cuh" - -template -infiniStatus_t awq_marlin_gemm_kernel( - const void *a, - void *c, - const void *b_q_weight, - void *b_bias, - void *b_scales, - void *a_scales, - void *global_scale, - void *b_zeros, - void *g_idx, - void *perm, - int64_t b_type_id, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float, - int size_m, - int size_k, - int size_n, - int b_q_size_0, - int b_q_size_1, - int a_stride_0, - int b_zeros_size_1, - int num_groups, - void *total_buffer, - cudaStream_t stream) { - // scalar_t *a, Tdata *b_scales - vllm::ScalarTypeId a_type_id, c_type_id, s_type_id; - - if constexpr (std::is_same::value) { - a_type_id = vllm::kFloat16.id(); - c_type_id = vllm::kFloat16.id(); - } else if constexpr (std::is_same::value) { - a_type_id = vllm::kBFloat16.id(); - c_type_id = vllm::kBFloat16.id(); - } else { - // 此时c和b_scales类型相同 - if constexpr (std::is_same::value) { - c_type_id = vllm::kFloat16.id(); - } else if constexpr (std::is_same::value) { - c_type_id = vllm::kBFloat16.id(); - } else { - c_type_id = vllm::kBFloat16.id(); - host::RuntimeCheck(c != nullptr, "c must be passed for W4A8-FP4\n"); - } - if constexpr (std::is_same::value) { - a_type_id = vllm::kFE4M3fn.id(); - } else if constexpr (std::is_same::value) { - a_type_id = vllm::kS8.id(); - } else { - host::RuntimeCheck(false, "unsupported `a` scalar_type\n"); - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - s_type_id = c_type_id; - if (b_type_id == vllm::kFE2M1f.id()) { - if constexpr (std::is_same::value) { - s_type_id = vllm::kFE4M3fn.id(); - } else if constexpr (std::is_same::value) { - s_type_id = vllm::kFE8M0fnu.id(); - } else { - host::RuntimeCheck(false, - "When b_type = float4_e2m1f, b_scale scalar type must be", - "float8_e4m3fn (for NVFP4) or float8_e8m0fnu (for MXFP4)."); - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - vllm::ScalarType a_type = vllm::ScalarType::from_id(a_type_id); - vllm::ScalarType b_type = vllm::ScalarType::from_id(b_type_id); - vllm::ScalarType c_type = vllm::ScalarType::from_id(c_type_id); - vllm::ScalarType s_type = vllm::ScalarType::from_id(s_type_id); - - int pack_factor = 32 / b_type.size_bits(); - - // Verify a = [size_m, size_k] - - // Verify b - host::RuntimeCheck( - size_k % MARLIN_NAMESPACE_NAME::tile_size == 0, "size_k = ", size_k, - " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); - host::RuntimeCheck((size_k / MARLIN_NAMESPACE_NAME::tile_size) == b_q_size_0, - "Shape mismatch: b_q_weight.size(0) = ", b_q_size_0, - ", size_k = ", size_k, - ", tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); - host::RuntimeCheck( - b_q_size_1 % MARLIN_NAMESPACE_NAME::tile_size == 0, - "b_q_weight.size(1) = ", b_q_size_1, - " is not divisible by tile_size = ", MARLIN_NAMESPACE_NAME::tile_size); - int actual_size_n = (b_q_size_1 / MARLIN_NAMESPACE_NAME::tile_size) * pack_factor; - host::RuntimeCheck(size_n == actual_size_n, "size_n = ", size_n, - ", actual_size_n = ", actual_size_n); - - // Verify device and strides - - // We use int4 (16 bytes) to load A, so A must aligned to 16 bytes - host::RuntimeCheck(a_stride_0 % 8 == 0, "A.stride(0) must divisible by 8"); - host::RuntimeCheck(reinterpret_cast(a) % 16 == 0, "A must aligned to 16 bytes"); - - if (a_scales != nullptr) { - host::RuntimeCheck(a_type.size_bits() == 8, - "a_scales can only be used for 8bit activation."); - } else { - host::RuntimeCheck(a_type.size_bits() != 8, - "the a_scales parameter must be passed for 8bit activation."); - } - - int device_id = 0; - cudaGetDevice(&device_id); - // thread_k: `k` size of a thread_tile in `weights` (can usually be left as - // auto -1) - int thread_k = -1; - // thread_n: `n` size of a thread_tile in `weights` (can usually be left as - // auto -1) - int thread_n = -1; - // sms: number of SMs to use for the kernel - int sms = -1; - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, device_id); - - // Alloc buffers - float *c_tmp = nullptr; - void *a_tmp = nullptr; - void *workspace = nullptr; - - int c_tmp_bytes = 0; - // Alloc C tmp buffer that is going to be used for the global reduce - - if (use_fp32_reduce) { - int max_m_block_size = (size_m + 16 - 1) / 16 * 16; - max_m_block_size = min(max_m_block_size, 64); - int max_c_tmp_size = sms * max_m_block_size * MARLIN_NAMESPACE_NAME::max_thread_n; - c_tmp_bytes = max_c_tmp_size * sizeof(float); - } - - // Detect groupsize and act_order - - // b_scales = [num_groups, size_n] - // g_idx.size(-1) == size_k && perm.size(-1) == size_k - int a_tmp_bytes = 0; - bool has_act_order = false; - - if (g_idx != nullptr && perm != nullptr) { - has_act_order = true; - } - int group_size = -1; - if (has_act_order) { - a_tmp_bytes = size_m * size_k * sizeof(scalar_t); - if (is_k_full) { - host::RuntimeCheck(num_groups > 1, "For act_order, num_groups must be > 1"); - host::RuntimeCheck(size_k % num_groups == 0, "size_k = ", size_k, - ", is not divisible by num_groups = ", num_groups); - group_size = size_k / num_groups; - } else { - group_size = 0; - } - } else { - if (num_groups > 1) { - host::RuntimeCheck( - size_k % num_groups == 0, "size_k = ", size_k, - ", is not divisible by b_scales.size(0) = ", num_groups); - group_size = size_k / num_groups; - } else { - group_size = -1; - } - } - - int workspace_bytes = sms * sizeof(int); - // ===================== 4. 手动切分指针（核心！） ===================== - uint8_t *ptr = reinterpret_cast(total_buffer); - // 分配 c_tmp - if (use_fp32_reduce && c_tmp_bytes > 0) { - c_tmp = reinterpret_cast(ptr); - ptr += c_tmp_bytes; - } - // 分配 a_tmp - if (has_act_order && a_tmp_bytes > 0) { - a_tmp = ptr; - ptr += a_tmp_bytes; - } - - // 分配 workspace - if (workspace_bytes > 0) { - workspace = ptr; - ptr += workspace_bytes; - } - - if (global_scale != nullptr) { - - host::RuntimeCheck(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn, - "global_scale can only be used for nvfp4 format."); - - } else { - host::RuntimeCheck(!(b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn), - "the global_scale parameter must be passed for nvfp4 format."); - } - // b_bias = [size_n, 1] - bool has_bias = (b_bias != nullptr); - - bool has_zp = (b_zeros != nullptr); - if (has_zp) { - host::RuntimeCheck( - b_type == vllm::kU4 || b_type == vllm::kU8, - "b_type must be u4 or u8 when has_zp = True. Got = ", b_type.str()); - - } else { - host::RuntimeCheck(b_type == vllm::kU4B8 || b_type == vllm::kU8B128 || b_type == vllm::kS4 || b_type == vllm::kS8 || b_type == vllm::kFE4M3fn || b_type == vllm::kFE2M1f, - "b_type must be uint4b8, uint8b128, int4, int8, " - "float8_e4m3fn or float4_e2m1f when has_zp = False. Got = ", - b_type.str()); - } - - if (has_zp && is_zp_float) { - if constexpr (!std::is_same::value) { - host::RuntimeCheck(false, "Computation a_type must be float16 (half) when using float zero points."); - } - } - - // Verify b_zeros - if (has_zp) { - if (is_zp_float) { - // b_zeros = [num_groups, size_n] - host::RuntimeCheck(b_zeros_size_1 == size_n, - "b_zeros dim 1 = ", b_zeros_size_1, - " is not size_n = ", size_n); - host::RuntimeCheck(num_groups != -1, "num_groups must be != -1"); - } else { - - host::RuntimeCheck(b_zeros_size_1 == size_n / pack_factor, - "b_zeros dim 1 = ", b_zeros_size_1, - " is not size_n / pack_factor = ", size_n / pack_factor); - } - } - - // Verify workspace size - host::RuntimeCheck(size_n % MARLIN_NAMESPACE_NAME::min_thread_n == 0, - "size_n = ", size_n, ", is not divisible by min_thread_n = ", - MARLIN_NAMESPACE_NAME::min_thread_n); - - // a_scales和global_scale都必须是float * - - if (a_type.size_bits() == 16) { - host::RuntimeCheck((a_type == c_type), "scalar type of a must be the same with c for 16 bit activation\n"); - } - - marlin::marlin_mm( - a, b_q_weight, c, c_tmp, - b_bias, a_scales, b_scales, - global_scale, b_zeros, g_idx, - perm, a_tmp, size_m, size_n, size_k, a_stride_0, - workspace, a_type, b_type, c_type, s_type, has_bias, - has_act_order, is_k_full, has_zp, num_groups, group_size, device_id, - stream, thread_k, thread_n, sms, - use_atomic_add, use_fp32_reduce, is_zp_float); - return INFINI_STATUS_SUCCESS; -} - -namespace op::awq_marlin_gemm::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t b_bias_desc, - infiniopTensorDescriptor_t b_scales_desc, - infiniopTensorDescriptor_t a_scales_desc, - infiniopTensorDescriptor_t global_scales_desc, - infiniopTensorDescriptor_t b_zeros_desc, - infiniopTensorDescriptor_t g_idx_desc, - infiniopTensorDescriptor_t perm_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = AwqMarlinGemmInfo::create(out_desc, a_desc, b_desc, b_bias_desc, b_scales_desc, a_scales_desc, global_scales_desc, b_zeros_desc, g_idx_desc, perm_desc); - size_t size_m = a_desc->dim(0); - size_t size_k = a_desc->dim(1); - int a_tmp_bytes = 0; - int c_tmp_bytes = 0; - - int device_id = 0; - cudaGetDevice(&device_id); - int sms = -1; - cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, device_id); - int workspace_bytes = sms * sizeof(int); - int max_m_block_size = (size_m + 16 - 1) / 16 * 16; - max_m_block_size = min(max_m_block_size, 64); - int max_c_tmp_size = sms * max_m_block_size * MARLIN_NAMESPACE_NAME::max_thread_n; - c_tmp_bytes = max_c_tmp_size * sizeof(float); - - if (g_idx_desc->numel() > 0 && perm_desc->numel() > 0) { - a_tmp_bytes = size_m * size_k * infiniSizeOf(a_desc->dtype()); - } - - size_t workspace_size = c_tmp_bytes + a_tmp_bytes + workspace_bytes; - - *desc_ptr = new Descriptor( - new Opaque{handle->internal()}, - result.take(), - workspace_size, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t -Descriptor::calculate( - void *workspace, size_t workspace_size, - void *c, - const void *a, - const void *b, - void *b_bias, - void *b_scales, - void *a_scales, - void *global_scales, - void *b_zeros, - void *g_idx, - void *perm, - int64_t b_q_type_id, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float, - void *stream_) const { - - cudaStream_t stream = (cudaStream_t)stream_; - int size_m = static_cast(_info.size_m); - int size_k = static_cast(_info.size_k); - int size_n = static_cast(_info.size_n); - int b_q_size_0 = static_cast(_info.b_q_size_0); - int b_q_size_1 = static_cast(_info.b_q_size_1); - int b_zeros_size_1 = static_cast(_info.b_zeros_size_1); - int a_stride_0 = static_cast(_info.a_stride_0); - int num_groups = _info.num_groups; - -#define MARLIN(SCALAR_T, TDATA) \ - awq_marlin_gemm_kernel(a, c, b, b_bias, b_scales, a_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float, size_m, size_k, size_n, b_q_size_0, b_q_size_1, a_stride_0, b_zeros_size_1, num_groups, workspace, stream) - - if (_info.a_dtype == INFINI_DTYPE_F16) { - return MARLIN(half, half); - } else if (_info.a_dtype == INFINI_DTYPE_BF16) { - return MARLIN(__nv_bfloat16, __nv_bfloat16); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::awq_marlin_gemm::nvidia -#endif diff --git a/src/infiniop/ops/awq_marlin_gemm/nvidia/awq_marlin_gemm_nvidia.cuh b/src/infiniop/ops/awq_marlin_gemm/nvidia/awq_marlin_gemm_nvidia.cuh deleted file mode 100644 index 2b4a20bbe..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/nvidia/awq_marlin_gemm_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __AWQ_MARLIN_GEMM_CUDA_CUH__ -#define __AWQ_MARLIN_GEMM_CUDA_CUH__ - -#include "../awq_marlin_gemm.h" - -DESCRIPTOR(nvidia) - -#endif // __AWQ_MARLIN_GEMM_CUDA_CUH__ diff --git a/src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py b/src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py deleted file mode 100644 index ab78d0ce2..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/nvidia/generate_kernels.py +++ /dev/null @@ -1,308 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import glob -import itertools -import os -import subprocess -import sys - -import jinja2 - -ARCHS = [] -SUPPORT_FP8 = False -SUPPORT_SM75 = False -SUPPORT_SM80 = False -for arch in sys.argv[1].split(","): - arch = arch[: arch.index(".") + 2].replace(".", "") - arch = int(arch) - # only SM89 and SM120 fully support - # mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32. - # SM90 and SM100 can use this PTX, but it’s simulated - # with FP16 MMA, so it cannot achieve any acceleration. - if arch in [89, 120]: - SUPPORT_FP8 = True - if arch >= 80: - SUPPORT_SM80 = True - if arch == 75: - SUPPORT_SM75 = True - -FILE_HEAD_COMMENT = """ -// auto generated by generate_kernels.py -// clang-format off -""".lstrip() - - -FILE_HEAD = ( - FILE_HEAD_COMMENT - + """ -#include "../marlin/kernel.h" -#include "../marlin/marlin_template.h" - -namespace MARLIN_NAMESPACE_NAME { -""" -) - -TEMPLATE = ( - "template __global__ void Marlin<" - "{{a_type_id}}, " - "{{b_type_id}}, " - "{{c_type_id}}, " - "{{s_type_id}}, " - "{{threads}}, " - "{{thread_m_blocks}}, " - "{{thread_n_blocks}}, " - "{{thread_k_blocks}}, " - "{{m_block_size_8}}, " - "{{stages}}, " - "{{group_blocks}}, " - "{{is_zp_float}}>" - "( MARLIN_KERNEL_PARAMS );" -) - -THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)] - -THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4] - -QUANT_CONFIGS = [ - # AWQ-INT4 - { - "b_type": "kU4", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 2, 4, 8], - }, - # GPTQ-INT4 - { - "b_type": "kU4B8", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 0, 2, 4, 8], - }, - # GPTQ-INT8 - { - "b_type": "kU8B128", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 0, 2, 4, 8], - }, - # FP8 - { - "b_type": "kFE4M3fn", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [-1, 8], - }, - # NVFP4 - { - "b_type": "kFE2M1f", - "s_type": "kFE4M3fn", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [1], - }, - # MXFP4 - { - "a_type": ["kBFloat16"], - "b_type": "kFE2M1f", - "s_type": "kFE8M0fnu", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": THREAD_M_BLOCKS, - "group_blocks": [2], - }, - # AWQ-INT4 with INT8 activation - { - "a_type": ["kS8"], - "b_type": "kU4", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # GPTQ-INT4 with INT8 activation - { - "a_type": ["kS8"], - "b_type": "kU4B8", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # GPTQ-INT4 with FP8 activation - { - "a_type": ["kFE4M3fn"], - "b_type": "kU4B8", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # AWQ-INT4 with FP8 activation - { - "a_type": ["kFE4M3fn"], - "b_type": "kU4", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [-1, 2, 4, 8], - }, - # MXFP4 with FP8 activation - { - "a_type": ["kFE4M3fn"], - "b_type": "kFE2M1f", - "c_type": ["kBFloat16"], - "s_type": "kFE8M0fnu", - "thread_configs": THREAD_CONFIGS, - "thread_m_blocks": [1, 2, 3, 4], - "group_blocks": [2], - }, -] - - -def remove_old_kernels(): - for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"): - subprocess.call(["rm", "-f", filename]) - - filename = os.path.dirname(__file__) + "/kernel_selector.h" - subprocess.call(["rm", "-f", filename]) - - -def generate_new_kernels(): - result_dict = {} - sm_75_result_dict = {} - - for quant_config in QUANT_CONFIGS: - c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"]) - a_types = quant_config.get("a_type", ["kFloat16", "kBFloat16"]) - b_type = quant_config["b_type"] - is_zp_float = quant_config.get("is_zp_float", False) - all_group_blocks = quant_config["group_blocks"] - all_m_blocks = quant_config["thread_m_blocks"] - all_thread_configs = quant_config["thread_configs"] - - for a_type, c_type in itertools.product(a_types, c_types): - if not SUPPORT_FP8 and a_type == "kFE4M3fn": - continue - if "16" in a_type and "16" in c_type and a_type != c_type: - continue - s_type = quant_config.get("s_type", c_type) - if (a_type, b_type, c_type) not in result_dict: - result_dict[(a_type, b_type, c_type)] = [] - if a_type in ["kFloat16", "kS8"] and c_type == "kFloat16": - sm_75_result_dict[(a_type, b_type, c_type)] = [] - - for group_blocks, m_blocks, thread_configs in itertools.product( - all_group_blocks, all_m_blocks, all_thread_configs - ): - thread_k, thread_n, threads = thread_configs - - if threads == 256: - # for small batch (m_blocks == 1), - # we only need (128, 128, 256) - # for large batch (m_blocks > 1), - # we only need (64, 256, 256) - if m_blocks <= 1 and (thread_k, thread_n) != (128, 128): - continue - if m_blocks > 1 and (thread_k, thread_n) != (64, 256): - continue - - config = { - "threads": threads, - "s_type": s_type, - "thread_m_blocks": max(m_blocks, 1), - "thread_k_blocks": thread_k // 16, - "thread_n_blocks": thread_n // 16, - "m_block_size_8": "true" if m_blocks == 0.5 else "false", - "stages": 4, - "group_blocks": group_blocks, - "is_zp_float": "true" if is_zp_float else "false", - } - - if SUPPORT_SM80: - result_dict[(a_type, b_type, c_type)].append(config) - if (a_type, b_type, c_type) in sm_75_result_dict and SUPPORT_SM75: - config_sm75 = config.copy() - config_sm75["stages"] = 2 - sm_75_result_dict[(a_type, b_type, c_type)].append(config_sm75) - - kernel_selector_str = FILE_HEAD_COMMENT - - for result_dict_tmp in [result_dict, sm_75_result_dict]: - for (a_type, b_type, c_type), config_list in result_dict_tmp.items(): - all_template_str_list = [] - if not config_list: - continue - for config in config_list: - s_type = config["s_type"] - template_str = jinja2.Template(TEMPLATE).render( - a_type_id=f"vllm::{a_type}.id()", - b_type_id=f"vllm::{b_type}.id()", - c_type_id=f"vllm::{c_type}.id()", - s_type_id=f"vllm::{s_type}.id()", - **config, - ) - all_template_str_list.append(template_str) - - conditions = [ - f"a_type == vllm::{a_type}", - f"b_type == vllm::{b_type}", - f"c_type == vllm::{c_type}", - f"s_type == vllm::{s_type}", - f"threads == {config['threads']}", - f"thread_m_blocks == {config['thread_m_blocks']}", - f"thread_n_blocks == {config['thread_n_blocks']}", - f"thread_k_blocks == {config['thread_k_blocks']}", - f"m_block_size_8 == {config['m_block_size_8']}", - f"stages == {config['stages']}", - f"group_blocks == {config['group_blocks']}", - f"is_zp_float == {config['is_zp_float']}", - ] - conditions = " && ".join(conditions) - - if kernel_selector_str == FILE_HEAD_COMMENT: - kernel_selector_str += f"if ({conditions})\n kernel = " - else: - kernel_selector_str += f"else if ({conditions})\n kernel = " - - kernel_template2 = ( - "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, " - "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, " - "{{thread_n_blocks}}, {{thread_k_blocks}}, " - "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, " - "{{is_zp_float}}>;" - ) - - kernel_selector_str += ( - jinja2.Template(kernel_template2).render( - a_type_id=f"vllm::{a_type}.id()", - b_type_id=f"vllm::{b_type}.id()", - c_type_id=f"vllm::{c_type}.id()", - s_type_id=f"vllm::{s_type}.id()", - **config, - ) - + "\n" - ) - - file_content = FILE_HEAD + "\n\n" - file_content += "\n\n".join(all_template_str_list) + "\n\n}\n" - if a_type == "kFE4M3fn": - filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" - elif result_dict_tmp is sm_75_result_dict: - filename = f"sm75_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" - else: - filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu" - - filename = filename.lower() - - with open(os.path.join(os.path.dirname(__file__), filename), "w") as f: - f.write(file_content) - - if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT: - kernel_selector_str += ( - "else if (a_type == vllm::kFE4M3fn)\n" - " host::RuntimeCheck(false, " - '"marlin kernel with fp8 activation is not built.");' - ) - - with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f: - f.write(kernel_selector_str) - - -if __name__ == "__main__": - remove_old_kernels() - generate_new_kernels() diff --git a/src/infiniop/ops/awq_marlin_gemm/nvidia/kernel.cuh b/src/infiniop/ops/awq_marlin_gemm/nvidia/kernel.cuh deleted file mode 100644 index 80b9fb5a3..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/nvidia/kernel.cuh +++ /dev/null @@ -1,549 +0,0 @@ -/* - * Modified by Neural Magic - * Copyright (C) Marlin.2024 Elias Frantar - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * Adapted from https://github.com/IST-DASLab/marlin - */ - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -#include "../core/utils.h" -#include "../marlin/kernel.h" - -#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ - static_assert(std::is_same::value || std::is_same::value, \ - "only float16 and bfloat16 is supported"); - -namespace marlin { - -__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; - -using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750 - -__global__ void permute_cols_kernel(int4 const *__restrict__ a_int4_ptr, - int const *__restrict__ perm_int_ptr, - int4 *__restrict__ out_int4_ptr, int size_m, - int size_k, int lda, int block_rows) {} - -} // namespace marlin - -template -void awq_marlin_gemm_kernel( - const void *a, - void *c, - const void *b_q_weight, - void *b_bias, - void *b_scales, - void *a_scales, - void *global_scale, - void *b_zeros, - void *g_idx, - void *perm, - int64_t b_type_id, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float, - int size_m, - int size_k, - int size_n, - int b_q_size_0, - int b_q_size_1, - int a_stride_0, - int b_zeros_size_1, - int num_groups, - cudaStream_t stream) { -} - -#else - -// For a given "a" of size [M,K] performs a permutation of the K columns based -// on the given "perm" indices. -__global__ void permute_cols_kernel(int4 const *__restrict__ a_int4_ptr, - int const *__restrict__ perm_int_ptr, - int4 *__restrict__ out_int4_ptr, int size_m, - int size_k, int lda, int block_rows) { - auto start_row = block_rows * blockIdx.x; - int finish_row = start_row + block_rows; - if (finish_row > size_m) { - finish_row = size_m; - } - int cur_block_rows = finish_row - start_row; - - int input_row_stride = lda * sizeof(half) / 16; - int output_row_stride = size_k * sizeof(half) / 16; - - auto permute_row = [&](int row) { - int iters = size_k / default_threads; - int rest = size_k % default_threads; - - int input_offset = row * input_row_stride; - int output_offset = row * output_row_stride; - - half const *a_row_half = reinterpret_cast(a_int4_ptr + input_offset); - half *out_half = reinterpret_cast(out_int4_ptr + output_offset); - - int base_k = 0; - - for (int i = 0; i < iters; i++) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - - base_k += default_threads; - } - - if (rest) { - if (threadIdx.x < rest) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - } - } - }; - - for (int i = 0; i < cur_block_rows; i++) { - int cur_row = start_row + i; - if (cur_row < size_m) { - permute_row(cur_row); - } - } -} - -typedef struct -{ - int thread_k; - int thread_n; - int num_threads; -} thread_config_t; - -thread_config_t small_batch_thread_configs[] = { - // Ordered by priority - - // thread_k, thread_n, num_threads - {128, 128, 256}, - {64, 128, 128}, - {128, 64, 128}}; - -thread_config_t large_batch_thread_configs[] = { - // Ordered by priority - - // thread_k, thread_n, num_threads - {64, 256, 256}, - {64, 128, 128}, - {128, 64, 128}}; - -typedef struct -{ - int blocks_per_sm; - thread_config_t tb_cfg; -} exec_config_t; - -int get_scales_cache_size(thread_config_t const &th_config, int prob_m, - int prob_n, int prob_k, int num_bits, int group_size, - bool has_act_order, bool is_k_full, int stages) { - bool cache_scales_chunk = has_act_order && !is_k_full; - - int tb_n = th_config.thread_n; - int tb_k = th_config.thread_k; - - // Get max scale groups per thread-block - int tb_groups; - if (group_size == -1) { - tb_groups = 1; - } else if (group_size == 0) { - tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size - } else { - tb_groups = div_ceil(tb_k, group_size); - } - - if (cache_scales_chunk) { - int load_groups = tb_groups * stages * 2; // Chunk size is 2x pipeline over dim K - load_groups = max(load_groups, 32); // We load at least 32 scale groups - return load_groups * tb_n * 2; - } else { - int tb_scales = tb_groups * tb_n * 2; - - return tb_scales * stages; - } -} - -int get_kernel_cache_size(thread_config_t const &th_config, int thread_m_blocks, - int prob_m, int prob_n, int prob_k, int num_bits, - int group_size, bool has_act_order, bool is_k_full, - int has_zp, bool is_zp_float, bool is_a_8bit, - int stages) { - int pack_factor = 32 / num_bits; - - // Get B size - int tb_k = th_config.thread_k; - int tb_n = th_config.thread_n; - int tb_m = thread_m_blocks * 16; - int sh_a_size = stages * (tb_m * tb_k) * (is_a_8bit ? 1 : 2); - int sh_b_size = stages * (tb_k * tb_n / pack_factor) * 4; - int sh_red_size = tb_m * (tb_n + 8) * 2; - int sh_bias_size = tb_n * 2; - int tmp_size = (sh_b_size > sh_red_size ? sh_red_size : sh_b_size) + sh_bias_size; - tmp_size = max(max(sh_b_size, sh_red_size), tmp_size); - - int sh_s_size = get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, - group_size, has_act_order, is_k_full, stages); - int sh_g_idx_size = has_act_order && !is_k_full ? stages * tb_k / 4 : 0; - int sh_zp_size = 0; - if (has_zp) { - if (is_zp_float) { - sh_zp_size = sh_s_size; - } else if (num_bits == 4) { - sh_zp_size = sh_s_size / 4; - } else if (num_bits == 8) { - sh_zp_size = sh_s_size / 2; - } - } - - int total_size = tmp_size + sh_a_size + sh_s_size + sh_zp_size + sh_g_idx_size; - - return total_size; -} - -bool is_valid_config(thread_config_t const &th_config, int thread_m_blocks, - int prob_m, int prob_n, int prob_k, int num_bits, - int group_size, bool has_act_order, bool is_k_full, - int has_zp, bool is_zp_float, bool is_a_8bit, int stages, - int max_shared_mem) { - // Sanity - if (th_config.thread_k == -1 || th_config.thread_n == -1 || th_config.num_threads == -1) { - return false; - } - - // Verify K/N are divisible by thread K/N - if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { - return false; - } - - // Verify min for thread K/N - if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) { - return false; - } - - // num_threads must be at least 128 (= 4 warps) - if (th_config.num_threads < 128) { - return false; - } - - // Check that pipeline fits into cache - int cache_size = get_kernel_cache_size( - th_config, thread_m_blocks, prob_m, prob_n, prob_k, num_bits, group_size, - has_act_order, is_k_full, has_zp, is_zp_float, is_a_8bit, stages); - return cache_size <= max_shared_mem; -} - -MarlinFuncPtr get_marlin_kernel( - const vllm::ScalarType a_type, const vllm::ScalarType b_type, - const vllm::ScalarType c_type, const vllm::ScalarType s_type, - int thread_m_blocks, int thread_n_blocks, int thread_k_blocks, - bool m_block_size_8, bool has_act_order, bool has_zp, int group_blocks, - int threads, bool is_zp_float, int stages) { - int num_bits = b_type.size_bits(); - auto kernel = MarlinDefault; - -#include "kernel_selector.h" - - return kernel; -} - -exec_config_t determine_exec_config( - const vllm::ScalarType &a_type, const vllm::ScalarType &b_type, - const vllm::ScalarType &c_type, const vllm::ScalarType &s_type, int prob_m, - int prob_n, int prob_k, int thread_m_blocks, bool m_block_size_8, - int num_bits, int group_size, bool has_act_order, bool is_k_full, - bool has_zp, bool is_zp_float, int is_a_8bit, int stages, - int max_shared_mem, int sms) { - exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}}; - thread_config_t *thread_configs = thread_m_blocks > 1 - ? large_batch_thread_configs - : small_batch_thread_configs; - int thread_configs_size = thread_m_blocks > 1 - ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t) - : sizeof(small_batch_thread_configs) / sizeof(thread_config_t); - - for (int i = 0; i < thread_configs_size; i++) { - thread_config_t th_config = thread_configs[i]; - - if (!is_valid_config(th_config, thread_m_blocks, prob_m, prob_n, prob_k, - num_bits, group_size, has_act_order, is_k_full, has_zp, - is_zp_float, is_a_8bit, stages, - max_shared_mem - 512)) { - continue; - } - - int cache_size = get_kernel_cache_size(th_config, thread_m_blocks, prob_m, - prob_n, prob_k, num_bits, group_size, - has_act_order, is_k_full, has_zp, - is_zp_float, is_a_8bit, stages); - - int group_blocks = 0; - if (!has_act_order) { - group_blocks = group_size == -1 ? -1 : group_size / 16; - } - - auto kernel = get_marlin_kernel(a_type, b_type, c_type, s_type, thread_m_blocks, - th_config.thread_n / 16, th_config.thread_k / 16, - m_block_size_8, has_act_order, has_zp, group_blocks, - th_config.num_threads, is_zp_float, stages); - - if (kernel == MarlinDefault) { - continue; - } - - return {1, th_config}; - } - - return exec_cfg; -} - -void marlin_mm(const void *A, const void *B, void *C, void *C_tmp, void *b_bias, - void *a_s, void *b_s, void *g_s, void *zp, void *g_idx, - void *perm, void *a_tmp, int prob_m, int prob_n, int prob_k, - int lda, void *workspace, vllm::ScalarType const &a_type, - vllm::ScalarType const &b_type, vllm::ScalarType const &c_type, - vllm::ScalarType const &s_type, bool has_bias, - bool has_act_order, bool is_k_full, bool has_zp, int num_groups, - int group_size, int dev, cudaStream_t stream, int thread_k_init, - int thread_n_init, int sms, bool use_atomic_add, - bool use_fp32_reduce, bool is_zp_float) { - bool is_a_8bit = a_type.size_bits() == 8; - host::RuntimeCheck(prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, - ", ", prob_n, ", ", prob_k, "]"); - - int group_blocks = 0; - if (has_act_order) { - if (is_k_full) { - host::RuntimeCheck(group_size != -1); - group_blocks = group_size / 16; - host::RuntimeCheck(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); - } else { - host::RuntimeCheck(group_size == 0); - group_blocks = 0; - } - } else { - if (group_size == -1) { - group_blocks = -1; - } else { - group_blocks = group_size / 16; - host::RuntimeCheck(prob_k % group_blocks == 0, "prob_k = ", prob_k, - " is not divisible by group_blocks = ", group_blocks); - } - } - - int num_bits = b_type.size_bits(); - const int4 *A_ptr = (const int4 *)A; - const int4 *B_ptr = (const int4 *)B; - int4 *C_ptr = (int4 *)C; - int4 *C_tmp_ptr = (int4 *)C_tmp; - - const int4 *bias_ptr = (const int4 *)b_bias; - const float *a_s_ptr = (const float *)a_s; - const int4 *b_s_ptr = (const int4 *)b_s; - const float *g_s_ptr = (const float *)g_s; - - const int4 *zp_ptr = (const int4 *)zp; - const int *g_idx_ptr = (const int *)g_idx; - const int *perm_ptr = (const int *)perm; - int4 *a_tmp_ptr = (int4 *)a_tmp; - int *locks = (int *)workspace; - - if (has_act_order) { - // Permute A columns - int block_rows = div_ceil(prob_m, sms); - // avoid ">>>" being formatted to "> > >" - // clang-format off - permute_cols_kernel<<>>( - A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, lda, block_rows); - // clang-format on - A_ptr = a_tmp_ptr; - lda = prob_k; - - // If we have a full K, then we can run the non-act-order version of Marlin - // (since the weight rows are reordered by increasing group ids, and by - // having a full K, we have full original groups) - if (is_k_full) { - has_act_order = false; - } - } - - int max_shared_mem = 0; - cudaDeviceGetAttribute(&max_shared_mem, - cudaDevAttrMaxSharedMemoryPerBlockOptin, dev); - host::RuntimeCheck(max_shared_mem > 0); - - int major_capability, minor_capability; - cudaDeviceGetAttribute(&major_capability, cudaDevAttrComputeCapabilityMajor, - dev); - cudaDeviceGetAttribute(&minor_capability, cudaDevAttrComputeCapabilityMinor, - dev); - host::RuntimeCheck(major_capability * 10 + minor_capability >= 75, - "marlin kernel only support Turing or newer GPUs."); - int stages = 4; - if (major_capability == 7 && minor_capability == 5) { - stages = 2; - host::RuntimeCheck(a_type == vllm::kFloat16 || a_type == vllm::kS8, - "Turing only support FP16 or INT8 activation."); - } - if (a_type == vllm::kFE4M3fn) { - host::RuntimeCheck( - major_capability * 10 + minor_capability == 89 || major_capability * 10 + minor_capability == 120, - "Marlin W4A8-FP8 only support SM89 or SM120 device (It is slower than " - "Marlin W4A16 on other devices)."); - } - - int max_par = 16; - if (prob_n <= 4096) { - max_par = 16 * 8; - } - int max_shared_mem_new = max_shared_mem; - int rest_m = prob_m; - int max_thread_m_blocks = 4; - while (rest_m) { - int par_count = rest_m / (max_thread_m_blocks * 16); - if (par_count > max_par) { - par_count = max_par; - } - int prob_m_split = par_count > 0 ? (par_count * (max_thread_m_blocks * 16)) : rest_m; - - int thread_k = thread_k_init; - int thread_n = thread_n_init; - - int thread_m_blocks = min(div_ceil(prob_m_split, 16), max_thread_m_blocks); - int m_block_size_8 = prob_m_split <= 8 && a_type.size_bits() == 16; - - // Set thread config - exec_config_t exec_cfg; - thread_config_t thread_tfg; - if (thread_k != -1 && thread_n != -1) { - thread_tfg = thread_config_t{thread_k, thread_n, default_threads}; - exec_cfg = exec_config_t{1, thread_tfg}; - host::RuntimeCheck(prob_n % thread_n == 0, "prob_n = ", prob_n, - " is not divisible by thread_n = ", thread_n); - host::RuntimeCheck(prob_k % thread_k == 0, "prob_k = ", prob_k, - " is not divisible by thread_k = ", thread_k); - } else { - // Auto config - exec_cfg = determine_exec_config( - a_type, b_type, c_type, s_type, prob_m_split, prob_n, prob_k, - thread_m_blocks, m_block_size_8, num_bits, group_size, has_act_order, - is_k_full, has_zp, is_zp_float, is_a_8bit, stages, max_shared_mem, - sms); - thread_tfg = exec_cfg.tb_cfg; - if (thread_tfg.thread_n != -1) { - if (prob_n / thread_tfg.thread_n * div_ceil(prob_m_split, thread_m_blocks * 16) * 4 <= sms) { - if (is_valid_config({128, 64, 128}, thread_m_blocks, prob_m_split, - prob_n, prob_k, num_bits, group_size, - has_act_order, is_k_full, has_zp, is_zp_float, - is_a_8bit, stages, max_shared_mem_new)) { - thread_tfg = {128, 64, 128}; - exec_cfg = {1, thread_tfg}; - } - } - } - - if (thread_tfg.thread_k == -1 && max_thread_m_blocks > 1) { - max_thread_m_blocks--; - continue; - } - } - - int num_threads = thread_tfg.num_threads; - thread_k = thread_tfg.thread_k; - thread_n = thread_tfg.thread_n; - int blocks = sms * exec_cfg.blocks_per_sm; - if (exec_cfg.blocks_per_sm > 1) { - max_shared_mem_new = max_shared_mem / exec_cfg.blocks_per_sm - 1024; - } - - int thread_k_blocks = thread_k / 16; - int thread_n_blocks = thread_n / 16; - - host::RuntimeCheck( - is_valid_config(thread_tfg, thread_m_blocks, prob_m_split, prob_n, - prob_k, num_bits, group_size, has_act_order, is_k_full, - has_zp, is_zp_float, is_a_8bit, stages, - max_shared_mem_new), - "Invalid thread config: thread_m_blocks = ", thread_m_blocks, - ", thread_k = ", thread_tfg.thread_k, - ", thread_n = ", thread_tfg.thread_n, - ", num_threads = ", thread_tfg.num_threads, " for MKN = [", prob_m, - ", ", prob_k, ", ", prob_n, "] and num_bits = ", num_bits, - ", prob_m_split = ", prob_m_split, ", group_size = ", group_size, - ", has_act_order = ", has_act_order, ", is_k_full = ", is_k_full, - ", has_zp = ", has_zp, ", is_zp_float = ", is_zp_float, - ", stages = ", stages, ", max_shared_mem_new = ", max_shared_mem_new); - - auto kernel = get_marlin_kernel( - a_type, b_type, c_type, s_type, thread_m_blocks, thread_n_blocks, - thread_k_blocks, m_block_size_8, has_act_order, has_zp, group_blocks, - num_threads, is_zp_float, stages); - - if (kernel == MarlinDefault) { - host::RuntimeCheck(false, "Unsupported shapes: MNK = [", prob_m, ", ", prob_n, - ", ", prob_k, "]", ", has_act_order = ", has_act_order, - ", num_groups = ", num_groups, ", group_size = ", group_size, - ", prob_m_split = ", prob_m_split, - ", thread_m_blocks = ", thread_m_blocks, - ", thread_n_blocks = ", thread_n_blocks, - ", thread_k_blocks = ", thread_k_blocks, - ", num_threads = ", num_threads, ", num_bits = ", num_bits); - } - - cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; - host::RuntimeCheck( - cudaStreamIsCapturing(stream, &capture_status) == cudaSuccess, - "cudaStreamIsCapturing failed"); - if (capture_status == cudaStreamCaptureStatusNone) { - host::RuntimeCheck( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem_new) == cudaSuccess, - "cudaFuncSetAttribute failed"); - } - - bool part_use_atomic_add = use_atomic_add && div_ceil(prob_m_split, 64) * prob_n <= 2048; - - // avoid ">>>" being formatted to "> > >" - // clang-format off - kernel<<>>( - A_ptr, B_ptr, C_ptr, C_tmp_ptr, bias_ptr, a_s_ptr, b_s_ptr, g_s_ptr, zp_ptr, - g_idx_ptr, num_groups, - prob_m_split, prob_n, prob_k, lda, locks, has_bias, part_use_atomic_add, - use_fp32_reduce, max_shared_mem_new); - // clang-format on - - bool is_a_8bit = a_type.size_bits() == 8; - A_ptr += prob_m_split * (lda / (is_a_8bit ? 16 : 8)); - a_s_ptr += prob_m_split; - C_ptr += prob_m_split * (prob_n / 8); - rest_m -= prob_m_split; - } -} - -} // namespace marlin - -#endif diff --git a/src/infiniop/ops/awq_marlin_gemm/operator.cc b/src/infiniop/ops/awq_marlin_gemm/operator.cc deleted file mode 100644 index 7ef8246b1..000000000 --- a/src/infiniop/ops/awq_marlin_gemm/operator.cc +++ /dev/null @@ -1,126 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/awq_marlin_gemm.h" - -#if defined ENABLE_NVIDIA_API -#include "nvidia/awq_marlin_gemm_nvidia.cuh" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAwqMarlinGemmDescriptor( - infiniopHandle_t handle, - infiniopAwqMarlinGemmDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t b_bias_desc, - infiniopTensorDescriptor_t b_scales_desc, - infiniopTensorDescriptor_t a_scales_desc, - infiniopTensorDescriptor_t global_scales_desc, - infiniopTensorDescriptor_t b_zeros_desc, - infiniopTensorDescriptor_t g_idx_desc, - infiniopTensorDescriptor_t perm_desc) { -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::awq_marlin_gemm::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - out_desc, \ - a_desc, \ - b_desc, \ - b_bias_desc, \ - b_scales_desc, \ - a_scales_desc, \ - global_scales_desc, \ - b_zeros_desc, \ - g_idx_desc, \ - perm_desc) - - switch (handle->device) { -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAwqMarlinGemmWorkspaceSize(infiniopAwqMarlinGemmDescriptor_t desc, - size_t *size) { -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopAwqMarlinGemm( - infiniopAwqMarlinGemmDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *b_bias, - void *b_scales, - void *a_scales, - void *global_scales, - void *b_zeros, - void *g_idx, - void *perm, - int64_t b_q_type_id, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, c, a, b, b_bias, b_scales, a_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float, stream) - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyAwqMarlinGemmDescriptor(infiniopAwqMarlinGemmDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} - -// #endif diff --git a/src/infiniop/ops/awq_marlin_repack/awq_marlin_repack.h b/src/infiniop/ops/awq_marlin_repack/awq_marlin_repack.h deleted file mode 100644 index e2768173f..000000000 --- a/src/infiniop/ops/awq_marlin_repack/awq_marlin_repack.h +++ /dev/null @@ -1,48 +0,0 @@ -#ifndef AWQ_MARLIN_REPACK_H -#define AWQ_MARLIN_REPACK_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::awq_marlin_repack::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AwqMarlinRepackInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - AwqMarlinRepackInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t output_desc, \ - infiniopTensorDescriptor_t input_desc, \ - int64_t num_bits, \ - bool is_a_8bit); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *output, \ - const void *input, \ - void *stream) const; \ - }; \ - } - -#endif // AWQ_MARLIN_REPACK_H diff --git a/src/infiniop/ops/awq_marlin_repack/cuda/kernel.cuh b/src/infiniop/ops/awq_marlin_repack/cuda/kernel.cuh deleted file mode 100644 index b7288b603..000000000 --- a/src/infiniop/ops/awq_marlin_repack/cuda/kernel.cuh +++ /dev/null @@ -1,197 +0,0 @@ -#include "../marlin/marlin.cuh" - -namespace marlin { - -template -__device__ void awq_marlin_repack_kernel( - uint32_t const *__restrict__ b_q_weight_ptr, uint32_t *__restrict__ out_ptr, - int size_k, int size_n) { - constexpr int pack_factor = 32 / num_bits; - - constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1); - constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1); - int k_tiles = size_k / target_tile_k_size; - int n_tiles = size_n / target_tile_n_size; - int block_k_tiles = div_ceil(k_tiles, gridDim.x); - - auto start_k_tile = blockIdx.x * block_k_tiles; - if (start_k_tile >= k_tiles) { - return; - } - - int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); - - // Wait until the next thread tile has been loaded to shared memory. - auto wait_for_stage = [&]() { - // We only have `stages - 2` active fetches since we are double buffering - // and can only issue the next fetch when it is guaranteed that the previous - // shared memory load is fully complete (as it may otherwise be - // overwritten). - cp_async_wait(); - __syncthreads(); - }; - - extern __shared__ int4 sh[]; - - constexpr int tile_n_ints = target_tile_n_size / pack_factor; - - constexpr int stage_n_threads = tile_n_ints / 4; - constexpr int stage_k_threads = target_tile_k_size; - constexpr int stage_size = stage_k_threads * stage_n_threads; - - auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - cp_async_fence(); - return; - } - - int first_n = n_tile_id * target_tile_n_size; - int first_n_packed = first_n / pack_factor; - - int4 *sh_ptr = sh + stage_size * pipe; - - if (threadIdx.x < stage_size) { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; - - int first_k = k_tile_id * target_tile_k_size; - - cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast( - &(b_q_weight_ptr[(first_k + k_id) * (size_n / pack_factor) + first_n_packed + (n_id * 4)]))); - } - - cp_async_fence(); - }; - - auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - return; - } - - auto warp_id = threadIdx.x / 32; - auto th_id = threadIdx.x % 32; - - if (warp_id >= 4) { - return; - } - - int tc_col = th_id / 4; - int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2); - - constexpr int tc_offsets[4] = {0, 1, 8, 9}; - - int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col; - int cur_n_packed = cur_n / pack_factor; - int cur_n_pos = cur_n % pack_factor; - - constexpr int sh_stride = tile_n_ints; - constexpr uint32_t mask = (1 << num_bits) - 1; - - int4 *sh_stage_ptr = sh + stage_size * pipe; - uint32_t *sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); - - // Undo interleaving - int cur_n_pos_unpacked; - if constexpr (num_bits == 4) { - constexpr int undo_pack[8] = {0, 4, 1, 5, 2, 6, 3, 7}; - cur_n_pos_unpacked = undo_pack[cur_n_pos]; - } else { - constexpr int undo_pack[4] = {0, 2, 1, 3}; - cur_n_pos_unpacked = undo_pack[cur_n_pos]; - } - - uint32_t vals[8]; -#pragma unroll - for (int i = 0; i < 4; i++) { - if constexpr (is_a_8bit) { - int cur_elem = tc_row + i; - - int packed_src_0 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) * (warp_id % 2) + sh_stride * cur_elem]; - int packed_src_1 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) * (warp_id % 2) + sh_stride * (cur_elem + 16)]; - - vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask; - vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask; - } else { - int cur_elem = tc_row + tc_offsets[i]; - - int packed_src_0 = sh_stage_int_ptr[cur_n_packed + sh_stride * cur_elem]; - int packed_src_1 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) + sh_stride * cur_elem]; - - vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask; - vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask; - } - } - - constexpr int tile_size = target_tile_k_size * target_tile_n_size / pack_factor; - int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size; - - // Result of: - // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h - if constexpr (!is_a_8bit && num_bits == 4) { - int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; - - uint32_t res = 0; -#pragma unroll - for (int i = 0; i < 8; i++) { - res |= vals[pack_idx[i]] << (i * 4); - } - - out_ptr[out_offset + th_id * 4 + warp_id] = res; - - } else if constexpr (is_a_8bit && num_bits == 4) { - int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7}; - - uint32_t res = 0; -#pragma unroll - for (int i = 0; i < 8; i++) { - res |= vals[pack_idx[i]] << (i * 4); - } - - out_ptr[out_offset + th_id * 4 + warp_id] = res; - - } else { - constexpr int pack_idx[4] = {0, 2, 1, 3}; - - uint32_t res1 = 0; - uint32_t res2 = 0; -#pragma unroll - for (int i = 0; i < 4; i++) { - const int ii = is_a_8bit ? i : pack_idx[i]; - res1 |= vals[ii] << (i * 8); - res2 |= vals[4 + ii] << (i * 8); - } - - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; - } - }; - - auto start_pipes = [&](int k_tile_id, int n_tile_id) { -#pragma unroll - for (int pipe = 0; pipe < repack_stages - 1; pipe++) { - fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); - } - - wait_for_stage(); - }; -#pragma unroll - for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) { - int n_tile_id = 0; - - start_pipes(k_tile_id, n_tile_id); - - while (n_tile_id < n_tiles) { -#pragma unroll - for (int pipe = 0; pipe < repack_stages; pipe++) { - fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, - n_tile_id + pipe + repack_stages - 1); - repack_tile(pipe, k_tile_id, n_tile_id + pipe); - wait_for_stage(); - } - n_tile_id += repack_stages; - } - } -} - -} // namespace marlin diff --git a/src/infiniop/ops/awq_marlin_repack/info.h b/src/infiniop/ops/awq_marlin_repack/info.h deleted file mode 100644 index c9dea93fe..000000000 --- a/src/infiniop/ops/awq_marlin_repack/info.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef __AWQ_MARLIN_REPACK_INFO_H__ -#define __AWQ_MARLIN_REPACK_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include "marlin/marlin.cuh" -#include - -#include - -namespace op::awq_marlin_repack { - -class AwqMarlinRepackInfo { - AwqMarlinRepackInfo() = default; - -public: - infiniDtype_t output_dtype, input_dtype; - size_t size_k, size_n; - int64_t num_bits; - bool is_a_8bit; - - static utils::Result create( - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - int64_t num_bits, - bool is_a_8bit) { - CHECK_OR_RETURN( - output_desc != nullptr && input_desc != nullptr, - INFINI_STATUS_NULL_POINTER); - const infiniDtype_t output_dtype = output_desc->dtype(); - const infiniDtype_t input_dtype = input_desc->dtype(); - CHECK_DTYPE(input_dtype, INFINI_DTYPE_I32); - CHECK_DTYPE(input_dtype, output_dtype); - - size_t size_k = input_desc->dim(0); - int const pack_factor = 32 / num_bits; - size_t size_n = input_desc->dim(1) * pack_factor; - - CHECK_OR_RETURN(size_k / marlin::tile_size == output_desc->dim(0) || size_n * marlin::tile_size / pack_factor == output_desc->dim(1), - INFINI_STATUS_BAD_TENSOR_SHAPE); - - return utils::Result( - AwqMarlinRepackInfo{output_dtype, input_dtype, size_k, size_n, num_bits, is_a_8bit}); - } -}; - -} // namespace op::awq_marlin_repack - -#endif // __AWQ_MARLIN_REPACK_INFO_H__ diff --git a/src/infiniop/ops/awq_marlin_repack/marlin/marlin.cuh b/src/infiniop/ops/awq_marlin_repack/marlin/marlin.cuh deleted file mode 100644 index f3d897d27..000000000 --- a/src/infiniop/ops/awq_marlin_repack/marlin/marlin.cuh +++ /dev/null @@ -1,178 +0,0 @@ -#pragma once - -#ifndef _marlin_cuh -#define _marlin_cuh - -#include -#include -#include -#include - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -template -__device__ __forceinline__ uint32_t __cvta_generic_to_shared(T *ptr) { - size_t smem_addr; - asm volatile( - "cvta.to.shared.u64 %0, %1;" - : "=l"(smem_addr) - : "l"(ptr)); - return static_cast(smem_addr); -} - -namespace MARLIN_NAMESPACE_NAME { - -// Marlin params - -// 8 warps are a good choice since every SM has 4 schedulers and having more -// than 1 warp per schedule allows some more latency hiding. At the same time, -// we want relatively few warps to have many registers per warp and small tiles. -static constexpr int default_threads = 256; - -static constexpr int pipe_stages = 4; // 4 pipeline stages fit into shared memory - -static constexpr int min_thread_n = 64; -static constexpr int min_thread_k = 64; -static constexpr int max_thread_n = 256; - -static constexpr int tile_size = 16; -static constexpr int max_par = 16; - -// Repack params -static constexpr int repack_stages = 8; - -static constexpr int repack_threads = 256; - -static constexpr int tile_k_size = tile_size; -static constexpr int tile_n_size = tile_k_size * 4; - -// Helpers -template -struct Vec { - T elems[n]; - __device__ T &operator[](int i) { return elems[i]; } -}; - -using I4 = Vec; - -constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; } - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - -__device__ inline void cp_async1_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async2_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4(void *smem_ptr, const void *glob_ptr) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; -} - -__device__ inline void cp_async_fence() {} - -template -__device__ inline void cp_async_wait() {} - -#else - -__device__ inline void cp_async1_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 4; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async2_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 8; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.cg.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4(void *smem_ptr, const void *glob_ptr) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " cp.async.cg.shared.global [%0], [%1], %2;\n" - "}\n" ::"r"(smem), - "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async_fence() { - asm volatile("cp.async.commit_group;\n" ::); -} - -template -__device__ inline void cp_async_wait() { - asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); -} - -#endif - -} // namespace MARLIN_NAMESPACE_NAME - -#endif diff --git a/src/infiniop/ops/awq_marlin_repack/nvidia/awq_marlin_repack_nvidia.cu b/src/infiniop/ops/awq_marlin_repack/nvidia/awq_marlin_repack_nvidia.cu deleted file mode 100644 index c4eb1d1df..000000000 --- a/src/infiniop/ops/awq_marlin_repack/nvidia/awq_marlin_repack_nvidia.cu +++ /dev/null @@ -1,122 +0,0 @@ -#if defined(ENABLE_NVIDIA_API) -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../cuda/kernel.cuh" -#include "awq_marlin_repack_nvidia.cuh" -#include - -template -INFINIOP_CUDA_KERNEL awqMarlinRepackKernel( - uint32_t const *__restrict__ b_q_weight_ptr, uint32_t *__restrict__ out_ptr, - int size_k, int size_n) { - marlin::awq_marlin_repack_kernel( - b_q_weight_ptr, out_ptr, - size_k, size_n); -} - -#define CALL_IF(NUM_BITS, IS_A_8BIT) \ - else if (num_bits == NUM_BITS && is_a_8bit == IS_A_8BIT) { \ - cudaFuncSetAttribute( \ - awqMarlinRepackKernel, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ - awqMarlinRepackKernel \ - <<>>( \ - b_q_weight_ptr, out_ptr, size_k, size_n); \ - } - -infiniStatus_t awqMarlinRepack(uint32_t *out_ptr, const uint32_t *b_q_weight_ptr, int64_t size_k, - int64_t size_n, int64_t num_bits, - bool is_a_8bit, cudaStream_t stream) { - // Verify compatibility with marlin tile of 16x64 - if (size_k % marlin::tile_k_size != 0) { - std::cout << "size_k = " << size_k << " is not divisible by tile_k_size = " << marlin::tile_k_size << std::endl; - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (size_n % marlin::tile_n_size != 0) { - std::cout << "size_n = " << size_n << " is not divisible by tile_n_size = " << marlin::tile_n_size << std::endl; - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (num_bits != 4 && num_bits != 8) { - std::cout << "num_bits must be 4 or 8. Got = " << num_bits << std::endl; - return INFINI_STATUS_BAD_PARAM; - } - - int const pack_factor = 32 / num_bits; - - // Get dev info - int device_id = 0; - - int blocks; - cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_id); - - int max_shared_mem = 0; - cudaDeviceGetAttribute(&max_shared_mem, - cudaDevAttrMaxSharedMemoryPerBlockOptin, device_id); - assert(max_shared_mem > 0 && "max_shared_mem must be greater than 0"); - - if (false) { - } - CALL_IF(4, false) - CALL_IF(8, false) - CALL_IF(4, true) - CALL_IF(8, true) - else { - assert(false && "Unsupported repack config: num_bits, is_a_8bit"); - } - - return INFINI_STATUS_SUCCESS; -} - -namespace op::awq_marlin_repack::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - int64_t num_bits, - bool is_a_8bit) { - - auto handle = reinterpret_cast(handle_); - auto result = AwqMarlinRepackInfo::create(output_desc, input_desc, num_bits, is_a_8bit); - - size_t workspace_size = 0; - - *desc_ptr = new Descriptor( - new Opaque{handle->internal()}, - result.take(), - workspace_size, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t -Descriptor::calculate( - void *workspace, size_t workspace_size, - void *output, - const void *input, - void *stream_) const { - - cudaStream_t stream = (cudaStream_t)stream_; - - int64_t size_k = static_cast(_info.size_k); - int64_t size_n = static_cast(_info.size_n); - int64_t num_bits = _info.num_bits; - bool is_a_8bit = _info.is_a_8bit; - - awqMarlinRepack((uint32_t *)output, (const uint32_t *)input, size_k, - size_n, num_bits, - is_a_8bit, stream); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::awq_marlin_repack::nvidia -#endif // ENABLE_NVIDIA_API diff --git a/src/infiniop/ops/awq_marlin_repack/nvidia/awq_marlin_repack_nvidia.cuh b/src/infiniop/ops/awq_marlin_repack/nvidia/awq_marlin_repack_nvidia.cuh deleted file mode 100644 index 3cbec6c66..000000000 --- a/src/infiniop/ops/awq_marlin_repack/nvidia/awq_marlin_repack_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __AWQ_MARLIN_REPACK_CUDA_CUH__ -#define __AWQ_MARLIN_REPACK_CUDA_CUH__ - -#include "../awq_marlin_repack.h" - -DESCRIPTOR(nvidia) - -#endif // __AWQ_MARLIN_REPACK_CUDA_CUH__ diff --git a/src/infiniop/ops/awq_marlin_repack/operator.cc b/src/infiniop/ops/awq_marlin_repack/operator.cc deleted file mode 100644 index ea02da87d..000000000 --- a/src/infiniop/ops/awq_marlin_repack/operator.cc +++ /dev/null @@ -1,101 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/awq_marlin_repack.h" - -#if defined ENABLE_NVIDIA_API -#include "nvidia/awq_marlin_repack_nvidia.cuh" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAwqMarlinRepackDescriptor( - infiniopHandle_t handle, - infiniopAwqMarlinRepackDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - int64_t num_bits, - bool is_a_8bit) { -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::awq_marlin_repack::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output_desc, \ - input_desc, \ - num_bits, \ - is_a_8bit) - - switch (handle->device) { -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAwqMarlinRepackWorkspaceSize(infiniopAwqMarlinRepackDescriptor_t desc, - size_t *size) { -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopAwqMarlinRepack( - infiniopAwqMarlinRepackDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, input, stream) - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyAwqMarlinRepackDescriptor(infiniopAwqMarlinRepackDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} - -// #endif diff --git a/src/infiniop/ops/axpy/axpy.h b/src/infiniop/ops/axpy/axpy.h deleted file mode 100644 index 617e9d01b..000000000 --- a/src/infiniop/ops/axpy/axpy.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef __AXPY_H__ -#define __AXPY_H__ - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::axpy::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - AxpyInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - AxpyInfo info, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(std::move(info)), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t alpha_desc, \ - infiniopTensorDescriptor_t x_desc, \ - infiniopTensorDescriptor_t y_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - const void *alpha, \ - const void *x, \ - void *y, \ - void *stream) const; \ - }; \ - } - -#endif // __AXPY_H__ diff --git a/src/infiniop/ops/axpy/bang/axpy_bang.h b/src/infiniop/ops/axpy/bang/axpy_bang.h deleted file mode 100644 index dbce7e5ca..000000000 --- a/src/infiniop/ops/axpy/bang/axpy_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __AXPY_BANG_H__ -#define __AXPY_BANG_H__ - -#include "../axpy.h" - -DESCRIPTOR(bang) - -#endif // __AXPY_BANG_H__ diff --git a/src/infiniop/ops/axpy/bang/axpy_bang.mlu b/src/infiniop/ops/axpy/bang/axpy_bang.mlu deleted file mode 100644 index a7cd75bf7..000000000 --- a/src/infiniop/ops/axpy/bang/axpy_bang.mlu +++ /dev/null @@ -1,104 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "axpy_bang.h" -#include "axpy_bang_kernel.mlu" - -namespace op::axpy::bang { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t alpha_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = AxpyInfo::createAxpyInfo(alpha_desc, x_desc, y_desc); - CHECK_RESULT(result); - - // Create descriptor - *desc_ptr = new Descriptor( - result.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateAxpy( - const AxpyInfo &info, - const Tdata *alpha, - const Tdata *x, - Tdata *y, - cnrtQueue_t queue) { - - const int n = utils::cast(info.n); - const int incx = utils::cast(info.incx); - const int incy = utils::cast(info.incy); - - cnrtDim3_t k_dim; - cnrtFunctionType_t k_type; - - k_dim.x = 4; - k_dim.y = 1; - k_dim.z = 1; - k_type = cnrtFuncTypeUnion1; - - if (incx == 1 && incy == 1) { - axpyKernelContiguous<<>>( - n, - alpha, - x, - y); - } else { - axpyKernelStrided<<>>( - n, - alpha, - x, - incx, - y, - incy); - } - - cnrtQueueSync(queue); - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_AXPY(TDATA) \ - calculateAxpy(_info, \ - (const TDATA *)alpha, \ - (const TDATA *)x, \ - (TDATA *)y, \ - (cnrtQueue_t)stream) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *alpha, - const void *x, - void *y, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - switch (_info.data_type) { - case INFINI_DTYPE_F16: - return CALCULATE_AXPY(half); - case INFINI_DTYPE_F32: - return CALCULATE_AXPY(float); - case INFINI_DTYPE_BF16: - return CALCULATE_AXPY(bfloat16_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE_AXPY - -} // namespace op::axpy::bang diff --git a/src/infiniop/ops/axpy/bang/axpy_bang_kernel.mlu b/src/infiniop/ops/axpy/bang/axpy_bang_kernel.mlu deleted file mode 100644 index df0409987..000000000 --- a/src/infiniop/ops/axpy/bang/axpy_bang_kernel.mlu +++ /dev/null @@ -1,85 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "axpy_bang.h" - -__nram__ char nram_buffer[NRAM_MAX_SIZE]; - -template -__mlu_global__ void axpyKernelContiguous( - int n, - const Tdata *alpha, - const Tdata *x, - Tdata *y) { - - char *nram_aligned = (char *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1)); - - size_t nram_usable = NRAM_MAX_SIZE - (nram_aligned - nram_buffer); - size_t max_chunk_elements = nram_usable / (2 * sizeof(Tdata)); - - size_t align_elements = ALIGN_SIZE / sizeof(Tdata); - if (align_elements == 0) { - align_elements = 1; - } - int chunk_size = (int)((max_chunk_elements / align_elements) * align_elements); - - Tdata *nram_x = (Tdata *)nram_aligned; - Tdata *nram_y = nram_x + chunk_size; - - int elements_per_core = n / taskDim; - int remain = n % taskDim; - int core_elements = elements_per_core + (taskId < remain ? 1 : 0); - int core_offset = taskId < remain ? taskId * core_elements : taskId * elements_per_core + remain; - - if (core_elements <= 0) { - return; - } - - int chunks = core_elements / chunk_size; - int chunk_rem = core_elements % chunk_size; - - for (int c = 0; c < chunks; c++) { - int current_offset = core_offset + c * chunk_size; - - __memcpy(nram_x, x + current_offset, chunk_size * sizeof(Tdata), GDRAM2NRAM); - __memcpy(nram_y, y + current_offset, chunk_size * sizeof(Tdata), GDRAM2NRAM); - - __bang_mul_scalar(nram_x, nram_x, alpha[0], chunk_size); - __bang_add(nram_y, nram_y, nram_x, chunk_size); - - __memcpy(y + current_offset, nram_y, chunk_size * sizeof(Tdata), NRAM2GDRAM); - } - - if (chunk_rem > 0) { - int current_offset = core_offset + chunks * chunk_size; - int align_rem = ((chunk_rem + align_elements - 1) / align_elements) * align_elements; - - __memcpy(nram_x, x + current_offset, chunk_rem * sizeof(Tdata), GDRAM2NRAM); - __memcpy(nram_y, y + current_offset, chunk_rem * sizeof(Tdata), GDRAM2NRAM); - - __bang_mul_scalar(nram_x, nram_x, alpha[0], align_rem); - __bang_add(nram_y, nram_y, nram_x, align_rem); - - __memcpy(y + current_offset, nram_y, chunk_rem * sizeof(Tdata), NRAM2GDRAM); - } -} - -template -__mlu_global__ void axpyKernelStrided( - int n, - const Tdata *alpha, - const Tdata *x, - int incx, - Tdata *y, - int incy) { - - int elements_per_core = n / taskDim; - int remain = n % taskDim; - int actual_tasks = elements_per_core + (taskId < remain ? 1 : 0); - int start_idx = taskId < remain ? taskId * actual_tasks : taskId * elements_per_core + remain; - - for (int i = start_idx; i < start_idx + actual_tasks; ++i) { - int idx_x = i * incx; - int idx_y = i * incy; - - y[idx_y] += alpha[0] * x[idx_x]; - } -} diff --git a/src/infiniop/ops/axpy/cpu/axpy_cpu.cc b/src/infiniop/ops/axpy/cpu/axpy_cpu.cc deleted file mode 100644 index d2f97f8df..000000000 --- a/src/infiniop/ops/axpy/cpu/axpy_cpu.cc +++ /dev/null @@ -1,95 +0,0 @@ -#include "axpy_cpu.h" -#include "../../../devices/cpu/common_cpu.h" - -namespace op::axpy::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t alpha_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = AxpyInfo::createAxpyInfo(alpha_desc, x_desc, y_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateAxpy( - const AxpyInfo &info, - const Tdata *alpha, - const Tdata *x, - Tdata *y) { - - const size_t n = info.n; - const ptrdiff_t incx = info.incx; - const ptrdiff_t incy = info.incy; - - if constexpr (std::is_same::value || std::is_same::value) { - const float alpha_f = utils::cast(alpha[0]); - for (size_t i = 0; i < n; ++i) { - const ptrdiff_t x_idx = utils::cast(i) * incx; - const ptrdiff_t y_idx = utils::cast(i) * incy; - const float x_f = utils::cast(x[x_idx]); - const float y_f = utils::cast(y[y_idx]); - y[y_idx] = utils::cast(alpha_f * x_f + y_f); - } - } else { - const Tdata alpha_v = alpha[0]; - for (size_t i = 0; i < n; ++i) { - const ptrdiff_t x_idx = utils::cast(i) * incx; - const ptrdiff_t y_idx = utils::cast(i) * incy; - y[y_idx] = alpha_v * x[x_idx] + y[y_idx]; - } - } - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_AXPY(TDATA) \ - calculateAxpy(_info, \ - (const TDATA *)alpha, \ - (const TDATA *)x, \ - (TDATA *)y) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *alpha, - const void *x, - void *y, - void *stream) const { - - (void)workspace; - (void)workspace_size; - (void)stream; - - switch (_info.data_type) { - case INFINI_DTYPE_F16: - return CALCULATE_AXPY(fp16_t); - case INFINI_DTYPE_F32: - return CALCULATE_AXPY(float); - case INFINI_DTYPE_F64: - return CALCULATE_AXPY(double); - case INFINI_DTYPE_BF16: - return CALCULATE_AXPY(bf16_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE_AXPY - -} // namespace op::axpy::cpu diff --git a/src/infiniop/ops/axpy/cpu/axpy_cpu.h b/src/infiniop/ops/axpy/cpu/axpy_cpu.h deleted file mode 100644 index f25b49ef3..000000000 --- a/src/infiniop/ops/axpy/cpu/axpy_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __AXPY_CPU_H__ -#define __AXPY_CPU_H__ - -#include "../axpy.h" - -DESCRIPTOR(cpu) - -#endif // __AXPY_CPU_H__ diff --git a/src/infiniop/ops/axpy/info.h b/src/infiniop/ops/axpy/info.h deleted file mode 100644 index 9dfaa6c07..000000000 --- a/src/infiniop/ops/axpy/info.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef __AXPY_INFO_H__ -#define __AXPY_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" - -class AxpyInfo { -private: - AxpyInfo() = default; - -public: - size_t n; - ptrdiff_t incx; - ptrdiff_t incy; - infiniDtype_t data_type; - - static utils::Result createAxpyInfo( - infiniopTensorDescriptor_t alpha_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc) { - - CHECK_OR_RETURN(alpha_desc != nullptr, INFINI_STATUS_NULL_POINTER); - CHECK_OR_RETURN(x_desc != nullptr, INFINI_STATUS_NULL_POINTER); - CHECK_OR_RETURN(y_desc != nullptr, INFINI_STATUS_NULL_POINTER); - - auto data_type = x_desc->dtype(); - - CHECK_OR_RETURN(alpha_desc->dtype() == data_type, INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(x_desc->dtype() == data_type, INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_DTYPE(data_type, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_OR_RETURN(alpha_desc->numel() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(x_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(y_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(x_desc->numel() == y_desc->numel(), INFINI_STATUS_BAD_TENSOR_SHAPE); - - auto n = x_desc->numel(); - auto incx = x_desc->stride(0); - auto incy = y_desc->stride(0); - - return utils::Result(AxpyInfo{ - n, - incx, - incy, - data_type}); - } -}; - -#endif // __AXPY_INFO_H__ diff --git a/src/infiniop/ops/axpy/metax/axpy_metax.cc b/src/infiniop/ops/axpy/metax/axpy_metax.cc deleted file mode 100644 index 55f173cb0..000000000 --- a/src/infiniop/ops/axpy/metax/axpy_metax.cc +++ /dev/null @@ -1,102 +0,0 @@ -#include "axpy_metax.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" - -namespace op::axpy::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t alpha_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = AxpyInfo::createAxpyInfo(alpha_desc, x_desc, y_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - new Opaque{handle->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *alpha, - const void *x, - void *y, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - const int n = utils::cast(_info.n); - const int incx = utils::cast(_info.incx); - const int incy = utils::cast(_info.incy); - const infiniDtype_t data_type = _info.data_type; - - hpccDataType alpha_type, x_type, y_type; - hpccDataType execution_type; - - switch (data_type) { - case INFINI_DTYPE_F16: - alpha_type = x_type = y_type = HPCC_R_16F; - execution_type = HPCC_R_32F; - break; - case INFINI_DTYPE_BF16: - alpha_type = x_type = y_type = HPCC_R_16BF; - execution_type = HPCC_R_32F; - break; - case INFINI_DTYPE_F32: - alpha_type = x_type = y_type = HPCC_R_32F; - execution_type = HPCC_R_32F; - break; - case INFINI_DTYPE_F64: - alpha_type = x_type = y_type = HPCC_R_64F; - execution_type = HPCC_R_64F; - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - CHECK_STATUS(_opaque->internal->useMcblas( - (hcStream_t)stream, - [&](hcblasHandle_t handle) { - CHECK_MCBLAS(hcblasSetPointerMode( - handle, - HCBLAS_POINTER_MODE_DEVICE)); - - CHECK_MCBLAS(hcblasAxpyEx( - handle, - n, - alpha, - alpha_type, - x, - x_type, - incx, - y, - y_type, - incy, - execution_type)); - - return INFINI_STATUS_SUCCESS; - })); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::axpy::metax diff --git a/src/infiniop/ops/axpy/metax/axpy_metax.h b/src/infiniop/ops/axpy/metax/axpy_metax.h deleted file mode 100644 index 8129ca2e4..000000000 --- a/src/infiniop/ops/axpy/metax/axpy_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __AXPY_METAX_H__ -#define __AXPY_METAX_H__ - -#include "../axpy.h" - -DESCRIPTOR(metax) - -#endif // __AXPY_METAX_H__ diff --git a/src/infiniop/ops/axpy/operator.cc b/src/infiniop/ops/axpy/operator.cc deleted file mode 100644 index cd57223e4..000000000 --- a/src/infiniop/ops/axpy/operator.cc +++ /dev/null @@ -1,127 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/axpy.h" - -#ifdef ENABLE_CPU_API -#include "cpu/axpy_cpu.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/axpy_metax.h" -#endif -#ifdef ENABLE_CAMBRICON_API -#include "bang/axpy_bang.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateAxpyDescriptor( - infiniopHandle_t handle, - infiniopAxpyDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t alpha_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::axpy::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - alpha_desc, \ - x_desc, \ - y_desc) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetAxpyWorkspaceSize(infiniopAxpyDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - GET(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopAxpy( - infiniopAxpyDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *alpha, - const void *x, - void *y, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, alpha, x, y, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyAxpyDescriptor(infiniopAxpyDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - DELETE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits.h b/src/infiniop/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits.h deleted file mode 100644 index a4da5cb8f..000000000 --- a/src/infiniop/ops/binary_cross_entropy_with_logits/binary_cross_entropy_with_logits.h +++ /dev/null @@ -1,67 +0,0 @@ -#ifndef __BINARY_CROSS_ENTROPY_WITH_LOGITS_H__ -#define __BINARY_CROSS_ENTROPY_WITH_LOGITS_H__ - -#include "../../operator.h" -#include "info.h" - -/** - * # 关于 `BCEWithLogits` 算子描述符的说明 - * * 采用 PImpl 设计模式，将不同硬件后端（如 CUDA 原生算子、CPU 循环、或是芯片厂商的专用库调用） - * 封装在 `Opaque` 结构中。 - * * 描述符在创建时会完成形状校验、步长分析，并确定最优的计算 Workspace 大小。 - */ - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::bce_with_logits::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - infiniDtype_t _dtype; \ - BCEWithLogitsInfo _info; /* 包含各输入输出张量的维度与步长 */ \ - size_t _workspace_size; \ - infiniopReduction_t _reduction; \ - \ - Descriptor( \ - infiniDtype_t dtype, \ - BCEWithLogitsInfo info, \ - infiniopReduction_t reduction, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _dtype(dtype), \ - _info(info), \ - _workspace_size(workspace_size_), \ - _reduction(reduction) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t logits_desc, \ - infiniopTensorDescriptor_t target_desc, \ - infiniopTensorDescriptor_t weight_desc, \ - infiniopTensorDescriptor_t pos_weight_desc, \ - infiniopReduction_t reduction); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *out, \ - const void *logits, \ - const void *target, \ - const void *weight, /* 可选，可为 nullptr */ \ - const void *pos_weight, /* 可选，可为 nullptr */ \ - void *stream) const; \ - }; \ - } - -#endif // __BINARY_CROSS_ENTROPY_WITH_LOGITS_H__ diff --git a/src/infiniop/ops/binary_cross_entropy_with_logits/cpu/binary_cross_entropy_with_logits_cpu.cc b/src/infiniop/ops/binary_cross_entropy_with_logits/cpu/binary_cross_entropy_with_logits_cpu.cc deleted file mode 100644 index 471e99d33..000000000 --- a/src/infiniop/ops/binary_cross_entropy_with_logits/cpu/binary_cross_entropy_with_logits_cpu.cc +++ /dev/null @@ -1,164 +0,0 @@ -#include "binary_cross_entropy_with_logits_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include -#include - -namespace op::bce_with_logits::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t logits_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t pos_weight_desc, - infiniopReduction_t reduction) { - - auto handle = reinterpret_cast(handle_); - auto dtype = logits_desc->dtype(); - - // 1. 类型检查 - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - // 2. 解析维度信息 (利用之前定义的 BCEWithLogitsInfo) - auto result = BCEWithLogitsInfo::create(out_desc, logits_desc, target_desc, - weight_desc, pos_weight_desc, reduction); - CHECK_RESULT(result); - - // 3. 实例化描述符 - *desc_ptr = new Descriptor( - dtype, result.take(), reduction, 0, - nullptr, - handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -/** - * 核心数值稳定逻辑：L = -w * [pw * y * log(sigmoid(x)) + (1-y) * log(1-sigmoid(x))] - * 变形为：L = w * [max(x, 0) - x * y * pw + (1 + (pw-1) * y) * log(1 + exp(-|x|))] - * 当 pw=1 时简化为：L = w * [max(x, 0) - x * y + log(1 + exp(-|x|))] - */ -template -void calculate_bce( - const BCEWithLogitsInfo &info, - void *out, - const void *logits, - const void *target, - const void *weight, - const void *pos_weight) { - - size_t n = info.num_elements; - float total_loss = 0.0f; - - // 获取各张量指针 - const Tdata *l_ptr = reinterpret_cast(logits); - const Tdata *t_ptr = reinterpret_cast(target); - const Tdata *w_ptr = reinterpret_cast(weight); - const Tdata *pw_ptr = reinterpret_cast(pos_weight); - Tdata *o_ptr = reinterpret_cast(out); - - auto &logits_info = info.logits; - auto &target_info = info.target; - auto &weight_info = info.weight; - auto &out_info = info.out; - -#pragma omp parallel for reduction(+ : total_loss) - for (ptrdiff_t i = 0; i < (ptrdiff_t)n; ++i) { - size_t idx = static_cast(i); - - // 使用 stride 计算实际内存偏移，支持任意内存布局 - size_t logits_offset = op::common_cpu::indexToOffset( - idx, - logits_info.ndim, - logits_info.dims.data(), - logits_info.stride.data()); - size_t target_offset = op::common_cpu::indexToOffset( - idx, - target_info.ndim, - target_info.dims.data(), - target_info.stride.data()); - - float x = utils::cast(l_ptr[logits_offset]); - float y = utils::cast(t_ptr[target_offset]); - - // 处理 pos_weight 广播 (假设 logits 形状 [..., C], pos_weight 为 [C] 且连续) - float pw = 1.0f; - if (pw_ptr && info.pos_weight.total_elements > 0) { - size_t c = idx % info.pos_weight.total_elements; - pw = utils::cast(pw_ptr[c]); - } - - // 处理 weight： - // - 如果与 logits 完全同形状，则按 stride 精确索引； - // - 如果为向量 [C]，则通过 indexToOffset 实现按最后一维广播。 - float w = 1.0f; - if (w_ptr && weight_info.ndim > 0) { - size_t weight_offset = op::common_cpu::indexToOffset( - idx, - weight_info.ndim, - weight_info.dims.data(), - weight_info.stride.data()); - w = utils::cast(w_ptr[weight_offset]); - } - - // 数值稳定的 BCEWithLogits 计算（对齐 PyTorch 实现）： - // max_val = max(-x, 0) - // log_weight = 1 + (pos_weight - 1) * y - // loss = (1 - y) * x + log_weight * (log(1 + exp(-|x|)) + max_val) - float max_val = std::max(-x, 0.0f); - float log_weight = 1.0f + (pw - 1.0f) * y; - float loss = (1.0f - y) * x + log_weight * (std::log1p(std::exp(-std::abs(x))) + max_val); - - loss *= w; - - if (info.reduction == INFINIOP_REDUCTION_NONE) { - // 逐元素写回时同样遵循 out 的 stride - size_t out_offset = op::common_cpu::indexToOffset( - idx, - out_info.ndim, - out_info.dims.data(), - out_info.stride.data()); - o_ptr[out_offset] = utils::cast(loss); - } else { - total_loss += loss; - } - } - - // 处理归约输出 - if (info.reduction == INFINIOP_REDUCTION_MEAN) { - o_ptr[0] = utils::cast(total_loss / n); - } else if (info.reduction == INFINIOP_REDUCTION_SUM) { - o_ptr[0] = utils::cast(total_loss); - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *logits, - const void *target, - const void *weight, - const void *pos_weight, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - cpu::calculate_bce(_info, out, logits, target, weight, pos_weight); - return INFINI_STATUS_SUCCESS; - case INFINI_DTYPE_BF16: - cpu::calculate_bce(_info, out, logits, target, weight, pos_weight); - return INFINI_STATUS_SUCCESS; - case INFINI_DTYPE_F32: - cpu::calculate_bce(_info, out, logits, target, weight, pos_weight); - return INFINI_STATUS_SUCCESS; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::bce_with_logits::cpu diff --git a/src/infiniop/ops/binary_cross_entropy_with_logits/cpu/binary_cross_entropy_with_logits_cpu.h b/src/infiniop/ops/binary_cross_entropy_with_logits/cpu/binary_cross_entropy_with_logits_cpu.h deleted file mode 100644 index f5a27c314..000000000 --- a/src/infiniop/ops/binary_cross_entropy_with_logits/cpu/binary_cross_entropy_with_logits_cpu.h +++ /dev/null @@ -1,16 +0,0 @@ -#ifndef __BINARY_CROSS_ENTROPY_WITH_LOGITS_CPU_H__ -#define __BINARY_CROSS_ENTROPY_WITH_LOGITS_CPU_H__ - -#include "../binary_cross_entropy_with_logits.h" - -/** - * 使用 bce_with_logits.h 中定义的 DESCRIPTOR 宏 - * * 这将自动在命名空间 op::bce_with_logits::cpu 中生成 Descriptor 类。 - * 该类将继承自 InfiniopDescriptor，并包含： - * - BCEWithLogitsInfo _info (存储校验后的维度和步长) - * - create() 静态方法 (负责 CPU 版描述符的实例化) - * - calculate() 方法 (负责 CPU 版数值稳定逻辑的执行) - */ -DESCRIPTOR(cpu) - -#endif // __BINARY_CROSS_ENTROPY_WITH_LOGITS_CPU_H__ diff --git a/src/infiniop/ops/binary_cross_entropy_with_logits/info.h b/src/infiniop/ops/binary_cross_entropy_with_logits/info.h deleted file mode 100644 index a2722426a..000000000 --- a/src/infiniop/ops/binary_cross_entropy_with_logits/info.h +++ /dev/null @@ -1,136 +0,0 @@ -#ifndef __BINARY_CROSS_ENTROPY_WITH_LOGITS_INFO_H__ -#define __BINARY_CROSS_ENTROPY_WITH_LOGITS_INFO_H__ - -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" -#include "infiniop/ops/binary_cross_entropy_with_logits.h" -#include -#include - -namespace op::bce_with_logits { - -/** - * 描述 BCE 算子中各张量的内存布局 - * 动态申请 dims 和 stride，支持任意维度的张量 - */ -struct BCETensorInfo { - size_t total_elements = 0; - size_t ndim = 0; - std::vector dims; // 动态存储维度 - std::vector stride; // 动态存储步长 - - BCETensorInfo() = default; - - static utils::Result create(infiniopTensorDescriptor_t desc) { - if (desc == nullptr) { - return INFINI_STATUS_SUCCESS; - } - - BCETensorInfo info; - info.ndim = desc->ndim(); - info.total_elements = 1; - - // 动态调整 vector 大小 - info.dims.reserve(info.ndim); - info.stride.reserve(info.ndim); - - for (size_t i = 0; i < info.ndim; ++i) { - size_t d = desc->dim(i); - info.dims.push_back(d); - info.stride.push_back(desc->stride(i)); - info.total_elements *= d; - } - return utils::Result(std::move(info)); - } - - // 辅助方法：获取最后一维大小（用于 pos_weight 校验） - size_t last_dim() const { - return ndim > 0 ? dims.back() : 0; - } -}; - -class BCEWithLogitsInfo { -public: - BCETensorInfo logits; - BCETensorInfo target; - BCETensorInfo weight; - BCETensorInfo pos_weight; - BCETensorInfo out; - - size_t num_elements; - infiniopReduction_t reduction; - - // 由于 BCETensorInfo 内部使用了 vector，BCEWithLogitsInfo 现在是可移动且安全的 - BCEWithLogitsInfo() = default; - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t logits_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t pos_weight_desc, - infiniopReduction_t reduction) { - - auto logits_res = BCETensorInfo::create(logits_desc); - CHECK_RESULT(logits_res); - auto target_res = BCETensorInfo::create(target_desc); - CHECK_RESULT(target_res); - auto out_res = BCETensorInfo::create(out_desc); - CHECK_RESULT(out_res); - - BCEWithLogitsInfo info; - info.logits = logits_res.take(); - info.target = target_res.take(); - info.out = out_res.take(); - info.reduction = reduction; - info.num_elements = info.logits.total_elements; - - // 1. 基本形状一致性校验 - if (info.logits.ndim != info.target.ndim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - for (size_t i = 0; i < info.logits.ndim; ++i) { - if (info.logits.dims[i] != info.target.dims[i]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - // 2. 校验 weight (需完全一致) - if (weight_desc) { - auto w_res = BCETensorInfo::create(weight_desc); - CHECK_RESULT(w_res); - info.weight = w_res.take(); - - // 允许两种情况： - // 1. 完全一致 - // 2. weight 是一个向量，且长度等于 logits 的最后一维 (常见广播场景) - bool is_full_match = (info.weight.total_elements == info.logits.total_elements); - bool is_last_dim_match = (info.weight.total_elements == info.logits.last_dim()); - - if (!is_full_match && !is_last_dim_match) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - // 3. 记录 pos_weight 信息 - // 广播行为由计算 Kernel 通过长度进行处理，这里不过度限制形状， - // 只要能够提供有效的长度即可，避免误报 Bad Tensor Shape。 - if (pos_weight_desc) { - auto pw_res = BCETensorInfo::create(pos_weight_desc); - CHECK_RESULT(pw_res); - info.pos_weight = pw_res.take(); - } - - // 4. 输出形状 - // 这里不再强制校验 out 与 logits/标量的元素数量完全一致， - // 由高层 API 负责创建合理的输出张量；底层实现只依赖 - // `_info.out` 的 stride 在 reduction==NONE 且逐元素写回时使用。 - - return utils::Result(std::move(info)); - } -}; - -} // namespace op::bce_with_logits - -#endif // __BINARY_CROSS_ENTROPY_WITH_LOGITS_INFO_H__ diff --git a/src/infiniop/ops/binary_cross_entropy_with_logits/metax/binary_cross_entropy_with_logits_metax.h b/src/infiniop/ops/binary_cross_entropy_with_logits/metax/binary_cross_entropy_with_logits_metax.h deleted file mode 100644 index 230bdcce5..000000000 --- a/src/infiniop/ops/binary_cross_entropy_with_logits/metax/binary_cross_entropy_with_logits_metax.h +++ /dev/null @@ -1,17 +0,0 @@ -#ifndef __BINARY_CROSS_ENTROPY_WITH_LOGITS_METAX_CUH__ -#define __BINARY_CROSS_ENTROPY_WITH_LOGITS_METAX_CUH__ - -#include "../binary_cross_entropy_with_logits.h" - -/** - * 使用 bce_with_logits.h 中定义的 DESCRIPTOR 宏。 - * 这将在命名空间 op::bce_with_logits::metax 中生成针对 METAX 设备的 Descriptor 类。 - * * * 在 METAX 端的实现（.cu 文件）中，Opaque 结构体通常包含： - * - cudnnHandle_t: 如果使用 cuDNN 的算子实现。 - * - cudnnTensorDescriptor_t: 用于描述各输入输出张量的 cuDNN 格式。 - * - KernelConfig: 用于自定义 METAX Kernel 的网格（Grid）和线程块（Block）配置。 - * - dataType: 对应的 METAX 数据类型（如 METAX_R_32F）。 - */ -DESCRIPTOR(metax) - -#endif // __BINARY_CROSS_ENTROPY_WITH_LOGITS_METAX_CUH__ diff --git a/src/infiniop/ops/binary_cross_entropy_with_logits/metax/binary_cross_entropy_with_logits_metax.maca b/src/infiniop/ops/binary_cross_entropy_with_logits/metax/binary_cross_entropy_with_logits_metax.maca deleted file mode 100644 index c14bb75bc..000000000 --- a/src/infiniop/ops/binary_cross_entropy_with_logits/metax/binary_cross_entropy_with_logits_metax.maca +++ /dev/null @@ -1,335 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" - -#include "binary_cross_entropy_with_logits_metax.h" - -#include - -namespace op::bce_with_logits::metax { - -using device::metax::indexToOffset; - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() = default; - -// 在 GPU 侧使用的简化张量信息（固定上限维度，支持 stride） -constexpr int BCE_MAX_DIMS = 8; - -struct BCETensorInfoDevice { - size_t ndim; - size_t shape[BCE_MAX_DIMS]; - ptrdiff_t strides[BCE_MAX_DIMS]; -}; - -static inline BCETensorInfoDevice make_device_info(const BCETensorInfo &info) { - BCETensorInfoDevice dev{}; - dev.ndim = info.ndim; - for (size_t i = 0; i < info.ndim && i < static_cast(BCE_MAX_DIMS); ++i) { - dev.shape[i] = info.dims[i]; - dev.strides[i] = info.stride[i]; - } - return dev; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t logits_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t pos_weight_desc, - infiniopReduction_t reduction) { - - auto handle = reinterpret_cast(handle_); - auto dtype = logits_desc->dtype(); - - // METAX 实现支持 F16 / F32 / BF16 - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = BCEWithLogitsInfo::create(out_desc, logits_desc, target_desc, - weight_desc, pos_weight_desc, reduction); - CHECK_RESULT(result); - - auto info = result.take(); - - // F16/BF16 在做归约时需要一个 float 标量 workspace 来累加 - size_t workspace_size = 0; - if (reduction != INFINIOP_REDUCTION_NONE && (dtype == INFINI_DTYPE_F16 || dtype == INFINI_DTYPE_BF16)) { - workspace_size = sizeof(float); - } - - *desc_ptr = new Descriptor( - dtype, std::move(info), reduction, workspace_size, - nullptr, - handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -// 将任意标量类型提升为 float -template -__device__ __forceinline__ float to_float(T x) { - if constexpr (std::is_same_v) { - return x; - } else if constexpr (std::is_same_v) { - return __half2float(x); - } else if constexpr (std::is_same_v) { - return __bfloat162float(x); - } else { - return static_cast(x); - } -} - -// 从 float 转回目标标量类型 -template -__device__ __forceinline__ T from_float(float x) { - if constexpr (std::is_same_v) { - return x; - } else if constexpr (std::is_same_v) { - return __float2half(x); - } else if constexpr (std::is_same_v) { - return __float2bfloat16(x); - } else { - return static_cast(x); - } -} - -// --- METAX Kernel: 支持 stride 的数值稳定 BCE 计算 --- -template -__global__ void bce_logits_kernel( - void *out_raw, - const Tdata *logits, - const Tdata *target, - const Tdata *weight, - const Tdata *pos_weight, - BCETensorInfoDevice logits_info, - BCETensorInfoDevice target_info, - BCETensorInfoDevice weight_info, - BCETensorInfoDevice out_info, - size_t n, - size_t pos_weight_len, - infiniopReduction_t reduction) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= n) { - return; - } - - // 计算逻辑索引在各张量中的偏移（支持任意 stride） - size_t logits_offset = indexToOffset(idx, logits_info.ndim, - logits_info.shape, logits_info.strides); - size_t target_offset = indexToOffset(idx, target_info.ndim, - target_info.shape, target_info.strides); - - float x = to_float(logits[logits_offset]); - float y = to_float(target[target_offset]); - - float pw = 1.0f; - if (pos_weight && pos_weight_len > 0) { - // 按最后一维广播：假设 pos_weight 是连续的一维张量 - size_t c = idx % pos_weight_len; - pw = to_float(pos_weight[c]); - } - - float w = 1.0f; - if (weight && weight_info.ndim > 0) { - size_t weight_offset = indexToOffset(idx, weight_info.ndim, - weight_info.shape, weight_info.strides); - w = to_float(weight[weight_offset]); - } - - // 数值稳定公式：max(x, 0) - x * y * pw + (1 + (pw - 1) * y) * log(1 + exp(-abs(x))) - float loss = (fmaxf(x, 0.0f) - x * y * pw + (1.0f + (pw - 1.0f) * y) * logf(1.0f + expf(-fabsf(x)))); - - loss *= w; - - if (reduction == INFINIOP_REDUCTION_NONE) { - // 写回逐元素 loss（支持 stride 的 out） - size_t out_offset = indexToOffset(idx, out_info.ndim, - out_info.shape, out_info.strides); - auto *out_ptr = static_cast(out_raw); - out_ptr[out_offset] = from_float(loss); - } else { - // 对于 mean 或 sum，使用 float 累加到标量位置 - auto *out_accum = static_cast(out_raw); - atomicAdd(out_accum, static_cast(loss)); - } -} - -// F32 mean 归约：对输出标量做除法 -__global__ void bce_logits_mean_scale_kernel_f32(float *val, size_t count) { - if (threadIdx.x == 0 && blockIdx.x == 0) { - *val /= static_cast(count); - } -} - -// F16/BF16 归约：从 float workspace 写回目标 dtype -template -__global__ void bce_logits_reduce_finalize_kernel( - Tdata *out, - float *workspace, - size_t count, - int is_mean) { - if (threadIdx.x == 0 && blockIdx.x == 0) { - float v = *workspace; - if (is_mean) { - v /= static_cast(count); - } - *out = from_float(v); - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *logits, - const void *target, - const void *weight, - const void *pos_weight, - void *stream) const { - - hcStream_t custream = (hcStream_t)stream; - size_t n = _info.num_elements; - - // F16/BF16 + 归约需要 float workspace - if (_reduction != INFINIOP_REDUCTION_NONE && (_dtype == INFINI_DTYPE_F16 || _dtype == INFINI_DTYPE_BF16)) { - if (workspace_size < sizeof(float)) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - } - - int block = 256; - int grid = static_cast((n + block - 1) / block); - - // 构造 GPU 侧的张量信息（含 stride） - BCETensorInfoDevice logits_info = make_device_info(_info.logits); - BCETensorInfoDevice target_info = make_device_info(_info.target); - BCETensorInfoDevice out_info = make_device_info(_info.out); - BCETensorInfoDevice weight_info = {}; - if (_info.weight.total_elements != 0) { - weight_info = make_device_info(_info.weight); - } - - size_t pos_weight_len = _info.pos_weight.total_elements; - - switch (_dtype) { - case INFINI_DTYPE_F32: { - // 如果是规约操作，计算前需将输出位置清零 - if (_reduction != INFINIOP_REDUCTION_NONE) { - hcMemsetAsync(out, 0, sizeof(float), custream); - } - - bce_logits_kernel<<>>( - out, - static_cast(logits), - static_cast(target), - static_cast(weight), - static_cast(pos_weight), - logits_info, - target_info, - weight_info, - out_info, - n, - pos_weight_len, - _reduction); - - if (_reduction == INFINIOP_REDUCTION_MEAN) { - bce_logits_mean_scale_kernel_f32<<<1, 1, 0, custream>>>( - static_cast(out), n); - } - break; - } - case INFINI_DTYPE_F16: { - auto *logits_h = static_cast(logits); - auto *target_h = static_cast(target); - auto *weight_h = static_cast(weight); - auto *pos_weight_h = static_cast(pos_weight); - - void *out_raw = nullptr; - float *workspace_f = nullptr; - - if (_reduction == INFINIOP_REDUCTION_NONE) { - out_raw = out; - } else { - workspace_f = static_cast(workspace); - hcMemsetAsync(workspace_f, 0, sizeof(float), custream); - out_raw = workspace_f; - } - - bce_logits_kernel<<>>( - out_raw, - logits_h, - target_h, - weight_h, - pos_weight_h, - logits_info, - target_info, - weight_info, - out_info, - n, - pos_weight_len, - _reduction); - - if (_reduction != INFINIOP_REDUCTION_NONE) { - int is_mean = (_reduction == INFINIOP_REDUCTION_MEAN) ? 1 : 0; - bce_logits_reduce_finalize_kernel<<<1, 1, 0, custream>>>( - static_cast(out), workspace_f, n, is_mean); - } - - break; - } - case INFINI_DTYPE_BF16: { - auto *logits_b = static_cast(logits); - auto *target_b = static_cast(target); - auto *weight_b = static_cast(weight); - auto *pos_weight_b = static_cast(pos_weight); - - void *out_raw = nullptr; - float *workspace_f = nullptr; - - if (_reduction == INFINIOP_REDUCTION_NONE) { - out_raw = out; - } else { - workspace_f = static_cast(workspace); - hcMemsetAsync(workspace_f, 0, sizeof(float), custream); - out_raw = workspace_f; - } - - bce_logits_kernel<<>>( - out_raw, - logits_b, - target_b, - weight_b, - pos_weight_b, - logits_info, - target_info, - weight_info, - out_info, - n, - pos_weight_len, - _reduction); - - if (_reduction != INFINIOP_REDUCTION_NONE) { - int is_mean = (_reduction == INFINIOP_REDUCTION_MEAN) ? 1 : 0; - bce_logits_reduce_finalize_kernel<<<1, 1, 0, custream>>>( - static_cast(out), workspace_f, n, is_mean); - } - - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - hcError_t err = hcGetLastError(); - if (err != hcSuccess) { - return INFINI_STATUS_INTERNAL_ERROR; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::bce_with_logits::metax diff --git a/src/infiniop/ops/binary_cross_entropy_with_logits/moore/binary_cross_entropy_with_logits_moore.h b/src/infiniop/ops/binary_cross_entropy_with_logits/moore/binary_cross_entropy_with_logits_moore.h deleted file mode 100644 index f7843f118..000000000 --- a/src/infiniop/ops/binary_cross_entropy_with_logits/moore/binary_cross_entropy_with_logits_moore.h +++ /dev/null @@ -1,16 +0,0 @@ -#ifndef __BINARY_CROSS_ENTROPY_WITH_LOGITS_MOORE_H__ -#define __BINARY_CROSS_ENTROPY_WITH_LOGITS_MOORE_H__ - -#include "../binary_cross_entropy_with_logits.h" - -/** - * 使用 bce_with_logits.h 中定义的 DESCRIPTOR 宏。 - * 这将在命名空间 op::bce_with_logits::moore 中生成针对 Moore 设备的 Descriptor 类。 - * * 在 Moore 端的实现（.mu 文件）中，Opaque 结构体通常包含： - * - musaHandle_t: 如果使用 MUSA 库的算子实现。 - * - KernelConfig: 用于 MUSA Kernel 的网格（Grid）和线程块（Block）配置。 - * - dataType: 对应的 MUSA 数据类型（如 MUSA_R_32F）。 - */ -DESCRIPTOR(moore) - -#endif // __BINARY_CROSS_ENTROPY_WITH_LOGITS_MOORE_H__ diff --git a/src/infiniop/ops/binary_cross_entropy_with_logits/moore/binary_cross_entropy_with_logits_moore.mu b/src/infiniop/ops/binary_cross_entropy_with_logits/moore/binary_cross_entropy_with_logits_moore.mu deleted file mode 100644 index cf1b4fe7b..000000000 --- a/src/infiniop/ops/binary_cross_entropy_with_logits/moore/binary_cross_entropy_with_logits_moore.mu +++ /dev/null @@ -1,320 +0,0 @@ -#include "../../../devices/moore/moore_handle.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "binary_cross_entropy_with_logits_moore.h" -#include - -namespace op::bce_with_logits::moore { - -using device::moore::indexToOffset; - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() = default; - -// 摩尔线程平台通常与 CUDA 保持一致的维度上限 -constexpr int BCE_MAX_DIMS = 8; - -struct BCETensorInfoDevice { - size_t ndim; - size_t shape[BCE_MAX_DIMS]; - ptrdiff_t strides[BCE_MAX_DIMS]; -}; - -static inline BCETensorInfoDevice make_device_info(const BCETensorInfo &info) { - BCETensorInfoDevice dev{}; - dev.ndim = info.ndim; - for (size_t i = 0; i < info.ndim && i < static_cast(BCE_MAX_DIMS); ++i) { - dev.shape[i] = info.dims[i]; - dev.strides[i] = info.stride[i]; - } - return dev; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t logits_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t pos_weight_desc, - infiniopReduction_t reduction) { - - auto handle = reinterpret_cast(handle_); - auto dtype = logits_desc->dtype(); - - // Moore 实现支持 F16 / F32 / BF16 - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = BCEWithLogitsInfo::create(out_desc, logits_desc, target_desc, - weight_desc, pos_weight_desc, reduction); - CHECK_RESULT(result); - - auto info = result.take(); - - size_t workspace_size = 0; - if (reduction != INFINIOP_REDUCTION_NONE && (dtype == INFINI_DTYPE_F16 || dtype == INFINI_DTYPE_BF16)) { - workspace_size = sizeof(float); - } - - *desc_ptr = new Descriptor( - dtype, std::move(info), reduction, workspace_size, - nullptr, - handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -// 针对 Moore 平台的类型提升逻辑 -template -__device__ __forceinline__ float to_float(T x) { - if constexpr (std::is_same_v) { - return x; - } else if constexpr (std::is_same_v) { - return __half2float(x); - } else if constexpr (std::is_same_v) { // MUSA 兼容 cuda_bfloat16 名称或使用内部 bf16 - return __bfloat162float(x); - } else { - return static_cast(x); - } -} - -template -__device__ __forceinline__ T from_float(float x) { - if constexpr (std::is_same_v) { - return x; - } else if constexpr (std::is_same_v) { - return __float2half(x); - } else if constexpr (std::is_same_v) { - return __float2bfloat16_rn(x); // Moore 平台推荐显式使用 _rn - } else { - return static_cast(x); - } -} - -// --- MUSA Kernel --- -template -__global__ void bce_logits_kernel( - void *out_raw, - const Tdata *logits, - const Tdata *target, - const Tdata *weight, - const Tdata *pos_weight, - BCETensorInfoDevice logits_info, - BCETensorInfoDevice target_info, - BCETensorInfoDevice weight_info, - BCETensorInfoDevice out_info, - size_t n, - size_t pos_weight_len, - infiniopReduction_t reduction) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= n) { - return; - } - - size_t logits_offset = indexToOffset(idx, logits_info.ndim, - logits_info.shape, logits_info.strides); - size_t target_offset = indexToOffset(idx, target_info.ndim, - target_info.shape, target_info.strides); - - float x = to_float(logits[logits_offset]); - float y = to_float(target[target_offset]); - - float pw = 1.0f; - if (pos_weight && pos_weight_len > 0) { - size_t c = idx % pos_weight_len; - pw = to_float(pos_weight[c]); - } - - float w = 1.0f; - if (weight && weight_info.ndim > 0) { - size_t weight_offset = indexToOffset(idx, weight_info.ndim, - weight_info.shape, weight_info.strides); - w = to_float(weight[weight_offset]); - } - - // 数值稳定的 BCEWithLogits 计算（对齐 PyTorch 实现）： - // max_val = max(-x, 0) - // log_weight = 1 + (pos_weight - 1) * y - // loss = (1 - y) * x + log_weight * (log(1 + exp(-|x|)) + max_val) - float max_val = fmaxf(-x, 0.0f); - float log_weight = 1.0f + (pw - 1.0f) * y; - float loss = (1.0f - y) * x + log_weight * (logf(1.0f + expf(-fabsf(x))) + max_val); - - loss *= w; - - if (reduction == INFINIOP_REDUCTION_NONE) { - size_t out_offset = indexToOffset(idx, out_info.ndim, - out_info.shape, out_info.strides); - auto *out_ptr = static_cast(out_raw); - out_ptr[out_offset] = from_float(loss); - } else { - auto *out_accum = static_cast(out_raw); - atomicAdd(out_accum, static_cast(loss)); - } -} - -__global__ void bce_logits_mean_scale_kernel_f32(float *val, size_t count) { - if (threadIdx.x == 0 && blockIdx.x == 0) { - *val /= static_cast(count); - } -} - -template -__global__ void bce_logits_reduce_finalize_kernel( - Tdata *out, - float *workspace, - size_t count, - int is_mean) { - if (threadIdx.x == 0 && blockIdx.x == 0) { - float v = *workspace; - if (is_mean) { - v /= static_cast(count); - } - *out = from_float(v); - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *logits, - const void *target, - const void *weight, - const void *pos_weight, - void *stream) const { - - musaStream_t mustream = (musaStream_t)stream; - size_t n = _info.num_elements; - - if (_reduction != INFINIOP_REDUCTION_NONE && (_dtype == INFINI_DTYPE_F16 || _dtype == INFINI_DTYPE_BF16)) { - if (workspace_size < sizeof(float)) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - } - - int block = 256; - int grid = static_cast((n + block - 1) / block); - - BCETensorInfoDevice logits_info = make_device_info(_info.logits); - BCETensorInfoDevice target_info = make_device_info(_info.target); - BCETensorInfoDevice out_info = make_device_info(_info.out); - BCETensorInfoDevice weight_info = {}; - if (_info.weight.total_elements != 0) { - weight_info = make_device_info(_info.weight); - } - - size_t pos_weight_len = _info.pos_weight.total_elements; - - switch (_dtype) { - case INFINI_DTYPE_F32: { - if (_reduction != INFINIOP_REDUCTION_NONE) { - musaMemsetAsync(out, 0, sizeof(float), mustream); - } - - bce_logits_kernel<<>>( - out, - static_cast(logits), - static_cast(target), - static_cast(weight), - static_cast(pos_weight), - logits_info, - target_info, - weight_info, - out_info, - n, - pos_weight_len, - _reduction); - - if (_reduction == INFINIOP_REDUCTION_MEAN) { - bce_logits_mean_scale_kernel_f32<<<1, 1, 0, mustream>>>( - static_cast(out), n); - } - break; - } - case INFINI_DTYPE_F16: { - auto *logits_h = static_cast(logits); - auto *target_h = static_cast(target); - auto *weight_h = static_cast(weight); - auto *pos_weight_h = static_cast(pos_weight); - - void *out_raw = nullptr; - float *workspace_f = nullptr; - - if (_reduction == INFINIOP_REDUCTION_NONE) { - out_raw = out; - } else { - workspace_f = static_cast(workspace); - musaMemsetAsync(workspace_f, 0, sizeof(float), mustream); - out_raw = workspace_f; - } - - bce_logits_kernel<<>>( - out_raw, - logits_h, - target_h, - weight_h, - pos_weight_h, - logits_info, - target_info, - weight_info, - out_info, - n, - pos_weight_len, - _reduction); - - if (_reduction != INFINIOP_REDUCTION_NONE) { - int is_mean = (_reduction == INFINIOP_REDUCTION_MEAN) ? 1 : 0; - bce_logits_reduce_finalize_kernel<<<1, 1, 0, mustream>>>( - static_cast(out), workspace_f, n, is_mean); - } - break; - } - case INFINI_DTYPE_BF16: { - auto *logits_b = static_cast(logits); - auto *target_b = static_cast(target); - auto *weight_b = static_cast(weight); - auto *pos_weight_b = static_cast(pos_weight); - - void *out_raw = nullptr; - float *workspace_f = nullptr; - - if (_reduction == INFINIOP_REDUCTION_NONE) { - out_raw = out; - } else { - workspace_f = static_cast(workspace); - musaMemsetAsync(workspace_f, 0, sizeof(float), mustream); - out_raw = workspace_f; - } - - bce_logits_kernel<<>>( - out_raw, - logits_b, - target_b, - weight_b, - pos_weight_b, - logits_info, - target_info, - weight_info, - out_info, - n, - pos_weight_len, - _reduction); - - if (_reduction != INFINIOP_REDUCTION_NONE) { - int is_mean = (_reduction == INFINIOP_REDUCTION_MEAN) ? 1 : 0; - bce_logits_reduce_finalize_kernel<<<1, 1, 0, mustream>>>( - static_cast(out), workspace_f, n, is_mean); - } - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::bce_with_logits::moore diff --git a/src/infiniop/ops/binary_cross_entropy_with_logits/nvidia/binary_cross_entropy_with_logits_nvidia.cu b/src/infiniop/ops/binary_cross_entropy_with_logits/nvidia/binary_cross_entropy_with_logits_nvidia.cu deleted file mode 100644 index 0bd5fec98..000000000 --- a/src/infiniop/ops/binary_cross_entropy_with_logits/nvidia/binary_cross_entropy_with_logits_nvidia.cu +++ /dev/null @@ -1,334 +0,0 @@ -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "binary_cross_entropy_with_logits_nvidia.cuh" -#include -#include -#include - -namespace op::bce_with_logits::nvidia { - -using device::nvidia::indexToOffset; - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() = default; - -// 在 GPU 侧使用的简化张量信息（固定上限维度，支持 stride） -constexpr int BCE_MAX_DIMS = 8; - -struct BCETensorInfoDevice { - size_t ndim; - size_t shape[BCE_MAX_DIMS]; - ptrdiff_t strides[BCE_MAX_DIMS]; -}; - -static inline BCETensorInfoDevice make_device_info(const BCETensorInfo &info) { - BCETensorInfoDevice dev{}; - dev.ndim = info.ndim; - for (size_t i = 0; i < info.ndim && i < static_cast(BCE_MAX_DIMS); ++i) { - dev.shape[i] = info.dims[i]; - dev.strides[i] = info.stride[i]; - } - return dev; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t logits_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t pos_weight_desc, - infiniopReduction_t reduction) { - - auto handle = reinterpret_cast(handle_); - auto dtype = logits_desc->dtype(); - - // NVIDIA 实现支持 F16 / F32 / BF16 - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = BCEWithLogitsInfo::create(out_desc, logits_desc, target_desc, - weight_desc, pos_weight_desc, reduction); - CHECK_RESULT(result); - - auto info = result.take(); - - // F16/BF16 在做归约时需要一个 float 标量 workspace 来累加 - size_t workspace_size = 0; - if (reduction != INFINIOP_REDUCTION_NONE && (dtype == INFINI_DTYPE_F16 || dtype == INFINI_DTYPE_BF16)) { - workspace_size = sizeof(float); - } - - *desc_ptr = new Descriptor( - dtype, std::move(info), reduction, workspace_size, - nullptr, - handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -// 将任意标量类型提升为 float -template -__device__ __forceinline__ float to_float(T x) { - if constexpr (std::is_same_v) { - return x; - } else if constexpr (std::is_same_v) { - return __half2float(x); - } else if constexpr (std::is_same_v) { - return __bfloat162float(x); - } else { - return static_cast(x); - } -} - -// 从 float 转回目标标量类型 -template -__device__ __forceinline__ T from_float(float x) { - if constexpr (std::is_same_v) { - return x; - } else if constexpr (std::is_same_v) { - return __float2half(x); - } else if constexpr (std::is_same_v) { - return __float2bfloat16(x); - } else { - return static_cast(x); - } -} - -// --- CUDA Kernel: 支持 stride 的数值稳定 BCE 计算 --- -template -__global__ void bce_logits_kernel( - void *out_raw, - const Tdata *logits, - const Tdata *target, - const Tdata *weight, - const Tdata *pos_weight, - BCETensorInfoDevice logits_info, - BCETensorInfoDevice target_info, - BCETensorInfoDevice weight_info, - BCETensorInfoDevice out_info, - size_t n, - size_t pos_weight_len, - infiniopReduction_t reduction) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= n) { - return; - } - - // 计算逻辑索引在各张量中的偏移（支持任意 stride） - size_t logits_offset = indexToOffset(idx, logits_info.ndim, - logits_info.shape, logits_info.strides); - size_t target_offset = indexToOffset(idx, target_info.ndim, - target_info.shape, target_info.strides); - - float x = to_float(logits[logits_offset]); - float y = to_float(target[target_offset]); - - float pw = 1.0f; - if (pos_weight && pos_weight_len > 0) { - // 按最后一维广播：假设 pos_weight 是连续的一维张量 - size_t c = idx % pos_weight_len; - pw = to_float(pos_weight[c]); - } - - float w = 1.0f; - if (weight && weight_info.ndim > 0) { - size_t weight_offset = indexToOffset(idx, weight_info.ndim, - weight_info.shape, weight_info.strides); - w = to_float(weight[weight_offset]); - } - - // 数值稳定公式：max(x, 0) - x * y * pw + (1 + (pw - 1) * y) * log(1 + exp(-abs(x))) - // max_val = max(-x, 0) - // log_weight = 1 + (pos_weight - 1) * y - // loss = (1 - y) * x + log_weight * (log1p(exp(-|x|)) + max_val) - float max_val = fmaxf(-x, 0.0f); - float log_weight = 1.0f + (pw - 1.0f) * y; - float loss = (1.0f - y) * x + log_weight * (log1pf(expf(-fabsf(x))) + max_val); - loss *= w; - if (reduction == INFINIOP_REDUCTION_NONE) { - // 写回逐元素 loss（支持 stride 的 out） - size_t out_offset = indexToOffset(idx, out_info.ndim, - out_info.shape, out_info.strides); - auto *out_ptr = static_cast(out_raw); - out_ptr[out_offset] = from_float(loss); - } else { - // 对于 mean 或 sum，使用 float 累加到标量位置 - auto *out_accum = static_cast(out_raw); - atomicAdd(out_accum, static_cast(loss)); - } -} - -// F32 mean 归约：对输出标量做除法 -__global__ void bce_logits_mean_scale_kernel_f32(float *val, size_t count) { - if (threadIdx.x == 0 && blockIdx.x == 0) { - *val /= static_cast(count); - } -} - -// F16/BF16 归约：从 float workspace 写回目标 dtype -template -__global__ void bce_logits_reduce_finalize_kernel( - Tdata *out, - float *workspace, - size_t count, - int is_mean) { - if (threadIdx.x == 0 && blockIdx.x == 0) { - float v = *workspace; - if (is_mean) { - v /= static_cast(count); - } - *out = from_float(v); - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *logits, - const void *target, - const void *weight, - const void *pos_weight, - void *stream) const { - - cudaStream_t custream = (cudaStream_t)stream; - size_t n = _info.num_elements; - - // F16/BF16 + 归约需要 float workspace - if (_reduction != INFINIOP_REDUCTION_NONE && (_dtype == INFINI_DTYPE_F16 || _dtype == INFINI_DTYPE_BF16)) { - if (workspace_size < sizeof(float)) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - } - - int block = 256; - int grid = static_cast((n + block - 1) / block); - - // 构造 GPU 侧的张量信息（含 stride） - BCETensorInfoDevice logits_info = make_device_info(_info.logits); - BCETensorInfoDevice target_info = make_device_info(_info.target); - BCETensorInfoDevice out_info = make_device_info(_info.out); - BCETensorInfoDevice weight_info = {}; - if (_info.weight.total_elements != 0) { - weight_info = make_device_info(_info.weight); - } - - size_t pos_weight_len = _info.pos_weight.total_elements; - - switch (_dtype) { - case INFINI_DTYPE_F32: { - // 如果是规约操作，计算前需将输出位置清零 - if (_reduction != INFINIOP_REDUCTION_NONE) { - cudaMemsetAsync(out, 0, sizeof(float), custream); - } - - bce_logits_kernel<<>>( - out, - static_cast(logits), - static_cast(target), - static_cast(weight), - static_cast(pos_weight), - logits_info, - target_info, - weight_info, - out_info, - n, - pos_weight_len, - _reduction); - - if (_reduction == INFINIOP_REDUCTION_MEAN) { - bce_logits_mean_scale_kernel_f32<<<1, 1, 0, custream>>>( - static_cast(out), n); - } - break; - } - case INFINI_DTYPE_F16: { - auto *logits_h = static_cast(logits); - auto *target_h = static_cast(target); - auto *weight_h = static_cast(weight); - auto *pos_weight_h = static_cast(pos_weight); - - void *out_raw = nullptr; - float *workspace_f = nullptr; - - if (_reduction == INFINIOP_REDUCTION_NONE) { - out_raw = out; - } else { - workspace_f = static_cast(workspace); - cudaMemsetAsync(workspace_f, 0, sizeof(float), custream); - out_raw = workspace_f; - } - - bce_logits_kernel<<>>( - out_raw, - logits_h, - target_h, - weight_h, - pos_weight_h, - logits_info, - target_info, - weight_info, - out_info, - n, - pos_weight_len, - _reduction); - - if (_reduction != INFINIOP_REDUCTION_NONE) { - int is_mean = (_reduction == INFINIOP_REDUCTION_MEAN) ? 1 : 0; - bce_logits_reduce_finalize_kernel<<<1, 1, 0, custream>>>( - static_cast(out), workspace_f, n, is_mean); - } - - break; - } - case INFINI_DTYPE_BF16: { - auto *logits_b = static_cast(logits); - auto *target_b = static_cast(target); - auto *weight_b = static_cast(weight); - auto *pos_weight_b = static_cast(pos_weight); - - void *out_raw = nullptr; - float *workspace_f = nullptr; - - if (_reduction == INFINIOP_REDUCTION_NONE) { - out_raw = out; - } else { - workspace_f = static_cast(workspace); - cudaMemsetAsync(workspace_f, 0, sizeof(float), custream); - out_raw = workspace_f; - } - - bce_logits_kernel<<>>( - out_raw, - logits_b, - target_b, - weight_b, - pos_weight_b, - logits_info, - target_info, - weight_info, - out_info, - n, - pos_weight_len, - _reduction); - - if (_reduction != INFINIOP_REDUCTION_NONE) { - int is_mean = (_reduction == INFINIOP_REDUCTION_MEAN) ? 1 : 0; - bce_logits_reduce_finalize_kernel<<<1, 1, 0, custream>>>( - static_cast(out), workspace_f, n, is_mean); - } - - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto err = cudaGetLastError(); - return (err == cudaSuccess) ? INFINI_STATUS_SUCCESS : INFINI_STATUS_INTERNAL_ERROR; -} - -} // namespace op::bce_with_logits::nvidia diff --git a/src/infiniop/ops/binary_cross_entropy_with_logits/nvidia/binary_cross_entropy_with_logits_nvidia.cuh b/src/infiniop/ops/binary_cross_entropy_with_logits/nvidia/binary_cross_entropy_with_logits_nvidia.cuh deleted file mode 100644 index 4c0872881..000000000 --- a/src/infiniop/ops/binary_cross_entropy_with_logits/nvidia/binary_cross_entropy_with_logits_nvidia.cuh +++ /dev/null @@ -1,17 +0,0 @@ -#ifndef __BINARY_CROSS_ENTROPY_WITH_LOGITS_NVIDIA_CUH__ -#define __BINARY_CROSS_ENTROPY_WITH_LOGITS_NVIDIA_CUH__ - -#include "../binary_cross_entropy_with_logits.h" - -/** - * 使用 bce_with_logits.h 中定义的 DESCRIPTOR 宏。 - * 这将在命名空间 op::bce_with_logits::nvidia 中生成针对 NVIDIA 设备的 Descriptor 类。 - * * * 在 NVIDIA 端的实现（.cu 文件）中，Opaque 结构体通常包含： - * - cudnnHandle_t: 如果使用 cuDNN 的算子实现。 - * - cudnnTensorDescriptor_t: 用于描述各输入输出张量的 cuDNN 格式。 - * - KernelConfig: 用于自定义 CUDA Kernel 的网格（Grid）和线程块（Block）配置。 - * - dataType: 对应的 CUDA 数据类型（如 CUDA_R_32F）。 - */ -DESCRIPTOR(nvidia) - -#endif // __BINARY_CROSS_ENTROPY_WITH_LOGITS_NVIDIA_CUH__ diff --git a/src/infiniop/ops/binary_cross_entropy_with_logits/operator.cc b/src/infiniop/ops/binary_cross_entropy_with_logits/operator.cc deleted file mode 100644 index 06180aa2c..000000000 --- a/src/infiniop/ops/binary_cross_entropy_with_logits/operator.cc +++ /dev/null @@ -1,188 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/binary_cross_entropy_with_logits.h" - -// 引入各硬件后端的 Descriptor 定义 -#ifdef ENABLE_CPU_API -#include "cpu/binary_cross_entropy_with_logits_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/binary_cross_entropy_with_logits_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/binary_cross_entropy_with_logits_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/binary_cross_entropy_with_logits_moore.h" -#endif - -// ----------------------------------------------------------------------------- -// 1. 创建描述符 -// ----------------------------------------------------------------------------- -__INFINI_C infiniStatus_t infiniopCreateBCEWithLogitsDescriptor( - infiniopHandle_t handle, - infiniopBCEWithLogitsDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t logits_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t pos_weight_desc, - infiniopReduction_t reduction) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::bce_with_logits::NAMESPACE::Descriptor::create(handle, \ - reinterpret_cast(desc_ptr), \ - out_desc, logits_desc, target_desc, weight_desc, pos_weight_desc, reduction) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ----------------------------------------------------------------------------- -// 2. 获取 Workspace 大小 -// ----------------------------------------------------------------------------- -__INFINI_C infiniStatus_t infiniopGetBCEWithLogitsWorkspaceSize( - infiniopBCEWithLogitsDescriptor_t desc, - size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -// ----------------------------------------------------------------------------- -// 3. 执行计算 -// ----------------------------------------------------------------------------- -__INFINI_C infiniStatus_t infiniopBCEWithLogits( - infiniopBCEWithLogitsDescriptor_t desc, - void *workspace, size_t workspace_size, - void *out, - const void *logits, - const void *target, - const void *weight, - const void *pos_weight, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, out, logits, target, weight, pos_weight, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -// ----------------------------------------------------------------------------- -// 4. 销毁描述符 -// ----------------------------------------------------------------------------- -__INFINI_C infiniStatus_t infiniopDestroyBCEWithLogitsDescriptor(infiniopBCEWithLogitsDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} diff --git a/src/infiniop/ops/bitwise_right_shift/cpu/bitwise_right_shift_cpu.cc b/src/infiniop/ops/bitwise_right_shift/cpu/bitwise_right_shift_cpu.cc deleted file mode 100644 index a2b1a71d9..000000000 --- a/src/infiniop/ops/bitwise_right_shift/cpu/bitwise_right_shift_cpu.cc +++ /dev/null @@ -1,66 +0,0 @@ -#include "bitwise_right_shift_cpu.h" - -namespace op::bitwise_right_shift::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &shift_desc = input_desc_vec.at(1); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_I8, INFINI_DTYPE_I16, INFINI_DTYPE_I32, INFINI_DTYPE_I64, - INFINI_DTYPE_U8, INFINI_DTYPE_U16, INFINI_DTYPE_U32, INFINI_DTYPE_U64); - - if (input_desc->dtype() != dtype || shift_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_I8: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_U8: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_U16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_U32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_U64: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::bitwise_right_shift::cpu diff --git a/src/infiniop/ops/bitwise_right_shift/cpu/bitwise_right_shift_cpu.h b/src/infiniop/ops/bitwise_right_shift/cpu/bitwise_right_shift_cpu.h deleted file mode 100644 index 71ec11913..000000000 --- a/src/infiniop/ops/bitwise_right_shift/cpu/bitwise_right_shift_cpu.h +++ /dev/null @@ -1,43 +0,0 @@ -#ifndef __BITWISE_RIGHT_SHIFT_CPU_H__ -#define __BITWISE_RIGHT_SHIFT_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" -#include -#include - -ELEMENTWISE_DESCRIPTOR(bitwise_right_shift, cpu) - -namespace op::bitwise_right_shift::cpu { -typedef struct BitwiseRightShiftOp { -public: - static constexpr size_t num_inputs = 2; - template - T operator()(const T &x, const T &shift) const { - constexpr unsigned kBits = static_cast(sizeof(T) * 8); - using WideUnsigned = std::conditional_t<(kBits <= 32), uint32_t, uint64_t>; - using WideSigned = std::conditional_t<(kBits <= 32), int32_t, int64_t>; - - if constexpr (std::is_signed_v) { - const WideSigned xw = static_cast(x); - const WideSigned sw = static_cast(shift); - - if (sw < 0 || sw >= static_cast(kBits)) { - return static_cast(xw < 0 ? WideSigned(-1) : WideSigned(0)); - } - - return static_cast(xw >> static_cast(sw)); - } else { - const WideUnsigned xw = static_cast(x); - const WideUnsigned sw = static_cast(shift); - - if (sw >= static_cast(kBits)) { - return static_cast(0); - } - - return static_cast(xw >> static_cast(sw)); - } - } -} BitwiseRightShiftOp; -} // namespace op::bitwise_right_shift::cpu - -#endif // __BITWISE_RIGHT_SHIFT_CPU_H__ diff --git a/src/infiniop/ops/bitwise_right_shift/cuda/kernel.cuh b/src/infiniop/ops/bitwise_right_shift/cuda/kernel.cuh deleted file mode 100644 index a183df974..000000000 --- a/src/infiniop/ops/bitwise_right_shift/cuda/kernel.cuh +++ /dev/null @@ -1,38 +0,0 @@ -#pragma once -#include -#include - -namespace op::bitwise_right_shift::cuda { - -struct BitwiseRightShiftOp { - static constexpr size_t num_inputs = 2; - - template - __device__ __forceinline__ T operator()(const T &x, const T &shift) const { - constexpr unsigned kBits = static_cast(sizeof(T) * 8); - using WideUnsigned = std::conditional_t<(kBits <= 32), uint32_t, uint64_t>; - using WideSigned = std::conditional_t<(kBits <= 32), int32_t, int64_t>; - - if constexpr (std::is_signed_v) { - const WideSigned xw = static_cast(x); - const WideSigned sw = static_cast(shift); - - if (sw < 0 || sw >= static_cast(kBits)) { - return static_cast(xw < 0 ? WideSigned(-1) : WideSigned(0)); - } - - return static_cast(xw >> static_cast(sw)); - } else { - const WideUnsigned xw = static_cast(x); - const WideUnsigned sw = static_cast(shift); - - if (sw >= static_cast(kBits)) { - return static_cast(0); - } - - return static_cast(xw >> static_cast(sw)); - } - } -}; - -} // namespace op::bitwise_right_shift::cuda diff --git a/src/infiniop/ops/bitwise_right_shift/metax/bitwise_right_shift_metax.h b/src/infiniop/ops/bitwise_right_shift/metax/bitwise_right_shift_metax.h deleted file mode 100644 index 2bfb7a4b3..000000000 --- a/src/infiniop/ops/bitwise_right_shift/metax/bitwise_right_shift_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BITWISE_RIGHT_SHIFT_METAX_API_H__ -#define __BITWISE_RIGHT_SHIFT_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(bitwise_right_shift, metax) - -#endif // __BITWISE_RIGHT_SHIFT_METAX_API_H__ diff --git a/src/infiniop/ops/bitwise_right_shift/metax/bitwise_right_shift_metax.maca b/src/infiniop/ops/bitwise_right_shift/metax/bitwise_right_shift_metax.maca deleted file mode 100644 index 34d470fa4..000000000 --- a/src/infiniop/ops/bitwise_right_shift/metax/bitwise_right_shift_metax.maca +++ /dev/null @@ -1,68 +0,0 @@ -#include "bitwise_right_shift_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::bitwise_right_shift::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_I8, INFINI_DTYPE_I16, INFINI_DTYPE_I32, INFINI_DTYPE_I64, - INFINI_DTYPE_U8, INFINI_DTYPE_U16, INFINI_DTYPE_U32, INFINI_DTYPE_U64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_I8: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int8_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I16: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int16_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int64_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U8: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint8_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U16: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint16_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U32: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U64: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint64_t>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::bitwise_right_shift::metax diff --git a/src/infiniop/ops/bitwise_right_shift/moore/bitwise_right_shift_moore.h b/src/infiniop/ops/bitwise_right_shift/moore/bitwise_right_shift_moore.h deleted file mode 100644 index 23170e15f..000000000 --- a/src/infiniop/ops/bitwise_right_shift/moore/bitwise_right_shift_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BITWISE_RIGHT_SHIFT_MOORE_API_H__ -#define __BITWISE_RIGHT_SHIFT_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(bitwise_right_shift, moore) - -#endif // __BITWISE_RIGHT_SHIFT_MOORE_API_H__ diff --git a/src/infiniop/ops/bitwise_right_shift/moore/bitwise_right_shift_moore.mu b/src/infiniop/ops/bitwise_right_shift/moore/bitwise_right_shift_moore.mu deleted file mode 100644 index 2026a9cd6..000000000 --- a/src/infiniop/ops/bitwise_right_shift/moore/bitwise_right_shift_moore.mu +++ /dev/null @@ -1,69 +0,0 @@ -#include "bitwise_right_shift_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "../cuda/kernel.cuh" - -namespace op::bitwise_right_shift::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_I8, INFINI_DTYPE_I16, INFINI_DTYPE_I32, INFINI_DTYPE_I64, - INFINI_DTYPE_U8, INFINI_DTYPE_U16, INFINI_DTYPE_U32, INFINI_DTYPE_U64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_I8: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int8_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I16: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int16_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int64_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U8: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint8_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U16: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint16_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U32: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U64: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint64_t>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::bitwise_right_shift::moore diff --git a/src/infiniop/ops/bitwise_right_shift/nvidia/bitwise_right_shift_nvidia.cu b/src/infiniop/ops/bitwise_right_shift/nvidia/bitwise_right_shift_nvidia.cu deleted file mode 100644 index 6a2972d4a..000000000 --- a/src/infiniop/ops/bitwise_right_shift/nvidia/bitwise_right_shift_nvidia.cu +++ /dev/null @@ -1,72 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "bitwise_right_shift_nvidia.cuh" - -namespace op::bitwise_right_shift::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &shift_desc = input_desc_vec.at(1); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_I8, INFINI_DTYPE_I16, INFINI_DTYPE_I32, INFINI_DTYPE_I64, - INFINI_DTYPE_U8, INFINI_DTYPE_U16, INFINI_DTYPE_U32, INFINI_DTYPE_U64); - - if (input_desc->dtype() != dtype || shift_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_I8: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int8_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I16: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int16_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int64_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U8: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint8_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U16: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint16_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U32: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U64: - return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint64_t>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::bitwise_right_shift::nvidia diff --git a/src/infiniop/ops/bitwise_right_shift/nvidia/bitwise_right_shift_nvidia.cuh b/src/infiniop/ops/bitwise_right_shift/nvidia/bitwise_right_shift_nvidia.cuh deleted file mode 100644 index 66cf6fcfd..000000000 --- a/src/infiniop/ops/bitwise_right_shift/nvidia/bitwise_right_shift_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BITWISE_RIGHT_SHIFT_NVIDIA_H__ -#define __BITWISE_RIGHT_SHIFT_NVIDIA_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(bitwise_right_shift, nvidia) - -#endif // __BITWISE_RIGHT_SHIFT_NVIDIA_H__ diff --git a/src/infiniop/ops/bitwise_right_shift/operator.cc b/src/infiniop/ops/bitwise_right_shift/operator.cc deleted file mode 100644 index b1f3d81ed..000000000 --- a/src/infiniop/ops/bitwise_right_shift/operator.cc +++ /dev/null @@ -1,171 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/bitwise_right_shift.h" - -#ifdef ENABLE_CPU_API -#include "cpu/bitwise_right_shift_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/bitwise_right_shift_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/bitwise_right_shift_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/bitwise_right_shift_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateBitwiseRightShiftDescriptor( - infiniopHandle_t handle, - infiniopBitwiseRightShiftDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::bitwise_right_shift::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - {x1_desc, x2_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetBitwiseRightShiftWorkspaceSize(infiniopBitwiseRightShiftDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopBitwiseRightShift( - infiniopBitwiseRightShiftDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, {x1, x2}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyBitwiseRightShiftDescriptor(infiniopBitwiseRightShiftDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/blas_amax/bang/blas_amax_bang.h b/src/infiniop/ops/blas_amax/bang/blas_amax_bang.h deleted file mode 100644 index 8bc6ca2c5..000000000 --- a/src/infiniop/ops/blas_amax/bang/blas_amax_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BLAS_AMAX_BANG_H__ -#define __BLAS_AMAX_BANG_H__ - -#include "../blas_amax.h" - -DESCRIPTOR(bang) - -#endif // __BLAS_AMAX_BANG_H__ diff --git a/src/infiniop/ops/blas_amax/bang/blas_amax_bang.mlu b/src/infiniop/ops/blas_amax/bang/blas_amax_bang.mlu deleted file mode 100644 index 26cc3c501..000000000 --- a/src/infiniop/ops/blas_amax/bang/blas_amax_bang.mlu +++ /dev/null @@ -1,96 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "blas_amax_bang.h" -#include "blas_amax_bang_kernel.mlu" - -namespace op::blas_amax::bang { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = BlasAmaxInfo::createBlasAmaxInfo(x_desc, result_desc); - CHECK_RESULT(result); - - // Create descriptor - *desc_ptr = new Descriptor( - result.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateBlasAmax( - const BlasAmaxInfo &info, - const Tdata *x, - int *result, - cnrtQueue_t queue) { - - const int n = utils::cast(info.n); - const int incx = utils::cast(info.incx); - - cnrtDim3_t k_dim; - cnrtFunctionType_t k_type; - - k_dim.x = 4; - k_dim.y = 1; - k_dim.z = 1; - k_type = cnrtFuncTypeUnion1; - - if (incx == 1) { - blasAmaxKernelContiguous<<>>( - n, - x, - result); - } else { - blasAmaxKernelStrided<<>>( - n, - x, - incx, - result); - } - - cnrtQueueSync(queue); - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_BLAS_AMAX(TDATA) \ - calculateBlasAmax(_info, \ - (const TDATA *)x, \ - (int *)result, \ - (cnrtQueue_t)stream) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - switch (_info.data_type) { - case INFINI_DTYPE_F16: - return CALCULATE_BLAS_AMAX(half); - case INFINI_DTYPE_F32: - return CALCULATE_BLAS_AMAX(float); - case INFINI_DTYPE_BF16: - return CALCULATE_BLAS_AMAX(bfloat16_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE_BLAS_AMAX - -} // namespace op::blas_amax::bang diff --git a/src/infiniop/ops/blas_amax/bang/blas_amax_bang_kernel.mlu b/src/infiniop/ops/blas_amax/bang/blas_amax_bang_kernel.mlu deleted file mode 100644 index 2939f2dd7..000000000 --- a/src/infiniop/ops/blas_amax/bang/blas_amax_bang_kernel.mlu +++ /dev/null @@ -1,136 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "blas_amax_bang.h" - -__nram__ char nram_buffer[NRAM_MAX_SIZE]; - -template -__mlu_global__ void blasAmaxKernelContiguous( - int n, - const Tdata *x, - int *result) { - - __mlu_shared__ int shared_max_index[4]; - __mlu_shared__ Tdata shared_max_value[4]; - - char *nram_aligned = (char *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1)); - - size_t nram_usable = NRAM_MAX_SIZE - (nram_aligned - nram_buffer); - size_t max_chunk_elements = nram_usable / sizeof(Tdata); - - size_t align_elements = ALIGN_SIZE / sizeof(Tdata); - if (align_elements == 0) { - align_elements = 1; - } - int chunk_size = (int)((max_chunk_elements / align_elements) * align_elements); - - Tdata *nram_x = (Tdata *)nram_aligned; - - int elements_per_core = n / taskDim; - int remain = n % taskDim; - int core_elements = elements_per_core + (taskId < remain ? 1 : 0); - int core_offset = taskId < remain ? taskId * core_elements : taskId * elements_per_core + remain; - - int chunks = core_elements / chunk_size; - int chunk_rem = core_elements % chunk_size; - - int max_index = -1; - Tdata max_value = static_cast(0); - - for (int c = 0; c < chunks; c++) { - int current_offset = core_offset + c * chunk_size; - __memcpy(nram_x, x + current_offset, chunk_size * sizeof(Tdata), GDRAM2NRAM); - - __bang_abs(nram_x, nram_x, chunk_size); - - for (int i = 0; i < chunk_size; i++) { - Tdata abs_val = nram_x[i]; - if (abs_val > max_value) { - max_value = abs_val; - max_index = current_offset + i; - } - } - } - - if (chunk_rem > 0) { - int current_offset = core_offset + chunks * chunk_size; - - __memcpy(nram_x, x + current_offset, chunk_rem * sizeof(Tdata), GDRAM2NRAM); - - __bang_abs(nram_x, nram_x, chunk_rem); - - for (int i = 0; i < chunk_rem; i++) { - Tdata abs_val = nram_x[i]; - if (abs_val > max_value) { - max_value = abs_val; - max_index = current_offset + i; - } - } - } - - shared_max_index[coreId] = max_index; - shared_max_value[coreId] = max_value; - - __sync_cluster(); - - if (coreId == 0) { - int cluster_max_index = -1; - Tdata cluster_max_value = static_cast(0); - - for (int i = 0; i < coreDim; i++) { - if (shared_max_value[i] > cluster_max_value) { - cluster_max_value = shared_max_value[i]; - cluster_max_index = shared_max_index[i]; - } - } - - result[0] = cluster_max_index + 1; // Convert to 1-based index - } -} - -template -__mlu_global__ void blasAmaxKernelStrided( - int n, - const Tdata *x, - int incx, - int *result) { - - __mlu_shared__ int shared_max_index[4]; - __mlu_shared__ Tdata shared_max_value[4]; - - int elements_per_core = n / taskDim; - int remain = n % taskDim; - int actual_tasks = elements_per_core + (taskId < remain ? 1 : 0); - int start_idx = taskId < remain ? taskId * actual_tasks : taskId * elements_per_core + remain; - - int max_index = -1; - Tdata max_value = static_cast(0); - - for (int i = start_idx; i < start_idx + actual_tasks; ++i) { - int offset = i * incx; - Tdata abs_val = x[offset] > static_cast(0) ? x[offset] : -x[offset]; - - if (abs_val > max_value) { - max_value = abs_val; - max_index = i; - } - } - - shared_max_index[coreId] = max_index; - shared_max_value[coreId] = max_value; - - __sync_cluster(); - - if (coreId == 0) { - int cluster_max_index = -1; - Tdata cluster_max_value = static_cast(0); - - for (int i = 0; i < coreDim; i++) { - if (shared_max_value[i] > cluster_max_value) { - cluster_max_value = shared_max_value[i]; - cluster_max_index = shared_max_index[i]; - } - } - - result[0] = cluster_max_index + 1; // Convert to 1-based index - } -} diff --git a/src/infiniop/ops/blas_amax/blas_amax.h b/src/infiniop/ops/blas_amax/blas_amax.h deleted file mode 100644 index 627b7b754..000000000 --- a/src/infiniop/ops/blas_amax/blas_amax.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef __BLAS_AMAX_H__ -#define __BLAS_AMAX_H__ - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::blas_amax::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - BlasAmaxInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - BlasAmaxInfo info, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(std::move(info)), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t x_desc, \ - infiniopTensorDescriptor_t result_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - const void *x, \ - void *result, \ - void *stream) const; \ - }; \ - } - -#endif // __BLAS_AMAX_H__ diff --git a/src/infiniop/ops/blas_amax/cpu/blas_amax_cpu.cc b/src/infiniop/ops/blas_amax/cpu/blas_amax_cpu.cc deleted file mode 100644 index 4b3553041..000000000 --- a/src/infiniop/ops/blas_amax/cpu/blas_amax_cpu.cc +++ /dev/null @@ -1,105 +0,0 @@ -#include "blas_amax_cpu.h" -#include "../../../devices/cpu/common_cpu.h" - -namespace op::blas_amax::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = BlasAmaxInfo::createBlasAmaxInfo(x_desc, result_desc); - CHECK_RESULT(result); - - // Create descriptor - *desc_ptr = new Descriptor( - result.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateBlasAmax( - const BlasAmaxInfo &info, - const Tdata *x, - int *result) { - - const size_t n = info.n; - const ptrdiff_t incx = info.incx; - - if (n < 1 || incx == 0) { - result[0] = 0; - return INFINI_STATUS_SUCCESS; - } - - size_t max_index = 0; - if constexpr (std::is_same::value || std::is_same::value) { - float max_value = std::abs(utils::cast(x[0])); - - for (size_t i = 1; i < n; ++i) { - const ptrdiff_t idx = utils::cast(i) * incx; - float current_value = std::abs(utils::cast(x[idx])); - if (current_value > max_value) { - max_value = current_value; - max_index = i; - } - } - } else { - Tdata max_value = std::abs(x[0]); - - for (size_t i = 1; i < n; ++i) { - const ptrdiff_t idx = utils::cast(i) * incx; - Tdata current_value = std::abs(x[idx]); - if (current_value > max_value) { - max_value = current_value; - max_index = i; - } - } - } - - result[0] = utils::cast(max_index) + 1; - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_BLAS_AMAX(TDATA) \ - calculateBlasAmax(_info, \ - (const TDATA *)x, \ - (int *)result) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream) const { - - (void)workspace; - (void)workspace_size; - (void)stream; - - switch (_info.data_type) { - case INFINI_DTYPE_F16: - return CALCULATE_BLAS_AMAX(fp16_t); - case INFINI_DTYPE_F32: - return CALCULATE_BLAS_AMAX(float); - case INFINI_DTYPE_F64: - return CALCULATE_BLAS_AMAX(double); - case INFINI_DTYPE_BF16: - return CALCULATE_BLAS_AMAX(bf16_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE_BLAS_AMAX - -} // namespace op::blas_amax::cpu diff --git a/src/infiniop/ops/blas_amax/cpu/blas_amax_cpu.h b/src/infiniop/ops/blas_amax/cpu/blas_amax_cpu.h deleted file mode 100644 index 66197c9c1..000000000 --- a/src/infiniop/ops/blas_amax/cpu/blas_amax_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BLAS_AMAX_CPU_H__ -#define __BLAS_AMAX_CPU_H__ - -#include "../blas_amax.h" - -DESCRIPTOR(cpu) - -#endif // __BLAS_AMAX_CPU_H__ diff --git a/src/infiniop/ops/blas_amax/info.h b/src/infiniop/ops/blas_amax/info.h deleted file mode 100644 index b10f84046..000000000 --- a/src/infiniop/ops/blas_amax/info.h +++ /dev/null @@ -1,41 +0,0 @@ -#ifndef __BLAS_AMAX_INFO_H__ -#define __BLAS_AMAX_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" - -class BlasAmaxInfo { -private: - BlasAmaxInfo() = default; - -public: - size_t n; - ptrdiff_t incx; - infiniDtype_t data_type; - - static utils::Result createBlasAmaxInfo( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - - CHECK_OR_RETURN(x_desc != nullptr, INFINI_STATUS_NULL_POINTER); - CHECK_OR_RETURN(result_desc != nullptr, INFINI_STATUS_NULL_POINTER); - - auto data_type = x_desc->dtype(); - auto itype = result_desc->dtype(); - CHECK_DTYPE(data_type, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - CHECK_DTYPE(itype, INFINI_DTYPE_I32); - - CHECK_OR_RETURN(x_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(result_desc->numel() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - - auto n = x_desc->numel(); - auto incx = x_desc->stride(0); - - return utils::Result(BlasAmaxInfo{ - n, - incx, - data_type}); - } -}; - -#endif // __BLAS_AMAX_INFO_H__ diff --git a/src/infiniop/ops/blas_amax/metax/blas_amax_metax.cc b/src/infiniop/ops/blas_amax/metax/blas_amax_metax.cc deleted file mode 100644 index 05816020d..000000000 --- a/src/infiniop/ops/blas_amax/metax/blas_amax_metax.cc +++ /dev/null @@ -1,71 +0,0 @@ -#include "blas_amax_metax.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" - -namespace op::blas_amax::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = BlasAmaxInfo::createBlasAmaxInfo(x_desc, result_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - new Opaque{handle->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - const int n = utils::cast(_info.n); - const int incx = utils::cast(_info.incx); - const infiniDtype_t data_type = _info.data_type; - - CHECK_STATUS(_opaque->internal->useMcblas( - (hcStream_t)stream, - [&](hcblasHandle_t handle) { - CHECK_MCBLAS(hcblasSetPointerMode(handle, HCBLAS_POINTER_MODE_DEVICE)); - - switch (data_type) { - case INFINI_DTYPE_F32: - CHECK_MCBLAS(hcblasIsamax(handle, n, (const float *)x, incx, (int *)result)); - break; - case INFINI_DTYPE_F64: - CHECK_MCBLAS(hcblasIdamax(handle, n, (const double *)x, incx, (int *)result)); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; - })); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::blas_amax::metax diff --git a/src/infiniop/ops/blas_amax/metax/blas_amax_metax.h b/src/infiniop/ops/blas_amax/metax/blas_amax_metax.h deleted file mode 100644 index 19e79851f..000000000 --- a/src/infiniop/ops/blas_amax/metax/blas_amax_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BLAS_AMAX_METAX_H__ -#define __BLAS_AMAX_METAX_H__ - -#include "../blas_amax.h" - -DESCRIPTOR(metax) - -#endif // __BLAS_AMAX_METAX_H__ diff --git a/src/infiniop/ops/blas_amax/operator.cc b/src/infiniop/ops/blas_amax/operator.cc deleted file mode 100644 index c6b48eeb4..000000000 --- a/src/infiniop/ops/blas_amax/operator.cc +++ /dev/null @@ -1,122 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/blas_amax.h" - -#ifdef ENABLE_CPU_API -#include "cpu/blas_amax_cpu.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/blas_amax_metax.h" -#endif -#ifdef ENABLE_CAMBRICON_API -#include "bang/blas_amax_bang.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateBlasAmaxDescriptor( - infiniopHandle_t handle, - infiniopBlasAmaxDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::blas_amax::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - x_desc, \ - result_desc) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetBlasAmaxWorkspaceSize(infiniopBlasAmaxDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - GET(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET -} - -__INFINI_C infiniStatus_t infiniopBlasAmax( - infiniopBlasAmaxDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, x, result, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyBlasAmaxDescriptor(infiniopBlasAmaxDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - DELETE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/blas_amin/bang/blas_amin_bang.h b/src/infiniop/ops/blas_amin/bang/blas_amin_bang.h deleted file mode 100644 index ba9dbaa21..000000000 --- a/src/infiniop/ops/blas_amin/bang/blas_amin_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BLAS_AMIN_BANG_H__ -#define __BLAS_AMIN_BANG_H__ - -#include "../blas_amin.h" - -DESCRIPTOR(bang) - -#endif // __BLAS_AMIN_BANG_H__ diff --git a/src/infiniop/ops/blas_amin/bang/blas_amin_bang.mlu b/src/infiniop/ops/blas_amin/bang/blas_amin_bang.mlu deleted file mode 100644 index dfd2a7e64..000000000 --- a/src/infiniop/ops/blas_amin/bang/blas_amin_bang.mlu +++ /dev/null @@ -1,96 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "blas_amin_bang.h" -#include "blas_amin_bang_kernel.mlu" - -namespace op::blas_amin::bang { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = BlasAminInfo::createBlasAminInfo(x_desc, result_desc); - CHECK_RESULT(result); - - // Create descriptor - *desc_ptr = new Descriptor( - result.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateBlasAmin( - const BlasAminInfo &info, - const Tdata *x, - int *result, - cnrtQueue_t queue) { - - const int n = utils::cast(info.n); - const int incx = utils::cast(info.incx); - - cnrtDim3_t k_dim; - cnrtFunctionType_t k_type; - - k_dim.x = 4; - k_dim.y = 1; - k_dim.z = 1; - k_type = cnrtFuncTypeUnion1; - - if (incx == 1) { - blasAminKernelContiguous<<>>( - n, - x, - result); - } else { - blasAminKernelStrided<<>>( - n, - x, - incx, - result); - } - - cnrtQueueSync(queue); - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_BLAS_AMIN(TDATA) \ - calculateBlasAmin(_info, \ - (const TDATA *)x, \ - (int *)result, \ - (cnrtQueue_t)stream) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - switch (_info.data_type) { - case INFINI_DTYPE_F16: - return CALCULATE_BLAS_AMIN(half); - case INFINI_DTYPE_F32: - return CALCULATE_BLAS_AMIN(float); - case INFINI_DTYPE_BF16: - return CALCULATE_BLAS_AMIN(bfloat16_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE_BLAS_AMIN - -} // namespace op::blas_amin::bang diff --git a/src/infiniop/ops/blas_amin/bang/blas_amin_bang_kernel.mlu b/src/infiniop/ops/blas_amin/bang/blas_amin_bang_kernel.mlu deleted file mode 100644 index 2c5ff7c48..000000000 --- a/src/infiniop/ops/blas_amin/bang/blas_amin_bang_kernel.mlu +++ /dev/null @@ -1,135 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "blas_amin_bang.h" - -__nram__ char nram_buffer[NRAM_MAX_SIZE]; - -template -__mlu_global__ void blasAminKernelContiguous( - int n, - const Tdata *x, - int *result) { - - __mlu_shared__ int shared_min_index[4]; - __mlu_shared__ Tdata shared_min_value[4]; - - char *nram_aligned = (char *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1)); - - size_t nram_usable = NRAM_MAX_SIZE - (nram_aligned - nram_buffer); - size_t max_chunk_elements = nram_usable / sizeof(Tdata); - - size_t align_elements = ALIGN_SIZE / sizeof(Tdata); - if (align_elements == 0) { - align_elements = 1; - } - int chunk_size = (int)((max_chunk_elements / align_elements) * align_elements); - - Tdata *nram_x = (Tdata *)nram_aligned; - - int elements_per_core = n / taskDim; - int remain = n % taskDim; - int core_elements = elements_per_core + (taskId < remain ? 1 : 0); - int core_offset = taskId < remain ? taskId * core_elements : taskId * elements_per_core + remain; - - int chunks = core_elements / chunk_size; - int chunk_rem = core_elements % chunk_size; - - int min_index = -1; - Tdata min_value = static_cast(0); - bool initialized = false; - - for (int c = 0; c < chunks; c++) { - int current_offset = core_offset + c * chunk_size; - __memcpy(nram_x, x + current_offset, chunk_size * sizeof(Tdata), GDRAM2NRAM); - - __bang_abs(nram_x, nram_x, chunk_size); - - for (int i = 0; i < chunk_size; i++) { - Tdata abs_val = nram_x[i]; - if (!initialized || abs_val < min_value) { - min_value = abs_val; - min_index = current_offset + i; - initialized = true; - } - } - } - - if (chunk_rem > 0) { - int current_offset = core_offset + chunks * chunk_size; - - __memcpy(nram_x, x + current_offset, chunk_rem * sizeof(Tdata), GDRAM2NRAM); - - __bang_abs(nram_x, nram_x, chunk_rem); - - for (int i = 0; i < chunk_rem; i++) { - Tdata abs_val = nram_x[i]; - if (!initialized || abs_val < min_value) { - min_value = abs_val; - min_index = current_offset + i; - initialized = true; - } - } - } - - shared_min_index[coreId] = min_index; - shared_min_value[coreId] = min_value; - - __sync_cluster(); - - if (coreId == 0) { - for (int i = 1; i < coreDim; i++) { - if (shared_min_index[i] >= 0 && shared_min_value[i] < min_value) { - min_value = shared_min_value[i]; - min_index = shared_min_index[i]; - } - } - - result[0] = min_index + 1; // Convert to 1-based index - } -} - -template -__mlu_global__ void blasAminKernelStrided( - int n, - const Tdata *x, - int incx, - int *result) { - - __mlu_shared__ int shared_min_index[4]; - __mlu_shared__ Tdata shared_min_value[4]; - - int elements_per_core = n / taskDim; - int remain = n % taskDim; - int actual_tasks = elements_per_core + (taskId < remain ? 1 : 0); - int start_idx = taskId < remain ? taskId * actual_tasks : taskId * elements_per_core + remain; - - int min_index = -1; - Tdata min_value = static_cast(0); - bool initialized = false; - - for (int i = start_idx; i < start_idx + actual_tasks; ++i) { - int offset = i * incx; - Tdata abs_val = x[offset] > static_cast(0) ? x[offset] : -x[offset]; - - if (!initialized || abs_val < min_value) { - min_value = abs_val; - min_index = i; - initialized = true; - } - } - - shared_min_index[coreId] = min_index; - shared_min_value[coreId] = min_value; - - __sync_cluster(); - - if (coreId == 0) { - for (int i = 1; i < coreDim; i++) { - if (shared_min_index[i] >= 0 && shared_min_value[i] < min_value) { - min_value = shared_min_value[i]; - min_index = shared_min_index[i]; - } - } - - result[0] = min_index + 1; // Convert to 1-based index - } -} diff --git a/src/infiniop/ops/blas_amin/blas_amin.h b/src/infiniop/ops/blas_amin/blas_amin.h deleted file mode 100644 index 5128fcf85..000000000 --- a/src/infiniop/ops/blas_amin/blas_amin.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef __BLAS_AMIN_H__ -#define __BLAS_AMIN_H__ - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::blas_amin::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - BlasAminInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - BlasAminInfo info, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(std::move(info)), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t x_desc, \ - infiniopTensorDescriptor_t result_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - const void *x, \ - void *result, \ - void *stream) const; \ - }; \ - } - -#endif // __BLAS_AMIN_H__ diff --git a/src/infiniop/ops/blas_amin/cpu/blas_amin_cpu.cc b/src/infiniop/ops/blas_amin/cpu/blas_amin_cpu.cc deleted file mode 100644 index 07cad0461..000000000 --- a/src/infiniop/ops/blas_amin/cpu/blas_amin_cpu.cc +++ /dev/null @@ -1,105 +0,0 @@ -#include "blas_amin_cpu.h" -#include "../../../devices/cpu/common_cpu.h" - -namespace op::blas_amin::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = BlasAminInfo::createBlasAminInfo(x_desc, result_desc); - CHECK_RESULT(result); - - // Create descriptor - *desc_ptr = new Descriptor( - result.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateBlasAmin( - const BlasAminInfo &info, - const Tdata *x, - int *result) { - - const size_t n = info.n; - const ptrdiff_t incx = info.incx; - - if (n < 1 || incx == 0) { - result[0] = 0; - return INFINI_STATUS_SUCCESS; - } - - size_t min_index = 0; - if constexpr (std::is_same::value || std::is_same::value) { - float min_value = std::abs(utils::cast(x[0])); - - for (size_t i = 1; i < n; ++i) { - const ptrdiff_t idx = utils::cast(i) * incx; - float current_value = std::abs(utils::cast(x[idx])); - if (current_value < min_value) { - min_value = current_value; - min_index = i; - } - } - } else { - Tdata min_value = std::abs(x[0]); - - for (size_t i = 1; i < n; ++i) { - const ptrdiff_t idx = utils::cast(i) * incx; - Tdata current_value = std::abs(x[idx]); - if (current_value < min_value) { - min_value = current_value; - min_index = i; - } - } - } - - result[0] = utils::cast(min_index) + 1; - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_BLAS_AMIN(TDATA) \ - calculateBlasAmin(_info, \ - (const TDATA *)x, \ - (int *)result) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream) const { - - (void)workspace; - (void)workspace_size; - (void)stream; - - switch (_info.data_type) { - case INFINI_DTYPE_F16: - return CALCULATE_BLAS_AMIN(fp16_t); - case INFINI_DTYPE_F32: - return CALCULATE_BLAS_AMIN(float); - case INFINI_DTYPE_F64: - return CALCULATE_BLAS_AMIN(double); - case INFINI_DTYPE_BF16: - return CALCULATE_BLAS_AMIN(bf16_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE_BLAS_AMIN - -} // namespace op::blas_amin::cpu diff --git a/src/infiniop/ops/blas_amin/cpu/blas_amin_cpu.h b/src/infiniop/ops/blas_amin/cpu/blas_amin_cpu.h deleted file mode 100644 index c5e4936d7..000000000 --- a/src/infiniop/ops/blas_amin/cpu/blas_amin_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BLAS_AMIN_CPU_H__ -#define __BLAS_AMIN_CPU_H__ - -#include "../blas_amin.h" - -DESCRIPTOR(cpu) - -#endif // __BLAS_AMIN_CPU_H__ diff --git a/src/infiniop/ops/blas_amin/info.h b/src/infiniop/ops/blas_amin/info.h deleted file mode 100644 index 0522edecb..000000000 --- a/src/infiniop/ops/blas_amin/info.h +++ /dev/null @@ -1,41 +0,0 @@ -#ifndef __BLAS_AMIN_INFO_H__ -#define __BLAS_AMIN_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" - -class BlasAminInfo { -private: - BlasAminInfo() = default; - -public: - size_t n; - ptrdiff_t incx; - infiniDtype_t data_type; - - static utils::Result createBlasAminInfo( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - - CHECK_OR_RETURN(x_desc != nullptr, INFINI_STATUS_NULL_POINTER); - CHECK_OR_RETURN(result_desc != nullptr, INFINI_STATUS_NULL_POINTER); - - auto data_type = x_desc->dtype(); - auto itype = result_desc->dtype(); - CHECK_DTYPE(data_type, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - CHECK_DTYPE(itype, INFINI_DTYPE_I32); - - CHECK_OR_RETURN(x_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(result_desc->numel() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - - auto n = x_desc->numel(); - auto incx = x_desc->stride(0); - - return utils::Result(BlasAminInfo{ - n, - incx, - data_type}); - } -}; - -#endif // __BLAS_AMIN_INFO_H__ diff --git a/src/infiniop/ops/blas_amin/metax/blas_amin_metax.cc b/src/infiniop/ops/blas_amin/metax/blas_amin_metax.cc deleted file mode 100644 index dadd2706d..000000000 --- a/src/infiniop/ops/blas_amin/metax/blas_amin_metax.cc +++ /dev/null @@ -1,71 +0,0 @@ -#include "blas_amin_metax.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" - -namespace op::blas_amin::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = BlasAminInfo::createBlasAminInfo(x_desc, result_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - new Opaque{handle->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - const int n = utils::cast(_info.n); - const int incx = utils::cast(_info.incx); - const infiniDtype_t data_type = _info.data_type; - - CHECK_STATUS(_opaque->internal->useMcblas( - (hcStream_t)stream, - [&](hcblasHandle_t handle) { - CHECK_MCBLAS(hcblasSetPointerMode(handle, HCBLAS_POINTER_MODE_DEVICE)); - - switch (data_type) { - case INFINI_DTYPE_F32: - CHECK_MCBLAS(hcblasIsamin(handle, n, (const float *)x, incx, (int *)result)); - break; - case INFINI_DTYPE_F64: - CHECK_MCBLAS(hcblasIdamin(handle, n, (const double *)x, incx, (int *)result)); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; - })); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::blas_amin::metax diff --git a/src/infiniop/ops/blas_amin/metax/blas_amin_metax.h b/src/infiniop/ops/blas_amin/metax/blas_amin_metax.h deleted file mode 100644 index 42a5b5fe9..000000000 --- a/src/infiniop/ops/blas_amin/metax/blas_amin_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BLAS_AMIN_METAX_H__ -#define __BLAS_AMIN_METAX_H__ - -#include "../blas_amin.h" - -DESCRIPTOR(metax) - -#endif // __BLAS_AMIN_METAX_H__ diff --git a/src/infiniop/ops/blas_amin/operator.cc b/src/infiniop/ops/blas_amin/operator.cc deleted file mode 100644 index 7f960b773..000000000 --- a/src/infiniop/ops/blas_amin/operator.cc +++ /dev/null @@ -1,122 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/blas_amin.h" - -#ifdef ENABLE_CPU_API -#include "cpu/blas_amin_cpu.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/blas_amin_metax.h" -#endif -#ifdef ENABLE_CAMBRICON_API -#include "bang/blas_amin_bang.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateBlasAminDescriptor( - infiniopHandle_t handle, - infiniopBlasAminDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t result_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::blas_amin::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - x_desc, \ - result_desc) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetBlasAminWorkspaceSize(infiniopBlasAminDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - GET(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET -} - -__INFINI_C infiniStatus_t infiniopBlasAmin( - infiniopBlasAminDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *x, - void *result, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, x, result, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyBlasAminDescriptor(infiniopBlasAminDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - DELETE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/blas_copy/bang/blas_copy_bang.h b/src/infiniop/ops/blas_copy/bang/blas_copy_bang.h deleted file mode 100644 index fb326fb3d..000000000 --- a/src/infiniop/ops/blas_copy/bang/blas_copy_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BLAS_COPY_BANG_H__ -#define __BLAS_COPY_BANG_H__ - -#include "../blas_copy.h" - -DESCRIPTOR(bang) - -#endif // __BLAS_COPY_BANG_H__ diff --git a/src/infiniop/ops/blas_copy/bang/blas_copy_bang.mlu b/src/infiniop/ops/blas_copy/bang/blas_copy_bang.mlu deleted file mode 100644 index fe47729c0..000000000 --- a/src/infiniop/ops/blas_copy/bang/blas_copy_bang.mlu +++ /dev/null @@ -1,96 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "blas_copy_bang.h" -#include "blas_copy_bang_kernel.mlu" - -namespace op::blas_copy::bang { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = BlasCopyInfo::createBlasCopyInfo(x_desc, y_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateBlasCopy( - const BlasCopyInfo &info, - const Tdata *x, - Tdata *y, - cnrtQueue_t queue) { - - const int n = utils::cast(info.n); - const int incx = utils::cast(info.incx); - const int incy = utils::cast(info.incy); - - cnrtDim3_t k_dim; - cnrtFunctionType_t k_type; - - k_dim.x = 4; - k_dim.y = 1; - k_dim.z = 1; - k_type = cnrtFuncTypeUnion1; - - if (incx == 1 && incy == 1) { - blasCopyKernelContiguous<<>>( - n, - x, - y); - } else { - blasCopyKernelStrided<<>>( - n, - x, - incx, - y, - incy); - } - - cnrtQueueSync(queue); - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_BLAS_COPY(TDATA) \ - calculateBlasCopy(_info, \ - (const TDATA *)x, \ - (TDATA *)y, \ - (cnrtQueue_t)stream) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - void *y, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - switch (_info.data_type) { - case INFINI_DTYPE_F16: - return CALCULATE_BLAS_COPY(half); - case INFINI_DTYPE_F32: - return CALCULATE_BLAS_COPY(float); - case INFINI_DTYPE_BF16: - return CALCULATE_BLAS_COPY(bfloat16_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE_BLAS_COPY - -} // namespace op::blas_copy::bang diff --git a/src/infiniop/ops/blas_copy/bang/blas_copy_bang_kernel.mlu b/src/infiniop/ops/blas_copy/bang/blas_copy_bang_kernel.mlu deleted file mode 100644 index 63bc22a66..000000000 --- a/src/infiniop/ops/blas_copy/bang/blas_copy_bang_kernel.mlu +++ /dev/null @@ -1,66 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "blas_copy_bang.h" - -__nram__ char nram_buffer[NRAM_MAX_SIZE]; - -template -__mlu_global__ void blasCopyKernelContiguous( - int n, - const Tdata *x, - Tdata *y) { - - char *nram_aligned = (char *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1)); - - size_t nram_usable = NRAM_MAX_SIZE - (nram_aligned - nram_buffer); - size_t max_chunk_elements = nram_usable / (2 * sizeof(Tdata)); - - size_t align_elements = ALIGN_SIZE / sizeof(Tdata); - if (align_elements == 0) { - align_elements = 1; - } - int chunk_size = (int)((max_chunk_elements / align_elements) * align_elements); - - Tdata *nram_x = (Tdata *)nram_aligned; - - int elements_per_core = n / taskDim; - int remain = n % taskDim; - int core_elements = elements_per_core + (taskId < remain ? 1 : 0); - int core_offset = taskId < remain ? taskId * core_elements : taskId * elements_per_core + remain; - - if (core_elements <= 0) { - return; - } - - int chunks = core_elements / chunk_size; - int chunk_rem = core_elements % chunk_size; - - for (int c = 0; c < chunks; c++) { - int current_offset = core_offset + c * chunk_size; - __memcpy(nram_x, x + current_offset, chunk_size * sizeof(Tdata), GDRAM2NRAM); - __memcpy(y + current_offset, nram_x, chunk_size * sizeof(Tdata), NRAM2GDRAM); - } - - if (chunk_rem > 0) { - int current_offset = core_offset + chunks * chunk_size; - __memcpy(nram_x, x + current_offset, chunk_rem * sizeof(Tdata), GDRAM2NRAM); - __memcpy(y + current_offset, nram_x, chunk_rem * sizeof(Tdata), NRAM2GDRAM); - } -} - -template -__mlu_global__ void blasCopyKernelStrided( - int n, - const Tdata *x, - int incx, - Tdata *y, - int incy) { - - int elements_per_core = n / taskDim; - int remain = n % taskDim; - int actual_tasks = elements_per_core + (taskId < remain ? 1 : 0); - int start_idx = taskId < remain ? taskId * actual_tasks : taskId * elements_per_core + remain; - - for (int i = start_idx; i < start_idx + actual_tasks; ++i) { - y[i * incy] = x[i * incx]; - } -} diff --git a/src/infiniop/ops/blas_copy/blas_copy.h b/src/infiniop/ops/blas_copy/blas_copy.h deleted file mode 100644 index 3670ba204..000000000 --- a/src/infiniop/ops/blas_copy/blas_copy.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef __BLAS_COPY_H__ -#define __BLAS_COPY_H__ - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::blas_copy::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - BlasCopyInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - BlasCopyInfo info, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(std::move(info)), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t x_desc, \ - infiniopTensorDescriptor_t y_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - const void *x, \ - void *y, \ - void *stream) const; \ - }; \ - } - -#endif // __BLAS_COPY_H__ diff --git a/src/infiniop/ops/blas_copy/cpu/blas_copy_cpu.cc b/src/infiniop/ops/blas_copy/cpu/blas_copy_cpu.cc deleted file mode 100644 index 47123d371..000000000 --- a/src/infiniop/ops/blas_copy/cpu/blas_copy_cpu.cc +++ /dev/null @@ -1,77 +0,0 @@ -#include "blas_copy_cpu.h" -#include "../../../devices/cpu/common_cpu.h" - -namespace op::blas_copy::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = BlasCopyInfo::createBlasCopyInfo(x_desc, y_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateBlasCopy( - const BlasCopyInfo &info, - const Tdata *x, - Tdata *y) { - - const size_t n = info.n; - - for (size_t i = 0; i < n; ++i) { - ptrdiff_t x_idx = utils::cast(i) * info.incx; - ptrdiff_t y_idx = utils::cast(i) * info.incy; - y[y_idx] = x[x_idx]; - } - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_BLAS_COPY(TDATA) \ - calculateBlasCopy(_info, \ - (const TDATA *)x, \ - (TDATA *)y) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - void *y, - void *stream) const { - - (void)workspace; - (void)workspace_size; - (void)stream; - - switch (_info.data_type) { - case INFINI_DTYPE_F16: - return CALCULATE_BLAS_COPY(fp16_t); - case INFINI_DTYPE_F32: - return CALCULATE_BLAS_COPY(float); - case INFINI_DTYPE_F64: - return CALCULATE_BLAS_COPY(double); - case INFINI_DTYPE_BF16: - return CALCULATE_BLAS_COPY(bf16_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE_BLAS_COPY - -} // namespace op::blas_copy::cpu diff --git a/src/infiniop/ops/blas_copy/cpu/blas_copy_cpu.h b/src/infiniop/ops/blas_copy/cpu/blas_copy_cpu.h deleted file mode 100644 index 7cfe48752..000000000 --- a/src/infiniop/ops/blas_copy/cpu/blas_copy_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BLAS_COPY_CPU_H__ -#define __BLAS_COPY_CPU_H__ - -#include "../blas_copy.h" - -DESCRIPTOR(cpu) - -#endif // __BLAS_COPY_CPU_H__ diff --git a/src/infiniop/ops/blas_copy/info.h b/src/infiniop/ops/blas_copy/info.h deleted file mode 100644 index 585138ab1..000000000 --- a/src/infiniop/ops/blas_copy/info.h +++ /dev/null @@ -1,45 +0,0 @@ -#ifndef __BLAS_COPY_INFO_H__ -#define __BLAS_COPY_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" - -class BlasCopyInfo { -private: - BlasCopyInfo() = default; - -public: - size_t n; - ptrdiff_t incx; - ptrdiff_t incy; - infiniDtype_t data_type; - - static utils::Result createBlasCopyInfo( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc) { - - CHECK_OR_RETURN(x_desc != nullptr, INFINI_STATUS_NULL_POINTER); - CHECK_OR_RETURN(y_desc != nullptr, INFINI_STATUS_NULL_POINTER); - - auto data_type = x_desc->dtype(); - - CHECK_OR_RETURN(y_desc->dtype() == data_type, INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_DTYPE(data_type, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_OR_RETURN(x_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(y_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(x_desc->numel() == y_desc->numel(), INFINI_STATUS_BAD_TENSOR_SHAPE); - - auto n = x_desc->numel(); - auto incx = x_desc->stride(0); - auto incy = y_desc->stride(0); - - return utils::Result(BlasCopyInfo{ - n, - incx, - incy, - data_type}); - } -}; - -#endif // __BLAS_COPY_INFO_H__ diff --git a/src/infiniop/ops/blas_copy/metax/blas_copy_metax.cc b/src/infiniop/ops/blas_copy/metax/blas_copy_metax.cc deleted file mode 100644 index e51ec95e9..000000000 --- a/src/infiniop/ops/blas_copy/metax/blas_copy_metax.cc +++ /dev/null @@ -1,72 +0,0 @@ -#include "blas_copy_metax.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" - -namespace op::blas_copy::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = BlasCopyInfo::createBlasCopyInfo(x_desc, y_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - new Opaque{handle->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - void *y, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - const int n = utils::cast(_info.n); - const int incx = utils::cast(_info.incx); - const int incy = utils::cast(_info.incy); - const infiniDtype_t data_type = _info.data_type; - - CHECK_STATUS(_opaque->internal->useMcblas( - (hcStream_t)stream, - [&](hcblasHandle_t handle) { - CHECK_MCBLAS(hcblasSetPointerMode(handle, HCBLAS_POINTER_MODE_DEVICE)); - - switch (data_type) { - case INFINI_DTYPE_F32: - CHECK_MCBLAS(hcblasScopy(handle, n, (const float *)x, incx, (float *)y, incy)); - break; - case INFINI_DTYPE_F64: - CHECK_MCBLAS(hcblasDcopy(handle, n, (const double *)x, incx, (double *)y, incy)); - break; - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - - return INFINI_STATUS_SUCCESS; - })); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::blas_copy::metax diff --git a/src/infiniop/ops/blas_copy/metax/blas_copy_metax.h b/src/infiniop/ops/blas_copy/metax/blas_copy_metax.h deleted file mode 100644 index 88f118dbf..000000000 --- a/src/infiniop/ops/blas_copy/metax/blas_copy_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BLAS_COPY_METAX_H__ -#define __BLAS_COPY_METAX_H__ - -#include "../blas_copy.h" - -DESCRIPTOR(metax) - -#endif // __BLAS_COPY_METAX_H__ diff --git a/src/infiniop/ops/blas_copy/operator.cc b/src/infiniop/ops/blas_copy/operator.cc deleted file mode 100644 index 394adc665..000000000 --- a/src/infiniop/ops/blas_copy/operator.cc +++ /dev/null @@ -1,121 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/blas_copy.h" - -#ifdef ENABLE_CPU_API -#include "cpu/blas_copy_cpu.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/blas_copy_metax.h" -#endif -#ifdef ENABLE_CAMBRICON_API -#include "bang/blas_copy_bang.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateBlasCopyDescriptor( - infiniopHandle_t handle, - infiniopBlasCopyDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::blas_copy::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - x_desc, y_desc) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetBlasCopyWorkspaceSize(infiniopBlasCopyDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - GET(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET -} - -__INFINI_C infiniStatus_t infiniopBlasCopy( - infiniopBlasCopyDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *x, - void *y, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, x, y, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyBlasCopyDescriptor(infiniopBlasCopyDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - DELETE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/blas_dot/bang/blas_dot_bang.h b/src/infiniop/ops/blas_dot/bang/blas_dot_bang.h deleted file mode 100644 index 1c2c18019..000000000 --- a/src/infiniop/ops/blas_dot/bang/blas_dot_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BLAS_DOT_BANG_H__ -#define __BLAS_DOT_BANG_H__ - -#include "../blas_dot.h" - -DESCRIPTOR(bang) - -#endif // __BLAS_DOT_BANG_H__ diff --git a/src/infiniop/ops/blas_dot/bang/blas_dot_bang.mlu b/src/infiniop/ops/blas_dot/bang/blas_dot_bang.mlu deleted file mode 100644 index 8d7f1b816..000000000 --- a/src/infiniop/ops/blas_dot/bang/blas_dot_bang.mlu +++ /dev/null @@ -1,103 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "blas_dot_bang.h" -#include "blas_dot_bang_kernel.mlu" - -namespace op::blas_dot::bang { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t result_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = BlasDotInfo::createBlasDotInfo(x_desc, y_desc, result_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateBlasDot( - const BlasDotInfo &info, - const Tdata *x, - const Tdata *y, - Tdata *result, - cnrtQueue_t queue) { - - const int n = utils::cast(info.n); - const int incx = utils::cast(info.incx); - const int incy = utils::cast(info.incy); - - cnrtDim3_t k_dim; - cnrtFunctionType_t k_type; - - k_dim.x = 4; - k_dim.y = 1; - k_dim.z = 1; - k_type = cnrtFuncTypeUnion1; - - if (incx == 1 && incy == 1) { - blasDotKernelContiguous<<>>( - n, - x, - y, - result); - } else { - blasDotKernelStrided<<>>( - n, - x, - incx, - y, - incy, - result); - } - - cnrtQueueSync(queue); - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_BLAS_DOT(TDATA) \ - calculateBlasDot(_info, \ - (const TDATA *)x, \ - (const TDATA *)y, \ - (TDATA *)result, \ - (cnrtQueue_t)stream) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - const void *y, - void *result, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - switch (_info.data_type) { - case INFINI_DTYPE_F16: - return CALCULATE_BLAS_DOT(half); - case INFINI_DTYPE_F32: - return CALCULATE_BLAS_DOT(float); - case INFINI_DTYPE_BF16: - return CALCULATE_BLAS_DOT(bfloat16_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE_BLAS_DOT - -} // namespace op::blas_dot::bang diff --git a/src/infiniop/ops/blas_dot/bang/blas_dot_bang_kernel.mlu b/src/infiniop/ops/blas_dot/bang/blas_dot_bang_kernel.mlu deleted file mode 100644 index bf249101e..000000000 --- a/src/infiniop/ops/blas_dot/bang/blas_dot_bang_kernel.mlu +++ /dev/null @@ -1,157 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "blas_dot_bang.h" - -#include - -__nram__ char nram_buffer[NRAM_MAX_SIZE]; - -template -__mlu_device__ void blasDotToCompute(float *dst, const Tdata *src, int size) { - if constexpr (std::is_same_v) { - __bang_half2float(dst, src, size); - } else if constexpr (std::is_same_v) { - __bang_bfloat162float(dst, src, size); - } else { - __memcpy(dst, src, size * sizeof(float), NRAM2NRAM); - } -} - -template -__mlu_device__ float blasDotToCompute(Tdata value) { - if constexpr (std::is_same_v) { - return __half2float(value); - } else if constexpr (std::is_same_v) { - return __bfloat162float(value); - } else { - return static_cast(value); - } -} - -template -__mlu_device__ void blasDotStoreResult(Tdata *result, Tdata *nram_result, float *nram_compute, float value) { - nram_compute[0] = value; - if constexpr (std::is_same_v) { - __bang_float2half(nram_result, nram_compute, 1); - result[0] = nram_result[0]; - } else if constexpr (std::is_same_v) { - __bang_float2bfloat16(nram_result, nram_compute, 1); - result[0] = nram_result[0]; - } else { - result[0] = nram_compute[0]; - } -} - -template -__mlu_global__ void blasDotKernelContiguous( - int n, - const Tdata *x, - const Tdata *y, - Tdata *result) { - - __mlu_shared__ float shared_partial_sum[4]; - - char *nram_aligned = (char *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1)); - - size_t nram_usable = NRAM_MAX_SIZE - (nram_aligned - nram_buffer); - size_t max_chunk_elements = nram_usable / (2 * sizeof(Tdata) + 2 * sizeof(float)); - - size_t align_elements = ALIGN_SIZE / sizeof(Tdata); - if (align_elements == 0) { - align_elements = 1; - } - int chunk_size = (int)((max_chunk_elements / align_elements) * align_elements); - - Tdata *nram_x = (Tdata *)nram_aligned; - Tdata *nram_y = nram_x + chunk_size; - float *nram_compute_x = (float *)(nram_y + chunk_size); - float *nram_compute_y = nram_compute_x + chunk_size; - - int elements_per_core = n / taskDim; - int remain = n % taskDim; - int core_elements = elements_per_core + (taskId < remain ? 1 : 0); - int core_offset = taskId < remain ? taskId * core_elements : taskId * elements_per_core + remain; - - int chunks = core_elements / chunk_size; - int chunk_rem = core_elements % chunk_size; - - float partial_sum = 0.0f; - - for (int c = 0; c < chunks; c++) { - int current_offset = core_offset + c * chunk_size; - __memcpy(nram_x, x + current_offset, chunk_size * sizeof(Tdata), GDRAM2NRAM); - __memcpy(nram_y, y + current_offset, chunk_size * sizeof(Tdata), GDRAM2NRAM); - - blasDotToCompute(nram_compute_x, nram_x, chunk_size); - blasDotToCompute(nram_compute_y, nram_y, chunk_size); - __bang_mul(nram_compute_x, nram_compute_x, nram_compute_y, chunk_size); - partial_sum += __bang_sum(nram_compute_x, chunk_size); - } - - if (chunk_rem > 0) { - int current_offset = core_offset + chunks * chunk_size; - __memcpy(nram_x, x + current_offset, chunk_rem * sizeof(Tdata), GDRAM2NRAM); - __memcpy(nram_y, y + current_offset, chunk_rem * sizeof(Tdata), GDRAM2NRAM); - - blasDotToCompute(nram_compute_x, nram_x, chunk_rem); - blasDotToCompute(nram_compute_y, nram_y, chunk_rem); - __bang_mul(nram_compute_x, nram_compute_x, nram_compute_y, chunk_rem); - partial_sum += __bang_sum(nram_compute_x, chunk_rem); - } - - shared_partial_sum[coreId] = partial_sum; - - __sync_cluster(); - - if (coreId == 0) { - float cluster_sum = 0.0f; - - for (int i = 0; i < coreDim; i++) { - cluster_sum += shared_partial_sum[i]; - } - - blasDotStoreResult(result, nram_x, nram_compute_x, cluster_sum); - } -} - -template -__mlu_global__ void blasDotKernelStrided( - int n, - const Tdata *x, - int incx, - const Tdata *y, - int incy, - Tdata *result) { - __mlu_shared__ float shared_partial_sum[4]; - - char *nram_aligned = (char *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1)); - - float *nram_compute = (float *)nram_aligned; - Tdata *nram_result = (Tdata *)(nram_compute + 1); - - int elements_per_core = n / taskDim; - int remain = n % taskDim; - int core_elements = elements_per_core + (taskId < remain ? 1 : 0); - int start_idx = taskId < remain ? taskId * core_elements : taskId * elements_per_core + remain; - - float partial_sum = 0.0f; - int x_offset = start_idx * incx; - int y_offset = start_idx * incy; - - for (int i = 0; i < core_elements; ++i) { - partial_sum += blasDotToCompute(x[x_offset]) * blasDotToCompute(y[y_offset]); - x_offset += incx; - y_offset += incy; - } - - shared_partial_sum[coreId] = partial_sum; - - __sync_cluster(); - - if (coreId == 0) { - float cluster_sum = 0.0f; - for (int i = 0; i < coreDim; ++i) { - cluster_sum += shared_partial_sum[i]; - } - blasDotStoreResult(result, nram_result, nram_compute, cluster_sum); - } -} diff --git a/src/infiniop/ops/blas_dot/blas_dot.h b/src/infiniop/ops/blas_dot/blas_dot.h deleted file mode 100644 index 09e81a7fc..000000000 --- a/src/infiniop/ops/blas_dot/blas_dot.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef __BLAS_DOT_H__ -#define __BLAS_DOT_H__ - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::blas_dot::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - BlasDotInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - BlasDotInfo info, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(std::move(info)), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t x_desc, \ - infiniopTensorDescriptor_t y_desc, \ - infiniopTensorDescriptor_t result_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - const void *x, \ - const void *y, \ - void *result, \ - void *stream) const; \ - }; \ - } - -#endif // __BLAS_DOT_H__ diff --git a/src/infiniop/ops/blas_dot/cpu/blas_dot_cpu.cc b/src/infiniop/ops/blas_dot/cpu/blas_dot_cpu.cc deleted file mode 100644 index e250aa2ee..000000000 --- a/src/infiniop/ops/blas_dot/cpu/blas_dot_cpu.cc +++ /dev/null @@ -1,102 +0,0 @@ -#include "blas_dot_cpu.h" -#include "../../../devices/cpu/common_cpu.h" - -namespace op::blas_dot::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t result_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = BlasDotInfo::createBlasDotInfo(x_desc, y_desc, result_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - nullptr, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateBlasDot( - const BlasDotInfo &info, - const Tdata *x, - const Tdata *y, - Tdata *result) { - - const size_t n = info.n; - const ptrdiff_t incx = info.incx; - const ptrdiff_t incy = info.incy; - - ptrdiff_t ix = (incx < 0) ? (1 - utils::cast(n)) * incx : 0; - ptrdiff_t iy = (incy < 0) ? (1 - utils::cast(n)) * incy : 0; - - if constexpr (std::is_same::value || std::is_same::value) { - float total = 0.0f; - - for (size_t i = 0; i < n; ++i) { - total += utils::cast(x[ix]) * utils::cast(y[iy]); - ix += incx; - iy += incy; - } - - result[0] = utils::cast(total); - } else { - Tdata total = utils::cast(0); - - for (size_t i = 0; i < n; ++i) { - total += x[ix] * y[iy]; - ix += incx; - iy += incy; - } - - result[0] = total; - } - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_BLAS_DOT(TDATA) \ - calculateBlasDot(_info, \ - (const TDATA *)x, \ - (const TDATA *)y, \ - (TDATA *)result) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - const void *y, - void *result, - void *stream) const { - - (void)workspace; - (void)workspace_size; - (void)stream; - - switch (_info.data_type) { - case INFINI_DTYPE_F16: - return CALCULATE_BLAS_DOT(fp16_t); - case INFINI_DTYPE_F32: - return CALCULATE_BLAS_DOT(float); - case INFINI_DTYPE_F64: - return CALCULATE_BLAS_DOT(double); - case INFINI_DTYPE_BF16: - return CALCULATE_BLAS_DOT(bf16_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -#undef CALCULATE_BLAS_DOT - -} // namespace op::blas_dot::cpu diff --git a/src/infiniop/ops/blas_dot/cpu/blas_dot_cpu.h b/src/infiniop/ops/blas_dot/cpu/blas_dot_cpu.h deleted file mode 100644 index 0f09f8d08..000000000 --- a/src/infiniop/ops/blas_dot/cpu/blas_dot_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BLAS_DOT_CPU_H__ -#define __BLAS_DOT_CPU_H__ - -#include "../blas_dot.h" - -DESCRIPTOR(cpu) - -#endif // __BLAS_DOT_CPU_H__ diff --git a/src/infiniop/ops/blas_dot/info.h b/src/infiniop/ops/blas_dot/info.h deleted file mode 100644 index 01e145f6e..000000000 --- a/src/infiniop/ops/blas_dot/info.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef __BLAS_DOT_INFO_H__ -#define __BLAS_DOT_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" - -class BlasDotInfo { -private: - BlasDotInfo() = default; - -public: - size_t n; - ptrdiff_t incx; - ptrdiff_t incy; - infiniDtype_t data_type; - - static utils::Result createBlasDotInfo( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t result_desc) { - - CHECK_OR_RETURN(x_desc != nullptr, INFINI_STATUS_NULL_POINTER); - CHECK_OR_RETURN(y_desc != nullptr, INFINI_STATUS_NULL_POINTER); - CHECK_OR_RETURN(result_desc != nullptr, INFINI_STATUS_NULL_POINTER); - - auto data_type = x_desc->dtype(); - - CHECK_OR_RETURN(y_desc->dtype() == data_type, INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(result_desc->dtype() == data_type, INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_DTYPE(data_type, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_OR_RETURN(x_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(y_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(x_desc->numel() == y_desc->numel(), INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(result_desc->numel() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - - auto n = x_desc->numel(); - auto incx = x_desc->stride(0); - auto incy = y_desc->stride(0); - - return utils::Result(BlasDotInfo{ - n, - incx, - incy, - data_type}); - } -}; - -#endif // __BLAS_DOT_INFO_H__ diff --git a/src/infiniop/ops/blas_dot/metax/blas_dot_metax.cc b/src/infiniop/ops/blas_dot/metax/blas_dot_metax.cc deleted file mode 100644 index 28e7f301a..000000000 --- a/src/infiniop/ops/blas_dot/metax/blas_dot_metax.cc +++ /dev/null @@ -1,102 +0,0 @@ -#include "blas_dot_metax.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" - -namespace op::blas_dot::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t result_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = BlasDotInfo::createBlasDotInfo(x_desc, y_desc, result_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - result.take(), - 0, - new Opaque{handle->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - const void *x, - const void *y, - void *result, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - const int n = utils::cast(_info.n); - const int incx = utils::cast(_info.incx); - const int incy = utils::cast(_info.incy); - const infiniDtype_t data_type = _info.data_type; - - hpccDataType x_type, y_type, result_type; - hpccDataType execution_type; - - switch (data_type) { - case INFINI_DTYPE_F16: - x_type = y_type = result_type = HPCC_R_16F; - execution_type = HPCC_R_32F; - break; - case INFINI_DTYPE_BF16: - x_type = y_type = result_type = HPCC_R_16BF; - execution_type = HPCC_R_32F; - break; - case INFINI_DTYPE_F32: - x_type = y_type = result_type = HPCC_R_32F; - execution_type = HPCC_R_32F; - break; - case INFINI_DTYPE_F64: - x_type = y_type = result_type = HPCC_R_64F; - execution_type = HPCC_R_64F; - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - CHECK_STATUS(_opaque->internal->useMcblas( - (hcStream_t)stream, - [&](hcblasHandle_t handle) { - CHECK_MCBLAS(hcblasSetPointerMode( - handle, - HCBLAS_POINTER_MODE_DEVICE)); - - CHECK_MCBLAS(hcblasDotEx( - handle, - n, - x, - x_type, - incx, - y, - y_type, - incy, - result, - result_type, - execution_type)); - - return INFINI_STATUS_SUCCESS; - })); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::blas_dot::metax diff --git a/src/infiniop/ops/blas_dot/metax/blas_dot_metax.h b/src/infiniop/ops/blas_dot/metax/blas_dot_metax.h deleted file mode 100644 index 0c5cefbf8..000000000 --- a/src/infiniop/ops/blas_dot/metax/blas_dot_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BLAS_DOT_METAX_H__ -#define __BLAS_DOT_METAX_H__ - -#include "../blas_dot.h" - -DESCRIPTOR(metax) - -#endif // __BLAS_DOT_METAX_H__ diff --git a/src/infiniop/ops/blas_dot/operator.cc b/src/infiniop/ops/blas_dot/operator.cc deleted file mode 100644 index e28943b8b..000000000 --- a/src/infiniop/ops/blas_dot/operator.cc +++ /dev/null @@ -1,125 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/blas_dot.h" - -#ifdef ENABLE_CPU_API -#include "cpu/blas_dot_cpu.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/blas_dot_metax.h" -#endif -#ifdef ENABLE_CAMBRICON_API -#include "bang/blas_dot_bang.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateBlasDotDescriptor( - infiniopHandle_t handle, - infiniopBlasDotDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t result_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::blas_dot::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - x_desc, \ - y_desc, \ - result_desc) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetBlasDotWorkspaceSize(infiniopBlasDotDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - GET(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET -} - -__INFINI_C infiniStatus_t infiniopBlasDot( - infiniopBlasDotDescriptor_t desc, - void *workspace, - size_t workspace_size, - const void *x, - const void *y, - void *result, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, x, y, result, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyBlasDotDescriptor(infiniopBlasDotDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_CAMBRICON_API - DELETE(INFINI_DEVICE_CAMBRICON, bang); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/block_diag/cpu/block_diag_cpu.cc b/src/infiniop/ops/block_diag/cpu/block_diag_cpu.cc deleted file mode 100644 index 4ae8de3ce..000000000 --- a/src/infiniop/ops/block_diag/cpu/block_diag_cpu.cc +++ /dev/null @@ -1,173 +0,0 @@ -#include "block_diag_cpu.h" -#include "../../../../utils.h" -#include "../../../tensor.h" -#include - -namespace op::block_diag::cpu { - -utils::Result BlockDiagInfo::create( - infiniopTensorDescriptor_t *input_descs, - size_t num_inputs, - infiniopTensorDescriptor_t y_desc) { - - if (num_inputs == 0) { - return INFINI_STATUS_BAD_PARAM; - } - - BlockDiagInfo info; - info.num_inputs = num_inputs; - info.input_shapes.resize(num_inputs); - info.input_stride0.resize(num_inputs); - info.input_stride1.resize(num_inputs); - info.row_offsets.resize(num_inputs); - info.col_offsets.resize(num_inputs); - - size_t total_rows = 0; - size_t total_cols = 0; - - // Process each input matrix - for (size_t i = 0; i < num_inputs; ++i) { - auto shape = input_descs[i]->shape(); - if (shape.size() != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - info.input_shapes[i] = shape; - info.input_stride0[i] = input_descs[i]->stride(0); - info.input_stride1[i] = input_descs[i]->stride(1); - info.row_offsets[i] = total_rows; - info.col_offsets[i] = total_cols; - total_rows += shape[0]; - total_cols += shape[1]; - } - - // Check output shape - auto y_shape = y_desc->shape(); - if (y_shape.size() != 2 || y_shape[0] != total_rows || y_shape[1] != total_cols) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - info.output_shape = y_shape; - auto y_strides = y_desc->strides(); - if (y_strides.size() != 2) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - info.output_stride0 = y_strides[0]; - info.output_stride1 = y_strides[1]; - info.output_size = y_desc->numel(); - - return utils::Result(std::move(info)); -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t *input_descs, - size_t num_inputs) { - - if (num_inputs == 0) { - return INFINI_STATUS_BAD_PARAM; - } - - auto dtype = input_descs[0]->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - if (y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // Check all inputs have same dtype - for (size_t i = 1; i < num_inputs; ++i) { - if (input_descs[i]->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - auto info_result = BlockDiagInfo::create(input_descs, num_inputs, y_desc); - CHECK_RESULT(info_result); - - *desc_ptr = new Descriptor(dtype, info_result.take(), handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -void block_diag_impl( - const BlockDiagInfo &info, - T *y, - const T **inputs) { - - const size_t out_rows = info.output_shape[0]; - const size_t out_cols = info.output_shape[1]; - const ptrdiff_t out_s0 = info.output_stride0; - const ptrdiff_t out_s1 = info.output_stride1; - - // Initialize output to zero (stride-aware). - for (size_t r = 0; r < out_rows; ++r) { - for (size_t c = 0; c < out_cols; ++c) { - ptrdiff_t out_off = static_cast(r) * out_s0 + static_cast(c) * out_s1; - y[out_off] = utils::cast(0.0f); - } - } - - // Place each input matrix at its diagonal position - for (size_t i = 0; i < info.num_inputs; ++i) { - size_t rows = info.input_shapes[i][0]; - size_t cols = info.input_shapes[i][1]; - size_t row_offset = info.row_offsets[i]; - size_t col_offset = info.col_offsets[i]; - const T *input = reinterpret_cast(inputs[i]); - const ptrdiff_t in_s0 = info.input_stride0[i]; - const ptrdiff_t in_s1 = info.input_stride1[i]; - - // Copy input matrix to output at diagonal position - for (size_t r = 0; r < rows; ++r) { - for (size_t c = 0; c < cols; ++c) { - size_t out_row = row_offset + r; - size_t out_col = col_offset + c; - - ptrdiff_t out_off = static_cast(out_row) * out_s0 + static_cast(out_col) * out_s1; - ptrdiff_t in_off = static_cast(r) * in_s0 + static_cast(c) * in_s1; - y[out_off] = input[in_off]; - } - } - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void **inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: { - const fp16_t **typed_inputs = reinterpret_cast(inputs); - block_diag_impl(_info, reinterpret_cast(y), typed_inputs); - break; - } - case INFINI_DTYPE_BF16: { - const bf16_t **typed_inputs = reinterpret_cast(inputs); - block_diag_impl(_info, reinterpret_cast(y), typed_inputs); - break; - } - case INFINI_DTYPE_F32: { - const float **typed_inputs = reinterpret_cast(inputs); - block_diag_impl(_info, reinterpret_cast(y), typed_inputs); - break; - } - case INFINI_DTYPE_F64: { - const double **typed_inputs = reinterpret_cast(inputs); - block_diag_impl(_info, reinterpret_cast(y), typed_inputs); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::block_diag::cpu diff --git a/src/infiniop/ops/block_diag/cpu/block_diag_cpu.h b/src/infiniop/ops/block_diag/cpu/block_diag_cpu.h deleted file mode 100644 index a7dc073b3..000000000 --- a/src/infiniop/ops/block_diag/cpu/block_diag_cpu.h +++ /dev/null @@ -1,60 +0,0 @@ -#ifndef __BLOCK_DIAG_CPU_H__ -#define __BLOCK_DIAG_CPU_H__ - -#include "../../../devices/cpu/common_cpu.h" -#include "../../../operator.h" -#include - -namespace op::block_diag::cpu { - -struct BlockDiagInfo { - size_t num_inputs; - std::vector> input_shapes; // Each input is 2D matrix - std::vector input_stride0; // stride(0) per input - std::vector input_stride1; // stride(1) per input - std::vector output_shape; // 2D output - ptrdiff_t output_stride0; - ptrdiff_t output_stride1; - std::vector row_offsets; // Row offset for each input matrix - std::vector col_offsets; // Column offset for each input matrix - size_t output_size; - - static utils::Result create( - infiniopTensorDescriptor_t *input_descs, - size_t num_inputs, - infiniopTensorDescriptor_t y_desc); -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - BlockDiagInfo _info; - - Descriptor(infiniDtype_t dtype, BlockDiagInfo info, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t *input_descs, - size_t num_inputs); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void **inputs, - void *stream) const; -}; - -} // namespace op::block_diag::cpu - -#endif // __BLOCK_DIAG_CPU_H__ diff --git a/src/infiniop/ops/block_diag/cuda/kernel.cuh b/src/infiniop/ops/block_diag/cuda/kernel.cuh deleted file mode 100644 index 34ce6adaf..000000000 --- a/src/infiniop/ops/block_diag/cuda/kernel.cuh +++ /dev/null @@ -1,55 +0,0 @@ -#include -#include - -namespace op::cuda { - -template -__global__ void block_diag_kernel( - T *output, - const T **inputs, - size_t num_inputs, - size_t output_rows, - size_t output_cols, - ptrdiff_t output_stride0, - ptrdiff_t output_stride1, - const size_t *row_offsets, - const size_t *col_offsets, - const size_t *input_rows, - const size_t *input_cols, - const ptrdiff_t *input_stride0, - const ptrdiff_t *input_stride1) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - size_t total = output_rows * output_cols; - - if (idx >= total) { - return; - } - - size_t out_row = idx / output_cols; - size_t out_col = idx % output_cols; - - const ptrdiff_t out_offset = static_cast(out_row) * output_stride0 + static_cast(out_col) * output_stride1; - - T value = T{}; - - // Find which input matrix this output position belongs to - for (size_t i = 0; i < num_inputs; ++i) { - size_t row_start = row_offsets[i]; - size_t row_end = row_start + input_rows[i]; - size_t col_start = col_offsets[i]; - size_t col_end = col_start + input_cols[i]; - - if (out_row >= row_start && out_row < row_end && out_col >= col_start && out_col < col_end) { - // This position belongs to input i - size_t in_row = out_row - row_start; - size_t in_col = out_col - col_start; - const ptrdiff_t in_offset = static_cast(in_row) * input_stride0[i] + static_cast(in_col) * input_stride1[i]; - value = inputs[i][in_offset]; - break; - } - } - output[out_offset] = value; -} - -} // namespace op::cuda diff --git a/src/infiniop/ops/block_diag/metax/block_diag_metax.h b/src/infiniop/ops/block_diag/metax/block_diag_metax.h deleted file mode 100644 index 01c51e6a8..000000000 --- a/src/infiniop/ops/block_diag/metax/block_diag_metax.h +++ /dev/null @@ -1,77 +0,0 @@ -#ifndef __BLOCK_DIAG_METAX_H__ -#define __BLOCK_DIAG_METAX_H__ - -#include "../../../operator.h" -#include -#include - -namespace op::block_diag::metax { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t num_inputs; - std::vector output_shape; - ptrdiff_t output_stride0; - ptrdiff_t output_stride1; - std::vector row_offsets; - std::vector col_offsets; - std::vector input_rows; - std::vector input_cols; - std::vector input_stride0; - std::vector input_stride1; - size_t output_size; - - Descriptor(infiniDtype_t dtype, size_t num_inputs, - std::vector output_shape, - ptrdiff_t output_stride0, - ptrdiff_t output_stride1, - std::vector row_offsets, - std::vector col_offsets, - std::vector input_rows, - std::vector input_cols, - std::vector input_stride0, - std::vector input_stride1, - size_t output_size, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - num_inputs(num_inputs), - output_shape(std::move(output_shape)), - output_stride0(output_stride0), - output_stride1(output_stride1), - row_offsets(std::move(row_offsets)), - col_offsets(std::move(col_offsets)), - input_rows(std::move(input_rows)), - input_cols(std::move(input_cols)), - input_stride0(std::move(input_stride0)), - input_stride1(std::move(input_stride1)), - output_size(output_size) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t *input_descs, - size_t num_inputs); - - size_t workspaceSize() const { - if (output_size == 0) { - return 0; - } - return 4 * num_inputs * sizeof(size_t) + 2 * num_inputs * sizeof(ptrdiff_t) + num_inputs * sizeof(void *); - } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void **inputs, - void *stream) const; -}; - -} // namespace op::block_diag::metax - -#endif // __BLOCK_DIAG_METAX_H__ diff --git a/src/infiniop/ops/block_diag/metax/block_diag_metax.maca b/src/infiniop/ops/block_diag/metax/block_diag_metax.maca deleted file mode 100644 index 39a0a3a29..000000000 --- a/src/infiniop/ops/block_diag/metax/block_diag_metax.maca +++ /dev/null @@ -1,168 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" - -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "block_diag_metax.h" - -namespace op::block_diag::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t *input_descs, - size_t num_inputs) { - - if (num_inputs == 0) { - return INFINI_STATUS_BAD_PARAM; - } - - auto dtype = input_descs[0]->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - if (y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - for (size_t i = 1; i < num_inputs; ++i) { - if (input_descs[i]->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - std::vector row_offsets(num_inputs); - std::vector col_offsets(num_inputs); - std::vector input_rows(num_inputs); - std::vector input_cols(num_inputs); - std::vector input_stride0(num_inputs); - std::vector input_stride1(num_inputs); - - size_t total_rows = 0; - size_t total_cols = 0; - - for (size_t i = 0; i < num_inputs; ++i) { - auto shape = input_descs[i]->shape(); - if (shape.size() != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - input_rows[i] = shape[0]; - input_cols[i] = shape[1]; - input_stride0[i] = input_descs[i]->stride(0); - input_stride1[i] = input_descs[i]->stride(1); - row_offsets[i] = total_rows; - col_offsets[i] = total_cols; - total_rows += shape[0]; - total_cols += shape[1]; - } - - auto y_shape = y_desc->shape(); - if (y_shape.size() != 2 || y_shape[0] != total_rows || y_shape[1] != total_cols) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - auto y_strides = y_desc->strides(); - if (y_strides.size() != 2) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - *desc_ptr = new Descriptor(dtype, num_inputs, y_shape, - y_strides[0], y_strides[1], - row_offsets, col_offsets, - std::move(input_rows), std::move(input_cols), - std::move(input_stride0), std::move(input_stride1), - y_desc->numel(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void **inputs, - void *stream) const { - - auto hc_stream = reinterpret_cast(stream); - - size_t output_bytes = output_size * infiniSizeOf(_dtype); - CHECK_METAX(hcMemsetAsync(y, 0, output_bytes, hc_stream)); - if (output_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - if (workspace_size < this->workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - size_t *d_row_offsets = reinterpret_cast(workspace); - size_t *d_col_offsets = d_row_offsets + num_inputs; - size_t *d_input_rows = d_col_offsets + num_inputs; - size_t *d_input_cols = d_input_rows + num_inputs; - ptrdiff_t *d_input_stride0 = reinterpret_cast(d_input_cols + num_inputs); - ptrdiff_t *d_input_stride1 = d_input_stride0 + num_inputs; - const void **d_inputs = reinterpret_cast(d_input_stride1 + num_inputs); - - CHECK_METAX(hcMemcpyAsync(d_row_offsets, row_offsets.data(), num_inputs * sizeof(size_t), hcMemcpyHostToDevice, hc_stream)); - CHECK_METAX(hcMemcpyAsync(d_col_offsets, col_offsets.data(), num_inputs * sizeof(size_t), hcMemcpyHostToDevice, hc_stream)); - CHECK_METAX(hcMemcpyAsync(d_input_rows, input_rows.data(), num_inputs * sizeof(size_t), hcMemcpyHostToDevice, hc_stream)); - CHECK_METAX(hcMemcpyAsync(d_input_cols, input_cols.data(), num_inputs * sizeof(size_t), hcMemcpyHostToDevice, hc_stream)); - CHECK_METAX(hcMemcpyAsync(d_input_stride0, input_stride0.data(), num_inputs * sizeof(ptrdiff_t), hcMemcpyHostToDevice, hc_stream)); - CHECK_METAX(hcMemcpyAsync(d_input_stride1, input_stride1.data(), num_inputs * sizeof(ptrdiff_t), hcMemcpyHostToDevice, hc_stream)); - CHECK_METAX(hcMemcpyAsync(d_inputs, inputs, num_inputs * sizeof(void *), hcMemcpyHostToDevice, hc_stream)); - - constexpr int BLOCK_SIZE = 256; - int num_blocks = (output_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - - switch (_dtype) { - case INFINI_DTYPE_F16: - cuda::block_diag_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(d_inputs), - num_inputs, - output_shape[0], output_shape[1], - output_stride0, output_stride1, - d_row_offsets, d_col_offsets, d_input_rows, d_input_cols, - d_input_stride0, d_input_stride1); - break; - case INFINI_DTYPE_BF16: - cuda::block_diag_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(d_inputs), - num_inputs, - output_shape[0], output_shape[1], - output_stride0, output_stride1, - d_row_offsets, d_col_offsets, d_input_rows, d_input_cols, - d_input_stride0, d_input_stride1); - break; - case INFINI_DTYPE_F32: - cuda::block_diag_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(d_inputs), - num_inputs, - output_shape[0], output_shape[1], - output_stride0, output_stride1, - d_row_offsets, d_col_offsets, d_input_rows, d_input_cols, - d_input_stride0, d_input_stride1); - break; - case INFINI_DTYPE_F64: - cuda::block_diag_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(d_inputs), - num_inputs, - output_shape[0], output_shape[1], - output_stride0, output_stride1, - d_row_offsets, d_col_offsets, d_input_rows, d_input_cols, - d_input_stride0, d_input_stride1); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::block_diag::metax diff --git a/src/infiniop/ops/block_diag/moore/block_diag_moore.h b/src/infiniop/ops/block_diag/moore/block_diag_moore.h deleted file mode 100644 index 7e26532f4..000000000 --- a/src/infiniop/ops/block_diag/moore/block_diag_moore.h +++ /dev/null @@ -1,77 +0,0 @@ -#ifndef __BLOCK_DIAG_MOORE_H__ -#define __BLOCK_DIAG_MOORE_H__ - -#include "../../../operator.h" -#include -#include - -namespace op::block_diag::moore { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t num_inputs; - std::vector output_shape; - ptrdiff_t output_stride0; - ptrdiff_t output_stride1; - std::vector row_offsets; - std::vector col_offsets; - std::vector input_rows; - std::vector input_cols; - std::vector input_stride0; - std::vector input_stride1; - size_t output_size; - - Descriptor(infiniDtype_t dtype, size_t num_inputs, - std::vector output_shape, - ptrdiff_t output_stride0, - ptrdiff_t output_stride1, - std::vector row_offsets, - std::vector col_offsets, - std::vector input_rows, - std::vector input_cols, - std::vector input_stride0, - std::vector input_stride1, - size_t output_size, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - num_inputs(num_inputs), - output_shape(std::move(output_shape)), - output_stride0(output_stride0), - output_stride1(output_stride1), - row_offsets(std::move(row_offsets)), - col_offsets(std::move(col_offsets)), - input_rows(std::move(input_rows)), - input_cols(std::move(input_cols)), - input_stride0(std::move(input_stride0)), - input_stride1(std::move(input_stride1)), - output_size(output_size) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t *input_descs, - size_t num_inputs); - - size_t workspaceSize() const { - if (output_size == 0) { - return 0; - } - return 4 * num_inputs * sizeof(size_t) + 2 * num_inputs * sizeof(ptrdiff_t) + num_inputs * sizeof(void *); - } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void **inputs, - void *stream) const; -}; - -} // namespace op::block_diag::moore - -#endif // __BLOCK_DIAG_MOORE_H__ diff --git a/src/infiniop/ops/block_diag/moore/block_diag_moore.mu b/src/infiniop/ops/block_diag/moore/block_diag_moore.mu deleted file mode 100644 index 4cec566a0..000000000 --- a/src/infiniop/ops/block_diag/moore/block_diag_moore.mu +++ /dev/null @@ -1,167 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "block_diag_moore.h" - -namespace op::block_diag::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t *input_descs, - size_t num_inputs) { - - if (num_inputs == 0) { - return INFINI_STATUS_BAD_PARAM; - } - - auto dtype = input_descs[0]->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - if (y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - for (size_t i = 1; i < num_inputs; ++i) { - if (input_descs[i]->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - std::vector row_offsets(num_inputs); - std::vector col_offsets(num_inputs); - std::vector input_rows(num_inputs); - std::vector input_cols(num_inputs); - std::vector input_stride0(num_inputs); - std::vector input_stride1(num_inputs); - - size_t total_rows = 0; - size_t total_cols = 0; - - for (size_t i = 0; i < num_inputs; ++i) { - auto shape = input_descs[i]->shape(); - if (shape.size() != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - input_rows[i] = shape[0]; - input_cols[i] = shape[1]; - input_stride0[i] = input_descs[i]->stride(0); - input_stride1[i] = input_descs[i]->stride(1); - row_offsets[i] = total_rows; - col_offsets[i] = total_cols; - total_rows += shape[0]; - total_cols += shape[1]; - } - - auto y_shape = y_desc->shape(); - if (y_shape.size() != 2 || y_shape[0] != total_rows || y_shape[1] != total_cols) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - auto y_strides = y_desc->strides(); - if (y_strides.size() != 2) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - *desc_ptr = new Descriptor(dtype, num_inputs, y_shape, - y_strides[0], y_strides[1], - row_offsets, col_offsets, - std::move(input_rows), std::move(input_cols), - std::move(input_stride0), std::move(input_stride1), - y_desc->numel(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void **inputs, - void *stream) const { - - auto musa_stream = reinterpret_cast(stream); - - size_t output_bytes = output_size * infiniSizeOf(_dtype); - CHECK_MOORE(musaMemsetAsync(y, 0, output_bytes, musa_stream)); - if (output_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - if (workspace_size < this->workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - size_t *d_row_offsets = reinterpret_cast(workspace); - size_t *d_col_offsets = d_row_offsets + num_inputs; - size_t *d_input_rows = d_col_offsets + num_inputs; - size_t *d_input_cols = d_input_rows + num_inputs; - ptrdiff_t *d_input_stride0 = reinterpret_cast(d_input_cols + num_inputs); - ptrdiff_t *d_input_stride1 = d_input_stride0 + num_inputs; - const void **d_inputs = reinterpret_cast(d_input_stride1 + num_inputs); - - CHECK_MOORE(musaMemcpyAsync(d_row_offsets, row_offsets.data(), num_inputs * sizeof(size_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(d_col_offsets, col_offsets.data(), num_inputs * sizeof(size_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(d_input_rows, input_rows.data(), num_inputs * sizeof(size_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(d_input_cols, input_cols.data(), num_inputs * sizeof(size_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(d_input_stride0, input_stride0.data(), num_inputs * sizeof(ptrdiff_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(d_input_stride1, input_stride1.data(), num_inputs * sizeof(ptrdiff_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(d_inputs, inputs, num_inputs * sizeof(void *), musaMemcpyHostToDevice, musa_stream)); - - constexpr int BLOCK_SIZE = 256; - int num_blocks = (output_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - - switch (_dtype) { - case INFINI_DTYPE_F16: - cuda::block_diag_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(d_inputs), - num_inputs, - output_shape[0], output_shape[1], - output_stride0, output_stride1, - d_row_offsets, d_col_offsets, d_input_rows, d_input_cols, - d_input_stride0, d_input_stride1); - break; - case INFINI_DTYPE_BF16: - cuda::block_diag_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(d_inputs), - num_inputs, - output_shape[0], output_shape[1], - output_stride0, output_stride1, - d_row_offsets, d_col_offsets, d_input_rows, d_input_cols, - d_input_stride0, d_input_stride1); - break; - case INFINI_DTYPE_F32: - cuda::block_diag_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(d_inputs), - num_inputs, - output_shape[0], output_shape[1], - output_stride0, output_stride1, - d_row_offsets, d_col_offsets, d_input_rows, d_input_cols, - d_input_stride0, d_input_stride1); - break; - case INFINI_DTYPE_F64: - cuda::block_diag_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(d_inputs), - num_inputs, - output_shape[0], output_shape[1], - output_stride0, output_stride1, - d_row_offsets, d_col_offsets, d_input_rows, d_input_cols, - d_input_stride0, d_input_stride1); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::block_diag::moore diff --git a/src/infiniop/ops/block_diag/nvidia/block_diag_nvidia.cu b/src/infiniop/ops/block_diag/nvidia/block_diag_nvidia.cu deleted file mode 100644 index b1d57cd7d..000000000 --- a/src/infiniop/ops/block_diag/nvidia/block_diag_nvidia.cu +++ /dev/null @@ -1,166 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "block_diag_nvidia.cuh" - -namespace op::block_diag::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t *input_descs, - size_t num_inputs) { - - if (num_inputs == 0) { - return INFINI_STATUS_BAD_PARAM; - } - - auto dtype = input_descs[0]->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - if (y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - for (size_t i = 1; i < num_inputs; ++i) { - if (input_descs[i]->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - std::vector row_offsets(num_inputs); - std::vector col_offsets(num_inputs); - std::vector> input_shapes(num_inputs); - std::vector input_rows(num_inputs); - std::vector input_cols(num_inputs); - std::vector input_stride0(num_inputs); - std::vector input_stride1(num_inputs); - - size_t total_rows = 0; - size_t total_cols = 0; - - for (size_t i = 0; i < num_inputs; ++i) { - auto shape = input_descs[i]->shape(); - if (shape.size() != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - input_shapes[i] = shape; - input_rows[i] = shape[0]; - input_cols[i] = shape[1]; - input_stride0[i] = input_descs[i]->stride(0); - input_stride1[i] = input_descs[i]->stride(1); - row_offsets[i] = total_rows; - col_offsets[i] = total_cols; - total_rows += shape[0]; - total_cols += shape[1]; - } - - auto y_shape = y_desc->shape(); - if (y_shape.size() != 2 || y_shape[0] != total_rows || y_shape[1] != total_cols) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - auto y_strides = y_desc->strides(); - if (y_strides.size() != 2) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - *desc_ptr = new Descriptor(dtype, num_inputs, y_shape, - y_strides[0], y_strides[1], - row_offsets, col_offsets, input_shapes, - input_rows, input_cols, - input_stride0, input_stride1, - y_desc->numel(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void **inputs, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - - if (workspace_size < this->workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - size_t *d_row_offsets = reinterpret_cast(workspace); - size_t *d_col_offsets = d_row_offsets + num_inputs; - size_t *d_input_rows = d_col_offsets + num_inputs; - size_t *d_input_cols = d_input_rows + num_inputs; - ptrdiff_t *d_input_stride0 = reinterpret_cast(d_input_cols + num_inputs); - ptrdiff_t *d_input_stride1 = d_input_stride0 + num_inputs; - const void **d_inputs = reinterpret_cast(d_input_stride1 + num_inputs); - - CHECK_CUDA(cudaMemcpyAsync(d_row_offsets, row_offsets.data(), num_inputs * sizeof(size_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(d_col_offsets, col_offsets.data(), num_inputs * sizeof(size_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(d_input_rows, input_rows.data(), num_inputs * sizeof(size_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(d_input_cols, input_cols.data(), num_inputs * sizeof(size_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(d_input_stride0, input_stride0.data(), num_inputs * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(d_input_stride1, input_stride1.data(), num_inputs * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(d_inputs, inputs, num_inputs * sizeof(void *), cudaMemcpyHostToDevice, cuda_stream)); - - constexpr int BLOCK_SIZE = 256; - int num_blocks = (output_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - if (num_blocks < 1) { - num_blocks = 1; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - cuda::block_diag_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(d_inputs), - num_inputs, - output_shape[0], output_shape[1], - output_stride0, output_stride1, - d_row_offsets, d_col_offsets, d_input_rows, d_input_cols, - d_input_stride0, d_input_stride1); - break; - case INFINI_DTYPE_BF16: - cuda::block_diag_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(d_inputs), - num_inputs, - output_shape[0], output_shape[1], - output_stride0, output_stride1, - d_row_offsets, d_col_offsets, d_input_rows, d_input_cols, - d_input_stride0, d_input_stride1); - break; - case INFINI_DTYPE_F32: - cuda::block_diag_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(d_inputs), - num_inputs, - output_shape[0], output_shape[1], - output_stride0, output_stride1, - d_row_offsets, d_col_offsets, d_input_rows, d_input_cols, - d_input_stride0, d_input_stride1); - break; - case INFINI_DTYPE_F64: - cuda::block_diag_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(d_inputs), - num_inputs, - output_shape[0], output_shape[1], - output_stride0, output_stride1, - d_row_offsets, d_col_offsets, d_input_rows, d_input_cols, - d_input_stride0, d_input_stride1); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::block_diag::nvidia diff --git a/src/infiniop/ops/block_diag/nvidia/block_diag_nvidia.cuh b/src/infiniop/ops/block_diag/nvidia/block_diag_nvidia.cuh deleted file mode 100644 index 10fb3ec47..000000000 --- a/src/infiniop/ops/block_diag/nvidia/block_diag_nvidia.cuh +++ /dev/null @@ -1,76 +0,0 @@ -#ifndef __BLOCK_DIAG_NVIDIA_H__ -#define __BLOCK_DIAG_NVIDIA_H__ - -#include "../../../operator.h" -#include - -namespace op::block_diag::nvidia { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t num_inputs; - std::vector output_shape; - ptrdiff_t output_stride0; - ptrdiff_t output_stride1; - std::vector row_offsets; - std::vector col_offsets; - std::vector> input_shapes; - std::vector input_rows; - std::vector input_cols; - std::vector input_stride0; - std::vector input_stride1; - size_t output_size; - - Descriptor(infiniDtype_t dtype, size_t num_inputs, - std::vector output_shape, - ptrdiff_t output_stride0, - ptrdiff_t output_stride1, - std::vector row_offsets, - std::vector col_offsets, - std::vector> input_shapes, - std::vector input_rows, - std::vector input_cols, - std::vector input_stride0, - std::vector input_stride1, - size_t output_size, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - num_inputs(num_inputs), - output_shape(std::move(output_shape)), - output_stride0(output_stride0), - output_stride1(output_stride1), - row_offsets(std::move(row_offsets)), - col_offsets(std::move(col_offsets)), - input_shapes(std::move(input_shapes)), - input_rows(std::move(input_rows)), - input_cols(std::move(input_cols)), - input_stride0(std::move(input_stride0)), - input_stride1(std::move(input_stride1)), - output_size(output_size) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t *input_descs, - size_t num_inputs); - - size_t workspaceSize() const { - return 4 * num_inputs * sizeof(size_t) + 2 * num_inputs * sizeof(ptrdiff_t) + num_inputs * sizeof(void *); - } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void **inputs, - void *stream) const; -}; - -} // namespace op::block_diag::nvidia - -#endif // __BLOCK_DIAG_NVIDIA_H__ diff --git a/src/infiniop/ops/block_diag/operator.cc b/src/infiniop/ops/block_diag/operator.cc deleted file mode 100644 index 31e6198f0..000000000 --- a/src/infiniop/ops/block_diag/operator.cc +++ /dev/null @@ -1,171 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/block_diag.h" - -#ifdef ENABLE_CPU_API -#include "cpu/block_diag_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/block_diag_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/block_diag_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/block_diag_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateBlockDiagDescriptor( - infiniopHandle_t handle, - infiniopBlockDiagDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t *input_descs, - size_t num_inputs) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::block_diag::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - input_descs, \ - num_inputs) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetBlockDiagWorkspaceSize(infiniopBlockDiagDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopBlockDiag( - infiniopBlockDiagDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void **inputs, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, inputs, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyBlockDiagDescriptor(infiniopBlockDiagDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/broadcast_to/broadcast_to.h b/src/infiniop/ops/broadcast_to/broadcast_to.h deleted file mode 100644 index 1518b2cdf..000000000 --- a/src/infiniop/ops/broadcast_to/broadcast_to.h +++ /dev/null @@ -1,48 +0,0 @@ -#ifndef __BROADCAST_TO_H__ -#define __BROADCAST_TO_H__ - -#include "../../operator.h" -#include "info.h" // 引用对应的 BroadcastToInfo 定义 -#include - -// 宏定义：用于生成不同命名空间下的 Descriptor 类 -#define DESCRIPTOR(NAMESPACE) \ - namespace op::broadcast_to::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - BroadcastToInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - BroadcastToInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - const std::vector &input_descs); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *output, \ - const std::vector &inputs, \ - void *stream) const; \ - }; \ - } - -#endif // __BROADCAST_TO_H__ diff --git a/src/infiniop/ops/broadcast_to/cpu/broadcast_to_cpu.cc b/src/infiniop/ops/broadcast_to/cpu/broadcast_to_cpu.cc deleted file mode 100644 index a131f9523..000000000 --- a/src/infiniop/ops/broadcast_to/cpu/broadcast_to_cpu.cc +++ /dev/null @@ -1,115 +0,0 @@ -#include "broadcast_to_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include -#include - -#include "../../../../utils/custom_types.h" - -namespace op::broadcast_to::cpu { - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - _opaque = nullptr; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - const std::vector &input_descs) { - - auto handle = reinterpret_cast(handle_); - auto result = BroadcastToInfo::create(out_desc, input_descs); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - new Opaque(), - result.take(), - 0, // CPU 实现不需要 workspace - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -void calculate_cpu_impl( - const BroadcastToInfo &info, - void *output, - const void *input) { - - size_t count = info.count(); - int ndim = info.ndim(); - - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - -// 并行遍历输出的每一个元素 -#pragma omp parallel for schedule(static) - for (ptrdiff_t i = 0; i < (ptrdiff_t)count; ++i) { - size_t temp_idx = i; - size_t input_offset = 0; - - // 坐标变换：Output Linear Index -> Coordinate -> Input Linear Offset - for (int d = 0; d < ndim; ++d) { - size_t out_stride = info._out_strides[d]; - size_t coord = temp_idx / out_stride; - temp_idx %= out_stride; - input_offset += coord * info._in_strides[d]; - } - - // 3. 赋值 - out_ptr[i] = in_ptr[input_offset]; - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const std::vector &inputs, - void *stream) const { - - if (inputs.size() != 1) { - return INFINI_STATUS_BAD_PARAM; - } - const void *input = inputs[0]; - auto dtype = _info.dtype(); - - switch (dtype) { - case INFINI_DTYPE_F32: - cpu::calculate_cpu_impl(_info, output, input); - break; - case INFINI_DTYPE_F64: - cpu::calculate_cpu_impl(_info, output, input); - break; - case INFINI_DTYPE_F16: - cpu::calculate_cpu_impl(_info, output, input); - break; - case INFINI_DTYPE_BF16: - cpu::calculate_cpu_impl(_info, output, input); - break; - case INFINI_DTYPE_I64: - cpu::calculate_cpu_impl(_info, output, input); - break; - case INFINI_DTYPE_I32: - cpu::calculate_cpu_impl(_info, output, input); - break; - case INFINI_DTYPE_U8: - cpu::calculate_cpu_impl(_info, output, input); - break; - case INFINI_DTYPE_I8: - cpu::calculate_cpu_impl(_info, output, input); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::broadcast_to::cpu diff --git a/src/infiniop/ops/broadcast_to/cpu/broadcast_to_cpu.h b/src/infiniop/ops/broadcast_to/cpu/broadcast_to_cpu.h deleted file mode 100644 index 496001c41..000000000 --- a/src/infiniop/ops/broadcast_to/cpu/broadcast_to_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BROADCAST_TO_CPU_H__ -#define __BROADCAST_TO_CPU_H__ - -#include "../broadcast_to.h" - -DESCRIPTOR(cpu) - -#endif // __BROADCAST_TO_CPU_H__ diff --git a/src/infiniop/ops/broadcast_to/cuda/kernel.cuh b/src/infiniop/ops/broadcast_to/cuda/kernel.cuh deleted file mode 100644 index 46e0e0bad..000000000 --- a/src/infiniop/ops/broadcast_to/cuda/kernel.cuh +++ /dev/null @@ -1,50 +0,0 @@ -#ifndef __BROADCAST_TO_CUDA_CUH__ -#define __BROADCAST_TO_CUDA_CUH__ - -#include -#include -#include - -namespace op::broadcast_to::cuda { - -// 最大维度定义，需与 BroadcastToInfo 中的保持一致 -static constexpr int MAX_DIM = 8; -struct BroadcastStrides { - int64_t out_strides[MAX_DIM]; - int64_t in_strides[MAX_DIM]; -}; -template -__global__ void broadcast_kernel( - T *__restrict__ output, // Output data pointer - const T *__restrict__ input, // Input data pointer - int ndim, - size_t count, // Total elements in output - BroadcastStrides strides) { // Strides passed by value - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < count) { - size_t temp_idx = idx; - size_t input_offset = 0; - -// 坐标变换与偏移计算 -#pragma unroll - for (int i = 0; i < MAX_DIM; ++i) { - if (i >= ndim) { - break; - } - - int64_t out_s = strides.out_strides[i]; - int64_t in_s = strides.in_strides[i]; - size_t coord = temp_idx / out_s; - temp_idx %= out_s; - input_offset += coord * in_s; - } - - output[idx] = input[input_offset]; - } -} - -} // namespace op::broadcast_to::cuda - -#endif // __BROADCAST_TO_CUDA_CUH__ diff --git a/src/infiniop/ops/broadcast_to/info.h b/src/infiniop/ops/broadcast_to/info.h deleted file mode 100644 index ecd86e7d9..000000000 --- a/src/infiniop/ops/broadcast_to/info.h +++ /dev/null @@ -1,105 +0,0 @@ -#ifndef __BROADCAST_TO_INFO_H__ -#define __BROADCAST_TO_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include // for std::max -#include - -namespace op::broadcast_to { - -class BroadcastToInfo { - BroadcastToInfo() = default; - -public: - static constexpr int MAX_DIM = 8; // 定义最大维度，方便做定长数组 - - int _dtype; - int _ndim; // 统一后的维度（等于输出维度） - size_t _count; - - // 存储对齐后的用于计算的信息 - int64_t _out_shape[MAX_DIM]; - int64_t _out_strides[MAX_DIM]; - int64_t _in_shape[MAX_DIM]; // 已经对齐并填充了1的输入Shape - int64_t _in_strides[MAX_DIM]; // 已经处理过广播（stride=0）的输入Stride - - int dtype() const { return _dtype; } - int ndim() const { return _ndim; } - size_t count() const { return _count; } - - // 构造函数 - BroadcastToInfo(int dtype, int ndim, size_t count) - : _dtype(dtype), _ndim(ndim), _count(count) {} - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - const std::vector &input_descs) { - - if (input_descs.size() != 1) { - return INFINI_STATUS_BAD_PARAM; - } - auto input_desc = input_descs[0]; - - if (out_desc->dtype() != input_desc->dtype()) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - if (out_desc->ndim() < input_desc->ndim()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (out_desc->ndim() > MAX_DIM) { - return INFINI_STATUS_BAD_PARAM; - } - - BroadcastToInfo info(out_desc->dtype(), int(out_desc->ndim()), 0); - - // 1. 计算总元素个数并拷贝 Output 信息 - size_t count = 1; - for (int i = 0; i < info._ndim; ++i) { - info._out_shape[i] = out_desc->shape()[i]; - info._out_strides[i] = out_desc->strides()[i]; - count *= out_desc->shape()[i]; - } - info._count = count; - - // 2. 维度对齐与广播规则检查 (Alignment & Broadcasting) - // 计算维度差：例如 out(2,3,4), in(3,4) -> offset = 1 - int offset = info._ndim - int(input_desc->ndim()); - - for (int i = 0; i < info._ndim; ++i) { - // i 是输出的维度索引 - // in_i 是对应的输入维度索引 - int in_i = i - offset; - - int64_t out_dim = info._out_shape[i]; - int64_t in_dim = 1; // 默认填充 1 (Input 维度不足时) - int64_t in_stride = 0; // 默认 Stride 0 (对应填充的 1) - - if (in_i >= 0) { - // 如果输入在这个维度有定义 - in_dim = input_desc->shape()[in_i]; - in_stride = input_desc->strides()[in_i]; - } - - // 检查规则 - if (in_dim != out_dim && in_dim != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // 保存对齐后的信息 - info._in_shape[i] = in_dim; - - if (in_dim == 1 && out_dim > 1) { - info._in_strides[i] = 0; - } else { - info._in_strides[i] = in_stride; - } - } - - return utils::Result(info); - } -}; - -} // namespace op::broadcast_to - -#endif // __BROADCAST_TO_INFO_H__ diff --git a/src/infiniop/ops/broadcast_to/metax/broadcast_to_metax.h b/src/infiniop/ops/broadcast_to/metax/broadcast_to_metax.h deleted file mode 100644 index 119686978..000000000 --- a/src/infiniop/ops/broadcast_to/metax/broadcast_to_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BROADCAST_TO_METAX_API_H__ -#define __BROADCAST_TO_METAX_API_H__ - -#include "../broadcast_to.h" - -DESCRIPTOR(metax) - -#endif // __BROADCAST_TO_METAX_API_H__ diff --git a/src/infiniop/ops/broadcast_to/metax/broadcast_to_metax.maca b/src/infiniop/ops/broadcast_to/metax/broadcast_to_metax.maca deleted file mode 100644 index 28fefe60d..000000000 --- a/src/infiniop/ops/broadcast_to/metax/broadcast_to_metax.maca +++ /dev/null @@ -1,174 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "broadcast_to_metax.h" - -#include -#include -#include -#include -#include - -namespace op::broadcast_to::metax { - -// ================================================================== -// 1. Kernel 定义 -// ================================================================== - -// 最大维度需与 BroadcastToInfo::MAX_DIM 保持一致 -static constexpr int MAX_DIM = 8; - -struct BroadcastStrides { - int64_t out_strides[MAX_DIM]; - int64_t in_strides[MAX_DIM]; -}; - -template -__global__ void broadcast_kernel( - T *__restrict__ output, // Output data pointer - const T *__restrict__ input, // Input data pointer - int ndim, - size_t count, // Total elements in output - BroadcastStrides strides) { // Strides passed by value - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < count) { - size_t temp_idx = idx; - size_t input_offset = 0; -#pragma unroll - for (int i = 0; i < MAX_DIM; ++i) { - if (i >= ndim) { - break; - } - - int64_t out_s = strides.out_strides[i]; - int64_t in_s = strides.in_strides[i]; - size_t coord = temp_idx / out_s; - temp_idx %= out_s; - input_offset += coord * in_s; - } - - output[idx] = input[input_offset]; - } -} - -// ================================================================== -// 2. Kernel Launch Logic -// ================================================================== -template -void launch_kernel( - void *output, - const void *input, - const BroadcastToInfo &info, - void *stream) { - - auto in_ptr = reinterpret_cast(input); - auto out_ptr = reinterpret_cast(output); - - auto hc_stream = reinterpret_cast(stream); - - BroadcastStrides strides; - for (int i = 0; i < BroadcastToInfo::MAX_DIM; ++i) { - strides.out_strides[i] = info._out_strides[i]; - strides.in_strides[i] = info._in_strides[i]; - } - - size_t count = info.count(); - size_t block_size = 256; - size_t grid_size = (count + block_size - 1) / block_size; - if (grid_size == 0) { - grid_size = 1; - } - - broadcast_kernel - <<>>( - out_ptr, - in_ptr, - info.ndim(), - count, - strides); -} - -// ================================================================== -// 3. Descriptor 实现 -// ================================================================== -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - const std::vector &input_descs) { - - auto handle = reinterpret_cast(handle_); - - auto info_result = BroadcastToInfo::create(out_desc, input_descs); - if (!info_result) { - return info_result.status(); - } - - size_t workspace_size = 0; // broadcast_to 不需要额外 workspace - - *desc_ptr = new Descriptor( - new Opaque(), - info_result.take(), - workspace_size, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const std::vector &inputs, - void *stream) const { - - if (inputs.size() != 1) { - return INFINI_STATUS_BAD_PARAM; - } - const void *input = inputs[0]; - - auto dtype = _info.dtype(); - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel<__half>(output, input, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_I64: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_I32: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_U8: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_I8: - launch_kernel(output, input, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::broadcast_to::metax diff --git a/src/infiniop/ops/broadcast_to/moore/broadcast_to_moore.h b/src/infiniop/ops/broadcast_to/moore/broadcast_to_moore.h deleted file mode 100644 index c0401398a..000000000 --- a/src/infiniop/ops/broadcast_to/moore/broadcast_to_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BROADCAST_TO_MOORE_H__ -#define __BROADCAST_TO_MOORE_H__ - -#include "../broadcast_to.h" - -DESCRIPTOR(moore) - -#endif // __BROADCAST_TO_MOORE_H__ diff --git a/src/infiniop/ops/broadcast_to/moore/broadcast_to_moore.mu b/src/infiniop/ops/broadcast_to/moore/broadcast_to_moore.mu deleted file mode 100644 index 2b3d5e886..000000000 --- a/src/infiniop/ops/broadcast_to/moore/broadcast_to_moore.mu +++ /dev/null @@ -1,124 +0,0 @@ -#include "../../../devices/moore/moore_handle.h" -#include "broadcast_to_moore.h" -#include "broadcast_to_moore_kernel.h" -#include -#include -#include - -namespace op::broadcast_to::moore { - -// ================================================================== -// Kernel Launch Logic -// ================================================================== -template -void launch_kernel( - void *output, - const void *input, - const BroadcastToInfo &info, - void *stream) { - - auto in_ptr = reinterpret_cast(input); - auto out_ptr = reinterpret_cast(output); - - auto musa_stream = reinterpret_cast(stream); - - // 复制 strides 到 Kernel 定义的结构体中 - op::broadcast_to::moore::BroadcastStrides strides; - for (int i = 0; i < BroadcastToInfo::MAX_DIM; ++i) { - strides.out_strides[i] = info._out_strides[i]; - strides.in_strides[i] = info._in_strides[i]; - } - - size_t count = info.count(); - size_t block_size = 256; - size_t grid_size = (count + block_size - 1) / block_size; - - op::broadcast_to::moore::broadcast_kernel - <<>>( - out_ptr, - in_ptr, - info.ndim(), - count, - strides); -} - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - const std::vector &input_descs) { - - auto handle = reinterpret_cast(handle_); - - auto info_result = BroadcastToInfo::create(out_desc, input_descs); - if (!info_result) { - return info_result.status(); - } - size_t workspace_size = 0; - - *desc_ptr = new Descriptor( - new Opaque(), - info_result.take(), - workspace_size, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const std::vector &inputs, - void *stream) const { - - if (inputs.size() != 1) { - return INFINI_STATUS_BAD_PARAM; - } - const void *input = inputs[0]; - - auto dtype = _info.dtype(); - - // 3. 根据数据类型分发 Kernel - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel<__mt_bfloat16>(output, input, _info, stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_I64: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_I32: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_U8: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_I8: - launch_kernel(output, input, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::broadcast_to::moore diff --git a/src/infiniop/ops/broadcast_to/moore/broadcast_to_moore_kernel.h b/src/infiniop/ops/broadcast_to/moore/broadcast_to_moore_kernel.h deleted file mode 100644 index 1d6995b98..000000000 --- a/src/infiniop/ops/broadcast_to/moore/broadcast_to_moore_kernel.h +++ /dev/null @@ -1,59 +0,0 @@ -#ifndef __BROADCAST_TO_MOORE_KERNEL_H__ -#define __BROADCAST_TO_MOORE_KERNEL_H__ - -#include -#include -#include - -#include -#include -#include - -namespace op::broadcast_to::moore { - -// 最大维度定义，需与 BroadcastToInfo 中的保持一致 -static constexpr int MAX_DIM = 8; - -struct BroadcastStrides { - int64_t out_strides[MAX_DIM]; - int64_t in_strides[MAX_DIM]; -}; - -template -__global__ void broadcast_kernel( - T *__restrict__ output, // Output data pointer - const T *__restrict__ input, // Input data pointer - int ndim, - size_t count, // Total elements in output - BroadcastStrides strides) { // Strides passed by value - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < count) { - size_t temp_idx = idx; - size_t input_offset = 0; - -// 坐标变换与偏移计算 -// 将扁平的 output index 转换为多维坐标，再根据 input strides 计算 input offset -#pragma unroll - for (int i = 0; i < MAX_DIM; ++i) { - if (i >= ndim) { - break; - } - - int64_t out_s = strides.out_strides[i]; - int64_t in_s = strides.in_strides[i]; - - // 计算当前维度的坐标 - size_t coord = temp_idx / out_s; - temp_idx %= out_s; - input_offset += coord * in_s; - } - - output[idx] = input[input_offset]; - } -} - -} // namespace op::broadcast_to::moore - -#endif // __BROADCAST_TO_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/broadcast_to/nvidia/broadcast_to_nvidia.cu b/src/infiniop/ops/broadcast_to/nvidia/broadcast_to_nvidia.cu deleted file mode 100644 index 7914fa9d5..000000000 --- a/src/infiniop/ops/broadcast_to/nvidia/broadcast_to_nvidia.cu +++ /dev/null @@ -1,119 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" - -#include "../cuda/kernel.cuh" -#include "broadcast_to_nvidia.cuh" -#include -#include -#include - -namespace op::broadcast_to::nvidia { - -// ================================================================== -// Kernel Launch Logic -// ================================================================== -template -void launch_kernel( - void *output, - const void *input, - const BroadcastToInfo &info, - void *stream) { - auto in_ptr = reinterpret_cast(input); - auto out_ptr = reinterpret_cast(output); - - auto cuda_stream = reinterpret_cast(stream); - op::broadcast_to::cuda::BroadcastStrides strides; - for (int i = 0; i < BroadcastToInfo::MAX_DIM; ++i) { - strides.out_strides[i] = info._out_strides[i]; - strides.in_strides[i] = info._in_strides[i]; - } - size_t count = info.count(); - size_t block_size = 256; - size_t grid_size = (count + block_size - 1) / block_size; - op::broadcast_to::cuda::broadcast_kernel - <<>>( - out_ptr, - in_ptr, - info.ndim(), - count, - strides); -} - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - const std::vector &input_descs) { - auto info_result = BroadcastToInfo::create(out_desc, input_descs); - if (!info_result) { - return info_result.status(); - } - size_t workspace_size = 0; - - *desc_ptr = new Descriptor( - new Opaque(), - info_result.take(), - workspace_size, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const std::vector &inputs, - void *stream) const { - - if (inputs.size() != 1) { - return INFINI_STATUS_BAD_PARAM; - } - const void *input = inputs[0]; - - auto dtype = _info.dtype(); - - // 3. 根据数据类型分发 Kernel - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_I64: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_I32: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_U8: - launch_kernel(output, input, _info, stream); - break; - case INFINI_DTYPE_I8: - launch_kernel(output, input, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::broadcast_to::nvidia diff --git a/src/infiniop/ops/broadcast_to/nvidia/broadcast_to_nvidia.cuh b/src/infiniop/ops/broadcast_to/nvidia/broadcast_to_nvidia.cuh deleted file mode 100644 index eb18075d0..000000000 --- a/src/infiniop/ops/broadcast_to/nvidia/broadcast_to_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __BROADCAST_TO_NVIDIA_CUH__ -#define __BROADCAST_TO_NVIDIA_CUH__ - -#include "../broadcast_to.h" - -DESCRIPTOR(nvidia) - -#endif // __BROADCAST_TO_NVIDIA_CUH__ diff --git a/src/infiniop/ops/broadcast_to/operator.cc b/src/infiniop/ops/broadcast_to/operator.cc deleted file mode 100644 index ea861f970..000000000 --- a/src/infiniop/ops/broadcast_to/operator.cc +++ /dev/null @@ -1,200 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/broadcast_to.h" - -// --- 后端实现头文件 --- -#ifdef ENABLE_CPU_API -#include "cpu/broadcast_to_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) -#include "nvidia/broadcast_to_nvidia.cuh" -#endif - -#ifdef ENABLE_METAX_API -#include "metax/broadcast_to_metax.h" -#endif - -#ifdef ENABLE_MOORE_API -#include "moore/broadcast_to_moore.h" -#endif - -extern "C" { - -// ======================================================================= -// 1. 创建算子描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopCreateBroadcastToDescriptor( - infiniopHandle_t handle, - infiniopBroadcastToDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::broadcast_to::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - {x_desc}) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ======================================================================= -// 2. 获取 Workspace 大小 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopGetBroadcastToWorkspaceSize(infiniopBroadcastToDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_ALI_API - GET(INFINI_DEVICE_ALI, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -// ======================================================================= -// 3. 执行计算 (Calculate) -// ======================================================================= -__INFINI_C infiniStatus_t infiniopBroadcastTo( - infiniopBroadcastToDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { - -// 注意：{x} 用于构造 std::vector -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, {x}, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -// ======================================================================= -// 4. 销毁算子描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopDestroyBroadcastToDescriptor(infiniopBroadcastToDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_ALI_API - DELETE(INFINI_DEVICE_ALI, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} - -} // extern "C" diff --git a/src/infiniop/ops/causal_conv1d/causal_conv1d.h b/src/infiniop/ops/causal_conv1d/causal_conv1d.h deleted file mode 100644 index c5d35c5ec..000000000 --- a/src/infiniop/ops/causal_conv1d/causal_conv1d.h +++ /dev/null @@ -1,50 +0,0 @@ -#ifndef __INFINIOP_CAUSAL_CONV1D_H__ -#define __INFINIOP_CAUSAL_CONV1D_H__ - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::causal_conv1d::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - CausalConv1dInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor(Opaque *opaque, CausalConv1dInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t conv_state_desc, \ - infiniopTensorDescriptor_t final_conv_state_desc, \ - infiniopTensorDescriptor_t qkv_desc, \ - infiniopTensorDescriptor_t weight_desc, \ - infiniopTensorDescriptor_t bias_desc, \ - infiniopTensorDescriptor_t cu_seqlens_desc, \ - infiniopTensorDescriptor_t initial_state_indices_desc, \ - infiniopTensorDescriptor_t final_state_indices_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *out, void *conv_state, void *final_conv_state, \ - const void *qkv, const void *weight, const void *bias, \ - const void *cu_seqlens, \ - const void *initial_state_indices, \ - const void *final_state_indices, \ - void *stream) const; \ - }; \ - } - -#endif // __INFINIOP_CAUSAL_CONV1D_H__ diff --git a/src/infiniop/ops/causal_conv1d/info.h b/src/infiniop/ops/causal_conv1d/info.h deleted file mode 100644 index 3fcec69e4..000000000 --- a/src/infiniop/ops/causal_conv1d/info.h +++ /dev/null @@ -1,193 +0,0 @@ -// infiniop/ops/causal_conv1d/info.h - -#ifndef __CAUSAL_CONV1D_INFO_H__ -#define __CAUSAL_CONV1D_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" - -#include - -namespace op { -namespace causal_conv1d { - -class CausalConv1dInfo { - CausalConv1dInfo() = default; - -public: - infiniDtype_t data_dtype; - infiniDtype_t cu_seqlens_dtype; - infiniDtype_t initial_state_indices_dtype; - infiniDtype_t final_state_indices_dtype; - - bool has_bias; - bool has_cu_seqlens; - bool has_initial_state_indices; - bool has_final_state_indices; - bool indexed_state_pool; - - size_t B; - size_t T; - size_t C; - size_t state_len; - size_t kernel_size; - size_t request_count; - size_t total_tokens; - size_t pool_size; - - std::vector out_strides; - std::vector conv_state_strides; - std::vector final_conv_state_strides; - std::vector qkv_strides; - std::vector weight_strides; - std::vector bias_strides; - - static utils::Result - create(infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t conv_state_desc, - infiniopTensorDescriptor_t final_conv_state_desc, - infiniopTensorDescriptor_t qkv_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t bias_desc, - infiniopTensorDescriptor_t cu_seqlens_desc, - infiniopTensorDescriptor_t initial_state_indices_desc, - infiniopTensorDescriptor_t final_state_indices_desc) { - - if (out_desc == nullptr || conv_state_desc == nullptr || qkv_desc == nullptr || weight_desc == nullptr) { - return INFINI_STATUS_NULL_POINTER; - } - - auto data_dtype = qkv_desc->dtype(); - CHECK_DTYPE(data_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - if (out_desc->dtype() != data_dtype || conv_state_desc->dtype() != data_dtype || weight_desc->dtype() != data_dtype || (final_conv_state_desc != nullptr && final_conv_state_desc->dtype() != data_dtype) || (bias_desc != nullptr && bias_desc->dtype() != data_dtype)) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - bool has_cu = cu_seqlens_desc != nullptr; - bool has_initial_indices = initial_state_indices_desc != nullptr; - bool has_final_indices = final_state_indices_desc != nullptr; - bool indexed_pool = has_initial_indices || has_final_indices; - - if (has_final_indices && final_conv_state_desc != nullptr) { - return INFINI_STATUS_BAD_PARAM; - } - if (!has_final_indices && final_conv_state_desc == nullptr) { - return INFINI_STATUS_NULL_POINTER; - } - - if (out_desc->ndim() != 3 || qkv_desc->ndim() != 3 || conv_state_desc->ndim() != 3 || weight_desc->ndim() != 3 || (!has_final_indices && final_conv_state_desc->ndim() != 3)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (bias_desc != nullptr && bias_desc->ndim() != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - auto qkv_shape = qkv_desc->shape(); - auto out_shape = out_desc->shape(); - auto state_shape = conv_state_desc->shape(); - auto weight_shape = weight_desc->shape(); - - size_t B = qkv_shape[0]; - size_t T = qkv_shape[1]; - size_t C = qkv_shape[2]; - size_t total_tokens = B * T; - size_t request_count = B; - - if (has_cu) { - if (cu_seqlens_desc->ndim() != 1 || cu_seqlens_desc->shape()[0] < 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - CHECK_DTYPE(cu_seqlens_desc->dtype(), INFINI_DTYPE_I32, INFINI_DTYPE_I64); - if (B != 1 || out_shape[0] != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - request_count = cu_seqlens_desc->shape()[0] - 1; - total_tokens = T; - } - - if (out_shape[0] != qkv_shape[0] || out_shape[1] != qkv_shape[1] || out_shape[2] != C) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (qkv_desc->strides()[2] != 1 || out_desc->strides()[2] != 1) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - if (weight_shape[0] != C || weight_shape[1] != 1 || weight_shape[2] < 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - size_t kernel_size = weight_shape[2]; - size_t state_len = kernel_size - 1; - if (kernel_size != 4) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (state_shape[1] != C || state_shape[2] != state_len) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (bias_desc != nullptr && bias_desc->shape()[0] != C) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t pool_size = state_shape[0]; - if (!indexed_pool && state_shape[0] != request_count) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (!has_final_indices) { - auto final_shape = final_conv_state_desc->shape(); - if (final_shape[0] != request_count || final_shape[1] != C || final_shape[2] != state_len) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - infiniDtype_t initial_indices_dtype = INFINI_DTYPE_INVALID; - infiniDtype_t final_indices_dtype = INFINI_DTYPE_INVALID; - if (has_initial_indices) { - if (initial_state_indices_desc->ndim() != 1 || initial_state_indices_desc->shape()[0] != request_count) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - initial_indices_dtype = initial_state_indices_desc->dtype(); - CHECK_DTYPE(initial_indices_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - } - if (has_final_indices) { - if (final_state_indices_desc->ndim() != 1 || final_state_indices_desc->shape()[0] != request_count) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - final_indices_dtype = final_state_indices_desc->dtype(); - CHECK_DTYPE(final_indices_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - } - - CausalConv1dInfo info; - info.data_dtype = data_dtype; - info.cu_seqlens_dtype = has_cu ? cu_seqlens_desc->dtype() : INFINI_DTYPE_INVALID; - info.initial_state_indices_dtype = initial_indices_dtype; - info.final_state_indices_dtype = final_indices_dtype; - info.has_bias = bias_desc != nullptr; - info.has_cu_seqlens = has_cu; - info.has_initial_state_indices = has_initial_indices; - info.has_final_state_indices = has_final_indices; - info.indexed_state_pool = indexed_pool; - info.B = B; - info.T = T; - info.C = C; - info.state_len = state_len; - info.kernel_size = kernel_size; - info.request_count = request_count; - info.total_tokens = total_tokens; - info.pool_size = pool_size; - info.out_strides = out_desc->strides(); - info.conv_state_strides = conv_state_desc->strides(); - if (final_conv_state_desc != nullptr) { - info.final_conv_state_strides = final_conv_state_desc->strides(); - } - info.qkv_strides = qkv_desc->strides(); - info.weight_strides = weight_desc->strides(); - if (bias_desc != nullptr) { - info.bias_strides = bias_desc->strides(); - } - return utils::Result(info); - } -}; - -} // namespace causal_conv1d -} // namespace op - -#endif // __CAUSAL_CONV1D_INFO_H__ diff --git a/src/infiniop/ops/causal_conv1d/metax/causal_conv1d_metax.h b/src/infiniop/ops/causal_conv1d/metax/causal_conv1d_metax.h deleted file mode 100644 index 9c3467a77..000000000 --- a/src/infiniop/ops/causal_conv1d/metax/causal_conv1d_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CAUSAL_CONV1D_METAX_H__ -#define __CAUSAL_CONV1D_METAX_H__ - -#include "../causal_conv1d.h" - -DESCRIPTOR(metax) - -#endif // __CAUSAL_CONV1D_METAX_H__ diff --git a/src/infiniop/ops/causal_conv1d/metax/causal_conv1d_metax.maca b/src/infiniop/ops/causal_conv1d/metax/causal_conv1d_metax.maca deleted file mode 100644 index 2ef458969..000000000 --- a/src/infiniop/ops/causal_conv1d/metax/causal_conv1d_metax.maca +++ /dev/null @@ -1,331 +0,0 @@ -// causal_conv1d_metax.maca - -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "causal_conv1d_metax.h" - -namespace op { -namespace causal_conv1d { -namespace metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t conv_state_desc, - infiniopTensorDescriptor_t final_conv_state_desc, - infiniopTensorDescriptor_t qkv_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t bias_desc, - infiniopTensorDescriptor_t cu_seqlens_desc, - infiniopTensorDescriptor_t initial_state_indices_desc, - infiniopTensorDescriptor_t final_state_indices_desc) { - - auto result = CausalConv1dInfo::create( - out_desc, conv_state_desc, final_conv_state_desc, qkv_desc, weight_desc, - bias_desc, cu_seqlens_desc, initial_state_indices_desc, final_state_indices_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - result.take(), - 0, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -__device__ __forceinline__ float load_as_float(const T *ptr, ptrdiff_t offset) { - return static_cast(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as_float(const half *ptr, ptrdiff_t offset) { - return __half2float(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as_float(const cuda_bfloat16 *ptr, ptrdiff_t offset) { - return __bfloat162float(ptr[offset]); -} - -template -__device__ __forceinline__ T cast_from_float(float x) { - return static_cast(x); -} - -template <> -__device__ __forceinline__ half cast_from_float(float x) { - return __float2half(x); -} - -template <> -__device__ __forceinline__ cuda_bfloat16 cast_from_float(float x) { - return __float2bfloat16(x); -} - -__device__ __forceinline__ int64_t load_optional_index( - const void *indices, - bool is_i64, - int idx, - int fallback) { - if (indices == nullptr) { - return static_cast(fallback); - } - return is_i64 - ? static_cast(indices)[idx] - : static_cast(static_cast(indices)[idx]); -} - -template -__device__ __forceinline__ float load_history_k4( - const T *conv_state, - const T *qkv, - int64_t history_pos, - int64_t token_begin, - int token_batch, - int channel, - ptrdiff_t state_base, - ptrdiff_t state_s2, - ptrdiff_t qkv_s0, - ptrdiff_t qkv_s1, - ptrdiff_t qkv_s2) { - constexpr int64_t STATE_LEN = 3; - if (history_pos < STATE_LEN) { - return load_as_float(conv_state, state_base + static_cast(history_pos) * state_s2); - } - - const int64_t token_idx = token_begin + history_pos - STATE_LEN; - const ptrdiff_t qkv_off = static_cast(token_batch) * qkv_s0 + static_cast(token_idx) * qkv_s1 + static_cast(channel) * qkv_s2; - return load_as_float(qkv, qkv_off); -} - -template -__global__ void causal_conv1d_k4_kernel( - T *out, - T *conv_state, - T *final_conv_state, - const T *qkv, - const T *weight, - const T *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - bool has_bias, - bool has_cu_seqlens, - bool cu_seqlens_i64, - bool initial_state_indices_i64, - bool final_state_indices_i64, - bool indexed_state_pool, - size_t T_tokens, - size_t C, - size_t total_tokens, - size_t pool_size, - ptrdiff_t out_s0, - ptrdiff_t out_s1, - ptrdiff_t out_s2, - ptrdiff_t state_s0, - ptrdiff_t state_s1, - ptrdiff_t state_s2, - ptrdiff_t final_s0, - ptrdiff_t final_s1, - ptrdiff_t final_s2, - ptrdiff_t qkv_s0, - ptrdiff_t qkv_s1, - ptrdiff_t qkv_s2, - ptrdiff_t weight_s0, - ptrdiff_t weight_s1, - ptrdiff_t weight_s2, - ptrdiff_t bias_s0) { - - const int request = blockIdx.x; - const int channel = blockIdx.y; - if (channel >= C || threadIdx.x != 0) { - return; - } - - int64_t token_begin = 0; - int64_t token_end = static_cast(T_tokens); - int token_batch = request; - if (has_cu_seqlens) { - token_begin = load_optional_index(cu_seqlens, cu_seqlens_i64, request, 0); - token_end = load_optional_index(cu_seqlens, cu_seqlens_i64, request + 1, 0); - token_batch = 0; - if (token_begin < 0 || token_end < token_begin || token_end > static_cast(total_tokens)) { - return; - } - } - const int64_t request_len = token_end - token_begin; - - int64_t read_slot = indexed_state_pool - ? load_optional_index(initial_state_indices, initial_state_indices_i64, request, request) - : static_cast(request); - int64_t write_slot = indexed_state_pool && final_state_indices != nullptr - ? load_optional_index(final_state_indices, final_state_indices_i64, request, request) - : static_cast(request); - if (read_slot < 0 || write_slot < 0 || read_slot >= static_cast(pool_size) || (final_state_indices != nullptr && write_slot >= static_cast(pool_size))) { - return; - } - - const ptrdiff_t state_base = static_cast(read_slot) * state_s0 + static_cast(channel) * state_s1; - const ptrdiff_t weight_base = static_cast(channel) * weight_s0; - (void)weight_s1; - - for (int64_t t = 0; t < request_len; ++t) { - float acc = load_as_float(weight, weight_base) * load_history_k4(conv_state, qkv, t, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2) + load_as_float(weight, weight_base + weight_s2) * load_history_k4(conv_state, qkv, t + 1, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2) + load_as_float(weight, weight_base + 2 * weight_s2) * load_history_k4(conv_state, qkv, t + 2, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2) + load_as_float(weight, weight_base + 3 * weight_s2) * load_history_k4(conv_state, qkv, t + 3, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2); - if (has_bias) { - acc += load_as_float(bias, static_cast(channel) * bias_s0); - } - const ptrdiff_t out_off = static_cast(token_batch) * out_s0 + static_cast(token_begin + t) * out_s1 + static_cast(channel) * out_s2; - out[out_off] = cast_from_float(acc); - } - - T *final_target = final_state_indices != nullptr ? conv_state : final_conv_state; - ptrdiff_t final_base; - ptrdiff_t final_stride0; - ptrdiff_t final_stride1; - ptrdiff_t final_stride2; - if (final_state_indices != nullptr) { - final_base = static_cast(write_slot) * state_s0 + static_cast(channel) * state_s1; - final_stride0 = state_s0; - final_stride1 = state_s1; - final_stride2 = state_s2; - } else { - final_base = static_cast(request) * final_s0 + static_cast(channel) * final_s1; - final_stride0 = final_s0; - final_stride1 = final_s1; - final_stride2 = final_s2; - } - (void)final_stride0; - (void)final_stride1; - - final_target[final_base] = cast_from_float( - load_history_k4(conv_state, qkv, request_len, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2)); - final_target[final_base + final_stride2] = cast_from_float( - load_history_k4(conv_state, qkv, request_len + 1, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2)); - final_target[final_base + 2 * final_stride2] = cast_from_float( - load_history_k4(conv_state, qkv, request_len + 2, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2)); -} - -template -infiniStatus_t launch_k4( - const CausalConv1dInfo &info, - void *out, - void *conv_state, - void *final_conv_state, - const void *qkv, - const void *weight, - const void *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - hcStream_t stream) { - - dim3 grid(static_cast(info.request_count), static_cast(info.C), 1); - causal_conv1d_k4_kernel<<>>( - static_cast(out), - static_cast(conv_state), - static_cast(final_conv_state), - static_cast(qkv), - static_cast(weight), - static_cast(bias), - cu_seqlens, - initial_state_indices, - final_state_indices, - info.has_bias, - info.has_cu_seqlens, - info.cu_seqlens_dtype == INFINI_DTYPE_I64, - info.initial_state_indices_dtype == INFINI_DTYPE_I64, - info.final_state_indices_dtype == INFINI_DTYPE_I64, - info.indexed_state_pool, - info.T, - info.C, - info.total_tokens, - info.pool_size, - info.out_strides[0], - info.out_strides[1], - info.out_strides[2], - info.conv_state_strides[0], - info.conv_state_strides[1], - info.conv_state_strides[2], - info.final_conv_state_strides.empty() ? 0 : info.final_conv_state_strides[0], - info.final_conv_state_strides.empty() ? 0 : info.final_conv_state_strides[1], - info.final_conv_state_strides.empty() ? 0 : info.final_conv_state_strides[2], - info.qkv_strides[0], - info.qkv_strides[1], - info.qkv_strides[2], - info.weight_strides[0], - info.weight_strides[1], - info.weight_strides[2], - info.bias_strides.empty() ? 0 : info.bias_strides[0]); - CHECK_METAX(hcGetLastError()); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t launch_by_kernel_size( - const CausalConv1dInfo &info, - void *out, - void *conv_state, - void *final_conv_state, - const void *qkv, - const void *weight, - const void *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - hcStream_t stream) { - - if (info.kernel_size != 4) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - switch (info.data_dtype) { - case INFINI_DTYPE_F32: - return launch_k4(info, out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, initial_state_indices, final_state_indices, stream); - case INFINI_DTYPE_F16: - return launch_k4(info, out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, initial_state_indices, final_state_indices, stream); - case INFINI_DTYPE_BF16: - return launch_k4(info, out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, initial_state_indices, final_state_indices, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - void *conv_state, - void *final_conv_state, - const void *qkv, - const void *weight, - const void *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - hcStream_t cuda_stream = reinterpret_cast(stream); - return launch_by_kernel_size(_info, out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, initial_state_indices, final_state_indices, cuda_stream); -} - -} // namespace metax -} // namespace causal_conv1d -} // namespace op diff --git a/src/infiniop/ops/causal_conv1d/moore/causal_conv1d_moore.h b/src/infiniop/ops/causal_conv1d/moore/causal_conv1d_moore.h deleted file mode 100644 index 65d8405cc..000000000 --- a/src/infiniop/ops/causal_conv1d/moore/causal_conv1d_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CAUSAL_CONV1D_MOORE_H__ -#define __CAUSAL_CONV1D_MOORE_H__ - -#include "../causal_conv1d.h" - -DESCRIPTOR(moore) - -#endif // __CAUSAL_CONV1D_MOORE_H__ diff --git a/src/infiniop/ops/causal_conv1d/moore/causal_conv1d_moore.mu b/src/infiniop/ops/causal_conv1d/moore/causal_conv1d_moore.mu deleted file mode 100644 index 5ebfe4b11..000000000 --- a/src/infiniop/ops/causal_conv1d/moore/causal_conv1d_moore.mu +++ /dev/null @@ -1,331 +0,0 @@ -// causal_conv1d_moore.mu - -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_handle.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "causal_conv1d_moore.h" - -namespace op { -namespace causal_conv1d { -namespace moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t conv_state_desc, - infiniopTensorDescriptor_t final_conv_state_desc, - infiniopTensorDescriptor_t qkv_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t bias_desc, - infiniopTensorDescriptor_t cu_seqlens_desc, - infiniopTensorDescriptor_t initial_state_indices_desc, - infiniopTensorDescriptor_t final_state_indices_desc) { - - auto result = CausalConv1dInfo::create( - out_desc, conv_state_desc, final_conv_state_desc, qkv_desc, weight_desc, - bias_desc, cu_seqlens_desc, initial_state_indices_desc, final_state_indices_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - result.take(), - 0, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -__device__ __forceinline__ float load_as_float(const T *ptr, ptrdiff_t offset) { - return static_cast(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as_float(const half *ptr, ptrdiff_t offset) { - return __half2float(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as_float(const cuda_bfloat16 *ptr, ptrdiff_t offset) { - return __bfloat162float(ptr[offset]); -} - -template -__device__ __forceinline__ T cast_from_float(float x) { - return static_cast(x); -} - -template <> -__device__ __forceinline__ half cast_from_float(float x) { - return __float2half(x); -} - -template <> -__device__ __forceinline__ cuda_bfloat16 cast_from_float(float x) { - return __float2bfloat16(x); -} - -__device__ __forceinline__ int64_t load_optional_index( - const void *indices, - bool is_i64, - int idx, - int fallback) { - if (indices == nullptr) { - return static_cast(fallback); - } - return is_i64 - ? static_cast(indices)[idx] - : static_cast(static_cast(indices)[idx]); -} - -template -__device__ __forceinline__ float load_history_k4( - const T *conv_state, - const T *qkv, - int64_t history_pos, - int64_t token_begin, - int token_batch, - int channel, - ptrdiff_t state_base, - ptrdiff_t state_s2, - ptrdiff_t qkv_s0, - ptrdiff_t qkv_s1, - ptrdiff_t qkv_s2) { - constexpr int64_t STATE_LEN = 3; - if (history_pos < STATE_LEN) { - return load_as_float(conv_state, state_base + static_cast(history_pos) * state_s2); - } - - const int64_t token_idx = token_begin + history_pos - STATE_LEN; - const ptrdiff_t qkv_off = static_cast(token_batch) * qkv_s0 + static_cast(token_idx) * qkv_s1 + static_cast(channel) * qkv_s2; - return load_as_float(qkv, qkv_off); -} - -template -__global__ void causal_conv1d_k4_kernel( - T *out, - T *conv_state, - T *final_conv_state, - const T *qkv, - const T *weight, - const T *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - bool has_bias, - bool has_cu_seqlens, - bool cu_seqlens_i64, - bool initial_state_indices_i64, - bool final_state_indices_i64, - bool indexed_state_pool, - size_t T_tokens, - size_t C, - size_t total_tokens, - size_t pool_size, - ptrdiff_t out_s0, - ptrdiff_t out_s1, - ptrdiff_t out_s2, - ptrdiff_t state_s0, - ptrdiff_t state_s1, - ptrdiff_t state_s2, - ptrdiff_t final_s0, - ptrdiff_t final_s1, - ptrdiff_t final_s2, - ptrdiff_t qkv_s0, - ptrdiff_t qkv_s1, - ptrdiff_t qkv_s2, - ptrdiff_t weight_s0, - ptrdiff_t weight_s1, - ptrdiff_t weight_s2, - ptrdiff_t bias_s0) { - - const int request = blockIdx.x; - const int channel = blockIdx.y; - if (channel >= C || threadIdx.x != 0) { - return; - } - - int64_t token_begin = 0; - int64_t token_end = static_cast(T_tokens); - int token_batch = request; - if (has_cu_seqlens) { - token_begin = load_optional_index(cu_seqlens, cu_seqlens_i64, request, 0); - token_end = load_optional_index(cu_seqlens, cu_seqlens_i64, request + 1, 0); - token_batch = 0; - if (token_begin < 0 || token_end < token_begin || token_end > static_cast(total_tokens)) { - return; - } - } - const int64_t request_len = token_end - token_begin; - - int64_t read_slot = indexed_state_pool - ? load_optional_index(initial_state_indices, initial_state_indices_i64, request, request) - : static_cast(request); - int64_t write_slot = indexed_state_pool && final_state_indices != nullptr - ? load_optional_index(final_state_indices, final_state_indices_i64, request, request) - : static_cast(request); - if (read_slot < 0 || write_slot < 0 || read_slot >= static_cast(pool_size) || (final_state_indices != nullptr && write_slot >= static_cast(pool_size))) { - return; - } - - const ptrdiff_t state_base = static_cast(read_slot) * state_s0 + static_cast(channel) * state_s1; - const ptrdiff_t weight_base = static_cast(channel) * weight_s0; - (void)weight_s1; - - for (int64_t t = 0; t < request_len; ++t) { - float acc = load_as_float(weight, weight_base) * load_history_k4(conv_state, qkv, t, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2) + load_as_float(weight, weight_base + weight_s2) * load_history_k4(conv_state, qkv, t + 1, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2) + load_as_float(weight, weight_base + 2 * weight_s2) * load_history_k4(conv_state, qkv, t + 2, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2) + load_as_float(weight, weight_base + 3 * weight_s2) * load_history_k4(conv_state, qkv, t + 3, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2); - if (has_bias) { - acc += load_as_float(bias, static_cast(channel) * bias_s0); - } - const ptrdiff_t out_off = static_cast(token_batch) * out_s0 + static_cast(token_begin + t) * out_s1 + static_cast(channel) * out_s2; - out[out_off] = cast_from_float(acc); - } - - T *final_target = final_state_indices != nullptr ? conv_state : final_conv_state; - ptrdiff_t final_base; - ptrdiff_t final_stride0; - ptrdiff_t final_stride1; - ptrdiff_t final_stride2; - if (final_state_indices != nullptr) { - final_base = static_cast(write_slot) * state_s0 + static_cast(channel) * state_s1; - final_stride0 = state_s0; - final_stride1 = state_s1; - final_stride2 = state_s2; - } else { - final_base = static_cast(request) * final_s0 + static_cast(channel) * final_s1; - final_stride0 = final_s0; - final_stride1 = final_s1; - final_stride2 = final_s2; - } - (void)final_stride0; - (void)final_stride1; - - final_target[final_base] = cast_from_float( - load_history_k4(conv_state, qkv, request_len, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2)); - final_target[final_base + final_stride2] = cast_from_float( - load_history_k4(conv_state, qkv, request_len + 1, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2)); - final_target[final_base + 2 * final_stride2] = cast_from_float( - load_history_k4(conv_state, qkv, request_len + 2, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2)); -} - -template -infiniStatus_t launch_k4( - const CausalConv1dInfo &info, - void *out, - void *conv_state, - void *final_conv_state, - const void *qkv, - const void *weight, - const void *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - musaStream_t stream) { - - dim3 grid(static_cast(info.request_count), static_cast(info.C), 1); - causal_conv1d_k4_kernel<<>>( - static_cast(out), - static_cast(conv_state), - static_cast(final_conv_state), - static_cast(qkv), - static_cast(weight), - static_cast(bias), - cu_seqlens, - initial_state_indices, - final_state_indices, - info.has_bias, - info.has_cu_seqlens, - info.cu_seqlens_dtype == INFINI_DTYPE_I64, - info.initial_state_indices_dtype == INFINI_DTYPE_I64, - info.final_state_indices_dtype == INFINI_DTYPE_I64, - info.indexed_state_pool, - info.T, - info.C, - info.total_tokens, - info.pool_size, - info.out_strides[0], - info.out_strides[1], - info.out_strides[2], - info.conv_state_strides[0], - info.conv_state_strides[1], - info.conv_state_strides[2], - info.final_conv_state_strides.empty() ? 0 : info.final_conv_state_strides[0], - info.final_conv_state_strides.empty() ? 0 : info.final_conv_state_strides[1], - info.final_conv_state_strides.empty() ? 0 : info.final_conv_state_strides[2], - info.qkv_strides[0], - info.qkv_strides[1], - info.qkv_strides[2], - info.weight_strides[0], - info.weight_strides[1], - info.weight_strides[2], - info.bias_strides.empty() ? 0 : info.bias_strides[0]); - CHECK_MOORE(musaGetLastError()); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t launch_by_kernel_size( - const CausalConv1dInfo &info, - void *out, - void *conv_state, - void *final_conv_state, - const void *qkv, - const void *weight, - const void *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - musaStream_t stream) { - - if (info.kernel_size != 4) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - switch (info.data_dtype) { - case INFINI_DTYPE_F32: - return launch_k4(info, out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, initial_state_indices, final_state_indices, stream); - case INFINI_DTYPE_F16: - return launch_k4(info, out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, initial_state_indices, final_state_indices, stream); - case INFINI_DTYPE_BF16: - return launch_k4(info, out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, initial_state_indices, final_state_indices, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - void *conv_state, - void *final_conv_state, - const void *qkv, - const void *weight, - const void *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - musaStream_t cuda_stream = reinterpret_cast(stream); - return launch_by_kernel_size(_info, out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, initial_state_indices, final_state_indices, cuda_stream); -} - -} // namespace moore -} // namespace causal_conv1d -} // namespace op diff --git a/src/infiniop/ops/causal_conv1d/nvidia/causal_conv1d_nvidia.cu b/src/infiniop/ops/causal_conv1d/nvidia/causal_conv1d_nvidia.cu deleted file mode 100644 index 701632fe5..000000000 --- a/src/infiniop/ops/causal_conv1d/nvidia/causal_conv1d_nvidia.cu +++ /dev/null @@ -1,334 +0,0 @@ -// causal_conv1d_nvidia.cu - -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "causal_conv1d_nvidia.cuh" - -#include -#include - -namespace op { -namespace causal_conv1d { -namespace nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t conv_state_desc, - infiniopTensorDescriptor_t final_conv_state_desc, - infiniopTensorDescriptor_t qkv_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t bias_desc, - infiniopTensorDescriptor_t cu_seqlens_desc, - infiniopTensorDescriptor_t initial_state_indices_desc, - infiniopTensorDescriptor_t final_state_indices_desc) { - - auto result = CausalConv1dInfo::create( - out_desc, conv_state_desc, final_conv_state_desc, qkv_desc, weight_desc, - bias_desc, cu_seqlens_desc, initial_state_indices_desc, final_state_indices_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - result.take(), - 0, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -__device__ __forceinline__ float load_as_float(const T *ptr, ptrdiff_t offset) { - return static_cast(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as_float(const half *ptr, ptrdiff_t offset) { - return __half2float(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as_float<__nv_bfloat16>(const __nv_bfloat16 *ptr, ptrdiff_t offset) { - return __bfloat162float(ptr[offset]); -} - -template -__device__ __forceinline__ T cast_from_float(float x) { - return static_cast(x); -} - -template <> -__device__ __forceinline__ half cast_from_float(float x) { - return __float2half(x); -} - -template <> -__device__ __forceinline__ __nv_bfloat16 cast_from_float<__nv_bfloat16>(float x) { - return __float2bfloat16(x); -} - -__device__ __forceinline__ int64_t load_optional_index( - const void *indices, - bool is_i64, - int idx, - int fallback) { - if (indices == nullptr) { - return static_cast(fallback); - } - return is_i64 - ? static_cast(indices)[idx] - : static_cast(static_cast(indices)[idx]); -} - -template -__device__ __forceinline__ float load_history_k4( - const T *conv_state, - const T *qkv, - int64_t history_pos, - int64_t token_begin, - int token_batch, - int channel, - ptrdiff_t state_base, - ptrdiff_t state_s2, - ptrdiff_t qkv_s0, - ptrdiff_t qkv_s1, - ptrdiff_t qkv_s2) { - constexpr int64_t STATE_LEN = 3; - if (history_pos < STATE_LEN) { - return load_as_float(conv_state, state_base + static_cast(history_pos) * state_s2); - } - - const int64_t token_idx = token_begin + history_pos - STATE_LEN; - const ptrdiff_t qkv_off = static_cast(token_batch) * qkv_s0 + static_cast(token_idx) * qkv_s1 + static_cast(channel) * qkv_s2; - return load_as_float(qkv, qkv_off); -} - -template -__global__ void causal_conv1d_k4_kernel( - T *out, - T *conv_state, - T *final_conv_state, - const T *qkv, - const T *weight, - const T *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - bool has_bias, - bool has_cu_seqlens, - bool cu_seqlens_i64, - bool initial_state_indices_i64, - bool final_state_indices_i64, - bool indexed_state_pool, - size_t T_tokens, - size_t C, - size_t total_tokens, - size_t pool_size, - ptrdiff_t out_s0, - ptrdiff_t out_s1, - ptrdiff_t out_s2, - ptrdiff_t state_s0, - ptrdiff_t state_s1, - ptrdiff_t state_s2, - ptrdiff_t final_s0, - ptrdiff_t final_s1, - ptrdiff_t final_s2, - ptrdiff_t qkv_s0, - ptrdiff_t qkv_s1, - ptrdiff_t qkv_s2, - ptrdiff_t weight_s0, - ptrdiff_t weight_s1, - ptrdiff_t weight_s2, - ptrdiff_t bias_s0) { - - const int request = blockIdx.x; - const int channel = blockIdx.y; - if (channel >= C || threadIdx.x != 0) { - return; - } - - int64_t token_begin = 0; - int64_t token_end = static_cast(T_tokens); - int token_batch = request; - if (has_cu_seqlens) { - token_begin = load_optional_index(cu_seqlens, cu_seqlens_i64, request, 0); - token_end = load_optional_index(cu_seqlens, cu_seqlens_i64, request + 1, 0); - token_batch = 0; - if (token_begin < 0 || token_end < token_begin || token_end > static_cast(total_tokens)) { - return; - } - } - const int64_t request_len = token_end - token_begin; - - int64_t read_slot = indexed_state_pool - ? load_optional_index(initial_state_indices, initial_state_indices_i64, request, request) - : static_cast(request); - int64_t write_slot = indexed_state_pool && final_state_indices != nullptr - ? load_optional_index(final_state_indices, final_state_indices_i64, request, request) - : static_cast(request); - if (read_slot < 0 || write_slot < 0 || read_slot >= static_cast(pool_size) || (final_state_indices != nullptr && write_slot >= static_cast(pool_size))) { - return; - } - - const ptrdiff_t state_base = static_cast(read_slot) * state_s0 + static_cast(channel) * state_s1; - const ptrdiff_t weight_base = static_cast(channel) * weight_s0; - (void)weight_s1; - - for (int64_t t = 0; t < request_len; ++t) { - float acc = load_as_float(weight, weight_base) * load_history_k4(conv_state, qkv, t, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2) + load_as_float(weight, weight_base + weight_s2) * load_history_k4(conv_state, qkv, t + 1, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2) + load_as_float(weight, weight_base + 2 * weight_s2) * load_history_k4(conv_state, qkv, t + 2, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2) + load_as_float(weight, weight_base + 3 * weight_s2) * load_history_k4(conv_state, qkv, t + 3, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2); - if (has_bias) { - acc += load_as_float(bias, static_cast(channel) * bias_s0); - } - const ptrdiff_t out_off = static_cast(token_batch) * out_s0 + static_cast(token_begin + t) * out_s1 + static_cast(channel) * out_s2; - out[out_off] = cast_from_float(acc); - } - - T *final_target = final_state_indices != nullptr ? conv_state : final_conv_state; - ptrdiff_t final_base; - ptrdiff_t final_stride0; - ptrdiff_t final_stride1; - ptrdiff_t final_stride2; - if (final_state_indices != nullptr) { - final_base = static_cast(write_slot) * state_s0 + static_cast(channel) * state_s1; - final_stride0 = state_s0; - final_stride1 = state_s1; - final_stride2 = state_s2; - } else { - final_base = static_cast(request) * final_s0 + static_cast(channel) * final_s1; - final_stride0 = final_s0; - final_stride1 = final_s1; - final_stride2 = final_s2; - } - (void)final_stride0; - (void)final_stride1; - - final_target[final_base] = cast_from_float( - load_history_k4(conv_state, qkv, request_len, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2)); - final_target[final_base + final_stride2] = cast_from_float( - load_history_k4(conv_state, qkv, request_len + 1, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2)); - final_target[final_base + 2 * final_stride2] = cast_from_float( - load_history_k4(conv_state, qkv, request_len + 2, token_begin, token_batch, channel, state_base, state_s2, qkv_s0, qkv_s1, qkv_s2)); -} - -template -infiniStatus_t launch_k4( - const CausalConv1dInfo &info, - void *out, - void *conv_state, - void *final_conv_state, - const void *qkv, - const void *weight, - const void *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - cudaStream_t stream) { - - dim3 grid(static_cast(info.request_count), static_cast(info.C), 1); - causal_conv1d_k4_kernel<<>>( - static_cast(out), - static_cast(conv_state), - static_cast(final_conv_state), - static_cast(qkv), - static_cast(weight), - static_cast(bias), - cu_seqlens, - initial_state_indices, - final_state_indices, - info.has_bias, - info.has_cu_seqlens, - info.cu_seqlens_dtype == INFINI_DTYPE_I64, - info.initial_state_indices_dtype == INFINI_DTYPE_I64, - info.final_state_indices_dtype == INFINI_DTYPE_I64, - info.indexed_state_pool, - info.T, - info.C, - info.total_tokens, - info.pool_size, - info.out_strides[0], - info.out_strides[1], - info.out_strides[2], - info.conv_state_strides[0], - info.conv_state_strides[1], - info.conv_state_strides[2], - info.final_conv_state_strides.empty() ? 0 : info.final_conv_state_strides[0], - info.final_conv_state_strides.empty() ? 0 : info.final_conv_state_strides[1], - info.final_conv_state_strides.empty() ? 0 : info.final_conv_state_strides[2], - info.qkv_strides[0], - info.qkv_strides[1], - info.qkv_strides[2], - info.weight_strides[0], - info.weight_strides[1], - info.weight_strides[2], - info.bias_strides.empty() ? 0 : info.bias_strides[0]); - CHECK_CUDA(cudaGetLastError()); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t launch_by_kernel_size( - const CausalConv1dInfo &info, - void *out, - void *conv_state, - void *final_conv_state, - const void *qkv, - const void *weight, - const void *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - cudaStream_t stream) { - - if (info.kernel_size != 4) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - switch (info.data_dtype) { - case INFINI_DTYPE_F32: - return launch_k4(info, out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, initial_state_indices, final_state_indices, stream); - case INFINI_DTYPE_F16: - return launch_k4(info, out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, initial_state_indices, final_state_indices, stream); - case INFINI_DTYPE_BF16: - return launch_k4<__nv_bfloat16>(info, out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, initial_state_indices, final_state_indices, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - void *conv_state, - void *final_conv_state, - const void *qkv, - const void *weight, - const void *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - cudaStream_t cuda_stream = reinterpret_cast(stream); - return launch_by_kernel_size(_info, out, conv_state, final_conv_state, qkv, weight, bias, cu_seqlens, initial_state_indices, final_state_indices, cuda_stream); -} - -} // namespace nvidia -} // namespace causal_conv1d -} // namespace op diff --git a/src/infiniop/ops/causal_conv1d/nvidia/causal_conv1d_nvidia.cuh b/src/infiniop/ops/causal_conv1d/nvidia/causal_conv1d_nvidia.cuh deleted file mode 100644 index 15a0b5f78..000000000 --- a/src/infiniop/ops/causal_conv1d/nvidia/causal_conv1d_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CAUSAL_CONV1D_NVIDIA_H__ -#define __CAUSAL_CONV1D_NVIDIA_H__ - -#include "../causal_conv1d.h" - -DESCRIPTOR(nvidia) - -#endif // __CAUSAL_CONV1D_NVIDIA_H__ diff --git a/src/infiniop/ops/causal_conv1d/operator.cc b/src/infiniop/ops/causal_conv1d/operator.cc deleted file mode 100644 index 8666d2e62..000000000 --- a/src/infiniop/ops/causal_conv1d/operator.cc +++ /dev/null @@ -1,146 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/causal_conv1d.h" - -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_HYGON_API) -#include "nvidia/causal_conv1d_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/causal_conv1d_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/causal_conv1d_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateCausalConv1dDescriptor( - infiniopHandle_t handle, - infiniopCausalConv1dDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t conv_state_desc, - infiniopTensorDescriptor_t final_conv_state_desc, - infiniopTensorDescriptor_t qkv_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t bias_desc, - infiniopTensorDescriptor_t cu_seqlens_desc, - infiniopTensorDescriptor_t initial_state_indices_desc, - infiniopTensorDescriptor_t final_state_indices_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::causal_conv1d::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - out_desc, conv_state_desc, final_conv_state_desc, \ - qkv_desc, weight_desc, bias_desc, cu_seqlens_desc, \ - initial_state_indices_desc, final_state_indices_desc); - - switch (handle->device) { -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia) -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetCausalConv1dWorkspaceSize( - infiniopCausalConv1dDescriptor_t desc, size_t *size) { -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopCausalConv1d( - infiniopCausalConv1dDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - void *conv_state, - void *final_conv_state, - const void *qkv, - const void *weight, - const void *bias, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - void *stream) { -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc)->calculate( \ - workspace, workspace_size, out, conv_state, final_conv_state, \ - qkv, weight, bias, cu_seqlens, initial_state_indices, \ - final_state_indices, stream); - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia) -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyCausalConv1dDescriptor( - infiniopCausalConv1dDescriptor_t desc) { -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_HYGON_API - DESTROY(INFINI_DEVICE_HYGON, nvidia) -#endif -#ifdef ENABLE_METAX_API - DESTROY(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - DESTROY(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DESTROY -} diff --git a/src/infiniop/ops/causal_softmax/ascend/causal_softmax_ascend.cc b/src/infiniop/ops/causal_softmax/ascend/causal_softmax_ascend.cc deleted file mode 100644 index 3b183ae0f..000000000 --- a/src/infiniop/ops/causal_softmax/ascend/causal_softmax_ascend.cc +++ /dev/null @@ -1,230 +0,0 @@ -#include "causal_softmax_ascend.h" -#include "../../../devices/ascend/common_ascend.h" -#include -#include -#include -#include - -namespace op::causal_softmax::ascend { - -namespace { - -bool isCompact(const CausalSoftmaxInfo &info, ptrdiff_t stride_b, ptrdiff_t stride_i, ptrdiff_t stride_j) { - return stride_j == 1 - && stride_i == static_cast(info.total_seq_len) - && (info.batch_size == 1 || stride_b == static_cast(info.seq_len * info.total_seq_len)); -} - -} // namespace - -struct Descriptor::Opaque { - aclnnTensorDescriptor_t x; - aclnnTensorDescriptor_t mask; - aclnnTensorDescriptor_t y; - aclnnTensorDescriptor_t value; - aclnnTensorDescriptor_t temp_x; - aclnnTensorDescriptor_t temp_y; - void *mask_addr; - void *value_addr; - void *temp_x_addr; - void *temp_y_addr; - size_t workspacesize; - aclOpExecutor *executor; - aclOpExecutor *temp_executor; - aclOpExecutor *copy_in_executor; - aclOpExecutor *copy_out_executor; - bool use_temp; - - ~Opaque() { - delete x; - delete mask; - delete y; - delete value; - delete temp_x; - delete temp_y; - - aclrtFree(mask_addr); - aclrtFree(value_addr); - aclrtFree(temp_x_addr); - aclrtFree(temp_y_addr); - - // Delete useless executor - aclDestroyAclOpExecutor(executor); - aclDestroyAclOpExecutor(temp_executor); - aclDestroyAclOpExecutor(copy_in_executor); - aclDestroyAclOpExecutor(copy_out_executor); - } -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - auto handle_ascend = reinterpret_cast(handle); - auto result = CausalSoftmaxInfo::create(y_desc, x_desc); - CHECK_RESULT(result); - CausalSoftmaxInfo info = result.take(); - - aclOpExecutor *executor = nullptr; - aclOpExecutor *temp_executor = nullptr; - aclOpExecutor *mask_executor = nullptr; - aclOpExecutor *copy_in_executor = nullptr; - aclOpExecutor *copy_out_executor = nullptr; - aclnnTensorDescriptor_t y = nullptr; - aclnnTensorDescriptor_t mask = nullptr; - aclnnTensorDescriptor_t x = nullptr; - aclnnTensorDescriptor_t value = nullptr; - aclnnTensorDescriptor_t temp_x = nullptr; - aclnnTensorDescriptor_t temp_y = nullptr; - void *mask_addr = nullptr; - void *value_addr = nullptr; - void *temp_x_addr = nullptr; - void *temp_y_addr = nullptr; - size_t workspacesize_softmax = 0; - size_t workspacesize_temp_softmax = 0; - size_t workspacesize_mask = 0; - size_t workspacesize_copy_in = 0; - size_t workspacesize_copy_out = 0; - - // Create Aclnn Tensor Descriptors for input, mask and output - std::vector shape = {static_cast(info.batch_size), static_cast(info.seq_len), static_cast(info.total_seq_len)}; - std::vector x_strides = {static_cast(info.x_stride_b), static_cast(info.x_stride_i), static_cast(info.x_stride_j)}; - std::vector y_strides = {static_cast(info.y_stride_b), static_cast(info.y_stride_i), static_cast(info.y_stride_j)}; - std::vector compact_strides = {static_cast(info.seq_len * info.total_seq_len), static_cast(info.total_seq_len), 1}; - y = new aclnnTensorDescriptor(toAclDataType(info.dtype), shape, y_strides); - x = new aclnnTensorDescriptor(toAclDataType(info.dtype), shape, x_strides); - temp_x = new aclnnTensorDescriptor(toAclDataType(info.dtype), shape, compact_strides); - temp_y = new aclnnTensorDescriptor(toAclDataType(info.dtype), shape, compact_strides); - mask = new aclnnTensorDescriptor(aclDataType::ACL_BOOL, {static_cast(info.seq_len), static_cast(info.total_seq_len)}, {static_cast(info.total_seq_len), 1}); - - // Initialize the value tensor with -inf - if (info.dtype == INFINI_DTYPE_F16) { - uint16_t mask_value = 0xfc00; - auto size = aclDataTypeSize(aclDataType::ACL_FLOAT16); - CHECK_ACL(aclrtMalloc(&value_addr, size, ACL_MEM_MALLOC_HUGE_FIRST)); - CHECK_ACL(aclrtMemcpy(value_addr, size, &mask_value, size, ACL_MEMCPY_HOST_TO_DEVICE)); - value = new aclnnTensorDescriptor(aclDataType::ACL_FLOAT16, {}, {}); - } else { - uint32_t mask_value = 0xff800000; - auto size = aclDataTypeSize(aclDataType::ACL_FLOAT); - CHECK_ACL(aclrtMalloc(&value_addr, size, ACL_MEM_MALLOC_HUGE_FIRST)); - CHECK_ACL(aclrtMemcpy(value_addr, size, &mask_value, size, ACL_MEMCPY_HOST_TO_DEVICE)); - value = new aclnnTensorDescriptor(aclDataType::ACL_FLOAT, {}, {}); - } - - // Fill Mask Tensor - std::vector mask_matrix(mask->numel(), 0); - for (size_t i = 0; i < info.seq_len; ++i) { - for (size_t j = info.total_seq_len - info.seq_len + i + 1; j < info.total_seq_len; ++j) { - size_t index = i * info.total_seq_len + j; - mask_matrix[index] = 1; - } - } - - auto size = mask->numel() * aclDataTypeSize(aclDataType::ACL_BOOL); - CHECK_ACL(aclrtMalloc(&mask_addr, size, ACL_MEM_MALLOC_HUGE_FIRST)); - CHECK_ACL(aclrtMemcpy(mask_addr, size, mask_matrix.data(), size, ACL_MEMCPY_HOST_TO_DEVICE)); - - // Get the workspace size for the op - aclTensor *tx = x->tensor; - aclTensor *ty = y->tensor; - aclTensor *ttemp_x = temp_x->tensor; - aclTensor *ttemp_y = temp_y->tensor; - aclTensor *tmask = mask->tensor; - aclTensor *tvalue = value->tensor; - - bool use_temp = !isCompact(info, info.x_stride_b, info.x_stride_i, info.x_stride_j) - || !isCompact(info, info.y_stride_b, info.y_stride_i, info.y_stride_j); - - if (use_temp) { - CHECK_ACL(aclnnInplaceCopyGetWorkspaceSize(ttemp_x, tx, &workspacesize_copy_in, ©_in_executor)); - aclSetAclOpExecutorRepeatable(copy_in_executor); - CHECK_ACL(aclnnInplaceCopyGetWorkspaceSize(ty, ttemp_y, &workspacesize_copy_out, ©_out_executor)); - aclSetAclOpExecutorRepeatable(copy_out_executor); - CHECK_ACL(aclnnInplaceMaskedFillTensorGetWorkspaceSize(ttemp_x, tmask, tvalue, &workspacesize_mask, &mask_executor)); - int64_t dim = 2; - CHECK_ACL(aclnnSoftmaxGetWorkspaceSize(ttemp_x, dim, ttemp_y, &workspacesize_temp_softmax, &temp_executor)); - aclSetAclOpExecutorRepeatable(temp_executor); - } else { - CHECK_ACL(aclnnInplaceMaskedFillTensorGetWorkspaceSize(tx, tmask, tvalue, &workspacesize_mask, &mask_executor)); - int64_t dim = 2; - CHECK_ACL(aclnnSoftmaxGetWorkspaceSize(tx, dim, ty, &workspacesize_softmax, &executor)); - // set executor reusable - aclSetAclOpExecutorRepeatable(executor); - } - - size_t op_workspace_size = std::max(std::max(workspacesize_softmax, workspacesize_temp_softmax), - std::max(workspacesize_mask, std::max(workspacesize_copy_in, workspacesize_copy_out))); - size_t all_workspacesize = op_workspace_size; - if (use_temp) { - size_t temp_bytes = temp_x->numel() * infiniSizeOf(info.dtype); - CHECK_ACL(aclrtMalloc(&temp_x_addr, temp_bytes, ACL_MEM_MALLOC_HUGE_FIRST)); - CHECK_ACL(aclrtMalloc(&temp_y_addr, temp_bytes, ACL_MEM_MALLOC_HUGE_FIRST)); - } - - *desc_ptr = new Descriptor(new Opaque{x, mask, y, value, temp_x, temp_y, mask_addr, value_addr, - temp_x_addr, temp_y_addr, op_workspace_size, executor, temp_executor, copy_in_executor, copy_out_executor, use_temp}, - std::move(info), all_workspacesize, handle_ascend->device, handle_ascend->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, void *y, const void *x, void *stream) const { - if (workspace_size < workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - auto tx = _opaque->x->tensor; - auto ty = _opaque->y->tensor; - auto tmask = _opaque->mask->tensor; - auto tvalue = _opaque->value->tensor; - - if (_opaque->use_temp) { - auto ttemp_x = _opaque->temp_x->tensor; - auto ttemp_y = _opaque->temp_y->tensor; - void *temp_x = _opaque->temp_x_addr; - void *temp_y = _opaque->temp_y_addr; - - AclSetTensorAddr(_opaque->copy_in_executor, 0, ttemp_x, temp_x); - AclSetTensorAddr(_opaque->copy_in_executor, 1, tx, (void *)x); - CHECK_ACL(aclnnInplaceCopy(workspace, _opaque->workspacesize, _opaque->copy_in_executor, stream)); - - aclOpExecutor *mask_executor = nullptr; - size_t workspacesize_mask = 0; - AclSetTensorAddr(mask_executor, 0, ttemp_x, temp_x); - AclSetTensorAddr(mask_executor, 1, tmask, _opaque->mask_addr); - AclSetTensorAddr(mask_executor, 2, tvalue, _opaque->value_addr); - CHECK_ACL(aclnnInplaceMaskedFillTensorGetWorkspaceSize(ttemp_x, tmask, tvalue, &workspacesize_mask, &mask_executor)); - CHECK_ACL(aclnnInplaceMaskedFillTensor(workspace, _opaque->workspacesize, mask_executor, stream)); - - AclSetTensorAddr(_opaque->temp_executor, 0, ttemp_x, temp_x); - AclSetTensorAddr(_opaque->temp_executor, 1, ttemp_y, temp_y); - CHECK_ACL(aclnnSoftmax(workspace, _opaque->workspacesize, _opaque->temp_executor, stream)); - - AclSetTensorAddr(_opaque->copy_out_executor, 0, ty, y); - AclSetTensorAddr(_opaque->copy_out_executor, 1, ttemp_y, temp_y); - CHECK_ACL(aclnnInplaceCopy(workspace, _opaque->workspacesize, _opaque->copy_out_executor, stream)); - return INFINI_STATUS_SUCCESS; - } - - aclOpExecutor *mask_executor = nullptr; - size_t workspacesize_mask = 0; - - AclSetTensorAddr(mask_executor, 0, tx, (void *)x); - AclSetTensorAddr(mask_executor, 1, tmask, _opaque->mask_addr); - AclSetTensorAddr(mask_executor, 2, tvalue, _opaque->value_addr); - CHECK_ACL(aclnnInplaceMaskedFillTensorGetWorkspaceSize(tx, tmask, tvalue, &workspacesize_mask, &mask_executor)); - CHECK_ACL(aclnnInplaceMaskedFillTensor(workspace, _opaque->workspacesize, mask_executor, stream)); - - AclSetTensorAddr(_opaque->executor, 0, tx, (void *)x); - AclSetTensorAddr(_opaque->executor, 1, ty, y); - CHECK_ACL(aclnnSoftmax(workspace, _opaque->workspacesize, _opaque->executor, stream)); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::causal_softmax::ascend diff --git a/src/infiniop/ops/causal_softmax/ascend/causal_softmax_ascend.h b/src/infiniop/ops/causal_softmax/ascend/causal_softmax_ascend.h deleted file mode 100644 index 23f380fe0..000000000 --- a/src/infiniop/ops/causal_softmax/ascend/causal_softmax_ascend.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __CAUSAL_SOFTMAX_ASCEND_H__ -#define __CAUSAL_SOFTMAX_ASCEND_H__ -#include "../causal_softmax.h" - -DESCRIPTOR(ascend) - -#endif diff --git a/src/infiniop/ops/causal_softmax/bang/causal_softmax_bang.h b/src/infiniop/ops/causal_softmax/bang/causal_softmax_bang.h deleted file mode 100644 index 6d89024e3..000000000 --- a/src/infiniop/ops/causal_softmax/bang/causal_softmax_bang.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __CAUSAL_SOFTMAX_BANG_H__ -#define __CAUSAL_SOFTMAX_BANG_H__ -#include "../causal_softmax.h" - -DESCRIPTOR(bang) - -#endif // __CAUSAL_SOFTMAX_BANG_H__ diff --git a/src/infiniop/ops/causal_softmax/bang/causal_softmax_bang.mlu b/src/infiniop/ops/causal_softmax/bang/causal_softmax_bang.mlu deleted file mode 100644 index 498330ed7..000000000 --- a/src/infiniop/ops/causal_softmax/bang/causal_softmax_bang.mlu +++ /dev/null @@ -1,200 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "../../../reduce/bang/reduce_bang.h" -#include "causal_softmax_bang.h" - -__nram__ char nram_buffer[NRAM_MAX_SIZE]; -const int SRC_MAX_SIZE = NRAM_MAX_SIZE / 4; - -template -__mlu_func__ void processSoftmaxStep(T *output, const T *input, float scalar, int num_elements, int stride, bool is_exp_phase) { - // Calculate buffer sizes (split between float and T buffers) - constexpr bool is_half = std::is_same_v; - constexpr bool is_bfloat16 = std::is_same_v; - constexpr bool is_float = !is_half && !is_bfloat16; - - const int chunk_size = SRC_MAX_SIZE / ((is_half || is_bfloat16) ? (2 * sizeof(float)) : sizeof(float)); - float *float_buffer = (float *)nram_buffer; - T *temp_buffer = is_float ? nullptr : (T *)(nram_buffer + chunk_size * sizeof(float)); - - // Common stride configurations - const int src_stride = stride * sizeof(T); - const int dst_stride = stride * sizeof(T); - - int processed = 0; - while (processed < num_elements) { - int curr_batch = std::min(chunk_size, num_elements - processed); - - // Gather input elements using 2D memcpy - if constexpr (is_float) { - __memcpy(float_buffer, (is_exp_phase ? input : output) + processed * stride, sizeof(float), - GDRAM2NRAM, sizeof(float), src_stride, curr_batch - 1); - } else { - __memcpy(temp_buffer, (is_exp_phase ? input : output) + processed * stride, sizeof(T), - GDRAM2NRAM, sizeof(T), src_stride, curr_batch - 1); - - // Convert to float - if constexpr (is_half) { - __bang_half2float(float_buffer, temp_buffer, curr_batch); - } else if constexpr (is_bfloat16) { - __bang_bfloat162float(float_buffer, temp_buffer, curr_batch); - } - } - - // Common processing for all types - if (is_exp_phase) { - __bang_sub_scalar(float_buffer, float_buffer, scalar, curr_batch); // scalar is max_val - __bang_active_exphp(float_buffer, float_buffer, curr_batch); - } else { - __bang_mul_scalar(float_buffer, float_buffer, scalar, curr_batch); // scalar is 1.0f/sum_val - } - - // Convert back and scatter output using 2D memcpy - if constexpr (is_float) { - __memcpy(output + processed * stride, float_buffer, sizeof(float), - NRAM2GDRAM, dst_stride, sizeof(float), curr_batch - 1); - } else { - // Convert back to original type - if constexpr (is_half) { - __bang_float2half(temp_buffer, float_buffer, curr_batch); - } else if constexpr (is_bfloat16) { - __bang_float2bfloat16(temp_buffer, float_buffer, curr_batch); - } - - // Scatter output - __memcpy(output + processed * stride, temp_buffer, sizeof(T), - NRAM2GDRAM, dst_stride, sizeof(T), curr_batch - 1); - } - - processed += curr_batch; - } -} - -template -__mlu_global__ void causalSoftmax(T *y, const T *x, - size_t batch_size, size_t seq_len, size_t total_seq_len, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_i, ptrdiff_t y_stride_j, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_i, ptrdiff_t x_stride_j) { - using namespace op::common_bang::reduce_op; - - // Get task information - size_t task_id = taskId; - size_t task_num = taskDimX * taskDimY; - - // Calculate elements per task with better load balancing - size_t total_tasks = batch_size * seq_len; - size_t tasks_per_core = (total_tasks + task_num - 1) / task_num; - size_t start = task_id * tasks_per_core; - size_t end = std::min(start + tasks_per_core, total_tasks); - - // Allocate NRAM buffers - const int max_batch = SRC_MAX_SIZE / sizeof(T); - T *src = (T *)nram_buffer; - float *dst = (float *)(nram_buffer + max_batch * sizeof(T)); - - for (size_t index = start; index < end; index++) { - size_t batch = index / seq_len; - size_t i = (index % seq_len); - ptrdiff_t y_offset = batch * y_stride_b + i * y_stride_i; - ptrdiff_t x_offset = batch * x_stride_b + i * x_stride_i; - T *y_ = y + y_offset; - const T *x_ = x + x_offset; - - // Calculate the valid sequence length for this position - size_t valid_len = total_seq_len - seq_len + i + 1; - - // Zero out future positions - for (size_t j = valid_len; j < total_seq_len; j++) { - y_[j * y_stride_j] = (T)0.0f; - } - - // Calculate max value using optimized reduction - float max_val = maxBatched(x_, src, dst, valid_len, max_batch); - - // Compute exp(x - max) - processSoftmaxStep(y_, x_, max_val, valid_len, x_stride_j, true); - - // Calculate sum of exponentials - float sum_val = sumBatched(y_, src, dst, valid_len, max_batch); - - // Normalize by sum - processSoftmaxStep(y_, y_, 1.0f / sum_val, valid_len, y_stride_j, false); - } -} - -template -void causalSoftmaxUnion(void *workspace, int core_per_cluster, int cluster_count, - cnrtQueue_t queue, void *y, const void *x, const op::causal_softmax::CausalSoftmaxInfo *info) { - cnrtDim3_t kernel_dim; - cnrtFunctionType_t kernel_type; - - // Configure kernel dimensions - kernel_dim.x = core_per_cluster; - kernel_dim.y = cluster_count; - kernel_dim.z = 1; - kernel_type = cnrtFuncTypeUnion1; - - // Launch kernel - causalSoftmax<<>>( - (T *)y, (const T *)x, - info->batch_size, info->seq_len, info->total_seq_len, - info->y_stride_b, info->y_stride_i, info->y_stride_j, - info->x_stride_b, info->x_stride_i, info->x_stride_j); - - cnrtQueueSync(queue); -} - -namespace op::causal_softmax::bang { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - auto handle = reinterpret_cast(handle_); - - auto result = CausalSoftmaxInfo::create(y_desc, x_desc); - CHECK_RESULT(result); - auto info = result.take(); - - *desc_ptr = new Descriptor( - new Descriptor::Opaque{static_cast(handle)->internal()}, - info, - 0, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *y, - const void *x, - void *stream) const { - auto queue = reinterpret_cast(stream); - int core_per_cluster = _opaque->internal->getCorePerCluster(); - int cluster_count = _opaque->internal->getClusterCount(); - - // Dispatch based on data type - if (_info.dtype == INFINI_DTYPE_F16) { - causalSoftmaxUnion(workspace, core_per_cluster, cluster_count, queue, y, x, &_info); - } else if (_info.dtype == INFINI_DTYPE_BF16) { - causalSoftmaxUnion(workspace, core_per_cluster, cluster_count, queue, y, x, &_info); - } else if (_info.dtype == INFINI_DTYPE_F32) { - causalSoftmaxUnion(workspace, core_per_cluster, cluster_count, queue, y, x, &_info); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::causal_softmax::bang diff --git a/src/infiniop/ops/causal_softmax/causal_softmax.h b/src/infiniop/ops/causal_softmax/causal_softmax.h deleted file mode 100644 index a31e99455..000000000 --- a/src/infiniop/ops/causal_softmax/causal_softmax.h +++ /dev/null @@ -1,46 +0,0 @@ -#ifndef CAUSAL_SOFTMAX_H -#define CAUSAL_SOFTMAX_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::causal_softmax::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - CausalSoftmaxInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - CausalSoftmaxInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t y_desc, \ - infiniopTensorDescriptor_t x_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *y, \ - const void *x, \ - void *stream) const; \ - }; \ - } - -#endif // CAUSAL_SOFTMAX_H diff --git a/src/infiniop/ops/causal_softmax/cpu/causal_softmax_cpu.cc b/src/infiniop/ops/causal_softmax/cpu/causal_softmax_cpu.cc deleted file mode 100644 index 540c9f43c..000000000 --- a/src/infiniop/ops/causal_softmax/cpu/causal_softmax_cpu.cc +++ /dev/null @@ -1,77 +0,0 @@ -#include "causal_softmax_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include "../../../reduce/cpu/reduce.h" - -namespace op::causal_softmax::cpu { - -Descriptor::~Descriptor() {} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - auto result = CausalSoftmaxInfo::create(y_desc, x_desc); - CHECK_RESULT(result); - *desc_ptr = new Descriptor(nullptr, result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t causal_softmax(const CausalSoftmaxInfo *info, T *y, const T *x) { -#pragma omp parallel for - for (ptrdiff_t index = 0; index < ptrdiff_t(info->batch_size * info->seq_len); index++) { - size_t batch = index / info->seq_len; - size_t i = (index % info->seq_len); - ptrdiff_t y_offset = batch * info->y_stride_b + i * info->y_stride_i; - ptrdiff_t x_offset = batch * info->x_stride_b + i * info->x_stride_i; - T *y_ = y + y_offset; - const T *x_ = x + x_offset; - - for (size_t j = info->total_seq_len - info->seq_len + i + 1; j < info->total_seq_len; j++) { - if constexpr (std::is_same::value || std::is_same::value) { - y_[j * info->y_stride_j] = utils::cast(0.0f); - } else { - y_[j * info->y_stride_j] = 0.0f; - } - } - float val = op::common_cpu::reduce_op::max(x_, info->total_seq_len - info->seq_len + i + 1, info->x_stride_j); - for (size_t j = 0; j <= info->total_seq_len - info->seq_len + i; j++) { - if constexpr (std::is_same::value || std::is_same::value) { - y_[j * info->y_stride_j] = utils::cast(std::exp(utils::cast(x_[j * info->x_stride_j]) - val)); - } else { - y_[j * info->y_stride_j] = std::exp(x_[j * info->x_stride_j] - val); - } - } - float sum = op::common_cpu::reduce_op::sum(y_, info->total_seq_len - info->seq_len + i + 1, info->y_stride_j); - for (size_t j = 0; j <= info->total_seq_len - info->seq_len + i; j++) { - if constexpr (std::is_same::value || std::is_same::value) { - y_[j * info->y_stride_j] = utils::cast(utils::cast(y_[j * info->y_stride_j]) / sum); - } else { - y_[j * info->y_stride_j] = y_[j * info->y_stride_j] / sum; - } - } - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - if (_info.dtype == INFINI_DTYPE_F16) { - CHECK_STATUS(causal_softmax(&_info, (fp16_t *)y, (const fp16_t *)x)); - } else if (_info.dtype == INFINI_DTYPE_BF16) { - CHECK_STATUS(causal_softmax(&_info, (bf16_t *)y, (const bf16_t *)x)); - } else if (_info.dtype == INFINI_DTYPE_F32) { - CHECK_STATUS(causal_softmax(&_info, (float *)y, (const float *)x)); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::causal_softmax::cpu diff --git a/src/infiniop/ops/causal_softmax/cpu/causal_softmax_cpu.h b/src/infiniop/ops/causal_softmax/cpu/causal_softmax_cpu.h deleted file mode 100644 index e2bf98096..000000000 --- a/src/infiniop/ops/causal_softmax/cpu/causal_softmax_cpu.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __CAUSAL_SOFTMAX_CPU_H__ -#define __CAUSAL_SOFTMAX_CPU_H__ -#include "../causal_softmax.h" - -DESCRIPTOR(cpu) - -#endif diff --git a/src/infiniop/ops/causal_softmax/cuda/kernel.cuh b/src/infiniop/ops/causal_softmax/cuda/kernel.cuh deleted file mode 100644 index 0808a9cd3..000000000 --- a/src/infiniop/ops/causal_softmax/cuda/kernel.cuh +++ /dev/null @@ -1,57 +0,0 @@ -#ifndef __CAUSAL_SOFTMAX_KERNEL_CUH__ -#define __CAUSAL_SOFTMAX_KERNEL_CUH__ - -template -__device__ void causalSoftmaxKernel( - Tdata *y_, const Tdata *x_, - size_t batch, size_t height, size_t width, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_h, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_h) { - - Tdata *y = y_ // threadIdx.x for col_id - + blockIdx.y * y_stride_b // gridDim.y for batch_id - + blockIdx.x * y_stride_h; // gridDim.x for row_id - const Tdata *x = x_ + blockIdx.y * x_stride_b + blockIdx.x * x_stride_h; - - // [Reduce] Find max value in each row and store in shared memory - __shared__ Tdata max_; - Tdata max_0 = op::common_cuda::reduce_op::max(x, width - height + 1 + blockIdx.x); - if (threadIdx.x == 0) { - max_ = max_0; - } - __syncthreads(); - - // [Elementwise] Subtract max value from each element and apply causal mask - for (size_t col = threadIdx.x; col < width; col += BLOCK_SIZE) { - // row_id ↓ |<- width ->| - // 0 | * * * ... * | - // 1 | * * * ... * * | - // 2 | * * * ... * * * | - // height: 3 col_id-> - if (width + blockIdx.x >= col + height) { - if constexpr (std::is_same_v || std::is_same_v) { - y[col] = hexp(x[col] - max_); - } else { - y[col] = exp(x[col] - max_); - } - } else { - y[col] = Tdata(0); - } - } - __syncthreads(); - - // [Reduce] Find the sum of each updated row and store in shared memory - __shared__ Tcompute sum_; - Tcompute sum_0 = op::common_cuda::reduce_op::sum(y, width); - if (threadIdx.x == 0) { - sum_ = sum_0; - } - __syncthreads(); - - // [Elementwise] Divide each element by the sum and store in shared memory - for (size_t col = threadIdx.x; col < width; col += BLOCK_SIZE) { - y[col] /= Tdata(sum_); - } -} - -#endif // __CAUSAL_SOFTMAX_KERNEL_CUH__ diff --git a/src/infiniop/ops/causal_softmax/info.h b/src/infiniop/ops/causal_softmax/info.h deleted file mode 100644 index 270642833..000000000 --- a/src/infiniop/ops/causal_softmax/info.h +++ /dev/null @@ -1,78 +0,0 @@ -#ifndef __CAUSAL_SOFTMAX_INFO_H__ -#define __CAUSAL_SOFTMAX_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::causal_softmax { - -class CausalSoftmaxInfo { - CausalSoftmaxInfo() = default; - -public: - infiniDtype_t dtype; - size_t batch_size; - size_t seq_len; - size_t total_seq_len; - - ptrdiff_t y_stride_b; - ptrdiff_t y_stride_i; - ptrdiff_t y_stride_j; - - ptrdiff_t x_stride_b; - ptrdiff_t x_stride_i; - ptrdiff_t x_stride_j; - - static utils::Result create(infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t x_desc) { - auto dtype = y_desc->dtype(); - if (dtype != x_desc->dtype()) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - - auto shape = y_desc->shape(); - CHECK_SAME_SHAPE(shape, x_desc->shape()); - - auto ndim = y_desc->ndim(); - if (ndim != 2 && ndim != 3) { - CHECK_STATUS(INFINI_STATUS_BAD_TENSOR_SHAPE); - } - - if (shape[ndim - 1] < shape[ndim - 2]) { - CHECK_STATUS(INFINI_STATUS_BAD_TENSOR_SHAPE); - } - - size_t batch_size = 1; - size_t seq_len = shape[ndim - 2]; - size_t total_seq_len = shape[ndim - 1]; - ptrdiff_t y_stride_b = 0, - y_stride_i = y_desc->stride(ndim - 2), - y_stride_j = y_desc->stride(ndim - 1); - ptrdiff_t x_stride_b = 0, - x_stride_i = x_desc->stride(ndim - 2), - x_stride_j = x_desc->stride(ndim - 1); - - if (ndim == 3) { - y_stride_b = y_desc->stride(0); - x_stride_b = x_desc->stride(0); - batch_size = shape[0]; - } - - return utils::Result(CausalSoftmaxInfo{ - dtype, - batch_size, - seq_len, - total_seq_len, - y_stride_b, - y_stride_i, - y_stride_j, - x_stride_b, - x_stride_i, - x_stride_j}); - } -}; - -} // namespace op::causal_softmax - -#endif // __CAUSAL_SOFTMAX_INFO_H__ diff --git a/src/infiniop/ops/causal_softmax/kunlun/causal_softmax_kunlun.h b/src/infiniop/ops/causal_softmax/kunlun/causal_softmax_kunlun.h deleted file mode 100644 index ef7e4654b..000000000 --- a/src/infiniop/ops/causal_softmax/kunlun/causal_softmax_kunlun.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CAUSAL_SOFTMAX_KUNLUN_H__ -#define __CAUSAL_SOFTMAX_KUNLUN_H__ - -#include "../causal_softmax.h" - -DESCRIPTOR(kunlun) - -#endif diff --git a/src/infiniop/ops/causal_softmax/kunlun/causal_softmax_kunlun.xpu b/src/infiniop/ops/causal_softmax/kunlun/causal_softmax_kunlun.xpu deleted file mode 100644 index 58324d162..000000000 --- a/src/infiniop/ops/causal_softmax/kunlun/causal_softmax_kunlun.xpu +++ /dev/null @@ -1,112 +0,0 @@ -#include "../../../devices/kunlun/kunlun_common.h" -#include "../../../devices/kunlun/kunlun_handle.h" -#include "../../../devices/kunlun/kunlun_kernel_common.h" -#include "causal_softmax_kunlun.h" -#include "kernel.h" - -template -__global__ void causalSoftmaxKernel( - Tdata *y, - const Tdata *x, - uint32_t batch, - uint32_t height, - uint32_t width, - int32_t y_stride_h, - int32_t x_stride_h) { - - __shared__ Tdata x_sm[SM_SIZE / sizeof(Tdata)]; - __shared__ Tdata y_sm[SM_SIZE / sizeof(Tdata)]; - - for (uint32_t row_id = cluster_id(); row_id < height; row_id += cluster_num()) { - __global_ptr__ Tdata *y_ = y + row_id * y_stride_h; - __global_ptr__ const Tdata *x_ = x + row_id * x_stride_h; - - if (core_id() == 0) { - GM2SM_ASYNC(x_, x_sm, width * sizeof(Tdata)); - } - sync_cluster(); - - causalSoftmaxBlock(y_sm, x_sm, height, width, row_id); - - if (core_id() == 0) { - SM2GM_ASYNC(y_sm, y_, width * sizeof(Tdata)); - } - sync_cluster(); - } -} - -namespace op::causal_softmax::kunlun { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - auto info = CausalSoftmaxInfo::create(y_desc, x_desc); - CHECK_RESULT(info); - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(void *y, const void *x, infiniDtype_t dtype, - size_t batch_size, size_t seq_len, size_t total_seq_len, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_h, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_h, - kunlunStream_t stream) { - - // Kunlunxin kernel dont support ptrdiff_t and size_t as parameters - uint32_t batch_size_ = static_cast(batch_size); - uint32_t seq_len_ = static_cast(seq_len); - uint32_t total_seq_len_ = static_cast(total_seq_len); - int32_t y_stride_b_ = static_cast(y_stride_b); - int32_t y_stride_h_ = static_cast(y_stride_h); - int32_t x_stride_b_ = static_cast(x_stride_b); - int32_t x_stride_h_ = static_cast(x_stride_h); - -#define LAUCH_KERNEL(Tdata, Tcompute) \ - for (uint32_t i = 0; i < batch_size_; ++i) { \ - causalSoftmaxKernel \ - <<<8, BLOCK_SIZE, stream>>>((Tdata *)y + i * y_stride_b_, (const Tdata *)x + i * x_stride_b_, \ - batch_size_, seq_len_, total_seq_len_, \ - y_stride_h_, x_stride_h_); \ - } - - if (dtype == INFINI_DTYPE_F16) { - LAUCH_KERNEL(half, float); - } else if (dtype == INFINI_DTYPE_BF16) { - LAUCH_KERNEL(bfloat16_t, float); - } else if (dtype == INFINI_DTYPE_F32) { - LAUCH_KERNEL(float, float); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; - -#undef LAUCH_KERNEL -} - -infiniStatus_t Descriptor::calculate(void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream_) const { - kunlunStream_t stream = (kunlunStream_t)stream_; - CHECK_STATUS(launchKernel<64>( - y, x, _info.dtype, _info.batch_size, _info.seq_len, _info.total_seq_len, - _info.y_stride_b, _info.y_stride_i, _info.x_stride_b, _info.x_stride_i, stream)); - xpu_wait(stream); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::causal_softmax::kunlun diff --git a/src/infiniop/ops/causal_softmax/kunlun/kernel.h b/src/infiniop/ops/causal_softmax/kunlun/kernel.h deleted file mode 100644 index 00664cf20..000000000 --- a/src/infiniop/ops/causal_softmax/kunlun/kernel.h +++ /dev/null @@ -1,73 +0,0 @@ -#ifndef __CAUSAL_SOFTMAX_KUNLUN_KERNEL_H__ -#define __CAUSAL_SOFTMAX_KUNLUN_KERNEL_H__ - -#include "../../../devices/kunlun/kunlun_kernel_common.h" -#include "../../../reduce/kunlun/reduce_kunlun.h" - -using namespace device::kunlun::kernel; - -template -__device__ void causalSoftmaxBlock( - __shared_ptr__ Tdata *y, - __shared_ptr__ const Tdata *x, - size_t height, - size_t width, - int row_id) { - - // Reduce max for each row and store in shared memory - __shared__ Tdata max_; - Tdata max_0 = op::common_kunlun::reduce_op::max(x, width - height + 1 + size_t(row_id)); - if (core_id() == 0) { - storeShared(&max_, max_0); - } - sync_cluster(); - - Tdata max_val = loadShared(&max_); - - // Elemetwise sub max for each element and apply causal softmax - for (size_t col = core_id(); col < width; col += BLOCK_SIZE) { - // row_id ↓ |<- width ->| - // 0 | * * * ... * | - // 1 | * * * ... * * | - // 2 | * * * ... * * * | - // height: 3 col_id-> - if (width + size_t(row_id) >= col + height) { - Tdata x_val = loadShared(x + col); - Tdata y_val; - if constexpr (std::is_same_v) { - y_val = hexp(x_val - max_val); - } else if constexpr (std::is_same_v) { - y_val = __float2bfloat16(exp(__bfloat162float(x_val) - __bfloat162float(max_val))); - } else { - y_val = exp(x_val - max_val); - } - storeShared(y + col, y_val); - } else { - storeShared(y + col, Tdata(0)); - } - } - sync_cluster(); - - // Reduce sum for each row - __shared__ Tcompute sum_; - Tcompute sum_0 = op::common_kunlun::reduce_op::sum(y, width); - if (core_id() == 0) { - storeShared(&sum_, sum_0); - } - sync_cluster(); - - Tcompute sum_val = loadShared(&sum_); - - // Apply softmax - for (size_t col = core_id(); col < width; col += BLOCK_SIZE) { - if (sum_val != 0) { - Tdata y_val = loadShared(y + col); - storeShared(y + col, Tdata(Tcompute(y_val) / sum_val)); - } else { - storeShared(y + col, Tdata(0)); - } - } - sync_cluster(); -} - -#endif diff --git a/src/infiniop/ops/causal_softmax/metax/causal_softmax_metax.h b/src/infiniop/ops/causal_softmax/metax/causal_softmax_metax.h deleted file mode 100644 index 1b68a8f5f..000000000 --- a/src/infiniop/ops/causal_softmax/metax/causal_softmax_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CAUSAL_SOFTMAX_METAX_H__ -#define __CAUSAL_SOFTMAX_METAX_H__ - -#include "../causal_softmax.h" - -DESCRIPTOR(metax) - -#endif diff --git a/src/infiniop/ops/causal_softmax/metax/causal_softmax_metax.maca b/src/infiniop/ops/causal_softmax/metax/causal_softmax_metax.maca deleted file mode 100644 index cde4009fe..000000000 --- a/src/infiniop/ops/causal_softmax/metax/causal_softmax_metax.maca +++ /dev/null @@ -1,97 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "causal_softmax_metax.h" - -#ifdef ENABLE_METAX_MC_API - #include -#else - #include -#endif -#include "../../../devices/metax/metax_kernel_common.h" - -#include "../../../reduce/cuda/reduce.cuh" - -#include "../cuda/kernel.cuh" - -template -INFINIOP_METAX_KERNEL causalSoftmax( - Tdata *y, const Tdata *x, - size_t batch, size_t height, size_t width, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_h, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_h) { - causalSoftmaxKernel(y, x, batch, height, width, y_stride_b, y_stride_h, x_stride_b, x_stride_h); -} - -namespace op::causal_softmax::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - auto info = CausalSoftmaxInfo::create(y_desc, x_desc); - CHECK_RESULT(info); - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(void *y, const void *x, infiniDtype_t dtype, - size_t batch_size, size_t seq_len, size_t total_seq_len, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_i, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_i, - hcStream_t stream) { - dim3 grid(uint32_t(seq_len), uint32_t(batch_size), 1); - if (dtype == INFINI_DTYPE_F16) { - causalSoftmax - <<>>((half *)y, (const half *)x, - batch_size, seq_len, total_seq_len, - y_stride_b, y_stride_i, - x_stride_b, x_stride_i); - } else if (dtype == INFINI_DTYPE_BF16) { - causalSoftmax - <<>>((__hpcc_bfloat16 *)y, (const __hpcc_bfloat16 *)x, - batch_size, seq_len, total_seq_len, - y_stride_b, y_stride_i, - x_stride_b, x_stride_i); - } else if (dtype == INFINI_DTYPE_F32) { - causalSoftmax - <<>>((float *)y, (const float *)x, - batch_size, seq_len, total_seq_len, - y_stride_b, y_stride_i, - x_stride_b, x_stride_i); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *y, - const void *x, - void *stream_) const { - hcStream_t stream = (hcStream_t)stream_; - if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.batch_size, _info.seq_len, _info.total_seq_len, - _info.y_stride_b, _info.y_stride_i, _info.x_stride_b, _info.x_stride_i, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.batch_size, _info.seq_len, _info.total_seq_len, - _info.y_stride_b, _info.y_stride_i, _info.x_stride_b, _info.x_stride_i, stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::causal_softmax::metax diff --git a/src/infiniop/ops/causal_softmax/moore/causal_softmax_kernel.h b/src/infiniop/ops/causal_softmax/moore/causal_softmax_kernel.h deleted file mode 100644 index 111e7cd4c..000000000 --- a/src/infiniop/ops/causal_softmax/moore/causal_softmax_kernel.h +++ /dev/null @@ -1,80 +0,0 @@ -#ifndef __CAUSAL_SOFTMAX_KERNEL_CUH__ -#define __CAUSAL_SOFTMAX_KERNEL_CUH__ - -template -__device__ void causalSoftmaxKernel( - Tdata *y_, const Tdata *x_, - size_t batch, size_t height, size_t width, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_h, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_h) { - - Tdata *y = y_ // threadIdx.x for col_id - + blockIdx.y * y_stride_b // gridDim.y for batch_id - + blockIdx.x * y_stride_h; // gridDim.x for row_id - const Tdata *x = x_ + blockIdx.y * x_stride_b + blockIdx.x * x_stride_h; - - // [Reduce] Find max value in each row and store in shared memory - __shared__ Tdata max_; - Tdata max_0 = op::common_cuda::reduce_op::max(x, width - height + 1 + blockIdx.x); - if (threadIdx.x == 0) { - max_ = max_0; - } - __syncthreads(); - - // [Elementwise] Subtract max value from each element and apply causal mask - for (size_t col = threadIdx.x; col < width; col += BLOCK_SIZE) { - // row_id ↓ |<- width ->| - // 0 | * * * ... * | - // 1 | * * * ... * * | - // 2 | * * * ... * * * | - // height: 3 col_id-> - if (width + blockIdx.x >= col + height) { - if constexpr (std::is_same_v || std::is_same_v) { - /* - * MUSA does not support CUDA's native `hexp` function. - * This code performs an explicit conversion: - * it casts the input to `float`, computes the exponential, and casts the result back. - * This ensures compatibility and correct behavior on the MUSA platform. - */ - float val = static_cast(x[col]) - static_cast(max_); - y[col] = static_cast(expf(val)); - } else { - y[col] = exp(x[col] - max_); - } - } else { - /* - * In MUSA, the `__mt_bfloat16` type has ambiguous constructors for integer literals (e.g., `0`), - * as it could be implicitly converted from either `float` or `double`. - * - * This differs from CUDA's `half` type, which can typically be initialized - * from integer literals without ambiguity. - * - * To resolve this, we use the float literal `0.0f` to explicitly - * specify the conversion path, ensuring platform compatibility. - */ - y[col] = Tdata(0.0f); - } - } - __syncthreads(); - - // [Reduce] Find the sum of each updated row and store in shared memory - __shared__ Tcompute sum_; - Tcompute sum_0 = op::common_cuda::reduce_op::sum(y, width); - if (threadIdx.x == 0) { - sum_ = sum_0; - } - __syncthreads(); - - // [Elementwise] Divide each element by the sum and store in shared memory - for (size_t col = threadIdx.x; col < width; col += BLOCK_SIZE) { - /* - * MUSA's bfloat16 type does not have a viable overloaded `/=` operator for float division. - * This change explicitly casts both operands to `float` before division, - * and then casts the result back to `Tdata`. - * This ensures the operation is performed correctly and avoids compilation errors. - */ - y[col] = static_cast(static_cast(y[col]) / static_cast(sum_)); - } -} - -#endif // __CAUSAL_SOFTMAX_KERNEL_CUH__ diff --git a/src/infiniop/ops/causal_softmax/moore/causal_softmax_moore.h b/src/infiniop/ops/causal_softmax/moore/causal_softmax_moore.h deleted file mode 100644 index 2c1763d98..000000000 --- a/src/infiniop/ops/causal_softmax/moore/causal_softmax_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CAUSAL_SOFTMAX_MOORE_H__ -#define __CAUSAL_SOFTMAX_MOORE_H__ - -#include "../causal_softmax.h" - -DESCRIPTOR(moore) - -#endif diff --git a/src/infiniop/ops/causal_softmax/moore/causal_softmax_moore.mu b/src/infiniop/ops/causal_softmax/moore/causal_softmax_moore.mu deleted file mode 100644 index 7547fdaa7..000000000 --- a/src/infiniop/ops/causal_softmax/moore/causal_softmax_moore.mu +++ /dev/null @@ -1,93 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "causal_softmax_moore.h" - -#include -#include "../../../devices/moore/moore_kernel_common.h" - -#include "../../../reduce/cuda/reduce.cuh" - -#include "causal_softmax_kernel.h" - -template -INFINIOP_MOORE_KERNEL causalSoftmax( - Tdata *y, const Tdata *x, - size_t batch, size_t height, size_t width, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_h, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_h) { - causalSoftmaxKernel(y, x, batch, height, width, y_stride_b, y_stride_h, x_stride_b, x_stride_h); -} - -namespace op::causal_softmax::moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - auto info = CausalSoftmaxInfo::create(y_desc, x_desc); - CHECK_RESULT(info); - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(void *y, const void *x, infiniDtype_t dtype, - size_t batch_size, size_t seq_len, size_t total_seq_len, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_i, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_i, - musaStream_t stream) { - dim3 grid(uint32_t(seq_len), uint32_t(batch_size), 1); - if (dtype == INFINI_DTYPE_F16) { - causalSoftmax - <<>>((half *)y, (const half *)x, - batch_size, seq_len, total_seq_len, - y_stride_b, y_stride_i, - x_stride_b, x_stride_i); - } else if (dtype == INFINI_DTYPE_BF16) { - causalSoftmax - <<>>((__mt_bfloat16 *)y, (const __mt_bfloat16 *)x, - batch_size, seq_len, total_seq_len, - y_stride_b, y_stride_i, - x_stride_b, x_stride_i); - } else if (dtype == INFINI_DTYPE_F32) { - causalSoftmax - <<>>((float *)y, (const float *)x, - batch_size, seq_len, total_seq_len, - y_stride_b, y_stride_i, - x_stride_b, x_stride_i); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *y, - const void *x, - void *stream_) const { - musaStream_t stream = (musaStream_t)stream_; - if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.batch_size, _info.seq_len, _info.total_seq_len, - _info.y_stride_b, _info.y_stride_i, _info.x_stride_b, _info.x_stride_i, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.batch_size, _info.seq_len, _info.total_seq_len, - _info.y_stride_b, _info.y_stride_i, _info.x_stride_b, _info.x_stride_i, stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::causal_softmax::moore diff --git a/src/infiniop/ops/causal_softmax/nvidia/causal_softmax_nvidia.cu b/src/infiniop/ops/causal_softmax/nvidia/causal_softmax_nvidia.cu deleted file mode 100644 index 6e671df1b..000000000 --- a/src/infiniop/ops/causal_softmax/nvidia/causal_softmax_nvidia.cu +++ /dev/null @@ -1,101 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "causal_softmax_nvidia.cuh" - -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include - -#include "../../../reduce/cuda/reduce.cuh" - -#include "../cuda/kernel.cuh" - -template -INFINIOP_CUDA_KERNEL causalSoftmax( - Tdata *y, const Tdata *x, - size_t batch, size_t height, size_t width, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_h, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_h) { - causalSoftmaxKernel(y, x, batch, height, width, y_stride_b, y_stride_h, x_stride_b, x_stride_h); -} - -namespace op::causal_softmax::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - auto info = CausalSoftmaxInfo::create(y_desc, x_desc); - CHECK_RESULT(info); - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(void *y, const void *x, infiniDtype_t dtype, - size_t batch_size, size_t seq_len, size_t total_seq_len, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_i, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_i, - cudaStream_t stream) { - dim3 grid(uint32_t(seq_len), uint32_t(batch_size), 1); - if (dtype == INFINI_DTYPE_F16) { - causalSoftmax - <<>>((half *)y, (const half *)x, - batch_size, seq_len, total_seq_len, - y_stride_b, y_stride_i, - x_stride_b, x_stride_i); - } else if (dtype == INFINI_DTYPE_BF16) { - causalSoftmax - <<>>((__nv_bfloat16 *)y, (const __nv_bfloat16 *)x, - batch_size, seq_len, total_seq_len, - y_stride_b, y_stride_i, - x_stride_b, x_stride_i); - } else if (dtype == INFINI_DTYPE_F32) { - causalSoftmax - <<>>((float *)y, (const float *)x, - batch_size, seq_len, total_seq_len, - y_stride_b, y_stride_i, - x_stride_b, x_stride_i); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *y, - const void *x, - void *stream_) const { - cudaStream_t stream = (cudaStream_t)stream_; - if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.batch_size, _info.seq_len, _info.total_seq_len, - _info.y_stride_b, _info.y_stride_i, _info.x_stride_b, _info.x_stride_i, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_2048) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.batch_size, _info.seq_len, _info.total_seq_len, - _info.y_stride_b, _info.y_stride_i, _info.x_stride_b, _info.x_stride_i, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.batch_size, _info.seq_len, _info.total_seq_len, - _info.y_stride_b, _info.y_stride_i, _info.x_stride_b, _info.x_stride_i, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - y, x, _info.dtype, _info.batch_size, _info.seq_len, _info.total_seq_len, - _info.y_stride_b, _info.y_stride_i, _info.x_stride_b, _info.x_stride_i, stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::causal_softmax::nvidia diff --git a/src/infiniop/ops/causal_softmax/nvidia/causal_softmax_nvidia.cuh b/src/infiniop/ops/causal_softmax/nvidia/causal_softmax_nvidia.cuh deleted file mode 100644 index 0362c43b1..000000000 --- a/src/infiniop/ops/causal_softmax/nvidia/causal_softmax_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CAUSAL_SOFTMAX_NVIDIA_H__ -#define __CAUSAL_SOFTMAX_NVIDIA_H__ - -#include "../causal_softmax.h" - -DESCRIPTOR(nvidia) - -#endif diff --git a/src/infiniop/ops/causal_softmax/operator.cc b/src/infiniop/ops/causal_softmax/operator.cc deleted file mode 100644 index f8d2bfba6..000000000 --- a/src/infiniop/ops/causal_softmax/operator.cc +++ /dev/null @@ -1,225 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/causal_softmax.h" - -#ifdef ENABLE_CPU_API -#include "cpu/causal_softmax_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) -#include "nvidia/causal_softmax_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/causal_softmax_metax.h" -#endif -#ifdef ENABLE_ASCEND_API -#include "ascend/causal_softmax_ascend.h" -#endif -#ifdef ENABLE_CAMBRICON_API -#include "bang/causal_softmax_bang.h" -#endif -#ifdef ENABLE_KUNLUN_API -#include "kunlun/causal_softmax_kunlun.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/causal_softmax_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateCausalSoftmaxDescriptor( - infiniopHandle_t handle, - infiniopCausalSoftmaxDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::causal_softmax::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - x_desc); - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang) -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_ASCEND_API - CREATE(INFINI_DEVICE_ASCEND, ascend) -#endif -#ifdef ENABLE_KUNLUN_API - CREATE(INFINI_DEVICE_KUNLUN, kunlun) -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetCausalSoftmaxWorkspaceSize(infiniopCausalSoftmaxDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - GET(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_ASCEND_API - GET(INFINI_DEVICE_ASCEND, ascend) -#endif -#ifdef ENABLE_CAMBRICON_API - GET(INFINI_DEVICE_CAMBRICON, bang) -#endif -#ifdef ENABLE_KUNLUN_API - GET(INFINI_DEVICE_KUNLUN, kunlun) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopCausalSoftmax( - infiniopCausalSoftmaxDescriptor_t desc, - void *workspace, size_t workspace_size, - void *y, - const void *x, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc)->calculate( \ - workspace, workspace_size, y, x, stream); - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang) -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_ASCEND_API - CALCULATE(INFINI_DEVICE_ASCEND, ascend) -#endif -#ifdef ENABLE_KUNLUN_API - CALCULATE(INFINI_DEVICE_KUNLUN, kunlun) -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyCausalSoftmaxDescriptor(infiniopCausalSoftmaxDescriptor_t desc) { - -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DESTROY(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - DESTROY(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - DESTROY(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DESTROY(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_CAMBRICON_API - DESTROY(INFINI_DEVICE_CAMBRICON, bang) -#endif -#ifdef ENABLE_METAX_API - DESTROY(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_ASCEND_API - DESTROY(INFINI_DEVICE_ASCEND, ascend) -#endif -#ifdef ENABLE_KUNLUN_API - DESTROY(INFINI_DEVICE_KUNLUN, kunlun) -#endif -#ifdef ENABLE_MOORE_API - DESTROY(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DESTROY -} diff --git a/src/infiniop/ops/cdist/cdist.h b/src/infiniop/ops/cdist/cdist.h deleted file mode 100644 index 15e5adb8c..000000000 --- a/src/infiniop/ops/cdist/cdist.h +++ /dev/null @@ -1,62 +0,0 @@ -#ifndef __CDIST_H__ -#define __CDIST_H__ - -#include "../../operator.h" -#include "info.h" - -/** - * # 关于 `cdist` 算子描述符的说明 - * * 仿照 GEMM 的 PImpl (Opaque) 设计模式，将硬件相关的执行上下文（如 CUDA Handle、计算流等） - * 隐藏在 `Opaque` 结构体中，确保头文件在不同后端（CPU/NVIDIA/Ascend）间的一致性。 - */ - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::cdist::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - infiniDtype_t _dtype; \ - CdistInfo _info; /* 包含 M, N, D 维度信息 */ \ - size_t _workspace_size; \ - double _p; /* 范数阶数，创建时固定 */ \ - \ - Descriptor( \ - infiniDtype_t dtype, \ - CdistInfo info, \ - double p, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _dtype(dtype), \ - _info(info), \ - _workspace_size(workspace_size_), \ - _p(p) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t y_desc, /* 输出 (M, N) */ \ - infiniopTensorDescriptor_t x1_desc, /* 输入 (M, D) */ \ - infiniopTensorDescriptor_t x2_desc, /* 输入 (N, D) */ \ - double p); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *y, /* 结果矩阵 */ \ - const void *x1, \ - const void *x2, \ - void *stream) const; \ - }; \ - } - -#endif // __CDIST_H__ diff --git a/src/infiniop/ops/cdist/cpu/cdist_cpu.cc b/src/infiniop/ops/cdist/cpu/cdist_cpu.cc deleted file mode 100644 index 6eb7ef6a5..000000000 --- a/src/infiniop/ops/cdist/cpu/cdist_cpu.cc +++ /dev/null @@ -1,126 +0,0 @@ -#include "cdist_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include - -namespace op::cdist::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p) { - - auto handle = reinterpret_cast(handle_); - auto dtype = y_desc->dtype(); - - // 1. 类型检查 - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - // 2. 解析维度信息 (CdistInfo 逻辑已在之前定义) - auto result = CdistInfo::create(y_desc, x1_desc, x2_desc); - CHECK_RESULT(result); - - // 3. 实例化描述符，CPU 版通常不需要 workspace - *desc_ptr = new Descriptor( - dtype, result.take(), p, 0, - nullptr, - handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -/** - * 核心计算模板：处理不同数据类型 - */ -template -void calculate_dist( - const CdistInfo &info, - void *y, - const void *x1, - const void *x2, - double p) { - - // Flatten loops: batch * m * n for OpenMP parallelization - const ptrdiff_t total = ptrdiff_t(info.batch) * ptrdiff_t(info.m) * ptrdiff_t(info.n); - -#pragma omp parallel for - for (ptrdiff_t idx = 0; idx < total; ++idx) { - ptrdiff_t b = idx / (info.m * info.n); - ptrdiff_t rem = idx % (info.m * info.n); - ptrdiff_t i = rem / info.n; - ptrdiff_t j = rem % info.n; - - // output pointer: y[b, i, j] - auto y_ptr = reinterpret_cast(y) - + b * info.y_matrix.stride - + i * info.y_matrix.row_stride - + j * info.y_matrix.col_stride; - - // input vectors: x1[b, i, :] and x2[b, j, :] - auto x1_vec = reinterpret_cast(x1) - + b * info.x1_matrix.stride - + i * info.x1_matrix.row_stride; - auto x2_vec = reinterpret_cast(x2) - + b * info.x2_matrix.stride - + j * info.x2_matrix.row_stride; - - double dist = 0.0; - - for (size_t k = 0; k < info.d; ++k) { - float v1 = utils::cast(*(x1_vec + k * info.x1_matrix.col_stride)); - float v2 = utils::cast(*(x2_vec + k * info.x2_matrix.col_stride)); - float diff = std::abs(v1 - v2); - - if (p == 1.0) { - dist += diff; - } else if (p == 2.0) { - dist += diff * diff; - } else if (std::isinf(p)) { - dist = std::max(dist, static_cast(diff)); - } else { - dist += std::pow(static_cast(diff), p); - } - } - - // final distance - if (p == 2.0) { - dist = std::sqrt(dist); - } else if (!std::isinf(p) && p != 1.0) { - dist = std::pow(dist, 1.0 / p); - } - - *y_ptr = utils::cast(static_cast(dist)); - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - cpu::calculate_dist(_info, y, x1, x2, _p); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_BF16: - cpu::calculate_dist(_info, y, x1, x2, _p); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_F32: - cpu::calculate_dist(_info, y, x1, x2, _p); - return INFINI_STATUS_SUCCESS; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::cdist::cpu diff --git a/src/infiniop/ops/cdist/cpu/cdist_cpu.h b/src/infiniop/ops/cdist/cpu/cdist_cpu.h deleted file mode 100644 index 1944c9d6e..000000000 --- a/src/infiniop/ops/cdist/cpu/cdist_cpu.h +++ /dev/null @@ -1,11 +0,0 @@ -#ifndef __CDIST_CPU_H__ -#define __CDIST_CPU_H__ - -#include "../cdist.h" - -// 使用 cdist.h 中定义的 DESCRIPTOR 宏 -// 这将在命名空间 op::cdist::cpu 中生成 Descriptor 类 -// 该类包含对 CdistInfo 的引用以及 create/calculate 等接口 -DESCRIPTOR(cpu) - -#endif // __CDIST_CPU_H__ diff --git a/src/infiniop/ops/cdist/info.h b/src/infiniop/ops/cdist/info.h deleted file mode 100644 index 93adace36..000000000 --- a/src/infiniop/ops/cdist/info.h +++ /dev/null @@ -1,113 +0,0 @@ -#ifndef __CDIST_INFO_H__ -#define __CDIST_INFO_H__ - -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" -#include - -namespace op::cdist { - -/** - * 借用 BlasMatrix 的概念来描述 cdist 的输入输出矩阵 - * x1: (Batch, M, D) - * x2: (Batch, N, D) - * y: (Batch, M, N) - */ -struct CdistMatrix { - size_t ndim; - size_t batch; - ptrdiff_t stride; // Batch 之间的步长 - size_t rows; // M 或 N - size_t cols; // D (特征维度) 或结果中的 N - ptrdiff_t row_stride; - ptrdiff_t col_stride; - - static utils::Result create(infiniopTensorDescriptor_t layout) { - CdistMatrix ans; - auto ndim = layout->ndim(); - - if (ndim == 2) { - ans.ndim = 2; - ans.batch = 1; - ans.stride = 0; - ans.rows = layout->dim(0); - ans.cols = layout->dim(1); - ans.row_stride = layout->stride(0); - ans.col_stride = layout->stride(1); - } else if (ndim == 3) { - ans.ndim = 3; - ans.batch = layout->dim(0); - ans.stride = ans.batch == 1 ? 0 : layout->stride(0); - ans.rows = layout->dim(1); - ans.cols = layout->dim(2); - ans.row_stride = layout->stride(1); - ans.col_stride = layout->stride(2); - } else { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - return utils::Result(ans); - } - - bool match_batch(size_t _batch) const { - return batch == _batch || batch == 1; - } -}; - -class CdistInfo { - CdistInfo() = default; - -public: - CdistMatrix x1_matrix; - CdistMatrix x2_matrix; - CdistMatrix y_matrix; - - size_t m, n, d, batch; - - static utils::Result create( - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc) { - - auto x1_res = CdistMatrix::create(x1_desc); - CHECK_RESULT(x1_res); - - auto x2_res = CdistMatrix::create(x2_desc); - CHECK_RESULT(x2_res); - - auto y_res = CdistMatrix::create(y_desc); - CHECK_RESULT(y_res); - - auto x1 = x1_res.take(); - auto x2 = x2_res.take(); - auto y = y_res.take(); - - // 1. 维度校验 - // x1(M, D), x2(N, D) -> y(M, N) - if (x1.cols != x2.cols) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; // 特征维度 D 必须一致 - } - if (y.rows != x1.rows || y.cols != x2.rows) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; // 输出形状必须为 M x N - } - - // 2. Batch 校验 - size_t batch_size = y.batch; - if (!x1.match_batch(batch_size) || !x2.match_batch(batch_size)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t m = x1.rows; - size_t n = x2.rows; - size_t d = x1.cols; - - return utils::Result(CdistInfo{ - x1, x2, y, - m, n, d, batch_size}); - } -}; - -} // namespace op::cdist - -#endif // __CDIST_INFO_H__ diff --git a/src/infiniop/ops/cdist/metax/cdist_metax.h b/src/infiniop/ops/cdist/metax/cdist_metax.h deleted file mode 100644 index 6a310e655..000000000 --- a/src/infiniop/ops/cdist/metax/cdist_metax.h +++ /dev/null @@ -1,16 +0,0 @@ -#ifndef __CDIST_METAX_CUH__ -#define __CDIST_METAX_CUH__ - -#include "../cdist.h" - -/** - * 使用 cdist.h 中定义的 DESCRIPTOR 宏。 - * 这将在命名空间 op::cdist::metax 中生成针对 METAX 设备的 Descriptor 类。 - * * 在 METAX 端的具体实现中，Opaque 结构体通常会存储： - * - cublasHandle_t: 用于 p=2.0 时的矩阵乘法加速。 - * - cudaStream_t: 当前执行的任务流。 - * - 自定义 Kernel 的配置参数。 - */ -DESCRIPTOR(metax) - -#endif // __CDIST_METAX_CUH__ diff --git a/src/infiniop/ops/cdist/metax/cdist_metax.maca b/src/infiniop/ops/cdist/metax/cdist_metax.maca deleted file mode 100644 index b0ebb0d81..000000000 --- a/src/infiniop/ops/cdist/metax/cdist_metax.maca +++ /dev/null @@ -1,163 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "cdist_metax.h" -#include -namespace op::cdist::metax { - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p) { - - auto handle = reinterpret_cast(handle_); - auto dtype = y_desc->dtype(); - - // 目前 METAX 后端仅支持 F32，测试也是 F32 - CHECK_DTYPE(dtype, INFINI_DTYPE_F32); - - auto result = CdistInfo::create(y_desc, x1_desc, x2_desc); - CHECK_RESULT(result); - - // 当前实现不使用 workspace - *desc_ptr = new Descriptor( - dtype, result.take(), p, 0, - nullptr, - handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -// --- Kernel 1: L2 Epilogue --- -// 保留占位，当前实现未使用 GEMM 加速路径 -template -__global__ void cdist_l2_epilogue_kernel(T *y, const T *x1_norm, const T *x2_norm, - int M, int N, int batch_stride_y) { - int j = blockIdx.x * blockDim.x + threadIdx.x; - int i = blockIdx.y * blockDim.y + threadIdx.y; - int b = blockIdx.z; - - if (i < M && j < N) { - int idx = b * batch_stride_y + i * N + j; - // GEMM 已经计算了 -2*x1*x2^T 并存入 y - float val = (float)x1_norm[b * M + i] + (float)x2_norm[b * N + j] + (float)y[idx]; - y[idx] = (T)sqrtf(fmaxf(val, 0.0f)); - } -} - -// --- Kernel 2: Generic P-Norm (F32, 支持通用步长) --- -__global__ void cdist_generic_kernel_f32( - float *y, - const float *x1, - const float *x2, - size_t m, - size_t n, - size_t d, - ptrdiff_t x1_stride, - ptrdiff_t x1_row_stride, - ptrdiff_t x1_col_stride, - ptrdiff_t x2_stride, - ptrdiff_t x2_row_stride, - ptrdiff_t x2_col_stride, - ptrdiff_t y_stride, - ptrdiff_t y_row_stride, - ptrdiff_t y_col_stride, - double p) { - - int j = blockIdx.x * blockDim.x + threadIdx.x; - int i = blockIdx.y * blockDim.y + threadIdx.y; - int b = blockIdx.z; - - if (i >= (int)m || j >= (int)n) { - return; - } - - // 定位输出位置 y[b, i, j] - float *y_ptr = y + b * y_stride + i * y_row_stride + j * y_col_stride; - - // 定位向量位置 x1[b, i, :] 和 x2[b, j, :] - const float *x1_vec = x1 + b * x1_stride + i * x1_row_stride; - const float *x2_vec = x2 + b * x2_stride + j * x2_row_stride; - - double dist = 0.0; - - for (size_t k = 0; k < d; ++k) { - float v1 = *(x1_vec + k * x1_col_stride); - float v2 = *(x2_vec + k * x2_col_stride); - float diff = fabsf(v1 - v2); - - if (p == 1.0) { - dist += diff; - } else if (p == 2.0) { - dist += diff * diff; - } else if (isinf(p)) { - dist = fmaxf((float)dist, diff); - } else { - dist += powf((float)diff, (float)p); - } - } - - if (p == 2.0) { - dist = sqrtf((float)dist); - } else if (!isinf(p) && p != 1.0) { - dist = powf((float)dist, 1.0f / (float)p); - } - - *y_ptr = (float)dist; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - if (_dtype != INFINI_DTYPE_F32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - hcStream_t custream = (hcStream_t)stream; - dim3 block(16, 16); - dim3 grid( - static_cast((_info.n + block.x - 1) / block.x), - static_cast((_info.m + block.y - 1) / block.y), - static_cast(_info.batch)); - - cdist_generic_kernel_f32<<>>( - static_cast(y), - static_cast(x1), - static_cast(x2), - _info.m, - _info.n, - _info.d, - _info.x1_matrix.stride, - _info.x1_matrix.row_stride, - _info.x1_matrix.col_stride, - _info.x2_matrix.stride, - _info.x2_matrix.row_stride, - _info.x2_matrix.col_stride, - _info.y_matrix.stride, - _info.y_matrix.row_stride, - _info.y_matrix.col_stride, - _p); - - auto err = hcGetLastError(); - if (err != hcSuccess) { - return INFINI_STATUS_INTERNAL_ERROR; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::cdist::metax diff --git a/src/infiniop/ops/cdist/moore/cdist_moore.h b/src/infiniop/ops/cdist/moore/cdist_moore.h deleted file mode 100644 index 864049402..000000000 --- a/src/infiniop/ops/cdist/moore/cdist_moore.h +++ /dev/null @@ -1,16 +0,0 @@ -#ifndef __CDIST_MOORE_H__ -#define __CDIST_MOORE_H__ - -#include "../cdist.h" - -/** - * 使用 cdist.h 中定义的 DESCRIPTOR 宏。 - * 这将在命名空间 op::cdist::moore 中生成针对 Moore 设备的 Descriptor 类。 - * * 在 Moore 端的具体实现中，Opaque 结构体通常会存储： - * - mublasHandle_t: 用于 p=2.0 时的矩阵乘法加速（对应 NVIDIA 的 cuBLAS）。 - * - musaStream_t: 当前执行的任务流。 - * - 自定义 Kernel 的配置参数。 - */ -DESCRIPTOR(moore) - -#endif // __CDIST_MOORE_H__ diff --git a/src/infiniop/ops/cdist/moore/cdist_moore.mu b/src/infiniop/ops/cdist/moore/cdist_moore.mu deleted file mode 100644 index 7b070f578..000000000 --- a/src/infiniop/ops/cdist/moore/cdist_moore.mu +++ /dev/null @@ -1,145 +0,0 @@ -#include "../../../devices/moore/moore_handle.h" -#include "cdist_moore.h" -#include -#include - -namespace op::cdist::moore { - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p) { - - // 1. 转换至 Moore 句柄 - auto handle = reinterpret_cast(handle_); - auto dtype = y_desc->dtype(); - - // 保持与原版一致，目前仅支持 F32 - CHECK_DTYPE(dtype, INFINI_DTYPE_F32); - - auto result = CdistInfo::create(y_desc, x1_desc, x2_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - dtype, result.take(), p, 0, - nullptr, - handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -// --- Kernel: Generic P-Norm (MUSA F32 实现) --- -__global__ void cdist_generic_kernel_f32( - float *y, - const float *x1, - const float *x2, - size_t m, - size_t n, - size_t d, - ptrdiff_t x1_stride, - ptrdiff_t x1_row_stride, - ptrdiff_t x1_col_stride, - ptrdiff_t x2_stride, - ptrdiff_t x2_row_stride, - ptrdiff_t x2_col_stride, - ptrdiff_t y_stride, - ptrdiff_t y_row_stride, - ptrdiff_t y_col_stride, - double p) { - - // 2. MUSA 同样支持 3D 线程索引 - int j = blockIdx.x * blockDim.x + threadIdx.x; - int i = blockIdx.y * blockDim.y + threadIdx.y; - int b = blockIdx.z; - - if (i >= (int)m || j >= (int)n) { - return; - } - - // 定位输出 y[b, i, j] - float *y_ptr = y + b * y_stride + i * y_row_stride + j * y_col_stride; - - // 定位输入向量 - const float *x1_vec = x1 + b * x1_stride + i * x1_row_stride; - const float *x2_vec = x2 + b * x2_stride + j * x2_row_stride; - - double dist = 0.0; - - for (size_t k = 0; k < d; ++k) { - float v1 = *(x1_vec + k * x1_col_stride); - float v2 = *(x2_vec + k * x2_col_stride); - float diff = fabsf(v1 - v2); - - if (p == 1.0) { - dist += (double)diff; - } else if (p == 2.0) { - dist += (double)diff * diff; - } else if (isinf(p)) { - dist = fmaxf((float)dist, diff); - } else { - dist += powf(diff, (float)p); - } - } - - if (p == 2.0) { - dist = sqrtf((float)dist); - } else if (!isinf(p) && p != 1.0) { - dist = powf((float)dist, 1.0f / (float)p); - } - - *y_ptr = (float)dist; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - if (_dtype != INFINI_DTYPE_F32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // 3. 切换至 musaStream_t - musaStream_t mustream = reinterpret_cast(stream); - - // 保持 16x16 的 Block 大小，这在 MUSA 架构上也是通用的 - dim3 block(16, 16); - dim3 grid( - static_cast((_info.n + block.x - 1) / block.x), - static_cast((_info.m + block.y - 1) / block.y), - static_cast(_info.batch)); - - cdist_generic_kernel_f32<<>>( - static_cast(y), - static_cast(x1), - static_cast(x2), - _info.m, - _info.n, - _info.d, - _info.x1_matrix.stride, - _info.x1_matrix.row_stride, - _info.x1_matrix.col_stride, - _info.x2_matrix.stride, - _info.x2_matrix.row_stride, - _info.x2_matrix.col_stride, - _info.y_matrix.stride, - _info.y_matrix.row_stride, - _info.y_matrix.col_stride, - _p); - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::cdist::moore diff --git a/src/infiniop/ops/cdist/nvidia/cdist_nvidia.cu b/src/infiniop/ops/cdist/nvidia/cdist_nvidia.cu deleted file mode 100644 index 7da35a7a4..000000000 --- a/src/infiniop/ops/cdist/nvidia/cdist_nvidia.cu +++ /dev/null @@ -1,161 +0,0 @@ -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "cdist_nvidia.cuh" -#include -namespace op::cdist::nvidia { - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p) { - - auto handle = reinterpret_cast(handle_); - auto dtype = y_desc->dtype(); - - // 目前 NVIDIA 后端仅支持 F32，测试也是 F32 - CHECK_DTYPE(dtype, INFINI_DTYPE_F32); - - auto result = CdistInfo::create(y_desc, x1_desc, x2_desc); - CHECK_RESULT(result); - - // 当前实现不使用 workspace - *desc_ptr = new Descriptor( - dtype, result.take(), p, 0, - nullptr, - handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -// --- Kernel 1: L2 Epilogue --- -// 保留占位，当前实现未使用 GEMM 加速路径 -template -__global__ void cdist_l2_epilogue_kernel(T *y, const T *x1_norm, const T *x2_norm, - int M, int N, int batch_stride_y) { - int j = blockIdx.x * blockDim.x + threadIdx.x; - int i = blockIdx.y * blockDim.y + threadIdx.y; - int b = blockIdx.z; - - if (i < M && j < N) { - int idx = b * batch_stride_y + i * N + j; - // GEMM 已经计算了 -2*x1*x2^T 并存入 y - float val = (float)x1_norm[b * M + i] + (float)x2_norm[b * N + j] + (float)y[idx]; - y[idx] = (T)sqrtf(fmaxf(val, 0.0f)); - } -} - -// --- Kernel 2: Generic P-Norm (F32, 支持通用步长) --- -__global__ void cdist_generic_kernel_f32( - float *y, - const float *x1, - const float *x2, - size_t m, - size_t n, - size_t d, - ptrdiff_t x1_stride, - ptrdiff_t x1_row_stride, - ptrdiff_t x1_col_stride, - ptrdiff_t x2_stride, - ptrdiff_t x2_row_stride, - ptrdiff_t x2_col_stride, - ptrdiff_t y_stride, - ptrdiff_t y_row_stride, - ptrdiff_t y_col_stride, - double p) { - - int j = blockIdx.x * blockDim.x + threadIdx.x; - int i = blockIdx.y * blockDim.y + threadIdx.y; - int b = blockIdx.z; - - if (i >= (int)m || j >= (int)n) { - return; - } - - // 定位输出位置 y[b, i, j] - float *y_ptr = y + b * y_stride + i * y_row_stride + j * y_col_stride; - - // 定位向量位置 x1[b, i, :] 和 x2[b, j, :] - const float *x1_vec = x1 + b * x1_stride + i * x1_row_stride; - const float *x2_vec = x2 + b * x2_stride + j * x2_row_stride; - - double dist = 0.0; - - for (size_t k = 0; k < d; ++k) { - float v1 = *(x1_vec + k * x1_col_stride); - float v2 = *(x2_vec + k * x2_col_stride); - float diff = fabsf(v1 - v2); - - if (p == 1.0) { - dist += diff; - } else if (p == 2.0) { - dist += diff * diff; - } else if (isinf(p)) { - dist = fmaxf((float)dist, diff); - } else { - dist += powf((float)diff, (float)p); - } - } - - if (p == 2.0) { - dist = sqrtf((float)dist); - } else if (!isinf(p) && p != 1.0) { - dist = powf((float)dist, 1.0f / (float)p); - } - - *y_ptr = (float)dist; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - if (_dtype != INFINI_DTYPE_F32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - cudaStream_t custream = (cudaStream_t)stream; - dim3 block(16, 16); - dim3 grid( - static_cast((_info.n + block.x - 1) / block.x), - static_cast((_info.m + block.y - 1) / block.y), - static_cast(_info.batch)); - - cdist_generic_kernel_f32<<>>( - static_cast(y), - static_cast(x1), - static_cast(x2), - _info.m, - _info.n, - _info.d, - _info.x1_matrix.stride, - _info.x1_matrix.row_stride, - _info.x1_matrix.col_stride, - _info.x2_matrix.stride, - _info.x2_matrix.row_stride, - _info.x2_matrix.col_stride, - _info.y_matrix.stride, - _info.y_matrix.row_stride, - _info.y_matrix.col_stride, - _p); - - auto err = cudaGetLastError(); - if (err != cudaSuccess) { - return INFINI_STATUS_INTERNAL_ERROR; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::cdist::nvidia diff --git a/src/infiniop/ops/cdist/nvidia/cdist_nvidia.cuh b/src/infiniop/ops/cdist/nvidia/cdist_nvidia.cuh deleted file mode 100644 index a4297f50a..000000000 --- a/src/infiniop/ops/cdist/nvidia/cdist_nvidia.cuh +++ /dev/null @@ -1,16 +0,0 @@ -#ifndef __CDIST_NVIDIA_CUH__ -#define __CDIST_NVIDIA_CUH__ - -#include "../cdist.h" - -/** - * 使用 cdist.h 中定义的 DESCRIPTOR 宏。 - * 这将在命名空间 op::cdist::nvidia 中生成针对 NVIDIA 设备的 Descriptor 类。 - * * 在 NVIDIA 端的具体实现中，Opaque 结构体通常会存储： - * - cublasHandle_t: 用于 p=2.0 时的矩阵乘法加速。 - * - cudaStream_t: 当前执行的任务流。 - * - 自定义 Kernel 的配置参数。 - */ -DESCRIPTOR(nvidia) - -#endif // __CDIST_NVIDIA_CUH__ diff --git a/src/infiniop/ops/cdist/operator.cc b/src/infiniop/ops/cdist/operator.cc deleted file mode 100644 index 0ed6541b7..000000000 --- a/src/infiniop/ops/cdist/operator.cc +++ /dev/null @@ -1,184 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/cdist.h" - -// 引入各硬件后端的 Descriptor 定义 -#ifdef ENABLE_CPU_API -#include "cpu/cdist_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/cdist_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/cdist_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/cdist_moore.h" -#endif - -// ----------------------------------------------------------------------------- -// 1. 创建描述符 -// ----------------------------------------------------------------------------- -__INFINI_C infiniStatus_t infiniopCreateCdistDescriptor( - infiniopHandle_t handle, - infiniopCdistDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::cdist::NAMESPACE::Descriptor::create(handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, x1_desc, x2_desc, p) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ----------------------------------------------------------------------------- -// 2. 获取 Workspace 大小 -// ----------------------------------------------------------------------------- -__INFINI_C infiniStatus_t infiniopGetCdistWorkspaceSize( - infiniopCdistDescriptor_t desc, - size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -// ----------------------------------------------------------------------------- -// 3. 执行计算 (计算成对距离) -// ----------------------------------------------------------------------------- -__INFINI_C infiniStatus_t infiniopCdist( - infiniopCdistDescriptor_t desc, - void *workspace, size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, x1, x2, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -// ----------------------------------------------------------------------------- -// 4. 销毁描述符 -// ----------------------------------------------------------------------------- -__INFINI_C infiniStatus_t infiniopDestroyCdistDescriptor(infiniopCdistDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} diff --git a/src/infiniop/ops/chunk_gated_delta_rule/chunk_gated_delta_rule.h b/src/infiniop/ops/chunk_gated_delta_rule/chunk_gated_delta_rule.h deleted file mode 100644 index 808e0e52c..000000000 --- a/src/infiniop/ops/chunk_gated_delta_rule/chunk_gated_delta_rule.h +++ /dev/null @@ -1,62 +0,0 @@ -// infiniop/ops/chunk_gated_delta_rule.h - -#ifndef __INFINIOP_CHUNK_GATED_DELTA_RULE_H__ -#define __INFINIOP_CHUNK_GATED_DELTA_RULE_H__ - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::chunk_gated_delta_rule::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - ChunkGatedDeltaRuleInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - ChunkGatedDeltaRuleInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t initial_state_desc, \ - infiniopTensorDescriptor_t final_state_desc, \ - infiniopTensorDescriptor_t q_desc, \ - infiniopTensorDescriptor_t k_desc, \ - infiniopTensorDescriptor_t v_desc, \ - infiniopTensorDescriptor_t g_desc, \ - infiniopTensorDescriptor_t beta_desc, \ - infiniopTensorDescriptor_t cu_seqlens_desc, \ - infiniopTensorDescriptor_t initial_state_indices_desc, \ - infiniopTensorDescriptor_t final_state_indices_desc, \ - bool use_qk_l2norm, \ - size_t chunk_size); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *out, void *initial_state, void *final_state, \ - const void *q, const void *k, const void *v, \ - const void *g, const void *beta, const void *cu_seqlens, \ - const void *initial_state_indices, \ - const void *final_state_indices, \ - void *stream) const; \ - }; \ - } - -#endif // __INFINIOP_CHUNK_GATED_DELTA_RULE_H__ diff --git a/src/infiniop/ops/chunk_gated_delta_rule/cuda/kernel.cuh b/src/infiniop/ops/chunk_gated_delta_rule/cuda/kernel.cuh deleted file mode 100644 index d9bac786d..000000000 --- a/src/infiniop/ops/chunk_gated_delta_rule/cuda/kernel.cuh +++ /dev/null @@ -1,649 +0,0 @@ -#ifndef __CHUNK_GATED_DELTA_RULE_KERNEL_CUH__ -#define __CHUNK_GATED_DELTA_RULE_KERNEL_CUH__ - -#include - -__device__ inline int64_t loadOptionalIndex( - const void *indices, - bool is_i64, - int idx, - int fallback) { - if (indices == nullptr) { - return static_cast(fallback); - } - return is_i64 - ? static_cast(indices)[idx] - : static_cast(static_cast(indices)[idx]); -} - -template -__device__ inline float loadAsFloat(const T *ptr, ptrdiff_t offset) { - return static_cast(ptr[offset]); -} - -template <> -__device__ inline float loadAsFloat(const half *ptr, ptrdiff_t offset) { - return __half2float(ptr[offset]); -} - -template <> -__device__ inline float loadAsFloat(const cuda_bfloat16 *ptr, ptrdiff_t offset) { - return __bfloat162float(ptr[offset]); -} - -#define CGDR_FOR(idx, n) \ - for (int idx = threadIdx.x; idx < static_cast(n); idx += blockDim.x) - -template -__device__ Tcompute blockReduceSum(Tcompute v) { - __shared__ Tcompute smem[NUM_THREADS]; - smem[threadIdx.x] = v; - __syncthreads(); - - for (int s = NUM_THREADS / 2; s > 0; s >>= 1) { - if (threadIdx.x < s) { - smem[threadIdx.x] += smem[threadIdx.x + s]; - } - __syncthreads(); - } - return smem[0]; -} - -template < - typename Tdata, - typename Tgate, - typename Tcompute, - size_t Dk, - size_t Dv, - size_t NUM_THREADS> -__device__ void chunkGatedDeltaRuleKernel( - Tcompute *state_workspace, - Tdata *out, - Tdata *initial_state, - Tdata *final_state, - const Tdata *q, - const Tdata *k, - const Tdata *v, - const Tgate *g, - const Tgate *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - bool cu_seqlens_i64, - bool initial_state_indices_i64, - bool final_state_indices_i64, - bool use_qk_l2norm, - bool has_cu_seqlens, - bool indexed_state_pool, - size_t T, - size_t chunk_size, - size_t pool_size, - size_t Hk, - size_t value_heads_per_key_head, - ptrdiff_t out_s0, - ptrdiff_t out_s1, - ptrdiff_t out_s2, - ptrdiff_t initial_s0, - ptrdiff_t initial_s1, - ptrdiff_t initial_s2, - ptrdiff_t initial_s3, - ptrdiff_t final_s0, - ptrdiff_t final_s1, - ptrdiff_t final_s2, - ptrdiff_t final_s3, - ptrdiff_t q_s0, - ptrdiff_t q_s1, - ptrdiff_t q_s2, - ptrdiff_t k_s0, - ptrdiff_t k_s1, - ptrdiff_t k_s2, - ptrdiff_t v_s0, - ptrdiff_t v_s1, - ptrdiff_t v_s2, - ptrdiff_t g_s0, - ptrdiff_t g_s1, - ptrdiff_t g_s2, - ptrdiff_t beta_s0, - ptrdiff_t beta_s1, - ptrdiff_t beta_s2) { - - const int batch_idx = blockIdx.x; - const int value_head_idx = blockIdx.y; - const int key_head_idx = value_head_idx / static_cast(value_heads_per_key_head); - - if (key_head_idx >= static_cast(Hk)) { - return; - } - - int64_t token_begin = 0; - int64_t token_end = static_cast(T); - - if (has_cu_seqlens) { - token_begin = loadOptionalIndex(cu_seqlens, cu_seqlens_i64, batch_idx, 0); - token_end = loadOptionalIndex(cu_seqlens, cu_seqlens_i64, batch_idx + 1, 0); - if (token_begin < 0 || token_end < token_begin || token_end > static_cast(T)) { - return; - } - } - - int64_t read_slot = batch_idx; - int64_t write_slot = batch_idx; - - if (indexed_state_pool) { - read_slot = loadOptionalIndex( - initial_state_indices, - initial_state_indices_i64, - batch_idx, - batch_idx); - - write_slot = final_state_indices == nullptr - ? static_cast(batch_idx) - : loadOptionalIndex( - final_state_indices, - final_state_indices_i64, - batch_idx, - batch_idx); - - if (read_slot < 0 || write_slot < 0 || read_slot >= static_cast(pool_size) || write_slot >= static_cast(pool_size)) { - return; - } - } - - const ptrdiff_t initial_base = indexed_state_pool - ? static_cast(read_slot) * initial_s0 + static_cast(value_head_idx) * initial_s1 - : static_cast(batch_idx) * initial_s0 + static_cast(value_head_idx) * initial_s1; - - Tdata *final_state_target = nullptr; - ptrdiff_t final_base = 0; - - if (indexed_state_pool && final_state_indices != nullptr) { - final_state_target = initial_state; - final_base = static_cast(write_slot) * initial_s0 + static_cast(value_head_idx) * initial_s1; - } else { - final_state_target = final_state; - final_base = static_cast(batch_idx) * final_s0 + static_cast(value_head_idx) * final_s1; - } - - const ptrdiff_t per_block_workspace = static_cast(Dk * Dv) + static_cast(chunk_size * Dk) * 3 + static_cast(chunk_size * Dv) * 3 + static_cast(chunk_size * chunk_size) + static_cast(chunk_size) * 3; - - const ptrdiff_t workspace_block = (static_cast(batch_idx) * gridDim.y + static_cast(value_head_idx)) * per_block_workspace; - - Tcompute *state_local = state_workspace + workspace_block; - Tcompute *q_buf = state_local + Dk * Dv; - Tcompute *k_buf = q_buf + chunk_size * Dk; - Tcompute *k_cumdecay = k_buf + chunk_size * Dk; - Tcompute *v_beta = k_cumdecay + chunk_size * Dk; - Tcompute *v_mid = v_beta + chunk_size * Dv; - Tcompute *v_new = v_mid + chunk_size * Dv; - Tcompute *attn = v_new + chunk_size * Dv; - Tcompute *g_cum = attn + chunk_size * chunk_size; - Tcompute *beta_buf = g_cum + chunk_size; - Tcompute *row_buf = beta_buf + chunk_size; - - const int token_batch = has_cu_seqlens ? 0 : batch_idx; - const Tcompute scale = rsqrtf(static_cast(Dk)); - - // Load initial state. - CGDR_FOR(i, Dk * Dv) { - int dk = i / Dv; - int dv = i % Dv; - - ptrdiff_t read_idx = initial_base + static_cast(dv) * initial_s2 + static_cast(dk) * initial_s3; - - state_local[i] = static_cast( - loadAsFloat(initial_state, read_idx)); - } - __syncthreads(); - - for (int64_t chunk_begin = token_begin; - chunk_begin < token_end; - chunk_begin += static_cast(chunk_size)) { - - const int64_t remaining = token_end - chunk_begin; - const int actual_len = static_cast( - remaining < static_cast(chunk_size) - ? remaining - : static_cast(chunk_size)); - - // Load beta and cumulative gate. - if (threadIdx.x == 0) { - Tcompute running_g = 0; - for (int t = 0; t < static_cast(chunk_size); ++t) { - if (t < actual_len) { - int64_t token_idx = chunk_begin + t; - ptrdiff_t gate_offset = static_cast(token_batch) * g_s0 + static_cast(token_idx) * g_s1 + static_cast(value_head_idx) * g_s2; - ptrdiff_t beta_offset = static_cast(token_batch) * beta_s0 + static_cast(token_idx) * beta_s1 + static_cast(value_head_idx) * beta_s2; - - running_g += static_cast( - loadAsFloat(g, gate_offset)); - beta_buf[t] = static_cast( - loadAsFloat(beta, beta_offset)); - g_cum[t] = running_g; - } else { - beta_buf[t] = 0; - g_cum[t] = running_g; - } - } - } - __syncthreads(); - - // Load q/k/v_beta. - CGDR_FOR(x, chunk_size * Dk) { - int t = x / Dk; - int dk = x % Dk; - - if (t < actual_len) { - int64_t token_idx = chunk_begin + t; - ptrdiff_t q_base = static_cast(token_batch) * q_s0 + static_cast(token_idx) * q_s1 + static_cast(key_head_idx) * q_s2; - ptrdiff_t k_base = static_cast(token_batch) * k_s0 + static_cast(token_idx) * k_s1 + static_cast(key_head_idx) * k_s2; - - q_buf[x] = static_cast(loadAsFloat(q, q_base + dk)) * scale; - k_buf[x] = static_cast(loadAsFloat(k, k_base + dk)); - } else { - q_buf[x] = 0; - k_buf[x] = 0; - } - } - - CGDR_FOR(x, chunk_size * Dv) { - int t = x / Dv; - int dv = x % Dv; - - if (t < actual_len) { - int64_t token_idx = chunk_begin + t; - ptrdiff_t v_base = static_cast(token_batch) * v_s0 + static_cast(token_idx) * v_s1 + static_cast(value_head_idx) * v_s2; - - v_beta[x] = static_cast(loadAsFloat(v, v_base + dv)) * beta_buf[t]; - } else { - v_beta[x] = 0; - } - } - __syncthreads(); - - // Optional q/k L2 norm. - if (use_qk_l2norm) { - for (int t = 0; t < static_cast(chunk_size); ++t) { - Tcompute q_sum = 0; - Tcompute k_sum = 0; - - for (int dk = threadIdx.x; dk < static_cast(Dk); dk += blockDim.x) { - q_sum += q_buf[t * Dk + dk] * q_buf[t * Dk + dk]; - k_sum += k_buf[t * Dk + dk] * k_buf[t * Dk + dk]; - } - - q_sum = blockReduceSum(q_sum); - k_sum = blockReduceSum(k_sum); - - Tcompute q_norm = rsqrtf(q_sum / (scale * scale) + 1e-6f); - Tcompute k_norm = rsqrtf(k_sum + 1e-6f); - - for (int dk = threadIdx.x; dk < static_cast(Dk); dk += blockDim.x) { - q_buf[t * Dk + dk] *= q_norm; - k_buf[t * Dk + dk] *= k_norm; - } - __syncthreads(); - } - } - - // Build lower-triangular attn. - CGDR_FOR(x, chunk_size * chunk_size) { - int i = x / chunk_size; - int j = x % chunk_size; - - if (j < i) { - Tcompute dot = 0; - for (int dk = 0; dk < static_cast(Dk); ++dk) { - dot += k_buf[i * Dk + dk] * beta_buf[i] * k_buf[j * Dk + dk]; - } - attn[x] = -dot * expf(g_cum[i] - g_cum[j]); - } else { - attn[x] = 0; - } - } - __syncthreads(); - - // Triangular solve-like correction. - // Sequential in i, parallel in j. - for (int i = 1; i < static_cast(chunk_size); ++i) { - CGDR_FOR(m, chunk_size) { - row_buf[m] = m < i ? attn[i * chunk_size + m] : 0; - } - __syncthreads(); - - for (int j = threadIdx.x; j < i; j += blockDim.x) { - Tcompute correction = 0; - for (int m = 0; m < i; ++m) { - correction += row_buf[m] * attn[m * chunk_size + j]; - } - attn[i * chunk_size + j] = row_buf[j] + correction; - } - __syncthreads(); - } - - CGDR_FOR(i, chunk_size) { - attn[i * chunk_size + i] = 1; - } - __syncthreads(); - - // v_mid = attn @ v_beta. - CGDR_FOR(x, chunk_size * Dv) { - int i = x / Dv; - int dv = x % Dv; - - Tcompute sum = 0; - for (int j = 0; j < static_cast(chunk_size); ++j) { - sum += attn[i * chunk_size + j] * v_beta[j * Dv + dv]; - } - v_mid[x] = sum; - } - - // k_cumdecay = attn @ (k * beta * exp(g)). - CGDR_FOR(x, chunk_size * Dk) { - int i = x / Dk; - int dk = x % Dk; - - Tcompute sum = 0; - for (int j = 0; j < static_cast(chunk_size); ++j) { - sum += attn[i * chunk_size + j] * k_buf[j * Dk + dk] * beta_buf[j] * expf(g_cum[j]); - } - k_cumdecay[x] = sum; - } - __syncthreads(); - - // v_new = v_mid - k_cumdecay @ state. - CGDR_FOR(x, chunk_size * Dv) { - int i = x / Dv; - int dv = x % Dv; - - Tcompute v_prime = 0; - for (int dk = 0; dk < static_cast(Dk); ++dk) { - v_prime += k_cumdecay[i * Dk + dk] * state_local[dk * Dv + dv]; - } - v_new[x] = v_mid[x] - v_prime; - } - __syncthreads(); - - // Output. - CGDR_FOR(x, actual_len * Dv) { - int i = x / Dv; - int dv = x % Dv; - - int64_t token_idx = chunk_begin + i; - ptrdiff_t out_base = static_cast(token_batch) * out_s0 + static_cast(token_idx) * out_s1 + static_cast(value_head_idx) * out_s2; - - Tcompute out_val = 0; - Tcompute q_decay = expf(g_cum[i]); - - for (int dk = 0; dk < static_cast(Dk); ++dk) { - out_val += q_buf[i * Dk + dk] * q_decay * state_local[dk * Dv + dv]; - } - - for (int j = 0; j <= i; ++j) { - Tcompute qk_attn = 0; - for (int dk = 0; dk < static_cast(Dk); ++dk) { - qk_attn += q_buf[i * Dk + dk] * k_buf[j * Dk + dk]; - } - - qk_attn *= expf(g_cum[i] - g_cum[j]); - out_val += qk_attn * v_new[j * Dv + dv]; - } - - out[out_base + dv] = static_cast(out_val); - } - __syncthreads(); - - // Update state. - const Tcompute last_decay = expf(g_cum[chunk_size - 1]); - - CGDR_FOR(x, Dk * Dv) { - int dk = x / Dv; - int dv = x % Dv; - - Tcompute next_state = state_local[x] * last_decay; - - for (int i = 0; i < static_cast(chunk_size); ++i) { - next_state += k_buf[i * Dk + dk] * expf(g_cum[chunk_size - 1] - g_cum[i]) * v_new[i * Dv + dv]; - } - - state_local[x] = next_state; - } - __syncthreads(); - } - - // Store final state. - CGDR_FOR(i, Dk * Dv) { - int dk = i / Dv; - int dv = i % Dv; - - const ptrdiff_t s2 = final_state_indices != nullptr ? initial_s2 : final_s2; - const ptrdiff_t s3 = final_state_indices != nullptr ? initial_s3 : final_s3; - ptrdiff_t write_idx = final_base + static_cast(dv) * s2 + static_cast(dk) * s3; - - final_state_target[write_idx] = static_cast(state_local[i]); - } -} - -template < - typename Tdata, - typename Tgate, - typename Tcompute, - size_t Dk, - size_t Dv, - size_t NUM_THREADS> -__device__ void chunkGatedDeltaRuleRecurrentKernel( - Tcompute *state_workspace, - Tdata *out, - Tdata *initial_state, - Tdata *final_state, - const Tdata *q, - const Tdata *k, - const Tdata *v, - const Tgate *g, - const Tgate *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - bool cu_seqlens_i64, - bool initial_state_indices_i64, - bool final_state_indices_i64, - bool use_qk_l2norm, - bool has_cu_seqlens, - bool indexed_state_pool, - size_t T, - size_t, - size_t pool_size, - size_t Hk, - size_t value_heads_per_key_head, - ptrdiff_t out_s0, - ptrdiff_t out_s1, - ptrdiff_t out_s2, - ptrdiff_t initial_s0, - ptrdiff_t initial_s1, - ptrdiff_t initial_s2, - ptrdiff_t initial_s3, - ptrdiff_t final_s0, - ptrdiff_t final_s1, - ptrdiff_t final_s2, - ptrdiff_t final_s3, - ptrdiff_t q_s0, - ptrdiff_t q_s1, - ptrdiff_t q_s2, - ptrdiff_t k_s0, - ptrdiff_t k_s1, - ptrdiff_t k_s2, - ptrdiff_t v_s0, - ptrdiff_t v_s1, - ptrdiff_t v_s2, - ptrdiff_t g_s0, - ptrdiff_t g_s1, - ptrdiff_t g_s2, - ptrdiff_t beta_s0, - ptrdiff_t beta_s1, - ptrdiff_t beta_s2) { - - const int batch_idx = blockIdx.x; - const int value_head_idx = blockIdx.y; - const int key_head_idx = value_head_idx / static_cast(value_heads_per_key_head); - - if (key_head_idx >= static_cast(Hk)) { - return; - } - - int64_t token_begin = 0; - int64_t token_end = static_cast(T); - - if (has_cu_seqlens) { - token_begin = loadOptionalIndex(cu_seqlens, cu_seqlens_i64, batch_idx, 0); - token_end = loadOptionalIndex(cu_seqlens, cu_seqlens_i64, batch_idx + 1, 0); - if (token_begin < 0 || token_end < token_begin || token_end > static_cast(T)) { - return; - } - } - - int64_t read_slot = batch_idx; - int64_t write_slot = batch_idx; - - if (indexed_state_pool) { - read_slot = loadOptionalIndex( - initial_state_indices, - initial_state_indices_i64, - batch_idx, - batch_idx); - - write_slot = final_state_indices == nullptr - ? static_cast(batch_idx) - : loadOptionalIndex( - final_state_indices, - final_state_indices_i64, - batch_idx, - batch_idx); - - if (read_slot < 0 || write_slot < 0 || read_slot >= static_cast(pool_size) || write_slot >= static_cast(pool_size)) { - return; - } - } - - const ptrdiff_t initial_base = indexed_state_pool - ? static_cast(read_slot) * initial_s0 + static_cast(value_head_idx) * initial_s1 - : static_cast(batch_idx) * initial_s0 + static_cast(value_head_idx) * initial_s1; - - Tdata *final_state_target = nullptr; - ptrdiff_t final_base = 0; - - if (indexed_state_pool && final_state_indices != nullptr) { - final_state_target = initial_state; - final_base = static_cast(write_slot) * initial_s0 + static_cast(value_head_idx) * initial_s1; - } else { - final_state_target = final_state; - final_base = static_cast(batch_idx) * final_s0 + static_cast(value_head_idx) * final_s1; - } - - const ptrdiff_t workspace_block = (static_cast(batch_idx) * gridDim.y + static_cast(value_head_idx)) * static_cast(Dk * Dv); - - Tcompute *state_local = state_workspace + workspace_block; - __shared__ Tcompute q_vec[Dk]; - __shared__ Tcompute k_vec[Dk]; - __shared__ Tcompute v_new[Dv]; - __shared__ Tcompute scalar_buf[3]; - - const int token_batch = has_cu_seqlens ? 0 : batch_idx; - const Tcompute scale = rsqrtf(static_cast(Dk)); - - CGDR_FOR(i, Dk * Dv) { - int dk = i / Dv; - int dv = i % Dv; - - ptrdiff_t read_idx = initial_base + static_cast(dv) * initial_s2 + static_cast(dk) * initial_s3; - - state_local[i] = static_cast( - loadAsFloat(initial_state, read_idx)); - } - __syncthreads(); - - for (int64_t token_idx = token_begin; token_idx < token_end; ++token_idx) { - ptrdiff_t q_base = static_cast(token_batch) * q_s0 + static_cast(token_idx) * q_s1 + static_cast(key_head_idx) * q_s2; - ptrdiff_t k_base = static_cast(token_batch) * k_s0 + static_cast(token_idx) * k_s1 + static_cast(key_head_idx) * k_s2; - - Tcompute q_sum = 0; - Tcompute k_sum = 0; - for (int dk = threadIdx.x; dk < static_cast(Dk); dk += blockDim.x) { - Tcompute q_raw = static_cast(loadAsFloat(q, q_base + dk)); - Tcompute k_raw = static_cast(loadAsFloat(k, k_base + dk)); - q_vec[dk] = q_raw; - k_vec[dk] = k_raw; - q_sum += q_raw * q_raw; - k_sum += k_raw * k_raw; - } - - q_sum = blockReduceSum(q_sum); - k_sum = blockReduceSum(k_sum); - - if (threadIdx.x == 0) { - ptrdiff_t gate_offset = static_cast(token_batch) * g_s0 + static_cast(token_idx) * g_s1 + static_cast(value_head_idx) * g_s2; - ptrdiff_t beta_offset = static_cast(token_batch) * beta_s0 + static_cast(token_idx) * beta_s1 + static_cast(value_head_idx) * beta_s2; - - scalar_buf[0] = expf(static_cast(loadAsFloat(g, gate_offset))); - scalar_buf[1] = static_cast(loadAsFloat(beta, beta_offset)); - scalar_buf[2] = use_qk_l2norm - ? rsqrtf(q_sum + static_cast(1e-6)) * scale - : scale; - } - __syncthreads(); - - const Tcompute decay = scalar_buf[0]; - const Tcompute beta_t = scalar_buf[1]; - const Tcompute q_scale = scalar_buf[2]; - const Tcompute k_scale = use_qk_l2norm - ? rsqrtf(k_sum + static_cast(1e-6)) - : static_cast(1); - - for (int dk = threadIdx.x; dk < static_cast(Dk); dk += blockDim.x) { - q_vec[dk] *= q_scale; - k_vec[dk] *= k_scale; - } - __syncthreads(); - - for (int dv = threadIdx.x; dv < static_cast(Dv); dv += blockDim.x) { - Tcompute projected = 0; - for (int dk = 0; dk < static_cast(Dk); ++dk) { - projected += k_vec[dk] * state_local[dk * Dv + dv]; - } - - ptrdiff_t v_base = static_cast(token_batch) * v_s0 + static_cast(token_idx) * v_s1 + static_cast(value_head_idx) * v_s2; - Tcompute v_raw = static_cast(loadAsFloat(v, v_base + dv)); - v_new[dv] = beta_t * (v_raw - decay * projected); - } - __syncthreads(); - - CGDR_FOR(x, Dk * Dv) { - int dk = x / Dv; - int dv = x % Dv; - state_local[x] = decay * state_local[x] + k_vec[dk] * v_new[dv]; - } - __syncthreads(); - - for (int dv = threadIdx.x; dv < static_cast(Dv); dv += blockDim.x) { - Tcompute out_val = 0; - for (int dk = 0; dk < static_cast(Dk); ++dk) { - out_val += q_vec[dk] * state_local[dk * Dv + dv]; - } - - ptrdiff_t out_base = static_cast(token_batch) * out_s0 + static_cast(token_idx) * out_s1 + static_cast(value_head_idx) * out_s2; - out[out_base + dv] = static_cast(out_val); - } - __syncthreads(); - } - - CGDR_FOR(i, Dk * Dv) { - int dk = i / Dv; - int dv = i % Dv; - - const ptrdiff_t s2 = final_state_indices != nullptr ? initial_s2 : final_s2; - const ptrdiff_t s3 = final_state_indices != nullptr ? initial_s3 : final_s3; - ptrdiff_t write_idx = final_base + static_cast(dv) * s2 + static_cast(dk) * s3; - - final_state_target[write_idx] = static_cast(state_local[i]); - } -} - -#endif // __CHUNK_GATED_DELTA_RULE_KERNEL_CUH__ diff --git a/src/infiniop/ops/chunk_gated_delta_rule/info.h b/src/infiniop/ops/chunk_gated_delta_rule/info.h deleted file mode 100644 index 101da34ae..000000000 --- a/src/infiniop/ops/chunk_gated_delta_rule/info.h +++ /dev/null @@ -1,207 +0,0 @@ -// infiniop/ops/chunk_gated_delta_rule/info.h - -#ifndef __CHUNK_GATED_DELTA_RULE_INFO_H__ -#define __CHUNK_GATED_DELTA_RULE_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op { -namespace chunk_gated_delta_rule { - -class ChunkGatedDeltaRuleInfo { - ChunkGatedDeltaRuleInfo() = default; - -public: - infiniDtype_t data_dtype; - infiniDtype_t gate_dtype; - infiniDtype_t cu_seqlens_dtype; - infiniDtype_t initial_state_indices_dtype; - infiniDtype_t final_state_indices_dtype; - - bool use_qk_l2norm; - bool has_cu_seqlens; - bool has_initial_state_indices; - bool has_final_state_indices; - bool indexed_state_pool; - - size_t B, T, total_tokens, Hk, Hv, Dk, Dv, chunk_size, pool_size, value_heads_per_key_head; - - std::vector out_strides; - std::vector initial_state_strides; - std::vector final_state_strides; - std::vector q_strides; - std::vector k_strides; - std::vector v_strides; - std::vector g_strides; - std::vector beta_strides; - - static utils::Result - create(infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t initial_state_desc, - infiniopTensorDescriptor_t final_state_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t g_desc, - infiniopTensorDescriptor_t beta_desc, - infiniopTensorDescriptor_t cu_seqlens_desc, - infiniopTensorDescriptor_t initial_state_indices_desc, - infiniopTensorDescriptor_t final_state_indices_desc, - bool use_qk_l2norm, - size_t chunk_size) { - - if (out_desc == nullptr || initial_state_desc == nullptr || q_desc == nullptr || k_desc == nullptr || v_desc == nullptr || g_desc == nullptr || beta_desc == nullptr) { - return INFINI_STATUS_NULL_POINTER; - } - - if (chunk_size == 0) { - return INFINI_STATUS_BAD_PARAM; - } - - auto data_dtype = q_desc->dtype(); - CHECK_DTYPE(data_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - if (k_desc->dtype() != data_dtype || v_desc->dtype() != data_dtype || out_desc->dtype() != data_dtype || initial_state_desc->dtype() != data_dtype || (final_state_desc != nullptr && final_state_desc->dtype() != data_dtype)) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto gate_dtype = g_desc->dtype(); - CHECK_DTYPE(gate_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - if (beta_desc->dtype() != gate_dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - bool has_cu = cu_seqlens_desc != nullptr; - bool has_initial_indices = initial_state_indices_desc != nullptr; - bool has_final_indices = final_state_indices_desc != nullptr; - bool indexed_pool = has_initial_indices || has_final_indices; - - if (has_final_indices && final_state_desc != nullptr) { - return INFINI_STATUS_BAD_PARAM; - } - if (!has_final_indices && final_state_desc == nullptr) { - return INFINI_STATUS_NULL_POINTER; - } - - if (q_desc->ndim() != 4 || k_desc->ndim() != 4 || v_desc->ndim() != 4 || out_desc->ndim() != 4 || g_desc->ndim() != 3 || beta_desc->ndim() != 3 || initial_state_desc->ndim() != 4 || (!has_final_indices && final_state_desc->ndim() != 4)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - auto q_shape = q_desc->shape(); - auto k_shape = k_desc->shape(); - auto v_shape = v_desc->shape(); - auto out_shape = out_desc->shape(); - auto g_shape = g_desc->shape(); - auto beta_shape = beta_desc->shape(); - - size_t B = q_shape[0], T = q_shape[1], Hk = q_shape[2], Dk = q_shape[3]; - size_t Hv = v_shape[2], Dv = v_shape[3], total_tokens = T; - - if (has_cu) { - if (cu_seqlens_desc->ndim() != 1 || cu_seqlens_desc->shape()[0] < 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - B = cu_seqlens_desc->shape()[0] - 1; - if (q_shape[0] != 1 || k_shape[0] != 1 || v_shape[0] != 1 || out_shape[0] != 1 || g_shape[0] != 1 || beta_shape[0] != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - total_tokens = q_shape[1]; - T = total_tokens; - } - - if (k_shape[0] != q_shape[0] || k_shape[1] != q_shape[1] || k_shape[2] != Hk || k_shape[3] != Dk || v_shape[0] != q_shape[0] || v_shape[1] != q_shape[1] || out_shape[0] != q_shape[0] || out_shape[1] != q_shape[1] || out_shape[2] != Hv || out_shape[3] != Dv || g_shape[0] != q_shape[0] || g_shape[1] != q_shape[1] || g_shape[2] != Hv || beta_shape[0] != q_shape[0] || beta_shape[1] != q_shape[1] || beta_shape[2] != Hv || Hk == 0 || Hv == 0 || Hv % Hk != 0) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (q_desc->strides()[3] != 1 || k_desc->strides()[3] != 1 || v_desc->strides()[3] != 1 || out_desc->strides()[3] != 1) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - auto initial_shape = initial_state_desc->shape(); - size_t pool_size = initial_shape[0]; - if (indexed_pool) { - if (initial_shape[1] != Hv || initial_shape[2] != Dv || initial_shape[3] != Dk) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } else { - if (initial_shape[0] != B || initial_shape[1] != Hv || initial_shape[2] != Dv || initial_shape[3] != Dk) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - if (!has_final_indices) { - auto final_shape = final_state_desc->shape(); - if (indexed_pool) { - if (final_shape[0] != B || final_shape[1] != Hv || final_shape[2] != Dv || final_shape[3] != Dk) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } else { - if (final_shape[0] != B || final_shape[1] != Hv || final_shape[2] != Dv || final_shape[3] != Dk) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - } - - infiniDtype_t cu_dtype = INFINI_DTYPE_INVALID; - infiniDtype_t initial_indices_dtype = INFINI_DTYPE_INVALID; - infiniDtype_t final_indices_dtype = INFINI_DTYPE_INVALID; - if (has_cu) { - cu_dtype = cu_seqlens_desc->dtype(); - CHECK_DTYPE(cu_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - } - if (has_initial_indices) { - if (initial_state_indices_desc->ndim() != 1 || initial_state_indices_desc->shape()[0] != B) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - initial_indices_dtype = initial_state_indices_desc->dtype(); - CHECK_DTYPE(initial_indices_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - } - if (has_final_indices) { - if (final_state_indices_desc->ndim() != 1 || final_state_indices_desc->shape()[0] != B) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - final_indices_dtype = final_state_indices_desc->dtype(); - CHECK_DTYPE(final_indices_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - } - - ChunkGatedDeltaRuleInfo info; - info.data_dtype = data_dtype; - info.gate_dtype = gate_dtype; - info.cu_seqlens_dtype = cu_dtype; - info.initial_state_indices_dtype = initial_indices_dtype; - info.final_state_indices_dtype = final_indices_dtype; - info.use_qk_l2norm = use_qk_l2norm; - info.has_cu_seqlens = has_cu; - info.has_initial_state_indices = has_initial_indices; - info.has_final_state_indices = has_final_indices; - info.indexed_state_pool = indexed_pool; - info.B = B; - info.T = T; - info.total_tokens = total_tokens; - info.Hk = Hk; - info.Hv = Hv; - info.Dk = Dk; - info.Dv = Dv; - info.chunk_size = chunk_size; - info.pool_size = pool_size; - info.value_heads_per_key_head = Hv / Hk; - info.out_strides = out_desc->strides(); - info.initial_state_strides = initial_state_desc->strides(); - if (final_state_desc != nullptr) { - info.final_state_strides = final_state_desc->strides(); - } - info.q_strides = q_desc->strides(); - info.k_strides = k_desc->strides(); - info.v_strides = v_desc->strides(); - info.g_strides = g_desc->strides(); - info.beta_strides = beta_desc->strides(); - - return utils::Result(info); - } -}; - -} // namespace chunk_gated_delta_rule -} // namespace op - -#endif // __CHUNK_GATED_DELTA_RULE_INFO_H__ diff --git a/src/infiniop/ops/chunk_gated_delta_rule/metax/chunk_gated_delta_rule_metax.h b/src/infiniop/ops/chunk_gated_delta_rule/metax/chunk_gated_delta_rule_metax.h deleted file mode 100644 index e9a161436..000000000 --- a/src/infiniop/ops/chunk_gated_delta_rule/metax/chunk_gated_delta_rule_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CHUNK_GATED_DELTA_RULE_METAX_H__ -#define __CHUNK_GATED_DELTA_RULE_METAX_H__ - -#include "../chunk_gated_delta_rule.h" - -DESCRIPTOR(metax) - -#endif diff --git a/src/infiniop/ops/chunk_gated_delta_rule/metax/chunk_gated_delta_rule_metax.maca b/src/infiniop/ops/chunk_gated_delta_rule/metax/chunk_gated_delta_rule_metax.maca deleted file mode 100644 index 991bac0db..000000000 --- a/src/infiniop/ops/chunk_gated_delta_rule/metax/chunk_gated_delta_rule_metax.maca +++ /dev/null @@ -1,304 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "chunk_gated_delta_rule_metax.h" - -#include "../cuda/kernel.cuh" - -template -INFINIOP_METAX_KERNEL chunkGatedDeltaRule( - Tcompute *state_workspace, - Tdata *out, - Tdata *initial_state, - Tdata *final_state, - const Tdata *q, - const Tdata *k, - const Tdata *v, - const Tgate *g, - const Tgate *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - bool cu_seqlens_i64, - bool initial_state_indices_i64, - bool final_state_indices_i64, - bool use_qk_l2norm, - bool has_cu_seqlens, - bool indexed_state_pool, - size_t T, - size_t chunk_size, - size_t pool_size, - size_t Hk, - size_t value_heads_per_key_head, - ptrdiff_t out_s0, - ptrdiff_t out_s1, - ptrdiff_t out_s2, - ptrdiff_t initial_s0, - ptrdiff_t initial_s1, - ptrdiff_t initial_s2, - ptrdiff_t initial_s3, - ptrdiff_t final_s0, - ptrdiff_t final_s1, - ptrdiff_t final_s2, - ptrdiff_t final_s3, - ptrdiff_t q_s0, - ptrdiff_t q_s1, - ptrdiff_t q_s2, - ptrdiff_t k_s0, - ptrdiff_t k_s1, - ptrdiff_t k_s2, - ptrdiff_t v_s0, - ptrdiff_t v_s1, - ptrdiff_t v_s2, - ptrdiff_t g_s0, - ptrdiff_t g_s1, - ptrdiff_t g_s2, - ptrdiff_t beta_s0, - ptrdiff_t beta_s1, - ptrdiff_t beta_s2) { - chunkGatedDeltaRuleRecurrentKernel( - state_workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, cu_seqlens_i64, - initial_state_indices_i64, final_state_indices_i64, use_qk_l2norm, - has_cu_seqlens, indexed_state_pool, T, chunk_size, pool_size, Hk, value_heads_per_key_head, - out_s0, out_s1, out_s2, initial_s0, initial_s1, initial_s2, initial_s3, - final_s0, final_s1, final_s2, final_s3, q_s0, q_s1, q_s2, k_s0, k_s1, - k_s2, v_s0, v_s1, v_s2, g_s0, g_s1, g_s2, beta_s0, beta_s1, beta_s2); -} - -template -INFINIOP_METAX_KERNEL chunkGatedDeltaRuleChunked( - Tcompute *state_workspace, - Tdata *out, - Tdata *initial_state, - Tdata *final_state, - const Tdata *q, - const Tdata *k, - const Tdata *v, - const Tgate *g, - const Tgate *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - bool cu_seqlens_i64, - bool initial_state_indices_i64, - bool final_state_indices_i64, - bool use_qk_l2norm, - bool has_cu_seqlens, - bool indexed_state_pool, - size_t T, - size_t chunk_size, - size_t pool_size, - size_t Hk, - size_t value_heads_per_key_head, - ptrdiff_t out_s0, - ptrdiff_t out_s1, - ptrdiff_t out_s2, - ptrdiff_t initial_s0, - ptrdiff_t initial_s1, - ptrdiff_t initial_s2, - ptrdiff_t initial_s3, - ptrdiff_t final_s0, - ptrdiff_t final_s1, - ptrdiff_t final_s2, - ptrdiff_t final_s3, - ptrdiff_t q_s0, - ptrdiff_t q_s1, - ptrdiff_t q_s2, - ptrdiff_t k_s0, - ptrdiff_t k_s1, - ptrdiff_t k_s2, - ptrdiff_t v_s0, - ptrdiff_t v_s1, - ptrdiff_t v_s2, - ptrdiff_t g_s0, - ptrdiff_t g_s1, - ptrdiff_t g_s2, - ptrdiff_t beta_s0, - ptrdiff_t beta_s1, - ptrdiff_t beta_s2) { - chunkGatedDeltaRuleKernel( - state_workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, cu_seqlens_i64, - initial_state_indices_i64, final_state_indices_i64, use_qk_l2norm, - has_cu_seqlens, indexed_state_pool, T, chunk_size, pool_size, Hk, value_heads_per_key_head, - out_s0, out_s1, out_s2, initial_s0, initial_s1, initial_s2, initial_s3, - final_s0, final_s1, final_s2, final_s3, q_s0, q_s1, q_s2, k_s0, k_s1, - k_s2, v_s0, v_s1, v_s2, g_s0, g_s1, g_s2, beta_s0, beta_s1, beta_s2); -} - -namespace op { -namespace chunk_gated_delta_rule { -namespace metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t initial_state_desc, - infiniopTensorDescriptor_t final_state_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t g_desc, - infiniopTensorDescriptor_t beta_desc, - infiniopTensorDescriptor_t cu_seqlens_desc, - infiniopTensorDescriptor_t initial_state_indices_desc, - infiniopTensorDescriptor_t final_state_indices_desc, - bool use_qk_l2norm, - size_t chunk_size) { - auto info = ChunkGatedDeltaRuleInfo::create( - out_desc, initial_state_desc, final_state_desc, q_desc, k_desc, v_desc, - g_desc, beta_desc, cu_seqlens_desc, initial_state_indices_desc, - final_state_indices_desc, use_qk_l2norm, chunk_size); - CHECK_RESULT(info); - - auto info_value = info.take(); - // We always want to use fast path, slow path is kept as a ref - const bool use_chunked_fallback = false; - const size_t per_block_workspace = use_chunked_fallback - ? info_value.Dk * info_value.Dv + info_value.chunk_size * info_value.Dk * 3 + info_value.chunk_size * info_value.Dv * 3 + info_value.chunk_size * info_value.chunk_size + info_value.chunk_size * 3 - : info_value.Dk * info_value.Dv; - const size_t workspace_size = info_value.B * info_value.Hv * per_block_workspace * sizeof(float); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info_value, workspace_size, handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernelWithGateDtype( - void *workspace, - void *out, - void *initial_state, - void *final_state, - const void *q, - const void *k, - const void *v, - const void *g, - const void *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - const ChunkGatedDeltaRuleInfo &info, - hcStream_t stream) { - -#define LAUNCH_ARGS(TYPE) \ - static_cast(workspace), static_cast(out), static_cast(initial_state), static_cast(final_state), \ - static_cast(q), static_cast(k), static_cast(v), \ - static_cast(g), static_cast(beta), cu_seqlens, initial_state_indices, \ - final_state_indices, info.cu_seqlens_dtype == INFINI_DTYPE_I64, \ - info.initial_state_indices_dtype == INFINI_DTYPE_I64, info.final_state_indices_dtype == INFINI_DTYPE_I64, \ - info.use_qk_l2norm, info.has_cu_seqlens, info.indexed_state_pool, info.T, info.chunk_size, info.pool_size, info.Hk, \ - info.value_heads_per_key_head, info.out_strides[0], info.out_strides[1], info.out_strides[2], \ - info.initial_state_strides[0], info.initial_state_strides[1], info.initial_state_strides[2], \ - info.initial_state_strides[3], info.final_state_strides.empty() ? 0 : info.final_state_strides[0], \ - info.final_state_strides.empty() ? 0 : info.final_state_strides[1], info.final_state_strides.empty() ? 0 : info.final_state_strides[2], \ - info.final_state_strides.empty() ? 0 : info.final_state_strides[3], info.q_strides[0], info.q_strides[1], \ - info.q_strides[2], info.k_strides[0], info.k_strides[1], info.k_strides[2], info.v_strides[0], \ - info.v_strides[1], info.v_strides[2], info.g_strides[0], info.g_strides[1], info.g_strides[2], \ - info.beta_strides[0], info.beta_strides[1], info.beta_strides[2] - -#define LAUNCH_GATE(TYPE) \ - do { \ - if (false) { \ - chunkGatedDeltaRuleChunked \ - <<>>(LAUNCH_ARGS(TYPE)); \ - } else { \ - chunkGatedDeltaRule \ - <<>>(LAUNCH_ARGS(TYPE)); \ - } \ - } while (0) - - if (info.gate_dtype == INFINI_DTYPE_F16) { - LAUNCH_GATE(half); - } else if (info.gate_dtype == INFINI_DTYPE_BF16) { - LAUNCH_GATE(cuda_bfloat16); - } else if (info.gate_dtype == INFINI_DTYPE_F32) { - LAUNCH_GATE(float); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -#undef LAUNCH_GATE -#undef LAUNCH_ARGS - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel( - void *workspace, - void *out, - void *initial_state, - void *final_state, - const void *q, - const void *k, - const void *v, - const void *g, - const void *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - const ChunkGatedDeltaRuleInfo &info, - hcStream_t stream) { - if (info.data_dtype == INFINI_DTYPE_F16) { - return launchKernelWithGateDtype( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, info, stream); - } - if (info.data_dtype == INFINI_DTYPE_BF16) { - return launchKernelWithGateDtype( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, info, stream); - } - if (info.data_dtype == INFINI_DTYPE_F32) { - return launchKernelWithGateDtype( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, info, stream); - } - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *out, void *initial_state, void *final_state, - const void *q, const void *k, const void *v, - const void *g, const void *beta, const void *cu_seqlens, - const void *initial_state_indices, const void *final_state_indices, - void *stream_) const { - hcStream_t stream = (hcStream_t)stream_; - if (workspace == nullptr || workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - if (_info.Dk == 128 && _info.Dv == 128) { - if (_opaque->internal->maxThreadsPerBlock() >= 128) { - return launchKernel<128, 128, 128>( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, _info, stream); - } - } else if (_info.Dk == 64 && _info.Dv == 64) { - if (_opaque->internal->maxThreadsPerBlock() >= 64) { - return launchKernel<64, 64, 64>( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, _info, stream); - } - } - - return INFINI_STATUS_BAD_TENSOR_SHAPE; -} - -} // namespace metax -} // namespace chunk_gated_delta_rule -} // namespace op diff --git a/src/infiniop/ops/chunk_gated_delta_rule/moore/chunk_gated_delta_rule_moore.h b/src/infiniop/ops/chunk_gated_delta_rule/moore/chunk_gated_delta_rule_moore.h deleted file mode 100644 index 87938a8d7..000000000 --- a/src/infiniop/ops/chunk_gated_delta_rule/moore/chunk_gated_delta_rule_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CHUNK_GATED_DELTA_RULE_MOORE_H__ -#define __CHUNK_GATED_DELTA_RULE_MOORE_H__ - -#include "../chunk_gated_delta_rule.h" - -DESCRIPTOR(moore) - -#endif diff --git a/src/infiniop/ops/chunk_gated_delta_rule/moore/chunk_gated_delta_rule_moore.mu b/src/infiniop/ops/chunk_gated_delta_rule/moore/chunk_gated_delta_rule_moore.mu deleted file mode 100644 index 15fafedde..000000000 --- a/src/infiniop/ops/chunk_gated_delta_rule/moore/chunk_gated_delta_rule_moore.mu +++ /dev/null @@ -1,304 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "chunk_gated_delta_rule_moore.h" - -#include "../cuda/kernel.cuh" - -template -INFINIOP_MOORE_KERNEL chunkGatedDeltaRule( - Tcompute *state_workspace, - Tdata *out, - Tdata *initial_state, - Tdata *final_state, - const Tdata *q, - const Tdata *k, - const Tdata *v, - const Tgate *g, - const Tgate *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - bool cu_seqlens_i64, - bool initial_state_indices_i64, - bool final_state_indices_i64, - bool use_qk_l2norm, - bool has_cu_seqlens, - bool indexed_state_pool, - size_t T, - size_t chunk_size, - size_t pool_size, - size_t Hk, - size_t value_heads_per_key_head, - ptrdiff_t out_s0, - ptrdiff_t out_s1, - ptrdiff_t out_s2, - ptrdiff_t initial_s0, - ptrdiff_t initial_s1, - ptrdiff_t initial_s2, - ptrdiff_t initial_s3, - ptrdiff_t final_s0, - ptrdiff_t final_s1, - ptrdiff_t final_s2, - ptrdiff_t final_s3, - ptrdiff_t q_s0, - ptrdiff_t q_s1, - ptrdiff_t q_s2, - ptrdiff_t k_s0, - ptrdiff_t k_s1, - ptrdiff_t k_s2, - ptrdiff_t v_s0, - ptrdiff_t v_s1, - ptrdiff_t v_s2, - ptrdiff_t g_s0, - ptrdiff_t g_s1, - ptrdiff_t g_s2, - ptrdiff_t beta_s0, - ptrdiff_t beta_s1, - ptrdiff_t beta_s2) { - chunkGatedDeltaRuleRecurrentKernel( - state_workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, cu_seqlens_i64, - initial_state_indices_i64, final_state_indices_i64, use_qk_l2norm, - has_cu_seqlens, indexed_state_pool, T, chunk_size, pool_size, Hk, value_heads_per_key_head, - out_s0, out_s1, out_s2, initial_s0, initial_s1, initial_s2, initial_s3, - final_s0, final_s1, final_s2, final_s3, q_s0, q_s1, q_s2, k_s0, k_s1, - k_s2, v_s0, v_s1, v_s2, g_s0, g_s1, g_s2, beta_s0, beta_s1, beta_s2); -} - -template -INFINIOP_MOORE_KERNEL chunkGatedDeltaRuleChunked( - Tcompute *state_workspace, - Tdata *out, - Tdata *initial_state, - Tdata *final_state, - const Tdata *q, - const Tdata *k, - const Tdata *v, - const Tgate *g, - const Tgate *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - bool cu_seqlens_i64, - bool initial_state_indices_i64, - bool final_state_indices_i64, - bool use_qk_l2norm, - bool has_cu_seqlens, - bool indexed_state_pool, - size_t T, - size_t chunk_size, - size_t pool_size, - size_t Hk, - size_t value_heads_per_key_head, - ptrdiff_t out_s0, - ptrdiff_t out_s1, - ptrdiff_t out_s2, - ptrdiff_t initial_s0, - ptrdiff_t initial_s1, - ptrdiff_t initial_s2, - ptrdiff_t initial_s3, - ptrdiff_t final_s0, - ptrdiff_t final_s1, - ptrdiff_t final_s2, - ptrdiff_t final_s3, - ptrdiff_t q_s0, - ptrdiff_t q_s1, - ptrdiff_t q_s2, - ptrdiff_t k_s0, - ptrdiff_t k_s1, - ptrdiff_t k_s2, - ptrdiff_t v_s0, - ptrdiff_t v_s1, - ptrdiff_t v_s2, - ptrdiff_t g_s0, - ptrdiff_t g_s1, - ptrdiff_t g_s2, - ptrdiff_t beta_s0, - ptrdiff_t beta_s1, - ptrdiff_t beta_s2) { - chunkGatedDeltaRuleKernel( - state_workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, cu_seqlens_i64, - initial_state_indices_i64, final_state_indices_i64, use_qk_l2norm, - has_cu_seqlens, indexed_state_pool, T, chunk_size, pool_size, Hk, value_heads_per_key_head, - out_s0, out_s1, out_s2, initial_s0, initial_s1, initial_s2, initial_s3, - final_s0, final_s1, final_s2, final_s3, q_s0, q_s1, q_s2, k_s0, k_s1, - k_s2, v_s0, v_s1, v_s2, g_s0, g_s1, g_s2, beta_s0, beta_s1, beta_s2); -} - -namespace op { -namespace chunk_gated_delta_rule { -namespace moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t initial_state_desc, - infiniopTensorDescriptor_t final_state_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t g_desc, - infiniopTensorDescriptor_t beta_desc, - infiniopTensorDescriptor_t cu_seqlens_desc, - infiniopTensorDescriptor_t initial_state_indices_desc, - infiniopTensorDescriptor_t final_state_indices_desc, - bool use_qk_l2norm, - size_t chunk_size) { - auto info = ChunkGatedDeltaRuleInfo::create( - out_desc, initial_state_desc, final_state_desc, q_desc, k_desc, v_desc, - g_desc, beta_desc, cu_seqlens_desc, initial_state_indices_desc, - final_state_indices_desc, use_qk_l2norm, chunk_size); - CHECK_RESULT(info); - - auto info_value = info.take(); - // We always want to use fast path, slow path is kept as a ref - const bool use_chunked_fallback = false; - const size_t per_block_workspace = use_chunked_fallback - ? info_value.Dk * info_value.Dv + info_value.chunk_size * info_value.Dk * 3 + info_value.chunk_size * info_value.Dv * 3 + info_value.chunk_size * info_value.chunk_size + info_value.chunk_size * 3 - : info_value.Dk * info_value.Dv; - const size_t workspace_size = info_value.B * info_value.Hv * per_block_workspace * sizeof(float); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info_value, workspace_size, handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernelWithGateDtype( - void *workspace, - void *out, - void *initial_state, - void *final_state, - const void *q, - const void *k, - const void *v, - const void *g, - const void *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - const ChunkGatedDeltaRuleInfo &info, - musaStream_t stream) { - -#define LAUNCH_ARGS(TYPE) \ - static_cast(workspace), static_cast(out), static_cast(initial_state), static_cast(final_state), \ - static_cast(q), static_cast(k), static_cast(v), \ - static_cast(g), static_cast(beta), cu_seqlens, initial_state_indices, \ - final_state_indices, info.cu_seqlens_dtype == INFINI_DTYPE_I64, \ - info.initial_state_indices_dtype == INFINI_DTYPE_I64, info.final_state_indices_dtype == INFINI_DTYPE_I64, \ - info.use_qk_l2norm, info.has_cu_seqlens, info.indexed_state_pool, info.T, info.chunk_size, info.pool_size, info.Hk, \ - info.value_heads_per_key_head, info.out_strides[0], info.out_strides[1], info.out_strides[2], \ - info.initial_state_strides[0], info.initial_state_strides[1], info.initial_state_strides[2], \ - info.initial_state_strides[3], info.final_state_strides.empty() ? 0 : info.final_state_strides[0], \ - info.final_state_strides.empty() ? 0 : info.final_state_strides[1], info.final_state_strides.empty() ? 0 : info.final_state_strides[2], \ - info.final_state_strides.empty() ? 0 : info.final_state_strides[3], info.q_strides[0], info.q_strides[1], \ - info.q_strides[2], info.k_strides[0], info.k_strides[1], info.k_strides[2], info.v_strides[0], \ - info.v_strides[1], info.v_strides[2], info.g_strides[0], info.g_strides[1], info.g_strides[2], \ - info.beta_strides[0], info.beta_strides[1], info.beta_strides[2] - -#define LAUNCH_GATE(TYPE) \ - do { \ - if (false) { \ - chunkGatedDeltaRuleChunked \ - <<>>(LAUNCH_ARGS(TYPE)); \ - } else { \ - chunkGatedDeltaRule \ - <<>>(LAUNCH_ARGS(TYPE)); \ - } \ - } while (0) - - if (info.gate_dtype == INFINI_DTYPE_F16) { - LAUNCH_GATE(half); - } else if (info.gate_dtype == INFINI_DTYPE_BF16) { - LAUNCH_GATE(cuda_bfloat16); - } else if (info.gate_dtype == INFINI_DTYPE_F32) { - LAUNCH_GATE(float); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -#undef LAUNCH_GATE -#undef LAUNCH_ARGS - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel( - void *workspace, - void *out, - void *initial_state, - void *final_state, - const void *q, - const void *k, - const void *v, - const void *g, - const void *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - const ChunkGatedDeltaRuleInfo &info, - musaStream_t stream) { - if (info.data_dtype == INFINI_DTYPE_F16) { - return launchKernelWithGateDtype( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, info, stream); - } - if (info.data_dtype == INFINI_DTYPE_BF16) { - return launchKernelWithGateDtype( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, info, stream); - } - if (info.data_dtype == INFINI_DTYPE_F32) { - return launchKernelWithGateDtype( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, info, stream); - } - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *out, void *initial_state, void *final_state, - const void *q, const void *k, const void *v, - const void *g, const void *beta, const void *cu_seqlens, - const void *initial_state_indices, const void *final_state_indices, - void *stream_) const { - musaStream_t stream = (musaStream_t)stream_; - if (workspace == nullptr || workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - if (_info.Dk == 128 && _info.Dv == 128) { - if (_opaque->internal->maxThreadsPerBlock() >= 128) { - return launchKernel<128, 128, 128>( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, _info, stream); - } - } else if (_info.Dk == 64 && _info.Dv == 64) { - if (_opaque->internal->maxThreadsPerBlock() >= 64) { - return launchKernel<64, 64, 64>( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, _info, stream); - } - } - - return INFINI_STATUS_BAD_TENSOR_SHAPE; -} - -} // namespace moore -} // namespace chunk_gated_delta_rule -} // namespace op diff --git a/src/infiniop/ops/chunk_gated_delta_rule/nvidia/chunk_gated_delta_rule_nvidia.cu b/src/infiniop/ops/chunk_gated_delta_rule/nvidia/chunk_gated_delta_rule_nvidia.cu deleted file mode 100644 index 461f95ad5..000000000 --- a/src/infiniop/ops/chunk_gated_delta_rule/nvidia/chunk_gated_delta_rule_nvidia.cu +++ /dev/null @@ -1,305 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "chunk_gated_delta_rule_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include - -template -INFINIOP_CUDA_KERNEL chunkGatedDeltaRule( - Tcompute *state_workspace, - Tdata *out, - Tdata *initial_state, - Tdata *final_state, - const Tdata *q, - const Tdata *k, - const Tdata *v, - const Tgate *g, - const Tgate *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - bool cu_seqlens_i64, - bool initial_state_indices_i64, - bool final_state_indices_i64, - bool use_qk_l2norm, - bool has_cu_seqlens, - bool indexed_state_pool, - size_t T, - size_t chunk_size, - size_t pool_size, - size_t Hk, - size_t value_heads_per_key_head, - ptrdiff_t out_s0, - ptrdiff_t out_s1, - ptrdiff_t out_s2, - ptrdiff_t initial_s0, - ptrdiff_t initial_s1, - ptrdiff_t initial_s2, - ptrdiff_t initial_s3, - ptrdiff_t final_s0, - ptrdiff_t final_s1, - ptrdiff_t final_s2, - ptrdiff_t final_s3, - ptrdiff_t q_s0, - ptrdiff_t q_s1, - ptrdiff_t q_s2, - ptrdiff_t k_s0, - ptrdiff_t k_s1, - ptrdiff_t k_s2, - ptrdiff_t v_s0, - ptrdiff_t v_s1, - ptrdiff_t v_s2, - ptrdiff_t g_s0, - ptrdiff_t g_s1, - ptrdiff_t g_s2, - ptrdiff_t beta_s0, - ptrdiff_t beta_s1, - ptrdiff_t beta_s2) { - chunkGatedDeltaRuleRecurrentKernel( - state_workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, cu_seqlens_i64, - initial_state_indices_i64, final_state_indices_i64, use_qk_l2norm, - has_cu_seqlens, indexed_state_pool, T, chunk_size, pool_size, Hk, value_heads_per_key_head, - out_s0, out_s1, out_s2, initial_s0, initial_s1, initial_s2, initial_s3, - final_s0, final_s1, final_s2, final_s3, q_s0, q_s1, q_s2, k_s0, k_s1, - k_s2, v_s0, v_s1, v_s2, g_s0, g_s1, g_s2, beta_s0, beta_s1, beta_s2); -} - -template -INFINIOP_CUDA_KERNEL chunkGatedDeltaRuleChunked( - Tcompute *state_workspace, - Tdata *out, - Tdata *initial_state, - Tdata *final_state, - const Tdata *q, - const Tdata *k, - const Tdata *v, - const Tgate *g, - const Tgate *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - bool cu_seqlens_i64, - bool initial_state_indices_i64, - bool final_state_indices_i64, - bool use_qk_l2norm, - bool has_cu_seqlens, - bool indexed_state_pool, - size_t T, - size_t chunk_size, - size_t pool_size, - size_t Hk, - size_t value_heads_per_key_head, - ptrdiff_t out_s0, - ptrdiff_t out_s1, - ptrdiff_t out_s2, - ptrdiff_t initial_s0, - ptrdiff_t initial_s1, - ptrdiff_t initial_s2, - ptrdiff_t initial_s3, - ptrdiff_t final_s0, - ptrdiff_t final_s1, - ptrdiff_t final_s2, - ptrdiff_t final_s3, - ptrdiff_t q_s0, - ptrdiff_t q_s1, - ptrdiff_t q_s2, - ptrdiff_t k_s0, - ptrdiff_t k_s1, - ptrdiff_t k_s2, - ptrdiff_t v_s0, - ptrdiff_t v_s1, - ptrdiff_t v_s2, - ptrdiff_t g_s0, - ptrdiff_t g_s1, - ptrdiff_t g_s2, - ptrdiff_t beta_s0, - ptrdiff_t beta_s1, - ptrdiff_t beta_s2) { - chunkGatedDeltaRuleKernel( - state_workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, cu_seqlens_i64, - initial_state_indices_i64, final_state_indices_i64, use_qk_l2norm, - has_cu_seqlens, indexed_state_pool, T, chunk_size, pool_size, Hk, value_heads_per_key_head, - out_s0, out_s1, out_s2, initial_s0, initial_s1, initial_s2, initial_s3, - final_s0, final_s1, final_s2, final_s3, q_s0, q_s1, q_s2, k_s0, k_s1, - k_s2, v_s0, v_s1, v_s2, g_s0, g_s1, g_s2, beta_s0, beta_s1, beta_s2); -} - -namespace op { -namespace chunk_gated_delta_rule { -namespace nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t initial_state_desc, - infiniopTensorDescriptor_t final_state_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t g_desc, - infiniopTensorDescriptor_t beta_desc, - infiniopTensorDescriptor_t cu_seqlens_desc, - infiniopTensorDescriptor_t initial_state_indices_desc, - infiniopTensorDescriptor_t final_state_indices_desc, - bool use_qk_l2norm, - size_t chunk_size) { - auto info = ChunkGatedDeltaRuleInfo::create( - out_desc, initial_state_desc, final_state_desc, q_desc, k_desc, v_desc, - g_desc, beta_desc, cu_seqlens_desc, initial_state_indices_desc, - final_state_indices_desc, use_qk_l2norm, chunk_size); - CHECK_RESULT(info); - - auto info_value = info.take(); - // We always want to use fast path, slow path is kept as a ref - const bool use_chunked_fallback = false; - const size_t per_block_workspace = use_chunked_fallback - ? info_value.Dk * info_value.Dv + info_value.chunk_size * info_value.Dk * 3 + info_value.chunk_size * info_value.Dv * 3 + info_value.chunk_size * info_value.chunk_size + info_value.chunk_size * 3 - : info_value.Dk * info_value.Dv; - const size_t workspace_size = info_value.B * info_value.Hv * per_block_workspace * sizeof(float); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info_value, workspace_size, handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernelWithGateDtype( - void *workspace, - void *out, - void *initial_state, - void *final_state, - const void *q, - const void *k, - const void *v, - const void *g, - const void *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - const ChunkGatedDeltaRuleInfo &info, - cudaStream_t stream) { - -#define LAUNCH_ARGS(TYPE) \ - static_cast(workspace), static_cast(out), static_cast(initial_state), static_cast(final_state), \ - static_cast(q), static_cast(k), static_cast(v), \ - static_cast(g), static_cast(beta), cu_seqlens, initial_state_indices, \ - final_state_indices, info.cu_seqlens_dtype == INFINI_DTYPE_I64, \ - info.initial_state_indices_dtype == INFINI_DTYPE_I64, info.final_state_indices_dtype == INFINI_DTYPE_I64, \ - info.use_qk_l2norm, info.has_cu_seqlens, info.indexed_state_pool, info.T, info.chunk_size, info.pool_size, info.Hk, \ - info.value_heads_per_key_head, info.out_strides[0], info.out_strides[1], info.out_strides[2], \ - info.initial_state_strides[0], info.initial_state_strides[1], info.initial_state_strides[2], \ - info.initial_state_strides[3], info.final_state_strides.empty() ? 0 : info.final_state_strides[0], \ - info.final_state_strides.empty() ? 0 : info.final_state_strides[1], info.final_state_strides.empty() ? 0 : info.final_state_strides[2], \ - info.final_state_strides.empty() ? 0 : info.final_state_strides[3], info.q_strides[0], info.q_strides[1], \ - info.q_strides[2], info.k_strides[0], info.k_strides[1], info.k_strides[2], info.v_strides[0], \ - info.v_strides[1], info.v_strides[2], info.g_strides[0], info.g_strides[1], info.g_strides[2], \ - info.beta_strides[0], info.beta_strides[1], info.beta_strides[2] - -#define LAUNCH_GATE(TYPE) \ - do { \ - if (false) { \ - chunkGatedDeltaRuleChunked \ - <<>>(LAUNCH_ARGS(TYPE)); \ - } else { \ - chunkGatedDeltaRule \ - <<>>(LAUNCH_ARGS(TYPE)); \ - } \ - } while (0) - - if (info.gate_dtype == INFINI_DTYPE_F16) { - LAUNCH_GATE(half); - } else if (info.gate_dtype == INFINI_DTYPE_BF16) { - LAUNCH_GATE(__nv_bfloat16); - } else if (info.gate_dtype == INFINI_DTYPE_F32) { - LAUNCH_GATE(float); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -#undef LAUNCH_GATE -#undef LAUNCH_ARGS - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel( - void *workspace, - void *out, - void *initial_state, - void *final_state, - const void *q, - const void *k, - const void *v, - const void *g, - const void *beta, - const void *cu_seqlens, - const void *initial_state_indices, - const void *final_state_indices, - const ChunkGatedDeltaRuleInfo &info, - cudaStream_t stream) { - if (info.data_dtype == INFINI_DTYPE_F16) { - return launchKernelWithGateDtype( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, info, stream); - } - if (info.data_dtype == INFINI_DTYPE_BF16) { - return launchKernelWithGateDtype<__nv_bfloat16, Dk, Dv, NUM_THREADS>( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, info, stream); - } - if (info.data_dtype == INFINI_DTYPE_F32) { - return launchKernelWithGateDtype( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, info, stream); - } - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *out, void *initial_state, void *final_state, - const void *q, const void *k, const void *v, - const void *g, const void *beta, const void *cu_seqlens, - const void *initial_state_indices, const void *final_state_indices, - void *stream_) const { - cudaStream_t stream = (cudaStream_t)stream_; - if (workspace == nullptr || workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - if (_info.Dk == 128 && _info.Dv == 128) { - if (_opaque->internal->maxThreadsPerBlock() >= 128) { - return launchKernel<128, 128, 128>( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, _info, stream); - } - } else if (_info.Dk == 64 && _info.Dv == 64) { - if (_opaque->internal->maxThreadsPerBlock() >= 64) { - return launchKernel<64, 64, 64>( - workspace, out, initial_state, final_state, q, k, v, g, beta, cu_seqlens, - initial_state_indices, final_state_indices, _info, stream); - } - } - - return INFINI_STATUS_BAD_TENSOR_SHAPE; -} - -} // namespace nvidia -} // namespace chunk_gated_delta_rule -} // namespace op diff --git a/src/infiniop/ops/chunk_gated_delta_rule/nvidia/chunk_gated_delta_rule_nvidia.cuh b/src/infiniop/ops/chunk_gated_delta_rule/nvidia/chunk_gated_delta_rule_nvidia.cuh deleted file mode 100644 index 5dde15f7d..000000000 --- a/src/infiniop/ops/chunk_gated_delta_rule/nvidia/chunk_gated_delta_rule_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CHUNK_GATED_DELTA_RULE_NVIDIA_H__ -#define __CHUNK_GATED_DELTA_RULE_NVIDIA_H__ - -#include "../chunk_gated_delta_rule.h" - -DESCRIPTOR(nvidia) - -#endif // __CHUNK_GATED_DELTA_RULE_NVIDIA_H__ diff --git a/src/infiniop/ops/chunk_gated_delta_rule/operator.cc b/src/infiniop/ops/chunk_gated_delta_rule/operator.cc deleted file mode 100644 index 19b5198bf..000000000 --- a/src/infiniop/ops/chunk_gated_delta_rule/operator.cc +++ /dev/null @@ -1,145 +0,0 @@ -// infiniop/ops/chunk_gated_delta_rule/operator.cc - -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/chunk_gated_delta_rule.h" - -#if defined(ENABLE_NVIDIA_API) -#include "nvidia/chunk_gated_delta_rule_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/chunk_gated_delta_rule_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/chunk_gated_delta_rule_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateChunkGatedDeltaRuleDescriptor( - infiniopHandle_t handle, - infiniopChunkGatedDeltaRuleDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t initial_state_desc, - infiniopTensorDescriptor_t final_state_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t g_desc, - infiniopTensorDescriptor_t beta_desc, - infiniopTensorDescriptor_t cu_seqlens_desc, - infiniopTensorDescriptor_t initial_state_indices_desc, - infiniopTensorDescriptor_t final_state_indices_desc, - bool use_qk_l2norm, - size_t chunk_size) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::chunk_gated_delta_rule::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast< \ - op::chunk_gated_delta_rule::NAMESPACE::Descriptor **>( \ - desc_ptr), \ - out_desc, initial_state_desc, final_state_desc, q_desc, \ - k_desc, v_desc, g_desc, beta_desc, cu_seqlens_desc, \ - initial_state_indices_desc, final_state_indices_desc, \ - use_qk_l2norm, chunk_size); - - switch (handle->device) { -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore) -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetChunkGatedDeltaRuleWorkspaceSize( - infiniopChunkGatedDeltaRuleDescriptor_t desc, size_t *size) { -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast< \ - op::chunk_gated_delta_rule::NAMESPACE::Descriptor *>(desc) \ - ->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopChunkGatedDeltaRule( - infiniopChunkGatedDeltaRuleDescriptor_t desc, - void *workspace, size_t workspace_size, - void *out, void *initial_state, void *final_state, - const void *q, const void *k, const void *v, - const void *g, const void *beta, const void *cu_seqlens, - const void *initial_state_indices, const void *final_state_indices, - void *stream) { -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast< \ - op::chunk_gated_delta_rule::NAMESPACE::Descriptor *>(desc) \ - ->calculate(workspace, workspace_size, out, initial_state, \ - final_state, q, k, v, g, beta, cu_seqlens, \ - initial_state_indices, final_state_indices, stream); - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore) -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyChunkGatedDeltaRuleDescriptor( - infiniopChunkGatedDeltaRuleDescriptor_t desc) { -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast< \ - op::chunk_gated_delta_rule::NAMESPACE::Descriptor *>(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_METAX_API - DESTROY(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - DESTROY(INFINI_DEVICE_MOORE, moore) -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DESTROY -} diff --git a/src/infiniop/ops/clip/cpu/clip_cpu.cc b/src/infiniop/ops/clip/cpu/clip_cpu.cc deleted file mode 100644 index fd651e9bf..000000000 --- a/src/infiniop/ops/clip/cpu/clip_cpu.cc +++ /dev/null @@ -1,56 +0,0 @@ -#include "clip_cpu.h" - -namespace op::clip::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &in_desc = input_desc_vec.at(0); - const auto &min_desc = input_desc_vec.at(1); - const auto &max_desc = input_desc_vec.at(2); - const auto &out_shape = out_desc->shape(); - const auto &in_shape = in_desc->shape(); - const auto &min_shape = min_desc->shape(); - const auto &max_shape = max_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - CHECK_SAME_SHAPE(out_shape, in_shape); - CHECK_SAME_SHAPE(out_shape, min_shape); - CHECK_SAME_SHAPE(out_shape, max_shape); - - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::clip::cpu diff --git a/src/infiniop/ops/clip/cpu/clip_cpu.h b/src/infiniop/ops/clip/cpu/clip_cpu.h deleted file mode 100644 index ab42fc6f4..000000000 --- a/src/infiniop/ops/clip/cpu/clip_cpu.h +++ /dev/null @@ -1,23 +0,0 @@ -#ifndef __CLIP_CPU_H__ -#define __CLIP_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" -#include "infiniop/ops/clip.h" - -ELEMENTWISE_DESCRIPTOR(clip, cpu) - -namespace op::clip::cpu { - -typedef struct ClipOp { -public: - static constexpr size_t num_inputs = 3; - - template - T operator()(const T &x, const T &min_val, const T &max_val) const { - return std::max(std::min(x, max_val), min_val); - } -} ClipOp; - -} // namespace op::clip::cpu - -#endif // __CLIP_CPU_H__ diff --git a/src/infiniop/ops/clip/cuda/kernel.cuh b/src/infiniop/ops/clip/cuda/kernel.cuh deleted file mode 100644 index bf3820ffd..000000000 --- a/src/infiniop/ops/clip/cuda/kernel.cuh +++ /dev/null @@ -1,29 +0,0 @@ -#ifndef __CLIP_CUDA_H__ -#define __CLIP_CUDA_H__ - -namespace op::clip::cuda { - -typedef struct ClipOp { -public: - static constexpr size_t num_inputs = 3; - - template - __device__ __forceinline__ T operator()(const T &x, const T &min_val, const T &max_val) const { - if constexpr (std::is_same_v || std::is_same_v) { -#ifndef ENABLE_ILUVATAR_API - return __hmax2(__hmin2(x, max_val), min_val); -#else - return {std::clamp(x.x, min_val.x, max_val.x), std::clamp(x.y, min_val.y, max_val.y)}; -#endif - } else { -#if defined(ENABLE_HYGON_API) - return x < min_val ? min_val : (max_val < x ? max_val : x); -#else - return std::clamp(x, min_val, max_val); -#endif - } - } -} ClipOp; -} // namespace op::clip::cuda - -#endif // __CLIP_CUDA_H__ diff --git a/src/infiniop/ops/clip/kunlun/clip_kunlun.h b/src/infiniop/ops/clip/kunlun/clip_kunlun.h deleted file mode 100644 index 193a40ec2..000000000 --- a/src/infiniop/ops/clip/kunlun/clip_kunlun.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CLIP_KUNLUN_API_H__ -#define __CLIP_KUNLUN_API_H__ - -#include "../../../elementwise/kunlun/elementwise_kunlun_api.h" - -ELEMENTWISE_DESCRIPTOR(clip, kunlun) - -#endif // __CLIP_KUNLUN_API_H__ diff --git a/src/infiniop/ops/clip/kunlun/clip_kunlun.xpu b/src/infiniop/ops/clip/kunlun/clip_kunlun.xpu deleted file mode 100644 index f52c36eca..000000000 --- a/src/infiniop/ops/clip/kunlun/clip_kunlun.xpu +++ /dev/null @@ -1,60 +0,0 @@ -#include "../../../elementwise/kunlun/elementwise_kunlun.h" -#include "clip_kunlun.h" -#include "kernel.h" - -namespace op::clip::kunlun { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &in_desc = input_desc_vec.at(0); - const auto &min_desc = input_desc_vec.at(1); - const auto &max_desc = input_desc_vec.at(2); - const auto &out_shape = out_desc->shape(); - const auto &in_shape = in_desc->shape(); - const auto &min_shape = min_desc->shape(); - const auto &max_shape = max_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - CHECK_SAME_SHAPE(out_shape, in_shape); - CHECK_SAME_SHAPE(out_shape, min_shape); - CHECK_SAME_SHAPE(out_shape, max_shape); - - CREATE_ELEMENTWISE_KUNLUN_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<8, ClipOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<8, ClipOp, bfloat16_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<8, ClipOp, float>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::clip::kunlun diff --git a/src/infiniop/ops/clip/kunlun/kernel.h b/src/infiniop/ops/clip/kunlun/kernel.h deleted file mode 100644 index 9a0d56857..000000000 --- a/src/infiniop/ops/clip/kunlun/kernel.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef __CLIP_KUNLUN_KERNEL_H__ -#define __CLIP_KUNLUN_KERNEL_H__ -#include - -namespace op::clip::kunlun { - -typedef struct ClipOp { -public: - static constexpr int num_inputs = 3; - template - inline __device__ T operator()(const T *inputs) const { - T x = inputs[0]; - T min_val = inputs[1]; - T max_val = inputs[2]; - return fmax(fmin(x, max_val), min_val); - } - - // bfloat16 特化版本（使用 float 计算精度） - inline __device__ bfloat16_t operator()(const bfloat16_t *inputs) const { - float x_f = __bfloat162float(inputs[0]); - float min_val_f = __bfloat162float(inputs[1]); - float max_val_f = __bfloat162float(inputs[2]); - float result_f = fmax(fmin(x_f, max_val_f), min_val_f); - return __float2bfloat16(result_f); - } -} ClipOp; - -} // namespace op::clip::kunlun - -#endif // __CLIP_KUNLUN_KERNEL_H__ diff --git a/src/infiniop/ops/clip/metax/clip_metax.h b/src/infiniop/ops/clip/metax/clip_metax.h deleted file mode 100644 index 1425cf3b8..000000000 --- a/src/infiniop/ops/clip/metax/clip_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CLIP_METAX_API_H__ -#define __CLIP_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(clip, metax) - -#endif // __CLIP_METAX_API_H__ diff --git a/src/infiniop/ops/clip/metax/clip_metax.maca b/src/infiniop/ops/clip/metax/clip_metax.maca deleted file mode 100644 index 7382d044c..000000000 --- a/src/infiniop/ops/clip/metax/clip_metax.maca +++ /dev/null @@ -1,63 +0,0 @@ -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" -#include "clip_metax.h" - -namespace op::clip::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &in_desc = input_desc_vec.at(0); - const auto &min_desc = input_desc_vec.at(1); - const auto &max_desc = input_desc_vec.at(2); - const auto &out_shape = out_desc->shape(); - const auto &in_shape = in_desc->shape(); - const auto &min_shape = min_desc->shape(); - const auto &max_shape = max_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - CHECK_SAME_SHAPE(out_shape, in_shape); - CHECK_SAME_SHAPE(out_shape, min_shape); - CHECK_SAME_SHAPE(out_shape, max_shape); - - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::ClipOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::ClipOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::ClipOp, double>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::ClipOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::clip::metax diff --git a/src/infiniop/ops/clip/nvidia/clip_nvidia.cu b/src/infiniop/ops/clip/nvidia/clip_nvidia.cu deleted file mode 100644 index ce5e8e95a..000000000 --- a/src/infiniop/ops/clip/nvidia/clip_nvidia.cu +++ /dev/null @@ -1,63 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "clip_nvidia.cuh" - -namespace op::clip::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &in_desc = input_desc_vec.at(0); - const auto &min_desc = input_desc_vec.at(1); - const auto &max_desc = input_desc_vec.at(2); - const auto &out_shape = out_desc->shape(); - const auto &in_shape = in_desc->shape(); - const auto &min_shape = min_desc->shape(); - const auto &max_shape = max_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - CHECK_SAME_SHAPE(out_shape, in_shape); - CHECK_SAME_SHAPE(out_shape, min_shape); - CHECK_SAME_SHAPE(out_shape, max_shape); - - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::ClipOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::ClipOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::ClipOp, double>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::ClipOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::clip::nvidia diff --git a/src/infiniop/ops/clip/nvidia/clip_nvidia.cuh b/src/infiniop/ops/clip/nvidia/clip_nvidia.cuh deleted file mode 100644 index 0b6c21d96..000000000 --- a/src/infiniop/ops/clip/nvidia/clip_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CLIP_CUDA_API_H__ -#define __CLIP_CUDA_API_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(clip, nvidia) - -#endif // __CLIP_CUDA_API_H__ diff --git a/src/infiniop/ops/clip/operator.cc b/src/infiniop/ops/clip/operator.cc deleted file mode 100644 index c933cafab..000000000 --- a/src/infiniop/ops/clip/operator.cc +++ /dev/null @@ -1,196 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/clip.h" - -#ifdef ENABLE_CPU_API -#include "cpu/clip_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) -#include "nvidia/clip_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/clip_metax.h" -#endif -#ifdef ENABLE_KUNLUN_API -#include "kunlun/clip_kunlun.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateClipDescriptor( - infiniopHandle_t handle, - infiniopClipDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t min_val, - infiniopTensorDescriptor_t max_val) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::clip::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y, \ - {x, min_val, max_val}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_KUNLUN_API - CREATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetClipWorkspaceSize(infiniopClipDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - GET(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_KUNLUN_API - GET(INFINI_DEVICE_KUNLUN, kunlun) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET -} - -__INFINI_C infiniStatus_t infiniopClip( - infiniopClipDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *min_val, - const void *max_val, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, {x, min_val, max_val}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_KUNLUN_API - CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyClipDescriptor(infiniopClipDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - DELETE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_KUNLUN_API - DELETE(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/conv/bang/conv_bang.cc b/src/infiniop/ops/conv/bang/conv_bang.cc deleted file mode 100644 index 16f5a665d..000000000 --- a/src/infiniop/ops/conv/bang/conv_bang.cc +++ /dev/null @@ -1,345 +0,0 @@ -#include "conv_bang.h" -#include "../../../devices/bang/common_bang.h" - -#include -#include -#include - -namespace op::conv::bang { - -static size_t alignSize(size_t value) { - return (value + ALIGN_SIZE - 1) / ALIGN_SIZE * ALIGN_SIZE; -} - -static size_t tensorElements(const std::vector &dims) { - size_t n = 1; - for (int d : dims) { - n *= static_cast(d); - } - return n; -} - -static infiniStatus_t setTensor(cnnlTensorDescriptor_t desc, - cnnlTensorLayout_t layout, - infiniDtype_t dtype, - const std::vector &dims) { - CHECK_BANG(cnnlSetTensorDescriptor( - desc, - layout, - device::bang::getCnnlDtype(dtype), - static_cast(dims.size()), - const_cast(dims.data()))); - return INFINI_STATUS_SUCCESS; -} - -struct ConvShapes { - std::vector x_orig; - std::vector x_cnnl; - std::vector w_orig; - std::vector w_cnnl; - std::vector y_orig; - std::vector y_cnnl; - std::vector pad; - std::vector stride; - std::vector dilation; - std::vector perm_to_cnnl; - std::vector perm_to_orig; - cnnlTensorLayout_t orig_layout; - cnnlTensorLayout_t cnnl_layout; - int conv_dim; -}; - -static ConvShapes makeShapes(const ConvInfo &info) { - ConvShapes s; - if (info.ndim() == 1) { - s.x_orig = {static_cast(info.batch()), static_cast(info.in_channels()), 1, static_cast(info.input_dim(0))}; - s.x_cnnl = {static_cast(info.batch()), 1, static_cast(info.input_dim(0)), static_cast(info.in_channels())}; - s.w_orig = {static_cast(info.out_channels()), static_cast(info.in_channels()), 1, static_cast(info.kernel_dim(0))}; - s.w_cnnl = {static_cast(info.out_channels()), 1, static_cast(info.kernel_dim(0)), static_cast(info.in_channels())}; - s.y_orig = {static_cast(info.batch()), static_cast(info.out_channels()), 1, static_cast(info.output_dim(0))}; - s.y_cnnl = {static_cast(info.batch()), 1, static_cast(info.output_dim(0)), static_cast(info.out_channels())}; - s.pad = {0, 0, static_cast(info.pad_info(0)), static_cast(info.pad_info(0))}; - s.stride = {1, static_cast(info.stride_info(0))}; - s.dilation = {1, static_cast(info.dilation_info(0))}; - s.perm_to_cnnl = {0, 2, 3, 1}; - s.perm_to_orig = {0, 3, 1, 2}; - s.orig_layout = CNNL_LAYOUT_NCHW; - s.cnnl_layout = CNNL_LAYOUT_NHWC; - s.conv_dim = 4; - } else if (info.ndim() == 2) { - s.x_orig = {static_cast(info.batch()), static_cast(info.in_channels()), static_cast(info.input_dim(0)), static_cast(info.input_dim(1))}; - s.x_cnnl = {static_cast(info.batch()), static_cast(info.input_dim(0)), static_cast(info.input_dim(1)), static_cast(info.in_channels())}; - s.w_orig = {static_cast(info.out_channels()), static_cast(info.in_channels()), static_cast(info.kernel_dim(0)), static_cast(info.kernel_dim(1))}; - s.w_cnnl = {static_cast(info.out_channels()), static_cast(info.kernel_dim(0)), static_cast(info.kernel_dim(1)), static_cast(info.in_channels())}; - s.y_orig = {static_cast(info.batch()), static_cast(info.out_channels()), static_cast(info.output_dim(0)), static_cast(info.output_dim(1))}; - s.y_cnnl = {static_cast(info.batch()), static_cast(info.output_dim(0)), static_cast(info.output_dim(1)), static_cast(info.out_channels())}; - s.pad = {static_cast(info.pad_info(0)), static_cast(info.pad_info(0)), static_cast(info.pad_info(1)), static_cast(info.pad_info(1))}; - s.stride = {static_cast(info.stride_info(0)), static_cast(info.stride_info(1))}; - s.dilation = {static_cast(info.dilation_info(0)), static_cast(info.dilation_info(1))}; - s.perm_to_cnnl = {0, 2, 3, 1}; - s.perm_to_orig = {0, 3, 1, 2}; - s.orig_layout = CNNL_LAYOUT_NCHW; - s.cnnl_layout = CNNL_LAYOUT_NHWC; - s.conv_dim = 4; - } else { - s.x_orig = {static_cast(info.batch()), static_cast(info.in_channels()), static_cast(info.input_dim(0)), static_cast(info.input_dim(1)), static_cast(info.input_dim(2))}; - s.x_cnnl = {static_cast(info.batch()), static_cast(info.input_dim(0)), static_cast(info.input_dim(1)), static_cast(info.input_dim(2)), static_cast(info.in_channels())}; - s.w_orig = {static_cast(info.out_channels()), static_cast(info.in_channels()), static_cast(info.kernel_dim(0)), static_cast(info.kernel_dim(1)), static_cast(info.kernel_dim(2))}; - s.w_cnnl = {static_cast(info.out_channels()), static_cast(info.kernel_dim(0)), static_cast(info.kernel_dim(1)), static_cast(info.kernel_dim(2)), static_cast(info.in_channels())}; - s.y_orig = {static_cast(info.batch()), static_cast(info.out_channels()), static_cast(info.output_dim(0)), static_cast(info.output_dim(1)), static_cast(info.output_dim(2))}; - s.y_cnnl = {static_cast(info.batch()), static_cast(info.output_dim(0)), static_cast(info.output_dim(1)), static_cast(info.output_dim(2)), static_cast(info.out_channels())}; - s.pad = {static_cast(info.pad_info(0)), static_cast(info.pad_info(0)), static_cast(info.pad_info(1)), static_cast(info.pad_info(1)), static_cast(info.pad_info(2)), static_cast(info.pad_info(2))}; - s.stride = {static_cast(info.stride_info(0)), static_cast(info.stride_info(1)), static_cast(info.stride_info(2))}; - s.dilation = {static_cast(info.dilation_info(0)), static_cast(info.dilation_info(1)), static_cast(info.dilation_info(2))}; - s.perm_to_cnnl = {0, 2, 3, 4, 1}; - s.perm_to_orig = {0, 4, 1, 2, 3}; - s.orig_layout = CNNL_LAYOUT_NCDHW; - s.cnnl_layout = CNNL_LAYOUT_NDHWC; - s.conv_dim = 5; - } - return s; -} - -struct Descriptor::Opaque { - std::shared_ptr internal; - cnnlTensorDescriptor_t x_orig_desc = nullptr; - cnnlTensorDescriptor_t x_cnnl_desc = nullptr; - cnnlTensorDescriptor_t w_orig_desc = nullptr; - cnnlTensorDescriptor_t w_cnnl_desc = nullptr; - cnnlTensorDescriptor_t y_orig_desc = nullptr; - cnnlTensorDescriptor_t y_cnnl_desc = nullptr; - cnnlTensorDescriptor_t bias_desc = nullptr; - cnnlConvolutionDescriptor_t conv_desc = nullptr; - cnnlQuantizeExDescriptor_t y_quant_desc = nullptr; - cnnlTransposeDescriptor_t x_trans_desc = nullptr; - cnnlTransposeDescriptor_t w_trans_desc = nullptr; - cnnlTransposeDescriptor_t y_trans_desc = nullptr; - cnnlConvolutionForwardAlgo_t algo = CNNL_CONVOLUTION_FWD_ALGO_DIRECT; - size_t x_cnnl_bytes = 0; - size_t w_cnnl_bytes = 0; - size_t y_cnnl_bytes = 0; - size_t transpose_workspace_size = 0; - size_t conv_workspace_size = 0; - - ~Opaque() { - if (x_orig_desc) { - cnnlDestroyTensorDescriptor(x_orig_desc); - } - if (x_cnnl_desc) { - cnnlDestroyTensorDescriptor(x_cnnl_desc); - } - if (w_orig_desc) { - cnnlDestroyTensorDescriptor(w_orig_desc); - } - if (w_cnnl_desc) { - cnnlDestroyTensorDescriptor(w_cnnl_desc); - } - if (y_orig_desc) { - cnnlDestroyTensorDescriptor(y_orig_desc); - } - if (y_cnnl_desc) { - cnnlDestroyTensorDescriptor(y_cnnl_desc); - } - if (bias_desc) { - cnnlDestroyTensorDescriptor(bias_desc); - } - if (conv_desc) { - cnnlDestroyConvolutionDescriptor(conv_desc); - } - if (y_quant_desc) { - cnnlDestroyQuantizeExDescriptor(y_quant_desc); - } - if (x_trans_desc) { - cnnlDestroyTransposeDescriptor(x_trans_desc); - } - if (w_trans_desc) { - cnnlDestroyTransposeDescriptor(w_trans_desc); - } - if (y_trans_desc) { - cnnlDestroyTransposeDescriptor(y_trans_desc); - } - } -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -static infiniStatus_t createTransposeDesc(cnnlTransposeDescriptor_t *desc, const std::vector &perm) { - CHECK_BANG(cnnlCreateTransposeDescriptor(desc)); - CHECK_BANG(cnnlSetTransposeDescriptor(*desc, static_cast(perm.size()), const_cast(perm.data()))); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t w_desc, - infiniopTensorDescriptor_t b_desc, - const void *pads, - const void *strides, - const void *dilations, - size_t n) { - auto handle = reinterpret_cast(handle_); - auto dtype = y_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = ConvInfo::create(handle_, y_desc, x_desc, w_desc, b_desc, pads, strides, dilations, n); - CHECK_RESULT(result); - auto info = result.take(); - if (info.ndim() < 1 || info.ndim() > 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - auto shapes = makeShapes(info); - auto opaque = new Opaque(); - opaque->internal = handle->internal(); - opaque->x_cnnl_bytes = alignSize(tensorElements(shapes.x_cnnl) * infiniSizeOf(dtype)); - opaque->w_cnnl_bytes = alignSize(tensorElements(shapes.w_cnnl) * infiniSizeOf(dtype)); - opaque->y_cnnl_bytes = alignSize(tensorElements(shapes.y_cnnl) * infiniSizeOf(dtype)); - - CHECK_BANG(cnnlCreateTensorDescriptor(&opaque->x_orig_desc)); - CHECK_BANG(cnnlCreateTensorDescriptor(&opaque->x_cnnl_desc)); - CHECK_BANG(cnnlCreateTensorDescriptor(&opaque->w_orig_desc)); - CHECK_BANG(cnnlCreateTensorDescriptor(&opaque->w_cnnl_desc)); - CHECK_BANG(cnnlCreateTensorDescriptor(&opaque->y_orig_desc)); - CHECK_BANG(cnnlCreateTensorDescriptor(&opaque->y_cnnl_desc)); - CHECK_STATUS(setTensor(opaque->x_orig_desc, shapes.orig_layout, dtype, shapes.x_orig)); - CHECK_STATUS(setTensor(opaque->x_cnnl_desc, shapes.cnnl_layout, dtype, shapes.x_cnnl)); - CHECK_STATUS(setTensor(opaque->w_orig_desc, shapes.orig_layout, dtype, shapes.w_orig)); - CHECK_STATUS(setTensor(opaque->w_cnnl_desc, shapes.cnnl_layout, dtype, shapes.w_cnnl)); - CHECK_STATUS(setTensor(opaque->y_orig_desc, shapes.orig_layout, dtype, shapes.y_orig)); - CHECK_STATUS(setTensor(opaque->y_cnnl_desc, shapes.cnnl_layout, dtype, shapes.y_cnnl)); - - if (b_desc != nullptr) { - int bias_dims[1] = {static_cast(info.out_channels())}; - CHECK_BANG(cnnlCreateTensorDescriptor(&opaque->bias_desc)); - CHECK_BANG(cnnlSetTensorDescriptor(opaque->bias_desc, CNNL_LAYOUT_ARRAY, device::bang::getCnnlDtype(dtype), 1, bias_dims)); - } - - CHECK_BANG(cnnlCreateConvolutionDescriptor(&opaque->conv_desc)); - cnnlDataType_t compute_dtype = dtype == INFINI_DTYPE_BF16 - ? CNNL_DTYPE_FLOAT - : device::bang::getCnnlDtype(dtype); - CHECK_BANG(cnnlSetConvolutionDescriptor( - opaque->conv_desc, - shapes.conv_dim, - shapes.pad.data(), - shapes.stride.data(), - shapes.dilation.data(), - 1, - compute_dtype)); - if (dtype == INFINI_DTYPE_BF16) { - CHECK_BANG(cnnlCreateQuantizeExDescriptor(&opaque->y_quant_desc)); - CHECK_BANG(cnnlSetQuantizeExDescriptorQuantSchemeAndDtype( - opaque->y_quant_desc, CNNL_QUANTIZE_NONE, CNNL_DTYPE_FLOAT)); - CHECK_BANG(cnnlSetConvolutionDescriptorQuant( - opaque->conv_desc, nullptr, nullptr, opaque->y_quant_desc)); - } - opaque->algo = CNNL_CONVOLUTION_FWD_ALGO_DIRECT; - - CHECK_STATUS(createTransposeDesc(&opaque->x_trans_desc, shapes.perm_to_cnnl)); - CHECK_STATUS(createTransposeDesc(&opaque->w_trans_desc, shapes.perm_to_cnnl)); - CHECK_STATUS(createTransposeDesc(&opaque->y_trans_desc, shapes.perm_to_orig)); - - size_t tx_ws = 0, tw_ws = 0, ty_ws = 0, conv_ws = 0; - CHECK_STATUS(opaque->internal->useCnnl( - nullptr, - [&](cnnlHandle_t cnnl_handle) { - CHECK_BANG(cnnlGetTransposeWorkspaceSize(cnnl_handle, opaque->x_orig_desc, opaque->x_trans_desc, &tx_ws)); - CHECK_BANG(cnnlGetTransposeWorkspaceSize(cnnl_handle, opaque->w_orig_desc, opaque->w_trans_desc, &tw_ws)); - CHECK_BANG(cnnlGetTransposeWorkspaceSize(cnnl_handle, opaque->y_cnnl_desc, opaque->y_trans_desc, &ty_ws)); - CHECK_BANG(cnnlGetConvolutionForwardWorkspaceSize( - cnnl_handle, - opaque->x_cnnl_desc, - opaque->w_cnnl_desc, - opaque->y_cnnl_desc, - opaque->bias_desc, - opaque->conv_desc, - opaque->algo, - &conv_ws)); - return INFINI_STATUS_SUCCESS; - })); - opaque->transpose_workspace_size = alignSize(std::max({tx_ws, tw_ws, ty_ws})); - opaque->conv_workspace_size = alignSize(conv_ws); - size_t workspace_size = opaque->x_cnnl_bytes + opaque->w_cnnl_bytes + opaque->y_cnnl_bytes + opaque->transpose_workspace_size + opaque->conv_workspace_size; - - *desc_ptr = new Descriptor(dtype, std::move(info), workspace_size, opaque, handle_->device, handle_->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *w, - const void *bias, - void *stream) const { - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - auto queue = reinterpret_cast(stream); - char *buffer = reinterpret_cast(workspace); - void *x_cnnl = buffer; - void *w_cnnl = buffer + _opaque->x_cnnl_bytes; - void *y_cnnl = buffer + _opaque->x_cnnl_bytes + _opaque->w_cnnl_bytes; - void *transpose_workspace = buffer + _opaque->x_cnnl_bytes + _opaque->w_cnnl_bytes + _opaque->y_cnnl_bytes; - void *conv_workspace = buffer + _opaque->x_cnnl_bytes + _opaque->w_cnnl_bytes + _opaque->y_cnnl_bytes + _opaque->transpose_workspace_size; - - CHECK_STATUS(_opaque->internal->useCnnl( - queue, - [&](cnnlHandle_t handle) { - CHECK_BANG(cnnlTranspose_v2( - handle, - _opaque->x_trans_desc, - _opaque->x_orig_desc, - x, - _opaque->x_cnnl_desc, - x_cnnl, - transpose_workspace, - _opaque->transpose_workspace_size)); - CHECK_BANG(cnnlTranspose_v2( - handle, - _opaque->w_trans_desc, - _opaque->w_orig_desc, - w, - _opaque->w_cnnl_desc, - w_cnnl, - transpose_workspace, - _opaque->transpose_workspace_size)); - CHECK_BANG(cnnlConvolutionForward( - handle, - _opaque->conv_desc, - _opaque->algo, - nullptr, - _opaque->x_cnnl_desc, - x_cnnl, - _opaque->w_cnnl_desc, - w_cnnl, - _opaque->bias_desc, - bias, - conv_workspace, - _opaque->conv_workspace_size, - nullptr, - _opaque->y_cnnl_desc, - y_cnnl)); - CHECK_BANG(cnnlTranspose_v2( - handle, - _opaque->y_trans_desc, - _opaque->y_cnnl_desc, - y_cnnl, - _opaque->y_orig_desc, - y, - transpose_workspace, - _opaque->transpose_workspace_size)); - return INFINI_STATUS_SUCCESS; - })); - cnrtQueueSync(queue); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::conv::bang diff --git a/src/infiniop/ops/conv/bang/conv_bang.h b/src/infiniop/ops/conv/bang/conv_bang.h deleted file mode 100644 index 6f1364ab8..000000000 --- a/src/infiniop/ops/conv/bang/conv_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CONV_BANG_H__ -#define __CONV_BANG_H__ - -#include "../conv.h" - -DESCRIPTOR(bang) - -#endif // __CONV_BANG_H__ diff --git a/src/infiniop/ops/conv/conv.h b/src/infiniop/ops/conv/conv.h deleted file mode 100644 index b1a73e491..000000000 --- a/src/infiniop/ops/conv/conv.h +++ /dev/null @@ -1,56 +0,0 @@ -#ifndef __CONV_H__ -#define __CONV_H__ - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::conv::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - infiniDtype_t _dtype; \ - ConvInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - infiniDtype_t dtype, \ - ConvInfo info, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _dtype(dtype), \ - _info(info), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t y, \ - infiniopTensorDescriptor_t x, \ - infiniopTensorDescriptor_t w, \ - infiniopTensorDescriptor_t b, \ - const void *pads, \ - const void *strides, \ - const void *dilations, \ - size_t n); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *y, \ - const void *x, \ - const void *w, \ - const void *bias, \ - void *stream) const; \ - }; \ - } -#endif // __CONV_H__ diff --git a/src/infiniop/ops/conv/cpu/conv_cpu.cc b/src/infiniop/ops/conv/cpu/conv_cpu.cc deleted file mode 100644 index 06bfbca3c..000000000 --- a/src/infiniop/ops/conv/cpu/conv_cpu.cc +++ /dev/null @@ -1,363 +0,0 @@ -#include "conv_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include - -namespace op::conv::cpu { - -inline size_t calculatePaddedInputSize(const ConvInfo &info) { - std::vector shape(info.ndim() + 2); - shape[0] = info.batch(); - shape[1] = info.in_channels(); - for (size_t i = 0; i < info.ndim(); ++i) { - shape[i + 2] = info.input_dim(i); - } - - return op::common_cpu::getPaddedSize(info.ndim() + 2, shape.data(), info.getPadsInfo()); -} - -inline size_t calculateOutputSize(const ConvInfo &info) { - size_t size = info.batch() * info.out_channels(); - for (size_t i = 0; i < info.ndim(); ++i) { - size *= info.output_dim(i); - } - return size; -} - -inline bool needsPadding(const ConvInfo &info) { - const size_t *pads = info.getPadsInfo(); - for (size_t i = 0; i < info.ndim(); ++i) { - if (pads[i] > 0) { - return true; - } - } - return false; -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t w_desc, - infiniopTensorDescriptor_t b_desc, - const void *pads, - const void *strides, - const void *dilations, - size_t n) { - auto handle = reinterpret_cast(handle_); - auto dtype = y_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = ConvInfo::create(handle_, y_desc, x_desc, w_desc, b_desc, - pads, strides, dilations, n); - CHECK_RESULT(result); - - size_t WorkSpaceSize = 0; - const ConvInfo &info = result.take(); - - if (needsPadding(info)) { - WorkSpaceSize += calculatePaddedInputSize(info) * infiniSizeOf(dtype); - } - - if (dtype == INFINI_DTYPE_F16 || dtype == INFINI_DTYPE_BF16) { - WorkSpaceSize += calculateOutputSize(info) * sizeof(float); - } - - *desc_ptr = new Descriptor( - dtype, std::move(info), WorkSpaceSize, - nullptr, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -void fillPaddedInput( - const ConvInfo &info, - const Tdata *x, - const size_t *padded_x_shape, - Tdata *padded_x, - size_t x_index, - size_t padded_x_index, - size_t ndim) { - size_t x_shape_val; - if (ndim == 0) { - x_shape_val = info.batch(); - } else if (ndim == 1) { - x_shape_val = info.in_channels(); - } else { - x_shape_val = info.input_dim(ndim - 2); - } - - const auto padded_x_shape_val = padded_x_shape[ndim]; - const auto x_base_index = x_index * x_shape_val; - - size_t pad_offset = 0; - if (ndim >= 2 && x_shape_val != padded_x_shape_val) { - pad_offset = info.pad_info(ndim - 2); - } - - const auto padded_x_base_index = padded_x_index * padded_x_shape_val + pad_offset; - - for (size_t i = 0; i < x_shape_val; ++i) { - if (ndim == info.ndim() + 2 - 1) { - padded_x[padded_x_base_index + i] = x[x_base_index + i]; - } else { - fillPaddedInput(info, x, padded_x_shape, padded_x, - x_base_index + i, padded_x_base_index + i, ndim + 1); - } - } -} - -template -void _applyConv( - const ConvInfo &info, - Ydata *y, - const Xdata *x, - const Xdata *w, - const size_t *x_shape, - size_t x_index, - size_t w_index, - size_t y_index, - size_t ndim) { - - size_t dim_size, kernel_size; - size_t dilation, stride; - - if (ndim < 2) { - return; - } else { - dim_size = x_shape[ndim]; - kernel_size = info.kernel_dim(ndim - 2); - dilation = info.dilation_info(ndim - 2); - stride = info.stride_info(ndim - 2); - } - if (stride == 0) { - std::cerr << "Error: stride cannot be zero." << std::endl; - } - const auto steps = (dim_size - dilation * (kernel_size - 1) - 1) / stride + 1; - x_index *= dim_size; - w_index *= kernel_size; - size_t y_stride; - if (ndim == 0) { - y_stride = info.out_channels(); - } else if (ndim == 1) { - y_stride = 1; - } else { - y_stride = info.output_dim(ndim - 2); - } - y_index *= y_stride; - for (size_t i = 0; i < steps; ++i, ++y_index) { - for (size_t k = 0; k < kernel_size; ++k) { - const auto curr_x_index = x_index + i * stride + k * dilation; - const auto curr_w_index = w_index + k; - if (ndim == info.ndim() + 1) { - if constexpr (std::is_same::value || std::is_same::value) { - y[y_index] += utils::cast(x[curr_x_index]) * utils::cast(w[curr_w_index]); - } else { - y[y_index] += x[curr_x_index] * w[curr_w_index]; - } - } else { - _applyConv(info, y, x, w, x_shape, curr_x_index, curr_w_index, - y_index, ndim + 1); - } - } - } -} - -template -void applyConv( - const ConvInfo &info, - Ydata *y, - const Xdata *x, - const Xdata *w, - const size_t *x_shape) { - const ptrdiff_t batch_size = static_cast(info.batch()); - const ptrdiff_t out_channels = static_cast(info.out_channels()); - const ptrdiff_t total_iterations = batch_size * out_channels; - -#pragma omp parallel for schedule(dynamic) - for (ptrdiff_t iter = 0; iter < total_iterations; ++iter) { - const ptrdiff_t i = iter / out_channels; // batch index - const ptrdiff_t j = iter % out_channels; // output channel index - - const size_t y_index = static_cast(i) * info.out_channels() + static_cast(j); - - // 内层循环：遍历输入通道 - for (size_t k = 0; k < info.in_channels(); ++k) { - const size_t x_index = static_cast(i) * info.in_channels() + k; - const size_t w_index = static_cast(j) * info.in_channels() + k; - _applyConv(info, y, x, w, x_shape, x_index, w_index, y_index, 2); - } - } -} - -template -void _conv_cpu( - const ConvInfo &info, - void *workspace, - size_t workspace_size, - Ydata *y, - const Xdata *x, - const Xdata *w) { - if (needsPadding(info)) { - auto padded_x = reinterpret_cast(workspace); - if constexpr (std::is_same::value) { - fp16_t zero_val = utils::cast(0.0f); - std::fill(padded_x, padded_x + calculatePaddedInputSize(info), zero_val); - } else if constexpr (std::is_same::value) { - bf16_t zero_val = utils::cast(0.0f); - std::fill(padded_x, padded_x + calculatePaddedInputSize(info), zero_val); - } else if constexpr (std::is_same::value) { - std::fill(padded_x, padded_x + calculatePaddedInputSize(info), 0.0f); - } - fillPaddedInput(info, x, info.getPaddedShape(), padded_x, 0, 0, 0); - - applyConv(info, y, padded_x, w, info.getPaddedShape()); - } else { - std::vector shape(info.ndim() + 2); - shape[0] = info.batch(); - shape[1] = info.in_channels(); - for (size_t i = 0; i < info.ndim(); ++i) { - shape[i + 2] = info.input_dim(i); - } - applyConv(info, y, x, w, shape.data()); - } -} - -template -infiniStatus_t conv_cpu( - const ConvInfo &info, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *w, - const void *bias) { - auto y_ptr = reinterpret_cast(y); - auto x_ptr = reinterpret_cast(x); - auto w_ptr = reinterpret_cast(w); - auto output_size = calculateOutputSize(info); - if constexpr (std::is_same::value) { - std::fill(y_ptr, y_ptr + output_size, 0.0f); - } else { - std::fill(y_ptr, y_ptr + output_size, static_cast(0)); - } - _conv_cpu(info, workspace, workspace_size, y_ptr, x_ptr, w_ptr); - if (bias != nullptr) { - auto bias_ptr = reinterpret_cast(bias); -#pragma omp parallel for - for (ptrdiff_t i = 0; i < static_cast(output_size); ++i) { - size_t channel_idx = (i / info.spatial_sizes()) % info.out_channels(); - y_ptr[i] += bias_ptr[channel_idx]; - } - } - return INFINI_STATUS_SUCCESS; -} - -template <> -infiniStatus_t conv_cpu( - const ConvInfo &info, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *w, - const void *bias) { - auto y_float = reinterpret_cast(workspace); - auto x_half = reinterpret_cast(x); - auto w_half = reinterpret_cast(w); - auto output_size = calculateOutputSize(info); - std::fill(y_float, y_float + output_size, 0.0f); - - void *conv_workspace = y_float + output_size; - size_t conv_workspace_size = workspace_size - output_size * sizeof(float); - - _conv_cpu(info, conv_workspace, conv_workspace_size, y_float, x_half, w_half); - - auto y_half = reinterpret_cast(y); - if (bias != nullptr) { - auto bias_half = reinterpret_cast(bias); -#pragma omp parallel for - for (ptrdiff_t i = 0; i < static_cast(output_size); ++i) { - size_t channel_idx = (i / info.spatial_sizes()) % info.out_channels(); - float bias_value = utils::cast(bias_half[channel_idx]); - y_float[i] += bias_value; - y_half[i] = utils::cast(y_float[i]); - } - } else { -#pragma omp parallel for - for (ptrdiff_t i = 0; i < static_cast(output_size); ++i) { - y_half[i] = utils::cast(y_float[i]); - } - } - - return INFINI_STATUS_SUCCESS; -} - -template <> -infiniStatus_t conv_cpu( - const ConvInfo &info, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *w, - const void *bias) { - auto y_float = reinterpret_cast(workspace); - auto x_half = reinterpret_cast(x); - auto w_half = reinterpret_cast(w); - auto output_size = calculateOutputSize(info); - std::fill(y_float, y_float + output_size, 0.0f); - - void *conv_workspace = y_float + output_size; - size_t conv_workspace_size = workspace_size - output_size * sizeof(float); - - _conv_cpu(info, conv_workspace, conv_workspace_size, y_float, x_half, w_half); - - auto y_half = reinterpret_cast(y); - if (bias != nullptr) { - auto bias_half = reinterpret_cast(bias); -#pragma omp parallel for - for (ptrdiff_t i = 0; i < static_cast(output_size); ++i) { - size_t channel_idx = (i / info.spatial_sizes()) % info.out_channels(); - float bias_value = utils::cast(bias_half[channel_idx]); - y_float[i] += bias_value; - y_half[i] = utils::cast(y_float[i]); - } - } else { -#pragma omp parallel for - for (ptrdiff_t i = 0; i < static_cast(output_size); ++i) { - y_half[i] = utils::cast(y_float[i]); - } - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *w, - const void *bias, - void *stream) const { - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - return conv_cpu(_info, workspace, workspace_size, y, x, w, bias); - case INFINI_DTYPE_F32: - return conv_cpu(_info, workspace, workspace_size, y, x, w, bias); - case INFINI_DTYPE_BF16: - return conv_cpu(_info, workspace, workspace_size, y, x, w, bias); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::conv::cpu diff --git a/src/infiniop/ops/conv/cpu/conv_cpu.h b/src/infiniop/ops/conv/cpu/conv_cpu.h deleted file mode 100644 index 22c0ad339..000000000 --- a/src/infiniop/ops/conv/cpu/conv_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CONV_CPU_H__ -#define __CONV_CPU_H__ - -#include "../conv.h" - -DESCRIPTOR(cpu) - -#endif // __CONV_CPU_H__ diff --git a/src/infiniop/ops/conv/info.h b/src/infiniop/ops/conv/info.h deleted file mode 100644 index 26c53a7b4..000000000 --- a/src/infiniop/ops/conv/info.h +++ /dev/null @@ -1,257 +0,0 @@ -#ifndef __CONV_INFO_H__ -#define __CONV_INFO_H__ - -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" - -#ifdef ENABLE_CUDA_API -#include "../../devices/nvidia/nvidia_handle.cuh" -#endif - -namespace op::conv { -class ConvInfo; -} // namespace op::conv - -namespace op::conv { - -class ConvInfo { -private: - std::vector _meta; - size_t _ndim; - size_t _batch; - size_t _in_channels; - size_t _out_channels; - size_t _spatial_sizes; - size_t _bias_dims_size; - size_t _padded_shape_size; - - ConvInfo(std::vector meta, - size_t ndim, - size_t batch, - size_t in_channels, - size_t out_channels, - size_t spatial_sizes, - size_t bias_dims_size, - size_t padded_shape_size) - : _meta(std::move(meta)), - _ndim(ndim), - _batch(batch), - _in_channels(in_channels), - _out_channels(out_channels), - _spatial_sizes(spatial_sizes), - _bias_dims_size(bias_dims_size), - _padded_shape_size(padded_shape_size) {} - -public: - inline size_t ndim() const { return _ndim; } - inline size_t batch() const { return _batch; } - inline size_t in_channels() const { return _in_channels; } - inline size_t out_channels() const { return _out_channels; } - inline size_t spatial_sizes() const { return _spatial_sizes; } - inline size_t bias_dims_size() const { return _bias_dims_size; } - inline size_t padded_shape_size() const { return _padded_shape_size; } - - inline size_t getMetaMemSize() const { - return _meta.size() * sizeof(size_t); - } - inline const int8_t *getMetaStart() const { - return reinterpret_cast(_meta.data()); - } - - inline const size_t *getInputDims() const { - return _meta.data(); - } - inline const size_t *getKernelDims() const { - return getInputDims() + _ndim; - } - inline const size_t *getOutputDims() const { - return getKernelDims() + _ndim; - } - inline const size_t *getBiasDims() const { - return getOutputDims() + _ndim; - } - inline const size_t *getPadsInfo() const { - return getBiasDims() + _bias_dims_size; - } - inline const ptrdiff_t *getStridesInfo() const { - return reinterpret_cast(getPadsInfo()) + _ndim; - } - inline const size_t *getDilationsInfo() const { - return reinterpret_cast(getStridesInfo()) + _ndim; - } - inline const size_t *getPaddedShape() const { - return getDilationsInfo() + _ndim; - } - - inline size_t input_dim(size_t i) const { - return i < _ndim ? getInputDims()[i] : 0; - } - inline size_t kernel_dim(size_t i) const { - return i < _ndim ? getKernelDims()[i] : 0; - } - inline size_t output_dim(size_t i) const { - return i < _ndim ? getOutputDims()[i] : 0; - } - inline size_t bias_dim(size_t i) const { - return i < _bias_dims_size ? getBiasDims()[i] : 0; - } - inline size_t pad_info(size_t i) const { - return i < _ndim ? getPadsInfo()[i] : 0; - } - inline ptrdiff_t stride_info(size_t i) const { - return i < _ndim ? getStridesInfo()[i] : 0; - } - inline size_t dilation_info(size_t i) const { - return i < _ndim ? getDilationsInfo()[i] : 0; - } - inline size_t padded_shape_dim(size_t i) const { - return i < _padded_shape_size ? getPaddedShape()[i] : 0; - } - - static utils::Result create( - infiniopHandle_t handle_, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t w_desc, - infiniopTensorDescriptor_t b_desc, - const void *pads, - const void *strides, - const void *dilations, - size_t n); -}; - -inline utils::Result calculateConvOutputSize( - size_t input_size, - size_t kernel_size, - size_t padding, - size_t stride, - size_t dilation) { - if (stride == 0) { - return utils::Result(INFINI_STATUS_BAD_TENSOR_SHAPE); - } - if (dilation == 0) { - return utils::Result(INFINI_STATUS_BAD_TENSOR_SHAPE); - } - if (kernel_size == 0) { - return utils::Result(INFINI_STATUS_BAD_TENSOR_SHAPE); - } - size_t effective_kernel = dilation * (kernel_size - 1) + 1; - - size_t padded_input = input_size + 2 * padding; - - if (padded_input < effective_kernel) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t output_size = (padded_input - effective_kernel) / stride + 1; - - return utils::Result(output_size); -} - -inline utils::Result ConvInfo::create( - infiniopHandle_t handle_, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t w_desc, - infiniopTensorDescriptor_t b_desc, - const void *pads, - const void *strides, - const void *dilations, - size_t n) { - - auto dtype = y_desc->dtype(); - if (dtype != x_desc->dtype() || dtype != w_desc->dtype()) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - size_t ndim = n; - size_t new_dims = n + 2; - - if (x_desc->ndim() < new_dims || y_desc->ndim() < new_dims || w_desc->ndim() < new_dims) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t batch = x_desc->shape()[0]; - size_t in_channels = x_desc->shape()[1]; - size_t out_channels = w_desc->shape()[0]; - - if (y_desc->shape()[0] != batch || y_desc->shape()[1] != out_channels || w_desc->shape()[1] != in_channels) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t bias_dims_size = (b_desc != nullptr) ? x_desc->ndim() : 0; - - const size_t *pads_ptr = reinterpret_cast(pads); - bool has_padding = false; - if (pads_ptr != nullptr) { - for (size_t i = 0; i < ndim; ++i) { - if (pads_ptr[i] > 0) { - has_padding = true; - break; - } - } - } - size_t padded_shape_size = has_padding ? (ndim + 2) : 0; - - // 计算meta总大小 - size_t meta_size = ndim * 6 + bias_dims_size + padded_shape_size; - std::vector meta(meta_size); - - size_t *input_dims = meta.data(); - size_t *kernel_dims = input_dims + ndim; - size_t *output_dims = kernel_dims + ndim; - size_t *bias_dims = output_dims + ndim; - size_t *pads_info = bias_dims + bias_dims_size; - ptrdiff_t *strides_info = reinterpret_cast(pads_info) + ndim; - size_t *dilations_info = reinterpret_cast(strides_info) + ndim; - size_t *padded_shape = dilations_info + ndim; - - const ptrdiff_t *strides_ptr = reinterpret_cast(strides); - const size_t *dilations_ptr = reinterpret_cast(dilations); - - size_t spatial_sizes = 1; - - for (size_t i = 0; i < ndim; i++) { - input_dims[i] = x_desc->shape()[i + 2]; - kernel_dims[i] = w_desc->shape()[i + 2]; - output_dims[i] = y_desc->shape()[i + 2]; - pads_info[i] = pads_ptr == nullptr ? 0 : pads_ptr[i]; - strides_info[i] = strides_ptr == nullptr ? 1 : strides_ptr[i]; - dilations_info[i] = dilations_ptr == nullptr ? 1 : dilations_ptr[i]; - spatial_sizes = spatial_sizes * output_dims[i]; - auto output_result = calculateConvOutputSize( - input_dims[i], - kernel_dims[i], - pads_info[i], - strides_info[i], - dilations_info[i]); - CHECK_RESULT(output_result); - size_t expected_output = output_result.take(); - if (output_dims[i] != expected_output) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - if (bias_dims_size > 0) { - std::fill(bias_dims, bias_dims + bias_dims_size, 1); - bias_dims[1] = b_desc->shape()[0]; - } - - if (padded_shape_size > 0) { - padded_shape[0] = batch; - padded_shape[1] = in_channels; - for (size_t i = 0; i < ndim; ++i) { - padded_shape[i + 2] = input_dims[i] + 2 * pads_info[i]; - } - } - - ConvInfo info(std::move(meta), ndim, batch, in_channels, out_channels, - spatial_sizes, bias_dims_size, padded_shape_size); - - return utils::Result(info); -} - -} // namespace op::conv - -#endif // __CONV_INFO_H__ diff --git a/src/infiniop/ops/conv/nvidia/conv_nvidia.cu b/src/infiniop/ops/conv/nvidia/conv_nvidia.cu deleted file mode 100644 index 15e6ef565..000000000 --- a/src/infiniop/ops/conv/nvidia/conv_nvidia.cu +++ /dev/null @@ -1,458 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "conv_nvidia.cuh" - -#define DESTROY_CUDNN_DESCRIPTOR(desc_ptr, destroy_func) \ - do { \ - if (desc_ptr) { \ - destroy_func(desc_ptr); \ - desc_ptr = nullptr; \ - } \ - } while (0) - -#define CLEANUP_CUDNN_DESCRIPTORS() \ - do { \ - DESTROY_CUDNN_DESCRIPTOR(x_desc, cudnnDestroyTensorDescriptor); \ - DESTROY_CUDNN_DESCRIPTOR(y_desc, cudnnDestroyTensorDescriptor); \ - DESTROY_CUDNN_DESCRIPTOR(w_desc, cudnnDestroyFilterDescriptor); \ - DESTROY_CUDNN_DESCRIPTOR(b_desc, cudnnDestroyTensorDescriptor); \ - DESTROY_CUDNN_DESCRIPTOR(act_desc, cudnnDestroyActivationDescriptor); \ - DESTROY_CUDNN_DESCRIPTOR(conv_desc, cudnnDestroyConvolutionDescriptor); \ - } while (0) - -namespace op::conv::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; - size_t workspace_size = 0; - -#ifdef ENABLE_CUDNN_API - cudnnTensorDescriptor_t x_desc = nullptr; - cudnnTensorDescriptor_t y_desc = nullptr; - cudnnFilterDescriptor_t w_desc = nullptr; - cudnnTensorDescriptor_t b_desc = nullptr; - cudnnActivationDescriptor_t act_desc = nullptr; - cudnnConvolutionDescriptor_t conv_desc = nullptr; - cudnnConvolutionFwdAlgo_t algo = CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM; -#endif - -private: - Opaque(std::shared_ptr internal_ptr) - : internal(internal_ptr) {} - -#ifdef ENABLE_CUDNN_API - void initializeDimensionArrays(const ConvInfo &info, - std::vector &input_dims, - std::vector &output_dims, - std::vector &filter_dims, - std::vector &input_strides, - std::vector &output_strides) const { - bool is_1d_conv = (info.ndim() == 1); - int actual_tensor_ndim = is_1d_conv ? 4 : static_cast(info.ndim() + 2); - - input_dims[0] = static_cast(info.batch()); - input_dims[1] = static_cast(info.in_channels()); - output_dims[0] = static_cast(info.batch()); - output_dims[1] = static_cast(info.out_channels()); - filter_dims[0] = static_cast(info.out_channels()); - filter_dims[1] = static_cast(info.in_channels()); - - if (is_1d_conv) { - input_dims[2] = 1; - input_dims[3] = static_cast(info.input_dim(0)); - output_dims[2] = 1; - output_dims[3] = static_cast(info.output_dim(0)); - filter_dims[2] = 1; - filter_dims[3] = static_cast(info.kernel_dim(0)); - } else { - for (size_t i = 0; i < info.ndim(); ++i) { - input_dims[i + 2] = static_cast(info.input_dim(i)); - output_dims[i + 2] = static_cast(info.output_dim(i)); - filter_dims[i + 2] = static_cast(info.kernel_dim(i)); - } - } - calculateStrides(input_dims, input_strides, actual_tensor_ndim); - calculateStrides(output_dims, output_strides, actual_tensor_ndim); - } - - void initializeConvolutionParams(const ConvInfo &info, - std::vector &pads, - std::vector &strides, - std::vector &dilations) const { - bool is_1d_conv = (info.ndim() == 1); - - if (is_1d_conv) { - pads[0] = 0; - pads[1] = static_cast(info.pad_info(0)); - strides[0] = 1; - strides[1] = static_cast(info.stride_info(0)); - dilations[0] = 1; - dilations[1] = static_cast(info.dilation_info(0)); - } else { - for (size_t i = 0; i < info.ndim(); ++i) { - pads[i] = static_cast(info.pad_info(i)); - strides[i] = static_cast(info.stride_info(i)); - dilations[i] = static_cast(info.dilation_info(i)); - } - } - } - - void calculateStrides(const std::vector &dims, - std::vector &strides, - int ndim) const { - strides[ndim - 1] = 1; - for (int d = ndim - 2; d >= 0; --d) { - strides[d] = strides[d + 1] * dims[d + 1]; - } - } - - infiniStatus_t getCudnnDataType(infiniDtype_t data_type, - cudnnDataType_t &cudnn_data_type) const { - if (data_type == INFINI_DTYPE_F16) { - cudnn_data_type = device::nvidia::getCudnnDtype(data_type); - } else if (data_type == INFINI_DTYPE_F32) { - cudnn_data_type = device::nvidia::getCudnnDtype(data_type); - } else if (data_type == INFINI_DTYPE_BF16) { - cudnn_data_type = device::nvidia::getCudnnDtype(data_type); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; - } - - infiniStatus_t createBasicDescriptors(const std::vector &input_dims, - const std::vector &output_dims, - const std::vector &filter_dims, - cudnnDataType_t cudnn_data_type, - int actual_tensor_ndim) { - CHECK_CUDNN(cudnnCreateTensorDescriptor(&x_desc)); - CHECK_CUDNN(cudnnCreateTensorDescriptor(&y_desc)); - CHECK_CUDNN(cudnnCreateFilterDescriptor(&w_desc)); - CHECK_CUDNN(cudnnCreateConvolutionDescriptor(&conv_desc)); - - CHECK_CUDNN(cudnnSetTensorNdDescriptorEx( - x_desc, CUDNN_TENSOR_NCHW, cudnn_data_type, - actual_tensor_ndim, input_dims.data())); - CHECK_CUDNN(cudnnSetTensorNdDescriptorEx( - y_desc, CUDNN_TENSOR_NCHW, cudnn_data_type, - actual_tensor_ndim, output_dims.data())); - CHECK_CUDNN(cudnnSetFilterNdDescriptor( - w_desc, cudnn_data_type, CUDNN_TENSOR_NCHW, - actual_tensor_ndim, filter_dims.data())); - - return INFINI_STATUS_SUCCESS; - } - - infiniStatus_t createBiasDescriptors(const ConvInfo &info, - cudnnDataType_t cudnn_data_type, - int actual_tensor_ndim) { - if (info.bias_dims_size() == 0) { - b_desc = nullptr; - act_desc = nullptr; - return INFINI_STATUS_SUCCESS; - } - - std::vector bias_dims_arr(actual_tensor_ndim); - std::vector bias_strides_arr(actual_tensor_ndim); - - bias_dims_arr[0] = 1; - bias_dims_arr[1] = static_cast(info.out_channels()); - for (int i = 2; i < actual_tensor_ndim; ++i) { - bias_dims_arr[i] = 1; - } - - if (actual_tensor_ndim == 4) { - bias_strides_arr[0] = static_cast(info.out_channels()); - bias_strides_arr[1] = 1; - bias_strides_arr[2] = 1; - bias_strides_arr[3] = 1; - } else { - calculateStrides(bias_dims_arr, bias_strides_arr, actual_tensor_ndim); - } - - CHECK_CUDNN(cudnnCreateTensorDescriptor(&b_desc)); - CHECK_CUDNN(cudnnSetTensorNdDescriptor( - b_desc, cudnn_data_type, static_cast(bias_dims_arr.size()), - bias_dims_arr.data(), bias_strides_arr.data())); - CHECK_CUDNN(cudnnCreateActivationDescriptor(&act_desc)); - CHECK_CUDNN(cudnnSetActivationDescriptor( - act_desc, CUDNN_ACTIVATION_IDENTITY, CUDNN_NOT_PROPAGATE_NAN, 0.0)); - - return INFINI_STATUS_SUCCESS; - } - - infiniStatus_t setupConvolutionDescriptor(const std::vector &pads, - const std::vector &strides, - const std::vector &dilations, - int spatial_ndim, - cudnnDataType_t compute_type) { - CHECK_CUDNN(cudnnSetConvolutionNdDescriptor( - conv_desc, - spatial_ndim, - pads.data(), - strides.data(), - dilations.data(), - CUDNN_CROSS_CORRELATION, - compute_type)); - - return INFINI_STATUS_SUCCESS; - } - - infiniStatus_t setupAlgorithmWithoutBias() { - algo = CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM; - CHECK_STATUS(internal->useCudnn( - nullptr, - [&](cudnnHandle_t handle) { - CHECK_CUDNN(cudnnGetConvolutionForwardWorkspaceSize( - handle, x_desc, w_desc, conv_desc, y_desc, - algo, &workspace_size)); - return INFINI_STATUS_SUCCESS; - })); - return INFINI_STATUS_SUCCESS; - } - - infiniStatus_t setupAlgorithmWithBias() { - int maxAlgoCount = 0; - CHECK_STATUS(internal->useCudnn( - nullptr, - [&](cudnnHandle_t handle) { - CHECK_CUDNN(cudnnGetConvolutionForwardAlgorithmMaxCount(handle, &maxAlgoCount)); - return INFINI_STATUS_SUCCESS; - })); - - if (maxAlgoCount <= 0) { - maxAlgoCount = 8; - } - - std::vector perf_results(maxAlgoCount); - int algoCounts = 0; - - CHECK_STATUS(internal->useCudnn( - nullptr, [&](cudnnHandle_t handle) { - CHECK_CUDNN(cudnnFindConvolutionForwardAlgorithm( - handle, x_desc, w_desc, conv_desc, y_desc, - maxAlgoCount, &algoCounts, perf_results.data())); - return INFINI_STATUS_SUCCESS; - })); - - if (algoCounts < 1) { - return INFINI_STATUS_BAD_PARAM; - } - - for (int i = 0; i < algoCounts; ++i) { - CHECK_STATUS(internal->useCudnn( - nullptr, - [&](cudnnHandle_t handle) { - CHECK_CUDNN(cudnnGetConvolutionForwardWorkspaceSize( - handle, x_desc, w_desc, conv_desc, y_desc, - perf_results[i].algo, &workspace_size)); - return INFINI_STATUS_SUCCESS; - })); - algo = perf_results[i].algo; - break; - } - - return INFINI_STATUS_SUCCESS; - } -#endif - -public: - Opaque(Opaque &&other) noexcept - : internal(std::move(other.internal)), - workspace_size(other.workspace_size) - // clang-format off -#ifdef ENABLE_CUDNN_API - , x_desc(other.x_desc) - , y_desc(other.y_desc) - , w_desc(other.w_desc) - , b_desc(other.b_desc) - , act_desc(other.act_desc) - , conv_desc(other.conv_desc) - , algo(other.algo) -#endif - // clang-format on - { -#ifdef ENABLE_CUDNN_API - other.x_desc = nullptr; - other.y_desc = nullptr; - other.w_desc = nullptr; - other.b_desc = nullptr; - other.act_desc = nullptr; - other.conv_desc = nullptr; -#endif - other.workspace_size = 0; - } - - ~Opaque() { -#ifdef ENABLE_CUDNN_API - CLEANUP_CUDNN_DESCRIPTORS(); -#endif - } - -#ifdef ENABLE_CUDNN_API - infiniStatus_t initializeCudnnContext(ConvInfo &info, - infiniDtype_t data_type, - cudnnDataType_t compute_type) { - bool is_1d_conv = (info.ndim() == 1); - int actual_tensor_ndim = is_1d_conv ? 4 : static_cast(info.ndim() + 2); - int spatial_ndim_for_conv_desc = is_1d_conv ? 2 : static_cast(info.ndim()); - - std::vector input_dims_arr(actual_tensor_ndim); - std::vector output_dims_arr(actual_tensor_ndim); - std::vector filter_dims_arr(actual_tensor_ndim); - std::vector input_strides_arr(actual_tensor_ndim); - std::vector output_strides_arr(actual_tensor_ndim); - - initializeDimensionArrays(info, input_dims_arr, output_dims_arr, - filter_dims_arr, input_strides_arr, output_strides_arr); - - std::vector pads_arr(spatial_ndim_for_conv_desc); - std::vector strides_arr(spatial_ndim_for_conv_desc); - std::vector dilations_arr(spatial_ndim_for_conv_desc); - - initializeConvolutionParams(info, pads_arr, strides_arr, dilations_arr); - - cudnnDataType_t cudnn_data_type; - CHECK_STATUS(getCudnnDataType(data_type, cudnn_data_type)); - - CHECK_STATUS(createBasicDescriptors(input_dims_arr, output_dims_arr, - filter_dims_arr, cudnn_data_type, actual_tensor_ndim)); - - CHECK_STATUS(createBiasDescriptors(info, cudnn_data_type, actual_tensor_ndim)); - - CHECK_STATUS(setupConvolutionDescriptor(pads_arr, strides_arr, dilations_arr, - spatial_ndim_for_conv_desc, compute_type)); - - // @todo: improve algo setup in the future - // if (info.bias_dims_size() == 0) { - // CHECK_STATUS(setupAlgorithmWithoutBias()); - // } else { - // CHECK_STATUS(setupAlgorithmWithBias()); - // } - CHECK_STATUS(setupAlgorithmWithoutBias()); - - return INFINI_STATUS_SUCCESS; - } -#endif - - static inline utils::Result create( - std::shared_ptr internal_ptr, - ConvInfo &info, - infiniDtype_t data_type) { -#ifdef ENABLE_CUDNN_API - Opaque opaque(internal_ptr); - auto status = opaque.initializeCudnnContext(info, data_type, CUDNN_DATA_FLOAT); - if (status != INFINI_STATUS_SUCCESS) { - return status; - } - return utils::Result(std::move(opaque)); -#else - return INFINI_STATUS_NOT_IMPLEMENTED; -#endif - } -}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t w_desc, - infiniopTensorDescriptor_t b_desc, - const void *pads, - const void *strides, - const void *dilations, - size_t n) { -#ifdef ENABLE_CUDNN_API - auto handle = reinterpret_cast(handle_); - auto dtype = y_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = ConvInfo::create(handle_, y_desc, x_desc, w_desc, b_desc, - pads, strides, dilations, n); - - CHECK_RESULT(result); - auto conv_info = result.take(); - auto opaque_result = Opaque::create(handle->internal(), conv_info, dtype); - CHECK_RESULT(opaque_result); - auto opaque = new Opaque(opaque_result.take()); - - *desc_ptr = new Descriptor( - dtype, - std::move(conv_info), - opaque->workspace_size, - opaque, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -#else - return INFINI_STATUS_NOT_IMPLEMENTED; -#endif -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *w, - const void *bias, - void *stream) const { -#ifdef ENABLE_CUDNN_API - const float alpha = 1.0f, beta = 0.0f; - if (bias != nullptr) { - CHECK_STATUS(_opaque->internal->useCudnn( - (cudaStream_t)stream, [&](cudnnHandle_t handle) { - CHECK_CUDNN(cudnnConvolutionBiasActivationForward( - handle, - &alpha, - _opaque->x_desc, - x, - _opaque->w_desc, - w, - _opaque->conv_desc, - _opaque->algo, - workspace, workspace_size, - &beta, - _opaque->y_desc, - y, - _opaque->b_desc, - bias, - _opaque->act_desc, - _opaque->y_desc, - y)); - return INFINI_STATUS_SUCCESS; - })); - } else { - CHECK_STATUS(_opaque->internal->useCudnn( - (cudaStream_t)stream, [&](cudnnHandle_t handle) { - CHECK_CUDNN(cudnnConvolutionForward( - handle, - &alpha, - _opaque->x_desc, - x, - _opaque->w_desc, - w, - _opaque->conv_desc, - _opaque->algo, - workspace, workspace_size, - &beta, - _opaque->y_desc, - y)); - return INFINI_STATUS_SUCCESS; - })); - } - - return INFINI_STATUS_SUCCESS; -#else - return INFINI_STATUS_NOT_IMPLEMENTED; -#endif -} -} // namespace op::conv::nvidia diff --git a/src/infiniop/ops/conv/nvidia/conv_nvidia.cuh b/src/infiniop/ops/conv/nvidia/conv_nvidia.cuh deleted file mode 100644 index 19bef0980..000000000 --- a/src/infiniop/ops/conv/nvidia/conv_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CONV_CUDA_CUH__ -#define __CONV_CUDA_CUH__ - -#include "../conv.h" - -DESCRIPTOR(nvidia) - -#endif // __GEMM_CUDA_CUH__ diff --git a/src/infiniop/ops/conv/operator.cc b/src/infiniop/ops/conv/operator.cc deleted file mode 100644 index 5d27c9232..000000000 --- a/src/infiniop/ops/conv/operator.cc +++ /dev/null @@ -1,190 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/conv.h" - -#ifdef ENABLE_CAMBRICON_API -#include "bang/conv_bang.h" -#endif - -#ifdef ENABLE_CPU_API -#include "cpu/conv_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) -#include "nvidia/conv_nvidia.cuh" -#endif - -__INFINI_C __export infiniStatus_t infiniopCreateConvDescriptor(infiniopHandle_t handle, - infiniopConvDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t w_desc, - infiniopTensorDescriptor_t b_desc, - void *pads, - void *strides, - void *dilations, - size_t n) { -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::conv::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - x_desc, \ - w_desc, \ - b_desc, \ - pads, \ - strides, \ - dilations, \ - n) - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -__INFINI_C infiniStatus_t -infiniopGetConvWorkspaceSize( - infiniopConvDescriptor_t desc, - size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ALI_API - GET(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_CAMBRICON_API - GET(INFINI_DEVICE_CAMBRICON, bang); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET -} - -__INFINI_C infiniStatus_t infiniopConv( - infiniopConvDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *w, - const void *bias, - void *stream) { -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, \ - y, \ - x, \ - w, \ - bias, \ - stream) - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyConvDescriptor(infiniopConvDescriptor_t desc) { -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ALI_API - DELETE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_CAMBRICON_API - DELETE(INFINI_DEVICE_CAMBRICON, bang); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} diff --git a/src/infiniop/ops/cross_entropy/cpu/cross_entropy_cpu.cc b/src/infiniop/ops/cross_entropy/cpu/cross_entropy_cpu.cc deleted file mode 100644 index c1098f3ee..000000000 --- a/src/infiniop/ops/cross_entropy/cpu/cross_entropy_cpu.cc +++ /dev/null @@ -1,99 +0,0 @@ -#include "cross_entropy_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include "../../../reduce/cpu/reduce.h" -#include -#include - -namespace op::cross_entropy::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t target_desc) { - - auto x_dtype = x_desc->dtype(); - auto t_dtype = target_desc->dtype(); - - CHECK_DTYPE(x_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - CHECK_DTYPE(t_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CrossEntropyInfo info{}; - info.dtype = x_dtype; - info.target_dtype = t_dtype; - - info.outer_size = target_desc->numel(); - - info.vocab_size = x_desc->shape().back(); - - info.x_stride = static_cast(info.vocab_size); - - *desc_ptr = new Descriptor(nullptr, info, 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t cross_entropy_kernel(const CrossEntropyInfo *info, - T *y, const T *x, const void *target) { - const Tidx *label = reinterpret_cast(target); - -#pragma omp parallel for - for (ptrdiff_t i = 0; i < ptrdiff_t(info->outer_size); ++i) { - const T *row = x + i * info->x_stride; - Tidx idx = label[i]; - - if (idx < 0 || static_cast(idx) >= info->vocab_size) { - y[i] = utils::cast(0.f); - continue; - } - - float max_val = op::common_cpu::reduce_op::max(row, info->vocab_size, 1); - - float sum_exp = 0.f; - for (size_t j = 0; j < info->vocab_size; ++j) { - sum_exp += std::exp(utils::cast(row[j]) - max_val); - } - - float log_term = std::log(sum_exp) + max_val; - float target_logit = utils::cast(row[idx]); - y[i] = utils::cast(log_term - target_logit); - } - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t dispatch_target_type(const CrossEntropyInfo *info, - T *y, const T *x, const void *target) { - - if (info->target_dtype == INFINI_DTYPE_I32) { - return cross_entropy_kernel(info, y, x, target); - } else if (info->target_dtype == INFINI_DTYPE_I64) { - return cross_entropy_kernel(info, y, x, target); - } - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *target, - void *stream) const { - - switch (_info.dtype) { - case INFINI_DTYPE_F16: - return dispatch_target_type(&_info, (fp16_t *)y, (const fp16_t *)x, target); - case INFINI_DTYPE_BF16: - return dispatch_target_type(&_info, (bf16_t *)y, (const bf16_t *)x, target); - case INFINI_DTYPE_F32: - return dispatch_target_type(&_info, (float *)y, (const float *)x, target); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::cross_entropy::cpu diff --git a/src/infiniop/ops/cross_entropy/cpu/cross_entropy_cpu.h b/src/infiniop/ops/cross_entropy/cpu/cross_entropy_cpu.h deleted file mode 100644 index e274efc9d..000000000 --- a/src/infiniop/ops/cross_entropy/cpu/cross_entropy_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CROSS_ENTROPY_CPU_H__ -#define __CROSS_ENTROPY_CPU_H__ - -#include "../cross_entropy.h" - -DESCRIPTOR(cpu) - -#endif diff --git a/src/infiniop/ops/cross_entropy/cross_entropy.h b/src/infiniop/ops/cross_entropy/cross_entropy.h deleted file mode 100644 index 075b17142..000000000 --- a/src/infiniop/ops/cross_entropy/cross_entropy.h +++ /dev/null @@ -1,42 +0,0 @@ -#ifndef CROSS_ENTROPY_H -#define CROSS_ENTROPY_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - namespace op::cross_entropy::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - CrossEntropyInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor(Opaque *opaque, \ - CrossEntropyInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - size_t workspaceSize() const { return _workspace_size; } \ - static infiniStatus_t create(infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t y_desc, \ - infiniopTensorDescriptor_t x_desc, \ - infiniopTensorDescriptor_t target_desc); \ - infiniStatus_t calculate(void *workspace, \ - size_t workspace_size, \ - void *y, \ - const void *x, \ - const void *target, \ - void *stream) const; \ - }; \ - } - -#endif diff --git a/src/infiniop/ops/cross_entropy/cuda/kernel.cuh b/src/infiniop/ops/cross_entropy/cuda/kernel.cuh deleted file mode 100644 index c048c1233..000000000 --- a/src/infiniop/ops/cross_entropy/cuda/kernel.cuh +++ /dev/null @@ -1,80 +0,0 @@ -#ifndef __CROSS_ENTROPY_KERNEL_CUH__ -#define __CROSS_ENTROPY_KERNEL_CUH__ - -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../reduce/cuda/reduce.cuh" - -template -__device__ void crossEntropyKernel( - Tdata *y_, - const Tdata *x_, - const void *target_, - size_t outer_size, - size_t vocab_size, - ptrdiff_t x_stride) { - - size_t row_idx = blockIdx.x; - if (row_idx >= outer_size) { - return; - } - - const Tdata *x = x_ + row_idx * x_stride; - const Tidx *target = reinterpret_cast(target_); - - Tidx label = target[row_idx]; - - Tdata max_val_raw = op::common_cuda::reduce_op::max(x, vocab_size); - __shared__ Tcompute max_val_shared; - if (threadIdx.x == 0) { - max_val_shared = static_cast(max_val_raw); - } - __syncthreads(); - Tcompute max_val = max_val_shared; - - Tcompute thread_sum = 0.0f; - for (size_t col = threadIdx.x; col < vocab_size; col += BLOCK_SIZE) { - Tcompute val = static_cast(x[col]); - thread_sum += expf(val - max_val); - } - - for (int offset = warpSize / 2; offset > 0; offset /= 2) { - thread_sum += __shfl_down_sync(0xffffffff, thread_sum, offset); - } - - static __shared__ Tcompute shared_sum[32]; - int lane = threadIdx.x % warpSize; - int warp = threadIdx.x / warpSize; - - if (lane == 0) { - shared_sum[warp] = thread_sum; - } - __syncthreads(); - - Tcompute block_sum = 0.0f; - if (warp == 0) { - - if (lane < (BLOCK_SIZE + warpSize - 1) / warpSize) { - block_sum = shared_sum[lane]; - } - for (int offset = warpSize / 2; offset > 0; offset /= 2) { - block_sum += __shfl_down_sync(0xffffffff, block_sum, offset); - } - } - - if (threadIdx.x == 0) { - Tcompute log_term = logf(block_sum) + max_val; - - Tcompute target_logit = 0.0f; - - if (label >= 0 && static_cast(label) < vocab_size) { - target_logit = static_cast(x[label]); - } else { - - log_term = 0.0f; - } - - y_[row_idx] = static_cast(log_term - target_logit); - } -} - -#endif diff --git a/src/infiniop/ops/cross_entropy/info.h b/src/infiniop/ops/cross_entropy/info.h deleted file mode 100644 index a915a4fe4..000000000 --- a/src/infiniop/ops/cross_entropy/info.h +++ /dev/null @@ -1,17 +0,0 @@ -#ifndef CROSS_ENTROPY_INFO_H -#define CROSS_ENTROPY_INFO_H -#include "../../../utils.h" -#include "../../tensor.h" -#include - -#include - -struct CrossEntropyInfo { - int dtype; - int target_dtype; - size_t outer_size; - size_t vocab_size; - ptrdiff_t x_stride; -}; - -#endif diff --git a/src/infiniop/ops/cross_entropy/metax/cross_entropy_metax.h b/src/infiniop/ops/cross_entropy/metax/cross_entropy_metax.h deleted file mode 100644 index 57bccea91..000000000 --- a/src/infiniop/ops/cross_entropy/metax/cross_entropy_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CROSS_ENTROPY_METAX_H__ -#define __CROSS_ENTROPY_METAX_H__ - -#include "../cross_entropy.h" - -DESCRIPTOR(metax) - -#endif // __CROSS_ENTROPY_METAX_H__ diff --git a/src/infiniop/ops/cross_entropy/metax/cross_entropy_metax.maca b/src/infiniop/ops/cross_entropy/metax/cross_entropy_metax.maca deleted file mode 100644 index efd791183..000000000 --- a/src/infiniop/ops/cross_entropy/metax/cross_entropy_metax.maca +++ /dev/null @@ -1,188 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "cross_entropy_metax.h" -#include "../../../devices/metax/metax_kernel_common.h" - -#include - -#include "../../../reduce/cuda/reduce.cuh" - -#include - -namespace { - -template -__device__ void crossEntropyKernel( - Tdata *y_, - const Tdata *x_, - const void *target_, - size_t outer_size, - size_t vocab_size, - ptrdiff_t x_stride) { - - size_t row_idx = blockIdx.x; - if (row_idx >= outer_size) { - return; - } - - const Tdata *x = x_ + row_idx * x_stride; - const Tidx *target = reinterpret_cast(target_); - - Tidx label = target[row_idx]; - - Tdata max_val_raw = op::common_cuda::reduce_op::max(x, vocab_size); - __shared__ Tcompute max_val_shared; - if (threadIdx.x == 0) { - max_val_shared = static_cast(max_val_raw); - } - __syncthreads(); - - Tcompute max_val = max_val_shared; - - Tcompute thread_sum = Tcompute(0); - for (size_t col = threadIdx.x; col < vocab_size; col += BLOCK_SIZE) { - Tcompute val = static_cast(x[col]); - thread_sum += expf(val - max_val); - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - Tcompute block_sum = BlockReduce(temp_storage).Sum(thread_sum); - - if (threadIdx.x == 0) { - if (label < 0 || static_cast(label) >= vocab_size) { - y_[row_idx] = static_cast(0.0f); - return; - } - Tcompute log_term = logf(block_sum) + max_val; - Tcompute target_logit = static_cast(x[label]); - y_[row_idx] = static_cast(log_term - target_logit); - } -} - -template -INFINIOP_METAX_KERNEL crossEntropy( - Tdata *y, const Tdata *x, const void *target, - size_t outer_size, size_t vocab_size, ptrdiff_t x_stride) { - crossEntropyKernel( - y, x, target, outer_size, vocab_size, x_stride); -} - -} // namespace - -namespace op::cross_entropy::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t target_desc) { - - (void)y_desc; - - auto x_dtype = x_desc->dtype(); - auto t_dtype = target_desc->dtype(); - - CHECK_DTYPE(x_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - CHECK_DTYPE(t_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CrossEntropyInfo info{}; - info.dtype = x_dtype; - info.target_dtype = t_dtype; - info.vocab_size = x_desc->shape().back(); - info.outer_size = target_desc->numel(); - info.x_stride = static_cast(info.vocab_size); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info, 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(void *y, const void *x, const void *target, - const CrossEntropyInfo &info, hcStream_t stream) { - dim3 grid(static_cast(info.outer_size), 1, 1); - - if (info.target_dtype == INFINI_DTYPE_I64) { - if (info.dtype == INFINI_DTYPE_F16) { - crossEntropy - <<>>( - (half *)y, (const half *)x, target, - info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_BF16) { - crossEntropy - <<>>( - (cuda_bfloat16 *)y, (const cuda_bfloat16 *)x, target, - info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_F32) { - crossEntropy - <<>>( - (float *)y, (const float *)x, target, - info.outer_size, info.vocab_size, info.x_stride); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (info.target_dtype == INFINI_DTYPE_I32) { - if (info.dtype == INFINI_DTYPE_F16) { - crossEntropy - <<>>( - (half *)y, (const half *)x, target, - info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_BF16) { - crossEntropy - <<>>( - (cuda_bfloat16 *)y, (const cuda_bfloat16 *)x, target, - info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_F32) { - crossEntropy - <<>>( - (float *)y, (const float *)x, target, - info.outer_size, info.vocab_size, info.x_stride); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *target, - void *stream_) const { - - (void)workspace; - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - auto stream = reinterpret_cast(stream_); - int max_threads = _opaque->internal->maxThreadsPerBlock(); - - if (max_threads >= METAX_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel(y, x, target, _info, stream)); - } else if (max_threads >= METAX_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel(y, x, target, _info, stream)); - } else { - CHECK_STATUS(launchKernel<256>(y, x, target, _info, stream)); - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::cross_entropy::metax diff --git a/src/infiniop/ops/cross_entropy/moore/cross_entropy_kernel.h b/src/infiniop/ops/cross_entropy/moore/cross_entropy_kernel.h deleted file mode 100644 index 6648b0e32..000000000 --- a/src/infiniop/ops/cross_entropy/moore/cross_entropy_kernel.h +++ /dev/null @@ -1,53 +0,0 @@ -#ifndef __CROSS_ENTROPY_KERNEL_CUH__ -#define __CROSS_ENTROPY_KERNEL_CUH__ - -template -__device__ void crossEntropyKernel( - Tdata *y_, - const Tdata *x_, - const void *target_, - size_t outer_size, - size_t vocab_size, - ptrdiff_t x_stride) { - - size_t row_idx = blockIdx.x; - if (row_idx >= outer_size) { - return; - } - - const Tdata *x = x_ + row_idx * x_stride; - const Tidx *target = reinterpret_cast(target_); - - Tidx label = target[row_idx]; - - Tdata max_val_raw = op::common_cuda::reduce_op::max(x, vocab_size); - __shared__ Tcompute max_val_shared; - if (threadIdx.x == 0) { - max_val_shared = static_cast(max_val_raw); - } - __syncthreads(); - - Tcompute max_val = max_val_shared; - - Tcompute thread_sum = Tcompute(0); - for (size_t col = threadIdx.x; col < vocab_size; col += BLOCK_SIZE) { - Tcompute val = static_cast(x[col]); - thread_sum += expf(val - max_val); - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - Tcompute block_sum = BlockReduce(temp_storage).Sum(thread_sum); - - if (threadIdx.x == 0) { - if (label < 0 || static_cast(label) >= vocab_size) { - y_[row_idx] = static_cast(0.0f); - return; - } - Tcompute log_term = logf(block_sum) + max_val; - Tcompute target_logit = static_cast(x[label]); - y_[row_idx] = static_cast(log_term - target_logit); - } -} - -#endif diff --git a/src/infiniop/ops/cross_entropy/moore/cross_entropy_moore.h b/src/infiniop/ops/cross_entropy/moore/cross_entropy_moore.h deleted file mode 100644 index 454b14617..000000000 --- a/src/infiniop/ops/cross_entropy/moore/cross_entropy_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CROSS_ENTROPY_MOORE_H__ -#define __CROSS_ENTROPY_MOORE_H__ - -#include "../cross_entropy.h" - -DESCRIPTOR(moore) - -#endif diff --git a/src/infiniop/ops/cross_entropy/moore/cross_entropy_moore.mu b/src/infiniop/ops/cross_entropy/moore/cross_entropy_moore.mu deleted file mode 100644 index 2535679dd..000000000 --- a/src/infiniop/ops/cross_entropy/moore/cross_entropy_moore.mu +++ /dev/null @@ -1,129 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "cross_entropy_moore.h" - -#include -#include "../../../devices/moore/moore_kernel_common.h" - -#include "../../../reduce/cuda/reduce.cuh" - -#include "cross_entropy_kernel.h" - -template -INFINIOP_MOORE_KERNEL crossEntropy( - Tdata *y, const Tdata *x, const void *target, - size_t outer_size, size_t vocab_size, ptrdiff_t x_stride) { - crossEntropyKernel( - y, x, target, outer_size, vocab_size, x_stride); -} - -namespace op::cross_entropy::moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t target_desc) { - - (void)y_desc; - - auto x_dtype = x_desc->dtype(); - auto t_dtype = target_desc->dtype(); - - CHECK_DTYPE(x_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - CHECK_DTYPE(t_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CrossEntropyInfo info{}; - info.dtype = x_dtype; - info.target_dtype = t_dtype; - info.vocab_size = x_desc->shape().back(); - info.outer_size = target_desc->numel(); - info.x_stride = static_cast(info.vocab_size); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info, 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(void *y, const void *x, const void *target, - const CrossEntropyInfo &info, musaStream_t stream) { - dim3 grid(static_cast(info.outer_size), 1, 1); - - if (info.target_dtype == INFINI_DTYPE_I64) { - if (info.dtype == INFINI_DTYPE_F16) { - crossEntropy - <<>>( - (half *)y, (const half *)x, target, - info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_BF16) { - crossEntropy - <<>>( - (__mt_bfloat16 *)y, (const __mt_bfloat16 *)x, target, - info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_F32) { - crossEntropy - <<>>( - (float *)y, (const float *)x, target, - info.outer_size, info.vocab_size, info.x_stride); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (info.target_dtype == INFINI_DTYPE_I32) { - if (info.dtype == INFINI_DTYPE_F16) { - crossEntropy - <<>>( - (half *)y, (const half *)x, target, - info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_BF16) { - crossEntropy - <<>>( - (__mt_bfloat16 *)y, (const __mt_bfloat16 *)x, target, - info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_F32) { - crossEntropy - <<>>( - (float *)y, (const float *)x, target, - info.outer_size, info.vocab_size, info.x_stride); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *y, - const void *x, - const void *target, - void *stream_) const { - musaStream_t stream = (musaStream_t)stream_; - (void)workspace; - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel(y, x, target, _info, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel(y, x, target, _info, stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::cross_entropy::moore diff --git a/src/infiniop/ops/cross_entropy/nvidia/cross_entropy_nvidia.cu b/src/infiniop/ops/cross_entropy/nvidia/cross_entropy_nvidia.cu deleted file mode 100644 index 77e3d2d58..000000000 --- a/src/infiniop/ops/cross_entropy/nvidia/cross_entropy_nvidia.cu +++ /dev/null @@ -1,107 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../cuda/kernel.cuh" -#include "cross_entropy_nvidia.cuh" - -template -INFINIOP_CUDA_KERNEL crossEntropy( - Tdata *y, const Tdata *x, const void *target, - size_t outer_size, size_t vocab_size, ptrdiff_t x_stride) { - - crossEntropyKernel( - y, x, target, outer_size, vocab_size, x_stride); -} - -namespace op::cross_entropy::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t target_desc) { - - auto x_dtype = x_desc->dtype(); - auto t_dtype = target_desc->dtype(); - - CrossEntropyInfo info; - info.dtype = x_dtype; - info.target_dtype = t_dtype; - - info.vocab_size = x_desc->shape().back(); - info.outer_size = target_desc->numel(); - info.x_stride = static_cast(info.vocab_size); - - auto internal = reinterpret_cast(handle)->internal(); - - *desc_ptr = new Descriptor( - new Opaque{internal}, - info, 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(void *y, const void *x, const void *target, - const CrossEntropyInfo &info, cudaStream_t stream) { - - dim3 grid(static_cast(info.outer_size), 1, 1); - - if (info.target_dtype == INFINI_DTYPE_I64) { - if (info.dtype == INFINI_DTYPE_F16) { - crossEntropy - <<>>((half *)y, (const half *)x, target, info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_BF16) { - crossEntropy - <<>>((__nv_bfloat16 *)y, (const __nv_bfloat16 *)x, target, info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_F32) { - crossEntropy - <<>>((float *)y, (const float *)x, target, info.outer_size, info.vocab_size, info.x_stride); - } - } else if (info.target_dtype == INFINI_DTYPE_I32) { - - if (info.dtype == INFINI_DTYPE_F16) { - crossEntropy - <<>>((half *)y, (const half *)x, target, info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_BF16) { - crossEntropy - <<>>((__nv_bfloat16 *)y, (const __nv_bfloat16 *)x, target, info.outer_size, info.vocab_size, info.x_stride); - } else if (info.dtype == INFINI_DTYPE_F32) { - crossEntropy - <<>>((float *)y, (const float *)x, target, info.outer_size, info.vocab_size, info.x_stride); - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *y, - const void *x, - const void *target, - void *stream_) const { - cudaStream_t stream = (cudaStream_t)stream_; - - int max_threads = _opaque->internal->maxThreadsPerBlock(); - - if (max_threads >= 1024) { - CHECK_STATUS(launchKernel<1024>(y, x, target, _info, stream)); - } else if (max_threads >= 512) { - CHECK_STATUS(launchKernel<512>(y, x, target, _info, stream)); - } else { - CHECK_STATUS(launchKernel<256>(y, x, target, _info, stream)); - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::cross_entropy::nvidia diff --git a/src/infiniop/ops/cross_entropy/nvidia/cross_entropy_nvidia.cuh b/src/infiniop/ops/cross_entropy/nvidia/cross_entropy_nvidia.cuh deleted file mode 100644 index 786e9d88f..000000000 --- a/src/infiniop/ops/cross_entropy/nvidia/cross_entropy_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __CROSS_ENTROPY_NVIDIA_H__ -#define __CROSS_ENTROPY_NVIDIA_H__ - -#include "../cross_entropy.h" - -DESCRIPTOR(nvidia) - -#endif diff --git a/src/infiniop/ops/cross_entropy/operator.cc b/src/infiniop/ops/cross_entropy/operator.cc deleted file mode 100644 index 75f35fcb7..000000000 --- a/src/infiniop/ops/cross_entropy/operator.cc +++ /dev/null @@ -1,174 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/cross_entropy.h" - -#ifdef ENABLE_CPU_API -#include "cpu/cross_entropy_cpu.h" -#endif - -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/cross_entropy_nvidia.cuh" -#endif - -#ifdef ENABLE_MOORE_API -#include "moore/cross_entropy_moore.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/cross_entropy_metax.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateCrossEntropyDescriptor( - infiniopHandle_t handle, - infiniopCrossEntropyDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t target_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::cross_entropy::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, x_desc, target_desc); - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia) -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia) -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetCrossEntropyWorkspaceSize( - infiniopCrossEntropyDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopCrossEntropy( - infiniopCrossEntropyDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *target, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, x, target, stream); - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia) -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia) -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyCrossEntropyDescriptor( - infiniopCrossEntropyDescriptor_t desc) { - -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DESTROY(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - DESTROY(INFINI_DEVICE_ILUVATAR, nvidia) -#endif -#ifdef ENABLE_QY_API - DESTROY(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_HYGON_API - DESTROY(INFINI_DEVICE_HYGON, nvidia) -#endif -#ifdef ENABLE_MOORE_API - DESTROY(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_METAX_API - DESTROY(INFINI_DEVICE_METAX, metax) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DESTROY -} diff --git a/src/infiniop/ops/deepseek_moe/deepseek_moe.h b/src/infiniop/ops/deepseek_moe/deepseek_moe.h deleted file mode 100644 index f3f246886..000000000 --- a/src/infiniop/ops/deepseek_moe/deepseek_moe.h +++ /dev/null @@ -1,42 +0,0 @@ -#ifndef DEEPSEEK_MOE_H -#define DEEPSEEK_MOE_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - namespace op::deepseek_moe::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - DeepseekMoeInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor(Opaque *opaque, DeepseekMoeInfo info, size_t workspace_size, \ - infiniDevice_t device_type, int device_id) \ - : InfiniopDescriptor{device_type, device_id}, _opaque(opaque), \ - _info(info), _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - size_t workspaceSize() const { return _workspace_size; } \ - static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t hidden_desc, \ - infiniopTensorDescriptor_t topk_indices_desc, \ - infiniopTensorDescriptor_t topk_weights_desc, \ - size_t intermediate_size, size_t num_experts); \ - infiniStatus_t calculate(void *workspace, size_t workspace_size, void *out, \ - const void *hidden, const void *topk_indices, \ - const void *topk_weights, const void *const *gate_weights, \ - const void *const *up_weights, const void *const *down_weights, \ - void *stream) const; \ - infiniStatus_t calculateWithDevicePtrs(void *workspace, size_t workspace_size, \ - void *out, const void *hidden, const void *topk_indices, \ - const void *topk_weights, const void *gate_weight_ptrs, \ - const void *up_weight_ptrs, const void *down_weight_ptrs, \ - void *stream) const; \ - }; \ - } - -#endif diff --git a/src/infiniop/ops/deepseek_moe/info.h b/src/infiniop/ops/deepseek_moe/info.h deleted file mode 100644 index be115d3ab..000000000 --- a/src/infiniop/ops/deepseek_moe/info.h +++ /dev/null @@ -1,65 +0,0 @@ -#ifndef __DEEPSEEK_MOE_INFO_H__ -#define __DEEPSEEK_MOE_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::deepseek_moe { - -class DeepseekMoeInfo { - DeepseekMoeInfo() = default; - -public: - infiniDtype_t dtype; - size_t ntokens; - size_t hidden_size; - size_t topk; - size_t intermediate_size; - size_t num_experts; - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t hidden_desc, - infiniopTensorDescriptor_t topk_indices_desc, - infiniopTensorDescriptor_t topk_weights_desc, - size_t intermediate_size, - size_t num_experts) { - - auto dtype = hidden_desc->dtype(); - if (dtype != INFINI_DTYPE_F16 && dtype != INFINI_DTYPE_BF16) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - if (out_desc->dtype() != dtype || topk_indices_desc->dtype() != INFINI_DTYPE_I32 || topk_weights_desc->dtype() != INFINI_DTYPE_F32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - if (out_desc->ndim() != 2 || hidden_desc->ndim() != 2 || topk_indices_desc->ndim() != 2 || topk_weights_desc->ndim() != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - auto hidden_shape = hidden_desc->shape(); - auto out_shape = out_desc->shape(); - auto indices_shape = topk_indices_desc->shape(); - auto weights_shape = topk_weights_desc->shape(); - if (out_shape != hidden_shape || indices_shape != weights_shape || indices_shape[0] != hidden_shape[0]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (intermediate_size == 0 || num_experts == 0 || indices_shape[1] == 0 || indices_shape[1] > num_experts) { - return INFINI_STATUS_BAD_PARAM; - } - if (hidden_desc->strides()[1] != 1 || out_desc->strides()[1] != 1 || topk_indices_desc->strides()[1] != 1 || topk_weights_desc->strides()[1] != 1) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - return utils::Result(DeepseekMoeInfo{ - dtype, - hidden_shape[0], - hidden_shape[1], - indices_shape[1], - intermediate_size, - num_experts}); - } -}; - -} // namespace op::deepseek_moe - -#endif diff --git a/src/infiniop/ops/deepseek_moe/metax/deepseek_moe_metax.cuh b/src/infiniop/ops/deepseek_moe/metax/deepseek_moe_metax.cuh deleted file mode 100644 index 1ce51737b..000000000 --- a/src/infiniop/ops/deepseek_moe/metax/deepseek_moe_metax.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef DEEPSEEK_MOE_METAX_CUH -#define DEEPSEEK_MOE_METAX_CUH - -#include "../deepseek_moe.h" - -DESCRIPTOR(metax) - -#endif diff --git a/src/infiniop/ops/deepseek_moe/metax/deepseek_moe_metax.maca b/src/infiniop/ops/deepseek_moe/metax/deepseek_moe_metax.maca deleted file mode 100644 index 25fe2dcc0..000000000 --- a/src/infiniop/ops/deepseek_moe/metax/deepseek_moe_metax.maca +++ /dev/null @@ -1,317 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "deepseek_moe_metax.cuh" -#include - -namespace op::deepseek_moe::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -namespace { - -constexpr size_t align_up(size_t value, size_t alignment) { - return (value + alignment - 1) / alignment * alignment; -} - -template -__device__ float to_float(T value) { - return static_cast(value); -} - -template <> -__device__ float to_float(half value) { - return __half2float(value); -} - -template <> -__device__ float to_float<__nv_bfloat16>(__nv_bfloat16 value) { - return __bfloat162float(value); -} - -template -__device__ T from_float(float value) { - return static_cast(value); -} - -template <> -__device__ half from_float(float value) { - return __float2half_rn(value); -} - -template <> -__device__ __nv_bfloat16 from_float<__nv_bfloat16>(float value) { - return __float2bfloat16_rn(value); -} - -template -__global__ void gate_up_kernel( - T *intermediate, - const T *hidden, - const int *topk_indices, - const float *topk_weights, - const void *const *gate_weights, - const void *const *up_weights, - size_t ntokens, - size_t hidden_size, - size_t topk, - size_t intermediate_size, - size_t num_experts) { - - const size_t block_linear = static_cast(blockIdx.x) + static_cast(blockIdx.y) * gridDim.x; - const size_t route = block_linear / intermediate_size; - const size_t j = block_linear - route * intermediate_size; - if (route >= ntokens * topk) { - return; - } - const int expert = topk_indices[route]; - if (expert < 0 || static_cast(expert) >= num_experts) { - return; - } - const size_t token = route / topk; - const T *x = hidden + token * hidden_size; - const T *gate = reinterpret_cast(gate_weights[expert]) + j * hidden_size; - const T *up = reinterpret_cast(up_weights[expert]) + j * hidden_size; - - float gate_sum = 0.0f; - float up_sum = 0.0f; - for (size_t h = threadIdx.x; h < hidden_size; h += blockDim.x) { - const float xv = to_float(x[h]); - gate_sum += xv * to_float(gate[h]); - up_sum += xv * to_float(up[h]); - } - - __shared__ float gate_shared[256]; - __shared__ float up_shared[256]; - gate_shared[threadIdx.x] = gate_sum; - up_shared[threadIdx.x] = up_sum; - __syncthreads(); - - for (unsigned int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - gate_shared[threadIdx.x] += gate_shared[threadIdx.x + stride]; - up_shared[threadIdx.x] += up_shared[threadIdx.x + stride]; - } - __syncthreads(); - } - - if (threadIdx.x == 0) { - const float g = gate_shared[0]; - const float silu = g / (1.0f + expf(-g)); - intermediate[route * intermediate_size + j] = from_float(silu * up_shared[0] * topk_weights[route]); - } -} - -template -__global__ void down_kernel( - T *out, - const T *intermediate, - const int *topk_indices, - const void *const *down_weights, - size_t ntokens, - size_t hidden_size, - size_t topk, - size_t intermediate_size, - size_t num_experts) { - - const size_t linear = static_cast(blockIdx.x) + static_cast(blockIdx.y) * gridDim.x; - const size_t token = linear / hidden_size; - const size_t h = linear - token * hidden_size; - if (token >= ntokens) { - return; - } - - float acc = 0.0f; - const size_t route_base = token * topk; - const size_t count = topk * intermediate_size; - for (size_t idx = threadIdx.x; idx < count; idx += blockDim.x) { - const size_t k = idx / intermediate_size; - const size_t j = idx - k * intermediate_size; - const size_t route = route_base + k; - const int expert = topk_indices[route]; - if (expert >= 0 && static_cast(expert) < num_experts) { - const T *down = reinterpret_cast(down_weights[expert]) + h * intermediate_size; - acc += to_float(intermediate[route * intermediate_size + j]) * to_float(down[j]); - } - } - - __shared__ float shared[256]; - shared[threadIdx.x] = acc; - __syncthreads(); - for (unsigned int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - shared[threadIdx.x] += shared[threadIdx.x + stride]; - } - __syncthreads(); - } - if (threadIdx.x == 0) { - out[token * hidden_size + h] = from_float(shared[0]); - } -} - -template -infiniStatus_t launch_typed( - void *workspace, - size_t workspace_size, - const DeepseekMoeInfo &info, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *const *gate_weights, - const void *const *up_weights, - const void *const *down_weights, - hcStream_t stream, - bool weight_ptrs_on_device) { - - const size_t ptr_bytes = align_up(info.num_experts * sizeof(void *), 256); - const size_t ptr_workspace = ptr_bytes * 3; - const size_t intermediate_offset = align_up(ptr_workspace, 256); - const size_t intermediate_bytes = info.ntokens * info.topk * info.intermediate_size * sizeof(T); - if (workspace_size < intermediate_offset + intermediate_bytes) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - auto *base = reinterpret_cast(workspace); - const void *const *gate_ptrs = reinterpret_cast(base); - const void *const *up_ptrs = reinterpret_cast(base + ptr_bytes); - const void *const *down_ptrs = reinterpret_cast(base + ptr_bytes * 2); - auto *intermediate = reinterpret_cast(base + intermediate_offset); - - if (weight_ptrs_on_device) { - gate_ptrs = gate_weights; - up_ptrs = up_weights; - down_ptrs = down_weights; - } else { - auto **gate_workspace = reinterpret_cast(base); - auto **up_workspace = reinterpret_cast(base + ptr_bytes); - auto **down_workspace = reinterpret_cast(base + ptr_bytes * 2); - CHECK_METAX(hcMemcpyAsync(gate_workspace, gate_weights, info.num_experts * sizeof(void *), hcMemcpyHostToDevice, stream)); - CHECK_METAX(hcMemcpyAsync(up_workspace, up_weights, info.num_experts * sizeof(void *), hcMemcpyHostToDevice, stream)); - CHECK_METAX(hcMemcpyAsync(down_workspace, down_weights, info.num_experts * sizeof(void *), hcMemcpyHostToDevice, stream)); - gate_ptrs = gate_workspace; - up_ptrs = up_workspace; - down_ptrs = down_workspace; - } - - constexpr int threads = 256; - constexpr size_t max_grid_x = 65535; - const size_t gate_work = info.ntokens * info.topk * info.intermediate_size; - const dim3 gate_blocks( - static_cast(std::min(gate_work, max_grid_x)), - static_cast((gate_work + max_grid_x - 1) / max_grid_x)); - gate_up_kernel<<>>( - intermediate, - reinterpret_cast(hidden), - reinterpret_cast(topk_indices), - reinterpret_cast(topk_weights), - gate_ptrs, - up_ptrs, - info.ntokens, - info.hidden_size, - info.topk, - info.intermediate_size, - info.num_experts); - - const size_t down_work = info.ntokens * info.hidden_size; - const dim3 down_blocks( - static_cast(std::min(down_work, max_grid_x)), - static_cast((down_work + max_grid_x - 1) / max_grid_x)); - down_kernel<<>>( - reinterpret_cast(out), - intermediate, - reinterpret_cast(topk_indices), - down_ptrs, - info.ntokens, - info.hidden_size, - info.topk, - info.intermediate_size, - info.num_experts); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t hidden_desc, - infiniopTensorDescriptor_t topk_indices_desc, - infiniopTensorDescriptor_t topk_weights_desc, - size_t intermediate_size, - size_t num_experts) { - - auto result = DeepseekMoeInfo::create(out_desc, hidden_desc, topk_indices_desc, topk_weights_desc, intermediate_size, num_experts); - CHECK_RESULT(result); - auto info = result.take(); - - const size_t dtype_size = info.dtype == INFINI_DTYPE_F16 ? sizeof(half) : sizeof(__nv_bfloat16); - const size_t ptr_bytes = align_up(info.num_experts * sizeof(void *), 256); - const size_t intermediate_offset = align_up(ptr_bytes * 3, 256); - const size_t intermediate_bytes = info.ntokens * info.topk * info.intermediate_size * dtype_size; - const size_t workspace_size = intermediate_offset + intermediate_bytes; - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info, - workspace_size, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *const *gate_weights, - const void *const *up_weights, - const void *const *down_weights, - void *stream_) const { - - auto stream = reinterpret_cast(stream_); - if (_info.dtype == INFINI_DTYPE_F16) { - return launch_typed(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, false); - } - if (_info.dtype == INFINI_DTYPE_BF16) { - return launch_typed<__nv_bfloat16>(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, false); - } - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -infiniStatus_t Descriptor::calculateWithDevicePtrs( - void *workspace, - size_t workspace_size, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *gate_weight_ptrs, - const void *up_weight_ptrs, - const void *down_weight_ptrs, - void *stream_) const { - - auto stream = reinterpret_cast(stream_); - auto gate_weights = reinterpret_cast(gate_weight_ptrs); - auto up_weights = reinterpret_cast(up_weight_ptrs); - auto down_weights = reinterpret_cast(down_weight_ptrs); - if (_info.dtype == INFINI_DTYPE_F16) { - return launch_typed(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, true); - } - if (_info.dtype == INFINI_DTYPE_BF16) { - return launch_typed<__nv_bfloat16>(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, true); - } - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -} // namespace op::deepseek_moe::metax diff --git a/src/infiniop/ops/deepseek_moe/moore/deepseek_moe_moore.h b/src/infiniop/ops/deepseek_moe/moore/deepseek_moe_moore.h deleted file mode 100644 index a0afabf2f..000000000 --- a/src/infiniop/ops/deepseek_moe/moore/deepseek_moe_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef DEEPSEEK_MOE_MOORE_H -#define DEEPSEEK_MOE_MOORE_H - -#include "../deepseek_moe.h" - -DESCRIPTOR(moore) - -#endif diff --git a/src/infiniop/ops/deepseek_moe/moore/deepseek_moe_moore.mu b/src/infiniop/ops/deepseek_moe/moore/deepseek_moe_moore.mu deleted file mode 100644 index 188ca2a9c..000000000 --- a/src/infiniop/ops/deepseek_moe/moore/deepseek_moe_moore.mu +++ /dev/null @@ -1,317 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "deepseek_moe_moore.h" -#include - -namespace op::deepseek_moe::moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -namespace { - -constexpr size_t align_up(size_t value, size_t alignment) { - return (value + alignment - 1) / alignment * alignment; -} - -template -__device__ float to_float(T value) { - return static_cast(value); -} - -template <> -__device__ float to_float(half value) { - return __half2float(value); -} - -template <> -__device__ float to_float<__mt_bfloat16>(__mt_bfloat16 value) { - return __bfloat162float(value); -} - -template -__device__ T from_float(float value) { - return static_cast(value); -} - -template <> -__device__ half from_float(float value) { - return __float2half_rn(value); -} - -template <> -__device__ __mt_bfloat16 from_float<__mt_bfloat16>(float value) { - return __float2bfloat16(value); -} - -template -__global__ void gate_up_kernel( - T *intermediate, - const T *hidden, - const int *topk_indices, - const float *topk_weights, - const void *const *gate_weights, - const void *const *up_weights, - size_t ntokens, - size_t hidden_size, - size_t topk, - size_t intermediate_size, - size_t num_experts) { - - const size_t block_linear = static_cast(blockIdx.x) + static_cast(blockIdx.y) * gridDim.x; - const size_t route = block_linear / intermediate_size; - const size_t j = block_linear - route * intermediate_size; - if (route >= ntokens * topk) { - return; - } - const int expert = topk_indices[route]; - if (expert < 0 || static_cast(expert) >= num_experts) { - return; - } - const size_t token = route / topk; - const T *x = hidden + token * hidden_size; - const T *gate = reinterpret_cast(gate_weights[expert]) + j * hidden_size; - const T *up = reinterpret_cast(up_weights[expert]) + j * hidden_size; - - float gate_sum = 0.0f; - float up_sum = 0.0f; - for (size_t h = threadIdx.x; h < hidden_size; h += blockDim.x) { - const float xv = to_float(x[h]); - gate_sum += xv * to_float(gate[h]); - up_sum += xv * to_float(up[h]); - } - - __shared__ float gate_shared[256]; - __shared__ float up_shared[256]; - gate_shared[threadIdx.x] = gate_sum; - up_shared[threadIdx.x] = up_sum; - __syncthreads(); - - for (unsigned int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - gate_shared[threadIdx.x] += gate_shared[threadIdx.x + stride]; - up_shared[threadIdx.x] += up_shared[threadIdx.x + stride]; - } - __syncthreads(); - } - - if (threadIdx.x == 0) { - const float g = gate_shared[0]; - const float silu = g / (1.0f + expf(-g)); - intermediate[route * intermediate_size + j] = from_float(silu * up_shared[0] * topk_weights[route]); - } -} - -template -__global__ void down_kernel( - T *out, - const T *intermediate, - const int *topk_indices, - const void *const *down_weights, - size_t ntokens, - size_t hidden_size, - size_t topk, - size_t intermediate_size, - size_t num_experts) { - - const size_t linear = static_cast(blockIdx.x) + static_cast(blockIdx.y) * gridDim.x; - const size_t token = linear / hidden_size; - const size_t h = linear - token * hidden_size; - if (token >= ntokens) { - return; - } - - float acc = 0.0f; - const size_t route_base = token * topk; - const size_t count = topk * intermediate_size; - for (size_t idx = threadIdx.x; idx < count; idx += blockDim.x) { - const size_t k = idx / intermediate_size; - const size_t j = idx - k * intermediate_size; - const size_t route = route_base + k; - const int expert = topk_indices[route]; - if (expert >= 0 && static_cast(expert) < num_experts) { - const T *down = reinterpret_cast(down_weights[expert]) + h * intermediate_size; - acc += to_float(intermediate[route * intermediate_size + j]) * to_float(down[j]); - } - } - - __shared__ float shared[256]; - shared[threadIdx.x] = acc; - __syncthreads(); - for (unsigned int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - shared[threadIdx.x] += shared[threadIdx.x + stride]; - } - __syncthreads(); - } - if (threadIdx.x == 0) { - out[token * hidden_size + h] = from_float(shared[0]); - } -} - -template -infiniStatus_t launch_typed( - void *workspace, - size_t workspace_size, - const DeepseekMoeInfo &info, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *const *gate_weights, - const void *const *up_weights, - const void *const *down_weights, - musaStream_t stream, - bool weight_ptrs_on_device) { - - const size_t ptr_bytes = align_up(info.num_experts * sizeof(void *), 256); - const size_t ptr_workspace = ptr_bytes * 3; - const size_t intermediate_offset = align_up(ptr_workspace, 256); - const size_t intermediate_bytes = info.ntokens * info.topk * info.intermediate_size * sizeof(T); - if (workspace_size < intermediate_offset + intermediate_bytes) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - auto *base = reinterpret_cast(workspace); - const void *const *gate_ptrs = reinterpret_cast(base); - const void *const *up_ptrs = reinterpret_cast(base + ptr_bytes); - const void *const *down_ptrs = reinterpret_cast(base + ptr_bytes * 2); - auto *intermediate = reinterpret_cast(base + intermediate_offset); - - if (weight_ptrs_on_device) { - gate_ptrs = gate_weights; - up_ptrs = up_weights; - down_ptrs = down_weights; - } else { - auto **gate_workspace = reinterpret_cast(base); - auto **up_workspace = reinterpret_cast(base + ptr_bytes); - auto **down_workspace = reinterpret_cast(base + ptr_bytes * 2); - CHECK_MOORE(musaMemcpyAsync(gate_workspace, gate_weights, info.num_experts * sizeof(void *), musaMemcpyHostToDevice, stream)); - CHECK_MOORE(musaMemcpyAsync(up_workspace, up_weights, info.num_experts * sizeof(void *), musaMemcpyHostToDevice, stream)); - CHECK_MOORE(musaMemcpyAsync(down_workspace, down_weights, info.num_experts * sizeof(void *), musaMemcpyHostToDevice, stream)); - gate_ptrs = gate_workspace; - up_ptrs = up_workspace; - down_ptrs = down_workspace; - } - - constexpr int threads = 256; - constexpr size_t max_grid_x = 65535; - const size_t gate_work = info.ntokens * info.topk * info.intermediate_size; - const dim3 gate_blocks( - static_cast(std::min(gate_work, max_grid_x)), - static_cast((gate_work + max_grid_x - 1) / max_grid_x)); - gate_up_kernel<<>>( - intermediate, - reinterpret_cast(hidden), - reinterpret_cast(topk_indices), - reinterpret_cast(topk_weights), - gate_ptrs, - up_ptrs, - info.ntokens, - info.hidden_size, - info.topk, - info.intermediate_size, - info.num_experts); - - const size_t down_work = info.ntokens * info.hidden_size; - const dim3 down_blocks( - static_cast(std::min(down_work, max_grid_x)), - static_cast((down_work + max_grid_x - 1) / max_grid_x)); - down_kernel<<>>( - reinterpret_cast(out), - intermediate, - reinterpret_cast(topk_indices), - down_ptrs, - info.ntokens, - info.hidden_size, - info.topk, - info.intermediate_size, - info.num_experts); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t hidden_desc, - infiniopTensorDescriptor_t topk_indices_desc, - infiniopTensorDescriptor_t topk_weights_desc, - size_t intermediate_size, - size_t num_experts) { - - auto result = DeepseekMoeInfo::create(out_desc, hidden_desc, topk_indices_desc, topk_weights_desc, intermediate_size, num_experts); - CHECK_RESULT(result); - auto info = result.take(); - - const size_t dtype_size = info.dtype == INFINI_DTYPE_F16 ? sizeof(half) : sizeof(__mt_bfloat16); - const size_t ptr_bytes = align_up(info.num_experts * sizeof(void *), 256); - const size_t intermediate_offset = align_up(ptr_bytes * 3, 256); - const size_t intermediate_bytes = info.ntokens * info.topk * info.intermediate_size * dtype_size; - const size_t workspace_size = intermediate_offset + intermediate_bytes; - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info, - workspace_size, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *const *gate_weights, - const void *const *up_weights, - const void *const *down_weights, - void *stream_) const { - - auto stream = reinterpret_cast(stream_); - if (_info.dtype == INFINI_DTYPE_F16) { - return launch_typed(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, false); - } - if (_info.dtype == INFINI_DTYPE_BF16) { - return launch_typed<__mt_bfloat16>(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, false); - } - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -infiniStatus_t Descriptor::calculateWithDevicePtrs( - void *workspace, - size_t workspace_size, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *gate_weight_ptrs, - const void *up_weight_ptrs, - const void *down_weight_ptrs, - void *stream_) const { - - auto stream = reinterpret_cast(stream_); - auto gate_weights = reinterpret_cast(gate_weight_ptrs); - auto up_weights = reinterpret_cast(up_weight_ptrs); - auto down_weights = reinterpret_cast(down_weight_ptrs); - if (_info.dtype == INFINI_DTYPE_F16) { - return launch_typed(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, true); - } - if (_info.dtype == INFINI_DTYPE_BF16) { - return launch_typed<__mt_bfloat16>(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, true); - } - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -} // namespace op::deepseek_moe::moore diff --git a/src/infiniop/ops/deepseek_moe/nvidia/deepseek_moe_nvidia.cu b/src/infiniop/ops/deepseek_moe/nvidia/deepseek_moe_nvidia.cu deleted file mode 100644 index e938d42f1..000000000 --- a/src/infiniop/ops/deepseek_moe/nvidia/deepseek_moe_nvidia.cu +++ /dev/null @@ -1,1165 +0,0 @@ -#include -#include -#include - -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "deepseek_moe_nvidia.cuh" - -#include -#include -#include - -namespace op::deepseek_moe::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -namespace { - -constexpr size_t ROUTE_BATCHED_GEMM_MAX_TOKENS = 16384; -constexpr size_t EXPERT_GROUPED_GEMM_MIN_TOKENS = 2048; -constexpr size_t FUSED_DECODE_MAX_TOKENS = 32; -constexpr int FUSED_KERNEL_THREADS = 64; -constexpr size_t CHUNKED_GROUPED_GEMM_TOKENS = 4096; - -constexpr size_t align_up(size_t value, size_t alignment) { - return (value + alignment - 1) / alignment * alignment; -} - -inline char *advance_workspace(char *&ptr, size_t bytes, size_t alignment = 256) { - ptr = reinterpret_cast(align_up(reinterpret_cast(ptr), alignment)); - char *out = ptr; - ptr += bytes; - return out; -} - -template -size_t route_batched_workspace_size(const DeepseekMoeInfo &info, size_t base_offset) { - const size_t routes = info.ntokens * info.topk; - size_t bytes = base_offset; - bytes = align_up(bytes, alignof(void *)); - bytes += routes * sizeof(void *) * 9; - bytes = align_up(bytes, 256); - bytes += routes * info.intermediate_size * sizeof(T) * 3; - bytes = align_up(bytes, 256); - bytes += routes * info.hidden_size * sizeof(T); - return align_up(bytes, 256); -} - -template -size_t expert_grouped_workspace_size(const DeepseekMoeInfo &info, size_t base_offset) { - const size_t routes = info.ntokens * info.topk; - size_t bytes = base_offset; - for (int i = 0; i < 3; ++i) { - bytes = align_up(bytes, 256); - bytes += (info.num_experts + 1) * sizeof(int); - } - for (int i = 0; i < 2; ++i) { - bytes = align_up(bytes, 256); - bytes += routes * sizeof(int); - } - bytes = align_up(bytes, 256); - bytes += routes * info.hidden_size * sizeof(T); - for (int i = 0; i < 2; ++i) { - bytes = align_up(bytes, 256); - bytes += routes * info.intermediate_size * sizeof(T); - } - return align_up(bytes, 256); -} - -template -__device__ float to_float(T value) { - return static_cast(value); -} - -template <> -__device__ float to_float(half value) { - return __half2float(value); -} - -template <> -__device__ float to_float<__nv_bfloat16>(__nv_bfloat16 value) { - return __bfloat162float(value); -} - -template -__device__ T from_float(float value) { - return static_cast(value); -} - -template <> -__device__ half from_float(float value) { - return __float2half_rn(value); -} - -template <> -__device__ __nv_bfloat16 from_float<__nv_bfloat16>(float value) { - return __float2bfloat16_rn(value); -} - -__device__ __forceinline__ float dot_bf16x8(const uint4 &a, const uint4 &b) { - const auto *a2 = reinterpret_cast(&a); - const auto *b2 = reinterpret_cast(&b); - float sum = 0.0f; -#pragma unroll - for (int pair = 0; pair < 4; ++pair) { - const float2 av = __bfloat1622float2(a2[pair]); - const float2 bv = __bfloat1622float2(b2[pair]); - sum += av.x * bv.x + av.y * bv.y; - } - return sum; -} - -template -__global__ void gate_up_kernel( - T *intermediate, - const T *hidden, - const int *topk_indices, - const float *topk_weights, - const void *const *gate_weights, - const void *const *up_weights, - size_t ntokens, - size_t hidden_size, - size_t topk, - size_t intermediate_size, - size_t num_experts) { - - const size_t route = blockIdx.x / intermediate_size; - const size_t j = blockIdx.x - route * intermediate_size; - if (route >= ntokens * topk) { - return; - } - const int expert = topk_indices[route]; - if (expert < 0 || static_cast(expert) >= num_experts) { - return; - } - const size_t token = route / topk; - const T *x = hidden + token * hidden_size; - const T *gate = reinterpret_cast(gate_weights[expert]) + j * hidden_size; - const T *up = reinterpret_cast(up_weights[expert]) + j * hidden_size; - - float gate_sum = 0.0f; - float up_sum = 0.0f; - if constexpr (std::is_same_v) { - if ((hidden_size & 7) == 0) { - const auto *x8 = reinterpret_cast(x); - const auto *gate8 = reinterpret_cast(gate); - const auto *up8 = reinterpret_cast(up); - const size_t hidden_chunks = hidden_size / 8; - for (size_t h8 = threadIdx.x; h8 < hidden_chunks; h8 += blockDim.x) { - const uint4 xv = x8[h8]; - gate_sum += dot_bf16x8(xv, gate8[h8]); - up_sum += dot_bf16x8(xv, up8[h8]); - } - } else if ((hidden_size & 1) == 0) { - const auto *x2 = reinterpret_cast(x); - const auto *gate2 = reinterpret_cast(gate); - const auto *up2 = reinterpret_cast(up); - const size_t hidden_pairs = hidden_size / 2; - for (size_t h2 = threadIdx.x; h2 < hidden_pairs; h2 += blockDim.x) { - const float2 xv = __bfloat1622float2(x2[h2]); - const float2 gv = __bfloat1622float2(gate2[h2]); - const float2 uv = __bfloat1622float2(up2[h2]); - gate_sum += xv.x * gv.x + xv.y * gv.y; - up_sum += xv.x * uv.x + xv.y * uv.y; - } - } else { - for (size_t h = threadIdx.x; h < hidden_size; h += blockDim.x) { - const float xv = to_float(x[h]); - gate_sum += xv * to_float(gate[h]); - up_sum += xv * to_float(up[h]); - } - } - } else { - for (size_t h = threadIdx.x; h < hidden_size; h += blockDim.x) { - const float xv = to_float(x[h]); - gate_sum += xv * to_float(gate[h]); - up_sum += xv * to_float(up[h]); - } - } - - __shared__ float gate_shared[FUSED_KERNEL_THREADS]; - __shared__ float up_shared[FUSED_KERNEL_THREADS]; - gate_shared[threadIdx.x] = gate_sum; - up_shared[threadIdx.x] = up_sum; - __syncthreads(); - - for (unsigned int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - gate_shared[threadIdx.x] += gate_shared[threadIdx.x + stride]; - up_shared[threadIdx.x] += up_shared[threadIdx.x + stride]; - } - __syncthreads(); - } - - if (threadIdx.x == 0) { - const float g = gate_shared[0]; - const float silu = g / (1.0f + __expf(-g)); - intermediate[route * intermediate_size + j] = from_float(silu * up_shared[0] * topk_weights[route]); - } -} - -template -__global__ void down_kernel( - T *out, - const T *intermediate, - const int *topk_indices, - const void *const *down_weights, - size_t ntokens, - size_t hidden_size, - size_t topk, - size_t intermediate_size, - size_t num_experts) { - - const size_t linear = blockIdx.x; - const size_t token = linear / hidden_size; - const size_t h = linear - token * hidden_size; - if (token >= ntokens) { - return; - } - - float acc = 0.0f; - const size_t route_base = token * topk; - const size_t count = topk * intermediate_size; - if constexpr (std::is_same_v) { - if ((intermediate_size & 7) == 0) { - const size_t chunks_per_expert = intermediate_size / 8; - const size_t chunk_count = topk * chunks_per_expert; - for (size_t chunk_idx = threadIdx.x; chunk_idx < chunk_count; chunk_idx += blockDim.x) { - const size_t k = chunk_idx / chunks_per_expert; - const size_t j8 = chunk_idx - k * chunks_per_expert; - const size_t route = route_base + k; - const int expert = topk_indices[route]; - if (expert >= 0 && static_cast(expert) < num_experts) { - const auto *intermediate8 = reinterpret_cast( - intermediate + route * intermediate_size); - const auto *down8 = reinterpret_cast( - reinterpret_cast(down_weights[expert]) + h * intermediate_size); - acc += dot_bf16x8(intermediate8[j8], down8[j8]); - } - } - } else if ((intermediate_size & 1) == 0) { - const size_t pairs_per_expert = intermediate_size / 2; - const size_t pair_count = topk * pairs_per_expert; - for (size_t pair_idx = threadIdx.x; pair_idx < pair_count; pair_idx += blockDim.x) { - const size_t k = pair_idx / pairs_per_expert; - const size_t j2 = pair_idx - k * pairs_per_expert; - const size_t route = route_base + k; - const int expert = topk_indices[route]; - if (expert >= 0 && static_cast(expert) < num_experts) { - const auto *intermediate2 = reinterpret_cast( - intermediate + route * intermediate_size); - const auto *down2 = reinterpret_cast( - reinterpret_cast(down_weights[expert]) + h * intermediate_size); - const float2 iv = __bfloat1622float2(intermediate2[j2]); - const float2 dv = __bfloat1622float2(down2[j2]); - acc += iv.x * dv.x + iv.y * dv.y; - } - } - } else { - for (size_t idx = threadIdx.x; idx < count; idx += blockDim.x) { - const size_t k = idx / intermediate_size; - const size_t j = idx - k * intermediate_size; - const size_t route = route_base + k; - const int expert = topk_indices[route]; - if (expert >= 0 && static_cast(expert) < num_experts) { - const T *down = reinterpret_cast(down_weights[expert]) + h * intermediate_size; - acc += to_float(intermediate[route * intermediate_size + j]) * to_float(down[j]); - } - } - } - } else { - for (size_t idx = threadIdx.x; idx < count; idx += blockDim.x) { - const size_t k = idx / intermediate_size; - const size_t j = idx - k * intermediate_size; - const size_t route = route_base + k; - const int expert = topk_indices[route]; - if (expert >= 0 && static_cast(expert) < num_experts) { - const T *down = reinterpret_cast(down_weights[expert]) + h * intermediate_size; - acc += to_float(intermediate[route * intermediate_size + j]) * to_float(down[j]); - } - } - } - - __shared__ float shared[FUSED_KERNEL_THREADS]; - shared[threadIdx.x] = acc; - __syncthreads(); - for (unsigned int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (threadIdx.x < stride) { - shared[threadIdx.x] += shared[threadIdx.x + stride]; - } - __syncthreads(); - } - if (threadIdx.x == 0) { - out[token * hidden_size + h] = from_float(shared[0]); - } -} - -__global__ void count_experts_kernel( - const int *topk_indices, - int *counts, - int routes, - int num_experts) { - int route = blockIdx.x * blockDim.x + threadIdx.x; - if (route >= routes) { - return; - } - int expert = topk_indices[route]; - if (expert >= 0 && expert < num_experts) { - atomicAdd(counts + expert, 1); - } -} - -__global__ void prefix_counts_kernel( - const int *counts, - int *offsets, - int num_experts) { - extern __shared__ int scan[]; - int tid = threadIdx.x; - if (tid < num_experts) { - scan[tid] = counts[tid]; - } else { - scan[tid] = 0; - } - __syncthreads(); - - for (int stride = 1; stride < blockDim.x; stride <<= 1) { - int value = 0; - if (tid >= stride) { - value = scan[tid - stride]; - } - __syncthreads(); - scan[tid] += value; - __syncthreads(); - } - - if (tid == 0) { - offsets[0] = 0; - } - if (tid < num_experts) { - offsets[tid + 1] = scan[tid]; - } -} - -template -__global__ void pack_grouped_hidden_kernel( - const T *hidden, - const int *topk_indices, - const int *offsets, - int *positions, - int *row_to_route, - int *route_to_row, - T *packed_hidden, - int routes, - int hidden_size, - int topk, - int num_experts) { - int route = blockIdx.x; - int tid = threadIdx.x; - if (route >= routes) { - return; - } - int expert = topk_indices[route]; - if (expert < 0 || expert >= num_experts) { - return; - } - __shared__ int row; - if (tid == 0) { - int local = atomicAdd(positions + expert, 1); - row = offsets[expert] + local; - row_to_route[row] = route; - route_to_row[route] = row; - } - __syncthreads(); - int token = route / topk; - for (int h = tid; h < hidden_size; h += blockDim.x) { - packed_hidden[static_cast(row) * hidden_size + h] = hidden[static_cast(token) * hidden_size + h]; - } -} - -template -__global__ void swiglu_weight_grouped_kernel( - T *activated, - const T *gate, - const T *up, - const int *row_to_route, - const float *topk_weights, - int routes, - int intermediate_size) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total = routes * intermediate_size; - if (idx >= total) { - return; - } - int row = idx / intermediate_size; - int route = row_to_route[row]; - float g = to_float(gate[idx]); - float u = to_float(up[idx]); - float silu = g / (1.0f + __expf(-g)); - activated[idx] = from_float(silu * u * topk_weights[route]); -} - -template -__global__ void sum_grouped_out_kernel( - T *out, - const T *packed_out, - const int *route_to_row, - int ntokens, - int hidden_size, - int topk) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total = ntokens * hidden_size; - if (idx >= total) { - return; - } - int token = idx / hidden_size; - int h = idx - token * hidden_size; - float acc = 0.0f; - for (int k = 0; k < topk; ++k) { - int route = token * topk + k; - int row = route_to_row[route]; - acc += to_float(packed_out[static_cast(row) * hidden_size + h]); - } - out[idx] = from_float(acc); -} - -template -__global__ void setup_route_batched_ptrs_kernel( - const void **a_array, - const void **b_array, - void **c_array, - const T *hidden, - const int *topk_indices, - const void *const *weights, - T *gemm_out, - int routes, - int hidden_size, - int topk, - int out_size, - int num_experts) { - int route = blockIdx.x * blockDim.x + threadIdx.x; - if (route >= routes) { - return; - } - int expert = topk_indices[route]; - if (expert < 0 || expert >= num_experts) { - expert = 0; - } - int token = route / topk; - a_array[route] = weights[expert]; - b_array[route] = hidden + static_cast(token) * hidden_size; - c_array[route] = gemm_out + static_cast(route) * out_size; -} - -template -__global__ void setup_down_batched_ptrs_kernel( - const void **a_array, - const void **b_array, - void **c_array, - const T *activated, - const int *topk_indices, - const void *const *down_weights, - T *route_out, - int routes, - int hidden_size, - int topk, - int intermediate_size, - int num_experts) { - int route = blockIdx.x * blockDim.x + threadIdx.x; - if (route >= routes) { - return; - } - int expert = topk_indices[route]; - if (expert < 0 || expert >= num_experts) { - expert = 0; - } - a_array[route] = down_weights[expert]; - b_array[route] = activated + static_cast(route) * intermediate_size; - c_array[route] = route_out + static_cast(route) * hidden_size; -} - -template -__global__ void setup_all_batched_ptrs_kernel( - const void **gate_a_array, - const void **gate_b_array, - void **gate_c_array, - const void **up_a_array, - const void **up_b_array, - void **up_c_array, - const void **down_a_array, - const void **down_b_array, - void **down_c_array, - const T *hidden, - const int *topk_indices, - const void *const *gate_weights, - const void *const *up_weights, - const void *const *down_weights, - T *gate_buf, - T *up_buf, - T *activated, - T *route_out, - int routes, - int hidden_size, - int topk, - int intermediate_size, - int num_experts) { - int route = blockIdx.x * blockDim.x + threadIdx.x; - if (route >= routes) { - return; - } - int expert = topk_indices[route]; - if (expert < 0 || expert >= num_experts) { - expert = 0; - } - int token = route / topk; - const T *token_hidden = hidden + static_cast(token) * hidden_size; - - gate_a_array[route] = gate_weights[expert]; - gate_b_array[route] = token_hidden; - gate_c_array[route] = gate_buf + static_cast(route) * intermediate_size; - - up_a_array[route] = up_weights[expert]; - up_b_array[route] = token_hidden; - up_c_array[route] = up_buf + static_cast(route) * intermediate_size; - - down_a_array[route] = down_weights[expert]; - down_b_array[route] = activated + static_cast(route) * intermediate_size; - down_c_array[route] = route_out + static_cast(route) * hidden_size; -} - -template -__global__ void swiglu_weight_kernel( - T *activated, - const T *gate, - const T *up, - const float *topk_weights, - int routes, - int intermediate_size) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total = routes * intermediate_size; - if (idx >= total) { - return; - } - int route = idx / intermediate_size; - float g = to_float(gate[idx]); - float u = to_float(up[idx]); - float silu = g / (1.0f + __expf(-g)); - activated[idx] = from_float(silu * u * topk_weights[route]); -} - -template -__global__ void sum_route_out_kernel( - T *out, - const T *route_out, - int ntokens, - int hidden_size, - int topk) { - int linear = blockIdx.x * blockDim.x + threadIdx.x; - int total = ntokens * hidden_size; - if (linear >= total) { - return; - } - int token = linear / hidden_size; - int h = linear - token * hidden_size; - float acc = 0.0f; - for (int k = 0; k < topk; ++k) { - int route = token * topk + k; - acc += to_float(route_out[static_cast(route) * hidden_size + h]); - } - out[linear] = from_float(acc); -} - -template -infiniStatus_t run_batched_gemm( - const std::shared_ptr &internal, - cudaStream_t stream, - const void **a_array, - const void **b_array, - void **c_array, - int m, - int n, - int k, - int lda, - int ldb, - int ldc, - int batch_count, - cudaDataType dtype) { - float alpha = 1.0f; - float beta = 0.0f; -#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) - cudaDataType compute_type = CUDA_R_32F; -#else - cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; -#endif - CHECK_STATUS(internal->useCublas( - stream, - [&](cublasHandle_t handle) { - CHECK_CUBLAS(cublasGemmBatchedEx( - handle, - CUBLAS_OP_T, - CUBLAS_OP_N, - m, - n, - k, - &alpha, - reinterpret_cast(a_array), - dtype, - lda, - reinterpret_cast(b_array), - dtype, - ldb, - &beta, - reinterpret_cast(c_array), - dtype, - ldc, - batch_count, - compute_type, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); - return INFINI_STATUS_SUCCESS; - })); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t run_gemm( - const std::shared_ptr &internal, - cudaStream_t stream, - const void *a, - const void *b, - void *c, - int m, - int n, - int k, - int lda, - int ldb, - int ldc, - cudaDataType dtype) { - float alpha = 1.0f; - float beta = 0.0f; -#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) - cudaDataType compute_type = CUDA_R_32F; -#else - cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; -#endif - CHECK_STATUS(internal->useCublas( - stream, - [&](cublasHandle_t handle) { - CHECK_CUBLAS(cublasGemmEx( - handle, - CUBLAS_OP_T, - CUBLAS_OP_N, - m, - n, - k, - &alpha, - a, - dtype, - lda, - b, - dtype, - ldb, - &beta, - c, - dtype, - ldc, - compute_type, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); - return INFINI_STATUS_SUCCESS; - })); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t run_chunked_grouped_gemm( - char *base, - size_t workspace_size, - size_t intermediate_offset, - const DeepseekMoeInfo &info, - T *out, - const T *hidden, - const int *topk_indices, - const float *topk_weights, - const void *const *gate_weights, - const void *const *up_weights, - const void *const *down_weights, - cudaStream_t stream, - bool weight_ptrs_on_device, - const std::shared_ptr &internal) { - - const int hidden_size = static_cast(info.hidden_size); - const int topk = static_cast(info.topk); - const int intermediate_size = static_cast(info.intermediate_size); - const int num_experts = static_cast(info.num_experts); - const cudaDataType blas_dtype = info.dtype == INFINI_DTYPE_F16 ? CUDA_R_16F : CUDA_R_16BF; - - std::vector host_counts(info.num_experts + 1); - std::vector host_offsets(info.num_experts + 1); - std::vector host_gate(info.num_experts); - std::vector host_up(info.num_experts); - std::vector host_down(info.num_experts); - if (weight_ptrs_on_device) { - CHECK_CUDA(cudaMemcpyAsync(host_gate.data(), gate_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); - CHECK_CUDA(cudaMemcpyAsync(host_up.data(), up_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); - CHECK_CUDA(cudaMemcpyAsync(host_down.data(), down_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); - } else { - for (size_t e = 0; e < info.num_experts; ++e) { - host_gate[e] = gate_weights[e]; - host_up[e] = up_weights[e]; - host_down[e] = down_weights[e]; - } - } - - for (size_t token_begin = 0; token_begin < info.ntokens; token_begin += CHUNKED_GROUPED_GEMM_TOKENS) { - const int chunk_tokens = static_cast(std::min( - CHUNKED_GROUPED_GEMM_TOKENS, info.ntokens - token_begin)); - const int routes = chunk_tokens * topk; - - char *gptr = base + intermediate_offset; - auto *counts = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); - auto *offsets = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); - auto *positions = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); - auto *row_to_route = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * sizeof(int))); - auto *route_to_row = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * sizeof(int))); - auto *packed_hidden = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.hidden_size * sizeof(T))); - auto *gate_buf = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.intermediate_size * sizeof(T))); - auto *up_buf = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.intermediate_size * sizeof(T))); - if (static_cast(gptr - base) > workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - const T *hidden_chunk = hidden + token_begin * info.hidden_size; - const int *indices_chunk = topk_indices + token_begin * info.topk; - const float *weights_chunk = topk_weights + token_begin * info.topk; - T *out_chunk = out + token_begin * info.hidden_size; - - CHECK_CUDA(cudaMemsetAsync(counts, 0, (info.num_experts + 1) * sizeof(int), stream)); - CHECK_CUDA(cudaMemsetAsync(positions, 0, (info.num_experts + 1) * sizeof(int), stream)); - count_experts_kernel<<<(routes + 255) / 256, 256, 0, stream>>>( - indices_chunk, counts, routes, num_experts); - prefix_counts_kernel<<<1, 256, 256 * sizeof(int), stream>>>( - counts, offsets, num_experts); - pack_grouped_hidden_kernel<<>>( - hidden_chunk, indices_chunk, offsets, positions, - row_to_route, route_to_row, packed_hidden, - routes, hidden_size, topk, num_experts); - - CHECK_CUDA(cudaMemcpyAsync(host_counts.data(), counts, (info.num_experts + 1) * sizeof(int), cudaMemcpyDeviceToHost, stream)); - CHECK_CUDA(cudaMemcpyAsync(host_offsets.data(), offsets, (info.num_experts + 1) * sizeof(int), cudaMemcpyDeviceToHost, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); - - for (int e = 0; e < num_experts; ++e) { - const int count = host_counts[e]; - if (count == 0) { - continue; - } - const int offset = host_offsets[e]; - CHECK_STATUS(run_gemm( - internal, stream, host_gate[e], packed_hidden + static_cast(offset) * hidden_size, - gate_buf + static_cast(offset) * intermediate_size, - intermediate_size, count, hidden_size, - hidden_size, hidden_size, intermediate_size, blas_dtype)); - CHECK_STATUS(run_gemm( - internal, stream, host_up[e], packed_hidden + static_cast(offset) * hidden_size, - up_buf + static_cast(offset) * intermediate_size, - intermediate_size, count, hidden_size, - hidden_size, hidden_size, intermediate_size, blas_dtype)); - } - - swiglu_weight_grouped_kernel<<<(static_cast(routes) * info.intermediate_size + 255) / 256, 256, 0, stream>>>( - gate_buf, gate_buf, up_buf, row_to_route, weights_chunk, routes, intermediate_size); - - for (int e = 0; e < num_experts; ++e) { - const int count = host_counts[e]; - if (count == 0) { - continue; - } - const int offset = host_offsets[e]; - CHECK_STATUS(run_gemm( - internal, stream, host_down[e], gate_buf + static_cast(offset) * intermediate_size, - packed_hidden + static_cast(offset) * hidden_size, - hidden_size, count, intermediate_size, - intermediate_size, intermediate_size, hidden_size, blas_dtype)); - } - - sum_grouped_out_kernel<<<(static_cast(chunk_tokens) * info.hidden_size + 255) / 256, 256, 0, stream>>>( - out_chunk, packed_hidden, route_to_row, chunk_tokens, hidden_size, topk); - } - - // Keep the large-prefill path memory-bounded across transformer layers. - // All earlier chunks are synchronized by the next chunk's routing copy; - // this final sync drains the last chunk before its workspace can be reused. - CHECK_CUDA(cudaStreamSynchronize(stream)); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launch_typed( - void *workspace, - size_t workspace_size, - const DeepseekMoeInfo &info, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *const *gate_weights, - const void *const *up_weights, - const void *const *down_weights, - cudaStream_t stream, - bool weight_ptrs_on_device, - const std::shared_ptr &internal) { - - const size_t ptr_bytes = align_up(info.num_experts * sizeof(void *), 256); - const size_t ptr_workspace = ptr_bytes * 3; - const size_t intermediate_offset = align_up(ptr_workspace, 256); - size_t intermediate_tokens = info.ntokens; - if (info.ntokens > ROUTE_BATCHED_GEMM_MAX_TOKENS && info.num_experts <= 256) { - intermediate_tokens = std::min(info.ntokens, CHUNKED_GROUPED_GEMM_TOKENS); - } - const size_t intermediate_bytes = intermediate_tokens * info.topk * info.intermediate_size * sizeof(T); - if (workspace_size < intermediate_offset + intermediate_bytes) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - auto *base = reinterpret_cast(workspace); - const void *const *gate_ptrs = reinterpret_cast(base); - const void *const *up_ptrs = reinterpret_cast(base + ptr_bytes); - const void *const *down_ptrs = reinterpret_cast(base + ptr_bytes * 2); - auto *intermediate = reinterpret_cast(base + intermediate_offset); - - if (weight_ptrs_on_device) { - gate_ptrs = gate_weights; - up_ptrs = up_weights; - down_ptrs = down_weights; - } else { - auto **gate_workspace = reinterpret_cast(base); - auto **up_workspace = reinterpret_cast(base + ptr_bytes); - auto **down_workspace = reinterpret_cast(base + ptr_bytes * 2); - CHECK_CUDA(cudaMemcpyAsync(gate_workspace, gate_weights, info.num_experts * sizeof(void *), cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemcpyAsync(up_workspace, up_weights, info.num_experts * sizeof(void *), cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemcpyAsync(down_workspace, down_weights, info.num_experts * sizeof(void *), cudaMemcpyHostToDevice, stream)); - gate_ptrs = gate_workspace; - up_ptrs = up_workspace; - down_ptrs = down_workspace; - } - - if (info.ntokens > ROUTE_BATCHED_GEMM_MAX_TOKENS && info.num_experts <= 256) { - return run_chunked_grouped_gemm( - base, workspace_size, intermediate_offset, info, - reinterpret_cast(out), reinterpret_cast(hidden), - reinterpret_cast(topk_indices), reinterpret_cast(topk_weights), - gate_weights, up_weights, down_weights, - stream, weight_ptrs_on_device, internal); - } - - if (info.ntokens > FUSED_DECODE_MAX_TOKENS - && info.ntokens <= ROUTE_BATCHED_GEMM_MAX_TOKENS) { - const int routes = static_cast(info.ntokens * info.topk); - const int hidden_size_i = static_cast(info.hidden_size); - const int topk_i = static_cast(info.topk); - const int intermediate_size_i = static_cast(info.intermediate_size); - const int num_experts_i = static_cast(info.num_experts); - const cudaDataType blas_dtype = info.dtype == INFINI_DTYPE_F16 ? CUDA_R_16F : CUDA_R_16BF; - - if (info.ntokens >= EXPERT_GROUPED_GEMM_MIN_TOKENS && info.num_experts <= 256) { - char *gptr = base + intermediate_offset; - auto *counts = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); - auto *offsets = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); - auto *positions = reinterpret_cast(advance_workspace(gptr, (info.num_experts + 1) * sizeof(int))); - auto *row_to_route = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * sizeof(int))); - auto *route_to_row = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * sizeof(int))); - auto *packed_hidden = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.hidden_size * sizeof(T))); - auto *gate_buf_g = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.intermediate_size * sizeof(T))); - auto *up_buf_g = reinterpret_cast(advance_workspace(gptr, static_cast(routes) * info.intermediate_size * sizeof(T))); - if (static_cast(gptr - base) <= workspace_size) { - CHECK_CUDA(cudaMemsetAsync(counts, 0, (info.num_experts + 1) * sizeof(int), stream)); - CHECK_CUDA(cudaMemsetAsync(positions, 0, (info.num_experts + 1) * sizeof(int), stream)); - count_experts_kernel<<<(routes + 255) / 256, 256, 0, stream>>>( - reinterpret_cast(topk_indices), counts, routes, num_experts_i); - prefix_counts_kernel<<<1, 256, 256 * sizeof(int), stream>>>( - counts, offsets, num_experts_i); - pack_grouped_hidden_kernel<<>>( - reinterpret_cast(hidden), reinterpret_cast(topk_indices), - offsets, positions, row_to_route, route_to_row, packed_hidden, - routes, hidden_size_i, topk_i, num_experts_i); - - std::vector host_counts(info.num_experts + 1); - std::vector host_offsets(info.num_experts + 1); - std::vector host_gate(info.num_experts); - std::vector host_up(info.num_experts); - std::vector host_down(info.num_experts); - CHECK_CUDA(cudaMemcpyAsync(host_counts.data(), counts, (info.num_experts + 1) * sizeof(int), cudaMemcpyDeviceToHost, stream)); - CHECK_CUDA(cudaMemcpyAsync(host_offsets.data(), offsets, (info.num_experts + 1) * sizeof(int), cudaMemcpyDeviceToHost, stream)); - if (weight_ptrs_on_device) { - CHECK_CUDA(cudaMemcpyAsync(host_gate.data(), gate_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); - CHECK_CUDA(cudaMemcpyAsync(host_up.data(), up_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); - CHECK_CUDA(cudaMemcpyAsync(host_down.data(), down_weights, info.num_experts * sizeof(void *), cudaMemcpyDeviceToHost, stream)); - } else { - for (size_t e = 0; e < info.num_experts; ++e) { - host_gate[e] = gate_weights[e]; - host_up[e] = up_weights[e]; - host_down[e] = down_weights[e]; - } - } - CHECK_CUDA(cudaStreamSynchronize(stream)); - - for (int e = 0; e < num_experts_i; ++e) { - int count = host_counts[e]; - if (count == 0) { - continue; - } - int offset = host_offsets[e]; - CHECK_STATUS(run_gemm( - internal, stream, host_gate[e], packed_hidden + static_cast(offset) * hidden_size_i, - gate_buf_g + static_cast(offset) * intermediate_size_i, - intermediate_size_i, count, hidden_size_i, - hidden_size_i, hidden_size_i, intermediate_size_i, blas_dtype)); - CHECK_STATUS(run_gemm( - internal, stream, host_up[e], packed_hidden + static_cast(offset) * hidden_size_i, - up_buf_g + static_cast(offset) * intermediate_size_i, - intermediate_size_i, count, hidden_size_i, - hidden_size_i, hidden_size_i, intermediate_size_i, blas_dtype)); - } - - swiglu_weight_grouped_kernel<<<(static_cast(routes) * info.intermediate_size + 255) / 256, 256, 0, stream>>>( - gate_buf_g, gate_buf_g, up_buf_g, row_to_route, reinterpret_cast(topk_weights), routes, intermediate_size_i); - - for (int e = 0; e < num_experts_i; ++e) { - int count = host_counts[e]; - if (count == 0) { - continue; - } - int offset = host_offsets[e]; - CHECK_STATUS(run_gemm( - internal, stream, host_down[e], gate_buf_g + static_cast(offset) * intermediate_size_i, - packed_hidden + static_cast(offset) * hidden_size_i, - hidden_size_i, count, intermediate_size_i, - intermediate_size_i, intermediate_size_i, hidden_size_i, blas_dtype)); - } - - sum_grouped_out_kernel<<<(info.ntokens * info.hidden_size + 255) / 256, 256, 0, stream>>>( - reinterpret_cast(out), packed_hidden, route_to_row, static_cast(info.ntokens), hidden_size_i, topk_i); - return INFINI_STATUS_SUCCESS; - } - } - - char *bptr = base + intermediate_offset; - auto **gate_a_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); - auto **gate_b_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); - auto **gate_c_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); - auto **up_a_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); - auto **up_b_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); - auto **up_c_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); - auto **down_a_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); - auto **down_b_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); - auto **down_c_array = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * sizeof(void *), alignof(void *))); - auto *gate_buf = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * info.intermediate_size * sizeof(T))); - auto *up_buf = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * info.intermediate_size * sizeof(T))); - auto *activated = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * info.intermediate_size * sizeof(T))); - auto *route_out = reinterpret_cast(advance_workspace(bptr, static_cast(routes) * info.hidden_size * sizeof(T))); - if (static_cast(bptr - base) <= workspace_size) { - const int blocks = (routes + 255) / 256; - if (info.ntokens >= EXPERT_GROUPED_GEMM_MIN_TOKENS) { - setup_all_batched_ptrs_kernel<<>>( - gate_a_array, gate_b_array, gate_c_array, - up_a_array, up_b_array, up_c_array, - down_a_array, down_b_array, down_c_array, - reinterpret_cast(hidden), reinterpret_cast(topk_indices), - gate_ptrs, up_ptrs, down_ptrs, gate_buf, up_buf, activated, route_out, - routes, hidden_size_i, topk_i, intermediate_size_i, num_experts_i); - CHECK_STATUS(run_batched_gemm( - internal, stream, gate_a_array, gate_b_array, gate_c_array, - intermediate_size_i, 1, hidden_size_i, hidden_size_i, hidden_size_i, intermediate_size_i, - routes, blas_dtype)); - - CHECK_STATUS(run_batched_gemm( - internal, stream, up_a_array, up_b_array, up_c_array, - intermediate_size_i, 1, hidden_size_i, hidden_size_i, hidden_size_i, intermediate_size_i, - routes, blas_dtype)); - - swiglu_weight_kernel<<<(static_cast(routes) * info.intermediate_size + 255) / 256, 256, 0, stream>>>( - activated, gate_buf, up_buf, reinterpret_cast(topk_weights), routes, intermediate_size_i); - - CHECK_STATUS(run_batched_gemm( - internal, stream, down_a_array, down_b_array, down_c_array, - hidden_size_i, 1, intermediate_size_i, intermediate_size_i, intermediate_size_i, hidden_size_i, - routes, blas_dtype)); - } else { - setup_route_batched_ptrs_kernel<<>>( - gate_a_array, gate_b_array, gate_c_array, - reinterpret_cast(hidden), reinterpret_cast(topk_indices), gate_ptrs, - gate_buf, routes, hidden_size_i, topk_i, intermediate_size_i, num_experts_i); - CHECK_STATUS(run_batched_gemm( - internal, stream, gate_a_array, gate_b_array, gate_c_array, - intermediate_size_i, 1, hidden_size_i, hidden_size_i, hidden_size_i, intermediate_size_i, - routes, blas_dtype)); - - setup_route_batched_ptrs_kernel<<>>( - gate_a_array, gate_b_array, gate_c_array, - reinterpret_cast(hidden), reinterpret_cast(topk_indices), up_ptrs, - up_buf, routes, hidden_size_i, topk_i, intermediate_size_i, num_experts_i); - CHECK_STATUS(run_batched_gemm( - internal, stream, gate_a_array, gate_b_array, gate_c_array, - intermediate_size_i, 1, hidden_size_i, hidden_size_i, hidden_size_i, intermediate_size_i, - routes, blas_dtype)); - - swiglu_weight_kernel<<<(static_cast(routes) * info.intermediate_size + 255) / 256, 256, 0, stream>>>( - activated, gate_buf, up_buf, reinterpret_cast(topk_weights), routes, intermediate_size_i); - - setup_down_batched_ptrs_kernel<<>>( - gate_a_array, gate_b_array, gate_c_array, - activated, reinterpret_cast(topk_indices), down_ptrs, - route_out, routes, hidden_size_i, topk_i, intermediate_size_i, num_experts_i); - CHECK_STATUS(run_batched_gemm( - internal, stream, gate_a_array, gate_b_array, gate_c_array, - hidden_size_i, 1, intermediate_size_i, intermediate_size_i, intermediate_size_i, hidden_size_i, - routes, blas_dtype)); - } - - sum_route_out_kernel<<<(info.ntokens * info.hidden_size + 255) / 256, 256, 0, stream>>>( - reinterpret_cast(out), route_out, static_cast(info.ntokens), hidden_size_i, topk_i); - return INFINI_STATUS_SUCCESS; - } - } - - constexpr int threads = FUSED_KERNEL_THREADS; - const dim3 gate_blocks(static_cast(info.ntokens * info.topk * info.intermediate_size)); - gate_up_kernel<<>>( - intermediate, - reinterpret_cast(hidden), - reinterpret_cast(topk_indices), - reinterpret_cast(topk_weights), - gate_ptrs, - up_ptrs, - info.ntokens, - info.hidden_size, - info.topk, - info.intermediate_size, - info.num_experts); - - const dim3 down_blocks(static_cast(info.ntokens * info.hidden_size)); - down_kernel<<>>( - reinterpret_cast(out), - intermediate, - reinterpret_cast(topk_indices), - down_ptrs, - info.ntokens, - info.hidden_size, - info.topk, - info.intermediate_size, - info.num_experts); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t hidden_desc, - infiniopTensorDescriptor_t topk_indices_desc, - infiniopTensorDescriptor_t topk_weights_desc, - size_t intermediate_size, - size_t num_experts) { - - auto result = DeepseekMoeInfo::create(out_desc, hidden_desc, topk_indices_desc, topk_weights_desc, intermediate_size, num_experts); - CHECK_RESULT(result); - auto info = result.take(); - - const size_t dtype_size = info.dtype == INFINI_DTYPE_F16 ? sizeof(half) : sizeof(__nv_bfloat16); - const size_t ptr_bytes = align_up(info.num_experts * sizeof(void *), 256); - const size_t intermediate_offset = align_up(ptr_bytes * 3, 256); - size_t workspace_tokens = info.ntokens; - if (info.ntokens > ROUTE_BATCHED_GEMM_MAX_TOKENS && info.num_experts <= 256) { - workspace_tokens = std::min(info.ntokens, CHUNKED_GROUPED_GEMM_TOKENS); - } - const size_t intermediate_bytes = workspace_tokens * info.topk * info.intermediate_size * dtype_size; - const size_t old_workspace_size = intermediate_offset + intermediate_bytes; - size_t batched_workspace_size = old_workspace_size; - size_t grouped_workspace_size = old_workspace_size; - if (info.ntokens > ROUTE_BATCHED_GEMM_MAX_TOKENS && info.num_experts <= 256) { - auto chunk_info = info; - chunk_info.ntokens = workspace_tokens; - if (info.dtype == INFINI_DTYPE_F16) { - grouped_workspace_size = expert_grouped_workspace_size(chunk_info, intermediate_offset); - } else { - grouped_workspace_size = expert_grouped_workspace_size<__nv_bfloat16>(chunk_info, intermediate_offset); - } - } else if (info.ntokens > FUSED_DECODE_MAX_TOKENS - && info.ntokens <= ROUTE_BATCHED_GEMM_MAX_TOKENS) { - if (info.dtype == INFINI_DTYPE_F16) { - batched_workspace_size = route_batched_workspace_size(info, intermediate_offset); - grouped_workspace_size = expert_grouped_workspace_size(info, intermediate_offset); - } else { - batched_workspace_size = route_batched_workspace_size<__nv_bfloat16>(info, intermediate_offset); - grouped_workspace_size = expert_grouped_workspace_size<__nv_bfloat16>(info, intermediate_offset); - } - } - const size_t workspace_size = std::max(old_workspace_size, std::max(batched_workspace_size, grouped_workspace_size)); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info, - workspace_size, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *const *gate_weights, - const void *const *up_weights, - const void *const *down_weights, - void *stream_) const { - - auto stream = reinterpret_cast(stream_); - if (_info.dtype == INFINI_DTYPE_F16) { - return launch_typed(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, false, _opaque->internal); - } - if (_info.dtype == INFINI_DTYPE_BF16) { - return launch_typed<__nv_bfloat16>(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, false, _opaque->internal); - } - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -infiniStatus_t Descriptor::calculateWithDevicePtrs( - void *workspace, - size_t workspace_size, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *gate_weight_ptrs, - const void *up_weight_ptrs, - const void *down_weight_ptrs, - void *stream_) const { - - auto stream = reinterpret_cast(stream_); - auto gate_weights = reinterpret_cast(gate_weight_ptrs); - auto up_weights = reinterpret_cast(up_weight_ptrs); - auto down_weights = reinterpret_cast(down_weight_ptrs); - if (_info.dtype == INFINI_DTYPE_F16) { - return launch_typed(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, true, _opaque->internal); - } - if (_info.dtype == INFINI_DTYPE_BF16) { - return launch_typed<__nv_bfloat16>(workspace, workspace_size, _info, out, hidden, topk_indices, topk_weights, gate_weights, up_weights, down_weights, stream, true, _opaque->internal); - } - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -} // namespace op::deepseek_moe::nvidia diff --git a/src/infiniop/ops/deepseek_moe/nvidia/deepseek_moe_nvidia.cuh b/src/infiniop/ops/deepseek_moe/nvidia/deepseek_moe_nvidia.cuh deleted file mode 100644 index f8c8c4de7..000000000 --- a/src/infiniop/ops/deepseek_moe/nvidia/deepseek_moe_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef DEEPSEEK_MOE_NVIDIA_CUH -#define DEEPSEEK_MOE_NVIDIA_CUH - -#include "../deepseek_moe.h" - -DESCRIPTOR(nvidia) - -#endif diff --git a/src/infiniop/ops/deepseek_moe/operator.cc b/src/infiniop/ops/deepseek_moe/operator.cc deleted file mode 100644 index d1ddfb58a..000000000 --- a/src/infiniop/ops/deepseek_moe/operator.cc +++ /dev/null @@ -1,211 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/deepseek_moe.h" - -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) \ - || defined(ENABLE_ALI_API) -#include "nvidia/deepseek_moe_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/deepseek_moe_metax.cuh" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/deepseek_moe_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateDeepseekMoeDescriptor( - infiniopHandle_t handle, - infiniopDeepseekMoeDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t hidden_desc, - infiniopTensorDescriptor_t topk_indices_desc, - infiniopTensorDescriptor_t topk_weights_desc, - size_t intermediate_size, - size_t num_experts) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::deepseek_moe::NAMESPACE::Descriptor::create( \ - handle, reinterpret_cast(desc_ptr), \ - out_desc, hidden_desc, topk_indices_desc, topk_weights_desc, \ - intermediate_size, num_experts) - - switch (handle->device) { -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetDeepseekMoeWorkspaceSize( - infiniopDeepseekMoeDescriptor_t desc, - size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_ALI_API - GET(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopDeepseekMoe( - infiniopDeepseekMoeDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *const *gate_weights, - const void *const *up_weights, - const void *const *down_weights, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc)->calculate( \ - workspace, workspace_size, out, hidden, topk_indices, topk_weights, \ - gate_weights, up_weights, down_weights, stream) - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDeepseekMoeWithDevicePtrs( - infiniopDeepseekMoeDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *hidden, - const void *topk_indices, - const void *topk_weights, - const void *gate_weight_ptrs, - const void *up_weight_ptrs, - const void *down_weight_ptrs, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc)->calculateWithDevicePtrs( \ - workspace, workspace_size, out, hidden, topk_indices, topk_weights, \ - gate_weight_ptrs, up_weight_ptrs, down_weight_ptrs, stream) - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyDeepseekMoeDescriptor( - infiniopDeepseekMoeDescriptor_t desc) { - -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DESTROY(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_ALI_API - DESTROY(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_METAX_API - DESTROY(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DESTROY(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DESTROY -} diff --git a/src/infiniop/ops/dequant/per_tensor_dequant_int8/cuda/kernel.cuh b/src/infiniop/ops/dequant/per_tensor_dequant_int8/cuda/kernel.cuh deleted file mode 100644 index c090270c2..000000000 --- a/src/infiniop/ops/dequant/per_tensor_dequant_int8/cuda/kernel.cuh +++ /dev/null @@ -1,36 +0,0 @@ -#ifndef __PER_TENSOR_DEQUANT_INT8_KERNEL_CUH__ -#define __PER_TENSOR_DEQUANT_INT8_KERNEL_CUH__ - -template -__device__ void perTensorDequantI8SymKernel( - Tout *x, const Tin *x_packed, const float *x_scale, - size_t batch_size, size_t channel, size_t hidden_dim, size_t width, - ptrdiff_t strides_0, ptrdiff_t strides_1, ptrdiff_t strides_2, ptrdiff_t strides_3, - ptrdiff_t p_strides_0, ptrdiff_t p_strides_1, ptrdiff_t p_strides_2, ptrdiff_t p_strides_3, - int num_elements) { - - unsigned int gid = blockIdx.x * blockDim.x + threadIdx.x; - const int grid_size = blockDim.x * gridDim.x; - float x_scale_val = x_scale[0]; - for (int ind = gid; ind < num_elements; ind += grid_size) { - int tid = ind; - int w = tid % (int)width; - tid = tid / (int)width; - - int h = tid % (int)hidden_dim; - tid = tid / (int)hidden_dim; - - int c = tid % (int)channel; - tid = tid / (int)channel; - - int b = tid % (int)batch_size; - - int index = w * (int)strides_3 + h * (int)strides_2 + c * (int)strides_1 + b * (int)strides_0; - int p_index = w * (int)p_strides_3 + h * (int)p_strides_2 + c * (int)p_strides_1 + b * (int)p_strides_0; - - float val = static_cast(x_packed[p_index]) * x_scale_val; - x[index] = static_cast(val); - } -} - -#endif // __PER_TENSOR_DEQUANT_INT8_KERNEL_CUH__ diff --git a/src/infiniop/ops/dequant/per_tensor_dequant_int8/info.h b/src/infiniop/ops/dequant/per_tensor_dequant_int8/info.h deleted file mode 100644 index 4614f5f40..000000000 --- a/src/infiniop/ops/dequant/per_tensor_dequant_int8/info.h +++ /dev/null @@ -1,76 +0,0 @@ -#ifndef __PER_TENSOR_DEQUANT_INT8_INFO_H__ -#define __PER_TENSOR_DEQUANT_INT8_INFO_H__ - -#include "../../../../utils.h" -#include "../../../operator.h" -#include "../../../tensor.h" - -namespace op::per_tensor_dequant_int8 { - -class PerTensorDequantI8Info { -private: - PerTensorDequantI8Info() = default; - -public: - infiniDtype_t dtype, packed_type; - size_t batch_size, channel, hidden_dim, width; - ptrdiff_t strides_0, strides_1, strides_2, strides_3; - ptrdiff_t p_strides_0, p_strides_1, p_strides_2, p_strides_3; - int num_elements; - - static utils::Result createPerTensorDequantI8Info( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t x_packed_desc, - infiniopTensorDescriptor_t x_scale_desc, - infiniopTensorDescriptor_t x_zero_desc) { - - CHECK_OR_RETURN( - x_packed_desc != nullptr && x_scale_desc != nullptr && x_desc != nullptr, - INFINI_STATUS_NULL_POINTER); - - const infiniDtype_t dtype = x_desc->dtype(); - const infiniDtype_t packed_type = x_packed_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - CHECK_DTYPE(packed_type, INFINI_DTYPE_I8); - - auto shape = x_desc->shape(); - CHECK_SAME_SHAPE(shape, x_packed_desc->shape()); - - auto ndim = x_desc->ndim(); - CHECK_OR_RETURN(ndim <= 4, - INFINI_STATUS_BAD_TENSOR_SHAPE); - - size_t width = shape[ndim - 1]; - size_t hidden_dim = (ndim > 1 ? shape[ndim - 2] : 1); - size_t channel = (ndim > 2 ? shape[ndim - 3] : 1); - size_t batch_size = (ndim > 3 ? shape[ndim - 4] : 1); - - ptrdiff_t strides_3 = x_desc->strides()[ndim - 1]; - ptrdiff_t strides_2 = (ndim > 1 ? x_desc->strides()[ndim - 2] : 0); - ptrdiff_t strides_1 = (ndim > 2 ? x_desc->strides()[ndim - 3] : 0); - ptrdiff_t strides_0 = (ndim > 3 ? x_desc->strides()[ndim - 4] : 0); - - ptrdiff_t p_strides_3 = x_packed_desc->strides()[ndim - 1]; - ptrdiff_t p_strides_2 = (ndim > 1 ? x_packed_desc->strides()[ndim - 2] : 0); - ptrdiff_t p_strides_1 = (ndim > 2 ? x_packed_desc->strides()[ndim - 3] : 0); - ptrdiff_t p_strides_0 = (ndim > 3 ? x_packed_desc->strides()[ndim - 4] : 0); - - int num_elements = 1; - for (int i = 0; i < (int)ndim; i++) { - num_elements *= static_cast(shape[i]); - } - - return utils::Result(PerTensorDequantI8Info{ - dtype, - packed_type, - batch_size, channel, hidden_dim, width, - strides_0, strides_1, strides_2, strides_3, - p_strides_0, p_strides_1, p_strides_2, p_strides_3, - num_elements}); - } -}; - -} // namespace op::per_tensor_dequant_int8 - -#endif // __PER_TENSOR_DEQUANT_INT8_INFO_H__ diff --git a/src/infiniop/ops/dequant/per_tensor_dequant_int8/nvidia/per_tensor_dequant_int8_nvidia.cu b/src/infiniop/ops/dequant/per_tensor_dequant_int8/nvidia/per_tensor_dequant_int8_nvidia.cu deleted file mode 100644 index 3f62e1b33..000000000 --- a/src/infiniop/ops/dequant/per_tensor_dequant_int8/nvidia/per_tensor_dequant_int8_nvidia.cu +++ /dev/null @@ -1,114 +0,0 @@ -#include "../../../../devices/nvidia/nvidia_common.cuh" -#include "per_tensor_dequant_int8_nvidia.cuh" - -#include "../../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../../reduce/cuda/reduce.cuh" -#include - -#include "../cuda/kernel.cuh" - -template -INFINIOP_CUDA_KERNEL perTensorDequantI8Sym( - Tout *x, const Tin *x_packed, const float *x_scale, - size_t batch_size, size_t channel, size_t hidden_dim, size_t width, - ptrdiff_t strides_0, ptrdiff_t strides_1, ptrdiff_t strides_2, ptrdiff_t strides_3, - ptrdiff_t p_strides_0, ptrdiff_t p_strides_1, ptrdiff_t p_strides_2, ptrdiff_t p_strides_3, - int num_elements) { - perTensorDequantI8SymKernel(x, x_packed, x_scale, - batch_size, channel, hidden_dim, width, - strides_0, strides_1, strides_2, strides_3, - p_strides_0, p_strides_1, p_strides_2, p_strides_3, - num_elements); -} - -namespace op::per_tensor_dequant_int8::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, Descriptor **desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t x_packed_desc, - infiniopTensorDescriptor_t x_scale_desc, - infiniopTensorDescriptor_t x_zero_desc) { - auto info = PerTensorDequantI8Info::createPerTensorDequantI8Info(x_desc, x_packed_desc, x_scale_desc, x_zero_desc); - CHECK_RESULT(info); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t per_tensor_dequant_int8Kernel(const PerTensorDequantI8Info &info, Tdata *x, const int8_t *x_packed, const float *x_scale, const float *x_zero, cudaStream_t stream) { - int num_elements = (int)info.num_elements; - int num_blocks = (num_elements + BLOCK_SIZE - 1) / BLOCK_SIZE; - - size_t batch_size = info.batch_size; - size_t channel = info.channel; - size_t hidden_dim = info.hidden_dim; - size_t width = info.width; - - ptrdiff_t strides_0 = info.strides_0; - ptrdiff_t strides_1 = info.strides_1; - ptrdiff_t strides_2 = info.strides_2; - ptrdiff_t strides_3 = info.strides_3; - - ptrdiff_t p_strides_0 = info.p_strides_0; - ptrdiff_t p_strides_1 = info.p_strides_1; - ptrdiff_t p_strides_2 = info.p_strides_2; - ptrdiff_t p_strides_3 = info.p_strides_3; - - if (x_zero == nullptr) { - perTensorDequantI8Sym - <<>>(x, x_packed, x_scale, - batch_size, channel, hidden_dim, width, - strides_0, strides_1, strides_2, strides_3, - p_strides_0, p_strides_1, p_strides_2, p_strides_3, - num_elements); - } else { - return INFINI_STATUS_BAD_PARAM; - } - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *x, - const void *x_packed, - const void *x_scale, - const void *x_zero, - void *stream_) const { - cudaStream_t stream = (cudaStream_t)stream_; -#define DEQUANT(BLOCK_SIZE, TDATA) \ - per_tensor_dequant_int8Kernel(_info, (TDATA *)x, (const int8_t *)x_packed, (const float *)x_scale, (const float *)x_zero, stream) -#define DEQUANT_WITH_BLOCK_SIZE(BLOCK_SIZE) \ - { \ - if (_info.dtype == INFINI_DTYPE_F16) \ - return DEQUANT(BLOCK_SIZE, half); \ - else if (_info.dtype == INFINI_DTYPE_F32) \ - return DEQUANT(BLOCK_SIZE, float); \ - else if (_info.dtype == INFINI_DTYPE_BF16) \ - return DEQUANT(BLOCK_SIZE, __nv_bfloat16); \ - else \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } - if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_1024) { - DEQUANT_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_1024) - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_512) { - DEQUANT_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_512) - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) { - DEQUANT_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_4096) - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::per_tensor_dequant_int8::nvidia diff --git a/src/infiniop/ops/dequant/per_tensor_dequant_int8/nvidia/per_tensor_dequant_int8_nvidia.cuh b/src/infiniop/ops/dequant/per_tensor_dequant_int8/nvidia/per_tensor_dequant_int8_nvidia.cuh deleted file mode 100644 index 66a7e5d03..000000000 --- a/src/infiniop/ops/dequant/per_tensor_dequant_int8/nvidia/per_tensor_dequant_int8_nvidia.cuh +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __PER_TENSOR_DEQUANT_INT8_NVIDIA_API_H__ -#define __PER_TENSOR_DEQUANT_INT8_NVIDIA_API_H__ -#include "../per_tensor_dequant_int8.h" - -DESCRIPTOR(nvidia) - -#endif // __PER_TENSOR_DEQUANT_INT8_NVIDIA_API_H__ diff --git a/src/infiniop/ops/dequant/per_tensor_dequant_int8/operator.cc b/src/infiniop/ops/dequant/per_tensor_dequant_int8/operator.cc deleted file mode 100644 index 48d416847..000000000 --- a/src/infiniop/ops/dequant/per_tensor_dequant_int8/operator.cc +++ /dev/null @@ -1,102 +0,0 @@ -#include "../../../operator.h" -#include "../../../handle.h" -#include "infiniop/ops/dequant/per_tensor_dequant_int8.h" - -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) -#include "nvidia/per_tensor_dequant_int8_nvidia.cuh" -#endif - -__INFINI_C infiniStatus_t infiniopCreatePerTensorDequantI8Descriptor(infiniopHandle_t handle, - infiniopPerTensorDequantI8Descriptor_t *desc_ptr, - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t x_packed_desc, - infiniopTensorDescriptor_t x_scale_desc, - infiniopTensorDescriptor_t x_zero_desc) { -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::per_tensor_dequant_int8::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - x_desc, \ - x_packed_desc, \ - x_scale_desc, \ - x_zero_desc); - switch (handle->device) { -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia) -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetPerTensorDequantI8WorkspaceSize(infiniopPerTensorDequantI8Descriptor_t desc, size_t *size) { - switch (desc->device_type) { -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->minWorkspaceSize(); \ - return INFINI_STATUS_SUCCESS; -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopPerTensorDequantI8(infiniopPerTensorDequantI8Descriptor_t desc, - void *workspace, - size_t workspace_size, - void *x, - const void *x_packed, - const void *x_scale, - const void *x_zero, - void *stream) { -#define DEQUANT(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc)->calculate( \ - workspace, workspace_size, x, x_packed, x_scale, x_zero, stream); - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - DEQUANT(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - DEQUANT(INFINI_DEVICE_QY, nvidia) -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DEQUANT -} - -__INFINI_C infiniStatus_t infiniopDestroyPerTensorDequantI8Descriptor(infiniopPerTensorDequantI8Descriptor_t desc) { -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - DESTROY(INFINI_DEVICE_QY, nvidia) -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DESTROY -} diff --git a/src/infiniop/ops/dequant/per_tensor_dequant_int8/per_tensor_dequant_int8.h b/src/infiniop/ops/dequant/per_tensor_dequant_int8/per_tensor_dequant_int8.h deleted file mode 100644 index 1ed54a8e5..000000000 --- a/src/infiniop/ops/dequant/per_tensor_dequant_int8/per_tensor_dequant_int8.h +++ /dev/null @@ -1,40 +0,0 @@ -#ifndef __PER_TENSOR_DEQUANT_I8_H__ -#define __PER_TENSOR_DEQUANT_I8_H__ - -#include "../../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::per_tensor_dequant_int8::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - PerTensorDequantI8Info _info; \ - size_t _workspace_size; \ - \ - Descriptor(Opaque *opaque, PerTensorDequantI8Info info, \ - size_t workspace_size, \ - infiniDevice_t device_type, int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), _info(info), _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t minWorkspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t x_desc, \ - infiniopTensorDescriptor_t x_packed_desc, \ - infiniopTensorDescriptor_t x_scale_desc, \ - infiniopTensorDescriptor_t x_zero_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *x, const void *x_packed, const void *x_scale, const void *x_zero, void *stream) const; \ - }; \ - } - -#endif // __PER_TENSOR_DEQUANT_I8_H__ diff --git a/src/infiniop/ops/dequantize_awq/dequantize_awq.h b/src/infiniop/ops/dequantize_awq/dequantize_awq.h deleted file mode 100644 index 21a0a8df6..000000000 --- a/src/infiniop/ops/dequantize_awq/dequantize_awq.h +++ /dev/null @@ -1,53 +0,0 @@ -#ifndef __DEQUANTIZE_AWQ_H__ -#define __DEQUANTIZE_AWQ_H__ - -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::dequantize_awq::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - DequantizeAWQInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - size_t workspace_size_, \ - Opaque *opaque, \ - DequantizeAWQInfo info, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t qweight_desc, \ - infiniopTensorDescriptor_t scales_desc, \ - infiniopTensorDescriptor_t zeros_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *out, \ - const void *qweight, \ - const void *scales, \ - const void *zeros, \ - void *stream) const; \ - }; \ - } - -#endif //__DEQUANTIZE_AWQ_H__ diff --git a/src/infiniop/ops/dequantize_awq/iluvatar/dequantize_w42f16_iluvatar.cu b/src/infiniop/ops/dequantize_awq/iluvatar/dequantize_w42f16_iluvatar.cu deleted file mode 100644 index 9425a0af2..000000000 --- a/src/infiniop/ops/dequantize_awq/iluvatar/dequantize_w42f16_iluvatar.cu +++ /dev/null @@ -1,127 +0,0 @@ -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "dequantize_w42f16_iluvatar.cuh" -#include "dequantize_w42f16_kernel.cuh" - -#include "../dequantize_awq.h" -#include - -__global__ void __launch_bounds__(64) - dequantize_weights_awq(int *__restrict__ B, half *__restrict__ scaling_factors, - int *__restrict__ zeros, half *__restrict__ C, int G) { - // static constexpr uint32_t ZERO = 0x0; - half B_shared[32 * (128 + 8)]; - - half *B_shared_ptr2 = B_shared; - - int N = blockDim.x * gridDim.x; // 2 - int col = (blockIdx.x * blockDim.x + threadIdx.x); - int row = (blockIdx.y * blockDim.y + threadIdx.y); - int index1 = 8 * col + 8 * row * N; - half *C_ptr2 = C + index1; - - int index2 = col + row * N; - int *B_ptr2 = B + index2; - - int index3 = col + (int)(row / G) * N; - int *zeros_ptr2 = zeros + index3; - int index4 = 8 * col + (int)(row / G) * N * 8; - half *scaling_factors_ptr2 = scaling_factors + index4; - - uint32_t zeros_loaded = *(uint32_t *)(zeros_ptr2); - uint4 B_loaded_zero = dequantize_s4_to_fp16x2_awq(zeros_loaded); - uint4 B_loaded_scale = *(uint4 *)(scaling_factors_ptr2); - - uint32_t B_loaded = *(uint32_t *)B_ptr2; - uint4 B_loaded_fp16 = dequantize_s4_to_fp16x2_awq(B_loaded); - - // Reinterpret uint4 components as __half2 - __half2 *B_loaded_fp16_h2 = reinterpret_cast<__half2 *>(&B_loaded_fp16); - __half2 *B_loaded_zero_h2 = reinterpret_cast<__half2 *>(&B_loaded_zero); - __half2 *B_loaded_scale_h2 = reinterpret_cast<__half2 *>(&B_loaded_scale); - - // Replace PTX sub.f16x2 with __hsub2 for each component - B_loaded_fp16_h2[0] = __hsub2(B_loaded_fp16_h2[0], B_loaded_zero_h2[0]); - B_loaded_fp16_h2[1] = __hsub2(B_loaded_fp16_h2[1], B_loaded_zero_h2[1]); - B_loaded_fp16_h2[2] = __hsub2(B_loaded_fp16_h2[2], B_loaded_zero_h2[2]); - B_loaded_fp16_h2[3] = __hsub2(B_loaded_fp16_h2[3], B_loaded_zero_h2[3]); - - // Replace PTX fma.rn.f16x2 with __hfma2 for each component - B_loaded_fp16_h2[0] = __hfma2(B_loaded_fp16_h2[0], B_loaded_scale_h2[0], __float2half2_rn(0.0f)); - B_loaded_fp16_h2[1] = __hfma2(B_loaded_fp16_h2[1], B_loaded_scale_h2[1], __float2half2_rn(0.0f)); - B_loaded_fp16_h2[2] = __hfma2(B_loaded_fp16_h2[2], B_loaded_scale_h2[2], __float2half2_rn(0.0f)); - B_loaded_fp16_h2[3] = __hfma2(B_loaded_fp16_h2[3], B_loaded_scale_h2[3], __float2half2_rn(0.0f)); - - // Store back to shared memory - *(uint4 *)B_shared_ptr2 = B_loaded_fp16; - - for (int i = 0; i < 8; ++i) { - *(C_ptr2 + i) = B_shared[i]; - } -} - -namespace op::dequantize_awq::iluvatar { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t qweight_desc, - infiniopTensorDescriptor_t scales_desc, - infiniopTensorDescriptor_t zeros_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = DequantizeAWQInfo::create(out_desc, qweight_desc, scales_desc, zeros_desc); - - *desc_ptr = new Descriptor( - 0, - new Opaque{handle->internal()}, - result.take(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t -Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *qweight, - const void *scales, - const void *zeros, - void *stream) const { - int in_features = _info.in_features(); - int out_features = _info.out_features(); - int group_size = in_features / _info.num_groups(); - - // ==================== 默认配置, 固定为 8 ==================== - constexpr int BLOCK_X = 8; - constexpr int BLOCK_Y = 8; - - int x_blocks = (out_features + BLOCK_X - 1) / BLOCK_X; - int y_blocks = (in_features + BLOCK_Y - 1) / BLOCK_Y; - - dim3 num_blocks(x_blocks, y_blocks); - dim3 threads_per_block(BLOCK_X, BLOCK_Y); - // ===================================================== - - half *out_ = reinterpret_cast(out); - - int *qweight_ = const_cast(reinterpret_cast(qweight)); - half *scales_ = const_cast(reinterpret_cast(scales)); - int *zeros_ = const_cast(reinterpret_cast(zeros)); - - dequantize_weights_awq<<(stream)>>>( - qweight_, scales_, zeros_, out_, group_size); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dequantize_awq::iluvatar diff --git a/src/infiniop/ops/dequantize_awq/iluvatar/dequantize_w42f16_iluvatar.cuh b/src/infiniop/ops/dequantize_awq/iluvatar/dequantize_w42f16_iluvatar.cuh deleted file mode 100644 index 253b7ed1e..000000000 --- a/src/infiniop/ops/dequantize_awq/iluvatar/dequantize_w42f16_iluvatar.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __DEQUANTIZE_AWQ_ILUVATAR_CUH__ -#define __DEQUANTIZE_AWQ_ILUVATAR_CUH__ - -#include "../dequantize_awq.h" - -DESCRIPTOR(iluvatar) - -#endif // __DEQUANTIZE_AWQ_ILUVATAR_CUH__ diff --git a/src/infiniop/ops/dequantize_awq/iluvatar/dequantize_w42f16_kernel.cuh b/src/infiniop/ops/dequantize_awq/iluvatar/dequantize_w42f16_kernel.cuh deleted file mode 100644 index 7acdfe122..000000000 --- a/src/infiniop/ops/dequantize_awq/iluvatar/dequantize_w42f16_kernel.cuh +++ /dev/null @@ -1,56 +0,0 @@ -#pragma once -#include // 需要此头文件来支持 __half 和 __half2 类型 - -/** - * @brief 将一个包含8个4-bit整数的uint32_t反量化为8个half精度浮点数。 - * - * 这是一个通用的 CUDA C++ 实现，用于替代原有的 PTX 汇编版本， - * 以便在不支持高级 PTX 指令（如 lop3.b32）的 GPU 上运行。 - * 输出顺序匹配 PTX 的交错打包：v0, v4, v1, v5, v2, v6, v3, v7（经 signed 调整后）。 - * - * @param source 输入的32位无符号整数，它打包了8个4-bit的数据。 - * @return 一个 uint4 变量，其中包含8个反量化后的 half 值。 - */ -__device__ __forceinline__ uint4 dequantize_s4_to_fp16x2_awq(uint32_t const &source) { - // 步骤 1: 从一个 32-bit 源数据中解包出 8 个 4-bit 无符号整数。 - // 源数据的内存布局被假定为 [v7, v6, v5, v4, v3, v2, v1, v0]， - // 其中每个 'v' 都是一个 4-bit 的半字节 (nibble)。 - const unsigned int v0 = (source >> 0) & 0x0F; - const unsigned int v1 = (source >> 4) & 0x0F; - const unsigned int v2 = (source >> 8) & 0x0F; - const unsigned int v3 = (source >> 12) & 0x0F; - const unsigned int v4 = (source >> 16) & 0x0F; - const unsigned int v5 = (source >> 20) & 0x0F; - const unsigned int v6 = (source >> 24) & 0x0F; - const unsigned int v7 = (source >> 28) & 0x0F; - - // 步骤 2: 对于 signed 4-bit (s4)，减去 8 以映射到 [-8, 7] 范围。 - // 定义偏移量 - __half offset = __half(8); - - // 计算 signed 值 - __half hv0 = __half(v0) - offset; - __half hv1 = __half(v1) - offset; - __half hv2 = __half(v2) - offset; - __half hv3 = __half(v3) - offset; - __half hv4 = __half(v4) - offset; - __half hv5 = __half(v5) - offset; - __half hv6 = __half(v6) - offset; - __half hv7 = __half(v7) - offset; - - // 步骤 3: 将 half 值按 PTX 交错顺序打包成 __half2 并存入 result 中。 - // 顺序：result_ptr[0]: low=hv0, high=hv4 - // result_ptr[1]: low=hv1, high=hv5 - // result_ptr[2]: low=hv2, high=hv6 - // result_ptr[3]: low=hv3, high=hv7 - // __halves2half2 函数：low 为第一个参数，high 为第二个参数。 - uint4 result; - __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); - - result_ptr[0] = __halves2half2(hv0, hv4); - result_ptr[1] = __halves2half2(hv1, hv5); - result_ptr[2] = __halves2half2(hv2, hv6); - result_ptr[3] = __halves2half2(hv3, hv7); - - return result; -} diff --git a/src/infiniop/ops/dequantize_awq/info.h b/src/infiniop/ops/dequantize_awq/info.h deleted file mode 100644 index b7770a963..000000000 --- a/src/infiniop/ops/dequantize_awq/info.h +++ /dev/null @@ -1,39 +0,0 @@ -#ifndef __DEQUANTIZE_AWQ_INFO_H__ -#define __DEQUANTIZE_AWQ_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::dequantize_awq { - -class DequantizeAWQInfo { - DequantizeAWQInfo() = default; - -public: - int _in_features, _out_features, _num_groups; - - int in_features() const { return _in_features; } - int out_features() const { return _out_features; } - int num_groups() const { return _num_groups; } - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t qweight_desc, - infiniopTensorDescriptor_t scales_desc, - infiniopTensorDescriptor_t zeros_desc) { - - int _in_features = qweight_desc->dim(0); - int _out_features = qweight_desc->dim(1); - int _num_groups = scales_desc->dim(0); - - return utils::Result(DequantizeAWQInfo{ - _in_features, - _out_features, - _num_groups}); - } -}; - -} // namespace op::dequantize_awq - -#endif // __DEQUANTIZE_AWQ_INFO_H__ diff --git a/src/infiniop/ops/dequantize_awq/moore/dequantize_w42f16_kernel.h b/src/infiniop/ops/dequantize_awq/moore/dequantize_w42f16_kernel.h deleted file mode 100644 index 5b926e892..000000000 --- a/src/infiniop/ops/dequantize_awq/moore/dequantize_w42f16_kernel.h +++ /dev/null @@ -1,56 +0,0 @@ -#pragma once -#include // 需要此头文件来支持 __half 和 __half2 类型 - -/** - * @brief 将一个包含8个4-bit整数的uint32_t反量化为8个half精度浮点数。 - * - * 这是一个通用的 CUDA C++ 实现，用于替代原有的 PTX 汇编版本， - * 以便在不支持高级 PTX 指令（如 lop3.b32）的 GPU 上运行。 - * 输出顺序匹配 PTX 的交错打包：v0, v4, v1, v5, v2, v6, v3, v7（经 signed 调整后）。 - * - * @param source 输入的32位无符号整数，它打包了8个4-bit的数据。 - * @return 一个 uint4 变量，其中包含8个反量化后的 half 值。 - */ -__device__ __forceinline__ uint4 dequantize_s4_to_fp16x2_awq(uint32_t const &source) { - // 步骤 1: 从一个 32-bit 源数据中解包出 8 个 4-bit 无符号整数。 - // 源数据的内存布局被假定为 [v7, v6, v5, v4, v3, v2, v1, v0]， - // 其中每个 'v' 都是一个 4-bit 的半字节 (nibble)。 - const unsigned int v0 = (source >> 0) & 0x0F; - const unsigned int v1 = (source >> 4) & 0x0F; - const unsigned int v2 = (source >> 8) & 0x0F; - const unsigned int v3 = (source >> 12) & 0x0F; - const unsigned int v4 = (source >> 16) & 0x0F; - const unsigned int v5 = (source >> 20) & 0x0F; - const unsigned int v6 = (source >> 24) & 0x0F; - const unsigned int v7 = (source >> 28) & 0x0F; - - // 步骤 2: 对于 signed 4-bit (s4)，减去 8 以映射到 [-8, 7] 范围。 - // 定义偏移量 - __half offset = __half(8); - - // 计算 signed 值 - __half hv0 = __half(v0) - offset; - __half hv1 = __half(v1) - offset; - __half hv2 = __half(v2) - offset; - __half hv3 = __half(v3) - offset; - __half hv4 = __half(v4) - offset; - __half hv5 = __half(v5) - offset; - __half hv6 = __half(v6) - offset; - __half hv7 = __half(v7) - offset; - - // 步骤 3: 将 half 值按 PTX 交错顺序打包成 __half2 并存入 result 中。 - // 顺序：result_ptr[0]: low=hv0, high=hv4 - // result_ptr[1]: low=hv1, high=hv5 - // result_ptr[2]: low=hv2, high=hv6 - // result_ptr[3]: low=hv3, high=hv7 - // __halves2half2 函数：low 为第一个参数，high 为第二个参数。 - uint4 result; - __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); - - result_ptr[0] = __halves2half2(hv0, hv4); - result_ptr[1] = __halves2half2(hv1, hv5); - result_ptr[2] = __halves2half2(hv2, hv6); - result_ptr[3] = __halves2half2(hv3, hv7); - - return result; -} diff --git a/src/infiniop/ops/dequantize_awq/moore/dequantize_w42f16_moore.h b/src/infiniop/ops/dequantize_awq/moore/dequantize_w42f16_moore.h deleted file mode 100644 index 8195b9781..000000000 --- a/src/infiniop/ops/dequantize_awq/moore/dequantize_w42f16_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __DEQUANTIZE_AWQ_MOORE_H__ -#define __DEQUANTIZE_AWQ_MOORE_H__ - -#include "../dequantize_awq.h" - -DESCRIPTOR(moore) - -#endif // __DEQUANTIZE_AWQ_MOORE_H__ diff --git a/src/infiniop/ops/dequantize_awq/moore/dequantize_w42f16_moore.mu b/src/infiniop/ops/dequantize_awq/moore/dequantize_w42f16_moore.mu deleted file mode 100644 index 970a58cd4..000000000 --- a/src/infiniop/ops/dequantize_awq/moore/dequantize_w42f16_moore.mu +++ /dev/null @@ -1,127 +0,0 @@ -#include "../../../devices/moore/moore_handle.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "dequantize_w42f16_moore.h" -#include "dequantize_w42f16_kernel.h" - -#include "../dequantize_awq.h" -#include - -__global__ void __launch_bounds__(64) - dequantize_weights_awq(int *__restrict__ B, half *__restrict__ scaling_factors, - int *__restrict__ zeros, half *__restrict__ C, int G) { - // static constexpr uint32_t ZERO = 0x0; - half B_shared[32 * (128 + 8)]; - - half *B_shared_ptr2 = B_shared; - - int N = blockDim.x * gridDim.x; // 2 - int col = (blockIdx.x * blockDim.x + threadIdx.x); - int row = (blockIdx.y * blockDim.y + threadIdx.y); - int index1 = 8 * col + 8 * row * N; - half *C_ptr2 = C + index1; - - int index2 = col + row * N; - int *B_ptr2 = B + index2; - - int index3 = col + (int)(row / G) * N; - int *zeros_ptr2 = zeros + index3; - int index4 = 8 * col + (int)(row / G) * N * 8; - half *scaling_factors_ptr2 = scaling_factors + index4; - - uint32_t zeros_loaded = *(uint32_t *)(zeros_ptr2); - uint4 B_loaded_zero = dequantize_s4_to_fp16x2_awq(zeros_loaded); - uint4 B_loaded_scale = *(uint4 *)(scaling_factors_ptr2); - - uint32_t B_loaded = *(uint32_t *)B_ptr2; - uint4 B_loaded_fp16 = dequantize_s4_to_fp16x2_awq(B_loaded); - - // Reinterpret uint4 components as __half2 - __half2 *B_loaded_fp16_h2 = reinterpret_cast<__half2 *>(&B_loaded_fp16); - __half2 *B_loaded_zero_h2 = reinterpret_cast<__half2 *>(&B_loaded_zero); - __half2 *B_loaded_scale_h2 = reinterpret_cast<__half2 *>(&B_loaded_scale); - - // Replace PTX sub.f16x2 with __hsub2 for each component - B_loaded_fp16_h2[0] = __hsub2(B_loaded_fp16_h2[0], B_loaded_zero_h2[0]); - B_loaded_fp16_h2[1] = __hsub2(B_loaded_fp16_h2[1], B_loaded_zero_h2[1]); - B_loaded_fp16_h2[2] = __hsub2(B_loaded_fp16_h2[2], B_loaded_zero_h2[2]); - B_loaded_fp16_h2[3] = __hsub2(B_loaded_fp16_h2[3], B_loaded_zero_h2[3]); - - // Replace PTX fma.rn.f16x2 with __hfma2 for each component - B_loaded_fp16_h2[0] = __hfma2(B_loaded_fp16_h2[0], B_loaded_scale_h2[0], __float2half2_rn(0.0f)); - B_loaded_fp16_h2[1] = __hfma2(B_loaded_fp16_h2[1], B_loaded_scale_h2[1], __float2half2_rn(0.0f)); - B_loaded_fp16_h2[2] = __hfma2(B_loaded_fp16_h2[2], B_loaded_scale_h2[2], __float2half2_rn(0.0f)); - B_loaded_fp16_h2[3] = __hfma2(B_loaded_fp16_h2[3], B_loaded_scale_h2[3], __float2half2_rn(0.0f)); - - // Store back to shared memory - *(uint4 *)B_shared_ptr2 = B_loaded_fp16; - - for (int i = 0; i < 8; ++i) { - *(C_ptr2 + i) = B_shared[i]; - } -} - -namespace op::dequantize_awq::moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t qweight_desc, - infiniopTensorDescriptor_t scales_desc, - infiniopTensorDescriptor_t zeros_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = DequantizeAWQInfo::create(out_desc, qweight_desc, scales_desc, zeros_desc); - - *desc_ptr = new Descriptor( - 0, - new Opaque{handle->internal()}, - result.take(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t -Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *qweight, - const void *scales, - const void *zeros, - void *stream) const { - int in_features = _info.in_features(); - int out_features = _info.out_features(); - int group_size = in_features / _info.num_groups(); - - // ==================== 默认配置, 固定为 8 ==================== - constexpr int BLOCK_X = 8; - constexpr int BLOCK_Y = 8; - - int x_blocks = (out_features + BLOCK_X - 1) / BLOCK_X; - int y_blocks = (in_features + BLOCK_Y - 1) / BLOCK_Y; - - dim3 num_blocks(x_blocks, y_blocks); - dim3 threads_per_block(BLOCK_X, BLOCK_Y); - // ===================================================== - - half *out_ = reinterpret_cast(out); - - int *qweight_ = const_cast(reinterpret_cast(qweight)); - half *scales_ = const_cast(reinterpret_cast(scales)); - int *zeros_ = const_cast(reinterpret_cast(zeros)); - - dequantize_weights_awq<<(stream)>>>( - qweight_, scales_, zeros_, out_, group_size); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dequantize_awq::moore diff --git a/src/infiniop/ops/dequantize_awq/nvidia/dequantize_w42f16_kernel.cuh b/src/infiniop/ops/dequantize_awq/nvidia/dequantize_w42f16_kernel.cuh deleted file mode 100644 index 8b397a5b0..000000000 --- a/src/infiniop/ops/dequantize_awq/nvidia/dequantize_w42f16_kernel.cuh +++ /dev/null @@ -1,125 +0,0 @@ -#pragma once - -__device__ uint4 dequantize_s4_to_fp16x2_awq(uint32_t const &source) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750 - // 步骤 1: 从一个 32-bit 源数据中解包出 8 个 4-bit 无符号整数。 - // 源数据的内存布局被假定为 [v7, v6, v5, v4, v3, v2, v1, v0]， - // 其中每个 'v' 都是一个 4-bit 的半字节 (nibble)。 - const unsigned int v0 = (source >> 0) & 0x0F; - const unsigned int v1 = (source >> 4) & 0x0F; - const unsigned int v2 = (source >> 8) & 0x0F; - const unsigned int v3 = (source >> 12) & 0x0F; - const unsigned int v4 = (source >> 16) & 0x0F; - const unsigned int v5 = (source >> 20) & 0x0F; - const unsigned int v6 = (source >> 24) & 0x0F; - const unsigned int v7 = (source >> 28) & 0x0F; - - // 步骤 2: 对于 signed 4-bit (s4)，减去 8 以映射到 [-8, 7] 范围。 - // 定义偏移量 - __half offset = __half(8); - - // 计算 signed 值 - __half hv0 = __half(v0) - offset; - __half hv1 = __half(v1) - offset; - __half hv2 = __half(v2) - offset; - __half hv3 = __half(v3) - offset; - __half hv4 = __half(v4) - offset; - __half hv5 = __half(v5) - offset; - __half hv6 = __half(v6) - offset; - __half hv7 = __half(v7) - offset; - - // 步骤 3: 将 half 值按 PTX 交错顺序打包成 __half2 并存入 result 中。 - // 顺序：result_ptr[0]: low=hv0, high=hv4 - // result_ptr[1]: low=hv1, high=hv5 - // result_ptr[2]: low=hv2, high=hv6 - // result_ptr[3]: low=hv3, high=hv7 - // __halves2half2 函数：low 为第一个参数，high 为第二个参数。 - uint4 result; - __half2 *result_ptr = reinterpret_cast<__half2 *>(&result); - - result_ptr[0] = __halves2half2(hv0, hv4); - result_ptr[1] = __halves2half2(hv1, hv5); - result_ptr[2] = __halves2half2(hv2, hv6); - result_ptr[3] = __halves2half2(hv3, hv7); - - return result; -#else - uint4 result; - - uint32_t *h = reinterpret_cast(&result); - uint32_t const i4s = reinterpret_cast(source); - - // First, we extract the i4s and construct an intermediate fp16 number. - static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; - static constexpr uint32_t BOTTOM_MASK = 0x000f000f; - static constexpr uint32_t TOP_MASK = 0x00f000f0; - static constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400; - - // Note that the entire sequence only requires 1 shift instruction. This is - // thanks to the register packing format and the fact that we force our - // integers to be unsigned, and account for this in the fp16 subtractions. In - // addition, I exploit the fact that sub and fma have the same throughput in - // order to convert elt_23 and elt_67 to fp16 without having to shift them to - // the bottom bits before hand. - - // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW - // dependency if we issue immediately before required. - const uint32_t top_i4s = i4s >> 8; - // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400 - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[0]) - : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), - "n"(immLut)); - // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400 - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[1]) - : "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), - "n"(immLut)); - // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400 - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[2]) - : "r"(top_i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), - "n"(immLut)); - // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400 - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(h[3]) - : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), - "n"(immLut)); - - // I use inline PTX below because I am not sure if the compiler will emit - // float2half instructions if I use the half2 ctor. In this case, I chose - // performance reliability over code readability. - - // This is the half2 {1032, 1032} represented as an integer. - // static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64086408; - // Haotian: subtract {1024, 1024} instead, we do not need to map to [-8, 7] - static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64006400; - // This is the half2 {1 / 16, 1 / 16} represented as an integer. - static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00; - // This is the half2 {-72, -72} represented as an integer. - // static constexpr uint32_t NEG_72 = 0xd480d480; - // Haotian: Let's use {-64, -64}. - static constexpr uint32_t NEG_64 = 0xd400d400; - - // Finally, we construct the output numbers. - // Convert elt_01 - asm volatile("sub.f16x2 %0, %1, %2;\n" - : "=r"(h[0]) - : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM)); - // Convert elt_23 - asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" - : "=r"(h[1]) - : "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); - // Convert elt_45 - asm volatile("sub.f16x2 %0, %1, %2;\n" - : "=r"(h[2]) - : "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM)); - // Convert elt_67 - asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" - : "=r"(h[3]) - : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); - - return result; -#endif - __builtin_unreachable(); // Suppress missing return statement warning -} diff --git a/src/infiniop/ops/dequantize_awq/nvidia/dequantize_w42f16_nvidia.cu b/src/infiniop/ops/dequantize_awq/nvidia/dequantize_w42f16_nvidia.cu deleted file mode 100644 index 0c5251142..000000000 --- a/src/infiniop/ops/dequantize_awq/nvidia/dequantize_w42f16_nvidia.cu +++ /dev/null @@ -1,203 +0,0 @@ -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) - -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "dequantize_w42f16_kernel.cuh" -#include "dequantize_w42f16_nvidia.cuh" - -#include "../dequantize_awq.h" -#include - -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 750) -__global__ void __launch_bounds__(64) - dequantize_weights_awq(int *__restrict__ B, half *__restrict__ scaling_factors, - int *__restrict__ zeros, half *__restrict__ C, int G, - int out_features, int in_features) { - // static constexpr uint32_t ZERO = 0x0; - - int col = (blockIdx.x * blockDim.x + threadIdx.x); - int row = (blockIdx.y * blockDim.y + threadIdx.y); - - // 边界检查，防止越界访问 - if (col >= out_features || row >= in_features) { - return; - } - - // 每个元素在输出中的起始地址：行主序，连续 8 个 half - int index1 = 8 * col + 8 * row * out_features; - half *C_ptr2 = C + index1; - - int index2 = col + row * out_features; - int *B_ptr2 = B + index2; - - int index3 = col + (int)(row / G) * out_features; - int *zeros_ptr2 = zeros + index3; - - int index4 = 8 * col + (int)(row / G) * out_features * 8; - half *scaling_factors_ptr2 = scaling_factors + index4; - - uint32_t zeros_loaded = *(uint32_t *)(zeros_ptr2); - uint4 B_loaded_zero = dequantize_s4_to_fp16x2_awq(zeros_loaded); - uint4 B_loaded_scale = *(uint4 *)(scaling_factors_ptr2); - - uint32_t B_loaded = *(uint32_t *)B_ptr2; - uint4 B_loaded_fp16 = dequantize_s4_to_fp16x2_awq(B_loaded); - - // Reinterpret uint4 components as __half2 - __half2 *B_loaded_fp16_h2 = reinterpret_cast<__half2 *>(&B_loaded_fp16); - __half2 *B_loaded_zero_h2 = reinterpret_cast<__half2 *>(&B_loaded_zero); - __half2 *B_loaded_scale_h2 = reinterpret_cast<__half2 *>(&B_loaded_scale); - - // Replace PTX sub.f16x2 with __hsub2 for each component - B_loaded_fp16_h2[0] = __hsub2(B_loaded_fp16_h2[0], B_loaded_zero_h2[0]); - B_loaded_fp16_h2[1] = __hsub2(B_loaded_fp16_h2[1], B_loaded_zero_h2[1]); - B_loaded_fp16_h2[2] = __hsub2(B_loaded_fp16_h2[2], B_loaded_zero_h2[2]); - B_loaded_fp16_h2[3] = __hsub2(B_loaded_fp16_h2[3], B_loaded_zero_h2[3]); - - // Replace PTX fma.rn.f16x2 with __hfma2 for each component - B_loaded_fp16_h2[0] = __hfma2(B_loaded_fp16_h2[0], B_loaded_scale_h2[0], __float2half2_rn(0.0f)); - B_loaded_fp16_h2[1] = __hfma2(B_loaded_fp16_h2[1], B_loaded_scale_h2[1], __float2half2_rn(0.0f)); - B_loaded_fp16_h2[2] = __hfma2(B_loaded_fp16_h2[2], B_loaded_scale_h2[2], __float2half2_rn(0.0f)); - B_loaded_fp16_h2[3] = __hfma2(B_loaded_fp16_h2[3], B_loaded_scale_h2[3], __float2half2_rn(0.0f)); - - // 直接写回全局内存输出 - half *out_vec = reinterpret_cast(&B_loaded_fp16); -#pragma unroll - for (int i = 0; i < 8; ++i) { - C_ptr2[i] = out_vec[i]; - } -} -#else -__global__ void __launch_bounds__(64) - dequantize_weights_awq(int *__restrict__ B, half *__restrict__ scaling_factors, - int *__restrict__ zeros, half *__restrict__ C, int group_size, - int out_features, int in_features) { - static constexpr uint32_t ZERO = 0x0; - - int col = (blockIdx.x * blockDim.x + threadIdx.x); - int row = blockIdx.y * blockDim.y + threadIdx.y; - - // 边界检查，防止越界访问 - if (col >= out_features || row >= in_features) { - return; - } - - int index1 = 8 * col + 8 * row * out_features; - half *C_ptr2 = C + index1; - - int index2 = col + row * out_features; - int *B_ptr2 = B + index2; - - int index3 = col + (int)(row / group_size) * out_features; - int *zeros_ptr2 = zeros + index3; - int index4 = 8 * col + (int)(row / group_size) * out_features * 8; - half *scaling_factors_ptr2 = scaling_factors + index4; - - uint32_t zeros_loaded = *(uint32_t *)(zeros_ptr2); - uint4 B_loaded_zero = dequantize_s4_to_fp16x2_awq(zeros_loaded); - uint4 B_loaded_scale = *(uint4 *)(scaling_factors_ptr2); - - uint32_t B_loaded = *(uint32_t *)B_ptr2; - uint4 B_loaded_fp16 = dequantize_s4_to_fp16x2_awq(B_loaded); - asm volatile("sub.f16x2 %0, %1, %2;\n" - : "=r"(B_loaded_fp16.x) - : "r"(B_loaded_fp16.x), "r"(B_loaded_zero.x)); - asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" - : "=r"(B_loaded_fp16.x) - : "r"(B_loaded_fp16.x), "r"(B_loaded_scale.x), "r"(ZERO)); - asm volatile("sub.f16x2 %0, %1, %2;\n" - : "=r"(B_loaded_fp16.y) - : "r"(B_loaded_fp16.y), "r"(B_loaded_zero.y)); - asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" - : "=r"(B_loaded_fp16.y) - : "r"(B_loaded_fp16.y), "r"(B_loaded_scale.y), "r"(ZERO)); - asm volatile("sub.f16x2 %0, %1, %2;\n" - : "=r"(B_loaded_fp16.z) - : "r"(B_loaded_fp16.z), "r"(B_loaded_zero.z)); - asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" - : "=r"(B_loaded_fp16.z) - : "r"(B_loaded_fp16.z), "r"(B_loaded_scale.z), "r"(ZERO)); - asm volatile("sub.f16x2 %0, %1, %2;\n" - : "=r"(B_loaded_fp16.w) - : "r"(B_loaded_fp16.w), "r"(B_loaded_zero.w)); - asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" - : "=r"(B_loaded_fp16.w) - : "r"(B_loaded_fp16.w), "r"(B_loaded_scale.w), "r"(ZERO)); - - // 直接写回全局内存输出 - half *out_vec = reinterpret_cast(&B_loaded_fp16); -#pragma unroll - for (int i = 0; i < 8; ++i) { - C_ptr2[i] = out_vec[i]; - } -} -#endif - -namespace op::dequantize_awq::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t qweight_desc, - infiniopTensorDescriptor_t scales_desc, - infiniopTensorDescriptor_t zeros_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = DequantizeAWQInfo::create(out_desc, qweight_desc, scales_desc, zeros_desc); - - *desc_ptr = new Descriptor( - 0, - new Opaque{handle->internal()}, - result.take(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t -Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *qweight, - const void *scales, - const void *zeros, - void *stream) const { - int in_features = _info.in_features(); - int out_features = _info.out_features(); - int group_size = in_features / _info.num_groups(); - - // ==================== 默认配置, 固定为 8 ==================== - constexpr int BLOCK_X = 8; - constexpr int BLOCK_Y = 8; - - int x_blocks = (out_features + BLOCK_X - 1) / BLOCK_X; - int y_blocks = (in_features + BLOCK_Y - 1) / BLOCK_Y; - - dim3 num_blocks(x_blocks, y_blocks); - dim3 threads_per_block(BLOCK_X, BLOCK_Y); - // ===================================================== - - half *out_ = reinterpret_cast(out); - - int *qweight_ = const_cast(reinterpret_cast(qweight)); - half *scales_ = const_cast(reinterpret_cast(scales)); - int *zeros_ = const_cast(reinterpret_cast(zeros)); - - dequantize_weights_awq<<(stream)>>>( - qweight_, scales_, zeros_, out_, group_size, out_features, in_features); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dequantize_awq::nvidia - -#endif diff --git a/src/infiniop/ops/dequantize_awq/nvidia/dequantize_w42f16_nvidia.cuh b/src/infiniop/ops/dequantize_awq/nvidia/dequantize_w42f16_nvidia.cuh deleted file mode 100644 index 2593c03f2..000000000 --- a/src/infiniop/ops/dequantize_awq/nvidia/dequantize_w42f16_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __DEQUANTIZE_AWQ_CUDA_CUH__ -#define __DEQUANTIZE_AWQ_CUDA_CUH__ - -#include "../dequantize_awq.h" - -DESCRIPTOR(nvidia) - -#endif // __DEQUANTIZE_AWQ_CUDA_CUH__ diff --git a/src/infiniop/ops/dequantize_awq/operator.cc b/src/infiniop/ops/dequantize_awq/operator.cc deleted file mode 100644 index c63aed0fd..000000000 --- a/src/infiniop/ops/dequantize_awq/operator.cc +++ /dev/null @@ -1,158 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/dequantize_awq.h" - -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) -#include "nvidia/dequantize_w42f16_nvidia.cuh" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/dequantize_w42f16_moore.h" -#endif -#ifdef ENABLE_ILUVATAR_API -#include "iluvatar/dequantize_w42f16_iluvatar.cuh" -#endif - -__INFINI_C infiniStatus_t infiniopCreateDequantizeAWQDescriptor( - infiniopHandle_t handle, - infiniopDequantizeAWQDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t qweight_desc, - infiniopTensorDescriptor_t scales_desc, - infiniopTensorDescriptor_t zeros_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::dequantize_awq::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - out_desc, \ - qweight_desc, \ - scales_desc, \ - zeros_desc) - - switch (handle->device) { -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, iluvatar); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetDequantizeAWQWorkspaceSize(infiniopDequantizeAWQDescriptor_t desc, - size_t *size) { -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, iluvatar); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ALI_API - GET(INFINI_DEVICE_ALI, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopDequantizeAWQ( - infiniopDequantizeAWQDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *qweight, - const void *scales, - const void *zeros, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, out, qweight, scales, zeros, stream) - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, iluvatar); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyDequantizeAWQDescriptor(infiniopDequantizeAWQDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, iluvatar); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ALI_API - DELETE(INFINI_DEVICE_ALI, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} - -// #endif diff --git a/src/infiniop/ops/dequantize_gptq/dequantize_gptq.h b/src/infiniop/ops/dequantize_gptq/dequantize_gptq.h deleted file mode 100644 index 27092a3b3..000000000 --- a/src/infiniop/ops/dequantize_gptq/dequantize_gptq.h +++ /dev/null @@ -1,55 +0,0 @@ -#ifndef __DEQUANTIZE_GPTQ_H__ -#define __DEQUANTIZE_GPTQ_H__ - -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::dequantize_gptq::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - DequantizeGPTQInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - size_t workspace_size_, \ - Opaque *opaque, \ - DequantizeGPTQInfo info, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t qweight_desc, \ - infiniopTensorDescriptor_t scales_desc, \ - infiniopTensorDescriptor_t zeros_desc, \ - infiniopTensorDescriptor_t g_idx_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *out, \ - const void *qweight, \ - const void *scales, \ - const void *zeros, \ - const void *g_idx, \ - void *stream) const; \ - }; \ - } - -#endif //__DEQUANTIZE_GPTQ_H__ diff --git a/src/infiniop/ops/dequantize_gptq/iluvatar/dequantize_w42f16_iluvatar.cu b/src/infiniop/ops/dequantize_gptq/iluvatar/dequantize_w42f16_iluvatar.cu deleted file mode 100644 index a853394bf..000000000 --- a/src/infiniop/ops/dequantize_gptq/iluvatar/dequantize_w42f16_iluvatar.cu +++ /dev/null @@ -1,129 +0,0 @@ -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "dequantize_w42f16_iluvatar.cuh" - -#include "../dequantize_gptq.h" -#include -#include - -namespace op::dequantize_gptq::iluvatar { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -// 对齐 nvidia 版：支持 g_idx -// qweight: [in_packed, out_features] packing 8 input channels per word -// zeros: [num_groups, out_packed] packing 8 output channels per word -// scales: [num_groups, out_features], g_idx: [in_features] -__global__ void __launch_bounds__(128) - dequantize_weights_gptq(const uint32_t *__restrict__ qweight, - const half *__restrict__ scales, - const uint32_t *__restrict__ zeros, - const int *__restrict__ g_idx, - half *__restrict__ out, - int in_features, - int out_features, - int out_packed, // ceil(out_features / 8) - int num_groups) { - const int col_pack = blockIdx.x * blockDim.x + threadIdx.x; // packed output column - const int row = blockIdx.y * blockDim.y + threadIdx.y; // real input row - if (col_pack >= out_packed || row >= in_features) { - return; - } - - const int gid_raw = g_idx ? g_idx[row] : 0; - const int gid = ((gid_raw % num_groups) + num_groups) % num_groups; - - const int pack_row = row >> 3; // packed input row (8 rows per pack) - const int q_shift = (row & 7) * 4; // nibble shift within uint32 - - const uint32_t zeros_loaded = zeros[gid * out_packed + col_pack]; - - const int col_base = col_pack << 3; // 8 real cols per pack - const int scale_base = gid * out_features + col_base; - -#pragma unroll - for (int j = 0; j < 8; ++j) { - const int col = col_base + j; - if (col >= out_features) { - break; - } - - const uint32_t q_loaded = qweight[pack_row * out_features + col]; - const int q_nib = (q_loaded >> q_shift) & 0xF; - - const int z_nib = (zeros_loaded >> (j * 4)) & 0xF; - const half scale = scales[scale_base + j]; - - // aligned with nvidia: (q - (z + 1)) * s - const float v = float(q_nib - (z_nib + 1)) * __half2float(scale); - out[row * out_features + col] = __float2half(v); - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t qweight_desc, - infiniopTensorDescriptor_t scales_desc, - infiniopTensorDescriptor_t zeros_desc, - infiniopTensorDescriptor_t g_idx_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = DequantizeGPTQInfo::create(out_desc, qweight_desc, scales_desc, zeros_desc, g_idx_desc); - - *desc_ptr = new Descriptor( - 0, - new Opaque{handle->internal()}, - result.take(), - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *qweight, - const void *scales, - const void *zeros, - const void *g_idx, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - const int in_features = _info.in_features(); - const int out_features = _info.out_features(); - const int out_packed = _info.out_packed(); - const int in_packed = _info.in_packed(); - const int num_groups = _info.num_groups(); - - if (num_groups <= 0 || in_features <= 0 || out_features <= 0 || out_packed <= 0 || in_packed <= 0) { - return INFINI_STATUS_BAD_PARAM; - } - - constexpr int BLOCK_X = 16; // packed columns - constexpr int BLOCK_Y = 4; // rows - dim3 threads(BLOCK_X, BLOCK_Y); - dim3 blocks((out_packed + BLOCK_X - 1) / BLOCK_X, - (in_features + BLOCK_Y - 1) / BLOCK_Y); - - dequantize_weights_gptq<<(stream)>>>( - reinterpret_cast(qweight), - reinterpret_cast(scales), - reinterpret_cast(zeros), - reinterpret_cast(g_idx), - reinterpret_cast(out), - in_features, out_features, out_packed, num_groups); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dequantize_gptq::iluvatar diff --git a/src/infiniop/ops/dequantize_gptq/iluvatar/dequantize_w42f16_iluvatar.cuh b/src/infiniop/ops/dequantize_gptq/iluvatar/dequantize_w42f16_iluvatar.cuh deleted file mode 100644 index 2916b0961..000000000 --- a/src/infiniop/ops/dequantize_gptq/iluvatar/dequantize_w42f16_iluvatar.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __DEQUANTIZE_GPTQ_ILUVATAR_CUH__ -#define __DEQUANTIZE_GPTQ_ILUVATAR_CUH__ - -#include "../dequantize_gptq.h" - -DESCRIPTOR(iluvatar) - -#endif // __DEQUANTIZE_GPTQ_ILUVATAR_CUH__ diff --git a/src/infiniop/ops/dequantize_gptq/info.h b/src/infiniop/ops/dequantize_gptq/info.h deleted file mode 100644 index aef668bb1..000000000 --- a/src/infiniop/ops/dequantize_gptq/info.h +++ /dev/null @@ -1,51 +0,0 @@ -#ifndef __DEQUANTIZE_GPTQ_INFO_H__ -#define __DEQUANTIZE_GPTQ_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -#include - -namespace op::dequantize_gptq { - -class DequantizeGPTQInfo { - DequantizeGPTQInfo() = default; - -public: - int _in_features, _out_features, _num_groups, _out_packed, _in_packed; - - int in_features() const { return _in_features; } - int out_features() const { return _out_features; } - int num_groups() const { return _num_groups; } - int out_packed() const { return _out_packed; } - int in_packed() const { return _in_packed; } - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t qweight_desc, - infiniopTensorDescriptor_t scales_desc, - infiniopTensorDescriptor_t zeros_desc, - infiniopTensorDescriptor_t g_idx_desc) { - - const int _in_features = g_idx_desc->dim(0); // real input channels - const int _in_packed = qweight_desc->dim(0); // ceil(in_features / 8) - const int _out_features = qweight_desc->dim(1); // real output channels - const int _num_groups = scales_desc->dim(0); // should be in_features / group_size - const int _out_packed = zeros_desc->dim(1); // ceil(out_features / 8) - - assert(out_desc->dim(0) == _in_features); - assert(out_desc->dim(1) == _out_features); - assert(_in_packed == (_in_features + 7) / 8); - assert(scales_desc->dim(1) == _out_features); - assert(_num_groups == zeros_desc->dim(0)); - assert(_out_packed == (_out_features + 7) / 8); - - return utils::Result( - DequantizeGPTQInfo{_in_features, _out_features, _num_groups, _out_packed, _in_packed}); - } -}; - -} // namespace op::dequantize_gptq - -#endif // __DEQUANTIZE_GPTQ_INFO_H__ diff --git a/src/infiniop/ops/dequantize_gptq/moore/dequantize_w42f16_moore.h b/src/infiniop/ops/dequantize_gptq/moore/dequantize_w42f16_moore.h deleted file mode 100644 index 689657e0a..000000000 --- a/src/infiniop/ops/dequantize_gptq/moore/dequantize_w42f16_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __DEQUANTIZE_GPTQ_MOORE_H__ -#define __DEQUANTIZE_GPTQ_MOORE_H__ - -#include "../dequantize_gptq.h" - -DESCRIPTOR(moore) - -#endif // __DEQUANTIZE_GPTQ_MOORE_H__ diff --git a/src/infiniop/ops/dequantize_gptq/moore/dequantize_w42f16_moore.mu b/src/infiniop/ops/dequantize_gptq/moore/dequantize_w42f16_moore.mu deleted file mode 100644 index d3c31dd52..000000000 --- a/src/infiniop/ops/dequantize_gptq/moore/dequantize_w42f16_moore.mu +++ /dev/null @@ -1,124 +0,0 @@ -#include "../../../devices/moore/moore_handle.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "dequantize_w42f16_moore.h" - -#include "../dequantize_gptq.h" -#include -#include - -// 对齐 nvidia 版：支持 g_idx -// qweight: [in_packed, out_features]，每个 uint32 打包 8 个输入通道的 4-bit -// zeros: [num_groups, out_packed]，每个 uint32 打包 8 个输出通道的 4-bit -// scales: [num_groups, out_features] -// g_idx: [in_features] -__global__ void __launch_bounds__(128) -dequantize_weights_gptq(const uint32_t *__restrict__ qweight, - const half *__restrict__ scales, - const uint32_t *__restrict__ zeros, - const int *__restrict__ g_idx, - half *__restrict__ out, - int in_features, - int out_features, - int out_packed, // ceil(out_features / 8) - int num_groups) { - const int col_pack = blockIdx.x * blockDim.x + threadIdx.x; // packed output column - const int row = blockIdx.y * blockDim.y + threadIdx.y; // real input row - if (col_pack >= out_packed || row >= in_features) return; - - // clamp gid to [0, num_groups) - const int gid_raw = g_idx ? g_idx[row] : 0; - const int gid = ((gid_raw % num_groups) + num_groups) % num_groups; - - const int pack_row = row >> 3; // packed input row (8 rows per pack) - const int q_shift = (row & 7) * 4; // nibble shift within uint32 - - const int zero_idx = gid * out_packed + col_pack; - const uint32_t zeros_loaded = zeros[zero_idx]; - - const int col_base = col_pack << 3; // 8 real cols per pack - const int scale_base = gid * out_features + col_base; - - #pragma unroll - for (int j = 0; j < 8; ++j) { - const int col = col_base + j; - if (col >= out_features) break; - - const uint32_t q_loaded = qweight[pack_row * out_features + col]; - const int q_nib = (q_loaded >> q_shift) & 0xF; - - const int z_nib = (zeros_loaded >> (j * 4)) & 0xF; - const half scale = scales[scale_base + j]; - - // 与 nvidia 版一致： (q - (z + 1)) * s - const float v = float(q_nib - (z_nib + 1)) * __half2float(scale); - out[row * out_features + col] = __float2half(v); - } -} - -namespace op::dequantize_gptq::moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t qweight_desc, - infiniopTensorDescriptor_t scales_desc, - infiniopTensorDescriptor_t zeros_desc, - infiniopTensorDescriptor_t g_idx_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = DequantizeGPTQInfo::create(out_desc, qweight_desc, scales_desc, zeros_desc, g_idx_desc); - - *desc_ptr = new Descriptor( - 0, - new Opaque{handle->internal()}, - result.take(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *out, - const void *qweight, - const void *scales, - const void *zeros, - const void *g_idx, - void *stream) const { - - const int in_features = _info.in_features(); - const int out_features = _info.out_features(); - const int out_packed = _info.out_packed(); - const int in_packed = _info.in_packed(); - const int num_groups = _info.num_groups(); - - if (num_groups <= 0 || in_features <= 0 || out_features <= 0 || out_packed <= 0 || in_packed <= 0) - return INFINI_STATUS_BAD_PARAM; - - constexpr int BLOCK_X = 16; // packed columns - constexpr int BLOCK_Y = 4; // rows - dim3 threads(BLOCK_X, BLOCK_Y); - dim3 blocks((out_packed + BLOCK_X - 1) / BLOCK_X, - (in_features + BLOCK_Y - 1) / BLOCK_Y); - - dequantize_weights_gptq<<(stream)>>>( - reinterpret_cast(qweight), - reinterpret_cast(scales), - reinterpret_cast(zeros), - reinterpret_cast(g_idx), - reinterpret_cast(out), - in_features, out_features, out_packed, num_groups); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dequantize_gptq::moore diff --git a/src/infiniop/ops/dequantize_gptq/nvidia/dequantize_w42f16_nvidia.cu b/src/infiniop/ops/dequantize_gptq/nvidia/dequantize_w42f16_nvidia.cu deleted file mode 100644 index 5e1d15eaf..000000000 --- a/src/infiniop/ops/dequantize_gptq/nvidia/dequantize_w42f16_nvidia.cu +++ /dev/null @@ -1,130 +0,0 @@ -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) - -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../dequantize_gptq.h" -#include "dequantize_w42f16_nvidia.cuh" -#include - -namespace op::dequantize_gptq::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t qweight_desc, - infiniopTensorDescriptor_t scales_desc, - infiniopTensorDescriptor_t zeros_desc, - infiniopTensorDescriptor_t g_idx_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = DequantizeGPTQInfo::create(out_desc, qweight_desc, scales_desc, zeros_desc, g_idx_desc); - - *desc_ptr = new Descriptor( - 0, - new Opaque{handle->internal()}, - result.take(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -// 仅保留这一版内核（支持 g_idx） -// qweight: [in_packed, out_features] packing 8 input channels per word -// zeros: [num_groups, out_packed] packing 8 output channels per word -// scales: [num_groups, out_features], g_idx: [in_features] -__global__ void __launch_bounds__(128) - dequantize_weights_gptq(const uint32_t *__restrict__ qweight, - const half *__restrict__ scales, - const uint32_t *__restrict__ zeros, - const int *__restrict__ g_idx, - half *__restrict__ out, - int in_features, - int out_features, - int out_packed, // ceil(out_features / 8) - int num_groups) { - // Each thread handles one packed output column (8 real output cols). - const int col_pack = blockIdx.x * blockDim.x + threadIdx.x; // packed output column - const int row = blockIdx.y * blockDim.y + threadIdx.y; // real input row - if (col_pack >= out_packed || row >= in_features) { - return; - } - - // Clamp gid to valid range - const int gid_raw = g_idx ? g_idx[row] : 0; - const int gid = ((gid_raw % num_groups) + num_groups) % num_groups; - - const int pack_row = row >> 3; // packed input row - - const int zero_idx = gid * out_packed + col_pack; // zeros layout: [num_groups, out_packed] - const uint32_t zeros_loaded = zeros[zero_idx]; - - const int q_shift = (row & 7) * 4; // qweight packs 8 input rows - const int col_base = col_pack << 3; // 8 real cols per pack - const int scale_base = gid * out_features + col_base; - -#pragma unroll - for (int j = 0; j < 8; ++j) { - const int col = col_base + j; - if (col >= out_features) { - break; - } - - const uint32_t q_loaded = qweight[pack_row * out_features + col]; - const int q_nib = (q_loaded >> q_shift) & 0xF; - - const int z_nib = (zeros_loaded >> (j * 4)) & 0xF; - const half scale = scales[scale_base + j]; - - // GPTQ quirk: The stored zero point is usually offset by 1. - // Standard formula: (q - (z + 1)) * s - const float v = float(q_nib - (z_nib + 1)) * __half2float(scale); - out[row * out_features + col] = __float2half(v); - } -} - -infiniStatus_t -Descriptor::calculate( - void *workspace, size_t workspace_size, - void *out, - const void *qweight, - const void *scales, - const void *zeros, - const void *g_idx, - void *stream) const { - - const int in_features = _info.in_features(); - const int out_features = _info.out_features(); - const int out_packed = _info.out_packed(); - const int in_packed = _info.in_packed(); - const int num_groups = _info.num_groups(); - - if (num_groups <= 0 || in_features <= 0 || out_features <= 0 || out_packed <= 0 || in_packed <= 0) { - return INFINI_STATUS_BAD_PARAM; - } - - constexpr int BLOCK_X = 16; // packed columns - constexpr int BLOCK_Y = 4; // rows - dim3 threads(BLOCK_X, BLOCK_Y); - dim3 blocks((out_packed + BLOCK_X - 1) / BLOCK_X, - (in_features + BLOCK_Y - 1) / BLOCK_Y); - - dequantize_weights_gptq<<(stream)>>>( - reinterpret_cast(qweight), - reinterpret_cast(scales), - reinterpret_cast(zeros), - reinterpret_cast(g_idx), - reinterpret_cast(out), - in_features, out_features, out_packed, num_groups); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dequantize_gptq::nvidia - -#endif diff --git a/src/infiniop/ops/dequantize_gptq/nvidia/dequantize_w42f16_nvidia.cuh b/src/infiniop/ops/dequantize_gptq/nvidia/dequantize_w42f16_nvidia.cuh deleted file mode 100644 index ccaa0e429..000000000 --- a/src/infiniop/ops/dequantize_gptq/nvidia/dequantize_w42f16_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __DEQUANTIZE_GPTQ_CUDA_CUH__ -#define __DEQUANTIZE_GPTQ_CUDA_CUH__ - -#include "../dequantize_gptq.h" - -DESCRIPTOR(nvidia) - -#endif // __DEQUANTIZE_GPTQ_CUDA_CUH__ diff --git a/src/infiniop/ops/dequantize_gptq/operator.cc b/src/infiniop/ops/dequantize_gptq/operator.cc deleted file mode 100644 index d48430445..000000000 --- a/src/infiniop/ops/dequantize_gptq/operator.cc +++ /dev/null @@ -1,143 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/dequantize_gptq.h" - -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) -#include "nvidia/dequantize_w42f16_nvidia.cuh" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/dequantize_w42f16_moore.h" -#endif -#ifdef ENABLE_ILUVATAR_API -#include "iluvatar/dequantize_w42f16_iluvatar.cuh" -#endif - -__INFINI_C infiniStatus_t infiniopCreateDequantizeGPTQDescriptor( - infiniopHandle_t handle, - infiniopDequantizeGPTQDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t qweight_desc, - infiniopTensorDescriptor_t scales_desc, - infiniopTensorDescriptor_t zeros_desc, - infiniopTensorDescriptor_t g_idx_desc) { // add g_idx -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::dequantize_gptq::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - out_desc, \ - qweight_desc, \ - scales_desc, \ - zeros_desc, g_idx_desc) - - switch (handle->device) { -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, iluvatar); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetDequantizeGPTQWorkspaceSize(infiniopDequantizeGPTQDescriptor_t desc, - size_t *size) { -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, iluvatar); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopDequantizeGPTQ( - infiniopDequantizeGPTQDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *qweight, - const void *scales, - const void *zeros, - const void *g_idx, // add g_idx - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, out, qweight, scales, zeros, g_idx, stream) - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, iluvatar); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyDequantizeGPTQDescriptor(infiniopDequantizeGPTQDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, iluvatar); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} - -// #endif diff --git a/src/infiniop/ops/diff/cpu/diff_cpu.cc b/src/infiniop/ops/diff/cpu/diff_cpu.cc deleted file mode 100644 index d3c7ea1ac..000000000 --- a/src/infiniop/ops/diff/cpu/diff_cpu.cc +++ /dev/null @@ -1,178 +0,0 @@ -#include "diff_cpu.h" -#include "../../../tensor.h" -#include -#include -#include - -namespace op::diff::cpu { - -utils::Result DiffInfo::create( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - int dim, - int n) { - - if (n <= 0) { - return INFINI_STATUS_BAD_PARAM; - } - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - size_t ndim = x_desc->ndim(); - - if (dim < 0) { - dim += static_cast(ndim); - } - if (dim < 0 || dim >= static_cast(ndim)) { - return INFINI_STATUS_BAD_PARAM; - } - - if (x_shape[dim] <= static_cast(n)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // Calculate output shape - std::vector expected_output_shape = x_shape; - expected_output_shape[dim] -= n; - - if (y_shape != expected_output_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - DiffInfo info; - info.ndim = ndim; - info.dim = dim; - info.n = n; - info.input_shape = x_shape; - info.output_shape = y_shape; - info.input_strides = x_desc->strides(); - info.output_strides = y_desc->strides(); - info.input_size = x_desc->numel(); - info.output_size = y_desc->numel(); - - return utils::Result(std::move(info)); -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int dim, - int n) { - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto info_result = DiffInfo::create(x_desc, y_desc, dim, n); - CHECK_RESULT(info_result); - - *desc_ptr = new Descriptor(dtype, info_result.take(), handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -void diff_impl( - const DiffInfo &info, - T *y, - const T *x) { - - // n-th order forward difference along `dim`: - // y[i] = sum_{k=0..n} (-1)^(n-k) * C(n,k) * x[i+k] - // Implemented directly to: - // - avoid intermediate buffers (and their size pitfalls for n>1) - // - respect input/output strides (tests cover as_strided cases) - - auto binom = [](int n, int k) -> double { - if (k < 0 || k > n) { - return 0.0; - } - k = std::min(k, n - k); - double res = 1.0; - for (int i = 1; i <= k; ++i) { - res *= static_cast(n - (k - i)); - res /= static_cast(i); - } - return res; - }; - - std::vector coeff(static_cast(info.n) + 1); - for (int k = 0; k <= info.n; ++k) { - double c = binom(info.n, k); - if (((info.n - k) & 1) != 0) { - c = -c; - } - coeff[static_cast(k)] = c; - } - - const auto &out_shape = info.output_shape; - const auto &in_strides = info.input_strides; - const auto &out_strides = info.output_strides; - const size_t out_numel = info.output_size; - const ptrdiff_t stride_dim = in_strides[static_cast(info.dim)]; - - auto unravel_index = [](size_t linear, const std::vector &shape, std::vector &idx) { - const size_t ndim = shape.size(); - for (size_t d = ndim; d-- > 0;) { - const size_t s = shape[d]; - idx[d] = linear % s; - linear /= s; - } - }; - -#pragma omp parallel - { - std::vector idx(info.ndim, 0); - -#pragma omp for - for (ptrdiff_t linear = 0; linear < static_cast(out_numel); ++linear) { - unravel_index(static_cast(linear), out_shape, idx); - - ptrdiff_t y_off = 0; - ptrdiff_t x_base_off = 0; - for (size_t d = 0; d < info.ndim; ++d) { - y_off += static_cast(idx[d]) * out_strides[d]; - x_base_off += static_cast(idx[d]) * in_strides[d]; - } - - double acc = 0.0; - for (int k = 0; k <= info.n; ++k) { - const ptrdiff_t x_off = x_base_off + static_cast(k) * stride_dim; - acc += coeff[static_cast(k)] * utils::cast(x[x_off]); - } - - y[y_off] = utils::cast(acc); - } - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - diff_impl(_info, reinterpret_cast(y), reinterpret_cast(x)); - break; - case INFINI_DTYPE_BF16: - diff_impl(_info, reinterpret_cast(y), reinterpret_cast(x)); - break; - case INFINI_DTYPE_F32: - diff_impl(_info, reinterpret_cast(y), reinterpret_cast(x)); - break; - case INFINI_DTYPE_F64: - diff_impl(_info, reinterpret_cast(y), reinterpret_cast(x)); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::diff::cpu diff --git a/src/infiniop/ops/diff/cpu/diff_cpu.h b/src/infiniop/ops/diff/cpu/diff_cpu.h deleted file mode 100644 index 3fdabe49d..000000000 --- a/src/infiniop/ops/diff/cpu/diff_cpu.h +++ /dev/null @@ -1,61 +0,0 @@ -#ifndef __DIFF_CPU_H__ -#define __DIFF_CPU_H__ - -#include "../../../devices/cpu/common_cpu.h" -#include "../../../operator.h" -#include - -namespace op::diff::cpu { - -struct DiffInfo { - size_t ndim; - int dim; - int n; - std::vector input_shape; - std::vector output_shape; - std::vector input_strides; - std::vector output_strides; - size_t input_size; - size_t output_size; - - static utils::Result create( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - int dim, - int n); -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - DiffInfo _info; - - Descriptor(infiniDtype_t dtype, DiffInfo info, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int dim, - int n); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::diff::cpu - -#endif // __DIFF_CPU_H__ diff --git a/src/infiniop/ops/diff/cuda/kernel.cuh b/src/infiniop/ops/diff/cuda/kernel.cuh deleted file mode 100644 index f371872e8..000000000 --- a/src/infiniop/ops/diff/cuda/kernel.cuh +++ /dev/null @@ -1,82 +0,0 @@ -#include - -struct Diff1Indexing { - static constexpr int kMaxNdim = 8; - - int ndim; - int dim; - int64_t out_shape[kMaxNdim]; - int64_t in_strides[kMaxNdim]; - int64_t out_strides[kMaxNdim]; -}; - -namespace op::cuda { - -template -__device__ __forceinline__ T from_f32(float v); - -template <> -__device__ __forceinline__ half from_f32(float v) { - return __float2half(v); -} - -template <> -__device__ __forceinline__ cuda_bfloat16 from_f32(float v) { - return __float2bfloat16_rn(v); -} - -template <> -__device__ __forceinline__ float from_f32(float v) { - return v; -} - -template -__global__ void diff1_strided_kernel( - T *out, - const T *in, - size_t out_numel, - Diff1Indexing indexing) { - - const size_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (linear >= out_numel) { - return; - } - - int64_t idx[Diff1Indexing::kMaxNdim] = {0}; - size_t tmp = linear; - for (int d = indexing.ndim - 1; d >= 0; --d) { - const int64_t s = indexing.out_shape[d]; - idx[d] = static_cast(tmp % static_cast(s)); - tmp /= static_cast(s); - } - - int64_t y_off = 0; - int64_t x_base_off = 0; - for (int d = 0; d < indexing.ndim; ++d) { - y_off += idx[d] * indexing.out_strides[d]; - x_base_off += idx[d] * indexing.in_strides[d]; - } - - const int64_t stride_dim = indexing.in_strides[indexing.dim]; - const int64_t x_off1 = x_base_off; - const int64_t x_off2 = x_base_off + stride_dim; - - if constexpr (std::is_same_v) { - out[y_off] = in[x_off2] - in[x_off1]; - } else { - float a; - float b; - if constexpr (std::is_same_v) { - a = __half2float(in[x_off1]); - b = __half2float(in[x_off2]); - } else if constexpr (std::is_same_v) { - a = __bfloat162float(in[x_off1]); - b = __bfloat162float(in[x_off2]); - } else { // float - a = static_cast(in[x_off1]); - b = static_cast(in[x_off2]); - } - out[y_off] = from_f32(b - a); - } -} -} // namespace op::cuda diff --git a/src/infiniop/ops/diff/metax/diff_metax.h b/src/infiniop/ops/diff/metax/diff_metax.h deleted file mode 100644 index c8d60c1db..000000000 --- a/src/infiniop/ops/diff/metax/diff_metax.h +++ /dev/null @@ -1,69 +0,0 @@ -#ifndef __DIFF_METAX_H__ -#define __DIFF_METAX_H__ - -#include "../../../operator.h" -#include - -namespace op::diff::metax { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _ndim; - int _dim; - int _n; - std::vector _input_shape; - std::vector _output_shape; - std::vector _input_strides; - std::vector _output_strides; - size_t _input_size; - size_t _output_size; - - Descriptor(infiniDtype_t dtype, size_t ndim, int dim, int n, - std::vector input_shape, std::vector output_shape, - std::vector input_strides, std::vector output_strides, - size_t input_size, size_t output_size, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _ndim(ndim), - _dim(dim), - _n(n), - _input_shape(std::move(input_shape)), - _output_shape(std::move(output_shape)), - _input_strides(std::move(input_strides)), - _output_strides(std::move(output_strides)), - _input_size(input_size), - _output_size(output_size) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int dim, - int n); - - size_t workspaceSize() const { - if (_n <= 1) { - return 0; - } - const size_t dim_size = _input_shape[static_cast(_dim)]; - const size_t outer = _input_size / dim_size; - const size_t max_intermediate = outer * (dim_size - 1); - return 2 * max_intermediate * infiniSizeOf(_dtype); - } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::diff::metax - -#endif // __DIFF_METAX_H__ diff --git a/src/infiniop/ops/diff/metax/diff_metax.maca b/src/infiniop/ops/diff/metax/diff_metax.maca deleted file mode 100644 index 6d1deac05..000000000 --- a/src/infiniop/ops/diff/metax/diff_metax.maca +++ /dev/null @@ -1,194 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../../../tensor.h" - -#include "../cuda/kernel.cuh" -#include "diff_metax.h" -#include -#include -#include -#include - -namespace op::diff::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int dim, - int n) { - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - if (n <= 0) { - return INFINI_STATUS_BAD_PARAM; - } - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - size_t ndim = x_desc->ndim(); - - if (dim < 0) { - dim += static_cast(ndim); - } - if (dim < 0 || dim >= static_cast(ndim)) { - return INFINI_STATUS_BAD_PARAM; - } - - if (x_shape[dim] <= static_cast(n)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - std::vector expected_output_shape = x_shape; - expected_output_shape[dim] -= n; - - if (y_shape != expected_output_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - *desc_ptr = new Descriptor(dtype, ndim, dim, n, x_shape, y_shape, - x_desc->strides(), y_desc->strides(), - x_desc->numel(), y_desc->numel(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - - constexpr int BLOCK_SIZE = 256; - - auto numel_of = [](const std::vector &shape) -> size_t { - return std::accumulate(shape.begin(), shape.end(), static_cast(1), std::multiplies{}); - }; - auto contiguous_strides = [](const std::vector &shape) -> std::vector { - std::vector strides(shape.size(), 1); - ptrdiff_t running = 1; - for (size_t d = shape.size(); d-- > 0;) { - strides[d] = running; - running *= static_cast(shape[d]); - } - return strides; - }; - auto fill_indexing = [&](Diff1Indexing &indexing, - const std::vector &out_shape, - const std::vector &in_strides, - const std::vector &out_strides) -> infiniStatus_t { - indexing.ndim = static_cast(_ndim); - indexing.dim = _dim; - if (indexing.ndim > Diff1Indexing::kMaxNdim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - for (int d = 0; d < Diff1Indexing::kMaxNdim; ++d) { - indexing.out_shape[d] = 1; - indexing.in_strides[d] = 0; - indexing.out_strides[d] = 0; - } - for (size_t d = 0; d < _ndim; ++d) { - indexing.out_shape[d] = static_cast(out_shape[d]); - indexing.in_strides[d] = static_cast(in_strides[d]); - indexing.out_strides[d] = static_cast(out_strides[d]); - } - return INFINI_STATUS_SUCCESS; - }; - - auto launch_diff1 = [&](void *out_ptr, - const void *in_ptr, - const std::vector &out_shape, - const std::vector &in_strides, - const std::vector &out_strides) -> infiniStatus_t { - const size_t out_numel = numel_of(out_shape); - const int blocks = static_cast((out_numel + BLOCK_SIZE - 1) / BLOCK_SIZE); - Diff1Indexing indexing{}; - auto st = fill_indexing(indexing, out_shape, in_strides, out_strides); - if (st != INFINI_STATUS_SUCCESS) { - return st; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - cuda::diff1_strided_kernel<<>>( - reinterpret_cast(out_ptr), reinterpret_cast(in_ptr), out_numel, indexing); - return INFINI_STATUS_SUCCESS; - case INFINI_DTYPE_BF16: - cuda::diff1_strided_kernel<<>>( - reinterpret_cast(out_ptr), reinterpret_cast(in_ptr), out_numel, indexing); - return INFINI_STATUS_SUCCESS; - case INFINI_DTYPE_F32: - cuda::diff1_strided_kernel<<>>( - reinterpret_cast(out_ptr), reinterpret_cast(in_ptr), out_numel, indexing); - return INFINI_STATUS_SUCCESS; - case INFINI_DTYPE_F64: - cuda::diff1_strided_kernel<<>>( - reinterpret_cast(out_ptr), reinterpret_cast(in_ptr), out_numel, indexing); - return INFINI_STATUS_SUCCESS; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - }; - - if (_n == 1) { - return launch_diff1(y, x, _output_shape, _input_strides, _output_strides); - } - - if (workspace_size < this->workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - const size_t elem_size = infiniSizeOf(_dtype); - const size_t dim_size = _input_shape[static_cast(_dim)]; - const size_t outer = _input_size / dim_size; - const size_t max_intermediate = outer * (dim_size - 1); - - auto *ws = reinterpret_cast(workspace); - void *buf_a = ws; - void *buf_b = ws + max_intermediate * elem_size; - - std::vector current_shape = _input_shape; - std::vector current_in_strides = _input_strides; - - std::vector out_shape = current_shape; - out_shape[static_cast(_dim)] -= 1; - std::vector out_strides = contiguous_strides(out_shape); - - auto st = launch_diff1(buf_a, x, out_shape, current_in_strides, out_strides); - if (st != INFINI_STATUS_SUCCESS) { - return st; - } - - current_shape = out_shape; - current_in_strides = out_strides; - bool a_is_input = true; - - for (int stage = 1; stage < _n - 1; ++stage) { - out_shape = current_shape; - out_shape[static_cast(_dim)] -= 1; - out_strides = contiguous_strides(out_shape); - - void *in_buf = a_is_input ? buf_a : buf_b; - void *out_buf = a_is_input ? buf_b : buf_a; - st = launch_diff1(out_buf, in_buf, out_shape, current_in_strides, out_strides); - if (st != INFINI_STATUS_SUCCESS) { - return st; - } - current_shape = out_shape; - current_in_strides = out_strides; - a_is_input = !a_is_input; - } - - void *in_buf = a_is_input ? buf_a : buf_b; - return launch_diff1(y, in_buf, _output_shape, current_in_strides, _output_strides); -} - -} // namespace op::diff::metax diff --git a/src/infiniop/ops/diff/moore/diff_moore.h b/src/infiniop/ops/diff/moore/diff_moore.h deleted file mode 100644 index bf4a0e11b..000000000 --- a/src/infiniop/ops/diff/moore/diff_moore.h +++ /dev/null @@ -1,67 +0,0 @@ -#ifndef __DIFF_MOORE_H__ -#define __DIFF_MOORE_H__ - -#include "../../../operator.h" - -namespace op::diff::moore { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _ndim; - int _dim; - int _n; - std::vector _input_shape; - std::vector _output_shape; - std::vector _input_strides; - std::vector _output_strides; - size_t _input_size; - size_t _output_size; - - Descriptor(infiniDtype_t dtype, size_t ndim, int dim, int n, - std::vector input_shape, std::vector output_shape, - std::vector input_strides, std::vector output_strides, - size_t input_size, size_t output_size, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _ndim(ndim), - _dim(dim), - _n(n), - _input_shape(std::move(input_shape)), - _output_shape(std::move(output_shape)), - _input_strides(std::move(input_strides)), - _output_strides(std::move(output_strides)), - _input_size(input_size), - _output_size(output_size) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int dim, - int n); - - size_t workspaceSize() const { - if (_n <= 1) { - return 0; - } - const size_t dim_size = _input_shape[static_cast(_dim)]; - const size_t outer = _input_size / dim_size; - const size_t max_intermediate = outer * (dim_size - 1); - return 2 * max_intermediate * infiniSizeOf(_dtype); - } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; -} // namespace op::diff::moore - -#endif // __DIFF_MOORE_H__ diff --git a/src/infiniop/ops/diff/moore/diff_moore.mu b/src/infiniop/ops/diff/moore/diff_moore.mu deleted file mode 100644 index c76311baf..000000000 --- a/src/infiniop/ops/diff/moore/diff_moore.mu +++ /dev/null @@ -1,193 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "../../../tensor.h" - -#include "../cuda/kernel.cuh" -#include "diff_moore.h" -#include -#include -#include -#include - -namespace op::diff::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int dim, - int n) { - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - if (n <= 0) { - return INFINI_STATUS_BAD_PARAM; - } - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - size_t ndim = x_desc->ndim(); - - if (dim < 0) { - dim += static_cast(ndim); - } - if (dim < 0 || dim >= static_cast(ndim)) { - return INFINI_STATUS_BAD_PARAM; - } - - if (x_shape[dim] <= static_cast(n)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - std::vector expected_output_shape = x_shape; - expected_output_shape[dim] -= n; - - if (y_shape != expected_output_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - *desc_ptr = new Descriptor(dtype, ndim, dim, n, x_shape, y_shape, - x_desc->strides(), y_desc->strides(), - x_desc->numel(), y_desc->numel(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - - constexpr int BLOCK_SIZE = 256; - - auto numel_of = [](const std::vector &shape) -> size_t { - return std::accumulate(shape.begin(), shape.end(), static_cast(1), std::multiplies{}); - }; - auto contiguous_strides = [](const std::vector &shape) -> std::vector { - std::vector strides(shape.size(), 1); - ptrdiff_t running = 1; - for (size_t d = shape.size(); d-- > 0;) { - strides[d] = running; - running *= static_cast(shape[d]); - } - return strides; - }; - auto fill_indexing = [&](Diff1Indexing &indexing, - const std::vector &out_shape, - const std::vector &in_strides, - const std::vector &out_strides) -> infiniStatus_t { - indexing.ndim = static_cast(_ndim); - indexing.dim = _dim; - if (indexing.ndim > Diff1Indexing::kMaxNdim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - for (int d = 0; d < Diff1Indexing::kMaxNdim; ++d) { - indexing.out_shape[d] = 1; - indexing.in_strides[d] = 0; - indexing.out_strides[d] = 0; - } - for (size_t d = 0; d < _ndim; ++d) { - indexing.out_shape[d] = static_cast(out_shape[d]); - indexing.in_strides[d] = static_cast(in_strides[d]); - indexing.out_strides[d] = static_cast(out_strides[d]); - } - return INFINI_STATUS_SUCCESS; - }; - - auto launch_diff1 = [&](void *out_ptr, - const void *in_ptr, - const std::vector &out_shape, - const std::vector &in_strides, - const std::vector &out_strides) -> infiniStatus_t { - const size_t out_numel = numel_of(out_shape); - const int blocks = static_cast((out_numel + BLOCK_SIZE - 1) / BLOCK_SIZE); - Diff1Indexing indexing{}; - auto st = fill_indexing(indexing, out_shape, in_strides, out_strides); - if (st != INFINI_STATUS_SUCCESS) { - return st; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - cuda::diff1_strided_kernel<<>>( - reinterpret_cast(out_ptr), reinterpret_cast(in_ptr), out_numel, indexing); - return INFINI_STATUS_SUCCESS; - case INFINI_DTYPE_BF16: - cuda::diff1_strided_kernel<<>>( - reinterpret_cast(out_ptr), reinterpret_cast(in_ptr), out_numel, indexing); - return INFINI_STATUS_SUCCESS; - case INFINI_DTYPE_F32: - cuda::diff1_strided_kernel<<>>( - reinterpret_cast(out_ptr), reinterpret_cast(in_ptr), out_numel, indexing); - return INFINI_STATUS_SUCCESS; - case INFINI_DTYPE_F64: - cuda::diff1_strided_kernel<<>>( - reinterpret_cast(out_ptr), reinterpret_cast(in_ptr), out_numel, indexing); - return INFINI_STATUS_SUCCESS; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - }; - - if (_n == 1) { - return launch_diff1(y, x, _output_shape, _input_strides, _output_strides); - } - - if (workspace_size < this->workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - const size_t elem_size = infiniSizeOf(_dtype); - const size_t dim_size = _input_shape[static_cast(_dim)]; - const size_t outer = _input_size / dim_size; - const size_t max_intermediate = outer * (dim_size - 1); - - auto *ws = reinterpret_cast(workspace); - void *buf_a = ws; - void *buf_b = ws + max_intermediate * elem_size; - - std::vector current_shape = _input_shape; - std::vector current_in_strides = _input_strides; - - std::vector out_shape = current_shape; - out_shape[static_cast(_dim)] -= 1; - std::vector out_strides = contiguous_strides(out_shape); - - auto st = launch_diff1(buf_a, x, out_shape, current_in_strides, out_strides); - if (st != INFINI_STATUS_SUCCESS) { - return st; - } - - current_shape = out_shape; - current_in_strides = out_strides; - bool a_is_input = true; - - for (int stage = 1; stage < _n - 1; ++stage) { - out_shape = current_shape; - out_shape[static_cast(_dim)] -= 1; - out_strides = contiguous_strides(out_shape); - - void *in_buf = a_is_input ? buf_a : buf_b; - void *out_buf = a_is_input ? buf_b : buf_a; - st = launch_diff1(out_buf, in_buf, out_shape, current_in_strides, out_strides); - if (st != INFINI_STATUS_SUCCESS) { - return st; - } - current_shape = out_shape; - current_in_strides = out_strides; - a_is_input = !a_is_input; - } - - void *in_buf = a_is_input ? buf_a : buf_b; - return launch_diff1(y, in_buf, _output_shape, current_in_strides, _output_strides); -} - -} // namespace op::diff::moore diff --git a/src/infiniop/ops/diff/nvidia/diff_nvidia.cu b/src/infiniop/ops/diff/nvidia/diff_nvidia.cu deleted file mode 100644 index 8bca78aaf..000000000 --- a/src/infiniop/ops/diff/nvidia/diff_nvidia.cu +++ /dev/null @@ -1,269 +0,0 @@ -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" -#include "../../../tensor.h" -#include "diff_nvidia.cuh" -#include -#include -#include -#include - -namespace op::diff::nvidia { - -Descriptor::~Descriptor() = default; - -template -__device__ __forceinline__ T from_f32(float v); - -template <> -__device__ __forceinline__ half from_f32(float v) { - return __float2half(v); -} - -template <> -__device__ __forceinline__ cuda_bfloat16 from_f32(float v) { - return __float2bfloat16_rn(v); -} - -template <> -__device__ __forceinline__ float from_f32(float v) { - return v; -} - -struct Diff1Indexing { - static constexpr int kMaxNdim = 8; - - int ndim; - int dim; - int64_t out_shape[kMaxNdim]; - int64_t in_strides[kMaxNdim]; - int64_t out_strides[kMaxNdim]; -}; - -template -__global__ void diff1_strided_kernel( - T *out, - const T *in, - size_t out_numel, - Diff1Indexing indexing) { - - const size_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (linear >= out_numel) { - return; - } - - int64_t idx[Diff1Indexing::kMaxNdim] = {0}; - size_t tmp = linear; - for (int d = indexing.ndim - 1; d >= 0; --d) { - const int64_t s = indexing.out_shape[d]; - idx[d] = static_cast(tmp % static_cast(s)); - tmp /= static_cast(s); - } - - int64_t y_off = 0; - int64_t x_base_off = 0; - for (int d = 0; d < indexing.ndim; ++d) { - y_off += idx[d] * indexing.out_strides[d]; - x_base_off += idx[d] * indexing.in_strides[d]; - } - - const int64_t stride_dim = indexing.in_strides[indexing.dim]; - const int64_t x_off1 = x_base_off; - const int64_t x_off2 = x_base_off + stride_dim; - - if constexpr (std::is_same_v) { - out[y_off] = in[x_off2] - in[x_off1]; - } else { - float a; - float b; - if constexpr (std::is_same_v) { - a = __half2float(in[x_off1]); - b = __half2float(in[x_off2]); - } else if constexpr (std::is_same_v) { - a = __bfloat162float(in[x_off1]); - b = __bfloat162float(in[x_off2]); - } else { // float - a = static_cast(in[x_off1]); - b = static_cast(in[x_off2]); - } - out[y_off] = from_f32(b - a); - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int dim, - int n) { - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - if (n <= 0) { - return INFINI_STATUS_BAD_PARAM; - } - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - size_t ndim = x_desc->ndim(); - - if (dim < 0) { - dim += static_cast(ndim); - } - if (dim < 0 || dim >= static_cast(ndim)) { - return INFINI_STATUS_BAD_PARAM; - } - - if (x_shape[dim] <= static_cast(n)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - std::vector expected_output_shape = x_shape; - expected_output_shape[dim] -= n; - - if (y_shape != expected_output_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - *desc_ptr = new Descriptor(dtype, ndim, dim, n, x_shape, y_shape, - x_desc->strides(), y_desc->strides(), - x_desc->numel(), y_desc->numel(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - - constexpr int BLOCK_SIZE = 256; - - auto numel_of = [](const std::vector &shape) -> size_t { - return std::accumulate(shape.begin(), shape.end(), static_cast(1), std::multiplies{}); - }; - auto contiguous_strides = [](const std::vector &shape) -> std::vector { - std::vector strides(shape.size(), 1); - ptrdiff_t running = 1; - for (size_t d = shape.size(); d-- > 0;) { - strides[d] = running; - running *= static_cast(shape[d]); - } - return strides; - }; - auto fill_indexing = [&](Diff1Indexing &indexing, - const std::vector &out_shape, - const std::vector &in_strides, - const std::vector &out_strides) -> infiniStatus_t { - indexing.ndim = static_cast(_ndim); - indexing.dim = _dim; - if (indexing.ndim > Diff1Indexing::kMaxNdim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - for (int d = 0; d < Diff1Indexing::kMaxNdim; ++d) { - indexing.out_shape[d] = 1; - indexing.in_strides[d] = 0; - indexing.out_strides[d] = 0; - } - for (size_t d = 0; d < _ndim; ++d) { - indexing.out_shape[d] = static_cast(out_shape[d]); - indexing.in_strides[d] = static_cast(in_strides[d]); - indexing.out_strides[d] = static_cast(out_strides[d]); - } - return INFINI_STATUS_SUCCESS; - }; - - auto launch_diff1 = [&](void *out_ptr, - const void *in_ptr, - const std::vector &out_shape, - const std::vector &in_strides, - const std::vector &out_strides) -> infiniStatus_t { - const size_t out_numel = numel_of(out_shape); - const int blocks = static_cast((out_numel + BLOCK_SIZE - 1) / BLOCK_SIZE); - Diff1Indexing indexing{}; - auto st = fill_indexing(indexing, out_shape, in_strides, out_strides); - if (st != INFINI_STATUS_SUCCESS) { - return st; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - diff1_strided_kernel<<>>( - reinterpret_cast(out_ptr), reinterpret_cast(in_ptr), out_numel, indexing); - return INFINI_STATUS_SUCCESS; - case INFINI_DTYPE_BF16: - diff1_strided_kernel<<>>( - reinterpret_cast(out_ptr), reinterpret_cast(in_ptr), out_numel, indexing); - return INFINI_STATUS_SUCCESS; - case INFINI_DTYPE_F32: - diff1_strided_kernel<<>>( - reinterpret_cast(out_ptr), reinterpret_cast(in_ptr), out_numel, indexing); - return INFINI_STATUS_SUCCESS; - case INFINI_DTYPE_F64: - diff1_strided_kernel<<>>( - reinterpret_cast(out_ptr), reinterpret_cast(in_ptr), out_numel, indexing); - return INFINI_STATUS_SUCCESS; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - }; - - if (_n == 1) { - return launch_diff1(y, x, _output_shape, _input_strides, _output_strides); - } - - if (workspace_size < this->workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - const size_t elem_size = infiniSizeOf(_dtype); - const size_t dim_size = _input_shape[static_cast(_dim)]; - const size_t outer = _input_size / dim_size; - const size_t max_intermediate = outer * (dim_size - 1); - - auto *ws = reinterpret_cast(workspace); - void *buf_a = ws; - void *buf_b = ws + max_intermediate * elem_size; - - std::vector current_shape = _input_shape; - std::vector current_in_strides = _input_strides; - - std::vector out_shape = current_shape; - out_shape[static_cast(_dim)] -= 1; - std::vector out_strides = contiguous_strides(out_shape); - - auto st = launch_diff1(buf_a, x, out_shape, current_in_strides, out_strides); - if (st != INFINI_STATUS_SUCCESS) { - return st; - } - - current_shape = out_shape; - current_in_strides = out_strides; - bool a_is_input = true; - - for (int stage = 1; stage < _n - 1; ++stage) { - out_shape = current_shape; - out_shape[static_cast(_dim)] -= 1; - out_strides = contiguous_strides(out_shape); - - void *in_buf = a_is_input ? buf_a : buf_b; - void *out_buf = a_is_input ? buf_b : buf_a; - st = launch_diff1(out_buf, in_buf, out_shape, current_in_strides, out_strides); - if (st != INFINI_STATUS_SUCCESS) { - return st; - } - current_shape = out_shape; - current_in_strides = out_strides; - a_is_input = !a_is_input; - } - - void *in_buf = a_is_input ? buf_a : buf_b; - return launch_diff1(y, in_buf, _output_shape, current_in_strides, _output_strides); -} - -} // namespace op::diff::nvidia diff --git a/src/infiniop/ops/diff/nvidia/diff_nvidia.cuh b/src/infiniop/ops/diff/nvidia/diff_nvidia.cuh deleted file mode 100644 index 83772d853..000000000 --- a/src/infiniop/ops/diff/nvidia/diff_nvidia.cuh +++ /dev/null @@ -1,69 +0,0 @@ -#ifndef __DIFF_NVIDIA_H__ -#define __DIFF_NVIDIA_H__ - -#include "../../../operator.h" -#include - -namespace op::diff::nvidia { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _ndim; - int _dim; - int _n; - std::vector _input_shape; - std::vector _output_shape; - std::vector _input_strides; - std::vector _output_strides; - size_t _input_size; - size_t _output_size; - - Descriptor(infiniDtype_t dtype, size_t ndim, int dim, int n, - std::vector input_shape, std::vector output_shape, - std::vector input_strides, std::vector output_strides, - size_t input_size, size_t output_size, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _ndim(ndim), - _dim(dim), - _n(n), - _input_shape(std::move(input_shape)), - _output_shape(std::move(output_shape)), - _input_strides(std::move(input_strides)), - _output_strides(std::move(output_strides)), - _input_size(input_size), - _output_size(output_size) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int dim, - int n); - - size_t workspaceSize() const { - if (_n <= 1) { - return 0; - } - const size_t dim_size = _input_shape[static_cast(_dim)]; - const size_t outer = _input_size / dim_size; - const size_t max_intermediate = outer * (dim_size - 1); - return 2 * max_intermediate * infiniSizeOf(_dtype); - } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::diff::nvidia - -#endif // __DIFF_NVIDIA_H__ diff --git a/src/infiniop/ops/diff/operator.cc b/src/infiniop/ops/diff/operator.cc deleted file mode 100644 index 3037dc3b5..000000000 --- a/src/infiniop/ops/diff/operator.cc +++ /dev/null @@ -1,173 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/diff.h" - -#ifdef ENABLE_CPU_API -#include "cpu/diff_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/diff_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/diff_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/diff_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateDiffDescriptor( - infiniopHandle_t handle, - infiniopDiffDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int dim, - int n) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::diff::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - x_desc, \ - dim, \ - n) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetDiffWorkspaceSize(infiniopDiffDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopDiff( - infiniopDiffDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, x, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyDiffDescriptor(infiniopDiffDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/digamma/cpu/digamma_cpu.cc b/src/infiniop/ops/digamma/cpu/digamma_cpu.cc deleted file mode 100644 index 84687b5c9..000000000 --- a/src/infiniop/ops/digamma/cpu/digamma_cpu.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "digamma_cpu.h" - -namespace op::digamma::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::digamma::cpu diff --git a/src/infiniop/ops/digamma/cpu/digamma_cpu.h b/src/infiniop/ops/digamma/cpu/digamma_cpu.h deleted file mode 100644 index f47be837e..000000000 --- a/src/infiniop/ops/digamma/cpu/digamma_cpu.h +++ /dev/null @@ -1,51 +0,0 @@ -#ifndef __DIGAMMA_CPU_H__ -#define __DIGAMMA_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" -#include -#include - -ELEMENTWISE_DESCRIPTOR(digamma, cpu) - -namespace op::digamma::cpu { - -// Digamma function implementation for x > 0 using recurrence + asymptotic series. -template -T digamma_impl(T x) { - if (x == static_cast(0)) { - return -std::numeric_limits::infinity(); - } - if (x < static_cast(0)) { - return std::numeric_limits::quiet_NaN(); - } - - T result = static_cast(0); - - // Recurrence to push x to a region where the asymptotic series is accurate. - while (x < static_cast(8)) { - result -= static_cast(1) / x; - x += static_cast(1); - } - - const T inv = static_cast(1) / x; - const T inv2 = inv * inv; - - // Asymptotic series: - // psi(x) = log(x) - 1/(2x) - 1/(12 x^2) + 1/(120 x^4) - 1/(252 x^6) + 1/(240 x^8) - 1/(132 x^10) + ... - const T series = inv2 * (static_cast(-1.0 / 12.0) + inv2 * (static_cast(1.0 / 120.0) + inv2 * (static_cast(-1.0 / 252.0) + inv2 * (static_cast(1.0 / 240.0) + inv2 * (static_cast(-1.0 / 132.0)))))); - - result += std::log(x) - static_cast(0.5) * inv + series; - return result; -} - -typedef struct DigammaOp { -public: - static constexpr size_t num_inputs = 1; - template - T operator()(const T &x) const { - return digamma_impl(x); - } -} DigammaOp; -} // namespace op::digamma::cpu - -#endif // __DIGAMMA_CPU_H__ diff --git a/src/infiniop/ops/digamma/cuda/kernel.cuh b/src/infiniop/ops/digamma/cuda/kernel.cuh deleted file mode 100644 index e99b389cf..000000000 --- a/src/infiniop/ops/digamma/cuda/kernel.cuh +++ /dev/null @@ -1,52 +0,0 @@ -#pragma once -#include -#include - -namespace op::cuda { - -// Digamma for x > 0 using recurrence + asymptotic series. -template -__device__ __forceinline__ T digamma_impl(T x) { - if (x == static_cast(0)) { - return static_cast(-INFINITY); - } - if (x < static_cast(0)) { - return static_cast(NAN); - } - - T result = static_cast(0); - while (x < static_cast(8)) { - result -= static_cast(1) / x; - x += static_cast(1); - } - - const T inv = static_cast(1) / x; - const T inv2 = inv * inv; - - const T series = inv2 * (static_cast(-1.0 / 12.0) + inv2 * (static_cast(1.0 / 120.0) + inv2 * (static_cast(-1.0 / 252.0) + inv2 * (static_cast(1.0 / 240.0) + inv2 * (static_cast(-1.0 / 132.0)))))); - - result += log(x) - static_cast(0.5) * inv + series; - return result; -} - -typedef struct DigammaOp { -public: - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(const T &x) const { - if constexpr (std::is_same_v) { - float xf = __half2float(x); - return __float2half(digamma_impl(xf)); - } else if constexpr (std::is_same_v) { - float xf = __bfloat162float(x); - return __float2bfloat16_rn(digamma_impl(xf)); - } else if constexpr (std::is_same_v) { - return digamma_impl(x); - } else { // double - return digamma_impl(static_cast(x)); - } - } -} DigammaOp; - -} // namespace op::cuda diff --git a/src/infiniop/ops/digamma/digamma.h b/src/infiniop/ops/digamma/digamma.h deleted file mode 100644 index f2c0e5beb..000000000 --- a/src/infiniop/ops/digamma/digamma.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __DIGAMMA_H__ -#define __DIGAMMA_H__ - -#include "../../elementwise/elementwise.h" - -#define DESCRIPTOR(NAMESPACE) ELEMENTWISE_DESCRIPTOR(digamma, NAMESPACE) - -#endif // __DIGAMMA_H__ diff --git a/src/infiniop/ops/digamma/metax/digamma_metax.h b/src/infiniop/ops/digamma/metax/digamma_metax.h deleted file mode 100644 index 26d8c6657..000000000 --- a/src/infiniop/ops/digamma/metax/digamma_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __DIGAMMA_METAX_API_H__ -#define __DIGAMMA_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(digamma, metax) - -#endif // __DIGAMMA_METAX_API_H__ diff --git a/src/infiniop/ops/digamma/metax/digamma_metax.maca b/src/infiniop/ops/digamma/metax/digamma_metax.maca deleted file mode 100644 index e25c4d9ad..000000000 --- a/src/infiniop/ops/digamma/metax/digamma_metax.maca +++ /dev/null @@ -1,59 +0,0 @@ -#include "digamma_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::digamma::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::DigammaOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::DigammaOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::DigammaOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::DigammaOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::digamma::metax diff --git a/src/infiniop/ops/digamma/moore/digamma_moore.h b/src/infiniop/ops/digamma/moore/digamma_moore.h deleted file mode 100644 index e78b4564b..000000000 --- a/src/infiniop/ops/digamma/moore/digamma_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __DIGAMMA_MOORE_API_H__ -#define __DIGAMMA_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(digamma, moore) - -#endif // __DIGAMMA_MOORE_API_H__ diff --git a/src/infiniop/ops/digamma/moore/digamma_moore.mu b/src/infiniop/ops/digamma/moore/digamma_moore.mu deleted file mode 100644 index ea8ed6112..000000000 --- a/src/infiniop/ops/digamma/moore/digamma_moore.mu +++ /dev/null @@ -1,60 +0,0 @@ -#include "digamma_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "../cuda/kernel.cuh" - -namespace op::digamma::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::DigammaOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::DigammaOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::DigammaOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::DigammaOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::digamma::moore diff --git a/src/infiniop/ops/digamma/nvidia/digamma_nvidia.cu b/src/infiniop/ops/digamma/nvidia/digamma_nvidia.cu deleted file mode 100644 index 0d0d4640f..000000000 --- a/src/infiniop/ops/digamma/nvidia/digamma_nvidia.cu +++ /dev/null @@ -1,58 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "digamma_nvidia.cuh" - -namespace op::digamma::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, op::cuda::DigammaOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, op::cuda::DigammaOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, op::cuda::DigammaOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, op::cuda::DigammaOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::digamma::nvidia diff --git a/src/infiniop/ops/digamma/nvidia/digamma_nvidia.cuh b/src/infiniop/ops/digamma/nvidia/digamma_nvidia.cuh deleted file mode 100644 index 452b690c9..000000000 --- a/src/infiniop/ops/digamma/nvidia/digamma_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __DIGAMMA_NVIDIA_H__ -#define __DIGAMMA_NVIDIA_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(digamma, nvidia) - -#endif // __DIGAMMA_NVIDIA_H__ diff --git a/src/infiniop/ops/digamma/operator.cc b/src/infiniop/ops/digamma/operator.cc deleted file mode 100644 index d6be706ca..000000000 --- a/src/infiniop/ops/digamma/operator.cc +++ /dev/null @@ -1,169 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/digamma.h" - -#ifdef ENABLE_CPU_API -#include "cpu/digamma_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/digamma_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/digamma_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/digamma_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateDigammaDescriptor( - infiniopHandle_t handle, - infiniopDigammaDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::digamma::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - {x_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetDigammaWorkspaceSize(infiniopDigammaDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopDigamma( - infiniopDigammaDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, {x}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyDigammaDescriptor(infiniopDigammaDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/dist/cpu/dist_cpu.cc b/src/infiniop/ops/dist/cpu/dist_cpu.cc deleted file mode 100644 index ec78fc21e..000000000 --- a/src/infiniop/ops/dist/cpu/dist_cpu.cc +++ /dev/null @@ -1,144 +0,0 @@ -#include "dist_cpu.h" -#include "../../../tensor.h" -#include -#include - -namespace op::dist::cpu { - -utils::Result DistInfo::create( - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - infiniopTensorDescriptor_t y_desc, - double p) { - - auto x1_shape = x1_desc->shape(); - auto x2_shape = x2_desc->shape(); - auto y_shape = y_desc->shape(); - - // Check that x1 and x2 have same shape - if (x1_shape != x2_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // Check that y is a scalar (0D tensor or shape [1]) - if (y_shape.size() != 0 && (y_shape.size() != 1 || y_shape[0] != 1)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - DistInfo info; - info.input_size = x1_desc->numel(); - info.p = p; - info.x1_strides = x1_desc->strides(); - info.x2_strides = x2_desc->strides(); - info.shape = x1_shape; - info.ndim = x1_desc->ndim(); - - return utils::Result(std::move(info)); -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p) { - - auto dtype = x1_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto info_result = DistInfo::create(x1_desc, x2_desc, y_desc, p); - CHECK_RESULT(info_result); - - *desc_ptr = new Descriptor(dtype, info_result.take(), handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -void dist_impl( - const DistInfo &info, - T *y, - const T *x1, - const T *x2) { - - double sum = 0.0; - const double p = info.p; - - for (size_t i = 0; i < info.input_size; ++i) { - size_t idx1 = info.x1_strides.size() == 1 && info.x1_strides[0] == 1 - ? i - : op::common_cpu::indexToOffset(i, info.ndim, info.shape.data(), info.x1_strides.data()); - size_t idx2 = info.x2_strides.size() == 1 && info.x2_strides[0] == 1 - ? i - : op::common_cpu::indexToOffset(i, info.ndim, info.shape.data(), info.x2_strides.data()); - - double diff = utils::cast(x1[idx1]) - utils::cast(x2[idx2]); - double abs_diff = std::abs(diff); - - if (p == 0.0) { - // L0 norm: count non-zero differences - if (abs_diff > 1e-10) { - sum += 1.0; - } - } else if (p == std::numeric_limits::infinity()) { - // L-infinity norm: max absolute difference - sum = std::max(sum, abs_diff); - } else { - // Lp norm: sum of |diff|^p - sum += std::pow(abs_diff, p); - } - } - - // Take p-th root (except for p=0 and p=inf) - if (p == 0.0) { - *y = utils::cast(sum); - } else if (p == std::numeric_limits::infinity()) { - *y = utils::cast(sum); - } else { - *y = utils::cast(std::pow(sum, 1.0 / p)); - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: { - dist_impl(_info, reinterpret_cast(y), - reinterpret_cast(x1), - reinterpret_cast(x2)); - break; - } - case INFINI_DTYPE_BF16: { - dist_impl(_info, reinterpret_cast(y), - reinterpret_cast(x1), - reinterpret_cast(x2)); - break; - } - case INFINI_DTYPE_F32: { - dist_impl(_info, reinterpret_cast(y), - reinterpret_cast(x1), - reinterpret_cast(x2)); - break; - } - case INFINI_DTYPE_F64: { - dist_impl(_info, reinterpret_cast(y), - reinterpret_cast(x1), - reinterpret_cast(x2)); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dist::cpu diff --git a/src/infiniop/ops/dist/cpu/dist_cpu.h b/src/infiniop/ops/dist/cpu/dist_cpu.h deleted file mode 100644 index 1bbbc02f4..000000000 --- a/src/infiniop/ops/dist/cpu/dist_cpu.h +++ /dev/null @@ -1,60 +0,0 @@ -#ifndef __DIST_CPU_H__ -#define __DIST_CPU_H__ - -#include "../../../devices/cpu/common_cpu.h" -#include "../../../operator.h" -#include -#include - -namespace op::dist::cpu { - -struct DistInfo { - size_t input_size; - double p; - std::vector x1_strides; - std::vector x2_strides; - std::vector shape; - size_t ndim; - - static utils::Result create( - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - infiniopTensorDescriptor_t y_desc, - double p); -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - DistInfo _info; - - Descriptor(infiniDtype_t dtype, DistInfo info, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) const; -}; - -} // namespace op::dist::cpu - -#endif // __DIST_CPU_H__ diff --git a/src/infiniop/ops/dist/cuda/kernel.cuh b/src/infiniop/ops/dist/cuda/kernel.cuh deleted file mode 100644 index 903c8d949..000000000 --- a/src/infiniop/ops/dist/cuda/kernel.cuh +++ /dev/null @@ -1,189 +0,0 @@ -#include "../../../reduce/cuda/reduce.cuh" -#include -#include - -struct DistIndexing { - static constexpr int kMaxNdim = 8; - - int ndim; - int64_t shape[kMaxNdim]; - int64_t x1_strides[kMaxNdim]; - int64_t x2_strides[kMaxNdim]; -}; - -namespace op::cuda { - -template -__device__ __forceinline__ float to_f32(T v) { - return static_cast(v); -} - -template <> -__device__ __forceinline__ float to_f32(half v) { - return __half2float(v); -} - -template <> -__device__ __forceinline__ float to_f32(cuda_bfloat16 v) { - return __bfloat162float(v); -} - -template -__device__ __forceinline__ Tdata cast_out(Tcompute v) { - return static_cast(v); -} - -template <> -__device__ __forceinline__ half cast_out(float v) { - return __float2half(v); -} - -template <> -__device__ __forceinline__ cuda_bfloat16 cast_out(float v) { - return __float2bfloat16_rn(v); -} - -template -__global__ void dist_strided_kernel( - Tcompute *result, - const Tdata *x1, - const Tdata *x2, - size_t n, - double p, - DistIndexing indexing) { - - Tcompute thread_val = static_cast(0); - - for (size_t linear = static_cast(threadIdx.x); linear < n; linear += BLOCK_SIZE) { - int64_t idx[DistIndexing::kMaxNdim] = {0}; - size_t tmp = linear; - for (int d = indexing.ndim - 1; d >= 0; --d) { - const int64_t s = indexing.shape[d]; - idx[d] = static_cast(tmp % static_cast(s)); - tmp /= static_cast(s); - } - - int64_t off1 = 0; - int64_t off2 = 0; - for (int d = 0; d < indexing.ndim; ++d) { - off1 += idx[d] * indexing.x1_strides[d]; - off2 += idx[d] * indexing.x2_strides[d]; - } - - Tcompute diff; - if constexpr (std::is_same_v) { - diff = static_cast(x1[off1]) - static_cast(x2[off2]); - } else { - diff = static_cast(to_f32(x1[off1]) - to_f32(x2[off2])); - } - const Tcompute abs_diff = fabs(diff); - - if (p == 0.0) { - if (abs_diff > static_cast(1e-10)) { - thread_val += static_cast(1); - } - } else if (isinf(p)) { - thread_val = fmax(thread_val, abs_diff); - } else { - thread_val += pow(abs_diff, static_cast(p)); - } - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - - if (isinf(p)) { - struct MaxOp { - __device__ __forceinline__ Tcompute operator()(Tcompute a, Tcompute b) const { - return a > b ? a : b; - } - }; - const Tcompute block_max = BlockReduce(temp_storage).Reduce(thread_val, MaxOp{}); - if (threadIdx.x == 0) { - *result = block_max; - } - return; - } - - const Tcompute block_sum = BlockReduce(temp_storage).Sum(thread_val); - if (threadIdx.x == 0) { - if (p == 0.0) { - *result = block_sum; - } else { - *result = pow(block_sum, static_cast(1.0 / p)); - } - } -} - -template -__global__ void dist_strided_out_kernel( - Tdata *out, - const Tdata *x1, - const Tdata *x2, - size_t n, - double p, - DistIndexing indexing) { - - Tcompute thread_val = static_cast(0); - - for (size_t linear = static_cast(threadIdx.x); linear < n; linear += BLOCK_SIZE) { - int64_t idx[DistIndexing::kMaxNdim] = {0}; - size_t tmp = linear; - for (int d = indexing.ndim - 1; d >= 0; --d) { - const int64_t s = indexing.shape[d]; - idx[d] = static_cast(tmp % static_cast(s)); - tmp /= static_cast(s); - } - - int64_t off1 = 0; - int64_t off2 = 0; - for (int d = 0; d < indexing.ndim; ++d) { - off1 += idx[d] * indexing.x1_strides[d]; - off2 += idx[d] * indexing.x2_strides[d]; - } - - Tcompute diff; - if constexpr (std::is_same_v) { - diff = static_cast(x1[off1]) - static_cast(x2[off2]); - } else { - diff = static_cast(to_f32(x1[off1]) - to_f32(x2[off2])); - } - const Tcompute abs_diff = fabs(diff); - - if (p == 0.0) { - if (abs_diff > static_cast(1e-10)) { - thread_val += static_cast(1); - } - } else if (isinf(p)) { - thread_val = fmax(thread_val, abs_diff); - } else { - thread_val += pow(abs_diff, static_cast(p)); - } - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - - if (isinf(p)) { - struct MaxOp { - __device__ __forceinline__ Tcompute operator()(Tcompute a, Tcompute b) const { - return a > b ? a : b; - } - }; - const Tcompute block_max = BlockReduce(temp_storage).Reduce(thread_val, MaxOp{}); - if (threadIdx.x == 0) { - *out = cast_out(block_max); - } - return; - } - - const Tcompute block_sum = BlockReduce(temp_storage).Sum(thread_val); - if (threadIdx.x == 0) { - if (p == 0.0) { - *out = cast_out(block_sum); - } else { - *out = cast_out(pow(block_sum, static_cast(1.0 / p))); - } - } -} -} // namespace op::cuda diff --git a/src/infiniop/ops/dist/metax/dist_metax.h b/src/infiniop/ops/dist/metax/dist_metax.h deleted file mode 100644 index 688683ec7..000000000 --- a/src/infiniop/ops/dist/metax/dist_metax.h +++ /dev/null @@ -1,54 +0,0 @@ -#ifndef __DIST_METAX_H__ -#define __DIST_METAX_H__ - -#include "../../../operator.h" - -namespace op::dist::metax { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _input_size; - double _p; - size_t _ndim; - std::vector _shape; - std::vector _x1_strides; - std::vector _x2_strides; - - Descriptor(infiniDtype_t dtype, size_t input_size, double p, - size_t ndim, std::vector shape, - std::vector x1_strides, std::vector x2_strides, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _input_size(input_size), - _p(p), - _ndim(ndim), - _shape(std::move(shape)), - _x1_strides(std::move(x1_strides)), - _x2_strides(std::move(x2_strides)) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) const; -}; - -} // namespace op::dist::metax - -#endif // __DIST_METAX_H__ diff --git a/src/infiniop/ops/dist/metax/dist_metax.maca b/src/infiniop/ops/dist/metax/dist_metax.maca deleted file mode 100644 index bbc9d7551..000000000 --- a/src/infiniop/ops/dist/metax/dist_metax.maca +++ /dev/null @@ -1,107 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../../../tensor.h" - -#include "../cuda/kernel.cuh" -#include "dist_metax.h" - -namespace op::dist::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p) { - - auto dtype = x1_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto x1_shape = x1_desc->shape(); - auto x2_shape = x2_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x1_shape != x2_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (y_shape.size() != 0 && (y_shape.size() != 1 || y_shape[0] != 1)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const size_t ndim = x1_desc->ndim(); - if (ndim > static_cast(DistIndexing::kMaxNdim)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t input_size = x1_desc->numel(); - *desc_ptr = new Descriptor(dtype, input_size, p, ndim, x1_shape, x1_desc->strides(), x2_desc->strides(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - constexpr int BLOCK_SIZE = 256; - - DistIndexing indexing{}; - indexing.ndim = static_cast(_ndim); - for (int d = 0; d < DistIndexing::kMaxNdim; ++d) { - indexing.shape[d] = 1; - indexing.x1_strides[d] = 0; - indexing.x2_strides[d] = 0; - } - for (size_t d = 0; d < _ndim; ++d) { - indexing.shape[d] = static_cast(_shape[d]); - indexing.x1_strides[d] = static_cast(_x1_strides[d]); - indexing.x2_strides[d] = static_cast(_x2_strides[d]); - } - - switch (_dtype) { - case INFINI_DTYPE_F16: { - cuda::dist_strided_out_kernel<<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(x1), reinterpret_cast(x2), - _input_size, _p, indexing); - break; - } - case INFINI_DTYPE_BF16: { - cuda::dist_strided_out_kernel<<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(x1), reinterpret_cast(x2), - _input_size, _p, indexing); - break; - } - case INFINI_DTYPE_F32: { - float *result_f = reinterpret_cast(y); - cuda::dist_strided_kernel<<<1, BLOCK_SIZE, 0, cuda_stream>>>( - result_f, reinterpret_cast(x1), reinterpret_cast(x2), - _input_size, _p, indexing); - break; - } - case INFINI_DTYPE_F64: { - double *result_d = reinterpret_cast(y); - cuda::dist_strided_kernel<<<1, BLOCK_SIZE, 0, cuda_stream>>>( - result_d, reinterpret_cast(x1), reinterpret_cast(x2), - _input_size, _p, indexing); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dist::metax diff --git a/src/infiniop/ops/dist/moore/dist_moore.h b/src/infiniop/ops/dist/moore/dist_moore.h deleted file mode 100644 index 326d4dd4f..000000000 --- a/src/infiniop/ops/dist/moore/dist_moore.h +++ /dev/null @@ -1,54 +0,0 @@ -#ifndef __DIST_MOORE_H__ -#define __DIST_MOORE_H__ - -#include "../../../operator.h" - -namespace op::dist::moore { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _input_size; - double _p; - size_t _ndim; - std::vector _shape; - std::vector _x1_strides; - std::vector _x2_strides; - - Descriptor(infiniDtype_t dtype, size_t input_size, double p, - size_t ndim, std::vector shape, - std::vector x1_strides, std::vector x2_strides, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _input_size(input_size), - _p(p), - _ndim(ndim), - _shape(std::move(shape)), - _x1_strides(std::move(x1_strides)), - _x2_strides(std::move(x2_strides)) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) const; -}; - -} // namespace op::dist::moore - -#endif // __DIST_MOORE_H__ diff --git a/src/infiniop/ops/dist/moore/dist_moore.mu b/src/infiniop/ops/dist/moore/dist_moore.mu deleted file mode 100644 index e2c1a1b61..000000000 --- a/src/infiniop/ops/dist/moore/dist_moore.mu +++ /dev/null @@ -1,107 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_handle.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "../../../tensor.h" - -#include "../cuda/kernel.cuh" -#include "dist_moore.h" - -namespace op::dist::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p) { - - auto dtype = x1_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto x1_shape = x1_desc->shape(); - auto x2_shape = x2_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x1_shape != x2_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (y_shape.size() != 0 && (y_shape.size() != 1 || y_shape[0] != 1)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const size_t ndim = x1_desc->ndim(); - if (ndim > static_cast(DistIndexing::kMaxNdim)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t input_size = x1_desc->numel(); - *desc_ptr = new Descriptor(dtype, input_size, p, ndim, x1_shape, x1_desc->strides(), x2_desc->strides(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - constexpr int BLOCK_SIZE = 256; - - DistIndexing indexing{}; - indexing.ndim = static_cast(_ndim); - for (int d = 0; d < DistIndexing::kMaxNdim; ++d) { - indexing.shape[d] = 1; - indexing.x1_strides[d] = 0; - indexing.x2_strides[d] = 0; - } - for (size_t d = 0; d < _ndim; ++d) { - indexing.shape[d] = static_cast(_shape[d]); - indexing.x1_strides[d] = static_cast(_x1_strides[d]); - indexing.x2_strides[d] = static_cast(_x2_strides[d]); - } - - switch (_dtype) { - case INFINI_DTYPE_F16: { - cuda::dist_strided_out_kernel<<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(x1), reinterpret_cast(x2), - _input_size, _p, indexing); - break; - } - case INFINI_DTYPE_BF16: { - cuda::dist_strided_out_kernel<<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(x1), reinterpret_cast(x2), - _input_size, _p, indexing); - break; - } - case INFINI_DTYPE_F32: { - float *result_f = reinterpret_cast(y); - cuda::dist_strided_kernel<<<1, BLOCK_SIZE, 0, cuda_stream>>>( - result_f, reinterpret_cast(x1), reinterpret_cast(x2), - _input_size, _p, indexing); - break; - } - case INFINI_DTYPE_F64: { - double *result_d = reinterpret_cast(y); - cuda::dist_strided_kernel<<<1, BLOCK_SIZE, 0, cuda_stream>>>( - result_d, reinterpret_cast(x1), reinterpret_cast(x2), - _input_size, _p, indexing); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dist::moore diff --git a/src/infiniop/ops/dist/nvidia/dist_nvidia.cu b/src/infiniop/ops/dist/nvidia/dist_nvidia.cu deleted file mode 100644 index b652b706d..000000000 --- a/src/infiniop/ops/dist/nvidia/dist_nvidia.cu +++ /dev/null @@ -1,290 +0,0 @@ -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../tensor.h" -#include "dist_nvidia.cuh" -#include -#include -#include -#include - -namespace op::dist::nvidia { - -Descriptor::~Descriptor() = default; - -struct DistIndexing { - static constexpr int kMaxNdim = 8; - - int ndim; - int64_t shape[kMaxNdim]; - int64_t x1_strides[kMaxNdim]; - int64_t x2_strides[kMaxNdim]; -}; - -template -__device__ __forceinline__ float to_f32(T v) { - return static_cast(v); -} - -template <> -__device__ __forceinline__ float to_f32(half v) { - return __half2float(v); -} - -template <> -__device__ __forceinline__ float to_f32(cuda_bfloat16 v) { - return __bfloat162float(v); -} - -template -__device__ __forceinline__ Tdata cast_out(Tcompute v) { - return static_cast(v); -} - -template <> -__device__ __forceinline__ half cast_out(float v) { - return __float2half(v); -} - -template <> -__device__ __forceinline__ cuda_bfloat16 cast_out(float v) { - return __float2bfloat16_rn(v); -} - -template -__global__ void dist_strided_kernel( - Tcompute *result, - const Tdata *x1, - const Tdata *x2, - size_t n, - double p, - DistIndexing indexing) { - - Tcompute thread_val = static_cast(0); - - for (size_t linear = static_cast(threadIdx.x); linear < n; linear += BLOCK_SIZE) { - int64_t idx[DistIndexing::kMaxNdim] = {0}; - size_t tmp = linear; - for (int d = indexing.ndim - 1; d >= 0; --d) { - const int64_t s = indexing.shape[d]; - idx[d] = static_cast(tmp % static_cast(s)); - tmp /= static_cast(s); - } - - int64_t off1 = 0; - int64_t off2 = 0; - for (int d = 0; d < indexing.ndim; ++d) { - off1 += idx[d] * indexing.x1_strides[d]; - off2 += idx[d] * indexing.x2_strides[d]; - } - - Tcompute diff; - if constexpr (std::is_same_v) { - diff = static_cast(x1[off1]) - static_cast(x2[off2]); - } else { - diff = static_cast(to_f32(x1[off1]) - to_f32(x2[off2])); - } - const Tcompute abs_diff = fabs(diff); - - if (p == 0.0) { - if (abs_diff > static_cast(1e-10)) { - thread_val += static_cast(1); - } - } else if (isinf(p)) { - thread_val = fmax(thread_val, abs_diff); - } else { - thread_val += pow(abs_diff, static_cast(p)); - } - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - - if (isinf(p)) { - struct MaxOp { - __device__ __forceinline__ Tcompute operator()(Tcompute a, Tcompute b) const { - return a > b ? a : b; - } - }; - const Tcompute block_max = BlockReduce(temp_storage).Reduce(thread_val, MaxOp{}); - if (threadIdx.x == 0) { - *result = block_max; - } - return; - } - - const Tcompute block_sum = BlockReduce(temp_storage).Sum(thread_val); - if (threadIdx.x == 0) { - if (p == 0.0) { - *result = block_sum; - } else { - *result = pow(block_sum, static_cast(1.0 / p)); - } - } -} - -template -__global__ void dist_strided_out_kernel( - Tdata *out, - const Tdata *x1, - const Tdata *x2, - size_t n, - double p, - DistIndexing indexing) { - - Tcompute thread_val = static_cast(0); - - for (size_t linear = static_cast(threadIdx.x); linear < n; linear += BLOCK_SIZE) { - int64_t idx[DistIndexing::kMaxNdim] = {0}; - size_t tmp = linear; - for (int d = indexing.ndim - 1; d >= 0; --d) { - const int64_t s = indexing.shape[d]; - idx[d] = static_cast(tmp % static_cast(s)); - tmp /= static_cast(s); - } - - int64_t off1 = 0; - int64_t off2 = 0; - for (int d = 0; d < indexing.ndim; ++d) { - off1 += idx[d] * indexing.x1_strides[d]; - off2 += idx[d] * indexing.x2_strides[d]; - } - - Tcompute diff; - if constexpr (std::is_same_v) { - diff = static_cast(x1[off1]) - static_cast(x2[off2]); - } else { - diff = static_cast(to_f32(x1[off1]) - to_f32(x2[off2])); - } - const Tcompute abs_diff = fabs(diff); - - if (p == 0.0) { - if (abs_diff > static_cast(1e-10)) { - thread_val += static_cast(1); - } - } else if (isinf(p)) { - thread_val = fmax(thread_val, abs_diff); - } else { - thread_val += pow(abs_diff, static_cast(p)); - } - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - - if (isinf(p)) { - struct MaxOp { - __device__ __forceinline__ Tcompute operator()(Tcompute a, Tcompute b) const { - return a > b ? a : b; - } - }; - const Tcompute block_max = BlockReduce(temp_storage).Reduce(thread_val, MaxOp{}); - if (threadIdx.x == 0) { - *out = cast_out(block_max); - } - return; - } - - const Tcompute block_sum = BlockReduce(temp_storage).Sum(thread_val); - if (threadIdx.x == 0) { - if (p == 0.0) { - *out = cast_out(block_sum); - } else { - *out = cast_out(pow(block_sum, static_cast(1.0 / p))); - } - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p) { - - auto dtype = x1_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto x1_shape = x1_desc->shape(); - auto x2_shape = x2_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x1_shape != x2_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (y_shape.size() != 0 && (y_shape.size() != 1 || y_shape[0] != 1)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const size_t ndim = x1_desc->ndim(); - if (ndim > static_cast(DistIndexing::kMaxNdim)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t input_size = x1_desc->numel(); - *desc_ptr = new Descriptor(dtype, input_size, p, ndim, x1_shape, x1_desc->strides(), x2_desc->strides(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - constexpr int BLOCK_SIZE = 256; - - DistIndexing indexing{}; - indexing.ndim = static_cast(_ndim); - for (int d = 0; d < DistIndexing::kMaxNdim; ++d) { - indexing.shape[d] = 1; - indexing.x1_strides[d] = 0; - indexing.x2_strides[d] = 0; - } - for (size_t d = 0; d < _ndim; ++d) { - indexing.shape[d] = static_cast(_shape[d]); - indexing.x1_strides[d] = static_cast(_x1_strides[d]); - indexing.x2_strides[d] = static_cast(_x2_strides[d]); - } - - switch (_dtype) { - case INFINI_DTYPE_F16: { - dist_strided_out_kernel<<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(x1), reinterpret_cast(x2), - _input_size, _p, indexing); - break; - } - case INFINI_DTYPE_BF16: { - dist_strided_out_kernel<<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(x1), reinterpret_cast(x2), - _input_size, _p, indexing); - break; - } - case INFINI_DTYPE_F32: { - float *result_f = reinterpret_cast(y); - dist_strided_kernel<<<1, BLOCK_SIZE, 0, cuda_stream>>>( - result_f, reinterpret_cast(x1), reinterpret_cast(x2), - _input_size, _p, indexing); - break; - } - case INFINI_DTYPE_F64: { - double *result_d = reinterpret_cast(y); - dist_strided_kernel<<<1, BLOCK_SIZE, 0, cuda_stream>>>( - result_d, reinterpret_cast(x1), reinterpret_cast(x2), - _input_size, _p, indexing); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dist::nvidia diff --git a/src/infiniop/ops/dist/nvidia/dist_nvidia.cuh b/src/infiniop/ops/dist/nvidia/dist_nvidia.cuh deleted file mode 100644 index 9e99dc759..000000000 --- a/src/infiniop/ops/dist/nvidia/dist_nvidia.cuh +++ /dev/null @@ -1,57 +0,0 @@ -#ifndef __DIST_NVIDIA_H__ -#define __DIST_NVIDIA_H__ - -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../operator.h" -#include -#include - -namespace op::dist::nvidia { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _input_size; - double _p; - size_t _ndim; - std::vector _shape; - std::vector _x1_strides; - std::vector _x2_strides; - - Descriptor(infiniDtype_t dtype, size_t input_size, double p, - size_t ndim, std::vector shape, - std::vector x1_strides, std::vector x2_strides, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _input_size(input_size), - _p(p), - _ndim(ndim), - _shape(std::move(shape)), - _x1_strides(std::move(x1_strides)), - _x2_strides(std::move(x2_strides)) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) const; -}; - -} // namespace op::dist::nvidia - -#endif // __DIST_NVIDIA_H__ diff --git a/src/infiniop/ops/dist/operator.cc b/src/infiniop/ops/dist/operator.cc deleted file mode 100644 index 6aba56b88..000000000 --- a/src/infiniop/ops/dist/operator.cc +++ /dev/null @@ -1,162 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/dist.h" - -#ifdef ENABLE_CPU_API -#include "cpu/dist_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) -#include "nvidia/dist_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/dist_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/dist_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateDistDescriptor( - infiniopHandle_t handle, - infiniopDistDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x1_desc, - infiniopTensorDescriptor_t x2_desc, - double p) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::dist::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - x1_desc, \ - x2_desc, \ - p) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetDistWorkspaceSize(infiniopDistDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopDist( - infiniopDistDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x1, - const void *x2, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, x1, x2, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyDistDescriptor(infiniopDistDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/dot/cpu/dot_cpu.cc b/src/infiniop/ops/dot/cpu/dot_cpu.cc deleted file mode 100644 index 2a45bc971..000000000 --- a/src/infiniop/ops/dot/cpu/dot_cpu.cc +++ /dev/null @@ -1,105 +0,0 @@ -#include "dot_cpu.h" -#include "../../../../utils.h" - -namespace op::dot::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - - auto dtype = a_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - // This op does not do implicit dtype conversion: y/a/b must match. - if (b_desc->dtype() != dtype || y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // Check that y is a scalar (0D tensor or shape [1]) - auto y_shape = y_desc->shape(); - if (y_shape.size() != 0 && (y_shape.size() != 1 || y_shape[0] != 1)) { - return INFINI_STATUS_BAD_PARAM; - } - - // Check that a and b are 1D vectors with same length - auto a_shape = a_desc->shape(); - auto b_shape = b_desc->shape(); - if (a_shape.size() != 1 || b_shape.size() != 1 || a_shape[0] != b_shape[0]) { - return INFINI_STATUS_BAD_PARAM; - } - - size_t n = a_shape[0]; - ptrdiff_t a_stride = a_desc->strides()[0]; - ptrdiff_t b_stride = b_desc->strides()[0]; - - // Negative/broadcasted strides are not supported without an explicit base offset. - if (a_stride <= 0 || b_stride <= 0) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - *desc_ptr = new Descriptor(dtype, n, a_stride, b_stride, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: { - const fp16_t *a_ptr = reinterpret_cast(a); - const fp16_t *b_ptr = reinterpret_cast(b); - float result = 0.0f; - for (size_t i = 0; i < _n; ++i) { - result += utils::cast(a_ptr[i * _a_stride]) * utils::cast(b_ptr[i * _b_stride]); - } - *reinterpret_cast(y) = utils::cast(result); - break; - } - case INFINI_DTYPE_BF16: { - const bf16_t *a_ptr = reinterpret_cast(a); - const bf16_t *b_ptr = reinterpret_cast(b); - float result = 0.0f; - for (size_t i = 0; i < _n; ++i) { - result += utils::cast(a_ptr[i * _a_stride]) * utils::cast(b_ptr[i * _b_stride]); - } - *reinterpret_cast(y) = utils::cast(result); - break; - } - case INFINI_DTYPE_F32: { - const float *a_ptr = reinterpret_cast(a); - const float *b_ptr = reinterpret_cast(b); - float result = 0.0f; - for (size_t i = 0; i < _n; ++i) { - result += a_ptr[i * _a_stride] * b_ptr[i * _b_stride]; - } - *reinterpret_cast(y) = result; - break; - } - case INFINI_DTYPE_F64: { - const double *a_ptr = reinterpret_cast(a); - const double *b_ptr = reinterpret_cast(b); - double result = 0.0; - for (size_t i = 0; i < _n; ++i) { - result += a_ptr[i * _a_stride] * b_ptr[i * _b_stride]; - } - *reinterpret_cast(y) = result; - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dot::cpu diff --git a/src/infiniop/ops/dot/cpu/dot_cpu.h b/src/infiniop/ops/dot/cpu/dot_cpu.h deleted file mode 100644 index 95a7a4d12..000000000 --- a/src/infiniop/ops/dot/cpu/dot_cpu.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef __DOT_CPU_H__ -#define __DOT_CPU_H__ - -#include "../../../devices/cpu/common_cpu.h" -#include "../../../operator.h" -#include "../../../reduce/cpu/reduce.h" -#include "../../../tensor.h" -#include - -namespace op::dot::cpu { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _n; - ptrdiff_t _a_stride; - ptrdiff_t _b_stride; - - Descriptor(infiniDtype_t dtype, size_t n, ptrdiff_t a_stride, ptrdiff_t b_stride, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _n(n), - _a_stride(a_stride), - _b_stride(b_stride) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const; -}; - -} // namespace op::dot::cpu - -#endif // __DOT_CPU_H__ diff --git a/src/infiniop/ops/dot/cuda/kernel.cuh b/src/infiniop/ops/dot/cuda/kernel.cuh deleted file mode 100644 index e044bf523..000000000 --- a/src/infiniop/ops/dot/cuda/kernel.cuh +++ /dev/null @@ -1,33 +0,0 @@ -#pragma once - -#include "../../../reduce/cuda/reduce.cuh" -#include - -namespace op::cuda { - -template -__global__ void dot_kernel( - Tout *result, - const Tdata *a, - const Tdata *b, - size_t n, - ptrdiff_t a_stride, - ptrdiff_t b_stride) { - - Tcompute sum = 0; - - for (size_t i = threadIdx.x; i < n; i += BLOCK_SIZE) { - sum += Tcompute(a[i * a_stride]) * Tcompute(b[i * b_stride]); - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - - Tcompute block_sum = BlockReduce(temp_storage).Sum(sum); - - if (threadIdx.x == 0) { - result[0] = static_cast(block_sum); - } -} - -} // namespace op::cuda diff --git a/src/infiniop/ops/dot/metax/dot_metax.h b/src/infiniop/ops/dot/metax/dot_metax.h deleted file mode 100644 index c9e49f636..000000000 --- a/src/infiniop/ops/dot/metax/dot_metax.h +++ /dev/null @@ -1,46 +0,0 @@ -#ifndef __DOT_METAX_H__ -#define __DOT_METAX_H__ - -#include "../../../operator.h" -#include "../../../tensor.h" - -namespace op::dot::metax { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _n; - ptrdiff_t _a_stride; - ptrdiff_t _b_stride; - - Descriptor(infiniDtype_t dtype, size_t n, ptrdiff_t a_stride, ptrdiff_t b_stride, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _n(n), - _a_stride(a_stride), - _b_stride(b_stride) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const; -}; - -} // namespace op::dot::metax - -#endif // __DOT_METAX_H__ diff --git a/src/infiniop/ops/dot/metax/dot_metax.maca b/src/infiniop/ops/dot/metax/dot_metax.maca deleted file mode 100644 index 38c99536f..000000000 --- a/src/infiniop/ops/dot/metax/dot_metax.maca +++ /dev/null @@ -1,98 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" - -#include "../cuda/kernel.cuh" -#include "dot_metax.h" - -namespace op::dot::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - - auto dtype = a_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - // Check that y is a scalar (0D tensor or shape [1]) - auto y_shape = y_desc->shape(); - if (y_shape.size() != 0 && (y_shape.size() != 1 || y_shape[0] != 1)) { - return INFINI_STATUS_BAD_PARAM; - } - - // Check that a and b are 1D vectors with same length - auto a_shape = a_desc->shape(); - auto b_shape = b_desc->shape(); - if (a_shape.size() != 1 || b_shape.size() != 1 || a_shape[0] != b_shape[0]) { - return INFINI_STATUS_BAD_PARAM; - } - - size_t n = a_shape[0]; - ptrdiff_t a_stride = a_desc->strides()[0]; - ptrdiff_t b_stride = b_desc->strides()[0]; - - *desc_ptr = new Descriptor(dtype, n, a_stride, b_stride, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const { - - auto hc_stream = reinterpret_cast(stream); - constexpr unsigned int BLOCK_SIZE = 256; - - switch (_dtype) { - - case INFINI_DTYPE_F32: - cuda::dot_kernel - <<<1, BLOCK_SIZE, 0, hc_stream>>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - _n, _a_stride, _b_stride); - break; - - case INFINI_DTYPE_F64: - cuda::dot_kernel - <<<1, BLOCK_SIZE, 0, hc_stream>>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - _n, _a_stride, _b_stride); - break; - - case INFINI_DTYPE_F16: - cuda::dot_kernel - <<<1, BLOCK_SIZE, 0, hc_stream>>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - _n, _a_stride, _b_stride); - break; - - case INFINI_DTYPE_BF16: - cuda::dot_kernel - <<<1, BLOCK_SIZE, 0, hc_stream>>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - _n, _a_stride, _b_stride); - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dot::metax diff --git a/src/infiniop/ops/dot/moore/dot_moore.h b/src/infiniop/ops/dot/moore/dot_moore.h deleted file mode 100644 index b11178e15..000000000 --- a/src/infiniop/ops/dot/moore/dot_moore.h +++ /dev/null @@ -1,46 +0,0 @@ -#ifndef __DOT_MOORE_H__ -#define __DOT_MOORE_H__ - -#include "../../../operator.h" -#include "../../../tensor.h" - -namespace op::dot::moore { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _n; - ptrdiff_t _a_stride; - ptrdiff_t _b_stride; - - Descriptor(infiniDtype_t dtype, size_t n, ptrdiff_t a_stride, ptrdiff_t b_stride, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _n(n), - _a_stride(a_stride), - _b_stride(b_stride) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const; -}; - -} // namespace op::dot::moore - -#endif // __DOT_MOORE_H__ diff --git a/src/infiniop/ops/dot/moore/dot_moore.mu b/src/infiniop/ops/dot/moore/dot_moore.mu deleted file mode 100644 index 7516f64bb..000000000 --- a/src/infiniop/ops/dot/moore/dot_moore.mu +++ /dev/null @@ -1,99 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" - -#include "../cuda/kernel.cuh" -#include "dot_moore.h" - -namespace op::dot::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - - auto dtype = a_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - // Check that y is a scalar (0D tensor or shape [1]) - auto y_shape = y_desc->shape(); - if (y_shape.size() != 0 && (y_shape.size() != 1 || y_shape[0] != 1)) { - return INFINI_STATUS_BAD_PARAM; - } - - // Check that a and b are 1D vectors with same length - auto a_shape = a_desc->shape(); - auto b_shape = b_desc->shape(); - if (a_shape.size() != 1 || b_shape.size() != 1 || a_shape[0] != b_shape[0]) { - return INFINI_STATUS_BAD_PARAM; - } - - size_t n = a_shape[0]; - ptrdiff_t a_stride = a_desc->strides()[0]; - ptrdiff_t b_stride = b_desc->strides()[0]; - - *desc_ptr = new Descriptor(dtype, n, a_stride, b_stride, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - constexpr unsigned int BLOCK_SIZE = 256; - - switch (_dtype) { - - case INFINI_DTYPE_F32: - cuda::dot_kernel - <<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - _n, _a_stride, _b_stride); - break; - - case INFINI_DTYPE_F64: - cuda::dot_kernel - <<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - _n, _a_stride, _b_stride); - break; - - case INFINI_DTYPE_F16: - cuda::dot_kernel - <<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - _n, _a_stride, _b_stride); - break; - - case INFINI_DTYPE_BF16: - cuda::dot_kernel - <<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - _n, _a_stride, _b_stride); - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dot::moore diff --git a/src/infiniop/ops/dot/nvidia/dot_nvidia.cu b/src/infiniop/ops/dot/nvidia/dot_nvidia.cu deleted file mode 100644 index d3a2e6f8e..000000000 --- a/src/infiniop/ops/dot/nvidia/dot_nvidia.cu +++ /dev/null @@ -1,121 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" - -#include "../cuda/kernel.cuh" -#include "dot_nvidia.cuh" - -namespace op::dot::nvidia { - -__global__ void store_half_from_f32(half *dst, const float *src) { - if (threadIdx.x == 0) { - dst[0] = __float2half(src[0]); - } -} - -__global__ void store_bf16_from_f32(cuda_bfloat16 *dst, const float *src) { - if (threadIdx.x == 0) { - dst[0] = __float2bfloat16_rn(src[0]); - } -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - - auto dtype = a_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - // This op does not do implicit dtype conversion: y/a/b must match. - if (b_desc->dtype() != dtype || y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // Check that y is a scalar (0D tensor or shape [1]) - auto y_shape = y_desc->shape(); - if (y_shape.size() != 0 && (y_shape.size() != 1 || y_shape[0] != 1)) { - return INFINI_STATUS_BAD_PARAM; - } - - // Check that a and b are 1D vectors with same length - auto a_shape = a_desc->shape(); - auto b_shape = b_desc->shape(); - if (a_shape.size() != 1 || b_shape.size() != 1 || a_shape[0] != b_shape[0]) { - return INFINI_STATUS_BAD_PARAM; - } - - size_t n = a_shape[0]; - ptrdiff_t a_stride = a_desc->strides()[0]; - ptrdiff_t b_stride = b_desc->strides()[0]; - - // Negative/broadcasted strides are not supported without an explicit base offset. - if (a_stride <= 0 || b_stride <= 0) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - *desc_ptr = new Descriptor(dtype, n, a_stride, b_stride, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - constexpr unsigned int BLOCK_SIZE = 256; - - switch (_dtype) { - - case INFINI_DTYPE_F32: - cuda::dot_kernel - <<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - _n, _a_stride, _b_stride); - break; - - case INFINI_DTYPE_F64: - cuda::dot_kernel - <<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - _n, _a_stride, _b_stride); - break; - - case INFINI_DTYPE_F16: - cuda::dot_kernel - <<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - _n, _a_stride, _b_stride); - break; - - case INFINI_DTYPE_BF16: - cuda::dot_kernel - <<<1, BLOCK_SIZE, 0, cuda_stream>>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - _n, _a_stride, _b_stride); - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::dot::nvidia diff --git a/src/infiniop/ops/dot/nvidia/dot_nvidia.cuh b/src/infiniop/ops/dot/nvidia/dot_nvidia.cuh deleted file mode 100644 index 8a15e50f1..000000000 --- a/src/infiniop/ops/dot/nvidia/dot_nvidia.cuh +++ /dev/null @@ -1,48 +0,0 @@ -#ifndef __DOT_NVIDIA_H__ -#define __DOT_NVIDIA_H__ - -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../operator.h" -#include "../../../tensor.h" -#include - -namespace op::dot::nvidia { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _n; - ptrdiff_t _a_stride; - ptrdiff_t _b_stride; - - Descriptor(infiniDtype_t dtype, size_t n, ptrdiff_t a_stride, ptrdiff_t b_stride, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _n(n), - _a_stride(a_stride), - _b_stride(b_stride) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const; -}; - -} // namespace op::dot::nvidia - -#endif // __DOT_NVIDIA_H__ diff --git a/src/infiniop/ops/dot/operator.cc b/src/infiniop/ops/dot/operator.cc deleted file mode 100644 index 10f9fda57..000000000 --- a/src/infiniop/ops/dot/operator.cc +++ /dev/null @@ -1,160 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/dot.h" - -#ifdef ENABLE_CPU_API -#include "cpu/dot_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) -#include "nvidia/dot_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/dot_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/dot_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateDotDescriptor( - infiniopHandle_t handle, - infiniopDotDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::dot::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - a_desc, \ - b_desc) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetDotWorkspaceSize(infiniopDotDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopDot( - infiniopDotDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, a, b, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyDotDescriptor(infiniopDotDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/embedding/ascend/embedding_ascend.cc b/src/infiniop/ops/embedding/ascend/embedding_ascend.cc deleted file mode 100644 index d54268d9e..000000000 --- a/src/infiniop/ops/embedding/ascend/embedding_ascend.cc +++ /dev/null @@ -1,105 +0,0 @@ -#include "embedding_ascend.h" -#include "../../../devices/ascend/common_ascend.h" -#include - -namespace op::embedding::ascend { - -struct Descriptor::Opaque { - aclnnTensorDescriptor_t output; - aclnnTensorDescriptor_t input; - aclnnTensorDescriptor_t weight; - void *workspace; - uint64_t workspace_size; - aclOpExecutor *executor; - - ~Opaque() { - delete output; - delete input; - delete weight; - if (workspace != nullptr) { - aclrtFree(workspace); - } - aclDestroyAclOpExecutor(executor); - } -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t weight_desc) { - - auto handle = reinterpret_cast(handle_); - - CHECK_API_OR(input_desc->dtype() == INFINI_DTYPE_I32 || input_desc->dtype() == INFINI_DTYPE_I64, true, - return INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_API_OR(output_desc->dtype() == weight_desc->dtype(), true, return INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_API_OR(weight_desc->ndim() == 2, true, return INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_API_OR(output_desc->ndim() == input_desc->ndim() + 1, true, return INFINI_STATUS_BAD_TENSOR_SHAPE); - - auto input_shape = input_desc->shape(); - auto output_shape = output_desc->shape(); - auto weight_shape = weight_desc->shape(); - for (size_t i = 0; i < input_desc->ndim(); ++i) { - CHECK_API_OR(output_shape[i] == input_shape[i], true, return INFINI_STATUS_BAD_TENSOR_SHAPE); - } - CHECK_API_OR(output_shape.back() == weight_shape[1], true, return INFINI_STATUS_BAD_TENSOR_SHAPE); - - size_t num_indices = 1; - for (auto dim : input_shape) { - num_indices *= dim; - } - - auto output = new aclnnTensorDescriptor(output_desc); - auto input = new aclnnTensorDescriptor(input_desc); - auto weight = new aclnnTensorDescriptor(weight_desc); - - uint64_t workspace_size = 0; - aclOpExecutor *executor = nullptr; - CHECK_ACL(aclnnEmbeddingGetWorkspaceSize(weight->tensor, input->tensor, output->tensor, - &workspace_size, &executor)); - aclSetAclOpExecutorRepeatable(executor); - - void *workspace = nullptr; - if (workspace_size != 0) { - CHECK_ACL(aclrtMalloc(&workspace, workspace_size, ACL_MEM_MALLOC_HUGE_FIRST)); - } - - *desc_ptr = new Descriptor( - num_indices, - weight_shape[1], - weight_shape[0], - input_desc->dtype(), - weight_desc->dtype(), - new Opaque{output, input, weight, workspace, workspace_size, executor}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *output, - const void *input, - const void *weight, - void *stream) const { - - auto tweight = _opaque->weight->tensor; - auto tinput = _opaque->input->tensor; - auto toutput = _opaque->output->tensor; - - AclSetTensorAddr(_opaque->executor, 0, tweight, const_cast(weight)); - AclSetTensorAddr(_opaque->executor, 1, tinput, const_cast(input)); - AclSetTensorAddr(_opaque->executor, 2, toutput, output); - - CHECK_ACL(aclnnEmbedding(_opaque->workspace, _opaque->workspace_size, - _opaque->executor, stream)); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::embedding::ascend diff --git a/src/infiniop/ops/embedding/ascend/embedding_ascend.h b/src/infiniop/ops/embedding/ascend/embedding_ascend.h deleted file mode 100644 index 7653a5866..000000000 --- a/src/infiniop/ops/embedding/ascend/embedding_ascend.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __EMBEDDING_ASCEND_H__ -#define __EMBEDDING_ASCEND_H__ - -#include "../embedding.h" - -DESCRIPTOR(ascend) - -#endif // __EMBEDDING_ASCEND_H__ diff --git a/src/infiniop/ops/embedding/bang/embedding_bang.h b/src/infiniop/ops/embedding/bang/embedding_bang.h deleted file mode 100644 index bc665a1a3..000000000 --- a/src/infiniop/ops/embedding/bang/embedding_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __EMBEDDING_BANG_H__ -#define __EMBEDDING_BANG_H__ - -#include "../embedding.h" - -DESCRIPTOR(bang) - -#endif // __EMBEDDING_BANG_H__ diff --git a/src/infiniop/ops/embedding/bang/embedding_bang.mlu b/src/infiniop/ops/embedding/bang/embedding_bang.mlu deleted file mode 100644 index 01ff3812d..000000000 --- a/src/infiniop/ops/embedding/bang/embedding_bang.mlu +++ /dev/null @@ -1,182 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "embedding_bang.h" - -#include - -__nram__ char nram_buffer[NRAM_MAX_SIZE]; - -template -__mlu_global__ void embeddingKernel( - T *__restrict__ output, - const IndexType *__restrict__ indices, - const T *__restrict__ weight, - size_t num_indices, - size_t embedding_dim, - size_t vocab_size) { - if (num_indices == 0 || embedding_dim == 0) { - return; - } - - size_t chunk_target = 512; - size_t parts_by_dim = (embedding_dim + chunk_target - 1) / chunk_target; - size_t max_parts_by_tasks = std::max(1, taskDim / num_indices); - size_t parts_per_row = std::max(1, std::min(parts_by_dim, max_parts_by_tasks)); - size_t part_size = (embedding_dim + parts_per_row - 1) / parts_per_row; - size_t logical_tasks = num_indices * parts_per_row; - size_t max_chunk = NRAM_MAX_SIZE / sizeof(T); - T *cache = reinterpret_cast(nram_buffer); - - for (size_t logical = taskId; logical < logical_tasks; logical += taskDim) { - size_t row = logical / parts_per_row; - size_t part = logical - row * parts_per_row; - size_t start = part * part_size; - size_t end = std::min(embedding_dim, start + part_size); - if (start >= end) { - continue; - } - - IndexType index_val = indices[row]; - if (index_val < 0 || static_cast(index_val) >= vocab_size) { - continue; - } - - const T *src = weight + static_cast(index_val) * embedding_dim + start; - T *dst = output + row * embedding_dim + start; - size_t processed = 0; - size_t len = end - start; - while (processed < len) { - size_t current = std::min(max_chunk, len - processed); - __memcpy(cache, src + processed, current * sizeof(T), GDRAM2NRAM); - __memcpy(dst + processed, cache, current * sizeof(T), NRAM2GDRAM); - processed += current; - } - } -} - -template -static infiniStatus_t launchEmbedding( - int core_per_cluster, - int cluster_count, - cnrtQueue_t queue, - void *output, - const void *input, - const void *weight, - size_t num_indices, - size_t embedding_dim, - size_t vocab_size) { - cnrtDim3_t kernel_dim; - kernel_dim.x = core_per_cluster; - kernel_dim.y = cluster_count; - kernel_dim.z = 1; - - embeddingKernel<<>>( - reinterpret_cast(output), - reinterpret_cast(input), - reinterpret_cast(weight), - num_indices, - embedding_dim, - vocab_size); - cnrtQueueSync(queue); - return INFINI_STATUS_SUCCESS; -} - -namespace op::embedding::bang { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t weight_desc) { - - auto input_shape = input_desc->shape(); - auto weight_shape = weight_desc->shape(); - - CHECK_OR_RETURN(weight_shape.size() == 2, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(output_desc->shape().size() == input_shape.size() + 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - - auto output_shape = output_desc->shape(); - size_t embedding_dim = weight_shape[1]; - CHECK_OR_RETURN(output_shape.back() == embedding_dim, INFINI_STATUS_BAD_TENSOR_SHAPE); - - for (size_t i = 0; i < input_shape.size(); ++i) { - CHECK_OR_RETURN(output_shape[i] == input_shape[i], INFINI_STATUS_BAD_TENSOR_SHAPE); - } - - auto input_dtype = input_desc->dtype(); - auto weight_dtype = weight_desc->dtype(); - CHECK_OR_RETURN(input_dtype == INFINI_DTYPE_I32 || input_dtype == INFINI_DTYPE_I64, - INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(weight_dtype == INFINI_DTYPE_F32 || weight_dtype == INFINI_DTYPE_F16 || weight_dtype == INFINI_DTYPE_BF16, - INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(output_desc->dtype() == weight_dtype, INFINI_STATUS_BAD_TENSOR_DTYPE); - - size_t num_indices = 1; - for (auto dim : input_shape) { - num_indices *= dim; - } - size_t vocab_size = weight_shape[0]; - - auto handle_bang = reinterpret_cast(handle); - *desc_ptr = new Descriptor( - num_indices, - embedding_dim, - vocab_size, - input_dtype, - weight_dtype, - new Opaque{handle_bang->internal()}, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *output, - const void *input, - const void *weight, - void *stream) const { - if (_num_indices == 0) { - return INFINI_STATUS_SUCCESS; - } - - auto queue = reinterpret_cast(stream); - int core_per_cluster = _opaque->internal->getCorePerCluster(); - int cluster_count = _opaque->internal->getClusterCount(); - -#define DISPATCH(T, IndexType) \ - return launchEmbedding( \ - core_per_cluster, cluster_count, queue, output, input, weight, \ - _num_indices, _embedding_dim, _vocab_size) - - if (_input_dtype == INFINI_DTYPE_I32) { - if (_weight_dtype == INFINI_DTYPE_F32) { - DISPATCH(float, int32_t); - } else if (_weight_dtype == INFINI_DTYPE_F16) { - DISPATCH(half, int32_t); - } else if (_weight_dtype == INFINI_DTYPE_BF16) { - DISPATCH(bfloat16_t, int32_t); - } - } else if (_input_dtype == INFINI_DTYPE_I64) { - if (_weight_dtype == INFINI_DTYPE_F32) { - DISPATCH(float, int64_t); - } else if (_weight_dtype == INFINI_DTYPE_F16) { - DISPATCH(half, int64_t); - } else if (_weight_dtype == INFINI_DTYPE_BF16) { - DISPATCH(bfloat16_t, int64_t); - } - } - -#undef DISPATCH - - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -} // namespace op::embedding::bang diff --git a/src/infiniop/ops/embedding/cpu/embedding_cpu.cc b/src/infiniop/ops/embedding/cpu/embedding_cpu.cc deleted file mode 100644 index 8e6648063..000000000 --- a/src/infiniop/ops/embedding/cpu/embedding_cpu.cc +++ /dev/null @@ -1,109 +0,0 @@ -#include "embedding_cpu.h" -#include "../../../../utils.h" -#include "../../../handle.h" -#include "../../../tensor.h" -#include - -namespace op::embedding::cpu { - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t weight_desc) { - - auto input_shape = input_desc->shape(); - auto weight_shape = weight_desc->shape(); - - CHECK_OR_RETURN(weight_shape.size() == 2, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(output_desc->shape().size() == input_shape.size() + 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - - auto output_shape = output_desc->shape(); - size_t embedding_dim = weight_shape[1]; - CHECK_OR_RETURN(output_shape.back() == embedding_dim, INFINI_STATUS_BAD_TENSOR_SHAPE); - - for (size_t i = 0; i < input_shape.size(); ++i) { - CHECK_OR_RETURN(output_shape[i] == input_shape[i], INFINI_STATUS_BAD_TENSOR_SHAPE); - } - - auto input_dtype = input_desc->dtype(); - auto weight_dtype = weight_desc->dtype(); - CHECK_OR_RETURN(input_dtype == INFINI_DTYPE_I32 || input_dtype == INFINI_DTYPE_I64, - INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(weight_dtype == INFINI_DTYPE_F32 || weight_dtype == INFINI_DTYPE_F16 || weight_dtype == INFINI_DTYPE_BF16, INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(output_desc->dtype() == weight_dtype, INFINI_STATUS_BAD_TENSOR_DTYPE); - - size_t num_indices = 1; - for (auto dim : input_shape) { - num_indices *= dim; - } - - size_t vocab_size = weight_shape[0]; - - *desc_ptr = new Descriptor( - num_indices, - embedding_dim, - vocab_size, - input_dtype, - weight_dtype, - new Opaque{}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *output, - const void *input, - const void *weight, - void *stream) const { - - if (_num_indices == 0) { - return INFINI_STATUS_SUCCESS; - } - - size_t element_size = infiniSizeOf(_weight_dtype); - size_t row_bytes = _embedding_dim * element_size; - - if (_input_dtype == INFINI_DTYPE_I32) { - const int32_t *indices_ptr = reinterpret_cast(input); - const std::byte *weight_ptr = reinterpret_cast(weight); - std::byte *out_ptr = reinterpret_cast(output); - - for (size_t i = 0; i < _num_indices; ++i) { - int32_t idx = indices_ptr[i]; - if (idx >= 0 && static_cast(idx) < _vocab_size) { - std::memcpy(out_ptr + i * row_bytes, - weight_ptr + static_cast(idx) * row_bytes, - row_bytes); - } - } - } else if (_input_dtype == INFINI_DTYPE_I64) { - const int64_t *indices_ptr = reinterpret_cast(input); - const std::byte *weight_ptr = reinterpret_cast(weight); - std::byte *out_ptr = reinterpret_cast(output); - - for (size_t i = 0; i < _num_indices; ++i) { - int64_t idx = indices_ptr[i]; - if (idx >= 0 && static_cast(idx) < _vocab_size) { - std::memcpy(out_ptr + i * row_bytes, - weight_ptr + static_cast(idx) * row_bytes, - row_bytes); - } - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::embedding::cpu diff --git a/src/infiniop/ops/embedding/cpu/embedding_cpu.h b/src/infiniop/ops/embedding/cpu/embedding_cpu.h deleted file mode 100644 index a5cc5b2d0..000000000 --- a/src/infiniop/ops/embedding/cpu/embedding_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __EMBEDDING_CPU_H__ -#define __EMBEDDING_CPU_H__ - -#include "../embedding.h" - -DESCRIPTOR(cpu) - -#endif // __EMBEDDING_CPU_H__ diff --git a/src/infiniop/ops/embedding/cuda/embedding_kernel.cuh b/src/infiniop/ops/embedding/cuda/embedding_kernel.cuh deleted file mode 100644 index 2914f06ed..000000000 --- a/src/infiniop/ops/embedding/cuda/embedding_kernel.cuh +++ /dev/null @@ -1,116 +0,0 @@ -#ifndef __EMBEDDING_CUDA_KERNEL_CUH__ -#define __EMBEDDING_CUDA_KERNEL_CUH__ - -#include - -// Helper function to check memory alignment -__forceinline__ __device__ bool is_aligned(const void *ptr, size_t alignment) { - // Use size_t for pointer arithmetic in device code (more compatible) - return (reinterpret_cast(ptr) % alignment == 0); -} - -// Vectorized copy for float type using float4 -template -__forceinline__ __device__ void copyVectorizedFloat4( - float *__restrict__ dst, - const float *__restrict__ src, - size_t embedding_dim) { - // Use float4 for vectorized access (16 bytes, 4 floats) - const float4 *src_vec = reinterpret_cast(src); - float4 *dst_vec = reinterpret_cast(dst); - size_t vec_count = embedding_dim / 4; - - // Vectorized copy using __ldg for read-only weight - for (size_t i = 0; i < vec_count; ++i) { - dst_vec[i] = __ldg(&src_vec[i]); - } - - // Copy remaining elements - size_t remaining = embedding_dim % 4; - if (remaining > 0) { - size_t offset = vec_count * 4; - for (size_t i = 0; i < remaining; ++i) { - dst[offset + i] = __ldg(&src[offset + i]); - } - } -} - -// Vectorized copy for float type using float2 (fallback when not aligned to 16 bytes) -template -__forceinline__ __device__ void copyVectorizedFloat2( - float *__restrict__ dst, - const float *__restrict__ src, - size_t embedding_dim) { - // Use float2 for vectorized access (8 bytes, 2 floats) - const float2 *src_vec = reinterpret_cast(src); - float2 *dst_vec = reinterpret_cast(dst); - size_t vec_count = embedding_dim / 2; - - // Vectorized copy using __ldg for read-only weight - for (size_t i = 0; i < vec_count; ++i) { - dst_vec[i] = __ldg(&src_vec[i]); - } - - // Copy remaining element if odd - if (embedding_dim % 2 != 0) { - dst[embedding_dim - 1] = __ldg(&src[embedding_dim - 1]); - } -} - -// Vectorized copy for half type using half2 -template -__forceinline__ __device__ void copyVectorizedHalf2( - half *__restrict__ dst, - const half *__restrict__ src, - size_t embedding_dim) { - // Use half2 for vectorized access (4 bytes, 2 halfs) - const half2 *src_vec = reinterpret_cast(src); - half2 *dst_vec = reinterpret_cast(dst); - size_t vec_count = embedding_dim / 2; - - // Vectorized copy using __ldg for read-only weight - for (size_t i = 0; i < vec_count; ++i) { - dst_vec[i] = __ldg(&src_vec[i]); - } - - // Copy remaining element if odd - if (embedding_dim % 2 != 0) { - dst[embedding_dim - 1] = __ldg(&src[embedding_dim - 1]); - } -} - -// Vectorized copy for bfloat16 type using bfloat162 -template -__forceinline__ __device__ void copyVectorizedBFloat162( - cuda_bfloat16 *__restrict__ dst, - const cuda_bfloat16 *__restrict__ src, - size_t embedding_dim) { - // Use bfloat162 for vectorized access (4 bytes, 2 bfloat16s) - const cuda_bfloat162 *src_vec = reinterpret_cast(src); - cuda_bfloat162 *dst_vec = reinterpret_cast(dst); - size_t vec_count = embedding_dim / 2; - - // Vectorized copy using __ldg for read-only weight - for (size_t i = 0; i < vec_count; ++i) { - dst_vec[i] = __ldg(&src_vec[i]); - } - - // Copy remaining element if odd - if (embedding_dim % 2 != 0) { - dst[embedding_dim - 1] = __ldg(&src[embedding_dim - 1]); - } -} - -// Scalar copy fallback with __ldg optimization -template -__forceinline__ __device__ void copyScalar( - T *__restrict__ dst, - const T *__restrict__ src, - size_t embedding_dim) { - // Scalar copy with __ldg for read-only weight - for (size_t i = 0; i < embedding_dim; ++i) { - dst[i] = __ldg(&src[i]); - } -} - -#endif // __EMBEDDING_CUDA_KERNEL_CUH__ diff --git a/src/infiniop/ops/embedding/embedding.h b/src/infiniop/ops/embedding/embedding.h deleted file mode 100644 index e0135dbfe..000000000 --- a/src/infiniop/ops/embedding/embedding.h +++ /dev/null @@ -1,54 +0,0 @@ -#ifndef __EMBEDDING_H__ -#define __EMBEDDING_H__ - -#include "../../../utils.h" -#include "../../operator.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::embedding::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - size_t _num_indices; \ - size_t _embedding_dim; \ - size_t _vocab_size; \ - infiniDtype_t _input_dtype; \ - infiniDtype_t _weight_dtype; \ - \ - Descriptor( \ - size_t num_indices, \ - size_t embedding_dim, \ - size_t vocab_size, \ - infiniDtype_t input_dtype, \ - infiniDtype_t weight_dtype, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _num_indices(num_indices), \ - _embedding_dim(embedding_dim), \ - _vocab_size(vocab_size), \ - _input_dtype(input_dtype), \ - _weight_dtype(weight_dtype) {} \ - \ - public: \ - ~Descriptor(); \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t output_desc, \ - infiniopTensorDescriptor_t input_desc, \ - infiniopTensorDescriptor_t weight_desc); \ - \ - infiniStatus_t calculate( \ - void *output, \ - const void *input, \ - const void *weight, \ - void *stream) const; \ - }; \ - } - -#endif // __EMBEDDING_H__ diff --git a/src/infiniop/ops/embedding/kunlun/embedding_kunlun.h b/src/infiniop/ops/embedding/kunlun/embedding_kunlun.h deleted file mode 100644 index 184791676..000000000 --- a/src/infiniop/ops/embedding/kunlun/embedding_kunlun.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __EMBEDDING_KUNLUN_H__ -#define __EMBEDDING_KUNLUN_H__ - -#include "../embedding.h" - -DESCRIPTOR(kunlun) - -#endif // __EMBEDDING_KUNLUN_H__ diff --git a/src/infiniop/ops/embedding/kunlun/embedding_kunlun.xpu b/src/infiniop/ops/embedding/kunlun/embedding_kunlun.xpu deleted file mode 100644 index 59ba93aa8..000000000 --- a/src/infiniop/ops/embedding/kunlun/embedding_kunlun.xpu +++ /dev/null @@ -1,169 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/kunlun/kunlun_common.h" -#include "../../../devices/kunlun/kunlun_handle.h" -#include "../../../devices/kunlun/kunlun_kernel_common.h" -#include "../../../tensor.h" -#include "embedding_kunlun.h" -#include - -using namespace device::kunlun::kernel; - -template -__global__ void embeddingKernel( - __global_ptr__ T *output, - __global_ptr__ const IndexT *indices, - __global_ptr__ const T *weight, - uint32_t total_size, - uint32_t embedding_dim, - uint32_t vocab_size) { - - int cid = core_id(); - int ncores = core_num(); - if (cid >= ncores) { - return; - } - - int thread_id = ncores * cluster_id() + cid; - int nthreads = ncores * cluster_num(); - - for (uint32_t linear = thread_id; linear < total_size; linear += nthreads) { - uint32_t token_pos = linear / embedding_dim; - uint32_t dim = linear - token_pos * embedding_dim; - IndexT token = indices[token_pos]; - - if (token >= 0 && static_cast(token) < static_cast(vocab_size)) { - __local__ T value; - GM2LM_ASYNC(weight + static_cast(token) * embedding_dim + dim, &value, sizeof(T)); - mfence(); - LM2GM_ASYNC(&value, output + linear, sizeof(T)); - } - } -} - -namespace op::embedding::kunlun { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t weight_desc) { - - auto input_shape = input_desc->shape(); - auto weight_shape = weight_desc->shape(); - auto output_shape = output_desc->shape(); - - CHECK_OR_RETURN(weight_shape.size() == 2, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(output_shape.size() == input_shape.size() + 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - - size_t embedding_dim = weight_shape[1]; - CHECK_OR_RETURN(output_shape.back() == embedding_dim, INFINI_STATUS_BAD_TENSOR_SHAPE); - for (size_t i = 0; i < input_shape.size(); ++i) { - CHECK_OR_RETURN(output_shape[i] == input_shape[i], INFINI_STATUS_BAD_TENSOR_SHAPE); - } - - auto input_dtype = input_desc->dtype(); - auto weight_dtype = weight_desc->dtype(); - CHECK_OR_RETURN(input_dtype == INFINI_DTYPE_I32 || input_dtype == INFINI_DTYPE_I64, - INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(weight_dtype == INFINI_DTYPE_F32 || weight_dtype == INFINI_DTYPE_F16 || weight_dtype == INFINI_DTYPE_BF16, - INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(output_desc->dtype() == weight_dtype, INFINI_STATUS_BAD_TENSOR_DTYPE); - - size_t num_indices = 1; - for (auto dim : input_shape) { - num_indices *= dim; - } - - *desc_ptr = new Descriptor( - num_indices, - embedding_dim, - weight_shape[0], - input_dtype, - weight_dtype, - new Opaque{reinterpret_cast(handle)->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchEmbeddingByWeightDtype( - void *output, - const void *input, - const void *weight, - size_t total_size, - size_t embedding_dim, - size_t vocab_size, - infiniDtype_t weight_dtype, - kunlunStream_t stream) { - - if (total_size > static_cast(UINT32_MAX) || - embedding_dim > static_cast(UINT32_MAX) || - vocab_size > static_cast(UINT32_MAX)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - -#define LAUNCH_KERNEL(T) \ - embeddingKernel<<<12, 64, stream>>>( \ - reinterpret_cast<__global_ptr__ T *>(output), \ - reinterpret_cast<__global_ptr__ const IndexT *>(input), \ - reinterpret_cast<__global_ptr__ const T *>(weight), \ - static_cast(total_size), \ - static_cast(embedding_dim), \ - static_cast(vocab_size)); - - switch (weight_dtype) { - case INFINI_DTYPE_F32: - LAUNCH_KERNEL(float); - break; - case INFINI_DTYPE_F16: - LAUNCH_KERNEL(half); - break; - case INFINI_DTYPE_BF16: - LAUNCH_KERNEL(bfloat16_t); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_KERNEL - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *output, - const void *input, - const void *weight, - void *stream) const { - - if (_num_indices == 0) { - return INFINI_STATUS_SUCCESS; - } - - auto total_size = _num_indices * _embedding_dim; - auto kunlun_stream = reinterpret_cast(stream); - - switch (_input_dtype) { - case INFINI_DTYPE_I32: - return launchEmbeddingByWeightDtype( - output, input, weight, total_size, _embedding_dim, _vocab_size, _weight_dtype, kunlun_stream); - case INFINI_DTYPE_I64: - return launchEmbeddingByWeightDtype( - output, input, weight, total_size, _embedding_dim, _vocab_size, _weight_dtype, kunlun_stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::embedding::kunlun diff --git a/src/infiniop/ops/embedding/metax/embedding_metax.cuh b/src/infiniop/ops/embedding/metax/embedding_metax.cuh deleted file mode 100644 index 7290fc918..000000000 --- a/src/infiniop/ops/embedding/metax/embedding_metax.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __EMBEDDING_METAX_H__ -#define __EMBEDDING_METAX_H__ - -#include "../embedding.h" - -DESCRIPTOR(metax) - -#endif // __EMBEDDING_METAX_H__ diff --git a/src/infiniop/ops/embedding/metax/embedding_metax.maca b/src/infiniop/ops/embedding/metax/embedding_metax.maca deleted file mode 100644 index 8a1b24ea0..000000000 --- a/src/infiniop/ops/embedding/metax/embedding_metax.maca +++ /dev/null @@ -1,217 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../../../tensor.h" -#include "../cuda/embedding_kernel.cuh" -#include "embedding_metax.cuh" - -template -INFINIOP_METAX_KERNEL embeddingKernel( - T *__restrict__ output, - const IndexType *__restrict__ indices, - const T *__restrict__ weight, - size_t num_indices, - size_t embedding_dim, - size_t vocab_size) { - // Calculate global thread index - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_indices) { - // Get the index value - IndexType index_val = __ldg(&indices[idx]); - - // Bounds check - handle negative indices gracefully - if (index_val >= 0 && static_cast(index_val) < vocab_size) { - // Copy embedding vector from weight to output - const T *src = weight + static_cast(index_val) * embedding_dim; - T *dst = output + idx * embedding_dim; - - // Choose optimal copy strategy based on type and alignment - if constexpr (std::is_same_v) { - // Check alignment for float4 (16 bytes) - bool aligned_16 = is_aligned(src, 16) && is_aligned(dst, 16); - if (aligned_16 && embedding_dim >= 4 && embedding_dim % 4 == 0) { - copyVectorizedFloat4(dst, src, embedding_dim); - } else if (embedding_dim >= 2 && embedding_dim % 2 == 0) { - // Try float2 if not aligned to 16 bytes - copyVectorizedFloat2(dst, src, embedding_dim); - } else { - copyScalar(dst, src, embedding_dim); - } - } else if constexpr (std::is_same_v) { - // Use half2 for vectorized access - if (embedding_dim >= 2 && embedding_dim % 2 == 0) { - copyVectorizedHalf2(dst, src, embedding_dim); - } else { - copyScalar(dst, src, embedding_dim); - } - } else if constexpr (std::is_same_v) { - // Use bfloat162 for vectorized access - if (embedding_dim >= 2 && embedding_dim % 2 == 0) { - copyVectorizedBFloat162(dst, src, embedding_dim); - } else { - copyScalar(dst, src, embedding_dim); - } - } else { - // Fallback to scalar copy with __ldg - copyScalar(dst, src, embedding_dim); - } - } - } -} - -namespace op::embedding::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t weight_desc) { - - auto input_shape = input_desc->shape(); - auto weight_shape = weight_desc->shape(); - - // Validate shapes - CHECK_OR_RETURN(weight_shape.size() == 2, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(output_desc->shape().size() == input_shape.size() + 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - - // Check output shape matches input shape + embedding_dim - auto output_shape = output_desc->shape(); - size_t embedding_dim = weight_shape[1]; - CHECK_OR_RETURN(output_shape.back() == embedding_dim, INFINI_STATUS_BAD_TENSOR_SHAPE); - - for (size_t i = 0; i < input_shape.size(); ++i) { - CHECK_OR_RETURN(output_shape[i] == input_shape[i], INFINI_STATUS_BAD_TENSOR_SHAPE); - } - - // Validate dtypes - auto input_dtype = input_desc->dtype(); - auto weight_dtype = weight_desc->dtype(); - CHECK_OR_RETURN(input_dtype == INFINI_DTYPE_I32 || input_dtype == INFINI_DTYPE_I64, - INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(weight_dtype == INFINI_DTYPE_F32 || weight_dtype == INFINI_DTYPE_F16 || - weight_dtype == INFINI_DTYPE_BF16, INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(output_desc->dtype() == weight_dtype, INFINI_STATUS_BAD_TENSOR_DTYPE); - - // Calculate number of indices (supporting batch dimension) - size_t num_indices = 1; - for (auto dim : input_shape) { - num_indices *= dim; - } - - size_t vocab_size = weight_shape[0]; - - *desc_ptr = new Descriptor( - num_indices, - embedding_dim, - vocab_size, - input_dtype, - weight_dtype, - new Opaque{reinterpret_cast(handle)->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *output, - const void *input, - const void *weight, - void *stream) const { - - if (_num_indices == 0) { - return INFINI_STATUS_SUCCESS; - } - - auto hc_stream = reinterpret_cast(stream); - - // Dynamic block size optimization based on embedding_dim for Metax platform - size_t block_size = 256; // Default block size for Metax - if (_embedding_dim <= 64) { - block_size = 512; // Small embedding_dim: use larger block for better occupancy - } else if (_embedding_dim >= 1024) { - block_size = 128; // Large embedding_dim: use smaller block to reduce register pressure - } - - size_t grid_size = (_num_indices + block_size - 1) / block_size; - - // Launch kernel based on dtypes for Metax platform - if (_input_dtype == INFINI_DTYPE_I32) { - const int32_t *indices_ptr = reinterpret_cast(input); - - if (_weight_dtype == INFINI_DTYPE_F32) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else if (_weight_dtype == INFINI_DTYPE_F16) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else if (_weight_dtype == INFINI_DTYPE_BF16) { - // Use Metax's bfloat16 type - embeddingKernel<__hpcc_bfloat16, int32_t><<>>( - reinterpret_cast<__hpcc_bfloat16 *>(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (_input_dtype == INFINI_DTYPE_I64) { - const int64_t *indices_ptr = reinterpret_cast(input); - - if (_weight_dtype == INFINI_DTYPE_F32) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else if (_weight_dtype == INFINI_DTYPE_F16) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else if (_weight_dtype == INFINI_DTYPE_BF16) { - embeddingKernel<__hpcc_bfloat16, int64_t><<>>( - reinterpret_cast<__hpcc_bfloat16 *>(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::embedding::metax diff --git a/src/infiniop/ops/embedding/moore/embedding_moore.h b/src/infiniop/ops/embedding/moore/embedding_moore.h deleted file mode 100644 index edc397be6..000000000 --- a/src/infiniop/ops/embedding/moore/embedding_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __EMBEDDING_MOORE_H__ -#define __EMBEDDING_MOORE_H__ - -#include "../embedding.h" - -DESCRIPTOR(moore) - -#endif // __EMBEDDING_MOORE_H__ diff --git a/src/infiniop/ops/embedding/moore/embedding_moore.mu b/src/infiniop/ops/embedding/moore/embedding_moore.mu deleted file mode 100644 index 147ac830f..000000000 --- a/src/infiniop/ops/embedding/moore/embedding_moore.mu +++ /dev/null @@ -1,227 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "../../../tensor.h" -#include "embedding_moore_kernel.h" -#include "embedding_moore.h" -#include - -template -INFINIOP_MOORE_KERNEL embeddingKernel( - T *__restrict__ output, - const IndexType *__restrict__ indices, - const T *__restrict__ weight, - size_t num_indices, - size_t embedding_dim, - size_t vocab_size) { - // Calculate global thread index - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_indices) { - // Get the index value with Moore-optimized memory access - IndexType index_val = indices[idx]; - - // Bounds check - handle negative indices gracefully - if (index_val >= 0 && static_cast(index_val) < vocab_size) { - // Copy embedding vector from weight to output - const T *src = weight + static_cast(index_val) * embedding_dim; - T *dst = output + idx * embedding_dim; - - // Choose optimal copy strategy based on type and alignment - if constexpr (std::is_same_v) { - // Check alignment for float4 (16 bytes) - bool aligned_16 = is_aligned(src, 16) && is_aligned(dst, 16); - if (aligned_16 && embedding_dim >= 4 && embedding_dim % 4 == 0) { - copyVectorizedFloat4(dst, src, embedding_dim); - } else if (embedding_dim >= 2 && embedding_dim % 2 == 0) { - // Try float2 if not aligned to 16 bytes - copyVectorizedFloat2(dst, src, embedding_dim); - } else { - copyScalar(dst, src, embedding_dim); - } - } else if constexpr (std::is_same_v) { - // Use half2 for vectorized access - if (embedding_dim >= 2 && embedding_dim % 2 == 0) { - copyVectorizedHalf2(dst, src, embedding_dim); - } else { - copyScalar(dst, src, embedding_dim); - } - } else if constexpr (std::is_same_v) { - // Use mt_bfloat162 for vectorized access (Moore-specific type) - if (embedding_dim >= 2 && embedding_dim % 2 == 0) { - copyVectorizedBFloat162(dst, src, embedding_dim); - } else { - copyScalar(dst, src, embedding_dim); - } - } else { - // Fallback to scalar copy with Moore-optimized memory access - copyScalar(dst, src, embedding_dim); - } - } - } -} - -namespace op::embedding::moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t weight_desc) { - - auto input_shape = input_desc->shape(); - auto weight_shape = weight_desc->shape(); - - // Validate shapes - CHECK_OR_RETURN(weight_shape.size() == 2, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(output_desc->shape().size() == input_shape.size() + 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - - // Check output shape matches input shape + embedding_dim - auto output_shape = output_desc->shape(); - size_t embedding_dim = weight_shape[1]; - CHECK_OR_RETURN(output_shape.back() == embedding_dim, INFINI_STATUS_BAD_TENSOR_SHAPE); - - for (size_t i = 0; i < input_shape.size(); ++i) { - CHECK_OR_RETURN(output_shape[i] == input_shape[i], INFINI_STATUS_BAD_TENSOR_SHAPE); - } - - // Validate dtypes - auto input_dtype = input_desc->dtype(); - auto weight_dtype = weight_desc->dtype(); - CHECK_OR_RETURN(input_dtype == INFINI_DTYPE_I32 || input_dtype == INFINI_DTYPE_I64, - INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(weight_dtype == INFINI_DTYPE_F32 || weight_dtype == INFINI_DTYPE_F16 || weight_dtype == INFINI_DTYPE_BF16, INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(output_desc->dtype() == weight_dtype, INFINI_STATUS_BAD_TENSOR_DTYPE); - - // Calculate number of indices (supporting batch dimension) - size_t num_indices = 1; - for (auto dim : input_shape) { - num_indices *= dim; - } - - size_t vocab_size = weight_shape[0]; - - *desc_ptr = new Descriptor( - num_indices, - embedding_dim, - vocab_size, - input_dtype, - weight_dtype, - new Opaque{reinterpret_cast(handle)->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *output, - const void *input, - const void *weight, - void *stream) const { - - if (_num_indices == 0) { - return INFINI_STATUS_SUCCESS; - } - - auto musa_stream = reinterpret_cast(stream); - - // Dynamic block size optimization based on embedding_dim - // Moore platform typically has different performance characteristics - size_t block_size = 256; // Default for Moore - if (_embedding_dim <= 64) { - block_size = 512; // Small embedding_dim: use larger block for better occupancy - } else if (_embedding_dim >= 1024) { - block_size = 128; // Large embedding_dim: use smaller block to reduce register pressure - } else if (_embedding_dim <= 256) { - block_size = 384; // Medium embedding_dim: balanced configuration - } - - size_t grid_size = (_num_indices + block_size - 1) / block_size; - - // Launch kernel based on dtypes - // Note: Moore uses __mt_bfloat16 instead of __nv_bfloat16 - if (_input_dtype == INFINI_DTYPE_I32) { - const int32_t *indices_ptr = reinterpret_cast(input); - - if (_weight_dtype == INFINI_DTYPE_F32) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else if (_weight_dtype == INFINI_DTYPE_F16) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else if (_weight_dtype == INFINI_DTYPE_BF16) { - // Use Moore's bfloat16 type - embeddingKernel<__mt_bfloat16, int32_t><<>>( - reinterpret_cast<__mt_bfloat16 *>(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (_input_dtype == INFINI_DTYPE_I64) { - const int64_t *indices_ptr = reinterpret_cast(input); - - if (_weight_dtype == INFINI_DTYPE_F32) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else if (_weight_dtype == INFINI_DTYPE_F16) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else if (_weight_dtype == INFINI_DTYPE_BF16) { - embeddingKernel<__mt_bfloat16, int64_t><<>>( - reinterpret_cast<__mt_bfloat16 *>(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // Check for kernel launch errors - musaError_t err = musaGetLastError(); - if (err != musaSuccess) { - return INFINI_STATUS_INTERNAL_ERROR; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::embedding::moore diff --git a/src/infiniop/ops/embedding/moore/embedding_moore_kernel.h b/src/infiniop/ops/embedding/moore/embedding_moore_kernel.h deleted file mode 100644 index 9a7427b05..000000000 --- a/src/infiniop/ops/embedding/moore/embedding_moore_kernel.h +++ /dev/null @@ -1,116 +0,0 @@ -#ifndef __EMBEDDING_MOORE_KERNEL_CUH__ -#define __EMBEDDING_MOORE_KERNEL_CUH__ - -#include - -// Helper function to check memory alignment -__forceinline__ __device__ bool is_aligned(const void *ptr, size_t alignment) { - // Use size_t for pointer arithmetic in device code (more compatible) - return (reinterpret_cast(ptr) % alignment == 0); -} - -// Vectorized copy for float type using float4 -template -__forceinline__ __device__ void copyVectorizedFloat4( - float *__restrict__ dst, - const float *__restrict__ src, - size_t embedding_dim) { - // Use float4 for vectorized access (16 bytes, 4 floats) - const float4 *src_vec = reinterpret_cast(src); - float4 *dst_vec = reinterpret_cast(dst); - size_t vec_count = embedding_dim / 4; - - // Vectorized copy with __ldg equivalent for Moore platform - for (size_t i = 0; i < vec_count; ++i) { - dst_vec[i] = src_vec[i]; - } - - // Copy remaining elements - size_t remaining = embedding_dim % 4; - if (remaining > 0) { - size_t offset = vec_count * 4; - for (size_t i = 0; i < remaining; ++i) { - dst[offset + i] = src[offset + i]; - } - } -} - -// Vectorized copy for float type using float2 (fallback when not aligned to 16 bytes) -template -__forceinline__ __device__ void copyVectorizedFloat2( - float *__restrict__ dst, - const float *__restrict__ src, - size_t embedding_dim) { - // Use float2 for vectorized access (8 bytes, 2 floats) - const float2 *src_vec = reinterpret_cast(src); - float2 *dst_vec = reinterpret_cast(dst); - size_t vec_count = embedding_dim / 2; - - // Vectorized copy with Moore-optimized memory access - for (size_t i = 0; i < vec_count; ++i) { - dst_vec[i] = src_vec[i]; - } - - // Copy remaining element if odd - if (embedding_dim % 2 != 0) { - dst[embedding_dim - 1] = src[embedding_dim - 1]; - } -} - -// Vectorized copy for half type using half2 -template -__forceinline__ __device__ void copyVectorizedHalf2( - half *__restrict__ dst, - const half *__restrict__ src, - size_t embedding_dim) { - // Use half2 for vectorized access (4 bytes, 2 halfs) - const half2 *src_vec = reinterpret_cast(src); - half2 *dst_vec = reinterpret_cast(dst); - size_t vec_count = embedding_dim / 2; - - // Vectorized copy optimized for Moore architecture - for (size_t i = 0; i < vec_count; ++i) { - dst_vec[i] = src_vec[i]; - } - - // Copy remaining element if odd - if (embedding_dim % 2 != 0) { - dst[embedding_dim - 1] = src[embedding_dim - 1]; - } -} - -// Vectorized copy for Moore bfloat16 type using bfloat162 -template -__forceinline__ __device__ void copyVectorizedBFloat162( - __mt_bfloat16 *__restrict__ dst, - const __mt_bfloat16 *__restrict__ src, - size_t embedding_dim) { - // Use mt_bfloat162 for vectorized access (4 bytes, 2 bfloat16s) - const __mt_bfloat162 *src_vec = reinterpret_cast(src); - __mt_bfloat162 *dst_vec = reinterpret_cast<__mt_bfloat162 *>(dst); - size_t vec_count = embedding_dim / 2; - - // Vectorized copy with Moore-specific optimization - for (size_t i = 0; i < vec_count; ++i) { - dst_vec[i] = src_vec[i]; - } - - // Copy remaining element if odd - if (embedding_dim % 2 != 0) { - dst[embedding_dim - 1] = src[embedding_dim - 1]; - } -} - -// Scalar copy fallback with Moore-optimized memory access -template -__forceinline__ __device__ void copyScalar( - T *__restrict__ dst, - const T *__restrict__ src, - size_t embedding_dim) { - // Scalar copy with Moore read-only weight optimization - for (size_t i = 0; i < embedding_dim; ++i) { - dst[i] = src[i]; - } -} - -#endif // __EMBEDDING_MOORE_KERNEL_CUH__ diff --git a/src/infiniop/ops/embedding/nvidia/embedding_nvidia.cu b/src/infiniop/ops/embedding/nvidia/embedding_nvidia.cu deleted file mode 100644 index 8414e187e..000000000 --- a/src/infiniop/ops/embedding/nvidia/embedding_nvidia.cu +++ /dev/null @@ -1,224 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../tensor.h" -#include "../cuda/embedding_kernel.cuh" -#include "embedding_nvidia.cuh" -#include - -template -INFINIOP_CUDA_KERNEL embeddingKernel( - T *__restrict__ output, - const IndexType *__restrict__ indices, - const T *__restrict__ weight, - size_t num_indices, - size_t embedding_dim, - size_t vocab_size) { - // Calculate global thread index - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_indices) { - // Get the index value - IndexType index_val = __ldg(&indices[idx]); - - // Bounds check - handle negative indices gracefully - if (index_val >= 0 && static_cast(index_val) < vocab_size) { - // Copy embedding vector from weight to output - const T *src = weight + static_cast(index_val) * embedding_dim; - T *dst = output + idx * embedding_dim; - - // Choose optimal copy strategy based on type and alignment - if constexpr (std::is_same_v) { - // Check alignment for float4 (16 bytes) - bool aligned_16 = is_aligned(src, 16) && is_aligned(dst, 16); - if (aligned_16 && embedding_dim >= 4 && embedding_dim % 4 == 0) { - copyVectorizedFloat4(dst, src, embedding_dim); - } else if (embedding_dim >= 2 && embedding_dim % 2 == 0) { - // Try float2 if not aligned to 16 bytes - copyVectorizedFloat2(dst, src, embedding_dim); - } else { - copyScalar(dst, src, embedding_dim); - } - } else if constexpr (std::is_same_v) { - // Use half2 for vectorized access - if (embedding_dim >= 2 && embedding_dim % 2 == 0) { - copyVectorizedHalf2(dst, src, embedding_dim); - } else { - copyScalar(dst, src, embedding_dim); - } - } else if constexpr (std::is_same_v) { - // Use bfloat162 for vectorized access - if (embedding_dim >= 2 && embedding_dim % 2 == 0) { - copyVectorizedBFloat162(dst, src, embedding_dim); - } else { - copyScalar(dst, src, embedding_dim); - } - } else { - // Fallback to scalar copy with __ldg - copyScalar(dst, src, embedding_dim); - } - } - } -} - -namespace op::embedding::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t weight_desc) { - - auto input_shape = input_desc->shape(); - auto weight_shape = weight_desc->shape(); - - // Validate shapes - CHECK_OR_RETURN(weight_shape.size() == 2, INFINI_STATUS_BAD_TENSOR_SHAPE); - CHECK_OR_RETURN(output_desc->shape().size() == input_shape.size() + 1, INFINI_STATUS_BAD_TENSOR_SHAPE); - - // Check output shape matches input shape + embedding_dim - auto output_shape = output_desc->shape(); - size_t embedding_dim = weight_shape[1]; - CHECK_OR_RETURN(output_shape.back() == embedding_dim, INFINI_STATUS_BAD_TENSOR_SHAPE); - - for (size_t i = 0; i < input_shape.size(); ++i) { - CHECK_OR_RETURN(output_shape[i] == input_shape[i], INFINI_STATUS_BAD_TENSOR_SHAPE); - } - - // Validate dtypes - auto input_dtype = input_desc->dtype(); - auto weight_dtype = weight_desc->dtype(); - CHECK_OR_RETURN(input_dtype == INFINI_DTYPE_I32 || input_dtype == INFINI_DTYPE_I64, - INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(weight_dtype == INFINI_DTYPE_F32 || weight_dtype == INFINI_DTYPE_F16 || weight_dtype == INFINI_DTYPE_BF16, INFINI_STATUS_BAD_TENSOR_DTYPE); - CHECK_OR_RETURN(output_desc->dtype() == weight_dtype, INFINI_STATUS_BAD_TENSOR_DTYPE); - - // Calculate number of indices (supporting batch dimension) - size_t num_indices = 1; - for (auto dim : input_shape) { - num_indices *= dim; - } - - size_t vocab_size = weight_shape[0]; - - *desc_ptr = new Descriptor( - num_indices, - embedding_dim, - vocab_size, - input_dtype, - weight_dtype, - new Opaque{reinterpret_cast(handle)->internal()}, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *output, - const void *input, - const void *weight, - void *stream) const { - - if (_num_indices == 0) { - return INFINI_STATUS_SUCCESS; - } - - auto cuda_stream = reinterpret_cast(stream); - - // Dynamic block size optimization based on embedding_dim - // Smaller embedding_dim benefits from larger block size (better occupancy) - // Larger embedding_dim benefits from smaller block size (more registers per thread) - size_t block_size = 256; // Default - if (_embedding_dim <= 64) { - block_size = 512; // Small embedding_dim: use larger block for better occupancy - } else if (_embedding_dim >= 1024) { - block_size = 128; // Large embedding_dim: use smaller block to reduce register pressure - } - - size_t grid_size = (_num_indices + block_size - 1) / block_size; - - // Launch kernel based on dtypes - if (_input_dtype == INFINI_DTYPE_I32) { - const int32_t *indices_ptr = reinterpret_cast(input); - - if (_weight_dtype == INFINI_DTYPE_F32) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else if (_weight_dtype == INFINI_DTYPE_F16) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else if (_weight_dtype == INFINI_DTYPE_BF16) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (_input_dtype == INFINI_DTYPE_I64) { - const int64_t *indices_ptr = reinterpret_cast(input); - - if (_weight_dtype == INFINI_DTYPE_F32) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else if (_weight_dtype == INFINI_DTYPE_F16) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else if (_weight_dtype == INFINI_DTYPE_BF16) { - embeddingKernel<<>>( - reinterpret_cast(output), - indices_ptr, - reinterpret_cast(weight), - _num_indices, - _embedding_dim, - _vocab_size); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // Check for kernel launch errors - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - return INFINI_STATUS_INTERNAL_ERROR; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::embedding::nvidia diff --git a/src/infiniop/ops/embedding/nvidia/embedding_nvidia.cuh b/src/infiniop/ops/embedding/nvidia/embedding_nvidia.cuh deleted file mode 100644 index c6b966d8d..000000000 --- a/src/infiniop/ops/embedding/nvidia/embedding_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __EMBEDDING_CUDA_H__ -#define __EMBEDDING_CUDA_H__ - -#include "../embedding.h" - -DESCRIPTOR(nvidia) - -#endif // __EMBEDDING_CUDA_H__ diff --git a/src/infiniop/ops/embedding/operator.cc b/src/infiniop/ops/embedding/operator.cc deleted file mode 100644 index 5301e57a5..000000000 --- a/src/infiniop/ops/embedding/operator.cc +++ /dev/null @@ -1,188 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/embedding.h" - -#ifdef ENABLE_CAMBRICON_API -#include "bang/embedding_bang.h" -#endif - -#ifdef ENABLE_CPU_API -#include "cpu/embedding_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) -#include "nvidia/embedding_nvidia.cuh" -#endif -#ifdef ENABLE_ASCEND_API -#include "ascend/embedding_ascend.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/embedding_metax.cuh" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/embedding_moore.h" -#endif -#ifdef ENABLE_KUNLUN_API -#include "kunlun/embedding_kunlun.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateEmbeddingDescriptor( - infiniopHandle_t handle, - infiniopEmbeddingDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t weight_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::embedding::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output_desc, \ - input_desc, \ - weight_desc) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_ASCEND_API - CREATE(INFINI_DEVICE_ASCEND, ascend); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_KUNLUN_API - CREATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopEmbedding( - infiniopEmbeddingDescriptor_t desc, - void *output, - const void *input, - const void *weight, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(output, input, weight, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_ASCEND_API - CALCULATE(INFINI_DEVICE_ASCEND, ascend); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_KUNLUN_API - CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyEmbeddingDescriptor(infiniopEmbeddingDescriptor_t desc) { - -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DESTROY(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - DESTROY(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - DESTROY(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DESTROY(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - DESTROY(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DESTROY(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_CAMBRICON_API - DESTROY(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_KUNLUN_API - DESTROY(INFINI_DEVICE_KUNLUN, kunlun); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DESTROY - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} diff --git a/src/infiniop/ops/equal/cpu/equal_cpu.cc b/src/infiniop/ops/equal/cpu/equal_cpu.cc deleted file mode 100644 index ff8ebe395..000000000 --- a/src/infiniop/ops/equal/cpu/equal_cpu.cc +++ /dev/null @@ -1,68 +0,0 @@ -#include -#include - -#include "equal_cpu.h" - -namespace op::equal::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - auto compute_dtype = a_desc->dtype(); - auto out_dtype = out_desc->dtype(); - - if (compute_dtype != b_desc->dtype()) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - CHECK_DTYPE(out_dtype, INFINI_DTYPE_BOOL); - - CHECK_DTYPE(compute_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, - INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, compute_dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::equal::cpu diff --git a/src/infiniop/ops/equal/cpu/equal_cpu.h b/src/infiniop/ops/equal/cpu/equal_cpu.h deleted file mode 100644 index fd811f4b0..000000000 --- a/src/infiniop/ops/equal/cpu/equal_cpu.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef __EQUAL_CPU_H__ -#define __EQUAL_CPU_H__ - -#include - -#include "../../../elementwise/cpu/elementwise_cpu.h" - -ELEMENTWISE_DESCRIPTOR(equal, cpu) - -namespace op::equal::cpu { - -typedef struct EqualOp { -public: - static constexpr size_t num_inputs = 2; - - template - bool operator()(const Tin0 &a, const Tin1 &b) { - if constexpr (std::is_same_v) { - return a == b; - } else { - return false; - } - } -} EqualOp; - -} // namespace op::equal::cpu - -#endif diff --git a/src/infiniop/ops/equal/cuda/kernel.cuh b/src/infiniop/ops/equal/cuda/kernel.cuh deleted file mode 100644 index 30b4ae4fb..000000000 --- a/src/infiniop/ops/equal/cuda/kernel.cuh +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef __EQUAL_CUDA_H__ -#define __EQUAL_CUDA_H__ - -#include - -namespace op::equal::cuda { - -typedef struct EqualOp { -public: - static constexpr size_t num_inputs = 2; - - template - __device__ __forceinline__ bool operator()(const Tin0 &a, const Tin1 &b) const { - if constexpr (std::is_same_v) { - if constexpr (std::is_same_v) { - static_assert(!std::is_same_v, "half2 is not supported for mixed output dtype"); - } else if constexpr (std::is_same_v) { - return static_cast(__heq(a, b)); - } else { - return static_cast(a == b); - } - } else { - return false; - } - } -} EqualOp; - -} // namespace op::equal::cuda - -#endif diff --git a/src/infiniop/ops/equal/metax/equal_metax.h b/src/infiniop/ops/equal/metax/equal_metax.h deleted file mode 100644 index 6e4cd64b9..000000000 --- a/src/infiniop/ops/equal/metax/equal_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __EQUAL_METAX_API_H__ -#define __EQUAL_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(equal, metax) - -#endif // __EQUAL_METAX_API_H__ diff --git a/src/infiniop/ops/equal/metax/equal_metax.maca b/src/infiniop/ops/equal/metax/equal_metax.maca deleted file mode 100644 index 265e5b5a6..000000000 --- a/src/infiniop/ops/equal/metax/equal_metax.maca +++ /dev/null @@ -1,69 +0,0 @@ -#include "equal_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::equal::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - - const auto &a_desc = input_desc_vec.at(0); - auto compute_dtype = a_desc->dtype(); - auto out_dtype = out_desc->dtype(); - - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(compute_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, - INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_F64); - - CHECK_DTYPE(out_dtype, INFINI_DTYPE_BOOL); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, compute_dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::EqualOp, bool, half, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::EqualOp, bool, cuda_bfloat16, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::EqualOp, bool, float, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, cuda::EqualOp, bool, int32_t, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, cuda::EqualOp, bool, int64_t, int64_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::EqualOp, bool, double, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::equal::metax diff --git a/src/infiniop/ops/equal/moore/equal_moore.h b/src/infiniop/ops/equal/moore/equal_moore.h deleted file mode 100644 index 2fed1bb40..000000000 --- a/src/infiniop/ops/equal/moore/equal_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __EQUAL_MOORE_API_H__ -#define __EQUAL_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(equal, moore) - -#endif // __EQUAL_MOORE_API_H__ diff --git a/src/infiniop/ops/equal/moore/equal_moore.mu b/src/infiniop/ops/equal/moore/equal_moore.mu deleted file mode 100644 index d0eb8395d..000000000 --- a/src/infiniop/ops/equal/moore/equal_moore.mu +++ /dev/null @@ -1,140 +0,0 @@ -#include "equal_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "equal_moore_kernel.h" - -namespace op::equal::moore { -namespace { - -inline bool can_use_contiguous_fast_path(const op::elementwise::ElementwiseInfo &info) { - if (!info.isOutputContiguous()) { - return false; - } - const bool *input_contiguous = info.getInputContiguous(); - const bool *input_broadcasted = info.getInputBroadcasted(); - for (size_t i = 0; i < 2; ++i) { - if (!input_contiguous[i] || input_broadcasted[i]) { - return false; - } - } - return true; -} - -template -INFINIOP_MOORE_KERNEL equal_contiguous_kernel(size_t numel, Tout *output, const Tin *a, const Tin *b) { - const auto op = op::equal::moore::EqualOp{}; - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - size_t stride = blockDim.x * gridDim.x; - for (; idx < numel; idx += stride) { - output[idx] = op.template operator()(a[idx], b[idx]); - } -} - -template -infiniStatus_t launch_fast_path(size_t numel, - void *output, - const std::vector &inputs, - void *stream) { - if (numel == 0) { - return INFINI_STATUS_SUCCESS; - } - - constexpr int kBlockSize = 256; - int grid = static_cast((numel + kBlockSize - 1) / kBlockSize); - if (grid > 65535) { - grid = 65535; - } - - auto musa_stream = reinterpret_cast(stream); - equal_contiguous_kernel<<>>( - numel, - reinterpret_cast(output), - reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[1])); - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - - const auto &a_desc = input_desc_vec.at(0); - auto compute_dtype = a_desc->dtype(); - auto out_dtype = out_desc->dtype(); - - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(compute_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, - INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_F64); - - CHECK_DTYPE(out_dtype, INFINI_DTYPE_BOOL); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create MOORE elementwise descriptor - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, compute_dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - if (can_use_contiguous_fast_path(_info)) { - size_t numel = _info.getOutputSize(); - switch (_dtype) { - case INFINI_DTYPE_F16: - return launch_fast_path(numel, output, inputs, stream); - case INFINI_DTYPE_BF16: - return launch_fast_path(numel, output, inputs, stream); - case INFINI_DTYPE_F32: - return launch_fast_path(numel, output, inputs, stream); - case INFINI_DTYPE_I32: - return launch_fast_path(numel, output, inputs, stream); - case INFINI_DTYPE_I64: - return launch_fast_path(numel, output, inputs, stream); - case INFINI_DTYPE_F64: - return launch_fast_path(numel, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, moore::EqualOp, bool, half, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, moore::EqualOp, bool, cuda_bfloat16, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, moore::EqualOp, bool, float, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, moore::EqualOp, bool, int32_t, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, moore::EqualOp, bool, int64_t, int64_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, moore::EqualOp, bool, double, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::equal::moore diff --git a/src/infiniop/ops/equal/moore/equal_moore_kernel.h b/src/infiniop/ops/equal/moore/equal_moore_kernel.h deleted file mode 100644 index a4e32880b..000000000 --- a/src/infiniop/ops/equal/moore/equal_moore_kernel.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef __EQUAL_MOORE_KERNEL_H__ -#define __EQUAL_MOORE_KERNEL_H__ - -#include - -namespace op::equal::moore { - -typedef struct EqualOp { -public: - static constexpr size_t num_inputs = 2; - - template - __device__ __forceinline__ bool operator()(const Tin0 &a, const Tin1 &b) const { - if constexpr (std::is_same_v) { - if constexpr (std::is_same_v) { - return __half2float(a) == __half2float(b); - } else if constexpr (std::is_same_v) { - return __bfloat162float(a) == __bfloat162float(b); - } else { - return a == b; - } - } else { - return false; - } - } -} EqualOp; - -} // namespace op::equal::moore - -#endif // __EQUAL_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/equal/nvidia/equal_nvidia.cu b/src/infiniop/ops/equal/nvidia/equal_nvidia.cu deleted file mode 100644 index 5bdf92e6c..000000000 --- a/src/infiniop/ops/equal/nvidia/equal_nvidia.cu +++ /dev/null @@ -1,137 +0,0 @@ -#include -#include -#include - -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "equal_nvidia.cuh" - -namespace { - -template -INFINIOP_CUDA_KERNEL FastEqualKernel(size_t n, Tout *output, const Tin *a, const Tin *b) { - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - size_t stride = blockDim.x * gridDim.x; - op::equal::cuda::EqualOp op{}; - for (; idx < n; idx += stride) { - output[idx] = op.template operator()(a[idx], b[idx]); - } -} - -template -infiniStatus_t launchFastEqualKernel(size_t numel, - void *output, - const std::vector &inputs, - void *stream) { - if (numel == 0) { - return INFINI_STATUS_SUCCESS; - } - constexpr int block = 256; - int grid = static_cast((numel + block - 1) / block); - grid = std::min(grid, 65535); - auto cuda_stream = reinterpret_cast(stream); - FastEqualKernel<<>>( - numel, - reinterpret_cast(output), - reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[1])); - auto err = cudaGetLastError(); - return err == cudaSuccess ? INFINI_STATUS_SUCCESS : INFINI_STATUS_INTERNAL_ERROR; -} - -} // namespace - -namespace op::equal::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - - const auto &a_desc = input_desc_vec.at(0); - auto compute_dtype = a_desc->dtype(); - auto out_dtype = out_desc->dtype(); - - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(compute_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, - INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_F64); - - CHECK_DTYPE(out_dtype, INFINI_DTYPE_BOOL, INFINI_DTYPE_U8, INFINI_DTYPE_I8); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, compute_dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - bool fast_path = _info.isOutputContiguous(); - if (fast_path) { - const bool *input_contiguous = _info.getInputContiguous(); - const bool *input_broadcasted = _info.getInputBroadcasted(); - for (size_t i = 0; i < 2; ++i) { - fast_path &= input_contiguous[i] && !input_broadcasted[i]; - } - } - - if (fast_path) { - size_t numel = _info.getOutputSize(); - switch (_dtype) { - case INFINI_DTYPE_F16: - return launchFastEqualKernel(numel, output, inputs, stream); - case INFINI_DTYPE_BF16: - return launchFastEqualKernel(numel, output, inputs, stream); - case INFINI_DTYPE_F32: - return launchFastEqualKernel(numel, output, inputs, stream); - case INFINI_DTYPE_I32: - return launchFastEqualKernel(numel, output, inputs, stream); - case INFINI_DTYPE_I64: - return launchFastEqualKernel(numel, output, inputs, stream); - case INFINI_DTYPE_F64: - return launchFastEqualKernel(numel, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::EqualOp, bool, half, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::EqualOp, bool, cuda_bfloat16, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::EqualOp, bool, float, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, cuda::EqualOp, bool, int32_t, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, cuda::EqualOp, bool, int64_t, int64_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::EqualOp, bool, double, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::equal::nvidia diff --git a/src/infiniop/ops/equal/nvidia/equal_nvidia.cuh b/src/infiniop/ops/equal/nvidia/equal_nvidia.cuh deleted file mode 100644 index 6565a80b5..000000000 --- a/src/infiniop/ops/equal/nvidia/equal_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __EQUAL_CUDA_API_H__ -#define __EQUAL_CUDA_API_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(equal, nvidia) - -#endif diff --git a/src/infiniop/ops/equal/operator.cc b/src/infiniop/ops/equal/operator.cc deleted file mode 100644 index 9c5866c1e..000000000 --- a/src/infiniop/ops/equal/operator.cc +++ /dev/null @@ -1,183 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/equal.h" - -#ifdef ENABLE_CPU_API -#include "cpu/equal_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/equal_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/equal_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/equal_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateEqualDescriptor( - infiniopHandle_t handle, - infiniopEqualDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::equal::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - c_desc, \ - {a_desc, b_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetEqualWorkspaceSize(infiniopEqualDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopEqual( - infiniopEqualDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, c, {a, b}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyEqualDescriptor(infiniopEqualDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/erf/cpu/erf_cpu.cc b/src/infiniop/ops/erf/cpu/erf_cpu.cc deleted file mode 100644 index 7d127bfae..000000000 --- a/src/infiniop/ops/erf/cpu/erf_cpu.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "erf_cpu.h" - -namespace op::erf::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::erf::cpu diff --git a/src/infiniop/ops/erf/cpu/erf_cpu.h b/src/infiniop/ops/erf/cpu/erf_cpu.h deleted file mode 100644 index 74ad19e57..000000000 --- a/src/infiniop/ops/erf/cpu/erf_cpu.h +++ /dev/null @@ -1,20 +0,0 @@ -#ifndef __ERF_CPU_H__ -#define __ERF_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" -#include - -ELEMENTWISE_DESCRIPTOR(erf, cpu) - -namespace op::erf::cpu { -typedef struct ErfOp { -public: - static constexpr size_t num_inputs = 1; - template - T operator()(const T &x) const { - return std::erf(x); - } -} ErfOp; -} // namespace op::erf::cpu - -#endif // __ERF_CPU_H__ diff --git a/src/infiniop/ops/erf/cuda/kernel.cuh b/src/infiniop/ops/erf/cuda/kernel.cuh deleted file mode 100644 index 8efe037e3..000000000 --- a/src/infiniop/ops/erf/cuda/kernel.cuh +++ /dev/null @@ -1,33 +0,0 @@ -#pragma once -#include -#include - -namespace op::cuda { - -struct ErfOp { - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(T x) const { - if constexpr (std::is_same_v) { - return erff(x); - } else if constexpr (std::is_same_v) { - return ::erf(x); - } else { - // For F16/BF16: promote to float, compute, then cast back - float xf; - if constexpr (std::is_same_v) { - xf = __half2float(x); - return __float2half_rn(erff(xf)); - } else if constexpr (std::is_same_v) { - xf = __bfloat162float(x); - return __float2bfloat16_rn(erff(xf)); - } else { - xf = static_cast(x); - return static_cast(erff(xf)); - } - } - } -}; - -} // namespace op::cuda diff --git a/src/infiniop/ops/erf/erf.h b/src/infiniop/ops/erf/erf.h deleted file mode 100644 index 7c967dea2..000000000 --- a/src/infiniop/ops/erf/erf.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ERF_H__ -#define __ERF_H__ - -#include "../../elementwise/elementwise.h" - -#define DESCRIPTOR(NAMESPACE) ELEMENTWISE_DESCRIPTOR(erf, NAMESPACE) - -#endif // __ERF_H__ diff --git a/src/infiniop/ops/erf/metax/erf_metax.h b/src/infiniop/ops/erf/metax/erf_metax.h deleted file mode 100644 index 5dfe23bbe..000000000 --- a/src/infiniop/ops/erf/metax/erf_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ERF_METAX_API_H__ -#define __ERF_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(erf, metax) - -#endif // __ERF_METAX_API_H__ diff --git a/src/infiniop/ops/erf/metax/erf_metax.maca b/src/infiniop/ops/erf/metax/erf_metax.maca deleted file mode 100644 index fcf956af2..000000000 --- a/src/infiniop/ops/erf/metax/erf_metax.maca +++ /dev/null @@ -1,60 +0,0 @@ -#include "erf_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::erf::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::ErfOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::ErfOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::ErfOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::ErfOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::erf::metax diff --git a/src/infiniop/ops/erf/moore/erf_moore.h b/src/infiniop/ops/erf/moore/erf_moore.h deleted file mode 100644 index 620055688..000000000 --- a/src/infiniop/ops/erf/moore/erf_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ERF_MOORE_API_H__ -#define __ERF_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(erf, moore) - -#endif // __ERF_MOORE_API_H__ diff --git a/src/infiniop/ops/erf/moore/erf_moore.mu b/src/infiniop/ops/erf/moore/erf_moore.mu deleted file mode 100644 index 856764c70..000000000 --- a/src/infiniop/ops/erf/moore/erf_moore.mu +++ /dev/null @@ -1,60 +0,0 @@ -#include "erf_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "../cuda/kernel.cuh" - -namespace op::erf::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::ErfOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::ErfOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::ErfOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::ErfOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::erf::moore diff --git a/src/infiniop/ops/erf/nvidia/erf_nvidia.cu b/src/infiniop/ops/erf/nvidia/erf_nvidia.cu deleted file mode 100644 index 03e14bb57..000000000 --- a/src/infiniop/ops/erf/nvidia/erf_nvidia.cu +++ /dev/null @@ -1,58 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "erf_nvidia.cuh" - -namespace op::erf::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::ErfOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::ErfOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::ErfOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::ErfOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::erf::nvidia diff --git a/src/infiniop/ops/erf/nvidia/erf_nvidia.cuh b/src/infiniop/ops/erf/nvidia/erf_nvidia.cuh deleted file mode 100644 index d20658027..000000000 --- a/src/infiniop/ops/erf/nvidia/erf_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ERF_NVIDIA_H__ -#define __ERF_NVIDIA_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(erf, nvidia) - -#endif // __ERF_NVIDIA_H__ diff --git a/src/infiniop/ops/erf/operator.cc b/src/infiniop/ops/erf/operator.cc deleted file mode 100644 index af1327c3b..000000000 --- a/src/infiniop/ops/erf/operator.cc +++ /dev/null @@ -1,169 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/erf.h" - -#ifdef ENABLE_CPU_API -#include "cpu/erf_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/erf_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/erf_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/erf_moore.h" -#endif - -__INFINI_C __export infiniStatus_t infiniopCreateErfDescriptor( - infiniopHandle_t handle, - infiniopErfDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::erf::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - {x_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C __export infiniStatus_t infiniopGetErfWorkspaceSize(infiniopErfDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C __export infiniStatus_t infiniopErf( - infiniopErfDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, {x}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C __export infiniStatus_t -infiniopDestroyErfDescriptor(infiniopErfDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/erfc/cpu/erfc_cpu.cc b/src/infiniop/ops/erfc/cpu/erfc_cpu.cc deleted file mode 100644 index 35b82c678..000000000 --- a/src/infiniop/ops/erfc/cpu/erfc_cpu.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "erfc_cpu.h" - -namespace op::erfc::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::erfc::cpu diff --git a/src/infiniop/ops/erfc/cpu/erfc_cpu.h b/src/infiniop/ops/erfc/cpu/erfc_cpu.h deleted file mode 100644 index dd6d69496..000000000 --- a/src/infiniop/ops/erfc/cpu/erfc_cpu.h +++ /dev/null @@ -1,20 +0,0 @@ -#ifndef __ERFC_CPU_H__ -#define __ERFC_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" -#include - -ELEMENTWISE_DESCRIPTOR(erfc, cpu) - -namespace op::erfc::cpu { -typedef struct ErfcOp { -public: - static constexpr size_t num_inputs = 1; - template - T operator()(const T &x) const { - return std::erfc(x); - } -} ErfcOp; -} // namespace op::erfc::cpu - -#endif // __ERFC_CPU_H__ diff --git a/src/infiniop/ops/erfc/cuda/kernel.cuh b/src/infiniop/ops/erfc/cuda/kernel.cuh deleted file mode 100644 index 5a448919d..000000000 --- a/src/infiniop/ops/erfc/cuda/kernel.cuh +++ /dev/null @@ -1,33 +0,0 @@ -#pragma once -#include -#include - -namespace op::cuda { - -struct ErfcOp { - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(T x) const { - if constexpr (std::is_same_v) { - return erfcf(x); - } else if constexpr (std::is_same_v) { - return ::erfc(x); - } else { - // For F16/BF16: promote to float, compute, then cast back - float xf; - if constexpr (std::is_same_v) { - xf = __half2float(x); - return __float2half_rn(erfcf(xf)); - } else if constexpr (std::is_same_v) { - xf = __bfloat162float(x); - return __float2bfloat16_rn(erfcf(xf)); - } else { - xf = static_cast(x); - return static_cast(erfcf(xf)); - } - } - } -}; - -} // namespace op::cuda diff --git a/src/infiniop/ops/erfc/erfc.h b/src/infiniop/ops/erfc/erfc.h deleted file mode 100644 index 9ee12fd43..000000000 --- a/src/infiniop/ops/erfc/erfc.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ERFC_H__ -#define __ERFC_H__ - -#include "../../elementwise/elementwise.h" - -#define DESCRIPTOR(NAMESPACE) ELEMENTWISE_DESCRIPTOR(erfc, NAMESPACE) - -#endif // __ERFC_H__ diff --git a/src/infiniop/ops/erfc/metax/erfc_metax.h b/src/infiniop/ops/erfc/metax/erfc_metax.h deleted file mode 100644 index 438f00095..000000000 --- a/src/infiniop/ops/erfc/metax/erfc_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ERFC_METAX_API_H__ -#define __ERFC_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(erfc, metax) - -#endif // __ERFC_METAX_API_H__ diff --git a/src/infiniop/ops/erfc/metax/erfc_metax.maca b/src/infiniop/ops/erfc/metax/erfc_metax.maca deleted file mode 100644 index 925e52ccf..000000000 --- a/src/infiniop/ops/erfc/metax/erfc_metax.maca +++ /dev/null @@ -1,60 +0,0 @@ -#include "erfc_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::erfc::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::ErfcOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::ErfcOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::ErfcOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::ErfcOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::erfc::metax diff --git a/src/infiniop/ops/erfc/moore/erfc_moore.h b/src/infiniop/ops/erfc/moore/erfc_moore.h deleted file mode 100644 index d032e4305..000000000 --- a/src/infiniop/ops/erfc/moore/erfc_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ERFC_MOORE_API_H__ -#define __ERFC_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(erfc, moore) - -#endif // __ERFC_MOORE_API_H__ diff --git a/src/infiniop/ops/erfc/moore/erfc_moore.mu b/src/infiniop/ops/erfc/moore/erfc_moore.mu deleted file mode 100644 index d58243e50..000000000 --- a/src/infiniop/ops/erfc/moore/erfc_moore.mu +++ /dev/null @@ -1,60 +0,0 @@ -#include "erfc_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "../cuda/kernel.cuh" - -namespace op::erfc::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::ErfcOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::ErfcOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::ErfcOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::ErfcOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::erfc::moore diff --git a/src/infiniop/ops/erfc/nvidia/erfc_nvidia.cu b/src/infiniop/ops/erfc/nvidia/erfc_nvidia.cu deleted file mode 100644 index 483f11a18..000000000 --- a/src/infiniop/ops/erfc/nvidia/erfc_nvidia.cu +++ /dev/null @@ -1,58 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "erfc_nvidia.cuh" - -namespace op::erfc::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::ErfcOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::ErfcOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::ErfcOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::ErfcOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::erfc::nvidia diff --git a/src/infiniop/ops/erfc/nvidia/erfc_nvidia.cuh b/src/infiniop/ops/erfc/nvidia/erfc_nvidia.cuh deleted file mode 100644 index 4d5321c9d..000000000 --- a/src/infiniop/ops/erfc/nvidia/erfc_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ERFC_NVIDIA_H__ -#define __ERFC_NVIDIA_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(erfc, nvidia) - -#endif // __ERFC_NVIDIA_H__ diff --git a/src/infiniop/ops/erfc/operator.cc b/src/infiniop/ops/erfc/operator.cc deleted file mode 100644 index ce283df19..000000000 --- a/src/infiniop/ops/erfc/operator.cc +++ /dev/null @@ -1,169 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/erfc.h" - -#ifdef ENABLE_CPU_API -#include "cpu/erfc_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/erfc_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/erfc_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/erfc_moore.h" -#endif - -__INFINI_C __export infiniStatus_t infiniopCreateErfcDescriptor( - infiniopHandle_t handle, - infiniopErfcDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::erfc::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - {x_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C __export infiniStatus_t infiniopGetErfcWorkspaceSize(infiniopErfcDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C __export infiniStatus_t infiniopErfc( - infiniopErfcDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, {x}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C __export infiniStatus_t -infiniopDestroyErfcDescriptor(infiniopErfcDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/erfinv/cpu/erfinv_cpu.cc b/src/infiniop/ops/erfinv/cpu/erfinv_cpu.cc deleted file mode 100644 index 16c5c8cba..000000000 --- a/src/infiniop/ops/erfinv/cpu/erfinv_cpu.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "erfinv_cpu.h" - -namespace op::erfinv::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::erfinv::cpu diff --git a/src/infiniop/ops/erfinv/cpu/erfinv_cpu.h b/src/infiniop/ops/erfinv/cpu/erfinv_cpu.h deleted file mode 100644 index e7b1b723d..000000000 --- a/src/infiniop/ops/erfinv/cpu/erfinv_cpu.h +++ /dev/null @@ -1,56 +0,0 @@ -#ifndef __ERFINV_CPU_H__ -#define __ERFINV_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" -#include -#include - -ELEMENTWISE_DESCRIPTOR(erfinv, cpu) - -namespace op::erfinv::cpu { - -// Inverse error function implementation using Newton's method -template -T erfinv_impl(T x) { - // Domain: x in (-1, 1) - if (x == 1.0) { - return std::numeric_limits::infinity(); - } - if (x == -1.0) { - return -std::numeric_limits::infinity(); - } - if (x > 1.0 || x < -1.0) { - return std::numeric_limits::quiet_NaN(); - } - if (x == 0.0) { - return 0.0; - } - - // Use Newton's method to solve erf(y) = x - T y = x; // Initial guess - const int max_iter = 10; - const T tol = static_cast(1e-10); - - for (int i = 0; i < max_iter; ++i) { - T erf_y = std::erf(y); - T derf_dy = T(2.0) / T(std::sqrt(3.14159265358979323846) * std::exp(-y * y)); - T error = erf_y - x; - if (std::abs(error) < tol) { - break; - } - y = y - error / derf_dy; - } - return y; -} - -typedef struct ErfinvOp { -public: - static constexpr size_t num_inputs = 1; - template - T operator()(const T &x) const { - return erfinv_impl(x); - } -} ErfinvOp; -} // namespace op::erfinv::cpu - -#endif // __ERFINV_CPU_H__ diff --git a/src/infiniop/ops/erfinv/cuda/kernel.cuh b/src/infiniop/ops/erfinv/cuda/kernel.cuh deleted file mode 100644 index f0ebcd52c..000000000 --- a/src/infiniop/ops/erfinv/cuda/kernel.cuh +++ /dev/null @@ -1,152 +0,0 @@ -#pragma once -#include -#include -#include - -namespace op::cuda { - -constexpr float PI_F = 3.14159265358979323846f; -constexpr double PI = 3.14159265358979323846; - -// Inverse error function. -__device__ __forceinline__ float erfinv_impl(float x) { - if (x == 1.0f) { -#if defined(ENABLE_HYGON_API) - return __builtin_huge_valf(); -#else - return std::numeric_limits::infinity(); -#endif - } - if (x == -1.0f) { -#if defined(ENABLE_HYGON_API) - return -__builtin_huge_valf(); -#else - return -std::numeric_limits::infinity(); -#endif - } - if (x > 1.0f || x < -1.0f) { -#if defined(ENABLE_HYGON_API) - return __builtin_nanf(""); -#else - return std::numeric_limits::quiet_NaN(); -#endif - } - if (x == 0.0f) { - return 0.0f; - } - - // Winitzki approximation - const float a = 0.147f; - const float ln = log1pf(-x * x); - const float t = 2.0f / (PI_F * a) + ln * 0.5f; - - float inside = t * t - ln / a; - inside = inside > 0.0f ? inside : 0.0f; - - float y0 = copysignf(sqrtf(sqrtf(inside) - t), x); - - float y = y0; - const float sqrt_pi_f = sqrtf(PI_F); - -#pragma unroll - for (int i = 0; i < 4; ++i) { - const float erf_y = erff(y); - const float derf_dy = 2.0f / sqrt_pi_f * expf(-y * y); - y = y - (erf_y - x) / derf_dy; - } - - // Slow path near |x| ~ 1 - const float ax = fabsf(x); - if (1.0f - ax < 1e-4f) { - const double xd = static_cast(x); - double yd = static_cast(y); - const double sqrt_pi = sqrt(PI); - -#pragma unroll - for (int i = 0; i < 4; ++i) { - const double erf_y = erf(yd); - const double derf_dy = 2.0 / sqrt_pi * exp(-yd * yd); - yd = yd - (erf_y - xd) / derf_dy; - } - y = static_cast(yd); - } - - return y; -} - -struct ErfinvOp { - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(T x) const { - if constexpr (std::is_same_v) { - return erfinv_impl(x); - - } else if constexpr (std::is_same_v) { - if (x == 1.0) { -#if defined(ENABLE_HYGON_API) - return __builtin_huge_val(); -#else - return std::numeric_limits::infinity(); -#endif - } - if (x == -1.0) { -#if defined(ENABLE_HYGON_API) - return -__builtin_huge_val(); -#else - return -std::numeric_limits::infinity(); -#endif - } - if (x > 1.0 || x < -1.0) { -#if defined(ENABLE_HYGON_API) - return __builtin_nan(""); -#else - return std::numeric_limits::quiet_NaN(); -#endif - } - if (x == 0.0) { - return 0.0; - } - - const double a = 0.147; - const double ln = log1p(-x * x); - const double t = 2.0 / (PI * a) + ln * 0.5; - - double inside = t * t - ln / a; - inside = inside > 0.0 ? inside : 0.0; - - double y = copysign(sqrt(sqrt(inside) - t), x); - - const int max_iter = 30; - const double tol = 1e-14; - const double sqrt_pi = sqrt(PI); - - for (int i = 0; i < max_iter; ++i) { - const double erf_y = erf(y); - const double error = erf_y - x; - if (fabs(error) < tol) { - break; - } - const double derf_dy = 2.0 / sqrt_pi * exp(-y * y); - y = y - error / derf_dy; - } - return y; - - } else { - // F16 / BF16 / other types - float xf; - if constexpr (std::is_same_v) { - xf = __half2float(x); - return __float2half_rn(erfinv_impl(xf)); - } else if constexpr (std::is_same_v) { - xf = __bfloat162float(x); - return __float2bfloat16_rn(erfinv_impl(xf)); - } else { - xf = static_cast(x); - return static_cast(erfinv_impl(xf)); - } - } - } -}; - -} // namespace op::cuda diff --git a/src/infiniop/ops/erfinv/erfinv.h b/src/infiniop/ops/erfinv/erfinv.h deleted file mode 100644 index f3ed9350f..000000000 --- a/src/infiniop/ops/erfinv/erfinv.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ERFINV_H__ -#define __ERFINV_H__ - -#include "../../elementwise/elementwise.h" - -#define DESCRIPTOR(NAMESPACE) ELEMENTWISE_DESCRIPTOR(erfinv, NAMESPACE) - -#endif // __ERFINV_H__ diff --git a/src/infiniop/ops/erfinv/metax/erfinv_metax.h b/src/infiniop/ops/erfinv/metax/erfinv_metax.h deleted file mode 100644 index 05058bfc6..000000000 --- a/src/infiniop/ops/erfinv/metax/erfinv_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ERFINV_METAX_API_H__ -#define __ERFINV_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(erfinv, metax) - -#endif // __ERFINV_METAX_API_H__ diff --git a/src/infiniop/ops/erfinv/metax/erfinv_metax.maca b/src/infiniop/ops/erfinv/metax/erfinv_metax.maca deleted file mode 100644 index 1e9144074..000000000 --- a/src/infiniop/ops/erfinv/metax/erfinv_metax.maca +++ /dev/null @@ -1,60 +0,0 @@ -#include "erfinv_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::erfinv::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::ErfinvOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::ErfinvOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::ErfinvOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::ErfinvOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::erfinv::metax diff --git a/src/infiniop/ops/erfinv/moore/erfinv_moore.h b/src/infiniop/ops/erfinv/moore/erfinv_moore.h deleted file mode 100644 index 9eed18024..000000000 --- a/src/infiniop/ops/erfinv/moore/erfinv_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ERFINV_MOORE_API_H__ -#define __ERFINV_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(erfinv, moore) - -#endif // __ERFINV_MOORE_API_H__ diff --git a/src/infiniop/ops/erfinv/moore/erfinv_moore.mu b/src/infiniop/ops/erfinv/moore/erfinv_moore.mu deleted file mode 100644 index 26febe69d..000000000 --- a/src/infiniop/ops/erfinv/moore/erfinv_moore.mu +++ /dev/null @@ -1,60 +0,0 @@ -#include "erfinv_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "../cuda/kernel.cuh" - -namespace op::erfinv::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::ErfinvOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::ErfinvOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::ErfinvOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::ErfinvOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::erfinv::moore diff --git a/src/infiniop/ops/erfinv/nvidia/erfinv_nvidia.cu b/src/infiniop/ops/erfinv/nvidia/erfinv_nvidia.cu deleted file mode 100644 index 35f5d3fe2..000000000 --- a/src/infiniop/ops/erfinv/nvidia/erfinv_nvidia.cu +++ /dev/null @@ -1,58 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "erfinv_nvidia.cuh" - -namespace op::erfinv::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::ErfinvOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::ErfinvOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::ErfinvOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::ErfinvOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::erfinv::nvidia diff --git a/src/infiniop/ops/erfinv/nvidia/erfinv_nvidia.cuh b/src/infiniop/ops/erfinv/nvidia/erfinv_nvidia.cuh deleted file mode 100644 index af80be12f..000000000 --- a/src/infiniop/ops/erfinv/nvidia/erfinv_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __ERFINV_NVIDIA_H__ -#define __ERFINV_NVIDIA_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(erfinv, nvidia) - -#endif // __ERFINV_NVIDIA_H__ diff --git a/src/infiniop/ops/erfinv/operator.cc b/src/infiniop/ops/erfinv/operator.cc deleted file mode 100644 index dc711c1ef..000000000 --- a/src/infiniop/ops/erfinv/operator.cc +++ /dev/null @@ -1,169 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/erfinv.h" - -#ifdef ENABLE_CPU_API -#include "cpu/erfinv_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/erfinv_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/erfinv_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/erfinv_moore.h" -#endif - -__INFINI_C __export infiniStatus_t infiniopCreateErfinvDescriptor( - infiniopHandle_t handle, - infiniopErfinvDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::erfinv::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - {x_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C __export infiniStatus_t infiniopGetErfinvWorkspaceSize(infiniopErfinvDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C __export infiniStatus_t infiniopErfinv( - infiniopErfinvDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, {x}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C __export infiniStatus_t -infiniopDestroyErfinvDescriptor(infiniopErfinvDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/flash_attention/ninetoothed/build.py b/src/infiniop/ops/flash_attention/ninetoothed/build.py deleted file mode 100644 index ef3235dfd..000000000 --- a/src/infiniop/ops/flash_attention/ninetoothed/build.py +++ /dev/null @@ -1,43 +0,0 @@ -import os - -import ninetoothed - -import infiniop.ninetoothed.build - -from . import flash_attention -from .flash_attention import CausalVariant - - -def build(): - - env_vars_to_check = ["MACA_HOME", "MACA_PATH", "MACA_ROOT"] - if any(var in os.environ for var in env_vars_to_check): - return - - with_kv_cache_values = (0,) - emb_dim_values = (16, 32, 64, 128, 256) - is_causal_values = (0, 1) - with_attn_mask_values = (0,) - causal_variant_values = (CausalVariant.UPPER_LEFT, CausalVariant.LOWER_RIGHT) - dtype_values = (ninetoothed.float16, ninetoothed.bfloat16, ninetoothed.float32) - block_size_m_values = (256,) - block_size_n_values = (64,) - - constexpr_param_grid = { - "with_kv_cache": with_kv_cache_values, - "emb_dim": emb_dim_values, - "is_causal": is_causal_values, - "with_attn_mask": with_attn_mask_values, - "causal_variant": causal_variant_values, - "dtype": dtype_values, - "block_size_m": block_size_m_values, - "block_size_n": block_size_n_values, - } - - infiniop.ninetoothed.build.build( - flash_attention.premake, - constexpr_param_grid, - caller="cuda", - op_name="flash_attention", - output_dir=infiniop.ninetoothed.build.BUILD_DIRECTORY_PATH, - ) diff --git a/src/infiniop/ops/flash_attention/ninetoothed/descriptor.h b/src/infiniop/ops/flash_attention/ninetoothed/descriptor.h deleted file mode 100644 index 0a6e9c1f8..000000000 --- a/src/infiniop/ops/flash_attention/ninetoothed/descriptor.h +++ /dev/null @@ -1,147 +0,0 @@ -#ifndef __FLASH_ATTENTION_DESCRIPTOR_H__ -#define __FLASH_ATTENTION_DESCRIPTOR_H__ - -#include "../../../handle.h" -#include "../../../operator.h" -#include "../../../tensor.h" - -#include "../../../../../build/ninetoothed/flash_attention.h" -#include "../../../ninetoothed/utils.h" - -namespace op::flash_attention::ninetoothed { - -class Descriptor final : public InfiniopDescriptor { -public: - Descriptor(infiniopHandle_t handle, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t total_kv_len, - double scale, - char is_causal) : InfiniopDescriptor{handle->device, handle->device_id}, - _query_shape{q_desc->shape()}, - _query_strides{q_desc->strides()}, - _key_shape{k_desc->shape()}, - _key_strides{k_desc->strides()}, - _value_shape{v_desc->shape()}, - _value_strides{v_desc->strides()}, - _total_kv_shape{total_kv_len->shape()}, - _total_kv_strides{total_kv_len->strides()}, - _output_strides{out_desc->strides()}, - _dtype{q_desc->dtype()}, - _scale{scale}, - _is_causal{is_causal} { - } - - ~Descriptor() = default; - - size_t get_workspace_size() const { - return 0; - } - - infiniStatus_t calculate(void *workspace, - size_t workspace_size, - void *out, - const void *q, - const void *k, - const void *v, - const void *total_kv_len, - void *stream) const { - uint64_t empty_shape[4]; - int64_t empty_strides[4]; - - auto query{::ninetoothed::Tensor{q, _query_shape, _query_strides}}; - auto key{::ninetoothed::Tensor{k, _key_shape, _key_strides}}; - auto value{::ninetoothed::Tensor{v, _value_shape, _value_strides}}; - auto total_kv_length{::ninetoothed::Tensor{total_kv_len, _total_kv_shape, _total_kv_strides}}; - - NineToothedTensor attn_mask{nullptr, empty_shape, empty_strides}; - NineToothedTensor is_causal; - NineToothedTensor scale{const_cast(&_scale), nullptr, nullptr}; - auto output{::ninetoothed::Tensor{out, _query_shape, _output_strides}}; - NineToothedTensor with_attn_mask; - NineToothedTensor causal_variant; - - const auto with_kv_cache_{0}; - const auto emb_dim_{_query_shape[3]}; - const auto is_causal_{_is_causal}; - const auto with_attn_mask_{0}; - const auto causal_variant_{2}; - const auto dtype_{_dtype}; - - constexpr auto block_size_m_{256}; - constexpr auto block_size_n_{64}; - - if (launch_flash_attention(stream, - query, - key, - value, - total_kv_length, - attn_mask, - is_causal, - scale, - output, - with_attn_mask, - causal_variant, - with_kv_cache_, - emb_dim_, - is_causal_, - with_attn_mask_, - causal_variant_, - dtype_, - block_size_m_, - block_size_n_)) { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - - return INFINI_STATUS_SUCCESS; - } - - static infiniStatus_t create(infiniopHandle_t handle, - Descriptor **desc, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t total_kv_len, - double scale, - char is_causal) { - *desc = new Descriptor{handle, out_desc, q_desc, k_desc, v_desc, total_kv_len, scale, is_causal}; - - return INFINI_STATUS_SUCCESS; - } - -private: - using Size = ::ninetoothed::Tensor<>::Size; - - using Stride = ::ninetoothed::Tensor<>::Stride; - - std::vector _query_shape; - - std::vector _query_strides; - - std::vector _key_shape; - - std::vector _key_strides; - - std::vector _value_shape; - - std::vector _value_strides; - - std::vector _total_kv_shape; - - std::vector _total_kv_strides; - - std::vector _output_strides; - - infiniDtype_t _dtype; - - double _scale; - - char _is_causal; -}; - -} // namespace op::flash_attention::ninetoothed - -#endif // __FLASH_ATTENTION_DESCRIPTOR_H__ diff --git a/src/infiniop/ops/flash_attention/ninetoothed/flash_attention.py b/src/infiniop/ops/flash_attention/ninetoothed/flash_attention.py deleted file mode 100644 index a4667ceb5..000000000 --- a/src/infiniop/ops/flash_attention/ninetoothed/flash_attention.py +++ /dev/null @@ -1,280 +0,0 @@ -import enum -import functools - -import ninetoothed -import ninetoothed.language as ntl -from ninetoothed import Tensor - -BLOCK_SIZE_M = ninetoothed.block_size() -BLOCK_SIZE_N = ninetoothed.block_size() - - -class CausalVariant(enum.IntEnum): - """Please refer to ``_.""" - - UPPER_LEFT = enum.auto() - - LOWER_RIGHT = enum.auto() - - -def arrangement( - query, - key, - value, - total_kv_len, - present_key, - present_value, - present_key_slot, - present_value_slot, - attn_mask, - is_causal, - scale, - output, - with_attn_mask, - causal_variant, - with_kv_cache, - block_size_m=None, - block_size_n=None, -): - def arrange_query_or_output(input): - arranged = input.tile((1, 1, block_size_m, -1)).tile( - (1, query.shape[-3] // key.shape[-3], 1, 1) - ) - arranged.dtype = arranged.dtype.squeeze((0, 2, 3)) - arranged.dtype.dtype = arranged.dtype.dtype.squeeze((0, 1)) - - return arranged - - def arrange_key_or_value(input): - arranged = ( - input.tile((1, 1, block_size_n, -1)) - .tile((1, 1, -1, -1)) - .expand((-1, -1, query_arranged.shape[-2], -1)) - ) - arranged.dtype = arranged.dtype.squeeze((0, 1, 3)) - arranged.dtype.dtype = arranged.dtype.dtype.squeeze((0, 1)) - - return arranged - - def arrange_total_kv_len(input, shape): - arranged = input.tile((1,)) - arranged = arranged.unsqueeze(1).unsqueeze(2).unsqueeze(3).expand(shape) - return arranged - - def arrange_present_key_or_present_value(input): - arranged = input.tile((1, 1, block_size_m, block_size_n)) - arranged.dtype = arranged.dtype.squeeze((0, 1)) - - return arranged - - def arrange_attn_mask(input): - arranged = input.tile((1, 1, block_size_m, block_size_n)).tile((1, 1, 1, -1)) - arranged.dtype = arranged.dtype.squeeze((0, 1, 2)) - arranged.dtype.dtype = arranged.dtype.dtype.squeeze((0, 1)) - - return arranged - - if block_size_m is None: - block_size_m = BLOCK_SIZE_M - - if block_size_n is None: - block_size_n = BLOCK_SIZE_N - - query_arranged = arrange_query_or_output(query) - key_arranged = arrange_key_or_value(key) - value_arranged = arrange_key_or_value(value) - total_kv_len_arranged = arrange_total_kv_len(total_kv_len, query_arranged.shape) - present_key_arranged = arrange_present_key_or_present_value(present_key) - present_value_arranged = arrange_present_key_or_present_value(present_value) - present_key_slot_arranged = arrange_present_key_or_present_value(present_key_slot) - present_value_slot_arranged = arrange_present_key_or_present_value( - present_value_slot - ) - attn_mask_arranged = arrange_attn_mask(attn_mask) - is_causal_arranged = is_causal - scale_arranged = scale - output_arranged = arrange_query_or_output(output) - with_attn_mask_arranged = with_attn_mask - causal_variant_arranged = causal_variant - - if with_kv_cache: - return ( - query_arranged, - key_arranged, - value_arranged, - total_kv_len_arranged, - present_key_arranged, - present_value_arranged, - present_key_slot_arranged, - present_value_slot_arranged, - attn_mask_arranged, - is_causal_arranged, - scale_arranged, - output_arranged, - with_attn_mask_arranged, - causal_variant_arranged, - ) - - return ( - query_arranged, - key_arranged, - value_arranged, - total_kv_len_arranged, - attn_mask_arranged, - is_causal_arranged, - scale_arranged, - output_arranged, - with_attn_mask_arranged, - causal_variant_arranged, - ) - - -def application_with_kv_cache( - query, - key, - value, - total_kv_len, - present_key, - present_value, - present_key_slot, - present_value_slot, - attn_mask, - is_causal, - scale, - output, - with_attn_mask, - causal_variant, -): - present_key_slot = present_key # noqa: F841 - present_value_slot = present_value # noqa: F841 - - application_without_kv_cache( - query, - key, - value, - total_kv_len, - attn_mask, - is_causal, - scale, - output, - with_attn_mask, - causal_variant, - ) - - -def application_without_kv_cache( - query, - key, - value, - total_kv_len, - attn_mask, - is_causal, - scale, - output, - with_attn_mask, - causal_variant, -): - actual_kv_len = total_kv_len[0] - - for i in range(query.shape[0]): - query_i = (1.4426950408889634 * scale * query[i]).to(query[i].dtype) - - acc = ntl.zeros((query_i.shape[-2], query_i.shape[-1]), dtype=ntl.float32) - lse = ntl.full((query_i.shape[-2],), 1, dtype=ntl.float32) - max = ntl.full((query_i.shape[-2],), float("-inf"), dtype=ntl.float32) - - for j in range(-(-actual_kv_len // key.dtype.shape[0])): - qk = ntl.dot(query_i, ntl.trans(key[j])) - - key_pos = key[j].offsets(-2) - qk = ntl.where(key_pos < actual_kv_len, qk, float("-inf")) - - if with_attn_mask: - qk += attn_mask[j] - - if is_causal: - query_pos = query[i].offsets(-2) - - if causal_variant == 2: # CausalVariant.LOWER_RIGHT: - mask = ( - query_pos[:, None] + actual_kv_len - query.source.shape[-2] - >= key_pos[None, :] - ) - else: - mask = query_pos[:, None] >= key_pos[None, :] - - qk = ntl.where(mask, qk, float("-inf")) - - next_max = ntl.maximum(max, ntl.max(qk, 1)) - stable_qk = ntl.exp2(qk - next_max[:, None]) - - alpha = ntl.exp2(max - next_max) - acc = acc * alpha[:, None] + ntl.dot(stable_qk.to(value[i].dtype), value[j]) - max = next_max - lse = lse * alpha + ntl.sum(stable_qk, 1) - - acc /= lse[:, None] - output[i] = acc # noqa: F841 - - -def premake( - with_kv_cache, - emb_dim=None, - is_causal=None, - with_attn_mask=None, - causal_variant=None, - dtype=None, - block_size_m=None, - block_size_n=None, -): - arrangement_ = functools.partial( - arrangement, - with_kv_cache=with_kv_cache, - block_size_m=block_size_m, - block_size_n=block_size_n, - ) - - query, key, value, attn_mask, output = ( - Tensor( - 4, - dtype=dtype, - shape_options=(None, None, None, {"constexpr": True, "upper_bound": 128}), - ) - for _ in range(5) - ) - total_kv_len = Tensor(1, dtype=ninetoothed.int32) - present_key, present_value, present_key_slot, present_value_slot = ( - Tensor(4, dtype=dtype) for _ in range(4) - ) - scale = Tensor(0, dtype=ninetoothed.float64) - is_causal = Tensor(0, constexpr=True, value=is_causal) - with_attn_mask = Tensor(0, constexpr=True, value=with_attn_mask) - causal_variant = Tensor(0, constexpr=True, value=causal_variant) - - if emb_dim is not None: - for tensor in (query, key, value, attn_mask, output): - tensor.shape = tensor.shape[:-1] + (emb_dim,) - - if with_kv_cache: - application = application_with_kv_cache - else: - application = application_without_kv_cache - - tensors = ( - query, - key, - value, - total_kv_len, - present_key, - present_value, - present_key_slot, - present_value_slot, - attn_mask, - is_causal, - scale, - output, - with_attn_mask, - causal_variant, - ) - - return arrangement_, application, tensors diff --git a/src/infiniop/ops/flash_attention/operator.cc b/src/infiniop/ops/flash_attention/operator.cc deleted file mode 100644 index 449c17adc..000000000 --- a/src/infiniop/ops/flash_attention/operator.cc +++ /dev/null @@ -1,121 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/flash_attention.h" - -#if defined(ENABLE_NINETOOTHED) -#if defined(ENABLE_NVIDIA_API) -#include "ninetoothed/descriptor.h" -#endif -#endif - -__INFINI_C infiniStatus_t infiniopCreateFlashAttentionDescriptor( - infiniopHandle_t handle, - infiniopFlashAttentionDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t q_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t total_kv_len, - float scale, - char is_causal) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::flash_attention::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - out_desc, \ - q_desc, \ - k_desc, \ - v_desc, \ - total_kv_len, \ - scale, \ - is_causal); - - switch (handle->device) { - -#if defined(ENABLE_NINETOOTHED) -#if defined(ENABLE_NVIDIA_API) - CREATE(INFINI_DEVICE_NVIDIA, ninetoothed); -#endif -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetFlashAttentionWorkspaceSize( - infiniopFlashAttentionDescriptor_t desc, - size_t *size) { - -#define GET_SIZE(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc) \ - ->get_workspace_size(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#if defined(ENABLE_NINETOOTHED) -#if defined(ENABLE_NVIDIA_API) - GET_SIZE(INFINI_DEVICE_NVIDIA, ninetoothed); -#endif -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET_SIZE -} - -__INFINI_C infiniStatus_t infiniopFlashAttention( - infiniopFlashAttentionDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *q, - const void *k, - const void *v, - const void *total_kv_len, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, out, q, k, v, total_kv_len, stream); - - switch (desc->device_type) { - -#if defined(ENABLE_NINETOOTHED) -#if defined(ENABLE_NVIDIA_API) - CALCULATE(INFINI_DEVICE_NVIDIA, ninetoothed); -#endif -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyFlashAttentionDescriptor( - infiniopFlashAttentionDescriptor_t desc) { - -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#if defined(ENABLE_NINETOOTHED) -#if defined(ENABLE_NVIDIA_API) - DESTROY(INFINI_DEVICE_NVIDIA, ninetoothed); -#endif -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DESTROY -} diff --git a/src/infiniop/ops/flipud/cpu/flipud_cpu.cc b/src/infiniop/ops/flipud/cpu/flipud_cpu.cc deleted file mode 100644 index 946737dec..000000000 --- a/src/infiniop/ops/flipud/cpu/flipud_cpu.cc +++ /dev/null @@ -1,170 +0,0 @@ -#include "flipud_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include -#include -#include -#include - -// 引用框架定义的 float16/bfloat16 类型支持 -#include "../../../../utils/custom_types.h" - -namespace op::flipud::cpu { - -// ================================================================== -// 0. 定义 Opaque 结构体 -// ================================================================== -struct Descriptor::Opaque { - std::vector shape; - std::vector in_strides; - std::vector out_strides; - int ndim; -}; - -// ================================================================== -// 1. 析构函数 -// ================================================================== -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - _opaque = nullptr; - } -} - -// ================================================================== -// 2. 创建描述符 -// ================================================================== -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc) { - - auto handle = reinterpret_cast(handle_); - - // 1. 创建 Info - auto result = FlipudInfo::create(out_desc, input_desc); - CHECK_RESULT(result); - - // 2. 创建并填充 Opaque - auto opaque = new Descriptor::Opaque(); - opaque->ndim = static_cast(input_desc->ndim()); - - const auto &shape = input_desc->shape(); - const auto &in_strides = input_desc->strides(); - const auto &out_strides = out_desc->strides(); - - for (int i = 0; i < opaque->ndim; ++i) { - opaque->shape.push_back(shape[i]); - opaque->in_strides.push_back(in_strides[i]); - opaque->out_strides.push_back(out_strides[i]); - } - - // 3. 创建 Descriptor - *desc_ptr = new Descriptor( - opaque, - result.take(), - 0, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -// ================================================================== -// 3. 核心计算逻辑 implementation -// ================================================================== -// [修正] 直接接收具体参数，避开 Descriptor::Opaque 的私有权限问题 -template -void calculate_cpu_impl( - int ndim, - const std::vector &shape, - const std::vector &in_strides, - const std::vector &out_strides, - size_t numel, - void *output, - const void *input) { - - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - - // 维度 0 的大小 - int64_t dim0_size = shape[0]; - -#pragma omp parallel for schedule(static) - for (ptrdiff_t i = 0; i < (ptrdiff_t)numel; ++i) { - // --- A. 坐标反解 --- - std::vector coords(ndim); - - size_t temp_idx = i; - for (int d = ndim - 1; d >= 0; --d) { - coords[d] = temp_idx % shape[d]; - temp_idx /= shape[d]; - } - - // --- B. 计算输出偏移量 --- - size_t out_offset = 0; - for (int d = 0; d < ndim; ++d) { - out_offset += coords[d] * out_strides[d]; - } - - // --- C. 翻转逻辑 (Flip Axis 0) --- - coords[0] = dim0_size - 1 - coords[0]; - - // --- D. 计算输入偏移量 --- - size_t in_offset = 0; - for (int d = 0; d < ndim; ++d) { - in_offset += coords[d] * in_strides[d]; - } - - // --- E. 数据搬运 --- - out_ptr[out_offset] = in_ptr[in_offset]; - } -} - -// ================================================================== -// 4. 执行计算 (Calculate 分发) -// ================================================================== -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) const { - - auto dtype = _info.dtype(); - size_t numel = _info.numel(); - - // 显式 Switch-Case 分发 - // 在这里解包 _opaque，因为 calculate 是成员函数，可以访问 private 的 _opaque - switch (dtype) { - case INFINI_DTYPE_F32: - cpu::calculate_cpu_impl( - _opaque->ndim, _opaque->shape, _opaque->in_strides, _opaque->out_strides, - numel, output, input); - break; - - case INFINI_DTYPE_F64: - cpu::calculate_cpu_impl( - _opaque->ndim, _opaque->shape, _opaque->in_strides, _opaque->out_strides, - numel, output, input); - break; - - case INFINI_DTYPE_F16: - cpu::calculate_cpu_impl( - _opaque->ndim, _opaque->shape, _opaque->in_strides, _opaque->out_strides, - numel, output, input); - break; - - case INFINI_DTYPE_BF16: - cpu::calculate_cpu_impl( - _opaque->ndim, _opaque->shape, _opaque->in_strides, _opaque->out_strides, - numel, output, input); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::flipud::cpu diff --git a/src/infiniop/ops/flipud/cpu/flipud_cpu.h b/src/infiniop/ops/flipud/cpu/flipud_cpu.h deleted file mode 100644 index a34908c30..000000000 --- a/src/infiniop/ops/flipud/cpu/flipud_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FLIPUD_CPU_H__ -#define __FLIPUD_CPU_H__ - -#include "../flipud.h" - -DESCRIPTOR(cpu) - -#endif // __FLIPUD_CPU_H__ diff --git a/src/infiniop/ops/flipud/cuda/kernel.cuh b/src/infiniop/ops/flipud/cuda/kernel.cuh deleted file mode 100644 index b3905e90b..000000000 --- a/src/infiniop/ops/flipud/cuda/kernel.cuh +++ /dev/null @@ -1,91 +0,0 @@ -#ifndef __FLIPUD_CUDA_CUH__ -#define __FLIPUD_CUDA_CUH__ - -#include - -namespace op::flipud::cuda { - -constexpr int MAX_DIMS = 8; - -template -struct alignas(sizeof(T) * N) Pack { - T val[N]; -}; - -struct TensorLayout { - int ndim; - size_t shape[MAX_DIMS]; - size_t in_strides[MAX_DIMS]; - size_t out_strides[MAX_DIMS]; -}; - -__device__ __forceinline__ void index_to_coords(size_t index, const TensorLayout &layout, size_t *coords) { - size_t temp = index; -#pragma unroll - for (int i = layout.ndim - 1; i >= 0; --i) { - coords[i] = temp % layout.shape[i]; - temp /= layout.shape[i]; - } -} - -__device__ __forceinline__ size_t coords_to_offset(const size_t *coords, const size_t *strides, int ndim) { - size_t offset = 0; -#pragma unroll - for (int i = 0; i < ndim; ++i) { - offset += coords[i] * strides[i]; - } - return offset; -} - -template -__global__ void flipud_kernel( - T *__restrict__ output, - const T *__restrict__ input, - size_t numel, - TensorLayout layout) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < numel) { - size_t coords[MAX_DIMS]; - index_to_coords(idx, layout, coords); - - size_t out_offset = coords_to_offset(coords, layout.out_strides, layout.ndim); - - coords[0] = layout.shape[0] - 1 - coords[0]; - - size_t in_offset = coords_to_offset(coords, layout.in_strides, layout.ndim); - - output[out_offset] = input[in_offset]; - } -} - -template -__global__ void flipud_kernel_vectorized( - T *__restrict__ output, - const T *__restrict__ input, - size_t num_packs, - TensorLayout layout) { - - using PackType = Pack; - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_packs) { - size_t scalar_idx = idx * PackSize; - size_t coords[MAX_DIMS]; - - index_to_coords(scalar_idx, layout, coords); - - size_t out_offset = coords_to_offset(coords, layout.out_strides, layout.ndim); - - coords[0] = layout.shape[0] - 1 - coords[0]; - - size_t in_offset = coords_to_offset(coords, layout.in_strides, layout.ndim); - - *reinterpret_cast(output + out_offset) = *reinterpret_cast(input + in_offset); - } -} - -} // namespace op::flipud::cuda - -#endif // __FLIPUD_CUDA_CUH__ diff --git a/src/infiniop/ops/flipud/flipud.h b/src/infiniop/ops/flipud/flipud.h deleted file mode 100644 index 22dafceed..000000000 --- a/src/infiniop/ops/flipud/flipud.h +++ /dev/null @@ -1,48 +0,0 @@ -#ifndef __FLIPUD_H__ -#define __FLIPUD_H__ - -#include "../../operator.h" -#include "info.h" - -// 宏定义：用于生成不同命名空间下的 Descriptor 类 -// 适配 Flipud 的单输入单输出模式 -#define DESCRIPTOR(NAMESPACE) \ - namespace op::flipud::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - FlipudInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - FlipudInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t input_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *output, \ - const void *input, \ - void *stream) const; \ - }; \ - } - -#endif // __FLIPUD_H__ diff --git a/src/infiniop/ops/flipud/info.h b/src/infiniop/ops/flipud/info.h deleted file mode 100644 index 1dd9fe813..000000000 --- a/src/infiniop/ops/flipud/info.h +++ /dev/null @@ -1,59 +0,0 @@ -#ifndef __FLIPUD_INFO_H__ -#define __FLIPUD_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::flipud { - -class FlipudInfo { - FlipudInfo() = default; - -public: - int _dtype; - int _ndim; - size_t _numel; - - int dtype() const { return _dtype; } - int ndim() const { return _ndim; } - size_t numel() const { return _numel; } - - FlipudInfo(int dtype, int ndim, size_t numel) - : _dtype(dtype), _ndim(ndim), _numel(numel) {} - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc) { - - if (out_desc->dtype() != input_desc->dtype()) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - if (out_desc->ndim() != input_desc->ndim()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (input_desc->ndim() < 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const auto &in_shape = input_desc->shape(); - const auto &out_shape = out_desc->shape(); - - for (size_t i = 0; i < input_desc->ndim(); ++i) { - if (in_shape[i] != out_shape[i]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - return utils::Result(FlipudInfo{ - input_desc->dtype(), - static_cast(input_desc->ndim()), - input_desc->numel()}); - } -}; - -} // namespace op::flipud - -#endif // __FLIPUD_INFO_H__ diff --git a/src/infiniop/ops/flipud/metax/flipud_metax.h b/src/infiniop/ops/flipud/metax/flipud_metax.h deleted file mode 100644 index 2c4739c65..000000000 --- a/src/infiniop/ops/flipud/metax/flipud_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FLIPUD_METAX_API_H__ -#define __FLIPUD_METAX_API_H__ - -#include "../flipud.h" - -DESCRIPTOR(metax) - -#endif // __FLIPUD_METAX_API_H__ diff --git a/src/infiniop/ops/flipud/metax/flipud_metax.maca b/src/infiniop/ops/flipud/metax/flipud_metax.maca deleted file mode 100644 index 08864440e..000000000 --- a/src/infiniop/ops/flipud/metax/flipud_metax.maca +++ /dev/null @@ -1,255 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "flipud_metax.h" -#include -#include -#include - -namespace op::flipud::metax { - -constexpr int MAX_DIMS = 4; - -struct TensorLayout { - int ndim; - int shape[MAX_DIMS]; - int in_strides[MAX_DIMS]; - int out_strides[MAX_DIMS]; -}; - -template -struct VectorType; -template <> -struct VectorType<16> { - using type = int4; -}; // 128-bit -template <> -struct VectorType<8> { - using type = int2; -}; // 64-bit -template <> -struct VectorType<4> { - using type = int; -}; // 32-bit - -// --------------------------------------------------------- - -__device__ inline size_t get_offset(int idx, const int *strides, int ndim, const int *shape) { - size_t offset = 0; - int rem = idx; -#pragma unroll - for (int i = ndim - 1; i >= 0; --i) { - int dim_sz = shape[i]; - int pos = rem % dim_sz; - rem /= dim_sz; - offset += pos * strides[i]; - } - return offset; -} - -__device__ inline size_t get_flipud_src_offset(int idx, const int *strides, int ndim, const int *shape) { - size_t offset = 0; - int rem = idx; -#pragma unroll - for (int i = ndim - 1; i >= 0; --i) { - int dim_sz = shape[i]; - int pos = rem % dim_sz; - rem /= dim_sz; - - if (i == 0) { - pos = dim_sz - 1 - pos; - } - offset += pos * strides[i]; - } - return offset; -} - -// 标量 Kernel -template -__global__ void flipud_kernel( - T *dst, const T *src, size_t n, TensorLayout layout) { - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= n) { - return; - } - - size_t dst_off = get_offset(idx, layout.out_strides, layout.ndim, layout.shape); - size_t src_off = get_flipud_src_offset(idx, layout.in_strides, layout.ndim, layout.shape); - - dst[dst_off] = src[src_off]; -} - -// 向量化 Kernel -template -__global__ void flipud_kernel_vectorized( - T *dst, const T *src, size_t num_packs, TensorLayout layout) { - // [修正] 使用 int4/int2 替代 aligned_storage - using VecT = typename VectorType::type; - - size_t pack_idx = blockIdx.x * blockDim.x + threadIdx.x; - if (pack_idx >= num_packs) { - return; - } - - int strides_in[MAX_DIMS], strides_out[MAX_DIMS], shape[MAX_DIMS]; - -#pragma unroll - for (int i = 0; i < MAX_DIMS; ++i) { - shape[i] = layout.shape[i]; - - // [逻辑正确] 确保所有步长都除以 PackSize - if (i == layout.ndim - 1) { - strides_in[i] = 1; - strides_out[i] = 1; - } else { - strides_in[i] = layout.in_strides[i] / PackSize; - strides_out[i] = layout.out_strides[i] / PackSize; - } - } - - if (layout.ndim > 0) { - shape[layout.ndim - 1] /= PackSize; - } - - size_t dst_pack_off = get_offset(pack_idx, strides_out, layout.ndim, shape); - size_t src_pack_off = get_flipud_src_offset(pack_idx, strides_in, layout.ndim, shape); - - // 强转为向量类型进行读写 - const VecT *src_vec = reinterpret_cast(src); - VecT *dst_vec = reinterpret_cast(dst); - - dst_vec[dst_pack_off] = src_vec[src_pack_off]; -} - -static inline bool is_pointer_aligned(const void *ptr, size_t alignment) { - return reinterpret_cast(ptr) % alignment == 0; -} - -struct Descriptor::Opaque { - TensorLayout layout; -}; - -template -void launch_kernel( - void *output, const void *input, - TensorLayout layout, - size_t numel, - void *stream) { - - auto in_ptr = reinterpret_cast(input); - auto out_ptr = reinterpret_cast(output); - auto hc_stream = reinterpret_cast(stream); - - constexpr int TotalBytes = 16; - constexpr int PackSize = TotalBytes / sizeof(T); - - // ---------------- Check Vectorization ---------------- - bool is_ptr_aligned = is_pointer_aligned(output, TotalBytes) && is_pointer_aligned(input, TotalBytes); - bool is_numel_divisible = (numel % PackSize == 0); - bool is_last_dim_aligned = (layout.ndim > 0) && (layout.shape[layout.ndim - 1] % PackSize == 0); - - bool is_inner_contiguous = false; - if (layout.ndim > 0) { - if (layout.in_strides[layout.ndim - 1] == 1 && layout.out_strides[layout.ndim - 1] == 1) { - is_inner_contiguous = true; - } - } - - bool is_stride_aligned = true; - for (int i = 0; i < layout.ndim - 1; ++i) { - if (layout.in_strides[i] % PackSize != 0 || layout.out_strides[i] % PackSize != 0) { - is_stride_aligned = false; - break; - } - } - - // [逻辑正确] 1D Tensor 禁止向量化 - bool is_dim_safe = (layout.ndim > 1); - - bool can_vectorize = (PackSize > 1) && is_ptr_aligned && is_numel_divisible && is_last_dim_aligned && is_inner_contiguous && is_stride_aligned && is_dim_safe; - - if (can_vectorize) { - size_t num_packs = numel / PackSize; - size_t block_size = 256; - size_t grid_size = (num_packs + block_size - 1) / block_size; - - flipud_kernel_vectorized - <<>>(out_ptr, in_ptr, num_packs, layout); - } else { - size_t block_size = 256; - size_t grid_size = (numel + block_size - 1) / block_size; - - flipud_kernel - <<>>(out_ptr, in_ptr, numel, layout); - } -} - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t input_desc) { - - auto handle = reinterpret_cast(handle_); - auto info_result = FlipudInfo::create(out_desc, input_desc); - if (!info_result) { - return info_result.status(); - } - - auto opaque = new Opaque(); - opaque->layout.ndim = static_cast(input_desc->ndim()); - - if (opaque->layout.ndim > MAX_DIMS) { - delete opaque; - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const auto &shape = input_desc->shape(); - const auto &in_strides = input_desc->strides(); - const auto &out_strides = out_desc->strides(); - - for (int i = 0; i < opaque->layout.ndim; ++i) { - opaque->layout.shape[i] = shape[i]; - opaque->layout.in_strides[i] = in_strides[i]; - opaque->layout.out_strides[i] = out_strides[i]; - } - - *desc_ptr = new Descriptor(opaque, info_result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, void *output, - const void *input, void *stream) const { - - auto dtype = _info.dtype(); - auto numel = _info.numel(); - - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel<__half>(output, input, _opaque->layout, numel, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel<__hpcc_bfloat16>(output, input, _opaque->layout, numel, stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(output, input, _opaque->layout, numel, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, _opaque->layout, numel, stream); - break; - case INFINI_DTYPE_I32: - launch_kernel(output, input, _opaque->layout, numel, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::flipud::metax diff --git a/src/infiniop/ops/flipud/moore/flipud_moore.h b/src/infiniop/ops/flipud/moore/flipud_moore.h deleted file mode 100644 index 3fc1364c7..000000000 --- a/src/infiniop/ops/flipud/moore/flipud_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FLIPUD_MOORE_H__ -#define __FLIPUD_MOORE_H__ - -#include "../flipud.h" - -DESCRIPTOR(moore) - -#endif // __FLIPUD_MOORE_H__ diff --git a/src/infiniop/ops/flipud/moore/flipud_moore.mu b/src/infiniop/ops/flipud/moore/flipud_moore.mu deleted file mode 100644 index e2adfbe01..000000000 --- a/src/infiniop/ops/flipud/moore/flipud_moore.mu +++ /dev/null @@ -1,153 +0,0 @@ -#include "../../../devices/moore/moore_handle.h" -#include "flipud_moore.h" -#include "flipud_moore_kernel.h" -#include -#include -#include - -namespace op::flipud::moore { - -// ================================================================== -// 辅助函数 -// ================================================================== -static inline bool is_pointer_aligned(const void *ptr, size_t alignment) { - return reinterpret_cast(ptr) % alignment == 0; -} - -// ================================================================== -// Opaque 定义：存储 Tensor Layout -// ================================================================== -struct Descriptor::Opaque { - op::flipud::moore::TensorLayout layout; -}; - -// ================================================================== -// Kernel Launch Logic -// ================================================================== -template -void launch_kernel( - void *output, const void *input, - op::flipud::moore::TensorLayout layout, - size_t numel, - void *stream) { - - auto in_ptr = reinterpret_cast(input); - auto out_ptr = reinterpret_cast(output); - auto musa_stream = reinterpret_cast(stream); - - constexpr int TotalBytes = 16; // 128-bit - constexpr int PackSize = TotalBytes / sizeof(T); - - // ------------------------------------------ - // 向量化判定 (Vectorization Check) - // ------------------------------------------ - // 1. 指针地址对齐 - bool is_ptr_aligned = is_pointer_aligned(output, TotalBytes) && is_pointer_aligned(input, TotalBytes); - - // 2. 元素总数必须是 PackSize 的倍数 - bool is_numel_divisible = (numel % PackSize == 0); - - // 3. 最后一维大小必须是 PackSize 的倍数 (保证 Pack 不会跨行读取) - bool is_last_dim_aligned = (layout.ndim > 0) && (layout.shape[layout.ndim - 1] % PackSize == 0); - - // 4. 连续性条件：维度 > 1 且 最内层在内存中是连续的 (stride=1) - bool is_inner_contiguous = (layout.ndim > 1) && (layout.in_strides[layout.ndim - 1] == 1) && (layout.out_strides[layout.ndim - 1] == 1); - - // 5. 步长对齐条件: 除非是最内层维度，否则所有 Stride 都必须是 PackSize 的倍数 - // 这样保证每个 Pack 读取的起始地址都是对齐的 - bool is_stride_aligned = true; - for (int i = 0; i < layout.ndim - 1; ++i) { - if (layout.in_strides[i] % PackSize != 0 || layout.out_strides[i] % PackSize != 0) { - is_stride_aligned = false; - break; - } - } - - bool can_vectorize = (PackSize > 1) && is_ptr_aligned && is_numel_divisible && is_last_dim_aligned && is_inner_contiguous && is_stride_aligned; - - if (can_vectorize) { - size_t num_packs = numel / PackSize; - size_t block_size = 256; - size_t grid_size = (num_packs + block_size - 1) / block_size; - - op::flipud::moore::flipud_kernel_vectorized - <<>>(out_ptr, in_ptr, num_packs, layout); - } else { - size_t block_size = 256; - size_t grid_size = (numel + block_size - 1) / block_size; - - op::flipud::moore::flipud_kernel - <<>>(out_ptr, in_ptr, numel, layout); - } -} - -// ================================================================== -// Descriptor 实现 -// ================================================================== -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t input_desc) { - - auto handle = reinterpret_cast(handle_); - - auto info_result = FlipudInfo::create(out_desc, input_desc); - if (!info_result) { - return info_result.status(); - } - - auto opaque = new Opaque(); - opaque->layout.ndim = static_cast(input_desc->ndim()); - - if (opaque->layout.ndim > op::flipud::moore::MAX_DIMS) { - delete opaque; - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const auto &shape = input_desc->shape(); - const auto &in_strides = input_desc->strides(); - const auto &out_strides = out_desc->strides(); - - for (int i = 0; i < opaque->layout.ndim; ++i) { - opaque->layout.shape[i] = shape[i]; - opaque->layout.in_strides[i] = in_strides[i]; - opaque->layout.out_strides[i] = out_strides[i]; - } - - *desc_ptr = new Descriptor(opaque, info_result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, void *output, - const void *input, void *stream) const { - - auto dtype = _info.dtype(); - auto numel = _info.numel(); - - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel(output, input, _opaque->layout, numel, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel<__mt_bfloat16>(output, input, _opaque->layout, numel, stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(output, input, _opaque->layout, numel, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, _opaque->layout, numel, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::flipud::moore diff --git a/src/infiniop/ops/flipud/moore/flipud_moore_kernel.h b/src/infiniop/ops/flipud/moore/flipud_moore_kernel.h deleted file mode 100644 index 616d7efc6..000000000 --- a/src/infiniop/ops/flipud/moore/flipud_moore_kernel.h +++ /dev/null @@ -1,96 +0,0 @@ -#ifndef __FLIPUD_MOORE_KERNEL_H__ -#define __FLIPUD_MOORE_KERNEL_H__ - -#include -#include -#include -#include - -namespace op::flipud::moore { - -constexpr int MAX_DIMS = 8; - -template -struct alignas(sizeof(T) * N) Pack { - T val[N]; -}; - -struct TensorLayout { - int ndim; - size_t shape[MAX_DIMS]; - size_t in_strides[MAX_DIMS]; - size_t out_strides[MAX_DIMS]; -}; - -__device__ __forceinline__ void index_to_coords(size_t index, const TensorLayout &layout, size_t *coords) { - size_t temp = index; -#pragma unroll - for (int i = layout.ndim - 1; i >= 0; --i) { - coords[i] = temp % layout.shape[i]; - temp /= layout.shape[i]; - } -} - -__device__ __forceinline__ size_t coords_to_offset(const size_t *coords, const size_t *strides, int ndim) { - size_t offset = 0; -#pragma unroll - for (int i = 0; i < ndim; ++i) { - offset += coords[i] * strides[i]; - } - return offset; -} - -template -__global__ void flipud_kernel( - T *__restrict__ output, - const T *__restrict__ input, - size_t numel, - TensorLayout layout) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < numel) { - size_t coords[MAX_DIMS]; - index_to_coords(idx, layout, coords); - - size_t out_offset = coords_to_offset(coords, layout.out_strides, layout.ndim); - - // Flip dimension 0 - coords[0] = layout.shape[0] - 1 - coords[0]; - - size_t in_offset = coords_to_offset(coords, layout.in_strides, layout.ndim); - - output[out_offset] = input[in_offset]; - } -} - -template -__global__ void flipud_kernel_vectorized( - T *__restrict__ output, - const T *__restrict__ input, - size_t num_packs, - TensorLayout layout) { - - using PackType = Pack; - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_packs) { - size_t scalar_idx = idx * PackSize; - size_t coords[MAX_DIMS]; - - index_to_coords(scalar_idx, layout, coords); - - size_t out_offset = coords_to_offset(coords, layout.out_strides, layout.ndim); - - // Flip dimension 0 - coords[0] = layout.shape[0] - 1 - coords[0]; - - size_t in_offset = coords_to_offset(coords, layout.in_strides, layout.ndim); - - *reinterpret_cast(output + out_offset) = *reinterpret_cast(input + in_offset); - } -} - -} // namespace op::flipud::moore - -#endif // __FLIPUD_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/flipud/nvidia/flipud_nvidia.cu b/src/infiniop/ops/flipud/nvidia/flipud_nvidia.cu deleted file mode 100644 index 34b9db8dc..000000000 --- a/src/infiniop/ops/flipud/nvidia/flipud_nvidia.cu +++ /dev/null @@ -1,154 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" - -#include "../cuda/kernel.cuh" -#include "flipud_nvidia.cuh" -#include -#include -#include - -namespace op::flipud::nvidia { - -// ================================================================== -// 辅助函数 -// ================================================================== -// [修改点 1] 去掉 template ，改为普通静态函数，避免解析错误 -// [修改点 2] 重命名为 is_pointer_aligned 避免潜在的命名冲突 -static inline bool is_pointer_aligned(const void *ptr, size_t alignment) { - return reinterpret_cast(ptr) % alignment == 0; -} - -// ================================================================== -// Opaque 定义：存储 Tensor Layout -// ================================================================== -// [关键] 必须在析构函数之前定义完整结构 -struct Descriptor::Opaque { - op::flipud::cuda::TensorLayout layout; -}; - -// ================================================================== -// Kernel Launch Logic -// ================================================================== -template -void launch_kernel( - void *output, const void *input, - op::flipud::cuda::TensorLayout layout, - size_t numel, - void *stream) { - - auto in_ptr = reinterpret_cast(input); - auto out_ptr = reinterpret_cast(output); - auto cuda_stream = reinterpret_cast(stream); - - constexpr int TotalBytes = 16; // 128-bit - constexpr int PackSize = TotalBytes / sizeof(T); - - // ------------------------------------------ - // 向量化判定 (Vectorization Check) - // ------------------------------------------ - bool is_ptr_aligned = is_pointer_aligned(output, TotalBytes) && is_pointer_aligned(input, TotalBytes); - - bool is_numel_divisible = (numel % PackSize == 0); - - bool is_last_dim_aligned = (layout.ndim > 0) && (layout.shape[layout.ndim - 1] % PackSize == 0); - - // 4. 连续性条件：维度 > 1 且 最内层连续 - bool is_inner_contiguous = (layout.ndim > 1) && (layout.in_strides[layout.ndim - 1] == 1) && (layout.out_strides[layout.ndim - 1] == 1); - - // 5. 步长对齐条件 - bool is_stride_aligned = true; - for (int i = 0; i < layout.ndim - 1; ++i) { - if (layout.in_strides[i] % PackSize != 0 || layout.out_strides[i] % PackSize != 0) { - is_stride_aligned = false; - break; - } - } - - bool can_vectorize = (PackSize > 1) && is_ptr_aligned && is_numel_divisible && is_last_dim_aligned && is_inner_contiguous && is_stride_aligned; - - if (can_vectorize) { - size_t num_packs = numel / PackSize; - size_t block_size = 256; - size_t grid_size = (num_packs + block_size - 1) / block_size; - - op::flipud::cuda::flipud_kernel_vectorized - <<>>(out_ptr, in_ptr, num_packs, layout); - } else { - size_t block_size = 256; - size_t grid_size = (numel + block_size - 1) / block_size; - - op::flipud::cuda::flipud_kernel - <<>>(out_ptr, in_ptr, numel, layout); - } -} - -// ================================================================== -// Descriptor 实现 -// ================================================================== -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, infiniopTensorDescriptor_t input_desc) { - - auto info_result = FlipudInfo::create(out_desc, input_desc); - if (!info_result) { - return info_result.status(); - } - - auto opaque = new Opaque(); - opaque->layout.ndim = static_cast(input_desc->ndim()); - - if (opaque->layout.ndim > op::flipud::cuda::MAX_DIMS) { - delete opaque; - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const auto &shape = input_desc->shape(); - const auto &in_strides = input_desc->strides(); - const auto &out_strides = out_desc->strides(); - - for (int i = 0; i < opaque->layout.ndim; ++i) { - opaque->layout.shape[i] = shape[i]; - opaque->layout.in_strides[i] = in_strides[i]; - opaque->layout.out_strides[i] = out_strides[i]; - } - - *desc_ptr = new Descriptor(opaque, info_result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, void *output, - const void *input, void *stream) const { - - auto dtype = _info.dtype(); - auto numel = _info.numel(); - - // 显式 Switch-Case 分发 - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel(output, input, _opaque->layout, numel, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, _opaque->layout, numel, stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(output, input, _opaque->layout, numel, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, _opaque->layout, numel, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::flipud::nvidia diff --git a/src/infiniop/ops/flipud/nvidia/flipud_nvidia.cuh b/src/infiniop/ops/flipud/nvidia/flipud_nvidia.cuh deleted file mode 100644 index cef257731..000000000 --- a/src/infiniop/ops/flipud/nvidia/flipud_nvidia.cuh +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __FLIPUD_NVIDIA_CUH__ -#define __FLIPUD_NVIDIA_CUH__ - -#include "../flipud.h" -DESCRIPTOR(nvidia) - -#endif // __FLIPUD_NVIDIA_CUH__ diff --git a/src/infiniop/ops/flipud/operator.cc b/src/infiniop/ops/flipud/operator.cc deleted file mode 100644 index 69a785ec6..000000000 --- a/src/infiniop/ops/flipud/operator.cc +++ /dev/null @@ -1,188 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/flipud.h" - -// --- 后端实现头文件 --- -#ifdef ENABLE_CPU_API -#include "cpu/flipud_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/flipud_nvidia.cuh" -#endif - -#ifdef ENABLE_METAX_API -#include "metax/flipud_metax.h" -#endif - -#ifdef ENABLE_MOORE_API -#include "moore/flipud_moore.h" -#endif - -extern "C" { - -// ======================================================================= -// 1. 创建算子描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopCreateFlipudDescriptor( - infiniopHandle_t handle, - infiniopFlipudDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::flipud::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output, \ - input) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ======================================================================= -// 2. 获取 Workspace 大小 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopGetFlipudWorkspaceSize(infiniopFlipudDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -// ======================================================================= -// 3. 执行计算 (Calculate) -// ======================================================================= -__INFINI_C infiniStatus_t infiniopFlipud( - infiniopFlipudDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, input, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -// ======================================================================= -// 4. 销毁描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopDestroyFlipudDescriptor(infiniopFlipudDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} - -} // extern "C" diff --git a/src/infiniop/ops/float_power/cpu/float_power_cpu.cc b/src/infiniop/ops/float_power/cpu/float_power_cpu.cc deleted file mode 100644 index f49e73394..000000000 --- a/src/infiniop/ops/float_power/cpu/float_power_cpu.cc +++ /dev/null @@ -1,147 +0,0 @@ -#include "float_power_cpu.h" -#include "../../../../utils/custom_types.h" -#include "../../../devices/cpu/common_cpu.h" -#include -#include -#include - -namespace op::float_power::cpu { - -Descriptor::~Descriptor() = default; - -// ================================================================== -// 创建描述符 -// ================================================================== -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t exponent, - float scalar_exponent) { - - auto handle = reinterpret_cast(handle_); - - // 创建 Info 对象进行校验 (Info 类已更新，支持混合精度和 Tensor 指数) - auto result = FloatPowerInfo::create(y, x, exponent, scalar_exponent); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - nullptr, - result.take(), - 0, // CPU 不需要 workspace - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -// ================================================================== -// 核心计算逻辑 -// 模板参数: T_OUT (输出类型), T_IN (输入类型) -// ================================================================== -template -void calculate_cpu_impl( - const FloatPowerInfo &info, - void *output, - const void *input, - const void *exponent_ptr) { - - size_t numel = info.num_elements(); - - // 获取指数模式 - bool is_scalar = info.is_scalar_exponent(); - float scalar_exp = info.scalar_exponent(); - - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - auto exp_ptr = reinterpret_cast(exponent_ptr); - - // 针对标量模式的简单优化标记 - bool is_square = is_scalar && (scalar_exp == 2.0f); - bool is_sqrt = is_scalar && (scalar_exp == 0.5f); - bool is_identity = is_scalar && (scalar_exp == 1.0f); - -#pragma omp parallel for schedule(static) - for (ptrdiff_t i = 0; i < ptrdiff_t(numel); ++i) { - // 1. 读取输入并转为 float - float in_val = utils::cast(in_ptr[i]); - float exp_val; - - // 2. 获取指数值 - if (is_scalar) { - exp_val = scalar_exp; - } else { - // Tensor 模式：读取对应位置的指数并转为 float - exp_val = utils::cast(exp_ptr[i]); - } - - // 3. 计算结果 - float result_val; - if (is_scalar && is_identity) { - result_val = in_val; - } else if (is_scalar && is_square) { - result_val = in_val * in_val; - } else if (is_scalar && is_sqrt) { - result_val = std::sqrt(in_val); - } else { - // 通用幂运算 - result_val = std::pow(in_val, exp_val); - } - - // 4. 转回输出类型 T_OUT 并存储 - out_ptr[i] = utils::cast(result_val); - } -} - -// ================================================================== -// 分发逻辑 -// ================================================================== -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *exponent, - void *stream) const { - - auto in_dtype = _info.input_dtype(); - auto out_dtype = _info.output_dtype(); - -// 定义内层宏：根据 Output 类型分发 -#define DISPATCH_OUT(IN_T) \ - switch (out_dtype) { \ - case INFINI_DTYPE_F32: \ - cpu::calculate_cpu_impl(_info, output, input, exponent); \ - return INFINI_STATUS_SUCCESS; \ - case INFINI_DTYPE_F64: \ - cpu::calculate_cpu_impl(_info, output, input, exponent); \ - return INFINI_STATUS_SUCCESS; \ - case INFINI_DTYPE_F16: \ - cpu::calculate_cpu_impl(_info, output, input, exponent); \ - return INFINI_STATUS_SUCCESS; \ - case INFINI_DTYPE_BF16: \ - cpu::calculate_cpu_impl(_info, output, input, exponent); \ - return INFINI_STATUS_SUCCESS; \ - default: \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } - - // 外层 Switch：根据 Input 类型分发 - switch (in_dtype) { - case INFINI_DTYPE_F32: - DISPATCH_OUT(float); - case INFINI_DTYPE_F64: - DISPATCH_OUT(double); - case INFINI_DTYPE_F16: - DISPATCH_OUT(fp16_t); - case INFINI_DTYPE_BF16: - DISPATCH_OUT(bf16_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef DISPATCH_OUT -} - -} // namespace op::float_power::cpu diff --git a/src/infiniop/ops/float_power/cpu/float_power_cpu.h b/src/infiniop/ops/float_power/cpu/float_power_cpu.h deleted file mode 100644 index 10841fbcd..000000000 --- a/src/infiniop/ops/float_power/cpu/float_power_cpu.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __FLOAT_POWER_CPU_H__ -#define __FLOAT_POWER_CPU_H__ - -#include "../float_power.h" -DESCRIPTOR(cpu) - -#endif // __FLOAT_POWER_CPU_H__ diff --git a/src/infiniop/ops/float_power/cuda/kernel.cuh b/src/infiniop/ops/float_power/cuda/kernel.cuh deleted file mode 100644 index 9749ea3ea..000000000 --- a/src/infiniop/ops/float_power/cuda/kernel.cuh +++ /dev/null @@ -1,104 +0,0 @@ -#ifndef __FLOAT_POWER_CUDA_CUH__ -#define __FLOAT_POWER_CUDA_CUH__ -#include - -namespace op::float_power::cuda { - -template -struct alignas(sizeof(T) * N) Pack { - T val[N]; -}; - -// ================================================================== -// Functor: 仅负责核心数学计算逻辑 -// ================================================================== -struct FloatPowerFunctor { - template - __device__ __forceinline__ float compute(const T_IN &input, float exponent_val) const { - // 将输入转为 float 参与计算，以保证计算精度和统一性 - float in_f = static_cast(input); - return powf(in_f, exponent_val); - } -}; -template -__global__ void float_power_kernel( - T_OUT *__restrict__ output, - const T_IN *__restrict__ input, - const T_EXP *__restrict__ exponent, - float scalar_exponent, - bool is_scalar, - size_t numel, - FloatPowerFunctor functor) { - for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - idx < numel; - idx += blockDim.x * gridDim.x) { - - float exp_val_f = is_scalar ? scalar_exponent : static_cast(exponent[idx]); - output[idx] = static_cast(functor.compute(input[idx], exp_val_f)); - } -} - -// ================================================================== -// 2. 标量模式向量化 Kernel -// ================================================================== -template -__global__ void float_power_kernel_vectorized_scalar( - T_OUT *__restrict__ output, - const T_IN *__restrict__ input, - float scalar_exponent, - size_t num_packs, - FloatPowerFunctor functor) { - - using PackTypeIn = Pack; - using PackTypeOut = Pack; - - auto in_vec = reinterpret_cast(input); - auto out_vec = reinterpret_cast(output); - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_packs) { - PackTypeIn in_pack = in_vec[idx]; - PackTypeOut out_pack; - -#pragma unroll - for (int i = 0; i < PackSize; ++i) { - out_pack.val[i] = static_cast(functor.compute(in_pack.val[i], scalar_exponent)); - } - out_vec[idx] = out_pack; - } -} -template -__global__ void float_power_kernel_vectorized_tensor( - T_OUT *__restrict__ output, - const T_IN *__restrict__ input, - const T_IN *__restrict__ exponent, - size_t num_packs, - FloatPowerFunctor functor) { - - using PackTypeIn = Pack; - using PackTypeOut = Pack; - - auto in_vec = reinterpret_cast(input); - auto exp_vec = reinterpret_cast(exponent); - auto out_vec = reinterpret_cast(output); - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_packs) { - PackTypeIn in_pack = in_vec[idx]; - PackTypeIn exp_pack = exp_vec[idx]; - PackTypeOut out_pack; - -#pragma unroll - for (int i = 0; i < PackSize; ++i) { - float e = static_cast(exp_pack.val[i]); - out_pack.val[i] = static_cast(functor.compute(in_pack.val[i], e)); - } - out_vec[idx] = out_pack; - } -} - -} // namespace op::float_power::cuda - -#endif // __FLOAT_POWER_CUDA_CUH__ diff --git a/src/infiniop/ops/float_power/float_power.h b/src/infiniop/ops/float_power/float_power.h deleted file mode 100644 index 63906fe53..000000000 --- a/src/infiniop/ops/float_power/float_power.h +++ /dev/null @@ -1,52 +0,0 @@ -#ifndef __FLOAT_POWER_H__ -#define __FLOAT_POWER_H__ - -#include "../../operator.h" -#include "info.h" - -// 宏定义：用于生成不同命名空间下的 Descriptor 类 -#define DESCRIPTOR(NAMESPACE) \ - namespace op::float_power::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - FloatPowerInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - FloatPowerInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - /* [修改] 增加 exponent 张量描述符 和 scalar_exponent */ \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t y, \ - infiniopTensorDescriptor_t x, \ - infiniopTensorDescriptor_t exponent, \ - float scalar_exponent); \ - \ - /* [修改] 增加 exponent 数据指针 */ \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *y, \ - const void *x, \ - const void *exponent, \ - void *stream) const; \ - }; \ - } - -#endif // __FLOAT_POWER_H__ diff --git a/src/infiniop/ops/float_power/info.h b/src/infiniop/ops/float_power/info.h deleted file mode 100644 index baae2cabe..000000000 --- a/src/infiniop/ops/float_power/info.h +++ /dev/null @@ -1,83 +0,0 @@ -#ifndef __FLOAT_POWER_INFO_H__ -#define __FLOAT_POWER_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::float_power { - -class FloatPowerInfo { - FloatPowerInfo() = default; - -public: - int _input_dtype; // 输入数据类型 - int _output_dtype; // 输出数据类型 - - bool _is_scalar_exponent; // 是否为标量指数 - float _scalar_exponent; // 标量指数的值 (仅当 _is_scalar_exponent 为 true 时有效) - - size_t _num_elements; // 元素总数 - - // Getters - int input_dtype() const { return _input_dtype; } - int output_dtype() const { return _output_dtype; } - bool is_scalar_exponent() const { return _is_scalar_exponent; } - float scalar_exponent() const { return _scalar_exponent; } - size_t num_elements() const { return _num_elements; } - - // 构造函数 - FloatPowerInfo(int in_dtype, int out_dtype, bool is_scalar, float scalar_exp, size_t numel) - : _input_dtype(in_dtype), _output_dtype(out_dtype), - _is_scalar_exponent(is_scalar), _scalar_exponent(scalar_exp), - _num_elements(numel) {} - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t exponent_desc, - float scalar_exponent) { - if (out_desc->ndim() != input_desc->ndim()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // 使用引用接收 vector，避免之前的编译错误 - const auto &in_shape = input_desc->shape(); - const auto &out_shape = out_desc->shape(); - size_t count = 1; - - for (size_t i = 0; i < input_desc->ndim(); ++i) { - if (in_shape[i] != out_shape[i]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - count *= in_shape[i]; - } - - // 3. 判断是标量模式还是张量模式 - bool is_scalar = (exponent_desc == nullptr); - - if (!is_scalar) { - if (exponent_desc->ndim() != input_desc->ndim()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - const auto &exp_shape = exponent_desc->shape(); - for (size_t i = 0; i < input_desc->ndim(); ++i) { - if (exp_shape[i] != in_shape[i]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - } - - // 构造 Info 对象 - return utils::Result(FloatPowerInfo{ - input_desc->dtype(), // Input Dtype - out_desc->dtype(), // Output Dtype (分开存储) - is_scalar, // Mode flag - scalar_exponent, // Scalar Value - count // Total elements - }); - } -}; - -} // namespace op::float_power - -#endif // __FLOAT_POWER_INFO_H__ diff --git a/src/infiniop/ops/float_power/metax/float_power_metax.h b/src/infiniop/ops/float_power/metax/float_power_metax.h deleted file mode 100644 index 5acf9d819..000000000 --- a/src/infiniop/ops/float_power/metax/float_power_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FLOAT_POWER_METAX_API_H__ -#define __FLOAT_POWER_METAX_API_H__ - -#include "../float_power.h" - -DESCRIPTOR(metax) - -#endif // __FLOAT_POWER_METAX_API_H__ diff --git a/src/infiniop/ops/float_power/metax/float_power_metax.maca b/src/infiniop/ops/float_power/metax/float_power_metax.maca deleted file mode 100644 index 4fbded158..000000000 --- a/src/infiniop/ops/float_power/metax/float_power_metax.maca +++ /dev/null @@ -1,307 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "float_power_metax.h" -#include -#include -#include - -namespace op::float_power::metax { - -// 基础定义: 向量化数据打包结构 -template -struct alignas(sizeof(T) * N) Pack { - T val[N]; -}; - -// Functor: 仅负责核心数学计算逻辑 -struct FloatPowerFunctor { - template - __device__ __forceinline__ float compute(const T_IN &input, float exponent_val) const { - // 将输入转为 float 参与计算 - float in_f = static_cast(input); - return powf(in_f, exponent_val); - } -}; - -// Kernel 1: 通用处理 (Grid-Stride Loop) -template -__global__ void float_power_kernel( - T_OUT *__restrict__ output, - const T_IN *__restrict__ input, - const T_EXP *__restrict__ exponent, - float scalar_exponent, - bool is_scalar, - size_t numel, - FloatPowerFunctor functor) { - - for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - idx < numel; - idx += blockDim.x * gridDim.x) { - - float exp_val_f = is_scalar ? scalar_exponent : static_cast(exponent[idx]); - output[idx] = static_cast(functor.compute(input[idx], exp_val_f)); - } -} - -// Kernel 2: 标量模式向量化 Kernel -template -__global__ void float_power_kernel_vectorized_scalar( - T_OUT *__restrict__ output, - const T_IN *__restrict__ input, - float scalar_exponent, - size_t num_packs, - FloatPowerFunctor functor) { - - using PackTypeIn = Pack; - using PackTypeOut = Pack; - - auto in_vec = reinterpret_cast(input); - auto out_vec = reinterpret_cast(output); - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_packs) { - PackTypeIn in_pack = in_vec[idx]; - PackTypeOut out_pack; - -#pragma unroll - for (int i = 0; i < PackSize; ++i) { - out_pack.val[i] = static_cast(functor.compute(in_pack.val[i], scalar_exponent)); - } - out_vec[idx] = out_pack; - } -} - -// Kernel 3: 张量模式向量化 Kernel -template -__global__ void float_power_kernel_vectorized_tensor( - T_OUT *__restrict__ output, - const T_IN *__restrict__ input, - const T_IN *__restrict__ exponent, - size_t num_packs, - FloatPowerFunctor functor) { - - using PackTypeIn = Pack; - using PackTypeOut = Pack; - - auto in_vec = reinterpret_cast(input); - auto exp_vec = reinterpret_cast(exponent); - auto out_vec = reinterpret_cast(output); - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_packs) { - PackTypeIn in_pack = in_vec[idx]; - PackTypeIn exp_pack = exp_vec[idx]; - PackTypeOut out_pack; - -#pragma unroll - for (int i = 0; i < PackSize; ++i) { - float e = static_cast(exp_pack.val[i]); - out_pack.val[i] = static_cast(functor.compute(in_pack.val[i], e)); - } - out_vec[idx] = out_pack; - } -} - -// ================================================================== -// 3. 辅助函数与 Launcher -// ================================================================== - -// 辅助函数: 检查内存地址对齐情况 -template -static inline bool is_aligned(const void *ptr, size_t alignment) { - return reinterpret_cast(ptr) % alignment == 0; -} - -// Launcher Implementation -template -void launch_kernel( - void *output, - const void *input, - const void *exponent, - const FloatPowerInfo &info, - void *stream) { - - size_t numel = info.num_elements(); - bool is_scalar = info.is_scalar_exponent(); - float scalar_exp = info.scalar_exponent(); - - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - // 假设指数 Tensor 的数据类型与输入 Tensor 一致 - auto exp_ptr = reinterpret_cast(exponent); - - auto hc_stream = reinterpret_cast(stream); - FloatPowerFunctor functor; - - // ------------------------------------------------------------------ - // 向量化分发路径 - // ------------------------------------------------------------------ - constexpr int AlignBytes = 16; - constexpr int PackSizeIn = AlignBytes / sizeof(T_IN); - - // 检查输入输出类型大小是否一致 - bool types_same_size = (sizeof(T_IN) == sizeof(T_OUT)); - - bool can_vectorize_base = types_same_size && (PackSizeIn > 1) && (numel % PackSizeIn == 0) && is_aligned(input, AlignBytes) && is_aligned(output, AlignBytes); - - if (can_vectorize_base) { - size_t num_packs = numel / PackSizeIn; - size_t block_size = 256; - size_t grid_size = (num_packs + block_size - 1) / block_size; - - if (is_scalar) { - // 路径 A1: 标量指数向量化 - float_power_kernel_vectorized_scalar - <<>>( - out_ptr, in_ptr, scalar_exp, num_packs, functor); - return; - } else if (is_aligned(exponent, AlignBytes)) { - // 路径 A2: 张量指数向量化 - float_power_kernel_vectorized_tensor - <<>>( - out_ptr, in_ptr, exp_ptr, num_packs, functor); - return; - } - } - - // ------------------------------------------------------------------ - // 通用回退路径 - // ------------------------------------------------------------------ - size_t block_size = 256; - size_t grid_size = (numel + block_size - 1) / block_size; - - float_power_kernel - <<>>( - out_ptr, in_ptr, exp_ptr, scalar_exp, is_scalar, numel, functor); -} - -// ================================================================== -// 4. Descriptor 接口实现 -// ================================================================== -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, Descriptor **desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t exponent, - float scalar_exponent) { - - auto handle = reinterpret_cast(handle_); - auto info_result = FloatPowerInfo::create(y, x, exponent, scalar_exponent); - if (!info_result) { - return info_result.status(); - } - - size_t workspace_size = 0; - *desc_ptr = new Descriptor(new Opaque(), info_result.take(), workspace_size, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, void *output, - const void *input, const void *exponent, - void *stream) const { - - auto in_dtype = _info.input_dtype(); - auto out_dtype = _info.output_dtype(); - - // ================================================================== - // 显式双重分发 (注意: half 和 cuda_bfloat16 已在上方适配) - // ================================================================== - - switch (in_dtype) { - - case INFINI_DTYPE_F32: - switch (out_dtype) { - case INFINI_DTYPE_F32: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F16: - launch_kernel<__half, float>(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, exponent, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - break; - - case INFINI_DTYPE_F64: - switch (out_dtype) { - case INFINI_DTYPE_F32: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F16: - launch_kernel<__half, double>(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, exponent, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - break; - - case INFINI_DTYPE_F16: - switch (out_dtype) { - case INFINI_DTYPE_F32: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F16: - launch_kernel<__half, __half>(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, exponent, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - break; - - case INFINI_DTYPE_BF16: - switch (out_dtype) { - case INFINI_DTYPE_F32: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F16: - launch_kernel<__half, cuda_bfloat16>(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, exponent, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::float_power::metax diff --git a/src/infiniop/ops/float_power/moore/float_power_moore.h b/src/infiniop/ops/float_power/moore/float_power_moore.h deleted file mode 100644 index 3028df383..000000000 --- a/src/infiniop/ops/float_power/moore/float_power_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FLOAT_POWER_MOORE_H__ -#define __FLOAT_POWER_MOORE_H__ - -#include "../float_power.h" - -DESCRIPTOR(moore) - -#endif // __FLOAT_POWER_MOORE_H__ diff --git a/src/infiniop/ops/float_power/moore/float_power_moore.mu b/src/infiniop/ops/float_power/moore/float_power_moore.mu deleted file mode 100644 index f79403713..000000000 --- a/src/infiniop/ops/float_power/moore/float_power_moore.mu +++ /dev/null @@ -1,203 +0,0 @@ -#include "../../../devices/moore/moore_handle.h" -#include "float_power_moore.h" -#include "float_power_moore_kernel.h" -#include -#include - -namespace op::float_power::moore { - -// ================================================================== -// 辅助函数: 检查内存地址对齐情况 -// ================================================================== -template -bool is_aligned(const void *ptr, size_t alignment) { - return reinterpret_cast(ptr) % alignment == 0; -} - -// ================================================================== -// Kernel Launch Logic -// ================================================================== -template -void launch_kernel( - void *output, - const void *input, - const void *exponent, - const FloatPowerInfo &info, - void *stream) { - - size_t numel = info.num_elements(); - bool is_scalar = info.is_scalar_exponent(); - float scalar_exp = info.scalar_exponent(); - - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - auto exp_ptr = reinterpret_cast(exponent); - - auto musa_stream = reinterpret_cast(stream); - op::float_power::moore::FloatPowerFunctor functor; - - // ------------------------------------------------------------------ - // 1. 向量化分发路径 (Vectorized Path) - // ------------------------------------------------------------------ - constexpr int AlignBytes = 16; - constexpr int PackSizeIn = AlignBytes / sizeof(T_IN); - - // 只有当输入和输出类型大小相同时，当前的 1:1 Pack 向量化逻辑才生效 - bool types_same_size = (sizeof(T_IN) == sizeof(T_OUT)); - - bool can_vectorize_base = types_same_size && (PackSizeIn > 1) && (numel % PackSizeIn == 0) && is_aligned(input, AlignBytes) && is_aligned(output, AlignBytes); - - if (can_vectorize_base) { - size_t num_packs = numel / PackSizeIn; - size_t block_size = 256; - size_t grid_size = (num_packs + block_size - 1) / block_size; - - if (is_scalar) { - // 路径 A1: 标量指数向量化 - op::float_power::moore::float_power_kernel_vectorized_scalar - <<>>( - out_ptr, in_ptr, scalar_exp, num_packs, functor); - return; - } else if (is_aligned(exponent, AlignBytes)) { - // 路径 A2: 张量指数向量化 - op::float_power::moore::float_power_kernel_vectorized_tensor - <<>>( - out_ptr, in_ptr, exp_ptr, num_packs, functor); - return; - } - } - - // ------------------------------------------------------------------ - // 2. 通用回退路径 (Fallback Path) - // ------------------------------------------------------------------ - size_t block_size = 256; - size_t grid_size = (numel + block_size - 1) / block_size; - - op::float_power::moore::float_power_kernel - <<>>( - out_ptr, in_ptr, exp_ptr, scalar_exp, is_scalar, numel, functor); -} - -// ================================================================== -// Descriptor 接口实现 -// ================================================================== -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, Descriptor **desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t exponent, - float scalar_exponent) { - - auto handle = reinterpret_cast(handle_); - - auto info_result = FloatPowerInfo::create(y, x, exponent, scalar_exponent); - if (!info_result) { - return info_result.status(); - } - - size_t workspace_size = 0; - *desc_ptr = new Descriptor(new Opaque(), info_result.take(), workspace_size, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, void *output, - const void *input, const void *exponent, - void *stream) const { - - auto in_dtype = _info.input_dtype(); - auto out_dtype = _info.output_dtype(); - - switch (in_dtype) { - - case INFINI_DTYPE_F32: - switch (out_dtype) { - case INFINI_DTYPE_F32: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F16: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel<__mt_bfloat16, float>(output, input, exponent, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - break; - - case INFINI_DTYPE_F64: - switch (out_dtype) { - case INFINI_DTYPE_F32: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F16: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel<__mt_bfloat16, double>(output, input, exponent, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - break; - - case INFINI_DTYPE_F16: - switch (out_dtype) { - case INFINI_DTYPE_F32: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F16: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel<__mt_bfloat16, half>(output, input, exponent, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - break; - - case INFINI_DTYPE_BF16: - switch (out_dtype) { - case INFINI_DTYPE_F32: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F16: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel<__mt_bfloat16, __mt_bfloat16>(output, input, exponent, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::float_power::moore diff --git a/src/infiniop/ops/float_power/moore/float_power_moore_kernel.h b/src/infiniop/ops/float_power/moore/float_power_moore_kernel.h deleted file mode 100644 index 1782f82ee..000000000 --- a/src/infiniop/ops/float_power/moore/float_power_moore_kernel.h +++ /dev/null @@ -1,147 +0,0 @@ -#ifndef __FLOAT_POWER_MOORE_KERNEL_H__ -#define __FLOAT_POWER_MOORE_KERNEL_H__ - -#include -#include -#include -#include -#include - -namespace op::float_power::moore { - -// ================================================================== -// 类型转换辅助函数 (适配 MUSA) -// ================================================================== -template -__device__ __forceinline__ float to_float(T val) { - if constexpr (std::is_same_v) { - return __half2float(val); - } else if constexpr (std::is_same_v) { - return __bfloat162float(val); - } else { - return static_cast(val); - } -} - -template -__device__ __forceinline__ T from_float(float val) { - if constexpr (std::is_same_v) { - return __float2half(val); - } else if constexpr (std::is_same_v) { - return __float2bfloat16(val); - } else { - return static_cast(val); - } -} - -// ================================================================== -// 基础定义: 向量化数据打包结构 -// ================================================================== -template -struct alignas(sizeof(T) * N) Pack { - T val[N]; -}; - -// ================================================================== -// Functor: 仅负责核心数学计算逻辑 -// ================================================================== -struct FloatPowerFunctor { - template - __device__ __forceinline__ float compute(const T_IN &input, float exponent_val) const { - // 使用 to_float 辅助函数处理 FP16/BF16 - float in_f = to_float(input); - return powf(in_f, exponent_val); - } -}; - -// ================================================================== -// 1. 通用处理 Kernel (Grid-Stride Loop) -// ================================================================== -template -__global__ void float_power_kernel( - T_OUT *__restrict__ output, - const T_IN *__restrict__ input, - const T_EXP *__restrict__ exponent, - float scalar_exponent, - bool is_scalar, - size_t numel, - FloatPowerFunctor functor) { - - // Grid-Stride Loop - for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - idx < numel; - idx += blockDim.x * gridDim.x) { - - float exp_val_f = is_scalar ? scalar_exponent : to_float(exponent[idx]); - output[idx] = from_float(functor.compute(input[idx], exp_val_f)); - } -} - -// ================================================================== -// 2. 标量模式向量化 Kernel -// ================================================================== -template -__global__ void float_power_kernel_vectorized_scalar( - T_OUT *__restrict__ output, - const T_IN *__restrict__ input, - float scalar_exponent, - size_t num_packs, - FloatPowerFunctor functor) { - - using PackTypeIn = Pack; - using PackTypeOut = Pack; - - auto in_vec = reinterpret_cast(input); - auto out_vec = reinterpret_cast(output); - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_packs) { - PackTypeIn in_pack = in_vec[idx]; - PackTypeOut out_pack; - -#pragma unroll - for (int i = 0; i < PackSize; ++i) { - out_pack.val[i] = from_float(functor.compute(in_pack.val[i], scalar_exponent)); - } - out_vec[idx] = out_pack; - } -} - -// ================================================================== -// 3. 张量模式向量化 Kernel -// ================================================================== -template -__global__ void float_power_kernel_vectorized_tensor( - T_OUT *__restrict__ output, - const T_IN *__restrict__ input, - const T_IN *__restrict__ exponent, - size_t num_packs, - FloatPowerFunctor functor) { - - using PackTypeIn = Pack; - using PackTypeOut = Pack; - - auto in_vec = reinterpret_cast(input); - auto exp_vec = reinterpret_cast(exponent); - auto out_vec = reinterpret_cast(output); - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_packs) { - PackTypeIn in_pack = in_vec[idx]; - PackTypeIn exp_pack = exp_vec[idx]; - PackTypeOut out_pack; - -#pragma unroll - for (int i = 0; i < PackSize; ++i) { - float e = to_float(exp_pack.val[i]); - out_pack.val[i] = from_float(functor.compute(in_pack.val[i], e)); - } - out_vec[idx] = out_pack; - } -} - -} // namespace op::float_power::moore - -#endif // __FLOAT_POWER_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/float_power/nvidia/float_power_nvidia.cu b/src/infiniop/ops/float_power/nvidia/float_power_nvidia.cu deleted file mode 100644 index ab1964eae..000000000 --- a/src/infiniop/ops/float_power/nvidia/float_power_nvidia.cu +++ /dev/null @@ -1,204 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" - -#include "../cuda/kernel.cuh" -#include "float_power_nvidia.cuh" -#include -#include - -namespace op::float_power::nvidia { - -// ================================================================== -// 辅助函数: 检查内存地址对齐情况 -// ================================================================== -template -bool is_aligned(const void *ptr, size_t alignment) { - return reinterpret_cast(ptr) % alignment == 0; -} -template -void launch_kernel( - void *output, - const void *input, - const void *exponent, - const FloatPowerInfo &info, - void *stream) { - - size_t numel = info.num_elements(); - bool is_scalar = info.is_scalar_exponent(); - float scalar_exp = info.scalar_exponent(); - - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - // 假设指数 Tensor 的数据类型与输入 Tensor 一致 - auto exp_ptr = reinterpret_cast(exponent); - - auto cuda_stream = reinterpret_cast(stream); - op::float_power::cuda::FloatPowerFunctor functor; - - // ------------------------------------------------------------------ - // 1. 向量化分发路径 (Vectorized Path) - // ------------------------------------------------------------------ - constexpr int AlignBytes = 16; // 16字节对齐是 CUDA 访存优化的标准 - constexpr int PackSizeIn = AlignBytes / sizeof(T_IN); - bool types_same_size = (sizeof(T_IN) == sizeof(T_OUT)); - - bool can_vectorize_base = types_same_size && (PackSizeIn > 1) && (numel % PackSizeIn == 0) && is_aligned(input, AlignBytes) && is_aligned(output, AlignBytes); - - if (can_vectorize_base) { - size_t num_packs = numel / PackSizeIn; - size_t block_size = 256; - size_t grid_size = (num_packs + block_size - 1) / block_size; - - if (is_scalar) { - // 路径 A1: 标量指数向量化（极快） - op::float_power::cuda::float_power_kernel_vectorized_scalar - <<>>( - out_ptr, in_ptr, scalar_exp, num_packs, functor); - return; - } else if (is_aligned(exponent, AlignBytes)) { - // 路径 A2: 张量指数向量化（解决 0.2x 倍速问题的核心） - op::float_power::cuda::float_power_kernel_vectorized_tensor - <<>>( - out_ptr, in_ptr, exp_ptr, num_packs, functor); - return; - } - } - - // ------------------------------------------------------------------ - // 2. 通用回退路径 (Fallback Path) - // 处理不对齐、非对称类型转换或小规模数据的场景 - // ------------------------------------------------------------------ - size_t block_size = 256; - size_t grid_size = (numel + block_size - 1) / block_size; - - op::float_power::cuda::float_power_kernel - <<>>( - out_ptr, in_ptr, exp_ptr, scalar_exp, is_scalar, numel, functor); -} - -// ================================================================== -// Descriptor 接口实现 -// ================================================================== -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, Descriptor **desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t exponent, - float scalar_exponent) { - - auto info_result = FloatPowerInfo::create(y, x, exponent, scalar_exponent); - if (!info_result) { - return info_result.status(); - } - - size_t workspace_size = 0; - *desc_ptr = new Descriptor(new Opaque(), info_result.take(), workspace_size, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, void *output, - const void *input, const void *exponent, - void *stream) const { - - auto in_dtype = _info.input_dtype(); - auto out_dtype = _info.output_dtype(); - - // ================================================================== - // 完全显式双重分发 (Fully Explicit Double Dispatch) - // ================================================================== - - switch (in_dtype) { - - case INFINI_DTYPE_F32: - switch (out_dtype) { - case INFINI_DTYPE_F32: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F16: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, exponent, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - break; - - case INFINI_DTYPE_F64: - switch (out_dtype) { - case INFINI_DTYPE_F32: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F16: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, exponent, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - break; - - case INFINI_DTYPE_F16: - switch (out_dtype) { - case INFINI_DTYPE_F32: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F16: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, exponent, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - break; - - case INFINI_DTYPE_BF16: - switch (out_dtype) { - case INFINI_DTYPE_F32: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_F16: - launch_kernel(output, input, exponent, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, exponent, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::float_power::nvidia diff --git a/src/infiniop/ops/float_power/nvidia/float_power_nvidia.cuh b/src/infiniop/ops/float_power/nvidia/float_power_nvidia.cuh deleted file mode 100644 index 071e2434a..000000000 --- a/src/infiniop/ops/float_power/nvidia/float_power_nvidia.cuh +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __FLOAT_POWER_NVIDIA_CUH__ -#define __FLOAT_POWER_NVIDIA_CUH__ - -#include "../float_power.h" -DESCRIPTOR(nvidia) - -#endif // __FLOAT_POWER_NVIDIA_CUH__ diff --git a/src/infiniop/ops/float_power/operator.cc b/src/infiniop/ops/float_power/operator.cc deleted file mode 100644 index de5485ef7..000000000 --- a/src/infiniop/ops/float_power/operator.cc +++ /dev/null @@ -1,192 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/float_power.h" - -#ifdef ENABLE_CPU_API -#include "cpu/float_power_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/float_power_nvidia.cuh" -#endif - -#ifdef ENABLE_METAX_API -#include "metax/float_power_metax.h" -#endif - -#ifdef ENABLE_MOORE_API -#include "moore/float_power_moore.h" -#endif - -extern "C" { - -// ======================================================================= -// 1. 创建算子描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopCreateFloatPowerDescriptor( - infiniopHandle_t handle, - infiniopFloatPowerDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y, - infiniopTensorDescriptor_t x, - infiniopTensorDescriptor_t exponent, - float scalar_exponent) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::float_power::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y, \ - x, \ - exponent, \ - scalar_exponent) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ======================================================================= -// 2. 获取 Workspace 大小 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopGetFloatPowerWorkspaceSize(infiniopFloatPowerDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -// ======================================================================= -// 3. 执行计算 (Calculate) -// ======================================================================= -__INFINI_C infiniStatus_t infiniopFloatPower( - infiniopFloatPowerDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - const void *exponent, // [新增参数] - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, x, exponent, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -// ======================================================================= -// 4. 销毁描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopDestroyFloatPowerDescriptor(infiniopFloatPowerDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} - -} // extern "C" diff --git a/src/infiniop/ops/floor/cpu/floor_cpu.cc b/src/infiniop/ops/floor/cpu/floor_cpu.cc deleted file mode 100644 index 6e68062d4..000000000 --- a/src/infiniop/ops/floor/cpu/floor_cpu.cc +++ /dev/null @@ -1,83 +0,0 @@ -// 【修改点 1】引用 Floor 专用的 CPU 头文件 -#include "floor_cpu.h" - -// 【修改点 2】命名空间必须是 floor，否则 operator.cc 找不到定义 -namespace op::floor::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - // 【修改点 3】Floor 算子通常支持浮点和整数 - // (整数做 floor 结果不变，但为了通用性建议加上) - CHECK_DTYPE(dtype, - INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, - INFINI_DTYPE_I8, INFINI_DTYPE_U8, - INFINI_DTYPE_I16, INFINI_DTYPE_U16, - INFINI_DTYPE_I32, INFINI_DTYPE_U32, - INFINI_DTYPE_I64, INFINI_DTYPE_U64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - // create CPU elementwise descriptor - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - // 【修改点 4】分发计算：将 GeluOp 替换为 FloorOp - switch (_dtype) { - // === 浮点类型 === - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - - // === 整数类型 (直接调用 FloorOp，因为 FloorOp 对整数是恒等映射) === - case INFINI_DTYPE_I8: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_U8: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_U16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_U32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_U64: - return _device_info->calculate(_info, output, inputs, stream); - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::floor::cpu diff --git a/src/infiniop/ops/floor/cpu/floor_cpu.h b/src/infiniop/ops/floor/cpu/floor_cpu.h deleted file mode 100644 index 5e077dd3b..000000000 --- a/src/infiniop/ops/floor/cpu/floor_cpu.h +++ /dev/null @@ -1,38 +0,0 @@ -#ifndef __FLOOR_CPU_H__ -#define __FLOOR_CPU_H__ - -// 引入基础宏定义 -#include "../../../elementwise/cpu/elementwise_cpu.h" - -// 使用宏声明 Descriptor 类 -ELEMENTWISE_DESCRIPTOR(floor, cpu) - -#include -#include - -namespace op::floor::cpu { - -typedef struct FloorOp { -public: - static constexpr size_t num_inputs = 1; - - template - T operator()(const T &x) const { - // 1. 整数类型：直接返回 - if constexpr (std::is_integral_v) { - return x; - } - // 2. 标准浮点类型 (float, double)：直接调用 std::floor，不降精度 - else if constexpr (std::is_same_v || std::is_same_v) { - return std::floor(x); - } - // 3. 半精度类型 (fp16, bf16)：先转 float 计算 - else { - return static_cast(std::floor(static_cast(x))); - } - } -} FloorOp; - -} // namespace op::floor::cpu - -#endif // __FLOOR_CPU_H__ diff --git a/src/infiniop/ops/floor/cuda/kernel.cuh b/src/infiniop/ops/floor/cuda/kernel.cuh deleted file mode 100644 index 66c4d719a..000000000 --- a/src/infiniop/ops/floor/cuda/kernel.cuh +++ /dev/null @@ -1,60 +0,0 @@ -#ifndef __FLOOR_CUDA_H__ -#define __FLOOR_CUDA_H__ - -#include -#include - -namespace op::floor::cuda { - -typedef struct FloorOp { -public: - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(const T &x) const { - - // 1. Half2 (向量化) - if constexpr (std::is_same_v) { - float2 vf = __half22float2(x); - float2 vr = make_float2(floorf(vf.x), floorf(vf.y)); - return __float22half2_rn(vr); - } - // 2. BFloat162 (向量化) - else if constexpr (std::is_same_v) { - float f0 = __bfloat162float(__low2bfloat16(x)); - float f1 = __bfloat162float(__high2bfloat16(x)); - // 已修复：使用 _rn 标准函数 - return __floats2bfloat162_rn(floorf(f0), floorf(f1)); - } - // 3. BFloat16 (标量) - else if constexpr (std::is_same_v) { - return __float2bfloat16(floorf(__bfloat162float(x))); - } - // 4. Half (标量) - else if constexpr (std::is_same_v) { - return __float2half(floorf(__half2float(x))); - } - // 5. Float - else if constexpr (std::is_same_v) { - return floorf(x); - } - // 6. Double - else if constexpr (std::is_same_v) { - // 【关键修复】使用 ::floor 避免与 namespace op::floor 冲突 - return ::floor(x); - } - // 7. 整数 - else if constexpr (std::is_integral_v) { - return x; - } - // 8. 兜底 - else { - // 【关键修复】使用 ::floor - return ::floor(x); - } - } -} FloorOp; - -} // namespace op::floor::cuda - -#endif // __FLOOR_CUDA_H__ diff --git a/src/infiniop/ops/floor/metax/floor_metax.h b/src/infiniop/ops/floor/metax/floor_metax.h deleted file mode 100644 index 3d293de25..000000000 --- a/src/infiniop/ops/floor/metax/floor_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FLOOR_METAX_API_H__ -#define __FLOOR_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(floor, metax) - -#endif // __FLOOR_METAX_API_H__ diff --git a/src/infiniop/ops/floor/metax/floor_metax.maca b/src/infiniop/ops/floor/metax/floor_metax.maca deleted file mode 100644 index 0c836c53e..000000000 --- a/src/infiniop/ops/floor/metax/floor_metax.maca +++ /dev/null @@ -1,69 +0,0 @@ -#include "floor_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" // 确保这里面定义了 cuda::FloorOp - -namespace op::floor::metax { // 1. 修改命名空间 - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - // 2. Floor 是单输入算子，只需要获取 input_desc_vec[0] - const auto &input_desc = input_desc_vec.at(0); - - const auto &out_shape = out_desc->shape(); - const auto &in_shape = input_desc->shape(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - // 4. 检查形状一致性 (Out == In) - CHECK_SAME_SHAPE(out_shape, in_shape); - - // create CUDA elementwise descriptor - // 这里的宏通常处理通用逻辑，直接复用即可 - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // 5. 调用 calculate 并传入 cuda::FloorOp - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::FloorOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::FloorOp, cuda_bfloat162>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::FloorOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::FloorOp, double>(_info, workspace, output, inputs, stream); - // 如果需要支持整数 (identity映射)，可以取消下面的注释，但前提是 cuda::FloorOp 支持整数类型 - case INFINI_DTYPE_I32: - return _device_info->calculate<256, cuda::FloorOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, cuda::FloorOp, int64_t>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::floor::metax diff --git a/src/infiniop/ops/floor/moore/floor_moore.h b/src/infiniop/ops/floor/moore/floor_moore.h deleted file mode 100644 index c7d20217c..000000000 --- a/src/infiniop/ops/floor/moore/floor_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FLOOR_MOORE_API_H__ -#define __FLOOR_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(floor, moore) - -#endif // __FLOOR_MOORE_API_H__ diff --git a/src/infiniop/ops/floor/moore/floor_moore.mu b/src/infiniop/ops/floor/moore/floor_moore.mu deleted file mode 100644 index 0c637298e..000000000 --- a/src/infiniop/ops/floor/moore/floor_moore.mu +++ /dev/null @@ -1,68 +0,0 @@ -#include "floor_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "floor_moore_kernel.h" - -namespace op::floor::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - // Floor is a unary operator, so we only look at the first input - const auto &in_desc = input_desc_vec.at(0); - const auto &out_shape = out_desc->shape(); - const auto &in_shape = in_desc->shape(); - - // Floor supports floating point types generally, and int types (though effectively no-op) - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - // Check if output shape matches input shape - CHECK_SAME_SHAPE(out_shape, in_shape); - - // create MOORE elementwise descriptor - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // Use moore::FloorOp template - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, moore::FloorOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, moore::FloorOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, moore::FloorOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, moore::FloorOp, double>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, moore::FloorOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, moore::FloorOp, int64_t>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::floor::moore diff --git a/src/infiniop/ops/floor/moore/floor_moore_kernel.h b/src/infiniop/ops/floor/moore/floor_moore_kernel.h deleted file mode 100644 index 0ffc2283c..000000000 --- a/src/infiniop/ops/floor/moore/floor_moore_kernel.h +++ /dev/null @@ -1,39 +0,0 @@ -#ifndef __FLOOR_MOORE_KERNEL_H__ -#define __FLOOR_MOORE_KERNEL_H__ - -/* - * This file contains the Floor operation implementation for the MUSA backend. - */ - -namespace op::floor::moore { -typedef struct FloorOp { -public: - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(const T &input) const { - if constexpr (std::is_same_v) { - // MUSA 环境可能缺失 __h2floor，改用拆分转 float 处理 - // 提取低位和高位浮点数 - float f1 = __low2float(input); - float f2 = __high2float(input); - // 分别向下取整，然后合并回 half2 - // 使用 __floats2half2_rn (round-to-nearest) 进行转换合并 - return __floats2half2_rn(::floorf(f1), ::floorf(f2)); - } else if constexpr (std::is_same_v) { - // MUSA 环境缺失 __hfloor，改用转 float 处理 - return __float2half(::floorf(__half2float(input))); - } else if constexpr (std::is_same_v) { - // Bfloat16 转 float 处理 - float val_f = __bfloat162float(input); - return __float2bfloat16(::floorf(val_f)); - } else if constexpr (std::is_same_v) { - return ::floorf(input); - } else { - return ::floor(input); - } - } -} FloorOp; -} // namespace op::floor::moore - -#endif // __FLOOR_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/floor/nvidia/floor_nvidia.cu b/src/infiniop/ops/floor/nvidia/floor_nvidia.cu deleted file mode 100644 index bc856e984..000000000 --- a/src/infiniop/ops/floor/nvidia/floor_nvidia.cu +++ /dev/null @@ -1,88 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -// 引入核心计算 Functor (我们在 src/infiniop/ops/floor/cuda/floor_cuda.h 中定义的) -#include "../cuda/kernel.cuh" -#include "floor_nvidia.cuh" - -namespace op::floor::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - CHECK_DTYPE(dtype, - INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, - INFINI_DTYPE_I8, INFINI_DTYPE_U8, - INFINI_DTYPE_I16, INFINI_DTYPE_U16, - INFINI_DTYPE_I32, INFINI_DTYPE_U32, - INFINI_DTYPE_I64, INFINI_DTYPE_U64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // ----------------------------------------------------------- - // 2. 算子分发：将 GeluOp 替换为 FloorOp - // 模板参数 <256, ...> 表示 CUDA Block Size - // ----------------------------------------------------------- - switch (_dtype) { - // === 浮点类型 === - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::FloorOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::FloorOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::FloorOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::FloorOp, double>(_info, workspace, output, inputs, stream); - - // === 整数类型 (调用 FloorOp 也会正确处理，直接返回原值) === - case INFINI_DTYPE_I8: - return _device_info->calculate<256, cuda::FloorOp, int8_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U8: - return _device_info->calculate<256, cuda::FloorOp, uint8_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I16: - return _device_info->calculate<256, cuda::FloorOp, int16_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U16: - return _device_info->calculate<256, cuda::FloorOp, uint16_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, cuda::FloorOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U32: - return _device_info->calculate<256, cuda::FloorOp, uint32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, cuda::FloorOp, int64_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_U64: - return _device_info->calculate<256, cuda::FloorOp, uint64_t>(_info, workspace, output, inputs, stream); - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::floor::nvidia diff --git a/src/infiniop/ops/floor/nvidia/floor_nvidia.cuh b/src/infiniop/ops/floor/nvidia/floor_nvidia.cuh deleted file mode 100644 index e81cebf05..000000000 --- a/src/infiniop/ops/floor/nvidia/floor_nvidia.cuh +++ /dev/null @@ -1,6 +0,0 @@ -#ifndef __FLOOR_NVIDIA_CUH__ -#define __FLOOR_NVIDIA_CUH__ -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" -ELEMENTWISE_DESCRIPTOR(floor, nvidia) - -#endif // __FLOOR_NVIDIA_CUH__ diff --git a/src/infiniop/ops/floor/operator.cc b/src/infiniop/ops/floor/operator.cc deleted file mode 100644 index 20a94f25c..000000000 --- a/src/infiniop/ops/floor/operator.cc +++ /dev/null @@ -1,207 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/floor.h" - -// --- 后端实现头文件 --- -#ifdef ENABLE_CPU_API -#include "cpu/floor_cpu.h" -#endif - -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/floor_nvidia.cuh" -#endif - -#ifdef ENABLE_METAX_API -#include "metax/floor_metax.h" -#endif - -// ========================================== -// 1. 添加 MOORE 头文件引用 -// ========================================== -#ifdef ENABLE_MOORE_API -#include "moore/floor_moore.h" -#endif -// ========================================== - -extern "C" { - -// ======================================================================= -// 1. 创建算子描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopCreateFloorDescriptor( - infiniopHandle_t handle, - infiniopFloorDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::floor::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output, \ - {input}) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -// ========================================== -// 添加 MOORE 分支 -// ========================================== -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ======================================================================= -// 2. 获取 Workspace 大小 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopGetFloorWorkspaceSize(infiniopFloorDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -// ========================================== -// 添加 MOORE 分支 -// ========================================== -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -// ======================================================================= -// 3. 执行计算 (Calculate) -// ======================================================================= -__INFINI_C infiniStatus_t infiniopFloor( - infiniopFloorDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, {input}, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -// ========================================== -// 添加 MOORE 分支 -// ========================================== -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -// ======================================================================= -// 4. 销毁描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopDestroyFloorDescriptor(infiniopFloorDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -// ========================================== -// 添加 MOORE 分支 -// ========================================== -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} - -} // extern "C" diff --git a/src/infiniop/ops/floor_divide/cpu/floor_divide_cpu.cc b/src/infiniop/ops/floor_divide/cpu/floor_divide_cpu.cc deleted file mode 100644 index e36261593..000000000 --- a/src/infiniop/ops/floor_divide/cpu/floor_divide_cpu.cc +++ /dev/null @@ -1,58 +0,0 @@ -#include "floor_divide_cpu.h" - -namespace op::floor_divide::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create CPU elementwise descriptor - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::floor_divide::cpu diff --git a/src/infiniop/ops/floor_divide/cpu/floor_divide_cpu.h b/src/infiniop/ops/floor_divide/cpu/floor_divide_cpu.h deleted file mode 100644 index ea828d297..000000000 --- a/src/infiniop/ops/floor_divide/cpu/floor_divide_cpu.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef __FLOOR_DIVIDE_CPU_H__ -#define __FLOOR_DIVIDE_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" -#include -#include - -ELEMENTWISE_DESCRIPTOR(floor_divide, cpu) - -namespace op::floor_divide::cpu { -typedef struct FloorDivideOp { -public: - static constexpr size_t num_inputs = 2; - template - T operator()(const T &a, const T &b) const { - if constexpr (std::is_floating_point_v) { - return std::floor(a / b); - } else { - T res = a / b; - T rem = a % b; - if (rem != 0 && ((a < 0) ^ (b < 0))) { - res -= 1; - } - return res; - } - } -} FloorDivideOp; -} // namespace op::floor_divide::cpu - -#endif // __FLOOR_DIVIDE_CPU_H__ diff --git a/src/infiniop/ops/floor_divide/cuda/kernel.cuh b/src/infiniop/ops/floor_divide/cuda/kernel.cuh deleted file mode 100644 index de3dc6f87..000000000 --- a/src/infiniop/ops/floor_divide/cuda/kernel.cuh +++ /dev/null @@ -1,35 +0,0 @@ -#ifndef __FLOOR_DIVIDE_CUDA_H__ -#define __FLOOR_DIVIDE_CUDA_H__ - -#include - -namespace op::floor_divide::cuda { -typedef struct FloorDivideOp { -public: - static constexpr size_t num_inputs = 2; - template - __device__ __forceinline__ T operator()(const T &a, const T &b) const { - if constexpr (std::is_same_v) { - return h2floor(__h2div(a, b)); - } else if constexpr (std::is_same_v) { - return hfloor(__hdiv(a, b)); - } else if constexpr (std::is_same_v) { - float val = __bfloat162float(a) / __bfloat162float(b); - return __float2bfloat16(floorf(val)); - } else if constexpr (std::is_same_v) { - return floorf(a / b); - } else if constexpr (std::is_same_v) { - return floor(a / b); - } else { - T res = a / b; - T rem = a % b; - if (rem != 0 && ((a < 0) ^ (b < 0))) { - res -= 1; - } - return res; - } - } -} FloorDivideOp; -} // namespace op::floor_divide::cuda - -#endif // __FLOOR_DIVIDE_CUDA_H__ diff --git a/src/infiniop/ops/floor_divide/metax/floor_divide_metax.h b/src/infiniop/ops/floor_divide/metax/floor_divide_metax.h deleted file mode 100644 index fa29a8b7c..000000000 --- a/src/infiniop/ops/floor_divide/metax/floor_divide_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FLOOR_DIVIDE_METAX_API_H__ -#define __FLOOR_DIVIDE_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(floor_divide, metax) - -#endif // __FLOOR_DIVIDE_METAX_API_H__ diff --git a/src/infiniop/ops/floor_divide/metax/floor_divide_metax.maca b/src/infiniop/ops/floor_divide/metax/floor_divide_metax.maca deleted file mode 100644 index bee3e4349..000000000 --- a/src/infiniop/ops/floor_divide/metax/floor_divide_metax.maca +++ /dev/null @@ -1,120 +0,0 @@ -#include "../../../elementwise/metax/elementwise_metax.h" -#include "floor_divide_metax.h" -#include - -namespace op::floor_divide::metax { - -struct FloorDivideOp { - static constexpr size_t num_inputs = 2; - - template - __device__ __forceinline__ T operator()(const T &a, const T &b) const { - // ------------------------------------------------ - // 1. Half2 向量化 - // ------------------------------------------------ - if constexpr (std::is_same_v) { - // MACA: 转为 float2 处理 - float2 fa = __half22float2(a); - float2 fb = __half22float2(b); - float2 res; - res.x = floorf(fa.x / fb.x); - res.y = floorf(fa.y / fb.y); - return __float22half2_rn(res); - } - // ------------------------------------------------ - // 2. Half 标量 - // ------------------------------------------------ - else if constexpr (std::is_same_v) { - return __float2half(floorf(__half2float(a) / __half2float(b))); - } - // ------------------------------------------------ - // 3. BFloat16 - // ------------------------------------------------ - else if constexpr (std::is_same_v) { - float val = __bfloat162float(a) / __bfloat162float(b); - return __float2bfloat16(floorf(val)); - } - // ------------------------------------------------ - // 4. Float / Double - // ------------------------------------------------ - else if constexpr (std::is_same_v) { - return floorf(a / b); - } else if constexpr (std::is_same_v) { - return floor(a / b); - } - // ------------------------------------------------ - // 5. 整数类型 (Int32 / Int64) - // ------------------------------------------------ - else { - // Python 语义: 向负无穷取整 - T res = a / b; - T rem = a % b; - if (rem != 0 && ((a < 0) ^ (b < 0))) { - res -= 1; - } - return res; - } - } -}; - -// ================================================================== -// 3. Descriptor 实现 -// ================================================================== -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create Metax elementwise descriptor - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, FloorDivideOp, __half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, FloorDivideOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, FloorDivideOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, FloorDivideOp, double>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, FloorDivideOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, FloorDivideOp, int64_t>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::floor_divide::metax diff --git a/src/infiniop/ops/floor_divide/moore/floor_divide_moore.h b/src/infiniop/ops/floor_divide/moore/floor_divide_moore.h deleted file mode 100644 index c87f8dd63..000000000 --- a/src/infiniop/ops/floor_divide/moore/floor_divide_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FLOOR_DIVIDE_MOORE_API_H__ -#define __FLOOR_DIVIDE_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(floor_divide, moore) - -#endif // __FLOOR_DIVIDE_MOORE_API_H__ diff --git a/src/infiniop/ops/floor_divide/moore/floor_divide_moore.mu b/src/infiniop/ops/floor_divide/moore/floor_divide_moore.mu deleted file mode 100644 index f707d0ac0..000000000 --- a/src/infiniop/ops/floor_divide/moore/floor_divide_moore.mu +++ /dev/null @@ -1,69 +0,0 @@ -#include "../../../devices/moore/moore_handle.h" -#include "../../../elementwise/moore/elementwise_moore.h" -#include "floor_divide_moore.h" -#include "floor_divide_moore_kernel.h" - -namespace op::floor_divide::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - // 检查支持的数据类型 - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_F64); - - // 检查 Shape 是否一致 (Moore Elementwise 框架通常也要求 Strict Shape 或由框架处理广播) - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // 创建 Moore Elementwise Descriptor - // 假设存在对应的宏 CREATE_ELEMENTWISE_MOORE_DESCRIPTOR 用于初始化 _device_info - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, FloorDivideOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - // Moore 架构通常使用 __mt_bfloat16 - return _device_info->calculate<256, FloorDivideOp, __mt_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, FloorDivideOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, FloorDivideOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, FloorDivideOp, int64_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, FloorDivideOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::floor_divide::moore diff --git a/src/infiniop/ops/floor_divide/moore/floor_divide_moore_kernel.h b/src/infiniop/ops/floor_divide/moore/floor_divide_moore_kernel.h deleted file mode 100644 index 5ea4a4ced..000000000 --- a/src/infiniop/ops/floor_divide/moore/floor_divide_moore_kernel.h +++ /dev/null @@ -1,39 +0,0 @@ -#ifndef __FLOOR_DIVIDE_MOORE_H__ -#define __FLOOR_DIVIDE_MOORE_H__ - -#include -#include -#include -#include - -namespace op::floor_divide::moore { -typedef struct FloorDivideOp { -public: - static constexpr size_t num_inputs = 2; - template - __device__ __forceinline__ T operator()(const T &a, const T &b) const { - if constexpr (std::is_same_v) { - return h2floor(__h2div(a, b)); - } else if constexpr (std::is_same_v) { - return hfloor(__hdiv(a, b)); - } else if constexpr (std::is_same_v) { - float val = __bfloat162float(a) / __bfloat162float(b); - return __float2bfloat16(floorf(val)); - } else if constexpr (std::is_same_v) { - return floorf(a / b); - } else if constexpr (std::is_same_v) { - return floor(a / b); - } else { - // Integer types - T res = a / b; - T rem = a % b; - if (rem != 0 && ((a < 0) ^ (b < 0))) { - res -= 1; - } - return res; - } - } -} FloorDivideOp; -} // namespace op::floor_divide::moore - -#endif // __FLOOR_DIVIDE_MOORE_H__ diff --git a/src/infiniop/ops/floor_divide/nvidia/floor_divide_nvidia.cu b/src/infiniop/ops/floor_divide/nvidia/floor_divide_nvidia.cu deleted file mode 100644 index 15206c9d3..000000000 --- a/src/infiniop/ops/floor_divide/nvidia/floor_divide_nvidia.cu +++ /dev/null @@ -1,65 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "floor_divide_nvidia.cuh" - -namespace op::floor_divide::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create CUDA elementwise descriptor - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::FloorDivideOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::FloorDivideOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::FloorDivideOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, cuda::FloorDivideOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, cuda::FloorDivideOp, int64_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::FloorDivideOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::floor_divide::nvidia diff --git a/src/infiniop/ops/floor_divide/nvidia/floor_divide_nvidia.cuh b/src/infiniop/ops/floor_divide/nvidia/floor_divide_nvidia.cuh deleted file mode 100644 index 1c70343cf..000000000 --- a/src/infiniop/ops/floor_divide/nvidia/floor_divide_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FLOOR_DIVIDE_CUDA_API_H__ -#define __FLOOR_DIVIDE_CUDA_API_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(floor_divide, nvidia) - -#endif // __FLOOR_DIVIDE_CUDA_API_H__ diff --git a/src/infiniop/ops/floor_divide/operator.cc b/src/infiniop/ops/floor_divide/operator.cc deleted file mode 100644 index d2644e944..000000000 --- a/src/infiniop/ops/floor_divide/operator.cc +++ /dev/null @@ -1,184 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/floor_divide.h" - -#ifdef ENABLE_CPU_API -#include "cpu/floor_divide_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/floor_divide_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/floor_divide_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/floor_divide_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateFloorDivideDescriptor( - infiniopHandle_t handle, - infiniopFloorDivideDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::floor_divide::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - c_desc, \ - {a_desc, \ - b_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetFloorDivideWorkspaceSize(infiniopFloorDivideDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopFloorDivide( - infiniopFloorDivideDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, c, {a, b}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyFloorDivideDescriptor(infiniopFloorDivideDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/fmin/cpu/fmin_cpu.cc b/src/infiniop/ops/fmin/cpu/fmin_cpu.cc deleted file mode 100644 index b9afb2eb4..000000000 --- a/src/infiniop/ops/fmin/cpu/fmin_cpu.cc +++ /dev/null @@ -1,58 +0,0 @@ -#include "fmin_cpu.h" - -namespace op::fmin::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create CPU elementwise descriptor - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::fmin::cpu diff --git a/src/infiniop/ops/fmin/cpu/fmin_cpu.h b/src/infiniop/ops/fmin/cpu/fmin_cpu.h deleted file mode 100644 index e7619c5ba..000000000 --- a/src/infiniop/ops/fmin/cpu/fmin_cpu.h +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef __FMIN_CPU_H__ -#define __FMIN_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" - -ELEMENTWISE_DESCRIPTOR(fmin, cpu) - -namespace op::fmin::cpu { -typedef struct FminOp { -public: - static constexpr size_t num_inputs = 2; - template - T operator()(const T &a, const T &b) const { - if constexpr (std::is_same_v || std::is_same_v) { - return utils::cast(std::fminf( - utils::cast(a), - utils::cast(b))); - } else if constexpr (std::is_floating_point_v) { - return std::fmin(a, b); - } else { - return std::min(a, b); - } - } -} FminOp; -} // namespace op::fmin::cpu - -#endif // __FMIN_CPU_H__ diff --git a/src/infiniop/ops/fmin/cuda/kernel.cuh b/src/infiniop/ops/fmin/cuda/kernel.cuh deleted file mode 100644 index 2915488fc..000000000 --- a/src/infiniop/ops/fmin/cuda/kernel.cuh +++ /dev/null @@ -1,33 +0,0 @@ -#ifndef __FMIN_CUDA_H__ -#define __FMIN_CUDA_H__ - -#include -namespace op::fmin::cuda { -typedef struct FminOp { -public: - static constexpr size_t num_inputs = 2; - template - __device__ __forceinline__ T operator()(const T &a, const T &b) const { -#if defined(ENABLE_ILUVATAR_API) - if constexpr (std::is_same_v) { - float a_f = __bfloat162float(a), b_f = __bfloat162float(b); - return __float2bfloat16(fminf(a_f, b_f)); - } else if constexpr (std::is_same_v) { - float a_f = __half2float(a), b_f = __half2float(b); - return __float2half(fminf(a_f, b_f)); - } -#elif defined(ENABLE_NVIDIA_API) - if constexpr (std::is_same_v || std::is_same_v) { - return __hmin(a, b); - } -#endif - if constexpr (std::is_same_v) { - return fminf(a, b); - } else { - return a < b ? a : b; - } - } -} FminOp; -} // namespace op::fmin::cuda - -#endif // __ADD_CUDA_H__ diff --git a/src/infiniop/ops/fmin/metax/fmin_metax.h b/src/infiniop/ops/fmin/metax/fmin_metax.h deleted file mode 100644 index 9348fa599..000000000 --- a/src/infiniop/ops/fmin/metax/fmin_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FMIN_METAX_API_H__ -#define __FMIN_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(fmin, metax) - -#endif // __FMIN_METAX_API_H__ diff --git a/src/infiniop/ops/fmin/metax/fmin_metax.maca b/src/infiniop/ops/fmin/metax/fmin_metax.maca deleted file mode 100644 index cebe4178a..000000000 --- a/src/infiniop/ops/fmin/metax/fmin_metax.maca +++ /dev/null @@ -1,69 +0,0 @@ -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" -#include "fmin_metax.h" - -namespace op::fmin::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t -Descriptor::create(infiniopHandle_t handle_, Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, - INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create CUDA elementwise descriptor - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::FminOp, half>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::FminOp, cuda_bfloat16>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::FminOp, float>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, cuda::FminOp, int32_t>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, cuda::FminOp, int64_t>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::FminOp, double>( - _info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::fmin::metax diff --git a/src/infiniop/ops/fmin/moore/fmin_moore.h b/src/infiniop/ops/fmin/moore/fmin_moore.h deleted file mode 100644 index f73d99e7d..000000000 --- a/src/infiniop/ops/fmin/moore/fmin_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FMIN_MOORE_API_H__ -#define __FMIN_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(fmin, moore) - -#endif // __FMIN_MOORE_API_H__ diff --git a/src/infiniop/ops/fmin/moore/fmin_moore.mu b/src/infiniop/ops/fmin/moore/fmin_moore.mu deleted file mode 100644 index d19071e06..000000000 --- a/src/infiniop/ops/fmin/moore/fmin_moore.mu +++ /dev/null @@ -1,66 +0,0 @@ -#include "fmin_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "fmin_moore_kernel.h" - -namespace op::fmin::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create MOORE elementwise descriptor - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, moore::FminOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, moore::FminOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, moore::FminOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, moore::FminOp, double>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, moore::FminOp, int32_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_I64: - return _device_info->calculate<256, moore::FminOp, int64_t>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::fmin::moore diff --git a/src/infiniop/ops/fmin/moore/fmin_moore_kernel.h b/src/infiniop/ops/fmin/moore/fmin_moore_kernel.h deleted file mode 100644 index cde58f26a..000000000 --- a/src/infiniop/ops/fmin/moore/fmin_moore_kernel.h +++ /dev/null @@ -1,27 +0,0 @@ -#ifndef __FMIN_MOORE_KERNEL_H__ -#define __FMIN_MOORE_KERNEL_H__ - -namespace op::fmin::moore { -typedef struct FminOp { -public: - static constexpr size_t num_inputs = 2; - template - __device__ __forceinline__ T operator()(const T &a, const T &b) const { - if constexpr (std::is_same_v) { - return __hmin2(a, b); - } else if constexpr (std::is_same_v) { - return __hmin(a, b); - } else if constexpr (std::is_same_v) { - float a_f = __bfloat162float(a); - float b_f = __bfloat162float(b); - return fminf(a_f, b_f); - } else if constexpr (std::is_same_v) { - return fminf(a, b); - } else { - return a < b ? a : b; - } - } -} FminOp; -} // namespace op::fmin::moore - -#endif // __FMIN_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/fmin/nvidia/fmin_nvidia.cu b/src/infiniop/ops/fmin/nvidia/fmin_nvidia.cu deleted file mode 100644 index bfe77c7bc..000000000 --- a/src/infiniop/ops/fmin/nvidia/fmin_nvidia.cu +++ /dev/null @@ -1,66 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "fmin_nvidia.cuh" - -namespace op::fmin::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t -Descriptor::create(infiniopHandle_t handle_, Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, - INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create CUDA elementwise descriptor - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::FminOp, half>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::FminOp, cuda_bfloat16>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::FminOp, float>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_I32: - return _device_info->calculate<256, cuda::FminOp, int32_t>( - _info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::FminOp, double>( - _info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::fmin::nvidia diff --git a/src/infiniop/ops/fmin/nvidia/fmin_nvidia.cuh b/src/infiniop/ops/fmin/nvidia/fmin_nvidia.cuh deleted file mode 100644 index f9def467b..000000000 --- a/src/infiniop/ops/fmin/nvidia/fmin_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FMIN_CUDA_API_H__ -#define __FMIN_CUDA_API_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(fmin, nvidia) - -#endif // __FMIN_CUDA_API_H__ diff --git a/src/infiniop/ops/fmin/operator.cc b/src/infiniop/ops/fmin/operator.cc deleted file mode 100644 index 4b41c93ae..000000000 --- a/src/infiniop/ops/fmin/operator.cc +++ /dev/null @@ -1,184 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/fmin.h" - -#ifdef ENABLE_CPU_API -#include "cpu/fmin_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/fmin_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/fmin_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/fmin_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateFminDescriptor( - infiniopHandle_t handle, - infiniopFminDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::fmin::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - c_desc, \ - {a_desc, \ - b_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetFminWorkspaceSize(infiniopFminDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopFmin( - infiniopFminDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, c, {a, b}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyFminDescriptor(infiniopFminDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/fmod/cpu/fmod_cpu.cc b/src/infiniop/ops/fmod/cpu/fmod_cpu.cc deleted file mode 100644 index adb6b43e6..000000000 --- a/src/infiniop/ops/fmod/cpu/fmod_cpu.cc +++ /dev/null @@ -1,53 +0,0 @@ -#include "fmod_cpu.h" - -namespace op::fmod__::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &out_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(out_shape, a_shape, b_shape); - - // create CPU elementwise descriptor - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} -} // namespace op::fmod__::cpu diff --git a/src/infiniop/ops/fmod/cpu/fmod_cpu.h b/src/infiniop/ops/fmod/cpu/fmod_cpu.h deleted file mode 100644 index f91622dbd..000000000 --- a/src/infiniop/ops/fmod/cpu/fmod_cpu.h +++ /dev/null @@ -1,19 +0,0 @@ -#ifndef _FMOD_CPU_H__ -#define _FMOD_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" - -ELEMENTWISE_DESCRIPTOR(fmod__, cpu) - -namespace op::fmod__::cpu { -typedef struct FmodOp { -public: - static constexpr size_t num_inputs = 2; - template - T operator()(const T &a, const T &b) const { - return std::fmod(a, b); - } -} FmodOp; -} // namespace op::fmod__::cpu - -#endif // _FMOD_CPU_H__ diff --git a/src/infiniop/ops/fmod/cuda/kernel.cuh b/src/infiniop/ops/fmod/cuda/kernel.cuh deleted file mode 100644 index 0830251bc..000000000 --- a/src/infiniop/ops/fmod/cuda/kernel.cuh +++ /dev/null @@ -1,48 +0,0 @@ -#ifndef __FMOD_CUDA_H__ -#define __FMOD_CUDA_H__ - -namespace op::fmod__::cuda { -typedef struct FmodOp { - static constexpr size_t num_inputs = 2; - template - __device__ __forceinline__ T operator()(const T &a, const T &b) const { - // fmod(a, b) = a - b * trunc(a / b) - if constexpr (std::is_same_v) { - // 对于 half2，转换为 float 计算后再转回 - float2 af = __half22float2(a); - float2 bf = __half22float2(b); - float2 result; - result.x = fmodf(af.x, bf.x); - result.y = fmodf(af.y, bf.y); - return __float22half2_rn(result); - } else if constexpr (std::is_same_v) { - // 对于 bfloat162，转换为 float 计算后再转回 - float af_low = __bfloat162float(__low2bfloat16(a)); - float af_high = __bfloat162float(__high2bfloat16(a)); - float bf_low = __bfloat162float(__low2bfloat16(b)); - float bf_high = __bfloat162float(__high2bfloat16(b)); - return __floats2bfloat162_rn(fmodf(af_low, bf_low), fmodf(af_high, bf_high)); - } else if constexpr (std::is_same_v) { - // 对于 half，转换为 float 计算后再转回 - float af = __half2float(a); - float bf = __half2float(b); - return __float2half(fmodf(af, bf)); - } else if constexpr (std::is_same_v) { - // 对于 bfloat16，转换为 float 计算后再转回 - float af = __bfloat162float(a); - float bf = __bfloat162float(b); - return __float2bfloat16(fmodf(af, bf)); - } else if constexpr (std::is_same_v) { - return fmodf(a, b); - } else if constexpr (std::is_same_v) { - return fmod(a, b); - } else { - // 整数类型使用 % 运算符 - return a % b; - } - } -} FmodOp; - -} // namespace op::fmod__::cuda - -#endif // __FMOD_CUDA_H__ diff --git a/src/infiniop/ops/fmod/metax/fmod_metax.h b/src/infiniop/ops/fmod/metax/fmod_metax.h deleted file mode 100644 index 27f628c68..000000000 --- a/src/infiniop/ops/fmod/metax/fmod_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FMOD_METAX_API_H__ -#define __FMOD_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(fmod__, metax) - -#endif // __FMOD_METAX_API_H__ diff --git a/src/infiniop/ops/fmod/metax/mul_metax.maca b/src/infiniop/ops/fmod/metax/mul_metax.maca deleted file mode 100644 index 4edb1e770..000000000 --- a/src/infiniop/ops/fmod/metax/mul_metax.maca +++ /dev/null @@ -1,61 +0,0 @@ -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -#include "fmod_metax.h" - -namespace op::fmod__::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::FmodOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::FmodOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::FmodOp, double>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::FmodOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::fmod__::metax diff --git a/src/infiniop/ops/fmod/moore/fmod_moore.h b/src/infiniop/ops/fmod/moore/fmod_moore.h deleted file mode 100644 index 090ed6efa..000000000 --- a/src/infiniop/ops/fmod/moore/fmod_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FMOD_MOORE_API_H__ -#define __FMOD_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(fmod__, moore) - -#endif // __FMOD_MOORE_API_H__ diff --git a/src/infiniop/ops/fmod/moore/fmod_moore.mu b/src/infiniop/ops/fmod/moore/fmod_moore.mu deleted file mode 100644 index 5c076b4fc..000000000 --- a/src/infiniop/ops/fmod/moore/fmod_moore.mu +++ /dev/null @@ -1,63 +0,0 @@ -#include "fmod_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "../cuda/kernel.cuh" - -namespace op::fmod__::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - // create MOORE elementwise descriptor - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::FmodOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::FmodOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::FmodOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::FmodOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::fmod__::moore diff --git a/src/infiniop/ops/fmod/nvidia/fmod_nvidia.cu b/src/infiniop/ops/fmod/nvidia/fmod_nvidia.cu deleted file mode 100644 index 669742848..000000000 --- a/src/infiniop/ops/fmod/nvidia/fmod_nvidia.cu +++ /dev/null @@ -1,59 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" -#include "../cuda/kernel.cuh" -#include "fmod_nvidia.cuh" - -namespace op::fmod__::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &a_desc = input_desc_vec.at(0); - const auto &b_desc = input_desc_vec.at(1); - const auto &c_shape = out_desc->shape(); - const auto &a_shape = a_desc->shape(); - const auto &b_shape = b_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(c_shape, a_shape, b_shape); - - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::FmodOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::FmodOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::FmodOp, double>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::FmodOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::fmod__::nvidia diff --git a/src/infiniop/ops/fmod/nvidia/fmod_nvidia.cuh b/src/infiniop/ops/fmod/nvidia/fmod_nvidia.cuh deleted file mode 100644 index 21009d104..000000000 --- a/src/infiniop/ops/fmod/nvidia/fmod_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FMOD_CUDA_API_H__ -#define __FMOD_CUDA_API_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(fmod__, nvidia) - -#endif // __FMOD_CUDA_API_H__ diff --git a/src/infiniop/ops/fmod/operator.cc b/src/infiniop/ops/fmod/operator.cc deleted file mode 100644 index 100257d38..000000000 --- a/src/infiniop/ops/fmod/operator.cc +++ /dev/null @@ -1,161 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/fmod.h" - -#ifdef ENABLE_CPU_API -#include "cpu/fmod_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) -#include "nvidia/fmod_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/fmod_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/fmod_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateFmodDescriptor( - infiniopHandle_t handle, - infiniopFmodDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::fmod__::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - c_desc, \ - {a_desc, \ - b_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetFmodWorkspaceSize(infiniopFmodDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopFmod( - infiniopFmodDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *c, - const void *a, - const void *b, - void *stream) { -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, c, {a, b}, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyFmodDescriptor(infiniopFmodDescriptor_t desc) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} diff --git a/src/infiniop/ops/fused_ffn/cpu/fused_ffn_cpu.cc b/src/infiniop/ops/fused_ffn/cpu/fused_ffn_cpu.cc deleted file mode 100644 index cf87c7630..000000000 --- a/src/infiniop/ops/fused_ffn/cpu/fused_ffn_cpu.cc +++ /dev/null @@ -1,209 +0,0 @@ -#include "fused_ffn_cpu.h" -#include "../../../../utils.h" -#include "../../../devices/cpu/common_cpu.h" -#include -#include - -namespace op::fused_ffn::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - infiniopTensorDescriptor_t residual_desc, - infiniopTensorDescriptor_t norm_weight_desc, - infiniopTensorDescriptor_t gate_up_weight_desc, - infiniopTensorDescriptor_t down_weight_desc, - float epsilon) { - - auto result = FusedFFNInfo::create( - out_desc, in_desc, residual_desc, - norm_weight_desc, gate_up_weight_desc, down_weight_desc, epsilon); - CHECK_RESULT(result); - auto info = result.take(); - - // Workspace size (same as NVIDIA implementation) - size_t dtype_size = infiniSizeOf(info.dtype); - size_t ntok = info.ntok(); - size_t d = info.d(); - size_t di = info.di(); - - size_t normalized_size = ntok * d * dtype_size; - size_t gate_up_size = ntok * 2 * di * dtype_size; - - size_t workspace_size = normalized_size + gate_up_size; - - *desc_ptr = new Descriptor( - nullptr, - std::move(info), - workspace_size, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculateTyped( - const FusedFFNInfo &info, - void *workspace, size_t workspace_size, - void *out, - const void *in, - const void *residual, - const void *norm_weight, - const void *gate_up_weight, - const void *down_weight) { - - size_t ntok = info.ntok(); - size_t d = info.d(); - size_t di = info.di(); - - // Partition workspace (no separate hidden_buf needed, SwiGLU is in-place) - char *ws_ptr = static_cast(workspace); - Tdata *normalized_buf = reinterpret_cast(ws_ptr); - ws_ptr += ntok * d * sizeof(Tdata); - Tdata *gate_up_buf = reinterpret_cast(ws_ptr); - - const Tdata *in_ptr = reinterpret_cast(in); - const Tdata *residual_ptr = reinterpret_cast(residual); - const TnormWeight *norm_w_ptr = reinterpret_cast(norm_weight); - const TmatWeight *gate_up_w_ptr = reinterpret_cast(gate_up_weight); - const TmatWeight *down_w_ptr = reinterpret_cast(down_weight); - Tdata *out_ptr = reinterpret_cast(out); - - // Stage 1: RMSNorm - for (size_t t = 0; t < ntok; t++) { - const Tdata *x = in_ptr + t * info.in_stride; - Tdata *norm = normalized_buf + t * d; - - // Compute variance - float sum_sq = 0.0f; - for (size_t i = 0; i < d; i++) { - float val = utils::cast(x[i]); - sum_sq += val * val; - } - - // Normalize - float rms = 1.0f / std::sqrt(sum_sq / d + info.epsilon); - for (size_t i = 0; i < d; i++) { - float val = utils::cast(x[i]) * utils::cast(norm_w_ptr[i]) * rms; - norm[i] = utils::cast(val); - } - } - - // Stage 2: GateUp GEMM (C = A @ B^T) - // normalized: [ntok, d], gate_up_weight: [2*di, d] -> gate_up: [ntok, 2*di] - for (size_t t = 0; t < ntok; t++) { - const Tdata *norm = normalized_buf + t * d; - Tdata *gate_up = gate_up_buf + t * 2 * di; - - for (size_t j = 0; j < 2 * di; j++) { - float sum = 0.0f; - for (size_t k = 0; k < d; k++) { - sum += utils::cast(norm[k]) * utils::cast(gate_up_w_ptr[j * d + k]); - } - gate_up[j] = utils::cast(sum); - } - } - - // Stage 3: SwiGLU (in-place, overwrites gate half of gate_up_buf) - for (size_t t = 0; t < ntok; t++) { - Tdata *gate_up = gate_up_buf + t * 2 * di; - - for (size_t i = 0; i < di; i++) { - float gate = utils::cast(gate_up[i]); - float up = utils::cast(gate_up[di + i]); - // SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x)) - float silu = gate / (1.0f + std::exp(-gate)); - gate_up[i] = utils::cast(silu * up); - } - } - - // Stage 4: Down GEMM (C = A @ B^T) + Residual Add (fused) - // Read from gate_up_buf (stride 2*di) to match non-fused path's buffer layout - { - bool fuse_residual = info.has_residual && (out_ptr == residual_ptr); - for (size_t t = 0; t < ntok; t++) { - const Tdata *hidden = gate_up_buf + t * 2 * di; // stride = 2*di to match non-fused - Tdata *o = out_ptr + t * info.out_stride; - - if (fuse_residual) { - const Tdata *res = residual_ptr + t * info.residual_stride; - for (size_t j = 0; j < d; j++) { - float sum = utils::cast(res[j]); - for (size_t k = 0; k < di; k++) { - sum += utils::cast(hidden[k]) * utils::cast(down_w_ptr[j * di + k]); - } - o[j] = utils::cast(sum); - } - } else { - for (size_t j = 0; j < d; j++) { - float sum = 0.0f; - for (size_t k = 0; k < di; k++) { - sum += utils::cast(hidden[k]) * utils::cast(down_w_ptr[j * di + k]); - } - o[j] = utils::cast(sum); - } - } - } - } - - // Stage 5: Residual Add (only when not fused into GEMM) - if (info.has_residual && out_ptr != residual_ptr) { - for (size_t t = 0; t < ntok; t++) { - Tdata *o = out_ptr + t * info.out_stride; - const Tdata *res = residual_ptr + t * info.residual_stride; - for (size_t i = 0; i < d; i++) { - float val = utils::cast(o[i]) + utils::cast(res[i]); - o[i] = utils::cast(val); - } - } - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *out, - const void *in, - const void *residual, - const void *norm_weight, - const void *gate_up_weight, - const void *down_weight, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // Dispatch based on dtype, wtype (norm weight), and mtype (matrix weight) - if (_info.dtype == INFINI_DTYPE_F16) { - if (_info.wtype == INFINI_DTYPE_F16 && _info.mtype == INFINI_DTYPE_F16) { - return calculateTyped(_info, workspace, workspace_size, out, in, residual, norm_weight, gate_up_weight, down_weight); - } else if (_info.wtype == INFINI_DTYPE_F32 && _info.mtype == INFINI_DTYPE_F16) { - return calculateTyped(_info, workspace, workspace_size, out, in, residual, norm_weight, gate_up_weight, down_weight); - } else if (_info.wtype == INFINI_DTYPE_F16 && _info.mtype == INFINI_DTYPE_F32) { - return calculateTyped(_info, workspace, workspace_size, out, in, residual, norm_weight, gate_up_weight, down_weight); - } else if (_info.wtype == INFINI_DTYPE_F32 && _info.mtype == INFINI_DTYPE_F32) { - return calculateTyped(_info, workspace, workspace_size, out, in, residual, norm_weight, gate_up_weight, down_weight); - } - } else if (_info.dtype == INFINI_DTYPE_BF16) { - if (_info.wtype == INFINI_DTYPE_BF16 && _info.mtype == INFINI_DTYPE_BF16) { - return calculateTyped(_info, workspace, workspace_size, out, in, residual, norm_weight, gate_up_weight, down_weight); - } else if (_info.wtype == INFINI_DTYPE_F32 && _info.mtype == INFINI_DTYPE_BF16) { - return calculateTyped(_info, workspace, workspace_size, out, in, residual, norm_weight, gate_up_weight, down_weight); - } else if (_info.wtype == INFINI_DTYPE_BF16 && _info.mtype == INFINI_DTYPE_F32) { - return calculateTyped(_info, workspace, workspace_size, out, in, residual, norm_weight, gate_up_weight, down_weight); - } else if (_info.wtype == INFINI_DTYPE_F32 && _info.mtype == INFINI_DTYPE_F32) { - return calculateTyped(_info, workspace, workspace_size, out, in, residual, norm_weight, gate_up_weight, down_weight); - } - } else if (_info.dtype == INFINI_DTYPE_F32) { - return calculateTyped(_info, workspace, workspace_size, out, in, residual, norm_weight, gate_up_weight, down_weight); - } - - return INFINI_STATUS_BAD_TENSOR_DTYPE; -} - -} // namespace op::fused_ffn::cpu diff --git a/src/infiniop/ops/fused_ffn/cpu/fused_ffn_cpu.h b/src/infiniop/ops/fused_ffn/cpu/fused_ffn_cpu.h deleted file mode 100644 index df7c450fb..000000000 --- a/src/infiniop/ops/fused_ffn/cpu/fused_ffn_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FUSED_FFN_CPU_H__ -#define __FUSED_FFN_CPU_H__ - -#include "../fused_ffn.h" - -DESCRIPTOR(cpu) - -#endif diff --git a/src/infiniop/ops/fused_ffn/fused_ffn.h b/src/infiniop/ops/fused_ffn/fused_ffn.h deleted file mode 100644 index 209bea1dc..000000000 --- a/src/infiniop/ops/fused_ffn/fused_ffn.h +++ /dev/null @@ -1,55 +0,0 @@ -#ifndef FUSED_FFN_H -#define FUSED_FFN_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::fused_ffn::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - FusedFFNInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - FusedFFNInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t in_desc, \ - infiniopTensorDescriptor_t residual_desc, \ - infiniopTensorDescriptor_t norm_weight_desc, \ - infiniopTensorDescriptor_t gate_up_weight_desc, \ - infiniopTensorDescriptor_t down_weight_desc, \ - float epsilon); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *out, \ - const void *in, \ - const void *residual, \ - const void *norm_weight, \ - const void *gate_up_weight, \ - const void *down_weight, \ - void *stream) const; \ - }; \ - } - -#endif // FUSED_FFN_H diff --git a/src/infiniop/ops/fused_ffn/info.h b/src/infiniop/ops/fused_ffn/info.h deleted file mode 100644 index 35298f11f..000000000 --- a/src/infiniop/ops/fused_ffn/info.h +++ /dev/null @@ -1,180 +0,0 @@ -#ifndef __FUSED_FFN_INFO_H__ -#define __FUSED_FFN_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::fused_ffn { - -class FusedFFNInfo { - FusedFFNInfo() = default; - -public: - infiniDtype_t dtype; - infiniDtype_t wtype; // norm weight dtype - infiniDtype_t mtype; // matrix weight dtype (gate_up, down) - float epsilon; - std::vector shape; - ptrdiff_t in_stride; - ptrdiff_t out_stride; - ptrdiff_t residual_stride; - size_t hidden_dim; - size_t intermediate_dim; - bool has_residual; - - size_t ntok() const { return shape[0]; } - size_t d() const { return hidden_dim; } - size_t di() const { return intermediate_dim; } - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - infiniopTensorDescriptor_t residual_desc, - infiniopTensorDescriptor_t norm_weight_desc, - infiniopTensorDescriptor_t gate_up_weight_desc, - infiniopTensorDescriptor_t down_weight_desc, - float epsilon) { - - auto dtype = out_desc->dtype(); - auto wtype = norm_weight_desc->dtype(); - auto mtype = gate_up_weight_desc->dtype(); - - // Check that input and output have the same dtype - if (in_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // Check matrix weight dtypes (gate_up and down must match each other) - if (down_weight_desc->dtype() != mtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // Validate norm weight dtype: for half-precision activations, can be same or FP32 - if (dtype == INFINI_DTYPE_F16 || dtype == INFINI_DTYPE_BF16) { - if (wtype != dtype && wtype != INFINI_DTYPE_F32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (dtype == INFINI_DTYPE_F32) { - if (wtype != INFINI_DTYPE_F32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // Validate matrix weight dtype: for half-precision activations, can be same or FP32 - if (dtype == INFINI_DTYPE_F16 || dtype == INFINI_DTYPE_BF16) { - if (mtype != dtype && mtype != INFINI_DTYPE_F32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (dtype == INFINI_DTYPE_F32) { - if (mtype != INFINI_DTYPE_F32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // Check tensor dimensions - const size_t out_ndim = out_desc->ndim(); - const size_t in_ndim = in_desc->ndim(); - - // Must be 2D tensors [ntok, hidden_dim] - if (out_ndim != 2 || in_ndim != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t ntok = out_desc->dim(0); - size_t hidden_dim = out_desc->dim(1); - - // Check input shape matches output - if (in_desc->dim(0) != ntok || in_desc->dim(1) != hidden_dim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // Check norm_weight is 1D [hidden_dim] - if (norm_weight_desc->ndim() != 1 || norm_weight_desc->dim(0) != hidden_dim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // Check gate_up_weight is 2D with shape either [2*intermediate_dim, hidden_dim] or [hidden_dim, 2*intermediate_dim] - if (gate_up_weight_desc->ndim() != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - size_t gate_up_dim0 = gate_up_weight_desc->dim(0); - size_t gate_up_dim1 = gate_up_weight_desc->dim(1); - size_t intermediate_dim; - if (gate_up_dim1 == hidden_dim && gate_up_dim0 % 2 == 0) { - // Layout A: [2*intermediate_dim, hidden_dim] - intermediate_dim = gate_up_dim0 / 2; - } else if (gate_up_dim0 == hidden_dim && gate_up_dim1 % 2 == 0) { - // Layout B: [hidden_dim, 2*intermediate_dim] - intermediate_dim = gate_up_dim1 / 2; - } else { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // Check down_weight is 2D with shape either [hidden_dim, intermediate_dim] or [intermediate_dim, hidden_dim] - if (down_weight_desc->ndim() != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - { - size_t dw_dim0 = down_weight_desc->dim(0); - size_t dw_dim1 = down_weight_desc->dim(1); - if (!((dw_dim0 == hidden_dim && dw_dim1 == intermediate_dim) || (dw_dim0 == intermediate_dim && dw_dim1 == hidden_dim))) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - // Check contiguity of the last dimension for activation tensors and norm weights - if (out_desc->stride(out_ndim - 1) != 1 || in_desc->stride(in_ndim - 1) != 1 || norm_weight_desc->stride(0) != 1) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - // For matrix weights, at least one stride dimension must be 1 (contiguous along one axis) - if (gate_up_weight_desc->stride(0) != 1 && gate_up_weight_desc->stride(1) != 1) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - if (down_weight_desc->stride(0) != 1 && down_weight_desc->stride(1) != 1) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - // Check residual if provided - bool has_residual = residual_desc != nullptr; - ptrdiff_t residual_stride = 0; - if (has_residual) { - if (residual_desc->ndim() != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (residual_desc->dim(0) != ntok || residual_desc->dim(1) != hidden_dim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (residual_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - if (residual_desc->stride(1) != 1) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - residual_stride = residual_desc->stride(0); - } - - FusedFFNInfo info; - info.dtype = dtype; - info.wtype = wtype; - info.mtype = mtype; - info.epsilon = epsilon; - info.shape = out_desc->shape(); - info.in_stride = in_desc->stride(0); - info.out_stride = out_desc->stride(0); - info.residual_stride = residual_stride; - info.hidden_dim = hidden_dim; - info.intermediate_dim = intermediate_dim; - info.has_residual = has_residual; - - return utils::Result(info); - } -}; - -} // namespace op::fused_ffn - -#endif // __FUSED_FFN_INFO_H__ diff --git a/src/infiniop/ops/fused_ffn/metax/fused_ffn_metax.cuh b/src/infiniop/ops/fused_ffn/metax/fused_ffn_metax.cuh deleted file mode 100644 index 5af0cae64..000000000 --- a/src/infiniop/ops/fused_ffn/metax/fused_ffn_metax.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FUSED_FFN_METAX_H__ -#define __FUSED_FFN_METAX_H__ - -#include "../fused_ffn.h" - -DESCRIPTOR(metax) - -#endif diff --git a/src/infiniop/ops/fused_ffn/metax/fused_ffn_metax.maca b/src/infiniop/ops/fused_ffn/metax/fused_ffn_metax.maca deleted file mode 100644 index d74c51b76..000000000 --- a/src/infiniop/ops/fused_ffn/metax/fused_ffn_metax.maca +++ /dev/null @@ -1,349 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "fused_ffn_metax.cuh" - -#include "../../../devices/metax/metax_kernel_common.h" -#include - -// Define CUDA kernel macro for compatibility with shared kernel code -#define INFINIOP_CUDA_KERNEL INFINIOP_METAX_KERNEL - -#include "../nvidia/kernel.cuh" - -namespace op::fused_ffn::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - infiniopTensorDescriptor_t residual_desc, - infiniopTensorDescriptor_t norm_weight_desc, - infiniopTensorDescriptor_t gate_up_weight_desc, - infiniopTensorDescriptor_t down_weight_desc, - float epsilon) { - - auto result = FusedFFNInfo::create( - out_desc, in_desc, residual_desc, - norm_weight_desc, gate_up_weight_desc, down_weight_desc, epsilon); - CHECK_RESULT(result); - auto info = result.take(); - - size_t dtype_size = infiniSizeOf(info.dtype); - size_t ntok = info.ntok(); - size_t d = info.d(); - size_t di = info.di(); - - size_t normalized_size = ntok * d * dtype_size; - size_t gate_up_size = ntok * 2 * di * dtype_size; - size_t workspace_size = normalized_size + gate_up_size; - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - std::move(info), - workspace_size, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -// Launch RMSNorm kernel with different data types -template -infiniStatus_t launchRmsnormKernel( - size_t ntok, size_t dim, - void *normalized, const void *x, const void *w, - infiniDtype_t dtype, infiniDtype_t wtype, - float epsilon, - ptrdiff_t in_stride, ptrdiff_t out_stride, - hcStream_t stream) { - -#define LAUNCH_KERNEL(Tdata, Tweight, Tcompute) \ - rmsnormKernel \ - <<>>( \ - reinterpret_cast(normalized), \ - reinterpret_cast(x), \ - reinterpret_cast(w), \ - ntok, dim, epsilon, in_stride, out_stride) - - if (dtype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(half, half, float); - } else if (dtype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(half, __hpcc_bfloat16, float); - } else if (dtype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(half, float, float); - } else if (dtype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(__hpcc_bfloat16, __hpcc_bfloat16, float); - } else if (dtype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(__hpcc_bfloat16, half, float); - } else if (dtype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(__hpcc_bfloat16, float, float); - } else if (dtype == INFINI_DTYPE_F32 && wtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(float, float, float); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_KERNEL - - return INFINI_STATUS_SUCCESS; -} - -// Launch SwiGLU kernel with different data types (in-place) -template -infiniStatus_t launchSwigluKernel( - size_t ntok, size_t intermediate_dim, - void *gate_up, - infiniDtype_t dtype, - ptrdiff_t stride, - hcStream_t stream) { - -#define LAUNCH_KERNEL(Tdata, Tcompute) \ - swigluKernel \ - <<>>( \ - reinterpret_cast(gate_up), \ - ntok, intermediate_dim, stride) - - if (dtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(half, float); - } else if (dtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(__hpcc_bfloat16, float); - } else if (dtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(float, float); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_KERNEL - - return INFINI_STATUS_SUCCESS; -} - -// Launch Residual Add kernel with different data types -template -infiniStatus_t launchResidualAddKernel( - size_t ntok, size_t dim, - void *out, const void *gemm_out, const void *residual, - infiniDtype_t dtype, - ptrdiff_t out_stride, ptrdiff_t residual_stride, - hcStream_t stream) { - -#define LAUNCH_KERNEL(Tdata, Tcompute) \ - residualAddKernel \ - <<>>( \ - reinterpret_cast(out), \ - reinterpret_cast(gemm_out), \ - reinterpret_cast(residual), \ - ntok, dim, out_stride, residual_stride) - - if (dtype == INFINI_DTYPE_F16) { - LAUNCH_KERNEL(half, float); - } else if (dtype == INFINI_DTYPE_BF16) { - LAUNCH_KERNEL(__hpcc_bfloat16, float); - } else if (dtype == INFINI_DTYPE_F32) { - LAUNCH_KERNEL(float, float); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_KERNEL - - return INFINI_STATUS_SUCCESS; -} - -// Helper function to get hcBLAS types for activation and weight matrices -static void getHcblasTypes(infiniDtype_t dtype, infiniDtype_t wtype, - hpccDataType &a_type, hpccDataType &b_type, hpccDataType &c_type, - hcblasComputeType_t &compute_type) { - switch (dtype) { - case INFINI_DTYPE_F16: - a_type = HPCC_R_16F; - c_type = HPCC_R_16F; - break; - case INFINI_DTYPE_BF16: - a_type = HPCC_R_16BF; - c_type = HPCC_R_16BF; - break; - case INFINI_DTYPE_F32: - a_type = HPCC_R_32F; - c_type = HPCC_R_32F; - break; - default: - break; - } - - switch (wtype) { - case INFINI_DTYPE_F16: - b_type = HPCC_R_16F; - break; - case INFINI_DTYPE_BF16: - b_type = HPCC_R_16BF; - break; - case INFINI_DTYPE_F32: - b_type = HPCC_R_32F; - break; - default: - break; - } - - switch (dtype) { - case INFINI_DTYPE_F16: - case INFINI_DTYPE_BF16: - compute_type = HCBLAS_COMPUTE_32F; - break; - case INFINI_DTYPE_F32: - compute_type = HCBLAS_COMPUTE_32F_FAST_TF32; - break; - default: - break; - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *out, - const void *in, - const void *residual, - const void *norm_weight, - const void *gate_up_weight, - const void *down_weight, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - auto hc_stream = reinterpret_cast(stream); - size_t ntok = _info.ntok(); - size_t d = _info.d(); - size_t di = _info.di(); - size_t dtype_size = infiniSizeOf(_info.dtype); - - // Partition workspace - char *ws_ptr = static_cast(workspace); - void *normalized_buf = ws_ptr; - ws_ptr += ntok * d * dtype_size; - void *gate_up_buf = ws_ptr; - - // Stage 1: RMSNorm - { - unsigned int max_threads = _opaque->internal->maxThreadsPerBlock(); - if (max_threads >= METAX_BLOCK_SIZE_1024) { - CHECK_STATUS(launchRmsnormKernel( - ntok, d, normalized_buf, in, norm_weight, - _info.dtype, _info.wtype, _info.epsilon, - _info.in_stride, d, hc_stream)); - } else { - CHECK_STATUS(launchRmsnormKernel( - ntok, d, normalized_buf, in, norm_weight, - _info.dtype, _info.wtype, _info.epsilon, - _info.in_stride, d, hc_stream)); - } - } - - // Stage 2: GateUp GEMM using hcBLAS - { - hpccDataType a_type, b_type, c_type; - hcblasComputeType_t compute_type; - getHcblasTypes(_info.dtype, _info.mtype, a_type, b_type, c_type, compute_type); - - float alpha = 1.0f, beta = 0.0f; - CHECK_STATUS(_opaque->internal->useMcblas( - hc_stream, - [&](hcblasHandle_t handle) { - CHECK_MCBLAS( - hcblasGemmStridedBatchedEx( - handle, - HCBLAS_OP_T, - HCBLAS_OP_N, - static_cast(2 * di), - static_cast(ntok), - static_cast(d), - &alpha, - gate_up_weight, b_type, static_cast(d), 0, - normalized_buf, a_type, static_cast(d), 0, - &beta, - gate_up_buf, c_type, static_cast(2 * di), 0, - 1, - compute_type, - HCBLAS_GEMM_DEFAULT_TENSOR_OP)); - return INFINI_STATUS_SUCCESS; - })); - } - - // Stage 3: SwiGLU transform (in-place) - if (_opaque->internal->maxThreadsPerBlock() >= METAX_BLOCK_SIZE_1024) { - CHECK_STATUS(launchSwigluKernel<1024>( - ntok, di, - gate_up_buf, - _info.dtype, - 2 * di, - hc_stream)); - } else { - CHECK_STATUS(launchSwigluKernel<512>( - ntok, di, - gate_up_buf, - _info.dtype, - 2 * di, - hc_stream)); - } - - // Stage 4: Down GEMM using hcBLAS - { - hpccDataType a_type, b_type, c_type; - hcblasComputeType_t compute_type; - getHcblasTypes(_info.dtype, _info.mtype, a_type, b_type, c_type, compute_type); - - bool fuse_residual = _info.has_residual && (out == residual); - float alpha = 1.0f, beta = fuse_residual ? 1.0f : 0.0f; - CHECK_STATUS(_opaque->internal->useMcblas( - hc_stream, - [&](hcblasHandle_t handle) { - CHECK_MCBLAS( - hcblasGemmStridedBatchedEx( - handle, - HCBLAS_OP_T, - HCBLAS_OP_N, - static_cast(d), - static_cast(ntok), - static_cast(di), - &alpha, - down_weight, b_type, static_cast(di), 0, - gate_up_buf, a_type, static_cast(2 * di), 0, - &beta, - out, c_type, static_cast(_info.out_stride), 0, - 1, - compute_type, - HCBLAS_GEMM_DEFAULT_TENSOR_OP)); - return INFINI_STATUS_SUCCESS; - })); - } - - // Stage 5: Residual Add (only when not fused into GEMM) - if (_info.has_residual && out != residual) { - if (_opaque->internal->maxThreadsPerBlock() >= METAX_BLOCK_SIZE_1024) { - CHECK_STATUS(launchResidualAddKernel<1024>( - ntok, d, - out, out, residual, - _info.dtype, - _info.out_stride, _info.residual_stride, - hc_stream)); - } else { - CHECK_STATUS(launchResidualAddKernel<512>( - ntok, d, - out, out, residual, - _info.dtype, - _info.out_stride, _info.residual_stride, - hc_stream)); - } - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::fused_ffn::metax diff --git a/src/infiniop/ops/fused_ffn/nvidia/fused_ffn_nvidia.cu b/src/infiniop/ops/fused_ffn/nvidia/fused_ffn_nvidia.cu deleted file mode 100644 index 18336b224..000000000 --- a/src/infiniop/ops/fused_ffn/nvidia/fused_ffn_nvidia.cu +++ /dev/null @@ -1,432 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "fused_ffn_nvidia.cuh" -#include "kernel.cuh" - -// Each op's public header re-defines a DESCRIPTOR(NAMESPACE) macro without -// guarding it, so including multiple sub-op headers in this TU clashes. -// We #undef between includes to sidestep the collision; order is arbitrary. -#undef DESCRIPTOR -#include "../../gemm/nvidia/gemm_nvidia.cuh" -#undef DESCRIPTOR -#include "../../rms_norm/nvidia/rms_norm_nvidia.cuh" - -// swiglu / add use a separate ELEMENTWISE_DESCRIPTOR macro so they do not -// clash with the DESCRIPTOR macro above. -#include "../../add/nvidia/add_nvidia.cuh" -#include "../../swiglu/nvidia/swiglu_nvidia.cuh" - -#include -#include -#include -#include - -namespace op::fused_ffn::nvidia { - -namespace { - -// 256-byte alignment is safe for both cublas working buffers and elementwise -// meta blobs; higher would waste workspace on small shapes. -constexpr size_t kWsAlign = 256; - -inline size_t alignUp(size_t x, size_t a) { - return (x + a - 1) & ~(a - 1); -} - -// Heap-allocate a 2-D tensor descriptor with explicit element strides. -// Returned pointer must be deleted by the caller. -inline infiniopTensorDescriptor_t make2D(infiniDtype_t dtype, - size_t d0, size_t d1, - ptrdiff_t s0, ptrdiff_t s1) { - const size_t shape[2] = {d0, d1}; - const ptrdiff_t strides[2] = {s0, s1}; - return new InfiniopTensorDescriptor(dtype, 2, shape, strides); -} - -// Synthesize a GEMM B-matrix view with logical shape [k, n] regardless of -// whether the original weight was stored as [n, k] (Layout A) or [k, n] -// (Layout B). FusedFFNInfo::create already guarantees one of the two strides -// is 1, so this only needs to decide which dim is which and swap accordingly. -inline infiniopTensorDescriptor_t makeWeightAsKN(infiniDtype_t dtype, - size_t k, size_t n, - infiniopTensorDescriptor_t orig) { - const size_t d0 = orig->dim(0); - const ptrdiff_t s0 = orig->stride(0); - const ptrdiff_t s1 = orig->stride(1); - if (d0 == n) { - // original [n, k] -> view as [k, n] by swapping axes - return make2D(dtype, k, n, s1, s0); - } - // original already [k, n] - return make2D(dtype, k, n, s0, s1); -} - -// RAII wrapper: owns a list of synthesized tensor descriptors and deletes -// them on scope exit. Sub-descriptor create() calls copy out what they need, -// so the temporaries only need to outlive the create() call. -class DescScope { - std::vector _owned; - -public: - ~DescScope() { - for (auto *t : _owned) { - delete t; - } - } - infiniopTensorDescriptor_t adopt(infiniopTensorDescriptor_t t) { - _owned.push_back(t); - return t; - } -}; - -} // namespace - -struct Descriptor::Opaque { - std::shared_ptr internal; - - // Workspace slab sizes (bytes), padded to kWsAlign. - size_t normalized_bytes = 0; - size_t gate_up_bytes = 0; - size_t hidden_bytes = 0; - size_t inner_ws_bytes = 0; // max of sub-descriptor workspaceSize() - - bool has_residual = false; - - // Deep-fused Gate-Up + SwiGLU (paper-style fusion) config. - // Decided at create-time via env var INFINIOP_FUSED_FFN_DEEP and - // a profile-driven scheduler that checks ntok against a threshold: - // INFINIOP_FUSED_FFN_DEEP=0 or unset -> always 5-stage (default) - // INFINIOP_FUSED_FFN_DEEP=1 -> scheduler: deep-fused when - // ntok <= threshold, else 5-stage. Threshold defaults to 4 and - // can be overridden with INFINIOP_FUSED_FFN_DEEP_MAX_NTOK. - // INFINIOP_FUSED_FFN_DEEP=2 -> force deep-fused always - // (debug/profiling only; will regress at large ntok) - bool use_deep_fused = false; - ptrdiff_t gate_up_w_k_stride = 0; - ptrdiff_t gate_up_w_col_stride = 0; - - // Sub-descriptors owned by this fused op; each one is a standard - // InfiniopDescriptor for the corresponding standalone operator. - std::unique_ptr rms_norm; - std::unique_ptr gate_up_gemm; - std::unique_ptr swiglu; - std::unique_ptr down_gemm; - std::unique_ptr residual_add; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - infiniopTensorDescriptor_t residual_desc, - infiniopTensorDescriptor_t norm_weight_desc, - infiniopTensorDescriptor_t gate_up_weight_desc, - infiniopTensorDescriptor_t down_weight_desc, - float epsilon) { - - auto info_result = FusedFFNInfo::create( - out_desc, in_desc, residual_desc, - norm_weight_desc, gate_up_weight_desc, down_weight_desc, epsilon); - CHECK_RESULT(info_result); - auto info = info_result.take(); - - auto handle = reinterpret_cast(handle_); - - auto opaque = std::make_unique(); - opaque->internal = handle->internal(); - opaque->has_residual = info.has_residual; - - const size_t ntok = info.ntok(); - const size_t d = info.d(); - const size_t di = info.di(); - const size_t dtype_sz = infiniSizeOf(info.dtype); - - // Profile-driven scheduler for the deep-fused kernel path. - // INFINIOP_FUSED_FFN_DEEP=0/unset -> always 5-stage (default) - // INFINIOP_FUSED_FFN_DEEP=1 -> scheduler: use deep-fused only - // when ntok <= max_ntok threshold (default 4); fall back to - // 5-stage for larger shapes where cuBLAS tensor-core GEMM wins. - // INFINIOP_FUSED_FFN_DEEP=2 -> force deep-fused always (debug) - // Threshold is tunable via INFINIOP_FUSED_FFN_DEEP_MAX_NTOK (default 4). - { - const char *env = std::getenv("INFINIOP_FUSED_FFN_DEEP"); - if (env != nullptr && env[0] == '2') { - // Mode 2: force deep-fused regardless of shape - opaque->use_deep_fused = true; - } else if (env != nullptr && env[0] == '1') { - // Mode 1: scheduler — deep-fused only for small ntok - size_t max_ntok = 4; - const char *thr = std::getenv("INFINIOP_FUSED_FFN_DEEP_MAX_NTOK"); - if (thr != nullptr) { - max_ntok = static_cast(std::atol(thr)); - if (max_ntok == 0) { - max_ntok = 4; - } - } - opaque->use_deep_fused = (ntok <= max_ntok); - } - // Mode 0 / unset: use_deep_fused stays false (5-stage) - } - - // Extract gate_up weight strides at create time so the deep-fused kernel - // can index [k, j] regardless of whether storage is Layout A [2*di, d] - // or Layout B [d, 2*di]. Layout is identified by which descriptor dim - // equals 2*di (the output column axis) vs d (the K axis). - { - const size_t gu_dim0 = gate_up_weight_desc->dim(0); - const ptrdiff_t gu_s0 = gate_up_weight_desc->stride(0); - const ptrdiff_t gu_s1 = gate_up_weight_desc->stride(1); - if (gu_dim0 == 2 * di) { - // Layout A: [2*di, d] — output column is dim0, K is dim1. - opaque->gate_up_w_col_stride = gu_s0; - opaque->gate_up_w_k_stride = gu_s1; - } else { - // Layout B: [d, 2*di] — K is dim0, output column is dim1. - opaque->gate_up_w_k_stride = gu_s0; - opaque->gate_up_w_col_stride = gu_s1; - } - } - - // ── Workspace layout ── - // normalized : [ntok, d] contiguous -> RMSNorm out, GateUp in - // gate_up : [ntok, 2*di] contiguous -> GateUp out, SwiGLU in - // hidden : [ntok, di] contiguous -> SwiGLU out, Down in - // inner_ws : max(sub->workspaceSize()) shared by sub-descriptors - // - // The compact hidden slab (stride=di instead of stride=2*di) gives the - // Down-GEMM a tightly packed K dimension, which matters on BIV150 where - // cuBLAS 10.2 tensor-core paths prefer aligned contiguous leading dims. - opaque->normalized_bytes = alignUp(ntok * d * dtype_sz, kWsAlign); - opaque->gate_up_bytes = alignUp(ntok * 2 * di * dtype_sz, kWsAlign); - opaque->hidden_bytes = alignUp(ntok * di * dtype_sz, kWsAlign); - - DescScope scope; - - // ── RMSNorm sub-descriptor ── - // Activation is 2-D [ntok, d]; weight is 1-D [d]. - auto normalized_desc = scope.adopt( - make2D(info.dtype, ntok, d, static_cast(d), 1)); - auto in_view = scope.adopt( - make2D(info.dtype, ntok, d, info.in_stride, 1)); - - { - op::rms_norm::nvidia::Descriptor *sub = nullptr; - CHECK_STATUS(op::rms_norm::nvidia::Descriptor::create( - handle_, &sub, - normalized_desc, in_view, norm_weight_desc, - info.epsilon)); - opaque->rms_norm.reset(sub); - opaque->inner_ws_bytes = std::max(opaque->inner_ws_bytes, sub->workspaceSize()); - } - - // ── GateUp GEMM sub-descriptor ── - // [ntok, 2*di] = [ntok, d] @ [d, 2*di] - auto gate_up_c_desc = scope.adopt( - make2D(info.dtype, ntok, 2 * di, static_cast(2 * di), 1)); - auto gate_up_b_desc = scope.adopt( - makeWeightAsKN(info.mtype, d, 2 * di, gate_up_weight_desc)); - - { - op::gemm::nvidia::Descriptor *sub = nullptr; - CHECK_STATUS(op::gemm::nvidia::Descriptor::create( - handle_, &sub, gate_up_c_desc, normalized_desc, gate_up_b_desc)); - opaque->gate_up_gemm.reset(sub); - opaque->inner_ws_bytes = std::max(opaque->inner_ws_bytes, sub->workspaceSize()); - } - - // ── SwiGLU sub-descriptor ── - // Operates on the interleaved [gate | up] buffer: - // logical inputs : up [ntok, di] row stride 2*di - // gate [ntok, di] row stride 2*di - // logical output : hidden [ntok, di] contiguous - // gate and up share identical shape/strides — only their base pointers - // differ at calculate time. - auto hidden_desc = scope.adopt( - make2D(info.dtype, ntok, di, static_cast(di), 1)); - auto half_desc = scope.adopt( - make2D(info.dtype, ntok, di, static_cast(2 * di), 1)); - - { - op::swiglu::nvidia::Descriptor *sub = nullptr; - CHECK_STATUS(op::swiglu::nvidia::Descriptor::create( - handle_, &sub, hidden_desc, {half_desc, half_desc})); - opaque->swiglu.reset(sub); - opaque->inner_ws_bytes = std::max(opaque->inner_ws_bytes, sub->workspaceSize()); - } - - // ── Down GEMM sub-descriptor ── - // out = [beta * out] + 1.0 * hidden @ down_weight - // The output matrix uses the user's out stride so the gemm writes - // directly into the caller's tensor. - auto out_view = scope.adopt( - make2D(info.dtype, ntok, d, info.out_stride, 1)); - auto down_b_desc = scope.adopt( - makeWeightAsKN(info.mtype, di, d, down_weight_desc)); - - { - op::gemm::nvidia::Descriptor *sub = nullptr; - CHECK_STATUS(op::gemm::nvidia::Descriptor::create( - handle_, &sub, out_view, hidden_desc, down_b_desc)); - opaque->down_gemm.reset(sub); - opaque->inner_ws_bytes = std::max(opaque->inner_ws_bytes, sub->workspaceSize()); - } - - // ── Residual add sub-descriptor (optional) ── - // Only used when residual is a distinct tensor from out; the - // (out == residual) case is fused into the Down-GEMM via beta=1 at - // calculate time. - if (info.has_residual) { - auto residual_view = scope.adopt( - make2D(info.dtype, ntok, d, info.residual_stride, 1)); - auto out_view_for_add = scope.adopt( - make2D(info.dtype, ntok, d, info.out_stride, 1)); - - op::add::nvidia::Descriptor *sub = nullptr; - CHECK_STATUS(op::add::nvidia::Descriptor::create( - handle_, &sub, - out_view_for_add, {out_view_for_add, residual_view})); - opaque->residual_add.reset(sub); - opaque->inner_ws_bytes = std::max(opaque->inner_ws_bytes, sub->workspaceSize()); - } - - const size_t workspace_size = opaque->normalized_bytes + opaque->gate_up_bytes + opaque->hidden_bytes + alignUp(opaque->inner_ws_bytes, kWsAlign); - - *desc_ptr = new Descriptor( - opaque.release(), - std::move(info), - workspace_size, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *out, - const void *in, - const void *residual, - const void *norm_weight, - const void *gate_up_weight, - const void *down_weight, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - const size_t di = _info.di(); - const size_t dtype_sz = infiniSizeOf(_info.dtype); - - // Partition the workspace into the three persistent slabs plus an - // inner scratch buffer shared by all sub-descriptors. - char *ws = static_cast(workspace); - void *normalized_buf = ws; - ws += _opaque->normalized_bytes; - void *gate_up_buf = ws; - ws += _opaque->gate_up_bytes; - void *hidden_buf = ws; - ws += _opaque->hidden_bytes; - void *inner_ws = ws; - const size_t inner_ws_size = _opaque->inner_ws_bytes; - - // gate and up are two halves of the interleaved gate_up buffer. - // Each token's layout is [gate[0..di) | up[0..di)], so gate starts at - // offset 0 and up starts at offset (di * dtype_sz) bytes within the - // first row; both use a row stride of 2*di elements (captured in the - // shared half_desc at create time). - const char *gu_bytes = static_cast(gate_up_buf); - const void *gate_ptr = gu_bytes; - const void *up_ptr = gu_bytes + di * dtype_sz; - - // Stage 1: RMSNorm - CHECK_STATUS(_opaque->rms_norm->calculate( - inner_ws, inner_ws_size, - normalized_buf, in, norm_weight, stream)); - - if (_opaque->use_deep_fused) { - // Stage 2+3 fused: one kernel produces hidden = SiLU(norm@Wg) * (norm@Wu) - // directly, eliminating the gate_up_buf HBM round-trip. Row strides - // for X and hidden are d and di respectively (contiguous buffers - // allocated at create time). - cudaStream_t cuda_stream = reinterpret_cast(stream); - const size_t ntok = _info.ntok(); - const size_t d = _info.d(); - const size_t di = _info.di(); - - constexpr unsigned int kBlock = 256; - dim3 grid(static_cast(ntok), static_cast(di)); - dim3 block(kBlock); - -#define DEEP_FUSED_LAUNCH(TD, TW) \ - deepFusedGateUpSiluKernel \ - <<>>( \ - reinterpret_cast(hidden_buf), \ - reinterpret_cast(normalized_buf), \ - reinterpret_cast(gate_up_weight), \ - ntok, d, di, \ - static_cast(d), \ - static_cast(di), \ - _opaque->gate_up_w_k_stride, \ - _opaque->gate_up_w_col_stride, \ - /*gate_col_base=*/0u, /*up_col_base=*/di) - - if (_info.dtype == INFINI_DTYPE_F16 && _info.mtype == INFINI_DTYPE_F16) { - DEEP_FUSED_LAUNCH(half, half); - } else if (_info.dtype == INFINI_DTYPE_BF16 && _info.mtype == INFINI_DTYPE_BF16) { - DEEP_FUSED_LAUNCH(__nv_bfloat16, __nv_bfloat16); - } else if (_info.dtype == INFINI_DTYPE_F32 && _info.mtype == INFINI_DTYPE_F32) { - DEEP_FUSED_LAUNCH(float, float); - } else if (_info.dtype == INFINI_DTYPE_F16 && _info.mtype == INFINI_DTYPE_F32) { - DEEP_FUSED_LAUNCH(half, float); - } else if (_info.dtype == INFINI_DTYPE_BF16 && _info.mtype == INFINI_DTYPE_F32) { - DEEP_FUSED_LAUNCH(__nv_bfloat16, float); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -#undef DEEP_FUSED_LAUNCH - // Avoid unused-variable warnings for the baseline-only pointers when - // the deep-fused branch takes over stage 2/3. - (void)gate_up_buf; - (void)gate_ptr; - (void)up_ptr; - } else { - // Stage 2: GateUp GEMM --> gate_up_buf = normalized_buf @ gate_up_weight - CHECK_STATUS(_opaque->gate_up_gemm->calculate( - inner_ws, inner_ws_size, - gate_up_buf, /*beta=*/0.f, - normalized_buf, gate_up_weight, - /*alpha=*/1.f, stream)); - - // Stage 3: SwiGLU --> hidden_buf = silu(gate) * up - // swiglu::nvidia expects inputs ordered {up, gate}; see - // swiglu_nvidia.cu input_desc_vec[0]=up, [1]=gate. - CHECK_STATUS(_opaque->swiglu->calculate( - inner_ws, inner_ws_size, - hidden_buf, {up_ptr, gate_ptr}, stream)); - } - - // Stage 4: Down GEMM, with optional in-place residual fuse via beta=1. - // fuse path : out = 1.0 * out + hidden_buf @ down_weight - // plain path: out = 0.0 * out + hidden_buf @ down_weight - const bool fuse_residual = _opaque->has_residual && (out == residual); - CHECK_STATUS(_opaque->down_gemm->calculate( - inner_ws, inner_ws_size, - out, /*beta=*/fuse_residual ? 1.f : 0.f, - hidden_buf, down_weight, - /*alpha=*/1.f, stream)); - - // Stage 5: Residual add (only when the in-place fuse did not apply). - if (_opaque->has_residual && !fuse_residual) { - CHECK_STATUS(_opaque->residual_add->calculate( - inner_ws, inner_ws_size, - out, out, residual, stream)); - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::fused_ffn::nvidia diff --git a/src/infiniop/ops/fused_ffn/nvidia/fused_ffn_nvidia.cuh b/src/infiniop/ops/fused_ffn/nvidia/fused_ffn_nvidia.cuh deleted file mode 100644 index 9ef1cd38a..000000000 --- a/src/infiniop/ops/fused_ffn/nvidia/fused_ffn_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FUSED_FFN_NVIDIA_CUDA_H__ -#define __FUSED_FFN_NVIDIA_CUDA_H__ - -#include "../fused_ffn.h" - -DESCRIPTOR(nvidia) - -#endif diff --git a/src/infiniop/ops/fused_ffn/nvidia/kernel.cuh b/src/infiniop/ops/fused_ffn/nvidia/kernel.cuh deleted file mode 100644 index 243b3b890..000000000 --- a/src/infiniop/ops/fused_ffn/nvidia/kernel.cuh +++ /dev/null @@ -1,226 +0,0 @@ -#ifndef __FUSED_FFN_CUDA_KERNEL_H__ -#define __FUSED_FFN_CUDA_KERNEL_H__ - -#include - -// RMSNorm preprocessing kernel -// Each block processes one token -// Computes: normalized = x * rsqrt(mean(x^2) + eps) * w -template -__device__ void rmsnormBlock( - Tdata *__restrict__ normalized, - const Tdata *__restrict__ x, - const Tweight *__restrict__ w, - size_t dim, - float epsilon, - ptrdiff_t in_stride, - ptrdiff_t out_stride) { - - size_t token_idx = blockIdx.x; - auto x_ptr = x + token_idx * in_stride; - auto norm_ptr = normalized + token_idx * out_stride; - - // Compute sum of squares - Tcompute sum_squared = 0; - for (size_t i = threadIdx.x; i < dim; i += BLOCK_SIZE) { - Tcompute val = Tcompute(x_ptr[i]); - sum_squared += val * val; - } - - // Block-reduce sum of squares - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - sum_squared = BlockReduce(temp_storage).Sum(sum_squared); - - // Thread 0 computes RMS = 1/sqrt(ss/dim + epsilon) - __shared__ Tcompute rms; - if (threadIdx.x == 0) { - rms = Tcompute(rsqrtf(sum_squared / Tcompute(dim) + epsilon)); - } - __syncthreads(); - - // Apply normalization: normalized = x * w * rms - for (size_t i = threadIdx.x; i < dim; i += BLOCK_SIZE) { - norm_ptr[i] = Tdata(Tcompute(x_ptr[i]) * Tcompute(w[i]) * rms); - } -} - -template -INFINIOP_CUDA_KERNEL rmsnormKernel( - Tdata *__restrict__ normalized, - const Tdata *__restrict__ x, - const Tweight *__restrict__ w, - size_t ntok, - size_t dim, - float epsilon, - ptrdiff_t in_stride, - ptrdiff_t out_stride) { - if (blockIdx.x < ntok) { - rmsnormBlock( - normalized, x, w, dim, epsilon, in_stride, out_stride); - } -} - -// SwiGLU transform kernel (in-place) -// Computes: gate_up[i] = silu(gate_up[i]) * gate_up[di+i] for i in [0, di) -// Writes result to the gate half of gate_up, overwriting gate values. -// This matches the non-fused path's buffer layout (hidden at stride 2*di). -template -__device__ void swigluBlock( - Tdata *__restrict__ gate_up, - size_t intermediate_dim, - ptrdiff_t stride) { - - size_t token_idx = blockIdx.x; - auto gate_up_ptr = gate_up + token_idx * stride; - - for (size_t i = threadIdx.x; i < intermediate_dim; i += BLOCK_SIZE) { - Tcompute gate = Tcompute(gate_up_ptr[i]); - Tcompute up = Tcompute(gate_up_ptr[intermediate_dim + i]); - - // SiLU(x) = x * sigmoid(x) = x / (1 + exp(-x)) - Tcompute sigmoid = Tcompute(1.0f) / (Tcompute(1.0f) + exp_(-gate)); - Tcompute silu = gate * sigmoid; - - gate_up_ptr[i] = Tdata(silu * up); - } -} - -template -INFINIOP_CUDA_KERNEL swigluKernel( - Tdata *__restrict__ gate_up, - size_t ntok, - size_t intermediate_dim, - ptrdiff_t stride) { - if (blockIdx.x < ntok) { - swigluBlock( - gate_up, intermediate_dim, stride); - } -} - -// Residual add kernel -// Computes: out = gemm_out + residual -template -__device__ void residualAddBlock( - Tdata *__restrict__ out, - const Tdata *__restrict__ gemm_out, - const Tdata *__restrict__ residual, - size_t dim, - ptrdiff_t out_stride, - ptrdiff_t residual_stride) { - - size_t token_idx = blockIdx.x; - auto out_ptr = out + token_idx * out_stride; - auto gemm_ptr = gemm_out + token_idx * out_stride; - auto residual_ptr = residual + token_idx * residual_stride; - - for (size_t i = threadIdx.x; i < dim; i += BLOCK_SIZE) { - out_ptr[i] = Tdata(Tcompute(gemm_ptr[i]) + Tcompute(residual_ptr[i])); - } -} - -template -INFINIOP_CUDA_KERNEL residualAddKernel( - Tdata *__restrict__ out, - const Tdata *__restrict__ gemm_out, - const Tdata *__restrict__ residual, - size_t ntok, - size_t dim, - ptrdiff_t out_stride, - ptrdiff_t residual_stride) { - if (blockIdx.x < ntok) { - residualAddBlock( - out, gemm_out, residual, dim, out_stride, residual_stride); - } -} - -// ── Deep-fused Gate-Up GEMM + SwiGLU ──────────────────────────────── -// -// Reproduces the paper-style fusion of DeepFusionKernel (Zhang et al., 2026) -// by streaming both matmul tiles and the SiLU-gated mul through registers, -// eliminating the gate_up_buf HBM round-trip that sits between the separate -// gate_up_gemm and swiglu kernels in the non-fused pipeline. -// -// For each output element hidden[token, col]: -// gate = Σ_k X[token, k] * W_gate[k, col] -// up = Σ_k X[token, k] * W_up [k, col] -// hidden[token, col] = (gate * sigmoid(gate)) * up -// -// Grid : (ntok, di) -// Block: BLOCK_SIZE threads cooperatively reducing over K=d. -// -// The weight pointer is the raw user pointer. w_k_stride and w_col_stride are -// the physical element strides when moving along K and along the output -// column axis respectively, so both [2*di, d] (layout A) and [d, 2*di] -// (layout B) storages are supported by choosing (stride_k, stride_col). -template -__device__ void deepFusedGateUpSiluBlock( - Tdata *__restrict__ hidden, - const Tdata *__restrict__ x, - const Tweight *__restrict__ w, - size_t d, - ptrdiff_t x_row_stride, - ptrdiff_t hidden_row_stride, - ptrdiff_t w_k_stride, - ptrdiff_t w_col_stride, - size_t gate_col_base, - size_t up_col_base) { - - const size_t token = blockIdx.x; - const size_t col = blockIdx.y; - - auto x_ptr = x + token * x_row_stride; - auto h_ptr = hidden + token * hidden_row_stride; - - const Tweight *w_gate_col = w + (gate_col_base + col) * w_col_stride; - const Tweight *w_up_col = w + (up_col_base + col) * w_col_stride; - - Tcompute gate_acc = Tcompute(0); - Tcompute up_acc = Tcompute(0); - - for (size_t k = threadIdx.x; k < d; k += BLOCK_SIZE) { - Tcompute x_val = Tcompute(x_ptr[k]); - Tcompute w_g = Tcompute(w_gate_col[k * w_k_stride]); - Tcompute w_u = Tcompute(w_up_col[k * w_k_stride]); - gate_acc += x_val * w_g; - up_acc += x_val * w_u; - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - - Tcompute gate_sum = BlockReduce(temp_storage).Sum(gate_acc); - __syncthreads(); - Tcompute up_sum = BlockReduce(temp_storage).Sum(up_acc); - - if (threadIdx.x == 0) { - Tcompute sig = Tcompute(1.0f) / (Tcompute(1.0f) + exp_(-gate_sum)); - Tcompute silu = gate_sum * sig; - h_ptr[col] = Tdata(silu * up_sum); - } -} - -template -INFINIOP_CUDA_KERNEL deepFusedGateUpSiluKernel( - Tdata *__restrict__ hidden, - const Tdata *__restrict__ x, - const Tweight *__restrict__ w, - size_t ntok, - size_t d, - size_t di, - ptrdiff_t x_row_stride, - ptrdiff_t hidden_row_stride, - ptrdiff_t w_k_stride, - ptrdiff_t w_col_stride, - size_t gate_col_base, - size_t up_col_base) { - if (blockIdx.x < ntok && blockIdx.y < di) { - deepFusedGateUpSiluBlock( - hidden, x, w, d, - x_row_stride, hidden_row_stride, - w_k_stride, w_col_stride, - gate_col_base, up_col_base); - } -} - -#endif // __FUSED_FFN_CUDA_KERNEL_H__ diff --git a/src/infiniop/ops/fused_ffn/operator.cc b/src/infiniop/ops/fused_ffn/operator.cc deleted file mode 100644 index 616a1957f..000000000 --- a/src/infiniop/ops/fused_ffn/operator.cc +++ /dev/null @@ -1,190 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/fused_ffn.h" - -#ifdef ENABLE_CPU_API -#include "cpu/fused_ffn_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) -#include "nvidia/fused_ffn_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/fused_ffn_metax.cuh" -#endif - -__INFINI_C infiniStatus_t infiniopCreateFusedFFNDescriptor( - infiniopHandle_t handle, - infiniopFusedFFNDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - infiniopTensorDescriptor_t residual_desc, - infiniopTensorDescriptor_t norm_weight_desc, - infiniopTensorDescriptor_t gate_up_weight_desc, - infiniopTensorDescriptor_t down_weight_desc, - float epsilon) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::fused_ffn::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - out_desc, \ - in_desc, \ - residual_desc, \ - norm_weight_desc, \ - gate_up_weight_desc, \ - down_weight_desc, \ - epsilon) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetFusedFFNWorkspaceSize( - infiniopFusedFFNDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - GET(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopFusedFFN( - infiniopFusedFFNDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *in, - const void *residual, - const void *norm_weight, - const void *gate_up_weight, - const void *down_weight, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, out, in, residual, \ - norm_weight, gate_up_weight, down_weight, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyFusedFFNDescriptor( - infiniopFusedFFNDescriptor_t desc) { - if (desc == nullptr) { - return INFINI_STATUS_SUCCESS; - } - -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DESTROY(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - DESTROY(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_METAX_API - DESTROY(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_QY_API - DESTROY(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DESTROY(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DESTROY -} diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating.h b/src/infiniop/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating.h deleted file mode 100644 index 6be2670e2..000000000 --- a/src/infiniop/ops/fused_gated_delta_net_gating/fused_gated_delta_net_gating.h +++ /dev/null @@ -1,35 +0,0 @@ -#ifndef __FUSED_GATED_DELTA_NET_GATING_H__ -#define __FUSED_GATED_DELTA_NET_GATING_H__ - -#include "../../operator.h" -#include "info.h" - -#include - -#define DESCRIPTOR(NAMESPACE) \ - namespace op::fused_gated_delta_net_gating::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - FusedGatedDeltaNetGatingInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor(Opaque *opaque, FusedGatedDeltaNetGatingInfo info, size_t workspace_size, infiniDevice_t device_type, int device_id) \ - : InfiniopDescriptor{device_type, device_id}, _opaque(opaque), _info(std::move(info)), _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create(infiniopHandle_t handle, Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t g_desc, infiniopTensorDescriptor_t beta_output_desc, \ - infiniopTensorDescriptor_t A_log_desc, infiniopTensorDescriptor_t a_desc, \ - infiniopTensorDescriptor_t b_desc, infiniopTensorDescriptor_t dt_bias_desc, \ - float beta, float threshold); \ - \ - infiniStatus_t calculate(void *workspace, size_t workspace_size, void *g, void *beta_output, \ - const void *A_log, const void *a, const void *b, const void *dt_bias, void *stream) const; \ - }; \ - } - -#endif diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/info.h b/src/infiniop/ops/fused_gated_delta_net_gating/info.h deleted file mode 100644 index 541751353..000000000 --- a/src/infiniop/ops/fused_gated_delta_net_gating/info.h +++ /dev/null @@ -1,90 +0,0 @@ -#ifndef __FUSED_GATED_DELTA_NET_GATING_INFO_H__ -#define __FUSED_GATED_DELTA_NET_GATING_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" - -#include - -namespace op::fused_gated_delta_net_gating { - -class FusedGatedDeltaNetGatingInfo { - FusedGatedDeltaNetGatingInfo() = default; - -public: - infiniDtype_t input_dtype; - size_t batch_size; - size_t seq_len; - size_t hidden; - std::vector g_strides; - std::vector beta_output_strides; - std::vector A_log_strides; - std::vector a_strides; - std::vector b_strides; - std::vector dt_bias_strides; - float beta; - float threshold; - - size_t numel() const { - return batch_size * seq_len * hidden; - } - - static utils::Result create( - infiniopTensorDescriptor_t g_desc, - infiniopTensorDescriptor_t beta_output_desc, - infiniopTensorDescriptor_t A_log_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t dt_bias_desc, - float beta, - float threshold) { - - if (g_desc->dtype() != INFINI_DTYPE_F32 || beta_output_desc->dtype() != INFINI_DTYPE_F32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto input_dtype = a_desc->dtype(); - if (input_dtype != INFINI_DTYPE_F32 && input_dtype != INFINI_DTYPE_F16 && input_dtype != INFINI_DTYPE_BF16) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - if (b_desc->dtype() != input_dtype || A_log_desc->dtype() != input_dtype || dt_bias_desc->dtype() != input_dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - if (g_desc->ndim() != 3 || beta_output_desc->ndim() != 3 || a_desc->ndim() != 3 || b_desc->ndim() != 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (A_log_desc->ndim() != 1 || dt_bias_desc->ndim() != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const auto &shape = a_desc->shape(); - if (shape != b_desc->shape() || shape != g_desc->shape() || shape != beta_output_desc->shape()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t hidden = shape[2]; - if (A_log_desc->shape()[0] != hidden || dt_bias_desc->shape()[0] != hidden) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - FusedGatedDeltaNetGatingInfo info; - info.input_dtype = input_dtype; - info.batch_size = shape[0]; - info.seq_len = shape[1]; - info.hidden = hidden; - info.g_strides = g_desc->strides(); - info.beta_output_strides = beta_output_desc->strides(); - info.A_log_strides = A_log_desc->strides(); - info.a_strides = a_desc->strides(); - info.b_strides = b_desc->strides(); - info.dt_bias_strides = dt_bias_desc->strides(); - info.beta = beta; - info.threshold = threshold; - return utils::Result(info); - } -}; - -} // namespace op::fused_gated_delta_net_gating - -#endif diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/metax/fused_gated_delta_net_gating_metax.h b/src/infiniop/ops/fused_gated_delta_net_gating/metax/fused_gated_delta_net_gating_metax.h deleted file mode 100644 index 4ab2eaff8..000000000 --- a/src/infiniop/ops/fused_gated_delta_net_gating/metax/fused_gated_delta_net_gating_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FUSED_GATED_DELTA_NET_GATING_METAX_H__ -#define __FUSED_GATED_DELTA_NET_GATING_METAX_H__ - -#include "../fused_gated_delta_net_gating.h" - -DESCRIPTOR(metax) - -#endif // __FUSED_GATED_DELTA_NET_GATING_METAX_H__ diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/metax/fused_gated_delta_net_gating_metax.maca b/src/infiniop/ops/fused_gated_delta_net_gating/metax/fused_gated_delta_net_gating_metax.maca deleted file mode 100644 index e4a3fcd5c..000000000 --- a/src/infiniop/ops/fused_gated_delta_net_gating/metax/fused_gated_delta_net_gating_metax.maca +++ /dev/null @@ -1,194 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "fused_gated_delta_net_gating_metax.h" - -namespace op::fused_gated_delta_net_gating::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t g_desc, - infiniopTensorDescriptor_t beta_output_desc, - infiniopTensorDescriptor_t A_log_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t dt_bias_desc, - float beta, - float threshold) { - - auto result = FusedGatedDeltaNetGatingInfo::create( - g_desc, beta_output_desc, A_log_desc, a_desc, b_desc, dt_bias_desc, beta, threshold); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - result.take(), - 0, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -__device__ __forceinline__ float load_as_float(const T *ptr, ptrdiff_t offset) { - return static_cast(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as_float(const half *ptr, ptrdiff_t offset) { - return __half2float(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as_float(const cuda_bfloat16 *ptr, ptrdiff_t offset) { - return __bfloat162float(ptr[offset]); -} - -__device__ __forceinline__ float sigmoidf_stable(float x) { - if (x >= 0.0f) { - float z = expf(-x); - return 1.0f / (1.0f + z); - } - float z = expf(x); - return z / (1.0f + z); -} - -__device__ __forceinline__ float softplus_beta_threshold(float x, float beta, float threshold) { - float bx = beta * x; - return bx <= threshold ? log1pf(expf(bx)) / beta : x; -} - -template -__global__ void fused_gated_delta_net_gating_kernel( - float *g, - float *beta_output, - const T *A_log, - const T *a, - const T *b, - const T *dt_bias, - size_t total, - size_t seq_len, - size_t hidden, - ptrdiff_t g_s0, - ptrdiff_t g_s1, - ptrdiff_t g_s2, - ptrdiff_t beta_s0, - ptrdiff_t beta_s1, - ptrdiff_t beta_s2, - ptrdiff_t A_log_s0, - ptrdiff_t a_s0, - ptrdiff_t a_s1, - ptrdiff_t a_s2, - ptrdiff_t b_s0, - ptrdiff_t b_s1, - ptrdiff_t b_s2, - ptrdiff_t dt_bias_s0, - float beta, - float threshold) { - - size_t linear = blockIdx.x * blockDim.x + threadIdx.x; - if (linear >= total) { - return; - } - - size_t h = linear % hidden; - size_t tmp = linear / hidden; - size_t s = tmp % seq_len; - size_t batch = tmp / seq_len; - - ptrdiff_t g_off = static_cast(batch) * g_s0 + static_cast(s) * g_s1 + static_cast(h) * g_s2; - ptrdiff_t beta_off = static_cast(batch) * beta_s0 + static_cast(s) * beta_s1 + static_cast(h) * beta_s2; - ptrdiff_t a_off = static_cast(batch) * a_s0 + static_cast(s) * a_s1 + static_cast(h) * a_s2; - ptrdiff_t b_off = static_cast(batch) * b_s0 + static_cast(s) * b_s1 + static_cast(h) * b_s2; - - float x = load_as_float(a, a_off) + load_as_float(dt_bias, static_cast(h) * dt_bias_s0); - g[g_off] = -expf(load_as_float(A_log, static_cast(h) * A_log_s0)) * softplus_beta_threshold(x, beta, threshold); - beta_output[beta_off] = sigmoidf_stable(load_as_float(b, b_off)); -} - -template -infiniStatus_t launch_kernel(const FusedGatedDeltaNetGatingInfo &info, - float *g, - float *beta_output, - const void *A_log, - const void *a, - const void *b, - const void *dt_bias, - hcStream_t stream) { - size_t total = info.numel(); - if (total == 0) { - return INFINI_STATUS_SUCCESS; - } - - constexpr int block = 256; - int grid = static_cast((total + block - 1) / block); - fused_gated_delta_net_gating_kernel<<>>( - g, - beta_output, - static_cast(A_log), - static_cast(a), - static_cast(b), - static_cast(dt_bias), - total, - info.seq_len, - info.hidden, - info.g_strides[0], - info.g_strides[1], - info.g_strides[2], - info.beta_output_strides[0], - info.beta_output_strides[1], - info.beta_output_strides[2], - info.A_log_strides[0], - info.a_strides[0], - info.a_strides[1], - info.a_strides[2], - info.b_strides[0], - info.b_strides[1], - info.b_strides[2], - info.dt_bias_strides[0], - info.beta, - info.threshold); - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *g, - void *beta_output, - const void *A_log, - const void *a, - const void *b, - const void *dt_bias, - void *stream) const { - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - hcStream_t cuda_stream = reinterpret_cast(stream); - switch (_info.input_dtype) { - case INFINI_DTYPE_F32: - return launch_kernel(_info, static_cast(g), static_cast(beta_output), A_log, a, b, dt_bias, cuda_stream); - case INFINI_DTYPE_F16: - return launch_kernel(_info, static_cast(g), static_cast(beta_output), A_log, a, b, dt_bias, cuda_stream); - case INFINI_DTYPE_BF16: - return launch_kernel(_info, static_cast(g), static_cast(beta_output), A_log, a, b, dt_bias, cuda_stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::fused_gated_delta_net_gating::metax diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/moore/fused_gated_delta_net_gating_moore.h b/src/infiniop/ops/fused_gated_delta_net_gating/moore/fused_gated_delta_net_gating_moore.h deleted file mode 100644 index 272df06e5..000000000 --- a/src/infiniop/ops/fused_gated_delta_net_gating/moore/fused_gated_delta_net_gating_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FUSED_GATED_DELTA_NET_GATING_MOORE_H__ -#define __FUSED_GATED_DELTA_NET_GATING_MOORE_H__ - -#include "../fused_gated_delta_net_gating.h" - -DESCRIPTOR(moore) - -#endif // __FUSED_GATED_DELTA_NET_GATING_MOORE_H__ diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/moore/fused_gated_delta_net_gating_moore.mu b/src/infiniop/ops/fused_gated_delta_net_gating/moore/fused_gated_delta_net_gating_moore.mu deleted file mode 100644 index fc2f7c45d..000000000 --- a/src/infiniop/ops/fused_gated_delta_net_gating/moore/fused_gated_delta_net_gating_moore.mu +++ /dev/null @@ -1,194 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_handle.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "fused_gated_delta_net_gating_moore.h" - -namespace op::fused_gated_delta_net_gating::moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t g_desc, - infiniopTensorDescriptor_t beta_output_desc, - infiniopTensorDescriptor_t A_log_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t dt_bias_desc, - float beta, - float threshold) { - - auto result = FusedGatedDeltaNetGatingInfo::create( - g_desc, beta_output_desc, A_log_desc, a_desc, b_desc, dt_bias_desc, beta, threshold); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - result.take(), - 0, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -__device__ __forceinline__ float load_as_float(const T *ptr, ptrdiff_t offset) { - return static_cast(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as_float(const half *ptr, ptrdiff_t offset) { - return __half2float(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as_float(const cuda_bfloat16 *ptr, ptrdiff_t offset) { - return __bfloat162float(ptr[offset]); -} - -__device__ __forceinline__ float sigmoidf_stable(float x) { - if (x >= 0.0f) { - float z = expf(-x); - return 1.0f / (1.0f + z); - } - float z = expf(x); - return z / (1.0f + z); -} - -__device__ __forceinline__ float softplus_beta_threshold(float x, float beta, float threshold) { - float bx = beta * x; - return bx <= threshold ? log1pf(expf(bx)) / beta : x; -} - -template -__global__ void fused_gated_delta_net_gating_kernel( - float *g, - float *beta_output, - const T *A_log, - const T *a, - const T *b, - const T *dt_bias, - size_t total, - size_t seq_len, - size_t hidden, - ptrdiff_t g_s0, - ptrdiff_t g_s1, - ptrdiff_t g_s2, - ptrdiff_t beta_s0, - ptrdiff_t beta_s1, - ptrdiff_t beta_s2, - ptrdiff_t A_log_s0, - ptrdiff_t a_s0, - ptrdiff_t a_s1, - ptrdiff_t a_s2, - ptrdiff_t b_s0, - ptrdiff_t b_s1, - ptrdiff_t b_s2, - ptrdiff_t dt_bias_s0, - float beta, - float threshold) { - - size_t linear = blockIdx.x * blockDim.x + threadIdx.x; - if (linear >= total) { - return; - } - - size_t h = linear % hidden; - size_t tmp = linear / hidden; - size_t s = tmp % seq_len; - size_t batch = tmp / seq_len; - - ptrdiff_t g_off = static_cast(batch) * g_s0 + static_cast(s) * g_s1 + static_cast(h) * g_s2; - ptrdiff_t beta_off = static_cast(batch) * beta_s0 + static_cast(s) * beta_s1 + static_cast(h) * beta_s2; - ptrdiff_t a_off = static_cast(batch) * a_s0 + static_cast(s) * a_s1 + static_cast(h) * a_s2; - ptrdiff_t b_off = static_cast(batch) * b_s0 + static_cast(s) * b_s1 + static_cast(h) * b_s2; - - float x = load_as_float(a, a_off) + load_as_float(dt_bias, static_cast(h) * dt_bias_s0); - g[g_off] = -expf(load_as_float(A_log, static_cast(h) * A_log_s0)) * softplus_beta_threshold(x, beta, threshold); - beta_output[beta_off] = sigmoidf_stable(load_as_float(b, b_off)); -} - -template -infiniStatus_t launch_kernel(const FusedGatedDeltaNetGatingInfo &info, - float *g, - float *beta_output, - const void *A_log, - const void *a, - const void *b, - const void *dt_bias, - musaStream_t stream) { - size_t total = info.numel(); - if (total == 0) { - return INFINI_STATUS_SUCCESS; - } - - constexpr int block = 256; - int grid = static_cast((total + block - 1) / block); - fused_gated_delta_net_gating_kernel<<>>( - g, - beta_output, - static_cast(A_log), - static_cast(a), - static_cast(b), - static_cast(dt_bias), - total, - info.seq_len, - info.hidden, - info.g_strides[0], - info.g_strides[1], - info.g_strides[2], - info.beta_output_strides[0], - info.beta_output_strides[1], - info.beta_output_strides[2], - info.A_log_strides[0], - info.a_strides[0], - info.a_strides[1], - info.a_strides[2], - info.b_strides[0], - info.b_strides[1], - info.b_strides[2], - info.dt_bias_strides[0], - info.beta, - info.threshold); - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *g, - void *beta_output, - const void *A_log, - const void *a, - const void *b, - const void *dt_bias, - void *stream) const { - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - musaStream_t cuda_stream = reinterpret_cast(stream); - switch (_info.input_dtype) { - case INFINI_DTYPE_F32: - return launch_kernel(_info, static_cast(g), static_cast(beta_output), A_log, a, b, dt_bias, cuda_stream); - case INFINI_DTYPE_F16: - return launch_kernel(_info, static_cast(g), static_cast(beta_output), A_log, a, b, dt_bias, cuda_stream); - case INFINI_DTYPE_BF16: - return launch_kernel(_info, static_cast(g), static_cast(beta_output), A_log, a, b, dt_bias, cuda_stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::fused_gated_delta_net_gating::moore diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/nvidia/fused_gated_delta_net_gating_nvidia.cu b/src/infiniop/ops/fused_gated_delta_net_gating/nvidia/fused_gated_delta_net_gating_nvidia.cu deleted file mode 100644 index 627f026bf..000000000 --- a/src/infiniop/ops/fused_gated_delta_net_gating/nvidia/fused_gated_delta_net_gating_nvidia.cu +++ /dev/null @@ -1,196 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "fused_gated_delta_net_gating_nvidia.cuh" - -#include -#include - -namespace op::fused_gated_delta_net_gating::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t g_desc, - infiniopTensorDescriptor_t beta_output_desc, - infiniopTensorDescriptor_t A_log_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t dt_bias_desc, - float beta, - float threshold) { - - auto result = FusedGatedDeltaNetGatingInfo::create( - g_desc, beta_output_desc, A_log_desc, a_desc, b_desc, dt_bias_desc, beta, threshold); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - result.take(), - 0, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -__device__ __forceinline__ float load_as_float(const T *ptr, ptrdiff_t offset) { - return static_cast(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as_float(const half *ptr, ptrdiff_t offset) { - return __half2float(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as_float<__nv_bfloat16>(const __nv_bfloat16 *ptr, ptrdiff_t offset) { - return __bfloat162float(ptr[offset]); -} - -__device__ __forceinline__ float sigmoidf_stable(float x) { - if (x >= 0.0f) { - float z = expf(-x); - return 1.0f / (1.0f + z); - } - float z = expf(x); - return z / (1.0f + z); -} - -__device__ __forceinline__ float softplus_beta_threshold(float x, float beta, float threshold) { - float bx = beta * x; - return bx <= threshold ? log1pf(expf(bx)) / beta : x; -} - -template -__global__ void fused_gated_delta_net_gating_kernel( - float *g, - float *beta_output, - const T *A_log, - const T *a, - const T *b, - const T *dt_bias, - size_t total, - size_t seq_len, - size_t hidden, - ptrdiff_t g_s0, - ptrdiff_t g_s1, - ptrdiff_t g_s2, - ptrdiff_t beta_s0, - ptrdiff_t beta_s1, - ptrdiff_t beta_s2, - ptrdiff_t A_log_s0, - ptrdiff_t a_s0, - ptrdiff_t a_s1, - ptrdiff_t a_s2, - ptrdiff_t b_s0, - ptrdiff_t b_s1, - ptrdiff_t b_s2, - ptrdiff_t dt_bias_s0, - float beta, - float threshold) { - - size_t linear = blockIdx.x * blockDim.x + threadIdx.x; - if (linear >= total) { - return; - } - - size_t h = linear % hidden; - size_t tmp = linear / hidden; - size_t s = tmp % seq_len; - size_t batch = tmp / seq_len; - - ptrdiff_t g_off = static_cast(batch) * g_s0 + static_cast(s) * g_s1 + static_cast(h) * g_s2; - ptrdiff_t beta_off = static_cast(batch) * beta_s0 + static_cast(s) * beta_s1 + static_cast(h) * beta_s2; - ptrdiff_t a_off = static_cast(batch) * a_s0 + static_cast(s) * a_s1 + static_cast(h) * a_s2; - ptrdiff_t b_off = static_cast(batch) * b_s0 + static_cast(s) * b_s1 + static_cast(h) * b_s2; - - float x = load_as_float(a, a_off) + load_as_float(dt_bias, static_cast(h) * dt_bias_s0); - g[g_off] = -expf(load_as_float(A_log, static_cast(h) * A_log_s0)) * softplus_beta_threshold(x, beta, threshold); - beta_output[beta_off] = sigmoidf_stable(load_as_float(b, b_off)); -} - -template -infiniStatus_t launch_kernel(const FusedGatedDeltaNetGatingInfo &info, - float *g, - float *beta_output, - const void *A_log, - const void *a, - const void *b, - const void *dt_bias, - cudaStream_t stream) { - size_t total = info.numel(); - if (total == 0) { - return INFINI_STATUS_SUCCESS; - } - - constexpr int block = 256; - int grid = static_cast((total + block - 1) / block); - fused_gated_delta_net_gating_kernel<<>>( - g, - beta_output, - static_cast(A_log), - static_cast(a), - static_cast(b), - static_cast(dt_bias), - total, - info.seq_len, - info.hidden, - info.g_strides[0], - info.g_strides[1], - info.g_strides[2], - info.beta_output_strides[0], - info.beta_output_strides[1], - info.beta_output_strides[2], - info.A_log_strides[0], - info.a_strides[0], - info.a_strides[1], - info.a_strides[2], - info.b_strides[0], - info.b_strides[1], - info.b_strides[2], - info.dt_bias_strides[0], - info.beta, - info.threshold); - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *g, - void *beta_output, - const void *A_log, - const void *a, - const void *b, - const void *dt_bias, - void *stream) const { - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - cudaStream_t cuda_stream = reinterpret_cast(stream); - switch (_info.input_dtype) { - case INFINI_DTYPE_F32: - return launch_kernel(_info, static_cast(g), static_cast(beta_output), A_log, a, b, dt_bias, cuda_stream); - case INFINI_DTYPE_F16: - return launch_kernel(_info, static_cast(g), static_cast(beta_output), A_log, a, b, dt_bias, cuda_stream); - case INFINI_DTYPE_BF16: - return launch_kernel<__nv_bfloat16>(_info, static_cast(g), static_cast(beta_output), A_log, a, b, dt_bias, cuda_stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::fused_gated_delta_net_gating::nvidia diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/nvidia/fused_gated_delta_net_gating_nvidia.cuh b/src/infiniop/ops/fused_gated_delta_net_gating/nvidia/fused_gated_delta_net_gating_nvidia.cuh deleted file mode 100644 index a427fde6a..000000000 --- a/src/infiniop/ops/fused_gated_delta_net_gating/nvidia/fused_gated_delta_net_gating_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __FUSED_GATED_DELTA_NET_GATING_NVIDIA_CUH__ -#define __FUSED_GATED_DELTA_NET_GATING_NVIDIA_CUH__ - -#include "../fused_gated_delta_net_gating.h" - -DESCRIPTOR(nvidia) - -#endif diff --git a/src/infiniop/ops/fused_gated_delta_net_gating/operator.cc b/src/infiniop/ops/fused_gated_delta_net_gating/operator.cc deleted file mode 100644 index 14296e846..000000000 --- a/src/infiniop/ops/fused_gated_delta_net_gating/operator.cc +++ /dev/null @@ -1,185 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/fused_gated_delta_net_gating.h" - -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) -#include "nvidia/fused_gated_delta_net_gating_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/fused_gated_delta_net_gating_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/fused_gated_delta_net_gating_moore.h" -#endif - -__INFINI_C __export infiniStatus_t -infiniopCreateFusedGatedDeltaNetGatingDescriptor( - infiniopHandle_t handle, - infiniopFusedGatedDeltaNetGatingDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t g_desc, - infiniopTensorDescriptor_t beta_output_desc, - infiniopTensorDescriptor_t A_log_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t dt_bias_desc, - float beta, - float threshold) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::fused_gated_delta_net_gating::NAMESPACE::Descriptor::create( \ - handle, reinterpret_cast(desc_ptr), \ - g_desc, beta_output_desc, A_log_desc, a_desc, b_desc, dt_bias_desc, beta, threshold) - - switch (handle->device) { -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C __export infiniStatus_t -infiniopGetFusedGatedDeltaNetGatingWorkspaceSize( - infiniopFusedGatedDeltaNetGatingDescriptor_t desc, - size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ALI_API - GET(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET -} - -__INFINI_C __export infiniStatus_t -infiniopFusedGatedDeltaNetGating( - infiniopFusedGatedDeltaNetGatingDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *g, - void *beta_output, - const void *A_log, - const void *a, - const void *b, - const void *dt_bias, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc)->calculate( \ - workspace, workspace_size, g, beta_output, A_log, a, b, dt_bias, stream) - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C __export infiniStatus_t -infiniopDestroyFusedGatedDeltaNetGatingDescriptor( - infiniopFusedGatedDeltaNetGatingDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ALI_API - DELETE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/gaussian_nll_loss/cpu/gaussian_nll_loss_cpu.cc b/src/infiniop/ops/gaussian_nll_loss/cpu/gaussian_nll_loss_cpu.cc deleted file mode 100644 index 80c153463..000000000 --- a/src/infiniop/ops/gaussian_nll_loss/cpu/gaussian_nll_loss_cpu.cc +++ /dev/null @@ -1,159 +0,0 @@ -#include "gaussian_nll_loss_cpu.h" -#include "../../../../utils.h" -#include - -namespace op::gaussian_nll_loss::cpu { - -utils::Result GaussianNllLossInfo::create( - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t var_desc, - infiniopTensorDescriptor_t y_desc, - int full, - double eps, - int reduction) { - - auto input_shape = input_desc->shape(); - auto target_shape = target_desc->shape(); - auto var_shape = var_desc->shape(); - auto y_shape = y_desc->shape(); - - if (input_shape != target_shape || input_shape != var_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - Reduction red = static_cast(reduction); - std::vector expected_y_shape; - if (red == Reduction::NONE) { - expected_y_shape = input_shape; - } else { - expected_y_shape = {}; - } - - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - GaussianNllLossInfo info; - info.input_size = input_desc->numel(); - info.full = full; - info.eps = eps; - info.reduction = red; - - return utils::Result(std::move(info)); -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t var_desc, - int full, - double eps, - int reduction) { - - auto dtype = input_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto info_result = GaussianNllLossInfo::create(input_desc, target_desc, var_desc, y_desc, full, eps, reduction); - CHECK_RESULT(info_result); - - *desc_ptr = new Descriptor(dtype, info_result.take(), handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -void gaussian_nll_loss_impl( - const GaussianNllLossInfo &info, - T *y, - const T *input, - const T *target, - const T *var) { - - size_t n = info.input_size; - const double eps_val = info.eps; - const double log_2pi = std::log(2.0 * 3.14159265358979323846); - - if (info.reduction == Reduction::NONE) { - // Element-wise loss - for (size_t i = 0; i < n; ++i) { - const double diff = utils::cast(input[i]) - utils::cast(target[i]); - double var_val = utils::cast(var[i]); - if (var_val < eps_val) { - var_val = eps_val; - } - double loss = 0.5 * (std::log(var_val) + (diff * diff) / var_val); - if (info.full) { - loss += 0.5 * log_2pi; - } - y[i] = utils::cast(loss); - } - } else { - // Sum or Mean - double sum = 0.0; - for (size_t i = 0; i < n; ++i) { - const double diff = utils::cast(input[i]) - utils::cast(target[i]); - double var_val = utils::cast(var[i]); - if (var_val < eps_val) { - var_val = eps_val; - } - double loss = 0.5 * (std::log(var_val) + (diff * diff) / var_val); - if (info.full) { - loss += 0.5 * log_2pi; - } - sum += loss; - } - if (info.reduction == Reduction::MEAN) { - y[0] = utils::cast(sum / static_cast(n)); - } else { - y[0] = utils::cast(sum); - } - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - const void *var, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - gaussian_nll_loss_impl(_info, reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var)); - break; - case INFINI_DTYPE_BF16: - gaussian_nll_loss_impl(_info, reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var)); - break; - case INFINI_DTYPE_F32: - gaussian_nll_loss_impl(_info, reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var)); - break; - case INFINI_DTYPE_F64: - gaussian_nll_loss_impl(_info, reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var)); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gaussian_nll_loss::cpu diff --git a/src/infiniop/ops/gaussian_nll_loss/cpu/gaussian_nll_loss_cpu.h b/src/infiniop/ops/gaussian_nll_loss/cpu/gaussian_nll_loss_cpu.h deleted file mode 100644 index 9205f8ad7..000000000 --- a/src/infiniop/ops/gaussian_nll_loss/cpu/gaussian_nll_loss_cpu.h +++ /dev/null @@ -1,71 +0,0 @@ -#ifndef __GAUSSIAN_NLL_LOSS_CPU_H__ -#define __GAUSSIAN_NLL_LOSS_CPU_H__ - -#include "../../../devices/cpu/common_cpu.h" -#include "../../../operator.h" -#include "../../../tensor.h" -#include - -namespace op::gaussian_nll_loss::cpu { - -enum class Reduction { - NONE = 0, - MEAN = 1, - SUM = 2 -}; - -struct GaussianNllLossInfo { - size_t input_size; - int full; - double eps; - Reduction reduction; - - static utils::Result create( - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t var_desc, - infiniopTensorDescriptor_t y_desc, - int full, - double eps, - int reduction); -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - GaussianNllLossInfo _info; - - Descriptor(infiniDtype_t dtype, GaussianNllLossInfo info, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t var_desc, - int full, - double eps, - int reduction); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - const void *var, - void *stream) const; -}; - -} // namespace op::gaussian_nll_loss::cpu - -#endif // __GAUSSIAN_NLL_LOSS_CPU_H__ diff --git a/src/infiniop/ops/gaussian_nll_loss/cuda/kernel.cuh b/src/infiniop/ops/gaussian_nll_loss/cuda/kernel.cuh deleted file mode 100644 index d1d887bd4..000000000 --- a/src/infiniop/ops/gaussian_nll_loss/cuda/kernel.cuh +++ /dev/null @@ -1,211 +0,0 @@ -#pragma once -#include "../../../reduce/cuda/reduce.cuh" -#include -#include - -namespace op::cuda { - -constexpr int kGaussianNllMaxDims = 8; - -struct GaussianNllTensorMeta { - int ndim; - size_t shape[kGaussianNllMaxDims]; - ptrdiff_t strides[kGaussianNllMaxDims]; // strides in elements -}; - -template -struct GaussianNllTypeTag {}; - -template -__device__ __forceinline__ Tcompute gaussian_nll_to_compute(const half v) { - return static_cast(__half2float(v)); -} - -template -__device__ __forceinline__ Tcompute gaussian_nll_to_compute(const cuda_bfloat16 v) { - return static_cast(__bfloat162float(v)); -} - -template -__device__ __forceinline__ Tcompute gaussian_nll_to_compute(const T v) { - return static_cast(v); -} - -__device__ __forceinline__ half gaussian_nll_from_compute(const float v, GaussianNllTypeTag) { - return __float2half_rn(v); -} - -__device__ __forceinline__ cuda_bfloat16 gaussian_nll_from_compute(const float v, GaussianNllTypeTag) { - return __float2bfloat16_rn(v); -} - -template -__device__ __forceinline__ T gaussian_nll_from_compute(const Tcompute v, GaussianNllTypeTag) { - return static_cast(v); -} - -__device__ __forceinline__ size_t gaussian_nll_offset(size_t flat, const GaussianNllTensorMeta &meta) { - size_t res = 0; - for (size_t i = meta.ndim; i-- > 0;) { - res += (flat % meta.shape[i]) * meta.strides[i]; - flat /= meta.shape[i]; - } - return res; -} - -template -__global__ void gaussian_nll_loss_kernel( - T *output, - const T *input, - const T *target, - const T *var, - size_t n, - GaussianNllTensorMeta out_meta, - GaussianNllTensorMeta in_meta, - GaussianNllTensorMeta tgt_meta, - GaussianNllTensorMeta var_meta, - Tcompute eps_val, - int full) { - - const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= n) { - return; - } - - const size_t out_off = gaussian_nll_offset(idx, out_meta); - const size_t in_off = gaussian_nll_offset(idx, in_meta); - const size_t tgt_off = gaussian_nll_offset(idx, tgt_meta); - const size_t var_off = gaussian_nll_offset(idx, var_meta); - - const Tcompute diff = gaussian_nll_to_compute(input[in_off]) - gaussian_nll_to_compute(target[tgt_off]); - Tcompute var_val = gaussian_nll_to_compute(var[var_off]); - if (var_val < eps_val) { - var_val = eps_val; - } - Tcompute loss = Tcompute(0.5) * (log(var_val) + (diff * diff) / var_val); - if (full) { - loss += Tcompute(0.9189385332046727); // log(2*pi)/2 - } - output[out_off] = gaussian_nll_from_compute(loss, GaussianNllTypeTag{}); -} - -template -__global__ void gaussian_nll_loss_reduce_kernel( - Tcompute *output, - const T *input, - const T *target, - const T *var, - size_t n, - GaussianNllTensorMeta in_meta, - GaussianNllTensorMeta tgt_meta, - GaussianNllTensorMeta var_meta, - Tcompute eps_val, - int full) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - Tcompute sum = 0; - - Tcompute log_2pi = full ? Tcompute(0.9189385332046727) : Tcompute(0.0); - - for (size_t i = idx; i < n; i += blockDim.x * gridDim.x) { - const size_t in_off = gaussian_nll_offset(i, in_meta); - const size_t tgt_off = gaussian_nll_offset(i, tgt_meta); - const size_t var_off = gaussian_nll_offset(i, var_meta); - - const Tcompute diff = gaussian_nll_to_compute(input[in_off]) - gaussian_nll_to_compute(target[tgt_off]); - Tcompute var_val = gaussian_nll_to_compute(var[var_off]); - if (var_val < eps_val) { - var_val = eps_val; - } - const Tcompute loss = Tcompute(0.5) * (log(var_val) + (diff * diff) / var_val) + log_2pi; - sum += loss; - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - Tcompute block_sum = BlockReduce(temp_storage).Sum(sum); - - if (threadIdx.x == 0) { - atomicAdd(output, block_sum); - } -} - -template -__global__ void gaussian_nll_loss_finalize_kernel( - Tout *output, - const Tcompute *accum, - Tcompute scale) { - if (blockIdx.x == 0 && threadIdx.x == 0) { - const Tcompute v = (*accum) * scale; - output[0] = gaussian_nll_from_compute(v, GaussianNllTypeTag{}); - } -} - -// --------------------------------------------------------------------------- -// Compatibility wrappers for backends that still reference the older contiguous -// kernel signatures. -// --------------------------------------------------------------------------- - -template -__global__ void gaussian_nll_loss_kernel( - T *output, - const T *input, - const T *target, - const T *var, - size_t n, - T eps_val, - int full) { - - const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= n) { - return; - } - - using Tcompute = std::conditional_t, double, float>; - const Tcompute eps_c = gaussian_nll_to_compute(eps_val); - const Tcompute diff = gaussian_nll_to_compute(input[idx]) - gaussian_nll_to_compute(target[idx]); - Tcompute var_val = gaussian_nll_to_compute(var[idx]); - if (var_val < eps_c) { - var_val = eps_c; - } - Tcompute loss = Tcompute(0.5) * (log(var_val) + (diff * diff) / var_val); - if (full) { - loss += Tcompute(0.9189385332046727); - } - output[idx] = gaussian_nll_from_compute(loss, GaussianNllTypeTag{}); -} - -template -__global__ void gaussian_nll_loss_reduce_kernel( - T *output, - const T *input, - const T *target, - const T *var, - size_t n, - Tcompute eps_val, - int full) { - - const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - Tcompute sum = 0; - - const Tcompute log_2pi = full ? Tcompute(0.9189385332046727) : Tcompute(0.0); - for (size_t i = idx; i < n; i += blockDim.x * gridDim.x) { - const Tcompute diff = gaussian_nll_to_compute(input[i]) - gaussian_nll_to_compute(target[i]); - Tcompute var_val = gaussian_nll_to_compute(var[i]); - if (var_val < eps_val) { - var_val = eps_val; - } - const Tcompute loss = Tcompute(0.5) * (log(var_val) + (diff * diff) / var_val) + log_2pi; - sum += loss; - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - const Tcompute block_sum = BlockReduce(temp_storage).Sum(sum); - - if (threadIdx.x == 0) { - atomicAdd(reinterpret_cast(output), block_sum); - } -} - -} // namespace op::cuda diff --git a/src/infiniop/ops/gaussian_nll_loss/metax/gaussian_nll_loss_metax.h b/src/infiniop/ops/gaussian_nll_loss/metax/gaussian_nll_loss_metax.h deleted file mode 100644 index 904d0e073..000000000 --- a/src/infiniop/ops/gaussian_nll_loss/metax/gaussian_nll_loss_metax.h +++ /dev/null @@ -1,83 +0,0 @@ -#ifndef __GAUSSIAN_NLL_LOSS_METAX_H__ -#define __GAUSSIAN_NLL_LOSS_METAX_H__ - -#include "../../../operator.h" - -namespace op::gaussian_nll_loss::metax { - -enum class Reduction { - NONE = 0, - MEAN = 1, - SUM = 2 -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t input_size; - size_t ndim; - std::vector shape; - std::vector y_strides; - std::vector input_strides; - std::vector target_strides; - std::vector var_strides; - int full; - double eps; - Reduction reduction; - void *reduce_buffer; - - Descriptor(infiniDtype_t dtype, - size_t input_size, - size_t ndim, - std::vector shape, - std::vector y_strides, - std::vector input_strides, - std::vector target_strides, - std::vector var_strides, - int full, - double eps, - Reduction reduction, - void *reduce_buffer, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - input_size(input_size), - ndim(ndim), - shape(std::move(shape)), - y_strides(std::move(y_strides)), - input_strides(std::move(input_strides)), - target_strides(std::move(target_strides)), - var_strides(std::move(var_strides)), - full(full), - eps(eps), - reduction(reduction), - reduce_buffer(reduce_buffer) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t var_desc, - int full, - double eps, - int reduction); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - const void *var, - void *stream) const; -}; - -} // namespace op::gaussian_nll_loss::metax - -#endif // __GAUSSIAN_NLL_LOSS_METAX_H__ diff --git a/src/infiniop/ops/gaussian_nll_loss/metax/gaussian_nll_loss_metax.maca b/src/infiniop/ops/gaussian_nll_loss/metax/gaussian_nll_loss_metax.maca deleted file mode 100644 index 09354825a..000000000 --- a/src/infiniop/ops/gaussian_nll_loss/metax/gaussian_nll_loss_metax.maca +++ /dev/null @@ -1,255 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "gaussian_nll_loss_metax.h" - -namespace op::gaussian_nll_loss::metax { - -Descriptor::~Descriptor() { - if (reduce_buffer != nullptr) { - hcFree(reduce_buffer); - reduce_buffer = nullptr; - } -} - -static bool build_meta( - op::cuda::GaussianNllTensorMeta &meta, - size_t ndim, - const std::vector &shape, - const std::vector &strides) { - - if (ndim > static_cast(op::cuda::kGaussianNllMaxDims)) { - return false; - } - - meta.ndim = static_cast(ndim); - for (size_t i = 0; i < static_cast(op::cuda::kGaussianNllMaxDims); ++i) { - meta.shape[i] = (i < ndim) ? shape[i] : 1; - meta.strides[i] = (i < ndim) ? strides[i] : 0; - } - return true; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t var_desc, - int full, - double eps, - int reduction) { - - auto dtype = input_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto input_shape = input_desc->shape(); - auto target_shape = target_desc->shape(); - auto var_shape = var_desc->shape(); - auto y_shape = y_desc->shape(); - - if (input_shape != target_shape || input_shape != var_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - Reduction red = static_cast(reduction); - std::vector expected_y_shape; - if (red == Reduction::NONE) { - expected_y_shape = input_shape; - } else { - expected_y_shape = {}; - } - - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - void *reduce_buffer = nullptr; - if (red != Reduction::NONE) { - if (dtype == INFINI_DTYPE_F16 || dtype == INFINI_DTYPE_BF16) { - CHECK_METAX(hcMalloc(&reduce_buffer, sizeof(float))); - } - } - - *desc_ptr = new Descriptor( - dtype, - input_desc->numel(), - input_desc->ndim(), - input_shape, - y_desc->strides(), - input_desc->strides(), - target_desc->strides(), - var_desc->strides(), - full, - eps, - red, - reduce_buffer, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - const void *var, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - constexpr int BLOCK_SIZE = 256; - int num_blocks = static_cast((input_size + BLOCK_SIZE - 1) / BLOCK_SIZE); - - if (reduction == Reduction::NONE) { - op::cuda::GaussianNllTensorMeta out_meta{}; - op::cuda::GaussianNllTensorMeta in_meta{}; - op::cuda::GaussianNllTensorMeta tgt_meta{}; - op::cuda::GaussianNllTensorMeta var_meta{}; - - if (!build_meta(out_meta, ndim, shape, y_strides) || !build_meta(in_meta, ndim, shape, input_strides) || !build_meta(tgt_meta, ndim, shape, target_strides) || !build_meta(var_meta, ndim, shape, var_strides)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: { - float eps_val = static_cast(eps); - cuda::gaussian_nll_loss_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, out_meta, in_meta, tgt_meta, var_meta, eps_val, full); - break; - } - case INFINI_DTYPE_BF16: { - float eps_val = static_cast(eps); - cuda::gaussian_nll_loss_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, out_meta, in_meta, tgt_meta, var_meta, eps_val, full); - break; - } - case INFINI_DTYPE_F32: { - float eps_val = static_cast(eps); - cuda::gaussian_nll_loss_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, out_meta, in_meta, tgt_meta, var_meta, eps_val, full); - break; - } - case INFINI_DTYPE_F64: { - double eps_val = eps; - cuda::gaussian_nll_loss_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, out_meta, in_meta, tgt_meta, var_meta, eps_val, full); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - // Sum or Mean reduction (scalar output) - op::cuda::GaussianNllTensorMeta in_meta{}; - op::cuda::GaussianNllTensorMeta tgt_meta{}; - op::cuda::GaussianNllTensorMeta var_meta{}; - - if (!build_meta(in_meta, ndim, shape, input_strides) || !build_meta(tgt_meta, ndim, shape, target_strides) || !build_meta(var_meta, ndim, shape, var_strides)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const bool is_mean = (reduction == Reduction::MEAN); - switch (_dtype) { - case INFINI_DTYPE_F16: { - float *accum = reinterpret_cast(reduce_buffer); - if (accum == nullptr) { - return INFINI_STATUS_INTERNAL_ERROR; - } - CHECK_METAX(hcMemsetAsync(accum, 0, sizeof(float), cuda_stream)); - float eps_val = static_cast(eps); - - const int reduce_blocks = std::min(num_blocks, 1024); - cuda::gaussian_nll_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, in_meta, tgt_meta, var_meta, eps_val, full); - - const float scale = is_mean ? (1.0f / static_cast(input_size)) : 1.0f; - cuda::gaussian_nll_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, scale); - break; - } - case INFINI_DTYPE_BF16: { - float *accum = reinterpret_cast(reduce_buffer); - if (accum == nullptr) { - return INFINI_STATUS_INTERNAL_ERROR; - } - CHECK_METAX(hcMemsetAsync(accum, 0, sizeof(float), cuda_stream)); - float eps_val = static_cast(eps); - - const int reduce_blocks = std::min(num_blocks, 1024); - cuda::gaussian_nll_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, in_meta, tgt_meta, var_meta, eps_val, full); - - const float scale = is_mean ? (1.0f / static_cast(input_size)) : 1.0f; - cuda::gaussian_nll_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, scale); - break; - } - case INFINI_DTYPE_F32: { - float eps_val = static_cast(eps); - float *accum = reinterpret_cast(y); - CHECK_METAX(hcMemsetAsync(accum, 0, sizeof(float), cuda_stream)); - const int reduce_blocks = std::min(num_blocks, 1024); - cuda::gaussian_nll_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, in_meta, tgt_meta, var_meta, eps_val, full); - const float scale = is_mean ? (1.0f / static_cast(input_size)) : 1.0f; - cuda::gaussian_nll_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, scale); - break; - } - case INFINI_DTYPE_F64: { - double eps_val = eps; - double *accum = reinterpret_cast(y); - CHECK_METAX(hcMemsetAsync(accum, 0, sizeof(double), cuda_stream)); - const int reduce_blocks = std::min(num_blocks, 1024); - cuda::gaussian_nll_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, in_meta, tgt_meta, var_meta, eps_val, full); - const double scale = is_mean ? (1.0 / static_cast(input_size)) : 1.0; - cuda::gaussian_nll_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, scale); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gaussian_nll_loss::metax diff --git a/src/infiniop/ops/gaussian_nll_loss/moore/gaussian_nll_loss_moore.h b/src/infiniop/ops/gaussian_nll_loss/moore/gaussian_nll_loss_moore.h deleted file mode 100644 index c5523714b..000000000 --- a/src/infiniop/ops/gaussian_nll_loss/moore/gaussian_nll_loss_moore.h +++ /dev/null @@ -1,82 +0,0 @@ -#ifndef __GAUSSIAN_NLL_LOSS_MOORE_H__ -#define __GAUSSIAN_NLL_LOSS_MOORE_H__ - -#include "../../../operator.h" - -namespace op::gaussian_nll_loss::moore { - -enum class Reduction { - NONE = 0, - MEAN = 1, - SUM = 2 -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t input_size; - size_t ndim; - std::vector shape; - std::vector y_strides; - std::vector input_strides; - std::vector target_strides; - std::vector var_strides; - int full; - double eps; - Reduction reduction; - void *reduce_buffer; - - Descriptor(infiniDtype_t dtype, - size_t input_size, - size_t ndim, - std::vector shape, - std::vector y_strides, - std::vector input_strides, - std::vector target_strides, - std::vector var_strides, - int full, - double eps, - Reduction reduction, - void *reduce_buffer, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - input_size(input_size), - ndim(ndim), - shape(std::move(shape)), - y_strides(std::move(y_strides)), - input_strides(std::move(input_strides)), - target_strides(std::move(target_strides)), - var_strides(std::move(var_strides)), - full(full), - eps(eps), - reduction(reduction), - reduce_buffer(reduce_buffer) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t var_desc, - int full, - double eps, - int reduction); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - const void *var, - void *stream) const; -}; -} // namespace op::gaussian_nll_loss::moore - -#endif // __GAUSSIAN_NLL_LOSS_MOORE_H__ diff --git a/src/infiniop/ops/gaussian_nll_loss/moore/gaussian_nll_loss_moore.mu b/src/infiniop/ops/gaussian_nll_loss/moore/gaussian_nll_loss_moore.mu deleted file mode 100644 index 76348518c..000000000 --- a/src/infiniop/ops/gaussian_nll_loss/moore/gaussian_nll_loss_moore.mu +++ /dev/null @@ -1,255 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "gaussian_nll_loss_moore.h" - -namespace op::gaussian_nll_loss::moore { - -Descriptor::~Descriptor() { - if (reduce_buffer != nullptr) { - musaFree(reduce_buffer); - reduce_buffer = nullptr; - } -} - -static bool build_meta( - op::cuda::GaussianNllTensorMeta &meta, - size_t ndim, - const std::vector &shape, - const std::vector &strides) { - - if (ndim > static_cast(op::cuda::kGaussianNllMaxDims)) { - return false; - } - - meta.ndim = static_cast(ndim); - for (size_t i = 0; i < static_cast(op::cuda::kGaussianNllMaxDims); ++i) { - meta.shape[i] = (i < ndim) ? shape[i] : 1; - meta.strides[i] = (i < ndim) ? strides[i] : 0; - } - return true; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t var_desc, - int full, - double eps, - int reduction) { - - auto dtype = input_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto input_shape = input_desc->shape(); - auto target_shape = target_desc->shape(); - auto var_shape = var_desc->shape(); - auto y_shape = y_desc->shape(); - - if (input_shape != target_shape || input_shape != var_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - Reduction red = static_cast(reduction); - std::vector expected_y_shape; - if (red == Reduction::NONE) { - expected_y_shape = input_shape; - } else { - expected_y_shape = {}; - } - - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - void *reduce_buffer = nullptr; - if (red != Reduction::NONE) { - if (dtype == INFINI_DTYPE_F16 || dtype == INFINI_DTYPE_BF16) { - CHECK_MOORE(musaMalloc(&reduce_buffer, sizeof(float))); - } - } - - *desc_ptr = new Descriptor( - dtype, - input_desc->numel(), - input_desc->ndim(), - input_shape, - y_desc->strides(), - input_desc->strides(), - target_desc->strides(), - var_desc->strides(), - full, - eps, - red, - reduce_buffer, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - const void *var, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - constexpr int BLOCK_SIZE = 256; - int num_blocks = static_cast((input_size + BLOCK_SIZE - 1) / BLOCK_SIZE); - - if (reduction == Reduction::NONE) { - op::cuda::GaussianNllTensorMeta out_meta{}; - op::cuda::GaussianNllTensorMeta in_meta{}; - op::cuda::GaussianNllTensorMeta tgt_meta{}; - op::cuda::GaussianNllTensorMeta var_meta{}; - - if (!build_meta(out_meta, ndim, shape, y_strides) || !build_meta(in_meta, ndim, shape, input_strides) || !build_meta(tgt_meta, ndim, shape, target_strides) || !build_meta(var_meta, ndim, shape, var_strides)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: { - float eps_val = static_cast(eps); - cuda::gaussian_nll_loss_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, out_meta, in_meta, tgt_meta, var_meta, eps_val, full); - break; - } - case INFINI_DTYPE_BF16: { - float eps_val = static_cast(eps); - cuda::gaussian_nll_loss_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, out_meta, in_meta, tgt_meta, var_meta, eps_val, full); - break; - } - case INFINI_DTYPE_F32: { - float eps_val = static_cast(eps); - cuda::gaussian_nll_loss_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, out_meta, in_meta, tgt_meta, var_meta, eps_val, full); - break; - } - case INFINI_DTYPE_F64: { - double eps_val = eps; - cuda::gaussian_nll_loss_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, out_meta, in_meta, tgt_meta, var_meta, eps_val, full); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - // Sum or Mean reduction (scalar output) - op::cuda::GaussianNllTensorMeta in_meta{}; - op::cuda::GaussianNllTensorMeta tgt_meta{}; - op::cuda::GaussianNllTensorMeta var_meta{}; - - if (!build_meta(in_meta, ndim, shape, input_strides) || !build_meta(tgt_meta, ndim, shape, target_strides) || !build_meta(var_meta, ndim, shape, var_strides)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const bool is_mean = (reduction == Reduction::MEAN); - switch (_dtype) { - case INFINI_DTYPE_F16: { - float *accum = reinterpret_cast(reduce_buffer); - if (accum == nullptr) { - return INFINI_STATUS_INTERNAL_ERROR; - } - CHECK_MOORE(musaMemsetAsync(accum, 0, sizeof(float), cuda_stream)); - float eps_val = static_cast(eps); - - const int reduce_blocks = std::min(num_blocks, 1024); - cuda::gaussian_nll_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, in_meta, tgt_meta, var_meta, eps_val, full); - - const float scale = is_mean ? (1.0f / static_cast(input_size)) : 1.0f; - cuda::gaussian_nll_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, scale); - break; - } - case INFINI_DTYPE_BF16: { - float *accum = reinterpret_cast(reduce_buffer); - if (accum == nullptr) { - return INFINI_STATUS_INTERNAL_ERROR; - } - CHECK_MOORE(musaMemsetAsync(accum, 0, sizeof(float), cuda_stream)); - float eps_val = static_cast(eps); - - const int reduce_blocks = std::min(num_blocks, 1024); - cuda::gaussian_nll_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, in_meta, tgt_meta, var_meta, eps_val, full); - - const float scale = is_mean ? (1.0f / static_cast(input_size)) : 1.0f; - cuda::gaussian_nll_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, scale); - break; - } - case INFINI_DTYPE_F32: { - float eps_val = static_cast(eps); - float *accum = reinterpret_cast(y); - CHECK_MOORE(musaMemsetAsync(accum, 0, sizeof(float), cuda_stream)); - const int reduce_blocks = std::min(num_blocks, 1024); - cuda::gaussian_nll_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, in_meta, tgt_meta, var_meta, eps_val, full); - const float scale = is_mean ? (1.0f / static_cast(input_size)) : 1.0f; - cuda::gaussian_nll_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, scale); - break; - } - case INFINI_DTYPE_F64: { - double eps_val = eps; - double *accum = reinterpret_cast(y); - CHECK_MOORE(musaMemsetAsync(accum, 0, sizeof(double), cuda_stream)); - const int reduce_blocks = std::min(num_blocks, 1024); - cuda::gaussian_nll_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, in_meta, tgt_meta, var_meta, eps_val, full); - const double scale = is_mean ? (1.0 / static_cast(input_size)) : 1.0; - cuda::gaussian_nll_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, scale); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gaussian_nll_loss::moore diff --git a/src/infiniop/ops/gaussian_nll_loss/nvidia/gaussian_nll_loss_nvidia.cu b/src/infiniop/ops/gaussian_nll_loss/nvidia/gaussian_nll_loss_nvidia.cu deleted file mode 100644 index 98b243190..000000000 --- a/src/infiniop/ops/gaussian_nll_loss/nvidia/gaussian_nll_loss_nvidia.cu +++ /dev/null @@ -1,255 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "gaussian_nll_loss_nvidia.cuh" - -namespace op::gaussian_nll_loss::nvidia { - -Descriptor::~Descriptor() { - if (reduce_buffer != nullptr) { - cudaFree(reduce_buffer); - reduce_buffer = nullptr; - } -} - -static bool build_meta( - op::cuda::GaussianNllTensorMeta &meta, - size_t ndim, - const std::vector &shape, - const std::vector &strides) { - - if (ndim > static_cast(op::cuda::kGaussianNllMaxDims)) { - return false; - } - - meta.ndim = static_cast(ndim); - for (size_t i = 0; i < static_cast(op::cuda::kGaussianNllMaxDims); ++i) { - meta.shape[i] = (i < ndim) ? shape[i] : 1; - meta.strides[i] = (i < ndim) ? strides[i] : 0; - } - return true; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t var_desc, - int full, - double eps, - int reduction) { - - auto dtype = input_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto input_shape = input_desc->shape(); - auto target_shape = target_desc->shape(); - auto var_shape = var_desc->shape(); - auto y_shape = y_desc->shape(); - - if (input_shape != target_shape || input_shape != var_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - Reduction red = static_cast(reduction); - std::vector expected_y_shape; - if (red == Reduction::NONE) { - expected_y_shape = input_shape; - } else { - expected_y_shape = {}; - } - - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - void *reduce_buffer = nullptr; - if (red != Reduction::NONE) { - if (dtype == INFINI_DTYPE_F16 || dtype == INFINI_DTYPE_BF16) { - CHECK_CUDA(cudaMalloc(&reduce_buffer, sizeof(float))); - } - } - - *desc_ptr = new Descriptor( - dtype, - input_desc->numel(), - input_desc->ndim(), - input_shape, - y_desc->strides(), - input_desc->strides(), - target_desc->strides(), - var_desc->strides(), - full, - eps, - red, - reduce_buffer, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - const void *var, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - constexpr int BLOCK_SIZE = 256; - int num_blocks = static_cast((input_size + BLOCK_SIZE - 1) / BLOCK_SIZE); - - if (reduction == Reduction::NONE) { - op::cuda::GaussianNllTensorMeta out_meta{}; - op::cuda::GaussianNllTensorMeta in_meta{}; - op::cuda::GaussianNllTensorMeta tgt_meta{}; - op::cuda::GaussianNllTensorMeta var_meta{}; - - if (!build_meta(out_meta, ndim, shape, y_strides) || !build_meta(in_meta, ndim, shape, input_strides) || !build_meta(tgt_meta, ndim, shape, target_strides) || !build_meta(var_meta, ndim, shape, var_strides)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: { - float eps_val = static_cast(eps); - cuda::gaussian_nll_loss_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, out_meta, in_meta, tgt_meta, var_meta, eps_val, full); - break; - } - case INFINI_DTYPE_BF16: { - float eps_val = static_cast(eps); - cuda::gaussian_nll_loss_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, out_meta, in_meta, tgt_meta, var_meta, eps_val, full); - break; - } - case INFINI_DTYPE_F32: { - float eps_val = static_cast(eps); - cuda::gaussian_nll_loss_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, out_meta, in_meta, tgt_meta, var_meta, eps_val, full); - break; - } - case INFINI_DTYPE_F64: { - double eps_val = eps; - cuda::gaussian_nll_loss_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, out_meta, in_meta, tgt_meta, var_meta, eps_val, full); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - // Sum or Mean reduction (scalar output) - op::cuda::GaussianNllTensorMeta in_meta{}; - op::cuda::GaussianNllTensorMeta tgt_meta{}; - op::cuda::GaussianNllTensorMeta var_meta{}; - - if (!build_meta(in_meta, ndim, shape, input_strides) || !build_meta(tgt_meta, ndim, shape, target_strides) || !build_meta(var_meta, ndim, shape, var_strides)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const bool is_mean = (reduction == Reduction::MEAN); - switch (_dtype) { - case INFINI_DTYPE_F16: { - float *accum = reinterpret_cast(reduce_buffer); - if (accum == nullptr) { - return INFINI_STATUS_INTERNAL_ERROR; - } - CHECK_CUDA(cudaMemsetAsync(accum, 0, sizeof(float), cuda_stream)); - float eps_val = static_cast(eps); - - const int reduce_blocks = std::min(num_blocks, 1024); - cuda::gaussian_nll_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, in_meta, tgt_meta, var_meta, eps_val, full); - - const float scale = is_mean ? (1.0f / static_cast(input_size)) : 1.0f; - cuda::gaussian_nll_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, scale); - break; - } - case INFINI_DTYPE_BF16: { - float *accum = reinterpret_cast(reduce_buffer); - if (accum == nullptr) { - return INFINI_STATUS_INTERNAL_ERROR; - } - CHECK_CUDA(cudaMemsetAsync(accum, 0, sizeof(float), cuda_stream)); - float eps_val = static_cast(eps); - - const int reduce_blocks = std::min(num_blocks, 1024); - cuda::gaussian_nll_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, in_meta, tgt_meta, var_meta, eps_val, full); - - const float scale = is_mean ? (1.0f / static_cast(input_size)) : 1.0f; - cuda::gaussian_nll_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, scale); - break; - } - case INFINI_DTYPE_F32: { - float eps_val = static_cast(eps); - float *accum = reinterpret_cast(y); - CHECK_CUDA(cudaMemsetAsync(accum, 0, sizeof(float), cuda_stream)); - const int reduce_blocks = std::min(num_blocks, 1024); - cuda::gaussian_nll_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, in_meta, tgt_meta, var_meta, eps_val, full); - const float scale = is_mean ? (1.0f / static_cast(input_size)) : 1.0f; - cuda::gaussian_nll_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, scale); - break; - } - case INFINI_DTYPE_F64: { - double eps_val = eps; - double *accum = reinterpret_cast(y); - CHECK_CUDA(cudaMemsetAsync(accum, 0, sizeof(double), cuda_stream)); - const int reduce_blocks = std::min(num_blocks, 1024); - cuda::gaussian_nll_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - reinterpret_cast(var), - input_size, in_meta, tgt_meta, var_meta, eps_val, full); - const double scale = is_mean ? (1.0 / static_cast(input_size)) : 1.0; - cuda::gaussian_nll_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, scale); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gaussian_nll_loss::nvidia diff --git a/src/infiniop/ops/gaussian_nll_loss/nvidia/gaussian_nll_loss_nvidia.cuh b/src/infiniop/ops/gaussian_nll_loss/nvidia/gaussian_nll_loss_nvidia.cuh deleted file mode 100644 index 72a7be07a..000000000 --- a/src/infiniop/ops/gaussian_nll_loss/nvidia/gaussian_nll_loss_nvidia.cuh +++ /dev/null @@ -1,85 +0,0 @@ -#ifndef __GAUSSIAN_NLL_LOSS_NVIDIA_H__ -#define __GAUSSIAN_NLL_LOSS_NVIDIA_H__ - -#include "../../../operator.h" -#include -#include - -namespace op::gaussian_nll_loss::nvidia { - -enum class Reduction { - NONE = 0, - MEAN = 1, - SUM = 2 -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t input_size; - size_t ndim; - std::vector shape; - std::vector y_strides; - std::vector input_strides; - std::vector target_strides; - std::vector var_strides; - int full; - double eps; - Reduction reduction; - void *reduce_buffer; - - Descriptor(infiniDtype_t dtype, - size_t input_size, - size_t ndim, - std::vector shape, - std::vector y_strides, - std::vector input_strides, - std::vector target_strides, - std::vector var_strides, - int full, - double eps, - Reduction reduction, - void *reduce_buffer, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - input_size(input_size), - ndim(ndim), - shape(std::move(shape)), - y_strides(std::move(y_strides)), - input_strides(std::move(input_strides)), - target_strides(std::move(target_strides)), - var_strides(std::move(var_strides)), - full(full), - eps(eps), - reduction(reduction), - reduce_buffer(reduce_buffer) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t var_desc, - int full, - double eps, - int reduction); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - const void *var, - void *stream) const; -}; - -} // namespace op::gaussian_nll_loss::nvidia - -#endif // __GAUSSIAN_NLL_LOSS_NVIDIA_H__ diff --git a/src/infiniop/ops/gaussian_nll_loss/operator.cc b/src/infiniop/ops/gaussian_nll_loss/operator.cc deleted file mode 100644 index 063319894..000000000 --- a/src/infiniop/ops/gaussian_nll_loss/operator.cc +++ /dev/null @@ -1,181 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/gaussian_nll_loss.h" - -#ifdef ENABLE_CPU_API -#include "cpu/gaussian_nll_loss_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/gaussian_nll_loss_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/gaussian_nll_loss_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/gaussian_nll_loss_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateGaussianNllLossDescriptor( - infiniopHandle_t handle, - infiniopGaussianNllLossDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t var_desc, - int full, - double eps, - int reduction) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::gaussian_nll_loss::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - input_desc, \ - target_desc, \ - var_desc, \ - full, \ - eps, \ - reduction) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetGaussianNllLossWorkspaceSize(infiniopGaussianNllLossDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopGaussianNllLoss( - infiniopGaussianNllLossDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - const void *var, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, input, target, var, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyGaussianNllLossDescriptor(infiniopGaussianNllLossDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/gelu/bang/gelu_bang.h b/src/infiniop/ops/gelu/bang/gelu_bang.h deleted file mode 100644 index 1b3df707a..000000000 --- a/src/infiniop/ops/gelu/bang/gelu_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GELU_BANG_H__ -#define __GELU_BANG_H__ - -#include "../../../elementwise/bang/elementwise_bang.h" - -ELEMENTWISE_DESCRIPTOR(gelu, bang) - -#endif // __GELU_BANG_H__ diff --git a/src/infiniop/ops/gelu/bang/gelu_bang.mlu b/src/infiniop/ops/gelu/bang/gelu_bang.mlu deleted file mode 100644 index 18c34324e..000000000 --- a/src/infiniop/ops/gelu/bang/gelu_bang.mlu +++ /dev/null @@ -1,56 +0,0 @@ -#include "gelu_bang.h" - -LAUNCH_ELEMENTWISE_KERNEL(Gelu) - -namespace op::gelu::bang { - -typedef struct GeluOp { - static constexpr size_t num_inputs = 1; - template - static infiniStatus_t launch(Args... args) { - launchGeluKernel(args...); - return INFINI_STATUS_SUCCESS; - } -} GeluOp; - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - CHECK_SAME_SHAPE(out_desc->shape(), input_desc->shape()); - - CREATE_ELEMENTWISE_BANG_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *queue) const { - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, workspace, output, inputs, queue); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, workspace, output, inputs, queue); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, workspace, output, inputs, queue); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::gelu::bang diff --git a/src/infiniop/ops/gelu/bang/gelu_bang_internal.mlu b/src/infiniop/ops/gelu/bang/gelu_bang_internal.mlu deleted file mode 100644 index be1001007..000000000 --- a/src/infiniop/ops/gelu/bang/gelu_bang_internal.mlu +++ /dev/null @@ -1,31 +0,0 @@ -#ifndef __GELU_BANG_INTERNAL_H__ -#define __GELU_BANG_INTERNAL_H__ - -#include "../../../elementwise/bang/elementwise_bang_kernel.h" -#include "bang.h" -#include "bang_device_functions.h" -#include -#include - -typedef struct GeluOp { -public: - static constexpr size_t num_inputs = 1; - template - __mlu_device__ void operator()(T *out, const T *input, const T *unused, size_t num_elements) const { - (void)unused; - constexpr float inv_sqrt2 = 0.70710678118654752440f; - for (size_t i = 0; i < num_elements; ++i) { - float x = static_cast(input[i]); - float y = 0.5f * x * (1.0f + erff(x * inv_sqrt2)); - out[i] = static_cast(y); - } - } -} GeluOp; - -LAUNCH_ELEMENTWISE_KERNEL_IMPL(Gelu, GeluOp) - -LAUNCH_ELEMENTWISE_KERNEL_INSTANTIATE(Gelu, half) -LAUNCH_ELEMENTWISE_KERNEL_INSTANTIATE(Gelu, bfloat16_t) -LAUNCH_ELEMENTWISE_KERNEL_INSTANTIATE(Gelu, float) - -#endif // __GELU_BANG_INTERNAL_H__ diff --git a/src/infiniop/ops/gelu/cpu/gelu_cpu.cc b/src/infiniop/ops/gelu/cpu/gelu_cpu.cc deleted file mode 100644 index a057ca4bc..000000000 --- a/src/infiniop/ops/gelu/cpu/gelu_cpu.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "gelu_cpu.h" - -namespace op::gelu::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - // create CPU elementwise descriptor - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::gelu::cpu diff --git a/src/infiniop/ops/gelu/cpu/gelu_cpu.h b/src/infiniop/ops/gelu/cpu/gelu_cpu.h deleted file mode 100644 index 5a2d3fa8b..000000000 --- a/src/infiniop/ops/gelu/cpu/gelu_cpu.h +++ /dev/null @@ -1,23 +0,0 @@ -#ifndef __GELU_CPU_H__ -#define __GELU_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" - -ELEMENTWISE_DESCRIPTOR(gelu, cpu) - -#include - -namespace op::gelu::cpu { -typedef struct GeluOp { -public: - static constexpr size_t num_inputs = 1; - - template - T operator()(const T &x) const { - return static_cast(0.5 * x * (1 + erf(x / sqrt(2.0f)))); - } -} GeluOp; - -} // namespace op::gelu::cpu - -#endif // __GELU_CPU_H__ diff --git a/src/infiniop/ops/gelu/cuda/kernel.cuh b/src/infiniop/ops/gelu/cuda/kernel.cuh deleted file mode 100644 index 31fa2b2be..000000000 --- a/src/infiniop/ops/gelu/cuda/kernel.cuh +++ /dev/null @@ -1,35 +0,0 @@ -#ifndef __GELU_CUDA_H__ -#define __GELU_CUDA_H__ - -#include - -namespace op::gelu::cuda { - -typedef struct GeluOp { -public: - static constexpr size_t num_inputs = 1; - template - __device__ __forceinline__ T operator()(const T &x) const { - - if constexpr (std::is_same_v) { - float x_f = __bfloat162float(x); - float result = 0.5 * x_f * (1 + erf(x_f / sqrt(2.0f))); - - return __float2bfloat16(result); - } else if constexpr (std::is_same_v) { - float x_f = __half2float(x); - float result = 0.5 * x_f * (1 + erf(x_f / sqrt(2.0f))); - - return __float2half(result); - } else if constexpr (std::is_same_v) { - - return 0.5 * x * (1 + erf(x / sqrt(2.0f))); - } else { - return 0.5 * x * (1 + erf(x / sqrt(2.0))); - } - } -} GeluOp; - -} // namespace op::gelu::cuda - -#endif // __GELU_CUDA_H__ diff --git a/src/infiniop/ops/gelu/kunlun/gelu_kunlun.h b/src/infiniop/ops/gelu/kunlun/gelu_kunlun.h deleted file mode 100644 index 71de9215e..000000000 --- a/src/infiniop/ops/gelu/kunlun/gelu_kunlun.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GELU_KUNLUN_H__ -#define __GELU_KUNLUN_H__ - -#include "../../../elementwise/kunlun/elementwise_kunlun_api.h" - -ELEMENTWISE_DESCRIPTOR(gelu, kunlun) - -#endif diff --git a/src/infiniop/ops/gelu/kunlun/gelu_kunlun.xpu b/src/infiniop/ops/gelu/kunlun/gelu_kunlun.xpu deleted file mode 100644 index 328d3f455..000000000 --- a/src/infiniop/ops/gelu/kunlun/gelu_kunlun.xpu +++ /dev/null @@ -1,56 +0,0 @@ -#include "../../../elementwise/kunlun/elementwise_kunlun.h" -#include "gelu_kunlun.h" -#include "kernel.h" - -namespace op::gelu::kunlun { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - // create Kunlun elementwise descriptor - CREATE_ELEMENTWISE_KUNLUN_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<8, GeluOp, bfloat16_t>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<8, GeluOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<8, GeluOp, float>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::gelu::kunlun diff --git a/src/infiniop/ops/gelu/kunlun/kernel.h b/src/infiniop/ops/gelu/kunlun/kernel.h deleted file mode 100644 index 853205ca1..000000000 --- a/src/infiniop/ops/gelu/kunlun/kernel.h +++ /dev/null @@ -1,32 +0,0 @@ -#ifndef __GELU_KUNLUN_KERNEL_H__ -#define __GELU_KUNLUN_KERNEL_H__ - -namespace op::gelu::kunlun { - -typedef struct GeluOp { -public: - static constexpr size_t num_inputs = 1; - template - inline __device__ T operator()(const T *x) const { - if constexpr (std::is_same_v) { - float x_f = __bfloat162float(x[0]); - float result = 0.5 * x_f * (1 + fast_erf(x_f / sqrt(2.0f))); - - return __float2bfloat16(result); - } else if constexpr (std::is_same_v) { - float x_f = __half2float(x[0]); - float result = 0.5 * x_f * (1 + fast_erf(x_f / sqrt(2.0f))); - - return __float2half(result); - } else if constexpr (std::is_same_v) { - - return 0.5 * x[0] * (1 + fast_erf(x[0] / sqrt(2.0f))); - } else { - return 0.5 * x[0] * (1 + fast_erf(x[0] / sqrt(2.0))); - } - } -} GeluOp; - -} // namespace op::gelu::kunlun - -#endif // __GELU_KUNLUN_H__ diff --git a/src/infiniop/ops/gelu/metax/gelu_meta.maca b/src/infiniop/ops/gelu/metax/gelu_meta.maca deleted file mode 100644 index 3a311530a..000000000 --- a/src/infiniop/ops/gelu/metax/gelu_meta.maca +++ /dev/null @@ -1,60 +0,0 @@ -#include "gelu_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::gelu::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - // create METAX elementwise descriptor - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::GeluOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::GeluOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::GeluOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::GeluOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::gelu::metax diff --git a/src/infiniop/ops/gelu/metax/gelu_metax.h b/src/infiniop/ops/gelu/metax/gelu_metax.h deleted file mode 100644 index 9385b7a27..000000000 --- a/src/infiniop/ops/gelu/metax/gelu_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GELU_METAX_API_H__ -#define __GELU_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(gelu, metax) - -#endif // __GELU_METAX_API_H__ diff --git a/src/infiniop/ops/gelu/nvidia/gelu_nvidia.cu b/src/infiniop/ops/gelu/nvidia/gelu_nvidia.cu deleted file mode 100644 index 4d42cf2df..000000000 --- a/src/infiniop/ops/gelu/nvidia/gelu_nvidia.cu +++ /dev/null @@ -1,59 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "gelu_nvidia.cuh" - -namespace op::gelu::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - // create CUDA elementwise descriptor - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::GeluOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::GeluOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::GeluOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::GeluOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::gelu::nvidia diff --git a/src/infiniop/ops/gelu/nvidia/gelu_nvidia.cuh b/src/infiniop/ops/gelu/nvidia/gelu_nvidia.cuh deleted file mode 100644 index 72dbbd4f0..000000000 --- a/src/infiniop/ops/gelu/nvidia/gelu_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GELU_CUDA_API_H__ -#define __GELU_CUDA_API_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(gelu, nvidia) - -#endif // __GELU_CUDA_API_H__ diff --git a/src/infiniop/ops/gelu/operator.cc b/src/infiniop/ops/gelu/operator.cc deleted file mode 100644 index 69b590d9a..000000000 --- a/src/infiniop/ops/gelu/operator.cc +++ /dev/null @@ -1,209 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/gelu.h" - -#ifdef ENABLE_CPU_API -#include "cpu/gelu_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) -#include "nvidia/gelu_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/gelu_metax.h" -#endif -#ifdef ENABLE_KUNLUN_API -#include "kunlun/gelu_kunlun.h" -#endif -#ifdef ENABLE_CAMBRICON_API -#include "bang/gelu_bang.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateGeluDescriptor( - infiniopHandle_t handle, - infiniopGeluDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::gelu::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output_desc, \ - {input_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_KUNLUN_API - CREATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetGeluWorkspaceSize(infiniopGeluDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_KUNLUN_API - GET(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_CAMBRICON_API - GET(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_ALI_API - GET(INFINI_DEVICE_ALI, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopGelu( - infiniopGeluDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, {input}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_KUNLUN_API - CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyGeluDescriptor(infiniopGeluDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_KUNLUN_API - DELETE(INFINI_DEVICE_KUNLUN, kunlun); -#endif -#ifdef ENABLE_CAMBRICON_API - DELETE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_ALI_API - DELETE(INFINI_DEVICE_ALI, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/gelutanh/cpu/gelutanh_cpu.cc b/src/infiniop/ops/gelutanh/cpu/gelutanh_cpu.cc deleted file mode 100644 index a9d41522c..000000000 --- a/src/infiniop/ops/gelutanh/cpu/gelutanh_cpu.cc +++ /dev/null @@ -1,52 +0,0 @@ -#include "gelutanh_cpu.h" - -namespace op::gelutanh::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &x_desc = input_desc_vec.at(0); - const auto &y_shape = out_desc->shape(); - const auto &x_shape = x_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - CHECK_SAME_SHAPE(y_shape, x_shape); - - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - (void)workspace; - (void)workspace_size; - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::gelutanh::cpu diff --git a/src/infiniop/ops/gelutanh/cpu/gelutanh_cpu.h b/src/infiniop/ops/gelutanh/cpu/gelutanh_cpu.h deleted file mode 100644 index 540ff01a8..000000000 --- a/src/infiniop/ops/gelutanh/cpu/gelutanh_cpu.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef __GELUTANH_CPU_H__ -#define __GELUTANH_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" - -#include - -ELEMENTWISE_DESCRIPTOR(gelutanh, cpu) - -namespace op::gelutanh::cpu { -typedef struct GeluTanhOp { -public: - static constexpr size_t num_inputs = 1; - - template - T operator()(const T &x) const { - // y = x * 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) - constexpr T alpha = static_cast(0.7978845608); // sqrt(2/pi) - constexpr T beta = static_cast(0.044715); - T inner = alpha * (x + beta * x * x * x); - return x * static_cast(0.5) * (static_cast(1) + std::tanh(inner)); - } -} GeluTanhOp; -} // namespace op::gelutanh::cpu - -#endif // __GELUTANH_CPU_H__ diff --git a/src/infiniop/ops/gelutanh/cuda/kernel.cuh b/src/infiniop/ops/gelutanh/cuda/kernel.cuh deleted file mode 100644 index eddfe1b5a..000000000 --- a/src/infiniop/ops/gelutanh/cuda/kernel.cuh +++ /dev/null @@ -1,58 +0,0 @@ -#ifndef __GELUTANH_CUDA_H__ -#define __GELUTANH_CUDA_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" -#include -#include -#include - -namespace op::gelutanh::cuda { - -typedef struct GeluTanhOp { -public: - static constexpr size_t num_inputs = 1; - - // GELU-Tanh constants - // static constexpr float alpha = std::sqrt(2.0 / M_PI); - // static constexpr float beta = 0.044715f; - static constexpr float alpha = 0.7978845608f; // sqrt(2/pi) - static constexpr float beta = 0.044715f; - // f32 tanh helper - __device__ __forceinline__ float tanh_f32_func(float x) const { - return tanhf(x); - } - - template - __device__ __forceinline__ T operator()(const T &x) const { - if constexpr (std::is_same_v) { - // half2 -> float2 - float2 vf = __half22float2(x); - float inner_x0 = alpha * (vf.x + beta * vf.x * vf.x * vf.x); - float inner_x1 = alpha * (vf.y + beta * vf.y * vf.y * vf.y); - float2 vr = make_float2(tanh_f32_func(inner_x0) * 0.5f + 0.5f, - tanh_f32_func(inner_x1) * 0.5f + 0.5f); - return __hmul2(x, __float22half2_rn(vr)); // y = x * 0.5 * (1 + tanh(...)) - } else if constexpr (std::is_same_v) { - float xf = __half2float(x); - float inner = alpha * (xf + beta * xf * xf * xf); - float yf = xf * 0.5f * (1.0f + tanh_f32_func(inner)); - return __float2half_rn(yf); - } else if constexpr (std::is_same_v) { - float xf = __bfloat162float(x); - float inner = alpha * (xf + beta * xf * xf * xf); - float yf = xf * 0.5f * (1.0f + tanh_f32_func(inner)); - return __float2bfloat16(yf); - } else if constexpr (std::is_same_v) { - float inner = alpha * (x + beta * x * x * x); - return x * 0.5f * (1.0f + tanh_f32_func(inner)); - } else { // double - double inner = alpha * (x + beta * x * x * x); - return x * 0.5 * (1.0 + std::tanh(inner)); - } - } - -} GeluTanhOp; - -} // namespace op::gelutanh::cuda - -#endif // __GELUTANH_CUDA_H__ diff --git a/src/infiniop/ops/gelutanh/metax/gelutanh_metax.h b/src/infiniop/ops/gelutanh/metax/gelutanh_metax.h deleted file mode 100644 index da1e678fa..000000000 --- a/src/infiniop/ops/gelutanh/metax/gelutanh_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GELUTANH_METAX_API_H__ -#define __GELUTANH_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(gelutanh, metax) - -#endif // __GELUTANH_METAX_API_H__ diff --git a/src/infiniop/ops/gelutanh/metax/gelutanh_metax.maca b/src/infiniop/ops/gelutanh/metax/gelutanh_metax.maca deleted file mode 100644 index ca151a8e4..000000000 --- a/src/infiniop/ops/gelutanh/metax/gelutanh_metax.maca +++ /dev/null @@ -1,80 +0,0 @@ -#include "../../../elementwise/metax/elementwise_metax.h" -#include "gelutanh_metax.h" - -#include -#include - -namespace op::gelutanh::metax { - -struct GeluTanhOp { - static constexpr size_t num_inputs = 1; - static constexpr float alpha = 0.7978845608f; - static constexpr float beta = 0.044715f; - - template - __device__ __forceinline__ T operator()(const T &x) const { - if constexpr (std::is_same_v) { - float xf = __half2float(x); - float inner = alpha * (xf + beta * xf * xf * xf); - return __float2half(xf * 0.5f * (1.0f + tanhf(inner))); - } else if constexpr (std::is_same_v) { - float xf = __bfloat162float(x); - float inner = alpha * (xf + beta * xf * xf * xf); - return __float2bfloat16(xf * 0.5f * (1.0f + tanhf(inner))); - } else if constexpr (std::is_same_v) { - float inner = alpha * (x + beta * x * x * x); - return x * 0.5f * (1.0f + tanhf(inner)); - } else { - double xd = static_cast(x); - double inner = static_cast(alpha) * (xd + static_cast(beta) * xd * xd * xd); - return static_cast(xd * 0.5 * (1.0 + tanh(inner))); - } - } -}; - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - const auto &x_desc = input_desc_vec.at(0); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - CHECK_SAME_SHAPE(out_desc->shape(), x_desc->shape()); - - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, GeluTanhOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, GeluTanhOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, GeluTanhOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, GeluTanhOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::gelutanh::metax diff --git a/src/infiniop/ops/gelutanh/moore/gelutanh_moore.h b/src/infiniop/ops/gelutanh/moore/gelutanh_moore.h deleted file mode 100644 index 70f47aa1f..000000000 --- a/src/infiniop/ops/gelutanh/moore/gelutanh_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GELUTANH_MOORE_API_H__ -#define __GELUTANH_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(gelutanh, moore) - -#endif // __GELUTANH_MOORE_API_H__ diff --git a/src/infiniop/ops/gelutanh/moore/gelutanh_moore.mu b/src/infiniop/ops/gelutanh/moore/gelutanh_moore.mu deleted file mode 100644 index d291970d5..000000000 --- a/src/infiniop/ops/gelutanh/moore/gelutanh_moore.mu +++ /dev/null @@ -1,85 +0,0 @@ -#include "../../../elementwise/moore/elementwise_moore.h" -#include "gelutanh_moore.h" - -#include -#include - -namespace op::gelutanh::moore { - -struct GeluTanhOp { - static constexpr size_t num_inputs = 1; - static constexpr float alpha = 0.7978845608f; - static constexpr float beta = 0.044715f; - - __device__ __forceinline__ float tanh_approx(float x) const { - return __fsub_rn(__fmul_rn(2.0f, __frcp_rn(__fadd_rn(1.0f, __expf(-2.0f * x)))), 1.0f); - } - - template - __device__ __forceinline__ T operator()(const T &x) const { - if constexpr (std::is_same_v) { - float xf = __half2float(x); - float inner = alpha * (xf + beta * xf * xf * xf); - return __float2half(xf * 0.5f * (1.0f + tanh_approx(inner))); - } else if constexpr (std::is_same_v) { - float xf = __bfloat162float(x); - float inner = alpha * (xf + beta * xf * xf * xf); - return __float2bfloat16_rn(xf * 0.5f * (1.0f + tanh_approx(inner))); - } else if constexpr (std::is_same_v) { - float inner = alpha * (x + beta * x * x * x); - return __fmul_rn(x, __fmul_rn(0.5f, __fadd_rn(1.0f, tanh_approx(inner)))); - } else { - double xd = static_cast(x); - double inner = static_cast(alpha) * (xd + static_cast(beta) * xd * xd * xd); - double tanh_inner = 2.0 / (1.0 + exp(-2.0 * inner)) - 1.0; - return static_cast(xd * 0.5 * (1.0 + tanh_inner)); - } - } -}; - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - const auto &x_desc = input_desc_vec.at(0); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - CHECK_SAME_SHAPE(out_desc->shape(), x_desc->shape()); - - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, GeluTanhOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, GeluTanhOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, GeluTanhOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, GeluTanhOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::gelutanh::moore diff --git a/src/infiniop/ops/gelutanh/nvidia/gelutanh_nvidia.cu b/src/infiniop/ops/gelutanh/nvidia/gelutanh_nvidia.cu deleted file mode 100644 index e51b09606..000000000 --- a/src/infiniop/ops/gelutanh/nvidia/gelutanh_nvidia.cu +++ /dev/null @@ -1,70 +0,0 @@ -#include "../cuda/kernel.cuh" -#include "gelutanh_nvidia.cuh" - -namespace op::gelutanh::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &x_desc = input_desc_vec.at(0); - const auto &y_shape = out_desc->shape(); - const auto &x_shape = x_desc->shape(); - - CHECK_DTYPE(dtype, - INFINI_DTYPE_F16, - INFINI_DTYPE_F32, - INFINI_DTYPE_F64, - INFINI_DTYPE_BF16); - - CHECK_SAME_SHAPE(y_shape, x_shape); - - // create CUDA elementwise descriptor - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::GeluTanhOp, half>( - _info, workspace, output, inputs, stream); - - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::GeluTanhOp, __nv_bfloat16>( - _info, workspace, output, inputs, stream); - - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::GeluTanhOp, float>( - _info, workspace, output, inputs, stream); - - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::GeluTanhOp, double>( - _info, workspace, output, inputs, stream); - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gelutanh::nvidia diff --git a/src/infiniop/ops/gelutanh/nvidia/gelutanh_nvidia.cuh b/src/infiniop/ops/gelutanh/nvidia/gelutanh_nvidia.cuh deleted file mode 100644 index 3155a7af1..000000000 --- a/src/infiniop/ops/gelutanh/nvidia/gelutanh_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GELUTANH_CUDA_API_H__ -#define __GELUTANH_CUDA_API_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(gelutanh, nvidia) - -#endif // __GELUTANH_CUDA_API_H__ diff --git a/src/infiniop/ops/gelutanh/operator.cc b/src/infiniop/ops/gelutanh/operator.cc deleted file mode 100644 index 91a213a24..000000000 --- a/src/infiniop/ops/gelutanh/operator.cc +++ /dev/null @@ -1,173 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/gelutanh.h" - -#ifdef ENABLE_CPU_API -#include "cpu/gelutanh_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/gelutanh_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/gelutanh_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/gelutanh_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateGeluTanhDescriptor( - infiniopHandle_t handle, - infiniopGeluTanhDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::gelutanh::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - {x_desc}) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetGeluTanhWorkspaceSize(infiniopGeluTanhDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET -} - -__INFINI_C infiniStatus_t infiniopGeluTanh( - infiniopGeluTanhDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, {x}, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyGeluTanhDescriptor(infiniopGeluTanhDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/gemm/ascend/gemm_ascend.cc b/src/infiniop/ops/gemm/ascend/gemm_ascend.cc deleted file mode 100644 index e8dfb2f0f..000000000 --- a/src/infiniop/ops/gemm/ascend/gemm_ascend.cc +++ /dev/null @@ -1,148 +0,0 @@ -#include "gemm_ascend.h" -#include "../../../devices/ascend/common_ascend.h" -#include -#include - -#include -#include - -// Custom hash function for alpha beta pair -struct FloatPairHash { - size_t operator()(const std::pair &p) const { - uint64_t combined; - std::memcpy(reinterpret_cast(&combined), &p.first, sizeof(float)); - std::memcpy(reinterpret_cast(&combined) + sizeof(float), &p.second, sizeof(float)); - - return std::hash()(combined); - } -}; - -struct FloatPairEqual { - bool operator()(const std::pair &a, const std::pair &b) const { - return a.first == b.first && a.second == b.second; - } -}; - -namespace op::gemm::ascend { - -struct Descriptor::Opaque { - aclnnTensorDescriptor_t c, a, b; - // cubeMathType - // see doc: - // https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/80RC3alpha002/apiref/appdevgapi/context/aclnnBatchMatMul.md - int8_t mt; - // alpha&beta hashmap - std::unordered_map, aclOpExecutor *, FloatPairHash, FloatPairEqual> lookup; - - ~Opaque() { - delete c; - delete a; - delete b; - for (auto &item : lookup) { - aclDestroyAclOpExecutor(item.second); - } - lookup.clear(); - } -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - auto handle = reinterpret_cast(handle_); - auto dtype = c_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - - auto result = MatmulInfo::create(c_desc, a_desc, b_desc, MatrixLayout::ROW_MAJOR); - CHECK_RESULT(result); - auto info = result.take(); - - auto c = new aclnnTensorDescriptor(toAclDataType(c_desc->dtype()), - {static_cast(info.m), static_cast(info.n)}, - {info.c_matrix.row_stride, info.c_matrix.col_stride}); - auto a = new aclnnTensorDescriptor(toAclDataType(a_desc->dtype()), - {static_cast(info.a_matrix.rows), static_cast(info.a_matrix.cols)}, - {info.a_matrix.row_stride, info.a_matrix.col_stride}); - auto b = new aclnnTensorDescriptor(toAclDataType(b_desc->dtype()), - {static_cast(info.b_matrix.rows), static_cast(info.b_matrix.cols)}, - {info.b_matrix.row_stride, info.b_matrix.col_stride}); - - auto tc = c->tensor, - ta = a->tensor, - tb = b->tensor; - - std::unordered_map, aclOpExecutor *, FloatPairHash, FloatPairEqual> lookup; - aclOpExecutor *executor = nullptr; - size_t workspace_size = 0; - int8_t mt = 1; - CHECK_ACL(aclnnGemmGetWorkspaceSize(ta, tb, tc, 1., 0., 0, 0, tc, mt, &workspace_size, &executor)); - CHECK_ACL(aclSetAclOpExecutorRepeatable(executor)); - lookup[std::make_pair(1.0f, 0.0f)] = executor; - CHECK_ACL(aclnnGemmGetWorkspaceSize(ta, tb, tc, 1., 1., 0, 0, tc, mt, &workspace_size, &executor)); - CHECK_ACL(aclSetAclOpExecutorRepeatable(executor)); - lookup[std::make_pair(1.0f, 1.0f)] = executor; - - *desc_ptr = new Descriptor( - dtype, info, workspace_size, - new Opaque{ - c, - a, - b, - mt, - std::move(lookup)}, - handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspaceSize_, - void *c, - float beta, - const void *a, - const void *b, - float alpha, - void *stream) const { - - auto tc = _opaque->c->tensor, - ta = _opaque->a->tensor, - tb = _opaque->b->tensor; - - size_t workspace_size = _workspace_size; - aclOpExecutor *executor; - auto key = std::make_pair(alpha, beta); - if (_opaque->lookup.find(key) != _opaque->lookup.end()) { - executor = _opaque->lookup[key]; - } else { - CHECK_ACL(aclnnGemmGetWorkspaceSize( - ta, tb, tc, alpha, beta, 0, 0, tc, _opaque->mt, - &workspace_size, &executor)); - CHECK_ACL(aclSetAclOpExecutorRepeatable(executor)); - _opaque->lookup[key] = executor; - } - - if (workspaceSize_ < workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - auto unit = infiniSizeOf(_dtype); - for (size_t i = 0; i < _info.batch; ++i) { - AclSetTensorAddr(executor, 0, ta, ((char *)a) + i * _info.a_matrix.stride * unit); - AclSetTensorAddr(executor, 1, tb, ((char *)b) + i * _info.b_matrix.stride * unit); - AclSetTensorAddr(executor, 2, tc, ((char *)c) + i * _info.c_matrix.stride * unit); - AclSetTensorAddr(executor, 3, tc, ((char *)c) + i * _info.c_matrix.stride * unit); - CHECK_ACL(aclnnGemm(workspace, workspace_size, executor, stream)); - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gemm::ascend diff --git a/src/infiniop/ops/gemm/ascend/gemm_ascend.h b/src/infiniop/ops/gemm/ascend/gemm_ascend.h deleted file mode 100644 index dab9f15fe..000000000 --- a/src/infiniop/ops/gemm/ascend/gemm_ascend.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GEMM_ASCEND_H__ -#define __GEMM_ASCEND_H__ - -#include "../gemm.h" - -DESCRIPTOR(ascend) - -#endif // __GEMM_ASCEND_H__ diff --git a/src/infiniop/ops/gemm/bang/gemm_bang.cc b/src/infiniop/ops/gemm/bang/gemm_bang.cc deleted file mode 100644 index 506bb2357..000000000 --- a/src/infiniop/ops/gemm/bang/gemm_bang.cc +++ /dev/null @@ -1,156 +0,0 @@ -#include "gemm_bang.h" -#include "../../../devices/bang/common_bang.h" -#include - -namespace op::gemm::bang { - -struct Descriptor::Opaque { - cnnlMatMulDescriptor_t op; - cnnlMatMulAlgo_t algo; - cnnlMatMulHeuristicResult_t algoResult; - cnnlTensorDescriptor_t a, b, c; - std::shared_ptr internal; - - ~Opaque() { - cnnlDestroyTensorDescriptor(a); - cnnlDestroyTensorDescriptor(b); - cnnlDestroyTensorDescriptor(c); - cnnlDestroyMatMulDescriptor(op); - cnnlDestroyMatMulAlgo(algo); - cnnlDestroyMatMulHeuristicResult(algoResult); - } -}; - -static infiniStatus_t setMatrixTensorEx( - cnnlTensorDescriptor_t desc, - const BlasMatrix &matrix, infiniDtype_t dtype, - bool trans = false) { - int ndim = matrix.ndim; - int batch = matrix.batch; - int stride = static_cast(matrix.stride); - int rows = matrix.rows; - int cols = matrix.cols; - int row_stride = matrix.row_stride; - int col_stride = matrix.col_stride; - - switch (ndim) { - case 3: { - std::vector dim_size = {batch, rows, cols}; - std::vector dim_stride = {stride, row_stride, col_stride}; - CHECK_BANG(cnnlSetTensorDescriptorEx( - desc, CNNL_LAYOUT_ARRAY, - device::bang::getCnnlDtype(dtype), dim_size.size(), - dim_size.data(), dim_stride.data())); - } break; - case 2: { - std::vector dim_size = {rows, cols}; - std::vector dim_stride = {row_stride, col_stride}; - CHECK_BANG(cnnlSetTensorDescriptorEx( - desc, CNNL_LAYOUT_ARRAY, - device::bang::getCnnlDtype(dtype), dim_size.size(), - dim_size.data(), dim_stride.data())); - } break; - } - return INFINI_STATUS_SUCCESS; -} - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - auto handle = reinterpret_cast(handle_); - auto dtype = c_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - - auto result = MatmulInfo::create(c_desc, a_desc, b_desc, MatrixLayout::ROW_MAJOR); - CHECK_RESULT(result); - auto info = result.take(); - - cnnlTensorDescriptor_t a, b, c; - CHECK_BANG(cnnlCreateTensorDescriptor(&a)); - CHECK_BANG(cnnlCreateTensorDescriptor(&b)); - CHECK_BANG(cnnlCreateTensorDescriptor(&c)); - - CHECK_STATUS(setMatrixTensorEx(a, info.a_matrix, a_desc->dtype())); - CHECK_STATUS(setMatrixTensorEx(b, info.b_matrix, b_desc->dtype())); - CHECK_STATUS(setMatrixTensorEx(c, info.c_matrix, c_desc->dtype())); - - cnnlMatMulDescriptor_t op; - cnnlMatMulAlgo_t algo; - cnnlMatMulHeuristicResult_t algoResult; - CHECK_BANG(cnnlCreateMatMulDescriptor(&op)); - CHECK_BANG(cnnlCreateMatMulAlgo(&algo)); - CHECK_BANG(cnnlCreateMatMulHeuristicResult(&algoResult)); - int32_t use_stride = true; - CHECK_BANG(cnnlSetMatMulDescAttr( - op, - CNNL_MATMUL_USE_STRIDE, - &use_stride, - sizeof(int32_t))); - int count = 0; - - CHECK_STATUS( - handle->internal()->useCnnl( - (cnrtQueue_t) nullptr, - [&](cnnlHandle_t _handle) { - CHECK_BANG( - cnnlGetBatchMatMulExAlgoHeuristic( - _handle, - op, a, b, c, - NULL, 1, &algoResult, &count)); - return INFINI_STATUS_SUCCESS; - })); - - size_t workspace_size; - CHECK_BANG(cnnlGetBatchMatMulExHeuristicResult(algoResult, algo, &workspace_size)); - - *desc_ptr = new Descriptor( - dtype, info, workspace_size, - new Opaque{ - op, algo, algoResult, a, b, c, handle->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *c, - float beta, - const void *a, - const void *b, - float alpha, - void *stream) const { - - if (_info.is_transed) { - std::swap(a, b); - } - CHECK_STATUS(_opaque->internal->useCnnl( - (cnrtQueue_t)stream, - [&](cnnlHandle_t handle) { - CHECK_BANG(cnnlBatchMatMulEx( - handle, - _opaque->op, - _opaque->algo, - &alpha, - _opaque->a, a, - _opaque->b, b, - &beta, - _opaque->c, c, - workspace, - workspace_size)); - return INFINI_STATUS_SUCCESS; - })); - cnrtQueueSync((cnrtQueue_t)stream); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gemm::bang diff --git a/src/infiniop/ops/gemm/bang/gemm_bang.h b/src/infiniop/ops/gemm/bang/gemm_bang.h deleted file mode 100644 index f3df849cc..000000000 --- a/src/infiniop/ops/gemm/bang/gemm_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GEMM_BANG_H__ -#define __GEMM_BANG_H__ - -#include "../gemm.h" - -DESCRIPTOR(bang) - -#endif // __GEMM_BANG_H__ diff --git a/src/infiniop/ops/gemm/cpu/gemm_cpu.cc b/src/infiniop/ops/gemm/cpu/gemm_cpu.cc deleted file mode 100644 index 6f7a2e3e0..000000000 --- a/src/infiniop/ops/gemm/cpu/gemm_cpu.cc +++ /dev/null @@ -1,104 +0,0 @@ -#include "gemm_cpu.h" -#include "../../../devices/cpu/common_cpu.h" - -namespace op::gemm::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - auto handle = reinterpret_cast(handle_); - auto dtype = c_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = MatmulInfo::create(c_desc, a_desc, b_desc, MatrixLayout::COL_MAJOR); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - dtype, result.take(), 0, - nullptr, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -void calculate( - const MatmulInfo &info, - void *c, - float beta, - const void *a, - const void *b, - float alpha) { - if (info.is_transed) { - std::swap(a, b); - } - -#pragma omp parallel for - for (ptrdiff_t index = 0; index < ptrdiff_t(info.batch * info.m * info.n); ++index) { - size_t ind = index; - size_t n_ = ind % info.n; - ind /= info.n; - size_t m_ = ind % info.m; - ind /= info.m; - size_t i = ind; - auto c_ = reinterpret_cast(c) + i * info.c_matrix.stride + m_ * info.c_matrix.row_stride + n_ * info.c_matrix.col_stride; - float sum = 0; - for (int k_ = 0; k_ < static_cast(info.k); ++k_) { - auto a_ = reinterpret_cast(a) + i * info.a_matrix.stride + m_ * info.a_matrix.row_stride + k_ * info.a_matrix.col_stride; - auto b_ = reinterpret_cast(b) + i * info.b_matrix.stride + n_ * info.b_matrix.col_stride + k_ * info.b_matrix.row_stride; - if constexpr (std::is_same::value || std::is_same::value) { - sum += utils::cast(*a_) * utils::cast(*b_); - } else { - sum += *a_ * (*b_); - } - } - if constexpr (std::is_same::value || std::is_same::value) { - if (beta == 0) { - *c_ = utils::cast(alpha * sum); - } else { - *c_ = utils::cast(beta * utils::cast(*c_) + alpha * sum); - } - } else { - if (beta == 0) { - *c_ = alpha * sum; - } else { - *c_ = beta * (*c_) + alpha * sum; - } - } - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *c, - float beta, - const void *a, - const void *b, - float alpha, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - cpu::calculate(_info, c, beta, a, b, alpha); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_BF16: - cpu::calculate(_info, c, beta, a, b, alpha); - return INFINI_STATUS_SUCCESS; - - case INFINI_DTYPE_F32: - cpu::calculate(_info, c, beta, a, b, alpha); - return INFINI_STATUS_SUCCESS; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::gemm::cpu diff --git a/src/infiniop/ops/gemm/cpu/gemm_cpu.h b/src/infiniop/ops/gemm/cpu/gemm_cpu.h deleted file mode 100644 index 6c8ced7b9..000000000 --- a/src/infiniop/ops/gemm/cpu/gemm_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GEMM_CPU_H__ -#define __GEMM_CPU_H__ - -#include "../gemm.h" - -DESCRIPTOR(cpu) - -#endif // __GEMM_CPU_H__ diff --git a/src/infiniop/ops/gemm/gemm.h b/src/infiniop/ops/gemm/gemm.h deleted file mode 100644 index 5f2503fb0..000000000 --- a/src/infiniop/ops/gemm/gemm.h +++ /dev/null @@ -1,93 +0,0 @@ -#ifndef __GEMM_H__ -#define __GEMM_H__ - -#include "../../operator.h" -#include "info.h" - -/** - * # 关于 `DESCRIPTOR(NAMESPACE)` 和 `struct Opaque;` 的说明 - * - * > - Data: 2025/02/20 - * > - Author: - * - * `DESCRIPTOR(NAMESPACE)` 宏定义了一个 `Descriptor` 类， - * 该类继承自 `InfiniopDescriptor`，并在内部定义了： - * - * - 描述任何硬件上的矩阵乘算子必要的信息 - * - 参数数据类型； - * - 矩阵乘运算信息（BMNK）； - * - 工作空间大小； - * - 特定硬件独有的描述信息保存在 `struct Opaque;` 中； - * - 私有的构造函数； - * - 公共接口 - * - 析构函数； - * - 静态的工厂函数； - * - 矩阵乘计算函数； - * - * 这个宏必须写成一个宏，因为静态成员函数是不可继承的，但每个不同硬件上有不同的实现。 - * 使用宏声明的不同硬件的矩阵乘描述符是不相关类型（并且它们也不必相关）。 - * - * 这些描述符的声明对 operator.cc 可见，而 operator.cc 文件编译到 infiniop 库的硬件无关部分中， - * 因此，描述符声明不可以出现硬件相关的类型。 - * - * 为了隐藏必须伴随算子描述符保存的硬件相关类型（例如，Ascend 上的 `executor`）， - * 这里使用了 `Opaque` 声明。 - * - * Opaque 表示“不透明的”，即其内部结构面型头文件的使用者隐藏。 - * 这种设计来自 C++ 的一种设计模式，其经典形式称为 [PImpl](https://zh.cppreference.com/w/cpp/language/pimpl) 模式。 - * 由于此处仅需要结构体的指针，因此结构体可以仅声明不定义，从而隐藏了结构体的成员。 - * 通过将结构体定义为 private，保证了外部不可能操作这个指针， - * 因此这个设计可以由类型自由控制其成员的生命周期， - * 同时不必将成员的形式暴露给头文件的使用者， - * 这是一种安全的封装。 - * - * 这个宏仅适用于矩阵乘，但这种模式很容易复制到其他算子，以简化和规范算子的声明。 - */ - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::gemm::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - infiniDtype_t _dtype; \ - MatmulInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - infiniDtype_t dtype, \ - MatmulInfo info, \ - size_t workspace_size_, \ - Opaque *opaque, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _dtype(dtype), \ - _info(info), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t c_desc, \ - infiniopTensorDescriptor_t a_desc, \ - infiniopTensorDescriptor_t b_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *c, \ - float beta, \ - const void *a, \ - const void *b, \ - float alpha, \ - void *stream) const; \ - }; \ - } - -#endif // __GEMM_H__ diff --git a/src/infiniop/ops/gemm/info.h b/src/infiniop/ops/gemm/info.h deleted file mode 100644 index 1f17d0525..000000000 --- a/src/infiniop/ops/gemm/info.h +++ /dev/null @@ -1,132 +0,0 @@ -#ifndef __GEMM_INFO_H__ -#define __GEMM_INFO_H__ - -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" -#include - -namespace op::gemm { - -struct BlasMatrix { - size_t ndim; - size_t batch; - ptrdiff_t stride; - size_t rows; - size_t cols; - ptrdiff_t row_stride; - ptrdiff_t col_stride; - - static utils::Result create(infiniopTensorDescriptor_t layout) { - BlasMatrix ans; - - if (layout->ndim() == 2) { - ans.ndim = 2; - ans.batch = 1; - ans.stride = 0; - ans.rows = layout->dim(0); - ans.cols = layout->dim(1); - ans.row_stride = layout->stride(0); - ans.col_stride = layout->stride(1); - } else if (layout->ndim() == 3) { - ans.ndim = 3; - ans.batch = layout->dim(0); - ans.stride = ans.batch == 1 ? 0 : layout->stride(0); - ans.rows = layout->dim(1); - ans.cols = layout->dim(2); - ans.row_stride = layout->stride(1); - ans.col_stride = layout->stride(2); - } else { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (ans.row_stride != 1 && ans.col_stride != 1) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - return utils::Result(ans); - } - - bool match_batch(size_t _batch) const { - return batch == _batch || batch == 1; - } - - void transpose() { - std::swap(rows, cols); - std::swap(row_stride, col_stride); - } - - ptrdiff_t ld() const { - return row_stride == 1 ? col_stride : row_stride; - } -}; - -enum class MatrixLayout : char { - COL_MAJOR, - ROW_MAJOR, -}; - -class MatmulInfo { - MatmulInfo() = default; - -public: - BlasMatrix a_matrix; - BlasMatrix b_matrix; - BlasMatrix c_matrix; - - size_t m, n, k, batch; - bool is_transed; - - static utils::Result create( - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - MatrixLayout layout) { - - auto a_matrix = BlasMatrix::create(a_desc); - CHECK_RESULT(a_matrix); - - auto b_matrix = BlasMatrix::create(b_desc); - CHECK_RESULT(b_matrix); - - auto c_matrix = BlasMatrix::create(c_desc); - CHECK_RESULT(c_matrix); - - if (c_matrix->rows != a_matrix->rows || c_matrix->cols != b_matrix->cols || a_matrix->cols != b_matrix->rows) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - auto batch = c_matrix->batch; - if (!a_matrix->match_batch(batch) || !b_matrix->match_batch(batch)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - auto is_transed = false; - if ((layout == MatrixLayout::COL_MAJOR && c_matrix->col_stride == 1) - || (layout == MatrixLayout::ROW_MAJOR && c_matrix->row_stride == 1)) { - c_matrix->transpose(); - b_matrix->transpose(); - a_matrix->transpose(); - std::swap(a_matrix, b_matrix); - is_transed = true; - } - - auto m = c_matrix->rows; - auto n = c_matrix->cols; - auto k = a_matrix->cols; - - return utils::Result(MatmulInfo{ - a_matrix.take(), - b_matrix.take(), - c_matrix.take(), - m, - n, - k, - batch, - is_transed}); - } -}; - -} // namespace op::gemm - -#endif // __GEMM_INFO_H__ diff --git a/src/infiniop/ops/gemm/kunlun/gemm_kunlun.cc b/src/infiniop/ops/gemm/kunlun/gemm_kunlun.cc deleted file mode 100644 index 0c195bd23..000000000 --- a/src/infiniop/ops/gemm/kunlun/gemm_kunlun.cc +++ /dev/null @@ -1,329 +0,0 @@ -#include "gemm_kunlun.h" -#include "../../../devices/kunlun/kunlun_common.h" -#include "../../../devices/kunlun/kunlun_xblas.h" -#include "gemm_kunlun_cast.h" -#include -#include - -namespace op::gemm::kunlun { - -typedef device::kunlun::blas::Handle::Internal HandleInternal; - -static bool isLargeBf16SkinnyGemm(const MatmulInfo &info) { - return info.is_transed && info.n == 1 && info.m > 2048; -} - -static bool useBf16Lt(const MatmulInfo &info) { - return !isLargeBf16SkinnyGemm(info); -} - -static size_t packedMatrixSize(size_t rows, size_t cols, size_t batch) { - return rows * cols * batch; -} - -static size_t bf16WorkspaceSize(const MatmulInfo &info) { - return (packedMatrixSize(info.m, info.k, info.batch) - + packedMatrixSize(info.k, info.n, info.batch) - + packedMatrixSize(info.m, info.n, info.batch)) - * sizeof(float); -} - -struct Descriptor::Opaque { - std::shared_ptr internal; - cublasLtHandle_t lt_handle = nullptr; - cublasLtMatmulDesc_t lt_desc = nullptr; - cublasLtMatrixLayout_t a_layout = nullptr; - cublasLtMatrixLayout_t b_layout = nullptr; - cublasLtMatrixLayout_t c_layout = nullptr; - void destroyLtDescriptors(); - bool createBf16LtDescriptors(const MatmulInfo &info); -}; - -static size_t ltLayoutRows(const BlasMatrix &matrix) { - return matrix.row_stride == 1 ? matrix.rows : matrix.cols; -} - -static size_t ltLayoutCols(const BlasMatrix &matrix) { - return matrix.row_stride == 1 ? matrix.cols : matrix.rows; -} - -static bool setLtLayoutBatch(cublasLtMatrixLayout_t layout, const BlasMatrix &matrix, size_t batch) { - int32_t batch_count = static_cast(batch); - int64_t stride = static_cast(matrix.stride); - return cublasLtMatrixLayoutSetAttribute( - layout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, - &batch_count, sizeof(batch_count)) - == CUBLAS_STATUS_SUCCESS - && cublasLtMatrixLayoutSetAttribute( - layout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, - &stride, sizeof(stride)) - == CUBLAS_STATUS_SUCCESS; -} - -void Descriptor::Opaque::destroyLtDescriptors() { - if (a_layout) { - cublasLtMatrixLayoutDestroy(a_layout); - } - if (b_layout) { - cublasLtMatrixLayoutDestroy(b_layout); - } - if (c_layout) { - cublasLtMatrixLayoutDestroy(c_layout); - } - if (lt_desc) { - cublasLtMatmulDescDestroy(lt_desc); - } - if (lt_handle) { - cublasLtDestroy(lt_handle); - } -} - -bool Descriptor::Opaque::createBf16LtDescriptors(const MatmulInfo &info) { - auto op_a = info.a_matrix.row_stride == 1 ? CUBLAS_OP_N : CUBLAS_OP_T; - auto op_b = info.b_matrix.row_stride == 1 ? CUBLAS_OP_N : CUBLAS_OP_T; - - if (cublasLtCreate(<_handle) != CUBLAS_STATUS_SUCCESS) { - return false; - } - if (cublasLtMatmulDescCreate(<_desc, CUBLAS_COMPUTE_32F, CUDA_R_32F) != CUBLAS_STATUS_SUCCESS) { - return false; - } - if (cublasLtMatmulDescSetAttribute( - lt_desc, CUBLASLT_MATMUL_DESC_TRANSA, - &op_a, sizeof(op_a)) - != CUBLAS_STATUS_SUCCESS) { - return false; - } - if (cublasLtMatmulDescSetAttribute( - lt_desc, CUBLASLT_MATMUL_DESC_TRANSB, - &op_b, sizeof(op_b)) - != CUBLAS_STATUS_SUCCESS) { - return false; - } - - if (cublasLtMatrixLayoutCreate( - &a_layout, CUDA_R_16BF, - ltLayoutRows(info.a_matrix), ltLayoutCols(info.a_matrix), - info.a_matrix.ld()) - != CUBLAS_STATUS_SUCCESS) { - return false; - } - if (cublasLtMatrixLayoutCreate( - &b_layout, CUDA_R_16BF, - ltLayoutRows(info.b_matrix), ltLayoutCols(info.b_matrix), - info.b_matrix.ld()) - != CUBLAS_STATUS_SUCCESS) { - return false; - } - if (cublasLtMatrixLayoutCreate( - &c_layout, CUDA_R_16BF, - ltLayoutRows(info.c_matrix), ltLayoutCols(info.c_matrix), - info.c_matrix.ld()) - != CUBLAS_STATUS_SUCCESS) { - return false; - } - - return setLtLayoutBatch(a_layout, info.a_matrix, info.batch) - && setLtLayoutBatch(b_layout, info.b_matrix, info.batch) - && setLtLayoutBatch(c_layout, info.c_matrix, info.batch); -} - -Descriptor::~Descriptor() { - if (_opaque) { - _opaque->destroyLtDescriptors(); - } - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - auto handle = reinterpret_cast(handle_); - auto dtype = c_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = MatmulInfo::create(c_desc, a_desc, b_desc, MatrixLayout::COL_MAJOR); - CHECK_RESULT(result); - - auto info = result.take(); - auto workspace_size = dtype == INFINI_DTYPE_BF16 ? bf16WorkspaceSize(info) : 0; - - auto opaque = new Opaque{handle->internal()}; - bool use_bf16_lt = dtype == INFINI_DTYPE_BF16 && useBf16Lt(info); - if (use_bf16_lt && !opaque->createBf16LtDescriptors(info)) { - opaque->destroyLtDescriptors(); - opaque->lt_handle = nullptr; - opaque->lt_desc = nullptr; - opaque->a_layout = nullptr; - opaque->b_layout = nullptr; - opaque->c_layout = nullptr; - } - - *desc_ptr = new Descriptor( - dtype, info, workspace_size, - opaque, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *c, - float beta, - const void *a, - const void *b, - float alpha, - void *stream) const { - - cudaDataType a_type, b_type, c_type; - cublasComputeType_t compute_type; - - switch (_dtype) { - case INFINI_DTYPE_F16: - a_type = b_type = c_type = CUDA_R_16F; - compute_type = CUBLAS_COMPUTE_32F; - break; - case INFINI_DTYPE_BF16: - a_type = b_type = c_type = CUDA_R_16BF; - compute_type = CUBLAS_COMPUTE_32F; - break; - case INFINI_DTYPE_F32: - a_type = b_type = c_type = CUDA_R_32F; - compute_type = CUBLAS_COMPUTE_32F_FAST_TF32; - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - if (_info.is_transed) { - std::swap(a, b); - } - - auto op_a = _info.a_matrix.row_stride == 1 ? CUBLAS_OP_N : CUBLAS_OP_T; - auto op_b = _info.b_matrix.row_stride == 1 ? CUBLAS_OP_N : CUBLAS_OP_T; - - if (_dtype == INFINI_DTYPE_BF16) { - auto use_bf16_lt = useBf16Lt(_info); - if (use_bf16_lt && _opaque->lt_handle && _opaque->lt_desc && _opaque->a_layout && _opaque->b_layout && _opaque->c_layout) { - auto lt_status = cublasLtMatmul( - _opaque->lt_handle, - _opaque->lt_desc, - &alpha, - a, - _opaque->a_layout, - b, - _opaque->b_layout, - &beta, - c, - _opaque->c_layout, - c, - _opaque->c_layout, - nullptr, - workspace, - workspace_size, - (cudaStream_t)stream); - if (lt_status == CUBLAS_STATUS_SUCCESS) { - xpu_wait(stream); - return INFINI_STATUS_SUCCESS; - } - } - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - auto temp_element_size = sizeof(float); - auto workspace_bytes = reinterpret_cast(workspace); - auto a_tmp = workspace_bytes; - auto b_tmp = a_tmp + packedMatrixSize(_info.m, _info.k, _info.batch) * temp_element_size; - auto c_tmp = b_tmp + packedMatrixSize(_info.k, _info.n, _info.batch) * temp_element_size; - auto temp_type = CUDA_R_32F; - auto a_stride = static_cast(packedMatrixSize(_info.m, _info.k, 1)); - auto b_stride = static_cast(packedMatrixSize(_info.k, _info.n, 1)); - auto c_stride = static_cast(packedMatrixSize(_info.m, _info.n, 1)); - - CHECK_STATUS(castBf16ToF32(a, a_tmp, _info.a_matrix, _info.batch, (kunlunStream_t)stream)); - CHECK_STATUS(castBf16ToF32(b, b_tmp, _info.b_matrix, _info.batch, (kunlunStream_t)stream)); - if (beta == 0.0f) { - CHECK_STATUS(zeroPackedBuffer(c_tmp, packedMatrixSize(_info.m, _info.n, _info.batch) * temp_element_size, (kunlunStream_t)stream)); - } else { - CHECK_STATUS(castBf16ToF32(c, c_tmp, _info.c_matrix, _info.batch, (kunlunStream_t)stream)); - } - - CHECK_STATUS(_opaque->internal->useCublas( - (cudaStream_t)stream, - [&](cublasHandle_t handle) { - CHECK_CUBLAS( - cublasGemmStridedBatchedEx( - handle, - CUBLAS_OP_N, - CUBLAS_OP_N, - static_cast(_info.m), - static_cast(_info.n), - static_cast(_info.k), - &alpha, - a_tmp, - temp_type, - static_cast(_info.m), - a_stride, - b_tmp, - temp_type, - static_cast(_info.k), - b_stride, - &beta, - c_tmp, - temp_type, - static_cast(_info.m), - c_stride, - static_cast(_info.batch), - CUBLAS_COMPUTE_32F, - CUBLAS_GEMM_DEFAULT)); - return INFINI_STATUS_SUCCESS; - })); - - CHECK_STATUS(castF32ToBf16(c_tmp, c, _info.c_matrix, _info.batch, (kunlunStream_t)stream)); - xpu_wait(stream); - return INFINI_STATUS_SUCCESS; - } - - CHECK_STATUS(_opaque->internal->useCublas( - (cudaStream_t)stream, - [&](cublasHandle_t handle) { - CHECK_CUBLAS( - cublasGemmStridedBatchedEx( - handle, - op_a, - op_b, - static_cast(_info.m), - static_cast(_info.n), - static_cast(_info.k), - &alpha, - a, - a_type, - static_cast(_info.a_matrix.ld()), - _info.a_matrix.stride, - b, - b_type, - static_cast(_info.b_matrix.ld()), - _info.b_matrix.stride, - &beta, - c, - c_type, - static_cast(_info.c_matrix.ld()), - _info.c_matrix.stride, - static_cast(_info.batch), - compute_type, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); - return INFINI_STATUS_SUCCESS; - })); - - xpu_wait(stream); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gemm::kunlun diff --git a/src/infiniop/ops/gemm/kunlun/gemm_kunlun.h b/src/infiniop/ops/gemm/kunlun/gemm_kunlun.h deleted file mode 100644 index e82deecf3..000000000 --- a/src/infiniop/ops/gemm/kunlun/gemm_kunlun.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GEMM_KUNLUN_H__ -#define __GEMM_KUNLUN_H__ - -#include "../gemm.h" - -DESCRIPTOR(kunlun) - -#endif // __GEMM_KUNLUN_H__ diff --git a/src/infiniop/ops/gemm/kunlun/gemm_kunlun_cast.h b/src/infiniop/ops/gemm/kunlun/gemm_kunlun_cast.h deleted file mode 100644 index ab7e39832..000000000 --- a/src/infiniop/ops/gemm/kunlun/gemm_kunlun_cast.h +++ /dev/null @@ -1,75 +0,0 @@ -#ifndef __GEMM_KUNLUN_CAST_H__ -#define __GEMM_KUNLUN_CAST_H__ - -#include "../../../devices/kunlun/kunlun_common.h" -#include "../info.h" - -namespace op::gemm::kunlun { - -infiniStatus_t castBf16ToF16( - const void *src, - void *dst, - const BlasMatrix &matrix, - size_t batch, - kunlunStream_t stream); - -infiniStatus_t castF16ToBf16( - const void *src, - void *dst, - const BlasMatrix &matrix, - size_t batch, - kunlunStream_t stream); - -infiniStatus_t castBf16ToF32( - const void *src, - void *dst, - const BlasMatrix &matrix, - size_t batch, - kunlunStream_t stream); - -infiniStatus_t castF32ToBf16( - const void *src, - void *dst, - const BlasMatrix &matrix, - size_t batch, - kunlunStream_t stream); - -infiniStatus_t castBf16ToF32Chunk( - const void *src, - void *dst, - const BlasMatrix &matrix, - size_t batch, - size_t row_offset, - size_t rows, - kunlunStream_t stream); - -infiniStatus_t castF32ToBf16Chunk( - const void *src, - void *dst, - const BlasMatrix &matrix, - size_t batch, - size_t row_offset, - size_t rows, - kunlunStream_t stream); - -infiniStatus_t zeroPackedBuffer( - void *dst, - size_t bytes, - kunlunStream_t stream); - -infiniStatus_t bf16SkinnyGemv( - const void *a, - const void *b, - void *c, - const BlasMatrix &a_matrix, - const BlasMatrix &b_matrix, - const BlasMatrix &c_matrix, - size_t m, - size_t k, - float alpha, - float beta, - kunlunStream_t stream); - -} // namespace op::gemm::kunlun - -#endif // __GEMM_KUNLUN_CAST_H__ diff --git a/src/infiniop/ops/gemm/kunlun/gemm_kunlun_cast.xpu b/src/infiniop/ops/gemm/kunlun/gemm_kunlun_cast.xpu deleted file mode 100644 index ad876e352..000000000 --- a/src/infiniop/ops/gemm/kunlun/gemm_kunlun_cast.xpu +++ /dev/null @@ -1,435 +0,0 @@ -#include "gemm_kunlun_cast.h" -#include "../../../devices/kunlun/kunlun_kernel_common.h" - -namespace op::gemm::kunlun { - -template -__global__ void packMatrixKernel( - const Tin *src, - Tout *dst, - long long batch, - long long rows, - long long cols, - long long stride, - long long row_stride, - long long col_stride) { - - int cid = core_id(); - int ncores = core_num(); - if (cid >= ncores) { - return; - } - int tid = ncores * cluster_id() + cid; - int nthreads = ncores * cluster_num(); - - long long matrix_elements = rows * cols; - long long total = batch * matrix_elements; - - for (long long idx = tid; idx < total; idx += nthreads) { - long long batch_id = idx / matrix_elements; - long long rem = idx - batch_id * matrix_elements; - long long row = rem % rows; - long long col = rem / rows; - long long src_batch_offset = stride == 0 ? 0 : batch_id * stride; - long long src_offset = src_batch_offset + row * row_stride + col * col_stride; - - Tin in; - Tout out; - GM2LM_ASYNC(src + src_offset, &in, sizeof(Tin)); - mfence(); - if constexpr (xpu_std::is_same::value && xpu_std::is_same::value) { - out = __float2half(__bfloat162float(in)); - } else if constexpr (xpu_std::is_same::value && xpu_std::is_same::value) { - out = __bfloat162float(in); - } else if constexpr (xpu_std::is_same::value && xpu_std::is_same::value) { - out = __float2bfloat16(__half2float(in)); - } else { - out = __float2bfloat16(in); - } - LM2GM_ASYNC(&out, dst + idx, sizeof(Tout)); - mfence(); - } - sync_cluster(); -} - -template -__global__ void unpackMatrixKernel( - const Tin *src, - Tout *dst, - long long batch, - long long rows, - long long cols, - long long stride, - long long row_stride, - long long col_stride) { - - int cid = core_id(); - int ncores = core_num(); - if (cid >= ncores) { - return; - } - int tid = ncores * cluster_id() + cid; - int nthreads = ncores * cluster_num(); - - long long matrix_elements = rows * cols; - long long total = batch * matrix_elements; - - for (long long idx = tid; idx < total; idx += nthreads) { - long long batch_id = idx / matrix_elements; - long long rem = idx - batch_id * matrix_elements; - long long row = rem % rows; - long long col = rem / rows; - long long dst_batch_offset = stride == 0 ? 0 : batch_id * stride; - long long dst_offset = dst_batch_offset + row * row_stride + col * col_stride; - - Tin in; - Tout out; - GM2LM_ASYNC(src + idx, &in, sizeof(Tin)); - mfence(); - if constexpr (xpu_std::is_same::value && xpu_std::is_same::value) { - out = __float2bfloat16(__half2float(in)); - } else if constexpr (xpu_std::is_same::value && xpu_std::is_same::value) { - out = __float2bfloat16(in); - } else if constexpr (xpu_std::is_same::value && xpu_std::is_same::value) { - out = __float2half(__bfloat162float(in)); - } else { - out = __bfloat162float(in); - } - LM2GM_ASYNC(&out, dst + dst_offset, sizeof(Tout)); - mfence(); - } - sync_cluster(); -} - -infiniStatus_t castBf16ToF16( - const void *src, - void *dst, - const BlasMatrix &matrix, - size_t batch, - kunlunStream_t stream) { - - packMatrixKernel - <<<12, 64, stream>>>( - reinterpret_cast(src), - reinterpret_cast(dst), - static_cast(batch), - static_cast(matrix.rows), - static_cast(matrix.cols), - static_cast(matrix.stride), - static_cast(matrix.row_stride), - static_cast(matrix.col_stride)); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t castF16ToBf16( - const void *src, - void *dst, - const BlasMatrix &matrix, - size_t batch, - kunlunStream_t stream) { - - unpackMatrixKernel - <<<12, 64, stream>>>( - reinterpret_cast(src), - reinterpret_cast(dst), - static_cast(batch), - static_cast(matrix.rows), - static_cast(matrix.cols), - static_cast(matrix.stride), - static_cast(matrix.row_stride), - static_cast(matrix.col_stride)); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t castBf16ToF32( - const void *src, - void *dst, - const BlasMatrix &matrix, - size_t batch, - kunlunStream_t stream) { - - packMatrixKernel - <<<12, 64, stream>>>( - reinterpret_cast(src), - reinterpret_cast(dst), - static_cast(batch), - static_cast(matrix.rows), - static_cast(matrix.cols), - static_cast(matrix.stride), - static_cast(matrix.row_stride), - static_cast(matrix.col_stride)); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t castF32ToBf16( - const void *src, - void *dst, - const BlasMatrix &matrix, - size_t batch, - kunlunStream_t stream) { - - unpackMatrixKernel - <<<12, 64, stream>>>( - reinterpret_cast(src), - reinterpret_cast(dst), - static_cast(batch), - static_cast(matrix.rows), - static_cast(matrix.cols), - static_cast(matrix.stride), - static_cast(matrix.row_stride), - static_cast(matrix.col_stride)); - return INFINI_STATUS_SUCCESS; -} - -template -__global__ void packMatrixRowsKernel( - const Tin *src, - Tout *dst, - long long batch, - long long rows, - long long cols, - long long stride, - long long row_stride, - long long col_stride, - long long row_offset) { - - int cid = core_id(); - int ncores = core_num(); - if (cid >= ncores) { - return; - } - int tid = ncores * cluster_id() + cid; - int nthreads = ncores * cluster_num(); - - long long matrix_elements = rows * cols; - long long total = batch * matrix_elements; - - for (long long idx = tid; idx < total; idx += nthreads) { - long long batch_id = idx / matrix_elements; - long long rem = idx - batch_id * matrix_elements; - long long row = rem % rows; - long long col = rem / rows; - long long src_batch_offset = stride == 0 ? 0 : batch_id * stride; - long long src_offset = src_batch_offset + (row_offset + row) * row_stride + col * col_stride; - - Tin in; - Tout out; - GM2LM_ASYNC(src + src_offset, &in, sizeof(Tin)); - mfence(); - if constexpr (xpu_std::is_same::value && xpu_std::is_same::value) { - out = __bfloat162float(in); - } else { - out = __float2bfloat16(in); - } - LM2GM_ASYNC(&out, dst + idx, sizeof(Tout)); - mfence(); - } - sync_cluster(); -} - -template -__global__ void unpackMatrixRowsKernel( - const Tin *src, - Tout *dst, - long long batch, - long long rows, - long long cols, - long long stride, - long long row_stride, - long long col_stride, - long long row_offset) { - - int cid = core_id(); - int ncores = core_num(); - if (cid >= ncores) { - return; - } - int tid = ncores * cluster_id() + cid; - int nthreads = ncores * cluster_num(); - - long long matrix_elements = rows * cols; - long long total = batch * matrix_elements; - - for (long long idx = tid; idx < total; idx += nthreads) { - long long batch_id = idx / matrix_elements; - long long rem = idx - batch_id * matrix_elements; - long long row = rem % rows; - long long col = rem / rows; - long long dst_batch_offset = stride == 0 ? 0 : batch_id * stride; - long long dst_offset = dst_batch_offset + (row_offset + row) * row_stride + col * col_stride; - - Tin in; - Tout out; - GM2LM_ASYNC(src + idx, &in, sizeof(Tin)); - mfence(); - if constexpr (xpu_std::is_same::value && xpu_std::is_same::value) { - out = __float2bfloat16(in); - } else { - out = __bfloat162float(in); - } - LM2GM_ASYNC(&out, dst + dst_offset, sizeof(Tout)); - mfence(); - } - sync_cluster(); -} - -infiniStatus_t castBf16ToF32Chunk( - const void *src, - void *dst, - const BlasMatrix &matrix, - size_t batch, - size_t row_offset, - size_t rows, - kunlunStream_t stream) { - - packMatrixRowsKernel - <<<12, 64, stream>>>( - reinterpret_cast(src), - reinterpret_cast(dst), - static_cast(batch), - static_cast(rows), - static_cast(matrix.cols), - static_cast(matrix.stride), - static_cast(matrix.row_stride), - static_cast(matrix.col_stride), - static_cast(row_offset)); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t castF32ToBf16Chunk( - const void *src, - void *dst, - const BlasMatrix &matrix, - size_t batch, - size_t row_offset, - size_t rows, - kunlunStream_t stream) { - - unpackMatrixRowsKernel - <<<12, 64, stream>>>( - reinterpret_cast(src), - reinterpret_cast(dst), - static_cast(batch), - static_cast(rows), - static_cast(matrix.cols), - static_cast(matrix.stride), - static_cast(matrix.row_stride), - static_cast(matrix.col_stride), - static_cast(row_offset)); - return INFINI_STATUS_SUCCESS; -} - - -__global__ void zeroPackedBufferKernel( - __global_ptr__ char *dst, - long long bytes) { - - int cid = core_id(); - int ncores = core_num(); - if (cid >= ncores) { - return; - } - int tid = ncores * cluster_id() + cid; - int nthreads = ncores * cluster_num(); - - for (long long idx = tid; idx < bytes; idx += nthreads) { - __local__ char zero = 0; - LM2GM_ASYNC(&zero, dst + idx, sizeof(char)); - } - mfence(); - sync_cluster(); -} - -infiniStatus_t zeroPackedBuffer( - void *dst, - size_t bytes, - kunlunStream_t stream) { - - zeroPackedBufferKernel<<<12, 64, stream>>>( - reinterpret_cast<__global_ptr__ char *>(dst), - static_cast(bytes)); - return INFINI_STATUS_SUCCESS; -} - -__global__ void bf16SkinnyGemvKernel( - __global_ptr__ const bfloat16_t *a, - __global_ptr__ const bfloat16_t *b, - __global_ptr__ bfloat16_t *c, - long long m, - long long k, - long long a_row_stride, - long long a_col_stride, - long long b_row_stride, - long long b_col_stride, - long long c_row_stride, - long long c_col_stride, - float alpha, - float beta) { - - int cid = core_id(); - int ncores = core_num(); - if (cid >= ncores) { - return; - } - int tid = ncores * cluster_id() + cid; - int nthreads = ncores * cluster_num(); - - for (long long row = tid; row < m; row += nthreads) { - float acc = 0.0f; - for (long long col = 0; col < k; ++col) { - __local__ bfloat16_t av; - __local__ bfloat16_t bv; - GM2LM_ASYNC(a + row * a_row_stride + col * a_col_stride, &av, sizeof(bfloat16_t)); - mfence(); - GM2LM_ASYNC(b + col * b_row_stride, &bv, sizeof(bfloat16_t)); - mfence(); - acc += __bfloat162float(av) * __bfloat162float(bv); - } - - auto c_ptr = c + row * c_row_stride; - if (beta != 0.0f) { - __local__ bfloat16_t cv; - GM2LM_ASYNC(c_ptr, &cv, sizeof(bfloat16_t)); - mfence(); - acc = alpha * acc + beta * __bfloat162float(cv); - } else { - acc *= alpha; - } - __local__ bfloat16_t out = __float2bfloat16(acc); - LM2GM_ASYNC(&out, c_ptr, sizeof(bfloat16_t)); - mfence(); - } - sync_cluster(); -} - -infiniStatus_t bf16SkinnyGemv( - const void *a, - const void *b, - void *c, - const BlasMatrix &a_matrix, - const BlasMatrix &b_matrix, - const BlasMatrix &c_matrix, - size_t m, - size_t k, - float alpha, - float beta, - kunlunStream_t stream) { - - bf16SkinnyGemvKernel<<<12, 64, stream>>>( - reinterpret_cast(a), - reinterpret_cast(b), - reinterpret_cast(c), - static_cast(m), - static_cast(k), - static_cast(a_matrix.row_stride), - static_cast(a_matrix.col_stride), - static_cast(b_matrix.row_stride), - static_cast(b_matrix.col_stride), - static_cast(c_matrix.row_stride), - static_cast(c_matrix.col_stride), - alpha, - beta); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gemm::kunlun diff --git a/src/infiniop/ops/gemm/metax/gemm_metax.cc b/src/infiniop/ops/gemm/metax/gemm_metax.cc deleted file mode 100644 index 9d45099dc..000000000 --- a/src/infiniop/ops/gemm/metax/gemm_metax.cc +++ /dev/null @@ -1,107 +0,0 @@ -#include "gemm_metax.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" - -namespace op::gemm::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - auto handle = reinterpret_cast(handle_); - auto dtype = c_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = MatmulInfo::create(c_desc, a_desc, b_desc, MatrixLayout::COL_MAJOR); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - dtype, result.take(), 0, - new Opaque{handle->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *c, - float beta, - const void *a, - const void *b, - float alpha, - void *stream) const { - - hpccDataType a_type, b_type, c_type; - hcblasComputeType_t compute_type; - - switch (_dtype) { - case INFINI_DTYPE_F16: - a_type = b_type = c_type = HPCC_R_16F; - compute_type = HCBLAS_COMPUTE_32F; - break; - case INFINI_DTYPE_BF16: - a_type = b_type = c_type = HPCC_R_16BF; - compute_type = HCBLAS_COMPUTE_32F; - break; - case INFINI_DTYPE_F32: - a_type = b_type = c_type = HPCC_R_32F; - compute_type = HCBLAS_COMPUTE_32F_FAST_TF32; - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - if (_info.is_transed) { - std::swap(a, b); - } - - auto op_a = _info.a_matrix.row_stride == 1 ? HCBLAS_OP_N : HCBLAS_OP_T; - auto op_b = _info.b_matrix.row_stride == 1 ? HCBLAS_OP_N : HCBLAS_OP_T; - - CHECK_STATUS(_opaque->internal->useMcblas( - (hcStream_t)stream, - [&](hcblasHandle_t handle) { - CHECK_MCBLAS( - hcblasGemmStridedBatchedEx( - handle, - op_a, - op_b, - static_cast(_info.m), - static_cast(_info.n), - static_cast(_info.k), - &alpha, - a, - a_type, - static_cast(_info.a_matrix.ld()), - _info.a_matrix.stride, - b, - b_type, - static_cast(_info.b_matrix.ld()), - _info.b_matrix.stride, - &beta, - c, - c_type, - static_cast(_info.c_matrix.ld()), - _info.c_matrix.stride, - static_cast(_info.batch), - compute_type, - HCBLAS_GEMM_DEFAULT_TENSOR_OP)); - return INFINI_STATUS_SUCCESS; - })); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gemm::metax diff --git a/src/infiniop/ops/gemm/metax/gemm_metax.h b/src/infiniop/ops/gemm/metax/gemm_metax.h deleted file mode 100644 index 42f651c77..000000000 --- a/src/infiniop/ops/gemm/metax/gemm_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GEMM_METAX_H__ -#define __GEMM_METAX_H__ - -#include "../gemm.h" - -DESCRIPTOR(metax) - -#endif // __GEMM_METAX_H__ diff --git a/src/infiniop/ops/gemm/moore/gemm_moore.h b/src/infiniop/ops/gemm/moore/gemm_moore.h deleted file mode 100644 index 03f399fe6..000000000 --- a/src/infiniop/ops/gemm/moore/gemm_moore.h +++ /dev/null @@ -1,104 +0,0 @@ -#ifndef __GEMM_MOORE_H__ -#define __GEMM_MOORE_H__ - -#include "mublas/gemm_mublas.h" -#include "mudnn/gemm_mudnn.h" - -namespace op::gemm::moore { - -// Descriptor class for GEMM operations on Moore devices. -// This class acts as a wrapper to select either mublas or mudnn backend. -// It encapsulates the backend-specific Descriptor implementation and provides -// a unified interface for workspace query and GEMM calculation. -class Descriptor final : public InfiniopDescriptor { -public: - // Destructor: deletes the backend-specific descriptor. - ~Descriptor() { - if (_backend == Backend::MUBLAS) { - delete reinterpret_cast(_impl); - } else { - delete reinterpret_cast(_impl); - } - } - - // Returns the required workspace size for the GEMM operation. - size_t workspaceSize() const { - if (_backend == Backend::MUBLAS) { - return reinterpret_cast(_impl)->workspaceSize(); - } else { - return reinterpret_cast(_impl)->workspaceSize(); - } - } - - // Static factory method to create a Descriptor instance. - // This method chooses the backend (mublas or mudnn) and constructs - // the corresponding implementation internally. - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - auto desc = new Descriptor(handle->device, handle->device_id); - - // Backend selection strategy: - // Currently defaulting to MUDNN. - // Can be modified to choose based on environment variables or runtime parameters. - desc->_backend = Backend::MUDNN; - - if (desc->_backend == Backend::MUBLAS) { - mublas::Descriptor *impl; - auto status = mublas::Descriptor::create(handle, &impl, c_desc, a_desc, b_desc); - if (status != INFINI_STATUS_SUCCESS) { - delete desc; - return status; - } - desc->_impl = impl; - } else { - mudnn::Descriptor *impl; - auto status = mudnn::Descriptor::create(handle, &impl, c_desc, a_desc, b_desc); - if (status != INFINI_STATUS_SUCCESS) { - delete desc; - return status; - } - desc->_impl = impl; - } - - *desc_ptr = desc; - return INFINI_STATUS_SUCCESS; - } - - // Unified GEMM calculation interface. - // Calls the corresponding backend's calculate function internally. - infiniStatus_t calculate( - void *workspace, size_t workspace_size, - void *c, float beta, - const void *a, const void *b, - float alpha, - void *stream) const { - if (_backend == Backend::MUBLAS) { - return reinterpret_cast(_impl) - ->calculate(workspace, workspace_size, c, beta, a, b, alpha, stream); - } else { - return reinterpret_cast(_impl) - ->calculate(workspace, workspace_size, c, beta, a, b, alpha, stream); - } - } - -private: - // Private constructor: ensures users cannot directly instantiate Descriptor. - // Instances must be created via the static create() factory method. - Descriptor(infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, _impl(nullptr) {} - - // Enum to indicate which backend is being used internally. - enum class Backend { MUBLAS, - MUDNN }; - - Backend _backend; // Currently selected MUBLAS/MUDNN backend - void *_impl; // Pointer to backend-specific descriptor (mublas::Descriptor* or mudnn::Descriptor*) -}; - -} // namespace op::gemm::moore - -#endif // __GEMM_MOORE_H__ diff --git a/src/infiniop/ops/gemm/moore/mublas/gemm_mublas.h b/src/infiniop/ops/gemm/moore/mublas/gemm_mublas.h deleted file mode 100644 index f09c1f909..000000000 --- a/src/infiniop/ops/gemm/moore/mublas/gemm_mublas.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GEMM_MUBLAS_H__ -#define __GEMM_MUBLAS_H__ - -#include "../../gemm.h" - -DESCRIPTOR(mublas) - -#endif // __GEMM_MUBLAS_H__ diff --git a/src/infiniop/ops/gemm/moore/mublas/gemm_mublas.mu b/src/infiniop/ops/gemm/moore/mublas/gemm_mublas.mu deleted file mode 100644 index e52b72acb..000000000 --- a/src/infiniop/ops/gemm/moore/mublas/gemm_mublas.mu +++ /dev/null @@ -1,125 +0,0 @@ -#include "../../../../devices/moore/moore_common.h" -#include "../../../../devices/moore/moore_handle.h" -#include "gemm_mublas.h" - -namespace op::gemm::mublas { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - auto handle = reinterpret_cast(handle_); - auto dtype = c_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = MatmulInfo::create(c_desc, a_desc, b_desc, MatrixLayout::COL_MAJOR); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - dtype, result.take(), 0, - new Opaque{handle->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *c, - float beta, - const void *a, - const void *b, - float alpha, - void *stream) const { - - musaDataType a_type, b_type, c_type; - mublasComputeType_t compute_type; - - // MUSA's GEMM operations require that the scalar values alpha and beta have the same data type as the matrices. - // This ensures correct computation during the muBLAS GEMM operation. - // Declare half-precision variables to handle F16 types. - half alpha_h, beta_h; - - // Initialize generic void pointers for alpha and beta. - // They point to the original float values - // It will be used directly when the GEMM operation is performed with F32 data. - const void *p_alpha = α - const void *p_beta = β - - switch (_dtype) { - case INFINI_DTYPE_F16: - a_type = b_type = c_type = MUSA_R_16F; - compute_type = MUBLAS_COMPUTE_16F; - - // Convert alpha/beta to half-precision and update the pointers. - alpha_h = __float2half(alpha); - beta_h = __float2half(beta); - p_alpha = &alpha_h; - p_beta = &beta_h; - - break; - case INFINI_DTYPE_BF16: - a_type = b_type = c_type = MUSA_R_16BF; - compute_type = MUBLAS_COMPUTE_32F; - break; - case INFINI_DTYPE_F32: - a_type = b_type = c_type = MUSA_R_32F; - compute_type = MUBLAS_COMPUTE_32F_FAST_TF32; - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - if (_info.is_transed) { - std::swap(a, b); - } - - auto op_a = _info.a_matrix.row_stride == 1 ? MUBLAS_OP_N : MUBLAS_OP_T; - auto op_b = _info.b_matrix.row_stride == 1 ? MUBLAS_OP_N : MUBLAS_OP_T; - - CHECK_STATUS(_opaque->internal->useMublas( - (musaStream_t)stream, - [&](mublasHandle_t handle) { - CHECK_MUBLAS( - mublasGemmStridedBatchedEx( - handle, - op_a, - op_b, - static_cast(_info.m), - static_cast(_info.n), - static_cast(_info.k), - p_alpha, - a, - a_type, - static_cast(_info.a_matrix.ld()), - _info.a_matrix.stride, - b, - b_type, - static_cast(_info.b_matrix.ld()), - _info.b_matrix.stride, - p_beta, - c, - c_type, - static_cast(_info.c_matrix.ld()), - _info.c_matrix.stride, - static_cast(_info.batch), - compute_type, - MUBLAS_GEMM_DEFAULT)); - return INFINI_STATUS_SUCCESS; - })); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gemm::mublas diff --git a/src/infiniop/ops/gemm/moore/mudnn/gemm_mudnn.h b/src/infiniop/ops/gemm/moore/mudnn/gemm_mudnn.h deleted file mode 100644 index 30a2c707d..000000000 --- a/src/infiniop/ops/gemm/moore/mudnn/gemm_mudnn.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GEMM_MUDNN_H__ -#define __GEMM_MUDNN_H__ - -#include "../../gemm.h" - -DESCRIPTOR(mudnn) - -#endif // __GEMM_MUDNN_H__ diff --git a/src/infiniop/ops/gemm/moore/mudnn/gemm_mudnn.mu b/src/infiniop/ops/gemm/moore/mudnn/gemm_mudnn.mu deleted file mode 100644 index 12acd7fc5..000000000 --- a/src/infiniop/ops/gemm/moore/mudnn/gemm_mudnn.mu +++ /dev/null @@ -1,198 +0,0 @@ -#include "../../../../devices/moore/moore_common.h" -#include "../../../../devices/moore/moore_handle.h" -#include "gemm_mudnn.h" - -#include - -namespace op::gemm::mudnn { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - auto handle = reinterpret_cast(handle_); - auto dtype = c_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = MatmulInfo::create(c_desc, a_desc, b_desc, MatrixLayout::ROW_MAJOR); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - dtype, result.take(), 0, - new Opaque{handle->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t calculate( - const MatmulInfo &info, - std::shared_ptr &_internal, - void *c, - float beta, - const void *a, - const void *b, - float alpha, - void *stream) -{ - // 0. For muDNN development, refer to the official documentation and the following headers: - // - /usr/local/musa/include/mudnn_base.h - // - /usr/local/musa/include/mudnn_math.h - // - /usr/local/musa/include/mudnn.h - - // 1. Create BatchMatMul operator - auto matmul_operator = std::make_unique<::musa::dnn::BatchMatMul>(); - matmul_operator->SetComputeMode(::musa::dnn::BatchMatMul::ComputeMode::TENSOR); - - // 2. Use _internal->useMudnn to manage muDNN handle - return _internal->useMudnn((musaStream_t)stream, [&](::musa::dnn::Handle &mudnn_handle) -> infiniStatus_t { - - // 3. Create BatchMatMul Tensor - ::musa::dnn::Tensor out, left, right; - - if constexpr (std::is_same::value) { - out.SetType(::musa::dnn::Tensor::Type::HALF); - left.SetType(::musa::dnn::Tensor::Type::HALF); - right.SetType(::musa::dnn::Tensor::Type::HALF); - } - else if constexpr (std::is_same::value){ - out.SetType(::musa::dnn::Tensor::Type::BFLOAT16); - left.SetType(::musa::dnn::Tensor::Type::BFLOAT16); - right.SetType(::musa::dnn::Tensor::Type::BFLOAT16); - } - else{ - out.SetType(::musa::dnn::Tensor::Type::FLOAT); - left.SetType(::musa::dnn::Tensor::Type::FLOAT); - right.SetType(::musa::dnn::Tensor::Type::FLOAT); - } - - // 4. Bind BatchMatMul Tensor addr - out.SetAddr(c); - left.SetAddr(a); - right.SetAddr(b); - - // 5. Config Tensor left - std::array a_dims_array; - std::array a_stride_array; - if (info.a_matrix.col_stride != 1) { - a_dims_array = { static_cast(info.batch), - static_cast(info.k), - static_cast(info.m) }; - } else { - a_dims_array = { static_cast(info.batch), - static_cast(info.m), - static_cast(info.k) }; - } - a_stride_array = { static_cast(info.a_matrix.stride), - static_cast(info.a_matrix.ld()), - 1 }; - left.SetNdInfo(static_cast(a_dims_array.size()), a_dims_array.data(), a_stride_array.data()); - - // 6. Config Tensor right - std::array b_dims_array; - std::array b_stride_array; - if (info.b_matrix.col_stride != 1) { - b_dims_array = { static_cast(info.batch), - static_cast(info.n), - static_cast(info.k) }; - } else { - b_dims_array = { static_cast(info.batch), - static_cast(info.k), - static_cast(info.n) }; - } - b_stride_array = { static_cast(info.b_matrix.stride), - static_cast(info.b_matrix.ld()), - 1 }; - right.SetNdInfo(static_cast(b_dims_array.size()), b_dims_array.data(), b_stride_array.data()); - - // 7. Confit Tensor out, muDNN BatchMatMul output only support row-major tensor - std::array c_dims_array = { static_cast(info.batch), - static_cast(info.m), - static_cast(info.n) }; - std::array c_stride_array = { static_cast(info.c_matrix.stride), - static_cast(info.c_matrix.ld()), - 1 }; - out.SetNdInfo(static_cast(c_dims_array.size()), c_dims_array.data(), c_stride_array.data()); - - // 8. Workspace Memory Handler - ::musa::dnn::MemoryMaintainer maintainer = [](size_t size) -> ::musa::dnn::MemoryHandler { - void* ptr = nullptr; - musaMalloc(&ptr, size); - return ::musa::dnn::MemoryHandler(ptr, [](void* p) { if(p) musaFree(p); }); - }; - - // 9. Tensor left and Tensor right transpose config - if (info.a_matrix.col_stride == 1 && info.b_matrix.col_stride != 1) - matmul_operator->SetTranspose(false, true); - else if (info.a_matrix.col_stride != 1 && info.b_matrix.col_stride == 1) - matmul_operator->SetTranspose(true, false); - else if (info.a_matrix.col_stride != 1 && info.b_matrix.col_stride != 1) - matmul_operator->SetTranspose(true, true); - else - matmul_operator->SetTranspose(false, false); - - // 10. BatchMatMul workspace config - size_t workspace_size_in_bytes = 0; - matmul_operator->GetWorkspaceSize(mudnn_handle, workspace_size_in_bytes, out, left, right); - - // 11. Alpha Beta Gamma - matmul_operator->SetAlpha(static_cast(alpha)); - matmul_operator->SetBeta(static_cast(beta)); - matmul_operator->SetGamma(0.0); - - // 12. Run - matmul_operator->Run( - mudnn_handle, - out, - left, - right, - static_cast(info.batch), - static_cast(info.m), - static_cast(info.n), - static_cast(info.k), - static_cast(info.a_matrix.ld()), - static_cast(info.b_matrix.ld()), - static_cast(info.c_matrix.ld()), - static_cast(info.a_matrix.stride), - static_cast(info.b_matrix.stride), - static_cast(info.c_matrix.stride), - maintainer - ); - - return INFINI_STATUS_SUCCESS; - }); -} - - -infiniStatus_t Descriptor::calculate(void *workspace, - size_t workspace_size, - void *c, - float beta, - const void *a, - const void *b, - float alpha, - void *stream) const { - switch (_dtype) { - case INFINI_DTYPE_F16: - return mudnn::calculate(_info, _opaque->internal, c, beta, a, b, alpha, stream); - case INFINI_DTYPE_F32: - return mudnn::calculate(_info,_opaque->internal, c, beta, a, b, alpha, stream); - case INFINI_DTYPE_BF16: - return mudnn::calculate<__mt_bfloat16>(_info,_opaque->internal, c, beta, a, b, alpha, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::gemm::mudnn diff --git a/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cu b/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cu deleted file mode 100644 index 0e0c65f2b..000000000 --- a/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cu +++ /dev/null @@ -1,122 +0,0 @@ -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "gemm_nvidia.cuh" - -namespace op::gemm::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - auto handle = reinterpret_cast(handle_); - auto dtype = c_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = MatmulInfo::create(c_desc, a_desc, b_desc, MatrixLayout::COL_MAJOR); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - dtype, result.take(), 0, - new Opaque{handle->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *c, - float beta, - const void *a, - const void *b, - float alpha, - void *stream) const { - - cudaDataType a_type, b_type, c_type; -#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) - cudaDataType compute_type; -#else - cublasComputeType_t compute_type; -#endif - - switch (_dtype) { - case INFINI_DTYPE_F16: - a_type = b_type = c_type = CUDA_R_16F; -#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) - compute_type = CUDA_R_32F; -#else - compute_type = CUBLAS_COMPUTE_32F; -#endif - break; - case INFINI_DTYPE_BF16: - a_type = b_type = c_type = CUDA_R_16BF; -#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) - compute_type = CUDA_R_32F; -#else - compute_type = CUBLAS_COMPUTE_32F; -#endif - break; - case INFINI_DTYPE_F32: - a_type = b_type = c_type = CUDA_R_32F; -#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) - compute_type = CUDA_R_32F; -#else - compute_type = CUBLAS_COMPUTE_32F_FAST_TF32; -#endif - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - if (_info.is_transed) { - std::swap(a, b); - } - - auto op_a = _info.a_matrix.row_stride == 1 ? CUBLAS_OP_N : CUBLAS_OP_T; - auto op_b = _info.b_matrix.row_stride == 1 ? CUBLAS_OP_N : CUBLAS_OP_T; - - CHECK_STATUS(_opaque->internal->useCublas( - (cudaStream_t)stream, - [&](cublasHandle_t handle) { - CHECK_CUBLAS( - cublasGemmStridedBatchedEx( - handle, - op_a, - op_b, - static_cast(_info.m), - static_cast(_info.n), - static_cast(_info.k), - &alpha, - a, - a_type, - static_cast(_info.a_matrix.ld()), - _info.a_matrix.stride, - b, - b_type, - static_cast(_info.b_matrix.ld()), - _info.b_matrix.stride, - &beta, - c, - c_type, - static_cast(_info.c_matrix.ld()), - _info.c_matrix.stride, - static_cast(_info.batch), - compute_type, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); - return INFINI_STATUS_SUCCESS; - })); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gemm::nvidia diff --git a/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cuh b/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cuh deleted file mode 100644 index 056069703..000000000 --- a/src/infiniop/ops/gemm/nvidia/gemm_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GEMM_CUDA_CUH__ -#define __GEMM_CUDA_CUH__ - -#include "../gemm.h" - -DESCRIPTOR(nvidia) - -#endif // __GEMM_CUDA_CUH__ diff --git a/src/infiniop/ops/gemm/operator.cc b/src/infiniop/ops/gemm/operator.cc deleted file mode 100644 index 81d9cb066..000000000 --- a/src/infiniop/ops/gemm/operator.cc +++ /dev/null @@ -1,252 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/gemm.h" - -#ifdef ENABLE_CPU_API -#include "cpu/gemm_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) -#include "nvidia/gemm_nvidia.cuh" -#endif -#ifdef ENABLE_CAMBRICON_API -#include "bang/gemm_bang.h" -#endif -#ifdef ENABLE_ASCEND_API -#include "ascend/gemm_ascend.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/gemm_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/gemm_moore.h" -#endif -#ifdef ENABLE_KUNLUN_API -#include "kunlun/gemm_kunlun.h" -#endif -#ifdef ENABLE_QY_API -#include "qy/gemm_qy.cuh" -#endif -__INFINI_C infiniStatus_t infiniopCreateGemmDescriptor( - infiniopHandle_t handle, - infiniopGemmDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::gemm::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - c_desc, \ - a_desc, \ - b_desc) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, qy); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_CAMBRICON_API - CREATE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_ASCEND_API - CREATE(INFINI_DEVICE_ASCEND, ascend); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - -#ifdef ENABLE_KUNLUN_API - CREATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t -infiniopGetGemmWorkspaceSize( - infiniopGemmDescriptor_t desc, - size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - GET(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, qy); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_CAMBRICON_API - GET(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_ASCEND_API - GET(INFINI_DEVICE_ASCEND, ascend); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_KUNLUN_API - GET(INFINI_DEVICE_KUNLUN, kunlun); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET -} - -__INFINI_C infiniStatus_t infiniopGemm( - infiniopGemmDescriptor_t desc, - void *workspace, size_t workspace_size, - void *c, - const void *a, - const void *b, - float alpha, - float beta, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, \ - c, beta, \ - a, b, alpha, \ - stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, qy); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_CAMBRICON_API - CALCULATE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_ASCEND_API - CALCULATE(INFINI_DEVICE_ASCEND, ascend); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_KUNLUN_API - CALCULATE(INFINI_DEVICE_KUNLUN, kunlun); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyGemmDescriptor(infiniopGemmDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - DELETE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, qy); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_CAMBRICON_API - DELETE(INFINI_DEVICE_CAMBRICON, bang); -#endif -#ifdef ENABLE_ASCEND_API - DELETE(INFINI_DEVICE_ASCEND, ascend); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_KUNLUN_API - DELETE(INFINI_DEVICE_KUNLUN, kunlun); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/gemm/qy/gemm_qy.cu b/src/infiniop/ops/gemm/qy/gemm_qy.cu deleted file mode 100644 index 34241a34f..000000000 --- a/src/infiniop/ops/gemm/qy/gemm_qy.cu +++ /dev/null @@ -1,134 +0,0 @@ -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "gemm_qy.cuh" - -namespace op::gemm::qy { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t c_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - auto handle = reinterpret_cast(handle_); - auto dtype = c_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - - auto result = MatmulInfo::create(c_desc, a_desc, b_desc, MatrixLayout::COL_MAJOR); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - dtype, result.take(), 0, - new Opaque{handle->internal()}, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *c, - float beta, - const void *a, - const void *b, - float alpha, - void *stream) const { - - cudaDataType a_type, b_type, c_type; -#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) - cudaDataType compute_type; -#else - cublasComputeType_t compute_type; -#endif - - switch (_dtype) { - case INFINI_DTYPE_F16: - a_type = b_type = c_type = CUDA_R_16F; -#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) - compute_type = CUDA_R_32F; -#else - compute_type = CUBLAS_COMPUTE_32F; -#endif - break; - case INFINI_DTYPE_BF16: - a_type = b_type = c_type = CUDA_R_16BF; -#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) - compute_type = CUDA_R_32F; -#else - compute_type = CUBLAS_COMPUTE_32F; -#endif - break; - case INFINI_DTYPE_F32: - a_type = b_type = c_type = CUDA_R_32F; -#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) - compute_type = CUDA_R_32F; -#else - compute_type = CUBLAS_COMPUTE_32F_FAST_TF32; -#endif - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - if (_info.is_transed) { - std::swap(a, b); - } - - auto op_a = _info.a_matrix.row_stride == 1 ? CUBLAS_OP_N : CUBLAS_OP_T; - auto op_b = _info.b_matrix.row_stride == 1 ? CUBLAS_OP_N : CUBLAS_OP_T; - - CHECK_STATUS(_opaque->internal->useCublas( - (cudaStream_t)stream, - [&](cublasHandle_t handle) { - // 1. 获取batch数和各矩阵的元素大小（字节） - const int batch = static_cast(_info.batch); - const size_t a_elem_size = 2; // getCudaDataTypeSize(a_type); - const size_t b_elem_size = 2; // getCudaDataTypeSize(b_type); - const size_t c_elem_size = 2; // getCudaDataTypeSize(c_type); - - // 2. 循环处理每个batch - for (int i = 0; i < batch; ++i) { - // 计算当前batch的A/B/C指针（stride是元素步长，转字节步长） - const void *a_batch = (const char *)a + i * _info.a_matrix.stride * a_elem_size; - const void *b_batch = (const char *)b + i * _info.b_matrix.stride * b_elem_size; - void *c_batch = (char *)c + i * _info.c_matrix.stride * c_elem_size; - - // 3. 调用单batch的cublasGemmEx（参数与原接口完全对齐） - CHECK_CUBLAS( - cublasGemmEx( - handle, - op_a, // 原op_a - op_b, // 原op_b - static_cast(_info.m), // 原m - static_cast(_info.n), // 原n - static_cast(_info.k), // 原k - &alpha, // 原alpha - a_batch, // 当前batch的A指针 - a_type, // 原a_type - static_cast(_info.a_matrix.ld()), // 原a的ld - b_batch, // 当前batch的B指针 - b_type, // 原b_type - static_cast(_info.b_matrix.ld()), // 原b的ld - &beta, // 原beta - c_batch, // 当前batch的C指针 - c_type, // 原c_type - static_cast(_info.c_matrix.ld()), // 原c的ld - compute_type, // 原compute_type - CUBLAS_GEMM_DEFAULT_TENSOR_OP // 原算法选择 - )); - } - return INFINI_STATUS_SUCCESS; - })); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gemm::qy diff --git a/src/infiniop/ops/gemm/qy/gemm_qy.cuh b/src/infiniop/ops/gemm/qy/gemm_qy.cuh deleted file mode 100644 index 0b218bac9..000000000 --- a/src/infiniop/ops/gemm/qy/gemm_qy.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GEMM_QY_CUH__ -#define __GEMM_QY_CUH__ - -#include "../gemm.h" - -DESCRIPTOR(qy) - -#endif // __GEMM_QY_CUH__ diff --git a/src/infiniop/ops/gptq_marlin_gemm/dlpack/dlpack.h b/src/infiniop/ops/gptq_marlin_gemm/dlpack/dlpack.h deleted file mode 100644 index 3f4ed201a..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/dlpack/dlpack.h +++ /dev/null @@ -1,647 +0,0 @@ -/*! - * Copyright (c) 2017 - by Contributors - * \file dlpack.h - * \brief The common header of DLPack. - */ -#ifndef DLPACK_DLPACK_H_ -#define DLPACK_DLPACK_H_ - -/** - * \brief Compatibility with C++ - */ -#ifdef __cplusplus -#define DLPACK_EXTERN_C extern "C" -#else -#define DLPACK_EXTERN_C -#endif - -/*! \brief The current major version of dlpack */ -#define DLPACK_MAJOR_VERSION 1 - -/*! \brief The current minor version of dlpack */ -#define DLPACK_MINOR_VERSION 3 - -/*! \brief DLPACK_DLL prefix for windows */ -#ifdef _WIN32 -#ifdef DLPACK_EXPORTS -#define DLPACK_DLL __declspec(dllexport) -#else -#define DLPACK_DLL __declspec(dllimport) -#endif -#else -#define DLPACK_DLL -#endif - -#include -#include - -#ifdef __cplusplus -extern "C" { -#endif - -/*! - * \brief The DLPack version. - * - * A change in major version indicates that we have changed the - * data layout of the ABI - DLManagedTensorVersioned. - * - * A change in minor version indicates that we have added new - * code, such as a new device type, but the ABI is kept the same. - * - * If an obtained DLPack tensor has a major version that disagrees - * with the version number specified in this header file - * (i.e. major != DLPACK_MAJOR_VERSION), the consumer must call the deleter - * (and it is safe to do so). It is not safe to access any other fields - * as the memory layout will have changed. - * - * In the case of a minor version mismatch, the tensor can be safely used as - * long as the consumer knows how to interpret all fields. Minor version - * updates indicate the addition of enumeration values. - */ -typedef struct { - /*! \brief DLPack major version. */ - uint32_t major; - /*! \brief DLPack minor version. */ - uint32_t minor; -} DLPackVersion; - -/*! - * \brief The device type in DLDevice. - */ -#ifdef __cplusplus -typedef enum : int32_t { -#else -typedef enum { -#endif - /*! \brief CPU device */ - kDLCPU = 1, - /*! \brief CUDA GPU device */ - kDLCUDA = 2, - /*! - * \brief Pinned CUDA CPU memory by cudaMallocHost - */ - kDLCUDAHost = 3, - /*! \brief OpenCL devices. */ - kDLOpenCL = 4, - /*! \brief Vulkan buffer for next generation graphics. */ - kDLVulkan = 7, - /*! \brief Metal for Apple GPU. */ - kDLMetal = 8, - /*! \brief Verilog simulator buffer */ - kDLVPI = 9, - /*! \brief ROCm GPUs for AMD GPUs */ - kDLROCM = 10, - /*! - * \brief Pinned ROCm CPU memory allocated by hipMallocHost - */ - kDLROCMHost = 11, - /*! - * \brief Reserved extension device type, - * used for quickly test extension device - * The semantics can differ depending on the implementation. - */ - kDLExtDev = 12, - /*! - * \brief CUDA managed/unified memory allocated by cudaMallocManaged - */ - kDLCUDAManaged = 13, - /*! - * \brief Unified shared memory allocated on a oneAPI non-partititioned - * device. Call to oneAPI runtime is required to determine the device - * type, the USM allocation type and the sycl context it is bound to. - * - */ - kDLOneAPI = 14, - /*! \brief GPU support for next generation WebGPU standard. */ - kDLWebGPU = 15, - /*! \brief Qualcomm Hexagon DSP */ - kDLHexagon = 16, - /*! \brief Microsoft MAIA devices */ - kDLMAIA = 17, - /*! \brief AWS Trainium */ - kDLTrn = 18, -} DLDeviceType; - -/*! - * \brief A Device for Tensor and operator. - */ -typedef struct { - /*! \brief The device type used in the device. */ - DLDeviceType device_type; - /*! - * \brief The device index. - * For vanilla CPU memory, pinned memory, or managed memory, this is set to 0. - */ - int32_t device_id; -} DLDevice; - -/*! - * \brief The type code options DLDataType. - */ -typedef enum { - /*! \brief signed integer */ - kDLInt = 0U, - /*! \brief unsigned integer */ - kDLUInt = 1U, - /*! \brief IEEE floating point */ - kDLFloat = 2U, - /*! - * \brief Opaque handle type, reserved for testing purposes. - * Frameworks need to agree on the handle data type for the exchange to be well-defined. - */ - kDLOpaqueHandle = 3U, - /*! \brief bfloat16 */ - kDLBfloat = 4U, - /*! - * \brief complex number - * (C/C++/Python layout: compact struct per complex number) - */ - kDLComplex = 5U, - /*! \brief boolean */ - kDLBool = 6U, - /*! \brief FP8 data types */ - kDLFloat8_e3m4 = 7U, - kDLFloat8_e4m3 = 8U, - kDLFloat8_e4m3b11fnuz = 9U, - kDLFloat8_e4m3fn = 10U, - kDLFloat8_e4m3fnuz = 11U, - kDLFloat8_e5m2 = 12U, - kDLFloat8_e5m2fnuz = 13U, - kDLFloat8_e8m0fnu = 14U, - /*! \brief FP6 data types - * Setting bits != 6 is currently unspecified, and the producer must ensure it is set - * while the consumer must stop importing if the value is unexpected. - */ - kDLFloat6_e2m3fn = 15U, - kDLFloat6_e3m2fn = 16U, - /*! \brief FP4 data types - * Setting bits != 4 is currently unspecified, and the producer must ensure it is set - * while the consumer must stop importing if the value is unexpected. - */ - kDLFloat4_e2m1fn = 17U, -} DLDataTypeCode; - -/*! - * \brief The data type the tensor can hold. The data type is assumed to follow the - * native endian-ness. An explicit error message should be raised when attempting to - * export an array with non-native endianness - * - * Examples - * - float: type_code = 2, bits = 32, lanes = 1 - * - float4(vectorized 4 float): type_code = 2, bits = 32, lanes = 4 - * - int8: type_code = 0, bits = 8, lanes = 1 - * - std::complex: type_code = 5, bits = 64, lanes = 1 - * - bool: type_code = 6, bits = 8, lanes = 1 (as per common array library convention, the underlying storage size of bool is 8 bits) - * - float8_e4m3: type_code = 8, bits = 8, lanes = 1 (packed in memory) - * - float6_e3m2fn: type_code = 16, bits = 6, lanes = 1 (packed in memory) - * - float4_e2m1fn: type_code = 17, bits = 4, lanes = 1 (packed in memory) - * - * When a sub-byte type is packed, DLPack requires the data to be in little bit-endian, i.e., - * for a packed data set D ((D >> (i * bits)) && bit_mask) stores the i-th element. - */ -typedef struct { - /*! - * \brief Type code of base types. - * We keep it uint8_t instead of DLDataTypeCode for minimal memory - * footprint, but the value should be one of DLDataTypeCode enum values. - * */ - uint8_t code; - /*! - * \brief Number of bits, common choices are 8, 16, 32. - */ - uint8_t bits; - /*! \brief Number of lanes in the type, used for vector types. */ - uint16_t lanes; -} DLDataType; - -/*! - * \brief Plain C Tensor object, does not manage memory. - */ -typedef struct { - /*! - * \brief The data pointer points to the allocated data. This will be CUDA - * device pointer or cl_mem handle in OpenCL. It may be opaque on some device - * types. This pointer is always aligned to 256 bytes as in CUDA. The - * `byte_offset` field should be used to point to the beginning of the data. - * - * Note that as of Nov 2021, multiple libraries (CuPy, PyTorch, TensorFlow, - * TVM, perhaps others) do not adhere to this 256 byte alignment requirement - * on CPU/CUDA/ROCm, and always use `byte_offset=0`. This must be fixed - * (after which this note will be updated); at the moment it is recommended - * to not rely on the data pointer being correctly aligned. - * - * For given DLTensor, the size of memory required to store the contents of - * data is calculated as follows: - * - * \code{.c} - * static inline size_t GetDataSize(const DLTensor* t) { - * size_t size = 1; - * for (tvm_index_t i = 0; i < t->ndim; ++i) { - * size *= t->shape[i]; - * } - * size *= (t->dtype.bits * t->dtype.lanes + 7) / 8; - * return size; - * } - * \endcode - * - * Note that if the tensor is of size zero, then the data pointer should be - * set to `NULL`. - */ - void *data; - /*! \brief The device of the tensor */ - DLDevice device; - /*! \brief Number of dimensions */ - int32_t ndim; - /*! \brief The data type of the pointer*/ - DLDataType dtype; - /*! - * \brief The shape of the tensor - * - * When ndim == 0, shape can be set to NULL. - */ - int64_t *shape; - /*! - * \brief strides of the tensor (in number of elements, not bytes), - * can not be NULL if ndim != 0, must points to - * an array of ndim elements that specifies the strides, - * so consumer can always rely on strides[dim] being valid for 0 <= dim < ndim. - * - * When ndim == 0, strides can be set to NULL. - * - * \note Before DLPack v1.2, strides can be NULL to indicate contiguous data. - * This is not allowed in DLPack v1.2 and later. The rationale - * is to simplify the consumer handling. - */ - int64_t *strides; - /*! \brief The offset in bytes to the beginning pointer to data */ - uint64_t byte_offset; -} DLTensor; - -/*! - * \brief C Tensor object, manage memory of DLTensor. This data structure is - * intended to facilitate the borrowing of DLTensor by another framework. It is - * not meant to transfer the tensor. When the borrowing framework doesn't need - * the tensor, it should call the deleter to notify the host that the resource - * is no longer needed. - * - * \note This data structure is used as Legacy DLManagedTensor - * in DLPack exchange and is deprecated after DLPack v0.8 - * Use DLManagedTensorVersioned instead. - * This data structure may get renamed or deleted in future versions. - * - * \sa DLManagedTensorVersioned - */ -typedef struct DLManagedTensor { - /*! \brief DLTensor which is being memory managed */ - DLTensor dl_tensor; - /*! \brief the context of the original host framework of DLManagedTensor in - * which DLManagedTensor is used in the framework. It can also be NULL. - */ - void *manager_ctx; - /*! - * \brief Destructor - this should be called - * to destruct the manager_ctx which backs the DLManagedTensor. It can be - * NULL if there is no way for the caller to provide a reasonable destructor. - * The destructor deletes the argument self as well. - */ - void (*deleter)(struct DLManagedTensor *self); -} DLManagedTensor; - -// bit masks used in the DLManagedTensorVersioned - -/*! \brief bit mask to indicate that the tensor is read only. */ -#define DLPACK_FLAG_BITMASK_READ_ONLY (1UL << 0UL) - -/*! - * \brief bit mask to indicate that the tensor is a copy made by the producer. - * - * If set, the tensor is considered solely owned throughout its lifetime by the - * consumer, until the producer-provided deleter is invoked. - */ -#define DLPACK_FLAG_BITMASK_IS_COPIED (1UL << 1UL) - -/*! - * \brief bit mask to indicate that whether a sub-byte type is packed or padded. - * - * The default for sub-byte types (ex: fp4/fp6) is assumed packed. This flag can - * be set by the producer to signal that a tensor of sub-byte type is padded. - */ -#define DLPACK_FLAG_BITMASK_IS_SUBBYTE_TYPE_PADDED (1UL << 2UL) - -/*! - * \brief A versioned and managed C Tensor object, manage memory of DLTensor. - * - * This data structure is intended to facilitate the borrowing of DLTensor by - * another framework. It is not meant to transfer the tensor. When the borrowing - * framework doesn't need the tensor, it should call the deleter to notify the - * host that the resource is no longer needed. - * - * \note This is the current standard DLPack exchange data structure. - */ -typedef struct DLManagedTensorVersioned { - /*! - * \brief The API and ABI version of the current managed Tensor - */ - DLPackVersion version; - /*! - * \brief the context of the original host framework. - * - * Stores DLManagedTensorVersioned is used in the - * framework. It can also be NULL. - */ - void *manager_ctx; - /*! - * \brief Destructor. - * - * This should be called to destruct manager_ctx which holds the DLManagedTensorVersioned. - * It can be NULL if there is no way for the caller to provide a reasonable - * destructor. The destructor deletes the argument self as well. - */ - void (*deleter)(struct DLManagedTensorVersioned *self); - /*! - * \brief Additional bitmask flags information about the tensor. - * - * By default the flags should be set to 0. - * - * \note Future ABI changes should keep everything until this field - * stable, to ensure that deleter can be correctly called. - * - * \sa DLPACK_FLAG_BITMASK_READ_ONLY - * \sa DLPACK_FLAG_BITMASK_IS_COPIED - */ - uint64_t flags; - /*! \brief DLTensor which is being memory managed */ - DLTensor dl_tensor; -} DLManagedTensorVersioned; - -//---------------------------------------------------------------------- -// DLPack `__dlpack_c_exchange_api__` fast exchange protocol definitions -//---------------------------------------------------------------------- -/*! - * \brief Request a producer library to create a new tensor. - * - * Create a new `DLManagedTensorVersioned` within the context of the producer - * library. The allocation is defined via the prototype DLTensor. - * - * This function is exposed by the framework through the DLPackExchangeAPI. - * - * \param prototype The prototype DLTensor. Only the dtype, ndim, shape, - * and device fields are used. - * \param out The output DLManagedTensorVersioned. - * \param error_ctx Context for `SetError`. - * \param SetError The function to set the error. - * \return The owning DLManagedTensorVersioned* or NULL on failure. - * SetError is called exactly when NULL is returned (the implementer - * must ensure this). - * \note - As a C function, must not thrown C++ exceptions. - * - Error propagation via SetError to avoid any direct need - * of Python API. Due to this `SetError` may have to ensure the GIL is - * held since it will presumably set a Python error. - * - * \sa DLPackExchangeAPI - */ -typedef int (*DLPackManagedTensorAllocator)( // - DLTensor *prototype, DLManagedTensorVersioned **out, void *error_ctx, // - void (*SetError)(void *error_ctx, const char *kind, const char *message) // -); - -/*! - * \brief Exports a PyObject* Tensor/NDArray to a DLManagedTensorVersioned. - * - * This function does not perform any stream synchronization. The consumer should query - * DLPackCurrentWorkStream to get the current work stream and launch kernels on it. - * - * This function is exposed by the framework through the DLPackExchangeAPI. - * - * \param py_object The Python object to convert. Must have the same type - * as the one the `DLPackExchangeAPI` was discovered from. - * \param out The output DLManagedTensorVersioned. - * \return The owning DLManagedTensorVersioned* or NULL on failure with a - * Python exception set. If the data cannot be described using DLPack - * this should be a BufferError if possible. - * \note - As a C function, must not thrown C++ exceptions. - * - * \sa DLPackExchangeAPI, DLPackCurrentWorkStream - */ -typedef int (*DLPackManagedTensorFromPyObjectNoSync)( // - void *py_object, // - DLManagedTensorVersioned **out // -); - -/*! - * \brief Exports a PyObject* Tensor/NDArray to a provided DLTensor. - * - * This function provides a faster interface for temporary, non-owning, exchange. - * The producer (implementer) still owns the memory of data, strides, shape. - * The liveness of the DLTensor and the data it views is only guaranteed until - * control is returned. - * - * This function currently assumes that the producer (implementer) can fill - * in the DLTensor shape and strides without the need for temporary allocations. - * - * This function does not perform any stream synchronization. The consumer should query - * DLPackCurrentWorkStream to get the current work stream and launch kernels on it. - * - * This function is exposed by the framework through the DLPackExchangeAPI. - * - * \param py_object The Python object to convert. Must have the same type - * as the one the `DLPackExchangeAPI` was discovered from. - * \param out The output DLTensor, whose space is pre-allocated on stack. - * \return 0 on success, -1 on failure with a Python exception set. - * \note - As a C function, must not thrown C++ exceptions. - * - * \sa DLPackExchangeAPI, DLPackCurrentWorkStream - */ -typedef int (*DLPackDLTensorFromPyObjectNoSync)( // - void *py_object, // - DLTensor *out // -); - -/*! - * \brief Obtain the current work stream of a device. - * - * Obtain the current work stream of a device from the producer framework. - * For example, it should map to torch.cuda.current_stream in PyTorch. - * - * When device_type is kDLCPU, the consumer do not have to query the stream - * and the producer can simply return NULL when queried. - * The consumer do not have to do anything on stream sync or setting. - * So CPU only framework can just provide a dummy implementation that - * always set out_current_stream[0] to NULL. - * - * \param device_type The device type. - * \param device_id The device id. - * \param out_current_stream The output current work stream. - * - * \return 0 on success, -1 on failure with a Python exception set. - * \note - As a C function, must not thrown C++ exceptions. - * - * \sa DLPackExchangeAPI - */ -typedef int (*DLPackCurrentWorkStream)( // - DLDeviceType device_type, // - int32_t device_id, // - void **out_current_stream // -); - -/*! - * \brief Imports a DLManagedTensorVersioned to a PyObject* Tensor/NDArray. - * - * Convert an owning DLManagedTensorVersioned* to the Python tensor of the - * producer (implementer) library with the correct type. - * - * This function does not perform any stream synchronization. - * - * This function is exposed by the framework through the DLPackExchangeAPI. - * - * \param tensor The DLManagedTensorVersioned to convert the ownership of the - * tensor is stolen. - * \param out_py_object The output Python object. - * \return 0 on success, -1 on failure with a Python exception set. - * - * \sa DLPackExchangeAPI - */ -typedef int (*DLPackManagedTensorToPyObjectNoSync)( // - DLManagedTensorVersioned *tensor, // - void **out_py_object // -); - -/*! - * \brief DLPackExchangeAPI stable header. - * \sa DLPackExchangeAPI - */ -typedef struct DLPackExchangeAPIHeader { - /*! - * \brief The provided DLPack version the consumer must check major version - * compatibility before using this struct. - */ - DLPackVersion version; - /*! - * \brief Optional pointer to an older DLPackExchangeAPI in the chain. - * - * It must be NULL if the framework does not support older versions. - * If the current major version is larger than the one supported by the - * consumer, the consumer may walk this to find an earlier supported version. - * - * \sa DLPackExchangeAPI - */ - struct DLPackExchangeAPIHeader *prev_api; -} DLPackExchangeAPIHeader; - -/*! - * \brief Framework-specific function pointers table for DLPack exchange. - * - * Additionally to `__dlpack__()` we define a C function table sharable by - * - * Python implementations via `__dlpack_c_exchange_api__`. - * This attribute must be set on the type as a Python PyCapsule - * with name "dlpack_exchange_api". - * - * A consumer library may use a pattern such as: - * - * \code - * - * PyObject *api_capsule = PyObject_GetAttrString( - * (PyObject *)Py_TYPE(tensor_obj), "__dlpack_c_exchange_api__") - * ); - * if (api_capsule == NULL) { goto handle_error; } - * MyDLPackExchangeAPI *api = (MyDLPackExchangeAPI *)PyCapsule_GetPointer( - * api_capsule, "dlpack_exchange_api" - * ); - * Py_DECREF(api_capsule); - * if (api == NULL) { goto handle_error; } - * - * \endcode - * - * Note that this must be defined on the type. The consumer should look up the - * attribute on the type and may cache the result for each unique type. - * - * The precise API table is given by: - * \code - * struct MyDLPackExchangeAPI : public DLPackExchangeAPI { - * MyDLPackExchangeAPI() { - * header.version.major = DLPACK_MAJOR_VERSION; - * header.version.minor = DLPACK_MINOR_VERSION; - * header.prev_version_api = nullptr; - * - * managed_tensor_allocator = MyDLPackManagedTensorAllocator; - * managed_tensor_from_py_object_no_sync = MyDLPackManagedTensorFromPyObjectNoSync; - * managed_tensor_to_py_object_no_sync = MyDLPackManagedTensorToPyObjectNoSync; - * dltensor_from_py_object_no_sync = MyDLPackDLTensorFromPyObjectNoSync; - * current_work_stream = MyDLPackCurrentWorkStream; - * } - * - * static const DLPackExchangeAPI* Global() { - * static MyDLPackExchangeAPI inst; - * return &inst; - * } - * }; - * \endcode - * - * Guidelines for leveraging DLPackExchangeAPI: - * - * There are generally two kinds of consumer needs for DLPack exchange: - * - N0: library support, where consumer.kernel(x, y, z) would like to run a kernel - * with the data from x, y, z. The consumer is also expected to run the kernel with the same - * stream context as the producer. For example, when x, y, z is torch.Tensor, - * consumer should query exchange_api->current_work_stream to get the - * current stream and launch the kernel with the same stream. - * This setup is necessary for no synchronization in kernel launch and maximum compatibility - * with CUDA graph capture in the producer. - * This is the desirable behavior for library extension support for frameworks like PyTorch. - * - N1: data ingestion and retention - * - * Note that obj.__dlpack__() API should provide useful ways for N1. - * The primary focus of the current DLPackExchangeAPI is to enable faster exchange N0 - * with the support of the function pointer current_work_stream. - * - * Array/Tensor libraries should statically create and initialize this structure - * then return a pointer to DLPackExchangeAPI as an int value in Tensor/Array. - * The DLPackExchangeAPI* must stay alive throughout the lifetime of the process. - * - * One simple way to do so is to create a static instance of DLPackExchangeAPI - * within the framework and return a pointer to it. The following code - * shows an example to do so in C++. It should also be reasonably easy - * to do so in other languages. - */ -typedef struct DLPackExchangeAPI { - /*! - * \brief The header that remains stable across versions. - */ - DLPackExchangeAPIHeader header; - /*! - * \brief Producer function pointer for DLPackManagedTensorAllocator - * This function must not be NULL. - * \sa DLPackManagedTensorAllocator - */ - DLPackManagedTensorAllocator managed_tensor_allocator; - /*! - * \brief Producer function pointer for DLPackManagedTensorFromPyObject - * This function must be not NULL. - * \sa DLPackManagedTensorFromPyObject - */ - DLPackManagedTensorFromPyObjectNoSync managed_tensor_from_py_object_no_sync; - /*! - * \brief Producer function pointer for DLPackManagedTensorToPyObject - * This function must be not NULL. - * \sa DLPackManagedTensorToPyObjectNoSync - */ - DLPackManagedTensorToPyObjectNoSync managed_tensor_to_py_object_no_sync; - /*! - * \brief Producer function pointer for DLPackDLTensorFromPyObject - * This function can be NULL when the producer does not support this function. - * \sa DLPackDLTensorFromPyObjectNoSync - */ - DLPackDLTensorFromPyObjectNoSync dltensor_from_py_object_no_sync; - /*! - * \brief Producer function pointer for DLPackCurrentWorkStream - * This function must be not NULL. - * \sa DLPackCurrentWorkStream - */ - DLPackCurrentWorkStream current_work_stream; -} DLPackExchangeAPI; - -#ifdef __cplusplus -} // DLPACK_EXTERN_C -#endif -#endif // DLPACK_DLPACK_H_ diff --git a/src/infiniop/ops/gptq_marlin_gemm/gptq_marlin_gemm.h b/src/infiniop/ops/gptq_marlin_gemm/gptq_marlin_gemm.h deleted file mode 100644 index ddfb3cbae..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/gptq_marlin_gemm.h +++ /dev/null @@ -1,66 +0,0 @@ -#ifndef __GPTQ_MARLIN_GEMM_H__ -#define __GPTQ_MARLIN_GEMM_H__ - -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::gptq_marlin_gemm::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - GptqMarlinGemmInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - size_t workspace_size_, \ - Opaque *opaque, \ - GptqMarlinGemmInfo info, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size_) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t a_desc, \ - infiniopTensorDescriptor_t b_desc, \ - infiniopTensorDescriptor_t b_scales_desc, \ - infiniopTensorDescriptor_t global_scales_desc, \ - infiniopTensorDescriptor_t b_zeros_desc, \ - infiniopTensorDescriptor_t g_idx_desc, \ - infiniopTensorDescriptor_t perm_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *out, \ - const void *a, \ - const void *b, \ - void *b_scales, \ - void *global_scales, \ - void *b_zeros, \ - void *g_idx, \ - void *perm, \ - int64_t b_q_type_id, \ - bool is_k_full, \ - bool use_atomic_add, \ - bool use_fp32_reduce, \ - bool is_zp_float, \ - void *stream) const; \ - }; \ - } - -#endif //__GPTQ_MARLIN_GEMM_H__ diff --git a/src/infiniop/ops/gptq_marlin_gemm/info.h b/src/infiniop/ops/gptq_marlin_gemm/info.h deleted file mode 100644 index 218c25e63..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/info.h +++ /dev/null @@ -1,59 +0,0 @@ -#ifndef __GPTQ_MARLIN_GEMM_INFO_H__ -#define __GPTQ_MARLIN_GEMM_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -#include - -namespace op::gptq_marlin_gemm { - -class GptqMarlinGemmInfo { - GptqMarlinGemmInfo() = default; - -public: - infiniDtype_t dtype; - size_t M, K, N, b_q_size_1; - int num_groups; - ptrdiff_t a_stride_0; - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t b_scales_desc, - infiniopTensorDescriptor_t global_scales_desc, - infiniopTensorDescriptor_t b_zeros_desc, - infiniopTensorDescriptor_t g_idx_desc, - infiniopTensorDescriptor_t perm_desc) { - CHECK_OR_RETURN( - out_desc != nullptr && a_desc != nullptr && b_desc != nullptr && b_scales_desc != nullptr, - INFINI_STATUS_NULL_POINTER); - const infiniDtype_t dtype = a_desc->dtype(); - size_t M = out_desc->dim(0); - size_t N = out_desc->dim(1); - size_t K = a_desc->dim(1); - size_t b_q_size_1 = b_desc->dim(1); - int num_groups = static_cast(b_scales_desc->dim(0)); - ptrdiff_t a_stride_0 = a_desc->strides()[0]; - - auto ndim = out_desc->ndim(); - CHECK_OR_RETURN(ndim == 2 - && a_desc->ndim() == ndim - && b_desc->ndim() == ndim - && b_scales_desc->ndim() == ndim, - INFINI_STATUS_BAD_TENSOR_SHAPE); - - CHECK_OR_RETURN(b_scales_desc->shape()[1] == N - && a_stride_0 % 8 == 0, - INFINI_STATUS_BAD_TENSOR_SHAPE); - - return utils::Result( - GptqMarlinGemmInfo{dtype, M, K, N, b_q_size_1, num_groups, a_stride_0}); - } -}; - -} // namespace op::gptq_marlin_gemm - -#endif // __GPTQ_MARLIN_GEMM_INFO_H__ diff --git a/src/infiniop/ops/gptq_marlin_gemm/marlin/dequant.h b/src/infiniop/ops/gptq_marlin_gemm/marlin/dequant.h deleted file mode 100644 index 6a0d90e5d..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/marlin/dequant.h +++ /dev/null @@ -1,508 +0,0 @@ -/* -Fast Dequantization (Converting INT4/INT8/FP4/FP8 to FP16/BF16) - -The process of fast dequantization can be summarized as a combination -of bitwise operations and floating-point computations: - -weight =>(bit_op / bitwise operations)=> -f16_value =>(flop / floating-point computation)=> -dequantized_weight - -Since the dequantized weights typically require subtracting the zero point and -applying a scale factor, the floating-point computation step can be fused with -the zero-point subtraction and scaling operations. - -The following are the parts that need to be modified for the fused operation -of zero-point subtraction and scaling. - -## INT4 => FP16/BF16 or INT8 => FP16 - -The floating-point computation is `__hsub2` - -If has zero points: - - flop(bit_op(weight)) - flop(bit_op(zp)) - = sub(bit_op(weight), bias) - sub(bit_op(zp), bias) - = bit_op(weight) - bit_op(zp) - -so we don't need additional modification. - -If has float zero points: - - flop(bit_op(weight)) - fzp - = sub(bit_op(weight), bias) - fzp - = bit_op(weight) - (fzp + bias) - -where the `fzp + bias` can be computed at weight loading. But this -may have accuracy issue, so we should not use this in most cases. - -If has not zero points: - - scale(flop(bit_op(weight))) - = scale(sub(bit_op(weight), bias)) - = scale(bit_op(weight)) - scale(bias) - = fma(bit_op(weight), scale_factor, scale(bias)) - -where the `scale(bias)` can be cached. But this may have accuracy issue, -so we should not use this in most cases. - - -## INT8 => BF16 - -INT8 => BF16 is a special case, it use byte_perm instead of flop. -We cannot fused byte_perm with scaling. - - -## FP4/FP8 => FP16/BF16 - - scale(flop(bit_op(weight))) - = scale(mul(bit_op(weight), multiplier)) - = mul(bit_op(weight), scale_factor * multiplier) - -where `scale_factor * multiplier` can be computed at weight loading. - -*/ - -#include "marlin_dtypes.cuh" - -namespace device::marlin { - -#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 -// Lookup-table based 3-input logical operation; explicitly used for -// dequantization as the compiler does not seem to automatically recognize it in -// all cases. -template -__device__ inline int lop3(int a, int b, int c) { - int res; - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" - : "=r"(res) - : "r"(a), "r"(b), "r"(c), "n"(lut)); - return res; -} - -// Constructs destination register by taking bytes from 2 sources (based on -// mask) -template -__device__ inline uint32_t prmt(uint32_t a) { - uint32_t res; - asm volatile("prmt.b32 %0, %1, %2, %3;\n" - : "=r"(res) - : "r"(a), "n"(start_byte), "n"(mask)); - return res; -} - -template -__device__ inline void dequant(int q, scalar_t2 *frag_b); - -// -// Efficiently dequantize 4bit values packed in an int32 value into a full -// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below, -// with some small changes: -// - FP16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287 -// - BF16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L327-L385 -// -template <> -__device__ inline void dequant(int q, half2 *frag_b) { - const int MASK = 0x000f000f; - const int EX = 0x64006400; - // Guarantee that the `(a & b) | c` operations are LOP3s. - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - q >>= 4; - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - - frag_b[0] = *reinterpret_cast(&lo); - frag_b[1] = *reinterpret_cast(&hi); -} - -template <> -__device__ inline void dequant(int q, half2 *frag_b) { - const int LO = 0x000f000f; - const int HI = 0x00f000f0; - const int EX = 0x64006400; - // Guarantee that the `(a & b) | c` operations are LOP3s. - // clang-format off - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); - // clang-format on - // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point - // directly into `SUB` and `ADD`. - const int SUB = 0x64086408; - const int MUL = 0x2c002c00; - const int ADD = 0xd480d480; - frag_b[0] = __hsub2(*reinterpret_cast(&lo), *reinterpret_cast(&SUB)); - frag_b[1] = __hfma2( - *reinterpret_cast(&hi), *reinterpret_cast(&MUL), *reinterpret_cast(&ADD)); -} - -template <> -__device__ inline void dequant(int q, half2 *frag_b) { - dequant(q, frag_b); -} - -template <> -__device__ inline void dequant(int q, half2 *frag_b) { - const int LO = 0x000f000f; - const int HI = 0x00f000f0; - const int EX = 0x64006400; - // Guarantee that the `(a & b) | c` operations are LOP3s. - // clang-format off - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX); - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX); - // clang-format on - // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point - // directly into `SUB` and `ADD`. - const int SUB = 0x64006400; - const int MUL = 0x2c002c00; - const int ADD = 0xd400d400; - frag_b[0] = __hsub2(*reinterpret_cast(&lo), *reinterpret_cast(&SUB)); - frag_b[1] = __hfma2( - *reinterpret_cast(&hi), *reinterpret_cast(&MUL), *reinterpret_cast(&ADD)); -} - -template <> -__device__ inline void dequant(int q, nv_bfloat162 *frag_b) { - static constexpr uint32_t MASK = 0x000f000f; - static constexpr uint32_t EX = 0x43004300; - - // Guarantee that the `(a & b) | c` operations are LOP3s. - // clang-format off - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - q >>= 4; - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, MASK, EX); - // clang-format on - - frag_b[0] = *reinterpret_cast(&lo); - frag_b[1] = *reinterpret_cast(&hi); -} - -template <> -__device__ inline void dequant(int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); - - static constexpr uint32_t SUB = 0x43084308; - - frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&SUB)); - frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&SUB)); -} - -template <> -__device__ inline void dequant(int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); -} - -template <> -__device__ inline void dequant(int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); - - static constexpr uint32_t SUB = 0x43004300; - - frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&SUB)); - frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&SUB)); -} - -// -// Fast Int8ToFp16/Int8ToBf16: Efficiently dequantize 8bit int values to fp16 or -// bf16 Reference: -// - FP16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85 -// - BF16: -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L125-L175 -// -template <> -__device__ inline void dequant(int q, half2 *frag_b) { - static constexpr uint32_t mask_for_elt_01 = 0x5250; - static constexpr uint32_t mask_for_elt_23 = 0x5351; - static constexpr uint32_t start_byte_for_fp16 = 0x64646464; - - uint32_t lo = prmt(q); - uint32_t hi = prmt(q); - - frag_b[0] = *reinterpret_cast(&lo); - frag_b[1] = *reinterpret_cast(&hi); -} - -template <> -__device__ inline void dequant(int q, half2 *frag_b) { - dequant(q, frag_b); - - static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480; - frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); - frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); -} - -template <> -__device__ inline void dequant(int q, half2 *frag_b) { - dequant(q, frag_b); -} - -template <> -__device__ inline void dequant(int q, half2 *frag_b) { - dequant(q, frag_b); - - static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400; - frag_b[0] = __hsub2(frag_b[0], *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); - frag_b[1] = __hsub2(frag_b[1], *reinterpret_cast(&I8s_TO_F16s_MAGIC_NUM)); -} - -template <> -__device__ inline void dequant(int q, nv_bfloat162 *frag_b) { - float fp32_intermediates[4]; - uint32_t *fp32_intermediates_casted = reinterpret_cast(fp32_intermediates); - - static constexpr uint32_t fp32_base = 0x4B000000; - fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); - fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); - fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); - fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); - - fp32_intermediates[0] -= 8388736.f; - fp32_intermediates[1] -= 8388736.f; - fp32_intermediates[2] -= 8388736.f; - fp32_intermediates[3] -= 8388736.f; - - uint32_t *bf16_result_ptr = reinterpret_cast(frag_b); - bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], fp32_intermediates_casted[1], 0x7632); - bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], fp32_intermediates_casted[3], 0x7632); -} - -template <> -__device__ inline void dequant(int q, nv_bfloat162 *frag_b) { - float fp32_intermediates[4]; - uint32_t *fp32_intermediates_casted = reinterpret_cast(fp32_intermediates); - - static constexpr uint32_t fp32_base = 0x4B000000; - fp32_intermediates_casted[0] = __byte_perm(q, fp32_base, 0x7650); - fp32_intermediates_casted[1] = __byte_perm(q, fp32_base, 0x7652); - fp32_intermediates_casted[2] = __byte_perm(q, fp32_base, 0x7651); - fp32_intermediates_casted[3] = __byte_perm(q, fp32_base, 0x7653); - - fp32_intermediates[0] -= 8388608.f; - fp32_intermediates[1] -= 8388608.f; - fp32_intermediates[2] -= 8388608.f; - fp32_intermediates[3] -= 8388608.f; - - uint32_t *bf16_result_ptr = reinterpret_cast(frag_b); - bf16_result_ptr[0] = __byte_perm(fp32_intermediates_casted[0], fp32_intermediates_casted[1], 0x7632); - bf16_result_ptr[1] = __byte_perm(fp32_intermediates_casted[2], fp32_intermediates_casted[3], 0x7632); -} - -template <> -__device__ inline void dequant(int q, half2 *frag_b) { - // Constants for FP8 (E4M3) and FP16 formats - constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5; - constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT; - constexpr int MASK = 0x7F007F00; - - // Extract and shift FP8 values to FP16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 8; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant(int q, half2 *frag_b) { - dequant(q, frag_b); - - // Constants for FP8 (E4M3) and FP16 formats - constexpr int FP8_EXPONENT = 4, FP16_EXPONENT = 5; - - // Construct and apply exponent bias - constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); - const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); - - // Convert to half2 and apply bias - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); -} - -template <> -__device__ inline void dequant(int q, nv_bfloat162 *frag_b) { - // Constants for FP8 (E4M3) and BF16 formats - constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; - constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; - - constexpr int MASK = 0x7F007F00; - - // Extract and shift FP8 values to BF16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 8; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant(int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); - - // Constants for FP8 (E4M3) and BF16 formats - constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; - - // Construct and apply exponent bias - constexpr int BIAS_OFFSET = (1 << (BF16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1)); - // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent - // position - constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; - const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); - - // Convert to bfloat162 and apply bias - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); -} - -template <> -__device__ inline void dequant(int q, half2 *frag_b) { - // Constants for FP4 (E2M1) and FP16 formats - constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5; - constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP4_EXPONENT; - constexpr int MASK = 0x70007000; - - // Extract and shift FP4 values to FP16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 4; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant(int q, half2 *frag_b) { - dequant(q, frag_b); - - // Constants for FP4 (E2M1) and FP16 formats - constexpr int FP4_EXPONENT = 2, FP16_EXPONENT = 5; - - // Construct and apply exponent bias - constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1)); - const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET)); - - // Convert to half2 and apply bias - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); -} - -template <> -__device__ inline void dequant(int q, nv_bfloat162 *frag_b) { - // Constants for FP4 (E2M1) and FP16 formats - constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8; - constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP4_EXPONENT; - constexpr int MASK = 0x70007000; - - // Extract and shift FP4 values to FP16 format - int Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 4; - int Out2 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant(int q, nv_bfloat162 *frag_b) { - dequant(q, frag_b); - - // Constants for FP4 (E2M1) and BF16 formats - constexpr int FP4_EXPONENT = 2, BF16_EXPONENT = 8; - - // Construct and apply exponent bias - constexpr int BIAS_OFFSET = (1 << (BF16_EXPONENT - 1)) - (1 << (FP4_EXPONENT - 1)); - // Add 127 (float exponent bias) to BIAS_OFFSET and shift to float exponent - // position - constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23; - const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast(&BIAS)); - - // Convert to half2 and apply bias - frag_b[1] = __hmul2(frag_b[1], bias_reg); - frag_b[0] = __hmul2(frag_b[0], bias_reg); -} - -template -__device__ inline void dequant_fp8_scales(int q, scalar_t2 *frag_b); - -template <> -__device__ inline void dequant_fp8_scales(int q, half2 *frag_b) { - int Out1 = (q & 0xFF00FF00) >> 1; - ; - q <<= 8; - int Out2 = (q & 0xFF00FF00) >> 1; - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -}; - -template <> -__device__ inline void dequant_fp8_scales(int q, nv_bfloat162 *frag_b) { - constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; - constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; - constexpr int MASK = 0x7F007F00; - - // Extract and shift FP8 values to BF16 format - int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 8; - int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -}; - -// New version with s_type_id parameter for marlin_moe_wna16_v2 -template -__device__ inline void dequant_fp8_scales(int q, scalar_t2 *frag_b); - -template <> -__device__ inline void dequant_fp8_scales(int q, half2 *frag_b) { - int Out1 = (q & 0xFF00FF00) >> 1; - ; - q <<= 8; - int Out2 = (q & 0xFF00FF00) >> 1; - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -}; - -template <> -__device__ inline void dequant_fp8_scales(int q, nv_bfloat162 *frag_b) { - constexpr int FP8_EXPONENT = 4, BF16_EXPONENT = 8; - constexpr int RIGHT_SHIFT = BF16_EXPONENT - FP8_EXPONENT; - constexpr int MASK = 0x7F007F00; - - // Extract and shift FP8 values to BF16 format - int Out1 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); - q <<= 8; - int Out2 = ((q & 0x80008000) >> 1) | ((q & MASK) >> RIGHT_SHIFT); - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -template <> -__device__ inline void dequant_fp8_scales(int q, nv_bfloat162 *frag_b) { - // In this conversion, 2 ** -127 in FP8E8M0 would become 0 in BF16, - // but we assume that such a extreme value would not occur in real models. - int Out1 = (q & 0xFF00FF00) >> 1; - q <<= 7; - int Out2 = q & 0x7F807F80; - - // Note: reverse indexing is intentional because weights are permuted - frag_b[1] = *reinterpret_cast(&Out1); - frag_b[0] = *reinterpret_cast(&Out2); -} - -#endif - -} // namespace device::marlin diff --git a/src/infiniop/ops/gptq_marlin_gemm/marlin/gptq_marlin.cuh b/src/infiniop/ops/gptq_marlin_gemm/marlin/gptq_marlin.cuh deleted file mode 100644 index ca85889f5..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/marlin/gptq_marlin.cuh +++ /dev/null @@ -1,1013 +0,0 @@ -/* - * Modified by Neural Magic - * Copyright (C) Marlin.2024 Elias Frantar - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * Adapted from https://github.com/IST-DASLab/marlin - */ - -#pragma once - -#include "../sgl_kernel/tensor.h" - -#include "../sgl_kernel/scalar_type.hpp" - -#include "kernel.h" -#include "marlin_template.h" - -namespace device::marlin { - -__global__ void MarlinDefault(MARLIN_KERNEL_PARAMS){}; - -using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - -__global__ void permute_cols_kernel( - int4 const *__restrict__ a_int4_ptr, - int const *__restrict__ perm_int_ptr, - int4 *__restrict__ out_int4_ptr, - int size_m, - int size_k, - int lda, - int block_rows) {} - -#else - -// For a given "a" of size [M,K] performs a permutation of the K columns based -// on the given "perm" indices. -__global__ void permute_cols_kernel( - int4 const *__restrict__ a_int4_ptr, - int const *__restrict__ perm_int_ptr, - int4 *__restrict__ out_int4_ptr, - int size_m, - int size_k, - int lda, - int block_rows) { - auto start_row = block_rows * blockIdx.x; - int finish_row = start_row + block_rows; - if (finish_row > size_m) { - finish_row = size_m; - } - int cur_block_rows = finish_row - start_row; - - int input_row_stride = lda * sizeof(half) / 16; - int output_row_stride = size_k * sizeof(half) / 16; - - auto permute_row = [&](int row) { - int iters = size_k / default_threads; - int rest = size_k % default_threads; - - int input_offset = row * input_row_stride; - int output_offset = row * output_row_stride; - - half const *a_row_half = reinterpret_cast(a_int4_ptr + input_offset); - half *out_half = reinterpret_cast(out_int4_ptr + output_offset); - - int base_k = 0; - - for (int i = 0; i < iters; i++) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - - base_k += default_threads; - } - - if (rest) { - if (threadIdx.x < rest) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - } - } - }; - - for (int i = 0; i < cur_block_rows; i++) { - int cur_row = start_row + i; - if (cur_row < size_m) { - permute_row(cur_row); - } - } -} - -typedef struct -{ - int thread_k; - int thread_n; - int num_threads; -} thread_config_t; - -thread_config_t small_batch_thread_configs[] = { - // Ordered by priority - - // thread_k, thread_n, num_threads - {128, 128, 256}, - {64, 128, 128}, - {128, 64, 128}}; - -thread_config_t large_batch_thread_configs[] = { - // Ordered by priority - - // thread_k, thread_n, num_threads - {64, 256, 256}, - {64, 128, 128}, - {128, 64, 128}}; - -typedef struct -{ - int blocks_per_sm; - thread_config_t tb_cfg; -} exec_config_t; - -int get_scales_cache_size( - thread_config_t const &th_config, - int prob_m, - int prob_n, - int prob_k, - int num_bits, - int group_size, - bool has_act_order, - bool is_k_full) { - bool cache_scales_chunk = has_act_order && !is_k_full; - - int tb_n = th_config.thread_n; - int tb_k = th_config.thread_k; - - // Get max scale groups per thread-block - int tb_groups; - if (group_size == -1) { - tb_groups = 1; - } else if (group_size == 0) { - tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size - } else { - tb_groups = div_ceil(tb_k, group_size); - } - - if (cache_scales_chunk) { - int load_groups = tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K - load_groups = max(load_groups, 32); // We load at least 32 scale groups - return load_groups * tb_n * 2; - } else { - int tb_scales = tb_groups * tb_n * 2; - - return tb_scales * pipe_stages; - } -} - -int get_kernel_cache_size( - thread_config_t const &th_config, - int thread_m_blocks, - int prob_m, - int prob_n, - int prob_k, - int num_bits, - int group_size, - bool has_act_order, - bool is_k_full, - int has_zp, - int is_zp_float) { - int pack_factor = 32 / num_bits; - - // Get B size - int tb_k = th_config.thread_k; - int tb_n = th_config.thread_n; - int tb_m = thread_m_blocks * 16; - int sh_a_size = pipe_stages * (tb_m * tb_k) * 2; - int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4; - int sh_red_size = tb_m * (tb_n + 8); - int sh_s_size = get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full); - int sh_g_idx_size = has_act_order && !is_k_full ? pipe_stages * tb_k / 4 : 0; - int sh_zp_size = 0; - if (has_zp) { - if (is_zp_float) { - sh_zp_size = sh_s_size; - } else if (num_bits == 4) { - sh_zp_size = sh_s_size / 4; - } else if (num_bits == 8) { - sh_zp_size = sh_s_size / 2; - } - } - - int total_size = max(sh_b_size, sh_red_size) + sh_a_size + sh_s_size + sh_zp_size + sh_g_idx_size; - - return total_size; -} - -bool is_valid_config( - thread_config_t const &th_config, - int thread_m_blocks, - int prob_m, - int prob_n, - int prob_k, - int num_bits, - int group_size, - bool has_act_order, - bool is_k_full, - int has_zp, - int is_zp_float, - int max_shared_mem) { - // Sanity - if (th_config.thread_k == -1 || th_config.thread_n == -1 || th_config.num_threads == -1) { - return false; - } - - // Verify K/N are divisible by thread K/N - if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { - return false; - } - - // Verify min for thread K/N - if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) { - return false; - } - - // num_threads must be at least 128 (= 4 warps) - if (th_config.num_threads < 128) { - return false; - } - - // Check that pipeline fits into cache - int cache_size = get_kernel_cache_size( - th_config, - thread_m_blocks, - prob_m, - prob_n, - prob_k, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float); - return cache_size <= max_shared_mem; -} - -#define _GET_IF( \ - W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \ - else if ( \ - q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && thread_n_blocks == THREAD_N_BLOCKS && thread_k_blocks == THREAD_K_BLOCKS && m_block_size_8 == M_BLOCK_SIZE_8 && group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && is_zp_float == IS_ZP_FLOAT) { \ - kernel = Marlin< \ - scalar_t, \ - W_TYPE.id(), \ - NUM_THREADS, \ - THREAD_M_BLOCKS, \ - THREAD_N_BLOCKS, \ - THREAD_K_BLOCKS, \ - M_BLOCK_SIZE_8, \ - pipe_stages, \ - GROUP_BLOCKS, \ - IS_ZP_FLOAT>; \ - } - -// COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false) -// this is the most common cases -// BIGGROUP: cases for big group size (group_blocks in [-1, 8]) -// FZP: cases for float-zero-point (is_zp_float = true) -// ACT: cases for act order case (group_blocks == 0) -// FP4: cases for nvfp4(e2m1) (group_blocks == 1) -#define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) - -#define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ - \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ - \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) - -#define COMMON_GET_IF(W_TYPE) \ - COMMON_GET_IF_M1(W_TYPE, 8, 8, 256) \ - COMMON_GET_IF_M1(W_TYPE, 8, 4, 128) \ - COMMON_GET_IF_M1(W_TYPE, 4, 8, 128) \ - COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \ - COMMON_GET_IF_M234(W_TYPE, 8, 4, 128) \ - COMMON_GET_IF_M234(W_TYPE, 4, 8, 128) - -#define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) - -#define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) - -#define BIGGROUP_GET_IF(W_TYPE) \ - BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256) \ - BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128) \ - BIGGROUP_GET_IF_M1(W_TYPE, 4, 8, 128) \ - BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \ - BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128) \ - BIGGROUP_GET_IF_M234(W_TYPE, 4, 8, 128) - -#define FP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) - -#define FP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) - -#define FP4_GET_IF(W_TYPE) \ - FP4_GET_IF_M1(W_TYPE, 8, 8, 256) \ - FP4_GET_IF_M1(W_TYPE, 8, 4, 128) \ - FP4_GET_IF_M1(W_TYPE, 4, 8, 128) \ - FP4_GET_IF_M234(W_TYPE, 16, 4, 256) \ - FP4_GET_IF_M234(W_TYPE, 8, 4, 128) \ - FP4_GET_IF_M234(W_TYPE, 4, 8, 128) - -// We currently have 4-bit models only with group_blocks == 4 -#define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) - -#define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) - -#define FZP_GET_IF(W_TYPE) \ - FZP_GET_IF_M1(W_TYPE, 8, 8, 256) \ - FZP_GET_IF_M1(W_TYPE, 8, 4, 128) \ - FZP_GET_IF_M1(W_TYPE, 4, 8, 128) \ - FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \ - FZP_GET_IF_M234(W_TYPE, 8, 4, 128) \ - FZP_GET_IF_M234(W_TYPE, 4, 8, 128) - -// We currently have 4-bit models only with group_blocks == 4 -#define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) - -#define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) - -#define ACT_GET_IF(W_TYPE) \ - ACT_GET_IF_M1(W_TYPE, 8, 8, 256) \ - ACT_GET_IF_M1(W_TYPE, 8, 4, 128) \ - ACT_GET_IF_M1(W_TYPE, 4, 8, 128) \ - ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \ - ACT_GET_IF_M234(W_TYPE, 8, 4, 128) \ - ACT_GET_IF_M234(W_TYPE, 4, 8, 128) - -template -MarlinFuncPtr get_marlin_kernel( - const host::ScalarType q_type, - int thread_m_blocks, - int thread_n_blocks, - int thread_k_blocks, - bool m_block_size_8, - bool has_act_order, - bool has_zp, - int group_blocks, - int num_threads, - bool is_zp_float) { - int num_bits = q_type.size_bits(); - auto kernel = MarlinDefault; - if (false) { - } - - COMMON_GET_IF(host::kU4) - COMMON_GET_IF(host::kU4B8) - COMMON_GET_IF(host::kU8B128) - - FP4_GET_IF(host::kFE2M1f) - - BIGGROUP_GET_IF(host::kFE4M3fn) - - ACT_GET_IF(host::kU4B8) - ACT_GET_IF(host::kU8B128) - - if (std::is_same::value) { - if (false) { - } - FZP_GET_IF(host::kU4) - } - - return kernel; -} - -template -exec_config_t determine_exec_config( - const host::ScalarType &q_type, - int prob_m, - int prob_n, - int prob_k, - int thread_m_blocks, - bool m_block_size_8, - int num_bits, - int group_size, - bool has_act_order, - bool is_k_full, - bool has_zp, - bool is_zp_float, - int max_shared_mem, - int sms) { - exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}}; - thread_config_t *thread_configs = thread_m_blocks > 1 ? large_batch_thread_configs : small_batch_thread_configs; - int thread_configs_size = thread_m_blocks > 1 ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t) - : sizeof(small_batch_thread_configs) / sizeof(thread_config_t); - - for (int i = 0; i < thread_configs_size; i++) { - thread_config_t th_config = thread_configs[i]; - - if (!is_valid_config( - th_config, - thread_m_blocks, - prob_m, - prob_n, - prob_k, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float, - max_shared_mem)) { - continue; - } - - int cache_size = get_kernel_cache_size( - th_config, - thread_m_blocks, - prob_m, - prob_n, - prob_k, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float); - - int group_blocks = 0; - if (!has_act_order) { - group_blocks = group_size == -1 ? -1 : group_size / 16; - } - - auto kernel = get_marlin_kernel( - q_type, - thread_m_blocks, - th_config.thread_n / 16, - th_config.thread_k / 16, - m_block_size_8, - has_act_order, - has_zp, - group_blocks, - th_config.num_threads, - is_zp_float); - - if (kernel == MarlinDefault) { - continue; - } - - // int m_tiles = div_ceil(prob_m, thread_m_blocks * 16); - // int n_tiles = prob_n / th_config.thread_n; - // int k_tiles = prob_k / th_config.thread_k; - - return {1, th_config}; - } - - return exec_cfg; -} - -template -void marlin_mm( - const void *A, - const void *B, - void *C, - void *C_tmp, - void *s, - void *s2, - void *zp, - void *g_idx, - void *perm, - void *a_tmp, - int prob_m, - int prob_n, - int prob_k, - int lda, - void *workspace, - host::ScalarType const &q_type, - bool has_act_order, - bool is_k_full, - bool has_zp, - int num_groups, - int group_size, - int dev, - cudaStream_t stream, - int thread_k_init, - int thread_n_init, - int sms, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float) { - if (has_zp) { - host::RuntimeCheck( - q_type == host::kU4 || q_type == host::kU8, "q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str()); - } else { - host::RuntimeCheck( - q_type == host::kU4B8 || q_type == host::kU8B128 || q_type == host::kFE4M3fn || q_type == host::kFE2M1f, - "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when " - "has_zp = False. Got = ", - q_type.str()); - } - - host::RuntimeCheck( - prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, ", ", prob_n, ", ", prob_k, "]"); - - int group_blocks = 0; - if (has_act_order) { - if (is_k_full) { - host::RuntimeCheck(group_size != -1); - group_blocks = group_size / 16; - host::RuntimeCheck( - prob_k % group_blocks == 0, "prob_k = ", prob_k, " is not divisible by group_blocks = ", group_blocks); - } else { - host::RuntimeCheck(group_size == 0); - group_blocks = 0; - } - } else { - if (group_size == -1) { - group_blocks = -1; - } else { - group_blocks = group_size / 16; - host::RuntimeCheck( - prob_k % group_blocks == 0, "prob_k = ", prob_k, " is not divisible by group_blocks = ", group_blocks); - } - } - - int num_bits = q_type.size_bits(); - const int4 *A_ptr = (const int4 *)A; - const int4 *B_ptr = (const int4 *)B; - int4 *C_ptr = (int4 *)C; - int4 *C_tmp_ptr = (int4 *)C_tmp; - const int4 *s_ptr = (const int4 *)s; - const uint16_t *s2_ptr = (const uint16_t *)s2; - const int4 *zp_ptr = (const int4 *)zp; - const int *g_idx_ptr = (const int *)g_idx; - const int *perm_ptr = (const int *)perm; - int4 *a_tmp_ptr = (int4 *)a_tmp; - - int *locks = (int *)workspace; - - if (has_act_order) { - // Permute A columns - int block_rows = div_ceil(prob_m, sms); - host::LaunchKernel(sms, default_threads, stream)( - permute_cols_kernel, A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, lda, block_rows); - A_ptr = a_tmp_ptr; - lda = prob_k; - - // If we have a full K, then we can run the non-act-order version of Marlin - // (since the weight rows are reordered by increasing group ids, and by - // having a full K, we have full original groups) - if (is_k_full) { - has_act_order = false; - } - } - - int max_shared_mem = 0; - host::RuntimeDeviceCheck(cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev)); - host::RuntimeCheck(max_shared_mem > 0); - - int max_par = 16; - if (prob_n <= 4096) { - max_par = 16 * 8; - } - int max_shared_mem_new = max_shared_mem; - int rest_m = prob_m; - int max_thread_m_blocks = 4; - while (rest_m) { - int par_count = rest_m / (max_thread_m_blocks * 16); - if (par_count > max_par) { - par_count = max_par; - } - int prob_m_split = par_count > 0 ? (par_count * (max_thread_m_blocks * 16)) : rest_m; - - int thread_k = thread_k_init; - int thread_n = thread_n_init; - - int thread_m_blocks = min(div_ceil(prob_m_split, 16), max_thread_m_blocks); - int m_block_size_8 = prob_m_split <= 8; - - // Set thread config - exec_config_t exec_cfg; - thread_config_t thread_tfg; - if (thread_k != -1 && thread_n != -1) { - thread_tfg = thread_config_t{thread_k, thread_n, default_threads}; - exec_cfg = exec_config_t{1, thread_tfg}; - host::RuntimeCheck(prob_n % thread_n == 0, "prob_n = ", prob_n, " is not divisible by thread_n = ", thread_n); - host::RuntimeCheck(prob_k % thread_k == 0, "prob_k = ", prob_k, " is not divisible by thread_k = ", thread_k); - } else { - // Auto config - exec_cfg = determine_exec_config( - q_type, - prob_m_split, - prob_n, - prob_k, - thread_m_blocks, - m_block_size_8, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float, - max_shared_mem, - sms); - thread_tfg = exec_cfg.tb_cfg; - if (thread_tfg.thread_k == -1 && max_thread_m_blocks > 1) { - max_thread_m_blocks--; - continue; - } - } - - int num_threads = thread_tfg.num_threads; - thread_k = thread_tfg.thread_k; - thread_n = thread_tfg.thread_n; - int blocks = sms * exec_cfg.blocks_per_sm; - if (exec_cfg.blocks_per_sm > 1) { - max_shared_mem_new = max_shared_mem / exec_cfg.blocks_per_sm - 1024; - } - - int thread_k_blocks = thread_k / 16; - int thread_n_blocks = thread_n / 16; - - host::RuntimeCheck( - is_valid_config( - thread_tfg, - thread_m_blocks, - prob_m_split, - prob_n, - prob_k, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float, - max_shared_mem_new), - "Invalid thread config: thread_m_blocks = ", - thread_m_blocks, - ", thread_k = ", - thread_tfg.thread_k, - ", thread_n = ", - thread_tfg.thread_n, - ", num_threads = ", - thread_tfg.num_threads, - " for MKN = [", - prob_m, - ", ", - prob_k, - ", ", - prob_n, - "] and num_bits = ", - num_bits, - ", prob_m_split = ", - prob_m_split, - ", group_size = ", - group_size, - ", has_act_order = ", - has_act_order, - ", is_k_full = ", - is_k_full, - ", has_zp = ", - has_zp, - ", is_zp_float = ", - is_zp_float, - ", max_shared_mem_new = ", - max_shared_mem_new); - - auto kernel = get_marlin_kernel( - q_type, - thread_m_blocks, - thread_n_blocks, - thread_k_blocks, - m_block_size_8, - has_act_order, - has_zp, - group_blocks, - num_threads, - is_zp_float); - - if (kernel == MarlinDefault) { - host::Panic( - "Unsupported shapes: MNK = [", - prob_m, - ", ", - prob_n, - ", ", - prob_k, - "]", - ", has_act_order = ", - has_act_order, - ", num_groups = ", - num_groups, - ", group_size = ", - group_size, - ", prob_m_split = ", - prob_m_split, - ", thread_m_blocks = ", - thread_m_blocks, - ", thread_n_blocks = ", - thread_n_blocks, - ", thread_k_blocks = ", - thread_k_blocks, - ", num_threads = ", - num_threads, - ", num_bits = ", - num_bits); - } - - host::RuntimeDeviceCheck( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem_new)); - - bool part_use_atomic_add = use_atomic_add && div_ceil(prob_m_split, 64) * prob_n <= 2048; - - host::LaunchKernel(blocks, num_threads, stream, max_shared_mem_new)( - kernel, - A_ptr, - B_ptr, - C_ptr, - C_tmp_ptr, - s_ptr, - s2_ptr, - zp_ptr, - g_idx_ptr, - num_groups, - prob_m_split, - prob_n, - prob_k, - lda, - locks, - part_use_atomic_add, - use_fp32_reduce, - max_shared_mem_new); - - A_ptr += prob_m_split * (lda / 8); - C_ptr += prob_m_split * (prob_n / 8); - rest_m -= prob_m_split; - } -} - -#endif - -} // namespace device::marlin - -template -void gptq_marlin_gemm( - tvm::ffi::TensorView a, - tvm::ffi::TensorView b_q_weight, - tvm::ffi::TensorView b_scales, - tvm::ffi::TensorView global_scale, - tvm::ffi::TensorView b_zeros, - tvm::ffi::TensorView g_idx, - tvm::ffi::TensorView perm, - tvm::ffi::TensorView c, - tvm::ffi::TensorView c_tmp, - tvm::ffi::TensorView a_tmp, - tvm::ffi::TensorView workspace, - int64_t b_q_type_id, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float) { - using namespace host; - - ScalarType const b_q_type = ScalarType::from_id(b_q_type_id); - int pack_factor = 32 / b_q_type.size_bits(); - - // Bind symbolic sizes - auto M = SymbolicSize{"M"}; - auto K = SymbolicSize{"K"}; - auto N = SymbolicSize{"N"}; - auto device = SymbolicDevice{}; - device.set_options(); - - // Verify a: [M, K] - auto lda = SymbolicSize{"lda"}; - TensorMatcher({M, K}).with_strides({lda, 1}).with_dtype().with_device(device).verify(a); - - int64_t size_m = M.unwrap(); - int64_t size_k = K.unwrap(); - - // Verify b_q_weight: [K/tile_size, packed_N] - RuntimeCheck( - size_k % device::marlin::tile_size == 0, - "size_k = ", - size_k, - " is not divisible by tile_size = ", - device::marlin::tile_size); - int64_t expected_bqw_dim0 = size_k / device::marlin::tile_size; - auto bqw_dim0 = SymbolicSize{"bqw_dim0"}; - auto bqw_dim1 = SymbolicSize{"bqw_dim1"}; - bqw_dim0.set_value(expected_bqw_dim0); - TensorMatcher({bqw_dim0, bqw_dim1}).with_dtype().with_device(device).verify(b_q_weight); - - RuntimeCheck( - b_q_weight.size(1) % device::marlin::tile_size == 0, - "b_q_weight.size(1) = ", - b_q_weight.size(1), - " is not divisible by tile_size = ", - device::marlin::tile_size); - int64_t actual_size_n = (b_q_weight.size(1) / device::marlin::tile_size) * pack_factor; - N.set_value(actual_size_n); - int64_t size_n = N.unwrap(); - - // Verify stride alignment - int64_t a_stride0 = a.stride(0); - RuntimeCheck(a_stride0 % 8 == 0, "a.stride(0) must be divisible by 8"); - - // Verify b_scales: [num_groups, N] - auto num_groups_sym = SymbolicSize{"num_groups"}; - TensorMatcher({num_groups_sym, N}).with_device(device).verify(b_scales); - int num_groups = static_cast(num_groups_sym.unwrap()); - - // Verify c: [M, N] - TensorMatcher({M, N}).with_dtype().with_device(device).verify(c); - - // Early return for zero-size M - if (size_m == 0) { - return; - } - - // Determine has_act_order from g_idx/perm sizes - int64_t g_idx_size = g_idx.size(0); - int64_t perm_size = perm.size(0); - bool has_act_order = g_idx_size > 0 && perm_size > 0; - - if (has_act_order) { - RuntimeCheck( - (g_idx_size == size_k && perm_size == size_k), - "Unexpected g_idx.size(0) = ", - g_idx_size, - " and perm.size(0) = ", - perm_size, - ", where size_k = ", - size_k); - } - - // Determine has_zp from b_zeros size - int64_t b_zeros_size = b_zeros.size(0); - bool has_zp = b_zeros_size > 0; - - if (has_zp) { - RuntimeCheck( - b_q_type == kU4 || b_q_type == kU8, "b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str()); - } else { - RuntimeCheck( - b_q_type == kU4B8 || b_q_type == kU8B128 || b_q_type == kFE4M3fn || b_q_type == kFE2M1f, - "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when " - "has_zp = False. Got = ", - b_q_type.str()); - } - - if (has_zp && is_zp_float) { - RuntimeCheck( - std::is_same::value, "Computation type must be float16 (half) when using float zero points."); - } - - // Verify b_zeros shape - if (has_zp) { - RuntimeCheck(b_zeros.dim() == 2, "b_zeros rank = ", b_zeros.dim(), " is not 2"); - if (is_zp_float) { - RuntimeCheck(b_zeros.size(1) == size_n, "b_zeros dim 1 = ", b_zeros.size(1), " is not size_n = ", size_n); - RuntimeCheck( - num_groups == b_zeros.size(0), "b_zeros dim 0 = ", b_zeros.size(0), " is not num_groups = ", num_groups); - RuntimeCheck(num_groups != -1, "num_groups must be != -1"); - } else { - RuntimeCheck( - b_zeros.size(0) == num_groups, "b_zeros dim 0 = ", b_zeros.size(0), " is not num_groups = ", num_groups); - RuntimeCheck( - b_zeros.size(1) == size_n / pack_factor, - "b_zeros dim 1 = ", - b_zeros.size(1), - " is not size_n / pack_factor = ", - size_n / pack_factor); - } - } - - // Verify global_scale - int64_t global_scale_size = global_scale.size(0); - if (global_scale_size > 0) { - RuntimeCheck(b_q_type == kFE2M1f, "global_scale can only be used for float4_e2m1f."); - } else { - RuntimeCheck(!(b_q_type == kFE2M1f), "the global_scale parameter must be passed for float4_e2m1f."); - } - - // Derive group_size - int group_size = -1; - if (has_act_order) { - if (is_k_full) { - RuntimeCheck(num_groups > 1, "For act_order, num_groups must be > 1"); - RuntimeCheck(size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by num_groups = ", num_groups); - group_size = static_cast(size_k / num_groups); - } else { - group_size = 0; - } - } else { - if (num_groups > 1) { - RuntimeCheck(size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by num_groups = ", num_groups); - group_size = static_cast(size_k / num_groups); - } else { - group_size = -1; - } - } - - // Verify workspace and get device info - RuntimeCheck( - size_n % device::marlin::min_thread_n == 0, - "size_n = ", - size_n, - ", is not divisible by min_thread_n = ", - device::marlin::min_thread_n); - - DLDevice dl_device = device.unwrap(); - int dev = dl_device.device_id; - cudaStream_t stream = LaunchKernel::resolve_device(dl_device); - - int sms = -1; - RuntimeDeviceCheck(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, dev)); - - RuntimeCheck( - workspace.size(0) >= sms, "workspace.size(0) = ", workspace.size(0), " is below min_workspace_size = ", sms); - - // Hardcoded defaults (auto config) - int thread_k_init = -1; - int thread_n_init = -1; - - // Compute c_tmp and a_tmp pointers - // c_tmp and a_tmp are pre-allocated by caller - - device::marlin::marlin_mm( - a.data_ptr(), - b_q_weight.data_ptr(), - c.data_ptr(), - c_tmp.data_ptr(), - b_scales.data_ptr(), - global_scale.data_ptr(), - b_zeros.data_ptr(), - g_idx.data_ptr(), - perm.data_ptr(), - a_tmp.data_ptr(), - static_cast(size_m), - static_cast(size_n), - static_cast(size_k), - static_cast(a_stride0), - workspace.data_ptr(), - b_q_type, - has_act_order, - is_k_full, - has_zp, - num_groups, - group_size, - dev, - stream, - thread_k_init, - thread_n_init, - sms, - use_atomic_add, - use_fp32_reduce, - is_zp_float); -} diff --git a/src/infiniop/ops/gptq_marlin_gemm/marlin/kernel.h b/src/infiniop/ops/gptq_marlin_gemm/marlin/kernel.h deleted file mode 100644 index 785d5e9b9..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/marlin/kernel.h +++ /dev/null @@ -1,33 +0,0 @@ - -#include "../sgl_kernel/scalar_type.hpp" - -#include "marlin.cuh" -#include "marlin_dtypes.cuh" - -#define MARLIN_KERNEL_PARAMS \ - const int4 *__restrict__ A, const int4 *__restrict__ B, int4 *__restrict__ C, int4 *__restrict__ C_tmp, \ - const int4 *__restrict__ scales_ptr, const uint16_t *__restrict__ scale2_ptr, const int4 *__restrict__ zp_ptr, \ - const int *__restrict__ g_idx, int num_groups, int prob_m, int prob_n, int prob_k, int lda, int *locks, \ - bool use_atomic_add, bool use_fp32_reduce, int max_shared_mem - -namespace device::marlin { -template < - typename scalar_t, // compute dtype, half or nv_float16 - const host::ScalarTypeId w_type_id, // weight ScalarType id - const int threads, // number of threads in a threadblock - const int thread_m_blocks, // number of 16x16 blocks in the m - // dimension (batchsize) of the - // threadblock - const int thread_n_blocks, // same for n dimension (output) - const int thread_k_blocks, // same for k dimension (reduction) - const bool m_block_size_8, // whether m_block_size == 8 - // only works when thread_m_blocks == 1 - const int stages, // number of stages for the async global->shared - // fetch pipeline - const int group_blocks, // number of consecutive 16x16 blocks - // with a separate quantization scale - const bool is_zp_float // is zero point of float16 type? - > -__global__ void Marlin(MARLIN_KERNEL_PARAMS); - -} // namespace device::marlin diff --git a/src/infiniop/ops/gptq_marlin_gemm/marlin/marlin.cuh b/src/infiniop/ops/gptq_marlin_gemm/marlin/marlin.cuh deleted file mode 100644 index 944ca3522..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/marlin/marlin.cuh +++ /dev/null @@ -1,85 +0,0 @@ -#pragma once - -#include "../sgl_kernel/utils.cuh" - -#include - -namespace device::marlin { -// Marlin params - -// 8 warps are a good choice since every SM has 4 schedulers and having more -// than 1 warp per schedule allows some more latency hiding. At the same time, -// we want relatively few warps to have many registers per warp and small tiles. -static constexpr int default_threads = 256; - -static constexpr int pipe_stages = 4; // 4 pipeline stages fit into shared memory - -static constexpr int min_thread_n = 64; -static constexpr int min_thread_k = 64; -static constexpr int max_thread_n = 256; - -static constexpr int tile_size = 16; -static constexpr int max_par = 16; - -// Repack params -static constexpr int repack_stages = 8; - -static constexpr int repack_threads = 256; - -static constexpr int tile_k_size = tile_size; -static constexpr int tile_n_size = tile_k_size * 4; - -// Helpers -template -struct Vec { - T elems[n]; - __device__ T &operator[](int i) { - return elems[i]; - } -}; - -using I4 = Vec; - -using host::div_ceil; - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 -// No support for async -#else - -__device__ inline void cp_async4_pred(void *smem_ptr, const void *glob_ptr, bool pred = true) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.cg.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), - "l"(glob_ptr), - "n"(BYTES)); -} - -__device__ inline void cp_async4(void *smem_ptr, const void *glob_ptr) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " cp.async.cg.shared.global [%0], [%1], %2;\n" - "}\n" ::"r"(smem), - "l"(glob_ptr), - "n"(BYTES)); -} - -__device__ inline void cp_async_fence() { - asm volatile("cp.async.commit_group;\n" ::); -} - -template -__device__ inline void cp_async_wait() { - asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); -} - -#endif - -} // namespace device::marlin diff --git a/src/infiniop/ops/gptq_marlin_gemm/marlin/marlin_dtypes.cuh b/src/infiniop/ops/gptq_marlin_gemm/marlin/marlin_dtypes.cuh deleted file mode 100644 index 783374ff2..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/marlin/marlin_dtypes.cuh +++ /dev/null @@ -1,78 +0,0 @@ -#ifndef _data_types_cuh -#define _data_types_cuh -#include "../sgl_kernel/utils.cuh" - -#include "marlin.cuh" - -namespace device::marlin { - -template -class ScalarType { -}; - -template <> -class ScalarType<__half> { -public: - using scalar_t = __half; - using scalar_t2 = fp16x2_t; - - // Matrix fragments for tensor core instructions; their precise layout is - // documented here: - // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type - using FragA = Vec; - using FragB = Vec; - using FragC = Vec; - using FragS = Vec; - using FragZP = Vec; - - static __device__ float inline num2float(const __half x) { - return __half2float(x); - } - - static __device__ fp16x2_t inline num2num2(const __half x) { - return __half2half2(x); - } - - static __device__ fp16x2_t inline nums2num2(const __half x1, const __half x2) { - return __halves2half2(x1, x2); - } - - static __host__ __device__ __half inline float2num(const float x) { - return __float2half(x); - } -}; - -template <> -class ScalarType<__nv_bfloat16> { -public: - using scalar_t = __nv_bfloat16; - using scalar_t2 = bf16x2_t; - - using FragA = Vec; - using FragB = Vec; - using FragC = Vec; - using FragS = Vec; - using FragZP = Vec; - -#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 - static __device__ float inline num2float(const __nv_bfloat16 x) { - return __bfloat162float(x); - } - - static __device__ bf16x2_t inline num2num2(const __nv_bfloat16 x) { - return __bfloat162bfloat162(x); - } - - static __device__ bf16x2_t inline nums2num2(const __nv_bfloat16 x1, const __nv_bfloat16 x2) { - return __halves2bfloat162(x1, x2); - } - - static __host__ __device__ __nv_bfloat16 inline float2num(const float x) { - return __float2bfloat16(x); - } -#endif -}; - -} // namespace device::marlin - -#endif diff --git a/src/infiniop/ops/gptq_marlin_gemm/marlin/marlin_template.h b/src/infiniop/ops/gptq_marlin_gemm/marlin/marlin_template.h deleted file mode 100644 index 4ea220265..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/marlin/marlin_template.h +++ /dev/null @@ -1,1645 +0,0 @@ -/* - * Modified by Neural Magic - * Copyright (C) Marlin.2024 Elias Frantar - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -/* - * Adapted from https://github.com/IST-DASLab/marlin - */ -#include "../sgl_kernel/scalar_type.hpp" - -#include "dequant.h" -#include "marlin.cuh" -#include "marlin_dtypes.cuh" - -#define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ - static_assert( \ - std::is_same::value || std::is_same::value, \ - "only float16 and bfloat16 is supported"); - -namespace device::marlin { - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - -template < - typename scalar_t, // compute dtype, half or nv_float16 - const host::ScalarTypeId w_type_id, // weight ScalarType id - const int threads, // number of threads in a threadblock - const int thread_m_blocks, // number of 16x16 blocks in the m - // dimension (batchsize) of the - // threadblock - const int thread_n_blocks, // same for n dimension (output) - const int thread_k_blocks, // same for k dimension (reduction) - const bool m_block_size_8, // whether m_block_size == 8 - // only works when thread_m_blocks == 1 - const int stages, // number of stages for the async global->shared - // fetch pipeline - const bool has_act_order, // whether act_order is enabled - const int group_blocks, // number of consecutive 16x16 blocks - // with a separate quantization scale - const bool is_zp_float // is zero point of float16 type? - > -__global__ void Marlin( - const int4 *__restrict__ A, // fp16 input matrix of shape mxk - const int4 *__restrict__ B, // 4bit quantized weight matrix of shape kxn - int4 *__restrict__ C, // fp16 output buffer of shape mxn - int4 *__restrict__ C_tmp, // fp32 tmp output buffer (for reduce) - const int4 *__restrict__ scales_ptr, // fp16 quantization scales of shape - // (k/groupsize)xn - const int *__restrict__ g_idx, // int32 group indices of shape k - int num_groups, // number of scale groups per output channel - int prob_m, // batch dimension m - int prob_n, // output dimension n - int prob_k, // reduction dimension k - int *locks, // extra global storage for barrier synchronization - bool use_fp32_reduce // whether to use fp32 global reduce -) { -} - -} // namespace device::marlin - -#else - -// m16n8k16 tensor core mma instruction with fp16 inputs and fp32 -// output/accumulation. -template -__device__ inline void -mma(const typename ScalarType::FragA &a_frag, - const typename ScalarType::FragB &frag_b, - typename ScalarType::FragC &frag_c) { - const uint32_t *a = reinterpret_cast(&a_frag); - const uint32_t *b = reinterpret_cast(&frag_b); - float *c = reinterpret_cast(&frag_c); - if constexpr (std::is_same::value) { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); - } else if constexpr (std::is_same::value) { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); - } else { - STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); - } -} - -template -__device__ inline void mma_trans( - const typename ScalarType::FragA &a_frag, - const typename ScalarType::FragB &frag_b, - const typename ScalarType::FragB &frag_b2, - typename ScalarType::FragC &frag_c) { - const uint32_t *a = reinterpret_cast(&a_frag); - const uint32_t *b = reinterpret_cast(&frag_b); - const uint32_t *b2 = reinterpret_cast(&frag_b2); - float *c = reinterpret_cast(&frag_c); - if constexpr (std::is_same::value) { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), - "r"(b2[0]), - "r"(b[1]), - "r"(b2[1]), - "r"(a[0]), - "r"(a[1]), - "f"(c[0]), - "f"(c[1]), - "f"(c[2]), - "f"(c[3])); - } else if constexpr (std::is_same::value) { - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n" - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3]) - : "r"(b[0]), - "r"(b2[0]), - "r"(b[1]), - "r"(b2[1]), - "r"(a[0]), - "r"(a[1]), - "f"(c[0]), - "f"(c[1]), - "f"(c[2]), - "f"(c[3])); - } else { - STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t); - } -} - -// Instruction for loading a full 16x16 matrix fragment of operand A from shared -// memory, directly in tensor core layout. -template -__device__ inline void ldsm(typename ScalarType::FragA &frag_a, const void *smem_ptr) { - uint32_t *a = reinterpret_cast(&frag_a); - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - if constexpr (count == 4) { - asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" - : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) - : "r"(smem)); - } else if constexpr (count == 2) { - asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" - : "=r"(a[0]), "=r"(a[1]) - : "r"(smem)); - } else if constexpr (count == 1) { - asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" - : "=r"(a[0]) - : "r"(smem)); - } else { - static_assert(count == 1 || count == 2 || count == 4, "invalid count"); - } -} - -// Multiply dequantized values by the corresponding quantization scale; used -// only for grouped quantization. -template -__device__ inline void -scale(typename ScalarType::FragB &frag_b, typename ScalarType::FragS &frag_s, int i) { - using scalar_t2 = typename ScalarType::scalar_t2; - scalar_t2 s = ScalarType::num2num2(reinterpret_cast(&frag_s)[i]); - frag_b[0] = __hmul2(frag_b[0], s); - frag_b[1] = __hmul2(frag_b[1], s); -} - -template -__device__ inline void scale_and_sub(typename ScalarType::FragB &frag_b, scalar_t s, scalar_t zp) { - using scalar_t2 = typename ScalarType::scalar_t2; - scalar_t2 s2 = ScalarType::num2num2(s); - scalar_t2 zp2 = ScalarType::num2num2(zp); - frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2)); - frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2)); -} - -template -__device__ inline void -sub_zp(typename ScalarType::FragB &frag_b, typename ScalarType::scalar_t2 &frag_zp, int i) { - using scalar_t2 = typename ScalarType::scalar_t2; - scalar_t2 zp = ScalarType::num2num2(reinterpret_cast(&frag_zp)[i]); - frag_b[0] = __hsub2(frag_b[0], zp); - frag_b[1] = __hsub2(frag_b[1], zp); -} - -// Same as above, but for act_order (each K is multiplied individually) -template -__device__ inline void scale4( - typename ScalarType::FragB &frag_b, - typename ScalarType::FragS &frag_s_1, - typename ScalarType::FragS &frag_s_2, - typename ScalarType::FragS &frag_s_3, - typename ScalarType::FragS &frag_s_4, - int i) { - using scalar_t2 = typename ScalarType::scalar_t2; - scalar_t2 s_val_1_2; - s_val_1_2.x = reinterpret_cast(&frag_s_1)[i]; - s_val_1_2.y = reinterpret_cast(&frag_s_2)[i]; - - scalar_t2 s_val_3_4; - s_val_3_4.x = reinterpret_cast(&frag_s_3)[i]; - s_val_3_4.y = reinterpret_cast(&frag_s_4)[i]; - - frag_b[0] = __hmul2(frag_b[0], s_val_1_2); - frag_b[1] = __hmul2(frag_b[1], s_val_3_4); -} - -// Given 2 floats multiply by 2 scales (halves) -template -__device__ inline void scale_float(float *c, typename ScalarType::FragS &s) { - scalar_t *s_ptr = reinterpret_cast(&s); - c[0] = __fmul_rn(c[0], ScalarType::num2float(s_ptr[0])); - c[1] = __fmul_rn(c[1], ScalarType::num2float(s_ptr[1])); -} - -// Wait until barrier reaches `count`, then lock for current threadblock. -__device__ inline void barrier_acquire(int *lock, int count) { - if (threadIdx.x == 0) { - int state = -1; - do { - // Guarantee that subsequent writes by this threadblock will be visible - // globally. - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" - : "=r"(state) - : "l"(lock)); - } while (state != count); - } - __syncthreads(); -} - -// Release barrier and increment visitation count. -__device__ inline void barrier_release(int *lock, bool reset = false) { - __syncthreads(); - if (threadIdx.x == 0) { - if (reset) { - lock[0] = 0; - return; - } - int val = 1; - // Make sure that all writes since acquiring this barrier are visible - // globally, while releasing the barrier. - asm volatile("fence.acq_rel.gpu;\n"); - asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" - : - : "l"(lock), "r"(val)); - } -} - -// Wait until value of lock to be negative, and then add 1 -__device__ inline void wait_negative_and_add(int *lock) { - if (threadIdx.x == 0) { - int state = 0; - do { - // Guarantee that subsequent writes by this threadblock will be visible - // globally. - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" - : "=r"(state) - : "l"(lock)); - } while (state >= 0); - atomicAdd(lock, 1); - } - __syncthreads(); -} - -template < - typename scalar_t, // compute dtype, half or nv_float16 - const host::ScalarTypeId w_type_id, // weight ScalarType id - const int threads, // number of threads in a threadblock - const int thread_m_blocks, // number of 16x16 blocks in the m - // dimension (batchsize) of the - // threadblock - const int thread_n_blocks, // same for n dimension (output) - const int thread_k_blocks, // same for k dimension (reduction) - const bool m_block_size_8, // whether m_block_size == 8 - // only works when thread_m_blocks == 1 - const int stages, // number of stages for the async global->shared - // fetch pipeline - const int group_blocks, // number of consecutive 16x16 blocks - // with a separate quantization scale - const bool is_zp_float // is zero point of float16 type? - > -__global__ void Marlin( - const int4 *__restrict__ A, // fp16 input matrix of shape mxk - const int4 *__restrict__ B, // 4bit quantized weight matrix of shape kxn - int4 *__restrict__ C, // fp16 output buffer of shape mxn - int4 *__restrict__ C_tmp, // fp32 tmp output buffer (for reduce) - const int4 *__restrict__ scales_ptr, // fp16 quantization scales of shape - // (k/groupsize)xn - const uint16_t *__restrict__ scale2_ptr, // fp16 global scale (for nvfp4 - // only) - const int4 *__restrict__ zp_ptr, // 4bit packed zero-points of shape - // (k/groupsize)x(n/pack_factor) - const int *__restrict__ g_idx, // int32 group indices of shape k - int num_groups, // number of scale groups per output channel - int prob_m, // batch dimension m - int prob_n, // output dimension n - int prob_k, // reduction dimension k - int lda, // A.stride(0), equal to prob_k is A is contiguous - int *locks, // extra global storage for barrier synchronization - bool use_atomic_add, // whether to use atomic add to reduce - bool use_fp32_reduce, // whether to use fp32 global reduce - int max_shared_mem) { - // Each threadblock processes one "stripe" of the B matrix with (roughly) the - // same size, which might involve multiple column "slices" (of width 16 * - // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM - // example: - // 0 1 3 - // 0 2 3 - // 1 2 4 - // While this kind of partitioning makes things somewhat more complicated, it - // ensures good utilization of all SMs for many kinds of shape and GPU - // configurations, while requiring as few slow global cross-threadblock - // reductions as possible. - using Dtype = ScalarType; - using scalar_t2 = typename ScalarType::scalar_t2; - using FragA = typename ScalarType::FragA; - using FragB = typename ScalarType::FragB; - using FragC = typename ScalarType::FragC; - using FragS = typename ScalarType::FragS; - using FragZP = typename ScalarType::FragZP; - - static constexpr auto w_type = host::ScalarType::from_id(w_type_id); - constexpr bool has_zp = w_type == host::kU4 || w_type == host::kU8; - constexpr bool is_int_type = w_type == host::kU4 || w_type == host::kU8 || w_type == host::kU4B8 || w_type == host::kU8B128; - // see comments of dequant.h for more details - constexpr bool dequant_skip_flop = !is_int_type || has_zp && !is_zp_float && !std::is_same::value || has_zp && !is_zp_float && !(w_type == host::kU8); - - scalar_t2 global_scale; - - if constexpr (w_type == host::kFE2M1f) { - uint16_t val = scale2_ptr[0]; - global_scale = Dtype::num2num2(*reinterpret_cast(&val)); - } - - constexpr bool has_act_order = group_blocks == 0; - constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); - - constexpr int pack_factor = 32 / w_type.size_bits(); - static_assert(thread_m_blocks == 1 || !m_block_size_8); - - // For larger GEMMs we run multiple batchsize 64 versions in parallel for a - // better partitioning with less reductions - int parallel = 1; - if (prob_m > m_block_size) { - parallel = prob_m / m_block_size; - prob_m = m_block_size; - } - - int k_tiles = prob_k / 16 / thread_k_blocks; - int n_tiles = prob_n / 16 / thread_n_blocks; - int iters = div_ceil(k_tiles * n_tiles * parallel, gridDim.x); - - if constexpr (!has_act_order && group_blocks != -1) { - if (group_blocks >= thread_k_blocks) { - // Ensure that the number of tiles in each stripe is a multiple of the - // groupsize; this avoids an annoying special case where a stripe starts - // in the middle of group. - iters = (group_blocks / thread_k_blocks) * div_ceil(iters, (group_blocks / thread_k_blocks)); - } - } - - int slice_row = (iters * blockIdx.x) % k_tiles; - int slice_col_par = (iters * blockIdx.x) / k_tiles; - int slice_col = slice_col_par; - int slice_iters; // number of threadblock tiles in the current slice - int slice_count = 0; // total number of active threadblocks in the current slice - int slice_idx; // index of threadblock in current slice; numbered bottom to - // top - - int par_id = 0; - int locks_off = 0; - - // We can easily implement parallel problem execution by just remapping - // indices and advancing global pointers - if (slice_col_par >= n_tiles) { - A += (slice_col_par / n_tiles) * 16 * thread_m_blocks * lda / 8; - C += (slice_col_par / n_tiles) * 16 * thread_m_blocks * prob_n / 8; - slice_col = slice_col_par % n_tiles; - par_id = slice_col_par / n_tiles; - } - if (parallel * n_tiles >= gridDim.x) { - // when parallel * n_tiles >= sms - // then there are at most $sms$ conflict tile blocks - locks_off = blockIdx.x; - } else { - locks_off = (iters * blockIdx.x) / k_tiles - 1; - } - - // Compute all information about the current slice which is required for - // synchronization. - auto init_slice = [&](bool first_init = false) { - slice_iters = iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); - if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) { - slice_iters = 0; - } - if (slice_iters == 0) { - return; - } - if (slice_row + slice_iters > k_tiles) { - slice_iters = k_tiles - slice_row; - } - slice_count = 1; - slice_idx = 0; - int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); - if (col_first <= k_tiles * (slice_col_par + 1)) { - int col_off = col_first - k_tiles * slice_col_par; - slice_count = div_ceil(k_tiles - col_off, iters); - if (col_off > 0) { - slice_count++; - } - int delta_first = iters * blockIdx.x - col_first; - if (delta_first < 0 || (col_off == 0 && delta_first == 0)) { - slice_idx = slice_count - 1; - } else { - slice_idx = slice_count - 1 - delta_first / iters; - if (col_off > 0) { - slice_idx--; - } - } - } - if (parallel * n_tiles >= gridDim.x) { - if (slice_count > 1 && slice_idx == slice_count - 1) { - locks_off++; - } - } else { - locks_off++; - } - - if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) { - constexpr int threads_per_m = 16 * thread_n_blocks / 8; - int m_per_thread = div_ceil(thread_m_blocks * 16, threads / threads_per_m); - if (m_block_size_8) { - m_per_thread = div_ceil(8, threads / threads_per_m); - } - for (int i = 0; i < m_per_thread; i++) { - int row = threads / threads_per_m * i + threadIdx.x / threads_per_m; - if (row < prob_m) { - int col = slice_col * 16 * thread_n_blocks / 8 + threadIdx.x % threads_per_m; - C[row * prob_n / 8 + col] = {0, 0, 0, 0}; - } - } - // After write zero to output, write a negative value to lock. - // Every SM that processes the same slice would wait for - // the negative value, and then atomicAdd 1 to it. - // After all SMs are processed, the lock value would back to 0 again. - __syncthreads(); - if (threadIdx.x == 0) { - locks[locks_off] = 1 - slice_count; - } - } - - if (slice_col == n_tiles) { - A += 16 * thread_m_blocks * lda / 8; - C += 16 * thread_m_blocks * prob_n / 8; - slice_col = 0; - par_id++; - } - }; - init_slice(true); - - // A sizes/strides - - // stride of the A matrix in global memory - int a_gl_stride = lda / 8; - // stride of an A matrix tile in shared memory - constexpr int a_sh_stride = 16 * thread_k_blocks / 8; - // delta between subsequent A tiles in global memory - constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8; - // between subsequent accesses within a tile - int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); - // between shared memory writes - constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); - // between shared memory tile reads - constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4)); - // within a shared memory tile - constexpr int a_sh_rd_delta_i = a_sh_stride * 16; - // overall size of a tile - constexpr int a_sh_stage = a_sh_stride * m_block_size; - // number of shared write iterations for a tile - constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); - - // B sizes/strides - int b_gl_stride = 16 * prob_n / (pack_factor * 4); - constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4; - constexpr int b_thread_vecs = w_type.size_bits() == 4 ? 1 : 2; - constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; - - int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks; - int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads); - constexpr int b_sh_wr_delta = threads * b_thread_vecs; - constexpr int b_sh_rd_delta = threads * b_thread_vecs; - constexpr int b_sh_stage = b_sh_stride * thread_k_blocks; - constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; - - // Scale sizes/strides without act_order - int s_gl_stride = prob_n / 8; - constexpr int s_sh_stride = 16 * thread_n_blocks / 8; - constexpr int s_tb_groups = !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks - ? thread_k_blocks / group_blocks / (w_type == host::kFE2M1f ? 2 : 1) - : 1; - constexpr int s_sh_stage = s_tb_groups * s_sh_stride; - int s_gl_rd_delta = s_gl_stride; - - // Scale size/strides with act_order - constexpr int tb_k = 16 * thread_k_blocks; - constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0; - // constexpr int act_s_row_stride = 1; - // int act_s_col_stride = act_s_row_stride * num_groups; - constexpr int act_s_max_num_groups = 32; - int act_s_col_stride = 1; - int act_s_col_warp_stride = act_s_col_stride * 8; - - int tb_n_warps = thread_n_blocks / 4; - int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps; - - // Zero-points sizes/strides - int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4; - constexpr int zp_sh_stride = is_zp_float ? 16 * thread_n_blocks / 8 : ((16 * thread_n_blocks) / pack_factor) / 4; - constexpr int zp_tb_groups = s_tb_groups; - constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0; - int zp_gl_rd_delta = zp_gl_stride; - - // Global A read index of current thread. - int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - a_gl_rd += a_gl_rd_delta_o * slice_row; - // Shared write index of current thread. - int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); - // Shared read index. - int a_sh_rd = a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) + (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1)); - a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4)); - - int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) + (threadIdx.x % b_sh_stride_threads) * b_thread_vecs; - b_gl_rd += b_sh_stride * slice_col; - b_gl_rd += b_gl_rd_delta_o * slice_row; - auto b_sh_wr = threadIdx.x * b_thread_vecs; - auto b_sh_rd = threadIdx.x * b_thread_vecs; - - // For act_order - constexpr int k_iter_size = tb_k / b_sh_wr_iters; - int slice_k_start = tb_k * slice_row; - int slice_k_finish = slice_k_start + tb_k * slice_iters; - int slice_k_start_shared_fetch = slice_k_start; - int slice_n_offset = act_s_col_tb_stride * slice_col; - - // No act_order - int s_gl_rd; - if constexpr (!has_act_order) { - if constexpr (group_blocks == -1) { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - } else { - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) / (w_type == host::kFE2M1f ? 2 : 1) + s_sh_stride * slice_col + threadIdx.x; - } - } - auto s_sh_wr = threadIdx.x; - bool s_sh_wr_pred = threadIdx.x < s_sh_stride; - - // Zero-points - int zp_gl_rd; - if constexpr (has_zp) { - if constexpr (group_blocks == -1) { - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; - } else { - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + zp_sh_stride * slice_col + threadIdx.x; - } - } - auto zp_sh_wr = threadIdx.x; - bool zp_sh_wr_pred = threadIdx.x < zp_sh_stride; - - // We use a different scale layout for grouped and column-wise quantization as - // we scale a `half2` tile in column-major layout in the former and in - // row-major in the latter case. - int s_sh_rd; - if constexpr (group_blocks != -1 && w_type == host::kFE2M1f) { - auto warp_id = threadIdx.x / 32; - int n_warps = thread_n_blocks / 4; - int warp_row = warp_id / n_warps; - - s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 4; - s_sh_rd = s_sh_rd * 2 + warp_row % 2; - } else if constexpr (group_blocks != -1) { - s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 4; - } else if constexpr (group_blocks == -1 && (m_block_size_8 || (has_zp && !dequant_skip_flop))) { - s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 8; - } else { - s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) % 4; - } - - // Zero-points have the same read layout as the scales - // (without column-wise case) - constexpr int num_col_threads = 8; - constexpr int num_row_threads = 4; - constexpr int num_ints_per_thread = 8 / pack_factor; - int zp_sh_rd; - if constexpr (has_zp) { - if constexpr (is_zp_float) { - if constexpr (group_blocks != -1) { - zp_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + (threadIdx.x % 32) / 4; - } - } else { - zp_sh_rd = num_ints_per_thread * num_col_threads * ((threadIdx.x / 32) % (thread_n_blocks / 4)) + num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); - } - } - - // Precompute which thread should not read memory in which iterations; this is - // needed if there are more threads than required for a certain tilesize or - // when the batchsize is not a multiple of 16. - bool a_sh_wr_pred[a_sh_wr_iters]; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) { - a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m; - } - - // To ensure that writing and reading A tiles to/from shared memory, the - // latter in fragment format, is fully bank conflict free, we need to use a - // rather fancy XOR-based layout. The key here is that neither reads nor - // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the - // same shared memory banks. Further, it seems (based on NSight-Compute) that - // each warp must also write a consecutive memory segment? - auto transform_a = [&](int i) { - int row = i / a_gl_rd_delta_o; - return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8); - }; - // Since the computation of this remapping is non-trivial and, due to our main - // loop unrolls, all shared memory accesses are static, we simply precompute - // both transformed reads and writes. - int a_sh_wr_trans[a_sh_wr_iters]; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) { - a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); - } - int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; -#pragma unroll - for (int i = 0; i < b_sh_wr_iters; i++) { -#pragma unroll - for (int j = 0; j < thread_m_blocks; j++) { - a_sh_rd_trans[i][j] = transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd); - } - } - - // Since B-accesses have non-constant stride they have to be computed at - // runtime; we break dependencies between subsequent accesses with a tile by - // maintining multiple pointers (we have enough registers), a tiny - // optimization. - const int4 *B_ptr[b_sh_wr_iters]; -#pragma unroll - for (int i = 0; i < b_sh_wr_iters; i++) { - B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd; - } - - extern __shared__ int4 sh[]; - // Shared memory storage for global fetch pipelines. - constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks; - constexpr int sh_b_size = stages * b_sh_stage; - int4 *sh_b = sh; - int4 *sh_red = sh; - int4 *sh_g_idx = sh_b + (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); - int4 *sh_zp = sh_g_idx + (stages * g_idx_stage); - constexpr int sh_s_size = has_act_order ? (act_s_max_num_groups * s_sh_stride) : (stages * s_sh_stage); - int4 *sh_s = sh_zp + (stages * zp_sh_stage); - // shared memory reused by reduction should be smaller than - // shared memory used by weight. - static_assert(thread_m_blocks * 16 * thread_n_blocks * 16 / 8 <= stages * b_sh_stage); - int4 *sh_a = sh_s + sh_s_size; - // constexpr int shm_size_used = - // stages * (g_idx_stage + zp_sh_stage) + sh_s_size + - // (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); - - // Register storage for double buffer of shared memory reads. - FragA frag_a[2][thread_m_blocks]; - I4 frag_b_quant[2][b_thread_vecs]; - FragC frag_c[thread_m_blocks][4][2]; - FragS frag_s[2][4]; // No act-order - FragS act_frag_s[2][4][4]; // For act-order - int frag_qzp[2][num_ints_per_thread]; // Zero-points - FragZP frag_zp; // Zero-points in fp16 - FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ - - // Zero accumulators. - auto zero_accums = [&]() { -#pragma unroll - for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) - reinterpret_cast(frag_c)[i] = 0; - }; - - int sh_first_group_id = -1; - int sh_num_groups = -1; - - auto fetch_act_order_scales_to_shared = [&](bool is_async, int first_group_id, int last_group_id) { - sh_first_group_id = first_group_id; - sh_num_groups = last_group_id - first_group_id + 1; - - if (sh_num_groups > act_s_max_num_groups) { - sh_num_groups = act_s_max_num_groups; - } - - if (sh_first_group_id + sh_num_groups > num_groups) { - sh_num_groups = num_groups - sh_first_group_id; - } - - int row_offset = first_group_id * s_gl_stride; - - if (is_async) { - for (int i = 0; i < sh_num_groups; i++) { - if (threadIdx.x < s_sh_stride) { - cp_async4_pred( - &sh_s[(i * s_sh_stride) + threadIdx.x], - &scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]); - } - } - } else { - for (int i = 0; i < sh_num_groups; i++) { - if (threadIdx.x < s_sh_stride) { - sh_s[(i * s_sh_stride) + threadIdx.x] = scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset + threadIdx.x]; - } - } - } - }; - // Asynchronously fetch the next A, B and s tile from global to the next - // shared memory pipeline location. - auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { - if (pred) { - int4 *sh_a_stage = sh_a + a_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < a_sh_wr_iters; i++) { - cp_async4_pred( - &sh_a_stage[a_sh_wr_trans[i]], - &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off], - a_sh_wr_pred[i]); - } - int4 *sh_b_stage = sh_b + b_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < b_sh_wr_iters; i++) { -#pragma unroll - for (int j = 0; j < b_thread_vecs; j++) { - cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j], B_ptr[i] + j); - } - - B_ptr[i] += b_gl_rd_delta_o; - } - - if constexpr (has_act_order) { - // Fetch g_idx thread-block portion - int full_pipe = a_off; - int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe; - if (cur_k < prob_k && cur_k < slice_k_finish) { - int4 *sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - - int4 const *cur_g_idx_stage_ptr = reinterpret_cast(&g_idx[cur_k]); - - if (threadIdx.x < g_idx_stage) { - cp_async4_pred(&sh_g_idx_stage[threadIdx.x], &cur_g_idx_stage_ptr[threadIdx.x]); - } - } - } else { - if constexpr (group_blocks != -1) { - int4 *sh_s_stage = sh_s + s_sh_stage * pipe; - - if constexpr (group_blocks >= thread_k_blocks) { - // Only fetch scales if this tile starts a new group - if (pipe % (group_blocks / thread_k_blocks) == 0) { - if (s_sh_wr_pred) { - cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); - } - s_gl_rd += s_gl_rd_delta; - } - } else { - for (int i = 0; i < s_tb_groups; i++) { - if (s_sh_wr_pred) { - cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr], &scales_ptr[s_gl_rd]); - } - s_gl_rd += s_gl_rd_delta; - } - } - } - - if constexpr (has_zp && group_blocks != -1) { - int4 *sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - if constexpr (group_blocks >= thread_k_blocks) { - // Only fetch zero-points if this tile starts a new group - if (pipe % (group_blocks / thread_k_blocks) == 0) { - if (zp_sh_wr_pred) { - cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - zp_gl_rd += zp_gl_rd_delta; - } - } else { - for (int i = 0; i < zp_tb_groups; i++) { - if (zp_sh_wr_pred) { - cp_async4(&sh_zp_stage[i * zp_sh_stride + zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - zp_gl_rd += zp_gl_rd_delta; - } - } - } - } - } - // Insert a fence even when we are winding down the pipeline to ensure that - // waiting is also correct at this point. - cp_async_fence(); - }; - - auto fetch_col_zp_to_shared = [&]() { - if (zp_sh_wr_pred) { - cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]); - } - }; - - auto fetch_col_scale_to_shared = [&]() { - if (s_sh_wr_pred) { - cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); - } - }; - - // Wait until the next thread tile has been loaded to shared memory. - auto wait_for_stage = [&]() { - // We only have `stages - 2` active fetches since we are double buffering - // and can only issue the next fetch when it is guaranteed that the previous - // shared memory load is fully complete (as it may otherwise be - // overwritten). - cp_async_wait(); - __syncthreads(); - }; - - // Load the next sub-tile from the current location in the shared memory pipe - // into the current register buffer. - auto fetch_to_registers = [&](int k, int pipe) { - int4 *sh_a_stage = sh_a + a_sh_stage * pipe; -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) - ldsm(frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); - int4 *sh_b_stage = sh_b + b_sh_stage * pipe; - -#pragma unroll - for (int i = 0; i < b_thread_vecs; i++) { - frag_b_quant[k % 2][i] = *reinterpret_cast(&sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]); - } - }; - - bool is_same_group[stages]; - int same_group_id[stages]; - - auto init_same_group = [&](int pipe) { - if constexpr (!has_act_order) { - return; - } - - int4 *sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int *sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - int group_id_1 = sh_g_idx_int_ptr[0]; - int group_id_2 = sh_g_idx_int_ptr[tb_k - 1]; - - is_same_group[pipe] = group_id_1 == group_id_2; - same_group_id[pipe] = group_id_1; - }; - - auto fetch_scales_to_registers = [&](int k, int full_pipe) { - int pipe = full_pipe % stages; - - if constexpr (!has_act_order) { - // No act-order case - if constexpr (group_blocks == -1) { - // load only when starting a new slice - if (k == 0 && full_pipe == 0) { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd]; - reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; - } - } else if constexpr (group_blocks != -1) { - if constexpr (group_blocks >= thread_k_blocks) { - if (k % b_sh_wr_iters == 0) { - int4 *sh_s_stage = sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); - reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; - } else { - reinterpret_cast(&frag_s[1])[0] = reinterpret_cast(&frag_s[0])[0]; - } - } else { - auto warp_id = threadIdx.x / 32; - int n_warps = thread_n_blocks / 4; - - int warp_row = warp_id / n_warps; - - int cur_k = warp_row * 16; - cur_k += k_iter_size * (k % b_sh_wr_iters); - - int k_blocks = cur_k / 16; - int cur_group_id = k_blocks / (group_blocks * (w_type == host::kFE2M1f ? 2 : 1)); - - int4 *sh_s_stage = sh_s + s_sh_stage * pipe; - - if constexpr (w_type_id != host::kFE2M1f.id()) { - reinterpret_cast(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; - } else { - reinterpret_cast(&frag_s[k % 2])[0] = reinterpret_cast(sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; - } - } - } - - return; - } - - // Act-order case - - // Determine K of the "current" thread-block - int cur_k = slice_k_start + tb_k * full_pipe; - if (cur_k >= prob_k || cur_k >= slice_k_finish) { - return; - } - - // Reset (to current thread-block) since we read g_idx portion from the - // shared memory - cur_k = 0; - - // Progress to current iteration - cur_k += k_iter_size * (k % b_sh_wr_iters); - - // Determine "position" inside the thread-block (based on warp and - // thread-id) - auto warp_id = threadIdx.x / 32; - int n_warps = thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N - - int warp_row = warp_id / n_warps; - int warp_col = warp_id % n_warps; - - cur_k += warp_row * 16; - - auto th_id = threadIdx.x % 32; - cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix - - int s_col_shift = - /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) + (th_id / 4) * act_s_col_stride; - - if (is_same_group[pipe]) { - if (k % 2 == 0) { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride + s_col_shift]; - } else { - *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))) = *(reinterpret_cast(&(act_frag_s[(k - 1) % 2][0][0]))); - } - - for (int i = 1; i < 4; i++) { - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = *(reinterpret_cast(&(act_frag_s[k % 2][0][0]))); - } - return; - } - - int4 *sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe; - int *sh_g_idx_int_ptr = reinterpret_cast(sh_g_idx_stage); - - constexpr int k_frag_offsets[4] = {0, 1, 8, 9}; // Tensor core offsets per thread - -#pragma unroll - for (int i = 0; i < 4; i++) { - int actual_k = cur_k + k_frag_offsets[i]; - - int group_id = sh_g_idx_int_ptr[actual_k]; - int rel_group_id = group_id - sh_first_group_id; - - *(reinterpret_cast(&(act_frag_s[k % 2][i][0]))) = sh_s[rel_group_id * s_sh_stride + s_col_shift]; - } - }; - - auto fetch_zp_to_registers = [&](int k, int full_pipe) { - // This code does not handle group_blocks == 0, - // which signifies act_order. - // has_zp implies AWQ, which doesn't have act_order, - static_assert(!has_zp || group_blocks != 0); - - if constexpr (has_zp && !is_zp_float) { - int pipe = full_pipe % stages; - - if constexpr (group_blocks == -1) { - // load only when starting a new slice - if (k == 0 && full_pipe == 0) { -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp))[zp_sh_rd + i]; - } - } - } else if constexpr (group_blocks >= thread_k_blocks) { - if (k % b_sh_wr_iters == 0) { - int4 *sh_zp_stage = sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; - } - } - } else { - auto warp_id = threadIdx.x / 32; - int n_warps = thread_n_blocks / 4; - - int warp_row = warp_id / n_warps; - - int cur_k = warp_row * 16; - cur_k += k_iter_size * (k % b_sh_wr_iters); - - int k_blocks = cur_k / 16; - int cur_group_id = 0; - - // Suppress bogus and persistent divide-by-zero warning -#pragma nv_diagnostic push -#pragma nv_diag_suppress divide_by_zero - cur_group_id = k_blocks / group_blocks; -#pragma nv_diagnostic pop - - int4 *sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - sh_zp_stage += cur_group_id * zp_sh_stride; - -#pragma unroll - for (int i = 0; i < num_ints_per_thread; i++) { - frag_qzp[k % 2][i] = (reinterpret_cast(sh_zp_stage))[zp_sh_rd + i]; - } - } - } - - else if constexpr (has_zp && is_zp_float) { - int pipe = full_pipe % stages; - - if constexpr (group_blocks != -1) { - if constexpr (group_blocks >= thread_k_blocks) { - if (k % b_sh_wr_iters == 0) { - int4 *sh_zp_stage = sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) * (pipe / (group_blocks / thread_k_blocks))); - reinterpret_cast(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd]; - } - } else { - auto warp_id = threadIdx.x / 32; - int n_warps = thread_n_blocks / 4; - - int warp_row = warp_id / n_warps; - - int cur_k = warp_row * 16; - cur_k += k_iter_size * (k % b_sh_wr_iters); - - int k_blocks = cur_k / 16; - // Suppress bogus and persistent divide-by-zero warning -#pragma nv_diagnostic push -#pragma nv_diag_suppress divide_by_zero - int cur_group_id = k_blocks / group_blocks; -#pragma nv_diagnostic pop - - int4 *sh_zp_stage = sh_zp + zp_sh_stage * pipe; - - reinterpret_cast(&frag_zpf[k % 2])[0] = sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride]; - } - } - } - }; - - auto dequant_data = [&](int q, scalar_t2 *frag_b_ptr) { - dequant(q, frag_b_ptr); - }; - - // Execute the actual tensor core matmul of a sub-tile. - bool is_first_matmul_in_slice = true; - auto matmul = [&](int k) { - int k2 = k % 2; - const bool is_new_zp = ((group_blocks != -1) && (group_blocks < thread_k_blocks || k == 0)) || (group_blocks == -1 && is_first_matmul_in_slice); - if constexpr (has_zp && !is_zp_float) { - if (is_new_zp) { - if constexpr (group_blocks == -1) { - is_first_matmul_in_slice = false; - } - FragB frag_zp_0; - FragB frag_zp_1; - int zp_quant_0, zp_quant_1; - - if constexpr (w_type.size_bits() == 4) { - zp_quant_0 = frag_qzp[k2][0]; - zp_quant_1 = zp_quant_0 >> 8; - } else { - static_assert(w_type.size_bits() == 8); - zp_quant_0 = frag_qzp[k2][0]; - zp_quant_1 = frag_qzp[k2][1]; - } - - dequant_data(zp_quant_0, reinterpret_cast(&frag_zp)); - dequant_data(zp_quant_1, reinterpret_cast(&frag_zp) + 2); - } - } - if constexpr (!dequant_skip_flop && has_zp && is_zp_float) { - if (is_new_zp) { - reinterpret_cast(&frag_zp)[0] = reinterpret_cast(&frag_zpf[k2])[0]; - } - } - - if constexpr (w_type == host::kFE2M1f) { - int s_quant_0 = reinterpret_cast(frag_s[k2])[0]; - int s_quant_1 = reinterpret_cast(frag_s[k2])[1]; - - dequant_fp8_scales(s_quant_0, reinterpret_cast(&frag_s[k2])); - dequant_fp8_scales(s_quant_1, reinterpret_cast(&frag_s[k2]) + 2); - } - -// We have the m dimension as the inner loop in order to encourage overlapping -// dequantization and matmul operations. -#pragma unroll - for (int j = 0; j < 4; j++) { - FragB frag_b0; - FragB frag_b1; - int b_quant_0, b_quant_1; - - if constexpr (w_type_id == host::kFE2M1f.id()) { - b_quant_1 = frag_b_quant[k2][0][j]; - b_quant_0 = b_quant_1 << 8; - } else if constexpr (w_type.size_bits() == 4) { - b_quant_0 = frag_b_quant[k2][0][j]; - b_quant_1 = b_quant_0 >> 8; - } else { - static_assert(w_type.size_bits() == 8); - int *frag_b_quant_ptr = reinterpret_cast(frag_b_quant[k2]); - b_quant_0 = frag_b_quant_ptr[j * 2 + 0]; - b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; - } - - dequant_data(b_quant_0, reinterpret_cast(&frag_b0)); - dequant_data(b_quant_1, reinterpret_cast(&frag_b1)); - - if constexpr (dequant_skip_flop && has_zp && !is_zp_float) { - sub_zp(frag_b0, frag_zp[j], 0); - sub_zp(frag_b1, frag_zp[j], 1); - } - - // Apply scale to frag_b0 - if constexpr (has_act_order) { - static_assert(group_blocks != -1); - scale4( - frag_b0, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k2][2][j], act_frag_s[k2][3][j], 0); - scale4( - frag_b1, act_frag_s[k2][0][j], act_frag_s[k2][1][j], act_frag_s[k2][2][j], act_frag_s[k2][3][j], 1); - } else if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && group_blocks == -1) { - int idx = (threadIdx.x / 4) % 2; - scalar_t2 s2 = Dtype::nums2num2( - reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 0])[idx], - reinterpret_cast(&frag_s[j / 2][j % 2 * 2 + 1])[idx]); - if (is_new_zp) - frag_zp[j] = __hmul2(frag_zp[j], s2); - scale_and_sub(frag_b0, s2.x, frag_zp[j].x); - scale_and_sub(frag_b1, s2.y, frag_zp[j].y); - } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1) { - if (is_new_zp) - frag_zp[j] = __hmul2(frag_zp[j], *reinterpret_cast(&frag_s[k2][j])); - scale_and_sub(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x); - scale_and_sub(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y); - } else if constexpr (group_blocks != -1) { - scale(frag_b0, frag_s[k2][j], 0); - scale(frag_b1, frag_s[k2][j], 1); - } - -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { - if constexpr (m_block_size_8) { - mma_trans(frag_a[k2][i], frag_b0, frag_b1, frag_c[i][j][0]); - } else { - mma(frag_a[k2][i], frag_b0, frag_c[i][j][0]); - mma(frag_a[k2][i], frag_b1, frag_c[i][j][1]); - } - } - } - }; - - // Since we slice across the k dimension of a tile in order to increase the - // number of warps while keeping the n dimension of a tile reasonable, we have - // multiple warps that accumulate their partial sums of the same output - // location; which we have to reduce over in the end. We do in shared memory. - auto thread_block_reduce = [&]() { - constexpr int red_off = threads / b_sh_stride_threads / 2; - if (red_off >= 1) { - auto red_idx = threadIdx.x / b_sh_stride_threads; - constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2; - constexpr int red_sh_delta = b_sh_stride_threads; - int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + (threadIdx.x % b_sh_stride_threads); - - // Parallel logarithmic shared memory reduction. We make sure to avoid any - // unnecessary read or write iterations, e.g., for two warps we write only - // once by warp 1 and read only once by warp 0. - -#pragma unroll - for (int m_block = 0; m_block < thread_m_blocks; m_block++) { -#pragma unroll - for (int i = red_off; i > 0; i /= 2) { - if (i <= red_idx && red_idx < 2 * i) { -#pragma unroll - for (int j = 0; j < 4 * 2; j += (m_block_size_8 ? 2 : 1)) { - int red_sh_wr = red_sh_delta * j + (red_sh_rd - red_sh_stride * i); - if (i < red_off) { - float *c_rd = reinterpret_cast(&sh_red[red_sh_delta * j + red_sh_rd]); - float *c_wr = reinterpret_cast(&sh_red[red_sh_wr]); -#pragma unroll - for (int k = 0; k < 4; k++) - reinterpret_cast(frag_c)[4 * 2 * m_block + j][k] += c_rd[k] + c_wr[k]; - } - sh_red[red_sh_wr] = reinterpret_cast(&frag_c)[4 * 2 * m_block + j]; - } - } - __syncthreads(); - } - if (red_idx == 0) { -#pragma unroll - for (int i = 0; i < 4 * 2; i += (m_block_size_8 ? 2 : 1)) { - float *c_rd = reinterpret_cast(&sh_red[red_sh_delta * i + red_sh_rd]); -#pragma unroll - for (int j = 0; j < 4; j++) - reinterpret_cast(frag_c)[4 * 2 * m_block + i][j] += c_rd[j]; - } - } - __syncthreads(); - } - } - }; - - // Since multiple threadblocks may process parts of the same column slice, we - // finally have to globally reduce over the results. As the striped - // partitioning minimizes the number of such reductions and our outputs are - // usually rather small, we perform this reduction serially in L2 cache. - auto global_reduce_fp16 = [&](bool first = false, bool last = false) { - // We are very careful here to reduce directly in the output buffer to - // maximize L2 cache utilization in this step. To do this, we write out - // results in FP16 (but still reduce with FP32 compute). - constexpr int active_threads = 32 * thread_n_blocks / 4; - if (threadIdx.x < active_threads) { - int c_gl_stride = prob_n / 8; - int c_gl_wr_delta_o = 8 * c_gl_stride; - int c_gl_wr_delta_i = 4 * (active_threads / 32); - int c_gl_wr; - if constexpr (m_block_size_8) { - c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + 4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8; - c_gl_wr += (2 * thread_n_blocks) * slice_col; - } else { - c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) + 4 * (threadIdx.x / 32) + threadIdx.x % 4; - c_gl_wr += (2 * thread_n_blocks) * slice_col; - } - constexpr int c_sh_wr_delta = active_threads; - auto c_sh_wr = threadIdx.x; - - int row = (threadIdx.x % 32) / 4; - - if (!first) { -// Interestingly, doing direct global accesses here really seems to mess up -// the compiler and lead to slowdowns, hence we also use async-copies even -// though these fetches are not actually asynchronous. -#pragma unroll - for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { - if constexpr (m_block_size_8) { - cp_async4_pred( - &sh_red[c_sh_wr + c_sh_wr_delta * i], - &C[c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i], - (threadIdx.x % 4) * 2 + i < prob_m); - } else { - cp_async4_pred( - &sh_red[c_sh_wr + c_sh_wr_delta * i], - &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)], - i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); - } - } - cp_async_fence(); - cp_async_wait<0>(); - } - -#pragma unroll - for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { - bool mask = (!m_block_size_8) && (i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m) || (m_block_size_8) && ((threadIdx.x % 4) * 2 + i < prob_m); - if (mask) { - if (!first) { - int4 c_red = sh_red[c_sh_wr + i * c_sh_wr_delta]; -#pragma unroll - for (int j = 0; j < 2 * 4; j++) { - int delta = 0; - if constexpr (m_block_size_8) { - delta = j % 2 == 1 ? -2 : 0; - } - reinterpret_cast(&frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta] += Dtype::num2float(reinterpret_cast(&c_red)[j]); - } - } - if (!last) { - int4 c; -#pragma unroll - for (int j = 0; j < 2 * 4; j++) { - int delta = 0; - if constexpr (m_block_size_8) { - delta = j % 2 == 1 ? -2 : 0; - } - reinterpret_cast(&c)[j] = Dtype::float2num(reinterpret_cast(&frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4) + delta]); - } - if constexpr (m_block_size_8) - C[c_gl_wr + i * c_gl_stride + (threadIdx.x % 8) / 4 * c_gl_wr_delta_i] = c; - else - C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2)] = c; - } - } - } - } - }; - - // Globally reduce over threadblocks that compute the same column block. - // We use a tmp C buffer to reduce in full fp32 precision. - auto global_reduce_fp32 = [&](bool first = false, bool last = false) { - constexpr int tb_m = thread_m_blocks * 16; - constexpr int tb_n = thread_n_blocks * 16; - - constexpr int c_size = tb_m * tb_n * sizeof(float) / 16; - - constexpr int active_threads = 32 * thread_n_blocks / 4; - bool is_th_active = threadIdx.x < active_threads; - - constexpr int num_floats = thread_m_blocks * 4 * 2 * 4; - constexpr int th_size = num_floats * sizeof(float) / 16; - - int c_cur_offset = locks_off * c_size; - - if (!is_th_active) { - return; - } - - if (!first) { - float *frag_c_ptr = reinterpret_cast(&frag_c); -#pragma unroll - for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) { - sh_red[threadIdx.x] = C_tmp[c_cur_offset + active_threads * k + threadIdx.x]; - - float *sh_c_ptr = reinterpret_cast(&sh_red[threadIdx.x]); -#pragma unroll - for (int f = 0; f < 4; f++) { - frag_c_ptr[k * 4 + f] += sh_c_ptr[f]; - } - } - } - - if (!last) { - int4 *frag_c_ptr = reinterpret_cast(&frag_c); -#pragma unroll - for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) { - C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k]; - } - } - }; - - // Write out the reduce final result in the correct layout. We only actually - // reshuffle matrix fragments in this step, the reduction above is performed - // in fragment layout. - auto write_result = [&]() { - int c_gl_stride = prob_n / 8; - constexpr int c_sh_stride = 2 * thread_n_blocks + 1; - int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); - constexpr int c_sh_rd_delta = c_sh_stride * (threads / (2 * thread_n_blocks)); - - int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); - c_gl_wr += (2 * thread_n_blocks) * slice_col; - int c_sh_wr; - if constexpr (m_block_size_8) { - c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) + (threadIdx.x % 32) / 4; - c_sh_wr += 64 * (threadIdx.x / 32); - } else { - c_sh_wr = (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; - c_sh_wr += 32 * (threadIdx.x / 32); - } - - int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + (threadIdx.x % (2 * thread_n_blocks)); - - int c_gl_wr_end = c_gl_stride * prob_m; - // We first reorder in shared memory to guarantee the most efficient final - // global write patterns - auto write = [&](int idx, float c0, float c1, FragS &s) { - scalar_t2 res = Dtype::nums2num2(Dtype::float2num(c0), Dtype::float2num(c1)); - - // For per-column quantization we finally apply the scale here (only for - // 4-bit) - if constexpr ( - !has_act_order && group_blocks == -1 && w_type.size_bits() == 4 && (has_zp && dequant_skip_flop || !has_zp)) { - res = __hmul2(res, s[0]); - } - - if constexpr (w_type == host::kFE2M1f) { - res = __hmul2(res, global_scale); - } - - if constexpr (m_block_size_8) { - ((scalar_t *)sh_red)[idx] = res.x; - ((scalar_t *)sh_red)[idx + 8 * c_sh_stride] = res.y; - } else { - ((scalar_t2 *)sh_red)[idx] = res; - } - }; - - if (threadIdx.x / 32 < thread_n_blocks / 4) { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int j = 0; j < 4; j++) { - if constexpr (m_block_size_8) { - int wr = c_sh_wr + 16 * j; - write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]); - write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 1]); - } else { - int wr = c_sh_wr + 8 * j; - write( - wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]); - write( - wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]); - write( - wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]); - write( - wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]); - } - } - c_sh_wr += 16 * (4 * c_sh_stride); - } - } - __syncthreads(); - -#pragma unroll - for (int i = 0; i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); i++) { - if (c_gl_wr < c_gl_wr_end) { - if (use_atomic_add && slice_count > 1) { - scalar_t2 *C_half2 = reinterpret_cast(&C[c_gl_wr]); - scalar_t2 *sh_red_half2 = reinterpret_cast(&sh_red[c_sh_rd]); -#pragma unroll - for (int a = 0; a < 4; a++) { - atomicAdd(&C_half2[a], sh_red_half2[a]); - } - } else { - C[c_gl_wr] = sh_red[c_sh_rd]; - } - c_gl_wr += c_gl_wr_delta; - c_sh_rd += c_sh_rd_delta; - } - } - __syncthreads(); - }; - - // Start global fetch and register load pipelines. - auto start_pipes = [&]() { - -#pragma unroll - for (int i = 0; i < stages - 1; i++) { - if (has_act_order && i == 0) { - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) { - last_g_idx = prob_k - 1; - } - fetch_act_order_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]); - } - - if constexpr (has_zp && !is_zp_float && group_blocks == -1) { - if (i == 0) { - fetch_col_zp_to_shared(); - if constexpr (!dequant_skip_flop) { - fetch_col_scale_to_shared(); - } - } - } - fetch_to_shared(i, i, i < slice_iters); - } - - zero_accums(); - wait_for_stage(); - init_same_group(0); - fetch_to_registers(0, 0); - fetch_scales_to_registers(0, 0); - fetch_zp_to_registers(0, 0); - a_gl_rd += a_gl_rd_delta_o * (stages - 1); - if constexpr (has_act_order) { - slice_k_start_shared_fetch += tb_k * (stages - 1); - } - }; - if (slice_iters) { - start_pipes(); - } - - // Main loop. - while (slice_iters) { - // We unroll over both the global fetch and the register load pipeline to - // ensure all shared memory accesses are static. Note that both pipelines - // have even length meaning that the next iteration will always start at - // index 0. - -#pragma unroll - for (int pipe = 0; pipe < stages;) { -#pragma unroll - for (int k = 0; k < b_sh_wr_iters; k++) { - fetch_to_registers(k + 1, pipe % stages); - fetch_scales_to_registers(k + 1, pipe); - fetch_zp_to_registers(k + 1, pipe); - if (k == b_sh_wr_iters - 2) { - fetch_to_shared((pipe + stages - 1) % stages, pipe, slice_iters >= stages); - pipe++; - wait_for_stage(); - init_same_group(pipe % stages); - } - matmul(k); - } - slice_iters--; - if (slice_iters == 0) { - break; - } - } - - a_gl_rd += a_gl_rd_delta_o * stages; - - if constexpr (has_act_order) { - slice_k_start += tb_k * stages; - - if (slice_k_start < prob_k) { - slice_k_start_shared_fetch += tb_k * stages; - int first_group_id = g_idx[slice_k_start]; - int last_g_idx = slice_k_start + stages * tb_k * 2; - if (last_g_idx >= prob_k) { - last_g_idx = prob_k - 1; - } - int last_group_id = g_idx[last_g_idx]; - if (last_group_id >= sh_first_group_id + sh_num_groups) { - fetch_act_order_scales_to_shared(false, first_group_id, last_group_id); - __syncthreads(); - } - } - } - - // Process results and, if necessary, proceed to the next column slice. - // While this pattern may not be the most readable, other ways of writing - // the loop seemed to noticeably worse performance after compilation. - if (slice_iters == 0) { - cp_async_wait<0>(); - bool last = slice_idx == slice_count - 1; - // For per-column scales, we only fetch them here in the final step before - // write-out - if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) { - if (w_type.size_bits() == 8 || (last || use_atomic_add)) { - if (s_sh_wr_pred) { - cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); - } - cp_async_fence(); - } - } - - thread_block_reduce(); - if constexpr (!has_act_order && group_blocks == -1 && (has_zp && dequant_skip_flop || !has_zp)) { - if (w_type.size_bits() == 8 || (last || use_atomic_add)) { - cp_async_wait<0>(); - __syncthreads(); - if (threadIdx.x / 32 < thread_n_blocks / 4) { - reinterpret_cast(&frag_s)[0] = sh_s[s_sh_rd + 0]; - reinterpret_cast(&frag_s)[1] = sh_s[s_sh_rd + 4]; - if constexpr (m_block_size_8) { - int idx = (threadIdx.x / 4) % 2; - scalar_t2 *frag_s_half2 = reinterpret_cast(frag_s); -#pragma unroll - for (int i = 0; i < 8; i++) { - frag_s_half2[i] = Dtype::num2num2(reinterpret_cast(&frag_s_half2[i])[idx]); - } - } - } - } - } - - // For 8-bit channelwise, we apply the scale before the global reduction - // that converts the fp32 results to fp16 (so that we avoid possible - // overflow in fp16) - if constexpr ( - !has_act_order && group_blocks == -1 && w_type.size_bits() == 8 && (has_zp && dequant_skip_flop || !has_zp)) { - if (threadIdx.x / 32 < thread_n_blocks / 4) { -#pragma unroll - for (int i = 0; i < thread_m_blocks; i++) { -#pragma unroll - for (int j = 0; j < 4; j++) { - scale_float(reinterpret_cast(&frag_c[i][j][0][0]), frag_s[j / 2][2 * (j % 2) + 0]); - scale_float( - reinterpret_cast(&frag_c[i][j][0][2]), frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]); - - if constexpr (!m_block_size_8) { - scale_float(reinterpret_cast(&frag_c[i][j][1][0]), frag_s[j / 2][2 * (j % 2) + 1]); - scale_float(reinterpret_cast(&frag_c[i][j][1][2]), frag_s[j / 2][2 * (j % 2) + 1]); - } - } - } - } - } - - if (slice_count > 1 && !use_atomic_add) { - // only globally reduce if there is more than one block in a slice - barrier_acquire(&locks[locks_off], slice_idx); - if (use_fp32_reduce) { - global_reduce_fp32(slice_idx == 0, last); - } else { - global_reduce_fp16(slice_idx == 0, last); - } - barrier_release(&locks[locks_off], last); - } - if (use_atomic_add && slice_count > 1 && slice_idx != 0) { - wait_negative_and_add(&locks[locks_off]); - } - if (last || use_atomic_add) { - // only the last block in a slice actually writes the result - write_result(); - } - slice_row = 0; - slice_col_par++; - slice_col++; - is_first_matmul_in_slice = true; - init_slice(); - - if (slice_iters) { - a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + (threadIdx.x % a_gl_rd_delta_o); -#pragma unroll - for (int i = 0; i < b_sh_wr_iters; i++) { - B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles; - } - if (slice_col == 0) { -#pragma unroll - for (int i = 0; i < b_sh_wr_iters; i++) { - B_ptr[i] -= b_gl_stride; - } - } - - // Update slice k/n for scales loading - if constexpr (has_act_order) { - slice_k_start = tb_k * slice_row; - slice_k_finish = slice_k_start + tb_k * slice_iters; - slice_k_start_shared_fetch = slice_k_start; - slice_n_offset = act_s_col_tb_stride * slice_col; - } else { - s_gl_rd = s_sh_stride * slice_col + threadIdx.x; - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; - } - - start_pipes(); - } - } - } -} - -} // namespace device::marlin - -#endif diff --git a/src/infiniop/ops/gptq_marlin_gemm/nvidia/gptq_marlin_gemm_nvidia.cu b/src/infiniop/ops/gptq_marlin_gemm/nvidia/gptq_marlin_gemm_nvidia.cu deleted file mode 100644 index edbe5fc8f..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/nvidia/gptq_marlin_gemm_nvidia.cu +++ /dev/null @@ -1,1141 +0,0 @@ -#if defined ENABLE_NVIDIA_API - -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../gptq_marlin_gemm.h" -#include "gptq_marlin_gemm_nvidia.cuh" -#if defined ENABLE_TVM_API -#include "../marlin/kernel.h" -#include "../marlin/marlin_template.h" -#include "../sgl_kernel/scalar_type.hpp" -#include "../sgl_kernel/tensor.h" - -namespace device::marlin { - -INFINIOP_CUDA_KERNEL MarlinDefault(MARLIN_KERNEL_PARAMS){}; - -using MarlinFuncPtr = void (*)(MARLIN_KERNEL_PARAMS); - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - -INFINIOP_CUDA_KERNEL permute_cols_kernel( - int4 const *__restrict__ a_int4_ptr, - int const *__restrict__ perm_int_ptr, - int4 *__restrict__ out_int4_ptr, - int size_m, - int size_k, - int lda, - int block_rows) {} - -#else - -// For a given "a" of size [M,K] performs a permutation of the K columns based -// on the given "perm" indices. -INFINIOP_CUDA_KERNEL permute_cols_kernel( - int4 const *__restrict__ a_int4_ptr, - int const *__restrict__ perm_int_ptr, - int4 *__restrict__ out_int4_ptr, - int size_m, - int size_k, - int lda, - int block_rows) { - auto start_row = block_rows * blockIdx.x; - int finish_row = start_row + block_rows; - if (finish_row > size_m) { - finish_row = size_m; - } - int cur_block_rows = finish_row - start_row; - - int input_row_stride = lda * sizeof(half) / 16; - int output_row_stride = size_k * sizeof(half) / 16; - - auto permute_row = [&](int row) { - int iters = size_k / default_threads; - int rest = size_k % default_threads; - - int input_offset = row * input_row_stride; - int output_offset = row * output_row_stride; - - half const *a_row_half = reinterpret_cast(a_int4_ptr + input_offset); - half *out_half = reinterpret_cast(out_int4_ptr + output_offset); - - int base_k = 0; - - for (int i = 0; i < iters; i++) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - - base_k += default_threads; - } - - if (rest) { - if (threadIdx.x < rest) { - auto cur_k = base_k + threadIdx.x; - int src_pos = perm_int_ptr[cur_k]; - - out_half[cur_k] = a_row_half[src_pos]; - } - } - }; - - for (int i = 0; i < cur_block_rows; i++) { - int cur_row = start_row + i; - if (cur_row < size_m) { - permute_row(cur_row); - } - } -} - -typedef struct -{ - int thread_k; - int thread_n; - int num_threads; -} thread_config_t; - -thread_config_t small_batch_thread_configs[] = { - // Ordered by priority - - // thread_k, thread_n, num_threads - {128, 128, 256}, - {64, 128, 128}, - {128, 64, 128}}; - -thread_config_t large_batch_thread_configs[] = { - // Ordered by priority - - // thread_k, thread_n, num_threads - {64, 256, 256}, - {64, 128, 128}, - {128, 64, 128}}; - -typedef struct -{ - int blocks_per_sm; - thread_config_t tb_cfg; -} exec_config_t; - -int get_scales_cache_size( - thread_config_t const &th_config, - int prob_m, - int prob_n, - int prob_k, - int num_bits, - int group_size, - bool has_act_order, - bool is_k_full) { - bool cache_scales_chunk = has_act_order && !is_k_full; - - int tb_n = th_config.thread_n; - int tb_k = th_config.thread_k; - - // Get max scale groups per thread-block - int tb_groups; - if (group_size == -1) { - tb_groups = 1; - } else if (group_size == 0) { - tb_groups = div_ceil(tb_k, 32); // Worst case is 32 group size - } else { - tb_groups = div_ceil(tb_k, group_size); - } - - if (cache_scales_chunk) { - int load_groups = tb_groups * pipe_stages * 2; // Chunk size is 2x pipeline over dim K - load_groups = max(load_groups, 32); // We load at least 32 scale groups - return load_groups * tb_n * 2; - } else { - int tb_scales = tb_groups * tb_n * 2; - - return tb_scales * pipe_stages; - } -} - -int get_kernel_cache_size( - thread_config_t const &th_config, - int thread_m_blocks, - int prob_m, - int prob_n, - int prob_k, - int num_bits, - int group_size, - bool has_act_order, - bool is_k_full, - int has_zp, - int is_zp_float) { - int pack_factor = 32 / num_bits; - - // Get B size - int tb_k = th_config.thread_k; - int tb_n = th_config.thread_n; - int tb_m = thread_m_blocks * 16; - int sh_a_size = pipe_stages * (tb_m * tb_k) * 2; - int sh_b_size = pipe_stages * (tb_k * tb_n / pack_factor) * 4; - int sh_red_size = tb_m * (tb_n + 8); - int sh_s_size = get_scales_cache_size(th_config, prob_m, prob_n, prob_k, num_bits, group_size, has_act_order, is_k_full); - int sh_g_idx_size = has_act_order && !is_k_full ? pipe_stages * tb_k / 4 : 0; - int sh_zp_size = 0; - if (has_zp) { - if (is_zp_float) { - sh_zp_size = sh_s_size; - } else if (num_bits == 4) { - sh_zp_size = sh_s_size / 4; - } else if (num_bits == 8) { - sh_zp_size = sh_s_size / 2; - } - } - - int total_size = max(sh_b_size, sh_red_size) + sh_a_size + sh_s_size + sh_zp_size + sh_g_idx_size; - - return total_size; -} - -bool is_valid_config( - thread_config_t const &th_config, - int thread_m_blocks, - int prob_m, - int prob_n, - int prob_k, - int num_bits, - int group_size, - bool has_act_order, - bool is_k_full, - int has_zp, - int is_zp_float, - int max_shared_mem) { - // Sanity - if (th_config.thread_k == -1 || th_config.thread_n == -1 || th_config.num_threads == -1) { - return false; - } - - // Verify K/N are divisible by thread K/N - if (prob_k % th_config.thread_k != 0 || prob_n % th_config.thread_n != 0) { - return false; - } - - // Verify min for thread K/N - if (th_config.thread_n < min_thread_n || th_config.thread_k < min_thread_k) { - return false; - } - - // num_threads must be at least 128 (= 4 warps) - if (th_config.num_threads < 128) { - return false; - } - - // Check that pipeline fits into cache - int cache_size = get_kernel_cache_size( - th_config, - thread_m_blocks, - prob_m, - prob_n, - prob_k, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float); - return cache_size <= max_shared_mem; -} - -#define _GET_IF( \ - W_TYPE, THREAD_M_BLOCKS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, M_BLOCK_SIZE_8, GROUP_BLOCKS, NUM_THREADS, IS_ZP_FLOAT) \ - else if ( \ - q_type == W_TYPE && thread_m_blocks == THREAD_M_BLOCKS && thread_n_blocks == THREAD_N_BLOCKS && thread_k_blocks == THREAD_K_BLOCKS && m_block_size_8 == M_BLOCK_SIZE_8 && group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS && is_zp_float == IS_ZP_FLOAT) { \ - kernel = Marlin< \ - scalar_t, \ - W_TYPE.id(), \ - NUM_THREADS, \ - THREAD_M_BLOCKS, \ - THREAD_N_BLOCKS, \ - THREAD_K_BLOCKS, \ - M_BLOCK_SIZE_8, \ - pipe_stages, \ - GROUP_BLOCKS, \ - IS_ZP_FLOAT>; \ - } - -// COMMON: cases for (group_blocks in [-1, 2, 4, 8] and is_zp_float == false) -// this is the most common cases -// BIGGROUP: cases for big group size (group_blocks in [-1, 8]) -// FZP: cases for float-zero-point (is_zp_float = true) -// ACT: cases for act order case (group_blocks == 0) -// FP4: cases for nvfp4(e2m1) (group_blocks == 1) -#define COMMON_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) - -#define COMMON_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ - \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ - \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 2, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) - -#define COMMON_GET_IF(W_TYPE) \ - COMMON_GET_IF_M1(W_TYPE, 8, 8, 256) \ - COMMON_GET_IF_M1(W_TYPE, 8, 4, 128) \ - COMMON_GET_IF_M1(W_TYPE, 4, 8, 128) \ - COMMON_GET_IF_M234(W_TYPE, 16, 4, 256) \ - COMMON_GET_IF_M234(W_TYPE, 8, 4, 128) \ - COMMON_GET_IF_M234(W_TYPE, 4, 8, 128) - -#define BIGGROUP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 8, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) - -#define BIGGROUP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, -1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 8, NUM_THREADS, false) - -#define BIGGROUP_GET_IF(W_TYPE) \ - BIGGROUP_GET_IF_M1(W_TYPE, 8, 8, 256) \ - BIGGROUP_GET_IF_M1(W_TYPE, 8, 4, 128) \ - BIGGROUP_GET_IF_M1(W_TYPE, 4, 8, 128) \ - BIGGROUP_GET_IF_M234(W_TYPE, 16, 4, 256) \ - BIGGROUP_GET_IF_M234(W_TYPE, 8, 4, 128) \ - BIGGROUP_GET_IF_M234(W_TYPE, 4, 8, 128) - -#define FP4_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) - -#define FP4_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 1, NUM_THREADS, false) - -#define FP4_GET_IF(W_TYPE) \ - FP4_GET_IF_M1(W_TYPE, 8, 8, 256) \ - FP4_GET_IF_M1(W_TYPE, 8, 4, 128) \ - FP4_GET_IF_M1(W_TYPE, 4, 8, 128) \ - FP4_GET_IF_M234(W_TYPE, 16, 4, 256) \ - FP4_GET_IF_M234(W_TYPE, 8, 4, 128) \ - FP4_GET_IF_M234(W_TYPE, 4, 8, 128) - -// We currently have 4-bit models only with group_blocks == 4 -#define FZP_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 4, NUM_THREADS, true) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) - -#define FZP_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 4, NUM_THREADS, true) - -#define FZP_GET_IF(W_TYPE) \ - FZP_GET_IF_M1(W_TYPE, 8, 8, 256) \ - FZP_GET_IF_M1(W_TYPE, 8, 4, 128) \ - FZP_GET_IF_M1(W_TYPE, 4, 8, 128) \ - FZP_GET_IF_M234(W_TYPE, 16, 4, 256) \ - FZP_GET_IF_M234(W_TYPE, 8, 4, 128) \ - FZP_GET_IF_M234(W_TYPE, 4, 8, 128) - -// We currently have 4-bit models only with group_blocks == 4 -#define ACT_GET_IF_M1(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, true, 0, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 1, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) - -#define ACT_GET_IF_M234(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \ - _GET_IF(W_TYPE, 2, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 3, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) \ - _GET_IF(W_TYPE, 4, N_BLOCKS, K_BLOCKS, false, 0, NUM_THREADS, false) - -#define ACT_GET_IF(W_TYPE) \ - ACT_GET_IF_M1(W_TYPE, 8, 8, 256) \ - ACT_GET_IF_M1(W_TYPE, 8, 4, 128) \ - ACT_GET_IF_M1(W_TYPE, 4, 8, 128) \ - ACT_GET_IF_M234(W_TYPE, 16, 4, 256) \ - ACT_GET_IF_M234(W_TYPE, 8, 4, 128) \ - ACT_GET_IF_M234(W_TYPE, 4, 8, 128) - -template -MarlinFuncPtr get_marlin_kernel( - const host::ScalarType q_type, - int thread_m_blocks, - int thread_n_blocks, - int thread_k_blocks, - bool m_block_size_8, - bool has_act_order, - bool has_zp, - int group_blocks, - int num_threads, - bool is_zp_float) { - int num_bits = q_type.size_bits(); - auto kernel = MarlinDefault; - if (false) { - } - - COMMON_GET_IF(host::kU4) - COMMON_GET_IF(host::kU4B8) - COMMON_GET_IF(host::kU8B128) - - FP4_GET_IF(host::kFE2M1f) - - BIGGROUP_GET_IF(host::kFE4M3fn) - - ACT_GET_IF(host::kU4B8) - ACT_GET_IF(host::kU8B128) - - if (std::is_same::value) { - if (false) { - } - FZP_GET_IF(host::kU4) - } - - return kernel; -} - -template -exec_config_t determine_exec_config( - const host::ScalarType &q_type, - int prob_m, - int prob_n, - int prob_k, - int thread_m_blocks, - bool m_block_size_8, - int num_bits, - int group_size, - bool has_act_order, - bool is_k_full, - bool has_zp, - bool is_zp_float, - int max_shared_mem, - int sms) { - exec_config_t exec_cfg = exec_config_t{1, thread_config_t{-1, -1, -1}}; - thread_config_t *thread_configs = thread_m_blocks > 1 ? large_batch_thread_configs : small_batch_thread_configs; - int thread_configs_size = thread_m_blocks > 1 ? sizeof(large_batch_thread_configs) / sizeof(thread_config_t) - : sizeof(small_batch_thread_configs) / sizeof(thread_config_t); - - for (int i = 0; i < thread_configs_size; i++) { - thread_config_t th_config = thread_configs[i]; - - if (!is_valid_config( - th_config, - thread_m_blocks, - prob_m, - prob_n, - prob_k, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float, - max_shared_mem)) { - continue; - } - - int cache_size = get_kernel_cache_size( - th_config, - thread_m_blocks, - prob_m, - prob_n, - prob_k, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float); - - int group_blocks = 0; - if (!has_act_order) { - group_blocks = group_size == -1 ? -1 : group_size / 16; - } - - auto kernel = get_marlin_kernel( - q_type, - thread_m_blocks, - th_config.thread_n / 16, - th_config.thread_k / 16, - m_block_size_8, - has_act_order, - has_zp, - group_blocks, - th_config.num_threads, - is_zp_float); - - if (kernel == MarlinDefault) { - continue; - } - - // int m_tiles = div_ceil(prob_m, thread_m_blocks * 16); - // int n_tiles = prob_n / th_config.thread_n; - // int k_tiles = prob_k / th_config.thread_k; - - return {1, th_config}; - } - - return exec_cfg; -} - -template -void marlin_mm( - const void *A, - const void *B, - void *C, - void *C_tmp, - void *s, - void *s2, - void *zp, - void *g_idx, - void *perm, - void *a_tmp, - int prob_m, - int prob_n, - int prob_k, - int lda, - void *workspace, - host::ScalarType const &q_type, - bool has_act_order, - bool is_k_full, - bool has_zp, - int num_groups, - int group_size, - int dev, - cudaStream_t stream, - int thread_k_init, - int thread_n_init, - int sms, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float) { - if (has_zp) { - host::RuntimeCheck( - q_type == host::kU4 || q_type == host::kU8, "q_type must be u4 or u8 when has_zp = True. Got = ", q_type.str()); - } else { - host::RuntimeCheck( - q_type == host::kU4B8 || q_type == host::kU8B128 || q_type == host::kFE4M3fn || q_type == host::kFE2M1f, - "q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when " - "has_zp = False. Got = ", - q_type.str()); - } - - host::RuntimeCheck( - prob_m > 0 && prob_n > 0 && prob_k > 0, "Invalid MNK = [", prob_m, ", ", prob_n, ", ", prob_k, "]"); - - int group_blocks = 0; - if (has_act_order) { - if (is_k_full) { - host::RuntimeCheck(group_size != -1); - group_blocks = group_size / 16; - host::RuntimeCheck( - prob_k % group_blocks == 0, "prob_k = ", prob_k, " is not divisible by group_blocks = ", group_blocks); - } else { - host::RuntimeCheck(group_size == 0); - group_blocks = 0; - } - } else { - if (group_size == -1) { - group_blocks = -1; - } else { - group_blocks = group_size / 16; - host::RuntimeCheck( - prob_k % group_blocks == 0, "prob_k = ", prob_k, " is not divisible by group_blocks = ", group_blocks); - } - } - - int num_bits = q_type.size_bits(); - const int4 *A_ptr = (const int4 *)A; - const int4 *B_ptr = (const int4 *)B; - int4 *C_ptr = (int4 *)C; - int4 *C_tmp_ptr = (int4 *)C_tmp; - const int4 *s_ptr = (const int4 *)s; - const uint16_t *s2_ptr = (const uint16_t *)s2; - const int4 *zp_ptr = (const int4 *)zp; - const int *g_idx_ptr = (const int *)g_idx; - const int *perm_ptr = (const int *)perm; - int4 *a_tmp_ptr = (int4 *)a_tmp; - - int *locks = (int *)workspace; - - if (has_act_order) { - // Permute A columns - int block_rows = div_ceil(prob_m, sms); - host::LaunchKernel(sms, default_threads, stream)( - permute_cols_kernel, A_ptr, perm_ptr, a_tmp_ptr, prob_m, prob_k, lda, block_rows); - A_ptr = a_tmp_ptr; - lda = prob_k; - - // If we have a full K, then we can run the non-act-order version of Marlin - // (since the weight rows are reordered by increasing group ids, and by - // having a full K, we have full original groups) - if (is_k_full) { - has_act_order = false; - } - } - - int max_shared_mem = 0; - host::RuntimeDeviceCheck(cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, dev)); - host::RuntimeCheck(max_shared_mem > 0); - - int max_par = 16; - if (prob_n <= 4096) { - max_par = 16 * 8; - } - int max_shared_mem_new = max_shared_mem; - int rest_m = prob_m; - int max_thread_m_blocks = 4; - while (rest_m) { - int par_count = rest_m / (max_thread_m_blocks * 16); - if (par_count > max_par) { - par_count = max_par; - } - int prob_m_split = par_count > 0 ? (par_count * (max_thread_m_blocks * 16)) : rest_m; - - int thread_k = thread_k_init; - int thread_n = thread_n_init; - - int thread_m_blocks = min(div_ceil(prob_m_split, 16), max_thread_m_blocks); - int m_block_size_8 = prob_m_split <= 8; - - // Set thread config - exec_config_t exec_cfg; - thread_config_t thread_tfg; - if (thread_k != -1 && thread_n != -1) { - thread_tfg = thread_config_t{thread_k, thread_n, default_threads}; - exec_cfg = exec_config_t{1, thread_tfg}; - host::RuntimeCheck(prob_n % thread_n == 0, "prob_n = ", prob_n, " is not divisible by thread_n = ", thread_n); - host::RuntimeCheck(prob_k % thread_k == 0, "prob_k = ", prob_k, " is not divisible by thread_k = ", thread_k); - } else { - // Auto config - exec_cfg = determine_exec_config( - q_type, - prob_m_split, - prob_n, - prob_k, - thread_m_blocks, - m_block_size_8, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float, - max_shared_mem, - sms); - thread_tfg = exec_cfg.tb_cfg; - if (thread_tfg.thread_k == -1 && max_thread_m_blocks > 1) { - max_thread_m_blocks--; - continue; - } - } - - int num_threads = thread_tfg.num_threads; - thread_k = thread_tfg.thread_k; - thread_n = thread_tfg.thread_n; - int blocks = sms * exec_cfg.blocks_per_sm; - if (exec_cfg.blocks_per_sm > 1) { - max_shared_mem_new = max_shared_mem / exec_cfg.blocks_per_sm - 1024; - } - - int thread_k_blocks = thread_k / 16; - int thread_n_blocks = thread_n / 16; - - host::RuntimeCheck( - is_valid_config( - thread_tfg, - thread_m_blocks, - prob_m_split, - prob_n, - prob_k, - num_bits, - group_size, - has_act_order, - is_k_full, - has_zp, - is_zp_float, - max_shared_mem_new), - "Invalid thread config: thread_m_blocks = ", - thread_m_blocks, - ", thread_k = ", - thread_tfg.thread_k, - ", thread_n = ", - thread_tfg.thread_n, - ", num_threads = ", - thread_tfg.num_threads, - " for MKN = [", - prob_m, - ", ", - prob_k, - ", ", - prob_n, - "] and num_bits = ", - num_bits, - ", prob_m_split = ", - prob_m_split, - ", group_size = ", - group_size, - ", has_act_order = ", - has_act_order, - ", is_k_full = ", - is_k_full, - ", has_zp = ", - has_zp, - ", is_zp_float = ", - is_zp_float, - ", max_shared_mem_new = ", - max_shared_mem_new); - - auto kernel = get_marlin_kernel( - q_type, - thread_m_blocks, - thread_n_blocks, - thread_k_blocks, - m_block_size_8, - has_act_order, - has_zp, - group_blocks, - num_threads, - is_zp_float); - - if (kernel == MarlinDefault) { - host::Panic( - "Unsupported shapes: MNK = [", - prob_m, - ", ", - prob_n, - ", ", - prob_k, - "]", - ", has_act_order = ", - has_act_order, - ", num_groups = ", - num_groups, - ", group_size = ", - group_size, - ", prob_m_split = ", - prob_m_split, - ", thread_m_blocks = ", - thread_m_blocks, - ", thread_n_blocks = ", - thread_n_blocks, - ", thread_k_blocks = ", - thread_k_blocks, - ", num_threads = ", - num_threads, - ", num_bits = ", - num_bits); - } - - cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; - host::RuntimeDeviceCheck(cudaStreamIsCapturing(stream, &capture_status)); - if (capture_status == cudaStreamCaptureStatusNone) { - host::RuntimeDeviceCheck( - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem_new)); - } - - bool part_use_atomic_add = use_atomic_add && div_ceil(prob_m_split, 64) * prob_n <= 2048; - - host::LaunchKernel(blocks, num_threads, stream, max_shared_mem_new)( - kernel, - A_ptr, - B_ptr, - C_ptr, - C_tmp_ptr, - s_ptr, - s2_ptr, - zp_ptr, - g_idx_ptr, - num_groups, - prob_m_split, - prob_n, - prob_k, - lda, - locks, - part_use_atomic_add, - use_fp32_reduce, - max_shared_mem_new); - - A_ptr += prob_m_split * (lda / 8); - C_ptr += prob_m_split * (prob_n / 8); - rest_m -= prob_m_split; - } -} - -#endif - -} // namespace device::marlin - -template -void gptq_marlin_gemm(const void *a, - const void *b_q_weight, - void *b_scales, - void *global_scales, - void *b_zeros, - void *g_idx, - void *perm, - void *c, - void *c_tmp, - void *a_tmp, - void *workspace, - int64_t b_q_type_id, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float, - int64_t size_m, - int64_t size_k, - int64_t size_n, - int64_t b_q_size_1, - int64_t a_stride0, - int num_groups, cudaStream_t stream) - -{ - using namespace host; - - // Verify a: [M, K] - ScalarType const b_q_type = ScalarType::from_id(b_q_type_id); - int pack_factor = 32 / b_q_type.size_bits(); - // Verify b_q_weight: [K/tile_size, packed_N] - RuntimeCheck( - size_k % device::marlin::tile_size == 0, - "size_k = ", - size_k, - " is not divisible by tile_size = ", - device::marlin::tile_size); - - int64_t expected_bqw_dim0 = size_k / device::marlin::tile_size; - RuntimeCheck( - b_q_size_1 % device::marlin::tile_size == 0, - "b_q_weight.size(1) = ", - b_q_size_1, - " is not divisible by tile_size = ", - device::marlin::tile_size); - int64_t actual_size_n = (b_q_size_1 / device::marlin::tile_size) * pack_factor; - RuntimeCheck(actual_size_n == size_n, "actual_size_n must = size_n"); - // size_n = actual_size_n - RuntimeCheck(a_stride0 % 8 == 0, "a.stride(0) must be divisible by 8"); - // Verify b_scales: [num_groups, N] - // Early return for zero-size M - if (size_m == 0) { - return; - } - - // int64_t g_idx_size = g_idx.size(0);// g_idx_size == size_k - // int64_t perm_size = perm.size(0);// perm_size == size_k - bool has_act_order = (g_idx != nullptr && perm != nullptr); - - // Determine has_zp from b_zeros size - // int64_t b_zeros_size = b_zeros.size(0); - bool has_zp = (b_zeros != nullptr); - - if (has_zp) { - RuntimeCheck( - b_q_type == kU4 || b_q_type == kU8, "b_q_type must be u4 or u8 when has_zp = True. Got = ", b_q_type.str()); - } else { - RuntimeCheck( - b_q_type == kU4B8 || b_q_type == kU8B128 || b_q_type == kFE4M3fn || b_q_type == kFE2M1f, - "b_q_type must be uint4b8, uint8b128, float8_e4m3fn or float4_e2m1f when " - "has_zp = False. Got = ", - b_q_type.str()); - } - - if (has_zp && is_zp_float) { - RuntimeCheck( - std::is_same::value, "Computation type must be float16 (half) when using float zero points."); - } - - // int64_t global_scales_size = global_scales.size(0); - if (global_scales != nullptr) { - RuntimeCheck(b_q_type == kFE2M1f, "global_scales can only be used for float4_e2m1f."); - } else { - RuntimeCheck(!(b_q_type == kFE2M1f), "the global_scales parameter must be passed for float4_e2m1f."); - } - // Derive group_size - int group_size = -1; - if (has_act_order) { - if (is_k_full) { - RuntimeCheck(num_groups > 1, "For act_order, num_groups must be > 1"); - RuntimeCheck(size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by num_groups = ", num_groups); - group_size = static_cast(size_k / num_groups); - } else { - group_size = 0; - } - } else { - if (num_groups > 1) { - RuntimeCheck(size_k % num_groups == 0, "size_k = ", size_k, ", is not divisible by num_groups = ", num_groups); - group_size = static_cast(size_k / num_groups); - } else { - group_size = -1; - } - } - - // Verify workspace and get device info - RuntimeCheck( - size_n % device::marlin::min_thread_n == 0, - "size_n = ", - size_n, - ", is not divisible by min_thread_n = ", - device::marlin::min_thread_n); - - int device_id = 0; - RuntimeDeviceCheck(cudaGetDevice(&device_id)); - int sms = -1; - RuntimeDeviceCheck(cudaDeviceGetAttribute(&sms, cudaDevAttrMultiProcessorCount, device_id)); - // RuntimeCheck(workspace.size(0) >= sms, "workspace.size(0) = ", workspace.size(0), " is below min_workspace_size = ", sms); - // Hardcoded defaults (auto config) - int thread_k_init = -1; - int thread_n_init = -1; - - // Compute c_tmp and a_tmp pointers - // c_tmp and a_tmp are pre-allocated by caller - - device::marlin::marlin_mm( - a, - b_q_weight, - c, - c_tmp, - b_scales, - global_scales, - b_zeros, - g_idx, - perm, - a_tmp, - static_cast(size_m), - static_cast(size_n), - static_cast(size_k), - static_cast(a_stride0), - workspace, - b_q_type, - has_act_order, - is_k_full, - has_zp, - num_groups, - group_size, - device_id, - stream, - thread_k_init, - thread_n_init, - sms, - use_atomic_add, - use_fp32_reduce, - is_zp_float); -} - -template -infiniStatus_t gptq_marlin_gemm_kernel(void *c, - const void *a, - const void *b_q_weight, - void *b_scales, - void *global_scales, - void *b_zeros, - void *g_idx, - void *perm, - int64_t b_q_type_id, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float, - int64_t size_m, - int64_t size_k, - int64_t size_n, - int64_t b_q_size_1, - int64_t a_stride0, - int num_groups, void *total_buffer, cudaStream_t stream) { - int _MAX_THREAD_N = 256; - int max_blocks_per_sm = 1; - float *c_tmp = nullptr; - void *a_tmp = nullptr; - void *workspace = nullptr; - - // 获取设备 SM 数量（只查询 1 次！） - int dev; - cudaGetDevice(&dev); // 获取当前设备号 - - cudaDeviceProp prop; - cudaGetDeviceProperties(&prop, dev); - const int sms = prop.multiProcessorCount; - - // ===================== 1. 计算每块内存大小 ===================== - size_t c_tmp_bytes = 0; - if (use_fp32_reduce) { - int max_m_block = ((size_m + 15) / 16) * 16; - max_m_block = min(max_m_block, 64); - const size_t c_elems = (size_t)sms * max_m_block * _MAX_THREAD_N; - c_tmp_bytes = c_elems * sizeof(float); - } - - size_t a_tmp_bytes = 0; - bool has_act_order = false; - if (g_idx != nullptr && perm != nullptr) { - has_act_order = true; - } - if (has_act_order) { - a_tmp_bytes = (size_t)size_m * size_k * sizeof(scalar_t); - } - - // workspace 大小（int 类型，必须分配） - const size_t workspace_elems = (size_t)sms * max_blocks_per_sm; - const size_t workspace_bytes = workspace_elems * sizeof(int); - - // ===================== 4. 手动切分指针（核心！） ===================== - uint8_t *ptr = reinterpret_cast(total_buffer); - - // 分配 c_tmp - if (use_fp32_reduce && c_tmp_bytes > 0) { - c_tmp = reinterpret_cast(ptr); - ptr += c_tmp_bytes; - } - - // 分配 a_tmp - if (has_act_order && a_tmp_bytes > 0) { - a_tmp = ptr; - ptr += a_tmp_bytes; - } - - // 分配 workspace - if (workspace_bytes > 0) { - workspace = ptr; - ptr += workspace_bytes; - } - - gptq_marlin_gemm( - a, - b_q_weight, - b_scales, - global_scales, - b_zeros, - g_idx, - perm, - c, - c_tmp, - a_tmp, - workspace, - b_q_type_id, - is_k_full, - use_atomic_add, - use_fp32_reduce, - is_zp_float, - size_m, - size_k, - size_n, - b_q_size_1, - a_stride0, - num_groups, - stream); - return INFINI_STATUS_SUCCESS; -} -#endif -int getCudaDeviceSMCount() { - int dev; - cudaGetDevice(&dev); - cudaDeviceProp prop; - cudaGetDeviceProperties(&prop, dev); - - return prop.multiProcessorCount; -} - -namespace op::gptq_marlin_gemm::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t b_scales_desc, - infiniopTensorDescriptor_t global_scales_desc, - infiniopTensorDescriptor_t b_zeros_desc, - infiniopTensorDescriptor_t g_idx_desc, - infiniopTensorDescriptor_t perm_desc) { - - auto handle = reinterpret_cast(handle_); - auto result = GptqMarlinGemmInfo::create(out_desc, a_desc, b_desc, b_scales_desc, global_scales_desc, b_zeros_desc, g_idx_desc, perm_desc); - - int sms = getCudaDeviceSMCount(); - int _MAX_THREAD_N = 256; - int max_blocks_per_sm = 1; - int max_m_block = ((out_desc->dim(0) + 15) / 16) * 16; - max_m_block = min(max_m_block, 64); - const size_t c_elems = (size_t)sms * max_m_block * _MAX_THREAD_N; - size_t c_tmp_bytes = c_elems * sizeof(float); - size_t a_tmp_bytes = 0; - if (g_idx_desc->numel() > 0 && perm_desc->numel() > 0) { - a_tmp_bytes = (size_t)a_desc->dim(0) * a_desc->dim(1) * infiniSizeOf(a_desc->dtype()); - } - const size_t workspace_elems = (size_t)sms * max_blocks_per_sm; - const size_t workspace_bytes = workspace_elems * sizeof(int); - size_t workspace_size = c_tmp_bytes + a_tmp_bytes + workspace_bytes; - - *desc_ptr = new Descriptor( - workspace_size, - new Opaque{handle->internal()}, - result.take(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t -Descriptor::calculate( - void *workspace, size_t workspace_size, - void *out, - const void *a, - const void *b, - void *b_scales, - void *global_scales, - void *b_zeros, - void *g_idx, - void *perm, - int64_t b_q_type_id, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float, - void *stream_) const { -#if defined ENABLE_TVM_API - cudaStream_t stream = (cudaStream_t)stream_; - int64_t M = static_cast(_info.M); - int64_t K = static_cast(_info.K); - int64_t N = static_cast(_info.N); - int64_t b_q_size_1 = static_cast(_info.b_q_size_1); - int64_t a_stride_0 = static_cast(_info.a_stride_0); - int num_groups = _info.num_groups; - -#define MARLIN(TDATA) \ - gptq_marlin_gemm_kernel(out, a, b, b_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float, M, K, N, b_q_size_1, a_stride_0, num_groups, workspace, stream) - - if (_info.dtype == INFINI_DTYPE_F16) { - return MARLIN(half); - } else if (_info.dtype == INFINI_DTYPE_BF16) { - return MARLIN(__nv_bfloat16); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -#endif - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gptq_marlin_gemm::nvidia - -#endif diff --git a/src/infiniop/ops/gptq_marlin_gemm/nvidia/gptq_marlin_gemm_nvidia.cuh b/src/infiniop/ops/gptq_marlin_gemm/nvidia/gptq_marlin_gemm_nvidia.cuh deleted file mode 100644 index f9c7eb6e9..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/nvidia/gptq_marlin_gemm_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GPTQ_MARLIN_GEMM_CUDA_CUH__ -#define __GPTQ_MARLIN_GEMM_CUDA_CUH__ - -#include "../gptq_marlin_gemm.h" - -DESCRIPTOR(nvidia) - -#endif // __GPTQ_MARLIN_GEMM_CUDA_CUH__ diff --git a/src/infiniop/ops/gptq_marlin_gemm/operator.cc b/src/infiniop/ops/gptq_marlin_gemm/operator.cc deleted file mode 100644 index 60c5e3cc7..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/operator.cc +++ /dev/null @@ -1,120 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/gptq_marlin_gemm.h" - -#if defined ENABLE_NVIDIA_API -#include "nvidia/gptq_marlin_gemm_nvidia.cuh" -#endif - -__INFINI_C infiniStatus_t infiniopCreateGptqMarlinGemmDescriptor( - infiniopHandle_t handle, - infiniopGptqMarlinGemmDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t b_scales_desc, - infiniopTensorDescriptor_t global_scales_desc, - infiniopTensorDescriptor_t b_zeros_desc, - infiniopTensorDescriptor_t g_idx_desc, - infiniopTensorDescriptor_t perm_desc) { -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::gptq_marlin_gemm::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - out_desc, \ - a_desc, \ - b_desc, \ - b_scales_desc, \ - global_scales_desc, \ - b_zeros_desc, \ - g_idx_desc, \ - perm_desc) - - switch (handle->device) { -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetGptqMarlinGemmWorkspaceSize(infiniopGptqMarlinGemmDescriptor_t desc, - size_t *size) { -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopGptqMarlinGemm( - infiniopGptqMarlinGemmDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *a, - const void *b, - void *b_scales, - void *global_scales, - void *b_zeros, - void *g_idx, - void *perm, - int64_t b_q_type_id, - bool is_k_full, - bool use_atomic_add, - bool use_fp32_reduce, - bool is_zp_float, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, out, a, b, b_scales, global_scales, b_zeros, g_idx, perm, b_q_type_id, is_k_full, use_atomic_add, use_fp32_reduce, is_zp_float, stream) - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyGptqMarlinGemmDescriptor(infiniopGptqMarlinGemmDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} - -// #endif diff --git a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/scalar_type.hpp b/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/scalar_type.hpp deleted file mode 100644 index 04ce7a537..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/scalar_type.hpp +++ /dev/null @@ -1,332 +0,0 @@ -#pragma once - -#include -#include -#ifndef __CUDACC__ -#include -#endif - -namespace host { - -// -// ScalarType can represent a wide range of floating point and integer types, -// in particular it can be used to represent sub-byte data types (something -// that torch.dtype currently does not support). -// -// The type definitions on the Python side can be found in: vllm/scalar_type.py -// these type definitions should be kept up to date with any Python API changes -// here. -// -class ScalarType { -public: - enum NanRepr : uint8_t { - NAN_NONE = 0, // nans are not supported - NAN_IEEE_754 = 1, // nans are: exp all 1s, mantissa not all 0s - NAN_EXTD_RANGE_MAX_MIN = 2, // nans are: exp all 1s, mantissa all 1s - - NAN_REPR_ID_MAX - }; - - constexpr ScalarType( - uint8_t exponent, - uint8_t mantissa, - bool signed_, - int32_t bias, - bool finite_values_only = false, - NanRepr nan_repr = NAN_IEEE_754) - : exponent(exponent), - mantissa(mantissa), - signed_(signed_), - bias(bias), - finite_values_only(finite_values_only), - nan_repr(nan_repr){}; - - static constexpr ScalarType int_(uint8_t size_bits, int32_t bias = 0) { - return ScalarType(0, size_bits - 1, true, bias); - } - - static constexpr ScalarType uint(uint8_t size_bits, int32_t bias = 0) { - return ScalarType(0, size_bits, false, bias); - } - - // IEEE 754 compliant floating point type - static constexpr ScalarType float_IEEE754(uint8_t exponent, uint8_t mantissa) { - assert(mantissa > 0 && exponent > 0); - return ScalarType(exponent, mantissa, true, 0, false, NAN_IEEE_754); - } - - // IEEE 754 non-compliant floating point type - static constexpr ScalarType float_(uint8_t exponent, uint8_t mantissa, bool finite_values_only, NanRepr nan_repr) { - assert(nan_repr < NAN_REPR_ID_MAX); - assert(mantissa > 0 && exponent > 0); - assert(nan_repr != NAN_IEEE_754); - return ScalarType(exponent, mantissa, true, 0, finite_values_only, nan_repr); - } - - uint8_t const exponent; // size of the exponent field (0 for integer types) - uint8_t const mantissa; // size of the mantissa field (size of the integer - // excluding the sign bit for integer types) - bool const signed_; // flag if the type supports negative numbers (i.e. has a - // sign bit) - int32_t const bias; // stored values equal value + bias, - // used for quantized type - - // Extra Floating point info - bool const finite_values_only; // i.e. no +/-inf if true - NanRepr const nan_repr; // how NaNs are represented - // (not applicable for integer types) - - using Id = int64_t; - -private: - // Field size in id - template - static constexpr size_t member_id_field_width() { - using T = std::decay_t; - return std::is_same_v ? 1 : sizeof(T) * 8; - } - - template - static constexpr auto reduce_members_helper(Fn f, Init val, Member member, Rest... rest) { - auto new_val = f(val, member); - if constexpr (sizeof...(rest) > 0) { - return reduce_members_helper(f, new_val, rest...); - } else { - return new_val; - }; - } - - template - constexpr auto reduce_members(Fn f, Init init) const { - // Should be in constructor order for `from_id` - return reduce_members_helper(f, init, exponent, mantissa, signed_, bias, finite_values_only, nan_repr); - }; - - template - static constexpr auto reduce_member_types(Fn f, Init init) { - constexpr auto dummy_type = ScalarType(0, 0, false, 0, false, NAN_NONE); - return dummy_type.reduce_members(f, init); - }; - - static constexpr auto id_size_bits() { - return reduce_member_types( - [](int acc, auto member) -> int { return acc + member_id_field_width(); }, 0); - } - -public: - // unique id for this scalar type that can be computed at compile time for - // c++17 template specialization this is not needed once we migrate to - // c++20 and can pass literal classes as template parameters - constexpr Id id() const { - static_assert(id_size_bits() <= sizeof(Id) * 8, "ScalarType id is too large to be stored"); - - auto or_and_advance = [](std::pair result, auto member) -> std::pair { - auto [id, bit_offset] = result; - auto constexpr bits = member_id_field_width(); - return {id | (int64_t(member) & ((uint64_t(1) << bits) - 1)) << bit_offset, bit_offset + bits}; - }; - return reduce_members(or_and_advance, std::pair{}).first; - } - - // create a ScalarType from an id, for c++17 template specialization, - // this is not needed once we migrate to c++20 and can pass literal - // classes as template parameters - static constexpr ScalarType from_id(Id id) { - auto extract_and_advance = [id](auto result, auto member) { - using T = decltype(member); - auto [tuple, bit_offset] = result; - auto constexpr bits = member_id_field_width(); - auto extracted_val = static_cast((int64_t(id) >> bit_offset) & ((uint64_t(1) << bits) - 1)); - auto new_tuple = std::tuple_cat(tuple, std::make_tuple(extracted_val)); - return std::pair{new_tuple, bit_offset + bits}; - }; - - auto [tuple_args, _] = reduce_member_types(extract_and_advance, std::pair, int>{}); - return std::apply([](auto... args) { return ScalarType(args...); }, tuple_args); - } - - constexpr int64_t size_bits() const { - return mantissa + exponent + is_signed(); - } - constexpr bool is_signed() const { - return signed_; - } - constexpr bool is_integer() const { - return exponent == 0; - } - constexpr bool is_floating_point() const { - return exponent > 0; - } - constexpr bool is_ieee_754() const { - return is_floating_point() && finite_values_only == false && nan_repr == NAN_IEEE_754; - } - constexpr bool has_nans() const { - return is_floating_point() && nan_repr != NAN_NONE; - } - constexpr bool has_infs() const { - return is_floating_point() && finite_values_only == false; - } - constexpr bool has_bias() const { - return bias != 0; - } - -#ifndef __CUDACC__ -private: - double _floating_point_max() const { - assert(mantissa <= 52 && exponent <= 11); - - uint64_t max_mantissa = (uint64_t(1) << mantissa) - 1; - if (nan_repr == NAN_EXTD_RANGE_MAX_MIN) { - max_mantissa -= 1; - } - - uint64_t max_exponent = (uint64_t(1) << exponent) - 2; - if (nan_repr == NAN_EXTD_RANGE_MAX_MIN || nan_repr == NAN_NONE) { - assert(exponent < 11); - max_exponent += 1; - } - - // adjust the exponent to match that of a double - // for now we assume the exponent bias is the standard 2^(e-1) -1, (where e - // is the exponent bits), there is some precedent for non-standard biases, - // example `float8_e4m3b11fnuz` here: https://github.com/jax-ml/ml_dtypes - // but to avoid premature over complication we are just assuming the - // standard exponent bias until there is a need to support non-standard - // biases - uint64_t exponent_bias = (uint64_t(1) << (exponent - 1)) - 1; - uint64_t exponent_bias_double = (uint64_t(1) << 10) - 1; // double e = 11 - - uint64_t max_exponent_double = max_exponent - exponent_bias + exponent_bias_double; - - // shift the mantissa into the position for a double and - // the exponent - uint64_t double_raw = (max_mantissa << (52 - mantissa)) | (max_exponent_double << 52); - - return *reinterpret_cast(&double_raw); - } - - constexpr std::variant _raw_max() const { - if (is_floating_point()) { - return {_floating_point_max()}; - } else { - assert(size_bits() < 64 || (size_bits() == 64 && is_signed())); - return {(int64_t(1) << mantissa) - 1}; - } - } - - constexpr std::variant _raw_min() const { - if (is_floating_point()) { - assert(is_signed()); - constexpr uint64_t sign_bit_double = (uint64_t(1) << 63); - - double max = _floating_point_max(); - uint64_t max_raw = *reinterpret_cast(&max); - uint64_t min_raw = max_raw | sign_bit_double; - return {*reinterpret_cast(&min_raw)}; - } else { - assert(!is_signed() || size_bits() <= 64); - if (is_signed()) { - // set the top bit to 1 (i.e. INT64_MIN) and the rest to 0 - // then perform an arithmetic shift right to set all the bits above - // (size_bits() - 1) to 1 - return {INT64_MIN >> (64 - size_bits())}; - } else { - return {int64_t(0)}; - } - } - } - -public: - // Max representable value for this scalar type. - // (accounting for bias if there is one) - constexpr std::variant max() const { - return std::visit([this](auto x) -> std::variant { return {x - bias}; }, _raw_max()); - } - - // Min representable value for this scalar type. - // (accounting for bias if there is one) - constexpr std::variant min() const { - return std::visit([this](auto x) -> std::variant { return {x - bias}; }, _raw_min()); - } -#endif // __CUDACC__ - -public: - std::string str() const { - /* naming generally follows: https://github.com/jax-ml/ml_dtypes - * for floating point types (leading f) the scheme is: - * `float_em[flags]` - * flags: - * - no-flags: means it follows IEEE 754 conventions - * - f: means finite values only (no infinities) - * - n: means nans are supported (non-standard encoding) - * for integer types the scheme is: - * `[u]int[b]` - * - if bias is not present it means its zero - */ - if (is_floating_point()) { - auto ret = "float" + std::to_string(size_bits()) + "_e" + std::to_string(exponent) + "m" + std::to_string(mantissa); - if (!is_ieee_754()) { - if (finite_values_only) { - ret += "f"; - } - if (nan_repr != NAN_NONE) { - ret += "n"; - } - } - return ret; - } else { - auto ret = ((is_signed()) ? "int" : "uint") + std::to_string(size_bits()); - if (has_bias()) { - ret += "b" + std::to_string(bias); - } - return ret; - } - } - - constexpr bool operator==(ScalarType const &other) const { - return mantissa == other.mantissa && exponent == other.exponent && bias == other.bias && signed_ == other.signed_ && finite_values_only == other.finite_values_only && nan_repr == other.nan_repr; - } -}; - -using ScalarTypeId = ScalarType::Id; - -// "rust style" names generally following: -// https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L60-L70 -static inline constexpr auto kS4 = ScalarType::int_(4); -static inline constexpr auto kU4 = ScalarType::uint(4); -static inline constexpr auto kU4B8 = ScalarType::uint(4, 8); -static inline constexpr auto kS8 = ScalarType::int_(8); -static inline constexpr auto kU8 = ScalarType::uint(8); -static inline constexpr auto kU8B128 = ScalarType::uint(8, 128); - -static inline constexpr auto kFE2M1f = ScalarType::float_(2, 1, true, ScalarType::NAN_NONE); -static inline constexpr auto kFE3M2f = ScalarType::float_(3, 2, true, ScalarType::NAN_NONE); -static inline constexpr auto kFE4M3fn = ScalarType::float_(4, 3, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); -static inline constexpr auto kFE8M0fnu = ScalarType(8, 0, false, 0, true, ScalarType::NAN_EXTD_RANGE_MAX_MIN); -static inline constexpr auto kFE5M2 = ScalarType::float_IEEE754(5, 2); -static inline constexpr auto kFE8M7 = ScalarType::float_IEEE754(8, 7); -static inline constexpr auto kFE5M10 = ScalarType::float_IEEE754(5, 10); - -// Fixed width style names, generally following: -// https://github.com/pytorch/pytorch/blob/6d9f74f0af54751311f0dd71f7e5c01a93260ab3/torch/csrc/api/include/torch/types.h#L47-L57 -static inline constexpr auto kInt4 = kS4; -static inline constexpr auto kUint4 = kU4; -static inline constexpr auto kUint4b8 = kU4B8; -static inline constexpr auto kInt8 = kS8; -static inline constexpr auto kUint8 = kU8; -static inline constexpr auto kUint8b128 = kU8B128; - -static inline constexpr auto kFloat4_e2m1f = kFE2M1f; -static inline constexpr auto kFloat6_e3m2f = kFE3M2f; -static inline constexpr auto kFloat8_e4m3fn = kFE4M3fn; -static inline constexpr auto kFloat8_e5m2 = kFE5M2; -static inline constexpr auto kFloat16_e8m7 = kFE8M7; -static inline constexpr auto kFloat16_e5m10 = kFE5M10; - -// colloquial names -static inline constexpr auto kHalf = kFE5M10; -static inline constexpr auto kFloat16 = kHalf; -static inline constexpr auto kBFloat16 = kFE8M7; - -static inline constexpr auto kFloat16Id = kFloat16.id(); -} // namespace host diff --git a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/source_location.h b/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/source_location.h deleted file mode 100644 index 9a06fb380..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/source_location.h +++ /dev/null @@ -1,40 +0,0 @@ -/// \file source_location.h -/// \brief Portable `source_location` wrapper. -/// -/// Uses `std::source_location` when available (C++20), otherwise falls -/// back to a minimal stub that returns empty/zero values. - -#pragma once -#include - -/// NOTE: fallback to a minimal source_location implementation -#if defined(__cpp_lib_source_location) -#include - -using source_location_t = std::source_location; - -#else - -struct source_location_fallback { -public: - static constexpr source_location_fallback current() noexcept { - return source_location_fallback{}; - } - constexpr source_location_fallback() noexcept = default; - constexpr unsigned line() const noexcept { - return 0; - } - constexpr unsigned column() const noexcept { - return 0; - } - constexpr const char *file_name() const noexcept { - return ""; - } - constexpr const char *function_name() const noexcept { - return ""; - } -}; - -using source_location_t = source_location_fallback; - -#endif diff --git a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/tensor.h b/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/tensor.h deleted file mode 100644 index 48c22c3f1..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/tensor.h +++ /dev/null @@ -1,547 +0,0 @@ -/// \file tensor.h -/// \brief Tensor validation and symbolic matching utilities. -#pragma once -#include "utils.h" - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#ifdef __CUDACC__ -#include "utils.cuh" -#endif - -namespace host { -struct SymbolicSize; -struct SymbolicDType; -struct SymbolicDevice; - -namespace details { -inline constexpr auto kAnyDeviceID = -1; -inline constexpr auto kAnySize = static_cast(-1); -inline constexpr auto kNullSize = static_cast(-1); -inline constexpr auto kNullDType = static_cast(18u); -inline constexpr auto kNullDevice = static_cast(-1); - -template -struct ArrayView { - const T *data; - size_t size; - - __host__ __device__ ArrayView() : data(nullptr), size(0) {} - __host__ __device__ ArrayView(const T *d, size_t s) : data(d), size(s) {} - - template - __host__ __device__ ArrayView(const std::array &arr) - : data(arr.data()), size(arr.size()) {} - - __host__ __device__ const T &operator[](size_t i) const { return data[i]; } - __host__ __device__ bool empty() const { return size == 0; } -}; - -template -struct PrintAbleSpan { - const T *data; - size_t length; - - PrintAbleSpan(const T *p, size_t l) : data(p), length(l) {} - size_t size() const { return length; } - const T &operator[](size_t i) const { return data[i]; } -}; - -inline constexpr const char *kDeviceStringMap[] = { - "", // 0 - "cpu", // 1 - "cuda", // 2 - "cuda_host", // 3 - "opencl", // 4 - "vulkan", // 5 - "metal", // 6 - "vpi", // 7 - "rocm", // 8 - "rocm_host", // 9 - "ext_dev", // 10 - "cuda_managed", // 11 - "oneapi", // 12 - "webgpu", // 13 - "hexagon", // 14 - "maia", // 15 - "trn", // 16 -}; - -constexpr int kMaxDeviceType = 16; - -template -struct _dtype_trait; - -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLInt, 8, 1}; -}; -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLInt, 16, 1}; -}; -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLInt, 32, 1}; -}; -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLInt, 64, 1}; -}; -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLUInt, 8, 1}; -}; -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLUInt, 16, 1}; -}; -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLUInt, 32, 1}; -}; -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLUInt, 64, 1}; -}; -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLFloat, 32, 1}; -}; -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLFloat, 64, 1}; -}; - -#ifdef __CUDACC__ -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLFloat, 16, 1}; -}; -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLBfloat, 16, 1}; -}; -template <> -struct _dtype_trait { - static constexpr DLDataType value = {kDLFloat8_e4m3fn, 8, 1}; -}; -#endif - -template -struct _device_trait { - static constexpr DLDevice value = {Code, kAnyDeviceID}; -}; - -template -inline constexpr std::array kDTypeList = { - _dtype_trait::value...}; - -template -inline constexpr std::array kDeviceList = { - _device_trait::value...}; - -} // namespace details - -inline std::ostream &operator<<(std::ostream &os, DLDevice device) { - int code = static_cast(device.device_type); - if (code < 1 || code > details::kMaxDeviceType) { - RuntimeCheck(false, "Unknown device: ", code); - } - os << details::kDeviceStringMap[code]; - if (device.device_id != details::kAnyDeviceID && device.device_type != kDLCPU) { - os << ":" << device.device_id; - } - return os; -} - -template -inline std::ostream &operator<<(std::ostream &os, const details::PrintAbleSpan &span) { - os << "["; - for (size_t i = 0; i < span.size(); ++i) { - if (i > 0) { - os << ", "; - } - os << span[i]; - } - os << "]"; - return os; -} - -// ============================================== -// SymbolicSize 完整定义 -// ============================================== -struct SymbolicSize { -public: - explicit SymbolicSize(std::string_view ann = {}) - : m_value(details::kNullSize), m_ann(ann) {} - - SymbolicSize(const SymbolicSize &) = delete; - SymbolicSize &operator=(const SymbolicSize &) = delete; - - std::string_view get_name() const { return m_ann; } - bool has_value() const { return m_value != details::kNullSize; } - - void set_value(int64_t v) { - RuntimeCheck(!has_value(), "Size already set"); - m_value = v; - } - - std::optional get_value() const { - return has_value() ? std::optional(m_value) : std::nullopt; - } - - int64_t unwrap(DebugInfo info = {}) const { - RuntimeCheck(info, has_value(), "Size not set"); - return m_value; - } - - void verify(int64_t v, const char *prefix, int64_t dim) { - if (has_value()) { - if (m_value != v) [[unlikely]] { - Panic("Size mismatch for ", m_name_str(prefix, dim), ": expected ", m_value, " got ", v); - } - } else { - set_value(v); - } - } - - std::string value_or_name(const char *prefix, int64_t dim) const { - if (auto v = get_value()) { - return std::to_string(*v); - } - return m_name_str(prefix, dim); - } - -private: - std::string m_name_str(const char *prefix, int64_t dim) const { - std::ostringstream os; - os << prefix << '#' << dim; - if (!m_ann.empty()) { - os << "('" << m_ann << "')"; - } - return os.str(); - } - - int64_t m_value; - std::string_view m_ann; -}; - -inline bool operator==(DLDevice a, DLDevice b) { - return a.device_type == b.device_type && a.device_id == b.device_id; -} - -// ============================================== -// SymbolicDType 完整定义 -// ============================================== -struct SymbolicDType { -public: - SymbolicDType() : m_value({details::kNullDType, 0, 0}) {} - SymbolicDType(const SymbolicDType &) = delete; - SymbolicDType &operator=(const SymbolicDType &) = delete; - - bool has_value() const { return m_value.code != details::kNullDType; } - - void set_value(DLDataType v) { - RuntimeCheck(!has_value(), "DType already set"); - RuntimeCheck(m_check(v), "DType not allowed: ", v); - m_value = v; - } - - std::optional get_value() const { - return has_value() ? std::optional(m_value) : std::nullopt; - } - - DLDataType unwrap(DebugInfo info = {}) const { - RuntimeCheck(info, has_value(), "DType not set"); - return m_value; - } - - void set_options(details::ArrayView opts) { m_opts = opts; } - - template - void set_options() { - m_opts = details::ArrayView(details::kDTypeList.data(), details::kDTypeList.size()); - } - - void verify(DLDataType dtype) { - if (has_value()) { - RuntimeCheck(m_value == dtype, "DType mismatch: expected ", m_value, " got ", dtype); - } else { - set_value(dtype); - } - } - - template - bool is_type() const { - return m_value == details::_dtype_trait::value; - } - -private: - bool m_check(DLDataType v) const { - if (m_opts.empty()) { - return true; - } - for (size_t i = 0; i < m_opts.size; ++i) { - if (m_opts[i] == v) { - return true; - } - } - return false; - } - - details::ArrayView m_opts; - DLDataType m_value; -}; - -// ============================================== -// SymbolicDevice 完整定义 -// ============================================== -struct SymbolicDevice { -public: - SymbolicDevice() : m_value({details::kNullDevice, details::kAnyDeviceID}) {} - SymbolicDevice(const SymbolicDevice &) = delete; - SymbolicDevice &operator=(const SymbolicDevice &) = delete; - - bool has_value() const { return m_value.device_type != details::kNullDevice; } - - void set_value(DLDevice v) { - RuntimeCheck(!has_value(), "Device already set"); - RuntimeCheck(m_check(v), "Device not allowed: ", details::PrintableDevice{v}); - m_value = v; - } - - std::optional get_value() const { - return has_value() ? std::optional(m_value) : std::nullopt; - } - - DLDevice unwrap(DebugInfo info = {}) const { - RuntimeCheck(info, has_value(), "Device not set"); - return m_value; - } - - void set_options(details::ArrayView opts) { m_opts = opts; } - - template - void set_options() { - m_opts = details::ArrayView(details::kDeviceList.data(), details::kDeviceList.size()); - } - - void verify(DLDevice dev) { - if (has_value()) { - RuntimeCheck(m_value == dev, "Device mismatch: expected ", - details::PrintableDevice{m_value}, " got ", details::PrintableDevice{dev}); - } else { - set_value(dev); - } - } - -private: - bool m_check(DLDevice v) const { - if (m_opts.empty()) { - return true; - } - for (size_t i = 0; i < m_opts.size; ++i) { - auto o = m_opts[i]; - if (o.device_type != v.device_type) { - continue; - } - if (o.device_id == details::kAnyDeviceID || o.device_id == v.device_id) { - return true; - } - } - return false; - } - - details::ArrayView m_opts; - DLDevice m_value; -}; - -// ============================================== -// BaseRef / Ref 类型（现在类型已完整定义） -// ============================================== -namespace details { -template -struct BaseRef { - BaseRef() : m_ref(&m_cache) {} - explicit BaseRef(T &r) : m_ref(&r) {} - BaseRef(const BaseRef &) = delete; - BaseRef &operator=(const BaseRef &) = delete; - - T *operator->() const { return m_ref; } - T &operator*() const { return *m_ref; } - void rebind(T &r) { m_ref = &r; } - -private: - T *m_ref; - T m_cache; -}; - -struct SizeRef : public BaseRef { - using BaseRef::BaseRef; - SizeRef(int64_t v); -}; - -struct DTypeRef : public BaseRef { - using BaseRef::BaseRef; - DTypeRef(DLDataType); - DTypeRef(std::initializer_list); - DTypeRef(ArrayView); -}; - -struct DeviceRef : public BaseRef { - using BaseRef::BaseRef; - DeviceRef(DLDevice); - DeviceRef(std::initializer_list); - DeviceRef(ArrayView); -}; - -inline SizeRef::SizeRef(int64_t v) { - if (v != kAnySize) { - (**this).set_value(v); - } -} -inline DTypeRef::DTypeRef(DLDataType v) { (**this).set_value(v); } -inline DTypeRef::DTypeRef(std::initializer_list l) : DTypeRef(ArrayView(l.begin(), l.size())) {} -inline DTypeRef::DTypeRef(ArrayView v) { (**this).set_options(v); } -inline DeviceRef::DeviceRef(DLDevice v) { (**this).set_value(v); } -inline DeviceRef::DeviceRef(std::initializer_list l) : DeviceRef(ArrayView(l.begin(), l.size())) {} -inline DeviceRef::DeviceRef(ArrayView v) { (**this).set_options(v); } - -} // namespace details - -template -inline bool is_type(DLDataType dtype) { - return dtype == details::_dtype_trait::value; -} - -// ============================================== -// TensorMatcher -// ============================================== -struct TensorMatcher { - using SizeRef = details::SizeRef; - using DTypeRef = details::DTypeRef; - using DeviceRef = details::DeviceRef; - - TensorMatcher(const TensorMatcher &) = delete; - TensorMatcher &operator=(const TensorMatcher &) = delete; - - explicit TensorMatcher(std::initializer_list s) - : m_shape(s.begin(), s.size()), m_strides(nullptr, 0) {} - - TensorMatcher &&with_strides(std::initializer_list s) && { - RuntimeCheck(m_strides.empty(), "Strides already set"); - RuntimeCheck(m_shape.size == s.size(), "Stride/shape size mismatch"); - m_strides = details::ArrayView(s.begin(), s.size()); - return std::move(*this); - } - - template - TensorMatcher &&with_dtype(DTypeRef &&d) && { - m_dtype.rebind(*d); - m_dtype->template set_options(); - return std::move(*this); - } - - template - TensorMatcher &&with_dtype() && { - m_dtype->template set_options(); - return std::move(*this); - } - - template - TensorMatcher &&with_device(DeviceRef &&d) && { - m_device.rebind(*d); - m_device->template set_options(); - return std::move(*this); - } - - template - TensorMatcher &&with_device() && { - m_device->template set_options(); - return std::move(*this); - } - - const TensorMatcher &&verify(tvm::ffi::TensorView, DebugInfo = {}) const &&; - -private: - static void s_print_tensor(std::ostringstream &, tvm::ffi::TensorView); - void m_verify_impl(tvm::ffi::TensorView) const; - - details::ArrayView m_shape; - details::ArrayView m_strides; - DTypeRef m_dtype; - DeviceRef m_device; -}; - -inline void TensorMatcher::s_print_tensor(std::ostringstream &os, tvm::ffi::TensorView v) { - os << "Tensor<"; - size_t d = 0; - for (int64_t s : v.shape()) { - if (d++) { - os << ", "; - } - os << s; - } - os << ">[strides=<"; - d = 0; - for (int64_t s : v.strides()) { - if (d++) { - os << ", "; - } - os << s; - } - os << ">, dtype=" << v.dtype(); - os << ", device=" << details::PrintableDevice{v.device()} << "]"; -} - -inline const TensorMatcher &&TensorMatcher::verify(tvm::ffi::TensorView v, DebugInfo info) const && { - try { - m_verify_impl(v); - } catch (PanicError &e) { - std::ostringstream os; - os << "Tensor match failed: "; - s_print_tensor(os, v); - os << " @ " << info.file_name() << ":" << info.line() << "\n- cause: " << e.root_cause(); - throw PanicError(os.str()); - } - return std::move(*this); -} - -inline void TensorMatcher::m_verify_impl(tvm::ffi::TensorView v) const { - size_t dim = static_cast(v.dim()); - RuntimeCheck(dim == m_shape.size, "Dim mismatch: expected ", m_shape.size, " got ", dim); - - for (size_t i = 0; i < dim; ++i) { - m_shape[i]->verify(v.size(i), "shape", (int64_t)i); - } - - if (!m_strides.empty()) { - for (size_t i = 0; i < dim; ++i) { - if (v.size(i) != 1 || !m_strides[i]->has_value()) { - m_strides[i]->verify(v.stride(i), "stride", (int64_t)i); - } - } - } else { - RuntimeCheck(v.is_contiguous(), "Tensor not contiguous"); - } - - m_dtype->verify(v.dtype()); - m_device->verify(v.device()); -} - -} // namespace host diff --git a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.cuh b/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.cuh deleted file mode 100644 index 18d5da7c3..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.cuh +++ /dev/null @@ -1,312 +0,0 @@ -/// \file utils.cuh -/// \brief Core CUDA/device utilities: type aliases, PDL helpers, -/// typed pointer access, kernel launch wrapper, and error checking. -/// -/// This header is included (directly or transitively) by nearly every -/// JIT kernel. It provides: -/// - Scalar/packed type aliases (`fp16_t`, `bf16_t`, `fp8_e4m3_t`, ...). -/// - `SGL_DEVICE` macro (forced-inline device function qualifier). -/// - `kWarpThreads` constant (32). -/// - PDL (Programmatic Dependent Launch) helpers for Hopper (sm_90+). -/// - Typed `load_as` / `store_as` for void-pointer access. -/// - `pointer::offset` for safe void-pointer arithmetic. -/// - `host::LaunchKernel` - kernel launcher with optional PDL. -/// - `host::RuntimeDeviceCheck` - CUDA error checking. - -#pragma once - -#include "utils.h" - -#include -#include - -#include -#include -#include -#ifndef USE_ROCM -#include -#include -#include -#include -#else -#include -#include -#include -#ifndef __grid_constant__ -#define __grid_constant__ -#endif -using cudaError_t = hipError_t; -using cudaStream_t = hipStream_t; -using cudaLaunchConfig_t = hipLaunchConfig_t; -using cudaLaunchAttribute = hipLaunchAttribute; -inline constexpr auto cudaSuccess = hipSuccess; -#define cudaStreamPerThread hipStreamPerThread -#define cudaGetErrorString hipGetErrorString -#define cudaGetLastError hipGetLastError -#define cudaLaunchKernel hipLaunchKernel -#endif - -#ifndef USE_ROCM -using fp32_t = float; -// using fp16_t = __half; -// using bf16_t = __nv_bfloat16; -using fp8_e4m3_t = __nv_fp8_e4m3; -using fp8_e5m2_t = __nv_fp8_e5m2; - -using fp32x2_t = float2; -using fp16x2_t = __half2; -using bf16x2_t = __nv_bfloat162; -using fp8x2_e4m3_t = __nv_fp8x2_e4m3; -using fp8x2_e5m2_t = __nv_fp8x2_e5m2; - -using fp32x4_t = float4; -#else -using fp32_t = float; -using fp16_t = __half; -using bf16_t = __hip_bfloat16; -using fp8_e4m3_t = uint8_t; -using fp8_e5m2_t = uint8_t; -using fp32x2_t = float2; -using fp16x2_t = half2; -using bf16x2_t = __hip_bfloat162; -using fp8x2_e4m3_t = uint16_t; -using fp8x2_e5m2_t = uint16_t; -using fp32x4_t = float4; -#endif - -/* - * LDG Support - */ -#ifndef USE_ROCM -#define SGLANG_LDG(arg) __ldg(arg) -#else -#define SGLANG_LDG(arg) *(arg) -#endif - -namespace device { - -/// \brief Macro: forced-inline device function qualifier. -#define SGL_DEVICE __forceinline__ __device__ - -// Architecture detection: SGL_CUDA_ARCH is injected by load_jit() and is -// available in both host and device compilation passes, whereas __CUDA_ARCH__ -// is only defined by nvcc during the device pass. -#if !defined(USE_ROCM) -#if !defined(SGL_CUDA_ARCH) -#error "SGL_CUDA_ARCH is not defined. JIT compilation must inject -DSGL_CUDA_ARCH via load_jit()." -#endif -#if defined(__CUDA_ARCH__) -static_assert( - __CUDA_ARCH__ == SGL_CUDA_ARCH, "SGL_CUDA_ARCH mismatch: injected arch flag does not match device target"); -#endif -#define SGL_ARCH_HOPPER_OR_GREATER (SGL_CUDA_ARCH >= 900) -#define SGL_ARCH_BLACKWELL_OR_GREATER ((SGL_CUDA_ARCH >= 1000) && (CUDA_VERSION >= 12090)) -#else // USE_ROCM -#define SGL_ARCH_HOPPER_OR_GREATER 0 -#define SGL_ARCH_BLACKWELL_OR_GREATER 0 -#endif - -// Maximum vector size in bytes supported by current architecture. -// Pre-Blackwell / AMD: 128-bit (16 bytes) -// Blackwell or greater: 256-bit (32 bytes) -inline constexpr std::size_t kMaxVecBytes = SGL_ARCH_BLACKWELL_OR_GREATER ? 32 : 16; - -/// \brief Number of threads per warp (always 32 on NVIDIA/AMD GPUs). -inline constexpr auto kWarpThreads = 32u; -/// \brief Full warp active mask (all 32 lanes). -inline constexpr auto kFullMask = 0xffffffffu; - -/** - * \brief PDL (Programmatic Dependent Launch): wait for the primary kernel. - * - * On Hopper (sm_90+), inserts a `griddepcontrol.wait` instruction to - * synchronize with a preceding kernel in the same stream. On older - * architectures or ROCm this is a no-op. - */ -template -SGL_DEVICE void PDLWaitPrimary() { -#if SGL_ARCH_HOPPER_OR_GREATER - if constexpr (kUsePDL) { - asm volatile("griddepcontrol.wait;" :: - : "memory"); - } -#endif -} - -/** - * \brief PDL: trigger dependent (secondary) kernel launch. - * - * On Hopper (sm_90+), inserts a `griddepcontrol.launch_dependents` - * instruction. On older architectures or ROCm this is a no-op. - */ -template -SGL_DEVICE void PDLTriggerSecondary() { -#if SGL_ARCH_HOPPER_OR_GREATER - if constexpr (kUsePDL) { - asm volatile("griddepcontrol.launch_dependents;" :: - :); - } -#endif -} - -template -SGL_DEVICE constexpr auto div_ceil(T a, U b) { - static_assert(std::is_integral::value && std::is_integral::value, - "div_ceil requires integer types"); - return (a + b - 1) / b; -} - -/** - * \brief Load data with the specified type and offset from a void pointer. - * \tparam T The type to load. - * \param ptr The base pointer. - * \param offset The offset in number of elements of type T. - */ -template -SGL_DEVICE T load_as(const void *ptr, int64_t offset = 0) { - return static_cast(ptr)[offset]; -} - -/** - * \brief Store data with the specified type and offset to a void pointer. - * \tparam T The type to store. - * \param ptr The base pointer. - * \param val The value to store. - * \param offset The offset in number of elements of type T. - * \note we use type_identity_t to force the caller to explicitly specify - * the template parameter `T`, which can avoid accidentally using the wrong type. - */ -template -SGL_DEVICE void store_as(void *ptr, T val, int64_t offset = 0) { - static_cast(ptr)[offset] = val; -} - -/// \brief Safe void-pointer arithmetic (byte-level by default). -namespace pointer { -// we only allow void * pointer arithmetic for safety - -template -SGL_DEVICE auto offset(void *ptr, U... offset) -> void * { - return static_cast(ptr) + (offset + ...); -} - -template -SGL_DEVICE auto offset(const void *ptr, U... offset) -> const void * { - return static_cast(ptr) + (offset + ...); -} - -} // namespace pointer - -} // namespace device - -namespace host { - -/** - * \brief Check the CUDA error code and panic with location info on failure. - */ -inline void RuntimeDeviceCheck(::cudaError_t error, DebugInfo location = {}) { - if (error != ::cudaSuccess) { - [[unlikely]]; - ::host::panic(location, "CUDA error: ", ::cudaGetErrorString(error)); - } -} - -/// \brief Check the last CUDA error (calls `cudaGetLastError`). -inline void RuntimeDeviceCheck(DebugInfo location = {}) { - return RuntimeDeviceCheck(::cudaGetLastError(), location); -} - -/** - * \brief Kernel launcher with automatic stream resolution and PDL support. - * - * Usage: - * \code - * host::LaunchKernel(grid, block, device) - * .enable_pdl(true) - * (my_kernel, arg1, arg2); - * \endcode - * - * The constructor resolves the CUDA stream from a `DLDevice` (via - * `TVMFFIEnvGetStream`) or accepts a raw `cudaStream_t`. The call - * operator launches the kernel and checks for errors. - */ -struct LaunchKernel { -public: - explicit LaunchKernel( - dim3 grid_dim, - dim3 block_dim, - DLDevice device, - std::size_t dynamic_shared_mem_bytes = 0, - DebugInfo location = {}) noexcept - : m_config(s_make_config(grid_dim, block_dim, resolve_device(device), dynamic_shared_mem_bytes)), - m_location(location) {} - - explicit LaunchKernel( - dim3 grid_dim, - dim3 block_dim, - cudaStream_t stream, - std::size_t dynamic_shared_mem_bytes = 0, - DebugInfo location = {}) noexcept - : m_config(s_make_config(grid_dim, block_dim, stream, dynamic_shared_mem_bytes)), m_location(location) {} - - LaunchKernel(const LaunchKernel &) = delete; - LaunchKernel &operator=(const LaunchKernel &) = delete; - - static auto resolve_device(DLDevice device) -> cudaStream_t { - return static_cast(::TVMFFIEnvGetStream(device.device_type, device.device_id)); - } - - auto enable_pdl(bool enabled = true) -> LaunchKernel & { -#ifdef USE_ROCM - (void)enabled; - m_config.numAttrs = 0; -#else - if (enabled) { - m_attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - m_attrs[0].val.programmaticStreamSerializationAllowed = true; - m_config.numAttrs = 1; - m_config.attrs = m_attrs; - } else { - m_config.numAttrs = 0; - } -#endif - return *this; - } - - template - auto operator()(T &&kernel, Args &&...args) const -> void { -#ifdef USE_ROCM - hipLaunchKernelGGL( - std::forward(kernel), - m_config.gridDim, - m_config.blockDim, - m_config.dynamicSmemBytes, - m_config.stream, - std::forward(args)...); - RuntimeDeviceCheck(m_location); -#else - RuntimeDeviceCheck(::cudaLaunchKernelEx(&m_config, kernel, std::forward(args)...), m_location); -#endif - } - -private: - static auto s_make_config( // Make a config for kernel launch - dim3 grid_dim, - dim3 block_dim, - cudaStream_t stream, - std::size_t smem) -> cudaLaunchConfig_t { - auto config = ::cudaLaunchConfig_t{}; - config.gridDim = grid_dim; - config.blockDim = block_dim; - config.dynamicSmemBytes = smem; - config.stream = stream; - config.numAttrs = 0; - return config; - } - - cudaLaunchConfig_t m_config; - const DebugInfo m_location; - cudaLaunchAttribute m_attrs[1]; -}; - -} // namespace host diff --git a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.h b/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.h deleted file mode 100644 index 855b32c36..000000000 --- a/src/infiniop/ops/gptq_marlin_gemm/sgl_kernel/utils.h +++ /dev/null @@ -1,245 +0,0 @@ -/// \file utils.h -/// \brief Host-side C++ utilities used by JIT kernel wrappers. -/// -/// Provides: -/// - `DebugInfo` - wraps `std::source_location` for error reporting. -/// - `RuntimeCheck` - runtime assertion with formatted error messages. -/// - `Panic` - unconditional abort with formatted error messages. -/// - `pointer::offset` - safe void-pointer arithmetic (host side). -/// - `div_ceil` - integer ceiling division. -/// - `dtype_bytes` - byte width of a `DLDataType`. -/// - `irange` - Python-style integer range for range-for loops. -/// - `PrintableDevice` - wrapper for outputting device info to streams (NVIDIA only). - -#pragma once - -// ref: https://forums.developer.nvidia.com/t/c-20s-source-location-compilation-error-when-using-nvcc-12-1/258026/3 -#ifdef __CUDACC__ -#include -#if CUDA_VERSION <= 12010 - -#pragma push_macro("__cpp_consteval") -#pragma push_macro("_NODISCARD") -#pragma push_macro("__builtin_LINE") - -#pragma clang diagnostic push -#pragma clang diagnostic ignored "-Wbuiltin-macro-redefined" -#define __cpp_consteval 201811L -#pragma clang diagnostic pop - -#ifdef _NODISCARD -#undef _NODISCARD -#define _NODISCARD -#endif - -#define consteval constexpr - -#include "source_location.h" - -#undef consteval -#pragma pop_macro("__cpp_consteval") -#pragma pop_macro("_NODISCARD") -#else // __CUDACC__ && CUDA_VERSION > 12010 -#include "source_location.h" -#endif -#else // no __CUDACC__ -#include "source_location.h" -#endif - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace host { - -template -inline constexpr bool dependent_false_v = false; - -#if defined(ENABLE_NVIDIA_API) - -namespace details { - -/// \brief PrintableDevice wrapper for outputting device info to streams. -struct PrintableDevice { - DLDevice device; -}; - -/// \brief Stream output operator for PrintableDevice. -inline std::ostream &operator<<(std::ostream &os, PrintableDevice pd) { - int code = static_cast(pd.device.device_type); - static const char *device_names[] = { - "", "cpu", "cuda", "cuda_host", "opencl", "vulkan", - "metal", "vpi", "rocm", "rocm_host", "ext_dev", - "cuda_managed", "oneapi", "webgpu", "hexagon", "maia", "trn"}; - if (code >= 0 && code <= 16) { - os << device_names[code]; - if (pd.device.device_id >= 0 && code != 1 && code != 0) { - os << ":" << pd.device.device_id; - } - } else { - os << "device(" << code << "," << pd.device.device_id << ")"; - } - return os; -} - -} // namespace details - -#endif // ENABLE_NVIDIA_API - -/// \brief Source-location wrapper for debug/error messages. -struct DebugInfo : public source_location_t { - DebugInfo(source_location_t loc = source_location_t::current()) : source_location_t(loc) {} -}; - -/// \brief Exception type thrown by `RuntimeCheck` and `Panic`. -struct PanicError : public std::runtime_error { -public: - explicit PanicError(std::string msg) : runtime_error(msg), m_message(std::move(msg)) {} - auto root_cause() const -> std::string_view { - const auto str = std::string_view{m_message}; - const auto pos = str.find(": "); - return pos == std::string_view::npos ? str : str.substr(pos + 2); - } - -private: - std::string m_message; -}; - -/// \brief Unconditionally abort with a formatted error message. -template -[[noreturn]] inline auto panic(DebugInfo location, Args &&...args) -> void { - std::ostringstream os; - os << "Runtime check failed at " << location.file_name() << ":" << location.line(); - if constexpr (sizeof...(args) > 0) { - os << ": "; - (os << ... << std::forward(args)); - } else { - os << " in " << location.function_name(); - } - throw PanicError(std::move(os).str()); -} - -/** - * \brief Runtime assertion: panics with a formatted message when `condition` - * is false. Extra `args` are streamed to the error message. - * - * Example: - * \code - * RuntimeCheck(n > 0, "n must be positive, got ", n); - * \endcode - */ -template -struct RuntimeCheck { - template - explicit RuntimeCheck(Cond &&condition, Args &&...args, DebugInfo location = {}) { - if (condition) { - return; - } - [[unlikely]] ::host::panic(location, std::forward(args)...); - } - template - explicit RuntimeCheck(DebugInfo location, Cond &&condition, Args &&...args) { - if (condition) { - return; - } - [[unlikely]] ::host::panic(location, std::forward(args)...); - } -}; - -template -struct Panic { - explicit Panic(Args &&...args, DebugInfo location = {}) { - ::host::panic(location, std::forward(args)...); - } - explicit Panic(DebugInfo location, Args &&...args) { - ::host::panic(location, std::forward(args)...); - } - [[noreturn]] ~Panic() { - std::terminate(); - } -}; - -template -explicit RuntimeCheck(Cond &&, Args &&...) -> RuntimeCheck; - -template -explicit RuntimeCheck(DebugInfo, Cond &&, Args &&...) -> RuntimeCheck; - -template -explicit Panic(Args &&...) -> Panic; - -template -explicit Panic(DebugInfo, Args &&...) -> Panic; - -namespace pointer { - -// we only allow void * pointer arithmetic for safety - -template ::value && ...)>> -inline auto offset(void *ptr, U... offset) -> void * { - return static_cast(ptr) + (... + offset); -} - -template ::value && ...)>> -inline auto offset(const void *ptr, U... offset) -> const void * { - return static_cast(ptr) + (... + offset); -} - -} // namespace pointer - -/// \brief Integer ceiling division: ceil(a / b). -template -inline constexpr auto div_ceil(T a, U b) { - static_assert(std::is_integral::value, "T must be integral"); - static_assert(std::is_integral::value, "U must be integral"); - return (a + b - 1) / b; -} - -/// \brief Returns the byte width of a DLPack data type. -inline auto dtype_bytes(DLDataType dtype) -> std::size_t { - return static_cast(dtype.bits / 8); -} - -// Pure C++11 compatible irange - removes std::ranges/std::integral for older CUDA compilers -template -struct IntegerRange { - T start_; - T end_; - - struct Iterator { - T value; - - T operator*() const { return value; } - Iterator &operator++() { - ++value; - return *this; - } - bool operator!=(const Iterator &other) const { - return value != other.value; - } - }; - - Iterator begin() const { return {start_}; } - Iterator end() const { return {end_}; } -}; - -/// Python-style integer range: irange(n) -> [0, n) -template -IntegerRange irange(T end) { - return {0, end}; -} - -/// Python-style integer range: irange(start, end) -> [start, end) -template -IntegerRange irange(T start, T end) { - return {start, end}; -} - -} // namespace host diff --git a/src/infiniop/ops/gptq_marlin_repack/cuda/kernel.cuh b/src/infiniop/ops/gptq_marlin_repack/cuda/kernel.cuh deleted file mode 100644 index 9424588b8..000000000 --- a/src/infiniop/ops/gptq_marlin_repack/cuda/kernel.cuh +++ /dev/null @@ -1,252 +0,0 @@ -#include "../marlin/marlin.cuh" - -namespace marlin { - -template -__device__ void gptq_marlin_repack_kernel( - uint32_t const *__restrict__ b_q_weight_ptr, - uint32_t const *__restrict__ perm_ptr, uint32_t *__restrict__ out_ptr, - int size_k, int size_n) { - constexpr int pack_factor = 32 / num_bits; - - constexpr int target_tile_n_size = tile_n_size / (is_a_8bit ? 2 : 1); - constexpr int target_tile_k_size = tile_k_size * (is_a_8bit ? 2 : 1); - int k_tiles = size_k / target_tile_k_size; - int n_tiles = size_n / target_tile_n_size; - int block_k_tiles = div_ceil(k_tiles, gridDim.x); - - auto start_k_tile = blockIdx.x * block_k_tiles; - if (start_k_tile >= k_tiles) { - return; - } - - int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); - - // Wait until the next thread tile has been loaded to shared memory. - auto wait_for_stage = [&]() { - // We only have `stages - 2` active fetches since we are double buffering - // and can only issue the next fetch when it is guaranteed that the previous - // shared memory load is fully complete (as it may otherwise be - // overwritten). - cp_async_wait(); - __syncthreads(); - }; - - extern __shared__ int4 sh[]; - - constexpr int perm_size = target_tile_k_size / 4; - - int4 *sh_perm_ptr = sh; - int4 *sh_pipe_ptr = sh_perm_ptr; - if constexpr (has_perm) { - sh_pipe_ptr += perm_size; - } - - constexpr int tile_ints = target_tile_k_size / pack_factor; - - constexpr int stage_n_threads = target_tile_n_size / 4; - constexpr int stage_k_threads = has_perm ? target_tile_k_size : tile_ints; - constexpr int stage_size = stage_k_threads * stage_n_threads; - - auto load_perm_to_shared = [&](int k_tile_id) { - int first_k_int4 = (k_tile_id * target_tile_k_size) / 4; - - int4 const *perm_int4_ptr = reinterpret_cast(perm_ptr); - - if (threadIdx.x < perm_size) { - sh_perm_ptr[threadIdx.x] = perm_int4_ptr[first_k_int4 + threadIdx.x]; - } - __syncthreads(); - }; - - auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - cp_async_fence(); - return; - } - - int first_n = n_tile_id * target_tile_n_size; - - int4 *sh_ptr = sh_pipe_ptr + stage_size * pipe; - - if constexpr (has_perm) { - if (threadIdx.x < stage_size) { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; - - uint32_t const *sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); - - int src_k = sh_perm_int_ptr[k_id]; - int src_k_packed = src_k / pack_factor; - - cp_async4( - &sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast(&( - b_q_weight_ptr[src_k_packed * size_n + first_n + (n_id * 4)]))); - } - } else { - if (threadIdx.x < stage_size) { - auto k_id = threadIdx.x / stage_n_threads; - auto n_id = threadIdx.x % stage_n_threads; - - int first_k = k_tile_id * target_tile_k_size; - int first_k_packed = first_k / pack_factor; - - cp_async4(&sh_ptr[k_id * stage_n_threads + n_id], - reinterpret_cast( - &(b_q_weight_ptr[(first_k_packed + k_id) * size_n + first_n + (n_id * 4)]))); - } - } - - cp_async_fence(); - }; - - auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { - if (n_tile_id >= n_tiles) { - return; - } - - auto warp_id = threadIdx.x / 32; - auto th_id = threadIdx.x % 32; - - if (warp_id >= 4) { - return; - } - - int tc_col = th_id / 4; - int tc_row = (th_id % 4) * (is_a_8bit ? 4 : 2); - - constexpr int tc_offsets[4] = {0, 1, 8, 9}; - - int cur_n = (warp_id / (is_a_8bit ? 2 : 1)) * 16 + tc_col; - - constexpr int sh_stride = target_tile_n_size; - constexpr uint32_t mask = (1 << num_bits) - 1; - - int4 *sh_stage_ptr = sh_pipe_ptr + stage_size * pipe; - uint32_t *sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); - - uint32_t *sh_perm_int_ptr = reinterpret_cast(sh_perm_ptr); - - uint32_t vals[8]; - - if constexpr (has_perm) { - static_assert(!is_a_8bit); - for (int i = 0; i < 4; i++) { - int k_idx = tc_row + tc_offsets[i]; - - uint32_t src_k = sh_perm_int_ptr[k_idx]; - uint32_t src_k_pos = src_k % pack_factor; - - uint32_t b1_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n]; - uint32_t b1_cur_val = (b1_val >> (src_k_pos * num_bits)) & mask; - - uint32_t b2_val = sh_stage_int_ptr[k_idx * sh_stride + cur_n + 8]; - uint32_t b2_cur_val = (b2_val >> (src_k_pos * num_bits)) & mask; - - vals[i] = b1_cur_val; - vals[4 + i] = b2_cur_val; - } - } else { - uint32_t b1_vals[tile_ints]; - uint32_t b2_vals[tile_ints]; - -#pragma unroll - for (int i = 0; i < tile_ints; i++) { - if constexpr (is_a_8bit) { - b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i + (warp_id % 2) * 8]; - } else { - b1_vals[i] = sh_stage_int_ptr[cur_n + sh_stride * i]; - b2_vals[i] = sh_stage_int_ptr[cur_n + 8 + sh_stride * i]; - } - } - -#pragma unroll - for (int i = 0; i < 4; i++) { - int cur_elem = tc_row + (is_a_8bit ? i : tc_offsets[i]); - int cur_int = cur_elem / pack_factor; - int cur_pos = cur_elem % pack_factor; - - vals[i] = (b1_vals[cur_int] >> (cur_pos * num_bits)) & mask; - if constexpr (is_a_8bit) - vals[4 + i] = (b1_vals[cur_int + tile_ints / 2] >> (cur_pos * num_bits)) & mask; - else - vals[4 + i] = (b2_vals[cur_int] >> (cur_pos * num_bits)) & mask; - } - } - - constexpr int tile_size = target_tile_k_size * target_tile_n_size / pack_factor; - int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size; - - // Result of: - // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h - if constexpr (!is_a_8bit && num_bits == 4) { - int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; - - uint32_t res = 0; -#pragma unroll - for (int i = 0; i < 8; i++) { - res |= vals[pack_idx[i]] << (i * 4); - } - - out_ptr[out_offset + th_id * 4 + warp_id] = res; - } else if constexpr (is_a_8bit && num_bits == 4) { - int pack_idx[8] = {0, 4, 1, 5, 2, 6, 3, 7}; - - uint32_t res = 0; -#pragma unroll - for (int i = 0; i < 8; i++) { - res |= vals[pack_idx[i]] << (i * 4); - } - - out_ptr[out_offset + th_id * 4 + warp_id] = res; - } else { - constexpr int pack_idx[4] = {0, 2, 1, 3}; - - uint32_t res1 = 0; - uint32_t res2 = 0; -#pragma unroll - for (int i = 0; i < 4; i++) { - const int ii = is_a_8bit ? i : pack_idx[i]; - res1 |= vals[ii] << (i * 8); - res2 |= vals[4 + ii] << (i * 8); - } - - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; - out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; - } - }; - - auto start_pipes = [&](int k_tile_id, int n_tile_id) { -#pragma unroll - for (int pipe = 0; pipe < repack_stages - 1; pipe++) { - fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); - } - - wait_for_stage(); - }; -#pragma unroll - for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) { - int n_tile_id = 0; - - if constexpr (has_perm) { - load_perm_to_shared(k_tile_id); - } - - start_pipes(k_tile_id, n_tile_id); - - while (n_tile_id < n_tiles) { -#pragma unroll - for (int pipe = 0; pipe < repack_stages; pipe++) { - fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, - n_tile_id + pipe + repack_stages - 1); - repack_tile(pipe, k_tile_id, n_tile_id + pipe); - wait_for_stage(); - } - n_tile_id += repack_stages; - } - } -} - -} // namespace marlin diff --git a/src/infiniop/ops/gptq_marlin_repack/gptq_marlin_repack.h b/src/infiniop/ops/gptq_marlin_repack/gptq_marlin_repack.h deleted file mode 100644 index 3a059b664..000000000 --- a/src/infiniop/ops/gptq_marlin_repack/gptq_marlin_repack.h +++ /dev/null @@ -1,50 +0,0 @@ -#ifndef GPTQ_MARLIN_REPACK_H -#define GPTQ_MARLIN_REPACK_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::gptq_marlin_repack::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - GptqMarlinRepackInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - GptqMarlinRepackInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t output_desc, \ - infiniopTensorDescriptor_t input_desc, \ - infiniopTensorDescriptor_t perm_desc, \ - int64_t num_bits, \ - bool is_a_8bit); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *output, \ - const void *input, \ - const void *perm, \ - void *stream) const; \ - }; \ - } - -#endif // GPTQ_MARLIN_REPACK_H diff --git a/src/infiniop/ops/gptq_marlin_repack/info.h b/src/infiniop/ops/gptq_marlin_repack/info.h deleted file mode 100644 index dd724d358..000000000 --- a/src/infiniop/ops/gptq_marlin_repack/info.h +++ /dev/null @@ -1,55 +0,0 @@ -#ifndef __AWQ_MARLIN_REPACK_INFO_H__ -#define __AWQ_MARLIN_REPACK_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include "marlin/marlin.cuh" -#include - -#include - -namespace op::gptq_marlin_repack { - -class GptqMarlinRepackInfo { - GptqMarlinRepackInfo() = default; - -public: - infiniDtype_t output_dtype, input_dtype; - size_t size_k, size_n; - int64_t num_bits; - bool is_a_8bit, has_perm; - - static utils::Result create( - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t perm_desc, - int64_t num_bits, - bool is_a_8bit) { - CHECK_OR_RETURN( - output_desc != nullptr && input_desc != nullptr, - INFINI_STATUS_NULL_POINTER); - const infiniDtype_t output_dtype = output_desc->dtype(); - const infiniDtype_t input_dtype = input_desc->dtype(); - - CHECK_DTYPE(input_dtype, output_dtype); - - int const pack_factor = 32 / num_bits; - size_t size_k = input_desc->dim(0) * pack_factor; - size_t size_n = input_desc->dim(1); - - CHECK_OR_RETURN(size_k / marlin::tile_size == output_desc->dim(0) || size_n * marlin::tile_size / pack_factor == output_desc->dim(1), - INFINI_STATUS_BAD_TENSOR_SHAPE); - bool has_perm = false; - - if (perm_desc != nullptr && perm_desc->dim(0) != 0) { - has_perm = true; - } - - return utils::Result( - GptqMarlinRepackInfo{output_dtype, input_dtype, size_k, size_n, num_bits, is_a_8bit, has_perm}); - } -}; - -} // namespace op::gptq_marlin_repack - -#endif // __AWQ_MARLIN_REPACK_INFO_H__ diff --git a/src/infiniop/ops/gptq_marlin_repack/marlin/marlin.cuh b/src/infiniop/ops/gptq_marlin_repack/marlin/marlin.cuh deleted file mode 100644 index f3d897d27..000000000 --- a/src/infiniop/ops/gptq_marlin_repack/marlin/marlin.cuh +++ /dev/null @@ -1,178 +0,0 @@ -#pragma once - -#ifndef _marlin_cuh -#define _marlin_cuh - -#include -#include -#include -#include - -#ifndef MARLIN_NAMESPACE_NAME -#define MARLIN_NAMESPACE_NAME marlin -#endif - -template -__device__ __forceinline__ uint32_t __cvta_generic_to_shared(T *ptr) { - size_t smem_addr; - asm volatile( - "cvta.to.shared.u64 %0, %1;" - : "=l"(smem_addr) - : "l"(ptr)); - return static_cast(smem_addr); -} - -namespace MARLIN_NAMESPACE_NAME { - -// Marlin params - -// 8 warps are a good choice since every SM has 4 schedulers and having more -// than 1 warp per schedule allows some more latency hiding. At the same time, -// we want relatively few warps to have many registers per warp and small tiles. -static constexpr int default_threads = 256; - -static constexpr int pipe_stages = 4; // 4 pipeline stages fit into shared memory - -static constexpr int min_thread_n = 64; -static constexpr int min_thread_k = 64; -static constexpr int max_thread_n = 256; - -static constexpr int tile_size = 16; -static constexpr int max_par = 16; - -// Repack params -static constexpr int repack_stages = 8; - -static constexpr int repack_threads = 256; - -static constexpr int tile_k_size = tile_size; -static constexpr int tile_n_size = tile_k_size * 4; - -// Helpers -template -struct Vec { - T elems[n]; - __device__ T &operator[](int i) { return elems[i]; } -}; - -using I4 = Vec; - -constexpr int div_ceil(int a, int b) { return (a + b - 1) / b; } - -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 - -__device__ inline void cp_async1_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async2_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - if (pred) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; - } -} - -__device__ inline void cp_async4(void *smem_ptr, const void *glob_ptr) { - reinterpret_cast(smem_ptr)[0] = reinterpret_cast(glob_ptr)[0]; -} - -__device__ inline void cp_async_fence() {} - -template -__device__ inline void cp_async_wait() {} - -#else - -__device__ inline void cp_async1_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 4; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async2_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 8; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4_ca_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.ca.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4_pred(void *smem_ptr, const void *glob_ptr, - bool pred = true) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " .reg .pred p;\n" - " setp.ne.b32 p, %0, 0;\n" - " @p cp.async.cg.shared.global [%1], [%2], %3;\n" - "}\n" ::"r"((int)pred), - "r"(smem), "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async4(void *smem_ptr, const void *glob_ptr) { - const int BYTES = 16; - uint32_t smem = static_cast(__cvta_generic_to_shared(smem_ptr)); - asm volatile( - "{\n" - " cp.async.cg.shared.global [%0], [%1], %2;\n" - "}\n" ::"r"(smem), - "l"(glob_ptr), "n"(BYTES)); -} - -__device__ inline void cp_async_fence() { - asm volatile("cp.async.commit_group;\n" ::); -} - -template -__device__ inline void cp_async_wait() { - asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); -} - -#endif - -} // namespace MARLIN_NAMESPACE_NAME - -#endif diff --git a/src/infiniop/ops/gptq_marlin_repack/nvidia/gptq_marlin_repack_nvidia.cu b/src/infiniop/ops/gptq_marlin_repack/nvidia/gptq_marlin_repack_nvidia.cu deleted file mode 100644 index 1c709fc4f..000000000 --- a/src/infiniop/ops/gptq_marlin_repack/nvidia/gptq_marlin_repack_nvidia.cu +++ /dev/null @@ -1,134 +0,0 @@ -#if defined(ENABLE_NVIDIA_API) -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../cuda/kernel.cuh" -#include "gptq_marlin_repack_nvidia.cuh" -#include - -template -INFINIOP_CUDA_KERNEL gptqMarlinRepackKernel( - const uint32_t *__restrict__ b_q_weight_ptr, - const uint32_t *__restrict__ perm_ptr, uint32_t *__restrict__ out_ptr, - int size_k, int size_n) { - - marlin::gptq_marlin_repack_kernel( - b_q_weight_ptr, perm_ptr, out_ptr, - size_k, size_n); -} - -#define CALL_IF(NUM_BITS, HAS_PERM, IS_A_8BIT) \ - else if (num_bits == NUM_BITS && has_perm == HAS_PERM && is_a_8bit == IS_A_8BIT) { \ - cudaFuncSetAttribute( \ - gptqMarlinRepackKernel, \ - cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \ - gptqMarlinRepackKernel \ - <<>>( \ - b_q_weight_ptr, perm_ptr, out_ptr, size_k, size_n); \ - } - -infiniStatus_t gptqMarlinRepack(uint32_t *out_ptr, const uint32_t *b_q_weight_ptr, const uint32_t *perm_ptr, - int64_t size_k, int64_t size_n, int64_t num_bits, - bool is_a_8bit, bool has_perm, cudaStream_t stream) { - - // Verify compatibility with marlin tile of 16x64 - if (size_k % marlin::tile_k_size != 0) { - std::cout << "size_k = " << size_k << " is not divisible by tile_k_size = " << marlin::tile_k_size << std::endl; - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (size_n % marlin::tile_n_size != 0) { - std::cout << "size_n = " << size_n << " is not divisible by tile_n_size = " << marlin::tile_n_size << std::endl; - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - if (num_bits != 4 && num_bits != 8) { - std::cout << "num_bits must be 4 or 8. Got = " << num_bits << std::endl; - return INFINI_STATUS_BAD_PARAM; - } - - // Get dev info - int device_id = 0; - - int blocks; - cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device_id); - - int max_shared_mem = 0; - cudaDeviceGetAttribute(&max_shared_mem, - cudaDevAttrMaxSharedMemoryPerBlockOptin, device_id); - assert(max_shared_mem > 0 && "max_shared_mem must be greater than 0"); - - if (false) { - } - CALL_IF(4, false, false) - CALL_IF(4, true, false) - CALL_IF(8, false, false) - CALL_IF(8, true, false) - - CALL_IF(4, false, true) - CALL_IF(8, false, true) - else { - fprintf(stderr, "Unsupported repack config: num_bits = %ld, has_perm = %s, is_a_8bit = %s\n", - num_bits, - has_perm ? "true" : "false", - is_a_8bit ? "true" : "false"); - assert(false); - } - - return INFINI_STATUS_SUCCESS; -} - -namespace op::gptq_marlin_repack::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { delete _opaque; } - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t perm_desc, - int64_t num_bits, - bool is_a_8bit) { - - auto handle = reinterpret_cast(handle_); - auto result = GptqMarlinRepackInfo::create(output_desc, input_desc, perm_desc, num_bits, is_a_8bit); - - size_t workspace_size = 0; - - *desc_ptr = new Descriptor( - new Opaque{handle->internal()}, - result.take(), - workspace_size, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t -Descriptor::calculate( - void *workspace, size_t workspace_size, - void *output, - const void *input, - const void *perm, - void *stream_) const { - - cudaStream_t stream = (cudaStream_t)stream_; - - int64_t size_k = static_cast(_info.size_k); - int64_t size_n = static_cast(_info.size_n); - int64_t num_bits = _info.num_bits; - bool is_a_8bit = _info.is_a_8bit; - bool has_perm = _info.has_perm; - - gptqMarlinRepack((uint32_t *)output, (const uint32_t *)input, (const uint32_t *)perm, - size_k, size_n, num_bits, - is_a_8bit, has_perm, stream); - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gptq_marlin_repack::nvidia -#endif // ENABLE_NVIDIA_API diff --git a/src/infiniop/ops/gptq_marlin_repack/nvidia/gptq_marlin_repack_nvidia.cuh b/src/infiniop/ops/gptq_marlin_repack/nvidia/gptq_marlin_repack_nvidia.cuh deleted file mode 100644 index 25a537c9e..000000000 --- a/src/infiniop/ops/gptq_marlin_repack/nvidia/gptq_marlin_repack_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __GPTQ_MARLIN_REPACK_CUDA_CUH__ -#define __GPTQ_MARLIN_REPACK_CUDA_CUH__ - -#include "../gptq_marlin_repack.h" - -DESCRIPTOR(nvidia) - -#endif // __GPTQ_MARLIN_REPACK_CUDA_CUH__ diff --git a/src/infiniop/ops/gptq_marlin_repack/operator.cc b/src/infiniop/ops/gptq_marlin_repack/operator.cc deleted file mode 100644 index 966574877..000000000 --- a/src/infiniop/ops/gptq_marlin_repack/operator.cc +++ /dev/null @@ -1,104 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/gptq_marlin_repack.h" - -#if defined ENABLE_NVIDIA_API -#include "nvidia/gptq_marlin_repack_nvidia.cuh" -#endif - -__INFINI_C infiniStatus_t infiniopCreateGptqMarlinRepackDescriptor( - infiniopHandle_t handle, - infiniopGptqMarlinRepackDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t perm_desc, - int64_t num_bits, - bool is_a_8bit) { -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::gptq_marlin_repack::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output_desc, \ - input_desc, \ - perm_desc, \ - num_bits, \ - is_a_8bit) - - switch (handle->device) { -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetGptqMarlinRepackWorkspaceSize(infiniopGptqMarlinRepackDescriptor_t desc, - size_t *size) { -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopGptqMarlinRepack( - infiniopGptqMarlinRepackDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *perm, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, input, perm, stream) - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyGptqMarlinRepackDescriptor(infiniopGptqMarlinRepackDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} - -// #endif diff --git a/src/infiniop/ops/gptq_qyblas_gemm/gptq_qyblas_gemm.h b/src/infiniop/ops/gptq_qyblas_gemm/gptq_qyblas_gemm.h deleted file mode 100644 index 456e540e1..000000000 --- a/src/infiniop/ops/gptq_qyblas_gemm/gptq_qyblas_gemm.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef GPTQ_QYBLAS_GEMM_H -#define GPTQ_QYBLAS_GEMM_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::gptq_qyblas_gemm::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - GptqQyblasGemmInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - GptqQyblasGemmInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t a_desc, \ - infiniopTensorDescriptor_t b_desc, \ - infiniopTensorDescriptor_t b_scales_desc, \ - infiniopTensorDescriptor_t b_zeros_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *out, \ - const void *a, const void *b, void *b_scale, void *b_zero, int64_t quant_type, int64_t bit, \ - void *stream) const; \ - }; \ - } - -#endif // GPTQ_QYBLAS_GEMM_H diff --git a/src/infiniop/ops/gptq_qyblas_gemm/info.h b/src/infiniop/ops/gptq_qyblas_gemm/info.h deleted file mode 100644 index 24820d53b..000000000 --- a/src/infiniop/ops/gptq_qyblas_gemm/info.h +++ /dev/null @@ -1,117 +0,0 @@ -#ifndef __GPTQ_QYBLAS_GEMM_INFO_H__ -#define __GPTQ_QYBLAS_GEMM_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include -#include - -inline void prepare_matrix_for_cublas( - infiniopTensorDescriptor_t tensor, - bool &transpose_tensor) { - - auto strides = tensor->strides(); - auto sizes = tensor->shape(); - - if ((strides[0] == 1) && (strides[1] >= std::max(1, sizes[0]))) { - - transpose_tensor = false; - return; - } - if ((strides[1] == 1) && (strides[0] >= std::max(1, sizes[1]))) { - - transpose_tensor = true; - return; - } - transpose_tensor = true; -} - -namespace op::gptq_qyblas_gemm { - -class GptqQyblasGemmInfo { - GptqQyblasGemmInfo() = default; - -public: - infiniDtype_t dtype, weight_dtype, scales_dtype, zeros_dtype, out_dtype; - size_t M, K, N, scales_size_0, scales_size_1; - ptrdiff_t lda, ldb, result_ld; - bool transpose_result; - char transa, transb; - - static utils::Result createGptqQyblasGemmInfo( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t b_scales_desc, - infiniopTensorDescriptor_t b_zeros_desc) { - - auto dtype = a_desc->dtype(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16); - auto out_dtype = out_desc->dtype(); - CHECK_DTYPE(dtype, out_dtype); - - const infiniDtype_t weight_dtype = b_desc->dtype(); - // CHECK_DTYPE(weight_dtype, INFINI_DTYPE_F8, INFINI_DTYPE_U8, INFINI_DTYPE_I8); - - const infiniDtype_t scales_dtype = b_scales_desc->dtype(); - const infiniDtype_t zeros_dtype = b_zeros_desc->dtype(); - - bool transpose_result = false; - bool transpose_mat_1 = false; - bool transpose_mat_2 = false; - - prepare_matrix_for_cublas(out_desc, transpose_result); - - auto mata = (transpose_result ? b_desc : a_desc); - prepare_matrix_for_cublas(transpose_result ? b_desc : a_desc, transpose_mat_1); - auto matb = (transpose_result ? a_desc : b_desc); - prepare_matrix_for_cublas(transpose_result ? a_desc : b_desc, transpose_mat_2); - - auto mat1_sizes = a_desc->shape(); - auto mat2_sizes = b_desc->shape(); - if (transpose_result) { - transpose_mat_1 = !transpose_mat_1; - transpose_mat_2 = !transpose_mat_2; - mat1_sizes = mata->shape(); - mat2_sizes = matb->shape(); - } - - size_t M = mat1_sizes[transpose_result ? 1 : 0]; - size_t K = mat1_sizes[transpose_result ? 0 : 1]; - size_t N = mat2_sizes[transpose_result ? 0 : 1]; - - size_t scales_size_0 = b_scales_desc->shape()[0]; - size_t scales_size_1 = b_scales_desc->shape()[1]; - - auto ndim = out_desc->ndim(); - CHECK_OR_RETURN(ndim == 2 - && a_desc->ndim() == ndim - && b_desc->ndim() == ndim - && b_scales_desc->ndim() == ndim - && b_zeros_desc->ndim() == ndim, - INFINI_STATUS_BAD_TENSOR_SHAPE); - - ptrdiff_t lda = mata->strides()[(transpose_mat_1 == transpose_result) - ? 1 - : 0]; - ptrdiff_t ldb = matb->strides()[(transpose_mat_2 == transpose_result) - ? 1 - : 0]; - ptrdiff_t result_ld = out_desc->strides()[transpose_result ? 0 : 1]; - - char transa = transpose_mat_1 ? 't' : 'n'; - char transb = transpose_mat_2 ? 't' : 'n'; - - return utils::Result(GptqQyblasGemmInfo{ - dtype, weight_dtype, scales_dtype, zeros_dtype, out_dtype, - M, K, N, scales_size_0, scales_size_1, - lda, ldb, result_ld, - transpose_result, - transa, transb}); - } -}; - -} // namespace op::gptq_qyblas_gemm - -#endif // __GPTQ_QYBLAS_GEMM_INFO_H__ diff --git a/src/infiniop/ops/gptq_qyblas_gemm/nvidia/gptq_qyblas_gemm_nvidia.cu b/src/infiniop/ops/gptq_qyblas_gemm/nvidia/gptq_qyblas_gemm_nvidia.cu deleted file mode 100644 index 9aab07aff..000000000 --- a/src/infiniop/ops/gptq_qyblas_gemm/nvidia/gptq_qyblas_gemm_nvidia.cu +++ /dev/null @@ -1,198 +0,0 @@ -#if defined ENABLE_QY_API -#include "../../../devices/nvidia/nvidia_handle.cuh" -#include "dlblas_ext.h" -#include "gptq_qyblas_gemm_nvidia.cuh" - -inline cudaDataType_t ScalarTypeToCudaDataType( - infiniDtype_t scalar_type) { - switch (scalar_type) { - case INFINI_DTYPE_U8: - return CUDA_R_8U; - case INFINI_DTYPE_I8: - return CUDA_R_8I; - case INFINI_DTYPE_I32: - return CUDA_R_32I; - case INFINI_DTYPE_F16: - return CUDA_R_16F; - case INFINI_DTYPE_F32: - return CUDA_R_32F; - case INFINI_DTYPE_F64: - return CUDA_R_64F; - case INFINI_DTYPE_I16: - return CUDA_R_16I; - case INFINI_DTYPE_I64: - return CUDA_R_64I; - case INFINI_DTYPE_BF16: - return CUDA_R_16BF; - case INFINI_DTYPE_F8: - return (cudaDataType_t)CUDA_R_8F_E4M3; - default: - fprintf(stderr, - "Cannot convert ScalarType %d\n", - (int)scalar_type); - abort(); - } -} -namespace op::gptq_qyblas_gemm::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t b_scales_desc, - infiniopTensorDescriptor_t b_zeros_desc) { - - auto info = GptqQyblasGemmInfo::createGptqQyblasGemmInfo(out_desc, a_desc, b_desc, b_scales_desc, b_zeros_desc); - - CHECK_RESULT(info); - - size_t workspace_size = 0; - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info.take(), workspace_size, handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, - size_t workspace_size, - void *out, - const void *a, - const void *b, - void *b_scales, - void *b_zeros, - int64_t quant_type, - int64_t bit, - void *stream) const { - int64_t K = static_cast(_info.K); - - cudaDataType_t computeType_ = (cudaDataType_t)CUDA_R_32F; - cudaDataType_t kernel_Atype_, kernel_Btype_, kernel_Ctype_, kernel_Stype_, kernel_Ztype_; - auto dtype = _info.dtype; - auto weight_dtype = _info.weight_dtype; - if (_info.transpose_result) { - std::swap(a, b); - std::swap(dtype, weight_dtype); - } - kernel_Atype_ = ScalarTypeToCudaDataType(dtype); - kernel_Btype_ = ScalarTypeToCudaDataType(weight_dtype); - - if (quant_type == 0) { - if (8 == bit) { - kernel_Atype_ = (cudaDataType_t)CUDA_R_8U; - } - - if (4 == bit) { - kernel_Atype_ = (cudaDataType_t)CUDA_R_4U; - K = K * 2; - } - } - - kernel_Ctype_ = ScalarTypeToCudaDataType(_info.out_dtype); - kernel_Stype_ = ScalarTypeToCudaDataType(_info.scales_dtype); - kernel_Ztype_ = ScalarTypeToCudaDataType(_info.zeros_dtype); - - float alpha = 1.0f; - float beta = 0.0f; - - int64_t M = static_cast(_info.M); - int64_t N = static_cast(_info.N); - int64_t lda = static_cast(_info.lda); - int64_t ldb = static_cast(_info.ldb); - - int64_t scales_size_0 = static_cast(_info.scales_size_0); - int64_t scales_size_1 = static_cast(_info.scales_size_1); - - int64_t result_ld = static_cast(_info.result_ld); - - dlblasExtQuantParametersV2_t extParameters; - - if (quant_type == 0) { - extParameters.a_group_size_m = M / scales_size_1; - extParameters.a_group_size_k = K / scales_size_0; - extParameters.a_zeropoints_type = kernel_Ztype_; - extParameters.a_zeropoints = b_zeros; - extParameters.a_scales_type = kernel_Stype_; - extParameters.a_scales = b_scales; - } else if (quant_type == 1) { - extParameters.a_group_size_m = 1; - extParameters.a_group_size_k = K; - extParameters.a_zeropoints = nullptr; - extParameters.a_scales_type = kernel_Stype_; - extParameters.a_scales = b_scales; - - } else if (quant_type == 2 || quant_type == 3) { - // calculate block_shape according weight/scales shape - int block_shape = 128; - while ((M + block_shape - 1) / block_shape < scales_size_0) { - block_shape /= 2; - if (block_shape < 32) { - fprintf(stderr, - "INTERNAL ASSERT FAILED: block_shape >= 32\n" - "Invalid fp blockwise linear arguments. Weight: [%d, %d]. Scales: [%d, %d].\n", - (int)M, (int)K, (int)scales_size_0, (int)scales_size_1); - abort(); - } - } - if (!((K + block_shape - 1) / block_shape == scales_size_1)) { - fprintf(stderr, - "CHECK FAILED: (K + block_shape - 1) / block_shape == scales_size_1\n"); - abort(); - } - extParameters.a_group_size_m = block_shape; - extParameters.a_group_size_k = block_shape; - extParameters.a_scales_type = kernel_Stype_; - extParameters.a_zeropoints = nullptr; - extParameters.a_scales = b_scales; - } - bool transpose_mat_1 = _info.transa == 't'; - bool transpose_mat_2 = _info.transb == 't'; - cublasOperation_t transa = transpose_mat_1 ? CUBLAS_OP_T : CUBLAS_OP_N; - cublasOperation_t transb = transpose_mat_2 ? CUBLAS_OP_T : CUBLAS_OP_N; - - if (_info.dtype == INFINI_DTYPE_F16 || _info.dtype == INFINI_DTYPE_BF16) { - CHECK_STATUS(_opaque->internal->useCublas( - (cudaStream_t)stream, - [&](cublasHandle_t handle) { - CHECK_CUBLAS( - dlblasGemmExV2(handle, - transa, - transb, - M, - N, - K, - &alpha, - a, - kernel_Atype_, - lda, - b, - kernel_Btype_, - ldb, - &beta, - out, - kernel_Ctype_, - result_ld, - computeType_, - CUBLAS_GEMM_DEFAULT_TENSOR_OP, - &extParameters)); - return INFINI_STATUS_SUCCESS; - })); - return INFINI_STATUS_SUCCESS; - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::gptq_qyblas_gemm::nvidia -#endif diff --git a/src/infiniop/ops/gptq_qyblas_gemm/nvidia/gptq_qyblas_gemm_nvidia.cuh b/src/infiniop/ops/gptq_qyblas_gemm/nvidia/gptq_qyblas_gemm_nvidia.cuh deleted file mode 100644 index b489858d9..000000000 --- a/src/infiniop/ops/gptq_qyblas_gemm/nvidia/gptq_qyblas_gemm_nvidia.cuh +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __GPTQ_QYBLAS_GEMM_NVIDIA_API_H__ -#define __GPTQ_QYBLAS_GEMM_NVIDIA_API_H__ -#include "../gptq_qyblas_gemm.h" - -DESCRIPTOR(nvidia) - -#endif // __GPTQ_QYBLAS_GEMM_NVIDIA_API_H__ diff --git a/src/infiniop/ops/gptq_qyblas_gemm/operator.cc b/src/infiniop/ops/gptq_qyblas_gemm/operator.cc deleted file mode 100644 index e7bdf7791..000000000 --- a/src/infiniop/ops/gptq_qyblas_gemm/operator.cc +++ /dev/null @@ -1,103 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/gptq_qyblas_gemm.h" - -#if defined ENABLE_QY_API -#include "nvidia/gptq_qyblas_gemm_nvidia.cuh" -#endif - -__INFINI_C infiniStatus_t infiniopCreateGptqQyblasGemmDescriptor( - infiniopHandle_t handle, - infiniopGptqQyblasGemmDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t b_scales_desc, - infiniopTensorDescriptor_t b_zeros_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::gptq_qyblas_gemm::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - out_desc, a_desc, b_desc, b_scales_desc, b_zeros_desc); - - switch (handle->device) { - -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia) -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetGptqQyblasGemmWorkspaceSize( - infiniopGptqQyblasGemmDescriptor_t desc, - size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopGptqQyblasGemm( - infiniopGptqQyblasGemmDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *out, - const void *a, - const void *b, - void *b_scale, - void *b_zero, - int64_t quant_type, - int64_t bit, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc)->calculate( \ - workspace, workspace_size, out, a, b, b_scale, b_zero, quant_type, bit, stream); - - switch (desc->device_type) { - -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyGptqQyblasGemmDescriptor( - infiniopGptqQyblasGemmDescriptor_t desc) { - -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_QY_API - DESTROY(INFINI_DEVICE_QY, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} diff --git a/src/infiniop/ops/hardswish/cpu/hardswish_cpu.cc b/src/infiniop/ops/hardswish/cpu/hardswish_cpu.cc deleted file mode 100644 index f47198580..000000000 --- a/src/infiniop/ops/hardswish/cpu/hardswish_cpu.cc +++ /dev/null @@ -1,91 +0,0 @@ -#include "hardswish_cpu.h" - -#include - -namespace op::hardswish::cpu { -namespace { - -inline bool can_use_contiguous_fast_path(const op::elementwise::ElementwiseInfo &info) { - return info.isOutputContiguous() && info.getInputSize() == 1 && info.getInputContiguous()[0] && !info.getInputBroadcasted()[0]; -} - -template -infiniStatus_t launch_contiguous_cpu(const op::elementwise::ElementwiseInfo &info, - void *output, - const std::vector &inputs) { - const T *in = reinterpret_cast(inputs[0]); - T *out = reinterpret_cast(output); - const ptrdiff_t size = static_cast(info.getOutputSize()); - -#pragma omp parallel for if (size > 1024) - for (ptrdiff_t i = 0; i < size; ++i) { - out[i] = HardSwishOp{}(in[i]); - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - const bool fast_path = can_use_contiguous_fast_path(_info); - if (fast_path) { - switch (_dtype) { - case INFINI_DTYPE_BF16: - return launch_contiguous_cpu(_info, output, inputs); - case INFINI_DTYPE_F16: - return launch_contiguous_cpu(_info, output, inputs); - case INFINI_DTYPE_F32: - return launch_contiguous_cpu(_info, output, inputs); - case INFINI_DTYPE_F64: - return launch_contiguous_cpu(_info, output, inputs); - default: - break; - } - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::hardswish::cpu diff --git a/src/infiniop/ops/hardswish/cpu/hardswish_cpu.h b/src/infiniop/ops/hardswish/cpu/hardswish_cpu.h deleted file mode 100644 index b853663aa..000000000 --- a/src/infiniop/ops/hardswish/cpu/hardswish_cpu.h +++ /dev/null @@ -1,50 +0,0 @@ -#ifndef __HARDSWISH_CPU_H__ -#define __HARDSWISH_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" - -ELEMENTWISE_DESCRIPTOR(hardswish, cpu) - -#include -#include - -namespace op::hardswish::cpu { - -typedef struct HardSwishOp { -public: - static constexpr size_t num_inputs = 1; - - template - T operator()(const T &x) const { - const float x_f = utils::cast(x); - const float clamped = std::min(std::max(x_f + 3.0f, 0.0f), 6.0f); - const float result = x_f * clamped * (1.0f / 6.0f); - return utils::cast(result); - } -} HardSwishOp; - -typedef struct HardSwishContiguousOp { -public: - static constexpr size_t num_inputs = 1; - - template - T operator()(const T &x) const { - - T three = static_cast(3); - T zero = static_cast(0); - T six = static_cast(6); - - T scale = static_cast(0.16666667f); - - T val = x + three; - - val = std::max(zero, val); - val = std::min(six, val); - - return x * val * scale; - } -} HardSwishContiguousOp; - -} // namespace op::hardswish::cpu - -#endif diff --git a/src/infiniop/ops/hardswish/cuda/kernel.cuh b/src/infiniop/ops/hardswish/cuda/kernel.cuh deleted file mode 100644 index 25dfd55a0..000000000 --- a/src/infiniop/ops/hardswish/cuda/kernel.cuh +++ /dev/null @@ -1,79 +0,0 @@ -#ifndef __HARDSWISH_CUDA_H__ -#define __HARDSWISH_CUDA_H__ - -#include - -namespace op::hardswish::cuda { - -typedef struct HardSwishOp { -public: - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(const T &x) const { - - if constexpr (std::is_same_v) { - - const half2 three = __float2half2_rn(3.0f); - const half2 scale = __float2half2_rn(0.16666667f); - - half2 val = __hadd2(x, three); - -#if defined(ENABLE_ILUVATAR_API) - - float2 val_f = __half22float2(val); - val_f.x = fminf(fmaxf(val_f.x, 0.0f), 6.0f); - val_f.y = fminf(fmaxf(val_f.y, 0.0f), 6.0f); - val = __floats2half2_rn(val_f.x, val_f.y); -#else - - const half2 zero = __float2half2_rn(0.0f); - const half2 six = __float2half2_rn(6.0f); - -#if __CUDA_ARCH__ >= 800 - - val = __hmin2(__hmax2(val, zero), six); -#else - - val = __hmax2(val, zero); - val = __hmin2(val, six); -#endif -#endif - - return __hmul2(__hmul2(x, val), scale); - - } - - else if constexpr (std::is_same_v) { - - const float x_f = __bfloat162float(x); - - const float val = fminf(fmaxf(x_f + 3.0f, 0.0f), 6.0f); - return __float2bfloat16(x_f * val * 0.16666667f); - - } - - else if constexpr (std::is_same_v) { - const float x_f = __half2float(x); - const float val = fminf(fmaxf(x_f + 3.0f, 0.0f), 6.0f); - return __float2half(x_f * val * 0.16666667f); - - } - - else if constexpr (std::is_same_v) { - - const float val = fminf(fmaxf(x + 3.0f, 0.0f), 6.0f); - return x * val * 0.16666667f; - - } - - else if constexpr (std::is_same_v) { - const double val = fmin(fmax(x + 3.0, 0.0), 6.0); - return x * val * (1.0 / 6.0); - } - } -} HardSwishOp; - -} // namespace op::hardswish::cuda - -#endif diff --git a/src/infiniop/ops/hardswish/metax/hardswish_metax.h b/src/infiniop/ops/hardswish/metax/hardswish_metax.h deleted file mode 100644 index 16b131aa9..000000000 --- a/src/infiniop/ops/hardswish/metax/hardswish_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __HARDSWISH_METAX_API_H__ -#define __HARDSWISH_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -ELEMENTWISE_DESCRIPTOR(hardswish, metax) - -#endif // __HARDSWISH_METAX_API_H__ diff --git a/src/infiniop/ops/hardswish/metax/hardswish_metax.maca b/src/infiniop/ops/hardswish/metax/hardswish_metax.maca deleted file mode 100644 index fc57a9b20..000000000 --- a/src/infiniop/ops/hardswish/metax/hardswish_metax.maca +++ /dev/null @@ -1,58 +0,0 @@ -#include "hardswish_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::hardswish::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::HardSwishOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::HardSwishOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::HardSwishOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::HardSwishOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::hardswish::metax diff --git a/src/infiniop/ops/hardswish/moore/hardswish_moore.h b/src/infiniop/ops/hardswish/moore/hardswish_moore.h deleted file mode 100644 index e5861a158..000000000 --- a/src/infiniop/ops/hardswish/moore/hardswish_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __HARDSWISH_MOORE_API_H__ -#define __HARDSWISH_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -ELEMENTWISE_DESCRIPTOR(hardswish, moore) - -#endif // __HARDSWISH_MOORE_API_H__ diff --git a/src/infiniop/ops/hardswish/moore/hardswish_moore.mu b/src/infiniop/ops/hardswish/moore/hardswish_moore.mu deleted file mode 100644 index 3a1290b35..000000000 --- a/src/infiniop/ops/hardswish/moore/hardswish_moore.mu +++ /dev/null @@ -1,118 +0,0 @@ -#include "hardswish_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "hardswish_moore_kernel.h" - -namespace op::hardswish::moore { -namespace { - -inline bool can_use_contiguous_fast_path(const op::elementwise::ElementwiseInfo &info) { - return info.isOutputContiguous() && info.getInputSize() == 1 && - info.getInputContiguous()[0] && !info.getInputBroadcasted()[0]; -} - -template -INFINIOP_MOORE_KERNEL hardswish_contiguous_kernel(size_t numel, T *out, const T *in) { - const auto op = op::hardswish::moore::HardSwishOp{}; - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - size_t stride = blockDim.x * gridDim.x; - for (; idx < numel; idx += stride) { - out[idx] = op(in[idx]); - } -} - -template -infiniStatus_t launch_fast_path(size_t numel, - void *output, - const std::vector &inputs, - void *stream) { - if (numel == 0) { - return INFINI_STATUS_SUCCESS; - } - - constexpr int kBlockSize = 256; - int grid = static_cast((numel + kBlockSize - 1) / kBlockSize); - if (grid > 65535) { - grid = 65535; - } - - auto musa_stream = reinterpret_cast(stream); - hardswish_contiguous_kernel<<>>( - numel, - reinterpret_cast(output), - reinterpret_cast(inputs[0])); - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - // create MOORE elementwise descriptor - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - const bool fast_path = can_use_contiguous_fast_path(_info); - if (fast_path) { - switch (_dtype) { - case INFINI_DTYPE_BF16: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream); - case INFINI_DTYPE_F16: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream); - case INFINI_DTYPE_F32: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream); - case INFINI_DTYPE_F64: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream); - default: - break; - } - } - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, moore::HardSwishOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, moore::HardSwishOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, moore::HardSwishOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, moore::HardSwishOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::hardswish::moore diff --git a/src/infiniop/ops/hardswish/moore/hardswish_moore_kernel.h b/src/infiniop/ops/hardswish/moore/hardswish_moore_kernel.h deleted file mode 100644 index 60e3dbc60..000000000 --- a/src/infiniop/ops/hardswish/moore/hardswish_moore_kernel.h +++ /dev/null @@ -1,39 +0,0 @@ -#ifndef __HARDSWISH_MOORE_KERNEL_H__ -#define __HARDSWISH_MOORE_KERNEL_H__ - -#include -#include - -namespace op::hardswish::moore { - -typedef struct HardSwishOp { -public: - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(const T &x) const { - if constexpr (std::is_same_v) { - float x_f = __half2float(x); - float val = fminf(fmaxf(x_f + 3.0f, 0.0f), 6.0f); - return __float2half(x_f * val * 0.16666667f); - } else if constexpr (std::is_same_v) { - float x_f = __bfloat162float(x); - float val = fminf(fmaxf(x_f + 3.0f, 0.0f), 6.0f); - return __float2bfloat16_rn(x_f * val * 0.16666667f); - } else if constexpr (std::is_same_v) { - float val = fminf(fmaxf(x + 3.0f, 0.0f), 6.0f); - return x * val * 0.16666667f; - } else if constexpr (std::is_same_v) { - double val = fmin(fmax(x + 3.0, 0.0), 6.0); - return x * val * (1.0 / 6.0); - } else { - float x_f = static_cast(x); - float val = fminf(fmaxf(x_f + 3.0f, 0.0f), 6.0f); - return static_cast(x_f * val * 0.16666667f); - } - } -} HardSwishOp; - -} // namespace op::hardswish::moore - -#endif // __HARDSWISH_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/hardswish/nvidia/hardswish_nvidia.cu b/src/infiniop/ops/hardswish/nvidia/hardswish_nvidia.cu deleted file mode 100644 index f7736a7fd..000000000 --- a/src/infiniop/ops/hardswish/nvidia/hardswish_nvidia.cu +++ /dev/null @@ -1,115 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "hardswish_nvidia.cuh" - -#include - -namespace op::hardswish::nvidia { -namespace { - -inline bool can_use_contiguous_fast_path(const op::elementwise::ElementwiseInfo &info) { - return info.isOutputContiguous() && info.getInputSize() == 1 && info.getInputContiguous()[0] && !info.getInputBroadcasted()[0]; -} - -template -__global__ void hardswish_contiguous_kernel(size_t numel, T *out, const T *in) { - const auto op = op::hardswish::cuda::HardSwishOp{}; - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - while (idx < numel) { - out[idx] = op(in[idx]); - idx += blockDim.x * gridDim.x; - } -} - -template -infiniStatus_t launch_fast_path(size_t numel, - void *output, - const std::vector &inputs, - void *stream) { - if (numel == 0) { - return INFINI_STATUS_SUCCESS; - } - - constexpr int BLOCK_SIZE = 256; - int grid = static_cast((numel + BLOCK_SIZE - 1) / BLOCK_SIZE); - grid = std::min(grid, 65535); - - auto *out_ptr = reinterpret_cast(output); - auto *in_ptr = reinterpret_cast(inputs[0]); - auto cuda_stream = reinterpret_cast(stream); - - hardswish_contiguous_kernel<<>>(numel, out_ptr, in_ptr); - cudaError_t err = cudaGetLastError(); - return err == cudaSuccess ? INFINI_STATUS_SUCCESS : INFINI_STATUS_INTERNAL_ERROR; -} - -} // namespace - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_shape); - - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - const bool fast_path = can_use_contiguous_fast_path(_info); - if (fast_path) { - switch (_dtype) { - case INFINI_DTYPE_BF16: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream); - case INFINI_DTYPE_F16: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream); - case INFINI_DTYPE_F32: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream); - case INFINI_DTYPE_F64: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream); - default: - break; - } - } - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::HardSwishOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::HardSwishOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::HardSwishOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::HardSwishOp, double>(_info, workspace, output, inputs, stream); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::hardswish::nvidia diff --git a/src/infiniop/ops/hardswish/nvidia/hardswish_nvidia.cuh b/src/infiniop/ops/hardswish/nvidia/hardswish_nvidia.cuh deleted file mode 100644 index eac0dd994..000000000 --- a/src/infiniop/ops/hardswish/nvidia/hardswish_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __HARDSWISH_CUDA_API_H__ -#define __HARDSWISH_CUDA_API_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(hardswish, nvidia) - -#endif diff --git a/src/infiniop/ops/hardswish/operator.cc b/src/infiniop/ops/hardswish/operator.cc deleted file mode 100644 index 918a6daf4..000000000 --- a/src/infiniop/ops/hardswish/operator.cc +++ /dev/null @@ -1,169 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/hardswish.h" - -#ifdef ENABLE_CPU_API -#include "cpu/hardswish_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) -#include "nvidia/hardswish_nvidia.cuh" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/hardswish_moore.h" -#endif -#ifdef ENABLE_METAX_API -#include "metax/hardswish_metax.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateHardSwishDescriptor( - infiniopHandle_t handle, - infiniopHardSwishDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::hardswish::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output_desc, \ - {input_desc}) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetHardSwishWorkspaceSize(infiniopHardSwishDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopHardSwish( - infiniopHardSwishDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, {input}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyHardSwishDescriptor(infiniopHardSwishDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/hardtanh/cpu/hardtanh_cpu.cc b/src/infiniop/ops/hardtanh/cpu/hardtanh_cpu.cc deleted file mode 100644 index 1bd276308..000000000 --- a/src/infiniop/ops/hardtanh/cpu/hardtanh_cpu.cc +++ /dev/null @@ -1,124 +0,0 @@ -#include "hardtanh_cpu.h" - -#include - -namespace op::hardtanh::cpu { - -Descriptor::Descriptor(infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - size_t workspace_size, - infiniDevice_t device_type, - int device_id, - float min_val, - float max_val) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)), - _workspace_size(workspace_size), - _min_val(min_val), - _max_val(max_val) {} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float min_val, - float max_val) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - CHECK_SAME_SHAPE(output_shape, input_shape); - - auto info_result = op::elementwise::ElementwiseInfo::create(out_desc, input_desc_vec); - CHECK_RESULT(info_result); - - *desc_ptr = new Descriptor( - dtype, - info_result.take(), - 0, - handle->device, - handle->device_id, - min_val, - max_val); - - return INFINI_STATUS_SUCCESS; -} - -template -static infiniStatus_t launchCpuHardTanh(const op::elementwise::ElementwiseInfo &info, - void *output, - const std::vector &inputs, - float min_val, - float max_val) { - if (inputs.empty()) { - return INFINI_STATUS_BAD_PARAM; - } - - T *out = reinterpret_cast(output); - const T *in = reinterpret_cast(inputs[0]); - const auto ndim = info.getNdim(); - const auto *output_shape = info.getOutputShape(); - const auto *output_strides = info.getOutputStrides(); - const auto *input_shape = info.getInputShape(0); - const auto *input_strides = info.getInputStrides(0); - const auto *input_contiguous = info.getInputContiguous(); - ptrdiff_t output_size = info.getOutputSize(); - -#pragma omp parallel for if (output_size > 1024) - for (ptrdiff_t i = 0; i < output_size; ++i) { - const size_t out_idx = info.isOutputContiguous() - ? static_cast(i) - : op::common_cpu::indexToOffset(i, ndim, output_shape, output_strides); - const size_t in_idx = input_contiguous[0] - ? static_cast(i) - : op::common_cpu::indexToOffset(i, ndim, input_shape, input_strides); - - if constexpr (std::is_same_v || std::is_same_v) { - float value = utils::cast(in[in_idx]); - float clamped = HardTanhOp{}(value, min_val, max_val); - out[out_idx] = utils::cast(clamped); - } else { - out[out_idx] = HardTanhOp{}(in[in_idx], min_val, max_val); - } - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - (void)workspace; - (void)workspace_size; - (void)stream; - - if (inputs.size() != 1) { - return INFINI_STATUS_BAD_PARAM; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return launchCpuHardTanh(_info, output, inputs, _min_val, _max_val); - case INFINI_DTYPE_F16: - return launchCpuHardTanh(_info, output, inputs, _min_val, _max_val); - case INFINI_DTYPE_F32: - return launchCpuHardTanh(_info, output, inputs, _min_val, _max_val); - case INFINI_DTYPE_F64: - return launchCpuHardTanh(_info, output, inputs, _min_val, _max_val); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} -} // namespace op::hardtanh::cpu diff --git a/src/infiniop/ops/hardtanh/cpu/hardtanh_cpu.h b/src/infiniop/ops/hardtanh/cpu/hardtanh_cpu.h deleted file mode 100644 index 09bfb340c..000000000 --- a/src/infiniop/ops/hardtanh/cpu/hardtanh_cpu.h +++ /dev/null @@ -1,63 +0,0 @@ -#ifndef __HARDTANH_CPU_H__ -#define __HARDTANH_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" -#include - -namespace op::hardtanh::cpu { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - op::elementwise::ElementwiseInfo _info; - size_t _workspace_size; - float _min_val; - float _max_val; - - Descriptor(infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - size_t workspace_size, - infiniDevice_t device_type, - int device_id, - float min_val, - float max_val); - -public: - ~Descriptor(); - - size_t workspaceSize() const { return _workspace_size; } - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float min_val, - float max_val); - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const; - - float minVal() const { return _min_val; } - float maxVal() const { return _max_val; } -}; - -typedef struct HardTanhOp { -public: - static constexpr size_t num_inputs = 1; - - template - T operator()(const T &x, float min_val, float max_val) const { - T low = static_cast(min_val); - T high = static_cast(max_val); - T val = x < low ? low : x; - return val > high ? high : val; - } -} HardTanhOp; - -} // namespace op::hardtanh::cpu - -#endif diff --git a/src/infiniop/ops/hardtanh/cuda/kernel.cuh b/src/infiniop/ops/hardtanh/cuda/kernel.cuh deleted file mode 100644 index fa8c3d130..000000000 --- a/src/infiniop/ops/hardtanh/cuda/kernel.cuh +++ /dev/null @@ -1,44 +0,0 @@ -#ifndef __HARDTANH_CUDA_H__ -#define __HARDTANH_CUDA_H__ - -#include - -namespace op::hardtanh::cuda { - -typedef struct HardTanhOp { -public: - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(const T &x, float min_val, float max_val) const { - if constexpr (std::is_same_v) { - - float2 x_f2 = __half22float2(x); - x_f2.x = fminf(max_val, fmaxf(min_val, x_f2.x)); - x_f2.y = fminf(max_val, fmaxf(min_val, x_f2.y)); - return __float22half2_rn(x_f2); - - } else if constexpr (std::is_same_v) { - - float x_f = __bfloat162float(x); - return __float2bfloat16(fminf(max_val, fmaxf(min_val, x_f))); - - } else if constexpr (std::is_same_v) { - - float x_f = __half2float(x); - return __float2half(fminf(max_val, fmaxf(min_val, x_f))); - - } else if constexpr (std::is_same_v) { - - return fminf(max_val, fmaxf(min_val, x)); - - } else if constexpr (std::is_same_v) { - - return fmin((double)max_val, fmax((double)min_val, x)); - } - } -} HardTanhOp; - -} // namespace op::hardtanh::cuda - -#endif diff --git a/src/infiniop/ops/hardtanh/metax/hardtanh_metax.h b/src/infiniop/ops/hardtanh/metax/hardtanh_metax.h deleted file mode 100644 index 182157116..000000000 --- a/src/infiniop/ops/hardtanh/metax/hardtanh_metax.h +++ /dev/null @@ -1,48 +0,0 @@ -#ifndef __HARDTANH_METAX_API_H__ -#define __HARDTANH_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -namespace op::hardtanh::metax { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - op::elementwise::ElementwiseInfo _info; - std::unique_ptr _device_info; - size_t _workspace_size; - float _min_val; - float _max_val; - - Descriptor(infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - op::elementwise::metax::DeviceImpl *device_info, - size_t workspace_size, - infiniDevice_t device_type, - int device_id, - float min_val, - float max_val); - -public: - ~Descriptor(); - - size_t workspaceSize() const { return _workspace_size; } - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float min_val, - float max_val); - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const; -}; - -} // namespace op::hardtanh::metax - -#endif // __HARDTANH_METAX_API_H__ diff --git a/src/infiniop/ops/hardtanh/metax/hardtanh_metax.maca b/src/infiniop/ops/hardtanh/metax/hardtanh_metax.maca deleted file mode 100644 index 596316e23..000000000 --- a/src/infiniop/ops/hardtanh/metax/hardtanh_metax.maca +++ /dev/null @@ -1,95 +0,0 @@ -#include "hardtanh_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::hardtanh::metax { - -Descriptor::Descriptor(infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - op::elementwise::metax::DeviceImpl *device_info, - size_t workspace_size, - infiniDevice_t device_type, - int device_id, - float min_val, - float max_val) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)), - _device_info(device_info), - _workspace_size(workspace_size), - _min_val(min_val), - _max_val(max_val) {} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float min_val, - float max_val) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - CHECK_SAME_SHAPE(output_shape, input_shape); - - auto info_result = op::elementwise::ElementwiseInfo::create(out_desc, input_desc_vec); - CHECK_RESULT(info_result); - auto info = info_result.take(); - auto workspace_size = info.getMetaMemSize() + info.getInputSize() * sizeof(void *); - - auto device_impl_result = op::elementwise::metax::DeviceImpl::create(handle->internal()); - CHECK_RESULT(device_impl_result); - - *desc_ptr = new Descriptor( - dtype, - std::move(info), - device_impl_result.take(), - workspace_size, - handle->device, - handle->device_id, - min_val, - max_val); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::HardTanhOp, cuda_bfloat16>( - _info, workspace, output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::HardTanhOp, half>( - _info, workspace, output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::HardTanhOp, float>( - _info, workspace, output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::HardTanhOp, double>( - _info, workspace, output, inputs, stream, _min_val, _max_val); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::hardtanh::metax diff --git a/src/infiniop/ops/hardtanh/moore/hardtanh_moore.h b/src/infiniop/ops/hardtanh/moore/hardtanh_moore.h deleted file mode 100644 index 470790d52..000000000 --- a/src/infiniop/ops/hardtanh/moore/hardtanh_moore.h +++ /dev/null @@ -1,51 +0,0 @@ -#ifndef __HARDTANH_MOORE_API_H__ -#define __HARDTANH_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" - -namespace op::hardtanh::moore { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - op::elementwise::ElementwiseInfo _info; - std::unique_ptr _device_info; - size_t _workspace_size; - float _min_val; - float _max_val; - - Descriptor(infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - op::elementwise::moore::DeviceImpl *device_info, - size_t workspace_size, - infiniDevice_t device_type, - int device_id, - float min_val, - float max_val); - -public: - ~Descriptor(); - - size_t workspaceSize() const { return _workspace_size; } - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float min_val, - float max_val); - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const; - - float minVal() const { return _min_val; } - float maxVal() const { return _max_val; } -}; - -} // namespace op::hardtanh::moore - -#endif // __HARDTANH_MOORE_API_H__ diff --git a/src/infiniop/ops/hardtanh/moore/hardtanh_moore.mu b/src/infiniop/ops/hardtanh/moore/hardtanh_moore.mu deleted file mode 100644 index 40e3dbe41..000000000 --- a/src/infiniop/ops/hardtanh/moore/hardtanh_moore.mu +++ /dev/null @@ -1,158 +0,0 @@ -#include "hardtanh_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" - -#include "hardtanh_moore_kernel.h" - -namespace op::hardtanh::moore { -namespace { - -inline bool can_use_contiguous_fast_path(const op::elementwise::ElementwiseInfo &info) { - return info.isOutputContiguous() && info.getInputSize() == 1 && - info.getInputContiguous()[0] && !info.getInputBroadcasted()[0]; -} - -template -INFINIOP_MOORE_KERNEL hardtanh_contiguous_kernel(size_t numel, - T *out, - const T *in, - float min_val, - float max_val) { - const auto op = op::hardtanh::moore::HardTanhOp{}; - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - size_t stride = blockDim.x * gridDim.x; - for (; idx < numel; idx += stride) { - out[idx] = op(in[idx], min_val, max_val); - } -} - -template -infiniStatus_t launch_fast_path(size_t numel, - void *output, - const std::vector &inputs, - void *stream, - float min_val, - float max_val) { - if (numel == 0) { - return INFINI_STATUS_SUCCESS; - } - - constexpr int kBlockSize = 256; - int grid = static_cast((numel + kBlockSize - 1) / kBlockSize); - if (grid > 65535) { - grid = 65535; - } - - auto musa_stream = reinterpret_cast(stream); - hardtanh_contiguous_kernel<<>>( - numel, - reinterpret_cast(output), - reinterpret_cast(inputs[0]), - min_val, - max_val); - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -Descriptor::Descriptor(infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - op::elementwise::moore::DeviceImpl *device_info, - size_t workspace_size, - infiniDevice_t device_type, - int device_id, - float min_val, - float max_val) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)), - _device_info(device_info), - _workspace_size(workspace_size), - _min_val(min_val), - _max_val(max_val) {} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float min_val, - float max_val) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - CHECK_SAME_SHAPE(output_shape, input_shape); - - auto info_result = op::elementwise::ElementwiseInfo::create(out_desc, input_desc_vec); - CHECK_RESULT(info_result); - auto info = info_result.take(); - auto workspace_size = info.getMetaMemSize() + info.getInputSize() * sizeof(void *); - - auto device_impl_result = op::elementwise::moore::DeviceImpl::create(handle->internal()); - CHECK_RESULT(device_impl_result); - - *desc_ptr = new Descriptor( - dtype, - std::move(info), - device_impl_result.take(), - workspace_size, - handle->device, - handle->device_id, - min_val, - max_val); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - const bool fast_path = can_use_contiguous_fast_path(_info); - if (fast_path) { - switch (_dtype) { - case INFINI_DTYPE_BF16: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F16: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F32: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F64: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream, _min_val, _max_val); - default: - break; - } - } - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, moore::HardTanhOp, cuda_bfloat16>(_info, workspace, output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, moore::HardTanhOp, half>(_info, workspace, output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, moore::HardTanhOp, float>(_info, workspace, output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, moore::HardTanhOp, double>(_info, workspace, output, inputs, stream, _min_val, _max_val); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::hardtanh::moore diff --git a/src/infiniop/ops/hardtanh/moore/hardtanh_moore_kernel.h b/src/infiniop/ops/hardtanh/moore/hardtanh_moore_kernel.h deleted file mode 100644 index db0a3c024..000000000 --- a/src/infiniop/ops/hardtanh/moore/hardtanh_moore_kernel.h +++ /dev/null @@ -1,34 +0,0 @@ -#ifndef __HARDTANH_MOORE_KERNEL_H__ -#define __HARDTANH_MOORE_KERNEL_H__ - -#include -#include - -namespace op::hardtanh::moore { - -typedef struct HardTanhOp { -public: - static constexpr size_t num_inputs = 1; - - template - __device__ __forceinline__ T operator()(const T &x, float min_val, float max_val) const { - if constexpr (std::is_same_v) { - float x_f = __half2float(x); - return __float2half(fminf(max_val, fmaxf(min_val, x_f))); - } else if constexpr (std::is_same_v) { - float x_f = __bfloat162float(x); - return __float2bfloat16_rn(fminf(max_val, fmaxf(min_val, x_f))); - } else if constexpr (std::is_same_v) { - return fminf(max_val, fmaxf(min_val, x)); - } else if constexpr (std::is_same_v) { - return fmin((double)max_val, fmax((double)min_val, x)); - } else { - float x_f = static_cast(x); - return static_cast(fminf(max_val, fmaxf(min_val, x_f))); - } - } -} HardTanhOp; - -} // namespace op::hardtanh::moore - -#endif // __HARDTANH_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/hardtanh/nvidia/hardtanh_nvidia.cu b/src/infiniop/ops/hardtanh/nvidia/hardtanh_nvidia.cu deleted file mode 100644 index 31ba489ab..000000000 --- a/src/infiniop/ops/hardtanh/nvidia/hardtanh_nvidia.cu +++ /dev/null @@ -1,150 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "hardtanh_nvidia.cuh" - -#include - -namespace op::hardtanh::nvidia { -namespace { - -inline bool can_use_contiguous_fast_path(const op::elementwise::ElementwiseInfo &info) { - return info.isOutputContiguous() && info.getInputSize() == 1 && info.getInputContiguous()[0] && !info.getInputBroadcasted()[0]; -} - -template -__global__ void hardtanh_contiguous_kernel(size_t numel, T *out, const T *in, float min_val, float max_val) { - const auto op = op::hardtanh::cuda::HardTanhOp{}; - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - while (idx < numel) { - out[idx] = op(in[idx], min_val, max_val); - idx += blockDim.x * gridDim.x; - } -} - -template -infiniStatus_t launch_fast_path(size_t numel, - void *output, - const std::vector &inputs, - void *stream, - float min_val, - float max_val) { - if (numel == 0) { - return INFINI_STATUS_SUCCESS; - } - - constexpr int BLOCK_SIZE = 256; - int grid = static_cast((numel + BLOCK_SIZE - 1) / BLOCK_SIZE); - grid = std::min(grid, 65535); - - auto *out_ptr = reinterpret_cast(output); - auto *in_ptr = reinterpret_cast(inputs[0]); - auto cuda_stream = reinterpret_cast(stream); - - hardtanh_contiguous_kernel<<>>(numel, out_ptr, in_ptr, min_val, max_val); - cudaError_t err = cudaGetLastError(); - return err == cudaSuccess ? INFINI_STATUS_SUCCESS : INFINI_STATUS_INTERNAL_ERROR; -} - -} // namespace - -Descriptor::Descriptor(infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - op::elementwise::nvidia::DeviceImpl *device_info, - size_t workspace_size, - infiniDevice_t device_type, - int device_id, - float min_val, - float max_val) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)), - _device_info(device_info), - _workspace_size(workspace_size), - _min_val(min_val), - _max_val(max_val) {} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float min_val, - float max_val) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - const auto &input_desc = input_desc_vec.at(0); - const auto &output_shape = out_desc->shape(); - const auto &input_shape = input_desc->shape(); - - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - CHECK_SAME_SHAPE(output_shape, input_shape); - - auto info_result = op::elementwise::ElementwiseInfo::create(out_desc, input_desc_vec); - CHECK_RESULT(info_result); - auto info = info_result.take(); - auto workspace_size = info.getMetaMemSize() + info.getInputSize() * sizeof(void *); - - auto device_impl_result = op::elementwise::nvidia::DeviceImpl::create(handle->internal()); - CHECK_RESULT(device_impl_result); - - *desc_ptr = new Descriptor( - dtype, - std::move(info), - device_impl_result.take(), - workspace_size, - handle->device, - handle->device_id, - min_val, - max_val); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - const bool fast_path = can_use_contiguous_fast_path(_info); - if (fast_path) { - switch (_dtype) { - case INFINI_DTYPE_BF16: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F16: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F32: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F64: - return launch_fast_path(_info.getOutputSize(), output, inputs, stream, _min_val, _max_val); - default: - break; - } - } - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::HardTanhOp, cuda_bfloat16>(_info, workspace, output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::HardTanhOp, half>(_info, workspace, output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::HardTanhOp, float>(_info, workspace, output, inputs, stream, _min_val, _max_val); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::HardTanhOp, double>(_info, workspace, output, inputs, stream, _min_val, _max_val); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::hardtanh::nvidia diff --git a/src/infiniop/ops/hardtanh/nvidia/hardtanh_nvidia.cuh b/src/infiniop/ops/hardtanh/nvidia/hardtanh_nvidia.cuh deleted file mode 100644 index ebd27d80e..000000000 --- a/src/infiniop/ops/hardtanh/nvidia/hardtanh_nvidia.cuh +++ /dev/null @@ -1,51 +0,0 @@ -#ifndef __HARDTANH_CUDA_API_H__ -#define __HARDTANH_CUDA_API_H__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -namespace op::hardtanh::nvidia { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - op::elementwise::ElementwiseInfo _info; - std::unique_ptr _device_info; - size_t _workspace_size; - float _min_val; - float _max_val; - - Descriptor(infiniDtype_t dtype, - op::elementwise::ElementwiseInfo info, - op::elementwise::nvidia::DeviceImpl *device_info, - size_t workspace_size, - infiniDevice_t device_type, - int device_id, - float min_val, - float max_val); - -public: - ~Descriptor(); - - size_t workspaceSize() const { return _workspace_size; } - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec, - float min_val, - float max_val); - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const; - - float minVal() const { return _min_val; } - float maxVal() const { return _max_val; } -}; - -} // namespace op::hardtanh::nvidia - -#endif diff --git a/src/infiniop/ops/hardtanh/operator.cc b/src/infiniop/ops/hardtanh/operator.cc deleted file mode 100644 index daf7581bb..000000000 --- a/src/infiniop/ops/hardtanh/operator.cc +++ /dev/null @@ -1,173 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/hardtanh.h" - -#ifdef ENABLE_CPU_API -#include "cpu/hardtanh_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) -#include "nvidia/hardtanh_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/hardtanh_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/hardtanh_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateHardTanhDescriptor( - infiniopHandle_t handle, - infiniopHardTanhDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_desc, - float min_val, - float max_val) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::hardtanh::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output_desc, \ - {input_desc}, \ - min_val, \ - max_val) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetHardTanhWorkspaceSize(infiniopHardTanhDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopHardTanh( - infiniopHardTanhDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, {input}, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyHardTanhDescriptor(infiniopHardTanhDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/hinge_embedding_loss/cpu/hinge_embedding_loss_cpu.cc b/src/infiniop/ops/hinge_embedding_loss/cpu/hinge_embedding_loss_cpu.cc deleted file mode 100644 index fb4fc1a09..000000000 --- a/src/infiniop/ops/hinge_embedding_loss/cpu/hinge_embedding_loss_cpu.cc +++ /dev/null @@ -1,179 +0,0 @@ -#include "hinge_embedding_loss_cpu.h" -#include "../../../../utils.h" -#include "../../../tensor.h" -#include -#include -#include -#include - -namespace op::hinge_embedding_loss::cpu { - -utils::Result HingeEmbeddingLossInfo::create( - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t y_desc, - double margin, - int reduction) { - - auto input_shape = input_desc->shape(); - auto target_shape = target_desc->shape(); - auto y_shape = y_desc->shape(); - - if (input_shape != target_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - Reduction red = static_cast(reduction); - std::vector expected_y_shape; - if (red == Reduction::NONE) { - expected_y_shape = input_shape; - } else { - // Mean or Sum: output is scalar - expected_y_shape = {}; - } - - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - HingeEmbeddingLossInfo info; - info.ndim = input_shape.size(); - info.shape = input_shape; - info.input_strides = input_desc->strides(); - info.target_strides = target_desc->strides(); - if (red == Reduction::NONE) { - info.y_strides = y_desc->strides(); - } - info.input_size = input_desc->numel(); - info.margin = margin; - info.reduction = red; - - return utils::Result(std::move(info)); -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - double margin, - int reduction) { - - auto dtype = input_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - if (target_desc->dtype() != dtype || y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto info_result = HingeEmbeddingLossInfo::create(input_desc, target_desc, y_desc, margin, reduction); - CHECK_RESULT(info_result); - - *desc_ptr = new Descriptor(dtype, info_result.take(), handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -void hinge_embedding_loss_impl( - const HingeEmbeddingLossInfo &info, - T *y, - const T *input, - const T *target) { - - size_t n = info.input_size; - using Tcompute = std::conditional_t, double, float>; - Tcompute margin_val = static_cast(info.margin); - - auto linearToOffset = [&](size_t linear, const std::vector &strides) -> ptrdiff_t { - ptrdiff_t off = 0; - for (size_t d = info.ndim; d-- > 0;) { - size_t coord = (info.shape.empty() ? 0 : (linear % info.shape[d])); - if (!info.shape.empty()) { - linear /= info.shape[d]; - } - off += static_cast(coord) * strides[d]; - } - return off; - }; - - auto loss_value = [&](Tcompute in, Tcompute t) -> Tcompute { - if (t == static_cast(1)) { - return in; - } - if (t == static_cast(-1)) { - return std::max(static_cast(0), margin_val - in); - } - // PyTorch defines a fallback behavior for unexpected target values - // (i.e., not exactly ±1): loss = max(input, margin). - return std::max(in, margin_val); - }; - - if (info.reduction == Reduction::NONE) { - // Element-wise loss - for (size_t i = 0; i < n; ++i) { - ptrdiff_t in_off = linearToOffset(i, info.input_strides); - ptrdiff_t t_off = linearToOffset(i, info.target_strides); - ptrdiff_t y_off = linearToOffset(i, info.y_strides); - Tcompute t = utils::cast(target[t_off]); - Tcompute in = utils::cast(input[in_off]); - y[y_off] = utils::cast(loss_value(in, t)); - } - } else { - // Sum or Mean - Tcompute sum = static_cast(0); - for (size_t i = 0; i < n; ++i) { - ptrdiff_t in_off = linearToOffset(i, info.input_strides); - ptrdiff_t t_off = linearToOffset(i, info.target_strides); - Tcompute t = utils::cast(target[t_off]); - Tcompute in = utils::cast(input[in_off]); - sum += loss_value(in, t); - } - if (info.reduction == Reduction::MEAN) { - const Tcompute mean_val = (n > 0) ? (sum / static_cast(n)) : std::numeric_limits::quiet_NaN(); - y[0] = utils::cast(mean_val); - } else { - y[0] = utils::cast(sum); - } - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - hinge_embedding_loss_impl(_info, reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target)); - break; - case INFINI_DTYPE_BF16: - hinge_embedding_loss_impl(_info, reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target)); - break; - case INFINI_DTYPE_F32: - hinge_embedding_loss_impl(_info, reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target)); - break; - case INFINI_DTYPE_F64: - hinge_embedding_loss_impl(_info, reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target)); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::hinge_embedding_loss::cpu diff --git a/src/infiniop/ops/hinge_embedding_loss/cpu/hinge_embedding_loss_cpu.h b/src/infiniop/ops/hinge_embedding_loss/cpu/hinge_embedding_loss_cpu.h deleted file mode 100644 index 8c93b886e..000000000 --- a/src/infiniop/ops/hinge_embedding_loss/cpu/hinge_embedding_loss_cpu.h +++ /dev/null @@ -1,69 +0,0 @@ -#ifndef __HINGE_EMBEDDING_LOSS_CPU_H__ -#define __HINGE_EMBEDDING_LOSS_CPU_H__ - -#include "../../../devices/cpu/common_cpu.h" -#include "../../../operator.h" -#include - -namespace op::hinge_embedding_loss::cpu { - -enum class Reduction { - NONE = 0, - MEAN = 1, - SUM = 2 -}; - -struct HingeEmbeddingLossInfo { - size_t ndim; - std::vector shape; - std::vector input_strides; - std::vector target_strides; - std::vector y_strides; - size_t input_size; - double margin; - Reduction reduction; - - static utils::Result create( - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - infiniopTensorDescriptor_t y_desc, - double margin, - int reduction); -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - HingeEmbeddingLossInfo _info; - - Descriptor(infiniDtype_t dtype, HingeEmbeddingLossInfo info, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - double margin, - int reduction); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - void *stream) const; -}; - -} // namespace op::hinge_embedding_loss::cpu - -#endif // __HINGE_EMBEDDING_LOSS_CPU_H__ diff --git a/src/infiniop/ops/hinge_embedding_loss/cuda/kernel.cuh b/src/infiniop/ops/hinge_embedding_loss/cuda/kernel.cuh deleted file mode 100644 index 9e3b6e248..000000000 --- a/src/infiniop/ops/hinge_embedding_loss/cuda/kernel.cuh +++ /dev/null @@ -1,155 +0,0 @@ -#include -#include - -#include "../../../reduce/cuda/reduce.cuh" - -namespace op::cuda { - -template -__device__ __forceinline__ T max_(T a, T b) { - return a > b ? a : b; -} - -template -__device__ __forceinline__ Tcompute load_as(const Tin *ptr, size_t offset) { - return static_cast(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as(const half *ptr, size_t offset) { - return __half2float(ptr[offset]); -} - -template <> -__device__ __forceinline__ float load_as(const cuda_bfloat16 *ptr, size_t offset) { - return __bfloat162float(ptr[offset]); -} - -template -__device__ __forceinline__ void store_from(Tout *ptr, size_t offset, Tcompute value) { - ptr[offset] = static_cast(value); -} - -template <> -__device__ __forceinline__ void store_from(half *ptr, size_t offset, float value) { - ptr[offset] = __float2half(value); -} - -template <> -__device__ __forceinline__ void store_from(cuda_bfloat16 *ptr, size_t offset, float value) { - ptr[offset] = __float2bfloat16_rn(value); -} - -template -__device__ __forceinline__ Tcompute hinge_embedding_loss_value(Tcompute input, Tcompute target, Tcompute margin) { - // Match PyTorch behavior: - // - target == 1 : loss = input - // - target == -1 : loss = max(0, margin - input) - // - else : loss = max(input, margin) - // - // Note: While the docs state targets should be ±1, PyTorch defines a fallback - // behavior for other values (e.g., randomly-initialized float targets). - // Example (PyTorch): input=0.2, target=0.5, margin=1.0 -> loss=1.0. - if (target == static_cast(1)) { - return input; - } - if (target == static_cast(-1)) { - return max_(static_cast(0), margin - input); - } - return max_(input, margin); -} - -template -__global__ void hinge_embedding_loss_none_kernel( - T *output, - const T *input, - const T *target, - size_t n, - size_t ndim, - const size_t *__restrict__ shape, - const ptrdiff_t *__restrict__ output_strides, - const ptrdiff_t *__restrict__ input_strides, - const ptrdiff_t *__restrict__ target_strides, - bool output_contiguous, - bool input_contiguous, - bool target_contiguous, - Tcompute margin_val) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= n) { - return; - } - - const size_t out_offset = output_contiguous ? idx : indexToOffset(idx, ndim, shape, output_strides); - const size_t in_offset = input_contiguous ? idx : indexToOffset(idx, ndim, shape, input_strides); - const size_t tgt_offset = target_contiguous ? idx : indexToOffset(idx, ndim, shape, target_strides); - - const Tcompute in_val = load_as(input, in_offset); - const Tcompute tgt_val = load_as(target, tgt_offset); - const Tcompute loss = hinge_embedding_loss_value(in_val, tgt_val, margin_val); - store_from(output, out_offset, loss); -} - -template -__global__ void hinge_embedding_loss_reduce_kernel( - Tcompute *accum, - const T *input, - const T *target, - size_t n, - size_t ndim, - const size_t *__restrict__ shape, - const ptrdiff_t *__restrict__ input_strides, - const ptrdiff_t *__restrict__ target_strides, - bool input_contiguous, - bool target_contiguous, - Tcompute margin_val) { - - const size_t tid = blockIdx.x * blockDim.x + threadIdx.x; - Tcompute sum = 0; - - for (size_t i = tid; i < n; i += blockDim.x * gridDim.x) { - const size_t in_offset = input_contiguous ? i : indexToOffset(i, ndim, shape, input_strides); - const size_t tgt_offset = target_contiguous ? i : indexToOffset(i, ndim, shape, target_strides); - - const Tcompute in_val = load_as(input, in_offset); - const Tcompute tgt_val = load_as(target, tgt_offset); - sum += hinge_embedding_loss_value(in_val, tgt_val, margin_val); - } - - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - const Tcompute block_sum = BlockReduce(temp_storage).Sum(sum); - - if (threadIdx.x == 0) { - atomicAdd(accum, block_sum); - } -} - -template -__global__ void hinge_embedding_loss_finalize_kernel( - Tout *output, - const Tcompute *accum, - size_t n, - bool mean) { - - if (blockIdx.x != 0 || threadIdx.x != 0) { - return; - } - - Tcompute result = *accum; - if (mean) { - // Match PyTorch behavior: mean reduction on an empty tensor yields NaN. - if (n == 0) { -#if defined(ENABLE_HYGON_API) - result = static_cast(__builtin_nanf("")); -#else - result = std::numeric_limits::quiet_NaN(); -#endif - } else { - result = result / static_cast(n); - } - } - store_from(output, 0, result); -} - -} // namespace op::cuda diff --git a/src/infiniop/ops/hinge_embedding_loss/metax/hinge_embedding_loss_metax.h b/src/infiniop/ops/hinge_embedding_loss/metax/hinge_embedding_loss_metax.h deleted file mode 100644 index eff9457b8..000000000 --- a/src/infiniop/ops/hinge_embedding_loss/metax/hinge_embedding_loss_metax.h +++ /dev/null @@ -1,85 +0,0 @@ -#ifndef __HINGE_EMBEDDING_LOSS_METAX_H__ -#define __HINGE_EMBEDDING_LOSS_METAX_H__ - -#include "../../../operator.h" - -namespace op::hinge_embedding_loss::metax { - -enum class Reduction { - NONE = 0, - MEAN = 1, - SUM = 2 -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t ndim; - size_t input_size; - std::vector shape; - std::vector input_strides; - std::vector target_strides; - std::vector output_strides; - bool input_contiguous; - bool target_contiguous; - bool output_contiguous; - double margin; - Reduction reduction; - - Descriptor(infiniDtype_t dtype, size_t ndim, size_t input_size, - std::vector shape, - std::vector input_strides, - std::vector target_strides, - std::vector output_strides, - bool input_contiguous, - bool target_contiguous, - bool output_contiguous, - double margin, Reduction reduction, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - ndim(ndim), - input_size(input_size), - shape(std::move(shape)), - input_strides(std::move(input_strides)), - target_strides(std::move(target_strides)), - output_strides(std::move(output_strides)), - input_contiguous(input_contiguous), - target_contiguous(target_contiguous), - output_contiguous(output_contiguous), - margin(margin), - reduction(reduction) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - double margin, - int reduction); - - size_t workspaceSize() const { - // For reduction != NONE, reserve 8 bytes for an aligned accumulation scalar. - const size_t acc_bytes = (reduction == Reduction::NONE) ? 0 : 8; - const size_t shape_bytes = ndim * sizeof(size_t); - const size_t stride_bytes = ndim * sizeof(ptrdiff_t); - const size_t in_target_bytes = 2 * stride_bytes; - const size_t out_stride_bytes = (reduction == Reduction::NONE) ? stride_bytes : 0; - return acc_bytes + shape_bytes + in_target_bytes + out_stride_bytes; - } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - void *stream) const; -}; - -} // namespace op::hinge_embedding_loss::metax - -#endif // __HINGE_EMBEDDING_LOSS_METAX_H__ diff --git a/src/infiniop/ops/hinge_embedding_loss/metax/hinge_embedding_loss_metax.maca b/src/infiniop/ops/hinge_embedding_loss/metax/hinge_embedding_loss_metax.maca deleted file mode 100644 index d1188f14a..000000000 --- a/src/infiniop/ops/hinge_embedding_loss/metax/hinge_embedding_loss_metax.maca +++ /dev/null @@ -1,289 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../../../tensor.h" - -#include "../cuda/kernel.cuh" -#include "hinge_embedding_loss_metax.h" -#include - -namespace op::hinge_embedding_loss::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - double margin, - int reduction) { - - auto dtype = input_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - if (target_desc->dtype() != dtype || y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto input_shape = input_desc->shape(); - auto target_shape = target_desc->shape(); - auto y_shape = y_desc->shape(); - - if (input_shape != target_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - Reduction red = static_cast(reduction); - std::vector expected_y_shape; - if (red == Reduction::NONE) { - expected_y_shape = input_shape; - } else { - expected_y_shape = {}; - } - - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const size_t ndim = input_desc->ndim(); - auto input_strides = input_desc->strides(); - auto target_strides = target_desc->strides(); - std::vector output_strides; - if (red == Reduction::NONE) { - output_strides = y_desc->strides(); - } - - *desc_ptr = new Descriptor(dtype, ndim, input_desc->numel(), - input_shape, - std::move(input_strides), - std::move(target_strides), - std::move(output_strides), - input_desc->isContiguous(), - target_desc->isContiguous(), - y_desc->isContiguous(), - margin, red, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - void *stream) const { - - auto metax_stream = reinterpret_cast(stream); - constexpr int BLOCK_SIZE = 256; - int num_blocks = static_cast((input_size + BLOCK_SIZE - 1) / BLOCK_SIZE); - if (num_blocks < 1) { - num_blocks = 1; - } - if (num_blocks > 1024) { - num_blocks = 1024; - } - - if (workspace_size < this->workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - if (reduction == Reduction::NONE) { - size_t *shape_d = nullptr; - ptrdiff_t *input_strides_d = nullptr; - ptrdiff_t *target_strides_d = nullptr; - ptrdiff_t *output_strides_d = nullptr; - if (ndim > 0) { - shape_d = reinterpret_cast(workspace); - input_strides_d = reinterpret_cast(shape_d + ndim); - target_strides_d = input_strides_d + ndim; - output_strides_d = target_strides_d + ndim; - - CHECK_METAX(hcMemcpyAsync(shape_d, shape.data(), ndim * sizeof(size_t), hcMemcpyHostToDevice, metax_stream)); - CHECK_METAX(hcMemcpyAsync(input_strides_d, input_strides.data(), ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, metax_stream)); - CHECK_METAX(hcMemcpyAsync(target_strides_d, target_strides.data(), ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, metax_stream)); - CHECK_METAX(hcMemcpyAsync(output_strides_d, output_strides.data(), ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, metax_stream)); - } - - // Element-wise loss - switch (_dtype) { - case INFINI_DTYPE_F16: { - float margin_val = static_cast(margin); - cuda::hinge_embedding_loss_none_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - output_strides_d, - input_strides_d, - target_strides_d, - output_contiguous, - input_contiguous, - target_contiguous, - margin_val); - break; - } - case INFINI_DTYPE_BF16: { - float margin_val = static_cast(margin); - cuda::hinge_embedding_loss_none_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - output_strides_d, - input_strides_d, - target_strides_d, - output_contiguous, - input_contiguous, - target_contiguous, - margin_val); - break; - } - case INFINI_DTYPE_F32: { - float margin_val = static_cast(margin); - cuda::hinge_embedding_loss_none_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - output_strides_d, - input_strides_d, - target_strides_d, - output_contiguous, - input_contiguous, - target_contiguous, - margin_val); - break; - } - case INFINI_DTYPE_F64: { - double margin_val = margin; - cuda::hinge_embedding_loss_none_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - output_strides_d, - input_strides_d, - target_strides_d, - output_contiguous, - input_contiguous, - target_contiguous, - margin_val); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - char *base = reinterpret_cast(workspace); - char *meta = base + 8; - size_t *shape_d = reinterpret_cast(meta); - ptrdiff_t *input_strides_d = reinterpret_cast(shape_d + ndim); - ptrdiff_t *target_strides_d = input_strides_d + ndim; - - CHECK_METAX(hcMemcpyAsync(shape_d, shape.data(), ndim * sizeof(size_t), hcMemcpyHostToDevice, metax_stream)); - CHECK_METAX(hcMemcpyAsync(input_strides_d, input_strides.data(), ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, metax_stream)); - CHECK_METAX(hcMemcpyAsync(target_strides_d, target_strides.data(), ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, metax_stream)); - - const bool mean = (reduction == Reduction::MEAN); - - // Sum or Mean: reduce into workspace accumulator, then write scalar output. - switch (_dtype) { - case INFINI_DTYPE_F16: { - float margin_val = static_cast(margin); - float *accum = reinterpret_cast(base); - CHECK_METAX(hcMemsetAsync(accum, 0, sizeof(float), metax_stream)); - cuda::hinge_embedding_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - input_strides_d, - target_strides_d, - input_contiguous, - target_contiguous, - margin_val); - cuda::hinge_embedding_loss_finalize_kernel<<<1, 1, 0, metax_stream>>>( - reinterpret_cast(y), accum, input_size, mean); - break; - } - case INFINI_DTYPE_BF16: { - float margin_val = static_cast(margin); - float *accum = reinterpret_cast(base); - CHECK_METAX(hcMemsetAsync(accum, 0, sizeof(float), metax_stream)); - cuda::hinge_embedding_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - input_strides_d, - target_strides_d, - input_contiguous, - target_contiguous, - margin_val); - cuda::hinge_embedding_loss_finalize_kernel<<<1, 1, 0, metax_stream>>>( - reinterpret_cast(y), accum, input_size, mean); - break; - } - case INFINI_DTYPE_F32: { - float margin_val = static_cast(margin); - float *accum = reinterpret_cast(base); - CHECK_METAX(hcMemsetAsync(accum, 0, sizeof(float), metax_stream)); - cuda::hinge_embedding_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - input_strides_d, - target_strides_d, - input_contiguous, - target_contiguous, - margin_val); - cuda::hinge_embedding_loss_finalize_kernel<<<1, 1, 0, metax_stream>>>( - reinterpret_cast(y), accum, input_size, mean); - break; - } - case INFINI_DTYPE_F64: { - double margin_val = margin; - double *accum = reinterpret_cast(base); - CHECK_METAX(hcMemsetAsync(accum, 0, sizeof(double), metax_stream)); - cuda::hinge_embedding_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - input_strides_d, - target_strides_d, - input_contiguous, - target_contiguous, - margin_val); - cuda::hinge_embedding_loss_finalize_kernel<<<1, 1, 0, metax_stream>>>( - reinterpret_cast(y), accum, input_size, mean); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::hinge_embedding_loss::metax diff --git a/src/infiniop/ops/hinge_embedding_loss/moore/hinge_embedding_loss_moore.h b/src/infiniop/ops/hinge_embedding_loss/moore/hinge_embedding_loss_moore.h deleted file mode 100644 index 7f8148497..000000000 --- a/src/infiniop/ops/hinge_embedding_loss/moore/hinge_embedding_loss_moore.h +++ /dev/null @@ -1,85 +0,0 @@ -#ifndef __HINGE_EMBEDDING_LOSS_MOORE_H__ -#define __HINGE_EMBEDDING_LOSS_MOORE_H__ - -#include "../../../operator.h" - -namespace op::hinge_embedding_loss::moore { - -enum class Reduction { - NONE = 0, - MEAN = 1, - SUM = 2 -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t ndim; - size_t input_size; - std::vector shape; - std::vector input_strides; - std::vector target_strides; - std::vector output_strides; - bool input_contiguous; - bool target_contiguous; - bool output_contiguous; - double margin; - Reduction reduction; - - Descriptor(infiniDtype_t dtype, size_t ndim, size_t input_size, - std::vector shape, - std::vector input_strides, - std::vector target_strides, - std::vector output_strides, - bool input_contiguous, - bool target_contiguous, - bool output_contiguous, - double margin, Reduction reduction, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - ndim(ndim), - input_size(input_size), - shape(std::move(shape)), - input_strides(std::move(input_strides)), - target_strides(std::move(target_strides)), - output_strides(std::move(output_strides)), - input_contiguous(input_contiguous), - target_contiguous(target_contiguous), - output_contiguous(output_contiguous), - margin(margin), - reduction(reduction) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - double margin, - int reduction); - - size_t workspaceSize() const { - // For reduction != NONE, reserve 8 bytes for an aligned accumulation scalar. - const size_t acc_bytes = (reduction == Reduction::NONE) ? 0 : 8; - const size_t shape_bytes = ndim * sizeof(size_t); - const size_t stride_bytes = ndim * sizeof(ptrdiff_t); - const size_t in_target_bytes = 2 * stride_bytes; - const size_t out_stride_bytes = (reduction == Reduction::NONE) ? stride_bytes : 0; - return acc_bytes + shape_bytes + in_target_bytes + out_stride_bytes; - } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - void *stream) const; -}; - -} // namespace op::hinge_embedding_loss::moore - -#endif // __HINGE_EMBEDDING_LOSS_MOORE_H__ diff --git a/src/infiniop/ops/hinge_embedding_loss/moore/hinge_embedding_loss_moore.mu b/src/infiniop/ops/hinge_embedding_loss/moore/hinge_embedding_loss_moore.mu deleted file mode 100644 index 553d6d5a6..000000000 --- a/src/infiniop/ops/hinge_embedding_loss/moore/hinge_embedding_loss_moore.mu +++ /dev/null @@ -1,288 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "../../../tensor.h" - -#include "../cuda/kernel.cuh" -#include "hinge_embedding_loss_moore.h" -#include - -namespace op::hinge_embedding_loss::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - double margin, - int reduction) { - - auto dtype = input_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - if (target_desc->dtype() != dtype || y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto input_shape = input_desc->shape(); - auto target_shape = target_desc->shape(); - auto y_shape = y_desc->shape(); - - if (input_shape != target_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - Reduction red = static_cast(reduction); - std::vector expected_y_shape; - if (red == Reduction::NONE) { - expected_y_shape = input_shape; - } else { - expected_y_shape = {}; - } - - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const size_t ndim = input_desc->ndim(); - auto input_strides = input_desc->strides(); - auto target_strides = target_desc->strides(); - std::vector output_strides; - if (red == Reduction::NONE) { - output_strides = y_desc->strides(); - } - - *desc_ptr = new Descriptor(dtype, ndim, input_desc->numel(), - input_shape, - std::move(input_strides), - std::move(target_strides), - std::move(output_strides), - input_desc->isContiguous(), - target_desc->isContiguous(), - y_desc->isContiguous(), - margin, red, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - void *stream) const { - - auto musa_stream = reinterpret_cast(stream); - constexpr int BLOCK_SIZE = 256; - int num_blocks = static_cast((input_size + BLOCK_SIZE - 1) / BLOCK_SIZE); - if (num_blocks < 1) { - num_blocks = 1; - } - if (num_blocks > 1024) { - num_blocks = 1024; - } - - if (workspace_size < this->workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - if (reduction == Reduction::NONE) { - size_t *shape_d = nullptr; - ptrdiff_t *input_strides_d = nullptr; - ptrdiff_t *target_strides_d = nullptr; - ptrdiff_t *output_strides_d = nullptr; - if (ndim > 0) { - shape_d = reinterpret_cast(workspace); - input_strides_d = reinterpret_cast(shape_d + ndim); - target_strides_d = input_strides_d + ndim; - output_strides_d = target_strides_d + ndim; - - CHECK_MOORE(musaMemcpyAsync(shape_d, shape.data(), ndim * sizeof(size_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(input_strides_d, input_strides.data(), ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(target_strides_d, target_strides.data(), ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(output_strides_d, output_strides.data(), ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, musa_stream)); - } - - // Element-wise loss - switch (_dtype) { - case INFINI_DTYPE_F16: { - float margin_val = static_cast(margin); - cuda::hinge_embedding_loss_none_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - output_strides_d, - input_strides_d, - target_strides_d, - output_contiguous, - input_contiguous, - target_contiguous, - margin_val); - break; - } - case INFINI_DTYPE_BF16: { - float margin_val = static_cast(margin); - cuda::hinge_embedding_loss_none_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - output_strides_d, - input_strides_d, - target_strides_d, - output_contiguous, - input_contiguous, - target_contiguous, - margin_val); - break; - } - case INFINI_DTYPE_F32: { - float margin_val = static_cast(margin); - cuda::hinge_embedding_loss_none_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - output_strides_d, - input_strides_d, - target_strides_d, - output_contiguous, - input_contiguous, - target_contiguous, - margin_val); - break; - } - case INFINI_DTYPE_F64: { - double margin_val = margin; - cuda::hinge_embedding_loss_none_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - output_strides_d, - input_strides_d, - target_strides_d, - output_contiguous, - input_contiguous, - target_contiguous, - margin_val); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - char *base = reinterpret_cast(workspace); - char *meta = base + 8; - size_t *shape_d = reinterpret_cast(meta); - ptrdiff_t *input_strides_d = reinterpret_cast(shape_d + ndim); - ptrdiff_t *target_strides_d = input_strides_d + ndim; - - CHECK_MOORE(musaMemcpyAsync(shape_d, shape.data(), ndim * sizeof(size_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(input_strides_d, input_strides.data(), ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(target_strides_d, target_strides.data(), ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, musa_stream)); - - const bool mean = (reduction == Reduction::MEAN); - - // Sum or Mean: reduce into workspace accumulator, then write scalar output. - switch (_dtype) { - case INFINI_DTYPE_F16: { - float margin_val = static_cast(margin); - float *accum = reinterpret_cast(base); - CHECK_MOORE(musaMemsetAsync(accum, 0, sizeof(float), musa_stream)); - cuda::hinge_embedding_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - input_strides_d, - target_strides_d, - input_contiguous, - target_contiguous, - margin_val); - cuda::hinge_embedding_loss_finalize_kernel<<<1, 1, 0, musa_stream>>>( - reinterpret_cast(y), accum, input_size, mean); - break; - } - case INFINI_DTYPE_BF16: { - float margin_val = static_cast(margin); - float *accum = reinterpret_cast(base); - CHECK_MOORE(musaMemsetAsync(accum, 0, sizeof(float), musa_stream)); - cuda::hinge_embedding_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - input_strides_d, - target_strides_d, - input_contiguous, - target_contiguous, - margin_val); - cuda::hinge_embedding_loss_finalize_kernel<<<1, 1, 0, musa_stream>>>( - reinterpret_cast(y), accum, input_size, mean); - break; - } - case INFINI_DTYPE_F32: { - float margin_val = static_cast(margin); - float *accum = reinterpret_cast(base); - CHECK_MOORE(musaMemsetAsync(accum, 0, sizeof(float), musa_stream)); - cuda::hinge_embedding_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - input_strides_d, - target_strides_d, - input_contiguous, - target_contiguous, - margin_val); - cuda::hinge_embedding_loss_finalize_kernel<<<1, 1, 0, musa_stream>>>( - reinterpret_cast(y), accum, input_size, mean); - break; - } - case INFINI_DTYPE_F64: { - double margin_val = margin; - double *accum = reinterpret_cast(base); - CHECK_MOORE(musaMemsetAsync(accum, 0, sizeof(double), musa_stream)); - cuda::hinge_embedding_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - input_strides_d, - target_strides_d, - input_contiguous, - target_contiguous, - margin_val); - cuda::hinge_embedding_loss_finalize_kernel<<<1, 1, 0, musa_stream>>>( - reinterpret_cast(y), accum, input_size, mean); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::hinge_embedding_loss::moore diff --git a/src/infiniop/ops/hinge_embedding_loss/nvidia/hinge_embedding_loss_nvidia.cu b/src/infiniop/ops/hinge_embedding_loss/nvidia/hinge_embedding_loss_nvidia.cu deleted file mode 100644 index ce42711df..000000000 --- a/src/infiniop/ops/hinge_embedding_loss/nvidia/hinge_embedding_loss_nvidia.cu +++ /dev/null @@ -1,287 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "hinge_embedding_loss_nvidia.cuh" - -namespace op::hinge_embedding_loss::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - double margin, - int reduction) { - - auto dtype = input_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - if (target_desc->dtype() != dtype || y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto input_shape = input_desc->shape(); - auto target_shape = target_desc->shape(); - auto y_shape = y_desc->shape(); - - if (input_shape != target_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - Reduction red = static_cast(reduction); - std::vector expected_y_shape; - if (red == Reduction::NONE) { - expected_y_shape = input_shape; - } else { - expected_y_shape = {}; - } - - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const size_t ndim = input_desc->ndim(); - auto input_strides = input_desc->strides(); - auto target_strides = target_desc->strides(); - std::vector output_strides; - if (red == Reduction::NONE) { - output_strides = y_desc->strides(); - } - - *desc_ptr = new Descriptor(dtype, ndim, input_desc->numel(), - input_shape, - std::move(input_strides), - std::move(target_strides), - std::move(output_strides), - input_desc->isContiguous(), - target_desc->isContiguous(), - y_desc->isContiguous(), - margin, red, - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - constexpr int BLOCK_SIZE = 256; - int num_blocks = static_cast((input_size + BLOCK_SIZE - 1) / BLOCK_SIZE); - if (num_blocks < 1) { - num_blocks = 1; - } - if (num_blocks > 1024) { - num_blocks = 1024; - } - - if (workspace_size < this->workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - if (reduction == Reduction::NONE) { - size_t *shape_d = nullptr; - ptrdiff_t *input_strides_d = nullptr; - ptrdiff_t *target_strides_d = nullptr; - ptrdiff_t *output_strides_d = nullptr; - if (ndim > 0) { - shape_d = reinterpret_cast(workspace); - input_strides_d = reinterpret_cast(shape_d + ndim); - target_strides_d = input_strides_d + ndim; - output_strides_d = target_strides_d + ndim; - - CHECK_CUDA(cudaMemcpyAsync(shape_d, shape.data(), ndim * sizeof(size_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(input_strides_d, input_strides.data(), ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(target_strides_d, target_strides.data(), ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(output_strides_d, output_strides.data(), ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, cuda_stream)); - } - - // Element-wise loss - switch (_dtype) { - case INFINI_DTYPE_F16: { - float margin_val = static_cast(margin); - cuda::hinge_embedding_loss_none_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - output_strides_d, - input_strides_d, - target_strides_d, - output_contiguous, - input_contiguous, - target_contiguous, - margin_val); - break; - } - case INFINI_DTYPE_BF16: { - float margin_val = static_cast(margin); - cuda::hinge_embedding_loss_none_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - output_strides_d, - input_strides_d, - target_strides_d, - output_contiguous, - input_contiguous, - target_contiguous, - margin_val); - break; - } - case INFINI_DTYPE_F32: { - float margin_val = static_cast(margin); - cuda::hinge_embedding_loss_none_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - output_strides_d, - input_strides_d, - target_strides_d, - output_contiguous, - input_contiguous, - target_contiguous, - margin_val); - break; - } - case INFINI_DTYPE_F64: { - double margin_val = margin; - cuda::hinge_embedding_loss_none_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - output_strides_d, - input_strides_d, - target_strides_d, - output_contiguous, - input_contiguous, - target_contiguous, - margin_val); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - char *base = reinterpret_cast(workspace); - char *meta = base + 8; - size_t *shape_d = reinterpret_cast(meta); - ptrdiff_t *input_strides_d = reinterpret_cast(shape_d + ndim); - ptrdiff_t *target_strides_d = input_strides_d + ndim; - - CHECK_CUDA(cudaMemcpyAsync(shape_d, shape.data(), ndim * sizeof(size_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(input_strides_d, input_strides.data(), ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(target_strides_d, target_strides.data(), ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, cuda_stream)); - - const bool mean = (reduction == Reduction::MEAN); - - // Sum or Mean: reduce into workspace accumulator, then write scalar output. - switch (_dtype) { - case INFINI_DTYPE_F16: { - float margin_val = static_cast(margin); - float *accum = reinterpret_cast(base); - CHECK_CUDA(cudaMemsetAsync(accum, 0, sizeof(float), cuda_stream)); - cuda::hinge_embedding_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - input_strides_d, - target_strides_d, - input_contiguous, - target_contiguous, - margin_val); - cuda::hinge_embedding_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, input_size, mean); - break; - } - case INFINI_DTYPE_BF16: { - float margin_val = static_cast(margin); - float *accum = reinterpret_cast(base); - CHECK_CUDA(cudaMemsetAsync(accum, 0, sizeof(float), cuda_stream)); - cuda::hinge_embedding_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - input_strides_d, - target_strides_d, - input_contiguous, - target_contiguous, - margin_val); - cuda::hinge_embedding_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, input_size, mean); - break; - } - case INFINI_DTYPE_F32: { - float margin_val = static_cast(margin); - float *accum = reinterpret_cast(base); - CHECK_CUDA(cudaMemsetAsync(accum, 0, sizeof(float), cuda_stream)); - cuda::hinge_embedding_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - input_strides_d, - target_strides_d, - input_contiguous, - target_contiguous, - margin_val); - cuda::hinge_embedding_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, input_size, mean); - break; - } - case INFINI_DTYPE_F64: { - double margin_val = margin; - double *accum = reinterpret_cast(base); - CHECK_CUDA(cudaMemsetAsync(accum, 0, sizeof(double), cuda_stream)); - cuda::hinge_embedding_loss_reduce_kernel<<>>( - accum, - reinterpret_cast(input), - reinterpret_cast(target), - input_size, - ndim, - shape_d, - input_strides_d, - target_strides_d, - input_contiguous, - target_contiguous, - margin_val); - cuda::hinge_embedding_loss_finalize_kernel<<<1, 1, 0, cuda_stream>>>( - reinterpret_cast(y), accum, input_size, mean); - break; - } - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::hinge_embedding_loss::nvidia diff --git a/src/infiniop/ops/hinge_embedding_loss/nvidia/hinge_embedding_loss_nvidia.cuh b/src/infiniop/ops/hinge_embedding_loss/nvidia/hinge_embedding_loss_nvidia.cuh deleted file mode 100644 index d50557d26..000000000 --- a/src/infiniop/ops/hinge_embedding_loss/nvidia/hinge_embedding_loss_nvidia.cuh +++ /dev/null @@ -1,86 +0,0 @@ -#ifndef __HINGE_EMBEDDING_LOSS_NVIDIA_H__ -#define __HINGE_EMBEDDING_LOSS_NVIDIA_H__ - -#include "../../../operator.h" -#include - -namespace op::hinge_embedding_loss::nvidia { - -enum class Reduction { - NONE = 0, - MEAN = 1, - SUM = 2 -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t ndim; - size_t input_size; - std::vector shape; - std::vector input_strides; - std::vector target_strides; - std::vector output_strides; - bool input_contiguous; - bool target_contiguous; - bool output_contiguous; - double margin; - Reduction reduction; - - Descriptor(infiniDtype_t dtype, size_t ndim, size_t input_size, - std::vector shape, - std::vector input_strides, - std::vector target_strides, - std::vector output_strides, - bool input_contiguous, - bool target_contiguous, - bool output_contiguous, - double margin, Reduction reduction, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - ndim(ndim), - input_size(input_size), - shape(std::move(shape)), - input_strides(std::move(input_strides)), - target_strides(std::move(target_strides)), - output_strides(std::move(output_strides)), - input_contiguous(input_contiguous), - target_contiguous(target_contiguous), - output_contiguous(output_contiguous), - margin(margin), - reduction(reduction) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - double margin, - int reduction); - - size_t workspaceSize() const { - // For reduction != NONE, reserve 8 bytes for an aligned accumulation scalar. - const size_t acc_bytes = (reduction == Reduction::NONE) ? 0 : 8; - const size_t shape_bytes = ndim * sizeof(size_t); - const size_t stride_bytes = ndim * sizeof(ptrdiff_t); - const size_t in_target_bytes = 2 * stride_bytes; - const size_t out_stride_bytes = (reduction == Reduction::NONE) ? stride_bytes : 0; - return acc_bytes + shape_bytes + in_target_bytes + out_stride_bytes; - } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - void *stream) const; -}; - -} // namespace op::hinge_embedding_loss::nvidia - -#endif // __HINGE_EMBEDDING_LOSS_NVIDIA_H__ diff --git a/src/infiniop/ops/hinge_embedding_loss/operator.cc b/src/infiniop/ops/hinge_embedding_loss/operator.cc deleted file mode 100644 index 6ad907a41..000000000 --- a/src/infiniop/ops/hinge_embedding_loss/operator.cc +++ /dev/null @@ -1,176 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/hinge_embedding_loss.h" - -#ifdef ENABLE_CPU_API -#include "cpu/hinge_embedding_loss_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/hinge_embedding_loss_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/hinge_embedding_loss_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/hinge_embedding_loss_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateHingeEmbeddingLossDescriptor( - infiniopHandle_t handle, - infiniopHingeEmbeddingLossDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - double margin, - int reduction) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::hinge_embedding_loss::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - input_desc, \ - target_desc, \ - margin, \ - reduction) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetHingeEmbeddingLossWorkspaceSize(infiniopHingeEmbeddingLossDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopHingeEmbeddingLoss( - infiniopHingeEmbeddingLossDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *input, - const void *target, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, input, target, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyHingeEmbeddingLossDescriptor(infiniopHingeEmbeddingLossDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/histc/cpu/histc_cpu.cc b/src/infiniop/ops/histc/cpu/histc_cpu.cc deleted file mode 100644 index d10e560f8..000000000 --- a/src/infiniop/ops/histc/cpu/histc_cpu.cc +++ /dev/null @@ -1,150 +0,0 @@ -#include "histc_cpu.h" -#include "../../../../utils.h" -#include -#include -#include - -namespace op::histc::cpu { - -utils::Result HistcInfo::create( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - int64_t bins, - double min_val, - double max_val) { - - if (bins <= 0) { - return INFINI_STATUS_BAD_PARAM; - } - - if (min_val >= max_val) { - return INFINI_STATUS_BAD_PARAM; - } - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - - // Input should be 1D - if (x_shape.size() != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // Output should be 1D with bins elements - if (y_shape.size() != 1 || y_shape[0] != static_cast(bins)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - HistcInfo info; - info.input_size = x_shape[0]; - info.bins = bins; - info.min_val = min_val; - info.max_val = max_val; - info.input_stride = x_desc->strides()[0]; - info.output_stride = y_desc->strides()[0]; - - // This implementation assumes x points to the first logical element and uses linear indexing. - // Negative (or broadcasted) strides would require an explicit base offset. - if (info.input_stride <= 0) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - // Writing a histogram into a broadcasted or negatively strided output is undefined. - if (info.output_stride <= 0) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - return utils::Result(std::move(info)); -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int64_t bins, - double min_val, - double max_val) { - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - // Histc output is always float32. - if (y_desc->dtype() != INFINI_DTYPE_F32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto info_result = HistcInfo::create(x_desc, y_desc, bins, min_val, max_val); - CHECK_RESULT(info_result); - - *desc_ptr = new Descriptor(dtype, info_result.take(), handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -void histc_impl( - const HistcInfo &info, - float *y, - const T *x) { - - // Initialize output to zero (supports non-unit stride). - for (int64_t b = 0; b < info.bins; ++b) { - y[b * info.output_stride] = 0.0f; - } - - const double bin_width = (info.max_val - info.min_val) / static_cast(info.bins); - - for (size_t i = 0; i < info.input_size; ++i) { - double val = utils::cast(x[i * info.input_stride]); - - // Skip values outside range - if (val < info.min_val || val > info.max_val) { - continue; - } - - // Calculate bin index - int64_t bin_idx = static_cast((val - info.min_val) / bin_width); - - // Handle edge case: max_val should go to last bin - if (bin_idx >= info.bins) { - bin_idx = info.bins - 1; - } - if (bin_idx < 0) { - bin_idx = 0; - } - - y[bin_idx * info.output_stride] += 1.0f; - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - float *y_ptr = reinterpret_cast(y); - - switch (_dtype) { - case INFINI_DTYPE_F16: - histc_impl(_info, y_ptr, reinterpret_cast(x)); - break; - case INFINI_DTYPE_BF16: - histc_impl(_info, y_ptr, reinterpret_cast(x)); - break; - case INFINI_DTYPE_F32: - histc_impl(_info, y_ptr, reinterpret_cast(x)); - break; - case INFINI_DTYPE_F64: - histc_impl(_info, y_ptr, reinterpret_cast(x)); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::histc::cpu diff --git a/src/infiniop/ops/histc/cpu/histc_cpu.h b/src/infiniop/ops/histc/cpu/histc_cpu.h deleted file mode 100644 index 4b4f1a315..000000000 --- a/src/infiniop/ops/histc/cpu/histc_cpu.h +++ /dev/null @@ -1,60 +0,0 @@ -#ifndef __HISTC_CPU_H__ -#define __HISTC_CPU_H__ - -#include "../../../devices/cpu/common_cpu.h" -#include "../../../operator.h" -#include "../../../tensor.h" - -namespace op::histc::cpu { - -struct HistcInfo { - size_t input_size; - int64_t bins; - double min_val; - double max_val; - ptrdiff_t input_stride; - ptrdiff_t output_stride; - - static utils::Result create( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - int64_t bins, - double min_val, - double max_val); -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - HistcInfo _info; - - Descriptor(infiniDtype_t dtype, HistcInfo info, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int64_t bins, - double min_val, - double max_val); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::histc::cpu - -#endif // __HISTC_CPU_H__ diff --git a/src/infiniop/ops/histc/cuda/kernel.cuh b/src/infiniop/ops/histc/cuda/kernel.cuh deleted file mode 100644 index 75a2474ff..000000000 --- a/src/infiniop/ops/histc/cuda/kernel.cuh +++ /dev/null @@ -1,45 +0,0 @@ -#pragma once -#include - -namespace op::cuda { - -template -__global__ void histc_kernel( - float *hist, - const T *input, - size_t input_size, - ptrdiff_t input_stride, - int64_t bins, - double min_val, - double max_val) { - - size_t idx = static_cast(blockIdx.x) * static_cast(blockDim.x) + static_cast(threadIdx.x); - size_t stride = static_cast(blockDim.x) * static_cast(gridDim.x); - size_t input_stride_u = static_cast(input_stride); - - double bin_width = (max_val - min_val) / static_cast(bins); - - for (size_t i = idx; i < input_size; i += stride) { - double val = static_cast(input[i * input_stride_u]); - - // Skip values outside range - if (val < min_val || val > max_val) { - continue; - } - - // Calculate bin index - int64_t bin_idx = static_cast((val - min_val) / bin_width); - - // Handle edge case: max_val should go to last bin - if (bin_idx >= bins) { - bin_idx = bins - 1; - } - if (bin_idx < 0) { - bin_idx = 0; - } - - atomicAdd(&hist[bin_idx], 1.0f); - } -} - -} // namespace op::cuda diff --git a/src/infiniop/ops/histc/metax/histc_metax.h b/src/infiniop/ops/histc/metax/histc_metax.h deleted file mode 100644 index 6bb74a87c..000000000 --- a/src/infiniop/ops/histc/metax/histc_metax.h +++ /dev/null @@ -1,52 +0,0 @@ -#ifndef __HISTC_METAX_H__ -#define __HISTC_METAX_H__ - -#include "../../../operator.h" -#include "../../../tensor.h" - -namespace op::histc::metax { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _input_size; - int64_t _bins; - double _min_val; - double _max_val; - ptrdiff_t _input_stride; - - Descriptor(infiniDtype_t dtype, size_t input_size, int64_t bins, - double min_val, double max_val, ptrdiff_t input_stride, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _input_size(input_size), - _bins(bins), - _min_val(min_val), - _max_val(max_val), - _input_stride(input_stride) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int64_t bins, - double min_val, - double max_val); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::histc::metax - -#endif // __HISTC_METAX_H__ diff --git a/src/infiniop/ops/histc/metax/histc_metax.maca b/src/infiniop/ops/histc/metax/histc_metax.maca deleted file mode 100644 index 7d290b7d2..000000000 --- a/src/infiniop/ops/histc/metax/histc_metax.maca +++ /dev/null @@ -1,106 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" - -#include "../cuda/kernel.cuh" -#include "histc_metax.h" - -namespace op::histc::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int64_t bins, - double min_val, - double max_val) { - - if (bins <= 0 || min_val >= max_val) { - return INFINI_STATUS_BAD_PARAM; - } - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x_shape.size() != 1 || y_shape.size() != 1 || y_shape[0] != static_cast(bins)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t input_size = x_shape[0]; - ptrdiff_t input_stride = x_desc->strides()[0]; - - *desc_ptr = new Descriptor(dtype, input_size, bins, min_val, max_val, - input_stride, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - auto hc_stream = reinterpret_cast(stream); - - // Initialize output to zero - CHECK_METAX(hcMemsetAsync(y, 0, _bins * sizeof(float), hc_stream)); - - constexpr int BLOCK_SIZE = 256; - int num_blocks = (_input_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - - switch (_dtype) { - case INFINI_DTYPE_F16: - cuda::histc_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(x), - _input_size, - _input_stride, - _bins, - _min_val, - _max_val); - break; - case INFINI_DTYPE_BF16: - cuda::histc_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(x), - _input_size, - _input_stride, - _bins, - _min_val, - _max_val); - break; - case INFINI_DTYPE_F32: - cuda::histc_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(x), - _input_size, - _input_stride, - _bins, - _min_val, - _max_val); - break; - case INFINI_DTYPE_F64: - cuda::histc_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(x), - _input_size, - _input_stride, - _bins, - _min_val, - _max_val); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::histc::metax diff --git a/src/infiniop/ops/histc/moore/histc_moore.h b/src/infiniop/ops/histc/moore/histc_moore.h deleted file mode 100644 index 2d9cb1d4f..000000000 --- a/src/infiniop/ops/histc/moore/histc_moore.h +++ /dev/null @@ -1,52 +0,0 @@ -#ifndef __HISTC_MOORE_H__ -#define __HISTC_MOORE_H__ - -#include "../../../operator.h" -#include "../../../tensor.h" - -namespace op::histc::moore { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _input_size; - int64_t _bins; - double _min_val; - double _max_val; - ptrdiff_t _input_stride; - - Descriptor(infiniDtype_t dtype, size_t input_size, int64_t bins, - double min_val, double max_val, ptrdiff_t input_stride, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _input_size(input_size), - _bins(bins), - _min_val(min_val), - _max_val(max_val), - _input_stride(input_stride) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int64_t bins, - double min_val, - double max_val); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::histc::moore - -#endif // __HISTC_MOORE_H__ diff --git a/src/infiniop/ops/histc/moore/histc_moore.mu b/src/infiniop/ops/histc/moore/histc_moore.mu deleted file mode 100644 index cc3173e94..000000000 --- a/src/infiniop/ops/histc/moore/histc_moore.mu +++ /dev/null @@ -1,106 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" - -#include "../cuda/kernel.cuh" -#include "histc_moore.h" - -namespace op::histc::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int64_t bins, - double min_val, - double max_val) { - - if (bins <= 0 || min_val >= max_val) { - return INFINI_STATUS_BAD_PARAM; - } - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x_shape.size() != 1 || y_shape.size() != 1 || y_shape[0] != static_cast(bins)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t input_size = x_shape[0]; - ptrdiff_t input_stride = x_desc->strides()[0]; - - *desc_ptr = new Descriptor(dtype, input_size, bins, min_val, max_val, - input_stride, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - auto musa_stream = reinterpret_cast(stream); - - // Initialize output to zero - CHECK_MOORE(musaMemsetAsync(y, 0, _bins * sizeof(float), musa_stream)); - - constexpr int BLOCK_SIZE = 256; - int num_blocks = (_input_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - - switch (_dtype) { - case INFINI_DTYPE_F16: - cuda::histc_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(x), - _input_size, - _input_stride, - _bins, - _min_val, - _max_val); - break; - case INFINI_DTYPE_BF16: - cuda::histc_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(x), - _input_size, - _input_stride, - _bins, - _min_val, - _max_val); - break; - case INFINI_DTYPE_F32: - cuda::histc_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(x), - _input_size, - _input_stride, - _bins, - _min_val, - _max_val); - break; - case INFINI_DTYPE_F64: - cuda::histc_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(x), - _input_size, - _input_stride, - _bins, - _min_val, - _max_val); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::histc::moore diff --git a/src/infiniop/ops/histc/nvidia/histc_nvidia.cu b/src/infiniop/ops/histc/nvidia/histc_nvidia.cu deleted file mode 100644 index 43daf7d0f..000000000 --- a/src/infiniop/ops/histc/nvidia/histc_nvidia.cu +++ /dev/null @@ -1,121 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" - -#include "../cuda/kernel.cuh" -#include "histc_nvidia.cuh" - -namespace op::histc::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int64_t bins, - double min_val, - double max_val) { - - if (bins <= 0 || min_val >= max_val) { - return INFINI_STATUS_BAD_PARAM; - } - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - // Histc output is always float32. This backend also requires a contiguous output. - if (y_desc->dtype() != INFINI_DTYPE_F32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x_shape.size() != 1 || y_shape.size() != 1 || y_shape[0] != static_cast(bins)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t input_size = x_shape[0]; - ptrdiff_t input_stride = x_desc->strides()[0]; - ptrdiff_t output_stride = y_desc->strides()[0]; - - // This implementation treats y as a contiguous `float*` buffer. - if (output_stride != 1) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - // Negative (or broadcasted) strides are not supported by this kernel without an explicit base offset. - if (input_stride <= 0) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - *desc_ptr = new Descriptor(dtype, input_size, bins, min_val, max_val, - input_stride, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - - // Initialize output to zero - CHECK_CUDA(cudaMemsetAsync(y, 0, _bins * sizeof(float), cuda_stream)); - - constexpr int BLOCK_SIZE = 256; - int num_blocks = (_input_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - - switch (_dtype) { - case INFINI_DTYPE_F16: - cuda::histc_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(x), - _input_size, - _input_stride, - _bins, - _min_val, - _max_val); - break; - case INFINI_DTYPE_BF16: - cuda::histc_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(x), - _input_size, - _input_stride, - _bins, - _min_val, - _max_val); - break; - case INFINI_DTYPE_F32: - cuda::histc_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(x), - _input_size, - _input_stride, - _bins, - _min_val, - _max_val); - break; - case INFINI_DTYPE_F64: - cuda::histc_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(x), - _input_size, - _input_stride, - _bins, - _min_val, - _max_val); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::histc::nvidia diff --git a/src/infiniop/ops/histc/nvidia/histc_nvidia.cuh b/src/infiniop/ops/histc/nvidia/histc_nvidia.cuh deleted file mode 100644 index 60acb6d3d..000000000 --- a/src/infiniop/ops/histc/nvidia/histc_nvidia.cuh +++ /dev/null @@ -1,53 +0,0 @@ -#ifndef __HISTC_NVIDIA_H__ -#define __HISTC_NVIDIA_H__ - -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../operator.h" -#include "../../../tensor.h" - -namespace op::histc::nvidia { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t _input_size; - int64_t _bins; - double _min_val; - double _max_val; - ptrdiff_t _input_stride; - - Descriptor(infiniDtype_t dtype, size_t input_size, int64_t bins, - double min_val, double max_val, ptrdiff_t input_stride, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _input_size(input_size), - _bins(bins), - _min_val(min_val), - _max_val(max_val), - _input_stride(input_stride) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int64_t bins, - double min_val, - double max_val); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::histc::nvidia - -#endif // __HISTC_NVIDIA_H__ diff --git a/src/infiniop/ops/histc/operator.cc b/src/infiniop/ops/histc/operator.cc deleted file mode 100644 index ed4d07662..000000000 --- a/src/infiniop/ops/histc/operator.cc +++ /dev/null @@ -1,163 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/histc.h" - -#ifdef ENABLE_CPU_API -#include "cpu/histc_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) -#include "nvidia/histc_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/histc_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/histc_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateHistcDescriptor( - infiniopHandle_t handle, - infiniopHistcDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - int64_t bins, - double min_val, - double max_val) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::histc::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - x_desc, \ - bins, \ - min_val, \ - max_val) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetHistcWorkspaceSize(infiniopHistcDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopHistc( - infiniopHistcDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, x, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyHistcDescriptor(infiniopHistcDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/huber_loss/cpu/huber_loss_cpu.cc b/src/infiniop/ops/huber_loss/cpu/huber_loss_cpu.cc deleted file mode 100644 index 005206c96..000000000 --- a/src/infiniop/ops/huber_loss/cpu/huber_loss_cpu.cc +++ /dev/null @@ -1,139 +0,0 @@ -#include "huber_loss_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include -#include -#include -#include - -#include "../../../../utils/custom_types.h" - -namespace op::huber_loss::cpu { - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - _opaque = nullptr; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - float delta, - int reduction) { - - auto handle = reinterpret_cast(handle_); - auto result = HuberLossInfo::create(out_desc, input_desc, target_desc, delta, reduction); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - new Opaque(), - result.take(), - 0, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -void calculate_cpu_impl( - const HuberLossInfo &info, - void *output, - const void *input, - const void *target) { - - size_t count = info.count(); - float delta = info.delta(); - int reduction = info.reduction(); - - auto out_ptr = reinterpret_cast(output); - auto in_ptr = reinterpret_cast(input); - // Huber Loss 中 target 是数值，类型与 input 一致 - auto tar_ptr = reinterpret_cast(target); - float half_delta = 0.5f * delta; - - if (reduction == 0) { // None -#pragma omp parallel for schedule(static) - for (ptrdiff_t i = 0; i < (ptrdiff_t)count; ++i) { - float val = utils::cast(in_ptr[i]); - float tgt = utils::cast(tar_ptr[i]); - - float diff = val - tgt; - float abs_diff = std::abs(diff); - float loss = 0.0f; - - if (abs_diff < delta) { - // 0.5 * (x - y)^2 - loss = 0.5f * diff * diff; - } else { - // delta * (|x - y| - 0.5 * delta) - loss = delta * (abs_diff - half_delta); - } - - out_ptr[i] = utils::cast(loss); - } - } else { // Mean or Sum - double total_loss = 0.0; - -#pragma omp parallel for reduction(+ : total_loss) schedule(static) - for (ptrdiff_t i = 0; i < (ptrdiff_t)count; ++i) { - float val = utils::cast(in_ptr[i]); - float tgt = utils::cast(tar_ptr[i]); - - float diff = val - tgt; - float abs_diff = std::abs(diff); - float loss = 0.0f; - - if (abs_diff < delta) { - loss = 0.5f * diff * diff; - } else { - loss = delta * (abs_diff - half_delta); - } - - total_loss += static_cast(loss); - } - - if (reduction == 1) { // Mean - total_loss /= static_cast(count); - } - - out_ptr[0] = utils::cast(static_cast(total_loss)); - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *target, - void *stream) const { - - auto dtype = _info.dtype(); - switch (dtype) { - case INFINI_DTYPE_F32: - cpu::calculate_cpu_impl(_info, output, input, target); - break; - case INFINI_DTYPE_F64: - cpu::calculate_cpu_impl(_info, output, input, target); - break; - case INFINI_DTYPE_F16: - cpu::calculate_cpu_impl(_info, output, input, target); - break; - case INFINI_DTYPE_BF16: - cpu::calculate_cpu_impl(_info, output, input, target); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::huber_loss::cpu diff --git a/src/infiniop/ops/huber_loss/cpu/huber_loss_cpu.h b/src/infiniop/ops/huber_loss/cpu/huber_loss_cpu.h deleted file mode 100644 index 4c2df88d8..000000000 --- a/src/infiniop/ops/huber_loss/cpu/huber_loss_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __HUBER_LOSS_CPU_H__ -#define __HUBER_LOSS_CPU_H__ - -#include "../huber_loss.h" - -DESCRIPTOR(cpu) - -#endif // __HUBER_LOSS_CPU_H__ diff --git a/src/infiniop/ops/huber_loss/cuda/kernel.cuh b/src/infiniop/ops/huber_loss/cuda/kernel.cuh deleted file mode 100644 index c0e42cba4..000000000 --- a/src/infiniop/ops/huber_loss/cuda/kernel.cuh +++ /dev/null @@ -1,107 +0,0 @@ -#ifndef __HUBER_LOSS_CUDA_CUH__ -#define __HUBER_LOSS_CUDA_CUH__ - -#include -#include - -namespace op::huber_loss::cuda { - -__device__ __forceinline__ float warpReduceSum(float val) { - unsigned int mask = 0xffffffff; - for (int offset = warpSize / 2; offset > 0; offset /= 2) { - val += __shfl_down_sync(mask, val, offset); - } - return val; -} - -__device__ __forceinline__ float blockReduceSum(float val) { - static __shared__ float shared[32]; - int lane = threadIdx.x % warpSize; - int wid = threadIdx.x / warpSize; - - val = warpReduceSum(val); - if (lane == 0) { - shared[wid] = val; - } - __syncthreads(); - - val = (threadIdx.x < blockDim.x / warpSize) ? shared[lane] : 0.0f; - if (wid == 0) { - val = warpReduceSum(val); - } - return val; -} - -struct HuberLossFunctor { - float delta; - float half_delta; - - __host__ __device__ HuberLossFunctor(float delta_val) - : delta(delta_val), half_delta(0.5f * delta_val) {} - - __device__ __forceinline__ float compute(float input_val, float target_val) const { - float diff = input_val - target_val; - float abs_diff = std::abs(diff); - - if (abs_diff < delta) { - return 0.5f * diff * diff; - } else { - return delta * (abs_diff - half_delta); - } - } -}; - -template -__global__ void huber_loss_kernel( - T *__restrict__ output, - const T *__restrict__ input, - const T *__restrict__ target, - size_t count, - HuberLossFunctor functor) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < count) { - float in_val = static_cast(input[idx]); - float tg_val = static_cast(target[idx]); - - float loss = functor.compute(in_val, tg_val); - - output[idx] = static_cast(loss); - } -} - -template -__global__ void huber_loss_reduce_kernel( - float *output, - const T *__restrict__ input, - const T *__restrict__ target, - size_t count, - HuberLossFunctor functor, - float scale) { - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - size_t stride = blockDim.x * gridDim.x; - float local_sum = 0.0f; - - for (size_t i = idx; i < count; i += stride) { - float in_val = static_cast(input[i]); - float tg_val = static_cast(target[i]); - - local_sum += functor.compute(in_val, tg_val); - } - - float block_sum = blockReduceSum(local_sum); - - if (threadIdx.x == 0) { - atomicAdd(output, block_sum * scale); - } -} - -template -__global__ void cast_float_to_t(T *output, const float *src) { - *output = static_cast(*src); -} - -} // namespace op::huber_loss::cuda - -#endif // __HUBER_LOSS_CUDA_CUH__ diff --git a/src/infiniop/ops/huber_loss/huber_loss.h b/src/infiniop/ops/huber_loss/huber_loss.h deleted file mode 100644 index c98345883..000000000 --- a/src/infiniop/ops/huber_loss/huber_loss.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef __HUBER_LOSS_H__ -#define __HUBER_LOSS_H__ - -#include "../../operator.h" -#include "info.h" // 引用对应的 HuberLossInfo 定义 -#define DESCRIPTOR(NAMESPACE) \ - namespace op::huber_loss::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - HuberLossInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - HuberLossInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t input_desc, \ - infiniopTensorDescriptor_t target_desc, \ - float delta, \ - int reduction); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *output, \ - const void *input, \ - const void *target, \ - void *stream) const; \ - }; \ - } - -#endif // __HUBER_LOSS_H__ diff --git a/src/infiniop/ops/huber_loss/info.h b/src/infiniop/ops/huber_loss/info.h deleted file mode 100644 index 00d0c883f..000000000 --- a/src/infiniop/ops/huber_loss/info.h +++ /dev/null @@ -1,77 +0,0 @@ -#ifndef __HUBER_LOSS_INFO_H__ -#define __HUBER_LOSS_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" - -namespace op::huber_loss { - -class HuberLossInfo { - HuberLossInfo() = default; - -public: - int _dtype; - float _delta; - int _reduction; - size_t _count; - - int dtype() const { return _dtype; } - float delta() const { return _delta; } - int reduction() const { return _reduction; } - size_t count() const { return _count; } - - HuberLossInfo(int dtype, float delta, int reduction, size_t count) - : _dtype(dtype), _delta(delta), _reduction(reduction), _count(count) {} - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - float delta, - int reduction) { - - if (input_desc->ndim() != target_desc->ndim()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t total_count = input_desc->numel(); - if (target_desc->numel() != total_count) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - for (size_t i = 0; i < input_desc->ndim(); ++i) { - if (input_desc->shape()[i] != target_desc->shape()[i]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - if (input_desc->dtype() != target_desc->dtype() || input_desc->dtype() != out_desc->dtype()) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - if (reduction == 0) { - if (out_desc->ndim() != input_desc->ndim()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - for (size_t i = 0; i < out_desc->ndim(); ++i) { - if (out_desc->shape()[i] != input_desc->shape()[i]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - } else { - if (out_desc->numel() != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - return utils::Result(HuberLossInfo{ - input_desc->dtype(), - delta, - reduction, - total_count}); - } -}; - -} // namespace op::huber_loss - -#endif // __HUBER_LOSS_INFO_H__ diff --git a/src/infiniop/ops/huber_loss/metax/huber_loss_metax.h b/src/infiniop/ops/huber_loss/metax/huber_loss_metax.h deleted file mode 100644 index 7091e9bef..000000000 --- a/src/infiniop/ops/huber_loss/metax/huber_loss_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __HUBER_LOSS_METAX_API_H__ -#define __HUBER_LOSS_METAX_API_H__ - -#include "../huber_loss.h" - -DESCRIPTOR(metax) - -#endif // __HUBER_LOSS_METAX_API_H__ diff --git a/src/infiniop/ops/huber_loss/metax/huber_loss_metax.maca b/src/infiniop/ops/huber_loss/metax/huber_loss_metax.maca deleted file mode 100644 index f6f1cab45..000000000 --- a/src/infiniop/ops/huber_loss/metax/huber_loss_metax.maca +++ /dev/null @@ -1,291 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "huber_loss_metax.h" - -#include -#include -#include -#include - -namespace op::huber_loss::metax { - -// ================================================================== -// 1. Functor: 核心数学逻辑 -// ================================================================== -struct HuberLossFunctor { - float delta; - float half_delta; - - __host__ __device__ HuberLossFunctor(float delta_val) - : delta(delta_val), half_delta(0.5f * delta_val) {} - - __device__ __forceinline__ float compute(float input_val, float target_val) const { - float diff = input_val - target_val; - float abs_diff = fabsf(diff); - - if (abs_diff < delta) { - return 0.5f * diff * diff; - } else { - return delta * (abs_diff - half_delta); - } - } -}; - -// ================================================================== -// 2. Kernel 定义 -// ================================================================== - -// ------------------------------------------------------------------ -// Kernel 1: Elementwise 模式 (Reduction = None) -// 输出形状同 input/target，长度为 count -// ------------------------------------------------------------------ -template -__global__ void huber_loss_kernel( - T *__restrict__ output, // [count] - const T *__restrict__ input, // [count] - const T *__restrict__ target, // [count] - size_t count, // 元素个数 - HuberLossFunctor functor) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < count) { - float in_val = static_cast(input[idx]); - float tg_val = static_cast(target[idx]); - - float loss = functor.compute(in_val, tg_val); - - output[idx] = static_cast(loss); - } -} - -// ------------------------------------------------------------------ -// Kernel 2: Reduction 模式 (Mean / Sum) -// ------------------------------------------------------------------ -template -__global__ void huber_loss_reduce_kernel( - float *output, // [1] Accumulator (Float) - const T *__restrict__ input, // [count] - const T *__restrict__ target, // [count] - size_t count, // 元素个数 - HuberLossFunctor functor, - float scale // Mean: 1/count, Sum: 1.0 -) { - __shared__ volatile float shared_mem[256]; - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - size_t stride = blockDim.x * gridDim.x; - float local_sum = 0.0f; - - // 1. Grid-Stride Loop 计算本线程负责的 loss 累加 - for (size_t i = idx; i < count; i += stride) { - float in_val = static_cast(input[i]); - float tg_val = static_cast(target[i]); - - local_sum += functor.compute(in_val, tg_val); - } - - // 2. 写入 shared memory - unsigned int tid = threadIdx.x; - if (tid < 256) { - shared_mem[tid] = local_sum; - } - __syncthreads(); - - // 3. Block 内树形归约 (Unrolled Reduction) - if (tid < 128) { - shared_mem[tid] += shared_mem[tid + 128]; - } - __syncthreads(); - if (tid < 64) { - shared_mem[tid] += shared_mem[tid + 64]; - } - __syncthreads(); - if (tid < 32) { - shared_mem[tid] += shared_mem[tid + 32]; - } - __syncthreads(); - if (tid < 16) { - shared_mem[tid] += shared_mem[tid + 16]; - } - __syncthreads(); - if (tid < 8) { - shared_mem[tid] += shared_mem[tid + 8]; - } - __syncthreads(); - if (tid < 4) { - shared_mem[tid] += shared_mem[tid + 4]; - } - __syncthreads(); - if (tid < 2) { - shared_mem[tid] += shared_mem[tid + 2]; - } - __syncthreads(); - if (tid < 1) { - shared_mem[tid] += shared_mem[tid + 1]; - } - __syncthreads(); - - // 4. Block 级结果累加到全局 output - if (tid == 0) { - float block_sum = shared_mem[0]; - atomicAdd(output, block_sum * scale); - } -} - -// ------------------------------------------------------------------ -// Kernel 3: 类型转换 (Float -> T) -// ------------------------------------------------------------------ -template -__global__ void cast_float_to_t(T *output, const float *src) { - *output = static_cast(*src); -} - -// ================================================================== -// 3. Kernel Launch Logic -// ================================================================== -template -void launch_kernel( - void *output, - const void *input, - const void *target, - void *workspace, - const HuberLossInfo &info, - void *stream) { - - auto in_ptr = reinterpret_cast(input); - auto tar_ptr = reinterpret_cast(target); - auto out_ptr = reinterpret_cast(output); - - auto hc_stream = reinterpret_cast(stream); - - size_t count = info.count(); - int reduction = info.reduction(); - - HuberLossFunctor functor(info.delta()); - - // -------------------------------------------------------------- - // Mode 1: Elementwise (Reduction = None) - // -------------------------------------------------------------- - if (reduction == 0) { - size_t block_size = 256; - size_t grid_size = (count + block_size - 1) / block_size; - if (grid_size == 0) { - grid_size = 1; - } - - huber_loss_kernel - <<>>( - out_ptr, in_ptr, tar_ptr, count, functor); - } - // -------------------------------------------------------------- - // Mode 2: Reduction (Mean / Sum) - // -------------------------------------------------------------- - else { - // workspace 用作浮点累加器 - float *acc_ptr = reinterpret_cast(workspace); - hcMemsetAsync(acc_ptr, 0, sizeof(float), hc_stream); - - float scale = (reduction == 1) - ? (1.0f / static_cast(count)) // reduction == 1: mean - : 1.0f; // reduction == 2: sum - - size_t block_size = 256; - size_t grid_size = std::min( - (count + block_size - 1) / block_size, - static_cast(1024)); - if (grid_size == 0) { - grid_size = 1; - } - - huber_loss_reduce_kernel - <<>>( - acc_ptr, in_ptr, tar_ptr, count, functor, scale); - - // 将 float 标量结果转换回目标类型 T - cast_float_to_t - <<<1, 1, 0, hc_stream>>>(out_ptr, acc_ptr); - } -} - -// ================================================================== -// 4. Descriptor Implementation -// ================================================================== -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - float delta, - int reduction) { - - auto handle = reinterpret_cast(handle_); - - auto info_result = HuberLossInfo::create( - out_desc, input_desc, target_desc, delta, reduction); - if (!info_result) { - return info_result.status(); - } - - // reduction 模式下需要一个 float workspace（用于 atomicAdd 累加） - size_t workspace_size = 0; - if (reduction != 0) { - workspace_size = sizeof(float); - } - - *desc_ptr = new Descriptor( - new Opaque(), - info_result.take(), - workspace_size, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *target, - void *stream) const { - - auto dtype = _info.dtype(); - int reduction = _info.reduction(); - - if (reduction != 0 && workspace_size < sizeof(float)) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel<__half>(output, input, target, workspace, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, target, workspace, _info, stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(output, input, target, workspace, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, target, workspace, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::huber_loss::metax diff --git a/src/infiniop/ops/huber_loss/moore/huber_loss_moore.h b/src/infiniop/ops/huber_loss/moore/huber_loss_moore.h deleted file mode 100644 index a3a2b4dde..000000000 --- a/src/infiniop/ops/huber_loss/moore/huber_loss_moore.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __HUBER_LOSS_MOORE_H__ -#define __HUBER_LOSS_MOORE_H__ - -#include "../huber_loss.h" - -DESCRIPTOR(moore) - -#endif // __HUBER_LOSS_MOORE_H__ diff --git a/src/infiniop/ops/huber_loss/moore/huber_loss_moore.mu b/src/infiniop/ops/huber_loss/moore/huber_loss_moore.mu deleted file mode 100644 index 05e6ce6e3..000000000 --- a/src/infiniop/ops/huber_loss/moore/huber_loss_moore.mu +++ /dev/null @@ -1,154 +0,0 @@ -#include "../../../devices/moore/moore_handle.h" -#include "huber_loss_moore.h" -#include "huber_loss_moore_kernel.h" -#include -#include - -namespace op::huber_loss::moore { - -template -static inline bool is_aligned(const void *ptr, size_t alignment) { - return reinterpret_cast(ptr) % alignment == 0; -} - -// ================================================================== -// Kernel Launch Logic -// ================================================================== -template -void launch_kernel( - void *output, - const void *input, - const void *target, - void *workspace, - const HuberLossInfo &info, - void *stream) { - - // 1. 准备指针 - auto in_ptr = reinterpret_cast(input); - auto out_ptr = reinterpret_cast(output); - // Huber Loss 中 Target 类型与 Input 一致 - auto tar_ptr = reinterpret_cast(target); - - auto musa_stream = reinterpret_cast(stream); - - // 2. 准备参数 - size_t count = info.count(); - int reduction = info.reduction(); - - op::huber_loss::moore::HuberLossFunctor functor(info.delta()); - - // ------------------------------------------ - // 模式 1: Elementwise (Reduction = None) - // ------------------------------------------ - if (reduction == 0) { - // 每个线程处理一个元素 - size_t block_size = 256; - size_t grid_size = (count + block_size - 1) / block_size; - - op::huber_loss::moore::huber_loss_kernel - <<>>( - out_ptr, in_ptr, tar_ptr, count, functor); - } - // ------------------------------------------ - // 模式 2: Reduction (Mean / Sum) - // ------------------------------------------ - else { - // 使用 workspace 作为临时的 float 累加器 - float *acc_ptr = reinterpret_cast(workspace); - musaMemsetAsync(acc_ptr, 0, sizeof(float), musa_stream); - - // 1=Mean, 2=Sum - float scale = (reduction == 1) ? (1.0f / static_cast(count)) : 1.0f; - - size_t block_size = 256; - size_t grid_size = std::min((count + block_size - 1) / block_size, static_cast(1024)); - - op::huber_loss::moore::huber_loss_reduce_kernel - <<>>( - acc_ptr, in_ptr, tar_ptr, count, functor, scale); - - // 将 float 累加结果转回 T 写入 output - op::huber_loss::moore::cast_float_to_t - <<<1, 1, 0, musa_stream>>>(out_ptr, acc_ptr); - } -} - -// ================================================================== -// Descriptor 实现 -// ================================================================== -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - float delta, - int reduction) { - - auto handle = reinterpret_cast(handle_); - - auto info_result = HuberLossInfo::create(out_desc, input_desc, target_desc, delta, reduction); - if (!info_result) { - return info_result.status(); - } - - size_t workspace_size = 0; - if (reduction != 0) { - workspace_size = sizeof(float); - } - - *desc_ptr = new Descriptor( - new Opaque(), - info_result.take(), - workspace_size, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *target, - void *stream) const { - - auto dtype = _info.dtype(); - int reduction = _info.reduction(); - - // 检查 workspace 是否够用 - if (reduction != 0 && workspace_size < sizeof(float)) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel(output, input, target, workspace, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel<__mt_bfloat16>(output, input, target, workspace, _info, stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(output, input, target, workspace, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, target, workspace, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::huber_loss::moore diff --git a/src/infiniop/ops/huber_loss/moore/huber_loss_moore_kernel.h b/src/infiniop/ops/huber_loss/moore/huber_loss_moore_kernel.h deleted file mode 100644 index 4a8f66977..000000000 --- a/src/infiniop/ops/huber_loss/moore/huber_loss_moore_kernel.h +++ /dev/null @@ -1,146 +0,0 @@ -#ifndef __HUBER_LOSS_MOORE_KERNEL_H__ -#define __HUBER_LOSS_MOORE_KERNEL_H__ - -#include -#include -#include -#include -#include -#include - -namespace op::huber_loss::moore { -template -__device__ __forceinline__ float to_float(T val) { - if constexpr (std::is_same_v) { - return __half2float(val); - } else if constexpr (std::is_same_v) { - return __bfloat162float(val); - } else { - return static_cast(val); - } -} - -template -__device__ __forceinline__ T from_float(float val) { - if constexpr (std::is_same_v) { - return __float2half(val); - } else if constexpr (std::is_same_v) { - return __float2bfloat16(val); - } else { - return static_cast(val); - } -} -__device__ __forceinline__ float warpReduceSum(float val) { - unsigned int mask = 0xffffffff; - for (int offset = warpSize / 2; offset > 0; offset /= 2) { - val += __shfl_down_sync(mask, val, offset); - } - return val; -} - -__device__ __forceinline__ float blockReduceSum(float val) { - static __shared__ float shared[32]; // Max 1024 threads / 32 warps - int lane = threadIdx.x % warpSize; - int wid = threadIdx.x / warpSize; - - val = warpReduceSum(val); - if (lane == 0) { - shared[wid] = val; - } - __syncthreads(); - val = (threadIdx.x < blockDim.x / warpSize) ? shared[lane] : 0.0f; - if (wid == 0) { - val = warpReduceSum(val); - } - return val; -} - -struct HuberLossFunctor { - float delta; - float half_delta; // 预计算 0.5 * delta - - __host__ __device__ HuberLossFunctor(float delta_val) - : delta(delta_val), half_delta(0.5f * delta_val) {} - - // Huber Loss 计算: - // if |x - y| < delta: 0.5 * (x - y)^2 - // else: delta * (|x - y| - 0.5 * delta) - __device__ __forceinline__ float compute(float input_val, float target_val) const { - float diff = input_val - target_val; - float abs_diff = std::abs(diff); - - if (abs_diff < delta) { - return 0.5f * diff * diff; - } else { - return delta * (abs_diff - half_delta); - } - } -}; - -// ================================================================== -// Kernel 1: Reduction = None (Element-wise output) -// ================================================================== -template -__global__ void huber_loss_kernel( - T *__restrict__ output, // [N] - const T *__restrict__ input, // [N] - const T *__restrict__ target, // [N] - size_t count, // Total elements (numel) - HuberLossFunctor functor) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < count) { - float in_val = to_float(input[idx]); - float tg_val = to_float(target[idx]); - - float loss = functor.compute(in_val, tg_val); - - output[idx] = from_float(loss); - } -} - -// ================================================================== -// Kernel 2: Reduction = Mean / Sum (Scalar output) -// ================================================================== -template -__global__ void huber_loss_reduce_kernel( - float *output, // [1] Accumulator (Float) - const T *__restrict__ input, // [N] - const T *__restrict__ target, // [N] - size_t count, // Total elements - HuberLossFunctor functor, - float scale // Mean模式传 1/N, Sum模式传 1.0 -) { - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - size_t stride = blockDim.x * gridDim.x; - float local_sum = 0.0f; - - // Grid-Stride Loop over all elements - for (size_t i = idx; i < count; i += stride) { - float in_val = to_float(input[i]); - float tg_val = to_float(target[i]); - - local_sum += functor.compute(in_val, tg_val); - } - - // Block Reduction - float block_sum = blockReduceSum(local_sum); - - // Global Atomic Add (Reduce to scalar) - if (threadIdx.x == 0) { - atomicAdd(output, block_sum * scale); - } -} - -// ================================================================== -// Helper: Cast float result to T (used for scalar output) -// ================================================================== -template -__global__ void cast_float_to_t(T *output, const float *src) { - *output = from_float(*src); -} - -} // namespace op::huber_loss::moore - -#endif // __HUBER_LOSS_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/huber_loss/nvidia/huber_loss_nvidia.cu b/src/infiniop/ops/huber_loss/nvidia/huber_loss_nvidia.cu deleted file mode 100644 index 52d924e67..000000000 --- a/src/infiniop/ops/huber_loss/nvidia/huber_loss_nvidia.cu +++ /dev/null @@ -1,123 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" - -#include "../cuda/kernel.cuh" -#include "huber_loss_nvidia.cuh" -#include -#include - -namespace op::huber_loss::nvidia { -template -static inline bool is_aligned(const void *ptr, size_t alignment) { - return reinterpret_cast(ptr) % alignment == 0; -} - -template -void launch_kernel( - void *output, - const void *input, - const void *target, - void *workspace, - const HuberLossInfo &info, - void *stream) { - - auto in_ptr = reinterpret_cast(input); - auto out_ptr = reinterpret_cast(output); - auto tar_ptr = reinterpret_cast(target); - - auto cuda_stream = reinterpret_cast(stream); - - size_t count = info.count(); - int reduction = info.reduction(); - - op::huber_loss::cuda::HuberLossFunctor functor(info.delta()); - - if (reduction == 0) { - size_t block_size = 256; - size_t grid_size = (count + block_size - 1) / block_size; - - op::huber_loss::cuda::huber_loss_kernel - <<>>( - out_ptr, in_ptr, tar_ptr, count, functor); - } else { - float *acc_ptr = reinterpret_cast(workspace); - cudaMemsetAsync(acc_ptr, 0, sizeof(float), cuda_stream); - float scale = (reduction == 1) ? (1.0f / static_cast(count)) : 1.0f; - - size_t block_size = 256; - size_t grid_size = std::min((count + block_size - 1) / block_size, static_cast(1024)); - - op::huber_loss::cuda::huber_loss_reduce_kernel - <<>>( - acc_ptr, in_ptr, tar_ptr, count, functor, scale); - op::huber_loss::cuda::cast_float_to_t - <<<1, 1, 0, cuda_stream>>>(out_ptr, acc_ptr); - } -} - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t target_desc, - float delta, - int reduction) { - - auto info_result = HuberLossInfo::create(out_desc, input_desc, target_desc, delta, reduction); - if (!info_result) { - return info_result.status(); - } - size_t workspace_size = 0; - if (reduction != 0) { - workspace_size = sizeof(float); - } - - *desc_ptr = new Descriptor(new Opaque(), info_result.take(), workspace_size, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *target, - void *stream) const { - - auto dtype = _info.dtype(); - int reduction = _info.reduction(); - - if (reduction != 0 && workspace_size < sizeof(float)) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel(output, input, target, workspace, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(output, input, target, workspace, _info, stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(output, input, target, workspace, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(output, input, target, workspace, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::huber_loss::nvidia diff --git a/src/infiniop/ops/huber_loss/nvidia/huber_loss_nvidia.cuh b/src/infiniop/ops/huber_loss/nvidia/huber_loss_nvidia.cuh deleted file mode 100644 index 236e5b7e4..000000000 --- a/src/infiniop/ops/huber_loss/nvidia/huber_loss_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __HUBER_LOSS_NVIDIA_CUH__ -#define __HUBER_LOSS_NVIDIA_CUH__ - -#include "../huber_loss.h" - -DESCRIPTOR(nvidia) - -#endif // __HUBER_LOSS_NVIDIA_CUH__ diff --git a/src/infiniop/ops/huber_loss/operator.cc b/src/infiniop/ops/huber_loss/operator.cc deleted file mode 100644 index beb0237de..000000000 --- a/src/infiniop/ops/huber_loss/operator.cc +++ /dev/null @@ -1,192 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/huber_loss.h" - -// --- 后端实现头文件 --- -#ifdef ENABLE_CPU_API -#include "cpu/huber_loss_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/huber_loss_nvidia.cuh" -#endif - -#ifdef ENABLE_METAX_API -#include "metax/huber_loss_metax.h" -#endif - -#ifdef ENABLE_MOORE_API -#include "moore/huber_loss_moore.h" -#endif - -extern "C" { - -// ======================================================================= -// 1. 创建算子描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopCreateHuberLossDescriptor( - infiniopHandle_t handle, - infiniopHuberLossDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - infiniopTensorDescriptor_t target, - float delta, - int reduction) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::huber_loss::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output, \ - input, \ - target, \ - delta, \ - reduction) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ======================================================================= -// 2. 获取 Workspace 大小 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopGetHuberLossWorkspaceSize(infiniopHuberLossDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -// ======================================================================= -// 3. 执行计算 (Calculate) -// ======================================================================= -__INFINI_C infiniStatus_t infiniopHuberLoss( - infiniopHuberLossDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *target, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, input, target, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyHuberLossDescriptor(infiniopHuberLossDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} - -} // extern "C" diff --git a/src/infiniop/ops/hypot/cpu/hypot_cpu.cc b/src/infiniop/ops/hypot/cpu/hypot_cpu.cc deleted file mode 100644 index 051b763ff..000000000 --- a/src/infiniop/ops/hypot/cpu/hypot_cpu.cc +++ /dev/null @@ -1,62 +0,0 @@ -// 引用 Hypot 专用的 CPU 头文件 -#include "hypot_cpu.h" - -namespace op::hypot::cpu { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - if (input_desc_vec.size() != 2) { - return INFINI_STATUS_BAD_PARAM; - } - - const auto &input_a_desc = input_desc_vec.at(0); - const auto &input_b_desc = input_desc_vec.at(1); - const auto &output_shape = out_desc->shape(); - - CHECK_DTYPE(dtype, - INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - // 检查两个输入的形状是否与输出一致 - CHECK_SAME_SHAPE(output_shape, input_a_desc->shape()); - CHECK_SAME_SHAPE(output_shape, input_b_desc->shape()); - - // create CPU elementwise descriptor - CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_BF16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate(_info, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate(_info, output, inputs, stream); - - default: - // 如果传入了整数类型或其他不支持的类型，将返回错误 - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::hypot::cpu diff --git a/src/infiniop/ops/hypot/cpu/hypot_cpu.h b/src/infiniop/ops/hypot/cpu/hypot_cpu.h deleted file mode 100644 index e41de0e02..000000000 --- a/src/infiniop/ops/hypot/cpu/hypot_cpu.h +++ /dev/null @@ -1,30 +0,0 @@ -#ifndef __HYPOT_CPU_H__ -#define __HYPOT_CPU_H__ - -#include "../../../elementwise/cpu/elementwise_cpu.h" - -ELEMENTWISE_DESCRIPTOR(hypot, cpu) - -#include -#include - -namespace op::hypot::cpu { - -typedef struct HypotOp { -public: - // Hypot 是二元算子，计算 sqrt(x^2 + y^2) - static constexpr size_t num_inputs = 2; - - template - T operator()(const T &x, const T &y) const { - if constexpr (std::is_same_v || std::is_same_v) { - return std::hypot(x, y); - } else { - return static_cast(std::hypot(static_cast(x), static_cast(y))); - } - } -} HypotOp; - -} // namespace op::hypot::cpu - -#endif // __HYPOT_CPU_H__ diff --git a/src/infiniop/ops/hypot/cuda/kernel.cuh b/src/infiniop/ops/hypot/cuda/kernel.cuh deleted file mode 100644 index a251dbc7f..000000000 --- a/src/infiniop/ops/hypot/cuda/kernel.cuh +++ /dev/null @@ -1,52 +0,0 @@ -#ifndef __HYPOT_CUDA_H__ -#define __HYPOT_CUDA_H__ - -#include -#include - -namespace op::hypot::cuda { - -typedef struct HypotOp { -public: - static constexpr size_t num_inputs = 2; - - template - __device__ __forceinline__ T operator()(const T &x, const T &y) const { - - if constexpr (std::is_same_v) { - return sqrtf(fmaf(x, x, y * y)); - } else if constexpr (std::is_same_v) { - half2 sq_sum = __hfma2(x, x, __hmul2(y, y)); - return h2sqrt(sq_sum); - } - - else if constexpr (std::is_same_v) { - return hsqrt(__hfma(x, x, __hmul(y, y))); - } else if constexpr (std::is_same_v) { - - float f0_x = __bfloat162float(__low2bfloat16(x)); - float f1_x = __bfloat162float(__high2bfloat16(x)); - - float f0_y = __bfloat162float(__low2bfloat16(y)); - float f1_y = __bfloat162float(__high2bfloat16(y)); - float res0 = sqrtf(fmaf(f0_x, f0_x, f0_y * f0_y)); - float res1 = sqrtf(fmaf(f1_x, f1_x, f1_y * f1_y)); - - return __floats2bfloat162_rn(res0, res1); - } else if constexpr (std::is_same_v) { - float fx = __bfloat162float(x); - float fy = __bfloat162float(y); - return __float2bfloat16(sqrtf(fmaf(fx, fx, fy * fy))); - } - - else if constexpr (std::is_same_v) { - return sqrt(fma(x, x, y * y)); - } else { - return static_cast(sqrt(fma(static_cast(x), static_cast(x), static_cast(y) * static_cast(y)))); - } - } -} HypotOp; - -} // namespace op::hypot::cuda - -#endif // __HYPOT_CUDA_H__ diff --git a/src/infiniop/ops/hypot/metax/hypot_metax.h b/src/infiniop/ops/hypot/metax/hypot_metax.h deleted file mode 100644 index f312df934..000000000 --- a/src/infiniop/ops/hypot/metax/hypot_metax.h +++ /dev/null @@ -1,9 +0,0 @@ -#ifndef __HYPOT_METAX_API_H__ -#define __HYPOT_METAX_API_H__ - -#include "../../../elementwise/metax/elementwise_metax_api.h" - -// 将第一个参数修改为 hypot -ELEMENTWISE_DESCRIPTOR(hypot, metax) - -#endif // __HYPOT_METAX_API_H__ diff --git a/src/infiniop/ops/hypot/metax/hypot_metax.maca b/src/infiniop/ops/hypot/metax/hypot_metax.maca deleted file mode 100644 index 3ce904fa3..000000000 --- a/src/infiniop/ops/hypot/metax/hypot_metax.maca +++ /dev/null @@ -1,78 +0,0 @@ -#include "hypot_metax.h" - -#include "../../../elementwise/metax/elementwise_metax.h" - -#include "../cuda/kernel.cuh" - -namespace op::hypot::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - // 3. Hypot 是二元算子，需要确保有两个输入 - if (input_desc_vec.size() != 2) { - return INFINI_STATUS_BAD_PARAM; - } - const auto &input_a_desc = input_desc_vec.at(0); - const auto &input_b_desc = input_desc_vec.at(1); - - const auto &out_shape = out_desc->shape(); - const auto &in_a_shape = input_a_desc->shape(); - const auto &in_b_shape = input_b_desc->shape(); - - // 检查数据类型 (通常 Hypot 支持浮点类型) - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - // 4. 检查形状一致性 - // 注意：如果框架支持 Broadcasting (广播)，这里可能不需要严格相等 - // 但参照你的 Floor 实现，这里先进行严格形状检查 - CHECK_SAME_SHAPE(out_shape, in_a_shape); - CHECK_SAME_SHAPE(out_shape, in_b_shape); - - // create CUDA elementwise descriptor - // 宏会自动处理 input_desc_vec 里的所有输入信息 - CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // 5. 调用 calculate 并传入 cuda::HypotOp - // HypotOp 通常定义为二元运算： sqrt(x*x + y*y) - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::HypotOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, cuda::HypotOp, cuda_bfloat162>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::HypotOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::HypotOp, double>(_info, workspace, output, inputs, stream); - - // Hypot 通常返回浮点数，如果是整数输入通常需要转为浮点计算， - // 这里暂时保持与 Floor 一致的浮点支持列表。 - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::hypot::metax diff --git a/src/infiniop/ops/hypot/moore/hypot_moore.h b/src/infiniop/ops/hypot/moore/hypot_moore.h deleted file mode 100644 index 09d2fd322..000000000 --- a/src/infiniop/ops/hypot/moore/hypot_moore.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __HYPOT_MOORE_API_H__ -#define __HYPOT_MOORE_API_H__ - -#include "../../../elementwise/moore/elementwise_moore_api.h" -ELEMENTWISE_DESCRIPTOR(hypot, moore) - -#endif // __HYPOT_MOORE_API_H__ diff --git a/src/infiniop/ops/hypot/moore/hypot_moore.mu b/src/infiniop/ops/hypot/moore/hypot_moore.mu deleted file mode 100644 index 382f58410..000000000 --- a/src/infiniop/ops/hypot/moore/hypot_moore.mu +++ /dev/null @@ -1,77 +0,0 @@ -#include "hypot_moore.h" - -#include "../../../elementwise/moore/elementwise_moore.h" -#include "hypot_moore_kernel.h" - -namespace op::hypot::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - // Hypot is a binary operator (z = hypot(x, y) = sqrt(x^2 + y^2)) - // 需要确保有两个输入 - if (input_desc_vec.size() != 2) { - return INFINI_STATUS_BAD_PARAM; - } - const auto &in_desc_0 = input_desc_vec.at(0); - const auto &in_desc_1 = input_desc_vec.at(1); - - const auto &out_shape = out_desc->shape(); - const auto &in_shape_0 = in_desc_0->shape(); - const auto &in_shape_1 = in_desc_1->shape(); - - // Hypot supports floating point types. - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - // Check if output shape matches input shapes - CHECK_SAME_SHAPE(out_shape, in_shape_0); - CHECK_SAME_SHAPE(out_shape, in_shape_1); - - // create MOORE elementwise descriptor - CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // Safety check for input count in calculate phase - if (inputs.size() != 2) { - return INFINI_STATUS_BAD_PARAM; - } - - // Use moore::HypotOp template defined in hypot_moore_kernel.h - switch (_dtype) { - case INFINI_DTYPE_F16: - return _device_info->calculate<256, moore::HypotOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_BF16: - return _device_info->calculate<256, moore::HypotOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, moore::HypotOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, moore::HypotOp, double>(_info, workspace, output, inputs, stream); - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} -} // namespace op::hypot::moore diff --git a/src/infiniop/ops/hypot/moore/hypot_moore_kernel.h b/src/infiniop/ops/hypot/moore/hypot_moore_kernel.h deleted file mode 100644 index 993b065ce..000000000 --- a/src/infiniop/ops/hypot/moore/hypot_moore_kernel.h +++ /dev/null @@ -1,60 +0,0 @@ -#ifndef __HYPOT_MOORE_KERNEL_H__ -#define __HYPOT_MOORE_KERNEL_H__ - -#include -#include -#include - -namespace op::hypot::moore { - -typedef struct HypotOp { -public: - static constexpr size_t num_inputs = 2; - - template - __device__ __forceinline__ T operator()(const T &x, const T &y) const { - // ----------------------------------------------------------------- - // 1. Half2 - // ----------------------------------------------------------------- - if constexpr (std::is_same_v) { - float x_low = __low2float(x); - float x_high = __high2float(x); - float y_low = __low2float(y); - float y_high = __high2float(y); - return __floats2half2_rn(::hypotf(x_low, y_low), ::hypotf(x_high, y_high)); - } - // ----------------------------------------------------------------- - // 2. Half - // ----------------------------------------------------------------- - else if constexpr (std::is_same_v) { - return __float2half(::hypotf(__half2float(x), __half2float(y))); - } - // ----------------------------------------------------------------- - // 3. Bfloat16 (__mt_bfloat16) - // ----------------------------------------------------------------- - else if constexpr (std::is_same_v) { - // 将 __mt_bfloat16 转为 float 计算 - float x_f = __bfloat162float(x); - float y_f = __bfloat162float(y); - - // 计算结果转回 __mt_bfloat16 - return __float2bfloat16(::hypotf(x_f, y_f)); - } - // ----------------------------------------------------------------- - // 4. Float32 - // ----------------------------------------------------------------- - else if constexpr (std::is_same_v) { - return ::hypotf(x, y); - } - // ----------------------------------------------------------------- - // 5. Double / Other - // ----------------------------------------------------------------- - else { - return ::hypot(x, y); - } - } -} HypotOp; - -} // namespace op::hypot::moore - -#endif // __HYPOT_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/hypot/nvidia/hypot_nvidia.cu b/src/infiniop/ops/hypot/nvidia/hypot_nvidia.cu deleted file mode 100644 index 5e0837557..000000000 --- a/src/infiniop/ops/hypot/nvidia/hypot_nvidia.cu +++ /dev/null @@ -1,74 +0,0 @@ -#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" - -#include "../cuda/kernel.cuh" -#include "hypot_nvidia.cuh" - -namespace op::hypot::nvidia { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - std::vector input_desc_vec) { - - auto handle = reinterpret_cast(handle_); - auto dtype = out_desc->dtype(); - - if (input_desc_vec.size() != 2) { - return INFINI_STATUS_BAD_PARAM; - } - - const auto &input_a_desc = input_desc_vec.at(0); - const auto &input_b_desc = input_desc_vec.at(1); - const auto &output_shape = out_desc->shape(); - - CHECK_DTYPE(dtype, - INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64); - - CHECK_SAME_SHAPE(output_shape, input_a_desc->shape()); - CHECK_SAME_SHAPE(output_shape, input_b_desc->shape()); - - // 创建描述符 - CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - std::vector inputs, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - // ----------------------------------------------------------- - // 算子分发：将 FloorOp 替换为 HypotOp - // ----------------------------------------------------------- - switch (_dtype) { - // === 浮点类型 === - case INFINI_DTYPE_BF16: - // 注意：cuda::HypotOp 对应我们在 hypot_cuda.h 中定义的 Functor - return _device_info->calculate<256, cuda::HypotOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F16: - return _device_info->calculate<256, cuda::HypotOp, half>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F32: - return _device_info->calculate<256, cuda::HypotOp, float>(_info, workspace, output, inputs, stream); - case INFINI_DTYPE_F64: - return _device_info->calculate<256, cuda::HypotOp, double>(_info, workspace, output, inputs, stream); - - // 【修改点 4】移除了整数类型的 Case - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::hypot::nvidia diff --git a/src/infiniop/ops/hypot/nvidia/hypot_nvidia.cuh b/src/infiniop/ops/hypot/nvidia/hypot_nvidia.cuh deleted file mode 100644 index c17b34863..000000000 --- a/src/infiniop/ops/hypot/nvidia/hypot_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __HYPOT_NVIDIA_CUH__ -#define __HYPOT_NVIDIA_CUH__ - -#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" - -ELEMENTWISE_DESCRIPTOR(hypot, nvidia) - -#endif // __HYPOT_NVIDIA_CUH__ diff --git a/src/infiniop/ops/hypot/operator.cc b/src/infiniop/ops/hypot/operator.cc deleted file mode 100644 index 53f82b8ef..000000000 --- a/src/infiniop/ops/hypot/operator.cc +++ /dev/null @@ -1,175 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/hypot.h" - -// --- 后端实现头文件 --- -#ifdef ENABLE_CPU_API -#include "cpu/hypot_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/hypot_nvidia.cuh" -#endif - -#ifdef ENABLE_METAX_API -#include "metax/hypot_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/hypot_moore.h" -#endif - -extern "C" { -__INFINI_C infiniStatus_t infiniopCreateHypotDescriptor( - infiniopHandle_t handle, - infiniopHypotDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input_x, - infiniopTensorDescriptor_t input_y) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::hypot::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output, \ - {input_x, input_y}) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetHypotWorkspaceSize(infiniopHypotDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -__INFINI_C infiniStatus_t infiniopHypot( - infiniopHypotDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input_x, - const void *input_y, - void *stream) { -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, {input_x, input_y}, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyHypotDescriptor(infiniopHypotDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} - -} // extern "C" diff --git a/src/infiniop/ops/index_add/cpu/index_add_cpu.cc b/src/infiniop/ops/index_add/cpu/index_add_cpu.cc deleted file mode 100644 index cc82b7510..000000000 --- a/src/infiniop/ops/index_add/cpu/index_add_cpu.cc +++ /dev/null @@ -1,152 +0,0 @@ -#include "index_add_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include -#include -#include // 必需：用于 std::conditional -#include - -namespace op::index_add::cpu { - -Descriptor::~Descriptor() = default; - -// ================================================================== -// 创建描述符 -// ================================================================== -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - int64_t dim, - infiniopTensorDescriptor_t index_desc, - infiniopTensorDescriptor_t source_desc, - float alpha) { - - auto handle = reinterpret_cast(handle_); - - auto result = IndexAddInfo::create(out_desc, in_desc, dim, index_desc, source_desc, alpha); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - nullptr, // Opaque* - result.take(), // Info - 0, // Workspace Size - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -void calculate_cpu_impl( - const IndexAddInfo &info, - void *output, - const void *input, - const void *index, - const void *source) { - - using CalcType = typename std::conditional::value, double, float>::type; - - // 转换 Alpha (使用 utils::cast 处理自定义类型) - CalcType alpha_val = utils::cast(info.alpha()); - size_t outer_size = info.outer_size(); - size_t inner_size = info.inner_size(); - size_t dim_size = info.dim_size(); - size_t index_len = info.index_len(); - - auto out_ptr = reinterpret_cast(output); - auto src_ptr = reinterpret_cast(source); - auto idx_ptr = reinterpret_cast(index); - - for (size_t o = 0; o < outer_size; ++o) { - for (size_t i = 0; i < index_len; ++i) { - - TIdx idx = idx_ptr[i]; - - // 处理负索引 - if (idx < 0) { - idx += static_cast(dim_size); - } - - // 边界检查 - if (idx < 0 || static_cast(idx) >= dim_size) { - continue; - } - - // 计算偏移 - size_t src_offset = o * index_len * inner_size + i * inner_size; - size_t out_offset = o * dim_size * inner_size + static_cast(idx) * inner_size; - - // Inner 维度循环 - for (size_t in = 0; in < inner_size; ++in) { - CalcType src_val = utils::cast(src_ptr[src_offset + in]); - CalcType out_old_val = utils::cast(out_ptr[out_offset + in]); - - // 2. 执行计算: out = out + src * alpha - CalcType result_val = out_old_val + src_val * alpha_val; - - // 3. 使用 utils::cast 转回 TData 并写入 - out_ptr[out_offset + in] = utils::cast(result_val); - } - } - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *index, - const void *source, - void *stream) const { - - auto dtype = _info.dtype(); - auto idx_dtype = _info.idx_dtype(); - -#define DISPATCH(TDATA, TIDX) \ - calculate_cpu_impl(_info, output, input, index, source); \ - return INFINI_STATUS_SUCCESS - - if (idx_dtype == INFINI_DTYPE_I32) { - switch (dtype) { - case INFINI_DTYPE_F32: - DISPATCH(float, int32_t); - case INFINI_DTYPE_F64: - DISPATCH(double, int32_t); - case INFINI_DTYPE_F16: - DISPATCH(fp16_t, int32_t); - case INFINI_DTYPE_BF16: - DISPATCH(bf16_t, int32_t); - case INFINI_DTYPE_I32: - DISPATCH(int32_t, int32_t); - case INFINI_DTYPE_I64: - DISPATCH(int64_t, int32_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (idx_dtype == INFINI_DTYPE_I64) { - switch (dtype) { - case INFINI_DTYPE_F32: - DISPATCH(float, int64_t); - case INFINI_DTYPE_F64: - DISPATCH(double, int64_t); - case INFINI_DTYPE_F16: - DISPATCH(fp16_t, int64_t); - case INFINI_DTYPE_BF16: - DISPATCH(bf16_t, int64_t); - case INFINI_DTYPE_I32: - DISPATCH(int32_t, int64_t); - case INFINI_DTYPE_I64: - DISPATCH(int64_t, int64_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - return INFINI_STATUS_BAD_TENSOR_DTYPE; - -#undef DISPATCH -} - -} // namespace op::index_add::cpu diff --git a/src/infiniop/ops/index_add/cpu/index_add_cpu.h b/src/infiniop/ops/index_add/cpu/index_add_cpu.h deleted file mode 100644 index ebf0fcd9d..000000000 --- a/src/infiniop/ops/index_add/cpu/index_add_cpu.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __INDEX_ADD_CPU_H__ -#define __INDEX_ADD_CPU_H__ - -#include "../index_add.h" -DESCRIPTOR(cpu) - -#endif // __INDEX_ADD_CPU_H__ diff --git a/src/infiniop/ops/index_add/cuda/kernel.cuh b/src/infiniop/ops/index_add/cuda/kernel.cuh deleted file mode 100644 index 929af819d..000000000 --- a/src/infiniop/ops/index_add/cuda/kernel.cuh +++ /dev/null @@ -1,186 +0,0 @@ -#ifndef __INDEX_ADD_CUDA_H__ -#define __INDEX_ADD_CUDA_H__ - -#include - -namespace op::index_add::cuda { - -__device__ __forceinline__ void atomic_add_custom(__half *address, __half val) { -#if __CUDA_ARCH__ >= 700 - atomicAdd(address, val); -#else - // Fallback for older architectures (< Volta) - unsigned int *address_as_ui = (unsigned int *)((char *)address - ((size_t)address & 2)); - unsigned int old = *address_as_ui; - unsigned int assumed; - - do { - assumed = old; - unsigned short old_val_raw = (size_t)address & 2 ? (old >> 16) : (old & 0xffff); - __half old_val = *reinterpret_cast<__half *>(&old_val_raw); - - __half new_val = old_val + val; - unsigned short new_val_raw = *reinterpret_cast(&new_val); - - unsigned int new_int = (size_t)address & 2 ? (old & 0xffff) | (new_val_raw << 16) - : (old & 0xffff0000) | new_val_raw; - - old = atomicCAS(address_as_ui, assumed, new_int); - } while (assumed != old); -#endif -} - -__device__ __forceinline__ void atomic_add_custom(cuda_bfloat16 *address, cuda_bfloat16 val) { -#if __CUDA_ARCH__ >= 800 - atomicAdd(address, val); -#else - // Fallback for older architectures (< Ampere) - unsigned int *address_as_ui = (unsigned int *)((char *)address - ((size_t)address & 2)); - unsigned int old = *address_as_ui; - unsigned int assumed; - - do { - assumed = old; - unsigned short old_val_raw = (size_t)address & 2 ? (old >> 16) : (old & 0xffff); - cuda_bfloat16 old_val = *reinterpret_cast(&old_val_raw); - - cuda_bfloat16 new_val = old_val + val; - unsigned short new_val_raw = *reinterpret_cast(&new_val); - - unsigned int new_int = (size_t)address & 2 ? (old & 0xffff) | (new_val_raw << 16) - : (old & 0xffff0000) | new_val_raw; - - old = atomicCAS(address_as_ui, assumed, new_int); - } while (assumed != old); -#endif -} - -__device__ __forceinline__ void atomic_add_custom(int64_t *address, int64_t val) { - atomicAdd(reinterpret_cast(address), static_cast(val)); -} - -// --- 通用模板 (float, double, int32 等) --- -template -__device__ __forceinline__ void atomic_add_custom(T *address, T val) { - atomicAdd(address, val); -} - -// ================================================================== -// 2. 定义向量化数据包 (Aligned Pack) -// ================================================================== -template -struct alignas(sizeof(T) * N) Pack { - T val[N]; -}; - -// ================================================================== -// 3. 标量版 Kernel (通用 fallback) -// ================================================================== -template -__global__ void index_add_kernel( - T *__restrict__ output, - const T *__restrict__ source, - const TIdx *__restrict__ indices, - size_t outer_size, // dim 左边的维度积 - size_t inner_size, // dim 右边的维度积 - size_t dim_size, // output 在 dim 维度的长度 - size_t index_len, // index 的长度 (source 在 dim 维度的长度) - size_t num_source, // source 的总元素数 - float alpha // 缩放因子 -) { - size_t tid = blockIdx.x * blockDim.x + threadIdx.x; - size_t stride = blockDim.x * gridDim.x; - - T alpha_val = static_cast(alpha); - - // Grid-Stride Loop 遍历 Source 张量 - for (size_t i = tid; i < num_source; i += stride) { - // 1. 将线性索引 i 转换为逻辑坐标 (outer, idx_idx, inner) - // Source Shape: [Outer, IndexLen, Inner] - size_t inner_idx = i % inner_size; - size_t tmp = i / inner_size; - size_t idx_idx = tmp % index_len; - size_t outer_idx = tmp / index_len; - - // 2. 读取索引值 - TIdx target_dim_idx = indices[idx_idx]; - - // 3. 处理负索引 (防御性) - if (target_dim_idx < 0) { - target_dim_idx += static_cast(dim_size); - } - - // 4. 边界检查与原子累加 - if (target_dim_idx >= 0 && target_dim_idx < static_cast(dim_size)) { - // 计算 Output 的线性偏移 - // Output Shape: [Outer, DimSize, Inner] - size_t out_offset = outer_idx * (dim_size * inner_size) + static_cast(target_dim_idx) * inner_size + inner_idx; - - // 使用自定义原子操作 - atomic_add_custom(&output[out_offset], source[i] * alpha_val); - } - } -} - -// ================================================================== -// 4. 向量化 Kernel (优化读取带宽) -// ================================================================== -template -__global__ void index_add_kernel_vectorized( - T *__restrict__ output, - const T *__restrict__ source, - const TIdx *__restrict__ indices, - size_t outer_size, - size_t inner_size, - size_t dim_size, - size_t index_len, - size_t num_packs, // Source 的 Pack 数量 - float alpha) { - // 将 source 强转为 Pack 指针，实现向量化读取 - using PackType = Pack; - const PackType *src_vec = reinterpret_cast(source); - - T alpha_val = static_cast(alpha); - - size_t tid = blockIdx.x * blockDim.x + threadIdx.x; - size_t stride = blockDim.x * gridDim.x; - - for (size_t i = tid; i < num_packs; i += stride) { - // 向量化读取 (LDG.128) - PackType reg_pack = src_vec[i]; - - // 当前 Pack 在 Source 中的起始线性索引 - size_t base_idx = i * PackSize; - -// 循环展开：处理 Pack 中的每一个元素 -#pragma unroll - for (int k = 0; k < PackSize; ++k) { - size_t curr_src_idx = base_idx + k; - - // 1. 坐标变换 - size_t inner_idx = curr_src_idx % inner_size; - size_t tmp = curr_src_idx / inner_size; - size_t idx_idx = tmp % index_len; - size_t outer_idx = tmp / index_len; - - // 2. 读取 Index - TIdx target_dim_idx = indices[idx_idx]; - - if (target_dim_idx < 0) { - target_dim_idx += static_cast(dim_size); - } - - // 3. 原子累加 - if (target_dim_idx >= 0 && target_dim_idx < static_cast(dim_size)) { - size_t out_offset = outer_idx * (dim_size * inner_size) + static_cast(target_dim_idx) * inner_size + inner_idx; - - // 使用自定义原子操作 - atomic_add_custom(&output[out_offset], reg_pack.val[k] * alpha_val); - } - } - } -} - -} // namespace op::index_add::cuda - -#endif // __INDEX_ADD_CUDA_H__ diff --git a/src/infiniop/ops/index_add/index_add.h b/src/infiniop/ops/index_add/index_add.h deleted file mode 100644 index bba78dcbf..000000000 --- a/src/infiniop/ops/index_add/index_add.h +++ /dev/null @@ -1,52 +0,0 @@ -#ifndef __INDEX_ADD_H__ -#define __INDEX_ADD_H__ - -#include "../../operator.h" -#include "info.h" // 引用 IndexAddInfo 定义 (需自行定义，包含 dim, alpha 等) -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::index_add::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - IndexAddInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - IndexAddInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t in_desc, \ - int64_t dim, \ - infiniopTensorDescriptor_t index_desc, \ - infiniopTensorDescriptor_t source_desc, \ - float alpha); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *output, \ - const void *input, \ - const void *index, \ - const void *source, \ - void *stream) const; \ - }; \ - } - -#endif // __INDEX_ADD_H__ diff --git a/src/infiniop/ops/index_add/info.h b/src/infiniop/ops/index_add/info.h deleted file mode 100644 index af7aa313e..000000000 --- a/src/infiniop/ops/index_add/info.h +++ /dev/null @@ -1,141 +0,0 @@ -#ifndef __INDEX_ADD_INFO_H__ -#define __INDEX_ADD_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::index_add { - -class IndexAddInfo { - IndexAddInfo() = default; - -public: - int _dtype; // 数据类型 (Input/Output/Source) - int _idx_dtype; // 索引类型 (int32, int64) - int64_t _dim; // 操作维度 - float _alpha; // 缩放因子 - - // 【新增】几何信息，用于计算内存偏移 - size_t _outer_size; // dim 左侧维度的乘积 - size_t _inner_size; // dim 右侧维度的乘积 - size_t _dim_size; // Input/Output 在 dim 维度的长度 - size_t _index_len; // Index 的长度 - - // 【修改】构造函数，初始化新增成员 - IndexAddInfo(int dtype, int idx_dtype, int64_t dim, float alpha, - size_t outer_size, size_t inner_size, size_t dim_size, size_t index_len) - : _dtype(dtype), _idx_dtype(idx_dtype), _dim(dim), _alpha(alpha), - _outer_size(outer_size), _inner_size(inner_size), _dim_size(dim_size), _index_len(index_len) {} - - int dtype() const { return _dtype; } - int idx_dtype() const { return _idx_dtype; } - int64_t dim() const { return _dim; } - float alpha() const { return _alpha; } - - // 【新增】Getter 方法 - size_t outer_size() const { return _outer_size; } - size_t inner_size() const { return _inner_size; } - size_t dim_size() const { return _dim_size; } - size_t index_len() const { return _index_len; } - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - int64_t dim, - infiniopTensorDescriptor_t index_desc, - infiniopTensorDescriptor_t source_desc, - float alpha) { - - // 1. 检查数据类型一致性 (Output vs Input vs Source) - int dtype = in_desc->dtype(); - if (out_desc->dtype() != dtype || source_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // 2. 检查索引数据类型 - int idx_dtype = index_desc->dtype(); - if (idx_dtype != INFINI_DTYPE_I32 && idx_dtype != INFINI_DTYPE_I64) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - // 3. 检查维度有效性 - int64_t ndim = static_cast(in_desc->ndim()); - if (dim < 0 || dim >= ndim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // 4. 检查 Index 形状 - if (index_desc->ndim() != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // 【新增】计算几何信息 - const auto &in_shape = in_desc->shape(); - - // outer_size: dim 之前所有维度的乘积 - size_t outer_size = 1; - for (int64_t i = 0; i < dim; ++i) { - outer_size *= in_shape[i]; - } - - // inner_size: dim 之后所有维度的乘积 (即 stride) - size_t inner_size = 1; - for (int64_t i = dim + 1; i < ndim; ++i) { - inner_size *= in_shape[i]; - } - - // dim_size - size_t dim_size = in_shape[dim]; - - // index_len - size_t index_len = index_desc->shape()[0]; - - // 5. 检查 Source 形状一致性 - // 规则: [Outer, IndexLen, Inner] - if (source_desc->ndim() != in_desc->ndim()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const auto &src_shape = source_desc->shape(); - - for (int64_t i = 0; i < ndim; ++i) { - if (i == dim) { - if (src_shape[i] != index_len) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } else { - if (src_shape[i] != in_shape[i]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - } - - // 6. 检查 Output 与 Input 形状一致性 - if (out_desc->ndim() != in_desc->ndim()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - const auto &out_shape = out_desc->shape(); - for (int64_t i = 0; i < ndim; ++i) { - if (out_shape[i] != in_shape[i]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - // 7. 返回 Info 对象 (包含计算好的几何信息) - return utils::Result(IndexAddInfo{ - dtype, - idx_dtype, - dim, - alpha, - outer_size, // pass - inner_size, // pass - dim_size, // pass - index_len // pass - }); - } -}; - -} // namespace op::index_add - -#endif // __INDEX_ADD_INFO_H__ diff --git a/src/infiniop/ops/index_add/metax/index_add_metax.h b/src/infiniop/ops/index_add/metax/index_add_metax.h deleted file mode 100644 index 7cc32e68b..000000000 --- a/src/infiniop/ops/index_add/metax/index_add_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __INDEX_ADD_METAX_H__ -#define __INDEX_ADD_METAX_H__ - -#include "../index_add.h" - -DESCRIPTOR(metax) - -#endif // __INDEX_ADD_METAX_H__ diff --git a/src/infiniop/ops/index_add/metax/index_add_metax.maca b/src/infiniop/ops/index_add/metax/index_add_metax.maca deleted file mode 100644 index 609599332..000000000 --- a/src/infiniop/ops/index_add/metax/index_add_metax.maca +++ /dev/null @@ -1,330 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../../../tensor.h" - -#include "../cuda/kernel.cuh" -#include "index_add_metax.h" -#include -#include - -#include - -namespace op::index_add::metax { - -// ================================================================== -// Atomic Helpers -// ================================================================== - -template -__device__ __forceinline__ void gpuAtomicAdd(T *address, T val) { - atomicAdd(address, val); -} - -template <> -__device__ __forceinline__ void gpuAtomicAdd( - int64_t *address, - int64_t val) { - atomicAdd( - reinterpret_cast(address), - static_cast(val)); -} - -template <> -__device__ __forceinline__ void gpuAtomicAdd( - cuda_bfloat16 *address, - cuda_bfloat16 val) { - unsigned int *addr = (unsigned int *)((char *)address - ((size_t)address & 2)); - - unsigned int old = *addr; - unsigned int assumed; - - do { - assumed = old; - - unsigned short old_val = ((size_t)address & 2) - ? (assumed >> 16) - : (assumed & 0xFFFF); - - cuda_bfloat16 sum = (cuda_bfloat16)((float)*reinterpret_cast(&old_val) + (float)val); - - unsigned short res = *reinterpret_cast(&sum); - - old = atomicCAS( - addr, - assumed, - ((size_t)address & 2) - ? ((assumed & 0xFFFF) | (res << 16)) - : ((assumed & 0xFFFF0000) | res)); - } while (assumed != old); -} - -// ================================================================== -// Kernel -// ================================================================== - -template -__global__ void index_add_kernel( - T *output, - const T *source, - const TIdx *indices, - int outer_size, - int inner_size, - int index_size, - int dim_size, - T alpha) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total_source = outer_size * index_size * inner_size; - - if (idx >= total_source) { - return; - } - - int inner = idx % inner_size; - int temp = idx / inner_size; - int i = temp % index_size; - int outer = temp / index_size; - - TIdx idx_pos = indices[i]; - if (idx_pos < 0) { - idx_pos += dim_size; - } - - if (idx_pos >= 0 && idx_pos < dim_size) { - int out_offset = outer * (dim_size * inner_size) + idx_pos * inner_size + inner; - - gpuAtomicAdd( - output + out_offset, - static_cast(source[idx] * alpha)); - } -} - -// ================================================================== -// Kernel Launcher -// ================================================================== - -template -void launch_kernel_impl( - void *output, - const void *source, - const void *indices, - int outer, - int inner, - int idx_size, - int dim_s, - float alpha, - void *stream) { - auto hc_stream = reinterpret_cast(stream); - - size_t total = (size_t)outer * idx_size * inner; - size_t block = 256; - size_t grid = (total + block - 1) / block; - - index_add_kernel - <<>>( - reinterpret_cast(output), - reinterpret_cast(source), - reinterpret_cast(indices), - outer, - inner, - idx_size, - dim_s, - static_cast(alpha)); -} - -// ================================================================== -// Utilities -// ================================================================== - -static size_t get_element_size(int dtype) { - if (dtype == INFINI_DTYPE_F64 || dtype == INFINI_DTYPE_I64) { - return 8; - } - if (dtype == INFINI_DTYPE_F32 || dtype == INFINI_DTYPE_I32) { - return 4; - } - return 2; -} - -struct Descriptor::Opaque { - std::shared_ptr internal; - float alpha; - int64_t dim; - int outer_size; - int inner_size; - int index_size; - int dim_size; - size_t total_bytes; -}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - int64_t dim, - infiniopTensorDescriptor_t index_desc, - infiniopTensorDescriptor_t source_desc, - float alpha) { - auto handle = reinterpret_cast(handle_); - - auto info_result = IndexAddInfo::create( - out_desc, in_desc, dim, index_desc, source_desc, alpha); - - if (!info_result) { - return info_result.status(); - } - - auto out_d = reinterpret_cast(out_desc); - auto idx_d = reinterpret_cast(index_desc); - - int ndim = out_d->ndim(); - int64_t real_dim = dim < 0 ? dim + ndim : dim; - - int outer = 1; - for (int i = 0; i < real_dim; ++i) { - outer *= out_d->shape()[i]; - } - - int inner = 1; - for (int i = real_dim + 1; i < ndim; ++i) { - inner *= out_d->shape()[i]; - } - - int dim_s = out_d->shape()[real_dim]; - - int idx_s = 1; - for (int i = 0; i < idx_d->ndim(); ++i) { - idx_s *= idx_d->shape()[i]; - } - - size_t bytes = (size_t)outer * dim_s * inner * get_element_size(out_d->dtype()); - - auto opaque = new Opaque{ - handle->internal(), - alpha, - dim, - outer, - inner, - idx_s, - dim_s, - bytes}; - - *desc_ptr = new Descriptor( - opaque, - info_result.take(), - 0, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *index, - const void *source, - void *stream) const { - auto hc_stream = reinterpret_cast(stream); - - hcMemcpyAsync( - output, - input, - _opaque->total_bytes, - hcMemcpyDeviceToDevice, - hc_stream); - - auto dtype = _info.dtype(); - auto idx_dtype = _info.idx_dtype(); - - int outer = _opaque->outer_size; - int inner = _opaque->inner_size; - int dim_s = _opaque->dim_size; - int idx_sz = _opaque->index_size; - float alpha = _opaque->alpha; - -#define LAUNCH(T, TIdx) \ - launch_kernel_impl( \ - output, source, index, \ - outer, inner, idx_sz, dim_s, alpha, stream) - - if (idx_dtype == INFINI_DTYPE_I32) { - switch (dtype) { - - case INFINI_DTYPE_F16: - LAUNCH(__half, int32_t); - break; - - case INFINI_DTYPE_BF16: - LAUNCH(cuda_bfloat16, int32_t); - break; - - case INFINI_DTYPE_F32: - LAUNCH(float, int32_t); - break; - - case INFINI_DTYPE_F64: - LAUNCH(double, int32_t); - break; - - case INFINI_DTYPE_I32: - LAUNCH(int32_t, int32_t); - break; - - case INFINI_DTYPE_I64: - LAUNCH(int64_t, int32_t); - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - } else if (idx_dtype == INFINI_DTYPE_I64) { - switch (dtype) { - - case INFINI_DTYPE_F16: - LAUNCH(__half, int64_t); - break; - - case INFINI_DTYPE_BF16: - LAUNCH(cuda_bfloat16, int64_t); - break; - - case INFINI_DTYPE_F32: - LAUNCH(float, int64_t); - break; - - case INFINI_DTYPE_F64: - LAUNCH(double, int64_t); - break; - - case INFINI_DTYPE_I32: - LAUNCH(int32_t, int64_t); - break; - - case INFINI_DTYPE_I64: - LAUNCH(int64_t, int64_t); - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::index_add::metax diff --git a/src/infiniop/ops/index_add/moore/index_add_moore.h b/src/infiniop/ops/index_add/moore/index_add_moore.h deleted file mode 100644 index 50978fe7c..000000000 --- a/src/infiniop/ops/index_add/moore/index_add_moore.h +++ /dev/null @@ -1,11 +0,0 @@ -#ifndef __INDEX_ADD_MOORE_API_H__ -#define __INDEX_ADD_MOORE_API_H__ - -// 引入上层定义的 Descriptor 宏和基础类 -#include "../index_add.h" - -// 使用 index_add.h 中定义的 DESCRIPTOR 宏 -// 这将自动生成 op::index_add::moore::Descriptor 类定义 -DESCRIPTOR(moore) - -#endif // __INDEX_ADD_MOORE_API_H__ diff --git a/src/infiniop/ops/index_add/moore/index_add_moore.mu b/src/infiniop/ops/index_add/moore/index_add_moore.mu deleted file mode 100644 index 119e4b4f5..000000000 --- a/src/infiniop/ops/index_add/moore/index_add_moore.mu +++ /dev/null @@ -1,201 +0,0 @@ -#include "index_add_moore.h" -#include "index_add_moore_kernel.h" // 包含 IndexAddOp Functor 定义 - -#include -#include -#include -#include - -#include "../../../devices/moore/moore_handle.h" - -namespace op::index_add::moore { - -// ================================================================== -// 1. Kernel Wrapper -// ================================================================== - -// 这是一个 Global Kernel 包装器，它调用 index_add_moore_kernel.h 中的 IndexAddOp Functor -template -__global__ void index_add_kernel( - const size_t num_elements, // Source 的总元素数量 (线程任务总量) - const size_t index_len, // Index 向量长度 - const size_t inner_size, // stride - const size_t dim_size, // Output 在 dim 维度的长度 - const float alpha, // 缩放因子 - const T *source, - const TIdx *indices, - T *output) { - - // idx 对应 Source 张量的线性索引 - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_elements) { - // 使用 Functor - IndexAddOp op; - op(idx, index_len, inner_size, dim_size, alpha, source, indices, output); - } -} - -// ================================================================== -// 2. Launcher Implementation -// ================================================================== - -template -void index_add_moore_launch( - const IndexAddInfo &info, - T *output, - const T *input, - const T *source, - const void *indices, // void* 传入，内部强转 - void *stream) { - - auto musa_stream = (musaStream_t)stream; - const TIdx *indices_ptr = static_cast(indices); - - // -------------------------------------------------------------- - // 步骤 1: Copy Input -> Output - // -------------------------------------------------------------- - // Output 初始化为 Input 的值。 - // 计算总元素数量: Output 形状与 Input 一致 - size_t total_out_elements = info.outer_size() * info.dim_size() * info.inner_size(); - - // 如果 input 和 output 指针不同，则执行拷贝 - if (output != input) { - musaMemcpyAsync(output, input, total_out_elements * sizeof(T), musaMemcpyDeviceToDevice, musa_stream); - } - - // -------------------------------------------------------------- - // 步骤 2: Scatter Add (Source -> Output) - // -------------------------------------------------------------- - // 线程并行度取决于 Source 的大小 - // Source 逻辑形状: [Outer, IndexLen, Inner] - size_t num_src_elements = info.outer_size() * info.index_len() * info.inner_size(); - - if (num_src_elements == 0) { - return; // 以此避免空 Kernel Launch - } - - int threads = 256; - int blocks = (num_src_elements + threads - 1) / threads; - - index_add_kernel<<>>( - num_src_elements, - info.index_len(), - info.inner_size(), - info.dim_size(), - info.alpha(), - source, - indices_ptr, - output); -} - -// ================================================================== -// 3. Descriptor Implementation -// ================================================================== - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - int64_t dim, - infiniopTensorDescriptor_t index_desc, - infiniopTensorDescriptor_t source_desc, - float alpha) { - - auto handle = reinterpret_cast(handle_); - - // 使用 Info 类校验形状和类型，并预计算几何参数 - auto info_result = IndexAddInfo::create(out_desc, in_desc, dim, index_desc, source_desc, alpha); - - if (!info_result) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - *desc_ptr = new Descriptor( - nullptr, - *info_result, - 0, // No workspace needed - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *index, - const void *source, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - -// -------------------------------------------------------------- -// 定义分发宏：解决 Data Type x Index Type 的组合爆炸 -// -------------------------------------------------------------- -#define LAUNCH_KERNEL(T) \ - do { \ - if (_info.idx_dtype() == INFINI_DTYPE_I32) { \ - index_add_moore_launch( \ - _info, \ - static_cast(output), \ - static_cast(input), \ - static_cast(source), \ - index, \ - stream); \ - } else if (_info.idx_dtype() == INFINI_DTYPE_I64) { \ - index_add_moore_launch( \ - _info, \ - static_cast(output), \ - static_cast(input), \ - static_cast(source), \ - index, \ - stream); \ - } else { \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } \ - } while (0) - - // -------------------------------------------------------------- - // 根据数据类型分发 - // -------------------------------------------------------------- - switch (_info.dtype()) { - case INFINI_DTYPE_F16: - LAUNCH_KERNEL(half); - break; - - case INFINI_DTYPE_BF16: - LAUNCH_KERNEL(__mt_bfloat16); - break; - - case INFINI_DTYPE_F32: - LAUNCH_KERNEL(float); - break; - - case INFINI_DTYPE_F64: - LAUNCH_KERNEL(double); - break; - - // 如果需要支持整数类型的 AtomicAdd，需要确保 Kernel 中有对应特化 - // 这里仅示例浮点类型 - case INFINI_DTYPE_I32: - LAUNCH_KERNEL(int32_t); - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_KERNEL - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::index_add::moore diff --git a/src/infiniop/ops/index_add/moore/index_add_moore_kernel.h b/src/infiniop/ops/index_add/moore/index_add_moore_kernel.h deleted file mode 100644 index 5bcd45f2d..000000000 --- a/src/infiniop/ops/index_add/moore/index_add_moore_kernel.h +++ /dev/null @@ -1,121 +0,0 @@ -#ifndef __INDEX_ADD_MOORE_KERNEL_H__ -#define __INDEX_ADD_MOORE_KERNEL_H__ - -#include -#include -#include -#include - -namespace op::index_add::moore { - -template -__device__ __forceinline__ void atomic_add_func(T *address, T val) { - atomicAdd(address, val); -} - -template <> -__device__ __forceinline__ void atomic_add_func(half *address, half val) { - // 将地址重解释为 unsigned short* 以便进行位操作 - unsigned short *address_as_us = reinterpret_cast(address); - unsigned short old = *address_as_us; - unsigned short assumed; - - do { - assumed = old; - half sum = __float2half(__half2float(*reinterpret_cast(&assumed)) + __half2float(val)); - - unsigned short sum_as_us = *reinterpret_cast(&sum); - - old = atomicCAS(address_as_us, assumed, sum_as_us); - - } while (assumed != old); -} - -template <> -__device__ __forceinline__ void atomic_add_func<__mt_bfloat16>(__mt_bfloat16 *address, __mt_bfloat16 val) { - unsigned short *address_as_us = reinterpret_cast(address); - unsigned short old = *address_as_us; - unsigned short assumed; - - do { - assumed = old; - // BF16 -> Float -> Add -> BF16 - float sum_f = __bfloat162float(*reinterpret_cast(&assumed)) + __bfloat162float(val); - - __mt_bfloat16 sum_bf = __float2bfloat16(sum_f); - unsigned short sum_as_us = *reinterpret_cast(&sum_bf); - - old = atomicCAS(address_as_us, assumed, sum_as_us); - - } while (assumed != old); -} - -// ================================================================== -// 2. Kernel Functor -// ================================================================== - -typedef struct IndexAddOp { -public: - template - __device__ __forceinline__ void operator()( - const size_t curr_idx, // Flattened index for Source - const size_t index_len, // Length of Index tensor - const size_t inner_size, // Stride of inner dims - const size_t dim_size, // Size of target dim in Output - const float alpha, // Scale factor - const T *source, // Source Tensor - const TIdx *indices, // Index Tensor - T *output // Output Tensor - ) const { - - size_t inner_idx = curr_idx % inner_size; - size_t tmp = curr_idx / inner_size; - size_t idx_in_indices = tmp % index_len; // 当前处理的是 Index 张量中的第几个索引 - size_t outer_idx = tmp / index_len; - - // --- 2. 读取 Source 并应用 Alpha --- - T src_val = source[curr_idx]; - float val_f; - - // 统一转 float 计算乘法 - if constexpr (std::is_same_v) { - val_f = __half2float(src_val); - } else if constexpr (std::is_same_v) { - val_f = __bfloat162float(src_val); - } else { - val_f = static_cast(src_val); - } - - val_f *= alpha; - - // 转回 T - T add_val; - if constexpr (std::is_same_v) { - add_val = __float2half(val_f); - } else if constexpr (std::is_same_v) { - add_val = __float2bfloat16(val_f); - } else { - add_val = static_cast(val_f); - } - - // --- 3. 读取 Index 并计算 Output 偏移 --- - TIdx target_dim_idx = indices[idx_in_indices]; - - // 处理 Python 风格负索引 - if (target_dim_idx < 0) { - target_dim_idx += static_cast(dim_size); - } - - // --- 4. 边界检查 & 原子累加 --- - if (target_dim_idx >= 0 && target_dim_idx < static_cast(dim_size)) { - // output_offset = outer * (dim_size * inner) + target_idx * inner + inner - size_t out_offset = outer_idx * (dim_size * inner_size) + static_cast(target_dim_idx) * inner_size + inner_idx; - atomic_add_func(output + out_offset, add_val); - } - } - -} IndexAddOp; - -} // namespace op::index_add::moore - -#endif // __INDEX_ADD_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/index_add/nvidia/index_add_nvidia.cu b/src/infiniop/ops/index_add/nvidia/index_add_nvidia.cu deleted file mode 100644 index d946b44d4..000000000 --- a/src/infiniop/ops/index_add/nvidia/index_add_nvidia.cu +++ /dev/null @@ -1,170 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" - -// Iluvatar does not support atomic add yet -#ifndef ENABLE_ILUVATAR_API -#include "../cuda/kernel.cuh" -#endif - -#include "index_add_nvidia.cuh" -#include - -namespace op::index_add::nvidia { - -template -bool is_aligned(const void *ptr, size_t alignment) { - return reinterpret_cast(ptr) % alignment == 0; -} - -// ================================================================== -// Kernel Launch Logic -// ================================================================== -template -void launch_kernel( - void *output, - const void *source, - const void *indices, - const IndexAddInfo &info, - void *stream) { -// Iluvatar does not support atomic add yet -#ifndef ENABLE_ILUVATAR_API - auto out_ptr = reinterpret_cast(output); - auto src_ptr = reinterpret_cast(source); - auto idx_ptr = reinterpret_cast(indices); - auto cuda_stream = reinterpret_cast(stream); - - // 获取几何信息 - size_t outer_size = info.outer_size(); - size_t inner_size = info.inner_size(); - size_t dim_size = info.dim_size(); - size_t index_len = info.index_len(); - float alpha = info.alpha(); - - // Source 总元素数 - size_t num_source = outer_size * index_len * inner_size; - - // --- 向量化参数配置 --- - // 目标：每个线程读取 128-bit (16 Bytes) Source 数据 - constexpr int TotalBytes = 16; - constexpr int PackSize = TotalBytes / sizeof(T); - bool can_vectorize = (PackSize > 1) && (num_source % PackSize == 0) && is_aligned(source, TotalBytes); - - if (can_vectorize) { - // === 路径 A: 向量化读取 Kernel === - size_t num_packs = num_source / PackSize; - - size_t block_size = 256; - size_t grid_size = (num_packs + block_size - 1) / block_size; - - op::index_add::cuda::index_add_kernel_vectorized - <<>>( - out_ptr, src_ptr, idx_ptr, - outer_size, inner_size, dim_size, index_len, - num_packs, alpha); - } else { - // === 路径 B: 标量 Kernel === - size_t block_size = 256; - size_t grid_size = (num_source + block_size - 1) / block_size; - - op::index_add::cuda::index_add_kernel - <<>>( - out_ptr, src_ptr, idx_ptr, - outer_size, inner_size, dim_size, index_len, - num_source, alpha); - } -#endif -} - -// ================================================================== -// Descriptor 实现 -// ================================================================== - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - int64_t dim, - infiniopTensorDescriptor_t index_desc, - infiniopTensorDescriptor_t source_desc, - float alpha) { - - // Info 创建 - auto info_result = IndexAddInfo::create(out_desc, in_desc, dim, index_desc, source_desc, alpha); - if (!info_result) { - return info_result.status(); - } - - *desc_ptr = new Descriptor( - new Opaque(), info_result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -// ================================================================== -// Calculate -// ================================================================== -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *index, - const void *source, - void *stream) const { - - auto dtype = _info.dtype(); - auto idx_dtype = _info.idx_dtype(); -#define LAUNCH_BY_SIZE(T_STORAGE) \ - switch (idx_dtype) { \ - case INFINI_DTYPE_I32: \ - launch_kernel(output, source, index, _info, stream); \ - break; \ - case INFINI_DTYPE_I64: \ - launch_kernel(output, source, index, _info, stream); \ - break; \ - default: \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } - - switch (dtype) { - // 32-bit Float - case INFINI_DTYPE_F32: - LAUNCH_BY_SIZE(float); - break; - // 64-bit Float - case INFINI_DTYPE_F64: - LAUNCH_BY_SIZE(double); - break; - // 16-bit Half (fp16) -> 使用 __half - case INFINI_DTYPE_F16: - LAUNCH_BY_SIZE(__half); - break; - // 16-bit BFloat16 (bf16) -> 使用 __nv_bfloat16 - case INFINI_DTYPE_BF16: - LAUNCH_BY_SIZE(__nv_bfloat16); - break; - // Integers - case INFINI_DTYPE_I32: - LAUNCH_BY_SIZE(int32_t); - break; - case INFINI_DTYPE_I64: - LAUNCH_BY_SIZE(int64_t); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_BY_SIZE - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::index_add::nvidia diff --git a/src/infiniop/ops/index_add/nvidia/index_add_nvidia.cuh b/src/infiniop/ops/index_add/nvidia/index_add_nvidia.cuh deleted file mode 100644 index f7cf75f22..000000000 --- a/src/infiniop/ops/index_add/nvidia/index_add_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __INDEX_ADD_NVIDIA_CUH__ -#define __INDEX_ADD_NVIDIA_CUH__ - -#include "../index_add.h" - -DESCRIPTOR(nvidia) - -#endif // __INDEX_ADD_NVIDIA_CUH__ diff --git a/src/infiniop/ops/index_add/operator.cc b/src/infiniop/ops/index_add/operator.cc deleted file mode 100644 index f655628fa..000000000 --- a/src/infiniop/ops/index_add/operator.cc +++ /dev/null @@ -1,190 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/index_add.h" - -#ifdef ENABLE_CPU_API -#include "cpu/index_add_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/index_add_nvidia.cuh" -#endif - -#ifdef ENABLE_METAX_API -#include "metax/index_add_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/index_add_moore.h" -#endif - -extern "C" { - -// ======================================================================= -// 1. 创建算子描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopCreateIndexAddDescriptor( - infiniopHandle_t handle, - infiniopIndexAddDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - int64_t dim, - infiniopTensorDescriptor_t index, - infiniopTensorDescriptor_t source, - float alpha) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::index_add::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output, \ - input, \ - dim, \ - index, \ - source, \ - alpha) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ======================================================================= -// 2. 获取 Workspace 大小 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopGetIndexAddWorkspaceSize(infiniopIndexAddDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -// 【关键修复】启用 Moore 分支 -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -// ======================================================================= -// 3. 执行计算 (Calculate) -// ======================================================================= -__INFINI_C infiniStatus_t infiniopIndexAdd( - infiniopIndexAddDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *index, - const void *source, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, input, index, source, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -// 【关键修复】启用 Moore 分支 -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -// ======================================================================= -// 4. 销毁描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopDestroyIndexAddDescriptor(infiniopIndexAddDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} - -} // extern "C" diff --git a/src/infiniop/ops/index_copy/cpu/index_copy_cpu.cc b/src/infiniop/ops/index_copy/cpu/index_copy_cpu.cc deleted file mode 100644 index 2e1a61f09..000000000 --- a/src/infiniop/ops/index_copy/cpu/index_copy_cpu.cc +++ /dev/null @@ -1,155 +0,0 @@ -#include "index_copy_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include -#include -#include - -namespace op::index_copy::cpu { - -Descriptor::~Descriptor() = default; - -// ================================================================== -// 创建描述符 -// ================================================================== -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - int64_t dim, - infiniopTensorDescriptor_t index_desc, - infiniopTensorDescriptor_t source_desc) { // 注意：移除了 float alpha - - auto handle = reinterpret_cast(handle_); - - // 创建 Info 对象 - auto result = IndexCopyInfo::create(out_desc, in_desc, dim, index_desc, source_desc); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - nullptr, // Opaque* - result.take(), // Info - 0, // Workspace Size - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -// ================================================================== -// 核心计算逻辑 (串行实现) -// ================================================================== -template -void calculate_cpu_impl( - const IndexCopyInfo &info, - void *output, - const void *input, - const void *index, - const void *source) { - - // IndexCopy 不需要 Alpha 也不需要提升精度进行计算，直接拷贝即可 - - // 1. 获取几何信息 - size_t outer_size = info.outer_size(); - size_t inner_size = info.inner_size(); - size_t dim_size = info.dim_size(); - size_t index_len = info.index_len(); - - auto out_ptr = reinterpret_cast(output); - auto src_ptr = reinterpret_cast(source); - auto idx_ptr = reinterpret_cast(index); - - // ----------------------------------------------------------- - // 串行循环逻辑 - // ----------------------------------------------------------- - for (size_t o = 0; o < outer_size; ++o) { - for (size_t i = 0; i < index_len; ++i) { - - TIdx idx = idx_ptr[i]; - - // 处理负索引 - if (idx < 0) { - idx += static_cast(dim_size); - } - - // 边界检查 - if (idx < 0 || static_cast(idx) >= dim_size) { - continue; - } - - // 计算偏移 - size_t src_offset = o * index_len * inner_size + i * inner_size; - size_t out_offset = o * dim_size * inner_size + static_cast(idx) * inner_size; - - // Inner 维度循环 - for (size_t in = 0; in < inner_size; ++in) { - // 【核心逻辑】 - // IndexCopy: output[idx] = source[i] - // 直接赋值，无需 utils::cast 提升精度 - out_ptr[out_offset + in] = src_ptr[src_offset + in]; - } - } - } -} - -// ================================================================== -// 执行函数 (分发逻辑) -// ================================================================== -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *index, - const void *source, - void *stream) const { - - auto dtype = _info.dtype(); - auto idx_dtype = _info.idx_dtype(); - -#define DISPATCH(TDATA, TIDX) \ - calculate_cpu_impl(_info, output, input, index, source); \ - return INFINI_STATUS_SUCCESS - - if (idx_dtype == INFINI_DTYPE_I32) { - switch (dtype) { - case INFINI_DTYPE_F32: - DISPATCH(float, int32_t); - case INFINI_DTYPE_F64: - DISPATCH(double, int32_t); - case INFINI_DTYPE_F16: - DISPATCH(fp16_t, int32_t); - case INFINI_DTYPE_BF16: - DISPATCH(bf16_t, int32_t); - case INFINI_DTYPE_I32: - DISPATCH(int32_t, int32_t); - case INFINI_DTYPE_I64: - DISPATCH(int64_t, int32_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (idx_dtype == INFINI_DTYPE_I64) { - switch (dtype) { - case INFINI_DTYPE_F32: - DISPATCH(float, int64_t); - case INFINI_DTYPE_F64: - DISPATCH(double, int64_t); - case INFINI_DTYPE_F16: - DISPATCH(fp16_t, int64_t); - case INFINI_DTYPE_BF16: - DISPATCH(bf16_t, int64_t); - case INFINI_DTYPE_I32: - DISPATCH(int32_t, int64_t); - case INFINI_DTYPE_I64: - DISPATCH(int64_t, int64_t); - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } - - return INFINI_STATUS_BAD_TENSOR_DTYPE; - -#undef DISPATCH -} - -} // namespace op::index_copy::cpu diff --git a/src/infiniop/ops/index_copy/cpu/index_copy_cpu.h b/src/infiniop/ops/index_copy/cpu/index_copy_cpu.h deleted file mode 100644 index 009098014..000000000 --- a/src/infiniop/ops/index_copy/cpu/index_copy_cpu.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __INDEX_COPY_CPU_H__ -#define __INDEX_COPY_CPU_H__ - -#include "../index_copy.h" -DESCRIPTOR(cpu) - -#endif // __INDEX_COPY_CPU_H__ diff --git a/src/infiniop/ops/index_copy/cuda/kernel.cuh b/src/infiniop/ops/index_copy/cuda/kernel.cuh deleted file mode 100644 index 3b7866564..000000000 --- a/src/infiniop/ops/index_copy/cuda/kernel.cuh +++ /dev/null @@ -1,126 +0,0 @@ -#ifndef __INDEX_COPY_CUDA_H__ -#define __INDEX_COPY_CUDA_H__ - -#include - -namespace op::index_copy::cuda { - -// ================================================================== -// 1. 定义向量化数据包 (Aligned Pack) -// ================================================================== -// 与 IndexAdd 保持一致，用于向量化读取 -template -struct alignas(sizeof(T) * N) Pack { - T val[N]; -}; - -// ================================================================== -// 2. 标量版 Kernel (通用 fallback) -// ================================================================== -template -__global__ void index_copy_kernel( - T *__restrict__ output, - const T *__restrict__ source, - const TIdx *__restrict__ indices, - size_t outer_size, // dim 左边的维度积 - size_t inner_size, // dim 右边的维度积 - size_t dim_size, // output 在 dim 维度的长度 - size_t index_len, // index 的长度 (source 在 dim 维度的长度) - size_t num_source // source 的总元素数 - // 注意：移除了 float alpha -) { - size_t tid = blockIdx.x * blockDim.x + threadIdx.x; - size_t stride = blockDim.x * gridDim.x; - - // Grid-Stride Loop 遍历 Source 张量 - for (size_t i = tid; i < num_source; i += stride) { - // 1. 将线性索引 i 转换为逻辑坐标 (outer, idx_idx, inner) - // Source Shape: [Outer, IndexLen, Inner] - size_t inner_idx = i % inner_size; - size_t tmp = i / inner_size; - size_t idx_idx = tmp % index_len; - size_t outer_idx = tmp / index_len; - - // 2. 读取索引值 - TIdx target_dim_idx = indices[idx_idx]; - - // 3. 处理负索引 (防御性) - if (target_dim_idx < 0) { - target_dim_idx += static_cast(dim_size); - } - - // 4. 边界检查与赋值 - if (target_dim_idx >= 0 && target_dim_idx < static_cast(dim_size)) { - // 计算 Output 的线性偏移 - // Output Shape: [Outer, DimSize, Inner] - size_t out_offset = outer_idx * (dim_size * inner_size) + static_cast(target_dim_idx) * inner_size + inner_idx; - - // 【核心修改】 - // IndexCopy 不需要原子操作，直接赋值。 - // 如果有多个索引指向同一个位置，结果由执行顺序决定（Race Condition），这是符合预期的行为。 - output[out_offset] = source[i]; - } - } -} - -// ================================================================== -// 3. 向量化 Kernel (优化读取带宽) -// ================================================================== -template -__global__ void index_copy_kernel_vectorized( - T *__restrict__ output, - const T *__restrict__ source, - const TIdx *__restrict__ indices, - size_t outer_size, - size_t inner_size, - size_t dim_size, - size_t index_len, - size_t num_packs // Source 的 Pack 数量 - // 注意：移除了 float alpha -) { - // 将 source 强转为 Pack 指针，实现向量化读取 - using PackType = Pack; - const PackType *src_vec = reinterpret_cast(source); - - size_t tid = blockIdx.x * blockDim.x + threadIdx.x; - size_t stride = blockDim.x * gridDim.x; - - for (size_t i = tid; i < num_packs; i += stride) { - // 向量化读取 (LDG.128) - PackType reg_pack = src_vec[i]; - - // 当前 Pack 在 Source 中的起始线性索引 - size_t base_idx = i * PackSize; - -// 循环展开：处理 Pack 中的每一个元素 -#pragma unroll - for (int k = 0; k < PackSize; ++k) { - size_t curr_src_idx = base_idx + k; - - // 1. 坐标变换 - size_t inner_idx = curr_src_idx % inner_size; - size_t tmp = curr_src_idx / inner_size; - size_t idx_idx = tmp % index_len; - size_t outer_idx = tmp / index_len; - - // 2. 读取 Index - TIdx target_dim_idx = indices[idx_idx]; - - if (target_dim_idx < 0) { - target_dim_idx += static_cast(dim_size); - } - - // 3. 赋值 - if (target_dim_idx >= 0 && target_dim_idx < static_cast(dim_size)) { - size_t out_offset = outer_idx * (dim_size * inner_size) + static_cast(target_dim_idx) * inner_size + inner_idx; - - // 【核心修改】直接赋值 - output[out_offset] = reg_pack.val[k]; - } - } - } -} - -} // namespace op::index_copy::cuda - -#endif // __INDEX_COPY_CUDA_H__ diff --git a/src/infiniop/ops/index_copy/index_copy.h b/src/infiniop/ops/index_copy/index_copy.h deleted file mode 100644 index 11d7548a0..000000000 --- a/src/infiniop/ops/index_copy/index_copy.h +++ /dev/null @@ -1,51 +0,0 @@ -#ifndef __INDEX_COPY_H__ -#define __INDEX_COPY_H__ - -#include "../../operator.h" -#include "info.h" // 引用 IndexCopyInfo 定义 -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::index_copy::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - IndexCopyInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - IndexCopyInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t in_desc, \ - int64_t dim, \ - infiniopTensorDescriptor_t index_desc, \ - infiniopTensorDescriptor_t source_desc); /* 注意：移除了 alpha */ \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *output, \ - const void *input, \ - const void *index, \ - const void *source, \ - void *stream) const; \ - }; \ - } - -#endif // __INDEX_COPY_H__ diff --git a/src/infiniop/ops/index_copy/info.h b/src/infiniop/ops/index_copy/info.h deleted file mode 100644 index deea831b7..000000000 --- a/src/infiniop/ops/index_copy/info.h +++ /dev/null @@ -1,119 +0,0 @@ -#ifndef __INDEX_COPY_INFO_H__ -#define __INDEX_COPY_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::index_copy { - -class IndexCopyInfo { - IndexCopyInfo() = default; - -public: - int _dtype; - int _idx_dtype; - int64_t _dim; - - size_t _outer_size; - size_t _inner_size; - size_t _dim_size; - size_t _index_len; - - IndexCopyInfo(int dtype, int idx_dtype, int64_t dim, - size_t outer_size, size_t inner_size, size_t dim_size, size_t index_len) - : _dtype(dtype), _idx_dtype(idx_dtype), _dim(dim), - _outer_size(outer_size), _inner_size(inner_size), _dim_size(dim_size), _index_len(index_len) {} - - int dtype() const { return _dtype; } - int idx_dtype() const { return _idx_dtype; } - int64_t dim() const { return _dim; } - - size_t outer_size() const { return _outer_size; } - size_t inner_size() const { return _inner_size; } - size_t dim_size() const { return _dim_size; } - size_t index_len() const { return _index_len; } - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - int64_t dim, - infiniopTensorDescriptor_t index_desc, - infiniopTensorDescriptor_t source_desc) { - - int dtype = in_desc->dtype(); - if (out_desc->dtype() != dtype || source_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - int idx_dtype = index_desc->dtype(); - if (idx_dtype != INFINI_DTYPE_I32 && idx_dtype != INFINI_DTYPE_I64) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - int64_t ndim = static_cast(in_desc->ndim()); - if (dim < 0 || dim >= ndim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (index_desc->ndim() != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const auto &in_shape = in_desc->shape(); - - size_t outer_size = 1; - for (int64_t i = 0; i < dim; ++i) { - outer_size *= in_shape[i]; - } - - size_t inner_size = 1; - for (int64_t i = dim + 1; i < ndim; ++i) { - inner_size *= in_shape[i]; - } - - size_t dim_size = in_shape[dim]; - size_t index_len = index_desc->shape()[0]; - - if (source_desc->ndim() != in_desc->ndim()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - const auto &src_shape = source_desc->shape(); - - for (int64_t i = 0; i < ndim; ++i) { - if (i == dim) { - if (src_shape[i] != index_len) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } else { - if (src_shape[i] != in_shape[i]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - } - - if (out_desc->ndim() != in_desc->ndim()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - const auto &out_shape = out_desc->shape(); - for (int64_t i = 0; i < ndim; ++i) { - if (out_shape[i] != in_shape[i]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - return utils::Result(IndexCopyInfo{ - dtype, - idx_dtype, - dim, - outer_size, - inner_size, - dim_size, - index_len}); - } -}; - -} // namespace op::index_copy - -#endif // __INDEX_COPY_INFO_H__ diff --git a/src/infiniop/ops/index_copy/metax/index_copy_metax.h b/src/infiniop/ops/index_copy/metax/index_copy_metax.h deleted file mode 100644 index 9d2ebdeb6..000000000 --- a/src/infiniop/ops/index_copy/metax/index_copy_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __INDEX_COPY_METAX_H__ -#define __INDEX_COPY_METAX_H__ - -#include "../index_copy.h" - -DESCRIPTOR(metax) - -#endif // __INDEX_COPY_METAX_H__ diff --git a/src/infiniop/ops/index_copy/metax/index_copy_metax.maca b/src/infiniop/ops/index_copy/metax/index_copy_metax.maca deleted file mode 100644 index 9618cb8de..000000000 --- a/src/infiniop/ops/index_copy/metax/index_copy_metax.maca +++ /dev/null @@ -1,195 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "index_copy_metax.h" -#include -#include -#include - -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" - -namespace op::index_copy::metax { - -template -__global__ void index_copy_kernel( - T *output, const T *source, const TIdx *indices, - int outer_size, int inner_size, int index_size, int dim_size) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - int total_source = outer_size * index_size * inner_size; - if (idx >= total_source) { - return; - } - - int inner = idx % inner_size; - int temp = idx / inner_size; - int i = temp % index_size; - int outer = temp / index_size; - - TIdx idx_pos = indices[i]; - if (idx_pos < 0) { - idx_pos += dim_size; - } - - if (idx_pos >= 0 && idx_pos < dim_size) { - int out_offset = outer * (dim_size * inner_size) + idx_pos * inner_size + inner; - output[out_offset] = source[idx]; - } -} - -template -void launch_kernel_impl( - void *output, const void *source, const void *indices, - int outer_size, int inner_size, int index_size, int dim_size, void *stream) { - auto hc_stream = reinterpret_cast(stream); - size_t total_elements = (size_t)outer_size * index_size * inner_size; - size_t block_size = 256; - size_t grid_size = (total_elements + block_size - 1) / block_size; - index_copy_kernel<<>>( - reinterpret_cast(output), reinterpret_cast(source), reinterpret_cast(indices), - outer_size, inner_size, index_size, dim_size); -} - -static size_t get_element_size(int dtype) { - if (dtype == INFINI_DTYPE_F64 || dtype == INFINI_DTYPE_I64) { - return 8; - } - if (dtype == INFINI_DTYPE_F32 || dtype == INFINI_DTYPE_I32) { - return 4; - } - return 2; -} - -struct Descriptor::Opaque { - std::shared_ptr internal; - int64_t dim; - int outer_size, inner_size, index_size, dim_size; - size_t total_bytes; -}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -// [修复 1] create 签名匹配头文件：dim 在 index_desc 之前 -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - int64_t dim, - infiniopTensorDescriptor_t index_desc, - infiniopTensorDescriptor_t source_desc) { - auto handle = reinterpret_cast(handle_); - // Info 创建顺序通常是 out, in, dim, index, source - auto info_result = IndexCopyInfo::create(out_desc, in_desc, dim, index_desc, source_desc); - if (!info_result) { - return info_result.status(); - } - - auto out_d = reinterpret_cast(out_desc); - auto idx_d = reinterpret_cast(index_desc); - - int ndim = out_d->ndim(); - int64_t real_dim = dim < 0 ? dim + ndim : dim; - - int outer = 1; - for (int i = 0; i < real_dim; ++i) { - outer *= out_d->shape()[i]; - } - int inner = 1; - for (int i = real_dim + 1; i < ndim; ++i) { - inner *= out_d->shape()[i]; - } - int dim_s = out_d->shape()[real_dim]; - int idx_s = 1; - for (int i = 0; i < idx_d->ndim(); ++i) { - idx_s *= idx_d->shape()[i]; - } - - size_t bytes = (size_t)outer * dim_s * inner * get_element_size(out_d->dtype()); - - auto opaque = new Opaque{handle->internal(), dim, outer, inner, idx_s, dim_s, bytes}; - *desc_ptr = new Descriptor(opaque, info_result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -// [修复 2] calculate 签名匹配头文件：显式参数 input, index, source -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *index, - const void *source, - void *stream) const { - auto hc_stream = reinterpret_cast(stream); - hcMemcpyAsync(output, input, _opaque->total_bytes, hcMemcpyDeviceToDevice, hc_stream); - - auto dtype = _info.dtype(); - auto idx_dtype = _info.idx_dtype(); - int outer = _opaque->outer_size; - int inner = _opaque->inner_size; - int dim_s = _opaque->dim_size; - int idx_sz = _opaque->index_size; - -#define LAUNCH(T, TIdx) launch_kernel_impl(output, source, index, outer, inner, idx_sz, dim_s, stream) - - if (idx_dtype == INFINI_DTYPE_I32) { - switch (dtype) { - case INFINI_DTYPE_F16: - LAUNCH(__half, int32_t); - break; - case INFINI_DTYPE_BF16: -#if defined(__MACA__) || defined(ENABLE_METAX_MC_API) - LAUNCH(cuda_bfloat16, int32_t); -#endif - break; - case INFINI_DTYPE_F32: - LAUNCH(float, int32_t); - break; - case INFINI_DTYPE_F64: - LAUNCH(double, int32_t); - break; - case INFINI_DTYPE_I32: - LAUNCH(int32_t, int32_t); - break; - case INFINI_DTYPE_I64: - LAUNCH(int64_t, int32_t); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (idx_dtype == INFINI_DTYPE_I64) { - switch (dtype) { - case INFINI_DTYPE_F16: - LAUNCH(__half, int64_t); - break; - case INFINI_DTYPE_BF16: -#if defined(__MACA__) || defined(ENABLE_METAX_MC_API) - LAUNCH(cuda_bfloat16, int64_t); -#endif - break; - case INFINI_DTYPE_F32: - LAUNCH(float, int64_t); - break; - case INFINI_DTYPE_F64: - LAUNCH(double, int64_t); - break; - case INFINI_DTYPE_I32: - LAUNCH(int32_t, int64_t); - break; - case INFINI_DTYPE_I64: - LAUNCH(int64_t, int64_t); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -#undef LAUNCH - return INFINI_STATUS_SUCCESS; -} -} // namespace op::index_copy::metax diff --git a/src/infiniop/ops/index_copy/moore/index_copy_moore.h b/src/infiniop/ops/index_copy/moore/index_copy_moore.h deleted file mode 100644 index b056ba2cb..000000000 --- a/src/infiniop/ops/index_copy/moore/index_copy_moore.h +++ /dev/null @@ -1,11 +0,0 @@ -#ifndef __INDEX_COPY_MOORE_API_H__ -#define __INDEX_COPY_MOORE_API_H__ - -// 引入上层定义的 Descriptor 宏和基础类 -#include "../index_copy.h" - -// 使用 index_copy.h 中定义的 DESCRIPTOR 宏 -// 这将自动生成 op::index_copy::moore::Descriptor 类定义 -DESCRIPTOR(moore) - -#endif // __INDEX_COPY_MOORE_API_H__ diff --git a/src/infiniop/ops/index_copy/moore/index_copy_moore.mu b/src/infiniop/ops/index_copy/moore/index_copy_moore.mu deleted file mode 100644 index 5f14a849f..000000000 --- a/src/infiniop/ops/index_copy/moore/index_copy_moore.mu +++ /dev/null @@ -1,205 +0,0 @@ -#include "index_copy_moore.h" -#include "index_copy_moore_kernel.h" // 包含 IndexCopyOp Functor 定义 - -#include -#include -#include - -#include "../../../devices/moore/moore_handle.h" - -namespace op::index_copy::moore { - -// ================================================================== -// 1. Kernel Wrapper Implementation -// ================================================================== - -template -__global__ void index_copy_kernel( - const size_t num_elements, // Source 元素总数 (线程任务总量) - const size_t index_len, // Index 长度 - const size_t inner_size, // Stride - const size_t dim_size, // Output 在 dim 维度的长度 - const T *source, - const TIdx *indices, - T *output) { - - // idx 对应 Source 张量的线性索引 - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - - if (idx < num_elements) { - // 使用 index_copy_moore_kernel.h 中定义的 Functor - IndexCopyOp op; - op(idx, index_len, inner_size, dim_size, source, indices, output); - } -} - -// ================================================================== -// 2. Launcher Implementation -// ================================================================== - -template -void index_copy_moore_launch( - const IndexCopyInfo &info, - T *output, - const T *input, - const T *source, - const void *indices, // void* 传入，内部强转 - void *stream) { - - auto musa_stream = (musaStream_t)stream; - const TIdx *indices_ptr = static_cast(indices); - - // -------------------------------------------------------------- - // 步骤 1: Copy Input -> Output - // -------------------------------------------------------------- - // Output 初始化为 Input 的值 - size_t total_out_elements = info.outer_size() * info.dim_size() * info.inner_size(); - - // 如果 input 和 output 指针不同，则执行拷贝 - if (output != input) { - musaMemcpyAsync(output, input, total_out_elements * sizeof(T), musaMemcpyDeviceToDevice, musa_stream); - } - - // -------------------------------------------------------------- - // 步骤 2: Scatter (Source -> Output) - // -------------------------------------------------------------- - // 线程并行度取决于 Source 的大小 - // Source 逻辑形状: [Outer, IndexLen, Inner] - size_t num_src_elements = info.outer_size() * info.index_len() * info.inner_size(); - - if (num_src_elements == 0) { - return; - } - - int threads = 256; - int blocks = (num_src_elements + threads - 1) / threads; - - index_copy_kernel<<>>( - num_src_elements, - info.index_len(), - info.inner_size(), - info.dim_size(), - source, - indices_ptr, - output); -} - -// ================================================================== -// 3. Descriptor Implementation -// ================================================================== - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - int64_t dim, - infiniopTensorDescriptor_t index_desc, - infiniopTensorDescriptor_t source_desc) { - - auto handle = reinterpret_cast(handle_); - - // 使用 Info 类校验形状和类型 - auto info_result = IndexCopyInfo::create(out_desc, in_desc, dim, index_desc, source_desc); - - if (!info_result) { - return info_result.status(); - } - - *desc_ptr = new Descriptor( - nullptr, - *info_result, - 0, // No workspace needed - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *index, - const void *source, - void *stream) const { - - if (workspace_size < _workspace_size) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - -// -------------------------------------------------------------- -// 定义分发宏：Data Type x Index Type -// -------------------------------------------------------------- -#define LAUNCH_KERNEL(T) \ - do { \ - if (_info.idx_dtype() == INFINI_DTYPE_I32) { \ - index_copy_moore_launch( \ - _info, \ - static_cast(output), \ - static_cast(input), \ - static_cast(source), \ - index, \ - stream); \ - } else if (_info.idx_dtype() == INFINI_DTYPE_I64) { \ - index_copy_moore_launch( \ - _info, \ - static_cast(output), \ - static_cast(input), \ - static_cast(source), \ - index, \ - stream); \ - } else { \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } \ - } while (0) - - // -------------------------------------------------------------- - // 根据数据类型分发 - // -------------------------------------------------------------- - switch (_info.dtype()) { - case INFINI_DTYPE_F16: - LAUNCH_KERNEL(half); - break; - - case INFINI_DTYPE_BF16: - LAUNCH_KERNEL(__mt_bfloat16); - break; - - case INFINI_DTYPE_F32: - LAUNCH_KERNEL(float); - break; - - case INFINI_DTYPE_F64: - LAUNCH_KERNEL(double); - break; - - case INFINI_DTYPE_I32: - LAUNCH_KERNEL(int32_t); - break; - - case INFINI_DTYPE_I64: - LAUNCH_KERNEL(int64_t); - break; - - case INFINI_DTYPE_I8: - LAUNCH_KERNEL(int8_t); - break; - - case INFINI_DTYPE_U8: - LAUNCH_KERNEL(uint8_t); - break; - - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_KERNEL - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::index_copy::moore diff --git a/src/infiniop/ops/index_copy/moore/index_copy_moore_kernel.h b/src/infiniop/ops/index_copy/moore/index_copy_moore_kernel.h deleted file mode 100644 index 522606832..000000000 --- a/src/infiniop/ops/index_copy/moore/index_copy_moore_kernel.h +++ /dev/null @@ -1,71 +0,0 @@ -#ifndef __INDEX_COPY_MOORE_KERNEL_H__ -#define __INDEX_COPY_MOORE_KERNEL_H__ - -#include // 必须包含，用于 __mt_bfloat16 定义 -#include -#include - -#include // 用于 std::is_same_v - -namespace op::index_copy::moore { - -typedef struct IndexCopyOp { -public: - // IndexCopy 算子涉及两种数据类型：数据本身 (T) 和 索引类型 (TIdx) - // 逻辑：Output[..., indices[i], ...] = Source[..., i, ...] - - template - __device__ __forceinline__ void operator()( - const size_t curr_idx, // 当前线程处理的 Source 元素的线性索引 - - // 几何参数 (来自 IndexCopyInfo) - const size_t index_len, // Index 向量的长度 (即 Source 在 dim 维度的长度) - const size_t inner_size, // dim 右侧维度的 stride - const size_t dim_size, // Output 在 dim 维度的长度 (用于计算偏移和边界检查) - - // 指针 - const T *source_data, // Source 数据指针 (Input source) - const TIdx *indices_data, // 索引数据指针 - T *output_data // Output 数据指针 (In/Out, 通常已由 Input 初始化) - ) const { - - // 1. 坐标映射 (Flat Index -> Multi-dim Index) - // Source 逻辑形状为: [Outer, IndexLen, Inner] - // 将 curr_idx 分解为 (outer, idx_in_index, inner) - - size_t inner_idx = curr_idx % inner_size; - size_t tmp = curr_idx / inner_size; - size_t idx_in_indices = tmp % index_len; // 当前处理的是 indices 张量中的第几个索引 - size_t outer_idx = tmp / index_len; - - // 2. 读取 Source 数据 - T src_val = source_data[curr_idx]; - - // 3. 获取目标 Index - TIdx target_dim_idx = indices_data[idx_in_indices]; - - // 4. 处理负索引 (支持 Python 风格，如 -1 代表最后一个元素) - if (target_dim_idx < 0) { - target_dim_idx += static_cast(dim_size); - } - - // 5. 边界检查与赋值 - if (target_dim_idx >= 0 && target_dim_idx < static_cast(dim_size)) { - // 计算 Output 的线性偏移 - // Output 逻辑形状: [Outer, DimSize, Inner] - // Offset = Outer_Pos + Dim_Pos + Inner_Pos - size_t out_offset = outer_idx * (dim_size * inner_size) + static_cast(target_dim_idx) * inner_size + inner_idx; - - // 直接赋值 (Scatter) - // MUSA 平台 half/bf16/float 均支持 operator= - // 对于重复索引的情况，通常行为是非确定性的 (Last write wins)，无需原子操作 - output_data[out_offset] = src_val; - } - // 如果索引越界，IndexCopy 通常忽略该操作，不做任何修改 - } - -} IndexCopyOp; - -} // namespace op::index_copy::moore - -#endif // __INDEX_COPY_MOORE_KERNEL_H__ diff --git a/src/infiniop/ops/index_copy/nvidia/index_copy_nvidia.cu b/src/infiniop/ops/index_copy/nvidia/index_copy_nvidia.cu deleted file mode 100644 index 5bf1208bb..000000000 --- a/src/infiniop/ops/index_copy/nvidia/index_copy_nvidia.cu +++ /dev/null @@ -1,180 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" - -#include "../cuda/kernel.cuh" -#include "index_copy_nvidia.cuh" -#include - -namespace op::index_copy::nvidia { - -// ================================================================== -// 辅助函数：检查指针内存对齐 -// ================================================================== -template -bool is_aligned(const void *ptr, size_t alignment) { - return reinterpret_cast(ptr) % alignment == 0; -} - -// ================================================================== -// Kernel Launch Logic -// ================================================================== -template -void launch_kernel( - void *output, - const void *source, - const void *indices, - const IndexCopyInfo &info, - void *stream) { - - auto out_ptr = reinterpret_cast(output); - auto src_ptr = reinterpret_cast(source); - auto idx_ptr = reinterpret_cast(indices); - auto cuda_stream = reinterpret_cast(stream); - - // 获取几何信息 (无 alpha) - size_t outer_size = info.outer_size(); - size_t inner_size = info.inner_size(); - size_t dim_size = info.dim_size(); - size_t index_len = info.index_len(); - - // Source 总元素数 - size_t num_source = outer_size * index_len * inner_size; - - // --- 向量化参数配置 --- - // 目标：每个线程读取 128-bit (16 Bytes) Source 数据 - // IndexCopy 是从 Source 读取并写入 Output，Source 是连续读取，适合向量化 Load - constexpr int TotalBytes = 16; - constexpr int PackSize = TotalBytes / sizeof(T); - - // 向量化条件检查： - // 1. PackSize > 1 - // 2. Source 总数能整除 PackSize (简化 tail 处理) - // 3. Source 指针地址对齐 (Load Vectorized 要求) - bool can_vectorize = (PackSize > 1) && (num_source % PackSize == 0) && is_aligned(source, TotalBytes); - - if (can_vectorize) { - // === 路径 A: 向量化读取 Kernel === - size_t num_packs = num_source / PackSize; - - size_t block_size = 256; - size_t grid_size = (num_packs + block_size - 1) / block_size; - - op::index_copy::cuda::index_copy_kernel_vectorized - <<>>( - out_ptr, src_ptr, idx_ptr, - outer_size, inner_size, dim_size, index_len, - num_packs - // 注意：移除了 alpha - ); - } else { - // === 路径 B: 标量 Kernel === - size_t block_size = 256; - size_t grid_size = (num_source + block_size - 1) / block_size; - - op::index_copy::cuda::index_copy_kernel - <<>>( - out_ptr, src_ptr, idx_ptr, - outer_size, inner_size, dim_size, index_len, - num_source - // 注意：移除了 alpha - ); - } -} - -// ================================================================== -// Descriptor 实现 -// ================================================================== - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t in_desc, - int64_t dim, - infiniopTensorDescriptor_t index_desc, - infiniopTensorDescriptor_t source_desc) { // 注意：移除了 alpha 参数 - - // Info 创建 - auto info_result = IndexCopyInfo::create(out_desc, in_desc, dim, index_desc, source_desc); // 无 alpha - if (!info_result) { - return info_result.status(); - } - - *desc_ptr = new Descriptor( - new Opaque(), info_result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -// ================================================================== -// Calculate -// ================================================================== -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *index, - const void *source, - void *stream) const { - - auto dtype = _info.dtype(); - auto idx_dtype = _info.idx_dtype(); - -// 宏：根据 T_STORAGE 类型实例化 launch_kernel -// T_STORAGE 将会是: float, double, int32_t, __half, cuda_bfloat16 -#define LAUNCH_BY_SIZE(T_STORAGE) \ - switch (idx_dtype) { \ - case INFINI_DTYPE_I32: \ - launch_kernel(output, source, index, _info, stream); \ - break; \ - case INFINI_DTYPE_I64: \ - launch_kernel(output, source, index, _info, stream); \ - break; \ - default: \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } - - switch (dtype) { - // 32-bit Float - case INFINI_DTYPE_F32: - LAUNCH_BY_SIZE(float); - break; - // 64-bit Float - case INFINI_DTYPE_F64: - LAUNCH_BY_SIZE(double); - break; - // 16-bit Half (fp16) -> 使用 __half - case INFINI_DTYPE_F16: - LAUNCH_BY_SIZE(__half); - break; - // 16-bit BFloat16 (bf16) -> 使用 cuda_bfloat16 - case INFINI_DTYPE_BF16: - LAUNCH_BY_SIZE(cuda_bfloat16); - break; - // Integers - case INFINI_DTYPE_I32: - LAUNCH_BY_SIZE(int32_t); - break; - case INFINI_DTYPE_I64: - LAUNCH_BY_SIZE(int64_t); - break; - - // 如果有其他整型需求 (I8, U8 等)，也在这里添加 case - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - -#undef LAUNCH_BY_SIZE - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::index_copy::nvidia diff --git a/src/infiniop/ops/index_copy/nvidia/index_copy_nvidia.cuh b/src/infiniop/ops/index_copy/nvidia/index_copy_nvidia.cuh deleted file mode 100644 index aa6afcfab..000000000 --- a/src/infiniop/ops/index_copy/nvidia/index_copy_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __INDEX_COPY_NVIDIA_CUH__ -#define __INDEX_COPY_NVIDIA_CUH__ - -#include "../index_copy.h" - -DESCRIPTOR(nvidia) - -#endif // __INDEX_COPY_NVIDIA_CUH__ diff --git a/src/infiniop/ops/index_copy/operator.cc b/src/infiniop/ops/index_copy/operator.cc deleted file mode 100644 index 97d078da4..000000000 --- a/src/infiniop/ops/index_copy/operator.cc +++ /dev/null @@ -1,196 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/index_copy.h" - -// --- 后端实现头文件 --- -#ifdef ENABLE_CPU_API -#include "cpu/index_copy_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/index_copy_nvidia.cuh" -#endif - -#ifdef ENABLE_METAX_API -#include "metax/index_copy_metax.h" -#endif - -#ifdef ENABLE_MOORE_API -#include "moore/index_copy_moore.h" -#endif - -extern "C" { - -// ======================================================================= -// 1. 创建算子描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopCreateIndexCopyDescriptor( - infiniopHandle_t handle, - infiniopIndexCopyDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t output, - infiniopTensorDescriptor_t input, - int64_t dim, - infiniopTensorDescriptor_t index, - infiniopTensorDescriptor_t source) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::index_copy::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - output, \ - input, \ - dim, \ - index, \ - source) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ======================================================================= -// 2. 获取 Workspace 大小 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopGetIndexCopyWorkspaceSize(infiniopIndexCopyDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -// ======================================================================= -// 3. 执行计算 (Calculate) -// ======================================================================= -__INFINI_C infiniStatus_t infiniopIndexCopy( - infiniopIndexCopyDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *output, - const void *input, - const void *index, - const void *source, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, output, input, index, source, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -// ======================================================================= -// 4. 销毁描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopDestroyIndexCopyDescriptor(infiniopIndexCopyDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} - -} // extern "C" diff --git a/src/infiniop/ops/inner/cpu/inner_cpu.cc b/src/infiniop/ops/inner/cpu/inner_cpu.cc deleted file mode 100644 index 085e46478..000000000 --- a/src/infiniop/ops/inner/cpu/inner_cpu.cc +++ /dev/null @@ -1,87 +0,0 @@ -#include "inner_cpu.h" -#include "../../../devices/cpu/common_cpu.h" - -namespace op::inner::cpu { - -Descriptor::~Descriptor() {} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t other_desc) { - - auto result = InnerInfo::create(out_desc, input_desc, other_desc); - CHECK_RESULT(result); - *desc_ptr = new Descriptor(nullptr, result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -infiniStatus_t inner(const InnerInfo *info, const T *input, const T *other, T *out) { -#pragma omp parallel for - for (ptrdiff_t out_index = 0; out_index < ptrdiff_t(info->total_elements); out_index++) { - - size_t out_dim_pos = info->out_ndim - 1; - size_t out_offset = op::common_cpu::indexToOffset(out_index, info->out_ndim, info->out_shape.data(), info->out_strides.data()); - size_t index = out_index; - - ptrdiff_t input_offset = 0; - ptrdiff_t other_offset = 0; - - for (int i = (int)info->other_ndim - 2; i >= 0; i--) { - other_offset += (index % info->out_shape[out_dim_pos]) * info->other_strides[i]; - index /= info->out_shape[out_dim_pos--]; - } - for (int i = (int)info->input_ndim - 2; i >= 0; i--) { - input_offset += (index % info->out_shape[out_dim_pos]) * info->input_strides[i]; - index /= info->out_shape[out_dim_pos--]; - } - - float tmp = 0.; - for (size_t i = 0; i < info->oper_len; i++) { - if constexpr (std::is_same::value || std::is_same::value) { - tmp += utils::cast(input[input_offset]) * utils::cast(other[other_offset]); - } else { - tmp += input[input_offset] * other[other_offset]; - } - input_offset += info->input_strides[info->input_ndim - 1]; - other_offset += info->other_strides[info->other_ndim - 1]; - } - - if constexpr (std::is_same::value || std::is_same::value) { - out[out_offset] = utils::cast(tmp); - } else { - out[out_offset] = tmp; - } - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *out, - const void *input, - const void *other, - void *stream) const { - - if (_info.dtype == INFINI_DTYPE_BF16) { - CHECK_STATUS(inner(&_info, (const bf16_t *)input, (const bf16_t *)other, (bf16_t *)out)); - } else if (_info.dtype == INFINI_DTYPE_F16) { - CHECK_STATUS(inner(&_info, (const fp16_t *)input, (const fp16_t *)other, (fp16_t *)out)); - } else if (_info.dtype == INFINI_DTYPE_F32) { - CHECK_STATUS(inner(&_info, (const float *)input, (const float *)other, (float *)out)); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::inner::cpu diff --git a/src/infiniop/ops/inner/cpu/inner_cpu.h b/src/infiniop/ops/inner/cpu/inner_cpu.h deleted file mode 100644 index 18b2e7b32..000000000 --- a/src/infiniop/ops/inner/cpu/inner_cpu.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __INNER_CPU_H__ -#define __INNER_CPU_H__ -#include "../inner.h" - -DESCRIPTOR(cpu); - -#endif diff --git a/src/infiniop/ops/inner/cuda/kernel.cuh b/src/infiniop/ops/inner/cuda/kernel.cuh deleted file mode 100644 index 928161ba8..000000000 --- a/src/infiniop/ops/inner/cuda/kernel.cuh +++ /dev/null @@ -1,43 +0,0 @@ -#ifndef __INNER_KERNEL_CUH__ -#define __INNER_KERNEL_CUH__ - -template -INFINIOP_CUDA_KERNEL innerKernel( - const T *input, const T *other, T *out, - size_t total_elements, size_t oper_len, size_t *out_shape, - ptrdiff_t *input_strides, ptrdiff_t *other_strides, ptrdiff_t *out_strides, - size_t input_ndim, size_t other_ndim, size_t out_ndim) { - - size_t out_index = blockDim.x * blockIdx.x + threadIdx.x; - if (out_index >= total_elements) { - return; - } - - size_t out_offset = device::nvidia::indexToOffset(out_index, out_ndim, out_shape, out_strides); - - size_t out_dim_pos = out_ndim - 1; - size_t index = out_index; - - ptrdiff_t input_offset = 0; - ptrdiff_t other_offset = 0; - - for (int i = (int)other_ndim - 2; i >= 0; i--) { - other_offset += (index % out_shape[out_dim_pos]) * other_strides[i]; - index /= out_shape[out_dim_pos--]; - } - for (int i = (int)input_ndim - 2; i >= 0; i--) { - input_offset += (index % out_shape[out_dim_pos]) * input_strides[i]; - index /= out_shape[out_dim_pos--]; - } - - T tmp = 0; - for (size_t i = 0; i < oper_len; i++) { - tmp += input[input_offset] * other[other_offset]; - input_offset += input_strides[input_ndim - 1]; - other_offset += other_strides[other_ndim - 1]; - } - - out[out_offset] = tmp; -} - -#endif // __INNER_KERNEL_CUH__ diff --git a/src/infiniop/ops/inner/info.h b/src/infiniop/ops/inner/info.h deleted file mode 100644 index 3262f9cf9..000000000 --- a/src/infiniop/ops/inner/info.h +++ /dev/null @@ -1,95 +0,0 @@ -#ifndef __INNER_INFO_H__ -#define __INNER_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::inner { - -class InnerInfo { - InnerInfo() = default; - -public: - infiniDtype_t dtype; - - size_t input_ndim; - size_t other_ndim; - size_t out_ndim; - - size_t oper_len; - size_t total_elements; - - std::vector out_shape; - - std::vector input_strides; - std::vector other_strides; - std::vector out_strides; - - static utils::Result create( - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t other_desc) { - - auto dtype = out_desc->dtype(); - if (dtype != input_desc->dtype() || dtype != other_desc->dtype()) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - CHECK_DTYPE(dtype, INFINI_DTYPE_BF16, INFINI_DTYPE_F16, INFINI_DTYPE_F32); - - auto input_ndim = input_desc->ndim(); - auto other_ndim = other_desc->ndim(); - auto out_ndim = out_desc->ndim(); - if (out_ndim + 2 != input_ndim + other_ndim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - auto input_shape = input_desc->shape(); - auto other_shape = other_desc->shape(); - auto out_shape = out_desc->shape(); - if (input_ndim && other_ndim) { - if (input_shape[input_ndim - 1] != other_shape[other_ndim - 1]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - size_t out_dim_pos = 0; - if (input_ndim) { - for (size_t i = 0; i < input_ndim - 1; i++) { - if (input_shape[i] != out_shape[out_dim_pos++]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - } - if (other_ndim) { - for (size_t i = 0; i < other_ndim - 1; i++) { - if (other_shape[i] != out_shape[out_dim_pos++]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - } - - size_t oper_len = 1; - if (input_ndim && other_ndim) { - oper_len = input_shape[input_ndim - 1]; - } - - size_t total_elements = 1; - for (size_t i = 0; i < out_ndim; i++) { - total_elements *= out_shape[i]; - } - - auto input_strides = input_desc->strides(); - auto other_strides = other_desc->strides(); - auto out_strides = out_desc->strides(); - - return utils::Result(InnerInfo{ - dtype, - input_ndim, other_ndim, out_ndim, - oper_len, total_elements, out_shape, - input_strides, other_strides, out_strides}); - } -}; - -} // namespace op::inner - -#endif // __INNER_INFO_H__ diff --git a/src/infiniop/ops/inner/inner.h b/src/infiniop/ops/inner/inner.h deleted file mode 100644 index c57e82708..000000000 --- a/src/infiniop/ops/inner/inner.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef Inner_H -#define Inner_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::inner::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - InnerInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - InnerInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t out_desc, \ - infiniopTensorDescriptor_t input_desc, \ - infiniopTensorDescriptor_t other_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *out, \ - const void *input, \ - const void *other, \ - void *stream) const; \ - }; \ - } - -#endif diff --git a/src/infiniop/ops/inner/metax/inner_metax.h b/src/infiniop/ops/inner/metax/inner_metax.h deleted file mode 100644 index 72d989a3d..000000000 --- a/src/infiniop/ops/inner/metax/inner_metax.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __INNER_METAX_H__ -#define __INNER_METAX_H__ -#include "../inner.h" - -DESCRIPTOR(metax); - -#endif diff --git a/src/infiniop/ops/inner/metax/inner_metax.maca b/src/infiniop/ops/inner/metax/inner_metax.maca deleted file mode 100644 index 95c8dcad7..000000000 --- a/src/infiniop/ops/inner/metax/inner_metax.maca +++ /dev/null @@ -1,122 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "inner_metax.h" - -#include "../../../devices/metax/metax_kernel_common.h" - -#include "inner_metax_kernel.h" - -namespace op::inner::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t other_desc) { - - auto result = InnerInfo::create(out_desc, input_desc, other_desc); - CHECK_RESULT(result); - auto info = result.take(); - size_t workspace_size = 0; - // out_shape - workspace_size += info.out_ndim * sizeof(size_t); - // input, other, out strides - workspace_size += (info.input_ndim + info.other_ndim + info.out_ndim) * sizeof(ptrdiff_t); - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info, workspace_size, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -infiniStatus_t launchKernel( - const InnerInfo &info, - const T *input, const T *other, T *out, - hcStream_t stream, void *workspace, size_t workspace_size) { - - size_t input_ndim = info.input_ndim; - size_t other_ndim = info.other_ndim; - size_t out_ndim = info.out_ndim; - size_t total_elements = info.total_elements; - size_t oper_len = info.oper_len; - - size_t workspace_offset = 0; - unsigned char *workspace_ptr = reinterpret_cast(workspace); - - size_t *out_shape_hc = reinterpret_cast(workspace_ptr + workspace_offset); - workspace_offset += out_ndim * sizeof(size_t); - - ptrdiff_t *input_strides_hc = reinterpret_cast(workspace_ptr + workspace_offset); - ptrdiff_t *other_strides_hc = input_strides_hc + input_ndim; - ptrdiff_t *out_strides_hc = other_strides_hc + other_ndim; - workspace_offset += (info.input_ndim + info.other_ndim + info.out_ndim) * sizeof(ptrdiff_t); - - assert(workspace_offset == workspace_size); - - CHECK_METAX(hcMemcpyAsync(out_shape_hc, info.out_shape.data(), out_ndim * sizeof(size_t), hcMemcpyHostToDevice, stream)); - CHECK_METAX(hcMemcpyAsync(input_strides_hc, info.input_strides.data(), input_ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, stream)); - CHECK_METAX(hcMemcpyAsync(other_strides_hc, info.other_strides.data(), other_ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, stream)); - CHECK_METAX(hcMemcpyAsync(out_strides_hc, info.out_strides.data(), out_ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, stream)); - - size_t block_size = BLOCK_SIZE; - size_t grid_size = (total_elements + BLOCK_SIZE - 1) / BLOCK_SIZE; - - innerKernel<<>>( - input, other, out, - total_elements, oper_len, out_shape_hc, - input_strides_hc, other_strides_hc, out_strides_hc, - input_ndim, other_ndim, out_ndim); - - // CHECK_METAX(hcDeviceSynchronize()); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *out, - const void *input, - const void *other, - void *stream_) const { - - hcStream_t stream = (hcStream_t)stream_; -#define CALCULATE_INNER(BLOCK_SIZE, T) \ - launchKernel( \ - _info, \ - (const T *)input, (const T *)other, (T *)out, \ - stream, workspace, workspace_size) -#define CALCULATE_INNER_WITH_BLOCK_SIZE(BLOCK_SIZE) \ - { \ - if (_info.dtype == INFINI_DTYPE_BF16) \ - return CALCULATE_INNER(BLOCK_SIZE, __hpcc_bfloat16); \ - else if (_info.dtype == INFINI_DTYPE_F16) \ - return CALCULATE_INNER(BLOCK_SIZE, half); \ - else if (_info.dtype == INFINI_DTYPE_F32) \ - return CALCULATE_INNER(BLOCK_SIZE, float); \ - else \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } - - if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_1024) { - CALCULATE_INNER_WITH_BLOCK_SIZE(METAX_BLOCK_SIZE_1024) - } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_512) { - CALCULATE_INNER_WITH_BLOCK_SIZE(METAX_BLOCK_SIZE_512) - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::inner::metax diff --git a/src/infiniop/ops/inner/metax/inner_metax_kernel.h b/src/infiniop/ops/inner/metax/inner_metax_kernel.h deleted file mode 100644 index 3e66a4590..000000000 --- a/src/infiniop/ops/inner/metax/inner_metax_kernel.h +++ /dev/null @@ -1,61 +0,0 @@ -#ifndef __INNER_METAX_KERNEL_H__ -#define __INNER_METAX_KERNEL_H__ - -template -INFINIOP_METAX_KERNEL innerKernel( - const T *input, const T *other, T *out, - size_t total_elements, size_t oper_len, size_t *out_shape, - ptrdiff_t *input_strides, ptrdiff_t *other_strides, ptrdiff_t *out_strides, - size_t input_ndim, size_t other_ndim, size_t out_ndim) { - - size_t out_index = blockDim.x * blockIdx.x + threadIdx.x; - if (out_index >= total_elements) { - return; - } - - size_t out_offset = device::metax::indexToOffset(out_index, out_ndim, out_shape, out_strides); - - size_t out_dim_pos = out_ndim - 1; - size_t index = out_index; - - ptrdiff_t input_offset = 0; - ptrdiff_t other_offset = 0; - - for (int i = (int)other_ndim - 2; i >= 0; i--) { - other_offset += (index % out_shape[out_dim_pos]) * other_strides[i]; - index /= out_shape[out_dim_pos--]; - } - for (int i = (int)input_ndim - 2; i >= 0; i--) { - input_offset += (index % out_shape[out_dim_pos]) * input_strides[i]; - index /= out_shape[out_dim_pos--]; - } - - // T tmp = input[input_offset] * other[other_offset]; - // input_offset += input_strides[input_ndim - 1]; - // other_offset += other_strides[other_ndim - 1]; - T tmp = 0.; - for (size_t i = 0; i < oper_len; i++) { - if constexpr (std::is_same_v || std::is_same_v) { - tmp += input[input_offset] * other[other_offset]; - } else { - // if (out_index == 0) { - // printf("input: %.10f\tother:%.10f\n", input[input_offset], other[other_offset]); - // } - T mul_a = __fmul_rn(input[input_offset], other[other_offset]); - // if (out_index == 0) { - // printf("mul answer: %.10f\n", mul_a); - // printf("tmp val: %.10f\n", tmp); - // } - tmp = __fadd_rd(tmp, mul_a); - // if (out_index == 0) - // printf("add answer: %.10f\n", tmp); - } - input_offset += input_strides[input_ndim - 1]; - other_offset += other_strides[other_ndim - 1]; - } - - // out[out_offset] = static_cast(tmp); - out[out_offset] = tmp; -} - -#endif // __INNER_KERNEL_CUH__ diff --git a/src/infiniop/ops/inner/moore/inner_moore.h b/src/infiniop/ops/inner/moore/inner_moore.h deleted file mode 100644 index e722f2adf..000000000 --- a/src/infiniop/ops/inner/moore/inner_moore.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __INNER_MOORE_H__ -#define __INNER_MOORE_H__ -#include "../inner.h" - -DESCRIPTOR(moore); - -#endif diff --git a/src/infiniop/ops/inner/moore/inner_moore.mu b/src/infiniop/ops/inner/moore/inner_moore.mu deleted file mode 100644 index 82abda6b4..000000000 --- a/src/infiniop/ops/inner/moore/inner_moore.mu +++ /dev/null @@ -1,122 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "inner_moore.h" - -#include "../../../devices/moore/moore_kernel_common.h" - -#include "inner_moore_kernel.h" - -namespace op::inner::moore { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t other_desc) { - - auto result = InnerInfo::create(out_desc, input_desc, other_desc); - CHECK_RESULT(result); - auto info = result.take(); - size_t workspace_size = 0; - // out_shape - workspace_size += info.out_ndim * sizeof(size_t); - // input, other, out strides - workspace_size += (info.input_ndim + info.other_ndim + info.out_ndim) * sizeof(ptrdiff_t); - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info, workspace_size, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -infiniStatus_t launchKernel( - const InnerInfo &info, - const T *input, const T *other, T *out, - musaStream_t stream, void *workspace, size_t workspace_size) { - - size_t input_ndim = info.input_ndim; - size_t other_ndim = info.other_ndim; - size_t out_ndim = info.out_ndim; - size_t total_elements = info.total_elements; - size_t oper_len = info.oper_len; - - size_t workspace_offset = 0; - unsigned char *workspace_ptr = reinterpret_cast(workspace); - - size_t *out_shape_musa = reinterpret_cast(workspace_ptr + workspace_offset); - workspace_offset += out_ndim * sizeof(size_t); - - ptrdiff_t *input_strides_musa = reinterpret_cast(workspace_ptr + workspace_offset); - ptrdiff_t *other_strides_musa = input_strides_musa + input_ndim; - ptrdiff_t *out_strides_musa = other_strides_musa + other_ndim; - workspace_offset += (info.input_ndim + info.other_ndim + info.out_ndim) * sizeof(ptrdiff_t); - - assert(workspace_offset == workspace_size); - - CHECK_MOORE(musaMemcpyAsync(out_shape_musa, info.out_shape.data(), out_ndim * sizeof(size_t), musaMemcpyHostToDevice, stream)); - CHECK_MOORE(musaMemcpyAsync(input_strides_musa, info.input_strides.data(), input_ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, stream)); - CHECK_MOORE(musaMemcpyAsync(other_strides_musa, info.other_strides.data(), other_ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, stream)); - CHECK_MOORE(musaMemcpyAsync(out_strides_musa, info.out_strides.data(), out_ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, stream)); - - size_t block_size = BLOCK_SIZE; - size_t grid_size = (total_elements + BLOCK_SIZE - 1) / BLOCK_SIZE; - - innerKernel<<>>( - input, other, out, - total_elements, oper_len, out_shape_musa, - input_strides_musa, other_strides_musa, out_strides_musa, - input_ndim, other_ndim, out_ndim); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *out, - const void *input, - const void *other, - void *stream_) const { - - musaStream_t stream = (musaStream_t)stream_; -#define CALCULATE_INNER(BLOCK_SIZE, T) \ - launchKernel( \ - _info, \ - (const T *)input, (const T *)other, (T *)out, \ - stream, workspace, workspace_size) -#define CALCULATE_INNER_WITH_BLOCK_SIZE(BLOCK_SIZE) \ - { \ - if (_info.dtype == INFINI_DTYPE_BF16) \ - return CALCULATE_INNER(BLOCK_SIZE, __mt_bfloat16); \ - else if (_info.dtype == INFINI_DTYPE_F16) \ - return CALCULATE_INNER(BLOCK_SIZE, half); \ - else if (_info.dtype == INFINI_DTYPE_F32) \ - return CALCULATE_INNER(BLOCK_SIZE, float); \ - else \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } - - if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_1024) { - CALCULATE_INNER_WITH_BLOCK_SIZE(MOORE_BLOCK_SIZE_1024) - } else if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_512) { - CALCULATE_INNER_WITH_BLOCK_SIZE(MOORE_BLOCK_SIZE_512) - } else if (_opaque->internal->maxThreadsPerBlock() == MOORE_BLOCK_SIZE_2048) { - CALCULATE_INNER_WITH_BLOCK_SIZE(MOORE_BLOCK_SIZE_2048) - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::inner::moore diff --git a/src/infiniop/ops/inner/moore/inner_moore_kernel.h b/src/infiniop/ops/inner/moore/inner_moore_kernel.h deleted file mode 100644 index fcb3f948f..000000000 --- a/src/infiniop/ops/inner/moore/inner_moore_kernel.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef __INNER_MOORE_KERNEL_H__ -#define __INNER_MOORE_KERNEL_H__ - -template -INFINIOP_MOORE_KERNEL innerKernel( - const T *input, const T *other, T *out, - size_t total_elements, size_t oper_len, size_t *out_shape, - ptrdiff_t *input_strides, ptrdiff_t *other_strides, ptrdiff_t *out_strides, - size_t input_ndim, size_t other_ndim, size_t out_ndim) { - - size_t out_index = blockDim.x * blockIdx.x + threadIdx.x; - if (out_index >= total_elements) { - return; - } - - size_t out_offset = device::moore::indexToOffset(out_index, out_ndim, out_shape, out_strides); - - size_t out_dim_pos = out_ndim - 1; - size_t index = out_index; - - ptrdiff_t input_offset = 0; - ptrdiff_t other_offset = 0; - - for (int i = (int)other_ndim - 2; i >= 0; i--) { - other_offset += (index % out_shape[out_dim_pos]) * other_strides[i]; - index /= out_shape[out_dim_pos--]; - } - for (int i = (int)input_ndim - 2; i >= 0; i--) { - input_offset += (index % out_shape[out_dim_pos]) * input_strides[i]; - index /= out_shape[out_dim_pos--]; - } - - float tmp = 0.; - for (size_t i = 0; i < oper_len; i++) { - // if constexpr (std::is_same_v || std::is_same_v) { - tmp += static_cast(input[input_offset]) * static_cast(other[other_offset]); - // } else { - // tmp += input[input_offset] * other[other_offset]; - // } - input_offset += input_strides[input_ndim - 1]; - other_offset += other_strides[other_ndim - 1]; - } - - out[out_offset] = static_cast(tmp); -} - -#endif // __INNER_KERNEL_CUH__ diff --git a/src/infiniop/ops/inner/nvidia/inner_nvidia.cu b/src/infiniop/ops/inner/nvidia/inner_nvidia.cu deleted file mode 100644 index a06e8955c..000000000 --- a/src/infiniop/ops/inner/nvidia/inner_nvidia.cu +++ /dev/null @@ -1,124 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "inner_nvidia.cuh" - -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" - -#include "../cuda/kernel.cuh" - -namespace op::inner::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t other_desc) { - - auto result = InnerInfo::create(out_desc, input_desc, other_desc); - CHECK_RESULT(result); - auto info = result.take(); - size_t workspace_size = 0; - // out_shape - workspace_size += info.out_ndim * sizeof(size_t); - // input, other, out strides - workspace_size += (info.input_ndim + info.other_ndim + info.out_ndim) * sizeof(ptrdiff_t); - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info, workspace_size, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -namespace { - -template -infiniStatus_t launchKernel( - const InnerInfo &info, - const T *input, const T *other, T *out, - cudaStream_t stream, void *workspace, size_t workspace_size) { - - size_t input_ndim = info.input_ndim; - size_t other_ndim = info.other_ndim; - size_t out_ndim = info.out_ndim; - size_t total_elements = info.total_elements; - size_t oper_len = info.oper_len; - - size_t workspace_offset = 0; - unsigned char *workspace_ptr = reinterpret_cast(workspace); - - size_t *out_shape_cuda = reinterpret_cast(workspace_ptr + workspace_offset); - workspace_offset += out_ndim * sizeof(size_t); - - ptrdiff_t *input_strides_cuda = reinterpret_cast(workspace_ptr + workspace_offset); - ptrdiff_t *other_strides_cuda = input_strides_cuda + input_ndim; - ptrdiff_t *out_strides_cuda = other_strides_cuda + other_ndim; - workspace_offset += (info.input_ndim + info.other_ndim + info.out_ndim) * sizeof(ptrdiff_t); - - CHECK_CUDA(cudaMemcpyAsync(out_shape_cuda, info.out_shape.data(), out_ndim * sizeof(size_t), cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemcpyAsync(input_strides_cuda, info.input_strides.data(), input_ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemcpyAsync(other_strides_cuda, info.other_strides.data(), other_ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemcpyAsync(out_strides_cuda, info.out_strides.data(), out_ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, stream)); - - size_t block_size = BLOCK_SIZE; - size_t grid_size = (total_elements + BLOCK_SIZE - 1) / BLOCK_SIZE; - - innerKernel<<>>( - input, other, out, - total_elements, oper_len, out_shape_cuda, - input_strides_cuda, other_strides_cuda, out_strides_cuda, - input_ndim, other_ndim, out_ndim); - - // CHECK_CUDA(cudaDeviceSynchronize()); - - return INFINI_STATUS_SUCCESS; -} - -} // namespace - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *out, - const void *input, - const void *other, - void *stream_) const { - - cudaStream_t stream = (cudaStream_t)stream_; -#define CALCULATE_INNER(BLOCK_SIZE, T) \ - launchKernel( \ - _info, \ - (const T *)input, (const T *)other, (T *)out, \ - stream, workspace, workspace_size) -#define CALCULATE_INNER_WITH_BLOCK_SIZE(BLOCK_SIZE) \ - { \ - if (_info.dtype == INFINI_DTYPE_BF16) \ - return CALCULATE_INNER(BLOCK_SIZE, __nv_bfloat16); \ - else if (_info.dtype == INFINI_DTYPE_F16) \ - return CALCULATE_INNER(BLOCK_SIZE, half); \ - else if (_info.dtype == INFINI_DTYPE_F32) \ - return CALCULATE_INNER(BLOCK_SIZE, float); \ - else \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } - - if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_1024) { - CALCULATE_INNER_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_1024) - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_512) { - CALCULATE_INNER_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_512) - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_2048) { - CALCULATE_INNER_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_2048) - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) { - CALCULATE_INNER_WITH_BLOCK_SIZE(CUDA_BLOCK_SIZE_4096) - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::inner::nvidia diff --git a/src/infiniop/ops/inner/nvidia/inner_nvidia.cuh b/src/infiniop/ops/inner/nvidia/inner_nvidia.cuh deleted file mode 100644 index 4a3b8cb95..000000000 --- a/src/infiniop/ops/inner/nvidia/inner_nvidia.cuh +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __INNER_NVIDIA_H__ -#define __INNER_NVIDIA_H__ -#include "../inner.h" - -DESCRIPTOR(nvidia); - -#endif diff --git a/src/infiniop/ops/inner/operator.cc b/src/infiniop/ops/inner/operator.cc deleted file mode 100644 index 54ed677c2..000000000 --- a/src/infiniop/ops/inner/operator.cc +++ /dev/null @@ -1,151 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/inner.h" - -#ifdef ENABLE_CPU_API -#include "cpu/inner_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) -#include "nvidia/inner_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/inner_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/inner_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateInnerDescriptor( - infiniopHandle_t handle, - infiniopInnerDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t out_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t other_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::inner::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - out_desc, \ - input_desc, \ - other_desc); - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore) -#endif - } - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopGetInnerWorkspaceSize(infiniopInnerDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif - } - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopInner( - infiniopInnerDescriptor_t desc, - void *workspace, size_t workspace_size, - void *out, - const void *input, - const void *other, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc)->calculate( \ - workspace, workspace_size, out, input, other, stream); - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore) -#endif - } - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopDestroyInnerDescriptor(infiniopInnerDescriptor_t desc) { - -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DESTROY(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DESTROY(INFINI_DEVICE_HYGON, nvidia); -#endif -#ifdef ENABLE_METAX_API - DESTROY(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - DESTROY(INFINI_DEVICE_MOORE, moore) -#endif - } - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} diff --git a/src/infiniop/ops/interpolate/cpu/interpolate_cpu.cc b/src/infiniop/ops/interpolate/cpu/interpolate_cpu.cc deleted file mode 100644 index 58e90bf49..000000000 --- a/src/infiniop/ops/interpolate/cpu/interpolate_cpu.cc +++ /dev/null @@ -1,415 +0,0 @@ -#include "interpolate_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include -#include -#include - -namespace op::interpolate::cpu { - -inline size_t offset_3d(const std::vector &s, - size_t b, size_t c, size_t x) { - return b * s[0] + c * s[1] + x * s[2]; -} - -inline size_t offset_4d(const std::vector &s, - size_t b, size_t c, size_t h, size_t w) { - return b * s[0] + c * s[1] + h * s[2] + w * s[3]; -} - -static bool try_parse_mode(const char *mode_str, InterpolateMode &mode) { - if (std::strcmp(mode_str, "nearest") == 0) { - mode = InterpolateMode::NEAREST; - return true; - } else if (std::strcmp(mode_str, "linear") == 0) { - mode = InterpolateMode::LINEAR; - return true; - } else if (std::strcmp(mode_str, "bilinear") == 0) { - mode = InterpolateMode::BILINEAR; - return true; - } else if (std::strcmp(mode_str, "trilinear") == 0) { - mode = InterpolateMode::TRILINEAR; - return true; - } else if (std::strcmp(mode_str, "area") == 0) { - mode = InterpolateMode::AREA; - return true; - } - return false; -} - -static double compute_scale(size_t in_size, size_t out_size, int align_corners) { - if (out_size == 0) { - return 0.0; - } - if (align_corners) { - return (out_size > 1) ? (static_cast(in_size) - 1.0) / (static_cast(out_size) - 1.0) : 0.0; - } - return static_cast(in_size) / static_cast(out_size); -} - -utils::Result InterpolateInfo::create( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - const char *mode_str, - void *size, - void *scale_factor, - int align_corners) { - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x_shape.size() < 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if ((size != nullptr) == (scale_factor != nullptr)) { - return INFINI_STATUS_BAD_PARAM; - } - - if (y_shape.size() != x_shape.size() || y_shape[0] != x_shape[0] || y_shape[1] != x_shape[1]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t ndim = x_shape.size() - 2; - - if (scale_factor != nullptr) { - const double *scale_array = reinterpret_cast(scale_factor); - const double scale = scale_array[0]; - std::vector expected_y_shape = x_shape; - for (size_t i = 0; i < ndim; ++i) { - expected_y_shape[i + 2] = static_cast(static_cast(x_shape[i + 2]) * scale); - } - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - InterpolateInfo info; - info.ndim = ndim; - info.input_shape = x_shape; - info.output_shape = y_shape; - if (!try_parse_mode(mode_str, info.mode)) { - return INFINI_STATUS_BAD_PARAM; - } - info.input_strides = x_desc->strides(); - info.output_strides = y_desc->strides(); - info.align_corners = align_corners; - info.input_size = x_desc->numel(); - info.output_size = y_desc->numel(); - - return utils::Result(std::move(info)); -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - const char *mode, - void *size, - void *scale_factor, - int align_corners) { - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto info_result = InterpolateInfo::create(x_desc, y_desc, mode, size, scale_factor, align_corners); - CHECK_RESULT(info_result); - - *desc_ptr = new Descriptor(dtype, info_result.take(), handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -void interpolate_nearest_1d( - const InterpolateInfo &info, - const T *input, T *output, - size_t batch, size_t channels, - size_t in_w, size_t out_w, - int align_corners) { - - double scale = compute_scale(in_w, out_w, align_corners); - - const auto &in_s = info.input_strides; - const auto &out_s = info.output_strides; - - for (size_t b = 0; b < batch; ++b) { - for (size_t c = 0; c < channels; ++c) { - for (size_t ow = 0; ow < out_w; ++ow) { - double src_x = align_corners ? ow * scale : (ow + 0.5) * scale - 0.5; - size_t ix = std::min(static_cast(std::round(src_x)), in_w - 1); - output[offset_3d(out_s, b, c, ow)] = input[offset_3d(in_s, b, c, ix)]; - } - } - } -} - -template -void interpolate_bilinear_2d( - const InterpolateInfo &info, - const T *input, T *output, - size_t batch, size_t channels, - size_t in_h, size_t in_w, - size_t out_h, size_t out_w, - int align_corners) { - - double scale_h = compute_scale(in_h, out_h, align_corners); - double scale_w = compute_scale(in_w, out_w, align_corners); - - const auto &in_s = info.input_strides; - const auto &out_s = info.output_strides; - - for (size_t b = 0; b < batch; ++b) { - for (size_t c = 0; c < channels; ++c) { - for (size_t oh = 0; oh < out_h; ++oh) { - for (size_t ow = 0; ow < out_w; ++ow) { - double src_y = align_corners ? oh * scale_h : (oh + 0.5) * scale_h - 0.5; - double src_x = align_corners ? ow * scale_w : (ow + 0.5) * scale_w - 0.5; - - src_y = std::max(0.0, std::min(src_y, (double)(in_h - 1))); - src_x = std::max(0.0, std::min(src_x, (double)(in_w - 1))); - - size_t y0 = (size_t)std::floor(src_y); - size_t y1 = std::min(y0 + 1, in_h - 1); - size_t x0 = (size_t)std::floor(src_x); - size_t x1 = std::min(x0 + 1, in_w - 1); - - double dy = src_y - y0; - double dx = src_x - x0; - - T v00 = input[offset_4d(in_s, b, c, y0, x0)]; - T v01 = input[offset_4d(in_s, b, c, y0, x1)]; - T v10 = input[offset_4d(in_s, b, c, y1, x0)]; - T v11 = input[offset_4d(in_s, b, c, y1, x1)]; - - T result = utils::cast( - (1 - dy) * (1 - dx) * utils::cast(v00) + (1 - dy) * dx * utils::cast(v01) + dy * (1 - dx) * utils::cast(v10) + dy * dx * utils::cast(v11)); - - output[offset_4d(out_s, b, c, oh, ow)] = result; - } - } - } - } -} - -template -void interpolate_trilinear_3d( - const InterpolateInfo &info, - const T *input, T *output, - size_t batch, size_t channels, - size_t in_d, size_t in_h, size_t in_w, - size_t out_d, size_t out_h, size_t out_w, - int align_corners) { - - double scale_d = compute_scale(in_d, out_d, align_corners); - double scale_h = compute_scale(in_h, out_h, align_corners); - double scale_w = compute_scale(in_w, out_w, align_corners); - - const auto &in_s = info.input_strides; - const auto &out_s = info.output_strides; - - auto offset_5d = [](const std::vector &s, - size_t b, size_t c, size_t d, size_t h, size_t w) { - return b * s[0] + c * s[1] + d * s[2] + h * s[3] + w * s[4]; - }; - - for (size_t b = 0; b < batch; ++b) { - for (size_t c = 0; c < channels; ++c) { - for (size_t od = 0; od < out_d; ++od) { - for (size_t oh = 0; oh < out_h; ++oh) { - for (size_t ow = 0; ow < out_w; ++ow) { - - double src_d = align_corners ? od * scale_d : (od + 0.5) * scale_d - 0.5; - double src_h = align_corners ? oh * scale_h : (oh + 0.5) * scale_h - 0.5; - double src_w = align_corners ? ow * scale_w : (ow + 0.5) * scale_w - 0.5; - - src_d = std::max(0.0, std::min(src_d, (double)(in_d - 1))); - src_h = std::max(0.0, std::min(src_h, (double)(in_h - 1))); - src_w = std::max(0.0, std::min(src_w, (double)(in_w - 1))); - - size_t d0 = (size_t)std::floor(src_d); - size_t d1 = std::min(d0 + 1, in_d - 1); - size_t h0 = (size_t)std::floor(src_h); - size_t h1 = std::min(h0 + 1, in_h - 1); - size_t w0 = (size_t)std::floor(src_w); - size_t w1 = std::min(w0 + 1, in_w - 1); - - double dd = src_d - d0; - double dh = src_h - h0; - double dw = src_w - w0; - - // 8 corners - T c000 = input[offset_5d(in_s, b, c, d0, h0, w0)]; - T c001 = input[offset_5d(in_s, b, c, d0, h0, w1)]; - T c010 = input[offset_5d(in_s, b, c, d0, h1, w0)]; - T c011 = input[offset_5d(in_s, b, c, d0, h1, w1)]; - T c100 = input[offset_5d(in_s, b, c, d1, h0, w0)]; - T c101 = input[offset_5d(in_s, b, c, d1, h0, w1)]; - T c110 = input[offset_5d(in_s, b, c, d1, h1, w0)]; - T c111 = input[offset_5d(in_s, b, c, d1, h1, w1)]; - - // trilinear interpolation - double v = (1 - dd) * (1 - dh) * (1 - dw) * utils::cast(c000) + (1 - dd) * (1 - dh) * dw * utils::cast(c001) + (1 - dd) * dh * (1 - dw) * utils::cast(c010) + (1 - dd) * dh * dw * utils::cast(c011) + dd * (1 - dh) * (1 - dw) * utils::cast(c100) + dd * (1 - dh) * dw * utils::cast(c101) + dd * dh * (1 - dw) * utils::cast(c110) + dd * dh * dw * utils::cast(c111); - - output[offset_5d(out_s, b, c, od, oh, ow)] = utils::cast(v); - } - } - } - } - } -} - -template -void interpolate_impl(const InterpolateInfo &info, T *y, const T *x) { - - size_t batch = info.input_shape[0]; - size_t channels = info.input_shape[1]; - - const auto &in_s = info.input_strides; - const auto &out_s = info.output_strides; - - if (info.mode == InterpolateMode::NEAREST) { - if (info.ndim == 1) { - interpolate_nearest_1d(info, x, y, batch, channels, - info.input_shape[2], info.output_shape[2], - info.align_corners); - } else if (info.ndim == 2) { - size_t in_h = info.input_shape[2]; - size_t in_w = info.input_shape[3]; - size_t out_h = info.output_shape[2]; - size_t out_w = info.output_shape[3]; - - double scale_h = compute_scale(in_h, out_h, info.align_corners); - double scale_w = compute_scale(in_w, out_w, info.align_corners); - - for (size_t b = 0; b < batch; ++b) { - for (size_t c = 0; c < channels; ++c) { - for (size_t oh = 0; oh < out_h; ++oh) { - for (size_t ow = 0; ow < out_w; ++ow) { - double src_y = info.align_corners ? oh * scale_h : (oh + 0.5) * scale_h - 0.5; - double src_x = info.align_corners ? ow * scale_w : (ow + 0.5) * scale_w - 0.5; - - size_t iy = std::min((size_t)std::round(src_y), in_h - 1); - size_t ix = std::min((size_t)std::round(src_x), in_w - 1); - - y[offset_4d(out_s, b, c, oh, ow)] = x[offset_4d(in_s, b, c, iy, ix)]; - } - } - } - } - } - } else if (info.mode == InterpolateMode::LINEAR || info.mode == InterpolateMode::BILINEAR || info.mode == InterpolateMode::TRILINEAR) { - if (info.ndim == 1) { - size_t in_w = info.input_shape[2]; - size_t out_w = info.output_shape[2]; - double scale = compute_scale(in_w, out_w, info.align_corners); - - for (size_t b = 0; b < batch; ++b) { - for (size_t c = 0; c < channels; ++c) { - for (size_t ow = 0; ow < out_w; ++ow) { - double src_x = info.align_corners ? ow * scale : (ow + 0.5) * scale - 0.5; - src_x = std::max(0.0, std::min(src_x, (double)(in_w - 1))); - - size_t x0 = (size_t)std::floor(src_x); - size_t x1 = std::min(x0 + 1, in_w - 1); - - double dx = src_x - x0; - - T v0 = x[offset_3d(in_s, b, c, x0)]; - T v1 = x[offset_3d(in_s, b, c, x1)]; - - y[offset_3d(out_s, b, c, ow)] = utils::cast((1 - dx) * utils::cast(v0) + dx * utils::cast(v1)); - } - } - } - } else if (info.ndim == 2) { - interpolate_bilinear_2d(info, x, y, batch, channels, - info.input_shape[2], info.input_shape[3], - info.output_shape[2], info.output_shape[3], - info.align_corners); - } else if (info.ndim == 3) { - size_t in_d = info.input_shape[2]; - size_t in_h = info.input_shape[3]; - size_t in_w = info.input_shape[4]; - size_t out_d = info.output_shape[2]; - size_t out_h = info.output_shape[3]; - size_t out_w = info.output_shape[4]; - - interpolate_trilinear_3d(info, x, y, batch, channels, - in_d, in_h, in_w, - out_d, out_h, out_w, - info.align_corners); - } - } else if (info.mode == InterpolateMode::AREA) { - size_t in_h = info.input_shape[2]; - size_t in_w = info.input_shape[3]; - size_t out_h = info.output_shape[2]; - size_t out_w = info.output_shape[3]; - - double scale_h = (double)in_h / out_h; - double scale_w = (double)in_w / out_w; - - for (size_t b = 0; b < batch; ++b) { - for (size_t c = 0; c < channels; ++c) { - for (size_t oh = 0; oh < out_h; ++oh) { - for (size_t ow = 0; ow < out_w; ++ow) { - - double start_h = oh * scale_h; - double end_h = (oh + 1) * scale_h; - double start_w = ow * scale_w; - double end_w = (ow + 1) * scale_w; - - size_t h0 = (size_t)std::floor(start_h); - size_t h1 = (size_t)std::ceil(end_h); - size_t w0 = (size_t)std::floor(start_w); - size_t w1 = (size_t)std::ceil(end_w); - - double sum = 0.0; - size_t count = 0; - - for (size_t ih = h0; ih < h1 && ih < in_h; ++ih) { - for (size_t iw = w0; iw < w1 && iw < in_w; ++iw) { - sum += utils::cast( - x[offset_4d(in_s, b, c, ih, iw)]); - count++; - } - } - - y[offset_4d(out_s, b, c, oh, ow)] = count > 0 ? utils::cast(sum / (double)count) - : utils::cast(0.0); - } - } - } - } - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - interpolate_impl(_info, reinterpret_cast(y), reinterpret_cast(x)); - break; - case INFINI_DTYPE_BF16: - interpolate_impl(_info, reinterpret_cast(y), reinterpret_cast(x)); - break; - case INFINI_DTYPE_F32: - interpolate_impl(_info, reinterpret_cast(y), reinterpret_cast(x)); - break; - case INFINI_DTYPE_F64: - interpolate_impl(_info, reinterpret_cast(y), reinterpret_cast(x)); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::interpolate::cpu diff --git a/src/infiniop/ops/interpolate/cpu/interpolate_cpu.h b/src/infiniop/ops/interpolate/cpu/interpolate_cpu.h deleted file mode 100644 index 72af07c64..000000000 --- a/src/infiniop/ops/interpolate/cpu/interpolate_cpu.h +++ /dev/null @@ -1,74 +0,0 @@ -#ifndef __INTERPOLATE_CPU_H__ -#define __INTERPOLATE_CPU_H__ - -#include "../../../operator.h" -#include "../../../tensor.h" -#include -#include - -namespace op::interpolate::cpu { - -enum class InterpolateMode { - NEAREST, - LINEAR, - BILINEAR, - TRILINEAR, - AREA -}; - -struct InterpolateInfo { - size_t ndim; - std::vector input_shape; - std::vector output_shape; - std::vector input_strides; - std::vector output_strides; - InterpolateMode mode; - int align_corners; - size_t input_size; - size_t output_size; - - static utils::Result create( - infiniopTensorDescriptor_t x_desc, - infiniopTensorDescriptor_t y_desc, - const char *mode_str, - void *size, - void *scale_factor, - int align_corners); -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - InterpolateInfo _info; - - Descriptor(infiniDtype_t dtype, InterpolateInfo info, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - const char *mode, - void *size, - void *scale_factor, - int align_corners); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::interpolate::cpu - -#endif // __INTERPOLATE_CPU_H__ diff --git a/src/infiniop/ops/interpolate/cuda/kernel.cuh b/src/infiniop/ops/interpolate/cuda/kernel.cuh deleted file mode 100644 index 493155be9..000000000 --- a/src/infiniop/ops/interpolate/cuda/kernel.cuh +++ /dev/null @@ -1,436 +0,0 @@ -#pragma once -#include -#include - -namespace op::interpolate::cuda { - -constexpr int kInterpolateMaxDims = 8; - -struct TensorMeta { - int ndim; - size_t shape[kInterpolateMaxDims]; - ptrdiff_t strides[kInterpolateMaxDims]; // strides in elements -}; - -template -struct TypeTag {}; - -template -__device__ __forceinline__ Tcompute to_compute(const half v) { - return static_cast(__half2float(v)); -} - -template -__device__ __forceinline__ Tcompute to_compute(const cuda_bfloat16 v) { - return static_cast(__bfloat162float(v)); -} - -template -__device__ __forceinline__ Tcompute to_compute(const T v) { - return static_cast(v); -} - -__device__ __forceinline__ half from_compute(const float v, TypeTag) { - return __float2half_rn(v); -} - -__device__ __forceinline__ cuda_bfloat16 from_compute(const float v, TypeTag) { - return __float2bfloat16_rn(v); -} - -template -__device__ __forceinline__ T from_compute(const Tcompute v, TypeTag) { - return static_cast(v); -} - -__device__ __forceinline__ double compute_scale(size_t in_size, size_t out_size, bool align_corners) { - if (out_size == 0) { - return 0.0; - } - if (align_corners) { - if (out_size == 1) { - return 0.0; - } - return (static_cast(in_size) - 1.0) / (static_cast(out_size) - 1.0); - } - return static_cast(in_size) / static_cast(out_size); -} - -__device__ __forceinline__ double compute_source_index( - int64_t out_idx, - double scale, - bool align_corners) { - if (align_corners) { - return static_cast(out_idx) * scale; - } - return (static_cast(out_idx) + 0.5) * scale - 0.5; -} - -template -__global__ void nearest_2d_kernel( - T *output, - const T *input, - TensorMeta out_meta, - TensorMeta in_meta) { - - const size_t total = out_meta.shape[0] * out_meta.shape[1] * out_meta.shape[2] * out_meta.shape[3]; - const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= total) { - return; - } - - const size_t out_w = out_meta.shape[3]; - const size_t out_h = out_meta.shape[2]; - const size_t c_stride = out_h * out_w; - const size_t n_stride = out_meta.shape[1] * c_stride; - - const size_t n = idx / n_stride; - size_t rem = idx - n * n_stride; - const size_t c = rem / c_stride; - rem -= c * c_stride; - const size_t oh = rem / out_w; - const size_t ow = rem - oh * out_w; - - const size_t in_h = in_meta.shape[2]; - const size_t in_w = in_meta.shape[3]; - - const double scale_h = static_cast(in_h) / static_cast(out_h); - const double scale_w = static_cast(in_w) / static_cast(out_w); - - const size_t ih_raw = static_cast(floor(static_cast(oh) * scale_h)); - const size_t iw_raw = static_cast(floor(static_cast(ow) * scale_w)); - const size_t ih = (ih_raw < in_h) ? ih_raw : (in_h - 1); - const size_t iw = (iw_raw < in_w) ? iw_raw : (in_w - 1); - - const size_t in_off = n * in_meta.strides[0] + c * in_meta.strides[1] + ih * in_meta.strides[2] + iw * in_meta.strides[3]; - const size_t out_off = n * out_meta.strides[0] + c * out_meta.strides[1] + oh * out_meta.strides[2] + ow * out_meta.strides[3]; - - output[out_off] = input[in_off]; -} - -template -__global__ void linear_1d_kernel( - T *output, - const T *input, - TensorMeta out_meta, - TensorMeta in_meta, - int align_corners) { - - const size_t total = out_meta.shape[0] * out_meta.shape[1] * out_meta.shape[2]; - const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= total) { - return; - } - - const size_t out_w = out_meta.shape[2]; - const size_t c_stride = out_w; - const size_t n_stride = out_meta.shape[1] * c_stride; - - const size_t n = idx / n_stride; - size_t rem = idx - n * n_stride; - const size_t c = rem / c_stride; - const size_t ow = rem - c * c_stride; - - const size_t in_w = in_meta.shape[2]; - const bool use_align = (align_corners != 0); - const double scale = compute_scale(in_w, out_w, use_align); - - double src = compute_source_index(static_cast(ow), scale, use_align); - src = fmax(0.0, fmin(src, static_cast(in_w - 1))); - - const size_t x0 = static_cast(floor(src)); - const size_t x1 = (x0 + 1 < in_w) ? (x0 + 1) : x0; - const double t = src - static_cast(x0); - - const size_t off0 = n * in_meta.strides[0] + c * in_meta.strides[1] + x0 * in_meta.strides[2]; - const size_t off1 = n * in_meta.strides[0] + c * in_meta.strides[1] + x1 * in_meta.strides[2]; - const Tcompute v0 = to_compute(input[off0]); - const Tcompute v1 = to_compute(input[off1]); - const Tcompute out = static_cast((1.0 - t) * static_cast(v0) + t * static_cast(v1)); - - const size_t out_off = n * out_meta.strides[0] + c * out_meta.strides[1] + ow * out_meta.strides[2]; - output[out_off] = from_compute(out, TypeTag{}); -} - -template -__global__ void bilinear_2d_kernel( - T *output, - const T *input, - TensorMeta out_meta, - TensorMeta in_meta, - int align_corners) { - - const size_t total = out_meta.shape[0] * out_meta.shape[1] * out_meta.shape[2] * out_meta.shape[3]; - const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= total) { - return; - } - - const size_t out_w = out_meta.shape[3]; - const size_t out_h = out_meta.shape[2]; - const size_t c_stride = out_h * out_w; - const size_t n_stride = out_meta.shape[1] * c_stride; - - const size_t n = idx / n_stride; - size_t rem = idx - n * n_stride; - const size_t c = rem / c_stride; - rem -= c * c_stride; - const size_t oh = rem / out_w; - const size_t ow = rem - oh * out_w; - - const size_t in_h = in_meta.shape[2]; - const size_t in_w = in_meta.shape[3]; - const bool use_align = (align_corners != 0); - const double scale_h = compute_scale(in_h, out_h, use_align); - const double scale_w = compute_scale(in_w, out_w, use_align); - - double src_y = compute_source_index(static_cast(oh), scale_h, use_align); - double src_x = compute_source_index(static_cast(ow), scale_w, use_align); - src_y = fmax(0.0, fmin(src_y, static_cast(in_h - 1))); - src_x = fmax(0.0, fmin(src_x, static_cast(in_w - 1))); - - const size_t y0 = static_cast(floor(src_y)); - const size_t x0 = static_cast(floor(src_x)); - const size_t y1 = (y0 + 1 < in_h) ? (y0 + 1) : y0; - const size_t x1 = (x0 + 1 < in_w) ? (x0 + 1) : x0; - const double wy = src_y - static_cast(y0); - const double wx = src_x - static_cast(x0); - - const size_t base = n * in_meta.strides[0] + c * in_meta.strides[1]; - const Tcompute v00 = to_compute(input[base + y0 * in_meta.strides[2] + x0 * in_meta.strides[3]]); - const Tcompute v01 = to_compute(input[base + y0 * in_meta.strides[2] + x1 * in_meta.strides[3]]); - const Tcompute v10 = to_compute(input[base + y1 * in_meta.strides[2] + x0 * in_meta.strides[3]]); - const Tcompute v11 = to_compute(input[base + y1 * in_meta.strides[2] + x1 * in_meta.strides[3]]); - - const double w00 = (1.0 - wy) * (1.0 - wx); - const double w01 = (1.0 - wy) * wx; - const double w10 = wy * (1.0 - wx); - const double w11 = wy * wx; - - const Tcompute out = static_cast( - w00 * static_cast(v00) + w01 * static_cast(v01) + w10 * static_cast(v10) + w11 * static_cast(v11)); - - const size_t out_off = n * out_meta.strides[0] + c * out_meta.strides[1] + oh * out_meta.strides[2] + ow * out_meta.strides[3]; - output[out_off] = from_compute(out, TypeTag{}); -} - -template -__global__ void trilinear_3d_kernel( - T *output, - const T *input, - TensorMeta out_meta, - TensorMeta in_meta, - int align_corners) { - - const size_t total = out_meta.shape[0] * out_meta.shape[1] * out_meta.shape[2] * out_meta.shape[3] * out_meta.shape[4]; - const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= total) { - return; - } - - const size_t out_d = out_meta.shape[2]; - const size_t out_h = out_meta.shape[3]; - const size_t out_w = out_meta.shape[4]; - const size_t c_stride = out_d * out_h * out_w; - const size_t n_stride = out_meta.shape[1] * c_stride; - - const size_t n = idx / n_stride; - size_t rem = idx - n * n_stride; - const size_t c = rem / c_stride; - rem -= c * c_stride; - - const size_t od = rem / (out_h * out_w); - rem -= od * (out_h * out_w); - const size_t oh = rem / out_w; - const size_t ow = rem - oh * out_w; - - const size_t in_d = in_meta.shape[2]; - const size_t in_h = in_meta.shape[3]; - const size_t in_w = in_meta.shape[4]; - const bool use_align = (align_corners != 0); - const double scale_d = compute_scale(in_d, out_d, use_align); - const double scale_h = compute_scale(in_h, out_h, use_align); - const double scale_w = compute_scale(in_w, out_w, use_align); - - double src_d = compute_source_index(static_cast(od), scale_d, use_align); - double src_h = compute_source_index(static_cast(oh), scale_h, use_align); - double src_w = compute_source_index(static_cast(ow), scale_w, use_align); - - src_d = fmax(0.0, fmin(src_d, static_cast(in_d - 1))); - src_h = fmax(0.0, fmin(src_h, static_cast(in_h - 1))); - src_w = fmax(0.0, fmin(src_w, static_cast(in_w - 1))); - - const size_t d0 = static_cast(floor(src_d)); - const size_t h0 = static_cast(floor(src_h)); - const size_t w0 = static_cast(floor(src_w)); - const size_t d1 = (d0 + 1 < in_d) ? (d0 + 1) : d0; - const size_t h1 = (h0 + 1 < in_h) ? (h0 + 1) : h0; - const size_t w1 = (w0 + 1 < in_w) ? (w0 + 1) : w0; - - const double td = src_d - static_cast(d0); - const double th = src_h - static_cast(h0); - const double tw = src_w - static_cast(w0); - - const size_t base = n * in_meta.strides[0] + c * in_meta.strides[1]; - auto load = [&](size_t d, size_t h, size_t w) -> Tcompute { - return to_compute(input[base + d * in_meta.strides[2] + h * in_meta.strides[3] + w * in_meta.strides[4]]); - }; - - const Tcompute v000 = load(d0, h0, w0); - const Tcompute v001 = load(d0, h0, w1); - const Tcompute v010 = load(d0, h1, w0); - const Tcompute v011 = load(d0, h1, w1); - const Tcompute v100 = load(d1, h0, w0); - const Tcompute v101 = load(d1, h0, w1); - const Tcompute v110 = load(d1, h1, w0); - const Tcompute v111 = load(d1, h1, w1); - - const double w000 = (1.0 - td) * (1.0 - th) * (1.0 - tw); - const double w001 = (1.0 - td) * (1.0 - th) * tw; - const double w010 = (1.0 - td) * th * (1.0 - tw); - const double w011 = (1.0 - td) * th * tw; - const double w100 = td * (1.0 - th) * (1.0 - tw); - const double w101 = td * (1.0 - th) * tw; - const double w110 = td * th * (1.0 - tw); - const double w111 = td * th * tw; - - const Tcompute out = static_cast( - w000 * static_cast(v000) + w001 * static_cast(v001) + w010 * static_cast(v010) + w011 * static_cast(v011) + w100 * static_cast(v100) + w101 * static_cast(v101) + w110 * static_cast(v110) + w111 * static_cast(v111)); - - const size_t out_off = n * out_meta.strides[0] + c * out_meta.strides[1] + od * out_meta.strides[2] + oh * out_meta.strides[3] + ow * out_meta.strides[4]; - output[out_off] = from_compute(out, TypeTag{}); -} - -template -__global__ void area_2d_kernel( - T *output, - const T *input, - TensorMeta out_meta, - TensorMeta in_meta) { - - const size_t total = out_meta.shape[0] * out_meta.shape[1] * out_meta.shape[2] * out_meta.shape[3]; - const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= total) { - return; - } - - const size_t out_w = out_meta.shape[3]; - const size_t out_h = out_meta.shape[2]; - const size_t c_stride = out_h * out_w; - const size_t n_stride = out_meta.shape[1] * c_stride; - - const size_t n = idx / n_stride; - size_t rem = idx - n * n_stride; - const size_t c = rem / c_stride; - rem -= c * c_stride; - const size_t oh = rem / out_w; - const size_t ow = rem - oh * out_w; - - const size_t in_h = in_meta.shape[2]; - const size_t in_w = in_meta.shape[3]; - - const double scale_h = static_cast(in_h) / static_cast(out_h); - const double scale_w = static_cast(in_w) / static_cast(out_w); - - const double h0 = static_cast(oh) * scale_h; - const double h1 = static_cast(oh + 1) * scale_h; - const double w0 = static_cast(ow) * scale_w; - const double w1 = static_cast(ow + 1) * scale_w; - - const size_t ih0 = static_cast(floor(h0)); - const size_t ih1 = static_cast(ceil(h1)); - const size_t iw0 = static_cast(floor(w0)); - const size_t iw1 = static_cast(ceil(w1)); - - const size_t base = n * in_meta.strides[0] + c * in_meta.strides[1]; - - double accum = 0.0; - double area = (h1 - h0) * (w1 - w0); - if (area <= 0.0) { - area = 1.0; - } - - for (size_t ih = ih0; ih < ih1 && ih < in_h; ++ih) { - const double wh = fmax(0.0, fmin(h1, static_cast(ih + 1)) - fmax(h0, static_cast(ih))); - for (size_t iw = iw0; iw < iw1 && iw < in_w; ++iw) { - const double ww = fmax(0.0, fmin(w1, static_cast(iw + 1)) - fmax(w0, static_cast(iw))); - const double w = wh * ww; - const Tcompute v = to_compute(input[base + ih * in_meta.strides[2] + iw * in_meta.strides[3]]); - accum += w * static_cast(v); - } - } - - const Tcompute out = static_cast(accum / area); - const size_t out_off = n * out_meta.strides[0] + c * out_meta.strides[1] + oh * out_meta.strides[2] + ow * out_meta.strides[3]; - output[out_off] = from_compute(out, TypeTag{}); -} - -} // namespace op::interpolate::cuda - -// --------------------------------------------------------------------------- -// Compatibility wrappers for backends that still reference `op::cuda::*`. -// These assume contiguous NCHW layout. -// --------------------------------------------------------------------------- - -namespace op::cuda { - -template -__global__ void interpolate_bilinear_2d_kernel( - T *output, - const T *input, - size_t batch, - size_t channels, - size_t in_h, - size_t in_w, - size_t out_h, - size_t out_w, - double scale_h, - double scale_w, - int align_corners) { - - const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - const size_t total = batch * channels * out_h * out_w; - if (idx >= total) { - return; - } - - const size_t n = idx / (channels * out_h * out_w); - size_t rem = idx - n * (channels * out_h * out_w); - const size_t c = rem / (out_h * out_w); - rem -= c * (out_h * out_w); - const size_t oh = rem / out_w; - const size_t ow = rem - oh * out_w; - - double src_y = align_corners ? static_cast(oh) * scale_h : (static_cast(oh) + 0.5) * scale_h - 0.5; - double src_x = align_corners ? static_cast(ow) * scale_w : (static_cast(ow) + 0.5) * scale_w - 0.5; - - src_y = fmax(0.0, fmin(src_y, static_cast(in_h - 1))); - src_x = fmax(0.0, fmin(src_x, static_cast(in_w - 1))); - - const size_t y0 = static_cast(floor(src_y)); - const size_t x0 = static_cast(floor(src_x)); - const size_t y1 = (y0 + 1 < in_h) ? (y0 + 1) : y0; - const size_t x1 = (x0 + 1 < in_w) ? (x0 + 1) : x0; - const double wy = src_y - static_cast(y0); - const double wx = src_x - static_cast(x0); - - using Tcompute = std::conditional_t, double, float>; - - const size_t base = ((n * channels + c) * in_h) * in_w; - const Tcompute v00 = op::interpolate::cuda::to_compute(input[base + y0 * in_w + x0]); - const Tcompute v01 = op::interpolate::cuda::to_compute(input[base + y0 * in_w + x1]); - const Tcompute v10 = op::interpolate::cuda::to_compute(input[base + y1 * in_w + x0]); - const Tcompute v11 = op::interpolate::cuda::to_compute(input[base + y1 * in_w + x1]); - - const double w00 = (1.0 - wy) * (1.0 - wx); - const double w01 = (1.0 - wy) * wx; - const double w10 = wy * (1.0 - wx); - const double w11 = wy * wx; - - const Tcompute out = static_cast( - w00 * static_cast(v00) + w01 * static_cast(v01) + w10 * static_cast(v10) + w11 * static_cast(v11)); - - output[((n * channels + c) * out_h + oh) * out_w + ow] = op::interpolate::cuda::from_compute(out, op::interpolate::cuda::TypeTag{}); -} - -} // namespace op::cuda diff --git a/src/infiniop/ops/interpolate/metax/interpolate_metax.h b/src/infiniop/ops/interpolate/metax/interpolate_metax.h deleted file mode 100644 index 083901f0b..000000000 --- a/src/infiniop/ops/interpolate/metax/interpolate_metax.h +++ /dev/null @@ -1,71 +0,0 @@ -#ifndef __INTERPOLATE_METAX_H__ -#define __INTERPOLATE_METAX_H__ - -#include "../../../operator.h" -#include - -namespace op::interpolate::metax { - -enum class InterpolateMode { - NEAREST, - LINEAR, - BILINEAR, - TRILINEAR, - AREA -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t ndim; - std::vector input_shape; - std::vector output_shape; - std::vector input_strides; - std::vector output_strides; - InterpolateMode mode; - int align_corners; - size_t input_size; - size_t output_size; - - Descriptor(infiniDtype_t dtype, size_t ndim, - std::vector input_shape, std::vector output_shape, - std::vector input_strides, std::vector output_strides, - InterpolateMode mode, int align_corners, - size_t input_size, size_t output_size, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - ndim(ndim), - input_shape(std::move(input_shape)), - output_shape(std::move(output_shape)), - input_strides(std::move(input_strides)), - output_strides(std::move(output_strides)), - mode(mode), - align_corners(align_corners), - input_size(input_size), - output_size(output_size) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - const char *mode, - void *size, - void *scale_factor, - int align_corners); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; -} // namespace op::interpolate::metax - -#endif // __INTERPOLATE_METAX_H__ diff --git a/src/infiniop/ops/interpolate/metax/interpolate_metax.maca b/src/infiniop/ops/interpolate/metax/interpolate_metax.maca deleted file mode 100644 index c82c5a45f..000000000 --- a/src/infiniop/ops/interpolate/metax/interpolate_metax.maca +++ /dev/null @@ -1,253 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "interpolate_metax.h" -#include - -namespace op::interpolate::metax { - -static bool try_parse_mode(const char *mode_str, InterpolateMode &mode) { - if (std::strcmp(mode_str, "nearest") == 0) { - mode = InterpolateMode::NEAREST; - return true; - } else if (std::strcmp(mode_str, "linear") == 0) { - mode = InterpolateMode::LINEAR; - return true; - } else if (std::strcmp(mode_str, "bilinear") == 0) { - mode = InterpolateMode::BILINEAR; - return true; - } else if (std::strcmp(mode_str, "trilinear") == 0) { - mode = InterpolateMode::TRILINEAR; - return true; - } else if (std::strcmp(mode_str, "area") == 0) { - mode = InterpolateMode::AREA; - return true; - } - return false; -} - -Descriptor::~Descriptor() = default; - -static bool build_meta( - op::interpolate::cuda::TensorMeta &meta, - const std::vector &shape, - const std::vector &strides) { - - const size_t ndim = shape.size(); - if (ndim > static_cast(op::interpolate::cuda::kInterpolateMaxDims)) { - return false; - } - if (strides.size() != ndim) { - return false; - } - - meta.ndim = static_cast(ndim); - for (size_t i = 0; i < static_cast(op::interpolate::cuda::kInterpolateMaxDims); ++i) { - meta.shape[i] = (i < ndim) ? shape[i] : 1; - meta.strides[i] = (i < ndim) ? strides[i] : 0; - } - return true; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - const char *mode, - void *size, - void *scale_factor, - int align_corners) { - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x_shape.size() < 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if ((size != nullptr) == (scale_factor != nullptr)) { - return INFINI_STATUS_BAD_PARAM; - } - if (y_shape.size() != x_shape.size() || y_shape[0] != x_shape[0] || y_shape[1] != x_shape[1]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t ndim = x_shape.size() - 2; - - if (scale_factor != nullptr) { - const double *scale_array = reinterpret_cast(scale_factor); - const double scale = scale_array[0]; - std::vector expected_y_shape = x_shape; - for (size_t i = 0; i < ndim; ++i) { - expected_y_shape[i + 2] = static_cast(static_cast(x_shape[i + 2]) * scale); - } - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - if (y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - InterpolateMode parsed_mode{}; - if (!try_parse_mode(mode, parsed_mode)) { - return INFINI_STATUS_BAD_PARAM; - } - - *desc_ptr = new Descriptor(dtype, ndim, x_shape, y_shape, x_desc->strides(), y_desc->strides(), parsed_mode, align_corners, - x_desc->numel(), y_desc->numel(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - constexpr int BLOCK_SIZE = 256; - const int num_blocks = static_cast((output_size + BLOCK_SIZE - 1) / BLOCK_SIZE); - - op::interpolate::cuda::TensorMeta in_meta{}; - op::interpolate::cuda::TensorMeta out_meta{}; - if (!build_meta(in_meta, input_shape, input_strides) || !build_meta(out_meta, output_shape, output_strides)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (mode == InterpolateMode::NEAREST) { - if (ndim != 2) { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::nearest_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::nearest_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::nearest_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::nearest_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (mode == InterpolateMode::LINEAR) { - if (ndim != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::linear_1d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::linear_1d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::linear_1d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::linear_1d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (mode == InterpolateMode::BILINEAR) { - if (ndim != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::bilinear_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::bilinear_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::bilinear_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::bilinear_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (mode == InterpolateMode::TRILINEAR) { - if (ndim != 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::trilinear_3d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::trilinear_3d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::trilinear_3d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::trilinear_3d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (mode == InterpolateMode::AREA) { - if (ndim != 2) { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::area_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::area_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::area_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::area_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::interpolate::metax diff --git a/src/infiniop/ops/interpolate/moore/interpolate_moore.h b/src/infiniop/ops/interpolate/moore/interpolate_moore.h deleted file mode 100644 index 509ac04b0..000000000 --- a/src/infiniop/ops/interpolate/moore/interpolate_moore.h +++ /dev/null @@ -1,71 +0,0 @@ -#ifndef __INTERPOLATE_MOORE_H__ -#define __INTERPOLATE_MOORE_H__ - -#include "../../../operator.h" -#include - -namespace op::interpolate::moore { - -enum class InterpolateMode { - NEAREST, - LINEAR, - BILINEAR, - TRILINEAR, - AREA -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t ndim; - std::vector input_shape; - std::vector output_shape; - std::vector input_strides; - std::vector output_strides; - InterpolateMode mode; - int align_corners; - size_t input_size; - size_t output_size; - - Descriptor(infiniDtype_t dtype, size_t ndim, - std::vector input_shape, std::vector output_shape, - std::vector input_strides, std::vector output_strides, - InterpolateMode mode, int align_corners, - size_t input_size, size_t output_size, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - ndim(ndim), - input_shape(std::move(input_shape)), - output_shape(std::move(output_shape)), - input_strides(std::move(input_strides)), - output_strides(std::move(output_strides)), - mode(mode), - align_corners(align_corners), - input_size(input_size), - output_size(output_size) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - const char *mode, - void *size, - void *scale_factor, - int align_corners); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; -} // namespace op::interpolate::moore - -#endif // __INTERPOLATE_MOORE_H__ diff --git a/src/infiniop/ops/interpolate/moore/interpolate_moore.mu b/src/infiniop/ops/interpolate/moore/interpolate_moore.mu deleted file mode 100644 index 75dc05bc0..000000000 --- a/src/infiniop/ops/interpolate/moore/interpolate_moore.mu +++ /dev/null @@ -1,253 +0,0 @@ -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "interpolate_moore.h" -#include - -namespace op::interpolate::moore { - -static bool try_parse_mode(const char *mode_str, InterpolateMode &mode) { - if (std::strcmp(mode_str, "nearest") == 0) { - mode = InterpolateMode::NEAREST; - return true; - } else if (std::strcmp(mode_str, "linear") == 0) { - mode = InterpolateMode::LINEAR; - return true; - } else if (std::strcmp(mode_str, "bilinear") == 0) { - mode = InterpolateMode::BILINEAR; - return true; - } else if (std::strcmp(mode_str, "trilinear") == 0) { - mode = InterpolateMode::TRILINEAR; - return true; - } else if (std::strcmp(mode_str, "area") == 0) { - mode = InterpolateMode::AREA; - return true; - } - return false; -} - -Descriptor::~Descriptor() = default; - -static bool build_meta( - op::interpolate::cuda::TensorMeta &meta, - const std::vector &shape, - const std::vector &strides) { - - const size_t ndim = shape.size(); - if (ndim > static_cast(op::interpolate::cuda::kInterpolateMaxDims)) { - return false; - } - if (strides.size() != ndim) { - return false; - } - - meta.ndim = static_cast(ndim); - for (size_t i = 0; i < static_cast(op::interpolate::cuda::kInterpolateMaxDims); ++i) { - meta.shape[i] = (i < ndim) ? shape[i] : 1; - meta.strides[i] = (i < ndim) ? strides[i] : 0; - } - return true; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - const char *mode, - void *size, - void *scale_factor, - int align_corners) { - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x_shape.size() < 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if ((size != nullptr) == (scale_factor != nullptr)) { - return INFINI_STATUS_BAD_PARAM; - } - if (y_shape.size() != x_shape.size() || y_shape[0] != x_shape[0] || y_shape[1] != x_shape[1]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t ndim = x_shape.size() - 2; - - if (scale_factor != nullptr) { - const double *scale_array = reinterpret_cast(scale_factor); - const double scale = scale_array[0]; - std::vector expected_y_shape = x_shape; - for (size_t i = 0; i < ndim; ++i) { - expected_y_shape[i + 2] = static_cast(static_cast(x_shape[i + 2]) * scale); - } - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - if (y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - InterpolateMode parsed_mode{}; - if (!try_parse_mode(mode, parsed_mode)) { - return INFINI_STATUS_BAD_PARAM; - } - - *desc_ptr = new Descriptor(dtype, ndim, x_shape, y_shape, x_desc->strides(), y_desc->strides(), parsed_mode, align_corners, - x_desc->numel(), y_desc->numel(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - constexpr int BLOCK_SIZE = 256; - const int num_blocks = static_cast((output_size + BLOCK_SIZE - 1) / BLOCK_SIZE); - - op::interpolate::cuda::TensorMeta in_meta{}; - op::interpolate::cuda::TensorMeta out_meta{}; - if (!build_meta(in_meta, input_shape, input_strides) || !build_meta(out_meta, output_shape, output_strides)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (mode == InterpolateMode::NEAREST) { - if (ndim != 2) { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::nearest_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::nearest_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::nearest_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::nearest_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (mode == InterpolateMode::LINEAR) { - if (ndim != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::linear_1d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::linear_1d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::linear_1d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::linear_1d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (mode == InterpolateMode::BILINEAR) { - if (ndim != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::bilinear_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::bilinear_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::bilinear_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::bilinear_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (mode == InterpolateMode::TRILINEAR) { - if (ndim != 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::trilinear_3d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::trilinear_3d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::trilinear_3d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::trilinear_3d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (mode == InterpolateMode::AREA) { - if (ndim != 2) { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::area_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::area_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::area_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::area_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::interpolate::moore diff --git a/src/infiniop/ops/interpolate/nvidia/interpolate_nvidia.cu b/src/infiniop/ops/interpolate/nvidia/interpolate_nvidia.cu deleted file mode 100644 index 22c99cf53..000000000 --- a/src/infiniop/ops/interpolate/nvidia/interpolate_nvidia.cu +++ /dev/null @@ -1,253 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "interpolate_nvidia.cuh" -#include - -namespace op::interpolate::nvidia { - -static bool try_parse_mode(const char *mode_str, InterpolateMode &mode) { - if (std::strcmp(mode_str, "nearest") == 0) { - mode = InterpolateMode::NEAREST; - return true; - } else if (std::strcmp(mode_str, "linear") == 0) { - mode = InterpolateMode::LINEAR; - return true; - } else if (std::strcmp(mode_str, "bilinear") == 0) { - mode = InterpolateMode::BILINEAR; - return true; - } else if (std::strcmp(mode_str, "trilinear") == 0) { - mode = InterpolateMode::TRILINEAR; - return true; - } else if (std::strcmp(mode_str, "area") == 0) { - mode = InterpolateMode::AREA; - return true; - } - return false; -} - -Descriptor::~Descriptor() = default; - -static bool build_meta( - op::interpolate::cuda::TensorMeta &meta, - const std::vector &shape, - const std::vector &strides) { - - const size_t ndim = shape.size(); - if (ndim > static_cast(op::interpolate::cuda::kInterpolateMaxDims)) { - return false; - } - if (strides.size() != ndim) { - return false; - } - - meta.ndim = static_cast(ndim); - for (size_t i = 0; i < static_cast(op::interpolate::cuda::kInterpolateMaxDims); ++i) { - meta.shape[i] = (i < ndim) ? shape[i] : 1; - meta.strides[i] = (i < ndim) ? strides[i] : 0; - } - return true; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - const char *mode, - void *size, - void *scale_factor, - int align_corners) { - - auto dtype = x_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - - if (x_shape.size() < 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if ((size != nullptr) == (scale_factor != nullptr)) { - return INFINI_STATUS_BAD_PARAM; - } - if (y_shape.size() != x_shape.size() || y_shape[0] != x_shape[0] || y_shape[1] != x_shape[1]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t ndim = x_shape.size() - 2; - - if (scale_factor != nullptr) { - const double *scale_array = reinterpret_cast(scale_factor); - const double scale = scale_array[0]; - std::vector expected_y_shape = x_shape; - for (size_t i = 0; i < ndim; ++i) { - expected_y_shape[i + 2] = static_cast(static_cast(x_shape[i + 2]) * scale); - } - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } - - if (y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - InterpolateMode parsed_mode{}; - if (!try_parse_mode(mode, parsed_mode)) { - return INFINI_STATUS_BAD_PARAM; - } - - *desc_ptr = new Descriptor(dtype, ndim, x_shape, y_shape, x_desc->strides(), y_desc->strides(), parsed_mode, align_corners, - x_desc->numel(), y_desc->numel(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - auto cuda_stream = reinterpret_cast(stream); - constexpr int BLOCK_SIZE = 256; - const int num_blocks = static_cast((output_size + BLOCK_SIZE - 1) / BLOCK_SIZE); - - op::interpolate::cuda::TensorMeta in_meta{}; - op::interpolate::cuda::TensorMeta out_meta{}; - if (!build_meta(in_meta, input_shape, input_strides) || !build_meta(out_meta, output_shape, output_strides)) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - if (mode == InterpolateMode::NEAREST) { - if (ndim != 2) { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::nearest_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::nearest_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::nearest_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::nearest_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (mode == InterpolateMode::LINEAR) { - if (ndim != 1) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::linear_1d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::linear_1d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::linear_1d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::linear_1d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (mode == InterpolateMode::BILINEAR) { - if (ndim != 2) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::bilinear_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::bilinear_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::bilinear_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::bilinear_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (mode == InterpolateMode::TRILINEAR) { - if (ndim != 3) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::trilinear_3d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::trilinear_3d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::trilinear_3d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::trilinear_3d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta, align_corners); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (mode == InterpolateMode::AREA) { - if (ndim != 2) { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - switch (_dtype) { - case INFINI_DTYPE_F16: - op::interpolate::cuda::area_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_BF16: - op::interpolate::cuda::area_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_F32: - op::interpolate::cuda::area_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - case INFINI_DTYPE_F64: - op::interpolate::cuda::area_2d_kernel<<>>( - reinterpret_cast(y), reinterpret_cast(x), out_meta, in_meta); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::interpolate::nvidia diff --git a/src/infiniop/ops/interpolate/nvidia/interpolate_nvidia.cuh b/src/infiniop/ops/interpolate/nvidia/interpolate_nvidia.cuh deleted file mode 100644 index c56eb1d69..000000000 --- a/src/infiniop/ops/interpolate/nvidia/interpolate_nvidia.cuh +++ /dev/null @@ -1,72 +0,0 @@ -#ifndef __INTERPOLATE_NVIDIA_H__ -#define __INTERPOLATE_NVIDIA_H__ - -#include "../../../operator.h" -#include - -namespace op::interpolate::nvidia { - -enum class InterpolateMode { - NEAREST, - LINEAR, - BILINEAR, - TRILINEAR, - AREA -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t ndim; - std::vector input_shape; - std::vector output_shape; - std::vector input_strides; - std::vector output_strides; - InterpolateMode mode; - int align_corners; - size_t input_size; - size_t output_size; - - Descriptor(infiniDtype_t dtype, size_t ndim, - std::vector input_shape, std::vector output_shape, - std::vector input_strides, std::vector output_strides, - InterpolateMode mode, int align_corners, - size_t input_size, size_t output_size, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - ndim(ndim), - input_shape(std::move(input_shape)), - output_shape(std::move(output_shape)), - input_strides(std::move(input_strides)), - output_strides(std::move(output_strides)), - mode(mode), - align_corners(align_corners), - input_size(input_size), - output_size(output_size) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - const char *mode, - void *size, - void *scale_factor, - int align_corners); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) const; -}; - -} // namespace op::interpolate::nvidia - -#endif // __INTERPOLATE_NVIDIA_H__ diff --git a/src/infiniop/ops/interpolate/operator.cc b/src/infiniop/ops/interpolate/operator.cc deleted file mode 100644 index 0197f8cbe..000000000 --- a/src/infiniop/ops/interpolate/operator.cc +++ /dev/null @@ -1,177 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/interpolate.h" - -#ifdef ENABLE_CPU_API -#include "cpu/interpolate_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/interpolate_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/interpolate_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/interpolate_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateInterpolateDescriptor( - infiniopHandle_t handle, - infiniopInterpolateDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc, - const char *mode, - void *size, - void *scale_factor, - int align_corners) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::interpolate::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - x_desc, \ - mode, \ - size, \ - scale_factor, \ - align_corners) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetInterpolateWorkspaceSize(infiniopInterpolateDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopInterpolate( - infiniopInterpolateDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *x, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, x, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyInterpolateDescriptor(infiniopInterpolateDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/kron/cpu/kron_cpu.cc b/src/infiniop/ops/kron/cpu/kron_cpu.cc deleted file mode 100644 index 1ed443f6a..000000000 --- a/src/infiniop/ops/kron/cpu/kron_cpu.cc +++ /dev/null @@ -1,152 +0,0 @@ -#include "kron_cpu.h" -#include "../../../../utils.h" -#include "../../../tensor.h" -#include - -namespace op::kron::cpu { - -utils::Result KronInfo::create( - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t y_desc) { - - auto a_shape = a_desc->shape(); - auto b_shape = b_desc->shape(); - auto y_shape = y_desc->shape(); - - // Kron requires same number of dimensions - if (a_shape.size() != b_shape.size()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t ndim = a_shape.size(); - - // Output shape: each dimension is a_shape[i] * b_shape[i] - std::vector expected_y_shape(ndim); - for (size_t i = 0; i < ndim; ++i) { - expected_y_shape[i] = a_shape[i] * b_shape[i]; - } - - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - KronInfo info; - info.ndim = ndim; - info.a_shape = a_shape; - info.b_shape = b_shape; - info.y_shape = y_shape; - info.a_strides = a_desc->strides(); - info.b_strides = b_desc->strides(); - info.y_strides = y_desc->strides(); - info.a_size = a_desc->numel(); - info.b_size = b_desc->numel(); - info.y_size = y_desc->numel(); - - return utils::Result(std::move(info)); -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - - auto dtype = a_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - if (b_desc->dtype() != dtype || y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto info_result = KronInfo::create(a_desc, b_desc, y_desc); - CHECK_RESULT(info_result); - - *desc_ptr = new Descriptor(dtype, info_result.take(), handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -void kron_impl( - const KronInfo &info, - T *y, - const T *a, - const T *b) { - - // Kronecker product: for each element a[i] in A, multiply it with entire B - // y[output_idx] = a[a_idx] * b[b_idx] - // where output coordinates are computed from a and b coordinates - - // Iterate over output tensor - std::vector y_coords(info.ndim, 0); - for (size_t y_idx = 0; y_idx < info.y_size; ++y_idx) { - // Convert linear index to coordinates - size_t temp = y_idx; - for (size_t d = info.ndim; d-- > 0;) { - y_coords[d] = temp % info.y_shape[d]; - temp /= info.y_shape[d]; - } - - // Compute corresponding a and b coordinates, then use descriptor strides - // to compute element offsets (supports non-contiguous / broadcasted strides). - ptrdiff_t a_off = 0; - ptrdiff_t b_off = 0; - ptrdiff_t out_off = 0; - for (size_t d = 0; d < info.ndim; ++d) { - size_t a_coord = y_coords[d] / info.b_shape[d]; - size_t b_coord = y_coords[d] % info.b_shape[d]; - a_off += static_cast(a_coord) * info.a_strides[d]; - b_off += static_cast(b_coord) * info.b_strides[d]; - out_off += static_cast(y_coords[d]) * info.y_strides[d]; - } - - if constexpr (std::is_same_v || std::is_same_v) { - float av = utils::cast(a[a_off]); - float bv = utils::cast(b[b_off]); - y[out_off] = utils::cast(av * bv); - } else { - y[out_off] = a[a_off] * b[b_off]; - } - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const { - - switch (_dtype) { - case INFINI_DTYPE_F16: - kron_impl(_info, reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b)); - break; - case INFINI_DTYPE_BF16: - kron_impl(_info, reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b)); - break; - case INFINI_DTYPE_F32: - kron_impl(_info, reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b)); - break; - case INFINI_DTYPE_F64: - kron_impl(_info, reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b)); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::kron::cpu diff --git a/src/infiniop/ops/kron/cpu/kron_cpu.h b/src/infiniop/ops/kron/cpu/kron_cpu.h deleted file mode 100644 index 95693b57f..000000000 --- a/src/infiniop/ops/kron/cpu/kron_cpu.h +++ /dev/null @@ -1,61 +0,0 @@ -#ifndef __KRON_CPU_H__ -#define __KRON_CPU_H__ - -#include "../../../devices/cpu/common_cpu.h" -#include "../../../operator.h" -#include - -namespace op::kron::cpu { - -struct KronInfo { - size_t ndim; - std::vector a_shape; - std::vector b_shape; - std::vector y_shape; - std::vector a_strides; - std::vector b_strides; - std::vector y_strides; - size_t a_size; - size_t b_size; - size_t y_size; - - static utils::Result create( - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc, - infiniopTensorDescriptor_t y_desc); -}; - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - KronInfo _info; - - Descriptor(infiniDtype_t dtype, KronInfo info, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - _info(std::move(info)) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc); - - size_t workspaceSize() const { return 0; } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const; -}; - -} // namespace op::kron::cpu - -#endif // __KRON_CPU_H__ diff --git a/src/infiniop/ops/kron/cuda/kernel.cuh b/src/infiniop/ops/kron/cuda/kernel.cuh deleted file mode 100644 index c79fafe89..000000000 --- a/src/infiniop/ops/kron/cuda/kernel.cuh +++ /dev/null @@ -1,50 +0,0 @@ -#include -#include - -namespace op::cuda { - -template -__global__ void kron_kernel( - T *output, - const T *a, - const T *b, - size_t total_output, - size_t ndim, - const size_t *a_shape, - const size_t *b_shape, - const size_t *y_shape, - const ptrdiff_t *a_strides, - const ptrdiff_t *b_strides, - const ptrdiff_t *y_strides) { - - size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx >= total_output) { - return; - } - - // Convert linear index to coordinates - size_t temp = idx; - constexpr size_t kMaxNdim = 8; - size_t y_coords[kMaxNdim]; // Max 8 dimensions - for (size_t d = ndim; d-- > 0;) { - y_coords[d] = temp % y_shape[d]; - temp /= y_shape[d]; - } - - ptrdiff_t a_offset = 0; - ptrdiff_t b_offset = 0; - ptrdiff_t y_offset = 0; - for (size_t d = 0; d < ndim; ++d) { - const auto y_coord = y_coords[d]; - const auto a_coord = y_coord / b_shape[d]; - const auto b_coord = y_coord % b_shape[d]; - - a_offset += static_cast(a_coord) * a_strides[d]; - b_offset += static_cast(b_coord) * b_strides[d]; - y_offset += static_cast(y_coord) * y_strides[d]; - } - - output[y_offset] = a[a_offset] * b[b_offset]; -} - -} // namespace op::cuda diff --git a/src/infiniop/ops/kron/metax/kron_metax.h b/src/infiniop/ops/kron/metax/kron_metax.h deleted file mode 100644 index 470814cb0..000000000 --- a/src/infiniop/ops/kron/metax/kron_metax.h +++ /dev/null @@ -1,68 +0,0 @@ -#ifndef __KRON_METAX_H__ -#define __KRON_METAX_H__ - -#include "../../../operator.h" -#include -#include - -namespace op::kron::metax { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t ndim; - std::vector a_shape; - std::vector b_shape; - std::vector y_shape; - std::vector a_strides; - std::vector b_strides; - std::vector y_strides; - size_t a_size; - size_t b_size; - size_t y_size; - - Descriptor(infiniDtype_t dtype, size_t ndim, - std::vector a_shape, std::vector b_shape, std::vector y_shape, - std::vector a_strides, - std::vector b_strides, - std::vector y_strides, - size_t a_size, size_t b_size, size_t y_size, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - ndim(ndim), - a_shape(std::move(a_shape)), - b_shape(std::move(b_shape)), - y_shape(std::move(y_shape)), - a_strides(std::move(a_strides)), - b_strides(std::move(b_strides)), - y_strides(std::move(y_strides)), - a_size(a_size), - b_size(b_size), - y_size(y_size) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc); - - size_t workspaceSize() const { - return (ndim * 3) * sizeof(size_t) + (ndim * 3) * sizeof(ptrdiff_t); - } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const; -}; - -} // namespace op::kron::metax - -#endif // __KRON_METAX_H__ diff --git a/src/infiniop/ops/kron/metax/kron_metax.maca b/src/infiniop/ops/kron/metax/kron_metax.maca deleted file mode 100644 index 27be951f8..000000000 --- a/src/infiniop/ops/kron/metax/kron_metax.maca +++ /dev/null @@ -1,148 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_kernel_common.h" -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "kron_metax.h" - -namespace op::kron::metax { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - - auto dtype = a_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - if (b_desc->dtype() != dtype || y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto a_shape = a_desc->shape(); - auto b_shape = b_desc->shape(); - auto y_shape = y_desc->shape(); - - if (a_shape.size() != b_shape.size()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t ndim = a_shape.size(); - if (ndim > 8) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - std::vector expected_y_shape(ndim); - for (size_t i = 0; i < ndim; ++i) { - expected_y_shape[i] = a_shape[i] * b_shape[i]; - } - - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - auto a_strides = a_desc->strides(); - auto b_strides = b_desc->strides(); - auto y_strides = y_desc->strides(); - if (a_strides.size() != ndim || b_strides.size() != ndim || y_strides.size() != ndim) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - *desc_ptr = new Descriptor(dtype, ndim, a_shape, b_shape, y_shape, - std::move(a_strides), std::move(b_strides), std::move(y_strides), - a_desc->numel(), b_desc->numel(), y_desc->numel(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const { - - if (workspace_size < this->workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - if (y_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - auto hc_stream = reinterpret_cast(stream); - size_t *a_shape_d = nullptr; - size_t *b_shape_d = nullptr; - size_t *y_shape_d = nullptr; - ptrdiff_t *a_strides_d = nullptr; - ptrdiff_t *b_strides_d = nullptr; - ptrdiff_t *y_strides_d = nullptr; - if (ndim > 0) { - size_t *shape_data = reinterpret_cast(workspace); - ptrdiff_t *stride_data = reinterpret_cast(shape_data + 3 * ndim); - a_shape_d = shape_data; - b_shape_d = shape_data + ndim; - y_shape_d = shape_data + 2 * ndim; - a_strides_d = stride_data; - b_strides_d = stride_data + ndim; - y_strides_d = stride_data + 2 * ndim; - - CHECK_METAX(hcMemcpyAsync(a_shape_d, a_shape.data(), ndim * sizeof(size_t), hcMemcpyHostToDevice, hc_stream)); - CHECK_METAX(hcMemcpyAsync(b_shape_d, b_shape.data(), ndim * sizeof(size_t), hcMemcpyHostToDevice, hc_stream)); - CHECK_METAX(hcMemcpyAsync(y_shape_d, y_shape.data(), ndim * sizeof(size_t), hcMemcpyHostToDevice, hc_stream)); - CHECK_METAX(hcMemcpyAsync(a_strides_d, a_strides.data(), ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, hc_stream)); - CHECK_METAX(hcMemcpyAsync(b_strides_d, b_strides.data(), ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, hc_stream)); - CHECK_METAX(hcMemcpyAsync(y_strides_d, y_strides.data(), ndim * sizeof(ptrdiff_t), hcMemcpyHostToDevice, hc_stream)); - } - - constexpr int BLOCK_SIZE = 256; - int num_blocks = (y_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - - switch (_dtype) { - case INFINI_DTYPE_F16: - cuda::kron_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - y_size, ndim, - a_shape_d, b_shape_d, y_shape_d, - a_strides_d, b_strides_d, y_strides_d); - break; - case INFINI_DTYPE_BF16: - cuda::kron_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - y_size, ndim, - a_shape_d, b_shape_d, y_shape_d, - a_strides_d, b_strides_d, y_strides_d); - break; - case INFINI_DTYPE_F32: - cuda::kron_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - y_size, ndim, - a_shape_d, b_shape_d, y_shape_d, - a_strides_d, b_strides_d, y_strides_d); - break; - case INFINI_DTYPE_F64: - cuda::kron_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - y_size, ndim, - a_shape_d, b_shape_d, y_shape_d, - a_strides_d, b_strides_d, y_strides_d); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::kron::metax diff --git a/src/infiniop/ops/kron/moore/kron_moore.h b/src/infiniop/ops/kron/moore/kron_moore.h deleted file mode 100644 index 72526640b..000000000 --- a/src/infiniop/ops/kron/moore/kron_moore.h +++ /dev/null @@ -1,68 +0,0 @@ -#ifndef __KRON_MOORE_H__ -#define __KRON_MOORE_H__ - -#include "../../../operator.h" -#include -#include - -namespace op::kron::moore { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t ndim; - std::vector a_shape; - std::vector b_shape; - std::vector y_shape; - std::vector a_strides; - std::vector b_strides; - std::vector y_strides; - size_t a_size; - size_t b_size; - size_t y_size; - - Descriptor(infiniDtype_t dtype, size_t ndim, - std::vector a_shape, std::vector b_shape, std::vector y_shape, - std::vector a_strides, - std::vector b_strides, - std::vector y_strides, - size_t a_size, size_t b_size, size_t y_size, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - ndim(ndim), - a_shape(std::move(a_shape)), - b_shape(std::move(b_shape)), - y_shape(std::move(y_shape)), - a_strides(std::move(a_strides)), - b_strides(std::move(b_strides)), - y_strides(std::move(y_strides)), - a_size(a_size), - b_size(b_size), - y_size(y_size) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc); - - size_t workspaceSize() const { - return (ndim * 3) * sizeof(size_t) + (ndim * 3) * sizeof(ptrdiff_t); - } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const; -}; - -} // namespace op::kron::moore - -#endif // __KRON_MOORE_H__ diff --git a/src/infiniop/ops/kron/moore/kron_moore.mu b/src/infiniop/ops/kron/moore/kron_moore.mu deleted file mode 100644 index 3e4199ce7..000000000 --- a/src/infiniop/ops/kron/moore/kron_moore.mu +++ /dev/null @@ -1,148 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/moore/moore_common.h" -#include "../../../devices/moore/moore_kernel_common.h" -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "kron_moore.h" - -namespace op::kron::moore { - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - - auto dtype = a_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - if (b_desc->dtype() != dtype || y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto a_shape = a_desc->shape(); - auto b_shape = b_desc->shape(); - auto y_shape = y_desc->shape(); - - if (a_shape.size() != b_shape.size()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t ndim = a_shape.size(); - if (ndim > 8) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - std::vector expected_y_shape(ndim); - for (size_t i = 0; i < ndim; ++i) { - expected_y_shape[i] = a_shape[i] * b_shape[i]; - } - - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - auto a_strides = a_desc->strides(); - auto b_strides = b_desc->strides(); - auto y_strides = y_desc->strides(); - if (a_strides.size() != ndim || b_strides.size() != ndim || y_strides.size() != ndim) { - return INFINI_STATUS_BAD_TENSOR_STRIDES; - } - - *desc_ptr = new Descriptor(dtype, ndim, a_shape, b_shape, y_shape, - std::move(a_strides), std::move(b_strides), std::move(y_strides), - a_desc->numel(), b_desc->numel(), y_desc->numel(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const { - - if (workspace_size < this->workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - if (y_size == 0) { - return INFINI_STATUS_SUCCESS; - } - - auto musa_stream = reinterpret_cast(stream); - size_t *a_shape_d = nullptr; - size_t *b_shape_d = nullptr; - size_t *y_shape_d = nullptr; - ptrdiff_t *a_strides_d = nullptr; - ptrdiff_t *b_strides_d = nullptr; - ptrdiff_t *y_strides_d = nullptr; - if (ndim > 0) { - size_t *shape_data = reinterpret_cast(workspace); - ptrdiff_t *stride_data = reinterpret_cast(shape_data + 3 * ndim); - a_shape_d = shape_data; - b_shape_d = shape_data + ndim; - y_shape_d = shape_data + 2 * ndim; - a_strides_d = stride_data; - b_strides_d = stride_data + ndim; - y_strides_d = stride_data + 2 * ndim; - - CHECK_MOORE(musaMemcpyAsync(a_shape_d, a_shape.data(), ndim * sizeof(size_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(b_shape_d, b_shape.data(), ndim * sizeof(size_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(y_shape_d, y_shape.data(), ndim * sizeof(size_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(a_strides_d, a_strides.data(), ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(b_strides_d, b_strides.data(), ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, musa_stream)); - CHECK_MOORE(musaMemcpyAsync(y_strides_d, y_strides.data(), ndim * sizeof(ptrdiff_t), musaMemcpyHostToDevice, musa_stream)); - } - - constexpr int BLOCK_SIZE = 256; - int num_blocks = (y_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - - switch (_dtype) { - case INFINI_DTYPE_F16: - cuda::kron_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - y_size, ndim, - a_shape_d, b_shape_d, y_shape_d, - a_strides_d, b_strides_d, y_strides_d); - break; - case INFINI_DTYPE_BF16: - cuda::kron_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - y_size, ndim, - a_shape_d, b_shape_d, y_shape_d, - a_strides_d, b_strides_d, y_strides_d); - break; - case INFINI_DTYPE_F32: - cuda::kron_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - y_size, ndim, - a_shape_d, b_shape_d, y_shape_d, - a_strides_d, b_strides_d, y_strides_d); - break; - case INFINI_DTYPE_F64: - cuda::kron_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - y_size, ndim, - a_shape_d, b_shape_d, y_shape_d, - a_strides_d, b_strides_d, y_strides_d); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::kron::moore diff --git a/src/infiniop/ops/kron/nvidia/kron_nvidia.cu b/src/infiniop/ops/kron/nvidia/kron_nvidia.cu deleted file mode 100644 index 43986c417..000000000 --- a/src/infiniop/ops/kron/nvidia/kron_nvidia.cu +++ /dev/null @@ -1,150 +0,0 @@ -#include "../../../../utils.h" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" -#include "../../../tensor.h" -#include "../cuda/kernel.cuh" -#include "kron_nvidia.cuh" - -namespace op::kron::nvidia { - -namespace { -constexpr size_t kMaxSupportedNdim = 8; -} // namespace - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - - auto dtype = a_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); - - if (b_desc->dtype() != dtype || y_desc->dtype() != dtype) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - auto a_shape = a_desc->shape(); - auto b_shape = b_desc->shape(); - auto y_shape = y_desc->shape(); - - if (a_shape.size() != b_shape.size()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - size_t ndim = a_shape.size(); - if (ndim > kMaxSupportedNdim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - std::vector expected_y_shape(ndim); - for (size_t i = 0; i < ndim; ++i) { - expected_y_shape[i] = a_shape[i] * b_shape[i]; - } - - if (y_shape != expected_y_shape) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - auto a_strides = a_desc->strides(); - auto b_strides = b_desc->strides(); - auto y_strides = y_desc->strides(); - - *desc_ptr = new Descriptor(dtype, ndim, a_shape, b_shape, y_shape, - a_strides, b_strides, y_strides, - a_desc->numel(), b_desc->numel(), y_desc->numel(), - handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const { - - if (workspace_size < this->workspaceSize()) { - return INFINI_STATUS_INSUFFICIENT_WORKSPACE; - } - - auto cuda_stream = reinterpret_cast(stream); - size_t *a_shape_d = nullptr; - size_t *b_shape_d = nullptr; - size_t *y_shape_d = nullptr; - ptrdiff_t *a_strides_d = nullptr; - ptrdiff_t *b_strides_d = nullptr; - ptrdiff_t *y_strides_d = nullptr; - if (ndim > 0) { - size_t *shape_data = reinterpret_cast(workspace); - ptrdiff_t *stride_data = reinterpret_cast(shape_data + 3 * ndim); - a_shape_d = shape_data; - b_shape_d = shape_data + ndim; - y_shape_d = shape_data + 2 * ndim; - a_strides_d = stride_data; - b_strides_d = stride_data + ndim; - y_strides_d = stride_data + 2 * ndim; - - CHECK_CUDA(cudaMemcpyAsync(a_shape_d, a_shape.data(), ndim * sizeof(size_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(b_shape_d, b_shape.data(), ndim * sizeof(size_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(y_shape_d, y_shape.data(), ndim * sizeof(size_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(a_strides_d, a_strides.data(), ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(b_strides_d, b_strides.data(), ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, cuda_stream)); - CHECK_CUDA(cudaMemcpyAsync(y_strides_d, y_strides.data(), ndim * sizeof(ptrdiff_t), cudaMemcpyHostToDevice, cuda_stream)); - } - - constexpr int BLOCK_SIZE = 256; - int num_blocks = (y_size + BLOCK_SIZE - 1) / BLOCK_SIZE; - if (num_blocks < 1) { - num_blocks = 1; - } - - switch (_dtype) { - case INFINI_DTYPE_F16: - cuda::kron_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - y_size, ndim, - a_shape_d, b_shape_d, y_shape_d, - a_strides_d, b_strides_d, y_strides_d); - break; - case INFINI_DTYPE_BF16: - cuda::kron_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - y_size, ndim, - a_shape_d, b_shape_d, y_shape_d, - a_strides_d, b_strides_d, y_strides_d); - break; - case INFINI_DTYPE_F32: - cuda::kron_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - y_size, ndim, - a_shape_d, b_shape_d, y_shape_d, - a_strides_d, b_strides_d, y_strides_d); - break; - case INFINI_DTYPE_F64: - cuda::kron_kernel<<>>( - reinterpret_cast(y), - reinterpret_cast(a), - reinterpret_cast(b), - y_size, ndim, - a_shape_d, b_shape_d, y_shape_d, - a_strides_d, b_strides_d, y_strides_d); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::kron::nvidia diff --git a/src/infiniop/ops/kron/nvidia/kron_nvidia.cuh b/src/infiniop/ops/kron/nvidia/kron_nvidia.cuh deleted file mode 100644 index 4ed49825f..000000000 --- a/src/infiniop/ops/kron/nvidia/kron_nvidia.cuh +++ /dev/null @@ -1,63 +0,0 @@ -#ifndef __KRON_NVIDIA_H__ -#define __KRON_NVIDIA_H__ - -#include "../../../operator.h" -#include - -namespace op::kron::nvidia { - -class Descriptor final : public InfiniopDescriptor { - infiniDtype_t _dtype; - size_t ndim; - std::vector a_shape; - std::vector b_shape; - std::vector y_shape; - std::vector a_strides; - std::vector b_strides; - std::vector y_strides; - size_t a_size; - size_t b_size; - size_t y_size; - - Descriptor(infiniDtype_t dtype, size_t ndim, - std::vector a_shape, std::vector b_shape, std::vector y_shape, - std::vector a_strides, std::vector b_strides, std::vector y_strides, - size_t a_size, size_t b_size, size_t y_size, - infiniDevice_t device_type, int device_id) - : InfiniopDescriptor{device_type, device_id}, - _dtype(dtype), - ndim(ndim), - a_shape(std::move(a_shape)), - b_shape(std::move(b_shape)), - y_shape(std::move(y_shape)), - a_strides(std::move(a_strides)), - b_strides(std::move(b_strides)), - y_strides(std::move(y_strides)), - a_size(a_size), - b_size(b_size), - y_size(y_size) {} - -public: - ~Descriptor(); - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc); - - size_t workspaceSize() const { return 3 * ndim * sizeof(size_t) + 3 * ndim * sizeof(ptrdiff_t); } - - infiniStatus_t calculate( - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) const; -}; - -} // namespace op::kron::nvidia - -#endif // __KRON_NVIDIA_H__ diff --git a/src/infiniop/ops/kron/operator.cc b/src/infiniop/ops/kron/operator.cc deleted file mode 100644 index f803f434a..000000000 --- a/src/infiniop/ops/kron/operator.cc +++ /dev/null @@ -1,172 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/kron.h" - -#ifdef ENABLE_CPU_API -#include "cpu/kron_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/kron_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -#include "metax/kron_metax.h" -#endif -#ifdef ENABLE_MOORE_API -#include "moore/kron_moore.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateKronDescriptor( - infiniopHandle_t handle, - infiniopKronDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t a_desc, - infiniopTensorDescriptor_t b_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::kron::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - a_desc, \ - b_desc) - - switch (handle->device) { - -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetKronWorkspaceSize(infiniopKronDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia) -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore) -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia) -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET - - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__INFINI_C infiniStatus_t infiniopKron( - infiniopKronDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *y, - const void *a, - const void *b, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, y, a, b, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t -infiniopDestroyKronDescriptor(infiniopKronDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/kthvalue/cpu/kthvalue_cpu.cc b/src/infiniop/ops/kthvalue/cpu/kthvalue_cpu.cc deleted file mode 100644 index 27a8d8104..000000000 --- a/src/infiniop/ops/kthvalue/cpu/kthvalue_cpu.cc +++ /dev/null @@ -1,157 +0,0 @@ -#include "kthvalue_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include -#include -#include -#include // 引入 type_traits 以支持 constexpr 判断 -#include -#include - -#include "../../../../utils/custom_types.h" - -namespace op::kthvalue::cpu { - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - _opaque = nullptr; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t values_desc, - infiniopTensorDescriptor_t indices_desc, - infiniopTensorDescriptor_t input_desc, - int k, - int dim, - int keepdim) { - - auto handle = reinterpret_cast(handle_); - - auto result = KthvalueInfo::create(values_desc, indices_desc, input_desc, k, dim, keepdim); - CHECK_RESULT(result); - - *desc_ptr = new Descriptor( - new Opaque(), - result.take(), - 0, - handle->device, - handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -template -void calculate_cpu_impl( - const KthvalueInfo &info, - void *values, - void *indices, - const void *input) { - - size_t outer_size = info.outer_size(); - size_t dim_size = info.dim_size(); - size_t inner_size = info.inner_size(); - int k = info.k(); // k is 1-based - - auto val_ptr = reinterpret_cast(values); - auto idx_ptr = reinterpret_cast(indices); - auto in_ptr = reinterpret_cast(input); - - size_t total_tasks = outer_size * inner_size; - - // k 在输入中是 1-based，转为 0-based 用于 vector索引 - int k_idx = k - 1; - -#pragma omp parallel for schedule(static) - for (ptrdiff_t task_id = 0; task_id < (ptrdiff_t)total_tasks; ++task_id) { - // 解算当前任务对应的外部索引和内部索引 - size_t o = task_id / inner_size; - size_t i = task_id % inner_size; - - // 计算输入数据的基地址偏移 - // Input layout logic: [outer, dim, inner] - // Offset = o * (dim_size * inner_size) + [0...dim_size-1] * inner_size + i - size_t input_base_offset = o * dim_size * inner_size + i; - size_t stride = inner_size; - - // 使用临时容器存储 (数值, 原始索引) - // 注意：这里在循环内分配内存，由于 dim_size 通常不大，对 CPU 来说尚可接受 - std::vector> row_data; - row_data.reserve(dim_size); - - for (size_t d = 0; d < dim_size; ++d) { - T val = in_ptr[input_base_offset + d * stride]; - row_data.push_back({val, static_cast(d)}); - } - - // 使用 nth_element 找到第 k 小的元素 (O(N) 复杂度) - // 修复: 使用 utils::cast 确保自定义类型(fp16/bf16)可以比较 - std::nth_element( - row_data.begin(), - row_data.begin() + k_idx, - row_data.end(), - [](const std::pair &a, const std::pair &b) { - // 如果是标准算术类型，直接比较；如果是自定义类型，转换为 float 比较 - if constexpr (std::is_arithmetic_v) { - return a.first < b.first; - } else { - return utils::cast(a.first) < utils::cast(b.first); - } - }); - - // 获取结果 - auto result_pair = row_data[k_idx]; - - // 写入输出 - val_ptr[task_id] = result_pair.first; - idx_ptr[task_id] = result_pair.second; - } -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *values, - void *indices, - const void *input, - void *stream) const { - - auto dtype = _info.dtype(); - - switch (dtype) { - case INFINI_DTYPE_F32: - cpu::calculate_cpu_impl(_info, values, indices, input); - break; - case INFINI_DTYPE_F64: - cpu::calculate_cpu_impl(_info, values, indices, input); - break; - case INFINI_DTYPE_F16: - cpu::calculate_cpu_impl(_info, values, indices, input); - break; - case INFINI_DTYPE_BF16: - cpu::calculate_cpu_impl(_info, values, indices, input); - break; - case INFINI_DTYPE_I32: - cpu::calculate_cpu_impl(_info, values, indices, input); - break; - case INFINI_DTYPE_I64: - cpu::calculate_cpu_impl(_info, values, indices, input); - break; - case INFINI_DTYPE_U32: - cpu::calculate_cpu_impl(_info, values, indices, input); - break; - case INFINI_DTYPE_U64: - cpu::calculate_cpu_impl(_info, values, indices, input); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::kthvalue::cpu diff --git a/src/infiniop/ops/kthvalue/cpu/kthvalue_cpu.h b/src/infiniop/ops/kthvalue/cpu/kthvalue_cpu.h deleted file mode 100644 index a1161f8c5..000000000 --- a/src/infiniop/ops/kthvalue/cpu/kthvalue_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __KTHVALUE_CPU_H__ -#define __KTHVALUE_CPU_H__ - -#include "../kthvalue.h" - -DESCRIPTOR(cpu) - -#endif // __KTHVALUE_CPU_H__ diff --git a/src/infiniop/ops/kthvalue/cuda/kernel.cuh b/src/infiniop/ops/kthvalue/cuda/kernel.cuh deleted file mode 100644 index c5965b6e8..000000000 --- a/src/infiniop/ops/kthvalue/cuda/kernel.cuh +++ /dev/null @@ -1,167 +0,0 @@ -#ifndef __KTHVALUE_CUDA_CUH__ -#define __KTHVALUE_CUDA_CUH__ - -#include -#include -#include - -namespace op::kthvalue::cuda { - -// ================================================================== -// 辅助结构: 键值对 (用于排序时携带索引) -// ================================================================== -template -struct alignas(sizeof(int64_t) * 2) KeyValuePair { // 确保对齐 - T val; - int64_t idx; - - __device__ __forceinline__ KeyValuePair() {} - __device__ __forceinline__ KeyValuePair(T v, int64_t i) : val(v), idx(i) {} - - // 获取用于排序的“无穷大”值，用于 Padding - __device__ __forceinline__ static KeyValuePair max_value() { - // 注意：这里需要根据 T 的具体类型返回最大值 - // 简单起见，对于浮点数我们使用 infinity，整数使用 max - // 在实际工程中可能需要针对 half/bf16 的特化 - if constexpr (std::is_floating_point_v) { - return {static_cast(INFINITY), -1}; - } else { - // 简单的回退策略，实际可能需要 std::numeric_limits 的 device 版特化 - // 这里假设 T 支持强制转换 huge value - return {static_cast(1e30), -1}; - } - } -}; - -// 针对 half/bf16 的比较辅助函数 -// 如果系统头文件未重载 < 运算符，可能需要在此处添加 -template -__device__ __forceinline__ bool is_smaller(const T &a, const T &b) { - return a < b; -} - -// ================================================================== -// Bitonic Sort Helpers (Shared Memory) -// ================================================================== -template -__device__ __forceinline__ void compare_and_swap(KeyValuePair &a, KeyValuePair &b, bool dir) { - // dir: true for ascending, false for descending - // 逻辑：如果 (a < b) != dir，说明顺序不对（或者 a > b 且 dir 为 true），则交换 - // 这里的 dir 含义：true 表示还需要保持 a < b - - // 自定义比较：先比值，值相同比索引（保持稳定性可选，这里简化为只比值） - bool smaller = is_smaller(a.val, b.val) || (a.val == b.val && a.idx < b.idx); - - if (smaller != dir) { - KeyValuePair tmp = a; - a = b; - b = tmp; - } -} - -// ================================================================== -// Kernel: 基于 Bitonic Sort 的 KthValue -// ================================================================== -// 假设: -// 1. Grid 处理 Outer * Inner 个 Slice -// 2. 每个 Block 处理 1 个 Slice (Dim 维度) -// 3. Shared Memory 大小为 power_of_2_dim * sizeof(KeyValuePair) -// 4. BlockDim.x 至少为 power_of_2_dim / 2 (用于并行比较) -template -__global__ void kthvalue_kernel( - T *__restrict__ out_values, // [Outer * Inner] (Flat) - int64_t *__restrict__ out_indices, // [Outer * Inner] (Flat) - const T *__restrict__ input, // [Outer, Dim, Inner] - size_t dim_size, - size_t inner_size, - int k, - size_t power_of_2_dim // 扩展到 2 的幂次的大小 -) { - // 动态共享内存 - extern __shared__ char smem[]; - auto s_data = reinterpret_cast *>(smem); - - unsigned int tid = threadIdx.x; - unsigned int bid = blockIdx.x; - - // 1. 计算当前 Slice 的基地址 - // Batch layout logic: flat_id -> (outer, inner) - // 假设 GridDim.x = Outer * Inner - size_t outer_idx = bid / inner_size; - size_t inner_idx = bid % inner_size; - - // Input layout: [outer, dim, inner] - // Base offset = outer * (dim_size * inner_size) + inner_idx - // Stride = inner_size - size_t input_base = outer_idx * dim_size * inner_size + inner_idx; - size_t stride = inner_size; - - // 2. 加载数据到 Shared Memory (处理 Padding) - // 循环加载，以支持 Dim > BlockDim 的情况 (虽然 Bitonic Sort 通常要求 threads >= N/2) - for (unsigned int i = tid; i < power_of_2_dim; i += blockDim.x) { - if (i < dim_size) { - // 读取输入 - T val = input[input_base + i * stride]; - s_data[i] = KeyValuePair(val, static_cast(i)); - } else { - // Padding 最大值，使其排序后位于末尾 - s_data[i] = KeyValuePair::max_value(); - } - } - __syncthreads(); - - // 3. 双调排序 (Bitonic Sort) - // 算法复杂度 O(log^2 N) - for (unsigned int size = 2; size <= power_of_2_dim; size <<= 1) { - // Bitonic Merge - // dir: 升序或降序交替，构造双调序列 - // bool dir = (tid & (size / 2)) == 0; - - // 这里的逻辑稍微复杂，为了简单和稳定，我们使用全升序排序逻辑 - // 标准 Bitonic Sort 代码如下： - - for (unsigned int stride_step = size >> 1; stride_step > 0; stride_step >>= 1) { - - // 确保线程在范围内 - // 我们需要对所有 pairs (i, i+stride) 进行比较 - // 映射逻辑： - // tid 0 处理: (0, stride), (2*stride, 3*stride)... - // 这种映射较复杂，常用如下方式： - // pos = 2*tid - (tid & (stride - 1)) ... 这种是 Butterfly 模式 - - unsigned int pos = 2 * tid - (tid & (stride_step - 1)); - - // 如果 pos + stride_step 在范围内 - if (pos + stride_step < power_of_2_dim) { // 边界检查，虽由 power_of_2_dim 保证 - unsigned int next_pos = pos + stride_step; - - // 计算比较方向 - // 在完整 Bitonic Sort 中，方向取决于 (pos & size) - // 但这里我们仅实现简单的升序 Sort， - // 需要更标准的 Bitonic Merge 网络: - bool direction = ((pos & size) == 0); - - compare_and_swap(s_data[pos], s_data[next_pos], direction); - } - __syncthreads(); - } - } - - // 4. 输出结果 - // 排序后，第 k 小的元素就在索引 k-1 处 (k is 1-based) - if (tid == 0) { - int target_k = k - 1; - // 简单保护 - if (target_k >= 0 && target_k < dim_size) { - out_values[bid] = s_data[target_k].val; - out_indices[bid] = s_data[target_k].idx; - } else { - // Should not happen if validated - // out_values[bid] = ...; - } - } -} - -} // namespace op::kthvalue::cuda - -#endif // __KTHVALUE_CUDA_CUH__ diff --git a/src/infiniop/ops/kthvalue/info.h b/src/infiniop/ops/kthvalue/info.h deleted file mode 100644 index 3ef30c264..000000000 --- a/src/infiniop/ops/kthvalue/info.h +++ /dev/null @@ -1,124 +0,0 @@ -#ifndef __KTHVALUE_INFO_H__ -#define __KTHVALUE_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::kthvalue { - -class KthvalueInfo { - KthvalueInfo() = default; - -public: - int _dtype; - int _indices_dtype; - int _k; - int _dim; - bool _keepdim; - - size_t _dim_size; - size_t _outer_size; - size_t _inner_size; - - int dtype() const { return _dtype; } - int indices_dtype() const { return _indices_dtype; } - int k() const { return _k; } - int dim() const { return _dim; } - bool keepdim() const { return _keepdim; } - size_t dim_size() const { return _dim_size; } - size_t outer_size() const { return _outer_size; } - size_t inner_size() const { return _inner_size; } - - KthvalueInfo(int dtype, int indices_dtype, int k, int dim, bool keepdim, - size_t dim_size, size_t outer_size, size_t inner_size) - : _dtype(dtype), _indices_dtype(indices_dtype), _k(k), _dim(dim), _keepdim(keepdim), - _dim_size(dim_size), _outer_size(outer_size), _inner_size(inner_size) {} - - static utils::Result create( - infiniopTensorDescriptor_t values_desc, - infiniopTensorDescriptor_t indices_desc, - infiniopTensorDescriptor_t input_desc, - int k, - int dim, - int keepdim) { - - int ndim = int(input_desc->ndim()); - - if (dim < 0) { - dim += ndim; - } - if (dim < 0 || dim >= ndim) { - return INFINI_STATUS_BAD_PARAM; - } - - size_t dim_size = input_desc->shape()[dim]; - if (k <= 0 || k > static_cast(dim_size)) { - return INFINI_STATUS_BAD_PARAM; - } - - size_t outer_size = 1; - for (int i = 0; i < dim; ++i) { - outer_size *= input_desc->shape()[i]; - } - - size_t inner_size = 1; - for (int i = dim + 1; i < ndim; ++i) { - inner_size *= input_desc->shape()[i]; - } - - if (values_desc->ndim() != indices_desc->ndim()) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - // 修复 1: 使用 size_t 避免有符号/无符号比较警告 - size_t expected_out_ndim = keepdim ? ndim : ndim - 1; - if (expected_out_ndim == 0) { - expected_out_ndim = 1; - } - - if (values_desc->ndim() != expected_out_ndim) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - - int out_idx = 0; - for (int i = 0; i < ndim; ++i) { - if (keepdim) { - size_t expected_size = (i == dim) ? 1 : input_desc->shape()[i]; - if (values_desc->shape()[i] != expected_size || indices_desc->shape()[i] != expected_size) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - } else { - if (i == dim) { - continue; - } - if (values_desc->shape()[out_idx] != input_desc->shape()[i] || indices_desc->shape()[out_idx] != input_desc->shape()[i]) { - return INFINI_STATUS_BAD_TENSOR_SHAPE; - } - out_idx++; - } - } - - if (values_desc->dtype() != input_desc->dtype()) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - if (indices_desc->dtype() != INFINI_DTYPE_I64 && indices_desc->dtype() != INFINI_DTYPE_I32) { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return utils::Result(KthvalueInfo{ - input_desc->dtype(), - indices_desc->dtype(), - k, - dim, - static_cast(keepdim), - dim_size, - outer_size, - inner_size}); - } -}; - -} // namespace op::kthvalue - -#endif diff --git a/src/infiniop/ops/kthvalue/kthvalue.h b/src/infiniop/ops/kthvalue/kthvalue.h deleted file mode 100644 index 737a99d0e..000000000 --- a/src/infiniop/ops/kthvalue/kthvalue.h +++ /dev/null @@ -1,52 +0,0 @@ -#ifndef __KTHVALUE_H__ -#define __KTHVALUE_H__ - -#include "../../operator.h" -#include "info.h" // 引用对应的 KthvalueInfo 定义 - -// 宏定义：用于生成不同命名空间下的 Descriptor 类 -#define DESCRIPTOR(NAMESPACE) \ - namespace op::kthvalue::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - KthvalueInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - KthvalueInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t values_desc, \ - infiniopTensorDescriptor_t indices_desc, \ - infiniopTensorDescriptor_t input_desc, \ - int k, \ - int dim, \ - int keepdim); \ - \ - infiniStatus_t calculate( \ - void *workspace, \ - size_t workspace_size, \ - void *values, \ - void *indices, \ - const void *input, \ - void *stream) const; \ - }; \ - } - -#endif // __KTHVALUE_H__ diff --git a/src/infiniop/ops/kthvalue/metax/kthvalue_metax.h b/src/infiniop/ops/kthvalue/metax/kthvalue_metax.h deleted file mode 100644 index 90eb4a157..000000000 --- a/src/infiniop/ops/kthvalue/metax/kthvalue_metax.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __KTHVALUE_METAX_H__ -#define __KTHVALUE_METAX_H__ - -#include "../kthvalue.h" - -DESCRIPTOR(metax) - -#endif // __KTHVALUE_METAX_H__ diff --git a/src/infiniop/ops/kthvalue/metax/kthvalue_metax.maca b/src/infiniop/ops/kthvalue/metax/kthvalue_metax.maca deleted file mode 100644 index bfef675a3..000000000 --- a/src/infiniop/ops/kthvalue/metax/kthvalue_metax.maca +++ /dev/null @@ -1,254 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "../../../devices/metax/metax_handle.h" -#include "../../../devices/metax/metax_kernel_common.h" - -#include "kthvalue_metax.h" -#include -#include -#include -#include -#include -#include - -namespace op::kthvalue::metax { - -template -__device__ __forceinline__ bool is_smaller(const T &a, const T &b) { - return a < b; -} - -__device__ __forceinline__ bool is_smaller(const cuda_bfloat16 &a, const cuda_bfloat16 &b) { - return __bfloat162float(a) < __bfloat162float(b); -} - -__device__ __forceinline__ bool is_smaller(const __half &a, const __half &b) { - return __half2float(a) < __half2float(b); -} - -template -struct alignas(sizeof(int64_t) * 2) KeyValuePair { - T val; - int64_t idx; - - __device__ __forceinline__ KeyValuePair() {} - __device__ __forceinline__ KeyValuePair(T v, int64_t i) : val(v), idx(i) {} - - __device__ __forceinline__ static KeyValuePair max_value() { - if constexpr (std::is_floating_point_v) { - return {static_cast(INFINITY), -1}; - } else if constexpr (std::is_same_v) { - unsigned short inf_val = 0x7C00; - return {*reinterpret_cast<__half *>(&inf_val), -1}; - } else if constexpr (std::is_same_v) { - return {__float2bfloat16(1e30f), -1}; - } else { - return {static_cast(1e30), -1}; - } - } -}; - -template -__device__ __forceinline__ void compare_and_swap(KeyValuePair &a, KeyValuePair &b, bool dir) { - bool smaller = is_smaller(a.val, b.val) || (a.val == b.val && a.idx < b.idx); - if (smaller != dir) { - KeyValuePair tmp = a; - a = b; - b = tmp; - } -} - -template -__global__ void kthvalue_kernel( - T *__restrict__ out_values, - int64_t *__restrict__ out_indices, - const T *__restrict__ input, - size_t dim_size, - size_t inner_size, - int k, - size_t power_of_2_dim) { - extern __shared__ char smem[]; - KeyValuePair *s_data = reinterpret_cast *>(smem); - - unsigned int tid = threadIdx.x; - unsigned int bid = blockIdx.x; - - size_t outer_idx = bid / inner_size; - size_t inner_idx = bid % inner_size; - size_t input_base = outer_idx * dim_size * inner_size + inner_idx; - size_t stride = inner_size; - - for (unsigned int i = tid; i < power_of_2_dim; i += blockDim.x) { - if (i < dim_size) { - T val = input[input_base + i * stride]; - s_data[i] = KeyValuePair(val, static_cast(i)); - } else { - s_data[i] = KeyValuePair::max_value(); - } - } - __syncthreads(); - - for (unsigned int size = 2; size <= power_of_2_dim; size <<= 1) { - bool dir = (tid & (size / 2)) == 0; - for (unsigned int stride_step = size >> 1; stride_step > 0; stride_step >>= 1) { - unsigned int pos = 2 * tid - (tid & (stride_step - 1)); - if (pos + stride_step < power_of_2_dim) { - unsigned int next_pos = pos + stride_step; - bool direction = ((pos & size) == 0); - compare_and_swap(s_data[pos], s_data[next_pos], direction); - } - __syncthreads(); - } - } - - if (tid == 0) { - int target_k = k - 1; - if (target_k >= 0 && target_k < dim_size) { - out_values[bid] = s_data[target_k].val; - out_indices[bid] = s_data[target_k].idx; - } - } -} - -static inline size_t next_power_of_2(size_t n) { - if (n == 0) { - return 1; - } - size_t p = 1; - while (p < n) { - p <<= 1; - } - return p; -} - -template -void launch_kernel( - void *values, - void *indices, - const void *input, - int outer_size, - int inner_size, - int dim_size, - int k, - void *stream) { - auto hc_stream = reinterpret_cast(stream); - - size_t power_of_2_dim = next_power_of_2(dim_size); - size_t total_slices = (size_t)outer_size * inner_size; - - unsigned int threads_per_block = std::max(1u, (unsigned int)(power_of_2_dim / 2)); - if (threads_per_block > 1024) { - threads_per_block = 1024; - } - - size_t smem_size = power_of_2_dim * sizeof(KeyValuePair); - - if (power_of_2_dim > 2048) { - return; - } - - kthvalue_kernel<<>>( - reinterpret_cast(values), - reinterpret_cast(indices), - reinterpret_cast(input), - dim_size, - inner_size, - k, - power_of_2_dim); -} - -struct Descriptor::Opaque { - std::shared_ptr internal; - int k; - int outer_size, inner_size, dim_size; -}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, Descriptor **desc_ptr, - infiniopTensorDescriptor_t values_desc, - infiniopTensorDescriptor_t indices_desc, - infiniopTensorDescriptor_t input_desc, - int k, - int dim, - int keepdim) { - auto handle = reinterpret_cast(handle_); - - auto info_result = KthvalueInfo::create(values_desc, indices_desc, input_desc, k, dim, keepdim); - if (!info_result) { - return info_result.status(); - } - - auto in_d = reinterpret_cast(input_desc); - int ndim = in_d->ndim(); - int64_t real_dim = dim < 0 ? dim + ndim : dim; - - int outer = 1; - for (int i = 0; i < real_dim; ++i) { - outer *= in_d->shape()[i]; - } - int inner = 1; - for (int i = real_dim + 1; i < ndim; ++i) { - inner *= in_d->shape()[i]; - } - int dim_s = in_d->shape()[real_dim]; - - if (next_power_of_2(dim_s) > 2048) { - return INFINI_STATUS_BAD_PARAM; - } - - auto opaque = new Opaque{handle->internal(), (int)k, outer, inner, dim_s}; - *desc_ptr = new Descriptor(opaque, info_result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *values, - void *indices, - const void *input, - void *stream) const { - auto dtype = _info.dtype(); - int k = _opaque->k; - int outer = _opaque->outer_size; - int inner = _opaque->inner_size; - int dim_s = _opaque->dim_size; - - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel<__half>(values, indices, input, outer, inner, dim_s, k, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(values, indices, input, outer, inner, dim_s, k, stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(values, indices, input, outer, inner, dim_s, k, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(values, indices, input, outer, inner, dim_s, k, stream); - break; - case INFINI_DTYPE_I32: - launch_kernel(values, indices, input, outer, inner, dim_s, k, stream); - break; - case INFINI_DTYPE_I64: - launch_kernel(values, indices, input, outer, inner, dim_s, k, stream); - break; - case INFINI_DTYPE_U32: - launch_kernel(values, indices, input, outer, inner, dim_s, k, stream); - break; - case INFINI_DTYPE_U64: - launch_kernel(values, indices, input, outer, inner, dim_s, k, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::kthvalue::metax diff --git a/src/infiniop/ops/kthvalue/moore/kthvalue_moore.h b/src/infiniop/ops/kthvalue/moore/kthvalue_moore.h deleted file mode 100644 index e512ad195..000000000 --- a/src/infiniop/ops/kthvalue/moore/kthvalue_moore.h +++ /dev/null @@ -1,6 +0,0 @@ -#ifndef __KTHVALUE_MOORE_API_H__ -#define __KTHVALUE_MOORE_API_H__ -#include "../kthvalue.h" -DESCRIPTOR(moore) - -#endif // __KTHVALUE_MOORE_API_H__ diff --git a/src/infiniop/ops/kthvalue/moore/kthvalue_moore_kernel.h b/src/infiniop/ops/kthvalue/moore/kthvalue_moore_kernel.h deleted file mode 100644 index 23f06d822..000000000 --- a/src/infiniop/ops/kthvalue/moore/kthvalue_moore_kernel.h +++ /dev/null @@ -1,103 +0,0 @@ -#ifndef __KTHVALUE_MOORE_H__ -#define __KTHVALUE_MOORE_H__ - -#include -#include -#include -#include -#include -#include - -namespace op::kthvalue::moore { - -template -struct alignas(sizeof(int64_t) * 2) KeyValuePair { - T val; - int64_t idx; - - __device__ __forceinline__ KeyValuePair() {} - __device__ __forceinline__ KeyValuePair(T v, int64_t i) : val(v), idx(i) {} - - __device__ __forceinline__ static KeyValuePair max_value() { - if constexpr (std::is_floating_point_v) { - return {static_cast(INFINITY), -1}; - } else { - return {static_cast(1e30), -1}; - } - } -}; - -template -__device__ __forceinline__ bool is_smaller(const T &a, const T &b) { - return a < b; -} - -template -__device__ __forceinline__ void compare_and_swap(KeyValuePair &a, KeyValuePair &b, bool dir) { - bool smaller = is_smaller(a.val, b.val) || (a.val == b.val && a.idx < b.idx); - - if (smaller != dir) { - KeyValuePair tmp = a; - a = b; - b = tmp; - } -} - -template -__global__ void kthvalue_kernel( - T *__restrict__ out_values, - int64_t *__restrict__ out_indices, - const T *__restrict__ input, - size_t dim_size, - size_t inner_size, - int k, - size_t power_of_2_dim) { - extern __shared__ char smem[]; - auto s_data = reinterpret_cast *>(smem); - - unsigned int tid = threadIdx.x; - unsigned int bid = blockIdx.x; - - size_t outer_idx = bid / inner_size; - size_t inner_idx = bid % inner_size; - - size_t input_base = outer_idx * dim_size * inner_size + inner_idx; - size_t stride = inner_size; - - for (unsigned int i = tid; i < power_of_2_dim; i += blockDim.x) { - if (i < dim_size) { - T val = input[input_base + i * stride]; - s_data[i] = KeyValuePair(val, static_cast(i)); - } else { - s_data[i] = KeyValuePair::max_value(); - } - } - __syncthreads(); - - for (unsigned int size = 2; size <= power_of_2_dim; size <<= 1) { - bool dir = (tid & (size / 2)) == 0; - - for (unsigned int stride_step = size >> 1; stride_step > 0; stride_step >>= 1) { - unsigned int pos = 2 * tid - (tid & (stride_step - 1)); - - if (pos + stride_step < power_of_2_dim) { - unsigned int next_pos = pos + stride_step; - bool direction = ((pos & size) == 0); - compare_and_swap(s_data[pos], s_data[next_pos], direction); - } - __syncthreads(); - } - } - - if (tid == 0) { - int target_k = k - 1; - if (target_k >= 0 && target_k < dim_size) { - out_values[bid] = s_data[target_k].val; - out_indices[bid] = s_data[target_k].idx; - } - } -} - -} // namespace op::kthvalue::moore - -#endif // __KTHVALUE_MOORE_H__ diff --git a/src/infiniop/ops/kthvalue/moore/kthvalue_moore_kernel.mu b/src/infiniop/ops/kthvalue/moore/kthvalue_moore_kernel.mu deleted file mode 100644 index 9b47fa1d2..000000000 --- a/src/infiniop/ops/kthvalue/moore/kthvalue_moore_kernel.mu +++ /dev/null @@ -1,143 +0,0 @@ -#include "../../../devices/moore/moore_handle.h" -#include "kthvalue_moore.h" -#include "kthvalue_moore_kernel.h" -#include -#include -#include -#include -#include - -namespace op::kthvalue::moore { - -template -static inline bool is_aligned(const void *ptr, size_t alignment) { - return reinterpret_cast(ptr) % alignment == 0; -} - -static inline size_t next_power_of_2(size_t n) { - if (n == 0) { - return 1; - } - size_t p = 1; - while (p < n) { - p <<= 1; - } - return p; -} - -template -void launch_kernel( - void *values, - void *indices, - const void *input, - const KthvalueInfo &info, - void *stream) { - - auto in_ptr = reinterpret_cast(input); - auto val_ptr = reinterpret_cast(values); - auto idx_ptr = reinterpret_cast(indices); - - auto musa_stream = reinterpret_cast(stream); - - size_t dim_size = info.dim_size(); - size_t outer_size = info.outer_size(); - size_t inner_size = info.inner_size(); - int k = info.k(); - - size_t power_of_2_dim = next_power_of_2(dim_size); - - size_t total_slices = outer_size * inner_size; - - unsigned int threads_per_block = std::max(1u, (unsigned int)(power_of_2_dim / 2)); - if (threads_per_block > 1024) { - threads_per_block = 1024; - } - - size_t smem_size = power_of_2_dim * sizeof(op::kthvalue::moore::KeyValuePair); - - if (power_of_2_dim > 2048) { - return; - } - - op::kthvalue::moore::kthvalue_kernel - <<>>( - val_ptr, - idx_ptr, - in_ptr, - dim_size, - inner_size, - k, - power_of_2_dim); -} - -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, Descriptor **desc_ptr, - infiniopTensorDescriptor_t values_desc, - infiniopTensorDescriptor_t indices_desc, - infiniopTensorDescriptor_t input_desc, - int k, - int dim, - int keepdim) { - - auto info_result = KthvalueInfo::create(values_desc, indices_desc, input_desc, k, dim, keepdim); - if (!info_result) { - return info_result.status(); - } - - size_t workspace_size = 0; - - *desc_ptr = new Descriptor(new Opaque(), info_result.take(), workspace_size, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *values, - void *indices, - const void *input, - void *stream) const { - - auto dtype = _info.dtype(); - - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel<__mt_bfloat16>(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_I32: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_I64: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_U32: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_U64: - launch_kernel(values, indices, input, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::kthvalue::moore diff --git a/src/infiniop/ops/kthvalue/nvidia/kthvalue_nvidia.cu b/src/infiniop/ops/kthvalue/nvidia/kthvalue_nvidia.cu deleted file mode 100644 index 69860f6e3..000000000 --- a/src/infiniop/ops/kthvalue/nvidia/kthvalue_nvidia.cu +++ /dev/null @@ -1,164 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include "../../../handle.h" - -#include "../cuda/kernel.cuh" -#include "kthvalue_nvidia.cuh" -#include -#include - -namespace op::kthvalue::nvidia { - -template -static inline bool is_aligned(const void *ptr, size_t alignment) { - return reinterpret_cast(ptr) % alignment == 0; -} - -// ================================================================== -// Helper: Next Power of 2 -// ================================================================== -static inline size_t next_power_of_2(size_t n) { - if (n == 0) { - return 1; - } - size_t p = 1; - while (p < n) { - p <<= 1; - } - return p; -} - -// ================================================================== -// Kernel Launch Logic -// ================================================================== -template -void launch_kernel( - void *values, - void *indices, - const void *input, - const KthvalueInfo &info, - void *stream) { - - // 1. 准备指针 - auto in_ptr = reinterpret_cast(input); - auto val_ptr = reinterpret_cast(values); - auto idx_ptr = reinterpret_cast(indices); - - auto cuda_stream = reinterpret_cast(stream); - - // 2. 准备形状参数 - size_t dim_size = info.dim_size(); - size_t outer_size = info.outer_size(); - size_t inner_size = info.inner_size(); - int k = info.k(); - - // 3. 计算 Grid/Block 和 Shared Memory - // Bitonic Sort 需要 padding 到 2 的幂次 - size_t power_of_2_dim = next_power_of_2(dim_size); - - // Grid: 总切片数 (Outer * Inner) - size_t total_slices = outer_size * inner_size; - - // Block: 至少 power_of_2_dim / 2 个线程用于比较 - // 限制最大线程数 1024 - unsigned int threads_per_block = std::max(1u, (unsigned int)(power_of_2_dim / 2)); - - // Shared Memory 大小 - size_t smem_size = power_of_2_dim * sizeof(op::kthvalue::cuda::KeyValuePair); - - // 硬件限制检查 (Shared Memory Sort 限制) - // 假设最大支持 Dim Size 为 2048 (对应 1024 线程) - // 如果超过此限制，需切换到 Global Memory Sort (此处简化处理，仅支持 Shared Mem Sort) - if (power_of_2_dim > 2048) { - // Log Error or Fallback? - // 在实际工程中应返回 Error Code，这里作为 void 函数假设上层已校验或接受限制 - return; - } - - // 4. 启动 Kernel - op::kthvalue::cuda::kthvalue_kernel - <<>>( - val_ptr, - idx_ptr, - in_ptr, - dim_size, - inner_size, - k, - power_of_2_dim); -} - -// ================================================================== -// Descriptor 实现 -// ================================================================== -struct Descriptor::Opaque {}; - -Descriptor::~Descriptor() { - if (_opaque) { - delete _opaque; - } -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, Descriptor **desc_ptr, - infiniopTensorDescriptor_t values_desc, - infiniopTensorDescriptor_t indices_desc, - infiniopTensorDescriptor_t input_desc, - int k, - int dim, - int keepdim) { - - auto info_result = KthvalueInfo::create(values_desc, indices_desc, input_desc, k, dim, keepdim); - if (!info_result) { - return info_result.status(); - } - - // 目前基于 Shared Memory 的实现不需要额外的 Workspace - size_t workspace_size = 0; - - *desc_ptr = new Descriptor(new Opaque(), info_result.take(), workspace_size, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *values, - void *indices, - const void *input, - void *stream) const { - - auto dtype = _info.dtype(); - - switch (dtype) { - case INFINI_DTYPE_F16: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_BF16: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_F32: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_F64: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_I32: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_I64: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_U32: - launch_kernel(values, indices, input, _info, stream); - break; - case INFINI_DTYPE_U64: - launch_kernel(values, indices, input, _info, stream); - break; - default: - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::kthvalue::nvidia diff --git a/src/infiniop/ops/kthvalue/nvidia/kthvalue_nvidia.cuh b/src/infiniop/ops/kthvalue/nvidia/kthvalue_nvidia.cuh deleted file mode 100644 index d89a1a5d1..000000000 --- a/src/infiniop/ops/kthvalue/nvidia/kthvalue_nvidia.cuh +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __KTHVALUE_NVIDIA_CUH__ -#define __KTHVALUE_NVIDIA_CUH__ - -#include "../kthvalue.h" -DESCRIPTOR(nvidia) - -#endif // __KTHVALUE_NVIDIA_CUH__ diff --git a/src/infiniop/ops/kthvalue/operator.cc b/src/infiniop/ops/kthvalue/operator.cc deleted file mode 100644 index 094d0a1e9..000000000 --- a/src/infiniop/ops/kthvalue/operator.cc +++ /dev/null @@ -1,197 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/kthvalue.h" - -// --- 后端实现头文件 --- -#ifdef ENABLE_CPU_API -#include "cpu/kthvalue_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) -#include "nvidia/kthvalue_nvidia.cuh" -#endif - -#ifdef ENABLE_METAX_API -#include "metax/kthvalue_metax.h" -#endif - -#ifdef ENABLE_MOORE_API -#include "moore/kthvalue_moore.h" -#endif - -extern "C" { - -// ======================================================================= -// 1. 创建算子描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopCreateKthvalueDescriptor( - infiniopHandle_t handle, - infiniopKthvalueDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t values, - infiniopTensorDescriptor_t indices, - infiniopTensorDescriptor_t input, - int k, - int dim, - int keepdim) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::kthvalue::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - values, \ - indices, \ - input, \ - k, \ - dim, \ - keepdim) - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CREATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CREATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CREATE -} - -// ======================================================================= -// 2. 获取 Workspace 大小 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopGetKthvalueWorkspaceSize(infiniopKthvalueDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - GET(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - GET(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - GET(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef GET -} - -// ======================================================================= -// 3. 执行计算 (Calculate) -// ======================================================================= -__INFINI_C infiniStatus_t infiniopKthvalue( - infiniopKthvalueDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *values, - void *indices, - const void *input, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, values, indices, input, stream) - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - CALCULATE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef CALCULATE -} - -// ======================================================================= -// 4. 销毁描述符 -// ======================================================================= -__INFINI_C infiniStatus_t infiniopDestroyKthvalueDescriptor(infiniopKthvalueDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DELETE(INFINI_DEVICE_CPU, cpu); -#endif -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - DELETE(INFINI_DEVICE_METAX, metax); -#endif -#ifdef ENABLE_MOORE_API - DELETE(INFINI_DEVICE_MOORE, moore); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } -#undef DELETE -} - -} // extern "C" diff --git a/src/infiniop/ops/kv_caching/cuda/kernel.cuh b/src/infiniop/ops/kv_caching/cuda/kernel.cuh deleted file mode 100644 index a4204db29..000000000 --- a/src/infiniop/ops/kv_caching/cuda/kernel.cuh +++ /dev/null @@ -1,63 +0,0 @@ -#ifndef __KV_CACHING_KERNEL_CUH__ -#define __KV_CACHING_KERNEL_CUH__ - -template -__device__ void kvCachingKernel( - Tdata *__restrict__ k_cache, - Tdata *__restrict__ v_cache, - const Tdata *__restrict__ k, - const Tdata *__restrict__ v, - const Tidx *__restrict__ past_kv_lengths, - int batch_size, - int num_kv_heads, - int max_seq_len, - int seq_len, - int hidden_dim, - ptrdiff_t k_cache_strides_0, - ptrdiff_t k_cache_strides_1, - ptrdiff_t k_cache_strides_2, - ptrdiff_t k_cache_strides_3, - ptrdiff_t v_cache_strides_0, - ptrdiff_t v_cache_strides_1, - ptrdiff_t v_cache_strides_2, - ptrdiff_t v_cache_strides_3, - ptrdiff_t k_strides_0, - ptrdiff_t k_strides_1, - ptrdiff_t k_strides_2, - ptrdiff_t k_strides_3, - ptrdiff_t v_strides_0, - ptrdiff_t v_strides_1, - ptrdiff_t v_strides_2, - ptrdiff_t v_strides_3) { - // num of ele = B * H * seq_len * D - int tid = blockIdx.x * blockDim.x + threadIdx.x; - int total = batch_size * num_kv_heads * seq_len * hidden_dim; - - const int grid_size = blockDim.x * gridDim.x; - - for (int idx = tid; idx < total; idx += grid_size) { - // unravel index - - int d = idx % hidden_dim; - idx /= hidden_dim; - - int s = idx % seq_len; - idx /= seq_len; - - int h = idx % num_kv_heads; - int b = idx / num_kv_heads; - - int past_len = static_cast(past_kv_lengths[b]); // Cast to int for both types - // write position - int cache_s = past_len + s; - int k_cache_offset = d * (int)k_cache_strides_3 + cache_s * (int)k_cache_strides_2 + h * (int)k_cache_strides_1 + b * (int)k_cache_strides_0; - int v_cache_offset = d * (int)v_cache_strides_3 + cache_s * (int)v_cache_strides_2 + h * (int)v_cache_strides_1 + b * (int)v_cache_strides_0; - - int k_src_offset = d * (int)k_strides_3 + s * (int)k_strides_2 + h * (int)k_strides_1 + b * (int)k_strides_0; - int v_src_offset = d * (int)v_strides_3 + s * (int)v_strides_2 + h * (int)v_strides_1 + b * (int)v_strides_0; - k_cache[k_cache_offset] = k[k_src_offset]; - v_cache[v_cache_offset] = v[v_src_offset]; - } -} - -#endif // __KV_CACHING_KERNEL_CUH__ diff --git a/src/infiniop/ops/kv_caching/info.h b/src/infiniop/ops/kv_caching/info.h deleted file mode 100644 index 8c384e48b..000000000 --- a/src/infiniop/ops/kv_caching/info.h +++ /dev/null @@ -1,109 +0,0 @@ -#ifndef __KV_CACHING_INFO_H__ -#define __KV_CACHING_INFO_H__ - -#include "../../../utils.h" -#include "../../operator.h" -#include "../../tensor.h" - -namespace op::kv_caching { - -class KVCachingInfo { -private: - KVCachingInfo() = default; - -public: - infiniDtype_t dtype; - infiniDtype_t past_len_dtype; - size_t batch_size, num_kv_heads, max_seq_len, seq_len, hidden_dim; - ptrdiff_t k_cache_strides_0, k_cache_strides_1, k_cache_strides_2, k_cache_strides_3; - ptrdiff_t v_cache_strides_0, v_cache_strides_1, v_cache_strides_2, v_cache_strides_3; - ptrdiff_t k_strides_0, k_strides_1, k_strides_2, k_strides_3; - ptrdiff_t v_strides_0, v_strides_1, v_strides_2, v_strides_3; - - static utils::Result createKVCachingInfo( - infiniopTensorDescriptor_t k_cache, - infiniopTensorDescriptor_t v_cache, - infiniopTensorDescriptor_t k, - infiniopTensorDescriptor_t v, - infiniopTensorDescriptor_t past_kv_lengths) { - - CHECK_OR_RETURN( - k_cache != nullptr && v_cache != nullptr && k != nullptr && v != nullptr && past_kv_lengths != nullptr, - INFINI_STATUS_NULL_POINTER); - - const infiniDtype_t dtype = k_cache->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - const infiniDtype_t past_len_dtype = past_kv_lengths->dtype(); - CHECK_DTYPE(past_len_dtype, INFINI_DTYPE_I32, INFINI_DTYPE_I64); - - CHECK_OR_RETURN(k_cache->ndim() == 4 - && v_cache->ndim() == 4 - && k->ndim() == 4 - && v->ndim() == 4, - INFINI_STATUS_BAD_TENSOR_SHAPE); - - auto shape = k_cache->shape(); - CHECK_SAME_SHAPE(shape, v_cache->shape()); - CHECK_SAME_SHAPE(k->shape(), v->shape()); - - size_t batch_size = shape[0]; - size_t num_kv_heads = shape[1]; - size_t max_seq_len = shape[2]; - size_t hidden_dim = shape[3]; - - size_t seq_len = k->shape()[2]; - - CHECK_OR_RETURN(batch_size == k->dim(0) - || num_kv_heads == k->dim(1) - || hidden_dim == k->dim(3), - INFINI_STATUS_BAD_TENSOR_SHAPE); - - ptrdiff_t k_cache_strides_0 = k_cache->strides()[0]; - ptrdiff_t k_cache_strides_1 = k_cache->strides()[1]; - ptrdiff_t k_cache_strides_2 = k_cache->strides()[2]; - ptrdiff_t k_cache_strides_3 = k_cache->strides()[3]; - - ptrdiff_t v_cache_strides_0 = v_cache->strides()[0]; - ptrdiff_t v_cache_strides_1 = v_cache->strides()[1]; - ptrdiff_t v_cache_strides_2 = v_cache->strides()[2]; - ptrdiff_t v_cache_strides_3 = v_cache->strides()[3]; - - ptrdiff_t k_strides_0 = k->strides()[0]; - ptrdiff_t k_strides_1 = k->strides()[1]; - ptrdiff_t k_strides_2 = k->strides()[2]; - ptrdiff_t k_strides_3 = k->strides()[3]; - - ptrdiff_t v_strides_0 = v->strides()[0]; - ptrdiff_t v_strides_1 = v->strides()[1]; - ptrdiff_t v_strides_2 = v->strides()[2]; - ptrdiff_t v_strides_3 = v->strides()[3]; - - return utils::Result(KVCachingInfo{ - dtype, - past_len_dtype, - batch_size, - num_kv_heads, - max_seq_len, - seq_len, - hidden_dim, - k_cache_strides_0, - k_cache_strides_1, - k_cache_strides_2, - k_cache_strides_3, - v_cache_strides_0, - v_cache_strides_1, - v_cache_strides_2, - v_cache_strides_3, - k_strides_0, - k_strides_1, - k_strides_2, - k_strides_3, - v_strides_0, - v_strides_1, - v_strides_2, - v_strides_3}); - } -}; -} // namespace op::kv_caching - -#endif // __KV_CACHING_INFO_H__ diff --git a/src/infiniop/ops/kv_caching/kv_caching.h b/src/infiniop/ops/kv_caching/kv_caching.h deleted file mode 100644 index e90a0db27..000000000 --- a/src/infiniop/ops/kv_caching/kv_caching.h +++ /dev/null @@ -1,49 +0,0 @@ -#ifndef KV_CACHING_H -#define KV_CACHING_H - -#include "../../operator.h" -#include "info.h" - -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::kv_caching::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - KVCachingInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - KVCachingInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t get_workspace_size() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t k_cache, \ - infiniopTensorDescriptor_t v_cache, \ - infiniopTensorDescriptor_t k, \ - infiniopTensorDescriptor_t v, \ - infiniopTensorDescriptor_t past_kv_lengths); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *k_cache, void *v_cache, \ - const void *k, const void *v, const void *past_kv_lengths, \ - void *stream) const; \ - }; \ - } - -#endif // KV_CACHING_H diff --git a/src/infiniop/ops/kv_caching/metax/kv_caching_metax.h b/src/infiniop/ops/kv_caching/metax/kv_caching_metax.h deleted file mode 100644 index 083e89de0..000000000 --- a/src/infiniop/ops/kv_caching/metax/kv_caching_metax.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __KV_CACHING_METAX_API_H__ -#define __KV_CACHING_METAX_API_H__ -#include "../kv_caching.h" - -DESCRIPTOR(metax) - -#endif // __KV_CACHING_METAX_API_H__ diff --git a/src/infiniop/ops/kv_caching/metax/kv_caching_metax.maca b/src/infiniop/ops/kv_caching/metax/kv_caching_metax.maca deleted file mode 100644 index bef30b4b8..000000000 --- a/src/infiniop/ops/kv_caching/metax/kv_caching_metax.maca +++ /dev/null @@ -1,173 +0,0 @@ -#include "../../../devices/metax/metax_common.h" -#include "kv_caching_metax.h" - -#include "../../../devices/metax/metax_kernel_common.h" -#include - -#include "../../../reduce/cuda/reduce.cuh" - -#include "../cuda/kernel.cuh" - -template -INFINIOP_METAX_KERNEL kvCaching( - Tdata *k_cache, - Tdata *v_cache, - const Tdata *k, - const Tdata *v, - const Tidx *past_kv_lengths, - int batch_size, - int num_kv_heads, - int max_seq_len, - int seq_len, - int hidden_dim, - ptrdiff_t k_cache_strides_0, - ptrdiff_t k_cache_strides_1, - ptrdiff_t k_cache_strides_2, - ptrdiff_t k_cache_strides_3, - ptrdiff_t v_cache_strides_0, - ptrdiff_t v_cache_strides_1, - ptrdiff_t v_cache_strides_2, - ptrdiff_t v_cache_strides_3, - ptrdiff_t k_strides_0, - ptrdiff_t k_strides_1, - ptrdiff_t k_strides_2, - ptrdiff_t k_strides_3, - ptrdiff_t v_strides_0, - ptrdiff_t v_strides_1, - ptrdiff_t v_strides_2, - ptrdiff_t v_strides_3) { - kvCachingKernel(k_cache, v_cache, k, v, past_kv_lengths, - batch_size, num_kv_heads, max_seq_len, seq_len, hidden_dim, - k_cache_strides_0, k_cache_strides_1, k_cache_strides_2, k_cache_strides_3, - v_cache_strides_0, v_cache_strides_1, v_cache_strides_2, v_cache_strides_3, - k_strides_0, k_strides_1, k_strides_2, k_strides_3, - v_strides_0, v_strides_1, v_strides_2, v_strides_3); -} - -namespace op::kv_caching::metax { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t k_cache, - infiniopTensorDescriptor_t v_cache, - infiniopTensorDescriptor_t k, - infiniopTensorDescriptor_t v, - infiniopTensorDescriptor_t past_kv_lengths) { - auto info = KVCachingInfo::createKVCachingInfo(k_cache, v_cache, k, v, past_kv_lengths); - CHECK_RESULT(info); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(const KVCachingInfo &info, - Tdata *k_cache, - Tdata *v_cache, - const Tdata *k, - const Tdata *v, - const Tidx *past_kv_lengths, - hcStream_t stream, void *workspace) { - - int batch_size = static_cast(info.batch_size); - int num_kv_heads = static_cast(info.num_kv_heads); - int max_seq_len = static_cast(info.max_seq_len); - int hidden_dim = static_cast(info.hidden_dim); - - int seq_len = static_cast(info.seq_len); - - int total = batch_size * num_kv_heads * seq_len * hidden_dim; - - ptrdiff_t k_cache_strides_0 = info.k_cache_strides_0; - ptrdiff_t k_cache_strides_1 = info.k_cache_strides_1; - ptrdiff_t k_cache_strides_2 = info.k_cache_strides_2; - ptrdiff_t k_cache_strides_3 = info.k_cache_strides_3; - - ptrdiff_t v_cache_strides_0 = info.v_cache_strides_0; - ptrdiff_t v_cache_strides_1 = info.v_cache_strides_1; - ptrdiff_t v_cache_strides_2 = info.v_cache_strides_2; - ptrdiff_t v_cache_strides_3 = info.v_cache_strides_3; - - ptrdiff_t k_strides_0 = info.k_strides_0; - ptrdiff_t k_strides_1 = info.k_strides_1; - ptrdiff_t k_strides_2 = info.k_strides_2; - ptrdiff_t k_strides_3 = info.k_strides_3; - - ptrdiff_t v_strides_0 = info.v_strides_0; - ptrdiff_t v_strides_1 = info.v_strides_1; - ptrdiff_t v_strides_2 = info.v_strides_2; - ptrdiff_t v_strides_3 = info.v_strides_3; - - int num_blocks = (total + BLOCK_SIZE - 1) / BLOCK_SIZE; - - kvCaching - <<>>(k_cache, v_cache, k, v, past_kv_lengths, - batch_size, num_kv_heads, max_seq_len, seq_len, hidden_dim, - k_cache_strides_0, k_cache_strides_1, k_cache_strides_2, k_cache_strides_3, - v_cache_strides_0, v_cache_strides_1, v_cache_strides_2, v_cache_strides_3, - k_strides_0, k_strides_1, k_strides_2, k_strides_3, - v_strides_0, v_strides_1, v_strides_2, v_strides_3); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *k_cache, - void *v_cache, - const void *k, - const void *v, - const void *past_kv_lengths, - void *stream_) const { - hcStream_t stream = (hcStream_t)stream_; - -#define LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DTYPE(BLOCK_SIZE, TDATA, TIDX) \ - launchKernel(_info, (TDATA *)k_cache, (TDATA *)v_cache, \ - (const TDATA *)k, (const TDATA *)v, \ - (const TIDX *)past_kv_lengths, stream, workspace) - -#define LAUNCH_KERNEL_WITH_BLOCK_SIZE(BLOCK_SIZE, TDATA) \ - if (_info.past_len_dtype == INFINI_DTYPE_I32) { \ - return LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DTYPE(BLOCK_SIZE, TDATA, int32_t); \ - } else { /* INFINI_DTYPE_I64 */ \ - return LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DTYPE(BLOCK_SIZE, TDATA, int64_t); \ - } - -#define LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(BLOCK_SIZE) \ - { \ - if (_info.dtype == INFINI_DTYPE_F16) { \ - LAUNCH_KERNEL_WITH_BLOCK_SIZE(BLOCK_SIZE, half) \ - } else if (_info.dtype == INFINI_DTYPE_F32) { \ - LAUNCH_KERNEL_WITH_BLOCK_SIZE(BLOCK_SIZE, float) \ - } else if (_info.dtype == INFINI_DTYPE_BF16) { \ - LAUNCH_KERNEL_WITH_BLOCK_SIZE(BLOCK_SIZE, __hpcc_bfloat16) \ - } else { \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } \ - } - - // Choose block size based on device capabilities - if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_1024) { - LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(METAX_BLOCK_SIZE_1024) - } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_512) { - LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(METAX_BLOCK_SIZE_512) - } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_2048) { - LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(METAX_BLOCK_SIZE_2048) - } else if (_opaque->internal->maxThreadsPerBlock() == METAX_BLOCK_SIZE_4096) { - LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(METAX_BLOCK_SIZE_4096) - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::kv_caching::metax diff --git a/src/infiniop/ops/kv_caching/ninetoothed/build.py b/src/infiniop/ops/kv_caching/ninetoothed/build.py deleted file mode 100644 index f76497434..000000000 --- a/src/infiniop/ops/kv_caching/ninetoothed/build.py +++ /dev/null @@ -1,28 +0,0 @@ -import ninetoothed - -import infiniop.ninetoothed.build - -from . import kv_caching - - -def build(): - dtype_values = ( - ninetoothed.float16, - ninetoothed.bfloat16, - ninetoothed.float32, - ) - - constexpr_param_grid = { - "emb_dim": (1, 16, 32, 64, 128, 256), - "dtype": dtype_values, - "block_size_m": (64,), - "block_size_n": (64,), - } - - infiniop.ninetoothed.build.build( - kv_caching.premake, - constexpr_param_grid, - caller="cuda", - op_name="kv_caching", - output_dir=infiniop.ninetoothed.build.BUILD_DIRECTORY_PATH, - ) diff --git a/src/infiniop/ops/kv_caching/ninetoothed/kv_caching.h b/src/infiniop/ops/kv_caching/ninetoothed/kv_caching.h deleted file mode 100644 index 43388f58d..000000000 --- a/src/infiniop/ops/kv_caching/ninetoothed/kv_caching.h +++ /dev/null @@ -1,101 +0,0 @@ -#ifndef KV_CACHING_H -#define KV_CACHING_H - -#include "../../../handle.h" -#include "../../../operator.h" -#include "../../../tensor.h" - -#include "../../../../../build/ninetoothed/kv_caching.h" -#include "../../../ninetoothed/utils.h" - -namespace op::kv_caching::ninetoothed { -class Descriptor final : public InfiniopDescriptor { - -public: - Descriptor( - infiniopHandle_t handle, - infiniopTensorDescriptor_t k_cache_desc, - infiniopTensorDescriptor_t v_cache_desc, - infiniopTensorDescriptor_t k_desc, - infiniopTensorDescriptor_t v_desc, - infiniopTensorDescriptor_t past_kv_lengths_desc) : InfiniopDescriptor{handle->device, handle->device_id}, - k_cache_shape_{k_cache_desc->shape()}, - k_cache_strides_{k_cache_desc->strides()}, - v_cache_shape_{v_cache_desc->shape()}, - v_cache_strides_{v_cache_desc->strides()}, - k_shape_{k_desc->shape()}, - k_strides_{k_desc->strides()}, - v_shape_{v_desc->shape()}, - v_strides_{v_desc->strides()}, - past_kv_lengths_shape_{past_kv_lengths_desc->shape()}, - past_kv_lengths_strides_{past_kv_lengths_desc->strides()}, - dtype_{k_desc->dtype()} {} - - ~Descriptor() = default; - - size_t get_workspace_size() const { return 0; }; - - static infiniStatus_t create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t k_cache, - infiniopTensorDescriptor_t v_cache, - infiniopTensorDescriptor_t k, - infiniopTensorDescriptor_t v, - infiniopTensorDescriptor_t past_kv_lengths) { - *desc_ptr = new Descriptor{handle, k_cache, v_cache, k, v, past_kv_lengths}; - return INFINI_STATUS_SUCCESS; - } - - infiniStatus_t calculate( - void *workspace, size_t workspace_size, - void *k_cache, - void *v_cache, - const void *k, - const void *v, - const void *past_kv_lengths, - void *stream) const { - auto k_cache_nt{::ninetoothed::Tensor{k_cache, k_cache_shape_, k_cache_strides_}}; - auto v_cache_nt{::ninetoothed::Tensor{v_cache, v_cache_shape_, v_cache_strides_}}; - auto k_nt{::ninetoothed::Tensor{k, k_shape_, k_strides_}}; - auto v_nt{::ninetoothed::Tensor{v, v_shape_, v_strides_}}; - auto past_kv_lengths_nt{::ninetoothed::Tensor{past_kv_lengths, past_kv_lengths_shape_, past_kv_lengths_strides_}}; - - if (launch_kv_caching(stream, - k_cache_nt, - v_cache_nt, - k_nt, - v_nt, - past_kv_lengths_nt, - k_shape_[3], - dtype_, - 64, 64)) { - return INFINI_STATUS_NOT_IMPLEMENTED; - } - - return INFINI_STATUS_SUCCESS; - } - -private: - using Size = ::ninetoothed::Tensor<>::Size; - using Stride = ::ninetoothed::Tensor<>::Stride; - - std::vector k_cache_shape_; - std::vector k_cache_strides_; - - std::vector v_cache_shape_; - std::vector v_cache_strides_; - - std::vector k_shape_; - std::vector k_strides_; - std::vector v_shape_; - std::vector v_strides_; - - std::vector past_kv_lengths_shape_; - std::vector past_kv_lengths_strides_; - - infiniDtype_t dtype_; -}; -} // namespace op::kv_caching::ninetoothed - -#endif // KV_CACHING_H diff --git a/src/infiniop/ops/kv_caching/ninetoothed/kv_caching.py b/src/infiniop/ops/kv_caching/ninetoothed/kv_caching.py deleted file mode 100644 index 006dbe377..000000000 --- a/src/infiniop/ops/kv_caching/ninetoothed/kv_caching.py +++ /dev/null @@ -1,67 +0,0 @@ -import functools - -import ninetoothed -from ninetoothed import Tensor - - -def arrangement( - k_cache, - v_cache, - k, - v, - past_lengths, - block_size_m=ninetoothed.block_size(), - block_size_n=ninetoothed.block_size(), -): - k_cache_arranged = k_cache.tile((1, block_size_m, 1, -1)).tile((1, 1, -1, 1)) - v_cache_arranged = v_cache.tile((1, block_size_m, 1, -1)).tile((1, 1, -1, 1)) - - k_arranged = k.tile((1, block_size_m, 1, -1)).tile((1, 1, -1, 1)) - v_arranged = v.tile((1, block_size_m, 1, -1)).tile((1, 1, -1, 1)) - - past_lengths_arranged = ( - past_lengths.tile((1,)) - .unsqueeze(1) - .unsqueeze(2) - .unsqueeze(3) - .unsqueeze(4) - .expand((-1, *k_arranged.shape)) - ) - - return ( - k_cache_arranged, - v_cache_arranged, - k_arranged, - v_arranged, - past_lengths_arranged, - ) - - -def application(k_cache, v_cache, k, v, past_lengths): - pos = past_lengths - - for i in range(k.shape[-2]): - k_cache[0, 0, pos + i, 0] = k[0, 0, i, 0] - v_cache[0, 0, pos + i, 0] = v[0, 0, i, 0] - - -def premake(emb_dim=None, dtype=None, block_size_m=None, block_size_n=None): - arrangement_ = functools.partial( - arrangement, block_size_m=block_size_m, block_size_n=block_size_n - ) - - shape_options = (None, None, None, {"constexpr": True, "upper_bound": 256}) - - tensors = ( - Tensor(4, dtype=dtype, shape_options=shape_options), - Tensor(4, dtype=dtype, shape_options=shape_options), - Tensor(4, dtype=dtype, shape_options=shape_options), - Tensor(4, dtype=dtype, shape_options=shape_options), - Tensor(1, dtype=ninetoothed.int64), - ) - - if emb_dim is not None: - for tensor in tensors: - tensor.shape = tensor.shape[:-1] + (emb_dim,) - - return arrangement_, application, tensors diff --git a/src/infiniop/ops/kv_caching/nvidia/kv_caching_nvidia.cu b/src/infiniop/ops/kv_caching/nvidia/kv_caching_nvidia.cu deleted file mode 100644 index 703230dd4..000000000 --- a/src/infiniop/ops/kv_caching/nvidia/kv_caching_nvidia.cu +++ /dev/null @@ -1,172 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "kv_caching_nvidia.cuh" - -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include - -#include "../../../reduce/cuda/reduce.cuh" - -#include "../cuda/kernel.cuh" - -template -INFINIOP_CUDA_KERNEL kvCaching( - Tdata *k_cache, - Tdata *v_cache, - const Tdata *k, - const Tdata *v, - const Tidx *past_kv_lengths, - int batch_size, - int num_kv_heads, - int max_seq_len, - int seq_len, - int hidden_dim, - ptrdiff_t k_cache_strides_0, - ptrdiff_t k_cache_strides_1, - ptrdiff_t k_cache_strides_2, - ptrdiff_t k_cache_strides_3, - ptrdiff_t v_cache_strides_0, - ptrdiff_t v_cache_strides_1, - ptrdiff_t v_cache_strides_2, - ptrdiff_t v_cache_strides_3, - ptrdiff_t k_strides_0, - ptrdiff_t k_strides_1, - ptrdiff_t k_strides_2, - ptrdiff_t k_strides_3, - ptrdiff_t v_strides_0, - ptrdiff_t v_strides_1, - ptrdiff_t v_strides_2, - ptrdiff_t v_strides_3) { - kvCachingKernel(k_cache, v_cache, k, v, past_kv_lengths, - batch_size, num_kv_heads, max_seq_len, seq_len, hidden_dim, - k_cache_strides_0, k_cache_strides_1, k_cache_strides_2, k_cache_strides_3, - v_cache_strides_0, v_cache_strides_1, v_cache_strides_2, v_cache_strides_3, - k_strides_0, k_strides_1, k_strides_2, k_strides_3, - v_strides_0, v_strides_1, v_strides_2, v_strides_3); -} - -namespace op::kv_caching::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t k_cache, - infiniopTensorDescriptor_t v_cache, - infiniopTensorDescriptor_t k, - infiniopTensorDescriptor_t v, - infiniopTensorDescriptor_t past_kv_lengths) { - auto info = KVCachingInfo::createKVCachingInfo(k_cache, v_cache, k, v, past_kv_lengths); - CHECK_RESULT(info); - - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(const KVCachingInfo &info, - Tdata *k_cache, - Tdata *v_cache, - const Tdata *k, - const Tdata *v, - const Tidx *past_kv_lengths, - cudaStream_t stream, void *workspace) { - - int batch_size = static_cast(info.batch_size); - int num_kv_heads = static_cast(info.num_kv_heads); - int max_seq_len = static_cast(info.max_seq_len); - int hidden_dim = static_cast(info.hidden_dim); - - int seq_len = static_cast(info.seq_len); - - int total = batch_size * num_kv_heads * seq_len * hidden_dim; - - ptrdiff_t k_cache_strides_0 = info.k_cache_strides_0; - ptrdiff_t k_cache_strides_1 = info.k_cache_strides_1; - ptrdiff_t k_cache_strides_2 = info.k_cache_strides_2; - ptrdiff_t k_cache_strides_3 = info.k_cache_strides_3; - - ptrdiff_t v_cache_strides_0 = info.v_cache_strides_0; - ptrdiff_t v_cache_strides_1 = info.v_cache_strides_1; - ptrdiff_t v_cache_strides_2 = info.v_cache_strides_2; - ptrdiff_t v_cache_strides_3 = info.v_cache_strides_3; - - ptrdiff_t k_strides_0 = info.k_strides_0; - ptrdiff_t k_strides_1 = info.k_strides_1; - ptrdiff_t k_strides_2 = info.k_strides_2; - ptrdiff_t k_strides_3 = info.k_strides_3; - - ptrdiff_t v_strides_0 = info.v_strides_0; - ptrdiff_t v_strides_1 = info.v_strides_1; - ptrdiff_t v_strides_2 = info.v_strides_2; - ptrdiff_t v_strides_3 = info.v_strides_3; - - int num_blocks = (total + BLOCK_SIZE - 1) / BLOCK_SIZE; - - kvCaching - <<>>(k_cache, v_cache, k, v, past_kv_lengths, - batch_size, num_kv_heads, max_seq_len, seq_len, hidden_dim, - k_cache_strides_0, k_cache_strides_1, k_cache_strides_2, k_cache_strides_3, - v_cache_strides_0, v_cache_strides_1, v_cache_strides_2, v_cache_strides_3, - k_strides_0, k_strides_1, k_strides_2, k_strides_3, - v_strides_0, v_strides_1, v_strides_2, v_strides_3); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *k_cache, - void *v_cache, - const void *k, - const void *v, - const void *past_kv_lengths, - void *stream_) const { - cudaStream_t stream = (cudaStream_t)stream_; - -#define LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DTYPE(BLOCK_SIZE, TDATA, TIDX) \ - launchKernel(_info, (TDATA *)k_cache, (TDATA *)v_cache, \ - (const TDATA *)k, (const TDATA *)v, \ - (const TIDX *)past_kv_lengths, stream, workspace) - -#define LAUNCH_KERNEL_WITH_BLOCK_SIZE(BLOCK_SIZE, TDATA) \ - if (_info.past_len_dtype == INFINI_DTYPE_I32) { \ - return LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DTYPE(BLOCK_SIZE, TDATA, int32_t); \ - } else { \ - return LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DTYPE(BLOCK_SIZE, TDATA, int64_t); \ - } - -#define LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(BLOCK_SIZE) \ - { \ - if (_info.dtype == INFINI_DTYPE_F16) { \ - LAUNCH_KERNEL_WITH_BLOCK_SIZE(BLOCK_SIZE, half) \ - } else if (_info.dtype == INFINI_DTYPE_F32) { \ - LAUNCH_KERNEL_WITH_BLOCK_SIZE(BLOCK_SIZE, float) \ - } else if (_info.dtype == INFINI_DTYPE_BF16) { \ - LAUNCH_KERNEL_WITH_BLOCK_SIZE(BLOCK_SIZE, __nv_bfloat16) \ - } else { \ - return INFINI_STATUS_BAD_TENSOR_DTYPE; \ - } \ - } - - if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_1024) { - LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(CUDA_BLOCK_SIZE_1024) - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_512) { - LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(CUDA_BLOCK_SIZE_512) - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_2048) { - LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(CUDA_BLOCK_SIZE_2048) - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) { - LAUNCH_KERNEL_WITH_BLOCK_SIZE_AND_DATA_TYPE(CUDA_BLOCK_SIZE_4096) - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::kv_caching::nvidia diff --git a/src/infiniop/ops/kv_caching/nvidia/kv_caching_nvidia.cuh b/src/infiniop/ops/kv_caching/nvidia/kv_caching_nvidia.cuh deleted file mode 100644 index 91eaa3f4b..000000000 --- a/src/infiniop/ops/kv_caching/nvidia/kv_caching_nvidia.cuh +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __KV_CACHING_NVIDIA_API_H__ -#define __KV_CACHING_NVIDIA_API_H__ -#include "../kv_caching.h" - -DESCRIPTOR(nvidia) - -#endif // __KV_CACHING_NVIDIA_API_H__ diff --git a/src/infiniop/ops/kv_caching/operator.cc b/src/infiniop/ops/kv_caching/operator.cc deleted file mode 100644 index e83bb1240..000000000 --- a/src/infiniop/ops/kv_caching/operator.cc +++ /dev/null @@ -1,176 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/kv_caching.h" - -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_QY_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_ALI_API) || defined(ENABLE_HYGON_API) -#include "nvidia/kv_caching_nvidia.cuh" -#endif -#if defined(ENABLE_METAX_API) -#include "metax/kv_caching_metax.h" -#endif - -__INFINI_C infiniStatus_t infiniopCreateKVCachingDescriptor( - infiniopHandle_t handle, - infiniopKVCachingDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t k_cache, - infiniopTensorDescriptor_t v_cache, - infiniopTensorDescriptor_t k, - infiniopTensorDescriptor_t v, - infiniopTensorDescriptor_t past_kv_lengths) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::kv_caching::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - k_cache, \ - v_cache, \ - k, \ - v, \ - past_kv_lengths) - - switch (handle->device) { - -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CREATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CREATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#if defined(ENABLE_METAX_API) - CREATE(INFINI_DEVICE_METAX, metax); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CREATE -} - -__INFINI_C infiniStatus_t infiniopGetKVCachingWorkspaceSize( - infiniopKVCachingDescriptor_t desc, - size_t *size) { - -#define GET_SIZE(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc) \ - ->get_workspace_size(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_NVIDIA_API - GET_SIZE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - GET_SIZE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - GET_SIZE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - GET_SIZE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_HYGON_API - GET_SIZE(INFINI_DEVICE_HYGON, nvidia); -#endif -#if defined(ENABLE_METAX_API) - GET_SIZE(INFINI_DEVICE_METAX, metax); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef GET_SIZE -} - -__INFINI_C infiniStatus_t infiniopKVCaching( - infiniopKVCachingDescriptor_t desc, - void *workspace, - size_t workspace_size, - void *k_cache, - void *v_cache, - const void *k, - const void *v, - const void *past_kv_lengths, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc) \ - ->calculate(workspace, workspace_size, k_cache, v_cache, k, v, past_kv_lengths, stream) - - switch (desc->device_type) { - -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - CALCULATE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_HYGON_API - CALCULATE(INFINI_DEVICE_HYGON, nvidia); -#endif -#if defined(ENABLE_METAX_API) - CALCULATE(INFINI_DEVICE_METAX, metax); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef CALCULATE -} - -__INFINI_C infiniStatus_t infiniopDestroyKVCachingDescriptor( - infiniopKVCachingDescriptor_t desc) { - -#define DELETE(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { - -#ifdef ENABLE_NVIDIA_API - DELETE(INFINI_DEVICE_NVIDIA, nvidia); -#endif -#ifdef ENABLE_QY_API - DELETE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_ILUVATAR_API - DELETE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_ALI_API - DELETE(INFINI_DEVICE_ALI, nvidia); -#endif -#ifdef ENABLE_HYGON_API - DELETE(INFINI_DEVICE_HYGON, nvidia); -#endif -#if defined(ENABLE_METAX_API) - DELETE(INFINI_DEVICE_METAX, metax); -#endif - - default: - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; - } - -#undef DELETE -} diff --git a/src/infiniop/ops/layer_norm/bang/layer_norm_bang.h b/src/infiniop/ops/layer_norm/bang/layer_norm_bang.h deleted file mode 100644 index 7c686bfb8..000000000 --- a/src/infiniop/ops/layer_norm/bang/layer_norm_bang.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __LAYER_NORM_BANG_H__ -#define __LAYER_NORM_BANG_H__ - -#include "../layer_norm.h" - -DESCRIPTOR(bang) - -#endif // __LAYER_NORM_BANG_H__ diff --git a/src/infiniop/ops/layer_norm/bang/layer_norm_bang.mlu b/src/infiniop/ops/layer_norm/bang/layer_norm_bang.mlu deleted file mode 100644 index d0379dce9..000000000 --- a/src/infiniop/ops/layer_norm/bang/layer_norm_bang.mlu +++ /dev/null @@ -1,288 +0,0 @@ -#include "../../../devices/bang/common_bang.h" -#include "../../../reduce/bang/reduce_bang.h" -#include "layer_norm_bang.h" - -#include -#include - -__nram__ char nram_buffer[NRAM_MAX_SIZE]; - -template -__mlu_device__ void loadToFloat(float *dst, T *cache, const T *src, size_t n) { - __memcpy(cache, src, n * sizeof(T), GDRAM2NRAM); - if constexpr (std::is_same::value) { - __bang_half2float(dst, cache, n); - } else if constexpr (std::is_same::value) { - __bang_bfloat162float(dst, cache, n); - } else { - __memcpy(dst, cache, n * sizeof(float), NRAM2NRAM); - } -} - -template -__mlu_device__ void storeFromFloat(T *dst, T *cache, float *src, size_t n) { - if constexpr (std::is_same::value) { - __bang_float2half(cache, src, n); - } else if constexpr (std::is_same::value) { - __bang_float2bfloat16(cache, src, n); - } else { - __memcpy(cache, src, n * sizeof(float), NRAM2NRAM); - } - __memcpy(dst, cache, n * sizeof(T), NRAM2GDRAM); -} - -template -__mlu_device__ void storeScalarFromFloat(T *dst, T *cache, float value) { - float *float_cache = reinterpret_cast(cache); - float_cache[0] = value; - storeFromFloat(dst, cache, float_cache, 1); -} - -__mlu_device__ int getMaxBatchSize(size_t dim, size_t data_size, size_t weight_size) { - constexpr int reduce_buffer_size = 128 / sizeof(float); - int max_batch_size = (NRAM_MAX_SIZE - 256 - reduce_buffer_size * sizeof(float)) / (4 * data_size + 2 * weight_size + 4 * sizeof(float)); - max_batch_size = std::min(max_batch_size, static_cast(dim)); - if (max_batch_size > 64) { - max_batch_size = (max_batch_size / 64) * 64; - } - return std::max(max_batch_size, 1); -} - -template -__mlu_global__ void layerNormKernel( - Tdata *__restrict__ output, - Tdata *__restrict__ standardization, - Tdata *__restrict__ std_deviation, - const Tdata *__restrict__ input, - const Tweight *__restrict__ weight, - const Tweight *__restrict__ bias, - const size_t *__restrict__ shape, - const ptrdiff_t *__restrict__ output_strides, - const ptrdiff_t *__restrict__ standardization_strides, - const ptrdiff_t *__restrict__ std_deviation_strides, - const ptrdiff_t *__restrict__ input_strides, - size_t rows, - size_t dim, - int ndim, - ptrdiff_t weight_stride, - ptrdiff_t bias_stride, - float eps, - bool has_bias) { - - constexpr int reduce_buffer_size = 128 / sizeof(float); - int max_batch_size = getMaxBatchSize(dim, sizeof(Tdata), sizeof(Tweight)); - - float *reduction_buffer = reinterpret_cast(nram_buffer); - Tdata *input_cache = reinterpret_cast(reduction_buffer + reduce_buffer_size); - Tdata *output_cache = input_cache + max_batch_size; - Tdata *std_cache = output_cache + max_batch_size; - Tweight *weight_cache = reinterpret_cast(std_cache + max_batch_size); - Tweight *bias_cache = weight_cache + max_batch_size; - float *input_float = reinterpret_cast(bias_cache + max_batch_size); - float *weight_float = input_float + max_batch_size; - float *bias_float = weight_float + max_batch_size; - - for (size_t row = taskId; row < rows; row += taskDim) { - ptrdiff_t input_offset = 0; - ptrdiff_t output_offset = 0; - ptrdiff_t standardization_offset = 0; - ptrdiff_t std_deviation_offset = 0; - size_t tmp = row; - for (int axis = ndim - 2; axis >= 0; --axis) { - size_t coord = tmp % shape[axis]; - tmp /= shape[axis]; - input_offset += static_cast(coord) * input_strides[axis]; - output_offset += static_cast(coord) * output_strides[axis]; - standardization_offset += static_cast(coord) * standardization_strides[axis]; - std_deviation_offset += static_cast(coord) * std_deviation_strides[axis]; - } - - const Tdata *input_row = input + input_offset; - Tdata *output_row = output + output_offset; - Tdata *standardization_row = standardization + standardization_offset; - - float sum = 0.0f; - size_t processed = 0; - while (processed < dim) { - size_t current = std::min(static_cast(max_batch_size), dim - processed); - loadToFloat(input_float, input_cache, input_row + processed * input_strides[ndim - 1], current); - if (current >= 128) { - op::common_bang::reduce_op::sumInternal(reduction_buffer, input_float, current); - sum += reduction_buffer[0]; - } else { - for (size_t i = 0; i < current; ++i) { - sum += input_float[i]; - } - } - processed += current; - } - - float mean = sum / static_cast(dim); - float variance_sum = 0.0f; - processed = 0; - while (processed < dim) { - size_t current = std::min(static_cast(max_batch_size), dim - processed); - loadToFloat(input_float, input_cache, input_row + processed * input_strides[ndim - 1], current); - __bang_sub_scalar(input_float, input_float, mean, current); - __bang_mul(input_float, input_float, input_float, current); - if (current >= 128) { - op::common_bang::reduce_op::sumInternal(reduction_buffer, input_float, current); - variance_sum += reduction_buffer[0]; - } else { - for (size_t i = 0; i < current; ++i) { - variance_sum += input_float[i]; - } - } - processed += current; - } - - float std_value = sqrtf(variance_sum / static_cast(dim) + eps); - float inv_std = 1.0f / std_value; - storeScalarFromFloat(std_deviation + std_deviation_offset, std_cache, std_value); - - processed = 0; - while (processed < dim) { - size_t current = std::min(static_cast(max_batch_size), dim - processed); - loadToFloat(input_float, input_cache, input_row + processed * input_strides[ndim - 1], current); - loadToFloat(weight_float, weight_cache, weight + processed * weight_stride, current); - if (has_bias) { - loadToFloat(bias_float, bias_cache, bias + processed * bias_stride, current); - } - - __bang_sub_scalar(input_float, input_float, mean, current); - __bang_mul_scalar(input_float, input_float, inv_std, current); - storeFromFloat(standardization_row + processed * standardization_strides[ndim - 1], std_cache, input_float, current); - - __bang_mul(input_float, input_float, weight_float, current); - if (has_bias) { - __bang_add(input_float, input_float, bias_float, current); - } - storeFromFloat(output_row + processed * output_strides[ndim - 1], output_cache, input_float, current); - processed += current; - } - } -} - -template -void launchLayerNorm( - int core_per_cluster, - int cluster_count, - cnrtQueue_t queue, - void *output, - void *standardization, - void *std_deviation, - const void *input, - const void *weight, - const void *bias, - void *workspace, - const op::layer_norm::LayerNormInfo &info) { - cnrtDim3_t kernel_dim; - kernel_dim.x = core_per_cluster; - kernel_dim.y = cluster_count; - kernel_dim.z = 1; - - char *tmp_device = reinterpret_cast(workspace); - size_t *mlu_shape = reinterpret_cast(tmp_device); - ptrdiff_t *mlu_output_strides = reinterpret_cast(mlu_shape + info.ndim); - ptrdiff_t *mlu_standardization_strides = mlu_output_strides + info.ndim; - ptrdiff_t *mlu_std_deviation_strides = mlu_standardization_strides + info.ndim; - ptrdiff_t *mlu_input_strides = mlu_std_deviation_strides + (info.ndim - 1); - - CNRT_CHECK(cnrtMemcpyAsync(mlu_shape, const_cast(info.input_shape.data()), info.ndim * sizeof(size_t), queue, cnrtMemcpyHostToDev)); - CNRT_CHECK(cnrtMemcpyAsync(mlu_output_strides, const_cast(info.output_strides.data()), info.ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev)); - CNRT_CHECK(cnrtMemcpyAsync(mlu_standardization_strides, const_cast(info.input_standardization_strides.data()), info.ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev)); - CNRT_CHECK(cnrtMemcpyAsync(mlu_std_deviation_strides, const_cast(info.input_std_deviation_strides.data()), (info.ndim - 1) * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev)); - CNRT_CHECK(cnrtMemcpyAsync(mlu_input_strides, const_cast(info.input_strides.data()), info.ndim * sizeof(ptrdiff_t), queue, cnrtMemcpyHostToDev)); - - layerNormKernel<<>>( - reinterpret_cast(output), - reinterpret_cast(standardization), - reinterpret_cast(std_deviation), - reinterpret_cast(input), - reinterpret_cast(weight), - reinterpret_cast(bias), - mlu_shape, - mlu_output_strides, - mlu_standardization_strides, - mlu_std_deviation_strides, - mlu_input_strides, - info.othersize, - info.normalized_size, - static_cast(info.ndim), - info.weight_strides.back(), - info.bias_exist ? info.bias_strides.back() : 0, - info.eps, - info.bias_exist); - cnrtQueueSync(queue); -} - -namespace op::layer_norm::bang { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_standardization_desc, - infiniopTensorDescriptor_t input_std_deviation_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t bias_desc, - float eps) { - auto handle = reinterpret_cast(handle_); - auto result = LayerNormInfo::createLayerNormInfo( - output_desc, - input_standardization_desc, - input_std_deviation_desc, - input_desc, - weight_desc, - bias_desc, - eps); - CHECK_RESULT(result); - auto info = result.take(); - - CHECK_DTYPE(info.dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - *desc_ptr = new Descriptor( - info.dtype, - info, - info.ndim * sizeof(size_t) + (4 * info.ndim - 1) * sizeof(ptrdiff_t), - new Descriptor::Opaque{handle->internal()}, - handle->device, - handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - void *input_standardization, - void *input_std_deviation, - const void *input, - const void *weight, - const void *bias, - void *stream) const { - auto queue = reinterpret_cast(stream); - int core_per_cluster = _opaque->internal->getCorePerCluster(); - int cluster_count = _opaque->internal->getClusterCount(); - - if (_info.dtype == INFINI_DTYPE_F16) { - launchLayerNorm(core_per_cluster, cluster_count, queue, output, input_standardization, input_std_deviation, input, weight, bias, workspace, _info); - } else if (_info.dtype == INFINI_DTYPE_BF16) { - launchLayerNorm(core_per_cluster, cluster_count, queue, output, input_standardization, input_std_deviation, input, weight, bias, workspace, _info); - } else if (_info.dtype == INFINI_DTYPE_F32) { - launchLayerNorm(core_per_cluster, cluster_count, queue, output, input_standardization, input_std_deviation, input, weight, bias, workspace, _info); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::layer_norm::bang diff --git a/src/infiniop/ops/layer_norm/cpu/layer_norm_cpu.cc b/src/infiniop/ops/layer_norm/cpu/layer_norm_cpu.cc deleted file mode 100644 index 7331514e5..000000000 --- a/src/infiniop/ops/layer_norm/cpu/layer_norm_cpu.cc +++ /dev/null @@ -1,218 +0,0 @@ -#include "layer_norm_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include "../../../reduce/cpu/reduce.h" -#include "../info.h" - -namespace op::layer_norm::cpu { - -template -infiniStatus_t calculate_layer_norm( - const LayerNormInfo &info, - Tdata *output, - Tdata *input_standardization, - Tdata *input_std_deviation, - const Tdata *input, - const Tdata *weight, - const Tdata *bias) { - - const size_t ndim = info.ndim; - const size_t norm_size = info.normalized_size; // last dim - const size_t othersize = info.othersize; - - const auto &shape = info.input_shape; - const auto &in_strides = info.input_strides; - const auto &out_strides = info.output_strides; - const auto &std_strides = info.input_standardization_strides; - const auto &stddev_strides = info.input_std_deviation_strides; - - // -------- Special fast path: 1D tensor -------- - if (ndim == 1) { - const Tdata *input_ptr = input; - Tdata *output_ptr = output; - Tdata *standard_ptr = input_standardization; - Tdata *std_ptr = input_std_deviation; - - float mean = op::common_cpu::reduce_op::sum( - input_ptr, - norm_size, - in_strides[0]) - / norm_size; - - float sum_sq = op::common_cpu::reduce_op::sumSquared( - input_ptr, - norm_size, - in_strides[0]); - - float var = sum_sq / norm_size - mean * mean; - float std_dev = std::sqrt(var + info.eps); - - *std_ptr = utils::cast(std_dev); - - for (size_t d = 0; d < norm_size; d++) { - float x = utils::cast( - *(input_ptr + d * in_strides[0])); - - float x_std = (x - mean) / std_dev; - - *(standard_ptr + d * std_strides[0]) = utils::cast(x_std); - - float w = utils::cast( - *(weight + d * info.weight_strides[0])); - - float bval = info.bias_exist - ? utils::cast( - *(bias + d * info.bias_strides[0])) - : 0.0f; - - *(output_ptr + d * out_strides[0]) = utils::cast(x_std * w + bval); - } - - return INFINI_STATUS_SUCCESS; - } - - // -------- General case: ndim >= 2 -------- - - // index for dims [0 ... ndim-2] - std::vector idx(ndim - 1, 0); - -#pragma omp parallel for - for (ptrdiff_t b = 0; b < (ptrdiff_t)othersize; b++) { - - // ---- compute base offsets ---- - ptrdiff_t in_offset = 0; - ptrdiff_t out_offset = 0; - ptrdiff_t std_offset = 0; - ptrdiff_t stddev_offset = 0; - - for (size_t d = 0; d < ndim - 1; d++) { - in_offset += idx[d] * in_strides[d]; - out_offset += idx[d] * out_strides[d]; - std_offset += idx[d] * std_strides[d]; - stddev_offset += idx[d] * stddev_strides[d]; - } - - const Tdata *input_ptr = input + in_offset; - Tdata *output_ptr = output + out_offset; - Tdata *standard_ptr = input_standardization + std_offset; - Tdata *std_ptr = input_std_deviation + stddev_offset; - - // ---- mean ---- - float mean = op::common_cpu::reduce_op::sum( - input_ptr, - norm_size, - in_strides[ndim - 1]) - / norm_size; - - // ---- variance ---- - float sum_sq = op::common_cpu::reduce_op::sumSquared( - input_ptr, - norm_size, - in_strides[ndim - 1]); - - float var = sum_sq / norm_size - mean * mean; - float std_dev = std::sqrt(var + info.eps); - - *std_ptr = utils::cast(std_dev); - - // ---- normalize ---- - for (size_t d = 0; d < norm_size; d++) { - float x = utils::cast( - *(input_ptr + d * in_strides[ndim - 1])); - - float x_std = (x - mean) / std_dev; - - *(standard_ptr + d * std_strides[ndim - 1]) = utils::cast(x_std); - - float w = utils::cast( - *(weight + d * info.weight_strides[0])); - - float bval = info.bias_exist - ? utils::cast( - *(bias + d * info.bias_strides[0])) - : 0.0f; - - *(output_ptr + d * out_strides[ndim - 1]) = utils::cast(x_std * w + bval); - } - - // ---- increment multi-dim index (odometer style) ---- - for (int d = (int)ndim - 2; d >= 0; d--) { - idx[d]++; - if (idx[d] < shape[d]) { - break; - } - idx[d] = 0; - } - } - - return INFINI_STATUS_SUCCESS; -} - -Descriptor::~Descriptor() = default; - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle_, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t output_desc, - infiniopTensorDescriptor_t input_standardization_desc, - infiniopTensorDescriptor_t input_std_deviation_desc, - infiniopTensorDescriptor_t input_desc, - infiniopTensorDescriptor_t weight_desc, - infiniopTensorDescriptor_t bias_desc, - float eps) { - auto handle = reinterpret_cast(handle_); - - // --------------------- start: check data type and calculate workspace size ---------------------- - auto dtype = input_desc->dtype(); - CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); - size_t WorkSpaceSize = 0; - - auto result = LayerNormInfo::createLayerNormInfo( - output_desc, - input_standardization_desc, - input_std_deviation_desc, - input_desc, - weight_desc, - bias_desc, - eps); - CHECK_RESULT(result); - const LayerNormInfo &info = result.take(); - - *desc_ptr = new Descriptor( - dtype, std::move(info), WorkSpaceSize, - nullptr, - handle->device, handle->device_id); - - return INFINI_STATUS_SUCCESS; -} - -#define CALCULATE_LAYER_NORM(TDATA) \ - CHECK_STATUS(calculate_layer_norm(_info, \ - (TDATA *)output, \ - (TDATA *)input_standardization, \ - (TDATA *)input_std_deviation, \ - (const TDATA *)input, \ - (const TDATA *)weight, \ - (const TDATA *)bias)) - -infiniStatus_t Descriptor::calculate( - void *workspace, - size_t workspace_size, - void *output, - void *input_standardization, - void *input_std_deviation, - const void *input, - const void *weight, - const void *bias, - void *stream) const { - if (_info.dtype == INFINI_DTYPE_F16) { - CALCULATE_LAYER_NORM(fp16_t); - } else if (_info.dtype == INFINI_DTYPE_BF16) { - CALCULATE_LAYER_NORM(bf16_t); - } else if (_info.dtype == INFINI_DTYPE_F32) { - CALCULATE_LAYER_NORM(float); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} -} // namespace op::layer_norm::cpu diff --git a/src/infiniop/ops/layer_norm/cpu/layer_norm_cpu.h b/src/infiniop/ops/layer_norm/cpu/layer_norm_cpu.h deleted file mode 100644 index 206025df6..000000000 --- a/src/infiniop/ops/layer_norm/cpu/layer_norm_cpu.h +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __LAYER_NORM_CPU_H__ -#define __LAYER_NORM_CPU_H__ - -#include "../layer_norm.h" - -DESCRIPTOR(cpu) - -#endif // __LAYER_NORM_CPU_H__ diff --git a/src/infiniop/ops/layer_norm/cuda/kernel.cuh b/src/infiniop/ops/layer_norm/cuda/kernel.cuh deleted file mode 100644 index 9381934ee..000000000 --- a/src/infiniop/ops/layer_norm/cuda/kernel.cuh +++ /dev/null @@ -1,157 +0,0 @@ -#ifndef __LAYER_NORM_KERNEL_CUH__ -#define __LAYER_NORM_KERNEL_CUH__ -#include - -template -__device__ void layerNormKernel( - Tdata *output, - Tdata *input_standardization, - Tdata *input_std_deviation, - const Tdata *input, - const Tdata *weight, - const Tdata *bias, - float eps, - size_t normalized_size, - const ptrdiff_t *output_strides, - const ptrdiff_t *input_standardization_strides, - const ptrdiff_t *input_std_deviation_strides, - const ptrdiff_t *input_strides, - ptrdiff_t weight_stride, - ptrdiff_t bias_stride, - bool bias_exist) { - size_t b0 = blockIdx.x, b1 = blockIdx.y; - - auto output_ptr = output + b0 * output_strides[0] + b1 * output_strides[1]; - auto input_ptr = input + b0 * input_strides[0] + b1 * input_strides[1]; - auto standard_ptr = input_standardization + b0 * input_standardization_strides[0] + b1 * input_standardization_strides[1]; - auto std_ptr = input_std_deviation + b0 * input_std_deviation_strides[0] + b1 * input_std_deviation_strides[1]; - Tcompute mean = op::common_cuda::reduce_op::sum( - input_ptr, - normalized_size) - / normalized_size; - Tcompute sum_squared = op::common_cuda::reduce_op::sumSquared( - input_ptr, - normalized_size); - - Tcompute var = sum_squared / normalized_size - mean * mean; - Tcompute std_deviation = sqrtf(var + Tcompute(eps)); - *std_ptr = std_deviation; - - for (size_t d = 0; d < normalized_size; d++) { - Tcompute x_standard = (Tcompute(input_ptr[d]) - mean) / std_deviation; - standard_ptr[d] = x_standard; - output_ptr[d] = x_standard * Tcompute(*(weight + d * weight_stride)) + (bias_exist ? Tcompute(*(bias + d * bias_stride)) : Tcompute(0)); - } -} - -template -__device__ void blockLayernormKernel(T *output, T const *input, T const *weight, T const *bias, float eps, int dimsize, - const ptrdiff_t *output_strides, - const ptrdiff_t *input_strides, - const size_t *shape, - ptrdiff_t weight_stride, - ptrdiff_t bias_stride, - int ndim, - bool bias_exist) { - // 只能处理axis=-1 - int ind_i = 0; // input id - int ind_o = 0; // output id - int tid = blockIdx.x; - for (int j = ndim - 2; j >= 0; j--) { - ind_i += (tid % (int)shape[j]) * (int)input_strides[j]; - ind_o += (tid % (int)shape[j]) * (int)output_strides[j]; - tid = tid / (int)shape[j]; - } - - float mu_partial = op::common_cuda::reduce_op::sum( - input + ind_i, - dimsize) - / dimsize; - __shared__ float mu; - if (threadIdx.x == 0) { - mu = mu_partial; - } // threadIdx.x = 0对应的是全局sum - typedef cub::BlockReduce BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - - __syncthreads(); - float sigma2_partial = 0.0f; - for (int id = threadIdx.x; id < dimsize; id += BLOCK_SIZE) { - sigma2_partial += (static_cast(input[ind_i + id]) - mu) * (static_cast(input[ind_i + id]) - mu); - } - - __shared__ float sigma2; - float sigma2_block = BlockReduce(temp_storage).Sum(sigma2_partial); - if (threadIdx.x == 0) { - float sigma_tmp = sqrt(sigma2_block * __fdividef(1.0F, dimsize) + eps); - sigma2 = __fdividef(1.0F, sigma_tmp); - } - __syncthreads(); - for (int id = threadIdx.x; id < dimsize; id += BLOCK_SIZE) { - output[ind_o + id] = static_cast(static_cast(weight[id * weight_stride]) * (static_cast(input[ind_i + id]) - mu) * sigma2 + (bias_exist ? static_cast(bias[id * bias_stride]) : 0.0f)); - } -} -template -struct SumOp { - __device__ __forceinline__ T operator()(const T &a, const T &b) const { - return a + b; - } -}; - -template