Skip to content

add support for reduced precision Jacobian #1862

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
merged 8 commits into from
Oct 30, 2025
Merged

add support for reduced precision Jacobian #1862

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
78487a4
start of support for reduced precision Jacobian
zingale Oct 19, 2025
13d0b7b
fix some types
zingale Oct 19, 2025
4c2e7d8
add control
zingale Oct 19, 2025
cb600d6
update defines
zingale Oct 19, 2025
f16ee32
add docs
zingale Oct 20, 2025
678b067
Merge branch 'development' into single_precision_jac
zingale Oct 20, 2025
b69058d
Merge branch 'development' into single_precision_jac
zingale Oct 27, 2025
8f141f8
Merge branch 'development' into single_precision_jac
zingale Oct 28, 2025
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
1 change: 1 addition & 0 deletions .github/workflows/good_defines.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ SCREENING
SCREEN_METHOD
SDC
SIMPLIFIED_SDC
SINGLE_PRECISION_JACOBIAN
STRANG
TRUE_SDC
_OPENMP
Expand Down
17 changes: 17 additions & 0 deletions Docs/source/ode_integrators.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,23 @@ For the other networks (usually pynucastro networks), the implementation is
provided in ``Microphysics/util/linpack.H`` and is templated on the number
of equations. Pivoting can be disabled by setting ``integrator.linalg_do_pivoting=0``.

.. index:: USE_SINGLE_PRECISION_JACOBIAN

.. tip::

The storage for the Jacobian can take up the most memory when
integrating the reaction system. It is possible to store the
Jacobian as single-precision, by building with:

::

USE_SINGLE_PRECISION_JACOBIAN=TRUE

This can speed up the integration prevent the code from running out
of memory when run on GPUs.



Integration errors
==================

Expand Down
2 changes: 1 addition & 1 deletion integration/BackwardEuler/be_type.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ struct be_t {
amrex::Real rtol_enuc;

amrex::Array1D<amrex::Real, 1, int_neqs> y;
ArrayUtil::MathArray2D<1, int_neqs, 1, int_neqs> jac;
ArrayUtil::MathArray2D<jac_t, 1, int_neqs, 1, int_neqs> jac;

short jacobian_type;
};
Expand Down
3 changes: 3 additions & 0 deletions integration/Make.package
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,9 @@ else
CEXE_headers += integrator_setup_strang.H
endif

ifeq ($(USE_SINGLE_PRECISION_JACOBIAN), TRUE)
DEFINES += -DSINGLE_PRECISION_JACOBIAN
endif

ifeq ($(USE_ALL_NSE), TRUE)
ifeq ($(USE_ALL_SDC), TRUE)
Expand Down
4 changes: 2 additions & 2 deletions integration/VODE/vode_type.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,11 +143,11 @@ struct dvode_t
amrex::Array1D<amrex::Real, 1, int_neqs> y;

// Jacobian
ArrayUtil::MathArray2D<1, int_neqs, 1, int_neqs> jac;
ArrayUtil::MathArray2D<jac_t, 1, int_neqs, 1, int_neqs> jac;

#ifdef ALLOW_JACOBIAN_CACHING
// Saved Jacobian
ArrayUtil::MathArray2D<1, int_neqs, 1, int_neqs> jac_save;
ArrayUtil::MathArray2D<jac_t, 1, int_neqs, 1, int_neqs> jac_save;
#endif

// the Nordsieck history array
Expand Down
2 changes: 1 addition & 1 deletion integration/integrator_data.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,6 @@ constexpr int integrator_neqs ()

using IArray1D = amrex::Array1D<short, 1, INT_NEQS>;
using RArray1D = amrex::Array1D<amrex::Real, 1, INT_NEQS>;
using RArray2D = ArrayUtil::MathArray2D<1, INT_NEQS, 1, INT_NEQS>;
using RArray2D = ArrayUtil::MathArray2D<jac_t, 1, INT_NEQS, 1, INT_NEQS>;

