Merge pull request #25047 from JuliaLang/rf/rand/range-52

faster rand!(::MersenneTwister, ::UnitRange)

Author: rfourquet

base/random/generation.jl

@@ -110,9 +110,13 @@ rand(r::AbstractRNG, ::SamplerTrivial{UInt52Raw{UInt64}}) =
 _rand52(r::AbstractRNG, ::Type{Float64}) = reinterpret(UInt64, rand(r, Close1Open2()))
 _rand52(r::AbstractRNG, ::Type{UInt64})  = rand(r, UInt64)
 
-rand(r::AbstractRNG, ::SamplerTrivial{UInt10{UInt16}}) = rand(r, UInt10Raw()) & 0x03ff
-rand(r::AbstractRNG, ::SamplerTrivial{UInt23{UInt32}}) = rand(r, UInt23Raw()) & 0x007fffff
-rand(r::AbstractRNG, ::SamplerTrivial{UInt52{UInt64}}) = rand(r, UInt52Raw()) & 0x000fffffffffffff
+rand(r::AbstractRNG, ::SamplerTrivial{UInt104Raw{UInt128}}) =
+    rand(r, UInt52Raw(UInt128)) << 52 ⊻ rand_inbounds(r, UInt52Raw(UInt128))
+
+rand(r::AbstractRNG, ::SamplerTrivial{UInt10{UInt16}})   = rand(r, UInt10Raw())  & 0x03ff
+rand(r::AbstractRNG, ::SamplerTrivial{UInt23{UInt32}})   = rand(r, UInt23Raw())  & 0x007fffff
+rand(r::AbstractRNG, ::SamplerTrivial{UInt52{UInt64}})   = rand(r, UInt52Raw())  & 0x000fffffffffffff
+rand(r::AbstractRNG, ::SamplerTrivial{UInt104{UInt128}}) = rand(r, UInt104Raw()) & 0x000000ffffffffffffffffffffffffff
 
 rand(r::AbstractRNG, sp::SamplerTrivial{<:UniformBits{T}}) where {T} =
         rand(r, uint_default(sp[])) % T
@@ -143,6 +147,12 @@ end
 # 2) "Default" which tries to use as few entropy bits as possible, at the cost of a
 #    a bigger upfront price associated with the creation of the sampler
 
+#### helper functions
+
+uint_sup(::Type{<:Union{Bool,BitInteger}}) = UInt32
+uint_sup(::Type{<:Union{Int64,UInt64}}) = UInt64
+uint_sup(::Type{<:Union{Int128,UInt128}}) = UInt128
+
 #### Fast
 
 struct SamplerRangeFast{U<:BitUnsigned,T<:Union{BitInteger,Bool}} <: Sampler
@@ -152,44 +162,37 @@ struct SamplerRangeFast{U<:BitUnsigned,T<:Union{BitInteger,Bool}} <: Sampler
     mask::U   # mask generated values before threshold rejection
 end
 
-function SamplerRangeFast(r::AbstractUnitRange{T}) where T<:Union{Base.BitInteger64,Bool}
-    isempty(r) && throw(ArgumentError("range must be non-empty"))
-    m = last(r) % UInt64 - first(r) % UInt64
-    bw = (64 - leading_zeros(m)) % UInt # bit-width
-    mask = (1 % UInt64 << bw) - (1 % UInt64)
-    SamplerRangeFast{UInt64,T}(first(r), bw, m, mask)
-end
+SamplerRangeFast(r::AbstractUnitRange{T}) where T<:Union{Bool,BitInteger} =
+    SamplerRangeFast(r, uint_sup(T))
 
-function SamplerRangeFast(r::AbstractUnitRange{T}) where T<:Union{Int128,UInt128}
+function SamplerRangeFast(r::AbstractUnitRange{T}, ::Type{U}) where {T,U}
     isempty(r) && throw(ArgumentError("range must be non-empty"))
-    m = (last(r)-first(r)) % UInt128
-    bw = (128 - leading_zeros(m)) % UInt # bit-width
-    mask = (1 % UInt128 << bw) - (1 % UInt128)
-    SamplerRangeFast{UInt128,T}(first(r), bw, m, mask)
+    m = (last(r) - first(r)) % U
+    bw = (sizeof(U) << 3 - leading_zeros(m)) % UInt # bit-width
+    mask = (1 % U << bw) - (1 % U)
+    SamplerRangeFast{U,T}(first(r), bw, m, mask)
 end
 
-function rand_lteq(r::AbstractRNG, S, u::U, mask::U) where U<:Integer
-    while true
-        x = rand(r, S) & mask
-        x <= u && return x
-    end
+function rand(rng::AbstractRNG, sp::SamplerRangeFast{UInt32,T}) where T
+    a, bw, m, mask = sp.a, sp.bw, sp.m, sp.mask
+    x = rand(rng, LessThan(m, Masked(mask, uniform(UInt32))))
+    (x + a % UInt32) % T
 end
 
