Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

⚡️ Avoid creating temporary expression in matrix multiplications #1866

Merged
merged 4 commits into from
Feb 15, 2019
Merged

⚡️ Avoid creating temporary expression in matrix multiplications #1866

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3fa87aa
⚡️ Avoid creating temporary expression in matrix multiplications
blegat Feb 15, 2019
fbb3319
Fix StackOverflow in Julia v1.1
blegat Feb 15, 2019
a7f7479
Add matrix_product speedtest
blegat Feb 15, 2019
9e841a5
Fix style in vector_speedtest
blegat Feb 15, 2019
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
46 changes: 37 additions & 9 deletions src/aff_expr.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,8 +185,11 @@ function add_to_expression! end

# TODO: add deprecations for Base.push! and Base.append!

function add_to_expression!(aff::GenericAffExpr{C,V}, new_coef::C, new_var::V) where {C,V}
_add_or_set!(aff.terms, new_var, new_coef)
# With one factor.

function add_to_expression!(aff::GenericAffExpr{C,V},
other::Real) where {C,V}
aff.constant += other
return aff
end

Expand All @@ -195,23 +198,48 @@ function add_to_expression!(aff::GenericAffExpr{C,V}, new_var::V) where {C,V}
return aff
end

function add_to_expression!(aff::GenericAffExpr{C,V}, other::GenericAffExpr{C,V}) where {C,V}
function add_to_expression!(aff::GenericAffExpr{C,V},
other::GenericAffExpr{C,V}) where {C,V}
# Note: merge!() doesn't appear to call sizehint!(). Is this important?
merge!(+, aff.terms, other.terms)
aff.constant += other.constant
return aff
end

function add_to_expression!(aff::GenericAffExpr{C,V}, other::C) where {C,V}
aff.constant += other
# With two factors.

function add_to_expression!(aff::GenericAffExpr{C,V}, new_coef::Real,
new_var::V) where {C,V}
_add_or_set!(aff.terms, new_var, convert(C, new_coef))
return aff
end
function add_to_expression!(aff::GenericAffExpr{C,V}, other::Real) where {C,V}
aff.constant += other

function add_to_expression!(aff::GenericAffExpr{C,V}, new_var::V,
new_coef::Real) where {C,V}
return add_to_expression!(aff, new_coef, new_var)
end

function add_to_expression!(aff::GenericAffExpr{C,V}, coef::Real,
other::GenericAffExpr{C,V}) where {C,V}
sizehint!(aff, length(linear_terms(aff)) + length(linear_terms(other)))
for (term_coef, var) in linear_terms(other)
_add_or_set!(aff.terms, var, coef * term_coef)
end
aff.constant += coef * other.constant
return aff
end

function Base.isequal(aff::GenericAffExpr{C,V},other::GenericAffExpr{C,V}) where {C,V}
return isequal(aff.constant, other.constant) && isequal(aff.terms, other.terms)
function add_to_expression!(aff::GenericAffExpr{C,V},
other::GenericAffExpr{C,V},
coef::Real) where {C,V}
return add_to_expression!(aff, coef, other)
end


function Base.isequal(aff::GenericAffExpr{C,V},
other::GenericAffExpr{C,V}) where {C,V}
return isequal(aff.constant, other.constant) &&
isequal(aff.terms, other.terms)
end

