[FSM] Add wait_until API

phanrahan · Mar 22, 2024 · ebf0c90 · ebf0c90
1 parent 8c4886f
commit ebf0c90
Show file tree

Hide file tree

Showing 4 changed files with 84 additions and 49 deletions.
diff --git a/magma/__init__.py b/magma/__init__.py
@@ -147,7 +147,7 @@ def set_mantle_target(t):
 from magma.debug import magma_helper_function
 import magma.mantle
 from magma.sum_type import Sum, match, case, Enum2
-from magma.syntax.fsm import fsm, FSMState, wait
+from magma.syntax.fsm import fsm, FSMState, wait, wait_until
 
 ################################################################################
 # BEGIN ALIASES

diff --git a/magma/syntax/fsm.py b/magma/syntax/fsm.py
@@ -43,12 +43,14 @@ def fsm(T: m.Kind, init: FSMState):
     return FSMContext(state_reg)
 
 
-def wait():
-    class State(FSMState):
-        WAIT = 0
-        DONE = 1
+def wait_until(cond: m.Bit):
     curr_fsm = _FSM_STACK.peek()
     with m.no_when():
         state_reg = m.Register(m.Bit, init=0)()
-    state_reg.I @= 1
+        state_reg.I @= 0
+    state_reg.I @= cond
     curr_fsm.cases[-1].exit_stack.enter_context(m.when(state_reg.O))
