Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 46 additions & 34 deletions src/Utilities/sparse_matrix.jl
Original file line number Diff line number Diff line change
Expand Up @@ -29,11 +29,11 @@ struct OneBasedIndexing <: AbstractIndexing end

_first_index(::ZeroBasedIndexing) = 0
_first_index(::OneBasedIndexing) = 1
_shift(x, ::ZeroBasedIndexing, ::ZeroBasedIndexing) = x

_shift(x, ::T, ::T) where {T<:AbstractIndexing} = x
_shift(x::Integer, ::ZeroBasedIndexing, ::OneBasedIndexing) = x + 1
_shift(x::Array{<:Integer}, ::ZeroBasedIndexing, ::OneBasedIndexing) = x .+ 1
_shift(x::Integer, ::OneBasedIndexing, ::ZeroBasedIndexing) = x - 1
_shift(x, ::OneBasedIndexing, ::OneBasedIndexing) = x
_shift(x::Array{<:Integer}, ::ZeroBasedIndexing, ::OneBasedIndexing) = x .+ 1

"""
mutable struct MutableSparseMatrixCSC{Tv,Ti<:Integer,I<:AbstractIndexing}
Expand Down Expand Up @@ -76,7 +76,8 @@ function MOI.empty!(A::MutableSparseMatrixCSC)
resize!(A.colptr, 1)
A.colptr[1] = 0
empty!(A.rowval)
return empty!(A.nzval)
empty!(A.nzval)
return
end

"""
Expand All @@ -103,7 +104,8 @@ function set_number_of_rows(A::MutableSparseMatrixCSC, num_rows)
A.colptr[i] += A.colptr[i-1]
end
resize!(A.rowval, A.colptr[end])
return resize!(A.nzval, A.colptr[end])
resize!(A.nzval, A.colptr[end])
return
end

"""
Expand All @@ -116,61 +118,71 @@ function final_touch(A::MutableSparseMatrixCSC)
for i in length(A.colptr):-1:2
A.colptr[i] = _shift(A.colptr[i-1], ZeroBasedIndexing(), A.indexing)
end
return A.colptr[1] = _first_index(A.indexing)
A.colptr[1] = _first_index(A.indexing)
return
end

_variable(term::MOI.ScalarAffineTerm) = term.variable
_variable(term::MOI.VectorAffineTerm) = _variable(term.scalar_term)
function _allocate_terms(colptr, index_map, terms)
for term in terms
colptr[index_map[_variable(term)].value+1] += 1
end
end
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@blegat what was the motivation to have these _allocate_terms function barriers?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It dates back to jump-dev/SCS.jl#192 where I benchmarked the two approaches and the function barrier seemed to be faster even if I didn't really get why.
It's maybe not the case anymore in Julia v1.6 or maybe I was mistaken when writing this.
We can revisit later. An easier speedup would be obtained by using @inbounds I guess.
The good thing with MatrixOfConstraints is that we will be able to spend time optimizing these pieces at once in MOI and it will affect all solver's copy_to performance so it might start being worth playing with @inbounds, etc...

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Given the massive amount of new complexity we're introducing, I'd rather that we merged a simplified version first, got it working and debugged, then wrote the various solvers, and only then revisited performance implications with benchmarks.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, definitely


"""
allocate_terms(A::MutableSparseMatrixCSC, index_map, func)
allocate_terms(
A::MutableSparseMatrixCSC,
index_map,
func::Union{MOI.ScalarAffineFunction,MOI.VectorAffineFunction},
)

