[Offloading] Embed the target features in the OffloadBinary

[jhuber6]

Summary

The target features are necessary for correctly compiling most programs
in LTO mode. Currently, these are derived in clang at link time and
passed as an arguemnt to the linker wrapper. This is problematic because
it requires knowing the required toolchain at link time, which should
not be necessry. Instead, these features should be embedded into the
offloading binary so we can unify them in the linker wrapper for LTO.
This also required changing the offload packager to interpret multiple
arguments as concatenation with a comma. This is so we can still use the
, separator for the argument list.

Depends on D127246

[tra]

clang/lib/Driver/ToolChains/Clang.cpp

8321–8327

I think this would look cleaner if expressed along these lines:

SmallVector<StringRef> Parts = {
  "--image=file=" + File,
  "triple=" + TC->getTripleString() ,
  "arch=" + Arch,
  "kind=" + Action::GetOffloadKindName(OffloadAction->getOffloadingDeviceKind()),
};

// Add "-feature=foo" arguments to `Parts` here.

llvm::join(",", Parts)

8330

This makes a couple of implicit assumptions that we should probably make more explicit.

• assert(!FeatureArgs.empty()), otherwise we may end up passing an empty feature=.
• that it will be appended to the comma-separated list, which is not obvious until we actually do it later.

clang/test/Driver/openmp-offload-gpu-new.c

122

This should probably be a bit more specific/verbose. E.g. I'd want to make sure that feature= is part of the --image argument and that ptx belongs to it and is not part of some other argument (or even a file name extension).

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

81–85

This looks like a homegrown version of getLastArg(). I wonder if there's a way to use the standard LLVM option parser for that.

[jhuber6]

Thanks for the comments, I'll try to address them.

clang/lib/Driver/ToolChains/Clang.cpp

8321–8327

