diff --git a/cuda_bindings/tests/test_nvvm.py b/cuda_bindings/tests/test_nvvm.py index 29e5d89a91..060a5268be 100644 --- a/cuda_bindings/tests/test_nvvm.py +++ b/cuda_bindings/tests/test_nvvm.py @@ -4,21 +4,12 @@ import binascii import re -import textwrap from contextlib import contextmanager import pytest from cuda.bindings import nvvm -MINIMAL_NVVMIR_FIXTURE_PARAMS = ["txt", "bitcode_static"] -try: - import llvmlite.binding as llvmlite_binding # Optional test dependency. -except ImportError: - llvmlite_binding = None -else: - MINIMAL_NVVMIR_FIXTURE_PARAMS.append("bitcode_dynamic") - -MINIMAL_NVVMIR_TXT = b"""\ +MINIMAL_NVVMIR_TXT_TEMPLATE = b"""\ target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-i128:128:128-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n16:32:64" target triple = "nvptx64-nvidia-cuda" @@ -130,43 +121,24 @@ "6e673e0000000000", } -MINIMAL_NVVMIR_CACHE = {} - -@pytest.fixture(params=MINIMAL_NVVMIR_FIXTURE_PARAMS) +@pytest.fixture(params=("txt", "bitcode_static")) def minimal_nvvmir(request): - for pass_counter in range(2): - nvvmir = MINIMAL_NVVMIR_CACHE.get(request.param, -1) - if nvvmir != -1: - if nvvmir is None: - pytest.skip(f"UNAVAILABLE: {request.param}") - return nvvmir - if pass_counter: - raise AssertionError("This code path is meant to be unreachable.") - # Build cache entries, then try again (above). - major, minor, debug_major, debug_minor = nvvm.ir_version() - txt = MINIMAL_NVVMIR_TXT % (major, debug_major) - if llvmlite_binding is None: - bitcode_dynamic = None - else: - bitcode_dynamic = llvmlite_binding.parse_assembly(txt.decode()).as_bitcode() - bitcode_static = MINIMAL_NVVMIR_BITCODE_STATIC.get((major, debug_major)) - if bitcode_static is not None: - bitcode_static = binascii.unhexlify(bitcode_static) - MINIMAL_NVVMIR_CACHE["txt"] = txt - MINIMAL_NVVMIR_CACHE["bitcode_dynamic"] = bitcode_dynamic - MINIMAL_NVVMIR_CACHE["bitcode_static"] = bitcode_static - if bitcode_static is None: - if bitcode_dynamic is None: - raise RuntimeError("Please `pip install llvmlite` to generate `bitcode_static` (see PR #443)") - bitcode_hex = binascii.hexlify(bitcode_dynamic).decode("ascii") - print("\n\nMINIMAL_NVVMIR_BITCODE_STATIC = { # PLEASE ADD TO test_nvvm.py") - print(f" ({major}, {debug_major}): # (major, debug_major)") - lines = textwrap.wrap(bitcode_hex, width=80) - for line in lines[:-1]: - print(f' "{line}"') - print(f' "{lines[-1]}",') - print("}\n", flush=True) + major, minor, debug_major, debug_minor = nvvm.ir_version() + + if request.param == "txt": + return MINIMAL_NVVMIR_TXT_TEMPLATE % (major, debug_major) + + bitcode_static_binascii = MINIMAL_NVVMIR_BITCODE_STATIC.get((major, debug_major)) + if bitcode_static_binascii: + return binascii.unhexlify(bitcode_static_binascii) + raise RuntimeError( + "Static bitcode for NVVM IR version " + f"{major}.{debug_major} is not available in this test.\n" + "Maintainers: Please run the helper script to generate it and add the " + "output to the MINIMAL_NVVMIR_BITCODE_STATIC dict:\n" + " ../../toolshed/build_static_bitcode_input.py" + ) @pytest.fixture(params=[nvvm.compile_program, nvvm.verify_program]) diff --git a/toolshed/build_static_bitcode_input.py b/toolshed/build_static_bitcode_input.py new file mode 100755 index 0000000000..95c3a610f6 --- /dev/null +++ b/toolshed/build_static_bitcode_input.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE + +""" +Helper to produce static bitcode input for test_nvvm.py. + +Usage: + python toolshed/build_static_bitcode_input.py + +It will print a ready-to-paste MINIMAL_NVVMIR_BITCODE_STATIC entry for the +current NVVM IR version detected at runtime. +""" + +import binascii +import os +import sys +import textwrap + +import llvmlite.binding # HINT: pip install llvmlite +from cuda.bindings import nvvm + + +def get_minimal_nvvmir_txt_template(): + cuda_bindings_tests_dir = os.path.normpath("cuda_bindings/tests") + assert os.path.isdir(cuda_bindings_tests_dir), ( + "Please run this helper script from the cuda-python top-level directory." + ) + sys.path.insert(0, os.path.abspath(cuda_bindings_tests_dir)) + import test_nvvm + + return test_nvvm.MINIMAL_NVVMIR_TXT_TEMPLATE + + +def main(): + major, _minor, debug_major, _debug_minor = nvvm.ir_version() + txt = get_minimal_nvvmir_txt_template() % (major, debug_major) + bitcode_dynamic = llvmlite.binding.parse_assembly(txt.decode()).as_bitcode() + bitcode_hex = binascii.hexlify(bitcode_dynamic).decode("ascii") + print("\n\nMINIMAL_NVVMIR_BITCODE_STATIC = { # PLEASE ADD TO test_nvvm.py") + print(f" ({major}, {debug_major}): # (major, debug_major)") + lines = textwrap.wrap(bitcode_hex, width=80) + for line in lines[:-1]: + print(f' "{line}"') + print(f' "{lines[-1]}",') + print("}\n", flush=True) + print() + + +if __name__ == "__main__": + assert len(sys.argv) == 1, "This helper script does not take any arguments." + main()