Annotation to elide function call if the return value is unused

Is there some annotation that I can use to give permission to the compiler to elide e.g. x = foo(y) if x is not used? And would it possibly work across function-call boundaries? I’d like to do something like

function foo(input)
    println("foo")
    2 * input
end

function bar(input)
    println("bar")
    2 + input
end

function baz(x, y)
    println("baz")
    x / y
end

function compute(input, callback)
  x = foo(input)
  y = bar(input)
  z = baz(x, y)
  return callback((;x, y, z))
end

and depending on what fields callback accesses, perhaps only foo is called. Or only bar. Or both. Or all three.
Concretely, currently:

julia> compute(10, nt -> nt.x)
foo
bar
baz
20

but I’d like to annotate the function calls so that only foo is printed.

I could turn all computations into thunks:

function foo_thunk(input)
    ()->(println("foo"); 2 * input())
end

function bar_thunk(input)
    ()->(println("bar"); 2 + input())
end

function baz_thunk(x, y)
    ()->(println("baz"); x() / y())
end

function compute_thunk(input, callback)
  x = foo_thunk(()->input)
  y = bar_thunk(()->input)
  z = baz_thunk(x, y)
  return callback((;x, y, z))
end
julia> compute_thunk(10, nt -> nt.x() + nt.z())
foo
baz
foo
bar
21.666666666666668

but I do not want foo to be computed twice. Caching/memoizing the thunks would be a solution, but I do not think I’m looking for that. I want to avoid the dynamic memory allocations and dynamic checks for whether values have already been computed. Otherwise I’m relying on other compiler optimizations like escape analysis, which would likely give suboptimal results due to the function-call boundary. (Though for any of this to make sense, callback has to be inlined anyway)

Truly, I’m looking for a solution that is basically equivalent to deleting any unused fields from the NamedTuple, and deleting the lines where the values for those fields are computed.

I don’t expect anything to be elided if there is e.g. any getproperty(nt, fieldname), for a fieldname that is not marked Core.Const.

I’m all but certain that Base.@pure is the wrong answer.

2 Likes

You can abuse some internals of lowering for this, but this may not always be guaranteed to work. Code wrapped in @elidable may be removed in lowering, if its result is unused.

julia> macro elidable(ex)
           return Expr(:unnecessary, ex)
       end
@elidable (macro with 1 method)

julia> function f(x)
           @elidable println("foo")
           x + 1
       end
f (generic function with 1 method)

julia> f(2)
3
2 Likes

:-o !

For reference: https://github.com/JuliaLang/julia/blob/442f727c65e27d794b0f0555c66ed07228b5546a/src/julia-syntax.scm#L4160-L4165

You can do this without direct help from the compiler, if you pass relevant closures to callback and use memoization to eliminate repeated calculations for the same input:

using Memoize
@memoize function foo(input)
    println("foo")
    2 * input
end

@memoize function bar(input)
    println("bar")
    2 + input
end

@memoize function baz(x, y)
    println("baz")
    x / y
end

function compute(input, callback)
  x() = foo(input)
  y() = bar(input)
  z() = baz(x(), y())
  return callback((;x, y, z))
end

julia> compute(1, nt -> nt.x())
foo
2

julia> compute(2, nt -> nt.x() + nt.z())
foo
bar
baz
5.0

The macro as written doesn’t resolve the variables correctly:

@macroexpand function compute(input, callback)
         x = @elidable foo(input)
         y = @elidable bar(input)
         z = @elidable baz(x, y)
         return callback((;x, y, z))
       end
:(function compute(input, callback)
      #= REPL[3]:1 =#
      #= REPL[3]:2 =#
      x = $(Expr(:unnecessary, :(Main.foo(Main.input))))
      #= REPL[3]:3 =#
      y = $(Expr(:unnecessary, :(Main.bar(Main.input))))
      #= REPL[3]:4 =#
      z = $(Expr(:unnecessary, :(Main.baz(Main.x, Main.y))))
      #= REPL[3]:5 =#
      return callback((; x, y, z))
  end)

defining it as

macro elidable(ex)
    @show ex
    return Expr(:unnecessary, esc(ex))
end

seems to do the trick.

I know that all bets are off at this point, but I’m not getting the desired behavior:

function foo(input)
    println("foo")
    2 * input
end

function bar(input)
    println("bar")
    2 + input
end

function baz(x, y)
    println("baz")
    x / y
end

function compute(input, callback)
  x = @elidable foo(input)
  y = @elidable bar(input)
  z = @elidable baz(x, y)
  return callback((;x, y, z))
end

@inline callback_xyz(nt) = nt.x + nt.y + nt.z
@inline callback_xy(nt) = nt.x + nt.y
@inline callback_xz(nt) = nt.x + nt.z
@inline callback_x(nt) = nt.x
julia> @code_typed compute(20, callback_x)
CodeInfo(
1 ─      invoke Main.println("foo"::String)::Any
│   %2 = Base.mul_int(2, input)::Int64
│        invoke Main.println("bar"::String)::Any
│        invoke Main.println("baz"::String)::Any
└──      return %2
) => Int64

(@yha as I mentioned in my original post, I’m not looking to do memoization.)

This behavior is expected, the elision is just a naive pass during lowering but inlining happens after that, so it won’t do anything across function boundaries.

1 Like