Add support for 5D data layout indexing

Signed-off-by: Giorgio Arena <giorgio.arena@arm.com>
Change-Id: Ib346bb6b90d2220ec5934c83a9a1f0cd540b8731
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/6377
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Gunes Bayir <gunes.bayir@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>

5 files changed, 72 insertions, 56 deletions

diff --git a/arm_compute/core/Helpers.h b/arm_compute/core/Helpers.h
index b6635aba6d..fd6e94c079 100644
--- a/arm_compute/core/Helpers.h
+++ b/arm_compute/core/Helpers.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2020 Arm Limited.
+ * Copyright (c) 2016-2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -180,6 +180,22 @@ inline Coordinates index2coords(const TensorShape &shape, int index);
  */
 inline int coords2index(const TensorShape &shape, const Coordinates &coord);
 
+/** Returns a static map used to find an index or dimension based on a data layout
+  *
+  * *** Layouts ***
+  *
+  * *** 4D ***
+  * [N C H W]
+  * [3 2 1 0]
+  * [N H W C]
+  *
+  * * *** 5D ***
+  * [N C D H W]
+  * [4 3 2 1 0]
+  * [N D H W C]
+  */
+const std::map<DataLayout, std::vector<DataLayoutDimension>> &get_layout_map();
+
 /** Get the index of the given dimension.
  *
  * @param[in] data_layout           The data layout.
@@ -187,7 +203,7 @@ inline int coords2index(const TensorShape &shape, const Coordinates &coord);
  *
  * @return The int conversion of the requested data layout index.
  */
-inline size_t get_data_layout_dimension_index(const DataLayout data_layout, const DataLayoutDimension data_layout_dimension);
+inline size_t get_data_layout_dimension_index(const DataLayout &data_layout, const DataLayoutDimension &data_layout_dimension);
 
 /** Get the DataLayoutDimension of a given index and layout.
  *
@@ -196,7 +212,7 @@ inline size_t get_data_layout_dimension_index(const DataLayout data_layout, cons
  *
  * @return The dimension which this index is requested for.
  */
-inline DataLayoutDimension get_index_data_layout_dimension(const DataLayout data_layout, const size_t index);
+inline DataLayoutDimension get_index_data_layout_dimension(const DataLayout &data_layout, const size_t index);
 
 /** Calculate the number of output tiles required by Winograd Convolution layer. This utility function can be used by the Winograd input transform
  *  to know the number of tiles on the x and y direction
diff --git a/arm_compute/core/Helpers.inl b/arm_compute/core/Helpers.inl
index a960876074..a910521f94 100644
--- a/arm_compute/core/Helpers.inl
+++ b/arm_compute/core/Helpers.inl
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2020 Arm Limited.
+ * Copyright (c) 2016-2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -190,61 +190,20 @@ inline int coords2index(const TensorShape &shape, const Coordinates &coord)
     return index;
 }
 
-inline size_t get_data_layout_dimension_index(const DataLayout data_layout, const DataLayoutDimension data_layout_dimension)
+inline size_t get_data_layout_dimension_index(const DataLayout &data_layout, const DataLayoutDimension &data_layout_dimension)
 {
     ARM_COMPUTE_ERROR_ON_MSG(data_layout == DataLayout::UNKNOWN, "Cannot retrieve the dimension index for an unknown layout!");
-
-    /* Return the index based on the data layout
-     * [N C H W]
-     * [3 2 1 0]
-     * [N H W C]
-    */
-    switch(data_layout_dimension)
-    {
-        case DataLayoutDimension::CHANNEL:
-            return (data_layout == DataLayout::NCHW) ? 2 : 0;
-            break;
-        case DataLayoutDimension::HEIGHT:
-            return (data_layout == DataLayout::NCHW) ? 1 : 2;
-            break;
-        case DataLayoutDimension::WIDTH:
-            return (data_layout == DataLayout::NCHW) ? 0 : 1;
-            break;
-        case DataLayoutDimension::BATCHES:
-            return 3;
-            break;
-        default:
-            break;
-    }
-    ARM_COMPUTE_ERROR("Data layout index not supported!");
+    const auto &dims = get_layout_map().at(data_layout);
+    const auto &it   = std::find(dims.cbegin(), dims.cend(), data_layout_dimension);
+    ARM_COMPUTE_ERROR_ON_MSG(it == dims.cend(), "Invalid dimension for the given layout.");
+    return it - dims.cbegin();
 }
 
