Add distribution of thread blocks to chiplets in old codegen by iomaganaris · Pull Request #2527 · spcl/dace

iomaganaris · 2026-08-26T08:02:01Z

The goal of this PR is to distribute the thread blocks of the GPU backend to AMD GPU chiplets in such way that the X dimension of the grid is divided by the number of chiplets, so each one of them gets a continuous range of the domain. The domains that execute on the same chiplet can use then the same L2 cache, improving cache hits when there are data that can be reused between thread blocks in the same region or data that live in one dimension in two dimensional kernels (i.e. ICON neighbor tables and other vertical level independent fields).
This change has a ~7% performance improvement in the icon4py dycore.
We have already discussed about the necessity of this solution with @ThrudPrimrose I was interested to implement it in the old codegen though because we still haven't switched to the new one

…plet_old_codegen_updateddace

Drop two changes that came along with the chiplet commits but are unrelated to the chiplet mapping and active regardless of it: * The `threadIdx.y` -> `threadIdx.z` remap for 2D thread-block maps. Block sizes are never remapped by the chiplet patch, so the second index of a 2D thread-block map would be read from `threadIdx.z`, which is 0 for a (bx, by, 1) block. Since `AddThreadBlockMap` is now applied everywhere, this fired on nearly every kernel. * The `/*block*/` and `/*threads*/` markers in the generated code. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GdEk8Xdh2bd53mE3tf5Thy

… generator Distributing the thread-blocks of a kernel over the chiplets of the GPU moves the second dimension of the grid to `blockIdx.z`, which the merged implementation did not account for in every case: * Grids using all three dimensions had their third dimension silently dropped, or were folded into an index expression asking for a fourth `blockIdx`, which raises. The distribution is now only applied when the third grid dimension is 1, and kernels that cannot use it keep their grid and report why. * Nested device maps and dynamic thread-block maps read `blockIdx` outside of `generate_kernel_scope`, where the chiplet ID would be taken for a work index. Both are excluded as well. * The first grid dimension is padded to a multiple of the number of chiplets, but the condition masking out the padded thread-blocks was elided whenever the map divided evenly by the block size, and was not generated at all for kernels with an inner thread-block map. It is now always generated. The default of `compiler.cuda.chiplet_number` is 1, i.e. the grid is left alone unless the distribution is asked for (`DACE_compiler_cuda_chiplet_number=6` on MI300A). The generated code with the default is identical to the code generated before the merge. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GdEk8Xdh2bd53mE3tf5Thy

ThrudPrimrose · 2026-08-26T08:34:16Z

The assignment of thread-blocks to physical locations is a matter of codegen, yet right now the choice is too implicit. It is also not providing the user the option to choose between options. (This change would force all kernels to be generated either the old way or the new chiplet-aware way)

I think it would make more sense to have a map-level attribute (to also not break the existing assumption of "GPU_Device" is a GPU kernel and no a new schedule type to update everything)

I see that, through the tests the existence of tblock maps it should be fine.

For this pattern (and new codegen enforces it):

for i,j in dace.map[0:D0:16, 0:D1:16] @ GPU_Device:
  for bi, bj in dace.map[0:16, 0:16] @ GPU_ThreadBlock: 
     <body>

Codegen only changes the tblocks are assigned to chiplets so it seems fine to me.

ThrudPrimrose · 2026-08-26T08:36:39Z

For test quality we should have a test where we check numerical equivalence between the two strategies (e.g. maybe a very small icon loopnest snippet written using the python frontend?). Also please add a test involving persistent schedule (so that we dont break it unknowingly)

…tiple chiplet numbers

…ight cases are tested

ThrudPrimrose

All previous comments are addressed.

The codegen is opt-in and I like it. We also have enough numerical correctness checks.

Only one question? Would it make sense, and justify the effort to auto-detect chiplet count if a value like 0 is provided?
We could have a simple dictionary between compute arch/model to chiplet count to make it simpler to run on multiple amd devises? What do you think?

iomaganaris · 2026-08-28T15:27:57Z

All previous comments are addressed.

The codegen is opt-in and I like it. We also have enough numerical correctness checks.

Only one question? Would it make sense, and justify the effort to auto-detect chiplet count if a value like 0 is provided? We could have a simple dictionary between compute arch/model to chiplet count to make it simpler to run on multiple amd devises? What do you think?

Thank you very much for the review.
I have some logic added here in gt4py to detect the number of XCDs based on the amdsmi python package if it exists. I would be okay moving this logic to DaCe or use a look up table. Whatever you prefer

