[Mlir-commits] [mlir] [mlir] [arith] Combine op for select (PR #234818)
Chuanqi Xu
llvmlistbot at llvm.org
Sun Oct 25 23:17:49 PDT 2026
https://github.com/ChuanqiXu9 updated https://github.com/llvm/llvm-project/pull/234818
>From 69bde683db5c3d86ad34b4a2c1014038224ca60a Mon Sep 17 00:00:00 2001
From: Chuanqi Xu <chuanqi.xcq at alibaba-inc.com>
Date: Fri, 23 Oct 2026 10:14:39 +0800
Subject: [PATCH] [mlir] [arith] Combine op for select
Fold select(c, F(a0, a1), F(b0, b1)) into F(select(c, a0, b0),
select(c, a1, b1))
Observed significant performance in our internal kernels.
AI assisted for testing generation.
---
mlir/lib/Dialect/Arith/IR/ArithOps.cpp | 153 ++++++++++++++++++
mlir/lib/Dialect/Arith/IR/CMakeLists.txt | 1 +
.../Arith/select-common-elementwise.mlir | 107 ++++++++++++
mlir/test/lib/Dialect/Test/TestOps.td | 8 +
4 files changed, 269 insertions(+)
create mode 100644 mlir/test/Dialect/Arith/select-common-elementwise.mlir
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 19e9caca6567a6..41519425b84fdf 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -22,11 +22,13 @@
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Support/LogicalResult.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/APInt.h"
#include "llvm/ADT/APSInt.h"
+#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/FloatingPointMode.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
@@ -3128,6 +3130,156 @@ struct SelectToExtUI : public OpRewritePattern<arith::SelectOp> {
}
};
+/// select(M, F(a0, a1, ..., a_n), F(b0, b1, ..., b_n))
+/// -> F(select(M, a0, b0), select(M, a1, b1), ...) for pure and
+/// elementwise F.
+///
+/// Correctness provement:
+///
+/// Let's discuss the cases ignoring poison first.
+///
+/// As F is elementwise, we have
+///
+/// F(a1, a2, ..., a_n)[p] = F(a1[p], a2[p], ..., a_n[p])
+///
+/// then for
+///
+/// L = select(M, F(a0, a1, ..., a_n), F(b0, b1, ..., b_n))
+/// R = F(select(M, a0, b0), select(M, a1, b1), ....)
+///
+/// for any position p, we have:
+///
+/// if M[p] == true, then L[p] == F(a)[p] == F(a[p]) == F(select(true, a[p], b[p])) == R[p]
+/// if M[p] == false, then L[p] == F(b)[p] == F(b[p]) == F(select(false, a[p], b[p])) == R[p]
+///
+/// And F is pure, so that we can move and reduce F.
+/// So select(M, F(a0, a1, ..., a_n), F(b0, b1, ..., b_n)) is equal to
+/// F(select(M, a0, b0), select(M, a1, b1), ...) ignoring poison.
+///
+/// Let's consider the poison value now. ArithBase.td says:
+///
+/// Unless otherwise stated, the operations within this dialect
+/// propagate poison values, i.e., if any of its inputs are poison, then the
+/// output is poison.
+///
+/// So most F, F(poison) = poison
+///
+/// Then for any position p, if a[p] is poison
+///
+/// if M[p] == true, L[p] == f(a[p]) == poison == F(select(true, a[p], b[p])) == R[p]
+/// if M[p] == false, L[p] == f(b[p]) == F(select(false, a[p], b[p])) == R[p]
+/// if M[p] == P, L[p] == poison == F(select(P, a, b)) == R[p]
+///
+/// Then the transformation doesn't change poison propagation.
+///
+/// For F which doesn't meet the assumption, this transformation may change the
+/// poison propagation, but the language allows us to replace poison with a undef
+/// or any value with the same type. So this is still valid transformation.
+struct SelectCommonElementwiseOp : public OpRewritePattern<arith::SelectOp> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(arith::SelectOp op,
+ PatternRewriter &rewriter) const override {
+ Operation *lhs = op.getTrueValue().getDefiningOp();
+ Operation *rhs = op.getFalseValue().getDefiningOp();
+ if (!lhs || !rhs || lhs == rhs)
+ return failure();
+
+ for (Operation *producer : {lhs, rhs}) {
+ // Conservative design. The motivation is we don't want to transfer:
+ //
+ // x = F(a)
+ // y = F(b)
+ // for i:
+ // r = select(c[i], x, y)
+ //
+ // into:
+ //
+ // for i:
+ // s = select(c[i], a, b)
+ // r = F(s)
+ //
+ // TODO: maybe we can handle other cases by introducing more CFG/Loop
+ // related analysis.
+ if (producer->getBlock() != op->getBlock())
+ return failure();
+
+ // Skipping hard cases.
+ //
+ // For regions, a counter example is,
+ //
+ // lhs = F(a) { v -> v + k0 }
+ // rhs = F(b) { v -> v + k1 }
+ // r = select(c, lhs, rhs)
+ //
+ // We can't rewrite this. The following check `isEquivalentTo` won't reject
+ // this as we passed ignoreValueEquivalence.
+ if (producer->getNumResults() != 1 || !producer->hasOneUse() ||
+ producer->getNumRegions() != 0 || producer->getNumSuccessors() != 0 ||
+ !producer->hasTrait<OpTrait::Elementwise>() || !isPure(producer))
+ return failure();
+ }
+
+ // Check if the lhs and the rhs are the same kind of op, have the
+ // same structure of inputs, types, and attributes.
+ if (!OperationEquivalence::isEquivalentTo(
+ lhs, rhs, OperationEquivalence::ignoreValueEquivalence, nullptr,
+ OperationEquivalence::IgnoreLocations))
+ return failure();
+
+ DenseMap<std::pair<Value, Value>, Value> selected;
+ SmallVector<Operation *> createdSelects;
+ SmallVector<Value> operands;
+ Type conditionType = op.getCondition().getType();
+ for (auto [a, b] : llvm::zip(lhs->getOperands(), rhs->getOperands())) {
+ // For select(c, F(a, k), F(b, k)), we should transform it into
+ // F(select(c, a, b), k) instead of F(select(c, a, b), select(c, k, k))
+ if (a == b) {
+ operands.push_back(a);
+ continue;
+ }
+ if (!conditionType.isSignlessInteger(1) &&
+ (!isa<VectorType, TensorType>(a.getType()) ||
+ cast<ShapedType>(a.getType()).clone(rewriter.getI1Type()) !=
+ conditionType)) {
+
+ for (Operation *created : llvm::reverse(createdSelects))
+ rewriter.eraseOp(created);
+ return failure();
+ }
+
+ // If we always create new pair, we may transfer
+ //
+ // select(c, F(a, a, a), F(b, b, b))
+ //
+ // into
+ //
+ // s0 = select(c, a, b)
+ // s1 = select(c, a, b)
+ // s2 = select(c, a, b)
+ // F(s0, s1, s2)
+ //
+ // instead of
+ //
+ // s = select(c, a, b)
+ // F = (s, s, s)
+ auto [it, inserted] = selected.try_emplace(std::make_pair(a, b));
+ if (inserted) {
+ auto select = arith::SelectOp::create(rewriter, op.getLoc(),
+ op.getCondition(), a, b);
+ it->second = select;
+ createdSelects.push_back(select);
+ }
+ operands.push_back(it->second);
+ }
+
+ Operation *common = rewriter.clone(*lhs);
+ common->setOperands(operands);
+ rewriter.replaceOp(op, common->getResults());
+ return success();
+ }
+};
+
void arith::SelectOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {
results.add<RedundantSelectFalse, RedundantSelectTrue, SelectNotCond,
@@ -3138,6 +3290,7 @@ void arith::SelectOp::getCanonicalizationPatterns(RewritePatternSet &results,
SelectCmpIUgtToMinUI, SelectCmpIUleToMaxUI, SelectCmpIUleToMinUI,
SelectCmpIUltToMaxUI, SelectCmpIUltToMinUI, SelectToExtUI>(
context);
+ results.add<SelectCommonElementwiseOp>(context);
}
OpFoldResult arith::SelectOp::fold(FoldAdaptor adaptor) {
diff --git a/mlir/lib/Dialect/Arith/IR/CMakeLists.txt b/mlir/lib/Dialect/Arith/IR/CMakeLists.txt
index 3423e11a7d0f0d..a6508946030a20 100644
--- a/mlir/lib/Dialect/Arith/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Arith/IR/CMakeLists.txt
@@ -31,6 +31,7 @@ add_mlir_dialect_library(MLIRArithDialect
MLIRInferIntRangeInterface
MLIRInferTypeOpInterface
MLIRIR
+ MLIRSideEffectInterfaces
MLIRUBDialect
)
diff --git a/mlir/test/Dialect/Arith/select-common-elementwise.mlir b/mlir/test/Dialect/Arith/select-common-elementwise.mlir
new file mode 100644
index 00000000000000..0236b9a6ed43c7
--- /dev/null
+++ b/mlir/test/Dialect/Arith/select-common-elementwise.mlir
@@ -0,0 +1,107 @@
+// RUN: mlir-opt %s --canonicalize="test-convergence" | FileCheck %s
+
+// Test select(M, F(a0, a1, ..., a_n), F(b0, b1, ..., b_n)) -> F(select(M, a0, b0), select(M, a1, b1), ...).
+
+// CHECK-LABEL: func.func @basic(
+// CHECK-NEXT: %[[A:.*]] = arith.select %arg0, %arg1, %arg3 : tensor<2x3xi1>, tensor<2x3xi32>
+// CHECK-NEXT: %[[B:.*]] = arith.select %arg0, %arg2, %arg4 : tensor<2x3xi1>, tensor<2x3xi32>
+// CHECK-NEXT: %[[M:.*]] = arith.subi %[[A]], %[[B]] : tensor<2x3xi32>
+// CHECK-NEXT: %[[R:.*]] = arith.extsi %[[M]] : tensor<2x3xi32> to tensor<2x3xi64>
+// CHECK-NEXT: return %[[R]] : tensor<2x3xi64>
+func.func @basic(%mask: tensor<2x3xi1>, %a0: tensor<2x3xi32>, %b0: tensor<2x3xi32>, %a1: tensor<2x3xi32>, %b1: tensor<2x3xi32>) -> tensor<2x3xi64> {
+ %x = arith.subi %a0, %b0 : tensor<2x3xi32>
+ %y = arith.subi %a1, %b1 : tensor<2x3xi32>
+ %ex = arith.extsi %x : tensor<2x3xi32> to tensor<2x3xi64>
+ %ey = arith.extsi %y : tensor<2x3xi32> to tensor<2x3xi64>
+ %r = arith.select %mask, %ex, %ey : tensor<2x3xi1>, tensor<2x3xi64>
+ return %r : tensor<2x3xi64>
+}
+
+// If the operands of the select op are not the same, do not transform
+// CHECK-LABEL: func.func @diff_op(
+// CHECK-NEXT: %[[X:.*]] = arith.muli %arg1, %arg3 overflow<nsw> : i32
+// CHECK-NEXT: %[[Y:.*]] = arith.muli %arg2, %arg3 overflow<nuw> : i32
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : i32
+// CHECK-NEXT: return %[[R]] : i32
+func.func @diff_op(%mask: i1, %a: i32, %b: i32, %k: i32) -> i32 {
+ %x = arith.muli %a, %k overflow<nsw> : i32
+ %y = arith.muli %b, %k overflow<nuw> : i32
+ %r = arith.select %mask, %x, %y : i32
+ return %r : i32
+}
+
+// CHECK-LABEL: func.func @non_elementwise(
+// CHECK-NEXT: %[[X:.*]] = vector.shuffle %arg1, %arg3 [1, 4, 3, 6] : vector<4xf32>, vector<4xf32>
+// CHECK-NEXT: %[[Y:.*]] = vector.shuffle %arg2, %arg3 [1, 4, 3, 6] : vector<4xf32>, vector<4xf32>
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : vector<4xi1>, vector<4xf32>
+// CHECK-NEXT: return %[[R]] : vector<4xf32>
+func.func @non_elementwise(%mask: vector<4xi1>, %a: vector<4xf32>, %b: vector<4xf32>, %k: vector<4xf32>) -> vector<4xf32> {
+ %x = vector.shuffle %a, %k [1, 4, 3, 6] : vector<4xf32>, vector<4xf32>
+ %y = vector.shuffle %b, %k [1, 4, 3, 6] : vector<4xf32>, vector<4xf32>
+ %r = arith.select %mask, %x, %y : vector<4xi1>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// If the regions capture different values, do not transform.
+// CHECK-LABEL: func.func @region(
+// CHECK-NEXT: %[[X:.*]] = "test.pure_elementwise_region"(%arg1) ({
+// CHECK-NEXT: ^bb0(%[[A:.*]]: f32):
+// CHECK-NEXT: %[[SUM0:.*]] = arith.addf %[[A]], %arg3 : f32
+// CHECK-NEXT: "test.return"(%[[SUM0]]) : (f32) -> ()
+// CHECK-NEXT: }) : (f32) -> f32
+// CHECK-NEXT: %[[Y:.*]] = "test.pure_elementwise_region"(%arg2) ({
+// CHECK-NEXT: ^bb0(%[[B:.*]]: f32):
+// CHECK-NEXT: %[[SUM1:.*]] = arith.addf %[[B]], %arg4 : f32
+// CHECK-NEXT: "test.return"(%[[SUM1]]) : (f32) -> ()
+// CHECK-NEXT: }) : (f32) -> f32
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : f32
+// CHECK-NEXT: return %[[R]] : f32
+func.func @region(%mask: i1, %a: f32, %b: f32, %k0: f32, %k1: f32) -> f32 {
+ %x = "test.pure_elementwise_region"(%a) ({
+ ^bb0(%v: f32):
+ %sum = arith.addf %v, %k0 : f32
+ "test.return"(%sum) : (f32) -> ()
+ }) : (f32) -> f32
+ %y = "test.pure_elementwise_region"(%b) ({
+ ^bb0(%v: f32):
+ %sum = arith.addf %v, %k1 : f32
+ "test.return"(%sum) : (f32) -> ()
+ }) : (f32) -> f32
+ %r = arith.select %mask, %x, %y : f32
+ return %r : f32
+}
+
+// If F is outside the loop, do not move it into the loop.
+// CHECK-LABEL: func.func @outside_loop(
+// CHECK-NEXT: %[[X:.*]] = math.sqrt %arg4 : f32
+// CHECK-NEXT: %[[Y:.*]] = math.sqrt %arg5 : f32
+// CHECK-NEXT: %[[R:.*]] = scf.for %[[I:.*]] = %arg0 to %arg1 step %arg2 iter_args(%[[ACC:.*]] = %arg6) -> (f32) {
+// CHECK-NEXT: %[[MASK:.*]] = arith.cmpi slt, %[[I]], %arg3 : index
+// CHECK-NEXT: %[[S:.*]] = arith.select %[[MASK]], %[[X]], %[[Y]] : f32
+// CHECK-NEXT: %[[SUM:.*]] = arith.addf %[[ACC]], %[[S]] : f32
+// CHECK-NEXT: scf.yield %[[SUM]] : f32
+// CHECK-NEXT: }
+// CHECK-NEXT: return %[[R]] : f32
+func.func @outside_loop(%lb: index, %ub: index, %step: index, %split: index, %a: f32, %b: f32, %init: f32) -> f32 {
+ %x = math.sqrt %a : f32
+ %y = math.sqrt %b : f32
+ %r = scf.for %i = %lb to %ub step %step iter_args(%acc = %init) -> f32 {
+ %mask = arith.cmpi slt, %i, %split : index
+ %s = arith.select %mask, %x, %y : f32
+ %sum = arith.addf %acc, %s : f32
+ scf.yield %sum : f32
+ }
+ return %r : f32
+}
+
+// Repeated operands share one select; a common operand needs no select.
+// CHECK-LABEL: func.func @repeated_operands(
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : f32
+// CHECK-NEXT: %[[R:.*]] = math.fma %[[S]], %[[S]], %arg3 : f32
+// CHECK-NEXT: return %[[R]] : f32
+func.func @repeated_operands(%mask: i1, %a: f32, %b: f32, %k: f32) -> f32 {
+ %x = math.fma %a, %a, %k : f32
+ %y = math.fma %b, %b, %k : f32
+ %r = arith.select %mask, %x, %y : f32
+ return %r : f32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 0a64885d8e2921..183f6afcdc5744 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -854,6 +854,14 @@ def ElementwiseMappableOp : TEST_Op<"elementwise_mappable",
let results = (outs Variadic<AnyType>);
}
+def PureElementwiseRegionOp : TEST_Op<"pure_elementwise_region",
+ [Elementwise, RecursivelySpeculatable, RecursiveMemoryEffects, SingleBlock]> {
+ let summary = "elementwise region op that can capture external values";
+ let arguments = (ins Variadic<AnyType>);
+ let results = (outs AnyType);
+ let regions = (region SizedRegion<1>:$body);
+}
+
def ArgAndResHaveFixedElementTypesOp :
TEST_Op<"arg_and_res_have_fixed_element_types",
[PredOpTrait<"fixed type combination",
More information about the Mlir-commits
mailing list