// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "orttraining/training_ops/rocm/math/div_grad.h"
#include "orttraining/training_ops/rocm/math/div_grad_impl.h"
#include "core/providers/rocm/math/binary_elementwise_ops.h"

using namespace onnxruntime::common;
namespace onnxruntime {
namespace rocm {

#define DIVGRAD_REGISTER_KERNEL_TYPED(T)                                                   \
  ONNX_OPERATOR_TYPED_KERNEL_EX(                                                           \
      DivGrad,                                                                             \
      kMSDomain,                                                                           \
      1,                                                                                   \
      T,                                                                                   \
      kRocmExecutionProvider,                                                              \
      (*KernelDefBuilder::Create()).TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
      DivGrad<T>);

DIVGRAD_REGISTER_KERNEL_TYPED(MLFloat16)
DIVGRAD_REGISTER_KERNEL_TYPED(float)
// DIVGRAD_REGISTER_KERNEL_TYPED(double)

TensorShapeVector prepended_dimension_1(const TensorShape& shape, size_t total_rank) {
  size_t input_rank = shape.NumDimensions();
  if (input_rank == total_rank)
    return shape.AsShapeVector();

  TensorShapeVector dims(total_rank, 1);

  // https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md
  // for property 3 of Multidirectional Broadcasting, we need to prepended with a dimension of length 1.
  if (input_rank > 0)
    std::copy(shape.GetDims().begin(), shape.GetDims().end(), &dims[total_rank - input_rank]);
  return dims;
}

template <typename T>
Status DivGrad<T>::ComputeInternal(OpKernelContext* context) const {
  typedef typename ToHipType<T>::MappedType HipT;

  const Tensor* dy_tensor = context->Input<Tensor>(0);
  const Tensor* a_tensor = context->Input<Tensor>(1);
  const Tensor* b_tensor = context->Input<Tensor>(2);
  const TensorShape& a_shape = a_tensor->Shape();
  const TensorShape& b_shape = b_tensor->Shape();
  const TensorShape& dy_shape = dy_tensor->Shape();

  // output shapes shall match its corresponding inputs
  Tensor* da_output_tensor = context->Output(0, a_shape);
  Tensor* db_output_tensor = context->Output(1, b_shape);
  if (!da_output_tensor && !db_output_tensor)
    return Status::OK();

  BinaryElementwisePreparation prepare;
  ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(a_tensor, b_tensor,
                                                        // TODO: BinaryElementwiseBroadcastPrepare shall take dy_tensor as const Tensor*.
                                                        const_cast<Tensor*>(dy_tensor), &prepare));
  const HipT* prepare_a_data = reinterpret_cast<const HipT*>(prepare.lhs_tensor->template Data<T>());
  const HipT* prepare_b_data = reinterpret_cast<const HipT*>(prepare.rhs_tensor->template Data<T>());
  const HipT* prepare_dy_data = reinterpret_cast<const HipT*>(prepare.output_tensor->template Data<T>());
  T* da_data = da_output_tensor ? da_output_tensor->template MutableData<T>() : nullptr;
  T* db_data = db_output_tensor ? db_output_tensor->template MutableData<T>() : nullptr;

