From eeb85154b00a9864d0d63e382e9c80ca8e294d5d Mon Sep 17 00:00:00 2001
From: "patrik.gustavsson" <patrik.gustavsson@arm.com>
Date: Mon, 21 Dec 2020 17:10:40 +0000
Subject: Revert "Revert "MLBEDSW-3645 4D class for op ifm/ofm shapes""

This reverts commit df0a5905177f3a1b836076bc3f9f39b2e86f1794.

Reason for revert: <INSERT REASONING HERE>

Change-Id: I891c66fb29db9d25e942947e8d1c29a10610de51
---
 ethosu/vela/high_level_command_stream_generator.py | 32 +++++++++-------------
 1 file changed, 13 insertions(+), 19 deletions(-)

(limited to 'ethosu/vela/high_level_command_stream_generator.py')

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)
-- 
cgit v1.2.1