Revert "Revert "MLBEDSW-3645 4D class for op ifm/ofm shapes""
This reverts commit df0a5905177f3a1b836076bc3f9f39b2e86f1794.
Reason for revert: <INSERT REASONING HERE>
Change-Id: I891c66fb29db9d25e942947e8d1c29a10610de51
diff --git a/ethosu/vela/shared_buffer_allocation.py b/ethosu/vela/shared_buffer_allocation.py
index 1f027d6..d8faf36 100644
--- a/ethosu/vela/shared_buffer_allocation.py
+++ b/ethosu/vela/shared_buffer_allocation.py
@@ -32,6 +32,7 @@
from .operation import NpuBlockType
from .range_set import MemoryRangeSet
from .register_command_stream_util import to_kernel
+from .shape4d import Shape4D
from .tensor import MemArea
@@ -195,14 +196,14 @@
ifm_bits = ifm_tensor.dtype.size_in_bits()
ifm_shape = ps.primary_op.ifm_shapes[0]
- if ifm_shape != []:
- ifm_depth = ifm_shape[-1]
+ if ifm_tensor.shape != []:
+ ifm_depth = ifm_shape.depth
if is_elementwise:
ifm_count = 2
if ifm_tensor.shape == []: # Scalar in ifm1
assert ifm2_tensor
- ifm_depth = ps.primary_op.ifm_shapes[1][-1]
+ ifm_depth = ps.primary_op.ifm_shapes[1].depth
ifm_count = 1
elif not ifm2_tensor or ifm2_tensor.shape == []: # Scalar in ifm2
ifm_count = 1
@@ -251,7 +252,7 @@
ifm_bits=ifm_bits,
ifm_depth=ifm_depth,
ifm_count=ifm_count,
- ofm_shape=ofm_shape,
+ ofm_shape=Shape4D(ofm_shape),
)
@@ -265,14 +266,9 @@
# Constrain the search space if the OFM is smaller than the max block size
# - Add other block search constraints here if required
- if len(alloc.ofm_shape) <= 2:
- max_block_height = max_block_width = alloc.ofm_shape[0]
- else:
- max_block_width = alloc.ofm_shape[-2]
- max_block_height = alloc.ofm_shape[-3]
-
- # Common block depth
- max_block_depth = alloc.ofm_shape[-1]
+ max_block_width = alloc.ofm_shape.width
+ max_block_height = alloc.ofm_shape.height
+ max_block_depth = alloc.ofm_shape.depth
# Constrain to valid ranges before search
max_block_width = min(arch.ofm_block_max.width, max_block_width)