Merge remote-tracking branch 'origin/develop' into stream-k-initial-impl

include/ck/tensor_operation/gpu/device/device_pool2d_fwd.hpp → include/ck/tensor_operation/gpu/device/device_pool_fwd.hpp

@@ -3,8 +3,7 @@
 
 #pragma once
 
-#include <iostream>
-#include <array>
+#include <vector>
 
 #include "ck/tensor_operation/gpu/device/device_base.hpp"
 #include "ck/utility/reduction_enums.hpp"
@@ -13,28 +12,33 @@ namespace ck {
 namespace tensor_operation {
 namespace device {
 
-template <ck::ReduceTensorOp ReduceOpId>
-struct DevicePool2dFwd : public BaseOperator
+template <index_t InOutRank,
+          index_t WindowRank,
+          typename InDataType,
+          typename OutDataType,
+          typename IndexDataType,
+          ReduceTensorOp ReduceOpId,
+          bool OutputIndex>
+struct DevicePoolFwd : public BaseOperator
 {
     virtual std::unique_ptr<BaseArgument>
-    MakeArgumentPointer(const void* in_dev,
-                        void* out_dev,
-                        void* out_indices_dev,
-                        ck::index_t N,
-                        ck::index_t C,
-                        std::array<ck::index_t, 2> input_spatial_lengths,
-                        std::array<ck::index_t, 2> window_spatial_lengths,
-                        std::array<ck::index_t, 2> output_spatial_lengths,
-                        std::array<ck::index_t, 2> window_strides,
-                        std::array<ck::index_t, 2> input_left_pads,
-                        std::array<ck::index_t, 2> input_right_pads) = 0;
+    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> input_left_pads,
+                        std::vector<ck::index_t> input_right_pads,
+                        std::vector<ck::index_t> pooling_dims) = 0;
 
     virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
 };
 
-template <ck::ReduceTensorOp ReduceOpId>
-using DevicePool2dFwdPtr = std::unique_ptr<DevicePool2dFwd<ReduceOpId>>;
-
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck