src/dynamic_fusion/sketch/gpu/ckw_driver/components/GpuCkwPool2d.cpp - ml/ComputeLibrary - Gitiles

 /*
  * Copyright (c) 2023 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 "src/dynamic_fusion/sketch/gpu/ckw_driver/components/GpuCkwPool2d.h"

 #include "arm_compute/core/Error.h"
 #include "arm_compute/core/Validate.h"
 #include "arm_compute/core/utils/helpers/AdjustVecSize.h"
 #include "ckw/TensorTileSampler.h"
 #include "src/core/helpers/WindowHelpers.h"
 #include "src/dynamic_fusion/sketch/gpu/GpuKernelArgument.h"
 #include "src/dynamic_fusion/sketch/gpu/GpuKernelComponentGroup.h"
 #include "src/dynamic_fusion/sketch/gpu/ckw_driver/GpuCkwKernelWriter.h"
 #include "src/dynamic_fusion/sketch/gpu/ckw_driver/GpuCkwScopedKernelWriter.h"
 #include "src/dynamic_fusion/sketch/gpu/ckw_driver/GpuCkwVariableTable.h"
 #include "src/dynamic_fusion/sketch/gpu/ckw_driver/components/utils/WriterHelper.h"
 #include "src/dynamic_fusion/sketch/gpu/ckw_driver/components/utils/type_converter/Common.h"

 using namespace ckw;

 namespace arm_compute
 {
 namespace experimental
 {
 namespace dynamic_fusion
 {
 GpuCkwPool2d::GpuCkwPool2d(ComponentId                      id,
                            const ArgumentPack<ITensorInfo> &tensors,
                            const Attributes                &attributes,
                            const Settings                  &settings)
     : IGpuCkwComponentDriver{ id, tensors },
       _src{},
       _dst{},
       _attributes{ attributes },
       _settings{ settings }

 {
     _src = this->tensors().get_const_tensor(TensorType::ACL_SRC_0);
     _dst = this->tensors().get_const_tensor(TensorType::ACL_DST_0);
     ARM_COMPUTE_ERROR_ON_NULLPTR(_src, _dst);
 }

 void GpuCkwPool2d::write_component_code(const ComponentGroup &comp_group, GpuCkwVariableTable &vtable, GpuCkwScopedKernelWriter writer) const
 {
     const auto         root_window = comp_group.get_root_component()->ckw_component_driver()->get_window();
     const unsigned int n0          = root_window.x().step();
     const unsigned int m0          = root_window.y().step();

     GpuCkwComponentArgument *src = vtable.declare_variable(comp_group, writer, _src, TensorStorageType::ClBufferUint8Ptr, "src");
     GpuCkwComponentArgument *dst = vtable.declare_variable(comp_group, writer, _dst, TensorStorageType::ClBufferUint8Ptr, "dst");

     TileOperand &gid_0 = writer->declare_tile("gid_0", ckw::DataType::Int32);
     TileOperand &gid_1 = writer->declare_tile("gid_1", ckw::DataType::Int32);
     TileOperand &gid_2 = writer->declare_tile("gid_2", ckw::DataType::Int32);

     writer->op_get_global_id(gid_0, 0);
     writer->op_get_global_id(gid_1, 1);
     writer->op_get_global_id(gid_2, 2);

     // Data Layout is NHWC
     constexpr int width_idx  = 1;
     constexpr int height_idx = 2;

     const int32_t pool_size_x   = static_cast<int32_t>(_attributes.pool_size().x());
     const int32_t pool_size_y   = static_cast<int32_t>(_attributes.pool_size().y());
     const int32_t pad_x         = static_cast<int32_t>(_attributes.pad().left);
     const int32_t pad_y         = static_cast<int32_t>(_attributes.pad().top);
     const int32_t src_width     = static_cast<int32_t>(_src->dimension(width_idx));
     const int32_t src_height    = static_cast<int32_t>(_src->dimension(height_idx));
     const auto    src_data_type = _src->data_type();

     // Check if this is global pooling path
     const bool is_global_pooling = (pool_size_x == src_width) && (pool_size_y == src_height) && (pad_x == 0) && (pad_y == 0);
     // Check if this a case of FP_MIXED_PRECISION
     const bool use_fp_mixed_precision = (src_data_type == DataType::F16) && _settings.mixed_precision() && _attributes.pool_type() != PoolingType::MAX;
     const auto acc_data_type          = (use_fp_mixed_precision) ? (DataType::F32) : (src_data_type);

     TileOperand       &const_0            = writer->declare_tile("0", 0);
     const TileOperand &const_1            = writer->declare_tile("1", 1);
     const TileOperand &const_lowest_value = writer->declare_tile("LOWEST_VALUE", std::numeric_limits<float>::lowest());
     const TileOperand &pool_size_x_tile   = writer->declare_tile("POOL_SIZE_X", pool_size_x);
     const TileOperand &pool_size_y_tile   = writer->declare_tile("POOL_SIZE_Y", pool_size_y);
     const TileOperand &stride_x_tile      = writer->declare_tile("STRIDE_X", static_cast<int32_t>(_attributes.stride().x()));
     const TileOperand &stride_y_tile      = writer->declare_tile("STRIDE_Y", static_cast<int32_t>(_attributes.stride().y()));
     const TileOperand &pad_x_tile         = writer->declare_tile("PAD_X", pad_x);
     const TileOperand &pad_y_tile         = writer->declare_tile("PAD_Y", pad_y);
     const TileOperand &dst_height_tile    = writer->declare_tile("DST_HEIGHT", static_cast<int32_t>(_dst->dimension(height_idx)));
     const TileOperand &src_height_tile    = writer->declare_tile("SRC_HEIGHT", src_height);
     const TileOperand &src_width_tile     = writer->declare_tile("SRC_WIDTH", src_width);

     TileOperand &idx_out_n = writer->declare_tile("idx_out_n", ckw::DataType::Int32);
     TileOperand &idx_out_h = writer->declare_tile("idx_out_h", ckw::DataType::Int32);
     TileOperand &idx_out_w = writer->declare_tile("idx_out_w", ckw::DataType::Int32);
     TileOperand &idx_out_c = writer->declare_tile("idx_out_c", ckw::DataType::Int32);

     const int32_t dst_partial_n0_v = _dst->tensor_shape()[0] % n0;

     get_coord(writer, idx_out_c, gid_0, n0, dst_partial_n0_v, "dst_x_", const_0);
     get_coord(writer, idx_out_w, gid_1, 1, 0, "dst_y_", const_0);

     writer->op_binary_expression(idx_out_h, gid_2, BinaryOp::Mod, dst_height_tile); // gid_2 % h
     writer->op_binary_expression(idx_out_n, gid_2, BinaryOp::Div, dst_height_tile); // gid_2 / h

     TensorTileSampler src_sampler;
     src_sampler.width(n0);
     src_sampler.height(m0);
     src_sampler.format(TensorSamplerFormat::C_W_H);
     src_sampler.address_mode_x(TensorSamplerAddressModeX::None);
     src_sampler.address_mode_y(TensorSamplerAddressModeY::None);
     src_sampler.address_mode_z(TensorSamplerAddressModeZ::None);
     src_sampler.x(idx_out_c);
     src_sampler.b(idx_out_n);

     TensorTileSampler dst_sampler;
     dst_sampler.width(n0);
     dst_sampler.height(m0);
     dst_sampler.format(TensorSamplerFormat::C_W_H);
     dst_sampler.address_mode_x(TensorSamplerAddressModeX::OverlappingMin);
     dst_sampler.address_mode_y(TensorSamplerAddressModeY::None);
     dst_sampler.address_mode_z(TensorSamplerAddressModeZ::None);
     dst_sampler.x(idx_out_c);
     dst_sampler.y(idx_out_w);
     dst_sampler.z(idx_out_h);
     dst_sampler.b(idx_out_n);

     // Prepare dst tensor and tile
     TileInfo dst_tile_info = TileInfo(to_ckw(src_data_type), m0, n0);
     if(!dst->has_tile())
     {
         TileOperand &dst_tile = writer->declare_tile("dst_tile", dst_tile_info);
         dst->init_virtual_tensor(dst_tile, dst_sampler);
     }
     const TileOperand &dst_tile = dst->tile();

     // A tile used to temporarily store results or as an accumulator in case of AVG and L2 pooling.
     const TileOperand &res_tile = writer->declare_tile("res_tile", TileInfo(to_ckw(acc_data_type), m0, n0));

     // Initialise result tile with appropriate value
     if(_attributes.pool_type() == PoolingType::MAX)
     {
         if(_settings.use_inf_as_limit())
         {
             TileContainer            minus_inf_tile_container;
             std::vector<std::string> value = std::vector<std::string>(n0, "(-INFINITY)");
             minus_inf_tile_container.push_back({ value });
             const TileOperand &minus_inf = writer->declare_tile("minus_inf_const", minus_inf_tile_container, to_ckw(acc_data_type));
             writer->op_assign(res_tile, minus_inf);
         }
         else
         {
             writer->op_assign(res_tile, const_lowest_value);
         }
     }
     else
     {
         writer->op_assign(res_tile, const_0);
     }

     // idx_in_w = idx_out_w * STRIDE_X - PAD_X
     TileOperand &idx_in_w = writer->declare_tile("idx_in_w", ckw::DataType::Int32);
     writer->op_binary_expression(idx_in_w, idx_out_w, BinaryOp::Mul, stride_x_tile);
     writer->op_binary_expression(idx_in_w, idx_in_w, BinaryOp::Sub, pad_x_tile);

     // idx_in_h = idx_out_h * STRIDE_Y - PAD_Y
     TileOperand &idx_in_h = writer->declare_tile("idx_in_h", ckw::DataType::Int32);
     writer->op_binary_expression(idx_in_h, idx_out_h, BinaryOp::Mul, stride_y_tile);
     writer->op_binary_expression(idx_in_h, idx_in_h, BinaryOp::Sub, pad_y_tile);

     TileOperand &minus_idx_in_w = writer->declare_tile("minus_idx_in_w", ckw::DataType::Int32);
     TileOperand &minus_idx_in_h = writer->declare_tile("minus_idx_in_h", ckw::DataType::Int32);

     writer->op_unary_expression(minus_idx_in_w, UnaryOp::Negate, idx_in_w);
     writer->op_unary_expression(minus_idx_in_h, UnaryOp::Negate, idx_in_h);

     // Pooling starting/ending offsets for X dim
     TileOperand &pool_x_s = writer->declare_tile("pool_x_s", ckw::DataType::Int32);
     TileOperand &pool_x_e = writer->declare_tile("pool_x_e", ckw::DataType::Int32);

     writer->op_binary_elementwise_function(pool_x_s, BinaryFunction::Max, const_0, minus_idx_in_w);
     writer->op_binary_expression(pool_x_e, src_width_tile, BinaryOp::Add, minus_idx_in_w);
     writer->op_binary_elementwise_function(pool_x_e, BinaryFunction::Min, pool_size_x_tile, pool_x_e);

     // Pooling starting/ending offsets for Y dim
     TileOperand &pool_y_s = writer->declare_tile("pool_y_s", ckw::DataType::Int32);
     TileOperand &pool_y_e = writer->declare_tile("pool_y_e", ckw::DataType::Int32);

     writer->op_binary_elementwise_function(pool_y_s, BinaryFunction::Max, const_0, minus_idx_in_h);
     writer->op_binary_expression(pool_y_e, src_height_tile, BinaryOp::Add, minus_idx_in_h);
     writer->op_binary_elementwise_function(pool_y_e, BinaryFunction::Min, pool_size_y_tile, pool_y_e);

     const TileOperand &filter_size = writer->declare_tile("filter_size", ckw::DataType::Int32);
     if(_attributes.exclude_padding())
     {
         const TileOperand &y_diff = writer->declare_tile("y_diff", ckw::DataType::Int32);
         const TileOperand &x_diff = writer->declare_tile("x_diff", ckw::DataType::Int32);

         writer->op_binary_expression(y_diff, pool_y_e, BinaryOp::Sub, pool_y_s);
         writer->op_binary_expression(x_diff, pool_x_e, BinaryOp::Sub, pool_x_s);

         writer->op_binary_expression(filter_size, y_diff, BinaryOp::Mul, x_diff);
     }
     else
     {
         writer->op_binary_expression(filter_size, pool_size_x_tile, BinaryOp::Mul, pool_size_y_tile);
     }

     const TileOperand &x = writer->declare_tile("x", ckw::DataType::Int32);
     const TileOperand &y = writer->declare_tile("y", ckw::DataType::Int32);

     if(is_global_pooling)
     {
         writer->op_assign(x, const_0);
         writer->op_assign(y, const_0);

         writer->op_assign(pool_y_e, pool_size_y_tile);
         writer->op_assign(pool_x_e, pool_size_x_tile);
     }
     else
     {
         writer->op_assign(x, pool_x_s);
         writer->op_assign(y, pool_y_s);
     }

     // Y dim for-loop
     writer->op_for_loop(y, BinaryOp::Less, pool_y_e, y, AssignmentOp::Increment, const_1, [&]()
     {
         // Reset the iterator for the inner loop
         if(is_global_pooling)
         {
             writer->op_assign(x, const_0);
         }
         else
         {
             writer->op_assign(x, pool_x_s);
         }

         TileOperand &a_y = writer->declare_tile("a_y", ckw::DataType::Int32);
         writer->op_binary_expression(a_y, idx_in_h, BinaryOp::Add, y);

         // X dim for-loop
         writer->op_for_loop(x, BinaryOp::Less, pool_x_e, x, AssignmentOp::Increment, const_1, [&]()
         {
             TileOperand &a_x = writer->declare_tile("a_x", ckw::DataType::Int32);
             writer->op_binary_expression(a_x, idx_in_w, BinaryOp::Add, x);

             TileOperand &src_tile = writer->declare_tile("src_tile", TileInfo(to_ckw(acc_data_type), m0, n0));

             src_sampler.y(a_x);
             src_sampler.z(a_y);

             // Load src tile
             if(use_fp_mixed_precision)
             {
                 TileOperand &src_uncasted_tile = writer->declare_tile("uncasted_src_tile", dst_tile_info);
                 writer->op_load(src_uncasted_tile, src->tensor(), src_sampler);
                 writer->op_cast_expression(src_tile, src_uncasted_tile, ckw::ConvertPolicy::None);
             }
             else
             {
                 writer->op_load(src_tile, src->tensor(), src_sampler);
             }

             // Take the square of the input, for L2 Pooling
             if(_attributes.pool_type() == PoolingType::L2)
             {
                 writer->op_binary_expression(src_tile, src_tile, BinaryOp::Mul, src_tile);
             }

             // Perfom Pooling op
             if(_attributes.pool_type() == PoolingType::MAX)
             {
                 writer->op_binary_elementwise_function(res_tile, BinaryFunction::Max, res_tile, src_tile);
             }
             else
             {
                 writer->op_binary_expression(res_tile, res_tile, BinaryOp::Add, src_tile);
             }
         });
     });

     if((_attributes.pool_type() == PoolingType::AVG) || (_attributes.pool_type() == PoolingType::L2))
     {
         // filter_size is automatically broadcasted in the operation
         writer->op_binary_expression(res_tile, res_tile, BinaryOp::Div, filter_size);
     }

     // Take square root of the result in L2 pooling
     if(_attributes.pool_type() == PoolingType::L2)
     {
         writer->op_unary_elementwise_function(res_tile, UnaryFunction::Sqrt, res_tile);
     }

     // Store the results and do casting if FP_MIXED_PRECISION
     if(use_fp_mixed_precision)
     {
         writer->op_cast_expression(dst_tile, res_tile, ckw::ConvertPolicy::None);
     }
     else
     {
         writer->op_assign(dst_tile, res_tile);
     }
 }

 Window GpuCkwPool2d::get_window() const
 {
     ARM_COMPUTE_ERROR_ON_MSG(_dst->tensor_shape().total_size() == 0U, "Destination tensor is not initialized");

     TensorShape        output_shape = _dst->tensor_shape();
     const unsigned int vec_size     = adjust_vec_size(((_dst->data_type() == DataType::F32) ? 2 : 4), _dst->dimension(0));
     // Create and configure kernel window
     auto win = calculate_max_window(output_shape, Steps(vec_size));
     win      = win.collapse_if_possible(win, Window::DimZ); // collapse window on batch size.
     return win;
 }

 std::string GpuCkwPool2d::get_name(const ComponentGroup &comp_group) const
 {
     ARM_COMPUTE_UNUSED(comp_group);

     return "pool2dMxN";
 }

 } // namespace dynamic_fusion
 } // namespace experimental
 } // namespace arm_compute
	/*
	* Copyright (c) 2023 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 "src/dynamic_fusion/sketch/gpu/ckw_driver/components/GpuCkwPool2d.h"

	#include "arm_compute/core/Error.h"
	#include "arm_compute/core/Validate.h"
	#include "arm_compute/core/utils/helpers/AdjustVecSize.h"
	#include "ckw/TensorTileSampler.h"
	#include "src/core/helpers/WindowHelpers.h"
	#include "src/dynamic_fusion/sketch/gpu/GpuKernelArgument.h"
	#include "src/dynamic_fusion/sketch/gpu/GpuKernelComponentGroup.h"
	#include "src/dynamic_fusion/sketch/gpu/ckw_driver/GpuCkwKernelWriter.h"
	#include "src/dynamic_fusion/sketch/gpu/ckw_driver/GpuCkwScopedKernelWriter.h"
	#include "src/dynamic_fusion/sketch/gpu/ckw_driver/GpuCkwVariableTable.h"
	#include "src/dynamic_fusion/sketch/gpu/ckw_driver/components/utils/WriterHelper.h"
	#include "src/dynamic_fusion/sketch/gpu/ckw_driver/components/utils/type_converter/Common.h"

	using namespace ckw;

	namespace arm_compute
	{
	namespace experimental
	{
	namespace dynamic_fusion
	{
	GpuCkwPool2d::GpuCkwPool2d(ComponentId id,
	const ArgumentPack<ITensorInfo> &tensors,
	const Attributes &attributes,
	const Settings &settings)
	: IGpuCkwComponentDriver{ id, tensors },
	_src{},
	_dst{},
	_attributes{ attributes },
	_settings{ settings }

	{
	_src = this->tensors().get_const_tensor(TensorType::ACL_SRC_0);
	_dst = this->tensors().get_const_tensor(TensorType::ACL_DST_0);
	ARM_COMPUTE_ERROR_ON_NULLPTR(_src, _dst);
	}

	void GpuCkwPool2d::write_component_code(const ComponentGroup &comp_group, GpuCkwVariableTable &vtable, GpuCkwScopedKernelWriter writer) const
	{
	const auto root_window = comp_group.get_root_component()->ckw_component_driver()->get_window();
	const unsigned int n0 = root_window.x().step();
	const unsigned int m0 = root_window.y().step();

	GpuCkwComponentArgument *src = vtable.declare_variable(comp_group, writer, _src, TensorStorageType::ClBufferUint8Ptr, "src");
	GpuCkwComponentArgument *dst = vtable.declare_variable(comp_group, writer, _dst, TensorStorageType::ClBufferUint8Ptr, "dst");

	TileOperand &gid_0 = writer->declare_tile("gid_0", ckw::DataType::Int32);
	TileOperand &gid_1 = writer->declare_tile("gid_1", ckw::DataType::Int32);
	TileOperand &gid_2 = writer->declare_tile("gid_2", ckw::DataType::Int32);

	writer->op_get_global_id(gid_0, 0);
	writer->op_get_global_id(gid_1, 1);
	writer->op_get_global_id(gid_2, 2);

	// Data Layout is NHWC
	constexpr int width_idx = 1;
	constexpr int height_idx = 2;

	const int32_t pool_size_x = static_cast<int32_t>(_attributes.pool_size().x());
	const int32_t pool_size_y = static_cast<int32_t>(_attributes.pool_size().y());
	const int32_t pad_x = static_cast<int32_t>(_attributes.pad().left);
	const int32_t pad_y = static_cast<int32_t>(_attributes.pad().top);
	const int32_t src_width = static_cast<int32_t>(_src->dimension(width_idx));
	const int32_t src_height = static_cast<int32_t>(_src->dimension(height_idx));
	const auto src_data_type = _src->data_type();

	// Check if this is global pooling path
	const bool is_global_pooling = (pool_size_x == src_width) && (pool_size_y == src_height) && (pad_x == 0) && (pad_y == 0);
	// Check if this a case of FP_MIXED_PRECISION
	const bool use_fp_mixed_precision = (src_data_type == DataType::F16) && _settings.mixed_precision() && _attributes.pool_type() != PoolingType::MAX;
	const auto acc_data_type = (use_fp_mixed_precision) ? (DataType::F32) : (src_data_type);

	TileOperand &const_0 = writer->declare_tile("0", 0);
	const TileOperand &const_1 = writer->declare_tile("1", 1);
	const TileOperand &const_lowest_value = writer->declare_tile("LOWEST_VALUE", std::numeric_limits<float>::lowest());
	const TileOperand &pool_size_x_tile = writer->declare_tile("POOL_SIZE_X", pool_size_x);
	const TileOperand &pool_size_y_tile = writer->declare_tile("POOL_SIZE_Y", pool_size_y);
	const TileOperand &stride_x_tile = writer->declare_tile("STRIDE_X", static_cast<int32_t>(_attributes.stride().x()));
	const TileOperand &stride_y_tile = writer->declare_tile("STRIDE_Y", static_cast<int32_t>(_attributes.stride().y()));
	const TileOperand &pad_x_tile = writer->declare_tile("PAD_X", pad_x);
	const TileOperand &pad_y_tile = writer->declare_tile("PAD_Y", pad_y);
	const TileOperand &dst_height_tile = writer->declare_tile("DST_HEIGHT", static_cast<int32_t>(_dst->dimension(height_idx)));
	const TileOperand &src_height_tile = writer->declare_tile("SRC_HEIGHT", src_height);
	const TileOperand &src_width_tile = writer->declare_tile("SRC_WIDTH", src_width);

	TileOperand &idx_out_n = writer->declare_tile("idx_out_n", ckw::DataType::Int32);
	TileOperand &idx_out_h = writer->declare_tile("idx_out_h", ckw::DataType::Int32);
	TileOperand &idx_out_w = writer->declare_tile("idx_out_w", ckw::DataType::Int32);
	TileOperand &idx_out_c = writer->declare_tile("idx_out_c", ckw::DataType::Int32);

	const int32_t dst_partial_n0_v = _dst->tensor_shape()[0] % n0;

	get_coord(writer, idx_out_c, gid_0, n0, dst_partial_n0_v, "dst_x_", const_0);
	get_coord(writer, idx_out_w, gid_1, 1, 0, "dst_y_", const_0);

	writer->op_binary_expression(idx_out_h, gid_2, BinaryOp::Mod, dst_height_tile); // gid_2 % h
	writer->op_binary_expression(idx_out_n, gid_2, BinaryOp::Div, dst_height_tile); // gid_2 / h

	TensorTileSampler src_sampler;
	src_sampler.width(n0);
	src_sampler.height(m0);
	src_sampler.format(TensorSamplerFormat::C_W_H);
	src_sampler.address_mode_x(TensorSamplerAddressModeX::None);
	src_sampler.address_mode_y(TensorSamplerAddressModeY::None);
	src_sampler.address_mode_z(TensorSamplerAddressModeZ::None);
	src_sampler.x(idx_out_c);
	src_sampler.b(idx_out_n);

	TensorTileSampler dst_sampler;
	dst_sampler.width(n0);
	dst_sampler.height(m0);
	dst_sampler.format(TensorSamplerFormat::C_W_H);
	dst_sampler.address_mode_x(TensorSamplerAddressModeX::OverlappingMin);
	dst_sampler.address_mode_y(TensorSamplerAddressModeY::None);
	dst_sampler.address_mode_z(TensorSamplerAddressModeZ::None);
	dst_sampler.x(idx_out_c);
	dst_sampler.y(idx_out_w);
	dst_sampler.z(idx_out_h);
	dst_sampler.b(idx_out_n);

	// Prepare dst tensor and tile
	TileInfo dst_tile_info = TileInfo(to_ckw(src_data_type), m0, n0);
	if(!dst->has_tile())
	{
	TileOperand &dst_tile = writer->declare_tile("dst_tile", dst_tile_info);
	dst->init_virtual_tensor(dst_tile, dst_sampler);
	}
	const TileOperand &dst_tile = dst->tile();

	// A tile used to temporarily store results or as an accumulator in case of AVG and L2 pooling.
	const TileOperand &res_tile = writer->declare_tile("res_tile", TileInfo(to_ckw(acc_data_type), m0, n0));

	// Initialise result tile with appropriate value
	if(_attributes.pool_type() == PoolingType::MAX)
	{
	if(_settings.use_inf_as_limit())
	{
	TileContainer minus_inf_tile_container;
	std::vector<std::string> value = std::vector<std::string>(n0, "(-INFINITY)");
	minus_inf_tile_container.push_back({ value });
	const TileOperand &minus_inf = writer->declare_tile("minus_inf_const", minus_inf_tile_container, to_ckw(acc_data_type));
	writer->op_assign(res_tile, minus_inf);
	}
	else
	{
	writer->op_assign(res_tile, const_lowest_value);
	}
	}
	else
	{
	writer->op_assign(res_tile, const_0);
	}

	// idx_in_w = idx_out_w * STRIDE_X - PAD_X
	TileOperand &idx_in_w = writer->declare_tile("idx_in_w", ckw::DataType::Int32);
	writer->op_binary_expression(idx_in_w, idx_out_w, BinaryOp::Mul, stride_x_tile);
	writer->op_binary_expression(idx_in_w, idx_in_w, BinaryOp::Sub, pad_x_tile);

	// idx_in_h = idx_out_h * STRIDE_Y - PAD_Y
	TileOperand &idx_in_h = writer->declare_tile("idx_in_h", ckw::DataType::Int32);
	writer->op_binary_expression(idx_in_h, idx_out_h, BinaryOp::Mul, stride_y_tile);
	writer->op_binary_expression(idx_in_h, idx_in_h, BinaryOp::Sub, pad_y_tile);

	TileOperand &minus_idx_in_w = writer->declare_tile("minus_idx_in_w", ckw::DataType::Int32);
	TileOperand &minus_idx_in_h = writer->declare_tile("minus_idx_in_h", ckw::DataType::Int32);

	writer->op_unary_expression(minus_idx_in_w, UnaryOp::Negate, idx_in_w);
	writer->op_unary_expression(minus_idx_in_h, UnaryOp::Negate, idx_in_h);

	// Pooling starting/ending offsets for X dim
	TileOperand &pool_x_s = writer->declare_tile("pool_x_s", ckw::DataType::Int32);
	TileOperand &pool_x_e = writer->declare_tile("pool_x_e", ckw::DataType::Int32);

	writer->op_binary_elementwise_function(pool_x_s, BinaryFunction::Max, const_0, minus_idx_in_w);
	writer->op_binary_expression(pool_x_e, src_width_tile, BinaryOp::Add, minus_idx_in_w);
	writer->op_binary_elementwise_function(pool_x_e, BinaryFunction::Min, pool_size_x_tile, pool_x_e);

	// Pooling starting/ending offsets for Y dim
	TileOperand &pool_y_s = writer->declare_tile("pool_y_s", ckw::DataType::Int32);
	TileOperand &pool_y_e = writer->declare_tile("pool_y_e", ckw::DataType::Int32);

	writer->op_binary_elementwise_function(pool_y_s, BinaryFunction::Max, const_0, minus_idx_in_h);
	writer->op_binary_expression(pool_y_e, src_height_tile, BinaryOp::Add, minus_idx_in_h);
	writer->op_binary_elementwise_function(pool_y_e, BinaryFunction::Min, pool_size_y_tile, pool_y_e);

	const TileOperand &filter_size = writer->declare_tile("filter_size", ckw::DataType::Int32);
	if(_attributes.exclude_padding())
	{
	const TileOperand &y_diff = writer->declare_tile("y_diff", ckw::DataType::Int32);
	const TileOperand &x_diff = writer->declare_tile("x_diff", ckw::DataType::Int32);

	writer->op_binary_expression(y_diff, pool_y_e, BinaryOp::Sub, pool_y_s);
	writer->op_binary_expression(x_diff, pool_x_e, BinaryOp::Sub, pool_x_s);

	writer->op_binary_expression(filter_size, y_diff, BinaryOp::Mul, x_diff);
	}
	else
	{
	writer->op_binary_expression(filter_size, pool_size_x_tile, BinaryOp::Mul, pool_size_y_tile);
	}

	const TileOperand &x = writer->declare_tile("x", ckw::DataType::Int32);
	const TileOperand &y = writer->declare_tile("y", ckw::DataType::Int32);

	if(is_global_pooling)
	{
	writer->op_assign(x, const_0);
	writer->op_assign(y, const_0);

	writer->op_assign(pool_y_e, pool_size_y_tile);
	writer->op_assign(pool_x_e, pool_size_x_tile);
	}
	else
	{
	writer->op_assign(x, pool_x_s);
	writer->op_assign(y, pool_y_s);
	}

	// Y dim for-loop
	writer->op_for_loop(y, BinaryOp::Less, pool_y_e, y, AssignmentOp::Increment, const_1, [&]()
	{
	// Reset the iterator for the inner loop
	if(is_global_pooling)
	{
	writer->op_assign(x, const_0);
	}
	else
	{
	writer->op_assign(x, pool_x_s);
	}

	TileOperand &a_y = writer->declare_tile("a_y", ckw::DataType::Int32);
	writer->op_binary_expression(a_y, idx_in_h, BinaryOp::Add, y);

	// X dim for-loop
	writer->op_for_loop(x, BinaryOp::Less, pool_x_e, x, AssignmentOp::Increment, const_1, [&]()
	{
	TileOperand &a_x = writer->declare_tile("a_x", ckw::DataType::Int32);
	writer->op_binary_expression(a_x, idx_in_w, BinaryOp::Add, x);

	TileOperand &src_tile = writer->declare_tile("src_tile", TileInfo(to_ckw(acc_data_type), m0, n0));

	src_sampler.y(a_x);
	src_sampler.z(a_y);

	// Load src tile
	if(use_fp_mixed_precision)
	{
	TileOperand &src_uncasted_tile = writer->declare_tile("uncasted_src_tile", dst_tile_info);
	writer->op_load(src_uncasted_tile, src->tensor(), src_sampler);
	writer->op_cast_expression(src_tile, src_uncasted_tile, ckw::ConvertPolicy::None);
	}
	else
	{
	writer->op_load(src_tile, src->tensor(), src_sampler);
	}

	// Take the square of the input, for L2 Pooling
	if(_attributes.pool_type() == PoolingType::L2)
	{
	writer->op_binary_expression(src_tile, src_tile, BinaryOp::Mul, src_tile);
	}

	// Perfom Pooling op
	if(_attributes.pool_type() == PoolingType::MAX)
	{
	writer->op_binary_elementwise_function(res_tile, BinaryFunction::Max, res_tile, src_tile);
	}
	else
	{
	writer->op_binary_expression(res_tile, res_tile, BinaryOp::Add, src_tile);
	}
	});
	});

	if((_attributes.pool_type() == PoolingType::AVG) \|\| (_attributes.pool_type() == PoolingType::L2))
	{
	// filter_size is automatically broadcasted in the operation
	writer->op_binary_expression(res_tile, res_tile, BinaryOp::Div, filter_size);
	}

	// Take square root of the result in L2 pooling
	if(_attributes.pool_type() == PoolingType::L2)
	{
	writer->op_unary_elementwise_function(res_tile, UnaryFunction::Sqrt, res_tile);
	}

	// Store the results and do casting if FP_MIXED_PRECISION
	if(use_fp_mixed_precision)
	{
	writer->op_cast_expression(dst_tile, res_tile, ckw::ConvertPolicy::None);
	}
	else
	{
	writer->op_assign(dst_tile, res_tile);
	}
	}

	Window GpuCkwPool2d::get_window() const
	{
	ARM_COMPUTE_ERROR_ON_MSG(_dst->tensor_shape().total_size() == 0U, "Destination tensor is not initialized");

	TensorShape output_shape = _dst->tensor_shape();
	const unsigned int vec_size = adjust_vec_size(((_dst->data_type() == DataType::F32) ? 2 : 4), _dst->dimension(0));
	// Create and configure kernel window
	auto win = calculate_max_window(output_shape, Steps(vec_size));
	win = win.collapse_if_possible(win, Window::DimZ); // collapse window on batch size.
	return win;
	}

	std::string GpuCkwPool2d::get_name(const ComponentGroup &comp_group) const
	{
	ARM_COMPUTE_UNUSED(comp_group);

	return "pool2dMxN";
	}

	} // namespace dynamic_fusion
	} // namespace experimental
	} // namespace arm_compute