MLBEDSW-3491: Fix index out of range in code gen
Usage of shape[-2] could cause index out of range.
Signed-off-by: Louis Verhaard <louis.verhaard@arm.com>
Change-Id: I1b64b117f8236ce9ba321ca03bdb25e5a03a6589
diff --git a/ethosu/vela/architecture_features.py b/ethosu/vela/architecture_features.py
index 6a02a4e..1eae79a 100644
--- a/ethosu/vela/architecture_features.py
+++ b/ethosu/vela/architecture_features.py
@@ -23,6 +23,7 @@
from .errors import OptionError
from .ethos_u55_regs.ethos_u55_regs import resampling_mode
+from .numeric_util import full_shape
from .numeric_util import round_up
from .numeric_util import round_up_divide
from .operation import Kernel
@@ -55,6 +56,13 @@
w, h, c = (int(v) for v in s.split("x"))
return cls(w, h, c)
+ @classmethod
+ def from_shape(cls, shape) -> "Block":
+ """Converts the shape to a Block"""
+ shp = full_shape(3, shape, 1)
+ # Note: index from end, as len(shp) may be > 3
+ return Block(shp[-2], shp[-3], shp[-1])
+
class Rect:
def __init__(self, x, y, z, x2, y2, z2):