# 分布式架构设计

## 术语缩写表

**编译与 IR**

| 缩写 | 英文全称 | 中文说明 |
| --- | --- | --- |
| IR | Intermediate Representation | 中间表示 |
| AST | Abstract Syntax Tree | 抽象语法树 |
| MLIR | Multi-Level Intermediate Representation | 多级中间表示，LLVM 生态下可扩展的编译器基础设施，Distributed Dialect / HIVM Dialect 均构建于其上 |
| TTIR | Triton IR | Triton 方言层面的中间表示，也是编译流程中一个 stage 的名字 |
| LLVM | （项目专名，不再对应缩写展开） | GPU 路径的下沉目标 IR / 后端基础设施 |
| HIVM | Ascend 后端核心 IR 方言名 |  `AscendNPU-IR`/`bishengir` 代码中作为专有名词使用，暂未见官方缩写全称展开 |
| DPS | Destination-Passing Style | 目标传递风格：通过预先分配的 init/输出值来传递计算结果的 IR 设计模式 |
| CSE | Common Subexpression Elimination | 公共子表达式消除 |
| LICM | Loop-Invariant Code Motion | 循环不变代码外提 |
| `TT_*`（如 `TT_Type`/`TT_Ptr`/`TT_IntLike`） | Triton Type | TableGen 中 Triton 方言类型约束的前缀 |

**硬件与执行单元**

| 缩写 | 英文全称 | 中文说明 |
| --- | --- | --- |
| NPU | Neural Processing Unit | 神经网络处理器，本文指华为昇腾 AI 处理器 |
| GPU | Graphics Processing Unit | 图形处理器 |
| GM | Global Memory | 全局内存，Device 侧全局可寻址内存 |
| UB | Unified Buffer | 统一缓冲区，AICore 片上临时缓冲区 |
| AIC | AI Cube Core | 矩阵单元，负责矩阵运算 |
| AIV | AI Vector Core | 向量单元，负责向量计算与低时延通信 |
| PE | Processing Element | 处理单元，对应一个 rank / 进程 / 设备参与方 |

**通信引擎与协议**

| 缩写 | 英文全称 | 中文说明 |
| --- | --- | --- |
| SHMEM | Symmetric Hierarchical Memory / Shared Memory | 对称内存通信库，本文指 `shmem`/`aclshmem` 对称堆通信库 |
| ACLSHMEM | ACL（AscendCL）+ SHMEM | 本文对称堆通信接口/符号的前缀命名 |
| MTE | Memory Transfer Engine | 内存传输引擎，昇腾 AICore 侧数据搬运引擎 |
| SDMA | System Direct Memory Access | 系统直接内存访问 |
| UDMA | UB Direct Memory Access | 统一直接内存访问 |
| RDMA | Remote Direct Memory Access | 远程直接内存访问 |
| RoCE / ROCE | RDMA over Converged Ethernet | 基于融合以太网的 RDMA |
| TLS | Transport Layer Security | 传输层安全协议 |
| QP | Queue Pair | 队列对 |
| SQ | Send Queue | 发送队列 |
| SQE | Send Queue Element | 发送队列元素 |
| WQE | Work Queue Element | 工作队列元素 |
| SGE | Scatter/Gather Element | 分散/聚合元素 |

**其他**

| 缩写 | 英文全称 | 中文说明 |
| --- | --- | --- |
| SIMD | Single Instruction Multiple Data | 单指令多数据 |
| SDK | Software Development Kit | 软件开发工具包 |

---

本章节介绍 Triton-distributed-ascend 在 Ascend 平台上的整体架构：从 Python 分布式原语到 NPU 硬件执行，中间经历了哪些 IR 层次、由哪些组件负责、以及对称内存运行时如何支撑跨卡通信。

---

## 一、整体架构

### 1.1 分层设计

![整体架构图](./images/architecture.png)

**核心设计理念**：Distributed Dialect 是**平台无关的中间表示**。各平台通过不同的 Conversion Pass 将其映射到各自的通信库，业务侧的 `dl.*` 代码不需要因平台而改写。

Ascend 平台的落点是 `hivm.custom` 操作——Distributed Dialect 的每个操作被转换为一条带符号名的 custom op，符号在设备端模板库中静态存在，由后端在链接期解析。**编译期结束后不需要任何运行时的 module patch**，这与 A/B 路径有本质区别（见 6.3）。

