feature: DataDimensionsField as argument to functions

[romanc]

Description

🧱 blocked by

• feat[cartesian]: lazy annotated gtscript functions GridTools/gt4py#2708
• refactor[cartesian]: skip LiftTrivalIf in sdfg.simplify() pipeline GridTools/gt4py#2710

with the following todos left

• update gt4py/dace submodules to a version that contains the above changes
• skip LiftTrivialIfs pass in NDSL calls to sdfg.simplify()
• with the gt4py/dace update we can also include c2bd78d

---

This PR adds support for data dimension fields as type annotations in functions (as we can do in stencils).

This PR depends on the following gt4py feature GridTools/gt4py#2708. I have updated that submodule in this PR to show how it will look like, but this PR will be in a draft state until the corresponding gt4py PR is merged.

Related issue: #504

How has this been tested?

New test cases added.

Checklist

• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation (e.g. add new modules to docs/docstrings/): N/A
• My changes generate no new warnings
• Any dependent changes have been merged and published in downstream modules
  Depends on feat[cartesian]: lazy annotated gtscript functions GridTools/gt4py#2708.
• New check tests, if applicable, are included

[romanc]

@CharlesKrop can you look at example usage in tests/quantity/test_data_dimensions_field.py and yell if this isn't what the issue was about?

[twicki]

just as a written trace of our discussion this morning:

1. The typehints are as far as we know dropped anyways, so we can make the function drop it so it does not bother the user? Or should we start caring about them?
2. Typehints of tempoaries are VERY limited in functions, so changing this would never affect that

[romanc]

More info on what actually complains: the gtscript parser restricts arguments of functions and stencils to _FieldDesriptor or a list of valid data types (including common booleans, integers, floats) in

https://github.com/GridTools/gt4py/blob/59a5260edd9f4d140ddf44de61590b3dcda8877e/src/gt4py/cartesian/frontend/gtscript_frontend.py#L2151-L2165

Since markup types aren't resolved at python parsing time (when the decorator kicks of function annotation), we end up with DataDimensionsMarkupType for DataDimensionsFields, which is rejected by the restriction linked above. I think it's fair game of gtscript to ensure that whatever is passed into a gtscript function/stencil can be handled by the stack that follows. Imo, blindly ignoring unknown types isn't wise.

More info to follow after lunch. Happy to discuss again.

[romanc]

Seconds attempt on gtscript.lazy_function() is much simpler branch comparison after I finally understood a bit better how to the preprocessing of stencils in the GTScriptParser is working.

[romanc]

Investigated the sudden increase in time it takes to parse / simplify (i.e. orchestrated acoustics took almost 5h (instead of the usual 30min) in this run after updating dace/gt4py). slowdown seems to be due to a new simplify pass LiftTrivialIf and isn't related to the lazy annotation of gtscript fuctions. To mitigate, we can just skip LiftTrivialIf for now.

[CharlesKrop]

Just getting back to this now, sorry for the delay. Looks like your tests cover #504 pretty well

[romanc]

Just getting back to this now, sorry for the delay. Looks like your tests cover #504 pretty well

Thanks for checking. And no worries, the PR is anyway blocked by downstream changes that need to be merged in GT4Py first.