Skip to content

Adapt weight for 1D Poisson solver #853

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
mkstoyanov merged 10 commits into from
Aug 11, 2026
Merged

Adapt weight for 1D Poisson solver #853

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
81f6ef9
Updated Doxygen documentation for two_stream_md
TTBurnett Jul 30, 2026
5fe1c9e
asgard_io.cpp now converts the electric field full level to raw data …
TTBurnett Aug 4, 2026
3f4c76c
Fixed poisson_1d interpolation on GPU
TTBurnett Aug 4, 2026
f4314e6
Bug fixes.
TTBurnett Aug 5, 2026
3b3d61a
commit so Miro can look at the code
TTBurnett Aug 6, 2026
bda089f
Misc cleanup.
TTBurnett Aug 6, 2026
0d76a64
Interpolation is only done for 1D electric field moments when needed.…
TTBurnett Aug 6, 2026
c145274
PR fixes
TTBurnett Aug 6, 2026
0807dec
Poisson solver is now setup on the root node if needed
TTBurnett Aug 10, 2026
16e978c
More PR fixes.
TTBurnett Aug 10, 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
1 change: 1 addition & 0 deletions doxygen/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@ doxygen_add_docs(asgard_doxygen
src/pde/spherical_diffusion.cpp
src/pde/sinwav.cpp
src/pde/two_stream.cpp
src/pde/two_stream_md.cpp
src/pde/vplb.cpp
src/pde/sod_shock.cpp
src/pde/burgers.cpp
Expand Down
2 changes: 1 addition & 1 deletion src/asgard_discretization.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -414,7 +414,7 @@ std::vector<precision> discretization_manager<precision>::get_moment(moment_id i
else
throw std::runtime_error("an electric moment was requested but no poisson solver is set");
}, get_poisson());
terms.moms.solve_poisson(terms.grid, terms.conn, terms.hier, terms.interp, terms.kwork, true);
terms.moms.solve_poisson(terms.grid, terms.conn, terms.hier, terms.interp, terms.kwork);
return terms.moms.get_cached_raw(id);
}
std::vector<precision> result;
Expand Down
2 changes: 1 addition & 1 deletion src/asgard_gpu_tensors_cuda.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -473,7 +473,7 @@ __global__ void kernel_moment_expand(int pos_block, int vel_block,
} else
for (int c = 0; c < num_cycles; c++) {
P const src = pos_data[i * pos_block + (threadIdx.x + c * cycle_size) / vel_block];
vals[j * pos_block * vel_block + threadIdx.x + (threadIdx.x + c * cycle_size)] = src;
vals[j * pos_block * vel_block + (threadIdx.x + c * cycle_size)] = src;
}

blk += gridDim.x * blockDim.y;
Expand Down
2 changes: 1 addition & 1 deletion src/asgard_gpu_tensors_rocm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -473,7 +473,7 @@ __global__ void kernel_moment_expand(int pos_block, int vel_block,
} else
for (int c = 0; c < num_cycles; c++) {
P const src = pos_data[i * pos_block + (threadIdx.x + c * cycle_size) / vel_block];
vals[j * pos_block * vel_block + threadIdx.x + (threadIdx.x + c * cycle_size)] = src;
vals[j * pos_block * vel_block + (threadIdx.x + c * cycle_size)] = src;
}

