Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
The table of contents is too big for display.
Diff view
Diff view
  •  
  •  
  •  
5 changes: 3 additions & 2 deletions .clang-format
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ BasedOnStyle: Microsoft

# Indentation
IndentWidth: 2 # 2 spaces per indent
ContinuationIndentWidth: 2
AccessModifierOffset: -2
IndentAccessModifiers: false # Align access modifiers to braces
NamespaceIndentation: All # Indent namspace contents
Expand All @@ -12,7 +13,6 @@ ConstructorInitializerIndentWidth: 2
# Comments which match this regex will be unformatted (and therefore can be longer or have more whitespace than other comments)
CommentPragmas: '^\*\*'


# Alignment
AlignConsecutiveAssignments:
Enabled: true
Expand All @@ -37,7 +37,8 @@ AlignConsecutiveMacros:
AcrossComments: true

# Newlines
ColumnLimit: 0
AlignAfterOpenBracket: BlockIndent # this will need to be replaced w/ clang-format 22
ColumnLimit: 80
BreakBeforeBraces: Allman # Braces on their own lines
SeparateDefinitionBlocks: Always # Separate definitions (functions etc.) with an empty line
AlwaysBreakTemplateDeclarations: true # Put template on their own lines
Expand Down
12 changes: 3 additions & 9 deletions GridKit/AutomaticDifferentiation/DependencyTracking/Variable.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -36,19 +36,15 @@ namespace GridKit
@brief Default constructor.
*/
Variable()
: value_(0.0),
variable_number_(INVALID_VAR_NUMBER),
is_fixed_(false)
: value_(0.0), variable_number_(INVALID_VAR_NUMBER), is_fixed_(false)
{
}

/**
@brief Constructor which initializes the value.
*/
explicit Variable(double value)
: value_(value),
variable_number_(INVALID_VAR_NUMBER),
is_fixed_(false)
: value_(value), variable_number_(INVALID_VAR_NUMBER), is_fixed_(false)
{
}

Expand All @@ -58,9 +54,7 @@ namespace GridKit

