[OpenMP] Make clang argument handling for the new driver more generic

[jhuber6]

Summary

In preparation for accepting other offloading kinds with the new driver,
this patch makes the way we handle offloading actions more generic. A
new field to get the associated device action's toolchain is used rather
than manually iterating a list. This makes building the arguments easier
and makes sure that we doin't rely on any implicit ordering.

[jhuber6]

ping

[jhuber6]

Rebase after making the new driver default.

[jhuber6]

Fix

[jhuber6]

Fix

[jhuber6]

Fix for real after rebasing failure...

[jhuber6]

Sorry for the spam

[jhuber6]

Fix after fix.

[jhuber6]

Hopefully this is the last time I need to update this

[yaxunl]

clang/lib/Driver/ToolChains/Clang.cpp

6949–6952

Linux path may contain ',' which cannot be passed by this option.

We've seen similar issues with -inputs= option of clang-offload-bundler. To solve that issue, we had to add support of multiple -input= options.

[jhuber6]

clang/lib/Driver/ToolChains/Clang.cpp

6949–6952

Ugh, you could probably split this starting from the end since we assume it will always end with at least three commas, but that's not a satisfying solution. I've been pondering creating a new tool to do this, but haven't quite decided how that should look. I'd imagine it would basically behave like how we invoke fatbinary.

This issue isn't related to this patch, so I'll probably just parse it in reverse upstream, thanks for bringing this up.

[yaxunl]

clang/lib/Driver/ToolChains/Clang.cpp

6949–6952

I agree the issue is orthogonal to this patch.

Parsing backward is an option. Another option is to put file name at the end of the option. Maybe the latter is less confusing.

[jhuber6]

ping

[yaxunl]

LGTM. Thanks.

clang/include/clang/Driver/Action.h

protected:
   unsigned ActiveOffloadKindMask = 0u;
 
   /// Offloading kind of the device.
   OffloadKind OffloadingDeviceKind = OFK_None;
 
   /// The Offloading architecture associated with this action.
   const char *OffloadingArch = nullptr;
 
+  /// The Offloading toolchain associated with this device action.
+  const ToolChain *OffloadingToolChain = nullptr;
+
   Action(ActionClass Kind, types::ID Type) : Action(Kind, ActionList(), Type) {}
   Action(ActionClass Kind, Action *Input, types::ID Type)
       : Action(Kind, ActionList({Input}), Type) {}
   Action(ActionClass Kind, Action *Input)
       : Action(Kind, ActionList({Input}), Input->getType()) {}
   Action(ActionClass Kind, const ActionList &Inputs, types::ID Type)
       : Kind(Kind), Type(Type), Inputs(Inputs) {}
 
GetOffloadingFileNamePrefix(OffloadKind Kind,
                               StringRef NormalizedTriple,
                               bool CreatePrefixForHost = false);
 
   /// Return a string containing a offload kind name.
   static StringRef GetOffloadKindName(OffloadKind Kind);
 
   /// Set the device offload info of this action and propagate it to its
   /// dependences.
-  void propagateDeviceOffloadInfo(OffloadKind OKind, const char *OArch);
+  void propagateDeviceOffloadInfo(OffloadKind OKind, const char *OArch,
+                                  const ToolChain *OToolChain);
 
   /// 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;
   }
 
   OffloadKind getOffloadingDeviceKind() const { return OffloadingDeviceKind; }
   const char *getOffloadingArch() const { return OffloadingArch; }
+  const ToolChain *getOffloadingToolChain() const {
+    return OffloadingToolChain;
+  }
 
   /// Check if this action have any offload kinds. Note that host offload kinds
   /// are only set if the action is a dependence to a host offload action.
-  bool isHostOffloading(OffloadKind OKind) const {
+  bool isHostOffloading(unsigned int OKind) const {
     return ActiveOffloadKindMask & OKind;
   }
   bool isDeviceOffloading(OffloadKind OKind) const {
     return OffloadingDeviceKind == OKind;
   }
   bool isOffloading(OffloadKind OKind) const {
     return isHostOffloading(OKind) || isDeviceOffloading(OKind);
   }

clang/lib/Driver/Action.cpp

case LinkerWrapperJobClass:
     return "clang-linker-wrapper";
   case StaticLibJobClass:
     return "static-lib-linker";
   }
 
   llvm_unreachable("invalid class");
 }
 
