aboutsummaryrefslogtreecommitdiff
path: root/ethosu/vela/graph_optimiser.py
diff options
context:
space:
mode:
authorPatrik Gustavsson <patrik.gustavsson@arm.com>2021-01-14 10:55:43 +0100
committerpatrik.gustavsson <patrik.gustavsson@arm.com>2021-01-14 15:42:45 +0000
commit224e99bd70a443e345d3ea454aedc51bf46cf261 (patch)
tree52dc8ba33017af8eca82f257f6c893551923dac1 /ethosu/vela/graph_optimiser.py
parent49b4f1c88d2825d8708c38092fa71103d5931d8c (diff)
downloadethos-u-vela-224e99bd70a443e345d3ea454aedc51bf46cf261.tar.gz
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
Diffstat (limited to 'ethosu/vela/graph_optimiser.py')
-rw-r--r--ethosu/vela/graph_optimiser.py6
1 files changed, 4 insertions, 2 deletions
diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py
index 6d2696c4..511ac954 100644
--- a/ethosu/vela/graph_optimiser.py
+++ b/ethosu/vela/graph_optimiser.py
@@ -106,7 +106,7 @@ def rewrite_concat(tens, arch, nng):
def rewrite_split(tens, arch, nng):
- if len(tens.ops) == 1 and tens.ops[0].type.is_split_op():
+ if len(tens.ops) == 1 and tens.ops[0].type.is_split_op() and tens.ops[0].type != Op.Unpack:
split_op = tens.ops[0]
# Not supported so leave it and run on CPU
@@ -125,6 +125,7 @@ def rewrite_split(tens, arch, nng):
# Get the start and end of the split
offset_start = [0] * 4
for idx, out in enumerate(outputs):
+ split_op.ofm_shapes[idx] = Shape4D(out.shape)
if out == tens:
break
if axis >= 0:
@@ -143,7 +144,8 @@ def rewrite_split(tens, arch, nng):
new_op.attrs["split_start"] = offset_start
new_op.run_on_npu = True
new_op.set_output_tensor(tens)
- new_op.set_ifm_ofm_shapes()
+ new_op.ifm_shapes.append(Shape4D(inp.shape))
+ new_op.ofm_shapes.append(Shape4D(full_shape(4, tens.shape, 1)))
DebugDatabase.add_optimised(split_op, new_op)
return tens