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/high_level_command_stream_generator.py b/ethosu/vela/high_level_command_stream_generator.py
index 18a419c..60e62aa 100644
--- a/ethosu/vela/high_level_command_stream_generator.py
+++ b/ethosu/vela/high_level_command_stream_generator.py
@@ -27,6 +27,7 @@
from .operation import create_activation_function
from .operation import NpuBlockType
from .operation import Op
+from .shape4d import Shape4D
from .tensor import TensorPurpose
@@ -90,8 +91,8 @@
weight_tensor = ps.weight_tensor
scale_tensor = ps.scale_tensor
- ofm_start = [0] * len(ofm_shape)
- ofm_end = list(ofm_shape)
+ ofm_start = [0, 0, 0, 0]
+ ofm_end = ofm_shape.as_list()
strides = None
skirt = None
@@ -100,9 +101,9 @@
strides = ps.primary_op.attrs.get("strides", None)
skirt = ps.primary_op.attrs.get("skirt", None)
if ps.primary_op.type == Op.Conv2DBackpropInputSwitchedBias:
- upscaling = ofm_shape[-3] // ifm_shape[-3]
+ upscaling = ofm_shape.height // ifm_shape.height
elif ps.primary_op.type == Op.ResizeBilinear:
- upscaling = round_up_divide(ofm_shape[-3], ifm_shape[-3])
+ upscaling = round_up_divide(ofm_shape.height, ifm_shape.height)
concat_axis = 0
concat_offset = 0
@@ -135,14 +136,7 @@
if ifm_shape is not None:
ifm_box, _, _ = ofm_box.transform_with_strides_and_skirt(
- strides,
- skirt,
- ifm_tensor.shape,
- npu_block_type,
- concat_axis,
- concat_offset,
- split_offsets[0],
- upscaling,
+ strides, skirt, ifm_shape, npu_block_type, concat_axis, concat_offset, split_offsets[0], upscaling,
)
else:
ifm_box = Box([], [])
@@ -163,7 +157,7 @@
intermediate_box, _, _ = ofm_box.transform_with_strides_and_skirt(
strides,
skirt,
- intermediate.shape,
+ Shape4D(intermediate.shape),
npu_block_type,
concat_axis,
concat_offset,
@@ -212,6 +206,7 @@
)
elif strat == SchedulingStrategy.IfmStream:
+ assert ifm_shape is not None
y_step = block_config[0]
y_start = ofm_start[-3]
y_dim = ofm_end[-3]
@@ -222,8 +217,7 @@
prev_pass_gen = generate_high_level_command_stream_for_pass(strat, passes, block_configs, idx - 1)
else:
ifm_y_present = 1
- if len(ifm_shape) >= 3:
- ifm_y_present = ifm_shape[-3]
+ ifm_y_present = ifm_shape.height
prev_pass_gen = []
prev_pass = None
@@ -276,7 +270,7 @@
intermediate_box, _, _ = ofm_box.transform_with_strides_and_skirt(
strides,
skirt,
- intermediate.shape,
+ Shape4D(intermediate.shape),
npu_block_type,
concat_axis,
concat_offset,
@@ -380,13 +374,13 @@
if cmd.is_npu_pass_command():
if cmd.is_first:
ifm_read = cmd.ifm_tensor.address_offset_for_coordinate(
- cmd.ifm_box.start_coord, shape=cmd.ps.ifm_shapes[0], is_top_box=False
+ cmd.ifm_box.start_coord, cmd.ps.ifm_shapes[0].as_list(), is_top_box=False
)
if ifm_read is None:
return 0
if cmd.is_last:
write_offset = cmd.ofm_tensor.address_offset_for_coordinate(
- cmd.ofm_box.end_coord, shape=cmd.ps.ofm_shapes[0], is_top_box=True
+ cmd.ofm_box.end_coord, cmd.ps.ofm_shapes[0].as_list(), is_top_box=True
)
if write_offset is None:
return 0
@@ -399,7 +393,7 @@
if cmd.is_first:
ifm_read = cmd.ifm_tensor.address_offset_for_coordinate(
- cmd.ifm_box.end_coord, shape=cmd.ps.ifm_shapes[0], is_top_box=True
+ cmd.ifm_box.end_coord, cmd.ps.ifm_shapes[0].as_list(), is_top_box=True
)
min_overlap = max(min_overlap, 0)