-void Action::propagateDeviceOffloadInfo(OffloadKind OKind, const char *OArch) {
+void Action::propagateDeviceOffloadInfo(OffloadKind OKind, const char *OArch,
+                                        const ToolChain *OToolChain) {
   // Offload action set its own kinds on their dependences.
   if (Kind == OffloadClass)
     return;
   // Unbundling actions use the host kinds.
   if (Kind == OffloadUnbundlingJobClass)
     return;
 
   assert((OffloadingDeviceKind == OKind || OffloadingDeviceKind == OFK_None) &&
          "Setting device kind to a different device??");
   assert(!ActiveOffloadKindMask && "Setting a device kind in a host action??");
   OffloadingDeviceKind = OKind;
   OffloadingArch = OArch;
+  OffloadingToolChain = OToolChain;
 
   for (auto *A : Inputs)
-    A->propagateDeviceOffloadInfo(OffloadingDeviceKind, OArch);
+    A->propagateDeviceOffloadInfo(OffloadingDeviceKind, OArch, OToolChain);
 }
 
 void Action::propagateHostOffloadInfo(unsigned OKinds, const char *OArch) {
   // Offload action set its own kinds on their dependences.
   if (Kind == OffloadClass)
     return;
 
   assert(OffloadingDeviceKind == OFK_None &&
          "Setting a host kind in a device action.");
   ActiveOffloadKindMask |= OKinds;
   OffloadingArch = OArch;
 
   for (auto *A : Inputs)
     A->propagateHostOffloadInfo(ActiveOffloadKindMask, OArch);
 }
 
 void Action::propagateOffloadInfo(const Action *A) {
   if (unsigned HK = A->getOffloadingHostActiveKinds())
     propagateHostOffloadInfo(HK, A->getOffloadingArch());
   else
     propagateDeviceOffloadInfo(A->getOffloadingDeviceKind(),
-                               A->getOffloadingArch());
+                               A->getOffloadingArch(),
+                               A->getOffloadingToolChain());
 }
 
 std::string Action::getOffloadingKindPrefix() const {
   switch (OffloadingDeviceKind) {
   case OFK_None:
     break;
   case OFK_Host:
     llvm_unreachable("Host kind is not an offloading device kind.");
HDep.getAction()->propagateHostOffloadInfo(HDep.getOffloadKinds(),
                                              HDep.getBoundArch());
 }
 
 OffloadAction::OffloadAction(const DeviceDependences &DDeps, types::ID Ty)
     : Action(OffloadClass, DDeps.getActions(), Ty),
       DevToolChains(DDeps.getToolChains()) {
   auto &OKinds = DDeps.getOffloadKinds();
   auto &BArchs = DDeps.getBoundArchs();
+  auto &OTCs = DDeps.getToolChains();
 
   // If all inputs agree on the same kind, use it also for this action.
   if (llvm::all_of(OKinds, [&](OffloadKind K) { return K == OKinds.front(); }))
     OffloadingDeviceKind = OKinds.front();
 
   // If we have a single dependency, inherit the architecture from it.
   if (OKinds.size() == 1)
     OffloadingArch = BArchs.front();
 
   // Propagate info to the dependencies.
   for (unsigned i = 0, e = getInputs().size(); i != e; ++i)
-    getInputs()[i]->propagateDeviceOffloadInfo(OKinds[i], BArchs[i]);
+    getInputs()[i]->propagateDeviceOffloadInfo(OKinds[i], BArchs[i], OTCs[i]);
 }
 
 OffloadAction::OffloadAction(const HostDependence &HDep,
                              const DeviceDependences &DDeps)
     : Action(OffloadClass, HDep.getAction()), HostTC(HDep.getToolChain()),
       DevToolChains(DDeps.getToolChains()) {
   // We use the kinds of the host dependence for this action.
   OffloadingArch = HDep.getBoundArch();
   ActiveOffloadKindMask = HDep.getOffloadKinds();
   HDep.getAction()->propagateHostOffloadInfo(HDep.getOffloadKinds(),
                                              HDep.getBoundArch());
 
   // Add device inputs and propagate info to the device actions. Do work only if
   // we have dependencies.
   for (unsigned i = 0, e = DDeps.getActions().size(); i != e; ++i)
     if (auto *A = DDeps.getActions()[i]) {
       getInputs().push_back(A);
       A->propagateDeviceOffloadInfo(DDeps.getOffloadKinds()[i],
-                                    DDeps.getBoundArchs()[i]);
+                                    DDeps.getBoundArchs()[i],
+                                    DDeps.getToolChains()[i]);
     }
 }
 
 void OffloadAction::doOnHostDependence(const OffloadActionWorkTy &Work) const {
   if (!HostTC)
     return;
   assert(!getInputs().empty() && "No dependencies for offload action??");
   auto *A = getInputs().front();

clang/lib/Driver/ToolChains/Clang.cpp

void Clang::ConstructJob(Compilation &C, const JobAction &JA,
   // include as part of the module. API extraction accepts a list of header
   // files whose API information is emitted in the output. All other jobs are
   // expected to have exactly one input.
   bool IsCuda = JA.isOffloading(Action::OFK_Cuda);
   bool IsCudaDevice = JA.isDeviceOffloading(Action::OFK_Cuda);
   bool IsHIP = JA.isOffloading(Action::OFK_HIP);
   bool IsHIPDevice = JA.isDeviceOffloading(Action::OFK_HIP);
   bool IsOpenMPDevice = JA.isDeviceOffloading(Action::OFK_OpenMP);
-  bool IsOpenMPHost = JA.isHostOffloading(Action::OFK_OpenMP);
   bool IsHeaderModulePrecompile = isa<HeaderModulePrecompileJobAction>(JA);
   bool IsExtractAPI = isa<ExtractAPIJobAction>(JA);
   bool IsDeviceOffloadAction = !(JA.isDeviceOffloading(Action::OFK_None) ||
                                  JA.isDeviceOffloading(Action::OFK_Host));
+  bool IsHostOffloadingAction =
+      JA.isHostOffloading(Action::OFK_OpenMP) &&
+      !Args.hasArg(options::OPT_fno_openmp_new_driver);
   bool IsUsingLTO = D.isUsingLTO(IsDeviceOffloadAction);
   auto LTOMode = D.getLTOMode(IsDeviceOffloadAction);
 
   // A header module compilation doesn't have a main input file, so invent a
   // fake one as a placeholder.
   const char *ModuleName = [&]{
     auto *ModuleNameArg = Args.getLastArg(options::OPT_fmodule_name_EQ);
     return ModuleNameArg ? ModuleNameArg->getValue() : "";
if (IsHeaderModulePrecompile)
       return HeaderModuleInput;
     if (IsExtractAPI)
       return ExtractAPIPlaceholderInput;
     return Inputs[0];
   }();
 
   InputInfoList ModuleHeaderInputs;
   InputInfoList ExtractAPIInputs;
-  InputInfoList OpenMPHostInputs;
+  InputInfoList HostOffloadingInputs;
   const InputInfo *CudaDeviceInput = nullptr;
   const InputInfo *OpenMPDeviceInput = nullptr;
   for (const InputInfo &I : Inputs) {
     if (&I == &Input) {
       // This is the primary input.
     } else if (IsHeaderModulePrecompile &&
                types::getPrecompiledType(I.getType()) == types::TY_PCH) {
       types::ID Expected = HeaderModuleInput.getType();
       if (I.getType() != Expected) {
         D.Diag(diag::err_drv_module_header_wrong_kind)
             << I.getFilename() << types::getTypeName(I.getType())
             << types::getTypeName(Expected);
       }
       ModuleHeaderInputs.push_back(I);
     } else if (IsExtractAPI) {
       auto ExpectedInputType = ExtractAPIPlaceholderInput.getType();
       if (I.getType() != ExpectedInputType) {
         D.Diag(diag::err_drv_extract_api_wrong_kind)
             << I.getFilename() << types::getTypeName(I.getType())
             << types::getTypeName(ExpectedInputType);
       }
       ExtractAPIInputs.push_back(I);
+    } else if (IsHostOffloadingAction) {
+      HostOffloadingInputs.push_back(I);
     } else if ((IsCuda || IsHIP) && !CudaDeviceInput) {
       CudaDeviceInput = &I;
     } else if (IsOpenMPDevice && !OpenMPDeviceInput) {
       OpenMPDeviceInput = &I;
-    } else if (IsOpenMPHost) {
-      OpenMPHostInputs.push_back(I);
     } else {
       llvm_unreachable("unexpectedly given multiple inputs");
     }
   }
 
   const llvm::Triple *AuxTriple =
       (IsCuda || IsHIP) ? TC.getAuxTriple() : nullptr;
   bool IsWindowsMSVC = RawTriple.isWindowsMSVCEnvironment();

   if (IsOpenMPDevice) {
     CmdArgs.push_back("-fopenmp-is-device");
     if (OpenMPDeviceInput) {
       CmdArgs.push_back("-fopenmp-host-ir-file-path");
       CmdArgs.push_back(Args.MakeArgString(OpenMPDeviceInput->getFilename()));
     }
   }
 
-  // Host-side OpenMP offloading recieves the device object files and embeds it
-  // in a named section including the associated target triple and architecture.
-  if (IsOpenMPHost && !OpenMPHostInputs.empty()) {
-    auto InputFile = OpenMPHostInputs.begin();
-    auto OpenMPTCs = C.getOffloadToolChains<Action::OFK_OpenMP>();
-    for (auto TI = OpenMPTCs.first, TE = OpenMPTCs.second; TI != TE;
-         ++TI, ++InputFile) {
-      const ToolChain *TC = TI->second;
-      const ArgList &TCArgs = C.getArgsForToolChain(TC, "", Action::OFK_OpenMP);
-      StringRef File =
-          C.getArgs().MakeArgString(TC->getInputFilename(*InputFile));
+  // Host-side offloading recieves the device object files and embeds it in a
+  // named section including the associated target triple and architecture.
+  for (const InputInfo Input : HostOffloadingInputs) {
+    const Action *OffloadAction = Input.getAction();
+    const ToolChain *TC = OffloadAction->getOffloadingToolChain();
+    const ArgList &TCArgs =
+        C.getArgsForToolChain(TC, OffloadAction->getOffloadingArch(),
+                              OffloadAction->getOffloadingDeviceKind());
+    StringRef File = C.getArgs().MakeArgString(TC->getInputFilename(Input));
+    StringRef Arch = (OffloadAction->getOffloadingArch())
+                         ? OffloadAction->getOffloadingArch()
+                         : TCArgs.getLastArgValue(options::OPT_march_EQ);
 
-      CmdArgs.push_back(
-          Args.MakeArgString("-fembed-offload-object=" + File + "," +
-                             Action::GetOffloadKindName(Action::OFK_OpenMP) +
-                             "," + TC->getTripleString() + "," +
-                             TCArgs.getLastArgValue(options::OPT_march_EQ)));
-    }
+    CmdArgs.push_back(Args.MakeArgString(
+        "-fembed-offload-object=" + File + "," +
+        Action::GetOffloadKindName(OffloadAction->getOffloadingDeviceKind()) +
+        "," + TC->getTripleString() + "," + Arch));
   }
 
   if (Triple.isAMDGPU()) {
     handleAMDGPUCodeObjectVersionOptions(D, Args, CmdArgs);
 
     Args.addOptInFlag(CmdArgs, options::OPT_munsafe_fp_atomics,
                       options::OPT_mno_unsafe_fp_atomics);
   }