Poseidon (cleaned)

.github/workflows/poseidon.yml

@@ -0,0 +1,56 @@
+name: Poseidon CI
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  merge_group:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
+jobs:
+  build-linux:
+    name: Poseidon CI LLVM ${{ matrix.llvm }} ${{ matrix.build }} ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      fail-fast: false
+      matrix:
+        llvm: ["20"]
+        build: ["Release", "Debug"]
+        os: [ubuntu-22.04]
+
+    timeout-minutes: 60
+
+    steps:
+    - name: Install dependencies
+      run: |
+          wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key|sudo apt-key add -
+          sudo apt-add-repository "deb http://apt.llvm.org/`lsb_release -c | cut -f2`/ llvm-toolchain-`lsb_release -c | cut -f2`-${{ matrix.llvm }} main" || true
+          sudo apt-get update
+          sudo apt-get install -y cmake gcc g++ llvm-${{ matrix.llvm }}-dev clang-${{ matrix.llvm }} libclang-${{ matrix.llvm }}-dev libzstd-dev libmpfr-dev
+          sudo python3 -m pip install --upgrade pip lit
+
+    - name: Install Racket
+      uses: Bogdanp/setup-racket@v1.12
+      with:
+        version: '8.15'
+
+    - uses: actions/checkout@v4
+
+    - name: Build Enzyme with Poseidon
+      run: |
+          mkdir build && cd build
+          cmake ../enzyme -DCMAKE_BUILD_TYPE=${{ matrix.build }} -DENABLE_POSEIDON=1 -DLLVM_EXTERNAL_LIT=$(which lit) -DLLVM_DIR=/usr/lib/llvm-${{ matrix.llvm }}/lib/cmake/llvm
+          make -j $(nproc)
+
+    - name: check-poseidon
+      working-directory: build
+      run: make -j $(nproc) check-poseidon
+
+    - name: check-poseidon-integration
+      working-directory: build
+      run: make -j3 check-poseidon-integration

enzyme/Enzyme/AdjointGenerator.h

@@ -48,6 +48,10 @@
 #include "TraceUtils.h"
 #include "TypeAnalysis/TBAA.h"
 
+#ifdef ENABLE_POSEIDON
+#include "Poseidon/Poseidon.h"
+#endif
+
 #define DEBUG_TYPE "enzyme"
 
 // Helper instruction visitor that generates adjoints
@@ -4610,7 +4614,8 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
                 .forceAnonymousTape = false,
                 .typeInfo = nextTypeInfo,
                 .runtimeActivity = gutils->runtimeActivity,
-                .strongZero = gutils->strongZero},
+                .strongZero = gutils->strongZero,
+                .profiled = gutils->profiled},
             TR.analyzer->interprocedural, subdata,
             /*omp*/ true);
 
@@ -5989,7 +5994,8 @@ class AdjointGenerator : public llvm::InstVisitor<AdjointGenerator> {
               .forceAnonymousTape = false,
               .typeInfo = nextTypeInfo,
               .runtimeActivity = gutils->runtimeActivity,
-              .strongZero = gutils->strongZero},
+              .strongZero = gutils->strongZero,
+              .profiled = gutils->profiled},
           TR.analyzer->interprocedural, subdata);
       if (!newcalled)
         return;

enzyme/Enzyme/CApi.cpp

@@ -717,7 +717,8 @@ LLVMValueRef EnzymeCreatePrimalAndGradient(
           .forceAnonymousTape = (bool)forceAnonymousTape,
           .typeInfo = eunwrap(typeInfo, cast<Function>(unwrap(todiff))),
           .runtimeActivity = (bool)runtimeActivity,
-          .strongZero = (bool)strongZero},
+          .strongZero = (bool)strongZero,
+          .profiled = false},
       eunwrap(TA), eunwrap(augmented)));
 }
 EnzymeAugmentedReturnPtr EnzymeCreateAugmentedPrimal(

enzyme/Enzyme/CMakeLists.txt

@@ -50,6 +50,54 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON)
 
 list(APPEND ENZYME_SRC TypeAnalysis/TypeTree.cpp TypeAnalysis/TypeAnalysis.cpp TypeAnalysis/TypeAnalysisPrinter.cpp TypeAnalysis/RustDebugInfo.cpp)
 
