cuda_stream.hpp

#pragma once
 
#include "cuda_error.hpp"

 
namespace tf {
 
 
// ----------------------------------------------------------------------------
// cudaEventBase
// ----------------------------------------------------------------------------
  
class cudaEventCreator {
    
  public:

  cudaEvent_t operator () () const {
    cudaEvent_t event;
    TF_CHECK_CUDA(cudaEventCreate(&event), "failed to create a CUDA event");
    return event;
  }
  
  cudaEvent_t operator () (unsigned int flag) const {
    cudaEvent_t event;
    TF_CHECK_CUDA(
      cudaEventCreateWithFlags(&event, flag),
      "failed to create a CUDA event with flag=", flag
    );
    return event;
  }
  
  cudaEvent_t operator () (cudaEvent_t event) const {
    return event;
  }
};

class cudaEventDeleter {
  public:
  void operator () (cudaEvent_t event) const {
    cudaEventDestroy(event);
  }
};

template <typename Creator, typename Deleter>
class cudaEventBase : public std::unique_ptr<std::remove_pointer_t<cudaEvent_t>, Deleter> {
 
  static_assert(std::is_pointer_v<cudaEvent_t>, "cudaEvent_t is not a pointer type");
 
  public:
  
  using base_type = std::unique_ptr<std::remove_pointer_t<cudaEvent_t>, Deleter>;

  template <typename... ArgsT>
  explicit cudaEventBase(ArgsT&& ... args) : base_type(
    Creator{}(std::forward<ArgsT>(args)...), Deleter()
  ) {
  }  
  
  cudaEventBase(cudaEventBase&&) = default;

  cudaEventBase& operator = (cudaEventBase&&) = default;
  
  private:
 
  cudaEventBase(const cudaEventBase&) = delete;
  cudaEventBase& operator = (const cudaEventBase&) = delete;
};

using cudaEvent = cudaEventBase<cudaEventCreator, cudaEventDeleter>;
 
// ----------------------------------------------------------------------------
// cudaStream
// ----------------------------------------------------------------------------

class cudaStreamCreator {
  
  public:

  cudaStream_t operator () () const {
    cudaStream_t stream;
    TF_CHECK_CUDA(cudaStreamCreate(&stream), "failed to create a CUDA stream");
    return stream;
  }
  
  cudaStream_t operator () (cudaStream_t stream) const {
    return stream;
  }
};

class cudaStreamDeleter {
 
  public:

  void operator () (cudaStream_t stream) const {
    cudaStreamDestroy(stream);
  }
};

template <typename Creator, typename Deleter>
class cudaStreamBase : public std::unique_ptr<std::remove_pointer_t<cudaStream_t>, Deleter> {
 
  static_assert(std::is_pointer_v<cudaStream_t>, "cudaStream_t is not a pointer type");
  
  public:
  
  using base_type = std::unique_ptr<std::remove_pointer_t<cudaStream_t>, Deleter>;

  template <typename... ArgsT>
  explicit cudaStreamBase(ArgsT&& ... args) : base_type(
    Creator{}(std::forward<ArgsT>(args)...), Deleter()
  ) {
  }  
  
  cudaStreamBase(cudaStreamBase&&) = default;

  cudaStreamBase& operator = (cudaStreamBase&&) = default;
  
  cudaStreamBase& synchronize() {
    TF_CHECK_CUDA(
      cudaStreamSynchronize(this->get()), "failed to synchronize a CUDA stream"
    );
    return *this;
  }
  
  void begin_capture(cudaStreamCaptureMode m = cudaStreamCaptureModeGlobal) const {
    TF_CHECK_CUDA(
      cudaStreamBeginCapture(this->get(), m), 
      "failed to begin capture on stream ", this->get(), " with thread mode ", m
    );
  }

  cudaGraph_t end_capture() const {
    cudaGraph_t native_g;
    TF_CHECK_CUDA(
      cudaStreamEndCapture(this->get(), &native_g), 
      "failed to end capture on stream ", this->get()
    );
    return native_g;
  }
  
  void record(cudaEvent_t event) const {
    TF_CHECK_CUDA(
      cudaEventRecord(event, this->get()), 
      "failed to record event ", event, " on stream ", this->get()
    );
  }

  void wait(cudaEvent_t event) const {
    TF_CHECK_CUDA(
      cudaStreamWaitEvent(this->get(), event, 0), 
      "failed to wait for event ", event, " on stream ", this->get()
    );
  }

  template <typename C, typename D>
  cudaStreamBase& run(const cudaGraphExecBase<C, D>& exec);

  cudaStreamBase& run(cudaGraphExec_t exec);
 
  private:
 
  cudaStreamBase(const cudaStreamBase&) = delete;
  cudaStreamBase& operator = (const cudaStreamBase&) = delete;
};

using cudaStream = cudaStreamBase<cudaStreamCreator, cudaStreamDeleter>;
 
}  // end of namespace tf -----------------------------------------------------