Add an opt-in profiler for model macros (#3981)

Author: odow

docs/src/tutorials/getting_started/debugging.jl

@@ -386,7 +386,39 @@ end
 #   you do not have enough memory to store the model. Use a computer with more
 #   RAM.
 
-# ### Strategies
+# ### Macro timing
+
+# JuMP has a built-in feature that can measure the time spent in each macro.
+# Turn it on using [`set_macro_timing`](@ref), build the model, and then use
+# [`print_macro_timing_summary`](@ref) to print a summary. Here's an example:
+
+begin
+    N = 200
+    demand = vcat(-1.0, zeros(N - 2), 1.0)
+    edges = [(i, j) for i in 1:N for j in 1:N if i < j]
+    model = Model()
+    set_macro_timing(model, true)
+    @variable(model, flows[e in edges] >= 0)
+    @constraint(
+        model,
+        [n in 1:N],
+        sum(flows[(i, j)] for (i, j) in edges if j == n) -
+        sum(flows[(i, j)] for (i, j) in edges if i == n) == demand[n]
+    )
+    print_macro_timing_summary(model)
+end
+
+# In this case, you can see that the `@constraint` call dominates the runtime.
+# If it isn't obvious why that is, read the [Performance problems with sum-if formulations](@ref)
+# tutorial.
+
+# Note that the macro timing feature measures only the time spent inside JuMP
+# macros. It does not measure regular Julia code outside the macros.
+
+# ### Other strategies
+
+# If the macro timing feature does not reveal the bottleneck, it means that your
+# issue is in regular Julia code that is not inside a JuMP macro.
 
 # The strategy to debug JuMP models that have performance problems depends on
 # how long your model takes to build.

src/Containers/macro.jl

@@ -324,16 +324,25 @@ end
 
 _strip_LineNumberNode(x) = x
 
+"""
+    build_macro_expression_string(macro_name, args)
+
+Reconstruct a string of the macro input expression from the name and arguments.
+"""
+function build_macro_expression_string(macro_name, args)
+    str_args = join(_strip_LineNumberNode.(args), ", ")
+    return string("`@", macro_name, "(", str_args, ")`")
+end
+
 """
     build_error_fn(macro_name, args, source)
 
 Return a function that can be used in place of `Base.error`, but which
 additionally prints the macro from which it was called.
 """
 function build_error_fn(macro_name, args, source)
-    str_args = join(_strip_LineNumberNode.(args), ", ")
-    msg = "At $(source.file):$(source.line): `@$macro_name($str_args)`: "
-    error_fn(str...) = error(msg, str...)
+    msg = build_macro_expression_string(macro_name, args)
+    error_fn(str...) = error("At $(source.file):$(source.line): $msg: ", str...)
     return error_fn
 end
 

src/JuMP.jl

@@ -126,6 +126,9 @@ mutable struct GenericModel{T<:Real} <: AbstractModel
     set_string_names_on_creation::Bool
     #
     variable_in_set_ref::Dict{Any,MOI.ConstraintIndex}
+    # A dictionary to store timing information from the JuMP macros.
+    enable_macro_timing::Bool
+    macro_times::Dict{Tuple{LineNumberNode,String},Float64}
 end
 
 value_type(::Type{GenericModel{T}}) where {T} = T
@@ -237,6 +240,8 @@ function direct_generic_model(
         Dict{Symbol,Any}(),
         true,
         Dict{Any,MOI.ConstraintIndex}(),
+        false,
+        Dict{Tuple{LineNumberNode,String},Float64}(),
     )
 end
 
@@ -937,6 +942,7 @@ function Base.empty!(model::GenericModel)::GenericModel
     model.nlp_model = nothing
     empty!(model.obj_dict)
     empty!(model.ext)
+    empty!(model.macro_times)
     model.is_model_dirty = false
     return model
 end

src/macros.jl

@@ -392,26 +392,52 @@ function _finalize_macro(
     source::LineNumberNode;
     register_name::Union{Nothing,Symbol} = nothing,
     wrap_let::Bool = false,
+    time_it::Union{Nothing,String} = nothing,
 )
     @assert Meta.isexpr(model, :escape)
-    if wrap_let && model.args[1] isa Symbol
-        code = quote
-            let $model = $model
+    ret = gensym()
+    code = if wrap_let && model.args[1] isa Symbol
+        quote
+            $ret = let $model = $model
                 $code
             end
         end
+    else
+        :($ret = $code)
     end
     if register_name !== nothing
         sym_name = Meta.quot(register_name)
         code = quote
             _error_if_cannot_register($model, $sym_name)
-            $(esc(register_name)) = $model[$sym_name] = $code
+            $code
+            $(esc(register_name)) = $model[$sym_name] = $ret
+        end
+    end
+    if time_it !== nothing
+        code = quote
+            start_time = time()
+            $code
+            _add_or_set_macro_time(
+                $model,
+                ($(QuoteNode(source)), $time_it),
+                time() - start_time,
+            )
+            $ret
         end
     end
     is_valid_code = :(_valid_model($model, $(Meta.quot(model.args[1]))))
     return Expr(:block, source, is_valid_code, code)
 end
 
+_add_or_set_macro_time(model::AbstractModel, key, value) = nothing
+
+function _add_or_set_macro_time(model::GenericModel, key, value)
+    if model.enable_macro_timing
+        model.macro_times[key] = get!(model.macro_times, key, 0.0) + value
+    end
+    return
+end
+
 function _error_if_cannot_register(model::AbstractModel, name::Symbol)
     obj_dict = object_dictionary(model)
     if haskey(obj_dict, name)
@@ -474,6 +500,100 @@ function _plural_macro_code(model, block, macro_sym)
     return code
 end
 
+"""
+    set_macro_timing(::GenericModel, value::Bool)
+
+Turn on (if `value`, or off, if `!value`) JuMP's built-in profiling of model
+construction macros.
+
+Use [`print_macro_timing_summary`](@ref) to display a summary.
+
+## Example
+
+```jldoctest; filter=[r"Total time inside macros: .+ seconds", r"├.+", r"└.+"]
+julia> begin
+           model = Model()
+           set_macro_timing(model, true)
+           @variable(model, x[1:2])
+           @objective(model, Min, sum(x))
+       end;
+
+julia> print_macro_timing_summary(model)
+Total time inside macros: 5.33690e-02 seconds
+│
+├ 2.96490e-02 s [55.55%]
+│ ├ REPL[8]:3
+│ └ `@variable(model, x[1:2])`
+│
+└ 2.37200e-02 s [44.45%]
+  ├ REPL[8]:4
+  └ `@objective(model, Min, sum(x))`
+```
+"""
+function set_macro_timing(model::GenericModel, value::Bool)
+    model.enable_macro_timing = value
+    return
+end
+
+function _string_summary(x)
+    if length(x) <= 75
+        return x
+    end
+    return x[1:32] * " [...] " * x[end-31:end]
+end
+
+_format_time(x::Float64) = string(_format(x), " seconds")
+
+"""
+    print_macro_timing_summary([io::IO = stdout], model::GenericModel)
+
+Print a summary of the runtime of each macro.
+
+Before calling this method, you must have enabled the macro timing feature using
+[`set_macro_timing`](@ref).
+
+## Example
+
+```jldoctest; filter=[r"Total time inside macros: .+ seconds", r"├.+", r"└.+"]
+julia> begin
+           model = Model()
+           set_macro_timing(model, true)
+           @variable(model, x[1:2])
+           @objective(model, Min, sum(x))
+       end;
+
+julia> print_macro_timing_summary(model)
+Total time inside macros: 5.33690e-02 seconds
+│
+├ 2.96490e-02 s [55.55%]
+│ ├ REPL[8]:3
+│ └ `@variable(model, x[1:2])`
+│
+└ 2.37200e-02 s [44.45%]
+  ├ REPL[8]:4
+  └ `@objective(model, Min, sum(x))`
+```
+"""
+function print_macro_timing_summary(io::IO, model::GenericModel)
+    total_time = sum(values(model.macro_times))
+    times = sort!(collect(model.macro_times); by = last, rev = true)
+    println(io, "Total time inside macros: ", _format_time(total_time))
+    for i in 1:length(times)
+        (source, expr), time = times[i]
+        percent = round(100 * time / total_time; digits = 2)
+        a, b = ifelse(i < length(times), ('├', '│'), ('└', ' '))
+        println(io, "│")
+        println(io, "$a ", _format(time), " s [$percent%]")
+        println(io, "$b ├ $(source.file):$(source.line)")
+        println(io, "$b └ ", replace(_string_summary(expr), "\n" => ""))
+    end
+    return
+end
+
+function print_macro_timing_summary(model::GenericModel)
+    return print_macro_timing_summary(stdout, model)
+end
+
 for file in readdir(joinpath(@__DIR__, "macros"))
     # The check for .jl is necessary because some users may have other files
     # like .cov from running code coverage. See JuMP.jl#3746.

src/macros/@constraint.jl

@@ -179,6 +179,10 @@ macro constraint(input_args...)
         __source__;
         register_name = name,
         wrap_let = true,
+        time_it = Containers.build_macro_expression_string(
+            :constraint,
+            input_args,
+        ),
     )
 end
 

