diff --git a/docs/designs/ptoas-memory-consistency-design.md b/docs/designs/ptoas-memory-consistency-design.md index d09c2176c7..f5699a7532 100644 --- a/docs/designs/ptoas-memory-consistency-design.md +++ b/docs/designs/ptoas-memory-consistency-design.md @@ -1,405 +1,69 @@ # PTOAS 内存一致性设计 -本文说明 PTOAS 如何建模并校验 GM payload 与 signal 之间的内存一致性要求。 +本文说明 PTOAS 当前暴露的内存一致性 IR，以及为什么暂时移除自动 +MemoryConsistency 分析 pass。 -这里讨论的是内存一致性，不是自动同步。自动同步负责 pipe 之间的执行顺序，例如 -`set_flag`、`wait_flag` 和 `pipe_barrier`。内存一致性负责回答另一个问题：当 signal -已经被对端观察到时，signal 之前发布的 payload 是否已经对正确的观察方可见。 +## 当前状态 -## 1. 背景 - -`pto.comm.tnotify` 用来发布一个 signal。对端通过 `pto.comm.twait` 或 -`pto.comm.ttest` 观察这个 signal，然后读取对应的 payload。 - -signal ready 不等价于 payload 一定已经可见。原因是 signal 和 payload 可能走不同的硬件路径： - -- signal 通常是一个较小的通信同步标记。 -- payload 通常是更大的 GM 数据，可能由 MTE3、FIX 或 comm macro op 写出。 -- 不同路径之间只靠源码顺序不一定形成完整的 GM 可见性关系。 - -因此，PTOAS 需要在发布 signal 前校验 release 侧动作，在消费 signal 后校验 acquire 侧动作。 - -## 2. 当前暴露的 IR - -当前 PR 只暴露两条内存一致性 IR： +PTOAS 目前保留显式内存一致性 IR 和 EmitC lowering： | PTO IR | 语义 | EmitC lowering | | --- | --- | --- | -| `pto.fence.barrier_all #pto.fence_scope` | GM visibility barrier，可用于 publish 或 acquire 边界 | `dsb(DSB_DDR)` | -| `pto.cmo.cacheinvalid all #pto.address_space` | whole-cache GM cache maintenance 边界，可用于 release 或 acquire | `dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE)` | -| `pto.cmo.cacheinvalid %addr single_cache_line : !pto.ptr` 或 `!pto.partition_tensor_view<...>` | 指定 GM payload 地址。cacheable 路径上是真实 cache maintenance 边界；non-cacheable 路径上可以作为 marker-only IR 驱动精准 pipe drain | cacheable 路径 lower 到 `dcci(addr, cache_line_t::SINGLE_CACHE_LINE)`；marker-only 路径由 pass 消费并消除 | - -对应 PyPTO 写法： - -```python -pto.FenceBarrierAllOp(pto.FenceScope.GM) -pto.CmoCacheInvalidOp(pto.AddressSpace.GM) -pto.CmoCacheInvalidOp(pto.AddressSpace.GM, addr=payload_ptr) -``` - -`FenceScopeAttr` 当前定义三个语义 scope： - -- `local_memory` -- `gm` -- `all` - -本 PR 只为 `gm` 和 `all` 提供 EmitC lowering；它们当前都 lower 成 `dsb(DSB_DDR)`。 -`local_memory` 预留给后续 UB/local-memory 语义。 - -`cmo.cacheinvalid` 支持两种粒度： - -```mlir -pto.cmo.cacheinvalid all #pto.address_space -pto.cmo.cacheinvalid %payload_ptr single_cache_line : !pto.ptr -``` - -`all` 表示 whole-cache 粒度，并且一定会 lower 成真实 `dcci`。在 release 侧， -它也是一个保守 payload marker：PTOAS 会认为用户希望发布到这条 CMO 为止已经 pending 的 -所有 GM payload 访问，并为这些访问补齐必要的 pipe drain。`single_cache_line` 表示 -`%payload_ptr` 所在 cache line，同时也是 PTOAS 用来精准识别 TNotify payload 的地址 -marker。当前还没有 range 形式；如果 payload 横跨多条 cache line，PyPTO 或用户可以 -生成多条 single-line CMO，或者使用 whole-cache 形式做保守处理。 - -对 non-cacheable 的 MTE2、MTE3 或 FIX payload path，`single_cache_line` 主要用作精准地址 -marker。MemoryConsistency pass 会用它和前序 pending payload access 做 alias 匹配，决定 -是否需要插入 `PIPE_MTE2`、`PIPE_MTE3` 或 `PIPE_FIX` drain。若匹配到的路径不经过 scalar -cache，EmitC lowering 会直接消除这条 `cmo.cacheinvalid`，不生成 `dcci`。`all` 形式不做 -地址匹配，而是保守选择该 op 之前已经 pending 的所有 GM payload access，并保留真实 -whole-cache `dcci`。 - -## 3. DCCI 语义边界 - -`dcci` 是 cache maintenance 指令，不是 pipe drain，也不是 GM visibility fence。 - -PTOAS 当前对齐 PTO-ISA 的写法，只生成 two-argument CCE builtin： - -```cpp -dcci(addr, cache_line); -``` - -其中 `cache_line_t::SINGLE_CACHE_LINE` 表示处理传入地址所在 cache line， -`cache_line_t::ENTIRE_DATA_CACHE` 表示处理整个 Data Cache。 - -公开 AscendC API 和部分底层 CCE 头文件还存在带 destination 参数的形式，例如 -`CACHELINE_OUT`。但是 PTO-ISA 主线的 `TNotify`、`TWait`、`TTest` 和 ready-queue -实现都使用 two-argument `dcci`。因此本 PR 不把三参数 `dcci(..., CACHELINE_OUT)` -作为 PTOAS 的生成代码契约，也不在 PTO IR 中暴露 destination 选择。 - -当前 public IR 使用 two-argument CCE builtin。whole-cache 形式： - -```cpp -dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE); -``` - -single-cache-line 形式： - -```cpp -dcci((__gm__ void*)addr, cache_line_t::SINGLE_CACHE_LINE); -``` - -它们都是 PTOAS 暴露给上游的显式 CMO 边界。虽然 op 名称沿用 -`cacheinvalid`，但当前 PTO-ISA 对齐的 two-argument `dcci` 同时服务两类上下文： +| `pto.cmo.cacheinvalid all #pto.address_space` | whole-cache GM cache maintenance | `dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE)` | +| `pto.cmo.cacheinvalid %addr single_cache_line : !pto.ptr` | 指定地址所在 cache line 的 GM cache maintenance | `dcci((__gm__ void*)addr, cache_line_t::SINGLE_CACHE_LINE)` | +| `pto.fence.barrier_all #pto.fence_scope` | GM visibility fence with conservative local producer drain | `pipe_barrier(PIPE_MTE2); pipe_barrier(PIPE_MTE3); pipe_barrier(PIPE_FIX); dsb(DSB_DDR)` | -- release 侧：放在 cacheable scalar GM store 之后、`fence.barrier_all` 之前，作为发布 - payload 前的 cache maintenance。 -- acquire 侧：放在 `TWait` 或 `TTest` 之后、cacheable GM load 之前，避免读取本地 stale - cache line。 +`pto.cmo.cacheinvalid` 和 `pto.fence.barrier_all` 由 PyPTO 或用户显式插入。 +PTOAS 不再通过独立 `pto-memory-consistency` pass 自动扫描 signal/payload +关系、消除 marker-only CMO，或诊断缺失的 CMO/fence。为了避免复杂分析带来的 +编译时退化，EmitC lowering 会在显式 GM fence 处保守 drain MTE2、MTE3 和 FIX。 -本 PR 不新增 `pto.cmo.clean`。原因是当前 release 侧和 acquire 侧都对齐到 PTO-ISA -two-argument `dcci`，新增一个 lowering 相同的 public op 会让 PyPTO 和用户难以区分。 -后续如果确认需要暴露不同 destination 或精确 writeback 语义，再单独扩展 CMO IR。 - -MemoryConsistency pass 在当前阶段不会证明 single-cache-line CMO 的地址是否覆盖所有 -pending payload。也就是说，`single_cache_line` 是一个精确 CMO 操作，但其正确使用需要 -PyPTO 或用户保证 `%addr` 确实落在需要发布或消费的 payload cache line 上。 - -## 4. MemoryConsistency Pass - -`pto-memory-consistency` 是一个 Module pass，运行在 shared mainline 上，因此 EmitC 和 -VPTO backend 都会先经过这一步。 - -这个 pass 的职责是： - -- 识别 signal publish 前是否存在 `cmo.cacheinvalid %addr` 精准 marker，或 - `cmo.cacheinvalid all #pto.address_space` 保守全量 marker。 -- 识别 signal acquire 后是否存在 cacheable GM payload read。 -- 校验用户或 PyPTO 是否已经插入必要的 `fence.barrier_all` 或 `cmo.cacheinvalid`。 -- 在显式 `barrier_all` 前自动补齐 marker 对应的 MTE3 或 FIX pipe drain。 -- 在 `TNotify` 前自动补齐 marker 对应的 MTE2 pipe drain。 -- 对缺失或顺序错误的场景报编译错误。 - -这个 pass 不负责分配 event id，也不属于 InsertSync 自动同步流水线。 - -## 5. 场景规则 - -### 5.1 MTE3、FIX 或 TPUT 写 Payload 后发布 Signal - -适用场景： - -- `TStore` 通过 `PIPE_MTE3` 写 GM。 -- `TStore` 通过 `PIPE_FIX` 写 GM，例如 ACC tile 写回 GM。 -- `TStoreFP` 通过 `PIPE_FIX` 写 GM。 -- `TPUT` macro op 内部通过 MTE3 写 peer GM。 -- 其他 macro op phase 中存在 MTE3 GM write。 - -PyPTO 或用户需要生成： - -```mlir -// payload producer -pto.cmo.cacheinvalid %payload single_cache_line : !pto.partition_tensor_view<...> -pto.fence.barrier_all #pto.fence_scope -pto.comm.tnotify ... -``` - -PTOAS 会用 `cmo.cacheinvalid %payload` 的地址和前序 pending payload access 做 alias 匹配。 -也可以使用 `pto.cmo.cacheinvalid all #pto.address_space`，表示不做精确地址指定， -保守发布到该 op 为止已经 pending 的全部 GM payload access。如果 marker 选择到 pending MTE3 或 -FIX GM write，就在 -`pto.fence.barrier_all #pto.fence_scope` 前自动插入对应 pipe 的 drain： - -```mlir -pto.barrier #pto.pipe -// or -pto.barrier #pto.pipe -``` - -最终 EmitC 关键顺序是： - -```cpp -pipe_barrier(PIPE_MTE3); -dsb(DSB_DDR); -pto::comm::TNOTIFY(...); -``` - -`pipe_barrier` 用来排空实际执行 GM write 的 pipe。`barrier_all` lower 出来的 -`dsb(DSB_DDR)` 用来保证这些 GM 写入在 signal 发布前进入 GM visibility domain。 - -如果 matching payload 缺少 `pto.fence.barrier_all #pto.fence_scope`，PTOAS 会报错。 -PTOAS 可以推导低层 pipe drain，但 payload marker 和 visibility fence 仍由 PyPTO 或用户 -显式表达。 - -### 5.2 MTE2 工作后发布 Signal - -`TLoad` 或其他 `PIPE_MTE2` 工作出现在 `TNotify` 前时，PyPTO 应生成 matching payload -marker： - -```mlir -pto.tload ... -pto.cmo.cacheinvalid %payload single_cache_line : !pto.partition_tensor_view<...> -pto.comm.tnotify ... -``` +## 为什么移除自动分析 pass -PTOAS 会用 marker 地址匹配前序 MTE2 GM read，并在 EmitC lowering 中生成 -`PIPE_MTE2` barrier。MTE2 是 GM read 方向，只需要 signal 前不要越过前序 MTE2 工作， -不需要 `barrier_all`。这条 marker 对 MTE2 non-cacheable path 不生成 `dcci`。 +issue #950 暴露了 `pto-memory-consistency` 在复杂单 block 控制流中的编译时 +间退化。该 pass 的 state `merge()` 直接追加状态向量；处理单块 `scf.if` 时， +子 region 返回值已经包含入口状态，父层又把它合并回原状态，导致每经过一个 +`if`，历史状态近似翻倍。 -### 5.3 TWait 或 TTest 后读取 Cacheable GM Payload +issue #950 的 repro 中包含约 80 个连续 `scf.if`、160 次 GM scalar load、163 次 +GM scalar store，因此状态规模接近指数增长，触发 ptoas 0.50 相比 0.48 的确定性 +编译超时。 -适用场景： +在重新设计 bounded、去重的数据流模型之前，默认 pipeline 不应继续运行这个 +自动分析 pass。当前策略是先回到显式 IR 模式，保证编译时间可控。 -- `TWait` 返回后执行 `load_scalar` 读取 GM payload。 -- `TTest` 成功观察到 signal 后执行 `load_scalar` 读取 GM payload。 +## 使用要求 -需要的顺序： +PyPTO 或用户需要显式表达内存一致性动作： ```mlir -pto.comm.twait ... -pto.cmo.cacheinvalid %payload_ptr single_cache_line : !pto.ptr -%value = pto.load_scalar ... -``` - -`cacheinvalid` 用来避免读取到本地 stale cache line。也可以使用 whole-cache 形式： -`pto.cmo.cacheinvalid all #pto.address_space`。如果缺少该 op，PTOAS 会报错。 - -Acquire 侧的 `single_cache_line` 形式是用户或 PyPTO 对 cache line 覆盖关系的显式承诺： -`%payload_ptr` 必须覆盖后续 `load_scalar` 实际读取的 GM 地址。PTOAS 当前只检查 -`TWait` 或 ready `TTest` 后、cacheable GM `load_scalar` 前存在 `cmo.cacheinvalid`， -不会证明该 CMO 地址和后续 load 地址是否 alias。如果生成方无法保证精确地址正确， -应使用 `pto.cmo.cacheinvalid all #pto.address_space` 做保守 whole-cache invalidate。 - -### 5.4 Cacheable Scalar GM Store 后发布 Signal - -适用场景： - -```mlir -pto.store_scalar %value, %payload[%idx] : !pto.ptr, i32 +// cacheable scalar GM store 发布 payload +pto.store_scalar %value, %payload[%idx] : !pto.ptr, i32 pto.cmo.cacheinvalid %payload single_cache_line : !pto.ptr pto.fence.barrier_all #pto.fence_scope pto.comm.tnotify ... ``` -最终 EmitC 关键顺序是： - -```cpp -payload[idx] = value; -dcci((__gm__ void*)payload, cache_line_t::SINGLE_CACHE_LINE); -dsb(DSB_DDR); -pto::comm::TNOTIFY(...); -``` - -这里 `cmo.cacheinvalid` 表示 publish 前的显式 CMO 边界，`fence.barrier_all` 表示 -GM visibility fence。二者顺序不能交换。如果 payload 覆盖多条 cache line，可以改用 -whole-cache 形式，或者生成多条 single-line CMO。如果缺少 `cmo.cacheinvalid`，或者把它 -放在 `fence.barrier_all` 后面，PTOAS 会报错。 - -## 6. PyPTO 对接说明 - -PyPTO 不需要手动生成 `pto.barrier #pto.pipe` 或 -`pto.barrier #pto.pipe`。这是低层 pipe drain 细节，由 PTOAS 根据 -`pto.cmo.cacheinvalid %payload single_cache_line` marker 匹配到的 GM payload access 自动插入。 -这样可以保证最终顺序是对应 pipe barrier 先于 `dsb(DSB_DDR)`，不会出现先 fence、后 drain -的错误顺序。没有 marker 的 memory access 不会被 PTOAS 当作 TNotify payload；也就是说， -marker 是 signal 和 payload 的显式关联。`single_cache_line` 只关联匹配地址的 payload； -`all` 是显式兜底 marker，会关联该 op 之前已经 pending 的所有 GM payload access。 - -PyPTO 生成规则： - -| 场景 | PyPTO 需要生成 | PTOAS 自动补齐 | -| --- | --- | --- | -| `TStore`、`TStoreFP` 或 `TPUT` 后发布 signal | `pto.cmo.cacheinvalid %payload single_cache_line` 作为精准 payload marker，或 `pto.cmo.cacheinvalid all #pto.address_space` 作为保守全量 marker，随后 `pto.fence.barrier_all #pto.fence_scope` | 若 marker 选择到 pending GM write，则补 `PIPE_MTE3` 或 `PIPE_FIX` drain；non-cacheable single-line marker 不生成 `dcci`，whole-cache marker 生成 whole-cache `dcci` | -| `TLoad` 后发布 signal | `pto.cmo.cacheinvalid %payload single_cache_line` 作为精准 payload marker，或 `pto.cmo.cacheinvalid all #pto.address_space` 作为保守全量 marker；不需要显式 fence | 若 marker 选择到 pending GM read，则补 `PIPE_MTE2` drain；non-cacheable single-line marker 不生成 `dcci`，whole-cache marker 生成 whole-cache `dcci` | -| `TWait` 后读取 cacheable scalar GM payload | payload load 前生成 `pto.cmo.cacheinvalid %addr single_cache_line : !pto.ptr`，或使用 whole-cache 形式；single-line 地址必须由 PyPTO 或用户保证覆盖后续 load | PTOAS 检查 CMO 存在和顺序，不校验 acquire CMO 与 load 的 alias 关系 | -| `TTest` ready path 后读取 cacheable scalar GM payload | 在 ready path 的 payload load 前生成 single-line 或 whole-cache `pto.cmo.cacheinvalid`；single-line 地址必须由 PyPTO 或用户保证覆盖后续 load | PTOAS 检查 CMO 存在和顺序，不校验 acquire CMO 与 load 的 alias 关系 | -| cacheable scalar GM store 后发布 signal | 在 payload store 后生成 `pto.cmo.cacheinvalid %payload single_cache_line`，或使用 `pto.cmo.cacheinvalid all #pto.address_space`；随后生成 `pto.fence.barrier_all #pto.fence_scope` | single-line 只覆盖指定 cache line；whole-cache 覆盖全部 scalar D-cache，并保守选择全部 pending GM payload | - -如果某个 TNotify 只是发布 signal，而不表示某块 GM payload 已经准备好，可以不生成 payload -marker。PTOAS 不会尝试从所有前序 memory op 里推断“可能相关”的 payload。 - -`pto.entry` launcher 可以调用多个 kernel 函数；每个 kernel 函数会被 -`pto-memory-consistency` 独立分析。kernel body 内部若通过 `func.call` 调用包含 payload -访问、CMO、fence 或 signal op 的 helper，PyPTO 应在 `pto-memory-consistency` 前将 helper -inline，或者把 payload、CMO、fence 和 signal 保持在同一个 caller 中。否则 pass 会报错， -避免 caller 侧 `TNotify` 或 `TWait` 看不到 callee 内部的 memory-consistency 状态。 - -### 6.1 Issue #872：TPUT 发布 Signal - -```mlir -pto.comm.tput ... -pto.cmo.cacheinvalid %peer_payload single_cache_line : !pto.partition_tensor_view<...> -pto.fence.barrier_all #pto.fence_scope -pto.comm.tnotify ... -``` - -对应 PyPTO 写法： - -```python -pto.TPutOp(...) -pto.CmoCacheInvalidOp(pto.AddressSpace.GM, addr=peer_payload) -pto.FenceBarrierAllOp(pto.FenceScope.GM) -pto.TNotifyOp(...) -``` - -这个形态对应 #872 中的 `TPUT -> TNotify` 问题。`TPUT` macro op 内部会通过 MTE3 写 -peer GM payload；如果直接发布 `TNotify`，receiver 可能先观察到 signal ready，但 payload -写入还没有完成 pipe drain 或进入 GM visibility domain。 - -PTOAS 会识别 `TPUT` macro model 中的 MTE3 GM write phase，并在 `barrier_all` 前自动补齐 -低层 pipe drain： - -```mlir -pto.comm.tput ... -pto.cmo.cacheinvalid %peer_payload single_cache_line : !pto.partition_tensor_view<...> -pto.barrier #pto.pipe -pto.fence.barrier_all #pto.fence_scope -pto.comm.tnotify ... -``` - -EmitC 最终生成的关键顺序是： - -```cpp -TPUT(...); -pipe_barrier(PIPE_MTE3); -dsb(DSB_DDR); -TNOTIFY(...); -``` - -这里的 `cmo.cacheinvalid %peer_payload` 对 `TPUT` 是 marker-only：它告诉 PTOAS 本次 -`TNotify` 发布的是哪一块 peer payload。因为 `TPUT` 内部 MTE3 store 是 non-cacheable -路径，最终不会生成 `dcci`，只会让 PTOAS 在 `fence.barrier_all` 前精准补上匹配 payload -的 MTE3 drain。 - -### 6.2 TWait 后读取 Scalar Payload - ```mlir +// TWait 后读取 cacheable scalar GM payload pto.comm.twait ... -pto.cmo.cacheinvalid %payload_ptr single_cache_line : !pto.ptr -%value = pto.load_scalar ... -``` - -对应 PyPTO 写法： - -```python -pto.TWaitOp(...) -pto.CmoCacheInvalidOp(pto.AddressSpace.GM, addr=payload_ptr) -pto.LoadScalarOp(...) -``` - -### 6.3 Scalar Store 发布 Signal - -```mlir -pto.store_scalar %value, %payload[%idx] : !pto.ptr, i32 pto.cmo.cacheinvalid %payload single_cache_line : !pto.ptr -pto.fence.barrier_all #pto.fence_scope -pto.comm.tnotify ... -``` - -对应 PyPTO 写法： - -```python -pto.StoreScalarOp(...) -pto.CmoCacheInvalidOp(pto.AddressSpace.GM, addr=payload_ptr) -pto.FenceBarrierAllOp(pto.FenceScope.GM) -pto.TNotifyOp(...) +%value = pto.load_scalar %payload[%idx] : !pto.ptr -> i32 ``` -这对应 cacheable scalar GM store 发布 payload 的场景。`CmoCacheInvalidOp` 和 -`FenceBarrierAllOp` 都是必需的，且顺序必须是 CMO 在 fence 之前。release 侧的 -`pto.cmo.cacheinvalid all #pto.address_space` 同时表示 whole-cache CMO 边界和保守全量 -payload marker。使用 whole-cache 形式时，不需要再额外生成 matching -`pto.cmo.cacheinvalid %payload single_cache_line`；代价是 PTOAS 会保守选择该 op 之前已经 -pending 的所有 GM payload access，并为这些访问补齐对应 pipe drain。 - -## 7. Backend Lowering 状态 - -### 7.1 EmitC - -EmitC backend 当前支持： - -- `pto.cmo.cacheinvalid all #pto.address_space` lower 到 `dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE)`。 -- `pto.cmo.cacheinvalid %addr single_cache_line` 在 cacheable scalar release 或 acquire 路径 lower 到 `dcci((__gm__ void*)addr, cache_line_t::SINGLE_CACHE_LINE)`。 -- `pto.cmo.cacheinvalid %addr single_cache_line` 在 non-cacheable TLoad、TStore、TStoreFP 或 comm macro release 路径只作为 payload marker，被 MemoryConsistency pass 消费后消除。 -- `pto.cmo.cacheinvalid all #pto.address_space` 在 release 路径作为保守全量 marker，选择该 op 之前已经 pending 的所有 GM payload access；它不会被消除。 -- `pto.fence.barrier_all #pto.fence_scope` lower 到 `dsb(DSB_DDR)`。 - -### 7.2 VPTO - -VPTO backend 当前已经支持低层 `pto.dsb` 和 `pto.dcci` 到 HIVM intrinsic 的 lowering。 - -本 PR 暴露的 `pto.cmo.cacheinvalid` 和 `pto.fence.barrier_all` 是较高层的 -signal/payload 一致性契约 IR。当前还没有在 VPTO pipeline 中把这两类高层 op 自动降成 -低层 `pto.dcci` 和 `pto.dsb`，因此它们进入 VPTO LLVM lowering 时仍然 fail-fast： - -- `pto.cmo.cacheinvalid` -- `pto.fence.barrier_all` - -如果这些 op 进入 VPTO LLVM lowering，PTOAS 会报错，提示 VPTO backend 尚不支持这些 -high-level memory-consistency op。后续需要补一层 VPTO memory-consistency lowering， -把明确的 CMO 和 fence 语义转换为 `pto.dcci` 与 `pto.dsb`。 - -## 8. 当前限制 +如果涉及 non-cacheable MTE/FIX/comm macro payload，当前 PTOAS 不再自动根据 +payload 类型做精确判断。上游只需要在 publish 点显式生成 +`pto.fence.barrier_all #pto.fence_scope`；EmitC lowering 会在 `dsb(DSB_DDR)` +前保守生成 `PIPE_MTE2`、`PIPE_MTE3` 和 `PIPE_FIX` 的 `pipe_barrier`。如果上游 +已经手写了同类 `pto.barrier`，可能会产生重复 barrier，因此一般不建议在 +`fence.barrier_all` 前再手写这些 pipe drain。 -- `cmo.cacheinvalid` 支持 whole-cache 和 single-cache-line 粒度，但还没有连续 range 形式。 -- MemoryConsistency pass 当前不证明 acquire 侧 single-line CMO 是否覆盖后续 `load_scalar` payload，地址正确性由 PyPTO 或用户保证；不确定时应使用 whole-cache `cmo.cacheinvalid all #pto.address_space`。 -- `TWait` 和 `TTest` acquire 侧当前只覆盖 `load_scalar`。 -- VPTO 暂不支持 high-level `cmo.cacheinvalid` 和 `fence.barrier_all` 的真实 lowering； - 低层 `pto.dcci` 和 `pto.dsb` 已有 VPTO lowering。 -- 对复杂 CFG 的分析仍是保守近似，不做完整 path-sensitive 数据流。 +## 后续重新引入自动分析的要求 -## 9. 后续工作 +后续若恢复 MemoryConsistency pass，需要先满足以下条件： -1. 在 VPTO pipeline 中把 high-level `cmo.cacheinvalid` 与 `fence.barrier_all` 降到 - low-level `pto.dcci` 与 `pto.dsb`。 -2. 将多条 single-line `cacheinvalid` 优化成精确 GM address range CMO。 -3. 如果后续确认 release writeback 需要不同 destination 或更精确语义，再决定是否引入新的 CMO op。 -4. 扩展 acquire 侧 consumer 范围，从 `load_scalar` 扩展到更多 cacheable GM read。 -5. 将 macro op phase 的 memory descriptor 做得更精细，减少误报。 +- state merge 必须 bounded，不能通过简单追加导致指数增长。 +- state key 至少应按 payload identity、pipe、cacheability 和 action kind 去重。 +- `scf.if` 的 region exit state 不能重复携带父层入口 state。 +- loop 和多 block region 需要明确固定点或保守 summary 上界。 +- 新 pass 需要包含 issue #950 规模的 compile-time 回归测试。 diff --git a/docs/designs/ptodsl-simt-micro-op-api-design.md b/docs/designs/ptodsl-simt-micro-op-api-design.md index ba0ef216c6..92d4e58513 100644 --- a/docs/designs/ptodsl-simt-micro-op-api-design.md +++ b/docs/designs/ptodsl-simt-micro-op-api-design.md @@ -34,7 +34,7 @@ math, conversion, sync, and state preservation. Current PTO-DSL already has a narrow SIMT surface: - `@pto.simt` decorator and `with pto.simt():` inline scope. -- `@pto.simt(max_threads=..., max_regs=...)` optional entry resource +- `@pto.simt(max_threads=...)` optional entry resource attributes. - `pto.store_vfsimt_info(dim_z, dim_y, dim_x)`. - `pto.get_tid_x()`, `pto.get_tid_y()`, `pto.get_tid_z()`. @@ -284,15 +284,14 @@ explicit peer function symbol. ### 5.6 `@pto.simt` Decorator Attributes -SIMT entry functions may carry optional VPTO attributes: +SIMT entry functions may carry an optional VPTO attribute: - `pto.simt_max_threads` -- `pto.simt_max_regs` -PTO-DSL exposes them through `@pto.simt`: +PTO-DSL exposes it through `@pto.simt`: ```python -@pto.simt(max_threads=256, max_regs=48) +@pto.simt(max_threads=256) def body(...): ... ``` @@ -302,14 +301,13 @@ Lowering: ```mlir func.func @body(...) attributes { pto.simt_entry, - pto.simt_max_threads = 256 : i32, - pto.simt_max_regs = 48 : i32 + pto.simt_max_threads = 256 : i32 } ``` -Lowering attaches these attributes to the generated specialized helper function, -not to the authored Python symbol. Omitting either argument emits no explicit -attribute and lets backend defaults apply. +Lowering attaches the attribute to the generated specialized helper function, +not to the authored Python symbol. Omitting the argument emits no explicit +attribute and lets the backend default (1024) apply. Validation: @@ -432,8 +430,10 @@ Minimum Python/frontend tests: produces distinct helper functions and distinct traced bodies. 8. `pto.simt_launch` callee attributes reference the actual generated helper symbols. -9. `@pto.simt(max_threads=..., max_regs=...)` emits `pto.simt_max_threads` and - `pto.simt_max_regs` on the generated helper function. +9. `@pto.simt(max_threads=...)` emits `pto.simt_max_threads` on the generated + helper function. The register budget (`simt-max-registers`) is automatically + derived from `max_threads` by the backend according to the hardware resource + partition and is no longer a configurable attribute. Suggested lit/frontend assertions: diff --git a/docs/isa/micro-isa/01-pipeline-sync.md b/docs/isa/micro-isa/01-pipeline-sync.md index c20baa14c3..31c9c10d68 100644 --- a/docs/isa/micro-isa/01-pipeline-sync.md +++ b/docs/isa/micro-isa/01-pipeline-sync.md @@ -58,12 +58,12 @@ pipe_barrier(pipe); ```mlir // Both stores target the same GM address — order matters! -pto.mte_ub_gm %ub_partial_0, %gm_result, ... +pto.mte_ub_gm %ub_partial_0, %gm_result, %len_burst ... // Without pipe_barrier, MTE3 could execute the second copy before the first // completes, producing a non-deterministic result at %gm_result. pto.pipe_barrier "PIPE_MTE3" // After barrier: first copy is guaranteed complete. Second copy overwrites deterministically. -pto.mte_ub_gm %ub_partial_1, %gm_result, ... +pto.mte_ub_gm %ub_partial_1, %gm_result, %len_burst ... ``` --- @@ -367,7 +367,7 @@ pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] // MTE3 waits until Vector's signal arrives pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] -pto.mte_ub_gm %ub_out, %gm_out, ... +pto.mte_ub_gm %ub_out, %gm_out, %len_burst ... ``` **Key property:** Every cross-pipeline edge is an explicit `(set_flag, wait_flag)` pair. Simple for straight-line code, but gets verbose in loops (see Example 3). @@ -407,7 +407,7 @@ pto.rls_buf "PIPE_V", %bufid_ub_out, %c0 : i64, i64 // ─── Stage 3: MTE3 stores result to GM ─── // MTE3 acquires ub_out — blocks until Vector releases it (RAW: V write → MTE3 read) pto.get_buf "PIPE_MTE3", %bufid_ub_out, %c0 : i64, i64 -pto.mte_ub_gm %ub_out, %gm_out, ... +pto.mte_ub_gm %ub_out, %gm_out, %len_burst ... // MTE3 done reading ub_out — release so Vector can reuse it in next iteration pto.rls_buf "PIPE_MTE3", %bufid_ub_out, %c0 : i64, i64 ``` @@ -483,7 +483,7 @@ scf.for %i = %c0 to %N step %c1 { // ── MTE3: store result from buf_out[i%2] to GM ── // RAW: wait for Vector to finish writing buf_out[i%2] pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVT_OUT_FWD_{pp}"] - pto.mte_ub_gm %ub_out[%pp], %gm_out[%i], ... + pto.mte_ub_gm %ub_out[%pp], %gm_out[%i], %len_burst ... // WAR: tell Vector "done reading buf_out[i%2]" pto.set_flag["PIPE_MTE3", "PIPE_V", "EVT_OUT_REV_{pp}"] } @@ -534,7 +534,7 @@ scf.for %i = %c0 to %N step %c1 { // ── MTE3: store result ── // Acquires out[i%2] — blocks until Vector releases it (RAW: automatic) pto.get_buf "PIPE_MTE3", %bufid_out[%pp], %c0 : i64, i64 - pto.mte_ub_gm %ub_out[%pp], %gm_out[%i], ... + pto.mte_ub_gm %ub_out[%pp], %gm_out[%i], %len_burst ... pto.rls_buf "PIPE_MTE3", %bufid_out[%pp], %c0 : i64, i64 } // No post-loop drain needed — last rls_buf completes the pipeline. diff --git a/docs/isa/micro-isa/02-dma-copy.md b/docs/isa/micro-isa/02-dma-copy.md index 4747237c04..cc93e3bd7b 100644 --- a/docs/isa/micro-isa/02-dma-copy.md +++ b/docs/isa/micro-isa/02-dma-copy.md @@ -73,12 +73,13 @@ pto.mte_gm_ub %gm_in, %ub_out, %cache, %len_burst - **syntax:** ```mlir pto.mte_ub_gm %ub_src, %gm_dst, %len_burst - nburst(%n_burst, %src_stride, %dst_stride) + nburst(%n_burst, %src_stride, %dst_stride) l2_cache_ctl(%l2_cache_ctl) [loop(%loop_count, %loop_src_stride, %loop_dst_stride)]* - : !pto.ptr, !pto.ptr, i64, i64, i64, i64, + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, [loop i64, i64, i64,]* ``` - **semantics:** Grouped UB→GM DMA transfer. `nburst(...)` defines the innermost repeated burst transfer, and optional `loop(...)` groups add outer repetition levels. + The `l2_cache_ctl(...)` group is optional in textual VPTO IR; when omitted, lowering uses `0`. **Parameter Table:** @@ -88,6 +89,7 @@ pto.mte_ub_gm %ub_src, %gm_dst, %len_burst | `%gm_dst` | ptr | GM destination pointer (`!pto.ptr`) | | `%len_burst` | 16 bits | Contiguous bytes transferred per burst row | | `nburst(%n_burst, %src_stride, %dst_stride)` | 16 bits / 21 bits / 40 bits | Required innermost burst group: count, UB source stride, GM destination stride | +| `l2_cache_ctl(%l2_cache_ctl)` | 4 bits | Optional GM store-side L2 cache control; omitted means `0` | | `loop(%loop_count, %loop_src_stride, %loop_dst_stride)` | 21 bits / 21 bits / 40 bits | Optional outer repetition group: count, UB source stride, GM destination stride | **Constraints:** @@ -103,10 +105,10 @@ pto.mte_ub_gm %ub_src, %gm_dst, %len_burst ```mlir pto.mte_ub_gm %ub_in, %gm_out, %len_burst - nburst(%rows, %ub_row_stride, %gm_row_stride) + nburst(%rows, %ub_row_stride, %gm_row_stride) l2_cache_ctl(%l2_cache_ctl) loop(%tiles, %ub_tile_stride, %gm_tile_stride) loop(%batches, %ub_batch_stride, %gm_batch_stride) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64, + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, loop i64, i64, i64, loop i64, i64, i64 ``` @@ -336,6 +338,7 @@ For a form ```mlir pto.mte_ub_gm %ub_src, %dst, %len_burst nburst(%n_burst, %src_stride, %dst_stride) + l2_cache_ctl(%l2_cache_ctl) loop(%c0, %s0, %d0) loop(%c1, %s1, %d1) ... @@ -507,8 +510,8 @@ GM (dest, 32 × 32 f32): ```mlir pto.mte_ub_gm %ub_out, %arg1, %c128_i64 - nburst(%c32_i64, %c128_i64, %c128_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + nburst(%c32_i64, %c128_i64, %c128_i64) l2_cache_ctl(%c0_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 ``` --- @@ -545,8 +548,8 @@ GM (dest, into 1024 × 512 matrix): ```mlir pto.mte_ub_gm %ub_ptr, %gm_ptr, %c256_i64 - nburst(%c64_i64, %c256_i64, %c1024_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + nburst(%c64_i64, %c256_i64, %c1024_i64) l2_cache_ctl(%c0_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 ``` --- diff --git a/docs/isa/micro-isa/17-simt.md b/docs/isa/micro-isa/17-simt.md index a260f8f2ec..1f57802d9b 100644 --- a/docs/isa/micro-isa/17-simt.md +++ b/docs/isa/micro-isa/17-simt.md @@ -48,24 +48,32 @@ The current PTO SIMT surface supports these operation families: | Conversion | `pto.convert` | | Entry synchronization and state | `pto.syncthreads`, `pto.threadfence`, `pto.threadfence_block`, `pto.keep`, `pto.resume` | -Two optional function attributes may be attached to a `pto.simt_entry` +One optional function attribute may be attached to a `pto.simt_entry` function: | Function attribute | Type | Default | Meaning | |--------------------|------|---------|---------| | `pto.simt_max_threads` | signless `i32` integer attribute | `1024` | Compile-time launch envelope. It should cover the largest `dim_x * dim_y * dim_z` launch count used for this entry. | -| `pto.simt_max_regs` | signless `i32` integer attribute | `32` | Compile-time scalar register budget per workitem. Lower values constrain scalar live state; higher values permit more scalar live values with higher resource pressure. | -Both attributes are optional. If present, they must be positive `i32` -attributes and may only appear on functions that also carry `pto.simt_entry`. -They do not launch work by themselves; the actual workitem count comes from -`pto.store_vfsimt_info` or `pto.simt_launch`. +`pto.simt_max_threads` may only appear on functions that also carry +`pto.simt_entry`. It must be a positive `i32` value no greater than 2048. The +thread envelope determines the emitted scalar register budget: + +| `pto.simt_max_threads` | Emitted `simt-max-registers` | +|------------------------|------------------------------| +| `1` to `256` | `128` | +| `257` to `512` | `64` | +| `513` to `1024` | `32` | +| `1025` to `2048` | `16` | + +The register budget is derived automatically and is not independently +configurable. `pto.simt_max_threads` does not launch work by itself; the actual +workitem count comes from `pto.store_vfsimt_info` or `pto.simt_launch`. ```mlir func.func @body(%dst: !pto.ptr) attributes {pto.simt_entry, - pto.simt_max_threads = 256 : i32, - pto.simt_max_regs = 48 : i32} { + pto.simt_max_threads = 256 : i32} { return } ``` @@ -1071,8 +1079,8 @@ func.call @body(%ub_out) : (!pto.ptr) -> () pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] pto.mte_ub_gm %ub_out, %gm_out, %len - nburst(%n, %src_stride, %dst_stride) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + nburst(%n, %src_stride, %dst_stride) l2_cache_ctl(%l2_cache_ctl) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 ``` For pipeline synchronization semantics, see @@ -1103,8 +1111,8 @@ module attributes {pto.target_arch = "a5", pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] pto.mte_ub_gm %ub_out, %out, %c128_i64 - nburst(%c32_i64, %c128_i64, %c128_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + nburst(%c32_i64, %c128_i64, %c128_i64) l2_cache_ctl(%c0_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 pto.barrier #pto.pipe return } diff --git a/docs/vpto-spec.md b/docs/vpto-spec.md index a3217572d4..4c13bc72ba 100644 --- a/docs/vpto-spec.md +++ b/docs/vpto-spec.md @@ -379,8 +379,8 @@ pto.vecscope { pto.set_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] pto.wait_flag["PIPE_V", "PIPE_MTE3", "EVENT_ID0"] pto.mte_ub_gm %8, %14, %c128_i64 - nburst(%c32_i64, %c128_i64, %c128_i64) - : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + nburst(%c32_i64, %c128_i64, %c128_i64) l2_cache_ctl(%c0_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 ``` ### Example: Strict VecScope @@ -479,6 +479,11 @@ pipeline** commits the signal, ensuring the preceding operation on that pipeline before the signal is issued. The receiver specifies which **local pipeline** should stall until the signal arrives. This is the fundamental IPC primitive for Cube–Vector cooperation on A5. +For the public `pto.sync.set` / `pto.sync.wait` surface on A5, event IDs are physical semaphore +IDs in the `0-31` range. AIV subblock 0 uses event IDs `0-15`, and AIV subblock 1 uses event +IDs `16-31`. Code that signals or waits for both AIV subblocks should explicitly emit both +the base ID and `base_id + 16`. + > **Note:** `pto.set_cross_core` / `pto.wait_cross_core` operate at **multi-cluster** scope and > are not used for intra-cluster communication. diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 6ed8a63774..c722d5406a 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -2962,16 +2962,13 @@ def CmoCacheInvalidOp : PTO_Op<"cmo.cacheinvalid"> { `dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE)`. The single-cache-line form carries the GM payload address selected by the - producer DSL/user. When it is required for a cacheable scalar path, it - lowers to `dcci(addr, cache_line_t::SINGLE_CACHE_LINE)`. When it only marks - a non-cacheable payload for signal release analysis, MemoryConsistency may - consume it as marker-only IR and erase it before EmitC lowering. + producer DSL/user and lowers to + `dcci(addr, cache_line_t::SINGLE_CACHE_LINE)`. - MemoryConsistency interprets both forms as explicit CMO boundaries on - producer release and consumer acquire paths. It also uses the addressed form - to match a TNotify payload against pending MTE2, MTE3, FIX, or scalar GM - accesses before deciding which pipe drains are required. Precise coverage of - the selected line is the producer DSL/user's responsibility in this phase. + This op is an explicit CMO boundary. PTOAS no longer runs the historical + MemoryConsistency analysis pass that consumed marker-only CMO ops or inferred + signal/payload pipe drains, so precise placement and coverage are the + producer DSL/user's responsibility. }]; let arguments = (ins @@ -2988,8 +2985,10 @@ def FenceBarrierAllOp : PTO_Op<"fence.barrier_all"> { let description = [{ Memory fence for a visibility scope. Around signal/payload communication it can be used as the producer-side publish barrier or the consumer-side - acquire barrier depending on placement. `scope = gm` lowers to - `dsb(DSB_DDR)`. + acquire barrier depending on placement. In EmitC, `scope = gm` and + `scope = all` conservatively drain MTE2, MTE3, and FIX before lowering to + `dsb(DSB_DDR)`, so users do not need to separately model those producer + pipe classes before a GM visibility fence. }]; let arguments = (ins PTO_FenceScopeAttr:$scope); diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index 905f4bb4fe..757034284e 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -2645,7 +2645,7 @@ def PTO_CopyUbufToGmOp : PTO_Op<"copy_ubuf_to_gm", [ I64:$sid, I64:$n_burst, I64:$len_burst, - I64:$reserved, + I64:$l2_cache_ctl, I64:$burst_dst_stride, I64:$burst_src_stride ); @@ -2656,9 +2656,9 @@ def PTO_CopyUbufToGmOp : PTO_Op<"copy_ubuf_to_gm", [ let assemblyFormat = [{ $source `,` $destination `,` $sid `,` $n_burst `,` $len_burst `,` - $reserved `,` $burst_dst_stride `,` $burst_src_stride + $l2_cache_ctl `,` $burst_dst_stride `,` $burst_src_stride attr-dict `:` type($source) `,` type($destination) `,` type($sid) `,` type($n_burst) `,` - type($len_burst) `,` type($reserved) `,` + type($len_burst) `,` type($l2_cache_ctl) `,` type($burst_dst_stride) `,` type($burst_src_stride) }]; } @@ -2674,6 +2674,7 @@ def PTO_MteUbGmOp : PTO_Op<"mte_ub_gm", [ I64:$n_burst, I64:$nburst_src_stride, I64:$nburst_dst_stride, + Optional:$l2_cache_ctl, Variadic:$loop_counts, Variadic:$loop_src_strides, Variadic:$loop_dst_strides @@ -2690,6 +2691,7 @@ def PTO_MteUbGmOp : PTO_Op<"mte_ub_gm", [ "::mlir::Value":$destination, "::mlir::Value":$lenBurst, "::mlir::pto::DmaLoopConfig":$nburst, + "::mlir::Value":$l2CacheCtl, "::llvm::ArrayRef<::mlir::pto::DmaLoopConfig>":$loops )>, OpBuilder<(ins @@ -2697,6 +2699,7 @@ def PTO_MteUbGmOp : PTO_Op<"mte_ub_gm", [ "::mlir::Value":$destination, "::mlir::Value":$lenBurst, "::mlir::pto::DmaLoopConfig":$nburst, + "::mlir::Value":$l2CacheCtl, "::std::optional<::mlir::pto::DmaLoopConfig>":$loop1, "::std::optional<::mlir::pto::DmaLoopConfig>":$loop2 )> diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index cdbd8e1c0a..3555735a06 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -77,7 +77,6 @@ std::unique_ptr createPTOValidateIntToPtrUsesPass(); std::unique_ptr createPTORematerializeFixpipeVectorQuantPass(); std::unique_ptr createPTOMaterializeTileHandlesPass(); std::unique_ptr createPTOResolveBufferSelectPass(); -std::unique_ptr createPTOMemoryConsistencyPass(); std::unique_ptr createInferPTOLayoutPass(); std::unique_ptr createPTOA5NormalizeTMovPass(); std::unique_ptr createPTORemoveIdentityTMovPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 61e6e4b00d..00b57ee697 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -782,25 +782,6 @@ def PTOResolveBufferSelect : Pass<"pto-resolve-buffer-select", "ModuleOp"> { ]; } -def PTOMemoryConsistency : Pass<"pto-memory-consistency", "ModuleOp"> { - let summary = "Annotate PTO memory consistency actions before backend lowering"; - let description = [{ - Analyzes signal/payload ordering requirements and annotates communication - signal ops and scalar GM consumers with release/acquire actions consumed by - backend lowering. It covers TNotify release actions for direct MTE - operations, macro-op MTE3 phases, cacheable scalar GM stores, and - conservative TWait/TTest acquire invalidation for scalar GM loads. - }]; - - let constructor = "mlir::pto::createPTOMemoryConsistencyPass()"; - - let dependentDialects = [ - "mlir::pto::PTODialect", - "mlir::func::FuncDialect", - "mlir::scf::SCFDialect" - ]; -} - def PTOUnrollSIMTFor : Pass<"pto-unroll-simt-for", "func::FuncOp"> { let summary = "Unroll explicitly annotated scf.for loops in pto.simt_entry functions"; diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 1b66634049..b8d1568331 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -3389,6 +3389,14 @@ LogicalResult mlir::pto::SyncSetOp::verify() { auto verifyA2A3 = [&]() -> LogicalResult { return success(); }; auto verifyA5 = [&]() -> LogicalResult { + if (IntegerAttr eventIdAttr = getEventIdAttr()) { + int64_t eventId = eventIdAttr.getInt(); + if (eventId < 0 || eventId > 31) { + return emitOpError() + << "A5 sync.set expects static event_id in [0, 31], but got " + << eventId; + } + } switch (getPipe().getPipe()) { case PIPE::PIPE_FIX: case PIPE::PIPE_MTE3: @@ -3552,6 +3560,13 @@ LogicalResult mlir::pto::SyncWaitOp::verify() { auto verifyA2A3 = [&]() -> LogicalResult { return success(); }; auto verifyA5 = [&]() -> LogicalResult { + if (IntegerAttr eventIdAttr = getEventIdAttr()) { + int64_t eventId = eventIdAttr.getInt(); + if (eventId < 0 || eventId > 31) + return emitOpError() + << "A5 sync.wait expects static physical event_id in [0, 31], but got " + << eventId; + } switch (getPipe().getPipe()) { case PIPE::PIPE_FIX: case PIPE::PIPE_MTE1: diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index 29ad4e648d..e76d1aee73 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -2578,6 +2578,105 @@ static bool isStructuredAccStoreInt32PreQuantMode(AccStoreQuantPreMode mode) { } } +enum class StructuredAccStoreDestinationFamily { + Any, + F32, + F16, + BF16, + I32, + I16, + I8, + I4, + FP8 +}; + +static StructuredAccStoreDestinationFamily +getStructuredAccStorePreQuantDestinationFamily(AccStoreQuantPreMode mode) { + switch (mode) { + case AccStoreQuantPreMode::F32F16: + case AccStoreQuantPreMode::QF322F16PreVec: + case AccStoreQuantPreMode::QF322F16PreScalar: + case AccStoreQuantPreMode::DEQF16Vec: + case AccStoreQuantPreMode::DEQF16Scalar: + return StructuredAccStoreDestinationFamily::F16; + case AccStoreQuantPreMode::F32BF16: + case AccStoreQuantPreMode::QF322BF16PreVec: + case AccStoreQuantPreMode::QF322BF16PreScalar: + case AccStoreQuantPreMode::QS322BF16PreVec: + case AccStoreQuantPreMode::QS322BF16PreScalar: + return StructuredAccStoreDestinationFamily::BF16; + case AccStoreQuantPreMode::QF322F32PreVec: + case AccStoreQuantPreMode::QF322F32PreScalar: + return StructuredAccStoreDestinationFamily::F32; + case AccStoreQuantPreMode::QF322HIF8PreVec: + case AccStoreQuantPreMode::QF322HIF8PreScalar: + case AccStoreQuantPreMode::QF322HIF8PreHybridVec: + case AccStoreQuantPreMode::QF322HIF8PreHybridScalar: + case AccStoreQuantPreMode::QF322FP8PreVec: + case AccStoreQuantPreMode::QF322FP8PreScalar: + return StructuredAccStoreDestinationFamily::FP8; + case AccStoreQuantPreMode::DEQS32IntVec: + case AccStoreQuantPreMode::DEQS32IntScalar: + return StructuredAccStoreDestinationFamily::I32; + case AccStoreQuantPreMode::QF162S16PreVec: + case AccStoreQuantPreMode::QF162S16PreScalar: + case AccStoreQuantPreMode::DEQS16Vec: + case AccStoreQuantPreMode::DEQS16Scalar: + return StructuredAccStoreDestinationFamily::I16; + case AccStoreQuantPreMode::QF162B8PreVec: + case AccStoreQuantPreMode::QF162B8PreScalar: + case AccStoreQuantPreMode::REQ8Vec: + case AccStoreQuantPreMode::REQ8Scalar: + case AccStoreQuantPreMode::QF322B8PreVec: + case AccStoreQuantPreMode::QF322B8PreScalar: + return StructuredAccStoreDestinationFamily::I8; + case AccStoreQuantPreMode::QF162S4PreVec: + case AccStoreQuantPreMode::QF162S4PreScalar: + case AccStoreQuantPreMode::REQ4Vec: + case AccStoreQuantPreMode::REQ4Scalar: + case AccStoreQuantPreMode::QF322S4PreVec: + case AccStoreQuantPreMode::QF322S4PreScalar: + return StructuredAccStoreDestinationFamily::I4; + case AccStoreQuantPreMode::NoConvert: + return StructuredAccStoreDestinationFamily::Any; + } + llvm_unreachable("unknown acc-store pre_quant mode"); +} + +static bool isStructuredAccStoreDestinationFamily( + Type type, StructuredAccStoreDestinationFamily family) { + switch (family) { + case StructuredAccStoreDestinationFamily::Any: + return true; + case StructuredAccStoreDestinationFamily::F32: + return type.isF32(); + case StructuredAccStoreDestinationFamily::F16: + return type.isF16(); + case StructuredAccStoreDestinationFamily::BF16: + return type.isBF16(); + case StructuredAccStoreDestinationFamily::I32: + if (auto intType = dyn_cast(type)) + return intType.getWidth() == 32; + return false; + case StructuredAccStoreDestinationFamily::I16: + if (auto intType = dyn_cast(type)) + return intType.getWidth() == 16 && !intType.isUnsigned(); + return false; + case StructuredAccStoreDestinationFamily::I8: + if (auto intType = dyn_cast(type)) + return intType.getWidth() == 8; + return false; + case StructuredAccStoreDestinationFamily::I4: + if (auto intType = dyn_cast(type)) + return intType.getWidth() == 4 && !intType.isUnsigned(); + return false; + case StructuredAccStoreDestinationFamily::FP8: + return pto::isPTOFloat8Type(type) || pto::isPTOHiFloat8Type(type) || + pto::isPTOHiFloat8x2Type(type); + } + llvm_unreachable("unknown acc-store destination family"); +} + static ParseResult parseStructuredAccStoreUnitFlag(OpAsmParser &parser, StructuredAccStoreAsmState &state) { if (state.unitFlag) @@ -2838,7 +2937,11 @@ static LogicalResult verifyStructuredAccStoreLike( if (static_cast(preQuant) != static_cast(preQuantMode)) return op->emitOpError("pre_quant requires payload and mode together"); if (preQuantMode) { - if (isStructuredAccStoreVectorQuantMode(*preQuantMode)) { + if (*preQuantMode == AccStoreQuantPreMode::NoConvert) { + // The no_convert keyword carries no quantization parameters; the + // syntactic payload operand is ignored for compatibility with the + // structured pre_quant clause form. + } else if (isStructuredAccStoreVectorQuantMode(*preQuantMode)) { if (!isStructuredAccStoreScalingPayload(preQuant)) return op->emitOpError("vector pre_quant mode requires scaling pointer payload"); if (!isStructuredAccStoreFloatScalarPayloadType( @@ -2857,16 +2960,24 @@ static LogicalResult verifyStructuredAccStoreLike( << " and destination element type " << destinationElementType; }; - if (isa(sourceElementType)) { - if (!isStructuredAccStoreFloatPreQuantMode(*preQuantMode)) - return emitIncompatibleQuantModeError(); - } else if (sourceElementType.isSignlessInteger(32)) { - if (!isStructuredAccStoreInt32PreQuantMode(*preQuantMode)) + if (*preQuantMode != AccStoreQuantPreMode::NoConvert) { + if (isa(sourceElementType)) { + if (!isStructuredAccStoreFloatPreQuantMode(*preQuantMode)) + return emitIncompatibleQuantModeError(); + } else if (sourceElementType.isSignlessInteger(32)) { + if (!isStructuredAccStoreInt32PreQuantMode(*preQuantMode)) + return emitIncompatibleQuantModeError(); + } else { + return op->emitOpError() + << "pre_quant requires source element type to be f32 or i32, got " + << sourceElementType; + } + + StructuredAccStoreDestinationFamily destinationFamily = + getStructuredAccStorePreQuantDestinationFamily(*preQuantMode); + if (!isStructuredAccStoreDestinationFamily(destinationElementType, + destinationFamily)) return emitIncompatibleQuantModeError(); - } else { - return op->emitOpError() - << "pre_quant requires source element type to be f32 or i32, got " - << sourceElementType; } } @@ -6592,10 +6703,13 @@ LogicalResult CopyUbufToGmOp::verify() { } void MteUbGmOp::build(OpBuilder &builder, OperationState &state, Value source, - Value destination, Value lenBurst, pto::DmaLoopConfig nburst, + Value destination, Value lenBurst, + pto::DmaLoopConfig nburst, Value l2CacheCtl, llvm::ArrayRef loops) { state.addOperands({source, destination, lenBurst, nburst.count, nburst.srcStride, nburst.dstStride}); + if (l2CacheCtl) + state.addOperands(l2CacheCtl); for (const pto::DmaLoopConfig &loop : loops) state.addOperands(loop.count); for (const pto::DmaLoopConfig &loop : loops) @@ -6606,14 +6720,15 @@ void MteUbGmOp::build(OpBuilder &builder, OperationState &state, Value source, state.addAttribute( getOperandSegmentSizeAttr(), builder.getDenseI32ArrayAttr( - {1, 1, 1, 1, 1, 1, + {1, 1, 1, 1, 1, 1, l2CacheCtl ? 1 : 0, static_cast(loops.size()), static_cast(loops.size()), static_cast(loops.size())})); } void MteUbGmOp::build(OpBuilder &builder, OperationState &state, Value source, - Value destination, Value lenBurst, pto::DmaLoopConfig nburst, + Value destination, Value lenBurst, + pto::DmaLoopConfig nburst, Value l2CacheCtl, std::optional loop1, std::optional loop2) { SmallVector loops; @@ -6621,11 +6736,13 @@ void MteUbGmOp::build(OpBuilder &builder, OperationState &state, Value source, loops.push_back(*loop1); if (loop2) loops.push_back(*loop2); - build(builder, state, source, destination, lenBurst, nburst, loops); + build(builder, state, source, destination, lenBurst, nburst, l2CacheCtl, + loops); } ParseResult MteUbGmOp::parse(OpAsmParser &parser, OperationState &result) { - OpAsmParser::UnresolvedOperand source, destination, lenBurst; + OpAsmParser::UnresolvedOperand source, destination, lenBurst, l2CacheCtl; + bool hasL2CacheCtl = false; SmallVector nburstOperands; SmallVector loopCountOperands; SmallVector loopSrcStrideOperands; @@ -6635,6 +6752,12 @@ ParseResult MteUbGmOp::parse(OpAsmParser &parser, OperationState &result) { parser.parseOperand(lenBurst) || parseDmaTripleGroup(parser, "nburst", nburstOperands)) return failure(); + if (succeeded(parser.parseOptionalKeyword("l2_cache_ctl"))) { + hasL2CacheCtl = true; + if (parser.parseLParen() || parser.parseOperand(l2CacheCtl) || + parser.parseRParen()) + return failure(); + } while (true) { StringRef parsedKeyword; SmallVector loopGroupOperands; @@ -6651,7 +6774,7 @@ ParseResult MteUbGmOp::parse(OpAsmParser &parser, OperationState &result) { if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()) return failure(); - Type sourceType, destinationType, lenBurstType; + Type sourceType, destinationType, lenBurstType, l2CacheCtlType; SmallVector nburstTypes, loopCountTypes, loopSrcStrideTypes, loopDstStrideTypes; if (parser.parseType(sourceType) || parser.parseComma() || @@ -6659,6 +6782,10 @@ ParseResult MteUbGmOp::parse(OpAsmParser &parser, OperationState &result) { parser.parseType(lenBurstType) || parser.parseComma() || parseDmaTripleTypes(parser, nburstTypes)) return failure(); + if (hasL2CacheCtl) { + if (parser.parseComma() || parser.parseType(l2CacheCtlType)) + return failure(); + } while (succeeded(parser.parseOptionalComma())) { StringRef keyword; if (parser.parseKeyword(&keyword)) @@ -6690,6 +6817,7 @@ ParseResult MteUbGmOp::parse(OpAsmParser &parser, OperationState &result) { auto &segments = result.getOrAddProperties().operandSegmentSizes; llvm::copy(ArrayRef{1, 1, 1, 1, 1, 1, + hasL2CacheCtl ? 1 : 0, loopGroupCount, loopGroupCount, loopGroupCount}, segments.begin()); @@ -6697,8 +6825,12 @@ ParseResult MteUbGmOp::parse(OpAsmParser &parser, OperationState &result) { parser.resolveOperand(destination, destinationType, result.operands) || parser.resolveOperand(lenBurst, lenBurstType, result.operands) || parser.resolveOperands(nburstOperands, nburstTypes, parser.getCurrentLocation(), - result.operands) || - parser.resolveOperands(loopCountOperands, loopCountTypes, + result.operands)) + return failure(); + if (hasL2CacheCtl && + parser.resolveOperand(l2CacheCtl, l2CacheCtlType, result.operands)) + return failure(); + if (parser.resolveOperands(loopCountOperands, loopCountTypes, parser.getCurrentLocation(), result.operands) || parser.resolveOperands(loopSrcStrideOperands, loopSrcStrideTypes, parser.getCurrentLocation(), result.operands) || @@ -6714,6 +6846,8 @@ void MteUbGmOp::print(OpAsmPrinter &printer) { << getLenBurst(); printDmaTripleGroup(printer, "nburst", getNBurst(), getNburstSrcStride(), getNburstDstStride()); + if (Value l2CacheCtl = getL2CacheCtl()) + printer << " l2_cache_ctl(" << l2CacheCtl << ")"; for (auto [count, srcStride, dstStride] : llvm::zip(getLoopCounts(), getLoopSrcStrides(), getLoopDstStrides())) printDmaTripleGroup(printer, "loop", count, srcStride, dstStride); @@ -6723,6 +6857,8 @@ void MteUbGmOp::print(OpAsmPrinter &printer) { << ", " << getNburstSrcStride().getType() << ", " << getNburstDstStride().getType(); + if (Value l2CacheCtl = getL2CacheCtl()) + printer << ", " << l2CacheCtl.getType(); for (auto [count, srcStride, dstStride] : llvm::zip(getLoopCounts(), getLoopSrcStrides(), getLoopDstStrides())) printDmaTripleTypes(printer, "loop", count.getType(), srcStride.getType(), @@ -6753,6 +6889,13 @@ LogicalResult MteUbGmOp::verify() { if (sourceElemBytes != destinationElemBytes) return emitOpError( "requires source and destination element byte widths to match"); + if (Value l2CacheCtlValue = getL2CacheCtl()) { + APInt l2CacheCtl; + if (matchPattern(l2CacheCtlValue, m_ConstantInt(&l2CacheCtl)) && + (l2CacheCtl.isNegative() || l2CacheCtl.ugt(15))) + return emitOpError( + "requires constant l2_cache_ctl to fit in range [0, 15]"); + } return verifyDmaLoadStoreLoopGroups( getOperation(), getLoopCounts(), getLoopSrcStrides(), getLoopDstStrides()); diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index d3986843e2..ccd135000a 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -61,7 +61,6 @@ add_mlir_dialect_library(PTOTransforms PTORemoveIdentityTMov.cpp PTOCanonicalizeIR.cpp PTOMaterializeTileHandles.cpp - PTOMemoryConsistency.cpp BufferizableOpInterfaceImpl.cpp ConvertToPTOOp.cpp PTOAssignDefaultFrontendPipeIdPass.cpp diff --git a/lib/PTO/Transforms/PTOMemoryConsistency.cpp b/lib/PTO/Transforms/PTOMemoryConsistency.cpp deleted file mode 100644 index ab9b4d0737..0000000000 --- a/lib/PTO/Transforms/PTOMemoryConsistency.cpp +++ /dev/null @@ -1,928 +0,0 @@ -// 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. - -#include "PTO/IR/PTO.h" -#include "PTO/Transforms/InsertSync/SyncMacroModel.h" -#include "PTO/Transforms/MemoryConsistencyAttrs.h" -#include "PTO/Transforms/Passes.h" -#include "Utils.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/SymbolTable.h" -#include "mlir/Pass/Pass.h" -#include "llvm/ADT/DenseSet.h" - -namespace mlir { -namespace pto { -#define GEN_PASS_DEF_PTOMEMORYCONSISTENCY -#include "PTO/Transforms/Passes.h.inc" -} // namespace pto -} // namespace mlir - -using namespace mlir; -using namespace mlir::pto; - -namespace { - -static bool isGmAddressSpace(pto::AddressSpace space) { - return space == pto::AddressSpace::GM || space == pto::AddressSpace::Zero; -} - -static Value getPayloadAliasRoot(Value value) { - for (unsigned depth = 0; value && depth < 64; ++depth) { - Operation *def = value.getDefiningOp(); - if (!def) - return value; - - Value next; - if (auto part = dyn_cast(def)) { - next = part.getSource(); - } else if (auto make = dyn_cast(def)) { - next = make.getPtr(); - } else if (auto addPtr = dyn_cast(def)) { - next = addPtr.getPtr(); - } else if (auto castPtr = dyn_cast(def)) { - next = castPtr.getInput(); - } else if (auto unrealized = dyn_cast(def)) { - auto result = dyn_cast(value); - if (result && result.getResultNumber() < unrealized.getNumOperands()) - next = unrealized.getOperand(result.getResultNumber()); - } else if (auto alias = pto::getOperationAliasInfo(def)) { - if (alias->first == value) - next = alias->second; - } - - if (!next || next == value) - return value; - value = next; - } - return value; -} - -static bool payloadMayAlias(Value lhs, Value rhs) { - if (!lhs || !rhs) - return true; - Value lhsRoot = getPayloadAliasRoot(lhs); - Value rhsRoot = getPayloadAliasRoot(rhs); - return lhsRoot == rhsRoot; -} - -struct PendingReleaseAccess { - pto::PIPE pipe = pto::PIPE::PIPE_UNASSIGNED; - Value payload; - bool drainPending = false; - bool needsDsbDdr = false; - bool needsGmCacheCmo = false; - bool dsbCovered = false; - bool cmoCovered = false; - bool selectedByWholeCacheCmo = false; -}; - -struct TNotifyReleaseState { - bool drainMte2 = false; - bool drainMte3 = false; - bool drainFix = false; - bool needsDsbDdr = false; - bool needsGmCacheCmo = false; - bool hasAddressedCmo = false; - bool hasWholeCacheCmo = false; - bool addressedDrainMte2 = false; - bool addressedDrainMte3 = false; - bool addressedDrainFix = false; - bool addressedNeedsDsbDdr = false; - bool addressedNeedsGmCacheCmo = false; - SmallVector pendingAccesses; - SmallVector addressedCmoPayloads; - - void merge(const TNotifyReleaseState &other) { - drainMte2 |= other.drainMte2; - drainMte3 |= other.drainMte3; - drainFix |= other.drainFix; - needsDsbDdr |= other.needsDsbDdr; - needsGmCacheCmo |= other.needsGmCacheCmo; - hasAddressedCmo |= other.hasAddressedCmo; - hasWholeCacheCmo |= other.hasWholeCacheCmo; - addressedDrainMte2 |= other.addressedDrainMte2; - addressedDrainMte3 |= other.addressedDrainMte3; - addressedDrainFix |= other.addressedDrainFix; - addressedNeedsDsbDdr |= other.addressedNeedsDsbDdr; - addressedNeedsGmCacheCmo |= other.addressedNeedsGmCacheCmo; - pendingAccesses.append(other.pendingAccesses.begin(), - other.pendingAccesses.end()); - addressedCmoPayloads.append(other.addressedCmoPayloads.begin(), - other.addressedCmoPayloads.end()); - refreshNeedsDsbDdr(); - refreshNeedsGmCacheCmo(); - if (hasReleaseCmoMarker()) - recomputeAddressedState(); - } - - void clear() { - drainMte2 = false; - drainMte3 = false; - drainFix = false; - needsDsbDdr = false; - needsGmCacheCmo = false; - hasAddressedCmo = false; - hasWholeCacheCmo = false; - addressedDrainMte2 = false; - addressedDrainMte3 = false; - addressedDrainFix = false; - addressedNeedsDsbDdr = false; - addressedNeedsGmCacheCmo = false; - pendingAccesses.clear(); - addressedCmoPayloads.clear(); - } - - bool hasReleaseCmoMarker() const { - return hasAddressedCmo || hasWholeCacheCmo; - } - - bool accessMatchesAddressedCmo(const PendingReleaseAccess &access) const { - for (Value payload : addressedCmoPayloads) - if (payloadMayAlias(access.payload, payload)) - return true; - return false; - } - - bool accessMatchesReleaseCmo(const PendingReleaseAccess &access) const { - return access.selectedByWholeCacheCmo || accessMatchesAddressedCmo(access); - } - - void refreshNeedsGmCacheCmo() { - needsGmCacheCmo = false; - for (const PendingReleaseAccess &access : pendingAccesses) { - if (access.needsGmCacheCmo && !access.cmoCovered) { - needsGmCacheCmo = true; - return; - } - } - } - - void refreshNeedsDsbDdr() { - needsDsbDdr = false; - for (const PendingReleaseAccess &access : pendingAccesses) { - if (access.needsDsbDdr && !access.dsbCovered) { - needsDsbDdr = true; - return; - } - } - } - - void recomputeAddressedState() { - addressedDrainMte2 = false; - addressedDrainMte3 = false; - addressedDrainFix = false; - addressedNeedsDsbDdr = false; - addressedNeedsGmCacheCmo = false; - - for (const PendingReleaseAccess &access : pendingAccesses) { - if (!accessMatchesReleaseCmo(access)) - continue; - if (access.drainPending) { - if (access.pipe == pto::PIPE::PIPE_MTE2) - addressedDrainMte2 = true; - if (access.pipe == pto::PIPE::PIPE_MTE3) - addressedDrainMte3 = true; - if (access.pipe == pto::PIPE::PIPE_FIX) - addressedDrainFix = true; - } - addressedNeedsDsbDdr |= access.needsDsbDdr && !access.dsbCovered; - if (access.needsGmCacheCmo && !access.cmoCovered) - addressedNeedsGmCacheCmo = true; - } - } - - void addAccess(pto::PIPE pipe, Value payload, bool needsDsb, - bool needsCmo) { - PendingReleaseAccess access; - access.pipe = pipe; - access.payload = payload; - access.drainPending = pipe == pto::PIPE::PIPE_MTE2 || - pipe == pto::PIPE::PIPE_MTE3 || - pipe == pto::PIPE::PIPE_FIX; - access.needsDsbDdr = needsDsb; - access.needsGmCacheCmo = needsCmo; - pendingAccesses.push_back(access); - - switch (pipe) { - case pto::PIPE::PIPE_MTE2: - drainMte2 = true; - break; - case pto::PIPE::PIPE_MTE3: - drainMte3 = true; - break; - case pto::PIPE::PIPE_FIX: - drainFix = true; - break; - default: - break; - } - needsDsbDdr |= needsDsb; - needsGmCacheCmo |= needsCmo; - if (hasReleaseCmoMarker()) - recomputeAddressedState(); - } - - void markPipeDrained(pto::PIPE pipe) { - auto pipeMatches = [&](pto::PIPE accessPipe) { - return pipe == pto::PIPE::PIPE_ALL || accessPipe == pipe; - }; - for (PendingReleaseAccess &access : pendingAccesses) - if (pipeMatches(access.pipe)) - access.drainPending = false; - - if (pipe == pto::PIPE::PIPE_MTE2 || pipe == pto::PIPE::PIPE_ALL) { - drainMte2 = false; - addressedDrainMte2 = false; - } - if (pipe == pto::PIPE::PIPE_MTE3 || pipe == pto::PIPE::PIPE_ALL) { - drainMte3 = false; - addressedDrainMte3 = false; - } - if (pipe == pto::PIPE::PIPE_FIX || pipe == pto::PIPE::PIPE_ALL) { - drainFix = false; - addressedDrainFix = false; - } - if (hasReleaseCmoMarker()) - recomputeAddressedState(); - } - - void applyBarrier(pto::PIPE pipe) { - markPipeDrained(pipe); - } - - void applyFenceBarrierAll(pto::FenceScope scope) { - if (scope != pto::FenceScope::GM && scope != pto::FenceScope::All) - return; - - for (PendingReleaseAccess &access : pendingAccesses) { - if (!access.needsDsbDdr) - continue; - if (hasReleaseCmoMarker() && !accessMatchesReleaseCmo(access)) - continue; - if (access.drainPending) - continue; - if (access.needsGmCacheCmo && !access.cmoCovered) - continue; - access.dsbCovered = true; - } - - refreshNeedsDsbDdr(); - if (hasReleaseCmoMarker()) - recomputeAddressedState(); - } - - void applyCmoCacheInvalid(pto::CmoCacheInvalidOp cmo) { - pto::AddressSpace space = cmo.getSpace().getAddressSpace(); - if (!isGmAddressSpace(space)) - return; - Value addr = cmo.getAddr(); - if (!addr) { - hasWholeCacheCmo = true; - for (PendingReleaseAccess &access : pendingAccesses) { - access.selectedByWholeCacheCmo = true; - if (access.needsGmCacheCmo) - access.cmoCovered = true; - } - refreshNeedsGmCacheCmo(); - recomputeAddressedState(); - cmo->removeAttr(kCmoCacheInvalidSkipLoweringAttrName); - return; - } - - hasAddressedCmo = true; - addressedCmoPayloads.push_back(addr); - bool needsRealCmo = false; - for (PendingReleaseAccess &access : pendingAccesses) { - if (!payloadMayAlias(access.payload, addr)) - continue; - if (access.needsGmCacheCmo && !access.cmoCovered) { - access.cmoCovered = true; - needsRealCmo = true; - } - } - refreshNeedsGmCacheCmo(); - recomputeAddressedState(); - - if (needsRealCmo) { - cmo->removeAttr(kCmoCacheInvalidSkipLoweringAttrName); - } else { - cmo->setAttr(kCmoCacheInvalidSkipLoweringAttrName, - UnitAttr::get(cmo.getContext())); - } - } -}; - -struct SignalAcquireState { - bool pendingInvalidateGmCache = false; - bool dirtyGmCache = false; - - void merge(const SignalAcquireState &other) { - pendingInvalidateGmCache |= other.pendingInvalidateGmCache; - dirtyGmCache |= other.dirtyGmCache; - } - - void consumeAcquire() { - pendingInvalidateGmCache = false; - dirtyGmCache = false; - } - - void applyCmoCacheInvalid(pto::AddressSpace space, - pto::CmoCacheInvalidOp cmo = nullptr) { - if (!isGmAddressSpace(space)) - return; - if (cmo && pendingInvalidateGmCache) - cmo->removeAttr(kCmoCacheInvalidSkipLoweringAttrName); - dirtyGmCache = false; - pendingInvalidateGmCache = false; - } -}; - -static bool isGmScalarMemory(Type type) { - if (auto ptrTy = dyn_cast(type)) { - pto::AddressSpace space = ptrTy.getMemorySpace().getAddressSpace(); - return isGmAddressSpace(space); - } - - if (auto memTy = dyn_cast(type)) { - auto spaceAttr = dyn_cast_or_null(memTy.getMemorySpace()); - return !spaceAttr || isGmAddressSpace(spaceAttr.getAddressSpace()); - } - - return false; -} - -static TNotifyReleaseState getMte2PayloadReadReleaseState(Value payload) { - TNotifyReleaseState state; - state.addAccess(pto::PIPE::PIPE_MTE2, payload, /*needsDsb=*/false, - /*needsCmo=*/false); - return state; -} - -static TNotifyReleaseState getMte3GmWriteReleaseState(Value payload) { - TNotifyReleaseState state; - state.addAccess(pto::PIPE::PIPE_MTE3, payload, /*needsDsb=*/true, - /*needsCmo=*/false); - return state; -} - -static TNotifyReleaseState getFixGmWriteReleaseState(Value payload) { - TNotifyReleaseState state; - state.addAccess(pto::PIPE::PIPE_FIX, payload, /*needsDsb=*/true, - /*needsCmo=*/false); - return state; -} - -static TNotifyReleaseState getCacheableGmStoreReleaseState(Value payload) { - TNotifyReleaseState state; - state.addAccess(pto::PIPE::PIPE_UNASSIGNED, payload, /*needsDsb=*/true, - /*needsCmo=*/true); - return state; -} - -static TNotifyReleaseState getReleaseStateForMacroModel(Operation *op) { - TNotifyReleaseState state; - auto model = getSyncMacroModel(op); - if (!model) - return state; - - for (const SyncMacroPhase &phase : model->phases) { - // Macro MTE3 phases write GM payloads internally. A following TNotify must - // publish its signal only after those stores are drained and DDR-visible. - if (phase.pipe == PipelineType::PIPE_MTE3) { - if (phase.defValues.empty()) { - state.merge(getMte3GmWriteReleaseState(Value{})); - } else { - for (Value value : phase.defValues) - state.merge(getMte3GmWriteReleaseState(value)); - } - } - } - return state; -} - -static TNotifyReleaseState getDirectTNotifyReleaseState(Operation *op) { - if (isa(op)) - return {}; - - if (auto store = dyn_cast(op)) { - if (isGmScalarMemory(store.getPtr().getType())) - return getCacheableGmStoreReleaseState(store.getPtr()); - } - - if (auto tload = dyn_cast(op)) - return getMte2PayloadReadReleaseState(tload.getSrc()); - - if (auto prefetch = dyn_cast(op)) - return getMte2PayloadReadReleaseState(prefetch.getSrc()); - - if (auto tstore = dyn_cast(op)) { - if (tstore.getPipe() == pto::PIPE::PIPE_MTE3) - return getMte3GmWriteReleaseState(tstore.getDst()); - if (tstore.getPipe() == pto::PIPE::PIPE_FIX) - return getFixGmWriteReleaseState(tstore.getDst()); - return {}; - } - - if (isa(op)) - return getFixGmWriteReleaseState(cast(op).getDst()); - - TNotifyReleaseState macroState = getReleaseStateForMacroModel(op); - if (macroState.drainMte3 || macroState.drainFix || - macroState.needsDsbDdr || macroState.needsGmCacheCmo) - return macroState; - - return {}; -} - -static TNotifyReleaseState collectTNotifyReleaseState(Operation *op) { - TNotifyReleaseState state = getDirectTNotifyReleaseState(op); - for (Region ®ion : op->getRegions()) - for (Block &block : region) - for (Operation &nested : block) - state.merge(collectTNotifyReleaseState(&nested)); - return state; -} - -static TNotifyReleaseState collectTNotifyReleaseState(Region ®ion) { - TNotifyReleaseState state; - for (Block &block : region) - for (Operation &nested : block) - state.merge(collectTNotifyReleaseState(&nested)); - return state; -} - -static void applyFenceBarrierAllForSummary(pto::FenceBarrierAllOp fence, - TNotifyReleaseState &state) { - if (fence.getScope().getScope() != pto::FenceScope::GM && - fence.getScope().getScope() != pto::FenceScope::All) - return; - - // The real annotation pass inserts the pending GM-write pipe drain before a - // barrier_all. Loop summaries must model that transfer without mutating IR, - // otherwise already-released loop-carried writes are reported again at the - // next iteration's TNotify. - state.markPipeDrained(pto::PIPE::PIPE_MTE3); - state.markPipeDrained(pto::PIPE::PIPE_FIX); - state.applyFenceBarrierAll(fence.getScope().getScope()); -} - -static TNotifyReleaseState getTNotifyReleaseExitStateForBlock( - Block &block, TNotifyReleaseState pendingState); - -static TNotifyReleaseState -getTNotifyReleaseExitState(Operation *op, - TNotifyReleaseState pendingState = {}) { - if (isa(op)) - pendingState.clear(); - - pendingState.merge(getDirectTNotifyReleaseState(op)); - - TNotifyReleaseState regionEntryState = pendingState; - TNotifyReleaseState combinedRegionExitState; - for (Region ®ion : op->getRegions()) { - if (region.hasOneBlock()) { - combinedRegionExitState.merge( - getTNotifyReleaseExitStateForBlock(region.front(), regionEntryState)); - continue; - } - - TNotifyReleaseState regionExitState = regionEntryState; - regionExitState.merge(collectTNotifyReleaseState(region)); - combinedRegionExitState.merge(regionExitState); - } - pendingState.merge(combinedRegionExitState); - - if (auto barrier = dyn_cast(op)) - pendingState.applyBarrier(barrier.getPipe().getPipe()); - if (auto cmo = dyn_cast(op)) - pendingState.applyCmoCacheInvalid(cmo); - if (auto fence = dyn_cast(op)) - applyFenceBarrierAllForSummary(fence, pendingState); - return pendingState; -} - -static TNotifyReleaseState getTNotifyReleaseExitStateForBlock( - Block &block, TNotifyReleaseState pendingState) { - for (Operation &op : block) - pendingState = getTNotifyReleaseExitState(&op, pendingState); - return pendingState; -} - -static bool isLoopLikeOp(Operation *op) { - return isa(op); -} - -static func::FuncOp lookupCallee(func::CallOp call) { - return SymbolTable::lookupNearestSymbolFrom( - call.getOperation(), call.getCalleeAttr()); -} - -static bool isMemoryConsistencyRelevantDirectOp(Operation *op) { - if (isa(op)) - return true; - - if (auto load = dyn_cast(op)) - return isGmScalarMemory(load.getPtr().getType()); - if (auto store = dyn_cast(op)) - return isGmScalarMemory(store.getPtr().getType()); - - TNotifyReleaseState macroState = getReleaseStateForMacroModel(op); - return macroState.drainMte2 || macroState.drainMte3 || - macroState.drainFix || macroState.needsDsbDdr || - macroState.needsGmCacheCmo; -} - -static bool calleeContainsMemoryConsistencyRelevantOps( - func::FuncOp callee, llvm::DenseSet &activeCallees) { - if (!callee || callee.isExternal()) - return false; - if (!activeCallees.insert(callee.getOperation()).second) - return false; - - WalkResult result = callee.walk([&](Operation *op) -> WalkResult { - if (op == callee.getOperation()) - return WalkResult::advance(); - - if (auto nestedCall = dyn_cast(op)) { - func::FuncOp nestedCallee = lookupCallee(nestedCall); - if (calleeContainsMemoryConsistencyRelevantOps(nestedCallee, - activeCallees)) - return WalkResult::interrupt(); - return WalkResult::advance(); - } - - if (isMemoryConsistencyRelevantDirectOp(op)) - return WalkResult::interrupt(); - return WalkResult::advance(); - }); - - activeCallees.erase(callee.getOperation()); - return result.wasInterrupted(); -} - -static bool diagnoseNonInlinedMemoryConsistencyCalls(ModuleOp module) { - bool hasFailure = false; - for (auto func : module.getOps()) { - if (func.isExternal()) - continue; - // Entry wrappers are launch/orchestration functions in EmitC tests. They - // can call several kernel functions that are analyzed independently by this - // module pass, so rejecting those calls would incorrectly forbid the normal - // multi-kernel entry shape. The unsafe case is a non-inlined call from - // inside an actual kernel body, where the caller-side signal/fence cannot - // see payload actions hidden in the callee. - if (func->hasAttr("pto.entry")) - continue; - - func.walk([&](func::CallOp call) { - func::FuncOp callee = lookupCallee(call); - if (!callee || callee.isExternal()) - return; - - llvm::DenseSet activeCallees; - if (!calleeContainsMemoryConsistencyRelevantOps(callee, activeCallees)) - return; - - call.emitOpError() - << "calls @" << callee.getSymName() - << ", which contains PTO memory consistency relevant operations; " - "inline the callee before `pto-memory-consistency` or keep " - "payload, CMO, fence, and signal operations in the caller"; - hasFailure = true; - }); - } - return hasFailure; -} - -static void setTNotifyReleaseAttrs(pto::TNotifyOp op, - const TNotifyReleaseState &state) { - op->removeAttr(kTNotifyDrainMte2AttrName); - op->removeAttr(kTNotifyDrainMte3AttrName); - if (state.drainMte2) - op->setAttr(kTNotifyDrainMte2AttrName, UnitAttr::get(op.getContext())); - if (state.drainMte3) - op->setAttr(kTNotifyDrainMte3AttrName, UnitAttr::get(op.getContext())); -} - -static void setTNotifyPipeDrainAttrs(pto::TNotifyOp op, - const TNotifyReleaseState &state) { - TNotifyReleaseState emitState; - emitState.drainMte2 = - state.hasReleaseCmoMarker() && state.addressedDrainMte2; - emitState.drainMte3 = - state.hasReleaseCmoMarker() && state.addressedDrainMte3; - setTNotifyReleaseAttrs(op, emitState); -} - -static void diagnoseTNotifyRelease(pto::TNotifyOp op, - const TNotifyReleaseState &state, - bool &hasFailure) { - if (!state.hasReleaseCmoMarker()) - return; - - if (state.addressedNeedsGmCacheCmo) { - op.emitOpError() - << "requires explicit `pto.cmo.cacheinvalid %addr " - "single_cache_line` or `pto.cmo.cacheinvalid all " - "#pto.address_space` after matching cacheable GM scalar stores " - "and before `pto.fence.barrier_all #pto.fence_scope`"; - hasFailure = true; - return; - } - if (state.addressedNeedsDsbDdr) { - op.emitOpError() - << "requires explicit `pto.fence.barrier_all #pto.fence_scope` " - "after the matching `pto.cmo.cacheinvalid` release marker and " - "before publishing the signal"; - hasFailure = true; - } -} - -static void insertDrainsBeforeBarrierAll(pto::FenceBarrierAllOp fence, - TNotifyReleaseState &state) { - if (fence.getScope().getScope() != pto::FenceScope::GM && - fence.getScope().getScope() != pto::FenceScope::All) - return; - OpBuilder builder(fence); - auto insertBarrier = [&](pto::PIPE pipe) { - builder.create( - fence.getLoc(), pto::PipeAttr::get(fence.getContext(), pipe)); - }; - const bool drainMte3 = - state.hasReleaseCmoMarker() && state.addressedDrainMte3; - const bool drainFix = - state.hasReleaseCmoMarker() && state.addressedDrainFix; - if (drainMte3) { - insertBarrier(pto::PIPE::PIPE_MTE3); - state.markPipeDrained(pto::PIPE::PIPE_MTE3); - } - if (drainFix) { - insertBarrier(pto::PIPE::PIPE_FIX); - state.markPipeDrained(pto::PIPE::PIPE_FIX); - } -} - -static void markNestedTNotifyWithState(Operation *op, - const TNotifyReleaseState &state, - bool &hasFailure) { - op->walk([&](pto::TNotifyOp notify) { - diagnoseTNotifyRelease(notify, state, hasFailure); - setTNotifyPipeDrainAttrs(notify, state); - }); -} - -static void markNestedTNotifyWithState(Region ®ion, - const TNotifyReleaseState &state, - bool &hasFailure) { - for (Block &block : region) { - for (Operation &nested : block) - markNestedTNotifyWithState(&nested, state, hasFailure); - } -} - -static TNotifyReleaseState -annotateTNotifyReleaseForBlock(Block &block, - TNotifyReleaseState entryPendingState, - TNotifyReleaseState loopCarriedState, - bool &hasFailure) { - TNotifyReleaseState pendingState = entryPendingState; - for (Operation &op : block) { - if (auto notify = dyn_cast(op)) { - TNotifyReleaseState notifyState = pendingState; - notifyState.merge(loopCarriedState); - diagnoseTNotifyRelease(notify, notifyState, hasFailure); - setTNotifyPipeDrainAttrs(notify, notifyState); - pendingState.clear(); - } - - pendingState.merge(getDirectTNotifyReleaseState(&op)); - - TNotifyReleaseState regionEntryState = pendingState; - TNotifyReleaseState combinedRegionExitState; - for (Region ®ion : op.getRegions()) { - TNotifyReleaseState nestedLoopCarriedState = loopCarriedState; - if (isLoopLikeOp(&op)) - nestedLoopCarriedState.merge(getTNotifyReleaseExitState(&op)); - - if (region.hasOneBlock()) { - combinedRegionExitState.merge(annotateTNotifyReleaseForBlock( - region.front(), regionEntryState, nestedLoopCarriedState, - hasFailure)); - } else { - TNotifyReleaseState regionState = collectTNotifyReleaseState(region); - TNotifyReleaseState nestedNotifyState = regionEntryState; - nestedNotifyState.merge(nestedLoopCarriedState); - nestedNotifyState.merge(regionState); - markNestedTNotifyWithState(region, nestedNotifyState, hasFailure); - - TNotifyReleaseState regionExitState = regionEntryState; - regionExitState.merge(regionState); - combinedRegionExitState.merge(regionExitState); - } - } - pendingState.merge(combinedRegionExitState); - - if (auto barrier = dyn_cast(op)) - pendingState.applyBarrier(barrier.getPipe().getPipe()); - if (auto cmo = dyn_cast(op)) - pendingState.applyCmoCacheInvalid(cmo); - if (auto fence = dyn_cast(op)) { - insertDrainsBeforeBarrierAll(fence, pendingState); - pendingState.applyFenceBarrierAll(fence.getScope().getScope()); - } - } - return pendingState; -} - -static bool annotateTNotifyRelease(ModuleOp module) { - bool hasFailure = false; - for (auto func : module.getOps()) { - if (func.isExternal()) - continue; - - if (func.getBody().hasOneBlock()) { - (void)annotateTNotifyReleaseForBlock(func.getBody().front(), - TNotifyReleaseState{}, - TNotifyReleaseState{}, - hasFailure); - continue; - } - - // Be conservative for pre-existing CFG: without a path-sensitive CFG data - // flow here, every TNotify may observe any release-relevant work in the - // function. - TNotifyReleaseState funcState = collectTNotifyReleaseState(func.getBody()); - markNestedTNotifyWithState(func.getBody(), funcState, hasFailure); - } - return hasFailure; -} - -static void clearAcquireAttrs(pto::LoadScalarOp op) { - op->removeAttr(kAcquireInvalidateGmCacheAttrName); -} - -static void diagnoseAcquireLoad(pto::LoadScalarOp op, - const SignalAcquireState &state, - bool &hasFailure) { - if (!state.pendingInvalidateGmCache || - !isGmScalarMemory(op.getPtr().getType())) - return; - if (state.dirtyGmCache) { - op.emitOpError() - << "cannot perform a cacheable GM load after signal acquire while " - "dirty GM cache may exist; insert explicit " - "`pto.cmo.cacheinvalid all #pto.address_space` before the load"; - hasFailure = true; - return; - } - op.emitOpError() - << "requires explicit `pto.cmo.cacheinvalid all #pto.address_space` " - "before a cacheable GM load after `pto.comm.twait` or successful " - "`pto.comm.ttest`"; - hasFailure = true; -} - -static void consumeAcquireAfterDiagnostic(SignalAcquireState &state) { - if (state.pendingInvalidateGmCache) - state.consumeAcquire(); -} - -static SignalAcquireState collectSignalAcquireState(Operation *op) { - SignalAcquireState state; - if (isa(op)) - state.pendingInvalidateGmCache = true; - if (auto store = dyn_cast(op); - store && isGmScalarMemory(store.getPtr().getType())) - state.dirtyGmCache = true; - if (auto cmo = dyn_cast(op)) - state.applyCmoCacheInvalid(cmo.getSpace().getAddressSpace()); - - for (Region ®ion : op->getRegions()) - for (Block &block : region) - for (Operation &nested : block) - state.merge(collectSignalAcquireState(&nested)); - return state; -} - -static SignalAcquireState collectSignalAcquireState(Region ®ion) { - SignalAcquireState state; - for (Block &block : region) - for (Operation &nested : block) - state.merge(collectSignalAcquireState(&nested)); - return state; -} - -static void markNestedAcquireLoadsWithState(Operation *op, - SignalAcquireState state, - bool &hasFailure) { - op->walk([&](pto::LoadScalarOp load) { - clearAcquireAttrs(load); - diagnoseAcquireLoad(load, state, hasFailure); - consumeAcquireAfterDiagnostic(state); - }); -} - -static void markNestedAcquireLoadsWithState(Region ®ion, - SignalAcquireState state, - bool &hasFailure) { - for (Block &block : region) { - for (Operation &nested : block) - markNestedAcquireLoadsWithState(&nested, state, hasFailure); - } -} - -static SignalAcquireState -annotateSignalAcquireForBlock(Block &block, SignalAcquireState entryState, - bool &hasFailure) { - SignalAcquireState state = entryState; - for (Operation &op : block) { - if (auto load = dyn_cast(op)) { - clearAcquireAttrs(load); - diagnoseAcquireLoad(load, state, hasFailure); - consumeAcquireAfterDiagnostic(state); - } - - if (auto store = dyn_cast(op); - store && isGmScalarMemory(store.getPtr().getType())) - state.dirtyGmCache = true; - - if (isa(op)) - state.pendingInvalidateGmCache = true; - - if (auto cmo = dyn_cast(op)) - state.applyCmoCacheInvalid(cmo.getSpace().getAddressSpace(), cmo); - - SignalAcquireState combinedRegionExitState; - for (Region ®ion : op.getRegions()) { - if (region.hasOneBlock()) { - combinedRegionExitState.merge( - annotateSignalAcquireForBlock(region.front(), state, hasFailure)); - } else { - markNestedAcquireLoadsWithState(region, state, hasFailure); - SignalAcquireState regionState = collectSignalAcquireState(region); - SignalAcquireState regionExitState = state; - regionExitState.merge(regionState); - combinedRegionExitState.merge(regionExitState); - } - } - - if (isLoopLikeOp(&op)) - combinedRegionExitState.merge(state); - state.merge(combinedRegionExitState); - } - return state; -} - -static bool annotateSignalAcquire(ModuleOp module) { - bool hasFailure = false; - for (auto func : module.getOps()) { - if (func.isExternal()) - continue; - - if (func.getBody().hasOneBlock()) { - (void)annotateSignalAcquireForBlock(func.getBody().front(), - SignalAcquireState{}, hasFailure); - continue; - } - - SignalAcquireState funcState = collectSignalAcquireState(func.getBody()); - markNestedAcquireLoadsWithState(func.getBody(), funcState, hasFailure); - } - return hasFailure; -} - -static void clearMemoryConsistencyInternalAttrs(ModuleOp module) { - module.walk([](pto::CmoCacheInvalidOp cmo) { - cmo->removeAttr(kCmoCacheInvalidSkipLoweringAttrName); - }); -} - -struct PTOMemoryConsistencyPass - : public mlir::pto::impl::PTOMemoryConsistencyBase< - PTOMemoryConsistencyPass> { - void runOnOperation() override { - ModuleOp module = getOperation(); - clearMemoryConsistencyInternalAttrs(module); - bool callFailed = diagnoseNonInlinedMemoryConsistencyCalls(module); - bool releaseFailed = annotateTNotifyRelease(module); - bool acquireFailed = annotateSignalAcquire(module); - if (callFailed || releaseFailed || acquireFailed) - signalPassFailure(); - } -}; - -} // namespace - -std::unique_ptr mlir::pto::createPTOMemoryConsistencyPass() { - return std::make_unique(); -} diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 2aa93a453a..8f94813bb2 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -5616,6 +5616,21 @@ static void emitDsbDdr(ConversionPatternRewriter &rewriter, Location loc) { ArrayAttr{}, ValueRange{}); } +static void emitPipeBarrier(ConversionPatternRewriter &rewriter, Location loc, + StringRef pipeTok) { + auto *ctx = rewriter.getContext(); + auto args = rewriter.getArrayAttr({emitc::OpaqueAttr::get(ctx, pipeTok)}); + rewriter.create(loc, TypeRange{}, "pipe_barrier", args, + ArrayAttr{}, ValueRange{}); +} + +static void emitConservativeGmFencePipeDrains( + ConversionPatternRewriter &rewriter, Location loc) { + emitPipeBarrier(rewriter, loc, "PIPE_MTE2"); + emitPipeBarrier(rewriter, loc, "PIPE_MTE3"); + emitPipeBarrier(rewriter, loc, "PIPE_FIX"); +} + struct PTOBarrierToEmitC : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -5668,6 +5683,7 @@ struct PTOFenceToEmitC : public OpConversionPattern { op.getScope().getScope() != pto::FenceScope::All) return rewriter.notifyMatchFailure(op, "unsupported fence scope"); + emitConservativeGmFencePipeDrains(rewriter, op.getLoc()); emitDsbDdr(rewriter, op.getLoc()); rewriter.eraseOp(op); return success(); @@ -6282,12 +6298,9 @@ struct PTOSyncSetToEmitC : public OpConversionPattern { return rewriter.notifyMatchFailure( op, "expects exactly one of static event_id attr or dynamic event_id operand"); - // A5 inter-core sync mirrors +16 only for cube-side producer (PIPE_FIX). - // Vec-side producer (PIPE_MTE3) emits a single set; hardware handles the - // subblock mapping in PTO-ISA custom flow. + // A5 sync uses physical semaphore IDs. Emit exactly the ID authored in IR; + // callers that need to notify both AIV subblocks must emit both IDs. if (targetArch == PTOArch::A5) { - pto::PIPE pipe = op.getPipe().getPipe(); - bool needsMirrorPlus16 = (pipe == pto::PIPE::PIPE_FIX); std::string pipeTok = pipeTokFromPipeAttr(op.getPipe()); auto emitSet = [&](Value eventOperand, IntegerAttr eventLiteral, bool isDynamic) { @@ -6314,21 +6327,9 @@ struct PTOSyncSetToEmitC : public OpConversionPattern { if (eventIdAttr) { emitSet(Value{}, eventIdAttr, /*isDynamic=*/false); - if (needsMirrorPlus16) { - auto plus16 = IntegerAttr::get(eventIdAttr.getType(), - getIntegerAttrSignedValue(eventIdAttr) + 16); - emitSet(Value{}, plus16, /*isDynamic=*/false); - } } else { Value eventI32 = castInterCoreEventIdToI32(rewriter, loc, eventIdDyn); emitSet(eventI32, IntegerAttr{}, /*isDynamic=*/true); - if (needsMirrorPlus16) { - auto i32Ty = emitc::OpaqueType::get(ctx, "int32_t"); - Value c16 = makeEmitCIntConstant(rewriter, loc, i32Ty, 16); - Value eventI32Plus16 = - rewriter.create(loc, i32Ty, eventI32, c16).getResult(); - emitSet(eventI32Plus16, IntegerAttr{}, /*isDynamic=*/true); - } } rewriter.eraseOp(op); @@ -7425,21 +7426,10 @@ static std::string notifyOpTok(pto::NotifyOp op) { return "pto::comm::NotifyOp::Set"; } -static void emitPipeBarrier(ConversionPatternRewriter &rewriter, Location loc, - StringRef pipeTok) { - auto *ctx = rewriter.getContext(); - auto args = rewriter.getArrayAttr({emitc::OpaqueAttr::get(ctx, pipeTok)}); - rewriter.create(loc, TypeRange{}, "pipe_barrier", args, - ArrayAttr{}, ValueRange{}); -} - -// Issue #711: TNOTIFY writes its signal on the scalar pipe, and -// TNOTIFY_IMPL's trailing pipe_barrier(PIPE_ALL) runs *after* that store. -// If prior MTE work is still in flight when the signal lands, the receiver's -// matching TWAIT can return before the producer-side payload operation is -// complete. MemoryConsistency now validates explicit CMO/fence operations for -// DDR visibility; lowering only keeps the pipe-drain actions that the pass may -// still annotate automatically. +// Historical hook for pre-annotated TNotify release drains. The automatic +// MemoryConsistency analysis pass that used to produce these attrs has been +// removed from the default pipeline; keeping the lowering hook is harmless for +// hand-authored or legacy IR that already carries the internal attrs. static void emitTNotifyReleaseActions(ConversionPatternRewriter &rewriter, Location loc, bool drainMte2, bool drainMte3) { diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index 5bb883eca5..f429f5470e 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -1574,8 +1574,8 @@ packCopyUbToGmConfig0(Operation *anchor, ValueRange operands) { Value sid = getI64Operand(2); Value nBurst = getI64Operand(3); Value lenBurst = getI64Operand(4); - Value reserved = getI64Operand(5); - if (!sid || !nBurst || !lenBurst || !reserved) + Value l2CacheCtl = getI64Operand(5); + if (!sid || !nBurst || !lenBurst || !l2CacheCtl) return failure(); auto shl = [&](Value value, uint64_t amount) -> Value { @@ -1589,7 +1589,7 @@ packCopyUbToGmConfig0(Operation *anchor, ValueRange operands) { Value config = sid; config = bitOr(config, shl(nBurst, 4)); config = bitOr(config, shl(lenBurst, 25)); - config = bitOr(config, shl(reserved, 60)); + config = bitOr(config, shl(l2CacheCtl, 60)); return config; } @@ -1602,12 +1602,12 @@ packCopyUbToGmConfig1(Operation *anchor, ValueRange operands) { [[maybe_unused]] static FailureOr packCopyUbToGmConfig0(Operation *anchor, Value sid, Value nBurst, - Value lenBurst, Value reserved) { + Value lenBurst, Value l2CacheCtl) { SmallVector operands(8); operands[2] = sid; operands[3] = nBurst; operands[4] = lenBurst; - operands[5] = reserved; + operands[5] = l2CacheCtl; return packCopyUbToGmConfig0(anchor, operands); } diff --git a/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp b/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp index 799cbc1b0b..d927915db5 100644 --- a/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp +++ b/lib/PTO/Transforms/VPTOExpandWrapperOps.cpp @@ -368,7 +368,9 @@ static void configureAccStoreScalarPreOps(Location loc, Value preQuant, } }; - if (preQuantMode && !isVectorQuantMode(*preQuantMode)) { + if (preQuantMode && + *preQuantMode != pto::AccStoreQuantPreMode::NoConvert && + !isVectorQuantMode(*preQuantMode)) { if (Value quantValue = materializeAccStoreScalarPayload(preQuant, rewriter, loc)) rewriter.create(loc, quantValue); } @@ -1196,9 +1198,10 @@ struct ExpandDmaStorePattern : public OpRewritePattern { Value source = offsetPointerByBytes(op.getSource(), srcOffset, rewriter, loc); Value destination = offsetPointerByBytes(op.getDestination(), dstOffset, rewriter, loc); + Value l2CacheCtl = op.getL2CacheCtl() ? op.getL2CacheCtl() : zero; rewriter.create( loc, source, destination, zero, op.getNBurst(), op.getLenBurst(), - zero, op.getNburstDstStride(), op.getNburstSrcStride()); + l2CacheCtl, op.getNburstDstStride(), op.getNburstSrcStride()); }); if (shouldRestoreDmaLoopSize(loop1Count, loop2Size)) rewriter.create(loc, one, one); diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 7951983ef5..ec8ac559de 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -1542,8 +1542,8 @@ packCopyUbToGmConfig0(Operation *anchor, ValueRange operands) { Value sid = getI64Operand(2); Value nBurst = getI64Operand(3); Value lenBurst = getI64Operand(4); - Value reserved = getI64Operand(5); - if (!sid || !nBurst || !lenBurst || !reserved) + Value l2CacheCtl = getI64Operand(5); + if (!sid || !nBurst || !lenBurst || !l2CacheCtl) return failure(); auto shl = [&](Value value, uint64_t amount) -> Value { @@ -1557,7 +1557,7 @@ packCopyUbToGmConfig0(Operation *anchor, ValueRange operands) { Value config = sid; config = bitOr(config, shl(nBurst, 4)); config = bitOr(config, shl(lenBurst, 25)); - config = bitOr(config, shl(reserved, 60)); + config = bitOr(config, shl(l2CacheCtl, 60)); return config; } @@ -1570,12 +1570,12 @@ packCopyUbToGmConfig1(Operation *anchor, ValueRange operands) { [[maybe_unused]] static FailureOr packCopyUbToGmConfig0(Operation *anchor, Value sid, Value nBurst, - Value lenBurst, Value reserved) { + Value lenBurst, Value l2CacheCtl) { SmallVector operands(8); operands[2] = sid; operands[3] = nBurst; operands[4] = lenBurst; - operands[5] = reserved; + operands[5] = l2CacheCtl; return packCopyUbToGmConfig0(anchor, operands); } diff --git a/lib/PTO/Transforms/VPTOLLVMEmitterHelper.cpp b/lib/PTO/Transforms/VPTOLLVMEmitterHelper.cpp index 8b5eb7a165..d4d4e0acba 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitterHelper.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitterHelper.cpp @@ -534,10 +534,19 @@ LogicalResult attachAIVectorScopeMetadata(llvm::Module &llvmModule, return success(); } +constexpr uint32_t getSimtMaxRegistersForThreads(uint32_t maxThreads) { + if (maxThreads > 1024) + return 16; + if (maxThreads > 512) + return 32; + if (maxThreads > 256) + return 64; + return 128; +} + void attachHIVMKernelAnnotations(llvm::Module &llvmModule, ModuleOp sourceModule) { constexpr uint32_t kDefaultSimtMaxThreads = 1024; - constexpr uint32_t kDefaultSimtMaxRegisters = 32; llvm::NamedMDNode *annotations = llvmModule.getOrInsertNamedMetadata("hivm.annotations"); @@ -545,7 +554,7 @@ void attachHIVMKernelAnnotations(llvm::Module &llvmModule, llvm::Type *i32Ty = llvm::Type::getInt32Ty(ctx); llvm::Constant *one = llvm::ConstantInt::get(i32Ty, 1); - llvm::StringMap> simtConfigByName; + llvm::StringMap simtMaxThreadsByName; llvm::StringSet ptoEntryFunctions; sourceModule.walk([&](LLVM::LLVMFuncOp funcOp) { @@ -558,15 +567,11 @@ void attachHIVMKernelAnnotations(llvm::Module &llvmModule, return; uint32_t maxThreads = kDefaultSimtMaxThreads; - uint32_t maxRegisters = kDefaultSimtMaxRegisters; if (auto attr = funcOp->getAttrOfType(pto::kPTOSimtMaxThreadsAttrName)) maxThreads = static_cast(attr.getInt()); - if (auto attr = funcOp->getAttrOfType( - pto::kPTOSimtMaxRegistersAttrName)) - maxRegisters = static_cast(attr.getInt()); - simtConfigByName[symName] = {maxThreads, maxRegisters}; + simtMaxThreadsByName[symName] = maxThreads; }); auto callsSimtEntry = [](llvm::Function &function) { @@ -611,12 +616,10 @@ void attachHIVMKernelAnnotations(llvm::Module &llvmModule, continue; if (function.getCallingConv() == llvm::CallingConv::SimtEntry) { uint32_t maxThreads = kDefaultSimtMaxThreads; - uint32_t maxRegisters = kDefaultSimtMaxRegisters; - if (auto it = simtConfigByName.find(function.getName()); - it != simtConfigByName.end()) { - maxThreads = it->second.first; - maxRegisters = it->second.second; - } + if (auto it = simtMaxThreadsByName.find(function.getName()); + it != simtMaxThreadsByName.end()) + maxThreads = it->second; + uint32_t maxRegisters = getSimtMaxRegistersForThreads(maxThreads); addLLVMFunctionI32Annotation(function, "simt-max-threads", maxThreads); addLLVMFunctionI32Annotation(function, "simt-max-registers", diff --git a/lib/TileOps/tload_template.py b/lib/TileOps/tload_template.py index 9e40bedb38..310c8d243e 100644 --- a/lib/TileOps/tload_template.py +++ b/lib/TileOps/tload_template.py @@ -355,14 +355,18 @@ def _constraint_tload_mat_nd2nz(src, dst) -> bool: return False # ND2NZ: source is in ND (row-major) format where the inner dimension (g4) # corresponds to the tile column count. Disambiguates from DN format where - # g4 corresponds to the tile row count. - if hasattr(src, 'rank') and src.rank == 5: - dst_valid_cols = dst.valid_shape[1] if hasattr(dst, 'valid_shape') and dst.valid_shape is not None else None - if dst_valid_cols is not None and hasattr(src, 'shape') and src.shape is not None: - src_inner = src.shape[4] if len(src.shape) >= 5 else None - if src_inner is not None: - if not _known_eq(dst_valid_cols, src_inner): - return False + # g4 corresponds to the tile row count. Use the *static* dst.shape + # (rows/cols from the tile_buf type) rather than dst.valid_shape, which + # may be dynamic ([null, null]) and would otherwise let both ND2NZ and + # DN2NZ match when the shape comparison is skipped. + if not (hasattr(src, 'rank') and src.rank == 5): + return False + dst_cols = dst.shape[1] if hasattr(dst, 'shape') and dst.shape is not None else None + if dst_cols is not None and hasattr(src, 'shape') and src.shape is not None: + src_inner = src.shape[4] if len(src.shape) >= 5 else None + if src_inner is not None: + if not _known_eq(dst_cols, src_inner): + return False return True @@ -383,14 +387,17 @@ def _constraint_tload_mat_dn2nz(src, dst) -> bool: return False # DN2NZ: source is in DN (col-major) format where the inner dimension (g4) # corresponds to the tile row count. Disambiguates from ND format where - # g4 corresponds to the tile column count. - if hasattr(src, 'rank') and src.rank == 5: - dst_valid_rows = dst.valid_shape[0] if hasattr(dst, 'valid_shape') and dst.valid_shape is not None else None - if dst_valid_rows is not None and hasattr(src, 'shape') and src.shape is not None: - src_inner = src.shape[4] if len(src.shape) >= 5 else None - if src_inner is not None: - if not _known_eq(dst_valid_rows, src_inner): - return False + # g4 corresponds to the tile column count. Use the *static* dst.shape + # (rows/cols from the tile_buf type) rather than dst.valid_shape, which + # may be dynamic and would otherwise let both templates match. + if not (hasattr(src, 'rank') and src.rank == 5): + return False + dst_rows = dst.shape[0] if hasattr(dst, 'shape') and dst.shape is not None else None + if dst_rows is not None and hasattr(src, 'shape') and src.shape is not None: + src_inner = src.shape[4] if len(src.shape) >= 5 else None + if src_inner is not None: + if not _known_eq(dst_rows, src_inner): + return False return True @@ -423,13 +430,23 @@ def template_tload_gm_to_mat_nd2nz(src: pto.PartitionTensorView, dst: pto.Tile): gm_ptr = src.as_ptr() mat_ptr = dst.as_ptr() + # Element byte width (src and dst share dtype per the template dtypes table). + elem_bytes = pto.bytewidth(dst.element_type) + + # rank-5 partition view metadata (g3 = M rows, g4 = K cols for ND source). + g0, g1, g2, g3, g4 = src.shape + s0, s1, s2, s3, s4 = src.strides + # ND2NZ parameter calculation # n_value = M (row count), d_value = K (column count) n_value = m d_value = k - # src_layout: inner stride = K (number of elements per row) - src_inner_stride = k + # src_layout: inner stride in BYTES = src physical row stride (s3). + # Hardware loop1_src_stride is in bytes (pto-isa a5/TLoad.hpp: + # GetByteSize(gStride3)); a strided GM slice has s3 != K, so the + # element stride must be scaled by elem_bytes. + src_inner_stride = s3 * elem_bytes # dst_group: (group_count, loop2_stride, loop3_stride, loop4_stride) # For simple single-block case: (1, 1, m, 0) @@ -478,6 +495,13 @@ def template_tload_gm_to_mat_dn2nz(src: pto.PartitionTensorView, dst: pto.Tile): gm_ptr = src.as_ptr() mat_ptr = dst.as_ptr() + # Element byte width (src and dst share dtype per the template dtypes table). + elem_bytes = pto.bytewidth(dst.element_type) + + # rank-5 partition view metadata (g3 = K, g4 = M for DN col-major source). + g0, g1, g2, g3, g4 = src.shape + s0, s1, s2, s3, s4 = src.strides + # DN2NZ parameter calculation # For DN format, the original shape is (K, M) -- no logical conversion # needed. dn2nz writes the same logical N x D result into NZ layout. @@ -485,8 +509,9 @@ def template_tload_gm_to_mat_dn2nz(src: pto.PartitionTensorView, dst: pto.Tile): n_value = k d_value = m - # src_layout: inner stride = M (number of elements per column) - src_inner_stride = m + # src_layout: inner stride in BYTES = src physical stride along the M + # direction (s4). pto-isa TLoadCubeND2NZ uses gStride4 when layout == DN. + src_inner_stride = s4 * elem_bytes # dst_group: (group_count, loop2_stride, loop3_stride, loop4_stride) # (1, 1, k, 0) diff --git a/lib/TileOps/tstore_template.py b/lib/TileOps/tstore_template.py index 59c2ee86cc..ef80400215 100644 --- a/lib/TileOps/tstore_template.py +++ b/lib/TileOps/tstore_template.py @@ -414,10 +414,20 @@ def template_tstore_acc_to_gm_nz2nd(src: pto.Tile, dst: pto.PartitionTensorView) acc_ptr = src.as_ptr() gm_ptr = dst.as_ptr() - # src_stride: ACC buffer stride (N under NZ format) - # dst_stride: GM stride (N under ND format) - src_stride = n - dst_stride = n + # rank-5 partition view metadata (ND: g3 = M rows, g4 = N cols). + g0, g1, g2, g3, g4 = dst.shape + s0, s1, s2, s3, s4 = dst.strides + + # Strides are in ELEMENT units (pto-isa a5/TStore.hpp: TStoreAccND passes + # srcStride/dstD to copy_matrix_cc_to_gm without GetByteSize; the hardware + # scales by dtype internally). The earlier `src_stride = dst_stride = n` + # filled both fields with the tile column count, which is wrong for any + # strided GM slice and produces sparse/shifted GM writes. + # + # src_stride = TileData::Rows = src.shape[0] (NZ row stride of L0C acc) + # dst_stride = gStride3 = s3 (GM physical row stride) + src_stride = src.shape[0] + dst_stride = s3 pto.mte_l0c_gm( acc_ptr, gm_ptr, @@ -459,10 +469,18 @@ def template_tstore_acc_to_gm_nz2dn(src: pto.Tile, dst: pto.PartitionTensorView) acc_ptr = src.as_ptr() gm_ptr = dst.as_ptr() - # NZ2DN requires an additional loop0_src_stride parameter - src_stride = n - dst_stride = m # Under DN format, stride is M - loop0_src_stride = 1 # loop0_src_stride for NZ2DN + # rank-5 partition view metadata (DN: g3 = M, g4 = N, col-major). + g0, g1, g2, g3, g4 = dst.shape + s0, s1, s2, s3, s4 = dst.strides + + # Strides are in ELEMENT units (see template_tstore_acc_to_gm_nz2nd). + # DN (col-major): adjacent rows step along the M direction, whose physical + # stride is s4 (PTOAS 5D DN view strides collapse to [M, M, M, 1, M]). + # src_stride = TileData::Rows = src.shape[0] (NZ row stride of L0C acc) + # dst_stride = s4 (GM physical stride along M) + src_stride = src.shape[0] + dst_stride = s4 + loop0_src_stride = 1 # loop0_src_stride for NZ2DN (NZ source) pto.mte_l0c_gm( acc_ptr, gm_ptr, @@ -503,7 +521,17 @@ def template_tstore_acc_to_gm_nz2nz(src: pto.Tile, dst: pto.PartitionTensorView) acc_ptr = src.as_ptr() gm_ptr = dst.as_ptr() - src_stride = n + # src_stride is in ELEMENT units (pto-isa a5/TStore.hpp: TStoreAccNZ sets + # srcStride = TileData::Rows). The earlier `src_stride = n` filled the L0C + # source stride with the column count, which is wrong whenever N differs + # from the acc tile row count. + # src_stride = TileData::Rows = src.shape[0] (NZ row stride of L0C acc) + # NOTE: dst_stride for NZ2NZ is the NZ-fractal physical row stride + # (ceil(M, FRACTAL_NZ_ROW) * c0 in pto-isa TStoreAccNZ). The compact NZ + # views used today set this equal to n, so the legacy value is retained + # until a strided NZ destination case is exercised to validate the exact + # stride index. + src_stride = src.shape[0] dst_stride = n split = 1 # NZ2NZ requires a split parameter @@ -616,8 +644,13 @@ def template_tstore_fp_acc_to_gm(src: pto.Tile, fp: pto.Tile, dst: pto.Partition # TODO: Replace with tstore_fp DSL surface once available. # Currently using mte_l0c_gm + pre_quant as a temporary workaround. - src_stride = n - dst_stride = n + # Strides are in ELEMENT units (see template_tstore_acc_to_gm_nz2nd). + # src_stride = TileData::Rows = src.shape[0] (NZ row stride of L0C acc) + # dst_stride = gStride3 = s3 (GM physical row stride, ND) + g0, g1, g2, g3, g4 = dst.shape + s0, s1, s2, s3, s4 = dst.strides + src_stride = src.shape[0] + dst_stride = s3 pto.mte_l0c_gm( acc_ptr, gm_ptr, 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 a75f15a680..e65b06aed1 100644 --- a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md +++ b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md @@ -986,7 +986,7 @@ Parameters are `Tile` references, typed UB pointers, and PTO scalars. The sub-kernel reads and writes individual elements through tile handles; results flow back to the caller via mutable tile parameters. -**Signature**: `@pto.simt(fn=None, *, name=None, target="a5", max_threads=None, max_regs=None)` +**Signature**: `@pto.simt(fn=None, *, name=None, target="a5", max_threads=None)` ```python @@ -1026,28 +1026,31 @@ in parallel, making it efficient for per-element operations. #### SIMT resource attributes -Optional `max_threads` and `max_regs` arguments attach VPTO resource attributes +An optional `max_threads` argument attaches a VPTO resource attribute to the generated `pto.simt_entry` helper. -**Signature**: `@pto.simt(fn=None, *, name=None, target="a5", max_threads=None, max_regs=None)` +**Signature**: `@pto.simt(fn=None, *, name=None, target="a5", max_threads=None)` **Parameters**: | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `max_threads` | positive Python `int` | backend default `1024` | Compile-time launch envelope for this SIMT helper | -| `max_regs` | positive Python `int` | backend default `32` | Scalar register budget per work-item | `max_threads` is not the launch size. The actual work-item count comes from the -SIMT launch dimensions. Both arguments must be Python integers known at trace -time, greater than zero, and fit in signless `i32`. They are only valid on +SIMT launch dimensions. The argument must be a Python integer known at trace +time, greater than zero, and fit in signless `i32`. It is only valid on decorated SIMT helper functions, not inline `with pto.simt():` scopes. +The per-workitem register budget (`simt-max-registers`) is automatically +derived from `max_threads` by the VPTO backend according to the hardware +resource partition. + **Example**: ```python -@pto.simt(max_threads=256, max_regs=48) +@pto.simt(max_threads=256) def write_tid(dst: pto.ptr(pto.i32, "gm")): tid = pto.get_tid_x() idx = scalar.index_cast(tid) diff --git a/ptodsl/docs/user_guide/07-data-movement-ops.md b/ptodsl/docs/user_guide/07-data-movement-ops.md index cabe44abdf..de748a0d92 100644 --- a/ptodsl/docs/user_guide/07-data-movement-ops.md +++ b/ptodsl/docs/user_guide/07-data-movement-ops.md @@ -166,7 +166,7 @@ pto.mte_gm_ub(gm_src, ub_dst, 0, 256, ### 7.2.2 UB → GM: `pto.mte_ub_gm` -#### `pto.mte_ub_gm(ub_src: PtrType, gm_dst: PtrType, len_burst: int, *, nburst: tuple[int, int, int], loops: list[tuple[int, int, int]] | None = None) -> None` +#### `pto.mte_ub_gm(ub_src: PtrType, gm_dst: PtrType, len_burst: int, *, nburst: tuple[int, int, int], loops: list[tuple[int, int, int]] | None = None, l2_cache: str = "nmfv") -> None` **Description**: Grouped DMA transfer from Unified Buffer to Global Memory. The MTE reads `len_burst` bytes from each UB row (skipping any padding), writing only valid data to GM. @@ -179,6 +179,7 @@ pto.mte_gm_ub(gm_src, ub_dst, 0, 256, | `len_burst` | `int` | Contiguous bytes transferred per burst row | | `nburst` | `tuple[int, int, int]` | `(n_burst, src_stride, dst_stride)` — innermost burst group (required) | | `loops` | `list[tuple[int, int, int]]` or `None` | Optional outer loop groups, ordered inner to outer | +| `l2_cache` | `str` | Store-side L2 cache policy token. Defaults to `"nmfv"` (raw control `0`); supported tokens are `"nmfv"`, `"nmlv"`, `"nmprs"`, `"nmred"`, `"naci"`, `"napw"`, `"napi"`, `"nared"`, `"wbhfv"`, `"wbhlv"`, `"wbhprs"`, `"wbhred"`, `"wtsfv"`, `"wtslv"`, `"wtsprs"`, and `"wtsred"` | **Returns**: None (side-effect operation). @@ -195,7 +196,8 @@ pto.mte_ub_gm(ub_src_f32, gm_dst_f32, 128, ```python pto.mte_ub_gm(ub_src, gm_dst, 256, - nburst=(64, 256, 1024)) + nburst=(64, 256, 1024), + l2_cache="nmfv") # UB: contiguous rows (256-byte stride). # GM: rows spaced at 1024-byte intervals (full matrix width). ``` @@ -1046,6 +1048,55 @@ pto.mte_l0c_ub(acc, ub, 16, 32, 16, 32, split=pto.SplitMode.N) # split N --- +#### `pre_quant` modes and destination types + +`pre_quant=(payload, mode)` applies a conversion or quantization step before +`pre_relu`, `clip`, saturation, and the final store. The selected mode must +match both the accumulator source element type and the destination element +type. + +Payload rules: + +- `_vec` modes require a scaling pointer payload with `f16`, `bf16`, or `f32` + elements. +- `_scalar` modes require an `f16`, `bf16`, or `f32` scalar payload. +- `f32_f16` and `f32_bf16` also require an `f16`, `bf16`, or `f32` scalar + payload; the payload is required by the public syntax but does not select + per-channel scaling. + +Legal mode families: + +| Mode family | Accumulator source dtype | Legal destination dtype | +|-------------|--------------------------|--------------------------| +| `f32_f16`, `qf322f16_pre_*`, `deqf16_*` | `f32` for `f32_f16` / `qf322f16_pre_*`; `i32` for `deqf16_*` | `f16` | +| `f32_bf16`, `qf322bf16_pre_*`, `qs322bf16_pre_*` | `f32` for `f32_bf16` / `qf322bf16_pre_*`; `i32` for `qs322bf16_pre_*` | `bf16` | +| `qf322f32_pre_*` | `f32` | `f32` | +| `qf322hif8_pre_*`, `qf322fp8_pre_*` | `f32` | PTO FP8 / HiFloat8 family | +| `deqs32_int_*` | `i32` | 32-bit integer | +| `qf162s16_pre_*`, `deqs16_*` | `i32` | signed or signless 16-bit integer | +| `qf162b8_pre_*`, `req8_*`, `qf322b8_pre_*` | `f32` for `qf162b8_pre_*` / `qf322b8_pre_*`; `i32` for `req8_*` | 8-bit integer | +| `qf162s4_pre_*`, `req4_*`, `qf322s4_pre_*` | `f32` for `qf162s4_pre_*` / `qf322s4_pre_*`; `i32` for `req4_*` | signed or signless 4-bit integer | + +Examples: + +```python +# Legal: f32 accumulator values are converted to bf16 on writeback. +pto.mte_l0c_gm( + acc_f32, dst_bf16, 16, 16, 16, 16, 0, 0, + pre_quant=(pto.bf16(1.0), "f32_bf16"), + layout="nz2nd", +) + +# Illegal: "f32_bf16" requires a bf16 destination, not f16. +pto.mte_l0c_gm( + acc_f32, dst_f16, 16, 16, 16, 16, 0, 0, + pre_quant=(pto.bf16(1.0), "f32_bf16"), + layout="nz2nd", +) +``` + +--- + ### Cube data movement quick reference | Data Flow | Operation | Src Space | Dst Space | diff --git a/ptodsl/docs/user_guide/10-sync-ops.md b/ptodsl/docs/user_guide/10-sync-ops.md index 29fb82dc08..aa68413ff7 100644 --- a/ptodsl/docs/user_guide/10-sync-ops.md +++ b/ptodsl/docs/user_guide/10-sync-ops.md @@ -405,6 +405,9 @@ pto.wait_cross_flag(pto.Pipe.FIX, 0) The Cube unit (matrix pipeline) has a dedicated synchronization channel separate from the standard pipe-flag mechanism used by MTE and Vector pipelines. `set_intra_flag` and `wait_intra_flag` synchronize Cube and Vector within the same block, ensuring that shared UB tile data is not accessed before the producer finishes. +Cross-core flags use the public `0`-`7` event range. A5 intra-block sync uses a separate +physical event range, described below. + Unlike `wait_cross_flag`, `wait_intra_flag` only stalls the specified pipeline — the SU and other pipelines continue executing. #### `pto.set_intra_flag(pipe, event_id)` @@ -415,8 +418,11 @@ Unlike `wait_cross_flag`, `wait_intra_flag` only stalls the specified pipeline | Parameter | Type | Description | |-----------|------|-------------| -| `pipe` | `Pipe` | Trigger endpoint for the synchronization event. The public DSL accepts `Pipe.MTE3` here. | -| `event_id` | `int` | Event identifier (`0`–`7`) | +| `pipe` | `Pipe` | Trigger endpoint for the synchronization event. The public DSL accepts `Pipe.FIX` and `Pipe.MTE3` here. | +| `event_id` | `int` or scalar expression | Physical event identifier (`0`–`31` for static IDs). | + +For A5 mixed C/V writeback code that must notify both AIV subblocks, emit two signals +explicitly: one for the base event ID and one for `base_id + 16`. **Returns**: None (side-effect operation). @@ -424,7 +430,7 @@ Unlike `wait_cross_flag`, `wait_intra_flag` only stalls the specified pipeline ```python -# Signal event ID0 from the MTE3-side endpoint +# Signal event ID0 from the MTE3-side endpoint. pto.set_intra_flag(pto.Pipe.MTE3, 0) ``` @@ -436,8 +442,11 @@ pto.set_intra_flag(pto.Pipe.MTE3, 0) | Parameter | Type | Description | |-----------|------|-------------| -| `pipe` | `Pipe` | Waiting endpoint for the synchronization event. The public DSL accepts `Pipe.V` here. | -| `event_id` | `int` | Event identifier to wait on (`0`–`7`) | +| `pipe` | `Pipe` | Waiting endpoint for the synchronization event. The public DSL accepts `Pipe.FIX` and `Pipe.V` here. | +| `event_id` | `int` or scalar expression | Physical event identifier to wait on (`0`–`31` for static IDs). | + +For A5 mixed C/V writeback code that must synchronize with both AIV subblocks, wait on both +the base event ID and `base_id + 16`. **Returns**: None (side-effect operation). diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 394d19bba8..c684c9c1ef 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -147,6 +147,13 @@ def _validate_static_buf_id(buf_id, *, context: str): raise ValueError(f"{context} expects static buf_id in [0, 31], got {buf_id}") +def _validate_static_event_id_range(event_id, *, context: str, lo: int, hi: int, meaning: str = "event_id"): + if isinstance(event_id, bool): + raise TypeError(f"{context} does not accept bool values") + if isinstance(event_id, int) and not lo <= event_id <= hi: + raise ValueError(f"{context} expects static {meaning} in [{lo}, {hi}], got {event_id}") + + def _validate_sync_pipe(pipe, *, context: str, allowed: tuple[str, ...]): canonical = _canonical_pipe_token(pipe) if canonical is None: @@ -154,6 +161,7 @@ def _validate_sync_pipe(pipe, *, context: str, allowed: tuple[str, ...]): if canonical not in allowed: expected = ", ".join(f"<{name}>" for name in allowed) raise ValueError(f"{context} expects pipe to be one of {expected}, got <{canonical}>") + return canonical def _require_explicit_mode(surface: str): @@ -4067,6 +4075,82 @@ def _acc_store_unit_flag_attr(unit_flag): context="acc store unit_flag", ) +_ACC_STORE_PRE_QUANT_MODES = frozenset({ + "no_convert", + "f32_f16", + "qf322hif8_pre_vec", + "qf322hif8_pre_scalar", + "qf322hif8_pre_hybrid_vec", + "qf322hif8_pre_hybrid_scalar", + "deqs32_int_vec", + "deqs32_int_scalar", + "req8_vec", + "req8_scalar", + "deqf16_vec", + "deqf16_scalar", + "qf322fp8_pre_vec", + "qf322fp8_pre_scalar", + "qf322f32_pre_vec", + "qf322f32_pre_scalar", + "f32_bf16", + "qf162b8_pre_vec", + "qf162b8_pre_scalar", + "qf162s4_pre_vec", + "qf162s4_pre_scalar", + "req4_vec", + "req4_scalar", + "qf322b8_pre_vec", + "qf322b8_pre_scalar", + "qf322s4_pre_vec", + "qf322s4_pre_scalar", + "deqs16_vec", + "deqs16_scalar", + "qf162s16_pre_vec", + "qf162s16_pre_scalar", + "qf322f16_pre_vec", + "qf322f16_pre_scalar", + "qf322bf16_pre_vec", + "qf322bf16_pre_scalar", + "qs322bf16_pre_vec", + "qs322bf16_pre_scalar", +}) + +_ACC_STORE_PRE_QUANT_SKIP_PAYLOAD_KIND_CHECK = frozenset({ + "no_convert", +}) + +_ACC_STORE_VECTOR_PRE_QUANT_MODES = frozenset({ + mode for mode in _ACC_STORE_PRE_QUANT_MODES if mode.endswith("_vec") +}) + + +def _is_acc_store_float_scalar_payload(value) -> bool: + raw_value = unwrap_surface_value(value) + return hasattr(raw_value, "type") and any( + cls.isinstance(raw_value.type) for cls in (F16Type, BF16Type, F32Type) + ) + + +def _is_acc_store_scaling_pointer_payload(value) -> bool: + raw_value = unwrap_surface_value(value) + if not hasattr(raw_value, "type"): + return False + try: + ptr_type = _pto.PtrType(raw_value.type) + except Exception: + return False + scaling_attr = _pto.AddressSpaceAttr.get(_pto.AddressSpace.SCALING) + memory_space = getattr(ptr_type, "memory_space", None) + if ( + memory_space != scaling_attr + and getattr(memory_space, "value", None) != _pto.AddressSpace.SCALING + ): + return False + return any( + cls.isinstance(ptr_type.element_type) + for cls in (F16Type, BF16Type, F32Type) + ) + def _acc_store_pre_quant(pre_quant): if pre_quant is None: @@ -4074,9 +4158,25 @@ def _acc_store_pre_quant(pre_quant): if not isinstance(pre_quant, tuple) or len(pre_quant) != 2: raise TypeError("acc store pre_quant expects (payload, mode)") payload, mode = pre_quant + normalized_mode = _normalize_token(mode, context="acc store pre_quant mode") + if normalized_mode not in _ACC_STORE_PRE_QUANT_MODES: + raise ValueError(f"unsupported acc store pre_quant mode: {normalized_mode}") + if normalized_mode in _ACC_STORE_PRE_QUANT_SKIP_PAYLOAD_KIND_CHECK: + pass + elif normalized_mode in _ACC_STORE_VECTOR_PRE_QUANT_MODES: + if not _is_acc_store_scaling_pointer_payload(payload): + raise TypeError( + "acc store vector pre_quant payload must be a scaling pointer " + "with f16, bf16, or f32 elements" + ) + else: + if not _is_acc_store_float_scalar_payload(payload): + raise TypeError( + "acc store scalar pre_quant payload must be an f16, bf16, or f32 scalar" + ) return ( unwrap_surface_value(payload), - Attribute.parse(f"#pto"), + Attribute.parse(f"#pto"), ) @@ -4440,7 +4540,15 @@ def mte_load(source, destination, l2_cache_ctl, len_burst, *, nburst, loops=None @_explicit_mode_only("pto.mte_store(...)") -def mte_store(source, destination, len_burst, *, nburst, loops=None): +def mte_store( + source, + destination, + len_burst, + *, + nburst, + loops=None, + l2_cache="nmfv", +): """Ptr-based UB->GM DMA wrapper aligned with the underlying ``pto.dma_store`` surface.""" n_burst, nburst_src_stride, nburst_dst_stride = _normalize_dma_group( "nburst", @@ -4461,6 +4569,10 @@ def mte_store(source, destination, len_burst, *, nburst, loops=None): loop_counts, loop_src_strides, loop_dst_strides, + l2_cache_ctl=_coerce_i64( + _normalize_mte_store_l2_cache(l2_cache, context="mte_store(...) l2_cache"), + context="mte_store l2 cache control", + ), ) @@ -4550,7 +4662,15 @@ def mte_gm_ub(source, destination, l2_cache_ctl, len_burst, *, nburst, loops=Non @_explicit_mode_only("pto.mte_ub_gm(...)") -def mte_ub_gm(source, destination, len_burst, *, nburst, loops=None): +def mte_ub_gm( + source, + destination, + len_burst, + *, + nburst, + loops=None, + l2_cache="nmfv", +): """``pto.mte_ub_gm`` – grouped UB-to-GM DMA surface.""" n_burst, nburst_src_stride, nburst_dst_stride = _normalize_dma_group( "nburst", @@ -4571,6 +4691,10 @@ def mte_ub_gm(source, destination, len_burst, *, nburst, loops=None): loop_counts, loop_src_strides, loop_dst_strides, + l2_cache_ctl=_coerce_i64( + _normalize_mte_store_l2_cache(l2_cache, context="mte_ub_gm(...) l2_cache"), + context="mte_ub_gm l2 cache control", + ), ) @@ -5232,6 +5356,24 @@ def get_lanemask_gt(): "wbhfv", "wbhlv", "wbhprs", "wbhred", "wtsfv", "wtslv", "wtsprs", "wtsred", } +_ST_L2_CACHE_CONTROL_VALUES = { + "nmfv": 0, + "nmlv": 1, + "nmprs": 2, + "nmred": 3, + "naci": 4, + "napw": 5, + "napi": 6, + "nared": 7, + "wbhfv": 8, + "wbhlv": 9, + "wbhprs": 10, + "wbhred": 11, + "wtsfv": 12, + "wtslv": 13, + "wtsprs": 14, + "wtsred": 15, +} _ROUNDING_TOKENS = {"r", "a", "f", "c", "z", "o", "h"} _SATURATION_TOKENS = {"sat", "nosat"} @@ -5260,6 +5402,14 @@ def _st_l2_cache_attr(value, *, context: str): return _simt_enum_attr("st_l2cache", value, supported=_ST_L2_CACHE_TOKENS, context=context) +def _normalize_mte_store_l2_cache(l2_cache, *, context: str): + token = _normalize_token(l2_cache, context=context) + if token not in _ST_L2_CACHE_CONTROL_VALUES: + expected = ", ".join(sorted(_ST_L2_CACHE_CONTROL_VALUES)) + raise ValueError(f"{context} does not support {l2_cache!r}; expected one of {expected}") + return _ST_L2_CACHE_CONTROL_VALUES[token] + + def _rounding_attr(value, *, context: str): return _simt_enum_attr("rounding", value, supported=_ROUNDING_TOKENS, context=context) @@ -5712,6 +5862,11 @@ def _sync_event_id_operand(event_id, *, context: str): return event_id if isinstance(event_id, int) else unwrap_surface_value(event_id) +def _sync_event_id_operand_in_range(event_id, *, context: str, lo: int, hi: int, meaning: str = "event_id"): + _validate_static_event_id_range(event_id, context=context, lo=lo, hi=hi, meaning=meaning) + return event_id if isinstance(event_id, int) else unwrap_surface_value(event_id) + + def _flag_event_id_operand(event_id, *, context: str): if isinstance(event_id, int): _validate_static_event_id(event_id, context=context) @@ -5735,15 +5890,35 @@ def wait_cross_flag(pipe, event_id): def set_intra_flag(pipe, event_id): """``pto.set_intra_flag(pipe, event_id)`` – intra-block sync facade for ``pto.sync.set``.""" - _validate_sync_pipe(pipe, context="set_intra_flag(pipe, event_id)", allowed=("PIPE_MTE3",)) - event_operand = _sync_event_id_operand(event_id, context="set_intra_flag(..., event_id=...)") + _validate_sync_pipe( + pipe, + context="set_intra_flag(pipe, event_id)", + allowed=("PIPE_FIX", "PIPE_MTE3"), + ) + event_operand = _sync_event_id_operand_in_range( + event_id, + context="set_intra_flag(..., event_id=...)", + lo=0, + hi=31, + meaning="physical event_id", + ) _pto.sync_set(_pipe_attr(pipe), event_operand) def wait_intra_flag(pipe, event_id): """``pto.wait_intra_flag(pipe, event_id)`` – intra-block sync facade for ``pto.sync.wait``.""" - _validate_sync_pipe(pipe, context="wait_intra_flag(pipe, event_id)", allowed=("PIPE_V",)) - event_operand = _sync_event_id_operand(event_id, context="wait_intra_flag(..., event_id=...)") + _validate_sync_pipe( + pipe, + context="wait_intra_flag(pipe, event_id)", + allowed=("PIPE_FIX", "PIPE_V"), + ) + event_operand = _sync_event_id_operand_in_range( + event_id, + context="wait_intra_flag(..., event_id=...)", + lo=0, + hi=31, + meaning="physical event_id", + ) _pto.sync_wait(_pipe_attr(pipe), event_operand) diff --git a/ptodsl/ptodsl/_subkernels.py b/ptodsl/ptodsl/_subkernels.py index 0f080b500e..da7e18b563 100644 --- a/ptodsl/ptodsl/_subkernels.py +++ b/ptodsl/ptodsl/_subkernels.py @@ -75,7 +75,6 @@ class SubkernelSpec: symbol_name: str target: str = "a5" simt_max_threads: int | None = None - simt_max_regs: int | None = None class SubkernelTemplate: @@ -404,7 +403,6 @@ def __init__( target: str = "a5", ast_rewrite: bool = True, simt_max_threads: int | None = None, - simt_max_regs: int | None = None, simt_inline_dims: tuple | None = None, ): self._role = role @@ -412,7 +410,6 @@ def __init__( self._target = target self._ast_rewrite = ast_rewrite self._simt_max_threads = simt_max_threads - self._simt_max_regs = simt_max_regs self._simt_inline_dims = simt_inline_dims self._session_cm = None @@ -425,17 +422,14 @@ def __call__(self, fn): symbol_name=self._name or fn.__name__, target=self._target, simt_max_threads=self._simt_max_threads, - simt_max_regs=self._simt_max_regs, ), fn, ast_rewrite=self._ast_rewrite, ) def __enter__(self): - if self._role == KernelRole.SIMT and ( - self._simt_max_threads is not None or self._simt_max_regs is not None - ): - raise TypeError("@pto.simt(max_threads=..., max_regs=...) is only supported as a function decorator") + if self._role == KernelRole.SIMT and self._simt_max_threads is not None: + raise TypeError("@pto.simt(max_threads=...) is only supported as a function decorator") runtime = current_runtime() if runtime is None: raise RuntimeError( @@ -469,7 +463,6 @@ def _subkernel_decorator( target: str = "a5", ast_rewrite: bool = True, simt_max_threads: int | None = None, - simt_max_regs: int | None = None, simt_inline_dims: tuple | None = None, ): return _SubkernelSurface( @@ -478,7 +471,6 @@ def _subkernel_decorator( target=target, ast_rewrite=ast_rewrite, simt_max_threads=simt_max_threads, - simt_max_regs=simt_max_regs, simt_inline_dims=simt_inline_dims, ) @@ -491,7 +483,6 @@ def _decorate_subkernel( target: str = "a5", ast_rewrite: bool = True, simt_max_threads: int | None = None, - simt_max_regs: int | None = None, simt_inline_dims: tuple | None = None, ): if fn is not None: @@ -501,7 +492,6 @@ def _decorate_subkernel( target=target, ast_rewrite=ast_rewrite, simt_max_threads=simt_max_threads, - simt_max_regs=simt_max_regs, simt_inline_dims=simt_inline_dims, )(fn) return _subkernel_decorator( @@ -510,7 +500,6 @@ def _decorate_subkernel( target=target, ast_rewrite=ast_rewrite, simt_max_threads=simt_max_threads, - simt_max_regs=simt_max_regs, simt_inline_dims=simt_inline_dims, ) @@ -542,10 +531,8 @@ def simt( target: str = "a5", ast_rewrite: bool = True, max_threads: int | None = None, - max_regs: int | None = None, ): max_threads = _validate_simt_resource_attr("max_threads", max_threads) - max_regs = _validate_simt_resource_attr("max_regs", max_regs) simt_inline_dims = None if fn is not None and not callable(fn): dims = (fn, *dims) @@ -561,7 +548,6 @@ def simt( target=target, ast_rewrite=ast_rewrite, simt_max_threads=max_threads, - simt_max_regs=max_regs, simt_inline_dims=simt_inline_dims, ) diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index e990a26a39..a7e8bcd17f 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -596,13 +596,6 @@ def _get_or_create_simt_helper_function(self, subkernel, *args, **kwargs): IntegerAttr.get(i32_attr_type, subkernel.spec.simt_max_threads), ) ) - if subkernel.spec.simt_max_regs is not None: - helper_attributes.append( - ( - "pto.simt_max_regs", - IntegerAttr.get(i32_attr_type, subkernel.spec.simt_max_regs), - ) - ) helper_fn = self._create_named_helper_function( helper_symbol, arg_types, diff --git a/ptodsl/ptodsl/tilelib/templates/a5/_load_store.py b/ptodsl/ptodsl/tilelib/templates/a5/_load_store.py index 01b216b556..ae05604f57 100644 --- a/ptodsl/ptodsl/tilelib/templates/a5/_load_store.py +++ b/ptodsl/ptodsl/tilelib/templates/a5/_load_store.py @@ -219,24 +219,28 @@ def tstore_nz_constraint(src_kind, src_shape, src_valid_shape, src_memory_space, ) -def tload_mat_nd2nz_constraint(src_kind, src_shape, src_memory_space, dst_kind, dst_valid_shape, dst_memory_space, dst_config, dst_dtype, **_): +def tload_mat_nd2nz_constraint(src_kind, src_shape, src_memory_space, dst_kind, dst_shape, dst_memory_space, dst_config, dst_dtype, **_): if src_kind != "view" or dst_kind != "tile" or src_memory_space != "gm" or dst_memory_space != "mat": return False if dst_config.b_layout != "col_major" or dst_config.s_layout != "row_major": return False if dst_dtype not in {"f16", "bf16", "f32"}: return False - return _view_rank(src_shape) != 5 or _known_eq(src_shape[4], dst_valid_shape[1]) + if _view_rank(src_shape) != 5: + return False + return _known_eq(src_shape[4], dst_shape[1]) -def tload_mat_dn2nz_constraint(src_kind, src_shape, src_memory_space, dst_kind, dst_valid_shape, dst_memory_space, dst_config, dst_dtype, **_): +def tload_mat_dn2nz_constraint(src_kind, src_shape, src_memory_space, dst_kind, dst_shape, dst_memory_space, dst_config, dst_dtype, **_): if src_kind != "view" or dst_kind != "tile" or src_memory_space != "gm" or dst_memory_space != "mat": return False if dst_config.b_layout != "col_major" or dst_config.s_layout != "row_major": return False if dst_dtype not in {"f16", "bf16", "f32"}: return False - return _view_rank(src_shape) != 5 or _known_eq(src_shape[4], dst_valid_shape[0]) + if _view_rank(src_shape) != 5: + return False + return _known_eq(src_shape[4], dst_shape[0]) def tstore_acc_base(src_kind, src_memory_space, src_dtype, dst_kind, dst_memory_space, **_): diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tload.py b/ptodsl/ptodsl/tilelib/templates/a5/tload.py index b6be85316c..db775edc5d 100644 --- a/ptodsl/ptodsl/tilelib/templates/a5/tload.py +++ b/ptodsl/ptodsl/tilelib/templates/a5/tload.py @@ -231,12 +231,17 @@ def template_tload_nz2nz(src: pto.PartitionTensorView, dst: pto.Tile): ) def template_tload_gm_to_mat_nd2nz(src: pto.PartitionTensorView, dst: pto.Tile): m, k = dst.valid_shape + elem_bytes = pto.bytewidth(dst.dtype) + src_inner_stride = k * elem_bytes + if len(src.shape) == 5 and src.strides is not None and src.strides[3] is not None: + src_inner_stride = src.strides[3] * elem_bytes + pto.mte_gm_l1_frac( src.as_ptr(), dst.as_ptr(), "nd2nz", shape=(m, k), - src_layout=(k,), + src_layout=(src_inner_stride,), dst_group=(1, 1, m, 0), ctrl=(0, False), ) @@ -258,12 +263,17 @@ def template_tload_gm_to_mat_nd2nz(src: pto.PartitionTensorView, dst: pto.Tile): ) def template_tload_gm_to_mat_dn2nz(src: pto.PartitionTensorView, dst: pto.Tile): m, k = dst.valid_shape + elem_bytes = pto.bytewidth(dst.dtype) + src_inner_stride = m * elem_bytes + if len(src.shape) == 5 and src.strides is not None and src.strides[4] is not None: + src_inner_stride = src.strides[4] * elem_bytes + pto.mte_gm_l1_frac( src.as_ptr(), dst.as_ptr(), "dn2nz", shape=(k, m), - src_layout=(m,), + src_layout=(src_inner_stride,), dst_group=(1, 1, k, 0), ctrl=(0, False), ) diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tstore.py b/ptodsl/ptodsl/tilelib/templates/a5/tstore.py index dce9bed211..3b216d31dc 100644 --- a/ptodsl/ptodsl/tilelib/templates/a5/tstore.py +++ b/ptodsl/ptodsl/tilelib/templates/a5/tstore.py @@ -206,13 +206,17 @@ def template_tstore_nz(src: pto.Tile, dst: pto.PartitionTensorView): ) def template_tstore_acc_to_gm_nz2nd(src: pto.Tile, dst: pto.PartitionTensorView): m, n = src.valid_shape + strides = dst.strides + src_stride = src.shape[0] + dst_stride = n if strides is None or strides[3] is None else strides[3] + pto.mte_l0c_gm( src.as_ptr(), dst.as_ptr(), m, n, - n, - n, + src_stride, + dst_stride, 0, 0, layout="nz2nd", @@ -236,16 +240,21 @@ def template_tstore_acc_to_gm_nz2nd(src: pto.Tile, dst: pto.PartitionTensorView) ) def template_tstore_acc_to_gm_nz2dn(src: pto.Tile, dst: pto.PartitionTensorView): m, n = src.valid_shape + strides = dst.strides + src_stride = src.shape[0] + dst_stride = m if strides is None or strides[4] is None else strides[4] + loop0_src_stride = 1 + pto.mte_l0c_gm( src.as_ptr(), dst.as_ptr(), m, n, - n, - m, + src_stride, + dst_stride, 0, 0, - layout=("nz2dn", 1), + layout=("nz2dn", loop0_src_stride), ) @@ -266,13 +275,16 @@ def template_tstore_acc_to_gm_nz2dn(src: pto.Tile, dst: pto.PartitionTensorView) ) def template_tstore_acc_to_gm_nz2nz(src: pto.Tile, dst: pto.PartitionTensorView): m, n = src.valid_shape + src_stride = src.shape[0] + dst_stride = n + pto.mte_l0c_gm( src.as_ptr(), dst.as_ptr(), m, n, - n, - n, + src_stride, + dst_stride, 0, 0, layout=("nz2nz", 1), @@ -295,14 +307,18 @@ def template_tstore_acc_to_gm_nz2nz(src: pto.Tile, dst: pto.PartitionTensorView) ) def template_tstore_fp_acc_to_gm(src: pto.Tile, fp: pto.Tile, dst: pto.PartitionTensorView): m, n = src.valid_shape + strides = dst.strides quant_mode = "qf322bf16_pre_vec" if str(fp.dtype) == "bf16" else "qf322f16_pre_vec" + src_stride = src.shape[0] + dst_stride = n if strides is None or strides[3] is None else strides[3] + pto.mte_l0c_gm( src.as_ptr(), dst.as_ptr(), m, n, - n, - n, + src_stride, + dst_stride, 0, 0, layout="nz2nd", diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 75adb03bce..2a9fa6c44e 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -838,7 +838,7 @@ def simt_grouped_query_probe(): pto.keep(grid_z, slot=8) -@pto.simt(max_threads=256, max_regs=48) +@pto.simt(max_threads=256) def simt_resource_attr_probe(): pto.get_tid_x() @@ -1960,6 +1960,96 @@ def public_cube_surface_probe( pto.mte_l0c_ub(acc_tile.as_ptr(), out_tile.as_ptr(), m, n, n, n, split=pto.SplitMode.M, layout="nz2nd") +@pto.jit(target="a5", mode="explicit") +def acc_store_pre_quant_surface_probe( + dst: pto.ptr(pto.bf16, "gm"), +): + size = pto.const(16) + zero = pto.const(0) + src = pto.alloc_tile( + shape=[16, 16], + dtype=pto.f32, + memory_space=pto.MemorySpace.ACC, + valid_shape=[16, 16], + blayout="ColMajor", + slayout="RowMajor", + ) + pto.mte_l0c_gm( + src.as_ptr(), + dst, + size, + size, + size, + size, + zero, + zero, + pre_quant=(pto.bf16(1.0), "f32_bf16"), + layout="nz2nd", + ) + + +@pto.jit(target="a5", mode="explicit") +def acc_store_bad_pre_quant_mode_probe( + dst: pto.ptr(pto.bf16, "gm"), +): + size = pto.const(16) + zero = pto.const(0) + src = pto.alloc_tile( + shape=[16, 16], + dtype=pto.f32, + memory_space=pto.MemorySpace.ACC, + valid_shape=[16, 16], + blayout="ColMajor", + slayout="RowMajor", + ) + pto.mte_l0c_gm( + src.as_ptr(), + dst, + size, + size, + size, + size, + zero, + zero, + pre_quant=(pto.bf16(1.0), "not_a_mode"), + layout="nz2nd", + ) + + +@pto.jit(target="a5", mode="explicit") +def acc_store_bad_pre_quant_payload_probe( + dst: pto.ptr(pto.bf16, "gm"), +): + size = pto.const(16) + zero = pto.const(0) + src = pto.alloc_tile( + shape=[16, 16], + dtype=pto.f32, + memory_space=pto.MemorySpace.ACC, + valid_shape=[16, 16], + blayout="ColMajor", + slayout="RowMajor", + ) + payload = pto.alloc_tile( + shape=[16], + dtype=pto.bf16, + memory_space=pto.MemorySpace.SCALING, + valid_shape=[16], + ) + pto.mte_l0c_gm( + src.as_ptr(), + dst, + size, + size, + size, + size, + zero, + zero, + pre_quant=(payload.as_ptr(), "f32_bf16"), + layout="nz2nd", + ) + + @pto.cube def public_cube_tile_mx_probe( mat_lhs: pto.Tile, @@ -2469,8 +2559,10 @@ def public_sync_surface_probe(): pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE3, event_id=dynamic_event) pto.set_cross_flag(pto.Pipe.FIX, 0) pto.set_intra_flag(pto.Pipe.MTE3, dynamic_event) + pto.set_intra_flag(pto.Pipe.FIX, 4) pto.wait_cross_flag(pto.Pipe.FIX, 0) pto.wait_intra_flag(pto.Pipe.V, dynamic_event) + pto.wait_intra_flag(pto.Pipe.FIX, 20) @pto.jit(target="a5") @@ -2527,6 +2619,8 @@ def public_data_movement_surface_probe(): pto.mte_gm_ub(gm_src, ub_dst, 0, 256, nburst=(8, 256, 256), loops=[(4, 2048, 2048)]) pto.mte_gm_ub(gm_src, ub_dst, 0, 200, nburst=(64, 200, 256), pad=(0.0, 0, 0)) pto.mte_ub_gm(ub_src, gm_dst, 256, nburst=(64, 256, 1024)) + pto.mte_ub_gm(ub_src, gm_dst, 128, nburst=(1, 128, 128), l2_cache="nared") + pto.mte_ub_gm(ub_src, gm_dst, 64, nburst=(1, 64, 64), l2_cache="wtsred") pto.mte_ub_ub(ub_src, ub_dst, 8, nburst=(16, 0, 4)) pto.mte_ub_l1(ub_src, l1_dst, 8, nburst=(16, 0, 4)) pto.mte_gm_l1(gm_src, l1_dst, 256, nburst=(8, 256, 256), loops=[(2, 2048, 2048)]) @@ -4560,21 +4654,16 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): expect_parse_roundtrip_and_verify(simt_resource_attr_text, "simt resource attr launch specialization") expect( re.search( - r"func\.func @simt_resource_attr_probe__simt_\d+\(\) attributes \{pto\.simt_entry, pto\.simt_max_regs = 48 : i32, pto\.simt_max_threads = 256 : i32\}", + r"func\.func @simt_resource_attr_probe__simt_\d+\(\) attributes \{pto\.simt_entry, pto\.simt_max_threads = 256 : i32\}", simt_resource_attr_text, ) is not None, - "@pto.simt(max_threads=..., max_regs=...) should attach resource attrs to the helper function", + "@pto.simt(max_threads=...) should attach resource attrs to the helper function", ) expect_raises( ValueError, lambda: pto.simt(max_threads=0)(lambda: None), "max_threads", ) - expect_raises( - TypeError, - lambda: pto.simt(max_regs=True)(lambda: None), - "max_regs", - ) def _enter_inline_simt_with_resource_attr(): with pto.simt(max_threads=256): @@ -5312,6 +5401,22 @@ def _enter_inline_simt_with_resource_attr(): public_surface_text = public_surface_exports_probe.compile().mlir_text() expect_parse_roundtrip_and_verify(public_surface_text, "public surface export specialization") + acc_store_pre_quant_text = acc_store_pre_quant_surface_probe.compile().mlir_text() + expect_parse_roundtrip_and_verify(acc_store_pre_quant_text, "acc-store pre_quant public mode specialization") + expect( + "pre_quant(%" in acc_store_pre_quant_text and "mode = f32_bf16" in acc_store_pre_quant_text, + "mte_l0c_gm pre_quant should preserve the f32_bf16 public mode", + ) + expect_raises( + ValueError, + lambda: acc_store_bad_pre_quant_mode_probe.compile().mlir_text(), + "unsupported acc store pre_quant mode", + ) + expect_raises( + TypeError, + lambda: acc_store_bad_pre_quant_payload_probe.compile().mlir_text(), + "acc store scalar pre_quant payload must be an f16, bf16, or f32 scalar", + ) compile_time_query_text = compile_time_query_probe.compile().mlir_text() expect_parse_roundtrip_and_verify(compile_time_query_text, "compile-time query specialization") eager_scalar_text = eager_scalar_constructor_probe.compile().mlir_text() @@ -5363,6 +5468,10 @@ def _enter_inline_simt_with_resource_attr(): expect_parse_roundtrip_and_verify(fixed_width_integer_text, "fixed-width integer specialization") expect("pto.mte_gm_ub" in public_surface_text, "mte_load(...) should lower to pto.mte_gm_ub") expect("pto.mte_ub_gm" in public_surface_text, "mte_store(...) should lower to pto.mte_ub_gm") + expect( + re.search(r"pto\.mte_ub_gm [^\n]+ nburst\([^)]+\) l2_cache_ctl\(%c0[^)\n]*\)", public_surface_text) is not None, + "mte_store(...) should pass default l2_cache_ctl through the pto.mte_ub_gm l2_cache_ctl group", + ) expect(public_surface_text.count("pto.mem_bar") >= 1, "mem_bar(...) should still lower explicit memory barriers") expect("pto.barrier " in public_surface_text, "pipe_barrier(Pipe.ALL) should lower to pto.barrier") expect("pto.vexp" in public_surface_text, "vexp(...) should lower to pto.vexp") @@ -5442,9 +5551,19 @@ def _enter_inline_simt_with_resource_attr(): expect("pto.sync.set , 0" in sync_surface_text, "set_cross_flag(Pipe.FIX, 0) should lower to pto.sync.set") expect("pto.sync.wait , 0" in sync_surface_text, "wait_cross_flag(Pipe.FIX, 0) should lower to pto.sync.wait") expect("pto.sync.set , %c3" in sync_surface_text, "set_intra_flag(Pipe.MTE3, dynamic_event) should lower dynamic event ids through pto.sync.set") + expect("pto.sync.set , 4" in sync_surface_text, "set_intra_flag(Pipe.FIX, 4) should lower physical event ids through pto.sync.set") expect("pto.sync.wait , %c3" in sync_surface_text, "wait_intra_flag(Pipe.V, dynamic_event) should lower dynamic event ids through pto.sync.wait") + expect("pto.sync.wait , 20" in sync_surface_text, "wait_intra_flag(Pipe.FIX, 20) should lower physical event ids through pto.sync.wait") expect(data_movement_surface_text.count("pto.mte_gm_ub") == 2, "public grouped GM->UB wrappers should lower to pto.mte_gm_ub") expect("pto.mte_ub_gm" in data_movement_surface_text, "public grouped UB->GM wrapper should lower to pto.mte_ub_gm") + expect( + re.search(r"pto\.mte_ub_gm [^\n]+ nburst\([^)]+\) l2_cache_ctl\(%c7[^)\n]*\)", data_movement_surface_text) is not None, + "mte_ub_gm(..., l2_cache='nared') should map to the l2_cache_ctl group", + ) + expect( + re.search(r"pto\.mte_ub_gm [^\n]+ nburst\([^)]+\) l2_cache_ctl\(%c15[^)\n]*\)", data_movement_surface_text) is not None, + "mte_ub_gm(..., l2_cache='wtsred') should map to the l2_cache_ctl group", + ) expect("pto.mte_ub_ub" in data_movement_surface_text, "public grouped UB->UB wrapper should lower to pto.mte_ub_ub") expect("pto.mte_ub_l1" in data_movement_surface_text, "public grouped UB->L1 wrapper should lower to pto.mte_ub_l1") expect("pto.mte_gm_l1" in data_movement_surface_text, "public grouped GM->L1 wrapper should lower to pto.mte_gm_l1") diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index cfa7c24a35..e58adc7d78 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -418,6 +418,15 @@ def test_cube_sat_modes_map_to_backend_tokens(self): self.assertEqual(_ops._mad_sat_attr(pto.SatMode.OFF), "#pto") self.assertEqual(_ops._acc_store_sat_attr(pto.SatMode.PRESERVE_NAN), "#pto") + def test_acc_store_no_convert_skips_payload_kind_check(self): + payload = object() + with patch.object(_ops, "Attribute") as attr: + attr.parse.side_effect = lambda text: text + value, mode = _ops._acc_store_pre_quant((payload, "no_convert")) + + self.assertIs(value, payload) + self.assertEqual(mode, "#pto") + def test_tile_selection_surface_exposes_optional_tmp(self): for func, expected in [ (_ops.tsel, ["mask", "src0", "src1", "dst", "tmp"]), @@ -574,8 +583,10 @@ def test_sync_event_id_rejects_out_of_range_static_values(self): (_ops.wait_flag, ("MTE2", "V"), {"event_id": -1}, "wait_flag(..., event_id=...)"), (_ops.set_cross_flag, (pto.Pipe.FIX, 8), {}, "set_cross_flag(..., event_id=...)"), (_ops.wait_cross_flag, (pto.Pipe.FIX, -1), {}, "wait_cross_flag(..., event_id=...)"), - (_ops.set_intra_flag, (pto.Pipe.MTE3, 9), {}, "set_intra_flag(..., event_id=...)"), - (_ops.wait_intra_flag, (pto.Pipe.V, -2), {}, "wait_intra_flag(..., event_id=...)"), + (_ops.set_intra_flag, (pto.Pipe.FIX, 32), {}, "set_intra_flag(..., event_id=...)", "[0, 31]"), + (_ops.set_intra_flag, (pto.Pipe.MTE3, 32), {}, "set_intra_flag(..., event_id=...)", "[0, 31]"), + (_ops.wait_intra_flag, (pto.Pipe.V, -2), {}, "wait_intra_flag(..., event_id=...)", "[0, 31]"), + (_ops.wait_intra_flag, (pto.Pipe.FIX, 32), {}, "wait_intra_flag(..., event_id=...)", "[0, 31]"), ] with patch.object(_ops._pto, "set_flag") as set_flag_op, \ @@ -584,13 +595,14 @@ def test_sync_event_id_rejects_out_of_range_static_values(self): patch.object(_ops._pto, "wait_flag_dyn") as wait_flag_dyn_op, \ patch.object(_ops._pto, "sync_set") as sync_set_op, \ patch.object(_ops._pto, "sync_wait") as sync_wait_op: - for func, args, kwargs, context in cases: + for case in cases: + func, args, kwargs, context, *expected_range = case with self.subTest(func=func.__name__, event_id=kwargs.get("event_id", args[-1])): with self.assertRaises(ValueError) as exc: func(*args, **kwargs) message = str(exc.exception) self.assertIn(context, message) - self.assertIn("[0, 7]", message) + self.assertIn(expected_range[0] if expected_range else "[0, 7]", message) set_flag_op.assert_not_called() set_flag_dyn_op.assert_not_called() @@ -603,8 +615,8 @@ def test_sync_facades_reject_illegal_pipe_endpoints(self): cases = [ (_ops.set_cross_flag, (pto.Pipe.V, 0), "set_cross_flag(pipe, event_id)", "", ""), (_ops.wait_cross_flag, (pto.Pipe.MTE3, 0), "wait_cross_flag(pipe, event_id)", "", ""), - (_ops.set_intra_flag, (pto.Pipe.FIX, 0), "set_intra_flag(pipe, event_id)", "", ""), - (_ops.wait_intra_flag, (pto.Pipe.MTE3, 0), "wait_intra_flag(pipe, event_id)", "", ""), + (_ops.set_intra_flag, (pto.Pipe.V, 0), "set_intra_flag(pipe, event_id)", ", ", ""), + (_ops.wait_intra_flag, (pto.Pipe.MTE3, 0), "wait_intra_flag(pipe, event_id)", ", ", ""), ] with patch.object(_ops._pto, "sync_set") as sync_set_op, \ @@ -621,6 +633,19 @@ def test_sync_facades_reject_illegal_pipe_endpoints(self): sync_set_op.assert_not_called() sync_wait_op.assert_not_called() + def test_intra_sync_mixed_writeback_event_ranges(self): + with patch.object(_ops, "_pipe_attr", side_effect=lambda pipe: f"pipe:{pipe}") as pipe_attr, \ + patch.object(_ops._pto, "sync_set") as sync_set_op, \ + patch.object(_ops._pto, "sync_wait") as sync_wait_op: + _ops.set_intra_flag(pto.Pipe.FIX, 31) + _ops.set_intra_flag(pto.Pipe.MTE3, 31) + _ops.wait_intra_flag(pto.Pipe.FIX, 16) + _ops.wait_intra_flag(pto.Pipe.V, 31) + + self.assertEqual(pipe_attr.call_count, 4) + self.assertEqual(sync_set_op.call_count, 2) + self.assertEqual(sync_wait_op.call_count, 2) + def test_pipe_namespace_and_buffer_helpers_are_exposed(self): names = [ "c2v", "v2c", "bidirectional", diff --git a/test/lit/pto/issue711_tnotify_mte_drain.pto b/test/lit/pto/issue711_tnotify_mte_drain.pto index 6ca67c6bee..71c3893b29 100644 --- a/test/lit/pto/issue711_tnotify_mte_drain.pto +++ b/test/lit/pto/issue711_tnotify_mte_drain.pto @@ -6,15 +6,11 @@ // 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. -// Regression for issue #711: pto.comm.tnotify lowering must drain MTE-side -// pipes before emitting pto::comm::TNOTIFY(...). TNOTIFY_IMPL writes the -// signal on the scalar pipe and only does pipe_barrier(PIPE_ALL) *after* the -// store, so without an MTE drain the signal can overtake an in-flight -// pto.tload / pto.tstore (local or peer) and the receiver's matching TWAIT -// returns before the data is visible. For prior MTE3 stores, the notification -// also needs a DDR-domain barrier_all before publishing the signal. The drain -// should be precise: only MTE pipes that may have prior in-flight work are -// drained before TNOTIFY, and MTE2-only cases do not emit a DDR fence. +// Regression for issue #711 under the explicit memory-consistency contract: +// TNotify release requires the producer side to drain local producer pipes +// before the GM fence and before pto.comm.tnotify. PTOAS no longer runs the +// historical MemoryConsistency analysis pass, so the explicit GM fence lowering +// conservatively drains MTE2, MTE3, and FIX before dsb(DSB_DDR). // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s @@ -49,8 +45,6 @@ module { outs(%tile : !pto.tile_buf) pto.tstore ins(%tile : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x32xf32>) - pto.barrier - pto.cmo.cacheinvalid %dst single_cache_line : !pto.partition_tensor_view<1x32xf32> pto.fence.barrier_all #pto.fence_scope %sig_view = pto.make_tensor_view %signal_ptr, @@ -85,7 +79,6 @@ module { pto.tstore ins(%acc : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<16x16xf32>) - pto.cmo.cacheinvalid %dst single_cache_line : !pto.partition_tensor_view<16x16xf32> pto.fence.barrier_all #pto.fence_scope %sig_view = pto.make_tensor_view %signal_ptr, @@ -120,7 +113,7 @@ module { pto.tload ins(%src : !pto.partition_tensor_view<1x32xf32>) outs(%tile : !pto.tile_buf) - pto.cmo.cacheinvalid %src single_cache_line : !pto.partition_tensor_view<1x32xf32> + pto.barrier %sig_view = pto.make_tensor_view %signal_ptr, shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> @@ -154,7 +147,6 @@ module { pto.tload ins(%src : !pto.partition_tensor_view<1x32xf32>) outs(%tile : !pto.tile_buf) - pto.cmo.cacheinvalid %src single_cache_line : !pto.partition_tensor_view<1x32xf32> pto.barrier %sig_view = pto.make_tensor_view %signal_ptr, @@ -167,8 +159,8 @@ module { return } - // A user/pass-provided MTE3 barrier drains the pending store; the explicit - // barrier_all completes the publish sequence. + // The explicit barrier_all conservatively drains MTE2, MTE3 and FIX, then + // completes the publish sequence with a DDR fence. func.func @tnotify_after_existing_mte3_barrier_and_release( %src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, @@ -199,8 +191,6 @@ module { pto.barrier pto.tstore ins(%tile : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x32xf32>) - pto.cmo.cacheinvalid %dst single_cache_line : !pto.partition_tensor_view<1x32xf32> - pto.barrier pto.fence.barrier_all #pto.fence_scope %sig_view = pto.make_tensor_view %signal_ptr, @@ -213,8 +203,8 @@ module { return } - // If the user already writes the complete release sequence, the pass should - // not emit a second DDR fence before TNotify. + // The GM fence lowering is conservative and self-contained; no analysis pass + // needs to add a second DDR fence before TNotify. func.func @tnotify_after_existing_mte3_barrier_and_fence( %dst_ptr: !pto.ptr, %signal_ptr: !pto.ptr) @@ -235,8 +225,6 @@ module { pto.tstore ins(%tile : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x32xf32>) - pto.cmo.cacheinvalid %dst single_cache_line : !pto.partition_tensor_view<1x32xf32> - pto.barrier pto.fence.barrier_all #pto.fence_scope %sig_view = pto.make_tensor_view %signal_ptr, @@ -294,12 +282,15 @@ module { // CHECK: TSTORE( // CHECK: pipe_barrier(PIPE_MTE2); // CHECK-NEXT: pipe_barrier(PIPE_MTE3); +// CHECK-NEXT: pipe_barrier(PIPE_FIX); // CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( // CHECK-LABEL: AICORE void tnotify_drain_after_acc_tstore( // CHECK: TSTORE -// CHECK: pipe_barrier(PIPE_FIX); +// CHECK: pipe_barrier(PIPE_MTE2); +// CHECK-NEXT: pipe_barrier(PIPE_MTE3); +// CHECK-NEXT: pipe_barrier(PIPE_FIX); // CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( @@ -318,13 +309,17 @@ module { // CHECK-LABEL: AICORE void tnotify_after_existing_mte3_barrier_and_release( // CHECK: TSTORE( -// CHECK: pipe_barrier(PIPE_MTE3); +// CHECK: pipe_barrier(PIPE_MTE2); +// CHECK-NEXT: pipe_barrier(PIPE_MTE3); +// CHECK-NEXT: pipe_barrier(PIPE_FIX); // CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( // CHECK-LABEL: AICORE void tnotify_after_existing_mte3_barrier_and_fence( // CHECK: TSTORE( -// CHECK: pipe_barrier(PIPE_MTE3); +// CHECK: pipe_barrier(PIPE_MTE2); +// CHECK-NEXT: pipe_barrier(PIPE_MTE3); +// CHECK-NEXT: pipe_barrier(PIPE_FIX); // CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( diff --git a/test/lit/pto/issue872_tput_tnotify_release.pto b/test/lit/pto/issue872_tput_tnotify_release.pto index 798984a9c5..ecd19f65c3 100644 --- a/test/lit/pto/issue872_tput_tnotify_release.pto +++ b/test/lit/pto/issue872_tput_tnotify_release.pto @@ -6,9 +6,11 @@ // 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. -// Regression for issue #872: pto.comm.tput is a macro that can issue MTE3 GM -// stores internally. A following TNotify publishes a cross-rank signal, so those -// TPUT payload stores must be drained and made DDR-visible before the signal. +// Regression for issue #872 under the explicit memory-consistency contract: +// pto.comm.tput/pto.comm.tbroadcast can issue MTE3 GM stores internally. A +// following TNotify publishes a cross-rank signal, so the producer side must +// execute a GM fence before the signal. The fence lowering conservatively +// drains MTE2, MTE3, and FIX before dsb(DSB_DDR). // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s @@ -45,7 +47,6 @@ module { !pto.partition_tensor_view<8x64xf32>, !pto.tile_buf) {atomicType = #pto} - pto.cmo.cacheinvalid %dst single_cache_line : !pto.partition_tensor_view<8x64xf32> pto.fence.barrier_all #pto.fence_scope %sig_view = pto.make_tensor_view %signal_ptr, @@ -90,7 +91,6 @@ module { !pto.partition_tensor_view<8x64xf32>, !pto.tile_buf) {atomicType = #pto} - pto.cmo.cacheinvalid %dst single_cache_line : !pto.partition_tensor_view<8x64xf32> pto.barrier pto.fence.barrier_all #pto.fence_scope @@ -133,7 +133,6 @@ module { !pto.partition_tensor_view<1x32xf32>, !pto.tile_buf, !pto.partition_tensor_view<1x32xf32>) {root = 0 : i32} - pto.cmo.cacheinvalid %peer single_cache_line : !pto.partition_tensor_view<1x32xf32> pto.fence.barrier_all #pto.fence_scope %sig_view = pto.make_tensor_view %signal_ptr, @@ -150,8 +149,9 @@ module { // CHECK-LABEL: AICORE void tput_tnotify_release( // CHECK: pto::comm::TPUT( // CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE -// CHECK-NOT: pipe_barrier(PIPE_MTE2); -// CHECK: pipe_barrier(PIPE_MTE3); +// CHECK: pipe_barrier(PIPE_MTE2); +// CHECK-NEXT: pipe_barrier(PIPE_MTE3); +// CHECK-NEXT: pipe_barrier(PIPE_FIX); // CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( @@ -159,15 +159,17 @@ module { // CHECK: pto::comm::TPUT( // CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE // CHECK-NEXT: pipe_barrier(PIPE_ALL); -// CHECK-NOT: pipe_barrier(PIPE_MTE2); -// CHECK-NOT: pipe_barrier(PIPE_MTE3); -// CHECK: dsb(DSB_DDR); +// CHECK-NEXT: pipe_barrier(PIPE_MTE2); +// CHECK-NEXT: pipe_barrier(PIPE_MTE3); +// CHECK-NEXT: pipe_barrier(PIPE_FIX); +// CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( // CHECK-LABEL: AICORE void tbroadcast_tnotify_release( // CHECK: pto::comm::TBROADCAST( // CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE -// CHECK-NOT: pipe_barrier(PIPE_MTE2); -// CHECK: pipe_barrier(PIPE_MTE3); +// CHECK: pipe_barrier(PIPE_MTE2); +// CHECK-NEXT: pipe_barrier(PIPE_MTE3); +// CHECK-NEXT: pipe_barrier(PIPE_FIX); // CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( diff --git a/test/lit/pto/memory_consistency_external_func.pto b/test/lit/pto/memory_consistency_external_func.pto deleted file mode 100644 index e24951f0cd..0000000000 --- a/test/lit/pto/memory_consistency_external_func.pto +++ /dev/null @@ -1,37 +0,0 @@ -// 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. - -// MemoryConsistency must skip external func declarations. They have no body to -// scan and should not affect release/acquire state in real kernels. - -// RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s - -module { - func.func private @external_consumer(!pto.ptr) - - func.func @external_func_decl_is_skipped(%signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %v_i32 = arith.constant 1 : i32 - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - return - } -} - -// CHECK-LABEL: AICORE void external_func_decl_is_skipped( -// CHECK-NOT: pipe_barrier( -// CHECK-NOT: dsb( -// CHECK: pto::comm::TNOTIFY( diff --git a/test/lit/pto/memory_consistency_invalid.pto b/test/lit/pto/memory_consistency_invalid.pto deleted file mode 100644 index a225c15163..0000000000 --- a/test/lit/pto/memory_consistency_invalid.pto +++ /dev/null @@ -1,103 +0,0 @@ -// 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: not ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s - -module { - func.func @scalar_release_cmo_after_fence( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %v_i32 = arith.constant 7 : i32 - - pto.store_scalar %v_i32, %payload_ptr[%c0] : !pto.ptr, i32 - pto.fence.barrier_all #pto.fence_scope - pto.cmo.cacheinvalid %payload_ptr single_cache_line : !pto.ptr - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - return - } - - func.func @scalar_release_addressed_cmo_before_store( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %v_i32 = arith.constant 7 : i32 - - pto.cmo.cacheinvalid %payload_ptr single_cache_line : !pto.ptr - pto.store_scalar %v_i32, %payload_ptr[%c0] : !pto.ptr, i32 - pto.fence.barrier_all #pto.fence_scope - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - return - } - - func.func @missing_acquire_invalidate( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %v_i32 = arith.constant 1 : i32 - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.twait(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {cmp = #pto} - - %val = pto.load_scalar %payload_ptr[%c0] : !pto.ptr -> i32 - pto.store_scalar %val, %payload_ptr[%c0] : !pto.ptr, i32 - return - } - - func.func @dirty_cache_before_acquire_load( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %v_i32 = arith.constant 1 : i32 - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.twait(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {cmp = #pto} - - pto.store_scalar %v_i32, %payload_ptr[%c0] : !pto.ptr, i32 - %val = pto.load_scalar %payload_ptr[%c0] : !pto.ptr -> i32 - pto.store_scalar %val, %payload_ptr[%c0] : !pto.ptr, i32 - return - } -} - -// CHECK: requires explicit `pto.fence.barrier_all #pto.fence_scope` after the matching `pto.cmo.cacheinvalid` release marker -// CHECK: requires explicit `pto.cmo.cacheinvalid %addr single_cache_line` or `pto.cmo.cacheinvalid all #pto.address_space` after matching cacheable GM scalar stores -// CHECK: requires explicit `pto.cmo.cacheinvalid all #pto.address_space` before a cacheable GM load -// CHECK: dirty GM cache may exist; insert explicit `pto.cmo.cacheinvalid all #pto.address_space` before the load diff --git a/test/lit/pto/memory_consistency_loop_release.pto b/test/lit/pto/memory_consistency_loop_release.pto deleted file mode 100644 index dcd99d253a..0000000000 --- a/test/lit/pto/memory_consistency_loop_release.pto +++ /dev/null @@ -1,88 +0,0 @@ -// 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 %s -o - 2>&1 | FileCheck %s - -// A loop-local barrier_all must clear the loop-carried GM-write pending state. -// Otherwise the next iteration's TNotify is falsely diagnosed as missing an -// explicit barrier_all even though each iteration already has one. - -module { - func.func @loop_tstore_release_tnotify( - %dst_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c2 = arith.constant 2 : index - %c32 = arith.constant 32 : index - %v_i32 = arith.constant 1 : i32 - - %tile = pto.alloc_tile : - !pto.tile_buf - - %dst_view = pto.make_tensor_view %dst_ptr, - shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view - %dst = pto.partition_view %dst_view, - offsets = [%c0, %c0], sizes = [%c1, %c32] - : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf32> - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - - scf.for %i = %c0 to %c2 step %c1 { - pto.tstore ins(%tile : !pto.tile_buf) - outs(%dst : !pto.partition_tensor_view<1x32xf32>) - pto.cmo.cacheinvalid %dst single_cache_line : !pto.partition_tensor_view<1x32xf32> - pto.fence.barrier_all #pto.fence_scope - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - } - return - } - - func.func @loop_scalar_store_release_tnotify( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c2 = arith.constant 2 : index - %v_i32 = arith.constant 1 : i32 - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - - scf.for %i = %c0 to %c2 step %c1 { - pto.store_scalar %v_i32, %payload_ptr[%c0] : !pto.ptr, i32 - pto.cmo.cacheinvalid %payload_ptr single_cache_line : !pto.ptr - pto.fence.barrier_all #pto.fence_scope - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - } - return - } -} - -// CHECK-LABEL: AICORE void loop_tstore_release_tnotify( -// CHECK: TSTORE( -// CHECK: pipe_barrier(PIPE_MTE3); -// CHECK-NEXT: dsb(DSB_DDR); -// CHECK: pto::comm::TNOTIFY( - -// CHECK-LABEL: AICORE void loop_scalar_store_release_tnotify( -// CHECK: {{.*}}[{{.*}}] = {{.*}}; -// CHECK: PTOAS__DCCI_SINGLE_CACHE_LINE({{.*}}); -// CHECK-NEXT: dsb(DSB_DDR); -// CHECK: pto::comm::TNOTIFY( diff --git a/test/lit/pto/memory_consistency_noninline_call_invalid.pto b/test/lit/pto/memory_consistency_noninline_call_invalid.pto deleted file mode 100644 index e6711173c3..0000000000 --- a/test/lit/pto/memory_consistency_noninline_call_invalid.pto +++ /dev/null @@ -1,61 +0,0 @@ -// 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: not ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s - -// Non-inlined calls are not context-sensitive: a caller-side TNotify cannot -// safely observe release-relevant payload writes hidden in a callee body. Such -// callees must be inlined before the memory consistency pass. - -module { - func.func private @producer( - %tile: !pto.tile_buf, - %dst_ptr: !pto.ptr) { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c32 = arith.constant 32 : index - %dst_view = pto.make_tensor_view %dst_ptr, - shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view - %dst = pto.partition_view %dst_view, - offsets = [%c0, %c0], sizes = [%c1, %c32] - : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf32> - pto.tstore ins(%tile : !pto.tile_buf) - outs(%dst : !pto.partition_tensor_view<1x32xf32>) - return - } - - func.func @call_hidden_payload_write( - %dst_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c32 = arith.constant 32 : index - %v_i32 = arith.constant 1 : i32 - - %tile = pto.alloc_tile : - !pto.tile_buf - - func.call @producer(%tile, %dst_ptr) : - (!pto.tile_buf, - !pto.ptr) -> () - pto.fence.barrier_all #pto.fence_scope - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - return - } -} - -// CHECK: calls @producer, which contains PTO memory consistency relevant operations -// CHECK: inline the callee before `pto-memory-consistency` diff --git a/test/lit/pto/signal_payload_cache_consistency.pto b/test/lit/pto/signal_payload_cache_consistency.pto index 640e16e755..3a6fa93191 100644 --- a/test/lit/pto/signal_payload_cache_consistency.pto +++ b/test/lit/pto/signal_payload_cache_consistency.pto @@ -6,321 +6,42 @@ // 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. -// Signal/payload memory-consistency regressions for the two public IRs kept by -// this PR: GM fence and cache invalidation. +// Explicit signal/payload memory-consistency IR lowering. PTOAS no longer runs +// the historical MemoryConsistency analysis pass, so this test only checks the +// public CMO/fence ops. The GM fence lowering conservatively drains MTE2, +// MTE3, and FIX before dsb(DSB_DDR). // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s module { - func.func @explicit_gm_fence_acquire() + func.func @explicit_gm_fence() attributes {pto.kernel_kind = #pto.kernel_kind} { pto.fence.barrier_all #pto.fence_scope return } - func.func @twait_load_scalar_acquire( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) + func.func @whole_cache_cmo() attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %v_i32 = arith.constant 1 : i32 - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.twait(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {cmp = #pto} - pto.cmo.cacheinvalid all #pto.address_space - %val = pto.load_scalar %payload_ptr[%c0] : !pto.ptr -> i32 - pto.store_scalar %val, %payload_ptr[%c0] : !pto.ptr, i32 return } - func.func @ttest_load_scalar_conservative_acquire( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) + func.func @single_line_ptr_cmo(%payload_ptr: !pto.ptr) attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %v_i32 = arith.constant 1 : i32 - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - %ready = pto.comm.ttest(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {cmp = #pto} -> i1 - - pto.cmo.cacheinvalid all #pto.address_space - %val = pto.load_scalar %payload_ptr[%c0] : !pto.ptr -> i32 - scf.if %ready { - pto.store_scalar %val, %payload_ptr[%c0] : !pto.ptr, i32 - } - return - } - - func.func @ttest_true_branch_acquire( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %v_i32 = arith.constant 1 : i32 - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - %ready = pto.comm.ttest(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {cmp = #pto} -> i1 - - scf.if %ready { - pto.cmo.cacheinvalid all #pto.address_space - %val = pto.load_scalar %payload_ptr[%c0] : !pto.ptr -> i32 - pto.store_scalar %val, %payload_ptr[%c0] : !pto.ptr, i32 - } - return - } - - func.func @ttest_false_branch_conservative_acquire( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %v_i32 = arith.constant 1 : i32 - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - %ready = pto.comm.ttest(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {cmp = #pto} -> i1 - - scf.if %ready { - } else { - pto.cmo.cacheinvalid all #pto.address_space - %val = pto.load_scalar %payload_ptr[%c0] : !pto.ptr -> i32 - pto.store_scalar %val, %payload_ptr[%c0] : !pto.ptr, i32 - } - return - } - - func.func @ttest_polling_loop_acquire( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %true = arith.constant true - %v_i32 = arith.constant 1 : i32 - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - - %done = scf.while (%keep_polling = %true) : (i1) -> i1 { - %ready = pto.comm.ttest(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {cmp = #pto} -> i1 - %not_ready = arith.xori %ready, %true : i1 - scf.condition(%not_ready) %not_ready : i1 - } do { - ^bb0(%keep_polling_body: i1): - scf.yield %keep_polling_body : i1 - } - - pto.cmo.cacheinvalid all #pto.address_space - %val = pto.load_scalar %payload_ptr[%c0] : !pto.ptr -> i32 - scf.if %done { - pto.store_scalar %val, %payload_ptr[%c0] : !pto.ptr, i32 - } - return - } - - func.func @store_scalar_tnotify_release( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %v_i32 = arith.constant 7 : i32 - - pto.store_scalar %v_i32, %payload_ptr[%c0] : !pto.ptr, i32 pto.cmo.cacheinvalid %payload_ptr single_cache_line : !pto.ptr - pto.fence.barrier_all #pto.fence_scope - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - return - } - - func.func @store_scalar_whole_cache_marker_tnotify_release( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %v_i32 = arith.constant 7 : i32 - - pto.store_scalar %v_i32, %payload_ptr[%c0] : !pto.ptr, i32 - pto.cmo.cacheinvalid all #pto.address_space - pto.fence.barrier_all #pto.fence_scope - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - return - } - - func.func @whole_cache_tstore_tnotify_release( - %src_ptr: !pto.ptr, - %dst_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c32 = arith.constant 32 : index - %v_i32 = arith.constant 7 : i32 - - %tile = pto.alloc_tile : - !pto.tile_buf - - %src_view = pto.make_tensor_view %src_ptr, - shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view - %src = pto.partition_view %src_view, - offsets = [%c0, %c0], sizes = [%c1, %c32] - : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf32> - %dst_view = pto.make_tensor_view %dst_ptr, - shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view - %dst = pto.partition_view %dst_view, - offsets = [%c0, %c0], sizes = [%c1, %c32] - : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf32> - - pto.tload ins(%src : !pto.partition_tensor_view<1x32xf32>) - outs(%tile : !pto.tile_buf) - pto.barrier - pto.tstore ins(%tile : !pto.tile_buf) - outs(%dst : !pto.partition_tensor_view<1x32xf32>) - pto.cmo.cacheinvalid all #pto.address_space - pto.fence.barrier_all #pto.fence_scope - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - return - } - - func.func @whole_cache_acc_tstore_tnotify_release( - %dst_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c16 = arith.constant 16 : index - %v_i32 = arith.constant 7 : i32 - - %acc = pto.alloc_tile : - !pto.tile_buf - - %dst_view = pto.make_tensor_view %dst_ptr, - shape = [%c16, %c16], strides = [%c16, %c1] : !pto.tensor_view - %dst = pto.partition_view %dst_view, - offsets = [%c0, %c0], sizes = [%c16, %c16] - : !pto.tensor_view -> !pto.partition_tensor_view<16x16xf32> - - pto.tstore ins(%acc : !pto.tile_buf) - outs(%dst : !pto.partition_tensor_view<16x16xf32>) - pto.cmo.cacheinvalid all #pto.address_space - pto.fence.barrier_all #pto.fence_scope - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - return - } - - func.func @whole_cache_tload_tnotify_release( - %src_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c32 = arith.constant 32 : index - %v_i32 = arith.constant 7 : i32 - - %tile = pto.alloc_tile : - !pto.tile_buf - - %src_view = pto.make_tensor_view %src_ptr, - shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view - %src = pto.partition_view %src_view, - offsets = [%c0, %c0], sizes = [%c1, %c32] - : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf32> - - pto.tload ins(%src : !pto.partition_tensor_view<1x32xf32>) - outs(%tile : !pto.tile_buf) - pto.cmo.cacheinvalid all #pto.address_space - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} return } - func.func @mixed_scalar_mte_release_tnotify( + func.func @scalar_release_sequence( %payload_ptr: !pto.ptr, - %dst_ptr: !pto.ptr, %signal_ptr: !pto.ptr) attributes {pto.kernel_kind = #pto.kernel_kind} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index - %c32 = arith.constant 32 : index %v_i32 = arith.constant 7 : i32 - %tile = pto.alloc_tile : - !pto.tile_buf - - %dst_view = pto.make_tensor_view %dst_ptr, - shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view - %dst = pto.partition_view %dst_view, - offsets = [%c0, %c0], sizes = [%c1, %c32] - : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf32> - pto.store_scalar %v_i32, %payload_ptr[%c0] : !pto.ptr, i32 - pto.tstore ins(%tile : !pto.tile_buf) - outs(%dst : !pto.partition_tensor_view<1x32xf32>) pto.cmo.cacheinvalid %payload_ptr single_cache_line : !pto.ptr - pto.cmo.cacheinvalid %dst single_cache_line : !pto.partition_tensor_view<1x32xf32> pto.fence.barrier_all #pto.fence_scope %sig_view = pto.make_tensor_view %signal_ptr, @@ -333,7 +54,7 @@ module { return } - func.func @single_line_twait_load_scalar_acquire( + func.func @scalar_acquire_sequence( %payload_ptr: !pto.ptr, %signal_ptr: !pto.ptr) attributes {pto.kernel_kind = #pto.kernel_kind} { @@ -355,29 +76,7 @@ module { return } - func.func @single_line_store_scalar_tnotify_release( - %payload_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %v_i32 = arith.constant 7 : i32 - - pto.store_scalar %v_i32, %payload_ptr[%c0] : !pto.ptr, i32 - pto.cmo.cacheinvalid %payload_ptr single_cache_line : !pto.ptr - pto.fence.barrier_all #pto.fence_scope - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - return - } - - func.func @single_line_tstore_marker_skips_dcci_release( + func.func @mte_release_sequence( %src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, %signal_ptr: !pto.ptr) @@ -395,48 +94,7 @@ module { %src = pto.partition_view %src_view, offsets = [%c0, %c0], sizes = [%c1, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf32> - %dst_view = pto.make_tensor_view %dst_ptr, - shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view - %dst = pto.partition_view %dst_view, - offsets = [%c0, %c0], sizes = [%c1, %c32] - : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf32> - - pto.tload ins(%src : !pto.partition_tensor_view<1x32xf32>) - outs(%tile : !pto.tile_buf) - pto.barrier - pto.tstore ins(%tile : !pto.tile_buf) - outs(%dst : !pto.partition_tensor_view<1x32xf32>) - pto.cmo.cacheinvalid %dst single_cache_line : !pto.partition_tensor_view<1x32xf32> - pto.fence.barrier_all #pto.fence_scope - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - return - } - - func.func @single_line_tstore_marker_before_producer_release( - %src_ptr: !pto.ptr, - %dst_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c32 = arith.constant 32 : index - %v_i32 = arith.constant 7 : i32 - %tile = pto.alloc_tile : - !pto.tile_buf - - %src_view = pto.make_tensor_view %src_ptr, - shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view - %src = pto.partition_view %src_view, - offsets = [%c0, %c0], sizes = [%c1, %c32] - : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf32> %dst_view = pto.make_tensor_view %dst_ptr, shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view %dst = pto.partition_view %dst_view, @@ -445,44 +103,8 @@ module { pto.tload ins(%src : !pto.partition_tensor_view<1x32xf32>) outs(%tile : !pto.tile_buf) - pto.barrier - pto.cmo.cacheinvalid %dst single_cache_line : !pto.partition_tensor_view<1x32xf32> - pto.tstore ins(%tile : !pto.tile_buf) - outs(%dst : !pto.partition_tensor_view<1x32xf32>) - pto.fence.barrier_all #pto.fence_scope - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - return - } - - func.func @single_line_unrelated_marker_no_mte3_drain( - %dst_ptr: !pto.ptr, - %other_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c32 = arith.constant 32 : index - %v_i32 = arith.constant 7 : i32 - - %tile = pto.alloc_tile : - !pto.tile_buf - - %dst_view = pto.make_tensor_view %dst_ptr, - shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view - %dst = pto.partition_view %dst_view, - offsets = [%c0, %c0], sizes = [%c1, %c32] - : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf32> - pto.tstore ins(%tile : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x32xf32>) - pto.cmo.cacheinvalid %other_ptr single_cache_line : !pto.ptr pto.fence.barrier_all #pto.fence_scope %sig_view = pto.make_tensor_view %signal_ptr, @@ -494,150 +116,38 @@ module { {notifyOp = #pto} return } - - func.func @single_line_tload_marker_skips_dcci_tnotify( - %src_ptr: !pto.ptr, - %signal_ptr: !pto.ptr) - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c32 = arith.constant 32 : index - %v_i32 = arith.constant 7 : i32 - - %tile = pto.alloc_tile : - !pto.tile_buf - - %src_view = pto.make_tensor_view %src_ptr, - shape = [%c1, %c32], strides = [%c32, %c1] : !pto.tensor_view - %src = pto.partition_view %src_view, - offsets = [%c0, %c0], sizes = [%c1, %c32] - : !pto.tensor_view -> !pto.partition_tensor_view<1x32xf32> - - pto.tload ins(%src : !pto.partition_tensor_view<1x32xf32>) - outs(%tile : !pto.tile_buf) - pto.cmo.cacheinvalid %src single_cache_line : !pto.partition_tensor_view<1x32xf32> - - %sig_view = pto.make_tensor_view %signal_ptr, - shape = [%c1], strides = [%c1] : !pto.tensor_view<1xi32> - %sig = pto.partition_view %sig_view, - offsets = [%c0], sizes = [%c1] - : !pto.tensor_view<1xi32> -> !pto.partition_tensor_view<1xi32> - pto.comm.tnotify(%sig, %v_i32 : !pto.partition_tensor_view<1xi32>, i32) - {notifyOp = #pto} - return - } } -// CHECK-LABEL: AICORE void explicit_gm_fence_acquire( -// CHECK: dsb(DSB_DDR); - -// CHECK-LABEL: AICORE void twait_load_scalar_acquire( -// CHECK: pto::comm::TWAIT( -// CHECK: dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE); -// CHECK: {{.*}} = {{.*}}[{{.*}}]; - -// CHECK-LABEL: AICORE void ttest_load_scalar_conservative_acquire( -// CHECK: pto::comm::TTEST( -// CHECK: dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE); -// CHECK: {{.*}} = {{.*}}[{{.*}}]; - -// CHECK-LABEL: AICORE void ttest_true_branch_acquire( -// CHECK: pto::comm::TTEST( -// CHECK: if ( -// CHECK: dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE); -// CHECK: {{.*}} = {{.*}}[{{.*}}]; - -// CHECK-LABEL: AICORE void ttest_false_branch_conservative_acquire( -// CHECK: pto::comm::TTEST( -// CHECK: else -// CHECK: dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE); -// CHECK: {{.*}} = {{.*}}[{{.*}}]; - -// CHECK-LABEL: AICORE void ttest_polling_loop_acquire( -// CHECK: pto::comm::TTEST( -// CHECK: dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE); -// CHECK: {{.*}} = {{.*}}[{{.*}}]; - -// CHECK-LABEL: AICORE void store_scalar_tnotify_release( -// CHECK: {{.*}}[{{.*}}] = {{.*}}; -// CHECK: PTOAS__DCCI_SINGLE_CACHE_LINE({{.*}}); -// CHECK-NEXT: dsb(DSB_DDR); -// CHECK: pto::comm::TNOTIFY( - -// CHECK-LABEL: AICORE void store_scalar_whole_cache_marker_tnotify_release( -// CHECK: {{.*}}[{{.*}}] = {{.*}}; -// CHECK: dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE); -// CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE -// CHECK-NEXT: dsb(DSB_DDR); -// CHECK: pto::comm::TNOTIFY( - -// CHECK-LABEL: AICORE void whole_cache_tstore_tnotify_release( -// CHECK: TSTORE( -// CHECK: dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE); -// CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE -// CHECK: pipe_barrier(PIPE_MTE3); +// CHECK-LABEL: AICORE void explicit_gm_fence( +// CHECK: pipe_barrier(PIPE_MTE2); +// CHECK-NEXT: pipe_barrier(PIPE_MTE3); +// CHECK-NEXT: pipe_barrier(PIPE_FIX); // CHECK-NEXT: dsb(DSB_DDR); -// CHECK: pto::comm::TNOTIFY( -// CHECK-LABEL: AICORE void whole_cache_acc_tstore_tnotify_release( -// CHECK: TSTORE +// CHECK-LABEL: AICORE void whole_cache_cmo( // CHECK: dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE); -// CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE -// CHECK: pipe_barrier(PIPE_FIX); -// CHECK-NEXT: dsb(DSB_DDR); -// CHECK: pto::comm::TNOTIFY( -// CHECK-LABEL: AICORE void whole_cache_tload_tnotify_release( -// CHECK: TLOAD( -// CHECK: dcci((__gm__ void*)0, cache_line_t::ENTIRE_DATA_CACHE); -// CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE -// CHECK: pipe_barrier(PIPE_MTE2); -// CHECK-NOT: dsb( -// CHECK: pto::comm::TNOTIFY( +// CHECK-LABEL: AICORE void single_line_ptr_cmo( +// CHECK: PTOAS__DCCI_SINGLE_CACHE_LINE({{.*}}); -// CHECK-LABEL: AICORE void mixed_scalar_mte_release_tnotify( +// CHECK-LABEL: AICORE void scalar_release_sequence( // CHECK: {{.*}}[{{.*}}] = {{.*}}; -// CHECK: TSTORE( // CHECK: PTOAS__DCCI_SINGLE_CACHE_LINE({{.*}}); +// CHECK-NEXT: pipe_barrier(PIPE_MTE2); // CHECK-NEXT: pipe_barrier(PIPE_MTE3); +// CHECK-NEXT: pipe_barrier(PIPE_FIX); // CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( -// CHECK-LABEL: AICORE void single_line_twait_load_scalar_acquire( +// CHECK-LABEL: AICORE void scalar_acquire_sequence( // CHECK: pto::comm::TWAIT( // CHECK: PTOAS__DCCI_SINGLE_CACHE_LINE({{.*}}); // CHECK: {{.*}} = {{.*}}[{{.*}}]; -// CHECK-LABEL: AICORE void single_line_store_scalar_tnotify_release( -// CHECK: {{.*}}[{{.*}}] = {{.*}}; -// CHECK: PTOAS__DCCI_SINGLE_CACHE_LINE({{.*}}); -// CHECK-NEXT: dsb(DSB_DDR); -// CHECK: pto::comm::TNOTIFY( - -// CHECK-LABEL: AICORE void single_line_tstore_marker_skips_dcci_release( +// CHECK-LABEL: AICORE void mte_release_sequence( // CHECK: TSTORE( -// CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE -// CHECK: pipe_barrier(PIPE_MTE3); -// CHECK-NEXT: dsb(DSB_DDR); -// CHECK: pto::comm::TNOTIFY( - -// CHECK-LABEL: AICORE void single_line_tstore_marker_before_producer_release( -// CHECK: TSTORE( -// CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE -// CHECK: pipe_barrier(PIPE_MTE3); -// CHECK-NEXT: dsb(DSB_DDR); -// CHECK: pto::comm::TNOTIFY( - -// CHECK-LABEL: AICORE void single_line_unrelated_marker_no_mte3_drain( -// CHECK: TSTORE( -// CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE -// CHECK-NOT: pipe_barrier(PIPE_MTE3) -// CHECK: dsb(DSB_DDR); -// CHECK: pto::comm::TNOTIFY( - -// CHECK-LABEL: AICORE void single_line_tload_marker_skips_dcci_tnotify( -// CHECK: TLOAD( -// CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE // CHECK: pipe_barrier(PIPE_MTE2); -// CHECK-NOT: dsb( +// CHECK-NEXT: pipe_barrier(PIPE_MTE3); +// CHECK-NEXT: pipe_barrier(PIPE_FIX); +// CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( diff --git a/test/lit/pto/sync_set_a5_emitc_physical_ids.pto b/test/lit/pto/sync_set_a5_emitc_physical_ids.pto new file mode 100644 index 0000000000..6f23f96906 --- /dev/null +++ b/test/lit/pto/sync_set_a5_emitc_physical_ids.pto @@ -0,0 +1,22 @@ +// 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 %s 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + func.func @sync_set_a5_emitc_physical_ids() attributes {pto.kernel} { + pto.sync.set , 0 + pto.sync.set , 16 + return + } +} + +// CHECK-LABEL: AICORE void sync_set_a5_emitc_physical_ids() +// CHECK: set_intra_block(PIPE_FIX, 0); +// CHECK-NEXT: set_intra_block(PIPE_FIX, 16); +// CHECK-NOT: set_intra_block(PIPE_FIX, 32); diff --git a/test/lit/vpto/acc_store_gm_f32_bf16.pto b/test/lit/vpto/acc_store_gm_f32_bf16.pto new file mode 100644 index 0000000000..378f09d673 --- /dev/null +++ b/test/lit/vpto/acc_store_gm_f32_bf16.pto @@ -0,0 +1,36 @@ +// 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 + +module attributes {"pto.target_arch" = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @acc_store_gm_f32_bf16(%src: !pto.ptr, + %dst: !pto.ptr) { + %c0_i64 = arith.constant 0 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c1_bf16 = arith.constant 1.000000e+00 : bf16 + + pto.mte_l0c_gm %src, %dst, %c16_i64, %c16_i64, %c16_i64, %c32_i64, + %c0_i64, %c0_i64, + pre_quant(%c1_bf16, mode = f32_bf16), + nz2nd + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64, + bf16 + + return + } +} + +// ROUNDTRIP-LABEL: func.func @acc_store_gm_f32_bf16( +// ROUNDTRIP: pto.mte_l0c_gm %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, pre_quant(%{{.*}}, mode = f32_bf16), nz2nd + +// EXPAND-LABEL: func.func @acc_store_gm_f32_bf16( +// EXPAND: pto.set_quant_pre +// EXPAND: pto.copy_matrix_cc_to_gm diff --git a/test/lit/vpto/acc_store_verify_invalid_f32_bf16_dst.pto b/test/lit/vpto/acc_store_verify_invalid_f32_bf16_dst.pto new file mode 100644 index 0000000000..2a882380dd --- /dev/null +++ b/test/lit/vpto/acc_store_verify_invalid_f32_bf16_dst.pto @@ -0,0 +1,30 @@ +// 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: not ptoas --pto-arch=a5 %s -o - 2>&1 | FileCheck %s + +module attributes {"pto.target_arch" = "a5"} { + func.func @acc_store_verify_invalid_f32_bf16_dst( + %src: !pto.ptr, %dst: !pto.ptr) { + %c0_i64 = arith.constant 0 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c1_bf16 = arith.constant 1.000000e+00 : bf16 + + pto.mte_l0c_gm %src, %dst, %c16_i64, %c16_i64, %c16_i64, %c32_i64, + %c0_i64, %c0_i64, + pre_quant(%c1_bf16, mode = f32_bf16), + nz2nd + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64, + bf16 + + return + } +} + +// CHECK: pre_quant mode f32_bf16 is incompatible with source element type 'f32' and destination element type 'f16' diff --git a/test/lit/vpto/acc_store_verify_invalid_f32_bf16_vec_payload.pto b/test/lit/vpto/acc_store_verify_invalid_f32_bf16_vec_payload.pto new file mode 100644 index 0000000000..f0b46d6ca2 --- /dev/null +++ b/test/lit/vpto/acc_store_verify_invalid_f32_bf16_vec_payload.pto @@ -0,0 +1,30 @@ +// 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: not ptoas --pto-arch=a5 %s -o - 2>&1 | FileCheck %s + +module attributes {"pto.target_arch" = "a5"} { + func.func @acc_store_verify_invalid_f32_bf16_vec_payload( + %src: !pto.ptr, %dst: !pto.ptr, + %payload: !pto.ptr) { + %c0_i64 = arith.constant 0 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + + pto.mte_l0c_gm %src, %dst, %c16_i64, %c16_i64, %c16_i64, %c32_i64, + %c0_i64, %c0_i64, + pre_quant(%payload, mode = f32_bf16), + nz2nd + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64, + !pto.ptr + + return + } +} + +// CHECK: scalar pre_quant mode requires f16/bf16/f32 payload diff --git a/test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto b/test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto new file mode 100644 index 0000000000..f91b6a0389 --- /dev/null +++ b/test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto @@ -0,0 +1,54 @@ +// 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 + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @mte_ub_gm_l2_cache_ctl(%out: !pto.ptr) attributes {pto.aicore} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c5_i64 = arith.constant 5 : i64 + %c64_i64 = arith.constant 64 : i64 + %ub = pto.castptr %c0_i64 : i64 -> !pto.ptr + + pto.mte_ub_gm %ub, %out, %c64_i64 + nburst(%c1_i64, %c64_i64, %c64_i64) l2_cache_ctl(%c5_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64 + return + } + + func.func @mte_ub_gm_l2_cache_ctl_default(%out: !pto.ptr) attributes {pto.aicore} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c64_i64 = arith.constant 64 : i64 + %ub = pto.castptr %c0_i64 : i64 -> !pto.ptr + + pto.mte_ub_gm %ub, %out, %c64_i64 + nburst(%c1_i64, %c64_i64, %c64_i64) + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 + return + } +} + +// ROUNDTRIP-LABEL: func.func @mte_ub_gm_l2_cache_ctl +// ROUNDTRIP: %[[L2:.*]] = arith.constant 5 : i64 +// ROUNDTRIP: pto.mte_ub_gm %{{[^,]+}}, %arg0, %{{[^ ]+}} nburst(%{{[^,]+}}, %{{[^,]+}}, %{{[^)]+}}) l2_cache_ctl(%[[L2]]) +// ROUNDTRIP-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64{{$}} +// ROUNDTRIP-LABEL: func.func @mte_ub_gm_l2_cache_ctl_default +// ROUNDTRIP: pto.mte_ub_gm %{{[^,]+}}, %arg0, %{{[^ ]+}} nburst(%{{[^,]+}}, %{{[^,]+}}, %{{[^)]+}}) +// ROUNDTRIP-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64{{$}} + +// EXPAND-LABEL: func.func @mte_ub_gm_l2_cache_ctl +// EXPAND: %[[L2:.*]] = arith.constant 5 : i64 +// EXPAND: pto.copy_ubuf_to_gm %{{[^,]+}}, %arg0, %{{[^,]+}}, %{{[^,]+}}, %{{[^,]+}}, %[[L2]], %{{[^,]+}}, %{{[^:]+}} +// EXPAND-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 +// EXPAND-LABEL: func.func @mte_ub_gm_l2_cache_ctl_default +// EXPAND: %[[ZERO:.*]] = arith.constant 0 : i64 +// EXPAND: pto.copy_ubuf_to_gm %{{[^,]+}}, %arg0, %[[ZERO]], %{{[^,]+}}, %{{[^,]+}}, %[[ZERO]], %{{[^,]+}}, %{{[^:]+}} +// EXPAND-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 diff --git a/test/lit/vpto/simt_lowlevel_launch_config_vpto_llvm.pto b/test/lit/vpto/simt_lowlevel_launch_config_vpto_llvm.pto index 99893cdc07..d9568b786a 100644 --- a/test/lit/vpto/simt_lowlevel_launch_config_vpto_llvm.pto +++ b/test/lit/vpto/simt_lowlevel_launch_config_vpto_llvm.pto @@ -18,7 +18,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) attributes {pto.simt_entry, pto.simt_max_threads = 256 : i32, pto.simt_max_regs = 48 : i32} { + func.func @launch_config_body(%dst: !pto.ptr) attributes {pto.simt_entry, pto.simt_max_threads = 256 : i32} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %tid_x = pto.get_tid_x : i32 @@ -30,5 +30,4 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind&1 | FileCheck %s +// RUN: ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-backend=vpto --emit-vpto-llvm-ir %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @kernel() attributes {pto.aicore} { + return + } + + func.func @default_config() attributes {pto.simt_entry} { + return + } + + func.func @threads_256() attributes {pto.simt_entry, pto.simt_max_threads = 256 : i32} { + return + } + + func.func @threads_257() attributes {pto.simt_entry, pto.simt_max_threads = 257 : i32} { + return + } + + func.func @threads_512() attributes {pto.simt_entry, pto.simt_max_threads = 512 : i32} { + return + } + + func.func @threads_513() attributes {pto.simt_entry, pto.simt_max_threads = 513 : i32} { + return + } + + func.func @threads_1024() attributes {pto.simt_entry, pto.simt_max_threads = 1024 : i32} { + return + } + + func.func @threads_1025() attributes {pto.simt_entry, pto.simt_max_threads = 1025 : i32} { + return + } + + func.func @threads_2048() attributes {pto.simt_entry, pto.simt_max_threads = 2048 : i32} { + return + } +} + +// CHECK-LABEL: define linkonce_odr simt_entry void @default_config() +// CHECK-SAME: !annotation ![[DEFAULT_THREADS:[0-9]+]] !annotation ![[DEFAULT_REGS:[0-9]+]] +// CHECK-LABEL: define linkonce_odr simt_entry void @threads_2048() +// CHECK-SAME: !annotation ![[THREADS_2048:[0-9]+]] !annotation ![[REGS_2048:[0-9]+]] + +// The module annotations derive one thread/register pair per SIMT entry. +// CHECK: distinct !{null, !"simt-max-threads", i32 1024} +// CHECK-NEXT: distinct !{null, !"simt-max-registers", i32 32} +// CHECK-NEXT: distinct !{null, !"simt-max-threads", i32 256} +// CHECK-NEXT: distinct !{null, !"simt-max-registers", i32 128} +// CHECK-NEXT: distinct !{null, !"simt-max-threads", i32 257} +// CHECK-NEXT: distinct !{null, !"simt-max-registers", i32 64} +// CHECK-NEXT: distinct !{null, !"simt-max-threads", i32 512} +// CHECK-NEXT: distinct !{null, !"simt-max-registers", i32 64} +// CHECK-NEXT: distinct !{null, !"simt-max-threads", i32 513} +// CHECK-NEXT: distinct !{null, !"simt-max-registers", i32 32} +// CHECK-NEXT: distinct !{null, !"simt-max-threads", i32 1024} +// CHECK-NEXT: distinct !{null, !"simt-max-registers", i32 32} +// CHECK-NEXT: distinct !{null, !"simt-max-threads", i32 1025} +// CHECK-NEXT: distinct !{null, !"simt-max-registers", i32 16} +// CHECK-NEXT: distinct !{null, !"simt-max-threads", i32 2048} +// CHECK-NEXT: distinct !{null, !"simt-max-registers", i32 16} + +// CHECK-DAG: ![[DEFAULT_THREADS]] = !{!"simt-max-threads", i32 1024} +// CHECK-DAG: ![[DEFAULT_REGS]] = !{!"simt-max-registers", i32 32} +// CHECK-DAG: ![[THREADS_2048]] = !{!"simt-max-threads", i32 2048} +// CHECK-DAG: ![[REGS_2048]] = !{!"simt-max-registers", i32 16} diff --git a/test/samples/Sync/test_intercore_sync_a5.py b/test/samples/Sync/test_intercore_sync_a5.py index c0355d7b1d..c3f9607ed4 100644 --- a/test/samples/Sync/test_intercore_sync_a5.py +++ b/test/samples/Sync/test_intercore_sync_a5.py @@ -48,6 +48,7 @@ def build(): sec_cube = pto.SectionCubeOp() with InsertionPoint(sec_cube.body.blocks.append()): pto.sync_set(pipe_fix, evt) + pto.sync_set(pipe_fix, evt + 16) sec_vec = pto.SectionVectorOp() with InsertionPoint(sec_vec.body.blocks.append()): diff --git a/test/samples/Sync/test_intercore_sync_a5_dyn.py b/test/samples/Sync/test_intercore_sync_a5_dyn.py index 71f75a788a..cc7b9d7535 100644 --- a/test/samples/Sync/test_intercore_sync_a5_dyn.py +++ b/test/samples/Sync/test_intercore_sync_a5_dyn.py @@ -29,6 +29,8 @@ def build(): with InsertionPoint(entry): c0 = arith.ConstantOp(idx, 0).result c0_evt = arith.ConstantOp(idx, 0).result + c16_evt = arith.ConstantOp(idx, 16).result + c16_evt_pair = arith.AddIOp(c0_evt, c16_evt).result two = arith.ConstantOp(f32, 2.0).result out = entry.arguments[0] pipe_fix = pto.PipeAttr.get(pto.PIPE.PIPE_FIX, ctx) @@ -37,6 +39,7 @@ def build(): sec_cube = pto.SectionCubeOp() with InsertionPoint(sec_cube.body.blocks.append()): pto.sync_set(pipe_fix, c0_evt) + pto.sync_set(pipe_fix, c16_evt_pair) sec_vec = pto.SectionVectorOp() with InsertionPoint(sec_vec.body.blocks.append()): diff --git a/test/samples/Sync/test_intercore_sync_a5_functional.py b/test/samples/Sync/test_intercore_sync_a5_functional.py index 328710e065..77d37495e5 100644 --- a/test/samples/Sync/test_intercore_sync_a5_functional.py +++ b/test/samples/Sync/test_intercore_sync_a5_functional.py @@ -43,6 +43,7 @@ def build(): sec_cube = pto.SectionCubeOp() with InsertionPoint(sec_cube.body.blocks.append()): pto.sync_set(pipe_fix, evt) + pto.sync_set(pipe_fix, evt + 16) sec_vec = pto.SectionVectorOp() with InsertionPoint(sec_vec.body.blocks.append()): diff --git a/test/samples/Sync/test_intercore_sync_a5_ptoisa_vec.py b/test/samples/Sync/test_intercore_sync_a5_ptoisa_vec.py index f77817fc42..68919478f1 100644 --- a/test/samples/Sync/test_intercore_sync_a5_ptoisa_vec.py +++ b/test/samples/Sync/test_intercore_sync_a5_ptoisa_vec.py @@ -34,7 +34,8 @@ def build(): c2 = arith.ConstantOp(f32, 2.0).result # PTO-ISA tmov_acc2vec/tmov_ub2l1 mix-kernel style: - # cube sets PIPE_FIX(syncId + syncId+16), vec waits PIPE_MTE3(syncId). + # cube explicitly sets PIPE_FIX(syncId) and PIPE_FIX(syncId+16); + # vec waits PIPE_MTE3(syncId). sync_id = 0 pipe_fix = pto.PipeAttr.get(pto.PIPE.PIPE_FIX, ctx) pipe_mte3 = pto.PipeAttr.get(pto.PIPE.PIPE_MTE3, ctx) @@ -42,6 +43,7 @@ def build(): sec_cube = pto.SectionCubeOp() with InsertionPoint(sec_cube.body.blocks.append()): pto.sync_set(pipe_fix, sync_id) + pto.sync_set(pipe_fix, sync_id + 16) sec_vec = pto.SectionVectorOp() with InsertionPoint(sec_vec.body.blocks.append()): diff --git a/test/samples/runop.sh b/test/samples/runop.sh index 9a48538bd5..7585fcd72e 100755 --- a/test/samples/runop.sh +++ b/test/samples/runop.sh @@ -642,7 +642,7 @@ process_one_dir() { continue fi if ! grep -Fq "set_intra_block(PIPE_FIX, 0)" "$cpp" || ! grep -Fq "set_intra_block(PIPE_FIX, 16)" "$cpp"; then - echo -e "${A}(${base}.py)\tFAIL\tmissing A5 cube-side mirrored set_intra_block(PIPE_FIX, id/id+16)" + echo -e "${A}(${base}.py)\tFAIL\tmissing A5 cube-side explicit set_intra_block(PIPE_FIX, id/id+16)" overall=1 continue fi @@ -664,7 +664,7 @@ process_one_dir() { continue fi if ! grep -Fq "set_intra_block(PIPE_FIX, 0)" "$cpp" || ! grep -Fq "set_intra_block(PIPE_FIX, 16)" "$cpp"; then - echo -e "${A}(${base}.py)\tFAIL\tmissing A5 cube-side mirrored set_intra_block(PIPE_FIX, id/id+16)" + echo -e "${A}(${base}.py)\tFAIL\tmissing A5 cube-side explicit set_intra_block(PIPE_FIX, id/id+16)" overall=1 continue fi @@ -686,7 +686,7 @@ process_one_dir() { continue fi if ! grep -Fq "set_intra_block(PIPE_FIX, 0)" "$cpp" || ! grep -Fq "set_intra_block(PIPE_FIX, 16)" "$cpp"; then - echo -e "${A}(${base}.py)\tFAIL\tmissing PTO-ISA-style cube-side mirrored set_intra_block(PIPE_FIX, id/id+16)" + echo -e "${A}(${base}.py)\tFAIL\tmissing PTO-ISA-style explicit set_intra_block(PIPE_FIX, id/id+16)" overall=1 continue fi @@ -753,7 +753,7 @@ process_one_dir() { continue fi if [[ "$set_count" -ne 2 ]]; then - echo -e "${A}(${base}.py)\tFAIL\tunexpected number of PIPE_FIX dynamic sync.set calls (expect 2: id and id+16)" + echo -e "${A}(${base}.py)\tFAIL\tunexpected number of PIPE_FIX dynamic sync.set calls (expect explicit id and id+16)" overall=1 continue fi diff --git a/test/vpto/cases/micro-op/cv-mixed/fixpipe-acc-store-dual-ub-cv/kernel.pto b/test/vpto/cases/micro-op/cv-mixed/fixpipe-acc-store-dual-ub-cv/kernel.pto index dea1462596..154415007d 100644 --- a/test/vpto/cases/micro-op/cv-mixed/fixpipe-acc-store-dual-ub-cv/kernel.pto +++ b/test/vpto/cases/micro-op/cv-mixed/fixpipe-acc-store-dual-ub-cv/kernel.pto @@ -111,7 +111,7 @@ module attributes {pto.target_arch = "a5"} { : !pto.ptr, !pto.ptr, i64, i64, i64, i64 } scf.if %is_subblock1 { - pto.sync.wait , 1 + pto.sync.wait , 17 pto.mte_ub_gm %ub_split_m, %out_split_m1_gm, %c256_i64 nburst(%c20_i64, %c256_i64, %c256_i64) : !pto.ptr, !pto.ptr, i64, i64, i64, i64 diff --git a/tilelang-dsl/docs/user_guide/08-sync-dma-operations.md b/tilelang-dsl/docs/user_guide/08-sync-dma-operations.md index 44d8a38165..b1da73e9fa 100644 --- a/tilelang-dsl/docs/user_guide/08-sync-dma-operations.md +++ b/tilelang-dsl/docs/user_guide/08-sync-dma-operations.md @@ -388,7 +388,7 @@ pto.mte_gm_ub( ) ``` -#### `pto.mte_ub_gm(ub_src, gm_dst, len_burst, *, nburst, loops=None) -> None` [Advanced Tier] +#### `pto.mte_ub_gm(ub_src, gm_dst, len_burst, *, nburst, loops=None, l2_cache="nmfv") -> None` [Advanced Tier] **Description**: Grouped UB→GM DMA transfer. @@ -400,6 +400,7 @@ pto.mte_gm_ub( | `len_burst` | `pto.i64` | Bytes transferred per burst row | | `nburst` | `tuple[i64, i64, i64]` | Required burst triple `(count, src_stride, dst_stride)` | | `loops` | `tuple[tuple[i64, i64, i64], ...] \| None` | Optional outer loop triples from inner to outer | +| `l2_cache` | `str` | Store-side L2 cache policy token. Defaults to `"nmfv"` (raw control `0`); supported tokens are `"nmfv"`, `"nmlv"`, `"nmprs"`, `"nmred"`, `"naci"`, `"napw"`, `"napi"`, `"nared"`, `"wbhfv"`, `"wbhlv"`, `"wbhprs"`, `"wbhred"`, `"wtsfv"`, `"wtslv"`, `"wtsprs"`, and `"wtsred"` | **Example**: ```python @@ -408,6 +409,7 @@ pto.mte_ub_gm( gm_ptr, 128, nburst=(32, 128, 128), + l2_cache="nmfv", ) ``` diff --git a/tilelang-dsl/python/tilelang_dsl/frontend_ast.py b/tilelang-dsl/python/tilelang_dsl/frontend_ast.py index 5801174339..0db5b7aec6 100644 --- a/tilelang-dsl/python/tilelang_dsl/frontend_ast.py +++ b/tilelang-dsl/python/tilelang_dsl/frontend_ast.py @@ -809,7 +809,7 @@ def _collect_reachable_inline_procs( _DMA_CALL_KEYWORDS: dict[str, frozenset[str]] = { "Tile": frozenset({"valid_shape", "blayout", "slayout", "fractal_size", "pad_value", "compact_mode", "addr"}), "mte_gm_ub": frozenset({"nburst", "loops", "pad"}), - "mte_ub_gm": frozenset({"nburst", "loops"}), + "mte_ub_gm": frozenset({"nburst", "loops", "l2_cache"}), "mte_ub_ub": frozenset({"nburst"}), "mte_ub_l1": frozenset({"nburst"}), "set_mov_pad_val": frozenset({"pad_value"}), diff --git a/tilelang-dsl/python/tilelang_dsl/lowering.py b/tilelang-dsl/python/tilelang_dsl/lowering.py index d57dcb60aa..f4a208cb22 100644 --- a/tilelang-dsl/python/tilelang_dsl/lowering.py +++ b/tilelang-dsl/python/tilelang_dsl/lowering.py @@ -2536,16 +2536,19 @@ def _render_mte_ub_gm( source = self._lower_expr(expr.args[0], env, indent=indent, into=into) destination = self._lower_expr(expr.args[1], env, indent=indent, into=into) len_burst = self._lower_to_i64(expr.args[2], env, indent=indent, into=into) - nburst = self._lower_cube_i64_tuple(expr.args[3], env, indent=indent, into=into, expected_len=3) - loop_groups = self._lower_cube_loop_groups(expr.args[4], env, indent=indent, into=into) + l2_cache_ctl = self._lower_to_i64(expr.args[3], env, indent=indent, into=into) + nburst = self._lower_cube_i64_tuple(expr.args[4], env, indent=indent, into=into, expected_len=3) + loop_groups = self._lower_cube_loop_groups(expr.args[5], env, indent=indent, into=into) op_text = ( f"pto.{expr.name} {source.name}, {destination.name}, {len_burst.name}" f" nburst({nburst[0].name}, {nburst[1].name}, {nburst[2].name})" + f" l2_cache_ctl({l2_cache_ctl.name})" ) type_text = ( f"{self._render_type(source.type)}, {self._render_type(destination.type)}, " - f"{self._render_type(len_burst.type)}, {self._render_type(nburst[0].type)}, " - f"{self._render_type(nburst[1].type)}, {self._render_type(nburst[2].type)}" + f"{self._render_type(len_burst.type)}, " + f"{self._render_type(nburst[0].type)}, {self._render_type(nburst[1].type)}, " + f"{self._render_type(nburst[2].type)}, {self._render_type(l2_cache_ctl.type)}" ) for count, src_stride, dst_stride in loop_groups: op_text += f" loop({count.name}, {src_stride.name}, {dst_stride.name})" diff --git a/tilelang-dsl/python/tilelang_dsl/semantic.py b/tilelang-dsl/python/tilelang_dsl/semantic.py index 2ff13fb8d4..12186df658 100644 --- a/tilelang-dsl/python/tilelang_dsl/semantic.py +++ b/tilelang-dsl/python/tilelang_dsl/semantic.py @@ -112,6 +112,26 @@ ) +_MTE_STORE_L2_CACHE_CONTROL_VALUES = { + "nmfv": 0, + "nmlv": 1, + "nmprs": 2, + "nmred": 3, + "naci": 4, + "napw": 5, + "napi": 6, + "nared": 7, + "wbhfv": 8, + "wbhlv": 9, + "wbhprs": 10, + "wbhred": 11, + "wtsfv": 12, + "wtslv": 13, + "wtsprs": 14, + "wtsred": 15, +} + + _DTYPE_SYMBOLS = { "i1": i1, "i8": i8, @@ -3891,6 +3911,19 @@ def _require_matching_cube_pointer_element_dtypes( return raise TypeError(f"{context} requires source/destination pointer element dtypes to match") + def _require_matching_pointer_element_dtypes( + self, + lhs: SemanticExpr, + rhs: SemanticExpr, + context: str, + ) -> None: + lhs_dtype = lhs.type.element_dtype + rhs_dtype = rhs.type.element_dtype + if lhs_dtype is None or rhs_dtype is None: + return + if lhs_dtype != rhs_dtype: + raise TypeError(f"{context} requires source/destination pointer element dtypes to match") + def _require_cube_i64_tuple( self, expr: SemanticExpr, @@ -4218,19 +4251,33 @@ def _analyze_mte_ub_gm( dst = self._require_pointer_expr(args[1], "pto.mte_ub_gm destination", memory_space="gm") self._require_matching_pointer_element_dtypes(src, dst, "pto.mte_ub_gm") self._require_i64_like_expr(args[2], "pto.mte_ub_gm len_burst") - allowed_keywords = {"nburst", "loops"} + allowed_keywords = {"nburst", "loops", "l2_cache"} unsupported = sorted(set(keywords) - allowed_keywords) if unsupported: raise TypeError( - "pto.mte_ub_gm only accepts keyword(s) nburst, loops in TileLang DSL v1; " + "pto.mte_ub_gm only accepts keyword(s) nburst, loops, l2_cache in TileLang DSL v1; " f"got unsupported keyword(s): {', '.join(unsupported)}" ) + if "l2_cache" in keywords: + l2_cache = self._require_string_expr(keywords["l2_cache"], "pto.mte_ub_gm l2_cache").strip().lower() + if l2_cache not in _MTE_STORE_L2_CACHE_CONTROL_VALUES: + expected = ", ".join(sorted(_MTE_STORE_L2_CACHE_CONTROL_VALUES)) + raise ValueError( + f"pto.mte_ub_gm l2_cache does not support {l2_cache!r}; expected one of {expected}" + ) + l2_cache_ctl = SemanticLiteralExpr( + value=_MTE_STORE_L2_CACHE_CONTROL_VALUES[l2_cache], + type=SemanticScalarType(dtype=ScalarType("i64")), + ) + else: + l2_cache_ctl = SemanticLiteralExpr(value=0, type=SemanticScalarType(dtype=ScalarType("i64"))) + self._require_i64_like_expr(l2_cache_ctl, "pto.mte_ub_gm l2_cache_ctl") nburst_expr = self._require_grouped_mte_nburst(keywords, "pto.mte_ub_gm") loops_expr = self._normalize_cube_loop_groups(keywords.get("loops"), "pto.mte_ub_gm loops") return SemanticCallExpr( namespace="pto", name="mte_ub_gm", - args=(args[0], args[1], args[2], nburst_expr, loops_expr), + args=(args[0], args[1], args[2], l2_cache_ctl, nburst_expr, loops_expr), type=None, ) diff --git a/tilelang-dsl/tests/test_tilelang_dsl_v1.py b/tilelang-dsl/tests/test_tilelang_dsl_v1.py index c61cb929c3..145a1ce1cb 100644 --- a/tilelang-dsl/tests/test_tilelang_dsl_v1.py +++ b/tilelang-dsl/tests/test_tilelang_dsl_v1.py @@ -3164,6 +3164,54 @@ def kernel(inp: pto.TensorView, part: pto.PartitionTensorView, l1: pto.Tile, lef r"%acc_ptr_\d+ = pto\.tile_buf_addr %arg5 : !pto\.tile_buf -> !pto\.ptr", ) + def test_ckernel_mte_ub_gm_lowers_l2_cache_ctl_operand(self) -> None: + @pto.vkernel(op="mte_ub_gm_l2_default_unique", dtypes=[(pto.f16,)], advanced=True) + def default_kernel(out: pto.TensorView): + ub = pto.Tile((1, 64), pto.f16, pto.MemorySpace.UB) + pto.mte_ub_gm(ub.as_ptr(), out.as_ptr(), 128, nburst=(1, 128, 128)) + return None + + default_text = pto.select_kernel( + "a5", + "mte_ub_gm_l2_default_unique", + (pto.f16,), + ).mlir_text() + self.assertRegex( + default_text, + r"pto\.mte_ub_gm %[A-Za-z0-9_]+, %[A-Za-z0-9_]+, %[A-Za-z0-9_]+ nburst" + r"\([^)]+\) l2_cache_ctl\(%c0_i64\)", + ) + self.assertRegex( + default_text, + r"pto\.mte_ub_gm [^\n]+ : !pto\.ptr, !pto\.ptr, i64, i64, i64, i64, i64", + ) + + @pto.vkernel(op="mte_ub_gm_l2_explicit_unique", dtypes=[(pto.f16,)], advanced=True) + def explicit_kernel(out: pto.TensorView): + ub = pto.Tile((1, 64), pto.f16, pto.MemorySpace.UB) + pto.mte_ub_gm(ub.as_ptr(), out.as_ptr(), 128, nburst=(1, 128, 128), l2_cache="nared") + return None + + explicit_text = pto.select_kernel( + "a5", + "mte_ub_gm_l2_explicit_unique", + (pto.f16,), + ).mlir_text() + self.assertRegex( + explicit_text, + r"pto\.mte_ub_gm %[A-Za-z0-9_]+, %[A-Za-z0-9_]+, %[A-Za-z0-9_]+ nburst" + r"\([^)]+\) l2_cache_ctl\(%c7_i64\)", + ) + + @pto.vkernel(op="mte_ub_gm_l2_invalid_unique", dtypes=[(pto.f16,)], advanced=True) + def invalid_kernel(out: pto.TensorView): + ub = pto.Tile((1, 64), pto.f16, pto.MemorySpace.UB) + pto.mte_ub_gm(ub.as_ptr(), out.as_ptr(), 128, nburst=(1, 128, 128), l2_cache="invalid") + return None + + with self.assertRaisesRegex(ValueError, "l2_cache does not support 'invalid'"): + pto.select_kernel("a5", "mte_ub_gm_l2_invalid_unique", (pto.f16,)).mlir_text() + def test_ckernel_full_pipeline_bridge_ops_lower_one_to_one_in_authoring_form(self) -> None: @pto.ckernel(op="cube_bridge_pipeline_query_unique", dtypes=[(pto.f16,)], name="cube_bridge_pipeline_unique") def kernel(inp: pto.TensorView): diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 8d997e7473..2c62a9005c 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -3072,7 +3072,6 @@ int mlir::pto::compilePTOASModule( pm.addPass(createNarrowUnusedMultiResultProvenancePass()); pm.addPass(createCanonicalizerPass()); pm.addPass(createCSEPass()); - pm.addPass(pto::createPTOMemoryConsistencyPass()); if (failed(applyConfiguredPassManagerCLOptions(pm, "main PTOAS pipeline"))) return 1;