From c5d5d1cf7dd0776c02664c559a98f9239df59463 Mon Sep 17 00:00:00 2001 From: peanutchan Date: Wed, 5 Aug 2026 14:40:14 +0800 Subject: [PATCH] feat(vmi): add VMIPredicateFold to DCE compile-time expert-pad Fold statically proven all-true/all-false vcmp into identity vsel and DCE dead pad work so frontends can always emit expert-pad masking when E is compile-time known (see #580). Co-authored-by: Cursor --- include/PTO/Transforms/Passes.h | 1 + include/PTO/Transforms/Passes.td | 22 + include/PTO/Transforms/VMIMaskUtils.h | 36 ++ lib/PTO/Transforms/CMakeLists.txt | 2 + .../Transforms/VMILowerUnifiedToLegacy.cpp | 20 +- lib/PTO/Transforms/VMIMaskUtils.cpp | 46 ++ lib/PTO/Transforms/VMIPredicateFold.cpp | 392 ++++++++++++++++++ test/lit/vmi_new/vmi_predicate_fold_pad.pto | 172 ++++++++ tools/ptoas/ptoas.cpp | 12 + 9 files changed, 684 insertions(+), 19 deletions(-) create mode 100644 include/PTO/Transforms/VMIMaskUtils.h create mode 100644 lib/PTO/Transforms/VMIMaskUtils.cpp create mode 100644 lib/PTO/Transforms/VMIPredicateFold.cpp create mode 100644 test/lit/vmi_new/vmi_predicate_fold_pad.pto diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index b951fec0e7..3fe3331146 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -126,6 +126,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 29a43cee84..e0d95c3b26 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -1098,6 +1098,28 @@ def VMINormalizeSignlessIntToUnsigned : let dependentDialects = ["pto::PTODialect"]; } +def VMIPredicateFold : Pass<"vmi-predicate-fold", "ModuleOp"> { + let summary = "Fold statically proven VMI vcmp/vsel predicates (pad DCE)"; + let description = [{ + Constant-proves lane ranges for `vci` / `vadds(vci(0), C)` / `vbrc(C)` + (and simple affine `iv*stride+C` bases with known `scf.for` bounds), then: + + * folds `vcmp`/`vcmps` whose result is all-true or all-false + * rewrites `vsel(all_true, t, f) → t` and `vsel(all_false, t, f) → f` + * rewrites `vsel(m, x, x) → x` + * DCEs unused compare / broadcast 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..0ba08f7a47 --- /dev/null +++ b/include/PTO/Transforms/VMIMaskUtils.h @@ -0,0 +1,36 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +//===- VMIMaskUtils.h - Shared VMI predicate / seed helpers -----*- C++ -*-===// +// +// Helpers shared by VMILowerUnifiedToLegacy and VMIPredicateFold for proving +// that a mask SSA value is statically all-active or all-inactive. +// +//===----------------------------------------------------------------------===// + +#ifndef PTO_TRANSFORMS_VMIMASKUTILS_H +#define PTO_TRANSFORMS_VMIMASKUTILS_H + +#include "mlir/IR/Value.h" + +namespace mlir { +namespace pto { + +/// 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); + +} // namespace pto +} // namespace mlir + +#endif // PTO_TRANSFORMS_VMIMASKUTILS_H diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index f82f31715d..bfde5a597c 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -44,6 +44,8 @@ add_mlir_dialect_library(PTOTransforms VMIControlFlowSupport.cpp VMILegalizeArithSelect.cpp VMIMaskGranularityAssignment.cpp + VMIMaskUtils.cpp + VMIPredicateFold.cpp VMILowerUnifiedToLegacy.cpp VMINormalizeSignlessIntToUnsigned.cpp VMILayoutAssignment.cpp diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp index 2a2344e01d..ef926be4f6 100644 --- a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp +++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp @@ -93,6 +93,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" @@ -304,25 +305,6 @@ 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; -} - /// Prepare a direct reduction result for a unit-stride group store. /// /// Explicit grouped reductions already produce one scalar per group. A full diff --git a/lib/PTO/Transforms/VMIMaskUtils.cpp b/lib/PTO/Transforms/VMIMaskUtils.cpp new file mode 100644 index 0000000000..7a86957eba --- /dev/null +++ b/lib/PTO/Transforms/VMIMaskUtils.cpp @@ -0,0 +1,46 @@ +// 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/IR/BuiltinAttributes.h" + +using namespace mlir; +using namespace mlir::pto; + +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; +} diff --git a/lib/PTO/Transforms/VMIPredicateFold.cpp b/lib/PTO/Transforms/VMIPredicateFold.cpp new file mode 100644 index 0000000000..762d20bb46 --- /dev/null +++ b/lib/PTO/Transforms/VMIPredicateFold.cpp @@ -0,0 +1,392 @@ +// 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 -------===// +// +// Constant-proves lane ranges for index vectors built from vci / vadds / +// vbrc (and simple affine scf.for bases), folds all-true / all-false vcmp +// results into identity / constant vsel, then DCEs dead defs. Primary +// consumer: expert-pad masking when num_experts is a compile-time constant. +// +//===----------------------------------------------------------------------===// + +#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" +#include "llvm/Support/MathExtras.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 { + +//===----------------------------------------------------------------------===// +// Scalar / vector range lattices +//===----------------------------------------------------------------------===// + +struct IntRange { + int64_t lo = 0; + int64_t hi = 0; // inclusive + + static IntRange splat(int64_t c) { return {c, c}; } +}; + +enum class MaskLattice { Unknown, AllTrue, AllFalse }; + +static std::optional matchConstantInt(Value v) { + APInt val; + if (matchPattern(v, m_ConstantInt(&val))) + return val.getSExtValue(); + return std::nullopt; +} + +/// Bound an integer SSA value over known constant / affine forms. +/// Recognizes: Imm, addi/subi with Imm, muli(iv|Imm, Imm), scf.for IV with +/// constant lb/ub/step (inclusive iteration set). +static std::optional matchAffineIntRange(Value v) { + if (auto c = matchConstantInt(v)) + return IntRange::splat(*c); + + // Peel casts that preserve integer magnitude (index ↔ i32/i64). + 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; + // [a,b] - [c,d] = [a-d, b-c] + 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)}; + } + + // scf.for induction variable with constant bounds. + 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; + // Last iterate: lb + n*step < ub. + int64_t last = *lb + ((*ub - 1 - *lb) / *step) * *step; + return IntRange{*lb, last}; + } + } + + return std::nullopt; +} + +/// Index vector lattice: every lane is in [lo, hi] (inclusive). +static std::optional matchVectorLaneRange(Value v) { + auto vty = dyn_cast(v.getType()); + if (!vty) + return std::nullopt; + int64_t vl = vty.getElementCount(); + + // vbrc(C) / broadcast(C) → splat + 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; + } + + // vci(base) / iota(base): continuous → [base, base+VL-1] + // with {group=G}: group size = VL/G → [base, base+GSize-1] + auto matchIotaLike = [&](Value base, std::optional group, + StringRef order) -> std::optional { + if (!order.empty() && order != "ASC") + return std::nullopt; // DESC not handled + auto baseR = matchAffineIntRange(base); + if (!baseR) + return std::nullopt; + // For a concrete or affine-bounded base, index covers + // [baseLo, baseHi + span - 1]. + 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.getGroup()) + group = *g; + 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.getGroup()) + group = *g; + StringRef order = iota.getOrder() ? *iota.getOrder() : StringRef("ASC"); + return matchIotaLike(iota.getBase(), group, order); + } + + // vadds(src, scalar, mask): if seed all-active (or unused merge), shift range + 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)) / similar not required for topk; skip. + 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") { + // Only when both sides are the same splat constant. + 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; +} + +static MaskLattice classifyMaskValue(Value mask) { + if (isAllActiveSeed(mask)) + return MaskLattice::AllTrue; + if (isAllInactiveSeed(mask)) + return MaskLattice::AllFalse; + + 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); + // Seed ANDs with the raw compare (pmode zeroing). + 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; +} + +static bool isTriviallyDeadPureOp(Operation *op) { + if (!op || op->getNumRegions() != 0) + return false; + if (op->hasTrait()) + return false; + if (isa(op)) + return false; + // Only drop side-effect-free ops (VMI compute / mask creators are Pure). + if (!isMemoryEffectFree(op)) + return false; + return llvm::all_of(op->getResults(), + [](Value r) { return r.use_empty(); }); +} + +/// Erase pure unused ops to a fixed point. Safer than recursive Value +/// erase chains when one def is reachable via multiple dead users. +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; + } + } +} + +//===----------------------------------------------------------------------===// +// Pass +//===----------------------------------------------------------------------===// + +struct VMIPredicateFoldPass + : public mlir::pto::impl::VMIPredicateFoldBase { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VMIPredicateFoldPass) + + void runOnOperation() override { + ModuleOp module = getOperation(); + SmallVector sels; + module.walk([&](VMIvSelOp op) { sels.push_back(op); }); + + for (VMIvSelOp sel : llvm::reverse(sels)) { + if (!sel->getBlock()) + continue; + + // vsel(m, x, x) → x + if (sel.getTrueValue() == sel.getFalseValue()) { + sel.getResult().replaceAllUsesWith(sel.getTrueValue()); + sel.erase(); + continue; + } + + MaskLattice lat = classifyMaskValue(sel.getMask()); + if (lat == MaskLattice::Unknown) + continue; + + Value replacement = + lat == MaskLattice::AllTrue ? sel.getTrueValue() : sel.getFalseValue(); + sel.getResult().replaceAllUsesWith(replacement); + sel.erase(); + } + + dcePureUnusedOps(module); + } +}; + +} // namespace + +std::unique_ptr mlir::pto::createVMIPredicateFoldPass() { + return std::make_unique(); +} 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/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index f31c065829..37f6b82a6c 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -590,6 +590,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"), @@ -2893,9 +2900,14 @@ static void appendVMISemanticPipeline(OpPassManager &pm) { // before any verifier, layout, or lowering pass sees them. 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 to legacy ops before layout assignment, // so downstream passes only see legacy ops. pm.addPass(pto::createVMILowerUnifiedToLegacyPass()); + pm.addPass(createCanonicalizerPass()); pm.addPass(pto::createVMILegalizeArithSelectPass()); pm.addPass(pto::createPTOValidateVMIIRPass());