Pool2d max/avg kernel in the BWD version (#1494)

* Add pool2d instance BWD AVG

* Add pool2d instance BWD MAX

* Fix: avg review

* Fix review: part2

* Fix - enable test when type is compiled

* Fix review part3

profiler/src/profile_max_pool2d_bwd.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <iostream>
+#include <vector>
+#include <unordered_map>
+
+#include "profiler/data_type_enum.hpp"
+#include "profiler/profile_max_pool2d_bwd_impl.hpp"
+#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
+#include "profiler_operation_registry.hpp"
+
+using ck::index_t;
+
+struct maxPoolbwdArgParser
+{
+    std::unordered_map<std::string, std::vector<int>> long_opts = {{"length", {}},
+                                                                   {"wsize", {}},
+                                                                   {"wstride", {}},
+                                                                   {"wdilation", {}},
+                                                                   {"pad1", {}},
+                                                                   {"pad2", {}}};
+
+    bool parse_opt(int argc, char* argv[], const std::string& key, int i)
+    {
+        if(std::string("--") + key == argv[i])
+        {
+            int pos = i;
+            while(++i < argc && argv[i][0] != '-') {}
+            int end = i;
+            for(int j = pos + 1; j < end; j++)
+            {
+                long_opts[key].push_back(std::stoi(argv[j]));
+            }
+            return true;
+        }
+        return false;
+    }
+
+    void operator()(int argc, char* argv[])
+    {
+        for(auto& kv : long_opts)
+        {
+            for(int i = 1; i < argc; i++)
+            {
+                if(parse_opt(argc, argv, kv.first, i))
+                    break;
+            }
+        }
+    }
+};
+
+void print_help_max_pool2d_bwd()
+{
+    std::cout << "arg1: data type (0: fp16; 1: fp32; 3: int8; 5: bf16)\n"
+              << "arg2: verification (0: no; 1: yes)\n"
+              << "arg3: initialization (0: no init; 1: integer value; 2: decimal value)\n"
+              << "arg4: print tensor value (0: no; 1: yes)\n"
+              << "arg5: time kernel (0=no, 1=yes)\n"
+              << "--length: input tensor length for NCHW(e.g, --length 2 32 30 30) \n"
+              << "--wsize: window size for YX (e.g, --wsize 2 2) \n"
+              << "--wstride: window stride for HW (e.g, --wstride 2 2) \n"
+              << "--wdilation: window dilation for HW (e.g, --wdilation 1 1) \n"
+              << "--pad1: left side of padding in HW (e.g, --pad1 1 1) \n"
+              << "--pad2: right side of padding in HW (e.g, --pad2 1 1) \n"
+              << "eg: ckProfiler max_pool2d_bwd 0 1 2 0 --length 2 32 30 30 --wsize 2 2 "
+                 "--wstride 2 2 --wdilation 1 1 --pad1 1 1 --pad2 1 1"
+              << std::endl;
+}
+
+int profile_max_pool2d_bwd(int argc, char* argv[])
+{
+    ck::DataTypeEnum data_type = ck::DataTypeEnum::Half;
+    bool do_verification       = true;
+    int init_method            = 2;
+    bool do_log                = false;
+    bool time_kernel           = true;
+
+    std::vector<index_t> in_length = {2, 32, 30, 30};
+    std::vector<index_t> wsize     = {2, 2};
+    std::vector<index_t> wstride   = {2, 2};
+    std::vector<index_t> wdilation = {1, 1};
+    std::vector<index_t> pad1      = {1, 1};
+    std::vector<index_t> pad2      = {1, 1};
+
+    if(argc != 2 && argc != 33)
+    {
+        print_help_max_pool2d_bwd();
+        return 0;
+    }
+    else if(argc == 33)
+    {
+        data_type       = static_cast<ck::DataTypeEnum>(std::stoi(argv[2]));
+        do_verification = std::stoi(argv[3]);
+        init_method     = std::stoi(argv[4]);
+        do_log          = std::stoi(argv[5]);
+        time_kernel     = std::stoi(argv[6]);
+
+        // parse the long options
+        maxPoolbwdArgParser arg_parser;
+        arg_parser(argc, argv);
+        in_length = arg_parser.long_opts["length"];
+        wsize     = arg_parser.long_opts["wsize"];
+        wstride   = arg_parser.long_opts["wstride"];
+        wdilation = arg_parser.long_opts["wdilation"];
+        pad1      = arg_parser.long_opts["pad1"];
+        pad2      = arg_parser.long_opts["pad2"];
+    }
+
+    using F16  = ck::half_t;
+    using BF16 = ck::bhalf_t;
+    using F32  = float;
+    using I8   = int8_t;
+    using I32  = int32_t;
+
+    if(data_type == ck::DataTypeEnum::Half)
+    {
+        ck::profiler::profile_max_pool2d_bwd_impl<F16, F16, I32, F16, F16, false>(do_verification,
+                                                                                  init_method,
+                                                                                  do_log,
+                                                                                  time_kernel,
+                                                                                  in_length,
+                                                                                  wsize,
+                                                                                  wstride,
+                                                                                  wdilation,
+                                                                                  pad1,
+                                                                                  pad2);
+    }
+    else if(data_type == ck::DataTypeEnum::BFloat16)
+    {
+        ck::profiler::profile_max_pool2d_bwd_impl<BF16, BF16, I32, BF16, BF16, false>(
+            do_verification,
+            init_method,
+            do_log,
+            time_kernel,
+            in_length,
+            wsize,
+            wstride,
+            wdilation,
+            pad1,
+            pad2);
+    }
+    else if(data_type == ck::DataTypeEnum::Float)
+    {
+        ck::profiler::profile_max_pool2d_bwd_impl<F32, F32, I32, F32, F32, false>(do_verification,
+                                                                                  init_method,
+                                                                                  do_log,
+                                                                                  time_kernel,
+                                                                                  in_length,
+                                                                                  wsize,
+                                                                                  wstride,
+                                                                                  wdilation,
+                                                                                  pad1,
+                                                                                  pad2);
+    }
+    else if(data_type == ck::DataTypeEnum::Int8)
+    {
+        ck::profiler::profile_max_pool2d_bwd_impl<I8, I8, I32, I8, I8, false>(do_verification,
+                                                                              init_method,
+                                                                              do_log,
+                                                                              time_kernel,
+                                                                              in_length,
+                                                                              wsize,
+                                                                              wstride,
+                                                                              wdilation,
+                                                                              pad1,
+                                                                              pad2);
+    }
+    else
+    {
+
+        throw std::runtime_error("not implemented yet");
+    }
+
+    return 0;
+}
+
+REGISTER_PROFILER_OPERATION("max_pool2d_bwd", "max_pool2d bwd", profile_max_pool2d_bwd);

test/pool/CMakeLists.txt

@@ -4,10 +4,14 @@ add_gtest_executable(test_avg_pool3d_bwd test_avg_pool3d_bwd.cpp)
 add_gtest_executable(test_max_pool3d_bwd test_max_pool3d_bwd.cpp)
 add_gtest_executable(test_avg_pool3d_fwd test_avg_pool3d_fwd.cpp)
 add_gtest_executable(test_max_pool3d_fwd test_max_pool3d_fwd.cpp)
+add_gtest_executable(test_avg_pool2d_bwd test_avg_pool2d_bwd.cpp)
+add_gtest_executable(test_max_pool2d_bwd test_max_pool2d_bwd.cpp)
 add_gtest_executable(test_avg_pool2d_fwd test_avg_pool2d_fwd.cpp)
 add_gtest_executable(test_max_pool2d_fwd test_max_pool2d_fwd.cpp)
 
 target_link_libraries(test_avg_pool3d_bwd PRIVATE utility device_avg_pool3d_bwd_instance)
+target_link_libraries(test_avg_pool2d_bwd PRIVATE utility device_avg_pool2d_bwd_instance)
+target_link_libraries(test_max_pool2d_bwd PRIVATE utility device_max_pool_bwd_instance)
 target_link_libraries(test_max_pool3d_bwd PRIVATE utility device_max_pool_bwd_instance)
 target_link_libraries(test_avg_pool3d_fwd PRIVATE utility device_pool3d_fwd_instance)
 target_link_libraries(test_max_pool3d_fwd PRIVATE utility device_pool3d_fwd_instance)
@@ -18,5 +22,7 @@ add_dependencies(test_pool test_avg_pool3d_bwd)
 add_dependencies(test_pool test_max_pool3d_bwd)
 add_dependencies(test_pool test_avg_pool3d_fwd)
 add_dependencies(test_pool test_max_pool3d_fwd)
+add_dependencies(test_pool test_avg_pool2d_bwd)
+add_dependencies(test_pool test_max_pool2d_bwd)
 add_dependencies(test_pool test_avg_pool2d_fwd)
 add_dependencies(test_pool test_max_pool2d_fwd)

test/pool/test_avg_pool2d_bwd.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "gtest/gtest.h"
+#include "profiler/profile_avg_pool2d_bwd_impl.hpp"
+#include "test_pool_fwd_common.hpp"
+
+template <typename T>
+class AvgPool2dBWDTest : public ::testing::Test
+{
+    protected:
+    using InDataType  = std::tuple_element_t<0, T>;
+    using OutDataType = std::tuple_element_t<1, T>;
+
+    static std::vector<PoolingParam> params;
+
+    void Run()
+    {
+        for(auto param : this->params)
+        {
+            bool success =
+                ck::profiler::profile_avg_pool2d_bwd_impl<InDataType, OutDataType, NHWC, NHWC>(
+                    true,
+                    2,
+                    false,
+                    false,
+                    param.length_,
+                    param.window_spatial_lengths_,
+                    param.window_strides_,
+                    param.window_dilations_,
+                    param.input_left_pads_,
+                    param.input_right_pads_);
+            EXPECT_TRUE(success);
+        }
+    }
+};
+
+template <typename T>
+std::vector<PoolingParam> AvgPool2dBWDTest<T>::params = {
+    {{1, 1, 1, 1}, {1, 1}, {1, 1}, {1, 1}, {0, 0}, {0, 0}},
+    {{1, 1, 64, 64}, {64, 64}, {1, 1}, {1, 1}, {0, 0}, {0, 0}},
+    {{1, 5, 7, 7}, {2, 2}, {2, 2}, {1, 1}, {2, 2}, {0, 0}},
+    {{1, 1, 8, 8}, {2, 2}, {2, 2}, {1, 1}, {2, 2}, {0, 0}},
+    {{1, 1, 8, 8}, {2, 2}, {1, 1}, {1, 1}, {1, 1}, {0, 0}},
+    {{2, 32, 30, 30}, {2, 2}, {2, 2}, {1, 1}, {1, 1}, {1, 1}},
+    {{1, 2, 30, 30}, {2, 2}, {2, 2}, {1, 1}, {0, 0}, {0, 0}}};
+
+using Avg_Pool_2D_f32_types  = ::testing::Types<std::tuple<F32, F32>>;
+using Avg_Pool_2D_int8_types = ::testing::Types<std::tuple<I8, I8>>;
+using Avg_Pool_2D_f16_types  = ::testing::Types<std::tuple<F16, F16>>;
+using Avg_Pool_2D_bf16_types = ::testing::Types<std::tuple<BF16, BF16>>;
+using Avg_Pool_2D_f8_types   = ::testing::Types<std::tuple<F8, F8>>;
+
+template <typename TType>
+class AvgPool2D_f32 : public AvgPool2dBWDTest<TType>
+{
+    protected:
+    void SetUp() override
+    {
+        if(!CK_ENABLE_FP32)
+        {
+            GTEST_SKIP() << "Skipping AvgPool2D_f32 tests because CK_ENABLE_FP32 is not enabled";
+        }
+    }
+};
+
+template <typename TType>
+class AvgPool2D_int8 : public AvgPool2dBWDTest<TType>
+{
+    protected:
+    void SetUp() override
+    {
+        if(!CK_ENABLE_INT8)
+        {
+            GTEST_SKIP() << "Skipping AvgPool2D_int8 tests because CK_ENABLE_INT8 is not enabled";
+        }
+    }
+};
+
+template <typename TType>
+class AvgPool2D_f16 : public AvgPool2dBWDTest<TType>
+{
+    protected:
+    void SetUp() override
+    {
+        if(!CK_ENABLE_FP16)
+        {
+            GTEST_SKIP() << "Skipping AvgPool2D_f16 because CK_ENABLE_FP16 is not enabled";
+        }
+    }
+};
+
+template <typename TType>
+class AvgPool2D_bf16 : public AvgPool2dBWDTest<TType>
+{
+    protected:
+    void SetUp() override
+    {
+        if(!CK_ENABLE_BF16)
+        {
+            GTEST_SKIP() << "Skipping AvgPool2D_bf16 tests because CK_ENABLE_BF16 is not enabled";
+        }
+    }
+};
+
+template <typename TType>
+class AvgPool2D_f8 : public AvgPool2dBWDTest<TType>
+{
+    protected:
+    void SetUp() override
+    {
+        if(!CK_ENABLE_FP8)
+        {
+            GTEST_SKIP() << "Skipping AvgPool2D_f8 tests because CK_ENABLE_FP8 is not enabled";
+        }
+    }
+};
+
+TYPED_TEST_SUITE(AvgPool2D_f32, Avg_Pool_2D_f32_types);
+TYPED_TEST_SUITE(AvgPool2D_int8, Avg_Pool_2D_int8_types);
+TYPED_TEST_SUITE(AvgPool2D_f16, Avg_Pool_2D_f16_types);
+TYPED_TEST_SUITE(AvgPool2D_bf16, Avg_Pool_2D_bf16_types);
+TYPED_TEST_SUITE(AvgPool2D_f8, Avg_Pool_2D_f8_types);
+
+TYPED_TEST(AvgPool2D_f32, AvgPool2DTest_f32) { this->Run(); }
+
+TYPED_TEST(AvgPool2D_int8, AvgPool2DTest_int8) { this->Run(); }
+
+TYPED_TEST(AvgPool2D_f16, AvgPool2DTest_f16) { this->Run(); }
+
+TYPED_TEST(AvgPool2D_bf16, AvgPool2DTest_bf16) { this->Run(); }
+
+TYPED_TEST(AvgPool2D_f8, AvgPool2DTest_f8) { this->Run(); }

test/pool/test_max_pool2d_bwd.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "gtest/gtest.h"
+#include "profiler/profile_max_pool2d_bwd_impl.hpp"
+#include "test_pool_fwd_common.hpp"
+
+template <typename T>
+class MaxPool2dBWDTest : public ::testing::Test
+{
+    protected:
+    using DOutDataType  = std::tuple_element_t<0, T>;
+    using DInDataType   = std::tuple_element_t<1, T>;
+    using IndexDataType = std::tuple_element_t<2, T>;
+
+    using InDataType  = DInDataType;
+    using OutDataType = DOutDataType;
+
+    static std::vector<PoolingParam> params;
+
+    void Run()
+    {
+        for(auto param : this->params)
+        {
+            bool success =
+                ck::profiler::profile_max_pool2d_bwd_impl<InDataType,
+                                                          OutDataType,
+                                                          IndexDataType,
+                                                          DOutDataType,
+                                                          DInDataType,
+                                                          false>(true,
+                                                                 2,
+                                                                 false,
+                                                                 false,
+                                                                 param.length_,
+                                                                 param.window_spatial_lengths_,
+                                                                 param.window_strides_,
+                                                                 param.window_dilations_,
+                                                                 param.input_left_pads_,
+                                                                 param.input_right_pads_);
+            EXPECT_TRUE(success);
+        }
+    }
+};
+
+template <typename T>
+std::vector<PoolingParam> MaxPool2dBWDTest<T>::params = {
+    {{1, 1, 1, 1}, {1, 1}, {1, 1}, {1, 1}, {0, 0}, {0, 0}},
+    {{2, 16, 64, 64}, {64, 64}, {1, 1}, {1, 1}, {0, 0}, {0, 0}},
+    {{2, 16, 64, 64}, {4, 4}, {4, 4}, {2, 2}, {0, 0}, {0, 0}},
+    {{2, 32, 30, 30}, {2, 2}, {2, 2}, {1, 1}, {1, 1}, {1, 1}},
+    {{2, 2, 30, 30}, {2, 2}, {2, 2}, {1, 1}, {1, 1}, {1, 1}}};
+
+using Max_Pool_2D_f32_types  = ::testing::Types<std::tuple<F32, F32, I32>>;
+using Max_Pool_2D_int8_types = ::testing::Types<std::tuple<I8, I8, I32>>;
+using Max_Pool_2D_f16_types  = ::testing::Types<std::tuple<F16, F16, I32>>;
+using Max_Pool_2D_bf16_types = ::testing::Types<std::tuple<BF16, BF16, I32>>;
+
+template <typename TType>
+class MaxPool2D_f32 : public MaxPool2dBWDTest<TType>
+{
+    protected:
+    void SetUp() override
+    {
+        if(!CK_ENABLE_FP32)
+        {
+            GTEST_SKIP() << "Skipping MaxPool2D_f32 tests because CK_ENABLE_FP32 is not enabled";
+        }
+    }
+};
+
+template <typename TType>
+class MaxPool2D_int8 : public MaxPool2dBWDTest<TType>
+{
+    protected:
+    void SetUp() override
+    {
+        if(!CK_ENABLE_INT8)
+        {
+            GTEST_SKIP() << "Skipping MaxPool2D_int8 tests because CK_ENABLE_INT8 is not enabled";
+        }
+    }
+};
+
+template <typename TType>
+class MaxPool2D_f16 : public MaxPool2dBWDTest<TType>
+{
+    protected:
+    void SetUp() override
+    {
+        if(!CK_ENABLE_FP16)
+        {
+            GTEST_SKIP() << "Skipping MaxPool2D_f16 because CK_ENABLE_FP16 is not enabled";
+        }
+    }
+};
+
+template <typename TType>
+class MaxPool2D_bf16 : public MaxPool2dBWDTest<TType>
+{
+    protected:
+    void SetUp() override
+    {
+        if(!CK_ENABLE_BF16)
+        {
+            GTEST_SKIP() << "Skipping MaxPool2D_bf16 tests because CK_ENABLE_BF16 is not enabled";
+        }
+    }
+};
+
+TYPED_TEST_SUITE(MaxPool2D_f32, Max_Pool_2D_f32_types);
+TYPED_TEST_SUITE(MaxPool2D_int8, Max_Pool_2D_int8_types);
+TYPED_TEST_SUITE(MaxPool2D_f16, Max_Pool_2D_f16_types);
+TYPED_TEST_SUITE(MaxPool2D_bf16, Max_Pool_2D_bf16_types);
+
+TYPED_TEST(MaxPool2D_f32, MaxPool2DTest_f32) { this->Run(); }
+
+TYPED_TEST(MaxPool2D_int8, MaxPool2DTest_int8) { this->Run(); }
+
+TYPED_TEST(MaxPool2D_f16, MaxPool2DTest_f16) { this->Run(); }
+
+TYPED_TEST(MaxPool2D_bf16, MaxPool2DTest_bf16) { this->Run(); }

test/pool/test_pool_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.
 
 #include "gtest/gtest.h"
 #include "ck/ck.hpp"
@@ -8,8 +8,11 @@ using F16  = ck::half_t;
 using BF16 = ck::bhalf_t;
 using F32  = float;
 using I32  = int32_t;
+using I8   = int8_t;
+using F8   = ck::f8_t;
 using ck::index_t;
 using NDHWC = ck::tensor_layout::convolution::NDHWC;
+using NHWC  = ck::tensor_layout::convolution::NHWC;
 
 struct PoolingParam
 {