Merge pull request #822

use correct wavefrontsize for device-side rng

Author: luraess

src/compiler/codegen.jl

@@ -118,14 +118,13 @@ end
 
 function compiler_config(dev::HIP.HIPDevice;
     name::Union{String, Nothing} = nothing, kernel::Bool = true,
-    unsafe_fp_atomics::Bool = true,
+    unsafe_fp_atomics::Bool = true, wavefrontsize64::Bool = HIP.wavefrontsize(dev) == 64,
 )
     dev_isa, features = parse_llvm_features(HIP.gcn_arch(dev))
     if !isempty(features)
         features = "$features,"
     end
 
-    wavefrontsize64 = HIP.wavefrontsize(dev) == 64
     features = if wavefrontsize64
         features * "-wavefrontsize32,+wavefrontsize64"
     else

src/device/random.jl

@@ -82,11 +82,10 @@ end
     threadId = workitemIdx().x +
         (workitemIdx().y - Int32(1)) * workgroupDim().x +
         (workitemIdx().z - Int32(1)) * workgroupDim().x * workgroupDim().y
-    warpId = (threadId - Int32(1)) >> 0x5 + Int32(1)  # fld1 by 32
+    wavefrontsize_log2 = ifelse(wavefrontsize() == UInt32(32), 0x5, 0x6)
+    warpId = (threadId - Int32(1)) >> wavefrontsize_log2 + Int32(1)  # fld1 by 32
 
-    if field === :seed
-        @inbounds global_random_seed()[1]
-    elseif field === :key
+    if field === :key
         @inbounds global_random_keys()[warpId]
     elseif field === :ctr1
         @inbounds global_random_counters()[warpId]
@@ -104,7 +103,8 @@ end
     threadId = workitemIdx().x +
         (workitemIdx().y - Int32(1)) * workgroupDim().x +
         (workitemIdx().z - Int32(1)) * workgroupDim().x * workgroupDim().y
-    warpId = (threadId - Int32(1)) >> 0x5 + Int32(1)  # fld1 by 32
+    wavefrontsize_log2 = ifelse(wavefrontsize() == UInt32(32), 0x5, 0x6)
+    warpId = (threadId - Int32(1)) >> wavefrontsize_log2 + Int32(1)  # fld1 by 32
 
     if field === :key
         @inbounds global_random_keys()[warpId] = x

src/highlevel.jl

@@ -93,7 +93,7 @@ register methods for the the `AMDGPU.Adaptor` type.
 rocconvert(arg) = adapt(Runtime.Adaptor(), arg)
 
 const MACRO_KWARGS = [:launch]
-const COMPILER_KWARGS = [:name, :unsafe_fp_atomics]
+const COMPILER_KWARGS = [:name, :unsafe_fp_atomics, :wavefrontsize64]
 const LAUNCH_KWARGS = [:gridsize, :groupsize, :shmem, :stream, :cooperative]
 
 """

test/device/random.jl

@@ -18,9 +18,11 @@ function apply_seed(seed)
     end
 end
 
-@testset "rand($T), seed $seed" for T in (
+wavefrontsize64 = AMDGPU.HIP.wavefrontsize(AMDGPU.device()) == 32 ? [false, true] : [true]
+
+@testset "rand($T), seed $seed, wavefrontsize64 = $wavefrontsize64" for T in (
     Int32, UInt32, Int64, UInt64, Int128, UInt128, Float16, Float32, Float64,
-), seed in (nothing, #=missing,=# 1234)
+), seed in (nothing, #=missing,=# 1234), wavefrontsize64 in wavefrontsize64
     # different kernel invocations should get different numbers
     @testset "across launches" begin
         function kernel(A::AbstractArray{T}, seed) where {T}
@@ -85,9 +87,9 @@ end
     end
 end
 
-@testset "basic randn($T), seed $seed" for T in (
+@testset "basic randn($T), seed $seed, wavefrontsize64 = $wavefrontsize64" for T in (
     Float16, Float32, Float64,
-), seed in (nothing, #=missing,=# 1234)
+), seed in (nothing, #=missing,=# 1234), wavefrontsize64 in wavefrontsize64
     function kernel(A::AbstractArray{T}, seed) where {T}
         apply_seed(seed)
         tid = workitemIdx().x
@@ -108,9 +110,9 @@ end
     end
 end
 
-@testset "basic randexp($T), seed $seed" for T in (
+@testset "basic randexp($T), seed $seed, wavefrontsize64 = $wavefrontsize64" for T in (
     Float16, Float32, Float64,
-), seed in (nothing, #=missing,=# 1234)
+), seed in (nothing, #=missing,=# 1234), wavefrontsize64 in wavefrontsize64
     function kernel(A::AbstractArray{T}, seed) where {T}
         apply_seed(seed)
         tid = workitemIdx().x

test/device/wavefront.jl

@@ -145,4 +145,8 @@ end
     @roc kernel(RX)
     AMDGPU.synchronize()
     @allowscalar @test RX[1] == wavefrontsize
+
+    @roc wavefrontsize64=true kernel(RX)
+    AMDGPU.synchronize()
+    @allowscalar @test RX[1] == 64
 end