### 1.2 与 GPU 路径的架构差异

| | GPU | Ascend |
| --- | --- | --- |
| 通信库符号如何绑定 | 编译出二进制后，host 侧把 SHMEM 上下文指针 **patch 进 device module** | 编译期生成 `hivm.custom` 符号，**链接期静态解析** |
| device 端如何找到通信状态 | 通过被 patch 进去的全局变量 | 位于**架构固定的 GM 地址**，无需传参 |
| 运行时初始化开销 | 每次加载 module 都要 patch | 无 |

这个差异决定了 Ascend 上 `dl.rank()` 这类操作可以做成无额外参数的纯函数——设备状态的地址是编译期已知的常量。

---

## 二、Distributed Dialect

### 2.1 Dialect 定义

Distributed Dialect 的 `dependentDialects` 被刻意置空，不依赖任何后端方言。这是它能同时下沉到 LLVM（GPU 路径）和 HIVM（Ascend 路径）的前提。

- Dialect 名：`distributed`
- C++ namespace：`::mlir::triton::distributed`

### 2.2 七个分布式op

| 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 或 TT_TensorDescType), `token`(TT_IntLike) | 同 `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，同类型) | `MemoryEffects`, `TypesMatchWith` |
| `distributed.notify` | `sigAddr`(TT_Ptr), `signalVal`(I64), `rank`(I32), `sigOp`, `commScope` | **无** | `MemoryEffects` |
| `distributed.extern_call` | `srcs`(Variadic), `libname`/`libpath`/`symbol`(StrAttr), `pure`(BoolAttr) | `result`(Variadic) | `MemoryEffects`, `ConditionallySpeculatable` |

### 2.3 属性枚举

**SignalOp**：

| 值 | 名称 | 语义 |
| --- | --- | --- |
| 1 | SET | 将信号值设置为指定值 |
| 2 | ADD | 将指定值加到当前信号值 |


### 2.4 内存副作用模型

副作用不在 TableGen 中声明，而是在 C++ 侧实现 `MemoryEffectsOpInterface`：

| Op | 副作用 |
| --- | --- |
| `WaitOp` | 只有 `Read` |
| `ConsumeTokenOp` | **空**——纯粹的 IR 依赖锚点 |
| `SymmAtOp` | 只有 `Read` |
| `NotifyOp` | `Read` **+** `Write` |
| `ExternCallOp` | `pure=true` 时无副作用，否则 `Write` + `Read` |

---

## 三、ConvertTritonDistributedToHIVM Pass

### 3.1 Pass 概览

- Pass 名：`convert-triton-distributed-to-hivm`
- 作用域：`ModuleOp`
- 依赖方言：`hivm::HIVMDialect`、`triton::distributed::DistributedDialect`

### 3.2 驱动逻辑

Pass 的 `runOnOperation()` 做四件事：

1. **扫描是否存在 `triton::DotOp` / `DotScaledOp`**，得到 `existDot` 标志。**没有 `tl.dot` 就说明这是纯 AIV kernel**，后续 core type 推导要用到这个信息。
2. 注册七个distributed op的重写模式。
3. 调用 **`applyPatternsAndFoldGreedily`** ，其余 Triton IR 要原样保留给下游的单卡编译阶段。

### 3.3 aclshmem接口OP生成

转换模板类为每个操作生成一个符号名：

| 源 op | 对应aclshmem接口名 |
| --- | --- |
| `SymmAtOp` | `aclshmem_ptr_<typename>`，如 `aclshmem_ptr_float` |
| `GetRankOp` | `aclshmem_my_pe` |
| `GetNumRanksOp` | `aclshmem_n_pes` |
| `NotifyOp`（i32 信号） | `aclshmemx_signal_op` |
| `NotifyOp`（ui64 / i64 信号） | `aclshmem_uint64_p` / `aclshmem_int64_p` |
| `ConsumeTokenOp` | `aclshmem_consume_token_<typename>`，如 `aclshmem_consume_token_float_ptr_1d` |
| `WaitOp` | `aclshmem_wait_<typename>`，如 `aclshmem_wait_int32` |
| `ExternCallOp` | 沿用 op 自带的 `symbol` 属性 |


### 3.4 属性设置

**TCoreType**（决定这条操作在哪类核上执行）：

推导顺序是：先查前缀映射表，未命中则按 `existDot` 决定。

| 符号前缀 | 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 |
| 其他（未命中前缀表） | `existDot ? CUBE_AND_VECTOR : VECTOR` |

`aclshmem_barrier_all` 被标为 CUBE_AND_VECTOR 这一点很重要：它意味着 barrier 在 kernel 被拆分为 cube func 与 vector func 之后**同时存在于两个 func 中**，这正是计算-通信重叠得以成立的机制。

**其余属性**：

| 属性 | 取值 | 说明 |
| --- | --- | --- |
| `PIPE` | `PIPE_S` | 见下方说明 |
| `VFMode` | `SIMD` | 无条件设置 |
| `hivm.is_distributed` | UnitAttr | 下游标记，供内存作用域推导、块指针分析、mix kernel 拆分等多个后续阶段识别分布式操作 |
| `symbol` | StrAttr | 裸符号名 |
| `no_side_effect` | UnitAttr（条件） | 仅当源 op 的内存副作用为空时设置 |
| `gm_addr_args_indices` | DenseI32Array | 收集「既是指针类型、又是 `tt.func` 入口块参数」的操作数下标。张量指针与 `tt.addptr` 的结果都不计入 |

> **关于 PIPE 的现状说明**：符号到流水线的映射表在代码中已声明但**当前为空**，因此所有 distributed 操作实际都落到 `PIPE_S`（标量流水）。这是一个预留的扩展点，而非"按操作类型分配流水线"的既成设计。阅读代码或做性能分析时请以此为准。

### 3.5 结果构造与替换

对每个 `RankedTensorType` 结果，先创建 `tensor::EmptyOp` 作为 DPS（destination-passing style）的 init 值，再构造 `hivm::CustomOp`。最后：无结果则 `eraseOp`，有结果则 `replaceOp`。

---

## 四、编译链路

### 4.1 完整链路与各阶段入口

```
  Python AST → Triton IR（用 DistributedOpBuilder 生成 distributed.* 操作）
  │     builder 绑定：create_distributed_wait / consume_token / get_rank /
  │                   get_num_ranks / symm_at / notify / extern_call
  ↓
