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

#pragma once

#include <cstring>
#include <float.h>
#include <cmath>

#ifndef SHARED_PROVIDER
#include "core/framework/tensor.h"
#include "core/framework/op_kernel.h"
#endif

#include "orttraining/core/graph/optimizer_config.h"
#include "orttraining/core/framework/distributed_run_context.h"

#ifdef ENABLE_CPU_FP16_TRAINING_OPS
#include "orttraining/core/framework/adasum/m256_utils.h"
#endif

namespace onnxruntime {
namespace training {

static inline bool IsPowerOfTwo(unsigned x) {
  return (x != 0) && ((x & (x - 1)) == 0);
}

// Interface for Adasum algorithm
template <typename Communicator_type>
class AdasumInterface {
 public:
  Status DispatchFusedAllreduce(void* grad_buffer, void* recv_buffer,
                                std::vector<int>& tensor_counts, int start_level,
                                Communicator_type communicator, int tag,
                                const Communicator_type* reduction_comms,
                                MLDataType data_type) {
    if (data_type == DataTypeImpl::GetType<MLFloat16>()) {
      FusedAllreduce((uint16_t*)grad_buffer, (uint16_t*)recv_buffer,
                     data_type, tensor_counts, start_level, communicator, tag,
                     reduction_comms);
    } else if (data_type == DataTypeImpl::GetType<float>()) {
      FusedAllreduce((float*)grad_buffer, (float*)recv_buffer,
                     data_type, tensor_counts, start_level, communicator, tag,
                     reduction_comms);
    } else if (data_type == DataTypeImpl::GetType<double>()) {
      FusedAllreduce((double*)grad_buffer, (double*)recv_buffer,
                     data_type, tensor_counts, start_level, communicator, tag,
                     reduction_comms);
    } else {
      // Shouldn't reach here
      ORT_THROW("Unsupported datatype for Adasum allreduce.");
    }
    return Status::OK();
  }

  virtual bool IsAdasumInitialized() = 0;

  virtual void InitializeVHDDReductionComms(WorkerGroupType worker_group) = 0;

  virtual const Communicator_type* GetReductionComms() = 0;

  virtual ~AdasumInterface() = default;

 protected:
  // Communication primitives required for Adasum algorithm
  virtual void PointToPointSendRecv(void* input_data_buffer,
                                    int64_t input_buffer_bytes,
                                    void* output_data_buffer,
                                    int64_t output_buffer_bytes,
                                    MLDataType data_type, int dst_src_rank,
                                    int tag, Communicator_type communicator) = 0;

  virtual void SumAllreduceWithComm(void* data, int num_elements,
                                    MLDataType data_type,
                                    Communicator_type comm) = 0;

  virtual int GetRankWithComm(Communicator_type communicator) = 0;

  virtual int GetSizeWithComm(Communicator_type communicator) = 0;

  virtual void DispatchComputeDotAndNormSqrds(const void* __restrict__ a,
                                              const void* __restrict__ b,
                                              MLDataType data_type,
                                              int count, double& dotProduct,
                                              double& anormsq, double& bnormsq) {
    if (data_type == DataTypeImpl::GetType<MLFloat16>()) {
      ComputeDotAndNormSqrdsfp16((uint16_t*)a, (uint16_t*)b, count, dotProduct,
                                 anormsq, bnormsq);
    } else if (data_type == DataTypeImpl::GetType<float>()) {
      ComputeDotAndNormSqrds((float*)a, (float*)b, count, dotProduct, anormsq,
                             bnormsq);
    } else if (data_type == DataTypeImpl::GetType<double>()) {
      ComputeDotAndNormSqrds((double*)a, (double*)b, count, dotProduct, anormsq,
                             bnormsq);
    } else {
      // Shouldn't reach here
      ORT_THROW("Unsupported datatype for Adasum allreduce.");
    }
  }

  virtual void DispatchScaledAdd(MLDataType data_type, int count,
                                 double acoeff, void* __restrict__ a,
                                 double bcoeff, void* __restrict__ b) {
    if (data_type == DataTypeImpl::GetType<MLFloat16>()) {
      ScaledAddfp16(count, acoeff, (uint16_t*)a, bcoeff, (uint16_t*)b);
    } else if (data_type == DataTypeImpl::GetType<float>()) {
      ScaledAdd(count, acoeff, (float*)a, bcoeff, (float*)b);
    } else if (data_type == DataTypeImpl::GetType<double>()) {
      ScaledAdd(count, acoeff, (double*)a, bcoeff, (double*)b);
    } else {
      // Shouldn't reach here
      ORT_THROW("Unsupported datatype for Adasum allreduce.");
    }
  }

 private:
  // Allocator for temporary buffer allocations
  AllocatorPtr allocator_ = nullptr;

  // Temp buffer used by Adasum operations
  uint8_t* recv_buffer_ = nullptr;

  // Keep track of current recv buffer length
  uint64_t current_recv_buffer_length = 0;

