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

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

Draft

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

Navigation Menu

Uh oh!

Add distribution of thread blocks to chiplets in old codegen - #2527

Add distribution of thread blocks to chiplets in old codegen#2527
iomaganaris wants to merge 16 commits into
spcl:mainfrom
GridTools:chiplet_old_codegen_updateddace

iomaganaris commented Aug 26, 2026

Uh oh!

ThrudPrimrose commented Aug 26, 2026 •

edited

Loading

Uh oh!

ThrudPrimrose commented Aug 26, 2026

Uh oh!

ThrudPrimrose left a comment

Uh oh!

iomaganaris commented Aug 28, 2026

Uh oh!

tbennun Aug 28, 2026

Uh oh!

iomaganaris Aug 31, 2026

Uh oh!

Uh oh!

tbennun Aug 28, 2026

Uh oh!

iomaganaris Aug 31, 2026

Uh oh!

iomaganaris Sep 3, 2026

Uh oh!

tbennun left a comment •

edited

Loading

Uh oh!

iomaganaris Aug 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ThrudPrimrose commented Aug 26, 2026 •

edited

Loading

tbennun left a comment •

edited

Loading