feat[next-dace]: eliminate write back buffers that are also consumed

[havogt]

A new pass, GT4PyWriteBackBufferElimination.

It matches a transient T that is defined in one state, copied in full into a non transient G in another state, and additionally read by other consumers. Every access to T is rewritten into an access to G, shifted by the offset of the copy, and the copy is removed, so producer and consumers work on G directly and the full array copy disappears.

Neither DistributedBufferRelocator nor GT4PyMapBufferElimination covers this pattern: the first needs out_degree(T) == 1, the second needs the copy to sit in the state where T is written and T to be unused afterwards.

The rewrite moves the write of G earlier, from the write back to the point where T is defined. That the range of G holding T is not modified in between is a hard requirement; the pass additionally rejects a modification after the write back, which is a simplification that avoids testing whether T is still read there. _has_conflicting_global_access() establishes both. Since rule 3 of ADR-18 takes precedence for global memory, and rules 6 to 10 may then be violated, nothing here assumes that a global has a single writer.

assume_pointwise asserts that a read of G in a state that defines T is elementwise, which ADR-18 rule 3 guarantees for a valid GT4Py program; the pass only checks that such a read is ordered before the definition of T.

Tests in test_write_back_buffer_elimination.py, one per rejection guard.

[philip-paul-mueller]

There are some aspects that needs some work.
If there are more question just drop me a line.

[edopao]

LGTM

[havogt]

Review round pushed as three commits, deliberately separated so the move is reviewable as a move.

• fdf38d4fb — pure move of GT4PyWriteBackBufferElimination into write_back_buffer_elimination.py. auto_optimize.py needed no change (it already went through gtx_transformations.…), and the test file was already named to match.
• 8cdf79004 — the documentation and naming points: the :1356 suggestion verbatim; the class docstring now says G must hold its original value between the definition of T and the write back; _find_candidates documents the (tmp_name, wb_edge, wb_state) triple and the conditions it guarantees; glob_read → glob_read_node; _has_conflicting_global_access rewritten to cover the write case and both read cases with the reason each disqualifies, plus the Todo: for the state-topology point; _only_feeds_tmp_producer rewritten to stand alone; and the swallowed ValueError removed.
• cdcad0213 — the rerouting rewrite. Details in the thread, but the headline: your :1538 comment found a real bug, and it produced wrong numbers rather than a validation error. It is specifically the data.data association assumption, not the "used inside Maps" part — inner memlet-tree edges were reached and translated correctly, which is why the pre-existing Map-consumer test passed. The failing shape is a memlet expressed relative to the other container, with the T-side range in other_subset.

Tests: test_write_back_buffer_elimination.py 4 → 7 passed; transformation_tests/ 282 → 285 passed, 3 xfailed. The 1 failure and 2 errors in the wider dace_tests/ run are the pre-existing test_dace_fastcall*[exec_alloc_descriptor1] GPU-allocator set, confirmed by reproducing them on a branch containing none of these changes. ruff, mypy and tach clean on the source files.

Both earlier correctness guards survive the rewrite: gt_propagate_strides_of is still the last statement of _eliminate, and _only_feeds_tmp_producer still narrows the assume_pointwise waiver.

I have also put this branch through an independent adversarial review pass; will follow up here with anything it turns up.

[havogt]

Still iterating on this one...

[havogt]

Second round pushed (49051c272..d09db8bdd), and it changes what I can honestly claim about this pass, so writing it up rather than burying it in commits.

An independent review found seven defects; all are fixed

Each was reproduced first and is verified by that reproduction:

