Support Trunc v2

fetch-repos.sh

@@ -28,9 +28,9 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-QONNX_COMMIT="0630ceaee17799096d1750abcfb5bbe0a2877888"
+QONNX_COMMIT="80b60068b021698f4425d4fb4119bd9156efd1fd"
 FINN_EXP_COMMIT="0724be21111a21f0d81a072fccc1c446e053f851"
-BREVITAS_COMMIT="4617f7bd136e96fa21c7f76e3c7e2e37fe563837"
+BREVITAS_COMMIT="b106358c4169d8a9b68cb2a531aa795417d74887"
 CNPY_COMMIT="8c82362372ce600bbd1cf11d64661ab69d38d7de"
 HLSLIB_COMMIT="a19482ba6886f6f26aff11b10126a82ce0dd7ab1"
 OMX_COMMIT="a5d48f93309b235fdd21556d16e86e6ef5db6e2e"

src/finn/custom_op/fpgadataflow/__init__.py

@@ -41,6 +41,8 @@ def register_custom_op(cls):
     assert issubclass(cls, HWCustomOp), f"{cls} must subclass {HWCustomOp}"
     # Insert the class into the custom_op dictionary by its name
     custom_op[cls.__name__] = cls
+    # Versioning
+    custom_op[cls.__name__ + "_v1"] = cls
     # Pass through the class unmodified
     return cls
 
@@ -108,3 +110,29 @@ def register_custom_op(cls):
 custom_op["StreamingDataWidthConverter"] = StreamingDataWidthConverter
 custom_op["StreamingEltwise"] = StreamingEltwise
 custom_op["UpsampleNearestNeighbour"] = UpsampleNearestNeighbour
+
+# Versioning
+custom_op["MVAU_v1"] = MVAU
+custom_op["StreamingFIFO_v1"] = StreamingFIFO
+custom_op["Thresholding_v1"] = Thresholding
+custom_op["VVAU_v1"] = VVAU
+custom_op["StreamingDataflowPartition_v1"] = StreamingDataflowPartition
+
+custom_op["AddStreams_v1"] = AddStreams
+custom_op["ChannelwiseOp_v1"] = ChannelwiseOp
+custom_op["ConvolutionInputGenerator_v1"] = ConvolutionInputGenerator
+custom_op["DuplicateStreams_v1"] = DuplicateStreams
+custom_op["FMPadding_v1"] = FMPadding
+custom_op["FMPadding_Pixel_v1"] = FMPadding_Pixel
+custom_op["GlobalAccPool_v1"] = GlobalAccPool
+custom_op["InnerShuffle_v1"] = InnerShuffle
+custom_op["LabelSelect_v1"] = LabelSelect
+custom_op["Lookup_v1"] = Lookup
+custom_op["OuterShuffle_v1"] = OuterShuffle
+custom_op["Pool_v1"] = Pool
+custom_op["Shuffle_v1"] = Shuffle
+custom_op["StreamingConcat_v1"] = StreamingConcat
+custom_op["StreamingSplit_v1"] = StreamingSplit
+custom_op["StreamingDataWidthConverter_v1"] = StreamingDataWidthConverter
+custom_op["StreamingEltwise_v1"] = StreamingEltwise
+custom_op["UpsampleNearestNeighbour_v1"] = UpsampleNearestNeighbour

src/finn/custom_op/fpgadataflow/channelwise_op.py

@@ -235,10 +235,7 @@ def execute_node(self, context, graph):
             outputs=[outp],
         )
 
-        opset_version = self.onnx_opset_version
-        opset_imports = [helper.make_opsetid("", opset_version)]
-        onnx_kwargs = {"opset_imports": opset_imports}
-        model_func = qonnx_make_model(graph_func, **onnx_kwargs)
+        model_func = qonnx_make_model(graph_func)
         idict = {node.input[0]: inp_values, node.input[1]: param_values}
         sess = rt.InferenceSession(model_func.SerializeToString())
         result = sess.run(None, idict)

src/finn/custom_op/fpgadataflow/convolutioninputgenerator.py

