test(cuda.core): share PDL overlap protocol via helper runner

Extract kernels and the same-stream / graph-capture overlap check into
run_pdl_overlap_check so launcher and GraphBuilder tests stay in sync.

Author: lijinf2

cuda_core/tests/graph/test_graph_builder.py

@@ -7,15 +7,15 @@
 import time
 import weakref
 
-import helpers
 import numpy as np
 import pytest
 from conftest import skipif_need_cuda_headers
 from cuda_python_test_helpers.marks import requires_module
 from helpers.graph_kernels import compile_common_kernels, compile_conditional_kernels
 from helpers.misc import try_create_condition
+from helpers.pdl_kernels import run_pdl_overlap_check
 
-from cuda.core import Device, LaunchConfig, LegacyPinnedMemoryResource, Program, ProgramOptions, launch
+from cuda.core import Device, LaunchConfig, LegacyPinnedMemoryResource, launch
 from cuda.core.graph import GraphBuilder, GraphDefinition
 from cuda.core.graph._graph_builder import (
     _capture_callback_with_tail_failure_for_testing,
@@ -783,85 +783,4 @@ def test_pdl_primary_secondary_overlap_graph_capture(init_cuda):
     but launches are captured into a CUDA graph (see CUDA Programming Guide,
     Programmatic Dependent Launch). Overlap is opportunistic → miss is xfail.
     """
-    dev = Device()
-    if dev.compute_capability < (9, 0):
-        pytest.skip("Programmatic Dependent Launch requires compute capability >= 9.0")
-    stream = dev.create_stream(options={"nonblocking": True})
-
-    # clock64 budgets are in GPU cycles; keep the post-trigger window long enough
-    # for the secondary to boot, but short enough for a unit test.
-    code = r"""
-    #include <cuda_device_runtime_api.h>
-
-    extern "C" __global__ void primary_kernel(int* secondary_started, int* overlapped) {
-        cudaTriggerProgrammaticLaunchCompletion();
-
-        const long long deadline = clock64() + 100000000LL;  // ~50ms @ ~2GHz
-        if (threadIdx.x == 0 && blockIdx.x == 0) {
-            while (clock64() < deadline) {
-                if (atomicAdd(secondary_started, 0) != 0) {
-                    atomicExch(overlapped, 1);
-                    return;
-                }
-                __nanosleep(1000);
-            }
-        }
-    }
-
-    extern "C" __global__ void secondary_kernel(int* secondary_started) {
-        if (threadIdx.x == 0 && blockIdx.x == 0) {
-            atomicExch(secondary_started, 1);
-        }
-    }
-    """
-
-    arch = "".join(f"{i}" for i in dev.compute_capability)
-    pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}", include_path=helpers.CUDA_INCLUDE_PATH)
-    prog = Program(code, code_type="c++", options=pro_opts)
-    mod = prog.compile("cubin")
-    primary = mod.get_kernel("primary_kernel")
-    secondary = mod.get_kernel("secondary_kernel")
-
-    mr = LegacyPinnedMemoryResource()
-    secondary_started = np.from_dlpack(mr.allocate(4)).view(np.int32)
-    overlapped = np.from_dlpack(mr.allocate(4)).view(np.int32)
-
-    primary_cfg = LaunchConfig(grid=1, block=1)
-    secondary_cfg = LaunchConfig(grid=1, block=1, programmatic_stream_serialization=True)
-    secondary_serial_cfg = LaunchConfig(grid=1, block=1)
-
-    def _run(secondary_launch_cfg: LaunchConfig) -> int:
-        secondary_started[0] = 0
-        overlapped[0] = 0
-        gb = stream.create_graph_builder().begin_building()
-        launch(gb, primary_cfg, primary, secondary_started.ctypes.data, overlapped.ctypes.data)
-        launch(gb, secondary_launch_cfg, secondary, secondary_started.ctypes.data)
-        graph = gb.end_building().complete()
-        try:
-            graph.launch(stream)
-            stream.sync()
-        finally:
-            graph.close()
-            gb.close()
-        return int(overlapped[0])
-
-    # Without the PDL attribute, same-stream kernels stay serialized even via graph.
-    assert _run(secondary_serial_cfg) == 0, "Expected no overlap when programmatic_stream_serialization is False"
-
-    # PDL overlap is opportunistic; retry a few times on a quiet GPU.
-    saw_overlap = False
-    for _ in range(5):
-        if _run(secondary_cfg) == 1:
-            saw_overlap = True
-            break
-
-    if not saw_overlap:
-        pytest.xfail(
-            "PDL (Programmatic Dependent Launch) graph-capture overlap was not observed. "
-            "If this keeps xfailing in CI, manually re-check on a quiet Hopper+ GPU."
-        )
-
-    print(
-        f"PDL graph-capture overlap verified on {dev.name} compute capability {dev.compute_capability}",
-        flush=True,
-    )
+    run_pdl_overlap_check(Device(), via_graph=True)

cuda_core/tests/helpers/pdl_kernels.py

@@ -0,0 +1,122 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Shared helpers for Programmatic Dependent Launch overlap tests."""
+
+import helpers
+import numpy as np
+import pytest
+
+from cuda.core import LaunchConfig, LegacyPinnedMemoryResource, Program, ProgramOptions, launch
+
+
+def compile_pdl_overlap_kernels(device):
+    """Compile primary/secondary kernels used to detect PDL same-stream overlap.
+
+    The primary triggers programmatic launch completion then spins briefly looking
+    for a flag written by the secondary. Seeing that flag proves both grids were
+    resident at once. clock64 budgets are in GPU cycles: long enough for the
+    secondary to boot, short enough for a unit test.
+
+    Returns:
+        (primary_kernel, secondary_kernel)
+    """
+    code = r"""
+    #include <cuda_device_runtime_api.h>
+
+    extern "C" __global__ void primary_kernel(int* secondary_started, int* overlapped) {
+        cudaTriggerProgrammaticLaunchCompletion();
+
+        const long long deadline = clock64() + 100000000LL;  // ~50ms @ ~2GHz
+        if (threadIdx.x == 0 && blockIdx.x == 0) {
+            while (clock64() < deadline) {
+                if (atomicAdd(secondary_started, 0) != 0) {
+                    atomicExch(overlapped, 1);
+                    return;
+                }
+                __nanosleep(1000);
+            }
+        }
+    }
+
+    extern "C" __global__ void secondary_kernel(int* secondary_started) {
+        if (threadIdx.x == 0 && blockIdx.x == 0) {
+            atomicExch(secondary_started, 1);
+        }
+    }
+    """
+    arch = "".join(f"{i}" for i in device.compute_capability)
+    pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}", include_path=helpers.CUDA_INCLUDE_PATH)
+    prog = Program(code, code_type="c++", options=pro_opts)
+    mod = prog.compile("cubin")
+    return mod.get_kernel("primary_kernel"), mod.get_kernel("secondary_kernel")
+
+
+def run_pdl_overlap_check(device, *, via_graph: bool = False):
+    """Run the shared primary/secondary PDL overlap protocol.
+
+    Asserts no overlap without ``programmatic_stream_serialization``, then retries
+    a few times with it enabled. Overlap is opportunistic → miss is xfail.
+
+    Args:
+        device: Current CUDA device (compute capability >= 9.0 required).
+        via_graph: If True, capture launches into a CUDA graph and launch the
+            graph; otherwise launch kernels directly on the stream.
+    """
+    if device.compute_capability < (9, 0):
+        pytest.skip("Programmatic Dependent Launch requires compute capability >= 9.0")
+
+    stream = device.create_stream(options={"nonblocking": True})
+    primary, secondary = compile_pdl_overlap_kernels(device)
+
+    mr = LegacyPinnedMemoryResource()
+    secondary_started = np.from_dlpack(mr.allocate(4)).view(np.int32)
+    overlapped = np.from_dlpack(mr.allocate(4)).view(np.int32)
+
+    primary_cfg = LaunchConfig(grid=1, block=1)
+    secondary_cfg = LaunchConfig(grid=1, block=1, programmatic_stream_serialization=True)
+    secondary_serial_cfg = LaunchConfig(grid=1, block=1)
+
+    def _run(secondary_launch_cfg: LaunchConfig) -> int:
+        secondary_started[0] = 0
+        overlapped[0] = 0
+        if via_graph:
+            gb = stream.create_graph_builder().begin_building()
+            launch(gb, primary_cfg, primary, secondary_started.ctypes.data, overlapped.ctypes.data)
+            launch(gb, secondary_launch_cfg, secondary, secondary_started.ctypes.data)
+            graph = gb.end_building().complete()
+            try:
+                graph.launch(stream)
+                stream.sync()
+            finally:
+                graph.close()
+                gb.close()
+        else:
+            launch(stream, primary_cfg, primary, secondary_started.ctypes.data, overlapped.ctypes.data)
+            launch(stream, secondary_launch_cfg, secondary, secondary_started.ctypes.data)
+            stream.sync()
+        return int(overlapped[0])
+
+    path = "graph-capture" if via_graph else "same-stream"
+    assert _run(secondary_serial_cfg) == 0, (
+        f"Expected no overlap when programmatic_stream_serialization is False ({path})"
+    )
+
+    saw_overlap = False
+    for _ in range(5):
+        if _run(secondary_cfg) == 1:
+            saw_overlap = True
+            break
+
+    if not saw_overlap:
+        # Overlap is never guaranteed by the driver, so a miss is reported as an
+        # expected failure rather than turning a busy GPU into a red CI run.
+        pytest.xfail(
+            f"PDL (Programmatic Dependent Launch) {path} overlap was not observed. "
+            "If this keeps xfailing in CI, manually re-check on a quiet Hopper+ GPU."
+        )
+
+    print(
+        f"PDL {path} overlap verified on {device.name} compute capability {device.compute_capability}",
+        flush=True,
+    )

