Skip to content

Commit

Permalink
it compiles
Browse files Browse the repository at this point in the history
  • Loading branch information
@K20shores
K20shores committed May 23, 2024
1 parent a4e7d79 commit 2fb3917
Show file tree
Hide file tree
Showing 14 changed files with 129 additions and 496 deletions.
16 changes: 7 additions & 9 deletions include/micm/solver/jit_linear_solver.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@ namespace micm
template<std::size_t L, class SparseMatrixPolicy, class LuDecompositionPolicy = JitLuDecomposition<L>>
class JitLinearSolver : public LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>
{
std::shared_ptr<JitCompiler> compiler_;
llvm::orc::ResourceTrackerSP solve_function_resource_tracker_;
void (*solve_function_)(const double*, double*, const double*, const double*);

Expand All @@ -39,8 +38,7 @@ namespace micm
/// @param matrix Block-diagonal sparse matrix to create solver for
/// @param initial_value Initial value for decomposed triangular matrix elements
JitLinearSolver(
std::shared_ptr<JitCompiler> compiler,
const SparseMatrix<double, SparseMatrixVectorOrdering<L>>& matrix,
const SparseMatrixPolicy& matrix,
double initial_value);

~JitLinearSolver();
Expand All @@ -49,17 +47,17 @@ namespace micm
/// @param matrix Matrix that will be factored into lower and upper triangular matrices
/// @param is_singular Flag that will be set to true if matrix is singular; false otherwise
void Factor(
SparseMatrix<double, SparseMatrixVectorOrdering<L>>& matrix,
SparseMatrix<double, SparseMatrixVectorOrdering<L>>& lower_matrix,
SparseMatrix<double, SparseMatrixVectorOrdering<L>>& upper_matrix,
SparseMatrixPolicy& matrix,
SparseMatrixPolicy& lower_matrix,
SparseMatrixPolicy& upper_matrix,
bool& is_singular);

/// @brief Decompose the matrix into upper and lower triangular matrices and general JIT functions
/// @param matrix Matrix that will be factored into lower and upper triangular matrices
void Factor(
SparseMatrix<double, SparseMatrixVectorOrdering<L>>& matrix,
SparseMatrix<double, SparseMatrixVectorOrdering<L>>& lower_matrix,
SparseMatrix<double, SparseMatrixVectorOrdering<L>>& upper_matrix);
SparseMatrixPolicy& matrix,
SparseMatrixPolicy& lower_matrix,
SparseMatrixPolicy& upper_matrix);

