[NVPTX] Fix NVPTX output name in the driver with `-save-temps`

Summary:
Currently, OpenMP and direct compilation uses an NVPTX toolchain to
directly invoke the CUDA tools from Clang to do the assembling and
linking of NVPTX codes. This breaks under `-save-temps` because of a
workaround. The `nvlink` linker does not accept `.o` files, so we need
to be selective when we output these. The previous logic keyed off of
presense in the temp files and wasn't a great solution. Change this to
just query the input args for `-c` to see if we stop at the assembler.

Fixes https://github.com/llvm/llvm-project/issues/60767

diff --git a/clang/lib/Driver/ToolChains/Cuda.cpp b/clang/lib/Driver/ToolChains/Cuda.cpp
index 6d710fd..95d932e 100644 (file)
--- a/clang/lib/Driver/ToolChains/Cuda.cpp
+++ b/clang/lib/Driver/ToolChains/Cuda.cpp
@@ -446,13 +446,8 @@ void NVPTX::Assembler::ConstructJob(Compilation &C, const JobAction &JA,
   std::string OutputFileName = TC.getInputFilename(Output);
 
   // If we are invoking `nvlink` internally we need to output a `.cubin` file.
-  // Checking if the output is a temporary is the cleanest way to determine
-  // this. Putting this logic in `getInputFilename` isn't an option because it
-  // relies on the compilation.
   // FIXME: This should hopefully be removed if NVIDIA updates their tooling.
-  if (Output.isFilename() &&
-      llvm::find(C.getTempFiles(), Output.getFilename()) !=
-          C.getTempFiles().end()) {
+  if (!C.getInputArgs().getLastArg(options::OPT_c)) {
     SmallString<256> Filename(Output.getFilename());
     llvm::sys::path::replace_extension(Filename, "cubin");
     OutputFileName = Filename.str();

diff --git a/clang/test/Driver/cuda-cross-compiling.c b/clang/test/Driver/cuda-cross-compiling.c
index 992fda8..9750f09 100644 (file)
--- a/clang/test/Driver/cuda-cross-compiling.c
+++ b/clang/test/Driver/cuda-cross-compiling.c
@@ -43,6 +43,8 @@
 //
 // RUN: %clang -target nvptx64-nvidia-cuda -march=sm_61 -c -### %s 2>&1 \
 // RUN:   | FileCheck -check-prefix=OBJECT %s
+// RUN: %clang -target nvptx64-nvidia-cuda -save-temps -march=sm_61 -c -### %s 2>&1 \
+// RUN:   | FileCheck -check-prefix=OBJECT %s
 
 //      OBJECT: -cc1" "-triple" "nvptx64-nvidia-cuda" "-S" {{.*}} "-target-cpu" "sm_61" "-target-feature" "+ptx{{[0-9]+}}" {{.*}} "-o" "[[PTX:.+]].s"
 // OBJECT-NEXT: ptxas{{.*}}"-m64" "-O0" "--gpu-name" "sm_61" "--output-file" "[[OBJ:.+]].o" "[[PTX]].s" "-c"