Decouple packToDma pass from AIR pipeline (nod-ai#773)

It's not necessary to go from `LinalgExt::PackOp` -> `air::DmaMemcpyNdOp` -> `AMDAIE::DmaCpyNdOp`. With `packToDma` pass decoupled from AIR pipeline, we can refactor the codes to convert directly from PackOp to AMDAIE::DmaCpyNdOp (will address in the next PR). For AIR pipeline, `DecomposeLinalgExtPackUnPackToAIR` can be used for both pack-peel and pad-pack pipelines. It doesn't have use `packToDma` pass. As a consequence, `AMDAIECanonicalizeDma` pass can also be eliminated, since similar logic already exists in MLIR-AIR and `AMDAIECanonicalizeDoublyStridedOp` for objectFifo.
newling · Sep 13, 2024 · e8d2ec6 · e8d2ec6
1 parent a87850e
commit e8d2ec6
Show file tree

Hide file tree

Showing 8 changed files with 1 addition and 203 deletions.
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIECanonicalizeDma.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIECanonicalizeDma.cpp
diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/CMakeLists.txt b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/CMakeLists.txt
@@ -50,7 +50,6 @@ iree_cc_library(
     "AMDAIEAssignLogicalObjectFifoDepth.cpp"
     "AMDAIEAssignNpuDmaBdIds.cpp"
     "AMDAIEBufferizeToAllocation.cpp"
-    "AMDAIECanonicalizeDma.cpp"
     "AMDAIECanonicalizeNpuDmaCpyNd.cpp"
     "AMDAIECanonicalizeDoublyStridedOp.cpp"
     "AMDAIECombineStridedOps.cpp"

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/PassDetail.h b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/PassDetail.h
@@ -27,7 +27,6 @@ namespace mlir::iree_compiler::AMDAIE {
 #define GEN_PASS_DEF_AMDAIEASSIGNNPUDMABDIDS
 #define GEN_PASS_DEF_AMDAIEBRIDGETOAIR
 #define GEN_PASS_DEF_AMDAIEBUFFERIZETOALLOCATION
-#define GEN_PASS_DEF_AMDAIECANONICALIZEDMA
 #define GEN_PASS_DEF_AMDAIECANONICALIZEDOUBLYSTRIDEDOP
 #define GEN_PASS_DEF_AMDAIECANONICALIZENPUDMACPYND
 #define GEN_PASS_DEF_AMDAIECLEANUP

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Passes.cpp b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Passes.cpp
@@ -670,11 +670,7 @@ void addMLIRAIRLoweringPasses(OpPassManager &passManager, AMDAIEDevice device) {
     passManager.addPass(createCanonicalizerPass());
   }
 
-  // TODO (Erwei): Figure out a way to work with AMDAIEPackToDmaPass.
-  if (clUseTilePipeline == TilePassPipeline::PackPeelPipeline)
-    passManager.addPass(createAMDAIEDecomposeLinalgExtPackUnPackToAIRPass());
-  else
-    passManager.addPass(createAMDAIEPackToDmaPass());
+  passManager.addPass(createAMDAIEDecomposeLinalgExtPackUnPackToAIRPass());
 
   // TODO(newling) adding createCanonicalizerPass introduces a dma copy lowering
   // failure. Understand and fix.
@@ -691,7 +687,6 @@ void addMLIRAIRLoweringPasses(OpPassManager &passManager, AMDAIEDevice device) {
   }
   passManager.addPass(createCanonicalizerPass());
   passManager.addPass(createCSEPass());
-  passManager.addPass(createAMDAIECanonicalizeDmaPass());
   passManager.addPass(xilinx::air::createCopyToDmaPass());
   passManager.addPass(createCanonicalizerPass());
   passManager.addPass(createCSEPass());

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Passes.h b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Passes.h
@@ -77,9 +77,6 @@ std::unique_ptr<Pass> createAMDAIEBridgeToAIRPass();
 std::unique_ptr<Pass> createAMDAIEBufferizeToAllocationPass(
     AMDAIEBufferizeToAllocationOptions options = {});
 
-/// Create pass to apply canonicalization to air.dma_memcpy_nd op's.
-std::unique_ptr<Pass> createAMDAIECanonicalizeDmaPass();
-
 /// Create pass to canonicalize `amdaie.npu.dma_cpy_nd` operations.
 std::unique_ptr<Pass> createAMDAIECanonicalizeNpuDmaCpyNdPass();
 

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Passes.td b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/Passes.td
@@ -82,12 +82,6 @@ def AMDAIEBufferizeToAllocation :
   ];
 }
 
-def AMDAIECanonicalizeDma :
-  Pass<"iree-amdaie-canonicalize-dma", ""> {
-  let summary = "Apply caonicaliztions to air.dma_memcpy_nd op's";
-  let constructor = "mlir::iree_compiler::AMDAIE::createAMDAIECanonicalizeDmaPass()";
-}
-
 def AMDAIECanonicalizeDoublyStridedOp :
     Pass<"iree-amdaie-canonicalize-doubly-strided-op", ""> {
   let summary = "Canonicalize doubly strided DMA operations.";
@@ -98,7 +92,6 @@ def AMDAIECanonicalizeDoublyStridedOp :
   ];
 }
 
-
 def AMDAIECanonicalizeNpuDmaCpyNd :
   Pass<"iree-amdaie-canonicalize-npu-dma-cpy-nd", "ModuleOp"> {
   let summary = "Canonicalize npu.dma_cpy_nd operations.";

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/CMakeLists.txt b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/CMakeLists.txt
@@ -16,7 +16,6 @@ iree_lit_test_suite(
     "assign_npu_dma_bd_ids.mlir"
     "bridge_to_air.mlir"
     "bufferize_to_allocation.mlir"
-    "canonicalize_dma.mlir"
     "canonicalize_doubly_strided_op.mlir"
     "canonicalize_npu_dma_cpy_nd.mlir"
     "combine_strided_ops.mlir"

diff --git a/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/canonicalize_dma.mlir b/compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/canonicalize_dma.mlir