Split L2 input and output objectFifos for memTile/shimTile distribution

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIELogicalObjFifoSplittingUtils.h

@@ -11,14 +11,18 @@
 
 namespace mlir::iree_compiler::AMDAIE {
 
-/// Utility to help fetch those input DmaCpyNd Ops which needs to be split.
-SmallVector<AMDAIE::DmaCpyNdOp> fetchDmaCpyNdOpsToSplitOrCombine(Operation *op);
-
-/// Utility to split logicalobjectfifos given a vector of L2->L1 dma ops.
-LogicalResult splitLogicalObjectFifos(
+/// Split logical objectfifos of the third input of a core given a vector of L2->L1 dma ops.
+LogicalResult splitThirdInputLogicalObjectFifos(
     IRRewriter &rewriter, SmallVector<AMDAIE::DmaCpyNdOp> &l2ToL1DmaOps,
     MLIRContext *context);
 
+/// Split L2 space input and output logical objectFifos.
+LogicalResult splitLogicalObjectFifo(IRRewriter &rewriter,
+                                     AMDAIE::LogicalObjectFifoFromMemrefOp op);
+
+/// Split DmaCpyNd ops between L2 and L3 memory spaces.
+LogicalResult splitDoublyStridedOp(IRRewriter &rewriter, AMDAIE::DmaCpyNdOp op);
+
 }  // namespace mlir::iree_compiler::AMDAIE
 
 #endif

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIESplitLogicalObjFifos.cpp

@@ -0,0 +1,78 @@
+// Copyright 2024 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#include "iree-amd-aie/IR/AMDAIEOps.h"
+#include "iree-amd-aie/Transforms/AMDAIELogicalObjFifoSplittingUtils.h"
+#include "iree-amd-aie/Transforms/Passes.h"
+#include "mlir/IR/Iterators.h"
+#include "mlir/Pass/Pass.h"
+
+#define DEBUG_TYPE "iree-amdaie-split-logical-objectfifos"
+
+namespace mlir::iree_compiler::AMDAIE {
+
+namespace {
+
+class AMDAIESplitLogicalObjFifosPass
+    : public impl::AMDAIESplitLogicalObjFifosBase<
+          AMDAIESplitLogicalObjFifosPass> {
+ public:
+  using AMDAIESplitLogicalObjFifosBase::AMDAIESplitLogicalObjFifosBase;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AMDAIEDialect>();
+  }
+  void runOnOperation() override;
+};
+
+void AMDAIESplitLogicalObjFifosPass::runOnOperation() {
+  ModuleOp moduleOp = getOperation();
+  MLIRContext *context = &getContext();
+  IRRewriter rewriter(context);
+
+  // Walk and collect all dma ops between L3 and L2.
+  SmallVector<AMDAIE::DmaCpyNdOp> l3L2DmaOps;
+  WalkResult res = moduleOp->walk([&](AMDAIE::DmaCpyNdOp op) {
+    std::optional<uint8_t> sourceMemSpace = op.getSourceMemorySpaceAsUInt();
+    std::optional<uint8_t> targetMemSpace = op.getTargetMemorySpaceAsUInt();
+    if (!sourceMemSpace || !targetMemSpace) {
+      op.emitOpError() << "expected a source and target memory space";
+      return WalkResult::interrupt();
+    }
+    if ((sourceMemSpace.value() == 1 && targetMemSpace.value() == 0) ||
+        (sourceMemSpace.value() == 0 && targetMemSpace.value() == 1)) {
+      l3L2DmaOps.push_back(op);
+    }
+    return WalkResult::advance();
+  });
+  if (res.wasInterrupted()) return signalPassFailure();
+
+  // Split the dma ops of L3->L2 / L2->L3.
+  for (AMDAIE::DmaCpyNdOp dmaOp : l3L2DmaOps) {
+    if (failed(splitDoublyStridedOp(rewriter, dmaOp))) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << "Failed to perform splitting of doubly strided op");
+      return signalPassFailure();
+    }
+  }
+
+  // Walk and split input and output objectfifos in L2 memory space.
+  res = moduleOp->walk([&](AMDAIE::LogicalObjectFifoFromMemrefOp op) {
+    if (op.getMemorySpaceAsUInt() != 1) return WalkResult::skip();
+    if (failed(splitLogicalObjectFifo(rewriter, op))) {
+      return WalkResult::interrupt();
+    }
+    return WalkResult::advance();
+  });
+  if (res.wasInterrupted()) return signalPassFailure();
+}
+
+}  // namespace
+
+std::unique_ptr<Pass> createAMDAIESplitLogicalObjFifosPass() {
+  return std::make_unique<AMDAIESplitLogicalObjFifosPass>();
+}
+
+}  // namespace mlir::iree_compiler::AMDAIE

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIESplitLogicalObjFifosForConnectionReuse.cpp

@@ -6,7 +6,6 @@
 #include "iree-amd-aie/IR/AMDAIEOps.h"
 #include "iree-amd-aie/Transforms/AMDAIELogicalObjFifoSplittingUtils.h"
 #include "iree-amd-aie/Transforms/Passes.h"
