[SYCL][thinLTO] Split early when thinLTO is enabled (#14259)

This change implements driver changes to call `sycl-post-link` during
`-c` when thinLTO is enabled, and updates `clang-offload-packager` to
accept response files (same format as temp file list), which will be
generated when we call `clang` on each output of `sycl-post-link`.

I will add a test for the `clang-offload-packager` changes later,
because we rely on changes in other PRs to test the full flow. However,
the code is already being hit because we compile `libdevice` with
thinLTO when compiling for the new offload model.

For a given file `a.cpp` that results in two total splits, the following
will happen:
1) We call `sycl-post-link` on `a.bc`
2) We extract the "Code" column from the table output of
`sycl-post-link`
3) We call `llvm-foreach` with the regular backend job that would be
called normally on each `bc` file from 2)
4) We call `clang-offload-packager` with the output of 3 which is a list
of the `bc` files from each `clang` run, with a '@' appended to the file
path
5) `clang-offload-packager` reads each entry inside 4) and adds each to
the fat object
6) `clang-linker-wrapper` extracts the fat object from 5) and finds all
of the files
7) ... to be implemented

---------

Signed-off-by: Sarnie, Nick <[email protected]>

Author: sarnex

clang/lib/Driver/Driver.cpp

@@ -8190,6 +8190,37 @@ Action *Driver::ConstructPhaseAction(
         TargetDeviceOffloadKind != Action::OFK_None) {
       types::ID Output =
           Args.hasArg(options::OPT_S) ? types::TY_LTO_IR : types::TY_LTO_BC;
+      if (getUseNewOffloadingDriver() &&
+          getLTOMode(/*IsDeviceOffloadAction=*/true) == LTOK_Thin &&
+          TargetDeviceOffloadKind == Action::OFK_SYCL) {
+        // For SYCL with thinLTO, run sycl-post-link, extract the BC files from
+        // the output table, run the backend on each output table.
+        llvm::Triple OffloadTriple =
+            Input->getOffloadingToolChain()->getTriple();
+        SYCLPostLinkJobAction *TypedPostLinkAction =
+            C.MakeAction<SYCLPostLinkJobAction>(Input, types::TY_Tempfiletable,
+                                                types::TY_Tempfiletable);
+        TypedPostLinkAction->setRTSetsSpecConstants(
+            OffloadTriple.isSPIROrSPIRV() && !OffloadTriple.isSPIRAOT());
+        auto *TypedExtractIRFilesAction = C.MakeAction<FileTableTformJobAction>(
+            TypedPostLinkAction, types::TY_Tempfilelist,
+            types::TY_Tempfilelist);
+
+        TypedExtractIRFilesAction->addExtractColumnTform(
+            FileTableTformJobAction::COL_CODE, false /*drop titles*/);
+        auto *OutputAction =
+            C.MakeAction<BackendJobAction>(TypedExtractIRFilesAction, Output);
+
+        auto *ForEach = C.MakeAction<ForEachWrappingAction>(
+            TypedExtractIRFilesAction, OutputAction);
+        // This final job is mostly a no-op, but we need it to set the Action
+        // type to Tempfilelist which is expected by clang-offload-packager.
+        auto *ExtractBCFiles = C.MakeAction<FileTableTformJobAction>(
+            ForEach, types::TY_Tempfilelist, types::TY_Tempfilelist);
+        ExtractBCFiles->addExtractColumnTform(FileTableTformJobAction::COL_ZERO,
+                                              false /*drop titles*/);
+        return ExtractBCFiles;
+      }
       return C.MakeAction<BackendJobAction>(Input, Output);
     }
     if (Args.hasArg(options::OPT_emit_llvm) ||

clang/lib/Driver/ToolChains/Clang.cpp

@@ -593,6 +593,11 @@ static void addDashXForInput(const ArgList &Args, const InputInfo &Input,
   if (Args.hasArg(options::OPT_verify_pch) && Input.getType() == types::TY_PCH)
     return;
 
+  // If the input is a Tempfilelist, this call is part for a
+  // llvm-foreach call and we should infer the type from the file extension.
+  if (Input.getType() == types::TY_Tempfilelist)
+    return;
+
   CmdArgs.push_back("-x");
   if (Args.hasArg(options::OPT_rewrite_objc))
     CmdArgs.push_back(types::getTypeName(types::TY_PP_ObjCXX));
@@ -10250,6 +10255,13 @@ void OffloadPackager::ConstructJob(Compilation &C, const JobAction &JA,
         C.getArgsForToolChain(TC, OffloadAction->getOffloadingArch(),
                               OffloadAction->getOffloadingDeviceKind());
     StringRef File = C.getArgs().MakeArgString(TC->getInputFilename(Input));
+
+    // If the input is a Tempfilelist, it is a response file
+    // which internally contains a list of files to be processed.
+    // Add an '@' so the tool knows to expand the response file.
+    if (Input.getType() == types::TY_Tempfilelist)
+      File = C.getArgs().MakeArgString("@" + File);
+
     StringRef Arch = OffloadAction->getOffloadingArch()
                          ? OffloadAction->getOffloadingArch()
                          : TCArgs.getLastArgValue(options::OPT_march_EQ);

clang/test/Driver/sycl-lto.cpp

@@ -8,7 +8,12 @@
 // RUN: not %clangxx -fsycl --offload-new-driver -foffload-lto=thin -fsycl-device-code-split=off %s -### 2>&1 | FileCheck -check-prefix=CHECK_SPLIT_ERROR %s
 // CHECK_SPLIT_ERROR: '-fsycl-device-code-split=off' is not supported when '-foffload-lto=thin' is set with '-fsycl'
 
-// Verify there's no error and we see the expected cc1 flags with the new offload driver.
+// Verify there's no error and we see the expected cc1 flags and tool invocations with the new offload driver.
 // RUN: %clangxx -fsycl --offload-new-driver -foffload-lto=thin %s -### 2>&1 | FileCheck -check-prefix=CHECK_SUPPORTED %s
 // CHECK_SUPPORTED: clang{{.*}} "-cc1" "-triple" "spir64-unknown-unknown" {{.*}} "-flto=thin" "-flto-unit"
+// CHECK_SUPPORTED: sycl-post-link{{.*}}
+// CHECK_SUPPORTED-NEXT: file-table-tform{{.*}}
+// CHECK_SUPPORTED-NEXT: llvm-foreach{{.*}} "--" {{.*}}clang{{.*}} "-fsycl-is-device"{{.*}} "-flto=thin" "-flto-unit"
+// CHECK_SUPPORTED-NEXT: file-table-tform{{.*}}
+// CHECK_SUPPORTED-NEXT: clang-offload-packager{{.*}} "-o" "{{.*}}" "--image=file=@{{.*}}"
 // CHECK_SUPPORTED: clang-linker-wrapper{{.*}} "-sycl-thin-lto"

clang/tools/clang-offload-packager/CMakeLists.txt

@@ -2,7 +2,8 @@ set(LLVM_LINK_COMPONENTS
   ${LLVM_TARGETS_TO_BUILD}
   BinaryFormat
   Object
-  Support)
+  Support
+  TargetParser)
 
 add_clang_tool(clang-offload-packager
   ClangOffloadPackager.cpp

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

@@ -26,6 +26,7 @@
 #include "llvm/Support/Signals.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/WithColor.h"
+#include "llvm/TargetParser/Host.h"
 
 using namespace llvm;
 using namespace llvm::object;
@@ -63,13 +64,35 @@ static void PrintVersion(raw_ostream &OS) {
   OS << clang::getClangToolFullVersion("clang-offload-packager") << '\n';
 }
 
+// For any response file arguments (those starting with '@'), expand them into
+// the contents of the response file, deliminated by commas.
+static StringRef expandResponseFileImageArguments(StringRef Arg,
+                                                  StringSaver &Saver) {
+  std::string FileStr = Arg.str();
+  SmallVector<const char *, 0> ExpandedFiles{FileStr.c_str()};
+  cl::ExpandResponseFiles(Saver,
+                          Triple(sys::getProcessTriple()).isOSWindows()
+                              ? cl::TokenizeWindowsCommandLine
+                              : cl::TokenizeGNUCommandLine,
+                          ExpandedFiles);
+  std::string NewValue;
+  for (size_t FileIdx = 0; FileIdx < ExpandedFiles.size(); FileIdx++) {
+    NewValue += ExpandedFiles[FileIdx];
+    if (FileIdx != ExpandedFiles.size() - 1)
+      NewValue += ',';
+  }
+  return Saver.save(NewValue);
+}
+
 // Get a map containing all the arguments for the image. Repeated arguments will
 // be placed in a comma separated list.
 static DenseMap<StringRef, StringRef> getImageArguments(StringRef Image,
                                                         StringSaver &Saver) {
   DenseMap<StringRef, StringRef> Args;
   for (StringRef Arg : llvm::split(Image, ",")) {
     auto [Key, Value] = Arg.split("=");
+    if (Key == "file" && Value[0] == '@')
+      Value = expandResponseFileImageArguments(Value, Saver);
     if (Args.count(Key))
       Args[Key] = Saver.save(Args[Key] + "," + Value);
     else