feat[next-dace]: let a SplitAccessNode fragment have several producers

[havogt]

Lets a SplitAccessNode fragment have several producers.

Before, a fragment was a single producer edge, so an AccessNode fed by several producers and read by several consumers could not be split, and the maps around it could not be fused. A fragment is now a set of producers joined by union find: a consumer that may intersect more than one producer merges them, so each fragment is a region that is written by a known set of producers and read by a known set of consumers, and the fragments are independent of each other.

Consequences worth knowing:

The check that everything a producer computes is also read applies only to a fragment with a single producer. With several producers no consumer is tied to an individual producer, so requiring the fragment to be read in full would reject the whole node, including the fragments that are fine. What stays unread is computed for nothing, but it is not a regression, without the split the same producer computes the same values and they are discarded just as well.

The producers of a fragment are merged with the adjacency only merger, since producers that overlap would describe the same memory twice.

split_node() now re-initializes the Memlet trees of the edges it is about to reroute. DaCe reruns Memlet propagation after every applied transformation, and it can leave a read on the edge of a NestedSDFG without a src_subset; the rerouting itself cannot repair that, because it runs while the old and the new edge are both present.

[philip-paul-mueller]

Kind of okay, just a few suggestions.

[Copilot]

Pull request overview

Extends DaCe SplitAccessNode to support fragments with multiple producers.

Changes:

• Builds producer fragments using union-find.
• Adds fragment coverage and overlap validation.
• Adds multi-producer and deduplication regression tests.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
split_access_nodes.py	Implements multi-producer fragment construction and validation.
test_split_access_node.py	Tests spanning consumers and merged-fragment behavior.

Suppressed comments (1)

src/gt4py/next/program_processors/runners/dace/transformations/split_access_nodes.py:519

• This does not actually test whether the consumers' union covers the fragment when their reads overlap. subset_merger() only joins exactly adjacent ranges (splitting_tools.py:668-673), so reads such as 0:10 and 5:15 remain separate and neither individually covers a 0:15 fragment; this incorrectly rejects a fully consumed multi-producer fragment. Please use a coverage calculation that handles overlapping subsets and add a regression case.

        merged = gtx_dace_split.subset_merger(
            [consumer_edge.data.src_subset for consumer_edge in consumer_edges]
        )
        return any(merged_subset.covers(subset) for merged_subset in merged)

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[havogt]

Added as test_partially_read_multi_producer_fragment_is_rejected(): a spanning consumer merges the second and third producer but stops short of what the third writes, with a first, fully read fragment present so the check is actually reached rather than short-circuited by the single-fragment early return.

Your suppressed comment on the same function was the more important one and it was right — see the separate note on the PR.

[havogt]

Review round pushed as two commits on top.

refactor[next-dace]: address review on SplitAccessNode fragments — @philip-paul-mueller's points: enumerate instead of range(len(...)), _is_covered_by_consumers → _is_fully_read_by_consumers, the two len(fragment.producers) == 1 conditions replaced by explicit multi-producer / single-producer branches (behaviour unchanged, reasoning inline above), and the test now asserts the fragment structure — producer counts per fragment and that each fragment reaches the consumer that reads it, traced through the map exit rather than matched on labels.

fix[next-dace]: decide fragment coverage with overlapping consumer reads — this one came out of a suppressed Copilot comment and was a real bug, so calling it out explicitly rather than leaving it buried.

_is_fully_read_by_consumers() used subset_merger(), which only joins subsets that are exactly adjacent (end1 + 1 == start2). Two consumers reading 0:10 and 5:15 therefore stay separate and neither covers 0:15 on its own:

adjacent     merged=['0:15']           covers(0:15)=True
overlapping  merged=['0:10', '5:15']   covers(0:15)=False   <- legal split rejected

Consumers reading overlapping regions is the normal case for a stencil reading a halo, so this hit precisely the situation multi producer fragments exist for. It fails in the safe direction — it declines a legal split rather than miscompiling — but it would have quietly defeated much of the point of this PR.

