add Random.fork(rng::Xoshiro) to split rng into a new instance

NEWS.md

@@ -47,6 +47,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

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

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)

stdlib/Random/src/RNGs.jl

@@ -306,3 +306,61 @@ function hash_seed(str::AbstractString, ctx::SHA_CTX)
     end
     SHA.update!(ctx, (0x05,))
 end
+
+
+## forking
+
+"""
+    Random.fork(rng::AbstractRNG)::typeof(rng)
+    Random.fork(rng::AbstractRNG, dims...)::Array{typeof(rng)}
+
+Return a new independent RNG derived from `rng`, of the same type.
+When `dims` is specified (as integers or a tuple of integers),
+return an array of independent RNGs.
+
+This is the recommended method to initialize reproducible RNG instances,
+especially 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 function generally mutates its `rng` argument, so concurrent
+calls on the same `rng` from multiple threads are unsafe.
+
+!!! note
+    Currently, this function is only implemented for `Xoshiro`, where it 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 requires Julia 1.13 or later.
+
+# Examples
+```jldoctest
+julia> x = Xoshiro(0)
+Xoshiro(0xdb2fa90498613fdf, 0x48d73dc42d195740, 0x8c49bc52dc8a77ea, 0x1911b814c02405e8, 0x22a21880af5dc689)
+
+julia> [x, Random.fork(x)] # x is mutated
+2-element Vector{Xoshiro}:
+ Xoshiro(0xdb2fa90498613fdf, 0x48d73dc42d195740, 0x8c49bc52dc8a77ea, 0x1911b814c02405e8, 0x26b589433d8074be)
+ Xoshiro(0x545f53b997598dfb, 0xac80f92d91bb35a5, 0xb6eb3382e8c409ca, 0xa9aa6968fbdd5e83, 0x26b589433d8074be)
+
+julia> Random.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...)))

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
 

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)...)

stdlib/Random/test/runtests.jl

@@ -1299,3 +1299,65 @@ end
 @testset "Docstrings" begin
     @test isempty(Docs.undocumented_names(Random))
 end
+
+@testset "fork" begin
+    # fork(::Xoshiro) uses the same algorithm as TaskLocalRNG() upon task spawning,
+    # so we test here that they behave the same, but this could change
+    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()
+    @test xx != x0 # the source is mutated in fork
+    @test x1 != xx
+    @test x1 != x0
+
+    # 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 = Random.fork(xx, 2)
+    @test y2 isa Vector{Xoshiro} && size(y2) == (2,)
+    y34 = Random.fork(xx, 3, 0x4)
+    @test y34 isa Matrix{Xoshiro} && size(y34) == (3, 4)
+    y123 = Random.fork(xx, (1, 2, 3))
+    @test y123 isa Array{Xoshiro, 3} && size(y123) == (1, 2, 3)
+    y0 = Random.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 Random.fork(rng)
+        @test_throws MethodError Random.fork(rng, (2, 3))
+    end
+end