Add polyester for SpMV kernels

.github/workflows/ci.yml

@@ -31,6 +31,35 @@ jobs:
           file: lcov.info
           verbose: true
 
+  threaded-test:
+    name: Test Threaded ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
+    runs-on: ${{ matrix.os }}
+    env:
+      JULIA_NUM_THREADS: 2
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.10'
+        os:
+          - ubuntu-latest
+        arch:
+          - x64
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v5
+        with:
+          file: lcov.info
+          verbose: true
+
   docs:
     name: Documentation
     permissions:

Project.toml

@@ -4,11 +4,15 @@ authors = ["Víctor Sande <vsande@cimne.upc.edu>", "Francesc Verdugo <fverdugo@c
 version = "0.6.10"
 
 [deps]
+Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
 
 [compat]
+Atomix = "1.1.2"
+Polyester = "0.7.19"
 julia = "1"
 
 [extras]

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

@@ -320,13 +320,37 @@ 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)
+  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]*α
+    end
+  end
+  return y
+end
+
+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]*α
@@ -336,6 +360,14 @@ function mul!(y::AbstractVector,A::SparseMatrixCSR,v::AbstractVector, α::Number
 end
 
 function mul!(y::AbstractVector,A::SparseMatrixCSR,v::AbstractVector)
+  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
+function mul!(y::AbstractVector,A::Adjoint{<:Any, <:SparseMatrixCSR},v::AbstractVector)
+  if Threads.nthreads() > 1
+    tmul!(y, A, v)
+  else
+    smul!(y, A, v)
+  end
+end
+
+function smul!(y::AbstractVector,A::Adjoint{<:Any, <:SparseMatrixCSR},v::AbstractVector)
   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{<:Any, <:SparseMatrixCSR},v::AbstractVector)
+  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)