wip

Author: simeonschaub

Project.toml

@@ -30,6 +30,9 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 [weakdeps]
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
+[sources]
+GPUCompiler = {rev = "tb/kernel_state_reference", url = "https://github.com/JuliaGPU/GPUCompiler.jl.git"}
+
 [extensions]
 SpecialFunctionsExt = "SpecialFunctions"
 

src/compiler/compilation.jl

@@ -8,7 +8,7 @@ GPUCompiler.runtime_module(::MetalCompilerJob) = Metal
 
 GPUCompiler.method_table(::MetalCompilerJob) = method_table
 
-GPUCompiler.kernel_state_type(job::MetalCompilerJob) = MtlRefValue{KernelState}
+GPUCompiler.kernel_state_type(job::MetalCompilerJob) = KernelState
 
 function GPUCompiler.finish_module!(@nospecialize(job::MetalCompilerJob),
                                     mod::LLVM.Module, entry::LLVM.Function)

src/device/runtime.jl

@@ -39,5 +39,4 @@ struct KernelState
     random_seed::UInt32
 end
 
-#@inline @generated kernel_state() = GPUCompiler.kernel_state_value(KernelState)
-@inline @generated kernel_state() = GPUCompiler.kernel_state_value(MtlRefValue{KernelState})[]
+@inline @generated kernel_state() = GPUCompiler.kernel_state_value(KernelState)

test/Project.toml

@@ -5,6 +5,7 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
@@ -20,3 +21,6 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+GPUCompiler = {rev = "tb/kernel_state_reference", url = "https://github.com/JuliaGPU/GPUCompiler.jl.git"}

test/device/random.jl

@@ -0,0 +1,13 @@
+@testset "device-side rng" begin
+    function rand_kernel!(a)
+        i = thread_position_in_grid_1d()
+        a[i] = rand(Float32)
+        return
+    end
+
+    n = 128
+    a = Metal.fill!(-1f0, n)
+    @metal threads=n rand_kernel!(a)
+    @test all(0 .<= a .< 1)
+    @test length(unique(Array(a))) == n
+end