Merge pull request #162 from fverdugo/spmv

Additional performance improvements in sparse matrix-vector product

Author: fverdugo

CHANGELOG.md

@@ -5,6 +5,16 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.5.1] - 2024-07-26
+
+### Added
+
+- Function `spmv!`.
+
+### Fixed
+
+- Performance improvements in sparse matrix-vector multiplication.
+
 ## [0.5.0] - 2024-07-26
 
 ### Changed

Project.toml

@@ -1,7 +1,7 @@
 name = "PartitionedArrays"
 uuid = "5a9dfac6-5c52-46f7-8278-5e2210713be9"
 authors = ["Francesc Verdugo <[email protected]> and contributors"]
-version = "0.5.0"
+version = "0.5.1"
 
 [deps]
 CircularArrays = "7a955b69-7140-5f4e-a0ed-f168c5e2e749"

src/PartitionedArrays.jl

@@ -153,6 +153,8 @@ export dense_diag
 export dense_diag!
 export rap
 export rap!
+export spmv!
+export spmtv!
 export spmm
 export spmm!
 export spmtm

src/p_sparse_matrix.jl

@@ -1829,7 +1829,7 @@ function LinearAlgebra.mul!(c::PVector,a::PSparseMatrix,b::PVector)
         return mul!(c,a,b,1,0)
     end
     t = consistent!(b)
-    foreach(mul!,own_values(c),own_own_values(a),own_values(b))
+    foreach(spmv!,own_values(c),own_own_values(a),own_values(b))
     wait(t)
     foreach(muladd!,own_values(c),own_ghost_values(a),ghost_values(b))
     c

src/sparse_utils.jl

@@ -597,3 +597,86 @@ function csrr_to_csc_step_2(
   nothing
 end
 
+@inline function spmv!(b,A,x)
+    mul!(b,A,x)
+end
+
+@inline function spmtv!(b,A,x)
+    mul!(b,transpose(A),x)
+end
+
+function spmv!(b,A::SparseMatrixCSR{1},x)
+    @boundscheck begin
+        @assert length(b) == size(A,1)
+        @assert length(x) == size(A,2)
+    end
+    spmv_csr!(b,x,A.rowptr,A.colval,A.nzval)
+end
+
+function spmtv!(b,A::SparseMatrixCSR{1},x)
+    @boundscheck begin
+        @assert length(b) == size(A,2)
+        @assert length(x) == size(A,1)
+    end
+    spmv_csc!(b,x,A.rowptr,A.colval,A.nzval)
+end
+
+function spmv!(b,A::SparseMatrixCSC,x)
+    @boundscheck begin
+        @assert length(b) == size(A,1)
+        @assert length(x) == size(A,2)
+    end
+    spmv_csc!(b,x,A.colptr,A.rowval,A.nzval)
+end
+
+function spmtv!(b,A::SparseMatrixCSC,x)
+    @boundscheck begin
+        @assert length(b) == size(A,2)
+        @assert length(x) == size(A,1)
+    end
+    spmv_csr!(b,x,A.colptr,A.rowval,A.nzval)
+end
+
+function spmv_csr!(b,x,rowptr_A,colval_A,nzval_A)
+    ncols = length(x)
+    nrows = length(b)
+    u = one(eltype(rowptr_A))
+    z = zero(eltype(b))
+    @inbounds for row in 1:nrows
+        pini = rowptr_A[row]
+        pend = rowptr_A[row+1]
+        bi = z
+        p = pini
+        while p < pend
+            aij = nzval_A[p]
+            col = colval_A[p]
+            xj = x[col]
+            bi += aij*xj
+            p += u
+        end
+        b[row] = bi
+    end
+    b
+end
+
+function spmv_csc!(b,x,colptr_A,rowval_A,nzval_A)
+    ncols = length(x)
+    nrows = length(b)
+    u = one(eltype(colptr_A))
+    z = zero(eltype(b))
+    fill!(b,z)
+    @inbounds for col in 1:ncols
+        pini = colptr_A[col]
+        pend = colptr_A[col+1]
+        p = pini
+        xj = x[col]
+        while p < pend
+            aij = nzval_A[p]
+            row = rowval_A[p]
+            b[row] += aij*xj
+            p += u
+        end
+    end
+    b
+end
+

test/sparse_utils_tests.jl

@@ -29,6 +29,21 @@ function test_mat(T)
   B = compresscoo(T,I,J,V,m,n)
   @test typeof(B) == T
   @test A == B
+
+  b1 = ones(Tv,size(B,1))
+  b2 = ones(Tv,size(B,1))
+  x = collect(Tv,1:size(B,2))
+  mul!(b1,B,x)
+  spmv!(b2,B,x)
+  @test norm(b1-b2)/norm(b1) + 1 ≈ 1
+
+  b1 = ones(Tv,size(B,2))
+  b2 = ones(Tv,size(B,2))
+  x = collect(Tv,1:size(B,1))
+  mul!(b1,transpose(B),x)
+  spmtv!(b2,B,x)
+  @test norm(b1-b2)/norm(b1) + 1 ≈ 1
+
 
   i,j,v = findnz(B)
   for (k,(ki,kj,kv)) in enumerate(nziterator(B))