From 95b6988bb238e85910e30053abd2f5ba4069bba8 Mon Sep 17 00:00:00 2001
From: James Newling <james.newling@gmail.com>
Date: Fri, 23 Aug 2024 09:55:25 -0700
Subject: [PATCH] [DistributeCoresAndObjectFifos][Windows flake] Don't erase op
 too early   (#697)

---
 .../AMDAIEDistributeCoresAndObjectFifos.cpp   | 73 ++++++++++++-------
 1 file changed, 45 insertions(+), 28 deletions(-)

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEDistributeCoresAndObjectFifos.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEDistributeCoresAndObjectFifos.cpp
index 5a268a204..bcc5db0a8 100644
--- a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEDistributeCoresAndObjectFifos.cpp
+++ b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEDistributeCoresAndObjectFifos.cpp
@@ -235,43 +235,60 @@ LogicalResult distributeLocalMemory(ModuleOp moduleOp) {
                 return success();
               })
               .Case<AMDAIE::LogicalObjectFifoFromMemrefOp>(
-                  [&rewriter, &newAlloc](
+                  [&rewriter, &newAlloc, &toBeErased](
                       AMDAIE::LogicalObjectFifoFromMemrefOp logicalObjectFifo) {
                     auto type = llvm::cast<MemRefType>(newAlloc.getType());
+
+                    // Collect all DmaCpyNdOps which have 'logicalObjectFifo' as
+                    // a source. Currently not handling the case of multiple.
+                    SmallVector<AMDAIE::DmaCpyNdOp> dmaOps;
                     for (Operation *objFifoUserOp :
                          logicalObjectFifo->getUsers()) {
                       if (auto dmaOp =
                               dyn_cast<AMDAIE::DmaCpyNdOp>(objFifoUserOp);
                           dmaOp.getSourceObjectFifo() == logicalObjectFifo) {
-                        SmallVector<Value> empty;
-                        rewriter.setInsertionPoint(dmaOp.getSourceObjectFifo());
-                        auto source =
-                            rewriter
-                                .create<AMDAIE::LogicalObjectFifoFromMemrefOp>(
-                                    rewriter.getUnknownLoc(),
-                                    LogicalObjectFifoType::get(type),
-                                    newAlloc.getResult());
-                        rewriter.replaceOp(dmaOp.getSourceObjectFifo(), source);
-                        rewriter.setInsertionPoint(dmaOp);
-                        auto newDmaOp = rewriter.create<AMDAIE::DmaCpyNdOp>(
-                            dmaOp.getLoc(), dmaOp.getTarget(),
-                            dmaOp.getTargetMixedOffsets(),
-                            dmaOp.getTargetMixedSizes(),
-                            dmaOp.getTargetMixedStrides(), source,
-                            dmaOp.getSourceMixedOffsets(),
-                            dmaOp.getSourceMixedSizes(),
-                            dmaOp.getSourceMixedStrides());
-                        rewriter.replaceOp(dmaOp, newDmaOp);
-                        // We have to discard non-zero offsets as subview has
-                        // been replaced by a dedicated allocated memref.
-                        SmallVector<int64_t> allocShape(type.getShape());
-                        (void)discardAllNonZeroOffsets<CopyOpOperateOn::Source>(
-                            rewriter,
-                            cast<AMDAIE::DoublyStridedOpInterface>(
-                                newDmaOp.getOperation()),
-                            allocShape);
+                        dmaOps.push_back(dmaOp);
                       }
                     }
+                    if (dmaOps.size() == 0) return success();
+                    if (dmaOps.size() > 1) {
+                      logicalObjectFifo->emitOpError(
+                          "Case of multiple DMA ops not handled yet (easy "
+                          "extension to logic here)");
+                      return failure();
+                    }
+                    AMDAIE::DmaCpyNdOp dmaOp = dmaOps[0];
+
+                    SmallVector<Value> empty;
+                    rewriter.setInsertionPoint(logicalObjectFifo);
+                    auto source =
+                        rewriter.create<AMDAIE::LogicalObjectFifoFromMemrefOp>(
+                            rewriter.getUnknownLoc(),
+                            LogicalObjectFifoType::get(type),
+                            newAlloc.getResult());
+                    rewriter.replaceAllUsesWith(logicalObjectFifo, source);
+                    toBeErased.push_back(logicalObjectFifo);
+                    rewriter.setInsertionPoint(dmaOp);
+                    auto newDmaOp = rewriter.create<AMDAIE::DmaCpyNdOp>(
+                        dmaOp.getLoc(), dmaOp.getTarget(),
+                        dmaOp.getTargetMixedOffsets(),
+                        dmaOp.getTargetMixedSizes(),
+                        dmaOp.getTargetMixedStrides(), source,
+                        dmaOp.getSourceMixedOffsets(),
+                        dmaOp.getSourceMixedSizes(),
+                        dmaOp.getSourceMixedStrides());
+                    rewriter.replaceAllUsesWith(dmaOp, newDmaOp);
+                    // TODO: maybe this should be left to a DCE somewhere,
+                    // instead of manually erasing unused ops?
+                    toBeErased.push_back(dmaOp);
+                    // We have to discard non-zero offsets as subview has
+                    // been replaced by a dedicated allocated memref.
+                    SmallVector<int64_t> allocShape(type.getShape());
+                    (void)discardAllNonZeroOffsets<CopyOpOperateOn::Source>(
+                        rewriter,
+                        cast<AMDAIE::DoublyStridedOpInterface>(
+                            newDmaOp.getOperation()),
+                        allocShape);
                     return success();
                   })
               .Default([&](Operation *userOp) {