CUSPARSE: Bugfixes for sparse vector broadcast. (#2780)

Author: maleadt

lib/cusparse/broadcast.jl

@@ -383,11 +383,12 @@ function sparse_to_sparse_broadcast_kernel(f::F, output::CuSparseDeviceVector{Tv
     row          = @inbounds row_and_ptrs[1]
     arg_ptrs     = @inbounds row_and_ptrs[2]
     vals = ntuple(Val(N)) do i
+        @inline
         arg = @inbounds args[i]
         # ptr is 0 if the sparse vector doesn't have an element at this row
         # ptr is 0 if the arg is a scalar AND f preserves zeros
         ptr = @inbounds arg_ptrs[i]
-        _getindex(arg, row, ptr)::Tv
+        _getindex(arg, row, ptr)
     end
     output_val = f(vals...)
     @inbounds output.iPtr[row_ix]  = row
@@ -470,12 +471,13 @@ function sparse_to_dense_broadcast_kernel(::Type{<:CuSparseVector}, f::F,
     row          = @inbounds row_and_ptrs[1]
     arg_ptrs     = @inbounds row_and_ptrs[2]
     vals = ntuple(Val(length(args))) do i
+        @inline
         arg = @inbounds args[i]
         # ptr is 0 if the sparse vector doesn't have an element at this row
         # ptr is row if the arg is dense OR a scalar with non-zero-preserving f
         # ptr is 0 if the arg is a scalar AND f preserves zeros
         ptr = @inbounds arg_ptrs[i]
-        _getindex(arg, row, ptr)::Tv
+        _getindex(arg, row, ptr)
     end
     @inbounds output[row] = f(vals...)
     return

test/libraries/cusparse/broadcast.jl

@@ -32,7 +32,7 @@ using CUDA.CUSPARSE, SparseArrays
         dz = dx .* dy .* elty(2)
         @test dz isa typ{elty}
         @test z == SparseMatrixCSC(dz)
-        
+
         # multiple inputs
         w = sprand(elty, m, n, p)
         dw = typ(w)
@@ -42,34 +42,34 @@ using CUDA.CUSPARSE, SparseArrays
         @test z == SparseMatrixCSC(dz)
     end
     @testset "$typ($elty)" for typ in [CuSparseVector,]
-        m = 64 
+        m = 64
         p = 0.5
         x = sprand(elty, m, p)
         dx = typ(x)
-        
+
         # zero-preserving
         y  = x .* elty(1)
         dy = dx .* elty(1)
         @test dy isa typ{elty}
-        @test collect(dy.iPtr) == collect(dx.iPtr) 
+        @test collect(dy.iPtr) == collect(dx.iPtr)
         @test collect(dy.iPtr) == y.nzind
         @test collect(dy.nzVal) == y.nzval
         @test y  == SparseVector(dy)
-        
+
         # not zero-preserving
         y = x .+ elty(1)
         dy = dx .+ elty(1)
         @test dy isa CuArray{elty}
         hy = Array(dy)
-        @test Array(y) == hy 
+        @test Array(y) == hy
 
         # involving something dense
         y = x .+ ones(elty, m)
         dy = dx .+ CUDA.ones(elty, m)
         @test dy isa CuArray{elty}
         @test y == Array(dy)
-         
-        # sparse to sparse 
+
+        # sparse to sparse
         dx = typ(x)
         y  = sprand(elty, m, p)
         dy = typ(y)
@@ -88,25 +88,41 @@ using CUDA.CUSPARSE, SparseArrays
         dz = @. dx * dy * dw
         @test dz isa typ{elty}
         @test z == SparseVector(dz)
-        
+
         y = sprand(elty, m, p)
         w = sprand(elty, m, p)
         dense_arr   = rand(elty, m)
-        d_dense_arr = CuArray(dense_arr) 
+        d_dense_arr = CuArray(dense_arr)
         dy = typ(y)
         dw = typ(w)
-        z  = @. x * y * w * dense_arr 
-        dz = @. dx * dy * dw * d_dense_arr 
+        z  = @. x * y * w * dense_arr
+        dz = @. dx * dy * dw * d_dense_arr
         @test dz isa CuArray{elty}
         @test z == Array(dz)
-        
+
         y = sprand(elty, m, p)
         dy = typ(y)
         dx = typ(x)
         z  = x .* y .* elty(2)
         dz = dx .* dy .* elty(2)
         @test dz isa typ{elty}
         @test z == SparseVector(dz)
+
+        # type-mismatching
+        ## non-zero-preserving
+        dx = typ(x)
+        dy = dx .+ 1
+        y = x .+ 1
+        @test dy isa CuArray{promote_type(elty, Int)}
+        @test y == Array(dy)
+        ## zero-preserving
+        dy = dx .* 1
+        y = x .* 1
+        @test dy isa typ{promote_type(elty, Int)}
+        @test collect(dy.iPtr) == collect(dx.iPtr)
+        @test collect(dy.iPtr) == y.nzind
+        @test collect(dy.nzVal) == y.nzval
+        @test y == SparseVector(dy)
     end
 end