Allow custom importing of files and syntactic sugar (#1752)

* Allow custom importing of files and syntactic sugar

* Fix build on older llvm vers

* Update sugar.cpp

* Update EnzymeClang.cpp

* fix

* fixup

* dump

* more printing

* print

* fixup

---------

Co-authored-by: Ivan Radanov Ivanov <ivanov.i.aa@m.titech.ac.jp>

enzyme/BUILD

@@ -142,6 +142,20 @@ gentbl(
  ],
 )
 
+gentbl(
+ name = "include-utils",
+ tbl_outs = [(
+ "-gen-header-strings",
+ "IncludeUtils.inc",
+ )],
+ tblgen = ":enzyme-tblgen",
+ td_file = "Enzyme/Clang/include_utils.td",
+ td_srcs = ["Enzyme/Clang/include_utils.td"],
+ deps = [
+ ":enzyme-tblgen",
+ ],
+)
+
 cc_library(
  name = "EnzymeStatic",
  srcs = glob(
@@ -167,6 +181,7 @@ cc_library(
  data = ["@llvm-project//clang:builtin_headers_gen"],
  visibility = ["//visibility:public"],
  deps = [
+ "include-utils",
  ":binop-derivatives",
  ":blas-attributor",
  ":blas-derivatives",

enzyme/Enzyme/CMakeLists.txt

@@ -37,6 +37,10 @@ add_public_tablegen_target(BlasDeclarationsIncGen)
 add_public_tablegen_target(BlasTAIncGen)
 add_public_tablegen_target(BlasDiffUseIncGen)
 
+set(LLVM_TARGET_DEFINITIONS Clang/include_utils.td)
+enzyme_tablegen(IncludeUtils.inc -gen-header-strings)
+add_public_tablegen_target(IncludeUtilsIncGen)
+
 include_directories(${CMAKE_CURRENT_BINARY_DIR})
 
 set(LLVM_LINK_COMPONENTS Demangle)
@@ -74,6 +78,7 @@ if (${Clang_FOUND})
  LLVM
  )
 target_compile_definitions(ClangEnzyme-${LLVM_VERSION_MAJOR} PUBLIC ENZYME_RUNPASS)
+add_dependencies(ClangEnzyme-${LLVM_VERSION_MAJOR} IncludeUtilsIncGen)
 endif()
  add_llvm_library( LLDEnzyme-${LLVM_VERSION_MAJOR}
  ${ENZYME_SRC} Clang/EnzymePassLoader.cpp
@@ -107,6 +112,7 @@ if (${Clang_FOUND})
  clang
  )
 target_compile_definitions(ClangEnzyme-${LLVM_VERSION_MAJOR} PUBLIC ENZYME_RUNPASS)
+add_dependencies(ClangEnzyme-${LLVM_VERSION_MAJOR} IncludeUtilsIncGen)
 endif()
  add_llvm_library( LLDEnzyme-${LLVM_VERSION_MAJOR}
  ${ENZYME_SRC} Clang/EnzymePassLoader.cpp

enzyme/Enzyme/Clang/EnzymeClang.cpp

@@ -25,16 +25,20 @@
 #include "clang/AST/Attr.h"
 #include "clang/AST/DeclGroup.h"
 #include "clang/AST/RecursiveASTVisitor.h"
+#include "clang/Basic/FileManager.h"
 #include "clang/Basic/MacroBuilder.h"
 #include "clang/Frontend/CompilerInstance.h"
 #include "clang/Frontend/FrontendAction.h"
 #include "clang/Frontend/FrontendPluginRegistry.h"
+#include "clang/Lex/HeaderSearch.h"
 #include "clang/Lex/PreprocessorOptions.h"
 #include "clang/Sema/Sema.h"
 #include "clang/Sema/SemaDiagnostic.h"
 
 #include "../Utils.h"
 
+#include "IncludeUtils.inc"
+
 using namespace clang;
 
 #if LLVM_VERSION_MAJOR >= 18
@@ -134,6 +138,39 @@ class EnzymePlugin final : public clang::ASTConsumer {
  Builder.defineMacro("ENZYME_VERSION_PATCH",
  std::to_string(ENZYME_VERSION_PATCH));
  CI.getPreprocessor().setPredefines(Predefines.str());
+
+ auto baseFS = &CI.getFileManager().getVirtualFileSystem();
+ llvm::vfs::OverlayFileSystem *fuseFS(
+ new llvm::vfs::OverlayFileSystem(baseFS));
+ IntrusiveRefCntPtr<llvm::vfs::InMemoryFileSystem> fs(
+ new llvm::vfs::InMemoryFileSystem());
+
+ struct tm y2k = {};
+
+ y2k.tm_hour = 0;
+ y2k.tm_min = 0;
+ y2k.tm_sec = 0;
+ y2k.tm_year = 100;
+ y2k.tm_mon = 0;
+ y2k.tm_mday = 1;
+ time_t timer = mktime(&y2k);
+ for (const auto &pair : include_headers) {
+ fs->addFile(StringRef(pair[0]), timer,
+ llvm::MemoryBuffer::getMemBuffer(
+ StringRef(pair[1]), StringRef(pair[0]),
+ /*RequiresNullTerminator*/ true));
+ }
+
+ fuseFS->pushOverlay(fs);
+ fuseFS->pushOverlay(baseFS);
+ CI.getFileManager().setVirtualFileSystem(fuseFS);
+
+ auto DE = CI.getFileManager().getDirectoryRef("/enzymeroot");
+ assert(DE);
+ auto DL = DirectoryLookup(*DE, SrcMgr::C_User,
+ /*isFramework=*/false);
+ CI.getPreprocessor().getHeaderSearchInfo().AddSearchPath(DL,
+ /*isAngled=*/true);
  }
  ~EnzymePlugin() {}
  void HandleTranslationUnit(ASTContext &context) override {}