cartesian[dace]: SDFG validation warnings for read after write

[CharlesKrop]

Under specific situations (see attached file for example stencil), an erroneous DaCe validation warning for "uninitialized transient" is raised. Insight from @romanc revealed that this is likely tied to some issue identifying input/output fields when building the SDFG.

See attached sdfg_validation_issue.py file for an example stencil which erroneously produces this warning, and the associated program.sdfgz for this stencil.

sdfg_validation_issue.py
program.sdfgz.zip

[romanc]

Current best guess: This is probably our fault in not considering read after write conditions when freezing the sdfg. To be investigated as part of cleanup in July.

[CharlesKrop]

ty for the clean up

[romanc]

Well ... turns out, DaCe was lying to us because we were relying on

class SDFG:
    ...
    def read_and_write_sets(self) -> Tuple[Set[AnyStr], Set[AnyStr]]:
        """
        Determines what data containers are read and written in this SDFG. Does
        not include reads to subsets of containers that have previously been
        written within the same state.
        """

source

which is just not true (anymore) since spcl/dace#1678 change the behavior of the "state-level" function to not consider read after write situations anymore as this note explains:

        # NOTE: In a previous version a _single_ read (i.e. leaving Memlet) that was
        #   fully covered by a single write (i.e. an incoming Memlet) was removed from
        #   the read set and only the write survived. However, this was never fully
        #   implemented nor correctly implemented and caused problems.
        #   So this filtering was removed.

source

[romanc]

I had another look and there at two things to note:

1. This is actually an NDSL-level issue.
2. DaCe is right to complain.

Consider the following minimal reproducer: (Yes, this is NDSL code: as said, this is an NDSL-level problem.)

from ndsl.dsl.gt4py import computation, interval, PARALLEL
from ndsl.boilerplate import get_factories_single_tile
from ndsl.constants import I_DIM, J_DIM, K_DIM
from ndsl.dsl.typing import FloatField
from ndsl.stencils.basic_operations import copy
from ndsl import StencilFactory, NDSLRuntime, QuantityFactory, Backend


class Code(NDSLRuntime):
    def __init__(self, stencil_factory: StencilFactory, quantity_factory: QuantityFactory):
        super().__init__(stencil_factory)
        self._copy_local = stencil_factory.from_dims_halo(
            func=copy,
            compute_dims=[I_DIM, J_DIM, K_DIM],
        )
        self._my_local = self.make_local(quantity_factory, [I_DIM, J_DIM, K_DIM])
        # bad access of local in non-orchestrated code path
        self._my_local[:] = 2.0

    def __call__(self, in_field: FloatField):
        self._copy_local(self._my_local, in_field)


if __name__ == "__main__":
    domain = (2, 2, 5)

    for backend in ["st:dace:cpu:IJK", "orch:dace:cpu:IJK"]:
        stencil_factory, quantity_factory = get_factories_single_tile(
            nx=domain[0], ny=domain[1], nz=domain[2], nhalo=0, backend=Backend(backend)
        )

        out_array = quantity_factory.ones(dims=[I_DIM, J_DIM, K_DIM], units="inputs")

        code = Code(stencil_factory, quantity_factory)
        code(out_array)

        assert (out_array.field == 2.0).all(), f"Backend '{backend}' not okay"

The backends orch:dace:cpu:IJK and st:dace:cpu:IJK will return different results in the above test case. This is because

• in orchestration, DaCe is in charge of memory management for Locals. And because locals are transient in DaCe lingo, they are only allocated on the cpu-side of things, which means that self._my_local is indeed undefined when used inside the __call__ function and DaCe has a right to complain about that.The 2.0 from the __init__ never makes it into the orchestrated path.
• in stencil mode, python keeps being in charge and the memory and the last written value (in that case 2.0) is there when used inside the call function.

NDSL PR NOAA-GFDL/NDSL#509 will deal with this problem by not allowing usage of Locals in non-orchestrated code paths.

[romanc]

Aaand ... there is a case for this issue to be in GT4Py: Consider NDSL runtime code which calls the same stencil multiple times. Furthermore, this stencil has 2d temporaries, but they aren't 2d temporaries yet, but instead passed from outside in as 2d locals, which are reset every time before the stencil starts. These locals are then (unexpectedly) marked as inputs of the stencils (because freezing doesn't consider read after write) and then the transient reads before write are reported in sdfg.validate() on the encompassing program.sdfgz (see attached repro from fillz before it was cleaned up to use 2d temporaries in NOAA-GFDL/pyFV3#150).

repro_fillz.zip