[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.
Xilinx · Dec 11, 2024 · 7d0dc15 · 7d0dc15
1 parent 45528b1
commit 7d0dc15
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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