From 12d05c6294599a4944d4b4d188cae191593cbae2 Mon Sep 17 00:00:00 2001
From: Luke Kenneth Casson Leighton <lkcl@lkcl.net>
Date: Sat, 30 Mar 2019 23:04:23 +0000
Subject: [PATCH] add pipeline context / stage management

---
 src/add/pipeline.py         | 65 ++++++++++++++++++++++++++++++++++++-
 src/add/pipeline_example.py | 58 +++++++++++++++++++++++++++------
 2 files changed, 112 insertions(+), 11 deletions(-)

diff --git a/src/add/pipeline.py b/src/add/pipeline.py
index fc5faa91..56e8e33c 100644
--- a/src/add/pipeline.py
+++ b/src/add/pipeline.py
@@ -4,6 +4,7 @@ from nmigen import Signal
 from nmigen.hdl.rec import Record
 from nmigen import tracer
 from nmigen.compat.fhdl.bitcontainer import value_bits_sign
+from contextlib import contextmanager
 
 from singlepipe import eq
 
@@ -69,7 +70,69 @@ class ObjectProxy:
         self._pipe.sync += eq(new_pipereg, value)
 
 
-class SimplePipeline(object):
+class PipelineStage:
+    """ Pipeline builder stage with auto generation of pipeline registers.
+    """
+
+    def __init__(self, pipe, prev=None):
+        self._pipe = pipe
+        self._preg_map = {}
+        self._prev_stage = prev
+        if prev:
+            print ("prev", prev._preg_map)
+            if prev._current_stage_num in prev._preg_map:
+                m = prev._preg_map[prev._current_stage_num]
+                self._preg_map[prev._current_stage_num] = m
+            self._current_stage_num = prev._current_stage_num + 1
+            if self._current_stage_num in prev._preg_map:
+                m = prev._preg_map[self._current_stage_num]
+                self._preg_map[self._current_stage_num] = m
+                print ("make current", m)
+        else:
+            self._current_stage_num = 0
+
+    def __getattr__(self, name):
+        try:
+            return self._preg_map[self._current_stage_num][name]
+        except KeyError:
+            raise AttributeError(
+                'error, no pipeline register "%s" defined for stage %d'
+                % (name, self._current_stage_num))
+
+    def __setattr__(self, name, value):
+        if name.startswith('_'):
+            # do not do anything tricky with variables starting with '_'
+            object.__setattr__(self, name, value)
+            return
+        next_stage = self._current_stage_num + 1
+        pipereg_id = str(self._current_stage_num) + 'to' + str(next_stage)
+        rname = 'pipereg_' + pipereg_id + '_' + name
+        #new_pipereg = Signal(value_bits_sign(value), name=rname,
+        #                     reset_less=True)
+        if isinstance(value, ObjectProxy):
+            new_pipereg = ObjectProxy.like(self._pipe, value,
+                                           name=rname, reset_less = True)
+        else:
+            new_pipereg = Signal.like(value, name=rname, reset_less = True)
+        if next_stage not in self._preg_map:
+            self._preg_map[next_stage] = {}
+        self._preg_map[next_stage][name] = new_pipereg
+        self._pipe.sync += eq(new_pipereg, value)
+
+
+class PipeManager:
+    def __init__(self, pipe):
+        self._pipe = pipe
+
+    @contextmanager
+    def Stage(self, prev=None):
+        stage = PipelineStage(self._pipe, prev)
+        try:
+            yield stage
+        finally:
+            pass
+
+class SimplePipeline:
     """ Pipeline builder with auto generation of pipeline registers.
     """
 
diff --git a/src/add/pipeline_example.py b/src/add/pipeline_example.py
index 6474a396..36295265 100644
--- a/src/add/pipeline_example.py
+++ b/src/add/pipeline_example.py
@@ -1,10 +1,10 @@
 """ Example 5: Making use of PyRTL and Introspection. """
 
 from nmigen import Module, Signal
-from nmigen.cli import main, verilog
+from nmigen.cli import main, verilog, rtlil
 
 
-from pipeline import SimplePipeline, ObjectProxy
+from pipeline import SimplePipeline, ObjectProxy, PipeManager
 
 
 class SimplePipelineExample(SimplePipeline):
@@ -37,7 +37,7 @@ class ObjectBasedPipelineExample(SimplePipeline):
     """ A very simple pipeline to show how registers are inferred. """
 
     def __init__(self, pipe):
-        ObjectBasedPipeline.__init__(self, pipe)
+        SimplePipeline.__init__(self, pipe)
         self._loopback = Signal(4)
         o = ObjectProxy(pipe)
         o.a = Signal(4)
@@ -82,12 +82,50 @@ class PipeModule:
     def get_fragment(self, platform=None):
         return self.m
 
+
+class PipelineStageExample(PipeManager):
+
+    def __init__(self):
+        self.m = Module()
+        self._loopback = Signal(4)
+        PipeManager.__init__(self, self.m.d)
+
+    def stage0(self):
+        self.n = ~self._loopback
+
+    def stage1(self):
+        self.n = self.n + 2
+
+    def stage2(self):
+        localv = Signal(4)
+        self._pipe.comb += localv.eq(2)
+        self.n = self.n << localv
+
+    def stage3(self):
+        self.n = ~self.n
+
+    def stage4(self):
+        self._pipe.sync += self._loopback.eq(self.n + 3)
+
+    def get_fragment(self, platform=None):
+
+        with self.Stage() as p:
+            p.n = ~self._loopback
+        with self.Stage(p) as p:
+            p.n = p.n + 2
+
+        return self.m
+
+
+
 if __name__ == "__main__":
     example = PipeModule()
-    main(example, ports=[
-                    example.p._loopback,
-        ])
-
-    #print(verilog.convert(example, ports=[
-    #           example.p._loopback,
-    #         ]))
+    with open("pipe_module.il", "w") as f:
+        f.write(rtlil.convert(example, ports=[
+               example.p._loopback,
+             ]))
+    example = PipelineStageExample()
+    with open("pipe_stage_module.il", "w") as f:
+        f.write(rtlil.convert(example, ports=[
+               example._loopback,
+             ]))
-- 
2.30.2