+set(ENABLE_POSEIDON 0 CACHE BOOL "Enable Poseidon")
+
+if(ENABLE_POSEIDON)
+  include(ExternalProject)
+    ExternalProject_Add(herbie
+        GIT_REPOSITORY https://github.com/sbrantq/herbie
+        GIT_TAG 193e2a4e6e4902fe5ee7fa8bd5747c19f70ce19b
+        UPDATE_COMMAND ""
+        CONFIGURE_COMMAND ""
+        BUILD_COMMAND make clean 
+            COMMAND raco pkg install --auto --update-deps fpbench || true 
+            COMMAND raco pkg install --auto --update-deps rival || true 
+            COMMAND cargo build --release --manifest-path=egg-herbie/Cargo.toml 
+            COMMAND raco pkg install ./egg-herbie 
+            COMMAND mkdir -p herbie-compiled/ 
+            COMMAND raco exe -o herbie --orig-exe --embed-dlls --vv src/main.rkt 
+            COMMAND raco distribute herbie-compiled herbie
+        BUILD_IN_SOURCE true
+        INSTALL_COMMAND COMMAND ${CMAKE_COMMAND} -E make_directory ${CMAKE_CURRENT_BINARY_DIR}/herbie/install
+            COMMAND ${CMAKE_COMMAND} -E copy_directory ${CMAKE_CURRENT_BINARY_DIR}/herbie-prefix/src/herbie/herbie-compiled ${CMAKE_CURRENT_BINARY_DIR}/herbie/install/herbie
+    )
+    list(APPEND ENZYME_SRC 
+        Poseidon/Poseidon.cpp
+        Poseidon/PoseidonEvaluators.cpp
+        Poseidon/PoseidonHerbieUtils.cpp
+        Poseidon/PoseidonProfUtils.cpp
+        Poseidon/PoseidonPrecUtils.cpp
+        Poseidon/PoseidonSolvers.cpp
+        Poseidon/PoseidonTypes.cpp
+        Poseidon/PoseidonUtils.cpp
+    )
+    add_compile_definitions(ENABLE_POSEIDON=1)
+    set_source_files_properties(Poseidon/PoseidonHerbieUtils.cpp PROPERTIES COMPILE_DEFINITIONS HERBIE_BINARY="${CMAKE_CURRENT_BINARY_DIR}/herbie/install/herbie/bin/herbie")
+
+    add_library(EnzymeFPProfile STATIC
+        Runtimes/FPProfiler/FPProfiler.cpp
+    )
+    target_compile_features(EnzymeFPProfile PRIVATE cxx_std_17)
+    target_include_directories(EnzymeFPProfile PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/Runtimes/FPProfiler)
+
+    install(TARGETS EnzymeFPProfile
+        EXPORT EnzymeTargets
+        ARCHIVE DESTINATION lib
+        COMPONENT dev
+    )
+endif()
+
+
 if (ENZYME_ENABLE_PLUGINS)
         # on windows `PLUGIN_TOOL` doesn't link against LLVM.dll
         if ((WIN32 OR CYGWIN) AND LLVM_LINK_LLVM_DYLIB)
@@ -119,6 +167,14 @@ if (ENZYME_ENABLE_PLUGINS)
                  )
              target_compile_definitions(LLDEnzyme-${LLVM_VERSION_MAJOR} PUBLIC ENZYME_RUNPASS)
         endif()
+
+        if(ENABLE_POSEIDON)
+            target_link_libraries(LLVMEnzyme-${LLVM_VERSION_MAJOR} PRIVATE mpfr)
+            if (${Clang_FOUND})
+                target_link_libraries(ClangEnzyme-${LLVM_VERSION_MAJOR} PRIVATE mpfr)
+            endif()
+            target_link_libraries(LLDEnzyme-${LLVM_VERSION_MAJOR} PRIVATE mpfr)
+        endif()
 endif()
 
 if (${ENZYME_STATIC_LIB})
@@ -129,6 +185,9 @@ if (${ENZYME_STATIC_LIB})
         DEPENDS
         intrinsics_gen
     )
