To what extent can JuMP's macro identify input expression and generate optimized backend codes?

Specific Domains Optimization (Mathematical)
,
1

JuMP’s macros are like magic black boxes to me. I’ve learned that one should generally make use of @variable, @constraints and @objective when constructing a model from scratch. But I have some small but nontrivial questions.

In the very introduction page of JuMP’s doc, there is this example code
@objective(model, Min, sum(residuals.^2))

To be honest I dislike the expression. Were residuals an Array of Float64, I would always write
sum(x -> x^2, residuals). I bet every experienced julia user would write it this style solely.

Now I’m curious: can I write @objective(model, Min, sum(x -> x^2, residuals)) when residuals is a Vector of decision variables? Can JuMP’s macro recognize this expression and generate optimized code to attain maximum performance?

According to my past experience, not all julia jargons are understood by JuMP’s macros. They may work, but left unoptimized. e.g. writing a mapreduce expression inside @expression (Am I right?).

Now return to the initial question: why did you choose to use the awkward expression sum(residuals.^2)? Because the JuMP macros can only recognize that particular form?

Another very similar question:
For numeric matrix it makes sense in julia to write sum(view(A, :, 1)) instead of sum(A[:, 1]). Can I also write sum(view(x, :, 1)) inside a macro (and get optimized)?

2

JuMP’s macros are like magic black boxes to me

I think I’m going to give a talk at JuMP-dev 2026 about the macros. It would be useful for the community to have a reference we can point people to about what is happening.

can I write @objective(model, Min, sum(x -> x^2, residuals)) when residuals is a Vector of decision variables

Yes.

Can JuMP’s macro recognize this expression and generate optimized code to attain maximum performance?

No

why did you choose to use the awkward expression sum(residuals.^2) ?

Because I don’t think it’s awkward :smile:

Can I also write sum(view(x, :, 1)) inside a macro (and get optimized)?

The answer to this, any many of your similar questions is: it depends. Try it out! And benchmark it! I’ve said this a few times in response to some of your example codes: don’t focus on performance. Write code that is simple to read first. Benchmark it to see if it is fast enough for your purpose. And only if it is a bottleneck should you try to be clever.

julia> function main(n)
           model = Model()
           @variable(model, x[1:n, 1:2])
           @time sum(view(x, :, 1))
           @time sum(x[:, 1])
           return
       end
main (generic function with 1 method)

julia> main(10_000)
  0.000350 seconds (48 allocations: 769.109 KiB)
  0.000339 seconds (51 allocations: 929.172 KiB)

julia> main(20_000)
  0.001558 seconds (55 allocations: 1.657 MiB)
  0.001718 seconds (58 allocations: 1.970 MiB)

julia> main(100_000)
  0.003989 seconds (66 allocations: 6.564 MiB)
  0.004013 seconds (69 allocations: 8.095 MiB)
2 Likes
3

I think it’s non-idiomatic and is misleading to people who come to the julia language primarily because of he wants to use JuMP (e.g. me). It is awkward to even use this style in the welcome page of JuMP.jl’s documentation.

julia> generate(N) = rand(N, N);

julia> idiomatic(A) = sum(x -> x^2, A);

julia> non_idiomatic(A) = sum(A .^ 2);

julia> function test(N)
           A = generate(N)
           @time non_idiomatic(A)
           @time idiomatic(A)
           @time non_idiomatic(A)
           @time idiomatic(A)
           @time non_idiomatic(A)
           @time idiomatic(A)
       end;

julia> test(10);
# omit this output

julia> test(10000);
  0.456077 seconds (3 allocations: 762.941 MiB, 5.45% gc time)
  0.030639 seconds
  0.495797 seconds (3 allocations: 762.941 MiB, 15.30% gc time)
  0.031325 seconds
  0.551592 seconds (3 allocations: 762.941 MiB, 21.33% gc time)
  0.030970 seconds

It makes sense for you to suggest me to do benchmarks. But I think it’s too expensive for me. Oftentimes I don’t have the bandwidth to do that—I need to first of all learn how to carry out benchmarks properly (in a fair manner). (I think I’m still relatively weak in this aspect).

4

sum(A.^2) is idiomatic, and the performance drawback doesn’t change that.

3 Likes
5

I have no authority to define idiomaticness. But it appears in julia it’s cleanest to write
sum(abs2, A).

(I’ll have to go off-line and write my own optimization code first🥲)

6

We already have odow’s (authoritative) answer to that, but here are my two cents.

I think one of JuMP’s most compelling selling points is that you can write Julia code that really looks like the mathematical expressions in the acompanying paper / documentation.

In this regard, I think sum(residuals.^2) is a great showcase of JuMP capabilities, worth figuring on the introduction of JuMP’s documentation.

Optimizing sum(residuals.^2) to something more performant falls outside the purpose of JuMP’s documentation IMO. It is regular Julia code, that can be optimized like anything else… but you don’t necessarily want to.

For example, how much time is optimize! going to take on the 10k variable model that took 500ms to build (instead of 3ms) ? If it’s more than a few seconds, you probably don’t want to spend time trying to optimize the expressions in your model, and are instead better off writing them in the most legible way for you or you co-developers.