• Views of T were never stride updated — gt_propagate_strides_of only descends into NestedSDFGs, and _gt_modify_strides_of_views_non_recursive skips views of non-transient data, which is exactly the assumption this pass invalidates. Globals carry symbolic strides and transients concrete ones, so it is guaranteed wrong when it happens. Now rejected in _find_candidates rather than repaired.
• _has_conflicting_global_access was unsound on a cyclic control flow path — a LoopRegion back edge makes a state that must see the old G look like it runs after the write back. Now rejected.
• A strided write back destination was accepted — Range("0:20:2").size() folds the step away and the offset came from min_element() alone, so b[i] was written instead of b[2i]. Note the asymmetry: the source check already rejected strides.
• A WCR on the write back edge was silently dropped, turning b += tmp into b = ....
• An empty dependency memlet became non empty via reroute_edge's full range fallback, producing an invalid SDFG from a valid one.
• modifies() under declared Descriptors/NestedSDFGs; added @explicit_cf_compatible.
• An ADR-18 rule 6 violating input was corrupted and then raised; now rejected up front.

Not fixed, but now stated accurately in the docstring: assume_pointwise does not hold for a non-zero offset. It is out of contract per ADR-18 rule 3, GT4PyMapBufferElimination has the identical hole, and it was verified not to be a regression — run_dace_cpu already miscompiles such a program with both passes disabled.

The pass only fires with static domain bounds

This was not previously understood, and it is the most important thing here for a reviewer.

Instrumenting _find_candidates on compute_perturbed_quantities_and_interpolation:

domain	candidates
dynamic	0
compile_time_domain	3 (CPU and GPU alike)

With a dynamic domain the transient is allocated with a clamped shape while the write back carries the unclamped extent, and Range.__eq__ compares sympy structurally, so 0:Max(0, hi - lo) never equals 0:hi - lo. With static bounds these fold to literals and the check succeeds. A sweep of 133 test invocations and 18 hand written program shapes with dynamic domains produced zero candidates, which is why this went unnoticed.

Arguably the source check is now too tight in the same way the destination check was too loose. I have deliberately not touched it: relaxing it would make the pass start firing in configurations where it currently does not, which is a behaviour change and wants a decision rather than being smuggled in with a defect fix.

New: a test on lowering produced input

test_write_back_buffer_elimination_lowering.py. Until now every test hand built an SDFG, so nothing verified the pass against what the lowering actually emits — plausibly why the defects above survived. The minimal firing program is eight lines, and a reduction ladder shows all three ingredients are necessary: a dual role output, a branch whose body reads it, and a static domain.

A second test pins the invariant that makes it safe not to update view strides, and guards against passing vacuously by first asserting that views do exist immediately after lowering.

Supporting evidence for that invariant, across 149 dycore stencil tests / 85 programs / 357 pass invocations: 858 View nodes exist at gt_auto_optimize entry and zero at pass time, every time. _construct_local_view is the only add_view call in the dace runner, and RemovePointwiseViews runs after this pass.

One number in the description I can no longer support

It says "4 cudaMemcpy2DAsync calls removed". The instrumentation finds 3 candidates, on both CPU and GPU. The production benchmark configuration differs from the stencil test, so both may be right, but I have not reconciled them and would rather flag it than leave an unverified figure standing.

Tests: test_write_back_buffer_elimination.py 7 → 15, plus 3 new integration tests; transformation_tests/ 285 → 293 passed, 3 xfailed.

[Copilot]

Pull request overview

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Adds a DaCe pass to eliminate “write-back buffers” (global transient + later copy back) when the transient is also consumed, avoiding extra device-to-device copies while preserving correctness (stride propagation + safety guards).

Changes:

• Introduces GT4PyWriteBackBufferElimination pass and integrates it into DaCe auto-optimization.
• Adds extensive unit + integration tests covering correctness guards (partial/strided writeback, loops, views, nested SDFGs/strides, unrelated readers, empty memlets).
• Exports the new pass from the transformations package.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/next_tests/unit_tests/program_processor_tests/runners_tests/dace_tests/transformation_tests/test_write_back_buffer_elimination.py	New unit tests for numerous SDFG patterns and safety conditions.
tests/next_tests/integration_tests/feature_tests/dace_tests/test_write_back_buffer_elimination_lowering.py	New integration tests ensuring the pattern is reachable from real lowering and checking View-related assumptions.
src/gt4py/next/program_processors/runners/dace/transformations/write_back_buffer_elimination.py	New transformation pass implementing write-back buffer elimination + safeguards + stride propagation.
src/gt4py/next/program_processors/runners/dace/transformations/auto_optimize.py	Runs the new pass during auto-optimization.
src/gt4py/next/program_processors/runners/dace/transformations/init.py	Exposes the new pass in __all__.

