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/backend/platforms/cuda.py

@@ -113,11 +113,11 @@ class Blockwise4DMapping(ThreadToIndexMapping):
     """Blockwise index mapping for up to 4D iteration spaces, where the outer three dimensions
     are mapped to block indices."""
 
-    _indices_in_loop_order = [  # slowest to fastest
-        BLOCK_IDX[2],
-        BLOCK_IDX[1],
-        BLOCK_IDX[0],
+    _indices_fastest_first = [  # slowest to fastest
         THREAD_IDX[0],
+        BLOCK_IDX[0],
+        BLOCK_IDX[1],
+        BLOCK_IDX[2]
     ]
 
     def __call__(self, ispace: IterationSpace) -> dict[PsSymbol, PsExpression]:
@@ -141,7 +141,7 @@ class Blockwise4DMapping(ThreadToIndexMapping):
         dimensions = ispace.dimensions_in_loop_order()
         idx_map: dict[PsSymbol, PsExpression] = dict()
 
-        for dim, tid in zip(dimensions, self._indices_in_loop_order):
+        for dim, tid in zip(dimensions[::-1], self._indices_fastest_first):
             idx_map[dim.counter] = dim.start + dim.step * PsCast(
                 deconstify(dim.counter.get_dtype()), tid
             )
@@ -152,7 +152,7 @@ class Blockwise4DMapping(ThreadToIndexMapping):
         self, ispace: SparseIterationSpace
     ) -> dict[PsSymbol, PsExpression]:
         sparse_ctr = PsExpression.make(ispace.sparse_counter)
-        thread_idx = self._indices_in_loop_order[-1]
+        thread_idx = self._indices_fastest_first[0]
         idx_map: dict[PsSymbol, PsExpression] = {
             ispace.sparse_counter: PsCast(
                 deconstify(sparse_ctr.get_dtype()), thread_idx