All that being said, I do sometimes encounter situations where this specific issue matters and the models are simple enough to optimize for construction times to become non-negligible. In that case, my personal preference goes to something like:

@objective(model, Min, sum(residual[i]^2 for i in 1:N))

which reads exactly like what I write in the accompanying documentation:

\text{min} \sum_{i=1}^N r_i^2.
5 Likes
7

:100:

2 Likes
8

This doc is questionable

image
image1066×1445 102 KB

The naive grammar is to write
x[1] + x[2] + x[3]
resolving into

a = x[1]
b = a + x[2]
c = b + x[3]

The a is merely a name, it’s not “constructed”. The description is imprecise.

And the allocation comparison contradicts the actual observation.

A related question arises when I want to define a helper function to extract the absolute output power:

julia> JuMP.@variable(model, p[t=1:24, s=1:1000, (:h, :l)]); # output power of a pwl-cost generator (Δ-formulation)

julia> power_abs(p, t, s) = p[t, s, :h] + p[t, s, :l]; # the naive method: there are few terms

julia> power_abs(p, t, s) = JuMP.@expression(model, p[t, s, :h] + p[t, s, :l]); # add a macro wrapper

The function power_abs is supposed to be used at many other places. In this case intuitively I think the naive methods work better?

9

is

power_abs(p, t, s) = p[t, s, :h] + p[t, s, :l]
JuMP.@expression(model, sum(power_abs(p, t, s) for t=1:24 for s=1:1000))

equivalent to

JuMP.@expression(model, sum(p[t, s, :h] + p[t, s, :l]) for t=1:24, s=1:1000)

? conceptually, or in terms of performance? I wonder.

10

is … equivalent to … conceptually, or in terms of performance?

The power_abs version will be slower. The macros do not recurse into function calls.

And the allocation comparison contradicts the actual observation.

I’ll have a look. This might be a compilation issue, or it might be due to a change in Julia.

Edit: it seems like this is only present in the dev docs: Performance tips · JuMP. stable gives the correct result: Performance tips · JuMP

1 Like
11

Yes, as I observed

import JuMP
const model = JuMP.Model();
const T = 24
const S = 10000
JuMP.@variable(model, p[t=1:T, s=1:S, (:h, :l)]);

a(p, t, s) = p[t, s, :h] + p[t, s, :l] # simplest
a(model, p, t, s) = JuMP.@expression(model, p[t, s, :h] + p[t, s, :l]) # contrived: add a @expression wrapper 

e_verbose(model, p, T, S) = JuMP.@expression(model, sum(p[t, s, :h] + p[t, s, :l] for t=1:T for s=1:S))
e_concise_function(model, p, T, S) = JuMP.@expression(model, sum(a(p, t, s) for t=1:T for s=1:S))
e_concise_macro_wrapped(model, p, T, S) = JuMP.@expression(model, sum(a(model, p, t, s) for t=1:T for s=1:S))

function test(model, p, T, S)
    println(@allocated(e_verbose(model, p, T, S)))
    println(@allocated(e_concise_function(model, p, T, S)))
    println(@allocated(e_concise_macro_wrapped(model, p, T, S)))
    println(@allocated(e_verbose(model, p, T, S)))
    println(@allocated(e_concise_function(model, p, T, S)))
    println(@allocated(e_concise_macro_wrapped(model, p, T, S)))
end;

julia> test(model, p, T, S)
36909440
136749408
136749408
36909440
136749408
136749408

But I just need a submodule like power_abs because that’s the actual physical quantity (the original p construct was defined to express the PWL costs…). I need to frequently refer to power_abs in a lot of other sorts of constraints. If I write verbose expressions manually at all places, my code doesn’t look clean. Any opinions?

12

Use power_abs and don’t worry about the performance.

1 Like
13

I think it could be explicitly modeled as a variable.

Code 
add_pwl_power(model, GD, T, UB) = JuMP.@variable(model, # GD = Reserve || Other
    [z=GD.Zone, g=eachindex(GD.N[z]), t=1:T, h=(:h, :l)],
    lower_bound = 0, upper_bound = UB[h][z][g]
)
add_power(model, GD, T, ::Symbol) = JuMP.@variable(model,
    [z=GD.Zone, g=eachindex(GD.N[z]), t=1:T],
)
prepare_ub(GD) = (
    h = [[(1-GD.pCost[z].KneeFraction[g])GD.PMax[z][g] for g=eachindex(GD.N[z])] for z=GD.Zone],
    l = [[    GD.pCost[z].KneeFraction[g]GD.PMax[z][g] for g=eachindex(GD.N[z])] for z=GD.Zone]
)
add_pwl_power(model, GD, T) = add_pwl_power(model, GD, T, prepare_ub(GD))
add_power(model, GD, T) = (
    i = add_power(model, GD, T, :identity),
    c = add_pwl_power(model, GD, T) 
)
add_c2i_constr(model, GD, T, p) = JuMP.@constraint(model,
    [z=GD.Zone, g=eachindex(GD.N[z]), t=1:T],
    p.i[z,g,t] == p.c[z,g,t,:h] + p.c[z,g,t,:l] # TODO: add a nonzero PMin
)
function add_power_and_c2i_constrs(model, GD, T)
    p = add_power(model, GD, T)
    add_c2i_constr(model, GD, T, p)
    p
end