src/macros/@expression.jl

@@ -94,6 +94,10 @@ macro expression(input_args...)
         __source__;
         register_name = name,
         wrap_let = true,
+        time_it = Containers.build_macro_expression_string(
+            :expression,
+            input_args,
+        ),
     )
 end
 

src/macros/@force_nonlinear.jl

@@ -84,10 +84,10 @@ julia> @expression(model, @force_nonlinear(x * 2.0 * (1 + x) * x))
 x * 2.0 * (1 + x) * x
 
 julia> @allocated @expression(model, x * 2.0 * (1 + x) * x)
-3584
+3680
 
 julia> @allocated @expression(model, @force_nonlinear(x * 2.0 * (1 + x) * x))
-672
+768
 ```
 """
 macro force_nonlinear(expr)

src/macros/@objective.jl

@@ -107,7 +107,15 @@ macro objective(input_args...)
         set_objective($esc_model, $sense, $expr)
         $expr
     end
-    return _finalize_macro(esc_model, code, __source__)
+    return _finalize_macro(
+        esc_model,
+        code,
+        __source__;
+        time_it = Containers.build_macro_expression_string(
+            :objective,
+            input_args,
+        ),
+    )
 end
 
 function _parse_moi_sense(error_fn::Function, sense)

src/macros/@variable.jl

@@ -320,6 +320,10 @@ macro variable(input_args...)
         __source__;
         register_name = name,
         wrap_let = true,
+        time_it = Containers.build_macro_expression_string(
+            :variable,
+            input_args,
+        ),
     )
 end
 

test/test_macros.jl

@@ -2527,4 +2527,44 @@ function test_malformed_kwarg()
     return
 end
 
+function test_print_macro_timing_summary()
+    model = Model()
+    @variable(model, x[1:2])  # This one should be excluded
+    set_macro_timing(model, true)
+    @objective(model, Min, sum(x))
+    @expression(model, x)
+    @constraint(model, [i in 1:2], x[i] <= i)
+    for i in 1:100
+        @constraint(model, sin(x[1]) <= i)
+    end
+    @variables(model, begin
+        a[1:2]
+        b, Bin, (start = 0)
+    end)
+    @expression(
+        model,
+        a_really_really_really_really_really_really_really_long_name,
+        x,
+    )
+    set_macro_timing(model, false)
+    @expression(model, expr, x[1] + x[2])  # So should this one
+    print_macro_timing_summary(model)
+    contents = sprint(print_macro_timing_summary, model)
+    needles = [
+        "@variable(model, a[1:2])",
+        "@variable(model, b, Bin, start = 0)",
+        "@objective(model, Min, sum(x))",
+        "@expression(model, x)",
+        "@constraint(model, [i in 1:2], x[i] <= i)",
+        "@constraint(model, sin(x[1]) <= i)",
+        "@expression(model, a_really_rea [...] lly_really_really_long_name, x)",
+    ]
+    @test all(n -> occursin(n, contents), needles)
+    @test length(model.macro_times) == length(needles)
+    @test all(>=(0), values(model.macro_times))
+    @test !occursin("@variable(model, x[1:2])", contents)
+    @test !occursin("@expression(model, expr, x[1] + x[2])", contents)
+    return
+end
+
 end  # module