Review comment address including offset/size/stride of linearized buffer

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIESplitBuffers.cpp → compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIESplitLogicalObjectFifos.cpp

@@ -7,35 +7,31 @@
 #include "iree-amd-aie/Transforms/AMDAIEDmaUtils.h"
 #include "iree-amd-aie/Transforms/Passes.h"
 #include "iree-amd-aie/Transforms/Transforms.h"
-#include "iree/compiler/Codegen/TransformStrategies/GPU/Common.h"
-#include "mlir/Dialect/Linalg/IR/Linalg.h"
-#include "mlir/Dialect/SCF/Transforms/Transforms.h"
-#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/Iterators.h"
 #include "mlir/Pass/Pass.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
 
-#define DEBUG_TYPE "iree-amdaie-split-buffers"
+#define DEBUG_TYPE "iree-amdaie-split-logical-objectfifos"
 
 namespace mlir::iree_compiler::AMDAIE {
 
 namespace {
 
-class AMDAIESplitBuffersPass
-    : public impl::AMDAIESplitBuffersBase<AMDAIESplitBuffersPass> {
+class AMDAIESplitLogicalObjectFifosPass
+    : public impl::AMDAIESplitLogicalObjectFifosBase<
+          AMDAIESplitLogicalObjectFifosPass> {
  public:
-  using AMDAIESplitBuffersBase::AMDAIESplitBuffersBase;
+  using AMDAIESplitLogicalObjectFifosBase::AMDAIESplitLogicalObjectFifosBase;
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<AMDAIEDialect>();
   }
   void runOnOperation() override;
 };
 
-void AMDAIESplitBuffersPass::runOnOperation() {
+void AMDAIESplitLogicalObjectFifosPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
-  IRRewriter rewriter(moduleOp.getContext());
+  MLIRContext *context = &getContext();
+  IRRewriter rewriter(context);
 
   SmallVector<AMDAIE::DmaCpyNdOp> l2ToL1DmaOps;
   // We are currently walking through CoreOps gathering 3rd Input DmaOp (if
@@ -60,7 +56,7 @@ void AMDAIESplitBuffersPass::runOnOperation() {
         l2ToL1DmaOp.getTargetObjectFifo();
     Value targetAllocOp = targetObjectFifo.getMemref();
 
-    // Now we'll create a narrowed L2 buffer.
+    // Now we'll create a narrowed linearized L2 buffer.
     rewriter.setInsertionPoint(sourceAllocOp);
     auto oldSourceMemRefType = cast<MemRefType>(sourceAllocOp.getType());
     auto targetMemRefType = cast<MemRefType>(targetAllocOp.getType());
@@ -93,11 +89,25 @@ void AMDAIESplitBuffersPass::runOnOperation() {
         rewriter.getUnknownLoc(), LogicalObjectFifoType::get(type),
         newAllocOp.getResult(), sourceObjectFifo.getTiles());
 
-    // Create new L3 -> L2 Dma Op.
+    // Create new L3 -> L2 Dma Op. Since the narrowed L2 buffer is linearized,
+    // we need to form offset/size/stride corresponding to the linearized
+    // buffer.
+    SmallVector<OpFoldResult, 4> staticOffsets(
+        4, getAsIndexOpFoldResult(context, 0));
+    SmallVector<OpFoldResult, 4> staticSizes(
+        4, getAsIndexOpFoldResult(context, 1));
+    SmallVector<OpFoldResult, 4> staticStrides(
+        4, getAsIndexOpFoldResult(context, 0));
+    OpFoldResult linearizedShape =
+        getAsIndexOpFoldResult(context, newAllocType.getNumElements());
+    staticSizes[staticSizes.size() - 1] = linearizedShape;
+    staticStrides[staticStrides.size() - 1] =
+        getAsIndexOpFoldResult(context, 1);
+    staticStrides[staticStrides.size() - 2] = linearizedShape;
     rewriter.setInsertionPoint(l3ToL2DmaOp);
     rewriter.create<AMDAIE::DmaCpyNdOp>(
-        l3ToL2DmaOp.getLoc(), source, l3ToL2DmaOp.getTargetMixedOffsets(),
-        l3ToL2DmaOp.getTargetMixedSizes(), l3ToL2DmaOp.getTargetMixedStrides(),
+        l3ToL2DmaOp.getLoc(), source, llvm::ArrayRef(staticOffsets),
+        llvm::ArrayRef(staticSizes), llvm::ArrayRef(staticStrides),
         l3ToL2DmaOp.getSource(), l3ToL2DmaOp.getSourceMixedOffsets(),
         l3ToL2DmaOp.getSourceMixedSizes(), l3ToL2DmaOp.getSourceMixedStrides());
 
@@ -107,8 +117,8 @@ void AMDAIESplitBuffersPass::runOnOperation() {
         l2ToL1DmaOp.getLoc(), l2ToL1DmaOp.getTarget(),
         l2ToL1DmaOp.getTargetMixedOffsets(), l2ToL1DmaOp.getTargetMixedSizes(),
         l2ToL1DmaOp.getTargetMixedStrides(), source,
-        l2ToL1DmaOp.getSourceMixedOffsets(), l2ToL1DmaOp.getSourceMixedSizes(),
-        l2ToL1DmaOp.getSourceMixedStrides());
+        llvm::ArrayRef(staticOffsets), llvm::ArrayRef(staticSizes),
+        llvm::ArrayRef(staticStrides));
     rewriter.replaceOp(l2ToL1DmaOp, newL2ToL1DmaOp);
     // We have to discard non-zero offsets as subview has been replaced by a
     // dedicated allocated memref.
@@ -138,8 +148,8 @@ void AMDAIESplitBuffersPass::runOnOperation() {
 
 }  // namespace
 
-std::unique_ptr<Pass> createAMDAIESplitBuffersPass() {
-  return std::make_unique<AMDAIESplitBuffersPass>();
+std::unique_ptr<Pass> createAMDAIESplitLogicalObjectFifosPass() {
+  return std::make_unique<AMDAIESplitLogicalObjectFifosPass>();
 }
 
 }  // namespace mlir::iree_compiler::AMDAIE

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

@@ -82,7 +82,7 @@ iree_cc_library(
     "AMDAIEPad.cpp"
     "AMDAIEPeelForLoop.cpp"
     "AMDAIEPropagateDataLayout.cpp"
-    "AMDAIESplitBuffers.cpp"
+    "AMDAIESplitLogicalObjectFifos.cpp"
     "AMDAIETile.cpp"
     "AMDAIETileAndFuse.cpp"
     "AMDAIEUtils.cpp"

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

@@ -62,7 +62,7 @@ namespace mlir::iree_compiler::AMDAIE {
 #define GEN_PASS_DEF_AMDAIEPAD
 #define GEN_PASS_DEF_AMDAIEPEELFORLOOP
 #define GEN_PASS_DEF_AMDAIEPROPAGATEDATALAYOUT
-#define GEN_PASS_DEF_AMDAIESPLITBUFFERS
+#define GEN_PASS_DEF_AMDAIESPLITLOGICALOBJECTFIFOS
 #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

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

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

@@ -198,7 +198,7 @@ std::unique_ptr<Pass> createAMDAIEPeelForLoopPass(
     AMDAIEPeelForLoopOptions options = {});
 
 /// Create a pass to split buffers.
-std::unique_ptr<Pass> createAMDAIESplitBuffersPass();
+std::unique_ptr<Pass> createAMDAIESplitLogicalObjectFifosPass();
 
 /// Create pass to tile TilingInterface operations.
 std::unique_ptr<Pass> createAMDAIETilePass(AMDAIETileOptions options = {});

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

@@ -402,10 +402,10 @@ def AMDAIEPropagateDataLayout :
   let constructor = "mlir::iree_compiler::AMDAIE::createAMDAIEPropagateDataLayoutPass()";
 }
 
-def AMDAIESplitBuffers :
-  Pass<"iree-amdaie-split-buffers", "ModuleOp"> {
+def AMDAIESplitLogicalObjectFifos :
+  Pass<"iree-amdaie-split-logical-objectfifos", "ModuleOp"> {
   let summary = "Pass to split L2 buffers to share inputs of Matmul and Elementwise operations.";
-  let constructor = "mlir::iree_compiler::AMDAIE::createAMDAIESplitBuffersPass()";
+  let constructor = "mlir::iree_compiler::AMDAIE::createAMDAIESplitLogicalObjectFifosPass()";
 }
 
 def AMDAIETile :

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

@@ -54,7 +54,7 @@ iree_lit_test_suite(
     "pad.mlir"
     "peel_for_loop.mlir"
     "propagate_data_layout.mlir"
-    "split_buffers.mlir"
+    "split_logical_objectfifos.mlir"
     "tile_and_fuse_using_scf_for.mlir"
     "tile_and_fuse_using_scf_forall.mlir"
     "tile_copy_using_scf_for.mlir"