-inline DataLayoutDimension get_index_data_layout_dimension(const DataLayout data_layout, const size_t index)
+inline DataLayoutDimension get_index_data_layout_dimension(const DataLayout &data_layout, const size_t index)
 {
-    ARM_COMPUTE_ERROR_ON_MSG(data_layout == DataLayout::UNKNOWN, "Cannot retrieve the dimension index for an unknown layout!");
-
-    /* Return the index based on the data layout
-    * [N C H W]
-    * [3 2 1 0]
-    * [N H W C]
-    */
-    switch(index)
-    {
-        case 0:
-            return (data_layout == DataLayout::NCHW) ? DataLayoutDimension::WIDTH : DataLayoutDimension::CHANNEL;
-            break;
-        case 1:
-            return (data_layout == DataLayout::NCHW) ? DataLayoutDimension::HEIGHT : DataLayoutDimension::WIDTH;
-            break;
-        case 2:
-            return (data_layout == DataLayout::NCHW) ? DataLayoutDimension::CHANNEL : DataLayoutDimension::HEIGHT;
-            break;
-        case 3:
-            return DataLayoutDimension::BATCHES;
-            break;
-        default:
-            ARM_COMPUTE_ERROR("Index value not supported!");
-            break;
-    }
+    ARM_COMPUTE_ERROR_ON_MSG(data_layout == DataLayout::UNKNOWN, "Cannot retrieve the layout dimension for an unknown layout!");
+    const auto &dims = get_layout_map().at(data_layout);
+    ARM_COMPUTE_ERROR_ON_MSG(index >= dims.size(), "Invalid index for the given layout.");
+    return dims[index];
 }
 } // namespace arm_compute
diff --git a/arm_compute/core/Types.h b/arm_compute/core/Types.h
index 31199e138b..37ba9f93bf 100644
--- a/arm_compute/core/Types.h
+++ b/arm_compute/core/Types.h
@@ -114,6 +114,7 @@ enum class DataLayout
     UNKNOWN, /**< Unknown data layout */
     NCHW,    /**< Num samples, channels, height, width */
     NHWC,    /**< Num samples, height, width, channels */
+    NCDHW,   /**< Num samples, channels, depth, height, width */
     NDHWC    /**< Num samples, depth, height, width, channels */
 };
 /** [DataLayout enum definition] **/
@@ -124,6 +125,7 @@ enum class DataLayoutDimension
     CHANNEL, /**< channel */
     HEIGHT,  /**< height */
     WIDTH,   /**< width */
+    DEPTH,   /**< depth */
     BATCHES  /**< batches */
 };
 
@@ -765,6 +767,20 @@ private:
 /** Padding information for 3D operations like Conv3d */
 struct Padding3D
 {
+    Padding3D()
+    {
+    }
+
+    Padding3D(size_t pad_x, size_t pad_y, size_t pad_z)
+        : left(pad_x), right(pad_x), top(pad_y), bottom(pad_y), front(pad_z), back(pad_z)
+    {
+    }
+
+    Padding3D(size_t left, size_t right, size_t top, size_t bottom, size_t front, size_t back)
+        : left(left), right(right), top(top), bottom(bottom), front(front), back(back)
+    {
+    }
+
     size_t left   = { 0 }; /**<  Padding across the width dimenstion on the left, in elements. */
     size_t right  = { 0 }; /**<  Padding across the width dimenstion on the right, in elements. */
     size_t top    = { 0 }; /**<  Padding across the height dimenstion  on the top, in elements. */
diff --git a/src/core/Helpers.cpp b/src/core/Helpers.cpp
index e692cc1e7c..28e7f4c1e5 100644
--- a/src/core/Helpers.cpp
+++ b/src/core/Helpers.cpp
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2020 Arm Limited.
+ * Copyright (c) 2016-2021 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -100,4 +100,23 @@ ValidRegion calculate_valid_region_scale(const ITensorInfo &src_info, const Tens
 
     return valid_region;
 }
+
+const std::map<DataLayout, std::vector<DataLayoutDimension>> &get_layout_map()
+{
+    constexpr DataLayoutDimension W = DataLayoutDimension::WIDTH;
+    constexpr DataLayoutDimension H = DataLayoutDimension::HEIGHT;
+    constexpr DataLayoutDimension C = DataLayoutDimension::CHANNEL;
+    constexpr DataLayoutDimension D = DataLayoutDimension::DEPTH;
+    constexpr DataLayoutDimension N = DataLayoutDimension::BATCHES;
+
+    static const std::map<DataLayout, std::vector<DataLayoutDimension>> layout_map =
+    {
+        { DataLayout::NDHWC, { C, W, H, D, N } },
+        { DataLayout::NCDHW, { W, H, D, C, N } },
+        { DataLayout::NHWC, { C, W, H, N } },
+        { DataLayout::NCHW, { W, H, C, N } }
+    };
+
+    return layout_map;
+}
 } // namespace arm_compute
\ No newline at end of file
diff --git a/utils/TypePrinter.h b/utils/TypePrinter.h
index 5fa92e6360..0bea408519 100644
--- a/utils/TypePrinter.h
+++ b/utils/TypePrinter.h
@@ -666,6 +666,9 @@ inline ::std::ostream &operator<<(::std::ostream &os, const DataLayout &data_lay
         case DataLayout::NDHWC:
             os << "NDHWC";
             break;
+        case DataLayout::NCDHW:
+            os << "NCDHW";
+            break;
         default:
             ARM_COMPUTE_ERROR("NOT_SUPPORTED!");
     }
@@ -707,6 +710,9 @@ inline ::std::ostream &operator<<(::std::ostream &os, const DataLayoutDimension
         case DataLayoutDimension::CHANNEL:
             os << "CHANNEL";
             break;
+        case DataLayoutDimension::DEPTH:
+            os << "DEPTH";
+            break;
         case DataLayoutDimension::BATCHES:
             os << "BATCHES";
             break;