Fix LaunchConfig.grid unit conversion when cluster is set

cuda_core/cuda/core/experimental/_launch_config.py

@@ -35,10 +35,20 @@ def _lazy_init():
 class LaunchConfig:
     """Customizable launch options.
 
+    Note
+    ----
+    When cluster is specified, the grid parameter represents the number of
+    clusters (not blocks). The hierarchy is: grid (clusters) -> cluster (blocks) -> 
+    block (threads). Each dimension in grid specifies clusters, each dimension in 
+    cluster specifies blocks per cluster, and each dimension in block specifies 
+    threads per block.
+
     Attributes
     ----------
     grid : Union[tuple, int]
-        Collection of threads that will execute a kernel function.
+        Collection of threads that will execute a kernel function. When cluster
+        is not specified, this represents the number of blocks. When cluster is
+        specified, this represents the number of clusters.
     cluster : Union[tuple, int]
         Group of blocks (Thread Block Cluster) that will execute on the same
         GPU Processing Cluster (GPC). Blocks within a cluster have access to
@@ -80,6 +90,13 @@ def __post_init__(self):
                     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)
+            # When cluster is set, grid should represent the number of clusters, not blocks.
+            # Convert grid dimensions from cluster units to block units by multiplying by cluster dimensions.
+            self.grid = (
+                self.grid[0] * self.cluster[0],
+                self.grid[1] * self.cluster[1], 
+                self.grid[2] * self.cluster[2]
+            )
         if self.shmem_size is None:
             self.shmem_size = 0
         if self.cooperative_launch and not Device().properties.cooperative_launch:

cuda_core/docs/source/release/0.3.3-notes.rst

@@ -0,0 +1,39 @@
+.. SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+.. SPDX-License-Identifier: Apache-2.0
+
+.. currentmodule:: cuda.core.experimental
+
+``cuda.core`` 0.3.3 Release Notes
+=================================
+
+Released on TBD
+
+
+Highlights
+----------
+
+- Fix for :class:`LaunchConfig` grid parameter unit conversion when thread block clusters are used.
+
+
+Breaking Changes
+----------------
+
+- **LaunchConfig grid parameter interpretation**: When :attr:`LaunchConfig.cluster` is specified, the :attr:`LaunchConfig.grid` parameter now correctly represents the number of clusters instead of blocks. Previously, the grid parameter was incorrectly interpreted as blocks, causing a mismatch with the expected C++ behavior. This change ensures that ``LaunchConfig(grid=4, cluster=2, block=32)`` correctly produces 4 clusters × 2 blocks/cluster = 8 total blocks, matching the C++ equivalent ``cudax::make_hierarchy(cudax::grid_dims(4), cudax::cluster_dims(2), cudax::block_dims(32))``.
+
+
+New features
+------------
+
+None.
+
+
+New examples
+------------
+
+None.
+
+
+Fixes and enhancements
+----------------------
+
+- Fix :class:`LaunchConfig` grid unit conversion when cluster is set (addresses issue #867).

cuda_core/tests/test_launcher.py

@@ -59,6 +59,60 @@ def test_launch_config_shmem_size():
     assert config.shmem_size == 0
 
 
+def test_launch_config_cluster_grid_conversion():
+    """Test that LaunchConfig correctly converts grid from cluster units to block units."""
+    # Mock the _lazy_init and device capability checks to avoid hardware requirements
+    import cuda.core.experimental._launch_config as lc_module
+
+    # Store original values
+    original_inited = lc_module._inited
+    original_use_ex = getattr(lc_module, '_use_ex', None)
+
+    try:
+        # Mock initialization state
+        lc_module._inited = True
+        lc_module._use_ex = True
+
+        # Mock Device class to avoid hardware checks
+        from unittest.mock import patch, MagicMock
+
+        mock_device = MagicMock()
+        mock_device.compute_capability = (9, 0)  # H100
+        mock_device.properties.cooperative_launch = True
+
+        with patch('cuda.core.experimental._launch_config.Device', return_value=mock_device):
+            # Test case 1: 1D - Issue #867 example
+            config = LaunchConfig(grid=4, cluster=2, block=32)
+            assert config.grid == (8, 1, 1), f"Expected (8, 1, 1), got {config.grid}"
+            assert config.cluster == (2, 1, 1), f"Expected (2, 1, 1), got {config.cluster}"
+            assert config.block == (32, 1, 1), f"Expected (32, 1, 1), got {config.block}"
+
+            # Test case 2: 2D grid and cluster
+            config = LaunchConfig(grid=(2, 3), cluster=(2, 2), block=32)
+            assert config.grid == (4, 6, 1), f"Expected (4, 6, 1), got {config.grid}"
+            assert config.cluster == (2, 2, 1), f"Expected (2, 2, 1), got {config.cluster}"
+
+            # Test case 3: 3D full specification
+            config = LaunchConfig(grid=(2, 2, 2), cluster=(3, 3, 3), block=(8, 8, 8))
+            assert config.grid == (6, 6, 6), f"Expected (6, 6, 6), got {config.grid}"
+            assert config.cluster == (3, 3, 3), f"Expected (3, 3, 3), got {config.cluster}"
+
+            # Test case 4: Identity case
+            config = LaunchConfig(grid=1, cluster=1, block=32)
+            assert config.grid == (1, 1, 1), f"Expected (1, 1, 1), got {config.grid}"
+
+            # Test case 5: No cluster (should not convert grid)
+            config = LaunchConfig(grid=4, block=32)
+            assert config.grid == (4, 1, 1), f"Expected (4, 1, 1), got {config.grid}"
+            assert config.cluster is None
+
+    finally:
+        # Restore original state
+        lc_module._inited = original_inited
+        if original_use_ex is not None:
+            lc_module._use_ex = original_use_ex
+
+
 def test_launch_invalid_values(init_cuda):
     code = 'extern "C" __global__ void my_kernel() {}'
     program = Program(code, "c++")