shared(AbstractArray) and sync_workgroup

Author: PhilipFackler

ext/JACCAMDGPU/JACCAMDGPU.jl

@@ -493,9 +493,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
@@ -539,6 +539,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/JACCCUDA/JACCCUDA.jl

@@ -504,9 +504,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
@@ -556,6 +556,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/JACCONEAPI/JACCONEAPI.jl

@@ -418,10 +418,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
@@ -465,6 +465,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

@@ -123,14 +123,22 @@ function parallel_reduce(
     return ret
 end
 
+function shared(::ThreadsBackend, x::AbstractArray)
+    return x
+end
+
 array_type(::ThreadsBackend) = Base.Array
 
 array(::ThreadsBackend, x::Base.Array) = x
 
 default_float(::Any) = Float64
 
-function shared(x::Base.Array{T, N}) where {T, N}
-    return x
+function shared(x::AbstractArray)
+    return shared(default_backend(), x)
+end
+
+function sync_workgroup()
+    return sync_workgroup(default_backend())
 end
 
 array_type() = array_type(default_backend())