architecture_design_and_core_features.md

Architecture Design and Core Features

1. Logical Architecture

Triton-Ascend Architecture Description

Core components:

• Ascend language extension: Triton language extension for Ascend
• compiler: Triton compiler for Ascend
• driver: Ascend device driver API

Component functions:

• Ascend language extension Syntax and semantic extensions for the Ascend NPU architecture are introduced based on the standard Triton language.

• compiler Receives the Triton IR (TTIR) file generated by the upper-layer Triton compiler and performs a series of transformations to adapt to Ascend hardware.

  Triton IR → Linalg IR → AscendNPU IR → triton_xxx_kernel.o

  Converts the Triton IR into the Linalg IR, and then generates the executable binary file triton_xxx_kernel.o for the Ascend NPU through the BiSheng Compiler.

• driver Provides the interconnection capability between the Triton runtime and the Ascend software stack (CANN), and loads the executable kernel file triton_xxx_kernel.o generated by the BiSheng Compiler on the device side.

2. Code Structure

2.1 Code Structure Principles

This project extends the support for Huawei Ascend NPU (using the CANN software stack) based on the standard Triton. The overall design complies with the following code principles:

• If the modification is target independent, it should be retained in the Triton core part (such as general modifications to the language and runtime).
• If the modification is target affinitive, it should be placed in the Triton-Ascend part.

2.2 Directory Structure and Function Description

Directory or File	Architecture Layer	Description
python/	Triton core	Contains the common Python implementation from standard Triton, including triton.language, JIT, runtime, cache, and tool entry points. Target-independent capabilities should live here first.
include/ and lib/	Triton core	Contains the common C++/MLIR infrastructure, dialects, passes, and conversion logic from standard Triton. Ascend-specific backend code is not placed here.
third_party/ascend/	Triton-Ascend	Root directory of the Ascend backend. It contains Ascend NPU, CANN, and BiSheng Compiler-specific language extensions, compiler backend, runtime driver, MLIR passes, examples, and tests.
third_party/ascend/language/	Ascend language extension	Contains Ascend language extensions. During installation, this directory is linked under triton.language.extra, so Triton kernels can use triton.language.extra.cann.
third_party/ascend/language/cann/libdevice.py	Ascend language extension	Provides the Ascend NPU-adapted Python libdevice interface, including math functions and low-level operator wrappers used by Triton kernels.
third_party/ascend/backend/compiler.py	compiler	Main entry of the Ascend compiler backend. It registers compiler options, organizes TTIR lowering to Ascend-adapted IR, Linalg, LLVM, and related stages, and invokes the downstream toolchain to generate executable binaries.
third_party/ascend/backend/driver.py	driver	Ascend runtime driver module. It connects Triton runtime with CANN/torch_npu runtime environments and launches compiled device-side executables.
third_party/ascend/include/ and third_party/ascend/lib/	compiler	Contains Ascend-specific MLIR dialects, passes, and conversions, such as TritonToLinalg, TritonToStructured, DynamicCVPipeline, and AutoBlockify.
third_party/ascend/AscendNPU-IR/	compiler	Contains Ascend NPU IR and BiSheng compilation-chain integration used by the Triton-Ascend pipeline when lowering further toward hardware code generation.
third_party/ascend/tutorials/ and third_party/ascend/unittest/	Examples and tests	Provides Triton examples on Ascend, migration examples, Python unit tests, and MLIR conversion tests for validating Ascend backend capabilities.

3. Modules

3.1 Triton Core Enhancement

3.1.1 Language Extension

No.	Operator	Description
1	tl.insert_slice(full, src, offsets, sizes, strides)	Inserts a tensor into another tensor according to the specified offset, size, and stride. Returns: target tensor. full: target tensor. The source tensor will be inserted into this tensor. src: source tensor. offsets: offset (integer tuple) on the target tensor. sizes: size (integer tuple) on the source tensor. strides: stride (integer tuple) on the target tensor.
2	tl.extract_slice(full, offsets, sizes, strides)	Extracts a slice tensor from another tensor according to the specified offset, size, and stride. Returns: slice tensor. full: source tensor. The slice is extracted from this tensor. offsets: offset (integer tuple) on the source tensor. sizes: size (integer tuple) of the slice tensor. strides: stride (integer tuple) on the source tensor.
3	tl.get_element(source, offset)	Reads a tensor with dimensions and returns a single element at the specified offset. source: source tensor. offset: offset (integer tuple) of the element to be extracted.

