calling gtscript functions from conditionals

[rheacangeo]

Perhaps this should not be supported, but when a gtscript function is called inside of a conditional, often an error raises "Temporary field {name} implicitly defined within run-time if-else region.", even if there are no temporaries in the function, if it's just using and assigning fields passed from the input arguments to the stencil... I assume using a function causes a temporary variable to be generated, and since that is not declared before the conditional, it is not allowed.
Here's and example:

@gtscript.function
def qcon_func(qcon, q0_liquid + q0_rain):
    qcon = q0_liquid  + q0_rain
    return qcon

then the stencil

with computation(BACKWARD), interval(0, -1):
    if ri < ri_ref:
        qcon = qcon_func(qcon, q0_liquid, q0_rain)

This works if instead of assigning qcon in the function

@gtscript.function
def qcon_func(qcon, q0_liquid, q0_rain):
    return q0_liquid  + q0_rain

So it is totally reasonable in this case to say do not assign qcon inside the function. But more complex functions can't be returned as one-line operations, and assign an input field that wouldn't cause problems inside the stencil spec. And the field being assigned is not temporary, it is just not known in the context of the function?

This is a lower order issue, but has tripped me up a couple of times so thought I'd bring it up! It's nice to be able to reuse functions, but of course less essential than having all the other functionality.

[ofuhrer]

Probably this issue is also the same issue. The following code does not work...

@gtscript.function
def abs_fn(a):
    b = a
    return b
...
    with computation(PARALLEL), interval(...):
        if a > 0:
            b = abs_fn(a)

[jdahm]

The following code works with the GTC backends, and the generated code looks correct, so I'll close for now. Let us know if you still have problems and we can re-open.

from gt4py import gtscript
import numpy as np


@gtscript.function
def abs_fn(a):
    b = a
    return b


@gtscript.stencil(backend="gtc:gt:cpu_ifirst")
def stencil(
    input_field: gtscript.Field[gtscript.IJK, np.int32],
    output_field: gtscript.Field[gtscript.IJK, np.int32],
    q_max: np.int32,
):
    with computation(PARALLEL), interval(...):
        if input_field > q_max:
            output_field = abs_fn(input_field)
        else:
            output_field = input_field