Coverage is now decided with a merger that also joins overlapping and contained subsets. I kept it local to split_access_nodes.py so the meaning of subset_merger() is unchanged: describing a split should only join adjacent subsets, deciding coverage should not. If you would rather have it live in splitting_tools.py next to subset_merger(), say so and I will move it.

test_overlapping_consumers_cover_multi_producer_fragment() pins it, and I verified it fails against the previous implementation before trusting it.

Testing: test_split_access_node.py 22 passed / 2 xfailed; the wider dace_tests suite 405 passed with one pre-existing failure and two pre-existing errors, all in test_dace_fastcall*[exec_alloc_descriptor1], confirmed pre-existing by reproducing them on a branch that contains no split_access_nodes.py change at all. mypy and ruff format clean.

[edopao]

I do not understand why the subset 0:5 is not split out.

[edopao]

I mean: a -> d should become an independent graph as in next test case.

[havogt]

Good question, and sorry for the slow reply — this one slipped past me.

0:5 is not split out because the check is all or nothing: _check_split_constraints() returns a single verdict for the whole node, so one fragment that fails rejects the split entirely, including the fragments that were fine. That is pre-existing behaviour rather than something this PR introduces — on main the same function return Falses out of a loop over all reassignments.

The reason it is structured that way is that split_node() is handed a partition of the entire AccessNode and every edge has to land in exactly one fragment. Splitting only 0:5 would need a residual fragment holding 5:15 that is deliberately left un-tight, and nothing currently constructs that.

So in this test: fragment 0:5 (producer a, consumer d) is fine, fragment 5:15 (producers b and c, consumer e reading only 5:12) is not, and the whole candidate is dropped. The test is checking the rejection, not arguing that rejecting everything is ideal.

You are right that it is more conservative than it needs to be — a partial split would be a real improvement, since 0:5 would then be served directly. I have not changed it here because it is a behaviour change to a pre-existing decision and would want its own reasoning about what happens to the residual, but I am happy to add a Todo on _check_split_constraints() recording it, or to take it on in a follow up if you would prefer.

[edopao]

Probably this case is unrealistic, because domain inference should shrink the domain of the intermediate field to [0:12] in the first place.

[havogt]

Note: I rebased this branch onto 721946037 (the uv lock file update) — apologies for the force-push on a branch you are reviewing, I should have merged instead. The reviewed code is unchanged by it; the only content change is one reformat for ruff 0.16.1 in a test file. No further force-pushes here.

[havogt]

Unrelated to this change, but noticed in the same file and it looks like a one word fix, so flagging rather than touching it here.

In SplitAccessNode.can_be_applied() the assume_single_use_data branch is overwritten by the following one:

```python
if self.assume_single_use_data:
single_use_data = {sdfg: {access_node.data}}
if self._single_use_data is None: # <- assigns again, unconditionally
find_single_use_data = dace_analysis.FindSingleUseData()
single_use_data = find_single_use_data.apply_pass(sdfg, None)
else:
single_use_data = self._single_use_data
```

So assume_single_use_data=True has no effect. That matters because gt_split_access_nodes() constructs the transformation with exactly that flag in order to avoid the scan — the comment there says "we set `assume_single_use_data` to `True` because we do this test outside" — and instead FindSingleUseData runs a full pass per candidate.

Purely a compile time cost, no behavioural consequence, and elif on the second if appears to be all that is needed. Happy to include it here if you would rather, or leave it for a separate change.

[philip-paul-mueller]

Some more points.

[philip-paul-mueller]

consider using oedge.data.get_src_subset(oedge, state) here.

[havogt]

Switched, here and for the consumer subset in the coverage test. Worth noting it is not only style: the raw property returns None when Memlet._is_data_src was never set, which is the failure mode in #2775.

[havogt]

I have to walk this one back: switching to get_src_subset() broke the icon4py run, and I pushed it without noticing. Reverted in 6f4e04cdf, with a NOTE recording why.

The accessor initializes the Memlet; the raw property does not. That difference is load bearing at this call site, because two lines below there is