  // Perform Adasum allreduce using a vector-halving, distance-doubling (VHDD)
  // approach. grad_buffer: holds the data to reduce and will hold the result.
  // recv_buffer: must point to a buffer of the same size as grad_buffer.
  // data_type: the element type of grad_buffer.
  // tensor_counts: is a list of how many elements grad_buffer contains for each
  // tensor
  //                involved in the allreduce. It should contain a 0 if this
  //                rank holds no data for the tensor (see start_level below for
  //                when this can happen).
  // start_level: set to 1 to perform all levels of the operation. When set to
  // n>1 the
  //              first n-1 levels are skipped. This is useful when the
  //              communication inside the node is implemented using another
  //              reduce-scatter algorithm, e.g. the one in NCCL, which may be
  //              desirable on some hardware configurations. When
  //              start_level>1, tensor_counts must be set according to the
  //              slices owned by this rank.
  // communicator: the communicator to reduce with.
  // tag: a value used as the message tag for each send/recv in this algorithm.
  // This is
  //      useful for multithreaded scenarios. Remember to also create separate
  //      reduction_comms instances when running with multiple threads.
  // reduction_comms: pointer to an array of communicators for computing dot
  // products and
  //                  norms for Adasum. The communicators should include exactly
  //                  the ranks that this rank has either directly or indirectly
  //                  communicated with after each level of VHDD.
  template <typename T>
  void FusedAllreduce(T* grad_buffer, T* recv_buffer, MLDataType data_type,
                      std::vector<int>& tensor_counts, int start_level,
                      Communicator_type communicator, int tag,
                      const Communicator_type* reduction_comms) {
    int per_element_size = data_type->Size();
    int rank = GetRankWithComm(communicator);
    int size = GetSizeWithComm(communicator);

    if (IsPowerOfTwo(size) == false) {
      ORT_THROW(
          "Adasum doesn't currently support reduction among non-power-of-2 number of ranks.");
    }

    std::vector<std::vector<int>> nghrCountVec;
    std::vector<double> normAndDots(tensor_counts.size() * 3 * 2);

    int nearest_power_2 = 1;
    for (nearest_power_2 = 1; (nearest_power_2 << 1) <= size;
         nearest_power_2 = (nearest_power_2 << 1)) {
    }
    int level;

    int nghrCountVec_index = 0;
    int orgSize = size;
    size = nearest_power_2;

    size_t total_counts_sum = 0;
    for (size_t i = 0; i < tensor_counts.size(); i++)
      total_counts_sum += tensor_counts[i];
    size_t myCount = total_counts_sum;
    int comm_index;
    for (level = 1, comm_index = 0; level < size;
         level = (level << 1), comm_index++) {
      if (level < start_level) {
        continue;
      }

      int neighbor_rank = rank ^ level;
      int nghrCount = 0;
      int sendOffset = 0;
      int recvOffset = 0;
      int firstHalfMyCount = (int)(myCount >> 1);
      int secondHalfMyCount = (int)myCount - firstHalfMyCount;

      nghrCountVec.emplace_back();
      nghrCountVec[nghrCountVec_index].resize(tensor_counts.size());

      size_t myCountSoFar = 0;
      int nghrCountSoFar = 0;
      if ((rank & level) != 0) {
        myCount = (size_t)secondHalfMyCount;
        nghrCount = firstHalfMyCount;
        sendOffset = 0;
        recvOffset = nghrCount;

        for (size_t i = 0; i < tensor_counts.size(); i++) {
          if (nghrCountSoFar <= nghrCount) {
            if (nghrCountSoFar + tensor_counts[i] <= nghrCount) {
              nghrCountVec[nghrCountVec_index][i] = tensor_counts[i];
              tensor_counts[i] = 0;
            } else {
              nghrCountVec[nghrCountVec_index][i] =
                  nghrCount - nghrCountSoFar;  // should not be negative
              tensor_counts[i] =
                  tensor_counts[i] -
                  (nghrCount - nghrCountSoFar);  // should not be negative
            }
          } else {
            tensor_counts[i] = tensor_counts[i];
            nghrCountVec[nghrCountVec_index][i] = 0;
          }
          nghrCountSoFar += nghrCountVec[nghrCountVec_index][i];
          myCountSoFar += tensor_counts[i];
        }
      } else {
        myCount = (size_t)firstHalfMyCount;
        nghrCount = secondHalfMyCount;
        sendOffset = myCount;
        recvOffset = 0;

        for (size_t i = 0; i < tensor_counts.size(); i++) {
          if (myCountSoFar <= myCount) {
            if (myCountSoFar + tensor_counts[i] <= myCount) {
              tensor_counts[i] = tensor_counts[i];
              nghrCountVec[nghrCountVec_index][i] = 0;
            } else {
              nghrCountVec[nghrCountVec_index][i] =
                  tensor_counts[i] -
                  (myCount - myCountSoFar);  // should not be negative
              tensor_counts[i] =
                  myCount - myCountSoFar;  // should not be negative
            }
          } else {
            nghrCountVec[nghrCountVec_index][i] = tensor_counts[i];
            tensor_counts[i] = 0;
          }
          nghrCountSoFar += nghrCountVec[nghrCountVec_index][i];
          myCountSoFar += tensor_counts[i];
        }
      }

      nghrCountVec_index++;

      this->PointToPointSendRecv(
          (uint8_t*)(&grad_buffer[sendOffset]), nghrCount * per_element_size,
          (uint8_t*)(&recv_buffer[recvOffset]), myCount * per_element_size,
          data_type, neighbor_rank, tag, communicator);
      if ((rank & level) != 0) {
        grad_buffer = &grad_buffer[nghrCount];
        recv_buffer = &recv_buffer[nghrCount];
      }
      FusedPairwiseReduceWithComm((uint8_t*)grad_buffer, (uint8_t*)recv_buffer,
                                  data_type, tensor_counts, reduction_comms[comm_index],
                                  (rank & level) == 0, normAndDots);
    }

    for (level = (size >> 1); level > 0; level = (level >> 1)) {
      if (level < start_level) {
        continue;
      }
      int neighbor_rank = rank ^ level;

      nghrCountVec_index--;
      int nghrCount = 0;
      for (size_t i = 0; i < tensor_counts.size(); i++) {
        nghrCount += nghrCountVec[nghrCountVec_index][i];
        tensor_counts[i] += nghrCountVec[nghrCountVec_index][i];
      }

      if ((rank & level) == 0) {
        recv_buffer = &grad_buffer[myCount];
      } else {
        recv_buffer = &grad_buffer[-nghrCount];
      }
      this->PointToPointSendRecv(grad_buffer, myCount * per_element_size,
                                 recv_buffer, nghrCount * per_element_size,
                                 data_type, neighbor_rank, tag,
                                 communicator);
      if ((rank & level) != 0) {
        grad_buffer = &grad_buffer[-nghrCount];
      }
      myCount += nghrCount;
    }
    size = orgSize;
  }