-# helper function, to turn types to values, should be removed once we can do rand(Uniform(UInt))
-rand(rng::AbstractRNG, ::Val{T}) where {T} = rand(rng, T)
-
 function rand(rng::AbstractRNG, sp::SamplerRangeFast{UInt64,T}) where T
     a, bw, m, mask = sp.a, sp.bw, sp.m, sp.mask
-    x = bw <= 52 ? rand_lteq(rng, UInt52Raw(), m, mask) :
-                   rand_lteq(rng, Val(UInt64), m, mask)
+    x = bw <= 52 ? rand(rng, LessThan(m, Masked(mask, UInt52Raw()))) :
+                   rand(rng, LessThan(m, Masked(mask, uniform(UInt64))))
     (x + a % UInt64) % T
 end
 
 function rand(rng::AbstractRNG, sp::SamplerRangeFast{UInt128,T}) where T
     a, bw, m, mask = sp.a, sp.bw, sp.m, sp.mask
-    x = bw <= 52  ? rand_lteq(rng, UInt52Raw(), m % UInt64, mask % UInt64) % UInt128 :
-        bw <= 104 ? rand_lteq(rng, UInt104Raw(), m, mask) :
-                    rand_lteq(rng, Val(UInt128), m, mask)
+    x = bw <= 52  ?
+        rand(rng, LessThan(m % UInt64, Masked(mask % UInt64, UInt52Raw()))) % UInt128 :
+    bw <= 104 ?
+        rand(rng, LessThan(m, Masked(mask, UInt104Raw()))) :
+        rand(rng, LessThan(m, Masked(mask, uniform(UInt128))))
     x % T + a
 end
 
@@ -199,60 +202,70 @@ end
 rem_knuth(a::UInt, b::UInt) = a % (b + (b == 0)) + a * (b == 0)
 rem_knuth(a::T, b::T) where {T<:Unsigned} = b != 0 ? a % b : a
 
-# maximum multiple of k <= 2^bits(T) decremented by one,
-# that is 0xFFFF...FFFF if k = typemax(T) - typemin(T) with intentional underflow
+# maximum multiple of k <= sup decremented by one,
+# that is 0xFFFF...FFFF if k = (typemax(T) - typemin(T)) + 1 and sup == typemax(T) - 1
+# with intentional underflow
 # see http://stackoverflow.com/questions/29182036/integer-arithmetic-add-1-to-uint-max-and-divide-by-n-without-overflow
-maxmultiple(k::T) where {T<:Unsigned} =
-    (div(typemax(T) - k + one(k), k + (k == 0))*k + k - one(k))::T
 
-# serves as rejection threshold
-_maxmultiple(k)  = maxmultiple(k)
+# sup == 0 means typemax(T) + 1
+maxmultiple(k::T, sup::T=zero(T)) where {T<:Unsigned} =
+    (div(sup - k, k + (k == 0))*k + k - one(k))::T
 
-# maximum multiple of k within 1:2^32 or 1:2^64 decremented by one, depending on size
-_maxmultiple(k::UInt64)::UInt64 = k >> 32 != 0 ?
-    maxmultiple(k) :
-    div(0x0000000100000000, k + (k == 0))*k - one(k)
+# similar but sup must not be equal to typemax(T)
+unsafe_maxmultiple(k::T, sup::T) where {T<:Unsigned} =
+    div(sup, k + (k == 0))*k - one(k)
 
 struct SamplerRangeInt{T<:Union{Bool,Integer},U<:Unsigned} <: Sampler
-    a::T   # first element of the range
-    k::U   # range length or zero for full range
-    u::U   # rejection threshold
+    a::T      # first element of the range
+    bw::Int   # bit width
+    k::U      # range length or zero for full range
+    u::U      # rejection threshold
 end
 
-function SamplerRangeInt(a::T, diff::U) where {T<:Union{Bool,Integer},U<:Unsigned}
-    k = diff+one(U)
-    SamplerRangeInt{T,U}(a, k, _maxmultiple(k)) # overflow ok
-end
 
-uint_sup(::Type{<:Union{Bool,BitInteger}}) = UInt32
-uint_sup(::Type{<:Union{Int64,UInt64}}) = UInt64
-uint_sup(::Type{<:Union{Int128,UInt128}}) = UInt128
+SamplerRangeInt(r::AbstractUnitRange{T}) where T<:Union{Bool,BitInteger} =
+    SamplerRangeInt(r, uint_sup(T))
 
