From 81a0744d33031d30021283c9f928b3a6fd4039df Mon Sep 17 00:00:00 2001 From: Christoph Berganski Date: Wed, 13 Dec 2023 17:13:47 +0100 Subject: [PATCH 1/6] Add rudimentary support for "arbitrary" dimensions in MultiThreshold This allows node execution of MultiThreshold operators with arbitrary number of dimensions, as long as the channel dimension is last. This is necessary to run some verification steps of attention operators which, at least for some intermediate steps, have 3 dimensional data layouts. This does not change the behavior of execution on the already existing 2d and 4d data layouts. --- src/qonnx/custom_op/general/multithreshold.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/src/qonnx/custom_op/general/multithreshold.py b/src/qonnx/custom_op/general/multithreshold.py index 57a2872d..6df58f95 100644 --- a/src/qonnx/custom_op/general/multithreshold.py +++ b/src/qonnx/custom_op/general/multithreshold.py @@ -133,6 +133,23 @@ def execute_node(self, context, graph): pass else: raise Exception("Unknown data_layout and input ndim" " combination for MultiThreshold.") + + # Remember whether the shape has been modified to handle 1d or 3d data + # layouts + orig_shape = None + # If the input tensor has dimensions not covered by the NC or NCWH data + # layouts, the shape needs to be adapted such that it can be handled by + # multithreshold. + # TODO: Seems like a rather sketchy solution to support arbitrary data + # layouts. This does not even validate the assumption of channel last + # layout. + if v.ndim not in {2, 4}: + # Remember the original shape to be restored later + orig_shape = v.shape + # Assume last dimension to be the channel dimension C and reshape + # into NC layout which is supported by multithreshold + v = v.reshape((-1, v.shape[-1])) + # calculate output output = multithreshold(v, thresholds, out_scale, out_bias) # setting context according to output @@ -145,6 +162,13 @@ def execute_node(self, context, graph): pass else: raise Exception("Unknown data_layout and output ndim" " combination for MultiThreshold.") + + # If the shape has been modified to support arbitrary layouts, restore + # the original shape + # TODO: Part of the rather sketchy solution above. + if orig_shape is not None: + output = output.reshape(orig_shape) + context[node.output[0]] = output def verify_node(self): From bfabd4d3db18217c3f26ff4d471daa1d4a917180 Mon Sep 17 00:00:00 2001 From: Christoph Berganski Date: Fri, 13 Sep 2024 15:29:39 +0200 Subject: [PATCH 2/6] [MultiThreshold] Generalize data layouts for node execution The relevant aspect of the data layout annotation seems to be which axis is labeled as the channel dimension "C": We do not actually have to care about the total number and ordering of the other axes, as long as we can find the index of the "C" axis and swap to have "C" at index 1 for node execution (and swap it back afterwards). Falls back to the default assumption that "C" is at index 1 if there is no layout annotation, which is equivalent to the "NCHW" or "NC" layouts. This is a rather experimental change which might break existing code and is currently still restricted to the well-known 2-, 3- and 4-dimensional layouts. --- src/qonnx/custom_op/general/multithreshold.py | 66 ++++++------------- 1 file changed, 20 insertions(+), 46 deletions(-) diff --git a/src/qonnx/custom_op/general/multithreshold.py b/src/qonnx/custom_op/general/multithreshold.py index 6df58f95..f1d985f6 100644 --- a/src/qonnx/custom_op/general/multithreshold.py +++ b/src/qonnx/custom_op/general/multithreshold.py @@ -92,7 +92,13 @@ def get_nodeattr_types(self): "out_dtype": ("s", True, ""), "out_scale": ("f", False, 1.0), "out_bias": ("f", False, 0.0), - "data_layout": ("s", False, "NCHW", {"NCHW", "NHWC"}), + # fmt: off + "data_layout": ("s", False, "NCHW", { + # TODO: Add more options or even remove the set of allowed data + # layouts - it is really the placement of the "C" that matters + "NCHW", "NHWC", "NC", "NWC", "NCW" + }) + # fmt: on } def make_shape_compatible_op(self, model): @@ -122,53 +128,21 @@ def execute_node(self, context, graph): # retrieve attributes if output scaling is used out_scale = self.get_nodeattr("out_scale") out_bias = self.get_nodeattr("out_bias") - # transpose input if NHWC data layout is chosen + + # Consider the data layout for transposing the input into the format + # accepted by the multithreshold function above, i.e, the channel + # dimension is along the axis with index 1. data_layout = self.get_nodeattr("data_layout") - if data_layout == "NHWC": - if v.ndim == 4: - # NHWC -> NCHW - v = np.transpose(v, (0, 3, 1, 2)) - elif v.ndim == 2: - # no HW dimension means NHWC and NCHW layouts are equivalent - pass - else: - raise Exception("Unknown data_layout and input ndim" " combination for MultiThreshold.") - - # Remember whether the shape has been modified to handle 1d or 3d data - # layouts - orig_shape = None - # If the input tensor has dimensions not covered by the NC or NCWH data - # layouts, the shape needs to be adapted such that it can be handled by - # multithreshold. - # TODO: Seems like a rather sketchy solution to support arbitrary data - # layouts. This does not even validate the assumption of channel last - # layout. - if v.ndim not in {2, 4}: - # Remember the original shape to be restored later - orig_shape = v.shape - # Assume last dimension to be the channel dimension C and reshape - # into NC layout which is supported by multithreshold - v = v.reshape((-1, v.shape[-1])) - - # calculate output + # Lookup the index of the channel dimension in the data layout + # Note: Assumes there is at most one "C" which denotes the channel + # dimension + cdim = data_layout.index("C") if "C" in data_layout else 1 + # Rearrange the input to the expected (N, C, ...) layout + v = v.swapaxes(cdim, 1) + # Now we can use the multithreshold function to calculate output output = multithreshold(v, thresholds, out_scale, out_bias) - # setting context according to output - if data_layout == "NHWC": - if output.ndim == 4: - # NCHW -> NHWC - output = np.transpose(output, (0, 2, 3, 1)) - elif output.ndim == 2: - # no HW dimension means NHWC and NCHW layouts are equivalent - pass - else: - raise Exception("Unknown data_layout and output ndim" " combination for MultiThreshold.") - - # If the shape has been modified to support arbitrary layouts, restore - # the original shape - # TODO: Part of the rather sketchy solution above. - if orig_shape is not None: - output = output.reshape(orig_shape) - + # Rearrange the output back to the original layout + output = output.swapaxes(cdim, 1) context[node.output[0]] = output def verify_node(self): From 4941c412ff3b882590c40f35fa5e835aee77099e Mon Sep 17 00:00:00 2001 From: Christoph Berganski Date: Fri, 13 Sep 2024 16:05:22 +0200 Subject: [PATCH 3/6] [Test] Add more allowed data layouts for MultiThreshold --- tests/core/test_custom_onnx_exec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/core/test_custom_onnx_exec.py b/tests/core/test_custom_onnx_exec.py index a763e268..c433909d 100644 --- a/tests/core/test_custom_onnx_exec.py +++ b/tests/core/test_custom_onnx_exec.py @@ -274,7 +274,7 @@ def test_execute_custom_node_multithreshold(): assert (execution_context["out"] == outputs_nhwc).all() # check the set of allowed values op_inst = getCustomOp(node_def) - assert op_inst.get_nodeattr_allowed_values("data_layout") == {"NCHW", "NHWC"} + assert op_inst.get_nodeattr_allowed_values("data_layout") == {"NCHW", "NHWC", "NC", "NWC", "NCW"} # exercise the allowed value checks # try to set attribute to non-allowed value, should raise an exception try: From 817a23ebdf3d38d0fae59bcb67b32ac77d714768 Mon Sep 17 00:00:00 2001 From: Christoph Berganski Date: Thu, 10 Oct 2024 16:36:59 +0200 Subject: [PATCH 4/6] [MultiThreshold] Remove set of allowed data layouts As we only really care for the index of the "C" axis there is no need to restrict the set of valid layouts here. --- src/qonnx/custom_op/general/multithreshold.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/qonnx/custom_op/general/multithreshold.py b/src/qonnx/custom_op/general/multithreshold.py index f1d985f6..708840e4 100644 --- a/src/qonnx/custom_op/general/multithreshold.py +++ b/src/qonnx/custom_op/general/multithreshold.py @@ -93,11 +93,7 @@ def get_nodeattr_types(self): "out_scale": ("f", False, 1.0), "out_bias": ("f", False, 0.0), # fmt: off - "data_layout": ("s", False, "NCHW", { - # TODO: Add more options or even remove the set of allowed data - # layouts - it is really the placement of the "C" that matters - "NCHW", "NHWC", "NC", "NWC", "NCW" - }) + "data_layout": ("s", False, "NCHW") # fmt: on } From 9d73e166fde765792224aa87becc4104e5140c15 Mon Sep 17 00:00:00 2001 From: Christoph Berganski Date: Thu, 10 Oct 2024 16:51:00 +0200 Subject: [PATCH 5/6] [Test] Remove "allowed values" check for data layouts CustomOp --- tests/core/test_custom_onnx_exec.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/core/test_custom_onnx_exec.py b/tests/core/test_custom_onnx_exec.py index c433909d..8eec7156 100644 --- a/tests/core/test_custom_onnx_exec.py +++ b/tests/core/test_custom_onnx_exec.py @@ -274,7 +274,9 @@ def test_execute_custom_node_multithreshold(): assert (execution_context["out"] == outputs_nhwc).all() # check the set of allowed values op_inst = getCustomOp(node_def) - assert op_inst.get_nodeattr_allowed_values("data_layout") == {"NCHW", "NHWC", "NC", "NWC", "NCW"} + # TODO: Removed this check to generalize the supported data layouts, but do + # we need some other check to verify the validity of data layouts? + # assert op_inst.get_nodeattr_allowed_values("data_layout") == {"NCHW", "NHWC", "NC", "NWC", "NCW"} # exercise the allowed value checks # try to set attribute to non-allowed value, should raise an exception try: From c0b45347b193cbd08186dc67dfc6c5c3912a614d Mon Sep 17 00:00:00 2001 From: Christoph Berganski Date: Thu, 24 Oct 2024 11:59:09 +0200 Subject: [PATCH 6/6] [MultiThreshold] Replace default data_layout by fallback in execute_node Note: Only covers data layouts for tensors with less than five axes --- src/qonnx/custom_op/general/multithreshold.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/qonnx/custom_op/general/multithreshold.py b/src/qonnx/custom_op/general/multithreshold.py index 708840e4..bcf0731d 100644 --- a/src/qonnx/custom_op/general/multithreshold.py +++ b/src/qonnx/custom_op/general/multithreshold.py @@ -92,9 +92,7 @@ def get_nodeattr_types(self): "out_dtype": ("s", True, ""), "out_scale": ("f", False, 1.0), "out_bias": ("f", False, 0.0), - # fmt: off - "data_layout": ("s", False, "NCHW") - # fmt: on + "data_layout": ("s", False, ""), } def make_shape_compatible_op(self, model): @@ -129,6 +127,13 @@ def execute_node(self, context, graph): # accepted by the multithreshold function above, i.e, the channel # dimension is along the axis with index 1. data_layout = self.get_nodeattr("data_layout") + # If there is no layout annotation, guess based on rank of the + # tensor + if not data_layout and len(v.shape) < 5: + # Maps tensor rank to layout annotation + rank_to_layout = {0: None, 1: "C", 2: "NC", 3: "NWC", 4: "NCHW"} + # Lookup the layout required by this input shape + data_layout = rank_to_layout[len(v.shape)] # Lookup the index of the channel dimension in the data layout # Note: Assumes there is at most one "C" which denotes the channel # dimension