Merge pull request #24885 from JuliaLang/rf/rand/ptr

fix some FIXMEs in random, and take 2 to #9174

Author: rfourquet

base/random/RNGs.jl

@@ -143,7 +143,7 @@ mt_setempty!(r::MersenneTwister) = r.idx = MTCacheLength
 mt_pop!(r::MersenneTwister) = @inbounds return r.vals[r.idx+=1]
 
 function gen_rand(r::MersenneTwister)
-    Base.@gc_preserve r dsfmt_fill_array_close1_open2!(r.state, pointer(r.vals), length(r.vals))
+    @gc_preserve r dsfmt_fill_array_close1_open2!(r.state, pointer(r.vals), length(r.vals))
     mt_setfull!(r)
 end
 
@@ -264,12 +264,14 @@ rand(r::MersenneTwister, ::SamplerType{Int128}) = rand(r, UInt128) % Int128
 
 #### arrays of floats
 
-function rand_AbstractArray_Float64!(r::MersenneTwister, A::AbstractArray{Float64},
-                                     region::AbstractUnitRange, I::FloatInterval_64)
-    # region should be a subset of linearindices(A)
+##### AbstractArray
+
+function rand!(r::MersenneTwister, A::AbstractArray{Float64},
+               I::SamplerTrivial{<:FloatInterval_64})
+    region = linearindices(A)
     # what follows is equivalent to this simple loop but more efficient:
     # for i=region
-    #     @inbounds A[i] = rand(r, I)
+    #     @inbounds A[i] = rand(r, I[])
     # end
     m = Base.checked_sub(first(region), 1)
     n = last(region)
