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

[yaxunl]

Summary

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.

[tra]

We should probably also check what happens when we specify compilation language explicitly: E.g. clang -x cuda a.cu -x c++ b.cc, clang -x cuda a.cu b.cc and clang a.cu -x cuda b.cc.

[yaxunl]

add more tests

clang/include/clang/Driver/Action.h

public:
   /// Set the device offload info of this action and propagate it to its
   /// dependences.
   void propagateDeviceOffloadInfo(OffloadKind OKind, const char *OArch);
 
   /// Append the host offload info of this action and propagate it to its
   /// 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);
 
   unsigned getOffloadingHostActiveKinds() const {
     return ActiveOffloadKindMask;
   }
 

clang/lib/Driver/Driver.cpp

class OffloadingActionBuilder final {
   bool IsValid = false;
 
   /// The compilation that is using this builder.
   Compilation &C;
 
   /// 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:
     typedef const llvm::SmallVectorImpl<phases::ID> PhasesTy;
 
     enum ActionBuilderReturnCode {
       // The builder acted successfully on the current action.
       ABRT_Success,
CanUseBundler =
         ValidBuilders && ValidBuilders == ValidBuildersSupportingBundling;
   }
 
   ~OffloadingActionBuilder() {
     for (auto *SB : SpecializedBuilders)
       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
   /// to be generated, return nullptr.
   Action *
   addDeviceDependencesToHostAction(Action *HostAction, const Arg *InputArg,
                                    phases::ID CurPhase, phases::ID FinalPhase,
                                    DeviceActionBuilder::PhasesTy &Phases) {
     if (!IsValid)
       return nullptr;
 
     if (SpecializedBuilders.empty())
       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
     // action. Also, keep track of the offloading kinds employed.
     auto &OffloadKind = InputArgToOffloadKindMap[InputArg];
     unsigned InactiveBuilders = 0u;
     unsigned IgnoringBuilders = 0u;
     for (auto *SB : SpecializedBuilders) {
public:
   /// Generate an action that adds a host dependence to a device action. The
   /// results will be kept in this action builder. Return true if an error was
   /// found.
   bool addHostDependenceToDeviceActions(Action *&HostAction,
                                         const Arg *InputArg) {
     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
     // is a bundle or not and if the input is not a bundle it assumes it is a
     // host file. Therefore it is safe to create an unbundling action even if
     // the input is not a bundle.
     if (CanUseBundler && isa<InputAction>(HostAction) &&
         InputArg->getOption().getKind() == llvm::opt::Option::InputClass &&
         (!types::isSrcFile(HostAction->getType()) ||
          HostAction->getType() == types::TY_PP_HIP)) {
       auto UnbundlingHostAction =
           C.MakeAction<OffloadUnbundlingJobAction>(HostAction);
       UnbundlingHostAction->registerDependentActionInfo(
           C.getSingleOffloadToolChain<Action::OFK_Host>(),
           /*BoundArch=*/StringRef(), Action::OFK_Host);
       HostAction = UnbundlingHostAction;
+      recordHostAction(HostAction, InputArg);
     }
 
     assert(HostAction && "Invalid host action!");
 
     // Register the offload kinds that are used.
     auto &OffloadKind = InputArgToOffloadKindMap[InputArg];
     for (auto *SB : SpecializedBuilders) {
       if (!SB->isValid())
bool addHostDependenceToDeviceActions(Action *&HostAction,
     return false;
   }
 
   /// Add the offloading top level actions to the provided action list. This
   /// function can replace the host action by a bundling action if the
   /// 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) {
       if (!SB->isValid())
         continue;
       SB->appendTopLevelActions(OffloadAL);
     }
 
     // If we can use the bundler, replace the host action by the bundling one in
     // the resulting list. Otherwise, just append the device actions. For
     // device only compilation, HostAction is a null pointer, therefore only do
     // this when HostAction is not a null pointer.
     if (CanUseBundler && HostAction &&
         HostAction->getType() != types::TY_Nothing && !OffloadAL.empty()) {
       // Add the host action to the list in order to create the bundling action.
       OffloadAL.push_back(HostAction);
 
       // We expect that the host action was just appended to the action list
       // 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());
 
     // Propagate to the current host action (if any) the offload information
     // associated with the current input.
     if (HostAction)
       HostAction->propagateHostOffloadInfo(InputArgToOffloadKindMap[InputArg],
if (DeviceAL.empty())
       return nullptr;
 
     // Let builders add host linking actions.
     Action* HA = nullptr;
     for (DeviceActionBuilder *SB : SpecializedBuilders) {
       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;
   }
 
   /// Processes the host linker action. This currently consists of replacing it
   /// with an offload action if there are device link objects and propagate to
   /// the host action all the offload kinds used in the current compilation. The
   /// resulting action is returned.
Action *processHostLinkAction(Action *HostAction) {
     // Calculate all the offload kinds used in the current compilation.
     unsigned ActiveOffloadKinds = 0u;
     for (auto &I : InputArgToOffloadKindMap)
       ActiveOffloadKinds |= I.second;
 
     // 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,
+      // 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;
     }
 
     // Create the offload action with all dependences. When an offload action
     // is created the kinds are propagated to the host action, so we don't have
     // to do that explicitly here.
     OffloadAction::HostDependence HDep(
         *HostAction, *C.getSingleOffloadToolChain<Action::OFK_Host>(),

clang/test/Driver/hip-phases.hip

 // PPELLVM2-DAG: [[P8:[0-9]+]]: backend, {[[P7]]}, ir, (device-[[T]], [[ARCH2]])
 // PPELLVM2-DAG: [[P9:[0-9]+]]: offload, "device-[[T]] (amdgcn-amd-amdhsa:[[ARCH2]])" {[[P8]]}, ir
 // PPELLVM2-DAG: [[P10:[0-9]+]]: clang-offload-bundler, {[[P4]], [[P9]]}, ir, (device-hip, )
 // PPELLVM2-NOT: host
 
 // 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++