@@ -251,8 +251,7 @@ def execute_node(self, context, graph):
             outputs=[outp],
         )
 
-        opset_version = self.onnx_opset_version
-        opset_imports = [helper.make_opsetid("", opset_version)]
+        opset_imports = [helper.make_opsetid("qonnx.custom_op.general", 1)]
         onnx_kwargs = {"opset_imports": opset_imports}
         model_im2col = ModelWrapper(qonnx_make_model(graph_im2col, **onnx_kwargs))
         model_im2col.set_tensor_datatype(node.input[0], self.get_input_datatype())

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

@@ -43,6 +43,8 @@ def register_custom_op(cls):
     assert issubclass(cls, HLSBackend), f"{cls} must subclass {HLSBackend}"
     # Insert the class into the custom_op dictionary by its name
     custom_op[cls.__name__] = cls
+    # Versioning
+    custom_op[cls.__name__ + "_v1"] = cls
     # Pass through the class unmodified
     return cls
 
@@ -98,3 +100,25 @@ def register_custom_op(cls):
 custom_op["MVAU_hls"] = MVAU_hls
 custom_op["VVAU_hls"] = VVAU_hls
 custom_op["OuterShuffle_hls"] = OuterShuffle_hls
+
+# Versioning
+custom_op["AddStreams_hls_v1"] = AddStreams_hls
+custom_op["ChannelwiseOp_hls_v1"] = ChannelwiseOp_hls
+custom_op["CheckSum_hls_v1"] = CheckSum_hls
+custom_op["DuplicateStreams_hls_v1"] = DuplicateStreams_hls
+custom_op["FMPadding_Pixel_hls_v1"] = FMPadding_Pixel_hls
+custom_op["GlobalAccPool_hls_v1"] = GlobalAccPool_hls
+custom_op["IODMA_hls_v1"] = IODMA_hls
+custom_op["LabelSelect_hls_v1"] = LabelSelect_hls
+custom_op["Lookup_hls_v1"] = Lookup_hls
+custom_op["Pool_hls_v1"] = Pool_hls
+custom_op["StreamingConcat_hls_v1"] = StreamingConcat_hls
+custom_op["StreamingSplit_hls_v1"] = StreamingSplit_hls
+custom_op["StreamingEltwise_hls_v1"] = StreamingEltwise_hls
+custom_op["StreamingDataWidthConverter_hls_v1"] = StreamingDataWidthConverter_hls
+custom_op["Thresholding_hls_v1"] = Thresholding_hls
+custom_op["TLastMarker_hls_v1"] = TLastMarker_hls
+custom_op["UpsampleNearestNeighbour_hls_v1"] = UpsampleNearestNeighbour_hls
+custom_op["MVAU_hls_v1"] = MVAU_hls
+custom_op["VVAU_hls_v1"] = VVAU_hls
+custom_op["OuterShuffle_hls_v1"] = OuterShuffle_hls

src/finn/custom_op/fpgadataflow/labelselect.py

@@ -152,10 +152,7 @@ def execute_node(self, context, graph):
             outputs=[val_outp, outp],
         )
 
-        opset_version = self.onnx_opset_version
-        opset_imports = [helper.make_opsetid("", opset_version)]
-        onnx_kwargs = {"opset_imports": opset_imports}
-        model_topk = qonnx_make_model(graph_topk, **onnx_kwargs)
+        model_topk = qonnx_make_model(graph_topk)
         idict = {node.input[0]: inp_values, "k_inp": [k]}
         sess = rt.InferenceSession(model_topk.SerializeToString())
         result = sess.run(None, idict)

src/finn/custom_op/fpgadataflow/lookup.py

@@ -180,8 +180,7 @@ def execute_node(self, context, graph):
             outputs=[outp],
         )
 
-        opset_version = 13
-        opset_imports = [helper.make_opsetid("", opset_version)]
+        opset_imports = [helper.make_opsetid("", 13)]
         onnx_kwargs = {"opset_imports": opset_imports}
         model_gather = qonnx_make_model(graph_gather, **onnx_kwargs)
         idict = {node.input[0]: inp_values, node.input[1]: data_values}

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