if consumer_subset is None:
    return None  # TODO(phimuell): Lift this.

So a Memlet with no side associated to it had no src_subset, and that None was acting as an implicit guard — an uninitialized Memlet meant "decline this candidate". Going through the accessor initializes it, the guard stops firing, and the transformation starts accepting nodes whose rerouting then fails inside reconfigure_dataflow_after_rerouting() — the subset_to_adjust is not None assertion, on a child of the Memlet tree inside a Map scope. With assertions disabled, as in the ICON runs, that surfaces as AttributeError: 'NoneType' object has no attribute 'offset'.

Bisected: main, and this branch up to and including bce97261a, pass test_compute_advection_in_horizontal_momentum_equation[compile_time_domain-apply_extra_diffusion_on_vn[True]] on dace_cpu. efa882017 fails it. That commit with only this line reverted passes. The other four changes from that round are untouched.

Two things I think are worth your view.

The guard should not be implicit. As it stands, split_node() relies on a Memlet being uninitialized to signal "do not touch this node", which is not something anyone would infer from the code, and it depends on whether the caller went through gt_split_access_nodes() — which does initialize the whole SDFG up front — or applied the transformation directly, as map_fusion_extended does. Your suggestion is the right API; what is missing is that the rerouting has to cope with those nodes before it can be adopted.

This is issue #2775 with a production traceback. I filed that assertion as reachable only from a hand constructed edge and said so explicitly; it is reachable from icon4py, and I have corrected the issue. The helper reads the raw {src,dst}_subset properties with a comment saying it cannot use the accessors because the SDFG is transiently invalid — I tried try_initialize() there and on the newly created bypass edge, and neither helps for an edge inside a Map scope, where neither endpoint is an AccessNode named after the Memlet's data. That is where I stopped, since the fix needs a decision about that helper rather than a guess from me.

Performance, measured after the revert (GH200, mch_icon-ch1_medium, 1800 calls, fresh translation caches): target program 0.7818 → 0.5527 with PR1 → 0.5202 with this PR, i.e. 0.89× → 1.26× → 1.33× against the OpenACC reference. No other program moved by more than 0.006 s against the pre-rework measurement, model total within the 0.020 s run to run spread.

[philip-paul-mueller]

This is absolutely wrong!
First src_subset is not the raw property it is already a "direction aware property" and a lot of issue we had in the past was that that thing was not updated correctly.
The raw property is either subset or other_subset, however, which one must be used highly depends on the state of the Memlet and how it is integrated into the SDFG.

To put it differently get_src_subset() is correct, the Memlet is wrong and you must find out where the Memlet has gone bad.

[havogt]

You are right on both counts, and thanks for pushing back. src_subset is direction aware, get_src_subset() is the correct API, and the Memlet is broken long before SplitAccessNode looks at it. Here is where it goes bad.

The corruption

unsqueeze_memlet() (dace/transformation/helpers.py) builds the outer Memlet of a NestedSDFG edge like this:

actual_result = Memlet.from_memlet(internal_memlet)   # copies _is_data_src
actual_result.data = external_memlet.data
...
else:
    actual_result.subset = result.subset

from_memlet() copies _is_data_src from the Memlet inside the nested SDFG, where it refers to a different data container, and then data is replaced by the external one. When the internal Memlet has no other_subset — a write expressed in terms of its destination, so _is_data_src is False, which is correct inside — the resulting outer Memlet claims that its data is on the destination side while carrying no other_subset. Its src_subset is None from then on.

Standalone, no gt4py, dace 2.0.0a5:

import dace
from dace.sdfg import propagation

inner = dace.SDFG("inner")
inner.add_scalar("__arg", dace.float64)
inner.add_scalar("__output", dace.float64)
istate = inner.add_state()
istate.add_nedge(
    istate.add_access("__arg"),
    istate.add_access("__output"),
    dace.Memlet(data="__output", subset="0"),   # no other_subset -> _is_data_src is False
)

