[Mlir-commits] [mlir] [mlir] [arith] Combine op for select (PR #234818)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 22 19:21:00 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-arith
Author: Chuanqi Xu (ChuanqiXu9)
<details>
<summary>Changes</summary>
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.
---
Patch is 35.00 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/234818.diff
5 Files Affected:
- (modified) mlir/lib/Dialect/Arith/IR/ArithOps.cpp (+102)
- (modified) mlir/lib/Dialect/Arith/IR/CMakeLists.txt (+1)
- (added) mlir/test/Dialect/Arith/select-common-elementwise-test-ops.mlir (+219)
- (added) mlir/test/Dialect/Arith/select-common-elementwise.mlir (+362)
- (modified) mlir/test/lib/Dialect/Test/TestOps.td (+21)
``````````diff
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 19e9caca6567a..2c8091597bd9a 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,105 @@ struct SelectToExtUI : public OpRewritePattern<arith::SelectOp> {
}
};
+/// select(c, F(a0, a1), F(b0, b1))
+/// -> F(select(c, a0, b0), select(c, a1, b1)).
+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();
+
+ // The equivalence check below also ignores SSA values inside regions,
+ // including captures that would remain unchanged when cloning the
+ // left-hand producer.
+ if (producer->getNumResults() != 1 || !producer->hasOneUse() ||
+ producer->getNumRegions() != 0 || producer->getNumSuccessors() != 0 ||
+ !producer->hasTrait<OpTrait::Elementwise>() || !isPure(producer))
+ return failure();
+ }
+
+ // If the lhs and the rhs are semantically equal, we can merge them.
+ 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 +3239,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 3423e11a7d0f0..a6508946030a2 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-test-ops.mlir b/mlir/test/Dialect/Arith/select-common-elementwise-test-ops.mlir
new file mode 100644
index 0000000000000..62948c66971ba
--- /dev/null
+++ b/mlir/test/Dialect/Arith/select-common-elementwise-test-ops.mlir
@@ -0,0 +1,219 @@
+// RUN: mlir-opt %s --canonicalize="test-convergence" | FileCheck %s
+// RUN: mlir-opt %s --canonicalize="test-convergence" --cse | FileCheck %s
+
+// A shared scalar parameter is legal with an elementwise tensor mask.
+// CHECK-LABEL: @shared_scalar
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : tensor<4xi1>, tensor<4xf32>
+// CHECK-NEXT: %[[R:.*]] = "test.pure_elementwise"(%[[S]], %arg3) {callee = @foo} : (tensor<4xf32>, i32) -> tensor<4xf32>
+// CHECK-NEXT: return %[[R]]
+func.func @shared_scalar(%c: tensor<4xi1>, %a: tensor<4xf32>, %b: tensor<4xf32>, %k: i32) -> tensor<4xf32> {
+ %x = "test.pure_elementwise"(%a, %k) {callee = @foo} : (tensor<4xf32>, i32) -> tensor<4xf32>
+ %y = "test.pure_elementwise"(%b, %k) {callee = @foo} : (tensor<4xf32>, i32) -> tensor<4xf32>
+ %r = arith.select %c, %x, %y : tensor<4xi1>, tensor<4xf32>
+ return %r : tensor<4xf32>
+}
+
+// A tensor mask cannot select differing scalar parameters. Erase the select
+// created for the first operand pair when the second pair fails validation.
+// CHECK-LABEL: @different_scalars_tensor
+// CHECK-NEXT: %[[X:.*]] = "test.pure_elementwise"(%arg1, %arg3)
+// CHECK-NEXT: %[[Y:.*]] = "test.pure_elementwise"(%arg2, %arg4)
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : tensor<4xi1>, tensor<4xf32>
+// CHECK-NEXT: return %[[R]]
+func.func @different_scalars_tensor(%c: tensor<4xi1>, %a: tensor<4xf32>, %b: tensor<4xf32>, %k: i32, %l: i32) -> tensor<4xf32> {
+ %x = "test.pure_elementwise"(%a, %k) : (tensor<4xf32>, i32) -> tensor<4xf32>
+ %y = "test.pure_elementwise"(%b, %l) : (tensor<4xf32>, i32) -> tensor<4xf32>
+ %r = arith.select %c, %x, %y : tensor<4xi1>, tensor<4xf32>
+ return %r : tensor<4xf32>
+}
+
+// CHECK-LABEL: @different_scalars_vector
+// CHECK-NEXT: %[[X:.*]] = "test.pure_elementwise"(%arg1, %arg3)
+// CHECK-NEXT: %[[Y:.*]] = "test.pure_elementwise"(%arg2, %arg4)
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : vector<4xi1>, vector<4xf32>
+// CHECK-NEXT: return %[[R]]
+func.func @different_scalars_vector(%c: vector<4xi1>, %a: vector<4xf32>, %b: vector<4xf32>, %k: i32, %l: i32) -> vector<4xf32> {
+ %x = "test.pure_elementwise"(%a, %k) : (vector<4xf32>, i32) -> vector<4xf32>
+ %y = "test.pure_elementwise"(%b, %l) : (vector<4xf32>, i32) -> vector<4xf32>
+ %r = arith.select %c, %x, %y : vector<4xi1>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// A scalar condition can select both the tensor and the scalar parameters.
+// CHECK-LABEL: @scalar_condition_mixed_operands
+// CHECK-NEXT: %[[S0:.*]] = arith.select %arg0, %arg1, %arg2 : tensor<4xf32>
+// CHECK-NEXT: %[[S1:.*]] = arith.select %arg0, %arg3, %arg4 : i32
+// CHECK-NEXT: %[[R:.*]] = "test.pure_elementwise"(%[[S0]], %[[S1]]) : (tensor<4xf32>, i32) -> tensor<4xf32>
+// CHECK-NEXT: return %[[R]]
+func.func @scalar_condition_mixed_operands(%c: i1, %a: tensor<4xf32>, %b: tensor<4xf32>, %k: i32, %l: i32) -> tensor<4xf32> {
+ %x = "test.pure_elementwise"(%a, %k) : (tensor<4xf32>, i32) -> tensor<4xf32>
+ %y = "test.pure_elementwise"(%b, %l) : (tensor<4xf32>, i32) -> tensor<4xf32>
+ %r = arith.select %c, %x, %y : tensor<4xf32>
+ return %r : tensor<4xf32>
+}
+
+// Elementwise permits compatible dynamic shapes; select requires exact types.
+// CHECK-LABEL: @different_mask_shape
+// CHECK-NEXT: %[[X:.*]] = "test.pure_elementwise"(%arg1)
+// CHECK-NEXT: %[[Y:.*]] = "test.pure_elementwise"(%arg2)
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : tensor<4xi1>, tensor<4xf32>
+// CHECK-NEXT: return %[[R]]
+func.func @different_mask_shape(%c: tensor<4xi1>, %a: tensor<?xf32>, %b: tensor<?xf32>) -> tensor<4xf32> {
+ %x = "test.pure_elementwise"(%a) : (tensor<?xf32>) -> tensor<4xf32>
+ %y = "test.pure_elementwise"(%b) : (tensor<?xf32>) -> tensor<4xf32>
+ %r = arith.select %c, %x, %y : tensor<4xi1>, tensor<4xf32>
+ return %r : tensor<4xf32>
+}
+
+// Matching dimensions alone are insufficient when encodings differ.
+// CHECK-LABEL: @different_mask_encoding
+// CHECK-NEXT: %[[X:.*]] = "test.pure_elementwise"(%arg1)
+// CHECK-NEXT: %[[Y:.*]] = "test.pure_elementwise"(%arg2)
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : tensor<4xi1, "out">, tensor<4xf32, "out">
+// CHECK-NEXT: return %[[R]]
+func.func @different_mask_encoding(%c: tensor<4xi1, "out">, %a: tensor<4xf32, "in">, %b: tensor<4xf32, "in">) -> tensor<4xf32, "out"> {
+ %x = "test.pure_elementwise"(%a) : (tensor<4xf32, "in">) -> tensor<4xf32, "out">
+ %y = "test.pure_elementwise"(%b) : (tensor<4xf32, "in">) -> tensor<4xf32, "out">
+ %r = arith.select %c, %x, %y : tensor<4xi1, "out">, tensor<4xf32, "out">
+ return %r : tensor<4xf32, "out">
+}
+
+// A shared input does not need to match the mask's encoding.
+// CHECK-LABEL: @shared_different_encoding
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : tensor<4xi1, "out">, tensor<4xf32, "out">
+// CHECK-NEXT: %[[R:.*]] = "test.pure_elementwise"(%[[S]], %arg3) : (tensor<4xf32, "out">, tensor<4xf32, "in">) -> tensor<4xf32, "out">
+// CHECK-NEXT: return %[[R]]
+func.func @shared_different_encoding(%c: tensor<4xi1, "out">, %a: tensor<4xf32, "out">, %b: tensor<4xf32, "out">, %k: tensor<4xf32, "in">) -> tensor<4xf32, "out"> {
+ %x = "test.pure_elementwise"(%a, %k) : (tensor<4xf32, "out">, tensor<4xf32, "in">) -> tensor<4xf32, "out">
+ %y = "test.pure_elementwise"(%b, %k) : (tensor<4xf32, "out">, tensor<4xf32, "in">) -> tensor<4xf32, "out">
+ %r = arith.select %c, %x, %y : tensor<4xi1, "out">, tensor<4xf32, "out">
+ return %r : tensor<4xf32, "out">
+}
+
+// CHECK-LABEL: @different_symbols
+// CHECK-NEXT: %[[X:.*]] = "test.pure_elementwise"(%arg1) {callee = @foo}
+// CHECK-NEXT: %[[Y:.*]] = "test.pure_elementwise"(%arg2) {callee = @bar}
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : f32
+// CHECK-NEXT: return %[[R]]
+func.func @different_symbols(%c: i1, %a: f32, %b: f32) -> f32 {
+ %x = "test.pure_elementwise"(%a) {callee = @foo} : (f32) -> f32
+ %y = "test.pure_elementwise"(%b) {callee = @bar} : (f32) -> f32
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// Speculatability alone is insufficient, even for an elementwise operation.
+// CHECK-LABEL: @elementwise_read
+// CHECK-NEXT: %[[X:.*]] = "test.read_elementwise"(%arg1) {callee = @foo}
+// CHECK-NEXT: %[[Y:.*]] = "test.read_elementwise"(%arg2) {callee = @foo}
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : f32
+// CHECK-NEXT: return %[[R]]
+func.func @elementwise_read(%c: i1, %a: f32, %b: f32) -> f32 {
+ %x = "test.read_elementwise"(%a) {callee = @foo} : (f32) -> f32
+ %y = "test.read_elementwise"(%b) {callee = @foo} : (f32) -> f32
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// Same operation name and result type do not imply matching operand types.
+// CHECK-LABEL: @different_operand_types
+// CHECK-NEXT: %[[X:.*]] = "test.pure_elementwise"(%arg1) : (f16) -> f32
+// CHECK-NEXT: %[[Y:.*]] = "test.pure_elementwise"(%arg2) : (bf16) -> f32
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : f32
+// CHECK-NEXT: return %[[R]]
+func.func @different_operand_types(%c: i1, %a: f16, %b: bf16) -> f32 {
+ %x = "test.pure_elementwise"(%a) : (f16) -> f32
+ %y = "test.pure_elementwise"(%b) : (bf16) -> f32
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// Multiple results remain outside the pattern even if there is only one use.
+// CHECK-LABEL: @multiple_results
+// CHECK-NEXT: %[[X:.*]]:2 = "test.pure_elementwise"(%arg1) : (f32) -> (f32, f32)
+// CHECK-NEXT: %[[Y:.*]]:2 = "test.pure_elementwise"(%arg2) : (f32) -> (f32, f32)
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]]#0, %[[Y]]#0 : f32
+// CHECK-NEXT: return %[[R]]
+func.func @multiple_results(%c: i1, %a: f32, %b: f32) -> f32 {
+ %x:2 = "test.pure_elementwise"(%a) : (f32) -> (f32, f32)
+ %y:2 = "test.pure_elementwise"(%b) : (f32) -> (f32, f32)
+ %r = arith.select %c, %x#0, %y#0 : f32
+ return %r : f32
+}
+
+// Corresponding operand encodings must match even with a scalar condition.
+// CHECK-LABEL: @different_operand_encodings
+// CHECK-NEXT: %[[X:.*]] = "test.pure_elementwise"(%arg1)
+// CHECK-NEXT: %[[Y:.*]] = "test.pure_elementwise"(%arg2)
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : tensor<4xf32, "out">
+// CHECK-NEXT: return %[[R]]
+func.func @different_operand_encodings(%c: i1, %a: tensor<4xf32, "a">, %b: tensor<4xf32, "b">) -> tensor<4xf32, "out"> {
+ %x = "test.pure_elementwise"(%a) : (tensor<4xf32, "a">) -> tensor<4xf32, "out">
+ %y = "test.pure_elementwise"(%b) : (tensor<4xf32, "b">) -> tensor<4xf32, "out">
+ %r = arith.select %c, %x, %y : tensor<4xf32, "out">
+ return %r : tensor<4xf32, "out">
+}
+
+// Comparing regions while ignoring SSA values would consider these bodies
+// equivalent. Replacing only the top-level operands cannot select k0 versus k1:
+// cloning the true producer would retain k0 even when c is false.
+// CHECK-LABEL: @different_region_captures
+// 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]]
+func.func @different_region_captures(%c: 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 %c, %x, %y : f32
+ return %r : f32
+}
+
+// Failure after creating two selects must erase each exactly once, even though
+// the first is reused. Keep existing operands and their defining operations.
+// CHECK-LABEL: @cleanup_partial_operands
+// CHECK-NEXT: %[[A:.*]] = math.sqrt %arg1 : tensor<4xf32>
+// CHECK-NEXT: %[[B:.*]] = math.sqrt %arg2 : tensor<4xf32>
+// CHECK-NEXT: %[[SHARED:.*]] = math.sqrt %arg5 : tensor<4xf32>
+// CHECK-NEXT: %[[X:.*]] = "test.pure_elementwise"(%[[SHARED]], %[[A]], %[[A]], %arg3, %arg6)
+// CHECK-NEXT: %[[Y:.*]] = "test.pure_elementwise"(%[[SHARED]], %[[B]], %[[B]], %arg4, %arg7)
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : tensor<4xi1>, tensor<4xf32>
+// CHECK-NEXT: return %[[R]]
+func.func @cleanup_partial_operands(%c: tensor<4xi1>, %a_in: tensor<4xf32>, %b_in: tensor<4xf32>, %d: tensor<4xf32>, %e: tensor<4xf32>, %shared_in: tensor<4xf32>, %k: i32, %l: i32) -> tensor<4xf32> {
+ %a = math.sqrt %a_in : tensor<4xf32>
+ %b = math.sqrt %b_in : tensor<4xf32>
+ %shared = math.sqrt %shared_in : tensor<4xf32>
+ %x = "test.pure_elementwise"(%shared, %a, %a, %d, %k) : (tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, i32) -> tensor<4xf32>
+ %y = "test.pure_elementwise"(%shared, %b, %b, %e, %l) : (tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, i32) -> tensor<4xf32>
+ %r = arith.select %c, %x, %y : tensor<4xi1>, tensor<4xf32>
+ return %r : tensor<4xf32>
+}
+
+// An encoding mismatch on a later pair must also clean up earlier selects.
+// CHECK-LABEL: @cleanup_on_encoding_mismatch
+// CHECK-NEXT: %[[X:.*]] = "test.pure_elementwise"(%arg1, %arg3)
+// CHECK-NEXT: %[[Y:.*]] = "test.pure_elementwise"(%arg2, %arg4)
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : tensor<4xi1, "out">, tensor<4xf32, "out">
+// CHECK-NEXT: return %[[R]]
+func.func @cleanup_on_encoding_mismatch(%c: tensor<4xi1, "out">, %a: tensor<4xf32, "out">, %b: tensor<4xf32, "out">, %k: tensor<4xf32, "in">, %l: tensor<4xf32, "in">) -> tensor<4xf32, "out"> {
+ %x = "test.pure_elementwise"(%a, %k) : (tensor<4xf32, "out">, tensor<4xf32, "in">) -> tensor<4xf32, "out">
+ %y = "test.pure_elementwise"(%b, %l) : (tensor<4xf32, "out">, tensor<4xf32, "in">) -> tensor<4xf32, "out">
+ %r = arith.select %c, %x, %y : tensor<4xi1, "out">, tensor<4xf32, "out">
+ return %r : tensor<4xf32, "out">
+}
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 0000000000000..12d3dbfd587c4
--- /dev/null
+++ b/mlir/test/Dialect/Arith/select-common-elementwise.mlir
@@ -0,0 +1,362 @@
+// RUN: mlir-opt %s --canonicalize="test-convergence" | FileCheck %s
+// RUN: mlir-opt %s --canonicalize="test-convergence" --cse | FileCheck %s
+
+// CHECK-LABEL: @scalar_unary
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : f32
+// CHECK-NEXT: %[[R:.*]] = math.sqrt %[[S]] : f32
+// CHECK-NEXT: return %[[R]] : f32
+func.func @scalar_unary(%c: i1, %a: f32, %b: f32) -> f32 {
+ %x = math.sqrt %a : f32 loc("true")
+ %y = math.sqrt %b : f32 loc("false")
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// CHECK-LABEL: @vector_binary
+// CHECK-NEXT: %[[S0:.*]] = arith.select %arg0, %arg1, %arg3 : vector<4xi1>, vector<4xf32>
+// CHECK-NEXT: %[[S1:.*]] = arith.select %arg0, %arg2, %arg4 : vector<4xi1>, vector<4xf32>
+// CHECK-NEXT: %[[R:.*]] = arith.mulf %[[S0]], %[[S1]] : vector<4xf32>
+// CHECK-NEXT: return %[[R]]
+func.func @vector_binary(%c: vector<4xi1>, %a: vector<4xf32>, %b: vector<4xf32>, %d: vector<4xf32>, %e: vector<4xf32>) -> vector<4xf32> {
+ %x = arith.mulf %a, %b : vector<4xf32>
+ %y = arith.mulf %d, %e : vector<4xf32>
+ %r = arith.select %c, %x, %y : vector<4xi1>, vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// CHECK-LABEL: @tensor_shared_operand
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : tensor<4xi1>, tensor<4xf32>
+// CHECK-NEXT: %[[R:.*]] = arith.mulf %[[S]], %arg3 : tensor<4xf32>
+// CHECK-NEXT: return %[[R]]
+func.func @tensor_shared_operand(%c: tensor<4xi1>, %a: tensor<4xf32>, %b: tensor<4xf32>, %k: tensor<4xf32>) -> tensor<4xf32> {
+ %x = arith.mulf %a, %k : tensor<4xf32>
+ %y = arith.mulf %b, %k : tensor<4xf32>
+ %r = arith.select %c, %x, %y : tensor<4xi1>, tensor<4xf32>
+ return %r : tensor<4xf32>
+}
+
+// Three differing operand positions need only one select.
+// CHECK-LABEL: @repeated_pair
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : f32
+// CHECK-NEXT: %[[R:.*]] = math.fma %[[S]], %[[S]], %[[S]] : f32
+// CHECK-NEXT: return %[[R]]
+func.func @repeated_pair(%c: i1, %a: f32, %b: f32) -> f32 {
+ %x = math.fma %a, %a, %a : f32
+ %y = math.fma %b, %b, %b : f32
+ %r = arith.select %c, %x, %y : f32
+ return %r : f...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/234818
More information about the Mlir-commits
mailing list