Skip to content

Commit 98f7226

Browse files
zingalezingale
authored
add support for reduced precision Jacobian (#1862)
1 parent 1a162ea commit 98f7226

12 files changed

Lines changed: 44 additions & 15 deletions

File tree

‎.github/workflows/good_defines.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ SCREENING
2424
SCREEN_METHOD
2525
SDC
2626
SIMPLIFIED_SDC
27+
SINGLE_PRECISION_JACOBIAN
2728
STRANG
2829
TRUE_SDC
2930
_OPENMP

‎Docs/source/ode_integrators.rst‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,23 @@ For the other networks (usually pynucastro networks), the implementation is
113113
provided in ``Microphysics/util/linpack.H`` and is templated on the number
114114
of equations. Pivoting can be disabled by setting ``integrator.linalg_do_pivoting=0``.
115115

116+
.. index:: USE_SINGLE_PRECISION_JACOBIAN
117+
118+
.. tip::
119+
120+
The storage for the Jacobian can take up the most memory when
121+
integrating the reaction system. It is possible to store the
122+
Jacobian as single-precision, by building with:
123+
124+
::
125+
126+
USE_SINGLE_PRECISION_JACOBIAN=TRUE
127+
128+
This can speed up the integration prevent the code from running out
129+
of memory when run on GPUs.
130+
131+
132+
116133
Integration errors
117134
==================
118135

‎integration/BackwardEuler/be_type.H‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ struct be_t {
4848
amrex::Real rtol_enuc;
4949

5050
amrex::Array1D<amrex::Real, 1, int_neqs> y;
51-
ArrayUtil::MathArray2D<1, int_neqs, 1, int_neqs> jac;
51+
ArrayUtil::MathArray2D<jac_t, 1, int_neqs, 1, int_neqs> jac;
5252

5353
short jacobian_type;
5454
};

‎integration/Make.package‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@ else
2626
CEXE_headers += integrator_setup_strang.H
2727
endif
2828

29+
ifeq ($(USE_SINGLE_PRECISION_JACOBIAN), TRUE)
30+
DEFINES += -DSINGLE_PRECISION_JACOBIAN
31+
endif
2932

3033
ifeq ($(USE_ALL_NSE), TRUE)
3134
ifeq ($(USE_ALL_SDC), TRUE)

‎integration/VODE/vode_type.H‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -143,11 +143,11 @@ struct dvode_t
143143
amrex::Array1D<amrex::Real, 1, int_neqs> y;
144144

145145
// Jacobian
146-
ArrayUtil::MathArray2D<1, int_neqs, 1, int_neqs> jac;
146+
ArrayUtil::MathArray2D<jac_t, 1, int_neqs, 1, int_neqs> jac;
147147

148148
#ifdef ALLOW_JACOBIAN_CACHING
149149
// Saved Jacobian
150-
ArrayUtil::MathArray2D<1, int_neqs, 1, int_neqs> jac_save;
150+
ArrayUtil::MathArray2D<jac_t, 1, int_neqs, 1, int_neqs> jac_save;
151151
#endif
152152

153153
// the Nordsieck history array

‎integration/integrator_data.H‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,6 @@ constexpr int integrator_neqs ()
5252

5353
using IArray1D = amrex::Array1D<short, 1, INT_NEQS>;
5454
using RArray1D = amrex::Array1D<amrex::Real, 1, INT_NEQS>;
55-
using RArray2D = ArrayUtil::MathArray2D<1, INT_NEQS, 1, INT_NEQS>;
55+
using RArray2D = ArrayUtil::MathArray2D<jac_t, 1, INT_NEQS, 1, INT_NEQS>;
5656

5757
#endif

‎integration/utils/circle_theorem.H‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
1717
void circle_theorem_sprad(const amrex::Real time, BurnT& state, T& int_state, amrex::Real& sprad)
1818
{
1919

20-
ArrayUtil::MathArray2D<1, INT_NEQS, 1, INT_NEQS> jac_array;
20+
ArrayUtil::MathArray2D<jac_t, 1, INT_NEQS, 1, INT_NEQS> jac_array;
2121

2222
if (integrator_rp::jacobian == 1) {
2323
jac(time, state, int_state, jac_array);

‎interfaces/ArrayUtilities.H‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ namespace ArrayUtil
3434
amrex::Real arr[(XHI-XLO+1)];
3535
};
3636

37-
template <int XLO, int XHI, int YLO, int YHI>
37+
template <typename T, int XLO, int XHI, int YLO, int YHI>
3838
struct MathArray2D
3939
{
4040
AMREX_GPU_HOST_DEVICE AMREX_INLINE
@@ -71,7 +71,7 @@ namespace ArrayUtil
7171
}
7272

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

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

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

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

101101
namespace Math

‎interfaces/burn_type.H‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,15 @@ const int SVAR = SEDEN+1;
5656
const int SVAR_EVOLVE = SFX;
5757

5858
// this is the data type of the dense Jacobian that the network wants.
59-
using JacNetArray2D = ArrayUtil::MathArray2D<1, neqs, 1, neqs>;
59+
60+
// the Jacobian can have a different type that our system
61+
#ifdef SINGLE_PRECISION_JACOBIAN
62+
using jac_t = float;
63+
#else
64+
using jac_t = amrex::Real;
65+
#endif
66+
67+
using JacNetArray2D = ArrayUtil::MathArray2D<jac_t, 1, neqs, 1, neqs>;
6068

6169
struct burn_t
6270
{

‎networks/rhs.H‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1357,7 +1357,7 @@ void rhs (const burn_t& burn_state, amrex::Array1D<amrex::Real, 1, nrhs>& ydot)
13571357

13581358
// Analytical Jacobian
13591359
AMREX_GPU_HOST_DEVICE AMREX_INLINE
1360-
void jac (const burn_t& burn_state, ArrayUtil::MathArray2D<1, neqs, 1, neqs>& jac)
1360+
void jac (const burn_t& burn_state, ArrayUtil::MathArray2D<jac_t, 1, neqs, 1, neqs>& jac)
13611361
{
13621362
rhs_state_t<autodiff::dual> rhs_state;
13631363

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

15231523
AMREX_GPU_HOST_DEVICE AMREX_INLINE
1524-
void actual_jac (const burn_t& state, ArrayUtil::MathArray2D<1, neqs, 1, neqs>& jac)
1524+
void actual_jac (const burn_t& state, ArrayUtil::MathArray2D<jac_t, 1, neqs, 1, neqs>& jac)
15251525
{
15261526
RHS::jac(state, jac);
15271527
}

0 commit comments

Comments
 (0)