[CUDA][OpenMP] Create generic offload toolchains

[sfantao]

Summary

This patch introduces the concept of offloading tool chain and offloading kind. Each tool chain may have associated an offloading kind that marks it as used in a given programming model that requires offloading.

It also adds the logic to iterate on the tool chains based on the kind. Currently, only CUDA is supported, but in general a programming model (an offloading kind) may have associated multiple tool chains that require supporting offloading.

This patch does not add tests - its goal is to keep the existing functionality.

This patch is the first of a series of three that attempts to make the current support of CUDA more generic and easier to extend to other programming models, namely OpenMP. It tries to capture the suggestions/improvements/concerns on the initial proposal in http://lists.llvm.org/pipermail/cfe-dev/2016-February/047547.html. It only tackles the more consensual part of the proposal, i.e.does not address the problem of intermediate files bundling yet.

[tra]

include/clang/Driver/Action.h

79

Nit: All-caps CUDA looks weird here. _Cuda may be better choice.
If you can shorten the prefix that would be nice, too. OK_ is already taken, though. OFK_?

include/clang/Driver/Compilation.h

42

This could be generalized into yet another toolchain kind and handled the same way as offload toolchains.

46

This could use a more descriptive name. ActiveOffloadMask?

51

Using std::multimap<OffloadKind, Toolchain> here would probably make specific_offload_kind_iterator unnecessary and would simplify toolchain lookup.

lib/Driver/Driver.cpp

406

llvm::any_of() could be useful here:

if (llvm::any_of(Inputs, [](auto &I) { return types::isCuda(I.first)})) {
  ...addOffloadDeviceToolchain();
}

419

TODO:

543

Get is somewhat misplaced here. Perhaps populate or initialize would work better.

[sfantao]

Address Art and Eric comments.

[sfantao]

Hi Art,

Thanks for the the review!

include/clang/Driver/Action.h

79

Ok, sure. I'm using OFK_ instead of OFFLOAD_ as suggested.

include/clang/Driver/Compilation.h

42

Ok, I created the OFK_Host kind to designate the host toolchain.

46

Using ActiveOffloadMask now.

51

Ok, that makes sense. I was trying to avoid the clients of this hook to iterate over a pair - something that a map of vectors would probably do without the need of the iterator.

In any case, I did as you say as it is a more clean solution.

lib/Driver/Driver.cpp

406

Ok, using any_of now.

419

Done!

543

I am using populate now.

[Hahnfeld]

Is there any progress on this one?

[Hahnfeld]

@tra any more comments on this?

[tra]

LGTM.

include/clang/Driver/Action.h

enum ActionClass {
     DsymutilJobClass,
     VerifyDebugInfoJobClass,
     VerifyPCHJobClass,
 
     JobClassFirst=PreprocessJobClass,
     JobClassLast=VerifyPCHJobClass
   };
 
+  // The offloading kind determines if this action is binded to a particular
+  // programming model. Each entry reserves one bit.
+  //
+  // FIXME: This is currently used to indicate that toolchains are used in a
+  // given programming as well, but will be used here as well once a generic
+  // offloading action is implemented.
+  enum OffloadKind {
+    OFFLOAD_None = 0x00,
+    OFFLOAD_CUDA = 0x01,
+  };
+
   static const char *getClassName(ActionClass AC);
 
 private:
   ActionClass Kind;
 
   /// The output type of this action.
   types::ID Type;
 

