ActsExamples/Utilities/CudaStream.hpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #pragma once
0010
0011 #include <sstream>
0012 #include <stdexcept>
0013 #include <utility>
0014
0015 #include <cuda_runtime_api.h>
0016
0017 namespace ActsExamples {
0018
0019 /// Owning CUDA stream with non-throwing cleanup.
0020 class CudaStream {
0021  public:
0022   CudaStream() {
0023     check(cudaStreamCreateWithFlags(&m_stream, cudaStreamNonBlocking), __FILE__,
0024           __LINE__);
0025   }
0026
0027   CudaStream(const CudaStream&) = delete;
0028   CudaStream& operator=(const CudaStream&) = delete;
0029
0030   CudaStream(CudaStream&& other) noexcept
0031       : m_stream{std::exchange(other.m_stream, nullptr)} {}
0032
0033   CudaStream& operator=(CudaStream&& other) noexcept {
0034     if (this != &other) {
0035       reset();
0036       m_stream = std::exchange(other.m_stream, nullptr);
0037     }
0038     return *this;
0039   }
0040
0041   ~CudaStream() noexcept { reset(); }
0042
0043   cudaStream_t get() const noexcept { return m_stream; }
0044
0045   void synchronize() const {
0046     check(cudaStreamSynchronize(m_stream), __FILE__, __LINE__);
0047   }
0048
0049  private:
0050   static void check(cudaError_t code, const char* file, int line) {
0051     if (code != cudaSuccess) {
0052       std::stringstream ss;
0053       ss << "CUDA error: " << cudaGetErrorString(code) << ", " << file << ":"
0054          << line;
0055       throw std::runtime_error(ss.str());
0056     }
0057   }
0058
0059   void reset() noexcept {
0060     if (m_stream != nullptr) {
0061       (void)cudaStreamDestroy(m_stream);
0062       m_stream = nullptr;
0063     }
0064   }
0065
0066   cudaStream_t m_stream = nullptr;
0067 };
0068
0069 }  // namespace ActsExamples