Native RNG: Support large arrays by avoiding counter overflowing.

Author: maleadt

src/random.jl

@@ -52,8 +52,8 @@ function Random.rand!(rng::RNG, A::AnyCuArray)
 
         # grid-stride loop
         threadId = threadIdx().x
-        window = blockDim().x * gridDim().x
-        offset = (blockIdx().x - 1) * blockDim().x
+        window = widemul(blockDim().x, gridDim().x)
+        offset = widemul(blockIdx().x - 1i32, blockDim().x)
         while offset < length(A)
             i = threadId + offset
             if i <= length(A)
@@ -96,8 +96,8 @@ function Random.randn!(rng::RNG, A::AnyCuArray{<:Union{AbstractFloat,Complex{<:A
 
         # grid-stride loop
         threadId = threadIdx().x
-        window = (blockDim().x - 1) * gridDim().x
-        offset = (blockIdx().x - 1) * blockDim().x
+        window = widemul(blockDim().x - 1i32, gridDim().x)
+        offset = widemul(blockIdx().x - 1i32, blockDim().x)
         while offset < length(A)
             i = threadId + offset
             j = threadId + offset + window
@@ -129,8 +129,8 @@ function Random.randn!(rng::RNG, A::AnyCuArray{<:Union{AbstractFloat,Complex{<:A
 
         # grid-stride loop
         threadId = threadIdx().x
-        window = (blockDim().x - 1) * gridDim().x
-        offset = (blockIdx().x - 1) * blockDim().x
+        window = widemul(blockDim().x - 1i32, gridDim().x)
+        offset = widemul(blockIdx().x - 1i32, blockDim().x)
         while offset < length(A)
             i = threadId + offset
             if i <= length(A)

test/base/random.jl

@@ -198,3 +198,9 @@ end
     end
 end
 
+@testset "counter overflow" begin
+    rng = CUDA.RNG()
+    c = CUDA.zeros(Float16, (64, 32, 512, 32, 64))
+    rand!(rng, c)
+    randn!(rng, c)
+end