-function SamplerRangeInt(r::AbstractUnitRange{T}) where T<:Union{Bool,BitInteger}
+function SamplerRangeInt(r::AbstractUnitRange{T}, ::Type{U}) where {T,U}
     isempty(r) && throw(ArgumentError("range must be non-empty"))
-    SamplerRangeInt(first(r), (last(r) - first(r)) % uint_sup(T))
+    a = first(r)
+    m = (last(r) - first(r)) % U
+    k = m + one(U)
+    bw = (sizeof(U) << 3 - leading_zeros(m)) % Int
+    mult = if U === UInt32
+        maxmultiple(k)
+    elseif U === UInt64
+        bw <= 52 ? unsafe_maxmultiple(k, one(UInt64) << 52) :
+                   maxmultiple(k)
+    else # U === UInt128
+        bw <= 52  ? unsafe_maxmultiple(k, one(UInt128) << 52) :
+        bw <= 104 ? unsafe_maxmultiple(k, one(UInt128) << 104) :
+                    maxmultiple(k)
+    end
+
+    SamplerRangeInt{T,U}(a, bw, k, mult) # overflow ok
 end
 
-Sampler(::AbstractRNG, r::AbstractUnitRange{T}, ::Repetition) where {T<:Union{Bool,BitInteger}} =
-    SamplerRangeInt(r)
+Sampler(::AbstractRNG, r::AbstractUnitRange{T},
+        ::Repetition) where {T<:Union{Bool,BitInteger}} = SamplerRangeInt(r)
+
+rand(rng::AbstractRNG, sp::SamplerRangeInt{T,UInt32}) where {T<:Union{Bool,BitInteger}} =
+    (unsigned(sp.a) + rem_knuth(rand(rng, LessThan(sp.u, uniform(UInt32))), sp.k)) % T
 
-function rand_lteq(rng::AbstractRNG, u::T)::T where T
-    while true
-        x = rand(rng, T)
-        x <= u && return x
-    end
-end
 
-# this function uses 32 bit entropy for small ranges of length <= typemax(UInt32) + 1
+# this function uses 52 bit entropy for small ranges of length <= 2^52
 function rand(rng::AbstractRNG, sp::SamplerRangeInt{T,UInt64}) where T<:BitInteger
-    x::UInt64 = (sp.k - 1) >> 32 == 0 ?
-        rand_lteq(rng, sp.u % UInt32) % UInt64 :
-        rand_lteq(rng, sp.u)
+    x = sp.bw <= 52 ? rand(rng, LessThan(sp.u, UInt52())) :
+                      rand(rng, LessThan(sp.u, uniform(UInt64)))
     return ((sp.a % UInt64) + rem_knuth(x, sp.k)) % T
 end
 
-rand(rng::AbstractRNG, sp::SamplerRangeInt{T,U}) where {T<:Union{Bool,BitInteger},U} =
-    (unsigned(sp.a) + rem_knuth(rand_lteq(rng, sp.u), sp.k)) % T
+function rand(rng::AbstractRNG, sp::SamplerRangeInt{T,UInt128}) where T<:BitInteger
+    x = sp.bw <= 52  ? rand(rng, LessThan(sp.u, UInt52(UInt128))) :
+        sp.bw <= 104 ? rand(rng, LessThan(sp.u, UInt104(UInt128))) :
+                       rand(rng, LessThan(sp.u, uniform(UInt128)))
+    return ((sp.a % UInt128) + rem_knuth(x, sp.k)) % T
+end
 
 
 ### BigInt

base/random/misc.jl

@@ -141,17 +141,10 @@ large.) Technically, this process is known as "Bernoulli sampling" of `A`.
 randsubseq(A::AbstractArray, p::Real) = randsubseq(GLOBAL_RNG, A, p)
 
 
-## rand_lt (helper function)
-
-"Return a random `Int` (masked with `mask`) in ``[0, n)``, when `n <= 2^52`."
-@inline function rand_lt(r::AbstractRNG, n::Int, mask::Int=nextpow2(n)-1)
-    # this duplicates the functionality of rand(1:n), to optimize this special case
-    while true
-        x = rand(r, UInt52Raw(Int)) & mask
-        x < n && return x
-    end
-end
+## rand Less Than Masked 52 bits (helper function)
 
+"Return a sampler generating a random `Int` (masked with `mask`) in ``[0, n)``, when `n <= 2^52`."
+ltm52(n::Int, mask::Int=nextpow2(n)-1) = LessThan(n-1, Masked(mask, UInt52Raw(Int)))
 
 ## shuffle & shuffle!
 
@@ -191,7 +184,7 @@ function shuffle!(r::AbstractRNG, a::AbstractArray)
     mask = nextpow2(n) - 1
     for i = n:-1:2
         (mask >> 1) == i && (mask >>= 1)
