feat: Adding Base.reverse / reverse! support using AcceleratedKernels.jl

Project.toml

@@ -6,6 +6,7 @@ version = "11.5.15"
 projects = ["lib/GPUArraysCore", "lib/JLArrays", "test", "docs"]
 
 [deps]
+AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
@@ -26,6 +27,7 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 JLD2Ext = "JLD2"
 
 [compat]
+AcceleratedKernels = "0.5"
 Adapt = "4.6.1"
 GPUArraysCore = "= 0.2.1"
 JLD2 = "0.4, 0.5, 0.6"

lib/JLArrays/src/JLArrays.jl

@@ -36,6 +36,7 @@ include("array.jl")
 include("sparse.jl")
 include("broadcast.jl")
 include("mapreduce.jl")
+include("reverse.jl")
 
 # KernelAbstractions
 include("JLKernels.jl")

lib/JLArrays/src/reverse.jl

@@ -0,0 +1,26 @@
+# reverse
+
+# JLArrays are backed by a regular Array, so reverse defers to Base on that storage. This keeps the
+# reference backend independent of AcceleratedKernels' GPU kernels. `reverse!` goes through the
+# out-of-place `reverse` so every `dims` form (integer, tuple, `:`) is handled by Base directly.
+# These mirror the GPUArrays generic signatures one-for-one so JLArrays wins by specificity.
+
+Base.reverse!(data::AnyJLArray; dims=:) =
+    (copyto!(typed_data(data), reverse(typed_data(data); dims)); data)
+
+Base.reverse(input::AnyJLArray; dims=:) = JLArray(reverse(typed_data(input); dims))
+
+Base.@propagate_inbounds function Base.reverse!(data::AnyJLVector, start::Integer,
+                                                stop::Integer=lastindex(data))
+    reverse!(typed_data(data), start, stop)
+    return data
+end
+
+Base.reverse!(data::AnyJLVector) = (reverse!(typed_data(data)); data)
+
+Base.@propagate_inbounds function Base.reverse(input::AnyJLVector, start::Integer,
+                                               stop::Integer=lastindex(input))
+    JLArray(reverse(typed_data(input), start, stop))
+end
+
+Base.reverse(input::AnyJLVector) = JLArray(reverse(typed_data(input)))

src/GPUArrays.jl

@@ -17,6 +17,8 @@ using Reexport
 
 using KernelAbstractions
 
+import AcceleratedKernels as AK
+
 # device functionality
 include("device/abstractarray.jl")
 include("device/sparse.jl")
@@ -27,6 +29,7 @@ include("host/construction.jl")
 ## integrations and specialized methods
 include("host/base.jl")
 include("host/indexing.jl")
+include("host/reverse.jl")
 include("host/broadcast.jl")
 include("host/mapreduce.jl")
 include("host/gemm.jl")

src/host/reverse.jl

@@ -0,0 +1,39 @@
+# reversing
+
+# GPUArrays' `reverse` / `reverse!` forward to AcceleratedKernels, the same way the other array
+# operations delegate to a single vendor-agnostic implementation: AK provides one KernelAbstractions
+# reverse (whole-array and per-`dims`) that every backend shares, so there is no per-backend kernel to
+# maintain here. A contiguous sub-range of a vector is reversed by handing AK a view of that range.
+
+
+# n-dimensional API
+
+function Base.reverse!(data::AnyGPUArray; dims=:)
+    AK.reverse!(data; dims)
+    return data
+end
+
+Base.reverse(input::AnyGPUArray; dims=:) = AK.reverse(input; dims)
+
+
+# 1-dimensional API (in-place)
+
+Base.@propagate_inbounds function Base.reverse!(data::AnyGPUVector, start::Integer,
+                                                stop::Integer=lastindex(data))
+    AK.reverse!(view(data, start:stop))
+    return data
+end
+
+Base.reverse!(data::AnyGPUVector) = (AK.reverse!(data); data)
+
+
+# 1-dimensional API (out-of-place)
+
+Base.@propagate_inbounds function Base.reverse(input::AnyGPUVector, start::Integer,
+                                               stop::Integer=lastindex(input))
+    output = copy(input)
+    AK.reverse!(view(output, start:stop))
+    return output
+end
+
+Base.reverse(input::AnyGPUVector) = AK.reverse(input)

test/testsuite/base.jl

@@ -502,6 +502,61 @@ end
       gA = reshape(AT(A),4)
     end
 
+    @testset "reverse" begin
+      # 1-d out-of-place
+      @test compare(x->reverse(x), AT, rand(Float32, 1000))
+      @test compare(x->reverse(x), AT, rand(Float32, 0))
+      @test compare(x->reverse(x), AT, rand(Float32, 1))
+      @test compare(x->reverse(x, 10), AT, rand(Float32, 1000))
+      @test compare(x->reverse(x, 10, 90), AT, rand(Float32, 1000))
+      @test compare(x->reverse(x, 3, 3), AT, Float32[1,2,3,4,5])
+
+      # 1-d in-place
+      @test compare(x->reverse!(x), AT, rand(Float32, 1000))
+      @test compare(x->reverse!(x), AT, rand(Float32, 0))
+      @test compare(x->reverse!(x), AT, rand(Float32, 1))
+      @test compare(x->reverse!(x, 10), AT, rand(Float32, 1000))
+      @test compare(x->reverse!(x, 10, 90), AT, rand(Float32, 1000))
+      @test compare(x->reverse!(x, 3, 3), AT, Float32[1,2,3,4,5])
+
+      # n-d out-of-place
+      for shape in ([1, 2, 4, 3], [4, 2], [5], [0], [1], [2^5, 2^5, 2^5]),
+          dim in 1:length(shape)
+          @testset "Shape: $shape, Dim: $dim" begin
+              @test compare(x->reverse(x; dims=dim), AT, rand(Float32, shape...))
+
+              cpu = rand(Float32, shape...)
+              gpu = AT(cpu)
+              reverse!(gpu; dims=dim)
+              @test Array(gpu) == reverse(cpu; dims=dim)
+          end
+      end
+
+      # supports multidimensional reverse
+      for shape in ([1,1,1,1], [1, 2, 4, 3], [2^5, 2^5, 2^5]),
+          dim in ((1,2),(2,3),(1,3),:)
+          @testset "Shape: $shape, Dim: $dim" begin
+              @test compare(x->reverse(x; dims=dim), AT, rand(Float32, shape...))
+
+              cpu = rand(Float32, shape...)
+              gpu = AT(cpu)
+              reverse!(gpu; dims=dim)
+              @test Array(gpu) == reverse(cpu; dims=dim)
+          end
+      end
+
+      # wrapped array
+      @test compare(x->reverse(x), AT, reshape(rand(Float32, 2,2), 4))
+
+      # error throwing
+      cpu = rand(Float32, 1,2,3,4)
+      gpu = AT(cpu)
+      @test_throws ArgumentError reverse!(gpu, dims=5)
+      @test_throws ArgumentError reverse!(gpu, dims=0)
+      @test_throws ArgumentError reverse(gpu, dims=5)
+      @test_throws ArgumentError reverse(gpu, dims=0)
+    end
+
     @testset "reinterpret" begin
       A = Int32[-1,-2,-3]
       dA = AT(A)