Skip to content

Commit

Permalink
Vectorizable matrix class (#87)
Browse files Browse the repository at this point in the history 
* add vector matrix class

* add matrix policy tests

* add initialization to matrix policy tests

* use matrix policy in State

* use matrix policy in ProcessSet

* use matrix policy in Rosenbrock solver

* use VectorMatrix in ProcessSet tests

* specialize vector matrix forcing function

* add VectorSize call to Vectorizable concept

* update CMake settings

* remove asserts from VectorMatrix

* use MatrixPolicy as template parameter

* patch for lack of NVidia c++20 support on gust

* fix issues related to nvhpc compilers

* remove c++20 flags

* remove last c++20 flag

---------

Co-authored-by: Matthew Dawson <mattdawson@gust02.hsn.gu.hpc.ucar.edu>
  • Loading branch information
@mattldawson
mattldawson and Matthew Dawson committed Jun 26, 2023
1 parent 840ef97 commit 6b5d220
Show file tree
Hide file tree
Showing 22 changed files with 763 additions and 271 deletions.
6 changes: 4 additions & 2 deletions include/micm/process/process.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,8 @@ namespace micm
/// @brief Update the solver state rate constants
/// @param processes The set of processes being solved
/// @param state The solver state to update
static void UpdateState(const std::vector<Process>& processes, State& state);
template<template<class> class MatrixPolicy>
static void UpdateState(const std::vector<Process>& processes, State<MatrixPolicy>& state);

friend class ProcessBuilder;
static ProcessBuilder create();
Expand Down Expand Up @@ -68,7 +69,8 @@ namespace micm
ProcessBuilder& phase(const Phase& phase);
};

void Process::UpdateState(const std::vector<Process>& processes, State& state)
template<template<class> class MatrixPolicy>
void Process::UpdateState(const std::vector<Process>& processes, State<MatrixPolicy>& state)
{
for (std::size_t i{}; i < state.custom_rate_parameters_.size(); ++i)
{
Expand Down
85 changes: 69 additions & 16 deletions include/micm/process/process_set.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,20 @@
#include <micm/solver/state.hpp>
#include <micm/util/matrix.hpp>
#include <micm/util/sparse_matrix.hpp>
#include <micm/util/vector_matrix.hpp>
#include <vector>

namespace micm
{

/// Concept for vectorizable matrices
template<typename T>
concept Vectorizable = requires(T t) {
t.BlockSize();
t.NumberOfBlocks();
t.VectorSize();
};

/// @brief Solver function calculators for a collection of processes
class ProcessSet
{
Expand All @@ -29,7 +38,8 @@ namespace micm
/// @brief Create a process set calculator for a given set of processes
/// @param processes Processes to create calculator for
/// @param state Solver state
ProcessSet(const std::vector<Process>& processes, const State& state);
template<template<class> class MatrixPolicy>
ProcessSet(const std::vector<Process>& processes, const State<MatrixPolicy>& state);

/// @brief Return the full set of non-zero Jacobian elements for the set of processes
/// @return Jacobian elements as a set of index pairs
Expand All @@ -43,22 +53,25 @@ namespace micm
/// @param rate_constants Current values for the process rate constants (grid cell, process)
/// @param state_variables Current state variable values (grid cell, state variable)
/// @param forcing Forcing terms for each state variable (grid cell, state variable)
void AddForcingTerms(
const Matrix<double>& rate_constants,
const Matrix<double>& state_variables,
Matrix<double>& forcing) const;
template<template<class> typename MatrixPolicy>
requires(!Vectorizable<MatrixPolicy<double>>)
void AddForcingTerms(const MatrixPolicy<double>& rate_constants, const MatrixPolicy<double>& state_variables, MatrixPolicy<double>& forcing) const;
template<template<class> typename MatrixPolicy>
requires Vectorizable<MatrixPolicy<double>>
void AddForcingTerms(const MatrixPolicy<double>& rate_constants, const MatrixPolicy<double>& state_variables, MatrixPolicy<double>& forcing)
const;

/// @brief Add Jacobian terms for the set of processes for the current conditions
/// @param rate_constants Current values for the process rate constants (grid cell, process)
/// @param state_variables Current state variable values (grid cell, state variable)
/// @param jacobian Jacobian matrix for the system (grid cell, dependent variable, independent variable)
void AddJacobianTerms(
const Matrix<double>& rate_constants,
const Matrix<double>& state_variables,
SparseMatrix<double>& jacobian) const;
template<template<class> class MatrixPolicy>
void AddJacobianTerms(const MatrixPolicy<double>& rate_constants, const MatrixPolicy<double>& state_variables, SparseMatrix<double>& jacobian)
const;
};

inline ProcessSet::ProcessSet(const std::vector<Process>& processes, const State& state)
template<template<class> class MatrixPolicy>
inline ProcessSet::ProcessSet(const std::vector<Process>& processes, const State<MatrixPolicy>& state)
: number_of_reactants_(),
reactant_ids_(),
number_of_products_(),
Expand Down Expand Up @@ -128,10 +141,10 @@ namespace micm
}
}

inline void ProcessSet::AddForcingTerms(
const Matrix<double>& rate_constants,
const Matrix<double>& state_variables,
Matrix<double>& forcing) const
template<template<class> typename MatrixPolicy>
requires(!Vectorizable<MatrixPolicy<double>>)
inline void
ProcessSet::AddForcingTerms(const MatrixPolicy<double>& rate_constants, const MatrixPolicy<double>& state_variables, MatrixPolicy<double>& forcing) const
{
// loop over grid cells
for (std::size_t i_cell = 0; i_cell < state_variables.size(); ++i_cell)
Expand All @@ -158,9 +171,49 @@ namespace micm
}
};