【stage "ttir"】
        标准 Triton 优化：inliner / combine / canonicalize / cse / licm / loop_unroll
  ↓
【stage "ttadapter"】     ← ★ 本项目的唯一插入点 ★
        add_convert_triton_distributed_to_hivm(pm)
        —— 在单卡编译的所有后续 pass 之前执行，
           distributed 操作在进入 TTIR→Linalg 转换前就已变成 hivm.custom
  ↓
【stage "npubin"】
        单卡 NPU 编译与二进制生成
        （TTIR → Linalg IR → AscendNPU IR → 机器码，
          由三方依赖 triton-ascend 完成，本文不展开其内部实现）
  ↓
  Ascend NPU 可执行 kernel
```

**关键点**：本项目对编译流程的介入**只有一处**——在 `ttadapter` 阶段的最开始插入 `ConvertTritonDistributedToHIVM` Pass。此后所有 distributed 操作都已变成 `hivm.custom`，下游的单卡编译能力由三方依赖 `triton-ascend` 提供。

`distributed` 模块在编译器中是**可选依赖**——通过 `try / except ImportError` 加载，未构建分布式支持时自动跳过该 Pass。

### 4.2 符号如何变成真实代码（AscendNPU-IR）

这是理解整个架构的最后一环：Pass 只生成了一个字符串符号名，它是怎么变成可执行代码的？

**第一步：符号名 mangle**

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

即 `symbol = "aclshmem_my_pe"` → 链接名 `_mlir_ciface_aclshmem_my_pe`。

**第二步：模板库中的静态实现**

Pass 生成的每一个符号，在设备端模板库中都有对应的 `_mlir_ciface_` 实现。几个代表性的例子：

- `_mlir_ciface_aclshmem_my_pe` / `_mlir_ciface_aclshmem_n_pes` —— 直接读设备状态
- `ACLSHMEM_PTR_WRAPPER` 宏 —— 展开出 12 种类型的 `aclshmem_ptr_<T>`，把 `aclshmem_ptr()` 的结果包回 memref 结构
- `ACLSHMEM_P_WRAPPER` 宏 —— 展开出 `aclshmem_int64_p` / `uint64_p` 等
- `ACLSHMEM_WAIT_WRAPPER` 宏 —— 展开出 `aclshmem_wait_int32/64` 等

**`dl.wait` 由于NPU架构不保证多核并发写入同一个cacheline（DataCache）的访存一致性（64B），遍历多个 barrier 时按 64 字节硬编码跨步**。这就是应用层 `SIGNAL_SLOT_STRIDE = 64 / sizeof(dtype)` 的由来——int32 信号是 16 个元素，int64 信号是 8 个元素。

`CONSUME_TOKEN_*_WRAPPER` 的实现体则**就是把输入原样返回**，是给编译器制造数据依赖的屏障，无任何运行时开销。

**第三步：编译为 bitcode 并链接**

模板库源码会被编译为四个变体（`aic` / `aiv` / `mix_aiv` / `mix_aic`）的 bitcode，在最终生成 kernel 二进制时链入。

**这就是"为什么 Pass 只需要生成一个字符串符号名"的完整答案：符号在模板库里静态存在，由后端在链接期解析。**

---

## 五、对称堆运行时

**对称堆**：所有 rank 上大小、相对布局一致的一段设备内存，堆内偏移在各 rank 间镜像对齐。**对称张量**：从对称堆分配出的普通 torch 张量，可直接参与运算和传参。二者的价值在于——任意 rank 都能用"本地偏移 + 目标 PE 堆基址"算出该张量在其他 rank 上的地址（见 5.3 `symm_at`），无需运行时地址交换即可跨卡读写。

### 5.1 Host 端初始化

Ascend 的对称堆初始化由用户代码直接调用 `shmem` Python 包完成（惯例 `import shmem as ash`），分五步：

```python
ash.set_conf_store_tls(False, "")                                  # 1. TLS 配置
attributes = ash.InitAttr()                                        # 2. 构造初始化属性
attributes.my_rank = rank
attributes.n_ranks = world_size
attributes.local_mem_size = 1024 * 1024 * 1024                     # 3. 对称堆大小
attributes.ip_port = "tcp://127.0.0.1:8666"                        #    bootstrap 地址
attributes.option_attr.data_op_engine_type = ash.OpEngineType.MTE  # 4. 传输引擎
ash.aclshmem_init(attributes)                                      # 5. 初始化
```

**传输引擎枚举**：`MTE` / `SDMA` / `ROCE` / `UDMA`，对应四条不同的数据通路。

### 5.2 对称张量分配

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

内部流程是 `calc_nbytes` → `aclshmem_malloc` → `construct_tensor_from_ptr`，返回一个正常的 torch 张量，可以直接参与 PyTorch 运算与 kernel 传参。

> **约束**：`aclshmem_malloc` / `aclshmem_free` 必须在**所有进程同步调用，且分配或释放相同大小的内存**。

### 5.3 `symm_at` 的远端地址解析（核心机制）

**答案：查表 + 偏移算术的混合方案，两者缺一不可。**

```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;
}
```

三个步骤：

1. 取设备状态；
2. 用本地指针减去**本地**堆基址，回退到"对称偏移"；
3. 加上**目标 PE 的**堆基址，得到远端地址。

设备状态结构中有**三套并行的基址表**——`p2p_device_heap_base`、`rdma_device_heap_base`、`sdma_device_heap_base`，运行时按目标 PE 的传输方式位掩码选择使用哪一套。

---

## 六、通信原语的运行时实现

### 6.1 机制总表

| 操作 | 底层机制 |
| --- | --- |
| signal SET（P2P / MTE） | **远端 GM 普通标量写 + cacheline 回写指令**，非硬件原子、非 doorbell |
| signal ADD（P2P） | **硬件原子加 + UB→GM DataCopy**，在写端做 read-modify-write |
| signal（跨节点 ROCE） | RDMA write + quiet，或 RDMA atomic |
| wait / wait_until | **本地自旋忙等**，每轮先失效 cacheline 再读 |
| barrier | 软件算法（集中式 pull），由 signal_set + spin 组合而成 |
| UDMA | **真 doorbell**：填 WQE/SGE 到 SQ ring → 写doorbell寄存器 |

### 6.2 signal / notify

**SET 路径**的实现是两步：

1. 普通标量存储 `*addr = val`；
2. `dcci_cacheline(addr)` —— 把数据缓存刷回 GM，注释直言其目的是 *"flush data cache to GM after signal to ensure it is visible to other ranks"*。

**ADD 路径**则要走一遍 UB：置值到 UB → 设置/等待 MTE3 标志 → 开启硬件原子加模式 → `copy_ub2gm` → 关闭原子模式。

`aclshmem_<T>_p`（点写）的机制与 signal SET 同构：`aclshmem_ptr` 求地址 + 标量存储 + cacheline 回写。

### 6.3 wait 的轮询机制

六个比较谓词（`_eq` / `_ne` / `_gt` / `_ge` / `_lt` / `_le`）形态完全一致：

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

**没有硬件阻塞或中断机制，正确性完全依赖每轮循环打掉 L2 cacheline**。`dcci_cacheline` 内部用 `DataCacheCleanAndInvalid`，前后各夹一条空的 `__asm__ __volatile__("")` 防止被编译器优化掉。

这也解释了细粒度同步的一个特性：`dl.wait` 是**非消耗**语义——它只轮询、不修改信号值，因此多个 waiter 等同一个槽是安全的。约束落在复位侧（谁负责把信号写回 0）。

### 6.4 barrier_all 的实现

`aclshmem_barrier_all()` 最终落到一个**集中式 pull barrier**：

1. 每个 PE 只写自己的 flag（`signal_set`）；
2. 然后**通过 `aclshmem_ptr(sync_pool, remote_pe)` 直接自旋读对端的 flag**——不需要对方主动 push；
3. 工作按可用 Vector Core 数分摊，复杂度 O(N/K)。

**一个防死锁的关键细节**：barrier 专用的等待谓词除了接受 `== cmp_val`，**还接受 `== cmp_val + 1`**，注释说明是 *"in case when peer pe enters next barrier"*——防止快的 PE 已经进入下一轮 barrier 导致慢的 PE 永远等不到。

### 6.5 UDMA

**UDMA 不是一组新的 `dl.*` 操作，而是 host 初始化时选定的传输引擎。** 同一份 `libshmem_device.putmem()` 调用，在 MTE 与 UDMA 两种配置下走的是完全不同的底层路径。

UDMA 的设备端发送流程是真正的 doorbell 机制：

```
QP fetch → slot 计算 → 填 SQE → 填 SGE → dcci_cachelines → 敲doorbell寄存器
```

---

## 七、跨平台分发机制

### 7.1 分发结构

```
dl.wait / dl.notify / dl.symm_at / dl.rank / dl.num_ranks / dl.consume_token
   → 静态，无 proxy → _builder.create_* → distributed 方言 → (Ascend) HIVM Pass
   ※ 唯一的后端条件分支：dl.notify 的 dtype 校验

