Skip to content

Commit

Permalink
Merge pull request #1866 from JuliaOpt/bl/tmpmatmul
Browse files Browse the repository at this point in the history 
⚡️ Avoid creating temporary expression in matrix multiplications
  • Loading branch information
@blegat
blegat authored Feb 15, 2019
2 parents 13cf483 + 9e841a5 commit 7426454
Show file tree
Hide file tree
Showing 5 changed files with 217 additions and 95 deletions.
46 changes: 37 additions & 9 deletions src/aff_expr.jl
View file
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
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
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

0 comments on commit 7426454

Please sign in to comment.