Examine fatbin input that contains LTOIR

isVoid · isVoid · commit bac11f68a9be · 2024-10-30T22:40:35.000-07:00
diff --git a/numba_cuda/numba/cuda/cudadrv/driver.py b/numba_cuda/numba/cuda/cudadrv/driver.py
@@ -21,6 +21,9 @@
 import traceback
 import asyncio
 import pathlib
+import subprocess
+import tempfile
+import re
 from itertools import product
 from abc import ABCMeta, abstractmethod
 from ctypes import (c_int, byref, c_size_t, c_char, c_char_p, addressof,
@@ -36,7 +39,7 @@
 from .drvapi import API_PROTOTYPES
 from .drvapi import cu_occupancy_b2d_size, cu_stream_callback_pyobj, cu_uuid
 from .mappings import FILE_EXTENSION_MAP
-from .linkable_code import LinkableCode, LTOIR
+from .linkable_code import LinkableCode, LTOIR, Fatbin, Object
 from numba.cuda.cudadrv import enums, drvapi, nvrtc
 
 USE_NV_BINDING = config.CUDA_USE_NVIDIA_BINDING
@@ -2710,12 +2713,25 @@ def add_file_guess_ext(self, path_or_code, ignore_nonlto=False):
                         "Don't know how to link file with extension "
                         f"{ext}"
                     )
-                if ignore_nonlto and kind != FILE_EXTENSION_MAP["ltoir"]:
-                    warnings.warn(
-                        f"Not adding {path_or_code} as it is not optimizable "
-                        "at link time, and `ignore_nonlto == True`."
-                    )
-                    return
+
+                if ignore_nonlto:
+                    warn_and_return = False
+                    if kind in (
+                        FILE_EXTENSION_MAP["fatbin"], FILE_EXTENSION_MAP["o"]
+                    ):
+                        entry_types = inspect_obj_content(path_or_code)
+                        if "nvvm" not in entry_types:
+                            warn_and_return = True
+                    elif kind != FILE_EXTENSION_MAP["ltoir"]:
+                        warn_and_return = True
+
+                    if warn_and_return:
+                        warnings.warn(
+                            f"Not adding {path_or_code} as it is not "
+                            "optimizable at link time, and `ignore_nonlto == "
+                            "True`."
+                        )
+                        return
 
                 self.add_file(path_or_code, kind)
             return
@@ -2729,12 +2745,24 @@ def add_file_guess_ext(self, path_or_code, ignore_nonlto=False):
             if path_or_code.kind == "cu":
                 self.add_cu(path_or_code.data, path_or_code.name)
             else:
