Skip to content

Use Kokkos::Parallel in ExtrapolatedSmootherTake ApplyAscOrtho #243

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

Merged
julianlitz merged 8 commits into from
May 12, 2026
Merged

Use Kokkos::Parallel in ExtrapolatedSmootherTake ApplyAscOrtho #243

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0800d6b
Use Kokkos::Parallel in ExtrapolatedSmootherTake ApplyAscOrtho
julianlitz May 2, 2026
50075ff
Update applyAscOrtho.inl
julianlitz May 6, 2026
c134300
Merge remote-tracking branch 'origin/main' into litz_011
EmilyBourne May 12, 2026
f2f4aad
Missing annotation
EmilyBourne May 12, 2026
4231c6c
Functions must be public
EmilyBourne May 12, 2026
437a57c
Run on host
EmilyBourne May 12, 2026
029128c
Clang formatting
EmilyBourne May 12, 2026
0147a4c
Merge branch 'main' into litz_011
EmilyBourne May 12, 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
20 changes: 11 additions & 9 deletions include/ExtrapolatedSmoother/ExtrapolatedSmootherGive/applyAscOrtho.inl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,11 @@
namespace extrapolated_smoother_give
{

static inline void nodeApplyAscOrthoCircleGive(int i_r, int i_theta, const PolarGrid& grid, bool DirBC_Interior,
SmootherColor smoother_color, ConstVector<double>& x,
ConstVector<double>& rhs, Vector<double>& result, double arr, double att,
double art, double detDF, double coeff_beta)
static KOKKOS_INLINE_FUNCTION void nodeApplyAscOrthoCircleGive(int i_r, int i_theta, const PolarGrid& grid,
bool DirBC_Interior, SmootherColor smoother_color,
ConstVector<double>& x, ConstVector<double>& rhs,
Vector<double>& result, double arr, double att,
double art, double detDF, double coeff_beta)
{
assert(i_r >= 0 && i_r <= grid.numberSmootherCircles());

Expand Down Expand Up @@ -400,10 +401,11 @@ static inline void nodeApplyAscOrthoCircleGive(int i_r, int i_theta, const Polar
namespace extrapolated_smoother_give
{

static inline void nodeApplyAscOrthoRadialGive(int i_r, int i_theta, const PolarGrid& grid, bool DirBC_Interior,
SmootherColor smoother_color, ConstVector<double>& x,
ConstVector<double>& rhs, Vector<double>& result, double arr, double att,
double art, double detDF, double coeff_beta)
static KOKKOS_INLINE_FUNCTION void nodeApplyAscOrthoRadialGive(int i_r, int i_theta, const PolarGrid& grid,
bool DirBC_Interior, SmootherColor smoother_color,
ConstVector<double>& x, ConstVector<double>& rhs,
Vector<double>& result, double arr, double att,
double art, double detDF, double coeff_beta)
{
assert(i_r >= grid.numberSmootherCircles() - 1 && i_r < grid.nr());

Expand Down Expand Up @@ -946,4 +948,4 @@ void ExtrapolatedSmootherGive<LevelCacheType>::applyAscOrthoRadialSection(int i_
nodeApplyAscOrthoRadialGive(i_r, i_theta, grid, DirBC_Interior, smoother_color, x, rhs, temp, arr, att, art,
detDF, coeff_beta);
}
}
}
105 changes: 72 additions & 33 deletions include/ExtrapolatedSmoother/ExtrapolatedSmootherTake/applyAscOrtho.inl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,12 @@
namespace extrapolated_smoother_take
{

static inline void nodeApplyAscOrthoCircleTake(int i_r, int i_theta, const PolarGrid& grid, bool DirBC_Interior,
ConstVector<double>& x, ConstVector<double>& rhs, Vector<double>& result,
ConstVector<double>& arr, ConstVector<double>& att,
ConstVector<double>& art, ConstVector<double>& detDF,
ConstVector<double>& coeff_beta)
static KOKKOS_INLINE_FUNCTION void nodeApplyAscOrthoCircleTake(int i_r, int i_theta, const PolarGrid& grid,
bool DirBC_Interior, ConstVector<double>& x,
ConstVector<double>& rhs, Vector<double>& result,
ConstVector<double>& arr, ConstVector<double>& att,
ConstVector<double>& art, ConstVector<double>& detDF,
ConstVector<double>& coeff_beta)
{
assert(i_r >= 0 && i_r <= grid.numberSmootherCircles());

Expand Down Expand Up @@ -185,11 +186,11 @@ static inline void nodeApplyAscOrthoCircleTake(int i_r, int i_theta, const Polar
}
}

static inline void nodeApplyAscOrthoRadialTake(int i_r, int i_theta, const PolarGrid& grid, bool DirBC_Interior,
ConstVector<double>& x, ConstVector<double>& rhs, Vector<double>& result,
ConstVector<double>& arr, const ConstVector<double>& att,
ConstVector<double>& art, const ConstVector<double>& detDF,
ConstVector<double>& coeff_beta)
static KOKKOS_INLINE_FUNCTION void
nodeApplyAscOrthoRadialTake(int i_r, int i_theta, const PolarGrid& grid, bool DirBC_Interior, ConstVector<double>& x,
ConstVector<double>& rhs, Vector<double>& result, ConstVector<double>& arr,
const ConstVector<double>& att, ConstVector<double>& art, const ConstVector<double>& detDF,
ConstVector<double>& coeff_beta)
{
assert(i_r >= grid.numberSmootherCircles() - 1 && i_r < grid.nr());

Expand Down Expand Up @@ -483,14 +484,22 @@ void ExtrapolatedSmootherTake<LevelCacheType>::applyAscOrthoBlackCircleSection(C

/* The outer most circle next to the radial section is defined to be black. */
const int start_black_circles = (grid.numberSmootherCircles() % 2 == 0) ? 1 : 0;

#pragma omp parallel for num_threads(num_omp_threads)
for (int i_r = start_black_circles; i_r < grid.numberSmootherCircles(); i_r += 2) {
for (int i_theta = 0; i_theta < grid.ntheta(); i_theta++) {
const int num_black_circles = (grid.numberSmootherCircles() - start_black_circles + 1) / 2;

Kokkos::parallel_for(
"ExtrapolatedSmootherTake: ApplyAscOrtho (Black Circular)",
Kokkos::MDRangePolicy<Kokkos::DefaultHostExecutionSpace, Kokkos::Rank<2>>( // Rank of the index space
{0, 0}, // Starting point of the index space
{num_black_circles, grid.ntheta()} // Ending point of the index space
),
// Kokkos lambda function to execute for each point in the index space
KOKKOS_LAMBDA(const int circle_task, const int i_theta) {
int i_r = start_black_circles + circle_task * 2;
nodeApplyAscOrthoCircleTake(i_r, i_theta, grid, DirBC_Interior, x, rhs, temp, arr, att, art, detDF,
coeff_beta);
}
}
});

Kokkos::fence();
}

template <class LevelCacheType>
Expand All @@ -516,14 +525,22 @@ void ExtrapolatedSmootherTake<LevelCacheType>::applyAscOrthoWhiteCircleSection(C

/* The outer most circle next to the radial section is defined to be black. */
const int start_white_circles = (grid.numberSmootherCircles() % 2 == 0) ? 0 : 1;

#pragma omp parallel for num_threads(num_omp_threads)
for (int i_r = start_white_circles; i_r < grid.numberSmootherCircles(); i_r += 2) {
for (int i_theta = 0; i_theta < grid.ntheta(); i_theta++) {
const int num_white_circles = (grid.numberSmootherCircles() - start_white_circles + 1) / 2;

Kokkos::parallel_for(
"ExtrapolatedSmootherTake: ApplyAscOrtho (White Circular)",
Kokkos::MDRangePolicy<Kokkos::DefaultHostExecutionSpace, Kokkos::Rank<2>>( // Rank of the index space
{0, 0}, // Starting point of the index space
{num_white_circles, grid.ntheta()} // Ending point of the index space
),
// Kokkos lambda function to execute for each point in the index space
KOKKOS_LAMBDA(const int circle_task, const int i_theta) {
const int i_r = start_white_circles + circle_task * 2;
nodeApplyAscOrthoCircleTake(i_r, i_theta, grid, DirBC_Interior, x, rhs, temp, arr, att, art, detDF,
coeff_beta);
}
}
});

Kokkos::fence();
}

template <class LevelCacheType>
Expand All @@ -547,13 +564,24 @@ void ExtrapolatedSmootherTake<LevelCacheType>::applyAscOrthoBlackRadialSection(C
ConstVector<double> detDF = level_cache.detDF();
ConstVector<double> coeff_beta = level_cache.coeff_beta();

#pragma omp parallel for num_threads(num_omp_threads)
for (int i_theta = 0; i_theta < grid.ntheta(); i_theta += 2) {
for (int i_r = grid.numberSmootherCircles(); i_r < grid.nr(); i_r++) {
assert(grid.ntheta() % 2 == 0);
const int start_black_radials = 0;
const int num_black_radial_lines = grid.ntheta() / 2;

Kokkos::parallel_for(
"ExtrapolatedSmootherTake: ApplyAscOrtho (Black Radial)",
Kokkos::MDRangePolicy<Kokkos::DefaultHostExecutionSpace, Kokkos::Rank<2>>( // Rank of the index space
{0, grid.numberSmootherCircles()}, // Starting point of the index space
{num_black_radial_lines, grid.nr()} // Ending point of the index space
),
// Kokkos lambda function to execute for each point in the index space
KOKKOS_LAMBDA(const int radial_task, const int i_r) {
const int i_theta = start_black_radials + radial_task * 2;
nodeApplyAscOrthoRadialTake(i_r, i_theta, grid, DirBC_Interior, x, rhs, temp, arr, att, art, detDF,
coeff_beta);
}
}
});

Kokkos::fence();
}

template <class LevelCacheType>
Expand All @@ -577,11 +605,22 @@ void ExtrapolatedSmootherTake<LevelCacheType>::applyAscOrthoWhiteRadialSection(C
ConstVector<double> detDF = level_cache.detDF();
ConstVector<double> coeff_beta = level_cache.coeff_beta();

#pragma omp parallel for num_threads(num_omp_threads)
for (int i_theta = 1; i_theta < grid.ntheta(); i_theta += 2) {
for (int i_r = grid.numberSmootherCircles(); i_r < grid.nr(); i_r++) {
assert(grid.ntheta() % 2 == 0);
const int start_white_radials = 1;
const int num_white_radial_lines = grid.ntheta() / 2;

Kokkos::parallel_for(
"ExtrapolatedSmootherTake: ApplyAscOrtho (White Radial)",
Kokkos::MDRangePolicy<Kokkos::DefaultHostExecutionSpace, Kokkos::Rank<2>>( // Rank of the index space
{0, grid.numberSmootherCircles()}, // Starting point of the index space
{num_white_radial_lines, grid.nr()} // Ending point of the index space
),
// Kokkos lambda function to execute for each point in the index space
KOKKOS_LAMBDA(const int radial_task, const int i_r) {
const int i_theta = start_white_radials + radial_task * 2;
nodeApplyAscOrthoRadialTake(i_r, i_theta, grid, DirBC_Interior, x, rhs, temp, arr, att, art, detDF,
coeff_beta);
}
}
}
});

Kokkos::fence();
}
4 changes: 2 additions & 2 deletions include/ExtrapolatedSmoother/ExtrapolatedSmootherTake/extrapolatedSmootherTake.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,8 @@ class ExtrapolatedSmootherTake : public ExtrapolatedSmoother<LevelCacheType>
/* Orthogonal application */
/* ---------------------- */

// Functions must be public due to cuda restriction
public:
// Compute temp = f_sc − A_sc^ortho * u_sc^ortho (precomputed right-hand side)
// where x = u_sc and rhs = f_sc
void applyAscOrthoBlackCircleSection(ConstVector<double> x, ConstVector<double> rhs, Vector<double> temp);
Expand All @@ -183,8 +185,6 @@ class ExtrapolatedSmootherTake : public ExtrapolatedSmoother<LevelCacheType>
/* Line-wise solvers */
/* ----------------- */

// Functions must be public due to cuda restriction
public:
// Solve the linear system:
// A_sc * u_sc = f_sc − A_sc^ortho * u_sc^ortho
// Parameter mapping:
Expand Down
20 changes: 11 additions & 9 deletions include/Smoother/SmootherGive/applyAscOrtho.inl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,11 @@
namespace smoother_give
{

static inline void nodeApplyAscOrthoCircleGive(int i_r, int i_theta, const PolarGrid& grid, bool DirBC_Interior,
SmootherColor smoother_color, ConstVector<double>& x,
ConstVector<double>& rhs, Vector<double>& result, double arr, double att,
double art, double detDF, double coeff_beta)
static KOKKOS_INLINE_FUNCTION void nodeApplyAscOrthoCircleGive(int i_r, int i_theta, const PolarGrid& grid,
bool DirBC_Interior, SmootherColor smoother_color,
ConstVector<double>& x, ConstVector<double>& rhs,
Vector<double>& result, double arr, double att,
double art, double detDF, double coeff_beta)
{
assert(i_r >= 0 && i_r <= grid.numberSmootherCircles());

Expand Down Expand Up @@ -187,10 +188,11 @@ static inline void nodeApplyAscOrthoCircleGive(int i_r, int i_theta, const Polar
}
}

static inline void nodeApplyAscOrthoRadialGive(int i_r, int i_theta, const PolarGrid& grid, bool DirBC_Interior,
SmootherColor smoother_color, ConstVector<double>& x,
ConstVector<double>& rhs, Vector<double>& result, double arr, double att,
double art, double detDF, double coeff_beta)
static KOKKOS_INLINE_FUNCTION void nodeApplyAscOrthoRadialGive(int i_r, int i_theta, const PolarGrid& grid,
bool DirBC_Interior, SmootherColor smoother_color,
ConstVector<double>& x, ConstVector<double>& rhs,
Vector<double>& result, double arr, double att,
double art, double detDF, double coeff_beta)
{
assert(i_r >= grid.numberSmootherCircles() - 1 && i_r < grid.nr());

Expand Down Expand Up @@ -479,4 +481,4 @@ void SmootherGive<LevelCacheType>::applyAscOrthoRadialSection(const int i_theta,
nodeApplyAscOrthoRadialGive(i_r, i_theta, grid, DirBC_Interior, smoother_color, x, rhs, temp, arr, att, art,
detDF, coeff_beta);
}
}
}
Loading