@@ -281,53 +283,101 @@ function rand_AbstractArray_Float64!(r::MersenneTwister, A::AbstractArray{Float6
         end
         m2 = min(n, m+s)
         for i=m+1:m2
-            @inbounds A[i] = rand_inbounds(r, I)
+            @inbounds A[i] = rand_inbounds(r, I[])
         end
         m = m2
     end
     A
 end
 
-rand!(r::MersenneTwister, A::AbstractArray{Float64}, I::SamplerTrivial{<:FloatInterval_64}) =
-    rand_AbstractArray_Float64!(r, A, linearindices(A), I[])
+
+##### Array : internal functions
+
+# internal array-like type to circumevent the lack of flexibility with reinterpret
+struct UnsafeView{T} <: DenseArray{T,1}
+    ptr::Ptr{T}
+    len::Int
+end
+
+Base.length(a::UnsafeView) = a.len
+Base.getindex(a::UnsafeView, i::Int) = unsafe_load(a.ptr, i)
+Base.setindex!(a::UnsafeView, x, i::Int) = unsafe_store!(a.ptr, x, i)
+Base.pointer(a::UnsafeView) = a.ptr
+Base.size(a::UnsafeView) = (a.len,)
+
+# this is essentially equivalent to rand!(r, ::AbstractArray{Float64}, I) above, but due to
+# optimizations which can't be done currently when working with pointers, we have to re-order
+# manually the computation flow to get the performance
+# (see https://discourse.julialang.org/t/unsafe-store-sometimes-slower-than-arrays-setindex)
+function _rand_max383!(r::MersenneTwister, A::UnsafeView{Float64}, I::FloatInterval_64)
+    n = length(A)
+    @assert n <= dsfmt_get_min_array_size()+1 # == 383
+    mt_avail(r) == 0 && gen_rand(r)
+    # from now on, at most one call to gen_rand(r) will be necessary
+    m = min(n, mt_avail(r))
+    @gc_preserve r unsafe_copy!(A.ptr, pointer(r.vals, r.idx+1), m)
+    if m == n
+        r.idx += m
+    else # m < n
+        gen_rand(r)
+        @gc_preserve r unsafe_copy!(A.ptr+m*sizeof(Float64), pointer(r.vals), n-m)
+        r.idx = n-m
+    end
+    if I isa CloseOpen
+        for i=1:n
+            A[i] -= 1.0
+        end
+    end
+    A
+end
+
 
 fill_array!(s::DSFMT_state, A::Ptr{Float64}, n::Int, ::CloseOpen_64) =
     dsfmt_fill_array_close_open!(s, A, n)
 
 fill_array!(s::DSFMT_state, A::Ptr{Float64}, n::Int, ::Close1Open2_64) =
     dsfmt_fill_array_close1_open2!(s, A, n)
 
-function _rand!(r::MersenneTwister, A::Array{Float64}, n::Int,
-               I::FloatInterval_64)
+
+function rand!(r::MersenneTwister, A::UnsafeView{Float64},
+               I::SamplerTrivial{<:FloatInterval_64})
     # depending on the alignment of A, the data written by fill_array! may have
     # to be left-shifted by up to 15 bytes (cf. unsafe_copy! below) for
     # reproducibility purposes;
     # so, even for well aligned arrays, fill_array! is used to generate only
     # the n-2 first values (or n-3 if n is odd), and the remaining values are
     # generated by the scalar version of rand
-    n > length(A) && throw(BoundsError(A, n))
+    n = length(A)
     n2 = (n-2) ÷ 2 * 2
-    if n2 < dsfmt_get_min_array_size()
-        rand_AbstractArray_Float64!(r, A, 1:n, I)
+    n2 < dsfmt_get_min_array_size() && return _rand_max383!(r, A, I[])
+
+    pA = A.ptr
+    align = Csize_t(pA) % 16
+    if align > 0
+        pA2 = pA + 16 - align
+        fill_array!(r.state, pA2, n2, I[]) # generate the data in-place, but shifted
+        unsafe_copy!(pA, pA2, n2) # move the data to the beginning of the array
     else
-        pA = pointer(A)
-        align = Csize_t(pA) % 16
-        Base.@gc_preserve A if align > 0
-            pA2 = pA + 16 - align
-            fill_array!(r.state, pA2, n2, I) # generate the data in-place, but shifted
-            unsafe_copy!(pA, pA2, n2) # move the data to the beginning of the array
-        else
-            fill_array!(r.state, pA, n2, I)
-        end
-        for i=n2+1:n
-            @inbounds A[i] = rand(r, I)
-        end
+        fill_array!(r.state, pA, n2, I[])
     end
+    for i=n2+1:n
+        A[i] = rand(r, I[])
+    end
+    A
+end
+
+# fills up A reinterpreted as an array of Float64 with n64 values
+function _rand!(r::MersenneTwister, A::Array{T}, n64::Int, I::FloatInterval_64) where T
+    # n64 is the length in terms of `Float64` of the target
+    @assert sizeof(Float64)*n64 <= sizeof(T)*length(A) && isbits(T)
+    @gc_preserve A rand!(r, UnsafeView{Float64}(pointer(A), n64), SamplerTrivial(I))
     A
 end
 
-rand!(r::MersenneTwister, A::Array{Float64}, sp::SamplerTrivial{<:FloatInterval_64}) =
-    _rand!(r, A, length(A), sp[])
+##### Array: Float64, Float16, Float32
+
+rand!(r::MersenneTwister, A::Array{Float64}, I::SamplerTrivial{<:FloatInterval_64}) =
+    _rand!(r, A, length(A), I[])
 
 mask128(u::UInt128, ::Type{Float16}) =
     (u & 0x03ff03ff03ff03ff03ff03ff03ff03ff) | 0x3c003c003c003c003c003c003c003c00
@@ -339,27 +389,27 @@ for T in (Float16, Float32)
     @eval function rand!(r::MersenneTwister, A::Array{$T}, ::SamplerTrivial{Close1Open2{$T}})
         n = length(A)
         n128 = n * sizeof($T) ÷ 16
-        Base.@gc_preserve A _rand!(r, unsafe_wrap(Array, convert(Ptr{Float64}, pointer(A)), 2*n128),
-                                   2*n128, Close1Open2())
-        # FIXME: This code is completely invalid!!!
-        A128 = unsafe_wrap(Array, convert(Ptr{UInt128}, pointer(A)), n128)
-        @inbounds for i in 1:n128
-            u = A128[i]
-            u ⊻= u << 26
-            # at this point, the 64 low bits of u, "k" being the k-th bit of A128[i] and "+"
-            # the bit xor, are:
-            # [..., 58+32,..., 53+27, 52+26, ..., 33+7, 32+6, ..., 27+1, 26, ..., 1]
-            # the bits needing to be random are
-            # [1:10, 17:26, 33:42, 49:58] (for Float16)
-            # [1:23, 33:55] (for Float32)
-            # this is obviously satisfied on the 32 low bits side, and on the high side,
-            # the entropy comes from bits 33:52 of A128[i] and then from bits 27:32
-            # (which are discarded on the low side)
-            # this is similar for the 64 high bits of u
-            A128[i] = mask128(u, $T)
+        _rand!(r, A, 2*n128, Close1Open2())
+        @gc_preserve A begin
+            A128 = UnsafeView{UInt128}(pointer(A), n128)
+            for i in 1:n128
+                u = A128[i]
+                u ⊻= u << 26
+                # at this point, the 64 low bits of u, "k" being the k-th bit of A128[i] and "+"
+                # the bit xor, are:
+                # [..., 58+32,..., 53+27, 52+26, ..., 33+7, 32+6, ..., 27+1, 26, ..., 1]
+                # the bits needing to be random are
+                # [1:10, 17:26, 33:42, 49:58] (for Float16)
+                # [1:23, 33:55] (for Float32)
+                # this is obviously satisfied on the 32 low bits side, and on the high side,
+                # the entropy comes from bits 33:52 of A128[i] and then from bits 27:32
+                # (which are discarded on the low side)
+                # this is similar for the 64 high bits of u
+                A128[i] = mask128(u, $T)
+            end
         end
         for i in 16*n128÷sizeof($T)+1:n
-            @inbounds A[i] = rand(r, $T) + oneunit($T)
+            @inbounds A[i] = rand(r, $T) + one($T)
         end
         A
     end
@@ -376,16 +426,14 @@ end
 
 #### arrays of integers
 
-function rand!(r::MersenneTwister, A::Array{UInt128}, ::SamplerType{UInt128})
+function rand!(r::MersenneTwister, A::UnsafeView{UInt128}, ::SamplerType{UInt128})
     n::Int=length(A)
-    # FIXME: This code is completely invalid!!!
-    Af = unsafe_wrap(Array, convert(Ptr{Float64}, pointer(A)), 2n)
     i = n
     while true
-        _rand!(r, Af, 2i, Close1Open2())
+        rand!(r, UnsafeView{Float64}(A.ptr, 2i), Close1Open2())
         n < 5 && break
         i = 0
-        @inbounds while n-i >= 5
+        while n-i >= 5
             u = A[i+=1]
             A[n]    ⊻= u << 48
             A[n-=1] ⊻= u << 36
@@ -397,19 +445,22 @@ function rand!(r::MersenneTwister, A::Array{UInt128}, ::SamplerType{UInt128})
     if n > 0
         u = rand(r, UInt2x52Raw())
         for i = 1:n
-            @inbounds A[i] ⊻= u << (12*i)
+            A[i] ⊻= u << (12*i)
         end
     end
     A
 end
 
 for T in BitInteger_types
-    T === UInt128 && continue
-    @eval function rand!(r::MersenneTwister, A::Array{$T}, ::SamplerType{$T})
+    @eval rand!(r::MersenneTwister, A::Array{$T}, sp::SamplerType{$T}) =
+        (@gc_preserve A rand!(r, UnsafeView(pointer(A), length(A)), sp); A)
+
+    T == UInt128 && continue
+
+    @eval function rand!(r::MersenneTwister, A::UnsafeView{$T}, ::SamplerType{$T})
         n = length(A)
         n128 = n * sizeof($T) ÷ 16
-        # FIXME: This code is completely invalid!!!
-        rand!(r, unsafe_wrap(Array, convert(Ptr{UInt128}, pointer(A)), n128))
+        rand!(r, UnsafeView{UInt128}(pointer(A), n128))
         for i = 16*n128÷sizeof($T)+1:n
             @inbounds A[i] = rand(r, $T)
         end

base/random/generation.jl

@@ -278,17 +278,19 @@ end
 
 function rand(rng::AbstractRNG, sp::SamplerBigInt)
     x = MPZ.realloc2(sp.nlimbsmax*8*sizeof(Limb))
-    limbs = unsafe_wrap(Array, x.d, sp.nlimbs)
-    while true
-        rand!(rng, limbs)
-        @inbounds limbs[end] &= sp.mask
-        MPZ.mpn_cmp(x, sp.m, sp.nlimbs) <= 0 && break
-    end
-    # adjust x.size (normally done by mpz_limbs_finish, in GMP version >= 6)
-    x.size = sp.nlimbs
-    while x.size > 0
-        @inbounds limbs[x.size] != 0 && break
-        x.size -= 1
+    @gc_preserve x begin
+        limbs = UnsafeView(x.d, sp.nlimbs)
+        while true
+            rand!(rng, limbs)
+            limbs[end] &= sp.mask
+            MPZ.mpn_cmp(x, sp.m, sp.nlimbs) <= 0 && break
+        end
+        # adjust x.size (normally done by mpz_limbs_finish, in GMP version >= 6)
+        x.size = sp.nlimbs
+        while x.size > 0
+            limbs[x.size] != 0 && break
+            x.size -= 1
+        end
     end
     MPZ.add!(x, sp.a)
 end

base/random/random.jl

@@ -4,7 +4,8 @@ module Random
 
 using Base.dSFMT
 using Base.GMP: Limb, MPZ
-using Base: BitInteger, BitInteger_types, BitUnsigned
+using Base: BitInteger, BitInteger_types, BitUnsigned, @gc_preserve
+
 import Base: copymutable, copy, copy!, ==, hash
 
 export srand,

test/random.jl

@@ -286,7 +286,8 @@ let mt = MersenneTwister(0)
     @test rand!(mt, AF64)[end] == 0.957735065345398
     @test rand!(mt, AF64)[end] == 0.6492481059865669
     resize!(AF64, 2*length(mt.vals))
-    @test Base.Random.rand_AbstractArray_Float64!(mt, AF64, linearindices(AF64), Base.Random.CloseOpen())[end]  == 0.432757268470779
+    @test invoke(rand!, Tuple{MersenneTwister,AbstractArray{Float64},Base.Random.SamplerTrivial{Base.Random.CloseOpen_64}},
+                 mt, AF64, Base.Random.SamplerTrivial(Base.Random.CloseOpen()))[end]  == 0.432757268470779
 end
 
 # Issue #9037