[CUDA][HIP] Fix offloading kind for linking C++ programs

When both CUDA or HIP programs and C++ programs are passed
to clang driver without -c, C++ programs are treated as CUDA
or HIP program, which is incorrect.

This is because action builder sets the offloading kind of input
job actions to the linking action to be the union of offloading
kind of the input job actions, i.e. if there is one HIP or CUDA
input to the linker, then all the input to the linker is marked
as HIP or CUDA.

To fix this issue, the offload action builder tracks the originating
input argument of each host action, which allows it to determine
the active offload kind of each host action. Then the offload
kind of each input action to the linker can be determined
individually.

Reviewed by: Artem Belevich

Differential Revision: https://reviews.llvm.org/D120911

Author: yxsamliu

clang/include/clang/Driver/Action.h

@@ -189,6 +189,11 @@ class Action {
   /// dependences.
   void propagateHostOffloadInfo(unsigned OKinds, const char *OArch);
 
+  void setHostOffloadInfo(unsigned OKinds, const char *OArch) {
+    ActiveOffloadKindMask |= OKinds;
+    OffloadingArch = OArch;
+  }
+
   /// Set the offload info of this action to be the same as the provided action,
   /// and propagate it to its dependences.
   void propagateOffloadInfo(const Action *A);

clang/lib/Driver/Driver.cpp

@@ -2465,6 +2465,9 @@ class OffloadingActionBuilder final {
   /// Map between an input argument and the offload kinds used to process it.
   std::map<const Arg *, unsigned> InputArgToOffloadKindMap;
 
+  /// Map between a host action and its originating input argument.
+  std::map<Action *, const Arg *> HostActionToInputArgMap;
+
   /// Builder interface. It doesn't build anything or keep any state.
   class DeviceActionBuilder {
   public:
@@ -3449,6 +3452,17 @@ class OffloadingActionBuilder final {
       delete SB;
   }
 
+  /// Record a host action and its originating input argument.
+  void recordHostAction(Action *HostAction, const Arg *InputArg) {
+    assert(HostAction && "Invalid host action");
+    assert(InputArg && "Invalid input argument");
+    auto Loc = HostActionToInputArgMap.find(HostAction);
+    if (Loc == HostActionToInputArgMap.end())
+      HostActionToInputArgMap[HostAction] = InputArg;
+    assert(HostActionToInputArgMap[HostAction] == InputArg &&
+           "host action mapped to multiple input arguments");
+  }
+
   /// Generate an action that adds device dependences (if any) to a host action.
   /// If no device dependence actions exist, just return the host action \a
   /// HostAction. If an error is found or if no builder requires the host action
@@ -3464,6 +3478,7 @@ class OffloadingActionBuilder final {
       return HostAction;
 
     assert(HostAction && "Invalid host action!");
+    recordHostAction(HostAction, InputArg);
 
     OffloadAction::DeviceDependences DDeps;
     // Check if all the programming models agree we should not emit the host
@@ -3517,6 +3532,8 @@ class OffloadingActionBuilder final {
     if (!IsValid)
       return true;
 
+    recordHostAction(HostAction, InputArg);
+
     // If we are supporting bundling/unbundling and the current action is an
     // input action of non-source file, we replace the host action by the
     // unbundling action. The bundler tool has the logic to detect if an input
@@ -3533,6 +3550,7 @@ class OffloadingActionBuilder final {
           C.getSingleOffloadToolChain<Action::OFK_Host>(),
           /*BoundArch=*/StringRef(), Action::OFK_Host);
       HostAction = UnbundlingHostAction;
+      recordHostAction(HostAction, InputArg);
     }
 
     assert(HostAction && "Invalid host action!");
@@ -3569,6 +3587,9 @@ class OffloadingActionBuilder final {
   /// programming models allow it.
   bool appendTopLevelActions(ActionList &AL, Action *HostAction,
                              const Arg *InputArg) {
+    if (HostAction)
+      recordHostAction(HostAction, InputArg);
+
     // Get the device actions to be appended.
     ActionList OffloadAL;
     for (auto *SB : SpecializedBuilders) {
@@ -3590,6 +3611,7 @@ class OffloadingActionBuilder final {
       // before this method was called.
       assert(HostAction == AL.back() && "Host action not in the list??");
       HostAction = C.MakeAction<OffloadBundlingJobAction>(OffloadAL);
+      recordHostAction(HostAction, InputArg);
       AL.back() = HostAction;
     } else
       AL.append(OffloadAL.begin(), OffloadAL.end());
@@ -3623,6 +3645,11 @@ class OffloadingActionBuilder final {
       if (!SB->isValid())
         continue;
       HA = SB->appendLinkHostActions(DeviceAL);
+      // This created host action has no originating input argument, therefore
+      // needs to set its offloading kind directly.
+      if (HA)
+        HA->propagateHostOffloadInfo(SB->getAssociatedOffloadKind(),
+                                     /*BoundArch=*/nullptr);
     }
     return HA;
   }
@@ -3649,10 +3676,22 @@ class OffloadingActionBuilder final {
     // If we don't have device dependencies, we don't have to create an offload
     // action.
     if (DDeps.getActions().empty()) {
-      // Propagate all the active kinds to host action. Given that it is a link
-      // action it is assumed to depend on all actions generated so far.
-      HostAction->propagateHostOffloadInfo(ActiveOffloadKinds,
-                                           /*BoundArch=*/nullptr);
+      // Set all the active offloading kinds to the link action. Given that it
+      // is a link action it is assumed to depend on all actions generated so
+      // far.
+      HostAction->setHostOffloadInfo(ActiveOffloadKinds,
+                                     /*BoundArch=*/nullptr);
+      // Propagate active offloading kinds for each input to the link action.
+      // Each input may have different active offloading kind.
+      for (auto A : HostAction->inputs()) {
+        auto ArgLoc = HostActionToInputArgMap.find(A);
+        if (ArgLoc == HostActionToInputArgMap.end())
+          continue;
+        auto OFKLoc = InputArgToOffloadKindMap.find(ArgLoc->second);
+        if (OFKLoc == InputArgToOffloadKindMap.end())
+          continue;
+        A->propagateHostOffloadInfo(OFKLoc->second, /*BoundArch=*/nullptr);
+      }
       return HostAction;
     }
 

clang/test/Driver/hip-phases.hip

@@ -459,17 +459,65 @@
 // Test mixed HIP and C++ compilation. HIP program should have HIP offload kind.
 // C++ program should have no offload kind.
 
+// Test compile empty.hip and empty.cpp.
 // RUN: %clang -target x86_64-unknown-linux-gnu \
 // RUN: -ccc-print-phases --cuda-gpu-arch=gfx803 --cuda-gpu-arch=gfx900 \
 // RUN: -c %S/Inputs/empty.hip %S/Inputs/empty.cpp 2>&1 | FileCheck -check-prefixes=MIXED %s
-
 // RUN: %clang -target x86_64-unknown-linux-gnu \
 // RUN: -ccc-print-phases --cuda-gpu-arch=gfx803 --cuda-gpu-arch=gfx900 \
 // RUN: -c %S/Inputs/empty.hip %S/Inputs/empty.cpp 2>&1 | FileCheck -check-prefixes=MIXED-NEG %s
 
+// Test compile and link empty.hip and empty.cpp.
+// RUN: %clang -target x86_64-unknown-linux-gnu \
+// RUN: -ccc-print-phases --cuda-gpu-arch=gfx803 --cuda-gpu-arch=gfx900 \
+// RUN: %S/Inputs/empty.hip %S/Inputs/empty.cpp 2>&1 | FileCheck -check-prefixes=MIXED %s
+// RUN: %clang -target x86_64-unknown-linux-gnu \
+// RUN: -ccc-print-phases --cuda-gpu-arch=gfx803 --cuda-gpu-arch=gfx900 \
+// RUN: %S/Inputs/empty.hip %S/Inputs/empty.cpp 2>&1 | FileCheck -check-prefixes=MIXED-NEG %s
+
+// Test compile and link empty.hip and empty.cpp with --hip-link -fgpu-rdc.
+// RUN: %clang -target x86_64-unknown-linux-gnu --hip-link -fgpu-rdc \
+// RUN: -ccc-print-phases --cuda-gpu-arch=gfx803 --cuda-gpu-arch=gfx900 \
+// RUN: %S/Inputs/empty.hip %S/Inputs/empty.cpp 2>&1 | FileCheck -check-prefixes=MIXED %s
+// RUN: %clang -target x86_64-unknown-linux-gnu --hip-link -fgpu-rdc \
+// RUN: -ccc-print-phases --cuda-gpu-arch=gfx803 --cuda-gpu-arch=gfx900 \
+// RUN: %S/Inputs/empty.hip %S/Inputs/empty.cpp 2>&1 | FileCheck -check-prefixes=MIXED-NEG %s
+
+// Test compile and link -x hip empty.hip and -x c++ empty.cpp.
+// RUN: %clang -target x86_64-unknown-linux-gnu \
+// RUN: -ccc-print-phases --cuda-gpu-arch=gfx803 --cuda-gpu-arch=gfx900 \
+// RUN: -x hip %S/Inputs/empty.hip -x c++ %S/Inputs/empty.cpp 2>&1 | FileCheck -check-prefixes=MIXED %s
+// RUN: %clang -target x86_64-unknown-linux-gnu \
+// RUN: -ccc-print-phases --cuda-gpu-arch=gfx803 --cuda-gpu-arch=gfx900 \
+// RUN: -x hip %S/Inputs/empty.hip -x c++ %S/Inputs/empty.cpp 2>&1 | FileCheck -check-prefixes=MIXED-NEG %s
+
+// Test compile and link -x hip empty.hip and empty.cpp.
+// RUN: %clang -target x86_64-unknown-linux-gnu \
+// RUN: -ccc-print-phases --cuda-gpu-arch=gfx803 --cuda-gpu-arch=gfx900 \
+// RUN: -x hip %S/Inputs/empty.hip %S/Inputs/empty.cpp 2>&1 | FileCheck -check-prefixes=MIXED2 %s
+// RUN: %clang -target x86_64-unknown-linux-gnu \
+// RUN: -ccc-print-phases --cuda-gpu-arch=gfx803 --cuda-gpu-arch=gfx900 \
+// RUN: -x hip %S/Inputs/empty.hip %S/Inputs/empty.cpp 2>&1 | FileCheck -check-prefixes=MIXED2-NEG %s
+
+// Test compile and link empty.hip and -x hip empty.cpp.
+// RUN: %clang -target x86_64-unknown-linux-gnu \
+// RUN: -ccc-print-phases --cuda-gpu-arch=gfx803 --cuda-gpu-arch=gfx900 \
+// RUN: %S/Inputs/empty.hip -x hip %S/Inputs/empty.cpp 2>&1 | FileCheck -check-prefixes=MIXED2 %s
+// RUN: %clang -target x86_64-unknown-linux-gnu \
+// RUN: -ccc-print-phases --cuda-gpu-arch=gfx803 --cuda-gpu-arch=gfx900 \
+// RUN: -x hip %S/Inputs/empty.hip %S/Inputs/empty.cpp 2>&1 | FileCheck -check-prefixes=MIXED2-NEG %s
+
 // MIXED-DAG: input, "{{.*}}empty.hip", hip, (host-hip)
 // MIXED-DAG: input, "{{.*}}empty.hip", hip, (device-hip, gfx803)
 // MIXED-DAG: input, "{{.*}}empty.hip", hip, (device-hip, gfx900)
 // MIXED-DAG: input, "{{.*}}empty.cpp", c++
 // MIXED-NEG-NOT: input, "{{.*}}empty.cpp", c++, (host-hip)
 // MIXED-NEG-NOT: input, "{{.*}}empty.cpp", c++, (device-hip
+
+// MIXED2-DAG: input, "{{.*}}empty.hip", hip, (host-hip)
+// MIXED2-DAG: input, "{{.*}}empty.hip", hip, (device-hip, gfx803)
+// MIXED2-DAG: input, "{{.*}}empty.hip", hip, (device-hip, gfx900)
+// MIXED2-DAG: input, "{{.*}}empty.cpp", hip, (host-hip)
+// MIXED2-DAG: input, "{{.*}}empty.cpp", hip, (device-hip, gfx803)
+// MIXED2-DAG: input, "{{.*}}empty.cpp", hip, (device-hip, gfx900)
+// MIXED2-NEG-NOT: input, "{{.*}}empty.cpp", c++