[OpenMP][NFC] Refactor new driver to be more general

[jhuber6]

Summary

This path refactors the new driver to be less dependent on OpenMP. This
is done in preparation for the new driver to be able to handle other
offloading kinds and compile them together.

[jdoerfert]

LG

clang/include/clang/Driver/Compilation.h

public:
   const Driver &getDriver() const { return TheDriver; }
 
   const ToolChain &getDefaultToolChain() const { return DefaultToolChain; }
 
   unsigned isOffloadingHostKind(Action::OffloadKind Kind) const {
     return ActiveOffloadMask & Kind;
   }
 
+  unsigned getActiveOffloadKinds() const { return ActiveOffloadMask; }
+
   /// Iterator that visits device toolchains of a given kind.
   using const_offload_toolchains_iterator =
       const std::multimap<Action::OffloadKind,
                           const ToolChain *>::const_iterator;
   using const_offload_toolchains_range =
       std::pair<const_offload_toolchains_iterator,
                 const_offload_toolchains_iterator>;
 
   template <Action::OffloadKind Kind>
   const_offload_toolchains_range getOffloadToolChains() const {
     return OrderedOffloadingToolchains.equal_range(Kind);
   }
 
+  const_offload_toolchains_range
+  getOffloadToolChains(Action::OffloadKind Kind) const {
+    return OrderedOffloadingToolchains.equal_range(Kind);
+  }
+
   /// Return true if an offloading tool chain of a given kind exists.
   template <Action::OffloadKind Kind> bool hasOffloadToolChain() const {
     return OrderedOffloadingToolchains.find(Kind) !=
            OrderedOffloadingToolchains.end();
   }
 
   /// Return an offload toolchain of the provided kind. Only one is expected to
   /// exist.

clang/lib/Driver/Driver.cpp

if (Arg *A = Args.getLastArg(options::OPT__SLASH_o)) {
     }
   }
 
   handleArguments(C, Args, Inputs, Actions);
 
   // Builder to be used to build offloading actions.
   OffloadingActionBuilder OffloadBuilder(C, Args, Inputs);
 
