[CustomOp] Update derive TAV method to allow for pre-hook for rtlsim

Author: auphelia

setup.cfg

@@ -132,6 +132,7 @@ markers =
     bnn_kv260: mark tests that execute KV260 BNN tests
     bnn_pynq: mark tests that execute Pynq-Z1 BNN tests
     bnn_zcu104: mark tests that execute ZCU104 BNN tests
+    node_tree_modeling: mark tests for analytical FIFO sizing tree models
 norecursedirs =
     dist
     build

src/finn/custom_op/fpgadataflow/hls/thresholding_hls.py

@@ -769,7 +769,17 @@ def ipgen_extra_directives(self):
 
         return ["config_compile -pipeline_style frp"]
 
-    def derive_characteristic_fxns(self, period):
+    def derive_token_access_vectors(
+        self,
+        model,
+        period,
+        strategy,
+        fpga_part,
+        clk_period,
+        op_type,
+        override_dict=None,
+        pre_hook=None,
+    ):
         n_inps = np.prod(self.get_folded_input_shape()[:-1])
         io_dict = {
             "inputs": {
@@ -782,7 +792,9 @@ def derive_characteristic_fxns(self, period):
             n_weight_inps = self.calc_tmem()
             num_w_reps = np.prod(self.get_nodeattr("numInputVectors"))
             io_dict["inputs"]["in1"] = [0 for i in range(num_w_reps * n_weight_inps)]
-        super().derive_characteristic_fxns(period, override_rtlsim_dict=io_dict)
+        super().derive_token_access_vectors(
+            model, period, strategy, fpga_part, clk_period, op_type, io_dict, pre_hook=pre_hook
+        )
 
     def minimize_weight_bit_width(self, model):
         """Minimize threshold datatype, with HLS-specific adjustments.

src/finn/custom_op/fpgadataflow/hwcustomop.py

@@ -319,7 +319,15 @@ def get_tree_model(self):
         return None
 
     def derive_token_access_vectors(
-        self, model, period, strategy, fpga_part, clk_period, op_type, override_dict=None
+        self,
+        model,
+        period,
+        strategy,
+        fpga_part,
+        clk_period,
+        op_type,
+        override_dict=None,
+        pre_hook=None,
     ):
         if override_dict is None:
             n_inps = np.prod(self.get_folded_input_shape()[:-1])
@@ -342,7 +350,9 @@ def derive_token_access_vectors(
         # there is a 20 clock marging added for when get_exp_cycles()
         # is underestimating the real operator runtime.
         period = self.get_exp_cycles() + 20
-        self.derive_token_access_vectors_using_rtlsim(model, period, fpga_part, clk_period, io_dict)
+        self.derive_token_access_vectors_using_rtlsim(
+            model, period, fpga_part, clk_period, io_dict, pre_hook=pre_hook
+        )
 
     def derive_token_access_vectors_using_tree_model(self, period, io_dict):
         # Analytical flow
@@ -469,9 +479,9 @@ def apply_micro_buffer_correction(start, txn_in, period):
 
     def generate_hdl_memstream(self, fpgapart, pumped_memory=0):
         """Helper function to generate verilog code for memstream component.
-        Currently utilized by MVAU, VVAU and HLS Thresholding layer."""
+        Currently utilized by MVAU, VVAU, HLS Thresholding and Elementwise layers."""
         ops = ["MVAU_hls", "MVAU_rtl", "VVAU_hls", "VVAU_rtl", "Thresholding_hls"]
-        if self.onnx_node.op_type in ops:
+        if self.onnx_node.op_type in ops or self.onnx_node.op_type.startswith("Elementwise"):
             template_path = (
                 os.environ["FINN_ROOT"] + "/finn-rtllib/memstream/hdl/memstream_wrapper_template.v"
             )
@@ -601,10 +611,16 @@ def generate_hdl_dynload(self):
             f.write(template_wrapper)
 
     def derive_token_access_vectors_using_rtlsim(
-        self, model, period, fpga_part, clk_period, override_rtlsim_dict=None
+        self, model, period, fpga_part, clk_period, override_rtlsim_dict=None, pre_hook=None
     ):
         """Return the token access vectors for this node using rtlsim.
-        Used by analytical FIFO sizing approach."""
+        Used by analytical FIFO sizing approach.
+
+        Args:
+            pre_hook: Optional callable that takes sim as argument, called after
+                      reset_rtlsim but before running the simulation. Used by
+                      FINNLoop to initialize MLO state.
+        """
         # ensure rtlsim is ready
 
         periods_to_simulate = 5
@@ -650,6 +666,8 @@ def derive_token_access_vectors_using_rtlsim(
         # signal name, note no underscore at the end (new finnxsi behavior)
         sname = "_V"
         self.reset_rtlsim(sim)
+        if pre_hook is not None:
+            pre_hook(sim)
 
         # create stream tracers for all input and output streams
         for k in txns_in.keys():

src/finn/custom_op/fpgadataflow/rtl/elementwise_binary_rtl.py

@@ -380,7 +380,17 @@ def instantiate_ip(self, cmd):
             "create_bd_cell -type hier -reference %s /%s/%s" % (top_module, node_name, node_name)
         )
 
-    def derive_characteristic_fxns(self, period, override_rtlsim_dict=None, pre_hook=None):
+    def derive_token_access_vectors(
+        self,
+        model,
+        period,
+        strategy,
+        fpga_part,
+        clk_period,
+        op_type,
+        override_dict=None,
+        pre_hook=None,
+    ):
         n_inps = np.prod(self.get_folded_input_shape(0)[:-1])
         io_dict = {
             "inputs": {
@@ -389,7 +399,9 @@ def derive_characteristic_fxns(self, period, override_rtlsim_dict=None, pre_hook
             },
             "outputs": {"out0": []},
         }
-        super().derive_characteristic_fxns(period, override_rtlsim_dict=io_dict, pre_hook=pre_hook)
+        super().derive_token_access_vectors(
+            model, period, strategy, fpga_part, clk_period, op_type, io_dict, pre_hook=pre_hook
+        )
 
     def execute_node(self, context, graph):
         mode = self.get_nodeattr("exec_mode")

src/finn/custom_op/fpgadataflow/rtl/finn_loop.py

@@ -293,9 +293,29 @@ def prepare_rtlsim(self, behav=False):
         sim_base, sim_rel = rtlsim_so
         self.set_nodeattr("rtlsim_so", sim_base + "/" + sim_rel)
 
-    def derive_characteristic_fxns(self, period):
+    def derive_token_access_vectors(
+        self,
+        model,
+        period,
+        strategy,
+        fpga_part,
+        clk_period,
+        op_type,
+        override_dict=None,
+        pre_hook=None,
+    ):
+        # FINNLoop always uses rtlsim strategy with MLO prehook
         mlo_prehook = mlo_prehook_func_factory(self.onnx_node)
-        super().derive_characteristic_fxns(period, pre_hook=mlo_prehook)
+        super().derive_token_access_vectors(
+            model,
+            period,
+            "rtlsim",
+            fpga_part,
+            clk_period,
+            op_type,
+            override_dict,
+            pre_hook=mlo_prehook,
+        )
 
     def execute_node(self, context, graph):
         node = self.onnx_node