Skip to content

Commit

Permalink
BDF: fixing types and template parameters in batched calls
Browse files Browse the repository at this point in the history 
Bascially we need template parameters to be more versatile
and cannot assume that all rank1 views will have the exact
same underlying type, for instance layouts can be different.
  • Loading branch information
@lucbv
lucbv committed Oct 25, 2023
1 parent 280d108 commit b96def6
Show file tree
Hide file tree
Showing 7 changed files with 85 additions and 46 deletions.
36 changes: 22 additions & 14 deletions batched/dense/impl/KokkosBatched_Gesv_Impl.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -366,20 +366,24 @@ KOKKOS_INLINE_FUNCTION void TeamVectorHadamard1D(const MemberType &member,
/// ===========
template <>
struct SerialGesv<Gesv::StaticPivoting> {
template <typename MatrixType, typename VectorType>
template <typename MatrixType, typename XVectorType, typename YVectorType>
KOKKOS_INLINE_FUNCTION static int invoke(const MatrixType A,
const VectorType X,
const VectorType Y,
const XVectorType X,
const YVectorType Y,
const MatrixType tmp) {
#if (KOKKOSKERNELS_DEBUG_LEVEL > 0)
static_assert(Kokkos::is_view<MatrixType>::value,
"KokkosBatched::gesv: MatrixType is not a Kokkos::View.");
static_assert(Kokkos::is_view<VectorType>::value,
"KokkosBatched::gesv: VectorType is not a Kokkos::View.");
static_assert(Kokkos::is_view<XVectorType>::value,
"KokkosBatched::gesv: XVectorType is not a Kokkos::View.");
static_assert(Kokkos::is_view<YVectorType>::value,
"KokkosBatched::gesv: YVectorType is not a Kokkos::View.");
static_assert(MatrixType::rank == 2,
"KokkosBatched::gesv: MatrixType must have rank 2.");
static_assert(VectorType::rank == 1,
"KokkosBatched::gesv: VectorType must have rank 1.");
static_assert(XVectorType::rank == 1,
"KokkosBatched::gesv: XVectorType must have rank 1.");
static_assert(YVectorType::rank == 1,
"KokkosBatched::gesv: YVectorType must have rank 1.");

// Check compatibility of dimensions at run time.

Expand Down Expand Up @@ -462,20 +466,24 @@ struct SerialGesv<Gesv::StaticPivoting> {

template <>
struct SerialGesv<Gesv::NoPivoting> {
template <typename MatrixType, typename VectorType>
template <typename MatrixType, typename XVectorType, typename YVectorType>
KOKKOS_INLINE_FUNCTION static int invoke(const MatrixType A,
const VectorType X,
const VectorType Y,
const XVectorType X,
const YVectorType Y,
const MatrixType /*tmp*/) {
#if (KOKKOSKERNELS_DEBUG_LEVEL > 0)
static_assert(Kokkos::is_view<MatrixType>::value,
"KokkosBatched::gesv: MatrixType is not a Kokkos::View.");
static_assert(Kokkos::is_view<VectorType>::value,
"KokkosBatched::gesv: VectorType is not a Kokkos::View.");
static_assert(Kokkos::is_view<XVectorType>::value,
"KokkosBatched::gesv: XVectorType is not a Kokkos::View.");
static_assert(Kokkos::is_view<YVectorType>::value,
"KokkosBatched::gesv: YVectorType is not a Kokkos::View.");
static_assert(MatrixType::rank == 2,
"KokkosBatched::gesv: MatrixType must have rank 2.");
static_assert(VectorType::rank == 1,
"KokkosBatched::gesv: VectorType must have rank 1.");
static_assert(XVectorType::rank == 1,
"KokkosBatched::gesv: XVectorType must have rank 1.");
static_assert(YVectorType::rank == 1,
"KokkosBatched::gesv: YVectorType must have rank 1.");

// Check compatibility of dimensions at run time.

Expand Down
11 changes: 10 additions & 1 deletion batched/dense/src/KokkosBatched_Gesv.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,11 +63,20 @@ struct Gesv {

template <typename ArgAlgo>
struct SerialGesv {
template <typename MatrixType, typename XVectorType, typename YVectorType>
KOKKOS_INLINE_FUNCTION static int invoke(const MatrixType A,
const XVectorType X,
const YVectorType Y,
const MatrixType tmp);

template <typename MatrixType, typename VectorType>
[[deprecated]]
KOKKOS_INLINE_FUNCTION static int invoke(const MatrixType A,
const VectorType X,
const VectorType Y,
const MatrixType tmp);
const MatrixType tmp) {
return invoke(A, X, Y, tmp);
}
};

/// \brief Team Batched GESV:
Expand Down
23 changes: 11 additions & 12 deletions ode/impl/KokkosODE_BDF_impl.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -141,8 +141,8 @@ struct BDF_system_wrapper2 {
t(t_),
dt(dt_) {}

template <class vec_type>
KOKKOS_FUNCTION void residual(const vec_type& y, const vec_type& f) const {
template <class YVectorType, class FVectorType>
KOKKOS_FUNCTION void residual(const YVectorType& y, const FVectorType& f) const {
// f = f(t+dt, y)
mySys.evaluate_function(t, dt, y, f);

Expand Down Expand Up @@ -228,21 +228,21 @@ KOKKOS_FUNCTION void update_D(const int order, const scalar_type factor,
KokkosBatched::Algo::Gemm::Blocked>::invoke(1.0, subTempD, U, 0.0, subD);
}

template <class ode_type, class mat_type, class vec_type, class scalar_type>
template <class ode_type, class mat_type, class vec_type, class res_type, class scalar_type>
KOKKOS_FUNCTION void initial_step_size(const ode_type ode, const int order, const scalar_type t0,
const scalar_type atol, const scalar_type rtol,
const vec_type& y0, const vec_type& f0,
const vec_type& y0, const res_type& f0,
const mat_type& temp, scalar_type& dt_ini) {
using KAT = Kokkos::ArithTraits<scalar_type>;

std::cout << "Estimating initial step size" << std::endl;
Kokkos::printf("Estimating initial step size\n");

// Extract subviews to store intermediate data
auto scale = Kokkos::subview(temp, Kokkos::ALL(), 0); // Scaling coefficients for error calculation
auto y1 = Kokkos::subview(temp, Kokkos::ALL(), 1); // Scaling coefficients for error calculation
auto f1 = Kokkos::subview(temp, Kokkos::ALL(), 2); // Scaling coefficients for error calculation

std::cout << "scale: rank=" << scale.rank() << ", extent(0)=" << scale.extent(0) << std::endl;
Kokkos::printf("scale: rank=%d, extent(0)=%d\n", (int)scale.rank(), scale.extent_int(0));

// Compute norms for y0 and f0
double n0 = KAT::zero(), n1 = KAT::zero(), dt0;
Expand Down Expand Up @@ -285,15 +285,15 @@ KOKKOS_FUNCTION void initial_step_size(const ode_type ode, const int order, cons
dt_ini = Kokkos::min(100 * dt0, dt_ini);
} // initial_step_size

template <class ode_type, class vec_type,
template <class ode_type, class vec_type, class res_type,
class mat_type, class scalar_type>
KOKKOS_FUNCTION void BDFStep(ode_type& ode, scalar_type& t, scalar_type& dt,
scalar_type t_end, int& order, int& num_equal_steps,
const int max_newton_iters,
const scalar_type atol, const scalar_type rtol,
const scalar_type min_factor,
const vec_type& y_old, const vec_type& y_new,
const vec_type& rhs, const vec_type& update,
const res_type& rhs, const res_type& update,
const mat_type& temp, const mat_type& temp2) {
using newton_params = KokkosODE::Experimental::Newton_params;

Expand All @@ -312,14 +312,13 @@ KOKKOS_FUNCTION void BDFStep(ode_type& ode, scalar_type& t, scalar_type& dt,
// const int numRows = temp.extent_int(0); const int numCols = temp.extent_int(1);
// std::cout << "numRows: " << numRows << ", numCols: " << numCols << std::endl;
// std::cout << "Extract subview from temp" << std::endl;
int offset = 0;
int offset = 2;
auto D = Kokkos::subview(temp, Kokkos::ALL(), Kokkos::pair<int, int>(offset, offset + 8)); // y and its derivatives
offset += 8;
auto tempD = Kokkos::subview(temp, Kokkos::ALL(), Kokkos::pair<int, int>(offset, offset + 8));
offset += 8;
auto scale = Kokkos::subview(temp, Kokkos::ALL(), offset + 1); ++offset; // Scaling coefficients for error calculation
auto y_predict = Kokkos::subview(temp, Kokkos::ALL(), offset + 1); ++offset; // Initial guess for y_{n+1}
// auto d = Kokkos::subview(temp, Kokkos::ALL(), offset + 1); ++offset; // Newton update: y_{n+1}^i - y_n
auto psi = Kokkos::subview(temp, Kokkos::ALL(), offset + 1); ++offset; // Higher order terms contribution to rhs
auto error = Kokkos::subview(temp, Kokkos::ALL(), offset + 1); ++offset; // Error estimate
auto jac = Kokkos::subview(temp, Kokkos::ALL(), Kokkos::pair<int, int>(offset, offset + ode.neqs)); // Jacobian matrix
Expand Down Expand Up @@ -383,8 +382,8 @@ KOKKOS_FUNCTION void BDFStep(ode_type& ode, scalar_type& t, scalar_type& dt,

sys.c = dt / alpha[order];
sys.jacobian(y_new, jac);
Kokkos::deep_copy(y_new, y_predict);
Kokkos::deep_copy(update, 0);
Kokkos::Experimental::local_deep_copy(y_new, y_predict);
Kokkos::Experimental::local_deep_copy(update, 0);
KokkosODE::Experimental::newton_solver_status newton_status =
KokkosODE::Experimental::Newton::Solve(sys, param, jac, tmp_gesv, y_new, rhs,
update);
Expand Down
8 changes: 4 additions & 4 deletions ode/impl/KokkosODE_Newton_impl.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,17 +30,17 @@
namespace KokkosODE {
namespace Impl {

template <class system_type, class mat_type, class vec_type>
template <class system_type, class mat_type, class ini_vec_type, class rhs_vec_type, class update_type>
KOKKOS_FUNCTION KokkosODE::Experimental::newton_solver_status NewtonSolve(
system_type& sys, const KokkosODE::Experimental::Newton_params& params,
mat_type& J, mat_type& tmp, vec_type& y0, vec_type& rhs, vec_type& update) {
mat_type& J, mat_type& tmp, ini_vec_type& y0, rhs_vec_type& rhs, update_type& update) {
using newton_solver_status = KokkosODE::Experimental::newton_solver_status;
using value_type = typename vec_type::non_const_value_type;
using value_type = typename ini_vec_type::non_const_value_type;

// Define the type returned by nrm2 to store
// the norm of the residual.
using norm_type = typename Kokkos::Details::InnerProductSpaceTraits<
typename vec_type::non_const_value_type>::mag_type;
typename ini_vec_type::non_const_value_type>::mag_type;
sys.residual(y0, rhs);
// std::cout << "y0=" << y0(0) << std::endl;
// std::cout << "rhs0= " << rhs(0) << std::endl;
Expand Down
20 changes: 18 additions & 2 deletions ode/src/KokkosODE_BDF.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,36 +41,50 @@ enum BDF_type : int {
template <BDF_type T>
struct BDF_coeff_helper {
using table_type = void;

BDF_coeff_helper() = default;
};

template <>
struct BDF_coeff_helper<BDF_type::BDF1> {
using table_type = KokkosODE::Impl::BDF_table<1>;

BDF_coeff_helper() = default;
};

template <>
struct BDF_coeff_helper<BDF_type::BDF2> {
using table_type = KokkosODE::Impl::BDF_table<2>;

BDF_coeff_helper() = default;
};

template <>
struct BDF_coeff_helper<BDF_type::BDF3> {
using table_type = KokkosODE::Impl::BDF_table<3>;

BDF_coeff_helper() = default;
};

template <>
struct BDF_coeff_helper<BDF_type::BDF4> {
using table_type = KokkosODE::Impl::BDF_table<4>;

BDF_coeff_helper() = default;
};

template <>
struct BDF_coeff_helper<BDF_type::BDF5> {
using table_type = KokkosODE::Impl::BDF_table<5>;

BDF_coeff_helper() = default;
};

template <>
struct BDF_coeff_helper<BDF_type::BDF6> {
using table_type = KokkosODE::Impl::BDF_table<6>;

BDF_coeff_helper() = default;
};

template <BDF_type T>
Expand All @@ -84,7 +98,7 @@ struct BDF {
const int num_steps, const vec_type& y0, const vec_type& y,
const vec_type& rhs, const vec_type& update, const mv_type& y_vecs,
const mv_type& kstack, const mat_type& temp, const mat_type& jac) {
const table_type table;
const table_type table{};

const double dt = (t_end - t_start) / num_steps;
double t = t_start;
Expand Down Expand Up @@ -159,7 +173,9 @@ KOKKOS_FUNCTION void BDFSolve(const ode_type& ode, const scalar_type t_start, co
using KAT = Kokkos::ArithTraits<scalar_type>;

// This needs to go away and be pulled out of temp instead...
vec_type rhs("rhs", ode.neqs), update("update", ode.neqs);
auto rhs = Kokkos::subview(temp, Kokkos::ALL(), 0);
auto update = Kokkos::subview(temp, Kokkos::ALL(), 1);
// vec_type rhs("rhs", ode.neqs), update("update", ode.neqs);
(void) max_step;

int order = 1, num_equal_steps = 0;
Expand Down
7 changes: 4 additions & 3 deletions ode/src/KokkosODE_Newton.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,11 +30,12 @@ namespace Experimental {

/// \brief Newton solver for non-linear system of equations
struct Newton {
template <class system_type, class mat_type, class vec_type>
template <class system_type, class mat_type, class ini_vec_type, class rhs_vec_type,
class update_type>
KOKKOS_FUNCTION static newton_solver_status Solve(
const system_type& sys, const Newton_params& params, const mat_type& J,
const mat_type& tmp, const vec_type& y0, const vec_type& rhs,
const vec_type& update) {
const mat_type& tmp, const ini_vec_type& y0, const rhs_vec_type& rhs,
const update_type& update) {
return KokkosODE::Impl::NewtonSolve(sys, params, J, tmp, y0, rhs, update);
}
};
Expand Down
26 changes: 16 additions & 10 deletions ode/unit_test/Test_ODE_BDF.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,8 +174,9 @@ struct BDFSolve_wrapper {
}
};

template <class execution_space, class scalar_type>
template <class device_type, class scalar_type>
void test_BDF_Logistic() {
using execution_space = typename device_type::execution_space;
using vec_type = Kokkos::View<scalar_type*, execution_space>;
using mv_type = Kokkos::View<scalar_type**, execution_space>;
using mat_type = Kokkos::View<scalar_type**, execution_space>;
Expand Down Expand Up @@ -347,8 +348,9 @@ void test_BDF_Logistic() {

} // test_BDF_Logistic

template <class execution_space, class scalar_type>
template <class device_type, class scalar_type>
void test_BDF_LotkaVolterra() {
using execution_space = typename device_type::execution_space;
using vec_type = Kokkos::View<scalar_type*, execution_space>;
using mv_type = Kokkos::View<scalar_type**, execution_space>;
using mat_type = Kokkos::View<scalar_type**, execution_space>;
Expand Down Expand Up @@ -376,8 +378,9 @@ void test_BDF_LotkaVolterra() {
Kokkos::parallel_for(myPolicy, solve_wrapper);
}

template <class execution_space, class scalar_type>
template <class device_type, class scalar_type>
void test_BDF_StiffChemistry() {
using execution_space = typename device_type::execution_space;
using vec_type = Kokkos::View<scalar_type*, execution_space>;
using mv_type = Kokkos::View<scalar_type**, execution_space>;
using mat_type = Kokkos::View<scalar_type**, execution_space>;
Expand Down Expand Up @@ -474,8 +477,9 @@ struct BDFSolve_parallel {
}
};

template <class execution_space, class scalar_type>
template <class device_type, class scalar_type>
void test_BDF_parallel() {
using execution_space = typename device_type::execution_space;
using vec_type = Kokkos::View<scalar_type*, execution_space>;
using mv_type = Kokkos::View<scalar_type**, execution_space>;
using mat_type = Kokkos::View<scalar_type**, execution_space>;
Expand Down Expand Up @@ -537,8 +541,9 @@ void update_D(const int order, const scalar_type factor, const mat_type& coeffs,
KokkosBatched::Algo::Gemm::Blocked>::invoke(1.0, U, subTempD, 0.0, subD);
}

template <class execution_space, class scalar_type>
template <class device_type, class scalar_type>
void test_Nordsieck() {
using execution_space = typename device_type::execution_space;
StiffChemistry mySys{};

Kokkos::View<scalar_type**, execution_space> R("coeffs", 6, 6), U("coeffs", 6, 6);
Expand Down Expand Up @@ -580,8 +585,9 @@ void test_Nordsieck() {
std::cout << " { " << D(1, 0) << ", " << D(1, 1) << ", " << D(1, 2) << " }" << std::endl;
}

template <class execution_space, class scalar_type>
template <class device_type, class scalar_type>
void test_adaptive_BDF() {
using execution_space = typename device_type::execution_space;
using vec_type = Kokkos::View<scalar_type*, execution_space>;
using mat_type = Kokkos::View<scalar_type**, execution_space>;

Expand All @@ -595,13 +601,13 @@ void test_adaptive_BDF() {

vec_type y0("initial conditions", mySys.neqs), y_new("solution", mySys.neqs);
vec_type rhs("rhs", mySys.neqs), update("update", mySys.neqs);
mat_type temp("buffer1", mySys.neqs, 21 + 2*mySys.neqs + 4), temp2("buffer2", 6, 7);
mat_type temp("buffer1", mySys.neqs, 23 + 2*mySys.neqs + 4), temp2("buffer2", 6, 7);

// Initial condition
Kokkos::deep_copy(y0, 0.5);

// Initialize D
auto D = Kokkos::subview(temp, Kokkos::ALL(), Kokkos::pair<int, int>(0, 8));
auto D = Kokkos::subview(temp, Kokkos::ALL(), Kokkos::pair<int, int>(2, 10));
D(0, 0) = y0(0);
mySys.evaluate_function(0, 0, y0, rhs);
D(0, 1) = dt*rhs(0);
Expand Down Expand Up @@ -674,7 +680,7 @@ void test_adaptive_BDF_v2() {
vec_type y0("initial conditions", mySys.neqs), y_new("solution", mySys.neqs);
Kokkos::deep_copy(y0, 0.5);

mat_type temp("buffer1", mySys.neqs, 21 + 2*mySys.neqs + 4), temp2("buffer2", 6, 7);
mat_type temp("buffer1", mySys.neqs, 23 + 2*mySys.neqs + 4), temp2("buffer2", 6, 7);

{
scalar_type dt = KAT::zero();
Expand Down Expand Up @@ -707,7 +713,7 @@ void test_BDF_adaptive_stiff() {
y0_h(2) = KAT::zero();
Kokkos::deep_copy(y0, y0_h);

mat_type temp("buffer1", mySys.neqs, 21 + 2*mySys.neqs + 4), temp2("buffer2", 6, 7);
mat_type temp("buffer1", mySys.neqs, 23 + 2*mySys.neqs + 4), temp2("buffer2", 6, 7);

{
vec_type f0("initial value f", mySys.neqs);
Expand Down

0 comments on commit b96def6

Please sign in to comment.