Add single precision option to the dycore (intermediate state)

[starkphi]

Adds the option to set wp/vpfloat to single via the env var FLOAT_PRECISION (previous options were 'double' and 'mixed') to run (at least) the dycore in single precision.

This PR was created to later reduce merge complexity of #886.

Changes

• Replace native python float types by eighter wpfloat or vpfloat
  • Change type annotations where there were native Python floats were present (which are by default passed as gtx.float64 to gtx.field_operators and gtx.programs) and add casts around float literals
  • Switch from wpfloat or vpfloat and vice versa if found inconsistency with other operators (to hopefully make fixing mixed precision easier in the future)
  • Directly import wpfloat and vpfloat (instead of ta.wpfloat -> more concise)
  • Replace DBL_EPS by WP_EPS and VP_EPS
• Setup pytests for single precision
  • Add pytest marker single_precision_ready (running a test with FLOAT_PRECISION=single deselects all tests without)
    • Add it to test_velocity_advection.py::test_compute_advection_in_vertical_momentum_equation and test_apply_diffusion_to_vn.py to have first two pytests for CI
  • Add nox argument --single-precision for CI
  • Remove pytest option --enable-mixed-precision (wpfloat or vpfloat are already set according to FLOAT_PRECISION at import time of a pytest file)

TODO

• minimal testing: add the .test_single_precision (default.yml) test to CI
• Find and solve bug that appears when running with the compiled self._calculate_divdamp_fields (introduced in Make calculate_divdamp_fields a compiled program #888) in combination with single precision.

[muellch]

cscs-ci run default

[egparedes]

It looks good in general although I have a few comments.

[egparedes]

Just a question: is there are any particular reason why the conversion is done here instead of at initialization in the post_init method?

[starkphi]

Not really. When I changed this I think I just wanted to get it to run in single as soon as possible and did not always check for the cleanest way to do things.
You mean I should change lines 144-154 to the following?

    def __post_init__(self, vct_a, vct_b):
        object.__setattr__(
            self,
            "_vct_a",
            gtx.astype(vct_a, wpfloat),
        )
        object.__setattr__(
            self,
            "_vct_b",
            gtx.astype(vct_b, wpfloat),
        )

I didn't work with dataclasses before and was not aware of __post_int__().

[egparedes]

Question: is this marker needed only in a transition phase? Otherwise, why is actually needed?

[starkphi]

I thought that it might be a good idea to be able to have all tests skipped that nobody prepared for single precision. So I suggest using this marker until all tests have suitable tolerances defined for single precision operation.

[OngChia]

Out of curiosity, why do you import vpfloat, wpfloat explicitly instead of doing the usual way of importing the module type_alias type_alias as ta, and then when you need vpfloat or wpfloat, you just do ta.vpfloat or ta.wpfloat? I think this is the most common way of using vpfloat and wpfloat in icon4py.
Except for some cases in field_operator, for example, broadcast(wpfloat("0.0", (dims.CellDim,))) where you have to use vpfloat or wpfloat directly.

[starkphi]

I just saw that in many dycore stencils it was done this way (directly using wpfloat/vpfloat) and in other cases the other. So we already have some kind of inconsistency in the way these types are imported in the repo. I preferred to use the shorter version because I think it looks less cluttered (especially in cases like self.cfl_w_limit: vpfloat = vpfloat(0.65)) and I don't see a lot of descriptive value in the additional ta. (type_alias).

[OngChia]

Could it be better with gtx.float64 | gtx.float32?

[starkphi]

I first changed it to only using wpfloat which is what the instance variable type is, but because the default arguments are in float, and also because ruff doesn't like function calls like fourth_order_divdamp_z: wpfloat = wpfloat(32500.0) in default arguments, I thought that like this it would show that it can also take native python floats.
In my opinion gtx.float64 | gtx.float32 could almost be a bit misleading.

[OngChia]

I think I would rather change instance variables in GlobalGridParams to wpfloat, because these grid properties and other physical constants should have working precision.

[starkphi]

You're right, that would be nicer. I don't know how these instances are constructed and where the casts to wpfloat would have to happen. Do you know how this can be done on dataclasses? (Assuming we don't change anything wherever these instances are constructed but just add casts to GlobalGridParams.)

[OngChia]

Here you can directly import type_alias from icon4py.model.common import type_alias as ta. Similarly in other places.

[starkphi]

I could, but as stated in my answer on about the same question in solve_nonhydro.py I think it looks less cluttered with a direct import.

[OngChia]

Similarly here, is it better with gtx.float64 | gtx.float32? It seems clearer in this way. So I know the initialization of diffusion granule accepts either single or double precision for these parameters.

[starkphi]

Why do you think it is clearer? Wouldn't you anyway directly give the respective wpfloat to it (gtx.float64 or gtx.float32 depending on precision)?

[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.