Skip to content
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.

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

Add par_reduce_inner functions #1147

Merged
merged 6 commits into from
Aug 26, 2024
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
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 CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
## Current develop

### Added (new features/APIs/variables/...)
- [[PR 1147]](https://github.com/parthenon-hpc-lab/parthenon/pull/1147) Add `par_reduce_inner` functions
- [[PR 1148]](https://github.com/parthenon-hpc-lab/parthenon/pull/1148) Add `GetPackDimension` to `StateDescriptor` for calculating pack sizes before `Mesh` initialization
- [[PR 1143]](https://github.com/parthenon-hpc-lab/parthenon/pull/1143) Add tensor indices to VariableState, add radiation constant to constants, add TypeLists, allow for arbitrary containers for solvers
- [[PR 1140]](https://github.com/parthenon-hpc-lab/parthenon/pull/1140) Allow for relative convergence tolerance in BiCGSTAB solver.
Expand Down
59 changes: 59 additions & 0 deletions src/kokkos_abstraction.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -976,6 +976,65 @@ KOKKOS_FORCEINLINE_FUNCTION void par_for_inner(team_mbr_t team_member, Args &&..
par_for_inner(DEFAULT_INNER_LOOP_PATTERN, team_member, std::forward<Args>(args)...);
}

// Inner reduction loops
template <typename Function, typename T>
KOKKOS_FORCEINLINE_FUNCTION void
par_reduce_inner(team_mbr_t team_member, const int kl, const int ku, const int jl,
const int ju, const int il, const int iu, const Function &function,
T reduction) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should these be

par_reduce_inner(InnerLoopPatternTTR, team_mbr_t team_member, const int kl, const int ku, const int jl,
                 const int ju, const int il, const int iu, const Function &function,
                 T reduction) {

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure. I didn't really think about it when dropping these in, and I don't think I fully understand the way loop pattern resolution is done anyway...
The way I understand it, since I use TVR_INNER_LOOP for ThreadVectorRange, these wouldn't resolve under that, right? So, we'd need a TVR and SIMD pattern for these too, if we wanted to make them selectable. Since these should be used pretty rarely, I don't know if the performance benefit of allowing customization is worth three implementations we're just going to stamp over with all the better general stuff you guys are writing for the next release.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I could see it both ways. With my suggestion, you would have to call parthenon::par_reduce_inner(inner_loop_pattern_ttr_tag, ...) instead of parthenon::par_reduce_inner(...), which would be more explicit that you aren't necessarily using the default inner loop pattern and conform to how the rest of the par_*_inner functions look. But this doesn't change anything about how this behaves.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, right, you can explicitly tag loops! It seems better to be explicit, then, let me give it a go today.

const int Nk = ku - kl + 1;
const int Nj = ju - jl + 1;
const int Ni = iu - il + 1;
const int NkNjNi = Nk * Nj * Ni;
const int NjNi = Nj * Ni;
Kokkos::parallel_reduce(
Kokkos::TeamThreadRange(team_member, NkNjNi),
[&](const int &idx, typename T::value_type &lreduce) {
int k = idx / NjNi;
int j = (idx - k * NjNi) / Ni;
int i = idx - k * NjNi - j * Ni;
k += kl;
j += jl;
i += il;
function(k, j, i, lreduce);
},
reduction);
}

template <typename Function, typename T>
KOKKOS_FORCEINLINE_FUNCTION void
par_reduce_inner(team_mbr_t team_member, const int jl, const int ju, const int il,
const int iu, const Function &function, T reduction) {
const int Nj = ju - jl + 1;
const int Ni = iu - il + 1;
const int NjNi = Nj * Ni;
Kokkos::parallel_reduce(
Kokkos::TeamThreadRange(team_member, NjNi),
[&](const int &idx, typename T::value_type &lreduce) {
int j = idx / Ni;
int i = idx - j * Ni;
j += jl;
i += il;
function(j, i, lreduce);
},
reduction);
}

template <typename Function, typename T>
KOKKOS_FORCEINLINE_FUNCTION void par_reduce_inner(team_mbr_t team_member, const int il,
const int iu, const Function &function,
T reduction) {
const int Ni = iu - il + 1;
Kokkos::parallel_reduce(
Kokkos::TeamThreadRange(team_member, Ni),
[&](const int &idx, typename T::value_type &lreduce) {
int i = idx;
i += il;
function(i, lreduce);
},
reduction);
}

// reused from kokoks/core/perf_test/PerfTest_ExecSpacePartitioning.cpp
// commit a0d011fb30022362c61b3bb000ae3de6906cb6a7
template <class ExecSpace>
Expand Down
Loading