Fix review notes

Author: arina-grovety

python/tvm/relay/backend/contrib/ethosu/codegen.py

@@ -32,7 +32,8 @@
 )
 from tvm.relay.backend.contrib.ethosu.legalize import LegalizeEthosU
 from tvm.relay.backend.contrib.ethosu import tir_to_cs_translator, util, vela_api
-from tvm.relay.expr_functor import ExprMutator, ExprVisitor
+from tvm.relay.expr_functor import ExprMutator, ExprVisitor, Call
+from tvm.relay import expr as _expr
 
 # pylint: disable=unused-import
 from tvm.relay.backend.contrib.ethosu.op import op_attrs

python/tvm/relay/op/contrib/ethosu.py

@@ -2341,7 +2341,7 @@ def partition_for_ethosu(
     mod : IRModule
         The partitioned IRModule with external global functions
     """
-    from tvm.relay.backend.contrib.ethosu import preprocess
+    from tvm.relay.backend.contrib.ethosu import preprocess, codegen
 
     if params:
         mod["main"] = bind_params_by_name(mod["main"], params)

python/tvm/relay/transform/__init__.py

@@ -18,7 +18,6 @@
 """The Relay IR namespace containing transformations."""
 # transformation passes
 from .transform import *
-from .replicate_pads_with_multiple_consumers import *
 from .recast import recast
 from . import fake_quantization_to_integer, mixed_precision
 from .flexible_shape import FlexibleShapeDispatch

python/tvm/relay/transform/replicate_pads_with_multiple_consumers.py

@@ -157,6 +157,76 @@ def conv2d_double(x):
     infra.compare_tvm_with_tflite(conv2d_double, [ifm_shape], accel_type)
 
 
+@pytest.mark.parametrize("accel_type", ACCEL_TYPES)
+@pytest.mark.parametrize("ifm_shape", [(1, 55, 32, 3)])
+@pytest.mark.parametrize(
+    "kernel_shape, activation_function",
+    [((3, 3), "RELU"), ((1, 2), "NONE")],
+)
+@pytest.mark.parametrize("strides, dilation", [((3, 2), (1, 1))])
+@pytest.mark.parametrize("op_padding", ["SAME", "VALID"])
+@pytest.mark.parametrize("sep_padding", [(0, 0, 1, 1), (7, 5, 4, 5)])
+@pytest.mark.parametrize(
+    "op_pairs", [("conv2d", "conv2d"), ("depthwise", "depthwise"), ("conv2d", "depthwise")]
+)
+def test_tflite_shared_pad(
+    accel_type,
+    ifm_shape,
+    kernel_shape,
+    activation_function,
+    strides,
+    dilation,
+    op_padding,
+    sep_padding,
+    op_pairs,
+):
+    np.random.seed(0)
+
+    @tf.function
+    def tf_function(x):
+        def make_depthwise_or_conv2d(pair_idx, x):
+            # The input strides to the TensorFlow API needs to be of shape 1x4
+            tf_strides = [1, strides[0], strides[1], 1]
+            if op_pairs[pair_idx] == "depthwise":
+                weight_shape = [kernel_shape[0], kernel_shape[1], ifm_shape[3], 1]
+                weight = tf.constant(np.random.uniform(size=weight_shape), dtype=tf.float32)
+                op =  tf.nn.depthwise_conv2d(
+                    x, weight, strides=tf_strides, padding=op_padding, dilations=dilation
+                )
+            else:
+                weight_shape = [kernel_shape[0], kernel_shape[1], ifm_shape[3], 3]
+                weight = tf.constant(np.random.uniform(size=weight_shape), dtype=tf.float32)
+                op =  tf.nn.conv2d(
+                    x,
+                    weight,
+                    strides=tf_strides,
+                    padding=op_padding,
+                    dilations=dilation,
+                )
+            if activation_function == "RELU":
+                op = tf.nn.relu(op)
+            return op
+
+        x = tf.pad(
+            x,
+            [
+                [0, 0],
+                [sep_padding[0], sep_padding[2]],
+                [sep_padding[1], sep_padding[3]],
+                [0, 0],
+            ],
+            "CONSTANT",
+        )
+
+        x1 = make_depthwise_or_conv2d(0, x)
+        x2 = make_depthwise_or_conv2d(1, x)
+
+        x3 = tf.math.add(x1, x2)
+        return x3
+
+    infra.compare_tvm_with_tflite(tf_function, [ifm_shape], accel_type)
+
+
 @pytest.mark.parametrize("weight_min, weight_max", [(0.0, 1e-11), (-1e10, 1e10)])
 def test_out_of_range_scaling(weight_min, weight_max):
     np.random.seed(0)

tests/python/contrib/test_ethosu/test_codegen.py

@@ -3671,6 +3671,22 @@ def create_tflite_graph():
         class Model(tf.Module):
             @tf.function
             def tf_function(self, x):
+                def make_depthwise_or_conv2d(pair_idx):
+                    if op_pairs[pair_idx] == "depthwise":
+                        weight_shape = [kernel_shape[0], kernel_shape[1], ifm_shape[3], 1]
+                        weight = tf.constant(np.random.uniform(size=weight_shape), dtype=tf.float32)
+                        return tf.nn.depthwise_conv2d(
+                            x, weight, strides=tf_strides, padding=op_padding, dilations=dilation
+                        )
+                    weight_shape = [kernel_shape[0], kernel_shape[1], ifm_shape[3], 3]
+                    weight = tf.constant(np.random.uniform(size=weight_shape), dtype=tf.float32)
+                    return tf.nn.conv2d(
+                        x,
+                        weight,
+                        strides=tf_strides,
+                        padding=op_padding,
+                        dilations=dilation,
+                    )
 
                 x = tf.pad(
                     x,
@@ -3686,39 +3702,8 @@ def tf_function(self, x):
                 # The input strides to the TensorFlow API needs to be of shape 1x4
                 tf_strides = [1, strides[0], strides[1], 1]
 
-                if op_pairs[0] == "depthwise":
-                    weight_shape = [kernel_shape[0], kernel_shape[1], ifm_shape[3], 1]
-                    weight = tf.constant(np.random.uniform(size=weight_shape), dtype=tf.float32)
-                    x1 = tf.nn.depthwise_conv2d(
-                        x, weight, strides=tf_strides, padding=op_padding, dilations=dilation
-                    )
-                else:
-                    weight_shape = [kernel_shape[0], kernel_shape[1], ifm_shape[3], 3]
-                    weight = tf.constant(np.random.uniform(size=weight_shape), dtype=tf.float32)
-                    x1 = tf.nn.conv2d(
-                        x,
-                        weight,
-                        strides=tf_strides,
-                        padding=op_padding,
-                        dilations=dilation,
-                    )
-
-                if op_pairs[1] == "depthwise":
-                    weight_shape = [kernel_shape[0], kernel_shape[1], ifm_shape[3], 1]
-                    weight = tf.constant(np.random.uniform(size=weight_shape), dtype=tf.float32)
-                    x2 = tf.nn.depthwise_conv2d(
-                        x, weight, strides=tf_strides, padding=op_padding, dilations=dilation
-                    )
-                else:
-                    weight_shape = [kernel_shape[0], kernel_shape[1], ifm_shape[3], 3]
-                    weight = tf.constant(np.random.uniform(size=weight_shape), dtype=tf.float32)
-                    x2 = tf.nn.conv2d(
-                        x,
-                        weight,
-                        strides=tf_strides,
-                        padding=op_padding,
-                        dilations=dilation,
-                    )
+                x1 = make_depthwise_or_conv2d(0)
+                x2 = make_depthwise_or_conv2d(1)
 
                 x3 = tf.math.add(x1, x2)
                 return x3