added pool2d fwd

example/13_pool2d_fwd/pool2d_fwd_common.hpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -101,7 +101,7 @@ bool pool_test(bool do_verification,
     switch(init_method)
     {
     case 0: break;
-    case 1: in_n_c_hi_wi.GenerateTensorValue(GeneratorTensor_1<InDataType>{1}); break;
+    case 1: in_n_c_hi_wi.GenerateTensorValue(GeneratorTensor_1<InDataType>{2}); break;
     case 2: in_n_c_hi_wi.GenerateTensorValue(GeneratorTensor_2<InDataType>{-5, 5}); break;
     default: in_n_c_hi_wi.GenerateTensorValue(GeneratorTensor_3<InDataType>{-5.0, 5.0});
     }

example/13_pool2d_fwd/pool2d_fwd_fp16.cpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #include <iostream>
 
@@ -52,7 +52,7 @@ int main(int argc, char* argv[])
     if(argc == 1)
     {
         do_verification = true;
-        init_method     = 1;
+        init_method     = 2;
         time_kernel     = true;
     }
     else if(argc == 4)

example/13_pool2d_fwd/pool2d_fwd_fp32.cpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #include <iostream>
 

include/ck/tensor_operation/gpu/device/impl/device_pool2d_fwd_nhwc_nhwc.hpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
-#include "ck/tensor_operation/gpu/device/impl/device_pool3d_fwd_ndhwc_ndhwc.hpp"
+#include <iostream>
+#include <sstream>
+
+#include "ck/tensor_description/tensor_descriptor.hpp"
+#include "ck/tensor_description/tensor_descriptor_helper.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/reduction_operator_mapping.hpp"
+#include "ck/tensor_operation/gpu/device/device_pool_fwd.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_reduce_common.hpp"
+#include "ck/tensor_operation/gpu/grid/gridwise_2d_reduction_threadwise.hpp"
+#include "ck/host_utility/device_prop.hpp"
+#include "ck/host_utility/kernel_launch.hpp"
 
 namespace ck {
 namespace tensor_operation {
@@ -16,95 +27,363 @@ template <typename InDataType,
           ck::ReduceTensorOp ReduceOpId,
           bool OutputIndex,
           ck::index_t BlockSize,
-          ck::index_t ReduceMThreadClusterSize,
-          ck::index_t ReduceKThreadClusterSize,
-          ck::index_t ReduceMThreadSliceSize,
-          ck::index_t ReduceKThreadSliceSize,
+          ck::index_t MThreadClusterSize,
+          ck::index_t KThreadClusterSize,
+          ck::index_t MThreadSliceSize,
+          ck::index_t KThreadSliceSize,
           ck::index_t InSrcOutDstVectorSize>
-struct DevicePool2dFwd_NHWC_NHWC : public DevicePool3dFwd_NDHWC_NDHWC<InDataType,
-                                                                      OutDataType,
-                                                                      IndexDataType,
-                                                                      ComputeDataType,
-                                                                      ReduceOpId,
-                                                                      OutputIndex,
-                                                                      BlockSize,
-                                                                      ReduceMThreadClusterSize,
-                                                                      ReduceKThreadClusterSize,
-                                                                      ReduceMThreadSliceSize,
-                                                                      ReduceKThreadSliceSize,
-                                                                      InSrcOutDstVectorSize>
+struct DevicePool2dFwd_NHWC_NHWC : public DevicePoolFwd<4,
+                                                        2,
+                                                        InDataType,
+                                                        OutDataType,
+                                                        IndexDataType,
+                                                        tensor_layout::convolution::NHWC,
+                                                        tensor_layout::convolution::NHWC,
+                                                        ReduceOpId,
+                                                        OutputIndex>
 {
-    using DevicePool3D = DevicePool3dFwd_NDHWC_NDHWC<InDataType,
+    static constexpr auto I0 = Number<0>{};
+    static constexpr auto I1 = Number<1>{};
+    static constexpr auto I2 = Number<2>{};
+    static constexpr auto I3 = Number<3>{};
+    static constexpr auto I4 = Number<4>{};
+    static constexpr auto I5 = Number<5>{};
+
+    static constexpr index_t InOutRank  = 4;
+    static constexpr index_t WindowRank = 2;
+
+    using ReduceOperation = typename reduce_binary_operator<ReduceOpId>::opType;
+
+    using InElementwiseOperation =
+        typename reduce_unary_operator<ReduceOpId, true, true>::InElementwiseOperation;
+
+    using AccElementwiseOperation =
+        typename reduce_unary_operator<ReduceOpId, true, true>::AccElementwiseOperation;
+
+    static constexpr ck::index_t M_BlockTileSize = MThreadClusterSize * MThreadSliceSize;
+    static constexpr ck::index_t K_BlockTileSize = KThreadClusterSize * KThreadSliceSize;
+
+    static auto MakeABGridDescriptor_A_M_K_B_M(std::vector<ck::index_t> input_nchw_lengths,
+                                               std::vector<ck::index_t> output_nchw_lengths,
+                                               std::vector<ck::index_t> input_nchw_stride,
+                                               std::vector<ck::index_t> output_nchw_stride,
+                                               std::vector<ck::index_t> window_spatial_yx_lengths,
+                                               std::vector<ck::index_t> window_yx_strides,
+                                               std::vector<ck::index_t> window_yx_dilations,
+                                               std::vector<ck::index_t> input_left_hw_pads,
+                                               std::vector<ck::index_t> input_right_hw_pads)
+    {
+        const index_t N  = input_nchw_lengths[0];
+        const index_t C  = input_nchw_lengths[1];
+        const index_t Hi = input_nchw_lengths[2];
+        const index_t Wi = input_nchw_lengths[3];
+
+        const index_t Ho = output_nchw_lengths[2];
+        const index_t Wo = output_nchw_lengths[3];
+        const index_t Y  = window_spatial_yx_lengths[0];
+        const index_t X  = window_spatial_yx_lengths[1];
+
+        const index_t WindowStrideH = window_yx_strides[0];
+        const index_t WindowStrideW = window_yx_strides[1];
+
+        const index_t WindowDilationH = window_yx_dilations[0];
+        const index_t WindowDilationW = window_yx_dilations[1];
+
+        const index_t InLeftPadH = input_left_hw_pads[0];
+        const index_t InLeftPadW = input_left_hw_pads[1];
+
+        const index_t InRightPadH = input_right_hw_pads[0];
+        const index_t InRightPadW = input_right_hw_pads[1];
+
+        const index_t MRaw = N * Ho * Wo * C;
+        const index_t MPad = math::integer_least_multiple(MRaw, M_BlockTileSize) - MRaw;
+
+        const index_t KRaw = Y * X;
+        const index_t KPad = math::integer_least_multiple(KRaw, K_BlockTileSize) - KRaw;
+
+        // A[ReduceM, ReduceK]
+        const index_t Ni_stride = input_nchw_stride[0];
+        const index_t Ci_stride = input_nchw_stride[1];
+        const index_t Hi_stride = input_nchw_stride[2];
+        const index_t Wi_stride = input_nchw_stride[3];
+
+        const auto in_grid_desc_n_hi_wi_c = make_naive_tensor_descriptor(
+            make_tuple(N, Hi, Wi, C), make_tuple(Ni_stride, Hi_stride, Wi_stride, Ci_stride));
+
+        const auto in_grid_desc_n_hip_wip_c = transform_tensor_descriptor(
+            in_grid_desc_n_hi_wi_c,
+            make_tuple(make_pass_through_transform(N),
+                       make_pad_transform(Hi, InLeftPadH, InRightPadH),
+                       make_pad_transform(Wi, InLeftPadW, InRightPadW),
+                       make_pass_through_transform(C)),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}));
+
+        const auto in_grid_desc_n_y_ho_x_wo_c = transform_tensor_descriptor(
+            in_grid_desc_n_hip_wip_c,
+            make_tuple(
+                make_pass_through_transform(N),
+                make_embed_transform(make_tuple(Y, Ho), make_tuple(WindowDilationH, WindowStrideH)),
+                make_embed_transform(make_tuple(X, Wo), make_tuple(WindowDilationW, WindowStrideW)),
+                make_pass_through_transform(C)),
+            make_tuple(Sequence<0>{}, Sequence<1>{}, Sequence<2>{}, Sequence<3>{}),
+            make_tuple(Sequence<0>{}, Sequence<1, 2>{}, Sequence<3, 4>{}, Sequence<5>{}));
+
+        const auto in_grid_desc_reducemraw_reducekraw =
+            transform_tensor_descriptor(in_grid_desc_n_y_ho_x_wo_c,
+                                        make_tuple(make_merge_transform(make_tuple(N, Ho, Wo, C)),
+                                                   make_merge_transform(make_tuple(Y, X))),
+                                        make_tuple(Sequence<0, 2, 4, 5>{}, Sequence<1, 3>{}),
+                                        make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        const auto in_grid_desc_reducem_reducek = transform_tensor_descriptor(
+            in_grid_desc_reducemraw_reducekraw,
+            make_tuple(make_right_pad_transform(MRaw, MPad), make_right_pad_transform(KRaw, KPad)),
+            make_tuple(Sequence<0>{}, Sequence<1>{}),
+            make_tuple(Sequence<0>{}, Sequence<1>{}));
+
+        // B[ReduceM]
+        const index_t No_stride = output_nchw_stride[0];
+        const index_t Co_stride = output_nchw_stride[1];
+        const index_t Ho_stride = output_nchw_stride[2];
+        const index_t Wo_stride = output_nchw_stride[3];
+
+        const auto out_grid_desc_n_ho_wo_c = make_naive_tensor_descriptor(
+            make_tuple(N, Hi, Wi, C), make_tuple(No_stride, Ho_stride, Wo_stride, Co_stride));
+
+        const auto out_grid_desc_reducemraw =
+            transform_tensor_descriptor(out_grid_desc_n_ho_wo_c,
+                                        make_tuple(make_merge_transform(make_tuple(N, Ho, Wo, C))),
+                                        make_tuple(Sequence<0, 1, 2, 3>{}),
+                                        make_tuple(Sequence<0>{}));
+
+        const auto out_grid_desc_reducem =
+            transform_tensor_descriptor(out_grid_desc_reducemraw,
+                                        make_tuple(make_right_pad_transform(MRaw, MPad)),
+                                        make_tuple(Sequence<0>{}),
+                                        make_tuple(Sequence<0>{}));
+
+        return make_tuple(in_grid_desc_reducem_reducek, out_grid_desc_reducem);
+    }
+
+    using ABGridDescs =
+        decltype(MakeABGridDescriptor_A_M_K_B_M({}, {}, {}, {}, {}, {}, {}, {}, {}));
+
+    using AGridDesc_M_K = remove_cvref_t<decltype(ABGridDescs{}[I0])>;
+    using BGridDesc_M   = remove_cvref_t<decltype(ABGridDescs{}[I1])>;
+
+    struct Argument : public BaseArgument
+    {
+        Argument(const InDataType* p_in_dev,
+                 OutDataType* p_out_dev,
+                 IndexDataType* p_out_indices_dev,
+                 std::vector<ck::index_t>& input_nchw_lengths,
+                 std::vector<ck::index_t>& output_nchw_lengths,
+                 std::vector<ck::index_t>& input_nchw_stride,
+                 std::vector<ck::index_t>& output_nchw_stride,
+                 std::vector<ck::index_t>&, // indices_nchw_stride
+                 std::vector<ck::index_t>& window_spatial_yx_lengths,
+                 std::vector<ck::index_t>& window_yx_strides,
+                 std::vector<ck::index_t>& window_yx_dilations,
+                 std::vector<ck::index_t>& input_left_hw_pads,
+                 std::vector<ck::index_t>& input_right_hw_pads)
+            : p_in_dev_{p_in_dev},
+              p_out_dev_{p_out_dev},
+              p_out_indices_dev_{p_out_indices_dev},
+              a_grid_desc_m_k_{},
+              b_grid_desc_m_{},
+              input_nchw_lengths_{input_nchw_lengths},
+              output_nchw_lengths_{output_nchw_lengths},
+              input_nchw_stride_{input_nchw_stride},
+              output_nchw_stride_{output_nchw_stride}
+        {
+            const auto descs = MakeABGridDescriptor_A_M_K_B_M(input_nchw_lengths,
+                                                              output_nchw_lengths,
+                                                              input_nchw_stride,
+                                                              output_nchw_stride,
+                                                              window_spatial_yx_lengths,
+                                                              window_yx_strides,
+                                                              window_yx_dilations,
+                                                              input_left_hw_pads,
+                                                              input_right_hw_pads);
+
+            a_grid_desc_m_k_ = descs[I0];
+            b_grid_desc_m_   = descs[I1];
+
+            int32_t reduceLength = window_spatial_yx_lengths[0] * window_spatial_yx_lengths[1];
+
+            std::tie(in_element_op_, acc_element_op_) =
+                reduce_unary_operator<ReduceOpId, true, true>::GetElementwiseOperator(reduceLength);
+        }
+
+        const InDataType* p_in_dev_;
+        OutDataType* p_out_dev_;
+        IndexDataType* p_out_indices_dev_;
+        AGridDesc_M_K a_grid_desc_m_k_;
+        BGridDesc_M b_grid_desc_m_;
+
+        InElementwiseOperation in_element_op_;
+        AccElementwiseOperation acc_element_op_;
+
+        // for checking vector load/store
+        std::vector<ck::index_t> input_nchw_lengths_;
+        std::vector<ck::index_t> output_nchw_lengths_;
+        std::vector<ck::index_t> input_nchw_stride_;
+        std::vector<ck::index_t> output_nchw_stride_;
+    };
+
+    struct Invoker : public BaseInvoker
+    {
+        float Run(const Argument& arg, const StreamConfig& stream_config = StreamConfig{})
+        {
+            // for NHWC, the dim C is the fastest dimension, and is not reduced.
+            // Hence, it is in M dimension for reduction kernel.
+            static constexpr index_t InSrcOutDstVectorDim = 0; // 0: M, 1: K
+
+            using gridwise_reduce =
+                GridwiseReduction_mk_to_m_threadwise<InDataType,
                                                      OutDataType,
-                                                     IndexDataType,
                                                      ComputeDataType,
-                                                     ReduceOpId,
-                                                     OutputIndex,
+                                                     IndexDataType,
+                                                     AGridDesc_M_K,
+                                                     BGridDesc_M,
+                                                     ReduceOperation,
+                                                     InElementwiseOperation,
+                                                     AccElementwiseOperation,
+                                                     InMemoryDataOperationEnum::Set,
+                                                     false, // propagate_nan
                                                      BlockSize,
-                                                     ReduceMThreadClusterSize,
-                                                     ReduceKThreadClusterSize,
-                                                     ReduceMThreadSliceSize,
-                                                     ReduceKThreadSliceSize,
+                                                     MThreadSliceSize,
+                                                     KThreadSliceSize,
+                                                     InSrcOutDstVectorDim,
+                                                     InSrcOutDstVectorSize,
                                                      InSrcOutDstVectorSize>;
 
-    std::unique_ptr<BaseArgument>
+            const auto kernel =
+                kernel_reduce_threadwise<gridwise_reduce,
+                                         OutputIndex,
+                                         true,  // pooling need to return global index
+                                         false, // don't have index input
+                                         InDataType,
+                                         OutDataType,
+                                         ComputeDataType,
+                                         IndexDataType,
+                                         AGridDesc_M_K,
+                                         BGridDesc_M,
+                                         InElementwiseOperation,
+                                         AccElementwiseOperation>;
+
+            ck::index_t M = arg.a_grid_desc_m_k_.GetLength(I0);
+
+            const index_t grid_size = (M / M_BlockTileSize);
+
+            return launch_and_time_kernel(stream_config,
+                                          kernel,
+                                          dim3(grid_size),
+                                          dim3(BlockSize),
+                                          0,
+                                          arg.a_grid_desc_m_k_,
+                                          arg.b_grid_desc_m_,
+                                          arg.in_element_op_,
+                                          arg.acc_element_op_,
+                                          float(1),
+                                          arg.p_in_dev_,
+                                          nullptr,
+                                          float(0),
+                                          arg.p_out_dev_,
+                                          arg.p_out_indices_dev_);
+        }
+
+        float Run(const BaseArgument* p_arg,
+                  const StreamConfig& stream_config = StreamConfig{}) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg), stream_config);
+        }
+    };
+
+    bool IsSupportedArgument(const BaseArgument* p_arg) override
+    {
+        const Argument* pArg = dynamic_cast<const Argument*>(p_arg);
+
+        // C should be fastest dimension
+        if(pArg->input_nchw_stride_[1] != 1)
+            return false;
+
+        for(int i = 0; i < InOutRank; ++i)
+        {
+            if(pArg->input_nchw_stride_[i] == 1 &&
+               pArg->input_nchw_lengths_[i] % InSrcOutDstVectorSize != 0)
+                return false;
+
+            if(pArg->output_nchw_stride_[i] == 1 &&
+               pArg->output_nchw_lengths_[i] % InSrcOutDstVectorSize != 0)
+                return false;
+        }
+
+        return true;
+    }
+
+    virtual std::unique_ptr<BaseArgument>
     MakeArgumentPointer(const void* p_in_dev,
                         void* p_out_dev,
                         void* p_out_indices_dev,
-                        std::vector<ck::index_t> input_lengths,
-                        std::vector<ck::index_t> window_lengths,
-                        std::vector<ck::index_t> output_lengths,
-                        std::vector<ck::index_t> input_stride,
-                        std::vector<ck::index_t> output_stride,
-                        std::vector<ck::index_t> indices_stride,
-                        std::vector<ck::index_t> window_strides,
-                        std::vector<ck::index_t> window_dilations,
-                        std::vector<ck::index_t> input_left_pads,
-                        std::vector<ck::index_t> input_right_pads,
+                        std::vector<ck::index_t> input_nchw_lengths,
+                        std::vector<ck::index_t> window_yx_lengths,
+                        std::vector<ck::index_t> output_nchw_lengths,
+                        std::vector<ck::index_t> input_nchw_stride,
+                        std::vector<ck::index_t> output_nchw_stride,
+                        std::vector<ck::index_t> indices_nchw_stride,
+                        std::vector<ck::index_t> window_yx_strides,
+                        std::vector<ck::index_t> window_yx_dilations,
+                        std::vector<ck::index_t> input_left_hw_pads,
+                        std::vector<ck::index_t> input_right_hw_pads,
                         std::vector<ck::index_t> pooling_dims) override
     {
-        static constexpr index_t InOutRank  = 4;
-        static constexpr index_t WindowRank = 2;
-
-        if(input_lengths.size() != InOutRank || window_lengths.size() != WindowRank ||
-           input_lengths.size() != InOutRank || window_strides.size() != WindowRank ||
-           window_dilations.size() != WindowRank || input_left_pads.size() != WindowRank ||
-           input_right_pads.size() != WindowRank)
+        if(input_nchw_lengths.size() != InOutRank || window_yx_lengths.size() != WindowRank ||
+           input_nchw_lengths.size() != InOutRank || window_yx_strides.size() != WindowRank ||
+           window_yx_dilations.size() != WindowRank || input_left_hw_pads.size() != WindowRank ||
+           input_right_hw_pads.size() != WindowRank)
             throw std::runtime_error("dimension is incorrect");
 
         if(pooling_dims != std::vector<ck::index_t>{2, 3})
             throw std::runtime_error("pooling_dims only support {2, 3} in pool2d so far");
 
-        // NCHW to NCDHW
-        input_lengths.insert(input_lengths.begin() + 2, 1);
-        output_lengths.insert(output_lengths.begin() + 2, 1);
-        input_stride.insert(input_stride.begin() + 2, 0);
-        output_stride.insert(output_stride.begin() + 2, 0);
-        indices_stride.insert(indices_stride.begin() + 2, 0);
-
-        // YX to ZYX
-        window_lengths.insert(window_lengths.begin(), 1);
-        window_strides.insert(window_strides.begin(), 0);
-        window_dilations.insert(window_dilations.begin(), 0);
-        input_left_pads.insert(input_left_pads.begin(), 0);
-        input_right_pads.insert(input_right_pads.begin(), 0);
-
-        pooling_dims = {2, 3, 4};
-
-        return DevicePool3D::MakeArgumentPointer(p_in_dev,
-                                                 p_out_dev,
-                                                 p_out_indices_dev,
-                                                 input_lengths,
-                                                 window_lengths,
-                                                 output_lengths,
-                                                 input_stride,
-                                                 output_stride,
-                                                 indices_stride,
-                                                 window_strides,
-                                                 window_dilations,
-                                                 input_left_pads,
-                                                 input_right_pads,
-                                                 pooling_dims);
+        if(output_nchw_stride != indices_nchw_stride)
+            throw std::runtime_error(
+                "output_nchw_stride need to be equal to indices_nchw_stride for now");
+
+        return std::make_unique<Argument>(static_cast<const InDataType*>(p_in_dev),
+                                          static_cast<OutDataType*>(p_out_dev),
+                                          static_cast<IndexDataType*>(p_out_indices_dev),
+                                          input_nchw_lengths,
+                                          output_nchw_lengths,
+                                          input_nchw_stride,
+                                          output_nchw_stride,
+                                          indices_nchw_stride,
+                                          window_yx_lengths,
+                                          window_yx_strides,
+                                          window_yx_dilations,
+                                          input_left_hw_pads,
+                                          input_right_hw_pads);
+    }
+
+    std::unique_ptr<BaseInvoker> MakeInvokerPointer() override
+    {
+        return std::make_unique<Invoker>(Invoker{});
+    }
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "DevicePool2dFwd_NHWC_NHWC<" << BlockSize << ",";
+        str << "M_C" << MThreadClusterSize << "_S" << MThreadSliceSize << ",";
+        str << "K_C" << KThreadClusterSize << "_S" << KThreadSliceSize << ",";
+        str <<"InSrcOutDstVectorSize_" << InSrcOutDstVectorSize << ">";
+        // clang-format on
+
+        return str.str();
     }
 };
 

