# 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

![Overall Architecture Diagram](./images/architecture.png)

**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**

```cpp
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:

```python
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

```python
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.**

```cpp
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:

```cpp
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 |