Combine divergence damping coefficient calculation with other stencils

[muellch]

No description provided.

[edopao]

@havogt I see an error in the gtfn lowering but not in dace, running model/atmosphere/dycore/tests/dycore/stencil_tests/test_dycore_utils.py::test_calculate_divdamp_fields :

self = <gt4py.next.program_processors.codegens.gtfn.itir_to_gtfn_ir._CannonicalizeUnstructuredDomain object at 0x12df50b20>
node = FunCall(fun=SymRef(id=SymbolRef('unstructured_domain')), args=[FunCall(fun=SymRef(id=SymbolRef('named_range')), args=[...hape=None)), SymRef(id=SymbolRef('out'))]), AxisLiteral(value='K', kind=<DimensionKind.VERTICAL: 'vertical'>)])])])])])

    def visit_FunCall(self, node: itir.FunCall) -> itir.FunCall:
        if node.fun == itir.SymRef(id="unstructured_domain"):
            # for no good reason, the domain arguments for unstructured need to be in order (horizontal, vertical)
            assert isinstance(node.args[0], itir.FunCall)
            first_axis_literal = node.args[0].args[0]
            assert isinstance(first_axis_literal, itir.AxisLiteral)
            if first_axis_literal.kind == itir.DimensionKind.VERTICAL:
>               assert len(node.args) == 2
E               AssertionError

../edopao/gt4py/src/gt4py/next/program_processors/codegens/gtfn/itir_to_gtfn_ir.py:215: AssertionError

[edopao]

I see a speedup in compute_advection_in_vertical_momentum_equation using the dace backend, it could be related to the watermark of the memory pool or some other side-effect of not launching the divdamp kernel:

[edopao]

I have updated gt4py to include this fix: GridTools/gt4py#2429

[edopao]

cscs-ci run default

[edopao]

@havogt I have been able to work around the GTFN issue by applying further inlining.

[edopao]

cscs-ci run default

[edopao]

cscs-ci run default

[edopao]

Please check if you agree with this change, from an application-code point of view. If you agree, the next step is to update the serialized data to fix the failing test.

[muellch]

Looks good!
I cannot approve since I am the original PR author.

[OngChia]

Can we simply call en_smag_fac_for_zero_nshift to compute interpolated_fourth_order_divdamp_factor (enh_divdamp_fac in icon) in this data test. Because, it is an intermediate variable to compute reduced_fourth_order_divdamp_coeff_at_nest_boundary and fourth_order_divdamp_scaling_coeff, and these two variables are already tested against serialized data in test_validate_divdamp_fields_against_savepoint_values, which implies thaten_smag_fac_for_zero_nshift is correct and we do not need to serialize enh_divdamp_fac

[edopao]

I will follow this suggestion so we can unblock this PR, @muellch please note it.

[edopao]

cscs-ci run default

[edopao]

cscs-ci run default

[edopao]

cscs-ci run distributed

[github-actions]

Mandatory Tests

Please make sure you run these tests via comment before you merge!

• cscs-ci run default
• cscs-ci run distributed

Optional Tests

To run benchmarks you can use:

• cscs-ci run benchmark-bencher

To run tests and benchmarks with the DaCe backend you can use:

• cscs-ci run dace

To run test levels ignored by the default test suite (mostly simple datatest for static fields computations) you can use:

• cscs-ci run extra

For more detailed information please look at CI in the EXCLAIM universe.

[edopao]

cscs-ci run default

[edopao]

cscs-ci run distributed