  void FusedPairwiseReduceWithComm(uint8_t* a, uint8_t* b,
                                   MLDataType data_type,
                                   std::vector<int>& tensor_counts,
                                   const Communicator_type& comm, bool isLeftNeighbor,
                                   std::vector<double>& normAndDots) {
    static double sqrt_double_min = std::sqrt(DBL_MIN);
    int per_element_size = data_type->Size();
    int bytesSoFar = 0;
    for (size_t i = 0; i < tensor_counts.size(); i++) {
      double dotProduct = 0.;
      double anormsq = 0.;
      double bnormsq = 0.;

      DispatchComputeDotAndNormSqrds(&a[bytesSoFar], &b[bytesSoFar],
                                     data_type, tensor_counts[i],
                                     dotProduct, anormsq, bnormsq);
      normAndDots[i * 3] = dotProduct;
      if (isLeftNeighbor) {
        normAndDots[i * 3 + 1] = anormsq;
        normAndDots[i * 3 + 2] = bnormsq;
      } else {
        normAndDots[i * 3 + 1] = bnormsq;
        normAndDots[i * 3 + 2] = anormsq;
      }
      bytesSoFar += tensor_counts[i] * per_element_size;
    }

    SumAllreduceWithComm((void*)normAndDots.data(),
                         3 * tensor_counts.size(), DataTypeImpl::GetType<double>(),
                         comm);

    bytesSoFar = 0;
    for (size_t i = 0; i < tensor_counts.size(); i++) {
      double dotProduct = normAndDots[i * 3];
      double anormsq;
      double bnormsq;
      if (isLeftNeighbor) {
        anormsq = normAndDots[i * 3 + 1];
        bnormsq = normAndDots[i * 3 + 2];
      } else {
        bnormsq = normAndDots[i * 3 + 1];
        anormsq = normAndDots[i * 3 + 2];
      }

      double acoeff = 1;
      double bcoeff = 1;
      if (anormsq >= sqrt_double_min) {
        acoeff = 1.0 - dotProduct / anormsq * 0.5;
      }
      if (bnormsq >= sqrt_double_min) {
        bcoeff = 1.0 - dotProduct / bnormsq * 0.5;
      }

      DispatchScaledAdd(data_type, tensor_counts[i], acoeff,
                        &a[bytesSoFar], bcoeff, &b[bytesSoFar]);
      bytesSoFar += tensor_counts[i] * per_element_size;
    }
  }

  // Given two vectors compute their dot product and the squared norm for each.
  template <typename T>
  void ComputeDotAndNormSqrds(const T* __restrict__ a, const T* __restrict__ b,
                              int count, double& dotProduct, double& anormsq,
                              double& bnormsq) {
    dotProduct = 0.;
    anormsq = 0.;
    bnormsq = 0.;

    for (int i = 0; i < count; i++) {
      dotProduct += (double)a[i] * (double)b[i];
      anormsq += (double)a[i] * (double)a[i];
      bnormsq += (double)b[i] * (double)b[i];
    }
  }

  // Update a vector to a linear combination of itself and another vector.
  template <typename T>
  void ScaledAdd(int n, double acoeff, T* __restrict__ a, double bcoeff,
                 T* __restrict__ b) {
    for (int i = 0; i < n; i++) {
      a[i] = acoeff * a[i] + bcoeff * b[i];
    }
  }
};

}  // namespace training
}  // namespace onnxruntime
