Skip to content

fix a leak in enzyme AD code with stack memory #239

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ebchin merged 7 commits into from
Aug 21, 2026
Merged

fix a leak in enzyme AD code with stack memory #239

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
7492e7c
fix a leak in enzyme AD code with stack memory
ebchin Aug 20, 2026
2a67f80
fix another leak
ebchin Aug 20, 2026
d498f47
formatting
ebchin Aug 20, 2026
71f9a54
fix naming, make delta from develop smaller
ebchin Aug 21, 2026
56c31f9
just change function signatures instead of add overloads
ebchin Aug 21, 2026
4ac2dc1
restore smooth bounds loop body
ebchin Aug 21, 2026
f0c6b4d
restored smooth_bound(), needed to avoid loop-local tape reuse in enz…
ebchin Aug 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 10 additions & 6 deletions src/tests/tribol_finite_diff_energy_mortar.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,11 +59,13 @@ FiniteDiffResult EnergyMortarCalculator::validate_g_tilde( const InterfacePair&
auto viewer2 = mesh2.getView();

auto projs0 = projections( pair, viewer1, viewer2 );
auto bounds0 = smoother_.bounds_from_projections( projs0, p_.del );
auto smooth_bounds0 = smoother_.smooth_bounds( bounds0, p_.del );
double bounds0[2];
smoother_.bounds_from_projections( projs0.data(), p_.del, bounds0 );
double smooth_bounds0[2];
smoother_.smooth_bounds( bounds0, p_.del, smooth_bounds0 );
QuadPoints qp0;
if ( !p_.enzyme_quadrature ) {
qp0 = compute_quadrature( smooth_bounds0, p_.N );
compute_quadrature( smooth_bounds0, p_.N, &qp0 );
}

auto [g1_base, g2_base] = eval_gtilde( pair, viewer1, viewer2 );
Expand Down Expand Up @@ -290,9 +292,11 @@ FiniteDiffResult EnergyMortarCalculator::validate_hessian( const InterfacePair&
QuadPoints qp0;
if ( !p_.enzyme_quadrature ) {
auto projs0 = projections( pair, viewer1, viewer2 );
auto bounds0 = smoother_.bounds_from_projections( projs0, p_.del );
auto smooth_bounds0 = smoother_.smooth_bounds( bounds0, p_.del );
qp0 = compute_quadrature( smooth_bounds0, p_.N );
double bounds0[2];
smoother_.bounds_from_projections( projs0.data(), p_.del, bounds0 );
double smooth_bounds0[2];
smoother_.smooth_bounds( bounds0, p_.del, smooth_bounds0 );
compute_quadrature( smooth_bounds0, p_.N, &qp0 );
}

auto eval_from_offsets = [&]( const std::array<double, 8>& du ) -> std::pair<double, double> {
Expand Down
191 changes: 99 additions & 92 deletions src/tribol/physics/EnergyMortar.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,9 +44,8 @@ TRIBOL_ENZYME_INLINE void find_normal( const double* coord1, const double* coord
}

// Gets the respective gauss-legendre nodes dependant on quadrature order
TRIBOL_ENZYME_INLINE void determine_legendre_nodes( int N, std::array<double, 3>& x )
TRIBOL_ENZYME_INLINE void determine_legendre_nodes( int N, double* x )
{
// x.resize( N );
if ( N == 1 ) {
x[0] = 0.0;
} else if ( N == 2 ) {
Expand Down Expand Up @@ -79,9 +78,8 @@ TRIBOL_ENZYME_INLINE void determine_legendre_nodes( int N, std::array<double, 3>
}

// Gets the respective gauss-legendre weights dependant on quadrature order
TRIBOL_ENZYME_INLINE void determine_legendre_weights( int N, std::array<double, 3>& W )
TRIBOL_ENZYME_INLINE void determine_legendre_weights( int N, double* W )
{
// W.resize( N );
if ( N == 1 ) {
W[0] = 2.0;
} else if ( N == 2 ) {
Expand Down Expand Up @@ -191,21 +189,16 @@ TRIBOL_ENZYME_INLINE void get_projections( const double* A0, const double* A1, c
const double dyA = A1[1] - A0[1];
const double len2A = dxA * dxA + dyA * dyA;

const double* B_endpoints[2] = { B0, B1 };
double q0[2] = { 0.0, 0.0 };
find_intersection( A0, A1, B0, nB, q0 );
// Convert the physical projection point on A to the local coordinate xi in [-0.5, 0.5].
const double alphaA0 = ( ( q0[0] - A0[0] ) * dxA + ( q0[1] - A0[1] ) * dyA ) / len2A;
const double xi0 = alphaA0 - 0.5;

double xi0 = 0.0, xi1 = 0.0;
for ( int i = 0; i < 2; ++i ) {
double q[2] = { 0.0, 0.0 };
find_intersection( A0, A1, B_endpoints[i], nB, q );
// Convert the physical projection point on A to the local coordinate xi in [-0.5, 0.5].
const double alphaA = ( ( q[0] - A0[0] ) * dxA + ( q[1] - A0[1] ) * dyA ) / len2A;
const double xiA = alphaA - 0.5;

if ( i == 0 )
xi0 = xiA;
else
xi1 = xiA;
}
double q1[2] = { 0.0, 0.0 };
find_intersection( A0, A1, B1, nB, q1 );
const double alphaA1 = ( ( q1[0] - A0[0] ) * dxA + ( q1[1] - A0[1] ) * dyA ) / len2A;
const double xi1 = alphaA1 - 0.5;

double xi_min = std::min( xi0, xi1 );
double xi_max = std::max( xi0, xi1 );
Expand All @@ -214,6 +207,35 @@ TRIBOL_ENZYME_INLINE void get_projections( const double* A0, const double* A1, c
projections[1] = xi_max;
}

// Isolate each endpoint to avoid incorrect loop-local tape reuse in Enzyme reverse mode.
TRIBOL_ENZYME_INLINE double smooth_bound( double bound, double del )
{
double xi = 0.0;
double xi_hat = 0.0;

// Shift from the local coordinate interval [-0.5, 0.5] to [0, 1].
xi = bound + 0.5;
if ( del == 0.0 ) {
xi_hat = xi;
} else {
// Apply quadratic ramps near the endpoints and leave the interior unchanged.
if ( 0.0 - del <= xi && xi <= del ) {
xi_hat = ( 1.0 / ( 4 * del ) ) * ( xi * xi ) + 0.5 * xi + del / 4.0;
} else if ( ( 1.0 - del ) <= xi && xi <= 1.0 + del ) {
double b = -1.0 / ( 4.0 * del );
double c = 0.5 + 1.0 / ( 2.0 * del );
double d = 1.0 - del + ( 1.0 / ( 4.0 * del ) ) * pow( 1.0 - del, 2 ) - 0.5 * ( 1.0 - del ) -
( 1.0 - del ) / ( 2.0 * del );

xi_hat = b * xi * xi + c * xi + d;
} else if ( del <= xi && xi <= ( 1.0 - del ) ) {
xi_hat = xi;
}
}
// Shift the smoothed coordinate back to [-0.5, 0.5].
return xi_hat - 0.5;
}

// Integrate the nodal smoothed gap and tributary area contributions over edge A.
// The quadrature rule is supplied through gp, allowing this kernel to be reused
// for both fixed-quadrature and geometry-dependent quadrature paths.
Expand Down Expand Up @@ -412,13 +434,14 @@ static void kernel_out_enzyme( const double* x, const void* kp_void, double* out

double projs[2] = { 0 };
get_projections( A0, A1, B0, B1, projs );
std::array<double, 2> projections = { projs[0], projs[1] };

// Recompute the integration bounds and quadrature from the current geometry.
auto bounds = ContactSmoothing::bounds_from_projections( projections, kp->del );
auto xi_bounds = ContactSmoothing::smooth_bounds( bounds, kp->del );

auto qp = EnergyMortarCalculator::compute_quadrature( xi_bounds, kp->N );
double bounds[2];
ContactSmoothing::bounds_from_projections( projs, kp->del, bounds );
double xi_bounds[2];
ContactSmoothing::smooth_bounds( bounds, kp->del, xi_bounds );
QuadPoints qp;
EnergyMortarCalculator::compute_quadrature( xi_bounds, kp->N, &qp );

Gparams gp;
for ( std::size_t i = 0; i < qp.qp.size(); ++i ) {
Expand Down Expand Up @@ -477,6 +500,25 @@ void d2_kernel( const double* x, const KernelParams* kp, double* H )
}
}

// Isolate loop-local arrays to avoid a leak in Enzyme's reverse-mode tape.
TRIBOL_ENZYME_INLINE double qp_penalty_kernel_qp_energy( double xiA, double w, const double* A0, const double* A1,
Comment thread
ebchin marked this conversation as resolved.
const double* B0, const double* B1, const double* nB,
double eta, double penalty, double J )
{
double x1[2];
iso_map( A0, A1, xiA, x1 );

double x2[2];
find_intersection( B0, B1, x1, nB, x2 );

const double dx = x1[0] - x2[0];
const double dy = x1[1] - x2[1];
const double gn = -( dx * nB[0] + dy * nB[1] );
const double gap = gn * eta;

return gap < 0.0 ? 0.5 * penalty * gap * gap * w * J : 0.0;
}

TRIBOL_ENZYME_INLINE void qp_penalty_kernel( const double* x, const KernelParams* kp, double* energy )
{
double A0[2] = { x[0], x[1] };
Expand All @@ -486,10 +528,12 @@ TRIBOL_ENZYME_INLINE void qp_penalty_kernel( const double* x, const KernelParams

double projs[2] = { 0.0, 0.0 };
get_projections( A0, A1, B0, B1, projs );
const std::array<double, 2> projections{ projs[0], projs[1] };
const auto bounds = ContactSmoothing::bounds_from_projections( projections, kp->del );
const auto xi_bounds = ContactSmoothing::smooth_bounds( bounds, kp->del );
const auto qp = EnergyMortarCalculator::compute_quadrature( xi_bounds, kp->N );
double bounds[2];
ContactSmoothing::bounds_from_projections( projs, kp->del, bounds );
double xi_bounds[2];
ContactSmoothing::smooth_bounds( bounds, kp->del, xi_bounds );
QuadPoints qp;
EnergyMortarCalculator::compute_quadrature( xi_bounds, kp->N, &qp );

double nB[2];
find_normal( B0, B1, nB );
Expand All @@ -504,19 +548,7 @@ TRIBOL_ENZYME_INLINE void qp_penalty_kernel( const double* x, const KernelParams

double value = 0.0;
for ( int i = 0; i < kp->N; ++i ) {
double x1[2];
iso_map( A0, A1, qp.qp[i], x1 );

double x2[2];
find_intersection( B0, B1, x1, nB, x2 );

const double dx = x1[0] - x2[0];
const double dy = x1[1] - x2[1];
const double gn = -( dx * nB[0] + dy * nB[1] );
const double gap = gn * eta;
if ( gap < 0.0 ) {
value += 0.5 * kp->k * gap * gap * qp.w[i] * J;
}
value += qp_penalty_kernel_qp_energy( qp.qp[i], qp.w[i], A0, A1, B0, B1, nB, eta, kp->k, J );
}

*energy = value;
Expand Down Expand Up @@ -582,10 +614,13 @@ Gparams EnergyMortarCalculator::construct_gparams( const InterfacePair& pair, co

// Build the smoothed integration bounds from the projection of edge B onto edge A.
auto projs = EnergyMortarCalculator::compute_projection_bounds( pair, mesh1, mesh2 );
auto bounds = smoother_.bounds_from_projections( projs, p_.del );
auto smooth_bounds = smoother_.smooth_bounds( bounds, p_.del );
double bounds[2];
smoother_.bounds_from_projections( projs.data(), p_.del, bounds );
double smooth_bounds[2];
smoother_.smooth_bounds( bounds, p_.del, smooth_bounds );

auto qp = EnergyMortarCalculator::compute_quadrature( smooth_bounds, p_.N );
QuadPoints qp;
EnergyMortarCalculator::compute_quadrature( smooth_bounds, p_.N, &qp );

const int N = static_cast<int>( qp.qp.size() );

Expand Down Expand Up @@ -631,11 +666,11 @@ std::array<double, 2> EnergyMortarCalculator::projections( const InterfacePair&
}

// Clamp the projection interval to the local smoothing support around edge A.
TRIBOL_ENZYME_INLINE std::array<double, 2> ContactSmoothing::bounds_from_projections( const std::array<double, 2>& proj,
double del )
TRIBOL_ENZYME_INLINE void ContactSmoothing::bounds_from_projections( const double* projections, double del,
double* bounds )
{
double xi_min = std::min( proj[0], proj[1] );
double xi_max = std::max( proj[0], proj[1] );
double xi_min = std::min( projections[0], projections[1] );
double xi_max = std::max( projections[0], projections[1] );

// Limit the integration interval to the extended range [-0.5 - del, 0.5 + del].
if ( xi_max < -0.5 - del ) {
Expand All @@ -651,56 +686,27 @@ TRIBOL_ENZYME_INLINE std::array<double, 2> ContactSmoothing::bounds_from_project
xi_max = 0.5 + del;
}

return { xi_min, xi_max };
bounds[0] = xi_min;
bounds[1] = xi_max;
}

// Smooth the integration bounds using a C1 ramp near the ends of edge A.
// Specific too the smoothing techniques in EnergyMortar. This smooths the
// Bounds of intergration by applying a quadratic ramping function near the ends of the paramteric
// space. The smooth region/length is defined by the input del. The returned 'bounds' is the new bounds
// of intergation that result after the quadratic ramping has been applied.
TRIBOL_ENZYME_INLINE std::array<double, 2> ContactSmoothing::smooth_bounds( const std::array<double, 2>& bounds,
double del )
TRIBOL_ENZYME_INLINE void ContactSmoothing::smooth_bounds( const double* bounds, double del, double* smooth_bounds )
{
std::array<double, 2> smooth_bounds;
for ( int i = 0; i < 2; ++i ) {
double xi = 0.0;
double xi_hat = 0.0;

// Shift from the local coordinate interval [-0.5, 0.5] to [0, 1].
xi = bounds[i] + 0.5;
if ( del == 0.0 ) {
xi_hat = xi;
} else {
// Apply quadratic ramps near the endpoints and leave the interior unchanged.
if ( 0.0 - del <= xi && xi <= del ) {
xi_hat = ( 1.0 / ( 4 * del ) ) * ( xi * xi ) + 0.5 * xi + del / 4.0;
} else if ( ( 1.0 - del ) <= xi && xi <= 1.0 + del ) {
double b = -1.0 / ( 4.0 * del );
double c = 0.5 + 1.0 / ( 2.0 * del );
double d = 1.0 - del + ( 1.0 / ( 4.0 * del ) ) * pow( 1.0 - del, 2 ) - 0.5 * ( 1.0 - del ) -
( 1.0 - del ) / ( 2.0 * del );

xi_hat = b * xi * xi + c * xi + d;
} else if ( del <= xi && xi <= ( 1.0 - del ) ) {
xi_hat = xi;
}
}
// Shift the smoothed coordinate back to [-0.5, 0.5].
smooth_bounds[i] = xi_hat - 0.5;
}

return smooth_bounds;
smooth_bounds[0] = smooth_bound( bounds[0], del );
smooth_bounds[1] = smooth_bound( bounds[1], del );
}

// Build a three-point Gauss-Legendre quadrature rule over the local integration bounds.
TRIBOL_ENZYME_INLINE QuadPoints EnergyMortarCalculator::compute_quadrature( const std::array<double, 2>& xi_bounds,
int N )
TRIBOL_ENZYME_INLINE void EnergyMortarCalculator::compute_quadrature( const double* xi_bounds, int N,
QuadPoints* quadrature )
{
QuadPoints out;

std::array<double, 3> qpoints;
std::array<double, 3> weights;
double qpoints[3] = { 0.0, 0.0, 0.0 };
double weights[3] = { 0.0, 0.0, 0.0 };

determine_legendre_nodes( N, qpoints );
determine_legendre_weights( N, weights );
Expand All @@ -711,11 +717,9 @@ TRIBOL_ENZYME_INLINE QuadPoints EnergyMortarCalculator::compute_quadrature( cons
const double J = 0.5 * ( xi_max - xi_min );

for ( int i = 0; i < N; ++i ) {
out.qp[i] = 0.5 * ( xi_max - xi_min ) * qpoints[i] + 0.5 * ( xi_max + xi_min );
out.w[i] = weights[i] * J;
quadrature->qp[i] = 0.5 * ( xi_max - xi_min ) * qpoints[i] + 0.5 * ( xi_max + xi_min );
quadrature->w[i] = weights[i] * J;
}

return out;
}

// Evaluate the weighted normal gap at local coordinate xiA on edge A.
Expand Down Expand Up @@ -764,10 +768,13 @@ NodalContactData EnergyMortarCalculator::compute_nodal_contact_data( const Inter
auto projs = projections( pair, mesh1, mesh2 );

// Build the smoothed integration interval from the projection bounds.
auto bounds = smoother_.bounds_from_projections( projs, p_.del );
auto smooth_bounds = smoother_.smooth_bounds( bounds, p_.del );
double bounds[2];
smoother_.bounds_from_projections( projs.data(), p_.del, bounds );
double smooth_bounds[2];
smoother_.smooth_bounds( bounds, p_.del, smooth_bounds );

auto qp = compute_quadrature( smooth_bounds, p_.N );
QuadPoints qp;
compute_quadrature( smooth_bounds, p_.N, &qp );

double g_tilde1 = 0.0;
double g_tilde2 = 0.0;
Expand Down
18 changes: 8 additions & 10 deletions src/tribol/physics/EnergyMortar.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,17 +91,15 @@ class ContactSmoothing {
public:
/// Clamp the projected overlap interval to the extended smoothing support.
///
/// The input `proj` contains the local projection bounds of edge B onto edge A.
/// The returned interval is restricted to the extended local range
/// `[-0.5 - del, 0.5 + del]`.
static std::array<double, 2> bounds_from_projections( const std::array<double, 2>& proj, double del );
/// The input `projections` contains the local projection bounds of edge B onto edge A.
/// The output `bounds` is restricted to the extended local range `[-0.5 - del, 0.5 + del]`.
static void bounds_from_projections( const double* projections, double del, double* bounds );

/// Smooth the integration bounds using the smoothing length `del`.
///
/// The returned bounds are obtained by applying the endpoint smoothing map to
/// the clamped integration interval. When `del = 0`, the bounds are returned
/// without smoothing.
static std::array<double, 2> smooth_bounds( const std::array<double, 2>& bounds, double del );
/// The output `smooth_bounds` is obtained by applying the endpoint smoothing map to
/// the clamped integration interval. When `del = 0`, the bounds are unchanged.
static void smooth_bounds( const double* bounds, double del, double* smooth_bounds );
};

/// Evaluates Energy Mortar contact quantities for a single interface pair.
Expand Down Expand Up @@ -132,8 +130,8 @@ class EnergyMortarCalculator {
///
/// The input bounds are local coordinates on edge A. The returned quadrature
/// points and weights are mapped from the reference interval to
/// `[xi_bounds[0], xi_bounds[1]]`
static QuadPoints compute_quadrature( const std::array<double, 2>& xi_bounds, int N );
/// `[xi_bounds[0], xi_bounds[1]]` and written to `quadrature`.
static void compute_quadrature( const double* xi_bounds, int N, QuadPoints* quadrature );

/// Compute the nodal smoothed gap integrals and tributary areas.
///
Expand Down
Loading