Add distribution of thread blocks to chiplets in old codegen

dace/codegen/targets/cuda.py

@@ -89,6 +89,9 @@ def __init__(self, frame_codegen: 'DaCeCodeGenerator', sdfg: SDFG):
         self._cpu_codegen: Optional['CPUCodeGen'] = None
         self._block_dims = None
         self._grid_dims = None
+        # Number of chiplets the grid of the kernel being generated is distributed over
+        # (1 when the distribution does not apply, see ``chiplet_count``)
+        self._kernel_chiplet_count = 1
         self._kernel_map = None
         self._kernel_state = None
         self._kernel_grid_conditions: List[str] = []
@@ -2071,6 +2074,67 @@ def get_tb_maps_recursive(self, subgraph):
                 res.append((node.map, {dace.symbol(k): dace.symbol(k) for k in node.map.range.free_symbols}))
         return res
 
+    def chiplet_count(self, kernelmap_entry: nodes.MapEntry, grid_size: List[symbolic.SymbolicType],
+                      is_persistent: bool, has_dtbmap: bool, extra_grid_dims: List[symbolic.SymbolicType]) -> int:
+        """
+        Returns the number of chiplets (XCDs on AMD GPUs) the grid of the given kernel is distributed over,
+        or 1 if the distribution does not apply to it.
+
+        The distribution moves the second grid dimension to ``blockIdx.z`` so that ``blockIdx.x`` is free to
+        carry the chiplet ID. It therefore only applies to kernels whose third grid dimension is 1, whose
+        kernel map has at most two dimensions of its own (a third dimension that only happens to size to a
+        single block is not enough: it is the map's dimensionality, not the derived grid size, that the
+        rest of the distribution logic in ``generate_kernel_scope`` assumes), and whose block indices are
+        generated by ``generate_kernel_scope`` alone. Kernels whose map has the
+        ``allow_chiplet_threadblock_distribution`` property set to False are left alone as well.
+
+        :param kernelmap_entry: Entry node of the kernel map.
+        :param grid_size: Size of the grid, in thread-blocks.
+        :param is_persistent: Whether the kernel uses a persistent grid.
+        :param has_dtbmap: Whether the kernel contains a dynamic thread-block map.
+        :param extra_grid_dims: Grid dimensions contributed by nested device maps, if any.
+        """
+        chiplets = int(Config.get('compiler', 'cuda', 'chiplet_number'))
+        if chiplets < 1:
+            raise ValueError(f'Invalid number of chiplets ({chiplets}) configured. Modify the '
+                             '`compiler.cuda.chiplet_number` configuration entry to a positive number.')
+
+        # A kernel that opts out of the distribution is left alone without any diagnostics
+        if not kernelmap_entry.map.allow_chiplet_threadblock_distribution:
+            return 1
+
+        if chiplets == 1:
+            if self.backend == 'hip':
+                warnings.warn(f'Not distributing the grid of kernel "{kernelmap_entry.map.label}" over the chiplets '
+                              'of the GPU because `compiler.cuda.chiplet_number` is set to 1. Set it to the number '
+                              'of chiplets of the GPU (6 on MI300A) to enable the distribution.')
+            return 1
+
+        if self.backend != 'hip':
+            warnings.warn(f'`compiler.cuda.chiplet_number` is set to {chiplets}, but the "{self.backend}" backend '
+                          'targets GPUs without chiplets. Distributing the grid over chiplets relies on the '
+                          'round-robin thread-block scheduling of multi-chiplet AMD GPUs.')
+
+        skip_reason = None
+        if is_persistent:
+            skip_reason = 'it uses a persistent grid'
+        elif has_dtbmap:
+            skip_reason = 'it contains a dynamic thread-block map'
+        elif extra_grid_dims:
+            skip_reason = 'it contains nested device maps'
+        elif symbolic.equal(grid_size[2], 1) is not True:
+            skip_reason = f'its third grid dimension ({grid_size[2]}) is not 1'
+        elif kernelmap_entry.map.range.dims() > 2:
+            skip_reason = f'its map has {kernelmap_entry.map.range.dims()} dimensions, but the distribution ' \
+                          'only supports up to two'
+
+        if skip_reason is not None:
+            warnings.warn(f'Not distributing the grid of kernel "{kernelmap_entry.map.label}" over {chiplets} '
+                          f'chiplets because {skip_reason}.')
+            return 1
+
+        return chiplets
+
     def get_kernel_dimensions(self, dfg_scope):
         """
         Determines a GPU kernel's grid/block dimensions from map scopes.
@@ -2281,6 +2345,22 @@ def get_kernel_dimensions(self, dfg_scope):
                              'thread-block size. To increase this limit, modify the '
                              '`compiler.cuda.block_size_lastdim_limit` configuration entry.')
 
+        # Distribute the thread-blocks of the first grid dimension over the chiplets of the GPU
+        self._kernel_chiplet_count = self.chiplet_count(kernelmap_entry, grid_size, is_persistent, has_dtbmap,
+                                                        extra_grid_dims)
+        if self._kernel_chiplet_count > 1:
+            # ``blockIdx.x`` carries the chiplet ID: the hardware dispatches thread-block ``f`` to
+            # chiplet ``f % gridDim.x`` in a round-robin fashion, and the flattened block index is
+            # ``blockIdx.x + blockIdx.y * gridDim.x + ...``, so making ``gridDim.x`` the number of
+            # chiplets pins each chiplet to one value of ``blockIdx.x``. Each chiplet then receives a
+            # 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]]
+            warnings.warn(f'Distributing the grid of kernel "{kernelmap_entry.map.label}" over '
+                          f'{self._kernel_chiplet_count} chiplets, adjusting its size from {original_grid_size} to '
+                          f'{grid_size}.')
+
         return grid_size, block_size, len(tb_maps_sym_map) > 0, has_dtbmap, extra_dim_offsets
 
     def generate_kernel_scope(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg_scope: ScopeSubgraphView, state_id: int,
@@ -2336,13 +2416,27 @@ def generate_kernel_scope(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg_scope: S
                 dfg_scope.source_nodes()[0])
 
         # do not generate an index if the kernel map is persistent
+        chiplet_count = self._kernel_chiplet_count
         if node.map.schedule != dtypes.ScheduleType.GPU_Persistent:
             # First three dimensions are evaluated directly
             for i in range(min(len(krange), 3)):
                 varname = kernel_map.params[-i - 1]
 
-                # If we defaulted to a fixed number of threads per block, offset by thread ID
-                block_expr = 'blockIdx.%s' % _named_idx(min(i, 2))
+                if chiplet_count > 1:
+                    if i == 0:
+                        # Contiguous partitioning: the chiplet ID is blockIdx.x and the slot within the
+                        # chiplet is blockIdx.y, so chiplet k owns the blocks
+                        # [k * gridDim.y .. (k + 1) * gridDim.y - 1] of this dimension
+                        block_expr = '(blockIdx.x * gridDim.y + blockIdx.y)'
+                    elif i == 1:
+                        # The second dimension was shifted to make room for the chiplet ID
+                        block_expr = 'blockIdx.z'
+                    else:
+                        # Dimensions beyond the second span a single block (see ``chiplet_count``)
+                        block_expr = '0'
+                else:
+                    # If we defaulted to a fixed number of threads per block, offset by thread ID
+                    block_expr = 'blockIdx.%s' % _named_idx(min(i, 2))
                 if not has_tbmap or has_dtbmap:
                     block_expr = '(%s * %s + threadIdx.%s)' % (block_expr, _topy(block_dims[i]), _named_idx(i))
 
@@ -2360,7 +2454,9 @@ def generate_kernel_scope(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg_scope: S
                 for i in range(3, len(krange)):
                     varname = kernel_map.params[-i - 1]
 
-                    block_expr = 'blockIdx.z'
+                    # With chiplets blockIdx.z holds the second dimension, and the dimensions folded
+                    # into the third one span a single block (see ``chiplet_count``)
+                    block_expr = '0' if chiplet_count > 1 else 'blockIdx.z'
                     if not has_tbmap or has_dtbmap:
                         block_expr = '(%s * %s + threadIdx.z)' % (block_expr, _topy(block_dims[2]))
 
@@ -2401,7 +2497,10 @@ def generate_kernel_scope(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg_scope: S
                 # Optimize conditions if they are always true
                 if i >= 3 or (dsym[i] >= minel) != True:
                     condition += '%s >= %s' % (v, _topy(minel))
-                if (i >= 3 or ((dsym_end[i] < maxel) != False and ((dsym_end[i] % self._block_dims[i]) != 0) == True)
+                # The grid of the distributed dimension is padded to a multiple of the number of
+                # chiplets, so its trailing blocks always have to be masked out
+                if (i >= 3 or (chiplet_count > 1 and i == 0)
+                        or ((dsym_end[i] < maxel) != False and ((dsym_end[i] % self._block_dims[i]) != 0) == True)
                         or (self._block_dims[i] > maxel) == True):
                     if len(condition) > 0:
                         condition += ' && '
@@ -2415,6 +2514,16 @@ def generate_kernel_scope(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg_scope: S
                     if not has_dtbmap:
                         kernel_stream.write('{', cfg, state_id, scope_entry)
 
+        # The conditions above are only generated when no inner thread-block map handles them. The padding
+        # of the distributed dimension is introduced at grid level, so its condition is needed either way.
+        # It is not added to ``_kernel_grid_conditions`` because those are only re-emitted around grid
+        # synchronization, which requires nested device maps, and those disable the distribution.
+        emit_chiplet_condition = (chiplet_count > 1 and has_tbmap and not has_dtbmap
+                                  and node.map.schedule != dtypes.ScheduleType.GPU_Persistent)
+        if emit_chiplet_condition:
+            kernel_stream.write('if (%s < %s) {' % (kernel_map.params[-1], _topy(krange.max_element()[0] + 1)), cfg,
+                                state_id, scope_entry)
+
         self._dispatcher.dispatch_subgraph(sdfg,
                                            cfg,
                                            dfg_scope,
@@ -2423,6 +2532,9 @@ def generate_kernel_scope(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg_scope: S
                                            kernel_stream,
                                            skip_entry_node=True)
 
+        if emit_chiplet_condition:
+            kernel_stream.write('}', cfg, state_id, node)
+
         if (not has_tbmap and not has_dtbmap and node.map.schedule != dtypes.ScheduleType.GPU_Persistent):
             for _ in kernel_map.params:
                 kernel_stream.write('}', cfg, state_id, node)
@@ -2432,6 +2544,7 @@ def generate_kernel_scope(self, sdfg: SDFG, cfg: ControlFlowRegion, dfg_scope: S
         self._kernel_state = None
         self._in_device_code = False
         self._grid_dims = None
+        self._kernel_chiplet_count = 1
         self.dynamic_tbmap_type = None
 
     def get_next_scope_entries(self, dfg, scope_entry):

dace/config_schema.yml

@@ -431,6 +431,20 @@ required:
                             Can be set to 'max' to maximize occupancy.
                         default: '32,1,1'
 
+                    chiplet_number:
+                        type: int
+                        title: Number of chiplets of the target GPU
+                        description: >
+                            Number of chiplets (XCDs) of the GPU the code is generated for. Set to 6 on MI300A,
+                            or 8 on MI300X.
+                            The thread-blocks of every kernel are then distributed over that many chiplets: the
+                            first dimension of the grid is split into as many contiguous chunks as there are
+                            chiplets, relying on the round-robin thread-block scheduling of multi-chiplet AMD
+                            GPUs. A value of 1 disables the distribution. Individual kernels can be excluded
+                            from it by setting the `allow_chiplet_threadblock_distribution` property of their
+                            map to False.
+                        default: 1
+
                     dynamic_map_block_size:
                         type: str
                         title: Thread-Block size for GPU_ThreadBlock_Dynamic

dace/sdfg/nodes.py

@@ -1068,6 +1068,13 @@ class Map(object):
 
     gpu_force_syncthreads = Property(dtype=bool, desc="Force a call to the __syncthreads for the map", default=False)
 
+    allow_chiplet_threadblock_distribution = Property(
+        dtype=bool,
+        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.ScheduleType.GPU_Device, dtypes.ScheduleType.GPU_ThreadBlock))
+
     def __init__(self,
                  label,
                  params,

dace/transformation/dataflow/add_threadblock_map.py

@@ -197,6 +197,8 @@ def apply(self, state: SDFGState, sdfg: SDFG):
         new_kernel_entry.map.gpu_launch_bounds = kernel_map_entry.map.gpu_launch_bounds
         new_kernel_entry.map.gpu_min_warps_per_eu = kernel_map_entry.map.gpu_min_warps_per_eu
         new_kernel_entry.map.gpu_maxnreg = kernel_map_entry.map.gpu_maxnreg
+        new_kernel_entry.map.allow_chiplet_threadblock_distribution = (
+            kernel_map_entry.map.allow_chiplet_threadblock_distribution)
 
     def preprocess_default_dims(self):
         """

