From 508e754523cbc3d5a1e17d483661d3a00db7f7ce Mon Sep 17 00:00:00 2001 From: Lilith Hafner Date: Tue, 5 Sep 2023 10:51:50 -0500 Subject: [PATCH 1/2] switch from internal searchsorted* to views --- src/linalg.jl | 42 +++++++++++++++++++++--------------------- src/solvers/cholmod.jl | 3 +-- src/sparsematrix.jl | 30 +++++++++++++++--------------- src/sparsevector.jl | 8 ++++---- 4 files changed, 41 insertions(+), 42 deletions(-) diff --git a/src/linalg.jl b/src/linalg.jl index f315300f..cb9093f3 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -453,7 +453,7 @@ function LinearAlgebra.generic_trimatmul!(C::StridedVecOrMat, uploc, isunitc, tf i1 = ia[j] i2 = ia[j + 1] - 1 done = unit - + bj = B[joff + j] for ii = i1:i2 jai = ja[ii] @@ -483,7 +483,7 @@ function LinearAlgebra.generic_trimatmul!(C::StridedVecOrMat, uploc, isunitc, tf i2 = ia[j + 1] - 1 akku = Z j0 = !unit ? j : j - 1 - + # loop through column j of A - only structural non-zeros for ii = i1:i2 jai = ja[ii] @@ -509,7 +509,7 @@ function LinearAlgebra.generic_trimatmul!(C::StridedVecOrMat, uploc, isunitc, tf i1 = ia[j] i2 = ia[j + 1] - 1 done = unit - + bj = B[joff + j] for ii = i2:-1:i1 jai = ja[ii] @@ -539,7 +539,7 @@ function LinearAlgebra.generic_trimatmul!(C::StridedVecOrMat, uploc, isunitc, tf i2 = ia[j + 1] - 1 akku = Z j0 = !unit ? j : j + 1 - + # loop through column j of A - only structural non-zeros for ii = i2:-1:i1 jai = ja[ii] @@ -582,7 +582,7 @@ function LinearAlgebra.generic_trimatmul!(C::StridedVecOrMat, uploc, isunitc, :: i1 = ia[j] i2 = ia[j + 1] - 1 done = unit - + bj = B[joff + j] for ii = i1:i2 jai = ja[ii] @@ -610,7 +610,7 @@ function LinearAlgebra.generic_trimatmul!(C::StridedVecOrMat, uploc, isunitc, :: i1 = ia[j] i2 = ia[j + 1] - 1 done = unit - + bj = B[joff + j] for ii = i2:-1:i1 jai = ja[ii] @@ -664,7 +664,7 @@ function LinearAlgebra.generic_trimatdiv!(C::StridedVecOrMat, uploc, isunitc, tf i2 = ia[j + 1] - one(eltype(ia)) # find diagonal element - ii = searchsortedfirst(ja, j, i1, i2, Base.Order.Forward) + ii = searchsortedfirst(view(ja, i1:i2), j) - i1 + 1 jai = ii > i2 ? zero(eltype(ja)) : ja[ii] cj = C[j,k] @@ -693,7 +693,7 @@ function LinearAlgebra.generic_trimatdiv!(C::StridedVecOrMat, uploc, isunitc, tf i2 = ia[j + 1] - 1 akku = B[j,k] done = false - + # loop through column j of A - only structural non-zeros for ii = i2:-1:i1 jai = ja[ii] @@ -721,11 +721,11 @@ function LinearAlgebra.generic_trimatdiv!(C::StridedVecOrMat, uploc, isunitc, tf for j = nrowB:-1:1 i1 = ia[j] i2 = ia[j + 1] - one(eltype(ia)) - + # find diagonal element - ii = searchsortedlast(ja, j, i1, i2, Base.Order.Forward) + ii = searchsortedlast(view(ja, i1:i2), j) - i1 + 1 jai = ii < i1 ? zero(eltype(ja)) : ja[ii] - + cj = C[j,k] # check for zero pivot and divide with pivot if jai == j @@ -737,7 +737,7 @@ function LinearAlgebra.generic_trimatdiv!(C::StridedVecOrMat, uploc, isunitc, tf elseif !unit throw(LinearAlgebra.SingularException(j)) end - + # update remaining part for i = ii:-1:i1 C[ja[i],k] -= cj * LinearAlgebra._ustrip(aa[i]) @@ -752,7 +752,7 @@ function LinearAlgebra.generic_trimatdiv!(C::StridedVecOrMat, uploc, isunitc, tf i2 = ia[j + 1] - 1 akku = B[j,k] done = false - + # loop through column j of A - only structural non-zeros for ii = i1:i2 jai = ja[ii] @@ -801,7 +801,7 @@ function LinearAlgebra.generic_trimatdiv!(C::StridedVecOrMat, uploc, isunitc, :: i2 = ia[j + 1] - one(eltype(ia)) # find diagonal element - ii = searchsortedfirst(ja, j, i1, i2, Base.Order.Forward) + ii = searchsortedfirst(view(ja, i1:i2), j) - i1 + 1 jai = ii > i2 ? zero(eltype(ja)) : ja[ii] cj = C[j,k] @@ -828,11 +828,11 @@ function LinearAlgebra.generic_trimatdiv!(C::StridedVecOrMat, uploc, isunitc, :: for j = nrowB:-1:1 i1 = ia[j] i2 = ia[j + 1] - one(eltype(ia)) - + # find diagonal element - ii = searchsortedlast(ja, j, i1, i2, Base.Order.Forward) + ii = searchsortedlast(view(ja, i1:i2), j) - i1 + 1 jai = ii < i1 ? zero(eltype(ja)) : ja[ii] - + cj = C[j,k] # check for zero pivot and divide with pivot if jai == j @@ -844,7 +844,7 @@ function LinearAlgebra.generic_trimatdiv!(C::StridedVecOrMat, uploc, isunitc, :: elseif !unit throw(LinearAlgebra.SingularException(j)) end - + # update remaining part for i = ii:-1:i1 C[ja[i],k] -= cj * LinearAlgebra._ustrip(conj(aa[i])) @@ -904,13 +904,13 @@ end function nzrangeup(A, i, excl=false) r = nzrange(A, i); r1 = r.start; r2 = r.stop rv = rowvals(A) - @inbounds r2 < r1 || rv[r2] <= i - excl ? r : r1:searchsortedlast(rv, i - excl, r1, r2, Forward) + @inbounds r2 < r1 || rv[r2] <= i - excl ? r : r1:(searchsortedlast(view(rv, r1:r2), i - excl) + r1-1) end # row range from diagonal (included if excl=false) to end function nzrangelo(A, i, excl=false) r = nzrange(A, i); r1 = r.start; r2 = r.stop rv = rowvals(A) - @inbounds r2 < r1 || rv[r1] >= i + excl ? r : searchsortedfirst(rv, i + excl, r1, r2, Forward):r2 + @inbounds r2 < r1 || rv[r1] >= i + excl ? r : (searchsortedfirst(view(rv, r1:r2), i + excl) + r1-1):r2 end dot(x::AbstractVector, A::RealHermSymComplexHerm{<:Any,<:AbstractSparseMatrixCSC}, y::AbstractVector) = @@ -985,7 +985,7 @@ function _dot(x::SparseVector, A::AbstractSparseMatrixCSC, y::SparseVector, rang r1 = Int(Acolptr[i]) r2 = Int(Acolptr[i+1]-1) r1 > r2 && continue - r1 = searchsortedfirst(Arowval, i, r1, r2, Forward) + r1 += searchsortedfirst(view(Arowval, r1:r2), i) - 1 ((r1 > r2) || (Arowval[r1] != i)) && continue r += dot(x[i], diagop(Anzval[r1]), y[i]) end diff --git a/src/solvers/cholmod.jl b/src/solvers/cholmod.jl index 9c832ef7..78247620 100644 --- a/src/solvers/cholmod.jl +++ b/src/solvers/cholmod.jl @@ -1180,8 +1180,7 @@ function getindex(A::Sparse{T}, i0::Integer, i1::Integer) where T r1 = Int(unsafe_load(s.p, i1) + 1) r2 = Int(unsafe_load(s.p, i1 + 1)) (r1 > r2) && return zero(T) - r1 = Int(searchsortedfirst(unsafe_wrap(Array, s.i, (s.nzmax,), own = false), - i0 - 1, r1, r2, Base.Order.Forward)) + r1 += Int(searchsortedfirst(view(unsafe_wrap(Array, s.i, (s.nzmax,), own = false), r1:r2), i0 - 1) - 1) ((r1 > r2) || (unsafe_load(s.i, r1) + 1 != i0)) ? zero(T) : unsafe_load(Ptr{T}(s.x), r1) end diff --git a/src/sparsematrix.jl b/src/sparsematrix.jl index 165a7046..c03a5809 100644 --- a/src/sparsematrix.jl +++ b/src/sparsematrix.jl @@ -1939,11 +1939,11 @@ function _sparse_findnextnz(m::AbstractSparseMatrixCSC, ij::CartesianIndex{2}) col > size(m, 2) && return nothing lo, hi = getcolptr(m)[col], getcolptr(m)[col+1] - n = searchsortedfirst(rowvals(m), row, lo, hi-1, Base.Order.Forward) + n = searchsortedfirst(view(rowvals(m), lo:hi-1), row) + lo - 1 if lo <= n <= hi-1 return CartesianIndex(rowvals(m)[n], col) end - nextcol = searchsortedfirst(getcolptr(m), hi + 1, col + 1, length(getcolptr(m)), Base.Order.Forward) + nextcol = searchsortedfirst(view(getcolptr(m), col+1:length(getcolptr(m))), hi + 1) + col nextcol > length(getcolptr(m)) && return nothing nextlo = getcolptr(m)[nextcol-1] return CartesianIndex(rowvals(m)[nextlo], nextcol - 1) @@ -1954,11 +1954,11 @@ function _sparse_findprevnz(m::AbstractSparseMatrixCSC, ij::CartesianIndex{2}) iszero(col) && return nothing lo, hi = getcolptr(m)[col], getcolptr(m)[col+1] - n = searchsortedlast(rowvals(m), row, lo, hi-1, Base.Order.Forward) + n = searchsortedlast(view(rowvals(m), lo:hi-1), row) - lo + 1 if lo <= n <= hi-1 return CartesianIndex(rowvals(m)[n], col) end - prevcol = searchsortedlast(getcolptr(m), lo - 1, 1, col - 1, Base.Order.Forward) + prevcol = searchsortedlast(view(getcolptr(m), 1:col-1), lo - 1) prevcol < 1 && return nothing prevhi = getcolptr(m)[prevcol+1] return CartesianIndex(rowvals(m)[prevhi-1], prevcol) @@ -2564,8 +2564,8 @@ function _findz(A::AbstractSparseMatrixCSC{Tv,Ti}, rows=1:size(A, 1), cols=1:siz r1::Int = colptr[col] r2::Int = colptr[col+1] - 1 if !allrows && (r1 <= r2) - r1 = searchsortedfirst(rowval, rowmin, r1, r2, Forward) - (r1 <= r2 ) && (r2 = searchsortedlast(rowval, rowmax, r1, r2, Forward)) + r1 += searchsortedfirst(view(rowval, r1:r2), rowmin) - 1 + (r1 <= r2 ) && (r2 = searchsortedlast(view(rowval, r1:r2), rowmax) + r1 - 1) end row = rowmin while (r1 <= r2) && (row == rowval[r1]) && _isnotzero(nzval[r1]) @@ -2677,7 +2677,7 @@ end r1 = Int(@inbounds getcolptr(A)[i1]) r2 = Int(@inbounds getcolptr(A)[i1+1]-1) (r1 > r2) && return zero(T) - r1 = searchsortedfirst(rowvals(A), i0, r1, r2, Forward) + r1 = searchsortedfirst(view(rowvals(A), r1:r2), i0) + r1 - 1 ((r1 > r2) || (rowvals(A)[r1] != i0)) ? zero(T) : nonzeros(A)[r1] end @@ -2813,7 +2813,7 @@ function getindex_I_sorted_bsearch_A(A::AbstractSparseMatrixCSC{Tv,Ti}, I::Abstr rowI = I[ptrI] ptrI += 1 (rowvalA[ptrA] > rowI) && continue - ptrA = searchsortedfirst(rowvalA, rowI, ptrA, stopA, Base.Order.Forward) + ptrA += searchsortedfirst(view(rowvalA, ptrA:stopA), rowI) - 1 (ptrA <= stopA) || break if rowvalA[ptrA] == rowI ptrS += 1 @@ -2837,7 +2837,7 @@ function getindex_I_sorted_bsearch_A(A::AbstractSparseMatrixCSC{Tv,Ti}, I::Abstr while ptrI <= nI rowI = I[ptrI] if rowvalA[ptrA] <= rowI - ptrA = searchsortedfirst(rowvalA, rowI, ptrA, stopA, Base.Order.Forward) + ptrA += searchsortedfirst(view(rowvalA, ptrA:stopA), rowI) - 1 (ptrA <= stopA) || break if rowvalA[ptrA] == rowI rowvalS[ptrS] = ptrI @@ -2941,7 +2941,7 @@ function getindex_I_sorted_bsearch_I(A::AbstractSparseMatrixCSC{Tv,Ti}, I::Abstr @inbounds for j = 1:m cval = cacheI[j] (cval == 0) && continue - ptrI = searchsortedfirst(I, j, ptrI, nI, Base.Order.Forward) + ptrI += searchsortedfirst(view(I, ptrI:nI), j) - 1 cacheI[j] = ptrI while ptrI <= nI && I[ptrI] == j ptrS += cval @@ -3129,7 +3129,7 @@ function _setindex_scalar!(A::AbstractSparseMatrixCSC{Tv,Ti}, _v, _i::Integer, _ end coljfirstk = Int(getcolptr(A)[j]) coljlastk = Int(getcolptr(A)[j+1] - 1) - searchk = searchsortedfirst(rowvals(A), i, coljfirstk, coljlastk, Base.Order.Forward) + searchk = searchsortedfirst(view(rowvals(A), coljfirstk:coljlastk), i) + coljfirstk - 1 if searchk <= coljlastk && rowvals(A)[searchk] == i # Column j contains entry A[i,j]. Update and return nonzeros(A)[searchk] = v @@ -3491,7 +3491,7 @@ function setindex!(A::AbstractSparseMatrixCSC, x::AbstractArray, I::AbstractMatr xidx += 1 if r1 <= r2 - copylen = searchsortedfirst(rowvalA, row, r1, r2, Forward) - r1 + copylen = searchsortedfirst(view(rowvalA, r1:r2), row) - 2r1 + 1 if (copylen > 0) if (nadd > 0) copyto!(rowvalB, bidx, rowvalA, r1, copylen) @@ -3621,7 +3621,7 @@ function setindex!(A::AbstractSparseMatrixCSC, x::AbstractArray, Ix::AbstractVec end if r1 <= r2 - copylen = searchsortedfirst(rowvalA, row, r1, r2, Forward) - r1 + copylen = searchsortedfirst(view(rowvalA, r1:r2), row) - 2r1 + 1 if (copylen > 0) if (nadd > 0) copyto!(rowvalB, bidx, rowvalA, r1, copylen) @@ -3705,7 +3705,7 @@ function dropstored!(A::AbstractSparseMatrixCSC, i::Integer, j::Integer) end coljfirstk = Int(getcolptr(A)[j]) coljlastk = Int(getcolptr(A)[j+1] - 1) - searchk = searchsortedfirst(rowvals(A), i, coljfirstk, coljlastk, Base.Order.Forward) + searchk = searchsortedfirst(view(rowvals(A), coljfirstk:coljlastk), i) - coljfirstk + 1 if searchk <= coljlastk && rowvals(A)[searchk] == i # Entry A[i,j] is stored. Drop and return. deleteat!(rowvals(A), searchk) @@ -4222,7 +4222,7 @@ function diag(A::AbstractSparseMatrixCSC{Tv,Ti}, d::Integer=0) where {Tv,Ti} r1 = Int(getcolptr(A)[c]) r2 = Int(getcolptr(A)[c+1]-1) r1 > r2 && continue - r1 = searchsortedfirst(rowvals(A), r, r1, r2, Forward) + r1 += searchsortedfirst(view(rowvals(A), r1:r2), r) + 1 ((r1 > r2) || (rowvals(A)[r1] != r)) && continue push!(ind, i) push!(val, nonzeros(A)[r1]) diff --git a/src/sparsevector.jl b/src/sparsevector.jl index 9fc8bec8..aa5b959b 100644 --- a/src/sparsevector.jl +++ b/src/sparsevector.jl @@ -639,8 +639,8 @@ function getindex(x::AbstractSparseMatrixCSC, I::AbstractUnitRange, j::Integer) c1 = convert(Int, getcolptr(x)[j]) c2 = convert(Int, getcolptr(x)[j+1]) - 1 # Restrict to the selected rows - r1 = searchsortedfirst(rowvals(x), first(I), c1, c2, Forward) - r2 = searchsortedlast(rowvals(x), last(I), c1, c2, Forward) + r1 = searchsortedfirst(view(rowvals(x), c1:c2), first(I)) + c1 - 1 + r2 = searchsortedlast(view(rowvals(x), c1:c2), last(I)) + c1 - 1 return @if_move_fixed x SparseVector(length(I), [rowvals(x)[i] - first(I) + 1 for i = r1:r2], nonzeros(x)[r1:r2]) end @@ -670,7 +670,7 @@ function Base.getindex(A::AbstractSparseMatrixCSC{Tv,Ti}, i::Integer, J::Abstrac stopA = Int(colptrA[col+1]-1) if ptrA <= stopA if rowvalA[ptrA] <= rowI - ptrA = searchsortedfirst(rowvalA, rowI, ptrA, stopA, Base.Order.Forward) + ptrA += searchsortedfirst(view(rowvalA, ptrA:stopA), rowI) - 1 if ptrA <= stopA && rowvalA[ptrA] == rowI push!(nzinds, j) push!(nzvals, nzvalA[ptrA]) @@ -959,7 +959,7 @@ function getindex(x::AbstractSparseVector{Tv,Ti}, I::AbstractUnitRange) where {T # locate the first j0, s.t. xnzind[j0] >= i0 j0 = searchsortedfirst(xnzind, i0) # locate the last j1, s.t. xnzind[j1] <= i1 - j1 = searchsortedlast(xnzind, i1, j0, m, Forward) + j1 = searchsortedlast(view(xnzind, j0:m), i1) + j0 - 1 # compute the number of non-zeros jrgn = j0:j1 From cf4aac57a888897eae5a4725897b48f2cf8500a6 Mon Sep 17 00:00:00 2001 From: Lilith Hafner Date: Tue, 5 Sep 2023 16:33:04 -0500 Subject: [PATCH 2/2] fixups --- src/linalg.jl | 8 ++++---- src/sparsematrix.jl | 10 +++++----- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/linalg.jl b/src/linalg.jl index cb9093f3..7670ce61 100644 --- a/src/linalg.jl +++ b/src/linalg.jl @@ -664,7 +664,7 @@ function LinearAlgebra.generic_trimatdiv!(C::StridedVecOrMat, uploc, isunitc, tf i2 = ia[j + 1] - one(eltype(ia)) # find diagonal element - ii = searchsortedfirst(view(ja, i1:i2), j) - i1 + 1 + ii = searchsortedfirst(view(ja, i1:i2), j) + i1 - 1 jai = ii > i2 ? zero(eltype(ja)) : ja[ii] cj = C[j,k] @@ -723,7 +723,7 @@ function LinearAlgebra.generic_trimatdiv!(C::StridedVecOrMat, uploc, isunitc, tf i2 = ia[j + 1] - one(eltype(ia)) # find diagonal element - ii = searchsortedlast(view(ja, i1:i2), j) - i1 + 1 + ii = searchsortedlast(view(ja, i1:i2), j) + i1 - 1 jai = ii < i1 ? zero(eltype(ja)) : ja[ii] cj = C[j,k] @@ -801,7 +801,7 @@ function LinearAlgebra.generic_trimatdiv!(C::StridedVecOrMat, uploc, isunitc, :: i2 = ia[j + 1] - one(eltype(ia)) # find diagonal element - ii = searchsortedfirst(view(ja, i1:i2), j) - i1 + 1 + ii = searchsortedfirst(view(ja, i1:i2), j) + i1 - 1 jai = ii > i2 ? zero(eltype(ja)) : ja[ii] cj = C[j,k] @@ -830,7 +830,7 @@ function LinearAlgebra.generic_trimatdiv!(C::StridedVecOrMat, uploc, isunitc, :: i2 = ia[j + 1] - one(eltype(ia)) # find diagonal element - ii = searchsortedlast(view(ja, i1:i2), j) - i1 + 1 + ii = searchsortedlast(view(ja, i1:i2), j) + i1 - 1 jai = ii < i1 ? zero(eltype(ja)) : ja[ii] cj = C[j,k] diff --git a/src/sparsematrix.jl b/src/sparsematrix.jl index c03a5809..e2dfacce 100644 --- a/src/sparsematrix.jl +++ b/src/sparsematrix.jl @@ -1954,7 +1954,7 @@ function _sparse_findprevnz(m::AbstractSparseMatrixCSC, ij::CartesianIndex{2}) iszero(col) && return nothing lo, hi = getcolptr(m)[col], getcolptr(m)[col+1] - n = searchsortedlast(view(rowvals(m), lo:hi-1), row) - lo + 1 + n = searchsortedlast(view(rowvals(m), lo:hi-1), row) + lo - 1 if lo <= n <= hi-1 return CartesianIndex(rowvals(m)[n], col) end @@ -3491,7 +3491,7 @@ function setindex!(A::AbstractSparseMatrixCSC, x::AbstractArray, I::AbstractMatr xidx += 1 if r1 <= r2 - copylen = searchsortedfirst(view(rowvalA, r1:r2), row) - 2r1 + 1 + copylen = searchsortedfirst(view(rowvalA, r1:r2), row) - 1 if (copylen > 0) if (nadd > 0) copyto!(rowvalB, bidx, rowvalA, r1, copylen) @@ -3621,7 +3621,7 @@ function setindex!(A::AbstractSparseMatrixCSC, x::AbstractArray, Ix::AbstractVec end if r1 <= r2 - copylen = searchsortedfirst(view(rowvalA, r1:r2), row) - 2r1 + 1 + copylen = searchsortedfirst(view(rowvalA, r1:r2), row) - 1 if (copylen > 0) if (nadd > 0) copyto!(rowvalB, bidx, rowvalA, r1, copylen) @@ -3705,7 +3705,7 @@ function dropstored!(A::AbstractSparseMatrixCSC, i::Integer, j::Integer) end coljfirstk = Int(getcolptr(A)[j]) coljlastk = Int(getcolptr(A)[j+1] - 1) - searchk = searchsortedfirst(view(rowvals(A), coljfirstk:coljlastk), i) - coljfirstk + 1 + searchk = searchsortedfirst(view(rowvals(A), coljfirstk:coljlastk), i) + coljfirstk - 1 if searchk <= coljlastk && rowvals(A)[searchk] == i # Entry A[i,j] is stored. Drop and return. deleteat!(rowvals(A), searchk) @@ -4222,7 +4222,7 @@ function diag(A::AbstractSparseMatrixCSC{Tv,Ti}, d::Integer=0) where {Tv,Ti} r1 = Int(getcolptr(A)[c]) r2 = Int(getcolptr(A)[c+1]-1) r1 > r2 && continue - r1 += searchsortedfirst(view(rowvals(A), r1:r2), r) + 1 + r1 += searchsortedfirst(view(rowvals(A), r1:r2), r) - 1 ((r1 > r2) || (rowvals(A)[r1] != r)) && continue push!(ind, i) push!(val, nonzeros(A)[r1])