  switch (prepare.output_rank_or_simple_broadcast) {
    case static_cast<int32_t>(SimpleBroadcast::NoBroadcast):
      ImplDivGradSimple<HipT>(
          Stream(context),
          SimpleBroadcast::NoBroadcast,
          prepare_a_data,
          prepare_b_data,
          prepare_dy_data,
          dy_shape.Size(),
          reinterpret_cast<HipT*>(da_data),
          reinterpret_cast<HipT*>(db_data));
      break;
    case static_cast<int32_t>(SimpleBroadcast::LeftScalar): {
      T* temp_da_data = nullptr;
      IAllocatorUniquePtr<T> temp_da_allocator;
      if (da_output_tensor) {
        temp_da_allocator = GetScratchBuffer<T>(dy_shape.Size(), context->GetComputeStream());
        temp_da_data = temp_da_allocator.get();
      }

      ImplDivGradSimple<HipT>(
          Stream(context),
          SimpleBroadcast::LeftScalar,
          prepare_a_data,
          prepare_b_data,
          prepare_dy_data,
          dy_shape.Size(),
          reinterpret_cast<HipT*>(temp_da_data),
          reinterpret_cast<HipT*>(db_data));

      if (da_output_tensor) {
        auto a_output_dims = prepended_dimension_1(a_shape, dy_shape.NumDimensions());
        ORT_RETURN_IF_ERROR((ReduceKernelShared<T, T, MIOPEN_REDUCE_TENSOR_NO_INDICES>(
            temp_da_data,
            dy_shape,
            da_data,
            TensorShape({}),
            MIOPEN_REDUCE_TENSOR_ADD,
            GetMiopenHandle(context),
            context->GetComputeStream(),
            a_output_dims)));
      }
      break;
    }
    case static_cast<int32_t>(SimpleBroadcast::RightScalar): {
      T* temp_db_data = nullptr;
      IAllocatorUniquePtr<T> temp_db_allocator;
      if (db_output_tensor) {
        temp_db_allocator = GetScratchBuffer<T>(dy_shape.Size(), context->GetComputeStream());
        temp_db_data = temp_db_allocator.get();
      }
      ImplDivGradSimple<HipT>(
          Stream(context),
          SimpleBroadcast::RightScalar,
          prepare_a_data,
          prepare_b_data,
          prepare_dy_data,
          dy_shape.Size(),
          reinterpret_cast<HipT*>(da_data),
          reinterpret_cast<HipT*>(temp_db_data));

      if (db_output_tensor) {
        auto b_output_dims = prepended_dimension_1(b_shape, dy_shape.NumDimensions());
        ORT_RETURN_IF_ERROR((ReduceKernelShared<T, T, MIOPEN_REDUCE_TENSOR_NO_INDICES>(
            temp_db_data,
            dy_shape,
            db_data,
            TensorShape({}),
            MIOPEN_REDUCE_TENSOR_ADD,
            GetMiopenHandle(context),
            context->GetComputeStream(),
            b_output_dims)));
      }
      break;
    }
    case static_cast<int32_t>(SimpleBroadcast::RightPerChannelBatch1):
    case static_cast<int32_t>(SimpleBroadcast::RightPerChannelBatchN): {
      T* temp_db_data = nullptr;
      IAllocatorUniquePtr<T> temp_db_allocator;
      if (db_output_tensor) {
        temp_db_allocator = GetScratchBuffer<T>(dy_shape.Size(), context->GetComputeStream());
        temp_db_data = temp_db_allocator.get();
      }
      if (prepare.output_rank_or_simple_broadcast == static_cast<int32_t>(SimpleBroadcast::RightPerChannelBatch1)) {
        // lhs(1,C,H) and rhs (C,1)
        ImplDivGradRhsPerChannelBatch1<HipT>(
            Stream(context),
            prepare_a_data,
            prepare_b_data,
            prepare_dy_data,
            dy_shape.Size(),
            prepare.fdm_H,
            reinterpret_cast<HipT*>(da_data),
            reinterpret_cast<HipT*>(temp_db_data));
      } else {
        // lhs(N,C,H) and rhs (C,1)
        ImplDivGradRhsPerChannelBatchN<HipT>(
            Stream(context),
            prepare_a_data,
            prepare_b_data,
            prepare_dy_data,
            dy_shape.Size(),
            prepare.fdm_H,
            prepare.fdm_C,
            reinterpret_cast<HipT*>(da_data),
            reinterpret_cast<HipT*>(temp_db_data));
      }

      if (db_output_tensor) {
        auto b_output_dims = prepended_dimension_1(b_shape, dy_shape.NumDimensions());
        ORT_RETURN_IF_ERROR((ReduceKernelShared<T, T, MIOPEN_REDUCE_TENSOR_NO_INDICES>(
            temp_db_data,
            dy_shape,
            db_data,
            b_shape,
            MIOPEN_REDUCE_TENSOR_ADD,
            GetMiopenHandle(context),
            context->GetComputeStream(),
            b_output_dims)));
      }
      break;
    }
    default: {
      bool need_reduce_da = da_output_tensor && a_shape.Size() != dy_shape.Size();
      bool need_reduce_db = db_output_tensor && b_shape.Size() != dy_shape.Size();
      IAllocatorUniquePtr<T> temp_da_allocator, temp_db_allocator;
      T* da_data_ref = nullptr;
      if (da_output_tensor) {
        if (need_reduce_da) {
          temp_da_allocator = GetScratchBuffer<T>(dy_shape.Size(), context->GetComputeStream());
          da_data_ref = temp_da_allocator.get();
        } else {
          da_data_ref = da_data;
        }
      }
      T* db_data_ref = nullptr;
      if (db_output_tensor) {
        if (need_reduce_db) {
          temp_db_allocator = GetScratchBuffer<T>(dy_shape.Size(), context->GetComputeStream());
          db_data_ref = temp_db_allocator.get();
        } else {
          db_data_ref = db_data;
        }
      }
      ImplDivGrad<HipT>(
          Stream(context),
          prepare.output_rank_or_simple_broadcast,
          prepare.lhs_padded_strides,
          prepare_a_data,
          prepare.rhs_padded_strides,
          prepare_b_data,
          prepare_dy_data,
          dy_shape.Size(),
          prepare.fdm_output_strides,
          reinterpret_cast<HipT*>(da_data_ref),
          reinterpret_cast<HipT*>(db_data_ref));

      if (need_reduce_da) {
        auto a_output_dims = prepended_dimension_1(a_shape, dy_shape.NumDimensions());
        ORT_RETURN_IF_ERROR((ReduceKernelShared<T, T, MIOPEN_REDUCE_TENSOR_NO_INDICES>(
            da_data_ref,
            dy_shape,
            da_data,
            a_shape,
            MIOPEN_REDUCE_TENSOR_ADD,
            GetMiopenHandle(context),
            context->GetComputeStream(),
            a_output_dims)));
      }

      if (need_reduce_db) {
        auto b_output_dims = prepended_dimension_1(b_shape, dy_shape.NumDimensions());
        ORT_RETURN_IF_ERROR((ReduceKernelShared<T, T, MIOPEN_REDUCE_TENSOR_NO_INDICES>(
            db_data_ref,
            dy_shape,
            db_data,
            b_shape,
            MIOPEN_REDUCE_TENSOR_ADD,
            GetMiopenHandle(context),
            context->GetComputeStream(),
            b_output_dims)));
      }
    }
  }
  return Status::OK();
}

}  // namespace rocm
}  // namespace onnxruntime
