Add allocations tracker for linear API

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
 
 """

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"])),
 )