架构设计与核心特性

1.逻辑架构

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 。

2.代码结构

2.1 代码结构原则

本项目在标准 Triton 基础上，扩展支持华为 Ascend NPU（通过 CANN 软件栈）。整体设计遵循以下代码原则：

若修改与目标硬件无关（target independent），应保留在 Triton core 部分（如language、runtime的通用修改）；

若修改与 Ascend 硬件强相关（target affinitive），应放在 Triton-Ascend 中。

2.2 目录结构与功能说明

目录或文件	对应架构层级	功能说明
`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/torch_npu 等运行时环境对接，加载并启动已编译的设备侧可执行文件。
`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 后端能力。

3. Modules

3.1 Triton core Enhancement

3.1.1 Language expansion

序号	算子名称	描述
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：元素提取位置的偏移量（整数元组）。

3.2 Triton-Ascend

3.2.1 Compiler Options

序号	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"`.

3.2.2 SIMD compiler

序号	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

3.2.2.1 TritonToStructured

处理指针表达式和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`的整数倍或其约数。
CreateAddpr	根据分析得到的 `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	使用由 `CreateAddpr` 生成的新指针表达式和由 `BuildMask` 生成的新掩码表达式，重新创建（替换）原始的 `tl.load` 操作，完成指令重写。	依赖于 `RewriteAddPtrOp`, `CreateAddpr`, `RewriteLoadOp`, `BuildMask` 等前置步骤均成功执行。
RewriteStoreOp	分析 `tl.store` 操作中的掩码 (`mask`) 表达式。其功能与 `RewriteLoadOp` 类似，将包含整除/取余的复杂掩码条件分解并建模为 `MaskState` 对象。	与 `RewriteLoadOp` 相同。
CreateStore	使用由 `CreateAddpr` 生成的新指针表达式和由 `BuildMask` 生成的新掩码表达式，重新创建（替换）原始的 `tl.store` 操作，完成指令重写。	依赖于 `RewriteAddPtrOp`, `CreateAddpr`, `RewriteStoreOp`, `BuildMask` 等前置步骤均成功执行。
RewriteAtomicRWMOp	处理原子读写修改操作（如 `atomic.add`, `atomic.max` 等）中的指针问题。	通常继承与 `RewriteAddPtrOp` 相同的局限性。对于某些特殊的、非连续或条件性的原子操作模式可能不支持。
RewriteAtomicCASOp	处理原子比较并交换操作 (`atomic.cas`) 中的指针线性化问题。分析其指针表达式，通过升维方法消除整除和取余操作，以匹配硬件原子指令的寻址要求。
RewriteWhile	处理 `while` 循环体内的指针叠加操作。	不支持循环体内包含条件分支 (`if`) 的复杂指针路径变换。
RewriteFor	处理 `for` 循环体内的指针叠加操作。

3.2.2.2 TritonToUnstructured

序号	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`顺序上移优化。这可以优化数据局部性，有些场景能消除转换后产生的不必要的循环，从而提升生成代码的执行效率。

3.2.2.2.1 discrete-mask-access-conversion

转换器名称	描述
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操作）

3.2.2.2.2 triton-to-unstructured

TritonToUnstructured Converters	描述
UnstructuredMemAccessConverter<triton::LoadOp>	将LoadOp转化为多重循环标量加载
UnstructuredMemAccessConverter<triton::StoreOp>	将StoreOp转化为多重循环标量存储
UnstructuredMemAccessConverter<triton::AtomicRMWOp>	将AtomicRMWOp转化为多重循环标量Atomic操作
UnstructuredMemAccessConverter<triton::AtomicCASOp>	将AtomicCASOp转化为多重循环标量Atomic操作

3.2.2.2.3 bubble-up-operation

转换器名称	描述
BubbleUpExtract<tensor::ExtractOp>	extract op顺序上移优化，在某些场景可以避免产生不必要的循环
BubbleUpExtract<tensor::ExtractSliceOp>	extract op/extract_slice顺序上移优化，在某些场景可以避免产生不必要的循环

3.2.2.3 TritonToLinalg

3.2.2.3.1 triton-to-linalg

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

3.2.2.4 other passes

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`

3.2.3 SIMT Compiler（Ascend 950）

昇腾 950 在 SIMD 路径之外增加 SIMT 能力，用于加速非结构化 / 离散访存（如间接索引的 load/store）。开发者通过 compile_mode 选择编译路径。

3.2.3.1 `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)

3.2.3.2 三种模式的编译分流

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 降级	不运行

3.2.3.3 混合模式：只对离散访存走 SIMT

混合模式不会把整个 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" 一致）。
TritonToLinalg
- 常规 Linalg IR 降级。

3.2.3.4 纯 SIMT（`simt_only`）

"simt_only" 直接下发 Triton IR 交给 AscendNPU IR 做纯 SIMT 编译。

3.2.4 Ascend affinitive Operators

序号	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`)，用于实现全核同步或特定模式的集体同步。