[Driver][SYCL] Update preprocessed file generation (#19849)

When generating preprocessed files for SYCL offloading, create a fully
packaged file that contains both the HOST and DEVICE binaries. This will
allow for consumption of these binary preprocessed files to be more
useful, as opposed to only being able to preprocess and keep the host
side of the offloading compilation

When the driver encounters preprocessed (file.ii) files on the command
line, these are processed in the following way:
 - Determines if the file is a packaged file (offload binary)
 - Extracts device side
 - Compiles device side, packages generated device into offload binary
 - Extracts host side
 - Compiles host side, embedding device binary

Offload binary determination is performed by checking the magic number
associated with the input file. The extraction is done via the
clang-offload-packager using a new JobAction.

When no output file is given, we will not package the preprocessed files
but will just perform the host preprocessing. When output to a file
(with an output file option), we will perform the host and device
compilation, package and output to that file.

Author: mdtoguchi

clang/include/clang/Driver/Action.h

@@ -76,6 +76,7 @@ class Action {
     OffloadUnbundlingJobClass,
     OffloadWrapperJobClass,
     OffloadPackagerJobClass,
+    OffloadPackagerExtractJobClass,
     OffloadDepsJobClass,
     SPIRVTranslatorJobClass,
     SYCLPostLinkJobClass,
@@ -719,6 +720,17 @@ class OffloadPackagerJobAction : public JobAction {
   }
 };
 
+class OffloadPackagerExtractJobAction : public JobAction {
+  void anchor() override;
+
+public:
+  OffloadPackagerExtractJobAction(ActionList &Inputs, types::ID Type);
+
+  static bool classof(const Action *A) {
+    return A->getKind() == OffloadPackagerExtractJobClass;
+  }
+};
+
 class OffloadDepsJobAction final : public JobAction {
   void anchor() override;
 

clang/include/clang/Driver/Driver.h

@@ -990,6 +990,9 @@ bool isObjectFile(std::string FileName);
 /// \return True if the filename has a static archive/lib extension.
 bool isStaticArchiveFile(const StringRef &FileName);
 
+/// \return True if the filename is an Offload Binary file.
+bool isOffloadBinaryFile(const StringRef &FileName);
+
 /// \return True if the argument combination will end up generating remarks.
 bool willEmitRemarks(const llvm::opt::ArgList &Args);
 

clang/include/clang/Driver/ToolChain.h

@@ -166,6 +166,7 @@ class ToolChain {
   mutable std::unique_ptr<Tool> OffloadBundler;
   mutable std::unique_ptr<Tool> OffloadWrapper;
   mutable std::unique_ptr<Tool> OffloadPackager;
+  mutable std::unique_ptr<Tool> OffloadPackagerExtract;
   mutable std::unique_ptr<Tool> OffloadDeps;
   mutable std::unique_ptr<Tool> SPIRVTranslator;
   mutable std::unique_ptr<Tool> SYCLPostLink;
@@ -185,6 +186,7 @@ class ToolChain {
   Tool *getOffloadBundler() const;
   Tool *getOffloadWrapper() const;
   Tool *getOffloadPackager() const;
+  Tool *getOffloadPackagerExtract() const;
   Tool *getOffloadDeps() const;
   Tool *getSPIRVTranslator() const;
   Tool *getSYCLPostLink() const;

clang/lib/Driver/Action.cpp

@@ -46,6 +46,8 @@ const char *Action::getClassName(ActionClass AC) {
     return "clang-offload-wrapper";
   case OffloadPackagerJobClass:
     return "clang-offload-packager";
+  case OffloadPackagerExtractJobClass:
+    return "clang-offload-packager-extract";
   case OffloadDepsJobClass:
     return "clang-offload-deps";
   case SPIRVTranslatorJobClass:
@@ -86,6 +88,15 @@ void Action::propagateDeviceOffloadInfo(OffloadKind OKind, const char *OArch,
   // Deps job uses the host kinds.
   if (Kind == OffloadDepsJobClass)
     return;
+  // Packaging actions can use host kinds for preprocessing.  When packaging
+  // preprocessed files, these packaged files will contain both host and device
+  // files, where the host side does not have any device info to propagate.
+  bool hasPreprocessJob =
+      std::any_of(Inputs.begin(), Inputs.end(), [](const Action *A) {
+        return A->getKind() == PreprocessJobClass;
+      });
+  if (Kind == OffloadPackagerJobClass && hasPreprocessJob)
+    return;
 
   assert((OffloadingDeviceKind == OKind || OffloadingDeviceKind == OFK_None) &&
          "Setting device kind to a different device??");
@@ -485,6 +496,12 @@ OffloadPackagerJobAction::OffloadPackagerJobAction(ActionList &Inputs,
                                                    types::ID Type)
     : JobAction(OffloadPackagerJobClass, Inputs, Type) {}
 
+void OffloadPackagerExtractJobAction::anchor() {}
+
+OffloadPackagerExtractJobAction::OffloadPackagerExtractJobAction(
+    ActionList &Inputs, types::ID Type)
+    : JobAction(OffloadPackagerExtractJobClass, Inputs, Type) {}
+
 void OffloadDepsJobAction::anchor() {}
 
 OffloadDepsJobAction::OffloadDepsJobAction(

clang/lib/Driver/Driver.cpp

@@ -7780,7 +7780,8 @@ Action *Driver::BuildOffloadingActions(Compilation &C,
                                        Action *HostAction) const {
   // Don't build offloading actions if explicitly disabled or we do not have a
   // valid source input.
-  if (offloadHostOnly() || !types::isSrcFile(Input.first))
+  if (offloadHostOnly() ||
+      !(types::isSrcFile(Input.first) || Input.first == types::TY_PP_CXX))
     return HostAction;
 
   bool HIPNoRDC =
@@ -7845,6 +7846,25 @@ Action *Driver::BuildOffloadingActions(Compilation &C,
     for (const ToolChain *TC : ToolChains) {
       for (StringRef Arch : OffloadArchs.lookup(TC)) {
         TCAndArchs.push_back(std::make_pair(TC, Arch));
+        // Check if the InputArg is a preprocessed file that is created by the
+        // clang-offload-packager.
+        if (InputType == types::TY_PP_CXX &&
+            isOffloadBinaryFile(InputArg->getAsString(Args))) {
+          // Extract the specific preprocessed file given the current arch
+          // and triple.  Add to DeviceActions if one was extracted.
+          ActionList PPActions;
+          OffloadAction::DeviceDependences DDep;
+          Action *IA = C.MakeAction<InputAction>(*InputArg, InputType, CUID);
+          PPActions.push_back(IA);
+          Action *PackagerAction =
+              C.MakeAction<OffloadPackagerExtractJobAction>(PPActions,
+                                                            types::TY_PP_CXX);
+          DDep.add(*PackagerAction,
+                   *C.getSingleOffloadToolChain<Action::OFK_Host>(), nullptr,
+                   C.getActiveOffloadKinds());
+          DeviceActions.push_back(PackagerAction);
+          continue;
+        }
         DeviceActions.push_back(
             C.MakeAction<InputAction>(*InputArg, InputType, CUID));
       }
@@ -8002,6 +8022,37 @@ Action *Driver::BuildOffloadingActions(Compilation &C,
     DDep.add(*LinkAction, *C.getSingleOffloadToolChain<Action::OFK_Host>(),
              nullptr, C.getActiveOffloadKinds());
     return C.MakeAction<OffloadAction>(DDep, types::TY_Nothing);
+  } else if (C.isOffloadingHostKind(Action::OFK_SYCL) &&
+             isa<PreprocessJobAction>(HostAction) &&
+             getFinalPhase(Args) == phases::Preprocess &&
+             Args.hasArg(options::OPT_o, options::OPT__SLASH_P,
+                         options::OPT__SLASH_o)) {
+    // Performing preprocessing only. Take the host and device preprocessed
+    // files and package them together.
+    ActionList PackagerActions;
+    // Only add the preprocess actions from the device side.  When one is
+    // found, add an additional compilation to generate the integration
+    // header/footer that is used for the host compile.
+    for (auto OA : OffloadActions) {
+      if (const OffloadAction *CurOA = dyn_cast<OffloadAction>(OA)) {
+        CurOA->doOnEachDependence(
+            [&](Action *A, const ToolChain *TC, const char *BoundArch) {
+              assert(TC && "Unknown toolchain");
+              if (isa<PreprocessJobAction>(A)) {
+                PackagerActions.push_back(OA);
+                A->setCannotBeCollapsedWithNextDependentAction();
+                Action *CompileAction =
+                    C.MakeAction<CompileJobAction>(A, types::TY_Nothing);
+                DDeps.add(*CompileAction, *TC, BoundArch, Action::OFK_SYCL);
+              }
+            });
+      }
+    }
+    PackagerActions.push_back(HostAction);
+    Action *PackagerAction = C.MakeAction<OffloadPackagerJobAction>(
+        PackagerActions, types::TY_PP_CXX);
+    DDeps.add(*PackagerAction, *C.getSingleOffloadToolChain<Action::OFK_Host>(),
+              nullptr, C.getActiveOffloadKinds());
   } else if (C.isOffloadingHostKind(Action::OFK_SYCL) &&
              Args.hasArg(options::OPT_fsycl_host_compiler_EQ)) {
     // -fsycl-host-compiler will create a bundled object instead of an
@@ -8150,6 +8201,23 @@ Action *Driver::ConstructPhaseAction(
       return C.MakeAction<VerifyPCHJobAction>(Input, types::TY_Nothing);
     if (Args.hasArg(options::OPT_extract_api))
       return C.MakeAction<ExtractAPIJobAction>(Input, types::TY_API_INFO);
+    // New offload driver enabled with a Preprocessed input file - check to make
+    // sure that the input file is an offload binary - if so, we need to
+    // extract the actual preprocessed file from the package, and that is what
+    // we will compile.
+    if (getUseNewOffloadingDriver() &&
+        TargetDeviceOffloadKind == Action::OFK_None &&
+        Input->getType() == types::TY_PP_CXX) {
+      const InputAction *IA = dyn_cast<InputAction>(Input);
+      if (IA && isOffloadBinaryFile(IA->getInputArg().getAsString(Args))) {
+        ActionList PPActions;
+        PPActions.push_back(Input);
+        Action *PackagerAction = C.MakeAction<OffloadPackagerExtractJobAction>(
+            PPActions, types::TY_PP_CXX);
+        return C.MakeAction<CompileJobAction>(PackagerAction,
+                                              types::TY_LLVM_BC);
+      }
+    }
     return C.MakeAction<CompileJobAction>(Input, types::TY_LLVM_BC);
   }
   case phases::Backend: {
@@ -9429,7 +9497,8 @@ const char *Driver::GetNamedOutputPath(Compilation &C, const JobAction &JA,
   // For /P, preprocess to file named after BaseInput.
   if (C.getArgs().hasArg(options::OPT__SLASH_P) &&
       ((AtTopLevel && isa<PreprocessJobAction>(JA)) ||
-       isa<OffloadBundlingJobAction>(JA))) {
+       isa<OffloadBundlingJobAction>(JA) ||
+       isa<OffloadPackagerJobAction>(JA))) {
     StringRef BaseName = llvm::sys::path::filename(BaseInput);
     StringRef NameArg;
     if (Arg *A = C.getArgs().getLastArg(options::OPT__SLASH_Fi))
@@ -9465,6 +9534,14 @@ const char *Driver::GetNamedOutputPath(Compilation &C, const JobAction &JA,
     }
   }
 
+  // When generating preprocessed files (-E) with offloading enabled, redirect
+  // the output to a properly named output file.
+  if (JA.getType() == types::TY_PP_CXX && isa<OffloadPackagerJobAction>(JA)) {
+    if (Arg *FinalOutput =
+            C.getArgs().getLastArg(options::OPT_o, options::OPT__SLASH_o))
+      return C.addResultFile(FinalOutput->getValue(), &JA);
+  }
+
   // Default to writing to stdout?
   if (AtTopLevel && !CCGenDiagnostics && HasPreprocessOutput(JA)) {
     return "-";
@@ -10437,6 +10514,12 @@ bool clang::driver::isStaticArchiveFile(const StringRef &FileName) {
   return (Magic == llvm::file_magic::archive);
 }
 
+bool clang::driver::isOffloadBinaryFile(const StringRef &FileName) {
+  llvm::file_magic Magic;
+  llvm::identify_magic(FileName, Magic);
+  return (Magic == llvm::file_magic::offload_binary);
+}
+
 bool clang::driver::willEmitRemarks(const ArgList &Args) {
   // -fsave-optimization-record enables it.
   if (Args.hasFlag(options::OPT_fsave_optimization_record,

clang/lib/Driver/ToolChain.cpp

@@ -617,6 +617,12 @@ Tool *ToolChain::getOffloadPackager() const {
   return OffloadPackager.get();
 }
 
+Tool *ToolChain::getOffloadPackagerExtract() const {
+  if (!OffloadPackagerExtract)
+    OffloadPackagerExtract.reset(new tools::OffloadPackagerExtract(*this));
+  return OffloadPackagerExtract.get();
+}
+
 Tool *ToolChain::getOffloadDeps() const {
   if (!OffloadDeps)
     OffloadDeps.reset(new tools::OffloadDeps(*this));
@@ -707,6 +713,8 @@ Tool *ToolChain::getTool(Action::ActionClass AC) const {
     return getOffloadWrapper();
   case Action::OffloadPackagerJobClass:
     return getOffloadPackager();
+  case Action::OffloadPackagerExtractJobClass:
+    return getOffloadPackagerExtract();
 
   case Action::OffloadDepsJobClass:
     return getOffloadDeps();

clang/lib/Driver/ToolChains/Clang.cpp

@@ -10300,6 +10300,8 @@ void OffloadPackager::ConstructJob(Compilation &C, const JobAction &JA,
   for (const InputInfo &Input : Inputs) {
     const Action *OffloadAction = Input.getAction();
     const ToolChain *TC = OffloadAction->getOffloadingToolChain();
+    if (!TC)
+      TC = &C.getDefaultToolChain();
     const ArgList &TCArgs =
         C.getArgsForToolChain(TC, OffloadAction->getOffloadingArch(),
                               OffloadAction->getOffloadingDeviceKind());
@@ -10389,6 +10391,50 @@ void OffloadPackager::ConstructJob(Compilation &C, const JobAction &JA,
       CmdArgs, Inputs, Output));
 }
 
+// Use the clang-offload-packager to extract binaries from a packaged
+// binary.  This currently only supports single input/single output.
+void OffloadPackagerExtract::ConstructJob(Compilation &C, const JobAction &JA,
+                                          const InputInfo &Output,
+                                          const InputInfoList &Inputs,
+                                          const llvm::opt::ArgList &Args,
+                                          const char *LinkingOutput) const {
+  ArgStringList CmdArgs;
+  const Action *OffloadAction = Inputs[0].getAction();
+  const ToolChain *TC = OffloadAction->getOffloadingToolChain();
+  if (!TC)
+    TC = &C.getDefaultToolChain();
+  const ArgList &TCArgs =
+      C.getArgsForToolChain(TC, OffloadAction->getOffloadingArch(),
+                            OffloadAction->getOffloadingDeviceKind());
+
+  // Input file name.
+  StringRef InFile = C.getArgs().MakeArgString(TC->getInputFilename(Inputs[0]));
+  CmdArgs.push_back(Args.MakeArgString(InFile));
+
+  // Generated --image option containing the output file name, triple, arch
+  // and associated offload kind.
+  assert(Output.isFilename() && "Invalid output.");
+  StringRef File = Output.getFilename();
+  StringRef Arch = OffloadAction->getOffloadingArch()
+                       ? OffloadAction->getOffloadingArch()
+                       : TCArgs.getLastArgValue(options::OPT_march_EQ);
+  StringRef Kind =
+      Action::GetOffloadKindName(OffloadAction->getOffloadingDeviceKind());
+
+  SmallVector<std::string> Parts{
+      "file=" + File.str(),
+      "triple=" + TC->getTripleString(),
+      "arch=" + (Arch.empty() ? "generic" : Arch.str()),
+      "kind=" + Kind.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));
+}
+
 // Begin OffloadDeps
 
 void OffloadDeps::constructJob(Compilation &C, const JobAction &JA,

clang/lib/Driver/ToolChains/Clang.h

@@ -190,6 +190,19 @@ class LLVM_LIBRARY_VISIBILITY OffloadPackager final : public Tool {
                     const char *LinkingOutput) const override;
 };
 
+/// Offload binary extract tool.
+class LLVM_LIBRARY_VISIBILITY OffloadPackagerExtract final : public Tool {
+public:
+  OffloadPackagerExtract(const ToolChain &TC)
+      : Tool("Offload::PackagerExtract", "clang-offload-packager", TC) {}
+
+  bool hasIntegratedCPP() const override { return false; }
+  void ConstructJob(Compilation &C, const JobAction &JA,
+                    const InputInfo &Output, const InputInfoList &Inputs,
+                    const llvm::opt::ArgList &TCArgs,
+                    const char *LinkingOutput) const override;
+};
+
 /// Offload deps tool.
 class LLVM_LIBRARY_VISIBILITY OffloadDeps final : public Tool {
   void constructJob(Compilation &C, const JobAction &JA,

clang/test/Driver/sycl-int-header-footer.cpp

@@ -17,7 +17,7 @@
 // FOOTER_PREPROC_GEN-SAME: "-dependency-filter" "[[INTHEADER]]"
 // FOOTER_PREPROC_GEN-SAME: "-include-internal-footer" "[[INTFOOTER]]"
 // FOOTER_PREPROC_GEN-SAME: "-dependency-filter" "[[INTFOOTER]]"
-// FOOTER_PREPROC_GEN-SAME: "-E"{{.*}} "-o" "-"
+// FOOTER_PREPROC_GEN-SAME: "-E"
 
 /// Preprocessed file use with integration footer
 // RUN: touch %t.ii

clang/test/Driver/sycl-preprocess.cpp

@@ -1,9 +1,10 @@
-/// Test preprocessing capabilities when using -fsycl
+// Test the behaviors when enabling SYCL offloading with preprocessed files.
+
 /// Creating a preprocessed file is expected to do an integration header
 /// creation step.
 // RUN: %clangxx -fsycl --offload-new-driver -E -o %t_output.ii %s -### 2>&1 \
 // RUN:  | FileCheck -check-prefix PREPROC_ONLY %s
-// RUN: %clang_cl -fsycl --offload-new-driver -P -Fi%t_output.ii %s -### 2>&1 \
+// RUN: %clang_cl -fsycl --offload-new-driver -P %s -### 2>&1 \
 // RUN:  | FileCheck -check-prefix PREPROC_ONLY %s
 // PREPROC_ONLY: clang{{.*}} "-fsycl-is-device"{{.*}} "-fsycl-int-header=[[INTHEADER:.+\.h]]" "-fsycl-int-footer=[[INTFOOTER:.+\.h]]"{{.*}} "-E"
 // PREPROC_ONLY: clang{{.*}} "-fsycl-is-host"{{.*}} "-include-internal-header" "[[INTHEADER]]"{{.*}} "-include-internal-footer" "[[INTFOOTER]]"{{.*}} "-o" "[[HOST_OUT:.+\.ii]]"
@@ -14,10 +15,31 @@
 // PREPROC_IN-NOT: "-fsycl-int-header={{.*}}"
 // PREPROC_IN: clang{{.*}} "-fsycl-is-host"
 
-// RUN: %clangxx -target x86_64-unknown-linux-gnu -fsycl --offload-new-driver -E %s -ccc-print-phases 2>&1 \
-// RUN:  | FileCheck -check-prefix PREPROC_PHASES %s
+/// When generating preprocessed files, verify the compilation phases.
+// RUN: %clangxx --target=x86_64-unknown-linux-gnu --offload-new-driver -fsycl -E %s -o %t.ii -ccc-print-phases 2>&1 \
+// RUN: | FileCheck %s -check-prefix PREPROC_PHASES -DTARGET=x86_64-unknown-linux-gnu
+// RUN: %clang_cl --target=x86_64-pc-windows-msvc --offload-new-driver -fsycl -P %s -Fi%t.ii -ccc-print-phases 2>&1 \
+// RUN: | FileCheck %s -check-prefix PREPROC_PHASES -DTARGET=x86_64-pc-windows-msvc
 // PREPROC_PHASES: 0: input, "[[INPUT:.+\.cpp]]", c++, (host-sycl)
 // PREPROC_PHASES: 1: preprocessor, {0}, c++-cpp-output, (host-sycl)
 // PREPROC_PHASES: 2: input, "[[INPUT]]", c++, (device-sycl)
 // PREPROC_PHASES: 3: preprocessor, {2}, c++-cpp-output, (device-sycl)
-// PREPROC_PHASES: 4: offload, "host-sycl (x86_64-unknown-linux-gnu)" {1}, "device-sycl (spir64-unknown-unknown)" {3}, c++-cpp-output
+// PREPROC_PHASES: 4: compiler, {3}, none, (device-sycl)
+// PREPROC_PHASES: 5: offload, "device-sycl (spir64-unknown-unknown)" {3}, c++-cpp-output
+// PREPROC_PHASES: 6: clang-offload-packager, {5, 1}, c++-cpp-output
+// PREPROC_PHASES: 7: offload, "host-sycl ([[TARGET]])" {1}, "device-sycl (spir64-unknown-unknown)" {3}, "device-sycl (spir64-unknown-unknown)" {4}, " ([[TARGET]])" {6}, c++-cpp-output
+
+/// When generating preprocessed files, verify the tools called and the expected
+/// output file name.
+// RUN: %clangxx --offload-new-driver -fsycl -E %s -o sycl-preprocess.ii -### 2>&1 \
+// RUN: | FileCheck %s -check-prefix PREPROC_TOOLS
+// RUN: %clang_cl --offload-new-driver -fsycl -P %s -Fisycl-preprocess.ii -### 2>&1 \
+// RUN: | FileCheck %s -check-prefix PREPROC_TOOLS
+// RUN: %clang_cl --offload-new-driver -fsycl -E %s -o sycl-preprocess.ii -### 2>&1 \
+// RUN: | FileCheck %s -check-prefix PREPROC_TOOLS
+// PREPROC_TOOLS: clang{{.*}} "-fsycl-is-device"
+// PREPROC_TOOLS-SAME: "-o" "[[DEVICE_PP_FILE:.+\.ii]]
+// PREPROC_TOOLS: clang{{.*}} "-fsycl-is-host"
+// PREPROC_TOOLS-SAME: "-o" "[[HOST_PP_FILE:.+\.ii]]
+// PREPROC_TOOLS: clang-offload-packager{{.*}} "-o" "sycl-preprocess.ii"
+// PREPROC_TOOLS-SAME: "--image=file=[[DEVICE_PP_FILE]],triple=spir64-unknown-unknown,arch=generic,kind=sycl{{.*}}" "--image=file=[[HOST_PP_FILE]],triple={{.*}},arch=generic,kind=host"