Distributed Architecture Design

Terminology and Abbreviations

Compilation and IR

Abbreviation	Full Name	Description
IR	Intermediate Representation	Intermediate representation
AST	Abstract Syntax Tree	Abstract syntax tree
MLIR	Multi-Level Intermediate Representation	Multi-level intermediate representation, an extensible compiler infrastructure in the LLVM ecosystem; both Distributed Dialect and HIVM Dialect are built on top of it
TTIR	Triton IR	Intermediate representation at the Triton dialect level, also the name of a stage in the compilation pipeline
LLVM	(Proper name, no longer corresponds to an acronym expansion)	The target IR / backend infrastructure for the GPU path
HIVM	Ascend backend core IR dialect name	Used as a proper noun in AscendNPU-IR/bishengir code, no official acronym expansion found
DPS	Destination-Passing Style	An IR design pattern that passes computation results through pre-allocated init/output values
CSE	Common Subexpression Elimination	Common subexpression elimination
LICM	Loop-Invariant Code Motion	Loop-invariant code motion
TT_* (e.g., TT_Type/TT_Ptr/TT_IntLike)	Triton Type	Prefix for Triton dialect type constraints in TableGen

Hardware and Execution Units

Abbreviation	Full Name	Description
NPU	Neural Processing Unit	Neural processing unit, referring to Huawei Ascend AI processors in this document
GPU	Graphics Processing Unit	Graphics processing unit
GM	Global Memory	Global memory, device-side globally addressable memory
UB	Unified Buffer	Unified buffer, AICore on-chip temporary buffer
AIC	AI Cube Core	Matrix unit, responsible for matrix operations
AIV	AI Vector Core	Vector unit, responsible for vector computation and low-latency communication
PE	Processing Element	Processing element, corresponding to a rank / process / device participant

Communication Engines and Protocols

Abbreviation	Full Name	Description
SHMEM	Symmetric Hierarchical Memory / Shared Memory	Symmetric memory communication library, referring to shmem/aclshmem symmetric heap communication library in this document
ACLSHMEM	ACL (AscendCL) + SHMEM	Prefix naming for symmetric heap communication interfaces/symbols in this document
MTE	Memory Transfer Engine	Memory transfer engine, data movement engine on Ascend AICore side
SDMA	System Direct Memory Access	System direct memory access
UDMA	UB Direct Memory Access	Unified direct memory access
RDMA	Remote Direct Memory Access	Remote direct memory access
RoCE / ROCE	RDMA over Converged Ethernet	RDMA over converged Ethernet
TLS	Transport Layer Security	Transport layer security protocol
QP	Queue Pair	Queue pair
SQ	Send Queue	Send queue
SQE	Send Queue Element	Send queue element
WQE	Work Queue Element	Work queue element
SGE	Scatter/Gather Element	Scatter/gather element

Other

Abbreviation	Full Name	Description
SIMD	Single Instruction Multiple Data	Single instruction multiple data
SDK	Software Development Kit	Software development kit

---

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.

---

1. Overall Architecture

1.1 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 6.3).

1.2 Architectural Differences from the GPU Path

	GPU	Ascend
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.

---

2. Distributed Dialect

2.1 Dialect Definition

The dependentDialects of Distributed Dialect is deliberately left empty, with no dependencies on any backend dialects. This is the prerequisite for it to lower to both LLVM (GPU path) and HIVM (Ascend path).

• Dialect name: distributed

• C++ namespace: ::mlir::triton::distributed

2.2 Seven Distributed Ops

Op	operands	results	traits
distributed.wait	barrierPtr(TT_PtrLike), numBarriers(TT_IntLike), waitValue(TT_Int), scope, semantic	token(TT_IntLike)	MemoryEffectsOpInterface, TypesMatchWith
distributed.consume_token	input(TT_Type or TT_TensorDescType), token(TT_IntLike)	Same type as input	Elementwise, MemoryEffects, InferTypeOpInterface
distributed.get_rank	axis(I32)	result(I32)	Pure
distributed.get_num_ranks	axis(I32)	result(I32)	Pure
distributed.symm_at	symmAddr(TT_Ptr), rank(I32)	remoteAddr(TT_Ptr, same type)	MemoryEffects, TypesMatchWith
distributed.notify	sigAddr(TT_Ptr), signalVal(I64), rank(I32), sigOp, commScope	None	MemoryEffects
distributed.extern_call	srcs(Variadic), libname/libpath/symbol(StrAttr), pure(BoolAttr)	result(Variadic)	MemoryEffects, ConditionallySpeculatable