sdfg = dace.SDFG("outer")
sdfg.add_array("A", [10], dace.float64)
sdfg.add_array("B", [10], dace.float64)
st = sdfg.add_state()
a_node, b_node = st.add_access("A"), st.add_access("B")

map_entry, map_exit = st.add_map("map", {"__i": "0:10"})
nsdfg = st.add_nested_sdfg(sdfg=inner, inputs={"__arg"}, outputs={"__output"})

map_entry.add_in_connector("IN_A")
map_entry.add_out_connector("OUT_A")
st.add_edge(a_node, None, map_entry, "IN_A", dace.Memlet(data="A", subset="0:10"))
st.add_edge(map_entry, "OUT_A", nsdfg, "__arg", dace.Memlet(data="A", subset="__i"))

map_exit.add_in_connector("IN_B")
map_exit.add_out_connector("OUT_B")
st.add_edge(nsdfg, "__output", map_exit, "IN_B", dace.Memlet(data="B", subset="__i"))
st.add_edge(map_exit, "OUT_B", b_node, None, dace.Memlet(data="B", subset="0:10"))
sdfg.validate()

edge = next(iter(st.in_edges(nsdfg)))
print("before:", edge.data, edge.data._is_data_src, edge.data.src_subset)
propagation.propagate_memlets_map_scope(sdfg, st, map_entry)
edge = next(iter(st.in_edges(nsdfg)))
print("after: ", edge.data, edge.data._is_data_src, edge.data.src_subset)
print("get_src_subset() ->", edge.data.get_src_subset(edge, st))

before: A[__i] True __i
after:  A[__i] False None
get_src_subset() -> __i

data and subset are untouched, only the association flips. Note the last line: the accessor repairs the Memlet.

Why it reached us, and why my earlier explanation was backwards

get_src_subset() never caused anything. It repaired the outer edge and let the split proceed, while the Memlets on the inner edges of the map scope stayed corrupt. reconfigure_dataflow_after_rerouting() walks exactly those inner edges and cannot use the accessor — it runs while the old and the new edge are both present, as the NOTE there says — so it read None and died. The raw property was declining the candidate by accident.

My first attempt at a fix was to repair after our own propagate_memlets_map_scope() call in split_maps(). That fixed one program and left two others failing, because DaCe propagates behind our back: PatternTransformation.apply_pattern() runs propagate_memlets_sdfg() after every applied transformation unless it annotates_memlets(). There is no point in guarding our own propagation calls.

What this PR does

The repair now sits at the entry of split_node(), which is the last moment where the graph still says where the data is: for every edge incident to the node being split, the Memlets of its whole Memlet tree are re-initialized. Those are exactly the edges the rerouting will later have to offset. _find_edge_reassignment() uses get_src_subset() again.

New test test_vertical_map_fusion_with_nested_sdfg_consumer: a Map that consumes the intermediate through a NestedSDFG. It hits assert subset_to_adjust is not None in ~3 s without the repair and passes with it. That is also the answer to why the suite stayed green while the model run aborted — no existing test builds an inner Memlet in that shape.

The real fix belongs in unsqueeze_memlet(), which should not carry _is_data_src across the SDFG boundary. That is DaCe, so I have not touched it; say the word and I will open an issue there with the snippet above.

[havogt]

Addressed in efa882017. Everything mechanical is applied; two points need your input.

• :357 — the union find now says what fragment_of_producer[i] holds.
• :371 — switched to oedge.data.get_src_subset(oedge, state), and did the same for the consumer subset in the coverage check. Worth noting the accessor matters for more than style: the raw property returns None when Memlet._is_data_src was never set, which is the failure mode in next[dace]: track_view raises on a lowering produced View, and reconfigure_dataflow_after_rerouting asserts on an uninitialised Memlet #2775.
• :378, :402, :423 — set comprehension plus sorted() where only determinism matters, find_fragment() precomputed, and your any(...) suggestion applied.
• :404 — added, and stated the asymmetry: the producer subsets must use the adjacency only merger because overlapping producers would describe the same memory twice, unlike the consumer subsets in _is_fully_read_by_consumers(), which only have to cover the fragment.
• :415 — "edge" is indeed the consumer edge; the comment says so now.