---

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

[philip-paul-mueller]

Legitimate point and we should use it at some point, but at this level not really present.

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

[edopao]

cartesian and next use the same dace version in latest gt4py, so the try/except below is not needed.

[havogt]

Thanks, applied.

[philip-paul-mueller]

The only reason why it does not apply for dynamic bounds is that SymPy does not know that lo <= hi.
However, it is possible to use assumptions, however, they are not properly implemented in DaCe and even if GT4Py sets them at some point, I am not sure that it will survive.
However, this should be seen more as a bug/limitation than intended behaviour.

[philip-paul-mueller]

Better, but I see more work and the need for clarifications.

[philip-paul-mueller]

Legitimate point and we should use it at some point, but at this level not really present.

[havogt]

Thanks — this was the right kind of review to get. I have treated it as four root causes rather than 24 fixes, because patching each comment would have produced another round.

Two commits: b732249bb re-derives the preconditions, 006e8ab8f rebuilds the tests.

1. The preconditions were derived empirically, not from ADR-18

That is the source of most of your comments, and you caught it being simultaneously too strict and too lax. I went back to the ADR and re-derived them. Every surviving condition now carries a comment naming the rule it enforces; anything I could not justify from the ADR is gone.

Deleted, with what made them unnecessary:

deleted	why
all T AccessNodes must be top level	rule 10 — an AccessNode in a Map scope needs AllocationLifetime.Scope, which an interstate T cannot have
T read and written through different nodes in one state	rules 6 + 8
is_view(tmp_desc), is_view(glob_desc), _find_viewed_data	subsumed by the whole-SDFG View bail-out
len(glob_desc.shape) != len(tmp_desc.shape)	subsumed by the per-dimension size check
scope_dict()[glob_node] is not None	rule 10
wb_state.out_degree(glob_node) != 0	simply wrong, as you noted
"any write to G anywhere"	narrowed to a write that may touch the write-back region
"any read in a state not reachable from wb_state"	narrowed to three exemptions: disjoint, strictly before every definition, or strictly after the write back

2. Cycles — you are right that the guard should go, but not for the stated reason

You asked "when do you have cycles, our SDFGs do not have them and ADR-18 forbids them?". Rule 5 forbids cycles in the state graph and, in the same sentence, mandates LoopRegion as the replacement — and a LoopRegion does produce a reachability cycle that this pass sees:

• sdfg.states() on the top-level SDFG returns states inside LoopRegions.
• ControlFlowBlockReachability._region_closure does closure.update(region.all_control_flow_blocks()) for a LoopRegion, commented "Any point inside the loop may reach any other point inside the loop again." It deliberately surfaces the cycle rather than hiding it.
• The lowering emits LoopRegions for scans.

Run on such an SDFG, the old code fired and produced wrong numbers — d one iteration late, [2,4,6,8] where the reference is [0,2,4,6].

But the guard was a symptom fix. The real defect is that reachable is a may-reach relation and the code used it as happens-after. Every use is now

def _runs_strictly_before(first, second, reachable):
    return second in reachable[first] and first not in reachable[second]

which is loop correct by construction and strictly more permissive than what it replaced — a conflicting access outside the loop is now accepted. The special case is gone; the loop test stays as the reproduction and now fails for the right reason.

3. Views: give up, as you asked

_find_viewed_data is gone. The pass now bails if the SDFG contains any View.

One deviation: I did not use get_all_view_nodes(). Once the answer is "give up on any View", there is nothing to resolve — that helper walks a View chain to its base array, and the base array cannot change the outcome. _has_view() is a nine line existence check. Happy to swap it if you would rather have the helper used for consistency.

4. The tests pinned defects, not the contract

