| 310 | } |
| 311 | |
| 312 | Status CudaDevice::Stream::WaitEvent(const Device::SyncEvent& event) { |
| 313 | auto cuda_event = |
| 314 | checked_cast<const CudaDevice::SyncEvent*, const Device::SyncEvent*>(&event); |
| 315 | if (!cuda_event) { |
| 316 | return Status::Invalid("CudaDevice::Stream cannot Wait on non-cuda event"); |
| 317 | } |
| 318 | |
| 319 | auto cu_event = cuda_event->value(); |
| 320 | if (!cu_event) { |
| 321 | return Status::Invalid("Cuda Stream cannot wait on null event"); |
| 322 | } |
| 323 | |
| 324 | ContextSaver set_temporary(reinterpret_cast<CUcontext>(context_.get()->handle())); |
| 325 | // we are currently building with CUDA toolkit 11.0.3 which doesn't have enum |
| 326 | // values for the flags yet. The "flags" param *must* be 0 for now. |
| 327 | CU_RETURN_NOT_OK("cuStreamWaitEvent", cuStreamWaitEvent(value(), cu_event, 0)); |
| 328 | return Status::OK(); |
| 329 | } |
| 330 | |
| 331 | Status CudaDevice::Stream::Synchronize() const { |
| 332 | ContextSaver set_temporary(reinterpret_cast<CUcontext>(context_.get()->handle())); |