Finally?

Author: christiangnrd

Project.toml

@@ -6,6 +6,7 @@ version = "11.2.3"
 AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
+GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -26,6 +27,7 @@ JLD2Ext = "JLD2"
 AcceleratedKernels = "0.4"
 Adapt = "4.0"
 GPUArraysCore = "= 0.2.0"
+GPUToolbox = "0.2, 0.3"
 JLD2 = "0.4, 0.5"
 KernelAbstractions = "0.9.28"
 LLVM = "3.9, 4, 5, 6, 7, 8, 9"

src/GPUArrays.jl

@@ -1,5 +1,6 @@
 module GPUArrays
 
+using GPUToolbox
 using KernelAbstractions
 using Serialization
 using Random

src/host/reverse.jl

@@ -16,7 +16,7 @@ function _reverse(input::AnyGPUArray{T, N}, output::AnyGPUArray{T, N};
     nd_idx = CartesianIndices(input)
 
     ## COV_EXCL_START
-    @kernel unsafe_indices=true function kernel(input::AbstractArray{T, N}, output::AbstractArray{T, N}) where {T, N}
+    @kernel unsafe_indices=true function kernel(input, output)
         offset_in = @groupsize()[1] * (@index(Group, Linear) - 1i32)
         index_in = offset_in + @index(Local, Linear)
 
@@ -32,7 +32,7 @@ function _reverse(input::AnyGPUArray{T, N}, output::AnyGPUArray{T, N};
     nthreads = 256
     nblocks = cld(length(input), nthreads)
 
-    kernel(get_backend(input), nblocks)(input, output; ndrange=length(nblocks))
+    kernel(get_backend(input), nblocks)(input, output; ndrange=length(input))
 end
 
 # in-place version, swapping elements on half the number of threads
@@ -52,10 +52,9 @@ function _reverse!(data::AnyGPUArray{T, N}; dims=1:ndims(data)) where {T, N}
     nd_idx = CartesianIndices(reduced_size)
 
     ## COV_EXCL_START
-    @kernel unsafe_indices=true function kernel(data::AbstractArray{T, N}) where {T, N}
+    @kernel unsafe_indices=true function kernel(data)
         offset_in = @groupsize()[1] * (@index(Group, Linear) - 1i32)
-
-        index_in = offset_in + threadIdx().x
+        index_in = offset_in + @index(Local, Linear)
 
         @inbounds if index_in <= reduced_length
             idx = Tuple(nd_idx[index_in])
@@ -80,7 +79,7 @@ function _reverse!(data::AnyGPUArray{T, N}; dims=1:ndims(data)) where {T, N}
     nthreads = 256
     nblocks = cld(prod(reduced_size), nthreads)
 
-    kernel(get_backend(input), nblocks)(input, output; ndrange=length(nblocks))
+    kernel(get_backend(data), nblocks)(data; ndrange=length(data))
 end