[Mlir-commits] [mlir] [MLIR][Vector] Add canonicalization for interleave/deinterleave chain (PR #196979)

[Artem Kroviakov]

================
@@ -8346,6 +8346,46 @@ Value mlir::vector::selectPassthru(OpBuilder &builder, Value mask,
 // InterleaveOp
 //===----------------------------------------------------------------------===//
 
+namespace {
+
+/// This canonicalization folds the following round-trip identity:
+///  interleave(deinterleave(x).even, deinterleave(x).odd) -> x
+struct InterleaveDeinterleaveFolder : public OpRewritePattern<InterleaveOp> {
+  using Base::Base;
+
+  LogicalResult matchAndRewrite(InterleaveOp interleaveOp,
+                                PatternRewriter &rewriter) const override {
+    auto lhsDefOp = interleaveOp.getLhs().getDefiningOp<DeinterleaveOp>();
+    auto rhsDefOp = interleaveOp.getRhs().getDefiningOp<DeinterleaveOp>();
+    if (!lhsDefOp || !rhsDefOp)
+      return rewriter.notifyMatchFailure(
+          interleaveOp,
+          "expected both operands to be defined by a deinterleave op");
+    if (lhsDefOp != rhsDefOp)
+      return rewriter.notifyMatchFailure(
+          interleaveOp,
+          "expected both operands to be defined by the same deinterleave op");
+    if (auto deinterleaveRes = cast<OpResult>(interleaveOp.getLhs());
+        deinterleaveRes.getResultNumber())
+      return rewriter.notifyMatchFailure(interleaveOp,
+                                         "expected the lhs operand to be the "
+                                         "first result of the deinterleave op");
+    if (auto deinterleaveRes = cast<OpResult>(interleaveOp.getRhs());
+        deinterleaveRes.getResultNumber() != 1)
+      return rewriter.notifyMatchFailure(
+          interleaveOp, "expected the rhs operand to be the second result of "
+                        "the deinterleave op");
+    rewriter.replaceOp(interleaveOp, lhsDefOp.getSource());
----------------
akroviakov wrote:

Using a for loop now

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