MLBEDSW-2779 Consider num dimensions, in check for NHCWB16
4 dimensions where assumed in check if NHCWB16 should be avoided.
Changed check so that if axis corresponds to C-dimension,
NHCWB16 should be avoided.
Signed-off-by: Patrik Gustavsson <patrik.gustavsson@arm.com>
Change-Id: I7784a7a813a3c3438d6142523bf0a3ba81742aca
diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py
index 3fe703e..daae915 100644
--- a/ethosu/vela/graph_optimiser.py
+++ b/ethosu/vela/graph_optimiser.py
@@ -69,11 +69,11 @@
tens.ops.append(new_op)
assert tens.shape[axis] == offset
- # If axis = 3, NHCWB16 can only be used in the output if all the concat_start's are a multiple of 16,
- # as it is only then the address offset for the ofm, for all operations, will be 16 byte aligned
- # For other values of axis the address offsets will be 16 byte aligned, as they are all based on c = 0
+ # If axis corresponds to C-dimension, NHCWB16 can only be used in the output if all the concat_start's are a
+ # multiple of 16. This as, it is only then the address offset for the ofm, for all operations, will be 16 byte
+ # aligned. For other values of axis the address offsets will be 16 byte aligned, as they are all based on c = 0
# and those addresses are always 16 byte aligned due to the NHCWB16 format.
- if axis == 3:
+ if axis == (len(tens.shape) - 1):
for op in tens.ops:
if op.attrs["concat_start"] % 16 != 0:
tens.avoid_NHCWB16 = True