tbennun · 2026-08-28T15:52:45Z

+        default=True,
+        desc="Allow the thread-blocks of this kernel to be distributed over the chiplets of the GPU "
+        "(see the `compiler.cuda.chiplet_number` configuration entry)",
+        serialize_if=lambda m: m.schedule in dtypes.GPU_SCHEDULES)


this should only be for GPU_Device and GPU_Persistent, right? not GPU_Threadblock[Dynamic]

If GPU_Persistent is enabled then based on the chiplet_count function nothing happens. I thought it would actually make sense probably to set this for the GPU_Threadblock schedule maybe?

tbennun · 2026-08-28T15:55:02Z

+            # contiguous chunk of ``ceil(grid_size[0] / chiplets)`` blocks of the first dimension,
+            # together with the full second dimension, which moves to ``blockIdx.z``.
+            original_grid_size = grid_size
+            grid_size = [self._kernel_chiplet_count, int_ceil(grid_size[0], self._kernel_chiplet_count), grid_size[1]]


I am not sure why we should change the grid size; that means that the grid sizes are effectively minimized because blockdim x is the one that can be the largest of the three.
I'd consider a straightforward code generation approach (changing the blocks based on a modulo operator rather than changing the actual map dimensions.

I deliberately wanted to avoid the module operations and the thread ID translation because they add extra resource usage and complicates further the generated code. Unless there is actual need to support 3D maps I would keep this simpler approach

Since I realized that different grid dimensions have different allowed ranges I have updated the logic to be codegen only with % and / usage

tbennun

I am a big fan of the idea, but I am not sure we should change the map dimensions to provide that. It should be a purely generated-code approach as it changes the mapping of logical thread-blocks to physical thread-block numbers.

…pdateddace

Co-authored-by: Tal Ben-Nun <tbennun@gmail.com>

iomaganaris · 2026-08-31T10:18:16Z

