| 71 | } |
| 72 | |
| 73 | bool GpuRng::Init() { |
| 74 | absl::MutexLock lock(&mu_); |
| 75 | CHECK(rng_ == nullptr); |
| 76 | |
| 77 | cuda::ScopedActivateExecutorContext sac(parent_); |
| 78 | curandStatus_t ret = curandCreateGenerator(&rng_, CURAND_RNG_PSEUDO_DEFAULT); |
| 79 | if (ret != CURAND_STATUS_SUCCESS) { |
| 80 | LOG(ERROR) << "failed to create random number generator: " << ret; |
| 81 | return false; |
| 82 | } |
| 83 | |
| 84 | CHECK(rng_ != nullptr); |
| 85 | return true; |
| 86 | } |
| 87 | |
| 88 | bool GpuRng::SetStream(Stream* stream) { |
| 89 | cuda::ScopedActivateExecutorContext sac(parent_); |