[ObjectFifo] Create a new pass to split logical objectFifos (#659)

-- This commit introduces a new pass
`--iree-amdaie-split-logical-objectfifos-for-connection-reuse` to
split logical objectFifos for dealing with Matmul+Elementwise.
-- Also contains a utility to check whether splitting can be performed.
-- It addresses sub-action 2 as well from
#644

Signed-off-by: Abhishek Varma <[email protected]>

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

@@ -0,0 +1,23 @@
+// 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
+
+#ifndef IREE_AMD_AIE_TRANSFORMS_AMDAIELOGICALOBJFIFOSPLITTINGUTILS_H_
+#define IREE_AMD_AIE_TRANSFORMS_AMDAIELOGICALOBJFIFOSPLITTINGUTILS_H_
+
+#include "iree-amd-aie/IR/AMDAIEOps.h"
+
+namespace mlir::iree_compiler::AMDAIE {
+
+/// Utility to split logicalobjectfifos given a struct
+/// `SplittingLogicalObjectFifoData` which contains all the required data to
+/// perform the splitting.
+LogicalResult splitLogicalObjectFifos(
+    IRRewriter &rewriter, SmallVector<AMDAIE::DmaCpyNdOp> &l2ToL1DmaOps,
+    MLIRContext *context);
+
+}  // namespace mlir::iree_compiler::AMDAIE
+
+#endif

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

@@ -0,0 +1,71 @@
+// 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 "llvm/Support/Debug.h"
+#include "mlir/IR/Iterators.h"
+#include "mlir/Pass/Pass.h"
+
+#define DEBUG_TYPE "iree-amdaie-split-logical-objectfifos-for-connection-reuse"
+
+namespace mlir::iree_compiler::AMDAIE {
+
+namespace {
+
+/// Utility to help fetch those input DmaCpyNd Ops which needs to be split.
+static SmallVector<AMDAIE::DmaCpyNdOp> fetchDmaCpyNdOpsToSplit(
+    ModuleOp moduleOp) {
+  SmallVector<AMDAIE::DmaCpyNdOp> l2ToL1DmaOps;
+  // We are currently walking through CoreOps gathering 3rd Input DmaOp (if
+  // applicable) from them.
+  // TODO(avarma): We will generalize this later.
+  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();
+  });
+  return l2ToL1DmaOps;
+}
+
+class AMDAIESplitLogicalObjFifosForConnectionReusePass
+    : public impl::AMDAIESplitLogicalObjFifosForConnectionReuseBase<
+          AMDAIESplitLogicalObjFifosForConnectionReusePass> {
+ public:
+  using AMDAIESplitLogicalObjFifosForConnectionReuseBase::
+      AMDAIESplitLogicalObjFifosForConnectionReuseBase;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AMDAIEDialect>();
+  }
+  void runOnOperation() override;
+};
+
+void AMDAIESplitLogicalObjFifosForConnectionReusePass::runOnOperation() {
+  ModuleOp moduleOp = getOperation();
+  MLIRContext *context = &getContext();
+  IRRewriter rewriter(context);
+
+  SmallVector<AMDAIE::DmaCpyNdOp> l2ToL1DmaOps =
+      fetchDmaCpyNdOpsToSplit(moduleOp);
+
+  if (failed(splitLogicalObjectFifos(rewriter, l2ToL1DmaOps, context))) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Failed to perform splitting of logicalobjectfifos");
+    return signalPassFailure();
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<Pass> createAMDAIESplitLogicalObjFifosForConnectionReusePass() {
+  return std::make_unique<AMDAIESplitLogicalObjFifosForConnectionReusePass>();
+}
+
+}  // namespace mlir::iree_compiler::AMDAIE

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

@@ -74,6 +74,7 @@ iree_cc_library(
     "AMDAIEInsertLoopsForVectorization.cpp"
     "AMDAIELinkExecutables.cpp"
     "AMDAIELocalizeLogicalObjectFifo.cpp"
+    "AMDAIELogicalObjFifoSplittingUtils.cpp"
     "AMDAIELowerExecutableTarget.cpp"
     "AMDAIELowerFuncArgs.cpp"
     "AMDAIELowerToAIE.cpp"
@@ -88,6 +89,7 @@ iree_cc_library(
     "AMDAIEPeelForLoop.cpp"
     "AMDAIEPropagateDataLayout.cpp"
     "AMDAIESinkIntoCore.cpp"
+    "AMDAIESplitLogicalObjFifosForConnectionReuse.cpp"
     "AMDAIETile.cpp"
     "AMDAIETileAndFuse.cpp"
     "AMDAIEUtils.cpp"

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

@@ -68,6 +68,7 @@ namespace mlir::iree_compiler::AMDAIE {
 #define GEN_PASS_DEF_AMDAIEPEELFORLOOP
 #define GEN_PASS_DEF_AMDAIEPROPAGATEDATALAYOUT
 #define GEN_PASS_DEF_AMDAIESINKINTOCORE
+#define GEN_PASS_DEF_AMDAIESPLITLOGICALOBJFIFOSFORCONNECTIONREUSE
 #define GEN_PASS_DEF_AMDAIETILE
 #define GEN_PASS_DEF_AMDAIETILEANDFUSE
 #define GEN_PASS_DEF_AMDAIEVECTORIZATION

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

@@ -594,6 +594,7 @@ void addAMDAIEObjectFifoLoweringPasses(OpPassManager &passManager) {
   passManager.addPass(createAMDAIEDistributeCoresAndObjectFifosPass());
   passManager.addPass(createCSEPass());
   passManager.addPass(createCanonicalizerPass());
+  passManager.addPass(createAMDAIESplitLogicalObjFifosForConnectionReusePass());
 
   passManager.addPass(createAMDAIEDmaToCircularDmaPass());
   passManager.addNestedPass<func::FuncOp>(createAMDAIECreateAIEWorkgroupPass());

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

@@ -224,6 +224,9 @@ std::unique_ptr<Pass> createAMDAIEPeelForLoopPass(
 /// Create a pass to sink all dependencies into `amdaie.core` operations.
 std::unique_ptr<Pass> createAMDAIESinkIntoCorePass();
 
+/// Create a pass to split logicalobjectfifos for connection reuse.
+std::unique_ptr<Pass> createAMDAIESplitLogicalObjFifosForConnectionReusePass();
+
 /// Create pass to tile TilingInterface operations.
 std::unique_ptr<Pass> createAMDAIETilePass(AMDAIETileOptions options = {});
 

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

@@ -470,6 +470,12 @@ def AMDAIESinkIntoCore :
   let constructor = "mlir::iree_compiler::AMDAIE::createAMDAIESinkIntoCorePass()";
 }
 
+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.";
+  let constructor = "mlir::iree_compiler::AMDAIE::createAMDAIESplitLogicalObjFifosForConnectionReusePass()";
+}
+
 def AMDAIETile :
     InterfacePass<"iree-amdaie-tile", "mlir::FunctionOpInterface"> {
   let summary = "Pass to tile TilingInterface operations.";

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

@@ -60,6 +60,7 @@ iree_lit_test_suite(
     "peel_for_loop.mlir"
     "propagate_data_layout.mlir"
     "sink_into_core.mlir"
+    "split_logicalobjfifos_for_connection_reuse.mlir"
     "tile_and_fuse_using_scf_for.mlir"
     "tile_and_fuse_using_scf_forall.mlir"
     "tile_copy_using_scf_for.mlir"