add module setup function in pipe stage

[ieee754fpu.git] / src / add / example_buf_pipe.py

diff --git a/src/add/example_buf_pipe.py b/src/add/example_buf_pipe.py
index 856fc7ab697ecf1ebbbb1dd726f6eb528f1a5613..24e83034da4ddc16fb2c8272895dc9e7e0c5e2d5 100644 (file)
--- a/src/add/example_buf_pipe.py
+++ b/src/add/example_buf_pipe.py
@@ -176,6 +176,8 @@ class BufferedPipeline:
 
     def elaborate(self, platform):
         m = Module()
+        if hasattr(self.stage, "setup"):
+            self.stage.setup(m, self.i.data)
 
         # establish some combinatorial temporaries
         o_n_validn = Signal(reset_less=True)
@@ -319,20 +321,27 @@ class CombPipe:
         # set up the input and output data
         self.i.data = stage.ispec() # input type
         self.r_data = stage.ispec() # input type
-        self.o_comb = stage.ospec() # output data
+        self.result = stage.ospec() # output data
         self.o.data = stage.ospec() # output type
         self.o.data.name = "outdata"
 
+    def set_input(self, i):
+        """ helper function to set the input data
+        """
+        return eq(self.i.data, i)
+
     def elaborate(self, platform):
         m = Module()
-        m.d.comb += eq(self.o_comb, self.stage.process(self.r_data))
+        if hasattr(self.stage, "setup"):
+            self.stage.setup(m, self.r_data)
+        m.d.comb += eq(self.result, self.stage.process(self.r_data))
         m.d.comb += self.o.n_valid.eq(self._data_valid)
         m.d.comb += self.o.p_ready.eq(~self._data_valid | self.i.n_ready)
         m.d.sync += self._data_valid.eq(self.i.p_valid | \
                                         (~self.i.n_ready & self._data_valid))
         with m.If(self.i.p_valid & self.o.p_ready):
             m.d.sync += eq(self.r_data, self.i.data)
-        m.d.comb += eq(self.o.data, self.o_comb)
+        m.d.comb += eq(self.o.data, self.result)
         return m
 
     def ports(self):