/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/
#ifndef CSRC_TORCH_API_COMPUTE_ALLREDUCE_COMM_OVERLAP_H_
#define CSRC_TORCH_API_COMPUTE_ALLREDUCE_COMM_OVERLAP_H_
#include <cstddef>
#include <cstdint>
#include <optional>
#include <vector>
#include "cncl.h"
#include "framework/core/MLUEvent.h"
#include "framework/core/MLUStream.h"
#include "torch/extension.h"
#include "torch_api/utils.h"
namespace tmo {
namespace torch_api {

cnclDataType_t getCnclDataType(const at::ScalarType &data_type);

template <typename MmOp>
struct ParallelAllReduce {
  static std::vector<torch_mlu::MLUEvent> event_;
  static std::optional<torch_mlu::MLUStream> stream_;
  cnclComm_t cncl_comm_;
  std::vector<at::Tensor> d_list_;

  ParallelAllReduce(int64_t cncl_comm) { cncl_comm_ = reinterpret_cast<cnclComm_t>(cncl_comm); }

  at::Tensor operator()(MmOp &mm) {
    auto compute_stream_ = torch_mlu::getCurrentMLUStream();
    if (!stream_.has_value()) {
      stream_.emplace(torch_mlu::getStreamFromPool());
    }
    auto comm_stream_ = stream_.value();
    uint64_t loop = mm.getLoopNum();
    if (event_.size() < loop + 1) {
      event_.resize(loop + 1);
    }
    for (uint64_t i = 0; i < loop; i++) {
      d_list_.push_back(mm.forward(i));
      event_[i].place(compute_stream_);
      event_[i].wait(comm_stream_);
      // reduce_sum d_send
      CNCL_CHECK(cnclAllReduce(getAtTensorPtr(d_list_[i]), getAtTensorPtr(d_list_[i]),
                               d_list_[i].numel(), getCnclDataType(d_list_[i].scalar_type()),
                               cnclSum, cncl_comm_, comm_stream_.stream()));
    }
    event_[loop].place(comm_stream_);
    event_[loop].wait(compute_stream_);
    return mm.getOutput();
  }
};

template <typename MmOp>
std::vector<torch_mlu::MLUEvent> ParallelAllReduce<MmOp>::event_;

template <typename MmOp>
std::optional<torch_mlu::MLUStream> ParallelAllReduce<MmOp>::stream_;
}  // namespace torch_api
}  // namespace tmo
#endif  // CSRC_TORCH_API_COMPUTE_ALLREDUCE_COMM_OVERLAP_H_
