Additional Enzyme examples and use cases (e.g DistributedGenerator)

cmake/FindEnzyme.cmake

@@ -101,7 +101,7 @@ message(STATUS "opt: ${GRIDKIT_OPT}")
 macro(enzyme_add_executable)
   set(options)
   set(oneValueArgs NAME)
-  set(multiValueArgs SOURCES LINK_LIBRARIES)
+  set(multiValueArgs SOURCES LINK_LIBRARIES INCLUDE_DIRECTORIES)
   cmake_parse_arguments(enzyme_add_executable "${options}" "${oneValueArgs}"
     "${multiValueArgs}" ${ARGN})
 
@@ -111,14 +111,32 @@ macro(enzyme_add_executable)
   set(PHASE5 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_add_executable_NAME}")
 
   set(OBJS "")
+  set(includes "${enzyme_add_executable_INCLUDE_DIRECTORIES}")
+
+  foreach(lib ${enzyme_add_executable_LINK_LIBRARIES})
+    get_target_property(include ${lib} INCLUDE_DIRECTORIES)
+    set(includes "${includes}" ${include})
+
+    get_target_property(libsource ${lib} SOURCES)
+    string(FIND "${libsource}" "TARGET" found)
+    if(NOT(${found} EQUAL -1))
+      list(APPEND LINKER_FLAGS "-Wl,${libsource}")
+    endif()
+  endforeach()
+
+  foreach(dir ${includes})
+    if(EXISTS ${dir})
+      list(APPEND INCLUDE_COMPILER_LIST "-I${dir}")
+    endif()
+  endforeach()
 
   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}
+      DEPENDS ${PHASE0} 
       OUTPUT ${PHASE1}
-      COMMAND ${CMAKE_CXX_COMPILER} -flto -c ${PHASE0} -O2 -fno-vectorize -ffast-math -fno-unroll-loops -o ${PHASE1}
+      COMMAND ${CMAKE_CXX_COMPILER} -flto -c ${PHASE0} ${INCLUDE_COMPILER_LIST} -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}")
@@ -148,9 +166,9 @@ macro(enzyme_add_executable)
     )
 
   add_custom_command(
-    DEPENDS ${PHASE4}
+    DEPENDS ${PHASE4} ${enzyme_add_executable_LINK_LIBRARIES}
     OUTPUT ${PHASE5}
-    COMMAND ${CMAKE_CXX_COMPILER} ${PHASE4} -o ${PHASE5}
+    COMMAND ${CMAKE_CXX_COMPILER} ${LINKER_FLAGS} ${PHASE4} -o ${PHASE5}
     )
 
   add_custom_target(

examples/CMakeLists.txt

@@ -56,7 +56,7 @@
 #
 
 add_subdirectory(MatPowerTesting)
-add_subdirectory(SparseTest)
+add_subdirectory(LinearAlgebra)
 add_subdirectory(DistributedGeneratorTest)
 
 if(TARGET SUNDIALS::kinsol)

examples/DistributedGeneratorTest/CMakeLists.txt

@@ -4,7 +4,7 @@
 
 add_executable(dgtest DGTest.cpp)
 target_link_libraries(dgtest GRIDKIT::powerelec_disgen
-                                GRIDKIT::powerelec_mircoline
+                                GRIDKIT::powerelec_microline
                                 GRIDKIT::powerelec_microload
                                 GRIDKIT::solvers_dyn)
 

examples/Enzyme/CMakeLists.txt

@@ -1,2 +1,3 @@
 add_subdirectory(Standalone)
 add_subdirectory(Library)
+add_subdirectory(PowerElectronics)

examples/Enzyme/Library/CMakeLists.txt

@@ -1,4 +1,2 @@
-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)
+add_subdirectory(Scalar)
+add_subdirectory(Vector)

examples/Enzyme/Library/Scalar/CMakeLists.txt

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

examples/Enzyme/Library/Scalar/EnzymeScalar.cpp

@@ -0,0 +1,23 @@
+#include <iostream>
+#include <limits>
+#include "ScalarModel.hpp"
+
+int main() {
+    int fail = 0;
+    ScalarModel scalar_model;
+    double var = 5.0;
+    scalar_model.setVariable(var);
+    scalar_model.evalFunction();
+    scalar_model.evalDerivative();
+    double sq = scalar_model.getFunctionValue();
+    double dsq = scalar_model.getDerivativeValue();
+
+    std::cout << "x = " << var << ", x^2 = " << sq << ", d(x^2)/dx = " << dsq << "\n"; 
+    if (std::abs(dsq - 2.0*var) > std::numeric_limits<double>::epsilon())
+    {
+        fail++;
+        std::cout << "Result incorrect\n";
+    }
+    std::cout << "Status: " << fail << "\n";
+    return fail;
+}

