Skip to content

Commit

Permalink
changes to make things work on gpu
Browse files Browse the repository at this point in the history
  • Loading branch information
@lroberts36
lroberts36 committed Oct 5, 2023
1 parent 8e5a6a7 commit 49d00ef
Show file tree
Hide file tree
Showing 3 changed files with 3 additions and 203 deletions.
1 change: 0 additions & 1 deletion example/poisson_gmg/poisson_equation.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,6 @@ class PoissonEquation {
return TaskStatus::complete;
}

private:
template <class var_t, bool only_md_level = false>
static parthenon::TaskStatus
CalculateFluxes(std::shared_ptr<parthenon::MeshData<Real>> &md) {
Expand Down
199 changes: 0 additions & 199 deletions example/poisson_gmg/poisson_package.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -88,205 +88,6 @@ TaskStatus PrintChosenValues(std::shared_ptr<MeshData<Real>> &md,
return TaskStatus::complete;
}

class flux_poisson {
public:
// std::vector<parthenon::Metadata> RequiredMetadata

template <class x_t, class out_t, bool only_md_level = false, class TL_t>
parthenon::TaskID Ax(TL_t &tl, parthenon::TaskID depends_on,
std::shared_ptr<parthenon::MeshData<Real>> &md, bool only_interior,
bool do_flux_cor = false) {
auto flux_res = tl.AddTask(depends_on, CalculateFluxes<x_t, only_md_level>, md);
if (do_flux_cor && !only_md_level) {
auto start_flxcor =
tl.AddTask(flux_res, parthenon::StartReceiveFluxCorrections, md);
auto send_flxcor = tl.AddTask(flux_res, parthenon::LoadAndSendFluxCorrections, md);
auto recv_flxcor = tl.AddTask(send_flxcor, parthenon::ReceiveFluxCorrections, md);
flux_res = tl.AddTask(recv_flxcor, parthenon::SetFluxCorrections, md);
}
return tl.AddTask(flux_res, FluxMultiplyMatrix<x_t, out_t, only_md_level>, md,
only_interior);
}

template <class diag_t, bool only_md_level = false>
parthenon::TaskStatus SetDiagonal(std::shared_ptr<parthenon::MeshData<Real>> &md) {
using namespace parthenon;
const int ndim = md->GetMeshPointer()->ndim;
using TE = parthenon::TopologicalElement;
TE te = TE::CC;
IndexRange ib = md->GetBoundsI(IndexDomain::interior, te);
IndexRange jb = md->GetBoundsJ(IndexDomain::interior, te);
IndexRange kb = md->GetBoundsK(IndexDomain::interior, te);

auto pkg = md->GetMeshPointer()->packages.Get("poisson_package");
const auto alpha = pkg->Param<Real>("diagonal_alpha");

int nblocks = md->NumBlocks();
std::vector<bool> include_block(nblocks, true);

if (only_md_level) {
for (int b = 0; b < nblocks; ++b)
include_block[b] = (md->grid.logical_level ==
md->GetBlockData(b)->GetBlockPointer()->loc.level());
}

auto desc = parthenon::MakePackDescriptor<diag_t, D>(md.get());
auto pack = desc.GetPack(md.get(), include_block);
parthenon::par_for(
DEFAULT_LOOP_PATTERN, "StoreDiagonal", DevExecSpace(), 0, pack.GetNBlocks() - 1,
kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
const auto &coords = pack.GetCoordinates(b);
// Build the unigrid diagonal of the matrix
Real dx1 = coords.template Dxc<X1DIR>(k, j, i);
Real diag_elem =
-(pack(b, TE::F1, D(), k, j, i) + pack(b, TE::F1, D(), k, j, i + 1)) /
(dx1 * dx1) -
alpha;
if (ndim > 1) {
Real dx2 = coords.template Dxc<X2DIR>(k, j, i);
diag_elem -=
(pack(b, TE::F2, D(), k, j, i) + pack(b, TE::F2, D(), k, j + 1, i)) /
(dx2 * dx2);
}
if (ndim > 2) {
Real dx3 = coords.template Dxc<X3DIR>(k, j, i);
diag_elem -=
(pack(b, TE::F3, D(), k, j, i) + pack(b, TE::F3, D(), k + 1, j, i)) /
(dx3 * dx3);
}
pack(b, te, diag_t(), k, j, i) = diag_elem;
});
return TaskStatus::complete;
}

private:
template <class var_t, bool only_md_level = false>
static parthenon::TaskStatus
CalculateFluxes(std::shared_ptr<parthenon::MeshData<Real>> &md) {
using namespace parthenon;
const int ndim = md->GetMeshPointer()->ndim;
using TE = parthenon::TopologicalElement;
TE te = TE::CC;
IndexRange ib = md->GetBoundsI(IndexDomain::interior, te);
IndexRange jb = md->GetBoundsJ(IndexDomain::interior, te);
IndexRange kb = md->GetBoundsK(IndexDomain::interior, te);

using TE = parthenon::TopologicalElement;

int nblocks = md->NumBlocks();
std::vector<bool> include_block(nblocks, true);

if (only_md_level) {
for (int b = 0; b < nblocks; ++b)
include_block[b] = (md->grid.logical_level ==
md->GetBlockData(b)->GetBlockPointer()->loc.level());
}

auto desc =
parthenon::MakePackDescriptor<var_t, D>(md.get(), {}, {PDOpt::WithFluxes});
auto pack = desc.GetPack(md.get(), include_block);
parthenon::par_for(
DEFAULT_LOOP_PATTERN, "CaclulateFluxes", DevExecSpace(), 0, pack.GetNBlocks() - 1,
kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
const auto &coords = pack.GetCoordinates(b);
Real dx1 = coords.template Dxc<X1DIR>(k, j, i);
pack.flux(b, X1DIR, var_t(), k, j, i) =
pack(b, TE::F1, D(), k, j, i) / dx1 *
(pack(b, te, var_t(), k, j, i - 1) - pack(b, te, var_t(), k, j, i));
if (i == ib.e)
pack.flux(b, X1DIR, var_t(), k, j, i + 1) =
pack(b, TE::F1, D(), k, j, i + 1) / dx1 *
(pack(b, te, var_t(), k, j, i) - pack(b, te, var_t(), k, j, i + 1));

if (ndim > 1) {
Real dx2 = coords.template Dxc<X2DIR>(k, j, i);
pack.flux(b, X2DIR, var_t(), k, j, i) =
pack(b, TE::F2, D(), k, j, i) *
(pack(b, te, var_t(), k, j - 1, i) - pack(b, te, var_t(), k, j, i)) / dx2;
if (j == jb.e)
pack.flux(b, X2DIR, var_t(), k, j + 1, i) =
pack(b, TE::F2, D(), k, j + 1, i) *
(pack(b, te, var_t(), k, j, i) - pack(b, te, var_t(), k, j + 1, i)) /
dx2;
}

if (ndim > 2) {
Real dx3 = coords.template Dxc<X3DIR>(k, j, i);
pack.flux(b, X3DIR, var_t(), k, j, i) =
pack(b, TE::F3, D(), k, j, i) *
(pack(b, te, var_t(), k - 1, j, i) - pack(b, te, var_t(), k, j, i)) / dx3;
if (k == kb.e)
pack.flux(b, X2DIR, var_t(), k + 1, j, i) =
pack(b, TE::F3, D(), k + 1, j, i) *
(pack(b, te, var_t(), k, j, i) - pack(b, te, var_t(), k + 1, j, i)) /
dx3;
}
});
return TaskStatus::complete;
}