blk += gridDim.x * blockDim.y;
Expand Down
8 changes: 2 additions & 6 deletions src/asgard_io.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,8 +141,7 @@ void h5manager<P>::write(prog_opts const &options, pde_domain<P> const &domain,
if (moms.has_poisson()) {
moment_id m0 = moms.find_id(moment::zero(moms.num_vel()));
moms.cache_moment(m0, grid, state);
constexpr bool raw_on_cpu = true;
moms.solve_poisson(terms.grid, terms.conn, terms.hier, terms.interp, terms.kwork, raw_on_cpu);
moms.solve_poisson(terms.grid, terms.conn, terms.hier, terms.interp, terms.kwork);
}
for (int i : iindexof(moms.num_moments()))
{
Expand All @@ -152,10 +151,7 @@ void h5manager<P>::write(prog_opts const &options, pde_domain<P> const &domain,
H5Easy::dump(file, auxstr.name,
std::string("__moment_") + moms.get_by_id(mid).to_string());
if (moms.needs_poisson(mid)) {
if (domain.num_pos() == 1)
write_vector(auxstr.data, moms.get_cached_level(mid)); // only the level is computed for poisson_1d
else
write_vector(auxstr.data, moms.get_cached_raw(mid));
write_vector(auxstr.data, moms.get_cached_raw(mid, terms.hier));
} else {
std::vector<P> vals;
moms.mcompute(grid, mid, state, vals);
Expand Down
219 changes: 109 additions & 110 deletions src/asgard_moment_manager.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -559,17 +559,19 @@ void moment_manager<P>::cache_moments(
interpolation_manager<P> const &interp, kronmult::workspace<P> &work) const
{
// Compute moments on hierarchical grid
bool has_poisson = false;
if (group == group_id::all()) { // do all moments
tools::time_event performance_("cache all moments");
for (auto mid : raw_moments_) {
if (mid == moment_id::unset() or needs_poisson(mid)) continue;
mcompute(grid, mid, state, raw_vals.get(mid));
if (mid == moment_id::unset()) continue;
if (needs_poisson(mid))
has_poisson = true;
else
mcompute(grid, mid, state, raw_vals.get(mid));
full_level.get(mid).resize(0);
}
solve_poisson(grid, conn, hier, interp, work);
} else {
tools::time_event performance_("cache moments (" + std::to_string(group()) + ")");
bool has_poisson = false;
for (auto mid = first_in(group, raw_moments_);
*mid != moment_id::unset(); mid++) {
if (needs_poisson(*mid))
Expand All @@ -578,8 +580,8 @@ void moment_manager<P>::cache_moments(
mcompute(grid, *mid, state, raw_vals.get(*mid));
full_level.get(*mid).resize(0);
}
if (has_poisson) solve_poisson(grid, conn, hier, interp, work);
}
if (has_poisson) solve_poisson(grid, conn, hier, interp, work);

// Poisson moments are already interpolated if they were computed on GPU
auto skip_poisson = [&](moment_id mid) -> bool {
Expand All @@ -596,13 +598,13 @@ void moment_manager<P>::cache_moments(
if (group == group_id::all()) {
for (auto mid : interp_moments_) {
if (skip_poisson(mid) or mid == moment_id::unset()) continue;
make_nodal(mid, interp, work, interp.it1);
make_nodal(mid, interp, hier, work, interp.it1);
}
} else {
for (auto mid = first_in(group, interp_moments_);
*mid != moment_id::unset(); mid++) {
if (skip_poisson(*mid)) continue;
make_nodal(*mid, interp, work, interp.it1);
make_nodal(*mid, interp, hier, work, interp.it1);
}
}
interp.it1.resize(num_entries);
Expand Down Expand Up @@ -641,99 +643,96 @@ void moment_manager<P>::complete_level(hierarchy_manipulator<P> const &hier,
template<typename P>
void moment_manager<P>::solve_poisson(sparse_grid const &grid, connection_patterns const &conn,
hierarchy_manipulator<P> const &hier, interpolation_manager<P> const &interp,
kronmult::workspace<P> &work, bool raw_on_cpu) const
kronmult::workspace<P> &work) const
{
std::visit([&](auto &p) {
if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<P>>) {
p.solve_periodic(get_cached_level(p.moment0(), hier), full_level.get(p.moment_electric()));
} else if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_md<P>>) {
#ifdef ASGARD_USE_GPU
compute->set_device(gpu::device{0});
int const num_entries = pos_grid.num_dof();
std::array<gpu::vector<P>, max_num_gpus> &work1 = interp.gpu_it1;
assert(work1[0].size() >= num_entries);
work1[0].copy_from_host(num_entries, raw_vals.get(p.moment0()).data());
for (moment_id const &mid : p.moments_electric()) {
if (mid == moment_id::unset()) continue;
gpu::vector<P> &res = gpu_interps[0][mid];
res.resize(full_block * grid.num_indexes());
}
solve_poisson_gpu(conn, hier, interp, work, raw_on_cpu);
auto find_mom = [this](moment_id mid) -> const mom_on_gpu &
{
for (auto const &mom : this->gpu_moments[0]) {
if (mom.mid == mid)
return mom;
}
throw std::runtime_error("could not find the moment");
};
for (moment_id const &mid : p.moments_electric()) {
if (mid == moment_id::unset()) continue;
mom_on_gpu const &mom = find_mom(mid);
if (mom.interp_on_cpu()) continue; // solve_poisson_gpu will already copy if interp_on_cpu() is true
gpu::vector<P> &res = gpu_interps[0][mid];
res.copy_to_host(interps.get(mid));
}
#else
std::ignore = interp;
std::ignore = grid;
std::ignore = raw_on_cpu;
update_position_grid_dsort();
p.solve_periodic(raw_vals.get(p.moment0()), raw_vals, pos_grid, conn, work);
#endif
}
}, poisson_solver);
if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<P>>) {
p.solve_periodic(get_cached_level(p.moment0(), hier), full_level.get(p.moment_electric()));
} else if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_md<P>>) {
#ifdef ASGARD_USE_GPU
compute->set_device(gpu::device{0});

// transfer density to GPU
int const num_entries = pos_grid.num_dof();
std::array<gpu::vector<P>, max_num_gpus> &work1 = interp.gpu_it1;
Comment thread
TTBurnett marked this conversation as resolved.
Outdated
assert(work1[0].size() >= num_entries);
work1[0].copy_from_host(num_entries, raw_vals.get(p.moment0()).data());
Comment thread
TTBurnett marked this conversation as resolved.
Outdated

// solve
constexpr bool result_to_cpu = true;
solve_poisson_gpu(gpu::device{0}, grid, conn, hier, interp, work, result_to_cpu);
#else
std::ignore = interp;
std::ignore = grid;
update_position_grid_dsort();
p.solve_periodic(raw_vals.get(p.moment0()), raw_vals, pos_grid, conn, work);
#endif
}
}, poisson_solver);
}

#ifdef ASGARD_USE_GPU
template<typename P>
void moment_manager<P>::solve_poisson_gpu(connection_patterns const &conn, hierarchy_manipulator<P> const &hier,
interpolation_manager<P> const &interp, kronmult::workspace<P> &work, bool raw_on_cpu) const
template<bool always_interp>
void moment_manager<P>::solve_poisson_gpu(gpu::device dev, sparse_grid const &grid, connection_patterns const &conn, hierarchy_manipulator<P> const &hier,
interpolation_manager<P> const &interp, kronmult::workspace<P> &work, bool result_to_cpu) const
{
// Setup
int64_t const num_entries = pos_grid.num_dof();
std::array<gpu::vector<P>, max_num_gpus> &work1 = interp.gpu_it1;
std::array<gpu::vector<P>, max_num_gpus> &work2 = interp.gpu_it2;
assert(work1[dev.id].size() >= num_entries);
assert(work2[dev.id].size() >= num_entries);
gpu::wrap_array<P> w1(work1[dev.id].data(), num_entries);
gpu::wrap_array<P> w2(work2[dev.id].data(), num_entries);

std::visit([&](auto &p) {
if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<P>>) {
// This is fast enough on CPU
int64_t const num_entries = pos_grid.num_dof();
std::array<gpu::vector<P>, max_num_gpus> &work1 = interp.gpu_it1;
assert(work1[0].size() >= num_entries);
gpu::wrap_array<P> w1(work1[0].data(), num_entries);
w1.vec.copy_to_host(raw_vals[p.moment0()]);
p.solve_periodic(get_cached_level(p.moment0(), hier), full_level.get(p.moment_electric()));
} else if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_md<P>>) {
// Setup
update_position_grid_dsort();
int64_t const num_entries = pos_grid.num_dof();
std::array<gpu::vector<P>, max_num_gpus> &work1 = interp.gpu_it1;
std::array<gpu::vector<P>, max_num_gpus> &work2 = interp.gpu_it2;
assert(work1[0].size() >= num_entries);
assert(work2[0].size() >= num_entries);
gpu::wrap_array<P> w1(work1[0].data(), num_entries);
gpu::wrap_array<P> w2(work2[0].data(), num_entries);
auto find_mom = [this](moment_id mid) -> const mom_on_gpu &
{
for (auto const &mom : this->gpu_moments[0]) {
if (mom.mid == mid)
return mom;
}
throw std::runtime_error("could not find the moment");
};
auto interp_func = [&](gpu::vector<P> const &efield, moment_id mid) -> void
{
if (raw_on_cpu) {
efield.copy_to_host(this->raw_vals[mid]);
this->full_level.get(mid).resize(0);
}
interp.pos2nodal(gpu::device{0}, this->pos_grid, efield.data(), this->wav_scale, w2.vec.data(), work);
gpu::vector<P> &res = this->gpu_interps[0][mid];
moment_expand(pdof, this->pos_grid.num_dims(), num_vel_, reduce_ij[0], w2.vec, res);
mom_on_gpu const &mom = find_mom(mid);
if (mom.interp_on_cpu()) res.copy_to_host(interps[mom.mid]);
};
// Solve
p.solve_periodic(w1.vec, pos_grid, conn, interp_func, work);
if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<P>>) {
// Solve on CPU
moment_id const m0 = p.moment0();
moment_id const melectric = p.moment_electric();
w1.vec.copy_to_host(raw_vals[m0]);
p.solve_periodic(get_cached_level(m0, hier), full_level.get(melectric));

// interpolate
if constexpr (always_interp) {
cache_raw_from_level(melectric, hier);
w1.vec.copy_from_host(raw_vals[melectric].size(), raw_vals[melectric].data());
interp.pos2nodal(dev, pos_grid, w1.vec.data(), wav_scale, w2.vec.data(), work);
gpu::vector<P> &res = gpu_interps[dev.id][melectric];
res.resize(full_block * grid.num_indexes());
moment_expand(pdof, pos_grid.num_dims(), num_vel_, reduce_ij[dev.id], w2.vec, res);
if (result_to_cpu) res.copy_to_host(interps[melectric]);
}
},
poisson_solver);
} else if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_md<P>>) {
// Setup
update_position_grid_dsort();
auto find_mom = [this](moment_id mid) -> const mom_on_gpu &
{
for (auto const &mom : this->gpu_moments[0]) {
if (mom.mid == mid)
return mom;
}
throw std::runtime_error("could not find the moment");
};
auto interp_func = [&](gpu::vector<P> const &efield, moment_id mid) -> void
{
mom_on_gpu const &mom = find_mom(mid);
if (result_to_cpu or mom.raw_on_cpu()) {
efield.copy_to_host(this->raw_vals[mid]);
this->full_level.get(mid).resize(0);
}
interp.pos2nodal(dev, this->pos_grid, efield.data(), this->wav_scale, w2.vec.data(), work);
gpu::vector<P> &res = this->gpu_interps[dev.id][mid];
res.resize(this->full_block * grid.num_indexes());
moment_expand(pdof, this->pos_grid.num_dims(), num_vel_, reduce_ij[dev.id], w2.vec, res);
if (result_to_cpu or mom.interp_on_cpu())
res.copy_to_host(interps[mid]);
};
// Solve
p.solve_periodic(w1.vec, pos_grid, conn, interp_func, work);
}
}, poisson_solver);
}
#endif // ASGARD_USE_GPU

