Distributed Architecture Design

This section introduces the overall architecture of Triton-distributed-ascend on the Ascend platform: from Python distributed primitives to NPU hardware execution, covering the IR levels involved, the components responsible for each, and how the symmetric memory runtime supports cross-device communication.

For definitions of professional terms and abbreviations used in this document, see Glossary.

Overall Architecture

Layered Design

Core Design Philosophy: Distributed Dialect is a platform-agnostic intermediate representation. Each platform maps it to its own communication library through different Conversion Passes, and business-side dl.* code does not need to be rewritten for different platforms.

The landing point for the Ascend platform is the hivm.custom operation—each Distributed Dialect operation is converted into a custom op with a symbol name, where the symbol statically exists in the device-side template library and is resolved by the backend during linking. No runtime module patching is required after compilation, which is fundamentally different from the A/B paths (see Runtime Implementation of Communication Primitives).

Architectural Differences from the GPU Path

	GPU	Ascend NPU
How communication library symbols are bound	After compiling the binary, the host side patches the SHMEM context pointer into the device module	During compilation, hivm.custom symbols are generated and statically resolved during linking
How the device side finds communication state	Through a global variable patched in	Located at a fixed GM address in the architecture, no parameter passing needed
Runtime initialization overhead	Module patching required every time it’s loaded	None

This difference determines that operations like dl.rank() on Ascend can be implemented as pure functions without additional parameters—the address of the device state is a compile-time known constant.

Distributed Dialect

Dialect Definition

See DistributedOps.md for the Distributed Dialect specification.

Memory Side Effect Model

Side effects are not declared in TableGen, but implemented in C++ via MemoryEffectsOpInterface:

Op	Side Effects
WaitOp	Only Read
ConsumeTokenOp	Empty—purely an IR dependency anchor
SymmAtOp	Only Read
NotifyOp	Read and Write
ExternCallOp	No side effects when pure=true, otherwise Write and Read

ConvertTritonDistributedToHIVM Pass

Pass Overview

• Pass name: convert-triton-distributed-to-hivm

• Scope: ModuleOp

• Dependent dialects: hivm::HIVMDialect, triton::distributed::DistributedDialect

Driver Logic

The runOnOperation() of the Pass does two things:

1. Scan for the presence of triton::DotOp / DotScaledOp, obtaining an existDot flag. The absence of tl.dot indicates this is a pure AIV kernel, and this information is used for subsequent core type inference.

2. Register rewrite patterns for the seven distributed ops.

aclshmem Interface OP Generation

The conversion template class generates a symbol name for each operation:

Source op	Corresponding aclshmem interface name
SymmAtOp	aclshmem_ptr_<typename>, e.g., aclshmem_ptr_float
GetRankOp	aclshmem_my_pe
GetNumRanksOp	aclshmem_n_pes
NotifyOp (i32 signal)	aclshmemx_signal_op
NotifyOp (ui64 / i64 signal)	aclshmem_uint64_p / aclshmem_int64_p
ConsumeTokenOp	aclshmem_consume_token_<typename>, e.g., aclshmem_consume_token_float_ptr_1d
WaitOp	aclshmem_wait_<typename>, e.g., aclshmem_wait_int32
ExternCallOp	Uses the op’s own symbol attribute

Attribute Configuration

TCoreType (determines which type of core this operation executes on):

The Triton-distributed Python frontend annotates each externalCallOp with a distributed.core_type string attribute, which is used to derive the hivm::TCoreType attribute consumed directly by downstream passes.

distributed.core_type	hivm::TCoreType
cube	CUBE
vector	VECTOR
cube_and_vector	CUBE_AND_VECTOR
cube_or_vector	Currently inferred from existDot

For an externalCallOp, if the distributed.core_type attribute is missing or has an unrecognized value, the pass reports an error immediately. This ensures that the Python frontend must be updated together with the pass; otherwise, the pass could infer an incorrect TCoreType.

The annotation is currently emitted by libaclshmem_device primitives. A distributed.extern_call produced through any other path must also specify distributed.core_type; otherwise, it cannot pass this conversion.

The core type for each primitive is shown below:

Python primitive	distributed.core_type
aclshmem_barrier_all / aclshmem_barrier	cube_and_vector
aclshmemx_barrier_all_vec / aclshmemx_barrier_vec	vector
rma & udma apis	vector
others	cube_or_vector

The marking of aclshmem_barrier_all as CUBE_AND_VECTOR is important: it means the barrier exists in both functions simultaneously after the kernel is split into cube func and vector func, which is the mechanism that enables computation-communication overlap.

distributed.core_type is a frontend-only annotation: once the op has been lowered to hivm.custom, the pass removes the attribute, leaving only the resolved hivm.tcore_type.

Other Attributes:

