diff --git a/docs/designs/assets/fig1-topology.png b/docs/designs/assets/fig1-topology.png new file mode 100644 index 0000000000..5e570cbbea Binary files /dev/null and b/docs/designs/assets/fig1-topology.png differ diff --git a/docs/designs/assets/fig2-shared-memory.png b/docs/designs/assets/fig2-shared-memory.png new file mode 100644 index 0000000000..7b25f3ec32 Binary files /dev/null and b/docs/designs/assets/fig2-shared-memory.png differ diff --git a/docs/designs/vpto-comm-model-design.md b/docs/designs/vpto-comm-model-design.md new file mode 100644 index 0000000000..b55cdd5042 --- /dev/null +++ b/docs/designs/vpto-comm-model-design.md @@ -0,0 +1,168 @@ +# VPTO 通信模型 + +本文描述跨 rank 通信的编程模型:对称共享内存、单边访问、完成与可见性约定, +以及 `comm_scope` 边界。 + +非目标:集合通信算法、Runtime HCCL 绑定细节、Tile 层 DSL、CCU。 + +当前 IR 已落地的只有 `pto.session_init` 与 `pto.sdma_gm_gm`,写法见 +[19. Async Communication](../isa/micro-isa/19-async-comm.md)。下文仍是模型约定。 +下列条目尚未进入当前 IR,不要当成可编写语法: + +- `!pto.async_session`、`pto.session_config` +- `#pto.remote` / `#pto.mr` 指针属性 +- `pto.comm_scope` +- `pto.urma_gm_gm`、`pto.rdma_gm_gm`、`pto.sdma_gm_l2c` +- 融合 `*_signal` / `*_counter` +- kick 返回 CQ 完成记录 + +对端地址目前就是普通 `!pto.ptr`,由 host 按 `windowsIn` 同偏移算好后 +作为 kernel 参数传入。A5 远端写需要 `{soft_put}`,该形态在 op 返回前完成拷贝。 + +## 1. 范式 + +采用 **PGAS / SHMEM** 式对称共享内存 + 单边访问:各 rank 共享段布局一致, +设备侧以「本端指针 + 目标 rank」读写对端同偏移数据,无需对端参与。 + +跨卡交换收敛为三件事:**寻址、搬运、显式同步**。Scale-up(节点内)与 +Scale-out(跨节点)只更换引擎与通路,不改变编程面。 + +![](assets/fig1-topology.png) + +一次通信分两阶段: + +- **Host 准备期**:建通信域、协商对称共享段、按需注册鉴权 MR、初始化异步引擎; + 随 launch 下发寻址上下文(只读)与引擎会话(有状态)。二者职责分离即可, + 字段级 ABI 不在本文展开。 +- **NPU 运行期**:kernel 内算址、发起搬运、用同步量或融合形态约定跨卡可见性。 + +```mermaid +flowchart LR + A["① Bootstrap
带外交换 root info"] --> B["② BuildComm
HCCL 建通信域"] + B --> C["③ 注册对称内存
Window / 鉴权 RMA MR
交换基址表 + token"] + B --> D["④ 逐引擎建 workspace
SDMA · URMA · RDMA
持久化于 device HBM"] + C --> E["CommDeviceContext
寻址上下文 · 只读"] + D --> F["引擎 workspace"] + F --> G["AsyncSession
有状态 · 引用 workspace"] + E --> K["launch 入参"] + G --> K + HC["CommContext
host only · 不下设备"] -.->|"X"| K +``` + +## 2. 共享内存与指针属性 + +跨卡地址空间一律 `gm`。远近与鉴权不另开地址空间,由可组合指针属性表达: + +| 形态 | 含义 | +|------|------| +| `!pto.ptr` | 本端、普通共享内存 | +| `!pto.ptr>` | 本端、已注册鉴权 RMA MR | +| `!pto.ptr` | 远端、普通共享内存 | +| `!pto.ptr, #pto.remote>` | 远端且已注册 | + +`#pto.remote` 管远近,`#pto.mr` 管鉴权;缺省分别为本端、未注册。 +同偏移算址由调用方用 `CommDeviceContext.windowsIn[]` 自行完成: + +```text +remote = windowsIn[peer] + (local − windowsIn[myRank]) +``` + +结果以 `pto.castptr` 等既有手段成型为 `!pto.ptr`(可与 +`#pto.mr` 组合)。不设专用 remap op。 + +![](assets/fig2-shared-memory.png) + +## 3. 完成与可见性(E2 / E3) + +| 事件 | 保证 | 观测 | +|------|------|------| +| **E2** | 本端 source 可复用 | 异步:轮询搬运返回的 CQ 完成记录;同步 MTE:指令/pipe 完成即成立 | +| **E3** | 对端可见本次 payload | 写远端同步量,或使用融合 `*_signal` / `*_counter` | + +E2 与 E3 相互独立:等到 E2 **不**代表对端可见。分离写法必须先到 E2 再发同步量; +融合形态同事务保证,对端观测到同步量即可读 payload。 + +```mermaid +sequenceDiagram + participant H as Host + participant D as rank i Kernel + participant Li as rank i 对称共享内存 + participant Rj as rank j 对称共享内存 + participant P as rank j Kernel + + H->>D: launch(寻址上下文, AsyncSession, 数据 buffer) + H->>P: launch(寻址上下文, AsyncSession, 数据 buffer) + D->>Li: 取得本端 payload / 同步量地址 + D->>D: windowsIn 同偏移算址 → #pto.remote 指针 + + alt 同步通路(MTE) + D->>Rj: 单边写 payload + Note over D,Rj: 返回即本端完成 + else 异步通路(DMA 引擎) + D->>Rj: kick 单边写 payload(不阻塞标量流) + D->>D: wait event → E2:本端源可复用 + end + + D->>Rj: 写 signal / atomic add counter → 发布 E3 + loop 同步量未满足 + P->>Rj: wait / test 本端 signal / counter + Rj-->>P: 未满足则继续轮询 + end + P->>Rj: 读取本端 payload + Rj-->>P: payload(已保证可见) +``` + +跨 rank E3 **不**复用 `cmo.cacheinvalid` / `fence.barrier_all`(核间粗栅栏)。 + +## 4. 同步量(内存约定,无新 op) + +跨 rank 同步量是对称段内用户自划的 `i32` 位置,不是专用指令族: + +| 用法 | 写者 | 发布 | +|------|------|------| +| **signal** | 单写者 | `stg` / `store` / 远端 `mte_ub_gm` | +| **counter** | 多写者汇合 | `atomic_add` | + +观测:本端 `dcci` + `ldg`;等待写成 IR 轮询。与片上 SC 信号量 +(`set_intra_core` 等)互不合并:核间用 SC,跨 rank 用 GM 同步量。 + +## 5. 融合搬运+同步 + +异步通路可将 E3 发布并进同一搬运事务:`*_gm_gm_signal` / `*_gm_gm_counter`。 +这是跨 rank 同步唯一新增的 mnemonic 族;独立发布仍用 §4 的普通访存。 +MTE 同步通路无融合形态。 + +## 6. `comm_scope` + +`comm_scope` 是 `section.vector` / `section.cube` 内的词法区域,给通信资源 +(session / 完成记录等)划寿命边界,并作为 sync 分析锚点。应对齐 +`pto.vecscope` 写在 `docs/vpto-spec.md` 的层级;本节暂存约定,后续迁入该处。 + +```mlir +pto.section.vector { + pto.vecscope { /* 计算 */ } + pto.comm_scope { + %dst = pto.castptr %remote_i64 : i64 -> !pto.ptr + %cq = pto.sdma_gm_gm %dst, %src, %nbytes session(%sess) + -> !pto.ptr + } +} +``` + +| 项 | 约定 | +|----|------| +| 位置 | `section.vector` / `section.cube`;**不**进入 `vecscope` | +| 职责 | 资源寿命边界 + Sync 分析锚点 | +| vs `vecscope` | 通信 kick、同步量读写、session、远端指针构造落在 `comm_scope` | +| 推断 | session/event 流可按 SSA 穿线推断;纯同步量流需显式书写 | +| cube | AIC 只发 GM↔GM kick 时 PlanMemory 锚点弱化;资源/Sync 锚点仍成立 | + +开放问题:出口是否默认强制 E2(与正确性正交,影响跨窗 overlap)。 + +## 7. 通路总览 + +| 通路 | 承载 | +|------|------| +| 同步远端 | 核内 MTE:`mte_gm_ub` / `mte_ub_gm` + `#pto.remote` | +| 异步 GM↔GM | SDMA / URMA / RDMA + session | +| 融合 notify | `*_gm_gm_signal` / `*_gm_gm_counter` | diff --git a/docs/isa/micro-isa/18-special-scalar.md b/docs/isa/micro-isa/18-special-scalar.md index a9d4370c0e..30c189a159 100644 --- a/docs/isa/micro-isa/18-special-scalar.md +++ b/docs/isa/micro-isa/18-special-scalar.md @@ -316,8 +316,8 @@ execution scope and cache-control contract are different. - no `l1cache` or `l2cache` policy attribute is accepted; - the op must appear in an ordinary AICore entry function, outside both a `pto.simt_entry` function and `pto.section.simt`; -- the supported target profile is A5 with CANN output version 9.0.0 official - or newer. +- `pto.ld_dev` is supported on A2/A3 and on A5 with CANN output version 9.0.0 + official or newer. `pto.st_dev` is A5-only with that same CANN profile. Both operations are non-atomic. They do not imply synchronization, memory ordering, cache invalidation, cache writeback, or an L2 cache policy. Programs diff --git a/docs/isa/micro-isa/19-async-comm.md b/docs/isa/micro-isa/19-async-comm.md new file mode 100644 index 0000000000..d19d33c4c2 --- /dev/null +++ b/docs/isa/micro-isa/19-async-comm.md @@ -0,0 +1,200 @@ +# 19. Async Communication + +> **Category:** Asynchronous GM↔GM engine transfers +> **Pipelines:** SDMA engine kick from an ordinary AICore scalar stream + +This group copies a contiguous GM range through the SDMA engine. The kick does +not wait for the engine except where `{soft_put}` is documented below. The op +does not publish a completion record; local drain and cross-rank visibility are +arranged by the caller. + +This document describes: + +- `pto.session_init` +- `pto.sdma_gm_gm` + +There is no `mte_gm_gm`. Synchronous GM↔UB copies remain in +[2. DMA Copy Programming](02-dma-copy.md). + +These ops must sit in an ordinary AICore entry function. They are illegal +inside `pto.simt_entry` functions and `pto.section.simt`. + +--- + +## Session + +A session cannot be a kernel argument: only `pto.declare_struct` may produce a +`!pto.struct`. The host therefore writes a GM template, and the kernel declares +its own struct and fills it with `pto.session_init`. + +The session type is fixed: + +```mlir +!pto.struct +``` + +The template uses one 8-byte slot per field, in the same order. Narrow fields +occupy the low half of their slot. Each core fills its own copy, so a session +is per-core even when the template is shared and read-only. + +After the fill, a kernel may retune individual fields with `pto.struct_set`. +The channel group is field 4, which is how a multi-core launch gives each core +its own queue without the host naming the core. + +--- + +## Operation Summary + +| Operation | Purpose | +|-----------|---------| +| `pto.session_init` | Copy the host template into a stack-local session struct | +| `pto.sdma_gm_gm` | Kick a contiguous GM→GM copy through the session | + +--- + +### `pto.session_init` + +- **Purpose:** Fill `session` in place from the host-written GM template. +- **Syntax:** + + ```mlir + pto.session_init %sess, %sess_gm + : !pto.struct, + !pto.ptr + ``` + +- **Operands:** + + | Operand | Type | Description | + |---------|------|-------------| + | `%sess` | the 13-field session struct | Destination; written in place, no result | + | `%sess_gm` | `!pto.ptr` | Base of the host template | + +- **Results:** None. +- **Constraints:** + - `%sess` must use the session struct type above. + - `%sess_gm` must be a GM pointer. + - Must be outside SIMT entry functions and `pto.section.simt`. + - Must be inside an ordinary AICore `pto.kernel` function. +- **Semantics:** Copy each template slot into the corresponding struct field. + The caller keeps using the value `pto.declare_struct` produced. + +```text +for i in 0 .. 13: + session[i] = template_slot[i] +``` + +- **Example:** + + ```mlir + %sess = pto.declare_struct + -> !pto.struct + pto.session_init %sess, %sess_gm + : !pto.struct, + !pto.ptr + ``` + +--- + +### `pto.sdma_gm_gm` + +- **Purpose:** Copy `%nbytes` contiguous bytes from `%src` to `%dst` through + the SDMA engine attached to `%sess`. +- **Syntax:** + + ```mlir + pto.sdma_gm_gm %dst, %src, %nbytes session(%sess) + {block_bytes = $block}? {channel_idx = $ch}? {soft_put}? + : !pto.ptr, !pto.ptr, i64, + !pto.struct + ``` + +- **Operands and attributes:** + + | Name | Type | Description | + |------|------|-------------| + | `%dst` | `!pto.ptr` | Destination range; may name peer memory by address | + | `%src` | `!pto.ptr` | Source range; may name peer memory by address | + | `%nbytes` | `i64` | Contiguous byte count | + | `session(%sess)` | the 13-field session struct | Required session | + | `block_bytes` | optional `i64` attr | Split size in bytes; omitted uses the session value | + | `channel_idx` | optional `i64` attr | Channel group for this kick; omitted uses the session value | + | `soft_put` | optional unit attr | A5 remote-write completion path; ignored on A2/A3 | + +- **Results:** None. +- **Constraints:** + - `%dst` and `%src` must be GM pointers. Element types need not match; the + transfer is counted in bytes. + - `%sess` must use the session struct type above. + - There is no stride or burst model. + - `block_bytes`, when present, must be positive and a multiple of 64. + - `channel_idx`, when present, must be less than 40. + - Must be outside SIMT entry functions and `pto.section.simt`. + - Must be inside an ordinary AICore `pto.kernel` function. +- **Semantics:** Post a copy of `%nbytes` bytes from `%src` to `%dst`. The + session supplies the engine connection, the default split, the channel group, + and the service class. Either pointer may address peer memory; peer-ness is + the numeric address, not a pointer attribute. + + Without `{soft_put}` the kick does not wait for the engine. Returning from + the kernel does not mean the destination is visible. The caller observes + completion by an agreed host-side check or a later sync object. + + `{soft_put}` is for a remote write on A5. That generation's engine does not + perform a remote write, so this attr makes the copy complete before the op + returns. A2/A3 ignore it and still post to the engine. + +```text +if soft_put and target is A5: + copy nbytes bytes from src to dst # finished when the op returns +else: + post the copy to the session's engine + return without waiting +``` + +- **Example (local copy):** + + ```mlir + %sess = pto.declare_struct + -> !pto.struct + pto.session_init %sess, %sess_gm + : !pto.struct, + !pto.ptr + pto.sdma_gm_gm %dst, %src, %nbytes session(%sess) + : !pto.ptr, !pto.ptr, i64, + !pto.struct + ``` + +- **Example (A5 remote write):** + + ```mlir + pto.sdma_gm_gm %dst, %src, %nbytes session(%sess) {soft_put} + : !pto.ptr, !pto.ptr, i64, + !pto.struct + ``` + +- **Example (per-core channel after init):** + + ```mlir + %bid = pto.get_block_idx + %bid32 = arith.trunci %bid : i64 to i32 + pto.session_init %sess, %sess_gm + : !pto.struct, + !pto.ptr + pto.struct_set %sess[4], %bid32 + : !pto.struct, i32 + pto.sdma_gm_gm %dst, %src, %nbytes session(%sess) + : !pto.ptr, !pto.ptr, i64, + !pto.struct + ``` + +--- + +## PTODSL + +PTODSL explicit mode exposes the same two operations as `pto.session_init` and +`pto.sdma_gm_gm`. The session type is `pto.async_session_type()`. A session still +cannot be a kernel argument: the host writes the GM template, and the kernel +declares its own struct then fills it. See +[7.7 GM↔GM SDMA](../../../ptodsl/docs/user_guide/07-data-movement-ops.md#77-gmgm-sdma-ptosession_init-and-ptosdma_gm_gm) +in the PTODSL user guide. diff --git a/docs/vpto-spec.md b/docs/vpto-spec.md index 730cb0ecb8..f6156a3534 100644 --- a/docs/vpto-spec.md +++ b/docs/vpto-spec.md @@ -1336,6 +1336,7 @@ This section provides a categorized overview of all PTO micro Instruction operat | 16 | [Cube Matrix Multiply](isa/micro-isa/16-cube-matmul.md) | GM↔L1 (`l1`/cbuf) staging, L1 raw fill and L1 (`l1`)↔UB/BT/FB side moves, L1→L0A/L0B loads, L0C (`l0c`) matmul, and FIXPIPE MTE writeback | 20 | `pto.mte_gm_l1`, `pto.raw_fill_l1`, `pto.mte_l1_ub`, `pto.mte_gm_l1_frac`, `pto.mte_l1_bt`, `pto.mte_l1_fb`, `pto.mte_l1_l0a`, `pto.mte_l1_l0b`, `pto.mte_l1_l0a_mx`, `pto.mte_l1_l0b_mx`, `pto.mad`, `pto.mad_acc`, `pto.mad_bias`, `pto.mad_mx`, `pto.mad_mx_acc`, `pto.mad_mx_bias`, `pto.mte_l0c_l1`, `pto.mte_l0c_gm`, `pto.mte_l0c_ub` | | 17 | [SIMT Ops](isa/micro-isa/17-simt.md) | SIMT launch, thread/lane queries, vote/shuffle/redux, scalar memory, atomics, scalar math, conversion, entry synchronization, and state preservation | ~65 | `pto.store_vfsimt_info`, `pto.simt_launch`, `pto.get_tid_x`, `pto.get_laneid`, `pto.vote_*`, `pto.shuffle_*`, `pto.redux_*`, `pto.load`, `pto.store`, `pto.atomic_*`, `pto.convert`, `pto.syncthreads`, `pto.keep`, `pto.resume`, etc. | | 18 | [Special Scalar Operations](isa/micro-isa/18-special-scalar.md) | PTO scalar kernel queries, typed pointer/address calculation, scalar-pipeline memory, and ordinary AICore GM L1-bypass access | 10 | `pto.get_block_idx`, `pto.get_subblock_idx`, `pto.get_block_num`, `pto.get_subblock_num`, `pto.castptr`, `pto.addptr`, `pto.load_scalar`, `pto.store_scalar`, `pto.ld_dev`, `pto.st_dev` | +| 19 | [Async Comm](isa/micro-isa/19-async-comm.md) | Async GM↔GM SDMA copy; session filled from a host template | 2 | `pto.session_init`, `pto.sdma_gm_gm` | --- @@ -1361,6 +1362,7 @@ This section provides a categorized overview of all PTO micro Instruction operat | Contiguous Store | 3 | `pto.vsts` with `NORM_B8` / `NORM_B16` / `NORM_B32` dist | | Scatter | 3 | `pto.vscatter` | | Scalar GM access bypassing local L1 data cache | 18 | `pto.ld_dev`, `pto.st_dev` | +| GM→GM SDMA copy | 19 | `pto.sdma_gm_gm` | ### Compute Operations @@ -1393,7 +1395,8 @@ This section provides a categorized overview of all PTO micro Instruction operat ### Scalar & Control Operations Group 14 covers shared MLIR scalar arithmetic. Group 18 catalogs PTO scalar -queries, pointer/address operations, and scalar-memory operations. SIMT scalar +queries, pointer/address operations, and scalar-memory operations. Group 19 +covers session fill and GM↔GM SDMA copies. SIMT scalar operations remain in Group 17, while shared structured-control semantics remain in Group 15. diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index 6ec1ccabec..aec6cee06c 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -3822,4 +3822,100 @@ def PTO_VsturOp : PTO_VectorMicroOp<"vstur", [ }]; } +//===----------------------------------------------------------------------===// +// Async communication +//===----------------------------------------------------------------------===// +// +// Asynchronous GM->GM transfers driven by a DMA engine. The kick does not block +// the scalar stream; ordering against the transfer is the caller's job. +// +// The session operand is the config half of an async session: a stack-local +// struct whose field positions are fixed by PTO/Support/AsyncSessionABI.h. The +// mutable half (queue rings, post ids) stays in the workspace that the session's +// context field points at, because it must survive across kernel launches. + +def SessionInitOp : PTO_MicroOp<"session_init", [ + DeclareOpInterfaceMethods + ]> { + let summary = "Fill a session config struct from a template in GM."; + let description = [{ + Copy the session the host prepared at `template_gm` into `session`, field by + field. The struct is written in place, so this produces no result and the + caller keeps using the value `pto.declare_struct` gave it. + + This exists because a session cannot be handed to a kernel directly: a + stack-local struct may be neither a function argument nor the result of any + op other than `pto.declare_struct`. A kernel therefore declares its own and + fills it, and without this op that means spelling out one load and one store + per field at every entry point. + + The template holds one 8-byte slot per field, indexed the same way the + struct is, so the host writes it through the same field enumeration the + expansion reads it with. Narrow fields occupy the low half of their slot. + + Each core fills its own copy, so a session is per-core state even though the + template it came from is shared and read-only. + }]; + + let arguments = (ins + StructType:$session, + PTO_BufferType:$template_gm + ); + + let results = (outs); + + let hasVerifier = 1; + + let assemblyFormat = [{ + $session `,` $template_gm attr-dict + `:` qualified(type($session)) `,` type($template_gm) + }]; +} + +def SdmaGmGmOp : PTO_MicroOp<"sdma_gm_gm", [ + DeclareOpInterfaceMethods + ]> { + let summary = "Kick an asynchronous SDMA copy between two GM ranges."; + let description = [{ + Post a transfer of `nbytes` contiguous bytes from `src` to `dst` to the SDMA + engine and continue without waiting. Either side may name peer memory. + + There is no stride or burst model: the transfer is one contiguous range. + Splitting it across engine posts is driven by the session's block size, so a + caller that wants a specific split sets `block_bytes`. + + `channel_idx` selects a channel group other than the session default, which + is the block index. Two cores sharing a group would corrupt each other's + queue rings, so an explicit index has to be chosen with that in mind. + + Service class is not a per-post knob: it is fixed for the session when the + host builds it, since it describes how these transfers share the memory + system rather than what any one of them does. + + The op reads the channel descriptors and advances the queue tail in the + workspace, then rings the engine doorbell. Completion is not observable from + the op itself; the caller polls the agreed synchronization word. + }]; + + let arguments = (ins + PTO_BufferType:$destination, + PTO_BufferType:$source, + I64:$nbytes, + StructType:$session, + OptionalAttr:$block_bytes, + OptionalAttr:$channel_idx, + OptionalAttr:$soft_put + ); + + let results = (outs); + + let hasVerifier = 1; + + let assemblyFormat = [{ + $destination `,` $source `,` $nbytes `session` `(` $session `)` attr-dict + `:` type($destination) `,` type($source) `,` type($nbytes) `,` + qualified(type($session)) + }]; +} + #endif // MLIR_DIALECT_PTO_IR_VPTOOPS diff --git a/include/PTO/Support/AsyncSessionABI.h b/include/PTO/Support/AsyncSessionABI.h new file mode 100644 index 0000000000..cc90b065f2 --- /dev/null +++ b/include/PTO/Support/AsyncSessionABI.h @@ -0,0 +1,417 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef PTO_SUPPORT_ASYNCSESSIONABI_H +#define PTO_SUPPORT_ASYNCSESSIONABI_H + +#include +#include + +// Layout contract shared by the host that populates an async workspace, the +// device code that reads it, and the compiler passes that generate that device +// code. Every hard-coded field position must come from here. +// +// The constants below fall into three groups that differ in who owns them, how +// often they change, and what can be done to catch a change. Keeping the groups +// apart is the point of this file's structure, because a check that suits one +// group is worthless for another. +// +// 1. Ours. The session config field order and the channel record layout are +// invented here; the host writes them, generated code reads them, and no +// outside party has to agree. Nothing can drift, so nothing is checked. +// +// 2. Foreign, produced elsewhere. The async workspace and its channel +// descriptors are written by the AICPU STARS query. We only read them, so +// the layout is not ours to fix and can move between CANN releases. It is +// therefore confined to the host, which reads it once at init, checks it +// against values it fed the query, and repacks into group 1. Generated +// code never sees it. +// +// 3. Hardware, per generation. The SQE is the DMA engine's queue entry +// format. It is fixed for a generation and changes only across them, which +// is what the a5/a2a3 split expresses. No check is possible at build or +// run time -- a wrong bit is simply wrong behaviour -- so the only +// assurance is running it on each generation's silicon. +// +// An async session splits in two, because one aggregate cannot serve both +// halves: +// +// - Config: immutable scalars, carried in a stack-local `!pto.struct` filled +// once per core at kernel entry so later reads cost no memory traffic. +// - Runtime: the queue tail/head pair, which tracks a hardware queue position +// and therefore has to survive across kernel launches; a stack copy would +// silently drop it. It stays where the engine keeps it, and the channel +// record names its address. +// +// `!pto.struct` also rejects array fields, which independently keeps the +// per-channel state out of the config struct. + +namespace mlir::pto::comm { + +//===----------------------------------------------------------------------===// +// Group 1: ours +//===----------------------------------------------------------------------===// + +//===----------------------------------------------------------------------===// +// Session config struct +//===----------------------------------------------------------------------===// + +// Field positions in the session config struct, spelled in IR as +// !pto.struct +// and addressed by `pto.struct_get` / `pto.struct_set`. +// +// Pointers are carried as `i64` because `!pto.struct` rejects pointer fields; +// consumers cast back with `pto.castptr`. +enum class SessionField : int64_t { + // Base of the channel record table in GM, which the host builds at init. It + // is not the async workspace itself: the workspace holds the descriptors the + // STARS query wrote, in a layout owned elsewhere, and this points at the + // repacked view of them. + ContextGm = 0, + // Base of the UB scratch used for staging. Required wherever the doorbell is + // only reachable by MTE, since the tail has to pass through UB to get there. + TmpBufAddr = 1, + TmpBufSize = 2, + // Pipe event id for that staging. It is a session field rather than a fixed + // id so a kernel already using MTE3 events can keep the two apart. + SyncId = 3, + // Channel group owned by this core. Defaults to the block index, which is + // what makes the queue state per-core rather than shared. + ChannelIdx = 4, + // Channels in the group. Bounds the per-channel descriptor indexing. + ChannelNum = 5, + // Bytes per engine post; drives how one transfer splits into SQEs. + BlockBytes = 6, + CommBlockOffset = 7, + Engine = 8, + DestRankId = 9, + QpIdx = 10, + Flags = 11, + // Memory-system service class for the transfers this session posts, in the + // MPAM sense: it shapes priority and bandwidth share, not correctness. The + // host sets it when it builds the session, matching how the surrounding stack + // treats QoS as a property of the communication domain rather than of one + // transfer. It is deliberately independent of any QoS the domain itself was + // configured with, because these posts bypass that path and drive the queue + // directly. + Qos = 12, + NumFields = 13, +}; + +constexpr int64_t sessionFieldIndex(SessionField field) { + return static_cast(field); +} + +constexpr unsigned kSessionNumFields = + static_cast(SessionField::NumFields); + +enum class SessionEngine : uint32_t { + Sdma = 0, + Urma = 1, + Rdma = 2, +}; + +// Bits in SessionField::Flags. +constexpr uint32_t kSessionFlagValid = 1u << 0; + +//===----------------------------------------------------------------------===// +// Session template (GM) +//===----------------------------------------------------------------------===// +// +// What the host writes so a kernel does not have to spell its session out in +// constants. The kernel copies it into its own stack-local struct at entry, so +// this is a read-only initial value, not shared state: each core gets its own +// copy and may then diverge. +// +// The layout is one 8-byte slot per field, so an offset is the field index +// scaled and neither side needs a table. That wastes a few bytes on the 32-bit +// fields, which is irrelevant for a per-launch template and buys the property +// that matters: host and generated code address this through the same +// SessionField enum, so there is no second list of fields to keep in step. + +namespace session_tmpl { + +constexpr size_t kSlotBytes = 8; +constexpr size_t kBytes = kSessionNumFields * kSlotBytes; + +constexpr size_t slotOffset(SessionField field) { + return static_cast(sessionFieldIndex(field)) * kSlotBytes; +} + +// Fills the template by field rather than by position, so adding a field +// cannot silently shift the ones after it. Narrow fields are written into the +// low half of their slot, which is where a 32-bit load looks on a +// little-endian target. +class Builder { +public: + Builder &set(SessionField field, uint64_t value) { + slots_[sessionFieldIndex(field)] = value; + return *this; + } + + uint64_t get(SessionField field) const { + return slots_[sessionFieldIndex(field)]; + } + + const void *data() const { return slots_; } + static constexpr size_t bytes() { return kBytes; } + +private: + uint64_t slots_[kSessionNumFields] = {}; +}; + +} // namespace session_tmpl + +//===----------------------------------------------------------------------===// +// Channel record (64 bytes) +//===----------------------------------------------------------------------===// +// +// Everything a post needs about one channel, in a layout this project defines. +// The host fills the table at init from the descriptors the STARS query wrote, +// after checking them; generated code reads only this. +// +// The head and tail are held as addresses rather than copies because the engine +// reads and advances them in place. Resolving them here is what keeps the +// foreign descriptor layout out of generated code: a release that moves those +// fields changes two host-side address computations and nothing else. + +namespace channel { + +constexpr size_t kRecordBytes = 64; + +constexpr size_t kSqBaseOffset = 0; // uint64, SQE ring base +constexpr size_t kDoorbellOffset = 8; // uint64, doorbell register base +constexpr size_t kTailAddrOffset = 16; // uint64, address of the live tail +constexpr size_t kHeadAddrOffset = 24; // uint64, address of the live head +// Queue depth less one, so wrapping is a mask rather than a division. The +// depth is a power of two and the host rejects it otherwise, which is what +// makes the two equivalent. Dividing here would be worse than slow: the +// divisor is data read from memory, and a zero one stops the core with the +// stream never completing, which takes the card down with it. A mask of zero +// merely aims every post at slot zero. +constexpr size_t kSlotMaskOffset = 32; // uint32 +constexpr size_t kStreamIdOffset = 36; // uint32 +// 40..63 reserved. + +// Record for one channel of one group, relative to SessionField::ContextGm. +constexpr size_t recordOffset(unsigned channelIdx, unsigned channelNum, + unsigned channelInGroup) { + return (static_cast(channelIdx) * channelNum + channelInGroup) * + kRecordBytes; +} + +} // namespace channel + +//===----------------------------------------------------------------------===// +// Group 2: foreign, host-only +//===----------------------------------------------------------------------===// +// +// The async workspace as the AICPU STARS query leaves it. Read once at init, +// checked, and repacked into the channel records above. +// +// Nothing here may be used by a compiler pass or reach generated code. These +// positions are a CANN-release fact, not a hardware one, and they have already +// moved once; baking them into an instruction stream turns a future move into a +// wrong address at run time instead of a failed check at init. + +namespace workspace { + +// Most channels the query is asked for. Descriptors past what it actually +// populated are left as allocated, so a session must not index beyond the count +// the header below reports. +constexpr unsigned kMaxChannels = 40; +constexpr unsigned kSqDepth = 2048; +// Shortest transfer the engine accepts. +constexpr unsigned kMinTransferBytes = 64; + +// Flag-info header that precedes the channel table. It carries no version or +// magic, so it says nothing about the layout, but totalQueueNum does report how +// far the table was populated -- the one thing here that need not be assumed. +constexpr size_t kFlagInfoBytes = 64; +constexpr size_t kFlagInfoFlagOffset = 0; // uint32 +constexpr size_t kFlagInfoTotalQueueNumOffset = 4; // uint32 + +// Region the AICPU STARS query fills in: the flag-info header followed by the +// full channel table, padded out. The host must allocate at least this much, +// because the query writes the whole table regardless of how many channels a +// session ends up using. +constexpr size_t kContextBytes = 16 * 1024; + +// Per-group payload area that follows the context region. Nothing here is used +// by a plain transfer; it backs the flag/signal variants. It is still part of +// the allocation so the layout matches what the rest of the stack expects. +constexpr size_t kFlagPayloadBytesPerGroup = 512; + +constexpr size_t kTotalBytes = + kContextBytes + kMaxChannels * kFlagPayloadBytesPerGroup; + +//===----------------------------------------------------------------------===// +// Channel descriptor (64 bytes). Holds both the engine-facing queue geometry +// and the mutable head/tail pair. Head and tail are adjacent and are persisted +// together as one 64-bit store, tail in the high half. +//===----------------------------------------------------------------------===// + +constexpr size_t kChannelDescBytes = 64; + +constexpr size_t kChannelSqHeadOffset = 0; // uint32 +constexpr size_t kChannelSqTailOffset = 4; // uint32 +constexpr size_t kChannelSqBaseOffset = 8; // uint64, SQE ring base +constexpr size_t kChannelSqRegBaseOffset = 16; // uint64, doorbell address +constexpr size_t kChannelSqDepthOffset = 24; // uint32 +constexpr size_t kChannelSqIdOffset = 28; // uint32 +constexpr size_t kChannelCqIdOffset = 32; // uint32 +constexpr size_t kChannelLogicCqIdOffset = 36; // uint32 +constexpr size_t kChannelCqeAddrOffset = 40; // uint64 +constexpr size_t kChannelReportCqeNumOffset = 48; // uint32 +constexpr size_t kChannelStreamIdOffset = 52; // uint32 +constexpr size_t kChannelDevIdOffset = 56; // uint32 + +// Descriptor for one channel of one group, relative to SessionField::ContextGm. +constexpr size_t channelDescOffset(unsigned channelIdx, unsigned channelNum, + unsigned channelInGroup) { + return kFlagInfoBytes + + (static_cast(channelIdx) * channelNum + channelInGroup) * + kChannelDescBytes; +} + +} // namespace workspace + +//===----------------------------------------------------------------------===// +// Group 3: hardware, per generation +//===----------------------------------------------------------------------===// +// +// SDMA SQE (64 bytes). +// +// A post writes only the fields a memcpy needs; the rest of the slot is left as +// the host initialized it. +// +// The two generations share the slot size, the SQE type, the stream/task word, +// the credit position, and the address pair, but differ everywhere else that +// matters: A5 moved the transfer length to offset 48, where A2/A3 keeps a link +// type, and A2/A3 puts an `ie2` bit ahead of sssv in word 4, shifting the four +// address-attribute bits up by one. +// +// This is the one group with no automatic protection. A wrong bit produces +// wrong engine behaviour with nothing to catch it, so each generation's +// constants are only as good as a run on that generation's hardware. A2/A3 has +// had one; A5 has not. + +namespace sqe { + +constexpr size_t kBytes = 64; + +// Word 1: rtStreamId:16 | taskId:16 +constexpr size_t kWord1Offset = 4; +constexpr unsigned kRtStreamIdShift = 0; +constexpr unsigned kTaskIdShift = 16; + +// Word 3 holds the credit at the same position on both generations. +constexpr size_t kWord3Offset = 12; +constexpr unsigned kKernelCreditShift = 16; + +constexpr size_t kWord0Offset = 0; +constexpr size_t kWord4Offset = 16; + +// Address pair. Low and high halves are adjacent, so each address is one +// 64-bit store. +constexpr size_t kSrcAddrOffset = 32; +constexpr size_t kDstAddrOffset = 40; + +constexpr uint32_t kTypeSdma = 11; +constexpr unsigned kTypeShift = 0; +constexpr unsigned kOpcodeShift = 0; + +// Four bits on both generations, so the same session value is valid either way +// even though the field lives in a different word. The default matches what the +// reference SDMA implementation posts. +constexpr uint32_t kQosMask = 0xF; +constexpr uint32_t kQosDefault = 6; + +namespace a5 { + +// Word 0: type:6 | lock:1 | unlock:1 | ie:1 | preP:1 | postP:1 | wrCqe:1 | +// ptrMode:1 | rttMode:1 | headUpdate:1 | reserved0:1 | numBlocks:16 +constexpr unsigned kWrCqeShift = 11; +constexpr unsigned kNumBlocksShift = 16; + +// Word 4: opcode:8 | sssv:1 | dssv:1 | sns:1 | dns:1 | sro:1 | dro:1 | +// stride:2 | ie2:1 | compEn:1 | res4:14 +constexpr unsigned kSssvShift = 8; +constexpr unsigned kDssvShift = 9; +constexpr unsigned kSnsShift = 10; +constexpr unsigned kDnsShift = 11; + +constexpr size_t kLengthOffset = 48; + +// Word 5: sqeId:16 | mpamPartId:8 | mpamns:1 | pmg:2 | qos:4 | d2dOffsetFlag:1 +// +// QoS sits in a different word here than on A2/A3, where it shares word 4 with +// the address attributes. Nothing else in word 5 is written, so the whole word +// is the QoS field shifted into place. +constexpr size_t kWord5Offset = 20; +constexpr unsigned kQosShift = 27; + +constexpr uint32_t kKernelCreditDefault = 254; + +// Source and destination are both "secure, non-shareable" virtual addresses. +constexpr uint32_t kWord4Memcpy = + (1u << kSssvShift) | (1u << kDssvShift) | (1u << kSnsShift) | (1u << kDnsShift); + +// SDMA type, request a CQE, single block. +constexpr uint32_t kWord0Memcpy = + (kTypeSdma << kTypeShift) | (1u << kWrCqeShift); + +constexpr uint32_t kWord3Memcpy = kKernelCreditDefault << kKernelCreditShift; + +} // namespace a5 + +namespace a2a3 { + +// Word 0: type:6 | res1:10 | blockDim:16. There is no wrCqe bit here; the +// A2/A3 SQE reports completion without being asked. +constexpr unsigned kBlockDimShift = 16; + +// Word 4: opcode:8 | ie2:1 | sssv:1 | dssv:1 | sns:1 | dns:1 | qos:4 | +// sro:1 | dro:1 | partid:8 | mpam:1 | res6:4 +constexpr unsigned kIe2Shift = 8; +constexpr unsigned kSssvShift = 9; +constexpr unsigned kDssvShift = 10; +constexpr unsigned kSnsShift = 11; +constexpr unsigned kDnsShift = 12; +constexpr unsigned kQosShift = 13; + +constexpr size_t kLengthOffset = 28; + +// Byte 48 starts linkType:8 followed by three reserved bytes, so the whole +// word is written at once. +constexpr size_t kLinkTypeOffset = 48; +constexpr uint32_t kLinkTypeNone = 255; + +constexpr uint32_t kKernelCreditDefault = 240; + +// QoS is not folded in here: it comes from the session, so the expansion ors it +// into this word at run time. +constexpr uint32_t kWord4Memcpy = + (1u << kSssvShift) | (1u << kDssvShift) | (1u << kSnsShift) | + (1u << kDnsShift); + +constexpr uint32_t kWord0Memcpy = kTypeSdma << kTypeShift; + +constexpr uint32_t kWord3Memcpy = kKernelCreditDefault << kKernelCreditShift; + +} // namespace a2a3 + +// Doorbell address relative to the channel's sq_reg_base. +constexpr size_t kDoorbellOffsetA5 = 0; +constexpr size_t kDoorbellOffsetA2A3 = 8; + +} // namespace sqe + +} // namespace mlir::pto::comm + +#endif // PTO_SUPPORT_ASYNCSESSIONABI_H diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index fbbc92e232..687ba63f04 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -17,6 +17,7 @@ #include "PTO/Support/CodeConstants.h" #include "PTO/IR/PTO.h" #include "PTO/IR/PTOTypeUtils.h" +#include "PTO/Support/AsyncSessionABI.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/SCF/IR/SCF.h" @@ -708,6 +709,121 @@ LogicalResult PTOStDevOp::verify() { getValue().getType()); } +// The config half of an async session. Field positions and widths are fixed by +// PTO/Support/AsyncSessionABI.h, because the expansion indexes them by position +// and the host fills the matching workspace layout. +static LogicalResult verifyAsyncSessionType(Operation *op, Type sessionType) { + auto structTy = dyn_cast(sessionType); + if (!structTy) + return op->emitOpError() << "requires a !pto.struct session operand"; + + ArrayRef fields = structTy.getFieldTypes(); + if (fields.size() != comm::kSessionNumFields) + return op->emitOpError() + << "session struct must have exactly " << comm::kSessionNumFields + << " fields to match the async session ABI, but has " + << fields.size(); + + auto expectWidth = [&](comm::SessionField field, + unsigned width) -> LogicalResult { + Type fieldTy = fields[comm::sessionFieldIndex(field)]; + auto intTy = dyn_cast(fieldTy); + if (!intTy || intTy.getWidth() != width) + return op->emitOpError() + << "session field " << comm::sessionFieldIndex(field) + << " must be i" << width << " per the async session ABI, but is " + << fieldTy; + return success(); + }; + + for (auto [field, width] : + {std::pair{comm::SessionField::ContextGm, 64u}, + std::pair{comm::SessionField::TmpBufAddr, 64u}, + std::pair{comm::SessionField::TmpBufSize, 32u}, + std::pair{comm::SessionField::SyncId, 32u}, + std::pair{comm::SessionField::ChannelIdx, 32u}, + std::pair{comm::SessionField::ChannelNum, 32u}, + std::pair{comm::SessionField::BlockBytes, 64u}, + std::pair{comm::SessionField::CommBlockOffset, 64u}, + std::pair{comm::SessionField::Engine, 32u}, + std::pair{comm::SessionField::DestRankId, 32u}, + std::pair{comm::SessionField::QpIdx, 32u}, + std::pair{comm::SessionField::Flags, 32u}, + std::pair{comm::SessionField::Qos, 32u}}) { + if (failed(expectWidth(field, width))) + return failure(); + } + return success(); +} + +static LogicalResult verifyAsyncTransferPtr(Operation *op, Type ptrType, + StringRef role) { + auto ptrTy = dyn_cast(ptrType); + if (!ptrTy) + return op->emitOpError() << role << " must be a !pto.ptr"; + if (ptrTy.getMemorySpace().getAddressSpace() != AddressSpace::GM) + return op->emitOpError() << role << " must be a GM pointer"; + return success(); +} + +LogicalResult SessionInitOp::verify() { + if (failed(verifyAsyncSessionType(getOperation(), getSession().getType()))) + return failure(); + if (failed(verifyAsyncTransferPtr(getOperation(), getTemplateGm().getType(), + "template"))) + return failure(); + + // The fill reads GM with pto.ld_dev, which a SIMT scope does not provide, and + // the session it fills is only usable by ops under the same restriction. + if (isInsideSimtExecutionScope(getOperation())) + return emitOpError() + << "must be outside pto.simt_entry functions and pto.section.simt"; + auto funcOp = getOperation()->getParentOfType(); + if (!funcOp || !pto::isPTOEntryFunction(funcOp)) + return emitOpError() + << "requires an enclosing ordinary AICore entry function"; + return success(); +} + +LogicalResult SdmaGmGmOp::verify() { + if (failed(verifyAsyncTransferPtr(getOperation(), getDestination().getType(), + "destination"))) + return failure(); + if (failed(verifyAsyncTransferPtr(getOperation(), getSource().getType(), + "source"))) + return failure(); + if (failed(verifyAsyncSessionType(getOperation(), getSession().getType()))) + return failure(); + + // The expansion posts to the queue with scalar GM stores and, on A5, rings the + // doorbell with pto.st_dev. None of that is legal under SIMT, so reject here + // rather than letting the expansion fail later with a less obvious diagnostic. + if (isInsideSimtExecutionScope(getOperation())) + return emitOpError() + << "must be outside pto.simt_entry functions and pto.section.simt"; + auto funcOp = getOperation()->getParentOfType(); + if (!funcOp || !pto::isPTOEntryFunction(funcOp)) + return emitOpError() + << "requires an enclosing ordinary AICore entry function"; + + if (auto blockBytes = getBlockBytes()) { + if (*blockBytes == 0) + return emitOpError() << "block_bytes must be positive"; + if (*blockBytes % comm::workspace::kMinTransferBytes != 0) + return emitOpError() << "block_bytes must be a multiple of " + << comm::workspace::kMinTransferBytes; + } + + // Only the first kMaxChannels descriptors are filled in; a group past that + // reads whatever the allocation happened to contain. + if (auto channelIdx = getChannelIdx()) { + if (*channelIdx >= comm::workspace::kMaxChannels) + return emitOpError() << "channel_idx must be less than " + << comm::workspace::kMaxChannels; + } + return success(); +} + LogicalResult ShuffleIdxOp::verify() { return verifyShuffleSemanticControl(getOperation(), getIndex().getType(), getWidthAttr(), "index"); @@ -852,6 +968,26 @@ void PTOStDevOp::getEffects( effects.emplace_back(MemoryEffects::Write::get(), &getPtrMutable()); } +void SessionInitOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getTemplateGmMutable()); + // Writing the session is what keeps this from being sunk past the posts that + // read it, or dropped when nothing appears to consume it. + effects.emplace_back(MemoryEffects::Write::get(), &getSessionMutable()); +} + +void SdmaGmGmOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSourceMutable()); + effects.emplace_back(MemoryEffects::Write::get(), &getDestinationMutable()); + // The queue tail in the channel descriptor is read and advanced, so two posts + // on one session must not be reordered or dropped. + effects.emplace_back(MemoryEffects::Read::get(), &getSessionMutable()); + effects.emplace_back(MemoryEffects::Write::get(), &getSessionMutable()); +} + template static void getAtomicEffects( OpTy op, diff --git a/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp b/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp index 529c26b54b..f4166a768d 100644 --- a/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp +++ b/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp @@ -9,6 +9,7 @@ #include "PTO/Support/CodeConstants.h" #include "PTO/IR/PTO.h" #include "PTO/IR/PTOTypeUtils.h" +#include "PTO/Support/AsyncSessionABI.h" #include "PTO/Transforms/Passes.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -2077,6 +2078,545 @@ struct ExpandAtomicConfigPattern } }; +//===----------------------------------------------------------------------===// +// Async SDMA post +//===----------------------------------------------------------------------===// + +namespace comm_abi = mlir::pto::comm; + +// Materialize a GM pointer to `elementType` at `baseAddr + byteOffset`. +// +// Descriptor and SQE fields sit at fixed byte offsets but have mixed widths, so +// each access folds its offset into the address and then loads or stores at +// element index zero. That keeps the element-offset operand of ld_dev/st_dev out +// of the picture, where a stale element size would silently move the access. +static Value gmFieldPointer(Location loc, PatternRewriter &rewriter, + Value baseAddr, int64_t byteOffset, + Type elementType) { + Value addr = baseAddr; + if (byteOffset != 0) { + Value offset = getI64Constant(loc, rewriter, byteOffset); + addr = rewriter.create(loc, addr, offset); + } + auto ptrType = pto::PtrType::get( + rewriter.getContext(), elementType, + pto::AddressSpaceAttr::get(rewriter.getContext(), pto::AddressSpace::GM)); + return rewriter.create(loc, ptrType, addr); +} + +static Value loadDevField(Location loc, PatternRewriter &rewriter, + Value baseAddr, int64_t byteOffset, + Type elementType) { + Value ptr = gmFieldPointer(loc, rewriter, baseAddr, byteOffset, elementType); + Value zero = rewriter.create(loc, 0); + return rewriter.create(loc, elementType, ptr, zero); +} + +static void storeDevField(Location loc, PatternRewriter &rewriter, + Value baseAddr, int64_t byteOffset, Value value) { + Value ptr = + gmFieldPointer(loc, rewriter, baseAddr, byteOffset, value.getType()); + Value zero = rewriter.create(loc, 0); + rewriter.create(loc, value, ptr, zero); +} + +// Descriptor and SQE writes go through an ordinary store, not st_dev. +// +// On 910B1 only the first handful of st_dev stores to HBM take effect and the +// rest are dropped, with barriers making no difference; an ordinary store lands +// every time. That is enough to disqualify st_dev here, since one post writes +// seven words per SQE. st_dev is still what rings the doorbell, which is a real +// device register rather than memory. +static void storeGmField(Location loc, PatternRewriter &rewriter, + Value baseAddr, int64_t byteOffset, Value value) { + Value ptr = + gmFieldPointer(loc, rewriter, baseAddr, byteOffset, value.getType()); + Value zero = rewriter.create(loc, 0); + rewriter.create(loc, ptr, zero, value); +} + +// A2/A3 rings the doorbell through UB instead of storing to it. +// +// sq_reg_base names a register rather than memory, and on this generation it +// only takes a value that arrives by MTE: st_dev has no effect there, and a +// scalar store to it faults the vector unit hard enough to leave the card in +// an unrecoverable RAS state. Staging the tail in UB and moving four bytes out +// is what the reference SDMA implementation does. +// +// The staging slot is the session's tmp_buf, which exists for exactly this. +static void ringDoorbellViaUb(Location loc, PatternRewriter &rewriter, + Value tmpBufAddr, Value syncId32, + Value doorbellAddr, int64_t byteOffset, + Value tail32) { + MLIRContext *ctx = rewriter.getContext(); + Type i32Type = rewriter.getI32Type(); + + auto ubPtrType = pto::PtrType::get( + ctx, i32Type, pto::AddressSpaceAttr::get(ctx, pto::AddressSpace::VEC)); + Value ubPtr = rewriter.create(loc, ubPtrType, tmpBufAddr); + Value zeroIdx = rewriter.create(loc, 0); + rewriter.create(loc, ubPtr, zeroIdx, tail32); + + // The scalar unit has to be done with the slot before MTE3 picks it up. The + // event id comes from the session so a caller who is already using MTE3 + // events elsewhere can keep this staging off them; a fixed id would collide + // silently. + auto pipe = [&](pto::PIPE p) { return pto::PipeAttr::get(ctx, p); }; + Value eventId = rewriter.create( + loc, rewriter.getIndexType(), syncId32); + rewriter.create(loc, pipe(pto::PIPE::PIPE_S), + pipe(pto::PIPE::PIPE_MTE3), eventId); + rewriter.create(loc, pipe(pto::PIPE::PIPE_S), + pipe(pto::PIPE::PIPE_MTE3), eventId); + + Value dbPtr = + gmFieldPointer(loc, rewriter, doorbellAddr, byteOffset, i32Type); + Value four = getI64Constant(loc, rewriter, 4); + Value one = getI64Constant(loc, rewriter, 1); + Value zero = getI64Constant(loc, rewriter, 0); + + // A doorbell takes one 32-bit write. The default c220 store carries its + // length in 32-byte blocks and would round these four bytes down to a + // transfer of nothing, leaving the engine waiting on a ring it was never + // told about, so this asks for the byte-granular path. + auto doorbellStore = rewriter.create( + loc, ubPtr, dbPtr, zero, one, four, zero, zero, zero); + doorbellStore->setAttr("vpto.byte_granular", rewriter.getUnitAttr()); +} + +// Those stores land in the data cache, and the engine reads memory, so the +// cache has to be pushed out before the doorbell is rung. One flush of the +// whole data cache covers every SQE of the post plus the descriptor, which is +// also what the reference SDMA implementation does. +static void writebackDataCache(Location loc, PatternRewriter &rewriter, + Value addr) { + Value ptr = gmFieldPointer(loc, rewriter, addr, 0, rewriter.getI8Type()); + rewriter.create( + loc, ptr, + pto::DcciCacheLineAttr::get(rewriter.getContext(), + pto::DcciCacheLine::ENTIRE_DATA_CACHE), + pto::DcciDstAttr{}); +} + +static Value getI32Constant(Location loc, PatternRewriter &rewriter, + int64_t value) { + return rewriter.create(loc, value, 32); +} + +// Read one session config field and widen it to i64 for address arithmetic. +static Value readSessionFieldI64(Location loc, PatternRewriter &rewriter, + Value session, comm_abi::SessionField field, + unsigned width) { + Type fieldType = rewriter.getIntegerType(width); + Value raw = rewriter.create( + loc, fieldType, session, + rewriter.getDenseI64ArrayAttr({comm_abi::sessionFieldIndex(field)})); + if (width == 64) + return raw; + return rewriter.create(loc, rewriter.getI64Type(), raw); +} + +// For fields that stay 32-bit all the way into an SQE word or an event id, +// where widening to i64 would only have to be undone. +static Value readSessionFieldI32(Location loc, PatternRewriter &rewriter, + Value session, comm_abi::SessionField field) { + return rewriter.create( + loc, rewriter.getI32Type(), session, + rewriter.getDenseI64ArrayAttr({comm_abi::sessionFieldIndex(field)})); +} + +// Expand a session fill into one load and one store per field. +// +// The template gives every field an 8-byte slot, so a field's address is its +// index scaled, and a narrow field is read at the base of its slot because the +// host wrote it into the low half. +struct ExpandSessionInitPattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(pto::SessionInitOp op, + PatternRewriter &rewriter) const override { + Location loc = op.getLoc(); + Value session = op.getSession(); + auto structType = cast(session.getType()); + + Value templateAddr = rewriter.create( + loc, rewriter.getI64Type(), op.getTemplateGm()); + + for (unsigned i = 0; i < comm_abi::kSessionNumFields; ++i) { + Type fieldType = structType.getFieldTypes()[i]; + const int64_t offset = + static_cast(i * comm_abi::session_tmpl::kSlotBytes); + Value value = loadDevField(loc, rewriter, templateAddr, offset, fieldType); + rewriter.create( + loc, session, rewriter.getDenseI64ArrayAttr({static_cast(i)}), + value); + } + + rewriter.eraseOp(op); + return success(); + } +}; + +// Expand one asynchronous SDMA post into descriptor reads, SQE writes, a +// release barrier, and a doorbell write. +// +// The transfer is split into at most `block_bytes` per SQE. Splitting runs as an +// scf.for whose trip count is only known at runtime, mirroring how the DMA +// wrapper ops expand their software loops. +// +// This version drives a single channel of the group. Spreading a post across the +// group is a scheduling policy on top of the same sequence and is left to a +// follow-up. +struct ExpandSdmaGmGmPattern : public OpRewritePattern { + ExpandSdmaGmGmPattern(MLIRContext *context, DmaArch dmaArch) + : OpRewritePattern(context), dmaArch(dmaArch) {} + + LogicalResult matchAndRewrite(pto::SdmaGmGmOp op, + PatternRewriter &rewriter) const override { + Location loc = op.getLoc(); + Value session = op.getSession(); + + // A5's engine will not PUT to a peer. The same kick becomes a synchronous + // GM→UB→GM copy, which is how the reference stack writes remotely there. + // Other generations ignore the attr and keep posting SQEs. + if (dmaArch == DmaArch::A5 && op.getSoftPutAttr()) + return expandA5SoftPut(op, rewriter); + + Value contextGm = readSessionFieldI64( + loc, rewriter, session, comm_abi::SessionField::ContextGm, 64); + Value commBlockOffset = readSessionFieldI64( + loc, rewriter, session, comm_abi::SessionField::CommBlockOffset, 64); + Value channelNum = readSessionFieldI64( + loc, rewriter, session, comm_abi::SessionField::ChannelNum, 32); + + // Service class is a session-wide property, so there is no per-post form to + // fall back from. + Value qos32 = readSessionFieldI32(loc, rewriter, session, + comm_abi::SessionField::Qos); + + // A per-post override wins over the session default for both knobs. + Value channelIdx; + if (auto attr = op.getChannelIdx()) + channelIdx = getI64Constant(loc, rewriter, *attr); + else + channelIdx = readSessionFieldI64( + loc, rewriter, session, comm_abi::SessionField::ChannelIdx, 32); + + Value blockBytes; + if (auto attr = op.getBlockBytes()) + blockBytes = getI64Constant(loc, rewriter, *attr); + else + blockBytes = readSessionFieldI64( + loc, rewriter, session, comm_abi::SessionField::BlockBytes, 64); + + // record = contextGm + (channelIdx * channelNum) * recordBytes + Value recordIndex = + rewriter.create(loc, channelIdx, channelNum); + Value recordBytes = getI64Constant( + loc, rewriter, + static_cast(comm_abi::channel::kRecordBytes)); + Value recordOffset = + rewriter.create(loc, recordIndex, recordBytes); + Value recordAddr = + rewriter.create(loc, contextGm, recordOffset); + + Type i32Type = rewriter.getI32Type(); + Type i64Type = rewriter.getI64Type(); + + Value sqBase = loadDevField(loc, rewriter, recordAddr, + comm_abi::channel::kSqBaseOffset, i64Type); + Value doorbellAddr = loadDevField( + loc, rewriter, recordAddr, comm_abi::channel::kDoorbellOffset, i64Type); + Value slotMask32 = loadDevField(loc, rewriter, recordAddr, + comm_abi::channel::kSlotMaskOffset, + i32Type); + Value streamId32 = loadDevField( + loc, rewriter, recordAddr, comm_abi::channel::kStreamIdOffset, i32Type); + + // The queue position stays where the engine keeps it, so the record hands + // over its address rather than a copy. + Value tailAddr = loadDevField(loc, rewriter, recordAddr, + comm_abi::channel::kTailAddrOffset, i64Type); + Value headAddr = loadDevField(loc, rewriter, recordAddr, + comm_abi::channel::kHeadAddrOffset, i64Type); + Value sqTail32 = loadDevField(loc, rewriter, tailAddr, 0, i32Type); + Value sqHead32 = loadDevField(loc, rewriter, headAddr, 0, i32Type); + Value sqHead = rewriter.create(loc, i64Type, sqHead32); + + Value slotMask = rewriter.create(loc, i64Type, slotMask32); + Value initialTail = rewriter.create(loc, i64Type, sqTail32); + + Value srcAddr = rewriter.create(loc, i64Type, op.getSource()); + Value dstAddr = + rewriter.create(loc, i64Type, op.getDestination()); + srcAddr = rewriter.create(loc, srcAddr, commBlockOffset); + dstAddr = rewriter.create(loc, dstAddr, commBlockOffset); + + // iterations = ceilDiv(nbytes, blockBytes) + // + // The block size can come from the session, so it is a runtime value that + // nothing has checked. Clamping it away from zero keeps a bad session to a + // wrong transfer: the trip count stays bounded by nbytes. Dividing by it + // raw would stop the core instead, and a core that never finishes takes the + // card with it. + Value nbytes = op.getNbytes(); + Value oneBlock = getI64Constant(loc, rewriter, 1); + blockBytes = rewriter.create(loc, blockBytes, oneBlock); + Value iterations = + rewriter.create(loc, nbytes, blockBytes); + Value iterationsIdx = rewriter.create( + loc, rewriter.getIndexType(), iterations); + Value zeroIdx = rewriter.create(loc, 0); + Value oneIdx = rewriter.create(loc, 1); + + // The tail advances once per SQE, so it is carried through the loop. + auto forOp = rewriter.create(loc, zeroIdx, iterationsIdx, oneIdx, + ValueRange{initialTail}); + { + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointToStart(forOp.getBody()); + Value iv = rewriter.create( + loc, i64Type, forOp.getInductionVar()); + Value tail = forOp.getRegionIterArg(0); + + Value chunkOffset = rewriter.create(loc, iv, blockBytes); + // The final chunk carries whatever is left of the transfer. + Value remaining = + rewriter.create(loc, nbytes, chunkOffset); + Value chunkBytes = + rewriter.create(loc, blockBytes, remaining); + + Value chunkSrc = + rewriter.create(loc, srcAddr, chunkOffset); + Value chunkDst = + rewriter.create(loc, dstAddr, chunkOffset); + + Value slot = rewriter.create(loc, tail, slotMask); + Value sqeBytes = getI64Constant( + loc, rewriter, static_cast(comm_abi::sqe::kBytes)); + Value slotOffset = rewriter.create(loc, slot, sqeBytes); + Value sqeAddr = rewriter.create(loc, sqBase, slotOffset); + + // The engine identifies a post by how far the queue has run ahead of + // what it has drained, so the task id is the outstanding depth. + Value taskId = rewriter.create(loc, tail, sqHead); + Value taskId32 = rewriter.create(loc, i32Type, taskId); + + writeMemcpySqe(loc, rewriter, sqeAddr, chunkSrc, chunkDst, chunkBytes, + streamId32, taskId32, qos32, dmaArch); + + Value one = getI64Constant(loc, rewriter, 1); + Value nextTail = rewriter.create(loc, tail, one); + nextTail = rewriter.create(loc, nextTail, slotMask); + rewriter.create(loc, ValueRange{nextTail}); + } + + Value finalTail = forOp.getResult(0); + Value finalTail32 = + rewriter.create(loc, i32Type, finalTail); + + // Publish only the tail. The head belongs to the engine, which advances it + // as it drains the queue, so writing back a head read before the SQE stores + // would roll that progress back. + storeGmField(loc, rewriter, tailAddr, 0, finalTail32); + + // Every SQE and the tail update must be visible to the engine before the + // doorbell tells it to look. + writebackDataCache(loc, rewriter, sqBase); + rewriter.create( + loc, pto::DsbMemAttr::get(rewriter.getContext(), pto::DsbMem::DDR)); + + // The doorbell is the one write that differs by generation. A5 takes it + // through st_dev, the device-register path it was meant for. A2/A3 accepts + // it only by MTE, so the tail goes out through UB there. + const int64_t doorbellOffset = dmaArch == DmaArch::A5 + ? comm_abi::sqe::kDoorbellOffsetA5 + : comm_abi::sqe::kDoorbellOffsetA2A3; + if (dmaArch == DmaArch::A5) { + storeDevField(loc, rewriter, doorbellAddr, doorbellOffset, finalTail32); + } else { + Value tmpBufAddr = readSessionFieldI64( + loc, rewriter, session, comm_abi::SessionField::TmpBufAddr, 64); + Value syncId32 = readSessionFieldI32(loc, rewriter, session, + comm_abi::SessionField::SyncId); + ringDoorbellViaUb(loc, rewriter, tmpBufAddr, syncId32, doorbellAddr, + doorbellOffset, finalTail32); + } + + rewriter.eraseOp(op); + return success(); + } + +private: + // A5 cannot post a remote write, so the bytes go through UB in chunks. + // The copy is finished when the op returns; there is no queue tail to poll. + static LogicalResult expandA5SoftPut(pto::SdmaGmGmOp op, + PatternRewriter &rewriter) { + Location loc = op.getLoc(); + MLIRContext *ctx = rewriter.getContext(); + Value session = op.getSession(); + + Value commBlockOffset = readSessionFieldI64( + loc, rewriter, session, comm_abi::SessionField::CommBlockOffset, 64); + Value syncId32 = readSessionFieldI32(loc, rewriter, session, + comm_abi::SessionField::SyncId); + Value tmpBufAddr = readSessionFieldI64( + loc, rewriter, session, comm_abi::SessionField::TmpBufAddr, 64); + + // Address arithmetic stays in i64. A same-type pto.castptr is illegal at + // emission, so do not go through offsetPointerByBytes once the pointers + // are already i8. + Type i64Type = rewriter.getI64Type(); + Value srcAddr = + rewriter.create(loc, i64Type, op.getSource()); + Value dstAddr = + rewriter.create(loc, i64Type, op.getDestination()); + srcAddr = rewriter.create(loc, srcAddr, commBlockOffset); + dstAddr = rewriter.create(loc, dstAddr, commBlockOffset); + + auto i8Type = rewriter.getI8Type(); + auto gmI8Type = pto::PtrType::get( + ctx, i8Type, pto::AddressSpaceAttr::get(ctx, pto::AddressSpace::GM)); + auto ubType = pto::PtrType::get( + ctx, i8Type, pto::AddressSpaceAttr::get(ctx, pto::AddressSpace::VEC)); + Value ub = rewriter.create(loc, ubType, tmpBufAddr); + + Value nbytes = op.getNbytes(); + Value chunkBytes = getI64Constant(loc, rewriter, 32768); + Value one = getI64Constant(loc, rewriter, 1); + chunkBytes = rewriter.create(loc, chunkBytes, one); + Value iterations = + rewriter.create(loc, nbytes, chunkBytes); + Value iterationsIdx = rewriter.create( + loc, rewriter.getIndexType(), iterations); + Value zeroIdx = rewriter.create(loc, 0); + Value oneIdx = rewriter.create(loc, 1); + Value zero64 = getI64Constant(loc, rewriter, 0); + Value falseBit = rewriter.create( + loc, rewriter.getI1Type(), rewriter.getBoolAttr(false)); + Value eventId = rewriter.create( + loc, rewriter.getIndexType(), syncId32); + auto pipe = [&](pto::PIPE p) { return pto::PipeAttr::get(ctx, p); }; + + auto forOp = rewriter.create(loc, zeroIdx, iterationsIdx, oneIdx); + { + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointToStart(forOp.getBody()); + Value iv = rewriter.create( + loc, i64Type, forOp.getInductionVar()); + Value chunkOffset = rewriter.create(loc, iv, chunkBytes); + Value remaining = + rewriter.create(loc, nbytes, chunkOffset); + Value thisBytes = + rewriter.create(loc, chunkBytes, remaining); + Value chunkSrcAddr = + rewriter.create(loc, srcAddr, chunkOffset); + Value chunkDstAddr = + rewriter.create(loc, dstAddr, chunkOffset); + Value chunkSrc = + rewriter.create(loc, gmI8Type, chunkSrcAddr); + Value chunkDst = + rewriter.create(loc, gmI8Type, chunkDstAddr); + + rewriter.create( + loc, chunkSrc, ub, zero64, one, thisBytes, zero64, zero64, falseBit, + zero64, zero64, zero64); + rewriter.create(loc, pipe(pto::PIPE::PIPE_MTE2), + pipe(pto::PIPE::PIPE_MTE3), eventId); + rewriter.create(loc, pipe(pto::PIPE::PIPE_MTE2), + pipe(pto::PIPE::PIPE_MTE3), eventId); + rewriter.create(loc, ub, chunkDst, zero64, one, + thisBytes, zero64, zero64, zero64); + rewriter.create(loc, pipe(pto::PIPE::PIPE_MTE3), + pipe(pto::PIPE::PIPE_MTE2), eventId); + rewriter.create(loc, pipe(pto::PIPE::PIPE_MTE3), + pipe(pto::PIPE::PIPE_MTE2), eventId); + } + + rewriter.create( + loc, pto::DsbMemAttr::get(ctx, pto::DsbMem::DDR)); + rewriter.eraseOp(op); + return success(); + } + + // Write the fields a memcpy post needs. The remaining bytes of the slot keep + // whatever the host initialized them to. + static void writeMemcpySqe(Location loc, PatternRewriter &rewriter, + Value sqeAddr, Value src, Value dst, Value bytes, + Value streamId32, Value taskId32, Value qos32, + DmaArch dmaArch) { + const bool isA5 = dmaArch == DmaArch::A5; + + // Four bits wide on both generations, so a session value that does not fit + // is truncated rather than allowed to run into a neighbouring field. + Value qos = rewriter.create( + loc, qos32, getI32Constant(loc, rewriter, comm_abi::sqe::kQosMask)); + + storeGmField(loc, rewriter, sqeAddr, comm_abi::sqe::kWord0Offset, + getI32Constant(loc, rewriter, + isA5 ? comm_abi::sqe::a5::kWord0Memcpy + : comm_abi::sqe::a2a3::kWord0Memcpy)); + + // Word 1 pairs a 16-bit stream id with a 16-bit task id. Both arrive as 32 + // bits, so mask each before packing or one would run into the other. + Value halfMask = getI32Constant(loc, rewriter, 0xFFFF); + Value rtStreamId = rewriter.create(loc, streamId32, halfMask); + Value taskId = rewriter.create(loc, taskId32, halfMask); + Value taskIdShift = + getI32Constant(loc, rewriter, comm_abi::sqe::kTaskIdShift); + Value taskIdField = + rewriter.create(loc, taskId, taskIdShift); + Value word1 = rewriter.create(loc, rtStreamId, taskIdField); + storeGmField(loc, rewriter, sqeAddr, comm_abi::sqe::kWord1Offset, word1); + + storeGmField(loc, rewriter, sqeAddr, comm_abi::sqe::kWord3Offset, + getI32Constant(loc, rewriter, + isA5 ? comm_abi::sqe::a5::kWord3Memcpy + : comm_abi::sqe::a2a3::kWord3Memcpy)); + + // QoS shares word 4 with the address attributes on A2/A3, but lives in + // word 5 on A5, so only one of the two words carries it. + Value word4 = getI32Constant(loc, rewriter, + isA5 ? comm_abi::sqe::a5::kWord4Memcpy + : comm_abi::sqe::a2a3::kWord4Memcpy); + if (!isA5) { + Value qosField = rewriter.create( + loc, qos, + getI32Constant(loc, rewriter, comm_abi::sqe::a2a3::kQosShift)); + word4 = rewriter.create(loc, word4, qosField); + } + storeGmField(loc, rewriter, sqeAddr, comm_abi::sqe::kWord4Offset, word4); + + if (isA5) { + // Nothing else in word 5 is set for a memcpy post, so the QoS field is + // the whole word. + Value word5 = rewriter.create( + loc, qos, getI32Constant(loc, rewriter, comm_abi::sqe::a5::kQosShift)); + storeGmField(loc, rewriter, sqeAddr, comm_abi::sqe::a5::kWord5Offset, + word5); + } + + storeGmField(loc, rewriter, sqeAddr, comm_abi::sqe::kSrcAddrOffset, src); + storeGmField(loc, rewriter, sqeAddr, comm_abi::sqe::kDstAddrOffset, dst); + + Value bytes32 = + rewriter.create(loc, rewriter.getI32Type(), bytes); + storeGmField(loc, rewriter, sqeAddr, + isA5 ? comm_abi::sqe::a5::kLengthOffset + : comm_abi::sqe::a2a3::kLengthOffset, + bytes32); + + // A2/A3 keeps a link type where A5 puts the length; an unlinked post has to + // say so explicitly. + if (!isA5) + storeGmField( + loc, rewriter, sqeAddr, comm_abi::sqe::a2a3::kLinkTypeOffset, + getI32Constant(loc, rewriter, comm_abi::sqe::a2a3::kLinkTypeNone)); + } + + DmaArch dmaArch; +}; + struct VPTOExpandWrapperOpsPass : public pto::impl::VPTOExpandWrapperOpsBase { using pto::impl::VPTOExpandWrapperOpsBase< @@ -2098,7 +2638,9 @@ struct VPTOExpandWrapperOpsPass RewritePatternSet patterns(&getContext()); patterns.add(std::make_unique(&getContext(), dmaArch)); patterns.add(std::make_unique(&getContext(), dmaArch)); - patterns.add(&getContext(), dmaArch)); + patterns.add +packCopyUbToGmCfgAlignV220(Operation *anchor, ValueRange operands) { + if (operands.size() != 8) + return failure(); + + OpBuilder builder(anchor); + builder.setInsertionPoint(anchor); + Location loc = anchor->getLoc(); + + auto getI64Operand = [&](unsigned idx) -> Value { + return castIntegerLikeTo(anchor, operands[idx], builder.getI64Type()); + }; + + Value sid = getI64Operand(2); + Value nBurst = getI64Operand(3); + Value lenBurst = getI64Operand(4); + if (!sid || !nBurst || !lenBurst) + return failure(); + + auto shl = [&](Value value, uint64_t amount) -> Value { + return builder.create(loc, value, + getI64Constant(builder, loc, amount)); + }; + auto bitOr = [&](Value lhs, Value rhs) -> Value { + return builder.create(loc, lhs, rhs); + }; + + Value cfg = sid; + cfg = bitOr(cfg, shl(nBurst, 4)); + cfg = bitOr(cfg, shl(lenBurst, 16)); + return cfg; +} + static FailureOr packCopyUbToGmCfgV220(Operation *anchor, ValueRange operands) { if (operands.size() != 8) @@ -5927,14 +5964,28 @@ class LowerCopyOpPattern final : public OpConversionPattern { } bool isC220 = march == "dav-c220-vec" || march == "dav-c220-cube"; + // The ordinary c220 UB->GM path carries its length in 32-byte units, so a + // transfer shorter than a block rounds down to nothing at all. Callers that + // need a byte-granular store, such as the 4-byte SDMA doorbell write, ask + // for the ALIGN intrinsic instead. + bool byteGranular = !isGmUb && isC220 && op->hasAttr("vpto.byte_granular"); + if (byteGranular) + calleeName = StringAttr::get(op.getContext(), + "llvm.hivm.MOV.UB.TO.OUT.ALIGN.b32.V220") + .getValue(); + bool useA3NonPadded = isC220 && isGmUb && !hasPadding; - bool useA3UbGm = isC220 && !isGmUb; + bool useA3UbGm = isC220 && !isGmUb && !byteGranular; bool useSingleConfig = useA3NonPadded || useA3UbGm; FailureOr config0 = failure(); FailureOr config1 = failure(); - if (useA3NonPadded) - { + if (byteGranular) { + config0 = packCopyUbToGmCfgAlignV220(op, adaptor.getOperands()); + // Only exercised with both strides at zero, which is all a single-burst + // store needs. + config1 = packCopyUbToGmConfig1(op, adaptor.getOperands()); + } else if (useA3NonPadded) { config0 = packCopyGmToUbCfgV220(op, adaptor.getOperands()); } else if (useA3UbGm) { config0 = packCopyUbToGmCfgV220(op, adaptor.getOperands()); @@ -13273,6 +13324,137 @@ class ConvertPtoStgOp final : public OpConversionPattern { LoweringState &state; }; +static std::string buildLdDevCalleeName(unsigned width) { + return "llvm.hivm.LD.DEV.u" + std::to_string(width) + ".GM"; +} + +static std::string buildStDevCalleeName(unsigned width) { + return "llvm.hivm.ST.DEV.u" + std::to_string(width); +} + +class ConvertPtoLdDevOp final : public OpConversionPattern { +public: + ConvertPtoLdDevOp(TypeConverter &typeConverter, MLIRContext *context, + LoweringState &state) + : OpConversionPattern(typeConverter, context), + state(state) {} + + LogicalResult + matchAndRewrite(pto::PTOLdDevOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto llvmPtrType = dyn_cast(adaptor.getPtr().getType()); + if (!llvmPtrType) + return rewriter.notifyMatchFailure(op, "expected LLVM pointer operand"); + + auto valueType = dyn_cast(op.getValue().getType()); + if (!valueType) + return rewriter.notifyMatchFailure(op, "expected integer result type"); + + Value offset = adaptor.getOffset(); + if (offset.getType().isIndex()) + offset = rewriter.create(op.getLoc(), + rewriter.getI64Type(), offset); + + Type convertedValueType = + getTypeConverter()->convertType(op.getValue().getType()); + if (!convertedValueType) + return rewriter.notifyMatchFailure(op, + "could not convert ld_dev result type"); + + Value elemPtr = adaptor.getPtr(); + if (!matchPattern(offset, m_Zero())) { + elemPtr = rewriter.create( + op.getLoc(), llvmPtrType, + normalizeGEPElementTypeForLLVMLowering(convertedValueType, rewriter), + adaptor.getPtr(), ValueRange{offset}); + } + + FailureOr gmPtr = reinterpretPointerToAddrSpace( + op, elemPtr, static_cast(pto::AddressSpace::GM)); + if (failed(gmPtr)) + return rewriter.notifyMatchFailure(op, "failed to map ld_dev GM pointer"); + + std::string calleeName = buildLdDevCalleeName(valueType.getWidth()); + Value intrinsicOffset = getI64Constant(rewriter, op.getLoc(), 0); + auto funcType = rewriter.getFunctionType( + TypeRange{gmPtr->getType(), rewriter.getI64Type()}, + TypeRange{rewriter.getI64Type()}); + auto call = rewriter.create( + op.getLoc(), calleeName, TypeRange{rewriter.getI64Type()}, + ValueRange{*gmPtr, intrinsicOffset}); + state.plannedDecls.push_back(PlannedDecl{calleeName, funcType}); + + Value result = call.getResult(0); + if (valueType.getWidth() < 64) + result = rewriter.create(op.getLoc(), convertedValueType, + result); + rewriter.replaceOp(op, result); + return success(); + } + +private: + LoweringState &state; +}; + +class ConvertPtoStDevOp final : public OpConversionPattern { +public: + ConvertPtoStDevOp(TypeConverter &typeConverter, MLIRContext *context, + LoweringState &state) + : OpConversionPattern(typeConverter, context), + state(state) {} + + LogicalResult + matchAndRewrite(pto::PTOStDevOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto llvmPtrType = dyn_cast(adaptor.getPtr().getType()); + if (!llvmPtrType) + return rewriter.notifyMatchFailure(op, "expected LLVM pointer operand"); + + auto valueType = dyn_cast(op.getValue().getType()); + if (!valueType) + return rewriter.notifyMatchFailure(op, "expected integer value type"); + + Value offset = adaptor.getOffset(); + if (offset.getType().isIndex()) + offset = rewriter.create(op.getLoc(), + rewriter.getI64Type(), offset); + + Value elemPtr = adaptor.getPtr(); + if (!matchPattern(offset, m_Zero())) { + elemPtr = rewriter.create( + op.getLoc(), llvmPtrType, + normalizeGEPElementTypeForLLVMLowering(adaptor.getValue().getType(), + rewriter), + adaptor.getPtr(), ValueRange{offset}); + } + + FailureOr gmPtr = reinterpretPointerToAddrSpace( + op, elemPtr, static_cast(pto::AddressSpace::GM)); + if (failed(gmPtr)) + return rewriter.notifyMatchFailure(op, "failed to map st_dev GM pointer"); + + Value payload = adaptor.getValue(); + if (valueType.getWidth() < 64) + payload = rewriter.create(op.getLoc(), + rewriter.getI64Type(), payload); + + std::string calleeName = buildStDevCalleeName(valueType.getWidth()); + Value intrinsicOffset = getI64Constant(rewriter, op.getLoc(), 0); + auto funcType = rewriter.getFunctionType( + TypeRange{rewriter.getI64Type(), gmPtr->getType(), + rewriter.getI64Type()}, + TypeRange{}); + rewriter.create(op.getLoc(), calleeName, TypeRange{}, + ValueRange{payload, *gmPtr, intrinsicOffset}); + state.plannedDecls.push_back(PlannedDecl{calleeName, funcType}); + rewriter.eraseOp(op); + return success(); + } + +private: + LoweringState &state; +}; + class ConvertVPTOTypedCarrierOp final : public ConversionPattern { public: ConvertVPTOTypedCarrierOp(TypeConverter &typeConverter, MLIRContext *context) @@ -13876,7 +14058,8 @@ static LogicalResult lowerVPTOTypes(ModuleOp module, llvm::raw_ostream &diagOS) }); target.addIllegalOp(); target.addDynamicallyLegalOp( [&](UnrealizedConversionCastOp op) { @@ -13904,7 +14087,7 @@ static LogicalResult lowerVPTOTypes(ModuleOp module, llvm::raw_ostream &diagOS) ConvertPtoStructGetOp, ConvertPtoStructSetOp, ConvertPtoStoreScalarOp>(typeConverter, context); patterns.add( + ConvertPtoStgOp, ConvertPtoLdDevOp, ConvertPtoStDevOp>( typeConverter, context, state); patterns.add(typeConverter, context); patterns.add(typeConverter, context); diff --git a/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp b/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp index 22fc492ed5..dddb756a21 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitterDispatcher.cpp @@ -20,32 +20,37 @@ static bool usesCANN900Lowering(const VPTOEmissionOptions &options) { options.cannVersion >= CANNVersion::release(mlir::pto::kValue9, 0, 0); } -static bool containsLdStDev(ModuleOp module) { +static bool containsOp(ModuleOp module, bool wantLdDev, bool wantStDev) { bool found = false; module.walk([&](Operation *op) { - if (isa(op)) { + if ((wantLdDev && isa(op)) || + (wantStDev && isa(op))) { found = true; } }); return found; } +// ld_dev is required on A2/A3 for SDMA descriptor reads. st_dev is A5-only: +// A2/A3 cannot use it to ring the SDMA doorbell, and stores through that path +// are not reliable on that generation. Other non-c220 targets still need the +// 9.0.0 official lowering; the 9.0.0-beta.1 intrinsic set has not been checked. static LogicalResult verifyLdStDevTarget(ModuleOp module, const VPTOEmissionOptions &options, llvm::raw_ostream &diagOS) { - if (!containsLdStDev(module) || usesCANN900Lowering(options)) { - return success(); - } - const bool isC220 = options.march == "dav-c220-vec" || options.march == "dav-c220-cube"; - if (isC220) { - diagOS << "VPTO LLVM emission failed: pto.ld_dev and pto.st_dev require " - "--pto-arch=a5\n"; - } else { - diagOS << "VPTO LLVM emission failed: pto.ld_dev and pto.st_dev require " - "CANN 9.0.0 or newer official lowering\n"; + if (isC220 && containsOp(module, false, true)) { + diagOS << "VPTO LLVM emission failed: pto.st_dev is not supported on A2/A3\n"; + return failure(); } + if (!containsOp(module, true, true) || isC220 || + usesCANN900Lowering(options)) { + return success(); + } + + diagOS << "VPTO LLVM emission failed: pto.ld_dev and pto.st_dev require " + "CANN 9.0.0 or newer official lowering\n"; return failure(); } diff --git a/ptodsl/docs/user_guide/01-introduction.md b/ptodsl/docs/user_guide/01-introduction.md index b7fd37a578..c1bcb867cb 100644 --- a/ptodsl/docs/user_guide/01-introduction.md +++ b/ptodsl/docs/user_guide/01-introduction.md @@ -64,6 +64,7 @@ Python Wrapper L0 user-facing wrapper (NumPy, torch-npu, pure Pyth │ └─ backend="emitc" EmitC backend, mode="auto" only ├─ Tile Ops tile.load, tile.store, tile.add, ... ├─ MTE Ops mte_load / mte_store / mte_gm_ub / ... + ├─ Async GM↔GM session_init / sdma_gm_gm ├─ @pto.tileop matrix products (mad, mte_l1_l0a, mte_l0c_ub, ...) ├─ @pto.tileop row-wise vector math (vlds, vadd, vexp, vsts, ...) └─ @pto.simt scalar-like compute (lds, sts, pointwise blends, ...) diff --git a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md index fb46df33f5..83c598dc3c 100644 --- a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md +++ b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md @@ -251,6 +251,73 @@ compiled[grid, stream](A.ctypes.data, O.ctypes.data, 4, 128) **Only `entry=True` kernels support `.compile()` and `[grid, stream]` launch.** Calling `.compile()` on an `entry=False` module raises an error. +### Host C++ in the kernel's library: `native_options` + +PTODSL builds a kernel into a shared library from two inputs it generates itself: +the launch code and the ptoas-produced kernel object. When a kernel's host side +already exists in C++ — a resource setup helper, an ABI the device code shares +with a host header — that closed set leaves only bad options: reimplement the C++ +in Python, or ship it as a second library the caller loads separately and keeps +in step by hand. + +`native_options` opens the build up. It names host C++ sources to compile and +link into the same library: + +```python +@pto.jit( + name="my_kernel", + target="a3", + mode="explicit", + native_options={ + # Compiled as host C++ and linked into this kernel's library. + "host_sources": ["MyHostHelper.cpp"], + # Include directories for those sources. + "include_dirs": [".", "../../include"], + # Extra link inputs; bare library names, as for -l. + "link_libraries": ["dl"], + "library_dirs": ["/opt/mylib/lib"], + }, +) +def my_kernel(buf: pto.ptr(pto.f32, "gm")): + ... +``` + +Then reach the host entry points through the library the build produced: + +```python +compiled = my_kernel.compile() +lib = compiled.native_library() # a ctypes.CDLL +lib.my_host_setup() # a symbol from MyHostHelper.cpp +compiled[grid, stream](buf) +``` + +Rules worth knowing: + +- **Relative paths resolve against the declaring file**, not the working + directory, so a kernel keeps building wherever it is invoked from. This is the + same rule `source=` uses. +- **Host sources are part of the build's identity.** Their contents are digested + into the cache key, so editing one rebuilds the library instead of silently + reusing the old one. +- **Host sources are compiled as host code**, with `-xc++` and no AI-core arch — + not as device code. They cannot contain kernel code. +- **`link_libraries` takes bare names.** Pass `"dl"`, not `"-ldl"` or a path to a + `.so`; directories belong in `library_dirs`, which also become rpath entries so + the library is found again at load time. Host sources also pull `-lstdc++`, + because `bisheng --cce-fatobj-link` does not add the C++ runtime by itself. +- **Undefined symbols remain an error.** The link keeps `-Wl,--no-undefined`, so a + library you forgot to name fails the build rather than the first launch. +- **`include_dirs` requires `host_sources`.** On its own it would apply to + nothing, so it is rejected rather than quietly ignored. +- Unknown keys are rejected when the decorator runs, not at build time. + +`native_library()` and `native_library_path()` are only for reaching host symbols; +launching needs neither. Both build the specialization if it is not built yet, and +both share one loaded library with every launch handle over that specialization. + +For a worked example, see `test/comm/` in the repository: a host helper that owns +an ABI shared with the device code, linked into the kernel that uses it. + ### Loading an existing PTO file Use `source=` when the kernel implementation already exists as hand-written PTO diff --git a/ptodsl/docs/user_guide/04-type-system-and-buffer.md b/ptodsl/docs/user_guide/04-type-system-and-buffer.md index f6471030f1..afc9b66a79 100644 --- a/ptodsl/docs/user_guide/04-type-system-and-buffer.md +++ b/ptodsl/docs/user_guide/04-type-system-and-buffer.md @@ -115,6 +115,9 @@ Python `int` literals can initialize integer and floating-point fields; Python ` Structs cannot be used as `@pto.jit`, `@pto.tileop`, or `@pto.simt` parameters, or as `pto.for_(...).carry(...)` state. A struct declared outside an ordinary `pto.for_` may be read and mutated inside that loop. Struct values also cannot be returned, yielded, or passed as function arguments from their declaring scope. +`pto.async_session_type()` is the 13-field session used by `pto.session_init` and +`pto.sdma_gm_gm`. See [7.7 GM↔GM SDMA](07-data-movement-ops.md#77-gmgm-sdma-ptosession_init-and-ptosdma_gm_gm). + ## 4.2 Vector register type Vector registers hold a fixed 256-byte payload. `pto.vreg(dtype)` infers the element count automatically: diff --git a/ptodsl/docs/user_guide/07-data-movement-ops.md b/ptodsl/docs/user_guide/07-data-movement-ops.md index b378212f36..2d96366bba 100644 --- a/ptodsl/docs/user_guide/07-data-movement-ops.md +++ b/ptodsl/docs/user_guide/07-data-movement-ops.md @@ -1,6 +1,6 @@ # 7. Data Movement Operations -This chapter covers every operation that moves data between memory spaces in PTODSL — tile-level transfers, DMA micro-instructions, vector loads and stores, and cube data movement. Operations are organized by abstraction level: tile ops for auto mode, DMA orchestration for explicit mode, vector memory ops on the SIMD unit, and cube memory ops on the Cube unit. +This chapter covers every operation that moves data between memory spaces in PTODSL — tile-level transfers, DMA micro-instructions, vector loads and stores, cube data movement, and GM↔GM engine copies. Operations are organized by abstraction level: tile ops for auto mode, DMA orchestration for explicit mode, vector memory ops on the SIMD unit, cube memory ops on the Cube unit, and session-driven GM↔GM copies. ## 7.1 Tile-level movement: tile.load and tile.store @@ -1600,3 +1600,123 @@ def vector_consumer( c2v.free(entry, split=0) pto.tile.store(b_tile, b_part) ``` + +## 7.7 GM↔GM SDMA: `pto.session_init` and `pto.sdma_gm_gm` + +These two operations copy a contiguous GM range through the SDMA engine. They +are explicit-mode only, must sit in an ordinary AICore `@pto.jit` body, and are +illegal inside `@pto.simt` or `pto.section.simt`. There is no stride or burst +model: the transfer is one contiguous byte count. + +A session cannot be a kernel argument. The host writes a GM template; the kernel +declares its own struct with `pto.async_session_type()` and fills it with +`pto.session_init`. Field order and widths match the ISA chapter +[19. Async Communication](../../../../docs/isa/micro-isa/19-async-comm.md). After +the fill, `pto.struct_set` may retune a field. Field 4 is the channel group: a +multi-core launch gives each core its own queue by writing that field rather +than by having the host name the core. + +The kick does not wait for the engine except when `soft_put=True` on A5. The +operation does not publish a completion record; local drain and cross-rank +visibility are arranged by the caller. + +#### `pto.async_session_type() -> StructTypeDescriptor` + +**Description**: The 13-field session type +`!pto.struct`. +Use it with `pto.declare_struct`. + +#### `pto.session_init(session, template_gm) -> None` + +**Description**: Copy the host template into `session` in place. The caller keeps +using the value `pto.declare_struct` produced. Each core fills its own copy, so +a session is per-core even when the template is shared and read-only. + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `session` | the 13-field session struct | Destination; written in place | +| `template_gm` | `PtrType` in GM | Base of the host template | + +**Returns**: None. + +**Constraints**: `session` must be `pto.async_session_type()`. `template_gm` must +be a GM pointer. Explicit mode only. + +#### `pto.sdma_gm_gm(destination, source, nbytes, *, session, block_bytes=None, channel_idx=None, soft_put=False) -> None` + +**Description**: Copy `nbytes` contiguous bytes from `source` to `destination` +through the session. Either pointer may address peer memory; peer-ness is the +numeric address, not a pointer attribute. Element types need not match; the +transfer is counted in bytes. + +When `block_bytes` is omitted, the split size comes from the session. When +`channel_idx` is omitted, the channel group comes from the session. `soft_put` +is for a remote write on A5: that generation's engine does not perform a remote +write, so this flag makes the copy complete before the call returns. A2/A3 +ignore it and still post to the engine. + +Without `soft_put`, returning from the kernel does not mean the destination is +visible. The caller observes completion by an agreed host-side check or a later +sync object. + +```text +if soft_put and target is A5: + copy nbytes bytes from source to destination # finished when the call returns +else: + post the copy to the session's engine + return without waiting +``` + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `destination` | `PtrType` in GM | Destination range | +| `source` | `PtrType` in GM | Source range | +| `nbytes` | `i64` | Contiguous byte count | +| `session` | the 13-field session struct | Required session | +| `block_bytes` | static `int` or `None` | Split size in bytes; omitted uses the session value | +| `channel_idx` | static `int` or `None` | Channel group for this kick; omitted uses the session value | +| `soft_put` | `bool` | A5 remote-write completion path; default `False` | + +**Returns**: None. + +**Constraints**: + +- `destination` and `source` must be GM pointers. +- `session` must be `pto.async_session_type()`. +- `block_bytes`, when present, must be a positive multiple of 64. +- `channel_idx`, when present, must be in `[0, 39]`. +- Explicit mode only; ordinary AICore entry, not SIMT. + +**Example (local copy):** + + +```python +sess = pto.declare_struct(pto.async_session_type()) +pto.session_init(sess, sess_gm) +pto.sdma_gm_gm(dst, src, nbytes, session=sess) +``` + +**Example (A5 remote write):** + + +```python +sess = pto.declare_struct(pto.async_session_type()) +pto.session_init(sess, sess_gm) +pto.sdma_gm_gm(dst, src, nbytes, session=sess, soft_put=True) +``` + +**Example (per-core channel after init):** + + +```python +sess = pto.declare_struct(pto.async_session_type()) +pto.session_init(sess, sess_gm) +# Field 4 is the channel group. A Python int or an i32 SSA value is legal; +# pto.get_block_idx() is i64 and must be narrowed first. +pto.struct_set(sess, 4, 0) +pto.sdma_gm_gm(dst, src, nbytes, session=sess) +``` diff --git a/ptodsl/ptodsl/_jit.py b/ptodsl/ptodsl/_jit.py index ec5f4baa8a..4893bfd267 100644 --- a/ptodsl/ptodsl/_jit.py +++ b/ptodsl/ptodsl/_jit.py @@ -21,6 +21,7 @@ ) from ._kernel_compilation import CompiledKernelHandle, KernelCompiler from ._kernel_signature import parse_jit_kernel_signature +from ._native_options import normalize_native_options from ._tracing import ( KernelModuleSpec, ModuleArtifact, @@ -231,6 +232,7 @@ def jit( insert_sync: bool | None = None, ast_rewrite: bool | None = None, frontend_options: Mapping | None = None, + native_options: Mapping | None = None, source: str | None = None, ): """ @@ -258,6 +260,14 @@ def jit( frontend_options: Reserved structured frontend options. Currently supports ``ast_rewrite`` and ``rewrite_part={"control_flow"}``. + native_options: + Host-side additions to the native build, for a kernel whose + host side already exists in C++. Supports ``host_sources`` + (C++ files compiled and linked into the kernel's shared + library), ``include_dirs`` for those sources, and + ``link_libraries`` / ``library_dirs`` for the link. Relative + paths resolve against the declaring file. Host sources are part + of the build's identity, so editing one rebuilds the library. source: Optional filesystem path to a PTO IR source file. When provided, PTODSL keeps the Python signature as the host ABI @@ -298,6 +308,11 @@ def decorator(fn): source_file = inspect.getsourcefile(fn) or inspect.getfile(fn) except (OSError, TypeError): source_file = None + normalized_native_options = normalize_native_options( + native_options, + declaring_file=source_file, + function_name=fn_name, + ) kernel_kind_explicit = kernel_kind is not _DEFAULT_KERNEL_KIND_SENTINEL effective_kernel_kind = kernel_kind if kernel_kind_explicit else _DEFAULT_KERNEL_KIND compiler = KernelCompiler( @@ -315,6 +330,7 @@ def decorator(fn): source_file=source_file, source_line=getattr(fn.__code__, "co_firstlineno", None), jit_source=source, + native_options=normalized_native_options, ), kernel_signature, fn, @@ -374,6 +390,7 @@ def __ptodsl_cache_signature__(self): module_spec.mode, module_spec.kernel_kind, module_spec.kernel_kind_explicit, + module_spec.native_options, ) def _build_default_module(self): diff --git a/ptodsl/ptodsl/_kernel_compilation.py b/ptodsl/ptodsl/_kernel_compilation.py index 60a589534f..c2b5f69dfb 100644 --- a/ptodsl/ptodsl/_kernel_compilation.py +++ b/ptodsl/ptodsl/_kernel_compilation.py @@ -16,7 +16,7 @@ kernel_module_compile_error, kernel_module_launch_error, ) -from ._runtime.launch import LaunchHandle, parse_launch_spec +from ._runtime.launch import LaunchHandle, build_and_load_native_library, parse_launch_spec from ._source_loader import SourceModuleLoader from ._tracing import ModuleArtifact, SignatureTracingRuntime @@ -39,6 +39,9 @@ def __init__( self._constexpr_bindings = dict(constexpr_bindings) self._module_spec = module_spec self._kernel_signature = kernel_signature + # Set on the first native build; shared by every launch handle over this + # specialization and by native_library(). + self._native_library_cache = None @property def specialization_key(self): @@ -57,6 +60,26 @@ def kernel_module_graph(self): """Return traced kernel-module import/dependency metadata for this build.""" return self.build_metadata().get("kernel_module_graph") + def native_library(self): + """Build this specialization if needed and return the loaded shared library. + + Launching does not need this. It is how a caller reaches a symbol that + ``@pto.jit(native_options={"host_sources": ...})`` compiled into the same + library, so host setup a kernel depends on can live in C++ next to the + kernel rather than in a separately built library. + """ + if self._module_spec.entry is False: + raise kernel_module_launch_error(self._py_name) + library, _path, _symbol = build_and_load_native_library(self) + return library + + def native_library_path(self): + """Build this specialization if needed and return its shared library path.""" + if self._module_spec.entry is False: + raise kernel_module_launch_error(self._py_name) + _library, path, _symbol = build_and_load_native_library(self) + return path + def __getitem__(self, launch_spec): if self._module_spec.entry is False: raise kernel_module_launch_error(self._py_name) diff --git a/ptodsl/ptodsl/_native_options.py b/ptodsl/ptodsl/_native_options.py new file mode 100644 index 0000000000..fc9bcc6971 --- /dev/null +++ b/ptodsl/ptodsl/_native_options.py @@ -0,0 +1,184 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""Host-side additions to the native build of one ``@pto.jit`` kernel. + +PTODSL builds a kernel into a shared library by compiling generated launch code +and the ptoas-produced kernel object, then linking the two. That is a closed set, +which is a problem for a kernel whose host side already exists in C++: the only +ways in were to reimplement it in Python or to ship it as a second library the +caller loads separately. + +``native_options`` opens that up. It names host C++ sources to compile and link +into the same library, the include directories they need, and the libraries to +link against. Those sources become part of the build's identity, so editing one +rebuilds the library rather than silently reusing it. + +The public surface is a plain mapping, matching ``frontend_options`` on the same +decorator. It is normalized here into a frozen, hashable record, because +``KernelModuleSpec`` is frozen and the build cache compares these values. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path + +_SUPPORTED_NATIVE_OPTION_KEYS = frozenset( + {"host_sources", "include_dirs", "link_libraries", "library_dirs"} +) + +# A library name goes onto the link line as -lNAME. Anything that could be read +# as a second argument, a path, or a flag is rejected rather than pasted. +_FORBIDDEN_LIBRARY_CHARS = frozenset({"/", "\\", " ", "\t", "\n", "\r", '"', "'", ";", "&", "|", "$"}) + + +@dataclass(frozen=True) +class NativeBuildOptions: + """Resolved host-side additions to one kernel's native build.""" + + host_sources: tuple[Path, ...] = () + include_dirs: tuple[Path, ...] = () + link_libraries: tuple[str, ...] = () + library_dirs: tuple[Path, ...] = () + + def is_empty(self) -> bool: + return not ( + self.host_sources or self.include_dirs or self.link_libraries or self.library_dirs + ) + + +EMPTY_NATIVE_OPTIONS = NativeBuildOptions() + + +def _require_path_sequence(key: str, value) -> tuple[str, ...]: + """Accept one path or an iterable of paths, and reject anything else clearly.""" + if isinstance(value, (str, Path)): + return (str(value),) + if isinstance(value, (bytes, bytearray)): + raise TypeError(f"@pto.jit native_options[{key!r}] must be str or os.PathLike, not bytes") + try: + items = list(value) + except TypeError as exc: + raise TypeError( + f"@pto.jit native_options[{key!r}] must be a path or an iterable of paths" + ) from exc + resolved = [] + for item in items: + if not isinstance(item, (str, Path)): + raise TypeError( + f"@pto.jit native_options[{key!r}] entries must be str or os.PathLike, " + f"got {type(item).__name__}" + ) + text = str(item) + if not text: + raise ValueError(f"@pto.jit native_options[{key!r}] contains an empty path") + resolved.append(text) + return tuple(resolved) + + +def _resolve_against(base_dir: Path | None, raw: str) -> Path: + """Anchor a relative path at the file that declared the kernel. + + Relative to the declaring module rather than the process working directory, + so a kernel keeps building wherever it is invoked from. This is the rule + ``@pto.jit(source=...)`` already uses for its IR path. + """ + candidate = Path(raw).expanduser() + if candidate.is_absolute() or base_dir is None: + return candidate.resolve() + return (base_dir / candidate).resolve() + + +def _normalize_paths(key: str, value, *, base_dir: Path | None) -> tuple[Path, ...]: + resolved = [] + for raw in _require_path_sequence(key, value): + path = _resolve_against(base_dir, raw) + if path not in resolved: + resolved.append(path) + return tuple(resolved) + + +def _normalize_libraries(value) -> tuple[str, ...]: + names = [] + for raw in _require_path_sequence("link_libraries", value): + # -l takes a bare name; the caller means library_dirs if they have a path. + bad = sorted(_FORBIDDEN_LIBRARY_CHARS.intersection(raw)) + if bad: + raise ValueError( + f"@pto.jit native_options['link_libraries'] entry {raw!r} contains " + f"unsupported characters {bad!r}; pass a bare library name such as 'dl' " + "and put directories in native_options['library_dirs']" + ) + if raw.startswith("-"): + raise ValueError( + f"@pto.jit native_options['link_libraries'] entry {raw!r} looks like a linker " + "flag; pass a bare library name such as 'dl'" + ) + if raw not in names: + names.append(raw) + return tuple(names) + + +def normalize_native_options( + native_options: Mapping | None, + *, + declaring_file: str | None = None, + function_name: str | None = None, +) -> NativeBuildOptions: + """Validate the public mapping and resolve it against the declaring file. + + Paths are resolved but not required to exist: a kernel is often declared + before its host sources are generated, and the build reports a missing source + with the command that needed it. + """ + if native_options is None: + return EMPTY_NATIVE_OPTIONS + if not isinstance(native_options, Mapping): + raise TypeError("@pto.jit native_options must be a mapping when provided") + + unknown = set(native_options) - _SUPPORTED_NATIVE_OPTION_KEYS + if unknown: + raise ValueError( + f"@pto.jit native_options has unsupported keys: {sorted(unknown)!r}; " + f"supported keys are {sorted(_SUPPORTED_NATIVE_OPTION_KEYS)!r}" + ) + + base_dir = Path(declaring_file).resolve().parent if declaring_file else None + + options = NativeBuildOptions( + host_sources=_normalize_paths( + "host_sources", native_options.get("host_sources", ()), base_dir=base_dir + ), + include_dirs=_normalize_paths( + "include_dirs", native_options.get("include_dirs", ()), base_dir=base_dir + ), + link_libraries=_normalize_libraries(native_options.get("link_libraries", ())), + library_dirs=_normalize_paths( + "library_dirs", native_options.get("library_dirs", ()), base_dir=base_dir + ), + ) + + # Include and library directories only mean something for sources being + # compiled or symbols being linked. Accepting them alone would silently do + # nothing, which is worth a diagnostic rather than a quiet no-op. + if options.include_dirs and not options.host_sources: + where = f" for {function_name}" if function_name else "" + raise ValueError( + f"@pto.jit native_options['include_dirs'] was set without 'host_sources'{where}; " + "include directories apply to the host sources this option compiles" + ) + + return options + + +__all__ = [ + "EMPTY_NATIVE_OPTIONS", + "NativeBuildOptions", + "normalize_native_options", +] diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 7308f1682c..1e1de4a2ff 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -575,6 +575,8 @@ mte_ub_l1, mte_ub_ub, raw_fill_l1, + sdma_gm_gm, + session_init, set_atomic_add, set_atomic_bf16, set_atomic_f16, @@ -687,6 +689,7 @@ __all__ = [ "const", "declare_struct", "struct_get", "struct_set", + "session_init", "sdma_gm_gm", "castptr", "addptr", "vlds", "vldas", "vldus", "vldsx2", "vsts", "vstsx2", "init_align", diff --git a/ptodsl/ptodsl/_ops_mte.py b/ptodsl/ptodsl/_ops_mte.py index 79a987998f..3d29858d4b 100644 --- a/ptodsl/ptodsl/_ops_mte.py +++ b/ptodsl/ptodsl/_ops_mte.py @@ -5,7 +5,7 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""Data-movement ops: MTE transfers, mad, accumulator-store attributes.""" +"""Data-movement ops: MTE transfers, async GM copies, mad, accumulator-store attributes.""" from functools import wraps import warnings @@ -45,6 +45,7 @@ wrap_surface_value, ) from ._types import ( + _ASYNC_SESSION_FIELD_WIDTHS, _is_struct_type, _isinstance_pto_type, _materialize_integer_literal, @@ -79,6 +80,9 @@ VectorType, ) +from ._ops_core import ( + _require_struct_value, +) from ._ops_common import ( _coerce_i1, _coerce_i32, @@ -529,6 +533,110 @@ def _require_pto_ptr_operand(value, *, context: str): return raw_value +def _require_gm_ptr(ptr_value, *, context: str): + raw_ptr = unwrap_surface_value(ptr_value) + try: + ptr_type = _pto.PtrType(raw_ptr.type) + except Exception as exc: + raise TypeError(f"{context} requires a typed PTO pointer") from exc + gm_space = _pto.AddressSpaceAttr.get(_pto.AddressSpace.GM) + if ptr_type.memory_space != gm_space: + raise TypeError(f"{context} requires a GM pointer, got {raw_ptr.type}") + return raw_ptr + + +def _require_async_session(session, *, op_name: str): + raw_session, struct_type = _require_struct_value(session, op_name=op_name) + fields = tuple(struct_type.field_types) + if len(fields) != len(_ASYNC_SESSION_FIELD_WIDTHS): + raise TypeError( + f"{op_name}: session must have {len(_ASYNC_SESSION_FIELD_WIDTHS)} fields " + f"to match the async session type, but has {len(fields)}" + ) + for index, (field, width) in enumerate(zip(fields, _ASYNC_SESSION_FIELD_WIDTHS)): + if not IntegerType.isinstance(field) or IntegerType(field).width != width: + raise TypeError( + f"{op_name}: session field {index} must be i{width}, got {field}" + ) + return raw_session + + +def _optional_static_i64_attr( + value, + *, + context: str, + minimum: int | None = None, + maximum: int | None = None, + alignment: int | None = None, +): + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{context} expects a static int or None, got {value!r}") + if minimum is not None and value < minimum: + raise ValueError(f"{context} expects a value >= {minimum}, got {value}") + if maximum is not None and value > maximum: + raise ValueError(f"{context} expects a value <= {maximum}, got {value}") + if alignment is not None and value % alignment != 0: + raise ValueError( + f"{context} expects a multiple of {alignment} bytes, got {value}" + ) + return IntegerAttr.get(IntegerType.get_signless(64), value) + + +@_explicit_mode_only("pto.session_init(...)") +def session_init(session, template_gm): + """Fill ``session`` in place from the host-written GM template.""" + raw_session = _require_async_session(session, op_name="pto.session_init(...)") + raw_template = _require_gm_ptr(template_gm, context="pto.session_init(...) template") + _pto.SessionInitOp(raw_session, raw_template) + + +@_explicit_mode_only("pto.sdma_gm_gm(...)") +def sdma_gm_gm( + destination, + source, + nbytes, + *, + session, + block_bytes=None, + channel_idx=None, + soft_put=False, +): + """Kick a contiguous GM→GM copy through the session.""" + if not isinstance(soft_put, bool): + raise TypeError("pto.sdma_gm_gm(...): soft_put expects a bool") + raw_dst = _require_gm_ptr(destination, context="pto.sdma_gm_gm(...) destination") + raw_src = _require_gm_ptr(source, context="pto.sdma_gm_gm(...) source") + raw_session = _require_async_session(session, op_name="pto.sdma_gm_gm(...)") + attrs = {} + block_bytes_attr = _optional_static_i64_attr( + block_bytes, + context="pto.sdma_gm_gm(...) block_bytes", + minimum=1, + alignment=64, + ) + channel_idx_attr = _optional_static_i64_attr( + channel_idx, + context="pto.sdma_gm_gm(...) channel_idx", + minimum=0, + maximum=39, + ) + if block_bytes_attr is not None: + attrs["block_bytes"] = block_bytes_attr + if channel_idx_attr is not None: + attrs["channel_idx"] = channel_idx_attr + if soft_put: + attrs["soft_put"] = UnitAttr.get() + _pto.SdmaGmGmOp( + raw_dst, + raw_src, + _coerce_i64(nbytes, context="pto.sdma_gm_gm(...) nbytes"), + raw_session, + **attrs, + ) + + @_explicit_mode_only("pto.mte_load(...)") def mte_load(source, destination, l2_cache_ctl, len_burst, *, nburst, loops=None, pad=None): """ diff --git a/ptodsl/ptodsl/_runtime/__init__.py b/ptodsl/ptodsl/_runtime/__init__.py index a83ab9eaff..ff1f688dc0 100644 --- a/ptodsl/ptodsl/_runtime/__init__.py +++ b/ptodsl/ptodsl/_runtime/__init__.py @@ -5,12 +5,34 @@ # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. -"""Private runtime helpers for ``@pto.jit`` launch.""" +"""Private runtime helpers for ``@pto.jit`` launch. -from .launch import LaunchHandle -from .native_build import build_native_library +Launch code talks to MLIR types. Native-build helpers do not. Keep this +package lazy so a test or caller that only wants command construction does +not need ``ptoas.mlir`` bindings. +""" + +from importlib import import_module __all__ = [ "LaunchHandle", + "build_and_load_native_library", "build_native_library", ] + +_EXPORTS = { + "LaunchHandle": (".launch", "LaunchHandle"), + "build_and_load_native_library": (".launch", "build_and_load_native_library"), + "build_native_library": (".native_build", "build_native_library"), +} + + +def __getattr__(name): + try: + module_name, attr_name = _EXPORTS[name] + except KeyError as exc: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc + module = import_module(module_name, __name__) + value = getattr(module, attr_name) + globals()[name] = value + return value diff --git a/ptodsl/ptodsl/_runtime/launch.py b/ptodsl/ptodsl/_runtime/launch.py index d6e5c26402..79fe64a26f 100644 --- a/ptodsl/ptodsl/_runtime/launch.py +++ b/ptodsl/ptodsl/_runtime/launch.py @@ -108,6 +108,29 @@ def _marshal_launch_args(kernel_signature, args): return marshaled +def build_and_load_native_library(compiled: CompiledKernelHandle): + """Build this specialization if needed and return ``(library, path, symbol)``. + + The loaded library is cached on the compiled handle. Several launch handles + over one specialization share it, and so does a caller reaching past the + launch entry for a symbol that ``native_options['host_sources']`` contributed. + """ + cached = getattr(compiled, "_native_library_cache", None) + if cached is not None: + return cached + + lib_path, launch_symbol = build_native_library( + py_name=compiled._py_name, + module_spec=compiled._module_spec, + kernel_signature=compiled._kernel_signature, + mlir_text=compiled.mlir_text(), + specialization_key=compiled.specialization_key, + ) + loaded = (ctypes.CDLL(str(lib_path)), lib_path, launch_symbol) + compiled._native_library_cache = loaded + return loaded + + class LaunchHandle: """Callable launch binding returned by ``compiled[grid, stream]``.""" @@ -124,14 +147,7 @@ def _ensure_launch_fn(self): if self._launch_fn is not None: return - lib_path, launch_symbol = build_native_library( - py_name=self._compiled._py_name, - module_spec=self._compiled._module_spec, - kernel_signature=self._compiled._kernel_signature, - mlir_text=self._compiled.mlir_text(), - specialization_key=self._compiled.specialization_key, - ) - lib = ctypes.CDLL(str(lib_path)) + lib, _lib_path, launch_symbol = build_and_load_native_library(self._compiled) fn = getattr(lib, launch_symbol) fn.argtypes = _launch_argtypes(self._compiled._kernel_signature) fn.restype = None @@ -174,5 +190,6 @@ def parse_launch_spec(launch_spec) -> tuple[int, object]: __all__ = [ "LaunchHandle", + "build_and_load_native_library", "parse_launch_spec", ] diff --git a/ptodsl/ptodsl/_runtime/native_build.py b/ptodsl/ptodsl/_runtime/native_build.py index 9bbe588575..97365695f0 100644 --- a/ptodsl/ptodsl/_runtime/native_build.py +++ b/ptodsl/ptodsl/_runtime/native_build.py @@ -20,7 +20,7 @@ is_native_build_current, write_manifest, ) -from .codegen import generate_launch_cpp, launch_symbol_name +from .._native_options import EMPTY_NATIVE_OPTIONS, NativeBuildOptions from .toolchain import ( aicore_arch_for_kernel_kind, common_include_flags, @@ -86,6 +86,37 @@ def _source_ptoas_overrides(module_spec) -> dict: return {"backend": module_spec.backend} +def _native_options_of(module_spec) -> NativeBuildOptions: + """Read the kernel's host-side build additions, defaulting to none. + + Read defensively because a module spec reaches this layer from several + frontends, and one that predates the option should build as it always did. + """ + return getattr(module_spec, "native_options", None) or EMPTY_NATIVE_OPTIONS + + +def _host_source_config_lines(native_options: NativeBuildOptions) -> list[str]: + """Cache identity for the host sources compiled into the library. + + Both the path and the file's contents, because a host source is an input to + this build like the MLIR is. Digesting the path alone would reuse a library + built from an older version of the same file, which is the failure mode a + caller iterating on a shim would hit first. + + A source that cannot be read yields no digest instead of an error: the build + reports the missing file with the command that needed it, and this function + only decides whether the previous build still counts. + """ + lines = [f"include_dir={path}" for path in native_options.include_dirs] + for source in native_options.host_sources: + try: + digest = _content_digest(source.read_text(encoding="utf-8")) + except OSError: + digest = "unreadable" + lines.append(f"host_source={source}:{digest}") + return lines + + def _compile_config_text( *, module_spec, @@ -102,12 +133,13 @@ def _compile_config_text( f"pto_level={effective_pto_level}", f"backend={ptoas_overrides.get('backend')}", "enable_tile_op_expand=True", + *_host_source_config_lines(_native_options_of(module_spec)), ] ) -def _host_compile_flags() -> list[str]: - return common_include_flags() + [ +def _host_compile_flags(include_dirs: tuple[Path, ...] = ()) -> list[str]: + return common_include_flags() + [f"-I{path}" for path in include_dirs] + [ "-std=gnu++17", "-O2", "-Wno-macro-redefined", @@ -170,12 +202,70 @@ def _compile_launch_cpp( ) +def _host_object_path(cache_dir: Path, index: int, source: Path) -> Path: + """Object path for one extra host source. + + Indexed as well as named, because two sources in different directories may + share a stem and would otherwise overwrite each other's object. + """ + return cache_dir / f"host_{index}_{source.stem}.o" + + +def _compile_host_sources( + native_options: NativeBuildOptions, + cache_dir: Path, +) -> list[Path]: + """Compile each ``native_options['host_sources']`` entry to an object file.""" + if not native_options.host_sources: + return [] + + # Every source is checked before the toolchain is resolved, so a mistyped + # path reports itself rather than whichever environment variable the + # compiler lookup happens to want first. The paths were resolved against the + # declaring file, so naming one in full is the useful part of the message. + missing = [source for source in native_options.host_sources if not source.is_file()] + if missing: + listed = ", ".join(str(path) for path in missing) + raise FileNotFoundError( + f"@pto.jit native_options['host_sources'] does not exist: {listed}" + ) + + bisheng = resolve_bisheng() + flags = _host_compile_flags(native_options.include_dirs) + objects = [] + for index, source in enumerate(native_options.host_sources): + host_object = _host_object_path(cache_dir, index, source) + _run([bisheng, *flags, "-c", str(source), "-o", str(host_object)]) + objects.append(host_object) + return objects + + +def _extra_link_flags(native_options: NativeBuildOptions) -> list[str]: + """Library search paths and library names contributed by ``native_options``. + + Each search path also becomes an rpath entry, matching how the CANN runtime + directories are added: a caller who had to name a directory to link against + needs it found again at load time. + """ + flags: list[str] = [] + for lib_dir in native_options.library_dirs: + flags.extend([f"-L{lib_dir}", f"-Wl,-rpath,{lib_dir}"]) + flags.extend(f"-l{name}" for name in native_options.link_libraries) + # bisheng --cce-fatobj-link does not pull the C++ runtime. Host C++ objects + # that use std::string, operator new, or static guards need it on the line. + if native_options.host_sources and "stdc++" not in native_options.link_libraries: + flags.append("-lstdc++") + return flags + + def _link_shared_library( launch_object: Path, kernel_object: Path, shared_library: Path, *, kernel_kind: str | None, + host_objects: list[Path] | None = None, + extra_link_flags: list[str] | None = None, ) -> None: bisheng = resolve_bisheng() soname = shared_library.name @@ -192,7 +282,13 @@ def _link_shared_library( str(shared_library), str(launch_object), str(kernel_object), + *[str(path) for path in host_objects or []], *runtime_library_flags(sim_mode=sim_mode), + # After the runtime flags, so a host source may depend on a runtime + # symbol. -Wl,--no-undefined is on, so anything still unresolved here + # is a missing entry in native_options rather than a load-time + # surprise. + *(extra_link_flags or []), ] ) @@ -210,8 +306,12 @@ def _native_build_config(module_spec): effective_pto_level=effective_pto_level, ptoas_overrides=ptoas_overrides, ) + native_options = _native_options_of(module_spec) + extra_link_flags = _extra_link_flags(native_options) sim_mode = bool(os.environ.get("MSPROF_SIMULATOR_MODE")) - link_config_text = "\n".join(runtime_library_flags(sim_mode=sim_mode)) + link_config_text = "\n".join( + runtime_library_flags(sim_mode=sim_mode) + extra_link_flags + ) return { "insert_sync": effective_insert_sync, "pto_level": effective_pto_level, @@ -238,11 +338,16 @@ def _compile_native_artifacts(artifacts, module_spec, *, effective_insert_sync, target_arch=module_spec.target_arch, export_macro=f"{module_spec.function_name}_EXPORTS", ) + native_options = _native_options_of(module_spec) + extra_link_flags = _extra_link_flags(native_options) + host_objects = _compile_host_sources(native_options, artifacts.cache_dir) _link_shared_library( launch_object, artifacts.kernel_object, artifacts.shared_library, kernel_kind=module_spec.kernel_kind, + host_objects=host_objects, + extra_link_flags=extra_link_flags, ) @@ -276,6 +381,8 @@ def build_native_library( specialization_key, ) -> tuple[Path, str]: """Build or reuse the shared library for one compiled specialization.""" + from .codegen import generate_launch_cpp, launch_symbol_name + ir_function_name = module_spec.function_name artifacts = artifact_paths(py_name, ir_function_name, specialization_key) launch_symbol = launch_symbol_name(ir_function_name) diff --git a/ptodsl/ptodsl/_tracing/module_builder.py b/ptodsl/ptodsl/_tracing/module_builder.py index 7d3e851a0c..7e453008d4 100644 --- a/ptodsl/ptodsl/_tracing/module_builder.py +++ b/ptodsl/ptodsl/_tracing/module_builder.py @@ -15,6 +15,8 @@ from ptoas.mlir.dialects import func from ptoas.mlir.ir import Attribute, InsertionPoint, Module, Operation, StringAttr, UnitAttr +from .._native_options import EMPTY_NATIVE_OPTIONS, NativeBuildOptions + class ModuleStyle(str, Enum): """Supported top-level PTODSL module layouts.""" @@ -40,6 +42,7 @@ class KernelModuleSpec: source_file: str | None = None source_line: int | None = None jit_source: str | None = None + native_options: NativeBuildOptions = EMPTY_NATIVE_OPTIONS def _build_flat_aicore_module(spec: KernelModuleSpec, arg_types): diff --git a/ptodsl/ptodsl/_types.py b/ptodsl/ptodsl/_types.py index 70dd2d56f5..d58cba6457 100644 --- a/ptodsl/ptodsl/_types.py +++ b/ptodsl/ptodsl/_types.py @@ -598,6 +598,18 @@ def struct_type(*field_types) -> _StructDescriptor: return _StructDescriptor(field_types) +# Field widths of the async-comm session, matching docs/isa/micro-isa/19-async-comm.md. +_ASYNC_SESSION_FIELD_WIDTHS = (64, 64, 32, 32, 32, 32, 64, 64, 32, 32, 32, 32, 32) + + +def async_session_type() -> _StructDescriptor: + """Return the 13-field session type used by ``pto.session_init`` / ``pto.sdma_gm_gm``.""" + fields = [] + for width in _ASYNC_SESSION_FIELD_WIDTHS: + fields.append(int64 if width == 64 else int32) + return struct_type(*fields) + + def vmi_vreg_type(lanes: int, elem) -> _VMIVRegDescriptor: """Return a lazy descriptor for ``!pto.vmi.vreg``.""" return _VMIVRegDescriptor(lanes, elem) @@ -714,7 +726,7 @@ def part_tensor_view_type_from_dims(dims, elem) -> Type: "si8", "si16", "si32", "si64", "ui8", "ui16", "ui32", "ui64", "index", - "ptr", "vreg_type", "vec_type", "mask_type", "struct_type", + "ptr", "vreg_type", "vec_type", "mask_type", "struct_type", "async_session_type", "vmi_vreg_type", "vmi_mask_type", "tile_buf_type", "tensor_view_type", "tensor_view_type_from_dims", "part_tensor_view_type", "part_tensor_view_type_from_dims", diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index 78ecd05493..1a982327bc 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -34,7 +34,7 @@ si8, si16, si32, si64, ui8, ui16, ui32, ui64, index, - ptr, vreg_type, vec_type, mask_type, struct_type, + ptr, vreg_type, vec_type, mask_type, struct_type, async_session_type, _resolve, ) from ._builtin_vector import Vec # noqa: F401 @@ -87,6 +87,7 @@ from ._ops import ( # noqa: F401 const, declare_struct, struct_get, struct_set, + session_init, sdma_gm_gm, get_op_attr, castptr, addptr, vlds, vldas, vldus, vldsx2, vsts, vstsx2, diff --git a/ptodsl/tests/support/docs_fragment_fixtures.py b/ptodsl/tests/support/docs_fragment_fixtures.py index fcd185e847..cb8fd7c968 100644 --- a/ptodsl/tests/support/docs_fragment_fixtures.py +++ b/ptodsl/tests/support/docs_fragment_fixtures.py @@ -1147,6 +1147,18 @@ def data_movement_grouped_dma_ptrs_probe(): {SNIPPET_PLACEHOLDER} """ ), + "data_movement.async_comm": _fixture( + f""" + @pto.jit(target="a5", mode="explicit") + def data_movement_async_comm_probe( + src: pto.ptr(pto.i8, "gm"), + sess_gm: pto.ptr(pto.i8, "gm"), + dst: pto.ptr(pto.i8, "gm"), + nbytes: pto.i64, + ): + {SNIPPET_PLACEHOLDER} + """ + ), "data_movement.low_precision_vector_memory": _fixture( f""" @pto.jit(target="a5", mode="explicit") diff --git a/ptodsl/tests/test_async_comm.py b/ptodsl/tests/test_async_comm.py new file mode 100644 index 0000000000..49e8d877f8 --- /dev/null +++ b/ptodsl/tests/test_async_comm.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import unittest + +from ptodsl import _types, pto +from ptodsl._context import make_context +from ptoas.mlir.ir import InsertionPoint, Location, Module + + +@pto.jit(target="a5", mode="explicit") +def sdma_gm_gm_surface_kernel( + src: pto.ptr(pto.i8, "gm"), + sess_gm: pto.ptr(pto.i8, "gm"), + dst: pto.ptr(pto.i8, "gm"), + nbytes: pto.i64, +): + sess = pto.declare_struct(pto.async_session_type()) + pto.session_init(sess, sess_gm) + pto.sdma_gm_gm(dst, src, nbytes, session=sess) + pto.sdma_gm_gm(dst, src, nbytes, session=sess, soft_put=True, block_bytes=64, channel_idx=0) + + +class AsyncCommSurfaceTest(unittest.TestCase): + def test_public_namespace_exports_the_surface(self): + for name in ("async_session_type", "session_init", "sdma_gm_gm"): + with self.subTest(name=name): + self.assertTrue(hasattr(pto, name), name) + self.assertIn("async_session_type", _types.__all__) + + def test_async_session_type_matches_the_isa_layout(self): + with make_context() as ctx, Location.unknown(ctx): + resolved = pto.async_session_type().resolve() + self.assertEqual( + str(resolved), + "!pto.struct", + ) + + def test_surface_emits_session_init_and_sdma_gm_gm(self): + text = sdma_gm_gm_surface_kernel.compile().mlir_text() + self.assertIn("pto.session_init", text) + self.assertIn("pto.sdma_gm_gm", text) + self.assertIn("soft_put", text) + self.assertIn("block_bytes = 64", text) + self.assertIn("channel_idx = 0", text) + + with make_context() as ctx: + module = Module.parse(text, ctx) + module.operation.verify() + + def test_rejects_wrong_session_and_pointer_spaces(self): + with make_context() as ctx, Location.unknown(ctx): + module = Module.create() + with InsertionPoint(module.body): + sess = pto.declare_struct(pto.async_session_type()) + other = pto.declare_struct(pto.struct_type(pto.i32)) + gm = pto.ptr(pto.i8, "gm") + ub = pto.ptr(pto.i8, "ub") + # Materialize dummy SSA pointers via castptr so the helpers see values. + zero = pto.const(0, dtype=pto.i64) + gm_ptr = pto.castptr(zero, gm) + ub_ptr = pto.castptr(zero, ub) + + with self.assertRaisesRegex(TypeError, "13 fields"): + pto.session_init(other, gm_ptr) + with self.assertRaisesRegex(TypeError, "GM pointer"): + pto.session_init(sess, ub_ptr) + with self.assertRaisesRegex(TypeError, "GM pointer"): + pto.sdma_gm_gm(ub_ptr, gm_ptr, 64, session=sess) + with self.assertRaisesRegex(ValueError, "multiple of 64"): + pto.sdma_gm_gm(gm_ptr, gm_ptr, 64, session=sess, block_bytes=32) + with self.assertRaisesRegex(ValueError, "<= 39"): + pto.sdma_gm_gm(gm_ptr, gm_ptr, 64, session=sess, channel_idx=40) + with self.assertRaisesRegex(TypeError, "soft_put"): + pto.sdma_gm_gm(gm_ptr, gm_ptr, 64, session=sess, soft_put=1) + + +if __name__ == "__main__": + unittest.main() diff --git a/ptodsl/tests/test_native_build_options.py b/ptodsl/tests/test_native_build_options.py new file mode 100644 index 0000000000..637ae1dd1f --- /dev/null +++ b/ptodsl/tests/test_native_build_options.py @@ -0,0 +1,386 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""``@pto.jit(native_options=...)`` normalization and native-build wiring. + +Everything here runs without a toolchain: the option layer is pure Python, and +the build layer is exercised by capturing the commands it would run. +""" + +import importlib.util +import tempfile +import types +import unittest +from pathlib import Path +from unittest import mock + +from ptodsl._native_options import ( + EMPTY_NATIVE_OPTIONS, + NativeBuildOptions, + normalize_native_options, +) +from ptodsl._runtime import native_build + +# KernelModuleSpec and @pto.jit pull MLIR bindings. The rest of this file +# only inspects option records and captured compiler command lines. +_HAS_PTOAS_MLIR = importlib.util.find_spec("ptoas.mlir") is not None + + +class NormalizeNativeOptionsTest(unittest.TestCase): + def test_none_and_empty_mapping_are_empty(self): + self.assertIs(normalize_native_options(None), EMPTY_NATIVE_OPTIONS) + self.assertTrue(normalize_native_options({}).is_empty()) + + def test_relative_paths_resolve_against_the_declaring_file(self): + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir).resolve() + declaring = root / "pkg" / "kernel.py" + declaring.parent.mkdir(parents=True) + declaring.touch() + + options = normalize_native_options( + {"host_sources": ["shim.cpp", "../other/helper.cpp"], "include_dirs": "inc"}, + declaring_file=str(declaring), + ) + + self.assertEqual( + options.host_sources, + (root / "pkg" / "shim.cpp", root / "other" / "helper.cpp"), + ) + self.assertEqual(options.include_dirs, (root / "pkg" / "inc",)) + + def test_absolute_paths_are_kept(self): + with tempfile.TemporaryDirectory() as temp_dir: + absolute = Path(temp_dir).resolve() / "shim.cpp" + options = normalize_native_options( + {"host_sources": [absolute]}, + declaring_file="/somewhere/else/kernel.py", + ) + self.assertEqual(options.host_sources, (absolute,)) + + def test_a_single_path_does_not_have_to_be_wrapped_in_a_list(self): + options = normalize_native_options({"host_sources": "/tmp/shim.cpp"}) + self.assertEqual(options.host_sources, (Path("/tmp/shim.cpp").resolve(),)) + + def test_duplicate_entries_collapse(self): + options = normalize_native_options( + {"host_sources": ["/tmp/a.cpp", "/tmp/a.cpp"], "link_libraries": ["dl", "dl"]} + ) + self.assertEqual(len(options.host_sources), 1) + self.assertEqual(options.link_libraries, ("dl",)) + + def test_options_are_hashable_so_a_kernel_spec_stays_frozen(self): + options = normalize_native_options( + {"host_sources": ["/tmp/a.cpp"], "link_libraries": ["dl"]} + ) + self.assertEqual(hash(options), hash(options)) + self.assertEqual(options, normalize_native_options( + {"host_sources": ["/tmp/a.cpp"], "link_libraries": ["dl"]} + )) + + def test_unknown_keys_are_rejected_and_named(self): + with self.assertRaises(ValueError) as caught: + normalize_native_options({"host_source": ["/tmp/a.cpp"]}) + self.assertIn("host_source", str(caught.exception)) + self.assertIn("host_sources", str(caught.exception)) + + def test_non_mapping_is_rejected(self): + with self.assertRaises(TypeError): + normalize_native_options(["/tmp/a.cpp"]) + + def test_library_names_reject_paths_and_flags(self): + for bad in ("/usr/lib/libdl.so", "-ldl", "dl; rm -rf /", "my lib"): + with self.assertRaises(ValueError, msg=f"{bad!r} should be rejected"): + normalize_native_options({"link_libraries": [bad]}) + + def test_library_dirs_do_not_have_to_accompany_host_sources(self): + # Linking against a prebuilt library needs no source of our own. + options = normalize_native_options( + {"link_libraries": ["dl"], "library_dirs": ["/opt/lib"]} + ) + self.assertEqual(options.link_libraries, ("dl",)) + self.assertEqual(options.library_dirs, (Path("/opt/lib"),)) + + def test_include_dirs_without_host_sources_is_rejected(self): + # They would apply to nothing, so a quiet no-op is worse than an error. + with self.assertRaises(ValueError) as caught: + normalize_native_options({"include_dirs": ["/opt/include"]}, function_name="k") + self.assertIn("host_sources", str(caught.exception)) + + def test_empty_path_is_rejected(self): + with self.assertRaises(ValueError): + normalize_native_options({"host_sources": [""]}) + + def test_wrong_entry_type_is_rejected(self): + with self.assertRaises(TypeError): + normalize_native_options({"host_sources": [123]}) + + +def _module_spec(native_options=EMPTY_NATIVE_OPTIONS): + return types.SimpleNamespace( + function_name="k", + target_arch="a3", + kernel_kind="vector", + mode="explicit", + backend="vpto", + insert_sync=None, + jit_source=None, + native_options=native_options, + ) + + +class HostSourceCacheIdentityTest(unittest.TestCase): + """Editing a host source has to invalidate the cached library.""" + + def _config_text(self, options): + return native_build._compile_config_text( + module_spec=_module_spec(options), + effective_insert_sync=False, + effective_pto_level="level3", + ptoas_overrides={"backend": "vpto"}, + ) + + def test_no_options_leaves_the_config_text_unchanged(self): + text = self._config_text(EMPTY_NATIVE_OPTIONS) + self.assertNotIn("host_source=", text) + self.assertIn("target_arch=a3", text) + + def test_editing_a_host_source_changes_the_config_text(self): + with tempfile.TemporaryDirectory() as temp_dir: + source = Path(temp_dir) / "shim.cpp" + source.write_text("int a() { return 1; }\n", encoding="utf-8") + options = NativeBuildOptions(host_sources=(source,)) + + before = self._config_text(options) + source.write_text("int a() { return 2; }\n", encoding="utf-8") + after = self._config_text(options) + + self.assertIn(str(source), before) + self.assertNotEqual(before, after) + + def test_an_unreadable_host_source_is_recorded_rather_than_raising(self): + options = NativeBuildOptions(host_sources=(Path("/nonexistent/shim.cpp"),)) + self.assertIn("unreadable", self._config_text(options)) + + def test_include_dirs_take_part_in_the_identity(self): + base = NativeBuildOptions(host_sources=(Path("/tmp/a.cpp"),)) + with_inc = NativeBuildOptions( + host_sources=(Path("/tmp/a.cpp"),), include_dirs=(Path("/opt/inc"),) + ) + self.assertNotEqual(self._config_text(base), self._config_text(with_inc)) + + def test_a_module_spec_without_the_option_still_builds(self): + legacy = types.SimpleNamespace( + function_name="k", target_arch="a3", kernel_kind="vector", mode="explicit", + backend="vpto", insert_sync=None, jit_source=None, + ) + self.assertIs(native_build._native_options_of(legacy), EMPTY_NATIVE_OPTIONS) + + +@unittest.skipUnless(_HAS_PTOAS_MLIR, "needs ptoas.mlir bindings") +class KernelModuleSpecTest(unittest.TestCase): + def test_the_spec_carries_the_options_and_stays_frozen(self): + from ptodsl._tracing import KernelModuleSpec + + options = NativeBuildOptions(host_sources=(Path("/tmp/a.cpp"),)) + spec = KernelModuleSpec( + function_name="k", target_arch="a3", kernel_kind="vector", native_options=options + ) + self.assertIs(spec.native_options, options) + self.assertIs(native_build._native_options_of(spec), options) + with self.assertRaises(Exception): + spec.native_options = EMPTY_NATIVE_OPTIONS + + def test_a_spec_without_the_option_defaults_to_none_of_it(self): + from ptodsl._tracing import KernelModuleSpec + + spec = KernelModuleSpec(function_name="k", target_arch="a3", kernel_kind="vector") + self.assertTrue(spec.native_options.is_empty()) + + +@unittest.skipUnless(_HAS_PTOAS_MLIR, "needs ptoas.mlir bindings") +class JitDecoratorSurfaceTest(unittest.TestCase): + """The decorator kwarg has to reach the spec, and reject bad input early.""" + + SOURCE = ( + 'module attributes {pto.target_arch = "a3"} {\n' + " func.func @native_option_probe(%arg0: !pto.ptr)" + " attributes {pto.kernel} {\n" + " return\n" + " }\n" + "}\n" + ) + + def _kernel(self, native_options): + from ptodsl import pto + + @pto.jit( + name="native_option_probe", + target="a3", + backend="vpto", + mode="explicit", + source=self.SOURCE, + native_options=native_options, + ) + def native_option_probe(buf: pto.ptr(pto.f32, "gm")): + pass + + return native_option_probe + + def test_options_reach_the_cache_signature(self): + kernel = self._kernel({"host_sources": ["shim.cpp"], "link_libraries": ["dl"]}) + options = kernel.__ptodsl_cache_signature__()[9] + self.assertEqual(options.link_libraries, ("dl",)) + self.assertEqual(len(options.host_sources), 1) + # Resolved against this test file, which is what the option promises. + self.assertEqual(options.host_sources[0], Path(__file__).resolve().parent / "shim.cpp") + + def test_omitting_the_option_changes_nothing(self): + kernel = self._kernel(None) + self.assertTrue(kernel.__ptodsl_cache_signature__()[9].is_empty()) + + def test_a_bad_option_is_rejected_at_decoration_time(self): + with self.assertRaises(ValueError): + self._kernel({"host_sourcez": ["shim.cpp"]}) + + +class ExtraLinkFlagsTest(unittest.TestCase): + def test_no_options_contributes_nothing(self): + self.assertEqual(native_build._extra_link_flags(EMPTY_NATIVE_OPTIONS), []) + + def test_library_dirs_become_search_paths_and_rpaths(self): + options = NativeBuildOptions( + link_libraries=("dl", "m"), library_dirs=(Path("/opt/lib"),) + ) + self.assertEqual( + native_build._extra_link_flags(options), + ["-L/opt/lib", "-Wl,-rpath,/opt/lib", "-ldl", "-lm"], + ) + + def test_search_paths_precede_the_libraries_that_need_them(self): + flags = native_build._extra_link_flags( + NativeBuildOptions(link_libraries=("custom",), library_dirs=(Path("/opt/lib"),)) + ) + self.assertLess(flags.index("-L/opt/lib"), flags.index("-lcustom")) + + def test_host_sources_pull_libstdcxx(self): + flags = native_build._extra_link_flags( + NativeBuildOptions(host_sources=(Path("/tmp/shim.cpp"),), link_libraries=("dl",)) + ) + self.assertEqual(flags, ["-ldl", "-lstdc++"]) + + def test_explicit_libstdcxx_is_not_duplicated(self): + flags = native_build._extra_link_flags( + NativeBuildOptions( + host_sources=(Path("/tmp/shim.cpp"),), link_libraries=("dl", "stdc++") + ) + ) + self.assertEqual(flags, ["-ldl", "-lstdc++"]) + + +class HostSourceCompileTest(unittest.TestCase): + def test_no_host_sources_runs_no_compiler(self): + with mock.patch.object(native_build, "_run") as run: + objects = native_build._compile_host_sources( + EMPTY_NATIVE_OPTIONS, Path("/tmp/cache") + ) + self.assertEqual(objects, []) + run.assert_not_called() + + def test_a_missing_host_source_is_named_before_the_toolchain_is_resolved(self): + # Deliberately without stubbing the toolchain: a mistyped path must report + # itself, not whichever environment variable the compiler lookup wants. + options = NativeBuildOptions(host_sources=(Path("/nonexistent/shim.cpp"),)) + with mock.patch.object(native_build, "_run") as run: + with self.assertRaises(FileNotFoundError) as caught: + native_build._compile_host_sources(options, Path("/tmp/cache")) + self.assertIn("/nonexistent/shim.cpp", str(caught.exception)) + run.assert_not_called() + + def test_each_source_is_compiled_as_host_c_plus_plus_with_its_include_dirs(self): + with tempfile.TemporaryDirectory() as temp_dir: + cache_dir = Path(temp_dir) / "cache" + source = Path(temp_dir) / "shim.cpp" + source.write_text("int a() { return 1; }\n", encoding="utf-8") + options = NativeBuildOptions( + host_sources=(source,), include_dirs=(Path("/opt/inc"),) + ) + + with mock.patch.object(native_build, "resolve_bisheng", return_value="bisheng"): + with mock.patch.object( + native_build, "common_include_flags", return_value=["-I/cann/include"] + ): + with mock.patch.object(native_build, "_run") as run: + objects = native_build._compile_host_sources(options, cache_dir) + + self.assertEqual(run.call_count, 1) + cmd = run.call_args[0][0] + self.assertEqual(cmd[0], "bisheng") + # Host code, not device code: no -xcce and no aicore arch. + self.assertIn("-xc++", cmd) + self.assertNotIn("-xcce", cmd) + self.assertFalse([arg for arg in cmd if arg.startswith("--cce-aicore-arch")]) + self.assertIn("-I/opt/inc", cmd) + self.assertIn("-I/cann/include", cmd) + self.assertIn("-fPIC", cmd) + self.assertEqual(cmd[-4:], ["-c", str(source), "-o", str(objects[0])]) + + def test_sources_sharing_a_stem_get_distinct_objects(self): + with tempfile.TemporaryDirectory() as temp_dir: + root = Path(temp_dir) + first = root / "a" / "shim.cpp" + second = root / "b" / "shim.cpp" + for path in (first, second): + path.parent.mkdir(parents=True) + path.write_text("int a() { return 1; }\n", encoding="utf-8") + options = NativeBuildOptions(host_sources=(first, second)) + + with mock.patch.object(native_build, "resolve_bisheng", return_value="bisheng"): + with mock.patch.object(native_build, "common_include_flags", return_value=[]): + with mock.patch.object(native_build, "_run"): + objects = native_build._compile_host_sources(options, root / "cache") + + self.assertEqual(len(set(objects)), 2) + + +class LinkLineTest(unittest.TestCase): + def _link(self, **kwargs): + with mock.patch.object(native_build, "resolve_bisheng", return_value="bisheng"): + with mock.patch.object( + native_build, "runtime_library_flags", return_value=["-lruntime"] + ): + with mock.patch.object(native_build, "_run") as run: + native_build._link_shared_library( + Path("/c/launch.o"), + Path("/c/kernel.o"), + Path("/c/libk.so"), + kernel_kind="vector", + **kwargs, + ) + return run.call_args[0][0] + + def test_without_options_the_link_line_is_unchanged(self): + cmd = self._link() + self.assertEqual(cmd[-3:], ["/c/launch.o", "/c/kernel.o", "-lruntime"]) + + def test_host_objects_are_linked_and_extra_flags_come_last(self): + cmd = self._link( + host_objects=[Path("/c/host_0_shim.o")], + extra_link_flags=["-L/opt/lib", "-ldl"], + ) + self.assertIn("/c/host_0_shim.o", cmd) + # The host object precedes the libraries that resolve its references. + self.assertLess(cmd.index("/c/host_0_shim.o"), cmd.index("-ldl")) + self.assertLess(cmd.index("-lruntime"), cmd.index("-ldl")) + self.assertEqual(cmd[-2:], ["-L/opt/lib", "-ldl"]) + # Undefined symbols stay an error, so a missing entry fails the build. + self.assertIn("-Wl,--no-undefined", cmd) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/comm/AsyncWorkspace.h b/test/comm/AsyncWorkspace.h new file mode 100644 index 0000000000..e08a335a98 --- /dev/null +++ b/test/comm/AsyncWorkspace.h @@ -0,0 +1,678 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#ifndef PTO_TEST_COMM_ASYNCWORKSPACE_H +#define PTO_TEST_COMM_ASYNCWORKSPACE_H + +// Test scaffolding: builds the async workspace that pto.sdma_gm_gm reads. +// +// Not a shipped interface, and deliberately not under include/. PTOAS owns the +// layout contract in PTO/Support/AsyncSessionABI.h, which is what code +// generation emits against; producing the workspace belongs to the +// communication library. The supported way to get one is HCCL, which acquires +// channels per engine (HcclChannelAcquire with COMM_ENGINE_AIV) and hands back +// an engine context. This file exists so the device path could be exercised on +// hardware before that binding was written. +// +// What it does, since none of it can be computed on the host -- the queue ring +// base and the doorbell address are STARS facts only an AICPU query reports: +// +// 1. create one device-only stream per channel and collect its ids +// 2. allocate and zero the workspace in device memory +// 3. stage the stream table where the operator can read it +// 4. run aclnnShmemSdmaStarsQuery, which fills the channel descriptors +// 5. check those descriptors and repack them into channel records +// 6. hand the record table to the kernel as the session context +// +// Step 5 is what keeps the descriptor layout from reaching generated code. It +// is a CANN-release fact rather than a hardware one -- the pto headers have +// already moved it once -- so it is resolved here, once, where a change fails +// the check instead of producing a wrong address inside a kernel. +// +// Everything is resolved with dlopen, so a program that never constructs an +// AsyncWorkspace links and runs unchanged on a machine without a toolkit. +// +// Two caveats that make this unsuitable for shipping. The stream-table and +// operator-descriptor structs below are transcribed rather than included, so a +// layout change on the CANN side corrupts memory silently instead of failing to +// build; and rtStreamGetSqid, rtStreamGetCqid and rtGetDeviceInfo are runtime +// internals rather than public ACL. pto-isa carries a near-identical copy of +// all this in sdma_workspace_manager.hpp; both are readers of a format neither +// owns, which is why agreement between them proves nothing and the check in +// step 5 compares against values fed to the query instead. +// +// Requires CANN 9.0.0 or newer at run time, for aclnnShmemSdmaStarsQuery. That +// is stricter than the device code needs: the generated kernel itself compiles +// and runs under CANN 8.x. + +#if defined(__CCE_KT_TEST__) || defined(__CCE__) +#error "AsyncWorkspace.h is a host-only header and cannot be included in device code." +#endif + +#include +#include +#include +#include +#include +#include + +#include + +#include "PTO/Support/AsyncSessionABI.h" + +namespace mlir::pto::comm { + +namespace detail { + +// ACL values this header depends on, spelled out rather than included so the +// header stays free of CANN at build time. Taken from acl_rt.h / acl_base.h. +constexpr uint32_t kAclStreamFastLaunch = 0x00000001U; +constexpr uint32_t kAclStreamFastSync = 0x00000002U; +constexpr uint32_t kAclStreamDeviceUseOnly = 0x00000020U; +constexpr int32_t kAclStreamAttrFailureMode = 1; +constexpr int32_t kAclMemMallocHugeFirst = 0; +constexpr int32_t kAclMemcpyHostToDevice = 1; +constexpr int32_t kAclMemcpyDeviceToHost = 2; +constexpr int32_t kAclUint64 = 10; +constexpr int32_t kAclFormatND = 2; + +// rtGetDeviceInfo selector for the physical die id. +constexpr int32_t kInfoTypePhyDieId = 19; + +// One entry of the stream table the STARS query consumes. The operator matches +// on the exact size, so the padding is load-bearing. +struct HostStreamInfo { + uint64_t stream; + uint64_t ctx; + int32_t streamId; + uint32_t sqId; + uint32_t cqId; + uint32_t logicCqId; + uint64_t cqeAddr; + int32_t devId; + uint8_t reserved[20]; +}; +static_assert(sizeof(HostStreamInfo) == 64, "HostStreamInfo must be 64 bytes"); + +// Descriptor the operator reads to find the stream table and the workspace. +struct OpResInfo { + uint64_t size; + uint64_t streamsAddr; + uint64_t workspaceAddr; + uint8_t reserved[40]; +}; +static_assert(sizeof(OpResInfo) == 64, "OpResInfo must be 64 bytes"); + +} // namespace detail + +class AsyncWorkspace { +public: + AsyncWorkspace() = default; + ~AsyncWorkspace() { finalize(); } + + AsyncWorkspace(const AsyncWorkspace &) = delete; + AsyncWorkspace &operator=(const AsyncWorkspace &) = delete; + + // Builds the workspace for `channels` channels on the current device, which + // must already be selected with aclrtSetDevice. Returns false and leaves + // `error()` set on failure. + bool init(unsigned channels = workspace::kMaxChannels) { + if (inited_) + return true; + if (channels == 0 || channels > workspace::kMaxChannels) + return fail("channel count must be in [1, " + + std::to_string(workspace::kMaxChannels) + "]"); + + if (!loadSymbols() || !createStreams(channels) || !allocWorkspace() || + !stageStreamTable() || !runStarsQuery() || !validateAndRepack()) { + finalize(); + return false; + } + inited_ = true; + return true; + } + + void finalize() { + if (recordsDevice_) { + aclrtFree_(recordsDevice_); + recordsDevice_ = nullptr; + } + if (streamsDevice_) { + aclrtFree_(streamsDevice_); + streamsDevice_ = nullptr; + } + if (opResDevice_) { + aclrtFree_(opResDevice_); + opResDevice_ = nullptr; + } + if (workspaceDevice_) { + aclrtFree_(workspaceDevice_); + workspaceDevice_ = nullptr; + } + for (auto &s : streams_) { + if (s.stream) + aclrtDestroyStream_(reinterpret_cast(s.stream)); + } + streams_.clear(); + closeLibs(); + inited_ = false; + } + + // Pass this to the kernel as SessionField::ContextGm. It addresses the + // channel records this class builds, not the workspace the query wrote. + void *contextGm() const { return recordsDevice_; } + + // The workspace itself, which stays allocated because the live head and tail + // remain in it and the records point at them. + void *workspaceGm() const { return workspaceDevice_; } + + unsigned channelCount() const { + return static_cast(streams_.size()); + } + + // Channels the query reported populating, or 0 if it reported nothing. Only + // meaningful after a successful init. + uint32_t reportedQueueNum() const { return reportedQueueNum_; } + + const std::string &error() const { return error_; } + + // Reads one channel descriptor back to the host. Only useful for checking + // that the STARS query actually populated the table. + bool readChannelDescriptor(unsigned channelIdx, unsigned channelNum, + unsigned channelInGroup, + uint8_t (&out)[workspace::kChannelDescBytes]) { + if (!inited_) + return fail("workspace is not initialized"); + size_t offset = + workspace::channelDescOffset(channelIdx, channelNum, channelInGroup); + if (offset + workspace::kChannelDescBytes > workspace::kContextBytes) + return fail("channel descriptor lies outside the context region"); + void *src = static_cast(workspaceDevice_) + offset; + if (aclrtMemcpy_(out, sizeof(out), src, sizeof(out), + detail::kAclMemcpyDeviceToHost) != 0) + return fail("aclrtMemcpy of channel descriptor failed"); + return true; + } + + // Reads one channel record back, which is what a kernel sees for that + // channel. Only useful for checking that the table reached the device. + bool readChannelRecord(unsigned index, + uint8_t (&out)[comm::channel::kRecordBytes]) { + if (!inited_ || index >= streams_.size()) + return fail("no such channel record"); + void *src = static_cast(recordsDevice_) + + comm::channel::recordOffset(index, 1, 0); + if (aclrtMemcpy_(out, sizeof(out), src, sizeof(out), + detail::kAclMemcpyDeviceToHost) != 0) + return fail("aclrtMemcpy of channel record failed"); + return true; + } + +private: + // Checks what the query wrote and repacks it into the channel records. + // + // The check is exact rather than a plausibility test, because the query was + // handed the stream, sq, cq and device ids and echoes them back: a descriptor + // whose ids sit where they are expected is at the layout this build assumes, + // and one byte of drift breaks the match. Only the fields the query invents + // -- the ring base, the doorbell and the depth -- fall back to a range check. + bool validateAndRepack() { + const unsigned channels = static_cast(streams_.size()); + + uint8_t header[workspace::kFlagInfoBytes]; + if (!readWorkspace(0, header, sizeof(header))) + return fail("reading the workspace flag-info header failed"); + std::memcpy(&reportedQueueNum_, + header + workspace::kFlagInfoTotalQueueNumOffset, + sizeof(reportedQueueNum_)); + + // Trust the reported count when it is sane; a query that reports fewer + // channels than asked has populated fewer descriptors, and indexing past + // them would read the zeroed allocation. + if (reportedQueueNum_ != 0 && reportedQueueNum_ < channels) + return fail("the query populated " + std::to_string(reportedQueueNum_) + + " channels but " + std::to_string(channels) + + " were requested"); + + std::vector records(static_cast(channels) * + comm::channel::kRecordBytes, + 0); + + for (unsigned i = 0; i < channels; ++i) { + const size_t descOffset = + workspace::kFlagInfoBytes + + static_cast(i) * workspace::kChannelDescBytes; + uint8_t desc[workspace::kChannelDescBytes]; + if (!readWorkspace(descOffset, desc, sizeof(desc))) + return fail("reading channel descriptor " + std::to_string(i) + + " failed"); + + const detail::HostStreamInfo &want = streams_[i]; + if (!expectU32(desc, workspace::kChannelStreamIdOffset, + static_cast(want.streamId), "stream_id", i) || + !expectU32(desc, workspace::kChannelSqIdOffset, want.sqId, "sq_id", + i) || + !expectU32(desc, workspace::kChannelCqIdOffset, want.cqId, "cq_id", + i) || + !expectU32(desc, workspace::kChannelLogicCqIdOffset, want.logicCqId, + "logic_cq_id", i) || + !expectU32(desc, workspace::kChannelDevIdOffset, + static_cast(want.devId), "dev_id", i)) + return false; + + uint64_t sqBase = 0, sqRegBase = 0; + uint32_t sqDepth = 0, sqHead = 0, sqTail = 0; + std::memcpy(&sqBase, desc + workspace::kChannelSqBaseOffset, + sizeof(sqBase)); + std::memcpy(&sqRegBase, desc + workspace::kChannelSqRegBaseOffset, + sizeof(sqRegBase)); + std::memcpy(&sqDepth, desc + workspace::kChannelSqDepthOffset, + sizeof(sqDepth)); + std::memcpy(&sqHead, desc + workspace::kChannelSqHeadOffset, + sizeof(sqHead)); + std::memcpy(&sqTail, desc + workspace::kChannelSqTailOffset, + sizeof(sqTail)); + + if (sqBase == 0 || (sqBase & (workspace::kChannelDescBytes - 1)) != 0) + return fail("channel " + std::to_string(i) + " has sq_base " + + hex(sqBase) + ", which is not a usable ring base"); + if (sqRegBase == 0) + return fail("channel " + std::to_string(i) + " has no doorbell address"); + // The ring the engine actually uses is a power of two; generated code + // wraps with a mask. A5's STARS query on CANN 9.2.0 reports 2049 for a + // 2048-slot ring, so a count that is one past a power of two is taken as + // that ring rather than rejected. Any other non-power-of-two stays a + // hard error: masking it would land posts in the wrong slot. + uint32_t ringDepth = sqDepth; + if (ringDepth > 1 && (ringDepth & (ringDepth - 1)) != 0 && + ((ringDepth - 1) & (ringDepth - 2)) == 0) + ringDepth = sqDepth - 1; + if (ringDepth == 0 || (ringDepth & (ringDepth - 1)) != 0) + return fail("channel " + std::to_string(i) + " reports depth " + + std::to_string(sqDepth) + + ", which is not a power of two and cannot be masked"); + if (sqHead >= ringDepth || sqTail >= ringDepth) + return fail("channel " + std::to_string(i) + " starts with head " + + std::to_string(sqHead) + " tail " + std::to_string(sqTail) + + " outside a queue of " + std::to_string(ringDepth)); + + const uint64_t descGm = + reinterpret_cast(workspaceDevice_) + descOffset; + uint8_t *rec = records.data() + comm::channel::recordOffset(i, 1, 0); + const uint64_t tailAddr = descGm + workspace::kChannelSqTailOffset; + const uint64_t headAddr = descGm + workspace::kChannelSqHeadOffset; + const uint32_t streamId = static_cast(want.streamId); + const uint32_t slotMask = ringDepth - 1; + std::memcpy(rec + comm::channel::kSqBaseOffset, &sqBase, sizeof(sqBase)); + std::memcpy(rec + comm::channel::kDoorbellOffset, &sqRegBase, + sizeof(sqRegBase)); + std::memcpy(rec + comm::channel::kTailAddrOffset, &tailAddr, + sizeof(tailAddr)); + std::memcpy(rec + comm::channel::kHeadAddrOffset, &headAddr, + sizeof(headAddr)); + std::memcpy(rec + comm::channel::kSlotMaskOffset, &slotMask, + sizeof(slotMask)); + std::memcpy(rec + comm::channel::kStreamIdOffset, &streamId, + sizeof(streamId)); + } + + if (aclrtMalloc_(&recordsDevice_, records.size(), + detail::kAclMemMallocHugeFirst) != 0) + return fail("aclrtMalloc of the channel record table failed"); + if (aclrtMemcpy_(recordsDevice_, records.size(), records.data(), + records.size(), detail::kAclMemcpyHostToDevice) != 0) + return fail("aclrtMemcpy of the channel record table failed"); + return true; + } + + bool readWorkspace(size_t offset, void *out, size_t bytes) { + if (offset + bytes > workspace::kContextBytes) + return false; + void *src = static_cast(workspaceDevice_) + offset; + return aclrtMemcpy_(out, bytes, src, bytes, + detail::kAclMemcpyDeviceToHost) == 0; + } + + // Compares one echoed id, and on a mismatch says where the value did land, so + // a layout that shifted reports the shift instead of just a bad number. + bool expectU32(const uint8_t *desc, size_t offset, uint32_t want, + const char *what, unsigned channel) { + uint32_t got = 0; + std::memcpy(&got, desc + offset, sizeof(got)); + if (got == want) + return true; + + std::string where; + for (size_t probe = 0; probe + sizeof(uint32_t) <= workspace::kChannelDescBytes; + probe += sizeof(uint32_t)) { + uint32_t at = 0; + std::memcpy(&at, desc + probe, sizeof(at)); + if (at == want) + where += (where.empty() ? "" : ", ") + std::to_string(probe); + } + return fail( + "channel " + std::to_string(channel) + ": expected " + what + " " + + std::to_string(want) + " at descriptor offset " + + std::to_string(offset) + " but found " + std::to_string(got) + + (where.empty() + ? "; it appears nowhere in the descriptor, so the query did not " + "run as expected" + : "; it appears at offset " + where + + ", so the descriptor layout has moved")); + } + + static std::string hex(uint64_t v) { + char buf[32]; + std::snprintf(buf, sizeof(buf), "%#llx", + static_cast(v)); + return buf; + } + + using AclrtCreateStreamWithConfigFn = int32_t (*)(void **, uint32_t, uint32_t); + using AclrtDestroyStreamFn = int32_t (*)(void *); + using AclrtStreamGetIdFn = int32_t (*)(void *, int32_t *); + using AclrtGetCurrentContextFn = int32_t (*)(void **); + using AclrtGetDeviceFn = int32_t (*)(int32_t *); + using AclrtMallocFn = int32_t (*)(void **, size_t, int32_t); + using AclrtFreeFn = int32_t (*)(void *); + using AclrtMemsetFn = int32_t (*)(void *, size_t, int32_t, size_t); + using AclrtMemcpyFn = int32_t (*)(void *, size_t, const void *, size_t, + int32_t); + using AclrtSynchronizeStreamFn = int32_t (*)(void *); + using AclrtSetStreamAttributeFn = int32_t (*)(void *, int32_t, const void *); + + using RtStreamGetSqidFn = int32_t (*)(const void *, uint32_t *); + using RtStreamGetCqidFn = int32_t (*)(const void *, uint32_t *, uint32_t *); + using RtGetDeviceInfoFn = int32_t (*)(uint32_t, int32_t, int32_t, int64_t *); + + using AclCreateTensorFn = void *(*)(const int64_t *, uint64_t, int32_t, + const int64_t *, int64_t, int32_t, + const int64_t *, uint64_t, void *); + using AclDestroyTensorFn = int32_t (*)(void *); + using StarsQueryGetWsSizeFn = int32_t (*)(const void *, void *, uint64_t *, + void **); + using StarsQueryExecFn = int32_t (*)(void *, uint64_t, void *, void *); + + bool fail(const std::string &what) { + if (error_.empty()) + error_ = what; + return false; + } + + template + bool bind(Fn &slot, void *handle, const char *name) { + slot = reinterpret_cast(dlsym(handle, name)); + if (!slot) + return fail(std::string("dlsym ") + name + " failed"); + return true; + } + + bool loadSymbols() { + aclHandle_ = dlopen("libascendcl.so", RTLD_NOW | RTLD_GLOBAL); + if (!aclHandle_) + return fail(std::string("dlopen libascendcl.so failed: ") + dlerror()); + rtHandle_ = dlopen("libruntime.so", RTLD_NOW); + if (!rtHandle_) + return fail(std::string("dlopen libruntime.so failed: ") + dlerror()); + opapiHandle_ = dlopen("libopapi.so", RTLD_NOW); + if (!opapiHandle_) + return fail(std::string("dlopen libopapi.so failed: ") + dlerror()); + nnopHandle_ = dlopen("libnnopbase.so", RTLD_NOW | RTLD_GLOBAL); + if (!nnopHandle_) + return fail(std::string("dlopen libnnopbase.so failed: ") + dlerror()); + + return bind(aclrtCreateStreamWithConfig_, aclHandle_, + "aclrtCreateStreamWithConfig") && + bind(aclrtDestroyStream_, aclHandle_, "aclrtDestroyStream") && + bind(aclrtStreamGetId_, aclHandle_, "aclrtStreamGetId") && + bind(aclrtGetCurrentContext_, aclHandle_, "aclrtGetCurrentContext") && + bind(aclrtGetDevice_, aclHandle_, "aclrtGetDevice") && + bind(aclrtMalloc_, aclHandle_, "aclrtMalloc") && + bind(aclrtFree_, aclHandle_, "aclrtFree") && + bind(aclrtMemset_, aclHandle_, "aclrtMemset") && + bind(aclrtMemcpy_, aclHandle_, "aclrtMemcpy") && + bind(aclrtSynchronizeStream_, aclHandle_, + "aclrtSynchronizeStream") && + bind(aclrtSetStreamAttribute_, aclHandle_, + "aclrtSetStreamAttribute") && + bind(rtStreamGetSqid_, rtHandle_, "rtStreamGetSqid") && + bind(rtStreamGetCqid_, rtHandle_, "rtStreamGetCqid") && + bind(rtGetDeviceInfo_, rtHandle_, "rtGetDeviceInfo") && + bind(aclCreateTensor_, nnopHandle_, "aclCreateTensor") && + bind(aclDestroyTensor_, nnopHandle_, "aclDestroyTensor") && + bind(starsQueryGetWsSize_, opapiHandle_, + "aclnnShmemSdmaStarsQueryGetWorkspaceSize") && + bind(starsQueryExec_, opapiHandle_, "aclnnShmemSdmaStarsQuery"); + } + + void closeLibs() { + for (void **h : {&opapiHandle_, &nnopHandle_, &rtHandle_, &aclHandle_}) { + if (*h) { + dlclose(*h); + *h = nullptr; + } + } + } + + bool createStreams(unsigned channels) { + int32_t deviceId = -1; + if (aclrtGetDevice_(&deviceId) != 0) + return fail("aclrtGetDevice failed; call aclrtSetDevice first"); + + int64_t dieId = -1; + if (rtGetDeviceInfo_(static_cast(deviceId), 0, + detail::kInfoTypePhyDieId, &dieId) != 0) + return fail("rtGetDeviceInfo(die id) failed"); + + streams_.resize(channels); + for (unsigned i = 0; i < channels; ++i) { + detail::HostStreamInfo &info = streams_[i]; + std::memset(&info, 0, sizeof(info)); + + void *stream = nullptr; + if (aclrtCreateStreamWithConfig_(&stream, 0, + detail::kAclStreamDeviceUseOnly) != 0) + return fail("aclrtCreateStreamWithConfig failed on channel " + + std::to_string(i)); + info.stream = reinterpret_cast(stream); + + int32_t streamId = 0; + if (aclrtStreamGetId_(stream, &streamId) != 0) + return fail("aclrtStreamGetId failed on channel " + std::to_string(i)); + + uint32_t sqId = 0; + if (rtStreamGetSqid_(stream, &sqId) != 0) + return fail("rtStreamGetSqid failed on channel " + std::to_string(i)); + + uint32_t cqId = 0; + uint32_t logicCqId = 0; + if (rtStreamGetCqid_(stream, &cqId, &logicCqId) != 0) + return fail("rtStreamGetCqid failed on channel " + std::to_string(i)); + + void *ctx = nullptr; + if (aclrtGetCurrentContext_(&ctx) != 0) + return fail("aclrtGetCurrentContext failed on channel " + + std::to_string(i)); + + info.ctx = reinterpret_cast(ctx); + info.streamId = streamId; + info.sqId = sqId; + info.cqId = cqId; + info.logicCqId = logicCqId; + info.devId = static_cast(dieId); + } + return true; + } + + bool allocWorkspace() { + if (aclrtMalloc_(&workspaceDevice_, workspace::kTotalBytes, + detail::kAclMemMallocHugeFirst) != 0) + return fail("aclrtMalloc of the async workspace failed"); + // The kernel reads descriptors the query does not touch, so start from a + // known state rather than whatever the allocator returned. + if (aclrtMemset_(workspaceDevice_, workspace::kTotalBytes, 0, + workspace::kTotalBytes) != 0) + return fail("aclrtMemset of the async workspace failed"); + return true; + } + + bool stageStreamTable() { + size_t streamsBytes = streams_.size() * sizeof(detail::HostStreamInfo); + if (aclrtMalloc_(&streamsDevice_, streamsBytes, + detail::kAclMemMallocHugeFirst) != 0) + return fail("aclrtMalloc of the stream table failed"); + if (aclrtMemcpy_(streamsDevice_, streamsBytes, streams_.data(), + streamsBytes, detail::kAclMemcpyHostToDevice) != 0) + return fail("aclrtMemcpy of the stream table failed"); + + detail::OpResInfo opRes{}; + opRes.size = streams_.size(); + opRes.streamsAddr = reinterpret_cast(streamsDevice_); + opRes.workspaceAddr = reinterpret_cast(workspaceDevice_); + + if (aclrtMalloc_(&opResDevice_, sizeof(opRes), + detail::kAclMemMallocHugeFirst) != 0) + return fail("aclrtMalloc of the operator descriptor failed"); + if (aclrtMemcpy_(opResDevice_, sizeof(opRes), &opRes, sizeof(opRes), + detail::kAclMemcpyHostToDevice) != 0) + return fail("aclrtMemcpy of the operator descriptor failed"); + return true; + } + + // A device buffer plus the tensor that wraps it, released together. + struct TensorHandle { + void *data{nullptr}; + void *tensor{nullptr}; + }; + + bool makeScalarTensor(const std::vector &values, + TensorHandle &out) { + const int64_t shape = static_cast(values.size()); + const int64_t stride = 1; + size_t bytes = values.size() * sizeof(uint64_t); + + if (aclrtMalloc_(&out.data, bytes, detail::kAclMemMallocHugeFirst) != 0) + return fail("aclrtMalloc of a query tensor failed"); + if (aclrtMemcpy_(out.data, bytes, values.data(), bytes, + detail::kAclMemcpyHostToDevice) != 0) + return fail("aclrtMemcpy of a query tensor failed"); + + out.tensor = aclCreateTensor_(&shape, 1, detail::kAclUint64, &stride, 0, + detail::kAclFormatND, &shape, 1, out.data); + if (!out.tensor) + return fail("aclCreateTensor failed"); + return true; + } + + void releaseTensor(TensorHandle &handle) { + if (handle.tensor) { + aclDestroyTensor_(handle.tensor); + handle.tensor = nullptr; + } + if (handle.data) { + aclrtFree_(handle.data); + handle.data = nullptr; + } + } + + bool runStarsQuery() { + void *stream = nullptr; + if (aclrtCreateStreamWithConfig_(&stream, 0, + detail::kAclStreamFastLaunch | + detail::kAclStreamFastSync) != 0) + return fail("aclrtCreateStreamWithConfig for the query stream failed"); + + // Ask the runtime to report a failed query rather than tearing the process + // down, so a missing STARS backend surfaces as a return code. + uint32_t failureMode = 1; + aclrtSetStreamAttribute_(stream, detail::kAclStreamAttrFailureMode, + &failureMode); + + TensorHandle input; + TensorHandle output; + void *aclnnWs = nullptr; + bool ok = false; + + do { + if (!makeScalarTensor({reinterpret_cast(opResDevice_), + reinterpret_cast(workspaceDevice_)}, + input)) + break; + if (!makeScalarTensor({0}, output)) + break; + + uint64_t wsSize = 0; + void *executor = nullptr; + if (starsQueryGetWsSize_(input.tensor, output.tensor, &wsSize, + &executor) != 0) { + fail("aclnnShmemSdmaStarsQueryGetWorkspaceSize failed"); + break; + } + if (wsSize > 0 && aclrtMalloc_(&aclnnWs, wsSize, + detail::kAclMemMallocHugeFirst) != 0) { + fail("aclrtMalloc of the operator workspace failed"); + break; + } + if (starsQueryExec_(aclnnWs, wsSize, executor, stream) != 0) { + fail("aclnnShmemSdmaStarsQuery failed"); + break; + } + if (aclrtSynchronizeStream_(stream) != 0) { + fail("aclrtSynchronizeStream after the query failed"); + break; + } + ok = true; + } while (false); + + if (aclnnWs) + aclrtFree_(aclnnWs); + releaseTensor(input); + releaseTensor(output); + aclrtDestroyStream_(stream); + return ok; + } + + bool inited_{false}; + std::string error_; + uint32_t reportedQueueNum_{0}; + std::vector streams_; + void *recordsDevice_{nullptr}; + void *workspaceDevice_{nullptr}; + void *streamsDevice_{nullptr}; + void *opResDevice_{nullptr}; + + void *aclHandle_{nullptr}; + void *rtHandle_{nullptr}; + void *opapiHandle_{nullptr}; + void *nnopHandle_{nullptr}; + + AclrtCreateStreamWithConfigFn aclrtCreateStreamWithConfig_{nullptr}; + AclrtDestroyStreamFn aclrtDestroyStream_{nullptr}; + AclrtStreamGetIdFn aclrtStreamGetId_{nullptr}; + AclrtGetCurrentContextFn aclrtGetCurrentContext_{nullptr}; + AclrtGetDeviceFn aclrtGetDevice_{nullptr}; + AclrtMallocFn aclrtMalloc_{nullptr}; + AclrtFreeFn aclrtFree_{nullptr}; + AclrtMemsetFn aclrtMemset_{nullptr}; + AclrtMemcpyFn aclrtMemcpy_{nullptr}; + AclrtSynchronizeStreamFn aclrtSynchronizeStream_{nullptr}; + AclrtSetStreamAttributeFn aclrtSetStreamAttribute_{nullptr}; + + RtStreamGetSqidFn rtStreamGetSqid_{nullptr}; + RtStreamGetCqidFn rtStreamGetCqid_{nullptr}; + RtGetDeviceInfoFn rtGetDeviceInfo_{nullptr}; + + AclCreateTensorFn aclCreateTensor_{nullptr}; + AclDestroyTensorFn aclDestroyTensor_{nullptr}; + StarsQueryGetWsSizeFn starsQueryGetWsSize_{nullptr}; + StarsQueryExecFn starsQueryExec_{nullptr}; +}; + +} // namespace mlir::pto::comm + +#endif // PTO_TEST_COMM_ASYNCWORKSPACE_H diff --git a/test/comm/AsyncWorkspaceShim.cpp b/test/comm/AsyncWorkspaceShim.cpp new file mode 100644 index 0000000000..dfe698f8cb --- /dev/null +++ b/test/comm/AsyncWorkspaceShim.cpp @@ -0,0 +1,213 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// A C entry surface over AsyncWorkspace, so a Python test can drive the engine +// path without a second copy of the session ABI. +// +// The point of this file is what it does not expose. Every field position, slot +// width and default in AsyncSessionABI.h stays on this side: a caller asks for a +// session by naming values, never offsets, and gets back an opaque byte image. +// Nothing here hands out a struct layout for a caller to mirror, because a +// mirrored layout is exactly the second list of fields that header exists to +// prevent. +// +// It also stays free of CANN headers. AsyncWorkspace resolves everything by +// dlopen, so this translation unit compiles with a plain host C++ compiler and +// needs only -ldl at link time; the toolkit is required to run, not to build. +// +// One workspace per process, addressed as a singleton. A test process needs one, +// the queue state it wraps is per-device, and a handle would add a parameter to +// every call while buying nothing. Cleanup rides on the static's destructor, so +// a caller that exits without calling finalize still releases the device memory +// and streams. + +#include +#include +#include + +#include "AsyncWorkspace.h" + +namespace { + +using mlir::pto::comm::AsyncWorkspace; +namespace comm = mlir::pto::comm; + +// Function-local rather than a namespace-scope object, so construction order +// across translation units never matters. +AsyncWorkspace &workspace() { + static AsyncWorkspace ws; + return ws; +} + +// Separate from AsyncWorkspace::error() so a shim-level rejection, such as a +// short output buffer, reads back through the same call as a workspace failure. +std::string &lastError() { + static std::string error; + return error; +} + +int fail(const char *what) { + lastError() = what; + return -1; +} + +} // namespace + +extern "C" { + +// --- lifecycle --------------------------------------------------------------- + +// Build the workspace for `channels` channels on the currently selected device. +// Returns 0 on success, -1 with pto_async_ws_error() set otherwise. +// +// The device must already be selected. Under the PTODSL ST harness that is +// torch_npu's doing, which is the arrangement this shim exists to test: the +// workspace has to attach to a device and context someone else set up. +int pto_async_ws_init(unsigned channels) { + lastError().clear(); + if (!workspace().init(channels)) { + lastError() = workspace().error(); + return -1; + } + return 0; +} + +void pto_async_ws_finalize() { + lastError().clear(); + workspace().finalize(); +} + +// Valid until the next call into this shim, which may overwrite the buffer. +// Callers are expected to copy it, as a ctypes c_char_p restype does. +const char *pto_async_ws_error() { + if (!lastError().empty()) { + return lastError().c_str(); + } + return workspace().error().c_str(); +} + +// --- what the kernel needs --------------------------------------------------- + +// SessionField::ContextGm, or 0 before a successful init. +// +// A device address has to cross this boundary as an integer, which is the same +// reason AsyncSessionABI.h carries pointers as i64: the session struct rejects +// pointer fields. It stays opaque either way -- it is produced here, consumed by +// pto_async_session_pack, and never interpreted by the caller. +uint64_t pto_async_ws_context_gm() { + return reinterpret_cast(workspace().contextGm()); +} + +unsigned pto_async_ws_channel_count() { return workspace().channelCount(); } + +// Channels the STARS query reported populating, which bounds what a session may +// index however many channels were asked for. +uint32_t pto_async_ws_reported_queue_num() { + return workspace().reportedQueueNum(); +} + +// --- observation ------------------------------------------------------------- + +// The live queue tail for one channel, or -1 on failure so a caller can tell a +// failed read from a tail that legitimately reads zero. +// +// This is the only way to see that a post happened the way it was asked to. +// Correct data at the destination says the bytes moved; one oversized entry +// moves the same bytes as a correct split, and only the tail separates them. +long long pto_async_ws_sq_tail(unsigned channel_idx, unsigned channel_num) { + static_assert(comm::workspace::kChannelSqTailOffset + sizeof(uint32_t) <= + comm::workspace::kChannelDescBytes, + "the queue tail must lie inside one channel descriptor"); + + lastError().clear(); + uint8_t desc[comm::workspace::kChannelDescBytes]; + if (!workspace().readChannelDescriptor(channel_idx, channel_num, 0, desc)) { + lastError() = workspace().error(); + return -1; + } + uint32_t tail = 0; + std::memcpy(&tail, desc + comm::workspace::kChannelSqTailOffset, + sizeof(tail)); + return static_cast(tail); +} + +// --- session template -------------------------------------------------------- + +// How large a buffer pto_async_session_pack needs. Asked for rather than +// assumed, so the slot width stays a fact of the C++ header. +uint32_t pto_async_session_bytes() { return comm::session_tmpl::kBytes; } + +// Values and limits a caller would otherwise hard-code. +uint32_t pto_async_session_qos_default() { return comm::sqe::kQosDefault; } + +uint32_t pto_async_ws_min_transfer_bytes() { + return comm::workspace::kMinTransferBytes; +} + +uint32_t pto_async_ws_max_channels() { return comm::workspace::kMaxChannels; } + +// Build one session image into `out`, which must hold at least +// pto_async_session_bytes() bytes. Returns 0 on success, -1 otherwise. +// +// Arguments are named values rather than a struct on purpose: a struct would +// have a layout the caller has to agree with, which is the duplication this shim +// avoids. ContextGm comes from the live workspace, and Engine and Flags have +// only one sensible value here, so none of the three is a parameter. +// +// tmp_buf_addr and sync_id matter on A2/A3, where the doorbell is reachable only +// by MTE and the tail therefore goes out through UB. A5 writes it with st_dev +// and reads neither. +int pto_async_session_pack(void *out, uint32_t out_bytes, uint32_t channel_idx, + uint32_t channel_num, uint64_t block_bytes, + uint64_t tmp_buf_addr, uint32_t tmp_buf_size, + uint32_t sync_id, uint64_t comm_block_offset, + uint32_t qos, uint32_t dest_rank_id) { + // The destination bound is checked below; this is the source side. + static_assert(comm::session_tmpl::Builder::bytes() == + comm::session_tmpl::kBytes, + "the builder must hold exactly one session template"); + + lastError().clear(); + if (out == nullptr) { + return fail("session output buffer is null"); + } + if (out_bytes < comm::session_tmpl::kBytes) { + return fail("session output buffer is shorter than the session template"); + } + + // A zero block size would make the engine expansion's trip count degenerate. + // It clamps the value at run time rather than dividing by it, so this is a + // caller error to report here instead of a wrong transfer to discover later. + if (block_bytes == 0) { + return fail("session block_bytes must be non-zero"); + } + if (channel_num == 0) { + return fail("session channel_num must be non-zero"); + } + + comm::session_tmpl::Builder tmpl; + tmpl.set(comm::SessionField::ContextGm, pto_async_ws_context_gm()) + .set(comm::SessionField::TmpBufAddr, tmp_buf_addr) + .set(comm::SessionField::TmpBufSize, tmp_buf_size) + .set(comm::SessionField::SyncId, sync_id) + .set(comm::SessionField::ChannelIdx, channel_idx) + .set(comm::SessionField::ChannelNum, channel_num) + .set(comm::SessionField::BlockBytes, block_bytes) + .set(comm::SessionField::CommBlockOffset, comm_block_offset) + .set(comm::SessionField::Engine, + static_cast(comm::SessionEngine::Sdma)) + .set(comm::SessionField::DestRankId, dest_rank_id) + .set(comm::SessionField::QpIdx, 0) + .set(comm::SessionField::Flags, comm::kSessionFlagValid) + .set(comm::SessionField::Qos, qos); + + std::memcpy(out, tmpl.data(), comm::session_tmpl::kBytes); + return 0; +} + +} // extern "C" diff --git a/test/comm/README.md b/test/comm/README.md new file mode 100644 index 0000000000..51cd5ace5e --- /dev/null +++ b/test/comm/README.md @@ -0,0 +1,156 @@ +# Async communication host helpers + +Host-side support for `pto.sdma_gm_gm` and the session template a kernel loads +with `pto.session_init`. + +| File | Role | +| --- | --- | +| `AsyncWorkspace.h` | Builds the async workspace: streams, the AICPU STARS query, and the channel record table a kernel reads. Resolves every CANN symbol with `dlopen`. | +| `AsyncWorkspaceShim.cpp` | `extern "C"` surface over that header, so Python can drive it without a second copy of the session ABI. | +| `async_workspace.py` | `ctypes` wrapper over the shim. | +| `spike_sdma_gm_gm_engine.py` | The engine-path spike described below. | +| `build_async_shim.sh` | Standalone build of the shim into `build/libpto_async_shim.so`. Not needed by the spike, which gets the shim through the kernel build; useful to check the shim compiles on a machine with no card. | + +## How the shim reaches Python + +The spike does not build or load a separate library. It lists the shim in its +kernel declaration: + +```python +@pto.jit( + name="sdma_gm_gm_engine", + native_options={ + "host_sources": ["AsyncWorkspaceShim.cpp"], + "include_dirs": [".", "../../include"], + "link_libraries": ["dl"], + }, + ... +) +``` + +PTODSL compiles that as host C++ and links it into the kernel's own shared +library, so `compiled.native_library()` hands back one library carrying both the +launch entry and the shim. One build, one artifact, and the shim's contents are +part of the build's cache key — editing it rebuilds rather than silently reusing. +See `ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md`, "Host C++ in the +kernel's library". + +## Why a shim rather than a Python port + +Every session field index, slot width, SQE constant and channel-record offset +lives in `include/PTO/Support/AsyncSessionABI.h`, and the expansion in +`lib/PTO/Transforms/VPTOExpandWrapperOps.cpp` reads them from there. A Python +test that filled a session template by index would be a second, unchecked copy +of that layout, and the header exists specifically to keep there from being two. + +So the shim takes named values and returns an opaque byte image. It asks the C++ +side how large a session template is; it never says. `async_workspace.py` +contains no offsets at all, and adding a session field cannot silently break it. + +The shim needs no CANN headers or libraries to build, because `AsyncWorkspace` +resolves the toolkit at run time through `dlopen`. `-ldl` is the only library +the shim names; the kernel build also links `-lstdc++` because the fat-object +link does not pull the C++ runtime. It still compiles on a machine with no +driver and no card. + +## The engine-path spike + +`test/vpto/cases/async-comm/sdma_gm_gm.py` covers the A5 `{soft_put}` form of +`pto.sdma_gm_gm` as an ordinary golden ST case. It can do that because +`{soft_put}` expands to a synchronous GM→UB→GM copy: the transfer is finished +when the kernel returns. + +The engine form is not like that. The kernel writes SQEs, publishes the queue +tail and rings a doorbell; the SDMA engine moves the bytes afterwards. Running it +under the PTODSL ST harness needs three things to be true, and none of them can +be checked without a card: + +1. `AsyncWorkspace` can attach to the device and context `torch_npu` already set + up, rather than needing to own `aclrtSetDevice` itself. +2. Memory the workspace allocates through the runtime is addressable by a kernel + the harness launched, so a raw device address can be handed over inside a + session template. +3. A destination the engine fills can be observed from the harness, which + synchronizes the kernel and knows nothing about the engine. + +`spike_sdma_gm_gm_engine.py` is what answers that. It is deliberately not under +`test/vpto/cases/`: `run_host_vpto_validation.sh` discovers and runs every `.py` +in that tree, and a case needing a hand-built shim and a CANN 9.0 toolkit would +fail the whole validation run everywhere else. It also has no skip path, because +a spike that quietly passes when it could not run answers nothing. + +If all three hold, the engine path needs no change to the ST framework, and the +C++ and shell harness these cases used to require has nothing left to do. If one +fails, the spike reports which. + +### Running it + +Requires CANN ≥ 9.0.0 for `aclnnShmemSdmaStarsQuery`, a real device, and a +`torch_npu` the harness can import. The shim is built as part of the kernel, so +there is no separate build step. + +`PTO_ASYNC_ARCH` has to name the generation of the card, and both are worth +running: the three assumptions above are generation-independent, but the doorbell +is not, so A5 and A2/A3 exercise different final writes. + +```bash +PTO_ASYNC_ARCH=a5 python3 test/comm/spike_sdma_gm_gm_engine.py +PTO_ASYNC_ARCH=a3 python3 test/comm/spike_sdma_gm_gm_engine.py +``` + +Start with A5. It is where `AsyncWorkspace` has already been through a live STARS +query — the 2049-for-2048 ring workaround in `AsyncWorkspace.h` was observed +there on CANN 9.2.0 — and the A2/A3 doorbell path is the one that can take a card +down if the sequence is wrong. + +To check that the shim itself compiles without a card or a kernel: + +```bash +test/comm/build_async_shim.sh +``` + +Useful before touching a card: + +```bash +PTO_ASYNC_ARCH=a5 python3 test/comm/spike_sdma_gm_gm_engine.py --list # case names +PTO_ASYNC_ARCH=a5 python3 test/comm/spike_sdma_gm_gm_engine.py --emit-mlir # the kernel, no device needed +``` + +Knobs: + +| Variable | Default | Meaning | +| --- | --- | --- | +| `PTO_ASYNC_ARCH` | **none, required** | `a2`, `a3` or `a5` — the generation of the card being pointed at. The doorbell is the one part of the post that differs: A2/A3 reaches it only by MTE and stages the tail in UB, A5 writes it with `st_dev` at a different offset. Deliberately not defaulted, because a store aimed at `sq_reg_base` the wrong way can leave an A2/A3 card in an unrecoverable RAS state. `a2` and `a3` generate the same code. | +| `PTO_ASYNC_POLL_TIMEOUT_MS` | `2000` | How long to wait for the engine to drain before calling the destination wrong. | +| `PTO_ASYNC_SHIM` | `test/comm/build/libpto_async_shim.so` | Shim path, for a caller loading a standalone build rather than taking it from a kernel library. | +| `CXX` | `g++` | Host compiler for `build_async_shim.sh`. | + +### What each case separates + +All four run the same kernel and differ only in the session template the host +writes, which is the point: the transfer is described by session data, not by the +kernel. `block_bytes` is what splits one transfer into SQEs, so it decides how +many entries the queue should gain. + +| Case | Transfer | Entries | +| --- | --- | --- | +| `single_entry` | 4096 bytes in 4096-byte blocks | 1 | +| `even_split` | 4096 bytes in 1024-byte blocks | 4 | +| `ragged_tail` | 4096 bytes in 1536-byte blocks | 3, the last one short | +| `comm_block_offset` | 2048 bytes at offset 2048 of an 8 KiB buffer | 1 | + +Each case checks the queue tail as well as the data. Correct bytes at the +destination only say the transfer happened; one oversized entry moves the same +bytes as a correct split, and the tail is what separates them. A correct entry +count with an unfilled destination is reported differently from a wrong entry +count, because the first means the kernel did its part and the engine did not. + +### If the spike passes + +Promoting it is a rename into `test/vpto/cases/async-comm/`; the shim needs no +step in `run_host_vpto_validation.sh`, because the kernel build already carries +it. After that, the two-device window shapes in `test/lit/vpto/async_*.pto` become +the next thing to give a runtime case. That needs `HcclWindows.h`, which went away +with the C++ cases that were its only caller; it comes back through the same shim +pattern and the same `native_options` entry, so nothing here has to change to +accommodate it. diff --git a/test/comm/async_workspace.py b/test/comm/async_workspace.py new file mode 100644 index 0000000000..d211acc18e --- /dev/null +++ b/test/comm/async_workspace.py @@ -0,0 +1,269 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""ctypes access to the AsyncWorkspace shim, for engine-path ST cases. + +This module deliberately contains no session field indices, slot widths, SQE +constants or channel-record offsets. Every one of those stays in +``include/PTO/Support/AsyncSessionABI.h`` and is reached through +``AsyncWorkspaceShim.cpp``, so this file cannot drift from the ABI: it asks the +shim how large a session template is and hands it named values to fill one. +""" + +from __future__ import annotations + +import ctypes +import os +from pathlib import Path + +import numpy as np + +_DEFAULT_LIB_RELPATH = Path("build") / "libpto_async_shim.so" +_LIB_PATH_ENV = "PTO_ASYNC_SHIM" + +# The A2/A3 doorbell stages one 32-bit word in UB before an MTE store carries it +# out. Only the address reaches the expansion today, but reserving a whole +# 32-byte block keeps the staging clear of anything a kernel puts next to it. +DOORBELL_STAGE_BYTES = 32 + + +class AsyncWorkspaceError(RuntimeError): + """A shim call failed, or the shim library could not be loaded.""" + + +def _shim_library_path(explicit: str | os.PathLike[str] | None = None) -> Path: + """Resolve the shim library path from an argument, the environment, or the default.""" + raw = explicit if explicit is not None else os.environ.get(_LIB_PATH_ENV) + if raw: + candidate = Path(raw).expanduser().resolve() + source = f"{_LIB_PATH_ENV}={raw}" if explicit is None else f"library_path={raw}" + else: + candidate = (Path(__file__).resolve().parent / _DEFAULT_LIB_RELPATH).resolve() + source = "default location" + if not candidate.is_file(): + raise AsyncWorkspaceError( + f"AsyncWorkspace shim not found at {candidate} ({source}); " + "build it with test/comm/build_async_shim.sh" + ) + return candidate + + +def _bind_signatures(lib: ctypes.CDLL) -> None: + """Declare every shim entry point, so ctypes never guesses a width.""" + lib.pto_async_ws_init.argtypes = [ctypes.c_uint] + lib.pto_async_ws_init.restype = ctypes.c_int + + lib.pto_async_ws_finalize.argtypes = [] + lib.pto_async_ws_finalize.restype = None + + lib.pto_async_ws_error.argtypes = [] + lib.pto_async_ws_error.restype = ctypes.c_char_p + + lib.pto_async_ws_context_gm.argtypes = [] + lib.pto_async_ws_context_gm.restype = ctypes.c_uint64 + + lib.pto_async_ws_channel_count.argtypes = [] + lib.pto_async_ws_channel_count.restype = ctypes.c_uint + + lib.pto_async_ws_reported_queue_num.argtypes = [] + lib.pto_async_ws_reported_queue_num.restype = ctypes.c_uint32 + + lib.pto_async_ws_sq_tail.argtypes = [ctypes.c_uint, ctypes.c_uint] + lib.pto_async_ws_sq_tail.restype = ctypes.c_longlong + + for getter in ( + lib.pto_async_session_bytes, + lib.pto_async_session_qos_default, + lib.pto_async_ws_min_transfer_bytes, + lib.pto_async_ws_max_channels, + ): + getter.argtypes = [] + getter.restype = ctypes.c_uint32 + + lib.pto_async_session_pack.argtypes = [ + ctypes.c_void_p, # out + ctypes.c_uint32, # out_bytes + ctypes.c_uint32, # channel_idx + ctypes.c_uint32, # channel_num + ctypes.c_uint64, # block_bytes + ctypes.c_uint64, # tmp_buf_addr + ctypes.c_uint32, # tmp_buf_size + ctypes.c_uint32, # sync_id + ctypes.c_uint64, # comm_block_offset + ctypes.c_uint32, # qos + ctypes.c_uint32, # dest_rank_id + ] + lib.pto_async_session_pack.restype = ctypes.c_int + + +class AsyncWorkspace: + """One async workspace on the device the caller has already selected. + + The device must be selected before ``init``. Under the PTODSL ST harness + that is torch_npu's doing, and whether a workspace can attach to a context + someone else created is the assumption these cases exist to check. + + The shim can arrive two ways. Pass ``library`` to use one already loaded, + which is how a kernel that listed the shim in + ``@pto.jit(native_options={"host_sources": ...})`` reaches it: the shim is + inside the kernel's own library, so there is one build and one artifact. + Pass ``library_path``, or neither, to load a standalone build from + ``build_async_shim.sh`` instead, which needs no kernel and so is also how the + shim gets exercised on a machine with no card. + """ + + def __init__( + self, + library_path: str | os.PathLike[str] | None = None, + *, + library: ctypes.CDLL | None = None, + ) -> None: + if library is not None and library_path is not None: + raise AsyncWorkspaceError("pass either library or library_path, not both") + if library is not None: + self._library_path: Path | None = None + self._lib: ctypes.CDLL = library + else: + self._library_path = _shim_library_path(library_path) + self._lib = ctypes.CDLL(str(self._library_path)) + self._channels: int = 0 + self._inited: bool = False + _bind_signatures(self._lib) + + # --- lifecycle ----------------------------------------------------------- + + def init(self, channels: int = 1) -> None: + if channels <= 0: + raise AsyncWorkspaceError(f"channel count must be positive, got {channels}") + limit = self.max_channels + if channels > limit: + raise AsyncWorkspaceError(f"channel count {channels} exceeds the shim limit {limit}") + if self._lib.pto_async_ws_init(ctypes.c_uint(channels)) != 0: + raise AsyncWorkspaceError(f"AsyncWorkspace init failed: {self._error()}") + self._channels = channels + self._inited = True + + def finalize(self) -> None: + if not self._inited: + return + self._lib.pto_async_ws_finalize() + self._channels = 0 + self._inited = False + + def _error(self) -> str: + raw = self._lib.pto_async_ws_error() + if raw is None: + return "no error reported" + return raw.decode("utf-8", errors="replace") + + def _require_inited(self) -> None: + if not self._inited: + raise AsyncWorkspaceError("workspace is not initialized; call init() first") + + # --- properties ---------------------------------------------------------- + + @property + def library_path(self) -> Path | None: + """Where the shim was loaded from, or None when it came in already loaded.""" + return self._library_path + + @property + def context_gm(self) -> int: + """SessionField::ContextGm, opaque here and only ever passed back to the shim.""" + self._require_inited() + return int(self._lib.pto_async_ws_context_gm()) + + @property + def channel_count(self) -> int: + self._require_inited() + return int(self._lib.pto_async_ws_channel_count()) + + @property + def reported_queue_num(self) -> int: + """Channels the STARS query reported populating, which bounds valid channel indices.""" + self._require_inited() + return int(self._lib.pto_async_ws_reported_queue_num()) + + @property + def session_bytes(self) -> int: + return int(self._lib.pto_async_session_bytes()) + + @property + def qos_default(self) -> int: + return int(self._lib.pto_async_session_qos_default()) + + @property + def min_transfer_bytes(self) -> int: + """Shortest transfer the engine accepts, so a block size below it is invalid.""" + return int(self._lib.pto_async_ws_min_transfer_bytes()) + + @property + def max_channels(self) -> int: + return int(self._lib.pto_async_ws_max_channels()) + + # --- observation --------------------------------------------------------- + + def sq_tail(self, channel_idx: int = 0, channel_num: int = 1) -> int: + """The live queue tail for one channel. + + Reading this before and after a launch is what separates "moved the + bytes" from "moved them the way the split implies": one oversized entry + transfers the same bytes as a correct split, and only the tail differs. + """ + self._require_inited() + tail = int(self._lib.pto_async_ws_sq_tail(ctypes.c_uint(channel_idx), ctypes.c_uint(channel_num))) + if tail < 0: + raise AsyncWorkspaceError(f"reading the queue tail failed: {self._error()}") + return tail + + # --- session template ---------------------------------------------------- + + def pack_session( + self, + *, + block_bytes: int, + channel_idx: int = 0, + channel_num: int = 1, + tmp_buf_addr: int = 0, + tmp_buf_size: int = DOORBELL_STAGE_BYTES, + sync_id: int = 0, + comm_block_offset: int = 0, + qos: int | None = None, + dest_rank_id: int = 0, + ) -> np.ndarray: + """Return the session image the kernel loads with ``pto.session_init``. + + The shim fills it, so the field order and slot width never appear here. + ``ContextGm`` is taken from this workspace and is not a parameter. + """ + self._require_inited() + size = self.session_bytes + image = np.zeros(size, dtype=np.uint8) + rc = self._lib.pto_async_session_pack( + image.ctypes.data_as(ctypes.c_void_p), + ctypes.c_uint32(size), + ctypes.c_uint32(channel_idx), + ctypes.c_uint32(channel_num), + ctypes.c_uint64(block_bytes), + ctypes.c_uint64(tmp_buf_addr), + ctypes.c_uint32(tmp_buf_size), + ctypes.c_uint32(sync_id), + ctypes.c_uint64(comm_block_offset), + ctypes.c_uint32(self.qos_default if qos is None else qos), + ctypes.c_uint32(dest_rank_id), + ) + if rc != 0: + raise AsyncWorkspaceError(f"packing the session template failed: {self._error()}") + return image + + +__all__ = [ + "AsyncWorkspace", + "AsyncWorkspaceError", + "DOORBELL_STAGE_BYTES", +] diff --git a/test/comm/build_async_shim.sh b/test/comm/build_async_shim.sh new file mode 100755 index 0000000000..5156506f44 --- /dev/null +++ b/test/comm/build_async_shim.sh @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# +# Build the host shim over AsyncWorkspace as a shared library that the PTODSL ST +# harness loads with ctypes. +# +# No CANN headers or libraries are needed to build this: AsyncWorkspace resolves +# every toolkit symbol with dlopen, so the only link dependency is -ldl. The +# toolkit is required to run the result, not to produce it, which means this +# builds on a machine with no driver and no card. +# +# test/comm/build_async_shim.sh [output_dir] +# +# Environment: +# CXX host C++ compiler; defaults to g++ + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)" + +OUT_DIR="${1:-${SCRIPT_DIR}/build}" +CXX_BIN="${CXX:-g++}" +LIB_NAME="libpto_async_shim.so" + +if ! command -v -- "${CXX_BIN}" >/dev/null 2>&1; then + echo "error: host C++ compiler not found: ${CXX_BIN} (override with CXX=)" >&2 + exit 1 +fi + +mkdir -p "${OUT_DIR}" + +# AsyncSessionABI.h comes from include/, AsyncWorkspace.h from this directory. +cxx_args=( + -std=c++17 + -O2 + -Wall + -Wextra + -fPIC + -shared + "-I${REPO_ROOT}/include" + "-I${SCRIPT_DIR}" + "${SCRIPT_DIR}/AsyncWorkspaceShim.cpp" + -o "${OUT_DIR}/${LIB_NAME}" + -Wl,-z,relro + -Wl,-z,now + -Wl,-z,noexecstack + -ldl +) + +"${CXX_BIN}" "${cxx_args[@]}" + +echo "built ${OUT_DIR}/${LIB_NAME}" diff --git a/test/comm/spike_sdma_gm_gm_engine.py b/test/comm/spike_sdma_gm_gm_engine.py new file mode 100644 index 0000000000..411806631b --- /dev/null +++ b/test/comm/spike_sdma_gm_gm_engine.py @@ -0,0 +1,374 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# ---------------------------------------------------------------------------- +# spike: sdma_gm_gm, engine path +# target_ops: pto.session_init, pto.sdma_gm_gm +# scenarios: SDMA engine post, SQE split, ragged tail, session-driven offset +# ---------------------------------------------------------------------------- +# +# The engine form of pto.sdma_gm_gm: the kernel writes SQEs, publishes the queue +# tail and rings a doorbell, and the SDMA engine moves the bytes afterwards. Its +# {soft_put} sibling in sdma_gm_gm.py is a synchronous copy and needs none of +# that, which is why the two are separate files rather than separate cases. +# +# This is a spike, and what it is really checking is not the transfer. The IR is +# pinned in test/lit/vpto and the data path is covered by the {soft_put} cases; +# what is unproven is whether the engine path can run under the PTODSL ST harness +# at all. Three assumptions have to hold together, and none is checkable offline: +# +# 1. AsyncWorkspace can attach to the device and context torch_npu already set +# up, instead of needing to own aclrtSetDevice itself. +# 2. Memory the workspace allocates through the runtime is addressable by a +# kernel this harness launched, so a raw device address can be handed over +# inside a session template. +# 3. A destination the engine fills can be observed from the harness, which +# synchronizes the kernel and knows nothing about the engine. +# +# If all three hold, the engine path needs no framework change and the C++ and +# shell harness these cases used to need has nothing left to do. If one fails, +# this file is where it shows up, and the failure says which. +# +# The host side comes along in the kernel's own shared library. AsyncWorkspaceShim +# is listed in native_options below, so PTODSL compiles and links it beside the +# generated launch code, and this file reaches it through +# compiled.native_library(). One build, one artifact, and the shim is part of the +# build's identity: editing it rebuilds rather than silently reusing. +# +# Requires CANN >= 9.0.0 for the AICPU STARS query and a real device. There is no +# skip path: a spike that quietly passes when it could not run would defeat its +# own purpose. +# +# Which is why it lives here and not under test/vpto/cases/async-comm/ next to its +# {soft_put} sibling. run_host_vpto_validation.sh discovers every .py in that tree +# and runs it, and a case needing a 9.0 toolkit would fail the whole validation +# run on every machine without one. Promoting it once the three assumptions above +# are confirmed is a rename, and that should be a deliberate step rather than +# something a spike does on its way in. Run it directly, naming the generation of +# the card it is pointed at: +# +# PTO_ASYNC_ARCH=a5 python3 test/comm/spike_sdma_gm_gm_engine.py +# PTO_ASYNC_ARCH=a3 python3 test/comm/spike_sdma_gm_gm_engine.py +# +# Do not add ``from __future__ import annotations``. ``@pto.jit`` reads the +# entry annotations at decoration time and needs the live ``pto.ptr`` objects, +# not their string forms. + +import os +from pathlib import Path +import sys +import time + +import numpy as np + + +def _bootstrap_paths() -> None: + """Put test/dsl-st and test/comm on sys.path, wherever this file is checked out.""" + here = Path(__file__).resolve() + for candidate in here.parents: + common_dir = candidate / "test" / "dsl-st" + comm_dir = candidate / "test" / "comm" + if (common_dir / "common.py").exists(): + sys.path.insert(0, str(common_dir)) + sys.path.insert(0, str(comm_dir)) + return + raise RuntimeError("Unable to locate test/dsl-st/common.py from spike_sdma_gm_gm_engine.py") + + +_bootstrap_paths() + +from async_workspace import DOORBELL_STAGE_BYTES, AsyncWorkspace +from common import auto_main +from ptodsl import pto + + +SEED = 31 + +# The doorbell is the one part of the post that differs by generation, so the +# arch is a knob rather than a constant. A2/A3 reaches it only by MTE and stages +# the tail in UB; A5 writes it with st_dev, at a different offset. Everything +# else is shared, and the compiler draws no distinction between a2 and a3. +# +# Required rather than defaulted. This is the only thing here that decides which +# card the result may be run on, and getting it wrong is not a benign mistake: on +# A2/A3 a store aimed at sq_reg_base the wrong way faults the vector unit hard +# enough to leave the card in an unrecoverable RAS state. A card that has to be +# reset because a variable was unset is worth one line of typing. +_SUPPORTED_ARCHES = ("a2", "a3", "a5") +ARCH = os.environ.get("PTO_ASYNC_ARCH") +if ARCH is None: + raise RuntimeError( + "PTO_ASYNC_ARCH must name the generation of the card this will run on, one of " + f"{_SUPPORTED_ARCHES}. There is no default: the doorbell sequence differs " + "between A2/A3 and A5, and posting the wrong one at a card can leave it in " + "an unrecoverable state." + ) +if ARCH not in _SUPPORTED_ARCHES: + raise RuntimeError(f"PTO_ASYNC_ARCH must be one of {_SUPPORTED_ARCHES}, got {ARCH!r}") + +# How long to wait for the engine to drain. The harness synchronizes the kernel, +# which returns once the doorbell is rung, so the destination is still being +# written when the check starts. +POLL_TIMEOUT_MS = int(os.environ.get("PTO_ASYNC_POLL_TIMEOUT_MS", "2000")) +POLL_INTERVAL_S = 0.001 + +# One channel, index zero. A group wide enough to spread a post across is a +# scheduling policy the expansion does not implement yet, and the tail read has +# to name the same geometry the session did. +CHANNEL_IDX = 0 +CHANNEL_NUM = 1 + +# UB base for the A2/A3 doorbell staging. UB is private to a core and these +# kernels touch nothing else in it, so the bottom of the buffer is free. +UB_STAGE_ADDR = 0 + +# Pipe event id for that staging, chosen so it cannot collide: these kernels +# raise no other MTE3 event. +DOORBELL_SYNC_ID = 0 + +_SESSION_TYPE = "!pto.struct" + +# The destination is the last pointer so the harness allocates it as the output +# buffer. Copy direction is still dst <- src. +ENGINE_SOURCE = f"""module attributes {{pto.target_arch = "{ARCH}", pto.kernel_kind = #pto.kernel_kind}} {{ + func.func @sdma_gm_gm_engine( + %src: !pto.ptr, %sess_gm: !pto.ptr, + %dst: !pto.ptr, %nbytes: i64) attributes {{pto.kernel}} {{ + %sess = pto.declare_struct -> {_SESSION_TYPE} + pto.session_init %sess, %sess_gm + : {_SESSION_TYPE}, !pto.ptr + pto.sdma_gm_gm %dst, %src, %nbytes session(%sess) + : !pto.ptr, !pto.ptr, i64, {_SESSION_TYPE} + return + }} +}} +""" + + +@pto.jit( + name="sdma_gm_gm_engine", + target=ARCH, + backend="vpto", + mode="explicit", + source=ENGINE_SOURCE, + native_options={ + # Paths resolve against this file. AsyncWorkspace.h sits beside the shim; + # AsyncSessionABI.h, which owns every field position the shim uses, comes + # from the repository include tree. + "host_sources": ["AsyncWorkspaceShim.cpp"], + "include_dirs": [".", "../../include"], + # AsyncWorkspace resolves the CANN runtime with dlopen rather than linking + # it, so this is the shim's only link dependency. + "link_libraries": ["dl"], + }, +) +def sdma_gm_gm_engine( + src: pto.ptr(pto.i8, "gm"), + sess_gm: pto.ptr(pto.i8, "gm"), + dst: pto.ptr(pto.i8, "gm"), + nbytes: pto.i64, +): + pass + + +_workspace: AsyncWorkspace | None = None + + +def workspace() -> AsyncWorkspace: + """The process-wide workspace, built on first use. + + Built lazily rather than at import: --list and --emit-mlir run before the + harness selects a device, and neither should need a card. + + Compiling here is what makes the shim reachable, since it lives in the + library this call produces. The harness compiles the same kernel again before + it launches; that hits the specialization cache and returns this same handle, + so the library is built once. + """ + global _workspace + if _workspace is None: + compiled = sdma_gm_gm_engine.compile() + ws = AsyncWorkspace(library=compiled.native_library()) + ws.init(channels=CHANNEL_NUM) + reported = ws.reported_queue_num + if reported <= CHANNEL_IDX: + raise RuntimeError( + f"the STARS query reported {reported} populated channels, " + f"which does not cover channel {CHANNEL_IDX}" + ) + print( + f"async workspace ready: arch={ARCH} " + f"library={compiled.native_library_path()} " + f"channels={ws.channel_count} reported={reported}" + ) + _workspace = ws + return _workspace + + +def make_source(nbytes: int) -> np.ndarray: + rng = np.random.default_rng(SEED) + return rng.integers(0, 256, size=nbytes, dtype=np.uint8) + + +def _tail_delta(before: int, after: int) -> int: + """SQEs posted between two tail reads. + + The expansion masks the tail into the ring on every entry, so a delta is only + a count while the ring did not wrap. These cases post single digits into a + ring thousands of entries deep, so a wrap here means something else is wrong + and is worth reporting rather than folding away. + """ + if after < before: + raise AssertionError( + f"the queue tail moved backwards, from {before} to {after}; " + "the ring wrapped or another writer is sharing the channel" + ) + return after - before + + +def _poll_for_match(device_output, golden: np.ndarray) -> tuple[np.ndarray, float]: + """Read the destination back until it matches the golden or the deadline passes.""" + started = time.monotonic() + deadline = started + POLL_TIMEOUT_MS / 1000.0 + actual = device_output.cpu().numpy() + while not np.array_equal(actual, golden): + if time.monotonic() >= deadline: + break + time.sleep(POLL_INTERVAL_S) + actual = device_output.cpu().numpy() + return actual, time.monotonic() - started + + +def engine_case( + name: str, + *, + buffer_bytes: int, + nbytes: int, + block_bytes: int, + comm_block_offset: int = 0, +): + """One engine transfer, described by the session the host writes for it. + + ``block_bytes`` is what splits the transfer into SQEs, so it is the knob that + decides how many entries the queue should gain. ``comm_block_offset`` shifts + both endpoints, so the moved window keeps its position in the buffer. + """ + if comm_block_offset + nbytes > buffer_bytes: + raise RuntimeError( + f"case {name!r} would move past the end of the buffer: " + f"{comm_block_offset} + {nbytes} > {buffer_bytes}" + ) + want_entries = -(-nbytes // block_bytes) + tail_chunk = nbytes - (want_entries - 1) * block_bytes + + def make_case(): + ws = workspace() + floor = ws.min_transfer_bytes + # Both the full block and the short last one become an engine transfer, + # so either being under the floor is a broken case rather than a finding. + if block_bytes < floor or tail_chunk < floor: + raise RuntimeError( + f"case {name!r} splits into chunks of {block_bytes} and {tail_chunk} bytes, " + f"below the {floor}-byte engine minimum" + ) + + src = make_source(buffer_bytes) + session = ws.pack_session( + block_bytes=block_bytes, + channel_idx=CHANNEL_IDX, + channel_num=CHANNEL_NUM, + tmp_buf_addr=UB_STAGE_ADDR, + tmp_buf_size=DOORBELL_STAGE_BYTES, + sync_id=DOORBELL_SYNC_ID, + comm_block_offset=comm_block_offset, + ) + + # Everything outside the moved window stays zero. A session read that + # dropped CommBlockOffset would copy from the base and shift the window, + # which shows up here rather than as a byte count that still adds up. + golden = np.zeros(buffer_bytes, dtype=np.uint8) + window = slice(comm_block_offset, comm_block_offset + nbytes) + golden[window] = src[window] + + out = np.zeros(buffer_bytes, dtype=np.uint8) + expected = { + "golden": golden, + "tail_before": ws.sq_tail(CHANNEL_IDX, CHANNEL_NUM), + } + return [src, session, out], expected, [nbytes] + + def check_case(device_inputs, expected): + actual, waited_s = _poll_for_match(device_inputs[-1], expected["golden"]) + entries = _tail_delta( + expected["tail_before"], workspace().sq_tail(CHANNEL_IDX, CHANNEL_NUM) + ) + matched = np.array_equal(actual, expected["golden"]) + + # Report the post before the data. A correct entry count with an unfilled + # destination means the kernel did its part and the engine did not, which + # is a different failure from an entry count that never matched. + if entries != want_entries: + raise AssertionError( + f"{name}: expected {want_entries} SQEs for {nbytes} bytes in " + f"{block_bytes}-byte blocks, the queue tail gained {entries}" + + ("" if matched else "; the destination did not match either") + ) + if not matched: + differing = int(np.count_nonzero(actual != expected["golden"])) + first = int(np.flatnonzero(actual != expected["golden"])[0]) + raise AssertionError( + f"{name}: {want_entries} SQEs were posted, but after {waited_s * 1000:.0f}ms " + f"{differing} of {actual.size} bytes still differ, first at byte {first}" + ) + + return { + "name": name, + "kernel": sdma_gm_gm_engine, + "make_case": make_case, + "check": check_case, + } + + +CASES = [ + # One SQE. The shortest post that exercises the record read, the SQE write, + # the tail publish and the doorbell together. + engine_case( + "sdma_gm_gm_engine_single_entry", + buffer_bytes=4096, + nbytes=4096, + block_bytes=4096, + ), + # An even split, so the loop runs more than once and every entry is full. + engine_case( + "sdma_gm_gm_engine_even_split", + buffer_bytes=4096, + nbytes=4096, + block_bytes=1024, + ), + # A block size that does not divide the transfer: two full entries and a + # short one. The tail entry has to carry what is left rather than a full + # block, which would run past the end of the buffer. + engine_case( + "sdma_gm_gm_engine_ragged_tail", + buffer_bytes=4096, + nbytes=4096, + block_bytes=1536, + ), + # The session, not the kernel, decides where the transfer lands. + engine_case( + "sdma_gm_gm_engine_comm_block_offset", + buffer_bytes=8192, + nbytes=2048, + block_bytes=2048, + comm_block_offset=2048, + ), +] + + +auto_main(globals()) diff --git a/test/dsl-st/README.md b/test/dsl-st/README.md index 442ba361d0..63e11f998c 100644 --- a/test/dsl-st/README.md +++ b/test/dsl-st/README.md @@ -130,6 +130,10 @@ auto_main(globals()) - 如果输出 dtype 需要和 golden 分开控制,可以显式指定 - `output_index` - 默认比较最后一个 tensor;如果不是最后一个,改这里 +- `grid` + - launch 的 block 数,默认 `1` + - 只有真的要覆盖多核行为(例如 kernel 里用 `pto.get_block_idx` 分片)才需要改 + - 此时 golden 要按“所有 block 合起来的结果”来写 - `rtol` / `atol` - 浮点结果建议显式写;位级结果一般用 `0.0` @@ -149,6 +153,8 @@ auto_main(globals()) - `make_case()` - `check(device_inputs, expected)` +自定义 case dict 里同样可以带 `"grid"`;不写就是 `1`。 + ## 当前参考用例 可以直接参考: diff --git a/test/dsl-st/common.py b/test/dsl-st/common.py index 422d8a9830..b4d0a96ff9 100644 --- a/test/dsl-st/common.py +++ b/test/dsl-st/common.py @@ -67,6 +67,7 @@ def golden_output_case( output_dtype=None, output_index: int = -1, launch_args=None, + grid: int = 1, rtol: float = 1e-5, atol: float = 1e-5, ): @@ -99,6 +100,7 @@ def check_case(device_inputs, golden): return { "name": name, "kernel": kernel, + "grid": grid, "make_case": make_case, "check": check_case, } @@ -190,6 +192,12 @@ def run_cases(cases: list[dict], *, emit_mlir_fn=None, argv=None) -> int: f"DSL ST case {name!r} make_case() must return 2 or 3 values, got {len(made_case)}" ) + grid = case.get("grid", 1) + if not isinstance(grid, int) or grid <= 0: + raise RuntimeError( + f"DSL ST case {name!r} grid must be a positive integer, got {grid!r}" + ) + device_inputs = [torch.from_numpy(array).to(_DEVICE) for array in inputs] stream = npu_stream(torch) @@ -198,12 +206,12 @@ def run_cases(cases: list[dict], *, emit_mlir_fn=None, argv=None) -> int: compile_s = time.perf_counter() - t0 t0 = time.perf_counter() - compiled[1, stream](*device_inputs, *launch_args) + compiled[grid, stream](*device_inputs, *launch_args) torch.npu.synchronize() launch_s = time.perf_counter() - t0 case["check"](device_inputs, expected) - print(f"PASS {name} compile={compile_s:.3f}s launch={launch_s:.3f}s") + print(f"PASS {name} grid={grid} compile={compile_s:.3f}s launch={launch_s:.3f}s") print("All cases passed.") return 0 diff --git a/test/lit/vpto/aicore_ld_st_dev_invalid.pto b/test/lit/vpto/aicore_ld_st_dev_invalid.pto index fcd403473a..5073cd1d72 100644 --- a/test/lit/vpto/aicore_ld_st_dev_invalid.pto +++ b/test/lit/vpto/aicore_ld_st_dev_invalid.pto @@ -13,7 +13,7 @@ // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/helper.pto -o %t/helper.o 2>&1 | FileCheck %s --check-prefix=HELPER // RUN: not ptoas --pto-arch=a5 --pto-backend=vpto %t/policy.pto -o %t/policy.o 2>&1 | FileCheck %s --check-prefix=POLICY // RUN: not ptoas --pto-arch=a5 --cann-output-version=9.0.0-beta.1 --pto-backend=vpto --emit-vpto-llvm-ir %t/target.pto -o - 2>&1 | FileCheck %s --check-prefix=BETA -// RUN: not ptoas --pto-arch=a3 --cann-output-version=9.0.0 --pto-backend=vpto --emit-vpto-llvm-ir %t/target.pto -o - 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: not ptoas --pto-arch=a3 --pto-backend=vpto --emit-vpto-llvm-ir %t/a3_st_dev.pto -o - 2>&1 | FileCheck %s --check-prefix=A3 // SIMT: 'pto.ld_dev' op must be outside pto.simt_entry functions and pto.section.simt // UB: 'pto.st_dev' op requires GM pointer @@ -21,7 +21,7 @@ // HELPER: 'pto.ld_dev' op requires an enclosing ordinary AICore entry function // POLICY: 'pto.ld_dev' op does not accept l1cache or l2cache policy attributes // BETA: pto.ld_dev and pto.st_dev require CANN 9.0.0 or newer official lowering -// A3: pto.ld_dev and pto.st_dev require --pto-arch=a5 +// A3: pto.st_dev is not supported on A2/A3 //--- simt.pto module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { @@ -71,6 +71,16 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @a3_st_dev(%gm: !pto.ptr) attributes {pto.aicore} { + %c0 = arith.constant 0 : index + %value = arith.constant 1 : i32 + pto.st_dev %value, %gm[%c0] : !pto.ptr, i32 + return + } +} + //--- policy.pto module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { func.func @policy(%gm: !pto.ptr) attributes {pto.aicore} { diff --git a/test/lit/vpto/aicore_ld_st_dev_vpto_llvm.pto b/test/lit/vpto/aicore_ld_st_dev_vpto_llvm.pto index a93bc49712..ea1b355d3c 100644 --- a/test/lit/vpto/aicore_ld_st_dev_vpto_llvm.pto +++ b/test/lit/vpto/aicore_ld_st_dev_vpto_llvm.pto @@ -7,6 +7,9 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-backend=vpto --emit-vpto-llvm-ir %s -o - 2>&1 | FileCheck %s +// A2/A3 must not emit ST.DEV. SDMA doorbell coverage for that split is in +// async_sdma_gm_gm.pto; A2/A3 rejection of pto.st_dev is in +// aicore_ld_st_dev_invalid.pto. module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { func.func @aicore_ld_st_dev( diff --git a/test/lit/vpto/async_sdma_gm_gm.pto b/test/lit/vpto/async_sdma_gm_gm.pto new file mode 100644 index 0000000000..89c78dc974 --- /dev/null +++ b/test/lit/vpto/async_sdma_gm_gm.pto @@ -0,0 +1,154 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-before=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ROUNDTRIP +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=EXPAND +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SQE-A5 +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SQE-A2A3 --implicit-check-not=pto.st_dev --implicit-check-not="arith.constant 3840 : i32" + +// The session config struct is built in the kernel from the pointer the host +// passes in, because a stack-local struct cannot be a kernel argument. Only the +// immutable scalars live here; the queue tail stays where the engine keeps it +// so it survives across launches, and the channel record names its address. + +module attributes {"pto.target_arch" = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @sdma_gm_gm_probe(%dst: !pto.ptr, + %src: !pto.ptr, + %ws: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %nbytes = arith.constant 4096 : i64 + // Ordinary AICore entry functions read GM scalars with ld_dev; pto.ldg is + // restricted to SIMT scopes. + %ctx = pto.ld_dev %ws[%c0] : !pto.ptr -> i64 + + %z64 = arith.constant 0 : i64 + %z32 = arith.constant 0 : i32 + %one32 = arith.constant 1 : i32 + %qos = arith.constant 6 : i32 + %block = arith.constant 1048576 : i64 + + %sess = pto.declare_struct + -> !pto.struct + + pto.struct_set %sess[0], %ctx + : !pto.struct, i64 + pto.struct_set %sess[1], %z64 + : !pto.struct, i64 + pto.struct_set %sess[2], %z32 + : !pto.struct, i32 + pto.struct_set %sess[3], %z32 + : !pto.struct, i32 + pto.struct_set %sess[4], %z32 + : !pto.struct, i32 + pto.struct_set %sess[5], %one32 + : !pto.struct, i32 + pto.struct_set %sess[6], %block + : !pto.struct, i64 + pto.struct_set %sess[7], %z64 + : !pto.struct, i64 + pto.struct_set %sess[8], %z32 + : !pto.struct, i32 + pto.struct_set %sess[9], %z32 + : !pto.struct, i32 + pto.struct_set %sess[10], %z32 + : !pto.struct, i32 + pto.struct_set %sess[11], %one32 + : !pto.struct, i32 + + // 12 qos: the service class these posts run at, fixed for the session. + pto.struct_set %sess[12], %qos + : !pto.struct, i32 + + pto.sdma_gm_gm %dst, %src, %nbytes session(%sess) + : !pto.ptr, !pto.ptr, i64, + !pto.struct + + return + } +} + +// ROUNDTRIP-LABEL: func.func @sdma_gm_gm_probe( +// ROUNDTRIP: pto.declare_struct +// ROUNDTRIP: pto.sdma_gm_gm %{{.*}}, %{{.*}}, %{{.*}} session(%{{.*}}) + +// The kick disappears, replaced by channel-record reads, a split loop that +// writes one SQE per block, a release barrier, and the doorbell store. +// +// The SQE words and the queue tail go out as ordinary stores on both +// generations, matching the reference SDMA implementation. What makes them +// visible to the engine is the cache flush and barrier that follow, so the +// order of those three is what this pins down. This run targets A5, where the +// doorbell is an st_dev. +// EXPAND-LABEL: func.func @sdma_gm_gm_probe( +// EXPAND-NOT: pto.sdma_gm_gm +// EXPAND: pto.struct_get %{{.*}}[0] +// EXPAND: pto.ld_dev +// EXPAND: scf.for +// EXPAND: pto.store +// EXPAND: scf.yield +// EXPAND: pto.store +// EXPAND: pto.dcci +// EXPAND: pto.dsb +// EXPAND: pto.st_dev + +// The SQE layout is not shared across generations, so the packed word values +// and the length position pin down which branch ran. A5 sets wrCqe in word 0, +// carries credit 254, puts sssv..dns at bit 8, and writes the length at byte 48. +// +// Service class comes from the session on both generations, but lands in a +// different word: A5 gives it a word of its own at byte 20, shifted to bit 27, +// while A2/A3 ors it into word 4 at bit 13. Reading field 12 and shifting by +// the generation's amount is what tells the two apart, since neither value is a +// constant the pass can fold. These share the group above because constants are +// hoisted to the top of the function, which leaves their relative order up to +// the builder. +// SQE-A5-DAG: arith.constant 2059 : i32 +// SQE-A5-DAG: arith.constant 16646144 : i32 +// SQE-A5-DAG: arith.constant 3840 : i32 +// SQE-A5-DAG: arith.constant 48 : i64 +// SQE-A5-DAG: arith.constant 27 : i32 +// SQE-A5-DAG: arith.constant 20 : i64 +// SQE-A5-DAG: pto.struct_get %{{.*}}[12] +// SQE-A5-NOT: arith.constant 7680 : i32 + +// A2/A3 has no wrCqe bit, uses credit 240, shifts sssv..dns up one bit to make +// room for ie2, keeps the length at byte 28, and has to declare the post +// unlinked because byte 48 holds a link type there. Word 4 is 7680 without the +// service class, which is ored in on top of it. +// SQE-A2A3-DAG: arith.constant 11 : i32 +// SQE-A2A3-DAG: arith.constant 15728640 : i32 +// SQE-A2A3-DAG: arith.constant 7680 : i32 +// SQE-A2A3-DAG: arith.constant 28 : i64 +// SQE-A2A3-DAG: arith.constant 255 : i32 +// SQE-A2A3-DAG: arith.constant 13 : i32 +// SQE-A2A3-DAG: pto.struct_get %{{.*}}[12] + +// The doorbell is the other thing that splits by generation, and it is the one +// place the two do not merely differ in bit layout. sq_reg_base names a +// register, not memory: A5 reaches it with st_dev, while A2/A3 takes it only +// from MTE, so the tail is staged in the session's tmp_buf and moved out four +// bytes at a time. Neither an st_dev nor a scalar store to that address works +// on A2/A3, and the store faults the vector unit hard enough to leave the card +// in an unrecoverable RAS state. That st_dev never appears on A2/A3, and that +// A5's word 4 never leaks into it, is enforced for the whole output by the +// implicit-check-not on the RUN line above. +// +// byte_granular is what makes those four bytes survive: the default c220 store +// counts its length in 32-byte blocks, so a four-byte doorbell write rounds +// down to a transfer of nothing and the engine is never told the ring moved. +// +// The event guarding that staging is the session's, not a fixed one, so it uses +// the dyn form of the flag ops. A kernel that already drives MTE3 events can +// then keep them apart from this one. +// SQE-A2A3: pto.struct_get %{{.*}}[1] +// SQE-A2A3: pto.struct_get %{{.*}}[3] +// SQE-A2A3: pto.castptr %{{.*}} : i64 -> !pto.ptr +// SQE-A2A3: pto.store %{{.*}}, %{{.*}} : !pto.ptr, i32 +// SQE-A2A3: pto.set_flag_dyn[, +// SQE-A2A3: pto.wait_flag_dyn[, +// SQE-A2A3: pto.copy_ubuf_to_gm {{.*}}{vpto.byte_granular} diff --git a/test/lit/vpto/async_sdma_gm_gm_a5_soft_put.pto b/test/lit/vpto/async_sdma_gm_gm_a5_soft_put.pto new file mode 100644 index 0000000000..5584a7b157 --- /dev/null +++ b/test/lit/vpto/async_sdma_gm_gm_a5_soft_put.pto @@ -0,0 +1,139 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=A5 +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=A3 + +// A5 cannot PUT through the engine. {soft_put} turns the kick into a +// GM→UB→GM copy. A2/A3 ignores the attr and still posts SQEs. + +module attributes {"pto.target_arch" = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @sdma_gm_gm_soft_put(%dst: !pto.ptr, + %src: !pto.ptr, + %ws: !pto.ptr) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %nbytes = arith.constant 4096 : i64 + %ctx = pto.ld_dev %ws[%c0] : !pto.ptr -> i64 + + %z64 = arith.constant 0 : i64 + %z32 = arith.constant 0 : i32 + %one32 = arith.constant 1 : i32 + %qos = arith.constant 6 : i32 + %block = arith.constant 1048576 : i64 + + %sess = pto.declare_struct + -> !pto.struct + + pto.struct_set %sess[0], %ctx + : !pto.struct, i64 + pto.struct_set %sess[1], %z64 + : !pto.struct, i64 + pto.struct_set %sess[2], %z32 + : !pto.struct, i32 + pto.struct_set %sess[3], %z32 + : !pto.struct, i32 + pto.struct_set %sess[4], %z32 + : !pto.struct, i32 + pto.struct_set %sess[5], %one32 + : !pto.struct, i32 + pto.struct_set %sess[6], %block + : !pto.struct, i64 + pto.struct_set %sess[7], %z64 + : !pto.struct, i64 + pto.struct_set %sess[8], %z32 + : !pto.struct, i32 + pto.struct_set %sess[9], %z32 + : !pto.struct, i32 + pto.struct_set %sess[10], %z32 + : !pto.struct, i32 + pto.struct_set %sess[11], %one32 + : !pto.struct, i32 + pto.struct_set %sess[12], %qos + : !pto.struct, i32 + + pto.sdma_gm_gm %dst, %src, %nbytes session(%sess) {soft_put} + : !pto.ptr, !pto.ptr, i64, + !pto.struct + + return + } + + // Peer kernels use i8 endpoints and a runtime byte count. A same-type + // pto.castptr is illegal at emission, so this function is the shape that + // used to fail VPTO validation. + func.func @sdma_gm_gm_soft_put_i8(%dst: !pto.ptr, + %src: !pto.ptr, + %ws: !pto.ptr, + %nbytes: i64) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %ctx = pto.ld_dev %ws[%c0] : !pto.ptr -> i64 + + %z64 = arith.constant 0 : i64 + %z32 = arith.constant 0 : i32 + %one32 = arith.constant 1 : i32 + %qos = arith.constant 6 : i32 + %block = arith.constant 1048576 : i64 + + %sess = pto.declare_struct + -> !pto.struct + + pto.struct_set %sess[0], %ctx + : !pto.struct, i64 + pto.struct_set %sess[1], %z64 + : !pto.struct, i64 + pto.struct_set %sess[2], %z32 + : !pto.struct, i32 + pto.struct_set %sess[3], %z32 + : !pto.struct, i32 + pto.struct_set %sess[4], %z32 + : !pto.struct, i32 + pto.struct_set %sess[5], %one32 + : !pto.struct, i32 + pto.struct_set %sess[6], %block + : !pto.struct, i64 + pto.struct_set %sess[7], %z64 + : !pto.struct, i64 + pto.struct_set %sess[8], %z32 + : !pto.struct, i32 + pto.struct_set %sess[9], %z32 + : !pto.struct, i32 + pto.struct_set %sess[10], %z32 + : !pto.struct, i32 + pto.struct_set %sess[11], %one32 + : !pto.struct, i32 + pto.struct_set %sess[12], %qos + : !pto.struct, i32 + + pto.sdma_gm_gm %dst, %src, %nbytes session(%sess) {soft_put} + : !pto.ptr, !pto.ptr, i64, + !pto.struct + + return + } +} + +// A5-LABEL: func.func @sdma_gm_gm_soft_put( +// A5-NOT: pto.sdma_gm_gm +// A5: pto.copy_gm_to_ubuf +// A5: pto.set_flag_dyn[, +// A5: pto.wait_flag_dyn[, +// A5: pto.copy_ubuf_to_gm +// A5-NOT: pto.st_dev +// A5-LABEL: func.func @sdma_gm_gm_soft_put_i8( +// A5-NOT: pto.sdma_gm_gm +// A5: pto.copy_gm_to_ubuf +// A5: pto.copy_ubuf_to_gm + +// A3 still posts the engine path; the attr is A5-only. +// A3-LABEL: func.func @sdma_gm_gm_soft_put( +// A3-NOT: pto.sdma_gm_gm +// A3: pto.store +// A3: pto.copy_ubuf_to_gm +// A3-LABEL: func.func @sdma_gm_gm_soft_put_i8( +// A3-NOT: pto.sdma_gm_gm +// A3: pto.store diff --git a/test/lit/vpto/async_sdma_gm_gm_a5_soft_put_session.pto b/test/lit/vpto/async_sdma_gm_gm_a5_soft_put_session.pto new file mode 100644 index 0000000000..de569a9700 --- /dev/null +++ b/test/lit/vpto/async_sdma_gm_gm_a5_soft_put_session.pto @@ -0,0 +1,86 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s + +// The two kernel shapes the PTODSL ST case ships +// (test/vpto/cases/async-comm/sdma_gm_gm.py). That case only runs where a +// device or simulator is available, so the shapes are pinned here as well. +// +// Both take their session from a host template instead of spelling the fields +// out, which the existing A5 {soft_put} coverage does not: async_session_init +// pairs pto.session_init with the A2/A3 engine path, and +// async_sdma_gm_gm_a5_soft_put builds its session by hand. + +module attributes {"pto.target_arch" = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @soft_put_from_template(%src: !pto.ptr, + %sess_gm: !pto.ptr, + %dst: !pto.ptr, + %nbytes: i64) attributes {pto.kernel} { + %sess = pto.declare_struct + -> !pto.struct + + pto.session_init %sess, %sess_gm + : !pto.struct, + !pto.ptr + + pto.sdma_gm_gm %dst, %src, %nbytes session(%sess) {soft_put} + : !pto.ptr, !pto.ptr, i64, + !pto.struct + + return + } + + // Per-core slicing through the session rather than the pointers: every core + // loads the same template and then retunes CommBlockOffset, which the + // expansion applies to both endpoints. + func.func @soft_put_per_core_offset(%src: !pto.ptr, + %sess_gm: !pto.ptr, + %dst: !pto.ptr, + %chunk: i64) attributes {pto.kernel} { + %sess = pto.declare_struct + -> !pto.struct + + pto.session_init %sess, %sess_gm + : !pto.struct, + !pto.ptr + + %bid = pto.get_block_idx + %off = arith.muli %bid, %chunk : i64 + pto.struct_set %sess[7], %off + : !pto.struct, i64 + + pto.sdma_gm_gm %dst, %src, %chunk session(%sess) {soft_put} + : !pto.ptr, !pto.ptr, i64, + !pto.struct + + return + } +} + +// The template load survives as one ld_dev per field, and the transfer becomes +// the GM->UB->GM copy. No doorbell is written, because nothing was posted. +// CHECK-LABEL: func.func @soft_put_from_template( +// CHECK-NOT: pto.session_init +// CHECK: pto.ld_dev +// CHECK: pto.struct_set %{{.*}}[0] +// CHECK-NOT: pto.sdma_gm_gm +// CHECK: pto.copy_gm_to_ubuf +// CHECK: pto.copy_ubuf_to_gm +// CHECK-NOT: pto.st_dev + +// The retuned offset is what the expansion reads back, so the block query has +// to reach field 7 before the copy is formed. +// CHECK-LABEL: func.func @soft_put_per_core_offset( +// CHECK: pto.get_block_idx +// CHECK: pto.struct_set %{{.*}}[7] +// CHECK: pto.struct_get %{{.*}}[7] +// CHECK-NOT: pto.sdma_gm_gm +// CHECK: pto.copy_gm_to_ubuf +// CHECK: pto.copy_ubuf_to_gm +// CHECK-NOT: pto.st_dev diff --git a/test/lit/vpto/async_session_init.pto b/test/lit/vpto/async_session_init.pto new file mode 100644 index 0000000000..e719d185a2 --- /dev/null +++ b/test/lit/vpto/async_session_init.pto @@ -0,0 +1,57 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto --emit-vpto --mlir-print-ir-before=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ROUNDTRIP +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=EXPAND + +// A session cannot reach a kernel as an argument, so the host leaves its values +// in GM and the kernel copies them into a struct it declares itself. Without +// this op that copy is one load and one store per field, written out at every +// entry point. + +module attributes {"pto.target_arch" = "a3", pto.kernel_kind = #pto.kernel_kind} { + func.func @session_init_probe(%dst: !pto.ptr, + %src: !pto.ptr, + %tmpl: !pto.ptr) attributes {pto.kernel} { + %nbytes = arith.constant 4096 : i64 + + %sess = pto.declare_struct + -> !pto.struct + + pto.session_init %sess, %tmpl + : !pto.struct, + !pto.ptr + + // A field may still be retuned after the copy, which is what makes the + // template a default rather than the whole story. + %grp = arith.constant 2 : i32 + pto.struct_set %sess[4], %grp + : !pto.struct, i32 + + pto.sdma_gm_gm %dst, %src, %nbytes session(%sess) + : !pto.ptr, !pto.ptr, i64, + !pto.struct + + return + } +} + +// ROUNDTRIP: pto.session_init %{{.*}}, %{{.*}} : !pto.struct, !pto.ptr + +// The op is gone and every field has been copied. Slots are eight bytes wide +// regardless of field width, so the offsets run 0, 8, 16 ... and a narrow field +// is read at the base of its slot rather than packed against its neighbour. +// EXPAND-NOT: pto.session_init +// EXPAND-DAG: arith.constant 8 : i64 +// EXPAND-DAG: arith.constant 16 : i64 +// EXPAND-DAG: arith.constant 96 : i64 + +// Field 0 is 64-bit and field 2 is 32-bit; both are read with ld_dev because a +// comm kernel is an ordinary entry function, where ldg is not available. +// EXPAND: pto.ld_dev %{{.*}} : !pto.ptr -> i64 +// EXPAND: pto.struct_set %{{.*}}[0] diff --git a/test/vpto/cases/async-comm/sdma_gm_gm.py b/test/vpto/cases/async-comm/sdma_gm_gm.py new file mode 100644 index 0000000000..5f468d59f3 --- /dev/null +++ b/test/vpto/cases/async-comm/sdma_gm_gm.py @@ -0,0 +1,261 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# ----------------------------------------------------------------------------- +# case: async-comm/sdma_gm_gm +# target_ops: pto.session_init, pto.sdma_gm_gm +# scenarios: A5 {soft_put} GM->UB->GM, chunk tail, session-driven offset, multicore +# ----------------------------------------------------------------------------- +# +# Only the A5 {soft_put} form of pto.sdma_gm_gm is exercised here, and that is +# deliberate. That expansion is a synchronous GM->UB->GM copy, so the transfer +# is complete when the kernel returns and the result is an ordinary golden +# comparison. The engine form instead posts SQEs and rings a doorbell, which +# needs an async workspace from the AICPU STARS query and a destination that is +# polled rather than read once; neither fits a golden ST harness. The engine +# path and the two-device window shapes stay in test/lit/vpto/async_*.pto. +# +# What the four cases separate is not four kernels but four things the +# expansion has to get right. Three of them share one kernel and differ only in +# the session template the host writes, which is the point: the transfer is +# described by session data, not by the kernel. + +from pathlib import Path +import sys + +import numpy as np + + +def _bootstrap_dsl_st_common() -> None: + here = Path(__file__).resolve() + for candidate in here.parents: + common_dir = candidate / "test" / "dsl-st" + if (common_dir / "common.py").exists(): + sys.path.insert(0, str(common_dir)) + return + raise RuntimeError("Unable to locate test/dsl-st/common.py from sdma_gm_gm.py") + + +_bootstrap_dsl_st_common() + +from common import auto_main, golden_output_case +from ptodsl import pto + + +SEED = 29 + +# Staging chunk the A5 {soft_put} expansion uses per iteration. It is fixed in +# the expansion rather than read from SessionField::BlockBytes, so a transfer +# longer than this is what makes the copy loop run more than once. +SOFT_PUT_CHUNK_BYTES = 32768 + +# Field slots and their values match include/PTO/Support/AsyncSessionABI.h. +# The template is one 8-byte slot per field, so a uint64 array indexed by the +# SessionField value is the whole layout. +_SESSION_FIELDS = 13 +_TMP_BUF_ADDR = 1 +_TMP_BUF_SIZE = 2 +_SYNC_ID = 3 +_CHANNEL_IDX = 4 +_CHANNEL_NUM = 5 +_BLOCK_BYTES = 6 +_COMM_BLOCK_OFFSET = 7 +_ENGINE = 8 +_FLAGS = 11 +_QOS = 12 + +_ENGINE_SDMA = 0 +_FLAG_VALID = 1 +_QOS_DEFAULT = 6 + +# UB base for the staging buffer. Nothing else in these kernels touches UB, and +# UB is private to a core, so every block can stage at the same address. +_UB_TMP_ADDR = 0 + +_SESSION_TYPE = ( + "!pto.struct" +) + +# The destination is the last pointer so the harness can allocate it as the +# golden output. Copy direction is still dst <- src. +SOFT_PUT_SOURCE = f"""module attributes {{pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind}} {{ + func.func @sdma_gm_gm_soft_put( + %src: !pto.ptr, %sess_gm: !pto.ptr, + %dst: !pto.ptr, %nbytes: i64) attributes {{pto.kernel}} {{ + %sess = pto.declare_struct -> {_SESSION_TYPE} + pto.session_init %sess, %sess_gm + : {_SESSION_TYPE}, !pto.ptr + pto.sdma_gm_gm %dst, %src, %nbytes session(%sess) {{soft_put}} + : !pto.ptr, !pto.ptr, i64, {_SESSION_TYPE} + return + }} +}} +""" + +# Each core takes one slice by overwriting SessionField::CommBlockOffset after +# the template load. The offset is written into the session rather than folded +# into the pointers because that is the field the expansion applies to both +# endpoints, so a core that read the template but failed to diverge lands on +# slice zero and is visible as a hole in the destination. +MULTICORE_SOURCE = f"""module attributes {{pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind}} {{ + func.func @sdma_gm_gm_soft_put_multicore( + %src: !pto.ptr, %sess_gm: !pto.ptr, + %dst: !pto.ptr, %chunk: i64) attributes {{pto.kernel}} {{ + %sess = pto.declare_struct -> {_SESSION_TYPE} + pto.session_init %sess, %sess_gm + : {_SESSION_TYPE}, !pto.ptr + %bid = pto.get_block_idx + %off = arith.muli %bid, %chunk : i64 + pto.struct_set %sess[{_COMM_BLOCK_OFFSET}], %off : {_SESSION_TYPE}, i64 + pto.sdma_gm_gm %dst, %src, %chunk session(%sess) {{soft_put}} + : !pto.ptr, !pto.ptr, i64, {_SESSION_TYPE} + return + }} +}} +""" + + +@pto.jit( + name="sdma_gm_gm_soft_put", + target="a5", + backend="vpto", + mode="explicit", + source=SOFT_PUT_SOURCE, +) +def sdma_gm_gm_soft_put( + src: pto.ptr(pto.i8, "gm"), + sess_gm: pto.ptr(pto.i8, "gm"), + dst: pto.ptr(pto.i8, "gm"), + nbytes: pto.i64, +): + pass + + +@pto.jit( + name="sdma_gm_gm_soft_put_multicore", + target="a5", + backend="vpto", + mode="explicit", + source=MULTICORE_SOURCE, +) +def sdma_gm_gm_soft_put_multicore( + src: pto.ptr(pto.i8, "gm"), + sess_gm: pto.ptr(pto.i8, "gm"), + dst: pto.ptr(pto.i8, "gm"), + chunk: pto.i64, +): + pass + + +def session_template(*, comm_block_offset: int = 0, sync_id: int = 0): + """Host-written session template, as one uint64 slot per SessionField.""" + slots = np.zeros(_SESSION_FIELDS, dtype=np.uint64) + # SessionField::ContextGm stays zero: it names the channel record table, and + # {soft_put} posts nothing, so no record is ever read. The engine path is + # what needs a real workspace address there. + slots[_TMP_BUF_ADDR] = _UB_TMP_ADDR + slots[_TMP_BUF_SIZE] = SOFT_PUT_CHUNK_BYTES + slots[_SYNC_ID] = sync_id + slots[_CHANNEL_IDX] = 0 + slots[_CHANNEL_NUM] = 1 + slots[_BLOCK_BYTES] = SOFT_PUT_CHUNK_BYTES + slots[_COMM_BLOCK_OFFSET] = comm_block_offset + slots[_ENGINE] = _ENGINE_SDMA + slots[_FLAGS] = _FLAG_VALID + slots[_QOS] = _QOS_DEFAULT + return slots.view(np.uint8).copy() + + +def make_source(nbytes: int): + rng = np.random.default_rng(SEED) + return rng.integers(0, 256, size=nbytes, dtype=np.uint8) + + +def whole_buffer_case(name, *, buffer_bytes, grid=1, kernel=sdma_gm_gm_soft_put): + """A transfer that ends up covering the whole destination buffer.""" + + def make_inputs(): + return [make_source(buffer_bytes), session_template()] + + def make_expected(src, _sess): + return src.copy() + + return golden_output_case( + name, + kernel, + inputs=make_inputs, + expected=make_expected, + # The trailing scalar is nbytes for the single-core kernel and the + # per-core chunk for the multicore one; both are bytes moved per launch + # of one block. + launch_args=lambda src, _sess: [int(src.size) // grid], + grid=grid, + rtol=0.0, + atol=0.0, + ) + + +def offset_case(name, *, buffer_bytes, offset, nbytes): + """A transfer displaced by SessionField::CommBlockOffset on both endpoints.""" + + def make_inputs(): + return [ + make_source(buffer_bytes), + session_template(comm_block_offset=offset), + ] + + def make_expected(src, _sess): + # The offset applies to source and destination alike, so the moved + # window keeps its position and everything outside it stays zero. A + # session read that dropped the field would copy from the base instead + # and shift the window, which this catches. + expected = np.zeros_like(src) + expected[offset : offset + nbytes] = src[offset : offset + nbytes] + return expected + + return golden_output_case( + name, + sdma_gm_gm_soft_put, + inputs=make_inputs, + expected=make_expected, + launch_args=[nbytes], + rtol=0.0, + atol=0.0, + ) + + +CASES = [ + # One staging iteration: the shortest thing that proves the template load, + # the UB round trip, and the MTE2/MTE3 handshake fit together. + whole_buffer_case("sdma_gm_gm_soft_put_single_chunk", buffer_bytes=4096), + # Two full chunks plus a short one, so the loop trip count is above one and + # the tail iteration is clamped to what is left rather than a full chunk. + whole_buffer_case( + "sdma_gm_gm_soft_put_chunk_tail", + buffer_bytes=2 * SOFT_PUT_CHUNK_BYTES + 1024, + ), + # The session, not the kernel, decides where the transfer lands. + offset_case( + "sdma_gm_gm_soft_put_comm_block_offset", + buffer_bytes=8192, + offset=2048, + nbytes=4096, + ), + # Four cores, four disjoint slices, one shared template. The destination is + # only whole if every core resolved its own offset. + whole_buffer_case( + "sdma_gm_gm_soft_put_multicore", + buffer_bytes=16384, + grid=4, + kernel=sdma_gm_gm_soft_put_multicore, + ), +] + + +auto_main(globals())