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

#include "paddle/phi/kernels/sync_batch_norm_kernel.h"
#include "paddle/common/enforce.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/sync_batch_norm_utils.h"

namespace phi {

template <typename T, typename Context>
void SyncBatchNormKernel(const Context& dev_ctx,
                         const DenseTensor& x,
                         const DenseTensor& mean,
                         const DenseTensor& variance,
                         const DenseTensor& scale,
                         const DenseTensor& bias,
                         bool is_test,
                         float momentum,
                         float epsilon_f,
                         const std::string& data_layout_str,
                         bool use_global_stats,
                         bool trainable_statistics,
                         DenseTensor* y,
                         DenseTensor* mean_out,
                         DenseTensor* variance_out,
                         DenseTensor* saved_mean,
                         DenseTensor* saved_variance,
                         DenseTensor* reserve_space) {
  PADDLE_ENFORCE_EQ(use_global_stats,
                    false,
                    common::errors::InvalidArgument(
                        "sync_batch_norm doesn't support "
                        "to set use_global_stats True. Please use batch_norm "
                        "in this case."));

  double epsilon = epsilon_f;
  const bool trainable_stats = trainable_statistics;
  const DataLayout layout = StringToDataLayout(data_layout_str);
  bool test_mode = is_test && (!trainable_statistics);
  const auto& x_dims = x.dims();
  PADDLE_ENFORCE_GE(x_dims.size(),
                    2,
                    common::errors::InvalidArgument(
                        "The Input dim size should be larger than 1."));
  PADDLE_ENFORCE_LE(x_dims.size(),
                    5,
                    common::errors::InvalidArgument(
                        "The Input dim size should be less than 6."));
  int N, C, H, W, D;
  funcs::ExtractNCWHD(x_dims, layout, &N, &C, &H, &W, &D);
  int64_t x_numel = x.numel();
  const int64_t fsize = static_cast<int64_t>(H) * W * D;

  const T* x_d = x.template data<T>();
  const auto* s_d = scale.template data<BatchNormParamType<T>>();
  const auto* b_d = bias.template data<BatchNormParamType<T>>();

  T* y_d = dev_ctx.template Alloc<T>(y);

  const BatchNormParamType<T>* mean_data = nullptr;
  const BatchNormParamType<T>* var_data = nullptr;

  auto stream = dev_ctx.stream();
  const int block = 512;
  int max_threads = dev_ctx.GetMaxPhysicalThreadCount();

  Allocator::AllocationPtr alloc_ptr{nullptr};

  if (test_mode) {
    mean_data = mean.template data<BatchNormParamType<T>>();
    var_data = variance.template data<BatchNormParamType<T>>();
  } else {
    // x, x^2, 1, here 1 is used to calc device num
    // device num also can be got from DeviceContextPool
    const int64_t bytes_64 =
        (static_cast<int64_t>(C) * 2 + 1) * sizeof(BatchNormParamType<T>);
    DenseTensor stats_tensor;
    stats_tensor.Resize({bytes_64});
    dev_ctx.template Alloc<BatchNormParamType<T>>(&stats_tensor);
    auto* stats_data = stats_tensor.data<BatchNormParamType<T>>();
    auto* stats = reinterpret_cast<BatchNormParamType<T>*>(stats_data);
    const int threads = 512;
    int grid = std::min(C, (max_threads + threads - 1) / threads);
    if (layout == DataLayout::NCHW) {
      KeLocalStats<T, threads, DataLayout::NCHW>
          <<<grid, threads, 0, stream>>>(x_d, N, fsize, C, stats);
    } else {
      KeLocalStats<T, threads, DataLayout::NHWC>
          <<<grid, threads, 0, stream>>>(x_d, N, fsize, C, stats);
    }

#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
    auto comm_ctx =
        static_cast<distributed::NCCLCommContext*>(dev_ctx.GetCommContext());
    if (comm_ctx) {
      comm_ctx->AllReduce(&stats_tensor, stats_tensor, ncclSum, stream);
    }
#endif

    auto* est_mean_data =
        dev_ctx.template Alloc<BatchNormParamType<T>>(mean_out);
    auto* est_var_data =
        dev_ctx.template Alloc<BatchNormParamType<T>>(variance_out);

    auto* sv_mean_data =
        dev_ctx.template Alloc<BatchNormParamType<T>>(saved_mean);
    auto* sv_inv_var_data =
        dev_ctx.template Alloc<BatchNormParamType<T>>(saved_variance);

    int64_t reserve_space_size = 0;
    DenseTensor tmp_reserve_space;
    if (reserve_space == nullptr) {
      reserve_space = &tmp_reserve_space;
    }
    reserve_space->Resize({reserve_space_size});
    dev_ctx.template Alloc<T>(reserve_space);

    // Note, Input('Mean')/Input('Variance') share variable with
    // Output('MeanOut')/Output('VarianceOut')
    KeSyncAndMovingStats<T>
        <<<(C + block - 1) / block, block, 0, stream>>>(stats,
                                                        stats + C,
                                                        stats + 2 * C,
                                                        C,
                                                        momentum,
                                                        epsilon,
                                                        sv_mean_data,
                                                        sv_inv_var_data,
                                                        est_mean_data,
                                                        est_var_data);

    mean_data = sv_mean_data;
    var_data = stats + C;
  }

  const int64_t grid2_64 =
      (std::min(x_numel, static_cast<int64_t>(max_threads)) + block - 1) /
      block;
  uint32_t grid2 = static_cast<uint32_t>(grid2_64);
  if (layout == DataLayout::NCHW) {
    KeNormAffine<T, DataLayout::NCHW><<<grid2, block, 0, stream>>>(
        x_d, s_d, b_d, mean_data, var_data, epsilon, C, fsize, x_numel, y_d);
  } else {
    KeNormAffine<T, DataLayout::NHWC><<<grid2, block, 0, stream>>>(
        x_d, s_d, b_d, mean_data, var_data, epsilon, C, fsize, x_numel, y_d);
  }
}

}  // namespace phi