/// @brief Solve for x in Ax = b
template<class MatrixPolicy>
Expand Down
28 changes: 12 additions & 16 deletions include/micm/solver/jit_linear_solver.inl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ namespace micm
template<std::size_t L, class SparseMatrixPolicy, class LuDecompositionPolicy>
inline JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::JitLinearSolver(JitLinearSolver &&other)
: LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>(std::move(other)),
compiler_(std::move(other.compiler_)),
solve_function_resource_tracker_(std::move(other.solve_function_resource_tracker_)),
solve_function_(std::move(other.solve_function_))
{
Expand All @@ -20,7 +19,6 @@ namespace micm
&JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::operator=(JitLinearSolver &&other)
{
LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>::operator=(std::move(other));
compiler_ = std::move(other.compiler_);
solve_function_resource_tracker_ = std::move(other.solve_function_resource_tracker_);
solve_function_ = std::move(other.solve_function_);
other.solve_function_ = NULL;
Expand All @@ -29,15 +27,13 @@ namespace micm

template<std::size_t L, class SparseMatrixPolicy, class LuDecompositionPolicy>
inline JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::JitLinearSolver(
std::shared_ptr<JitCompiler> compiler,
const SparseMatrix<SparseMatrixVectorOrdering<L>> &matrix,
const SparseMatrixPolicy &matrix,
double initial_value)
: LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>(
matrix,
initial_value,
[&](const SparseMatrixPolicy &m) -> LuDecompositionPolicy
{ return LuDecompositionPolicy(compiler, m); }),
compiler_(compiler)
{ return LuDecompositionPolicy(m); })
{
solve_function_ = NULL;
if (matrix.NumberOfBlocks() != L || matrix.GroupVectorSize() != L)
Expand All @@ -63,19 +59,19 @@ namespace micm

template<std::size_t L, class SparseMatrixPolicy, class LuDecompositionPolicy>
inline void JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::Factor(
SparseMatrix<double, SparseMatrixVectorOrdering<L>> &matrix,
SparseMatrix<double, SparseMatrixVectorOrdering<L>> &lower_matrix,
SparseMatrix<double, SparseMatrixVectorOrdering<L>> &upper_matrix,
SparseMatrixPolicy &matrix,
SparseMatrixPolicy &lower_matrix,
SparseMatrixPolicy &upper_matrix,
bool &is_singular)
{
LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>::Factor(matrix, lower_matrix, upper_matrix, is_singular);
}

template<std::size_t L, class SparseMatrixPolicy, class LuDecompositionPolicy>
inline void JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::Factor(
SparseMatrix<double, SparseMatrixVectorOrdering<L>> &matrix,
SparseMatrix<double, SparseMatrixVectorOrdering<L>> &lower_matrix,
SparseMatrix<double, SparseMatrixVectorOrdering<L>> &upper_matrix)
SparseMatrixPolicy &matrix,
SparseMatrixPolicy &lower_matrix,
SparseMatrixPolicy &upper_matrix)
{
LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>::Factor(matrix, lower_matrix, upper_matrix);
}
Expand All @@ -96,18 +92,18 @@ namespace micm
inline void JitLinearSolver<L, SparseMatrixPolicy, LuDecompositionPolicy>::GenerateSolveFunction()
{
std::string function_name = "linear_solve_" + GenerateRandomString();
JitFunction func = JitFunction::Create(compiler_)
JitFunction func = JitFunction::Create()
.SetName(function_name)
.SetArguments({ { "b", JitType::DoublePtr },
{ "x", JitType::DoublePtr },
{ "L", JitType::DoublePtr },
{ "U", JitType::DoublePtr } })
.SetReturnType(JitType::Void);
llvm::Type *double_type = func.GetType(JitType::Double);
auto Lij_yj = LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>::Lij_yj_.begin();
auto Uij_xj = LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>::Uij_xj_.begin();
auto Lij_yj = LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>::Lij_yj_.begin();
auto Uij_xj = LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>::Uij_xj_.begin();
std::size_t offset = 0;
for (auto &nLij_Lii : LinearSolver<double, SparseMatrixPolicy, LuDecompositionPolicy>::nLij_Lii_)
for (auto &nLij_Lii : LinearSolver<SparseMatrixPolicy, LuDecompositionPolicy>::nLij_Lii_)
{
// the x vector is used for y values to conserve memory
{
Expand Down
19 changes: 7 additions & 12 deletions include/micm/solver/jit_lu_decomposition.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
*/
#pragma once

#include <micm/jit/jit_compiler.hpp>
#include <micm/jit/jit_function.hpp>
#include <micm/solver/lu_decomposition.hpp>
#include <micm/util/random_string.hpp>
Expand All @@ -20,7 +19,6 @@ namespace micm
template<std::size_t L = MICM_DEFAULT_VECTOR_SIZE>
class JitLuDecomposition : public LuDecomposition
{
std::shared_ptr<JitCompiler> compiler_;
llvm::orc::ResourceTrackerSP decompose_function_resource_tracker_;
void (*decompose_function_)(const double *, double *, double *);

Expand All @@ -33,11 +31,8 @@ namespace micm
JitLuDecomposition &operator=(JitLuDecomposition &&);

/// @brief Create a JITed LU decomposer for a given sparse matrix structure
/// @param compiler JIT compiler
/// @param matrix Sparse matrix to create LU decomposer for
JitLuDecomposition(
std::shared_ptr<JitCompiler> compiler,
const SparseMatrix<double, SparseMatrixVectorOrdering<L>> &matrix);
JitLuDecomposition(const SparseMatrix<double, SparseMatrixVectorOrdering<L>> &matrix);

~JitLuDecomposition();

Expand All @@ -46,19 +41,19 @@ namespace micm
/// @param lower The lower triangular matrix created by decomposition
/// @param upper The upper triangular matrix created by decomposition
/// @param is_singular Flag that will be set to true if A is singular; false otherwise
template<typename T, template<class> class SparseMatrixPolicy>
template<class SparseMatrixPolicy>
void Decompose(
const SparseMatrixPolicy<T> &A,
SparseMatrixPolicy<T> &lower,
SparseMatrixPolicy<T> &upper,
const SparseMatrixPolicy &A,
SparseMatrixPolicy &lower,
SparseMatrixPolicy &upper,
bool &is_singular) const;

/// @brief Create sparse L and U matrices for a given A matrix
/// @param A Sparse matrix that will be decomposed
/// @param lower The lower triangular matrix created by decomposition
/// @param upper The upper triangular matrix created by decomposition
template<typename T, template<class> class SparseMatrixPolicy>
void Decompose(const SparseMatrixPolicy<T> &A, SparseMatrixPolicy<T> &lower, SparseMatrixPolicy<T> &upper) const;
template<class SparseMatrixPolicy>
void Decompose(const SparseMatrixPolicy &A, SparseMatrixPolicy &lower, SparseMatrixPolicy &upper) const;

private:
/// @brief Generates a function to perform the LU decomposition for a specific matrix sparsity structure
Expand Down
26 changes: 11 additions & 15 deletions include/micm/solver/jit_lu_decomposition.inl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ namespace micm
template<std::size_t L>
inline JitLuDecomposition<L>::JitLuDecomposition(JitLuDecomposition &&other)
: LuDecomposition(std::move(other)),
compiler_(std::move(other.compiler_)),
decompose_function_resource_tracker_(std::move(other.decompose_function_resource_tracker_)),
decompose_function_(std::move(other.decompose_function_))
{
Expand All @@ -19,7 +18,6 @@ namespace micm
inline JitLuDecomposition<L> &JitLuDecomposition<L>::operator=(JitLuDecomposition &&other)
{
LuDecomposition::operator=(std::move(other));
compiler_ = std::move(other.compiler_);
decompose_function_resource_tracker_ = std::move(other.decompose_function_resource_tracker_);
decompose_function_ = std::move(other.decompose_function_);
other.decompose_function_ = NULL;
Expand All @@ -28,10 +26,8 @@ namespace micm

template<std::size_t L>
inline JitLuDecomposition<L>::JitLuDecomposition(
std::shared_ptr<JitCompiler> compiler,
const SparseMatrix<double, SparseMatrixVectorOrdering<L>> &matrix)
: LuDecomposition(LuDecomposition::Create<double, SparseMatrix<double, SparseMatrixVectorOrdering<L>>>(matrix)),
compiler_(compiler)
: LuDecomposition(LuDecomposition::Create<SparseMatrix<double, SparseMatrixVectorOrdering<L>>>(matrix))
{
decompose_function_ = NULL;
if (matrix.NumberOfBlocks() > L)
Expand Down Expand Up @@ -59,7 +55,7 @@ namespace micm
void JitLuDecomposition<L>::GenerateDecomposeFunction()
{
std::string function_name = "lu_decompose_" + GenerateRandomString();
JitFunction func = JitFunction::Create(compiler_)
JitFunction func = JitFunction::Create()
.SetName(function_name)
.SetArguments({ { "A matrix", JitType::DoublePtr },
{ "lower matrix", JitType::DoublePtr },
Expand Down Expand Up @@ -190,22 +186,22 @@ namespace micm
}

template<std::size_t L>
template<typename T, template<class> class SparseMatrixPolicy>
template<class SparseMatrixPolicy>
void JitLuDecomposition<L>::Decompose(
const SparseMatrixPolicy<T> &A,
SparseMatrixPolicy<T> &lower,
SparseMatrixPolicy<T> &upper,
const SparseMatrixPolicy &A,
SparseMatrixPolicy &lower,
SparseMatrixPolicy &upper,
bool &is_singular) const
{
LuDecomposition::Decompose<T, SparseMatrixPolicy>(A, lower, upper, is_singular);
LuDecomposition::Decompose<SparseMatrixPolicy>(A, lower, upper, is_singular);
}

template<std::size_t L>
template<typename T, template<class> class SparseMatrixPolicy>
template<class SparseMatrixPolicy>
void JitLuDecomposition<L>::Decompose(
const SparseMatrixPolicy<T> &A,
SparseMatrixPolicy<T> &lower,
SparseMatrixPolicy<T> &upper) const
const SparseMatrixPolicy &A,
SparseMatrixPolicy &lower,
SparseMatrixPolicy &upper) const
{
decompose_function_(A.AsVector().data(), lower.AsVector().data(), upper.AsVector().data());
}
Expand Down
20 changes: 7 additions & 13 deletions include/micm/solver/jit_rosenbrock.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
*/
#pragma once

#include <micm/jit/jit_compiler.hpp>
#include <micm/jit/jit_function.hpp>
#include <micm/process/jit_process_set.hpp>
#include <micm/solver/jit_linear_solver.hpp>
Expand All @@ -33,12 +32,11 @@ namespace micm

template<
template<class> class MatrixPolicy = VectorMatrix,
template<class> class SparseMatrixPolicy = VectorSparseMatrix,
class SparseMatrixPolicy = DefaultVectorSparseMatrix,
class LinearSolverPolicy = JitLinearSolver<MICM_DEFAULT_VECTOR_SIZE, SparseMatrixPolicy>,
class ProcessSetPolicy = JitProcessSet<MICM_DEFAULT_VECTOR_SIZE>>
class JitRosenbrockSolver : public RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy, LinearSolverPolicy, ProcessSetPolicy>
{
std::shared_ptr<JitCompiler> compiler_;
llvm::orc::ResourceTrackerSP function_resource_tracker_;
using FuncPtr = void (*)(double*, const double);
FuncPtr alpha_minus_jacobian_ = nullptr;
Expand All @@ -48,7 +46,6 @@ namespace micm
JitRosenbrockSolver& operator=(const JitRosenbrockSolver&) = delete;
JitRosenbrockSolver(JitRosenbrockSolver&& other)
: RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy, LinearSolverPolicy, ProcessSetPolicy>(std::move(other)),
compiler_(std::move(other.compiler_)),
function_resource_tracker_(std::move(other.function_resource_tracker_)),
alpha_minus_jacobian_(std::move(other.alpha_minus_jacobian_))
{
Expand All @@ -58,7 +55,6 @@ namespace micm
JitRosenbrockSolver& operator=(JitRosenbrockSolver&& other)
{
RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy, LinearSolverPolicy, ProcessSetPolicy>::operator=(std::move(other));
compiler_ = std::move(other.compiler_);
function_resource_tracker_ = std::move(other.function_resource_tracker_);
alpha_minus_jacobian_ = std::move(other.alpha_minus_jacobian_);
other.alpha_minus_jacobian_ = NULL;
Expand All @@ -69,22 +65,20 @@ namespace micm
/// @param system The chemical system to create the solver for
/// @param processes The collection of chemical processes that will be applied during solving
JitRosenbrockSolver(
std::shared_ptr<JitCompiler> compiler,
const System& system,
const std::vector<Process>& processes,
const RosenbrockSolverParameters& parameters)
: RosenbrockSolver<MatrixPolicy, SparseMatrixPolicy, LinearSolverPolicy, ProcessSetPolicy>(
system,
processes,
parameters,
[&](const SparseMatrixPolicy<double>& matrix, double initial_value) -> LinearSolverPolicy {
return LinearSolverPolicy{ compiler, matrix, initial_value };
[&](const SparseMatrixPolicy& matrix, double initial_value) -> LinearSolverPolicy {
return LinearSolverPolicy{ matrix, initial_value };
},
[&](const std::vector<Process>& processes,
const std::map<std::string, std::size_t>& variable_map) -> ProcessSetPolicy {
return ProcessSetPolicy{ compiler, processes, variable_map };
}),
compiler_(compiler)
return ProcessSetPolicy{ processes, variable_map };
})
{
MatrixPolicy<double> temp{};
if (temp.GroupVectorSize() != parameters.number_of_grid_cells_)
Expand All @@ -110,7 +104,7 @@ namespace micm
/// @brief compute [alpha * I - dforce_dy]
/// @param jacobian Jacobian matrix (dforce_dy)
/// @param alpha
void AlphaMinusJacobian(SparseMatrixPolicy<double>& jacobian, const double& alpha) const
void AlphaMinusJacobian(SparseMatrixPolicy& jacobian, const double& alpha) const
{
double a = alpha;
if (alpha_minus_jacobian_)
Expand All @@ -133,7 +127,7 @@ namespace micm
std::size_t number_of_nonzero_jacobian_elements = jacobian.AsVector().size();

// Create the JitFunction with the modified name
JitFunction func = JitFunction::Create(compiler_)
JitFunction func = JitFunction::Create()
.SetName("alpha_minus_jacobian")
.SetArguments({ { "jacobian", JitType::DoublePtr }, { "alpha", JitType::Double } })
.SetReturnType(JitType::Void);
Expand Down
3 changes: 1 addition & 2 deletions include/micm/util/sparse_matrix_vector_ordering.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,6 @@ namespace micm
};

// Default vectorized SparseMatrix
template<class T>
using VectorSparseMatrix = SparseMatrix<T, SparseMatrixVectorOrdering<MICM_DEFAULT_VECTOR_SIZE>>;
using DefaultVectorSparseMatrix = SparseMatrix<double, SparseMatrixVectorOrdering<MICM_DEFAULT_VECTOR_SIZE>>;

} // namespace micm
Loading

0 comments on commit 2fb3917

Please sign in to comment.