torch/lib/c10d/ProcessGroupNCCL.cpp
// Helper that checks if the NCCL kernels are completed on the GPUs
Helper that checks if the NCCL kernels are completed on the GPUs
Helper that checks if the NCCL kernels are completed on the GPUs
bool ProcessGroupNCCL::WorkNCCL::finishedGPUExecution() const {
  at::DeviceGuard gpuGuard;
  for (size_t i = 0; i < devices_.size(); ++i) {
    gpuGuard.set_index(devices_[i].index());
    auto& cudaEvent = cudaEvents_[i];
    auto ret = cudaEventQuery(cudaEvent.getEvent());
    if (ret != cudaSuccess && ret != cudaErrorNotReady) {
      C10D_CUDA_CHECK(ret);
    }
    if (ret == cudaErrorNotReady) {
      return false;
    }
  }
  return true;
}
bool ProcessGroupNCCL :: WorkNCCL :: finishedGPUExecution ( ) const { at :: DeviceGuard gpuGuard ; for ( size_t i = 0 ; i < devices_ . size ( ) ; ++ i ) { gpuGuard . set_index ( devices_ [ i ] . index ( ) ) ; auto & cudaEvent = cudaEvents_ [ i ] ; auto ret = cudaEventQuery ( cudaEvent . getEvent ( ) ) ; if ( ret != cudaSuccess && ret != cudaErrorNotReady ) { C10D_CUDA_CHECK ( ret ) ; } if ( ret == cudaErrorNotReady ) { return false ; } } return true ; }
Helper that checks if the NCCL kernels are completed on the GPUs
Helper that checks if the NCCL kernels are completed on the GPUs
// Checking the work's corresponding CUDA events' status