libshmem_device.<f>  → ModuleProxy 动态分发
   → (Ascend) libaclshmem_device.<f> → extern_call → ExternCallOp → aclshmem_* 符号
```

**核心的 `dl.*` 原语走静态路径**——它们直接创建 MLIR 操作，平台差异完全由 Conversion Pass 承担，Python 层没有分支。这是跨平台可移植性的基础。

**`libshmem_device.*` 走动态路径**——因为不同平台的 SHMEM 库函数集不完全相同，需要运行时按平台选择模块。

### 7.2 ModuleProxy

分发的唯一原语是 `ModuleProxy`，逻辑很简单：

1. 构造时对 `[(predicate, module), ...]` 逐一求值，断言**恰好有一个** predicate 为真；
2. `__getattr__` 把属性访问透传到选中的模块；
3. `dispatch` 装饰器让被装饰函数在**调用时**才按 `func.__name__` 去 active module 查找同名函数。

### 7.3 后端信息获取

后端判定全部基于 `shutil.which(<厂商 smi 工具>)` ：

| 函数 | 判据 |
| --- | --- |
| `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 分支清单

| 位置 | Ascend 分支做了什么 |
| --- | --- |
| `libshmem_device.py` | 把 `ascend.libaclshmem_device` 加入 proxy 列表 |
| `language_extra.py` | 把 `ascend.language_extra` 加入 proxy 列表（当前为空模块） |
| `distributed_ops.py` | `dl.notify` dtype 校验：Ascend 允许 int32（set + add），int64/uint64 **仅 set**；其他平台只允许 int64/uint64 |
| `utils.py` | 导入期 Ascend 分支为空，不导入任何厂商 SDK |
| 编译器 | 在 `ttadapter` 阶段插入 HIVM Pass |
| pybind 层 | 构建 Ascend 版本时挂载 `ascend_passes` 子模块 |