| 31 | |
| 32 | public: |
| 33 | cudaStream_t create(int device) { |
| 34 | std::lock_guard<std::mutex> l(mutex_); |
| 35 | HMP_REQUIRE(device < MaxDevices, |
| 36 | "CUDAStreamCache: device index({}) is out of range {}", |
| 37 | device, MaxDevices); |
| 38 | |
| 39 | cudaStream_t stream = 0; |
| 40 | auto &cache = streamCache_[device]; |
| 41 | if (!cache.empty()) { |
| 42 | stream = cache.back(); |
| 43 | cache.pop_back(); |
| 44 | } else { |
| 45 | int oldDevice; |
| 46 | HMP_CUDA_CHECK(cudaGetDevice(&oldDevice)); |
| 47 | HMP_CUDA_CHECK(cudaSetDevice(device)); |
| 48 | HMP_CUDA_CHECK(cudaStreamCreate(&stream)); |
| 49 | HMP_CUDA_CHECK(cudaSetDevice(oldDevice)); |
| 50 | } |
| 51 | |
| 52 | return stream; |
| 53 | } |
| 54 | |
| 55 | void destroy(cudaStream_t stream, int device) { |
| 56 | HMP_REQUIRE(device < MaxDevices, |