Add distribution of thread blocks to chiplets in old codegen

dace/codegen/common.py

@@ -185,6 +185,44 @@ def get_gpu_runtime() -> gpu_runtime.GPURuntime:
     return gpu_runtime.GPURuntime(backend, libpath)
 
 
+@lru_cache()
+def get_gpu_chiplet_count() -> Optional[int]:
+    """
+    Returns the number of chiplets (XCDs) of the GPU of this machine, or None if it cannot be determined.
+
+    The result is cached: the device does not change within a process, and ``amdsmi`` has to be
+    initialized and shut down around the query, which must happen exactly once. The warning on
+    failure is emitted from here, so that it is emitted only once per process as well.
+
+    :note: ``amdsmi`` is not a DaCe dependency, it ships with ROCm. Importing it loads
+           ``libamd_smi.so`` through ``ctypes``, which fails on machines without ROCm and not
+           necessarily with an ``ImportError``, hence the broad exception handling.
+    :note: ``amdsmi`` ignores ``ROCR_VISIBLE_DEVICES`` and ``HIP_VISIBLE_DEVICES`` and enumerates all
+           physical GPUs, so the first processor handle is not necessarily the device that HIP uses.
+           This is only accurate on nodes where all GPUs are of the same model.
+    """
+    try:
+        import amdsmi
+
+        amdsmi.amdsmi_init()
+        try:
+            processor_handles = amdsmi.amdsmi_get_processor_handles()
+            if not processor_handles:
+                raise RuntimeError('`amdsmi` did not report any GPU.')
+            chiplets = int(amdsmi.amdsmi_get_gpu_xcd_counter(processor_handles[0]))
+            if chiplets < 1:
+                raise RuntimeError(f'`amdsmi` reported an invalid number of chiplets ({chiplets}).')
+            return chiplets
+        finally:
+            amdsmi.amdsmi_shut_down()
+    except Exception as e:
+        warnings.warn(f'Could not determine the number of GPU chiplets through `amdsmi`: {e}. The '
+                      'distribution of thread-blocks over chiplets is disabled. Set the '
+                      '`compiler.cuda.chiplet_number` configuration entry to the number of chiplets of '
+                      'the GPU (6 on MI300A) to enable it.')
+        return None
+
+
 def platform_library_name(libname: str) -> str:
     """ Get the filename of a library.
 

dace/codegen/targets/cuda.py

@@ -89,6 +89,11 @@ 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
+        # Number of thread-blocks of the first grid dimension every chiplet owns (see ``chiplet_count``)
+        self._kernel_chiplet_chunk = 1
         self._kernel_map = None
         self._kernel_state = None
         self._kernel_grid_conditions: List[str] = []
@@ -2071,6 +2076,83 @@ 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 detected_chiplet_count(self) -> int:
+        """
+        Returns the number of chiplets of the GPU of this machine, or 1 if it cannot be determined.
+
+        Called when the ``compiler.cuda.chiplet_number`` configuration entry is left at its default of 0,
+        which means that the number of chiplets has not been configured by the user. A detected number is
+        written back to that entry, so that the rest of the process sees the number the code is generated
+        for. Only the HIP backend is queried, chiplets being a feature of AMD GPUs, and a failure to
+        determine the number leaves the entry alone and disables the distribution.
+        """
+        if self.backend != 'hip':
+            return 1
+
+        chiplets = common.get_gpu_chiplet_count()
+        if chiplets is None:
+            return 1
+
+        Config.set('compiler', 'cuda', 'chiplet_number', value=chiplets)
+        return chiplets
+
+    def chiplet_count(self, kernelmap_entry: nodes.MapEntry, 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 pads the first grid dimension to a multiple of the number of chiplets and permutes
+        it within itself, leaving the second and third grid dimensions untouched. It therefore applies to
+        kernels of any dimensionality, and only steps aside for kernels whose block indices are not
+        generated by ``generate_kernel_scope`` alone. Kernels whose map has the
+        ``allow_chiplet_threadblock_distribution`` property set to False are left alone as well.
+
+        The number of chiplets comes from the ``compiler.cuda.chiplet_number`` configuration entry, whose
+        default of 0 means that it is detected automatically (see ``detected_chiplet_count``).
+
+        :param kernelmap_entry: Entry node of the kernel map.
+        :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 < 0:
+            raise ValueError(f'Invalid number of chiplets ({chiplets}) configured. Modify the '
+                             '`compiler.cuda.chiplet_number` configuration entry to a positive number, or to 0 '
+                             'to detect the number of chiplets of the GPU automatically.')
+
+        # A kernel that opts out of the distribution is left alone without any diagnostics, and without
+        # querying the GPU for the number of its chiplets
+        if not kernelmap_entry.map.allow_chiplet_threadblock_distribution:
+            return 1
+
+        if chiplets == 0:
+            chiplets = self.detected_chiplet_count()
+
+        if chiplets == 1:
+            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'
+
+        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 +2363,25 @@ 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, is_persistent, has_dtbmap, extra_grid_dims)
+        if self._kernel_chiplet_count > 1:
+            # The hardware dispatches thread-block ``f`` to chiplet ``f % chiplets`` in a round-robin
+            # fashion, and the flattened block index is ``blockIdx.x + blockIdx.y * gridDim.x +
+            # blockIdx.z * gridDim.x * gridDim.y``. Padding ``gridDim.x`` to a multiple of the number of
+            # chiplets makes the two trailing terms vanish modulo that number, so the chiplet of a block
+            # is ``blockIdx.x % chiplets`` regardless of ``blockIdx.y`` and ``blockIdx.z``, which the
+            # distribution therefore leaves untouched. Each chiplet receives a contiguous chunk of
+            # ``ceil(grid_size[0] / chiplets)`` blocks of the first dimension (see the block index
+            # computed in ``generate_kernel_scope``), together with the whole of the other dimensions.
+            self._kernel_chiplet_chunk = int_ceil(grid_size[0], self._kernel_chiplet_count)
+            original_grid_size = grid_size
+            grid_size = [self._kernel_chiplet_chunk * self._kernel_chiplet_count] + grid_size[1:]
+            if Config.get_bool('debugprint'):
+                print(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 +2437,22 @@ 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
+        chiplet_chunk = self._kernel_chiplet_chunk
         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 and i == 0:
+                    # Contiguous partitioning: the chiplet of a block is ``blockIdx.x % chiplets`` and its
+                    # slot within that chiplet is ``blockIdx.x / chiplets``, so chiplet k owns the blocks
+                    # [k * chunk .. (k + 1) * chunk - 1] of this dimension (see ``get_kernel_dimensions``)
+                    block_expr = '((blockIdx.x %% %d) * %s + blockIdx.x / %d)' % (chiplet_count, _topy(chiplet_chunk),
+                                                                                  chiplet_count)
+                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))
 
@@ -2401,7 +2511,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 +2528,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 +2546,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 +2558,8 @@ 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._kernel_chiplet_chunk = 1
         self.dynamic_tbmap_type = None
 
     def get_next_scope_entries(self, dfg, scope_entry):

dace/config_schema.yml

@@ -431,6 +431,22 @@ 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, 6 on MI300A or 8 on
+                            MI300X. The thread-blocks of every kernel are distributed over that many chiplets:
+                            the first dimension of the grid is padded to a multiple of that number and split
+                            into as many contiguous chunks as there are chiplets, relying on the round-robin
+                            thread-block scheduling of multi-chiplet AMD GPUs. The other two grid dimensions
+                            are left untouched. The default of 0 detects the number of chiplets of the GPU of
+                            this machine when targeting HIP, and disables the distribution when it cannot be
+                            determined. A value of 1 disables the distribution as well. Individual kernels can
+                            be excluded from it by setting the `allow_chiplet_threadblock_distribution`
+                            property of their map to False.
+                        default: 0
+
                     dynamic_map_block_size:
                         type: str
                         title: Thread-Block size for GPU_ThreadBlock_Dynamic

dace/sdfg/nodes.py

@@ -1049,6 +1049,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,59 @@ 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.
+
+``compiler.cuda.chiplet_number``, 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: that dimension is padded to a multiple of the
+number of chiplets, so the grid becomes ``(ceil(grid_x / chiplets) * chiplets, grid_y, grid_z)``. Since the flattened
+block index is ``blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y``, a ``gridDim.x`` that is a
+multiple of the number of chiplets reduces the chiplet a block is dispatched to to ``blockIdx.x % chiplets``, whatever
+its other two indices are. The blocks of the first dimension are then permuted to
+``(blockIdx.x % chiplets) * ceil(grid_x / chiplets) + blockIdx.x / chiplets``, so that every chiplet works on a
+contiguous chunk of that dimension, together with the whole of the other two. The blocks that the padding adds beyond
+the range of the map are masked out. Setting the entry to 1 leaves the grid untouched.
+
+The entry is left at 0 by default, which makes the code generator determine the number of chiplets of the GPU of this
+machine when targeting HIP, through the ``amdsmi`` module that ships with ROCm, so that the distribution applies to a
+multi-chiplet AMD GPU without any configuration. The number is that of the first GPU ``amdsmi`` reports, which is
+accurate on nodes whose GPUs are all of the same model. If it cannot be determined, when generating code on a machine
+without ROCm for instance, a warning is issued once and the grids are left untouched.
+
+The detected number can be overridden, which is what generating code for a GPU other than the one of this machine
+calls for:
+
+.. 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 reshapes the first dimension of the grid only, and leaves the second and third ones on their own
+grid dimension, so it applies to kernels of any dimensionality. It is inapplicable to kernels using a persistent grid,
+a dynamic thread-block map, or nested device maps, whose block indices are not derived from the grid alone. 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
+
 Optimizing GPU SDFGs
 --------------------