[Mlir-commits] [mlir] [mlir] [arith] Combine op for select (PR #234818)
Chuanqi Xu
llvmlistbot at llvm.org
Thu Oct 22 19:20:19 PDT 2026
https://github.com/ChuanqiXu9 created https://github.com/llvm/llvm-project/pull/234818
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.
>From a2ad0485ad636e92df2b78055b620512a46df506 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 | 102 +++++
mlir/lib/Dialect/Arith/IR/CMakeLists.txt | 1 +
.../select-common-elementwise-test-ops.mlir | 219 +++++++++++
.../Arith/select-common-elementwise.mlir | 362 ++++++++++++++++++
mlir/test/lib/Dialect/Test/TestOps.td | 21 +
5 files changed, 705 insertions(+)
create mode 100644 mlir/test/Dialect/Arith/select-common-elementwise-test-ops.mlir
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..2c8091597bd9a2 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 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-test-ops.mlir b/mlir/test/Dialect/Arith/select-common-elementwise-test-ops.mlir
new file mode 100644
index 00000000000000..62948c66971ba0
--- /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 00000000000000..12d3dbfd587c47
--- /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 : f32
+}
+
+// The full ordered pair is the key, not just the true operand.
+// CHECK-LABEL: @same_true_different_false
+// CHECK-NEXT: %[[S0:.*]] = arith.select %arg0, %arg1, %arg2 : f32
+// CHECK-NEXT: %[[S1:.*]] = arith.select %arg0, %arg1, %arg3 : f32
+// CHECK-NEXT: %[[R:.*]] = math.fma %[[S0]], %[[S1]], %arg4 : f32
+// CHECK-NEXT: return %[[R]]
+func.func @same_true_different_false(%c: i1, %a: f32, %b: f32, %d: f32, %k: f32) -> f32 {
+ %x = math.fma %a, %a, %k : f32
+ %y = math.fma %b, %d, %k : f32
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// CHECK-LABEL: @reversed_pairs
+// CHECK-NEXT: %[[S0:.*]] = arith.select %arg0, %arg1, %arg2 : f32
+// CHECK-NEXT: %[[S1:.*]] = arith.select %arg0, %arg2, %arg1 : f32
+// CHECK-NEXT: %[[R:.*]] = arith.divf %[[S0]], %[[S1]] : f32
+// CHECK-NEXT: return %[[R]]
+func.func @reversed_pairs(%c: i1, %a: f32, %b: f32) -> f32 {
+ %x = arith.divf %a, %b : f32
+ %y = arith.divf %b, %a : f32
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// Select the input type, even when its width differs from the result width.
+// CHECK-LABEL: @conversion
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : vector<4xi1>, vector<4xf32>
+// CHECK-NEXT: %[[R:.*]] = arith.fptosi %[[S]] : vector<4xf32> to vector<4xi8>
+// CHECK-NEXT: return %[[R]]
+func.func @conversion(%c: vector<4xi1>, %a: vector<4xf32>, %b: vector<4xf32>) -> vector<4xi8> {
+ %x = arith.fptosi %a : vector<4xf32> to vector<4xi8>
+ %y = arith.fptosi %b : vector<4xf32> to vector<4xi8>
+ %r = arith.select %c, %x, %y : vector<4xi1>, vector<4xi8>
+ return %r : vector<4xi8>
+}
+
+// CHECK-LABEL: @same_fastmath
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : f32
+// CHECK-NEXT: %[[R:.*]] = math.sqrt %[[S]] fastmath<nnan,ninf> : f32
+// CHECK-NEXT: return %[[R]]
+func.func @same_fastmath(%c: i1, %a: f32, %b: f32) -> f32 {
+ %x = math.sqrt %a fastmath<nnan,ninf> : f32
+ %y = math.sqrt %b fastmath<nnan,ninf> : f32
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// CHECK-LABEL: @different_fastmath
+// CHECK-NEXT: %[[X:.*]] = math.sqrt %arg1 fastmath<nnan> : f32
+// CHECK-NEXT: %[[Y:.*]] = math.sqrt %arg2 fastmath<ninf> : f32
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : f32
+// CHECK-NEXT: return %[[R]]
+func.func @different_fastmath(%c: i1, %a: f32, %b: f32) -> f32 {
+ %x = math.sqrt %a fastmath<nnan> : f32
+ %y = math.sqrt %b fastmath<ninf> : f32
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// Overflow flags are operation properties and must match as well.
+// CHECK-LABEL: @different_overflow
+// 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]]
+func.func @different_overflow(%c: 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 %c, %x, %y : i32
+ return %r : i32
+}
+
+// CHECK-LABEL: @multiple_uses_true
+// CHECK-NEXT: %[[X:.*]] = math.sqrt %arg1 : f32
+// CHECK-NEXT: %[[Y:.*]] = math.sqrt %arg2 : f32
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : f32
+// CHECK-NEXT: return %[[R]], %[[X]]
+func.func @multiple_uses_true(%c: i1, %a: f32, %b: f32) -> (f32, f32) {
+ %x = math.sqrt %a : f32
+ %y = math.sqrt %b : f32
+ %r = arith.select %c, %x, %y : f32
+ return %r, %x : f32, f32
+}
+
+// CHECK-LABEL: @multiple_uses_false
+// CHECK-NEXT: %[[X:.*]] = math.sqrt %arg1 : f32
+// CHECK-NEXT: %[[Y:.*]] = math.sqrt %arg2 : f32
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : f32
+// CHECK-NEXT: return %[[R]], %[[Y]]
+func.func @multiple_uses_false(%c: i1, %a: f32, %b: f32) -> (f32, f32) {
+ %x = math.sqrt %a : f32
+ %y = math.sqrt %b : f32
+ %r = arith.select %c, %x, %y : f32
+ return %r, %y : f32, f32
+}
+
+// CHECK-LABEL: @three_distinct_pairs
+// CHECK-NEXT: %[[S0:.*]] = arith.select %arg0, %arg1, %arg4 : f32
+// CHECK-NEXT: %[[S1:.*]] = arith.select %arg0, %arg2, %arg5 : f32
+// CHECK-NEXT: %[[S2:.*]] = arith.select %arg0, %arg3, %arg6 : f32
+// CHECK-NEXT: %[[R:.*]] = math.fma %[[S0]], %[[S1]], %[[S2]] : f32
+// CHECK-NEXT: return %[[R]]
+func.func @three_distinct_pairs(%c: i1, %a: f32, %b: f32, %d: f32, %e: f32, %f: f32, %g: f32) -> f32 {
+ %x = math.fma %a, %b, %d : f32
+ %y = math.fma %e, %f, %g : f32
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// Memory-effect freedom alone is insufficient: dynamic integer division cannot
+// be speculated because the divisor can be zero.
+// CHECK-LABEL: @not_speculatable
+// CHECK-NEXT: %[[X:.*]] = arith.divsi %arg1, %arg3 : i32
+// CHECK-NEXT: %[[Y:.*]] = arith.divsi %arg2, %arg3 : i32
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : i32
+// CHECK-NEXT: return %[[R]]
+func.func @not_speculatable(%c: i1, %a: i32, %b: i32, %k: i32) -> i32 {
+ %x = arith.divsi %a, %k : i32
+ %y = arith.divsi %b, %k : i32
+ %r = arith.select %c, %x, %y : i32
+ return %r : i32
+}
+
+// CHECK-LABEL: @different_ops
+// CHECK-NEXT: %[[X:.*]] = math.sqrt %arg1 : f32
+// CHECK-NEXT: %[[Y:.*]] = math.exp %arg2 : f32
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : f32
+// CHECK-NEXT: return %[[R]]
+func.func @different_ops(%c: i1, %a: f32, %b: f32) -> f32 {
+ %x = math.sqrt %a : f32
+ %y = math.exp %b : f32
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// CHECK-LABEL: @scalar_condition_vector
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : vector<4xf32>
+// CHECK-NEXT: %[[R:.*]] = math.sqrt %[[S]] : vector<4xf32>
+// CHECK-NEXT: return %[[R]]
+func.func @scalar_condition_vector(%c: i1, %a: vector<4xf32>, %b: vector<4xf32>) -> vector<4xf32> {
+ %x = math.sqrt %a : vector<4xf32>
+ %y = math.sqrt %b : vector<4xf32>
+ %r = arith.select %c, %x, %y : vector<4xf32>
+ return %r : vector<4xf32>
+}
+
+// CHECK-LABEL: @dynamic_tensor
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : tensor<?x4xi1>, tensor<?x4xf32>
+// CHECK-NEXT: %[[R:.*]] = math.sqrt %[[S]] : tensor<?x4xf32>
+// CHECK-NEXT: return %[[R]]
+func.func @dynamic_tensor(%c: tensor<?x4xi1>, %a: tensor<?x4xf32>, %b: tensor<?x4xf32>) -> tensor<?x4xf32> {
+ %x = math.sqrt %a : tensor<?x4xf32>
+ %y = math.sqrt %b : tensor<?x4xf32>
+ %r = arith.select %c, %x, %y : tensor<?x4xi1>, tensor<?x4xf32>
+ return %r : tensor<?x4xf32>
+}
+
+// CHECK-LABEL: @scalable_vector
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : vector<2x[4]xi1>, vector<2x[4]xf32>
+// CHECK-NEXT: %[[R:.*]] = math.sqrt %[[S]] : vector<2x[4]xf32>
+// CHECK-NEXT: return %[[R]]
+func.func @scalable_vector(%c: vector<2x[4]xi1>, %a: vector<2x[4]xf32>, %b: vector<2x[4]xf32>) -> vector<2x[4]xf32> {
+ %x = math.sqrt %a : vector<2x[4]xf32>
+ %y = math.sqrt %b : vector<2x[4]xf32>
+ %r = arith.select %c, %x, %y : vector<2x[4]xi1>, vector<2x[4]xf32>
+ return %r : vector<2x[4]xf32>
+}
+
+// CHECK-LABEL: @tensor_encoding
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : tensor<4xi1, "layout">, tensor<4xf32, "layout">
+// CHECK-NEXT: %[[R:.*]] = math.sqrt %[[S]] : tensor<4xf32, "layout">
+// CHECK-NEXT: return %[[R]]
+func.func @tensor_encoding(%c: tensor<4xi1, "layout">, %a: tensor<4xf32, "layout">, %b: tensor<4xf32, "layout">) -> tensor<4xf32, "layout"> {
+ %x = math.sqrt %a : tensor<4xf32, "layout">
+ %y = math.sqrt %b : tensor<4xf32, "layout">
+ %r = arith.select %c, %x, %y : tensor<4xi1, "layout">, tensor<4xf32, "layout">
+ return %r : tensor<4xf32, "layout">
+}
+
+// Greedy canonicalization must remove dead conversions and then factor the
+// muls, without relying on a subsequent CSE pass.
+// CHECK-LABEL: @conversion_multiply_chain
+// CHECK-NEXT: %[[S0:.*]] = arith.select %arg0, %arg1, %arg3 : tensor<4xi1>, tensor<4xf32>
+// CHECK-NEXT: %[[S1:.*]] = arith.select %arg0, %arg2, %arg4 : tensor<4xi1>, tensor<4xf32>
+// CHECK-NEXT: %[[M:.*]] = arith.mulf %[[S0]], %[[S1]] : tensor<4xf32>
+// CHECK-NEXT: %[[R:.*]] = arith.fptosi %[[M]] : tensor<4xf32> to tensor<4xi8>
+// CHECK-NEXT: return %[[R]]
+func.func @conversion_multiply_chain(%c: tensor<4xi1>, %a: tensor<4xf32>, %alpha: tensor<4xf32>, %b: tensor<4xf32>, %beta: tensor<4xf32>) -> tensor<4xi8> {
+ %x = arith.mulf %a, %alpha : tensor<4xf32>
+ %y = arith.mulf %b, %beta : tensor<4xf32>
+ %ix = arith.fptosi %x : tensor<4xf32> to tensor<4xi8>
+ %iy = arith.fptosi %y : tensor<4xf32> to tensor<4xi8>
+ %r = arith.select %c, %ix, %iy : tensor<4xi1>, tensor<4xi8>
+ return %r : tensor<4xi8>
+}
+
+// Select producers can also be factored through the outer select.
+// CHECK-LABEL: @nested_selects
+// CHECK-NEXT: %[[C:.*]] = arith.select %arg0, %arg1, %arg2 : i1
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg4, %arg5 : f32
+// CHECK-NEXT: %[[R:.*]] = arith.select %[[C]], %arg3, %[[S]] : f32
+// CHECK-NEXT: return %[[R]]
+func.func @nested_selects(%c: i1, %d: i1, %e: i1, %a: f32, %b: f32, %k: f32) -> f32 {
+ %x = arith.select %d, %a, %b : f32
+ %y = arith.select %e, %a, %k : f32
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// CHECK-LABEL: @memory_reads
+// CHECK-NEXT: %[[X:.*]] = memref.load %arg1[%arg3] : memref<?xf32>
+// CHECK-NEXT: %[[Y:.*]] = memref.load %arg2[%arg3] : memref<?xf32>
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : f32
+// CHECK-NEXT: return %[[R]]
+func.func @memory_reads(%c: i1, %a: memref<?xf32>, %b: memref<?xf32>, %i: index) -> f32 {
+ %x = memref.load %a[%i] : memref<?xf32>
+ %y = memref.load %b[%i] : memref<?xf32>
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// Pure reductions are not elementwise.
+// CHECK-LABEL: @reductions
+// CHECK-NEXT: %[[X:.*]] = vector.reduction <add>, %arg1 : vector<4xf32> into f32
+// CHECK-NEXT: %[[Y:.*]] = vector.reduction <add>, %arg2 : vector<4xf32> into f32
+// CHECK-NEXT: %[[R:.*]] = arith.select %arg0, %[[X]], %[[Y]] : f32
+// CHECK-NEXT: return %[[R]]
+func.func @reductions(%c: i1, %a: vector<4xf32>, %b: vector<4xf32>) -> f32 {
+ %x = vector.reduction <add>, %a : vector<4xf32> into f32
+ %y = vector.reduction <add>, %b : vector<4xf32> into f32
+ %r = arith.select %c, %x, %y : f32
+ return %r : f32
+}
+
+// CHECK-LABEL: @same_overflow
+// CHECK-NEXT: %[[S:.*]] = arith.select %arg0, %arg1, %arg2 : i32
+// CHECK-NEXT: %[[R:.*]] = arith.muli %[[S]], %arg3 overflow<nsw> : i32
+// CHECK-NEXT: return %[[R]]
+func.func @same_overflow(%c: i1, %a: i32, %b: i32, %k: i32) -> i32 {
+ %x = arith.muli %a, %k overflow<nsw> : i32
+ %y = arith.muli %b, %k overflow<nsw> : i32
+ %r = arith.select %c, %x, %y : i32
+ return %r : i32
+}
+
+// Single-use producers outside a loop must not be sunk into the loop.
+// CHECK-LABEL: @producers_outside_loop
+// CHECK-NEXT: %[[X:.*]] = math.sqrt %arg4 : f32
+// CHECK-NEXT: %[[Y:.*]] = math.sqrt %arg5 : f32
+// CHECK-NEXT: %[[LOOP:.*]] = scf.for %[[I:.*]] = %arg0 to %arg1 step %arg2 iter_args(%[[ACC:.*]] = %arg6) -> (f32) {
+// CHECK-NEXT: %[[C:.*]] = arith.cmpi slt, %[[I]], %arg3 : index
+// CHECK-NEXT: %[[S:.*]] = arith.select %[[C]], %[[X]], %[[Y]] : f32
+// CHECK-NEXT: %[[SUM:.*]] = arith.addf %[[ACC]], %[[S]] : f32
+// CHECK-NEXT: scf.yield %[[SUM]] : f32
+// CHECK-NEXT: }
+// CHECK-NEXT: return %[[LOOP]] : f32
+func.func @producers_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
+ %result = scf.for %i = %lb to %ub step %step iter_args(%acc = %init) -> (f32) {
+ %c = arith.cmpi slt, %i, %split : index
+ %r = arith.select %c, %x, %y : f32
+ %sum = arith.addf %acc, %r : f32
+ scf.yield %sum : f32
+ }
+ return %result : f32
+}
+
+// Check the block of each producer, not just whether the producers share one.
+// CHECK-LABEL: @false_producer_outside_loop
+// CHECK-NEXT: %[[Y:.*]] = math.sqrt %arg5 : f32
+// CHECK-NEXT: %[[LOOP:.*]] = scf.for %[[I:.*]] = %arg0 to %arg1 step %arg2 iter_args(%[[ACC:.*]] = %arg6) -> (f32) {
+// CHECK-NEXT: %[[X:.*]] = math.sqrt %arg4 : f32
+// CHECK-NEXT: %[[C:.*]] = arith.cmpi slt, %[[I]], %arg3 : index
+// CHECK-NEXT: %[[S:.*]] = arith.select %[[C]], %[[X]], %[[Y]] : f32
+// CHECK-NEXT: %[[SUM:.*]] = arith.addf %[[ACC]], %[[S]] : f32
+// CHECK-NEXT: scf.yield %[[SUM]] : f32
+// CHECK-NEXT: }
+// CHECK-NEXT: return %[[LOOP]] : f32
+func.func @false_producer_outside_loop(%lb: index, %ub: index, %step: index, %split: index, %a: f32, %b: f32, %init: f32) -> f32 {
+ %y = math.sqrt %b : f32
+ %result = scf.for %i = %lb to %ub step %step iter_args(%acc = %init) -> (f32) {
+ %x = math.sqrt %a : f32
+ %c = arith.cmpi slt, %i, %split : index
+ %r = arith.select %c, %x, %y : f32
+ %sum = arith.addf %acc, %r : f32
+ scf.yield %sum : f32
+ }
+ return %result : f32
+}
+
+// Factoring inside a loop is allowed when all three operations share a block.
+// CHECK-LABEL: @producers_inside_loop
+// CHECK-NEXT: %[[LOOP:.*]] = scf.for %[[I:.*]] = %arg0 to %arg1 step %arg2 iter_args(%[[ACC:.*]] = %arg6) -> (f32) {
+// CHECK-NEXT: %[[C:.*]] = arith.cmpi slt, %[[I]], %arg3 : index
+// CHECK-NEXT: %[[S:.*]] = arith.select %[[C]], %arg4, %arg5 : f32
+// CHECK-NEXT: %[[SQRT:.*]] = math.sqrt %[[S]] : f32
+// CHECK-NEXT: %[[SUM:.*]] = arith.addf %[[ACC]], %[[SQRT]] : f32
+// CHECK-NEXT: scf.yield %[[SUM]] : f32
+// CHECK-NEXT: }
+// CHECK-NEXT: return %[[LOOP]] : f32
+func.func @producers_inside_loop(%lb: index, %ub: index, %step: index, %split: index, %a: f32, %b: f32, %init: f32) -> f32 {
+ %result = scf.for %i = %lb to %ub step %step iter_args(%acc = %init) -> (f32) {
+ %x = math.sqrt %a : f32
+ %y = math.sqrt %b : f32
+ %c = arith.cmpi slt, %i, %split : index
+ %r = arith.select %c, %x, %y : f32
+ %sum = arith.addf %acc, %r : f32
+ scf.yield %sum : f32
+ }
+ return %result : f32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 0a64885d8e2921..e16ab0b677d96e 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -854,6 +854,27 @@ def ElementwiseMappableOp : TEST_Op<"elementwise_mappable",
let results = (outs Variadic<AnyType>);
}
+def PureElementwiseOp : TEST_Op<"pure_elementwise", [Elementwise, Pure]> {
+ let summary = "pure elementwise op with arbitrary operand and result types";
+ let arguments = (ins Variadic<AnyType>);
+ let results = (outs Variadic<AnyType>);
+}
+
+def ReadElementwiseOp : TEST_Op<"read_elementwise",
+ [Elementwise, AlwaysSpeculatable, MemoryEffects<[MemRead]>]> {
+ let summary = "speculatable elementwise op with a memory read effect";
+ let arguments = (ins Variadic<AnyType>);
+ 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