Implement kernel random state using CuArray.

JuliaGPU · Mar 22, 2021 · 84d94d2 · 84d94d2
1 parent 9157693
commit 84d94d2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 14 deletions.
diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl
@@ -205,15 +205,10 @@ mutable struct HostKernel{F,TT} <: AbstractKernel{F,TT}
     mod::CuModule
     fun::CuFunction
 
-    random_state::Union{Nothing,Missing,Mem.DeviceBuffer}
+    random_state::Union{Nothing,Missing,CuVector{UInt32}}
 
     function HostKernel{F,TT}(ctx::CuContext, mod::CuModule, fun::CuFunction, random_state) where {F,TT}
         kernel = new{F,TT}(ctx, mod, fun, random_state)
-        finalizer(kernel) do k
-            if k.random_state isa Mem.DeviceBuffer
-                @context! skip_destroyed=true k.ctx Mem.free(k.random_state; stream_ordered=false)
-            end
-        end
     end
 end
 

diff --git a/src/device/intrinsics/random.jl b/src/device/intrinsics/random.jl
@@ -17,17 +17,12 @@ struct ThreadLocalXorshift32 <: RandomNumbers.AbstractRNG{UInt32}
 end
 
 function init_random_state!(kernel, len)
-    required_size = sizeof(UInt32) * len
-
-    if kernel.random_state === missing
-        kernel.random_state = Mem.alloc(Mem.Device, required_size; async=true)
-    elseif sizeof(kernel.random_state) < required_size
-        Mem.free(kernel.random_state, async=true)
-        kernel.random_state = Mem.alloc(Mem.Device, required_size; async=true)
+    if kernel.random_state === missing || length(kernel.random_state) < len
+        kernel.random_state = CuVector{UInt32}(undef, len)
     end
 
     random_state_ptr = CuGlobal{Ptr{Cvoid}}(kernel.mod, "global_random_state")
-    random_state_ptr[] = reinterpret(Ptr{Cvoid}, convert(CuPtr{Cvoid}, kernel.random_state))
+    random_state_ptr[] = reinterpret(Ptr{Cvoid}, pointer(kernel.random_state))
 end
 
 @eval @inline function global_random_state()