template <class in_t, class out_t, bool only_md_level = false>
static parthenon::TaskStatus
FluxMultiplyMatrix(std::shared_ptr<parthenon::MeshData<Real>> &md, bool only_interior) {
using namespace parthenon;
const int ndim = md->GetMeshPointer()->ndim;
using TE = parthenon::TopologicalElement;
TE te = TE::CC;
IndexRange ib = md->GetBoundsI(IndexDomain::interior, te);
IndexRange jb = md->GetBoundsJ(IndexDomain::interior, te);
IndexRange kb = md->GetBoundsK(IndexDomain::interior, te);

auto pkg = md->GetMeshPointer()->packages.Get("poisson_package");
const auto alpha = pkg->Param<Real>("diagonal_alpha");

int nblocks = md->NumBlocks();
std::vector<bool> include_block(nblocks, true);
if (only_interior) {
for (int b = 0; b < nblocks; ++b)
include_block[b] = md->GetBlockData(b)->GetBlockPointer()->neighbors.size() == 0;
}

if (only_md_level) {
for (int b = 0; b < nblocks; ++b)
include_block[b] =
include_block[b] && (md->grid.logical_level ==
md->GetBlockData(b)->GetBlockPointer()->loc.level());
}

auto desc =
parthenon::MakePackDescriptor<in_t, out_t>(md.get(), {}, {PDOpt::WithFluxes});
auto pack = desc.GetPack(md.get(), include_block);
parthenon::par_for(
DEFAULT_LOOP_PATTERN, "FluxMultiplyMatrix", DevExecSpace(), 0,
pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, ib.s, ib.e,
KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) {
const auto &coords = pack.GetCoordinates(b);
Real dx1 = coords.template Dxc<X1DIR>(k, j, i);
pack(b, te, out_t(), k, j, i) = -alpha * pack(b, te, in_t(), k, j, i);
pack(b, te, out_t(), k, j, i) += (pack.flux(b, X1DIR, in_t(), k, j, i) -
pack.flux(b, X1DIR, in_t(), k, j, i + 1)) /
dx1;

if (ndim > 1) {
Real dx2 = coords.template Dxc<X2DIR>(k, j, i);
pack(b, te, out_t(), k, j, i) += (pack.flux(b, X2DIR, in_t(), k, j, i) -
pack.flux(b, X2DIR, in_t(), k, j + 1, i)) /
dx2;
}

if (ndim > 2) {
Real dx3 = coords.template Dxc<X3DIR>(k, j, i);
pack(b, te, out_t(), k, j, i) += (pack.flux(b, X3DIR, in_t(), k, j, i) -
pack.flux(b, X3DIR, in_t(), k + 1, j, i)) /
dx3;
}
});
return TaskStatus::complete;
}
};

} // namespace poisson_package

#endif // EXAMPLE_POISSON_GMG_POISSON_PACKAGE_HPP_
6 changes: 3 additions & 3 deletions src/solvers/mg_solver.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -128,15 +128,16 @@ class MGSolver {
int GetCurrentIterations() const { return iter_counter; }
Real GetFinalResidual() const { return final_residual; }
int GetFinalIterations() const { return final_iteration; }

protected:
MGParams params_;
int iter_counter;
AllReduce<Real> residual;
equations eqs_;
Real final_residual;
int final_iteration;

// These functions apparently have to be public to compile with cuda since
// they contain device side lambdas
public:
enum class GSType { all, red, black };

template <class rhs_t, class Axold_t, class D_t, class xold_t, class xnew_t,
Expand Down Expand Up @@ -253,7 +254,6 @@ class MGSolver {
PARTHENON_FAIL("Unknown solver type.");
}

int ndim = pmesh->ndim;
bool multilevel = (level != min_level);
TaskID last_task;

Expand Down

0 comments on commit 49d00ef

Please sign in to comment.