cuda.core: add tests for ObjectCode.from_library (NVIDIA#2281)

* initial version of testing ObjectCode.from_library

* use pytest parametrize and merge boilerplate code

* rephrase to be accurate

Author: lijinf2

.github/workflows/build-wheel.yml

@@ -469,7 +469,10 @@ jobs:
         uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a  # v7.0.1
         with:
           name: ${{ env.CUDA_CORE_ARTIFACT_NAME }}-test-binaries
-          path: ${{ env.CUDA_CORE_TEST_BINARIES_DIR }}/*.o
+          path: |
+            ${{ env.CUDA_CORE_TEST_BINARIES_DIR }}/*.o
+            ${{ env.CUDA_CORE_TEST_BINARIES_DIR }}/*.a
+            ${{ env.CUDA_CORE_TEST_BINARIES_DIR }}/*.lib
           if-no-files-found: error
 
       - name: Download cuda.bindings build artifacts from the prior branch

.gitignore

@@ -21,6 +21,8 @@ cache_nvrtc
 
 # cuda.core test object fixtures built locally / downloaded as CI artifacts
 cuda_core/tests/test_binaries/*.o
+cuda_core/tests/test_binaries/*.a
+cuda_core/tests/test_binaries/*.lib
 
 # CUDA Python specific (auto-generated)
 cuda_bindings/cuda/bindings/_bindings/cyruntime.pxd

cuda_core/tests/test_binaries/build_test_binaries.sh

@@ -19,4 +19,10 @@ NVCC="${NVCC:-nvcc}"
 "${NVCC}" -dc "${NVCC_EXTRA_FLAGS[@]}" -arch=all-major \
     -o "${SCRIPTPATH}/saxpy.o" "${SCRIPTPATH}/saxpy.cu"
 
-ls -lah "${SCRIPTPATH}/saxpy.o"
+if [[ "${OS:-}" == "Windows_NT" ]]; then
+    nvcc -lib -o "${SCRIPTPATH}/saxpy.lib" "${SCRIPTPATH}/saxpy.o"
+    ls -lah "${SCRIPTPATH}/saxpy.o" "${SCRIPTPATH}/saxpy.lib"
+else
+    nvcc -lib -o "${SCRIPTPATH}/saxpy.a" "${SCRIPTPATH}/saxpy.o"
+    ls -lah "${SCRIPTPATH}/saxpy.o" "${SCRIPTPATH}/saxpy.a"
+fi

cuda_core/tests/test_module.py

@@ -177,22 +177,26 @@ 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.
+def _read_saxpy_rdc(kind: str) -> bytes:
+    """Read a pre-built saxpy RDC object or library.
 
-    In CI: produced by build stage into a test wheel file.
+    In CI: produced by the build stage.
     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.
+    saxpy.cu, remove stale RDC files (i.e. saxpy.o, saxpy.a, or saxpy.lib) to force a rebuild.
     """
     binaries_dir = Path(__file__).parent / "test_binaries"
-    obj_path = binaries_dir / "saxpy.o"
+    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 obj_path.is_file():
+    if not rdc_path.is_file():
         nvcc_path = find_nvidia_binary_utility("nvcc")
         if nvcc_path is None:
             pytest.skip(
-                f"saxpy.o not found at {obj_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."
             )
         env = os.environ.copy()
@@ -202,8 +206,7 @@ def get_saxpy_object():
             check=True,
             env=env,
         )
-
-    return obj_path.read_bytes()
+    return rdc_path.read_bytes()
 
 
 def test_get_kernel(init_cuda):
@@ -364,26 +367,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"'):
+@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_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)
+@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 == "object"
+    assert mod_obj.code_type == kind
 
 
-def test_object_code_load_object_with_linker(get_saxpy_object, 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"""
@@ -397,7 +421,7 @@ def test_object_code_load_object_with_linker(get_saxpy_object, init_cuda):
     ).compile("cubin")
     linked = Linker(
         kernel_code,
-        ObjectCode.from_object(get_saxpy_object),
+        from_fn(_read_saxpy_rdc(kind)),
         options=LinkerOptions(arch=arch),
     ).link("cubin")
     kernel = linked.get_kernel("linked_kernel")