COMPMID-1710: Fixes in StrideSlice calculations.

Change-Id: I66eb922f1ff15142de278bf4439a61c979f98ba7
Reviewed-on: https://review.mlplatform.org/382
Reviewed-by: Matthew Bentham <matthew.bentham@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Pablo Marquez <pablo.tello@arm.com>

1 files changed, 102 insertions, 90 deletions

diff --git a/src/core/utils/helpers/tensor_transform.cpp b/src/core/utils/helpers/tensor_transform.cpp
index a4bce5da5a..08803c7fb0 100644
--- a/src/core/utils/helpers/tensor_transform.cpp
+++ b/src/core/utils/helpers/tensor_transform.cpp
@@ -23,143 +23,155 @@
  */
 #include "arm_compute/core/utils/helpers/tensor_transform.h"
 
+#include "arm_compute/core/utils/helpers/bit_ops.h"
+
 namespace arm_compute
 {
 namespace helpers
 {
 namespace tensor_transform
 {
-Coordinates slice_absolute_end_coords(TensorShape input_shape, Coordinates ends)
-{
-    // Create end mask
-    int32_t end_mask = 0;
-    for(unsigned int i = 0; i < ends.num_dimensions(); ++i)
-    {
-        if(ends[i] < 0)
-        {
-            end_mask |= 1 << i;
-        }
-    }
-    // Get unit strides
-    const BiStrides unit_strides = strided_slice_strides(input_shape, BiStrides());
-
-    return strided_slice_absolute_end_coords(input_shape, Coordinates(), ends, unit_strides, end_mask);
-}
-
-TensorShape compute_slice_output_shape(TensorShape input_shape, Coordinates starts, Coordinates ends_abs)
+int calculate_stride_on_index(int index, Coordinates strides)
 {
-    // Get unit strides
-    const BiStrides unit_strides = strided_slice_strides(input_shape, BiStrides());
-    return compute_strided_slice_output_shape(input_shape, starts, ends_abs, unit_strides);
+    return index >= static_cast<int>(strides.num_dimensions()) ? 1 : strides[index];
 }
 
-Coordinates strided_slice_absolute_start_coords(TensorShape input_shape, Coordinates starts, Coordinates strides, int32_t begin_mask)
+int calculate_start_on_index(TensorShape input_shape, int index, Coordinates starts, Coordinates strides, int32_t begin_mask)
 {
-    Coordinates starts_abs;
-    for(unsigned int i = 0; i < starts.num_dimensions(); ++i)
+    // Early exit
+    if(index >= static_cast<int>(starts.num_dimensions()))
     {
-        // Get start index
-        int start_i = starts[i];
+        return 0;
+    }
 
-        // Reset in case of begin mask present
-        if((begin_mask & 1 << i) != 0)
-        {
-            start_i = strides[i] > 0 ? std::numeric_limits<int>::lowest() : std::numeric_limits<int>::max();
-        }
+    // Get stride
+    const int stride = calculate_stride_on_index(index, strides);
 
-        // Account negative start points
-        const int dim_size = input_shape[i];
-        if(start_i < 0)
-        {
-            start_i += dim_size;
-        }
+    // Calculate start
+    int start = starts[index];
 
-        // Final clamp
-        start_i = utility::clamp(start_i, 0, dim_size - 1);
-        starts_abs.set(i, start_i);
+    // Reset in case of begin mask present
+    if(arm_compute::helpers::bit_ops::is_bit_set(begin_mask, index))
+    {
+        start = stride > 0 ? std::numeric_limits<int>::lowest() : std::numeric_limits<int>::max();
     }
 
-    // Fill remaining
-    for(unsigned int i = starts_abs.num_dimensions(); i < input_shape.num_dimensions(); ++i)
+    // Account negative start points
+    const int dim_size = input_shape[index];
+    if(start < 0)
     {
-        starts_abs.set(i, 0);
+        start += dim_size;
     }
 
-    return starts_abs;
+    // Final clamp
+    start = utility::clamp(start, 0, dim_size - 1);
+
+    return start;
 }
 
-Coordinates strided_slice_absolute_end_coords(TensorShape input_shape, Coordinates starts_abs, Coordinates ends, Coordinates strides,
-                                              int32_t end_mask, int32_t shrink_axis_mask)
+int calculate_end_on_index(TensorShape input_shape, int index, int start_on_index,
+                           Coordinates ends, Coordinates strides,
+                           int32_t end_mask, int32_t shrink_axis_mask)
 {
-    Coordinates ends_abs;
-    for(unsigned int i = 0; i < ends.num_dimensions(); ++i)
+    // Early exit
+    if(index >= static_cast<int>(ends.num_dimensions()))
     {
-        // Get end index
-        int stop_i = ends[i];
+        return input_shape[index];
+    }
 
-        // Shrink dimension
-        if((shrink_axis_mask & (1 << i)) != 0)
-        {
-            stop_i = starts_abs[i] + 1;
-        }
+    const int  stride      = calculate_stride_on_index(index, strides);
+    const bool shrink_axis = arm_compute::helpers::bit_ops::is_bit_set(shrink_axis_mask, index);
 
-        // Reset in case of begin mask present
-        if((end_mask & 1 << i) != 0)
-        {
-            stop_i = (strides[i] > 0) ? std::numeric_limits<int>::max() : std::numeric_limits<int>::lowest();
-        }
+    // Calculate start
+    int stop = ends[index];
 
-        // Account negative end points
-        const int dim_size = input_shape[i];
-        if(stop_i < 0)
-        {
-            stop_i += dim_size;
-        }
+    // Shrink dimension
+    if(shrink_axis)
+    {
+        stop = start_on_index + 1;
+    }
 
-        // Final clamp
-        stop_i = (strides[i] > 0) ? utility::clamp(stop_i, 0, dim_size) : utility::clamp(stop_i, -1, dim_size - 1);
-        ends_abs.set(i, stop_i);
+    // Reset in case of begin mask present
+    if(arm_compute::helpers::bit_ops::is_bit_set(end_mask, index) && !shrink_axis)
+    {
+        stop = (stride > 0) ? std::numeric_limits<int>::max() : std::numeric_limits<int>::lowest();
     }
 
-    // Fill remaining ends
-    for(unsigned int i = ends_abs.num_dimensions(); i < input_shape.num_dimensions(); ++i)
+    // Account negative end points
+    const int dim_size = input_shape[index];
+    if(stop < 0)
     {
-        ends_abs.set(i, input_shape[i]);
+        stop += dim_size;
     }
 
-    return ends_abs;
+    // Final clamp
+    stop = (stride > 0) ? utility::clamp(stop, 0, dim_size) : utility::clamp(stop, -1, dim_size - 1);
+
+    return stop;
 }
 
-Coordinates strided_slice_strides(TensorShape input_shape, Coordinates strides)
+std::tuple<Coordinates, Coordinates, Coordinates> calculate_strided_slice_coords(TensorShape input_shape,
+                                                                                 Coordinates starts, Coordinates ends, Coordinates strides,
+                                                                                 int32_t begin_mask, int32_t end_mask, int32_t shrink_axis_mask)
 {
-    for(unsigned int i = strides.num_dimensions(); i < input_shape.num_dimensions(); ++i)
+    Coordinates starts_abs, ends_abs, final_strides;
+    for(unsigned int i = 0; i < input_shape.num_dimensions(); ++i)
     {
-        strides.set(i, 1);
+        const int start_i = calculate_start_on_index(input_shape, i, starts, strides, begin_mask);
+        starts_abs.set(i, start_i);
+        ends_abs.set(i, calculate_end_on_index(input_shape, i, start_i, ends, strides, end_mask, shrink_axis_mask));
+        final_strides.set(i, calculate_stride_on_index(i, strides));
     }
-    return strides;
+
+    return std::make_tuple(starts_abs, ends_abs, final_strides);
 }
 
-TensorShape compute_strided_slice_output_shape(TensorShape input_shape, Coordinates starts_abs, Coordinates ends_abs, Coordinates final_strides)
+TensorShape compute_strided_slice_output_shape(TensorShape input_shape, Coordinates starts, Coordinates ends, Coordinates strides,
+                                               int32_t begin_mask, int32_t end_mask, int32_t shrink_axis_mask, bool return_unshrinked)
 {
-    TensorShape output_shape = input_shape;
+    unsigned int index = 0;
+
+    TensorShape output_shape;
     for(unsigned int i = 0; i < input_shape.num_dimensions(); ++i)
     {
-        const int stride_i = final_strides[i];
-        const int range    = ends_abs[i] - starts_abs[i];
-        if((range == 0) ||                 // Zero range
-           (range < 0 && stride_i >= 0) || // Negative range with positive stride
-           (range > 0 && stride_i <= 0))   // Positive range with negative stride
+        const int stride = calculate_stride_on_index(index, strides);
+        const int start  = calculate_start_on_index(input_shape, i, starts, strides, begin_mask);
+        const int end    = calculate_end_on_index(input_shape, i, start, ends, strides, end_mask, shrink_axis_mask);
+        const int range  = end - start;
+
+        const bool is_shrink = arm_compute::helpers::bit_ops::is_bit_set(shrink_axis_mask, i);
+        if(return_unshrinked || !is_shrink)
         {
-            output_shape.set(i, 0);
-            return output_shape;
+            if((range == 0) ||               // Zero range
+               (range < 0 && stride >= 0) || // Negative range with positive stride
+               (range > 0 && stride <= 0))   // Positive range with negative stride
+            {
+                output_shape.set(index, 0);
+                return output_shape;
+            }
+            else
+            {
+                int dim = range / stride + (range % stride != 0 ? 1 : 0);
+                output_shape.set(index++, dim);
+            }
         }
-        else
+    }
+    return output_shape;
+}
+
+int32_t construct_slice_end_mask(Coordinates ends)
+{
+    // Create end mask
+    int32_t end_mask = 0;
+    for(unsigned int i = 0; i < ends.num_dimensions(); ++i)
+    {
+        if(ends[i] < 0)
         {
-            int dim = range / stride_i + (range % stride_i != 0 ? 1 : 0);
-            output_shape.set(i, dim);
+            end_mask |= 1 << i;
         }
     }
-    return output_shape;
+
+    return end_mask;
 }
 } // namespace tensor_transform
 } // namespace helpers