[WIP] Eliding transfers when not required.

compiler/src/iree/compiler/Dialect/Stream/Transforms/ElideAsyncCopies.cpp

@@ -396,6 +396,7 @@ static bool isSafeToElideCloneOp(IREE::Stream::AsyncCloneOp cloneOp,
       llvm::cast<IREE::Stream::ResourceType>(cloneOp.getSource().getType());
   auto targetType =
       llvm::cast<IREE::Stream::ResourceType>(cloneOp.getResult().getType());
+  // DO NOT SUBMIT
   if (sourceType != targetType &&
       sourceType.getLifetime() == IREE::Stream::Lifetime::Constant) {
     LLVM_DEBUG(llvm::dbgs()
@@ -628,6 +629,75 @@ static void elideSliceOp(IREE::Stream::AsyncSliceOp sliceOp) {
   sliceOp.erase();
 }
 
+//===----------------------------------------------------------------------===//
+// IREE::Stream::AsyncTransferOp elision
+//===----------------------------------------------------------------------===//
+
+static bool isTransferRequired(IREE::Stream::AffinityAttr sourceAffinityAttr,
+                               IREE::Stream::AffinityAttr targetAffinityAttr) {
+  if (sourceAffinityAttr == targetAffinityAttr) {
+    return false;
+  }
+  return false;
+}
+
+// DO NOT SUBMIT
+static bool isSafeToDemoteTransferOp(IREE::Stream::AsyncTransferOp transferOp,
+                                     ElisionAnalysis &analysis) {
+  LLVM_DEBUG({
+    llvm::dbgs() << "isSafeToDemoteTransferOp:\n";
+    llvm::dbgs() << "  ";
+    transferOp.print(llvm::dbgs(), analysis.getAsmState());
+    llvm::dbgs() << "\n";
+  });
+
+  // DO NOT SUBMIT busted types
+  // return false;
+
+  // DO NOT SUBMIT if staging on either side
+  auto sourceType =
+      llvm::cast<IREE::Stream::ResourceType>(transferOp.getSource().getType());
+  auto targetType =
+      llvm::cast<IREE::Stream::ResourceType>(transferOp.getResult().getType());
+  if (sourceType.getLifetime() == IREE::Stream::Lifetime::Staging ||
+      targetType.getLifetime() == IREE::Stream::Lifetime::Staging) {
+    LLVM_DEBUG(llvm::dbgs() << "  - transfer required for staging from "
+                            << sourceType << " to " << targetType << "\n");
+    return false;
+  }
+
+  if (isTransferRequired(transferOp.getSourceAffinityAttr(),
+                         transferOp.getResultAffinityAttr())) {
+    LLVM_DEBUG(llvm::dbgs() << "  - transfer required from "
+                            << transferOp.getSourceAffinityAttr() << " to "
+                            << transferOp.getResultAffinityAttr() << "\n");
+    return false;
+  }
+
+  // DO NOT SUBMIT
+  LLVM_DEBUG(llvm::dbgs() << "  + transfer not required from "
+                          << transferOp.getSourceAffinityAttr() << " to "
+                          << transferOp.getResultAffinityAttr()
+                          << "; can demote to clone\n");
+  return true;
+
+  // Not safe.
+  LLVM_DEBUG(llvm::dbgs() << "  - transfer cannot be demoted\n");
+  return false;
+}
+
+// (Potentially) elides a stream.async.transfer op by replacing it with a clone
+// when the transfer is not required. Subsequent iteration may elide the clone.
+static void demoteTransferOp(IREE::Stream::AsyncTransferOp transferOp) {
+  OpBuilder builder(transferOp);
+  auto cloneOp = builder.create<IREE::Stream::AsyncCloneOp>(
+      transferOp.getLoc(), transferOp.getResult().getType(),
+      transferOp.getSource(), transferOp.getSourceSize(),
+      transferOp.getResultSize(), transferOp.getResultAffinityAttr());
+  transferOp.replaceAllUsesWith(cloneOp.getResult());
+  transferOp.erase();
+}
+
 //===----------------------------------------------------------------------===//
 // --iree-stream-elide-async-copies
 //===----------------------------------------------------------------------===//
@@ -654,6 +724,13 @@ static bool tryElideAsyncCopiesInRegion(Region &region,
             }
             return WalkResult::advance();
           })
+          .Case<IREE::Stream::AsyncTransferOp>([&](auto transferOp) {
+            if (isSafeToDemoteTransferOp(transferOp, analysis)) {
+              demoteTransferOp(transferOp);
+              didChange = true;
+            }
+            return WalkResult::advance();
+          })
           .Default([&](auto *op) { return WalkResult::advance(); });
     });
   }