chunk_stream_pair.cuh

Go to the documentation of this file.

#pragma once
 
#include <cuda_runtime.h>
#include <cuda/stream>
 
#include <array>
#include <utility>
 
#include <cusbf/error.hpp>
#include <cusbf/helpers.cuh>
 
namespace cusbf::detail {
 
class ChunkStreamPair {
   public:
    ChunkStreamPair() {
        for (cudaStream_t& stream : streams_) {
            CUSBF_CUDA_ABORT(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
        }
    }
 
    ChunkStreamPair(const ChunkStreamPair&) = delete;
    ChunkStreamPair& operator=(const ChunkStreamPair&) = delete;
 
    ChunkStreamPair(ChunkStreamPair&& other) noexcept : streams_{other.streams_} {
        for (cudaStream_t& stream : other.streams_) {
            stream = nullptr;
        }
    }
 
    ChunkStreamPair& operator=(ChunkStreamPair&& other) noexcept {
        if (this != &other) {
            destroy();
            streams_ = other.streams_;
            for (cudaStream_t& stream : other.streams_) {
                stream = nullptr;
            }
        }
        return *this;
    }
 
    ~ChunkStreamPair() {
        destroy();
    }
 
    [[nodiscard]] cuda::stream_ref operator[](size_t index) const noexcept {
        return {streams_.at(index)};
    }
 
    [[nodiscard]] Result<void> sync_all() const {
        for (cudaStream_t stream : streams_) {
            if (stream != nullptr) {
                CUSBF_CUDA_TRY(cudaStreamSynchronize(stream));
            }
        }
        return {};
    }
 
   private:
    std::array<cudaStream_t, 2> streams_{nullptr, nullptr};
 
    void destroy() {
        for (cudaStream_t& stream : streams_) {
            if (stream != nullptr) {
                CUSBF_CUDA_ABORT(cudaStreamDestroy(stream));
                stream = nullptr;
            }
        }
    }
};
 
}  // namespace cusbf::detail