feat[next-dace]: fold single iteration Map dimensions by havogt · Pull Request #73 · havogt/gt4py

havogt · 2026-08-10T11:51:27Z

Staged on the fork — not yet proposed upstream.

A Map dimension i = c:c+1:1 leaves Memlets referring to i symbolically. Range.covers() does not know the Map range, so a producer writing a[c] is not seen to cover a consumer reading a[i-o], and a legal fusion is rejected.

Folding the constant into the scope subgraph removes the blind spot. The dimension is kept: dace's own TrivialMapElimination drops it and then refuses GPU schedules.

In the SDFG

From compute_perturbed_quantities_and_interpolation with the static bounds the model run bakes in. i_K_gtx_vertical = 8 is the surface level, so the Map has a single iteration in K while keeping the dimension.

map_18_fieldop, range 1:17, 8 before and after — only the Memlets change:

                       before                                    after
MapEntry -> Tasklet    gtir_tmp_32[i_Cell - 1, i_K - 8]          gtir_tmp_32[i_Cell - 1, 0]
MapEntry -> Tasklet    gtir_tmp_38[i_Cell - 1, i_K - 8]          gtir_tmp_38[i_Cell - 1, 0]
Tasklet  -> MapExit    gtir_tmp_40[i_Cell - 1, i_K - 8]          gtir_tmp_40[i_Cell - 1, 0]

The folded value need not be 0, and symbols that are not Map parameters are left alone. map_14_fieldop, same range:

                       before                                    after
MapEntry -> Tasklet    wgtfacq_c[.., -__wgtfacq_c_K_range_0       wgtfacq_c[.., 7 - __wgtfacq_c_K_range_0]
                                     + i_K - 1]
MapEntry -> Tasklet    gtir_tmp_14[i_Cell - 1, i_K - 1]          gtir_tmp_14[i_Cell - 1, 7]
Tasklet  -> MapExit    gtir_tmp_32[i_Cell - 1, i_K - 8]          gtir_tmp_32[i_Cell - 1, 0]

i_Cell is untouched throughout: its dimension has more than one iteration.

What it changes in the final SDFG

At the point in _gt_auto_process_top_level_maps where the pass first matches, the program has 48 top level Maps, 11 of them single iteration in K over the surface level (1:17, 80 with 80 levels). Those 11 are what the pass unblocks.

Top level Maps of the whole program after gt_auto_optimize, by range, with and without the pass. Identical at 8 and at 80 levels:

range           without   with
0, 1:15               1      1
0, 1:17               1      1
0, 15:18              1      1
0, 17                 1      1
0:8, 1:17             1      1
1:4, 1:15             2      2
1:4, 15:17            1      1
1:4, 15:18            1      1
1:8, 17               1      1
4, 1:15               2      2
4, 15:17              2      2
4:8, 1:15             1      1
4:8, 15:18            1      1
5:8, 1:15             3      3
5:8, 15:17            3      3
8, 1:17               2      1   <==
TOTAL                24     23

The only difference is the surface: the two Maps over K = 8 become one. Nothing else in the program moves, which is what the pass is meant to do — it removes a comparison failure, it does not restructure anything.

Measurements

icon4py compute_perturbed_quantities_and_interpolation, GH200, dace_gpu. Measured before PR GridTools#2763 and GridTools#2764 landed, where it collapsed the three surface kernels into one, 11 → 9 kernels. On the current pipeline the remaining effect is the one above, two surface Maps into one.

It also removes a real non-determinism: the same program compiled to 10 kernels alone and 11 in a 3-variant session.

The wall-clock contribution measured on its own was small (~0.01 s on the target program), which is why this was deprioritised relative to the other changes. The determinism argument is the stronger one now, given GridTools#2779, GridTools#2780 and GridTools#2786.

Testing

test_trivial_map_dimension_folding.py, 3 passed. Merged to main at 0d8fb34d3; the full dace_tests unit suite passes on the merge, with one pre-existing environment failure (test_dace_fastcall[exec_alloc_descriptor1], missing cupy) that reproduces on unmodified main.

Status

Independent of everything else in the stack. Not reviewed by anyone yet.