2.3 Attribute Enumerations

SignalOp:

Value	Name	Semantics
1	SET	Set the signal value to the specified value
2	ADD	Add the specified value to the current signal value

2.4 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 + Write
ExternCallOp	No side effects when pure=true, otherwise Write + Read

---

3. ConvertTritonDistributedToHIVM Pass

3.1 Pass Overview

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

• Scope: ModuleOp

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

3.2 Driver Logic

The runOnOperation() of the Pass does four 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.

3. Call applyPatternsAndFoldGreedily; other Triton IR should be preserved as-is for downstream single-device compilation stages.

3.3 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

3.4 Attribute Configuration

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

The inference order is: first check the prefix mapping table, if not matched then decide based on existDot.

Symbol Prefix	Core Type
aclshmem_barrier_all	CUBE_AND_VECTOR
aclshmemx_barrier_all_vec	VECTOR
aclshmem_putmem / getmem / putmem_nbi / getmem_nbi / putmem_signal / putmem_signal_nbi	VECTOR
Other (not in prefix table)	existDot ? CUBE_AND_VECTOR : 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.

Other Attributes:

Attribute	Value	Description
PIPE	PIPE_S	See explanation below
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

Note on the current state of PIPE: The symbol-to-pipeline mapping table is declared in the code but currently empty, so all distributed operations actually fall to PIPE_S (scalar pipeline). This is a reserved extension point, not an established “assign pipeline by operation type” design. Please use this as a reference when reading code or doing performance analysis.

3.5 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.

---

4. Compilation Pipeline

4.1 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"】     ← ★ This project's only insertion point ★
        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,
          completed by third-party dependency triton-ascend, internal implementation not expanded here)
  ↓
  Ascend NPU executable kernel

Key Point: This project’s intervention in the compilation pipeline has only one location—inserting the ConvertTritonDistributedToHIVM Pass at the very beginning of the ttadapter stage. After this, all distributed operations have already become hivm.custom, and downstream single-device compilation capabilities are provided by the third-party dependency triton-ascend.

The distributed module is an optional dependency in the compiler—loaded via try / except ImportError, automatically skipping this Pass when distributed support is not built.

4.2 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

std::string prefix = concreteOp.getSymbol();
if (!hasMemrefInArgOrRet()) { prefix = "_mlir_ciface_" + prefix; }
return prefix + callNameMangleSuffix(op);

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.

dl.wait uses a 64-byte hardcoded stride when iterating over multiple barriers because the NPU architecture does not guarantee memory consistency (DataCache) for concurrent writes from multiple cores to the same cacheline (64B). This is the origin of the application-level SIGNAL_SLOT_STRIDE = 64 / sizeof(dtype)—int32 signals are 16 elements, int64 signals are 8 elements.

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.

---

5. Symmetric Heap Runtime

Symmetric Heap: A segment of device memory with consistent size and relative layout across all ranks, where heap offsets are mirror-aligned across ranks. Symmetric Tensor: An ordinary torch tensor allocated from the symmetric heap, which can directly participate in computations and be passed as parameters. The value of both is that—any rank can calculate the address of that tensor on other ranks using “local offset + target PE heap base address” (see 5.3 symm_at), enabling cross-device read/write without runtime address exchange.

5.1 Host-side Initialization

Ascend’s symmetric heap initialization is directly called by user code via the shmem Python package (conventionally import shmem as ash), in five steps:

ash.set_conf_store_tls(False, "")                                  # 1. TLS configuration
attributes = ash.InitAttr()                                        # 2. Construct initialization attributes
attributes.my_rank = rank
attributes.n_ranks = world_size
attributes.local_mem_size = 1024 * 1024 * 1024                     # 3. Symmetric heap size
attributes.ip_port = "tcp://127.0.0.1:8666"                        #    Bootstrap address
attributes.option_attr.data_op_engine_type = ash.OpEngineType.MTE  # 4. Transport engine
ash.aclshmem_init(attributes)                                      # 5. Initialize

Transport Engine Enumeration: MTE / SDMA / ROCE / UDMA, corresponding to four different data paths.

5.2 Symmetric Tensor Allocation

peer_mem = ash.aclshmem_create_tensor(shape, dtype=torch.float16, device_id=rank)
...
ash.aclshmem_free_tensor(peer_mem)

