Fix propagation of PyTorch CUDA flags for multiple archs

Author: stotko

src/charonload/cmake/charonload-config.cmake

@@ -81,22 +81,32 @@ if(charonload_FIND_QUIETLY)
     set(CUDNN_FIND_QUIETLY 1)
 endif()
 
-# Back up CUDA_NVCC_FLAGS for later restoring
+# Back up CUDA_NVCC_FLAGS and CMAKE_CUDA_FLAGS for later restoring
 set(CHARONLOAD_CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS})
+set(CHARONLOAD_CMAKE_CUDA_FLAGS ${CMAKE_CUDA_FLAGS})
 
 find_dependency(Torch)
 
 list(POP_BACK CMAKE_MESSAGE_INDENT)
 
 if(Torch_FOUND)
     # 1. CUDA flag patching
-    if(NOT CHARONLOAD_CUDA_NVCC_FLAGS STREQUAL CUDA_NVCC_FLAGS AND TARGET torch_cuda)
+    message(STATUS "${CUDA_NVCC_FLAGS}")
+    message(STATUS "${CHARONLOAD_CUDA_NVCC_FLAGS}")
+    if((NOT CHARONLOAD_CUDA_NVCC_FLAGS STREQUAL CUDA_NVCC_FLAGS OR NOT CHARONLOAD_CUDA_NVCC_FLAGS STREQUAL CUDA_NVCC_FLAGS) AND TARGET torch_cuda)
+        # Extract modified flags
+        string(REPLACE "${CHARONLOAD_CUDA_NVCC_FLAGS}" "" CHARONLOAD_CUDA_NVCC_FLAGS_MODIFIED "${CUDA_NVCC_FLAGS}")
+        string(REPLACE ";" " " CHARONLOAD_CUDA_NVCC_FLAGS_MODIFIED "${CHARONLOAD_CUDA_NVCC_FLAGS_MODIFIED}")
+        string(STRIP "${CHARONLOAD_CUDA_NVCC_FLAGS_MODIFIED}" CHARONLOAD_CUDA_NVCC_FLAGS_MODIFIED)
+
         # Use modified CUDA_NVCC_FLAGS
-        target_compile_options(torch_cuda INTERFACE $<$<COMPILE_LANGUAGE:CUDA>:${CUDA_NVCC_FLAGS}>)
+        target_compile_options(torch_cuda INTERFACE "$<$<COMPILE_LANGUAGE:CUDA>:SHELL:${CHARONLOAD_CUDA_NVCC_FLAGS_MODIFIED}>")
+        unset(CHARONLOAD_CUDA_NVCC_FLAGS_MODIFIED)
 
-        # Restore CUDA_NVCC_FLAGS
+        # Restore CUDA_NVCC_FLAGS and CMAKE_CUDA_FLAGS
         set(CUDA_NVCC_FLAGS ${CHARONLOAD_CUDA_NVCC_FLAGS})
-        message(STATUS "Patched target \"torch_cuda\" with modified \"CUDA_NVCC_FLAGS\" settings and rolled back the variable modifications.")
+        set(CMAKE_CUDA_FLAGS ${CHARONLOAD_CMAKE_CUDA_FLAGS})
+        message(STATUS "Patched target \"torch_cuda\" with modified \"CUDA_NVCC_FLAGS\"/\"CMAKE_CUDA_FLAGS\" settings and rolled back the variable modifications.")
     endif()
 
     # 2. Python bindings library
@@ -120,6 +130,7 @@ endif()
 
 # Clean up backup variable
 unset(CHARONLOAD_CUDA_NVCC_FLAGS)
+unset(CHARONLOAD_CMAKE_CUDA_FLAGS)
 
 
 include("${CMAKE_CURRENT_LIST_DIR}/torch/cxx_standard.cmake")

tests/test_finder.py

@@ -5,6 +5,7 @@
 import importlib.util
 import io
 import multiprocessing
+import os
 import pathlib
 import platform
 import shutil
@@ -87,6 +88,43 @@ def test_torch_cuda(shared_datadir: pathlib.Path, tmp_path: pathlib.Path) -> Non
     assert torch.equal(t_output, 2 * t_input)
 
 
+def _torch_cuda_custom_archs(shared_datadir: pathlib.Path, tmp_path: pathlib.Path) -> None:
+    os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0 8.6"  # Use 2 different archs
+
+    project_directory = shared_datadir / "torch_cuda"
+    build_directory = tmp_path / "build"
+
+    charonload.module_config["test_torch_cuda_custom_archs"] = charonload.Config(
+        project_directory,
+        build_directory,
+        stubs_directory=VSCODE_STUBS_DIRECTORY,
+    )
+
+    import test_torch_cuda_custom_archs as test_torch
+
+    t_input = torch.randint(0, 10, size=(3, 3, 3), dtype=torch.float, device="cuda")
+    t_output = test_torch.two_times(t_input)
+
+    assert t_output.device == t_input.device
+    assert t_output.shape == t_input.shape
+    assert torch.equal(t_output, 2 * t_input)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_torch_cuda_custom_archs(shared_datadir: pathlib.Path, tmp_path: pathlib.Path) -> None:
+    p = multiprocessing.get_context("spawn").Process(
+        target=_torch_cuda_custom_archs,
+        args=(
+            shared_datadir,
+            tmp_path,
+        ),
+    )
+
+    p.start()
+    p.join()
+    assert p.exitcode == 0
+
+
 def test_torch_common_static(shared_datadir: pathlib.Path, tmp_path: pathlib.Path) -> None:
     project_directory = shared_datadir / "torch_common_static"
     build_directory = tmp_path / "build"