[SYCL RTC] Introduce --auto-pch support (#20226)

Compilation of `#include <sycl/sycl.hpp>` is slow and that's especially
problematic for SYCL RTC (run-time compilation). One way to overcome
this is fine-grained includes that are being pursued separately. Another
way is to employ clang's precompiled headers support which this PR is
doing. Those two approaches can be combined, and this PR adds
`test-e2e/PerformanceTests/KernelCompiler/auto-pch.cpp` that gives some
idea of the PCH impact. The test shows PCH benefits when compiling some
of the fine-grained includes on top of absolute minimum required to
compiled SYCL RTC's "Hello world". From one of the CI runs:

  | Extra Headers | Without PCH | With auto-PCH
  | -|-|-
  | <none> | 176ms 137ms 136ms 136ms 136ms | 226ms 64ms 64ms 64ms 64ms 
| sycl/half_type.hpp | 165ms 165ms 165ms 165ms 165ms | 267ms 71ms 72ms
72ms 72ms
| sycl/ext/oneapi/bfloat16.hpp | 174ms 173ms 173ms 173ms 173ms | 279ms
76ms 73ms 73ms 74ms
| sycl/marray.hpp | 142ms 143ms 142ms 142ms 143ms | 235ms 66ms 66ms 66ms
66ms
| sycl/vector.hpp | 296ms 290ms 290ms 290ms 290ms | 487ms 124ms 125ms
125ms 125ms
| sycl/multi_ptr.hpp | 278ms 278ms 276ms 275ms 274ms | 441ms 125ms 125ms
125ms 125ms
| sycl/builtins.hpp | 537ms 533ms 531ms 531ms 531ms | 883ms 218ms 218ms
219ms 218ms

It misses `sycl/sycl.hpp` line because that currently crashes FE when
reading the generated PCH, the crash is being investigated/fixed
separately.

Implementation-wise I'm reusing existing upstream
`clang::PrecompiledPreamble` with one minor modification. It seems that
`PrecompiledPreamble`'s main usage is for things like `clangd` so it
ignores errors in the code. I've modified it so that those errors would
break pch-generation the same way normal compilation would break. I'm
also not sure if we'd want that long-term, because it seems that making
such "auto-pch" persistent would deviate from the upstream version of
`PrecompiledPreamble` even more. I can imagine that in some near future
we'd need to "fork" it into a separate utility. Still, seems to be fine
for the first step.

Driver modifications are for the `--auto-pch` option support that should
only be present on the SYCL RTC path and not for the regular `clang`
invocations from the command line. I'm relatively confident those will
stay in future.

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,9 @@ def FC1Option : OptionVisibility;
 // are made available when the driver is running in DXC compatibility mode.
 def DXCOption : OptionVisibility;
 
+// SYCLRTCOnlyOption - only acceptable for the SYCL RTC (Run Time Compilation).
+def SYCLRTCOnlyOption : OptionVisibility;
+
 /////////
 // Docs
 
@@ -195,6 +198,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]>;
@@ -7543,6 +7551,15 @@ def fsyclbin : Flag<["-"], "fsyclbin">, Alias<fsyclbin_EQ>,
   AliasArgs<["executable"]>;
 } // let Group = sycl_Group
 
+// Options specific to the SYCL RTC and only available for JIT compilation (not
+// through regular `clang++ -fsycl` in command line):
+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 {
         IntrusiveRefCntPtr<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

@@ -247,9 +247,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;
@@ -285,17 +286,19 @@ 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,
-                             const CodeGenOptions &CodeGenOpts)
+                             const CodeGenOptions &CodeGenOpts,
+                             bool AllowASTWithErrors = true)
       : PCHGenerator(PP, ModCache, "", isysroot, std::move(Buffer), CodeGenOpts,
                      ArrayRef<std::shared_ptr<ModuleFileExtension>>(),
-                     /*AllowASTWithErrors=*/true),
+                     AllowASTWithErrors),
         Action(Action) {}
 
   bool HandleTopLevelDecl(DeclGroupRef DG) override {
@@ -337,7 +340,7 @@ PrecompilePreambleAction::CreateASTConsumer(CompilerInstance &CI,
 
   return std::make_unique<PrecompilePreambleConsumer>(
       *this, CI.getPreprocessor(), CI.getModuleCache(), Sysroot, Buffer,
-      CI.getCodeGenOpts());
+      CI.getCodeGenOpts(), AllowASTWithErrors);
 }
 
 template <class T> bool moveOnNoError(llvm::ErrorOr<T> Val, T &Output) {
@@ -415,7 +418,8 @@ llvm::ErrorOr<PrecompiledPreamble> PrecompiledPreamble::Build(
     IntrusiveRefCntPtr<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);
@@ -512,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

@@ -12,6 +12,7 @@
 #include "Resource.h"
 #include "translation/Translation.h"
 
+#include "clang/Lex/PreprocessorOptions.h"
 #include <clang/Basic/DiagnosticDriver.h>
 #include <clang/Basic/Version.h>
 #include <clang/CodeGen/CodeGenAction.h>
@@ -25,6 +26,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 +80,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 +148,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 +167,83 @@ 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()));
+
+      assert(Invocation->getPreprocessorOpts().RetainRemappedFileBuffers ==
+             false);
+      // `PreprocessorOptions::RetainRemappedFileBuffers` defaults to false, so
+      // MemoryBuffer will be cleaned up by the CompilerInstance, thus
+      // `std::unique_ptr::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 +253,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 +267,21 @@ class SYCLToolchain {
     auto Files = llvm::makeIntrusiveRefCnt<clang::FileManager>(
         clang::FileSystemOptions{"." /* WorkingDir */}, FS);
 
