Integrate Loop Rolling and add BERT example

brainsmith/primitives/transforms/extract_shell_integration_metadata.py

@@ -4,9 +4,12 @@
 """Shell integration metadata extraction transform."""
 
 import json
+import os
+import shutil
+import numpy as np
 from qonnx.transformation.base import Transformation
 import qonnx.custom_op.registry as registry
-
+from finn.util.mlo_sim import dat_file_to_numpy_array
 
 class ExtractShellIntegrationMetadata(Transformation):
     """Walks the ONNX graph and extracts all relevant metadata for shell integration
@@ -19,26 +22,72 @@ def __init__(self, metadata_file: str):
     def apply(self, model):
         graph = model.graph
 
+        # destination dir to copy artifacts
+        dirname = os.path.dirname(self.metadata_file)
+
+        # Search for FINNLoop ops (Does not currently support nested FINNLoops)
+        finn_loops={}
+        mlo = False
+        for node in model.graph.node:
+            if node.op_type == "FINNLoop":
+                finnloop_op = registry.getCustomOp(node)
+                finnloop_body = finnloop_op.get_nodeattr("body");
+
+                mvau_hbm_weights = {}
+                extern_idx = 0
+                for idx, lb_inp in enumerate(finnloop_body.graph.input):
+                    downstream = finnloop_body.find_consumer(lb_inp.name)
+                    if downstream.op_type.startswith("MVAU"):
+                        mlo = True
+                        mvau_hbm_weights[idx] = {}
+                        mvau_hbm_weights[idx]["name"] = lb_inp.name
+                        datfile = (
+                            f"{finnloop_op.get_nodeattr('code_gen_dir_ipgen')}/memblock_MVAU_rtl_id_{idx}.dat"
+                        )
+
+                        # Save the weights as a numpy file
+                        np_dat = dat_file_to_numpy_array(datfile)
+                        mvau_hbm_weights[idx]["weight_npy"] = f"memblock_MVAU_rtl_id_{idx}.npy"
+                        np.save(f"{dirname}/{mvau_hbm_weights[idx]['weight_npy']}", np_dat)
+
+                        # Copy to the destination dir
+                        mvau_hbm_weights[idx]["extern_idx"] = extern_idx
+                        mvau_hbm_weights[idx]["extern_name"] = f"m_axi_MVAU_id_{idx}"
+                        mlo_mvau = registry.getCustomOp(downstream)
+                        mvau_hbm_weights[idx]["PE"] = mlo_mvau.get_nodeattr("PE")
+                        mvau_hbm_weights[idx]["SIMD"] = mlo_mvau.get_nodeattr("SIMD")
+                        mvau_hbm_weights[idx]["MH"] = mlo_mvau.get_nodeattr("MH")
+                        mvau_hbm_weights[idx]["MW"] = mlo_mvau.get_nodeattr("MW")
+                        mvau_hbm_weights[idx]["weightDataType"] = mlo_mvau.get_nodeattr("weightDataType")
+                        extern_idx += 1
+                finn_loops[node.name] = mvau_hbm_weights
+        self.md["mlo"] = mlo
+        self.md["finn_loops"] = finn_loops
+
+
+        # Extract instream widths
         instreams = {}
         for input_tensor in graph.input:
             consumer = model.find_consumer(input_tensor.name)
             inst = registry.getCustomOp(consumer)
-            instreams[input_tensor.name] = {
-                'width': inst.get_instream_width(),
-                'shape': inst.get_normal_input_shape()
-            }
-        self.md['instreams'] = instreams
+            instream = {}
+            instream['width'] = inst.get_instream_width()
+            instreams[input_tensor.name] = instream
+            instream['shape'] = inst.get_normal_input_shape()
+            instream['datatype'] = inst.get_input_datatype().name
+        self.md['insteams'] = instreams
 
         outstreams = {}
         for output_tensor in graph.output:
             producer = model.find_producer(output_tensor.name)
             inst = registry.getCustomOp(producer)
-            outstreams[output_tensor.name] = {
-                'width': inst.get_outstream_width(),
-                'shape': inst.get_normal_output_shape()
-            }
-        self.md['outstreams'] = outstreams
-
+            outstream = {}
+            outstream['width'] = inst.get_outstream_width()
+            outstreams[output_tensor.name] = outstream
+            outstream['shape'] = inst.get_normal_output_shape()
+            outstream['datatype'] = inst.get_output_datatype().name
+        self.md['outsteams'] = outstreams
+
         static_matmuls = {}
         for node in graph.node:
             if (node.op_type == "MVAU_rtl"):
@@ -54,4 +103,4 @@ def apply(self, model):
         with open(self.metadata_file, "w") as fp:
             json.dump(self.md, fp, indent=4)
 
-        return (model, False)
+        return(model, False)

brainsmith/steps/bert_custom_steps.py

@@ -89,7 +89,8 @@ def bert_streamlining_step(model: Any, cfg: Any) -> Any:
         MoveScalarMulPastMatMul(),
         MoveScalarLinearPastInvariants(),
         AbsorbMulIntoMultiThreshold(),
-        AbsorbAddIntoMultiThreshold()
+        AbsorbAddIntoMultiThreshold(),
+        RoundAndClipThresholds()
     ]:
         model = model.transform(transform)
 

docker/fetch-repos.sh

@@ -77,9 +77,21 @@ fi
 
 # Define our Git dependencies - URLs and revisions
 declare -A GIT_DEPS=(
+<<<<<<< HEAD
+    ["brevitas"]="https://github.com/Xilinx/brevitas.git@95edaa0bdc8e639e39b1164466278c59df4877be"
+    ["qonnx"]="https://github.com/fastmachinelearning/qonnx.git@custom/brainsmith"
+    ["finn"]="https://github.com/tafk7/finn.git@custom/transformer"
+    ["onnxscript"]="https://github.com/jsmonson/onnxscript.git@62c7110aba46554432ce8e82ba2d8a086bd6227c"
+||||||| c35477d
+    ["brevitas"]="https://github.com/Xilinx/brevitas.git@95edaa0bdc8e639e39b1164466278c59df4877be"
+    ["qonnx"]="https://github.com/fastmachinelearning/qonnx.git@f2c4ccd3e71795c9f116ee5a0c87a7dfd590c6d0"
+    ["finn"]="https://github.com/tafk7/finn.git@custom/transformer"
+    ["onnxscript"]="https://github.com/jsmonson/onnxscript.git@62c7110aba46554432ce8e82ba2d8a086bd6227c"
+=======
     ["brevitas"]="https://github.com/Xilinx/brevitas.git@c10ef8764967e9cacc60347ce185be14e4ad97c4"
     ["qonnx"]="https://github.com/fastmachinelearning/qonnx.git@f2c4ccd3e71795c9f116ee5a0c87a7dfd590c6d0"
     ["finn"]="https://github.com/tafk7/finn.git@feature/logging-integration-transformer"
+>>>>>>> develop
     ["finn-experimental"]="https://github.com/Xilinx/finn-experimental.git@0724be21111a21f0d81a072fccc1c446e053f851"
     ["dataset-loading"]="https://github.com/fbcotter/dataset_loading.git@0.0.4"
 )