diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIECreateLogicalObjectFifoLink.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIECreateLogicalObjectFifoLink.cpp index 3ffcd6e6f..6533a4e01 100644 --- a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIECreateLogicalObjectFifoLink.cpp +++ b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIECreateLogicalObjectFifoLink.cpp @@ -34,7 +34,6 @@ namespace mlir::iree_compiler::AMDAIE { template LogicalResult checkForContiguousAccessPatterns( ArrayRef> stridedOps) { - for (auto &&[i, stridedOpAndOffset] : llvm::enumerate(stridedOps)) { DoublyStridedCopyOpInterface stridedOp = stridedOpAndOffset.first; std::optional extent; @@ -77,7 +76,8 @@ LogicalResult checkForContiguousAccessPatterns( /// block and an error will be emitted if that's not the case. LogicalResult createLogicalObjectFifoLink( RewriterBase &rewriter, - AMDAIE::LogicalObjectFifoFromMemrefOp logicalObjectFifo) { + AMDAIE::LogicalObjectFifoFromMemrefOp logicalObjectFifo, + SmallVector &newLinkOps) { Attribute memSpace = logicalObjectFifo.getMemorySpace(); if (!memSpace || dyn_cast(memSpace).getInt() != 1) { return success(); @@ -140,12 +140,10 @@ LogicalResult createLogicalObjectFifoLink( // Check that access patterns are not overlapping between consumers // respectively producers. - if (failed( - checkForContiguousAccessPatterns(ins))) { + if (failed(checkForContiguousAccessPatterns(ins))) { return failure(); } - if (failed( - checkForContiguousAccessPatterns(outs))) { + if (failed(checkForContiguousAccessPatterns(outs))) { return failure(); } @@ -161,8 +159,30 @@ LogicalResult createLogicalObjectFifoLink( // Insert the `LogicalObjectFifoLink` after the last user operation. if (lastUserOp) { rewriter.setInsertionPointAfter(lastUserOp); - rewriter.create(rewriter.getUnknownLoc(), - inResults, outResults); + auto linkOp = rewriter.create( + rewriter.getUnknownLoc(), inResults, outResults); + newLinkOps.push_back(linkOp); + } + return success(); +} + +LogicalResult discardLinkNonZeroOffsets(RewriterBase &rewriter, + AMDAIE::LogicalObjectFifoLink linkOp) { + for (Value input : linkOp.getIns()) { + if (auto stridedOp = dyn_cast( + input.getDefiningOp())) { + SmallVector shape; + (void)discardAllNonZeroOffsets(rewriter, + stridedOp, shape); + } + } + for (Value output : linkOp.getOuts()) { + if (auto stridedOp = dyn_cast( + output.getDefiningOp())) { + SmallVector shape; + (void)discardAllNonZeroOffsets(rewriter, + stridedOp, shape); + } } return success(); } @@ -180,10 +200,11 @@ struct AMDAIECreateLogicalObjectFifoLinkPass Operation *parentOp = getOperation(); IRRewriter rewriter(parentOp->getContext()); + SmallVector newLinkOps; WalkResult res = parentOp->walk( [&](AMDAIE::LogicalObjectFifoFromMemrefOp logicalObjectFifo) { - if (failed( - createLogicalObjectFifoLink(rewriter, logicalObjectFifo))) { + if (failed(createLogicalObjectFifoLink(rewriter, logicalObjectFifo, + newLinkOps))) { logicalObjectFifo.emitError() << "couldn't create a link operation"; return WalkResult::interrupt(); } @@ -192,24 +213,10 @@ struct AMDAIECreateLogicalObjectFifoLinkPass if (res.wasInterrupted()) return signalPassFailure(); // Remove all non-zero offsets. - parentOp->walk([&](AMDAIE::LogicalObjectFifoLink linkOp) { - for (Value input : linkOp.getIns()) { - if (auto stridedOp = dyn_cast( - input.getDefiningOp())) { - SmallVector shape; - (void)discardAllNonZeroOffsets( - rewriter, stridedOp, shape); - } - } - for (Value output : linkOp.getOuts()) { - if (auto stridedOp = dyn_cast( - output.getDefiningOp())) { - SmallVector shape; - (void)discardAllNonZeroOffsets( - rewriter, stridedOp, shape); - } - } - }); + for (AMDAIE::LogicalObjectFifoLink linkOp : newLinkOps) { + if (failed(discardLinkNonZeroOffsets(rewriter, linkOp))) + return signalPassFailure(); + } } }; diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/create_logical_objectfifo_link.mlir b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/create_logical_objectfifo_link.mlir index c59dedd93..8da2315a3 100644 --- a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/create_logical_objectfifo_link.mlir +++ b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/create_logical_objectfifo_link.mlir @@ -1,4 +1,4 @@ -// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(func.func(iree-amdaie-create-logical-objectfifo-link, cse))" --verify-diagnostics %s | FileCheck %s +// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(func.func(iree-amdaie-create-logical-objectfifo-link,cse,canonicalize))" --verify-diagnostics %s | FileCheck %s // CHECK-LABEL: func.func @link // CHECK: %[[DMA0:.+]] = amdaie.circular_dma_cpy_nd