/*
 * Copyright (c) 2020, NVIDIA CORPORATION.
 *
 * 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

#include <nccl.h>
#include <raft/error.hpp>
#include <string>

/**
 * @brief Error checking macro for NCCL runtime API functions.
 *
 * Invokes a NCCL runtime API function call, if the call does not return ncclSuccess, throws an
 * exception detailing the NCCL error that occurred
 */
#define NCCL_TRY(call)                                                        \
  do {                                                                        \
    ncclResult_t const status = (call);                                       \
    if (ncclSuccess != status) {                                              \
      std::string msg{};                                                      \
      SET_ERROR_MSG(msg,                                                      \
                    "NCCL error encountered at: ", "call='%s', Reason=%d:%s", \
                    #call, status, ncclGetErrorString(status));               \
      throw raft::logic_error(msg);                                           \
    }                                                                         \
  } while (0);

#define NCCL_TRY_NO_THROW(call)                           \
  do {                                                    \
    ncclResult_t status = call;                           \
    if (ncclSuccess != status) {                          \
      printf("NCCL call='%s' failed. Reason:%s\n", #call, \
             ncclGetErrorString(status));                 \
    }                                                     \
  } while (0)

namespace raft {
namespace comms {

constexpr size_t get_datatype_size(const datatype_t datatype) {
  switch (datatype) {
    case datatype_t::CHAR:
      return sizeof(char);
    case datatype_t::UINT8:
      return sizeof(uint8_t);
    case datatype_t::INT32:
      return sizeof(int);
    case datatype_t::UINT32:
      return sizeof(unsigned int);
    case datatype_t::INT64:
      return sizeof(int64_t);
    case datatype_t::UINT64:
      return sizeof(uint64_t);
    case datatype_t::FLOAT32:
      return sizeof(float);
    case datatype_t::FLOAT64:
      return sizeof(double);
    default:
      throw "Unsupported datatype";
  }
}

constexpr ncclDataType_t get_nccl_datatype(const datatype_t datatype) {
  switch (datatype) {
    case datatype_t::CHAR:
      return ncclChar;
    case datatype_t::UINT8:
      return ncclUint8;
    case datatype_t::INT32:
      return ncclInt;
    case datatype_t::UINT32:
      return ncclUint32;
    case datatype_t::INT64:
      return ncclInt64;
    case datatype_t::UINT64:
      return ncclUint64;
    case datatype_t::FLOAT32:
      return ncclFloat;
    case datatype_t::FLOAT64:
      return ncclDouble;
    default:
      throw "Unsupported datatype";
  }
}

constexpr ncclRedOp_t get_nccl_op(const op_t op) {
  switch (op) {
    case op_t::SUM:
      return ncclSum;
    case op_t::PROD:
      return ncclProd;
    case op_t::MIN:
      return ncclMin;
    case op_t::MAX:
      return ncclMax;
    default:
      throw "Unsupported datatype";
  }
}
};  // namespace comms
};  // namespace raft
