[Mlir-commits] [mlir] [mlir][vector] Refine select uses on uniform reduction paths (PR #235640)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 26 00:52:37 PDT 2026
https://github.com/meinie0826 created https://github.com/llvm/llvm-project/pull/235640
## Summary
Refine uses of an elementwise `arith.select` when a surrounding region can only execute with a uniform mask:
- An unmasked `vector.reduction <and>` of an `i1` mask is true only when every lane is true, so the select's true operand can be used on that path.
- An unmasked `vector.reduction <or>` is false only when every lane is false, so the false operand can be used on that path.
```python
%all = vector.reduction <and>, %mask : vector<4xi1> into i1
%x = arith.select %mask, %a, %b : vector<4xi1>, vector<4xi32>
scf.if %all {
func.call @consume(%x) : (vector<4xi32>) -> ()
} else {
func.call @consume(%x) : (vector<4xi32>) -> ()
}
```
The then call can use `%a`. The else call must retain `%x`: a false AND reduction does not distinguish an all-false mask from a mixed mask. The OR case works only on the false path. Outside uses keep the elementwise select.
## Implementation
Register a select-rooted pattern with the Vector dialect canonicalizer. It uses the existing `RegionBranchOpInterface` dependency, not SCF operation classes or a new dialect dependency. For an enclosing region-branch operation with the reduction result as an operand, query entry successors with that operand set to the opposite boolean value, and traverse region-to-region successors conservatively. Replace the use only if its region is unreachable with the opposite value and reachable with the useful value.
Following region edges matters: a while body may be reached through another region or on a later iteration, even when it is not an entry successor. An initial reduction result forwarded as a loop argument is not automatically a persistent guard. The interface traversal leaves this case unchanged. Positive and negative query results are cached per region for each select.
Matching requires the exact SSA mask, an `i1` AND/OR reduction, no execution mask and no accumulator. A masked reduction constrains only participating lanes and cannot establish whole-mask uniformity. Accumulators are conservatively excluded to keep the scope narrow. A frozen reduction result is not matched. Unhandled or imprecise interface control flow simply prevents the optimization. No reduction, branch, computation, or select definition is moved or cloned; only individual select-result uses are replaced.
## Tests
Add regression cases to the existing Vector `canonicalize.mlir`: branch-local and captured selects, outside and opposite-branch uses, unrelated masks, nested regions, invariant and changed loop masks, masked/accumulator reductions, frozen predicates, and a while-entry counterexample.
Built with assertions enabled using system Clang 14 from upstream base `aaa2c083d57b597673e95eccba536f74c588ae86`.
- 21 focused cases pass with the patch; FileCheck fails with the unmodified reduction-path optimizer.
- Vector/Tensor/Arith/SCF and canonicalize/CSE/SCCP tests: 212 passed, 1 expected failure, no unexpected failures.
- Changed C++ regions pass clang-format 22.1.0; `git diff --check` passes.
These are local compiler regression results, not upstream CI or measured performance improvements.
AI assistance: implementation, tests and description were prepared with OpenAI Codex assistance.
>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] [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>)
More information about the Mlir-commits
mailing list