[PyCDE] Import MLIR: provide hook to preprocess ops (#8624)

Adds the ability for users to load the MLIR file and preprocess
or skip pre-lowering ops.

Author: teqdruid

frontends/PyCDE/src/pycde/system.py

@@ -164,6 +164,8 @@ def import_mlir(self,
                   lowering: Optional[List[str]] = None,
                   output_filename: Optional[str] = None,
                   importer: Optional[Callable] = None,
+                  preprocess_op: Optional[Callable[[ir.OpView],
+                                                   Optional[ir.OpView]]] = None,
                   debug: bool = False) -> Dict[str, Any]:
     """Import mlir asm created elsewhere into our space. Exactly one of the
     arguments module_str or file must be provided.
@@ -176,6 +178,8 @@ def import_mlir(self,
       output_filename: Optional filename to set on the imported ops.
       importer: Optional custom importer function which takes a System and the
                 op. Must return a 3-tuple of (name, obj, proxy).
+      preprocess_op: Optional preprocessing function which takes an Operation and
+                     returns a modified Operation or None to skip.
       debug: Whether to enable debug output for the import process.
     """
 
@@ -234,6 +238,12 @@ def import_builtin(
 
     ret: Dict[str, Any] = {}
     for op in compat_mod.body:
+      if preprocess_op is not None:
+        op = preprocess_op(op)
+        if op is None:
+          # If the op was preprocessed to None, skip it.
+          continue
+
       # TODO: handle symbolrefs pointing to potentially renamed symbols.
       imported_obj = None
       imported_name = None

frontends/PyCDE/test/test_import_hw_modules.py

@@ -1,6 +1,5 @@
 # RUN: %PYTHON% %s %t | FileCheck %s
 
-from pycde.circt.ir import Module as IrModule
 from pycde.circt.dialects import hw
 
 from pycde import Input, Output, System, generator, Module, types
@@ -22,6 +21,8 @@
   hw.output %0 : i1
 }
 
+hw.module.extern @skip()
+
 hw.type_scope @otherScope {}
 """
 
@@ -54,8 +55,15 @@ def generate(ports):
     ports.out1 = outs[1]
 
 
+def preprocess_op(op):
+  """Preprocess the operation to skip the `skip` module."""
+  if isinstance(op, hw.HWModuleExternOp) and op.name.value == "skip":
+    return None
+  return op
+
+
 system = System([Top], output_directory=sys.argv[1])
-imported_modules = system.import_mlir(mlir_module)
+imported_modules = system.import_mlir(mlir_module, preprocess_op=preprocess_op)
 imported_modules["add"].add_metadata(Metadata(name="add"))
 system.generate()
 TypeAlias.declare_aliases(system.mod)