[MLIR] Add support for compile guard select

Author: rsetaluri

magma/backend/mlir/hardware_module.py

@@ -52,7 +52,7 @@
 from magma.circuit import AnonymousCircuitType, CircuitKind, DefineCircuitKind
 from magma.clock import Reset, ResetN, AsyncReset, AsyncResetN
 from magma.common import filter_by_key, assert_false
-from magma.compile_guard import get_compile_guard_data
+from magma.compile_guard import get_compile_guard_data, CompileGuardSelect
 from magma.digital import Digital, DigitalMeta
 from magma.inline_verilog_expression import InlineVerilogExpression
 from magma.inline_verilog2 import InlineVerilog2
@@ -880,6 +880,27 @@ def visit_magma_xmr_source(self, module: ModuleWrapper) -> bool:
             sv.ReadInOutOp(operands=[in_out], results=[result])
         return True
 
+    @wrap_with_not_implemented_error
+    def visit_magma_compile_guard_select(self, module: ModuleWrapper) -> bool:
+        inst = module.module
+        defn = type(inst)
+        assert isinstance(defn, CompileGuardSelect)
+        assert len(defn.keys) + 1 == len(module.operands)
+        assert len(module.results) == 1
+        result = module.results[0]
+        mlir_type = magma_type_to_mlir_type(defn.T)
+        reg = self.ctx.new_value(hw.InOutType(mlir_type))
+        sv.RegOp(results=[reg])
+        with contextlib.ExitStack() as stack:
+            for i, key in enumerate(defn.keys):
+                if_def = sv.IfDefOp(key)
+                stack.enter_context(push_block(if_def.then_block))
+                sv.AssignOp(operands=[reg, module.operands[i]])
+                stack.enter_context(push_block(if_def.else_block))
+            sv.AssignOp(operands=[reg, module.operands[-1]])
+        sv.ReadInOutOp(operands=[reg], results=[result])
+        return True
+
     @wrap_with_not_implemented_error
     def visit_inline_verilog(self, module: ModuleWrapper) -> bool:
         inst = module.module
@@ -921,6 +942,8 @@ def visit_instance(self, module: ModuleWrapper) -> bool:
             return self.visit_magma_xmr_sink(module)
         if isinstance(defn, XMRSource):
             return self.visit_magma_xmr_source(module)
+        if isinstance(defn, CompileGuardSelect):
+            return self.visit_magma_compile_guard_select(module)
         if getattr(defn, "inline_verilog_strs", []):
             return self.visit_inline_verilog(module)
         if isprimitive(defn):

tests/test_backend/test_mlir/examples.py

@@ -488,12 +488,15 @@ class xmr_bind_asserts(m.Circuit):
 class simple_compile_guard(m.Circuit):
     io = m.IO(I=m.In(m.Bit), O=m.Out(m.Bit)) + m.ClockIO()
     with m.compile_guard(
-            "COND1", defn_name="COND1_compile_guard", type="defined"):
-        out = m.Register(m.Bit)()(io.I)
+            "COND1", defn_name="COND1_compile_guard", type="defined"
+    ):
+        m.Register(m.Bit)()(io.I)
     with m.compile_guard(
-            "COND2", defn_name="COND2_compile_guard", type="undefined"):
-        out = m.Register(m.Bit)()(io.I)
-    io.O @= io.I
+            "COND2", defn_name="COND2_compile_guard", type="undefined"
+    ):
+        m.Register(m.Bit)()(io.I)
+    out = m.compile_guard_select(COND1=io.I, COND2=~io.I, default=0)
+    io.O @= out
 
 
 m.passes.clock.WireClockPass(simple_compile_guard).run()

tests/test_backend/test_mlir/golds/simple_compile_guard.mlir

@@ -22,13 +22,27 @@ module attributes {circt.loweringOptions = "locationInfoStyle=none"} {
         %0 = sv.read_inout %1 : !hw.inout<i1>
     }
     hw.module @simple_compile_guard(%I: i1, %CLK: i1) -> (O: i1) {
+        %1 = hw.constant -1 : i1
+        %0 = comb.xor %1, %I : i1
+        %2 = hw.constant 0 : i1
+        %4 = sv.reg : !hw.inout<i1>
+        sv.ifdef "COND1" {
+            sv.assign %4, %I : i1
+        } else {
+            sv.ifdef "COND2" {
+                sv.assign %4, %0 : i1
+            } else {
+                sv.assign %4, %2 : i1
+            }
+        }
+        %3 = sv.read_inout %4 : !hw.inout<i1>
         sv.ifdef "COND1" {
             hw.instance "COND1_compile_guard" @COND1_compile_guard(port_0: %I: i1, port_1: %CLK: i1) -> ()
         }
         sv.ifdef "COND2" {
         } else {
             hw.instance "COND2_compile_guard" @COND2_compile_guard(port_0: %I: i1, port_1: %CLK: i1) -> ()
         }
-        hw.output %I : i1
+        hw.output %3 : i1
     }
 }

tests/test_backend/test_mlir/golds/simple_compile_guard.v

@@ -29,18 +29,26 @@ module simple_compile_guard(
   output O
 );
 
+  reg _GEN;
   `ifdef COND1
+    assign _GEN = I;
     COND1_compile_guard COND1_compile_guard (
       .port_0 (I),
       .port_1 (CLK)
     );
+  `else  // COND1
+    `ifdef COND2
+      assign _GEN = ~I;
+    `else  // COND2
+      assign _GEN = 1'h0;
+    `endif // COND2
   `endif // COND1
   `ifndef COND2
     COND2_compile_guard COND2_compile_guard (
       .port_0 (I),
       .port_1 (CLK)
     );
   `endif // not def COND2
-  assign O = I;
+  assign O = _GEN;
 endmodule