+
+
+def wait():
+    wait_until(1)
diff --git a/tests/test_syntax/gold/test_fsm_wait.mlir b/tests/test_syntax/gold/test_fsm_wait.mlir
@@ -19,70 +19,71 @@ module attributes {circt.loweringOptions = "locationInfoStyle=none,omitVersionCo
         %11 = hw.constant 2 : i2
         %12 = comb.icmp eq %6, %11 : i2
         %13 = hw.constant 57005 : i16
-        %14 = hw.constant 1 : i1
-        %17 = sv.reg name "Register_inst1" : !hw.inout<i1>
+        %14 = hw.constant 0 : i1
+        %15 = hw.constant 1 : i1
+        %18 = sv.reg name "Register_inst1" : !hw.inout<i1>
         sv.alwaysff(posedge %CLK) {
-            sv.passign %17, %15 : i1
+            sv.passign %18, %16 : i1
         }
-        %18 = hw.constant 0 : i1
+        %19 = hw.constant 0 : i1
         sv.initial {
-            sv.bpassign %17, %18 : i1
+            sv.bpassign %18, %19 : i1
         }
-        %16 = sv.read_inout %17 : !hw.inout<i1>
-        %19 = hw.constant 48879 : i16
-        %20 = hw.constant 1 : i2
+        %17 = sv.read_inout %18 : !hw.inout<i1>
+        %20 = hw.constant 48879 : i16
         %21 = hw.constant 1 : i2
-        %22 = comb.icmp eq %6, %21 : i2
-        %23 = hw.constant 48879 : i16
-        %26 = sv.reg name "Register_inst2" : !hw.inout<i1>
+        %22 = hw.constant 1 : i2
+        %23 = comb.icmp eq %6, %22 : i2
+        %24 = hw.constant 48879 : i16
+        %27 = sv.reg name "Register_inst2" : !hw.inout<i1>
         sv.alwaysff(posedge %CLK) {
-            sv.passign %26, %24 : i1
+            sv.passign %27, %25 : i1
         }
         sv.initial {
-            sv.bpassign %26, %18 : i1
+            sv.bpassign %27, %19 : i1
         }
-        %25 = sv.read_inout %26 : !hw.inout<i1>
-        %27 = hw.constant 57005 : i16
-        %28 = hw.constant 3 : i2
-        %29 = hw.constant 57069 : i16
-        %31 = sv.reg : !hw.inout<i16>
-        %30 = sv.read_inout %31 : !hw.inout<i16>
-        %32 = sv.reg : !hw.inout<i2>
-        %0 = sv.read_inout %32 : !hw.inout<i2>
-        %33 = sv.reg : !hw.inout<i1>
-        %15 = sv.read_inout %33 : !hw.inout<i1>
+        %26 = sv.read_inout %27 : !hw.inout<i1>
+        %28 = hw.constant 57005 : i16
+        %29 = hw.constant 3 : i2
+        %30 = hw.constant 57069 : i16
+        %32 = sv.reg : !hw.inout<i16>
+        %31 = sv.read_inout %32 : !hw.inout<i16>
+        %33 = sv.reg : !hw.inout<i2>
+        %0 = sv.read_inout %33 : !hw.inout<i2>
         %34 = sv.reg : !hw.inout<i1>
-        %24 = sv.read_inout %34 : !hw.inout<i1>
+        %16 = sv.read_inout %34 : !hw.inout<i1>
+        %35 = sv.reg : !hw.inout<i1>
+        %25 = sv.read_inout %35 : !hw.inout<i1>
         sv.alwayscomb {
-            sv.bpassign %32, %6 : i2
-            sv.bpassign %33, %16 : i1
-            sv.bpassign %34, %25 : i1
+            sv.bpassign %34, %14 : i1
+            sv.bpassign %35, %14 : i1
+            sv.bpassign %33, %6 : i2
             sv.if %8 {
-                sv.bpassign %31, %9 : i16
-                sv.bpassign %32, %10 : i2
+                sv.bpassign %32, %9 : i16
+                sv.bpassign %33, %10 : i2
             } else {
                 sv.if %12 {
-                    sv.bpassign %31, %13 : i16
-                    sv.bpassign %33, %14 : i1
-                    sv.if %16 {
-                        sv.bpassign %31, %19 : i16
-                        sv.bpassign %32, %20 : i2
+                    sv.bpassign %32, %13 : i16
+                    sv.bpassign %34, %15 : i1
+                    sv.if %17 {
+                        sv.bpassign %32, %20 : i16
+                        sv.bpassign %33, %21 : i2
                     }
                 } else {
-                    sv.if %22 {
-                        sv.bpassign %31, %23 : i16
-                        sv.bpassign %34, %14 : i1
-                        sv.if %25 {
-                            sv.bpassign %31, %27 : i16
-                            sv.bpassign %32, %28 : i2
+                    sv.if %23 {
+                        sv.bpassign %32, %24 : i16
+                        sv.bpassign %35, %15 : i1
+                        sv.if %26 {
+                            sv.bpassign %32, %28 : i16
+                            sv.bpassign %33, %29 : i2
                         }
                     } else {
-                        sv.bpassign %31, %29 : i16
-                        sv.bpassign %32, %28 : i2
+                        sv.bpassign %32, %30 : i16
+                        sv.bpassign %33, %29 : i2
                     }
                 }
             }
         }
-        hw.output %30 : i16
+        hw.output %31 : i16
     }
 }
diff --git a/tests/test_syntax/test_fsm.py b/tests/test_syntax/test_fsm.py
@@ -52,3 +52,35 @@ class Foo(m.Circuit):
 
     m.compile("build/test_fsm_loop_unroll", Foo, output="mlir")
     assert check_gold(__file__, "test_fsm_loop_unroll.mlir")
+
+
+def test_fsm_wait_until():
+    class State(m.FSMState):
+        INIT = 0
+        RUN = 1
+        WAIT = 2
+        DONE = 3
+
+    class Foo(m.Circuit):
+        io = m.IO(I=m.In(m.Bit), O=m.Out(m.Bits[16]))
+        with m.fsm(State, init=State.INIT) as state:
+            with m.case(State.INIT):
+                io.O @= 0xFEED
+                state.next @= State.WAIT
+            with m.case(State.WAIT):
+                io.O @= 0xDEAD
+                m.wait_until(io.I)
+                io.O @= 0xBEEF
+                state.next @= State.RUN
+            with m.case(State.RUN):
+                io.O @= 0xBEEF
+                # Assumes state.next is not assigned before when
+                m.wait()
+                io.O @= 0xDEAD
+                state.next @= State.DONE
+            with m.case(State.DONE):
+                io.O @= 0xDEED
+                state.next @= State.DONE
+
+    m.compile("build/test_fsm_wait_until", Foo, output="mlir-verilog")
+    assert check_gold(__file__, "test_fsm_wait_until.mlir")