use pytest parametrize and merge boilerplate code

Author: lijinf2

cuda_core/tests/test_module.py

@@ -177,57 +177,33 @@ def get_saxpy_fatbin(init_cuda):
     return bytes(fatbin), sym_map
 
 
-@pytest.fixture(scope="module")
-def get_saxpy_object():
-    """Read the pre-built saxpy.o.
-
-    In CI: produced by build stage into a test wheel file.
-    In local dev: auto-built on demand if nvcc is available; if you edit
-    saxpy.cu, remove the stale saxpy.o to force a rebuild.
-    """
-    binaries_dir = Path(__file__).parent / "test_binaries"
-    obj_path = binaries_dir / "saxpy.o"
-
-    if not obj_path.is_file():
-        if find_nvidia_binary_utility("nvcc") is None:
-            pytest.skip(
-                f"saxpy.o not found at {obj_path} and nvcc is unavailable. "
-                "In CI this is downloaded from the build stage."
-            )
-        subprocess.run(  # noqa: S603
-            ["bash", str(binaries_dir / "build_test_binaries.sh")],  # noqa: S607
-            check=True,
-            env=os.environ,
-        )
+def _read_saxpy_rdc(kind: str) -> bytes:
+    """Read a pre-built saxpy RDC object or library.
 
-    return obj_path.read_bytes()
-
-
-@pytest.fixture(scope="module")
-def get_saxpy_library():
-    """Read the pre-built saxpy.a (Linux) or saxpy.lib (Windows).
-
-    In CI: produced by build stage alongside saxpy.o.
+    In CI: produced by the build stage.
     In local dev: auto-built on demand if nvcc is available; if you edit
