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>
EnzymeAD · Feb 27, 2024 · ca06ef3 · ca06ef3
1 parent 744ede0
commit ca06ef3
Show file tree

Hide file tree

Showing 10 changed files with 916 additions and 70 deletions.
diff --git a/enzyme/BUILD b/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",

diff --git a/enzyme/Enzyme/CMakeLists.txt b/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

diff --git a/enzyme/Enzyme/Clang/EnzymeClang.cpp b/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 {}