Added support for CuArray via Adapt.jl

Project.toml

@@ -3,8 +3,17 @@ uuid = "1277b4bf-5013-50f5-be3d-901d8477a67a"
 repo = "https://github.com/JuliaArrays/ShiftedArrays.jl.git"
 version = "2.0.0"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+
+[extensions]
+CUDASupportExt = ["CUDA", "Adapt"]
+
 [compat]
-julia = "1"
+CUDA = "5.2, 5.3, 5.4, 5.5, 5.6"
+Adapt = "3.7, 4.0, 4.1"
+julia = "1.9, 1.10, 1.11"
 
 [extras]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

ext/CUDASupportExt.jl

@@ -0,0 +1,27 @@
+module CUDASupportExt
+using CUDA 
+using Adapt
+using ShiftedArrays
+using Base # to allow displaying such arrays without causing the single indexing CUDA error
+
+Adapt.adapt_structure(to, x::CircShiftedArray{T, D}) where {T, D} = CircShiftedArray(adapt(to, parent(x)), shifts(x));
+parent_type(::Type{CircShiftedArray{T, N, S}})  where {T, N, S} = S
+Base.Broadcast.BroadcastStyle(::Type{T})  where {T<:CircShiftedArray} = Base.Broadcast.BroadcastStyle(parent_type(T))
+
+# cu_storage_type(::Type{T}) where {CT,CN,CD,T<:CuArray{CT,CN,CD}} = CD
+# lets do this for the ShiftedArray type
+Adapt.adapt_structure(to, x::ShiftedArray{T, M, N}) where {T, M, N} = ShiftedArray(adapt(to, parent(x)), shifts(x); default=ShiftedArrays.default(x));
+
+# function Base.Broadcast.BroadcastStyle(::Type{T})  where (CT,CN,CD,T<: ShiftedArray{<:Any,<:Any,<:Any,<:CuArray})
+function Base.Broadcast.BroadcastStyle(::Type{T})  where {T2, N, CD, T<:ShiftedArray{<:Any,<:Any,<:Any,<:CuArray{T2,N,CD}}}
+    CUDA.CuArrayStyle{N,CD}()
+end
+
+# function Base.show(io::IO, mm::MIME"text/plain", cs::CircShiftedArray) 
+#     CUDA.@allowscalar invoke(Base.show, Tuple{IO, typeof(mm), AbstractArray}, io, mm, cs) 
+# end
+
+# function Base.show(io::IO, mm::MIME"text/plain", cs::ShiftedArray) 
+#     CUDA.@allowscalar invoke(Base.show, Tuple{IO, typeof(mm), AbstractArray}, io, mm, cs) 
+# end
+end

test/runtests.jl

@@ -1,8 +1,22 @@
 using ShiftedArrays, Test
 using AbstractFFTs 
+use_cuda = false;  # set this to true to test ShiftedArrays for the CuArray datatype
+if (use_cuda)
+    using CUDA
+    CUDA.allowscalar(true); # needed for some of the comparisons
+end
+
+function opt_convert(v)
+    if (use_cuda)
+        CuArray(v)
+    else
+        v
+    end
+end
 
 @testset "ShiftedVector" begin
     v = [1, 3, 5, 4]
+    v = opt_convert(v);
     @test all(v .== ShiftedVector(v))
     sv = ShiftedVector(v, -1)
     @test isequal(sv, ShiftedVector(v, (-1,)))
@@ -28,6 +42,7 @@ end
 
 @testset "ShiftedArray" begin
     v = reshape(1:16, 4, 4)
+    v = opt_convert(v);
     @test all(v .== ShiftedArray(v))
     sv = ShiftedArray(v, (-2, 0))
     @test length(sv) == 16
@@ -64,6 +79,7 @@ end
 
 @testset "padded_tuple" begin
     v = rand(2, 2)
+    v = opt_convert(v);
     @test (1, 0) == @inferred ShiftedArrays.padded_tuple(v, 1)
     @test (0, 0) == @inferred ShiftedArrays.padded_tuple(v, ())
     @test (3, 0) == @inferred ShiftedArrays.padded_tuple(v, (3,))
@@ -82,11 +98,12 @@ end
 
 @testset "CircShiftedVector" begin
     v = [1, 3, 5, 4]
+    v = opt_convert(v);
     @test all(v .== CircShiftedVector(v))
     sv = CircShiftedVector(v, -1)
     @test isequal(sv, CircShiftedVector(v, (-1,)))
     @test length(sv) == 4
-    @test all(sv .== [3, 5, 4, 1])
+    @test all(sv .== opt_convert([3, 5, 4, 1]))
     diff = v .- sv
     @test diff == [-2, -2, 1, 3]
     @test shifts(sv) == (3,)
@@ -110,6 +127,7 @@ end
 
 @testset "CircShiftedArray" begin
     v = reshape(1:16, 4, 4)
+    v = opt_convert(v);
     @test all(v .== CircShiftedArray(v))
     sv = CircShiftedArray(v, (-2, 0))
     @test length(sv) == 16
@@ -130,6 +148,7 @@ end
 
 @testset "circshift" begin
     v = reshape(1:16, 4, 4)
+    v = opt_convert(v);
     @test all(circshift(v, (1, -1)) .== ShiftedArrays.circshift(v, (1, -1)))
     @test all(circshift(v, (1,)) .== ShiftedArrays.circshift(v, (1,)))
     @test all(circshift(v, 3) .== ShiftedArrays.circshift(v, 3))
@@ -163,6 +182,7 @@ end
 
 @testset "laglead" begin
     v = [1, 3, 8, 12]
+    v = opt_convert(v);
     diff = v .- ShiftedArrays.lag(v)
     @test isequal(diff, [missing, 2, 5, 4])