commit 2349d429d926e258e9a61d34c7fd97660ab9fb98
author Patrik Gustavsson <patrik.gustavsson@arm.com> Tue Dec 01 16:02:29 2020 +0100
committer Patrik Gustavsson <patrik.gustavsson@arm.com> Fri Dec 18 16:33:32 2020 +0100
tree b5151d0f12428e47d64b1fb2ce4f2f8c19304a0d
parent 528a56df829b65f7a2c61953650b123c461095f7

MLBEDSW-3654 Add/use op ifm/ofm shapes

Add ifm/ofm shapes to op
Changed to rely on these shapes

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

ethosu/vela/pass_packing.py

diff --git a/ethosu/vela/pass_packing.py b/ethosu/vela/pass_packing.py
index 9bc04f2..095a78d 100644
--- a/ethosu/vela/pass_packing.py
+++ b/ethosu/vela/pass_packing.py
@@ -397,11 +397,28 @@
 
             if len(ps.inputs) > 2:
                 ps.ifm_tensor = ps.inputs[-2]
+
+            # Get the corresponding ifm_shapes
+            for op in input_ops_list + [primary_op]:
+                if ps.ifm_tensor == op.ifm:
+                    ps.ifm_shapes.append(op.ifm_shapes[0])
+                elif ps.ifm_tensor == op.ifm2:
+                    ps.ifm_shapes.append(op.ifm_shapes[1])
+            for op in input_ops_list + [primary_op]:
+                if ps.ifm2_tensor == op.ifm:
+                    ps.ifm_shapes.append(op.ifm_shapes[0])
+                elif ps.ifm2_tensor == op.ifm2:
+                    ps.ifm_shapes.append(op.ifm_shapes[1])
         else:
             ps.ifm_tensor = ifm_tensor
             ps.ifm2_tensor = None
+            if ps.primary_op is not None:
+                ps.ifm_shapes.append(ps.primary_op.ifm_shapes[0])
 
         ps.ofm_tensor = ofm_tensor
+        if ps.primary_op is not None:
+            ps.ofm_shapes.append(ps.primary_op.ofm_shapes[0])
+
         assert ps.placement != PassPlacement.Npu or ps.ofm_tensor is not None
         ps.weight_tensor = ps.get_primary_op_ifm_weights()[1]
         ps.scale_tensor = ps.get_primary_op_ifm_weights_biases_ofm()[2]
@@ -436,6 +453,8 @@
             avgpool_out = inp.clone("_avgpooled")
             avgpool_out.consumer_list.append(op)
             avgpool_op.set_output_tensor(avgpool_out)
+            avgpool_op.ifm_shapes = op.ifm_shapes
+            avgpool_op.ofm_shapes = op.ofm_shapes
 
             op.inputs[0] = avgpool_out
             op_list.insert(0, avgpool_op)