WIP herbie

enzyme/CMakeLists.txt

@@ -265,6 +265,7 @@ find_library(MPFR_LIB_PATH mpfr)
 CHECK_INCLUDE_FILE("mpfr.h" HAS_MPFR_H)
 message("MPFR lib: " ${MPFR_LIB_PATH})
 message("MPFR header: " ${HAS_MPFR_H})
+link_libraries(mpfr)
 
 file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/include/SCEV/ScalarEvolutionExpander.h" "${INPUT_TEXT}")
 

enzyme/Enzyme/ActivityAnalysis.cpp

@@ -204,6 +204,7 @@ const StringSet<> KnownInactiveFunctions = {
     "memcmp",
     "memchr",
     "gettimeofday",
+    "getrusage",
     "stat",
     "mkdir",
     "compress2",
@@ -1024,7 +1025,8 @@ bool isValuePotentiallyUsedAsPointer(llvm::Value *val) {
     for (auto u : cur->users()) {
       if (isa<ReturnInst>(u))
         return true;
-      if (!cast<Instruction>(u)->mayReadOrWriteMemory()) {
+      if (isa<ConstantExpr>(u) ||
+          !cast<Instruction>(u)->mayReadOrWriteMemory()) {
         todo.push_back(u);
         continue;
       }

enzyme/Enzyme/CMakeLists.txt

@@ -44,12 +44,33 @@ set(LLVM_LINK_COMPONENTS Demangle)
 file(GLOB ENZYME_SRC CONFIGURE_DEPENDS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR}
     "*.cpp"
 )
-list(REMOVE_ITEM ENZYME_SRC "eopt.cpp")
+list(REMOVE_ITEM ENZYME_SRC "eopt.cpp" "Herbie.cpp")
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 
 list(APPEND ENZYME_SRC TypeAnalysis/TypeTree.cpp TypeAnalysis/TypeAnalysis.cpp TypeAnalysis/TypeAnalysisPrinter.cpp TypeAnalysis/RustDebugInfo.cpp)
 
+# set(ENZYME_LINK_TARGETS)
+set(ENZYME_ENABLE_FPOPT 0 CACHE BOOL "Enable FPOpt")
+
+if(ENZYME_ENABLE_FPOPT)
+  include(ExternalProject)
+    ExternalProject_Add(herbie
+        GIT_REPOSITORY https://github.com/sbrantq/herbie
+        GIT_TAG 9a81e9ecd01515ed88fb4efb77e4c0ae7be9cba5
+        UPDATE_COMMAND ""
+        CONFIGURE_COMMAND ""
+        BUILD_COMMAND make egg-herbie && make update && raco exe -o herbie --orig-exe --embed-dlls --vv src/main.rkt
+        BUILD_IN_SOURCE true
+        INSTALL_COMMAND COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_CURRENT_BINARY_DIR}/herbie/install
+            COMMAND ${CMAKE_COMMAND} -E copy_if_different ${CMAKE_CURRENT_BINARY_DIR}/herbie-prefix/src/herbie/herbie ${CMAKE_CURRENT_BINARY_DIR}/herbie/install/herbie
+    )
+    list(APPEND ENZYME_SRC Herbie.cpp)
+    add_compile_definitions(ENZYME_ENABLE_FPOPT=1)
+    set_source_files_properties(Herbie.cpp PROPERTIES COMPILE_DEFINITIONS HERBIE_BINARY="${CMAKE_CURRENT_BINARY_DIR}/herbie/install/herbie")
+endif()
+
+
 # on windows `PLUGIN_TOOL` doesn't link against LLVM.dll
 if ((WIN32 OR CYGWIN) AND LLVM_LINK_LLVM_DYLIB)
     add_llvm_library( LLVMEnzyme-${LLVM_VERSION_MAJOR}
@@ -71,6 +92,7 @@ if (${Clang_FOUND})
         intrinsics_gen
 	LINK_COMPONENTS
 	LLVM
+    ${ENZYME_LINK_TARGETS}
     )
 target_compile_definitions(ClangEnzyme-${LLVM_VERSION_MAJOR} PUBLIC ENZYME_RUNPASS)
 endif()
@@ -185,17 +207,17 @@ target_link_options(LLDEnzymePrintFlags INTERFACE "SHELL: -Wl,-mllvm -Wl,-enzyme
 add_library(LLDEnzymeNoStrictAliasingFlags INTERFACE)
 target_link_options(LLDEnzymeNoStrictAliasingFlags INTERFACE "SHELL: -Wl,-mllvm -Wl,-enzyme-strict-aliasing=0")
 
-# this custom target exists to prevent CMake from incorrectly assuming that 
-# targets that link depend on LLDEnzyme-XX can be built at the same time or 
+# this custom target exists to prevent CMake from incorrectly assuming that
+# targets that link depend on LLDEnzyme-XX can be built at the same time or
 # before LLDEnzyme-XX has finished.
 add_custom_target(LLDEnzymeDummy "" DEPENDS LLDEnzyme-${LLVM_VERSION_MAJOR})
 add_dependencies(LLDEnzymeFlags LLDEnzymeDummy)
 
 add_library(ClangEnzymeFlags INTERFACE)
 target_compile_options(ClangEnzymeFlags INTERFACE "-fplugin=$<TARGET_FILE:ClangEnzyme-${LLVM_VERSION_MAJOR}>")
 
-# this custom target exists to prevent CMake from incorrectly assuming that 
-# targets that link depend on ClangEnzyme-XX can be built at the same time or 
+# this custom target exists to prevent CMake from incorrectly assuming that
+# targets that link depend on ClangEnzyme-XX can be built at the same time or
 # before ClangEnzyme-XX has finished.
 add_custom_target(ClangEnzymeDummy "" DEPENDS ClangEnzyme-${LLVM_VERSION_MAJOR})
 add_dependencies(ClangEnzymeFlags ClangEnzymeDummy)

enzyme/Enzyme/DiffeGradientUtils.cpp

@@ -107,11 +107,13 @@ DiffeGradientUtils *DiffeGradientUtils::CreateFromClone(
   std::string prefix;
 
   switch (mode) {
-  case DerivativeMode::ForwardModeError:
   case DerivativeMode::ForwardMode:
   case DerivativeMode::ForwardModeSplit:
     prefix = "fwddiffe";
     break;
+  case DerivativeMode::ForwardModeError:
+    prefix = "fwderr";
+    break;
   case DerivativeMode::ReverseModeCombined:
   case DerivativeMode::ReverseModeGradient:
     prefix = "diffe";

enzyme/Enzyme/DifferentialUseAnalysis.h

@@ -123,14 +123,17 @@ inline bool is_value_needed_in_reverse(
         }
       }
     }
-    if (gutils->mode == DerivativeMode::ForwardModeError &&
+#ifdef ENZYME_ENABLE_FPOPT
+    if ((hasFPOptLogger(gutils->oldFunc->getParent()) ||
+         gutils->mode == DerivativeMode::ForwardModeError) &&
         !gutils->isConstantValue(const_cast<Value *>(inst))) {
       if (EnzymePrintDiffUse)
         llvm::errs()
             << " Need: " << to_string(VT) << " of " << *inst
             << " in reverse as forward mode error always needs result\n";
       return seen[idx] = true;
     }
+#endif
   }
 
   if (auto CI = dyn_cast<CallInst>(inst)) {

enzyme/Enzyme/Enzyme.cpp

@@ -3272,6 +3272,9 @@ AnalysisKey EnzymeNewPM::Key;
 #include "ActivityAnalysisPrinter.h"
 #include "JLInstSimplify.h"
 #include "PreserveNVVM.h"
+#ifdef ENZYME_ENABLE_FPOPT
+#include "Herbie.h"
+#endif
 #include "TypeAnalysis/TypeAnalysisPrinter.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Transforms/AggressiveInstCombine/AggressiveInstCombine.h"
@@ -3385,6 +3388,10 @@ void augmentPassBuilder(llvm::PassBuilder &PB) {
     OptimizerPM.addPass(llvm::SROAPass());
 #endif
     MPM.addPass(createModuleToFunctionPassAdaptor(std::move(OptimizerPM)));
+#ifdef ENZYME_ENABLE_FPOPT
+    if (EnzymeEnableFPOpt)
+      MPM.addPass(FPOptNewPM());
+#endif
     MPM.addPass(EnzymeNewPM(/*PostOpt=*/true));
     MPM.addPass(PreserveNVVMNewPM(/*Begin*/ false));
 #if LLVM_VERSION_MAJOR >= 16
@@ -3669,6 +3676,12 @@ extern "C" void registerEnzyme(llvm::PassBuilder &PB) {
           MPM.addPass(EnzymeNewPM());
           return true;
         }
+#ifdef ENZYME_ENABLE_FPOPT
+        if (Name == "fp-opt") {
+          MPM.addPass(FPOptNewPM());
+          return true;
+        }
+#endif
         if (Name == "preserve-nvvm") {
           MPM.addPass(PreserveNVVMNewPM(/*Begin*/ true));
           return true;

enzyme/Enzyme/EnzymeLogic.cpp

@@ -6395,6 +6395,7 @@ llvm::Function *EnzymeLogic::CreateNoFree(RequestContext context, Function *F) {
                         "vprintf",
                         "fprintf",
                         "fputc",
+                        "getrusage",
                          "memchr",
                          "time",
                          "strlen",

enzyme/Enzyme/FunctionUtils.cpp

@@ -114,6 +114,10 @@
 
 #include "CacheUtility.h"
 
+#ifdef ENZYME_ENABLE_FPOPT
+#include "Herbie.h"
+#endif
+
 #define addAttribute addAttributeAtIndex
 #define removeAttribute removeAttributeAtIndex
 #define getAttribute getAttributeAtIndex
@@ -1418,6 +1422,27 @@ Function *PreProcessCache::preprocessForClone(Function *F,
         /*ModuleLevelChanges*/ CloneFunctionChangeType::LocalChangesOnly,
         Returns, "", nullptr);
   }
+#ifdef ENZYME_ENABLE_FPOPT
+  if (hasFPOptLogger(F->getParent())) {
+    for (const auto &pair : VMap) {
+      if (auto *before = dyn_cast<Instruction>(pair.first)) {
+        if (!before->getType()->isFloatingPointTy()) {
+          continue;
+        }
+        auto *after = cast<Instruction>(pair.second);
+        attachFPOptMetadata(after, before);
+      } else if (auto *beforeBB = dyn_cast<BasicBlock>(pair.first)) {
+        auto *afterBB = cast<BasicBlock>(pair.second);
+        for (const auto &[before, after] : zip(*beforeBB, *afterBB)) {
+          if (!before.getType()->isFloatingPointTy()) {
+            continue;
+          }
+          attachFPOptMetadata(&after, &before);
+        }
+      }
+    }
+  }
+#endif
   CloneOrigin[NewF] = F;
   NewF->setAttributes(F->getAttributes());
   if (EnzymeNoAlias)
@@ -1431,7 +1456,13 @@ Function *PreProcessCache::preprocessForClone(Function *F,
   setFullWillReturn(NewF);
 
   if (EnzymePreopt) {
+#ifdef ENZYME_ENABLE_FPOPT
+    // Disable recursive inlining since no FPOpt metadata is attached
+    // to inlined instructions
+    if (!hasFPOptLogger(F->getParent()) && EnzymeInline) {
+#else
     if (EnzymeInline) {
+#endif
       ForceRecursiveInlining(NewF, /*Limit*/ EnzymeInlineCount);
       setFullWillReturn(NewF);
       PreservedAnalyses PA;
@@ -1800,7 +1831,8 @@ Function *PreProcessCache::preprocessForClone(Function *F,
       FAM.invalidate(*NewF, PA);
     }
 
-    if (mode != DerivativeMode::ForwardMode)
+    if (mode != DerivativeMode::ForwardMode &&
+        mode != DerivativeMode::ForwardModeError)
       ReplaceReallocs(NewF);
 
     {
@@ -1836,7 +1868,8 @@ Function *PreProcessCache::preprocessForClone(Function *F,
     PA.preserve<PhiValuesAnalysis>();
   }
 
-  if (mode != DerivativeMode::ForwardMode)
+  if (mode != DerivativeMode::ForwardMode &&
+      mode != DerivativeMode::ForwardModeError)
     ReplaceReallocs(NewF);
 
   if (mode == DerivativeMode::ReverseModePrimal ||