GPU Indexing Schemes and Launch Configurations

Closes #119 (closed)

This MR contributes

• Lambdas (#119 (closed)), which act as wrappers around one-line backend expressions that can be evaluatated
• More flexible GPU indexing schemes via modular work item mappings
• Launch Configurations as encapsulated, user-configurable objects

Lambdas

Add codegen.functions.Lambda, a simple wrapper around a backend expression tree that can be exported to the user and evaluated by the runtime system

Indexing Schemes in Config

• Make the two indexing schemes Linear3D and Blockwise4D available through the config (they correspond to the pystencils 1.3.x block and line indexing options)
• Permit users to set a default block size for Linear3D

GPU Thread Index Mapping

• Remove GpuThreadsRange; the thread range is no longer being computed by the CUDA and SYCL platforms
• Extend the Cuda platform to receive its mapping from thread indices to iteration space points via a callback object
• Implement thread indexing for Linear3D and Blockwise4D

Launch Configurations

• Encapsulate GPU launch configurations into objects of the base class codegen.gpu_indexing.GpuLaunchConfiguration, which is evaluated by the JIT in the context of the kernel launch.
• Add subclasses for fully manual and fully automatic lauch configurations, as well as a launch config with dynamic block size for the Linear3D indexing scheme. Enable the user to set launch config parameters on the compiled kernel object.
• Update JIT to adhere to the new launch config interface.

Codegen Driver Updates

• Introduce the GpuIndexing class which acts as a factory for launch configurations and thread indexing objects depending on the user-provided configuration

src/pystencils/codegen/gpu_indexing.py

@@ -19,7 +19,6 @@ from ..backend.ast.expressions import PsExpression
 
 
 dim3 = tuple[int, int, int]
-_Dim3Params = tuple[Parameter, Parameter, Parameter]
 _Dim3Lambda = tuple[Lambda, Lambda, Lambda]
 
 
@@ -73,15 +72,9 @@ class AutomaticLaunchConfiguration(GpuLaunchConfiguration):
 
     @property
     def parameters(self) -> frozenset[Parameter]:
-        """Parameters of this launch configuration"""
         return self._params
 
     def evaluate(self, **kwargs) -> tuple[dim3, dim3]:
-        """Compute block and grid size for a kernel launch.
-
-        Args:
-            kwargs: Valuation providing a value for each parameter listed in `parameters`
-        """
         block_size = tuple(int(bs(**kwargs)) for bs in self._block_size)
         grid_size = tuple(int(gs(**kwargs)) for gs in self._grid_size)
         return cast(dim3, block_size), cast(dim3, grid_size)
@@ -136,27 +129,38 @@ class ManualLaunchConfiguration(GpuLaunchConfiguration):
 
 
 class DynamicBlockSizeLaunchConfiguration(GpuLaunchConfiguration):
-    """GPU launch configuration that permits the user to set a block size dynamically."""
+    """GPU launch configuration that permits the user to set a block size and dynamically computes the grid size.
+    
+    The actual launch grid size is computed from the user-defined ``user_block_size`` and the number of work items
+    in the kernel's iteration space as follows.
+    For each dimension :math:`c \\in \\{ x, y, z \\}`,
+
+    - if ``user_block_size.c > num_work_items.c``, ``block_size = num_work_items.c`` and ``grid_size.c = 1``;
+    - otherwise, ``block_size.c = user_block_size.c`` and ``grid_size.c = ceil(num_work_items.c / block_size.c)``.
+    """
 
     def __init__(
         self,
-        block_size_expr: _Dim3Lambda,
-        grid_size_expr: _Dim3Lambda,
-        block_size_params: _Dim3Params,
+        num_work_items: _Dim3Lambda,
         default_block_size: dim3 | None = None,
     ) -> None:
-        self._block_size_expr = block_size_expr
-        self._grid_size_expr = grid_size_expr
+        self._num_work_items = num_work_items
 
-        self._block_size_params = block_size_params
         self._block_size: dim3 | None = default_block_size
 
         self._params: frozenset[Parameter] = frozenset().union(
-            *(lb.parameters for lb in chain(block_size_expr, grid_size_expr))
-        ) - set(self._block_size_params)
+            *(wit.parameters for wit in num_work_items)
+        )
+
+    @property
+    def num_work_items(self) -> _Dim3Lambda:
+        """Lambda expressions that compute the number of work items in each iteration space
+        dimension from kernel parameters."""
+        return self._num_work_items
 
     @property
     def block_size(self) -> dim3 | None:
