diff --git a/NEWS.md b/NEWS.md index 1bd8f16f13b76..1640889e9c78d 100644 --- a/NEWS.md +++ b/NEWS.md @@ -44,6 +44,11 @@ Standard library changes #### Profile +#### Random + +* It's now possible to efficiently create a new `Xoshiro` instance from an existing one via `Random.fork`, which + can be useful in parallel computations ([#58193]). + #### REPL #### Test diff --git a/src/task.c b/src/task.c index ae36c98066d01..09cbe80387186 100644 --- a/src/task.c +++ b/src/task.c @@ -1035,6 +1035,7 @@ main RNG state collision. [4]: https://discourse.julialang.org/t/linear-relationship-between-xoshiro-tasks/110454 */ +// if this code is updated, the julia equivalent should be updated as well in Random.fork() void jl_rng_split(uint64_t dst[JL_RNG_SIZE], uint64_t src[JL_RNG_SIZE]) JL_NOTSAFEPOINT { // load and advance the internal LCG state diff --git a/stdlib/Random/docs/src/index.md b/stdlib/Random/docs/src/index.md index 9ef86bb0d94f8..5a998084b874a 100644 --- a/stdlib/Random/docs/src/index.md +++ b/stdlib/Random/docs/src/index.md @@ -83,6 +83,7 @@ Random.TaskLocalRNG Random.Xoshiro Random.MersenneTwister Random.RandomDevice +Random.fork ``` ## [Hooking into the `Random` API](@id rand-api-hook) diff --git a/stdlib/Random/src/RNGs.jl b/stdlib/Random/src/RNGs.jl index b8f71f3e2f3c8..1dd50e5da083c 100644 --- a/stdlib/Random/src/RNGs.jl +++ b/stdlib/Random/src/RNGs.jl @@ -920,3 +920,61 @@ function advance!(r::MersenneTwister, adv_jump, adv, adv_vals, idxF, adv_ints, i _advance_to!(r, adv, work) r end + + +## forking + +""" + Random.fork(rng::AbstractRNG, [dims...])::typeof(rng) + +Create deterministically a new independent RNG object from an existing one, of the same type. +When `dims` is specified (as integers or a tuple of integers), +an array of such independent RNGs is created. + +This is the recommended way to initialize fresh RNG instances when reproducibility is +required, and can be useful in particular in multi-threaded contexts, where race conditions +must be avoided. For example: +```julia +function dotask(rng::Xoshiro) + num_subtasks = 10 + rngs = Random.fork(rng, num_subtasks) + result = Vector(undef, num_subtasks) + Threads.@threads for i=1:num_subtasks + result[i] = dosubtask(i, rngs[i]) + end + result +end +``` +Note that this functions does mutate its `rng` argument, so calling it on the same `rng` +must not be done concurrently from multiple threads. + +Currently, a method is only implemented for `Xoshiro`. + +!!! note + `Random.fork(::Xoshiro)` uses the same algorithm as when the task-local RNG + of a new task is created from the task-local RNG of the parent task. + +!!! compat "Julia 1.13" + This function was introduced in Julia 1.13. + +# Examples +```jldoctest +julia> x = Xoshiro(0) +Xoshiro(0xdb2fa90498613fdf, 0x48d73dc42d195740, 0x8c49bc52dc8a77ea, 0x1911b814c02405e8, 0x22a21880af5dc689) + +julia> [x, fork(x)] # x is mutated +2-element Vector{Xoshiro}: + Xoshiro(0xdb2fa90498613fdf, 0x48d73dc42d195740, 0x8c49bc52dc8a77ea, 0x1911b814c02405e8, 0x26b589433d8074be) + Xoshiro(0x545f53b997598dfb, 0xac80f92d91bb35a5, 0xb6eb3382e8c409ca, 0xa9aa6968fbdd5e83, 0x26b589433d8074be) + +julia> fork(x, 3) +3-element Vector{Xoshiro}: + Xoshiro(0xfc86733bafa6df6d, 0x5ff051ecb5937fcf, 0x935c4e55a82ca686, 0xa57b44768cdb84e9, 0x845f4ebfc53d5497) + Xoshiro(0x9c49327a59542654, 0x7c2f821b7716e6b7, 0x586a3fe58fed92f7, 0x28bbf526c1aca281, 0x425e2dc6f55934e4) + Xoshiro(0xd5cbe598083243e0, 0xfba09a94aaf998af, 0xa674c207f3796c54, 0x084d4986ad49c4eb, 0x52a722b1a914a4b5) +``` +""" +fork + +fork(rng::AbstractRNG, dims::Dims) = typeof(rng)[fork(rng) for _=LinearIndices(dims)] +fork(rng::AbstractRNG, d1::Integer, dims::Integer...) = fork(rng, Dims((d1, dims...))) diff --git a/stdlib/Random/src/Random.jl b/stdlib/Random/src/Random.jl index 796b3edfe6321..20b847a45007a 100644 --- a/stdlib/Random/src/Random.jl +++ b/stdlib/Random/src/Random.jl @@ -29,7 +29,7 @@ export rand!, randn!, randcycle, randcycle!, AbstractRNG, MersenneTwister, RandomDevice, TaskLocalRNG, Xoshiro -public seed!, default_rng, Sampler, SamplerType, SamplerTrivial, SamplerSimple +public fork, seed!, default_rng, Sampler, SamplerType, SamplerTrivial, SamplerSimple ## general definitions diff --git a/stdlib/Random/src/Xoshiro.jl b/stdlib/Random/src/Xoshiro.jl index 63a027046c545..c8bbbb9680d4d 100644 --- a/stdlib/Random/src/Xoshiro.jl +++ b/stdlib/Random/src/Xoshiro.jl @@ -246,6 +246,9 @@ hash(x::Union{TaskLocalRNG, Xoshiro}, h::UInt) = hash(getstate(x), h + 0x49a62c2 seed!(rng::Union{TaskLocalRNG, Xoshiro}, seeder::AbstractRNG) = initstate!(rng, rand(seeder, NTuple{4, UInt64})) +# when seeder is a Xoshiro, use forking instead of regular seeding, to avoid correlations +seed!(rng::Union{TaskLocalRNG, Xoshiro}, seeder::Union{TaskLocalRNG, Xoshiro}) = + setstate!(rng, _fork(seeder)) @inline function rand(x::Union{TaskLocalRNG, Xoshiro}, ::SamplerType{UInt64}) s0, s1, s2, s3 = getstate(x) @@ -296,3 +299,40 @@ for FT in (Float16, Float32, Float64) @eval rand(r::Union{TaskLocalRNG, Xoshiro}, ::SamplerTrivial{CloseOpen01{$(FT)}}) = _uint2float(rand(r, $(UT)), $(FT)) end + + +## fork Xoshiro RNGs, in the same way that TaskLocalRNG() is forked upon task spawning +## cf. jl_rng_split in src/task.c + +const xoshiro_split_a = ( + 0x214c146c88e47cb7, + 0xa66d8cc21285aafa, + 0x68c7ef2d7b1a54d4, + 0xb053a7d7aa238c61, +) + +const xoshiro_split_m = ( + 0xaef17502108ef2d9, + 0xf34026eeb86766af, + 0x38fd70ad58dd9fbb, + 0x6677f9b93ab0c04d, +) + +function _fork(src::Union{Xoshiro, TaskLocalRNG}) + s0, s1, s2, s3, s4 = getstate(src) + x = s4 + s4 = x * 0xd1342543de82ef95 + 1 + + state = map((s0, s1, s2, s3), xoshiro_split_a, xoshiro_split_m) do c, ai, mi + w = x ⊻ ai + c += w * (2*c + 1) + c ⊻= c >> ((c >> 59) + 5) + c *= mi + c ⊻= c >> 43 + c + end + setstate!(src, (s0, s1, s2, s3, s4)) # only for s4, other values are unchanged + (state..., s4) +end + +fork(src::Xoshiro) = Xoshiro(_fork(src)...) diff --git a/stdlib/Random/test/runtests.jl b/stdlib/Random/test/runtests.jl index cbd8eb41f5515..029d5833b5b0b 100644 --- a/stdlib/Random/test/runtests.jl +++ b/stdlib/Random/test/runtests.jl @@ -1299,3 +1299,60 @@ end @testset "Docstrings" begin @test isempty(Docs.undocumented_names(Random)) end + +@testset "fork" begin + xx = copy(TaskLocalRNG()) + x0 = copy(xx) + x1 = Random.fork(xx) + x2 = fetch(@async copy(TaskLocalRNG())) + @test x1 isa Xoshiro && x2 isa Xoshiro + @test x1 == x2 # currently, equality involves all 5 UInt64 words of the state + @test xx == TaskLocalRNG() + + # fork is deterministic + copy!(xx, x0) + x3 = Random.fork(xx) + @test x3 == x2 + @test rand(x3, UInt128) == rand(x2, UInt128) + + # seed! uses the same mechanism + copy!(xx, x0) + x4 = Xoshiro(0, 0, 0, 0, 0) + @test x4 === Random.seed!(x4, xx) + @test x4 == x1 + copy!(TaskLocalRNG(), x0) + x5 = Xoshiro(0, 0, 0, 0, 0) + @test x5 === Random.seed!(x5, TaskLocalRNG()) + @test x5 == x1 + @test xx == TaskLocalRNG() # both are in the same state after being forked off + + @test TaskLocalRNG() == Random.seed!(TaskLocalRNG(), copy!(xx, x0)) + @test TaskLocalRNG() == x4 + copy!(TaskLocalRNG(), x0) + @test TaskLocalRNG() == Random.seed!(TaskLocalRNG(), TaskLocalRNG()) + + # self-seeding + copy!(xx, x0) + copy!(TaskLocalRNG(), x0) + @test xx === Random.seed!(xx, xx) + @test TaskLocalRNG() === Random.seed!(TaskLocalRNG(), TaskLocalRNG()) + @test xx == TaskLocalRNG() + @test xx == x1 + + # arrays + y2 = fork(xx, 2) + @test y2 isa Vector{Xoshiro} && size(y2) == (2,) + y34 = fork(xx, 3, 0x4) + @test y34 isa Matrix{Xoshiro} && size(y34) == (3, 4) + y123 = fork(xx, (1, 2, 3)) + @test y123 isa Array{Xoshiro, 3} && size(y123) == (1, 2, 3) + y0 = fork(xx, ()) + @test y0 isa Array{Xoshiro, 0} && size(y0) == () + @test allunique([y2..., y34..., y123..., y0...]) + + # fork is currently unsupported for other RNGs + for rng = (RandomDevice(), MersenneTwister(0), TaskLocalRNG(), SeedHasher(0)) + @test_throws MethodError fork(rng) + @test_throws MethodError fork(rng, (2, 3)) + end +end