Skip to content

Commit

Permalink
Add jump functions (jump, long_jump) for Xoshiro (#47743)
Browse files Browse the repository at this point in the history
Straightforward implementations given
https://xoshiro.di.unimi.it/xoshiro256plusplus.c

This is useful functionality which does not interfere with any existing
code. Utility aside, it is arguable that the jump functions are
necessary to complete the implementation of `xoshiro256++` tooling. In
essence, given that the `xoshiro` family supports jump functions, we
would be remiss to neglect this capability.

For most users, I expect that the existing `TaskLocalRNG` is more than
sufficient (hence, why I propose that this not be exported, just like
`seed!`).

However, if/when one does want to total control, jump functions are a
requirement. Use cases arise when one wishes to utilize a single seed
(with state advanced sufficiently far as to give non-overlapping
subsequences) as the basis for a parallel computation. The alternative
is manual seeding, which
lacks the flexibility required for testing (imagine a program which
requires a variable number of sub-sequences, one for each parallel
portion).

If further justification is needed, [good
precedent](https://docs.rs/rand_xoshiro/latest/rand_xoshiro/struct.Xoshiro256PlusPlus.html)
exists.
  • Loading branch information
andrewjradcliffe committed Sep 28, 2023
1 parent ed891d6 commit 2d93567
Show file tree
Hide file tree
Showing 2 changed files with 165 additions and 0 deletions.
102 changes: 102 additions & 0 deletions stdlib/Random/src/Xoshiro.jl
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,108 @@ end

rng_native_52(::Xoshiro) = UInt64

# Jump functions from: https://xoshiro.di.unimi.it/xoshiro256plusplus.c

for (fname, JUMP) in ((:jump_128, (0x180ec6d33cfd0aba, 0xd5a61266f0c9392c, 0xa9582618e03fc9aa, 0x39abdc4529b1661c)),
(:jump_192, (0x76e15d3efefdcbbf, 0xc5004e441c522fb3, 0x77710069854ee241, 0x39109bb02acbe635)))
local fname! = Symbol(fname, :!)
@eval function $fname!(rng::Xoshiro)
_s0 = 0x0000000000000000
_s1 = 0x0000000000000000
_s2 = 0x0000000000000000
_s3 = 0x0000000000000000
s0, s1, s2, s3 = rng.s0, rng.s1, rng.s2, rng.s3
for j in $JUMP
for b in 0x0000000000000000:0x000000000000003f
if (j & 0x0000000000000001 << b) != 0
_s0 ⊻= s0
_s1 ⊻= s1
_s2 ⊻= s2
_s3 ⊻= s3
end
t = s1 << 17
s2 = xor(s2, s0)
s3 = xor(s3, s1)
s1 = xor(s1, s2)
s0 = xor(s0, s3)
s2 = xor(s2, t)
s3 = s3 << 45 | s3 >> 19
end
end
setstate!(rng, (_s0, _s1, _s2, _s3, nothing))
end
@eval $fname(rng::Xoshiro) = $fname!(copy(rng))

@eval function $fname!(rng::Xoshiro, n::Integer)
n < 0 && throw(DomainError(n, "the number of jumps must be ≥ 0"))
i = zero(n)
while i < n
$fname!(rng)
i += one(n)
end
rng
end

@eval $fname(rng::Xoshiro, n::Integer) = $fname!(copy(rng), n)
end

for (fname, sz) in ((:jump_128, 128), (:jump_192, 192))
local fname! = Symbol(fname, :!)
local see_other = Symbol(fname === :jump_128 ? :jump_192 : :jump_128)
local see_other! = Symbol(see_other, :!)
local seq_pow = 256 - sz
@eval begin
"""
$($fname!)(rng::Xoshiro, [n::Integer=1])
Jump forward, advancing the state equivalent to `2^$($sz)` calls which consume
8 bytes (i.e. a full `UInt64`) each.
If `n > 0` is provided, the state is advanced equivalent to `n * 2^$($sz)` calls; if `n = 0`,
the state remains unchanged.
This can be used to generate `2^$($seq_pow)` non-overlapping subsequences for parallel computations.
See also: [`$($fname)`](@ref), [`$($see_other!)`](@ref)
# Examples
```julia-repl
julia> $($fname!)($($fname!)(Xoshiro(1))) == $($fname!)(Xoshiro(1), 2)
true
```
"""
function $fname! end
end

@eval begin
"""
$($fname)(rng::Xoshiro, [n::Integer=1])
Return a copy of `rng` with the state advanced equivalent to `n * 2^$($sz)` calls which consume
8 bytes (i.e. a full `UInt64`) each; if `n = 0`, the state of the returned copy will be
identical to `rng`.
This can be used to generate `2^$($seq_pow)` non-overlapping subsequences for parallel computations.
See also: [`$($fname!)`](@ref), [`$($see_other)`](@ref)
# Examples
```julia-repl
julia> x = Xoshiro(1);
julia> $($fname)($($fname)(x)) == $($fname)(x, 2)
true
julia> $($fname)(x, 0) == x
true
julia> $($fname)(x, 0) === x
false
```
"""
function $fname end
end
end

## Task local RNG

Expand Down
63 changes: 63 additions & 0 deletions stdlib/Random/test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ using Random
using Random.DSFMT

using Random: Sampler, SamplerRangeFast, SamplerRangeInt, SamplerRangeNDL, MT_CACHE_F, MT_CACHE_I
using Random: jump_128, jump_192, jump_128!, jump_192!

import Future # randjump

Expand Down Expand Up @@ -1105,3 +1106,65 @@ end
@test TaskLocalRNG() == rng3
end
end

# Xoshiro jumps
@testset "Xoshiro jump, basic" begin
x1 = Xoshiro(1)
x2 = Xoshiro(1)

@test jump_128!(jump_128!(x1)) == jump_128!(x1, 2)

xo1 = Xoshiro(0xfff0241072ddab67, 0xc53bc12f4c3f0b4e, 0x56d451780b2dd4ba, 0x50a4aa153d208dd8)
@test rand(jump_128(xo1), UInt64) == 0x87c158da8c35824d
@test rand(jump_192(xo1), UInt64) == 0xcaecd5afdd0847d5

@test rand(jump_128(xo1, 98765), UInt64) == 0xcbec1d5053142608
@test rand(jump_192(xo1, 98765), UInt64) == 0x3b97a94c44d66216

# Throws where appropriate
@test_throws DomainError jump_128(Xoshiro(1), -1)
@test_throws DomainError jump_128!(Xoshiro(1), -1)
@test_throws DomainError jump_192(Xoshiro(1), -1)
@test_throws DomainError jump_192!(Xoshiro(1), -1)

# clean copy when non-mut and no state advance
x = Xoshiro(1)
@test jump_128(x, 0) == x
@test jump_128(x, 0) !== x
@test jump_192(x, 0) == x
@test jump_192(x, 0) !== x

y = Xoshiro(1)
@test jump_128!(x, 0) == y
@test jump_192!(x, 0) == y
end

@testset "Xoshiro jump_128, various seeds" begin
for seed in (0, 1, 0xa0a3f09d0cecd878, 0x7ff8)
x = Xoshiro(seed)
@test jump_128(jump_128(jump_128(x))) == jump_128(x, 3)
x1 = Xoshiro(seed)
@test jump_128!(jump_128!(jump_128!(x1))) == jump_128(x, 3)
jump_128!(x1, 997)
x2 = jump_128!(Xoshiro(seed), 1000)
for T (Float64, UInt64, Int, Char, Bool)
@test rand(x1, T, 5) == rand(x2, T, 5)
@test rand(jump_128!(x1), T, 5) == rand(jump_128!(x2), T, 5)
end
end
end

@testset "Xoshiro jump_192, various seeds" begin
for seed in (0, 1, 0xa0a3f09d0cecd878, 0x7ff8)
x = Xoshiro(seed)
@test jump_192(jump_192(jump_192(x))) == jump_192(x, 3)
x1 = Xoshiro(seed)
@test jump_192!(jump_192!(jump_192!(x1))) == jump_192(x, 3)
jump_192!(x1, 997)
x2 = jump_192!(Xoshiro(seed), 1000)
for T (Float64, UInt64, Int, Char, Bool)
@test rand(x1, T, 5) == rand(x2, T, 5)
@test rand(jump_192!(x1), T, 5) == rand(jump_192!(x2), T, 5)
end
end
end

2 comments on commit 2d93567

@nanosoldier
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Executing the daily package evaluation, I will reply here when finished:

@nanosoldier runtests(isdaily = true)

@nanosoldier
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The package evaluation job you requested has completed - possible new issues were detected.
The full report is available.

Please sign in to comment.