[DRAFT][SYCL RTC] Implement --auto-pch support

Author: aelovikov-intel

clang/include/clang/Driver/Options.h

@@ -41,6 +41,7 @@ enum ClangVisibility {
   FlangOption = (1 << 4),
   FC1Option = (1 << 5),
   DXCOption = (1 << 6),
+  SYCLRTCOnlyOption = (1 << 7),
 };
 
 enum ID {

clang/include/clang/Driver/Options.td

@@ -107,6 +107,8 @@ def FC1Option : OptionVisibility;
 // are made available when the driver is running in DXC compatibility mode.
 def DXCOption : OptionVisibility;
 
+def SYCLRTCOnlyOption : OptionVisibility;
+
 /////////
 // Docs
 
@@ -195,6 +197,11 @@ def sycl_Group : OptionGroup<"<SYCL group>">, Group<f_Group>,
                  DocName<"SYCL options">,
                  Visibility<[ClangOption, CLOption]>;
 
+def sycl_rtc_only_Group : OptionGroup<"<SYCL RTC only group">,
+                          Group<f_Group>,
+                          DocName<"SYCL RTC specific options">,
+                          Visibility<[SYCLRTCOnlyOption]>;
+
 def cuda_Group : OptionGroup<"<CUDA group>">, Group<f_Group>,
                    DocName<"CUDA options">,
                    Visibility<[ClangOption, CLOption]>;
@@ -7460,6 +7467,13 @@ def fsyclbin : Flag<["-"], "fsyclbin">, Alias<fsyclbin_EQ>,
   AliasArgs<["executable"]>;
 } // let Group = sycl_Group
 
+let Visibility = [SYCLRTCOnlyOption] in {
+  let Group = sycl_rtc_only_Group in {
+    def auto_pch : Flag<["--"], "auto-pch">,
+                   HelpText<"Enable Auto-PCH for SYCL RTC Compilation">;
+  } // let Group = sycl_rtc_only_Group
+} // let Visibility = [SYCLRTCOnlyOption]
+
 // FIXME: -fsycl-explicit-simd is deprecated. remove it when support is dropped.
 def : Flag<["-"], "fsycl-explicit-simd">, Flags<[Deprecated]>,
   Group<clang_ignored_legacy_options_Group>,

clang/include/clang/Frontend/PrecompiledPreamble.h

@@ -87,8 +87,8 @@ class PrecompiledPreamble {
         DiagnosticsEngine &Diagnostics,
         IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
         std::shared_ptr<PCHContainerOperations> PCHContainerOps,
-        bool StoreInMemory, StringRef StoragePath,
-        PreambleCallbacks &Callbacks);
+        bool StoreInMemory, StringRef StoragePath, PreambleCallbacks &Callbacks,
+        bool AllowASTWithErrors = true);
 
   PrecompiledPreamble(PrecompiledPreamble &&);
   PrecompiledPreamble &operator=(PrecompiledPreamble &&);

clang/lib/Frontend/PrecompiledPreamble.cpp

@@ -249,9 +249,10 @@ class TempPCHFile {
 class PrecompilePreambleAction : public ASTFrontendAction {
 public:
   PrecompilePreambleAction(std::shared_ptr<PCHBuffer> Buffer, bool WritePCHFile,
-                           PreambleCallbacks &Callbacks)
+                           PreambleCallbacks &Callbacks,
+                           bool AllowASTWithErrors = true)
       : Buffer(std::move(Buffer)), WritePCHFile(WritePCHFile),
-        Callbacks(Callbacks) {}
+        Callbacks(Callbacks), AllowASTWithErrors(AllowASTWithErrors) {}
 
   std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                  StringRef InFile) override;