-        j = 1 + rand_lt(r, i, mask)
+        j = 1 + rand(r, ltm52(i, mask))
         a[i], a[j] = a[j], a[i]
     end
     return a
@@ -277,7 +270,7 @@ function randperm!(r::AbstractRNG, a::Array{<:Integer})
     a[1] = 1
     mask = 3
     @inbounds for i = 2:n
-        j = 1 + rand_lt(r, i, mask)
+        j = 1 + rand(r, ltm52(i, mask))
         if i != j # a[i] is uninitialized (and could be #undef)
             a[i] = a[j]
         end
@@ -339,7 +332,7 @@ function randcycle!(r::AbstractRNG, a::Array{<:Integer})
     a[1] = 1
     mask = 3
     @inbounds for i = 2:n
-        j = 1 + rand_lt(r, i-1, mask)
+        j = 1 + rand(r, ltm52(i-1, mask))
         a[i] = a[j]
         a[j] = i
         i == 1+mask && (mask = 2mask + 1)

base/random/random.jl

@@ -106,7 +106,7 @@ Sampler(rng::AbstractRNG, sp::Sampler, ::Repetition) =
 Sampler(rng::AbstractRNG, X) = Sampler(rng, X, Val(Inf))
 Sampler(rng::AbstractRNG, ::Type{X}) where {X} = Sampler(rng, X, Val(Inf))
 
-#### pre-defined useful Sampler subtypes
+#### pre-defined useful Sampler types
 
 # default fall-back for types
 struct SamplerType{T} <: Sampler end
@@ -139,6 +139,38 @@ struct SamplerTag{T,S} <: Sampler
 end
 
 
+#### helper samplers
+
+##### Adapter to generate a randome value in [0, n]
+
+struct LessThan{T<:Integer,S} <: Sampler
+    sup::T
+    s::S    # the scalar specification/sampler to feed to rand
+end
+
+function rand(rng::AbstractRNG, sp::LessThan)
+    while true
+        x = rand(rng, sp.s)
+        x <= sp.sup && return x
+    end
+end
+
+struct Masked{T<:Integer,S} <: Sampler
+    mask::T
+    s::S
+end
+
+rand(rng::AbstractRNG, sp::Masked) = rand(rng, sp.s) & sp.mask
+
+##### Uniform
+
+struct UniformT{T} <: Sampler end
+
+uniform(::Type{T}) where {T} = UniformT{T}()
+
+rand(rng::AbstractRNG, ::UniformT{T}) where {T} = rand(rng, T)
+
+
 ### machinery for generation with Sampler
 
 # This describes how to generate random scalars or arrays, by generating a Sampler

test/random.jl

@@ -102,12 +102,13 @@ if sizeof(Int32) < sizeof(Int)
     local r = rand(Int32(-1):typemax(Int32))
     @test typeof(r) == Int32
     @test -1 <= r <= typemax(Int32)
-    @test all([div(0x00010000000000000000,k)*k - 1 == Base.Random.RangeGenerator(map(UInt64,1:k)).u for k in 13 .+ Int64(2).^(32:62)])
-    @test all([div(0x00010000000000000000,k)*k - 1 == Base.Random.RangeGenerator(map(Int64,1:k)).u for k in 13 .+ Int64(2).^(32:61)])
+    for U = (Int64, UInt64)
+        @test all(div(one(UInt128) << 52, k)*k - 1 == SamplerRangeInt(map(U, 1:k)).u
+                  for k in 13 .+ Int64(2).^(32:51))
+        @test all(div(one(UInt128) << 64, k)*k - 1 == SamplerRangeInt(map(U, 1:k)).u
+                  for k in 13 .+ Int64(2).^(52:62))
+    end
 
-    @test Base.Random._maxmultiple(0x000100000000) === 0xffffffffffffffff
-    @test Base.Random._maxmultiple(0x0000FFFFFFFF) === 0x00000000fffffffe
-    @test Base.Random._maxmultiple(0x000000000000) === 0xffffffffffffffff
 end
 
 # BigInt specific
@@ -224,8 +225,11 @@ guardsrand() do
     @test r == rand(map(UInt64, 97:122))
 end
 
-@test all([div(0x000100000000,k)*k - 1 == Base.Random.RangeGenerator(map(UInt64,1:k)).u for k in 13 .+ Int64(2).^(1:30)])
-@test all([div(0x000100000000,k)*k - 1 == Base.Random.RangeGenerator(map(Int64,1:k)).u for k in 13 .+ Int64(2).^(1:30)])
+for U in (Int64, UInt64)
+    @test all(div(one(UInt64) << 52, k)*k - 1 == SamplerRangeInt(map(U, 1:k)).u
+              for k in 13 .+ Int64(2).^(1:30))
+end
+
 
 import Base.Random: uuid1, uuid4, UUID, uuid_version