Fixes in gen_msb

calyxir · Aug 16, 2023 · 272e29c · 272e29c
1 parent 6742d36
commit 272e29c
Showing 1 changed file with 12 additions and 7 deletions.
diff --git a/calyx-py/calyx/gen_msb.py b/calyx-py/calyx/gen_msb.py
@@ -17,7 +17,14 @@
     Control,
     CombGroup,
 )
-from calyx.builder import Builder, ComponentBuilder, const, HI, while_
+from calyx.builder import (
+    Builder,
+    CellGroupBuilder,
+    ComponentBuilder,
+    const,
+    HI,
+    while_with,
+)
 
 
 def gen_msb_calc(width: int, int_width: int) -> List[Component]:
@@ -116,17 +123,15 @@ def gen_msb_calc(width: int, int_width: int) -> List[Component]:
 
     comp.control += [
         wr_cur_val,
-        while_(
-            neq.out,
-            cur_val_cond,
+        while_with(
+            CellGroupBuilder(neq, cur_val_cond),
             [incr_count, shift_cur_val],
         ),
         decr_count,
         wr_count,
         wr_val_build,
-        while_(
-            neq.out,
-            count_cond,
+        while_with(
+            CellGroupBuilder(neq, count_cond),
             [decr_count, shift_val_build],
         ),
         wr_val,