Skip to content

Move from omp to kokkos loops in DirectSolver #229

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

Draft
AbdelhadiKara wants to merge 13 commits into
base: main
Choose a base branch
from
Draft

Move from omp to kokkos loops in DirectSolver #229

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
1d7f8b4
attempt to change the first loop
AbdelhadiKara Apr 24, 2026
0d65a25
make methods const
AbdelhadiKara Apr 28, 2026
767007c
copy pass PolarGrid
AbdelhadiKara Apr 28, 2026
725bc32
use atomic add operation, does not fix execution issue
AbdelhadiKara Apr 28, 2026
5b41fb0
pass const reference, use atomic operations
AbdelhadiKara Apr 29, 2026
fba67b3
use deep copy instead of assign, use kokkos function for some methods
AbdelhadiKara Apr 29, 2026
e12b139
pass solver matrix as const reference
AbdelhadiKara Apr 29, 2026
788a40d
remove use of atomic_store, fix index error
AbdelhadiKara Apr 29, 2026
178d80b
revert changes, remove kokkos function use
AbdelhadiKara Apr 29, 2026
8ef3c03
move methods to static inline functions
AbdelhadiKara Apr 30, 2026
1759adf
remove Stecil relaed methods
AbdelhadiKara Apr 30, 2026
978963c
extract stencils from class
AbdelhadiKara Apr 30, 2026
dfb57b5
remove commented arguments
AbdelhadiKara Apr 30, 2026
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
167 changes: 118 additions & 49 deletions include/DirectSolver/DirectSolver-CSR-LU-Give/buildSolverMatrix.inl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,72 @@ static inline void updateMatrixElement(SparseMatrixCSR<double>& matrix, int offs
matrix.row_nz_entry(row, offset) += val;
}

} // namespace direct_solver_csr_lu_give
static inline int getStencilSize(int global_index, const PolarGrid& grid, const bool DirBC_Interior)
{
//const PolarGrid& grid = DirectSolver<LevelCacheType>::grid_;
// const bool DirBC_Interior = DirectSolver<LevelCacheType>::DirBC_Interior_;

template <class LevelCacheType>
void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r, int i_theta, const PolarGrid& grid,
const bool DirBC_Interior,
SparseMatrixCSR<double>& solver_matrix,
double arr, double att, double art,
double detDF, double coeff_beta)
int i_r, i_theta;
grid.multiIndex(global_index, i_r, i_theta);

const int size_stencil_inner_boundary = DirBC_Interior ? 1 : 7;
const int size_stencil_next_inner_boundary = DirBC_Interior ? 9 : 9;
const int size_stencil_interior = 9;
const int size_stencil_next_outer_boundary = 9;
const int size_stencil_outer_boundary = 1;

if ((i_r > 1 && i_r < grid.nr() - 2) || (i_r == 1 && !DirBC_Interior)) {
return size_stencil_interior;
}
else if (i_r == 0 && !DirBC_Interior) {
return size_stencil_inner_boundary;
}
else if ((i_r == 0 && DirBC_Interior) || i_r == grid.nr() - 1) {
return size_stencil_outer_boundary;
}
else if (i_r == 1 && DirBC_Interior) {
return size_stencil_next_inner_boundary;
}
else if (i_r == grid.nr() - 2) {
return size_stencil_next_outer_boundary;
}
throw std::out_of_range("Invalid index for stencil");
}

static inline const Stencil& getStencil(int i_r, const PolarGrid& grid, const bool DirBC_Interior)
{
//const PolarGrid& grid = DirectSolver<LevelCacheType>::grid_;
//const bool DirBC_Interior = DirectSolver<LevelCacheType>::DirBC_Interior_;

assert(0 <= i_r && i_r < grid.nr());
assert(grid.nr() >= 4);

if ((i_r > 1 && i_r < grid.nr() - 2) || (i_r == 1 && !DirBC_Interior)) {
return stencil_interior_;
}
else if (i_r == 0 && !DirBC_Interior) {
return stencil_across_origin_;
}
else if ((i_r == 0 && DirBC_Interior) || i_r == grid.nr() - 1) {
return stencil_DB_;
}
else if (i_r == 1 && DirBC_Interior) {
return stencil_next_inner_DB_;
}
else if (i_r == grid.nr() - 2) {
return stencil_next_outer_DB_;
}
throw std::out_of_range("Invalid index for stencil");
}

static inline void nodeBuildSolverMatrixGive(int i_r, int i_theta, SparseMatrixCSR<double>& solver_matrix,
const PolarGrid grid, const bool DirBC_Interior, double arr, double att,
double art, double detDF, double coeff_beta)
{
using direct_solver_csr_lu_give::getStencil;
using direct_solver_csr_lu_give::getStencilSize;
using direct_solver_csr_lu_give::updateMatrixElement;

int offset;
int row, col;
double val;
Expand Down Expand Up @@ -48,7 +104,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i,j) */
row = center_index;

