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

NEWS.md

@@ -39,6 +39,11 @@ Standard library changes
 
 #### Profile
 
+#### Random
+
+* It's now possible to efficiently create a new `Xoshiro` RNG 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,8 @@ Random.TaskLocalRNG
 Random.Xoshiro
 Random.MersenneTwister
 Random.RandomDevice
+Random.fork!
+Random.fork
 ```
 
 ## [Hooking into the `Random` API](@id rand-api-hook)

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, fork!, seed!, default_rng, Sampler, SamplerType, SamplerTrivial, SamplerSimple
 
 ## general definitions
 

stdlib/Random/src/Xoshiro.jl

@@ -309,3 +309,69 @@ 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
+
+"""
+    Random.fork!(dst::Union{Xoshiro, TaskLocalRNG}, src::Union{Xoshiro, TaskLocalRNG} = TaskLocalRNG()) -> dst
+
+Equivalent to `copy!(dst, fork(src))`.
+See also [`fork`](@ref).
+
+!!! compat "Julia 1.13"
+    This function was introduced in Julia 1.13.
+"""
+fork!(dst::Union{Xoshiro, TaskLocalRNG}, src::Union{Xoshiro, TaskLocalRNG}=TaskLocalRNG()) =
+    setstate!(dst, _fork(src))
+
+"""
+    Random.fork(src::Union{Xoshiro, TaskLocalRNG} = TaskLocalRNG())::Xoshiro
+
+Create a new `Xoshiro` object from `src`, in the same way that the task local RNG of a new
+task is created from the task local RNG of the parent task.
+This is the recommended way to initialize a fresh RNG from an existing one.
+
+!!! note
+    When `src` is of type `TaskLocalRNG`, this function is guaranteed to return an RNG of
+    type `Xoshiro` only as long as `Xoshiro` and `TaskLocalRNG` are implementing the same
+    underlying generator. It may change in the future.
+
+!!! compat "Julia 1.13"
+    This function was introduced in Julia 1.13.
+
+See also [`Random.fork!`](@ref).
+"""
+fork(src::Union{Xoshiro, TaskLocalRNG}=TaskLocalRNG()) = Xoshiro(_fork(src)...)

stdlib/Random/test/runtests.jl

@@ -1263,3 +1263,43 @@ end
 @testset "Docstrings" begin
     @test isempty(Docs.undocumented_names(Random))
 end
+
+@testset "fork" begin
+    xx = copy(TaskLocalRNG())
+    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()
+
+    x3 = Random.fork(TaskLocalRNG())
+    @test x3 isa Xoshiro
+    copy!(TaskLocalRNG(), xx) # reset its state
+    x4 = Random.fork(xx)
+    @test x4 isa Xoshiro
+    @test x3 == x4
+    copy!(xx, TaskLocalRNG())
+
+    @test xx == TaskLocalRNG() # check assumptions
+    x5 = Random.fork()
+    @test xx != TaskLocalRNG() # TaskLocalRNG() was forked off
+    copy!(TaskLocalRNG(), xx)
+    @test x5 == x4
+
+    x6 = Xoshiro(0, 0, 0, 0, 0)
+    @test x6 === Random.fork!(x6, xx)
+    copy!(xx, TaskLocalRNG())
+    @test x6 == x5
+    @test x6 === Random.fork!(x6, TaskLocalRNG())
+    copy!(TaskLocalRNG(), xx)
+    @test x6 == x5
+    @test xx == TaskLocalRNG() # check assumptions
+    @test x6 === Random.fork!(x6)
+    @test xx != TaskLocalRNG()
+    copy!(TaskLocalRNG(), xx)
+    @test x6 == x5
+
+    @test TaskLocalRNG() === Random.fork!(TaskLocalRNG(), copy(xx))
+    @test x6 === Random.fork!(x6, copy(xx))
+    @test TaskLocalRNG() == x6
+end