Add preliminary support for hierarchy in sequential

magma/syntax/sequential.py

@@ -27,7 +27,7 @@ def visit_Return(self, node):
         return node
 
 
-def get_initial_value_map(init_func):
+def get_initial_value_map(init_func, defn_env):
     """
     Parses __init__ funciton of the form
 
@@ -47,6 +47,8 @@ def __init__(self):
     """
     initial_value_map = {}
     init_def = get_ast(init_func).body[0]
+    init_def = ExecuteEscapedPythonExpressions(defn_env).visit(init_def)
+    init_def = SpecializeConstantInts(defn_env).visit(init_def)
     for stmt in init_def.body:
         # We only support basic assignments of the form
         #     self.x: m.Bits(2) = m.bits(0, 2)
@@ -121,17 +123,22 @@ def gen_register_instances(initial_value_map):
     """
     register_instances = ""
     for name, (value, type_) in initial_value_map.items():
-        # TODO: Only support m.bits(x, y) for now
-        assert isinstance(value, ast.Call) and \
-            isinstance(value.func, ast.Attribute) and \
-            isinstance(value.func.value, ast.Name) and \
-            value.func.value.id == "m" and \
-            value.func.attr == "bits"
-        assert isinstance(value.args[0], ast.Num)
-        assert isinstance(value.args[1], ast.Num)
-        n = value.args[1].n
-        init = value.args[0].n
-        register_instances += f"        {name} = Register({n}, init={init})\n"
+        if isinstance(value, EscapedExpression):
+            orig = astor.to_source(value.orig.elts[0]).rstrip()
+            register_instances += f"        {name} = {orig}\n"
+            print(register_instances)
+        else:
+            # TODO: Only support m.bits(x, y) for now
+            assert isinstance(value, ast.Call) and \
+                isinstance(value.func, ast.Attribute) and \
+                isinstance(value.func.value, ast.Name) and \
+                value.func.value.id == "m" and \
+                value.func.attr == "bits"
+            assert isinstance(value.args[0], ast.Num)
+            assert isinstance(value.args[1], ast.Num)
+            n = value.args[1].n
+            init = value.args[0].n
+            register_instances += f"        {name} = Register({n}, init={init})\n"
     return register_instances
 
 
@@ -145,12 +152,41 @@ def gen_io_list(inputs, output_type):
     return io_list + "]"
 
 
+class EscapedExpression(ast.AST):
+    def __init__(self, value, orig):
+        self.value = value
+        self.orig = orig
+
+
+class ExecuteEscapedPythonExpressions(ast.NodeTransformer):
+    def __init__(self, defn_env):
+        self.defn_env = defn_env
+
+    def visit_List(self, node):
+        assert len(node.elts) == 1, "Expected only list literatals that " \
+            "contain a single expression"
+        result = eval(astor.to_source(node.elts[0]).rstrip(), self.defn_env)
+        return EscapedExpression(result, node)
+
+
+class SpecializeConstantInts(ast.NodeTransformer):
+    def __init__(self, defn_env):
+        self.defn_env = defn_env
+
+    def visit_Name(self, node):
+        if node.id in self.defn_env:
+            value = self.defn_env[node.id]
+            if isinstance(value, int):
+                return ast.Num(value)
+        return node
+
+
 @ast_utils.inspect_enclosing_env
-def sequential(defn_env : dict, cls):
+def sequential(defn_env: dict, cls):
     if not inspect.isclass(cls):
         raise ValueError("sequential decorator only works with classes")
 
-    initial_value_map = get_initial_value_map(cls.__init__)
+    initial_value_map = get_initial_value_map(cls.__init__, defn_env)
 
     call_def = get_ast(cls.__call__).body[0]
     inputs, output_type = get_io(call_def)

setup.cfg

@@ -0,0 +1,10 @@
+[pycodestyle]
+max-line-length = 80
+
+# E741: do not use variables named ‘l’, ‘O’, or ‘I’
+# We ignore this because I and O is currently idiomatic magma for port names
+ignore = E741
+
+[tool:pytest]
+codestyle_ignore = E741
+codestyle_max_line_length = 80

tests/test_syntax/gold/test_seq_hierarchy.json

@@ -0,0 +1,114 @@
+{"top":"global.ShiftRegister",
+"namespaces":{
+  "global":{
+    "modules":{
+      "Register":{
+        "type":["Record",[
+          ["I",["Array",2,"BitIn"]],
+          ["O",["Array",2,"Bit"]]
+        ]],
+        "instances":{
+          "Register2_inst0":{
+            "modref":"global.Register2"
+          },
+          "Register_comb_inst0":{
+            "modref":"global.Register_comb"
+          }
+        },
+        "connections":[
+          ["Register_comb_inst0.O0","Register2_inst0.I"],
+          ["Register_comb_inst0.self_value_0","Register2_inst0.O"],
+          ["self.I","Register_comb_inst0.I_0"],
+          ["self.O","Register_comb_inst0.O1"]
+        ]
+      },
+      "Register2":{
+        "type":["Record",[
+          ["I",["Array",2,"BitIn"]],
+          ["O",["Array",2,"Bit"]],
+          ["CLK",["Named","coreir.clkIn"]]
+        ]]
+      },
+      "Register2_0001":{
+        "type":["Record",[
+          ["I",["Array",2,"BitIn"]],
+          ["O",["Array",2,"Bit"]],
+          ["CLK",["Named","coreir.clkIn"]]
+        ]]
+      },
+      "Register_comb":{
+        "type":["Record",[
+          ["I_0",["Array",2,"BitIn"]],
+          ["self_value_0",["Array",2,"BitIn"]],
+          ["O0",["Array",2,"Bit"]],
+          ["O1",["Array",2,"Bit"]]
+        ]],
+        "connections":[
+          ["self.O0","self.I_0"],
+          ["self.self_value_0","self.O1"]
+        ]
+      },
+      "Register_unq1":{
+        "type":["Record",[
+          ["I",["Array",2,"BitIn"]],
+          ["O",["Array",2,"Bit"]]
+        ]],
+        "instances":{
+          "Register2_0001_inst0":{
+            "modref":"global.Register2_0001"
+          },
+          "Register_comb_inst0":{
+            "modref":"global.Register_comb"
+          }
+        },
+        "connections":[
+          ["Register_comb_inst0.O0","Register2_0001_inst0.I"],
+          ["Register_comb_inst0.self_value_0","Register2_0001_inst0.O"],
+          ["self.I","Register_comb_inst0.I_0"],
+          ["self.O","Register_comb_inst0.O1"]
+        ]
+      },
+      "ShiftRegister":{
+        "type":["Record",[
+          ["I",["Array",2,"BitIn"]],
+          ["O",["Array",2,"Bit"]]
+        ]],
+        "instances":{
+          "Register_inst0":{
+            "modref":"global.Register"
+          },
+          "Register_inst1":{
+            "modref":"global.Register_unq1"
+          },
+          "ShiftRegister_comb_inst0":{
+            "modref":"global.ShiftRegister_comb"
+          }
+        },
+        "connections":[
+          ["ShiftRegister_comb_inst0.O0","Register_inst0.I"],
+          ["ShiftRegister_comb_inst0.self_x_0","Register_inst0.O"],
+          ["ShiftRegister_comb_inst0.O1","Register_inst1.I"],
+          ["ShiftRegister_comb_inst0.self_y_0","Register_inst1.O"],
+          ["self.I","ShiftRegister_comb_inst0.I_0"],
+          ["self.O","ShiftRegister_comb_inst0.O2"]
+        ]
+      },
+      "ShiftRegister_comb":{
+        "type":["Record",[
+          ["I_0",["Array",2,"BitIn"]],
+          ["self_x_0",["Array",2,"BitIn"]],
+          ["self_y_0",["Array",2,"BitIn"]],
+          ["O0",["Array",2,"Bit"]],
+          ["O1",["Array",2,"Bit"]],
+          ["O2",["Array",2,"Bit"]]
+        ]],
+        "connections":[
+          ["self.O0","self.I_0"],
+          ["self.self_x_0","self.O1"],
+          ["self.self_y_0","self.O2"]
+        ]
+      }
+    }
+  }
+}
+}

tests/test_syntax/test_sequential.py

@@ -60,3 +60,41 @@ def __call__(self, I: m.Bits(2)) -> m.Bits(2):
             return O
 
     compile_and_check("test_seq_simple", Basic, target)
+
+
+def test_seq_hierarchy(target):
+    def DefineCustomRegister(width, init=0):
+        @m.circuit.sequential
+        class Register:
+            def __init__(self):
+                self.value: m.Bits(width) = m.bits(init, width)
+
+            def __call__(self, I: m.Bits(width)) -> m.Bits(width):
+                O = self.value
+                self.value = I
+                return O
+
+        return Register
+
+    @m.circuit.sequential
+    class ShiftRegister:
+        def __init__(self):
+            # [<exp>] means execute the Python expression during the first
+            # stage of compilation
+
+            # The type of a hierarchal output must have the same signature for
+            # the input and output (can have a tuple for multiple inputs). This
+            # reflects the ability to read and write to the variable with the
+            # same meaning
+
+            # Only supported in __init__ body for now
+            self.x: m.Bits(2) = [DefineCustomRegister(2, init=0)()]
+            self.y: m.Bits(2) = [DefineCustomRegister(2, init=1)()]
+
+        def __call__(self, I: m.Bits(2)) -> m.Bits(2):
+            O = self.y
+            self.y = self.x
+            self.x = I
+            return O
+
+    compile_and_check("test_seq_hierarchy", ShiftRegister, target)