+    if(ENABLE_POSEIDON)
+        target_link_libraries(EnzymeStatic-${LLVM_VERSION_MAJOR} PRIVATE mpfr)
+    endif()
 endif()
 
 if (${ENZYME_EXTERNAL_SHARED_LIB})
@@ -154,6 +213,10 @@ if (${ENZYME_EXTERNAL_SHARED_LIB})
         target_link_libraries(Enzyme-${LLVM_VERSION_MAJOR} LLVM)
     endif()
 
+    if(ENABLE_POSEIDON)
+        target_link_libraries(Enzyme-${LLVM_VERSION_MAJOR} PRIVATE mpfr)
+    endif()
+
     install(TARGETS Enzyme-${LLVM_VERSION_MAJOR}
         EXPORT EnzymeTargets
         LIBRARY DESTINATION lib COMPONENT shlib

enzyme/Enzyme/DiffeGradientUtils.cpp

@@ -89,7 +89,8 @@ DiffeGradientUtils *DiffeGradientUtils::CreateFromClone(
     bool strongZero, unsigned width, Function *todiff, TargetLibraryInfo &TLI,
     TypeAnalysis &TA, FnTypeInfo &oldTypeInfo, DIFFE_TYPE retType,
     bool shadowReturn, bool diffeReturnArg, ArrayRef<DIFFE_TYPE> constant_args,
-    bool returnTape, bool returnPrimal, Type *additionalArg, bool omp) {
+    bool returnTape, bool returnPrimal, Type *additionalArg, bool omp,
+    bool profiled) {
   Function *oldFunc = todiff;
   assert(mode == DerivativeMode::ReverseModeGradient ||
          mode == DerivativeMode::ReverseModeCombined ||
@@ -108,14 +109,16 @@ 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";
+    prefix = profiled ? "instr" : "diffe";
     break;
   case DerivativeMode::ReverseModePrimal:
     llvm_unreachable("invalid DerivativeMode: ReverseModePrimal\n");
@@ -129,7 +132,7 @@ DiffeGradientUtils *DiffeGradientUtils::CreateFromClone(
       nonconstant_values, returnvals, returnTape, returnPrimal,
       (mode == DerivativeMode::ReverseModeGradient) ? false : shadowReturn,
       prefix + oldFunc->getName(), &originalToNew,
-      /*diffeReturnArg*/ diffeReturnArg, additionalArg);
+      /*diffeReturnArg*/ diffeReturnArg, additionalArg, profiled);
 
   // Convert overwritten args from the input function to the preprocessed
   // function
@@ -165,6 +168,7 @@ DiffeGradientUtils *DiffeGradientUtils::CreateFromClone(
       Logic, newFunc, oldFunc, TLI, TA, TR, invertedPointers, constant_values,
       nonconstant_values, retType, shadowReturn, constant_args, originalToNew,
       mode, runtimeActivity, strongZero, width, omp);
+  res->profiled = profiled;
 
   return res;
 }

enzyme/Enzyme/DiffeGradientUtils.h

@@ -84,7 +84,8 @@ class DiffeGradientUtils final : public GradientUtils {
                   FnTypeInfo &oldTypeInfo, DIFFE_TYPE retType,
                   bool shadowReturnArg, bool diffeReturnArg,
                   llvm::ArrayRef<DIFFE_TYPE> constant_args, bool returnTape,
-                  bool returnPrimal, llvm::Type *additionalArg, bool omp);
+                  bool returnPrimal, llvm::Type *additionalArg, bool omp,
+                  bool profiled = false);
 
   llvm::AllocaInst *getDifferential(llvm::Value *val);
 

enzyme/Enzyme/DifferentialUseAnalysis.h

@@ -123,12 +123,13 @@ inline bool is_value_needed_in_reverse(
         }
       }
     }
-    if (gutils->mode == DerivativeMode::ForwardModeError &&
+    if ((gutils->profiled ||
+         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";
+        llvm::errs() << " Need: " << to_string(VT) << " of " << *inst
+                     << " in reverse as profiling mode or forward mode "
+                        "error always needs result\n";
       return seen[idx] = true;
     }
   }