cythonize _launch_config.py (#1221)

* cythonize _launch_config

* pre-commit

* moving back

* removing comment

* reverting list code back to original

* pre-commit

* removing unused import

* Update cuda_core/cuda/core/experimental/_launch_config.pyx

Co-authored-by: Copilot <[email protected]>

---------

Co-authored-by: Copilot <[email protected]>

Author: rparolin

cuda_core/cuda/core/experimental/_launch_config.pxd

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+cdef bint _inited
+cdef bint _use_ex
+
+cdef void _lazy_init() except *
+
+cdef class LaunchConfig:
+    """Customizable launch options."""
+    cdef public tuple grid
+    cdef public tuple cluster
+    cdef public tuple block
+    cdef public int shmem_size
+    cdef public bint cooperative_launch
+
+cpdef object _to_native_launch_config(LaunchConfig config)

cuda_core/cuda/core/experimental/_launch_config.py renamed to cuda_core/cuda/core/experimental/_launch_config.pyx

@@ -2,9 +2,6 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from dataclasses import dataclass
-from typing import Union
-
 from cuda.core.experimental._device import Device
 from cuda.core.experimental._utils.cuda_utils import (
     CUDAError,
@@ -15,24 +12,27 @@
 )
 
 # TODO: revisit this treatment for py313t builds
-_inited = False
+cdef bint _inited = False
+cdef bint _use_ex = False
 
 
-def _lazy_init():
-    global _inited
+cdef void _lazy_init() except *:
+    """Initialize module-level globals for driver version checks."""
+    global _inited, _use_ex
     if _inited:
         return
 
-    global _use_ex
+    cdef tuple _py_major_minor
+    cdef int _driver_ver
+
     # binding availability depends on cuda-python version
     _py_major_minor = get_binding_version()
     _driver_ver = handle_return(driver.cuDriverGetVersion())
     _use_ex = (_driver_ver >= 11080) and (_py_major_minor >= (11, 8))
     _inited = True
 
 
-@dataclass
-class LaunchConfig:
+cdef class LaunchConfig:
     """Customizable launch options.
 
     Note
@@ -65,21 +65,36 @@ class LaunchConfig:
     """
 
     # TODO: expand LaunchConfig to include other attributes
-    grid: Union[tuple, int] = None
-    cluster: Union[tuple, int] = None
-    block: Union[tuple, int] = None
-    shmem_size: int | None = None
-    cooperative_launch: bool | None = False
-
-    def __post_init__(self):
+    # Note: attributes are declared in _launch_config.pxd
+
+    def __init__(self, grid=None, cluster=None, block=None,
+                 shmem_size=None, cooperative_launch=False):
+        """Initialize LaunchConfig with validation.
+
+        Parameters
+        ----------
+        grid : Union[tuple, int], optional
+            Grid dimensions (number of blocks or clusters if cluster is specified)
+        cluster : Union[tuple, int], optional
+            Cluster dimensions (Thread Block Cluster)
+        block : Union[tuple, int], optional
+            Block dimensions (threads per block)
+        shmem_size : int, optional
+            Dynamic shared memory size in bytes (default: 0)
+        cooperative_launch : bool, optional
+            Whether to launch as cooperative kernel (default: False)
+        """
         _lazy_init()
-        self.grid = cast_to_3_tuple("LaunchConfig.grid", self.grid)
-        self.block = cast_to_3_tuple("LaunchConfig.block", self.block)
+
+        # Convert and validate grid and block dimensions
+        self.grid = cast_to_3_tuple("LaunchConfig.grid", grid)
+        self.block = cast_to_3_tuple("LaunchConfig.block", block)
+
         # FIXME: Calling Device() strictly speaking is not quite right; we should instead
         # look up the device from stream. We probably need to defer the checks related to
         # device compute capability or attributes.
         # thread block clusters are supported starting H100
-        if self.cluster is not None:
+        if cluster is not None:
             if not _use_ex:
                 err, drvers = driver.cuDriverGetVersion()
                 drvers_fmt = f" (got driver version {drvers})" if err == driver.CUresult.CUDA_SUCCESS else ""
@@ -89,19 +104,53 @@ def __post_init__(self):
                 raise CUDAError(
                     f"thread block clusters are not supported on devices with compute capability < 9.0 (got {cc})"
                 )
-            self.cluster = cast_to_3_tuple("LaunchConfig.cluster", self.cluster)
-        if self.shmem_size is None:
+            self.cluster = cast_to_3_tuple("LaunchConfig.cluster", cluster)
+        else:
+            self.cluster = None
+
+        # Handle shmem_size default
+        if shmem_size is None:
             self.shmem_size = 0
+        else:
+            self.shmem_size = shmem_size
+
+        # Handle cooperative_launch
+        self.cooperative_launch = cooperative_launch
+
+        # Validate cooperative launch support
         if self.cooperative_launch and not Device().properties.cooperative_launch:
             raise CUDAError("cooperative kernels are not supported on this device")
 
+    def __repr__(self):
+        """Return string representation of LaunchConfig."""
+        return (f"LaunchConfig(grid={self.grid}, cluster={self.cluster}, "
+                f"block={self.block}, shmem_size={self.shmem_size}, "
+                f"cooperative_launch={self.cooperative_launch})")
+
 
-def _to_native_launch_config(config: LaunchConfig) -> driver.CUlaunchConfig:
+cpdef object _to_native_launch_config(LaunchConfig config):
+    """Convert LaunchConfig to native driver CUlaunchConfig.
+
+    Parameters
+    ----------
+    config : LaunchConfig
+        High-level launch configuration
+
+    Returns
+    -------
+    driver.CUlaunchConfig
+        Native CUDA driver launch configuration
+    """
     _lazy_init()
-    drv_cfg = driver.CUlaunchConfig()
+
+    cdef object drv_cfg = driver.CUlaunchConfig()
+    cdef list attrs
+    cdef object attr
+    cdef object dim
+    cdef tuple grid_blocks
 
     # Handle grid dimensions and cluster configuration
-    if config.cluster:
+    if config.cluster is not None:
         # Convert grid from cluster units to block units
         grid_blocks = (
             config.grid[0] * config.cluster[0],
@@ -122,11 +171,14 @@ def _to_native_launch_config(config: LaunchConfig) -> driver.CUlaunchConfig:
 
     drv_cfg.blockDimX, drv_cfg.blockDimY, drv_cfg.blockDimZ = config.block
     drv_cfg.sharedMemBytes = config.shmem_size
+
     if config.cooperative_launch:
         attr = driver.CUlaunchAttribute()
         attr.id = driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_COOPERATIVE
         attr.value.cooperative = 1
         attrs.append(attr)
+
     drv_cfg.numAttrs = len(attrs)
     drv_cfg.attrs = attrs
+
     return drv_cfg

cuda_core/cuda/core/experimental/_launcher.pyx

@@ -9,7 +9,7 @@ from cuda.core.experimental._stream cimport _try_to_get_stream_ptr
 from typing import Union
 
 from cuda.core.experimental._kernel_arg_handler import ParamHolder
-from cuda.core.experimental._launch_config import LaunchConfig, _to_native_launch_config
+from cuda.core.experimental._launch_config cimport LaunchConfig, _to_native_launch_config
 from cuda.core.experimental._module import Kernel
 from cuda.core.experimental._stream import IsStreamT, Stream
 from cuda.core.experimental._utils.clear_error_support import assert_type

cuda_core/tests/test_launcher.py

@@ -131,7 +131,7 @@ def test_launch_invalid_values(init_cuda):
     ker = mod.get_kernel("my_kernel")
     config = LaunchConfig(grid=(1, 1, 1), block=(1, 1, 1), shmem_size=0)
 
-    with pytest.raises(ValueError):
+    with pytest.raises(TypeError):
         launch(None, ker, config)
 
     with pytest.raises(TypeError):