CUDAUtilities/interface/StreamCache.h

0001 #ifndef HeterogeneousCore_CUDAUtilities_StreamCache_h
0002 #define HeterogeneousCore_CUDAUtilities_StreamCache_h
0003
0004 #include <vector>
0005
0006 #include <cuda_runtime.h>
0007
0008 #include "FWCore/Utilities/interface/ReusableObjectHolder.h"
0009 #include "HeterogeneousCore/CUDAUtilities/interface/SharedStreamPtr.h"
0010
0011 class CUDAService;
0012
0013 namespace cms {
0014   namespace cuda {
0015     class StreamCache {
0016     public:
0017       using BareStream = SharedStreamPtr::element_type;
0018
0019       StreamCache();
0020
0021       // Gets a (cached) CUDA stream for the current device. The stream
0022       // will be returned to the cache by the shared_ptr destructor.
0023       // This function is thread safe
0024       SharedStreamPtr get();
0025
0026     private:
0027       friend class ::CUDAService;
0028       // not thread safe, intended to be called only from CUDAService destructor
0029       void clear();
0030
0031       class Deleter {
0032       public:
0033         Deleter() = default;
0034         Deleter(int d) : device_{d} {}
0035         void operator()(cudaStream_t stream) const;
0036
0037       private:
0038         int device_ = -1;
0039       };
0040
0041       std::vector<edm::ReusableObjectHolder<BareStream, Deleter>> cache_;
0042     };
0043
0044     // Gets the global instance of a StreamCache
0045     // This function is thread safe
0046     StreamCache& getStreamCache();
0047   }  // namespace cuda
0048 }  // namespace cms
0049
0050 #endif