Triton-Ascend 架构说明
核心组件：
Ascend language extension：适配 Ascend 的 Triton 语言扩展compiler：适配 Ascend 的 Triton 编译器driver：适配 Ascend 的设备驱动接口
组件功能：
-
Ascend language extension在标准 Triton 语言基础上，引入针对 Ascend NPU 架构的语法与语义扩展。 -
compiler接收来自上层 Triton compiler 生成的中间表示文件TTIR（Triton IR），执行一系列适配昇腾硬件的转换。Triton IR → Linalg IR → AscendNPU IR → triton_xxx_kernel.o
Triton IR 转换为 Linalg IR，再经 BiSheng Compiler 生成面向 Ascend NPU 的可执行二进制文件
triton_xxx_kernel.o。 -
driver提供 Triton 运行时与 Ascend 软件栈（CANN）之间的对接能力， 加载由 BiSheng Compiler 生成的设备侧可执行内核triton_xxx_kernel.o。
本项目在标准 Triton 基础上，扩展支持华为 Ascend NPU（通过 CANN 软件栈）。整体设计遵循以下代码原则：
- 若修改与目标硬件无关（target independent），应保留在 Triton core 部分（如language、runtime的通用修改）；
- 若修改与 Ascend 硬件强相关（target affinitive），应放在 Triton-Ascend 中。
| 目录或文件 | 对应架构层级 | 功能说明 |
|---|---|---|
python/ |
Triton core | 保留标准 Triton 的 Python 侧通用实现，包括 triton.language、JIT、运行时、缓存、工具链入口等。与硬件无关的通用能力优先放在该目录。 |
include/ 和 lib/ |
Triton core | 保留标准 Triton 的通用 C++/MLIR 基础设施、Dialect、Pass 和转换逻辑。这里不承载 Ascend 专属后端实现。 |
third_party/ascend/ |
Triton-Ascend | Ascend 后端的根目录，集中放置与 Ascend NPU、CANN、BiSheng Compiler 强相关的语言扩展、编译后端、运行时驱动、MLIR Pass、示例和测试。 |
third_party/ascend/language/ |
Ascend language extension | Ascend 语言扩展目录，安装后会链接到 triton.language.extra 下，供 Triton kernel 通过 triton.language.extra.cann 使用。 |
third_party/ascend/language/cann/libdevice.py |
Ascend language extension | 适配 Ascend NPU 的 libdevice Python 接口，提供数学函数和底层算子封装，供 Triton 算子调用。 |
third_party/ascend/backend/compiler.py |
compiler | Ascend 编译器后端主入口，负责注册编译选项、组织 TTIR 到 Ascend 适配 IR、Linalg/LLVM 等阶段的转换，并调用后续工具链生成可执行二进制文件。 |
third_party/ascend/backend/driver.py |
driver | Ascend 运行时驱动模块，负责与 CANN/TorchNPU 等运行时环境对接，加载并启动已编译的设备侧可执行文件。 |
third_party/ascend/include/ 和 third_party/ascend/lib/ |
compiler | Ascend 专属 MLIR Dialect、Pass 和转换实现，例如 TritonToLinalg、TritonToStructured、DynamicCVPipeline、AutoBlockify 等。 |
third_party/ascend/AscendNPU-IR/ |
compiler | Ascend NPU 相关 IR 与 BiSheng 编译链适配内容，是从 Triton-Ascend 编译流程继续下沉到硬件侧代码生成的重要组成部分。 |
third_party/ascend/tutorials/ 和 third_party/ascend/unittest/ |
示例与测试 | 提供 Ascend 平台上的 Triton 示例、迁移样例、Python 单元测试和 MLIR 转换测试，用于验证 Ascend 后端能力。 |
| 序号 | 算子名称 | 描述 |
|---|---|---|
| 1 | tl.insert_slice(full, src, offsets, sizes, strides) |
按照指定的偏移量（offsets）、尺寸（sizes）和步幅（strides）参数，将一个张量插入到另一个张量中。 返回值：目标张量。 full：目标张量，源张量将被插入到此张量中。 src：源张量。 offsets：目标张量上的偏移量（整数元组）。 sizes：源张量上的尺寸（整数元组）。 strides：目标张量上的步幅（整数元组）。 |
| 2 | tl.extract_slice(full, offsets, sizes, strides) |
按照指定的偏移量（offsets）、尺寸（sizes）和步幅（strides）参数，从另一个张量中提取一个切片张量。 返回值：切片张量。 full：源张量，从此张量中提取切片。 offsets：源张量上的偏移量（整数元组）。 sizes：切片张量的尺寸（整数元组）。 strides：源张量上的步幅（整数元组）。 |
| 3 | tl.get_element(source, offset) |
读取一个具有维度的张量，并返回指定偏移量处的单个元素。 source：源张量。 offset：元素提取位置的偏移量（整数元组）。 |
| 序号 | NPUOptions | 硬件平台 | 用途 |
|---|---|---|---|
| 1 | multibuffer | NPU | Autotune Option: Enable or disable ping-pong pipeline. Enabled by default. |
| 2 | enable_auto_bind_sub_block | NPU | Autotune option (CV-fused kernels only): Enable or disable auto-binding of sub-blocks. |
| 3 | enable_hivm_auto_cv_balance | NPU | Autotune option (CV-fused kernels only): Enable or disable automatic CV balancing. |
| 4 | sync_solver | NPU | Autotune option (CV-fused kernels only): Enable or disable the synchronization solver. |
| 5 | unit_flag | NPU | Autotune option: Enable or disable the sync unit flag. |
| 6 | inject_barrier_all | NPU | Autotune option: Enable or disable automatic injection of barriers for all operations. |
| 7 | inject_block_all | NPU | Autotune option: Enable or disable automatic injection of blocks for all operations. |
| 8 | limit_auto_multi_buffer_only_for_local_buffer | NPU | Autotune option: Restrict automatic multi-buffering only to local buffers. |
| 9 | limit_auto_multi_buffer_of_local_buffer | NPU | Autotune option: Enable or disable automatic multi-buffering for local buffers. |
| 10 | set_workspace_multibuffer | NPU | Autotune option: Enable or disable multi-buffering for the workspace. |
| 11 | tile_mix_vector_loop | NPU | Autotune option (CV-fused kernels only): Enable or disable tiling for vector loops. |
| 12 | tile_mix_cube_loop | NPU | Autotune option (CV-fused kernels only): Enable or disable tiling for cube loops. |
| 13 | disable_auto_inject_block_sync | NPU | Autotune option (CV-fused kernels only): Enable or disable automatic injection of block synchronizations. |
| 14 | stream | NPU | Optional: Inform the compiler about the NPU stream to use. |
| 15 | enable_linearize | NPU | Autotune option: Enable or disable the linearization pass. |
| 16 | enable_nd2nz_on_vector | NPU | Autotune option (CV-fused kernels only): Enable or disable the ND (n-dimensional) to NZ (non-zero) layout transformation. |
| 17 | auto_blockify_size | NPU | Autotune option: Enable or disable AutoBlockify pass. It is ignored when TRITON_ALL_BLOCKS_PARALLEL is not set |
| 18 | compile_mode | NPU (950) | Compilation mode: "unstructured_in_simt" (default) / "simd" / "simt_only". |
| 序号 | Pass | 目的 | IR 转换 |
|---|---|---|---|
| 1 | triton-to-structured | linearize | ttir->ttir |
| 2 | triton-to-unstructured | convert indirect axis to loop | ttir->ttir |
| 3 | triton-to-linalg | memory/reduction/view/creation/math/arith/linear algebra to linalgir | ttir->linalgir |
| 4 | triton-to-other | ttir->hivm/hfusion/llvm | ttir->hivm/hfusion/llvm |
处理指针表达式和mask表达式中的整除取余，通过升维的方法，去除整除取余后重新生成load/store 等OP。
| Converter | 功能 | 局限性 |
|---|---|---|
| RewriteAddPtrOp | 分析 tl.load, tl.store等操作中的指针表达式 (AddPtrOp)。将原始的指针偏移计算分解并建模为包含各维度（轴）具体偏移信息的 PtrState 对象。例如，对于形如 ptr + x // 1024 * 4096 + x % 1024 * 4 + y 的表达式，分析出 x 和 y 轴的贡献与关系。 |
1. 所涉及的原始迭代轴（如x）必须能被分裂轴（如1024）整除。2. 外部的 XBLOCK 大小必须是分裂轴divisor的整数倍或其约数。 |
| CreateAddPtr | 根据分析得到的 PtrState 对象，重新构造一个新的 AddPtrOp 指针计算操作。新生成的指针表达式将消除原表达式中的整数除法 (//) 和取模 (%) 操作。 |
依赖于 RewriteAddPtrOp 成功生成的、合法的 PtrState。 |
| RewriteLoadOp | 分析 tl.load 操作中的掩码 (mask) 表达式。将包含整除/取余的复杂掩码条件分解并建模为包含各维度边界信息的 MaskState 对象。例如，对于 mask = x // 1024 < 8 and x % 1024 < 1024 and y < 4，分析出各维度的独立约束条件。 |
1. 所涉及的原始迭代轴（如x）必须能被分裂轴（如1024）整除。2. 外部的 XBLOCK 大小必须是分裂轴divisor的整数倍或其约数。 |
| BuildMask | 根据分析得到的 MaskState 对象，重新构造一个新的掩码 (mask) 表达式。新掩码将消除原表达式中的整数除法 (//) 和取模 (%) 操作。 |
仅处理由 RewriteLoadOp 或 RewriteStoreOp 生成的 MaskState。无法处理任意复杂的、非规范化的掩码表达式。 |
| CreateLoad | 使用由 CreateAddPtr 生成的新指针表达式和由 BuildMask 生成的新掩码表达式，重新创建（替换）原始的 tl.load 操作，完成指令重写。 |
依赖于 RewriteAddPtrOp, CreateAddPtr, RewriteLoadOp, BuildMask 等前置步骤均成功执行。 |
| RewriteStoreOp | 分析 tl.store 操作中的掩码 (mask) 表达式。其功能与 RewriteLoadOp 类似，将包含整除/取余的复杂掩码条件分解并建模为 MaskState 对象。 |
与 RewriteLoadOp 相同。 |
| CreateStore | 使用由 CreateAddPtr 生成的新指针表达式和由 BuildMask 生成的新掩码表达式，重新创建（替换）原始的 tl.store 操作，完成指令重写。 |
依赖于 RewriteAddPtrOp, CreateAddPtr, RewriteStoreOp, BuildMask 等前置步骤均成功执行。 |
| RewriteAtomicRWMOp | 处理原子读写修改操作（如 atomic.add, atomic.max 等）中的指针问题。 |
通常继承与 RewriteAddPtrOp 相同的局限性。对于某些特殊的、非连续或条件性的原子操作模式可能不支持。 |
| RewriteAtomicCASOp | 处理原子比较并交换操作 (atomic.cas) 中的指针线性化问题。分析其指针表达式，通过升维方法消除整除和取余操作，以匹配硬件原子指令的寻址要求。 |
|
| RewriteWhile | 处理 while 循环体内的指针叠加操作。 |
不支持循环体内包含条件分支 (if) 的复杂指针路径变换。 |
| RewriteFor | 处理 for 循环体内的指针叠加操作。 |
| 序号 | Pass / 转换器 | 描述 |
|---|---|---|
| 1 | discrete-mask-access-conversion | 将Triton中基于离散索引掩码（Discrete Mask）的内存访问模式（如triton.language.load带非连续mask）进行分析与转换，为后续将离散轴展开为循环做准备。该Pass识别出那些无法被后端硬件高效处理的、非规律性的或稀疏的访问模式。 |
| 2 | triton-to-unstructured | 将经过discrete-mask-access-conversion识别出的、包含离散轴（Discrete Axes）的张量操作，转换为基于显式标量循环的标量访存。 |
| 3 | bubble-up-operation | 主要对extract op/extract_slice顺序上移优化。这可以优化数据局部性，有些场景能消除转换后产生的不必要的循环，从而提升生成代码的执行效率。 |
| 转换器名称 | 描述 |
|---|---|
| DiscreteMaskStoreConversion | 首先进行mask分析，如果mask分析结果是非连续的，将原始的store操作转化为以下序列： 1. load（加载目标存储地址的内容） 2. select（根据mask挑选目标存储内容和待存储的value内容） 3. store（将select的结果存储回目标地址） |
| DiscreteMaskLoadConversion | 首先进行mask分析，如果mask分析结果是非连续的，将原始的load操作转化为以下序列： 1. load（加载源tensor的所有内容） 2. select（根据mask挑选源tensor内容，被掩盖部分设置为other值） |
| DiscreteMaskAtomicAddConversion | 首先进行mask分析，如果mask分析结果是非连续的，将原始的atomic_add操作转化为以下序列： 1. select（根据mask挑选value的值，被掩盖部分设为0） 2. atomic_add（使用select后的结果重新生成atomic_add操作） |
| TritonToUnstructured Converters | 描述 |
|---|---|
| UnstructuredMemAccessConverter<triton::LoadOp> | 将LoadOp转化为多重循环标量加载 |
| UnstructuredMemAccessConverter<triton::StoreOp> | 将StoreOp转化为多重循环标量存储 |
| UnstructuredMemAccessConverter<triton::AtomicRMWOp> | 将AtomicRMWOp转化为多重循环标量Atomic操作 |
| UnstructuredMemAccessConverter<triton::AtomicCASOp> | 将AtomicCASOp转化为多重循环标量Atomic操作 |
| 转换器名称 | 描述 |
|---|---|
| BubbleUpExtract<tensor::ExtractOp> | extract op顺序上移优化，在某些场景可以避免产生不必要的循环 |
| BubbleUpExtract<tensor::ExtractSliceOp> | extract op/extract_slice顺序上移优化，在某些场景可以避免产生不必要的循环 |
TritonToLinalg converts ttir to linalg ir.
| Converter | 描述 |
|---|---|
| StoreConverter | triton::StoreOp to memref::copy |
| AddPtrConverter | triton::AddPtrOp to memref::ReinterpretCastOp |
| GetProgramIDConverter | triton::GetProgramIdOp to a param in functionOp |
| GetNumProgramsConverter | triton::GetNumProgramsOp to a param in functionOp |
| LoadConverter | triton::LoadOp to memref::copy and bufferization::ToTensorOp |
| AtomicRMWConverter | triton::AtomicRMWOp to linalg::GenericOp |
| AtomicCASConverter | triton::AtomicCASOp to linalg::GenericOp |
| MakeRangeConverter | triton::MakeRangeOp to linalg::GenericOp |
| SplatConverter | triton::SplatOp to linalg::FillOp |
| ClampFConverter | triton::ClampFOp to tensor::EmptyOp, linalg::FillOp |
| PreciseDivConverter | triton::PreciseDivFOp to arith::DivFOp |
| ArgMinConverter | triton::ArgMinOp to linalg::ReduceOp |
| ArgMaxConverter | triton::ArgMaxOp to linalg::ReduceOp |
| ReduceConverter | triton::ReduceOp to linalg::ReduceOp |
| ScanConverter | triton::ScanOp to func::CallOp |
| ReshapeConverter | triton::ReshapeOp to tensor::ReshapeOp |
| ExpandDimsConverter | triton::ExpandDimsOp to tensor::ExpandShapeOp |
| BroadcastConverter | triton::BroadcastOp to linalg::BroadcastOp |
| DenseConstantConverter | arith::ConstantOp to linalg::FillOp |
| ExternElementwiseClOpConverter | triton::ExternElementwiseOp to linalg::MapOp |
| TritonMulhiuiConverter | triton::MulhiUIOp to arith::MulSIExtendedOp |
| TritonPreciseSqrtConverter | triton::PreciseSqrtOp to math::SqrtOp |
| AdvanceConverter | triton::AdvanceOp to memref::ReinterpretCastOp |
| TransposeConverter | triton::TransOp to linalg::TransposeOp |
| SplitConverter | triton::SplitOp to tensor::ExtractSliceOp |
| JoinConverter | triton::JoinOp to tensor::InsertSliceOp |
| CatConverter | triton::CatOp to tensor::InsertSliceOp |
| BitcastConverter | triton::BitcastOp to arith::BitcastOp |
| LoopConverter<scf::ForOp> | scf::ForOp to scf::ForOp |
| LoopConverter<scf::WhileOp> | scf::WhileOp to scf::WhileOp |
| YieldConverter | scf::YieldOp to scf::YieldOp |
| GatherConverter | triton::GatherOp to func::FuncOp |
| GatherLoadConverter | triton::GatherLoadOp to scf::ForOp |
| DeviceAssertConverter | triton::AssertOp to func::FuncOp |
| DevicePrintConverter | triton::PrintOp to func::FuncOp |
| MatmulConverter | triton::DotOp to linalg::MatmulOp |
| SortOpConverter | triton::SortOp to func::FuncOp |
| DotScaledConverter | triton::DotScaledOp to linalg::MatmulOp |
| PtrToIntConverter | triton::PtrToIntOp |
| MakeTensorPtrConverter | triton::PtrToIntOp to arith::IndexCastOp |
| Pass名称 | 功能描述 | 核心转换器 | 转换器描述 |
|---|---|---|---|
| triton-to-annotation | 处理Ascend NPU特有的编译提示指令 (tl.compile_hint)，将其转换为后端的Annotation方言，用于指导后续的硬件特定优化或资源配置。 |
TritonAnnotationConversion | 将 triton::AnnotationOp 转换为 annotation::MarkOp，实现高级编译提示信息向底层注释标记的传递。 |
| triton-to-hfusion | 将Triton中的TTIR转换为Ascend NPU硬件加速器HFusion方言中的对应操作 |
TritonHistogramToHFusionConversion | 将 triton::HistogramOp 转换为 hfusion::HistogramOp，使能在NPU的专用硬件上高效执行。 |
| triton-to-hivm | 处理Triton的块同步操作 (tl.sync_block_all, tl.sync_block_set, tl.sync_block_wait)，将其转换为Ascend NPU的HIVM方言中的跨核心同步指令。这些指令用于管理多核流水线中的同步与数据依赖，是流水优化的关键。 |
TritonCustomOpToHIVMSyncOpConversion | 实现Triton同步指令到HIVM同步指令的转换： • sync_block_all：全局块同步• sync_block_set：设置同步点• sync_block_wait：等待同步点 |
| triton-to-llvm | 将Triton中的内联汇编操作 (tl.inline_assembly) 转换为LLVM方言的内联汇编，并最终映射为Ascend NPU的CCE硬件固有函数（Intrinsics） |
ElementwiseInlineAsmOpConversion | 将 triton::ElementwiseInlineAsmOp 转换为 LLVM::InlineAsmOp |
昇腾 950 在 SIMD 路径之外增加 SIMT 能力，用于加速非结构化 / 离散访存（如间接索引的 load/store）。
开发者通过 compile_mode 选择编译路径。
compile_mode |
含义 | 编译路径 |
|---|---|---|
"simd" |
纯 SIMD：结构化访存走 DMA；非结构化走标量循环 | Triton IR → Linalg IR → AscendNPU IR |
"unstructured_in_simt"（默认） |
混合：结构化仍走 SIMD；离散访存尽量走 SIMT 模板 | Triton IR → Linalg IR → AscendNPU IR |
"simt_only" |
纯 SIMT：直接下发 Triton IR 给NPU IR处理 | Triton IR → AscendNPU IR |
用法示例：
# 纯 SIMD
kernel[grid](..., compile_mode="simd")
# 混合（默认；950 上离散访存优先走 SIMT）
kernel[grid](..., compile_mode="unstructured_in_simt")
# 纯 SIMT
kernel[grid](..., compile_mode="simt_only", num_warps=32)flowchart TD
A[compile_mode] --> B["simd"]
A --> C["unstructured_in_simt"]
A --> D["simt_only"]
%% simt_only 分支
D --> D1[直接下发 Triton IR，纯SIMT编译，Triton IR → AscendNPU IR]
%% simd 完整链路
B --> B1[discrete-mask-access-conversion]
B1 --> B2[拆成连续/离散后用 SIMD 方式处理]
B2 --> B3[triton-to-unstructured]
B3 --> B4[离散访存展开为标量循环]
B4 --> B5[TritonToLinalg]
B5 --> B6[AscendNPU IR]
%% unstructured_in_simt 完整链路
C --> C1[discrete-mask-access-conversion]
C1 --> C2[满足条件打上标记，下发下层SIMT处理]
C2 --> C3[triton-to-unstructured]
C3 --> C4{离散访存可转为 indirect_load/store SIMT 模板？}
C4 -- 是 --> B5
C4 -- 否 --> C5[回退标量循环]
C5 --> B5
%% 样式定义
classDef root fill:#e6f7ff,stroke:#1890ff
classDef pass fill:#fff7e6,stroke:#fa8c16,stroke-width:2px
classDef logic fill:#f0fff4,stroke:#52c41a
classDef simtOnly fill:#f0f2f5,stroke:#8c8c8c
%% 绑定样式
class A root
class B1,C1,B3,C3,B5,B6 pass
class B2,B4,C2,C4,C5 logic
class D,D1 simtOnly
| 阶段 | "simd" |
"unstructured_in_simt" |
"simt_only" |
|---|---|---|---|
| 离散 mask 处理 | 拆成连续/离散边界，用 load + select / store 处理 | Ascend 950 且张量维数 ≤ 5：标记后交给下游；否则同左 | 不运行 |
| 非结构化访存 | 展开为标量循环 | 尽量转为 SIMT 间接访存（维数 ≤ 5）；失败则回退标量循环 | 不运行 |
| TritonToLinalg | 常规 Linalg IR 降级 | 常规 Linalg IR 降级 | 不运行 |
混合模式不会把整个 kernel 切到 SIMT，只对离散 / 非结构化访存点走 SIMT，其余仍走 SIMD：
-
离散 mask 处理
- 若判定为非连续 mask，且满足 Ascend 950、混合模式、维数 ≤ 5：不改写 IR，只标记「下游走 SIMT」。
- 否则（纯 SIMD 或不满足条件）：将 mask 拆成连续 / 离散部分，用连续边界限定全局内存访问，再通过 select 合并结果。
-
非结构化访存处理
- 在 Ascend 950 混合模式下，对非结构化访存或已标记的离散访存走 SIMT 快速通道：
load/store→indirect_load/indirect_store（维数 ≤ 5）atomic操作 →hivm.custom(symbol="__builtin_indirect_atomic")
- 不满足条件则回退为标量循环（与
"simd"一致）。
- 在 Ascend 950 混合模式下，对非结构化访存或已标记的离散访存走 SIMT 快速通道：
-
TritonToLinalg
- 常规 Linalg IR 降级。
"simt_only" 直接下发 Triton IR 交给 AscendNPU IR 做纯 SIMT 编译。
| 序号 | Operator | 功能描述 |
|---|---|---|
| 1 | tl.custom_op | Ascend NPU扩展的自定义算子集，用于支持硬件特定的内存访问与数据搬运模式，例如： • index_select: 基于索引选择数据• index_put: 基于索引放置数据• gather_out_to_ub: 将外部数据收集到Unified Buffer (UB)• scatter_ub_to_out: 将UB中的数据分散输出• indirect_load: 间接地址加载• indirect_store: 间接地址存储 |
| 2 | tl.compile_hint | 向编译器传递硬件特定的编译提示信息，用于指导后端优化策略、资源分配或内核配置。 |
| 3 | tl.sync_block_wait(sender, receiver, event_id) |
块同步等待操作。指定接收核 (receiver) 等待由发送核 (sender) 发出的事件信号 (event_id)，用于管理跨核流水线中的数据依赖与执行顺序。 |
| 4 | tl.sync_block_set(sender, receiver, event_id) |
块同步设置操作。指定发送核 (sender) 向接收核 (receiver) 发出一个事件信号 (event_id)，表明某个执行阶段或数据已准备就绪。 |
| 5 | tl.sync_block_all(mode, event_id) |
全局块同步操作。根据指定的同步模式 (mode)，向所有相关的接收核广播一个事件信号 (event_id)，用于实现全核同步或特定模式的集体同步。 |