MLBEDSW-2379: Fix 1-element tensors that were marked as scalars

Tensors that are the result of an operation were incorrectly marked
as scalars.

Also fixes a bug for IFM2 of shape [*,*,*,1] in elementwise operations.

Signed-off-by: Louis Verhaard <louis.verhaard@arm.com>
Change-Id: I82a0e643b12e93c7158e4aca3185415c59033a73

2 files changed, 14 insertions, 2 deletions

diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py
index ca8b89fc..8a393a2e 100644
--- a/ethosu/vela/graph_optimiser.py
+++ b/ethosu/vela/graph_optimiser.py
@@ -24,10 +24,10 @@ from . import rewrite_graph
 from .data_type import DataType
 from .errors import UnsupportedFeatureError
 from .ethos_u55_regs.ethos_u55_regs import resampling_mode
+from .numeric_util import full_shape
 from .operation import NpuBlockType
 from .operation import Operation
 from .tensor import Tensor
-from .numeric_util import full_shape
 
 passthrough_nodes = set(("Identity",))
 
@@ -448,17 +448,27 @@ def fixup_act_reorder(op, arch):
             op.type = "Identity"
     return op
 
+
 def fixup_elementwise_with_scalars(op, arch):
     if op.type in binary_elementwise_op:
-        ifm_tensor, ifm2_tensor, _, ofm_tensor = op.get_ifm_ifm2_weights_ofm()
+        ifm_tensor, ifm2_tensor, _, _ = op.get_ifm_ifm2_weights_ofm()
         if ifm2_tensor.shape != [] and ifm_tensor.shape != []:
             diff = len(ifm_tensor.shape) - len(ifm2_tensor.shape)
             if diff > 0:
                 ifm2_tensor.shape = full_shape(len(ifm_tensor.shape), ifm2_tensor.shape, 1)
             elif diff < 0:
                 ifm_tensor.shape = full_shape(len(ifm2_tensor.shape), ifm_tensor.shape, 1)
+        elif ifm_tensor.shape == [] and ifm_tensor.quant_values is None:
+            # IFM is marked as a scalar, but is a result of an operation; change it to a shape of size 1
+            ifm_tensor.shape = len(ifm2_tensor.shape) * [1]
+            ifm_tensor.storage_shape = ifm_tensor.shape
+        elif ifm2_tensor.shape == [] and ifm2_tensor.quant_values is None:
+            # IFM2 is marked as a scalar, but is a result of an operation; change it to a shape of size 1
+            ifm2_tensor.shape = len(ifm_tensor.shape) * [1]
+            ifm2_tensor.storage_shape = ifm2_tensor.shape
     return op
 
+
 # Set input/output tensor equivalence to the same id for memory operations
 def set_tensor_equivalence(op, arch):
     if op.type == "Reshape":
diff --git a/ethosu/vela/high_level_command_stream.py b/ethosu/vela/high_level_command_stream.py
index 2c77e10c..9b0298bb 100644
--- a/ethosu/vela/high_level_command_stream.py
+++ b/ethosu/vela/high_level_command_stream.py
@@ -52,6 +52,8 @@ class Box:
             new_start_coord[-1] = 0
             new_end_coord[-1] = ifm_shape[-1]
 
+        if npu_block_type == NpuBlockType.ElementWise and min(len(new_end_coord), len(ifm_shape)) >= 1:
+            new_end_coord[-1] = min(new_end_coord[-1], ifm_shape[-1])
         if min(len(new_end_coord), len(ifm_shape)) >= 2:
             new_end_coord[-2] = min(new_end_coord[-2], ifm_shape[-2])
         if min(len(new_end_coord), len(ifm_shape)) >= 3: