MLBEDSW-7077: Store original tensor shape in optimized file
- CPU side always needs to work we the original tensor shape.
Due to a bypass memory optimization the IFM, produced by CPU,
was stored with the wrong shape in the optimized file.
- Store the original tensor shape so it can be correctly
written to the optimized file.
Change-Id: I666dbcb0acd806ad208c0f925a51dfc25421688b
Signed-off-by: Johan Alfven <johan.alfven@arm.com>
diff --git a/ethosu/vela/tensor.py b/ethosu/vela/tensor.py
index ba38588..673208a 100644
--- a/ethosu/vela/tensor.py
+++ b/ethosu/vela/tensor.py
@@ -340,6 +340,7 @@
class Tensor:
__slots__ = (
"shape",
+ "_original_shape",
"storage_shape",
"bandwidth_shape",
"dtype",
@@ -379,6 +380,7 @@
def __init__(self, shape: Shape, dtype: DataType, name: str):
self.shape = shape
+ self._original_shape = shape
self.storage_shape = shape
self.bandwidth_shape = shape
self.dtype = dtype
@@ -425,6 +427,10 @@
self.src_tensor: Optional[Tensor] = None
@property
+ def original_shape(self):
+ return self._original_shape
+
+ @property
def address(self) -> int:
return TensorAddressMap.get_address_for_tens(self.equivalence_id, self.mem_type)
diff --git a/ethosu/vela/tflite_writer.py b/ethosu/vela/tflite_writer.py
index ce53f9b..6fdfe01 100644
--- a/ethosu/vela/tflite_writer.py
+++ b/ethosu/vela/tflite_writer.py
@@ -26,6 +26,7 @@
from .operation import Op
from .reader_util import align_inputs_indices
from .tensor import MemType
+from .tensor import shape_num_elements
from .tensor import TensorPurpose
from .tflite import Buffer
from .tflite import Metadata
@@ -248,7 +249,12 @@
def serialise_tensor(self, tens):
builder = self.builder
- tens_shape = tens.shape
+ if shape_num_elements(tens.original_shape) != shape_num_elements(tens.shape):
+ # shapes have changed size, therefore assume that the latest (modified) shape is correct
+ tens_shape = tens.shape
+ else:
+ # shapes have not changed size, therefore the original shape is valid
+ tens_shape = tens.original_shape
values = tens.values
if values is None: