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, 55 insertions, 37 deletions

diff --git a/arm_compute/core/utils/helpers/tensor_transform.h b/arm_compute/core/utils/helpers/tensor_transform.h
index 966c1f1fdf..aa359ad119 100644
--- a/arm_compute/core/utils/helpers/tensor_transform.h
+++ b/arm_compute/core/utils/helpers/tensor_transform.h
@@ -32,45 +32,33 @@ namespace helpers
 {
 namespace tensor_transform
 {
-/** Returns the absolute ends coordinates of slice
+/** Computes stride of a given index
  *
- * @param[in] input_shape Input tensor shape
- * @param[in] ends        End coordinates
+ * @param[in] index   Index of tensor to calculate absolute start position
+ * @param[in] strides Slice strides
  *
- * @return Absolute end coordinate
+ * @return Stride at a given index
  */
-Coordinates slice_absolute_end_coords(TensorShape input_shape, Coordinates ends);
+int calculate_stride_on_index(int index, Coordinates strides);
 
-/** Computes output shape of slice
- *
- * @warning Ends must be non-negative
- *
- * @param[in] input_shape Input tensor shape
- * @param[in] starts      Start coordinates
- * @param[in] ends_abs    Absolute end coordinates
- *
- * @return The output tensor shape
- */
-TensorShape compute_slice_output_shape(TensorShape input_shape, Coordinates starts, Coordinates ends_abs);
-
-/** Returns the absolute start coordinates of strided slice
+/** Computes absolute start position of a given index for a strided slice operation
  *
  * @param[in] input_shape Input tensor shape
+ * @param[in] index       Index of tensor to calculate absolute start position
  * @param[in] starts      Start coordinates
  * @param[in] strides     Slice strides
  * @param[in] begin_mask  (Optional) If the ith bit of begin_mask is set, starts[i] is ignored and
  *                        the fullest possible range in that dimension is used instead.
  *
- * @return Absolute start coordinates
+ * @return Absolute start position of a given index
  */
-Coordinates strided_slice_absolute_start_coords(TensorShape input_shape, Coordinates starts, Coordinates strides, int32_t begin_mask = 0);
+int calculate_start_on_index(TensorShape input_shape, int index, Coordinates starts, Coordinates strides, int32_t begin_mask);
 
-/** Returns the absolute ends coordinates of strided slice
- *
- * @warning Starts must be non-negative
+/** Returns the absolute end position of a given index for a strided slice operation
  *
  * @param[in] input_shape      Input tensor shape
- * @param[in] starts_abs       Absolute start coordinates
+ * @param[in] index            Index of tensor to calculate absolute start position
+ * @param[in] start_on_index   Absolute start coordinate for given index
  * @param[in] ends             End coordinates
  * @param[in] strides          Slice strides
  * @param[in] end_mask         (Optional) If the ith bit of end_mask is set, end[i] is ignored and
@@ -78,32 +66,62 @@ Coordinates strided_slice_absolute_start_coords(TensorShape input_shape, Coordin
  * @param[in] shrink_axis_mask (Optional) If the ith bit of shrink_axis_mask is set, it implies that the ith specification shrinks the dimensionality by 1.
  *                             A slice of size 1 starting from starts[i] in the dimension must be preserved.
  *
- * @return Absolute end coordinates
+ * @return Absolute end position of a given index
  */
-Coordinates strided_slice_absolute_end_coords(TensorShape input_shape, Coordinates starts_abs, Coordinates ends, Coordinates strides,
-                                              int32_t end_mask = 0, int32_t shrink_axis_mask = 0);
-/** Returns the final strides of strided slice
+int calculate_end_on_index(TensorShape input_shape, int index, int start_on_index, Coordinates ends, Coordinates strides,
+                           int32_t end_mask = 0, int32_t shrink_axis_mask = 0);
+
+/** Calculate start, end and stride coordinates for a strided slice
  *
- * @param[in] input_shape Input tensor shape
- * @param[in] strides     Slice strides
+ * @param[in] input_shape      Input tensor shape
+ * @param[in] starts           Start coordinates
+ * @param[in] ends             End coordinates
+ * @param[in] strides          Slice strides
+ * @param[in] begin_mask       (Optional) If the ith bit of begin_mask is set, starts[i] is ignored and
+ *                             the fullest possible range in that dimension is used instead.
+ * @param[in] end_mask         (Optional) If the ith bit of end_mask is set, end[i] is ignored and
+ *                             the fullest possible range in that dimension is used instead.
+ * @param[in] shrink_axis_mask (Optional) If the ith bit of shrink_axis_mask is set, it implies that the ith specification shrinks the dimensionality by 1.
+ *                             A slice of size 1 starting from starts[i] in the dimension must be preserved.
  *
- * @return The final strides need by strided slice
+ * @return A tuple with <Start,End,Strides>
  */
-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 = 0, int32_t end_mask = 0, int32_t shrink_axis_mask = 0);
 
 /** Computes output shape of strided slice
  *
  * @warning Starts and ends must be non-negative
  * @warning Starts, ends and final strides should have the same dimensions as the input shape
  *
- * @param[in] input_shape   Input tensor shape
- * @param[in] starts_abs    Absolute start coordinates
- * @param[in] ends_abs      Absolute end coordinates
- * @param[in] final_strides Slice strides
+ * @param[in] input_shape       Input tensor shape
+ * @param[in] starts            Absolute start coordinates
+ * @param[in] ends              Absolute end coordinates
+ * @param[in] strides           Slice strides
+ * @param[in] begin_mask        (Optional) If the ith bit of begin_mask is set, starts[i] is ignored and
+ *                              the fullest possible range in that dimension is used instead.
+ * @param[in] end_mask          (Optional) If the ith bit of end_mask is set, end[i] is ignored and
+ *                              the fullest possible range in that dimension is used instead.
+ * @param[in] shrink_axis_mask  (Optional) If the ith bit of shrink_axis_mask is set, it implies that the ith specification shrinks the dimensionality by 1.
+ *                              A slice of size 1 starting from starts[i] in the dimension must be preserved.
+ * @param[in] return_unshrinked (Optional) Returns un-shrinked shape
  *
  * @return The output tensor shape
  */
-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 = 0, int32_t end_mask = 0, int32_t shrink_axis_mask = 0,
+                                               bool return_unshrinked = false);
+
+/** Constructs end mask in case we want to perform a slice operation using the strided slice interface
+ *
+ * @note Ends are inclusive in slice operations that is why construction an end mask is needed
+ *
+ * @param[in] ends End coordinates
+ *
+ * @return End mask
+ */
+int32_t construct_slice_end_mask(Coordinates ends);
 } // namespace tensor_tranform
 } // namespace helpers
 } // namespace arm_compute