corecel/sys/Stream.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/Stream.hh
0007 //---------------------------------------------------------------------------//
0008 #pragma once
0009
0010 #include "corecel/DeviceRuntimeApi.hh"
0011
0012 #include "corecel/Assert.hh"
0013 #include "corecel/Macros.hh"
0014 #include "corecel/Types.hh"
0015
0016 namespace celeritas
0017 {
0018 //---------------------------------------------------------------------------//
0019 #if !CELER_USE_DEVICE
0020 struct MockStream_st;
0021 template<class Pointer>
0022 struct MockMemoryResource
0023 {
0024     virtual Pointer do_allocate(std::size_t, std::size_t) = 0;
0025
0026     virtual void do_deallocate(Pointer, std::size_t, std::size_t) = 0;
0027 };
0028 #endif
0029
0030 //---------------------------------------------------------------------------//
0031 /*!
0032  * Thrust async memory resource associated with a Stream.
0033  */
0034 template<class Pointer>
0035 #if CELER_USE_DEVICE
0036 class AsyncMemoryResource final : public thrust::mr::memory_resource<Pointer>
0037 #else
0038 class AsyncMemoryResource final : public MockMemoryResource<Pointer>
0039 #endif
0040 {
0041   public:
0042     //!@{
0043     //! \name Type aliases
0044     using pointer = Pointer;
0045 #if CELER_USE_DEVICE
0046     using StreamT = CELER_DEVICE_PREFIX(Stream_t);
0047 #else
0048     using StreamT = MockStream_st*;
0049 #endif
0050     //!@}
0051
0052     // Construct memory resource for the stream
0053     explicit AsyncMemoryResource(StreamT stream) : stream_{stream} {}
0054
0055     // Construct with default Stream
0056     AsyncMemoryResource() = default;
0057
0058     // Allocate device memory
0059     pointer do_allocate(std::size_t bytes, std::size_t) final;
0060
0061     // Deallocate device memory
0062     void do_deallocate(pointer p, std::size_t, std::size_t) final;
0063
0064   private:
0065     StreamT stream_{nullptr};
0066 };
0067
0068 //---------------------------------------------------------------------------//
0069 /*!
0070  * CUDA or HIP stream.
0071  *
0072  * This creates/destroys a stream on construction/destruction and provides
0073  * accessors to low-level stream-related functionality. This class will
0074  * typically be accessed only by low-level device implementations.
0075  */
0076 class Stream
0077 {
0078   public:
0079     //!@{
0080     //! \name Type aliases
0081 #if CELER_USE_DEVICE
0082     using StreamT = CELER_DEVICE_PREFIX(Stream_t);
0083 #else
0084     using StreamT = MockStream_st*;
0085 #endif
0086     using ResourceT = AsyncMemoryResource<void*>;
0087     //!@}
0088
0089   public:
0090     // Whether asynchronous operations are supported
0091     static bool async();
0092
0093     // Construct by creating a stream
0094     Stream();
0095
0096     // Construct with the default stream
0097     Stream(std::nullptr_t) {}
0098
0099     // Destroy the stream
0100     ~Stream();
0101
0102     // Move construct and assign
0103     Stream(Stream const&) = delete;
0104     Stream& operator=(Stream const&) = delete;
0105     Stream(Stream&&) noexcept;
0106     Stream& operator=(Stream&&) noexcept;
0107     void swap(Stream& other) noexcept;
0108
0109     // Access the stream
0110     StreamT get() const { return stream_; }
0111
0112     // Access the thrust resource allocator associated with the stream
0113     ResourceT& memory_resource() { return memory_resource_; }
0114
0115     // Allocate memory asynchronously on this stream if possible
0116     void* malloc_async(std::size_t bytes) const;
0117
0118     // Free memory asynchronously on this stream if possible
0119     void free_async(void* ptr) const;
0120
0121   private:
0122     StreamT stream_{nullptr};
0123     ResourceT memory_resource_;
0124 };
0125
0126 //---------------------------------------------------------------------------//
0127 }  // namespace celeritas