Enzyme support example (#29)

* Preliminary support for automatic differentiation using Enzyme

* Build GridKit with Enzyme as dependency

* Add simple Enzyme tests

* Add Spack package for GridKit.

* gridkit +enzyme in CI.

* Update README with Enzyme dependency.

* Changed Enzyme add_custom_target name to fix build issues with Ninja.

---------

Co-authored-by: Asher Mancinelli <[email protected]>
Co-authored-by: Slaven Peles <[email protected]>

.github/workflows/spack_default_build.yaml

@@ -7,15 +7,15 @@ env:
   # Our repo name contains upper case characters, so we can't use ${{ github.repository }}
   IMAGE_NAME: gridkit
   USERNAME: gridkit-bot
-  BASE_VERSION: ubuntu-22.04-fortran-v0.1.0
+  BASE_VERSION: ubuntu-24.04-fortran-v0.2.2
 
 # Until we remove the need to clone submodules to build, this should on be in PRs
 on: [pull_request]
 
 jobs:
   base_image_build:
     name: Build base image
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     permissions:
       contents: read
       packages: write
@@ -53,7 +53,7 @@ jobs:
         if: ${{ env.IMAGE_EXISTS == 'false' }}
         run: |
           cat << EOF > Dockerfile
-          FROM ubuntu:22.04
+          FROM ubuntu:24.04
           RUN apt-get update && \
               apt-get install -y --no-install-recommends \
                 software-properties-common \
@@ -87,16 +87,17 @@ jobs:
 
   gridkit_spack_builds:
     needs: base_image_build
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-24.04
     permissions:
       packages: write
       contents: read
+    timeout-minutes: 30
 
     strategy:
       matrix:
         # Minimal Build(s) - GHCR mirror speeds these up a lot!
         spack_spec:
-          - gridkit@develop
+          - gridkit@develop +enzyme
 
     name: Build gridkit with Spack
     steps:
@@ -125,6 +126,11 @@ jobs:
             mirrors:
               local-buildcache: oci://${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
               spack: https://binaries.spack.io/develop
+            packages:
+              llvm:
+                externals:
+                - spec: llvm@18
+                  prefix: /usr/lib/llvm-18
           EOF
 
       - name: Configure GHCR mirror
@@ -134,7 +140,7 @@ jobs:
         run: spack -e . buildcache keys --install --trust
 
       - name: Find external packages
-        run: spack -e . external find --all
+        run: spack -e . external find --all --exclude llvm
 
       - name: Spack develop gridkit
         run: spack -e . develop --path=$(pwd) gridkit@develop

.gitmodules

@@ -1,4 +1,4 @@
 [submodule "Buildsystem/spack/spack"]
 	path = Buildsystem/spack/spack
 	url = https://github.com/nkoukpaizan/spack.git
-	branch = Gridkit-package-dev
+	branch = Gridkit-package-dev+enzyme

CMake/FindEnzyme.cmake

@@ -0,0 +1,194 @@
+# 
+# Copyright (c) 2017, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by Slaven Peles <[email protected]>.
+# LLNL-CODE-718378.
+# All rights reserved.
+# 
+# This file is part of GridKit™. For details, see github.com/LLNL/GridKit 
+# Please also read the LICENSE file. 
+# 
+# Redistribution and use in source and binary forms, with or without 
+# modification, are permitted provided that the following conditions are met:
+# - Redistributions of source code must retain the above copyright notice, 
+#   this list of conditions and the disclaimer below.
+# - Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the disclaimer (as noted below) in the 
+#   documentation and/or other materials provided with the distribution.
+# - Neither the name of the LLNS/LLNL nor the names of its contributors may 
+#   be used to endorse or promote products derived from this software without 
+#   specific prior written permission.
+# 
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
+# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
+# INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
+# MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL LAWRENCE LIVERMORE NATIONAL
+# SECURITY, LLC, THE U.S. DEPARTMENT OF ENERGY OR CONTRIBUTORS BE
+# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY,
+# OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
+# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISINGIN ANY 
+# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
+# THE POSSIBILITY OF SUCH DAMAGE.
+# 
+# Lawrence Livermore National Laboratory is operated by Lawrence Livermore 
+# National Security, LLC, for the U.S. Department of Energy, National 
+# Nuclear Security Administration under Contract DE-AC52-07NA27344.
+# 
+# This document was prepared as an account of work sponsored by an agency 
+# of the United States government. Neither the United States government nor 
+# Lawrence Livermore National Security, LLC, nor any of their employees 
+# makes any warranty, expressed or implied, or assumes any legal liability 
+# or responsibility for the accuracy, completeness, or usefulness of any 
+# information, apparatus, product, or process disclosed, or represents that 
+# its use would not infringe privately owned rights. Reference herein to 
+# any specific commercial product, process, or service by trade name, 
+# trademark, manufacturer, or otherwise does not necessarily constitute or 
+# imply its endorsement, recommendation, or favoring by the United States 
+# government or Lawrence Livermore National Security, LLC. The views and 
+# opinions of authors expressed herein do not necessarily state or reflect 
+# those of the United States government or Lawrence Livermore National 
+# Security, LLC, and shall not be used for advertising or product 
+# endorsement purposes. 
+# 
+
+#[[
+
+Finds Enzyme Clang plugin
+
+User may set:
+- ENZYME_DIR
+
+Author(s):
+- Asher Mancinelli <[email protected]>
+
+]]
+
+if(NOT ${CMAKE_C_COMPILER_ID} STREQUAL "Clang" OR NOT ${CMAKE_CXX_COMPILER_ID} STREQUAL "Clang")
+  message(FATAL_ERROR "Enzyme may only be enabled if building with Clang")
+endif()
+
+find_library(ENZYME_LLVM_PLUGIN_LIBRARY
+  NAMES
+  LLVMEnzyme-19.so
+  LLVMEnzyme-18.so
+  LLVMEnzyme-17.so
+  LLVMEnzyme-16.so
+  LLVMEnzyme-15.so
+  LLVMEnzyme-14.so
+  LLVMEnzyme-13.so
+  LLVMEnzyme-12.so
+  LLVMEnzyme-11.so
+  LLVMEnzyme-10.so
+  LLVMEnzyme-9.so
+  PATHS
+  ${ENZYME_DIR} $ENV{ENZYME_DIR}
+  ENV LD_LIBRARY_PATH ENV DYLD_LIBRARY_PATH
+  PATH_SUFFIXES
+  lib64 lib)
+
+find_library(ENZYME_CLANG_PLUGIN_LIBRARY
+  NAMES
+  ClangEnzyme-19.so
+  ClangEnzyme-18.so
+  ClangEnzyme-17.so
+  ClangEnzyme-16.so
+  ClangEnzyme-15.so
+  ClangEnzyme-14.so
+  ClangEnzyme-13.so
+  ClangEnzyme-12.so
+  ClangEnzyme-11.so
+  ClangEnzyme-10.so
+  ClangEnzyme-9.so
+  PATHS
+  ${ENZYME_DIR} $ENV{ENZYME_DIR}
+  ENV LD_LIBRARY_PATH ENV DYLD_LIBRARY_PATH
+  PATH_SUFFIXES
+  lib64 lib)
+
+find_program(GRIDKIT_LLVM_LINK llvm-link)
+find_program(GRIDKIT_OPT opt)
+
+if("${GRIDKIT_LLVM_LINK}" STREQUAL "GRIDKIT_LLVM_LINK-NOTFOUND")
+  message(FATAL_ERROR "Could not find llvm-link! Will not be able to build Enzyme targets.")
+endif()
+
+if("${GRIDKIT_OPT}" STREQUAL "GRIDKIT_OPT-NOTFOUND")
+  message(FATAL_ERROR "Could not find opt! Will not be able to build Enzyme targets.")
+endif()
+
+message(STATUS "${ENZYME_CLANG_PLUGIN_LIBRARY};${ENZYME_LLVM_PLUGIN_LIBRARY}")
+if((NOT ${ENZYME_LLVM_PLUGIN_LIBRARY}) OR (NOT ${ENZYME_CLANG_PLUGIN_LIBRARY}))
+  set(ENZYME_CLANG_PLUGIN_LIBRARY CACHE FILEPATH "Path to Enzyme Clang plugin library")
+  set(ENZYME_LLVM_PLUGIN_LIBRARY CACHE FILEPATH "Path to Enzyme LLVM plugin library")
+  message(STATUS "Found Enzyme clang plugin: ${ENZYME_CLANG_PLUGIN_LIBRARY}")
+  message(STATUS "Found Enzyme LLVM plugin: ${ENZYME_LLVM_PLUGIN_LIBRARY}")
+  get_filename_component(ENZYME_LIBRARY_DIR ${ENZYME_CLANG_PLUGIN_LIBRARY} DIRECTORY CACHE "Enzyme library directory")
+else()
+  message(FATAL_ERROR "Enzyme library not found!"
+    " Set ENZYME_DIR to Enzyme's installation prefix.")
+endif()
+
+macro(enzyme_add_executable)
+  set(options)
+  set(oneValueArgs NAME)
+  set(multiValueArgs SOURCES LINK_LIBRARIES)
+  cmake_parse_arguments(enzyme_add_executable "${options}" "${oneValueArgs}"
+    "${multiValueArgs}" ${ARGN})
+
+  set(PHASE2 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_add_executable_NAME}.bc")
+  set(PHASE3 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_add_executable_NAME}_enzyme.ll")
+  set(PHASE4 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_add_executable_NAME}_opt.ll")
+  set(PHASE5 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_add_executable_NAME}")
+
+  set(OBJS "")
+
+  foreach(SRC ${enzyme_add_executable_SOURCES})
+    set(PHASE0 "${CMAKE_CURRENT_SOURCE_DIR}/${SRC}")
+    set(PHASE1 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_add_executable_NAME}_${SRC}_compile.o")
+    add_custom_command(
+      DEPENDS ${PHASE0}
+      OUTPUT ${PHASE1}
+      COMMAND ${CMAKE_CXX_COMPILER} -flto -c ${PHASE0} -O2 -fno-vectorize -ffast-math -fno-unroll-loops -o ${PHASE1}
+      COMMENT "Compiling ${SRC} to object file for target ${enzyme_add_executable_NAME}"
+      )
+    set(OBJS "${OBJS} ${PHASE1}")
+  endforeach()
+
+  cmake_language(EVAL CODE "
+  add_custom_command(
+    DEPENDS ${OBJS}
+    OUTPUT ${PHASE2}
+    COMMAND ${GRIDKIT_LLVM_LINK} ${OBJS} -o ${PHASE2}
+    COMMENT \"Linking object files to LLVM bytecode for target ${enzyme_add_executable_NAME}\"
+    )
+  ")
+
+  add_custom_command(
+    DEPENDS ${PHASE2}
+    OUTPUT ${PHASE3}
+    COMMAND ${GRIDKIT_OPT} ${PHASE2} -load-pass-plugin=${ENZYME_LLVM_PLUGIN_LIBRARY} -passes=enzyme -o ${PHASE3} -S
+    COMMENT "Running Enzyme opt pass on target ${enzyme_add_executable_NAME}"
+    )
+
+  add_custom_command(
+    DEPENDS ${PHASE3}
+    OUTPUT ${PHASE4}
+    COMMAND ${GRIDKIT_OPT} ${PHASE3} -O2 -o ${PHASE4} -S
+    COMMENT "Running remaining opt passes on target ${enzyme_add_executable_NAME}"
+    )
+
+  add_custom_command(
+    DEPENDS ${PHASE4}
+    OUTPUT ${PHASE5}
+    COMMAND ${CMAKE_CXX_COMPILER} ${PHASE4} -o ${PHASE5}
+    )
+
+  add_custom_target(
+    "${enzyme_add_executable_NAME}_target" ALL
+    DEPENDS ${PHASE5}
+    )
+endmacro()

CMakeLists.txt

@@ -76,14 +76,16 @@ set(PACKAGE_VERSION_PATCH "7")
 set(PACKAGE_VERSION "${PACKAGE_VERSION_MAJOR}.${PACKAGE_VERSION_MINOR}.${PACKAGE_VERSION_PATCH}")
 
 # Ipopt support is disabled by default
-option(GRIDKIT_ENABLE_IPOPT "Enable Ipopt support" ON)
+option(GRIDKIT_ENABLE_IPOPT "Enable Ipopt support" OFF)
 
 # SUNDIALS support is enabled by default
 option(GRIDKIT_ENABLE_SUNDIALS "Enable SUNDIALS support" ON)
 
 # Enable KLU
 option(GRIDKIT_ENABLE_SUNDIALS_SPARSE "Enable SUNDIALS sparse linear solvers" ON)
 
+# Enzyme support is disabled by default
+option(GRIDKIT_ENABLE_ENZYME "Enable automatic differentiation with Enzyme" OFF)
 
 set(CMAKE_MACOSX_RPATH 1)
 
@@ -133,6 +135,10 @@ if(GRIDKIT_ENABLE_SUNDIALS_SPARSE)
   include(FindSuiteSparse)
 endif()
 
+if(${GRIDKIT_ENABLE_ENZYME})
+  include(FindEnzyme)
+endif()
+
 # Macro that adds libraries
 include(GridkitAddLibrary)
 

Examples/CMakeLists.txt

@@ -75,3 +75,7 @@ if(TARGET SUNDIALS::idas)
     add_subdirectory(ParameterEstimation)
   endif()
 endif()
+
+if(GRIDKIT_ENABLE_ENZYME)
+  add_subdirectory(Enzyme)
+endif()

Examples/Enzyme/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_subdirectory(Standalone)
+add_subdirectory(Library)

Examples/Enzyme/Library/CMakeLists.txt

@@ -0,0 +1,4 @@
+enzyme_add_executable(NAME EnzymeLibCheck SOURCES main.cpp library.cpp)
+#install(TARGETS ${CMAKE_CURRENT_BINARY_DIR}/EnzymeLibCheck DESTINATION bin)
+
+add_test(NAME "EnzymeLibCheck" COMMAND ${CMAKE_CURRENT_BINARY_DIR}/EnzymeLibCheck)

Examples/Enzyme/Library/library.cpp

@@ -0,0 +1,3 @@
+double square(double x) {
+  return x * x;
+}

Examples/Enzyme/Library/library.hpp

@@ -0,0 +1,2 @@
+#pragma once
+double square(double x);

Examples/Enzyme/Library/main.cpp

@@ -0,0 +1,20 @@
+#include <iostream>
+#include <limits>
+#include "library.hpp"
+
+double __enzyme_autodiff(double (*)(double), ...);
+
+int main() {
+  int fail = 0;
+  double sq = square(5.0);
+  double dsq = __enzyme_autodiff(square, 5.0);
+
+  std::cout << "x = 5, x^2 = " << sq << ", d(x^2)/dx = " << dsq << "\n"; 
+  if (std::abs(dsq - 10.0) > std::numeric_limits<double>::epsilon())
+  {
+    fail++;
+    std::cout << "Result incorrect\n";
+  }
+  std::cout << "Status: " << fail << "\n";
+  return fail;
+}

Examples/Enzyme/Standalone/CMakeLists.txt

@@ -0,0 +1,6 @@
+enzyme_add_executable(
+  NAME EnzymeStandaloneCheck
+  SOURCES main.cpp
+  )
+
+add_test(NAME "EnzymeStandaloneCheck" COMMAND ${CMAKE_CURRENT_BINARY_DIR}/EnzymeStandaloneCheck)

Examples/Enzyme/Standalone/main.cpp

@@ -0,0 +1,27 @@
+#include <iostream>
+#include <limits>
+
+double square(double x) {
+  return x * x;
+}
+
+double __enzyme_autodiff(double(*)(double), ...);
+double dsquare(double x) {
+  return __enzyme_autodiff(square, x);
+}
+
+double sq  = square(5.0);
+double dsq = dsquare(5.0);
+
+int main()
+{
+  int fail = 0;
+  std::cout << "x = 5, x^2 = " << sq << ", d(x^2)/dx = " << dsq << "\n"; 
+  if (std::abs(dsq - 10.0) > std::numeric_limits<double>::epsilon())
+  {
+    fail++;
+    std::cout << "Result incorrect\n";
+  }
+  std::cout << "Status: " << fail << "\n";
+  return fail;
+}

README.md

@@ -17,6 +17,8 @@ You should have all of the following installed before installing GridKit™
 	- [SUNDIALS](https://github.com/LLNL/sundials) >= 7.0.0
 	- [Suitesparse](https://github.com/DrTimothyAldenDavis/SuiteSparse) >= 5.x (optional)
 	- [Ipopt](https://github.com/coin-or/Ipopt) >= 3.x (optional)
+    - [Enzyme](https://github.com/EnzymeAD/Enzyme) >=0.0.131 (optional)
+        - [LLVM](https://github.com/llvm/llvm-project) >= 18.x 
 - [CMake](https://cmake.org/) >= 3.12
 - C++ 17 compliant compiler