LoopBlocking produces non-deterministic SDFGs (memory address in a Tasklet name, set-ordered node processing)

[havogt]

LoopBlocking
(src/gt4py/next/program_processors/runners/dace/transformations/loop_blocking.py)
produces a different SDFG from one process to the next, for two independent reasons.
Both are in _rewire_map_scope, and both reach the serialized SDFG, which gt4py's
compile cache is keyed on.

1. A memory address ends up in a Tasklet name

copy_tlet = state.add_tasklet(
    name=f"loop_blocking_copy_tlet_{independent_node.data}_{id(out_edge)}",
    ...
)

id() is the object's address, so the name differs on every run. A Tasklet's label is
a serialized property, so it lands in the SDFG.

TASKLET_NAME: ['loop_blocking_copy_tlet_t_130121870367616']
TASKLET_NAME: ['loop_blocking_copy_tlet_t_136689034664928']
TASKLET_NAME: ['loop_blocking_copy_tlet_t_131992464874080']

Reproduction (three runs, three names)

import dace
from dace.sdfg import nodes as dace_nodes
from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations

sdfg = dace.SDFG("sdfg_with_direct_output_access_node")
state = sdfg.add_state(is_start_block=True)
sdfg.add_array("A", shape=(40,), dtype=dace.float64, transient=False)
sdfg.add_array("B", shape=(40, 10), dtype=dace.float64, transient=False)
sdfg.add_scalar("t", dtype=dace.float64, transient=True)
A, B, t = (state.add_access(name) for name in "ABt")
me, mx = state.add_map("main_comp", ndrange={"__i0": "0:40", "__i1": "0:10"})
tlet = state.add_tasklet("inner_tasklet", inputs={"__in"}, outputs={"__out"},
                         code="__out = __in + 10.0")
state.add_edge(A, None, me, "IN_A", dace.Memlet("A[0:40]"))
state.add_edge(me, "OUT_A", tlet, "__in", dace.Memlet("A[__i0]"))
me.add_scope_connectors("A")
state.add_edge(tlet, "__out", t, None, dace.Memlet("t[0]"))
state.add_edge(t, None, mx, "IN_B", dace.Memlet("B[__i0, __i1]"))
state.add_edge(mx, "OUT_B", B, None, dace.Memlet("B[0:40, 0:10]"))
mx.add_scope_connectors("B")
sdfg.validate()

sdfg.apply_transformations_repeated(
    gtx_transformations.LoopBlocking(blocking_size=2, blocking_parameters=["__i1"]),
    validate=True, validate_all=True)

print(sorted(n.label for n in state.nodes()
             if isinstance(n, dace_nodes.Tasklet)
             and n.label.startswith("loop_blocking_copy_tlet_")))

A counter, or a name derived from the edge's endpoints and connectors, would be stable.

2. Independent nodes are processed in set order

_independent_nodes: Optional[set[dace_nodes.AccessNode]]
...
for independent_node in self._independent_nodes:
    ...
    copy_tlet = state.add_tasklet(...)
    out_edge = state.add_edge(...)

_independent_nodes is a set of AccessNodes, and the loop adds nodes and edges, so
its iteration order becomes their order in the state. AccessNode does not override
__hash__, so the order follows id() and varies per process — note this is
independent of PYTHONHASHSEED.

With two independent scalars, six runs gave two different orders (run 4 differs):

COPY_TASKLET_ORDER: ['loop_blocking_copy_tlet_t2', 'loop_blocking_copy_tlet_t1']
COPY_TASKLET_ORDER: ['loop_blocking_copy_tlet_t2', 'loop_blocking_copy_tlet_t1']
COPY_TASKLET_ORDER: ['loop_blocking_copy_tlet_t2', 'loop_blocking_copy_tlet_t1']
COPY_TASKLET_ORDER: ['loop_blocking_copy_tlet_t1', 'loop_blocking_copy_tlet_t2']
COPY_TASKLET_ORDER: ['loop_blocking_copy_tlet_t2', 'loop_blocking_copy_tlet_t1']
COPY_TASKLET_ORDER: ['loop_blocking_copy_tlet_t2', 'loop_blocking_copy_tlet_t1']

Reproduction

import dace
from dace.sdfg import nodes as dace_nodes
from gt4py.next.program_processors.runners.dace import transformations as gtx_transformations

sdfg = dace.SDFG("two_independent_scalars")
state = sdfg.add_state(is_start_block=True)
sdfg.add_array("A", shape=(40,), dtype=dace.float64, transient=False)
for out in ("B", "C"):
    sdfg.add_array(out, shape=(40, 10), dtype=dace.float64, transient=False)
for s in ("t1", "t2"):
    sdfg.add_scalar(s, dtype=dace.float64, transient=True)

A, B, C, t1, t2 = (state.add_access(n) for n in ("A", "B", "C", "t1", "t2"))
me, mx = state.add_map("main_comp", ndrange={"__i0": "0:40", "__i1": "0:10"})
state.add_edge(A, None, me, "IN_A", dace.Memlet("A[0:40]"))
me.add_scope_connectors("A")
for scalar, acc, out in ((t1, "t1", "B"), (t2, "t2", "C")):
    tl = state.add_tasklet(f"tlet_{acc}", inputs={"__in"}, outputs={"__out"},
                           code="__out = __in + 10.0")
    state.add_edge(me, "OUT_A", tl, "__in", dace.Memlet("A[__i0]"))
    state.add_edge(tl, "__out", scalar, None, dace.Memlet(f"{acc}[0]"))
    state.add_edge(scalar, None, mx, f"IN_{out}", dace.Memlet(f"{out}[__i0, __i1]"))
    state.add_edge(mx, f"OUT_{out}", {"B": B, "C": C}[out], None,
                   dace.Memlet(f"{out}[0:40, 0:10]"))
    mx.add_scope_connectors(out)
sdfg.validate()

sdfg.apply_transformations_repeated(
    gtx_transformations.LoopBlocking(blocking_size=2, blocking_parameters=["__i1"]),
    validate=True, validate_all=True)

print([n.label.rsplit("_", 1)[0] for n in state.nodes()
       if isinstance(n, dace_nodes.Tasklet)
       and n.label.startswith("loop_blocking_copy_tlet_")])

An OrderedSet would fix this, in line with #2779 / #2780 and spcl/dace#2445.

Impact

The SDFGs are semantically equivalent; the cost is that the serialization, and therefore
the compile-cache key, is unstable. Both are gated on loop blocking being enabled
(blocking_dims), so a default icon4py run does not hit them.

Found while auditing the dace transformations for non-deterministic ordering (see #2779,
#2780). Verified on main @ 6ec5244.

[edopao]

Accepted.