+        """The desired GPU block size."""
         return self._block_size
 
     @block_size.setter
@@ -172,16 +176,23 @@ class DynamicBlockSizeLaunchConfiguration(GpuLaunchConfiguration):
         if self._block_size is None:
             raise AttributeError("No GPU block size was specified by the user!")
 
-        kwargs.update(
-            {
-                param.name: value
-                for param, value in zip(self._block_size_params, self._block_size)
-            }
+        from ..utils import div_ceil
+
+        num_work_items = cast(
+            dim3, tuple(int(wit(**kwargs)) for wit in self._num_work_items)
+        )
+        reduced_block_size = cast(
+            dim3,
+            tuple(min(wit, bs) for wit, bs in zip(num_work_items, self._block_size)),
+        )
+        grid_size = cast(
+            dim3,
+            tuple(
+                div_ceil(wit, bs) for wit, bs in zip(num_work_items, reduced_block_size)
+            ),
         )
 
-        block_size = tuple(int(bs(**kwargs)) for bs in self._block_size_expr)
-        grid_size = tuple(int(gs(**kwargs)) for gs in self._grid_size_expr)
-        return cast(dim3, block_size), cast(dim3, grid_size)
+        return reduced_block_size, grid_size
 
     def jit_cache_key(self) -> Any:
         return self._block_size
@@ -226,50 +237,17 @@ class GpuIndexing(ABC):
     def _get_linear3d_config_factory(
         self,
     ) -> Callable[[], DynamicBlockSizeLaunchConfiguration]:
-        work_items = self._get_work_items()
-        rank = len(work_items)
-
-        from ..backend.constants import PsConstant
-        from ..backend.ast.expressions import PsExpression, PsIntDiv
-
-        block_size_symbols = [
-            self._ctx.get_new_symbol(f"gpuBlockSize_{c}", self._ctx.index_dtype) for c in range(rank)
-        ]
-
-        block_size = [
-            Lambda.from_expression(self._ctx, self._factory.parse_index(bs_symb))
-            for bs_symb in block_size_symbols
-        ] + [
-            Lambda.from_expression(self._ctx, self._factory.parse_index(1))
-            for _ in range(3 - rank)
-        ]
-
-        def div_ceil(a: PsExpression, b: PsExpression):
-            return self._factory.parse_index(
-                PsIntDiv(a + b - PsExpression.make(PsConstant(1)), b)
-            )
-
-        grid_size = [
-            Lambda.from_expression(
-                self._ctx, div_ceil(witems, PsExpression.make(bsize))
-            )
-            for witems, bsize in zip(work_items, block_size_symbols)
-        ] + [
-            Lambda.from_expression(self._ctx, self._factory.parse_index(1))
-            for _ in range(3 - rank)
-        ]
-
-        from .driver import _symbol_to_param
+        work_items_expr = self._get_work_items()
+        rank = len(work_items_expr)
 
-        block_size_params = tuple(
-            _symbol_to_param(self._ctx, s) for s in block_size_symbols
+        num_work_items = cast(
+            _Dim3Lambda,
+            tuple(Lambda.from_expression(self._ctx, wit) for wit in work_items_expr),
         )
 
         def factory():
             return DynamicBlockSizeLaunchConfiguration(
-                cast(_Dim3Lambda, tuple(block_size)),
-                cast(_Dim3Lambda, tuple(grid_size)),
-                cast(_Dim3Params, block_size_params),
+                num_work_items,
                 self._get_default_block_size(rank),
             )