Handle sg input and output for Squeeze operator3.1.0.rc2

Update to handle the case when the Squeeze Op ifm/ofm are the
subgraph ifm/ofm, to facilitate the removal of the Squeeze Op.
Adding NOP to maintain the original tensors.
Updated pytests for squeeze operator.

Signed-off-by: Jonas Ohlsson <jonas.ohlsson@arm.com>
Change-Id: I623cae05e696fb16ccf29dedc42fd822601e9fd9

2 files changed, 90 insertions, 54 deletions

diff --git a/ethosu/vela/test/test_graph_optimiser.py b/ethosu/vela/test/test_graph_optimiser.py
index e0eedd6..2f96572 100644
--- a/ethosu/vela/test/test_graph_optimiser.py
+++ b/ethosu/vela/test/test_graph_optimiser.py
@@ -323,62 +323,71 @@ def test_pad_followed_by_avg_pool(k_size, padding, expect_pad_removed):
         assert pool_op.attrs["padding"] == Padding.VALID
 
 
+# Setup network to test removal of op with op_type Op.Reshape or Op.Squeeze
+# op_type should be Op.Reshape or Op.Squeeze
+def setup_network(op_type):
+    assert op_type == Op.Reshape or op_type == Op.Squeeze
+    if op_type == Op.Reshape:
+        op_str = "reshape"
+    elif op_type == Op.Squeeze:
+        op_str = "squeeze"
+
+    quant = testutil.default_quant_params()
+    # create reshape1 op
+    ifm_shape = [64, 16]
+    reshape1_ofm_shape = [1, 4, 16, 16]
+    reshape1_ifm = create_const_tensor(f"{op_str}1_in", ifm_shape, DataType.uint8, np.zeros(ifm_shape))
+    reshape1_ifm.quantization = quant
+    reshape1_ofm = create_const_tensor(
+        f"{op_str}1_out", reshape1_ofm_shape, DataType.uint8, np.zeros(reshape1_ofm_shape)
+    )
+    reshape1_ofm.quantization = quant
+    shape_tens = create_const_tensor(f"{op_str}1_shape", [1], DataType.int32, reshape1_ofm_shape)
+    reshape1_op = testutil.create_op(op_type, [reshape1_ifm, shape_tens], reshape1_ofm, set_ifm_ofm_shapes=False)
+    reshape1_op.attrs["new_shape"] = reshape1_ofm_shape
+    reshape1_op.run_on_npu = True
+
+    # create conv op
+    conv_ofm = Tensor([1, 8, 8, 16], DataType.uint8, "output")
+    conv_ofm.quantization = quant.clone()
+    weight_tens = Tensor([1, 1, 16, 16], DataType.uint8, "weights")
+    weight_tens.values = np.zeros(weight_tens.shape, np.uint8)
+    weight_tens.quantization = quant.clone()
+    bias_tens = Tensor([16], DataType.int32, "biases")
+
+    attrs = {"padding": Padding.SAME, "stride_w": 1, "stride_h": 1, "dilation_w_factor": 1, "dilation_h_factor": 1}
+    attrs["strides"] = (1, attrs["stride_h"], attrs["stride_w"], 1)
+
+    conv2d_op = testutil.create_op(
+        Op.Conv2D, [reshape1_ofm, weight_tens, bias_tens], conv_ofm, attrs=attrs, set_ifm_ofm_shapes=False
+    )
+    conv2d_op.run_on_npu = True
+
+    # create reshape2 op
+    ofm_shape = [8, 8, 16]
+    reshape2_ofm = create_const_tensor(f"{op_str}2_out", ofm_shape, DataType.uint8, np.zeros(ofm_shape))
+    reshape2_ofm.quantization = quant
+    shape_tens = create_const_tensor(f"{op_str}2_shape", [1], DataType.int32, ofm_shape)
+    reshape2_op = testutil.create_op(op_type, [conv_ofm, shape_tens], reshape2_ofm, set_ifm_ofm_shapes=False)
+    reshape2_op.attrs["new_shape"] = ofm_shape
+    reshape2_op.run_on_npu = True
+    nng = Graph()
+    sg = testutil.create_subgraph([reshape1_op, conv2d_op, reshape2_op])
+    nng.subgraphs.append(sg)
+
+    return nng, reshape1_op, conv2d_op, reshape2_op
+
+
 def test_remove_reshape():
     """
     Tests that the expected reshape are removed in graph_optimisation
     """
 
