@@ -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