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

[github-actions]

Mandatory Tests

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

• cscs-ci run default

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