cuda_core/tests/test_launcher.py

@@ -6,6 +6,7 @@
 import helpers
 from cuda_python_test_helpers.marks import requires_module
 from helpers.misc import StreamWrapper
+from helpers.pdl_kernels import run_pdl_overlap_check
 
 try:
     import cupy as cp
@@ -203,7 +204,8 @@ def test_to_native_launch_config_pdl():
 
 
 @skipif_need_cuda_headers
-def test_pdl_primary_secondary_overlap_same_stream():
+@requires_module(np, "2.2.5", reason="need numpy 2.2.5+ (numpy GH #28632)")
+def test_pdl_primary_secondary_overlap_same_stream(init_cuda):
     """Primary + secondary PDL launch on one stream can overlap on Hopper+.
 
     Secondary is launched with ``programmatic_stream_serialization=True``. After
@@ -214,84 +216,7 @@ def test_pdl_primary_secondary_overlap_same_stream():
     Note concurrency is opportunistic, so a missing overlap execution is reported as
     an expected failure.
     """
-    dev = Device()
-    if dev.compute_capability < (9, 0):
-        pytest.skip("Programmatic Dependent Launch requires compute capability >= 9.0")
-    dev.set_current()
-    stream = dev.create_stream(options={"nonblocking": True})
-
-    # clock64 budgets are in GPU cycles; keep the post-trigger window long enough
-    # for the secondary to boot, but short enough for a unit test.
-    code = r"""
-    #include <cuda_device_runtime_api.h>
-
-    extern "C" __global__ void primary_kernel(int* secondary_started, int* overlapped) {
-        cudaTriggerProgrammaticLaunchCompletion();
-
-        const long long deadline = clock64() + 100000000LL;  // ~50ms @ ~2GHz
-        if (threadIdx.x == 0 && blockIdx.x == 0) {
-            while (clock64() < deadline) {
-                if (atomicAdd(secondary_started, 0) != 0) {
-                    atomicExch(overlapped, 1);
-                    return;
-                }
-                __nanosleep(1000);
-            }
-        }
-    }
-
-    extern "C" __global__ void secondary_kernel(int* secondary_started) {
-        if (threadIdx.x == 0 && blockIdx.x == 0) {
-            atomicExch(secondary_started, 1);
-        }
-    }
-    """
-
-    arch = "".join(f"{i}" for i in dev.compute_capability)
-    pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}", include_path=helpers.CUDA_INCLUDE_PATH)
-    prog = Program(code, code_type="c++", options=pro_opts)
-    mod = prog.compile("cubin")
-    primary = mod.get_kernel("primary_kernel")
-    secondary = mod.get_kernel("secondary_kernel")
-
-    mr = LegacyPinnedMemoryResource()
-    secondary_started = np.from_dlpack(mr.allocate(4)).view(np.int32)
-    overlapped = np.from_dlpack(mr.allocate(4)).view(np.int32)
-
-    primary_cfg = LaunchConfig(grid=1, block=1)
-    secondary_cfg = LaunchConfig(grid=1, block=1, programmatic_stream_serialization=True)
-    secondary_serial_cfg = LaunchConfig(grid=1, block=1)
-
-    def _run(secondary_launch_cfg: LaunchConfig) -> int:
-        secondary_started[0] = 0
-        overlapped[0] = 0
-        launch(stream, primary_cfg, primary, secondary_started.ctypes.data, overlapped.ctypes.data)
-        launch(stream, secondary_launch_cfg, secondary, secondary_started.ctypes.data)
-        stream.sync()
-        return int(overlapped[0])
-
-    # Without the PDL attribute, same-stream kernels stay serialized.
-    assert _run(secondary_serial_cfg) == 0, "Expected no overlap when programmatic_stream_serialization is False"
-
-    # PDL overlap is opportunistic; retry a few times on a quiet GPU.
-    saw_overlap = False
-    for _ in range(5):
-        if _run(secondary_cfg) == 1:
-            saw_overlap = True
-            break
-
-    if not saw_overlap:
-        # Overlap is never guaranteed by the driver, so a miss is reported as an
-        # expected failure rather than turning a busy GPU into a red CI run.
-        pytest.xfail(
-            "PDL (Programmatic Dependent Launch) overlap was not observed. "
-            "If this keeps xfailing in CI, manually re-check on a quiet Hopper+ GPU."
-        )
-
-    print(
-        f"PDL (Programmatic Dependent Launch) overlap verified on {dev.name} compute capability {dev.compute_capability}",
-        flush=True,
-    )
+    run_pdl_overlap_check(Device(), via_graph=False)
 
 
 def test_launch_config_cluster_accepts_hopper_cc(monkeypatch):