Rebuilt around a matrix of how a consumer can reference T: direct_copy | map_scope | nested_sdfg x memlet named after T or after the other container x offset (0,0) or (11,22). Every defect found in this PR so far is one cell of it. Each cell runs the SDFG before and after and compares all non-transient arrays.

• :110 — the stride test is replaced by the NestedSDFG x non-zero-offset cells. b is 100x100 against a 10x10 tmp, so a missing gt_propagate_strides_of shows up as wrong numbers rather than as an assertion about a mock.
• :217 — the empty Memlet test now asserts the empty edge starts at the same b node the producer writes, which is the property whose loss would lose the ordering.
• :273 — deleted, it tested that an invalid SDFG stays invalid.
• :408 — the View tests now expect a bail-out.

Your other points

:361, assume_pointwise. You were right and this changed shape. _only_feeds_tmp_producer rejected exactly your G → Map1 → TT → Map2 → T case. It is now _is_read_by_tmp_producer: for every out-edge of the G read, all T-writing AccessNodes must be downstream of that edge's destination. The one-hop formulation matters — a naive "T is reachable from the G node" wrongly accepts a shared AccessNode, which is now a test.

Going further than you asked: a shifted write back combined with a producer read of G is now rejected rather than documented as a limitation. Rule 3 is about the same memory, and G[i] → G[i+off] clobbers across iterations.

Reusing a G node (:1538). You are right, and my rule was wrong in one direction. It is now decided per T node: written → its own AccessNode (rules 3 and 8 permit the extra output node, and reusing the read would make the state cyclic); read only → join an existing G node that has in_degree == 0 and is top level. Requiring in-degree zero makes the reused node a source, so joining it cannot create a cycle. Both branches have tests.

:179 and :236, the missing cases. Neither needs a condition, and the docstring now says why. From the write back onwards T and G[wb_region] hold the same value, and rule 6 forbids a second AccessNode writing G downstream of glob_node. Other consumers of T are rerouted to G identically by _eliminate(), before or after the write back. Both are covered numerically.

:86. Correct and stronger than what I had, but only for writes — reads after the write back see identical content either way. For writes the extended interval is free, since rule 6 forbids a writer downstream of glob_node outright. The docstring states both halves.

:285, the optional cyclic test on the def states. Declined. Two def states reaching each other means two definitions of T on one path, which rule 6 forbids — that is a defence against invalid input, which is the habit this round is trying to break. It also would not catch anything the strictly-before requirement already lets through.

:273 (the step check). Kept, as you allowed, now with a comment on what it establishes. I could not work out which "much stronger check below" subsumes it — the step check is orthogonal, and DaCe's memlet validation constrains total volume rather than the per-dimension shape the offset shift needs. Happy to drop it if you point at the one you meant.

:227 is moot — the loop it suggested a fix for is deleted.

One risk I am accepting deliberately, so it is a decision and not an oversight

Deleting the top-level check on T's AccessNodes is the deletion with teeth. If a T node did sit in a Map scope — a rule 10 violation — _eliminate would now attach a top-level G node to nodes inside the scope and produce an invalid SDFG rather than declining. _replacement_node refuses to reuse a non-top-level G node, which limits it, but the creation path is unguarded. That is the "assume the contract" principle applied consistently; say the word if you would rather keep the check.

Testing

test_write_back_buffer_elimination.py 33 plus the 3 lowering tests: 36 passed, and the pass still fires on lowering produced input. Full dace_tests/: 435 passed, with the pre-existing test_dace_fastcall*[exec_alloc_descriptor1] GPU allocator failure and its two companion errors. ruff clean, mypy reports nothing new.

[philip-paul-mueller]

You see I was maybe not so clear, this transformation only makes sense to apply to T on top level.
So you should ignore everything else.

[philip-paul-mueller]

I see more work.

[philip-paul-mueller]

Mostly minor changes.

[philip-paul-mueller]

We are almost there.

[philip-paul-mueller]

This is something that belongs to the Notes section or at the end of the doc string.