:476, the tightness comment. Two things.

The CopyChainRemover justification is not mine — it is on main today, at split_access_nodes.py:461-463, directly above your TODO(phimuell): Lift this limitation. I lifted it when hoisting the comment and generalised it to the fragment case, which is how it ended up somewhere it reads oddly. I have removed that claim from my comment and left the original where it sits, so the question of what tightness is actually protecting is yours to answer — I did not want to invent a new justification for a requirement I did not introduce.

On the producer [0:10] / consumer [0:5] case: it is rejected on both paths, and the comment now says so. With several producers _is_fully_read_by_consumers() requires the consumers to cover the fragment, and [0:5] does not cover [0:10]. With one producer the covers() test below rejects it directly. The comment previously explained only why the granularity differs between the two branches and never said that the requirement itself is unchanged, which I think is what made it read as if something was being weakened.

:525, where the overlap aware merger should live. Agreed it should move; I have no strong view on where. Three options as I see them, in the order I would pick them:

1. splitting_tools.py beside subset_merger(), as a second entry point with a name that says it answers coverage rather than describing a split.
2. A parameter on subset_merger() itself — smaller surface, but it makes one function mean two different things, and the adjacency only behaviour is load bearing for the producer side.
3. Leave it private here until a second caller appears.

Happy to move it wherever you prefer; #2776 has the reasoning if it is easier to decide there.

Tests: 22 passed, 2 xfailed. ruff clean.

[edopao]

The code agent seems to have missed my comment:
#2764 (comment)

[philip-paul-mueller]

LGTM

New code changes.

[philip-paul-mueller]

This is absolutely wrong!
First src_subset is not the raw property it is already a "direction aware property" and a lot of issue we had in the past was that that thing was not updated correctly.
The raw property is either subset or other_subset, however, which one must be used highly depends on the state of the Memlet and how it is integrated into the SDFG.

To put it differently get_src_subset() is correct, the Memlet is wrong and you must find out where the Memlet has gone bad.

[philip-paul-mueller]

We are close.

[philip-paul-mueller]

This function should be called _find_fragments() or so.

[havogt]

Renamed to _find_fragments().

[philip-paul-mueller]

This is a bit different.
"Dead Data" was fully eliminated by the Dead Dataflow pass, but now this data is generated and written to memory, but it sits around and is never read.

To ensure that no data is generated that is never read, even with multiple producers/consumers.
You simply merge the consumers together (this is either done with your new overlapping subset merger or with dace.subset.bounding_box_union()) and then check if it covers _Fragement.subset.
However, I am not fully sure if this restriction is useful, but it would preserve the old invariant the transformation used to have.

[havogt]

You are right about the mechanism and my comment was wrong. I checked both shapes against DeadDataflowElimination: it removes a fragment with no consumer entirely (AccessNode, Tasklet, MapEntry and MapExit all go), and it removes nothing from a partially read one, because the array is read, so the producer keeps computing values that are then discarded. The comment now says that instead of equating the two.

On the check itself, three findings argue against implementing it as proposed:

1. bounding_box_union() can not express it. Consumers [0:2] and [4:6] give [0:6], which covers a fragment [0:6] although 2:4 is never read. In 2D it is worse, since the per dimension special cases fire independently: [0:2,0:2] and [2:4,2:4] give [0:4,0:4], with 8 of 16 elements read. The test would accept exactly the fragments it is meant to reject. (It lives in dace.subsets and is binary, so it needs a fold; SubsetUnion.covers() is any(...) and fails the other way.)
2. subset_merger() can not either. It merges only exactly adjacent subsets, so ['0:10','5:15'] stays split and coverage of 0:15 reports False — it would reject legal splits. The overlap aware merger that could decide this is the one removed in bac31c364, which also answered your question at :525 about where that functionality should live.
3. It would not restore the invariant main had. On main the condition is evaluated per producer edge, and three of the six cases impose nothing at all (zero consumers, non transient source, multi consumer transient source). Where it does apply — MapExit producers, and single consumer transient AccessNode producers — it requires consumer_subset == producer_subset exactly, per consumer, which is strictly stronger than "the merged consumers cover the fragment". I ran a three producer, two consumer node on main: it applies and only warns about the unread producer, so produced-but-unread data was tolerated there as well.

