// Copyright (c) 2021 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

// CUDA and HIP use ReduceGpuKernel API, XPU use ReduceKernel API.
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
    defined(PADDLE_WITH_XPU_KP)

#include "paddle/phi/core/visit_type.h"
#include "paddle/phi/kernels/funcs/reduce_function.h"
#include "paddle/phi/kernels/funcs/reduce_gpu_kernel.h"

namespace phi {

template <typename T,
          template <typename>
          class ReduceOp,
          template <typename, typename>
          class TransformOp,
          bool IsMean = false>
void Reduce(const KPDevice& dev_ctx,
            const DenseTensor& x,
            bool reduce_all,
            const std::vector<int64_t>& dims,
            bool keep_dim,  // unused
            DataType out_dtype,
            DenseTensor* out) {
  reduce_all = recompute_reduce_all(x, dims, reduce_all);
  std::vector<int> reduce_dims =
      funcs::details::GetReduceDim(dims, x.dims().size(), reduce_all);

  int64_t reduce_num = 1;
  for (auto i : reduce_dims) {
    reduce_num *= (x.dims())[i];
  }
#ifdef PADDLE_WITH_XPU_KP
  using MPType = typename MPTypeTrait<T>::Type;
  funcs::ReduceKernel<T, T, ReduceOp, TransformOp<T, MPType>, IsMean>(
      dev_ctx, x, out, TransformOp<T, MPType>(reduce_num), reduce_dims);
#else
  if (out_dtype != DataType::UNDEFINED && out_dtype != x.dtype()) {
    auto tmp_tensor = Cast<T>(dev_ctx, x, out_dtype);
    PD_VISIT_BOOL_AND_FLOATING_AND_COMPLEX_AND_4_TYPES(
        DataType::INT32,
        DataType::INT64,
        DataType::FLOAT16,
        DataType::BFLOAT16,
        out_dtype,
        "ReduceKernel",
        ([&] {
          using MPType = typename MPTypeTrait<data_t>::Type;
          funcs::ReduceKernel<data_t,
                              data_t,
                              ReduceOp,
                              TransformOp<data_t, MPType>,
                              IsMean>(dev_ctx,
                                      tmp_tensor,
                                      out,
                                      TransformOp<data_t, MPType>(reduce_num),
                                      reduce_dims);
        }));
  } else {
    using MPType = typename MPTypeTrait<T>::Type;
    funcs::ReduceKernel<T, T, ReduceOp, TransformOp<T, MPType>, IsMean>(
        dev_ctx, x, out, TransformOp<T, MPType>(reduce_num), reduce_dims);
  }
#endif
}

#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
template <typename T, template <typename, typename, typename> class ReduceOp>
void Reduce(const KPDevice& dev_ctx,
            const DenseTensor& x,
            bool reduce_all,
            const std::vector<int64_t>& dims,
            DataType out_dtype,
            DenseTensor* out) {
  reduce_all = recompute_reduce_all(x, dims, reduce_all);
  std::vector<int> reduce_dims =
      funcs::details::GetReduceDim(dims, x.dims().size(), reduce_all);

  int64_t reduce_num = 1;
  for (auto i : reduce_dims) {
    reduce_num *= (x.dims())[i];
  }

  if (out_dtype != DataType::UNDEFINED && out_dtype != x.dtype()) {
    if (x.dtype() == DataType::BFLOAT16 && out_dtype == DataType::FLOAT32) {
      phi::funcs::ReduceGpuKernel<phi::bfloat16, float, ReduceOp>(
          dev_ctx, x, out, reduce_dims);
    } else if (x.dtype() == DataType::FLOAT16 &&
               out_dtype == DataType::FLOAT32) {
      phi::funcs::ReduceGpuKernel<phi::float16, float, ReduceOp>(
          dev_ctx, x, out, reduce_dims);
    } else {
      auto tmp_tensor = Cast<T>(dev_ctx, x, out_dtype);
      tmp_tensor.set_strides(x.strides());

      PD_VISIT_BOOL_AND_FLOATING_AND_COMPLEX_AND_4_TYPES(
          DataType::INT32,
          DataType::INT64,
          DataType::FLOAT16,
          DataType::BFLOAT16,
          out_dtype,
          "ReduceGpuKernel",
          ([&] {
            using MPType = typename MPTypeTrait<data_t>::Type;
            phi::funcs::ReduceGpuKernel<data_t, data_t, ReduceOp>(
                dev_ctx, tmp_tensor, out, reduce_dims);
          }));
    }
  } else {
    using MPType = typename MPTypeTrait<T>::Type;
    phi::funcs::ReduceGpuKernel<T, T, ReduceOp>(dev_ctx, x, out, reduce_dims);
  }
}
#endif

}  // namespace phi

#endif