[Driver] Use clang-offload-deps tool for generating device dependencies

clang/include/clang/Driver/Action.h

@@ -74,6 +74,7 @@ class Action {
     OffloadBundlingJobClass,
     OffloadUnbundlingJobClass,
     OffloadWrapperJobClass,
+    OffloadDepsJobClass,
     SPIRVTranslatorJobClass,
     SPIRCheckJobClass,
     SYCLPostLinkJobClass,
@@ -648,6 +649,60 @@ class OffloadWrapperJobAction : public JobAction {
   }
 };
 
+class OffloadDepsJobAction final : public JobAction {
+  void anchor() override;
+
+public:
+  /// Type that provides information about the actions that depend on this
+  /// offload deps action.
+  struct DependentActionInfo final {
+    /// The tool chain of the dependent action.
+    const ToolChain *DependentToolChain = nullptr;
+
+    /// The bound architecture of the dependent action.
+    StringRef DependentBoundArch;
+
+    /// The offload kind of the dependent action.
+    const OffloadKind DependentOffloadKind = OFK_None;
+
+    DependentActionInfo(const ToolChain *DependentToolChain,
+                        StringRef DependentBoundArch,
+                        const OffloadKind DependentOffloadKind)
+        : DependentToolChain(DependentToolChain),
+          DependentBoundArch(DependentBoundArch),
+          DependentOffloadKind(DependentOffloadKind) {}
+  };
+
+private:
+  /// The host offloading toolchain that should be used with the action.
+  const ToolChain *HostTC = nullptr;
+
+  /// Container that keeps information about each dependence of this deps
+  /// action.
+  SmallVector<DependentActionInfo, 6> DependentActionInfoArray;
+
+public:
+  OffloadDepsJobAction(const OffloadAction::HostDependence &HDep,
+                       types::ID Type);
+
+  /// Register information about a dependent action.
+  void registerDependentActionInfo(const ToolChain *TC, StringRef BoundArch,
+                                   OffloadKind Kind) {
+    DependentActionInfoArray.push_back({TC, BoundArch, Kind});
+  }
+
+  /// Return the information about all depending actions.
+  ArrayRef<DependentActionInfo> getDependentActionsInfo() const {
+    return DependentActionInfoArray;
+  }
+
+  const ToolChain *getHostTC() const { return HostTC; }
+
+  static bool classof(const Action *A) {
+    return A->getKind() == OffloadDepsJobClass;
+  }
+};
+
 class SPIRVTranslatorJobAction : public JobAction {
   void anchor() override;
 

clang/include/clang/Driver/ToolChain.h

@@ -144,6 +144,7 @@ class ToolChain {
   mutable std::unique_ptr<Tool> IfsMerge;
   mutable std::unique_ptr<Tool> OffloadBundler;
   mutable std::unique_ptr<Tool> OffloadWrapper;
+  mutable std::unique_ptr<Tool> OffloadDeps;
   mutable std::unique_ptr<Tool> SPIRVTranslator;
   mutable std::unique_ptr<Tool> SPIRCheck;
   mutable std::unique_ptr<Tool> SYCLPostLink;
@@ -160,6 +161,7 @@ class ToolChain {
   Tool *getClangAs() const;
   Tool *getOffloadBundler() const;
   Tool *getOffloadWrapper() const;
+  Tool *getOffloadDeps() const;
   Tool *getSPIRVTranslator() const;
   Tool *getSPIRCheck() const;
   Tool *getSYCLPostLink() const;

clang/lib/Driver/Action.cpp

@@ -43,6 +43,8 @@ const char *Action::getClassName(ActionClass AC) {
     return "clang-offload-unbundler";
   case OffloadWrapperJobClass:
     return "clang-offload-wrapper";
+  case OffloadDepsJobClass:
+    return "clang-offload-deps";
   case SPIRVTranslatorJobClass:
     return "llvm-spirv";
   case SPIRCheckJobClass:
@@ -69,6 +71,9 @@ void Action::propagateDeviceOffloadInfo(OffloadKind OKind, const char *OArch) {
   // Unbundling actions use the host kinds.
   if (Kind == OffloadUnbundlingJobClass)
     return;
+  // Deps job uses the host kinds.
+  if (Kind == OffloadDepsJobClass)
+    return;
 
   assert((OffloadingDeviceKind == OKind || OffloadingDeviceKind == OFK_None) &&
          "Setting device kind to a different device??");
@@ -444,6 +449,18 @@ OffloadWrapperJobAction::OffloadWrapperJobAction(Action *Input,
                                                  types::ID Type)
     : JobAction(OffloadWrapperJobClass, Input, Type) {}
 
+void OffloadDepsJobAction::anchor() {}
+
+OffloadDepsJobAction::OffloadDepsJobAction(
+    const OffloadAction::HostDependence &HDep, types::ID Type)
+    : JobAction(OffloadDepsJobClass, HDep.getAction(), Type),
+      HostTC(HDep.getToolChain()) {
+  OffloadingArch = HDep.getBoundArch();
+  ActiveOffloadKindMask = HDep.getOffloadKinds();
+  HDep.getAction()->propagateHostOffloadInfo(HDep.getOffloadKinds(),
+                                             HDep.getBoundArch());
+}
+
 void SPIRVTranslatorJobAction::anchor() {}
 
 SPIRVTranslatorJobAction::SPIRVTranslatorJobAction(Action *Input,

clang/lib/Driver/Driver.cpp

@@ -2886,6 +2886,9 @@ class OffloadingActionBuilder final {
     /// Append linker actions generated by the builder.
     virtual void appendLinkDependences(OffloadAction::DeviceDependences &DA) {}
 
+    /// Append linker actions generated by the builder.
+    virtual void addDeviceLinkDependencies(OffloadDepsJobAction *DA) {}
+
     /// Initialize the builder. Return true if any initialization errors are
     /// found.
     virtual bool initialize() { return false; }
@@ -3607,6 +3610,24 @@ class OffloadingActionBuilder final {
 
     void appendLinkDependences(OffloadAction::DeviceDependences &DA) override {}
 
+    void addDeviceLinkDependencies(OffloadDepsJobAction *DA) override {
+      for (unsigned I = 0; I < ToolChains.size(); ++I) {
+        // Register dependent toolchain.
+        DA->registerDependentActionInfo(
+            ToolChains[I], /*BoundArch=*/StringRef(), Action::OFK_OpenMP);
+
+        if (!ToolChains[I]->getTriple().isSPIR()) {
+          // Create object from the deps bitcode.
+          auto *BA = C.MakeAction<BackendJobAction>(DA, types::TY_PP_Asm);
+          auto *AA = C.MakeAction<AssembleJobAction>(BA, types::TY_Object);
+
+          // Add deps object to linker inputs.
+          DeviceLinkerInputs[I].push_back(AA);
+        } else
+          DeviceLinkerInputs[I].push_back(DA);
+      }
+    }
+
     bool initialize() override {
       // Get the OpenMP toolchains. If we don't get any, the action builder will
       // know there is nothing to do related to OpenMP offloading.
@@ -4340,6 +4361,17 @@ class OffloadingActionBuilder final {
       }
     }
 
+    void addDeviceLinkDependencies(OffloadDepsJobAction *DA) override {
+      for (unsigned I = 0; I < ToolChains.size(); ++I) {
+        // Register dependent toolchain.
+        DA->registerDependentActionInfo(
+            ToolChains[I], /*BoundArch=*/StringRef(), Action::OFK_SYCL);
+
+        // Add deps output to linker inputs.
+        DeviceLinkerInputs[I].push_back(DA);
+      }
+    }
+
     /// Initialize the GPU architecture list from arguments - this populates `GpuArchList` from
     /// `--cuda-gpu-arch` flags. Only relevant if compiling to CUDA. Return true if any
     /// initialization errors are found.
@@ -4742,6 +4774,32 @@ class OffloadingActionBuilder final {
     return false;
   }
 
+  /// Create link job from the given host inputs and feed the result to offload
+  /// deps job which fetches device dependencies from the linked host image.
+  /// Offload deps output is then forwarded to active device action builders so
+  /// they can add it to the device linker inputs.
+  void addDeviceLinkDependenciesFromHost(ActionList &LinkerInputs) {
+    // Link image for reading dependencies from it.
+    auto *LA = C.MakeAction<LinkJobAction>(LinkerInputs, types::TY_Image);
+
+    // Calculate all the offload kinds used in the current compilation.
+    unsigned ActiveOffloadKinds = 0u;
+    for (auto &I : InputArgToOffloadKindMap)
+      ActiveOffloadKinds |= I.second;
+
+    OffloadAction::HostDependence HDep(
+        *LA, *C.getSingleOffloadToolChain<Action::OFK_Host>(),
+        /*BoundArch*/ nullptr, ActiveOffloadKinds);
+
+    auto *DA = C.MakeAction<OffloadDepsJobAction>(HDep, types::TY_LLVM_BC);
+
+    for (auto *SB : SpecializedBuilders) {
+      if (!SB->isValid())
+        continue;
+      SB->addDeviceLinkDependencies(DA);
+    }
+  }
+
   void makeHostLinkAction(ActionList &LinkerInputs) {
     // Build a list of device linking actions.
     ActionList DeviceAL;
@@ -5066,6 +5124,11 @@ void Driver::BuildActions(Compilation &C, DerivedArgList &Args,
 
   OffloadBuilder.appendTopLevelLinkAction(Actions);
 
+  // With static fat archives we need to create additional steps for
+  // generating dependence objects for device link actions.
+  if (!LinkerInputs.empty() && C.getDriver().getOffloadStaticLibSeen())
+    OffloadBuilder.addDeviceLinkDependenciesFromHost(LinkerInputs);
+
   // Go through all of the args, and create a Linker specific argument list.
   // When dealing with fat static archives each archive is individually
   // unbundled.
@@ -5967,7 +6030,10 @@ InputInfo Driver::BuildJobsForActionNoCache(
   const JobAction *JA = cast<JobAction>(A);
   ActionList CollapsedOffloadActions;
 
-  ToolSelector TS(JA, *TC, C, isSaveTempsEnabled(),
+  auto *DA = dyn_cast<OffloadDepsJobAction>(JA);
+  const ToolChain *JATC = DA ? DA->getHostTC() : TC;
+
+  ToolSelector TS(JA, *JATC, C, isSaveTempsEnabled(),
                   embedBitcodeInObject() && !isUsingLTO());
   const Tool *T = TS.getTool(Inputs, CollapsedOffloadActions);
 
@@ -5996,8 +6062,9 @@ InputInfo Driver::BuildJobsForActionNoCache(
     bool SubJobAtTopLevel =
         AtTopLevel && (isa<DsymutilJobAction>(A) || isa<VerifyJobAction>(A));
     InputInfos.push_back(BuildJobsForAction(
-        C, Input, TC, BoundArch, SubJobAtTopLevel, MultipleArchs, LinkingOutput,
-        CachedResults, A->getOffloadingDeviceKind()));
+        C, Input, JATC, DA ? DA->getOffloadingArch() : BoundArch,
+        SubJobAtTopLevel, MultipleArchs, LinkingOutput, CachedResults,
+        A->getOffloadingDeviceKind()));
   }
   // Check if we are in sub-work for preprocessing for host side. If so we will
   // add another job to print information to terminal later.
@@ -6166,6 +6233,46 @@ InputInfo Driver::BuildJobsForActionNoCache(
              "Result does not exist??");
       Result = CachedResults[ActionTC];
     }
+  } else if (auto *DA = dyn_cast<OffloadDepsJobAction>(JA)) {
+    for (auto &DI : DA->getDependentActionsInfo()) {
+      assert(DI.DependentOffloadKind != Action::OFK_None &&
+             "Deps job with no offloading");
+
+      std::string OffloadingPrefix = Action::GetOffloadingFileNamePrefix(
+          DI.DependentOffloadKind,
+          DI.DependentToolChain->getTriple().normalize(),
+          /*CreatePrefixForHost=*/true);
+      auto CurI = InputInfo(
+          DA,
+          GetNamedOutputPath(C, *DA, BaseInput, DI.DependentBoundArch,
+                             /*AtTopLevel=*/false,
+                             MultipleArchs ||
+                                 DI.DependentOffloadKind == Action::OFK_HIP,
+                             OffloadingPrefix),
+          BaseInput);
+      // Save the result.
+      UnbundlingResults.push_back(CurI);
+
+      // Get the unique string identifier for this dependence and cache the
+      // result.
+      StringRef Arch = TargetDeviceOffloadKind == Action::OFK_HIP
+                           ? DI.DependentOffloadKind == Action::OFK_Host
+                                 ? StringRef()
+                                 : DI.DependentBoundArch
+                           : BoundArch;
+
+      CachedResults[{A, GetTriplePlusArchString(DI.DependentToolChain, Arch,
+                                                DI.DependentOffloadKind)}] =
+          CurI;
+    }
+
+    // Now that we have all the results generated, select the one that should be
+    // returned for the current depending action.
+    std::pair<const Action *, std::string> ActionTC = {
+        A, GetTriplePlusArchString(TC, BoundArch, TargetDeviceOffloadKind)};
+    auto It = CachedResults.find(ActionTC);
+    assert(It != CachedResults.end() && "Result does not exist??");
+    Result = It->second;
   } else if (JA->getType() == types::TY_Nothing)
     Result = InputInfo(A, BaseInput);
   else {

clang/lib/Driver/ToolChain.cpp

@@ -328,6 +328,12 @@ Tool *ToolChain::getOffloadWrapper() const {
   return OffloadWrapper.get();
 }
 
+Tool *ToolChain::getOffloadDeps() const {
+  if (!OffloadDeps)
+    OffloadDeps.reset(new tools::OffloadDeps(*this));
+  return OffloadDeps.get();
+}
+
 Tool *ToolChain::getSPIRVTranslator() const {
   if (!SPIRVTranslator)
     SPIRVTranslator.reset(new tools::SPIRVTranslator(*this));
@@ -403,6 +409,9 @@ Tool *ToolChain::getTool(Action::ActionClass AC) const {
   case Action::OffloadWrapperJobClass:
     return getOffloadWrapper();
 
+  case Action::OffloadDepsJobClass:
+    return getOffloadDeps();
+
   case Action::SPIRVTranslatorJobClass:
     return getSPIRVTranslator();
 

clang/lib/Driver/ToolChains/Clang.cpp

@@ -7988,6 +7988,71 @@ void OffloadWrapper::ConstructJob(Compilation &C, const JobAction &JA,
       CmdArgs, Inputs));
 }
 
+// Begin OffloadDeps
+
+void OffloadDeps::constructJob(Compilation &C, const JobAction &JA,
+                               ArrayRef<InputInfo> Outputs,
+                               ArrayRef<InputInfo> Inputs,
+                               const llvm::opt::ArgList &TCArgs,
+                               const char *LinkingOutput) const {
+  auto &DA = cast<OffloadDepsJobAction>(JA);
+
+  ArgStringList CmdArgs;
+
+  // Get the targets.
+  SmallString<128> Targets{"-targets="};
+  auto DepInfo = DA.getDependentActionsInfo();
+  for (unsigned I = 0; I < DepInfo.size(); ++I) {
+    auto &Dep = DepInfo[I];
+    if (I)
+      Targets += ',';
+    Targets += Action::GetOffloadKindName(Dep.DependentOffloadKind);
+    Targets += '-';
+    Targets += Dep.DependentToolChain->getTriple().normalize();
+    if (Dep.DependentOffloadKind == Action::OFK_HIP &&
+        !Dep.DependentBoundArch.empty()) {
+      Targets += '-';
+      Targets += Dep.DependentBoundArch;
+    }
+  }
+  CmdArgs.push_back(TCArgs.MakeArgString(Targets));
+
+  // Prepare outputs.
+  SmallString<128> Outs{"-outputs="};
+  for (unsigned I = 0; I < Outputs.size(); ++I) {
+    if (I)
+      Outs += ',';
+    Outs += DepInfo[I].DependentToolChain->getInputFilename(Outputs[I]);
+  }
+  CmdArgs.push_back(TCArgs.MakeArgString(Outs));
+
+  // Add input file.
+  CmdArgs.push_back(Inputs.front().getFilename());
+
+  // All the inputs are encoded as commands.
+  C.addCommand(std::make_unique<Command>(
+      JA, *this, ResponseFileSupport::None(),
+      TCArgs.MakeArgString(getToolChain().GetProgramPath(getShortName())),
+      CmdArgs, None, Outputs));
+}
+
+void OffloadDeps::ConstructJob(Compilation &C, const JobAction &JA,
+                               const InputInfo &Output,
+                               const InputInfoList &Inputs,
+                               const llvm::opt::ArgList &TCArgs,
+                               const char *LinkingOutput) const {
+  constructJob(C, JA, Output, Inputs, TCArgs, LinkingOutput);
+}
+
+void OffloadDeps::ConstructJobMultipleOutputs(Compilation &C,
+                                              const JobAction &JA,
+                                              const InputInfoList &Outputs,
+                                              const InputInfoList &Inputs,
+                                              const llvm::opt::ArgList &TCArgs,
+                                              const char *LinkingOutput) const {
+  constructJob(C, JA, Outputs, Inputs, TCArgs, LinkingOutput);
+}
+
 // Begin SPIRVTranslator
 
 void SPIRVTranslator::ConstructJob(Compilation &C, const JobAction &JA,