CuSparseMatrixCSR plus and minus other sparse format

Author: yuehhua

lib/cusparse/array.jl

@@ -45,6 +45,8 @@ mutable struct CuSparseMatrixCSC{Tv} <: AbstractCuSparseMatrix{Tv}
     end
 end
 
+CuSparseMatrixCSC(A::CuSparseMatrixCSC) = A
+
 function CUDA.unsafe_free!(xs::CuSparseMatrixCSC)
     unsafe_free!(xs.colPtr)
     unsafe_free!(rowvals(xs))
@@ -72,6 +74,8 @@ mutable struct CuSparseMatrixCSR{Tv} <: AbstractCuSparseMatrix{Tv}
     end
 end
 
+CuSparseMatrixCSR(A::CuSparseMatrixCSR) = A
+
 function CUDA.unsafe_free!(xs::CuSparseMatrixCSR)
     unsafe_free!(xs.rowPtr)
     unsafe_free!(xs.colVal)
@@ -100,6 +104,8 @@ mutable struct CuSparseMatrixBSR{Tv} <: AbstractCuSparseMatrix{Tv}
     end
 end
 
+CuSparseMatrixBSR(A::CuSparseMatrixBSR) = A
+
 function CUDA.unsafe_free!(xs::CuSparseMatrixBSR)
     unsafe_free!(xs.rowPtr)
     unsafe_free!(xs.colVal)
@@ -126,6 +132,8 @@ mutable struct CuSparseMatrixCOO{Tv} <: AbstractCuSparseMatrix{Tv}
     end
 end
 
+CuSparseMatrixCOO(A::CuSparseMatrixCOO) = A
+
 """
 Utility union type of [`CuSparseMatrixCSC`](@ref), [`CuSparseMatrixCSR`](@ref),
 [`CuSparseMatrixBSR`](@ref), [`CuSparseMatrixCOO`](@ref).

lib/cusparse/interfaces.jl

@@ -131,22 +131,22 @@ function Base.:(-)(A::Transpose{T,<:CuSparseMatrixCSR}, B::Transpose{T,<:CuSpars
     return CuSparseMatrixCSR(cscC.colPtr, cscC.rowVal, cscC.nzVal, size(cscC))
 end
 
-function Base.:(+)(A::CuSparseMatrixCSR, B::CuSparseMatrixCSC)
+function Base.:(+)(A::CuSparseMatrixCSR, B::CuSparseMatrix)
     csrB = CuSparseMatrixCSR(B)
     return geam(one(eltype(A)), A, one(eltype(A)), csrB, 'O')
 end
 
-function Base.:(-)(A::CuSparseMatrixCSR, B::CuSparseMatrixCSC)
+function Base.:(-)(A::CuSparseMatrixCSR, B::CuSparseMatrix)
     csrB = CuSparseMatrixCSR(B)
     return geam(one(eltype(A)), A, -one(eltype(A)), csrB, 'O')
 end
 
-function Base.:(+)(A::CuSparseMatrixCSC, B::CuSparseMatrixCSR)
+function Base.:(+)(A::CuSparseMatrix, B::CuSparseMatrixCSR)
     csrA = CuSparseMatrixCSR(A)
     return geam(one(eltype(A)), csrA, one(eltype(A)), B, 'O')
 end
 
-function Base.:(-)(A::CuSparseMatrixCSC, B::CuSparseMatrixCSR)
+function Base.:(-)(A::CuSparseMatrix, B::CuSparseMatrixCSR)
     csrA = CuSparseMatrixCSR(A)
     return geam(one(eltype(A)), csrA, -one(eltype(A)), B, 'O')
 end

test/cusparse/interfaces.jl

@@ -24,6 +24,34 @@ using LinearAlgebra, SparseArrays
         @test C ≈ collect(dC)
     end
 
+    @testset "A±$f(B) $elty" for elty in [Float32, Float64, ComplexF32, ComplexF64],
+                                 f in (CuSparseMatrixCSR, CuSparseMatrixCSC, CuSparseMatrixCOO, x->CuSparseMatrixBSR(x,1))
+        n = 10
+        alpha = rand()
+        beta = rand()
+        A = sprand(elty, n, n, rand())
+        B = sprand(elty, n, n, rand())
+
+        dA = CuSparseMatrixCSR(A)
+        dB = CuSparseMatrixCSR(B)
+
+        C = A + B
+        dC = dA + f(dB)
+        @test C ≈ collect(dC)
+
+        C = B + A
+        dC = f(dB) + dA
+        @test C ≈ collect(dC)
+
+        C = A - B
+        dC = dA - f(dB)
+        @test C ≈ collect(dC)
+
+        C = B - A
+        dC = f(dB) - dA
+        @test C ≈ collect(dC)
+    end
+
     @testset "dense(A)$(op)sparse(B) $elty" for elty in [Float32, Float64, ComplexF32, ComplexF64],
                                                 op in [+, -]
         n = 10