[SYCL-RTC] Use ToolInvocation directly instead of via ClangTool (intel#19922)

We only run SYCL JIT on a single TU at a time, so using `ClangTool` is a
bit awkward, as it is primarily used to run the same action across a set
of files:


https://github.com/intel/llvm/blob/357f96b7e19d8acb972eb2f1fb276dbc6aa2060b/clang/include/clang/Tooling/Tooling.h#L310-L317

Using `ToolInvocation` better matches our scenario of always doing a
single clang invocation:


https://github.com/intel/llvm/blob/357f96b7e19d8acb972eb2f1fb276dbc6aa2060b/clang/include/clang/Tooling/Tooling.h#L244-L245

Another benefit is that we have more control over the virtual file
system which I'm planning to use in a subsequent PR to have the SYCL
toolchain headers distributed inside `libsycl-jit.so` and then put into
an `llvm::vfs::InMemoryFileSystem` once to be re-used across all
compilation queries.

I'm also simplifying the inheritance scheme around `clang::ToolAction`.
Instead of having both hashing/compiling doing that, I'm providing a
single helper that accepts a reference to the `FrontendAction` that can
be kept on the caller's stack, reducing the amount of boilerplate
helpers necessary, i.e.
`RTCToolActionBase`/`GetSourceHashAction`/`GetLLVMModuleAction` before
vs. single `SYCLToolchain::Action` after.

Author: aelovikov-intel

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

@@ -177,101 +177,76 @@ class HashPreprocessedAction : public PreprocessorFrontendAction {
   bool Executed = false;
 };
 
-class RTCToolActionBase : public ToolAction {
-public:
-  // Code adapted from `FrontendActionFactory::runInvocation`.
-  bool runInvocation(std::shared_ptr<CompilerInvocation> Invocation,
-                     FileManager *Files,
-                     std::shared_ptr<PCHContainerOperations> PCHContainerOps,
-                     DiagnosticConsumer *DiagConsumer) override {
-    assert(!hasExecuted() && "Action should only be invoked on a single file");
-
-    // Create a compiler instance to handle the actual work.
-    CompilerInstance Compiler(std::move(Invocation), std::move(PCHContainerOps));
-    Compiler.setFileManager(Files);
-    // Suppress summary with number of warnings and errors being printed to
-    // stdout.
-    Compiler.setVerboseOutputStream(std::make_unique<llvm::raw_null_ostream>());
-
-    // Create the compiler's actual diagnostics engine.
-    Compiler.createDiagnostics(Files->getVirtualFileSystem(), DiagConsumer,
-                               /*ShouldOwnClient=*/false);
-    if (!Compiler.hasDiagnostics()) {
-      return false;
-    }
-
-    Compiler.createSourceManager(*Files);
-
-    return executeAction(Compiler, Files);
-  }
-
-  virtual ~RTCToolActionBase() = default;
-
-protected:
-  virtual bool hasExecuted() = 0;
-  virtual bool executeAction(CompilerInstance &, FileManager *) = 0;
-};
-
-class GetSourceHashAction : public RTCToolActionBase {
-protected:
-  bool executeAction(CompilerInstance &CI, FileManager *Files) override {
-    HashPreprocessedAction HPA;
-    const bool Success = CI.ExecuteAction(HPA);
-    Files->clearStatCache();
-    if (!Success) {
-      return false;
+class SYCLToolchain {
+  SYCLToolchain() {}
+
+  // 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`.
+  class Action : public ToolAction {
+    FrontendAction &FEAction;
+
+  public:
+    Action(FrontendAction &FEAction) : FEAction(FEAction) {}
+    ~Action() override = default;
+
+    // Code adapted from `FrontendActionFactory::runInvocation`:
+    bool runInvocation(std::shared_ptr<CompilerInvocation> Invocation,
+                       FileManager *Files,
+                       std::shared_ptr<PCHContainerOperations> PCHContainerOps,
+                       DiagnosticConsumer *DiagConsumer) override {
+      // Create a compiler instance to handle the actual work.
+      CompilerInstance Compiler(std::move(Invocation),
+                                std::move(PCHContainerOps));
+      Compiler.setFileManager(Files);
+      // Suppress summary with number of warnings and errors being printed to
+      // stdout.
+      Compiler.setVerboseOutputStream(
+          std::make_unique<llvm::raw_null_ostream>());
+
+      // Create the compiler's actual diagnostics engine.
+      Compiler.createDiagnostics(Files->getVirtualFileSystem(), DiagConsumer,
+                                 /*ShouldOwnClient=*/false);
+      if (!Compiler.hasDiagnostics())
+        return false;
+
+      Compiler.createSourceManager(*Files);
+
+      const bool Success = Compiler.ExecuteAction(FEAction);
+
+      Files->clearStatCache();
+      return Success;
     }
-
-    Hash = HPA.takeHash();
-    Executed = true;
-    return true;
-  }
-
-  bool hasExecuted() override { return Executed; }
+  };
 
 public:
-  BLAKE3Result<> takeHash() {
-    assert(Executed);
-    Executed = false;
-    return std::move(Hash);
+  static SYCLToolchain &instance() {
+    static SYCLToolchain Instance;
+    return Instance;
   }
 
-private:
-  BLAKE3Result<> Hash;
-  bool Executed = false;
-};
+  bool run(const std::vector<std::string> &CommandLine,
+           FrontendAction &FEAction,
+           IntrusiveRefCntPtr<FileSystem> FSOverlay = nullptr,
+           DiagnosticConsumer *DiagConsumer = nullptr) {
+    auto FS = llvm::makeIntrusiveRefCnt<llvm::vfs::OverlayFileSystem>(
+        llvm::vfs::getRealFileSystem());
+    if (FSOverlay)
+      FS->pushOverlay(FSOverlay);
 
-struct GetLLVMModuleAction : public RTCToolActionBase {
-protected:
-  bool executeAction(CompilerInstance &CI, FileManager *Files) override {
-    // Ignore `Compiler.getFrontendOpts().ProgramAction` (would be `EmitBC`) and
-    // create/execute an `EmitLLVMOnlyAction` (= codegen to LLVM module without
-    // emitting anything) instead.
-    EmitLLVMOnlyAction ELOA{&Context};
-    const bool Success = CI.ExecuteAction(ELOA);
-    Files->clearStatCache();
-    if (!Success) {
-      return false;
-    }
+    auto Files = llvm::makeIntrusiveRefCnt<clang::FileManager>(
+        clang::FileSystemOptions{"." /* WorkingDir */}, FS);
 
-    // Take the module to extend its lifetime.
-    Module = ELOA.takeModule();
+    Action A{FEAction};
+    ToolInvocation TI{CommandLine, &A, Files.get(),
+                      std::make_shared<PCHContainerOperations>()};
+    TI.setDiagnosticConsumer(DiagConsumer ? DiagConsumer : &IgnoreDiag);
 
-    return true;
-  }
-
-  bool hasExecuted() override { return static_cast<bool>(Module); }
-
-public:
-  GetLLVMModuleAction(LLVMContext &Context) : Context{Context}, Module{} {}
-  ModuleUPtr takeModule() {
-    assert(Module);
-    return std::move(Module);
+    return TI.run();
   }
 
 private:
-  LLVMContext &Context;
-  ModuleUPtr Module;
+  clang::IgnoringDiagConsumer IgnoreDiag;
 };
 
 class ClangDiagnosticWrapper {
@@ -320,9 +295,9 @@ class LLVMDiagnosticWrapper : public llvm::DiagnosticHandler {
 
 } // anonymous namespace
 
-static void adjustArgs(const InputArgList &UserArgList,
-                       const std::string &DPCPPRoot, BinaryFormat Format,
-                       SmallVectorImpl<std::string> &CommandLine) {
+static std::vector<std::string>
+createCommandLine(const InputArgList &UserArgList, std::string_view DPCPPRoot,
+                  BinaryFormat Format, std::string_view SourceFilePath) {
   DerivedArgList DAL{UserArgList};
   const auto &OptTable = getDriverOptTable();
   DAL.AddFlagArg(nullptr, OptTable.getOption(OPT_fsycl_device_only));
@@ -349,36 +324,30 @@ static void adjustArgs(const InputArgList &UserArgList,
   for_each(DAL, [&DAL, &ASL](Arg *A) { A->render(DAL, ASL); });
   for_each(UserArgList,
            [&UserArgList, &ASL](Arg *A) { A->render(UserArgList, ASL); });
+
+  std::vector<std::string> CommandLine;
+  CommandLine.reserve(ASL.size() + 2);
+  CommandLine.emplace_back((DPCPPRoot + "/bin/clang++").str());
   transform(ASL, std::back_inserter(CommandLine),
             [](const char *AS) { return std::string{AS}; });
+  CommandLine.emplace_back(SourceFilePath);
+  return CommandLine;
 }
 
-static void setupTool(ClangTool &Tool, const std::string &DPCPPRoot,
-                      InMemoryFile SourceFile, View<InMemoryFile> IncludeFiles,
-                      DiagnosticConsumer *Consumer) {
-  Tool.setDiagnosticConsumer(Consumer);
-  // Suppress message "Error while processing" being printed to stdout.
-  Tool.setPrintErrorMessage(false);
-
-  // Set up in-memory filesystem.
-  Tool.mapVirtualFile(SourceFile.Path, SourceFile.Contents);
-  for (const auto &IF : IncludeFiles) {
-    Tool.mapVirtualFile(IF.Path, IF.Contents);
-  }
+static llvm::IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem>
+getInMemoryFS(InMemoryFile SourceFile, View<InMemoryFile> IncludeFiles) {
+  auto InMemoryFS = llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
+
+  InMemoryFS->setCurrentWorkingDirectory(
+      *llvm::vfs::getRealFileSystem()->getCurrentWorkingDirectory());
 
-  // Reset argument adjusters to drop the `-fsyntax-only` flag which is added by
-  // default by this API.
-  Tool.clearArgumentsAdjusters();
-  // Then, modify argv[0] so that the driver picks up the correct SYCL
-  // environment. We've already set the resource directory above.
-  Tool.appendArgumentsAdjuster(
-      [&DPCPPRoot](const CommandLineArguments &Args,
-                   StringRef Filename) -> CommandLineArguments {
-        (void)Filename;
-        CommandLineArguments NewArgs = Args;
-        NewArgs[0] = (Twine(DPCPPRoot) + "/bin/clang++").str();
-        return NewArgs;
-      });
+  InMemoryFS->addFile(SourceFile.Path, 0,
+                      llvm::MemoryBuffer::getMemBuffer(SourceFile.Contents));
+  for (InMemoryFile F : IncludeFiles)
+    InMemoryFS->addFile(F.Path, 0,
+                        llvm::MemoryBuffer::getMemBuffer(F.Contents));
+
+  return InMemoryFS;
 }
 
 Expected<std::string> jit_compiler::calculateHash(
@@ -391,20 +360,20 @@ Expected<std::string> jit_compiler::calculateHash(
     return createStringError("Could not locate DPCPP root directory");
   }
 
-  SmallVector<std::string> CommandLine;
-  adjustArgs(UserArgList, DPCPPRoot, Format, CommandLine);
+  std::vector<std::string> CommandLine =
+      createCommandLine(UserArgList, DPCPPRoot, Format, SourceFile.Path);
 
-  FixedCompilationDatabase DB{".", CommandLine};
-  ClangTool Tool{DB, {SourceFile.Path}};
+  HashPreprocessedAction HashAction;
 
-  clang::IgnoringDiagConsumer DiagConsumer;
-  setupTool(Tool, DPCPPRoot, SourceFile, IncludeFiles, &DiagConsumer);
+  if (SYCLToolchain::instance().run(CommandLine, HashAction,
+                                    getInMemoryFS(SourceFile, IncludeFiles))) {
+    BLAKE3Result<> SourceHash = HashAction.takeHash();
+    // Last argument is the source file in the format `rtc_N.cpp` which is
+    // unique for each query, so drop it:
+    CommandLine.pop_back();
 
-  GetSourceHashAction Action;
-  if (!Tool.run(&Action)) {
-    BLAKE3Result<> SourceHash = Action.takeHash();
-    // The adjusted command line contains the DPCPP root and clang major
-    // version.
+    // The command line contains the DPCPP root and clang major version in
+    // "-resource-dir=<...>" argument.
     BLAKE3Result<> CommandLineHash =
         BLAKE3::hash(arrayRefFromStringRef(join(CommandLine, ",")));
 
@@ -413,9 +382,10 @@ Expected<std::string> jit_compiler::calculateHash(
     // Make the encoding filesystem-friendly.
     std::replace(EncodedHash.begin(), EncodedHash.end(), '/', '-');
     return std::move(EncodedHash);
-  }
 
-  return createStringError("Calculating source hash failed");
+  } else {
+    return createStringError("Calculating source hash failed");
+  }
 }
 
 Expected<ModuleUPtr> jit_compiler::compileDeviceCode(
@@ -429,23 +399,17 @@ Expected<ModuleUPtr> jit_compiler::compileDeviceCode(
     return createStringError("Could not locate DPCPP root directory");
   }
 
-  SmallVector<std::string> CommandLine;
-  adjustArgs(UserArgList, DPCPPRoot, Format, CommandLine);
-
-  FixedCompilationDatabase DB{".", CommandLine};
-  ClangTool Tool{DB, {SourceFile.Path}};
-
+  EmitLLVMOnlyAction ELOA{&Context};
   DiagnosticOptions DiagOpts;
   ClangDiagnosticWrapper Wrapper(BuildLog, &DiagOpts);
 
-  setupTool(Tool, DPCPPRoot, SourceFile, IncludeFiles, Wrapper.consumer());
-
-  GetLLVMModuleAction Action{Context};
-  if (!Tool.run(&Action)) {
-    return Action.takeModule();
+  if (SYCLToolchain::instance().run(
+          createCommandLine(UserArgList, DPCPPRoot, Format, SourceFile.Path),
+          ELOA, getInMemoryFS(SourceFile, IncludeFiles), Wrapper.consumer())) {
+    return ELOA.takeModule();
+  } else {
+    return createStringError(BuildLog);
   }
-
-  return createStringError(BuildLog);
 }
 
 // This function is a simplified copy of the device library selection process