[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:58 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-tensor
Author: meinie (meinie0826)
<details>
<summary>Changes</summary>
## 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.
## 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.
---
Full diff: https://github.com/llvm/llvm-project/pull/235593.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/Vector/IR/Vector.td (+1)
- (modified) mlir/lib/Dialect/Tensor/IR/TensorOps.cpp (+18-1)
- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+27)
- (modified) mlir/test/Dialect/Tensor/canonicalize.mlir (+97)
- (modified) mlir/test/Dialect/Vector/canonicalize.mlir (+98)
``````````diff
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..0ac890b74a5ddf 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -8692,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
//===----------------------------------------------------------------------===//
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>
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/235593
More information about the Mlir-commits
mailing list