CUDAUtilities/src/StreamCache.cc

0001 #include "FWCore/Utilities/interface/thread_safety_macros.h"
0002 #include "HeterogeneousCore/CUDAUtilities/interface/StreamCache.h"
0003 #include "HeterogeneousCore/CUDAUtilities/interface/cudaCheck.h"
0004 #include "HeterogeneousCore/CUDAUtilities/interface/currentDevice.h"
0005 #include "HeterogeneousCore/CUDAUtilities/interface/deviceCount.h"
0006 #include "HeterogeneousCore/CUDAUtilities/interface/ScopedSetDevice.h"
0007
0008 namespace cms::cuda {
0009   void StreamCache::Deleter::operator()(cudaStream_t stream) const {
0010     if (device_ != -1) {
0011       ScopedSetDevice deviceGuard{device_};
0012       cudaCheck(cudaStreamDestroy(stream));
0013     }
0014   }
0015
0016   // StreamCache should be constructed by the first call to
0017   // getStreamCache() only if we have CUDA devices present
0018   StreamCache::StreamCache() : cache_(deviceCount()) {}
0019
0020   SharedStreamPtr StreamCache::get() {
0021     const auto dev = currentDevice();
0022     return cache_[dev].makeOrGet([dev]() {
0023       cudaStream_t stream;
0024       cudaCheck(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
0025       return std::unique_ptr<BareStream, Deleter>(stream, Deleter{dev});
0026     });
0027   }
0028
0029   void StreamCache::clear() {
0030     // Reset the contents of the caches, but leave an
0031     // edm::ReusableObjectHolder alive for each device. This is needed
0032     // mostly for the unit tests, where the function-static
0033     // StreamCache lives through multiple tests (and go through
0034     // multiple shutdowns of the framework).
0035     cache_.clear();
0036     cache_.resize(deviceCount());
0037   }
0038
0039   StreamCache& getStreamCache() {
0040     // the public interface is thread safe
0041     CMS_THREAD_SAFE static StreamCache cache;
0042     return cache;
0043   }
0044 }  // namespace cms::cuda