*/
Variable(double value, size_t variable_number)
: value_(value),
variable_number_(variable_number),
is_fixed_(false)
: value_(value), variable_number_(variable_number), is_fixed_(false)
{
dependencies_[variable_number_] = 1.0;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -367,33 +367,41 @@ namespace GridKit
namespace std
{

#define IMPL_FUN_1(FUN, DER) \
inline GridKit::DependencyTracking::Variable FUN(const GridKit::DependencyTracking::Variable& x) \
{ \
double val = FUN(x()); \
double der = DER(x()); \
GridKit::DependencyTracking::Variable res(x); /* copy derivatives of x*/ \
res.setValue(val); /* set function value f(x) */ \
res.scaleDependencies(der); /* compute derivatives of f(x) */ \
return res; \
#define IMPL_FUN_1(FUN, DER) \
inline GridKit::DependencyTracking::Variable FUN( \
const GridKit::DependencyTracking::Variable& x \
) \
{ \
double val = FUN(x()); \
double der = DER(x()); \
GridKit::DependencyTracking::Variable res(x); /* copy derivatives of x*/ \
res.setValue(val); /* set function value f(x) */ \
res.scaleDependencies(der); /* compute derivatives of f(x) */ \
return res; \
}

#define IMPL_FUN_2(FUN, DER1, DER2) \
inline GridKit::DependencyTracking::Variable FUN(const GridKit::DependencyTracking::Variable& x, \
const GridKit::DependencyTracking::Variable& y) \
{ \
double val = FUN(x(), y()); \
double der1 = DER1(x(), y()); \
double der2 = DER2(x(), y()); \
GridKit::DependencyTracking::Variable res1(x); /* copy derivatives of x*/ \
GridKit::DependencyTracking::Variable res2(y); /* copy derivatives of y*/ \
res1.setValue(val); /* set function value f(x, y) */ \
res2.setValue(val); /* set function value f(x, y) */ \
res1.scaleDependencies(der1); /* compute derivatives of f(x, y) with respect to x */ \
res2.scaleDependencies(der2); /* compute derivatives of f(x, y) with respect to y */ \
GridKit::DependencyTracking::Variable res(std::move(res1)); /* relocate derivatives w.r.t. x into return */ \
res.addDependencies(res2); /* add derivatives with respect to y to return */ \
return res; \
#define IMPL_FUN_2(FUN, DER1, DER2) \
inline GridKit::DependencyTracking::Variable FUN( \
const GridKit::DependencyTracking::Variable& x, \
const GridKit::DependencyTracking::Variable& y \
) \
{ \
double val = FUN(x(), y()); \
double der1 = DER1(x(), y()); \
double der2 = DER2(x(), y()); \
GridKit::DependencyTracking::Variable res1(x); /* copy derivatives of x*/ \
GridKit::DependencyTracking::Variable res2(y); /* copy derivatives of y*/ \
res1.setValue(val); /* set function value f(x, y) */ \
res2.setValue(val); /* set function value f(x, y) */ \
res1.scaleDependencies(der1 \
); /* compute derivatives of f(x, y) with respect to x */ \
res2.scaleDependencies(der2 \
); /* compute derivatives of f(x, y) with respect to y */ \
GridKit::DependencyTracking::Variable res(std::move(res1) \
); /* relocate derivatives w.r.t. x into return */ \
res.addDependencies(res2 \
); /* add derivatives with respect to y to return */ \
return res; \
}

IMPL_FUN_1(sin, GridKit::DependencyTracking::sin_derivative)
Expand All @@ -412,7 +420,11 @@ namespace std
IMPL_FUN_1(sqrt, GridKit::DependencyTracking::sqrt_derivative)
IMPL_FUN_1(abs, GridKit::DependencyTracking::abs_derivative)

IMPL_FUN_2(atan2, GridKit::DependencyTracking::atan2_derivative1, GridKit::DependencyTracking::atan2_derivative2)
IMPL_FUN_2(
atan2,
GridKit::DependencyTracking::atan2_derivative1,
GridKit::DependencyTracking::atan2_derivative2
)

#undef IMPL_FUN_1
#undef IMPL_FUN_2
Expand Down
132 changes: 74 additions & 58 deletions GridKit/AutomaticDifferentiation/Enzyme/D2LDx2.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,9 @@ namespace GridKit
namespace Sparse
{
/**
* @brief Enzyme sparse Hessian of \f$L = \sigma f + \lambda^T g\f$, by one
* forward sweep per local variable over the reverse gradient of \f$L\f$
* @brief Enzyme sparse Hessian of \f$L = \sigma f + \lambda^T g\f$, by
* one forward sweep per local variable over the reverse gradient of
* \f$L\f$
*
* @tparam ModelT - optimization component model
*/
Expand All @@ -40,48 +41,56 @@ namespace GridKit
* @param[out] nnz - Number of nonzeros
*/
template <typename ScalarT, typename IdxT, typename RealT>
static void eval(const ModelT* model,
const size_t n_var,
const IdxT* variable_indices,
const ScalarT* x,
const RealT sigma,
const RealT* lambda,
IdxT* rows,
IdxT* cols,
ScalarT* vals,
IdxT& nnz)
static void eval(
const ModelT* model,
const size_t n_var,
const IdxT* variable_indices,
const ScalarT* x,
const RealT sigma,
const RealT* lambda,
IdxT* rows,
IdxT* cols,
ScalarT* vals,
IdxT& nnz
)
{
std::array<ScalarT, ObjectiveT::NUM_INPUTS> gradient{};
for (size_t var_i = 0; var_i < n_var; ++var_i)
{
// Sparse storage. @see LowerSparseStorage.hpp
ScalarT* seed = __enzyme_todense<ScalarT*>((void*) ident_load<ScalarT, IdxT>,
(void*) ident_store<ScalarT, IdxT>,
var_i);
ScalarT* dgradient = __enzyme_todense<ScalarT*>((void*) sparse_load<ScalarT, IdxT>,
(void*) sparse_store<ScalarT, IdxT>,
var_i,
1.0, // value scaling
variable_indices,
variable_indices,
rows,
cols,
vals,
&nnz);
ScalarT* seed = __enzyme_todense<ScalarT*>(
(void*) ident_load<ScalarT, IdxT>,
(void*) ident_store<ScalarT, IdxT>,
var_i
);
ScalarT* dgradient = __enzyme_todense<ScalarT*>(
(void*) sparse_load<ScalarT, IdxT>,
(void*) sparse_store<ScalarT, IdxT>,
var_i,
1.0, // value scaling
variable_indices,
variable_indices,
rows,
cols,
vals,
&nnz
);

__enzyme_fwddiff<void>((void*) lagrangianGradient<ScalarT, RealT>,
enzyme_const,
model,
enzyme_dup,
x,
seed,
enzyme_dupnoneed,
gradient.data(),
dgradient,
enzyme_const,
sigma,
enzyme_const,
lambda);
__enzyme_fwddiff<void>(
(void*) lagrangianGradient<ScalarT, RealT>,
enzyme_const,
model,
enzyme_dup,
x,
seed,
enzyme_dupnoneed,
gradient.data(),
dgradient,
enzyme_const,
sigma,
enzyme_const,
lambda
);
}
}

Expand All @@ -90,10 +99,12 @@ namespace GridKit
using ConstraintsT = ModelWrapper<ModelT, MemberFunctions::Constraints>;

template <typename ScalarT, typename RealT>
[[gnu::always_inline]] inline static ScalarT lagrangian(const ModelT* model,
const ScalarT* x,
const RealT sigma,
const RealT* lambda)
[[gnu::always_inline]] inline static ScalarT lagrangian(
const ModelT* model,
const ScalarT* x,
const RealT sigma,
const RealT* lambda
)
{
std::array<ScalarT, ConstraintsT::NUM_OUTPUTS> g{};
ConstraintsT::eval(model, x, g.data());
Expand All @@ -106,24 +117,29 @@ namespace GridKit
return value;
}

/// Add the Lagrangian gradient to `gradient`, whose tangent alone is used
/// Add the Lagrangian gradient to `gradient`, whose tangent alone is
/// used
template <typename ScalarT, typename RealT>
[[gnu::always_inline]] inline static void lagrangianGradient(const ModelT* model,
const ScalarT* x,
ScalarT* gradient,
const RealT sigma,
const RealT* lambda)
[[gnu::always_inline]] inline static void lagrangianGradient(
const ModelT* model,
const ScalarT* x,
ScalarT* gradient,
const RealT sigma,
const RealT* lambda
)
{
__enzyme_autodiff<void>((void*) lagrangian<ScalarT, RealT>,
enzyme_const,
model,
enzyme_dup,
x,
gradient,
enzyme_const,
sigma,
enzyme_const,
lambda);
__enzyme_autodiff<void>(
(void*) lagrangian<ScalarT, RealT>,
enzyme_const,
model,
enzyme_dup,
x,
gradient,
enzyme_const,
sigma,
enzyme_const,
lambda
);
}
};
} // namespace Sparse
Expand Down
Loading