Pull upstream changes into qonnx_quant_op

README.md

@@ -1,4 +1,4 @@
-## <img src=https://raw.githubusercontent.com/Xilinx/finn/master/docs/img/finn-logo.png width=128/> Core Components for Quantized Neural Network Inference
+## <img src=https://raw.githubusercontent.com/Xilinx/finn/github-pages/docs/img/finn-logo.png width=128/> Core Components for Quantized Neural Network Inference
 
 [![Gitter](https://badges.gitter.im/xilinx-finn/community.svg)](https://gitter.im/xilinx-finn/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge)
 [![ReadTheDocs](https://readthedocs.org/projects/finn-base/badge/?version=latest&style=plastic)](http://finn-base.readthedocs.io/)

docker/Dockerfile

@@ -73,7 +73,7 @@ RUN python -mpip install --upgrade pip && \
     rm requirements.txt
 
 # Install custom fork of pyverilator
-RUN pip install git+https://github.com/maltanar/pyverilator.git#egg=pyverilator
+RUN pip install git+https://github.com/maltanar/pyverilator.git@0c3eb9343500fc1352a02c020a736c8c2db47e8e
 
 # Install pytest-xdist (not in requirements, only for faster testing in Docker)
 RUN pip install pytest-xdist==2.0.0

src/finn/core/data_layout.py

@@ -31,5 +31,19 @@
 
 NHWC = ["N", "H", "W", "C"]
 NCHW = ["N", "C", "H", "W"]
+NCW = ["N", "C", "W"]
+NWC = ["N", "W", "C"]
 NC = ["N", "C"]
 UNKNOWN = []
+
+
+def is_channels_last(layout):
+    return layout[-1] == "C"
+
+
+def get_channels_last_layout_for_ndims(ndims):
+    return {4: NHWC, 3: NWC, 2: NC}[ndims]
+
+
+def get_channels_first_layout_for_ndims(ndims):
+    return {4: NCHW, 3: NCW, 2: NC}[ndims]

src/finn/core/modelwrapper.py

@@ -58,7 +58,9 @@ def __init__(self, onnx_model_proto, make_deepcopy=False):
         is made internally.
         """
         if isinstance(onnx_model_proto, str):
-            assert os.path.isfile(onnx_model_proto)
+            assert os.path.isfile(
+                onnx_model_proto
+            ), f"File not found: {onnx_model_proto}"
             self._model_proto = onnx.load(onnx_model_proto)
         elif isinstance(onnx_model_proto, bytes):
             self._model_proto = onnx.load_from_string(onnx_model_proto)
@@ -217,7 +219,11 @@ def get_tensor_valueinfo(self, tensor_name):
         vi_names += [(x.name, x) for x in graph.output]
         vi_names += [(x.name, x) for x in graph.value_info]
         try:
-            vi_ind = [x[0] for x in vi_names].index(tensor_name)
+            vi_t_names = [x[0] for x in vi_names]
+            assert vi_t_names.count(tensor_name) <= 1, (
+                "Multiple ValueInfoProto found for " + tensor_name
+            )
+            vi_ind = vi_t_names.index(tensor_name)
             vi = vi_names[vi_ind][1]
             return vi
         except ValueError:
@@ -230,7 +236,11 @@ def get_tensor_shape(self, tensor_name):
         vi_names += [(x.name, x) for x in graph.output]
         vi_names += [(x.name, x) for x in graph.value_info]
         try:
-            vi_ind = [x[0] for x in vi_names].index(tensor_name)
+            vi_t_names = [x[0] for x in vi_names]
+            assert vi_t_names.count(tensor_name) <= 1, (
+                "Multiple ValueInfoProto found for " + tensor_name
+            )
+            vi_ind = vi_t_names.index(tensor_name)
             vi = vi_names[vi_ind][1]
             dims = [x.dim_value for x in vi.type.tensor_type.shape.dim]
             return dims
@@ -240,6 +250,8 @@ def get_tensor_shape(self, tensor_name):
     def set_tensor_shape(self, tensor_name, tensor_shape, dtype=TensorProto.FLOAT):
         """Assigns shape in ValueInfoProto for tensor with given name."""
         new_vi = oh.make_tensor_value_info(tensor_name, dtype, tensor_shape)
+        # call get_tensor_shape to catch multiple ValueInfoProto cases
+        self.get_tensor_shape(tensor_name)
         # find what container tis tensor's ValueInfo lives in
         # if not found anywhere, we assume it's a new value_info
         target_container = self.graph.value_info
@@ -534,13 +546,7 @@ def get_tensor_layout(self, tensor_name):
     def set_tensor_layout(self, tensor_name, data_layout):
         """Sets the data layout annotation of tensor with given name. See
         get_tensor_layout for examples."""
-        tensor_shape = self.get_tensor_shape(tensor_name)
         assert type(data_layout) == list, "data_layout must be a list"
-        if tensor_shape is not None:
-            assert len(tensor_shape) == len(
-                data_layout
-            ), """Mismatch between number
-            of dimensions of tensor shape and data layout annotation."""
         graph = self._model_proto.graph
         qnt_annotations = graph.quantization_annotation
         ret = util.get_by_name(qnt_annotations, tensor_name, "tensor_name")

src/finn/core/onnx_exec.py

@@ -51,7 +51,7 @@ def execute_node(node, context, graph, return_full_exec_context=False):
 
     Input/output provided via context."""
 
-    if node.op_type == "GenericPartition":
+    if node.op_type in ["GenericPartition", "StreamingDataflowPartition"]:
         partition_node = getCustomOp(node)
         model = ModelWrapper(partition_node.get_nodeattr("model"))
         inp_ctx = dict(filter(lambda x: x[0] in node.input, context.items()))
@@ -71,32 +71,6 @@ def execute_node(node, context, graph, return_full_exec_context=False):
             for tname in ret.keys():
                 if tname not in [x.name for x in model.graph.output]:
                     context[node.name + "_" + tname] = ret[tname]
-    elif node.op_type == "StreamingDataflowPartition":
-        sdp_node = getCustomOp(node)
-        model = ModelWrapper(sdp_node.get_nodeattr("model"))
-        inp_ctx = dict(filter(lambda x: x[0] in node.input, context.items()))
-        # input may have been renamed in partition
-        assert len(inp_ctx) == 1
-        old_iname = node.input[0]
-        new_iname = model.graph.input[0].name
-        if old_iname != new_iname:
-            inp_ctx[new_iname] = inp_ctx[old_iname]
-            del inp_ctx[old_iname]
-        ret = execute_onnx(model, inp_ctx, return_full_exec_context)
-        # if the model was in ip-stitched rtlsim mode, may get annotation
-        # for numbet of elapsed cycles, save again
-        if model.get_metadata_prop("exec_mode") == "rtlsim":
-            model.save(sdp_node.get_nodeattr("model"))
-        # output may have been renamed in partition
-        assert len(model.graph.output) == 1
-        node_oname = node.output[0]
-        model_oname = model.graph.output[0].name
-        context[node_oname] = ret[model_oname]
-        # prefix and insert exec context entries
-        if return_full_exec_context:
-            for tname in ret.keys():
-                if tname != model_oname:
-                    context[node.name + "_" + tname] = ret[tname]
     else:
         if is_finn_op(node.domain):
             ex_cu_node.execute_custom_node(node, context, graph)
@@ -108,7 +82,9 @@ def execute_node(node, context, graph, return_full_exec_context=False):
             # graph.value_info as well as graph.output or graph.input
             # nodes with multiple outputs that are a mix of value_info and
             # input/outputs may get them reordered below
+            # note: a node's input may (also) be a top-level input or output
             node_inputs = list(filter(lambda x: x.name in node.input, graph.input))
+            node_inputs += list(filter(lambda x: x.name in node.input, graph.output))
             node_inputs += list(
                 filter(lambda x: x.name in node.input, graph.value_info)
             )

src/finn/custom_op/general/quantavgpool2d.py

@@ -94,7 +94,7 @@ def make_shape_compatible_op(self, model):
     def infer_node_datatype(self, model):
         node = self.onnx_node
         bw = self.get_nodeattr("obits")
-        if bw in [2, 4, 8, 16, 32]:
+        if bw in range(2, 33):
             if self.get_nodeattr("signed") == 0:
                 dtype = DataType["UINT%d" % bw]
             else:

src/finn/transformation/change_3d_tensors_to_4d.py

@@ -56,6 +56,12 @@ def _find_invalid_nodes(model):
         "Transpose",
         "LogSoftmax",
         "ArgMax",
+        "Div",
+        "TopK",
+        "MatMul",
+        "Flatten",
+        "Reshape",
+        "MaxPool",
     ]
     invalid_nodes = []
     for n in model.graph.node:
@@ -96,7 +102,7 @@ def apply(self, model):
         model = model.transform(RemoveUnusedTensors())
 
         # This list contains all nodes with initializers that need to be converted
-        nodes_with_initializers = ["Mul", "Conv", "Add"]
+        nodes_with_initializers = ["Mul", "Conv", "Add", "Div", "Reshape"]
         # Obtain a list of initializer names (used to filter out only value infos)
         initializers_names = [x.name for x in model.graph.initializer]
 
@@ -118,8 +124,7 @@ def apply(self, model):
                 if x.name not in initializers_names
             },
         }
-        # Extract only initializers from Conv, Mul and Add nodes (which are the
-        # only ones relevant for conversion)
+        # Extract only initializers from nodes that are relevant for conversion
         all_tensors = {
             **all_tensors,
             **{
@@ -143,10 +148,11 @@ def apply(self, model):
         tensors_reduced_dimension = []
         for n in model.graph.node:
             node_op_type = n.op_type
+            input_shape = model.get_tensor_shape(n.input[0])
             # Find tensors that are the output of nodes that reduce the dimension
             if node_op_type == "ArgMax":
                 keep_dims = get_by_name(n.attribute, "keepdims", "name").i
-                if keep_dims == 0:
+                if len(input_shape) == 3 and keep_dims == 0:
                     node_out = n.output
                     for n_o in node_out:
                         tensors_reduced_dimension.append(n_o)
@@ -158,10 +164,10 @@ def apply(self, model):
                     len(perm) == 3
                 ):  # Meaning that the transpose operation was on a 3D tensor
                     perm.append(3)  # append 4th dimension
-            elif node_op_type == "ArgMax" or node_op_type == "LogSoftMax":
+            elif node_op_type in ["ArgMax", "LogSoftMax", "TopK", "Flatten"]:
                 axis = get_by_name(n.attribute, "axis", "name")
-                if axis.i == -1:
-                    axis.i = 2  # argmax is now on the second-to-last axis
+                if len(input_shape) == 3 and axis.i < 0:
+                    axis.i = 3 + axis.i  # count dimensions from the front
             elif node_op_type == "Conv":
                 dilations = get_by_name(n.attribute, "dilations", "name").ints
                 kernel_shape = get_by_name(n.attribute, "kernel_shape", "name").ints
@@ -180,6 +186,19 @@ def apply(self, model):
                     pads.append(0)
                 if len(strides) == 1:  # strides = [stride_h, stride_w]
                     strides.append(1)
+            elif node_op_type == "MaxPool":
+                kernel_shape = get_by_name(n.attribute, "kernel_shape", "name").ints
+                pads = get_by_name(n.attribute, "pads", "name").ints
+                strides = get_by_name(n.attribute, "strides", "name").ints
+                if len(kernel_shape) == 1:  # we must add another dimension to it
+                    kernel_shape.append(1)
+                if (
+                    len(pads) == 2
+                ):  # pads = [x1_begin, x1_end] --> [x1_begin, x2_begin, x1_end, x2_end]
+                    pads.insert(1, 0)
+                    pads.append(0)
+                if len(strides) == 1:  # strides = [stride_h, stride_w]
+                    strides.append(1)
 
         # Change format of each input/value_info/output tensor
         for k, v in all_tensors.items():

src/finn/transformation/create_generic_partitions.py

@@ -131,33 +131,23 @@ def apply(self, model):
                 to_check = next_to_check
 
             # set p graph in/out to be p_in/p_out
-            for x in p_model.graph.input:
-                p_model.graph.input.remove(x)
+            while len(p_model.graph.input) > 0:
+                p_model.graph.input.pop()
             for i in p_in_vi:
                 p_model.graph.input.append(i)
 
-            for x in p_model.graph.output:
-                p_model.graph.output.remove(x)
+            while len(p_model.graph.output) > 0:
+                p_model.graph.output.pop()
             for o in p_out_vi:
                 p_model.graph.output.append(o)
 
             # remove redundant input and output value_info entries
             for i in p_in_vi:
-                # the tensor can be both an input and value_info, so we also have to
-                # ensure that the tensor is not a relevant value_info before removing
-                if (
-                    i in p_model.graph.value_info
-                    and p_model.find_producer(i.name) is None
-                ):
+                if i in p_model.graph.value_info:
                     p_model.graph.value_info.remove(i)
 
             for o in p_out_vi:
-                # the tensor can both an output and value_info, so we also have to
-                # ensure that the tensor is not a relevant value_info before removing
-                if (
-                    o in p_model.graph.value_info
-                    and p_model.find_consumers(o.name) is None
-                ):
+                if o in p_model.graph.value_info:
                     p_model.graph.value_info.remove(o)
 
             # save partition model

src/finn/transformation/general.py

@@ -134,9 +134,13 @@ def apply(self, model):
                 if model.get_initializer(i) is not None:
                     model.rename_tensor(i, "%s_param%d" % (n.name, init_in_num))
                     init_in_num += 1
-        # give special names to the main model input and output
-        model.rename_tensor(model.graph.input[0].name, "global_in")
-        model.rename_tensor(model.graph.output[0].name, "global_out")
+        # give special names to the model inputs and outputs
+        for i, inp in enumerate(model.graph.input):
+            iname = "global_in" if i == 0 else "global_in_%d" % i
+            model.rename_tensor(inp.name, iname)
+        for i, outp in enumerate(model.graph.output):
+            oname = "global_out" if i == 0 else "global_out_%d" % i
+            model.rename_tensor(outp.name, oname)
         # return model_was_changed = False as single iteration is always enough
         return (model, False)
 

src/finn/transformation/infer_data_layouts.py

@@ -63,35 +63,40 @@ def _infer_node_data_layout(model, node):
     """Infer output data layout annotation(s) for a particular node.
     Returns True if any changes were made."""
     old_layouts = list(map(lambda x: model.get_tensor_layout(x), node.output))
-    if is_finn_op(node.domain):
-        # try to guess based on number of output dims
-        for o in node.output:
-            ndims = len(model.get_tensor_shape(o))
-            new_layout = _dims_to_layout(model, node, ndims)
-            model.set_tensor_layout(o, new_layout)
-    else:
-        if node.op_type == "Transpose":
-            # grab input annotation and switch it around using perm
-            perm = get_by_name(node.attribute, "perm").ints
-            inp_layout = model.get_tensor_layout(node.input[0])
-            out_layout = [inp_layout[i] for i in perm]
-            model.set_tensor_layout(node.output[0], out_layout)
-        elif node.op_type == "Unsqueeze":
-            inp_layout = model.get_tensor_layout(node.input[0])
-            # add dummy dimension at the output
-            out_layout = inp_layout + ["x"]
-            model.set_tensor_layout(node.output[0], out_layout)
-        elif node.op_type == "Squeeze":
-            inp_layout = model.get_tensor_layout(node.input[0])
-            assert inp_layout[-1] == "x"
-            # remove dummy dimension
-            out_layout = inp_layout[:-1]
-            model.set_tensor_layout(node.output[0], out_layout)
-        else:
+    try:
+        if is_finn_op(node.domain):
             # try to guess based on number of output dims
             for o in node.output:
                 ndims = len(model.get_tensor_shape(o))
-                model.set_tensor_layout(o, _dims_to_layout(model, node, ndims))
+                new_layout = _dims_to_layout(model, node, ndims)
+                model.set_tensor_layout(o, new_layout)
+        else:
+            if node.op_type == "Transpose":
+                # grab input annotation and switch it around using perm
+                perm = get_by_name(node.attribute, "perm").ints
+                inp_layout = model.get_tensor_layout(node.input[0])
+                out_layout = [inp_layout[i] for i in perm]
+                model.set_tensor_layout(node.output[0], out_layout)
+            elif node.op_type == "Unsqueeze":
+                inp_layout = model.get_tensor_layout(node.input[0])
+                # add dummy dimension at the output
+                out_layout = inp_layout + ["x"]
+                model.set_tensor_layout(node.output[0], out_layout)
+            elif node.op_type == "Squeeze":
+                inp_layout = model.get_tensor_layout(node.input[0])
+                assert inp_layout[-1] == "x"
+                # remove dummy dimension
+                out_layout = inp_layout[:-1]
+                model.set_tensor_layout(node.output[0], out_layout)
+            else:
+                # try to guess based on number of output dims
+                for o in node.output:
+                    ndims = len(model.get_tensor_shape(o))
+                    model.set_tensor_layout(o, _dims_to_layout(model, node, ndims))
+    except Exception:
+        for o in node.output:
+            model.set_tensor_layout(o, DataLayout.UNKNOWN)
+
     # compare old and new output dtypes to see if anything changed
     new_layouts = list(map(lambda x: model.get_tensor_layout(x), node.output))
     graph_modified = new_layouts != old_layouts