#ifdef PADDLE_WITH_HIP
PD_REGISTER_KERNEL(sync_batch_norm,
                   GPU,
                   ALL_LAYOUT,
                   phi::SyncBatchNormKernel,
                   float,
                   phi::float16) {
  if (kernel_key.dtype() == phi::DataType::FLOAT16) {
    kernel->InputAt(1).SetDataType(phi::DataType::FLOAT32);
    kernel->InputAt(2).SetDataType(phi::DataType::FLOAT32);
    kernel->InputAt(3).SetDataType(phi::DataType::FLOAT32);
    kernel->InputAt(4).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
  }
}
#else
#if CUDNN_VERSION_MIN(8, 1, 0)
PD_REGISTER_KERNEL(sync_batch_norm,
                   GPU,
                   ALL_LAYOUT,
                   phi::SyncBatchNormKernel,
                   float,
                   double,
                   phi::float16,
                   phi::bfloat16) {
  if (kernel_key.dtype() == phi::DataType::FLOAT16 ||
      kernel_key.dtype() == phi::DataType::BFLOAT16) {
    kernel->InputAt(1).SetDataType(phi::DataType::FLOAT32);
    kernel->InputAt(2).SetDataType(phi::DataType::FLOAT32);
    kernel->InputAt(3).SetDataType(phi::DataType::FLOAT32);
    kernel->InputAt(4).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
  }
}
#else
PD_REGISTER_KERNEL(sync_batch_norm,
                   GPU,
                   ALL_LAYOUT,
                   phi::SyncBatchNormKernel,
                   float,
                   double,
                   phi::float16) {
  if (kernel_key.dtype() == phi::DataType::FLOAT16) {
    kernel->InputAt(1).SetDataType(phi::DataType::FLOAT32);
    kernel->InputAt(2).SetDataType(phi::DataType::FLOAT32);
    kernel->InputAt(3).SetDataType(phi::DataType::FLOAT32);
    kernel->InputAt(4).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(1).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(3).SetDataType(phi::DataType::FLOAT32);
    kernel->OutputAt(4).SetDataType(phi::DataType::FLOAT32);
  }
}
#endif
#endif