Skip to content

Commit 94139a0

Browse files
Andy-Jostclaude
andauthored
test(cuda.core): compile the scalar-argument kernel once per module in test_launcher.py (#3042)
`test_launch_scalar_argument` and `test_launch_numpy_scalar_subclass_fallback` compiled the same `write_scalar<T>` template once per parametrized item. On a local-toolkit row that is 41 header-bearing NVRTC compiles, each parsing `cuda_fp16.h` and `cuda/std/complex`, at about 1,031 ms each on L4 against 22 ms for the same template without the two CCCL headers (CI step timestamps of two full matrix runs). In a single measurement run on one L40S GPU (not a CI runner) the same items took about 1.9 s each against about 55 ms for the header-free twin. Co-authored-by: Claude Fable 5.1 <[email protected]>
1 parent e39c1df commit 94139a0

1 file changed

Lines changed: 38 additions & 48 deletions

File tree

‎cuda_core/tests/test_launcher.py‎

Lines changed: 38 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -559,11 +559,41 @@ def test_launch_invalid_values(init_cuda):
559559
(complex, "cuda::std::complex<double>", 5 - 7j),
560560
)
561561

562+
# Kernel source shared by test_launch_scalar_argument and
563+
# test_launch_numpy_scalar_subclass_fallback. The write_scalar_module fixture
564+
# below compiles it once per module.
565+
_WRITE_SCALAR_SRC = r"""
566+
template <typename T>
567+
__global__ void write_scalar(T* arr, T val) {
568+
arr[0] = val;
569+
}
570+
"""
571+
_CCCL_HEADERS = r"""
572+
#include <cuda_fp16.h>
573+
#include <cuda/std/complex>
574+
"""
575+
576+
577+
@pytest.fixture(scope="module")
578+
def write_scalar_module():
579+
# One compile for the header-bearing scalar items; each item still launches
580+
# its own instantiation, so the per-type argument paths stay covered.
581+
dev = Device(0)
582+
arch = "".join(f"{i}" for i in dev.compute_capability)
583+
code = _WRITE_SCALAR_SRC
584+
if helpers.CCCL_INCLUDE_PATHS is not None:
585+
code = _CCCL_HEADERS + code
586+
names = {f"write_scalar<{cpp_type}>" for _, cpp_type, _ in PARAMS}
587+
names |= {f"write_scalar<{cpp_type}>" for _, _, cpp_type, _ in _NUMPY_SUBCLASS_FALLBACK_PARAMS}
588+
pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}", include_path=helpers.CCCL_INCLUDE_PATHS)
589+
prog = Program(code, code_type="c++", options=pro_opts)
590+
return prog.compile("cubin", name_expressions=tuple(sorted(names)))
591+
562592

563593
@pytest.mark.parametrize("python_type, cpp_type, init_value", PARAMS)
564594
@requires_module(np, "2.2.5", reason="need numpy 2.2.5+ (numpy GH #28632)")
565-
def test_launch_scalar_argument(python_type, cpp_type, init_value):
566-
dev = Device()
595+
def test_launch_scalar_argument(write_scalar_module, python_type, cpp_type, init_value):
596+
dev = Device(0)
567597
dev.set_current()
568598

569599
# Prepare pinned host array
@@ -575,29 +605,8 @@ def test_launch_scalar_argument(python_type, cpp_type, init_value):
575605
# Prepare scalar argument in Python
576606
scalar = python_type(init_value)
577607

578-
# CUDA kernel templated on type T
579-
code = r"""
580-
template <typename T>
581-
__global__ void write_scalar(T* arr, T val) {
582-
arr[0] = val;
583-
}
584-
"""
585-
586-
# Compile and force instantiation for this type
587-
arch = "".join(f"{i}" for i in dev.compute_capability)
588-
if helpers.CCCL_INCLUDE_PATHS is not None:
589-
code = (
590-
r"""
591-
#include <cuda_fp16.h>
592-
#include <cuda/std/complex>
593-
"""
594-
+ code
595-
)
596-
pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}", include_path=helpers.CCCL_INCLUDE_PATHS)
597-
prog = Program(code, code_type="c++", options=pro_opts)
598-
ker_name = f"write_scalar<{cpp_type}>"
599-
mod = prog.compile("cubin", name_expressions=(ker_name,))
600-
ker = mod.get_kernel(ker_name)
608+
# The kernel is compiled once per module; fetch this type's instantiation.
609+
ker = write_scalar_module.get_kernel(f"write_scalar<{cpp_type}>")
601610

602611
# Launch with 1 thread
603612
stream = dev.default_stream
@@ -930,7 +939,7 @@ class MyComplex(complex):
930939
ids=_NUMPY_SUBCLASS_FALLBACK_IDS,
931940
)
932941
@pytest.mark.agent_authored(model="gpt-5.6-sol")
933-
def test_launch_numpy_scalar_subclass_fallback(base_type, np_dtype, cpp_type, raw_value):
942+
def test_launch_numpy_scalar_subclass_fallback(write_scalar_module, base_type, np_dtype, cpp_type, raw_value):
934943
"""Subclassed numpy scalars take prepare_numpy_arg's isinstance fallback and reach the kernel (readback)."""
935944

936945
class Subclassed(base_type):
@@ -939,35 +948,16 @@ class Subclassed(base_type):
939948
scalar = Subclassed(raw_value)
940949
expected = np_dtype(raw_value)
941950

942-
dev = Device()
951+
dev = Device(0)
943952
dev.set_current()
944953

945954
mr = LegacyPinnedMemoryResource()
946955
b = mr.allocate(np.dtype(np_dtype).itemsize)
947956
arr = np.from_dlpack(b).view(np_dtype)
948957
arr[:] = 0
949958

950-
code = r"""
951-
template <typename T>
952-
__global__ void write_scalar(T* arr, T val) {
953-
arr[0] = val;
954-
}
955-
"""
956-
if helpers.CCCL_INCLUDE_PATHS is not None:
957-
code = (
958-
r"""
959-
#include <cuda_fp16.h>
960-
#include <cuda/std/complex>
961-
"""
962-
+ code
963-
)
964-
965-
arch = "".join(f"{i}" for i in dev.compute_capability)
966-
pro_opts = ProgramOptions(std="c++17", arch=f"sm_{arch}", include_path=helpers.CCCL_INCLUDE_PATHS)
967-
prog = Program(code, code_type="c++", options=pro_opts)
968-
ker_name = f"write_scalar<{cpp_type}>"
969-
mod = prog.compile("cubin", name_expressions=(ker_name,))
970-
ker = mod.get_kernel(ker_name)
959+
# The kernel is compiled once per module; fetch this type's instantiation.
960+
ker = write_scalar_module.get_kernel(f"write_scalar<{cpp_type}>")
971961

972962
stream = dev.default_stream
973963
config = LaunchConfig(grid=1, block=1)

0 commit comments

Comments
 (0)