@@ -51,3 +51,13 @@
 custom_op["VVAU_rtl"] = VVAU_rtl
 custom_op["Thresholding_rtl"] = Thresholding_rtl
 custom_op["InnerShuffle_rtl"] = InnerShuffle_rtl
+
+# Versioning
+custom_op["ConvolutionInputGenerator_rtl_v1"] = ConvolutionInputGenerator_rtl
+custom_op["FMPadding_rtl_v1"] = FMPadding_rtl
+custom_op["StreamingDataWidthConverter_rtl_v1"] = StreamingDataWidthConverter_rtl
+custom_op["StreamingFIFO_rtl_v1"] = StreamingFIFO_rtl
+custom_op["MVAU_rtl_v1"] = MVAU_rtl
+custom_op["VVAU_rtl_v1"] = VVAU_rtl
+custom_op["Thresholding_rtl_v1"] = Thresholding_rtl
+custom_op["InnerShuffle_rtl_v1"] = InnerShuffle_rtl

src/finn/custom_op/fpgadataflow/upsampler.py

@@ -152,8 +152,7 @@ def execute_node(self, context, graph):
             outputs=[outp],
         )
 
-        opset_version = 13
-        opset_imports = [helper.make_opsetid("", opset_version)]
+        opset_imports = [helper.make_opsetid("", 13)]
         onnx_kwargs = {"opset_imports": opset_imports}
         model_resize = qonnx_make_model(graph_resize, **onnx_kwargs)
         idict = {node.input[0]: inp_values, "scales": scales_val}

src/finn/transformation/qonnx/fold_quant_weights.py

@@ -34,6 +34,7 @@
 from qonnx.transformation.infer_shapes import InferShapes
 from qonnx.transformation.quant_constant_folding import FoldTransposeIntoQuantInit
 from qonnx.transformation.remove import remove_node_and_rewire
+from qonnx.util.basic import get_preferred_qonnx_opset
 
 
 class FoldQuantWeights(Transformation):
@@ -46,6 +47,7 @@ def apply(self, model):
         node_ind = 0
         graph_modified = False
         execution_context = model.make_empty_exec_context()
+        opset_imports = model.get_opset_imports()
         for n in graph.node:
             node_ind += 1
             if n.op_type == "Quant" or n.op_type == "BipolarQuant":
@@ -89,7 +91,11 @@ def apply(self, model):
                     unity_scale = (scale.flatten() == 1.0).all()
                     # this node has no dynamic inputs, only constant ones -- so we can
                     # do constant folding.
-                    oxe.execute_node(n, execution_context, graph)
+                    if n.domain in opset_imports:
+                        opset_version = opset_imports[n.domain]
+                    else:
+                        opset_version = get_preferred_qonnx_opset()
+                    oxe.execute_node(n, execution_context, graph, opset_version)
                     q_node_output = execution_context[node_out]
                     # Check we can directly constant fold
                     if unity_scale:

src/finn/transformation/qonnx/infer_quant_avg_pool_2d.py

@@ -28,6 +28,7 @@
 
 
 import math
+import numpy as np
 from onnx import TensorProto, helper
 from qonnx.core.datatype import DataType
 from qonnx.custom_op.registry import getCustomOp