#endif
2 changes: 1 addition & 1 deletion integration/utils/circle_theorem.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
void circle_theorem_sprad(const amrex::Real time, BurnT& state, T& int_state, amrex::Real& sprad)
{

ArrayUtil::MathArray2D<1, INT_NEQS, 1, INT_NEQS> jac_array;
ArrayUtil::MathArray2D<jac_t, 1, INT_NEQS, 1, INT_NEQS> jac_array;

if (integrator_rp::jacobian == 1) {
jac(time, state, int_state, jac_array);
Expand Down
10 changes: 5 additions & 5 deletions interfaces/ArrayUtilities.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ namespace ArrayUtil
amrex::Real arr[(XHI-XLO+1)];
};

template <int XLO, int XHI, int YLO, int YHI>
template <typename T, int XLO, int XHI, int YLO, int YHI>
struct MathArray2D
{
AMREX_GPU_HOST_DEVICE AMREX_INLINE
Expand Down Expand Up @@ -71,7 +71,7 @@ namespace ArrayUtil
}

[[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
amrex::Real get (const int i, const int j) const noexcept {
T get (const int i, const int j) const noexcept {
AMREX_ASSERT(i >= XLO && i <= XHI && j >= YLO && j <= YHI);
return arr[i+j*(XHI-XLO+1)-(YLO*(XHI-XLO+1)+XLO)];
}
Expand All @@ -84,18 +84,18 @@ namespace ArrayUtil
}

AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
const amrex::Real& operator() (int i, int j) const noexcept {
const T& operator() (int i, int j) const noexcept {
AMREX_ASSERT(i >= XLO && i <= XHI && j >= YLO && j <= YHI);
return arr[i+j*(XHI-XLO+1)-(YLO*(XHI-XLO+1)+XLO)];
}

AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
amrex::Real& operator() (int i, int j) noexcept {
T& operator() (int i, int j) noexcept {
AMREX_ASSERT(i >= XLO && i <= XHI && j >= YLO && j <= YHI);
return arr[i+j*(XHI-XLO+1)-(YLO*(XHI-XLO+1)+XLO)];
}

amrex::Real arr[(XHI-XLO+1)*(YHI-YLO+1)];
T arr[(XHI-XLO+1)*(YHI-YLO+1)];
};

namespace Math
Expand Down
10 changes: 9 additions & 1 deletion interfaces/burn_type.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,15 @@ const int SVAR = SEDEN+1;
const int SVAR_EVOLVE = SFX;

// this is the data type of the dense Jacobian that the network wants.
using JacNetArray2D = ArrayUtil::MathArray2D<1, neqs, 1, neqs>;

// the Jacobian can have a different type that our system
#ifdef SINGLE_PRECISION_JACOBIAN
using jac_t = float;
#else
using jac_t = amrex::Real;
#endif

using JacNetArray2D = ArrayUtil::MathArray2D<jac_t, 1, neqs, 1, neqs>;

struct burn_t
{
Expand Down
4 changes: 2 additions & 2 deletions networks/rhs.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1357,7 +1357,7 @@ void rhs (const burn_t& burn_state, amrex::Array1D<amrex::Real, 1, nrhs>& ydot)

// Analytical Jacobian
AMREX_GPU_HOST_DEVICE AMREX_INLINE
void jac (const burn_t& burn_state, ArrayUtil::MathArray2D<1, neqs, 1, neqs>& jac)
void jac (const burn_t& burn_state, ArrayUtil::MathArray2D<jac_t, 1, neqs, 1, neqs>& jac)
{
rhs_state_t<autodiff::dual> rhs_state;

Expand Down Expand Up @@ -1521,7 +1521,7 @@ void actual_rhs (const burn_t& state, amrex::Array1D<amrex::Real, 1, neqs>& ydot
}

AMREX_GPU_HOST_DEVICE AMREX_INLINE
void actual_jac (const burn_t& state, ArrayUtil::MathArray2D<1, neqs, 1, neqs>& jac)
void actual_jac (const burn_t& state, ArrayUtil::MathArray2D<jac_t, 1, neqs, 1, neqs>& jac)
{
RHS::jac(state, jac);
}
Expand Down
2 changes: 1 addition & 1 deletion unit_test/test_linear_algebra/test_linear_algebra.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,7 @@ void linear_algebra() {

// now try to output a Jacobian mask based on the actual Jacobian

ArrayUtil::MathArray2D<1, INT_NEQS, 1, INT_NEQS> jac;
ArrayUtil::MathArray2D<jac_t, 1, INT_NEQS, 1, INT_NEQS> jac;

burn_t burn_state;
burn_state.rho = 1.e8;
Expand Down
2 changes: 1 addition & 1 deletion unit_test/test_rhs/rhs_zones.H
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ bool do_rhs (int i, int j, int k, amrex::Array4<amrex::Real> const& state, const

amrex::Array1D<amrex::Real, 1, neqs> ydot;

ArrayUtil::MathArray2D<1, neqs, 1, neqs> jac;
ArrayUtil::MathArray2D<jac_t, 1, neqs, 1, neqs> jac;

#ifdef NEW_NETWORK_IMPLEMENTATION
RHS::rhs(burn_state, ydot);
Expand Down
Loading