Skip to content

Commit 2d93567

Browse files
Add jump functions (jump, long_jump) for Xoshiro (#47743)
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.
1 parent ed891d6 commit 2d93567

File tree

2 files changed

+165
-0
lines changed

2 files changed

+165
-0
lines changed

stdlib/Random/src/Xoshiro.jl

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,108 @@ end
7070

7171
rng_native_52(::Xoshiro) = UInt64
7272

73+
# Jump functions from: https://xoshiro.di.unimi.it/xoshiro256plusplus.c
74+
75+
for (fname, JUMP) in ((:jump_128, (0x180ec6d33cfd0aba, 0xd5a61266f0c9392c, 0xa9582618e03fc9aa, 0x39abdc4529b1661c)),
76+
(:jump_192, (0x76e15d3efefdcbbf, 0xc5004e441c522fb3, 0x77710069854ee241, 0x39109bb02acbe635)))
77+
local fname! = Symbol(fname, :!)
78+
@eval function $fname!(rng::Xoshiro)
79+
_s0 = 0x0000000000000000
80+
_s1 = 0x0000000000000000
81+
_s2 = 0x0000000000000000
82+
_s3 = 0x0000000000000000
83+
s0, s1, s2, s3 = rng.s0, rng.s1, rng.s2, rng.s3
84+
for j in $JUMP
85+
for b in 0x0000000000000000:0x000000000000003f
86+
if (j & 0x0000000000000001 << b) != 0
87+
_s0 ⊻= s0
88+
_s1 ⊻= s1
89+
_s2 ⊻= s2
90+
_s3 ⊻= s3
91+
end
92+
t = s1 << 17
93+
s2 = xor(s2, s0)
94+
s3 = xor(s3, s1)
95+
s1 = xor(s1, s2)
96+
s0 = xor(s0, s3)
97+
s2 = xor(s2, t)
98+
s3 = s3 << 45 | s3 >> 19
99+
end
100+
end
101+
setstate!(rng, (_s0, _s1, _s2, _s3, nothing))
102+
end
103+
@eval $fname(rng::Xoshiro) = $fname!(copy(rng))
104+
105+
@eval function $fname!(rng::Xoshiro, n::Integer)
106+
n < 0 && throw(DomainError(n, "the number of jumps must be ≥ 0"))
107+
i = zero(n)
108+
while i < n
109+
$fname!(rng)
110+
i += one(n)
111+
end
112+
rng
113+
end
114+
115+
@eval $fname(rng::Xoshiro, n::Integer) = $fname!(copy(rng), n)
116+
end
117+
118+
for (fname, sz) in ((:jump_128, 128), (:jump_192, 192))
119+
local fname! = Symbol(fname, :!)
120+
local see_other = Symbol(fname === :jump_128 ? :jump_192 : :jump_128)
121+
local see_other! = Symbol(see_other, :!)
122+
local seq_pow = 256 - sz
123+
@eval begin
124+
"""
125+
$($fname!)(rng::Xoshiro, [n::Integer=1])
126+
127+
Jump forward, advancing the state equivalent to `2^$($sz)` calls which consume
128+
8 bytes (i.e. a full `UInt64`) each.
129+
130+
If `n > 0` is provided, the state is advanced equivalent to `n * 2^$($sz)` calls; if `n = 0`,
131+
the state remains unchanged.
132+
133+
This can be used to generate `2^$($seq_pow)` non-overlapping subsequences for parallel computations.
134+
135+
See also: [`$($fname)`](@ref), [`$($see_other!)`](@ref)
136+
137+
# Examples
138+
```julia-repl
139+
julia> $($fname!)($($fname!)(Xoshiro(1))) == $($fname!)(Xoshiro(1), 2)
140+
true
141+
```
142+
"""
143+
function $fname! end
144+
end
145+
146+
@eval begin
147+
"""
148+
$($fname)(rng::Xoshiro, [n::Integer=1])
149+
150+
Return a copy of `rng` with the state advanced equivalent to `n * 2^$($sz)` calls which consume
151+
8 bytes (i.e. a full `UInt64`) each; if `n = 0`, the state of the returned copy will be
152+
identical to `rng`.
153+
154+
This can be used to generate `2^$($seq_pow)` non-overlapping subsequences for parallel computations.
155+
156+
See also: [`$($fname!)`](@ref), [`$($see_other)`](@ref)
157+
158+
# Examples
159+
```julia-repl
160+
julia> x = Xoshiro(1);
161+
162+
julia> $($fname)($($fname)(x)) == $($fname)(x, 2)
163+
true
164+
165+
julia> $($fname)(x, 0) == x
166+
true
167+
168+
julia> $($fname)(x, 0) === x
169+
false
170+
```
171+
"""
172+
function $fname end
173+
end
174+
end
73175

74176
## Task local RNG
75177

stdlib/Random/test/runtests.jl

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ using Random
1111
using Random.DSFMT
1212

1313
using Random: Sampler, SamplerRangeFast, SamplerRangeInt, SamplerRangeNDL, MT_CACHE_F, MT_CACHE_I
14+
using Random: jump_128, jump_192, jump_128!, jump_192!
1415

1516
import Future # randjump
1617

@@ -1105,3 +1106,65 @@ end
11051106
@test TaskLocalRNG() == rng3
11061107
end
11071108
end
1109+
1110+
# Xoshiro jumps
1111+
@testset "Xoshiro jump, basic" begin
1112+
x1 = Xoshiro(1)
1113+
x2 = Xoshiro(1)
1114+
1115+
@test jump_128!(jump_128!(x1)) == jump_128!(x1, 2)
1116+
1117+
xo1 = Xoshiro(0xfff0241072ddab67, 0xc53bc12f4c3f0b4e, 0x56d451780b2dd4ba, 0x50a4aa153d208dd8)
1118+
@test rand(jump_128(xo1), UInt64) == 0x87c158da8c35824d
1119+
@test rand(jump_192(xo1), UInt64) == 0xcaecd5afdd0847d5
1120+
1121+
@test rand(jump_128(xo1, 98765), UInt64) == 0xcbec1d5053142608
1122+
@test rand(jump_192(xo1, 98765), UInt64) == 0x3b97a94c44d66216
1123+
1124+
# Throws where appropriate
1125+
@test_throws DomainError jump_128(Xoshiro(1), -1)
1126+
@test_throws DomainError jump_128!(Xoshiro(1), -1)
1127+
@test_throws DomainError jump_192(Xoshiro(1), -1)
1128+
@test_throws DomainError jump_192!(Xoshiro(1), -1)
1129+
1130+
# clean copy when non-mut and no state advance
1131+
x = Xoshiro(1)
1132+
@test jump_128(x, 0) == x
1133+
@test jump_128(x, 0) !== x
1134+
@test jump_192(x, 0) == x
1135+
@test jump_192(x, 0) !== x
1136+
1137+
y = Xoshiro(1)
1138+
@test jump_128!(x, 0) == y
1139+
@test jump_192!(x, 0) == y
1140+
end
1141+
1142+
@testset "Xoshiro jump_128, various seeds" begin
1143+
for seed in (0, 1, 0xa0a3f09d0cecd878, 0x7ff8)
1144+
x = Xoshiro(seed)
1145+
@test jump_128(jump_128(jump_128(x))) == jump_128(x, 3)
1146+
x1 = Xoshiro(seed)
1147+
@test jump_128!(jump_128!(jump_128!(x1))) == jump_128(x, 3)
1148+
jump_128!(x1, 997)
1149+
x2 = jump_128!(Xoshiro(seed), 1000)
1150+
for T (Float64, UInt64, Int, Char, Bool)
1151+
@test rand(x1, T, 5) == rand(x2, T, 5)
1152+
@test rand(jump_128!(x1), T, 5) == rand(jump_128!(x2), T, 5)
1153+
end
1154+
end
1155+
end
1156+
1157+
@testset "Xoshiro jump_192, various seeds" begin
1158+
for seed in (0, 1, 0xa0a3f09d0cecd878, 0x7ff8)
1159+
x = Xoshiro(seed)
1160+
@test jump_192(jump_192(jump_192(x))) == jump_192(x, 3)
1161+
x1 = Xoshiro(seed)
1162+
@test jump_192!(jump_192!(jump_192!(x1))) == jump_192(x, 3)
1163+
jump_192!(x1, 997)
1164+
x2 = jump_192!(Xoshiro(seed), 1000)
1165+
for T (Float64, UInt64, Int, Char, Bool)
1166+
@test rand(x1, T, 5) == rand(x2, T, 5)
1167+
@test rand(jump_192!(x1), T, 5) == rand(jump_192!(x2), T, 5)
1168+
end
1169+
end
1170+
end

0 commit comments

Comments
 (0)