Merge pull request #13 from birkmichael/master

Similar, indexing CuSparseArrays, Vector reshaping

lib/cusparse/array.jl

@@ -322,24 +322,56 @@ Base.getindex(A::AbstractCuSparseMatrix, i, ::Colon)       = getindex(A, i, 1:si
 Base.getindex(A::AbstractCuSparseMatrix, ::Colon, i)       = getindex(A, 1:size(A, 1), i)
 Base.getindex(A::AbstractCuSparseMatrix, I::Tuple{Integer,Integer}) = getindex(A, I[1], I[2])
 
-# column slices
-function Base.getindex(x::CuSparseMatrixCSC, ::Colon, j::Integer)
-    checkbounds(x, :, j)
-    r1 = convert(Int, x.colPtr[j])
-    r2 = convert(Int, x.colPtr[j+1]) - 1
-    CuSparseVector(rowvals(x)[r1:r2], nonzeros(x)[r1:r2], size(x, 1))
+Base.getindex(A::CuSparseMatrixCSC,I::AbstractVector,J::AbstractVector) = CuSparseMatrixCSC(getindex(CuSparseMatrixCSR(getindex(A,:,J)),I,:))
+Base.getindex(A::CuSparseMatrixCSR,I::AbstractVector,J::AbstractVector) = CuSparseMatrixCSR(getindex(CuSparseMatrixCSC(getindex(A,I,:)),:,J))
+
+function Base.getindex(A::CuSparseMatrixCSR,I::AbstractVector,::Colon)
+    m,n = size(A)
+    I_sorted = sort(I);
+    if !(I_sorted[end] <= m && 1 <= I_sorted[1])
+        throw(BoundsError())
+    end
+
+    rowptr = getrowptr(A)
+    vals = nonzeros(A)
+    cols = colvals(A)
+
+    diffVec = vcat(1,diff(rowptr)[I_sorted])
+    newrowptr = cumsum(diffVec)
+
+    indices = map((x,y)->range(x,stop=y),rowptr[I_sorted],rowptr[I_sorted.+1].-1)
+    @CUDA.allowscalar indices = vcat(indices...) #TODO: figure out how to avoid this step
+    newval = vals[indices]
+    newcol = cols[indices]
+    return CuSparseMatrixCSR(newrowptr,newcol,newval,(length(I_sorted),n))
 end
 
-function Base.getindex(x::CuSparseMatrixCSR, i::Integer, ::Colon)
-    checkbounds(x, i, :)
-    c1 = convert(Int, x.rowPtr[i])
-    c2 = convert(Int, x.rowPtr[i+1]) - 1
-    CuSparseVector(x.colVal[c1:c2], nonzeros(x)[c1:c2], size(x, 2))
+function Base.getindex(A::CuSparseMatrixCSC,::Colon,I::AbstractVector)
+    m,n = size(A)
+    I_sorted = sort(I);
+    if !(I_sorted[end] <= n && 1 <= I_sorted[1])
+        throw(BoundsError())
+    end
+
+    colptr = getcolptr(A)
+    vals = nonzeros(A)
+    rows = rowvals(A)
+
+    diffVec = vcat(1,diff(colptr)[I_sorted])
+    newcolptr = cumsum(diffVec)
+
+    indices = map((x,y)->range(x,stop=y),colptr[I_sorted],colptr[I_sorted.+1].-1)
+    @CUDA.allowscalar indices = vcat(indices...) #TODO: figure out how to avoid this step
+    newval = vals[indices]
+    newrow = rows[indices]
+    CuSparseMatrixCSC(newcolptr,newrow,newval,(m,length(I)))
+    return CuSparseMatrixCSC(newcolptr,newrow,newval,(m,length(I)))
 end
 
-# row slices
-Base.getindex(A::CuSparseMatrixCSC, i::Integer, ::Colon) = CuSparseVector(sparse(A[i, 1:end]))  # TODO: optimize
-Base.getindex(A::CuSparseMatrixCSR, ::Colon, j::Integer) = CuSparseVector(sparse(A[1:end, j]))  # TODO: optimize
+Base.getindex(x::CuSparseMatrixCSCR, ::Colon, j::Integer) = CuSparseVector(getindex(x,:,[j]))
+Base.getindex(x::CuSparseMatrixCSCR, i::Integer ,::Colon) = CuSparseVector(getindex(x,[i],:))
+Base.getindex(x::CuSparseMatrixCSCR, I::AbstractVector, j::Integer) = CuSparseVector(getindex(x, I, [j]))
+Base.getindex(x::CuSparseMatrixCSCR, i::Integer, J::AbstractVector) = CuSparseVector(getindex(x, [i], J))
 
 function Base.getindex(A::CuSparseMatrixCSC{T}, i0::Integer, i1::Integer) where T
     m, n = size(A)
@@ -541,7 +573,6 @@ for (gpu, cpu) in [CuSparseMatrixCSC => SparseMatrixCSC,
     end
 end
 
-
 # interop with device arrays
 
 function Adapt.adapt_structure(to::CUDA.Adaptor, x::CuSparseVector)

lib/cusparse/conversions.jl

@@ -41,8 +41,24 @@ function SparseArrays.sparse(I::CuVector{Cint}, J::CuVector{Cint}, V::CuVector{T
     end
 end
 
-## Matrix to vector, no direct conversion
-CuSparseVector(Mat::AbstractCuSparseMatrix) = CuSparseVector(sparse(reshape(Mat,:)))
+# CSC to vector
+function CuSparseVector(Mat::CuSparseMatrixCSC)
+    n,m = size(Mat)
+    dim = n*m
+    I = 1:m
+    nzval = nonzeros(Mat)
+    nzind = rowvals(Mat)
+    colptr = Array(getcolptr(Mat)) #TODO: figure out how to avoid this step
+    indices = map((x,y)->range(x,stop=y),colptr[I],colptr[I.+1].-1)
+    indices = cu.(collect.(indices))
+    nzind_array = map(.+,map(inds->nzind[inds],indices) , n.*(0:m-1)) #Vector of cuArrays with indices that should be vectorised into nzind
+    nzind = vcat(nzind_array...)
+    return CuSparseVector(nzind,nzval,dim)
+end
+
+# CSR, COO, BSR to vector
+CuSparseVector(Mat::AbstractCuSparseMatrix) = CuSparseVector(CuSparseMatrixCSC(Mat))
+
 
 ## CSR to CSC