[havogt]

Moved into Notes.

[philip-paul-mueller]

You can remove that with the "shifted by" it is kind of obvious.

[havogt]

Removed "The producer of T and all of its consumers thus operate on G directly", which is what I read "that" as pointing at. Say so if you meant the full coverage sentence instead.

[philip-paul-mueller]

The first part "that the region of G is not modified until the wb" is a hard requirement, the part that T is not used beyond that point is the simplification.

[havogt]

You are right, they are not the same kind of constraint, and calling both a simplification was wrong. Not modifying the range of G between the definition of T and the write back is hard: otherwise the preponed write either ends up behind that write or in front of it, and the result is wrong. Extending it past the write back is the simplification, since a later modification only matters for a consumer of T that reads after it, and rejecting it outright avoids testing whether T is still read there. The docstring says that now.

[philip-paul-mueller]

Suggested change

	if isinstance(edge.dst, dace_nodes.AccessNode) and not edge.dst.desc(sdfg).transient
	if isinstance(edge.dst, dace_nodes.AccessNode) and not edge.dst.desc(sdfg).transient and state.scope_dict()[edge.src] is not None

Ensures that the thing only applies to nodes that are on the top level.
As pointed out above, only they are relevant.
But I can life without it.

[havogt]

I think the polarity is the other way round. state.scope_dict()[n] is None when n is at the top level and the enclosing MapEntry when it is inside a scope, so is not None selects the nodes inside a scope. I checked on a small SDFG to be sure:

top-level AccessNode  -> scope_dict: None
in-scope AccessNode   -> scope_dict: MapEntry

Since T is a top level transient by construction, the suggestion as written would most likely make the pass match nothing. What you describe needs is None, and state.scope_dict() should come out of the comprehension, where it is recomputed per edge. Happy to add it in that form, but as you say it is not needed, so I left it out for now rather than guess.

[philip-paul-mueller]

Yes, you are right it is indeed is None.

[havogt]

Added in 7faefef95, with is None and scope_dict() hoisted out of the comprehension:

write_backs = []
for state, (tmp_reads, _) in tmp_access.items():
    scope_dict = state.scope_dict()
    write_backs.extend(
        (edge, state)
        for node in tmp_reads
        if scope_dict[node] is None
        for edge in state.out_edges(node)
        if isinstance(edge.dst, dace_nodes.AccessNode)
        and not edge.dst.desc(sdfg).transient
    )

The check sits in the node loop, so it runs once per candidate node rather than once per out edge, and scope_dict() once per state. The comment gives ADR-18 rule 10 as the reason: an AccessNode inside a Map scope has Scope lifetime and can therefore not refer to the data this pass removes.

33 unit tests and the whole dace transformation suite are green; the icon4py dycore run is still going, I will report if it turns up anything.

[philip-paul-mueller]

Suggested change

	# Equal sizes do not imply equal ranges, `Range.size()` divides by the
	# step, so `[0:8:2]` and `[0:4]` have the same size. The source is the
	# full array and therefore has unit steps, and shifting an access can not
	# turn it into a scatter, so the destination must have them too.

[havogt]

Removed. One thing it takes with it: the comment recorded why testing only the destination is enough, namely that src_subset has to equal Range.from_array(tmp_desc) and therefore has unit steps. Tell me if you want that as a one liner instead.

[philip-paul-mueller]

Move it up since it is a super cheap test.
General rule: Order the tests in such a way that you can bail out as fast as possible.

Also the rule below should be moved up.

[havogt]

Both moved up, right after the write back edge is picked, together with glob_node = wb_edge.dst. They are an attribute read and a degree lookup and do not depend on the subset work, so this is only a reordering.

[philip-paul-mueller]

Consider putting this into an else branch.
Then you would not have two "remove old edge" but also not this continue.

[havogt]

Done, else branch, one remove_edge() and no continue.

[philip-paul-mueller]

LGTM

With the additional is None check, thanks for spotting.

[philip-paul-mueller]

Yes, you are right it is indeed is None.