Informs `A` that the terms of the function `func` where the variable
indices are mapped with `index_map` will be loaded with [`load_terms`](@ref).
The function `func` should be canonicalized, see [`is_canonical`](@ref).
"""
function allocate_terms(A::MutableSparseMatrixCSC, index_map, func)
return _allocate_terms(A.colptr, index_map, func.terms)
function allocate_terms(
A::MutableSparseMatrixCSC,
index_map,
func::Union{MOI.ScalarAffineFunction,MOI.VectorAffineFunction},
)
for term in func.terms
A.colptr[index_map[_variable(term)].value+1] += 1
end
return
end

_output_index(::MOI.ScalarAffineTerm) = 1
_output_index(term::MOI.VectorAffineTerm) = term.output_index
_coefficient(term::MOI.ScalarAffineTerm) = term.coefficient
_coefficient(term::MOI.VectorAffineTerm) = _coefficient(term.scalar_term)
function _load_terms(colptr, rowval, nzval, index_map, terms, offset)
for term in terms
ptr = colptr[index_map[_variable(term)].value] += 1
rowval[ptr] = offset + _output_index(term)
nzval[ptr] = _coefficient(term)
end
end

"""
load_terms(A::MutableSparseMatrixCSC, index_map, func, offset)
load_terms(
A::MutableSparseMatrixCSC,
index_map,
func::Union{MOI.ScalarAffineFunction,MOI.VectorAffineFunction},
offset::Int,
)

Loads the terms of `func` to `A` mapping the variable indices with `index_map`.
The `i`th dimension of `func` is loaded at the `(offset + i)`th row of `A`. The
function should be allocated first with [`allocate_terms`](@ref). The function
`func` should be canonicalized, see [`is_canonical`](@ref).
"""
function load_terms(A::MutableSparseMatrixCSC, index_map, func, offset)
return _load_terms(
A.colptr,
A.rowval,
A.nzval,
index_map,
func.terms,
_shift(offset, OneBasedIndexing(), A.indexing),
)
function load_terms(
A::MutableSparseMatrixCSC,
index_map,
func::Union{MOI.ScalarAffineFunction,MOI.VectorAffineFunction},
offset::Int,
)
new_offset = _shift(offset, OneBasedIndexing(), A.indexing)
for term in func.terms
ptr = A.colptr[index_map[_variable(term)].value] += 1
A.rowval[ptr] = new_offset + _output_index(term)
A.nzval[ptr] = MOI.coefficient(term)
end
return
end

"""
Base.convert(::Type{SparseMatrixCSC{Tv, Ti}}, A::MutableSparseMatrixCSC{Tv, Ti, I}) where {Tv, Ti, I}
Base.convert(
::Type{SparseMatrixCSC{Tv,Ti}},
A::MutableSparseMatrixCSC{Tv,Ti,I},
) where {Tv,Ti,I}

Converts `A` to a `SparseMatrixCSC`. Note that the field `A.nzval` is **not
copied** so if `A` is modified after the call of this function, it can still
Expand Down
176 changes: 176 additions & 0 deletions test/Utilities/sparse_matrix.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,176 @@
module TestSparseMatrix

import SparseArrays
using Test

import MathOptInterface
const MOI = MathOptInterface

function runtests()
for name in names(@__MODULE__; all = true)
if startswith("$(name)", "test_")
@testset "$(name)" begin
getfield(@__MODULE__, name)()
end
end
end
return
end

"""
test_empty_size()

Julia 1.7 is going to outlaw constructing empty sparse matrices. Test that we
can still do so.
"""
function test_empty_size()
A = MOI.Utilities.MutableSparseMatrixCSC{
Float64,
Int,
MOI.Utilities.ZeroBasedIndexing,
}()
MOI.empty!(A)
@test A.rowval == Int[]
@test A.nzval == Float64[]
@test A.colptr == Int[0]
end

function test_empty_large()
A = MOI.Utilities.MutableSparseMatrixCSC{
Float64,
Int,
MOI.Utilities.ZeroBasedIndexing,
}()
MOI.empty!(A)
MOI.Utilities.add_column(A)
MOI.Utilities.add_column(A)
MOI.Utilities.set_number_of_rows(A, 5)
B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
@test size(B) == (5, 2)
end

function test_VectorAffine_ZeroBased()
A = MOI.Utilities.MutableSparseMatrixCSC{
Float64,
Int,
MOI.Utilities.ZeroBasedIndexing,
}()
MOI.empty!(A)
x = MOI.VariableIndex.(1:3)
f = MOI.VectorAffineFunction(
vcat(
MOI.VectorAffineTerm.(1, MOI.ScalarAffineTerm.(1.0, x)),
MOI.VectorAffineTerm.(2, MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3])),
),
[0.5, 1.2],
)

