diff options
Diffstat (limited to 'ethosu/vela/high_level_command_stream_generator.py')
| -rw-r--r-- | ethosu/vela/high_level_command_stream_generator.py | 32 |
1 files changed, 13 insertions, 19 deletions
diff --git a/ethosu/vela/high_level_command_stream_generator.py b/ethosu/vela/high_level_command_stream_generator.py index 18a419c0..60e62aa6 100644 --- a/ethosu/vela/high_level_command_stream_generator.py +++ b/ethosu/vela/high_level_command_stream_generator.py @@ -27,6 +27,7 @@ from .numeric_util import round_up_divide 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 @@ def generate_high_level_command_stream_for_pass(strat, passes, block_configs, id 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 @@ def generate_high_level_command_stream_for_pass(strat, passes, block_configs, id 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 @@ def generate_high_level_command_stream_for_pass(strat, passes, block_configs, id 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 @@ def generate_high_level_command_stream_for_pass(strat, passes, block_configs, id 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 @@ def generate_high_level_command_stream_for_pass(strat, passes, block_configs, id ) 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 @@ def generate_high_level_command_stream_for_pass(strat, passes, block_configs, id 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 @@ def generate_high_level_command_stream_for_pass(strat, passes, block_configs, id 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 @@ def calc_allowed_ofm_ifm_overlap_for_pass_list(strat, passes, block_configs): 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 @@ def calc_allowed_ofm_ifm_overlap_for_pass_list(strat, passes, block_configs): 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) |