[Mlir-commits] [mlir] [mlir][vector][tensor] Scalarize select conditions from broadcast and splat (PR #235593)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Oct 25 22:17:41 PDT 2026
https://github.com/meinie0826 created https://github.com/llvm/llvm-project/pull/235593
## Summary
Canonicalize `arith.select` with a mask produced by `vector.broadcast` of a scalar `i1`, or by `tensor.splat`, to a select on the scalar condition.
```python
%mask = vector.broadcast %p : i1 to vector<4xi1>
%result = arith.select %mask, %a, %b : vector<4xi1>, vector<4xi32>
```
becomes:
```python
%result = arith.select %p, %a, %b : vector<4xi32>
```
Every mask element is the same scalar condition, so selecting elementwise is equivalent to selecting the entire vector or tensor. The select keeps its two value operands, result type and attributes. Other mask uses are untouched, and ordinary DCE can remove an unused broadcast/splat.
## Implementation
Register select-rooted patterns with the Vector and Tensor dialect canonicalizers. This keeps producer-specific matching out of Arith and SCF and uses their existing Arith dependencies. The vector pattern requires a scalar `i1` broadcast source: broadcasting a nonuniform vector mask is not sufficient. The tensor pattern recognizes `tensor.splat`, whose input is scalar `i1` when its result is a valid select condition.
The rule does not require an `scf.if`, does not move computation or infer facts from `all/any`, and is independent of #235433. It can expose scalar selects to other simplifications, but this PR does not claim memory-planning or runtime improvements.
## Tests
Add cases to the existing Vector and Tensor `canonicalize.mlir` files. They cover floating-point and integer results, multidimensional and rank-zero shapes, scalable vectors, dynamic tensors, shared mask uses, nonuniform masks, and the existing fully-poisoned-arm refinement.
Built `mlir-opt` from upstream base `aaa2c083d57b597673e95eccba536f74c588ae86` with assertions enabled and system Clang 14.
- Both focused regression files fail FileCheck on the optimizer without these rules and pass with the change (14 cases).
- 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 using the repository style; `git diff --check` passes.
These are local tests, not upstream CI or measured performance results.
## AI assistance
Implementation, tests, and this description were prepared with OpenAI Codex assistance.
>From 53cdd85f56ca7e7deee478d1a11e2086865cb7cc Mon Sep 17 00:00:00 2001
From: meinie0826 <929449726 at qq.com>
Date: Mon, 26 Oct 2026 13:17:23 +0800
Subject: [PATCH] [mlir][vector][tensor] Scalarize select conditions from
broadcast and splat
Select whole vectors and tensors with the scalar predicate when every mask element is its broadcast or splat. Preserve nonuniform masks and shared producer uses.
Assisted-by: OpenAI Codex
---
mlir/include/mlir/Dialect/Vector/IR/Vector.td | 1 +
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 19 +++-
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 23 +++++
mlir/test/Dialect/Tensor/canonicalize.mlir | 97 ++++++++++++++++++
mlir/test/Dialect/Vector/canonicalize.mlir | 98 +++++++++++++++++++
5 files changed, 237 insertions(+), 1 deletion(-)
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/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 2ee6c93b313e69..94b38816c0a742 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -4196,13 +4196,30 @@ struct FoldTensorCastProducerOp
}
};
+namespace {
+/// A splat of a scalar condition selects the same arm in every element.
+struct SelectSplatCondition : OpRewritePattern<arith::SelectOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(arith::SelectOp select,
+ PatternRewriter &rewriter) const override {
+ auto splat = select.getCondition().getDefiningOp<SplatOp>();
+ if (!splat)
+ return failure();
+ rewriter.modifyOpInPlace(
+ select, [&] { select.getConditionMutable().assign(splat.getInput()); });
+ return success();
+ }
+};
+} // namespace
+
//===----------------------------------------------------------------------===//
// TensorDialect
//===----------------------------------------------------------------------===//
void TensorDialect::getCanonicalizationPatterns(
RewritePatternSet &results) const {
- results.add<FoldTensorCastProducerOp>(getContext());
+ results.add<FoldTensorCastProducerOp, SelectSplatCondition>(getContext());
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index eedf483df2d312..c42c5e62a218f7 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -574,6 +574,29 @@ void VectorDialect::initialize() {
declarePromisedInterface<ConvertToLLVMPatternInterface, VectorDialect>();
}
+namespace {
+/// A broadcast of a scalar condition selects the same arm in every lane.
+struct SelectBroadcastCondition : OpRewritePattern<arith::SelectOp> {
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(arith::SelectOp select,
+ PatternRewriter &rewriter) const override {
+ auto broadcast = select.getCondition().getDefiningOp<BroadcastOp>();
+ if (!broadcast || !broadcast.getSourceType().isSignlessInteger(1))
+ return failure();
+ rewriter.modifyOpInPlace(select, [&] {
+ select.getConditionMutable().assign(broadcast.getSource());
+ });
+ return success();
+ }
+};
+} // namespace
+
+void VectorDialect::getCanonicalizationPatterns(
+ RewritePatternSet &results) const {
+ results.add<SelectBroadcastCondition>(getContext());
+}
+
/// Materialize a single constant operation from a given attribute value with
/// the desired resultant type.
Operation *VectorDialect::materializeConstant(OpBuilder &builder,
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index f8d333fc0a90ac..dfe31828f80a57 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -2973,3 +2973,100 @@ func.func @sink_expand_of_cast_preserves_sparse_encoding(
: tensor<?xf32, #sparse_v> into tensor<?x?xf32, #sparse_m>
return %e : tensor<?x?xf32, #sparse_m>
}
+
+// -----
+
+// CHECK-LABEL: func.func @select_splat_condition(
+// CHECK-SAME: %[[P:[^:]+]]: i1, %[[A:[^:]+]]: tensor<4xf32>, %[[B:[^:]+]]: tensor<4xf32>
+// CHECK-NOT: tensor.splat
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : tensor<4xf32>
+// CHECK-NOT: arith.select
+// CHECK: return %[[X]]
+func.func @select_splat_condition(%p: i1, %a: tensor<4xf32>, %b: tensor<4xf32>) -> tensor<4xf32> {
+ %m = tensor.splat %p : tensor<4xi1>
+ %x = arith.select %m, %a, %b : tensor<4xi1>, tensor<4xf32>
+ return %x : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @select_splat_multidimensional(
+// CHECK-SAME: %[[P:[^:]+]]: i1, %[[A:[^:]+]]: tensor<2x4xi32>, %[[B:[^:]+]]: tensor<2x4xi32>
+// CHECK-NOT: tensor.splat
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : tensor<2x4xi32>
+// CHECK-NOT: arith.select
+// CHECK: return %[[X]]
+func.func @select_splat_multidimensional(%p: i1, %a: tensor<2x4xi32>, %b: tensor<2x4xi32>) -> tensor<2x4xi32> {
+ %m = tensor.splat %p : tensor<2x4xi1>
+ %x = arith.select %m, %a, %b : tensor<2x4xi1>, tensor<2x4xi32>
+ return %x : tensor<2x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @select_splat_dynamic(
+// CHECK-SAME: %[[P:[^:]+]]: i1, %[[A:[^:]+]]: tensor<?x4xi32>, %[[B:[^:]+]]: tensor<?x4xi32>
+// CHECK-NOT: tensor.splat
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : tensor<?x4xi32>
+// CHECK-NOT: arith.select
+// CHECK: return %[[X]]
+func.func @select_splat_dynamic(%p: i1, %a: tensor<?x4xi32>, %b: tensor<?x4xi32>, %n: index) -> tensor<?x4xi32> {
+ %m = tensor.splat %p[%n] : tensor<?x4xi1>
+ %x = arith.select %m, %a, %b : tensor<?x4xi1>, tensor<?x4xi32>
+ return %x : tensor<?x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @select_splat_rank_zero(
+// CHECK-SAME: %[[P:[^:]+]]: i1, %[[A:[^:]+]]: tensor<i32>, %[[B:[^:]+]]: tensor<i32>
+// CHECK-NOT: tensor.splat
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : tensor<i32>
+// CHECK-NOT: arith.select
+// CHECK: return %[[X]]
+func.func @select_splat_rank_zero(%p: i1, %a: tensor<i32>, %b: tensor<i32>) -> tensor<i32> {
+ %m = tensor.splat %p : tensor<i1>
+ %x = arith.select %m, %a, %b : tensor<i1>, tensor<i32>
+ return %x : tensor<i32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @select_splat_shared_mask(
+// CHECK-SAME: %[[P:[^:]+]]: i1, %[[A:[^:]+]]: tensor<4xi32>, %[[B:[^:]+]]: tensor<4xi32>
+// CHECK: %[[M:.*]] = tensor.splat %[[P]]
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : tensor<4xi32>
+// CHECK-NOT: arith.select
+// CHECK: return %[[X]], %[[M]]
+func.func @select_splat_shared_mask(%p: i1, %a: tensor<4xi32>, %b: tensor<4xi32>) -> (tensor<4xi32>, tensor<4xi1>) {
+ %m = tensor.splat %p : tensor<4xi1>
+ %x = arith.select %m, %a, %b : tensor<4xi1>, tensor<4xi32>
+ return %x, %m : tensor<4xi32>, tensor<4xi1>
+}
+
+// -----
+
+// An arbitrary tensor mask is not known to be a splat.
+// CHECK-LABEL: func.func @select_nonuniform_tensor(
+// CHECK-SAME: %[[P:[^:]+]]: tensor<4xi1>
+// CHECK: arith.select %[[P]], %{{.*}}, %{{.*}} : tensor<4xi1>, tensor<4xi32>
+func.func @select_nonuniform_tensor(%p: tensor<4xi1>, %a: tensor<4xi32>, %b: tensor<4xi32>) -> tensor<4xi32> {
+ %x = arith.select %p, %a, %b : tensor<4xi1>, tensor<4xi32>
+ return %x : tensor<4xi32>
+}
+
+// -----
+
+// Preserve the existing refinement of a select with a fully poisoned arm.
+// CHECK-LABEL: func.func @select_uniform_unselected_poison_tensor(
+// CHECK-SAME: %[[P:[^:]+]]: i1
+// CHECK-SAME: %[[A:[^:]+]]: tensor<4xi32>
+// CHECK-NOT: tensor.splat
+// CHECK-NOT: arith.select
+// CHECK: return %[[A]] : tensor<4xi32>
+func.func @select_uniform_unselected_poison_tensor(%p: i1, %a: tensor<4xi32>) -> tensor<4xi32> {
+ %bad = ub.poison : tensor<4xi32>
+ %m = tensor.splat %p : tensor<4xi1>
+ %x = arith.select %m, %a, %bad : tensor<4xi1>, tensor<4xi32>
+ return %x : tensor<4xi32>
+}
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 4eb3988c417ec7..5211116576f3e1 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -4677,3 +4677,101 @@ func.func @interleave_splat() -> vector<[4]xf32> {
%0 = vector.interleave %cst, %cst : vector<[2]xf32> -> vector<[4]xf32>
return %0 : vector<[4]xf32>
}
+
+// -----
+
+// CHECK-LABEL: func.func @select_broadcast_condition(
+// CHECK-SAME: %[[P:[^:]+]]: i1, %[[A:[^:]+]]: vector<4xf32>, %[[B:[^:]+]]: vector<4xf32>
+// CHECK-NOT: vector.broadcast
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : vector<4xf32>
+// CHECK-NOT: arith.select
+// CHECK: return %[[X]]
+func.func @select_broadcast_condition(%p: i1, %a: vector<4xf32>, %b: vector<4xf32>) -> vector<4xf32> {
+ %m = vector.broadcast %p : i1 to vector<4xi1>
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xf32>
+ return %x : vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @select_broadcast_multidimensional(
+// CHECK-SAME: %[[P:[^:]+]]: i1, %[[A:[^:]+]]: vector<2x4xi32>, %[[B:[^:]+]]: vector<2x4xi32>
+// CHECK-NOT: vector.broadcast
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : vector<2x4xi32>
+// CHECK-NOT: arith.select
+// CHECK: return %[[X]]
+func.func @select_broadcast_multidimensional(%p: i1, %a: vector<2x4xi32>, %b: vector<2x4xi32>) -> vector<2x4xi32> {
+ %m = vector.broadcast %p : i1 to vector<2x4xi1>
+ %x = arith.select %m, %a, %b : vector<2x4xi1>, vector<2x4xi32>
+ return %x : vector<2x4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @select_broadcast_scalable(
+// CHECK-SAME: %[[P:[^:]+]]: i1, %[[A:[^:]+]]: vector<[4]xi32>, %[[B:[^:]+]]: vector<[4]xi32>
+// CHECK-NOT: vector.broadcast
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : vector<[4]xi32>
+// CHECK-NOT: arith.select
+// CHECK: return %[[X]]
+func.func @select_broadcast_scalable(%p: i1, %a: vector<[4]xi32>, %b: vector<[4]xi32>) -> vector<[4]xi32> {
+ %m = vector.broadcast %p : i1 to vector<[4]xi1>
+ %x = arith.select %m, %a, %b : vector<[4]xi1>, vector<[4]xi32>
+ return %x : vector<[4]xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @select_broadcast_rank_zero(
+// CHECK-SAME: %[[P:[^:]+]]: i1, %[[A:[^:]+]]: vector<i32>, %[[B:[^:]+]]: vector<i32>
+// CHECK-NOT: vector.broadcast
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : vector<i32>
+// CHECK-NOT: arith.select
+// CHECK: return %[[X]]
+func.func @select_broadcast_rank_zero(%p: i1, %a: vector<i32>, %b: vector<i32>) -> vector<i32> {
+ %m = vector.broadcast %p : i1 to vector<i1>
+ %x = arith.select %m, %a, %b : vector<i1>, vector<i32>
+ return %x : vector<i32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @select_broadcast_shared_mask(
+// CHECK-SAME: %[[P:[^:]+]]: i1, %[[A:[^:]+]]: vector<4xi32>, %[[B:[^:]+]]: vector<4xi32>
+// CHECK: %[[M:.*]] = vector.broadcast %[[P]]
+// CHECK: %[[X:.*]] = arith.select %[[P]], %[[A]], %[[B]] : vector<4xi32>
+// CHECK-NOT: arith.select
+// CHECK: return %[[X]], %[[M]]
+func.func @select_broadcast_shared_mask(%p: i1, %a: vector<4xi32>, %b: vector<4xi32>) -> (vector<4xi32>, vector<4xi1>) {
+ %m = vector.broadcast %p : i1 to vector<4xi1>
+ %x = arith.select %m, %a, %b : vector<4xi1>, vector<4xi32>
+ return %x, %m : vector<4xi32>, vector<4xi1>
+}
+
+// -----
+
+// A vector source need not be uniform; do not replace its broadcast mask.
+// CHECK-LABEL: func.func @select_nonuniform_broadcast(
+// CHECK: %[[M:.*]] = vector.broadcast %{{.*}} : vector<4xi1> to vector<2x4xi1>
+// CHECK: arith.select %[[M]], %{{.*}}, %{{.*}} : vector<2x4xi1>, vector<2x4xi32>
+func.func @select_nonuniform_broadcast(%p: vector<4xi1>, %a: vector<2x4xi32>, %b: vector<2x4xi32>) -> vector<2x4xi32> {
+ %m = vector.broadcast %p : vector<4xi1> to vector<2x4xi1>
+ %x = arith.select %m, %a, %b : vector<2x4xi1>, vector<2x4xi32>
+ return %x : vector<2x4xi32>
+}
+
+// -----
+
+// Preserve the existing refinement of a select with a fully poisoned arm.
+// CHECK-LABEL: func.func @select_uniform_unselected_poison_vector(
+// CHECK-SAME: %[[P:[^:]+]]: i1
+// CHECK-SAME: %[[A:[^:]+]]: vector<4xi32>
+// CHECK-NOT: vector.broadcast
+// CHECK-NOT: arith.select
+// CHECK: return %[[A]] : vector<4xi32>
+func.func @select_uniform_unselected_poison_vector(%p: i1, %a: vector<4xi32>) -> vector<4xi32> {
+ %bad = ub.poison : vector<4xi32>
+ %m = vector.broadcast %p : i1 to vector<4xi1>
+ %x = arith.select %m, %a, %bad : vector<4xi1>, vector<4xi32>
+ return %x : vector<4xi32>
+}
More information about the Mlir-commits
mailing list