library/include/ck/library/tensor_operation_instance/gpu/pool2d_fwd.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <cstdlib>
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "ck/tensor_operation/gpu/device/device_pool_fwd.hpp"
+#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
+
+#include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+static constexpr auto InOutRank  = 4;
+static constexpr auto WindowRank = 2;
+
+static constexpr auto MaxOp = ck::ReduceTensorOp::MAX;
+static constexpr auto AvgOp = ck::ReduceTensorOp::AVG;
+#ifdef CK_ENABLE_FP16
+// FP16
+void add_device_pool2d_fwd_nhwc_f16_instances(
+    std::vector<std::unique_ptr<
+        DevicePoolFwd<InOutRank, WindowRank, F16, F16, I32, NHWC, NHWC, MaxOp, false>>>&);
+
+void add_device_pool2d_fwd_nhwc_f16_instances(
+    std::vector<std::unique_ptr<
+        DevicePoolFwd<InOutRank, WindowRank, F16, F16, I32, NHWC, NHWC, AvgOp, false>>>&);
+
+// FP16 - return index
+void add_device_pool2d_fwd_nhwc_index_f16_instances(
+    std::vector<std::unique_ptr<
+        DevicePoolFwd<InOutRank, WindowRank, F16, F16, I32, NHWC, NHWC, MaxOp, true>>>&);
+#endif
+#ifdef CK_ENABLE_BF16
+// BF16
+void add_device_pool2d_fwd_nhwc_bf16_instances(
+    std::vector<std::unique_ptr<
+        DevicePoolFwd<InOutRank, WindowRank, BF16, BF16, I32, NHWC, NHWC, MaxOp, false>>>&);
+
+void add_device_pool2d_fwd_nhwc_bf16_instances(
+    std::vector<std::unique_ptr<
+        DevicePoolFwd<InOutRank, WindowRank, BF16, BF16, I32, NHWC, NHWC, AvgOp, false>>>&);
+
+// BF16 - return index
+void add_device_pool2d_fwd_nhwc_index_bf16_instances(
+    std::vector<std::unique_ptr<
+        DevicePoolFwd<InOutRank, WindowRank, BF16, BF16, I32, NHWC, NHWC, MaxOp, true>>>&);
+#endif
+#ifdef CK_ENABLE_FP32
+// FP32
+void add_device_pool2d_fwd_nhwc_f32_instances(
+    std::vector<std::unique_ptr<
+        DevicePoolFwd<InOutRank, WindowRank, F32, F32, I32, NHWC, NHWC, MaxOp, false>>>&);
+
+void add_device_pool2d_fwd_nhwc_f32_instances(
+    std::vector<std::unique_ptr<
+        DevicePoolFwd<InOutRank, WindowRank, F32, F32, I32, NHWC, NHWC, AvgOp, false>>>&);
+
+// FP32 - return index
+void add_device_pool2d_fwd_nhwc_index_f32_instances(
+    std::vector<std::unique_ptr<
+        DevicePoolFwd<InOutRank, WindowRank, F32, F32, I32, NHWC, NHWC, MaxOp, true>>>&);
+#endif
+template <typename InDataType,
+          typename OutDataType,
+          typename IndexDataType,
+          typename InLayout,
+          typename OutLayout,
+          ck::ReduceTensorOp ReduceOpId,
+          bool OutputIndex>
+struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DevicePoolFwd<InOutRank,
+                                                                                  WindowRank,
+                                                                                  InDataType,
+                                                                                  OutDataType,
+                                                                                  IndexDataType,
+                                                                                  InLayout,
+                                                                                  OutLayout,
+                                                                                  ReduceOpId,
+                                                                                  OutputIndex>>
+{
+    using DeviceOp = DevicePoolFwd<InOutRank,
+                                   WindowRank,
+                                   InDataType,
+                                   OutDataType,
+                                   IndexDataType,
+                                   InLayout,
+                                   OutLayout,
+                                   ReduceOpId,
+                                   OutputIndex>;
+
+    static auto GetInstances()
+    {
+        std::vector<std::unique_ptr<DeviceOp>> op_ptrs;
+        if constexpr(is_same_v<InLayout, NHWC> && is_same_v<OutLayout, NHWC>)
+        {
+#ifdef CK_ENABLE_FP16
+            if constexpr(is_same_v<InDataType, F16> && is_same_v<OutDataType, F16> &&
+                         is_same_v<IndexDataType, I32>)
+            {
+                if constexpr(OutputIndex && ReduceOpId == MaxOp)
+                {
+                    add_device_pool2d_fwd_nhwc_index_f16_instances(op_ptrs);
+                }
+                else
+                {
+                    add_device_pool2d_fwd_nhwc_f16_instances(op_ptrs);
+                }
+            }
+#endif
+#ifdef CK_ENABLE_BF16
+            else if constexpr(is_same_v<InDataType, BF16> && is_same_v<OutDataType, BF16> &&
+                              is_same_v<IndexDataType, I32>)
+            {
+                if constexpr(OutputIndex && ReduceOpId == MaxOp)
+                {
+                    add_device_pool2d_fwd_nhwc_index_bf16_instances(op_ptrs);
+                }
+                else
+                {
+                    add_device_pool2d_fwd_nhwc_bf16_instances(op_ptrs);
+                }
+            }
+#endif
+#ifdef CK_ENABLE_FP32
+            else if constexpr(is_same_v<InDataType, F32> && is_same_v<OutDataType, F32> &&
+                              is_same_v<IndexDataType, I32>)
+            {
+                if constexpr(OutputIndex && ReduceOpId == MaxOp)
+                {
+                    add_device_pool2d_fwd_nhwc_index_f32_instances(op_ptrs);
+                }
+                else
+                {
+                    add_device_pool2d_fwd_nhwc_f32_instances(op_ptrs);
+                }
+            }
+#endif
+        }
+
+        return op_ptrs;
+    }
+};
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/pool2d_fwd/CMakeLists.txt

