Skip to content

Commit 2627d10

Browse files
committed
* Make length(A.nzval)==nnz(A) and add strict buffer checking #30662.
* Add sizehint!(::SparseMatrixCSC, args...), * Fix illegal SparseMatrixCSC construction in cholmod and linalg. * Remove tests targetting now illegal buffers * Fix invalid buffer creation in kron and more
1 parent 3d08f82 commit 2627d10

File tree

9 files changed

+212
-203
lines changed

9 files changed

+212
-203
lines changed

NEWS.md

+1
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,7 @@ Standard library changes
111111

112112
#### SparseArrays
113113

114+
* new `sizehint!(::SparseMatrixCSC, ::Integer)` method ([#30676]).
114115

115116
#### Dates
116117

stdlib/SparseArrays/src/higherorderfns.jl

+29-33
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import Base: map, map!, broadcast, copy, copyto!, _extrema_dims, _extrema_itr
88

99
using Base: front, tail, to_shape
1010
using ..SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseVector, AbstractSparseMatrixCSC,
11-
AbstractSparseMatrix, AbstractSparseArray, indtype, nnz, nzrange,
11+
AbstractSparseMatrix, AbstractSparseArray, indtype, nnz, nzrange, spzeros,
1212
SparseVectorUnion, AdjOrTransSparseVectorUnion, nonzeroinds, nonzeros, rowvals, getcolptr, widelength
1313
using Base.Broadcast: BroadcastStyle, Broadcasted, flatten
1414
using LinearAlgebra
@@ -132,12 +132,17 @@ function trimstorage!(A::SparseVecOrMat, maxstored)
132132
resize!(storedvals(A), maxstored)
133133
return maxstored
134134
end
135+
135136
function expandstorage!(A::SparseVecOrMat, maxstored)
136-
length(storedinds(A)) < maxstored && resize!(storedinds(A), maxstored)
137-
length(storedvals(A)) < maxstored && resize!(storedvals(A), maxstored)
137+
if length(storedinds(A)) < maxstored
138+
resize!(storedinds(A), maxstored)
139+
resize!(storedvals(A), maxstored)
140+
end
138141
return maxstored
139142
end
140143

144+
_checkbuffers(S::SparseMatrixCSC) = (@assert length(getcolptr(S)) == size(S, 2) + 1 && getcolptr(S)[end] - 1 == length(rowvals(S)) == length(nonzeros(S)); S)
145+
_checkbuffers(S::SparseVector) = (@assert length(storedvals(S)) == length(storedinds(S)); S)
141146

142147
# (2) map[!] entry points
143148
map(f::Tf, A::SparseVector) where {Tf} = _noshapecheck_map(f, A)
@@ -181,7 +186,7 @@ copy(bc::SpBroadcasted1) = _noshapecheck_map(bc.f, bc.args[1])
181186
storedvals(C)[1] = fofnoargs
182187
broadcast!(f, view(storedvals(C), 2:length(storedvals(C))))
183188
end
184-
return C
189+
return _checkbuffers(C)
185190
end
186191

187192
function _diffshape_broadcast(f::Tf, A::SparseVecOrMat, Bs::Vararg{SparseVecOrMat,N}) where {Tf,N}
@@ -224,22 +229,13 @@ _maxnnzfrom(shape::NTuple{2}, A::AbstractSparseMatrixCSC) = nnz(A) * div(shape[1
224229
@inline _unchecked_maxnnzbcres(shape, As...) = _unchecked_maxnnzbcres(shape, As)
225230
@inline _checked_maxnnzbcres(shape::NTuple{1}, As...) = shape[1] != 0 ? _unchecked_maxnnzbcres(shape, As) : 0
226231
@inline _checked_maxnnzbcres(shape::NTuple{2}, As...) = shape[1] != 0 && shape[2] != 0 ? _unchecked_maxnnzbcres(shape, As) : 0
227-
@inline function _allocres(shape::NTuple{1}, indextype, entrytype, maxnnz)
228-
storedinds = Vector{indextype}(undef, maxnnz)
229-
storedvals = Vector{entrytype}(undef, maxnnz)
230-
return SparseVector(shape..., storedinds, storedvals)
231-
end
232-
@inline function _allocres(shape::NTuple{2}, indextype, entrytype, maxnnz)
233-
pointers = ones(indextype, shape[2] + 1)
234-
storedinds = Vector{indextype}(undef, maxnnz)
235-
storedvals = Vector{entrytype}(undef, maxnnz)
236-
return SparseMatrixCSC(shape..., pointers, storedinds, storedvals)
237-
end
232+
@inline _allocres(shape::NTuple{1}, indextype, entrytype, maxnnz) = sizehint!(SparseVector(shape[1], indextype[], entrytype[]), maxnnz)
233+
@inline _allocres(shape::NTuple{2}, indextype, entrytype, maxnnz) = sizehint!(SparseMatrixCSC(shape..., fill(indextype(1), shape[2]+1), indextype[], entrytype[]), maxnnz)
238234

239235
# (4) _map_zeropres!/_map_notzeropres! specialized for a single sparse vector/matrix
240236
"Stores only the nonzero entries of `map(f, Array(A))` in `C`."
241237
function _map_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat) where Tf
242-
spaceC::Int = min(length(storedinds(C)), length(storedvals(C)))
238+
spaceC::Int = length(nonzeros(C))
243239
Ck = 1
244240
@inbounds for j in columns(C)
245241
setcolptr!(C, j, Ck)
@@ -255,7 +251,7 @@ function _map_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat) where Tf
255251
end
256252
@inbounds setcolptr!(C, numcols(C) + 1, Ck)
257253
trimstorage!(C, Ck - 1)
258-
return C
254+
return _checkbuffers(C)
259255
end
260256
"""
261257
Densifies `C`, storing `fillvalue` in place of each unstored entry in `A` and
@@ -274,7 +270,7 @@ function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMa
274270
end
275271
# NOTE: Combining the fill! above into the loop above to avoid multiple sweeps over /
276272
# nonsequential access of storedvals(C) does not appear to improve performance.
277-
return C
273+
return _checkbuffers(C)
278274
end
279275
# helper functions for these methods and some of those below
280276
@inline _densecoloffsets(A::SparseVector) = 0
@@ -297,7 +293,7 @@ end
297293

298294
# (5) _map_zeropres!/_map_notzeropres! specialized for a pair of sparse vectors/matrices
299295
function _map_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::SparseVecOrMat) where Tf
300-
spaceC::Int = min(length(storedinds(C)), length(storedvals(C)))
296+
spaceC::Int = length(nonzeros(C))
301297
rowsentinelA = convert(indtype(A), numrows(C) + 1)
302298
rowsentinelB = convert(indtype(B), numrows(C) + 1)
303299
Ck = 1
@@ -336,7 +332,7 @@ function _map_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::SparseVe
336332
end
337333
@inbounds setcolptr!(C, numcols(C) + 1, Ck)
338334
trimstorage!(C, Ck - 1)
339-
return C
335+
return _checkbuffers(C)
340336
end
341337
function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMat, B::SparseVecOrMat) where Tf
342338
# Build dense matrix structure in C, expanding storage if necessary
@@ -368,13 +364,13 @@ function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMa
368364
Cx != fillvalue && (storedvals(C)[jo + Ci] = Cx)
369365
end
370366
end
371-
return C
367+
return _checkbuffers(C)
372368
end
373369

374370

375371
# (6) _map_zeropres!/_map_notzeropres! for more than two sparse matrices / vectors
376372
function _map_zeropres!(f::Tf, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N}) where {Tf,N}
377-
spaceC::Int = min(length(storedinds(C)), length(storedvals(C)))
373+
spaceC::Int = length(nonzeros(C))
378374
rowsentinel = numrows(C) + 1
379375
Ck = 1
380376
stopks = _colstartind_all(1, As)
@@ -398,7 +394,7 @@ function _map_zeropres!(f::Tf, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N})
398394
end
399395
@inbounds setcolptr!(C, numcols(C) + 1, Ck)
400396
trimstorage!(C, Ck - 1)
401-
return C
397+
return _checkbuffers(C)
402398
end
403399
function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N}) where {Tf,N}
404400
# Build dense matrix structure in C, expanding storage if necessary
@@ -421,7 +417,7 @@ function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg{Spars
421417
activerow = min(rows...)
422418
end
423419
end
424-
return C
420+
return _checkbuffers(C)
425421
end
426422

427423
# helper methods for map/map! methods just above
@@ -462,7 +458,7 @@ end
462458
# (7) _broadcast_zeropres!/_broadcast_notzeropres! specialized for a single (input) sparse vector/matrix
463459
function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat) where Tf
464460
isempty(C) && return _finishempty!(C)
465-
spaceC::Int = min(length(storedinds(C)), length(storedvals(C)))
461+
spaceC::Int = length(nonzeros(C))
466462
# C and A cannot have the same shape, as we directed that case to map in broadcast's
467463
# entry point; here we need efficiently handle only heterogeneous C-A combinations where
468464
# one or both of C and A has at least one singleton dimension.
@@ -509,7 +505,7 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat) where
509505
end
510506
@inbounds setcolptr!(C, numcols(C) + 1, Ck)
511507
trimstorage!(C, Ck - 1)
512-
return C
508+
return _checkbuffers(C)
513509
end
514510
function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMat) where Tf
515511
# For information on this code, see comments in similar code in _broadcast_zeropres! above
@@ -540,14 +536,14 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
540536
end
541537
end
542538
end
543-
return C
539+
return _checkbuffers(C)
544540
end
545541

546542

547543
# (8) _broadcast_zeropres!/_broadcast_notzeropres! specialized for a pair of (input) sparse vectors/matrices
548544
function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::SparseVecOrMat) where Tf
549545
isempty(C) && return _finishempty!(C)
550-
spaceC::Int = min(length(storedinds(C)), length(storedvals(C)))
546+
spaceC::Int = length(nonzeros(C))
551547
rowsentinelA = convert(indtype(A), numrows(C) + 1)
552548
rowsentinelB = convert(indtype(B), numrows(C) + 1)
553549
# C, A, and B cannot all have the same shape, as we directed that case to map in broadcast's
@@ -711,7 +707,7 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::Sp
711707
end
712708
@inbounds setcolptr!(C, numcols(C) + 1, Ck)
713709
trimstorage!(C, Ck - 1)
714-
return C
710+
return _checkbuffers(C)
715711
end
716712
function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMat, B::SparseVecOrMat) where Tf
717713
# For information on this code, see comments in similar code in _broadcast_zeropres! above
@@ -810,7 +806,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
810806
end
811807
end
812808
end
813-
return C
809+
return _checkbuffers(C)
814810
end
815811
_finishempty!(C::SparseVector) = C
816812
_finishempty!(C::AbstractSparseMatrixCSC) = (fill!(getcolptr(C), 1); C)
@@ -861,7 +857,7 @@ end
861857
# (9) _broadcast_zeropres!/_broadcast_notzeropres! for more than two (input) sparse vectors/matrices
862858
function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N}) where {Tf,N}
863859
isempty(C) && return _finishempty!(C)
864-
spaceC::Int = min(length(storedinds(C)), length(storedvals(C)))
860+
spaceC::Int = length(nonzeros(C))
865861
expandsverts = _expandsvert_all(C, As)
866862
expandshorzs = _expandshorz_all(C, As)
867863
rowsentinel = numrows(C) + 1
@@ -909,7 +905,7 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, As::Vararg{SparseVecOrMa
909905
end
910906
@inbounds setcolptr!(C, numcols(C) + 1, Ck)
911907
trimstorage!(C, Ck - 1)
912-
return C
908+
return _checkbuffers(C)
913909
end
914910
function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N}) where {Tf,N}
915911
isempty(C) && return _finishempty!(C)
@@ -950,7 +946,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg
950946
end
951947
end
952948
end
953-
return C
949+
return _checkbuffers(C)
954950
end
955951

956952
# helper method for broadcast/broadcast! methods just above

stdlib/SparseArrays/src/linalg.jl

+18-31
Original file line numberDiff line numberDiff line change
@@ -935,16 +935,15 @@ function triu(S::AbstractSparseMatrixCSC{Tv,Ti}, k::Integer=0) where {Tv,Ti}
935935
end
936936
rowval = Vector{Ti}(undef, nnz)
937937
nzval = Vector{Tv}(undef, nnz)
938-
A = SparseMatrixCSC(m, n, colptr, rowval, nzval)
939938
for col = max(k+1,1) : n
940939
c1 = getcolptr(S)[col]
941-
for c2 in nzrange(A, col)
942-
rowvals(A)[c2] = rowvals(S)[c1]
943-
nonzeros(A)[c2] = nonzeros(S)[c1]
940+
for c2 in colptr[col]:colptr[col+1]-1
941+
rowval[c2] = rowvals(S)[c1]
942+
nzval[c2] = nonzeros(S)[c1]
944943
c1 += 1
945944
end
946945
end
947-
A
946+
SparseMatrixCSC(m, n, colptr, rowval, nzval)
948947
end
949948

950949
function tril(S::AbstractSparseMatrixCSC{Tv,Ti}, k::Integer=0) where {Tv,Ti}
@@ -965,17 +964,16 @@ function tril(S::AbstractSparseMatrixCSC{Tv,Ti}, k::Integer=0) where {Tv,Ti}
965964
end
966965
rowval = Vector{Ti}(undef, nnz)
967966
nzval = Vector{Tv}(undef, nnz)
968-
A = SparseMatrixCSC(m, n, colptr, rowval, nzval)
969967
for col = 1 : min(n, m+k)
970968
c1 = getcolptr(S)[col+1]-1
971-
l2 = getcolptr(A)[col+1]-1
972-
for c2 = 0 : l2 - getcolptr(A)[col]
973-
rowvals(A)[l2 - c2] = rowvals(S)[c1]
974-
nonzeros(A)[l2 - c2] = nonzeros(S)[c1]
969+
l2 = colptr[col+1]-1
970+
for c2 = 0 : l2 - colptr[col]
971+
rowval[l2 - c2] = rowvals(S)[c1]
972+
nzval[l2 - c2] = nonzeros(S)[c1]
975973
c1 -= 1
976974
end
977975
end
978-
A
976+
SparseMatrixCSC(m, n, colptr, rowval, nzval)
979977
end
980978

981979
## diff
@@ -1340,19 +1338,16 @@ end
13401338

13411339
## kron
13421340
@inline function kron!(C::SparseMatrixCSC, A::AbstractSparseMatrixCSC, B::AbstractSparseMatrixCSC)
1343-
nnzC = nnz(A)*nnz(B)
13441341
mA, nA = size(A); mB, nB = size(B)
13451342
mC, nC = mA*mB, nA*nB
13461343

13471344
rowvalC = rowvals(C)
13481345
nzvalC = nonzeros(C)
13491346
colptrC = getcolptr(C)
13501347

1351-
@boundscheck begin
1352-
length(colptrC) == nC+1 || throw(DimensionMismatch("expect C to be preallocated with $(nC+1) colptrs "))
1353-
length(rowvalC) == nnzC || throw(DimensionMismatch("expect C to be preallocated with $(nnzC) rowvals"))
1354-
length(nzvalC) == nnzC || throw(DimensionMismatch("expect C to be preallocated with $(nnzC) nzvals"))
1355-
end
1348+
nnzC = nnz(A)*nnz(B)
1349+
resize!(nzvalC, nnzC)
1350+
resize!(rowvalC, nnzC)
13561351

13571352
col = 1
13581353
@inbounds for j = 1:nA
@@ -1381,16 +1376,13 @@ end
13811376
end
13821377

13831378
@inline function kron!(z::SparseVector, x::SparseVector, y::SparseVector)
1384-
nnzx = nnz(x); nnzy = nnz(y); nnzz = nnz(z);
1379+
nnzx = nnz(x); nnzy = nnz(y);
13851380
nzind = nonzeroinds(z)
13861381
nzval = nonzeros(z)
13871382

1388-
@boundscheck begin
1389-
nnzval = length(nzval); nnzind = length(nzind)
1390-
nnzz = nnzx*nnzy
1391-
nnzval == nnzz || throw(DimensionMismatch("expect z to be preallocated with $nnzz nonzeros"))
1392-
nnzind == nnzz || throw(DimensionMismatch("expect z to be preallocated with $nnzz nonzeros"))
1393-
end
1383+
nnzz = nnzx*nnzy
1384+
resize!(nzind, nnzz)
1385+
resize!(nzval, nnzz)
13941386

13951387
@inbounds for i = 1:nnzx, j = 1:nnzy
13961388
this_ind = (i-1)*nnzy+j
@@ -1402,17 +1394,12 @@ end
14021394

14031395
# sparse matrix ⊗ sparse matrix
14041396
function kron(A::AbstractSparseMatrixCSC{T1,S1}, B::AbstractSparseMatrixCSC{T2,S2}) where {T1,S1,T2,S2}
1405-
nnzC = nnz(A)*nnz(B)
14061397
mA, nA = size(A); mB, nB = size(B)
14071398
mC, nC = mA*mB, nA*nB
14081399
Tv = typeof(one(T1)*one(T2))
14091400
Ti = promote_type(S1,S2)
1410-
colptrC = Vector{Ti}(undef, nC+1)
1411-
rowvalC = Vector{Ti}(undef, nnzC)
1412-
nzvalC = Vector{Tv}(undef, nnzC)
1413-
colptrC[1] = 1
1414-
# skip sparse_check
1415-
C = SparseMatrixCSC{Tv, Ti}(mC, nC, colptrC, rowvalC, nzvalC)
1401+
C = spzeros(Tv, Ti, mC, nC)
1402+
sizehint!(C, nnz(A)*nnz(B))
14161403
return @inbounds kron!(C, A, B)
14171404
end
14181405

0 commit comments

Comments
 (0)