some fixes in in-place sparse mul!'s, removed comments

fverdugo · Sep 3, 2024 · 7924078 · 7924078
1 parent 774f07c
commit 7924078
Show file tree

Hide file tree

Showing 7 changed files with 52 additions and 175 deletions.
diff --git a/src/PartitionedArrays.jl b/src/PartitionedArrays.jl
@@ -6,7 +6,8 @@ using LinearAlgebra
 using Printf
 using CircularArrays
 using StaticArrays
-import Base: +,-,*,/
+using Base
+import Base: +,-,*,/,copy
 import MPI
 import IterativeSolvers
 import Distances

diff --git a/src/sequential_implementations.jl b/src/sequential_implementations.jl
@@ -1,4 +1,4 @@
-function *(A::SparseMatrixCSR{Bi,Tv,Ti},B::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
+function Base.:*(A::SparseMatrixCSR{Bi,Tv,Ti},B::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
     p,q = size(A)
     r,s = size(B)
     if q != r && throw(DimensionMismatch("A has dimensions ($(p),$(q)) but B has dimensions ($(p),$(q))"));end
@@ -8,7 +8,7 @@ function *(A::SparseMatrixCSR{Bi,Tv,Ti},B::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,
     SparseMatrixCSR{Bi}(p, s, Ccsc.colptr, rowvals(Ccsc), nonzeros(Ccsc)) 
 end
 
-function *(At::Transpose{Tv, SparseMatrixCSR{Bi,Tv,Ti}},B::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
+function Base.:*(At::Transpose{Tv, SparseMatrixCSR{Bi,Tv,Ti}},B::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
     p,q = size(At)
     r,s = size(B)
     if q != r && throw(DimensionMismatch("At has dimensions ($(p),$(q)) but B has dimensions ($(p),$(q))"));end
@@ -19,7 +19,7 @@ function *(At::Transpose{Tv, SparseMatrixCSR{Bi,Tv,Ti}},B::SparseMatrixCSR{Bi,Tv
     SparseMatrixCSR{Bi}(p, s, Ccsc.colptr, rowvals(Ccsc), nonzeros(Ccsc)) 
 end
 
-function *(A::SparseMatrixCSR{Bi,Tv,Ti},Bt::Transpose{Tv, SparseMatrixCSR{Bi,Tv,Ti}}) where {Bi,Tv,Ti}
+function Base.:*(A::SparseMatrixCSR{Bi,Tv,Ti},Bt::Transpose{Tv, SparseMatrixCSR{Bi,Tv,Ti}}) where {Bi,Tv,Ti}
     p,q = size(A)
     r,s = size(Bt)
     if q != r && throw(DimensionMismatch("A has dimensions ($(p),$(q)) but B has dimensions ($(p),$(q))"));end
@@ -30,7 +30,7 @@ function *(A::SparseMatrixCSR{Bi,Tv,Ti},Bt::Transpose{Tv, SparseMatrixCSR{Bi,Tv,
     SparseMatrixCSR{Bi}(p, s, Ccsc.colptr, rowvals(Ccsc), nonzeros(Ccsc)) 
 end
 
-function *(At::Transpose{Tv,SparseMatrixCSR{Bi,Tv,Ti}},Bt::Transpose{Tv, SparseMatrixCSR{Bi,Tv,Ti}}) where {Bi,Tv,Ti}
+function Base.:*(At::Transpose{Tv,SparseMatrixCSR{Bi,Tv,Ti}},Bt::Transpose{Tv, SparseMatrixCSR{Bi,Tv,Ti}}) where {Bi,Tv,Ti}
     p,q = size(At)
     r,s = size(Bt)
     if q != r && throw(DimensionMismatch("A has dimensions ($(p),$(q)) but B has dimensions ($(p),$(q))"));end
@@ -42,12 +42,12 @@ function *(At::Transpose{Tv,SparseMatrixCSR{Bi,Tv,Ti}},Bt::Transpose{Tv, SparseM
     SparseMatrixCSR{Bi}(p, s, Ccsc.colptr, rowvals(Ccsc), nonzeros(Ccsc)) 
 end
 
-function *(x::Number,A::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
+function Base.:*(x::Number,A::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
     SparseMatrixCSR{Bi}(size(A)..., copy(A.rowptr), copy(A.colval), map(a -> x*a, A.nzval))
 end
-function *(A::SparseMatrixCSR,x::Number) *(x,A) end
+function Base.:*(A::SparseMatrixCSR,x::Number) *(x,A) end
 
-function /(A::SparseMatrixCSR{Bi,Tv,Ti},x::Number) where {Bi,Tv,Ti}
+function Base.:/(A::SparseMatrixCSR{Bi,Tv,Ti},x::Number) where {Bi,Tv,Ti}
     SparseMatrixCSR{Bi}(size(A)..., copy(A.rowptr), copy(A.colval), map(a -> a/x, A.nzval))
 end
 
@@ -87,7 +87,9 @@ function LinearAlgebra.mul!(C::SparseMatrixCSC{Tv,Ti},
         end
         for ip in nzrange(C,j)
             i = IC[ip]
-            VC[ip] = x[i]
+            if xb[i] == j
+                VC[ip] = x[i]
+            end
         end
     end
     C
@@ -132,7 +134,9 @@ function LinearAlgebra.mul!(C::SparseMatrixCSC{Tv,Ti},
         end
         for ip in nzrange(C,j)
             i = IC[ip]
-            VC[ip] += α*x[i]
+            if xb[i] == j
+                VC[ip] += α*x[i]
+            end
         end
     end
     C
@@ -482,12 +486,6 @@ function LinearAlgebra.mul!(C::SparseMatrixCSC{Tv,Ti},
     C
 end
 
-# function LinearAlgebra.mul!(C::SparseMatrixCSC{Tv,Ti},At::Transpose{Tv,SparseMatrixCSC{Tv,Ti}},Bt::Transpose{Tv,SparseMatrixCSC{Tv,Ti}},α::Number,β::Number) where {Tv,Ti}
-#     mul!(C,Bt.parent,At.parent,α,β)
-#     C
-# end
-
-
 function LinearAlgebra.mul!(C::SparseMatrixCSR{Bi,Tv,Ti},
                             A::SparseMatrixCSR{Bi,Tv,Ti},
                             B::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
@@ -553,7 +551,7 @@ function LinearAlgebra.mul!(C::SparseMatrixCSR{Bi,Tv,Ti},
 end
 
 # Alternative to lazy csr to csc for matrix addition that does not drop structural zeros.
-function +(A::SparseMatrixCSR{Bi,Tv,Ti},B::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
+function Base.:+(A::SparseMatrixCSR{Bi,Tv,Ti},B::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
     if size(A) == size(B) || throw(DimensionMismatch("Size of B $(size(B)) must match size of A $(size(A))"));end
     p,q = size(A)
     nnz_C_upperbound = nnz(A) + nnz(B)
@@ -611,7 +609,7 @@ function +(A::SparseMatrixCSR{Bi,Tv,Ti},B::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,
 end
 
 # Alternative to lazy csr to csc for matrix subtraction that does not drop structural zeros. Subtracts B from A, i.e. A - B.
-function -(A::SparseMatrixCSR{Bi,Tv,Ti},B::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
+function Base.:-(A::SparseMatrixCSR{Bi,Tv,Ti},B::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
     if size(A) == size(B) || throw(DimensionMismatch("Size of B $(size(B)) must match size of A $(size(A))"));end
     nnz_C_upperbound = nnz(A) + nnz(B)
     p,r = size(A)
@@ -668,7 +666,7 @@ function -(A::SparseMatrixCSR{Bi,Tv,Ti},B::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,
     SparseMatrixCSR{Bi}(p,r,IC,JC,VC)   # A += B
 end
 
-function -(A::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
+function Base.:-(A::SparseMatrixCSR{Bi,Tv,Ti}) where {Bi,Tv,Ti}
     SparseMatrixCSR{Bi}(size(A)..., copy(A.rowptr), copy(A.colval), map(a->-a, A.nzval))
 end
 
@@ -731,7 +729,7 @@ function +(A::SparseMatrixCSC{Tv,Ti},B::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
 end
 
 # Alternative to lazy csr to csc for matrix subtraction that does not drop structural zeros. Subtracts B from A, i.e. A - B.
-function -(A::SparseMatrixCSC{Tv,Ti},B::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
+function Base.:-(A::SparseMatrixCSC{Tv,Ti},B::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
     if size(A) == size(B) || throw(DimensionMismatch("Size of B $(size(B)) must match size of A $(size(A))"));end
     p,q = size(A)
     nnz_C_upperbound = nnz(A) + nnz(B)
@@ -788,7 +786,7 @@ function -(A::SparseMatrixCSC{Tv,Ti},B::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
     SparseMatrixCSC{Tv,Ti}(p,q,JC,IC,VC)
 end
 
-function -(A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
+function Base.:-(A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
     SparseMatrixCSC{Tv,Ti}(size(A)..., copy(A.colptr), copy(A.rowval), map(a->-a, A.nzval))
 end
 
@@ -1083,8 +1081,7 @@ function RAP!(C::SparseMatrixCSR{Bi,Tv,Ti}, Plt::Transpose{Tv,SparseMatrixCSR{Bi
             v = nonzeros(Pl)[kp]
             for jp in nzrange(C,k)
                 j = JC[jp]
-                if xb[j] == i 
-                    # println("HUHHH")
+                if xb[j] == i
                     VC[jp] += v*x[j]
                 end 
             end
@@ -1183,15 +1180,12 @@ function RAP!(C::SparseMatrixCSR{Bi,Tv,Ti}, Plt::Transpose{Tv,SparseMatrixCSR{Bi
                 end
             end
         end
-        # @show((i,xb))
-        # @show((x))
         for kp in nzrange(Pl, i)
             k = colvals(Pl)[kp] # rowvals when transposed conceptually
             v = nonzeros(Pl)[kp]
             for jp in nzrange(C,k)
                 j = JC[jp]
                 if xb[j] == i
-                    # println("EEEEEEEEEEEEEEEEEEHHH")
                     VC[jp] += α*v*x[j]
                 end 
             end
@@ -1312,7 +1306,9 @@ function RAP(Pl::SparseMatrixCSR{Bi,Tv,Ti}, A::SparseMatrixCSR{Bi,Tv,Ti}, Pr::Sp
             end
             for ind in nzrange(C,i)
                 j = JC[ind]
-                VC[ind] = xC[j]
+                if xbC[j] == i
+                    VC[ind] = xC[j]
+                end
             end
         end
     end
@@ -1381,7 +1377,9 @@ function RAP!(C::SparseMatrixCSR{Bi,Tv,Ti},Pl::SparseMatrixCSR{Bi,Tv,Ti}, A::Spa
         end
         for ind in nzrange(C,i)
             j = JC[ind]
-            VC[ind] = xC[j]
+            if xbC[j] == i
+                VC[ind] = xC[j]
+            end
         end
     end
     C
@@ -1441,7 +1439,9 @@ function RAP!(C::SparseMatrixCSR{Bi,Tv,Ti},Pl::SparseMatrixCSR{Bi,Tv,Ti}, A::Spa
         end
         for ind in nzrange(C,i)
             j = JC[ind]
-            VC[ind] += α*xC[j]
+            if xbC[j] == i
+                VC[ind] += α*xC[j]
+            end
         end
     end
     C
@@ -1457,142 +1457,6 @@ function RAP(Pl::SparseMatrixCSR{Bi,Tv,Ti}, A::SparseMatrixCSR{Bi,Tv,Ti}, Prt::T
     RAP(Pl,A,halfperm(Prt.parent))
 end
 
-# # Not worth it, complexity of N^2, very slow for small problems
-# function RAP(Pl::SparseMatrixCSR{Bi,Tv,Ti}, A::SparseMatrixCSR{Bi,Tv,Ti}, Prt::Transpose{Tv,SparseMatrixCSR{Bi,Tv,Ti}}) where {Bi,Tv,Ti}
-#     p,q = size(Pl)
-#     m,r = size(A)
-#     n,s = size(Prt)
-#     if q == m || throw(DimensionMismatch("Invalid dimensions for R*A: ($p,$q)*($m,$r),"));end
-#     if r == n || throw(DimensionMismatch("Invalid dimensions: RA*P: ($p,$r)*($n,$s)"));end
-#     # find max row length of transposed matrix
-#     function find_row_lengths!(At,xb)
-#         foreach(colvals(At.parent)) do i
-#             xb[i] += 1
-#         end
-#         xb
-#     end
-#     function RAP_symbolic!(Pl,A,Prt)
-#         Pr = Prt.parent
-#         JPl = colvals(Pl)
-#         JA = colvals(A)
-#         IPr = colvals(Pr) # colvals can be interpreted as rowvals when Pr is virtually transposed.
-#         xb = zeros(Ti, q)
-#         x = similar(xb, Tv) # sparse accumulator
-#         max_rPl = find_max_row_length(Pl)
-#         max_rA = find_max_row_length(A)
-#         find_row_lengths!(Prt, xb)
-
-#         max_rRA = max_rPl*max_rA
-#         JRA = Vector{Ti}(undef,max_rRA)
-#         JRA_permutation = collect(Ti, 1:max_rRA)
-#         nnz_C_upperbound = sum(l->max_rRA*l,xb)#p*max_rPl*max_rA*max_rPl
-#         # @show(nnz_C_upperbound)
-#         xb .= 0
-#         IC = Vector{Ti}(undef,p+1)
-#         JC = Vector{Ti}(undef, nnz_C_upperbound)
-#         nnz_C = 1
-#         IC[1] = nnz_C
-#         lp = Ref{Ti}()
-#         for i in 1:p
-#             lp[] = 0
-#              # local column pointer, refresh every row, start at 0 to allow empty rows
-#             # loop over columns "j" in row i of A
-#             for jp in nzrange(Pl, i)
-#                 j = JPl[jp]
-#                 # loop over columns "k" in row j of B
-#                 for kp in nzrange(A, j)
-#                     k = JA[kp]
-#                     # since C is constructed rowwise, xb tracks if a column index is present in a new row in C.
-#                     if xb[k] != i
-#                         lp[] += 1
-#                         JRA[lp[]] = k
-#                         xb[k] = i
-#                     end
-#                 end
-#             end
-#             sort!(JRA_permutation,alg=QuickSort,by=i -> i <= lp[] ? JRA[i] : typemax(i))
-#             j_min = JRA[JRA_permutation[1]]
-#             j_max = JRA[JRA_permutation[lp[]]]
-#             for j in 1:size(Prt,2)
-#                 ip_range = nzrange(Pr,j)
-#                 ip = ip_range.start
-#                 ip_stop = ip_range.stop
-#                 i_min = IPr[ip]
-#                 i_max = IPr[ip_stop]
-#                 if i_min > j_max || j_min > i_max # no intersection
-#                     continue
-#                 end
-#                 while ip <= ip_stop
-#                     iPr = IPr[ip]
-#                     if xb[iPr] == i
-#                         JC[nnz_C] = j
-#                         nnz_C += 1
-#                         break
-#                     end
-#                     ip +=1
-#                 end
-#             end
-#             IC[i+1] = nnz_C
-#         end
-#         nnz_C -= 1
-#         resize!(JC,nnz_C)
-#         VC = zeros(Tv,nnz_C)
-#         cache = (xb,x,JRA)
-#         SparseMatrixCSR{Bi}(p,s,IC,JC,VC), cache # values not yet initialized
-#     end
-
-#     function RAP_numeric!(C,Pl,A,Prt,cache)
-#         Pr = Prt.parent
-#         JPl = colvals(Pl)
-#         VPl = nonzeros(Pl)
-#         JA = colvals(A)
-#         VA = nonzeros(A)
-#         IPr = colvals(Pr) # colvals can be interpreted as rowvals when Pr is virtually transposed.
-#         VPr = nonzeros(Pr)
-#         JC = colvals(C)
-#         VC = nonzeros(C)
-#         (xb,x,_) = cache
-#         for i in 1:p
-#             # loop over columns "j" in row i of A
-#             for jp in nzrange(Pl, i)
-#                 j = JPl[jp]
-#                 # loop over columns "k" in row j of B
-#                 for kp in nzrange(A, j)
-#                     k = JA[kp]
-#                     # since C is constructed rowwise, xb tracks if a column index is present in a new row in C.
-#                     if xb[k] != i
-#                         xb[k] = i
-#                         x[k] = VPl[jp] * VA[kp]
-#                     else
-#                         x[k] += VPl[jp] * VA[kp]
-#                     end
-#                 end
-#             end
-#             for col_ptr in nzrange(C,i)
-#                 Pr_col = JC[col_ptr]
-#                 Pr_rows_range = nzrange(Pr,Pr_col) # column l of Pr^T
-#                 for ip in Pr_rows_range
-#                     Pr_row = IPr[ip]
-#                     if xb[Pr_row] == i
-#                         VC[col_ptr] += x[Pr_row]*VPr[ip]
-#                     end
-#                 end
-#             end
-#         end
-#     end
-#     function _RAP(Pl,A,Prt)
-#         C,cache = RAP_symbolic!(Pl,A,Prt)
-#         # @code_warntype RAP_symbolic!(Pl,A,Prt)
-#         xb = cache[1]
-#         xb.=0
-#         RAP_numeric!(C,Pl,A,Prt,cache)
-#         # @code_warntype RAP_numeric!(C,Pl,A,Prt,cache)
-
-#         C,cache
-#     end
-#     _RAP(Pl,A,Prt)
-# end
-
 function RAP!(C::SparseMatrixCSR{Bi,Tv,Ti}, Pl::SparseMatrixCSR{Bi,Tv,Ti}, A::SparseMatrixCSR{Bi,Tv,Ti}, Prt::Transpose{Tv,SparseMatrixCSR{Bi,Tv,Ti}},cache) where {Bi,Tv,Ti}
     p,q = size(Pl)
     m,r = size(A)
@@ -1659,8 +1523,6 @@ function RAP!(C::SparseMatrixCSR{Bi,Tv,Ti}, Pl::SparseMatrixCSR{Bi,Tv,Ti}, A::Sp
     VC .*= β
     # some cache items are present with the regular RAP product in mind, which is how the allocating verison is performed
     (xb,_,x,_,_) = cache
-    # yes = 0
-    # no = 0
     for i in 1:p
         # loop over columns "j" in row i of A
         for jp in nzrange(Pl, i)
@@ -1685,16 +1547,12 @@ function RAP!(C::SparseMatrixCSR{Bi,Tv,Ti}, Pl::SparseMatrixCSR{Bi,Tv,Ti}, A::Sp
                 iPr = IPr[ip]
                 if xb[iPr] == i
                     v += x[iPr]*VPr[ip]
-                #     yes += 1
-                # else
-                #     no += 1
                 end
             end
 
             VC[jpPr] += α*v
         end
     end
-    # @show((yes,no,yes/(yes+no),length(nonzeros(C))/p))
     C
 end
 

diff --git a/src/sparse_helpers.jl b/src/sparse_helpers.jl
@@ -177,14 +177,22 @@ end
 
 function approx_equivalent(A::SparseMatrixCSR,B::SparseMatrixCSC) approx_equivalent_equivalent(B,A) end
 
-function copy(At::Transpose{Tv,SparseMatrixCSR{Bi,Tv,Ti}}) where {Bi,Tv,Ti}
+function Base.copy(At::Transpose{Tv,SparseMatrixCSR{Bi,Tv,Ti}}) where {Bi,Tv,Ti}
     A = At.parent
     p,q = size(A)
     Acsc = SparseMatrixCSC{Tv, Ti}(q, p, A.rowptr, A.colval, A.nzval)
     Acsc_T = copy(transpose(Acsc)) # materialize SparseMAtrixCSC transpose
     SparseMatrixCSR{Bi}(q, p, Acsc_T.colptr, rowvals(Acsc_T), nonzeros(Acsc_T))
 end
 
+function Base.similar(A::SparseMatrixCSR{Bi}, m::Integer, n::Integer) where Bi
+    return SparseMatrixCSR{1}(m, n, ones(eltype(A.rowptr), m+1), eltype(A.colval)[], eltype(A.nzval)[])
+end
+
+function Base.similar(A::SparseMatrixCSR{Bi}) where Bi
+    return SparseMatrixCSR{Bi}(size(A)..., copy(A.rowptr), copy(colvals(A)), similar(nonzeros(A)))
+end
+
 function pointer_array(A::SparseMatrixCSR)
     A.rowptr
 end

diff --git a/test/p_sparse_matrix_tests.jl b/test/p_sparse_matrix_tests.jl
@@ -377,6 +377,7 @@ function p_sparse_matrix_tests(distribute)
     end
     A_seq = centralize(A)
     spmm!(B,Z,A,cacheB)
+
     @test centralize(B) ≈ Z_seq*(A_seq)
     B = transpose(Z)*A
     @test centralize(B) ≈ transpose(Z_seq)*A_seq