Fill in missing KA functionality (KA.functional + sparse matrices adaption from CUDAbackend) (#2740)

Author: Abdelrahman912

Diff for: src/CUDAKernels.jl

@@ -1,11 +1,12 @@
 module CUDAKernels
 
 using ..CUDA
-using ..CUDA: @device_override
+using ..CUDA: @device_override, CUSPARSE
 
 import KernelAbstractions as KA
 
 import StaticArrays
+import SparseArrays: AbstractSparseArray
 
 import Adapt
 
@@ -28,9 +29,11 @@ KA.get_backend(::CuArray) = CUDABackend()
 KA.get_backend(::CUSPARSE.AbstractCuSparseArray) = CUDABackend()
 KA.synchronize(::CUDABackend) = synchronize()
 
-Adapt.adapt_storage(::CUDABackend, a::Array) = Adapt.adapt(CuArray, a)
-Adapt.adapt_storage(::CUDABackend, a::CuArray) = a
-Adapt.adapt_storage(::KA.CPU, a::CuArray) = convert(Array, a)
+KA.functional(::CUDABackend) = CUDA.functional()
+
+Adapt.adapt_storage(::CUDABackend, a::AbstractArray) = Adapt.adapt(CuArray, a)
+Adapt.adapt_storage(::CUDABackend, a::Union{CuArray,CUSPARSE.AbstractCuSparseArray}) = a
+Adapt.adapt_storage(::KA.CPU, a::Union{CuArray,CUSPARSE.AbstractCuSparseArray}) = Adapt.adapt(Array, a)
 
 ## memory operations
 

Diff for: test/base/kernelabstractions.jl

@@ -1,7 +1,29 @@
 import KernelAbstractions
+import KernelAbstractions as KA
+using SparseArrays
+using Adapt
+
 include(joinpath(dirname(pathof(KernelAbstractions)), "..", "test", "testsuite.jl"))
 
 Testsuite.testsuite(()->CUDABackend(false, false), "CUDA", CUDA, CuArray, CuDeviceArray)
 for (PreferBlocks, AlwaysInline) in Iterators.product((true, false), (true, false))
     Testsuite.unittest_testsuite(()->CUDABackend(PreferBlocks, AlwaysInline), "CUDA", CUDA, CuDeviceArray)
 end
+
+@testset "KA.functional" begin
+    @test KA.functional(CUDABackend()) == CUDA.functional()
+end
+
+@testset "CUDA Backend Adapt Tests" begin
+    # CPU → GPU
+    A = sprand(Float32, 10, 10, 0.5) #CSC
+    A_d = adapt(CUDABackend(), A) 
+    @test A_d isa CUSPARSE.CuSparseMatrixCSC
+    @test adapt(CUDABackend(), A_d) |> typeof == typeof(A_d)
+
+    # GPU → CPU
+    B_d = A |> cu # CuCSC
+    B = adapt(KA.CPU(), A_d)
+    @test B isa SparseMatrixCSC
+    @test adapt(KA.CPU(), B) |> typeof == typeof(B) 
+end