chore: import upstream snapshot with attribution

paddle/phi/kernels/impl/pool_kernel_impl.h

@@ -0,0 +1,528 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+
+#include <algorithm>
+
+#include "paddle/common/ddim.h"
+#include "paddle/phi/kernels/full_kernel.h"
+#include "paddle/phi/kernels/funcs/pooling.h"
+#include "paddle/phi/kernels/pool_kernel.h"
+
+#if defined(__HIPCC__) || defined(__NVCC__)
+#include "paddle/phi/kernels/gpu/reduce.h"
+#endif
+
+namespace phi {
+
+inline int64_t GetReduceNum(const DenseTensor& input,
+                            const DenseTensor* output,
+                            const bool channel_last,
+                            std::vector<int>* reduce_dim) {
+  int64_t reduce_num = 0;
+  const int output_height =
+      channel_last ? output->dims()[1] : output->dims()[2];
+  const int output_width = channel_last ? output->dims()[2] : output->dims()[3];
+  if ((output_height == 1) && (output_width == 1)) {
+    if (channel_last) {
+      reduce_dim->push_back(1);
+      reduce_dim->push_back(2);
+      reduce_num = input.dims()[1] * input.dims()[2];
+    } else {
+      reduce_dim->push_back(2);
+      reduce_dim->push_back(3);
+      reduce_num = input.dims()[2] * input.dims()[3];
+    }
+  }
+  return reduce_num;
+}
+
+template <typename T, typename Context>
+void PoolRawKernel(const Context& dev_ctx,
+                   const DenseTensor& x,
+                   const std::vector<int64_t>& kernel_size,
+                   const std::vector<int64_t>& strides,
+                   const std::vector<int64_t>& paddings,
+                   bool exclusive,
+                   const std::string& data_format,
+                   const std::string& pooling_type,
+                   bool global_pooling,
+                   bool adaptive,
+                   const std::string& padding_algorithm,
+                   const float norm_type,
+                   DenseTensor* out) {
+  if (x.numel() == 0) {
+    Full<T, Context>(dev_ctx, out->dims(), NAN, out);
+    return;
+  }
+  const bool channel_last = (data_format == "NHWC" || data_format == "NDHWC");
+  std::vector<int64_t> paddings_ = paddings;
+  std::vector<int64_t> kernel_size_ = kernel_size;
+
+  // update paddings
+  auto x_dims = x.dims();
+  DDim data_dims;
+  if (channel_last) {
+    data_dims = slice_ddim(x_dims, 1, x_dims.size() - 1);
+  } else {
+    data_dims = slice_ddim(x_dims, 2, x_dims.size());
+  }
+
+  std::string true_type;
+  if (norm_type == INFINITY)
+    true_type = "max";
+  else
+    true_type = pooling_type;
+  if (true_type == "lp" && norm_type == 0)
+    PADDLE_THROW(
+        errors::InvalidArgument("norm_type of LPPool op cannot be 0."));
+
+  funcs::UpdatePadding(&paddings_,
+                       global_pooling,
+                       adaptive,
+                       padding_algorithm,
+                       data_dims,
+                       strides,
+                       kernel_size_);
+
+  if (data_dims.size() * 2 == static_cast<int>(paddings_.size())) {
+    for (int i = 0; i < data_dims.size(); ++i) {
+      paddings_.erase(paddings_.begin() + i + 1);
+    }
+  }
+
+  if (global_pooling) {
+    funcs::UpdateKernelSize(&kernel_size_, data_dims);
+  }
+
+  switch (kernel_size_.size()) {
+    case 2: {
+      if (true_type == "max") {
+        funcs::Pool2dFunctor<Context, funcs::MaxPool<T>, T> pool2d_forward;
+        funcs::MaxPool<T> pool_process;
+        pool2d_forward(dev_ctx,
+                       x,
+                       kernel_size_,
+                       strides,
+                       paddings_,
+                       data_format,
+                       true,
+                       false,
+                       out,
+                       pool_process);
+
+      } else if (true_type == "avg") {
+        std::vector<int> reduce_dim;
+        int64_t reduce_num = GetReduceNum(x, out, channel_last, &reduce_dim);
+        if (reduce_num > 0 &&
+            adaptive) {  // for adaptive_avg_pool2d && output_size == 1
+#if defined(__HIPCC__) || defined(__NVCC__)
+          auto stream = dev_ctx.stream();
+          funcs::ReduceGpuKernel<T, T, kps::MeanOps>(
+              dev_ctx, x, out, reduce_dim);
+#else  // for cpu
+          funcs::Pool2dFunctor<Context, funcs::AvgPool<T>, T> pool2d_forward;
+          funcs::AvgPool<T> pool_process;
+          pool2d_forward(dev_ctx,
+                         x,
+                         kernel_size_,
+                         strides,
+                         paddings_,
+                         data_format,
+                         exclusive,
+                         adaptive,
+                         out,
+                         pool_process);
+#endif
+        } else {  // avgpool_2d or  adaptive_avg_pool2d && output_size != 1
+          funcs::Pool2dFunctor<Context, funcs::AvgPool<T>, T> pool2d_forward;
+          funcs::AvgPool<T> pool_process;
+          pool2d_forward(dev_ctx,
+                         x,
+                         kernel_size_,
+                         strides,
+                         paddings_,
+                         data_format,
+                         exclusive,
+                         adaptive,
+                         out,
+                         pool_process);
+        }
+      } else {  // lp_pool2d
+        funcs::Pool2dFunctor<Context, funcs::LPPool<T>, T> pool2d_forward;
+        funcs::LPPool<T> pool_process;
+        pool_process.setNormType(norm_type);
+        pool2d_forward(dev_ctx,
+                       x,
+                       kernel_size_,
+                       strides,
+                       paddings_,
+                       data_format,
+                       exclusive,
+                       adaptive,
+                       out,
+                       pool_process);
+      }
+    } break;
+    case 3: {
+      if (true_type == "max") {
+        funcs::Pool3dFunctor<Context, funcs::MaxPool<T>, T> pool3d_forward;
+        funcs::MaxPool<T> pool_process;
+        pool3d_forward(dev_ctx,
+                       x,
+                       kernel_size_,
+                       strides,
+                       paddings_,
+                       data_format,
+                       true,
+                       false,
+                       out,
+                       pool_process);
+      } else if (true_type == "avg") {
+        funcs::Pool3dFunctor<Context, funcs::AvgPool<T>, T> pool3d_forward;
+        funcs::AvgPool<T> pool_process;
+        pool3d_forward(dev_ctx,
+                       x,
+                       kernel_size_,
+                       strides,
+                       paddings_,
+                       data_format,
+                       exclusive,
+                       adaptive,
+                       out,
+                       pool_process);
+      } else {  // lp_pool3d
+        PADDLE_THROW(
+            errors::InvalidArgument("LPPool op only supports 2D input."));
+      }
+    } break;
+    default: {
+      PADDLE_THROW(
+          errors::InvalidArgument("Pool op only supports 2D and 3D input."));
+    }
+  }
+}
+
+template <typename Context, typename T1, typename T2 = int>
+void MaxPoolWithIndexRawKernel(const Context& dev_ctx,
+                               const DenseTensor& x,
+                               const std::vector<int>& kernel_size,
+                               const std::vector<int>& strides,
+                               const std::vector<int>& paddings,
+                               const std::vector<int>& dilations,
+                               bool global_pooling,
+                               bool adaptive,
+                               DenseTensor* out,
+                               DenseTensor* mask) {
+  if (x.numel() == 0) {
+    if (out) {
+      Full<T1, Context>(dev_ctx, out->dims(), NAN, out);
+    }
+    if (mask) {
+      Full<T2, Context>(dev_ctx, mask->dims(), 0, mask);
+    }
+    return;
+  }
+
+  std::vector<int64_t> paddings_(paddings.begin(), paddings.end());
+  std::vector<int64_t> kernel_size_(kernel_size.begin(), kernel_size.end());
+  std::vector<int64_t> strides_(strides.begin(), strides.end());
+  std::vector<int64_t> dilations_(dilations.begin(), dilations.end());
+
+  if (global_pooling) {
+    for (size_t i = 0; i < kernel_size_.size(); ++i) {
+      paddings_[i] = 0;
+      kernel_size_[i] = static_cast<int>(x.dims()[i + 2]);
+      dilations_[i] = 1;  // Reset dilation for global pooling
+    }
+  }
+
+  switch (kernel_size_.size()) {
+    case 2: {
+      funcs::MaxPool2dWithIndexFunctor<Context, T1, T2> pool2d_forward;
+      pool2d_forward(dev_ctx,
+                     x,
+                     kernel_size_,
+                     strides_,
+                     paddings_,
+                     dilations_,
+                     adaptive,
+                     out,
+                     mask);
+    } break;
+    case 3: {
+      funcs::MaxPool3dWithIndexFunctor<Context, T1, T2> pool3d_forward;
+      pool3d_forward(dev_ctx,
+                     x,
+                     kernel_size_,
+                     strides_,
+                     paddings_,
+                     dilations_,
+                     adaptive,
+                     out,
+                     mask);
+    } break;
+    default: {
+      PADDLE_THROW(
+          errors::InvalidArgument("Pool op only supports 2D and 3D input."));
+    }
+  }
+}
+
+template <typename T, typename Context>
+void Pool2dKernel(const Context& dev_ctx,
+                  const DenseTensor& x,
+                  const IntArray& kernel_size,
+                  const std::vector<int64_t>& strides,
+                  const std::vector<int64_t>& paddings,
+                  bool ceil_mode UNUSED,
+                  bool exclusive,
+                  const std::string& data_format,
+                  const std::string& pooling_type,
+                  bool global_pooling,
+                  bool adaptive,
+                  const std::string& padding_algorithm,
+                  DenseTensor* out) {
+  if (x.numel() == 0) {
+    if (pooling_type == "max") {
+      Full<T, Context>(dev_ctx, out->dims(), 0, out);
+    } else {  // for pooling_type == "avg"
+      Full<T, Context>(dev_ctx, out->dims(), NAN, out);
+    }
+    return;
+  }
+  PoolRawKernel<T, Context>(dev_ctx,
+                            x,
+                            kernel_size.GetData(),
+                            strides,
+                            paddings,
+                            exclusive,
+                            data_format,
+                            pooling_type,
+                            global_pooling,
+                            adaptive,
+                            padding_algorithm,
+                            0,
+                            out);
+}
+
+template <typename T, typename Context>
+void LPPool2dKernel(const Context& dev_ctx,
+                    const DenseTensor& x,
+                    const IntArray& kernel_size,
+                    const std::vector<int64_t>& strides,
+                    const std::vector<int64_t>& paddings,
+                    bool ceil_mode UNUSED,
+                    bool exclusive,
+                    const std::string& data_format,
+                    const std::string& pooling_type,
+                    bool global_pooling,
+                    bool adaptive,
+                    const std::string& padding_algorithm,
+                    const float norm_type,
+                    DenseTensor* out) {
+  int dims = x.dims().size();
+  if (x.numel() == 0 && dims) {
+    bool need_zero = false;
+    for (int i = 1; i < dims; i++) {
+      if (x.dims()[i] == 0) {
+        need_zero = true;
+        break;
+      }
+    }
+    if (need_zero) {
+      Full<T, Context>(dev_ctx, out->dims(), 0, out);
+      return;
+    }
+  }
+  PoolRawKernel<T, Context>(dev_ctx,
+                            x,
+                            kernel_size.GetData(),
+                            strides,
+                            paddings,
+                            exclusive,
+                            data_format,
+                            pooling_type,
+                            global_pooling,
+                            adaptive,
+                            padding_algorithm,
+                            norm_type,
+                            out);
+}
+
+template <typename T, typename Context>
+void MaxPool2dWithIndexKernel(const Context& dev_ctx,
+                              const DenseTensor& x,
+                              const std::vector<int>& kernel_size,
+                              const std::vector<int>& strides,
+                              const std::vector<int>& paddings,
+                              const std::vector<int>& dilations,
+                              bool global_pooling,
+                              bool adaptive,
+                              bool ceil_mode UNUSED,
+                              DenseTensor* out,
+                              DenseTensor* mask) {
+  MaxPoolWithIndexRawKernel<Context, T>(dev_ctx,
+                                        x,
+                                        kernel_size,
+                                        strides,
+                                        paddings,
+                                        dilations,
+                                        global_pooling,
+                                        adaptive,
+                                        out,
+                                        mask);
+}
+
+template <typename T, typename Context>
+void Pool3dKernel(const Context& dev_ctx,
+                  const DenseTensor& x,
+                  const std::vector<int64_t>& kernel_size,
+                  const std::vector<int64_t>& strides,
+                  const std::vector<int64_t>& paddings,
+                  bool ceil_mode UNUSED,
+                  bool exclusive,
+                  const std::string& data_format,
+                  const std::string& pooling_type,
+                  bool global_pooling,
+                  bool adaptive,
+                  const std::string& padding_algorithm,
+                  DenseTensor* out) {
+  if (x.numel() == 0) {
+    if (pooling_type == "max" || (!adaptive && pooling_type == "avg")) {
+      Full<T, Context>(dev_ctx, out->dims(), 0, out);
+    } else {
+      Full<T, Context>(dev_ctx, out->dims(), NAN, out);
+    }
+    return;
+  }
+  PoolRawKernel<T, Context>(dev_ctx,
+                            x,
+                            kernel_size,
+                            strides,
+                            paddings,
+                            exclusive,
+                            data_format,
+                            pooling_type,
+                            global_pooling,
+                            adaptive,
+                            padding_algorithm,
+                            0,
+                            out);
+}
+
+template <typename T, typename Context>
+void MaxPool3dWithIndexKernel(const Context& dev_ctx,
+                              const DenseTensor& x,
+                              const std::vector<int>& kernel_size,
+                              const std::vector<int>& strides,
+                              const std::vector<int>& paddings,
+                              const std::vector<int>& dilations,
+                              bool global_pooling,
+                              bool adaptive,
+                              bool ceil_mode UNUSED,
+                              DenseTensor* out,
+                              DenseTensor* mask) {
+  MaxPoolWithIndexRawKernel<Context, T>(dev_ctx,
+                                        x,
+                                        kernel_size,
+                                        strides,
+                                        paddings,
+                                        dilations,
+                                        global_pooling,
+                                        adaptive,
+                                        out,
+                                        mask);
+}
+
+template <typename Context, typename T1, typename T2 = int>
+void FractionalMaxPoolRawKernel(const Context& dev_ctx,
+                                const DenseTensor& x,
+                                const std::vector<int>& output_size,
+                                const std::vector<int>& kernel_size,
+                                float random_u,
+                                bool return_mask,
+                                DenseTensor* out,
+                                DenseTensor* mask) {
+  if (x.numel() == 0) {
+    if (out) {
+      Full<T1, Context>(dev_ctx, out->dims(), NAN, out);
+    }
+    if (mask) {
+      Full<T2, Context>(dev_ctx, mask->dims(), 0, mask);
+    }
+    return;
+  }
+
+  std::vector<int64_t> output_size_(output_size.begin(), output_size.end());
+  std::vector<int64_t> kernel_size_(kernel_size.begin(), kernel_size.end());
+
+  switch (output_size_.size()) {
+    case 2: {
+      funcs::FractionalMaxPool2dFunctor<Context, T1, T2> pool2d_forward;
+      pool2d_forward(dev_ctx,
+                     x,
+                     output_size_,
+                     kernel_size_,
+                     random_u,
+                     return_mask,
+                     out,
+                     mask);
+    } break;
+    case 3: {
+      funcs::FractionalMaxPool3dFunctor<Context, T1, T2> pool3d_forward;
+      pool3d_forward(dev_ctx,
+                     x,
+                     output_size_,
+                     kernel_size_,
+                     random_u,
+                     return_mask,
+                     out,
+                     mask);
+    } break;
+    default: {
+      PADDLE_THROW(
+          errors::InvalidArgument("Pool op only supports 2D and 3D input."));
+    }
+  }
+}
+
+template <typename T, typename Context>
+void FractionalMaxPool2dKernel(const Context& dev_ctx,
+                               const DenseTensor& x,
+                               const std::vector<int>& output_size,
+                               const std::vector<int>& kernel_size,
+                               float random_u,
+                               bool return_mask,
+                               DenseTensor* out,
+                               DenseTensor* mask) {
+  FractionalMaxPoolRawKernel<Context, T>(
+      dev_ctx, x, output_size, kernel_size, random_u, return_mask, out, mask);
+}
+
+template <typename T, typename Context>
+void FractionalMaxPool3dKernel(const Context& dev_ctx,
+                               const DenseTensor& x,
+                               const std::vector<int>& output_size,
+                               const std::vector<int>& kernel_size,
+                               float random_u,
+                               bool return_mask,
+                               DenseTensor* out,
+                               DenseTensor* mask) {
+  FractionalMaxPoolRawKernel<Context, T>(
+      dev_ctx, x, output_size, kernel_size, random_u, return_mask, out, mask);
+}
+
+}  // namespace phi