Skip to content

Commit 49272c3

Browse files
nkoukpaizanWiktoria Zielinska
nkoukpaizan
authored and
Wiktoria Zielinska
committed
Enzyme Jacobian for PhasorDynamics::Load within GridKit (#131)
* Separated out Enzyme macros from FindEnzyme and added a EnzymeAddLibrary macro. * Enzyme wrapper within GridKit. * Fixes to rebase. * More interesting vector residual for Enzyme examples. * Fix some warnings. * Address one more warning. * Apply pre-commmit fixes * Empty commit to trigger CI after pre-commit fixes. * Successful use of target_compiile_options to pass enzyme-auto-sparsity flags in examples. * More flags needed to get the correct answer to machine precision in EnzymePowerElectronicsCheck. * EnzymeAddLibrary not currently needed. * Apply pre-commmit fixes * Empty commit to trigger CI after pre-commit fixes. * Updated documentation in EnzymeAddLibrary. * Fix warnings previously hidden by manual builds through enzyme_add_executable CMake macro. * Use GridKit::Testing::isEqual in Enzyme examples. * Fixes to rebase. * Apply pre-commmit fixes * Empty commit to trigger CI after pre-commit fixes. * residual_wrapper can now be placed within GridKit::Enzyme namespace. --------- Co-authored-by: nkoukpaizan <nkoukpaizan@users.noreply.github.com>
1 parent a24e6fa commit 49272c3

25 files changed

Lines changed: 784 additions & 395 deletions

‎CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,8 @@ endif()
103103

104104
if(${GRIDKIT_ENABLE_ENZYME})
105105
include(FindEnzyme)
106+
# todo Add a centralized configuration file
107+
add_definitions(-DGRIDKIT_ENABLE_ENZYME)
106108
endif()
107109

108110
# Macro that adds libraries

‎cmake/EnzymeAddLibrary.cmake‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
#
2+
#[[
3+
4+
Macro to manually compile with Enzyme
5+
6+
Author(s):
7+
- Asher Mancinelli <ashermancinelli@gmail.com>
8+
- Nicholson Koukpaizan <koukpaizannk@ornl.gov>
9+
10+
]]
11+
12+
macro(enzyme_build_object)
13+
set(options)
14+
set(oneValueArgs NAME)
15+
set(multiValueArgs SOURCES LINK_LIBRARIES INCLUDE_DIRECTORIES)
16+
cmake_parse_arguments(enzyme_build_object "${options}" "${oneValueArgs}"
17+
"${multiValueArgs}" ${ARGN})
18+
19+
set(PHASE2 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_build_object_NAME}.bc")
20+
set(PHASE3 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_build_object_NAME}_enzyme.ll")
21+
set(PHASE4 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_build_object_NAME}_opt.ll")
22+
set(PHASE5 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_build_object_NAME}")
23+
24+
set(OBJS "")
25+
set(includes "${enzyme_build_object_INCLUDE_DIRECTORIES}")
26+
27+
foreach(lib ${enzyme_build_object_LINK_LIBRARIES})
28+
get_target_property(include ${lib} INCLUDE_DIRECTORIES)
29+
set(includes "${includes}" ${include})
30+
31+
get_target_property(libsource ${lib} SOURCES)
32+
string(FIND "${libsource}" "TARGET" found)
33+
if(NOT(${found} EQUAL -1))
34+
list(APPEND LINKER_FLAGS "-Wl,${libsource}")
35+
endif()
36+
endforeach()
37+
38+
foreach(dir ${includes})
39+
if(EXISTS ${dir})
40+
list(APPEND INCLUDE_COMPILER_LIST "-I${dir}")
41+
endif()
42+
endforeach()
43+
44+
foreach(SRC ${enzyme_build_object_SOURCES})
45+
set(PHASE0 "${CMAKE_CURRENT_SOURCE_DIR}/${SRC}")
46+
set(PHASE1 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_build_object_NAME}_${SRC}_compile.o")
47+
add_custom_command(
48+
DEPENDS ${PHASE0}
49+
OUTPUT ${PHASE1}
50+
COMMAND ${CMAKE_CXX_COMPILER} -flto -c ${PHASE0} ${INCLUDE_COMPILER_LIST} -O2 -fno-vectorize -ffast-math -fno-unroll-loops -fpass-plugin=${ENZYME_CLANG_PLUGIN_LIBRARY} -Xclang -load -Xclang ${ENZYME_CLANG_PLUGIN_LIBRARY} -mllvm -enable-load-pre=0 -mllvm -enzyme-auto-sparsity=1 -o ${PHASE1}
51+
COMMENT "Compiling ${SRC} to object file for target ${enzyme_build_object_NAME}"
52+
)
53+
set(OBJS "${OBJS} ${PHASE1}")
54+
endforeach()
55+
56+
cmake_language(EVAL CODE "
57+
add_custom_command(
58+
DEPENDS ${OBJS}
59+
OUTPUT ${PHASE2}
60+
COMMAND ${GRIDKIT_LLVM_LINK} ${OBJS} -o ${PHASE2}
61+
COMMENT \"Linking object files to LLVM bytecode for target ${enzyme_build_object_NAME}\"
62+
)
63+
")
64+
65+
add_custom_command(
66+
DEPENDS ${PHASE2}
67+
OUTPUT ${PHASE3}
68+
COMMAND ${GRIDKIT_OPT} ${PHASE2} -load-pass-plugin=${ENZYME_LLVM_PLUGIN_LIBRARY} -passes=enzyme -o ${PHASE3} -S
69+
COMMENT "Running Enzyme opt pass on target ${enzyme_build_object_NAME}"
70+
)
71+
72+
add_custom_command(
73+
DEPENDS ${PHASE3}
74+
OUTPUT ${PHASE4}
75+
COMMAND ${GRIDKIT_OPT} ${PHASE3} -O2 -o ${PHASE4} -S
76+
COMMENT "Running remaining opt passes on target ${enzyme_build_object_NAME}"
77+
)
78+
79+
add_custom_command(
80+
DEPENDS ${PHASE4}
81+
OUTPUT ${PHASE5}
82+
COMMAND ${CMAKE_CXX_COMPILER} -c ${PHASE4} -o ${PHASE5}
83+
COMMENT "Generating optimized object file for target ${enzyme_build_object_NAME}"
84+
)
85+
endmacro()
86+
87+
macro(enzyme_add_executable)
88+
set(options)
89+
set(oneValueArgs NAME)
90+
set(multiValueArgs SOURCES LINK_LIBRARIES INCLUDE_DIRECTORIES)
91+
cmake_parse_arguments(enzyme_add_executable "${options}" "${oneValueArgs}"
92+
"${multiValueArgs}" ${ARGN})
93+
94+
enzyme_build_object(
95+
NAME "${enzyme_add_executable_NAME}.o"
96+
SOURCES ${enzyme_add_executable_SOURCES}
97+
LINK_LIBRARIES ${enzyme_add_executable_LINK_LIBRARIES}
98+
INCLUDE_DIRECTORIES ${enzyme_add_executable_INCLUDE_DIRECTORIES}
99+
)
100+
101+
add_executable("${enzyme_add_executable_NAME}" "${enzyme_add_executable_NAME}.o")
102+
set_target_properties("${enzyme_add_executable_NAME}" PROPERTIES LINKER_LANGUAGE CXX)
103+
target_link_libraries("${enzyme_add_executable_NAME}" ${enzyme_add_executable_LINK_LIBRARIES})
104+
endmacro()
105+
106+
macro(enzyme_add_library)
107+
set(options)
108+
set(oneValueArgs NAME)
109+
set(multiValueArgs SOURCES LINK_LIBRARIES INCLUDE_DIRECTORIES)
110+
cmake_parse_arguments(enzyme_add_library "${options}" "${oneValueArgs}"
111+
"${multiValueArgs}" ${ARGN})
112+
113+
enzyme_build_object(
114+
NAME "${enzyme_add_library_NAME}.o"
115+
SOURCES ${enzyme_add_library_SOURCES}
116+
LINK_LIBRARIES ${enzyme_add_library_LINK_LIBRARIES}
117+
INCLUDE_DIRECTORIES ${enzyme_add_library_INCLUDE_DIRECTORIES}
118+
)
119+
120+
add_library("${enzyme_add_library_NAME}" "${enzyme_add_library_NAME}.o")
121+
set_target_properties("${enzyme_add_library_NAME}" PROPERTIES LINKER_LANGUAGE CXX)
122+
target_link_libraries("${enzyme_add_library_NAME}" ${enzyme_add_library_LINK_LIBRARIES})
123+
endmacro()

‎cmake/FindEnzyme.cmake‎

Lines changed: 0 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -59,82 +59,3 @@ find_program(GRIDKIT_OPT opt
5959
bin
6060
REQUIRED)
6161
message(STATUS "opt: ${GRIDKIT_OPT}")
62-
63-
macro(enzyme_add_executable)
64-
set(options)
65-
set(oneValueArgs NAME)
66-
set(multiValueArgs SOURCES LINK_LIBRARIES INCLUDE_DIRECTORIES)
67-
cmake_parse_arguments(enzyme_add_executable "${options}" "${oneValueArgs}"
68-
"${multiValueArgs}" ${ARGN})
69-
70-
set(PHASE2 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_add_executable_NAME}.bc")
71-
set(PHASE3 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_add_executable_NAME}_enzyme.ll")
72-
set(PHASE4 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_add_executable_NAME}_opt.ll")
73-
set(PHASE5 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_add_executable_NAME}")
74-
75-
set(OBJS "")
76-
set(includes "${enzyme_add_executable_INCLUDE_DIRECTORIES}")
77-
78-
foreach(lib ${enzyme_add_executable_LINK_LIBRARIES})
79-
get_target_property(include ${lib} INCLUDE_DIRECTORIES)
80-
set(includes "${includes}" ${include})
81-
82-
get_target_property(libsource ${lib} SOURCES)
83-
string(FIND "${libsource}" "TARGET" found)
84-
if(NOT(${found} EQUAL -1))
85-
list(APPEND LINKER_FLAGS "-Wl,${libsource}")
86-
endif()
87-
endforeach()
88-
89-
foreach(dir ${includes})
90-
if(EXISTS ${dir})
91-
list(APPEND INCLUDE_COMPILER_LIST "-I${dir}")
92-
endif()
93-
endforeach()
94-
95-
foreach(SRC ${enzyme_add_executable_SOURCES})
96-
set(PHASE0 "${CMAKE_CURRENT_SOURCE_DIR}/${SRC}")
97-
set(PHASE1 "${CMAKE_CURRENT_BINARY_DIR}/${enzyme_add_executable_NAME}_${SRC}_compile.o")
98-
add_custom_command(
99-
DEPENDS ${PHASE0}
100-
OUTPUT ${PHASE1}
101-
COMMAND ${CMAKE_CXX_COMPILER} -flto -c ${PHASE0} ${INCLUDE_COMPILER_LIST} -O2 -fno-vectorize -ffast-math -fno-unroll-loops -fpass-plugin=${ENZYME_CLANG_PLUGIN_LIBRARY} -Xclang -load -Xclang ${ENZYME_CLANG_PLUGIN_LIBRARY} -mllvm -enable-load-pre=0 -mllvm -enzyme-auto-sparsity=1 -o ${PHASE1}
102-
COMMENT "Compiling ${SRC} to object file for target ${enzyme_add_executable_NAME}"
103-
)
104-
set(OBJS "${OBJS} ${PHASE1}")
105-
endforeach()
106-
107-
cmake_language(EVAL CODE "
108-
add_custom_command(
109-
DEPENDS ${OBJS}
110-
OUTPUT ${PHASE2}
111-
COMMAND ${GRIDKIT_LLVM_LINK} ${OBJS} -o ${PHASE2}
112-
COMMENT \"Linking object files to LLVM bytecode for target ${enzyme_add_executable_NAME}\"
113-
)
114-
")
115-
116-
add_custom_command(
117-
DEPENDS ${PHASE2}
118-
OUTPUT ${PHASE3}
119-
COMMAND ${GRIDKIT_OPT} ${PHASE2} -load-pass-plugin=${ENZYME_LLVM_PLUGIN_LIBRARY} -passes=enzyme -o ${PHASE3} -S
120-
COMMENT "Running Enzyme opt pass on target ${enzyme_add_executable_NAME}"
121-
)
122-
123-
add_custom_command(
124-
DEPENDS ${PHASE3}
125-
OUTPUT ${PHASE4}
126-
COMMAND ${GRIDKIT_OPT} ${PHASE3} -O2 -o ${PHASE4} -S
127-
COMMENT "Running remaining opt passes on target ${enzyme_add_executable_NAME}"
128-
)
129-
130-
add_custom_command(
131-
DEPENDS ${PHASE4} ${enzyme_add_executable_LINK_LIBRARIES}
132-
OUTPUT ${PHASE5}
133-
COMMAND ${CMAKE_CXX_COMPILER} ${LINKER_FLAGS} ${PHASE4} -o ${PHASE5}
134-
)
135-
136-
add_custom_target(
137-
"${enzyme_add_executable_NAME}_target" ALL
138-
DEPENDS ${PHASE5}
139-
)
140-
endmacro()

‎examples/Enzyme/Library/Scalar/CMakeLists.txt‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,4 @@
1-
enzyme_add_executable(
2-
NAME EnzymeLibScalarCheck
3-
SOURCES EnzymeScalar.cpp ScalarModel.cpp
4-
)
1+
add_executable(EnzymeLibScalarCheck EnzymeScalar.cpp ScalarModel.cpp)
2+
target_link_libraries(EnzymeLibScalarCheck ClangEnzymeFlags GRIDKIT::Utilities)
53

64
add_test(NAME "EnzymeLibScalarCheck" COMMAND ${CMAKE_CURRENT_BINARY_DIR}/EnzymeLibScalarCheck)

‎examples/Enzyme/Library/Scalar/EnzymeScalar.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
#include <iostream>
2-
#include <limits>
32

43
#include "ScalarModel.hpp"
4+
#include <Utilities/Testing.hpp>
55

66
/**
77
* @brief Example that computes the derivative of a library function
@@ -23,7 +23,7 @@ int main()
2323
double dsq = scalar_model.getDerivativeValue();
2424

2525
std::cout << "x = " << var << ", x^2 = " << sq << ", d(x^2)/dx = " << dsq << "\n";
26-
if (std::abs(dsq - 2.0 * var) > std::numeric_limits<double>::epsilon())
26+
if (!GridKit::Testing::isEqual(dsq, 2.0 * var))
2727
{
2828
fail++;
2929
std::cout << "Result incorrect\n";

‎examples/Enzyme/Library/Vector/CMakeLists.txt‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,4 @@
1-
enzyme_add_executable(
2-
NAME EnzymeLibVectorCheck
3-
SOURCES EnzymeVector.cpp VectorModel.cpp
4-
LINK_LIBRARIES GRIDKIT::DenseMatrix
5-
)
1+
add_executable(EnzymeLibVectorCheck EnzymeVector.cpp VectorModel.cpp)
2+
target_link_libraries(EnzymeLibVectorCheck ClangEnzymeFlags GRIDKIT::DenseMatrix GRIDKIT::Utilities)
63

74
add_test(NAME "EnzymeLibVectorCheck" COMMAND ${CMAKE_CURRENT_BINARY_DIR}/EnzymeLibVectorCheck)

‎examples/Enzyme/Library/Vector/EnzymeVector.cpp‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
#include <iostream>
2-
#include <limits>
32

43
#include "VectorModel.hpp"
4+
#include <Utilities/Testing.hpp>
55

66
/**
77
* @brief Example that computes the Jacobian of a vector-valued residual
@@ -20,12 +20,12 @@ inline double dsquare_ref_scalar(double x)
2020
DenseMatrix dsquare_ref(std::vector<double> x, std::vector<double> y)
2121
{
2222
DenseMatrix jac(x.size(), y.size());
23-
for (int idy = 0; idy < y.size(); ++idy)
23+
for (size_t idy = 0; idy < y.size(); ++idy)
2424
{
25-
for (int idx = 0; idx < x.size(); ++idx)
25+
for (size_t idx = 0; idx < x.size(); ++idx)
2626
{
27-
if (idx == idy)
28-
jac.setValue(idx, idy, dsquare_ref_scalar(x[idx]));
27+
if (idy <= idx)
28+
jac.setValue(idx, idy, dsquare_ref_scalar(x[idy]));
2929
}
3030
}
3131
return jac;
@@ -34,12 +34,12 @@ DenseMatrix dsquare_ref(std::vector<double> x, std::vector<double> y)
3434
int main()
3535
{
3636
// Size and variable declarations
37-
constexpr int n = 10;
37+
constexpr size_t n = 10;
3838
std::vector<double> var(n);
3939

4040
// Random input values
41-
srand(time(NULL));
42-
for (int idx = 0; idx < var.size(); ++idx)
41+
srand(static_cast<unsigned int>(time(NULL)));
42+
for (size_t idx = 0; idx < var.size(); ++idx)
4343
{
4444
var[idx] = rand();
4545
}
@@ -59,11 +59,11 @@ int main()
5959
// Check
6060
int fail = 0;
6161
bool verbose = true;
62-
for (int idy = 0; idy < res.size(); ++idy)
62+
for (size_t idy = 0; idy < res.size(); ++idy)
6363
{
64-
for (int idx = 0; idx < var.size(); ++idx)
64+
for (size_t idx = 0; idx < var.size(); ++idx)
6565
{
66-
if (std::abs(jac.getValue(idx, idy) - jac_ref.getValue(idx, idy)) > std::numeric_limits<double>::epsilon())
66+
if (!GridKit::Testing::isEqual(jac.getValue(idx, idy), jac_ref.getValue(idx, idy)))
6767
{
6868
fail++;
6969
if (verbose)

‎examples/Enzyme/Library/Vector/VectorModel.cpp‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
#include "EnzymeWrapper.hpp"
66

7-
VectorModel::VectorModel(int n)
7+
VectorModel::VectorModel(size_t n)
88
: x_(n),
99
f_(n),
1010
df_dx_(n, n)
@@ -18,15 +18,19 @@ inline double VectorModel::square_scalar(double x)
1818

1919
void VectorModel::square(std::vector<double>& x, std::vector<double>& y)
2020
{
21-
for (int idx = 0; idx < x.size(); ++idx)
21+
for (size_t idx = 0; idx < x.size(); ++idx)
2222
{
23-
y[idx] = this->square_scalar(x[idx]);
23+
y[idx] = 0.0;
24+
for (size_t idy = 0; idy <= idx; idy++)
25+
{
26+
y[idx] += this->square_scalar(x[idy]);
27+
}
2428
}
2529
}
2630

2731
void VectorModel::setVariable(std::vector<double> x)
2832
{
29-
for (int idx = 0; idx < x.size(); ++idx)
33+
for (size_t idx = 0; idx < x.size(); ++idx)
3034
{
3135
x_[idx] = x[idx];
3236
}
@@ -39,13 +43,13 @@ void VectorModel::evalResidual()
3943

4044
void VectorModel::evalJacobian()
4145
{
42-
const int n = x_.size();
46+
const size_t n = x_.size();
4347
std::vector<double> v(n);
4448
VectorModel d_vector_model(n);
45-
for (int idy = 0; idy < n; ++idy)
49+
for (size_t idy = 0; idy < n; ++idy)
4650
{
4751
// Elementary vector for Jacobian-vector product
48-
for (int idx = 0; idx < n; ++idx)
52+
for (size_t idx = 0; idx < n; ++idx)
4953
{
5054
v[idx] = 0.0;
5155
}
@@ -60,7 +64,7 @@ void VectorModel::evalJacobian()
6064
&d_vector_model);
6165

6266
// Store result
63-
for (int idx = 0; idx < n; ++idx)
67+
for (size_t idx = 0; idx < n; ++idx)
6468
{
6569
df_dx_.setValue(idx, idy, d_res[idx]);
6670
}

‎examples/Enzyme/Library/Vector/VectorModel.hpp‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ class VectorModel
1717
void square(std::vector<double>&, std::vector<double>&);
1818

1919
public:
20-
VectorModel(int);
20+
VectorModel(size_t);
2121
void setVariable(std::vector<double>);
2222
void evalResidual();
2323
void evalJacobian();

‎examples/Enzyme/PowerElectronics/CMakeLists.txt‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
1-
enzyme_add_executable(
2-
NAME EnzymePowerElectronicsCheck
3-
SOURCES main.cpp
4-
LINK_LIBRARIES GRIDKIT::DenseMatrix GRIDKIT::power_elec_disgen
5-
)
1+
add_executable(EnzymePowerElectronicsCheck main.cpp)
2+
target_compile_options(EnzymePowerElectronicsCheck PUBLIC -fno-vectorize -ffast-math -fno-unroll-loops)
3+
target_link_libraries(EnzymePowerElectronicsCheck ClangEnzymeFlags GRIDKIT::DenseMatrix GRIDKIT::power_elec_disgen GRIDKIT::Utilities)
64

75
add_test(NAME "EnzymePowerElectronicsCheck" COMMAND ${CMAKE_CURRENT_BINARY_DIR}/EnzymePowerElectronicsCheck)

0 commit comments

Comments
 (0)