examples/Enzyme/Library/Scalar/EnzymeWrapper.hpp

@@ -0,0 +1,15 @@
+#pragma once
+
+int enzyme_dup;
+int enzyme_dupnoneed;
+int enzyme_out;
+int enzyme_const;
+
+template <typename return_type, typename ... T>
+return_type __enzyme_fwddiff(return_type*, int, T* ... );
+
+template <typename return_type, typename T>
+return_type wrapper(T* obj) {
+    obj->evalFunction();
+    return obj->getFunctionValue();
+}

examples/Enzyme/Library/Scalar/ScalarModel.cpp

@@ -0,0 +1,34 @@
+#include "EnzymeWrapper.hpp"
+#include "ScalarModel.hpp"
+#include <iostream>
+
+inline
+double ScalarModel::square(double x) {
+    return x * x;
+}
+
+void ScalarModel::setVariable(double x) {
+    x_ = x;
+}
+
+void ScalarModel::evalFunction() {
+    f_ = square(x_);
+}
+
+void ScalarModel::evalDerivative() {
+    ScalarModel d_scalar_model;
+    d_scalar_model.setVariable(1.0);
+    dfdx_ = __enzyme_fwddiff<double, ScalarModel>((double*)wrapper<double, ScalarModel>, enzyme_dup, this, &d_scalar_model);
+}
+
+double ScalarModel::getVariable() const {
+    return x_;
+}
+
+double ScalarModel::getFunctionValue() const {
+    return f_;
+}
+
+double ScalarModel::getDerivativeValue() const {
+    return dfdx_;
+}

examples/Enzyme/Library/Scalar/ScalarModel.hpp

@@ -0,0 +1,17 @@
+#pragma once
+
+class ScalarModel {
+private:
+    double x_, f_, dfdx_;
+    inline double square(double);
+
+public:
+    ScalarModel() {};
+    void setVariable(double);
+    void evalFunction();
+    void evalDerivative();
+    double getVariable() const;
+    double getFunctionValue() const;
+    double getDerivativeValue() const;
+    ~ScalarModel() {};
+};

examples/Enzyme/Library/Vector/CMakeLists.txt

@@ -0,0 +1,7 @@
+enzyme_add_executable(
+  NAME EnzymeLibVectorCheck 
+  SOURCES EnzymeVector.cpp VectorModel.cpp
+  LINK_LIBRARIES GRIDKIT::DenseMatrix
+)
+
+add_test(NAME "EnzymeLibVectorCheck" COMMAND ${CMAKE_CURRENT_BINARY_DIR}/EnzymeLibVectorCheck)

examples/Enzyme/Library/Vector/EnzymeVector.cpp

@@ -0,0 +1,76 @@
+#include <iostream>
+#include <limits>
+#include "VectorModel.hpp"
+
+inline
+double dsquare_ref_scalar(double x) {
+    return 2.0 * x;
+}
+
+DenseMatrix dsquare_ref(std::vector<double> x, std::vector<double> y) {
+    DenseMatrix jac(x.size(), y.size());
+    for (int idy = 0; idy < y.size(); ++idy)
+    {
+        for (int idx = 0; idx < x.size(); ++idx)
+        {
+            if (idx == idy) 
+                jac.setValue(idx, idy, dsquare_ref_scalar(x[idx]));
+        }
+    }
+    return jac;
+}
+
+int main() {
+    // Size and variable declarations
+    constexpr int N = 10;
+    std::vector<double> var(N);
+
+    // Random input values
+    srand(time(NULL));
+    for (int idx = 0; idx < var.size(); ++idx)
+    {
+        var[idx] = rand();
+    }
+
+    // Model
+    VectorModel* vector_model = new VectorModel(N);
+    vector_model->setVariable(var);
+    vector_model->evalResidual();
+    vector_model->evalJacobian();
+    std::vector<double> var_temp = vector_model->getVariable();
+    std::vector<double> res = vector_model->getResidual();
+    DenseMatrix jac = vector_model->getJacobian();
+
+    // Reference Jacobian
+    DenseMatrix jac_ref = dsquare_ref(var, res);
+
+    // Check
+    int fail = 0;
+    bool verbose = true;
+    for (int idy = 0; idy < res.size(); ++idy)
+    {
+        for (int idx = 0; idx < var.size(); ++idx)
+        {
+            if (std::abs(jac.getValue(idx, idy) - jac_ref.getValue(idx, idy)) > std::numeric_limits<double>::epsilon())
+            {
+                fail++;
+                if (verbose)
+                {
+                    std::cout << "Result incorrect at line = " << idy << ", column = " << idx << "\n";
+                    std::cout << "x = " << var_temp[idx] << ", x^2 = " << res[idx] << ", d(x^2)/dx = " << jac.getValue(idx, idy) << "\n"; 
+                }
+            }
+        }
+    }
+    if (verbose)
+    {
+        jac.printMatrix("Autodiff Jacobian");
+        jac_ref.printMatrix("Reference Jacobian");
+    }
+    std::cout << "Status: " << fail << "\n";
+
+    // Cleanup
+    delete vector_model;
+
+    return fail;
+}

