[Driver][SYCL] Update preprocessed file generation

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"