/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

#pragma once

#include <cuda.h>
#include <cuda_runtime_api.h>

#include <stdexcept>
#include <string>

#include "nccl.h"

namespace kineto_stress_test {

inline void checkCudaStatus(cudaError_t status, int lineNumber = -1) {
  if (status != cudaSuccess) {
    printf(
        "PID %d --> CUDA API failed with status %d: %s at line %d\n",
        getpid(),
        status,
        cudaGetErrorString(status),
        lineNumber);
    exit(EXIT_FAILURE);
  }
}

#define CUDA_CHECK(EXPR)                            \
  do {                                              \
    const cudaError_t err = EXPR;                   \
    if (err == cudaSuccess) {                       \
      break;                                        \
    }                                               \
    std::string error_message;                      \
    error_message.append(__FILE__);                 \
    error_message.append(":");                      \
    error_message.append(std::to_string(__LINE__)); \
    error_message.append(" CUDA error: ");          \
    error_message.append(cudaGetErrorString(err));  \
    throw std::runtime_error(error_message);        \
  } while (0)

#define CUDA_KERNEL_LAUNCH_CHECK() CUDA_CHECK(cudaGetLastError())

#define MPICHECK(cmd)                                  \
  do {                                                 \
    int e = cmd;                                       \
    if (e != MPI_SUCCESS) {                            \
      printf(                                          \
          "PID %d --> Failed: MPI error %s:%d '%d'\n", \
          getpid(),                                    \
          __FILE__,                                    \
          __LINE__,                                    \
          e);                                          \
      exit(EXIT_FAILURE);                              \
    }                                                  \
  } while (0)

#define NCCLCHECK(cmd)                                  \
  do {                                                  \
    ncclResult_t r = cmd;                               \
    if (r != ncclSuccess) {                             \
      printf(                                           \
          "PID %d --> Failed, NCCL error %s:%d '%s'\n", \
          getpid(),                                     \
          __FILE__,                                     \
          __LINE__,                                     \
          ncclGetErrorString(r));                       \
      exit(EXIT_FAILURE);                               \
    }                                                   \
  } while (0)

} // namespace kineto_stress_test