Add polyester for SpMV kernels

Author: termi-official

src/SparseMatricesCSR.jl

@@ -14,6 +14,9 @@ export SymSparseMatrixCSR
 export sparsecsr, symsparsecsr
 export colvals, getBi, getoffset
 
+import Polyester: @batch
+import Atomix: @atomic
+
 include("SparseMatrixCSR.jl")
 
 include("SymSparseMatrixCSR.jl")

src/SparseMatrixCSR.jl

@@ -319,14 +319,22 @@ If  `pred` not given, it counts the number of `true` values.
 count(pred, S::SparseMatrixCSR) = count(pred, nzvalview(S))
 count(S::SparseMatrixCSR) = count(i->true, nzvalview(S))
 
-function mul!(y::AbstractVector,A::SparseMatrixCSR,v::AbstractVector, α::Number, β::Number)
+function mul!(y::AbstractVector,A::Adjoint{T, <:SparseMatrixCSR},v::AbstractVector) where T
+  if Threads.nthreads() > 1
+    tmul!(y, A, v)
+  else
+    smul!(y, A, v)
+  end
+end
+
+function smul!(y::AbstractVector,A::SparseMatrixCSR,v::AbstractVector, α::Number, β::Number)
   A.n == size(v, 1) || throw(DimensionMismatch())
   A.m == size(y, 1) || throw(DimensionMismatch())
   if β != 1
     β != 0 ? rmul!(y, β) : fill!(y, zero(eltype(y)))
   end
   o = getoffset(A)
-  for row = 1:size(y, 1)
+  @batch for row = 1:size(y, 1)
     @inbounds for nz in nzrange(A,row)
       col = A.colval[nz]+o
       y[row] += A.nzval[nz]*v[col]*α
@@ -335,7 +343,31 @@ function mul!(y::AbstractVector,A::SparseMatrixCSR,v::AbstractVector, α::Number
   return y
 end
 
-function mul!(y::AbstractVector,A::SparseMatrixCSR,v::AbstractVector)
+function tmul!(y::AbstractVector,A::SparseMatrixCSR,v::AbstractVector, α::Number, β::Number)
+  A.n == size(v, 1) || throw(DimensionMismatch())
+  A.m == size(y, 1) || throw(DimensionMismatch())
+  if β != 1
+    β != 0 ? rmul!(y, β) : fill!(y, zero(eltype(y)))
+  end
+  o = getoffset(A)
+  @batch for row = 1:size(y, 1)
+    @inbounds for nz in nzrange(A,row)
+      col = A.colval[nz]+o
+      y[row] += A.nzval[nz]*v[col]*α
+    end
+  end
+  return y
+end
+
+function mul!(y::AbstractVector,A::SparseMatrixCSR,v::AbstractVector) where T
+  if Threads.nthreads() > 1
+    tmul!(y, A, v)
+  else
+    smul!(y, A, v)
+  end
+end
+
+function smul!(y::AbstractVector,A::SparseMatrixCSR,v::AbstractVector)
   A.n == size(v, 1) || throw(DimensionMismatch())
   A.m == size(y, 1) || throw(DimensionMismatch())
   fill!(y, zero(eltype(y)))
@@ -349,9 +381,31 @@ function mul!(y::AbstractVector,A::SparseMatrixCSR,v::AbstractVector)
   return y
 end
 
+function tmul!(y::AbstractVector,A::SparseMatrixCSR,v::AbstractVector)
+  A.n == size(v, 1) || throw(DimensionMismatch())
+  A.m == size(y, 1) || throw(DimensionMismatch())
+  fill!(y, zero(eltype(y)))
+  o = getoffset(A)
+  @batch for row = 1:size(y, 1)
+    @inbounds for nz in nzrange(A,row)
+      col = A.colval[nz]+o
+      y[row] += A.nzval[nz]*v[col]
+    end
+  end
+  return y
+end
+
 *(A::SparseMatrixCSR, v::Vector) = (y = similar(v,size(A,1));mul!(y,A,v))
 
 function mul!(y::AbstractVector,A::Adjoint{T, <:SparseMatrixCSR},v::AbstractVector) where T
+  if Threads.nthreads() > 1
+    tmul!(y, A, v)
+  else
+    smul!(y, A, v)
+  end
+end
+
+function smul!(y::AbstractVector,A::Adjoint{T, <:SparseMatrixCSR},v::AbstractVector) where T
   P = A.parent
   P.n == size(y, 1) || throw(DimensionMismatch())
   P.m == size(v, 1) || throw(DimensionMismatch())
@@ -366,6 +420,21 @@ function mul!(y::AbstractVector,A::Adjoint{T, <:SparseMatrixCSR},v::AbstractVect
   return y
 end
 
+function tmul!(y::AbstractVector,A::Adjoint{T, <:SparseMatrixCSR},v::AbstractVector) where T
+  P = A.parent
+  P.n == size(y, 1) || throw(DimensionMismatch())
+  P.m == size(v, 1) || throw(DimensionMismatch())
+  fill!(y,zero(eltype(y)))
+  o = getoffset(P)
+  @batch for row = 1:size(P, 1)
+    for nz in nzrange(P,row)
+      col = P.colval[nz]+o
+      y[col] += P.nzval[nz]*v[row]
+    end
+  end
+  return y
+end
+
 *(A::Adjoint{T, <:SparseMatrixCSR}, v::AbstractVector) where T = (y = similar(v, promote_type(eltype(v),T), size(A,1)); mul!(y, A, v))
 
 function show(io::IO, ::MIME"text/plain", S::SparseMatrixCSR)