-    saxpy.cu, remove stale saxpy.o / saxpy.a to force a rebuild.
+    saxpy.cu, remove stale RDC files to force a rebuild.
     """
     binaries_dir = Path(__file__).parent / "test_binaries"
-    lib_name = "saxpy.lib" if os.name == "nt" else "saxpy.a"
-    lib_path = binaries_dir / lib_name
+    if kind == "object":
+        rdc_path = binaries_dir / "saxpy.o"
+    elif kind == "library":
+        rdc_path = binaries_dir / ("saxpy.lib" if os.name == "nt" else "saxpy.a")
+    else:
+        raise ValueError(f"unknown saxpy RDC kind: {kind!r}")
 
-    if not lib_path.is_file():
+    if not rdc_path.is_file():
         if find_nvidia_binary_utility("nvcc") is None:
             pytest.skip(
-                f"{lib_name} not found at {lib_path} and nvcc is unavailable. "
+                f"{rdc_path.name} not found at {rdc_path} and nvcc is unavailable. "
                 "In CI this is downloaded from the build stage."
             )
         subprocess.run(  # noqa: S603
             ["bash", str(binaries_dir / "build_test_binaries.sh")],  # noqa: S607
             check=True,
             env=os.environ,
         )
-
-    return lib_path.read_bytes()
+    return rdc_path.read_bytes()
 
 
 def test_get_kernel(init_cuda):
@@ -388,86 +364,47 @@ def test_object_code_load_fatbin_from_file(get_saxpy_fatbin, tmp_path, convert_p
     mod_obj.get_kernel("saxpy<double>")  # force loading
 
 
-def test_object_code_load_object(get_saxpy_object):
-    obj = get_saxpy_object
-    assert isinstance(obj, bytes)
-    mod_obj = ObjectCode.from_object(obj)
-    assert mod_obj.code == obj
-    assert mod_obj.code_type == "object"
-    with pytest.raises(RuntimeError, match=r'Unsupported code type "object"'):
-        mod_obj.get_kernel("saxpy<float>")
-
-
-def test_object_code_load_object_from_file(get_saxpy_object, tmp_path):
-    obj_file = tmp_path / "test.o"
-    obj_file.write_bytes(get_saxpy_object)
-    arg = str(obj_file)
-    mod_obj = ObjectCode.from_object(arg)
-    assert mod_obj.code == arg
-    assert mod_obj.code_type == "object"
-
-
-def test_object_code_load_object_with_linker(get_saxpy_object, init_cuda):
-    arch = f"sm_{init_cuda.arch}"
-    kernel_code = Program(
-        r"""
-        extern __device__ float saxpy_step(float a, float x, float y);
-        extern "C" __global__ void linked_kernel(float a, float x, float y, float* out) {
-            if (threadIdx.x == 0 && blockIdx.x == 0) *out = saxpy_step(a, x, y);
-        }
-        """,
-        "c++",
-        ProgramOptions(relocatable_device_code=True, arch=arch),
-    ).compile("cubin")
-    linked = Linker(
-        kernel_code,
-        ObjectCode.from_object(get_saxpy_object),
-        options=LinkerOptions(arch=arch),
-    ).link("cubin")
-    kernel = linked.get_kernel("linked_kernel")
-
-    stream = init_cuda.create_stream()
-    host_buf = cuda.core.LegacyPinnedMemoryResource().allocate(4)
-    result = np.from_dlpack(host_buf).view(np.float32)
-    result[:] = 0.0
-    dev_buf = init_cuda.memory_resource.allocate(4, stream=init_cuda.default_stream)
-
-    cuda.core.launch(
-        stream,
-        cuda.core.LaunchConfig(grid=1, block=1),
-        kernel,
-        np.float32(2.0),
-        np.float32(3.0),
-        np.float32(4.0),
-        dev_buf,
-    )
-    dev_buf.copy_to(host_buf, stream=stream)
-    stream.sync()
-
-    assert result[0] == 10.0
-
-
-def test_object_code_load_library(get_saxpy_library):
-    lib = get_saxpy_library
-    assert isinstance(lib, bytes)
-    mod_obj = ObjectCode.from_library(lib)
-    assert mod_obj.code == lib
-    assert mod_obj.code_type == "library"
-    with pytest.raises(RuntimeError, match=r'Unsupported code type "library"'):
+@pytest.mark.parametrize(
+    ("kind", "from_fn"),
+    [
+        ("object", ObjectCode.from_object),
+        ("library", ObjectCode.from_library),
+    ],
+)
+def test_object_code_load_rdc(kind, from_fn):
+    data = _read_saxpy_rdc(kind)
+    assert isinstance(data, bytes)
+    mod_obj = from_fn(data)
+    assert mod_obj.code == data
+    assert mod_obj.code_type == kind
+    with pytest.raises(RuntimeError, match=rf'Unsupported code type "{kind}"'):
         mod_obj.get_kernel("saxpy<float>")
 
 
-def test_object_code_load_library_from_file(get_saxpy_library, tmp_path):
-    lib_ext = ".lib" if os.name == "nt" else ".a"
-    lib_file = tmp_path / f"test{lib_ext}"
-    lib_file.write_bytes(get_saxpy_library)
-    arg = str(lib_file)
-    mod_obj = ObjectCode.from_library(arg)
+@pytest.mark.parametrize(
+    ("kind", "from_fn", "suffix"),
+    [
+        ("object", ObjectCode.from_object, ".o"),
+        ("library", ObjectCode.from_library, ".lib" if os.name == "nt" else ".a"),
+    ],
+)
+def test_object_code_load_rdc_from_file(kind, from_fn, suffix, tmp_path):
+    rdc_file = tmp_path / f"test{suffix}"
+    rdc_file.write_bytes(_read_saxpy_rdc(kind))
+    arg = str(rdc_file)
+    mod_obj = from_fn(arg)
     assert mod_obj.code == arg
-    assert mod_obj.code_type == "library"
+    assert mod_obj.code_type == kind
 
 
-def test_object_code_load_library_with_linker(get_saxpy_library, init_cuda):
+@pytest.mark.parametrize(
+    ("kind", "from_fn"),
+    [
+        ("object", ObjectCode.from_object),
+        ("library", ObjectCode.from_library),
+    ],
+)
+def test_object_code_load_rdc_with_linker(kind, from_fn, init_cuda):
     arch = f"sm_{init_cuda.arch}"
     kernel_code = Program(
         r"""
@@ -481,7 +418,7 @@ def test_object_code_load_library_with_linker(get_saxpy_library, init_cuda):
     ).compile("cubin")
     linked = Linker(
         kernel_code,
-        ObjectCode.from_library(get_saxpy_library),
+        from_fn(_read_saxpy_rdc(kind)),
         options=LinkerOptions(arch=arch),
     ).link("cubin")
     kernel = linked.get_kernel("linked_kernel")