-  // Offload kinds active for this compilation.
-  unsigned OffloadKinds = Action::OFK_None;
-  if (C.hasOffloadToolChain<Action::OFK_OpenMP>())
-    OffloadKinds |= Action::OFK_OpenMP;
-
   // Construct the actions to perform.
   HeaderModulePrecompileJobAction *HeaderModuleAction = nullptr;
   ActionList LinkerInputs;
   ActionList MergerInputs;
 
   for (auto &I : Inputs) {
     types::ID InputType = I.first;
     const Arg *InputArg = I.second;
for (auto &I : Inputs) {
     // If we ended with something, add to the output list.
     if (Current)
       Actions.push_back(Current);
 
     // Add any top level actions generated for offloading.
     if (!Args.hasArg(options::OPT_fopenmp_new_driver))
       OffloadBuilder.appendTopLevelActions(Actions, Current, InputArg);
     else if (Current)
-      Current->propagateHostOffloadInfo(OffloadKinds,
+      Current->propagateHostOffloadInfo(C.getActiveOffloadKinds(),
                                         /*BoundArch=*/nullptr);
   }
 
   // Add a link action if necessary.
 
   if (LinkerInputs.empty()) {
     Arg *FinalPhaseArg;
     if (getFinalPhase(Args, &FinalPhaseArg) == phases::Link)
       OffloadBuilder.appendDeviceLinkActions(Actions);
   }
 
   if (!LinkerInputs.empty()) {
     if (!Args.hasArg(options::OPT_fopenmp_new_driver))
       if (Action *Wrapper = OffloadBuilder.makeHostLinkAction())
         LinkerInputs.push_back(Wrapper);
     Action *LA;
     // Check if this Linker Job should emit a static library.
     if (ShouldEmitStaticLibrary(Args)) {
       LA = C.MakeAction<StaticLibJobAction>(LinkerInputs, types::TY_Image);
     } else if (Args.hasArg(options::OPT_fopenmp_new_driver) &&
-               OffloadKinds != Action::OFK_None) {
+               C.getActiveOffloadKinds() != Action::OFK_None) {
       LA = C.MakeAction<LinkerWrapperJobAction>(LinkerInputs, types::TY_Image);
-      LA->propagateHostOffloadInfo(OffloadKinds,
+      LA->propagateHostOffloadInfo(C.getActiveOffloadKinds(),
                                    /*BoundArch=*/nullptr);
     } else {
       LA = C.MakeAction<LinkJobAction>(LinkerInputs, types::TY_Image);
     }
     if (!Args.hasArg(options::OPT_fopenmp_new_driver))
       LA = OffloadBuilder.processHostLinkAction(LA);
     Actions.push_back(LA);
   }

 
 Action *Driver::BuildOffloadingActions(Compilation &C,
                                        llvm::opt::DerivedArgList &Args,
                                        const InputTy &Input,
                                        Action *HostAction) const {
   if (!isa<CompileJobAction>(HostAction))
     return HostAction;
 
-  SmallVector<const ToolChain *, 2> ToolChains;
-  ActionList DeviceActions;
+  OffloadAction::DeviceDependences DDeps;
 
   types::ID InputType = Input.first;
   const Arg *InputArg = Input.second;
 
-  auto OpenMPTCRange = C.getOffloadToolChains<Action::OFK_OpenMP>();
-  for (auto TI = OpenMPTCRange.first, TE = OpenMPTCRange.second; TI != TE; ++TI)
+  const Action::OffloadKind OffloadKinds[] = {Action::OFK_OpenMP};
+
+  for (Action::OffloadKind Kind : OffloadKinds) {
+    SmallVector<const ToolChain *, 2> ToolChains;
+    ActionList DeviceActions;
+
+    auto TCRange = C.getOffloadToolChains(Kind);
+    for (auto TI = TCRange.first, TE = TCRange.second; TI != TE; ++TI)
       ToolChains.push_back(TI->second);
 
+    if (ToolChains.empty())
+      continue;
+
     for (unsigned I = 0; I < ToolChains.size(); ++I)
       DeviceActions.push_back(C.MakeAction<InputAction>(*InputArg, InputType));
 
     if (DeviceActions.empty())
       return HostAction;
 
     auto PL = types::getCompilationPhases(*this, Args, InputType);
 
     for (phases::ID Phase : PL) {
       if (Phase == phases::Link) {
         assert(Phase == PL.back() && "linking must be final compilation step.");
         break;
       }
 
       auto TC = ToolChains.begin();
       for (Action *&A : DeviceActions) {
-      A = ConstructPhaseAction(C, Args, Phase, A, Action::OFK_OpenMP);
+        A = ConstructPhaseAction(C, Args, Phase, A, Kind);
 
-      if (isa<CompileJobAction>(A)) {
+        if (isa<CompileJobAction>(A) && Kind == Action::OFK_OpenMP) {
           HostAction->setCannotBeCollapsedWithNextDependentAction();
           OffloadAction::HostDependence HDep(
               *HostAction, *C.getSingleOffloadToolChain<Action::OFK_Host>(),
               /*BourdArch=*/nullptr, Action::OFK_OpenMP);
           OffloadAction::DeviceDependences DDep;
-        DDep.add(*A, **TC, /*BoundArch=*/nullptr, Action::OFK_OpenMP);
+          DDep.add(*A, **TC, /*BoundArch=*/nullptr, Kind);
           A = C.MakeAction<OffloadAction>(HDep, DDep);
         }
         ++TC;
       }
     }
 
-  OffloadAction::DeviceDependences DDeps;
-
     auto TC = ToolChains.begin();
     for (Action *A : DeviceActions) {
-    DDeps.add(*A, **TC, /*BoundArch=*/nullptr, Action::OFK_OpenMP);
+      DDeps.add(*A, **TC, /*BoundArch=*/nullptr, Kind);
       TC++;
     }
+  }
 
   OffloadAction::HostDependence HDep(
       *HostAction, *C.getSingleOffloadToolChain<Action::OFK_Host>(),
       /*BoundArch=*/nullptr, DDeps);
   return C.MakeAction<OffloadAction>(HDep, DDeps);
 }
 
 Action *Driver::ConstructPhaseAction(