The internal flow is calc_nbytes → aclshmem_malloc → construct_tensor_from_ptr, returning a normal torch tensor that can directly participate in PyTorch operations and kernel parameter passing.

Constraint: aclshmem_malloc / aclshmem_free must be synchronously called by all processes, allocating or freeing the same size of memory.

5.3 Remote Address Resolution of symm_at (Core Mechanism)

Answer: A hybrid scheme of lookup table + offset arithmetic, both are indispensable.

ACLSHMEM_DEVICE __gm__ void *aclshmem_ptr(__gm__ void *ptr, int pe)
{
    __gm__ aclshmem_device_host_state_t *device_state = aclshmemi_get_state();
    ptrdiff_t offset = (uintptr_t)ptr - (uintptr_t)device_state->heap_base;
    uintptr_t remote_ptr = (uintptr_t)device_state->p2p_device_heap_base[pe] + offset;
    return (__gm__ void *)remote_ptr;
}

Three steps:

1. Fetch device state;

2. Subtract the local heap base address from the local pointer, regressing to “symmetric offset”;

3. Add the heap base address of the target PE, obtaining the remote address.

The device state structure has three parallel base address tables—p2p_device_heap_base, rdma_device_heap_base, sdma_device_heap_base, and the runtime selects which one to use based on the transport mode bitmask of the target PE.

---

6. Runtime Implementation of Communication Primitives

6.1 Mechanism Summary Table

Operation	Underlying Mechanism
signal SET (P2P / MTE)	Remote GM ordinary scalar write + cacheline writeback instruction, not hardware atomic, not doorbell
signal ADD (P2P)	Hardware atomic add + UB→GM DataCopy, read-modify-write at the write end
signal (cross-node ROCE)	RDMA write + quiet, or RDMA atomic
wait / wait_until	Local spin busy-wait, invalidate cacheline before reading each round
barrier	Software algorithm (centralized pull), composed of signal_set + spin
UDMA	True doorbell: Fill WQE/SGE to SQ ring → write doorbell register

6.2 signal / notify

The SET path implementation has two steps:

1. Ordinary scalar store *addr = val;

2. dcci_cacheline(addr) —— Flush data cache back to GM, the comment explicitly states its purpose is “flush data cache to GM after signal to ensure it is visible to other ranks”.

The ADD path requires a round trip through UB: set value to UB → set/wait for MTE3 flag → enable hardware atomic add mode → copy_ub2gm → disable atomic mode.

The mechanism of aclshmem_<T>_p (point write) is isomorphic to signal SET: aclshmem_ptr obtains address + scalar store + cacheline writeback.

6.3 Wait Polling Mechanism

Six comparison predicates (_eq / _ne / _gt / _ge / _lt / _le) have identical forms:

do { dcci_cacheline((__gm__ uint8_t *)sig_addr); } while (*sig_addr != cmp_val);

There is no hardware blocking or interrupt mechanism; correctness depends entirely on invalidating the L2 cacheline each loop iteration. dcci_cacheline uses DataCacheCleanAndInvalid internally, with an empty __asm__ __volatile__("") before and after to prevent compiler optimization.

This also explains a characteristic of fine-grained synchronization: dl.wait has non-consuming semantics—it only polls without modifying the signal value, so multiple waiters on the same slot is safe. The constraint falls on the reset side (who is responsible for writing the signal back to 0).

6.4 barrier_all Implementation

aclshmem_barrier_all() ultimately falls to a centralized pull barrier:

1. Each PE only writes its own flag (signal_set);

2. Then directly spin-reads the peer’s flag via aclshmem_ptr(sync_pool, remote_pe)—no need for the peer to actively push;

3. Work is distributed across available Vector Cores, with O(N/K) complexity.

A key detail to prevent deadlock: The barrier-specific wait predicate accepts not only == cmp_val, but also == cmp_val + 1, with a comment explaining it’s “in case when peer pe enters next barrier”—preventing fast PEs that have already entered the next barrier from causing slow PEs to wait forever.

6.5 UDMA

UDMA is not a new set of dl.* operations, but a transport engine selected during host initialization. The same libshmem_device.putmem() call takes completely different underlying paths under MTE and UDMA configurations.

The UDMA device-side send flow is a true doorbell mechanism:

QP fetch → slot calculation → fill SQE → fill SGE → dcci_cachelines → ring doorbell register

---

7. Cross-platform Dispatch Mechanism

7.1 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.

7.2 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__.

7.3 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")

7.4 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 + 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