| 251 | } |
| 252 | |
| 253 | bool GpuRng::SetSeed(Stream* stream, const uint8* seed, uint64 seed_bytes) { |
| 254 | absl::MutexLock lock{&mu_}; |
| 255 | CHECK(rng_ != nullptr); |
| 256 | |
| 257 | if (!CheckSeed(seed, seed_bytes)) { |
| 258 | return false; |
| 259 | } |
| 260 | |
| 261 | if (!SetStream(stream)) { |
| 262 | return false; |
| 263 | } |
| 264 | |
| 265 | // Requires 8 bytes of seed data; checked in RngSupport::CheckSeed (above) |
| 266 | // (which itself requires 16 for API consistency with host RNG fallbacks). |
| 267 | hiprandStatus_t ret = wrap::hiprandSetPseudoRandomGeneratorSeed( |
| 268 | parent_, rng_, *(reinterpret_cast<const uint64*>(seed))); |
| 269 | if (ret != HIPRAND_STATUS_SUCCESS) { |
| 270 | LOG(ERROR) << "failed to set rng seed: " << ret; |
| 271 | return false; |
| 272 | } |
| 273 | |
| 274 | ret = wrap::hiprandSetGeneratorOffset(parent_, rng_, 0); |
| 275 | if (ret != HIPRAND_STATUS_SUCCESS) { |
| 276 | LOG(ERROR) << "failed to reset rng position: " << ret; |
| 277 | return false; |
| 278 | } |
| 279 | return true; |
| 280 | } |
| 281 | |
| 282 | } // namespace gpu |
| 283 | |