template<template<class> typename MatrixPolicy>
requires Vectorizable<MatrixPolicy<double>>
inline void
ProcessSet::AddForcingTerms(const MatrixPolicy<double>& rate_constants, const MatrixPolicy<double>& state_variables, MatrixPolicy<double>& forcing) const
{
const auto& v_rate_constants = rate_constants.AsVector();
const auto& v_state_variables = state_variables.AsVector();
auto& v_forcing = forcing.AsVector();
// loop over all rows
for (std::size_t i_block = 0; i_block < state_variables.NumberOfBlocks(); ++i_block)
{
std::size_t L = rate_constants.VectorSize();
auto react_id = reactant_ids_.begin();
auto prod_id = product_ids_.begin();
auto yield = yields_.begin();
std::size_t offset_rc = i_block * rate_constants.BlockSize();
std::size_t offset_state = i_block * state_variables.BlockSize();
std::size_t offset_forcing = i_block * forcing.BlockSize();
for (std::size_t i_rxn = 0; i_rxn < number_of_reactants_.size(); ++i_rxn)
{
double rate[L];
for (std::size_t i_cell = 0; i_cell < L; ++i_cell)
rate[i_cell] = v_rate_constants[offset_rc + i_rxn * L + i_cell];
for (std::size_t i_react = 0; i_react < number_of_reactants_[i_rxn]; ++i_react)
for (std::size_t i_cell = 0; i_cell < L; ++i_cell)
rate[i_cell] *= v_state_variables[offset_state + react_id[i_react] * L + i_cell];
for (std::size_t i_react = 0; i_react < number_of_reactants_[i_rxn]; ++i_react)
for (std::size_t i_cell = 0; i_cell < L; ++i_cell)
v_forcing[offset_forcing + react_id[i_react] * L + i_cell] -= rate[i_cell];
for (std::size_t i_prod = 0; i_prod < number_of_products_[i_rxn]; ++i_prod)
for (std::size_t i_cell = 0; i_cell < L; ++i_cell)
v_forcing[offset_forcing + prod_id[i_prod] * L + i_cell] += yield[i_prod] * rate[i_cell];
react_id += number_of_reactants_[i_rxn];
prod_id += number_of_products_[i_rxn];
yield += number_of_products_[i_rxn];
}
}
}

template<template<class> class MatrixPolicy>
inline void ProcessSet::AddJacobianTerms(
const Matrix<double>& rate_constants,
const Matrix<double>& state_variables,
const MatrixPolicy<double>& rate_constants,
const MatrixPolicy<double>& state_variables,
SparseMatrix<double>& jacobian) const
{
auto cell_jacobian = jacobian.AsVector().begin();
Expand Down
18 changes: 9 additions & 9 deletions include/micm/solver/chapman_ode_solver.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ namespace micm
* @brief An implementation of the Chapman mechnanism solver
*
*/
class ChapmanODESolver : public RosenbrockSolver
class ChapmanODESolver : public RosenbrockSolver<micm::Matrix>
{
std::size_t number_sparse_factor_elements_ = 23;
public:
Expand All @@ -32,7 +32,7 @@ namespace micm

/// Returns a state variable for the Chapman system
/// @return State variable for Chapman
State GetState() const override;
State<> GetState() const;

/// @brief A virtual function to be defined by any solver baseclass
/// @param time_start Time step to start at
Expand All @@ -42,7 +42,7 @@ namespace micm
Solver::SolverResult Solve(
double time_start,
double time_end,
State& state) noexcept;
State<>& state) noexcept;

/// @brief Returns a list of reaction names
/// @return vector of strings
Expand Down Expand Up @@ -81,7 +81,7 @@ namespace micm

/// @brief Update the rate constants for the environment state
/// @param state The current state of the chemical system
void UpdateState(State& state) override;
void UpdateState(State<>& state);

/// @brief Solve the system
/// @param K idk, something
Expand Down Expand Up @@ -132,15 +132,15 @@ namespace micm
{
}

inline State ChapmanODESolver::GetState() const
inline State<> ChapmanODESolver::GetState() const
{
return State{ 9, 3, 7 };
return State<Matrix>{ 9, 3, 7 };
}

inline Solver::SolverResult ChapmanODESolver::Solve(
double time_start,
double time_end,
State& state) noexcept
State<>& state) noexcept
{
std::vector<std::vector<double>> K(parameters_.stages_, std::vector<double>(parameters_.N_, 0));
std::vector<double> Y(state.variables_[0]);
Expand Down Expand Up @@ -565,7 +565,7 @@ namespace micm
return x;
}

inline void ChapmanODESolver::UpdateState(State& state)
inline void ChapmanODESolver::UpdateState(State<>& state)
{
double temperature = state.conditions_[0].temperature_;
double pressure = state.conditions_[0].pressure_;
Expand Down Expand Up @@ -787,4 +787,4 @@ namespace micm

return jacobian;
}
} // namespace micm
} // namespace micm
Loading

0 comments on commit 6b5d220

Please sign in to comment.