Merge pull request #1410 from JuliaGPU/tb/sparse_int

Support limited sparse integer arrays by bitcasting to floating point.

Author: maleadt

lib/cusparse/conversions.jl

@@ -74,7 +74,6 @@ function CuSparseMatrixCSR{T}(S::Adjoint{T, <:CuSparseMatrixCSC{T}}) where {T <:
     return CuSparseMatrixCSR{T}(csc.colPtr, csc.rowVal, conj.(csc.nzVal), size(csc))
 end
 
-
 # by flipping rows and columns, we can use that to get CSC to CSR too
 for (fname,elty) in ((:cusparseScsr2csc, :Float32),
                      (:cusparseDcsr2csc, :Float64),
@@ -141,6 +140,8 @@ for (fname,elty) in ((:cusparseScsr2csc, :Float32),
     end
 end
 
+# implement Float16 conversions using wider types
+# TODO: Float16 is sometimes natively supported
 for (elty, welty) in ((:Float16, :Float32),
                       (:ComplexF16, :ComplexF32))
     @eval begin
@@ -204,6 +205,46 @@ for (elty, welty) in ((:Float16, :Float32),
     end
 end
 
+# implement Int conversions using reinterpreted Float
+for (elty, felty) in ((:Int16, :Float16),
+                      (:Int32, :Float32),
+                      (:Int64, :Float64),
+                      (:Int128, :ComplexF64))
+    @eval begin
+        function CuSparseMatrixCSR{$elty}(csc::CuSparseMatrixCSC{$elty})
+            csc_compat = CuSparseMatrixCSC(
+                csc.colPtr,
+                csc.rowVal,
+                reinterpret($felty, csc.nzVal),
+                size(csc)
+            )
+            csr_compat = CuSparseMatrixCSR(csc_compat)
+            CuSparseMatrixCSR(
+                csr_compat.rowPtr,
+                csr_compat.colVal,
+                reinterpret($elty, csr_compat.nzVal),
+                size(csr_compat)
+            )
+        end
+
+        function CuSparseMatrixCSC{$elty}(csr::CuSparseMatrixCSR{$elty})
+            csr_compat = CuSparseMatrixCSR(
+                csr.rowPtr,
+                csr.colVal,
+                reinterpret($felty, csr.nzVal),
+                size(csr)
+            )
+            csc_compat = CuSparseMatrixCSC(csr_compat)
+            CuSparseMatrixCSC(
+                csc_compat.colPtr,
+                csc_compat.rowVal,
+                reinterpret($elty, csc_compat.nzVal),
+                size(csc_compat)
+            )
+        end
+    end
+end
+
 ## CSR to BSR and vice-versa
 
 for (fname,elty) in ((:cusparseScsr2bsr, :Float32),
@@ -260,6 +301,52 @@ for (fname,elty) in ((:cusparseSbsr2csr, :Float32),
     end
 end
 
+# implement Int conversions using reinterpreted Float
+for (elty, felty) in ((:Int16, :Float16),
+                      (:Int32, :Float32),
+                      (:Int64, :Float64),
+                      (:Int128, :ComplexF64))
+    @eval begin
+        function CuSparseMatrixCSR{$elty}(bsr::CuSparseMatrixBSR{$elty})
+            bsr_compat = CuSparseMatrixBSR(
+                bsr.rowPtr,
+                bsr.colVal,
+                reinterpret($felty, bsr.nzVal),
+                bsr.blockDim,
+                bsr.dir,
+                bsr.nnzb,
+                size(bsr)
+            )
+            csr_compat = CuSparseMatrixCSR(bsr_compat)
+            CuSparseMatrixCSR(
+                csr_compat.rowPtr,
+                csr_compat.colVal,
+                reinterpret($elty, csr_compat.nzVal),
+                size(csr_compat)
+            )
+        end
+
+        function CuSparseMatrixBSR{$elty}(csr::CuSparseMatrixCSR{$elty}, blockDim)
+            csr_compat = CuSparseMatrixCSR(
+                csr.rowPtr,
+                csr.colVal,
+                reinterpret($felty, csr.nzVal),
+                size(csr)
+            )
+            bsr_compat = CuSparseMatrixBSR(csr_compat, blockDim)
+            CuSparseMatrixBSR(
+                bsr_compat.rowPtr,
+                bsr_compat.colVal,
+                reinterpret($elty, bsr_compat.nzVal),
+                bsr_compat.blockDim,
+                bsr_compat.dir,
+                bsr_compat.nnzb,
+                size(bsr_compat)
+            )
+        end
+    end
+end
+
 ## CSR to COO and vice-versa
 
 function CuSparseMatrixCSR(coo::CuSparseMatrixCOO{Tv}, ind::SparseChar='O') where {Tv}

test/cusparse.jl

@@ -108,7 +108,7 @@ blockdim = 5
 end
 
 @testset "construction" begin
-    @testset for elty in [Float32, Float64, ComplexF32, ComplexF64]
+    @testset for elty in [Int32, Int64, Float32, Float64, ComplexF32, ComplexF64]
         @testset "vector" begin
             x = sprand(elty,m, 0.2)
             d_x = CuSparseVector(x)

test/cusparse/broadcast.jl

@@ -3,35 +3,34 @@ using CUDA.CUSPARSE, SparseArrays
 m,n = 2,3
 p = 0.5
 
-# we rely on CUSPARSE conversions, so only test supported types
-for elty in [Float32, Float64]
-    @testset "$typ" for typ in [CuSparseMatrixCSR, CuSparseMatrixCSC]
+for elty in [Int32, Int64, Float32, Float64]
+    @testset "$typ($elty)" for typ in [CuSparseMatrixCSR, CuSparseMatrixCSC]
         x = sprand(elty, m, n, p)
         dx = typ(x)
 
         # zero-preserving
-        y = x .* 1
-        dy = dx .* 1
+        y = x .* elty(1)
+        dy = dx .* elty(1)
         @test dy isa typ{elty}
         @test y == SparseMatrixCSC(dy)
 
         # not zero-preserving
-        y = x .+ 1
-        dy = dx .+ 1
+        y = x .+ elty(1)
+        dy = dx .+ elty(1)
         @test dy isa CuArray{elty}
         @test y == Array(dy)
 
         # involving something dense
-        y = x .* ones(m, n)
-        dy = dx .* CUDA.ones(m, n)
+        y = x .* ones(elty, m, n)
+        dy = dx .* CUDA.ones(elty, m, n)
         @test dy isa CuArray{elty}
         @test y == Array(dy)
 
         # multiple inputs
         y = sprand(elty, m, n, p)
         dy = typ(y)
-        z = x .* y .* 2
-        dz = dx .* dy .* 2
+        z = x .* y .* elty(2)
+        dz = dx .* dy .* elty(2)
         @test dz isa typ{elty}
         @test z == SparseMatrixCSC(dz)
     end