A Map dimension `i = c:c+1:1` can only ever take the value `c`, but the Memlets inside its scope keep referring to `i` symbolically. `MapFusionVertical` compares producer and consumer subsets with `Range.covers()`, which does not know the Map range, so a producer writing `a[c]` is not recognized as covering a consumer reading `a[i - o]` and a legal fusion is rejected. The new transformation folds the value into the scope and runs before the top level fusion rounds. The dimension itself is kept, unlike DaCe's `TrivialMapElimination` which removes it and thereby prevents the Map from being scheduled on the GPU. On `compute_perturbed_quantities_and_interpolation` this collapses the three surface level kernels into one, 10 -> 9 kernels. It also removes a non-determinism: the same program compiled to 10 or 11 kernels depending on which other variants were compiled in the same session. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…-dimensions

iomaganaris

I'm not sure how it can help in the large stencil (didn't run it locally to check) but the transformation looks reasonable

havogt · 2026-08-14T14:41:04Z

cpqi_sdfgs_pr73.zip

edopao · 2026-08-14T14:53:05Z

+        return {
+            param: rng[0]
+            for param, rng in zip(map_.params, map_.range.ranges)
+            if (rng[0] == rng[1]) == True and (rng[2] == 1) == True  # noqa: E712 [true-false-comparison]  # SymPy comparison


Probably it is not needed to check (rng[2] == 1) == True. This is the step between two elements to read/write, but we just need to check that we access a single element ((rng[0] == rng[1]) == True).

I agree the rng[2] == 1 test is unneeded.

Agreed, dropped in bda6175. With start == end the dimension has a single iteration for any step, so the step carries no information here.

Dropped in bda6175.

edopao · 2026-08-14T14:56:03Z

The transformation is interesting. I am surprised that dace cannot do this on its own.

edopao · 2026-08-14T14:59:03Z

The transformation is interesting. I am surprised that dace cannot do this on its own.

Maybe because we are skipping ConstantPropagation in our optimization pipeline?

This pull request includes updates to dependencies, configuration handling for CUDA architecture detection, and a CI job setting. The most significant changes are an upgrade to the required DaCe version, an improved approach to configuring CUDA architecture detection, and a minor CI configuration adjustment. **Dependency and configuration updates:** * Upgraded the required `dace` package version from `2.0.0a5` to `2.0.0a6` in `pyproject.toml`, ensuring compatibility with recent DaCe features and fixes. * Changed how CUDA architecture detection is handled in `set_dace_config` in `common.py`: now sets `compiler.cuda.cuda_arch` directly instead of passing a CMake argument, aligning with best practices and recent DaCe updates. **CI/CD configuration:** * Removed `allow_failure: true` from the `build_cscs_amd_rocm` CI job in `ci/cscs-ci-dace-determinism.yml`, making failures in this job block the pipeline.

philip-paul-mueller

The transformation looks okay, I just have some small things.

philip-paul-mueller · 2026-08-17T05:36:15Z

+        return {
+            param: rng[0]
+            for param, rng in zip(map_.params, map_.range.ranges)
+            if (rng[0] == rng[1]) == True and (rng[2] == 1) == True  # noqa: E712 [true-false-comparison]  # SymPy comparison


I agree the rng[2] == 1 test is unneeded.

philip-paul-mueller · 2026-08-17T05:43:00Z

+            ):
+                return True
+        for node in scope.nodes():
+            if node is not map_entry and any(str(sym) in replacements for sym in node.free_symbols):


What is about the associated map exit?

It is covered, on both counts.

scope_subgraph(map_entry, include_entry=True, include_exit=True) contains the MapExit, so its edges are checked in can_be_applied and rewritten in apply. Its own free_symbols turns out to be empty (so is the MapEntry's), so the node is not map_entry guard is about the general case rather than this one.

The restore reaches it too, because MapEntry and MapExit share one Map object:

MapEntry.map is MapExit.map: True after mutating entry.map.params, exit sees: ['X', 'Y']

so restoring through map_entry.map restores what the exit sees.

philip-paul-mueller · 2026-08-17T05:46:28Z

+        The dimension itself is kept, in contrast to DaCe's `TrivialMapElimination`
+        which removes it. Removing it changes the shape of the Map and prevents it
+        from being scheduled on the GPU, whereas folding only rewrites the uses of
+        the parameter and leaves the iteration space, the schedule and any blocking


The problem with the GPU scheduling is, that you can access data that resides on the GPU only inside a Map, whereas data on CPU can be accessed everywhere except a GPU kernel (with the exception of scalar data, which should be possible to access there as well).

This note, while interesting should be shortened to something like:
Unlike the DaCe native TrivialMapElimination transformation the dimension is not removed only the single value is folded into the body of the Map.

The rest are external consequences.

Thanks, shortened in bda6175 to:

Note: Unlike DaCe's native `TrivialMapElimination` the dimension is not removed, only the single value is folded into the body of the Map.

philip-paul-mueller · 2026-08-17T05:49:23Z

+        for param, value in replacements.items():
+            scope.replace(param, value)


There is also replace_dict() which takes a dict and operates on it as a whole.
Should be more efficient.

Done in bda6175, one call instead of one per parameter.

philip-paul-mueller · 2026-08-17T05:55:05Z

+        replacements = self._single_iteration_parameters(map_entry.map)
+        scope = graph.scope_subgraph(map_entry, include_entry=True, include_exit=True)
+
+        # `replace()` would also rewrite the Map's own parameters and range, which


This is more my curiosity.
Before the parameters and the ranges of the Map were not subject to replacement, this is why MapFusion does it manually.
So I am wondering if the replace function does it now or not?

Don't get me wrong, you should not remove the code that does it, as this "not replacement" is a bug in my opinion.

Not yet, same as when MapFusion started doing it by hand. Measured on this DaCe version:

after replace: params ['__i', '__j'] -> ['__i', '__j'] after replace: range 0:20, 7 -> 0:20, 7

So the save and restore is currently dead code. Keeping it as you suggest, since it is the behaviour we want to be robust against if that ever changes.

philip-paul-mueller · 2026-08-17T05:56:52Z

+    assert map_entry.map.params == ["__i", "__j"]
+    assert str(map_entry.map.range[1][0]) == "7"
+    assert all(
+        "__j" not in str(edge.data.subset)


To be complete there is also other_subset that should be tested.
However, I expect that all of them are none anyway, but technically testing only subset is not enough.

Added in bda6175, the assertion now iterates over both subset and other_subset of every Memlet. As you expected they are all None here, so it passes either way.

philip-paul-mueller · 2026-08-17T05:58:20Z

+    assert all(np.allclose(ref[name], res[name]) for name in ref.keys())
+
+
+def test_trivial_map_dimension_folding_is_a_fixpoint():


What is the difference between this test and the test above.
Both run the transformation using apply_transformations_repeated() and check if it returns 1, so both are a fix point, the test above, just also tests if the result is correct.
Either delete this test or remove the _reapeat above.

You are right, they overlap. Deleted the separate fixpoint test and kept apply_transformations_repeated in the main one, with a comment saying that a single application under repeated application is what makes it a fixpoint.

…rfaced by the Python 3.12 floor (GridTools#2755) Audit of the codebase after the floor moved to 3.12 (GridTools#2326). Bugs fixed: - `cartesian`: `...` index detection matched the deprecated `ast.Ellipsis` on 3.12/3.13 and nothing at all on 3.14. Now matched via `_is_ellipsis_node`, which also fixes the regression from GridTools#2772 (`types.EllipsisType` is the type of the `...` object, not of its AST node). - `next`: `_type_conversion_helper` raised on PEP 604 unions and the builtin `tuple` spelling; it now dispatches on `get_origin()`. - `eve`: `frozen="strict"` called a non-existent `xtyping.is_hashable_type` and raised at class definition. The check now decomposes composite annotations, resolves PEP 695 aliases, and names the offending fields. - `next`: `test_typing_exports` picked up the project `[tool.mypy]` table, so the 3.13/3.14 runs duplicated 3.12. Adds `typing_tests/mypy.ini`. - `next`: `wait_for_compilation` flattened multiple failures into one `RuntimeError`; they are now raised as a group, each keeping its traceback. Deprecated and dead code removed: - `eve`: the PEP 585 `typing` aliases (`Dict`, `FrozenSet`, `List`, `Set`, `Tuple`, `Type`) are dropped from the namespace and rejected explicitly, and the use sites move to the builtin generics. Forward-reference resolution in user annotations is unchanged. - Dead `sys.version_info` / `typing_extensions.Any` guards, and the `add_note` override in `next/errors/exceptions.py`, dead since PEP 678 landed in 3.11. Diagnostics: `try*` is named as a construct, and constructs newer than the floor (`ast.TemplateStr`) are registered by name and skipped where absent. Build and docs: `.gitpod.Dockerfile` Python version, the noxfile PEP 723 `requires-python`, the floor stated in `AGENTS.md`, an explicit `[tool.mypy] python_version`, and stale mypy / isort overrides.

… folding - drop the `rng[2] == 1` step check: with `start == end` the dimension has a single iteration for any step, so the step carries no information here. - fold with a single `replace_dict()` instead of one `replace()` per parameter. - shorten the note contrasting the transformation with `TrivialMapElimination` to what distinguishes them, the dimension being kept rather than removed. - test the `other_subset` of every Memlet as well, not only the `subset`. - drop the separate fixpoint test: the main test already applies the transformation repeatedly and asserts a single application. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

havogt · 2026-08-17T08:45:14Z

Good question — I checked, and ConstantPropagation is not the reason.

It propagates symbols that were assigned on interstate edges; its should_apply() scans sdfg.all_interstate_edges() for assignments. A Map parameter is not one of those, it is bound by the Map scope, so the pass has nothing to look at. Run against the SDFG at the point where this transformation first matches (the before file in the zip above):

before ConstantPropagation    single iteration maps= 11   memlets still using the param= 30
  should_apply(sdfg) = False
  apply_pass returned: None
after  ConstantPropagation    single iteration maps= 11   memlets still using the param= 30

It early exits before doing any work. Note we also do not avoid it entirely: it is in GT_SIMPLIFY_DEFAULT_SKIP_SET, but simplify.py constructs and runs it explicitly with initial_symbols=repl for compile time symbol substitution.

What DaCe does have for this is TrivialMapElimination, and it fires 11 times on the same SDFG:

before:  params=['i_Cell_gtx_horizontal', 'i_K_gtx_vertical']   range=1:17, 80
after:   params=['i_Cell_gtx_horizontal']                        range=1:17

It removes the dimension, which is what we cannot use here, hence folding instead. So DaCe can do it, just not in a form that survives the GPU schedule.

…-dimensions

- point the property description at the class docstring - drop the note on `include_entry`, it holds for `include_exit` just as much Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

havogt force-pushed the dace-fold-trivial-map-dimensions branch from e1061f3 to 84c5978 Compare August 10, 2026 13:37

havogt changed the base branch from main to dace-split-access-node-multi-producer August 10, 2026 13:44

havogt changed the base branch from dace-split-access-node-multi-producer to main August 10, 2026 13:44

Merge remote-tracking branch 'origin/main' into dace-fold-trivial-map…

8f7a3cc

…-dimensions

iomaganaris approved these changes Aug 14, 2026

View reviewed changes

edopao reviewed Aug 14, 2026

View reviewed changes

philip-paul-mueller reviewed Aug 17, 2026

View reviewed changes

egparedes and others added 2 commits August 17, 2026 08:56

Merge remote-tracking branch 'origin/main' into dace-fold-trivial-map…

9187c2e

…-dimensions

havogt mentioned this pull request Aug 17, 2026

feat[next-dace]: fold single iteration Map dimensions GridTools/gt4py#2799

Merged

1 of 2 tasks

refactor[next-dace]: address review comments on 2799

586c5b4

- point the property description at the class docstring - drop the note on `include_entry`, it holds for `include_exit` just as much Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

		for param, value in replacements.items():
		scope.replace(param, value)

		assert all(np.allclose(ref[name], res[name]) for name in ref.keys())


		def test_trivial_map_dimension_folding_is_a_fixpoint():

Conversation

havogt commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

In the SDFG

What it changes in the final SDFG

Measurements

Testing

Status

Uh oh!

iomaganaris left a comment

Choose a reason for hiding this comment

Uh oh!

havogt commented Aug 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

edopao commented Aug 14, 2026

Uh oh!

edopao commented Aug 14, 2026

Uh oh!

philip-paul-mueller left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

havogt commented Aug 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

havogt commented Aug 10, 2026 •

edited

Loading