abstract_cuda_task.h

// NIST-developed software is provided by NIST as a public service. You may use, copy and distribute copies of the
// software in any medium, provided that you keep intact this entire notice. You may improve, modify and create
// derivative works of the software or any portion of the software, and you may copy and distribute such modifications
// or works. Modified works should carry a notice stating that you changed the software and should note the date and
// nature of any such change. Please explicitly acknowledge the National Institute of Standards and Technology as the
// source of the software. NIST-developed software is expressly provided "AS IS." NIST MAKES NO WARRANTY OF ANY KIND,
// EXPRESS, IMPLIED, IN FACT OR ARISING BY OPERATION OF LAW, INCLUDING, WITHOUT LIMITATION, THE IMPLIED WARRANTY OF
// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, NON-INFRINGEMENT AND DATA ACCURACY. NIST NEITHER REPRESENTS NOR
// WARRANTS THAT THE OPERATION OF THE SOFTWARE WILL BE UNINTERRUPTED OR ERROR-FREE, OR THAT ANY DEFECTS WILL BE
// CORRECTED. NIST DOES NOT WARRANT OR MAKE ANY REPRESENTATIONS REGARDING THE USE OF THE SOFTWARE OR THE RESULTS
// THEREOF, INCLUDING BUT NOT LIMITED TO THE CORRECTNESS, ACCURACY, RELIABILITY, OR USEFULNESS OF THE SOFTWARE. You
// are solely responsible for determining the appropriateness of using and distributing the software and you assume
// all risks associated with its use, including but not limited to the risks and costs of program errors, compliance
// with applicable laws, damage to or loss of data, programs or equipment, and the unavailability or interruption of 
// operation. This software is not intended to be used in any situation where a failure could cause risk of injury or
// damage to property. The software developed by NIST employees is not subject to copyright protection within the
// United States.
//

#ifndef HEDGEHOG_ABSTRACT_CUDA_TASK_H
#define HEDGEHOG_ABSTRACT_CUDA_TASK_H
//#ifdef HH_USE_CUDA
#include <unordered_set>
#include <cublas.h>
#include <cuda_runtime.h>
#include "abstract_task.h"

namespace hh {
#ifndef checkCudaErrors
inline void __checkCudaErrors(cudaError_t err, const char *file, const int line) {
  if (cudaSuccess != err) {
    std::cerr << "checkCudaErrors() Cuda error = "
              << err
              << "\"" << cudaGetErrorString(err) << " \" from "
              << file << ":" << line << std::endl;
    exit(43);
  }
}

inline void __checkCudaErrors(cublasStatus_t status, const char *file, const int line) {
  if (CUBLAS_STATUS_SUCCESS != status) {
    std::cerr << "checkCudaErrors() Status Error = "
              << status << " from "
              << file << ":" << line << std::endl;
    exit(44);
  }
}
#define checkCudaErrors(err) __checkCudaErrors(err, __FILE__, __LINE__)
#endif

template<class TaskOutput, class ... TaskInputs>
class AbstractCUDATask : public AbstractTask<TaskOutput, TaskInputs...> {
  static_assert(traits::isUnique < TaskInputs...>, "A Task can't accept multiple inputs with the same type.");
  static_assert(sizeof... (TaskInputs) >= 1, "A node need to have one output type and at least one output type.");

 private:
  bool enablePeerAccess_ = false;               
  std::unordered_set<int> peerDeviceIds_ = {};  
  cudaStream_t stream_ = {};                    

 public:
  explicit AbstractCUDATask(size_t numberThreads = 1)
      : AbstractTask<TaskOutput, TaskInputs...>("CudaTask", numberThreads, false),
        enablePeerAccess_(true) {
    this->core()->isCudaRelated(true);
  }

  explicit AbstractCUDATask(std::string_view const &name, size_t numberThreads = 1)
      : AbstractTask<TaskOutput, TaskInputs...>(name, numberThreads, false),
        enablePeerAccess_(true) {
    this->core()->isCudaRelated(true);
  }

  AbstractCUDATask(std::string_view const &name, size_t numberThreads, bool automaticStart, bool enablePeerAccess)
      : AbstractTask<TaskOutput, TaskInputs...>(name, numberThreads, automaticStart),
        enablePeerAccess_(enablePeerAccess) {
    this->core()->isCudaRelated(true);
  }

  void initialize() final {
    int numGpus = 0;
    int canAccess = 0;
    checkCudaErrors(cudaGetDeviceCount(&numGpus));
    assert(this->deviceId() < numGpus);
    checkCudaErrors(cudaSetDevice(this->deviceId()));
    checkCudaErrors(cudaStreamCreate(&stream_));

    if (enablePeerAccess_) {
      for (int i = 0; i < numGpus; ++i) {
        if (i != this->deviceId()) {
          checkCudaErrors(cudaDeviceCanAccessPeer(&canAccess, this->deviceId(), i));

          if (canAccess) {
            auto ret = cudaDeviceEnablePeerAccess(i, 0);
            if (ret != cudaErrorPeerAccessAlreadyEnabled) {
              checkCudaErrors(ret);
            }
            peerDeviceIds_.insert(i);
          }
        }
      }
    }
    auto ret = cudaGetLastError();
    if (ret != cudaErrorPeerAccessAlreadyEnabled) {
      checkCudaErrors(ret);
    }
    this->initializeCuda();
  }

  void shutdown() final {
    this->shutdownCuda();
    checkCudaErrors(cudaStreamDestroy(stream_));
  }

  virtual void initializeCuda() {}

  virtual void shutdownCuda() {}

  bool enablePeerAccess() const { return enablePeerAccess_; }

  cudaStream_t stream() const { return stream_; }

  bool hasPeerAccess(int peerDeviceId) { return peerDeviceIds_.find(peerDeviceId) != peerDeviceIds_.end(); }
};
}

//#endif //HH_USE_CUDA

#endif //HEDGEHOG_ABSTRACT_CUDA_TASK_H