[OM] Preserve non-OM operations in OM LinkModules (#8109)

This commit changes OM LinkModule pass to preserve non-OM operations.

lib/Dialect/OM/Transforms/LinkModules.cpp

@@ -13,6 +13,7 @@
 #include "circt/Dialect/OM/OMOps.h"
 #include "circt/Dialect/OM/OMPasses.h"
 #include "circt/Support/Namespace.h"
+#include "mlir/IR/Block.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/SymbolTable.h"
@@ -39,7 +40,9 @@ using SymMappingTy =
     llvm::DenseMap<std::pair<ModuleOp, StringAttr>, StringAttr>;
 
 struct ModuleInfo {
-  ModuleInfo(mlir::ModuleOp module) : module(module) {}
+  ModuleInfo(mlir::ModuleOp module) : module(module) {
+    block = std::make_unique<Block>();
+  }
 
   // Populate `symbolToClasses`.
   LogicalResult initialize();
@@ -52,6 +55,9 @@ struct ModuleInfo {
 
   // A target module.
   ModuleOp module;
+
+  // A block to store original operations.
+  std::unique_ptr<Block> block;
 };
 
 struct LinkModulesPass
@@ -65,8 +71,10 @@ LogicalResult ModuleInfo::initialize() {
   for (auto &op : llvm::make_early_inc_range(module.getOps())) {
     if (auto classLike = dyn_cast<ClassLike>(op))
       symbolToClasses.insert({classLike.getSymNameAttr(), classLike});
-    else
-      op.erase();
+    else {
+      // Keep the op.
+      op.moveBefore(block.get(), block->end());
+    }
   }
   return success();
 }
@@ -293,11 +301,13 @@ void LinkModulesPass::runOnOperation() {
 
   // Finally move operations to the toplevel module.
   auto *block = toplevelModule.getBody();
-  for (auto info : modules) {
+  for (auto &info : modules) {
     block->getOperations().splice(block->end(),
                                   info.module.getBody()->getOperations());
-    // Erase the module.
-    info.module.erase();
+    // Restore non-OM operations.
+    assert(info.module.getBody()->empty());
+    info.module.getBody()->getOperations().splice(
+        info.module.getBody()->begin(), info.block->getOperations());
   }
 }
 

test/Dialect/OM/link-modules.mlir

@@ -2,7 +2,7 @@
 
 module {
   // CHECK-LABEL: module
-  // CHECK-NOT:   module
+  // CHECK-NEXT:   module
   // CHECK-NOT:   om.class.extern
   // CHECK-LABEL: om.class @A
   // CHECK-LABEL: om.class @Conflict_A
@@ -61,12 +61,12 @@ module {
 
 // -----
 
-// Check that OM ops are deleted.  Make the "delete-me" op a landmine that will
-// cause a symbol collision if it is _not_ deleted.
+// Check that OM ops are not deleted.  Make the "dont-delete-me" op a landmine that will
+// cause a symbol collision if it is moved to the top level.
 module {
   module {
-    // CHECK-NOT: delete-me
-    "delete-me"() {sym_name = "Bar"} : () -> ()
+    // CHECK: dont-delete-me
+    "dont-delete-me"() {sym_name = "Bar"} : () -> ()
     om.class @Foo() {
       om.class.fields
     }

test/om-linker/Inputs/a.mlir

@@ -5,4 +5,5 @@ module {
   om.class @Conflict(){
     om.class.fields
   }
+  hw.module @hello() {}
 }

test/om-linker/Inputs/b.mlir

@@ -6,4 +6,5 @@ module {
   om.class @Conflict(){
     om.class.fields
   }
+  hw.module.extern @hello()
 }

test/om-linker/link.mlir

@@ -1,5 +1,20 @@
 // RUN: om-linker %S/Inputs/a.mlir %S/Inputs/b.mlir %S/Inputs/other.mlir | FileCheck %s
 // CHECK:      module {
+// CHECK-NEXT:   module attributes {om.namespace = "a"} {
+// CHECK-NEXT:     hw.module @hello() {
+// CHECK-NEXT:       hw.output
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
+// CHECK-NEXT:   module attributes {om.namespace = "b"} {
+// CHECK-NEXT:     hw.module.extern @hello()
+// CHECK-NEXT:   }
+// CHECK-NEXT:   module attributes {om.namespace = "other"} {
+// CHECK-NEXT:     hw.module @HW(in %a : i1, out b : i1) {
+// CHECK:            hw.output
+// CHECK-NEXT:     }
+// CHECK-NEXT:     emit.file "foo.sv" sym @Emit {
+// CHECK-NEXT:     }
+// CHECK-NEXT:   }
 // CHECK-NEXT:   om.class @A(%arg: i1) {
 // CHECK-NEXT:     om.class.fields
 // CHECK-NEXT:   }