-    def setup_network():
-        quant = testutil.default_quant_params()
-        # create reshape1 op
-        ifm_shape = [64, 16]
-        reshape1_ofm_shape = [1, 4, 16, 16]
-        reshape1_ifm = create_const_tensor("reshape1_in", ifm_shape, DataType.uint8, np.zeros(ifm_shape))
-        reshape1_ifm.quantization = quant
-        reshape1_ofm = create_const_tensor(
-            "reshape1_out", reshape1_ofm_shape, DataType.uint8, np.zeros(reshape1_ofm_shape)
-        )
-        reshape1_ofm.quantization = quant
-        shape_tens = create_const_tensor("reshape1_shape", [1], DataType.int32, reshape1_ofm_shape)
-        reshape1_op = testutil.create_op(Op.Reshape, [reshape1_ifm, shape_tens], reshape1_ofm, set_ifm_ofm_shapes=False)
-        reshape1_op.attrs["new_shape"] = reshape1_ofm_shape
-        reshape1_op.run_on_npu = True
-
-        # create conv op
-        conv_ofm = Tensor([1, 8, 8, 16], DataType.uint8, "output")
-        conv_ofm.quantization = quant.clone()
-        weight_tens = Tensor([1, 1, 16, 16], DataType.uint8, "weights")
-        weight_tens.values = np.zeros(weight_tens.shape, np.uint8)
-        weight_tens.quantization = quant.clone()
-        bias_tens = Tensor([16], DataType.int32, "biases")
-
-        attrs = {"padding": Padding.SAME, "stride_w": 1, "stride_h": 1, "dilation_w_factor": 1, "dilation_h_factor": 1}
-        attrs["strides"] = (1, attrs["stride_h"], attrs["stride_w"], 1)
-
-        conv2d_op = testutil.create_op(
-            Op.Conv2D, [reshape1_ofm, weight_tens, bias_tens], conv_ofm, attrs=attrs, set_ifm_ofm_shapes=False
-        )
-        conv2d_op.run_on_npu = True
-
-        # create reshape2 op
-        ofm_shape = [8, 8, 16]
-        reshape2_ofm = create_const_tensor("reshape2_out", ofm_shape, DataType.uint8, np.zeros(ofm_shape))
-        reshape2_ofm.quantization = quant
-        shape_tens = create_const_tensor("reshape2_shape", [1], DataType.int32, ofm_shape)
-        reshape2_op = testutil.create_op(Op.Reshape, [conv_ofm, shape_tens], reshape2_ofm, set_ifm_ofm_shapes=False)
-        reshape2_op.attrs["new_shape"] = ofm_shape
-        reshape2_op.run_on_npu = True
-        nng = Graph()
-        sg = testutil.create_subgraph([reshape1_op, conv2d_op, reshape2_op])
-        nng.subgraphs.append(sg)
-
-        return nng, reshape1_op, conv2d_op, reshape2_op
-
     # Test1 no Reshape op is expected to remain in the NPU subgrapgh
     # but first one will be put on CPU
     # Network is Reshape-Conv-Reshape
     # Result is Conv
-    nng, reshape1_op, conv2d_op, reshape2_op = setup_network()
+    nng, reshape1_op, conv2d_op, reshape2_op = setup_network(Op.Reshape)
     arch = testutil.create_arch()
     assert verify_graph_health(nng)
     nng = optimise_graph(nng, arch, NetworkType.TFLite)
@@ -387,10 +396,37 @@ def test_remove_reshape():
     # Test2 reshape1 with different quantisation, this Reshape op is expected to remain
     # Network is Reshape-Conv-Reshape
     # expected is Reshape-Conv
-    nng, reshape1_op, conv2d_op, reshape2_op = setup_network()
+    nng, reshape1_op, conv2d_op, reshape2_op = setup_network(Op.Reshape)
     quant_zp32 = testutil.default_quant_params()
     quant_zp32.zero_point = 32
     reshape1_op.ofm.quantization = quant_zp32
     assert verify_graph_health(nng)
     nng = optimise_graph(nng, arch, NetworkType.TFLite)
     assert verify_graph_health(nng)
