Add Pool3d reference implementation

This patch
  - adds the reference implementation for the 3D pooling layer
  - supports FP32/FP16 and INT8/UINT8 types
  - adds a function to calculate the output shape for 3D pooling
  - adds a new type for describing pool 3d info (Pool3DInfo)

Resolves: COMPMID-4659

Change-Id: I22a18fa30625c98fa827ef1b50781db6893ba9c4
Signed-off-by: Gunes Bayir <gunes.bayir@arm.com>
Reviewed-on: https://review.mlplatform.org/c/ml/ComputeLibrary/+/7219
Reviewed-by: Gian Marco Iodice <gianmarco.iodice@arm.com>
Comments-Addressed: Arm Jenkins <bsgcomp@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>

4 files changed, 365 insertions, 2 deletions

diff --git a/arm_compute/core/Types.h b/arm_compute/core/Types.h
index 9615811349..6c49cc35c0 100644
--- a/arm_compute/core/Types.h
+++ b/arm_compute/core/Types.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2021 Arm Limited.
+ * Copyright (c) 2016-2022 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -1307,6 +1307,42 @@ private:
     unsigned int _sampling_ratio;
 };
 
+struct Pool3DInfo
+{
+    Pool3DInfo() = default;
+
+    /** Constructor
+     *
+     * @param[in] pool_type         Pooling type @ref PoolingType.
+     * @param[in] pool_size         Pooling size, in elements, across  x, y and z @ref Size3D
+     * @param[in] padding           Paddings in x, y and z dimensions
+     * @param[in] strides           Strides in x, y and z dimensions @ref Size3D
+     * @param[in] round_type        Dimension rounding type (ceil or floor)
+     * @param[in] exclude_padding   Strategy when accounting padding in calculations.
+     *                              True will exclude padding while false will not (Used in AVG/L2 pooling to determine the pooling area).
+     * @param[in] is_global_pooling Sets the pool size to the input size if True
+     */
+    Pool3DInfo(const PoolingType           pool_type,
+               const Size3D                pool_size,
+               const Padding3D             padding,
+               const Size3D                strides,
+               const DimensionRoundingType round_type,
+               bool                        exclude_padding,
+               bool                        is_global_pooling)
+        : pool_type(pool_type), pool_size(pool_size), padding(padding), strides(strides), round_type(round_type), exclude_padding(exclude_padding), is_global_pooling(is_global_pooling)
+
+    {
+    }
+
+    PoolingType           pool_type{ PoolingType::MAX };
+    Size3D                pool_size{ 1U, 1U, 1U };
+    Padding3D             padding{};
+    Size3D                strides{ 1U, 1U, 1U };
+    DimensionRoundingType round_type{ DimensionRoundingType::FLOOR };
+    bool                  exclude_padding{ false };
+    bool                  is_global_pooling{ false };
+};
+
 /** Generate Proposals Information class */
 class GenerateProposalsInfo
 {
diff --git a/arm_compute/core/utils/misc/ShapeCalculator.h b/arm_compute/core/utils/misc/ShapeCalculator.h
index 3e8b024f82..ee4fe0c02f 100644
--- a/arm_compute/core/utils/misc/ShapeCalculator.h
+++ b/arm_compute/core/utils/misc/ShapeCalculator.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2021 Arm Limited.
+ * Copyright (c) 2017-2022 Arm Limited.
  *
  * SPDX-License-Identifier: MIT
  *
@@ -1460,6 +1460,63 @@ inline TensorShape compute_conv3d_shape(const TensorShape &src, const TensorShap
     return output_shape;
 }
 
+/** Calculate the output pool3d shape of a tensor
+ *
+ * @param[in] src         Input tensor info
+ * @param[in] pool3d_info Pooling layer info
+ *
+ * @return the calculated shape
+ */
+inline TensorShape compute_pool3d_shape(const TensorShape &src, Pool3DInfo pool3d_info)
+{
+    TensorShape output_shape{ src };
+
+    const int idx_width          = 1;
+    const int idx_height         = 2;
+    const int idx_depth          = 3;
+    const int pool_size_width    = pool3d_info.is_global_pooling ? src[idx_width] : pool3d_info.pool_size.width;
+    const int pool_size_height   = pool3d_info.is_global_pooling ? src[idx_height] : pool3d_info.pool_size.height;
+    const int pool_size_depth    = pool3d_info.is_global_pooling ? src[idx_depth] : pool3d_info.pool_size.depth;
+    const int pool_stride_width  = pool3d_info.strides.width;
+    const int pool_stride_height = pool3d_info.strides.height;
+    const int pool_stride_depth  = pool3d_info.strides.depth;
+
+    int output_width_size  = 0;
+    int output_height_size = 0;
+    int output_depth_size  = 0;
+
+    const size_t pad_left   = pool3d_info.padding.left;
+    const size_t pad_right  = pool3d_info.padding.right;
+    const size_t pad_top    = pool3d_info.padding.top;
+    const size_t pad_bottom = pool3d_info.padding.bottom;
+    const size_t pad_front  = pool3d_info.padding.front;
+    const size_t pad_back   = pool3d_info.padding.back;
+
+    switch(pool3d_info.round_type)
+    {
+        case DimensionRoundingType::FLOOR:
+            output_width_size  = static_cast<int>(std::floor((static_cast<float>(src[idx_width] + pad_left + pad_right - pool_size_width)) / pool_stride_width) + 1);
+            output_height_size = static_cast<int>(std::floor((static_cast<float>(src[idx_height] + pad_top + pad_bottom - pool_size_height)) / pool_stride_height) + 1);
+            output_depth_size  = static_cast<int>(std::floor((static_cast<float>(src[idx_depth] + pad_front + pad_back - pool_size_depth)) / pool_stride_depth) + 1);
+            break;
+        case DimensionRoundingType::CEIL:
+            output_width_size  = static_cast<int>(std::ceil((static_cast<float>(src[idx_width] + pad_left + pad_right - pool_size_width)) / pool_stride_width) + 1);
+            output_height_size = static_cast<int>(std::ceil((static_cast<float>(src[idx_height] + pad_top + pad_bottom - pool_size_height)) / pool_stride_height) + 1);
+            output_depth_size  = static_cast<int>(std::ceil((static_cast<float>(src[idx_depth] + pad_front + pad_back - pool_size_depth)) / pool_stride_depth) + 1);
+            break;
+        default:
+            ARM_COMPUTE_ERROR("Unsupported rounding type");
+    }
+
+    ARM_COMPUTE_ERROR_ON_MSG((output_width_size < 1 || output_height_size < 1 || output_depth_size < 1), "Calculated output dimension size is invalid");
+
+    output_shape.set(idx_width, static_cast<size_t>(output_width_size));
+    output_shape.set(idx_height, static_cast<size_t>(output_height_size));
+    output_shape.set(idx_depth, static_cast<size_t>(output_depth_size));
+
+    return output_shape;
+}
+
 inline TensorShape compute_gather_shape(const TensorShape &input_shape, const TensorShape &indices_shape, uint32_t actual_axis)
 {
     ARM_COMPUTE_ERROR_ON(indices_shape.num_dimensions() > 1);
diff --git a/tests/validation/reference/Pool3D.cpp b/tests/validation/reference/Pool3D.cpp
new file mode 100644
index 0000000000..85a594e262
--- /dev/null
+++ b/tests/validation/reference/Pool3D.cpp
@@ -0,0 +1,221 @@
+/*
+ * Copyright (c) 2022 Arm Limited.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+
+#include "Pool3D.h"
+#include "arm_compute/core/utils/misc/ShapeCalculator.h"
+#include "tests/validation/Helpers.h"
+
+namespace arm_compute
+{
+namespace test
+{
+namespace validation
+{
+namespace reference
+{
+using namespace arm_compute::misc::shape_calculator;
+
+template <typename T>
+SimpleTensor<T> pool3d_internal(const SimpleTensor<T> &src, const Pool3DInfo &pool3d_info, SimpleTensor<uint32_t> *indices)
+{
+    TensorShape     pooled_shape = compute_pool3d_shape(src.shape(), pool3d_info);
+    SimpleTensor<T> dst{ pooled_shape, src.data_type(), 1 };
+
+    if(indices != nullptr)
+    {
+        *indices = SimpleTensor<uint32_t> { pooled_shape, DataType::U32, 1 };
+    }
+
+    const int idx_channel = 0;
+    const int idx_width   = 1;
+    const int idx_height  = 2;
+    const int idx_depth   = 3;
+    const int idx_batch   = 4;
+
+    const int pool_size_width  = pool3d_info.is_global_pooling ? src.shape()[idx_width] : pool3d_info.pool_size.width;
+    const int pool_size_height = pool3d_info.is_global_pooling ? src.shape()[idx_height] : pool3d_info.pool_size.height;
+    const int pool_size_depth  = pool3d_info.is_global_pooling ? src.shape()[idx_depth] : pool3d_info.pool_size.depth;
+
+    const int pool_stride_width  = static_cast<int>(pool3d_info.strides.width);
+    const int pool_stride_height = static_cast<int>(pool3d_info.strides.height);
+    const int pool_stride_depth  = static_cast<int>(pool3d_info.strides.depth);
+
+    const int pad_left  = static_cast<int>(pool3d_info.padding.left);
+    const int pad_top   = static_cast<int>(pool3d_info.padding.top);
+    const int pad_front = static_cast<int>(pool3d_info.padding.front);
+
+    const int pad_right  = static_cast<int>(pool3d_info.padding.right);
+    const int pad_bottom = static_cast<int>(pool3d_info.padding.bottom);
+    const int pad_back   = static_cast<int>(pool3d_info.padding.back);
+
+    const int num_channels = static_cast<int>(src.shape()[idx_channel]);
+    const int num_batches  = static_cast<int>(src.shape()[idx_batch]);
+
+    ARM_COMPUTE_ERROR_ON(num_channels != static_cast<int>(dst.shape()[idx_channel]));
+    ARM_COMPUTE_ERROR_ON(num_batches != static_cast<int>(dst.shape()[idx_batch]));
+
+    const int w_src = static_cast<int>(src.shape()[idx_width]);
+    const int h_src = static_cast<int>(src.shape()[idx_height]);
+    const int d_src = static_cast<int>(src.shape()[idx_depth]);
+    const int w_dst = static_cast<int>(dst.shape()[idx_width]);
+    const int h_dst = static_cast<int>(dst.shape()[idx_height]);
+    const int d_dst = static_cast<int>(dst.shape()[idx_depth]);
+
+    const bool exclude_padding = pool3d_info.exclude_padding;
+
+    const int height_stride_src = num_channels * w_src;
+    const int depth_stride_src  = height_stride_src * h_src;
+    const int batch_stride_src  = depth_stride_src * d_src;
+    const int height_stride_dst = num_channels * w_dst;
+    const int depth_stride_dst  = height_stride_dst * h_dst;
+    const int batch_stride_dst  = depth_stride_dst * d_dst;
+
+    for(int b = 0; b < num_batches; ++b)
+    {
+        const int batch_offset_dst = b * batch_stride_dst;
+        const int batch_offset_src = b * batch_stride_src;
+        for(int c = 0; c < num_channels; ++c)
+        {
+            for(int d = 0; d < d_dst; ++d)
+            {
+                const int depth_offset_dst = d * depth_stride_dst;
+                for(int h = 0; h < h_dst; ++h)
+                {
+                    const int height_offset_dst = h * height_stride_dst;
+                    for(int w = 0; w < w_dst; ++w)
+                    {
+                        int wstart = w * pool_stride_width - pad_left;
+                        int hstart = h * pool_stride_height - pad_top;
+                        int dstart = d * pool_stride_depth - pad_front;
+                        int wend   = std::min(wstart + pool_size_width, w_src + pad_right);
+                        int hend   = std::min(hstart + pool_size_height, h_src + pad_bottom);
+                        int dend   = std::min(dstart + pool_size_depth, d_src + pad_back);
+
+                        // this may not be equal to pool_w * pool_h * pool_d because of
+                        // DimensionRoundingType choice (CEIL)
+                        int pool_size = (dend - dstart) * (hend - hstart) * (wend - wstart);
+
+                        // limit [start, end) to [0, w_src)
+                        wstart = std::max(wstart, 0);
+                        hstart = std::max(hstart, 0);
+                        dstart = std::max(dstart, 0);
+                        wend   = std::min(wend, w_src);
+                        hend   = std::min(hend, h_src);
+                        dend   = std::min(dend, d_src);
+
+                        auto max_val = -std::numeric_limits<T>::infinity();
+                        int  max_index{ 0 };
+                        T    avg_val = static_cast<T>(0.f);
+                        T    l2_val  = static_cast<T>(0.f);
+
+                        if(exclude_padding)
+                        {
+                            pool_size = (dend - dstart) * (hend - hstart) * (wend - wstart);
+                        }
+
+                        for(int z = dstart; z < dend; ++z)
+                        {
+                            const int depth_offset_src = z * depth_stride_src;
+                            for(int y = hstart; y < hend; ++y)
+                            {
+                                const int height_offset_src = y * height_stride_src;
+                                for(int x = wstart; x < wend; ++x)
+                                {
+                                    const auto val = static_cast<T>(
+                                                         src[batch_offset_src + depth_offset_src + height_offset_src + x * num_channels + c]);
+
+                                    if(val > max_val)
+                                    {
+                                        max_val   = val;
+                                        max_index = coord2index(src.shape(), Coordinates(c, x, y, z, 0));
+                                    }
+
+                                    avg_val += val;
+                                    l2_val += val * val;
+                                }
+                            }
+                        }
+
+                        avg_val /= pool_size;
+                        l2_val = static_cast<T>(std::sqrt(l2_val / pool_size));
+
+                        int dst_index = batch_offset_dst + depth_offset_dst + height_offset_dst + w * num_channels + c;
+                        switch(pool3d_info.pool_type)
+                        {
+                            case PoolingType::MAX:
+                                dst[dst_index] = static_cast<T>(max_val);
+                                break;
+                            case PoolingType::AVG:
+                                dst[dst_index] = static_cast<T>(avg_val);
+                                break;
+                            case PoolingType::L2:
+                                dst[dst_index] = static_cast<T>(l2_val);
+                                break;
+                            default:
+                                ARM_COMPUTE_ERROR("Pooling Type should be either MAX, AVG or L2");
+                        }
+
+                        if(indices != nullptr)
+                        {
+                            (*indices)[dst_index] = max_index;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    return dst;
+}
+
+template SimpleTensor<float> pool3d(const SimpleTensor<float> &src, const Pool3DInfo &pool3d_info, const QuantizationInfo &output_qinfo, SimpleTensor<uint32_t> *indices);
+template SimpleTensor<half> pool3d(const SimpleTensor<half> &src, const Pool3DInfo &pool3d_info, const QuantizationInfo &output_qinfo, SimpleTensor<uint32_t> *indices);
+
+template <typename T>
+SimpleTensor<T> pool3d(const SimpleTensor<T> &src, const Pool3DInfo &pool3d_info, const QuantizationInfo &output_qinfo, SimpleTensor<uint32_t> *indices)
+{
+    ARM_COMPUTE_UNUSED(output_qinfo);
+    return pool3d_internal<T>(src, pool3d_info, indices);
+}
+
+template <>
+SimpleTensor<int8_t> pool3d<int8_t>(const SimpleTensor<int8_t> &src, const Pool3DInfo &pool3d_info, const QuantizationInfo &output_qinfo, SimpleTensor<uint32_t> *indices)
+{
+    SimpleTensor<float> src_tmp = convert_from_asymmetric(src);
+    SimpleTensor<float> dst_tmp = pool3d_internal<float>(src_tmp, pool3d_info, indices);
+    return convert_to_asymmetric<int8_t>(dst_tmp, output_qinfo);
+}
+
+template <>
+SimpleTensor<uint8_t> pool3d<uint8_t>(const SimpleTensor<uint8_t> &src, const Pool3DInfo &pool3d_info, const QuantizationInfo &output_qinfo, SimpleTensor<uint32_t> *indices)
+{
+    SimpleTensor<float> src_tmp = convert_from_asymmetric(src);
+    SimpleTensor<float> dst_tmp = pool3d_internal<float>(src_tmp, pool3d_info, indices);
+    return convert_to_asymmetric<uint8_t>(dst_tmp, output_qinfo);
+}
+
+} // namespace reference
+} // namespace validation
+} // namespace test
+} // namespace arm_compute
\ No newline at end of file
diff --git a/tests/validation/reference/Pool3D.h b/tests/validation/reference/Pool3D.h
new file mode 100644
index 0000000000..bdb5744ecc
--- /dev/null
+++ b/tests/validation/reference/Pool3D.h
@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) 2022 Arm Limited.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to
+ * deal in the Software without restriction, including without limitation the
+ * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+ * sell copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ */
+#ifndef ARM_COMPUTE_TEST_POOL3D_LAYER_H
+#define ARM_COMPUTE_TEST_POOL3D_LAYER_H
+
+#include "Utils.h"
+#include "arm_compute/core/Types.h"
+#include "tests/SimpleTensor.h"
+#include "tests/validation/Helpers.h"
+
+namespace arm_compute
+{
+namespace test
+{
+namespace validation
+{
+namespace reference
+{
+template <typename T>
+SimpleTensor<T> pool3d_internal(const SimpleTensor<T> &src, const Pool3DInfo &pool3d_info, SimpleTensor<uint32_t> *indices);
+
+template <typename T>
+SimpleTensor<T> pool3d(const SimpleTensor<T> &src, const Pool3DInfo &pool3d_info, const QuantizationInfo &output_qinfo, SimpleTensor<uint32_t> *indices);
+} // namespace reference
+} // namespace validation
+} // namespace test
+} // namespace arm_compute
+#endif /* ARM_COMPUTE_TEST_POOL3D_LAYER_H */