[SYCL][Driver] Add device code split job

clang/include/clang/Driver/Action.h

@@ -75,6 +75,7 @@ class Action {
     OffloadWrapperJobClass,
     SPIRVTranslatorJobClass,
     SPIRCheckJobClass,
+    SYCLPostLinkJobClass,
     BackendCompileJobClass,
 
     JobClassFirst = PreprocessJobClass,
@@ -667,6 +668,17 @@ class SPIRCheckJobAction : public JobAction {
   }
 };
 
+class SYCLPostLinkJobAction : public JobAction {
+  void anchor() override;
+
+public:
+  SYCLPostLinkJobAction(Action *Input, types::ID OutputType);
+
+  static bool classof(const Action *A) {
+    return A->getKind() == SYCLPostLinkJobClass;
+  }
+};
+
 class BackendCompileJobAction : public JobAction {
   void anchor() override;
 

clang/include/clang/Driver/Options.td

@@ -1822,6 +1822,13 @@ def fsycl_add_targets_EQ : CommaJoined<["-"], "fsycl-add-targets=">, Flags<[Driv
   HelpText<"Specify comma-separated list of triple and device binary image pairs to add to the final SYCL binary">;
 def fsycl_link_targets_EQ : CommaJoined<["-"], "fsycl-link-targets=">, Flags<[DriverOption, CC1Option, CoreOption]>,
   HelpText<"Specify comma-separated list of triples SYCL offloading targets to produce linked device images">;
+def fsycl_device_code_split_EQ : Joined<["-"], "fsycl-device-code-split=">,
+   Flags<[CC1Option, CoreOption]>, HelpText<"Perform SYCL device code split: per_kernel (device code module is "
+  "created for each SYCL kernel) | per_source (device code module is created for each source (translation unit)). "
+  "Default is no device code split - all kernels go into a single module`">, Values<"per_source,per_kernel">;
+def fsycl_device_code_split : Flag<["-"], "fsycl-device-code-split">, Alias<fsycl_device_code_split_EQ>,
+  AliasArgs<["per_source"]>, Flags<[CC1Option, CoreOption]>,
+  HelpText<"Perform SYCL device code split in the per_source mode i.e. create a device code module for each source (translation unit)">;
 def fsycl_use_bitcode : Flag<["-"], "fsycl-use-bitcode">,
   Flags<[CC1Option, CoreOption]>, HelpText<"Use LLVM bitcode instead of SPIR-V in fat objects">;
 def fno_sycl_use_bitcode : Flag<["-"], "fno-sycl-use-bitcode">,

clang/include/clang/Driver/ToolChain.h

@@ -142,6 +142,7 @@ class ToolChain {
   mutable std::unique_ptr<Tool> OffloadWrapper;
   mutable std::unique_ptr<Tool> SPIRVTranslator;
   mutable std::unique_ptr<Tool> SPIRCheck;
+  mutable std::unique_ptr<Tool> SYCLPostLink;
   mutable std::unique_ptr<Tool> BackendCompiler;
 
   Tool *getClang() const;
@@ -154,6 +155,7 @@ class ToolChain {
   Tool *getOffloadWrapper() const;
   Tool *getSPIRVTranslator() const;
   Tool *getSPIRCheck() const;
+  Tool *getSYCLPostLink() const;
   Tool *getBackendCompiler() const;
 
   mutable std::unique_ptr<SanitizerArgs> SanitizerArguments;

clang/include/clang/Driver/Types.def

@@ -102,7 +102,8 @@ TYPE("hip-fatbin",               HIP_FATBIN,   INVALID,         "hipfb",  phases
 TYPE("spirv",                    SPIRV,        INVALID,         "spv",    phases::Compile, phases::Backend, phases::Assemble, phases::Link)
 TYPE("sycl-header",              SYCL_Header,  INVALID,         "h",      phases::Compile, phases::Backend, phases::Assemble, phases::Link)
 TYPE("sycl-fatbin",              SYCL_FATBIN,  INVALID,         nullptr,  phases::Compile, phases::Backend, phases::Assemble, phases::Link)
-TYPE("tempfilelist",             Tempfilelist, INVALID,         nullptr,  phases::Compile, phases::Backend, phases::Assemble, phases::Link)
+TYPE("tempfilelist",             Tempfilelist, INVALID,         "txt",    phases::Compile, phases::Backend, phases::Assemble, phases::Link)
+TYPE("tempentriesfilelist",      TempEntriesfilelist, INVALID,  "txt",    phases::Compile, phases::Backend, phases::Assemble, phases::Link)
 TYPE("archive",                  Archive,      INVALID,         "a",      phases::Compile, phases::Backend, phases::Assemble, phases::Link)
 TYPE("fpga-aocx",                FPGA_AOCX,    INVALID,         "aocx",   phases::Compile, phases::Backend, phases::Assemble, phases::Link)
 TYPE("fpga-aocr",                FPGA_AOCR,    INVALID,         "aocr",   phases::Compile, phases::Backend, phases::Assemble, phases::Link)

clang/lib/Driver/Action.cpp

@@ -47,6 +47,8 @@ const char *Action::getClassName(ActionClass AC) {
     return "llvm-spirv";
   case SPIRCheckJobClass:
     return "llvm-no-spir-kernel";
+  case SYCLPostLinkJobClass:
+    return "sycl-post-link";
   case BackendCompileJobClass:
     return "backend-compiler";
   }
@@ -443,6 +445,11 @@ void SPIRCheckJobAction::anchor() {}
 SPIRCheckJobAction::SPIRCheckJobAction(Action *Input, types::ID Type)
     : JobAction(SPIRCheckJobClass, Input, Type) {}
 
+void SYCLPostLinkJobAction::anchor() {}
+
+SYCLPostLinkJobAction::SYCLPostLinkJobAction(Action *Input, types::ID Type)
+    : JobAction(SYCLPostLinkJobClass, Input, Type) {}
+
 void BackendCompileJobAction::anchor() {}
 
 BackendCompileJobAction::BackendCompileJobAction(Action *Input,

clang/lib/Driver/Compilation.cpp

@@ -135,7 +135,8 @@ bool Compilation::CleanupFileList(const TempFileList &Files,
   for (const auto &File: Files) {
     // Temporary file lists contain files that need to be cleaned. The
     // file containing the information is also removed
-    if (File.second == types::TY_Tempfilelist) {
+    if (File.second == types::TY_Tempfilelist ||
+        File.second == types::TY_TempEntriesfilelist) {
       std::ifstream ListFile(File.first);
       if (ListFile) {
         // These are temporary files and need to be removed.

clang/lib/Driver/Driver.cpp

@@ -3195,9 +3195,12 @@ class OffloadingActionBuilder final {
     /// Flag to signal if the user requested device-only compilation.
     bool CompileDeviceOnly = false;
 
-    /// Flag to signal if the user requested the device object to be wrapped
+    /// Flag to signal if the user requested the device object to be wrapped.
     bool WrapDeviceOnlyBinary = false;
 
+    /// Flag to signal if the user requested device code split.
+    bool DeviceCodeSplit = false;
+
     /// The SYCL actions for the current input.
     ActionList SYCLDeviceActions;
 
@@ -3237,7 +3240,7 @@ class OffloadingActionBuilder final {
                          phases::ID CurPhase, phases::ID FinalPhase,
                          PhasesTy &Phases) override {
 
-      // FIXME: This adds the integrated header generation pass before the
+      // FIXME: This adds the integration header generation pass before the
       // Host compilation pass so the Host can use the header generated.  This
       // can be improved upon to where the header generation and spv generation
       // is done in the same step.  Currently, its not too efficient.
@@ -3279,9 +3282,12 @@ class OffloadingActionBuilder final {
             SYCLLinkBinary =
                 C.MakeAction<OffloadWrapperJobAction>(DeviceLinkAction,
                                                       types::TY_Object);
-          } else
-            SYCLLinkBinary = C.MakeAction<LinkJobAction>(SYCLLinkBinaryList,
+          } else {
+            auto *Link = C.MakeAction<LinkJobAction>(SYCLLinkBinaryList,
                                                          types::TY_Image);
+            SYCLLinkBinary = C.MakeAction<SPIRVTranslatorJobAction>(
+                Link, types::TY_Image);
+          }
 
           // Remove the SYCL actions as they are already connected to an host
           // action or fat binary.
@@ -3458,19 +3464,38 @@ class OffloadingActionBuilder final {
             }
           }
           if (!DeviceObjects.empty()) {
-            // link and wrap the device binary, but do not perform the
-            // backend compile.
+            // When aocx or aocr is found, there is an expectation that none of
+            // the other objects processed have any kernel. So, there
+            // is no need in device code split and backend compile here. Just
+            // link and wrap the device binary.
             auto *DeviceLinkAction =
-                C.MakeAction<LinkJobAction>(DeviceObjects, types::TY_SPIRV);
+                C.MakeAction<LinkJobAction>(DeviceObjects, types::TY_LLVM_BC);
+            auto *SPIRVTranslateAction = C.MakeAction<SPIRVTranslatorJobAction>(
+                DeviceLinkAction, types::TY_SPIRV);
             auto *DeviceWrappingAction = C.MakeAction<OffloadWrapperJobAction>(
-                DeviceLinkAction, types::TY_Object);
+                SPIRVTranslateAction, types::TY_Object);
             DA.add(*DeviceWrappingAction, **TC, /*BoundArch=*/nullptr,
                    Action::OFK_SYCL);
           }
           continue;
         }
         auto *DeviceLinkAction =
-            C.MakeAction<LinkJobAction>(LI, types::TY_SPIRV);
+            C.MakeAction<LinkJobAction>(LI, types::TY_LLVM_BC);
+        ActionList WrapperInputs;
+        Action *SPIRVInput = DeviceLinkAction;
+        types::ID OutType = types::TY_SPIRV;
+        if (DeviceCodeSplit) {
+          auto *SplitAction = C.MakeAction<SYCLPostLinkJobAction>(
+              DeviceLinkAction, types::TY_Tempfilelist);
+          auto *EntryGenAction = C.MakeAction<SYCLPostLinkJobAction>(
+              DeviceLinkAction, types::TY_TempEntriesfilelist);
+          SPIRVInput = SplitAction;
+          WrapperInputs.push_back(EntryGenAction);
+          OutType = types::TY_Tempfilelist;
+        }
+        auto *SPIRVTranslateAction =
+            C.MakeAction<SPIRVTranslatorJobAction>(SPIRVInput, OutType);
+
         auto TT = SYCLTripleList[I];
         bool SYCLAOTCompile =
             (TT.getSubArch() != llvm::Triple::NoSubArch &&
@@ -3480,21 +3505,26 @@ class OffloadingActionBuilder final {
 
         // After the Link, wrap the files before the final host link
         if (SYCLAOTCompile) {
-          types::ID OutType = types::TY_Image;
-          if (TT.getSubArch() == llvm::Triple::SPIRSubArch_fpga)
-            OutType = FPGAOutType;
+          OutType = types::TY_Tempfilelist;
+          if (!DeviceCodeSplit) {
+            OutType = (TT.getSubArch() == llvm::Triple::SPIRSubArch_fpga)
+                          ? FPGAOutType
+                          : types::TY_Image;
+          }
           // Do the additional Ahead of Time compilation when the specific
           // triple calls for it (provided a valid subarch).
-          auto *DeviceBECompileAction =
-              C.MakeAction<BackendCompileJobAction>(DeviceLinkAction, OutType);
+          auto *DeviceBECompileAction = C.MakeAction<BackendCompileJobAction>(
+              SPIRVTranslateAction, OutType);
 
+          WrapperInputs.push_back(DeviceBECompileAction);
           auto *DeviceWrappingAction = C.MakeAction<OffloadWrapperJobAction>(
-              DeviceBECompileAction, types::TY_Object);
+              WrapperInputs, types::TY_Object);
           DA.add(*DeviceWrappingAction, **TC, /*BoundArch=*/nullptr,
                  Action::OFK_SYCL);
         } else {
+          WrapperInputs.push_back(SPIRVTranslateAction);
           auto *DeviceWrappingAction = C.MakeAction<OffloadWrapperJobAction>(
-              DeviceLinkAction, types::TY_Object);
+              WrapperInputs, types::TY_Object);
           DA.add(*DeviceWrappingAction, **TC, /*BoundArch=*/nullptr,
                  Action::OFK_SYCL);
         }
@@ -3538,6 +3568,7 @@ class OffloadingActionBuilder final {
       Arg *SYCLLinkTargets = Args.getLastArg(
                                   options::OPT_fsycl_link_targets_EQ);
       WrapDeviceOnlyBinary = Args.hasArg(options::OPT_fsycl_link_EQ);
+      DeviceCodeSplit = Args.hasArg(options::OPT_fsycl_device_code_split);
       // Device only compilation for -fsycl-link (no FPGA) and
       // -fsycl-link-targets
       CompileDeviceOnly =
@@ -5407,7 +5438,7 @@ const char *Driver::GetNamedOutputPath(Compilation &C, const JobAction &JA,
     } else {
       TmpName = GetTemporaryPath(Split.first, Suffix);
     }
-    return C.addTempFile(C.getArgs().MakeArgString(TmpName));
+    return C.addTempFile(C.getArgs().MakeArgString(TmpName), JA.getType());
   }
 
   SmallString<128> BasePath(BaseInput);

clang/lib/Driver/ToolChain.cpp

@@ -322,6 +322,12 @@ Tool *ToolChain::getSPIRCheck() const {
   return SPIRCheck.get();
 }
 
+Tool *ToolChain::getSYCLPostLink() const {
+  if (!SYCLPostLink)
+    SYCLPostLink.reset(new tools::SYCLPostLink(*this));
+  return SYCLPostLink.get();
+}
+
 Tool *ToolChain::getBackendCompiler() const {
   if (!BackendCompiler)
     BackendCompiler.reset(buildBackendCompiler());
@@ -370,6 +376,9 @@ Tool *ToolChain::getTool(Action::ActionClass AC) const {
   case Action::SPIRCheckJobClass:
     return getSPIRCheck();
 
+  case Action::SYCLPostLinkJobClass:
+    return getSYCLPostLink();
+
   case Action::BackendCompileJobClass:
     return getBackendCompiler();
   }