[Mlir-commits] [mlir] [mlir] [arith] Combine op for select (PR #234818)

[Renato Golin]

================
@@ -3128,6 +3130,119 @@ struct SelectToExtUI : public OpRewritePattern<arith::SelectOp> {
   }
 };
 
+///   select(c, F(a0, a1), F(b0, b1))
+///     -> F(select(c, a0, b0), select(c, a1, b1)).
+/// For a non-poison condition, each lane of F receives exactly the operands of
+/// the selected producer, so poison in the other producer remains masked. For
+/// a poison condition, the original result is poison; the replacement may
+/// refine it, but must not introduce immediate undefined behavior. In
+/// particular, speculatability must also hold with the new operands.
+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.
----------------
rengolin wrote:

is this really necessary and sufficient? 

https://github.com/llvm/llvm-project/pull/234818