3 files changed, 77 insertions, 4 deletions

diff --git a/clang/lib/Driver/Action.cpp b/clang/lib/Driver/Action.cpp
index a351429ad7b..90bc149fd33 100644
--- a/clang/lib/Driver/Action.cpp
+++ b/clang/lib/Driver/Action.cpp
@@ -36,6 +36,8 @@ const char *Action::getClassName(ActionClass AC) {
   case DsymutilJobClass: return "dsymutil";
   case VerifyDebugInfoJobClass: return "verify-debug-info";
   case VerifyPCHJobClass: return "verify-pch";
+  case OffloadBundlingJobClass:
+    return "clang-offload-bundler";
   }
 
   llvm_unreachable("invalid class");
@@ -346,3 +348,8 @@ void VerifyPCHJobAction::anchor() {}
 
 VerifyPCHJobAction::VerifyPCHJobAction(Action *Input, types::ID Type)
     : VerifyJobAction(VerifyPCHJobClass, Input, Type) {}
+
+void OffloadBundlingJobAction::anchor() {}
+
+OffloadBundlingJobAction::OffloadBundlingJobAction(ActionList &Inputs)
+    : JobAction(OffloadBundlingJobClass, Inputs, Inputs.front()->getType()) {}
diff --git a/clang/lib/Driver/Driver.cpp b/clang/lib/Driver/Driver.cpp
index fe9605279f3..8d96284a69b 100644
--- a/clang/lib/Driver/Driver.cpp
+++ b/clang/lib/Driver/Driver.cpp
@@ -1568,6 +1568,9 @@ class OffloadingActionBuilder final {
     /// found.
     virtual bool initialize() { return false; }
 
+    /// Return true if the builder can use bundling/unbundling.
+    virtual bool canUseBundlerUnbundler() const { return false; }
+
     /// Return true if this builder is valid. We have a valid builder if we have
     /// associated device tool chains.
     bool isValid() { return !ToolChains.empty(); }
@@ -1911,6 +1914,26 @@ class OffloadingActionBuilder final {
       return ABRT_Success;
     }
 
+    void appendTopLevelActions(ActionList &AL) override {
+      if (OpenMPDeviceActions.empty())
+        return;
+
+      // We should always have an action for each input.
+      assert(OpenMPDeviceActions.size() == ToolChains.size() &&
+             "Number of OpenMP actions and toolchains do not match.");
+
+      // Append all device actions followed by the proper offload action.
+      auto TI = ToolChains.begin();
+      for (auto *A : OpenMPDeviceActions) {
+        OffloadAction::DeviceDependences Dep;
+        Dep.add(*A, **TI, /*BoundArch=*/nullptr, Action::OFK_OpenMP);
+        AL.push_back(C.MakeAction<OffloadAction>(Dep, A->getType()));
+        ++TI;
+      }
+      // We no longer need the action stored in this builder.
+      OpenMPDeviceActions.clear();
+    }
+
     void appendLinkDependences(OffloadAction::DeviceDependences &DA) override {
       assert(ToolChains.size() == DeviceLinkerInputs.size() &&
              "Toolchains and linker inputs sizes do not match.");
@@ -1937,6 +1960,11 @@ class OffloadingActionBuilder final {
       DeviceLinkerInputs.resize(ToolChains.size());
       return false;
     }
+
+    bool canUseBundlerUnbundler() const override {
+      // OpenMP should use bundled files whenever possible.
+      return true;
+    }
   };
 
   ///
@@ -1946,6 +1974,9 @@ class OffloadingActionBuilder final {
   /// Specialized builders being used by this offloading action builder.
   SmallVector<DeviceActionBuilder *, 4> SpecializedBuilders;
 
+  /// Flag set to true if all valid builders allow file bundling/unbundling.
+  bool CanUseBundler;
+
 public:
   OffloadingActionBuilder(Compilation &C, DerivedArgList &Args,
                           const Driver::InputList &Inputs)
@@ -1964,9 +1995,22 @@ public:
     // TODO: Build other specialized builders here.
     //
 
-    // Initialize all the builders, keeping track of errors.
-    for (auto *SB : SpecializedBuilders)
+    // Initialize all the builders, keeping track of errors. If all valid
+    // builders agree that we can use bundling, set the flag to true.
+    unsigned ValidBuilders = 0u;
+    unsigned ValidBuildersSupportingBundling = 0u;
+    for (auto *SB : SpecializedBuilders) {
       IsValid = IsValid && !SB->initialize();
+
+      // Update the counters if the builder is valid.
+      if (SB->isValid()) {
+        ++ValidBuilders;
+        if (SB->canUseBundlerUnbundler())
+          ++ValidBuildersSupportingBundling;
+      }
+    }
+    CanUseBundler =
+        ValidBuilders && ValidBuilders == ValidBuildersSupportingBundling;
   }
 
   ~OffloadingActionBuilder() {
@@ -2066,15 +2110,33 @@ public:
     return false;
   }
 
-  /// Add the offloading top level actions to the provided action list.
+  /// 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) {
+    // Get the device actions to be appended.
+    ActionList OffloadAL;
     for (auto *SB : SpecializedBuilders) {
       if (!SB->isValid())
         continue;
-      SB->appendTopLevelActions(AL);
+      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.
+    if (CanUseBundler && !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);
+      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)
diff --git a/clang/lib/Driver/ToolChain.cpp b/clang/lib/Driver/ToolChain.cpp
index 90232659659..a82ebfe6d98 100644
--- a/clang/lib/Driver/ToolChain.cpp
+++ b/clang/lib/Driver/ToolChain.cpp
@@ -263,6 +263,10 @@ Tool *ToolChain::getTool(Action::ActionClass AC) const {
   case Action::VerifyPCHJobClass:
   case Action::BackendJobClass:
     return getClang();
+
+  case Action::OffloadBundlingJobClass:
+    // FIXME: Add a tool for the bundling actions.
+    return nullptr;
   }
 
   llvm_unreachable("Invalid tool kind.");