So I left the behaviour alone and fixed only the comment. If you do want the restriction, the honest form is an exact coverage test, which means bringing the overlap aware merger back — happy to do that if you say so.

[havogt]

Some changes on the way during review killed the performance gain that this should have introduced, claude is back to debugging...

[havogt]

ok, probably unlucky dace indeterminism.

[philip-paul-mueller]

LGTM.

[philip-paul-mueller]

I still do not see why it is a regression.
Consider the following case you have the following three producers [0:10], [10:20] and [20:30] and one consumer [0:15].
The old transformation would have refused to work, but the new one does not, instead it will create two fragments [0:20] and [20:30], the last one with zero consumers and the first fragment is only partially read.
I agree, the situation before and after the split are equivalent, in the sense that there is data that is written but not read.

However, this check alters the behaviour of the transformation in a very strange way.
Before if there was a single consumer of a fragment then the fragment had to be consumed in full.
Now if there is a single consumer of a fragment the fragment must only be consumed in full if it has a single producer.
This is slightly different than before and caused by the new definition of a fragment.
And the question is, does this matter?

This rule was established, as previous comments noted, to ensure that CopyChainReomver works that assumes some tightness.
And we should still work to make that happy.

Let's make the example from above more concrete:

• P1: tmp[0:10] = G[a:b]
• P2: tmp[10:20] = G[c:d]
• P3: tmp[20:30] = ... (case not important).

Again P1 and P2 will end up in the same fragment, furthermore they are fed by the same data G (which is not necessarily global).
If we assume b == c then they can be represented as one edge, that was, at some point, split for whatever reasons.
If this edge merge happens after CopyChainRemover then we might have created something that it does not like.

As I said, it is now slightly different and the question is: does it matter.
After some thinking I came to the conclusion that it does not matter, at least I consider the cases where it might matter to be unlikely.

However, this discussion has showed us, that we should probably rework CopyChainRemover.
I think it should be split into two transformations and I think for the Pull case we can sharpen some constraints a bit.

[philip-paul-mueller]

Suggested change

	# TODO(phimuell): Replace this with a cheaper repair. The cost is not the Memlet
	# tree, which is scope local, but `try_initialize()`, which calls `memlet_path()`
	# once per edge; that is also where the "undefined for more than one path"
	# caveat of `memlet_path()` applies.
	# TODO(phimuell): Replace this with a cheaper repair, that fixes the Memlet
	# trees directly instead of every single Memlet.The cost is not the Memlet.

[havogt]

Applied, with one deviation: the suggestion block covers only the first two lines of the TODO, so taking it verbatim would have left ...instead of every single Memlet.The cost is not the Memlet. followed by the old third line once per edge; ..., i.e. a sentence cut in half. I kept your wording and made the rest follow from it:

    # TODO(phimuell): Replace this with a cheaper repair, that fixes the Memlet trees
    #   directly instead of every single Memlet. The cost is not the Memlet tree, which
    #   is scope local, but `try_initialize()`, which calls `memlet_path()` once per
    #   edge; that is also where the "undefined for more than one path" caveat of
    #   `memlet_path()` applies.

Say the word if you meant to drop the second half entirely.

[havogt]

Correction to my previous reply: I misread the suggestion as covering only the first two lines. It spans 282-285, i.e. the whole TODO, so "The cost is not the Memlet." is a complete sentence and nothing was cut. Applied verbatim in ae122182b (only adding the missing space after the period).

One thing to confirm you meant: it drops the note that the cost and the "undefined for more than one path" caveat both sit in try_initialize() -> memlet_path() rather than in memlet_tree(), which is the part I had measured. Fine by me either way, it is your TODO — say so and I put that one line back.

[philip-paul-mueller]

Yes this is correct.