Attribute	Value	Description
PIPE	PIPE_S	Set unconditionally
VFMode	SIMD	Set unconditionally
hivm.is_distributed	UnitAttr	Downstream marker for identifying distributed operations in multiple subsequent stages such as memory scope inference, block pointer analysis, and mix kernel splitting
symbol	StrAttr	Bare symbol name
no_side_effect	UnitAttr (conditional)	Set only when the source op’s memory side effects are empty
gm_addr_args_indices	DenseI32Array	Collects the operand indices that are both pointer types and entry block arguments of tt.func. Tensor pointers and results of tt.addptr are not included

Result Construction and Replacement

For each RankedTensorType result, first create a tensor::EmptyOp as the init value for DPS (destination-passing style), then construct hivm::CustomOp. Finally: if no results, eraseOp; if results, replaceOp.

Compilation Pipeline

Complete Pipeline and Stage Entry Points

  Python AST → Triton IR (generate distributed.* operations using DistributedOpBuilder)
  │     builder bindings: create_distributed_wait / consume_token / get_rank /
  │                       get_num_ranks / symm_at / notify / extern_call
  ↓
【stage "ttir"】
        Standard Triton optimizations: inliner / combine / canonicalize / cse / licm / loop_unroll
  ↓
【stage "ttadapter"】
        add_convert_triton_distributed_to_hivm(pm)
        —— Executes before all subsequent passes of single-device compilation,
           distributed operations have already become hivm.custom before entering TTIR→Linalg conversion
  ↓
【stage "npubin"】
        Single-device NPU compilation and binary generation
        (TTIR → Linalg IR → AscendNPU IR → machine code)
  ↓
  Ascend NPU executable kernel

How Symbols Become Real Code (AscendNPU-IR)

This is the final link in understanding the entire architecture: the Pass only generates a string symbol name, how does it become executable code?

Step 1: Symbol Name Mangling

The mlir_ciface prefix is added; that is, symbol = "aclshmem_my_pe" → link name _mlir_ciface_aclshmem_my_pe.

Step 2: Static Implementation in Template Library

Every symbol generated by the Pass has a corresponding _mlir_ciface_ implementation in the device-side template library. A few representative examples:

• _mlir_ciface_aclshmem_my_pe / _mlir_ciface_aclshmem_n_pes —— Directly read device state

• ACLSHMEM_PTR_WRAPPER macro —— Expands to 12 type variants of aclshmem_ptr_<T>, wrapping the result of aclshmem_ptr() back into a memref structure

• ACLSHMEM_P_WRAPPER macro —— Expands to aclshmem_int64_p / uint64_p, etc.

• ACLSHMEM_WAIT_WRAPPER macro —— Expands to aclshmem_wait_int32/64, etc.

The implementation body of CONSUME_TOKEN_*_WRAPPER simply returns the input as-is, serving as a barrier to create data dependencies for the compiler, with no runtime overhead.

Step 3: Compile to Bitcode and Link

The template library source code is compiled into four variants (aic / aiv / mix_aiv / mix_aic) of bitcode, which are linked in when generating the final kernel binary.

This is the complete answer to “why the Pass only needs to generate a string symbol name”: the symbol statically exists in the template library and is resolved by the backend during linking.

Cross-platform Dispatch Mechanism

Dispatch Structure

dl.wait / dl.notify / dl.symm_at / dl.rank / dl.num_ranks / dl.consume_token
   → Static, no proxy → _builder.create_* → distributed dialect → (Ascend) HIVM Pass
   ※ Only backend conditional branch: dl.notify dtype validation

libshmem_device.<f>  → ModuleProxy dynamic dispatch
   → (Ascend) libaclshmem_device.<f> → extern_call → ExternCallOp → aclshmem_* symbol

Core dl.* primitives take the static path—they directly create MLIR operations, with platform differences entirely handled by Conversion Passes; the Python layer has no branches. This is the foundation of cross-platform portability.

libshmem_device.* takes the dynamic path—because different platforms have non-identical SHMEM library function sets, requiring runtime platform-based module selection.

ModuleProxy

The only dispatch primitive is ModuleProxy, with simple logic:

1. During construction, evaluate [(predicate, module), ...] one by one, asserting exactly one predicate is true;

2. __getattr__ transparently forwards attribute access to the selected module;

3. The dispatch decorator makes the decorated function look up the same-named function in the active module at call time by func.__name__.

Backend Information Retrieval

Backend determination is entirely based on shutil.which(<vendor smi tool>):

Function	Criterion
is_ascend()	shutil.which("npu-smi")
is_cuda()	shutil.which("nvidia-smi")
is_hip()	shutil.which("rocm-smi")
is_maca()	shutil.which("mx-smi")

Ascend Branch Checklist

Location	What the Ascend Branch Does
libshmem_device.py	Adds ascend.libaclshmem_device to the proxy list
language_extra.py	Adds ascend.language_extra to the proxy list (currently empty module)
distributed_ops.py	dl.notify dtype validation: Ascend allows int32 (set and add), int64/uint64 set only; other platforms only allow int64/uint64
utils.py	Ascend branch empty during import, does not import any vendor SDK
Compiler	Inserts HIVM Pass in the ttadapter stage
pybind layer	Mounts ascend_passes submodule when building Ascend version