[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:34:25 PDT 2026
https://github.com/meinie0826 updated https://github.com/llvm/llvm-project/pull/235593
>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 1/2] [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>
+}
>From 804b0c2059d490a6b4c31ccdb91ced22478338e7 Mon Sep 17 00:00:00 2001
From: meinie0826 <929449726 at qq.com>
Date: Mon, 26 Oct 2026 13:34:07 +0800
Subject: [PATCH 2/2] [mlir][vector] Group dialect canonicalization near the
end of VectorOps.cpp
Keep select-condition patterns and their registration together without changing the rewrite.
Assisted-by: OpenAI Codex
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 50 +++++++++++++-----------
1 file changed, 27 insertions(+), 23 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index c42c5e62a218f7..0ac890b74a5ddf 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -574,29 +574,6 @@ 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,
@@ -8715,6 +8692,33 @@ DeinterleaveOp::getIterationShapeToUnroll() {
return llvm::to_vector<4>(getResultVectorType().getShape());
}
+//===----------------------------------------------------------------------===//
+// VectorDialect canonicalization
+//===----------------------------------------------------------------------===//
+
+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());
+}
+
//===----------------------------------------------------------------------===//
// TableGen'd op method definitions
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list