Elide chained transfers at Flow level

Signed-off-by: Alex Vasile <48962821+Alex-Vasile@users.noreply.github.com>

Author: Alex-Vasile

compiler/src/iree/compiler/Dialect/Flow/IR/FlowOpFolders.cpp

@@ -1192,11 +1192,35 @@ struct ElideRedundantTransfer : public OpRewritePattern<TensorTransferOp> {
   }
 };
 
+// Attempts to identify trivial case of chained transfer ops (A -> B -> C) and rewrite it as (A -> C)
+// Writes it as A -> B and A -> C relying on dead code elimination to remove the unused A -> B transfer.
+struct ElideChainedTransfer : public OpRewritePattern<TensorTransferOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(TensorTransferOp currTransferOp,
+                                PatternRewriter &rewriter) const override {
+    auto baseValue =
+        IREE::Util::TiedOpInterface::findTiedBaseValue(currTransferOp.getOperand());
+    if (auto prevTransferOp = dyn_cast_if_present<IREE::Flow::TensorTransferOp>(
+            baseValue.getDefiningOp())) {
+      rewriter.replaceOpWithNewOp<TensorTransferOp>(
+        currTransferOp, 
+        currTransferOp->getResultTypes(),
+        prevTransferOp.getOperand(),
+        currTransferOp.getOperandDims(),
+        currTransferOp.getTarget());
+      return success();
+    }
+    return failure();
+  } 
+};
+
+
 } // namespace
 
 void TensorTransferOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                    MLIRContext *context) {
   results.insert<ElideRedundantTransfer>(context);
+  results.insert<ElideChainedTransfer>(context);
 }
 
 //===----------------------------------------------------------------------===//

compiler/src/iree/compiler/Dialect/Flow/IR/test/tensor_folding.mlir

@@ -424,6 +424,36 @@ util.func public @ElideRedundantTransfer(%operand: tensor<4x?xf32>, %dim: index)
 
 // -----
 
+// CHECK-LABEL: @ElideChainedTransferTwoTransfers
+//  CHECK-SAME: (%[[OPERAND:.+]]: tensor<1xf16>)
+util.func public @ElideChainedTransferTwoTransfers(%operand: tensor<1xf16>) -> tensor<1xf16> {
+  // CHECK-NOT: flow.tensor.transfer
+  %redundant = flow.tensor.transfer %operand : tensor<1xf16> to "target1"
+  // CHECK: %[[RESULT:.+]] = flow.tensor.transfer %[[OPERAND]]
+  %result = flow.tensor.transfer %redundant : tensor<1xf16> to "target2"
+  // CHECK-NEXT: util.return %[[RESULT]]
+  util.return %result : tensor<1xf16>
+}
+
+// -----
+
+// CHECK-LABEL: @ElideChainedTransferFourTransfers
+//  CHECK-SAME: (%[[OPERAND:.+]]: tensor<1xf16>)
+util.func public @ElideChainedTransferFourTransfers(%operand: tensor<1xf16>) -> tensor<1xf16> {
+  // CHECK-NOT: flow.tensor.transfer
+  %redundant = flow.tensor.transfer %operand : tensor<1xf16> to "target1"
+  // CHECK-NOT: flow.tensor.transfer
+  %redundant2 = flow.tensor.transfer %redundant : tensor<1xf16> to "target2"
+  // CHECK-NOT: flow.tensor.transfer
+  %redundant3 = flow.tensor.transfer %redundant2 : tensor<1xf16> to "target3"
+  // CHECK: %[[RESULT:.+]] = flow.tensor.transfer %[[OPERAND]]
+  %result = flow.tensor.transfer %redundant3 : tensor<1xf16> to "target4"
+  // CHECK-NEXT: util.return %[[RESULT]]
+  util.return %result : tensor<1xf16>
+}
+
+// -----
+
 // CHECK-LABEL: @sliceConst0D
 util.func public @sliceConst0D() -> tensor<i32> {
   %0 = arith.constant dense<0> : tensor<i32>