Merge pull request #2651 from JuliaGPU/ksh/mixedsparse

Re-enable mixed precision sparse mv

Author: kshyatt

lib/cusparse/interfaces.jl

@@ -62,19 +62,19 @@ op_wrappers = ((identity, T -> 'N', identity),
                (T -> :(HermOrSym{T, <:$T}), T -> 'N', A -> :(parent($A))))
 
 # legacy methods with final MulAddMul argument
-LinearAlgebra.generic_matvecmul!(C::CuVector{T}, tA::AbstractChar, A::CuSparseMatrix{T}, B::DenseCuVector{T}, _add::MulAddMul) where {T <: Union{Float16, ComplexF16, BlasFloat}} =
+LinearAlgebra.generic_matvecmul!(C::CuVector{T}, tA::AbstractChar, A::CuSparseMatrix{S}, B::DenseCuVector{T}, _add::MulAddMul) where {T <: Union{Float16, ComplexF16, BlasFloat}, S <: Union{Float16, ComplexF16, BlasFloat}} =
     LinearAlgebra.generic_matvecmul!(C, tA, A, B, _add.alpha, _add.beta)
-LinearAlgebra.generic_matvecmul!(C::CuVector{T}, tA::AbstractChar, A::CuSparseMatrix{T}, B::CuSparseVector{T}, _add::MulAddMul) where {T <: Union{Float16, ComplexF16, BlasFloat}} =
+LinearAlgebra.generic_matvecmul!(C::CuVector{T}, tA::AbstractChar, A::CuSparseMatrix{S}, B::CuSparseVector{T}, _add::MulAddMul) where {T <: Union{Float16, ComplexF16, BlasFloat}, S <: Union{Float16, ComplexF16, BlasFloat}} =
     LinearAlgebra.generic_matvecmul!(C, tA, A, B, _add.alpha, _add.beta)
 LinearAlgebra.generic_matmatmul!(C::CuMatrix{T}, tA, tB, A::CuSparseMatrix{T}, B::DenseCuMatrix{T}, _add::MulAddMul) where {T <: Union{Float16, ComplexF16, BlasFloat}} =
     LinearAlgebra.generic_matmatmul!(C, tA, tB, A, B, _add.alpha, _add.beta)
 
-function LinearAlgebra.generic_matvecmul!(C::CuVector{T}, tA::AbstractChar, A::CuSparseMatrix{T}, B::DenseCuVector{T}, alpha::Number, beta::Number) where {T <: Union{Float16, ComplexF16, BlasFloat}}
+function LinearAlgebra.generic_matvecmul!(C::CuVector{T}, tA::AbstractChar, A::CuSparseMatrix{S}, B::DenseCuVector{T}, alpha::Number, beta::Number) where {T <: Union{Float16, ComplexF16, BlasFloat}, S <: Union{Float16, ComplexF16, BlasFloat}}
     tA = tA in ('S', 's', 'H', 'h') ? 'N' : tA
     mv_wrapper(tA, alpha, A, B, beta, C)
 end
 
-function LinearAlgebra.generic_matvecmul!(C::CuVector{T}, tA::AbstractChar, A::CuSparseMatrix{T}, B::CuSparseVector{T}, alpha::Number, beta::Number) where {T <: Union{Float16, ComplexF16, BlasFloat}}
+function LinearAlgebra.generic_matvecmul!(C::CuVector{T}, tA::AbstractChar, A::CuSparseMatrix{S}, B::CuSparseVector{T}, alpha::Number, beta::Number) where {T <: Union{Float16, ComplexF16, BlasFloat}, S <: Union{Float16, ComplexF16, BlasFloat}}
     tA = tA in ('S', 's', 'H', 'h') ? 'N' : tA
     mv_wrapper(tA, alpha, A, CuVector{T}(B), beta, C)
 end

test/libraries/cusparse/interfaces.jl

@@ -128,6 +128,25 @@ using LinearAlgebra, SparseArrays
         end
     end
 
+    @testset "CuSparseMatrix * CuVector -- mul!(c, A, b) mixed $eltys" for eltys in ((Float32, ComplexF32), (Float64, ComplexF64))
+        eltya, eltyb = eltys
+        for opa in (identity, transpose, adjoint)
+            n = 10
+            m = 20
+            A = opa == identity ? sprand(eltya, n, m, 0.5) : sprand(eltya, m, n, 0.5)
+            b = rand(eltyb, m)
+            c = rand(eltyb, n)
+
+            dA = CuSparseMatrixCSR(A)
+            db = CuArray(b)
+            dc = CuArray(c)
+
+            mul!(c, opa(A), b)
+            mul!(dc, opa(dA), db)
+            @test c ≈ collect(dc)
+        end
+    end
+
     for SparseMatrixType in (CuSparseMatrixCSC, CuSparseMatrixCSR, CuSparseMatrixCOO, CuSparseMatrixBSR)
 
         if CUSPARSE.version() >= v"11.7.4"
@@ -565,4 +584,4 @@ end
 
         @test ref_cuda_sparse.colPtr == cuda_spdiagm.colPtr
     end
-end
+end