Skip to content

Refactor `@with_signature` to allow iterations over vectors and custom ranges

Florian Atteneder requested to merge fa/rework-with-signature into main

Atm @with_signature functions are mainly used through broadcast_[volume,face,bdry]. The latter methods do the unpacking and also control the for loops over the relevant indices. The @with_signature functions only accept tuples of arguments to process data.

Benchmarks showed that this usage of @with_signature bears noticeable overhead, due to tuple packing and unpacking needed to pass and receive values when calling @with_signature functions.

A more straight forward implementation of the same equations (only considered llf, av_nflux_2d) by copy pasting directly the function bodies of @with_signature into the loops showed worthwhile performance improvements.

This PR reworks (or extends) the @with_signature macro so that it can now generate the loops around the function bodies.

Further improvements:

  • allow to customizing @with_signature defs through options like @with_signature [opt1=true,...] function ...
  • allow to add options to @accepts, @returns keywords, e.g. @accepts [Prefix(lhs),bdry] u1, u2
  • for the moment, keep the current behavior of @with_signature working; the new one can be enabled using @with_signature [legacy=false]
  • deprecate the State options for the @accepts macro for the non-legacy version of @with_signature
  • allow to iterate over two different sets of indices at the same time (usually volume and bdry indices); e.g.
@with_signature [legacy=false] function flux(eq::Equation)
   @accepts fx, fy, u, v
   @accepts [Prefix(out),bdry] fx, fy, u, v
   @accepts [bdry] nx, ny
   nflx = nx*fx + ny*fy
   out_nflx = nx*out_fx + ny*out_fy
   nflx = LLF(nflx, out_nflx, u, out_u, absmax(v, out_v))
   @returns [bdry] nflx
end

This can be called with flux(eq, [fx, fy, u, v, fx, fy, u, v, nx, ny], [nflx], #=volume indices=# 1:100, #=bdry indices=#100:-1:1).

It remains to rework the broadcast methods, but I will do this in a separate PR.

Merge request reports