+set(DEVICE_POOL2D_FWD_INSTANCES)
+list(APPEND DEVICE_POOL2D_FWD_INSTANCES device_avg_pool2d_fwd_nhwc_f16_instance.cpp
+                                        device_max_pool2d_fwd_nhwc_f16_instance.cpp
+                                        device_avg_pool2d_fwd_nhwc_f32_instance.cpp
+                                        device_max_pool2d_fwd_nhwc_f32_instance.cpp
+                                        device_avg_pool2d_fwd_nhwc_bf16_instance.cpp
+                                        device_max_pool2d_fwd_nhwc_bf16_instance.cpp)
+add_instance_library(device_pool2d_fwd_instance ${DEVICE_POOL2D_FWD_INSTANCES})

library/src/tensor_operation_instance/gpu/pool2d_fwd/device_avg_pool2d_fwd_nhwc_bf16_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "pool2d_fwd_instance_common.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+static constexpr auto ReduceOpId = ck::ReduceTensorOp::AVG;
+
+void add_device_pool2d_fwd_nhwc_bf16_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<4, 2, BF16, BF16, I32, NHWC, NHWC, ReduceOpId, false>>>&
+        instances)
+{
+    add_device_operation_instances(
+        instances, device_pool2d_fwd_nhwc_instances<BF16, BF16, I32, F32, ReduceOpId, false>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/pool2d_fwd/device_avg_pool2d_fwd_nhwc_f16_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "pool2d_fwd_instance_common.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+static constexpr auto ReduceOpId = ck::ReduceTensorOp::AVG;
+
+void add_device_pool2d_fwd_nhwc_f16_instances(
+    std::vector<std::unique_ptr<DevicePoolFwd<4, 2, F16, F16, I32, NHWC, NHWC, ReduceOpId, false>>>&
+        instances)
+{
+    add_device_operation_instances(
+        instances, device_pool2d_fwd_nhwc_instances<F16, F16, I32, F32, ReduceOpId, false>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/pool2d_fwd/device_avg_pool2d_fwd_nhwc_f32_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "pool2d_fwd_instance_common.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+static constexpr auto ReduceOpId = ck::ReduceTensorOp::AVG;
+
+void add_device_pool2d_fwd_nhwc_f32_instances(
+    std::vector<std::unique_ptr<DevicePoolFwd<4, 2, F32, F32, I32, NHWC, NHWC, ReduceOpId, false>>>&
+        instances)
+{
+    add_device_operation_instances(
+        instances, device_pool2d_fwd_nhwc_instances<F32, F32, I32, F32, ReduceOpId, false>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/pool2d_fwd/device_max_pool2d_fwd_nhwc_bf16_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "pool2d_fwd_instance_common.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+static constexpr auto ReduceOpId = ck::ReduceTensorOp::MAX;
+
+void add_device_pool2d_fwd_nhwc_bf16_instances(
+    std::vector<
+        std::unique_ptr<DevicePoolFwd<4, 2, BF16, BF16, I32, NHWC, NHWC, ReduceOpId, false>>>&
+        instances)
+{
+    add_device_operation_instances(
+        instances, device_pool2d_fwd_nhwc_instances<BF16, BF16, I32, F32, ReduceOpId, false>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/pool2d_fwd/device_max_pool2d_fwd_nhwc_f16_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "pool2d_fwd_instance_common.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+static constexpr auto ReduceOpId = ck::ReduceTensorOp::MAX;
+
+void add_device_pool2d_fwd_nhwc_f16_instances(
+    std::vector<std::unique_ptr<DevicePoolFwd<4, 2, F16, F16, I32, NHWC, NHWC, ReduceOpId, false>>>&
+        instances)
+{
+    add_device_operation_instances(
+        instances, device_pool2d_fwd_nhwc_instances<F16, F16, I32, F32, ReduceOpId, false>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/pool2d_fwd/device_max_pool2d_fwd_nhwc_f32_instance.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "pool2d_fwd_instance_common.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+static constexpr auto ReduceOpId = ck::ReduceTensorOp::MAX;
+
+void add_device_pool2d_fwd_nhwc_f32_instances(
+    std::vector<std::unique_ptr<DevicePoolFwd<4, 2, F32, F32, I32, NHWC, NHWC, ReduceOpId, false>>>&
+        instances)
+{
+    add_device_operation_instances(
+        instances, device_pool2d_fwd_nhwc_instances<F32, F32, I32, F32, ReduceOpId, false>{});
+}
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck

library/src/tensor_operation_instance/gpu/pool2d_fwd/pool2d_fwd_instance_common.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck/ck.hpp"
+#include "ck/tensor_operation/gpu/device/impl/device_pool2d_fwd_nhwc_nhwc.hpp"
+#include "ck/utility/data_type.hpp"
+
+#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace device {
+namespace instance {
+
+using I32  = int32_t;
+using F16  = ck::half_t;
+using BF16 = ck::bhalf_t;
+using F32  = float;
+using NHWC = ck::tensor_layout::convolution::NHWC;
+
+template <typename InDataType,
+          typename OutDataType,
+          typename IndexDataType,
+          typename ComputeDataType,
+          ReduceTensorOp ReduceOpId,
+          bool OutputIndex>
+using device_pool2d_fwd_nhwc_instances =
+    // clang-format off
+    std::tuple <
+        DevicePool2dFwd_NHWC_NHWC<InDataType, OutDataType, IndexDataType, ComputeDataType, ReduceOpId, OutputIndex, 256, 256, 1, 1, 1, 1>,
+        DevicePool2dFwd_NHWC_NHWC<InDataType, OutDataType, IndexDataType, ComputeDataType, ReduceOpId, OutputIndex, 256, 256, 1, 2, 1, 2>,
+        DevicePool2dFwd_NHWC_NHWC<InDataType, OutDataType, IndexDataType, ComputeDataType, ReduceOpId, OutputIndex, 256, 256, 1, 4, 1, 4>
+               // clang-format on
+               >;
+
+} // namespace instance
+} // namespace device
+} // namespace tensor_operation
+} // namespace ck