diff --git a/docs/vmi/vmi_predicate_fold_generalization_report.md b/docs/vmi/vmi_predicate_fold_generalization_report.md new file mode 100644 index 0000000000..2ad0df949f --- /dev/null +++ b/docs/vmi/vmi_predicate_fold_generalization_report.md @@ -0,0 +1,311 @@ +# VMI Predicate / Neutral-Element Fold — Implementation Report + +**Date**: 2026-08-08 +**Scope**: Generalize `VMIPredicateFold` + skip compiler-synthesized neutral reduce combines in `VMIToVPTO`. +**Binary-equal**: folds are algebraically identity (AllTrue demask, AllFalse passthru, `max(x,-inf)=x`, `add(x,0)=x`, `vdhist(acc,*,F)=acc`). Camodel ACL runs of the four R4 reduce kernels all **PASS** `np.allclose` vs reference. Device NPU (`torch.npu`) was not available on this host. + +## What shipped + +| ID | Rule | Where | +|----|------|--------| +| **R1** | AllTrue demask `create_mask(VL)` on compute | `VMIPredicateFold` | +| **R2** | AllTrue pad `vsel` → true arm | `VMIPredicateFold` (existing + kept) | +| **R3** | AllFalse pad `vsel` → false arm | `VMIPredicateFold` (existing + kept) | +| **R4** | Skip `vadd`/`vmax`/`vmin(reduced, 0/-inf/+inf)` after `vcadd`/`vcmax`/`vcmin` | `VMIToVPTO` | +| **R5** | Unrolled `vmax(-inf,x)` / `vmin(+inf,x)` / `vadd(0,x)` → `x` | `VMIPredicateFold` | +| **R6** | AllFalse `vdhist(acc,src,F) → acc` | `VMIPredicateFold` | +| — | Always-emit pad (drop `need_pad`) | `topk_gate_vmi_w128.py` | + +Shared analysis: `VMIMaskUtils` (`IntRange`, `MaskLattice`, affine/vci ranges, `mask_and/or/xor/not`). + +Lit green: `vmi_predicate_fold_pad.pto`, `vmi_predicate_fold_general.pto`, all updated `vmi_to_vpto_reduce_*.pto` FileChecks. + +--- + +## Top-3 IR gains per rule + +Gain ranking is **static IR cost** (ops removed on the hot path), ordered by corpus hit rate / documented cycle gaps where known. + +### R1 — AllTrue demask (largest corpus surface) + +Hits: quant `create_mask(VL)` on almost every VF iter; dsl ~362 full-mask sites. + +#### #1 `vmul` full-mask (quant scale / swiglu-style) + +**Before** +```mlir +%c64 = arith.constant 64 : index +%m = pto.vmi.create_mask %c64 : index -> !pto.vmi.mask<64xpred> +%out = pto.vmi.vmul %a, %b, %m + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xf32> +``` + +**After** +```mlir +%out = pto.vmi.vmul %a, %b + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> +``` + +#### #2 `vadd` full-mask (dsl elementwise) + +**Before** +```mlir +%c64 = arith.constant 64 : index +%m = pto.vmi.create_mask %c64 : index -> !pto.vmi.mask<64xpred> +%out = pto.vmi.vadd %a, %b, %m + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xf32> +``` + +**After** +```mlir +%out = pto.vmi.vadd %a, %b + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> +``` + +#### #3 `vmax` full-mask (amax / tree reduce leaves) + +**Before** +```mlir +%c64 = arith.constant 64 : index +%m = pto.vmi.create_mask %c64 : index -> !pto.vmi.mask<64xpred> +%out = pto.vmi.vmax %a, %b, %m + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xf32> +``` + +**After** +```mlir +%out = pto.vmi.vmax %a, %b + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> +``` + +--- + +### R2 — AllTrue pad `vsel` (topk when `E % VL == 0` or affine-proven) + +Documented in this week’s pad fold; always-emit pad + fold removes Python `need_pad` on w128. + +#### #1 `pad_all_true_vci` (E = VL = 64) + +**Before**: `vci` + `vbrc(E)` + `vcmp lt` + `vsel(m, score, -inf)` +**After**: `return %score` (cmp/sel DCE’d) + +#### #2 `pad_affine_multipass_all_true` (multipass index affine) + +**Before**: loop body `vci`/`vadds`/`vcmp`/`vsel` per iter +**After**: loop yields identity; pad ops removed + +#### #3 `vsel_same_arms` (degenerate pad) + +**Before**: `vsel %m, %x, %x` +**After**: `return %x` + +--- + +### R3 — AllFalse pad `vsel` (tail / empty expert chunk) + +#### #1 `pad_all_false_vci` (E = 0) + +**Before**: `vcmp` + `vsel(m, score, -inf)` +**After**: `return %-inf` + +#### #2 `pad_vadds_all_false` (index past VL) + +**Before**: `vadds` + `vcmp` + `vsel` +**After**: `return %-inf` + +#### #3 topk rem chunk with proven empty window + +Same rewrite as #1/#2 when range analysis proves all lanes fail `idx < E`. + +--- + +### R4 — Skip neutral reduce combine (RowMax / SoftmaxGrad / EuclideanNorm) + +From [`performance_analysis_0804.md`](../../../pto-vmi/docs/performance_analysis_0804.md): RowMax +33%, SoftmaxGrad +17%, EuclideanNorm +21% attributed in part to DSL `vmax`/`vadd` cleanup after `vcmax`/`vcadd`. + +#### #1 `vcadd` f16 (SoftmaxGrad / EuclideanNorm style) + +**Before** (old lowering) +```mlir +%init = pto.vdup %c0, %pall : f16, !pto.mask -> !pto.vreg<128xf16> +%vl1 = pto.pset_b16 "PAT_VL1" : !pto.mask +%red = pto.vcadd %src, %mask : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<128xf16> +%out = pto.vadd %red, %init, %vl1 : ... -> !pto.vreg<128xf16> // dead: +0 +``` + +**After** +```mlir +// init may remain dead until later DCE; combine + PAT_VL1 gone +%red = pto.vcadd %src, %mask : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<128xf16> +return %red +``` + +#### #2 `vcadd` i32 / f32 single-chunk + +Same shape as #1: `return %vcadd` with no `vadd(reduced, 0)`. + +#### #3 `vcmax` multichunk (RowMax across physical parts) + +**Before**: `vcmax` × N → `vmax(reduced, -inf)` then inter-chunk `vmax` +**After**: seed acc from first `vcmax`; only inter-chunk `vmax` remains (no combine with `-inf`) + +```mlir +%3 = pto.vcmax %arg0, %arg2 +%4 = pto.vcmax %arg1, %arg3 +%5 = pto.vmax %4, %3, %vl1 // inter-chunk only +return %5 +``` + +--- + +### R5 — Unrolled first-iter neutral splat + +Hits: topk `acc_max = -inf` tree; quant `acc = 0` sum trees. + +#### #1 `vmax(vbrc(-inf), x)` → `x` + +**Before** +```mlir +%neg = pto.vmi.vbrc %neginf : f32 -> !pto.vmi.vreg<64xf32> +%out = pto.vmi.vmax %neg, %x, %m : ... +``` + +**After**: `return %x` + +#### #2 `vadd(vbrc(0), x)` → `x` + +**Before** +```mlir +%zero = pto.vmi.vbrc %c0 : f32 -> !pto.vmi.vreg<64xf32> +%out = pto.vmi.vadd %zero, %x, %m : ... +``` + +**After**: `return %x` + +#### #3 `vmin(vbrc(+inf), x)` → `x` + +**Before** +```mlir +%pos = pto.vmi.vbrc %posinf : f32 -> !pto.vmi.vreg<64xf32> +%out = pto.vmi.vmin %pos, %x, %m : ... +``` + +**After**: `return %x` + +--- + +### R6 — AllFalse `vdhist` + +Hits: rem/empty histogram updates in group-count style kernels. + +#### #1 `vdhist` with `create_mask(0)` + +**Before** +```mlir +%m = pto.vmi.create_mask %c0 : index -> !pto.vmi.mask<256xpred> +%out = pto.vmi.vdhist %acc, %src, %m : ... +``` + +**After**: `return %acc` + +#### #2 / #3 + +Same rewrite whenever the mask lattice proves AllFalse (range-`vcmp` or `mask_and` with AllFalse). Additional call sites collapse identically to #1; corpus density is lower than R1/R4. + +--- + +## Camodel A/B (R4) — Ascend950PR_9599 / CANN 9.1.0-beta.3 + +Measured 2026-08-08 on `edgexpert-59a6` (aarch64). +**Before** = stock `ptoas` 0.53; **After** = local `ptoas` 0.56 with neutral-combine skip. +All cases: `*_real_float_Rows_128_Cols_64.py`, ACL path, `PASS`. + +| Kernel | CCE (0804) | DSL before | DSL after | Δ ticks | Δ RVECEX | Dead op removed | +|--------|------------|------------|-----------|---------|----------|-----------------| +| **RowMaxKernel** | 358 | 477 | **358** | −119 (−25%) | 388→257 (−131) | 128× `RV_VMAX` | +| **RowMinKernel** | 358 | 477 | **358** | −119 (−25%) | 388→257 (−131) | 128× `RV_VMIN` | +| **SoftmaxGradKernel** | 692 | 813 | **692** | −121 (−15%) | 772→641 (−131) | 128× `RV_VADD` | +| **EuclideanNormVfKernel** | 523 | 632 | **523** | −109 (−17%) | 644→513 (−131) | 128× `RV_VADD` | + +After R4, all four match the CCE `vf_real_execute_time` from [`performance_analysis_0804.md`](../../../pto-vmi/docs/performance_analysis_0804.md) (the prior +17…+33% DSL gaps attributed to neutral `vadd`/`vmax`/`vmin` cleanup). + +Small-shape sanity (`RowMaxKernel.case1_float_Rows_2_Cols_64`): vf 74→71, RVECEX 10→5, `RV_VMAX` 2→0. + +--- + +## Camodel A/B (R1–R3, R5–R6) — Ascend950PR_9599 / CANN 9.1.0-beta.3 + +Measured after rebuilding `libPTOASCompiler.so` (fold was in `pto-test-opt` / `.o` but not linked into the CLI shared lib until 2026-08-08). +Harness: local `ptoas` 0.56; **Before** = `PTO_FLAGS=--disable-vmi-predicate-fold` (now honored by `ptodsl` `native_build`); **After** = fold on. Full `topk_gate` camodel still blocked on `wait_flag` uncovered-section normalize under 0.56 — pad / peep coverage uses VL=64 micros with the same `vcmp+vsel` / `vmax(-inf)` / `vdhist` shapes. + +### R1 — AllTrue demask (dsl real) + +| Case | Result | vf_real | Signal | +|------|--------|---------|--------| +| SoftmaxGradKernel real 128×64 | PASS | 813→**692** | Dominated by **R4** (−128 `RV_VADD`); `RV_PSET` 3→1 | +| confusionSoftmaxGradArKernel real half 128×64 | PASS | 673→**553** | Same R4-class −128 `RV_VADD`; `RV_PSET` 3→1 | +| BlkScaleMulKernel real f16 128×128 | PASS | 183→183 | Identical asm (no R1 tick win on this shape) | + +R1 alone is typically ADD-pipe / setup; largest corpus surface, smaller tick delta than R4. + +### R2 — AllTrue pad `vsel` (E covers lanes) + +| Case | Result | vf | Before ops | After ops | +|------|--------|-----|------------|-----------| +| pad E=64 (VL=64) | PASS | 53→**50** | `VCI+VCMP+VSEL` | **0** `RV_VSEL`/`RV_VCMP` (load→store) | +| pad E=256 | PASS | 53→**50** | same | **0** pad vsel/vcmp | +| pad E=384 | PASS | 53→**50** | same | **0** pad vsel/vcmp | + +### R3 — AllFalse pad `vsel` + +| Case | Result | vf | After behavior | +|------|--------|-----|----------------| +| pad E=300 (idx base ≥ E → AllFalse) | PASS | 53→**49** | `VDUPS(-inf)`+store; **0** vsel/vcmp | +| pad E=100 (same AllFalse class) | PASS | 53→**49** | same | +| pad partial E=48 | PASS | 53→53 | Keeps `VCMP+VSEL` (correct non-fold) | + +### R5 — Neutral `vmax(-inf,x)` + +| Case | Result | vf | Asm | +|------|--------|-----|-----| +| micro `vmax(vbrc(-inf), x)` | PASS | 52→**50** | Before: `VDUPS+VMAX`; After: **0** both (identity) | +| VFShiftVectorKernel real 128×64 | PASS | 439→439 | No `-inf` seed peep site; `RV_PSET` 2→1 only | + +### R6 — AllFalse `vdhist → acc` + +| Case | Result | vf | Asm | +|------|--------|-----|-----| +| micro `vdhist(acc,src,create_mask(0))` | PASS | 61→**50** | Before: 1× `RV_DHIST`; After: **0** (acc passthrough) | +| ExpertTokenHistVfKernel case1 | PASS | ~73–75 | No AllFalse site (2× `DHIST` both sides); `PSET` 3→2 | +| IndexStatisticInt32VfKernel real | PASS | ~1600 | Still 128× `DHIST`; `PSET` 5→3 (R1-ish) | + +### Env notes + +- Rebuild `ninja PTOASCompiler` after editing `VMIPredicateFold` — `python/pto/ptoas.so` can be stale; CLI uses `python/ptoas/mlir/_mlir_libs/libPTOASCompiler.so`. +- `PTO_FLAGS` is forwarded by `ptodsl/_runtime/native_build.py` and included in the compile-config cache key. + +## Verification checklist + +| Check | Result | +|-------|--------| +| `vmi_predicate_fold_pad.pto` | PASS | +| `vmi_predicate_fold_general.pto` | PASS | +| `vmi_to_vpto_reduce_{addf,addf_f16,addi,minf,*_multichunk}.pto` | PASS | +| Camodel R4 RowMax / RowMin / SoftmaxGrad / EuclideanNorm real 128×64 | PASS + CCE-parity ticks | +| Camodel R1 SoftmaxGrad / CSG / BlkScale | PASS | +| Camodel R2 AllTrue pad E∈{64,256,384} | PASS + 0 pad `VSEL`/`VCMP` | +| Camodel R3 AllFalse pad E∈{300,100} + partial E=48 | PASS + false-arm / keep-mixed | +| Camodel R5 neutral vmax micro (+ VFShiftVector) | PASS + `VDUPS+VMAX` gone on micro | +| Camodel R6 AllFalse vdhist micro (+ ExpertTokenHist / IndexStatistic) | PASS + 0 `DHIST` on AllFalse site | +| Full `topk_gate` camodel | **Blocked** — `wait_flag` uncovered tile section under local 0.56 | +| Device NPU smoke (`test_topk_gate.py`) | **Blocked** — no `torch.npu` here | + +## Follow-ups (optional) + +1. DCE leftover `vdup(0/-inf)` after R4 skips the combine (init still materialized by LowerUnified; already not in VF body for these cases). +2. Memory AllFalse merge-store erase (Tier B). +3. On-device topk / MoE binary-equal under TileKernels-vmi (`env_npu.sh`, CANN 9.1b3). +4. Unblock full `topk_gate` camodel (`PTONormalizeUncoveredTileSections` / `wait_flag`) for production E shapes. diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index 1441697f0d..ace6ea37b8 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -136,6 +136,7 @@ std::unique_ptr createVMILayoutFoldPass(); std::unique_ptr createVMILayoutRematerializePass(); std::unique_ptr createVMILayoutSinkMaterializationPass(); std::unique_ptr createVMILegalizeArithSelectPass(); +std::unique_ptr createVMIPredicateFoldPass(); std::unique_ptr createVMILowerUnifiedToLegacyPass(); std::unique_ptr createVMINormalizeSignlessIntToUnsignedPass(); std::unique_ptr createVMIToVPTOPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index b3ee537f03..cd61d9c95e 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -1132,6 +1132,30 @@ def VMINormalizeSignlessIntToUnsigned : let dependentDialects = ["pto::PTODialect"]; } +def VMIPredicateFold : Pass<"vmi-predicate-fold", "ModuleOp"> { + let summary = "Fold statically proven VMI predicates and neutral identities"; + let description = [{ + Shared VMI predicate / neutral-element simplifier on unified IR: + + * proves AllTrue / AllFalse masks from `create_mask`, range-proven + `vcmp`/`vcmps`, and mask algebra + * folds `vsel`/`select` identity / constant arms + * demasks AllTrue Variadic-mask compute; folds AllFalse pure ops + (`vdhist→acc`, merge passthru, reduce nils) + * peeps unrolled `vmax(-inf,x)` / `vadd(0,x)` first-iter identities + * materializes proven masks to `create_mask(VL|0)` and DCEs dead defs + + Enables frontends to always emit expert-pad `vcmp_lt`+`vsel` and rely on + the compiler when `num_experts` covers the index span at compile time. + }]; + let constructor = "mlir::pto::createVMIPredicateFoldPass()"; + let dependentDialects = [ + "mlir::pto::PTODialect", + "mlir::arith::ArithDialect", + "mlir::scf::SCFDialect" + ]; +} + def VMILowerUnifiedToLegacy : Pass<"vmi-lower-unified-to-legacy", "ModuleOp"> { let summary = "Lower unified VMI ops to legacy equivalents before layout assignment"; let description = [{ diff --git a/include/PTO/Transforms/VMIMaskUtils.h b/include/PTO/Transforms/VMIMaskUtils.h new file mode 100644 index 0000000000..ed1fdb9a7a --- /dev/null +++ b/include/PTO/Transforms/VMIMaskUtils.h @@ -0,0 +1,58 @@ +// 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. + +//===- VMIMaskUtils.h - Shared VMI predicate / seed helpers -----*- C++ -*-===// +// +// Helpers shared by VMILowerUnifiedToLegacy, VMIPredicateFold, and related +// passes for proving mask shape and classifying compile-time predicates. +// +//===----------------------------------------------------------------------===// + +#ifndef PTO_TRANSFORMS_VMIMASKUTILS_H +#define PTO_TRANSFORMS_VMIMASKUTILS_H + +#include "mlir/IR/Value.h" +#include + +namespace mlir { +namespace pto { + +/// Inclusive integer range used for affine / lane-index proofs. +struct IntRange { + int64_t lo = 0; + int64_t hi = 0; + + static IntRange splat(int64_t c) { return {c, c}; } +}; + +/// Compile-time mask lattice. +enum class MaskLattice { Unknown, AllTrue, AllFalse }; + +/// Returns true if `seed` is provably an all-active mask (every lane active), +/// so `mask_and(x, seed)` is the identity. Covers a `pset` and a +/// `create_mask` whose active_lanes is a constant >= the mask lane count. +bool isAllActiveSeed(Value seed); + +/// Returns true if `seed` is provably an all-inactive mask (every lane +/// inactive). Covers `create_mask(0)`. +bool isAllInactiveSeed(Value seed); + +/// Bound an integer SSA value over known constant / affine forms. +std::optional matchAffineIntRange(Value v); + +/// Bound every lane of a VMI vector to an inclusive integer range when the +/// producer is a statically analyzable index form (vci / vadds / vbrc / …). +std::optional matchVectorLaneRange(Value v); + +/// Classify a VMI mask SSA value as AllTrue / AllFalse / Unknown. +MaskLattice classifyMaskValue(Value mask); + +} // namespace pto +} // namespace mlir + +#endif // PTO_TRANSFORMS_VMIMASKUTILS_H diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index 7f2ec906d1..5c034781d2 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -66,6 +66,8 @@ add_mlir_dialect_library(PTOTransforms VMIControlFlowSupport.cpp VMILegalizeArithSelect.cpp VMIMaskGranularityAssignment.cpp + VMIMaskUtils.cpp + VMIPredicateFold.cpp VMILowerUnifiedToLegacy.cpp VMINormalizeSignlessIntToUnsigned.cpp VMILayoutRematerializeWeakProducers.cpp diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp index 9854de1e94..cd7921c1f7 100644 --- a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp +++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp @@ -88,6 +88,7 @@ #include "PTO/IR/PTO.h" #include "PTO/IR/PTOTypeUtils.h" #include "PTO/Transforms/Passes.h" +#include "PTO/Transforms/VMIMaskUtils.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/IR/BuiltinOps.h" @@ -280,28 +281,7 @@ lowerMaskedUnary(UnifiedOp op, OpBuilder &builder, // Category C1 helpers: vcmp / vcmps //===----------------------------------------------------------------------===// -/// Returns true if `seed` is provably an all-active mask (every lane active), -/// so `mask_and(x, seed)` is the identity and the AND can be skipped. Covers a -/// `pset` (all lanes active by definition) and a `create_mask` whose -/// active_lanes is a constant >= the mask lane count. -static bool isAllActiveSeed(Value seed) { - Operation *def = seed.getDefiningOp(); - if (!def) { - return false; - } - if (isa(def)) { - return true; - } - if (auto cm = dyn_cast(def)) { - auto maskTy = cast(cm.getResult().getType()); - if (auto cst = cm.getActiveLanes().getDefiningOp()) { - if (auto ia = dyn_cast(cst.getValue())) { - return ia.getInt() >= maskTy.getElementCount(); - } - } - } - return false; -} +// isAllActiveSeed lives in VMIMaskUtils (shared with VMIPredicateFold). static bool isCompactGroupCount(int64_t count) { return count == kSingleGroupCount || count == mlir::pto::kValue2 || diff --git a/lib/PTO/Transforms/VMIMaskUtils.cpp b/lib/PTO/Transforms/VMIMaskUtils.cpp new file mode 100644 index 0000000000..b5e83a7e71 --- /dev/null +++ b/lib/PTO/Transforms/VMIMaskUtils.cpp @@ -0,0 +1,368 @@ +// 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. + +//===- VMIMaskUtils.cpp - Shared VMI predicate / seed helpers -------------===// + +#include "PTO/Transforms/VMIMaskUtils.h" + +#include "PTO/IR/PTO.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/Matchers.h" +#include "llvm/Support/MathExtras.h" + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +static std::optional matchConstantInt(Value v) { + APInt val; + if (matchPattern(v, m_ConstantInt(&val))) + return val.getSExtValue(); + return std::nullopt; +} + +static MaskLattice classifyCompare(StringRef cmp, const IntRange &lhs, + const IntRange &rhs) { + if (cmp == "lt" || cmp == "olt") { + if (lhs.hi < rhs.lo) + return MaskLattice::AllTrue; + if (lhs.lo >= rhs.hi) + return MaskLattice::AllFalse; + return MaskLattice::Unknown; + } + if (cmp == "le" || cmp == "ole") { + if (lhs.hi <= rhs.lo) + return MaskLattice::AllTrue; + if (lhs.lo > rhs.hi) + return MaskLattice::AllFalse; + return MaskLattice::Unknown; + } + if (cmp == "gt" || cmp == "ogt") { + if (lhs.lo > rhs.hi) + return MaskLattice::AllTrue; + if (lhs.hi <= rhs.lo) + return MaskLattice::AllFalse; + return MaskLattice::Unknown; + } + if (cmp == "ge" || cmp == "oge") { + if (lhs.lo >= rhs.hi) + return MaskLattice::AllTrue; + if (lhs.hi < rhs.lo) + return MaskLattice::AllFalse; + return MaskLattice::Unknown; + } + if (cmp == "eq" || cmp == "oeq") { + if (lhs.lo == lhs.hi && rhs.lo == rhs.hi && lhs.lo == rhs.lo) + return MaskLattice::AllTrue; + if (lhs.hi < rhs.lo || lhs.lo > rhs.hi) + return MaskLattice::AllFalse; + return MaskLattice::Unknown; + } + if (cmp == "ne" || cmp == "one") { + if (lhs.hi < rhs.lo || lhs.lo > rhs.hi) + return MaskLattice::AllTrue; + if (lhs.lo == lhs.hi && rhs.lo == rhs.hi && lhs.lo == rhs.lo) + return MaskLattice::AllFalse; + return MaskLattice::Unknown; + } + return MaskLattice::Unknown; +} + +} // namespace + +bool mlir::pto::isAllActiveSeed(Value seed) { + Operation *def = seed.getDefiningOp(); + if (!def) + return false; + if (isa(def)) + return true; + if (auto cm = dyn_cast(def)) { + auto maskTy = cast(cm.getResult().getType()); + if (auto cst = cm.getActiveLanes().getDefiningOp()) + if (auto ia = dyn_cast(cst.getValue())) + return ia.getInt() >= maskTy.getElementCount(); + } + return false; +} + +bool mlir::pto::isAllInactiveSeed(Value seed) { + Operation *def = seed.getDefiningOp(); + if (!def) + return false; + if (auto cm = dyn_cast(def)) { + if (auto cst = cm.getActiveLanes().getDefiningOp()) + if (auto ia = dyn_cast(cst.getValue())) + return ia.getInt() <= 0; + } + return false; +} + +std::optional mlir::pto::matchAffineIntRange(Value v) { + if (auto c = matchConstantInt(v)) + return IntRange::splat(*c); + + if (auto cast = v.getDefiningOp()) + return matchAffineIntRange(cast.getIn()); + if (auto cast = v.getDefiningOp()) + return matchAffineIntRange(cast.getIn()); + if (auto cast = v.getDefiningOp()) + return matchAffineIntRange(cast.getIn()); + if (auto cast = v.getDefiningOp()) + return matchAffineIntRange(cast.getIn()); + if (auto cast = v.getDefiningOp()) + return matchAffineIntRange(cast.getIn()); + + if (auto add = v.getDefiningOp()) { + auto lhs = matchAffineIntRange(add.getLhs()); + auto rhs = matchAffineIntRange(add.getRhs()); + if (!lhs || !rhs) + return std::nullopt; + int64_t lo, hi; + if (llvm::AddOverflow(lhs->lo, rhs->lo, lo) || + llvm::AddOverflow(lhs->hi, rhs->hi, hi)) + return std::nullopt; + return IntRange{lo, hi}; + } + + if (auto sub = v.getDefiningOp()) { + auto lhs = matchAffineIntRange(sub.getLhs()); + auto rhs = matchAffineIntRange(sub.getRhs()); + if (!lhs || !rhs) + return std::nullopt; + int64_t lo, hi; + if (llvm::SubOverflow(lhs->lo, rhs->hi, lo) || + llvm::SubOverflow(lhs->hi, rhs->lo, hi)) + return std::nullopt; + return IntRange{lo, hi}; + } + + if (auto mul = v.getDefiningOp()) { + auto lhsC = matchConstantInt(mul.getLhs()); + auto rhsC = matchConstantInt(mul.getRhs()); + if (lhsC && rhsC) { + int64_t prod; + if (llvm::MulOverflow(*lhsC, *rhsC, prod)) + return std::nullopt; + return IntRange::splat(prod); + } + + Value dyn = lhsC ? mul.getRhs() : mul.getLhs(); + auto factorOpt = lhsC ? lhsC : rhsC; + if (!factorOpt) + return std::nullopt; + int64_t factor = *factorOpt; + auto dynR = matchAffineIntRange(dyn); + if (!dynR) + return std::nullopt; + int64_t a, b; + if (llvm::MulOverflow(dynR->lo, factor, a) || + llvm::MulOverflow(dynR->hi, factor, b)) + return std::nullopt; + return IntRange{std::min(a, b), std::max(a, b)}; + } + + if (auto blockArg = dyn_cast(v)) { + if (auto forOp = dyn_cast(blockArg.getOwner()->getParentOp())) { + if (blockArg != forOp.getInductionVar()) + return std::nullopt; + auto lb = matchConstantInt(forOp.getLowerBound()); + auto ub = matchConstantInt(forOp.getUpperBound()); + auto step = matchConstantInt(forOp.getStep()); + if (!lb || !ub || !step || *step <= 0 || *lb >= *ub) + return std::nullopt; + int64_t last = *lb + ((*ub - 1 - *lb) / *step) * *step; + return IntRange{*lb, last}; + } + } + + return std::nullopt; +} + +std::optional mlir::pto::matchVectorLaneRange(Value v) { + auto vty = dyn_cast(v.getType()); + if (!vty) + return std::nullopt; + int64_t vl = vty.getElementCount(); + + if (auto brc = v.getDefiningOp()) { + if (auto c = matchConstantInt(brc.getValue())) + return IntRange::splat(*c); + return std::nullopt; + } + if (auto brc = v.getDefiningOp()) { + if (auto c = matchConstantInt(brc.getValue())) + return IntRange::splat(*c); + return std::nullopt; + } + + auto matchIotaLike = [&](Value base, std::optional group, + StringRef order) -> std::optional { + if (!order.empty() && order != "ASC") + return std::nullopt; + auto baseR = matchAffineIntRange(base); + if (!baseR) + return std::nullopt; + int64_t span = vl; + if (group && *group > 0) { + if (vl % *group != 0) + return std::nullopt; + span = vl / *group; + } + int64_t lo = baseR->lo; + int64_t hi; + if (llvm::AddOverflow(baseR->hi, span - 1, hi)) + return std::nullopt; + return IntRange{lo, hi}; + }; + + if (auto vci = v.getDefiningOp()) { + std::optional group; + if (auto g = vci->getAttrOfType("group")) + group = g.getInt(); + StringRef order = vci.getOrder() ? *vci.getOrder() : StringRef("ASC"); + return matchIotaLike(vci.getBase(), group, order); + } + if (auto iota = v.getDefiningOp()) { + std::optional group; + if (auto g = iota->getAttrOfType("group")) + group = g.getInt(); + StringRef order = iota.getOrder() ? *iota.getOrder() : StringRef("ASC"); + return matchIotaLike(iota.getBase(), group, order); + } + + if (auto vadds = v.getDefiningOp()) { + if (!isAllActiveSeed(vadds.getMask())) + return std::nullopt; + auto srcR = matchVectorLaneRange(vadds.getSrc()); + auto sc = matchConstantInt(vadds.getScalar()); + if (!srcR || !sc) + return std::nullopt; + int64_t lo, hi; + if (llvm::AddOverflow(srcR->lo, *sc, lo) || + llvm::AddOverflow(srcR->hi, *sc, hi)) + return std::nullopt; + return IntRange{lo, hi}; + } + + // vadd(v, vbrc(C)) / vadd(vbrc(C), v) with all-active (or absent) mask. + if (auto vadd = v.getDefiningOp()) { + if (!vadd.getMask().empty() && !isAllActiveSeed(vadd.getMask().front())) + return std::nullopt; + Value lhs = vadd.getLhs(); + Value rhs = vadd.getRhs(); + auto tryShift = [&](Value src, Value brcCand) -> std::optional { + auto srcR = matchVectorLaneRange(src); + if (!srcR) + return std::nullopt; + std::optional sc; + if (auto vb = brcCand.getDefiningOp()) + sc = matchConstantInt(vb.getValue()); + else if (auto vb = brcCand.getDefiningOp()) + sc = matchConstantInt(vb.getValue()); + if (!sc) + return std::nullopt; + int64_t lo, hi; + if (llvm::AddOverflow(srcR->lo, *sc, lo) || + llvm::AddOverflow(srcR->hi, *sc, hi)) + return std::nullopt; + return IntRange{lo, hi}; + }; + if (auto r = tryShift(lhs, rhs)) + return r; + if (auto r = tryShift(rhs, lhs)) + return r; + } + + return std::nullopt; +} + +MaskLattice mlir::pto::classifyMaskValue(Value mask) { + if (isAllActiveSeed(mask)) + return MaskLattice::AllTrue; + if (isAllInactiveSeed(mask)) + return MaskLattice::AllFalse; + + if (auto mand = mask.getDefiningOp()) { + MaskLattice lhs = classifyMaskValue(mand.getLhs()); + MaskLattice rhs = classifyMaskValue(mand.getRhs()); + if (lhs == MaskLattice::AllFalse || rhs == MaskLattice::AllFalse) + return MaskLattice::AllFalse; + if (lhs == MaskLattice::AllTrue && rhs == MaskLattice::AllTrue) + return MaskLattice::AllTrue; + if (lhs == MaskLattice::AllTrue) + return rhs; + if (rhs == MaskLattice::AllTrue) + return lhs; + return MaskLattice::Unknown; + } + if (auto mor = mask.getDefiningOp()) { + MaskLattice lhs = classifyMaskValue(mor.getLhs()); + MaskLattice rhs = classifyMaskValue(mor.getRhs()); + if (lhs == MaskLattice::AllTrue || rhs == MaskLattice::AllTrue) + return MaskLattice::AllTrue; + if (lhs == MaskLattice::AllFalse && rhs == MaskLattice::AllFalse) + return MaskLattice::AllFalse; + if (lhs == MaskLattice::AllFalse) + return rhs; + if (rhs == MaskLattice::AllFalse) + return lhs; + return MaskLattice::Unknown; + } + if (auto mxor = mask.getDefiningOp()) { + MaskLattice lhs = classifyMaskValue(mxor.getLhs()); + MaskLattice rhs = classifyMaskValue(mxor.getRhs()); + if (lhs == MaskLattice::Unknown || rhs == MaskLattice::Unknown) + return MaskLattice::Unknown; + if (lhs == rhs) + return MaskLattice::AllFalse; + return MaskLattice::AllTrue; + } + if (auto mnot = mask.getDefiningOp()) { + MaskLattice src = classifyMaskValue(mnot.getSource()); + if (src == MaskLattice::AllTrue) + return MaskLattice::AllFalse; + if (src == MaskLattice::AllFalse) + return MaskLattice::AllTrue; + return MaskLattice::Unknown; + } + + if (auto vcmp = mask.getDefiningOp()) { + auto lhs = matchVectorLaneRange(vcmp.getLhs()); + auto rhs = matchVectorLaneRange(vcmp.getRhs()); + if (!lhs || !rhs) + return MaskLattice::Unknown; + MaskLattice raw = classifyCompare(vcmp.getCmp(), *lhs, *rhs); + MaskLattice seedLat = classifyMaskValue(vcmp.getSeed()); + if (raw == MaskLattice::AllFalse || seedLat == MaskLattice::AllFalse) + return MaskLattice::AllFalse; + if (raw == MaskLattice::AllTrue && seedLat == MaskLattice::AllTrue) + return MaskLattice::AllTrue; + return MaskLattice::Unknown; + } + + if (auto vcmps = mask.getDefiningOp()) { + auto lhs = matchVectorLaneRange(vcmps.getSrc()); + auto sc = matchConstantInt(vcmps.getScalar()); + if (!lhs || !sc) + return MaskLattice::Unknown; + MaskLattice raw = + classifyCompare(vcmps.getCmp(), *lhs, IntRange::splat(*sc)); + MaskLattice seedLat = classifyMaskValue(vcmps.getSeed()); + if (raw == MaskLattice::AllFalse || seedLat == MaskLattice::AllFalse) + return MaskLattice::AllFalse; + if (raw == MaskLattice::AllTrue && seedLat == MaskLattice::AllTrue) + return MaskLattice::AllTrue; + return MaskLattice::Unknown; + } + + return MaskLattice::Unknown; +} diff --git a/lib/PTO/Transforms/VMIPredicateFold.cpp b/lib/PTO/Transforms/VMIPredicateFold.cpp new file mode 100644 index 0000000000..00da5d014b --- /dev/null +++ b/lib/PTO/Transforms/VMIPredicateFold.cpp @@ -0,0 +1,475 @@ +// 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. + +//===- VMIPredicateFold.cpp - Fold statically proven VMI predicates -------===// +// + // Predicate-algebra simplifier for unified VMI IR: + // * prove AllTrue / AllFalse masks (create_mask, vcmp ranges, mask algebra) + // * fold vsel / select + // * demask AllTrue consumers (Variadic-mask compute) + // * fold AllFalse pure consumers (vdhist → acc, merge → passthru) + // * fold first-iter neutral splat peeps (vmax(-inf,x) / vadd(0,x)) + // * materialize proven masks to create_mask(VL|0) + // * DCE pure unused defs +// +//===----------------------------------------------------------------------===// + +#include "PTO/IR/PTO.h" +#include "PTO/Transforms/Passes.h" +#include "PTO/Transforms/VMIMaskUtils.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Matchers.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace pto { +#define GEN_PASS_DEF_VMIPREDICATEFOLD +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +static std::optional matchConstantInt(Value v) { + APInt val; + if (matchPattern(v, m_ConstantInt(&val))) + return val.getSExtValue(); + return std::nullopt; +} + +static StringRef getPmodeOrDefault(Operation *op) { + if (auto attr = op->getAttrOfType("pmode")) + return attr.getValue(); + return "merge"; +} + +static bool isZeroPmode(Operation *op) { return getPmodeOrDefault(op) == "zero"; } + +/// Match a splat constant vector: vbrc/broadcast of Imm, or dense constant. +enum class SplatKind { None, Zero, NegInf, PosInf, Other }; + +static SplatKind classifySplat(Value v, Type *elemTyOut = nullptr) { + auto vty = dyn_cast(v.getType()); + if (!vty) + return SplatKind::None; + if (elemTyOut) + *elemTyOut = vty.getElementType(); + + auto fromAPFloat = [&](const APFloat &f) -> SplatKind { + if (f.isZero()) + return SplatKind::Zero; + if (f.isInfinity()) + return f.isNegative() ? SplatKind::NegInf : SplatKind::PosInf; + return SplatKind::Other; + }; + auto fromAPInt = [&](const APInt &i) -> SplatKind { + if (i.isZero()) + return SplatKind::Zero; + return SplatKind::Other; + }; + + auto fromScalar = [&](Value s) -> SplatKind { + if (auto c = matchConstantInt(s)) + return *c == 0 ? SplatKind::Zero : SplatKind::Other; + Attribute attr; + if (!matchPattern(s, m_Constant(&attr))) + return SplatKind::None; + if (auto fa = dyn_cast(attr)) + return fromAPFloat(fa.getValue()); + if (auto ia = dyn_cast(attr)) + return fromAPInt(ia.getValue()); + return SplatKind::None; + }; + + if (auto brc = v.getDefiningOp()) + return fromScalar(brc.getValue()); + if (auto brc = v.getDefiningOp()) + return fromScalar(brc.getValue()); + if (auto cst = v.getDefiningOp()) { + auto dense = dyn_cast(cst.getValue()); + if (!dense || dense.getNumElements() == 0) + return SplatKind::None; + auto fvals = dense.tryGetValues(); + if (succeeded(fvals)) { + auto it = fvals->begin(); + APFloat first = *it; + for (APFloat x : *fvals) + if (!x.bitwiseIsEqual(first)) + return SplatKind::None; + return fromAPFloat(first); + } + auto ivals = dense.tryGetValues(); + if (succeeded(ivals)) { + auto it = ivals->begin(); + APInt first = *it; + for (APInt x : *ivals) + if (x != first) + return SplatKind::None; + return fromAPInt(first); + } + } + return SplatKind::None; +} + +static Value createSplatConstant(OpBuilder &builder, Location loc, + VMIVRegType vty, SplatKind kind) { + Type elem = vty.getElementType(); + int64_t lanes = vty.getElementCount(); + auto shaped = RankedTensorType::get({lanes}, elem); + DenseElementsAttr attr; + if (auto floatTy = dyn_cast(elem)) { + APFloat val = APFloat::getZero(floatTy.getFloatSemantics()); + if (kind == SplatKind::NegInf) + val = APFloat::getInf(floatTy.getFloatSemantics(), /*Negative=*/true); + else if (kind == SplatKind::PosInf) + val = APFloat::getInf(floatTy.getFloatSemantics(), /*Negative=*/false); + attr = DenseElementsAttr::get(shaped, val); + } else { + auto intTy = cast(elem); + APInt val = APInt::getZero(intTy.getWidth()); + if (kind == SplatKind::NegInf) + val = APInt::getSignedMinValue(intTy.getWidth()); + else if (kind == SplatKind::PosInf) + val = APInt::getSignedMaxValue(intTy.getWidth()); + attr = DenseElementsAttr::get(shaped, val); + } + return builder.create(loc, vty, attr).getResult(); +} + +static Value materializeCanonicalMask(OpBuilder &builder, Location loc, + VMIMaskType maskTy, MaskLattice lat) { + int64_t lanes = maskTy.getElementCount(); + int64_t active = lat == MaskLattice::AllTrue ? lanes : 0; + Value n = builder.create(loc, active); + return builder.create(loc, maskTy, n).getResult(); +} + +static bool isTriviallyDeadPureOp(Operation *op) { + if (!op || op->getNumRegions() != 0) + return false; + if (op->hasTrait()) + return false; + if (isa(op)) + return false; + if (!isMemoryEffectFree(op)) + return false; + return llvm::all_of(op->getResults(), + [](Value r) { return r.use_empty(); }); +} + +static void dcePureUnusedOps(ModuleOp module) { + bool changed = true; + while (changed) { + changed = false; + SmallVector dead; + module.walk([&](Operation *op) { + if (isTriviallyDeadPureOp(op)) + dead.push_back(op); + }); + for (Operation *op : dead) { + op->erase(); + changed = true; + } + } +} + +//===----------------------------------------------------------------------===// +// Rewrites +//===----------------------------------------------------------------------===// + +static bool foldSelectLike(Value mask, Value t, Value f, Value result, + Operation *op) { + if (t == f) { + result.replaceAllUsesWith(t); + op->erase(); + return true; + } + MaskLattice lat = classifyMaskValue(mask); + if (lat == MaskLattice::Unknown) + return false; + // Default vsel pmode is merge-like (false arm). Explicit zero → 0 splat. + if (lat == MaskLattice::AllTrue) { + result.replaceAllUsesWith(t); + } else if (isZeroPmode(op)) { + OpBuilder b(op); + auto vty = cast(result.getType()); + result.replaceAllUsesWith( + createSplatConstant(b, op->getLoc(), vty, SplatKind::Zero)); + } else { + result.replaceAllUsesWith(f); + } + op->erase(); + return true; +} + +static bool foldNeutralBinary(Value lhs, Value rhs, Value result, Operation *op, + SplatKind identityOnLhs, SplatKind identityOnRhs) { + // vmax(neg_inf, x) / vmax(x, neg_inf) → x; vadd(0, x) → x; etc. + if (classifySplat(lhs) == identityOnLhs) { + result.replaceAllUsesWith(rhs); + op->erase(); + return true; + } + if (classifySplat(rhs) == identityOnRhs) { + result.replaceAllUsesWith(lhs); + op->erase(); + return true; + } + return false; +} + +template +static bool demaskBinaryAllTrue(OpTy op) { + if (op.getMask().empty()) + return false; + if (classifyMaskValue(op.getMask().front()) != MaskLattice::AllTrue) + return false; + OpBuilder b(op); + auto neu = + b.create(op.getLoc(), op.getResult().getType(), op.getLhs(), + op.getRhs(), ValueRange{}, op.getPmodeAttr()); + op.getResult().replaceAllUsesWith(neu.getResult()); + op.erase(); + return true; +} + +template +static bool demaskUnaryAllTrue(OpTy op) { + if (op.getMask().empty()) + return false; + if (classifyMaskValue(op.getMask().front()) != MaskLattice::AllTrue) + return false; + OpBuilder b(op); + auto neu = + b.create(op.getLoc(), op.getResult().getType(), op.getSource(), + ValueRange{}, op.getPmodeAttr()); + op.getResult().replaceAllUsesWith(neu.getResult()); + op.erase(); + return true; +} + +template +static bool foldBinaryAllFalse(OpTy op) { + if (op.getMask().empty()) + return false; + if (classifyMaskValue(op.getMask().front()) != MaskLattice::AllFalse) + return false; + OpBuilder b(op); + auto vty = cast(op.getResult().getType()); + if (isZeroPmode(op)) { + op.getResult().replaceAllUsesWith( + createSplatConstant(b, op.getLoc(), vty, SplatKind::Zero)); + } else { + // merge (default): inactive lanes pass lhs / source convention → lhs + op.getResult().replaceAllUsesWith(op.getLhs()); + } + op.erase(); + return true; +} + +template +static bool foldUnaryAllFalse(OpTy op) { + if (op.getMask().empty()) + return false; + if (classifyMaskValue(op.getMask().front()) != MaskLattice::AllFalse) + return false; + OpBuilder b(op); + auto vty = cast(op.getResult().getType()); + if (isZeroPmode(op)) { + op.getResult().replaceAllUsesWith( + createSplatConstant(b, op.getLoc(), vty, SplatKind::Zero)); + } else { + op.getResult().replaceAllUsesWith(op.getSource()); + } + op.erase(); + return true; +} + +static bool foldMaskAlgebra(Value result, Operation *op) { + MaskLattice lat = classifyMaskValue(result); + if (lat == MaskLattice::Unknown) + return false; + // Only rewrite pure mask producers that are not already canonical. + if (isa(op)) + return false; + OpBuilder b(op); + auto maskTy = cast(result.getType()); + Value canon = materializeCanonicalMask(b, op->getLoc(), maskTy, lat); + result.replaceAllUsesWith(canon); + op->erase(); + return true; +} + +static bool foldHistAllFalse(Value acc, Value mask, Value result, + Operation *op) { + if (classifyMaskValue(mask) != MaskLattice::AllFalse) + return false; + result.replaceAllUsesWith(acc); + op->erase(); + return true; +} + +static bool foldReduceAllFalse(Value mask, Value result, Operation *op, + SplatKind nilKind) { + if (classifyMaskValue(mask) != MaskLattice::AllFalse) + return false; + OpBuilder b(op); + auto vty = cast(result.getType()); + result.replaceAllUsesWith(createSplatConstant(b, op->getLoc(), vty, nilKind)); + op->erase(); + return true; +} + +//===----------------------------------------------------------------------===// +// Pass +//===----------------------------------------------------------------------===// + +struct VMIPredicateFoldPass + : public mlir::pto::impl::VMIPredicateFoldBase { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VMIPredicateFoldPass) + + void runOnOperation() override { + ModuleOp module = getOperation(); + bool changed = true; + while (changed) { + changed = false; + + // 1) vsel / select + SmallVector sels; + module.walk([&](Operation *op) { + if (isa(op)) + sels.push_back(op); + }); + for (Operation *op : llvm::reverse(sels)) { + if (!op->getBlock()) + continue; + if (auto sel = dyn_cast(op)) { + changed |= foldSelectLike(sel.getMask(), sel.getTrueValue(), + sel.getFalseValue(), sel.getResult(), op); + } else if (auto sel = dyn_cast(op)) { + changed |= foldSelectLike(sel.getMask(), sel.getTrueValue(), + sel.getFalseValue(), sel.getResult(), op); + } + } + + // 2) Neutral splat peeps on unrolled accumulators + SmallVector binOps; + module.walk([&](Operation *op) { + if (isa(op)) + binOps.push_back(op); + }); + for (Operation *op : llvm::reverse(binOps)) { + if (!op->getBlock()) + continue; + if (auto vmax = dyn_cast(op)) { + // Only fold when mask absent or AllTrue (identity under full lanes). + if (!vmax.getMask().empty() && + classifyMaskValue(vmax.getMask().front()) != MaskLattice::AllTrue) + continue; + changed |= foldNeutralBinary(vmax.getLhs(), vmax.getRhs(), + vmax.getResult(), op, SplatKind::NegInf, + SplatKind::NegInf); + } else if (auto vmin = dyn_cast(op)) { + if (!vmin.getMask().empty() && + classifyMaskValue(vmin.getMask().front()) != MaskLattice::AllTrue) + continue; + changed |= foldNeutralBinary(vmin.getLhs(), vmin.getRhs(), + vmin.getResult(), op, SplatKind::PosInf, + SplatKind::PosInf); + } else if (auto vadd = dyn_cast(op)) { + if (!vadd.getMask().empty() && + classifyMaskValue(vadd.getMask().front()) != MaskLattice::AllTrue) + continue; + changed |= foldNeutralBinary(vadd.getLhs(), vadd.getRhs(), + vadd.getResult(), op, SplatKind::Zero, + SplatKind::Zero); + } + } + + // 3) AllTrue demask / AllFalse fold on compute + SmallVector compute; + module.walk([&](Operation *op) { + if (isa(op)) + compute.push_back(op); + }); + for (Operation *op : llvm::reverse(compute)) { + if (!op->getBlock()) + continue; + if (auto o = dyn_cast(op)) { + changed |= demaskBinaryAllTrue(o) || foldBinaryAllFalse(o); + } else if (auto o = dyn_cast(op)) { + changed |= demaskBinaryAllTrue(o) || foldBinaryAllFalse(o); + } else if (auto o = dyn_cast(op)) { + changed |= demaskBinaryAllTrue(o) || foldBinaryAllFalse(o); + } else if (auto o = dyn_cast(op)) { + changed |= demaskBinaryAllTrue(o) || foldBinaryAllFalse(o); + } else if (auto o = dyn_cast(op)) { + changed |= demaskBinaryAllTrue(o) || foldBinaryAllFalse(o); + } else if (auto o = dyn_cast(op)) { + changed |= demaskBinaryAllTrue(o) || foldBinaryAllFalse(o); + } else if (auto o = dyn_cast(op)) { + changed |= demaskUnaryAllTrue(o) || foldUnaryAllFalse(o); + } else if (auto o = dyn_cast(op)) { + changed |= demaskUnaryAllTrue(o) || foldUnaryAllFalse(o); + } else if (auto o = dyn_cast(op)) { + changed |= + foldHistAllFalse(o.getAcc(), o.getMask(), o.getResult(), op); + } else if (auto o = dyn_cast(op)) { + changed |= + foldHistAllFalse(o.getAcc(), o.getMask(), o.getResult(), op); + } else if (auto o = dyn_cast(op)) { + changed |= + foldReduceAllFalse(o.getMask(), o.getResult(), op, SplatKind::Zero); + } else if (auto o = dyn_cast(op)) { + changed |= foldReduceAllFalse(o.getMask(), o.getResult(), op, + SplatKind::NegInf); + } else if (auto o = dyn_cast(op)) { + changed |= foldReduceAllFalse(o.getMask(), o.getResult(), op, + SplatKind::PosInf); + } + } + + // 4) Materialize proven mask algebra / vcmp results to create_mask + SmallVector masks; + module.walk([&](Operation *op) { + if (isa(op)) + masks.push_back(op); + }); + for (Operation *op : llvm::reverse(masks)) { + if (!op->getBlock() || op->use_empty()) + continue; + changed |= foldMaskAlgebra(op->getResult(0), op); + } + + if (changed) + dcePureUnusedOps(module); + } + + dcePureUnusedOps(module); + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createVMIPredicateFoldPass() { + return std::make_unique(); +} diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 24f8011cb2..2fbf7bd1aa 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -9642,6 +9642,53 @@ struct OneToNVMICompressStoreOpPattern } }; +enum class ReduceNeutralKind { Add, Max, Min }; + +/// True when `init` is a compile-time splat of the reduction's algebraic +/// identity (0 / -inf / +inf). Unified vcadd/vcmax invent this init in +/// LowerUnified; combining it after the hardware reduce is a no-op. +static bool isNeutralReduceInit(Value init, ReduceNeutralKind kind) { + auto cst = init.getDefiningOp(); + if (!cst) + return false; + auto dense = dyn_cast(cst.getValue()); + if (!dense || dense.getNumElements() == 0) + return false; + + auto floats = dense.tryGetValues(); + if (succeeded(floats)) { + APFloat first = *floats->begin(); + for (APFloat v : *floats) + if (!v.bitwiseIsEqual(first)) + return false; + switch (kind) { + case ReduceNeutralKind::Add: + return first.isZero(); + case ReduceNeutralKind::Max: + return first.isInfinity() && first.isNegative(); + case ReduceNeutralKind::Min: + return first.isInfinity() && !first.isNegative(); + } + } + + auto ints = dense.tryGetValues(); + if (succeeded(ints)) { + APInt first = *ints->begin(); + for (APInt v : *ints) + if (v != first) + return false; + switch (kind) { + case ReduceNeutralKind::Add: + return first.isZero(); + case ReduceNeutralKind::Max: + return first.isMinSignedValue(); + case ReduceNeutralKind::Min: + return first.isMaxSignedValue(); + } + } + return false; +} + struct OneToNVMIReduceAddIOpPattern : OneToNOpConversionPattern { using OneToNOpConversionPattern::OneToNOpConversionPattern; @@ -9680,6 +9727,13 @@ struct OneToNVMIReduceAddIOpPattern op, "reduce_addi requires every mask chunk to have the same " "predicate type"); + const bool skipNeutralCombine = + isNeutralReduceInit(op.getInit(), ReduceNeutralKind::Add); + + auto ensureFirstLaneMask = [&]() -> FailureOr { + return createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); + }; + FailureOr combined = combineEquivalentMaskedParts( op.getLoc(), sourceParts, maskParts, resultType, rewriter); if (succeeded(combined)) { @@ -9688,34 +9742,47 @@ struct OneToNVMIReduceAddIOpPattern .create(op.getLoc(), resultType, *combined, maskParts.front()) .getResult(); + Value result = reduced; + if (!skipNeutralCombine) { + FailureOr firstLaneMask = ensureFirstLaneMask(); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create reduce_addi first-lane mask"); + result = rewriter + .create(op.getLoc(), resultType, reduced, + initParts.front(), *firstLaneMask) + .getResult(); + } replaceOpWithFlatConvertedValues( rewriter, op, SmallVector{reduced}, *this->getTypeConverter()); return success(); } - Value accumulator = rewriter - .create(op.getLoc(), resultType, - sourceParts.front(), - maskParts.front()) - .getResult(); - if (sourceParts.size() == 1) { - replaceOpWithFlatConvertedValues( - rewriter, op, SmallVector{accumulator}, - *this->getTypeConverter()); - return success(); - } - FailureOr firstLaneMask = - createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); - if (failed(firstLaneMask)) - return rewriter.notifyMatchFailure( - op, "failed to create reduce_addi first-lane mask"); - for (size_t part = 1; part < sourceParts.size(); ++part) { + Value accumulator; + bool haveAcc = false; + FailureOr firstLaneMask; + bool haveFirstLaneMask = false; + for (auto [sourcePart, maskPart] : + llvm::zip_equal(sourceParts, maskParts)) { Value reduced = rewriter .create(op.getLoc(), resultType, sourceParts[part], maskParts[part]) .getResult(); + if (!haveAcc) { + accumulator = skipNeutralCombine ? reduced : initParts.front(); + haveAcc = true; + if (skipNeutralCombine) + continue; + } + if (!haveFirstLaneMask) { + firstLaneMask = ensureFirstLaneMask(); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create reduce_addi first-lane mask"); + haveFirstLaneMask = true; + } accumulator = rewriter .create(op.getLoc(), resultType, reduced, accumulator, *firstLaneMask) @@ -9767,6 +9834,13 @@ struct OneToNVMIReduceAddFOpPattern op, "reduce_addf requires every mask chunk to have the same " "predicate type"); + const bool skipNeutralCombine = + isNeutralReduceInit(op.getInit(), ReduceNeutralKind::Add); + + auto ensureFirstLaneMask = [&]() -> FailureOr { + return createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); + }; + FailureOr combined = combineEquivalentMaskedParts( op.getLoc(), sourceParts, maskParts, resultType, rewriter); if (succeeded(combined)) { @@ -9775,34 +9849,47 @@ struct OneToNVMIReduceAddFOpPattern .create(op.getLoc(), resultType, *combined, maskParts.front()) .getResult(); + Value result = reduced; + if (!skipNeutralCombine) { + FailureOr firstLaneMask = ensureFirstLaneMask(); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create reduce_addf first-lane mask"); + result = rewriter + .create(op.getLoc(), resultType, reduced, + initParts.front(), *firstLaneMask) + .getResult(); + } replaceOpWithFlatConvertedValues( rewriter, op, SmallVector{reduced}, *this->getTypeConverter()); return success(); } - Value accumulator = rewriter - .create(op.getLoc(), resultType, - sourceParts.front(), - maskParts.front()) - .getResult(); - if (sourceParts.size() == 1) { - replaceOpWithFlatConvertedValues( - rewriter, op, SmallVector{accumulator}, - *this->getTypeConverter()); - return success(); - } - FailureOr firstLaneMask = - createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); - if (failed(firstLaneMask)) - return rewriter.notifyMatchFailure( - op, "failed to create reduce_addf first-lane mask"); - for (size_t part = 1; part < sourceParts.size(); ++part) { + Value accumulator; + bool haveAcc = false; + FailureOr firstLaneMask; + bool haveFirstLaneMask = false; + for (auto [sourcePart, maskPart] : + llvm::zip_equal(sourceParts, maskParts)) { Value reduced = rewriter .create(op.getLoc(), resultType, sourceParts[part], maskParts[part]) .getResult(); + if (!haveAcc) { + accumulator = skipNeutralCombine ? reduced : initParts.front(); + haveAcc = true; + if (skipNeutralCombine) + continue; + } + if (!haveFirstLaneMask) { + firstLaneMask = ensureFirstLaneMask(); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create reduce_addf first-lane mask"); + haveFirstLaneMask = true; + } accumulator = rewriter .create(op.getLoc(), resultType, reduced, accumulator, *firstLaneMask) @@ -10503,6 +10590,16 @@ struct OneToNVMIReduceMinMaxOpPattern : OneToNOpConversionPattern { op, "min/max reduction requires every mask chunk to have " "the same predicate type"); + constexpr bool isMax = + std::is_same_v; + const bool skipNeutralCombine = isNeutralReduceInit( + op.getInit(), + isMax ? ReduceNeutralKind::Max : ReduceNeutralKind::Min); + + auto ensureFirstLaneMask = [&]() -> FailureOr { + return createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); + }; + FailureOr combined = combineEquivalentMaskedParts( op.getLoc(), sourceParts, maskParts, resultType, rewriter); if (succeeded(combined)) { @@ -10511,33 +10608,46 @@ struct OneToNVMIReduceMinMaxOpPattern : OneToNOpConversionPattern { .create(op.getLoc(), resultType, *combined, maskParts.front()) .getResult(); + Value result = reduced; + if (!skipNeutralCombine) { + FailureOr firstLaneMask = ensureFirstLaneMask(); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create min/max reduction first-lane mask"); + result = rewriter + .create(op.getLoc(), resultType, reduced, + initParts.front(), *firstLaneMask) + .getResult(); + } replaceOpWithFlatConvertedValues( rewriter, op, SmallVector{reduced}, *this->getTypeConverter()); return success(); } - Value accumulator = rewriter - .create(op.getLoc(), resultType, - sourceParts.front(), - maskParts.front()) - .getResult(); - if (sourceParts.size() == 1) { - replaceOpWithFlatConvertedValues( - rewriter, op, SmallVector{accumulator}, - *this->getTypeConverter()); - return success(); - } - FailureOr firstLaneMask = - createPrefixMask(op.getLoc(), maskType, "PAT_VL1", rewriter); - if (failed(firstLaneMask)) - return rewriter.notifyMatchFailure( - op, "failed to create min/max reduction first-lane mask"); - for (size_t part = 1; part < sourceParts.size(); ++part) { + Value accumulator; + bool haveAcc = false; + FailureOr firstLaneMask; + bool haveFirstLaneMask = false; + for (auto [sourcePart, maskPart] : + llvm::zip_equal(sourceParts, maskParts)) { Value reduced = rewriter .create(op.getLoc(), resultType, sourceParts[part], maskParts[part]) .getResult(); + if (!haveAcc) { + accumulator = skipNeutralCombine ? reduced : initParts.front(); + haveAcc = true; + if (skipNeutralCombine) + continue; + } + if (!haveFirstLaneMask) { + firstLaneMask = ensureFirstLaneMask(); + if (failed(firstLaneMask)) + return rewriter.notifyMatchFailure( + op, "failed to create min/max reduction first-lane mask"); + haveFirstLaneMask = true; + } accumulator = rewriter .create(op.getLoc(), resultType, reduced, accumulator, *firstLaneMask) diff --git a/ptodsl/ptodsl/_runtime/native_build.py b/ptodsl/ptodsl/_runtime/native_build.py index 22516c7eb7..aba2c803c6 100644 --- a/ptodsl/ptodsl/_runtime/native_build.py +++ b/ptodsl/ptodsl/_runtime/native_build.py @@ -39,6 +39,14 @@ def _run(cmd: list[str], *, cwd: Path | None = None) -> None: ) +def _pto_flags_from_env() -> list[str]: + """Extra ptoas CLI flags from PTO_FLAGS (space-separated), for A/B harnesses.""" + raw = os.environ.get("PTO_FLAGS", "").strip() + if not raw: + return [] + return raw.split() + + def _run_ptoas( mlir_path: Path, kernel_object: Path, @@ -59,6 +67,7 @@ def _run_ptoas( cmd.append(f"--pto-level={pto_level}") if insert_sync is True: cmd.append("--enable-insert-sync") + cmd.extend(_pto_flags_from_env()) cmd.extend([ "--enable-tile-op-expand", str(mlir_path), @@ -102,6 +111,7 @@ def _compile_config_text( f"pto_level={effective_pto_level}", f"backend={ptoas_overrides.get('backend')}", "enable_tile_op_expand=True", + f"pto_flags={' '.join(_pto_flags_from_env())}", ] ) diff --git a/test/lit/vmi_new/vmi_predicate_fold_general.pto b/test/lit/vmi_new/vmi_predicate_fold_general.pto new file mode 100644 index 0000000000..3e588e905a --- /dev/null +++ b/test/lit/vmi_new/vmi_predicate_fold_general.pto @@ -0,0 +1,94 @@ +// 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: pto-test-opt %s -vmi-predicate-fold | FileCheck %s + +// Generalized predicate / neutral folds beyond pad vsel. + +module { + // CHECK-LABEL: func.func @demask_vadd_all_true + // CHECK-NOT: create_mask + // CHECK: %[[OUT:.*]] = pto.vmi.vadd %{{.*}}, %{{.*}} : + // CHECK-SAME: !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> -> !pto.vmi.vreg<64xf32> + // CHECK: return %[[OUT]] + func.func @demask_vadd_all_true( + %a: !pto.vmi.vreg<64xf32>, + %b: !pto.vmi.vreg<64xf32>) + -> !pto.vmi.vreg<64xf32> { + %c64i = arith.constant 64 : index + %m = pto.vmi.create_mask %c64i : index -> !pto.vmi.mask<64xpred> + %out = pto.vmi.vadd %a, %b, %m + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xf32> + return %out : !pto.vmi.vreg<64xf32> + } + + // CHECK-LABEL: func.func @neutral_vmax_neginf + // CHECK-NOT: vbrc + // CHECK-NOT: vmax + // CHECK: return %[[X:.*]] : !pto.vmi.vreg<64xf32> + func.func @neutral_vmax_neginf( + %x: !pto.vmi.vreg<64xf32>) + -> !pto.vmi.vreg<64xf32> { + %c64i = arith.constant 64 : index + %m = pto.vmi.create_mask %c64i : index -> !pto.vmi.mask<64xpred> + %ni = arith.constant 0xFF800000 : f32 + %neg = pto.vmi.vbrc %ni : f32 -> !pto.vmi.vreg<64xf32> + %out = pto.vmi.vmax %neg, %x, %m + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xf32> + return %out : !pto.vmi.vreg<64xf32> + } + + // CHECK-LABEL: func.func @neutral_vadd_zero + // CHECK-NOT: vbrc + // CHECK-NOT: vadd + // CHECK: return %[[X:.*]] : !pto.vmi.vreg<64xf32> + func.func @neutral_vadd_zero( + %x: !pto.vmi.vreg<64xf32>) + -> !pto.vmi.vreg<64xf32> { + %c64i = arith.constant 64 : index + %m = pto.vmi.create_mask %c64i : index -> !pto.vmi.mask<64xpred> + %z = arith.constant 0.0 : f32 + %zero = pto.vmi.vbrc %z : f32 -> !pto.vmi.vreg<64xf32> + %out = pto.vmi.vadd %zero, %x, %m + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xf32> + return %out : !pto.vmi.vreg<64xf32> + } + + // CHECK-LABEL: func.func @vdhist_all_false + // CHECK-NOT: vdhist + // CHECK: return %[[ACC:.*]] : !pto.vmi.vreg<128xui16> + func.func @vdhist_all_false( + %acc: !pto.vmi.vreg<128xui16>, + %src: !pto.vmi.vreg<256xui8>) + -> !pto.vmi.vreg<128xui16> { + %c0 = arith.constant 0 : index + %m = pto.vmi.create_mask %c0 : index -> !pto.vmi.mask<256xpred> + %out = pto.vmi.vdhist %acc, %src, %m + : !pto.vmi.vreg<128xui16>, !pto.vmi.vreg<256xui8>, !pto.vmi.mask<256xpred> + -> !pto.vmi.vreg<128xui16> + return %out : !pto.vmi.vreg<128xui16> + } + + // CHECK-LABEL: func.func @partial_mask_no_demask + // CHECK: create_mask + // CHECK: vadd + func.func @partial_mask_no_demask( + %a: !pto.vmi.vreg<64xf32>, + %b: !pto.vmi.vreg<64xf32>) + -> !pto.vmi.vreg<64xf32> { + %c48 = arith.constant 48 : index + %m = pto.vmi.create_mask %c48 : index -> !pto.vmi.mask<64xpred> + %out = pto.vmi.vadd %a, %b, %m + : !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xf32> + return %out : !pto.vmi.vreg<64xf32> + } +} diff --git a/test/lit/vmi_new/vmi_predicate_fold_pad.pto b/test/lit/vmi_new/vmi_predicate_fold_pad.pto new file mode 100644 index 0000000000..e1c0e11fb4 --- /dev/null +++ b/test/lit/vmi_new/vmi_predicate_fold_pad.pto @@ -0,0 +1,172 @@ +// 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: pto-test-opt %s -vmi-predicate-fold | FileCheck %s + +// Expert-pad style folds: vsel(vcmp(vci/vadds, vbrc(E), lt), score, neg_inf) + +module { + // CHECK-LABEL: func.func @pad_all_true_vci + // CHECK-NOT: pto.vmi.vcmp + // CHECK-NOT: pto.vmi.vsel + // CHECK: return %[[SCORE:.*]] : !pto.vmi.vreg<64xf32> + func.func @pad_all_true_vci( + %score: !pto.vmi.vreg<64xf32>, + %neg_inf: !pto.vmi.vreg<64xf32>) + -> !pto.vmi.vreg<64xf32> { + %c0 = arith.constant 0 : i32 + %c64 = arith.constant 64 : i32 + %c64i = arith.constant 64 : index + %seed = pto.vmi.create_mask %c64i : index -> !pto.vmi.mask<64xpred> + %idx = pto.vmi.vci %c0 : i32 -> !pto.vmi.vreg<64xi32> + %num_exp = pto.vmi.vbrc %c64 : i32 -> !pto.vmi.vreg<64xi32> + %m = pto.vmi.vcmp %idx, %num_exp, %seed {cmp = "lt"} + : !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.mask<64xpred> + %out = pto.vmi.vsel %m, %score, %neg_inf + : !pto.vmi.mask<64xpred>, !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> + -> !pto.vmi.vreg<64xf32> + return %out : !pto.vmi.vreg<64xf32> + } + + // CHECK-LABEL: func.func @pad_all_false_vci + // CHECK-NOT: pto.vmi.vcmp + // CHECK-NOT: pto.vmi.vsel + // CHECK: return %[[NEG:.*]] : !pto.vmi.vreg<64xf32> + func.func @pad_all_false_vci( + %score: !pto.vmi.vreg<64xf32>, + %neg_inf: !pto.vmi.vreg<64xf32>) + -> !pto.vmi.vreg<64xf32> { + %c64 = arith.constant 64 : i32 + %c64i = arith.constant 64 : index + %seed = pto.vmi.create_mask %c64i : index -> !pto.vmi.mask<64xpred> + %idx = pto.vmi.vci %c64 : i32 -> !pto.vmi.vreg<64xi32> + %num_exp = pto.vmi.vbrc %c64 : i32 -> !pto.vmi.vreg<64xi32> + %m = pto.vmi.vcmp %idx, %num_exp, %seed {cmp = "lt"} + : !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.mask<64xpred> + %out = pto.vmi.vsel %m, %score, %neg_inf + : !pto.vmi.mask<64xpred>, !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> + -> !pto.vmi.vreg<64xf32> + return %out : !pto.vmi.vreg<64xf32> + } + + // CHECK-LABEL: func.func @pad_mixed_no_fold + // CHECK: pto.vmi.vcmp + // CHECK: pto.vmi.vsel + func.func @pad_mixed_no_fold( + %score: !pto.vmi.vreg<64xf32>, + %neg_inf: !pto.vmi.vreg<64xf32>) + -> !pto.vmi.vreg<64xf32> { + %c32 = arith.constant 32 : i32 + %c64 = arith.constant 64 : i32 + %c64i = arith.constant 64 : index + %seed = pto.vmi.create_mask %c64i : index -> !pto.vmi.mask<64xpred> + %idx = pto.vmi.vci %c32 : i32 -> !pto.vmi.vreg<64xi32> + %num_exp = pto.vmi.vbrc %c64 : i32 -> !pto.vmi.vreg<64xi32> + %m = pto.vmi.vcmp %idx, %num_exp, %seed {cmp = "lt"} + : !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.mask<64xpred> + %out = pto.vmi.vsel %m, %score, %neg_inf + : !pto.vmi.mask<64xpred>, !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> + -> !pto.vmi.vreg<64xf32> + return %out : !pto.vmi.vreg<64xf32> + } + + // CHECK-LABEL: func.func @pad_vadds_all_false + // CHECK-NOT: pto.vmi.vcmp + // CHECK-NOT: pto.vmi.vsel + // CHECK: return %[[NEG:.*]] : !pto.vmi.vreg<64xf32> + func.func @pad_vadds_all_false( + %score: !pto.vmi.vreg<64xf32>, + %neg_inf: !pto.vmi.vreg<64xf32>) + -> !pto.vmi.vreg<64xf32> { + %c0 = arith.constant 0 : i32 + %c384 = arith.constant 384 : i32 + %c64i = arith.constant 64 : index + %seed = pto.vmi.create_mask %c64i : index -> !pto.vmi.mask<64xpred> + %iota0 = pto.vmi.vci %c0 : i32 -> !pto.vmi.vreg<64xi32> + %idx = pto.vmi.vadds %iota0, %c384, %seed + : !pto.vmi.vreg<64xi32>, i32, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xi32> + %num_exp = pto.vmi.vbrc %c384 : i32 -> !pto.vmi.vreg<64xi32> + %m = pto.vmi.vcmp %idx, %num_exp, %seed {cmp = "lt"} + : !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.mask<64xpred> + %out = pto.vmi.vsel %m, %score, %neg_inf + : !pto.vmi.mask<64xpred>, !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> + -> !pto.vmi.vreg<64xf32> + return %out : !pto.vmi.vreg<64xf32> + } + + // CHECK-LABEL: func.func @pad_affine_multipass_all_true + // iv in [0,1], pass_base = iv*384, index = vci(pass_base) → [0..447] < 768 + // CHECK-NOT: pto.vmi.vcmp + // CHECK-NOT: pto.vmi.vsel + func.func @pad_affine_multipass_all_true( + %score: !pto.vmi.vreg<64xf32>, + %neg_inf: !pto.vmi.vreg<64xf32>) + -> !pto.vmi.vreg<64xf32> { + %c0 = arith.constant 0 : index + %c2 = arith.constant 2 : index + %c1 = arith.constant 1 : index + %c384 = arith.constant 384 : i32 + %c768 = arith.constant 768 : i32 + %c64i = arith.constant 64 : index + %seed = pto.vmi.create_mask %c64i : index -> !pto.vmi.mask<64xpred> + %num_exp = pto.vmi.vbrc %c768 : i32 -> !pto.vmi.vreg<64xi32> + %out = scf.for %iv = %c0 to %c2 step %c1 + iter_args(%acc = %score) -> (!pto.vmi.vreg<64xf32>) { + %iv_i32 = arith.index_cast %iv : index to i32 + %pass_base = arith.muli %iv_i32, %c384 : i32 + %idx = pto.vmi.vci %pass_base : i32 -> !pto.vmi.vreg<64xi32> + %m = pto.vmi.vcmp %idx, %num_exp, %seed {cmp = "lt"} + : !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.mask<64xpred> + %sel = pto.vmi.vsel %m, %acc, %neg_inf + : !pto.vmi.mask<64xpred>, !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> + -> !pto.vmi.vreg<64xf32> + scf.yield %sel : !pto.vmi.vreg<64xf32> + } + return %out : !pto.vmi.vreg<64xf32> + } + + // CHECK-LABEL: func.func @pad_dynamic_num_exp_no_fold + // CHECK: pto.vmi.vcmp + // CHECK: pto.vmi.vsel + func.func @pad_dynamic_num_exp_no_fold( + %score: !pto.vmi.vreg<64xf32>, + %neg_inf: !pto.vmi.vreg<64xf32>, + %e: i32) + -> !pto.vmi.vreg<64xf32> { + %c0 = arith.constant 0 : i32 + %c64i = arith.constant 64 : index + %seed = pto.vmi.create_mask %c64i : index -> !pto.vmi.mask<64xpred> + %idx = pto.vmi.vci %c0 : i32 -> !pto.vmi.vreg<64xi32> + %num_exp = pto.vmi.vbrc %e : i32 -> !pto.vmi.vreg<64xi32> + %m = pto.vmi.vcmp %idx, %num_exp, %seed {cmp = "lt"} + : !pto.vmi.vreg<64xi32>, !pto.vmi.vreg<64xi32>, !pto.vmi.mask<64xpred> + -> !pto.vmi.mask<64xpred> + %out = pto.vmi.vsel %m, %score, %neg_inf + : !pto.vmi.mask<64xpred>, !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> + -> !pto.vmi.vreg<64xf32> + return %out : !pto.vmi.vreg<64xf32> + } + + // CHECK-LABEL: func.func @vsel_same_arms + // CHECK-NOT: pto.vmi.vsel + func.func @vsel_same_arms( + %m: !pto.vmi.mask<64xpred>, + %x: !pto.vmi.vreg<64xf32>) + -> !pto.vmi.vreg<64xf32> { + %out = pto.vmi.vsel %m, %x, %x + : !pto.vmi.mask<64xpred>, !pto.vmi.vreg<64xf32>, !pto.vmi.vreg<64xf32> + -> !pto.vmi.vreg<64xf32> + return %out : !pto.vmi.vreg<64xf32> + } +} diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_addf.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_addf.pto index 17c72f287f..d5634d6452 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_addf.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_addf.pto @@ -26,7 +26,7 @@ module { // CHECK-LABEL: func.func @vmi_to_vpto_reduce_addf( // CHECK: %[[REDUCED:.*]] = pto.vcadd %arg0, %arg1 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: return %[[REDUCED]] : !pto.vreg<64xf32> +// CHECK: return %[[REDUCED]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_addf_f16.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_addf_f16.pto index 5c14eb4b93..8fc3d80fae 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_addf_f16.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_addf_f16.pto @@ -27,7 +27,8 @@ module { // CHECK-LABEL: func.func @vmi_to_vpto_reduce_addf_f16( // CHECK: %[[REDUCED:.*]] = pto.vcadd %arg0, %arg1 // CHECK-SAME: !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<128xf16> -// CHECK: return {{.*}} : !pto.vreg<128xf16> +// CHECK-NOT: pto.vadd %[[REDUCED]] +// CHECK: return %[[REDUCED]] : !pto.vreg<128xf16> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_addf_multichunk.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_addf_multichunk.pto index d2e6803f4b..8a981f171a 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_addf_multichunk.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_addf_multichunk.pto @@ -42,9 +42,8 @@ module { // CHECK-LABEL: func.func @vmi_to_vpto_reduce_addf_multichunk( // CHECK: %[[RED0:.*]] = pto.vcadd %arg0, %arg2 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[RED1:.*]] = pto.vcadd %arg1, %arg3 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: pto.vadd %[[RED1]], %[[RED0]], %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: pto.vadd %[[RED1]], %[[RED0]], {{.*}} : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast @@ -53,7 +52,7 @@ module { // CHECK: %[[MASK0:.*]] = pto.pset_b32 "PAT_ALL" : !pto.mask // CHECK: %[[MERGED:.*]] = pto.vadd %arg{{[01]}}, %arg{{[01]}}, %[[MASK0]] // CHECK: %[[REDUCED:.*]] = pto.vcadd %[[MERGED]], %[[MASK0]] -// CHECK-NOT: pto.vadd %[[REDUCED]], +// CHECK: return %[[REDUCED]] // CHECK-NOT: pto.vcadd // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_addi.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_addi.pto index dfbab711d5..403416b0c5 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_addi.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_addi.pto @@ -26,7 +26,7 @@ module { // CHECK-LABEL: func.func @vmi_to_vpto_reduce_addi( // CHECK: %[[REDUCED:.*]] = pto.vcadd %arg0, %arg1 : !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> -// CHECK: return %[[REDUCED]] : !pto.vreg<64xi32> +// CHECK: return %[[REDUCED]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_addi_multichunk.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_addi_multichunk.pto index 65af87b292..c31272a355 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_addi_multichunk.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_addi_multichunk.pto @@ -42,9 +42,8 @@ module { // CHECK-LABEL: func.func @vmi_to_vpto_reduce_addi_multichunk( // CHECK: %[[RED0:.*]] = pto.vcadd %arg0, %arg2 : !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> -// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[RED1:.*]] = pto.vcadd %arg1, %arg3 : !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> -// CHECK: pto.vadd %[[RED1]], %[[RED0]], %[[FIRST]] : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> +// CHECK: pto.vadd %[[RED1]], %[[RED0]], {{.*}} : !pto.vreg<64xi32>, !pto.vreg<64xi32>, !pto.mask -> !pto.vreg<64xi32> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast @@ -53,7 +52,7 @@ module { // CHECK: %[[MASK0:.*]] = pto.pset_b32 "PAT_ALL" : !pto.mask // CHECK: %[[MERGED:.*]] = pto.vadd %arg{{[01]}}, %arg{{[01]}}, %[[MASK0]] // CHECK: %[[REDUCED:.*]] = pto.vcadd %[[MERGED]], %[[MASK0]] -// CHECK-NOT: pto.vadd %[[REDUCED]], +// CHECK: return %[[REDUCED]] // CHECK-NOT: pto.vcadd // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_maxf_multichunk.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_maxf_multichunk.pto index b01c2ee675..44111fa2f6 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_maxf_multichunk.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_maxf_multichunk.pto @@ -56,18 +56,16 @@ module { // CHECK-LABEL: func.func @vmi_to_vpto_reduce_maxf_multichunk( // CHECK: %[[RED0:.*]] = pto.vcmax %arg0, %arg2 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[RED1:.*]] = pto.vcmax %arg1, %arg3 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: pto.vmax %[[RED1]], %[[RED0]], %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: pto.vmax %[[RED1]], %[[RED0]], {{.*}} : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast // CHECK-LABEL: func.func @vmi_to_vpto_reduce_minf_multichunk( // CHECK: %[[RED0:.*]] = pto.vcmin %arg0, %arg2 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: %[[FIRST:.*]] = pto.pset_b32 "PAT_VL1" : !pto.mask // CHECK: %[[RED1:.*]] = pto.vcmin %arg1, %arg3 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> -// CHECK: pto.vmin %[[RED1]], %[[RED0]], %[[FIRST]] : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> +// CHECK: pto.vmin %[[RED1]], %[[RED0]], {{.*}} : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast @@ -76,7 +74,7 @@ module { // CHECK: %[[MASK0:.*]] = pto.pset_b32 "PAT_ALL" : !pto.mask // CHECK: %[[MERGED:.*]] = pto.vmax %arg{{[01]}}, %arg{{[01]}}, %[[MASK0]] // CHECK: %[[REDUCED:.*]] = pto.vcmax %[[MERGED]], %[[MASK0]] -// CHECK-NOT: pto.vmax %[[REDUCED]], +// CHECK: return %[[REDUCED]] // CHECK-NOT: pto.vcmax // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_reduce_minf.pto b/test/lit/vmi_new/vmi_to_vpto_reduce_minf.pto index d5be772ee9..f8c35e9232 100644 --- a/test/lit/vmi_new/vmi_to_vpto_reduce_minf.pto +++ b/test/lit/vmi_new/vmi_to_vpto_reduce_minf.pto @@ -26,7 +26,7 @@ module { // CHECK-LABEL: func.func @vmi_to_vpto_reduce_minf( // CHECK: %[[REDUCED:.*]] = pto.vcmin %arg0, %arg1 : !pto.vreg<128xf16>, !pto.mask -> !pto.vreg<128xf16> -// CHECK: return %[[REDUCED]] : !pto.vreg<128xf16> +// CHECK: return %[[REDUCED]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 8ebd2b9e42..9de8b3802f 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -596,6 +596,13 @@ static llvm::cl::opt disableInferLayout( llvm::cl::desc("Disable PTO layout inference pass (static-only)"), llvm::cl::init(false)); +static llvm::cl::opt disableVMIPredicateFold( + "disable-vmi-predicate-fold", + llvm::cl::desc( + "Disable VMIPredicateFold (A/B: keep statically-proven pad " + "vcmp/vsel that would otherwise DCE)"), + llvm::cl::init(false)); + static llvm::cl::opt enableSoftPostUpdate( "enable-vpto-soft-postupdate", llvm::cl::desc("Enable VPTO soft post-update optimization"), @@ -3265,10 +3272,15 @@ static void appendVMISemanticPipeline(OpPassManager &pm) { // verifier, layout, or lowering pass sees signless integer element types. pm.addNestedPass( pto::createVMINormalizeSignlessIntToUnsignedPass()); + // Fold statically proven vcmp/vsel (e.g. expert-pad when E covers indices) + // before unified→legacy lowering so dead pad work never reaches layout. + if (!disableVMIPredicateFold) + pm.addPass(pto::createVMIPredicateFoldPass()); // Expand unified VMI ops before layout assignment so grouped vci becomes // the contiguous-only legacy group_iota producer. Layout assignment can // then materialize any consumer-requested non-contiguous use explicitly. pm.addPass(pto::createVMILowerUnifiedToLegacyPass()); + pm.addPass(createCanonicalizerPass()); pm.addPass(pto::createVMILegalizeArithSelectPass()); pm.addPass(pto::createPTOValidateVMIIRPass());