-                if ignore_nonlto and not isinstance(path_or_code.kind, LTOIR):
-                    warnings.warn(
-                        f"Not adding {path_or_code.name} as it is not "
-                        "optimizable at link time, and `ignore_nonlto == True`."
-                    )
-                    return
+                if ignore_nonlto:
+                    warn_and_return = False
+                    if isinstance(path_or_code, (Fatbin, Object)):
+                        with tempfile.NamedTemporaryFile("w") as fp:
+                            fp.write(path_or_code.data)
+                            entry_types = inspect_obj_content(fp.name)
+                        if "nvvm" not in entry_types:
+                            warn_and_return = True
+                    elif not isinstance(path_or_code, LTOIR):
+                        warn_and_return = True
+
+                    if warn_and_return:
+                        warnings.warn(
+                            f"Not adding {path_or_code.name} as it is not "
+                            "optimizable at link time, and `ignore_nonlto == "
+                            "True`."
+                        )
+                        return
 
                 self.add_data(
                     path_or_code.data, path_or_code.kind, path_or_code.name
@@ -3411,3 +3439,16 @@ def get_version():
     Return the driver version as a tuple of (major, minor)
     """
     return driver.get_version()
+
+
+def inspect_obj_content(objpath: str):
+    code_types :set[str] = set()
+
+    out = subprocess.run(["cuobjdump", objpath], capture_output=True)
+    objtable = out.stdout.decode()
+    entry_pattern = r"Fatbin (.*) code"
+    for line in objtable.split("\n"):
+        if match := re.match(entry_pattern, line):
+            code_types.add(match.group(1))
+
+    return code_types
diff --git a/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py b/numba_cuda/numba/cuda/tests/cudadrv/test_nvjitlink.py
@@ -27,6 +27,9 @@
     test_device_functions_fatbin = os.path.join(
         TEST_BIN_DIR, "test_device_functions.fatbin"
     )
+    test_device_functions_fatbin_multi = os.path.join(
+        TEST_BIN_DIR, "test_device_functions_multi.fatbin"
+    )
     test_device_functions_o = os.path.join(
         TEST_BIN_DIR, "test_device_functions.o"
     )
@@ -178,6 +181,7 @@ def test_nvjitlink_jit_with_linkable_code_lto_dump_assembly(self):
         files = [
             test_device_functions_cu,
             test_device_functions_ltoir,
+            test_device_functions_fatbin_multi
         ]
 
         config.DUMP_ASSEMBLY = True
diff --git a/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile b/numba_cuda/numba/cuda/tests/test_binary_generation/Makefile
@@ -14,9 +14,14 @@ endif
 # Gencode flags suitable for most tests
 GENCODE := -gencode arch=compute_$(GPU_CC),code=sm_$(GPU_CC)
 
+MULTI_GENCODE := -gencode arch=compute_$(GPU_CC),code=[sm_$(GPU_CC),lto_$(GPU_CC)]
+
 # Fatbin tests need to generate code for an additional compute capability
 FATBIN_GENCODE := $(GENCODE) -gencode arch=compute_$(ALT_CC),code=sm_$(ALT_CC)
 
+# Fatbin that contains both LTO, SASS for multiple architectures
+MULTI_FATBIN_GENCODE := $(MULTI_GENCODE) -gencode arch=compute_$(ALT_CC),code=[sm_$(ALT_CC),lto_$(ALT_CC)]
+
 # LTO-IR tests need to generate for the LTO "architecture" instead
 LTOIR_GENCODE := -gencode arch=lto_$(GPU_CC),code=lto_$(GPU_CC)
 
@@ -30,6 +35,7 @@ PTX_FLAGS := $(GENCODE) -ptx
 OBJECT_FLAGS := $(GENCODE) -dc
 LIBRARY_FLAGS := $(GENCODE) -lib
 FATBIN_FLAGS := $(FATBIN_GENCODE) --fatbin
+MULTI_FATBIN_FLAGS := $(MULTI_FATBIN_GENCODE) --fatbin
 LTOIR_FLAGS := $(LTOIR_GENCODE) -dc
 
 OUTPUT_DIR := ./
@@ -41,6 +47,7 @@ all:
 	nvcc $(NVCC_FLAGS) $(CUBIN_FLAGS) -o $(OUTPUT_DIR)/undefined_extern.cubin undefined_extern.cu
 	nvcc $(NVCC_FLAGS) $(CUBIN_FLAGS) -o $(OUTPUT_DIR)/test_device_functions.cubin test_device_functions.cu
 	nvcc $(NVCC_FLAGS) $(FATBIN_FLAGS) -o $(OUTPUT_DIR)/test_device_functions.fatbin test_device_functions.cu
+	nvcc $(NVCC_FLAGS) $(MULTI_FATBIN_FLAGS) -o $(OUTPUT_DIR)/test_device_functions_multi.fatbin test_device_functions.cu
 	nvcc $(NVCC_FLAGS) $(PTX_FLAGS) -o $(OUTPUT_DIR)/test_device_functions.ptx test_device_functions.cu
 	nvcc $(NVCC_FLAGS) $(OBJECT_FLAGS) -o $(OUTPUT_DIR)/test_device_functions.o test_device_functions.cu
 	nvcc $(NVCC_FLAGS) $(LIBRARY_FLAGS) -o $(OUTPUT_DIR)/test_device_functions.a test_device_functions.cu

Original file line number	Diff line number	Diff line change
`@@ -27,6 +27,9 @@`
`27`	`27`	`test_device_functions_fatbin = os.path.join(`
`28`	`28`	`TEST_BIN_DIR, "test_device_functions.fatbin"`
`29`	`29`	`)`
	`30`	`+ test_device_functions_fatbin_multi = os.path.join(`
	`31`	`+ TEST_BIN_DIR, "test_device_functions_multi.fatbin"`
	`32`	`+ )`
`30`	`33`	`test_device_functions_o = os.path.join(`
`31`	`34`	`TEST_BIN_DIR, "test_device_functions.o"`
`32`	`35`	`)`
`@@ -178,6 +181,7 @@ def test_nvjitlink_jit_with_linkable_code_lto_dump_assembly(self):`
`178`	`181`	`files = [`
`179`	`182`	`test_device_functions_cu,`
`180`	`183`	`test_device_functions_ltoir,`
	`184`	`+ test_device_functions_fatbin_multi`
`181`	`185`	`]`
`182`	`186`
`183`	`187`	`config.DUMP_ASSEMBLY = True`