Base.hash(aff::GenericAffExpr, h::UInt) = hash(aff.constant, hash(aff.terms, h))
Expand Down
73 changes: 15 additions & 58 deletions src/operators.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,11 +156,7 @@ function Base.:+(lhs::GenericAffExpr{C, V}, rhs::GenericAffExpr{C, V}) where {C,
operator_warn(owner_model(first(linear_terms(lhs))[2]))
end
end
result_terms = copy(lhs.terms)
# merge() returns a Dict(), so we need to call merge!() instead.
# Note: merge!() doesn't appear to call sizehint!(). Is this important?
merge!(+, result_terms, rhs.terms)
return GenericAffExpr(lhs.constant + rhs.constant, result_terms)
return add_to_expression!(copy(lhs), rhs)
end

function Base.:-(lhs::GenericAffExpr{C, V}, rhs::GenericAffExpr{C, V}) where {C, V<:_JuMPTypes}
Expand All @@ -175,44 +171,7 @@ end

function Base.:*(lhs::GenericAffExpr{C, V}, rhs::GenericAffExpr{C, V}) where {C, V<:_JuMPTypes}
result = zero(GenericQuadExpr{C, V})

lhs_length = length(linear_terms(lhs))
rhs_length = length(linear_terms(rhs))

# Quadratic terms
for (lhscoef, lhsvar) in linear_terms(lhs)
for (rhscoef, rhsvar) in linear_terms(rhs)
add_to_expression!(result, lhscoef*rhscoef, lhsvar, rhsvar)
end
end

# Try to preallocate space for aff
if !iszero(lhs.constant) && !iszero(rhs.constant)
sizehint!(result.aff, lhs_length + rhs_length)
elseif !iszero(lhs.constant)
sizehint!(result.aff, rhs_length)
elseif !iszero(rhs.constant)
sizehint!(result.aff, lhs_length)
end

# [LHS constant] * [RHS linear terms]
if !iszero(lhs.constant)
c = lhs.constant
for (rhscoef, rhsvar) in linear_terms(rhs)
add_to_expression!(result.aff, c*rhscoef, rhsvar)
end
end

# [RHS constant] * [LHS linear terms]
if !iszero(rhs.constant)
c = rhs.constant
for (lhscoef, lhsvar) in linear_terms(lhs)
add_to_expression!(result.aff, c*lhscoef, lhsvar)
end
end

result.aff.constant = lhs.constant * rhs.constant

add_to_expression!(result, lhs, rhs)
return result
end
# GenericAffExpr--GenericQuadExpr
Expand Down Expand Up @@ -413,11 +372,10 @@ function _A_mul_B!(ret::AbstractArray{T}, A, B) where {T <: _JuMPTypes}
q = ret[i, j]
_sizehint_expr!(q, size(A, 2))
for k ∈ 1:size(A, 2)
tmp = convert(T, A[i, k] * B[k, j])
add_to_expression!(q, tmp)
add_to_expression!(q, A[i, k], B[k, j])
end
end
ret
return ret
end

function _A_mul_B!(ret::AbstractArray{<:GenericAffOrQuadExpr}, A::SparseMatrixCSC, B)
Expand All @@ -426,11 +384,11 @@ function _A_mul_B!(ret::AbstractArray{<:GenericAffOrQuadExpr}, A::SparseMatrixCS
for col ∈ 1:size(A, 2)
for k ∈ 1:size(ret, 2)
for j ∈ nzrange(A, col)
add_to_expression!(ret[rv[j], k], nzv[j] * B[col, k])
add_to_expression!(ret[rv[j], k], nzv[j], B[col, k])
end
end
end
ret
return ret
end

function _A_mul_B!(ret::AbstractArray{<:GenericAffOrQuadExpr}, A::AbstractMatrix, B::SparseMatrixCSC)
Expand All @@ -442,11 +400,11 @@ function _A_mul_B!(ret::AbstractArray{<:GenericAffOrQuadExpr}, A::AbstractMatrix
q = ret[multivec_row, col]
_sizehint_expr!(q, length(idxset))
for k ∈ idxset
add_to_expression!(q, A[multivec_row, rowval[k]] * nzval[k])
add_to_expression!(q, A[multivec_row, rowval[k]], nzval[k])
end
end
end
ret
return ret
end

# TODO: Implement sparse * sparse code as in base/sparse/linalg.jl (spmatmul).
Expand All @@ -471,8 +429,7 @@ function _At_mul_B!(ret::AbstractArray{T}, A, B) where {T <: _JuMPTypes}
q = ret[i, j]
_sizehint_expr!(q, size(A, 1))
for k ∈ 1:size(A, 1)
tmp = convert(T, A[k, i] * B[k, j]) # transpose
add_to_expression!(q, tmp)
add_to_expression!(q, A[k, i], B[k, j]) # transpose
end
end
ret
Expand Down Expand Up @@ -517,23 +474,23 @@ function _At_mul_B!(ret::StridedVecOrMat{<:GenericAffOrQuadExpr}, A::SparseMatri
size(B, 2) == size(ret, 2) || throw(DimensionMismatch())
nzv = A.nzval
rv = A.rowval
# ret is already filled with zeros by _return_arrayt.
# `ret` is already filled with zeros by `_return_arrayt`.
for k = 1:size(ret, 2)
@inbounds for col = 1:A.n
tmp = zero(eltype(ret))
for j = A.colptr[col]:(A.colptr[col + 1] - 1)
tmp += adjoint(nzv[j]) * B[rv[j],k]
add_to_expression!(tmp, adjoint(nzv[j]), B[rv[j],k])
end
ret[col,k] += tmp
ret[col, k] += tmp
end
end
ret
return ret
end
function _At_mul_B(A, B)
size(A, 1) == size(B, 1) || error("Incompatible sizes")
ret = _A_mul_B_ret(transpose(A), B)
_At_mul_B!(ret, A, B)
ret
return ret
end

# TODO: Implement sparse * sparse code as in base/sparse/linalg.jl (spmatmul).
Expand Down Expand Up @@ -617,7 +574,7 @@ function Base.:-(x::AbstractArray{T}) where {T <: _JuMPTypes}
for I in eachindex(ret)
ret[I] = -x[I]
end
ret
return ret
end

###############################################################################
Expand Down
134 changes: 119 additions & 15 deletions src/quad_expr.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -126,33 +126,137 @@ function Base.eltype(qti::QuadTermIterator{GenericQuadExpr{C, V}}
return Tuple{C, V, V}
end

function add_to_expression!(quad::GenericQuadExpr{C,V}, new_coef::C, new_var1::V, new_var2::V) where {C,V}
# Node: OrderedDict updates the *key* as well. That is, if there was a
# previous value for UnorderedPair(new_var2, new_var1), it's key will now be
# UnorderedPair(new_var1, new_var2) (because these are defined as equal).
key = UnorderedPair(new_var1, new_var2)
_add_or_set!(quad.terms, key, new_coef)
return quad
end
# With one factor.

function add_to_expression!(quad::GenericQuadExpr{C, V}, new_coef::C, new_var::V) where {C,V}
add_to_expression!(quad.aff, new_coef, new_var)
return quad
function add_to_expression!(quad::GenericQuadExpr{C}, other::C) where C
return add_to_expression!(quad.aff, other)
end

function add_to_expression!(q::GenericQuadExpr{T,S}, other::GenericAffExpr{T,S}) where {T,S}
function add_to_expression!(q::GenericQuadExpr{T,S},
other::GenericAffExpr{T,S}) where {T,S}
add_to_expression!(q.aff, other)
return q
end

function add_to_expression!(q::GenericQuadExpr{T,S}, other::GenericQuadExpr{T,S}) where {T,S}
function add_to_expression!(q::GenericQuadExpr{T,S},
other::GenericQuadExpr{T,S}) where {T,S}
merge!(+, q.terms, other.terms)
add_to_expression!(q.aff, other.aff)
return q
end

function add_to_expression!(quad::GenericQuadExpr{C}, other::C) where C
return add_to_expression!(quad.aff, other)
# With two factors.

function add_to_expression!(quad::GenericQuadExpr{C, V},
new_coef::Real,
new_var::V) where {C,V}
add_to_expression!(quad.aff, new_coef, new_var)
return quad
end

function add_to_expression!(quad::GenericQuadExpr{C, V},
new_var::Union{V, GenericAffExpr{C, V}},
new_coef::Real) where {C,V}
return add_to_expression!(quad, new_coef, new_var)
end

function add_to_expression!(quad::GenericQuadExpr{C},
new_coef::Real,
new_aff::GenericAffExpr{C}) where {C}
add_to_expression!(quad.aff, new_coef, new_aff)
return quad
end

function add_to_expression!(quad::GenericQuadExpr{C, V}, coef::Real,
other::GenericQuadExpr{C, V}) where {C, V}
for (key, term_coef) in other.terms
_add_or_set!(quad.terms, key, coef * term_coef)
end
return add_to_expression!(quad, coef, other.aff)
end

function add_to_expression!(quad::GenericQuadExpr{C, V},
other::GenericQuadExpr{C, V},
coef::Real) where {C, V}
return add_to_expression!(quad, coef, other)
end

function add_to_expression!(quad::GenericQuadExpr{C},
var_1::AbstractVariableRef,
var_2::AbstractVariableRef) where {C}
return add_to_expression!(quad, one(C), var_1, var_2)
end

function add_to_expression!(quad::GenericQuadExpr{C,V},
var::V,
aff::GenericAffExpr{C,V}) where {C,V}
for (coef, term_var) in linear_terms(aff)
key = UnorderedPair(var, term_var)
_add_or_set!(quad.terms, key, coef)
end
return add_to_expression!(quad, var, aff.constant)
end

function add_to_expression!(quad::GenericQuadExpr{C,V},
aff::GenericAffExpr{C,V},
var::V) where {C,V}
return add_to_expression!(quad, var, aff)
end

function add_to_expression!(quad::GenericQuadExpr{C,V},
lhs::GenericAffExpr{C,V},
rhs::GenericAffExpr{C,V}) where {C,V}
lhs_length = length(linear_terms(lhs))
rhs_length = length(linear_terms(rhs))

# Quadratic terms
for (lhscoef, lhsvar) in linear_terms(lhs)
for (rhscoef, rhsvar) in linear_terms(rhs)
add_to_expression!(quad, lhscoef*rhscoef, lhsvar, rhsvar)
end
end

# Try to preallocate space for aff
cur = length(linear_terms(quad))
if !iszero(lhs.constant) && !iszero(rhs.constant)
sizehint!(quad.aff, cur + lhs_length + rhs_length)
elseif !iszero(lhs.constant)
sizehint!(quad.aff, cur + rhs_length)
elseif !iszero(rhs.constant)
sizehint!(quad.aff, cur + lhs_length)
end

# [LHS constant] * [RHS linear terms]
if !iszero(lhs.constant)
c = lhs.constant
for (rhscoef, rhsvar) in linear_terms(rhs)
add_to_expression!(quad.aff, c*rhscoef, rhsvar)
end
end

# [RHS constant] * [LHS linear terms]
if !iszero(rhs.constant)
c = rhs.constant
for (lhscoef, lhsvar) in linear_terms(lhs)
add_to_expression!(quad.aff, c*lhscoef, lhsvar)
end
end

quad.aff.constant += lhs.constant * rhs.constant

return quad
end

# With three factors.

function add_to_expression!(quad::GenericQuadExpr{C,V}, new_coef::C,
new_var1::V, new_var2::V) where {C,V}
# Node: OrderedDict updates the *key* as well. That is, if there was a
# previous value for UnorderedPair(new_var2, new_var1), it's key will now be
# UnorderedPair(new_var1, new_var2) (because these are defined as equal).
key = UnorderedPair(new_var1, new_var2)
_add_or_set!(quad.terms, key, new_coef)
return quad
end


Expand Down
Loading