+
+
+def test_remove_squeeze():
+    """
+    Tests that the expected squeeze are removed in graph_optimisation
+    """
+
+    # Test1 no Squeeze op is expected to remain in the NPU subgrapgh
+    # but first one will be put on CPU
+    # Network is Squeeze-Conv-Squeeze
+    # Result is Conv
+    nng, squeeze1_op, conv2d_op, squeeze2_op = setup_network(Op.Squeeze)
+    arch = testutil.create_arch()
+    assert verify_graph_health(nng)
+    nng = optimise_graph(nng, arch, NetworkType.TFLite)
+    assert verify_graph_health(nng)
+
+    # Test2 squeeze1 with different quantisation, this Squeeze op is expected to remain
+    # Network is Squeeze-Conv-Squeeze
+    # expected is Squeeze-Conv
+    nng, squeeze1_op, conv2d_op, squeeze2_op = setup_network(Op.Squeeze)
+    quant_zp32 = testutil.default_quant_params()
+    quant_zp32.zero_point = 32
+    squeeze1_op.ofm.quantization = quant_zp32
+    assert verify_graph_health(nng)
+    nng = optimise_graph(nng, arch, NetworkType.TFLite)
+    assert verify_graph_health(nng)
diff --git a/ethosu/vela/tflite_graph_optimiser.py b/ethosu/vela/tflite_graph_optimiser.py
index 6c85bb4..ff2f5a0 100644
--- a/ethosu/vela/tflite_graph_optimiser.py
+++ b/ethosu/vela/tflite_graph_optimiser.py
@@ -185,7 +185,7 @@ def remove_SplitSliceRead(op, arch):
             len(op.ofm.consumer_list) == 1
             and op.ofm.consumer_list[0] is not None
             and op.ofm.consumer_list[0].run_on_npu
-            and op.ofm.consumer_list[0].type != Op.Reshape
+            and op.ofm.consumer_list[0].type not in (Op.Reshape, Op.Squeeze)
             and op.ofm_shapes[0] == Shape4D.from_list(op.ofm.shape)
         ):
             # SplitSliceRead can be performed by tensor consumer
@@ -245,10 +245,10 @@ def insert_copy_op_after_tens(tens):
 
 
 def fix_sg_input_output(op, arch, nng):
-    if not op.run_on_npu or op.type != Op.Reshape:
+    if not op.run_on_npu or op.type not in (Op.Reshape, Op.Squeeze):
         return op
 
-    # For the Reshape operators we want to remove, tensors are removed.
+    # For the Reshape/Squeeze operators we want to remove, tensors are removed.
     # But in order to to do this, they cannot be outputs of the sg,
     # this need to be fixed prior to the removal.
     # Solution is to add a avgpool NOP, to maintain the original tensor.
@@ -259,12 +259,12 @@ def fix_sg_input_output(op, arch, nng):
     ifm_is_sg_ifm = op.ifm.ops[0].type in (Op.Placeholder, Op.SubgraphInput, Op.Const)
     ifm_is_sg_ofm = any(ifm_cons is None for ifm_cons in op.ifm.consumer_list)
     ofm_is_sg_ofm = any(ofm_cons is None for ofm_cons in op.ofm.consumer_list)
-    # Check if ifm/ofm is produced repectivly consumed by CPU
+    # Check if ifm/ofm is produced respectively consumed by CPU
     ifm_is_cpu_produced = any(ifm_prod is not None and not ifm_prod.run_on_npu for ifm_prod in op.ifm.ops)
     ofm_is_cpu_consumed = any(ofm_cons is not None and not ofm_cons.run_on_npu for ofm_cons in op.ofm.consumer_list)
 
     if (ifm_is_sg_ofm or ifm_is_sg_ifm or ifm_is_cpu_produced) and (ofm_is_sg_ofm or ofm_is_cpu_consumed):
-        # Both ifm and ofm need to persist, but only ifm need a copy, in order to remove the Reshape
+        # Both ifm and ofm need to persist, but only ifm need a copy, in order to remove the Reshape/Squeeze
         insert_copy_op_after_tens(op.ifm)
 
     return op
@@ -1062,7 +1062,7 @@ def convert_tanh_sigmoid_to_lut(op, arch, nng):
 
 
 def remove_reshape_and_squeeze_ops(op, arch):
-    if op.run_on_npu and (op.type == Op.Reshape or op.type == Op.Squeeze):
+    if op.run_on_npu and op.type in (Op.Reshape, Op.Squeeze):
         ofm = op.ofm
         ifm = op.ifm