+            # together with the full second dimension, which moves to ``blockIdx.z``.
+            original_grid_size = grid_size
+            grid_size = [self._kernel_chiplet_count, int_ceil(grid_size[0], self._kernel_chiplet_count), grid_size[1]]
+            warnings.warn(f'Distributing the grid of kernel "{kernelmap_entry.map.label}" over '


This warning prints a lot of warnings like:

/capstor/scratch/cscs/ioannmag/cycle38/icon4py-benchmarks/venv_mi300/lib/python3.12/site-packages/dace/codegen/targets/cuda.py:2321: UserWarning: Distributing the grid of kernel "map_0_fieldop" over 6 chiplets, adjusting its size from [649, 1, 1] to [6, 109, 1].

when the option is enabled.
I am not really sure if this is beneficial or not

…pdateddace

…on package

ThrudPrimrose · 2026-09-09T07:50:57Z

@tbennun I think the latest stage is good can re-review?

tbennun · 2026-09-09T15:22:10Z

cscs-ci run

tbennun · 2026-09-09T15:22:52Z

Thanks, will review

tbennun · 2026-09-10T03:12:40Z

@iomaganaris looks like the test is failing

…l#2538)

A map body carried into a nested SDFG got a local descriptor sized from the number of elements of the outer subset, while the inner memlet was only shifted by the outer start. For a strided outer subset such as A[15*i:15*i+28:3] this left an inner range of 0:28:3 against a shape-(10,) array, and parsing failed with "Memlet subset out-of-bounds".

`StripMining._create_strided_range` over-approximated a non-divides-evenly tile as `tile_size` indices wide while its exact bound spans `tile_size * td_step`, so `Range.size()` reported a strided map's tile as `tile_size / td_step` iterations -- tiling a step-2 map by 128 produced a thread-block map that `InferGPUGridAndBlockSize` measured at 64 threads and rejected against the kernel's declared `gpu_block_size`. The approximation now carries the step, matching the exact bound and the `divides_evenly` branch one line above, and is identical for `td_step == 1` so no existing caller changes. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

iomaganaris · 2026-09-10T07:32:54Z

cscs-ci run

iomaganaris · 2026-09-10T07:46:52Z

@tbennun the test is now fixed. Could you please have another look?

@ThrudPrimrose

The goal of this PR is to distribute the thread blocks of the GPU backend to AMD GPU chiplets in such way that the X dimension of the grid is divided by the number of chiplets, so each one of them gets a continuous range of the domain. The domains that execute on the same chiplet can use then the same L2 cache, improving cache hits when there are data that can be reused between thread blocks in the same region or data that live in one dimension in two dimensional kernels (i.e. ICON neighbor tables and other vertical level independent fields). This change has a ~7% performance improvement in the `icon4py` dycore. We have already discussed about the necessity of this solution with @ThrudPrimrose I was interested to implement it in the old codegen though because we still haven't switched to the new one --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: Tal Ben-Nun <tbennun@gmail.com> Co-authored-by: Yakup Koray Budanaz <budanaz.yakup@gmail.com>

iomaganaris and others added 6 commits April 29, 2026 13:36

Initial changes in grid dimensions

3746bbb

Fixes in grid size

294157b

Small fix in code generation

3d5d463

Merge remote-tracking branch 'gridtools/chiplet_old_codegen' into chi…

3e29f19

…plet_old_codegen_updateddace

iomaganaris added 2 commits August 26, 2026 18:00

Added explicit test for GPU_ThreadBlock

f003ac3

Added couldsc_tidy_branch test for correctness in CPU and GPU for mul…

ca6ec4f

…tiple chiplet numbers

iomaganaris mentioned this pull request Aug 27, 2026

perf[next-dace]: Added automatic setting of chiplet number using amdsmi GridTools/gt4py#2832

Closed

Clear cache before and after each sdfg compilation to make sure the r…

a149159

…ight cases are tested

ThrudPrimrose reviewed Aug 28, 2026

View reviewed changes

tbennun reviewed Aug 28, 2026

View reviewed changes

Comment thread dace/config_schema.yml Outdated

tbennun reviewed Aug 28, 2026

View reviewed changes

tbennun requested changes Aug 28, 2026

View reviewed changes

iomaganaris and others added 3 commits August 31, 2026 11:38

Merge remote-tracking branch 'origin/main' into chiplet_old_codegen_u…

0ad7260

…pdateddace

Added check to make sure that the map range is only 2D

accef3c

Address Tal's review comments

5115128

Co-authored-by: Tal Ben-Nun <tbennun@gmail.com>

iomaganaris commented Aug 31, 2026

View reviewed changes

iomaganaris added 4 commits September 1, 2026 15:36

Merge remote-tracking branch 'origin/main' into chiplet_old_codegen_u…

42f7f5c

…pdateddace

Updated chiplet distribution with module to avoid 3D kernel limitation

8c7c82e

Enable automatic configuration of XCD number based on the amdsmi pyth…

67edb12

…on package

Merge branch 'main' into chiplet_old_codegen_updateddace

b2adcb0

ThrudPrimrose self-requested a review September 9, 2026 07:17

ThrudPrimrose approved these changes Sep 9, 2026

View reviewed changes

Merge branch 'main' into chiplet_old_codegen_updateddace

2753e97

ThrudPrimrose requested a review from tbennun September 9, 2026 07:50

iomaganaris and others added 5 commits September 10, 2026 09:31

Disable report generation for InstrumentationType.GPU_TX_MARKERS (spc…

dd64f7b

…l#2538)

Fix cloudsc test with chiplets

6ed22ef

Merge branch 'main' into chiplet_old_codegen_updateddace

8a61b5d

tbennun approved these changes Sep 11, 2026

View reviewed changes

tbennun added this pull request to the merge queue Sep 11, 2026

Merged via the queue into spcl:main with commit 7e4b35e Sep 11, 2026
17 checks passed

Conversation

iomaganaris commented Aug 26, 2026

Uh oh!

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

Uh oh!

Uh oh!

ThrudPrimrose commented Aug 26, 2026

Uh oh!

ThrudPrimrose left a comment

Choose a reason for hiding this comment

Uh oh!

iomaganaris commented Aug 28, 2026

Uh oh!

tbennun Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

iomaganaris Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tbennun Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

iomaganaris Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

iomaganaris Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

tbennun left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

iomaganaris Aug 31, 2026

Choose a reason for hiding this comment

Uh oh!

ThrudPrimrose commented Sep 9, 2026

Uh oh!

tbennun commented Sep 9, 2026

Uh oh!

tbennun commented Sep 9, 2026

Uh oh!

tbennun commented Sep 10, 2026

Uh oh!

iomaganaris commented Sep 10, 2026

Uh oh!

iomaganaris commented Sep 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ThrudPrimrose commented Aug 26, 2026 •

edited

Loading

tbennun left a comment •

edited

Loading