Add adapt_storage methods for explicit sparse formats in CuSparse arrays (#3129)

Author: albertomercurio

lib/cusparse/src/array.jl

@@ -683,7 +683,19 @@ SparseArrays.SparseMatrixCSC(x::CuSparseMatrixCOO) = SparseMatrixCSC(CuSparseMat
 # GPU array adaptor
 Adapt.adapt_storage(::Type{CuArray}, xs::SparseVector) = CuSparseVector(xs)
 Adapt.adapt_storage(::Type{CuArray}, xs::SparseMatrixCSC) = CuSparseMatrixCSC(xs)
+
+# Explicit sparse formats
+Adapt.adapt_storage(::Type{<:CuSparseVector}, xs::AbstractSparseVector) = CuSparseVector(xs)
+Adapt.adapt_storage(::Type{<:CuSparseMatrixCSC}, xs::AbstractSparseMatrix) = CuSparseMatrixCSC(xs)
+Adapt.adapt_storage(::Type{<:CuSparseMatrixCSR}, xs::AbstractSparseMatrix) = CuSparseMatrixCSR(xs)
+Adapt.adapt_storage(::Type{<:CuSparseMatrixCOO}, xs::AbstractSparseMatrix) = CuSparseMatrixCOO(xs)
+
 ## preserve type parameters
+Adapt.adapt_storage(::Type{<:CuSparseVector{T}}, xs::AbstractSparseVector) where {T} = CuSparseVector{T}(xs)
+Adapt.adapt_storage(::Type{<:CuSparseMatrixCSC{T}}, xs::AbstractSparseMatrix) where {T} = CuSparseMatrixCSC{T}(xs)
+Adapt.adapt_storage(::Type{<:CuSparseMatrixCSR{T}}, xs::AbstractSparseMatrix) where {T} = CuSparseMatrixCSR{T}(xs)
+Adapt.adapt_storage(::Type{<:CuSparseMatrixCOO{T}}, xs::AbstractSparseMatrix) where {T} = CuSparseMatrixCOO{T}(xs)
+
 Adapt.adapt_storage(::Type{<:CuArray{T}}, xs::SparseVector) where {T} = CuSparseVector{T}(xs)
 Adapt.adapt_storage(::Type{<:CuArray{T}}, xs::SparseMatrixCSC) where {T} = CuSparseMatrixCSC{T}(xs)
 Adapt.adapt_storage(::Type{<:CuArray{T, N}}, xs::SparseVector) where {T, N} = CuSparseVector{T}(xs)

lib/cusparse/test/array.jl

@@ -1,3 +1,5 @@
+using Adapt
+
 @testset "array" begin
     x = sprand(m,0.2)
     d_x = CuSparseVector(x)
@@ -366,4 +368,30 @@
     let A = sprand(Float64, m, n, 0.2)
         @test CuSparseMatrixCSC(A) isa CuSparseMatrixCSC{Float64}
     end
+
+    @testset "adapt" begin
+        # CPU sparse array to specific GPU sparse array
+        let A = sprand(Float64, m, n, 0.2)
+            @test Adapt.adapt(CuSparseMatrixCSC, A) isa CuSparseMatrixCSC{Float64}
+            @test Adapt.adapt(CuSparseMatrixCSR, A) isa CuSparseMatrixCSR{Float64}
+            @test Adapt.adapt(CuSparseMatrixCOO, A) isa CuSparseMatrixCOO{Float64}
+            @test Adapt.adapt(CuSparseMatrixCSC{Float32}, A) isa CuSparseMatrixCSC{Float32}
+            @test Adapt.adapt(CuSparseMatrixCSR{Float32}, A) isa CuSparseMatrixCSR{Float32}
+            @test Adapt.adapt(CuSparseMatrixCOO{Float32}, A) isa CuSparseMatrixCOO{Float32}
+        end
+        let v = sprand(Float64, m, 0.2)
+            @test Adapt.adapt(CuSparseVector, v) isa CuSparseVector{Float64}
+            @test Adapt.adapt(CuSparseVector{Float32}, v) isa CuSparseVector{Float32}
+        end
+
+        # CuArray target keeps defaults
+        let A = sprand(Float64, m, n, 0.2)
+            @test Adapt.adapt(CuArray, A) isa CuSparseMatrixCSC{Float64}
+            @test Adapt.adapt(CuArray{Float32}, A) isa CuSparseMatrixCSC{Float32}
+        end
+        let v = sprand(Float64, m, 0.2)
+            @test Adapt.adapt(CuArray, v) isa CuSparseVector{Float64}
+            @test Adapt.adapt(CuArray{Float32}, v) isa CuSparseVector{Float32}
+        end
+    end
 end