Fix inlining of generated functions

[dennisYatunin]

This PR fixes the inlining of generated functions, which will hopefully prevent compilation from hanging in the non-orographic gravity wave parametrization test in ClimaAtmos.

Normally, inlining works by inserting an inline annotation at the beginning of a function definition:

julia> @macroexpand @inline add_one_expr(x) = Expr(:call, :+, :x, 1)
:(add_one_expr(x) = begin
          $(Expr(:meta, :inline))
          #= REPL[13]:1 =#
          Expr(:call, :+, :x, 1)
      end)

When the function is generated, though, the @inline macro inserts two such annotations (one at the start of the generated function and one in the expression that generates it):

julia> @macroexpand @inline @generated add_one(x) = Expr(:call, :+, :x, 1)
:(add_one(x) = begin
          $(Expr(:meta, :inline))
          if $(Expr(:generated))
              $(Expr(:meta, :inline))
              #= REPL[10]:1 =#
              Expr(:call, :+, :x, 1)
          else
              $(Expr(:meta, :generated_only))
              return
          end
      end)

This type of expression is probably not supported by the compiler, and I think it is responsible for the strange behavior we have been seeing in ClimaAtmos. Specifically, we have observed that replacing the definition

@inline @generated function val_unrolled_reduce(op, ::Val{N}, init) where {N}
    return foldl((:init, 1:N...)) do prev_op_expr, item_expr
        :(op($prev_op_expr, $item_expr))
    end
end

with an equivalent version that uses assignment syntax

@inline @generated val_unrolled_reduce(op, ::Val{N}, init) where {N} =
    foldl((:init, 1:N...)) do prev_op_expr, item_expr
        :(op($prev_op_expr, $item_expr))
    end

causes compilation to hang. Having the double inline annotation here is probably forcing the compiler to enter some unhandled infinite loop or wait condition for this complex test (though I don't think we're seeing a compiler bug here because such a use case probably isn't supported).

In order to correctly inline generated functions, we need to only insert the inline annotation at the start of the generated function expression. For example,

julia> @macroexpand @generated add_one(x) = Expr(:block, Expr(:meta, :inline), Expr(:call, :+, :x, 1))
:(add_one(x) = begin
          #= REPL[14]:1 =#
          if $(Expr(:generated))
              #= REPL[14]:1 =#
              Expr(:block, Expr(:meta, :inline), Expr(:call, :+, :x, 1))
          else
              $(Expr(:meta, :generated_only))
              return
          end
      end)

[charleskawczynski]

This looks fine to me.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 57.48%. Comparing base (aa04164) to head (8443ddf).

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #13      +/-   ##
==========================================
+ Coverage   56.57%   57.48%   +0.91%     
==========================================
  Files           8        8              
  Lines         251      247       -4     
==========================================
  Hits          142      142              
+ Misses        109      105       -4     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[dennisYatunin]

I have tried this out in CliMA/ClimaAtmos.jl#3313, and unfortunately this does not actually solve the strange compilation issue mentioned above. I originally defined the first method of the val_unrolled_reduce function in this PR as follows:

@generated val_unrolled_reduce(op, ::Val{N}, init) where {N} = Expr(
    :block,
    Expr(:meta, :inline),
    foldl((prev_op_expr, item_expr) -> :(op($prev_op_expr, $item_expr)), (:init, 1:N...)),
)

This caused compilation to hang for the non-orographic gravity wave parameterization tests in Buildkite Run 20622. I then tried rewriting this using function ... end form, instead of assignment form:

@generated function val_unrolled_reduce(op, ::Val{N}, init) where {N}
    return Expr(
        :block,
        Expr(:meta, :inline),
        foldl((prev_op_expr, item_expr) -> :(op($prev_op_expr, $item_expr)), (:init, 1:N...)),
    )
end

This also caused compilation to hang for the non-orographic gravity wave parameterization tests in Buildkite Run 20623. Finally, I returned this function to the implementation it had before this PR (but without the @inline annotation):

@generated function val_unrolled_reduce(op, ::Val{N}, init) where {N}
    return foldl((:init, 1:N...)) do prev_op_expr, item_expr
        :(op($prev_op_expr, $item_expr))
    end
end

This version took less than a minute to compile in Buildkite Run 20624.

Although this PR does not resolve the compilation mystery, this is still a change that we should make in order to correctly inline generated functions. I'll go ahead and merge this in.