-// #include "llvm/Support/Debug.h"
 #include "mlir/IR/Iterators.h"
 #include "mlir/Pass/Pass.h"
 
@@ -34,10 +33,20 @@ void AMDAIESplitLogicalObjFifosForConnectionReusePass::runOnOperation() {
   MLIRContext *context = &getContext();
   IRRewriter rewriter(context);
 
-  SmallVector<AMDAIE::DmaCpyNdOp> l2ToL1DmaOps =
-      fetchDmaCpyNdOpsToSplitOrCombine(moduleOp);
-
-  if (failed(splitLogicalObjectFifos(rewriter, l2ToL1DmaOps, context))) {
+  // Walk through CoreOps gathering 3rd input DmaOps (if applicable) which will
+  // be used to split L2 objectFifos of elementwise input for connection reuse.
+  SmallVector<AMDAIE::DmaCpyNdOp> l2ToL1DmaOps;
+  moduleOp->walk([&](AMDAIE::CoreOp coreOp) {
+    SmallVector<Value> inputDmas = coreOp.getInputDmas();
+    if (inputDmas.size() != 3) return WalkResult::skip();
+    auto dmaCpyNdOp = inputDmas[2].getDefiningOp<AMDAIE::DmaCpyNdOp>();
+    assert(dmaCpyNdOp && "expected an amdaie.dma_cpy_nd op");
+    l2ToL1DmaOps.push_back(dmaCpyNdOp);
+    return WalkResult::advance();
+  });
+
+  if (failed(
+          splitThirdInputLogicalObjectFifos(rewriter, l2ToL1DmaOps, context))) {
     LLVM_DEBUG(llvm::dbgs()
                << "Failed to perform splitting of logicalobjectfifos");
     return signalPassFailure();

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/CMakeLists.txt

@@ -99,6 +99,7 @@ iree_cc_library(
     "AMDAIEPropagateDataLayout.cpp"
     "AMDAIERemoveMemorySpace.cpp"
     "AMDAIESinkIntoCore.cpp"
+    "AMDAIESplitLogicalObjFifos.cpp"
     "AMDAIESplitLogicalObjFifosForConnectionReuse.cpp"
     "AMDAIETemporaryAllocBufferization.cpp"
     "AMDAIETile.cpp"

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/PassDetail.h

@@ -78,6 +78,7 @@ namespace mlir::iree_compiler::AMDAIE {
 #define GEN_PASS_DEF_AMDAIEPROPAGATEDATALAYOUT
 #define GEN_PASS_DEF_AMDAIEREMOVEMEMORYSPACE
 #define GEN_PASS_DEF_AMDAIESINKINTOCORE
+#define GEN_PASS_DEF_AMDAIESPLITLOGICALOBJFIFOS
 #define GEN_PASS_DEF_AMDAIESPLITLOGICALOBJFIFOSFORCONNECTIONREUSE
 #define GEN_PASS_DEF_AMDAIETEMPORARYALLOCBUFFERIZATION
 #define GEN_PASS_DEF_AMDAIETILE

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Passes.h

@@ -266,6 +266,9 @@ std::unique_ptr<Pass> createAMDAIERemoveMemorySpacePass();
 /// Create a pass to sink all dependencies into `amdaie.core` operations.
 std::unique_ptr<Pass> createAMDAIESinkIntoCorePass();
 
+/// Create a pass to split logicalobjectfifos for shimTile/memTile distribution.
+std::unique_ptr<Pass> createAMDAIESplitLogicalObjFifosPass();
+
 /// Create a pass to split logicalobjectfifos for connection reuse.
 std::unique_ptr<Pass> createAMDAIESplitLogicalObjFifosForConnectionReusePass();
 

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Passes.td

@@ -629,6 +629,12 @@ def AMDAIESinkIntoCore :
   let constructor = "mlir::iree_compiler::AMDAIE::createAMDAIESinkIntoCorePass()";
 }
 
+def AMDAIESplitLogicalObjFifos :
+  Pass<"iree-amdaie-split-logical-objectfifos", "ModuleOp"> {
+  let summary = "Pass to split L2 buffers to distribute on multiple shimTiles and memTiles.";
+  let constructor = "mlir::iree_compiler::AMDAIE::createAMDAIESplitLogicalObjFifosPass()";
+}
+
 def AMDAIESplitLogicalObjFifosForConnectionReuse :
   Pass<"iree-amdaie-split-logical-objectfifos-for-connection-reuse", "ModuleOp"> {
   let summary = "Pass to split L2 buffers to share inputs of Matmul and Elementwise operations.";

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/CMakeLists.txt

@@ -71,6 +71,7 @@ iree_lit_test_suite(
     "propagate_data_layout.mlir"
     "remove_memory_space.mlir"
     "sink_into_core.mlir"
+    "split_logicalobjfifos.mlir"
     "split_logicalobjfifos_for_connection_reuse.mlir"
     "temporary_alloc_bufferization.mlir"
     "tile_and_fuse_using_scf_for.mlir"