This is an automated email from the ASF dual-hosted git repository.

ekalda pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new fa223b077f [Bugfix][TVMC] Fix tvmc option for printing which operators 
are offloaded to the Ethos-U (#14994)
fa223b077f is described below

commit fa223b077f3fb375e42f1e0d4563330daa730f18
Author: sergio-grovety <[email protected]>
AuthorDate: Mon Jun 19 10:17:04 2023 +0300

    [Bugfix][TVMC] Fix tvmc option for printing which operators are offloaded 
to the Ethos-U (#14994)
    
    Lines from the initial Relay, which don't correspond to the relay.Call now 
are printed to the output with the --dump-offloads option enabled. Thus we get 
rid of the incomprehensible gaps in the line numbers of the initial relay, 
which occurred before.
    
    ---------
    
    Co-authored-by: Sergey Smirnov 
<[email protected]>
    Co-authored-by: Arina.Naumova <[email protected]>
---
 python/tvm/driver/tvmc/compiler.py                 |  44 +++++---
 .../tvm/relay/analysis/operations_distribution.py  |  38 +++----
 python/tvm/relay/transform/suffixes.py             |   4 +-
 .../test_pass_operations_distribution.py           |  48 ++++----
 tests/python/driver/tvmc/test_compiler.py          | 123 ++++++++++++++-------
 5 files changed, 163 insertions(+), 94 deletions(-)

diff --git a/python/tvm/driver/tvmc/compiler.py 
b/python/tvm/driver/tvmc/compiler.py
index c42974593a..91f56a3031 100644
--- a/python/tvm/driver/tvmc/compiler.py
+++ b/python/tvm/driver/tvmc/compiler.py
@@ -572,22 +572,27 @@ def dump_operation_offloads(mod: tvm.ir.IRModule, 
initial_mod: tvm.ir.IRModule,
                 # to the Relay, in which case we can not associate the initial 
Relay string
                 # with the resulting Relay call
                 source_name = x.span.source_name.name
-                if source_name in operations_distribution:
-                    compiler_name, op_name, func_id = 
operations_distribution[source_name]
+                suffix = tvm.relay.transform.suffixes.SUFFIX_STRING
+                result = re.search(r"(.*)(" + suffix + r")(.*)", source_name)
+                func_id = result.group(1)
+                if func_id in operations_distribution:
+                    compiler_name, op_name = operations_distribution[func_id]
                     ret = (
                         f", compiler_name: {compiler_name}, op_name: 
{op_name}, "
                         f"func_id: {func_id}"
                     )
                 else:
                     ret = ", compiler_name: unknown, op_name: unknown, 
func_id: unknown"
+            elif isinstance(x, (relay.Tuple, relay.TupleGetItem)):
+                ret = ", compiler_name: none, op_name: none, func_id: none"
+
             return ret
 
         initial_relay_astext = initial_mod.astext(show_meta_data=False, 
annotate=annotate_f).split(
             "\n"
         )
 
-        # funcs_list is a list of internal composite/function IDs
-        # generated by analyze_operations_distribution().
+        # funcs_list is a list of internal composite/function IDs.
         # funcs_list helps keep the order of lines from the initial Relay.
         funcs_list = []
 
@@ -598,14 +603,20 @@ def dump_operation_offloads(mod: tvm.ir.IRModule, 
initial_mod: tvm.ir.IRModule,
         # funcs_dict is a mapping of the generated 
analyze_operations_distribution
         # internal composite/function IDs to a list, where:
         # 1st element is
-        #   (1a): target name - it could be "generic" or "unknown" or
+        #   (1a): "generic"|"unknown"|"none"* or
         #   (1b): specific target name, like "ethos-u" or "cmsis-nn"
         # 2nd element is
         #   (2a): corresponding initial Relay line for the case (1a) or
         #   (2b): the name of the target composite functon in the other case 
(1b)
         # 3rd element or subsequent ones are presented only for the case (2b)
-        # and are the initial Relay lines included in the corresponding
+        # and are the initial Relay's lines included in the corresponding
         # target composite functon
+        #
+        # *Description of what is meant by "generic"|"unknown"|"none":
+        # "generic" means that operation will be run on a host
+        # "unknown" means that unique identifier of this Relay line not found 
in the partitioned
+        #           Relay and therefore not present in the 
operations_distribution dictionary
+        # "none" means that this Relay line is not relay.Call
         funcs_dict = {}
 
         # Here we group together initial Relay lines from the one composite
@@ -618,19 +629,20 @@ def dump_operation_offloads(mod: tvm.ir.IRModule, 
initial_mod: tvm.ir.IRModule,
                 target_name = result.group(2)
                 op_name = result.group(4)
                 func_id = result.group(6)
-                s = re.sub(r", compiler_name: (.*)", "", s).lstrip()
-                target_statistic[target_name] += 1
+                if target_name != "none":
+                    target_statistic[target_name] += 1
 
-                # create an identifier for each "unknown" case to keep the 
lines order
-                if func_id == "unknown":
+                # create an identifier for each "unknown" or "none" case to 
keep the lines order
+                if func_id == "unknown" or func_id == "none" or target_name == 
"generic":
                     func_id = str(next(counter) * -1)
 
                 if func_id not in funcs_dict:
                     funcs_list.append(func_id)
                     funcs_dict[func_id] = [target_name]
-                    if target_name not in ["unknown", "generic"]:
+                    if target_name not in ["unknown", "generic", "none"]:
                         funcs_dict[func_id].append(op_name)
 
+                s = re.sub(r", compiler_name: (.*)", "", s).lstrip()
                 funcs_dict[func_id].append(s)
 
         # Here we prepare the output for printing.
@@ -648,7 +660,12 @@ def dump_operation_offloads(mod: tvm.ir.IRModule, 
initial_mod: tvm.ir.IRModule,
 
         for func_id in funcs_list:
             _list = funcs_dict[func_id]
-            output.append(f"{_list[0]:10}     <-     {_list[1]}")
+
+            if _list[0] != "none":
+                output.append(f"{_list[0]:<15}<-{' ':5}{_list[1]}")
+            else:
+                output.append(f"{' ':>22}{_list[1]}")
+
             if _list[0] == "unknown":
                 output.append(
                     "Warning: The above line means that some pass(es) \
@@ -660,7 +677,7 @@ def dump_operation_offloads(mod: tvm.ir.IRModule, 
initial_mod: tvm.ir.IRModule,
                               with the resulting Relay"
                 )
             for el in _list[2:]:
-                output.append(f"{_list[0]:10}     <-          {el}")
+                output.append(f"{_list[0]:<15}<-{' ':10}{el}")
 
         if print_to_console:
             print("\n" + "\n".join(output))
@@ -669,3 +686,4 @@ def dump_operation_offloads(mod: tvm.ir.IRModule, 
initial_mod: tvm.ir.IRModule,
             os.makedirs(os.path.dirname(file_path), exist_ok=True)
             with open(file_path, "w") as f:
                 f.write("\n".join(output))
+                f.write("\n")
diff --git a/python/tvm/relay/analysis/operations_distribution.py 
b/python/tvm/relay/analysis/operations_distribution.py
index fc983c8e7e..769f9ee88f 100644
--- a/python/tvm/relay/analysis/operations_distribution.py
+++ b/python/tvm/relay/analysis/operations_distribution.py
@@ -17,6 +17,8 @@
 """Utilities that enable analyze Relay and get mappings for
 the unique identifier of the Relay line to the tuple of
 compiler name, composite name and composite/function identifier."""
+import re
+
 import tvm
 from tvm import relay
 from tvm.relay.expr_functor import ExprVisitor
@@ -24,23 +26,23 @@ from tvm.relay.expr_functor import ExprVisitor
 
 class AnalyzeOperationsDistribution(ExprVisitor):
     """A visitor pass that maintains the dictionary unique_op_ids where
-    the tuple (compiler name, composite name, composite/function identifier)
-    corresponds to the unique identifier of the Relay line.
-    TVMC compiler adds a unique Relay line identifier as a suffix
-    to the call span field using the tag_suffixes pass
-    if the --dump-offloads option is specified.
+    the tuple (compiler name, composite name) corresponds to the unique
+    identifier of the Relay line. The identifier will allow us to link
+    the lines of the initial Relay with the information about operators
+    offloading, which is present in the partitioned Relay
+    TVMC compiler adds a unique Relay line identifier as a suffix to the
+    call span field using the tag_suffixes pass if the --dump-offloads
+    option is specified.
 
     Attributes
     ----------
-    unique_op_ids : Dict[str, str, int]
+    unique_op_ids : Dict[str, str]
         Mapping the unique identifier of the Relay line obtained from
         the "span" field of the Call and the tuple of compiler name,
-        composite name and internal composite/function identifier.
+        composite name.
     func_name : str
-        The name of the composite name in the partitioned Relay or
+        The name of the composite in the partitioned Relay or
         'generic' in case the Call has not been included in any composite.
-    func_id : int
-        Internal(inside unique_op_ids) composite/function identifier.
     compiler_name : str
         A name of the compiler (e.g. 'ethos-u' or 'cmsis-nn') or 'generic'
         in case the Call has not been included in any composite.
@@ -49,7 +51,6 @@ class AnalyzeOperationsDistribution(ExprVisitor):
     def __init__(self):
         self.unique_op_ids = {}
         self.func_name = ""
-        self.func_id = 1
         self.compiler_name = ""
         super().__init__()
 
@@ -64,12 +65,12 @@ class AnalyzeOperationsDistribution(ExprVisitor):
         if isinstance(call.op, tvm.ir.Op):
             if call.span:
                 src = call.span.source_name.name
-                self.unique_op_ids[src] = [self.compiler_name, self.func_name, 
self.func_id]
-                if self.func_name == "generic":
-                    self.func_id += 1
+                suffix = tvm.relay.transform.suffixes.SUFFIX_STRING
+                result = re.search(r"(.*)(" + suffix + r")(.*)", src)
+                res = result.group(1)
+                self.unique_op_ids[res] = [self.compiler_name, self.func_name]
         if isinstance(call.op, relay.Function):
             self.func_name = call.op.attrs["Composite"]
-            self.func_id += 1
         super().visit_call(call)
 
 
@@ -78,7 +79,7 @@ def analyze_operations_distribution(mod):
     of the Relay line from the Call's span field.
     The result is maintained in the dictionary unique_op_ids where
     the unique indicator obtained from the op's span corresponds to
-    the tuple (compiler name, composite name, composite/function identifier).
+    the tuple (compiler name, composite name).
     With this information we can annotate the textual representation
     of the initial Relay by indicating into which target composite
     and function the operators are converted
@@ -91,10 +92,9 @@ def analyze_operations_distribution(mod):
 
     Returns
     -------
-    unique_op_ids : Dict[str, str, int]
+    unique_op_ids : Dict[str, str]
         Mapping from the unique identifier of the Relay line to the tuple of
-        compiler name, composite name, internal composite/function
-        identifier.
+        compiler name, composite name.
     """
     analyze = AnalyzeOperationsDistribution()
     for _, func in mod.functions.items():
diff --git a/python/tvm/relay/transform/suffixes.py 
b/python/tvm/relay/transform/suffixes.py
index e2f7a3c224..eaca116070 100644
--- a/python/tvm/relay/transform/suffixes.py
+++ b/python/tvm/relay/transform/suffixes.py
@@ -22,6 +22,8 @@ import tvm
 from ..expr_functor import ExprMutator
 from .. import expr as _expr
 
+SUFFIX_STRING = r"_PART_"
+
 
 class _SuffixTagger(ExprMutator):
     """A pass to traverse the Relay graph to add suffix to the call's span 
fields.
@@ -34,7 +36,7 @@ class _SuffixTagger(ExprMutator):
         ExprMutator.__init__(self)
         # key: span or source name, value: counter, indexed from 0
         self.lookup = defaultdict(int)
-        self.suffix = "_PART_"
+        self.suffix = SUFFIX_STRING
         # a set to record hashes of an expressions which spans have been 
already rewritten
         self.hashes = set()
 
diff --git 
a/tests/python/contrib/test_ethosu/test_pass_operations_distribution.py 
b/tests/python/contrib/test_ethosu/test_pass_operations_distribution.py
index 17f973d8b3..8b127de327 100644
--- a/tests/python/contrib/test_ethosu/test_pass_operations_distribution.py
+++ b/tests/python/contrib/test_ethosu/test_pass_operations_distribution.py
@@ -45,6 +45,9 @@ def test_operations_distribution_ethos():
     def simple_net(x):
         weight_shape = [kernel_shape[0], kernel_shape[1], input_shape[3], 3]
         weights = tf.constant(np.random.uniform(size=weight_shape), 
dtype=tf.float32)
+        weight_shape[2] = 3
+        weights1 = tf.constant(np.random.uniform(size=weight_shape), 
dtype=tf.float32)
+        weights2 = tf.constant(np.random.uniform(size=weight_shape), 
dtype=tf.float32)
         op = tf.nn.conv2d(
             x,
             filters=weights,
@@ -53,21 +56,29 @@ def test_operations_distribution_ethos():
             data_format="NHWC",
             dilations=1,
         )
-        op = tf.pad(
+        op1 = tf.nn.conv2d(
             op,
-            [[0, 0], [padding[0], padding_out[2]], [padding_out[1], 
padding[3]], [0, 0]],
-            "CONSTANT",
+            filters=weights1,
+            strides=1,
+            padding="SAME",
+            data_format="NHWC",
+            dilations=1,
         )
-        op = tf.pad(
+        op2 = tf.nn.conv2d(
             op,
-            [[0, 0], [padding[0], padding[2]], [padding[1], padding[3]], [0, 
0]],
-            "CONSTANT",
+            filters=weights2,
+            strides=1,
+            padding="SAME",
+            data_format="NHWC",
+            dilations=1,
         )
-        return tf.pad(
+        op = tf.concat([op1, op2], 1)
+        op = tf.pad(
             op,
-            [[0, 0], [padding_out[0], padding[2]], [padding[1], 
padding_out[3]], [0, 0]],
+            [[0, 0], [padding[0], padding_out[1]], [padding_out[2], 
padding[3]], [0, 0]],
             "CONSTANT",
         )
+        return op
 
     _, tflite_graph = get_tflite_graph(simple_net, [input_shape])
 
@@ -88,12 +99,11 @@ def test_operations_distribution_ethos():
     operations_distribution = analyze_operations_distribution(mod)
 
     expected = {
-        "Pad_PART_0": ["generic", "generic", 1],
-        "Conv2D2_PART_2": ["ethos-u", "ethos-u.qnn_conv2d", 3],
-        "Conv2D2_PART_1": ["ethos-u", "ethos-u.qnn_conv2d", 3],
-        "Conv2D2_PART_0": ["ethos-u", "ethos-u.qnn_conv2d", 3],
-        "Identity_PART_0": ["ethos-u", "ethos-u.pad2d", 4],
-        "Pad_1_PART_0": ["ethos-u", "ethos-u.pad2d", 5],
+        "Identity": ["generic", "generic"],
+        "concat": ["ethos-u", "ethos-u.concat"],
+        "Conv2D_11": ["ethos-u", "ethos-u.qnn_conv2d"],
+        "Conv2D1": ["ethos-u", "ethos-u.qnn_conv2d"],
+        "Conv2D_221": ["ethos-u", "ethos-u.qnn_conv2d"],
     }
 
     assert operations_distribution == expected
@@ -160,12 +170,10 @@ def test_operations_distribution_generic():
     operations_distribution = analyze_operations_distribution(mod)
 
     expected = {
-        "Identity_PART_0": ["generic", "generic", 1],
-        "Pad_1_PART_0": ["generic", "generic", 2],
-        "Pad_PART_0": ["generic", "generic", 3],
-        "Conv2D2_PART_2": ["generic", "generic", 4],
-        "Conv2D2_PART_1": ["generic", "generic", 5],
-        "Conv2D2_PART_0": ["generic", "generic", 6],
+        "Identity": ["generic", "generic"],
+        "Pad_1": ["generic", "generic"],
+        "Pad": ["generic", "generic"],
+        "Conv2D2": ["generic", "generic"],
     }
 
     assert operations_distribution == expected
diff --git a/tests/python/driver/tvmc/test_compiler.py 
b/tests/python/driver/tvmc/test_compiler.py
index f624984481..6165099558 100644
--- a/tests/python/driver/tvmc/test_compiler.py
+++ b/tests/python/driver/tvmc/test_compiler.py
@@ -98,7 +98,7 @@ def test_save_dump_offloads_ethosu(tmp_path_factory):
             data_format="NHWC",
             dilations=1,
         )
-        op = tf.math.add(op1, op2)
+        op = tf.concat([op1, op2], 1)
         op = tf.pad(
             op,
             [[0, 0], [padding[0], padding_out[1]], [padding_out[2], 
padding[3]], [0, 0]],
@@ -163,9 +163,12 @@ def test_save_dump_offloads_ethosu(tmp_path_factory):
         r'ethos-u        <-          %5 = qnn.conv2d(%2, %v_param_5, -128, 0, 
0.11364f, meta[relay.Constant][4], padding=[1, 1, 1, 1], channels=3, 
kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO", 
out_dtype="int32")',
         r"ethos-u        <-          %6 = nn.bias_add(%5, %v_param_6, axis=3)",
         r'ethos-u        <-          %8 = qnn.requantize(%6, 
meta[relay.Constant][5], 0, 1.20538f, -128, axis=3, out_dtype="int8")',
-        r"ethos-u        <-     ethos-u.add",
-        r"ethos-u        <-          %9 = qnn.add(%7, %8, 1.56803f, -128, 
1.20538f, -128, 2.77341f, -128)",
-        r"generic        <-     nn.pad(%9, -128f, pad_width=[[0, 0], [1, 33], 
[33, 1], [0, 0]])",
+        r"                      %9 = (%7, %8)",
+        r"                      %10 = (1.59778f, 1.59778f)",
+        r"                      %11 = (-128, -128)",
+        r"ethos-u        <-     ethos-u.concat",
+        r"ethos-u        <-          %12 = qnn.concatenate(%9, %10, %11, 
1.59778f, -128, axis=1)",
+        r"generic        <-     nn.pad(%12, -128f, pad_width=[[0, 0], [1, 33], 
[33, 1], [0, 0]])",
     ]
 
     file_path = os.path.abspath(output_file_name)
@@ -174,7 +177,7 @@ def test_save_dump_offloads_ethosu(tmp_path_factory):
     with open(file_path, "r") as f:
         for i, file_string in enumerate(f):
             r_output = re.search(r"(.*)\(", file_string.strip(), re.DOTALL)
-            r_expected = re.search(r"(.*)\(", expected[i], re.DOTALL)
+            r_expected = re.search(r"(.*)\(", expected[i].strip(), re.DOTALL)
             # check that there is the same sequence of operations and 
composites,
             # combined with target names
             if r_output and r_expected:
@@ -202,6 +205,9 @@ def test_save_dump_offloads_cmsis(tmp_path_factory):
     def simple_net(x):
         weight_shape = [kernel_shape[0], kernel_shape[1], input_shape[3], 3]
         weights = tf.constant(np.random.uniform(size=weight_shape), 
dtype=tf.float32)
+        weight_shape[2] = 3
+        weights1 = tf.constant(np.random.uniform(size=weight_shape), 
dtype=tf.float32)
+        weights2 = tf.constant(np.random.uniform(size=weight_shape), 
dtype=tf.float32)
         op = tf.nn.conv2d(
             x,
             filters=weights,
@@ -210,22 +216,29 @@ def test_save_dump_offloads_cmsis(tmp_path_factory):
             data_format="NHWC",
             dilations=1,
         )
-        op = tf.nn.relu(op)
-        op = tf.pad(
+        op1 = tf.nn.conv2d(
             op,
-            [[0, 0], [padding[0], padding_out[2]], [padding_out[1], 
padding[3]], [0, 0]],
-            "CONSTANT",
+            filters=weights1,
+            strides=1,
+            padding="SAME",
+            data_format="NHWC",
+            dilations=1,
         )
-        op = tf.pad(
+        op2 = tf.nn.conv2d(
             op,
-            [[0, 0], [padding[0], padding[2]], [padding[1], padding[3]], [0, 
0]],
-            "CONSTANT",
+            filters=weights2,
+            strides=1,
+            padding="SAME",
+            data_format="NHWC",
+            dilations=1,
         )
-        return tf.pad(
+        op = tf.concat([op1, op2], 1)
+        op = tf.pad(
             op,
-            [[0, 0], [padding_out[0], padding[2]], [padding[1], 
padding_out[3]], [0, 0]],
+            [[0, 0], [padding[0], padding_out[1]], [padding_out[2], 
padding[3]], [0, 0]],
             "CONSTANT",
         )
+        return op
 
     from tests.python.contrib.test_ethosu.infra import get_tflite_graph
 
@@ -265,18 +278,27 @@ def test_save_dump_offloads_cmsis(tmp_path_factory):
 
     expected = [
         r"Total number of operators and distribution by targets",
-        r"Total: 7",
-        r"cmsis-nn: 4",
-        r"generic: 3",
+        r"Total: 11",
+        r"cmsis-nn: 9",
+        r"generic: 2",
         r"",
         r"cmsis-nn       <-     cmsis-nn.qnn_conv2d",
         r'cmsis-nn       <-          %0 = qnn.conv2d(%x, %v_param_1, -128, 0, 
0.00392157f, meta[relay.Constant][0], padding=[1, 1, 1, 1], channels=3, 
kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO", 
out_dtype="int32")',
         r"cmsis-nn       <-          %1 = nn.bias_add(%0, %v_param_2, axis=3)",
-        r'cmsis-nn       <-          %2 = qnn.requantize(%1, 
meta[relay.Constant][1], 0, 0.113405f, -128, axis=3, out_dtype="int8")',
-        r"cmsis-nn       <-          %3 = clip(%2, a_min=-128f, a_max=127f)",
-        r"generic        <-     %4 = nn.pad(%3, -128f, pad_width=[[0, 0], [1, 
33], [33, 1], [0, 0]])",
-        r"generic        <-     %5 = nn.pad(%4, -128f, pad_width=[[0, 0], [1, 
1], [1, 1], [0, 0]])",
-        r"generic        <-     nn.pad(%5, -128f, pad_width=[[0, 0], [1, 1], 
[1, 1], [0, 0]])",
+        r'cmsis-nn       <-          %2 = qnn.requantize(%1, 
meta[relay.Constant][1], 0, 0.115114f, -128, axis=3, out_dtype="int8")',
+        r"cmsis-nn       <-     cmsis-nn.qnn_conv2d",
+        r'cmsis-nn       <-          %3 = qnn.conv2d(%2, %v_param_3, -128, 0, 
0.115114f, meta[relay.Constant][2], padding=[1, 1, 1, 1], channels=3, 
kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO", 
out_dtype="int32")',
+        r"cmsis-nn       <-          %4 = nn.bias_add(%3, %v_param_4, axis=3)",
+        r'cmsis-nn       <-          %7 = qnn.requantize(%4, 
meta[relay.Constant][3], 0, 1.59328f, -128, axis=3, out_dtype="int8")',
+        r"cmsis-nn       <-     cmsis-nn.qnn_conv2d",
+        r'cmsis-nn       <-          %5 = qnn.conv2d(%2, %v_param_5, -128, 0, 
0.115114f, meta[relay.Constant][4], padding=[1, 1, 1, 1], channels=3, 
kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO", 
out_dtype="int32")',
+        r"cmsis-nn       <-          %6 = nn.bias_add(%5, %v_param_6, axis=3)",
+        r'cmsis-nn       <-          %8 = qnn.requantize(%6, 
meta[relay.Constant][5], 0, 1.59328f, -128, axis=3, out_dtype="int8")',
+        r"                      %9 = (%7, %8)",
+        r"                      %10 = (1.59328f, 1.59328f)",
+        r"                      %11 = (-128, -128)",
+        r"generic        <-     %12 = qnn.concatenate(%9, %10, %11, 1.59328f, 
-128, axis=1)",
+        r"generic        <-     nn.pad(%12, -128f, pad_width=[[0, 0], [1, 33], 
[33, 1], [0, 0]])",
     ]
 
     file_path = os.path.abspath(output_file_name)
@@ -284,14 +306,14 @@ def test_save_dump_offloads_cmsis(tmp_path_factory):
     assert os.path.exists(file_path)
     with open(file_path, "r") as f:
         for i, file_string in enumerate(f):
-            r_output = re.search(r"(.*)\(", file_string.strip(), re.DOTALL)
+            r_output = re.search(r"(.*)\(", file_string.replace("\n", ""), 
re.DOTALL)
             r_expected = re.search(r"(.*)\(", expected[i], re.DOTALL)
             # check that there is the same sequence of operations and 
composites,
             # combined with target names
             if r_output and r_expected:
                 assert r_output.group(0) == r_expected.group(0)
             else:
-                assert r_output == r_expected
+                assert file_string.replace("\n", "") == expected[i]
 
 
 def test_save_dump_offloads_generic(tmp_path_factory):
@@ -313,6 +335,9 @@ def test_save_dump_offloads_generic(tmp_path_factory):
     def simple_net(x):
         weight_shape = [kernel_shape[0], kernel_shape[1], input_shape[3], 3]
         weights = tf.constant(np.random.uniform(size=weight_shape), 
dtype=tf.float32)
+        weight_shape[2] = 3
+        weights1 = tf.constant(np.random.uniform(size=weight_shape), 
dtype=tf.float32)
+        weights2 = tf.constant(np.random.uniform(size=weight_shape), 
dtype=tf.float32)
         op = tf.nn.conv2d(
             x,
             filters=weights,
@@ -321,21 +346,29 @@ def test_save_dump_offloads_generic(tmp_path_factory):
             data_format="NHWC",
             dilations=1,
         )
-        op = tf.pad(
+        op1 = tf.nn.conv2d(
             op,
-            [[0, 0], [padding[0], padding_out[2]], [padding_out[1], 
padding[3]], [0, 0]],
-            "CONSTANT",
+            filters=weights1,
+            strides=1,
+            padding="SAME",
+            data_format="NHWC",
+            dilations=1,
         )
-        op = tf.pad(
+        op2 = tf.nn.conv2d(
             op,
-            [[0, 0], [padding[0], padding[2]], [padding[1], padding[3]], [0, 
0]],
-            "CONSTANT",
+            filters=weights2,
+            strides=1,
+            padding="SAME",
+            data_format="NHWC",
+            dilations=1,
         )
-        return tf.pad(
+        op = tf.concat([op1, op2], 1)
+        op = tf.pad(
             op,
-            [[0, 0], [padding_out[0], padding[2]], [padding[1], 
padding_out[3]], [0, 0]],
+            [[0, 0], [padding[0], padding_out[1]], [padding_out[2], 
padding[3]], [0, 0]],
             "CONSTANT",
         )
+        return op
 
     from tests.python.contrib.test_ethosu.infra import get_tflite_graph
 
@@ -374,15 +407,23 @@ def test_save_dump_offloads_generic(tmp_path_factory):
 
     expected = [
         r"Total number of operators and distribution by targets",
-        r"Total: 6",
-        r"generic: 6",
+        r"Total: 11",
+        r"generic: 11",
         r"",
-        r'generic        <-     %0 = qnn.conv2d(%x, %v_param_1, -128, 0, 
0.00392156f, meta[relay.Constant][0], padding=[1, 1, 1, 1], channels=3, 
kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO", 
out_dtype="int32")',
+        r'generic        <-     %0 = qnn.conv2d(%x, %v_param_1, -128, 0, 
0.00392157f, meta[relay.Constant][0], padding=[1, 1, 1, 1], channels=3, 
kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO", 
out_dtype="int32")',
         r"generic        <-     %1 = nn.bias_add(%0, %v_param_2, axis=3)",
-        r'generic        <-     %2 = qnn.requantize(%1, 
meta[relay.Constant][1], 0, 0.103975f, -128, axis=3, out_dtype="int8")',
-        r"generic        <-     %3 = nn.pad(%2, -128f, pad_width=[[0, 0], [1, 
33], [33, 1], [0, 0]])",
-        r"generic        <-     %4 = nn.pad(%3, -128f, pad_width=[[0, 0], [1, 
1], [1, 1], [0, 0]])",
-        r"generic        <-     nn.pad(%4, -128f, pad_width=[[0, 0], [1, 1], 
[1, 1], [0, 0]])",
+        r'generic        <-     %2 = qnn.requantize(%1, 
meta[relay.Constant][1], 0, 0.109484f, -128, axis=3, out_dtype="int8")',
+        r'generic        <-     %3 = qnn.conv2d(%2, %v_param_3, -128, 0, 
0.109484f, meta[relay.Constant][2], padding=[1, 1, 1, 1], channels=3, 
kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO", 
out_dtype="int32")',
+        r"generic        <-     %4 = nn.bias_add(%3, %v_param_4, axis=3)",
+        r'generic        <-     %5 = qnn.conv2d(%2, %v_param_5, -128, 0, 
0.109484f, meta[relay.Constant][4], padding=[1, 1, 1, 1], channels=3, 
kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO", 
out_dtype="int32")',
+        r"generic        <-     %6 = nn.bias_add(%5, %v_param_6, axis=3)",
+        r'generic        <-     %7 = qnn.requantize(%4, 
meta[relay.Constant][3], 0, 1.45572f, -128, axis=3, out_dtype="int8")',
+        r'generic        <-     %8 = qnn.requantize(%6, 
meta[relay.Constant][5], 0, 1.45572f, -128, axis=3, out_dtype="int8")',
+        r"                      %9 = (%7, %8)",
+        r"                      %10 = (1.45572f, 1.45572f)",
+        r"                      %11 = (-128, -128)",
+        r"generic        <-     %12 = qnn.concatenate(%9, %10, %11, 1.45572f, 
-128, axis=1)",
+        r"generic        <-     nn.pad(%12, -128f, pad_width=[[0, 0], [1, 33], 
[33, 1], [0, 0]])",
     ]
 
     file_path = os.path.abspath(output_file_name)
@@ -390,14 +431,14 @@ def test_save_dump_offloads_generic(tmp_path_factory):
     assert os.path.exists(file_path)
     with open(file_path, "r") as f:
         for i, file_string in enumerate(f):
-            r_output = re.search(r"(.*)\(", file_string.strip(), re.DOTALL)
+            r_output = re.search(r"(.*)\(", file_string.replace("\n", ""), 
re.DOTALL)
             r_expected = re.search(r"(.*)\(", expected[i], re.DOTALL)
             # check that there is the same sequence of operations and 
composites,
             # combined with target names
             if r_output and r_expected:
                 assert r_output.group(0) == r_expected.group(0)
             else:
-                assert r_output == r_expected
+                assert file_string.replace("\n", "") == expected[i]
 
 
 # End to end tests for compilation

Reply via email to