@@ -287,16 +288,18 @@ class PrecompilePreambleAction : public ASTFrontendAction {
   bool WritePCHFile; // otherwise the PCH is written into the PCHBuffer only.
   std::unique_ptr<llvm::raw_pwrite_stream> FileOS; // null if in-memory
   PreambleCallbacks &Callbacks;
+  bool AllowASTWithErrors;
 };
 
 class PrecompilePreambleConsumer : public PCHGenerator {
 public:
   PrecompilePreambleConsumer(PrecompilePreambleAction &Action, Preprocessor &PP,
                              ModuleCache &ModCache, StringRef isysroot,
-                             std::shared_ptr<PCHBuffer> Buffer)
+                             std::shared_ptr<PCHBuffer> Buffer,
+                             bool AllowASTWithErrors = true)
       : PCHGenerator(PP, ModCache, "", isysroot, std::move(Buffer),
                      ArrayRef<std::shared_ptr<ModuleFileExtension>>(),
-                     /*AllowASTWithErrors=*/true),
+                     AllowASTWithErrors),
         Action(Action) {}
 
   bool HandleTopLevelDecl(DeclGroupRef DG) override {
@@ -337,7 +340,8 @@ PrecompilePreambleAction::CreateASTConsumer(CompilerInstance &CI,
     Sysroot.clear();
 
   return std::make_unique<PrecompilePreambleConsumer>(
-      *this, CI.getPreprocessor(), CI.getModuleCache(), Sysroot, Buffer);
+      *this, CI.getPreprocessor(), CI.getModuleCache(), Sysroot, Buffer,
+      AllowASTWithErrors);
 }
 
 template <class T> bool moveOnNoError(llvm::ErrorOr<T> Val, T &Output) {
@@ -415,7 +419,8 @@ llvm::ErrorOr<PrecompiledPreamble> PrecompiledPreamble::Build(
     DiagnosticsEngine &Diagnostics,
     IntrusiveRefCntPtr<llvm::vfs::FileSystem> VFS,
     std::shared_ptr<PCHContainerOperations> PCHContainerOps, bool StoreInMemory,
-    StringRef StoragePath, PreambleCallbacks &Callbacks) {
+    StringRef StoragePath, PreambleCallbacks &Callbacks,
+    bool AllowASTWithErrors) {
   assert(VFS && "VFS is null");
 
   auto PreambleInvocation = std::make_shared<CompilerInvocation>(Invocation);
@@ -511,7 +516,7 @@ llvm::ErrorOr<PrecompiledPreamble> PrecompiledPreamble::Build(
   auto Act = std::make_unique<PrecompilePreambleAction>(
       std::move(Buffer),
       /*WritePCHFile=*/Storage->getKind() == PCHStorage::Kind::TempFile,
-      Callbacks);
+      Callbacks, AllowASTWithErrors);
   if (!Act->BeginSourceFile(*Clang, Clang->getFrontendOpts().Inputs[0]))
     return BuildPreambleError::BeginSourceFileFailed;
 

clang/test/Driver/sycl-unsupported.cpp

@@ -64,6 +64,15 @@
 // UNSUPPORTED_OPT-NOT: clang{{.*}} "-fsycl-is-device"{{.*}} "[[OPT_CC1]]{{.*}}"
 // UNSUPPORTED_OPT: clang{{.*}} "-fsycl-is-host"{{.*}} "[[OPT_CC1]]{{.*}}"
 
+// "--auto-pch" should only be enabled for SYCL RTC compilations, regular driver
+// shouldn't know about it:
+//
+// RUN: not %clangxx                    -### %s --auto-pch 2>&1 | FileCheck %s --check-prefix AUTO_PCH
+// RUN: not %clangxx -fsycl-device-only -### %s --auto-pch 2>&1 | FileCheck %s --check-prefix AUTO_PCH
+// RUN: not %clangxx -fsycl             -### %s --auto-pch 2>&1 | FileCheck %s --check-prefix AUTO_PCH
+//
+// AUTO_PCH: error: unknown argument: '--auto-pch'
+
 // FPGA support has been removed, usage of any FPGA specific options and any
 // options that have FPGA specific arguments should emit a specific error
 // diagnostic.

sycl-jit/jit-compiler/lib/rtc/DeviceCompilation.cpp

@@ -25,6 +25,7 @@
 #include <clang/Frontend/ChainedDiagnosticConsumer.h>
 #include <clang/Frontend/CompilerInstance.h>
 #include <clang/Frontend/FrontendActions.h>
+#include <clang/Frontend/PrecompiledPreamble.h>
 #include <clang/Frontend/TextDiagnosticBuffer.h>
 #include <clang/Frontend/TextDiagnosticPrinter.h>
 #include <clang/Frontend/Utils.h>
@@ -78,6 +79,12 @@ class SYCLToolchain {
     }
   }
 
+  struct PrecompiledPreambles {
+    using key = std::pair<std::string /*Opts*/, std::string /*Preamble*/>;
+    std::mutex Mutex;
+    std::map<key, std::shared_ptr<PrecompiledPreamble>> PreamblesMap;
+  };
+
   // Similar to FrontendActionFactory, but we don't take ownership of
   // `FrontendAction`, nor do we create copies of it as we only perform a single
   // `ToolInvocation`.
@@ -140,9 +147,15 @@ class SYCLToolchain {
     }
 
     ArgStringList ASL;
-    for_each(DAL, [&DAL, &ASL](Arg *A) { A->render(DAL, ASL); });
-    for_each(UserArgList,
-             [&UserArgList, &ASL](Arg *A) { A->render(UserArgList, ASL); });
+    for (Arg *A : DAL)
+      A->render(DAL, ASL);
+    for (Arg *A : UserArgList) {
+      Option Group = A->getOption().getGroup();
+      if (Group.isValid() && Group.getID() == OPT_sycl_rtc_only_Group)
+        continue;
+
+      A->render(UserArgList, ASL);
+    }
 
     std::vector<std::string> CommandLine;
     CommandLine.reserve(ASL.size() + 2);
@@ -153,6 +166,79 @@ class SYCLToolchain {
     return CommandLine;
   }
 
+  class ActionWithPCHPreamble : public Action {
+    std::string CmdLineOpts;
+
+  public:
+    ActionWithPCHPreamble(FrontendAction &FEAction, std::string &&CmdLineOpts)
+        : Action(FEAction), CmdLineOpts(std::move(CmdLineOpts)) {}
+
+    bool runInvocation(std::shared_ptr<CompilerInvocation> Invocation,
+                       FileManager *Files,
+                       std::shared_ptr<PCHContainerOperations> PCHContainerOps,
+                       DiagnosticConsumer *DiagConsumer) override {
+      auto MainFilePath = Invocation->getFrontendOpts().Inputs[0].getFile();
+      auto MainFileBuffer = Files->getBufferForFile(MainFilePath);
+      assert(MainFileBuffer && "Can't get memory buffer for in-memory source?");
+
+      PreambleBounds Bounds = ComputePreambleBounds(
+          Invocation->getLangOpts(), **MainFileBuffer, 100 /* MaxLines */);
+
+      PrecompiledPreambles::key key{
+          std::move(CmdLineOpts),
+          (*MainFileBuffer)->getBuffer().substr(0, Bounds.Size).str()};
+
+      std::shared_ptr<PrecompiledPreamble> Preamble;
+      {
+        PrecompiledPreambles &Preambles = SYCLToolchain::instance().Preambles;
+        std::lock_guard<std::mutex> Lock{Preambles.Mutex};
+        auto [It, Inserted] = Preambles.PreamblesMap.try_emplace(key);
+
+        if (Inserted) {
+          PreambleCallbacks Callbacks;
+          auto DiagIds = llvm::makeIntrusiveRefCnt<DiagnosticIDs>();
+          auto DiagOpts = Invocation->getDiagnosticOpts();
+          auto Diags = llvm::makeIntrusiveRefCnt<DiagnosticsEngine>(
+              DiagIds, DiagOpts, DiagConsumer, false);
+
+          static std::string StoragePath =
+              (SYCLToolchain::instance().getPrefix() + "/preambles").str();
+          llvm::ErrorOr<PrecompiledPreamble> NewPreamble =
+              PrecompiledPreamble::Build(
+                  *Invocation, MainFileBuffer->get(), Bounds, *Diags,
+                  Files->getVirtualFileSystemPtr(), PCHContainerOps,
+                  /*StorePreamblesInMemory*/ true, StoragePath, Callbacks,
+                  /*AllowASTWithErrors=*/false);
+
+          if (!NewPreamble)
+            return false;
+
+          It->second = std::make_shared<PrecompiledPreamble>(
+              std::move(NewPreamble.get()));
+        }
+
+        Preamble = It->second;
+      } // End lock
+
+      assert(Preamble);
+      assert(Preamble->CanReuse(*Invocation, **MainFileBuffer, Bounds,
+                                Files->getVirtualFileSystem()));
+
+      // FIXME: WHY release????
+      auto Buf = llvm::MemoryBuffer::getMemBufferCopy(
+                     (*MainFileBuffer)->getBuffer(), MainFilePath)
+                     .release();
+
+      auto VFS = Files->getVirtualFileSystemPtr();
+      Preamble->AddImplicitPreamble(*Invocation, VFS, Buf);
+      auto NewFiles = makeIntrusiveRefCnt<FileManager>(
+          Files->getFileSystemOpts(), std::move(VFS));
+
+      return Action::runInvocation(std::move(Invocation), NewFiles.get(),
+                                   std::move(PCHContainerOps), DiagConsumer);
+    }
+  };
+
 public:
   static SYCLToolchain &instance() {
     static SYCLToolchain Instance;
@@ -162,7 +248,8 @@ class SYCLToolchain {
   bool run(const InputArgList &UserArgList, BinaryFormat Format,
            const char *SourceFilePath, FrontendAction &FEAction,
            IntrusiveRefCntPtr<FileSystem> FSOverlay = nullptr,
-           DiagnosticConsumer *DiagConsumer = nullptr) {
+           DiagnosticConsumer *DiagConsumer = nullptr,
+           bool UseAutoPCH = false) {
     std::vector<std::string> CommandLine =
         createCommandLine(UserArgList, Format, SourceFilePath);
 
@@ -175,9 +262,14 @@ class SYCLToolchain {
     auto Files = llvm::makeIntrusiveRefCnt<clang::FileManager>(
         clang::FileSystemOptions{"." /* WorkingDir */}, FS);
 
-    Action A{FEAction};
-    ToolInvocation TI{CommandLine, &A, Files.get(),
-                      std::make_shared<PCHContainerOperations>()};
+    Action Normal{FEAction};
+    ActionWithPCHPreamble WithPreamble{FEAction,
+                                       join(drop_end(CommandLine, 1), " ")};
+    ToolInvocation TI{CommandLine,
+                      UseAutoPCH ? static_cast<Action *>(&WithPreamble)
+                                 : &Normal,
+                      Files.get(), std::make_shared<PCHContainerOperations>()};
+
     TI.setDiagnosticConsumer(DiagConsumer ? DiagConsumer : &IgnoreDiag);
 
     return TI.run();
@@ -217,6 +309,8 @@ class SYCLToolchain {
   std::string ClangXXExe = (Prefix + "/bin/clang++").str();
   llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> ToolchainFS =
       llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
+
+  PrecompiledPreambles Preambles;
 };
 
 class ClangDiagnosticWrapper {
@@ -348,9 +442,11 @@ Expected<ModuleUPtr> jit_compiler::compileDeviceCode(
   DiagnosticOptions DiagOpts;
   ClangDiagnosticWrapper Wrapper(BuildLog, &DiagOpts);
 
+  bool AutoPCH = UserArgList.hasArg(OPT_auto_pch);
+
   if (SYCLToolchain::instance().run(UserArgList, Format, SourceFile.Path, ELOA,
                                     getInMemoryFS(SourceFile, IncludeFiles),
-                                    Wrapper.consumer())) {
+                                    Wrapper.consumer(), AutoPCH)) {
     return ELOA.takeModule();
   } else {
     return createStringError(BuildLog);

sycl/doc/extensions/experimental/sycl_ext_oneapi_kernel_compiler.asciidoc

@@ -1127,6 +1127,58 @@ build_options{{
 Relax the requirement that parameter types for free-function kernels must be
 forward-declarable.
 
+===== `--auto-pch`
+
+Enable auto-detection of the preamble and use it as a pre-compiled header to
+speed up subsequent compilations of TUs matching the preamble/compilation
+options. Example of the code that can benefit from this:
+
+[source,c++]
+----
+#include <sycl/ext/oneapi/free_function_queries.hpp>
+#include <sycl/ext/oneapi/kernel_properties/properties.hpp>
+
+// Auto-detected preamble ends before next line:
+namespace syclext = sycl::ext::oneapi;
+namespace syclexp = sycl::ext::oneapi::experimental;
+
+extern "C"
+SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<1>))
+void iota(float start, float *ptr) {
+    size_t id = syclext::this_work_item::get_nd_item<1>().get_global_linear_id();
+    ptr[id] = start + static_cast<float>(id);
+}
+----
+
+Limitations:
+
+* Preamble detection is done at the Lexer level and can't handle code like
+
+[source,c++]
+----
+#if 1
+#include <sycl/sycl.hpp>
+#else
+// Auto-detected preamble ends in the middle of `#else` and would fail to compile.
+void foo() {}
+#endif
+----
+
+* Any changes in either preamble or compilation options (including
+  `-DSOMETHING`!) result in a creation of a new pre-compiled header/preamble.
+
+* No support (including not reporting any errors) for `+__DATE__+`/`+__TIME__+`
+  macros inside auto-detected preamble (transitively in regards to the
+  includes).
+
+* Files used inside preamble must not change between different compilations (at
+  least for the same auto-detected preamble).
+
+* Auto-generated pre-compiled headers/preambles are stored in memory only. That means:
+  - No persistency between invocations
+  - Currently there is no eviction mechanism, so application is expected to use
+     the option only when number of preambles is limited.
+
 === Known issues and limitations when the language is `sycl`
 
 ==== Changing the compiler action or output