const Stencil& CenterStencil = getStencil(i_r);
const Stencil& CenterStencil = getStencil(i_r, grid, DirBC_Interior);

offset = CenterStencil[StencilPosition::Center];
col = center_index;
Expand Down Expand Up @@ -83,7 +139,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i-1,j) */
row = left_index;

const Stencil& LeftStencil = getStencil(i_r - 1);
const Stencil& LeftStencil = getStencil(i_r - 1, grid, DirBC_Interior);

offset = LeftStencil[StencilPosition::Right];
col = center_index;
Expand All @@ -108,7 +164,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i+1,j) */
row = right_index;

const Stencil& RightStencil = getStencil(i_r + 1);
const Stencil& RightStencil = getStencil(i_r + 1, grid, DirBC_Interior);

offset = RightStencil[StencilPosition::Left];
col = center_index;
Expand Down Expand Up @@ -205,7 +261,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i,j) */
row = center_index;

const Stencil& CenterStencil = getStencil(i_r);
const Stencil& CenterStencil = getStencil(i_r, grid, DirBC_Interior);

offset = CenterStencil[StencilPosition::Center];
col = center_index;
Expand All @@ -215,7 +271,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i+1,j) */
row = right_index;

const Stencil& RightStencil = getStencil(i_r + 1);
const Stencil& RightStencil = getStencil(i_r + 1, grid, DirBC_Interior);

offset = RightStencil[StencilPosition::Left];
col = center_index;
Expand Down Expand Up @@ -268,7 +324,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i,j) */
row = center_index;

const Stencil& CenterStencil = getStencil(i_r);
const Stencil& CenterStencil = getStencil(i_r, grid, DirBC_Interior);

offset = CenterStencil[StencilPosition::Center];
col = center_index;
Expand Down Expand Up @@ -324,7 +380,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i+1,j) */
row = right_index;

const Stencil& RightStencil = getStencil(i_r + 1);
const Stencil& RightStencil = getStencil(i_r + 1, grid, DirBC_Interior);

offset = RightStencil[StencilPosition::Left];
col = center_index;
Expand Down Expand Up @@ -417,7 +473,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i,j) */
row = center_index;

const Stencil& CenterStencil = getStencil(i_r);
const Stencil& CenterStencil = getStencil(i_r, grid, DirBC_Interior);

offset = CenterStencil[StencilPosition::Center];
col = center_index;
Expand Down Expand Up @@ -453,7 +509,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i-1,j) */
row = left_index;

const Stencil& LeftStencil = getStencil(i_r - 1);
const Stencil& LeftStencil = getStencil(i_r - 1, grid, DirBC_Interior);

offset = LeftStencil[StencilPosition::Right];
col = center_index;
Expand All @@ -479,7 +535,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i+1,j) */
row = right_index;

const Stencil& RightStencil = getStencil(i_r + 1);
const Stencil& RightStencil = getStencil(i_r + 1, grid, DirBC_Interior);

offset = RightStencil[StencilPosition::Left];
col = center_index;
Expand Down Expand Up @@ -577,7 +633,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i,j) */
row = center_index;

const Stencil& CenterStencil = getStencil(i_r);
const Stencil& CenterStencil = getStencil(i_r, grid, DirBC_Interior);

offset = CenterStencil[StencilPosition::Center];
col = center_index;
Expand Down Expand Up @@ -612,7 +668,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i-1,j) */
row = left_index;

const Stencil& LeftStencil = getStencil(i_r - 1);
const Stencil& LeftStencil = getStencil(i_r - 1, grid, DirBC_Interior);

offset = LeftStencil[StencilPosition::Right];
col = center_index;
Expand Down Expand Up @@ -707,7 +763,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i,j) */
row = center_index;

const Stencil& CenterStencil = getStencil(i_r);
const Stencil& CenterStencil = getStencil(i_r, grid, DirBC_Interior);

offset = CenterStencil[StencilPosition::Center];
col = center_index;
Expand All @@ -718,7 +774,7 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
/* Fill matrix row of (i-1,j) */
row = left_index;

const Stencil& LeftStencil = getStencil(i_r - 1);
const Stencil& LeftStencil = getStencil(i_r - 1, grid, DirBC_Interior);

offset = LeftStencil[StencilPosition::Right];
col = center_index;
Expand All @@ -743,12 +799,10 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::nodeBuildSolverMatrixGive(int i_r
}