3.2 Triton-Ascend

3.2.1 Compiler Options

No.	NPU Option	Hardware Platform	Description
1	multibuffer	NPU	Autotune option. It enables or disables the ping-pong pipeline. Enabled by default.
2	enable_auto_bind_sub_block	NPU	Autotune option (CV-fused kernels only). It enables or disables auto-binding of sub-blocks.
3	enable_hivm_auto_cv_balance	NPU	Autotune option (CV-fused kernels only). It enables or disables automatic CV balancing.
4	sync_solver	NPU	Autotune option (CV-fused kernels only). It enables or disables the synchronization solver.
5	unit_flag	NPU	Autotune option. It enables or disables the sync unit flag.
6	inject_barrier_all	NPU	Autotune option. It enables or disables automatic injection of barriers for all operations.
7	inject_block_all	NPU	Autotune option. It enables or disables automatic injection of blocks for all operations.
8	limit_auto_multi_buffer_only_for_local_buffer	NPU	Autotune option. It restricts automatic multi-buffering only to local buffers.
9	limit_auto_multi_buffer_of_local_buffer	NPU	Autotune option. It enables or disables automatic multi-buffering for local buffers.
10	set_workspace_multibuffer	NPU	Autotune option. It enables or disables multi-buffering for the workspace.
11	tile_mix_vector_loop	NPU	Autotune option (CV-fused kernels only). It enables or disables tiling for vector loops.
12	tile_mix_cube_loop	NPU	Autotune option (CV-fused kernels only). It enables or disables tiling for cube loops.
13	disable_auto_inject_block_sync	NPU	Autotune option (CV-fused kernels only). It enables or disables automatic injection of block synchronizations.
14	stream	NPU	(Optional) Informs the compiler about the NPU stream to use.
15	enable_linearize	NPU	Autotune option. It enables or disables the linearization pass.
16	enable_nd2nz_on_vector	NPU	Autotune option (CV-fused kernels only). It enables or disables the ND (n-dimensional) to NZ (non-zero) layout transformation.
17	auto_blockify_size	NPU	Autotune option. It enables or disables AutoBlockify pass. It is ignored when TRITON_ALL_BLOCKS_PARALLEL is not set

3.2.2 SIMD Compiler

No.	Pass	Purpose	IR Conversion
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

The integer division and modulo operations in pointer expressions and mask expressions are converted to tensor operations to regenerate OPs such as load and store.

