diff --git a/src/allocs_model.jl b/src/allocs_model.jl index c59e06e..514e0d4 100644 --- a/src/allocs_model.jl +++ b/src/allocs_model.jl @@ -1,13 +1,14 @@ export test_allocs_nlpmodels, test_allocs_nlsmodels, print_nlp_allocations """ - test_allocs_nlpmodels(nlp::AbstractNLPModel; exclude = []) + test_allocs_nlpmodels(nlp::AbstractNLPModel; linear_api = false, exclude = []) Returns a `Dict` containing allocations of the in-place functions of NLPModel API. The keyword `exclude` takes a Array of Function to be excluded from the tests. Use `hess` (resp. `jac`) to exclude `hess_coord` and `hess_structure` (resp. `jac_coord` and `jac_structure`). +Setting `linear_api` to `true` will also checks the functions specific to linear and nonlinear constraints. """ -function test_allocs_nlpmodels(nlp::AbstractNLPModel; exclude = []) +function test_allocs_nlpmodels(nlp::AbstractNLPModel; linear_api = false, exclude = []) nlp_allocations = Dict( :obj => NaN, :grad! => NaN, @@ -90,6 +91,7 @@ function test_allocs_nlpmodels(nlp::AbstractNLPModel; exclude = []) if get_ncon(nlp) > 0 && !(jac in exclude) rows = Vector{Int}(undef, nlp.meta.nnzj) cols = Vector{Int}(undef, nlp.meta.nnzj) + @show "Do we pass here?" jac_structure!(nlp, rows, cols) nlp_allocations[:jac_structure!] = @allocated jac_structure!(nlp, rows, cols) x = get_x0(nlp) @@ -124,6 +126,71 @@ function test_allocs_nlpmodels(nlp::AbstractNLPModel; exclude = []) mul!(Jtv, J', w) nlp_allocations[:jac_op_transpose_prod!] = @allocated mul!(Jtv, J', w) end + + for type in (:nln, :lin) + nn = type == :lin ? nlp.meta.nlin : nlp.meta.nnln + nnzj = type == :lin ? nlp.meta.lin_nnzj : nlp.meta.nln_nnzj + if !linear_api || (nn == 0) + continue + end + if !(cons in exclude) + x = get_x0(nlp) + c = Vector{eltype(x)}(undef, nn) + fun = Symbol(:cons_, type, :!) + eval(fun)(nlp, x, c) + nlp_allocations[fun] = @allocated eval(fun)(nlp, x, c) + end + if !(jac in exclude) + rows = Vector{Int}(undef, nnzj) + cols = Vector{Int}(undef, nnzj) + fun = type == :lin ? jac_lin_structure! : jac_nln_structure! # eval(fun) would allocate here + fun(nlp, rows, cols) + nlp_allocations[Symbol(fun)] = @allocated fun(nlp, rows, cols) + x = get_x0(nlp) + vals = Vector{eltype(x)}(undef, nnzj) + fun = Symbol(:jac_, type, :_coord!) + eval(fun)(nlp, x, vals) + nlp_allocations[fun] = @allocated eval(fun)(nlp, x, vals) + end + if !(jprod in exclude) + x = get_x0(nlp) + v = copy(x) + Jv = Vector{eltype(x)}(undef, nn) + fun = Symbol(:jprod_, type, :!) + eval(fun)(nlp, x, v, Jv) + nlp_allocations[fun] = @allocated eval(fun)(nlp, x, v, Jv) + end + if !(jtprod in exclude) + x = get_x0(nlp) + v = copy(get_y0(nlp)[1:nn]) + Jtv = similar(x) + fun = Symbol(:jtprod_, type, :!) + eval(fun)(nlp, x, v, Jtv) + nlp_allocations[fun] = @allocated eval(fun)(nlp, x, v, Jtv) + end + if !(jac_op in exclude) + x = get_x0(nlp) + Jtv = similar(x) + Jv = Vector{eltype(x)}(undef, nn) + + v = copy(x) + w = randn(eltype(x), nn) + fun = Symbol(:jac_, type, :_op!) + if type == :lin + J = jac_lin_op!(nlp, x, Jv, Jtv) + mul!(Jv, J, v) + nlp_allocations[Symbol(:jac_lin_op_prod!)] = @allocated mul!(Jv, J, v) + mul!(Jtv, J', w) + nlp_allocations[Symbol(:jac_lin_op_transpose_prod!)] = @allocated mul!(Jtv, J', w) + else + J = jac_nln_op!(nlp, x, Jv, Jtv) + mul!(Jv, J, v) + nlp_allocations[Symbol(:jac_nln_op_prod!)] = @allocated mul!(Jv, J, v) + mul!(Jtv, J', w) + nlp_allocations[Symbol(:jac_nln_op_transpose_prod!)] = @allocated mul!(Jtv, J', w) + end + end + end return nlp_allocations end @@ -229,7 +296,7 @@ end function NLPModels.histline(s, v, maxv) @assert 0 ≤ v ≤ maxv λ = maxv == 0 ? 0 : ceil(Int, 20 * v / maxv) - return @sprintf("%22s: %s %-6s", s, "█"^λ * "⋅"^(20 - λ), v) + return @sprintf("%27s: %s %-6s", s, "█"^λ * "⋅"^(20 - λ), v) end """ diff --git a/test/runtests.jl b/test/runtests.jl index 0c7acf0..cbe5000 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -52,12 +52,12 @@ pmap(nls_tests, NLPModelsTest.nls_problems) io = IOBuffer(); map( - nlp -> print_nlp_allocations(io, nlp, test_allocs_nlpmodels(nlp)), + nlp -> print_nlp_allocations(io, nlp, test_allocs_nlpmodels(nlp, linear_api = true)), map(x -> eval(Symbol(x))(), NLPModelsTest.nlp_problems), ) -print_nlp_allocations(io, LLS(), test_allocs_nlpmodels(LLS(), exclude = [hess])) +print_nlp_allocations(io, LLS(), test_allocs_nlpmodels(LLS(), linear_api = true, exclude = [hess])) map( - nlp -> print_nlp_allocations(io, nlp, test_allocs_nlpmodels(nlp)), + nlp -> print_nlp_allocations(io, nlp, test_allocs_nlpmodels(nlp, linear_api = true)), map(x -> eval(Symbol(x))(), setdiff(NLPModelsTest.nls_problems, ["LLS"])), )