- Shaped by: Till Ehrengruber
- Appetite (FTEs, weeks):
- Developers: <!-- Filled in at the betting table unless someone is specifically required here -->
## Problem
<!-- The raw idea, a use case, or something we’ve seen that motivates us to work on this -->
Some configuration values of a model are only available at runtime, but beside that constant. Examples are a particular choice of advection scheme to use, whether a limited area model is to be used and many more. For performance reasons and to allow the use of features that are not fully supported by the compiled backends (e.g. multi-line if-statements) we want to be able to partially apply these values to a `@program`. That way they can be constant folded by the toolchain, giving better performance and allowing the elimination of some constructs that for the case of if-statements make the result compatible with the compiled backend.
The propose syntax to bind arguments is as follows:
```python
@program
def my_program(use_lam: bool, inp: Field[...], ...):
my_field_op(use_lam, ...)
my_lam_program = my_program.with_bound_args(use_lam=True) # note: only scalar should allowed
my_lam_program(inp, ...) # `use_lam` is not part of the signature anymore
# semantically equivalent to
my_lam_program = functools.partial(my_lam_program, use_lam=True)
my_lam_program(inp, ...)
```
## Appetite
<!-- Explain how much time we want to spend and how that constrains the solution -->
1-2 weeks. For someone with only little or no knowledge of ITIR around 2-3 weeks.
## Solution
<!-- The core elements we came up with, presented in a form that’s easy for people to immediately understand -->
The proposed solution consists of two components that can be developed in isolation (and merged as seperate PRs).
### Component 1 - `ProgramWithBoundArgs`
First the frontend needs to be extended such that `Program.with_bound_args` returns an callable that behaves like a program with the arguments being bound. We propose to introduce a new class `ProgramWithBoundArgs` that inherits from `gt4py.next.ffront.decorator.Program`. The bound arguments in `_process_args` are forwarded such that the signature checking etc. works as usual. The `itir` method then just binds the arguments using a let-statement, effectively shadowing the fencil arguments.
```python
class Program:
...
def with_bound_args(**kwargs) -> ProgramWithBoundArgs:
# TODO
class ProgramWithBoundArgs(Program): # TODO: maybe inheritance is not the best solution
def _process_args(*args, **kwargs):
# TODO: enhance arguments by bound arguments and forward them
# ...
return super()._process_args(*new_args, **new_kwargs)
def itir():
itir = super().itir()
# TODO: transform itir such that argument is bound (see idea below)
return new_itir
```
Where `itir()` transforms the original ITIR from (lisp s-expr notation)
```
(lambda (use_lam inp ...)
body)
```
into
```
(lambda (use_lam inp ...)
(let (use_lam True)
body))
```
### Component 2 - Constant folding pass
To optimize the ITIR again a constant folding pass/transformation should be added.
__Position in the pipeline__
Should be executed after the first function inliner (such that the bound arguments are propagated), but before `collapse_tuple`, `inline_lifts` passes.
__Supported expressions__
Most importantly if statements should be folded, e.g:<br>
`if_(True, true_branch, false_branch)` -> `true_branch`
Additionally the pass should allow constant folding for calls where all arguments are literals, e.g.:<br>
`1 + 1` -> `2`
### Add-on - Allow bindings argument of field operators
If time permits an addition that permits bindings arguments of field operators could be implemented. Please peer with someone for the design before implementing.
```python
@field_operator
def my_field(use_lam, inp, ...):
...
my_lam_field_op = my_field_op.with_bound_args(use_lam=True)
```
## Rabbit holes
<!-- Details about the solution worth calling out to avoid problems -->
## No-gos
<!-- Anything specifically excluded from the concept: functionality or use cases we intentionally aren’t covering to fit the ## appetite or make the problem tractable -->