Converter	Description	Limitation
RewriteAddPtrOp	Analyzes the pointer expression (AddPtrOp) in operations such as tl.load and tl.store. The original pointer offset is calculated and modeled into a PtrState object that contains the specific offset information of each dimension (axis). For example, for an expression like ptr + x // 1024 * 4096 + x % 1024 * 4 + y, the contributions and relationships of the x and y axes are analyzed.	1. The original iteration axis (such as x) must be exactly divisible by the split axis (such as 1024). 2. The size of the external XBLOCK must be an integer multiple or divisor of the split axis divisor.
CreateAddPtr	Reconstructs a new AddPtrOp pointer calculation operation based on the analyzed PtrState object to eliminate the integer division (//) and modulo (%) operations in the original expression.	It depends on the valid PtrState object successfully generated by RewriteAddPtrOp.
RewriteLoadOp	Analyzes the mask expression in the tl.load operation and decomposes and models the complex mask conditions involving integer division and modulo into a MaskState object that contains the boundary information of each dimension. For example, for mask = x // 1024 < 8 and x % 1024 < 1024 and y < 4, an independent constraint condition of each dimension is obtained through analysis.	1. The original iteration axis (such as x) must be exactly divisible by the split axis (such as 1024). 2. The size of the external XBLOCK must be an integer multiple or divisor of the split axis divisor.
BuildMask	Reconstructs a new mask expression based on the analyzed MaskState object to eliminate the integer division (//) and modulo (%) operations in the original expression.	Only the MaskState object generated by RewriteLoadOp or RewriteStoreOp is processed. Any complex and non-standardized mask expressions cannot be processed.
CreateLoad	Re-creates/Replaces the original tl.load operation using the new pointer expression generated by CreateAddPtr and the new mask expression generated by BuildMask to complete instruction rewriting.	The prerequisite steps, such as RewriteAddPtrOp, CreateAddPtr, RewriteLoadOp, BuildMask, must have been successfully executed.
RewriteStoreOp	Analyzes the mask expression in the tl.store operation. Its function is similar to that of RewriteLoadOp. It decomposes and models the complex mask conditions involving integer division and modulo into a MaskState object.	Same as RewriteLoadOp.
CreateStore	Re-creates/Replaces the original tl.store operation using the new pointer expression generated by CreateAddPtr and the new mask expression generated by BuildMask to complete instruction rewriting.	The prerequisite steps, such as RewriteAddPtrOp, CreateAddPtr, RewriteStoreOp, BuildMask, must have been successfully executed.
RewriteAtomicRWMOp	Handles pointer issues in atomic read/write/modify operations (such as atomic.add and atomic.max).	Generally, it inherits the same limitations as RewriteAddPtrOp. Some special, non-contiguous, or conditional atomic operation modes may not be supported.
RewriteAtomicCASOp	Handles pointer linearization issues in atomic compare-and-swap operations (such as atomic.cas) and analyzes the pointer expression to eliminate the integer division and modulo operations by tensorization to meet the addressing requirements of hardware atomic instructions.
RewriteWhile	Handles pointer overlapping operations in the while loop body.	Complex pointer path transformation that contains conditional branches (if) in the loop body is not supported.
RewriteFor	Handles pointer overlapping operations in the for loop body.

3.2.2.2 TritonToUnstructured

No.	Pass/Converter	Description
1	discrete-mask-access-conversion	Analyzes and converts the memory access pattern based on the discrete index mask in Triton (for example, triton.language.load with a non-contiguous mask) to prepare for the subsequent expansion of discrete axes into loops. This pass identifies irregular or sparse access patterns that cannot be efficiently processed by the backend hardware.
2	triton-to-unstructured	Converts the tensor operations identified by discrete-mask-access-conversion and containing discrete axes into scalar memory access based on explicit scalar loops.
3	bubble-up-operation	Bubbles up extract op/extract_slice for optimization. This can optimize data locality. In some scenarios, unnecessary loops generated after the transformation can be eliminated, thereby improving the execution efficiency of the generated code.

3.2.2.2.1 discrete-mask-access-conversion

Converter	Description
DiscreteMaskStoreConversion	Analyzes the mask and converts the original store operation into the following sequences if it is non-contiguous: 1. load (loading the content at the target address) 2. select (selecting the target content and the value to be stored based on the mask) 3. store (storing the select result back to the target address)
DiscreteMaskLoadConversion	Analyzes the mask and converts the original load operation into the following sequences if it is non-contiguous: 1. load (loading all content of the source tensor) 2. select (selecting the source tensor content based on the mask, and setting the masked part to the other value)
DiscreteMaskAtomicAddConversion	Analyzes the mask and converts the original atomic_add operation into the following sequences if it is non-contiguous: 1. select (selecting the value based on the mask and setting the masked part to 0) 2. atomic_add (regenerating the atomic_add operation using the select result)

3.2.2.2.2 triton-to-unstructured

TritonToUnstructured Converter	Description
UnstructuredMemAccessConverter<triton::LoadOp>	Converts LoadOp into a multi-loop scalar load operation.
UnstructuredMemAccessConverter<triton::StoreOp>	Converts StoreOp into a multi-loop scalar store operation.
UnstructuredMemAccessConverter<triton::AtomicRMWOp>	Converts AtomicRMWOp into a multi-loop scalar atomic operation.
UnstructuredMemAccessConverter<triton::AtomicCASOp>	Converts AtomicCASOp into a multi-loop scalar atomic operation.

3.2.2.2.3 bubble-up-operation

Converter	Description
BubbleUpExtract<tensor::ExtractOp>	Bubbles up extract op, avoiding unnecessary loops in some scenarios.
BubbleUpExtract<tensor::ExtractSliceOp>	Bubbles up extract op/extract_slice, avoiding unnecessary loops in some scenarios.

3.2.2.3 TritonToLinalg

3.2.2.3.1 triton-to-linalg

TritonToLinalg converts ttir to linalg ir.

Converter	Description
StoreConverter	triton::StoreOp to memref::copy
AddPtrConverter	triton::AddPtrOp to memref::ReinterpretCast
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
IndexSelectSimdConverter	triton::ascend::IndexSelectSimdOp 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 to memref::ExtractAlignedPointerAsIndexOp + arith::IndexCastOp
MakeTensorPtrConverter	triton::MakeTensorPtrOp to memref::ReinterpretCastOp

3.2.2.4 other passes

Pass	Description	Core Converter	Description
triton-to-annotation	Converts Ascend NPU-specific compilation hint (tl.compile_hint) into backend annotation dialects, which are used to guide subsequent hardware-specific optimization or resource configuration.	TritonAnnotationConversion	Converts triton::AnnotationOp into annotation::MarkOp to transfer advanced compilation hints to the underlying annotation marks.
triton-to-hfusion	Converts TTIR in Triton into the corresponding operation in the HFusion dialect of the Ascend NPU hardware accelerator.	TritonHistogramToHFusionConversion	Converts triton::HistogramOp into hfusion::HistogramOp to enable efficient execution on the dedicated NPU hardware.
triton-to-hivm	Processes the block synchronization operations (tl.sync_block_all, tl.sync_block_set, and tl.sync_block_wait) of Triton and converts them into the cross-core synchronization instruction in the HIVM dialect of Ascend NPU. These instructions are used to manage synchronization and data dependencies in the multi-core pipeline, which is the key to pipeline optimization.	TritonCustomOpToHIVMSyncOpConversion	Converts Triton synchronization instructions to HIVM synchronization instructions. • sync_block_all: synchronizes blocks globally. • sync_block_set: sets a synchronization point. • sync_block_wait: waits for a synchronization point.
triton-to-llvm	Converts the inline assembly operation (tl.inline_assembly) in Triton to the inline assembly in the LLVM dialect, and finally maps it to a CCE hardware intrinsic function of Ascend NPU.	ElementwiseInlineAsmOpConversion	Converts triton::ElementwiseInlineAsmOp to LLVM::InlineAsmOp.

3.2.3 Ascend affinitive Operators

No.	Operator	Description
1	tl.custom_op	A set of custom operators extended by Ascend NPU, used to support hardware-specific memory access and data movement patterns. For example: • index_select: selects data based on an index. • index_put: places data based on an index. • gather_out_to_ub: collects external data to the unified buffer (UB). • scatter_ub_to_out: scatters data from the UB to the output. • indirect_load: loads content from an indirect address. • indirect_store: stores content to an indirect address.
2	tl.compile_hint	Provides hardware-specific compilation hints to the compiler, which are used to guide the backend optimization policy, resource allocation, or kernel configuration.
3	tl.sync_block_wait(sender, receiver, event_id)	Waits for block synchronization. The receiver waits for the event signal (event_id) sent by the sender, which is used to manage data dependencies and execution sequence in the cross-core pipeline.
4	tl.sync_block_set(sender, receiver, event_id)	Sets block synchronization. The sender sends an event signal (event_id) to the receiver, indicating that an execution phase or data is ready.
5	tl.sync_block_all(mode, event_id)	Globally synchronizes blocks. The sender broadcasts an event signal (event_id) to all related receivers according to the specified synchronization mode (mode) to implement full-core synchronization or collective synchronization in a specific mode.