examples/Enzyme/Library/Vector/EnzymeWrapper.hpp

@@ -0,0 +1,17 @@
+#pragma once
+
+#include <vector>
+
+int enzyme_dup;
+int enzyme_dupnoneed;
+int enzyme_out;
+int enzyme_const;
+
+template <typename T>
+std::vector<double> __enzyme_fwddiff(std::vector<double>*, int, T*, T*);
+
+template <typename T>
+std::vector<double> wrapper(T* obj) {
+    obj->evalResidual();
+    return obj->getResidual();
+}

examples/Enzyme/Library/Vector/VectorModel.cpp

@@ -0,0 +1,72 @@
+#include "EnzymeWrapper.hpp"
+#include "VectorModel.hpp"
+#include <iostream>
+
+VectorModel::VectorModel(int N) :
+   x_(N),
+   f_(N),
+   dfdx_(N, N)
+{
+};
+
+inline
+double VectorModel::square_scalar(double x) {
+    return x * x;
+}
+
+void VectorModel::square(std::vector<double>& x, std::vector<double>& y) {
+    for (int idx = 0; idx < x.size(); ++idx)
+    {
+        y[idx] = this->square_scalar(x[idx]);
+    }
+}
+
+void VectorModel::setVariable(std::vector<double> x) {
+    for (int idx = 0; idx < x.size(); ++idx)
+    {
+        x_[idx] = x[idx];
+    }
+}
+
+void VectorModel::evalResidual() {
+    square(x_, f_);
+}
+
+void VectorModel::evalJacobian() {
+    const int N = x_.size();
+    std::vector<double> v(N);
+    VectorModel d_vector_model(N);
+    for (int idy = 0; idy < N; ++idy)
+    {
+        // Elementary vector for Jacobian-vector product
+        for (int idx = 0; idx < N; ++idx)
+        {
+            v[idx] = 0.0;
+        }
+        v[idy] = 1.0;
+        d_vector_model.setVariable(v);
+
+        // Autodiff
+        std::vector<double> d_res = __enzyme_fwddiff<VectorModel>(
+                                      (std::vector<double>*)wrapper<VectorModel>, 
+                                      enzyme_dup, this, &d_vector_model);
+
+        // Store result
+        for (int idx = 0; idx < N; ++idx)
+        {
+            dfdx_.setValue(idx, idy, d_res[idx]);
+        }
+    }
+}
+
+std::vector<double>& VectorModel::getVariable() {
+    return x_;
+}
+
+std::vector<double>& VectorModel::getResidual() {
+    return f_;
+}
+
+DenseMatrix& VectorModel::getJacobian() {
+    return dfdx_;
+}

examples/Enzyme/Library/Vector/VectorModel.hpp

@@ -0,0 +1,23 @@
+#pragma once
+
+#include <LinearAlgebra/DenseMatrix/DenseMatrix.hpp>
+
+using DenseMatrix = GridKit::LinearAlgebra::DenseMatrix<double, size_t>;
+
+class VectorModel {
+private:
+    std::vector<double> x_, f_;
+    DenseMatrix dfdx_;
+    inline double square_scalar(double);
+    void square(std::vector<double>&, std::vector<double>&);
+
+public:
+    VectorModel(int);
+    void setVariable(std::vector<double>);
+    void evalResidual();
+    void evalJacobian();
+    std::vector<double>& getVariable();
+    std::vector<double>& getResidual();
+    DenseMatrix& getJacobian();
+    ~VectorModel() {};
+};