Yeah, it would probably be best to organize something like this. Though we would either need to make them strings, use a string saver, or store a Twine (which probably isn't recommended). I'll try to clean it up like this.

8330

I think I should add some appropriate logic to not add this string and empty comma if the features are empty. I'll try to include that in a cleanup here.

clang/test/Driver/openmp-offload-gpu-new.c

122

Sure, I was just hesitant to make it super specific since the specific feature will change depending on the CUDA installation (or lack thereof).

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

81–85

It's copy-pasted from unifyTargetFeatures in Clang, which does the same thing. I may be able to simplify this, but for now I just wanted to copy something that ostensibly worked. Using LLVM option parsing stuff is another thing on my list. I want to pass arguments to the linker wrapper via some --offload-opt similar to --plugin-opt for LTO. Since my ultimate goal is to incorporate this into a linker completely.

[tra]

clang/test/Driver/openmp-offload-gpu-new.c

122

I'm all for concise tests as long as they:

• express what you want to verify in a way that the reader would be able to understand w/o having to look at the source code.
• are reasonably robust and do not produce false positives. .* wildcards make it very easy to match things unintentionally. Their use should be carefully restricted.

I wish FileCheck would have some sort of nested match check. E.g.

; CHECK : [[CANDIDATE:--option.*?\s+]]
; CHECK:  [[SUBMATCH1: match something within [[CANDIDATE]]]]
; CHECK:  [[SUBMATCH2: match something within [[SUBMATCH1]]]]

[jhuber6]

Adjust how we generate arguments.

[tra]

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

735

// Record the index of the last occurence of the feature.

741

// Populate UnifiedFeatures only with last mentions of specific feature.

743–748

auto FeatureName = Name.drop_front(1); 
if (!LastOpt.count(FeatureName)) // could be just assert(LastOpt.count(FeatureName))
  continue
if (LastOpt[FeatureName] != I) // could be merged with the `if` above
  continue;

[jhuber6]

Does this approach work? I'm just using the reverse iterator and only adding the argument if it hasn't been seen yet.

[jhuber6]

ping.

[tra]

LGTM with a nit.

clang/test/Driver/openmp-offload-gpu-new.c

122

Nit: ^^^ looks like this one slipped through the cracks and wasn't addressed.

clang/lib/Driver/ToolChains/Clang.cpp

for (const InputInfo &Input : Inputs) {
     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);
+    StringRef Kind =
+      Action::GetOffloadKindName(OffloadAction->getOffloadingDeviceKind());
 
-    CmdArgs.push_back(Args.MakeArgString(
-        "--image=file=" + File + "," + "triple=" + TC->getTripleString() + "," +
-        "arch=" + Arch + "," + "kind=" +
-        Action::GetOffloadKindName(OffloadAction->getOffloadingDeviceKind())));
+    ArgStringList Features;
+    SmallVector<StringRef> FeatureArgs;
+    getTargetFeatures(TC->getDriver(), TC->getTriple(), Args, Features, false);
+    llvm::copy_if(Features, std::back_inserter(FeatureArgs),
+                  [](StringRef Arg) { return !Arg.startswith("-target"); });
+
+    SmallVector<std::string> Parts{
+        "file=" + File.str(),
+        "triple=" + TC->getTripleString(),
+        "arch=" + Arch.str(),
+        "kind=" + Kind.str(),
+    };
+
+    if (TC->getDriver().isUsingLTO(/* IsOffload */ true))
+      for (StringRef Feature : FeatureArgs)
+        Parts.emplace_back("feature=" + Feature.str());
+
+    CmdArgs.push_back(Args.MakeArgString("--image=" + llvm::join(Parts, ",")));
   }
 
   C.addCommand(std::make_unique<Command>(
       JA, *this, ResponseFileSupport::None(),
       Args.MakeArgString(getToolChain().GetProgramPath(getShortName())),
       CmdArgs, Inputs, Output));
 }
 
 void LinkerWrapper::ConstructJob(Compilation &C, const JobAction &JA,
                                  const InputInfo &Output,
                                  const InputInfoList &Inputs,
                                  const ArgList &Args,
for (auto &I : llvm::make_range(OpenMPTCRange.first, OpenMPTCRange.second)) {
 
     for (StringRef LibName : BCLibs)
       CmdArgs.push_back(Args.MakeArgString(
           "-target-library=" + Action::GetOffloadKindName(Action::OFK_OpenMP) +
           "-" + TC->getTripleString() + "-" + Arch + "=" + LibName));
   }
 
   if (D.isUsingLTO(/* IsOffload */ true)) {
-    // Pass in target features for each toolchain.
-    for (auto &I :
-         llvm::make_range(OpenMPTCRange.first, OpenMPTCRange.second)) {
-      const ToolChain *TC = I.second;
-      const ArgList &TCArgs = C.getArgsForToolChain(TC, "", Action::OFK_OpenMP);
-      ArgStringList FeatureArgs;
-      TC->addClangTargetOptions(TCArgs, FeatureArgs, Action::OFK_OpenMP);
-      auto FeatureIt = llvm::find(FeatureArgs, "-target-feature");
-      if (FeatureIt != FeatureArgs.end())
-        CmdArgs.push_back(
-            Args.MakeArgString("-target-feature=" + TC->getTripleString() +
-                               "=" + *(FeatureIt + 1)));
-    }
-
     // Pass in the optimization level to use for LTO.
     if (const Arg *A = Args.getLastArg(options::OPT_O_Group)) {
       StringRef OOpt;
       if (A->getOption().matches(options::OPT_O4) ||
           A->getOption().matches(options::OPT_Ofast))
         OOpt = "3";
       else if (A->getOption().matches(options::OPT_O)) {
         OOpt = A->getValue();

clang/test/Driver/openmp-offload-gpu-new.c

 
 // CHECK-LTO-LIBRARY: {{.*}}-lomptarget{{.*}}-lomptarget.devicertl
 
 // RUN:   %clang -### --target=x86_64-unknown-linux-gnu -fopenmp --offload-arch=sm_52 -nogpulib \
 // RUN:     -Xoffload-linker a -Xoffload-linker-nvptx64-nvidia-cuda b -Xoffload-linker-nvptx64 c \
 // RUN:     %s 2>&1 | FileCheck --check-prefix=CHECK-XLINKER %s
 
 // CHECK-XLINKER: -device-linker=a{{.*}}-device-linker=nvptx64-nvidia-cuda=b{{.*}}-device-linker=nvptx64-nvidia-cuda=c{{.*}}--
+
+// RUN:   %clang -### --target=x86_64-unknown-linux-gnu -fopenmp --offload-arch=sm_52 -nogpulib \
+// RUN:     -foffload-lto %s 2>&1 | FileCheck --check-prefix=CHECK-LTO-FEATURES %s
+
+// CHECK-LTO-FEATURES: clang-offload-packager{{.*}}--image={{.*}}feature=+ptx{{[0-9]+}}

clang/tools/clang-linker-wrapper/ClangLinkerWrapper.cpp

 // will be hidden.
 static cl::OptionCategory
     ClangLinkerWrapperCategory("clang-linker-wrapper options");
 
 static cl::opt<std::string> LinkerUserPath("linker-path", cl::Required,
                                            cl::desc("Path of linker binary"),
                                            cl::cat(ClangLinkerWrapperCategory));
 
-static cl::opt<std::string>
-    TargetFeatures("target-feature", cl::desc("Target features for triple"),
-                   cl::cat(ClangLinkerWrapperCategory));
-
 static cl::opt<std::string> OptLevel("opt-level",
                                      cl::desc("Optimization level for LTO"),
                                      cl::init("O2"),
                                      cl::cat(ClangLinkerWrapperCategory));
 
 static cl::list<std::string>
     BitcodeLibraries("target-library",
                      cl::desc("Path for the target bitcode library"),
case DS_Note:
     WithColor::note(errs(), LinkerExecutable) << ErrStorage << "\n";
     break;
   case DS_Remark:
     WithColor::remark(errs()) << ErrStorage << "\n";
     break;
   }
 }
 
-// Get the target features passed in from the driver as <triple>=<features>.
-std::vector<std::string> getTargetFeatures(const Triple &TheTriple) {
-  std::vector<std::string> Features;
-  auto TargetAndFeatures = StringRef(TargetFeatures).split('=');
-  if (TargetAndFeatures.first != TheTriple.getTriple())
-    return Features;
-
-  for (auto Feature : llvm::split(TargetAndFeatures.second, ','))
-    Features.push_back(Feature.str());
-  return Features;
+// Get the list of target features from the input file and unify them such that
+// if there are multiple +xxx or -xxx features we only keep the last one.
+std::vector<std::string> getTargetFeatures(ArrayRef<OffloadFile> InputFiles) {
+  SmallVector<StringRef> Features;
+  for (const OffloadFile &File : InputFiles) {
+    for (auto Arg : llvm::split(File.getBinary()->getString("feature"), ","))
+      Features.emplace_back(Arg);
+  }
+
+  // Only add a feature if it hasn't been seen before starting from the end.
+  std::vector<std::string> UnifiedFeatures;
+  DenseSet<StringRef> UsedFeatures;
+  for (StringRef Feature : llvm::reverse(Features)) {
+    if (UsedFeatures.insert(Feature.drop_front()).second)
+      UnifiedFeatures.push_back(Feature.str());
+  }
+
+  return UnifiedFeatures;
 }
 
 CodeGenOpt::Level getCGOptLevel(unsigned OptLevel) {
   switch (OptLevel) {
   case 0:
     return CodeGenOpt::None;
   case 1:
     return CodeGenOpt::Less;
   case 2:
     return CodeGenOpt::Default;
   case 3:
     return CodeGenOpt::Aggressive;
   }
   llvm_unreachable("Invalid optimization level");
 }
 
 template <typename ModuleHook = function_ref<bool(size_t, const Module &)>>
 std::unique_ptr<lto::LTO> createLTO(
     const Triple &TheTriple, StringRef Arch, bool WholeProgram,
+    const std::vector<std::string> &Features,
     ModuleHook Hook = [](size_t, const Module &) { return true; }) {
   lto::Config Conf;
   lto::ThinBackend Backend;
   // TODO: Handle index-only thin-LTO
   Backend =
       lto::createInProcessThinBackend(llvm::heavyweight_hardware_concurrency());
 
   Conf.CPU = Arch.str();
   Conf.Options = codegen::InitTargetOptionsFromCodeGenFlags(TheTriple);
 
-  Conf.MAttrs = getTargetFeatures(TheTriple);
+  Conf.MAttrs = Features;
   Conf.CGOptLevel = getCGOptLevel(OptLevel[1] - '0');
   Conf.OptLevel = OptLevel[1] - '0';
   if (Conf.OptLevel > 0)
     Conf.UseDefaultPipeline = true;
   Conf.DefaultTriple = TheTriple.getTriple();
   Conf.DiagHandler = diagnosticHandler;
 
   Conf.PTO.LoopVectorization = Conf.OptLevel > 1;
auto OutputBitcode = [&](size_t Task, const Module &M) {
     if (EC)
       HandleError(errorCodeToError(EC));
     WriteBitcodeToFile(M, LinkedBitcode);
     BitcodeOutput.push_back(*TempFileOrErr);
     return false;
   };
 
   // We assume visibility of the whole program if every input file was bitcode.
+  auto Features = getTargetFeatures(BitcodeInputFiles);
   bool WholeProgram = InputFiles.empty();
   auto LTOBackend =
-      (EmbedBitcode) ? createLTO(TheTriple, Arch, WholeProgram, OutputBitcode)
-                     : createLTO(TheTriple, Arch, WholeProgram);
+      (EmbedBitcode)
+          ? createLTO(TheTriple, Arch, WholeProgram, Features, OutputBitcode)
+          : createLTO(TheTriple, Arch, WholeProgram, Features);
 
   // We need to resolve the symbols so the LTO backend knows which symbols need
   // to be kept or can be internalized. This is a simplified symbol resolution
   // scheme to approximate the full resolution a linker would do.
   DenseSet<StringRef> PrevailingSymbols;
   for (auto &BitcodeInput : BitcodeInputFiles) {
     MemoryBufferRef Buffer =
         MemoryBufferRef(BitcodeInput.getBinary()->getImage(), "");

clang/tools/clang-offload-packager/ClangOffloadPackager.cpp

 #include "llvm/Object/Binary.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Object/OffloadBinary.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileOutputBuffer.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/Signals.h"
+#include "llvm/Support/StringSaver.h"
 #include "llvm/Support/WithColor.h"
 
 using namespace llvm;
 using namespace llvm::object;
 
 static cl::opt<bool> Help("h", cl::desc("Alias for -help"), cl::Hidden);
 
 static cl::OptionCategory
int main(int argc, const char **argv) {
   auto reportError = [argv](Error E) {
     logAllUnhandledErrors(std::move(E), WithColor::error(errs(), argv[0]));
     return EXIT_FAILURE;
   };
 
   SmallVector<char, 1024> BinaryData;
   raw_svector_ostream OS(BinaryData);
   for (StringRef Image : DeviceImages) {
+    BumpPtrAllocator Alloc;
+    StringSaver Saver(Alloc);
+
     StringMap<StringRef> Args;
-    for (StringRef Arg : llvm::split(Image, ","))
-      Args.insert(Arg.split("="));
+    for (StringRef Arg : llvm::split(Image, ",")) {
+      auto KeyAndValue = Arg.split("=");
+      if (Args.count(KeyAndValue.first))
+        Args[KeyAndValue.first] =
+            Saver.save(Args[KeyAndValue.first] + "," + KeyAndValue.second);
+      else
+        Args[KeyAndValue.first] = KeyAndValue.second;
+    }
 
     if (!Args.count("triple") || !Args.count("file"))
       return reportError(createStringError(
           inconvertibleErrorCode(),
           "'file' and 'triple' are required image arguments"));
 
     OffloadBinary::OffloadingImage ImageBinary{};
     std::unique_ptr<llvm::MemoryBuffer> DeviceImage;