Use RandomNumbers.jl to provide rand() for other types.

Manifest.toml

@@ -190,6 +190,12 @@ uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 deps = ["Serialization"]
 uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
+[[RandomNumbers]]
+deps = ["Random", "Requires"]
+git-tree-sha1 = "441e6fc35597524ada7f85e13df1f4e10137d16f"
+uuid = "e6cf234a-135c-5ec9-84dd-332b85af5143"
+version = "1.4.0"
+
 [[Reexport]]
 git-tree-sha1 = "57d8440b0c7d98fc4f889e478e80f268d534c9d5"
 uuid = "189a3867-3050-52da-a836-e630ba90ab69"

Project.toml

@@ -21,6 +21,7 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Memoize = "c03570c3-d221-55d1-a50c-7939bbd78826"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+RandomNumbers = "e6cf234a-135c-5ec9-84dd-332b85af5143"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

src/device/intrinsics/random.jl

@@ -1,6 +1,7 @@
 ## random number generation
 
 using Random
+import RandomNumbers
 
 
 # helpers
@@ -11,7 +12,7 @@ global_index() = (threadIdx().x, threadIdx().y, threadIdx().z,
 
 # global state
 
-struct ThreadLocalRNG <: AbstractRNG
+struct ThreadLocalXorshift32 <: RandomNumbers.AbstractRNG{UInt32}
     vals::CuDeviceArray{UInt32, 6, AS.Generic}
 end
 
@@ -37,13 +38,13 @@ end
     CuDeviceArray(dims, ptr)
 end
 
-@device_override Random.default_rng() = ThreadLocalRNG(global_random_state())
+@device_override Random.default_rng() = ThreadLocalXorshift32(global_random_state())
 
 @device_override Random.make_seed() = clock(UInt32)
 
-function Random.seed!(rng::ThreadLocalRNG, seed::Integer)
+function Random.seed!(rng::ThreadLocalXorshift32, seed::Integer)
     index = global_index()
-    rng.vals[index...] = seed
+    rng.vals[index...] = seed % UInt32
     return
 end
 
@@ -57,29 +58,14 @@ function xorshift(x::UInt32)::UInt32
     return x
 end
 
-function get_thread_word(rng::ThreadLocalRNG)
+function Random.rand(rng::ThreadLocalXorshift32, ::Type{UInt32})
     # NOTE: we add the current linear index to the local state, to make sure threads get
     #       different random numbers when unseeded (initial state = 0 for all threads)
     index = global_index()
     offset = LinearIndices(rng.vals)[index...]
     state = rng.vals[index...] + UInt32(offset)
 
-    new_state = generate_next_state(state)
+    new_state = xorshift(state)
     rng.vals[index...] = new_state
-
-    return new_state    # FIXME: return old state?
-end
-
-function generate_next_state(state::UInt32)
-    new_val = xorshift(state)
-    return UInt32(new_val)
-end
-
-# TODO: support for more types (can we reuse more of the Random standard library?)
-#       see RandomNumbers.jl
-
-function Random.rand(rng::ThreadLocalRNG, ::Type{Float32})
-    word = get_thread_word(rng)
-    res = (word >> 9) | reinterpret(UInt32, 1f0)
-    return reinterpret(Float32, res) - 1.0f0
+    return new_state
 end

test/device/intrinsics.jl

@@ -1197,15 +1197,20 @@ n = 256
         return nothing
     end
 
-    a = CUDA.zeros(Float32, n)
-    b = CUDA.zeros(Float32, n)
+    @testset for T in (Int32, UInt32, Int64, UInt64, Int128, UInt128,
+                       Float32, Float64)
+        a = CUDA.zeros(T, n)
+        b = CUDA.zeros(T, n)
 
-    @cuda threads=n kernel(a, b)
+        @cuda threads=n kernel(a, b)
 
-    # FIXME(?): all these tests have a (very small) chance to fail, but that's somewhat inherent to rand() without a seed
-    @test allunique(Array(a))
-    @test allunique(Array(b))
-    @test Array(a) != Array(b)
+        @test all(Array(a) .!= Array(b))
+
+        if T == Float64
+            @test allunique(Array(a))
+            @test allunique(Array(b))
+        end
+    end
 end
 
 @testset "custom seed" begin
@@ -1216,13 +1221,16 @@ end
         return nothing
     end
 
-    a = CUDA.zeros(Float32, n)
-    b = CUDA.zeros(Float32, n)
+    @testset for T in (Int32, UInt32, Int64, UInt64, Int128, UInt128,
+                       Float32, Float64)
+        a = CUDA.zeros(T, n)
+        b = CUDA.zeros(T, n)
 
-    @cuda threads=n kernel(a)
-    @cuda threads=n kernel(b)
+        @cuda threads=n kernel(a)
+        @cuda threads=n kernel(b)
 
-    @test Array(a) == Array(b)
+        @test Array(a) == Array(b)
+    end
 end
 
 end