shared(AbstractArray) and sync_workgroup

Author: PhilipFackler

ext/AMDGPUExt/AMDGPUExt.jl

@@ -488,9 +488,9 @@ function reduce_kernel_amdgpu_MN((M, N), op, red, ret)
     return nothing
 end
 
-function JACC.shared(x::ROCDeviceArray{T, N}) where {T, N}
+function JACC.shared(::AMDGPUBackend, x::AbstractArray)
     size = length(x)
-    shmem = @ROCDynamicLocalArray(T, size)
+    shmem = @ROCDynamicLocalArray(eltype(x), size)
     num_threads = workgroupDim().x * workgroupDim().y
     if (size <= num_threads)
         if workgroupDim().y == 1
@@ -534,6 +534,8 @@ function JACC.shared(x::ROCDeviceArray{T, N}) where {T, N}
     return shmem
 end
 
+JACC.sync_workgroup(::AMDGPUBackend) = AMDGPU.sync_workgroup()
+
 JACC.array_type(::AMDGPUBackend) = AMDGPU.ROCArray
 
 JACC.array(::AMDGPUBackend, x::Base.Array) = AMDGPU.ROCArray(x)

ext/CUDAExt/CUDAExt.jl

@@ -493,9 +493,9 @@ function reduce_kernel_cuda_MN((M, N), op, red, ret)
     return nothing
 end
 
-function JACC.shared(x::CuDeviceArray{T, N}) where {T, N}
+function JACC.shared(::CUDABackend, x::AbstractArray)
     size = length(x)
-    shmem = CuDynamicSharedArray(T, size)
+    shmem = CuDynamicSharedArray(eltype(x), size)
     num_threads = blockDim().x * blockDim().y
     if (size <= num_threads)
         if blockDim().y == 1
@@ -545,6 +545,8 @@ function JACC.shared(x::CuDeviceArray{T, N}) where {T, N}
     return shmem
 end
 
+JACC.sync_workgroup(::CUDABackend) = CUDA.sync_threads()
+
 JACC.array_type(::CUDABackend) = CUDA.CuArray
 
 JACC.array(::CUDABackend, x::Base.Array) = CUDA.CuArray(x)

ext/oneAPIExt/oneAPIExt.jl

@@ -375,10 +375,10 @@ function reduce_kernel_oneapi_MN((M, N), op, red, ret)
     return nothing
 end
 
-function JACC.shared(x::oneDeviceArray{T, N}) where {T, N}
+function JACC.shared(::oneAPIBackend, x::AbstractArray)
     size::Int32 = length(x)
     # This is wrong, we should use size not 512 ...
-    shmem = oneLocalArray(T, 512)
+    shmem = oneLocalArray(eltype(x), 512)
     num_threads = get_local_size(0) * get_local_size(1)
     if (size <= num_threads)
         if get_local_size(1) == 1
@@ -422,6 +422,8 @@ function JACC.shared(x::oneDeviceArray{T, N}) where {T, N}
     return shmem
 end
 
+JACC.sync_workgroup(::oneAPIBackend) = oneAPI.barrier()
+
 JACC.array_type(::oneAPIBackend) = oneAPI.oneArray
 
 JACC.array(::oneAPIBackend, x::Base.Array) = oneAPI.oneArray(x)

src/JACC.jl

@@ -47,9 +47,9 @@ launch_spec(; kw...) = LaunchSpec{typeof(default_backend())}(; kw...)
 
 default_float(::Any) = Float64
 
-function shared(x::Base.Array{T, N}) where {T, N}
-    return x
-end
+shared(x::AbstractArray) = shared(default_backend(), x)
+
+sync_workgroup() = sync_workgroup(default_backend())
 
 array_type() = array_type(default_backend())
 

src/threads/threads.jl

@@ -105,4 +105,6 @@ JACC.array_type(::ThreadsBackend) = Base.Array
 
 JACC.array(::ThreadsBackend, x::Base.Array) = x
 
+JACC.shared(::ThreadsBackend, x::AbstractArray) = x
+
 end