commit 224e99bd70a443e345d3ea454aedc51bf46cf261
author Patrik Gustavsson <patrik.gustavsson@arm.com> Thu Jan 14 10:55:43 2021 +0100
committer patrik.gustavsson <patrik.gustavsson@arm.com> Thu Jan 14 15:42:45 2021 +0000
tree 52dc8ba33017af8eca82f257f6c893551923dac1
parent 49b4f1c88d2825d8708c38092fa71103d5931d8c

MLBEDSW-3654 Fix for split/concat ops

Fix for split/concat ops
- set correct ifm_shapes in pass packing

Signed-off-by: Patrik Gustavsson <patrik.gustavsson@arm.com>
Change-Id: I7373b1743e4511b6c1dfaa398b927fbb1b454f60

ethosu/vela/pass_packing.py

diff --git a/ethosu/vela/pass_packing.py b/ethosu/vela/pass_packing.py
index 7abf3b2..ee0d712 100644
--- a/ethosu/vela/pass_packing.py
+++ b/ethosu/vela/pass_packing.py
@@ -244,6 +244,7 @@
         input_set = set()
         ifm_tensor = None
         primary_op = None
+        ifm_shapes = None
 
         to_process = collections.deque()
         for start_op in start_ops_to_process:
@@ -279,6 +280,7 @@
                             ):
                                 assert len(curr_op.inputs) >= 1
                                 ifm_tensor = curr_op.ifm
+                                ifm_shapes = curr_op.ifm_shapes.copy()
                                 assert ifm_tensor is not None, "IFM missing in {}".format(curr_op)
                                 assert ifm_tensor.purpose == TensorPurpose.FeatureMap
 
@@ -417,7 +419,7 @@
             ps.ifm_tensor = ifm_tensor
             ps.ifm2_tensor = None
             if ps.primary_op is not None and ps.primary_op.run_on_npu:
-                ps.ifm_shapes.append(ps.primary_op.ifm_shapes[0])
+                ps.ifm_shapes.append(ifm_shapes[0])
 
         ps.ofm_tensor = ofm_tensor
         ps.ofm_shapes.append(ofm_shape)