doc/optimization/gpu.rst

@@ -173,6 +173,48 @@ Note that if you are using CuPy, install its appropriate HIP/ROCm version.
     If you find a feature that is not supported in DaCe, please open an issue on GitHub.
 
 
+Distributing thread-blocks over chiplets
+----------------------------------------
+
+Multi-chiplet AMD GPUs, such as the MI300 series, are partitioned into chiplets (XCDs), each with its own L2 cache.
+Thread-blocks are dispatched to them in a round-robin fashion, so consecutive blocks of a kernel land on different
+chiplets and the data they share has to be replicated in every L2 cache.
+
+Setting ``compiler.cuda.chiplet_number`` to the number of chiplets of the GPU (6 on MI300A) makes the code generator
+distribute the first dimension of the grid over the chiplets instead: the grid becomes
+``(chiplets, ceil(grid_x / chiplets), grid_y)``, so that ``blockIdx.x`` is the chiplet ID and every chiplet works on a
+contiguous chunk of the first dimension, together with the whole second dimension. The default value of 1 leaves the
+grid untouched.
+
+.. code-block:: yaml
+
+    compiler:
+      cuda:
+        chiplet_number: 6
+
+The setting can also be given through the environment, without changing ``.dace.conf``:
+
+.. code-block:: bash
+
+    $ DACE_compiler_cuda_chiplet_number=6 python my_program.py
+
+The distribution moves the second dimension of the grid to ``blockIdx.z``, which makes it inapplicable to kernels
+whose third grid dimension is not 1, as well as to kernels using a persistent grid, a dynamic thread-block map, or
+nested device maps. Such kernels keep their original grid, and a warning naming the kernel is issued.
+
+The setting describes the GPU, and applies to every kernel that can use it. A single kernel can be excluded from the
+distribution by setting the :attr:`~dace.sdfg.nodes.Map.allow_chiplet_threadblock_distribution` attribute of its map
+to ``False``, in which case its grid is left untouched and no warning is issued for it:
+
+.. code-block:: python
+
+    for node, _ in sdfg.all_nodes_recursive():
+        if isinstance(node, dace.nodes.MapEntry) and node.map.label == 'my_kernel':
+            node.map.allow_chiplet_threadblock_distribution = False
+
+Conversely, when a map allows the distribution but ``compiler.cuda.chiplet_number`` is left at 1 while targeting HIP,
+a warning naming the kernel points out that the number of chiplets has not been configured.
+
 Optimizing GPU SDFGs
 --------------------