@@ -117,6 +118,34 @@ class AvgPoolAndTruncToQuantAvgPool(Transformation):
     To the FINN op: QuantAvgPool2d
     """
 
+    def apply(self, model):
+        opset_imports = model.get_opset_imports()
+        if "qonnx.custom_op.general" in opset_imports:
+            trunc_opset = opset_imports["qonnx.custom_op.general"]
+        elif "onnx.brevitas" in opset_imports:
+            trunc_opset = opset_imports["onnx.brevitas"]
+        else:
+            trunc_opset = 1  # Default to v1 if no opset found
+        if trunc_opset == 1:
+            model = model.transform(AvgPoolAndTruncv1ToQuantAvgPool())
+            return model, False
+        elif trunc_opset == 2:
+            model = model.transform(AvgPoolAndTruncv2ToQuantAvgPool())
+            return model, False
+        else:
+            raise NotImplementedError(
+                f"AvgPoolAndTruncToQuantAvgPool not implemented for "
+                f"Trunc opset version {trunc_opset}."
+            )
+
+
+class AvgPoolAndTruncv1ToQuantAvgPool(Transformation):
+    """
+    Convert a section of nodes of the pattern:
+    AveragePool -> Mul (scalar) -> Trunc (v1)
+    To the FINN op: Div -> QuantAvgPool2d -> Mul
+    """
+
     def apply(self, model):
         graph = model.graph
         node_ind = 0
@@ -282,3 +311,171 @@ def apply(self, model):
                         return model, True
 
         return model, False
+
+
+class AvgPoolAndTruncv2ToQuantAvgPool(Transformation):
+    """
+    Convert a section of nodes of the pattern:
+    AveragePool -> Trunc (v2)
+    To the FINN op: Div -> QuantAvgPool2d -> Mul
+    """
+
+    def apply(self, model):
+        graph = model.graph
+        node_ind = 0
+        for node in graph.node:
+            node_ind += 1
+            if node.op_type == "AveragePool":
+                t_node = model.find_direct_successors(node)
+                if t_node is not None and len(t_node) == 1 and t_node[0].op_type == "Trunc":
+                    t_node = t_node[0]
+                    running_node_index = node_ind
+                    # Check node for compatibility
+                    # Avg pooling node
+                    k_s = get_by_name(node.attribute, "kernel_shape")
+                    if k_s is None or len(k_s.ints) != 2 or len(set(k_s.ints)) != 1:
+                        raise ValueError(
+                            "FINN only supports average pooling with " "2D square kernels."
+                        )
+                    k_s = k_s.ints[0]
+
+                    pads = get_by_name(node.attribute, "pads")
+                    if pads is None or len(set(pads.ints)) != 1 or pads.ints[0] != 0:
+                        raise ValueError("FINN dosn't support padding for average pooling.")
+
+                    stride = get_by_name(node.attribute, "strides")
+                    if stride is None or len(stride.ints) != 2 or len(set(stride.ints)) != 1:
+                        raise ValueError(
+                            "FINN only supports 2D strides with equal values in " "each direction."
+                        )
+                    stride = stride.ints[0]
+
+                    # Trunc node
+                    rounding_mode = get_by_name(t_node.attribute, "rounding_mode")
+                    normalized_mode_string = rounding_mode.s.upper()
+                    if rounding_mode is None or normalized_mode_string != b"FLOOR":
+                        raise ValueError(
+                            "The Trunc node must have the rounding_mode " "set to 'FLOOR'."
+                        )
+                    for inp in t_node.input[1:]:
+                        if model.get_initializer(inp) is None:
+                            raise ValueError(
+                                f"All inputs of the Trunc node, "
+                                f"except the first, must be statically "
+                                f"initialized. However, {inp} is not."
+                            )
+                    zero_pt = model.get_initializer(t_node.input[2])
+                    if len(zero_pt.shape) != 0 or zero_pt != 0:
+                        raise ValueError(
+                            f"Finn only supports 0 as the zero point for "
+                            f"the Trunc node, it currently is {zero_pt}."
+                        )
+                    scale = model.get_initializer(t_node.input[1]).flatten()
+                    out_scale = model.get_initializer(t_node.input[4]).flatten()
+
+                    trunc_in_bits = model.get_initializer(t_node.input[3]).flatten()
+                    trunc_out_bits = model.get_initializer(t_node.input[5]).flatten()
+                    if len(trunc_in_bits.shape) != 1 or len(trunc_out_bits.shape) != 1:
+                        raise ValueError(
+                            f"Finn only supports scalar bit widths "
+                            f"for the Trunc node. The input bit width "
+                            f"currently is: {trunc_in_bits}, "
+                            f"while the output bit width is: {trunc_out_bits}."
+                        )
+                    trunc_in_bits = int(trunc_in_bits[0])
+                    trunc_out_bits = int(trunc_out_bits[0])
+                    if np.round(np.log2(out_scale / scale) != trunc_in_bits - trunc_out_bits):
+                        raise ValueError(
+                            f"The scale values for the Trunc node are not "
+                            f"compatible with the specified bit widths. "
+                            f"Input scale: {scale}, output scale: {out_scale}, "
+                            f"input bits: {trunc_in_bits}, output bits: {trunc_out_bits}."
+                        )
+
+                    # Calculate parameters for the QuantAvgPool2d node,
+                    # Calculate input bit width. Basically this backwards:
+                    # https://github.com/Xilinx/finn-base/blob/
+                    # 7c2603a95e90e4de2575020e575c24eab6a15889/src/finn/custom_op/
+                    # general/quantavgpool2d.py#L94
+                    ibits = math.floor(math.log(2**trunc_in_bits / (k_s * k_s), 2))
+                    # Get sign
+                    signed = _get_signed_from_upstream(model, t_node)
+                    # ToDo: Change this to NHWC,
+                    #  when the channels last layout comes around.
+                    data_layout = "NCHW"
+
+                    # Insert scale nodes, QuantAvgPool2d node and required tensors
+                    scale = model.get_initializer(t_node.input[1])
+                    # for Trunc v2 update input scale by receptive field
+                    scale = (scale * k_s * k_s).astype(scale.dtype)
+                    scale_div_tensor = helper.make_tensor_value_info(
+                        model.make_new_valueinfo_name(),
+                        TensorProto.FLOAT,
+                        None,
+                    )
+                    graph.value_info.append(scale_div_tensor)
+                    model.set_initializer(scale_div_tensor.name, scale)
+
+                    act_scale_div_tensor = helper.make_tensor_value_info(
+                        model.make_new_valueinfo_name(),
+                        TensorProto.FLOAT,
+                        None,
+                    )
+                    graph.value_info.append(act_scale_div_tensor)
+
+                    scale_div_node = helper.make_node(
+                        "Div",
+                        [node.input[0], scale_div_tensor.name],
+                        [act_scale_div_tensor.name],
+                    )
+                    graph.node.insert(running_node_index, scale_div_node)
+                    running_node_index += 1
+
+                    act_scale_mul_tensor = helper.make_tensor_value_info(
+                        model.make_new_valueinfo_name(),
+                        TensorProto.FLOAT,
+                        None,
+                    )
+                    graph.value_info.append(act_scale_mul_tensor)
+                    QuantAvgPool2d_node = helper.make_node(
+                        "QuantAvgPool2d",
+                        [act_scale_div_tensor.name],
+                        [act_scale_mul_tensor.name],
+                        domain="qonnx.custom_op.general",
+                        stride=stride,
+                        kernel=k_s,
+                        ibits=ibits,
+                        obits=trunc_out_bits,
+                        signed=int(signed),
+                        data_layout=data_layout,
+                    )
+                    graph.node.insert(running_node_index, QuantAvgPool2d_node)
+                    running_node_index += 1
+
+                    scale_mul_tensor = helper.make_tensor_value_info(
+                        model.make_new_valueinfo_name(),
+                        TensorProto.FLOAT,
+                        None,
+                    )
+                    graph.value_info.append(scale_mul_tensor)
+                    model.set_initializer(scale_mul_tensor.name, out_scale)
+
+                    scale_mul_node = helper.make_node(
+                        "Mul",
+                        [act_scale_mul_tensor.name, scale_mul_tensor.name],
+                        [t_node.output[0]],
+                    )
+                    graph.node.insert(running_node_index, scale_mul_node)
+                    running_node_index += 1
+
+                    # Remove old nodes
+                    graph.node.remove(node)
+                    graph.node.remove(t_node)
+
+                    # Recompute shapes and datatypes
+                    model = model.transform(InferShapes())
+                    model = model.transform(InferDataTypes())
+
+                    return model, True
+
+        return model, False