diff --git a/docs/designs/ptodsl-loop-unroll-hint-design.md b/docs/designs/ptodsl-loop-unroll-hint-design.md index 4029a42d8b..928029b42c 100644 --- a/docs/designs/ptodsl-loop-unroll-hint-design.md +++ b/docs/designs/ptodsl-loop-unroll-hint-design.md @@ -1,15 +1,30 @@ # PTODSL / PTOAS Loop Unroll Hint 设计文档 -> **修订记录（v2）**：初版设计为双阶段——Pass A 原生展开 + Pass B -> （`pto-lower-loop-hints`）把残留 hint 透传为 LLVM loop metadata（含 -> `unroll="enable"`/`"disable"` 两个 cost-model 委托 hint）。评审后决定 -> **移除阶段一（Pass B）**：`full`/`unroll_factor` 的原生展开已覆盖 -> #1242/#1000 需要的 unroll 接口能力，`enable`/`disable` 随之删除， -> 前端遇到这两个值直接报错。本文档已更新为 v2 状态；涉及 Pass B 的 -> 历史分析（如 §2.2 的 metadata 下传链路）保留作为决策记录。 - +> **修订记录（v4）**：评审指出，原实现只局部降级带注解 loop 会把新生成的 +> cf blocks 留在外层 single-block region（无 hint 的外层 `scf.for`、`scf.if` +> 等）内，触发 SingleBlock verifier 失败。现改为**单 pass 完成整个函数的 +> SCF→CF 转换**——上游 conversion patterns + 更高 benefit 的带注解 loop +> pattern,该 pass 因此**替代**两条 emitter pipeline 中的 +> `createConvertSCFToCFPass`;pass 名相应改为 +> `pto-convert-scf-to-cf-with-loop-hints`（它已不只处理 hint,而是真正承担 +> 控制流转换）。另修复 Pass A 中 `enable` 判断排在 native-unroll guard 之后 +> 导致空 body / 非 index loop 的 hint 被静默丢弃的问题。 +> +> **修订记录（v3）**：v2 移除整个阶段一（Pass B 与 `enable`/`disable`） +> 后，#1242 Req2 的验收标准（"public `pto.for_` 可以携带 unroll-enable +> hint；hint 能到达 LLVM/BiSheng 而不是前端强制 full unroll"）无法闭 +> 环。v3 **恢复 `unroll="enable"` 单一 metadata 通道**：精简版 +> `pto-convert-scf-to-cf-with-loop-hints` 只负责 `enable` → `!llvm.loop.unroll.enable`， +> `disable` 不再恢复，`full`/`factor` 的丢弃语义不变。 +> +> ~~修订记录（v2）~~：初版设计为双阶段——Pass A 原生展开 + Pass B 把残 +> 留 hint 透传为 LLVM loop metadata（含 `enable`/`disable`）。评审后决 +> 定移除阶段一：`full`/`unroll_factor` 的原生展开已覆盖 unroll 接口能 +> 力，`enable`/`disable` 随之删除。v3 重新评估后部分回退了该决定（仅 +> 恢复 `enable`）。 +> > 关联 Issue： -> - [Issue #1242](https://github.com/hw-native-sys/PTOAS/issues/1242) Requirement 2 —— `pto.for_` loop-unroll hint（v1 的 `unroll="enable"` 已随阶段一移除，由 `unroll="full"` / `unroll_factor=N` 覆盖该接口诉求） +> - [Issue #1242](https://github.com/hw-native-sys/PTOAS/issues/1242) Requirement 2 —— `pto.for_` loop-unroll hint（`unroll="enable"` 经 `pto-convert-scf-to-cf-with-loop-hints` 转为 `llvm.loop.unroll.enable`） > - [Issue #1000](https://github.com/hw-native-sys/PTOAS/issues/1000) —— 支持 Loop Unroll Hint（含 `pto.range`、factor unroll、两阶段计划） > - [PR #838](https://github.com/hw-native-sys/PTOAS/pull/838) —— `PTOUnrollSIMTForPass`（规避 BiSheng AICore 后端 bug 的临时方案，本次一并重构） @@ -107,6 +122,7 @@ PTODSL 目前**没有任何 public API** 设置 `pto.unroll` attr——该 attr | 前端写法 | `scf.for` attr | 语义 | Native 处理（Pass A） | 无法原生展开时 | |---|---|---|---|---| | （无 hint） | 无 | 现状不变 | 不处理 | — | +| `unroll="enable"` | `pto.unroll = "enable"` | 保留 loop，编译器 cost model 决定 full/partial | 不处理（在所有 native-unroll guard 之前跳过，留给 Pass B） | —（必然转成 metadata） | | `unroll="full"` | `pto.unroll = "full"` | PTOAS 强制全展开 | `loopUnrollByFactor(tripCount)`，loop 消失 | 动态 trip：丢 hint + remark，保留 loop | | `unroll_factor=N` | `pto.unroll_factor = N`（i32） | PTOAS 按 N 展开 | `loopUnrollByFactor(N)`，生成 main + epilogue | 动态 step / 超上限 / N=1：丢 hint + remark，保留 loop | @@ -130,11 +146,14 @@ PTODSL 目前**没有任何 public API** 设置 `pto.unroll` attr——该 attr - Pass A 的展开 fixpoint 不设轮数上限：每轮重新 walk 拾取外层展开克隆出的内层带注解 loop，直到某轮不再有任何变化为止。固定轮数上限会让超过该深度的嵌套 `full` hint 静默残留并丢失，违反 `full` 对静态循环强制 native 展开的契约； - Pass A 不使用 greedy pattern driver，而是按 post-order（内层先处理）手动驱动 `loopUnrollByFactor`：该 util 用内部 `IRRewriter` 删除被展开的 loop，绕过 driver 的 listener，会使 driver worklist 中的指针悬空；post-order 保证 erase 外层 loop 时其内层已全部处理完，不存在悬空指针； - factor 上限：`max-unroll-factor`（默认 1024）限制 native 展开的 factor，超限丢 hint + remark，防止巨大 factor 导致编译器挂死/OOM； -- v1 曾设计 metadata 透传路径（Pass B），其中两个实现要点已不再需要，但值得记录：(a) LLVM 19 的 `convert-scf-to-cf` 不会把 `llvm.loop_annotation` 从 `scf.for` 传到 latch `cf.br`（上游新版本才支持），v1 因此让 Pass B 自行降级带注解 loop 并以 ODS 裸名 `loop_annotation` 挂到 latch（MLIR→LLVM IR 翻译经 `BrOp::getLoopAnnotationAttr()` 按裸名查找）；(b) v1 的 Pass B 用 `applyOpPatternsAndFold` + `ExistingOps` 把改写限制在带注解 loop 上，避免全函数 greedy 顺带折叠无关 op（如 ub→llvm config word 的 `arith.ori` 链）。v2 移除该路径后，Pass A 手动驱动天然满足同样的"不动无关 IR"约束。 +- metadata 通道的两个实现要点：(a) LLVM 19 的 `convert-scf-to-cf` 不会把 `llvm.loop_annotation` 从 `scf.for` 传到 latch `cf.br`（上游新版本才支持），因此 Pass B 用自定义 pattern 降级带注解 loop 并以 ODS 裸名 `loop_annotation` 挂到 latch（MLIR→LLVM IR 翻译经 `BrOp::getLoopAnnotationAttr()` 按裸名查找）;(b) 该降级**必须覆盖整个函数**而非只处理带注解的 loop:只降级内层带注解 loop 会把新生成的 condition/body/latch/exit blocks 留在外层 single-block region 内（无 hint 的外层 `scf.for`、`scf.if`、`scf.while` 等），在 stock 转换运行前就触发该 op 的 SingleBlock verifier 失败。因此 Pass B 组合上游 conversion patterns（处理无 hint 结构）与自定义 pattern(benefit=2，处理带注解 loop),用 `applyPartialConversion` 一次转换完毕；这也顺带满足"不折叠无关 IR"的约束——conversion driver 本就不做 folding。 + +### 3.2 重复展开的 by-construction 排除 -### 3.2 重复展开问题随阶段一移除而消失 +#1000 担心"阶段一 CCE bypass 与阶段二 native unroll 对同一循环重复展开"。本设计中每个 loop 的 attr 只被消费一次： -#1000 担心"阶段一 CCE bypass 与阶段二 native unroll 对同一循环重复展开"。v2 只有原生展开一个阶段：每个 loop 的 attr 由 Pass A 消费一次（展开或丢弃），不存在第二条消费通道，该问题在构造上消失。 +- `full`/`factor`：只被 Pass A 消费（展开或丢弃），Pass B 只认 `enable`，对其他 attr 一律不触碰； +- `enable`：Pass A 在所有 native-unroll guard 之前直接跳过（不消费），由 Pass B 翻译成 metadata。 factor 展开生成的 epilogue loop 不带任何 hint attr，fixpoint 与后续管线都不会再处理它，天然防止二次展开。 @@ -144,20 +163,25 @@ factor 展开生成的 epilogue loop 不带任何 hint attr，fixpoint 与后续 PTODSL 前端 PTOAS 后端 ───────────── ───────────────────────────────────────────── pto.for_(..., unroll=...) prepareVPTOForEmission: -pto.range(...) (AST rewrite) pto-unroll-loops（唯一的 hint 消费者） +pto.range(...) (AST rewrite) [Pass A] pto-unroll-loops │ ├─ "full" / factor 可展开 → 原生展开，attr 移除 - ▼ ├─ 无法展开（动态 trip/step、超上限、factor=1、 -scf.for {pto.unroll = "full", 空 body）→ 丢 hint + remark，loop 保留 - pto.unroll_factor = N} └─ 非法 hint（enable/disable/未知值/互斥/ - │ 非 i32 factor）→ 硬错误 - └──────────────────────────▶ SCCP / canonicalize / CSE ← 折叠展开后的常量分支 - ...（其余 VPTO 优化 pass） - ───────────────────────────── - convert-scf-to-cf / LLVM lowering - （attr 已全部消费，无需任何 metadata 通道） + ▼ ├─ 无法展开 → 丢 hint + remark，loop 保留 +scf.for {pto.unroll = "full"/"enable", ├─ "enable" → 跳过（留给 Pass B） + pto.unroll_factor = N} └─ 非法 hint → 硬错误 + │ SCCP / canonicalize / CSE + └──────────────────────────▶ ───────────────────────────── + VPTO emission pipeline（两个 emitter）: + [Pass B] pto-convert-scf-to-cf-with-loop-hints + ├─ "enable" → #llvm.loop_annotation， + │ 自定义 pattern 降级并挂 latch cf.br + └─ 其余 scf.for/if/while → 上游 patterns + （本 pass 替代 convert-scf-to-cf） + convert-cf-to-llvm / translateModuleToLLVMIR + → !llvm.loop.unroll.enable metadata │ ▼ - BiSheng + BiSheng（enable loop 由其 cost model 决定 + full/partial；其余 loop 走默认估价） ``` --- @@ -173,7 +197,7 @@ def for_(start, stop, *, step, unroll=None, unroll_factor=None): ... ``` -- `unroll`：取值 `None | "full"`（v1 的 `"enable"`/`"disable"` 已随阶段一移除）； +- `unroll`：取值 `None | "full" | "enable"`（`"enable"` 为 metadata hint，见 §5.3；`"disable"` 不支持）； - `unroll_factor`：`None` 或 ≥ 1 的 `int`； - 入口参数校验：非法取值 / 互斥冲突 / 非正整数 factor 抛 `TypeError` 或 `ValueError`，诊断信息可定位到调用点； - hint 沿 `_ForBuilder` → `_ForCM.__enter__`（`_control_flow.py`）传递，`scf.ForOp` 创建后立即通过共享 helper（`_tracing/control_flow.py` 的 `apply_unroll_hint`，配套校验函数 `normalize_unroll_hint`）挂 attr： @@ -225,7 +249,7 @@ for i in pto.range(0, N, unroll_factor=4): - **新 pass 名**：`pto-unroll-loops`；保留 `pto-unroll-simt-for` 作为 alias（两个现存测试通过 `--mlir-print-ir-after=pto-unroll-simt-for` 引用，行为不变，零回归）； - **位置不变**：`prepareVPTOForEmission` 内、SCCP/canonicalize/CSE 之前（`tools/ptoas/ptoas.cpp`），保留 #838 "展开后常量分支被折叠"的收益； - **两阶段结构**：先 walk 全函数校验所有 hint（收集全部诊断后统一失败——函数 pass adaptor 在某个函数失败后可能跳过其余函数，诊断必须函数内完备），合法再进入展开 fixpoint； -- **校验**（硬错误）：`pto.unroll` 非 `"full"` 值（含已删除的 `"enable"`/`"disable"`）、两 attr 同现、factor 不符合 signless i32 正数契约（`isValidUnrollFactorAttr`）； +- **校验**（硬错误）：`pto.unroll` 非 `"full"`/`"enable"` 值（`"disable"` 及未知值）、两 attr 同现、factor 不符合 signless i32 正数契约（`isValidUnrollFactorAttr`）；`"enable"` 不消费、原样留给 Pass B(该判断排在所有 native-unroll guard 之前，否则空 body / 非 index loop 的 hint 会被静默丢弃); - **处理逻辑**（校验通过后）： - `pto.unroll = "full"`：静态 lb/ub/step、正 step、可计算 trip count → `loopUnrollByFactor(tripCount)` 全展开，loop 与 attr 一并消失；动态 trip 无法展开 → 丢 hint + remark，loop 保留； - `pto.unroll_factor = N`：`loopUnrollByFactor(N)`（动态 bounds 同样支持，上游 util 自动生成 epilogue 并穿线 live-out carry）；成功后 attr 移除；N=1、动态 step、超过 `max-unroll-factor`（默认 1024）→ 丢 hint + remark，loop 保留； @@ -233,14 +257,16 @@ for i in pto.range(0, N, unroll_factor=4): - **放开 SIMT-context 限制**：#838 的 auto-detect（trip count ≤ 64 自动展开）已移除，现存逻辑只认显式 attr——显式 attr 即用户意图，在非 SIMT 函数中静默忽略反而违反直觉。删除 `isInSIMTContext` 检查，pass 文档同步更新； - **可选护栏**：full unroll 静态 trip count 超过阈值（默认 1024，可用 pass option 调整）时 emit warning，防止 IR 体积爆炸。 -### 5.3 （已移除）Pass B：`PTOLowerLoopHints` +### 5.3 Pass B：`PTOConvertSCFToCFWithLoopHints`（承担 SCF→CF 转换 + `enable` 透传） -v1 曾设计 `pto-lower-loop-hints` 把残留 hint 翻译成 `llvm.loop_annotation` 并经自定义 SCF→CF 降级挂到 latch `cf.br`（LLVM 19 的 `convert-scf-to-cf` 不传该注解），最终生成 `!llvm.loop.unroll.*` metadata。评审后认定阶段二原生展开已覆盖需求，该 pass 及 `enable`/`disable` hint 一并移除： +v2 曾整体移除该 pass;v3 为满足 #1242 Req2 的 enable 验收标准恢复其最小子集;v4 按评审意见把它从"局部降级"改为**完整的 PTOAS 专用 SCF→CF 转换**,并相应改名（原名 `pto-lower-loop-hints` 已不能反映其职责）: -- hint 无法原生展开时的行为从"降级为 metadata"改为"丢 hint + remark"； -- 诊断职责全部收回 Pass A； -- epilogue 不再需要 `pto.unroll = "disable"` 盖章（没有下游消费者，fixpoint 不处理无 hint 的 loop）；这同时根除了 v1 中"promoteIfSingleIteration splice 出的嵌套 loop 被误盖 disable"的正确性问题——盖章逻辑已随消费者一并删除； -- 两个 emitter pipeline 恢复原样（v1 各插过一行 `addNestedPass`）。 +- **pass 名**：`pto-convert-scf-to-cf-with-loop-hints`,func-level; +- **插入点**：两个 emitter pipeline 中原 `createConvertSCFToCFPass()` 的位置——本 pass **替代**它，两者不可同时运行（会冗余）;必须排在所有 structured-loop 变换之后，确保没有后续 pass 克隆 loop 时丢失 hint; +- **翻译**：`{pto.unroll = "enable"}` → `#llvm.loop_annotation>`（即 `!llvm.loop.unroll.enable`,LLVM ForceEnable 语义)；loop 上已有 `llvm.loop_annotation` 时合并 unroll 字段（已有 unroll 条目被覆盖时 warning);其余 attr 一律不触碰(full/factor 归 Pass A,理论上到不了这里); +- **转换机制**:`populateSCFToControlFlowConversionPatterns`（上游全套 for/if/while/forall/parallel）+ 自定义 `LowerAnnotatedForPattern`(benefit=2,覆盖上游 `ForLowering`),经 `applyPartialConversion` 一次完成（scf 系全部标记 illegal),与上游 `SCFToControlFlowPass` 同构。带注解 loop 的注解以 ODS 裸名 `loop_annotation` 挂到 latch `cf.br`(MLIR→LLVM IR 翻译经 `BrOp::getLoopAnnotationAttr()` 按裸名查找,`convert-cf-to-llvm` 原样转发分支属性); +- **为什么必须整体转换**：只降级带注解 loop 会把新生成的 condition/body/latch/exit blocks 留在外层 single-block region 内，触发 `scf.for`/`scf.if`/`scf.while` 的 SingleBlock verifier 失败（emitter 的 PassManager 开启 `enableVerifier()`,失败发生在 stock 转换运行之前）; +- **副作用差异**:conversion driver 不做 folding,因此不存在 greedy driver 那种"顺带折叠无关 op"的问题；空 body 的 enable loop 也按上游行为正常降级（不再被当死代码删除）。 ### 5.4 与 #838 bug 规避语义的关系 @@ -260,7 +286,7 @@ v1 曾设计 `pto-lower-loop-hints` 把残留 hint 翻译成 `llvm.loop_annotati | 普通路径 `range(...)` / `pto.range(...)` 使用常量非正 step | PTODSL 前端 `PTODSLAstRewriteError`（负 step 仅带 break/continue 的 `pto._while` 路径支持） | | `pto.range` 在非 AST-rewrite 上下文被调用 | `RuntimeError`（提示仅用于 rewrite 场景） | | 手写 IR 中 attr 种类错误（`pto.unroll` 非 string / `pto.unroll_factor` 非 integer） | Pass A `emitError`（否则 typed getter 返回空，malformed hint 会静默留在 IR 中） | -| 手写 IR 中 `pto.unroll` 未知字符串（含已删除的 `"enable"`/`"disable"`） | Pass A `emitError` | +| 手写 IR 中 `pto.unroll` 未知字符串（含不支持的 `"disable"`；`"enable"` 合法） | Pass A `emitError` | | 手写 IR 中 `pto.unroll_factor` 类型/范围不合约（非 signless i32 或非正） | Pass A `emitError` | | 手写 IR 中 `pto.unroll` 与 `pto.unroll_factor` 同时出现在一个 loop 上 | Pass A `emitError`（互斥） | | `"full"` / factor 无法原生展开（动态 trip / 动态 step / 超 `max-unroll-factor` / factor=1 / 空 body / 静态空迭代区间 `ub <= lb` / 非 index 归纳变量） | Pass A emit remark + 丢 hint，loop 保留，编译继续 | @@ -272,7 +298,7 @@ v1 曾设计 `pto-lower-loop-hints` 把残留 hint 翻译成 `llvm.loop_annotati ### 7.1 PTODSL 前端测试（`ptodsl/tests/`） -- `for_(..., unroll="full")` / `unroll_factor=4` 生成的 `scf.for` 携带正确 attr；`unroll="enable"`/`"disable"` 被拒绝； +- `for_(..., unroll="full")` / `unroll="enable"` / `unroll_factor=4` 生成的 `scf.for` 携带正确 attr；`unroll="disable"` 被拒绝； - `for i in pto.range(...)` 与 `with pto.for_(...)` 生成相同 IR（bounds / step / attr / SSA 语义逐字节一致）； - `.carry(...)` 循环携带 hint 并正确编译（live-out carry 值正确）； - `range` / `pto.range` / `pto.for_` 无 hint 时 IR 完全一致； @@ -294,7 +320,7 @@ v1 曾设计 `pto-lower-loop-hints` 把残留 hint 翻译成 `llvm.loop_annotati **Hint 丢弃与诊断（Pass A）**： - 动态 trip 的 `"full"`、动态 step / 超上限 / factor=1 的 factor → remark + 丢 hint、loop 保留； -- 非法 hint（未知值含已删除的 `enable`/`disable`、互斥、非 i32 factor）→ error 而非静默通过；单函数内多个非法 loop 的诊断一次性全部发出（不受函数级并行调度影响）； +- 非法 hint（未知值含不支持的 `disable`、互斥、非 i32 factor；`enable` 合法）→ error 而非静默通过；单函数内多个非法 loop 的诊断一次性全部发出（不受函数级并行调度影响）； - 无 hint 循环的 IR 与最终产物逐字节不变。 **回归**： @@ -327,16 +353,32 @@ v1 曾设计 `pto-lower-loop-hints` 把残留 hint 翻译成 `llvm.loop_annotati # PTODSL / PTOAS Loop Unroll Hint — Design Document -> **Revision history (v2)**: the initial design had two stages - Pass A -> (native unrolling) plus Pass B (`pto-lower-loop-hints`) forwarding leftover -> hints as LLVM loop metadata (including the cost-model-delegating -> `unroll="enable"`/`"disable"` hints). After review, **stage 1 (Pass B) was -> removed**: native unrolling of `full`/`unroll_factor` already covers the -> unroll-interface needs of #1242/#1000, and `enable`/`disable` were removed -> with it - the frontend rejects both values. This document reflects v2; -> analyses that motivated Pass B (e.g. the metadata delivery chain in §2.2) -> are kept as decision records. - +> **Revision history (v4)**: review pointed out that lowering only the +> annotated loops leaves the freshly created cf blocks inside an enclosing +> single-block region (an unannotated outer `scf.for`, an `scf.if`, ...) and +> trips its SingleBlock verifier. The pass now performs the **complete +> SCF-to-CF conversion for the function** - upstream conversion patterns plus +> a higher-benefit pattern for annotated loops - and therefore **replaces** +> `createConvertSCFToCFPass` in both emitter pipelines. It was renamed to +> `pto-convert-scf-to-cf-with-loop-hints` accordingly (it no longer just +> handles hints; it owns the control-flow conversion). Also fixes Pass A, +> where the `enable` check sat behind the native-unroll guards and silently +> dropped the hint on empty-body and non-index loops. +> +> **Revision history (v3)**: after v2 removed stage 1 entirely (Pass B and +> `enable`/`disable`), #1242 Req2's acceptance criteria ("public `pto.for_` +> can carry an unroll-enable hint; the hint reaches LLVM/BiSheng instead of +> being force-unrolled by the frontend") could not be satisfied. v3 +> **restores `unroll="enable"` as the single metadata channel**: a slimmed +> `pto-convert-scf-to-cf-with-loop-hints` handles only `enable` → +> `!llvm.loop.unroll.enable`; `disable` stays removed, and the drop semantics +> of `full`/`factor` is unchanged. +> +> ~~Revision history (v2)~~: the initial design had two stages - Pass A +> (native unrolling) plus Pass B forwarding leftover hints as LLVM loop +> metadata. After review, stage 1 was removed and `enable`/`disable` with +> it. v3 partially reverted that decision (restoring only `enable`). +> > Related issues: > - [Issue #1242](https://github.com/hw-native-sys/PTOAS/issues/1242) Requirement 2 — `pto.for_` loop-unroll hint (v1's `unroll="enable"` was removed with stage 1; the interface need is covered by `unroll="full"` / `unroll_factor=N`) > - [Issue #1000](https://github.com/hw-native-sys/PTOAS/issues/1000) — Loop Unroll Hint support (incl. `pto.range`, factor unroll, two-phase plan) @@ -441,6 +483,7 @@ One attribute encoding (discardable attrs on `scf.for`), one consumption path (n | Frontend syntax | `scf.for` attr | Semantics | Native handling (Pass A) | When native unrolling is impossible | |---|---|---|---|---| | (no hint) | none | unchanged | not handled | — | +| `unroll="enable"` | `pto.unroll = "enable"` | keep the loop; the compiler's cost model decides full/partial | not handled (skipped before every native-unroll guard, left for Pass B) | — (always becomes metadata) | | `unroll="full"` | `pto.unroll = "full"` | PTOAS forced full unroll | `loopUnrollByFactor(tripCount)`; the loop disappears | dynamic trip: drop hint + remark, keep the loop | | `unroll_factor=N` | `pto.unroll_factor = N` (i32) | PTOAS unrolls by N | `loopUnrollByFactor(N)`; main + epilogue loops | dynamic step / over-cap / N=1: drop hint + remark, keep the loop | @@ -492,26 +535,31 @@ Edge cases (implementation decisions): - factor cap: `max-unroll-factor` (default 1024) bounds native factor unrolling; beyond it the hint is dropped with a remark, preventing a huge factor from hanging/OOMing the compiler; -- v1 designed a metadata forwarding path (Pass B); two of its implementation - points are no longer needed but worth recording: (a) LLVM 19's +- two implementation points of the metadata channel: (a) LLVM 19's `convert-scf-to-cf` does not propagate `llvm.loop_annotation` from - `scf.for` to the latch `cf.br` (only newer upstream versions do), so v1's - Pass B lowered annotated loops to control flow itself and attached the - annotation to the latch under the bare ODS name `loop_annotation` (the - MLIR-to-LLVM-IR translation looks it up via - `BrOp::getLoopAnnotationAttr()`); (b) v1's Pass B restricted its rewrite to - annotated loops via `applyOpPatternsAndFold` + `ExistingOps` so a - function-wide greedy run could not fold unrelated ops (e.g. the - `arith.ori` chains of the ub-to-llvm config words). With the path removed - in v2, Pass A's manual driving satisfies the same - "don't touch unrelated IR" constraint by construction. - -### 3.2 Double unrolling disappears with stage 1 + `scf.for` to the latch `cf.br` (only newer upstream versions do), so Pass B + lowers annotated loops with a custom pattern and stores the annotation on + the latch under the bare ODS name `loop_annotation` (the MLIR-to-LLVM-IR + translation looks it up via `BrOp::getLoopAnnotationAttr()`); (b) that + lowering **must cover the whole function** rather than just the annotated + loops: lowering an inner annotated loop alone leaves its new + condition/body/latch/exit blocks inside the enclosing single-block region + (an unannotated outer `scf.for`, `scf.if`, `scf.while`, ...) and fails that + op's SingleBlock verifier before the stock conversion ever runs. Pass B + therefore combines the upstream conversion patterns (for unannotated + structures) with the custom pattern (benefit=2, for annotated loops) in a + single `applyPartialConversion`; this also satisfies the "do not fold + unrelated IR" constraint for free - the conversion driver does no folding + +### 3.2 Double unrolling excluded by construction #1000 worried about "phase-1 CCE bypass and phase-2 native unroll both -unrolling the same loop". v2 has a single stage: each loop's attribute is -consumed exactly once by Pass A (unrolled or dropped); there is no second -consumption channel, so the problem disappears by construction. +unrolling the same loop". In this design each loop's attribute is consumed +exactly once: + +- `full`/`factor`: consumed only by Pass A (unrolled or dropped); Pass B + only recognizes `enable` and never touches anything else; +- `enable`: skipped by Pass A before every native-unroll guard, and translated into metadata by Pass B. The epilogue loop produced by a factor unroll carries no hint attribute, so neither the fixpoint nor any later pass touches it again - re-unrolling is @@ -523,24 +571,30 @@ prevented without any tagging. PTODSL frontend PTOAS backend ───────────── ───────────────────────────────────────────── pto.for_(..., unroll=...) prepareVPTOForEmission: -pto.range(...) (AST rewrite) pto-unroll-loops (the only hint consumer) +pto.range(...) (AST rewrite) [Pass A] pto-unroll-loops │ ├─ "full" / factor unrollable → unroll ▼ │ natively, attribute removed -scf.for {pto.unroll = "full", ├─ not unrollable (dynamic trip/step, - pto.unroll_factor = N} │ over-cap factor, factor=1, empty - │ │ body) → drop hint + remark, loop kept - └──────────────────────────▶ └─ malformed hint (enable/disable/unknown - value, conflicting attrs, non-i32 - factor) → hard error - SCCP / canonicalize / CSE (fold constant - ... branches exposed by unrolling) - ───────────────────────────── - convert-scf-to-cf / LLVM lowering - (attributes fully consumed; no metadata - channel needed) - │ - ▼ - BiSheng +scf.for {pto.unroll = "full"/"enable", ├─ not unrollable → drop hint + remark + pto.unroll_factor = N} ├─ "enable" → skipped (left for Pass B) + │ └─ malformed hint → hard error + └──────────────────────────▶ SCCP / canonicalize / CSE + ───────────────────────────── + VPTO emission pipeline (both emitters): + [Pass B] pto-convert-scf-to-cf-with-loop-hints + ├─ "enable" → #llvm.loop_annotation, + │ custom pattern, annotation on + │ the latch cf.br + └─ every other scf.for/if/while → + upstream patterns (this pass + replaces convert-scf-to-cf) + convert-cf-to-llvm / translateModuleToLLVMIR + → !llvm.loop.unroll.enable metadata + │ + ▼ + BiSheng (the enable loop's full/partial + unroll is chosen by its cost model; + everything else gets the default + evaluation) ``` --- @@ -608,7 +662,7 @@ Promote `kUnrollAttrName` / `kUnrollFullValue` (previously private in `PTOUnroll - **New pass name**: `pto-unroll-loops`; keep `pto-unroll-simt-for` as an alias (two existing tests reference it via `--mlir-print-ir-after=pto-unroll-simt-for`; behavior is unchanged, zero regression); - **Position unchanged**: inside `prepareVPTOForEmission`, before SCCP/canonicalize/CSE (`tools/ptoas/ptoas.cpp`), preserving the #838 benefit of folding constant branches after unrolling; - **Two-phase structure**: first walk the function and validate every hint (collecting all diagnostics before failing once - the function pass adaptor may stop scheduling functions after the first failure, so diagnostics must be complete per function), then run the unroll fixpoint; -- **Validation** (hard errors): a `pto.unroll` value other than `"full"` (including the removed `"enable"`/`"disable"`), both attributes on one loop, or a factor violating the signless-i32 positive-factor contract (`isValidUnrollFactorAttr`); +- **Validation** (hard errors): a `pto.unroll` value other than `"full"`/`"enable"` (`"disable"` and unknown values), both attributes on one loop, or a factor violating the signless-i32 positive-factor contract (`isValidUnrollFactorAttr`); `"enable"` is passed through untouched for Pass B (this check runs before every native-unroll guard - otherwise an empty-body or non-index loop would silently lose the hint); - **Handling logic** (after validation): - `pto.unroll = "full"`: static lb/ub/step, positive step, computable trip count → fully unroll via `loopUnrollByFactor(tripCount)`; the loop and the attribute disappear. A dynamic trip count cannot be unrolled → drop the hint with a remark and keep the loop; - `pto.unroll_factor = N`: unroll via `loopUnrollByFactor(N)` (dynamic bounds supported; the upstream utility generates the epilogue and threads live-out carries); on success the attribute is removed. N=1, a dynamic step, or a factor above `max-unroll-factor` (default 1024) → drop the hint with a remark and keep the loop; @@ -616,26 +670,43 @@ Promote `kUnrollAttrName` / `kUnrollFullValue` (previously private in `PTOUnroll - **Lift the SIMT-context restriction**: #838's auto-detection (auto-unroll for trip count ≤ 64) has already been removed; only explicit attributes remain — an explicit attribute is user intent, and silently ignoring it outside SIMT contexts would be counterintuitive. The `isInSIMTContext` check is removed and the pass documentation updated; - **Optional guardrail**: warn when a full-unroll static trip count exceeds a threshold (default 1024, tunable via pass option) to prevent IR explosion. -### 5.3 (Removed) Pass B: `PTOLowerLoopHints` - -v1 designed `pto-lower-loop-hints` to translate leftover hints into -`llvm.loop_annotation` and attach them to the latch `cf.br` via a custom -SCF→CF lowering (LLVM 19's `convert-scf-to-cf` does not propagate the -annotation), producing `!llvm.loop.unroll.*` metadata. Review concluded -that native unrolling (stage 2) already covers the requirements, so the pass -and the `enable`/`disable` hints were removed: - -- hints that cannot be unrolled natively are now dropped with a remark - instead of being degraded to metadata; -- all diagnostics moved into Pass A; -- the epilogue no longer needs a `pto.unroll = "disable"` stamp (no - downstream consumer exists, and the fixpoint never touches hint-free - loops); this also eliminates v1's correctness hazard where - `promoteIfSingleIteration` could splice nested loops into the parent block - and get them mis-stamped - the stamping logic was deleted together with - its consumer; -- both emitter pipelines are back to their original shape (v1 inserted one - `addNestedPass` line in each). +### 5.3 Pass B: `PTOConvertSCFToCFWithLoopHints` (owns SCF-to-CF plus `enable` forwarding) + +v2 removed this pass entirely; v3 restored its minimal subset for #1242 +Req2's enable criteria; v4 turned it from a local lowering into the +**complete PTOAS-specific SCF-to-CF conversion** per review, and renamed it +accordingly (the old `pto-lower-loop-hints` no longer described its job): + +- **Pass name**: `pto-convert-scf-to-cf-with-loop-hints`, func-level; +- **Insertion point**: where `createConvertSCFToCFPass()` used to sit in both + emitter pipelines - this pass **replaces** it, and the two must not both + run (that would be redundant). It must stay after every structured-loop + transformation so no later pass can clone a loop and lose its hint; +- **Translation**: `{pto.unroll = "enable"}` → + `#llvm.loop_annotation>` (i.e. + `!llvm.loop.unroll.enable`, LLVM's ForceEnable semantics). An existing + `llvm.loop_annotation` is merged (an existing unroll entry is overwritten + with a warning). Every other attribute is untouched (`full`/factor belong + to Pass A and cannot legitimately reach here); +- **Conversion mechanism**: + `populateSCFToControlFlowConversionPatterns` (the full upstream set: + for/if/while/forall/parallel) plus a custom `LowerAnnotatedForPattern` + (benefit=2, overriding the upstream `ForLowering`), driven by a single + `applyPartialConversion` with the scf ops marked illegal - structurally + identical to upstream `SCFToControlFlowPass`. The annotation is stored on + the latch `cf.br` under the bare ODS name `loop_annotation` (the + MLIR-to-LLVM-IR translation looks it up via + `BrOp::getLoopAnnotationAttr()`; `convert-cf-to-llvm` forwards branch + attributes verbatim); +- **Why the whole function**: lowering only the annotated loops leaves the + new condition/body/latch/exit blocks inside the enclosing single-block + region and fails the `scf.for`/`scf.if`/`scf.while` SingleBlock verifier - + the emitter's PassManager has `enableVerifier()`, so this happens before + the stock conversion would run; +- **Side-effect difference**: the conversion driver does no folding, so + there is no "accidentally folds unrelated ops" hazard like the greedy + driver had, and an empty-body enable loop is lowered like upstream would + (it is no longer removed as dead code). ### 5.4 Relationship to the #838 bug-workaround semantics @@ -655,7 +726,7 @@ and the `enable`/`disable` hints were removed: | Constant non-positive step on the plain `range(...)` / `pto.range(...)` path | PTODSL frontend `PTODSLAstRewriteError` (negative steps are only supported on the break/continue `pto._while` path) | | `pto.range` called outside an AST-rewrite context | `RuntimeError` (rewrite-only hint) | | Hand-written IR with a wrongly *kinded* attribute (`pto.unroll` not a string / `pto.unroll_factor` not an integer) | Pass A `emitError` (otherwise the typed getters return null and the malformed hint survives silently) | -| Hand-written IR with unknown `pto.unroll` string (including the removed `"enable"`/`"disable"`) | Pass A `emitError` | +| Hand-written IR with unknown `pto.unroll` string (including the unsupported `"disable"`; `"enable"` is legal) | Pass A `emitError` | | Hand-written IR with an out-of-contract `pto.unroll_factor` (not signless i32, or non-positive) | Pass A `emitError` | | Hand-written IR with both `pto.unroll` and `pto.unroll_factor` on one loop | Pass A `emitError` (mutual exclusion) | | `"full"` / factor cannot be unrolled natively (dynamic trip / dynamic step / above `max-unroll-factor` / factor=1 / empty body / statically empty iteration space `ub <= lb` / non-index induction variable) | Pass A emits a remark + drops the hint; the loop is kept; compilation continues | @@ -667,7 +738,7 @@ and the `enable`/`disable` hints were removed: ### 7.1 PTODSL frontend tests (`ptodsl/tests/`) -- `for_(..., unroll="full")` / `unroll_factor=4` produce `scf.for` with the correct attributes; `unroll="enable"`/`"disable"` are rejected; +- `for_(..., unroll="full")` / `unroll="enable"` / `unroll_factor=4` produce `scf.for` with the correct attributes; `unroll="disable"` is rejected; - `for i in pto.range(...)` and `with pto.for_(...)` produce identical IR (bounds / step / attributes / SSA semantics, byte-for-byte); - `.carry(...)` loops carry the hint and compile correctly (correct live-out carry values); - `range` / `pto.range` / `pto.for_` without hints produce identical IR; @@ -689,7 +760,7 @@ and the `enable`/`disable` hints were removed: **Hint dropping and diagnostics (Pass A)**: - dynamic-trip `"full"`, dynamic-step / over-cap / factor=1 → remark + hint dropped, loop kept; -- malformed hints (unknown values including the removed `enable`/`disable`, conflicting attrs, non-i32 factor) → errors, never silently accepted; multiple malformed loops in one function produce all diagnostics at once (immune to function-level parallel scheduling); +- malformed hints (unknown values including the unsupported `disable`, conflicting attrs, non-i32 factor; `enable` is legal) → errors, never silently accepted; multiple malformed loops in one function produce all diagnostics at once (immune to function-level parallel scheduling); - unhinted loops produce byte-identical IR. **Regression**: diff --git a/include/PTO/IR/PTO.h b/include/PTO/IR/PTO.h index 7b3d380efa..db2cfd82a9 100644 --- a/include/PTO/IR/PTO.h +++ b/include/PTO/IR/PTO.h @@ -213,20 +213,26 @@ inline constexpr llvm::StringLiteral kPTODSLLogicalNameAttrName = /// Loop-unroll hint attributes carried on `scf.for` as discardable attrs. /// -/// `pto.unroll` is a string attribute; only "full" is supported. +/// `pto.unroll` is a string attribute; "full" and "enable" are supported. /// `pto.unroll_factor` is an integer attribute holding a positive unroll /// factor. The two attributes are mutually exclusive on one loop. /// -/// Consumption contract (`pto-unroll-loops` is the only consumer): -/// - "full": unrolled natively when the trip count is a positive constant; -/// otherwise the hint is dropped with a remark and the loop is kept. -/// - `pto.unroll_factor`: unrolled natively when the value satisfies -/// `isValidUnrollFactorAttr`, the step is a positive constant, and the -/// factor does not exceed the pass's max-unroll-factor cap; otherwise the -/// hint is dropped with a remark. Malformed hints (unknown pto.unroll -/// value, both attributes on one loop, out-of-contract factor) are hard -/// errors reported by the pass. +/// Consumption contract: +/// - "full": `pto-unroll-loops` unrolls natively when the trip count is a +/// positive constant; otherwise the hint is dropped with a remark and the +/// loop is kept. +/// - "enable": never unrolled natively. `pto-convert-scf-to-cf-with-loop-hints` translates +/// it into an llvm.loop_annotation that becomes !llvm.loop.unroll.enable +/// metadata, delegating the unroll decision to the compiler's cost model +/// (LLVM's ForceEnable semantics). +/// - `pto.unroll_factor`: unrolled natively by `pto-unroll-loops` when the +/// value satisfies `isValidUnrollFactorAttr`, the step is a positive +/// constant, and the factor does not exceed the pass's max-unroll-factor +/// cap; otherwise the hint is dropped with a remark. Malformed hints +/// (unknown pto.unroll value, both attributes on one loop, out-of-contract +/// factor) are hard errors reported by `pto-unroll-loops`. inline constexpr llvm::StringLiteral kUnrollAttrName = "pto.unroll"; +inline constexpr llvm::StringLiteral kUnrollEnableValue = "enable"; inline constexpr llvm::StringLiteral kUnrollFullValue = "full"; inline constexpr llvm::StringLiteral kUnrollFactorAttrName = "pto.unroll_factor"; diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index b115f4a901..3b982fba8c 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -96,6 +96,7 @@ LogicalResult validateIntToPtrUses(func::FuncOp func); std::unique_ptr createPTOUnrollLoopsPass(); /// Backward-compatible alias of createPTOUnrollLoopsPass(). std::unique_ptr createPTOUnrollSIMTForPass(); +std::unique_ptr createPTOConvertSCFToCFWithLoopHintsPass(); std::unique_ptr createPTONarrowVPTOLoopCountersPass(); std::unique_ptr createPTOAnalyzeSIMTPersistentFragmentPass(); std::unique_ptr createPTOMaterializeSIMTPersistentFragmentPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index ca5e10b4f0..0921994969 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -855,6 +855,49 @@ def PTOUnrollSIMTFor : Pass<"pto-unroll-simt-for", "func::FuncOp"> { ]; } +def PTOConvertSCFToCFWithLoopHints : Pass<"pto-convert-scf-to-cf-with-loop-hints", "func::FuncOp"> { + let summary = + "Convert SCF to CF, preserving the enable loop hint as an LLVM loop " + "annotation"; + let description = [{ + Translates `{pto.unroll = "enable"}` attributes on `scf.for` into + `#llvm.loop_annotation>` attributes so that + the `!llvm.loop.unroll.enable` metadata reaches the emitted LLVM IR and + the downstream compiler's cost model decides whether and how to unroll + (LLVM's ForceEnable semantics). All other unroll hints are owned by + `pto-unroll-loops` and are left untouched. + + An existing `llvm.loop_annotation` attribute on the loop is merged rather + than overwritten (a pre-existing unroll entry is replaced with a + warning). + + The stock LLVM 19 `convert-scf-to-cf` does not propagate + `llvm.loop_annotation` from `scf.for` to the loop latch, so this pass + owns the SCF-to-CF conversion for the function: it runs the upstream + conversion patterns together with a higher-benefit pattern that lowers + annotated `scf.for` loops (mirroring `ForLowering`) and attaches the + annotation to the latch `cf.br`. + + Converting the whole function is required for correctness - lowering an + annotated loop in isolation would leave several blocks inside whatever + enclosing single-block region held it (an outer `scf.for`, an `scf.if`, + ...) and fail that op's verifier. + + Consequently this pass **replaces** `createConvertSCFToCFPass` in the + pipelines that run it; it must sit at the same position, after every + structured-loop transformation, so no later pass can clone a loop and + lose its hint. + }]; + let constructor = "mlir::pto::createPTOConvertSCFToCFWithLoopHintsPass()"; + let dependentDialects = [ + "mlir::func::FuncDialect", + "mlir::scf::SCFDialect", + "mlir::arith::ArithDialect", + "mlir::cf::ControlFlowDialect", + "mlir::LLVM::LLVMDialect" + ]; +} + def PTONarrowVPTOLoopCounters : Pass<"pto-narrow-vpto-loop-counters", "func::FuncOp"> { let summary = diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index 05067d9e7c..9419efe865 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -62,6 +62,7 @@ add_mlir_dialect_library(PTOTransforms VPTOBufferMaterialization.cpp PTOValidateVPTOIR.cpp PTONarrowVPTOLoopCounters.cpp + PTOConvertSCFToCFWithLoopHintsPass.cpp PTOUnrollLoopsPass.cpp PTOValidateVMIIR.cpp VMIPreAssignmentCombine.cpp diff --git a/lib/PTO/Transforms/PTOConvertSCFToCFWithLoopHintsPass.cpp b/lib/PTO/Transforms/PTOConvertSCFToCFWithLoopHintsPass.cpp new file mode 100644 index 0000000000..063fa8234a --- /dev/null +++ b/lib/PTO/Transforms/PTOConvertSCFToCFWithLoopHintsPass.cpp @@ -0,0 +1,283 @@ +// 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. + +//===- PTOConvertSCFToCFWithLoopHintsPass.cpp -----------------------------===// +// +// PTOAS-specific SCF-to-CF conversion that preserves loop unroll hints. +// +// This is the pipeline's SCF-to-CF conversion (it replaces +// createConvertSCFToCFPass), extended with the one thing the stock pass +// cannot do on LLVM 19: carry {pto.unroll = "enable"} over to the loop latch +// as LLVM loop metadata. +// +// "enable" is the only hint with a metadata channel: it asks the downstream +// compiler's cost model to unroll (LLVM's ForceEnable semantics - the way of +// unrolling is chosen by the cost model, the budget veto is lifted). The +// "full" / factor hints are consumed natively by pto-unroll-loops; anything +// this pass sees carrying {pto.unroll = "enable"} is forwarded as +// +// #llvm.loop_annotation> +// +// which the MLIR-to-LLVM-IR translation turns into !llvm.loop.unroll.enable +// metadata. +// +// LLVM 19's convert-scf-to-cf does not propagate llvm.loop_annotation from +// scf.for to the loop latch (that upstream support only exists in newer +// MLIR), so this pass owns the SCF-to-CF conversion for the whole function: +// it runs the upstream conversion patterns together with a higher-benefit +// pattern that lowers annotated scf.for loops and attaches the annotation to +// the latch cf.br. Downstream CF->LLVM lowering preserves branch attributes +// on llvm.br, and the MLIR-to-LLVM-IR translation attaches the metadata. +// +// Converting the whole function (rather than just the annotated loops) is +// required for correctness: lowering an annotated loop in isolation leaves +// the freshly created condition/body/latch/exit blocks inside whatever +// enclosing single-block region held the loop (an outer unannotated scf.for, +// an scf.if, ...), which immediately fails that op's SingleBlock verifier. +// +// Because this pass performs the full conversion, the pipelines that run it +// must NOT also run createConvertSCFToCFPass afterwards. It replaces that +// pass and must run at the same position - after every structured-loop +// transformation, so no later pass can clone a loop and lose its hint. +// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/Transforms/Passes.h" + +#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/DialectConversion.h" + +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_PTOCONVERTSCFTOCFWITHLOOPHINTS +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; + +#define DEBUG_TYPE "pto-convert-scf-to-cf-with-loop-hints" + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +namespace { + +/// Name of the LLVM annotation attribute as it appears on scf.for (and as +/// LoopAnnotationAttr's ODS name on cf.br; the MLIR-to-LLVM-IR translation +/// looks it up under the bare name via BrOp::getLoopAnnotationAttr()). +static constexpr llvm::StringLiteral kLoopAnnotationAttrName = + "llvm.loop_annotation"; +static constexpr llvm::StringLiteral kBranchLoopAnnotationAttrName = + "loop_annotation"; + +/// Merge the enable unroll entry into the loop's existing +/// llvm.loop_annotation (if any) and set the merged attribute on *forOp*. +static void setMergedLoopAnnotation(scf::ForOp forOp) { + MLIRContext *ctx = forOp.getContext(); + // disableNonforced = false prints as `unroll = `, which + // the MLIR-to-LLVM-IR translation maps to !llvm.loop.unroll.enable. + LLVM::LoopUnrollAttr unroll = LLVM::LoopUnrollAttr::get( + ctx, BoolAttr::get(ctx, false), {}, {}, {}, {}, {}, {}); + auto existing = + forOp->getAttrOfType(kLoopAnnotationAttrName); + + LLVM::LoopAnnotationAttr merged; + if (!existing) { + merged = LLVM::LoopAnnotationAttr::get(ctx, {}, {}, {}, unroll, {}, {}, {}, + {}, {}, {}, {}, {}, {}, {}, {}); + } else { + if (existing.getUnroll()) { + forOp.emitWarning() << "overwriting an existing unroll entry in '" + << kLoopAnnotationAttrName << "'"; + } + merged = LLVM::LoopAnnotationAttr::get( + ctx, existing.getDisableNonforced(), existing.getVectorize(), + existing.getInterleave(), unroll, existing.getUnrollAndJam(), + existing.getLicm(), existing.getDistribute(), existing.getPipeline(), + existing.getPeeled(), existing.getUnswitch(), + existing.getMustProgress(), existing.getIsVectorized(), + existing.getStartLoc(), existing.getEndLoc(), + existing.getParallelAccesses()); + } + forOp->setAttr(kLoopAnnotationAttrName, merged); +} + +/// Translate the enable hint on one loop into an llvm.loop_annotation +/// attribute. Only {pto.unroll = "enable"} is consumed here; every other +/// attribute belongs to pto-unroll-loops and is left untouched. +static LogicalResult translateLoopHint(scf::ForOp forOp) { + auto unrollAttr = forOp->getAttrOfType(pto::kUnrollAttrName); + StringRef hintValue = unrollAttr ? unrollAttr.getValue() : ""; + if (hintValue != pto::kUnrollEnableValue) { + return success(); + } + + LLVM_DEBUG(llvm::dbgs() << "PTOConvertSCFToCFWithLoopHints: forwarding enable hint at " + << forOp.getLoc() << "\n"); + setMergedLoopAnnotation(forOp); + forOp->removeAttr(pto::kUnrollAttrName); + return success(); +} + +/// Lower one annotated scf.for to control-flow ops, attaching its +/// LLVM-dialect attributes (llvm.loop_annotation, stored on the latch under +/// the bare ODS name loop_annotation) to the latch cf.br. +/// +/// This mirrors convert-scf-to-cf's ForLowering; the latch-attribute copy +/// backports the behavior that upstream MLIR only provides in newer +/// versions. Registered with a higher benefit than the upstream pattern so +/// it wins for annotated loops; unannotated loops fall through to the +/// upstream ForLowering. +struct LowerAnnotatedForPattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(scf::ForOp forOp, + PatternRewriter &rewriter) const override { + if (!forOp->hasAttr(kLoopAnnotationAttrName)) { + return failure(); + } + + Location loc = forOp.getLoc(); + + // Start by splitting the block containing the 'scf.for' into two parts. + // The part before will get the init code, the part after will be the end + // point. + auto *initBlock = rewriter.getInsertionBlock(); + auto initPosition = rewriter.getInsertionPoint(); + auto *endBlock = rewriter.splitBlock(initBlock, initPosition); + + // Use the first block of the loop body as the condition block since it is + // the block that has the induction variable and loop-carried values as + // arguments. Split out all operations from the first block into a new + // block. Move all body blocks from the loop body region to the region + // containing the loop. + auto *conditionBlock = &forOp.getRegion().front(); + auto *firstBodyBlock = + rewriter.splitBlock(conditionBlock, conditionBlock->begin()); + auto *lastBodyBlock = &forOp.getRegion().back(); + rewriter.inlineRegionBefore(forOp.getRegion(), endBlock); + auto iv = conditionBlock->getArgument(0); + + // Append the induction variable stepping logic to the last body block and + // branch back to the condition block. Loop-carried values are taken from + // the operands of the loop terminator. + Operation *terminator = lastBodyBlock->getTerminator(); + rewriter.setInsertionPointToEnd(lastBodyBlock); + Value stepped = rewriter.create(loc, iv, forOp.getStep()); + + SmallVector loopCarried; + loopCarried.push_back(stepped); + loopCarried.append(terminator->operand_begin(), terminator->operand_end()); + auto latchBranch = + rewriter.create(loc, conditionBlock, loopCarried); + + // Attach the LLVM attributes of the scf.for to the latch branch: LLVM + // requires loop metadata on the backedge. The loop annotation is stored + // under its bare ODS name ("loop_annotation") so that the MLIR-to-LLVM-IR + // translation picks it up via BrOp::getLoopAnnotationAttr(). + for (const NamedAttribute &attr : forOp->getAttrs()) { + if (!isa(attr.getValue().getDialect())) { + continue; + } + StringRef name = attr.getName().getValue(); + if (name == kLoopAnnotationAttrName) { + name = kBranchLoopAnnotationAttrName; + } + latchBranch->setAttr(name, attr.getValue()); + } + + rewriter.eraseOp(terminator); + + // Compute loop bounds before branching to the condition. + rewriter.setInsertionPointToEnd(initBlock); + Value lowerBound = forOp.getLowerBound(); + Value upperBound = forOp.getUpperBound(); + + // The initial values of loop-carried values are obtained from the + // operands of the loop operation. + SmallVector destOperands; + destOperands.push_back(lowerBound); + llvm::append_range(destOperands, forOp.getInitArgs()); + rewriter.create(loc, conditionBlock, destOperands); + + // With the body block done, we can fill in the condition block. + rewriter.setInsertionPointToEnd(conditionBlock); + auto comparison = rewriter.create( + loc, arith::CmpIPredicate::slt, iv, upperBound); + + rewriter.create(loc, comparison, firstBodyBlock, + ArrayRef(), endBlock, + ArrayRef()); + + // The result of the loop operation is the values of the condition block + // arguments except the induction variable on the last iteration. + rewriter.replaceOp(forOp, conditionBlock->getArguments().drop_front()); + return success(); + } +}; + +struct PTOConvertSCFToCFWithLoopHints + : public pto::impl::PTOConvertSCFToCFWithLoopHintsBase { + using pto::impl::PTOConvertSCFToCFWithLoopHintsBase< + PTOConvertSCFToCFWithLoopHints>::PTOConvertSCFToCFWithLoopHintsBase; + + void runOnOperation() override { + func::FuncOp func = getOperation(); + + // Step 1: translate {pto.unroll = "enable"} attributes into + // llvm.loop_annotation attributes on scf.for. + func.walk([&](scf::ForOp forOp) { (void)translateLoopHint(forOp); }); + + // Step 2: run the complete SCF-to-CF conversion for the function, with + // the annotated-loop lowering taking precedence over the upstream + // ForLowering. Converting everything in one pass is what keeps the IR + // verifiable: a partially lowered loop would leave multiple blocks inside + // an enclosing single-block region (outer scf.for, scf.if, ...). + RewritePatternSet patterns(&getContext()); + populateSCFToControlFlowConversionPatterns(patterns); + patterns.add(patterns.getContext(), + /*benefit=*/2); + + ConversionTarget target(getContext()); + target.addIllegalOp(); + target.markUnknownOpDynamicallyLegal([](Operation *) { return true; }); + if (mlir::failed( + applyPartialConversion(func, target, std::move(patterns)))) { + signalPassFailure(); + } + } +}; + +} // namespace + +// --------------------------------------------------------------------------- +// Pass constructor +// --------------------------------------------------------------------------- + +std::unique_ptr mlir::pto::createPTOConvertSCFToCFWithLoopHintsPass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/PTOUnrollLoopsPass.cpp b/lib/PTO/Transforms/PTOUnrollLoopsPass.cpp index c89a72e0cc..c4be13cc93 100644 --- a/lib/PTO/Transforms/PTOUnrollLoopsPass.cpp +++ b/lib/PTO/Transforms/PTOUnrollLoopsPass.cpp @@ -26,10 +26,16 @@ // attribute is removed with a remark in those // cases. // -// Anything malformed is a hard error reported here (this pass is the only -// consumer of the hint attributes): an unknown pto.unroll value ("enable" and -// "disable" are not supported), both attributes on one loop, or an -// out-of-contract factor (wrong type/width, non-positive) all fail the pass. +// {pto.unroll = "enable"} is never unrolled here; it is left untouched for +// pto-convert-scf-to-cf-with-loop-hints, which forwards it to the compiler's cost model as +// !llvm.loop.unroll.enable metadata. It is recognized before every +// native-unroll guard (empty body, non-index induction variable, ...): +// those guards exist because loopUnrollByFactor cannot handle such loops, +// which is irrelevant for a hint that only becomes metadata. +// +// Anything malformed is a hard error reported here: an unknown pto.unroll +// value, both attributes on one loop, or an out-of-contract factor (wrong +// type/width, non-positive) all fail the pass. // // Loops without any unroll annotation are never modified. // @@ -265,13 +271,14 @@ struct PTOUnrollLoopsImpl { return failure(); } - if (unrollAttr && unrollAttr.getValue() != pto::kUnrollFullValue) { + StringRef unrollValue = unrollAttr ? unrollAttr.getValue() : ""; + if (unrollAttr && unrollValue != pto::kUnrollFullValue && + unrollValue != pto::kUnrollEnableValue) { forOp.emitError() << "unknown '" << pto::kUnrollAttrName << "' value '" << unrollAttr.getValue() - << "'; only \"full\" is supported (hint metadata " - "forwarding was removed; use '" - << pto::kUnrollFactorAttrName << "' to ask for an " - "explicit unroll factor)"; + << "'; expected \"full\" (native full unroll) or " + "\"enable\" (forwarded to the compiler's cost " + "model by pto-convert-scf-to-cf-with-loop-hints)"; return failure(); } @@ -299,12 +306,25 @@ struct PTOUnrollLoopsImpl { auto factorAttr = forOp->getAttrOfType(pto::kUnrollFactorAttrName); + // "enable" is the metadata hint owned by pto-convert-scf-to-cf-with-loop-hints: it never + // reaches the native-unroll utility, so none of the guards below apply + // to it. This check must come first - dropping the hint on an + // empty-body or non-index loop would break the "enable is consumed only + // by the metadata pass" contract and silently lose the annotation. + StringRef unrollValue = unrollAttr ? unrollAttr.getValue() : ""; + if (unrollValue == pto::kUnrollEnableValue) { + return UnrollOutcome::Unchanged; + } + + // Everything below only concerns the native-unroll hints ("full" and + // pto.unroll_factor), so only those attributes are dropped. + // loopUnrollByFactor reports success on empty-body loops without // changing them, which would make the fixpoint below loop forever. // Drop the hint on such loops instead. if (llvm::hasSingleElement(forOp.getBody()->getOperations())) { - forOp.emitRemark() - << "loop with an unroll hint has an empty body; dropping the hint"; + forOp.emitRemark() << "loop with a native unroll hint has an empty " + "body; dropping the hint"; forOp->removeAttr(pto::kUnrollAttrName); forOp->removeAttr(pto::kUnrollFactorAttrName); return UnrollOutcome::Unchanged; @@ -315,11 +335,12 @@ struct PTOUnrollLoopsImpl { // arith::ConstantIndexOp unconditionally. Unrolling an i16/i32 loop // would therefore emit mixed-type ops (e.g. arith.muli(i16, index)) and // an scf.for whose step no longer matches its bounds, both of which fail - // the verifier. Only index loops can be unrolled here; anything else - // keeps its loop and drops the hint. + // the verifier. Only index loops can be unrolled natively; anything + // else keeps its loop and drops the hint. if (!forOp.getInductionVar().getType().isIndex()) { forOp.emitRemark() - << "loop with an unroll hint has a non-index induction variable (" + << "loop with a native unroll hint has a non-index induction " + "variable (" << forOp.getInductionVar().getType() << "); native unrolling only supports index loops, dropping the " "hint"; @@ -328,8 +349,9 @@ struct PTOUnrollLoopsImpl { return UnrollOutcome::Unchanged; } - if (unrollAttr) + if (unrollAttr) { return tryFullUnroll(forOp); + } if (factorAttr) { if (factorAttr.getInt() == 1) { @@ -386,9 +408,15 @@ struct PTOUnrollLoopsImpl { return success(); bool changed = false; - for (scf::ForOp forOp : annotated) - if (tryUnrollAnnotated(forOp) == UnrollOutcome::Changed) + for (scf::ForOp forOp : annotated) { + UnrollOutcome outcome = tryUnrollAnnotated(forOp); + if (outcome == UnrollOutcome::Error) { + return failure(); + } + if (outcome == UnrollOutcome::Changed) { changed = true; + } + } if (!changed) return success(); } diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index 0019778dd1..db85c0cfef 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -11920,7 +11920,13 @@ static LogicalResult runPipeline(ModuleOp module, llvm::raw_ostream &diagOS, kernelModulePM.addPass( std::make_unique()); kernelModulePM.addPass(arith::createArithExpandOpsPass()); - kernelModulePM.addPass(createConvertSCFToCFPass()); + // pto-convert-scf-to-cf-with-loop-hints performs the SCF-to-CF conversion for this pipeline: + // it runs the upstream conversion patterns plus a higher-benefit lowering + // for {pto.unroll = "enable"} loops that attaches llvm.loop_annotation to + // the latch, so the !llvm.loop.unroll.enable metadata survives into the + // emitted LLVM IR. It replaces createConvertSCFToCFPass here; running both + // would be redundant. + kernelModulePM.addNestedPass(pto::createPTOConvertSCFToCFWithLoopHintsPass()); kernelModulePM.addPass(createArithToLLVMConversionPass()); kernelModulePM.addPass(createConvertIndexToLLVMPass()); kernelModulePM.addPass(createFinalizeMemRefToLLVMConversionPass()); diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 811ccdb26d..1c03390c30 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -14208,7 +14208,13 @@ static LogicalResult runPipeline(ModuleOp module, const std::string &march, kernelModulePM.addPass( std::make_unique()); kernelModulePM.addPass(arith::createArithExpandOpsPass()); - kernelModulePM.addPass(createConvertSCFToCFPass()); + // pto-convert-scf-to-cf-with-loop-hints performs the SCF-to-CF conversion for this pipeline: + // it runs the upstream conversion patterns plus a higher-benefit lowering + // for {pto.unroll = "enable"} loops that attaches llvm.loop_annotation to + // the latch, so the !llvm.loop.unroll.enable metadata survives into the + // emitted LLVM IR. It replaces createConvertSCFToCFPass here; running both + // would be redundant. + kernelModulePM.addNestedPass(pto::createPTOConvertSCFToCFWithLoopHintsPass()); kernelModulePM.addPass(createArithToLLVMConversionPass()); kernelModulePM.addPass(createConvertIndexToLLVMPass()); kernelModulePM.addPass(createFinalizeMemRefToLLVMConversionPass()); diff --git a/ptodsl/docs/user_guide/05-control-flow.md b/ptodsl/docs/user_guide/05-control-flow.md index 3005ad9387..e2c28fe08b 100644 --- a/ptodsl/docs/user_guide/05-control-flow.md +++ b/ptodsl/docs/user_guide/05-control-flow.md @@ -169,13 +169,14 @@ def unroll_hint_probe(*, BLOCK: pto.const_expr = 8): for i in pto.range(BLOCK, unroll="full"): acc = acc + pto.const(1, dtype=pto.i32) # Explicit pto.for_ path (the .carry(...) form takes the same keywords). - with pto.for_(0, BLOCK, step=1, unroll_factor=2) as i: + with pto.for_(0, BLOCK, step=1, unroll="enable") as i: acc = acc + pto.const(2, dtype=pto.i32) _ = acc ``` | Hint | Meaning | |---|---| +| `unroll="enable"` | Keep the loop and emit `llvm.loop.unroll.enable` metadata; the compiler's cost model decides whether/how to unroll (equivalent to a no-factor `#pragma unroll`). | | `unroll="full"` | Unroll completely when the trip count is a compile-time constant; otherwise the hint is dropped with a remark. | | `unroll_factor=N` | Unroll by N when the step is a compile-time constant (dynamic upper bounds are supported and produce an epilogue loop); otherwise the hint is dropped with a remark. | diff --git a/ptodsl/ptodsl/_control_flow.py b/ptodsl/ptodsl/_control_flow.py index 9c79c12e25..9ee5f11b1f 100644 --- a/ptodsl/ptodsl/_control_flow.py +++ b/ptodsl/ptodsl/_control_flow.py @@ -223,17 +223,20 @@ def for_(start, stop, *, step, unroll=None, unroll_factor=None): loop.update(acc=cur) out = loop.final("acc") - An optional loop-unroll hint asks PTOAS to unroll the loop natively - before LLVM lowering:: + An optional loop-unroll hint is forwarded to the compiler:: - with pto.for_(c0, c16, step=c1, unroll="full") as i: + with pto.for_(c0, c16, step=c1, unroll="enable") as i: ... - ``unroll="full"`` unrolls the loop completely when the trip count is a - compile-time constant (otherwise the hint is dropped with a remark). - ``unroll_factor=N`` unrolls by ``N`` (an epilogue loop handles the - remainder; dynamic upper bounds are supported). The two arguments are - mutually exclusive. Loops without a hint are unchanged. + ``unroll="enable"`` keeps the loop and forwards + ``llvm.loop.unroll.enable`` metadata, letting the compiler's cost model + decide whether and how to unroll (equivalent to a no-factor + ``#pragma unroll``). ``unroll="full"`` asks PTOAS to unroll the loop + completely when the trip count is a compile-time constant (otherwise the + hint is dropped with a remark). ``unroll_factor=N`` asks PTOAS to + unroll by ``N`` (an epilogue loop handles the remainder; dynamic upper + bounds are supported). The two arguments are mutually exclusive. + Loops without a hint are unchanged. """ normalize_unroll_hint(unroll, unroll_factor, context="pto.for_(...)") return _ForBuilder(start, stop, step, unroll=unroll, unroll_factor=unroll_factor) diff --git a/ptodsl/ptodsl/_tracing/control_flow.py b/ptodsl/ptodsl/_tracing/control_flow.py index aeec9abe2e..d4079e0ff6 100644 --- a/ptodsl/ptodsl/_tracing/control_flow.py +++ b/ptodsl/ptodsl/_tracing/control_flow.py @@ -22,7 +22,7 @@ # ── loop-unroll hints ───────────────────────────────────────────────────────── -_UNROLL_HINT_VALUES = ("full",) +_UNROLL_HINT_VALUES = ("full", "enable") # pto.unroll_factor is encoded as a signless i32 attribute and read back as a # signed value by the backend, so the factor must fit in [1, INT32_MAX]. @@ -32,9 +32,9 @@ def normalize_unroll_hint(unroll, unroll_factor, *, context="pto.for_(...)"): """Validate one loop-unroll hint pair and return it unchanged. - ``unroll`` must be "full" (the only supported value); ``unroll_factor`` - must be a positive Python int that fits the signless i32 attribute - encoding (<= 2**31 - 1). The two are mutually exclusive. + ``unroll`` must be "full" or "enable"; ``unroll_factor`` must be a + positive Python int that fits the signless i32 attribute encoding + (<= 2**31 - 1). The two are mutually exclusive. """ if unroll is not None: if not isinstance(unroll, str) or unroll not in _UNROLL_HINT_VALUES: diff --git a/ptodsl/tests/test_loop_unroll_hints.py b/ptodsl/tests/test_loop_unroll_hints.py index b058b36cc3..5c652b57ae 100644 --- a/ptodsl/tests/test_loop_unroll_hints.py +++ b/ptodsl/tests/test_loop_unroll_hints.py @@ -41,6 +41,14 @@ def range_hint_full(): _ = acc +@pto.jit(target="a5") +def range_hint_enable(): + acc = pto.const(0, dtype=pto.i32) + for i in pto.range(4, unroll="enable"): + acc = acc + pto.const(1, dtype=pto.i32) + _ = acc + + @pto.jit(target="a5") def range_hint_factor(): acc = pto.const(0, dtype=pto.i32) @@ -208,6 +216,10 @@ def main(): text = range_hint_full.compile().mlir_text() _assert_loop_attr(text, r'pto\.unroll\s*=\s*"full"', "pto.range unroll=full") + # 1b. pto.range(unroll="enable") (metadata hint owned by pto-convert-scf-to-cf-with-loop-hints). + text = range_hint_enable.compile().mlir_text() + _assert_loop_attr(text, r'pto\.unroll\s*=\s*"enable"', "pto.range unroll=enable") + # 2. pto.range(0, 8, 2, unroll_factor=4). text = range_hint_factor.compile().mlir_text() _assert_loop_attr(text, r"pto\.unroll_factor\s*=\s*4", "pto.range unroll_factor=4") @@ -320,19 +332,13 @@ def main(): text = plain_range_true_step.compile().mlir_text() _assert_loop_attr(text, r"scf\.for", "plain range step=True compiles") - # 11b. "enable"/"disable" were removed with the metadata forwarding path; - # only "full" remains a valid unroll= value. - _expect_raises( - (ValueError,), - lambda: pto.for_(0, 4, step=1, unroll="enable"), - "unroll= expects one of", - "removed enable hint", - ) + # 11b. "disable" remains unsupported ("enable" was restored with the + # metadata forwarding pass). _expect_raises( (ValueError,), lambda: pto.for_(0, 4, step=1, unroll="disable"), "unroll= expects one of", - "removed disable hint", + "unsupported disable hint", ) # 12. Type traps at eager validation: bool is an int subclass and must not diff --git a/test/dsl-st/unroll_hint_numeric.py b/test/dsl-st/unroll_hint_numeric.py index 247bcecee4..4cbbc5ec36 100644 --- a/test/dsl-st/unroll_hint_numeric.py +++ b/test/dsl-st/unroll_hint_numeric.py @@ -16,7 +16,10 @@ native main loop + epilogue (pto-unroll-loops). - out[2]: pto.range(..., unroll_factor=3) with an odd factor -> another main loop + epilogue combination; the loop must still compute correctly. -- out[3]: nested unroll="full" -> unrolling the outer loop clones the inner +- out[3]: pto.range(..., unroll="enable") -> the loop is kept and forwarded + as llvm.loop.unroll.enable metadata (pto-convert-scf-to-cf-with-loop-hints); it must still + execute correctly. +- out[4]: nested unroll="full" -> unrolling the outer loop clones the inner one; the clones must be consumed as well (review fix). """ @@ -43,13 +46,18 @@ def unroll_hint_body(output_ptr: pto.ptr(pto.i32, "gm")) -> None: odd = odd + scalar.index_cast(pto.i32, i) scalar.store(odd, output_ptr, 2) + meta = pto.const(0, dtype=pto.i32) + for i in pto.range(8, unroll="enable"): + meta = meta + scalar.index_cast(pto.i32, i) + scalar.store(meta, output_ptr, 3) + nested = pto.const(0, dtype=pto.i32) ten = pto.const(10, dtype=pto.i32) for i in pto.range(2, unroll="full"): i32 = scalar.index_cast(pto.i32, i) for j in pto.range(3, unroll="full"): nested = nested + i32 * ten + scalar.index_cast(pto.i32, j) - scalar.store(nested, output_ptr, 3) + scalar.store(nested, output_ptr, 4) @pto.jit( @@ -71,8 +79,8 @@ def unroll_hint_numeric( "unroll_hint_numeric", unroll_hint_numeric, inputs=lambda: [], - expected=np.array([6, 45, 28, 36], dtype=np.int32), - output_shape=(4,), + expected=np.array([6, 45, 28, 28, 36], dtype=np.int32), + output_shape=(5,), output_dtype=np.int32, ) ] diff --git a/test/lit/vpto/convert_scf_to_cf_with_loop_hints.pto b/test/lit/vpto/convert_scf_to_cf_with_loop_hints.pto new file mode 100644 index 0000000000..e04c28ae69 --- /dev/null +++ b/test/lit/vpto/convert_scf_to_cf_with_loop_hints.pto @@ -0,0 +1,249 @@ +// 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. + +// Test pto-convert-scf-to-cf-with-loop-hints pass: pto.unroll = "enable" translation into the +// llvm.loop.unroll.enable metadata hint. +// +// RUN-LINE-1 prints the IR right after pto-convert-scf-to-cf-with-loop-hints: annotated loops +// are lowered to cf.* by the pass itself (LLVM 19's convert-scf-to-cf would +// drop the annotation), with the annotation attached to the latch cf.br under +// the bare ODS name "loop_annotation" (the MLIR-to-LLVM-IR translation looks +// it up via BrOp::getLoopAnnotationAttr()). +// RUN-LINE-2 checks the final translated LLVM IR carries the +// !llvm.loop.unroll.enable metadata. +// +// Covers: +// 1. pto.unroll="enable" -> unroll = -> llvm.loop.unroll.enable +// 2. enable with iter_args -> loop-carried values are threaded through the +// cf lowering and the latch carries the annotation +// 3. Nested enable loops -> both latches carry the annotation +// 4. An existing llvm.loop_annotation with an unroll entry -> warning + +// the pto hint wins; without an unroll entry -> merged silently +// 5. An empty-body enable loop -> lowered like any other; the empty body +// leaves only the induction-variable step in the latch +// 6. Only the INNER loop is annotated -> the pass converts the whole +// function, so no multi-block region is left inside the unannotated +// outer loop (lowering just the annotated loop would break its +// SingleBlock verifier) +// 7. An annotated loop inside an scf.if -> same, via the upstream +// IfLowering +// 8. An annotated loop inside an scf.while body -> same, via the upstream +// WhileLowering +// 9. A function with no hints at all is still fully converted to cf +// 10. An enable loop with a non-index induction variable is lowered and +// annotated like any other (the index-only restriction belongs to +// native unrolling in pto-unroll-loops, not to this metadata path) + +// RUN: pto-test-opt --pto-convert-scf-to-cf-with-loop-hints %s 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto-llvm-ir %s -o - 2>&1 | FileCheck %s --check-prefix=LLVMIR + +// The unroll entry prints as a numbered alias at the top. The merge warning +// is emitted while the pass runs and prints before the IR, so its CHECK line +// lives here at the top. +// CHECK: warning: overwriting an existing unroll entry in 'llvm.loop_annotation' +// CHECK: #[[UNROLL:.*]] = #llvm.loop_unroll + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + // Case 1: enable. + // CHECK-LABEL: func.func @hint_enable + // CHECK-NOT: scf.for + // CHECK: cf.br {{.*}} {loop_annotation = #[[ANNOT:.*]]} + // LLVMIR: br {{.*}}!llvm.loop + // LLVMIR: !{!"llvm.loop.unroll.enable"} + func.func @hint_enable(%dst: !pto.ptr, %n: index) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + scf.for %i = %c0 to %n step %c1 { + %value = arith.index_cast %i : index to i32 + pto.store %value, %dst[%i] : !pto.ptr, i32 + } {pto.unroll = "enable"} + return + } + + // Case 2: enable with a loop-carried value. The custom SCF->CF lowering + // must thread the iter_arg correctly (init operand into the condition + // block, latch feeds back the yielded value) while the latch carries the + // annotation. + // CHECK-LABEL: func.func @hint_enable_iter_args + // CHECK-NOT: scf.for + // CHECK: cf.br ^{{.*}}(%c0, %c0 : index, index) + // CHECK: cf.br ^{{.*}}(%{{.*}}, %{{.*}} : index, index) {loop_annotation = #[[ANNOT]]} + func.func @hint_enable_iter_args(%n: index) -> index attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %acc = scf.for %i = %c0 to %n step %c1 iter_args(%carry = %c0) -> (index) { + %next = arith.addi %carry, %i : index + scf.yield %next : index + } {pto.unroll = "enable"} + return %acc : index + } + + // Case 3: nested enable loops. The outer loop is lowered first (the pass + // collects annotated loops pre-order), moving the inner scf.for into the + // new body blocks; the inner loop is then lowered too. Both latches keep + // the annotation. + // CHECK-LABEL: func.func @hint_nested + // CHECK-NOT: scf.for + // CHECK: cf.br {{.*}} {loop_annotation = #[[ANNOT]]} + // CHECK: cf.br {{.*}} {loop_annotation = #[[ANNOT]]} + func.func @hint_nested(%dst: !pto.ptr, %n: index, %m: index) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + scf.for %i = %c0 to %n step %c1 { + scf.for %j = %c0 to %m step %c1 { + %value = arith.index_cast %j : index to i32 + pto.store %value, %dst[%j] : !pto.ptr, i32 + } {pto.unroll = "enable"} + } {pto.unroll = "enable"} + return + } + + // Case 4a: the loop already carries an llvm.loop_annotation WITH an unroll + // entry: the pto hint wins and a warning is emitted (the warning CHECK line + // is at the top of the file - diagnostics print before the IR). + // CHECK-LABEL: func.func @hint_merge_overwrite + // CHECK-NOT: scf.for + // CHECK: cf.br {{.*}} {loop_annotation = #[[ANNOT]]} + func.func @hint_merge_overwrite(%dst: !pto.ptr, %n: index) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + scf.for %i = %c0 to %n step %c1 { + %value = arith.index_cast %i : index to i32 + pto.store %value, %dst[%i] : !pto.ptr, i32 + } {llvm.loop_annotation = #llvm.loop_annotation>, pto.unroll = "enable"} + return + } + + // Case 4b: an existing annotation WITHOUT an unroll entry (here: + // mustProgress) is preserved while the pto hint is merged in. The merged + // attribute differs from case 1's, so it prints under a different alias. + // CHECK-LABEL: func.func @hint_merge_keep_others + // CHECK-NOT: scf.for + // CHECK: cf.br {{.*}} {loop_annotation = #{{.*}}} + // LLVMIR: !{!"llvm.loop.mustprogress"} + func.func @hint_merge_keep_others(%dst: !pto.ptr, %n: index) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + scf.for %i = %c0 to %n step %c1 { + %value = arith.index_cast %i : index to i32 + pto.store %value, %dst[%i] : !pto.ptr, i32 + } {llvm.loop_annotation = #llvm.loop_annotation, pto.unroll = "enable"} + return + } + + // Case 5: an empty-body enable loop. The hint is translated and the loop + // is lowered like any other (the conversion driver does not fold IR away, + // matching what the stock scf->cf conversion produces); the empty body + // leaves a latch that only steps the induction variable. + // CHECK-LABEL: func.func @hint_empty_body + // CHECK-NOT: scf.for + // CHECK: cf.br {{.*}} {loop_annotation = #[[ANNOT]]} + // CHECK: return + func.func @hint_empty_body(%n: index) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + scf.for %i = %c0 to %n step %c1 { + } {pto.unroll = "enable"} + return + } + + // Case 6: only the inner loop carries the hint. Lowering it in isolation + // would leave the inner loop's condition/body/latch/exit blocks inside the + // unannotated outer scf.for, violating its SingleBlock constraint and + // failing the verifier before convert-scf-to-cf ever runs. Because this + // pass converts the whole function, the outer loop becomes cf too. + // CHECK-LABEL: func.func @hint_inner_only + // CHECK-NOT: scf.for + // CHECK: cf.br {{.*}} {loop_annotation = #[[ANNOT]]} + func.func @hint_inner_only(%dst: !pto.ptr, %n: index, %m: index) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + scf.for %i = %c0 to %n step %c1 { + scf.for %j = %c0 to %m step %c1 { + %value = arith.index_cast %j : index to i32 + pto.store %value, %dst[%j] : !pto.ptr, i32 + } {pto.unroll = "enable"} + } + return + } + + // Case 7: an annotated loop nested in an scf.if - the same single-block + // hazard, resolved by the upstream IfLowering running in this pass. + // CHECK-LABEL: func.func @hint_inside_if + // CHECK-NOT: scf.if + // CHECK-NOT: scf.for + // CHECK: cf.br {{.*}} {loop_annotation = #[[ANNOT]]} + func.func @hint_inside_if(%dst: !pto.ptr, %n: index, %cond: i1) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + scf.if %cond { + scf.for %i = %c0 to %n step %c1 { + %value = arith.index_cast %i : index to i32 + pto.store %value, %dst[%i] : !pto.ptr, i32 + } {pto.unroll = "enable"} + } + return + } + + // Case 8: an annotated loop nested in an scf.while body - resolved by the + // upstream WhileLowering. + // CHECK-LABEL: func.func @hint_inside_while + // CHECK-NOT: scf.while + // CHECK-NOT: scf.for + // CHECK: cf.br {{.*}} {loop_annotation = #[[ANNOT]]} + func.func @hint_inside_while(%dst: !pto.ptr, %n: index) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %last = scf.while (%cur = %c0) : (index) -> index { + %cmp = arith.cmpi slt, %cur, %n : index + scf.condition(%cmp) %cur : index + } do { + ^bb0(%cur: index): + scf.for %i = %c0 to %n step %c1 { + %value = arith.index_cast %i : index to i32 + pto.store %value, %dst[%i] : !pto.ptr, i32 + } {pto.unroll = "enable"} + %next = arith.addi %cur, %c1 : index + scf.yield %next : index + } + return + } + + // Case 9: a hint-free function is converted like the stock pass would do - + // this pass replaces convert-scf-to-cf, so it must handle plain IR too. + // CHECK-LABEL: func.func @no_hint_still_converted + // CHECK-NOT: scf.for + // CHECK: cf.cond_br + func.func @no_hint_still_converted(%dst: !pto.ptr, %n: index) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + scf.for %i = %c0 to %n step %c1 { + %value = arith.index_cast %i : index to i32 + pto.store %value, %dst[%i] : !pto.ptr, i32 + } + return + } + + // Case 10: an i32-typed loop with the enable hint. pto-unroll-loops + // leaves it alone (native unrolling is index-only), and this pass lowers + // it via the upstream ForLowering machinery with the annotation attached, + // so the metadata reaches the LLVM IR just like for an index loop. + // CHECK-LABEL: func.func @hint_non_index_loop + // CHECK-NOT: scf.for + // CHECK: cf.br {{.*}} {loop_annotation = #[[ANNOT]]} + func.func @hint_non_index_loop(%dst: !pto.ptr) attributes {pto.kernel} { + %idx0 = arith.constant 0 : index + %c0 = arith.constant 0 : i32 + %c4 = arith.constant 4 : i32 + %c1 = arith.constant 1 : i32 + scf.for %i = %c0 to %c4 step %c1 : i32 { + pto.store %i, %dst[%idx0] : !pto.ptr, i32 + } {pto.unroll = "enable"} + return + } +} diff --git a/test/lit/vpto/unroll_loops.pto b/test/lit/vpto/unroll_loops.pto index 2fd528c1af..50c36a8823 100644 --- a/test/lit/vpto/unroll_loops.pto +++ b/test/lit/vpto/unroll_loops.pto @@ -53,6 +53,12 @@ // 17. Loops over a signless integer induction variable are left alone: // loopUnrollByFactor builds index-typed arithmetic unconditionally and // would emit verifier-invalid mixed-type IR. Both hint kinds drop. +// 18. pto.unroll="enable" is never unrolled here: the loop and the +// attribute are left untouched for pto-convert-scf-to-cf-with-loop-hints, including on +// loops the native-unroll guards would reject (empty body, non-index +// induction variable) - those guards only exist because +// loopUnrollByFactor cannot handle such loops, which is irrelevant for +// a hint that only becomes metadata. // RUN: pto-test-opt --pto-unroll-loops %s 2>&1 | FileCheck %s @@ -61,7 +67,7 @@ // CHECK: remark: 'pto.unroll = "full"' loop has no constant trip count // CHECK: remark: 'pto.unroll_factor' = 1 is a no-op; dropping the hint // CHECK-COUNT-2: remark: 'pto.unroll_factor' loop never iterates; dropping the hint -// CHECK-COUNT-2: remark: loop with an unroll hint has a non-index induction variable +// CHECK-COUNT-2: remark: loop with a native unroll hint has a non-index induction variable module { // Case 1: divisible constant trip count -> single unrolled loop, 4 stores. @@ -436,4 +442,51 @@ module { } {pto.unroll = "full"} return } + + // Case 18: pto.unroll="enable" is the metadata hint owned by + // pto-convert-scf-to-cf-with-loop-hints; this pass must leave the loop and the attribute + // completely untouched. + // CHECK-LABEL: func.func @enable_untouched + // CHECK: scf.for + // CHECK: pto.unroll = "enable" + func.func @enable_untouched(%n: index) { + %buf = memref.alloc() : memref<1xindex> + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + scf.for %i = %c0 to %n step %c1 { + memref.store %i, %buf[%c0] : memref<1xindex> + } {pto.unroll = "enable"} + return + } + + // Case 19: an empty-body loop with the enable hint. The native-unroll + // guard for empty bodies must not fire here: the hint never reaches + // loopUnrollByFactor, so it stays for pto-convert-scf-to-cf-with-loop-hints. + // CHECK-LABEL: func.func @enable_empty_body_untouched + // CHECK: scf.for + // CHECK: pto.unroll = "enable" + func.func @enable_empty_body_untouched(%n: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + scf.for %i = %c0 to %n step %c1 { + } {pto.unroll = "enable"} + return + } + + // Case 20: a non-index (i32) loop with the enable hint. Same reasoning: + // the index-only restriction belongs to native unrolling, not to the + // metadata hint. + // CHECK-LABEL: func.func @enable_non_index_untouched + // CHECK: scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} : i32 + // CHECK: pto.unroll = "enable" + func.func @enable_non_index_untouched(%buf: memref<4xi32>) { + %idx0 = arith.constant 0 : index + %c0 = arith.constant 0 : i32 + %c4 = arith.constant 4 : i32 + %c1 = arith.constant 1 : i32 + scf.for %i = %c0 to %c4 step %c1 : i32 { + memref.store %i, %buf[%idx0] : memref<4xi32> + } {pto.unroll = "enable"} + return + } } diff --git a/test/lit/vpto/unroll_loops_invalid.pto b/test/lit/vpto/unroll_loops_invalid.pto index 123ebe85d3..f846f67b31 100644 --- a/test/lit/vpto/unroll_loops_invalid.pto +++ b/test/lit/vpto/unroll_loops_invalid.pto @@ -24,18 +24,14 @@ module { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index - // CHECK: error: unknown 'pto.unroll' value 'sometimes'; only "full" is supported + // CHECK: error: unknown 'pto.unroll' value 'sometimes'; expected "full" (native full unroll) or "enable" scf.for %i = %c0 to %n step %c1 { memref.store %i, %buf[%c0] : memref<1xindex> } {pto.unroll = "sometimes"} - // "enable" was removed together with the metadata forwarding path. - // CHECK: error: unknown 'pto.unroll' value 'enable'; only "full" is supported - scf.for %i = %c0 to %n step %c1 { - memref.store %i, %buf[%c0] : memref<1xindex> - } {pto.unroll = "enable"} - - // CHECK: error: unknown 'pto.unroll' value 'disable'; only "full" is supported + // "disable" remains unsupported ("enable" is legal and owned by + // pto-convert-scf-to-cf-with-loop-hints). + // CHECK: error: unknown 'pto.unroll' value 'disable'; expected "full" (native full unroll) or "enable" scf.for %i = %c0 to %n step %c1 { memref.store %i, %buf[%c0] : memref<1xindex> } {pto.unroll = "disable"} diff --git a/tools/pto-test-opt/pto-test-opt.cpp b/tools/pto-test-opt/pto-test-opt.cpp index 6ec1dc70ef..27fd85d8e3 100644 --- a/tools/pto-test-opt/pto-test-opt.cpp +++ b/tools/pto-test-opt/pto-test-opt.cpp @@ -15,6 +15,7 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/DialectRegistry.h" @@ -25,7 +26,8 @@ int main(int argc, char **argv) { mlir::DialectRegistry registry; registry.insert(); + mlir::scf::SCFDialect, mlir::cf::ControlFlowDialect, + mlir::LLVM::LLVMDialect>(); mlir::registerAllPasses(); mlir::pto::registerPTOPasses();