template <class LevelCacheType>
void DirectSolver_CSR_LU_Give<LevelCacheType>::buildSolverMatrixCircleSection(const int i_r,
SparseMatrixCSR<double>& solver_matrix)
static inline void buildSolverMatrixCircleSection(const int i_r, SparseMatrixCSR<double>& solver_matrix,
const PolarGrid& grid, const LevelCacheType& level_cache,
const bool DirBC_Interior)
{
const PolarGrid& grid = DirectSolver<LevelCacheType>::grid_;
const LevelCacheType& level_cache = DirectSolver<LevelCacheType>::level_cache_;
const bool DirBC_Interior = DirectSolver<LevelCacheType>::DirBC_Interior_;

const double r = grid.radius(i_r);
for (int i_theta = 0; i_theta < grid.ntheta(); i_theta++) {
Expand All @@ -759,10 +813,13 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::buildSolverMatrixCircleSection(co
level_cache.obtainValues(i_r, i_theta, global_index, r, theta, coeff_beta, arr, att, art, detDF);

// Build solver matrix at the current node
nodeBuildSolverMatrixGive(i_r, i_theta, grid, DirBC_Interior, solver_matrix, arr, att, art, detDF, coeff_beta);
direct_solver_csr_lu_give::nodeBuildSolverMatrixGive(i_r, i_theta, solver_matrix, grid, DirBC_Interior, arr,
att, art, detDF, coeff_beta);
}
}

} // namespace direct_solver_csr_lu_give

template <class LevelCacheType>
void DirectSolver_CSR_LU_Give<LevelCacheType>::buildSolverMatrixRadialSection(const int i_theta,
SparseMatrixCSR<double>& solver_matrix)
Expand All @@ -780,28 +837,35 @@ void DirectSolver_CSR_LU_Give<LevelCacheType>::buildSolverMatrixRadialSection(co
level_cache.obtainValues(i_r, i_theta, global_index, r, theta, coeff_beta, arr, att, art, detDF);

// Build solver matrix at the current node
nodeBuildSolverMatrixGive(i_r, i_theta, grid, DirBC_Interior, solver_matrix, arr, att, art, detDF, coeff_beta);
direct_solver_csr_lu_give::nodeBuildSolverMatrixGive(i_r, i_theta, solver_matrix, grid, DirBC_Interior, arr,
att, art, detDF, coeff_beta);
}
}

