From 72f2c657be9eb1d17c67cef3897999243d53154f Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 4 Mar 2025 15:58:51 +0530 Subject: [PATCH 01/12] 5-term mul! with Diagonal --- src/linalg.jl | 84 +++++++++++++++++++++++++++++++++----------------- test/linalg.jl | 11 +++++++ 2 files changed, 67 insertions(+), 28 deletions(-) diff --git a/src/linalg.jl b/src/linalg.jl index a7dbf350..5a6ce54b 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -1877,47 +1877,75 @@ inv(A::AbstractSparseMatrixCSC) = error("The inverse of a sparse matrix can ofte ## scale methods # Copy colptr and rowval from one sparse matrix to another -function copyinds!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC) - if getcolptr(C) !== getcolptr(A) +function copyinds!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC; copy_rows=true, copy_cols=true) + if copy_cols && getcolptr(C) !== getcolptr(A) resize!(getcolptr(C), length(getcolptr(A))) copyto!(getcolptr(C), getcolptr(A)) end - if rowvals(C) !== rowvals(A) + if copy_rows && rowvals(C) !== rowvals(A) resize!(rowvals(C), length(rowvals(A))) copyto!(rowvals(C), rowvals(A)) end end # multiply by diagonal matrix as vector -function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, D::Diagonal) - m, n = size(A) - b = D.diag - lb = length(b) - n == lb || throw(DimensionMismatch("A has size ($m, $n) but D has size ($lb, $lb)")) - size(A)==size(C) || throw(DimensionMismatch("A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))")) - copyinds!(C, A) - Cnzval = nonzeros(C) - Anzval = nonzeros(A) - resize!(Cnzval, length(Anzval)) - for col in axes(A,2), p in nzrange(A, col) - @inbounds Cnzval[p] = Anzval[p] * b[col] +function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, D::Diagonal, alpha::Number, beta::Number) + beta_is_zero = iszero(beta) + rows_match = rowvals(C) == rowvals(A) + cols_match = getcolptr(C) == getcolptr(A) + identical_nzinds = rows_match && cols_match + if beta_is_zero || identical_nzinds + m, n = size(A) + b = D.diag + lb = length(b) + n == lb || throw(DimensionMismatch(lazy"A has size ($m, $n) but D has size ($lb, $lb)")) + size(A)==size(C) || throw(DimensionMismatch(lazy"A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))")) + identical_nzinds || copyinds!(C, A, copy_rows = !rows_match, copy_cols = !cols_match) + Cnzval = nonzeros(C) + Anzval = nonzeros(A) + resize!(Cnzval, length(Anzval)) + if beta_is_zero + for col in axes(A,2), p in nzrange(A, col) + @inbounds Cnzval[p] = Anzval[p] * b[col] * alpha + end + else + for col in axes(A,2), p in nzrange(A, col) + @inbounds Cnzval[p] = Anzval[p] * b[col] * alpha + Cnzval[p] * beta + end + end + else + @invoke mul!(C::AbstractMatrix, A::AbstractMatrix, D::Diagonal, alpha, beta) end C end -function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AbstractSparseMatrixCSC) - m, n = size(A) - b = D.diag - lb = length(b) - m == lb || throw(DimensionMismatch("D has size ($lb, $lb) but A has size ($m, $n)")) - size(A)==size(C) || throw(DimensionMismatch("A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))")) - copyinds!(C, A) - Cnzval = nonzeros(C) - Anzval = nonzeros(A) - Arowval = rowvals(A) - resize!(Cnzval, length(Anzval)) - for col in axes(A,2), p in nzrange(A, col) - @inbounds Cnzval[p] = b[Arowval[p]] * Anzval[p] +function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AbstractSparseMatrixCSC, alpha::Number, beta::Number) + beta_is_zero = iszero(beta) + rows_match = rowvals(C) == rowvals(A) + cols_match = getcolptr(C) == getcolptr(A) + identical_nzinds = rows_match && cols_match + if beta_is_zero || identical_nzinds + m, n = size(A) + b = D.diag + lb = length(b) + m == lb || throw(DimensionMismatch(lazy"D has size ($lb, $lb) but A has size ($m, $n)")) + size(A)==size(C) || throw(DimensionMismatch(lazy"A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))")) + identical_nzinds || copyinds!(C, A, copy_rows = !rows_match, copy_cols = !cols_match) + Cnzval = nonzeros(C) + Anzval = nonzeros(A) + Arowval = rowvals(A) + resize!(Cnzval, length(Anzval)) + if beta_is_zero + for col in axes(A,2), p in nzrange(A, col) + @inbounds Cnzval[p] = b[Arowval[p]] * Anzval[p] * alpha + end + else + for col in axes(A,2), p in nzrange(A, col) + @inbounds Cnzval[p] = b[Arowval[p]] * Anzval[p] * alpha + Cnzval[p] * beta + end + end + else + @invoke mul!(C::AbstractMatrix, D::Diagonal, A::AbstractMatrix, alpha, beta) end C end diff --git a/test/linalg.jl b/test/linalg.jl index b6113c50..bdd1867a 100644 --- a/test/linalg.jl +++ b/test/linalg.jl @@ -580,6 +580,17 @@ end @test lmul!(D, copy(sA)) ≈ D * dA @test mul!(sC, D, copy(sA)) ≈ D * dA end + + @testset "5-arg mul!" begin + sA2 = similar(sA) + nonzeros(sA2) .= 1 + sA3 = copy(sA2) + D = Diagonal(rand(size(sA,2))) + @test mul!(sA2, sA, D, 3, 2) ≈ dA * D * 3 + sA3 * 2 + nonzeros(sA2) .= 1 + D = Diagonal(rand(size(sA,1))) + @test mul!(sA2, D, sA, 3, 2) ≈ D * dA * 3 + sA3 * 2 + end end @testset "conj" begin From bf73793f2567b47c8187aa096dc4c1bedca1ee4c Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 4 Mar 2025 16:03:28 +0530 Subject: [PATCH 02/12] Add tests --- test/linalg.jl | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/test/linalg.jl b/test/linalg.jl index bdd1867a..e281f542 100644 --- a/test/linalg.jl +++ b/test/linalg.jl @@ -583,13 +583,15 @@ end @testset "5-arg mul!" begin sA2 = similar(sA) - nonzeros(sA2) .= 1 - sA3 = copy(sA2) - D = Diagonal(rand(size(sA,2))) - @test mul!(sA2, sA, D, 3, 2) ≈ dA * D * 3 + sA3 * 2 - nonzeros(sA2) .= 1 - D = Diagonal(rand(size(sA,1))) - @test mul!(sA2, D, sA, 3, 2) ≈ D * dA * 3 + sA3 * 2 + @testset for (alpha, beta) in [(true, false), (true, true), (2,3)] + nonzeros(sA2) .= 1 + sA3 = copy(sA2) + D = Diagonal(rand(size(sA,2))) + @test mul!(sA2, sA, D, alpha, beta) ≈ dA * D * alpha + sA3 * beta + nonzeros(sA2) .= 1 + D = Diagonal(rand(size(sA,1))) + @test mul!(sA2, D, sA, alpha, beta) ≈ D * dA * alpha + sA3 * beta + end end end From 5f36fc00aac059f48eb95b84f69bde9f3972e896 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 4 Mar 2025 16:07:35 +0530 Subject: [PATCH 03/12] Mismatched indices in dest --- test/linalg.jl | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/test/linalg.jl b/test/linalg.jl index e281f542..bcfd1eee 100644 --- a/test/linalg.jl +++ b/test/linalg.jl @@ -582,15 +582,16 @@ end end @testset "5-arg mul!" begin - sA2 = similar(sA) - @testset for (alpha, beta) in [(true, false), (true, true), (2,3)] - nonzeros(sA2) .= 1 - sA3 = copy(sA2) - D = Diagonal(rand(size(sA,2))) - @test mul!(sA2, sA, D, alpha, beta) ≈ dA * D * alpha + sA3 * beta - nonzeros(sA2) .= 1 - D = Diagonal(rand(size(sA,1))) - @test mul!(sA2, D, sA, alpha, beta) ≈ D * dA * alpha + sA3 * beta + for sA2 in (similar(sA), sprand(size(sA)..., 0.1)) + @testset for (alpha, beta) in [(true, false), (true, true), (2,3)] + nonzeros(sA2) .= 1 + sA3 = copy(sA2) + D = Diagonal(rand(size(sA,2))) + @test mul!(sA2, sA, D, alpha, beta) ≈ dA * D * alpha + sA3 * beta + nonzeros(sA2) .= 1 + D = Diagonal(rand(size(sA,1))) + @test mul!(sA2, D, sA, alpha, beta) ≈ D * dA * alpha + sA3 * beta + end end end end From 1e81032c626dd140de18b25219c9377692e9cbdc Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 4 Mar 2025 16:10:16 +0530 Subject: [PATCH 04/12] Specify types of alpha and beta in invoke --- src/linalg.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/linalg.jl b/src/linalg.jl index 5a6ce54b..70a60778 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -1914,7 +1914,7 @@ function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, D::Diagona end end else - @invoke mul!(C::AbstractMatrix, A::AbstractMatrix, D::Diagonal, alpha, beta) + @invoke mul!(C::AbstractMatrix, A::AbstractMatrix, D::Diagonal, alpha::Number, beta::Number) end C end @@ -1945,7 +1945,7 @@ function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AbstractSparseMatrixCS end end else - @invoke mul!(C::AbstractMatrix, D::Diagonal, A::AbstractMatrix, alpha, beta) + @invoke mul!(C::AbstractMatrix, D::Diagonal, A::AbstractMatrix, alpha::Number, beta::Number) end C end From 18c0872b3d324fecfdcb32dac0cf7c51dc60b1d2 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 4 Mar 2025 18:26:31 +0530 Subject: [PATCH 05/12] update mul! module test to 5-arg --- test/issues.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/issues.jl b/test/issues.jl index 9eabc818..f198812a 100644 --- a/test/issues.jl +++ b/test/issues.jl @@ -445,8 +445,8 @@ end A = sprand(5,5,0.5) D = Diagonal(rand(5)) C = copy(A) - m1 = @which mul!(C,A,D) - m2 = @which mul!(C,D,A) + m1 = @which mul!(C,A,D,true,false) + m2 = @which mul!(C,D,A,true,false) @test m1.module == SparseArrays @test m2.module == SparseArrays end From 31782e51d5f98420b8d0b3d602a60fee06ceb0bd Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 4 Mar 2025 21:49:24 +0530 Subject: [PATCH 06/12] Faster path for non-zero beta --- src/linalg.jl | 83 ++++++++++++++++++++++++++++++++++++++++----------- 1 file changed, 66 insertions(+), 17 deletions(-) diff --git a/src/linalg.jl b/src/linalg.jl index 70a60778..03ca7208 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -1888,21 +1888,52 @@ function copyinds!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC; copy_ end end +@inline function rowcheck_index(A::AbstractSparseMatrixCSC, row::Integer, col::Integer) + nzinds = nzrange(A, col) + rows_col = @view rowvals(A)[nzinds] + # faster implementation of row ∈ rows_col, assuming that rows_col is sorted + row_ind_col = searchsortedfirst(rows_col, row) + row_exists = row_ind_col ∈ axes(rows_col,1) && rows_col[row_ind_col] == row + row_ind = row_ind_col + first(nzinds) - firstindex(nzinds) + row_exists, row_ind +end + +function mergeinds!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC) + C_colptr = getcolptr(C) + for col in axes(A,2) + n_extra = 0 + for ind in nzrange(A, col) + row = rowvals(A)[ind] + row_exists, ind = rowcheck_index(C, row, col) + if !row_exists + n_extra += 1 + insert!(rowvals(C), ind, row) + insert!(nonzeros(C), ind, zero(eltype(C))) + C_colptr[col+1] += 1 + end + end + if !iszero(n_extra) + @views C_colptr[col+2:end] .+= n_extra + end + end + C +end + # multiply by diagonal matrix as vector function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, D::Diagonal, alpha::Number, beta::Number) + m, n = size(A) + b = D.diag + lb = length(b) + n == lb || throw(DimensionMismatch(lazy"A has size ($m, $n) but D has size ($lb, $lb)")) + size(A)==size(C) || throw(DimensionMismatch(lazy"A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))")) beta_is_zero = iszero(beta) rows_match = rowvals(C) == rowvals(A) cols_match = getcolptr(C) == getcolptr(A) identical_nzinds = rows_match && cols_match + Cnzval = nonzeros(C) + Anzval = nonzeros(A) if beta_is_zero || identical_nzinds - m, n = size(A) - b = D.diag - lb = length(b) - n == lb || throw(DimensionMismatch(lazy"A has size ($m, $n) but D has size ($lb, $lb)")) - size(A)==size(C) || throw(DimensionMismatch(lazy"A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))")) identical_nzinds || copyinds!(C, A, copy_rows = !rows_match, copy_cols = !cols_match) - Cnzval = nonzeros(C) - Anzval = nonzeros(A) resize!(Cnzval, length(Anzval)) if beta_is_zero for col in axes(A,2), p in nzrange(A, col) @@ -1914,26 +1945,35 @@ function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, D::Diagona end end else - @invoke mul!(C::AbstractMatrix, A::AbstractMatrix, D::Diagonal, alpha::Number, beta::Number) + mergeinds!(C, A) + for col in axes(C,2), p in nzrange(C, col) + row = rowvals(C)[p] + row_exists, row_ind_A = rowcheck_index(A, row, col) + if row_exists + @inbounds Cnzval[p] = Anzval[row_ind_A] * b[col] * alpha + Cnzval[p] * beta + else + @inbounds Cnzval[p] = Cnzval[p] * beta + end + end end C end function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AbstractSparseMatrixCSC, alpha::Number, beta::Number) + m, n = size(A) + b = D.diag + lb = length(b) + m == lb || throw(DimensionMismatch(lazy"D has size ($lb, $lb) but A has size ($m, $n)")) + size(A)==size(C) || throw(DimensionMismatch(lazy"A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))")) beta_is_zero = iszero(beta) rows_match = rowvals(C) == rowvals(A) cols_match = getcolptr(C) == getcolptr(A) identical_nzinds = rows_match && cols_match + Cnzval = nonzeros(C) + Anzval = nonzeros(A) + Arowval = rowvals(A) if beta_is_zero || identical_nzinds - m, n = size(A) - b = D.diag - lb = length(b) - m == lb || throw(DimensionMismatch(lazy"D has size ($lb, $lb) but A has size ($m, $n)")) - size(A)==size(C) || throw(DimensionMismatch(lazy"A has size ($m, $n), D has size ($lb, $lb), C has size $(size(C))")) identical_nzinds || copyinds!(C, A, copy_rows = !rows_match, copy_cols = !cols_match) - Cnzval = nonzeros(C) - Anzval = nonzeros(A) - Arowval = rowvals(A) resize!(Cnzval, length(Anzval)) if beta_is_zero for col in axes(A,2), p in nzrange(A, col) @@ -1945,7 +1985,16 @@ function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AbstractSparseMatrixCS end end else - @invoke mul!(C::AbstractMatrix, D::Diagonal, A::AbstractMatrix, alpha::Number, beta::Number) + mergeinds!(C, A) + for col in axes(C,2), p in nzrange(C, col) + row = rowvals(C)[p] + row_exists, row_ind_A = rowcheck_index(A, row, col) + if row_exists + @inbounds Cnzval[p] = b[row_ind_A] * Anzval[row_ind_A] * alpha + Cnzval[p] * beta + else + @inbounds Cnzval[p] = Cnzval[p] * beta + end + end end C end From db23499723188f06f892db2995db0ae9e6628534 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 4 Mar 2025 21:56:10 +0530 Subject: [PATCH 07/12] Use faster insert --- src/linalg.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/linalg.jl b/src/linalg.jl index 03ca7208..1b735fa8 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -1907,8 +1907,9 @@ function mergeinds!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC) row_exists, ind = rowcheck_index(C, row, col) if !row_exists n_extra += 1 - insert!(rowvals(C), ind, row) - insert!(nonzeros(C), ind, zero(eltype(C))) + nz = getcolptr(C)[end] + _insert!(rowvals(C), ind, row, nz) + _insert!(nonzeros(C), ind, zero(eltype(C)), nz) C_colptr[col+1] += 1 end end From dcf437abf07b1ccf7650a488ad6a031b0e27c4c7 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 4 Mar 2025 22:31:51 +0530 Subject: [PATCH 08/12] Don't assume sorted rows --- src/linalg.jl | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/src/linalg.jl b/src/linalg.jl index 1b735fa8..1c5d0ef6 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -1891,10 +1891,13 @@ end @inline function rowcheck_index(A::AbstractSparseMatrixCSC, row::Integer, col::Integer) nzinds = nzrange(A, col) rows_col = @view rowvals(A)[nzinds] - # faster implementation of row ∈ rows_col, assuming that rows_col is sorted - row_ind_col = searchsortedfirst(rows_col, row) - row_exists = row_ind_col ∈ axes(rows_col,1) && rows_col[row_ind_col] == row - row_ind = row_ind_col + first(nzinds) - firstindex(nzinds) + row_ind_col = findfirst(==(row), rows_col) + row_exists = !isnothing(row_ind_col) + # faster implementation of row ∈ rows_col and obtaining the index, + # assuming that rows_col is sorted + # row_ind_col = searchsortedfirst(rows_col, row) + # row_exists = row_ind_col ∈ axes(rows_col,1) && rows_col[row_ind_col] == row + row_ind = (row_exists ? something(row_ind_col) : length(rows_col) + 1) + first(nzinds) - firstindex(nzinds) row_exists, row_ind end @@ -1907,9 +1910,8 @@ function mergeinds!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC) row_exists, ind = rowcheck_index(C, row, col) if !row_exists n_extra += 1 - nz = getcolptr(C)[end] - _insert!(rowvals(C), ind, row, nz) - _insert!(nonzeros(C), ind, zero(eltype(C)), nz) + insert!(rowvals(C), ind, row) + insert!(nonzeros(C), ind, zero(eltype(C))) C_colptr[col+1] += 1 end end From 798d2574d89dcffe90ba36a5edc550f79495986c Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Tue, 11 Mar 2025 23:12:44 +0530 Subject: [PATCH 09/12] Use `searchsortedfirst` to fix indexing; add `mergeinds!` tests --- src/linalg.jl | 17 ++++++++++++----- test/linalg.jl | 23 +++++++++++++++++++++++ 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/src/linalg.jl b/src/linalg.jl index 1c5d0ef6..d1350297 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -1888,16 +1888,23 @@ function copyinds!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC; copy_ end end +""" + rowcheck_index(A::AbstractSparseMatrixCSC, row::Integer, col::Integer) + +Check if A[row, col] is a stored value, and return the index of the row in `rowvals(A)`. +Returns `(row_exists, row_ind)`, where `row_exists::Bool` signifies +whether the corresponding index is populated, and `row_ind` is the index. +If `row_exists` is `false`, the `row_ind` is the index where the value should be inserted into +`rowvals(A)` such that the subarray `@view rowvals(A)[nzrange(A, col)]` remains sorted. +""" @inline function rowcheck_index(A::AbstractSparseMatrixCSC, row::Integer, col::Integer) nzinds = nzrange(A, col) rows_col = @view rowvals(A)[nzinds] - row_ind_col = findfirst(==(row), rows_col) - row_exists = !isnothing(row_ind_col) # faster implementation of row ∈ rows_col and obtaining the index, # assuming that rows_col is sorted - # row_ind_col = searchsortedfirst(rows_col, row) - # row_exists = row_ind_col ∈ axes(rows_col,1) && rows_col[row_ind_col] == row - row_ind = (row_exists ? something(row_ind_col) : length(rows_col) + 1) + first(nzinds) - firstindex(nzinds) + row_ind_col = searchsortedfirst(rows_col, row) + row_exists = row_ind_col ∈ axes(rows_col,1) && rows_col[row_ind_col] == row + row_ind = row_ind_col + first(nzinds) - firstindex(nzinds) row_exists, row_ind end diff --git a/test/linalg.jl b/test/linalg.jl index bcfd1eee..fe73662d 100644 --- a/test/linalg.jl +++ b/test/linalg.jl @@ -582,6 +582,29 @@ end end @testset "5-arg mul!" begin + @testset "merge indices" begin + # for zero arrays, merge and copy are identical + A = spzeros(size(sA)) + SparseArrays.mergeinds!(A, sA) + B = spzeros(size(sA)) + SparseArrays.copyinds!(B, sA) + @test all(col -> nzrange(A, col) == nzrange(B, col), axes(A,2)) + # for arrays with different indices populated, merge should combine these + A = spzeros(5,5) + A[diagind(A,1)] .= 5 + B = spzeros(5,5) + B[diagind(A,-1)] .= 10 + SparseArrays.mergeinds!(B, A) + @test rowvals(B) == [2, 1,3, 2,4, 3,5, 4] + @test [nzrange(B,col) for col in axes(B,2)] == [1:1, 2:3, 4:5, 6:7, 8:8] + @test nonzeros(B) == [10, 0,10, 0,10, 0,10, 0] + # for arrays with overlapping indices, merge should only add the extra ones + A[diagind(A,2)] .= 5 + SparseArrays.mergeinds!(B, A) + @test rowvals(B) == [2, 1,3, 1,2,4, 2,3,5, 3,4] + @test [nzrange(B,col) for col in axes(B,2)] == [1:1, 2:3, 4:6, 7:9, 10:11] + @test nonzeros(B) == [10, 0,10, 0,0,10, 0,0,10, 0,0] + end for sA2 in (similar(sA), sprand(size(sA)..., 0.1)) @testset for (alpha, beta) in [(true, false), (true, true), (2,3)] nonzeros(sA2) .= 1 From 8020cdcb48ed318895f77e206d8fbb11f0866964 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Wed, 12 Mar 2025 16:11:33 +0530 Subject: [PATCH 10/12] Fix pre-multiplication --- src/linalg.jl | 2 +- test/linalg.jl | 8 +++----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/src/linalg.jl b/src/linalg.jl index d1350297..2db355b0 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -2000,7 +2000,7 @@ function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AbstractSparseMatrixCS row = rowvals(C)[p] row_exists, row_ind_A = rowcheck_index(A, row, col) if row_exists - @inbounds Cnzval[p] = b[row_ind_A] * Anzval[row_ind_A] * alpha + Cnzval[p] * beta + @inbounds Cnzval[p] = b[row] * Anzval[row_ind_A] * alpha + Cnzval[p] * beta else @inbounds Cnzval[p] = Cnzval[p] * beta end diff --git a/test/linalg.jl b/test/linalg.jl index fe73662d..a0fa9d74 100644 --- a/test/linalg.jl +++ b/test/linalg.jl @@ -606,14 +606,12 @@ end @test nonzeros(B) == [10, 0,10, 0,0,10, 0,0,10, 0,0] end for sA2 in (similar(sA), sprand(size(sA)..., 0.1)) + nonzeros(sA2) .= 1 @testset for (alpha, beta) in [(true, false), (true, true), (2,3)] - nonzeros(sA2) .= 1 - sA3 = copy(sA2) D = Diagonal(rand(size(sA,2))) - @test mul!(sA2, sA, D, alpha, beta) ≈ dA * D * alpha + sA3 * beta - nonzeros(sA2) .= 1 + @test mul!(copy(sA2), sA, D, alpha, beta) ≈ dA * D * alpha + sA2 * beta D = Diagonal(rand(size(sA,1))) - @test mul!(sA2, D, sA, alpha, beta) ≈ D * dA * alpha + sA3 * beta + @test mul!(copy(sA2), D, sA, alpha, beta) ≈ D * dA * alpha + sA2 * beta end end end From a83d27c6830f4495e35bc1d411e820735e5e84b3 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Wed, 12 Mar 2025 16:26:53 +0530 Subject: [PATCH 11/12] Add comments and a docstring to `mergeinds!` --- src/linalg.jl | 42 ++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 40 insertions(+), 2 deletions(-) diff --git a/src/linalg.jl b/src/linalg.jl index 2db355b0..120f59b2 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -1908,6 +1908,42 @@ If `row_exists` is `false`, the `row_ind` is the index where the value should be row_exists, row_ind end +""" + mergeinds!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC) + +Update `C` to contain stored values corresponding to the stored indices of `A`. +Stored indices common to `C` and `A` are not touched. Indices of `A` at which +`C` did not have a stored value are populated with zeros after the call. + +# Examples +```jldoctest +julia> A = spzeros(3,3); + +julia> A[4:4:8] .= 1; + +julia> A +3×3 SparseMatrixCSC{Float64, Int64} with 2 stored entries: + ⋅ 1.0 ⋅ + ⋅ ⋅ 1.0 + ⋅ ⋅ ⋅ + +julia> C = spzeros(3,3); + +julia> C[2:4:6] .= 2; + +julia> C +3×3 SparseMatrixCSC{Float64, Int64} with 2 stored entries: + ⋅ ⋅ ⋅ + 2.0 ⋅ ⋅ + ⋅ 2.0 ⋅ + +julia> SparseArrays.mergeinds!(C, A) +3×3 SparseMatrixCSC{Float64, Int64} with 4 stored entries: + ⋅ 0.0 ⋅ + 2.0 ⋅ 0.0 + ⋅ 2.0 ⋅ +``` +""" function mergeinds!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC) C_colptr = getcolptr(C) for col in axes(A,2) @@ -1958,10 +1994,11 @@ function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, D::Diagona mergeinds!(C, A) for col in axes(C,2), p in nzrange(C, col) row = rowvals(C)[p] + # check if the index (row, col) is stored in A row_exists, row_ind_A = rowcheck_index(A, row, col) if row_exists @inbounds Cnzval[p] = Anzval[row_ind_A] * b[col] * alpha + Cnzval[p] * beta - else + else # A[row,col] == 0 @inbounds Cnzval[p] = Cnzval[p] * beta end end @@ -1998,10 +2035,11 @@ function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AbstractSparseMatrixCS mergeinds!(C, A) for col in axes(C,2), p in nzrange(C, col) row = rowvals(C)[p] + # check if the index (row, col) is stored in A row_exists, row_ind_A = rowcheck_index(A, row, col) if row_exists @inbounds Cnzval[p] = b[row] * Anzval[row_ind_A] * alpha + Cnzval[p] * beta - else + else # A[row,col] == 0 @inbounds Cnzval[p] = Cnzval[p] * beta end end From 07da37559f3011ae793b98efd1f2eed3e5b8a172 Mon Sep 17 00:00:00 2001 From: Jishnu Bhattacharya Date: Mon, 31 Mar 2025 19:30:13 +0530 Subject: [PATCH 12/12] Skip multiplication with alpha if isone --- src/linalg.jl | 52 +++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 42 insertions(+), 10 deletions(-) diff --git a/src/linalg.jl b/src/linalg.jl index 120f59b2..82770210 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -1982,12 +1982,24 @@ function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, D::Diagona identical_nzinds || copyinds!(C, A, copy_rows = !rows_match, copy_cols = !cols_match) resize!(Cnzval, length(Anzval)) if beta_is_zero - for col in axes(A,2), p in nzrange(A, col) - @inbounds Cnzval[p] = Anzval[p] * b[col] * alpha + if isone(alpha) + for col in axes(A,2), p in nzrange(A, col) + @inbounds Cnzval[p] = Anzval[p] * b[col] + end + else + for col in axes(A,2), p in nzrange(A, col) + @inbounds Cnzval[p] = Anzval[p] * b[col] * alpha + end end else - for col in axes(A,2), p in nzrange(A, col) - @inbounds Cnzval[p] = Anzval[p] * b[col] * alpha + Cnzval[p] * beta + if isone(alpha) + for col in axes(A,2), p in nzrange(A, col) + @inbounds Cnzval[p] = Anzval[p] * b[col] + Cnzval[p] * beta + end + else + for col in axes(A,2), p in nzrange(A, col) + @inbounds Cnzval[p] = Anzval[p] * b[col] * alpha + Cnzval[p] * beta + end end end else @@ -1997,7 +2009,11 @@ function mul!(C::AbstractSparseMatrixCSC, A::AbstractSparseMatrixCSC, D::Diagona # check if the index (row, col) is stored in A row_exists, row_ind_A = rowcheck_index(A, row, col) if row_exists - @inbounds Cnzval[p] = Anzval[row_ind_A] * b[col] * alpha + Cnzval[p] * beta + if isone(alpha) + @inbounds Cnzval[p] = Anzval[row_ind_A] * b[col] + Cnzval[p] * beta + else + @inbounds Cnzval[p] = Anzval[row_ind_A] * b[col] * alpha + Cnzval[p] * beta + end else # A[row,col] == 0 @inbounds Cnzval[p] = Cnzval[p] * beta end @@ -2023,12 +2039,24 @@ function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AbstractSparseMatrixCS identical_nzinds || copyinds!(C, A, copy_rows = !rows_match, copy_cols = !cols_match) resize!(Cnzval, length(Anzval)) if beta_is_zero - for col in axes(A,2), p in nzrange(A, col) - @inbounds Cnzval[p] = b[Arowval[p]] * Anzval[p] * alpha + if isone(alpha) + for col in axes(A,2), p in nzrange(A, col) + @inbounds Cnzval[p] = b[Arowval[p]] * Anzval[p] + end + else + for col in axes(A,2), p in nzrange(A, col) + @inbounds Cnzval[p] = b[Arowval[p]] * Anzval[p] * alpha + end end else - for col in axes(A,2), p in nzrange(A, col) - @inbounds Cnzval[p] = b[Arowval[p]] * Anzval[p] * alpha + Cnzval[p] * beta + if isone(alpha) + for col in axes(A,2), p in nzrange(A, col) + @inbounds Cnzval[p] = b[Arowval[p]] * Anzval[p] + Cnzval[p] * beta + end + else + for col in axes(A,2), p in nzrange(A, col) + @inbounds Cnzval[p] = b[Arowval[p]] * Anzval[p] * alpha + Cnzval[p] * beta + end end end else @@ -2038,7 +2066,11 @@ function mul!(C::AbstractSparseMatrixCSC, D::Diagonal, A::AbstractSparseMatrixCS # check if the index (row, col) is stored in A row_exists, row_ind_A = rowcheck_index(A, row, col) if row_exists - @inbounds Cnzval[p] = b[row] * Anzval[row_ind_A] * alpha + Cnzval[p] * beta + if isone(alpha) + @inbounds Cnzval[p] = b[row] * Anzval[row_ind_A] + Cnzval[p] * beta + else + @inbounds Cnzval[p] = b[row] * Anzval[row_ind_A] * alpha + Cnzval[p] * beta + end else # A[row,col] == 0 @inbounds Cnzval[p] = Cnzval[p] * beta end