Expand All @@ -751,12 +750,14 @@ void moment_manager<P>::update_position_grid_dsort() const
template<typename P>
void moment_manager<P>::make_nodal(
moment_id id, interpolation_manager<P> const &interp,
hierarchy_manipulator<P> const &hier,
kronmult::workspace<P> &kwork, std::vector<P> &workspace) const
{
assert(not (num_pos_ == 1 and needs_poisson(id))); // poisson_1d solves the moment for the full level
// so there is no raw moment to interpolate
update_position_grid_dsort();

if (num_pos_ == 1 and needs_poisson(id))
cache_raw_from_level(id, hier);

interp.pos2nodal(pos_grid, raw_vals[id].data(), wav_scale, workspace, kwork);

interps[id].resize(pntr.back() * full_block);
Expand Down Expand Up @@ -799,7 +800,7 @@ void moment_manager<P>::compute_interps(
solve_poisson(grid, conn, hier, interp, work);
}
for (auto const &id : ids)
make_nodal(id, interp, work, interp.it1);
make_nodal(id, interp, hier, work, interp.it1);
interp.it1.resize(num_entries);
}

Expand Down Expand Up @@ -954,12 +955,7 @@ void moment_manager<P>::prepare_pos_grid_gpu(group_id group, sparse_grid const &

if (not has_interp.empty()) {
if (group == group_id::all() or has_interp[group()]) {
if (dsort_generation != pos_grid.generation()) {
pos_grid.dsort_ = dimension_sort(pos_grid.iset_);
bool constexpr skip_indexes = true; // already loaded in reduce_grid()
pos_grid.gpu_sync<skip_indexes>();
dsort_generation = pos_grid.generation();
}
update_position_grid_dsort();
}
}
}
Expand Down Expand Up @@ -1012,13 +1008,13 @@ void moment_manager<P>::compute_moments(

moment const mom = mlist[im->mid]; // using this to get the necessary powers

if (mom.is_electric()) continue; // skip electric field moments

gpu::vector<P> &res = gpu_interps[g][im->mid];
if (not im->skip_interp()) {
res.resize(full_block * grid.num_indexes());
}

if (mom.is_electric()) continue; // skip electric field moments

std::array<P const *, max_mom_dims> itg =
{gpu_integ[g][0].data() + mom.pows[0] * integ[0].stride(), nullptr, nullptr};
for (int i = 1; i < num_vel_; i++)
Expand Down Expand Up @@ -1048,10 +1044,8 @@ void moment_manager<P>::compute_moments(
}

// solve poisson equation while w1 holds density
if (mom.is_zero()) {
assert(g == 0); // poisson uses gpu 0 for solve
solve_poisson_gpu(conn, hier, interp, kwork);
}
if (mom.is_zero())
solve_poisson_gpu(gpu::device{g}, grid, conn, hier, interp, kwork);

if (im->skip_interp()) continue;

Expand Down Expand Up @@ -1083,6 +1077,7 @@ void moment_manager<P>::compute_moments(
// which will also allow to avoid the spread when doing term-apply

prepare_pos_grid_gpu(group_id::all(), grid);
update_position_grid_dsort(); // interpolation is always done in this overload of compute_moments()

int const pos_block = pos_grid.block_size();

Expand All @@ -1108,11 +1103,11 @@ void moment_manager<P>::compute_moments(
{
moment const mom = mlist[im]; // using this to get the necessary powers

if (mom.is_electric()) continue; // skip electric field moments

gpu::vector<P> &res = gpu_interps[0][im];
res.resize(full_block * grid.num_indexes());

if (mom.is_electric()) continue; // skip electric field moments

std::array<P const *, max_mom_dims> itg =
{gpu_integ[0][0].data() + mom.pows[0] * integ[0].stride(), nullptr, nullptr};
for (int i = 1; i < num_vel_; i++)
Expand Down Expand Up @@ -1142,11 +1137,15 @@ void moment_manager<P>::compute_moments(

moment_expand(pdof, pos_grid.num_dims(), num_vel_, reduce_ij[0], w2.vec, res);

if (result_to_cpu) res.copy_to_host(interps[im]);
if (result_to_cpu) {
res.copy_to_host(interps[im]);
full_level.get(im).resize(0);
}

// solve poisson equation while w1 holds density
constexpr bool always_interp = true;
if (mom.is_zero() and has_electric)
solve_poisson_gpu(conn, hier, interp, kwork);
solve_poisson_gpu<always_interp>(gpu::device{0}, grid, conn, hier, interp, kwork, result_to_cpu);
}
}
#endif
Expand Down
Loading
Loading