template <class LevelCacheType>
SparseMatrixCSR<double> DirectSolver_CSR_LU_Give<LevelCacheType>::buildSolverMatrix()
{
const PolarGrid& grid = DirectSolver<LevelCacheType>::grid_;
const int num_omp_threads = DirectSolver<LevelCacheType>::num_omp_threads_;
using direct_solver_csr_lu_give::buildSolverMatrixCircleSection;
using direct_solver_csr_lu_give::getStencilSize;

const PolarGrid& grid = DirectSolver<LevelCacheType>::grid_;
const LevelCacheType& level_cache = DirectSolver<LevelCacheType>::level_cache_;
const bool DirBC_Interior = DirectSolver<LevelCacheType>::DirBC_Interior_;
const int num_omp_threads = DirectSolver<LevelCacheType>::num_omp_threads_;

const int n = grid.numberOfNodes();

std::function<int(int)> nnz_per_row = [&](int global_index) {
return getStencilSize(global_index);
std::function<int(int)> nnz_per_row = KOKKOS_CLASS_LAMBDA(int global_index)
{
return getStencilSize(global_index, grid, DirBC_Interior);
};

SparseMatrixCSR<double> solver_matrix(n, n, nnz_per_row);

if (num_omp_threads == 1) {
/* Single-threaded execution */
for (int i_r = 0; i_r < grid.numberSmootherCircles(); i_r++) {
buildSolverMatrixCircleSection(i_r, solver_matrix);
buildSolverMatrixCircleSection(i_r, solver_matrix, grid, level_cache, DirBC_Interior);
}
for (int i_theta = 0; i_theta < grid.ntheta(); i_theta++) {
buildSolverMatrixRadialSection(i_theta, solver_matrix);
Expand All @@ -812,25 +876,30 @@ SparseMatrixCSR<double> DirectSolver_CSR_LU_Give<LevelCacheType>::buildSolverMat
const int num_circle_tasks = grid.numberSmootherCircles();
const int additional_radial_tasks = grid.ntheta() % 3;
const int num_radial_tasks = grid.ntheta() - additional_radial_tasks;
const int stride = 3;
const int num_steps = (num_circle_tasks + stride - 1) / stride;

//for (int circle_task = 0; circle_task < num_circle_tasks; circle_task += 3) {
Kokkos::parallel_for(
num_steps, KOKKOS_LAMBDA(const int circle_task) {
int circle_idx = stride * circle_task;
int i_r = grid.numberSmootherCircles() - circle_idx - 1;
buildSolverMatrixCircleSection(i_r, solver_matrix, grid, level_cache, DirBC_Interior);
});
Kokkos::fence();

for (int circle_task = 1; circle_task < num_circle_tasks; circle_task += 3) {
int i_r = grid.numberSmootherCircles() - circle_task - 1;
buildSolverMatrixCircleSection(i_r, solver_matrix, grid, level_cache, DirBC_Interior);
;
}

for (int circle_task = 2; circle_task < num_circle_tasks; circle_task += 3) {
int i_r = grid.numberSmootherCircles() - circle_task - 1;
buildSolverMatrixCircleSection(i_r, solver_matrix, grid, level_cache, DirBC_Interior);
}
#pragma omp parallel num_threads(num_omp_threads)
{
#pragma omp for
for (int circle_task = 0; circle_task < num_circle_tasks; circle_task += 3) {
int i_r = grid.numberSmootherCircles() - circle_task - 1;
buildSolverMatrixCircleSection(i_r, solver_matrix);
}
#pragma omp for
for (int circle_task = 1; circle_task < num_circle_tasks; circle_task += 3) {
int i_r = grid.numberSmootherCircles() - circle_task - 1;
buildSolverMatrixCircleSection(i_r, solver_matrix);
}
#pragma omp for nowait
for (int circle_task = 2; circle_task < num_circle_tasks; circle_task += 3) {
int i_r = grid.numberSmootherCircles() - circle_task - 1;
buildSolverMatrixCircleSection(i_r, solver_matrix);
}

#pragma omp for
for (int radial_task = 0; radial_task < num_radial_tasks; radial_task += 3) {
if (radial_task > 0) {
Expand Down
42 changes: 21 additions & 21 deletions include/DirectSolver/DirectSolver-CSR-LU-Give/directSolverGiveCustomLU.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,22 +5,7 @@
namespace gmgpolar
{

template <class LevelCacheType>
class DirectSolver_CSR_LU_Give : public DirectSolver<LevelCacheType>
{
public:
explicit DirectSolver_CSR_LU_Give(const PolarGrid& grid, const LevelCacheType& level_cache, bool DirBC_Interior,
int num_omp_threads);

// Note: The rhs (right-hand side) vector gets overwritten with the solution.
void solveInPlace(Vector<double> solution) override;

private:
// Solver matrix and solver structure
SparseMatrixCSR<double> solver_matrix_;
SparseLUSolver<double> lu_solver_;

// clang-format off
// clang-format off
const Stencil stencil_interior_ = {
7, 4, 8,
1, 0, 2,
Expand All @@ -46,10 +31,25 @@ class DirectSolver_CSR_LU_Give : public DirectSolver<LevelCacheType>
1, 0, 2,
5, 3, 6
};
// clang-format on
// clang-format on

template <class LevelCacheType>
class DirectSolver_CSR_LU_Give : public DirectSolver<LevelCacheType>
{
public:
explicit DirectSolver_CSR_LU_Give(const PolarGrid& grid, const LevelCacheType& level_cache, bool DirBC_Interior,
int num_omp_threads);

// Note: The rhs (right-hand side) vector gets overwritten with the solution.
void solveInPlace(Vector<double> solution) override;

private:
// Solver matrix and solver structure
SparseMatrixCSR<double> solver_matrix_;
SparseLUSolver<double> lu_solver_;

SparseMatrixCSR<double> buildSolverMatrix();
void buildSolverMatrixCircleSection(const int i_r, SparseMatrixCSR<double>& solver_matrix);
//void buildSolverMatrixCircleSection(const int i_r, SparseMatrixCSR<double>& solver_matrix) const;
void buildSolverMatrixRadialSection(const int i_theta, SparseMatrixCSR<double>& solver_matrix);

// Returns the total number of non-zero elements in the solver matrix.
Expand All @@ -59,9 +59,9 @@ class DirectSolver_CSR_LU_Give : public DirectSolver<LevelCacheType>

int getStencilSize(int global_index) const;

void nodeBuildSolverMatrixGive(int i_r, int i_theta, const PolarGrid& grid, const bool DirBC_Interior,
SparseMatrixCSR<double>& solver_matrix, double arr, double att, double art,
double detDF, double coeff_beta);
/* void nodeBuildSolverMatrixGive(int i_r, int i_theta, const PolarGrid grid, const bool DirBC_Interior,
const SparseMatrixCSR<double>& solver_matrix, double arr, double att, double art,
double detDF, double coeff_beta) const;*/
};

#include "buildSolverMatrix.inl"
Expand Down
Loading