-    Action A{FEAction};
-    ToolInvocation TI{std::move(CommandLine), &A, Files.get(),
-                      std::make_shared<PCHContainerOperations>()};
+    Action Normal{FEAction};
+
+    // User compilation options must be part of the key in the preambles map. We
+    // can either use "raw" user options or the "processed" from
+    // `createCommandLine` as long as we're consistent in what we're using.
+    // Current internal APIs pass `InputArgList` around instead of a single
+    // `std::string`, so it's easier to use `CommandLine`. Just make sure to
+    // drop `rtc_N.cpp` that is always different:
+    ActionWithPCHPreamble WithPreamble{FEAction,
+                                       join(drop_end(CommandLine, 1), " ")};
+    ToolInvocation TI{std::move(CommandLine),
+                      UseAutoPCH ? static_cast<Action *>(&WithPreamble)
+                                 : &Normal,
+                      Files.get(), std::make_shared<PCHContainerOperations>()};
+
     TI.setDiagnosticConsumer(DiagConsumer ? DiagConsumer : &IgnoreDiag);
 
     return TI.run();
@@ -217,6 +321,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 +454,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,60 @@ build_options{{
 Relax the requirement that parameter types for free-function kernels must be
 forward-declarable.
 
+===== `--auto-pch`
+
+The first time this option is passed, the compiler finds the initial set of
+preprocessor directives (e.g., `#define`/`#include`) and comments in the
+compiled source string (the preamble) and pre-compiles it. Essentialy, it
+behaves like a precompiled header containing that preamble. On subsequent
+compilations, if the compiled source string has the same preamble and the same
+compilation options are used, the precompiled preamble is used, which speeds up
+compilation.
+
+If the compiled source string has a different preamble or compilation options
+differ, a new precompiled preamble is generated, and that preamble can also be
+used to speed up subsequent compilations. These precompiled preambles are stored
+internally in memory, so they do not persist from one execution of the
+application to the next.
+
+The preamble ends with the first statement that is not a preprocessor directive
+or a comment. For example, in the code below, the preamble ends immediately
+before the namespace syclext = statement.
+
+[source,c++]
+----
+#define SYCL_SIMPLE_SWIZZLES
+#include <sycl/sycl.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(sycl::vec<int, 2> *p) {
+    size_t id = syclext::this_work_item::get_nd_item<1>().get_global_linear_id();
+    p[id] = p[id].xx();
+}
+----
+
+The compiler uses the following factors when deciding whether a previously
+generated precompiled preamble can be used:
+
+* The preamble must exactly match (including whitespace and comments).
+* The compilation options must match (including the same order and the same spelling).
+* There are also certain restrictions that the user must avoid:
+
+  - The content of each header file in the preamble must not change from one
+    compilation to another.
+  - It is not recommended to use the `+__DATE__+` or `+__TIME__+` macros in the
+    preamble header files. Depending on the circumstances, these macros may be
+    replaced with the date / time that corresponds to the time at which the
+    precompiled preamble was generated, rather than the time at which the source
+    string is compiled. See also the clang compiler options `-Wpch-date-time`
+    and `-Werror=pch-date-time`, which cause the compiler to diagnose a warning
+    or error in this scenario.
+
 === Known issues and limitations when the language is `sycl`
 
 ==== Changing the compiler action or output