[Mlir-commits] [mlir] [mlir][vector] Refine select uses on uniform reduction paths (PR #235640)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 29 12:48:57 PDT 2026
https://github.com/meinie0826 updated https://github.com/llvm/llvm-project/pull/235640
>From 2c7b87c5b00ea984d95feac3e4d542b8da8f5d19 Mon Sep 17 00:00:00 2001
From: meinie0826 <929449726 at qq.com>
Date: Mon, 26 Oct 2026 15:52:17 +0800
Subject: [PATCH 1/2] [mlir][vector] Refine select uses on uniform reduction
paths
Use region-branch reachability to prove AND-true and OR-false mask uniformity at individual select uses, preserving opposite and loop-carried paths.
---
mlir/include/mlir/Dialect/Vector/IR/Vector.td | 1 +
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 104 ++++
mlir/test/Dialect/Vector/canonicalize.mlir | 554 ++++++++++++++++++
3 files changed, 659 insertions(+)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/Vector.td b/mlir/include/mlir/Dialect/Vector/IR/Vector.td
index 5125ae7c137178..45ebb7e83d8d48 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/Vector.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/Vector.td
@@ -21,6 +21,7 @@ def Vector_Dialect : Dialect {
let useDefaultAttributePrinterParser = 1;
let hasConstantMaterializer = 1;
+ let hasCanonicalizer = 1;
let dependentDialects = [
"arith::ArithDialect",
"ub::UBDialect"
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index eedf483df2d312..872227cc0f9f33 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -37,6 +37,7 @@
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/ValueRange.h"
+#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/SubsetOpInterface.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "mlir/Support/LLVM.h"
@@ -44,6 +45,7 @@
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/Repeated.h"
#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/ADT/StringSet.h"
@@ -8692,6 +8694,108 @@ DeinterleaveOp::getIterationShapeToUnroll() {
return llvm::to_vector<4>(getResultVectorType().getShape());
}
+//===----------------------------------------------------------------------===//
+// VectorDialect canonicalization
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// Conservatively compute the regions reachable from an operation entry.
+/// Region-to-region edges are included: an entry-only check is insufficient for
+/// loops and operations that enter one region through another.
+static bool isRegionReachable(RegionBranchOpInterface branch, Region *target,
+ ArrayRef<Attribute> operands) {
+ SmallVector<RegionSuccessor> successors;
+ branch.getEntrySuccessorRegions(operands, successors);
+ SmallVector<Region *> worklist;
+ llvm::SmallPtrSet<Region *, 4> visited;
+ auto append = [&] {
+ for (RegionSuccessor successor : successors)
+ if (Region *region = successor.getSuccessor())
+ if (region->getParentOp() == branch.getOperation() &&
+ visited.insert(region).second)
+ worklist.push_back(region);
+ };
+ append();
+ while (!worklist.empty()) {
+ Region *region = worklist.pop_back_val();
+ if (region == target)
+ return true;
+ successors.clear();
+ branch.getSuccessorRegions(*region, successors);
+ append();
+ }
+ return false;
+}
+
+/// Infer a reduction's value only when the opposite value cannot reach the use.
+struct SelectReductionBranch : OpRewritePattern<arith::SelectOp> {
+ using OpRewritePattern::OpRewritePattern;
+ LogicalResult matchAndRewrite(arith::SelectOp select,
+ PatternRewriter &rewriter) const override {
+ if (!isa<VectorType>(select.getCondition().getType()))
+ return failure();
+ DenseMap<Region *, Value> replacements;
+ bool changed = false;
+ for (OpOperand &use :
+ llvm::make_early_inc_range(select.getResult().getUses())) {
+ for (Operation *child = use.getOwner();
+ auto *parent = child->getParentOp(); child = parent) {
+ auto branch = dyn_cast<RegionBranchOpInterface>(parent);
+ if (!branch)
+ continue;
+ Region *useRegion = child->getParentRegion();
+ if (!replacements.contains(useRegion)) {
+ replacements[useRegion] = Value();
+ for (Value operand : parent->getOperands()) {
+ auto reduction = operand.getDefiningOp<ReductionOp>();
+ if (!reduction ||
+ !reduction.getDest().getType().isSignlessInteger(1) ||
+ reduction.getAcc() || reduction.isMasked() ||
+ reduction.getVector() != select.getCondition())
+ continue;
+ bool value;
+ if (reduction.getKind() == CombiningKind::AND)
+ value = true;
+ else if (reduction.getKind() == CombiningKind::OR)
+ value = false;
+ else
+ continue;
+ SmallVector<Attribute> knownOperands(parent->getNumOperands());
+ // Every occurrence of this SSA value must have the same constant.
+ for (auto [i, v] : llvm::enumerate(parent->getOperands()))
+ if (v == operand)
+ knownOperands[i] = rewriter.getBoolAttr(!value);
+ if (isRegionReachable(branch, useRegion, knownOperands))
+ continue;
+ for (auto [i, v] : llvm::enumerate(parent->getOperands()))
+ if (v == operand)
+ knownOperands[i] = rewriter.getBoolAttr(value);
+ if (!isRegionReachable(branch, useRegion, knownOperands))
+ continue;
+ Value replacement =
+ value ? select.getTrueValue() : select.getFalseValue();
+ replacements[useRegion] = replacement;
+ break;
+ }
+ }
+ if (Value replacement = replacements.lookup(useRegion)) {
+ rewriter.modifyOpInPlace(use.getOwner(),
+ [&] { use.set(replacement); });
+ changed = true;
+ break;
+ }
+ }
+ }
+ return success(changed);
+ }
+};
+} // namespace
+
+void VectorDialect::getCanonicalizationPatterns(
+ RewritePatternSet &results) const {
+ results.add<SelectReductionBranch>(getContext());
+}
+
//===----------------------------------------------------------------------===//
// TableGen'd op method definitions
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 4eb3988c417ec7..32b7094ecb6f7c 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -4677,3 +4677,557 @@ func.func @interleave_splat() -> vector<[4]xf32> {
%0 = vector.interleave %cst, %cst : vector<[2]xf32> -> vector<[4]xf32>
return %0 : vector<[4]xf32>
}
+
+// -----
+
+// Reduction predicates refine select uses only on paths that require a
+// uniform mask. Opposite branches and outside uses retain elementwise selects.
+// CHECK-LABEL: func.func @and_acc(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if
+// CHECK: %[[LOCAL:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[LOCAL]])
+// CHECK: } else {
+// CHECK: %[[ELSE:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[ELSE]])
+// CHECK: return %{{.*}} : vector<4xi32>
+func.func @and_acc(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <and>, %m, %acc : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ } else {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @and_changed_loop_mask(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: scf.for
+// CHECK: %[[X:.*]] = arith.select %{{.*}}, %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[X]])
+// CHECK: return %{{.*}} : vector<4xi32>
+func.func @and_changed_loop_mask(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <and>, %m : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ scf.if %r {
+ %result = scf.for %i = %c0 to %n step %c1 iter_args(%iter = %m) -> (vector<4xi1>) {
+ %x = arith.select %iter, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ scf.yield %other : vector<4xi1>
+ }
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @and_execute(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[R:.*]] = vector.reduction <and>, %[[M]]
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if %[[R]]
+// CHECK: func.call @consume(%[[A]])
+// CHECK: } else {
+// CHECK: func.call @consume(%{{.*}})
+func.func @and_execute(%m: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>) {
+ %r = vector.reduction <and>, %m : vector<4xi1> into i1
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ scf.execute_region {
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ scf.yield
+ }
+ } else {
+ scf.execute_region {
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ scf.yield
+ }
+ }
+ return
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @and_freeze(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if
+// CHECK: %[[LOCAL:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[LOCAL]])
+// CHECK: } else {
+// CHECK: %[[ELSE:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[ELSE]])
+// CHECK: return %{{.*}} : vector<4xi32>
+func.func @and_freeze(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <and>, %m : vector<4xi1> into i1
+ %f = "llvm.freeze"(%r) : (i1) -> i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %f {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ } else {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @and_inside(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[R:.*]] = vector.reduction <and>, %[[M]]
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if %[[R]]
+// CHECK: func.call @consume(%[[A]])
+// CHECK: } else {
+// CHECK: %[[LOCAL:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[LOCAL]])
+// CHECK: return %[[X]] : vector<4xi32>
+func.func @and_inside(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <and>, %m : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ } else {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @and_loop(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[R:.*]] = vector.reduction <and>, %[[M]]
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if %[[R]]
+// CHECK: func.call @consume(%[[A]])
+// CHECK: } else {
+// CHECK: %[[LOCAL:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[LOCAL]])
+// CHECK: return %[[X]] : vector<4xi32>
+func.func @and_loop(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <and>, %m : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ scf.for %i = %c0 to %n step %c1 {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ } else {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ scf.for %i = %c0 to %n step %c1 {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @and_masked(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if
+// CHECK: %[[LOCAL:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[LOCAL]])
+// CHECK: } else {
+// CHECK: %[[ELSE:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[ELSE]])
+// CHECK: return %{{.*}} : vector<4xi32>
+func.func @and_masked(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.mask %other { vector.reduction <and>, %m : vector<4xi1> into i1 } : vector<4xi1> -> i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ } else {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @and_nested(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[R:.*]] = vector.reduction <and>, %[[M]]
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if %[[R]]
+// CHECK: func.call @consume(%[[A]])
+// CHECK: } else {
+// CHECK: %[[LOCAL:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[LOCAL]])
+// CHECK: return %[[X]] : vector<4xi32>
+func.func @and_nested(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <and>, %m : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ scf.if %acc {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ } else {
+ scf.if %acc {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @and_other_mask(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[X:.*]] = arith.select %[[OTHER]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[X]])
+// CHECK: func.call @consume(%[[X]])
+// CHECK: return %{{.*}} : vector<4xi32>
+func.func @and_other_mask(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <and>, %m : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ %x = arith.select %other, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ } else {
+ %x = arith.select %other, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @and_outside(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[R:.*]] = vector.reduction <and>, %[[M]]
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if %[[R]]
+// CHECK: func.call @consume(%[[A]])
+// CHECK: } else {
+// CHECK: func.call @consume(%{{.*}})
+// CHECK: return %[[X]] : vector<4xi32>
+func.func @and_outside(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <and>, %m : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ func.call @consume(%outside) : (vector<4xi32>) -> ()
+ } else {
+ func.call @consume(%outside) : (vector<4xi32>) -> ()
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @while_entry(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[X]])
+func.func @while_entry(%m: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %next: i1) {
+ %r = vector.reduction <and>, %m : vector<4xi1> into i1
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ %result = scf.while (%arg = %r) : (i1) -> i1 {
+ scf.condition(%arg) %arg : i1
+ } do {
+ ^bb0(%arg: i1):
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ scf.yield %next : i1
+ }
+ return
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @or_acc(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if
+// CHECK: %[[LOCAL:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[LOCAL]])
+// CHECK: } else {
+// CHECK: %[[ELSE:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[ELSE]])
+// CHECK: return %{{.*}} : vector<4xi32>
+func.func @or_acc(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <or>, %m, %acc : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ } else {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @or_changed_loop_mask(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: scf.for
+// CHECK: %[[X:.*]] = arith.select %{{.*}}, %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[X]])
+// CHECK: return %{{.*}} : vector<4xi32>
+func.func @or_changed_loop_mask(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <or>, %m : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ scf.if %r {
+ %result = scf.for %i = %c0 to %n step %c1 iter_args(%iter = %m) -> (vector<4xi1>) {
+ %x = arith.select %iter, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ scf.yield %other : vector<4xi1>
+ }
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @or_execute(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[R:.*]] = vector.reduction <or>, %[[M]]
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if %[[R]]
+// CHECK: func.call @consume(%{{.*}})
+// CHECK: } else {
+// CHECK: func.call @consume(%[[B]])
+func.func @or_execute(%m: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>) {
+ %r = vector.reduction <or>, %m : vector<4xi1> into i1
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ scf.execute_region {
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ scf.yield
+ }
+ } else {
+ scf.execute_region {
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ scf.yield
+ }
+ }
+ return
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @or_freeze(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if
+// CHECK: %[[LOCAL:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[LOCAL]])
+// CHECK: } else {
+// CHECK: %[[ELSE:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[ELSE]])
+// CHECK: return %{{.*}} : vector<4xi32>
+func.func @or_freeze(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <or>, %m : vector<4xi1> into i1
+ %f = "llvm.freeze"(%r) : (i1) -> i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %f {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ } else {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @or_inside(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[R:.*]] = vector.reduction <or>, %[[M]]
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if %[[R]]
+// CHECK: %[[LOCAL:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[LOCAL]])
+// CHECK: } else {
+// CHECK: func.call @consume(%[[B]])
+// CHECK: return %[[X]] : vector<4xi32>
+func.func @or_inside(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <or>, %m : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ } else {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @or_loop(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[R:.*]] = vector.reduction <or>, %[[M]]
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if %[[R]]
+// CHECK: %[[LOCAL:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[LOCAL]])
+// CHECK: } else {
+// CHECK: func.call @consume(%[[B]])
+// CHECK: return %[[X]] : vector<4xi32>
+func.func @or_loop(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <or>, %m : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ scf.for %i = %c0 to %n step %c1 {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ } else {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ scf.for %i = %c0 to %n step %c1 {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @or_masked(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if
+// CHECK: %[[LOCAL:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[LOCAL]])
+// CHECK: } else {
+// CHECK: %[[ELSE:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[ELSE]])
+// CHECK: return %{{.*}} : vector<4xi32>
+func.func @or_masked(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.mask %other { vector.reduction <or>, %m : vector<4xi1> into i1 } : vector<4xi1> -> i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ } else {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @or_nested(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[R:.*]] = vector.reduction <or>, %[[M]]
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if %[[R]]
+// CHECK: %[[LOCAL:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[LOCAL]])
+// CHECK: } else {
+// CHECK: func.call @consume(%[[B]])
+// CHECK: return %[[X]] : vector<4xi32>
+func.func @or_nested(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <or>, %m : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ scf.if %acc {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ } else {
+ scf.if %acc {
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @or_other_mask(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[X:.*]] = arith.select %[[OTHER]], %[[A]], %[[B]]
+// CHECK: func.call @consume(%[[X]])
+// CHECK: func.call @consume(%[[X]])
+// CHECK: return %{{.*}} : vector<4xi32>
+func.func @or_other_mask(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <or>, %m : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ %x = arith.select %other, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ } else {
+ %x = arith.select %other, %a, %b : vector<4xi1>, vector<4xi32>
+ func.call @consume(%x) : (vector<4xi32>) -> ()
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
+
+// -----
+
+// CHECK-LABEL: func.func @or_outside(
+// CHECK-SAME: %[[M:[^:]+]]: vector<4xi1>, %[[OTHER:[^:]+]]: vector<4xi1>, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[R:.*]] = vector.reduction <or>, %[[M]]
+// CHECK: %[[X:.*]] = arith.select %[[M]], %[[A]], %[[B]]
+// CHECK: scf.if %[[R]]
+// CHECK: func.call @consume(%{{.*}})
+// CHECK: } else {
+// CHECK: func.call @consume(%[[B]])
+// CHECK: return %[[X]] : vector<4xi32>
+func.func @or_outside(%m: vector<4xi1>, %other: vector<4xi1>, %a: vector<4xi32>, %b: vector<4xi32>, %acc: i1, %n: index) -> vector<4xi32> {
+ %r = vector.reduction <or>, %m : vector<4xi1> into i1
+ %outside = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ scf.if %r {
+ func.call @consume(%outside) : (vector<4xi32>) -> ()
+ } else {
+ func.call @consume(%outside) : (vector<4xi32>) -> ()
+ }
+ return %outside : vector<4xi32>
+}
+func.func private @consume(vector<4xi32>)
>From f778545376b97cb87c250e0de5b8f6a9087cc177 Mon Sep 17 00:00:00 2001
From: meinie0826 <929449726 at qq.com>
Date: Mon, 26 Oct 2026 16:44:38 +0800
Subject: [PATCH 2/2] [mlir][vector] Use try_emplace for region replacement
cache
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 8 ++++----
1 file changed, 4 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 872227cc0f9f33..86676c3a7d4966 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -8744,8 +8744,8 @@ struct SelectReductionBranch : OpRewritePattern<arith::SelectOp> {
if (!branch)
continue;
Region *useRegion = child->getParentRegion();
- if (!replacements.contains(useRegion)) {
- replacements[useRegion] = Value();
+ auto [it, inserted] = replacements.try_emplace(useRegion);
+ if (inserted) {
for (Value operand : parent->getOperands()) {
auto reduction = operand.getDefiningOp<ReductionOp>();
if (!reduction ||
@@ -8774,11 +8774,11 @@ struct SelectReductionBranch : OpRewritePattern<arith::SelectOp> {
continue;
Value replacement =
value ? select.getTrueValue() : select.getFalseValue();
- replacements[useRegion] = replacement;
+ it->second = replacement;
break;
}
}
- if (Value replacement = replacements.lookup(useRegion)) {
+ if (Value replacement = it->second) {
rewriter.modifyOpInPlace(use.getOwner(),
[&] { use.set(replacement); });
changed = true;
More information about the Mlir-commits
mailing list