index_map = MOI.Utilities.IndexMap()
for i in 1:3
MOI.Utilities.add_column(A)
index_map[x[i]] = x[i]
end
MOI.Utilities.allocate_terms(A, index_map, f)
MOI.Utilities.set_number_of_rows(A, 2)
MOI.Utilities.load_terms(A, index_map, f, 0)
MOI.Utilities.final_touch(A)
B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
@test B == [1.0 1.0 1.0; 0.0 2.0 3.0]
@test A.rowval == [0, 0, 1, 0, 1]
@test A.nzval == [1.0, 1.0, 2.0, 1.0, 3.0]
@test A.colptr == [0, 1, 3, 5]
end

function test_VectorAffine_OneBased()
A = MOI.Utilities.MutableSparseMatrixCSC{
Float64,
Int,
MOI.Utilities.OneBasedIndexing,
}()
MOI.empty!(A)
x = MOI.VariableIndex.(1:3)
f = MOI.VectorAffineFunction(
vcat(
MOI.VectorAffineTerm.(1, MOI.ScalarAffineTerm.(1.0, x)),
MOI.VectorAffineTerm.(2, MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3])),
),
[0.5, 1.2],
)

index_map = MOI.Utilities.IndexMap()
for i in 1:3
MOI.Utilities.add_column(A)
index_map[x[i]] = x[i]
end
MOI.Utilities.allocate_terms(A, index_map, f)
MOI.Utilities.set_number_of_rows(A, 2)
MOI.Utilities.load_terms(A, index_map, f, 0)
MOI.Utilities.final_touch(A)
B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
@test B == [1.0 1.0 1.0; 0.0 2.0 3.0]
@test A.rowval == [1, 1, 2, 1, 2]
@test A.nzval == [1.0, 1.0, 2.0, 1.0, 3.0]
@test A.colptr == [1, 2, 4, 6]
end

function test_ScalarAffine_ZeroBased()
A = MOI.Utilities.MutableSparseMatrixCSC{
Float64,
Int,
MOI.Utilities.ZeroBasedIndexing,
}()
MOI.empty!(A)
x = MOI.VariableIndex.(1:3)
f1 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.5)
f2 =
MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3]), 1.2)
index_map = MOI.Utilities.IndexMap()
for i in 1:3
MOI.Utilities.add_column(A)
index_map[x[i]] = x[i]
end
MOI.Utilities.allocate_terms(A, index_map, f1)
MOI.Utilities.allocate_terms(A, index_map, f2)
MOI.Utilities.set_number_of_rows(A, 2)
MOI.Utilities.load_terms(A, index_map, f1, 0)
MOI.Utilities.load_terms(A, index_map, f2, 1)
MOI.Utilities.final_touch(A)
B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
@test B == [1.0 1.0 1.0; 0.0 2.0 3.0]
@test A.rowval == [0, 0, 1, 0, 1]
@test A.nzval == [1.0, 1.0, 2.0, 1.0, 3.0]
@test A.colptr == [0, 1, 3, 5]
end

function test_ScalarAffine_OneBased()
A = MOI.Utilities.MutableSparseMatrixCSC{
Float64,
Int,
MOI.Utilities.OneBasedIndexing,
}()
MOI.empty!(A)
x = MOI.VariableIndex.(1:3)
f1 = MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.5)
f2 =
MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.([2.0, 3.0], x[2:3]), 1.2)
index_map = MOI.Utilities.IndexMap()
for i in 1:3
MOI.Utilities.add_column(A)
index_map[x[i]] = x[i]
end
MOI.Utilities.allocate_terms(A, index_map, f1)
MOI.Utilities.allocate_terms(A, index_map, f2)
MOI.Utilities.set_number_of_rows(A, 2)
MOI.Utilities.load_terms(A, index_map, f1, 0)
MOI.Utilities.load_terms(A, index_map, f2, 1)
MOI.Utilities.final_touch(A)
B = convert(SparseArrays.SparseMatrixCSC{Float64,Int}, A)
@test B == [1.0 1.0 1.0; 0.0 2.0 3.0]
@test A.rowval == [1, 1, 2, 1, 2]
@test A.nzval == [1.0, 1.0, 2.0, 1.0, 3.0]
@test A.colptr == [1, 2, 4, 6]
end

end

TestSparseMatrix.runtests()