include/clang/Driver/Compilation.h

 /// invocation.
 class Compilation {
   /// The driver we were created by.
   const Driver &TheDriver;
 
   /// The default tool chain.
   const ToolChain &DefaultToolChain;
 
-  const ToolChain *CudaHostToolChain;
-  const ToolChain *CudaDeviceToolChain;
+  /// The tool chain of the offload host.
+  const ToolChain *OffloadHostToolChain;
+
+  /// The host offload kinds, it will be a mask of all the programming models
+  /// the host has to support.
+  unsigned OffloadHostKinds;
+
+  /// Array with the toolchains of offloading devices in the order they were
+  /// requested by the user.
+  typedef std::pair<const ToolChain *, Action::OffloadKind> OffloadToolChainTy;
+  SmallVector<OffloadToolChainTy, 4> OrderedOffloadingToolchains;
 
   /// The original (untranslated) input argument list.
   llvm::opt::InputArgList *Args;
 
   /// The driver translated arguments. Note that toolchains may perform their
   /// own argument translation.
   llvm::opt::DerivedArgList *TranslatedArgs;
 
public:
   Compilation(const Driver &D, const ToolChain &DefaultToolChain,
               llvm::opt::InputArgList *Args,
               llvm::opt::DerivedArgList *TranslatedArgs);
   ~Compilation();
 
   const Driver &getDriver() const { return TheDriver; }
 
   const ToolChain &getDefaultToolChain() const { return DefaultToolChain; }
-  const ToolChain *getCudaHostToolChain() const { return CudaHostToolChain; }
-  const ToolChain *getCudaDeviceToolChain() const {
-    return CudaDeviceToolChain;
+  const ToolChain *getOffloadingHostToolChain() const {
+    return OffloadHostToolChain;
+  }
+  unsigned isOffloadingHostKind(Action::OffloadKind Kind) const {
+    return OffloadHostKinds & Kind;
+  }
+
+  /// Iterator that visits device toolchains of a given kind.
+  template <Action::OffloadKind Kind>
+  class specific_offload_kind_iterator
+      : public llvm::iterator_adaptor_base<
+            specific_offload_kind_iterator<Kind>,
+            ArrayRef<OffloadToolChainTy>::const_iterator,
+            std::forward_iterator_tag, OffloadToolChainTy, ptrdiff_t,
+            OffloadToolChainTy, OffloadToolChainTy> {
+    ArrayRef<OffloadToolChainTy>::const_iterator End;
+
+    void SkipKinds() {
+      while (this->I != End && this->I->second != Kind)
+        ++this->I;
     }
 
-  void setCudaHostToolChain(const ToolChain *HostToolChain) {
-    CudaHostToolChain = HostToolChain;
+  public:
+    explicit specific_offload_kind_iterator(ArrayRef<OffloadToolChainTy> TCs)
+        : specific_offload_kind_iterator::iterator_adaptor_base(TCs.begin()),
+          End(TCs.end()) {
+      SkipKinds();
+    }
+
+    const ToolChain *operator*() const { return this->I->first; }
+    const ToolChain *operator->() const { return **this; }
+
+    specific_offload_kind_iterator &operator++() {
+      ++this->I;
+      SkipKinds();
+      return *this;
     }
-  void setCudaDeviceToolChain(const ToolChain *DeviceToolChain) {
-    CudaDeviceToolChain = DeviceToolChain;
+  };
+
+  template <Action::OffloadKind Kind>
+  llvm::iterator_range<specific_offload_kind_iterator<Kind>>
+  getOffloadDeviceToolChains() const {
+    return {specific_offload_kind_iterator<Kind>(OrderedOffloadingToolchains),
+            specific_offload_kind_iterator<Kind>(
+                llvm::makeArrayRef(OrderedOffloadingToolchains.end(), 0))};
+  }
+
+  // Return an offload device toolchain of the provided kind. Only one is
+  // expected to exist. If we can't match any toolchain, return nullptr.
+  template <Action::OffloadKind Kind>
+  const ToolChain *getSingleOffloadDeviceToolChain() const {
+    auto TCs = getOffloadDeviceToolChains<Kind>();
+
+    if (TCs.begin() != TCs.end()) {
+      assert(std::next(TCs.begin()) == TCs.end() &&
+             "More than one tool chain of the this kind exist.");
+      return *TCs.begin();
+    }
+    return nullptr;
+  }
+
+  void addOffloadDeviceToolChain(const ToolChain *DeviceToolChain,
+                                 Action::OffloadKind OffloadKind) {
+    // Update the host offload kind to also contain this kind.
+    OffloadHostKinds |= OffloadKind;
+    OrderedOffloadingToolchains.push_back(
+        std::make_pair(DeviceToolChain, OffloadKind));
   }
 
   const llvm::opt::InputArgList &getInputArgs() const { return *Args; }
 
   const llvm::opt::DerivedArgList &getArgs() const { return *TranslatedArgs; }
 
   llvm::opt::DerivedArgList &getArgs() { return *TranslatedArgs; }
 

include/clang/Driver/Driver.h

public:
 
   bool embedBitcodeEnabled() const { return BitcodeEmbed == EmbedBitcode; }
   bool embedBitcodeMarkerOnly() const { return BitcodeEmbed == EmbedMarker; }
 
   /// @}
   /// @name Primary Functionality
   /// @{
 
+  /// CreateOffloadingDeviceToolChains - create all the toolchains required to
+  /// support offloading devices given the programming models specified in the
+  /// current compilation. Also, update the host tool chain kind accordingly.
+  void CreateOffloadingDeviceToolChains(Compilation &C, InputList &Inputs);
+
   /// BuildCompilation - Construct a compilation object for a command
   /// line argument vector.
   ///
   /// \return A compilation, or 0 if none was built for the given
   /// argument vector. A null return value does not necessarily
   /// indicate an error condition, the diagnostics should be queried
   /// to determine if an error occurred.
   Compilation *BuildCompilation(ArrayRef<const char *> Args);

lib/Driver/Compilation.cpp

 
 using namespace clang::driver;
 using namespace clang;
 using namespace llvm::opt;
 
 Compilation::Compilation(const Driver &D, const ToolChain &_DefaultToolChain,
                          InputArgList *_Args, DerivedArgList *_TranslatedArgs)
     : TheDriver(D), DefaultToolChain(_DefaultToolChain),
-      CudaHostToolChain(&DefaultToolChain), CudaDeviceToolChain(nullptr),
+      OffloadHostToolChain(&DefaultToolChain), OffloadHostKinds(0u),
       Args(_Args), TranslatedArgs(_TranslatedArgs), Redirects(nullptr),
       ForDiagnostics(false) {}
 
 Compilation::~Compilation() {
   delete TranslatedArgs;
   delete Args;
 
   // Free any derived arg lists.

lib/Driver/Driver.cpp

void Driver::setLTOMode(const llvm::opt::ArgList &Args) {
 
   if (LTOMode == LTOK_Unknown) {
     assert(A);
     Diag(diag::err_drv_unsupported_option_argument) << A->getOption().getName()
                                                     << A->getValue();
   }
 }
 
+void Driver::CreateOffloadingDeviceToolChains(Compilation &C,
+                                              InputList &Inputs) {
+
+  //
+  // CUDA
+  //
+  // We need to generate a CUDA toolchain if any of the inputs has a CUDA type.
+  for (auto &I : Inputs)
+    // Have we founs a CUDA file? If so generate the toolchain.
+    if (types::isCuda(I.first)) {
+      const ToolChain &TC = getToolChain(
+          C.getInputArgs(),
+          llvm::Triple(C.getOffloadingHostToolChain()->getTriple().isArch64Bit()
+                           ? "nvptx64-nvidia-cuda"
+                           : "nvptx-nvidia-cuda"));
+      C.addOffloadDeviceToolChain(&TC, Action::OFFLOAD_CUDA);
+      break;
+    }
+
+  //
+  // Add support for other offloading programming models here.
+  //
+
+  return;
+}
+
 Compilation *Driver::BuildCompilation(ArrayRef<const char *> ArgList) {
   llvm::PrettyStackTraceString CrashInfo("Compilation construction");
 
   // FIXME: Handle environment options which affect driver behavior, somewhere
   // (client?). GCC_EXEC_PREFIX, LPATH, CC_PRINT_OPTIONS.
 
   if (char *env = ::getenv("COMPILER_PATH")) {
     StringRef CompilerPath = env;
Compilation *Driver::BuildCompilation(ArrayRef<const char *> ArgList) {
 
   // Owned by the host.
   const ToolChain &TC =
       getToolChain(*UArgs, computeTargetTriple(DefaultTargetTriple, *UArgs));
 
   // The compilation takes ownership of Args.
   Compilation *C = new Compilation(*this, TC, UArgs.release(), TranslatedArgs);
 
-  C->setCudaDeviceToolChain(
-      &getToolChain(C->getArgs(), llvm::Triple(TC.getTriple().isArch64Bit()
-                                                   ? "nvptx64-nvidia-cuda"
-                                                   : "nvptx-nvidia-cuda")));
   if (!HandleImmediateArgs(*C))
     return C;
 
   // Construct the list of inputs.
   InputList Inputs;
   BuildInputs(C->getDefaultToolChain(), *TranslatedArgs, Inputs);
 
+  // Get the toolchains for the offloading devices, if any.
+  CreateOffloadingDeviceToolChains(*C, Inputs);
+
   // Construct the list of abstract actions to perform for this compilation. On
   // MachO targets this uses the driver-driver and universal actions.
   if (TC.getTriple().isOSBinFormatMachO())
     BuildUniversalActions(*C, C->getDefaultToolChain(), Inputs);
   else
     BuildActions(*C, C->getArgs(), Inputs, C->getActions());
 
   if (CCCPrintPhases) {
if (GpuArchList.empty())
     GpuArchList.push_back("sm_20");
 
   // Replicate inputs for each GPU architecture.
   Driver::InputList CudaDeviceInputs;
   for (unsigned I = 0, E = GpuArchList.size(); I != E; ++I)
     CudaDeviceInputs.push_back(std::make_pair(types::TY_CUDA_DEVICE, InputArg));
 
   // Build actions for all device inputs.
-  assert(C.getCudaDeviceToolChain() &&
+  assert(C.getSingleOffloadDeviceToolChain<Action::OFFLOAD_CUDA>() &&
          "Missing toolchain for device-side compilation.");
   ActionList CudaDeviceActions;
   C.getDriver().BuildActions(C, Args, CudaDeviceInputs, CudaDeviceActions);
   assert(GpuArchList.size() == CudaDeviceActions.size() &&
          "Failed to create actions for all devices");
 
   // Check whether any of device actions stopped before they could generate PTX.
   bool PartialCompilation =
if (const BindArchAction *BAA = dyn_cast<BindArchAction>(A)) {
     return BuildJobsForAction(C, *BAA->input_begin(), TC, ArchName, AtTopLevel,
                               MultipleArchs, LinkingOutput, CachedResults);
   }
 
   if (const CudaDeviceAction *CDA = dyn_cast<CudaDeviceAction>(A)) {
     // Initial processing of CudaDeviceAction carries host params.
     // Call BuildJobsForAction() again, now with correct device parameters.
     InputInfo II = BuildJobsForAction(
-        C, *CDA->input_begin(), C.getCudaDeviceToolChain(),
+        C, *CDA->input_begin(),
+        C.getSingleOffloadDeviceToolChain<Action::OFFLOAD_CUDA>(),
         CDA->getGpuArchName(), CDA->isAtTopLevel(), /*MultipleArchs=*/true,
         LinkingOutput, CachedResults);
     // Currently II's Action is *CDA->input_begin().  Set it to CDA instead, so
     // that one can retrieve II's GPU arch.
     II.setAction(A);
     return II;
   }
 

lib/Driver/Tools.cpp

void Clang::ConstructJob(Compilation &C, const JobAction &JA,
   CmdArgs.push_back(Args.MakeArgString(TripleStr));
 
   const ToolChain *AuxToolChain = nullptr;
   if (IsCuda) {
     // FIXME: We need a (better) way to pass information about
     // particular compilation pass we're constructing here. For now we
     // can check which toolchain we're using and pick the other one to
     // extract the triple.
-    if (&getToolChain() == C.getCudaDeviceToolChain())
-      AuxToolChain = C.getCudaHostToolChain();
-    else if (&getToolChain() == C.getCudaHostToolChain())
-      AuxToolChain = C.getCudaDeviceToolChain();
+    if (&getToolChain() ==
+        C.getSingleOffloadDeviceToolChain<Action::OFFLOAD_CUDA>())
+      AuxToolChain = C.getOffloadingHostToolChain();
+    else if (&getToolChain() == C.getOffloadingHostToolChain())
+      AuxToolChain = C.getSingleOffloadDeviceToolChain<Action::OFFLOAD_CUDA>();
     else
       llvm_unreachable("Can't figure out CUDA compilation mode.");
     assert(AuxToolChain != nullptr && "No aux toolchain.");
     CmdArgs.push_back("-aux-triple");
     CmdArgs.push_back(Args.MakeArgString(AuxToolChain->getTriple().str()));
     CmdArgs.push_back("-fcuda-target-overloads");
     CmdArgs.push_back("-fcuda-disable-target-call-checks");
   }