sys/detail/StreamStorage.hh

0001 //----------------------------------*-C++-*----------------------------------//
0002 // Copyright 2023-2024 UT-Battelle, LLC, and other Celeritas developers.
0003 // See the top-level COPYRIGHT file for details.
0004 // SPDX-License-Identifier: (Apache-2.0 OR MIT)
0005 //---------------------------------------------------------------------------//
0006 //! \file corecel/sys/detail/StreamStorage.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include <vector>
0011
0012 #include "corecel/Assert.hh"
0013 #include "corecel/Types.hh"
0014
0015 #include "../Stream.hh"
0016 #include "../ThreadId.hh"
0017
0018 namespace celeritas
0019 {
0020 namespace detail
0021 {
0022 //---------------------------------------------------------------------------//
0023 /*!
0024  * Helper class for managing CUDA/HIP streams.
0025  */
0026 class StreamStorage
0027 {
0028   public:
0029     //!@{
0030     //! \name Type aliases
0031     using VecStream = std::vector<Stream>;
0032     //!@}
0033
0034   public:
0035     //! Construct with the default stream
0036     StreamStorage() = default;
0037
0038     // Construct by creating the given number of streams
0039     explicit StreamStorage(size_type num_streams) : streams_(num_streams) {}
0040
0041     //! Number of streams allocated
0042     size_type size() const { return streams_.size(); }
0043
0044     // Access a stream
0045     inline Stream& get(StreamId);
0046
0047   private:
0048     VecStream streams_;
0049     Stream default_stream_{nullptr};
0050 };
0051
0052 //---------------------------------------------------------------------------//
0053 // INLINE DEFINITIONS
0054 //---------------------------------------------------------------------------//
0055 /*!
0056  * Access a stream.
0057  *
0058  * If no streams have been created the default stream is returned.
0059  */
0060 Stream& StreamStorage::get(StreamId id)
0061 {
0062     if (!streams_.empty())
0063     {
0064         CELER_ASSERT(id < streams_.size());
0065         return streams_[id.get()];
0066     }
0067     return default_stream_;
0068 }
0069
0070 //---------------------------------------------------------------------------//
0071 }  // namespace detail
0072 }  // namespace celeritas