triton-lang/triton-ext

GitHub: triton-lang/triton-ext

为 Triton 编译器提供树外扩展开发框架,支持在不修改核心代码库的前提下动态加载自定义 MLIR pass、dialect、后端和语言扩展。

Stars: 30 | Forks: 9

# Triton Extensions build status discord chat 一个为 Triton 编译器提供的树外扩展(out-of-tree extensions)集合,包括 pass、dialect、backend 以及语言扩展。 ## 概述 本代码库提供了一个框架,用于开发和构建 Triton 编译器扩展,这些扩展无需修改核心 Triton 代码库即可扩展功能。扩展被构建为共享库,可以在运行时由 Triton 动态加载。 扩展构建于 [上游][triton-plugins] 文档中记载的 Triton 插件基础设施之上。 2026 年 1 月 Triton 社区Meetup 的幻灯片: [Triton Community Meetup: Triton Extensions Jan 2026](https://docs.google.com/presentation/d/1dnm8uhvabdwqsQAsaPM7IRpEh2tktQ91E1d40r91n1M/edit?usp=sharing) ### 结构 每个子目录的 `CMakeLists.txt` 负责构建其各自的扩展。 - **[`backend/`](./backend/)**: 用于实现 backend 扩展 (目前仅为脚手架)。 - **[`dialect/`](./dialect/)**: 用于自定义 MLIR dialect 扩展 (目前仅为脚手架)。 - **[`language/`](./language/)**: 用于语言扩展的实现 (目前仅为脚手架)。 - **[`pass/`](./pass/)**: 包含 MLIR pass 扩展。每个 pass 扩展都 实现为可以动态加载的共享库。Pass 扩展包含一个 `triton-ext.toml` 清单,用于指定扩展 名称和状态。 - **[`extensions/`](./extensions/)**: 包含独立的插件扩展,将 dialect、pass 和语言绑定打包为自包含的共享库,可由 Triton 在运行时加载。 - **[`utlx` (µTLX)](./extensions/utlx/)**: Triton 语言扩展插件, 提供了 Meta 的大部分 TLX 功能,而无需修改 Triton fork。包含本地内存操作(`local_alloc`、`local_view`、 `local_store`、`local_load`、`alloc_barriers`)、自定义 pass(例如 PingPong、PruneUnusedBarriers)、TLX dialect、转换模式以及 Python DSL。构建生成 `libutlx.so`。 - **[`support/`](./support/)**: 包含扩展基础设施代码,用于 向 Triton 自动注册扩展。 ## 前置条件 - 支持 C++17 的 C++ 编译器 - CMake - GitHub CLI ([`gh`]),用于下载预构建的依赖项(可选) - Ninja - Python 3,用于测试和构建脚本;使用以下命令安装依赖 `pip install -r requirements.txt` - Triton,使用 `TRITON_EXT_ENABLED=ON` 构建 注意:在 3.7 及以上版本的 Triton 发行版中,默认已启用扩展。 ## 构建 此扩展代码库设计为支持树外构建。它需要指定 LLVM (`LLVM_INSTALL_DIR`) 和 Triton (`TRITON_INSTALL_DIR`) 的路径。 这两个目录都应为安装目录,即先构建,然后 通过 `make install` 打包。 要构建扩展: 1. **构建 LLVM**:将 LLVM 构建为共享库并将其安装到一个已知 位置;参见 CI [action][build-llvm] 作为参考。或者, 从 GitHub 下载预构建的 LLVM 二进制文件:运行 `ci/download-artifact.py llvm` [^list-artifacts]。 2. **构建 Triton**:构建 Triton 并将其安装到一个已知 位置;参见 CI [action][build-triton] 作为参考。或者, 从 GitHub 下载预构建的 Triton 二进制文件:运行 `ci/download-artifact.py triton` [^list-artifacts]。 \[^list-artifacts\]: GitHub artifacts 仅对有限范围内的 commit、操作系统和硬件架构可用。要列出可用的 artifacts,请运行 `ci/fetch-artifacts.py`。 1. **构建扩展**: export LLVM_INSTALL_DIR=/path/to/llvm/install export TRITON_INSTALL_DIR=/path/to/triton/install make build 请注意,如果未设置 `LLVM_INSTALL_DIR` 和 `TRITON_INSTALL_DIR`, `Makefile` 将会自动在项目目录中[搜索][search]它们。参见 CI [workflow](./.github/workflows/ci.yml) 作为参考。扩展被构建为 位于 `build/lib` 下的共享库。 ## 测试 运行测试套件以验证扩展是否正常工作: ``` make test ``` ## 使用 Triton 在运行时使用 `TRITON_PLUGIN_PATHS` 环境变量加载扩展(参见 [Triton 插件][triton-plugins]): ``` export TRITON_PLUGIN_PATHS=/path/to/libmy_pass.so python your_script.py ``` 某些扩展可以从 Python 中访问:例如,µTLX 插件在被导入时会自动 注册为 `triton.language.extra.tlx`,因此不需要文件系统 软链接: ``` export TRITON_PLUGIN_PATHS=/path/to/triton-ext/build/lib/libutlx.so python your_script.py ``` 要加载多个插件,请使用 `:` 分隔路径: ``` export TRITON_PLUGIN_PATHS=build/lib/libutlx.so:build/lib/libother_plugin.so ```
标签:Bash脚本, MLIR, SOC Prime, Triton, 开发工具, 编译器, 编译器扩展, 逆向工具