From 7d0dc1507eb54ee1dad9b3ec8d33ebbdcd432143 Mon Sep 17 00:00:00 2001
From: Ronan Keryell <ronan.keryell@amd.com>
Date: Tue, 10 Dec 2024 18:05:18 -0800
Subject: [PATCH] [CIR] Add new --cir-keep-aie-device pass

Extract the aie.device operation as the top-module operation.
Anything else is removed.
For now just keep the first aie.device.
---
 include/aie/CIR/CIRToAIEPasses.h  |  3 +++
 include/aie/CIR/CIRToAIEPasses.td | 16 ++++++++++++++++
 lib/CIR/CIRToAIEPasses.cpp        | 32 +++++++++++++++++++++++++++++++
 3 files changed, 51 insertions(+)

diff --git a/include/aie/CIR/CIRToAIEPasses.h b/include/aie/CIR/CIRToAIEPasses.h
index 489deb82da..a9367655d5 100644
--- a/include/aie/CIR/CIRToAIEPasses.h
+++ b/include/aie/CIR/CIRToAIEPasses.h
@@ -31,6 +31,9 @@ createCIRToAIEInlineKernelLambdaPass();
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
 createCIRToAIEDecaptureKernelPass();
 
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
+createCIRKeepAIEDevice();
+
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "aie/CIR/CIRToAIEPasses.h.inc"
diff --git a/include/aie/CIR/CIRToAIEPasses.td b/include/aie/CIR/CIRToAIEPasses.td
index 5a6a8ac3a8..0486308a26 100644
--- a/include/aie/CIR/CIRToAIEPasses.td
+++ b/include/aie/CIR/CIRToAIEPasses.td
@@ -89,4 +89,20 @@ def CIRToAIEDecaptureKernel : Pass<"cir-to-aie-decapture-kernel", "mlir::ModuleO
   ];
 }
 
+def CIRKeepAIEDevice : Pass<"cir-keep-aie-device", "mlir::ModuleOp"> {
+
+  let summary = "Remove everything but the first aie.device";
+
+  let description = [{ Extract the aie.device from a CIR module by removing everything but the first aie.device.
+
+  TODO: handle multiple devices.
+  }];
+
+  let constructor = "xilinx::AIE::CIR::createCIRKeepAIEDevice()";
+  let dependentDialects = [
+    "cir::CIRDialect",
+    "xilinx::AIE::AIEDialect",
+  ];
+}
+
 #endif
diff --git a/lib/CIR/CIRToAIEPasses.cpp b/lib/CIR/CIRToAIEPasses.cpp
index 9e23f752bc..367519fcf3 100644
--- a/lib/CIR/CIRToAIEPasses.cpp
+++ b/lib/CIR/CIRToAIEPasses.cpp
@@ -987,6 +987,34 @@ struct CIRToAIEDecaptureKernel
   }
 };
 
+struct CIRKeepAIEDevice : CIRKeepAIEDeviceBase<CIRKeepAIEDevice> {
+  void runOnOperation() override {
+    auto module = getOperation();
+    mlir::OpBuilder b{module};
+    llvm::SmallVector<mlir::Operation *> opToDelete;
+    xilinx::AIE::DeviceOp d;
+    // Use pre-order walk for early exit to pick the first aie.device for now
+    module->walk<mlir::WalkOrder::PreOrder>([&](mlir::Operation *op) {
+      if (!mlir::isa<xilinx::AIE::DeviceOp>(op))
+        return mlir::WalkResult::advance();
+      d = mlir::cast<xilinx::AIE::DeviceOp>(op);
+      return mlir::WalkResult::interrupt();
+    });
+    // Extract the aie.device from its function up to the top first
+    // operation of the module
+    if (d)
+      d->moveBefore(&module.front());
+    // Erase all the top-module operations but the aie.device
+    for (auto &op : module) {
+      if (auto dev = mlir::dyn_cast<xilinx::AIE::DeviceOp>(op))
+        if (dev == d)
+          continue;
+      opToDelete.push_back(&op);
+    }
+    eraseOpsAndUsers(opToDelete);
+  }
+};
+
 } // namespace
 
 std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
@@ -1008,4 +1036,8 @@ createCIRToAIEDecaptureKernelPass() {
   return std::make_unique<CIRToAIEDecaptureKernel>();
 }
 
+std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>> createCIRKeepAIEDevice() {
+  return std::make_unique<CIRKeepAIEDevice>();
+}
+
 } // namespace xilinx::AIE::CIR