Skip to content

Commit

Permalink
Fix inconsistency in team_member.reduce sample code snippet (#634)
Browse files Browse the repository at this point in the history
Co-authored-by: Daniel Arndt <[email protected]>
  • Loading branch information
mabruzzo and masterleinad authored Feb 26, 2025
1 parent f2526c2 commit 8b5a09f
Showing 1 changed file with 12 additions and 4 deletions.
16 changes: 12 additions & 4 deletions docs/source/ProgrammingGuide/HierarchicalParallelism.md
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,9 @@ Kokkos::TeamPolicy<SomeTag, ExecutionSpace>
The team policy's `member_type` provides the necessary functionality to use teams within a parallel kernel. It allows access to thread identifiers such as the league rank and size, and the team rank and size. It also provides team-synchronous actions such as team barriers, reductions and scans.
```c++
using Kokkos::atomic_add;
using Kokkos::PerTeam;
using Kokkos::Sum;
using Kokkos::TeamPolicy;
using Kokkos::parallel_for;
Expand All @@ -73,12 +76,17 @@ TeamPolicy<ExecutionSpace> policy (league_size, Kokkos::AUTO() );
// Launch a kernel
parallel_for (policy, KOKKOS_LAMBDA (member_type team_member) {
// Calculate a global thread id
int k = team_member.league_rank () * team_member.team_size () +
int k = team_member.league_rank () * team_member.team_size () +
team_member.team_rank ();
// Calculate the sum of the global thread ids of this team
int team_sum = team_member.reduce (k);
// Atomically add the value to a global value
a() += team_sum;
int team_sum = k;
team_member.team_reduce(Sum<int, typename ExecutionSpace::memory_space>(team_sum));
// Atomically add the computed sum to a global value
Kokkos::single (PerTeam (team_member), [=] () {
atomic_add(&global_value(), team_sum);
});
});
```

Expand Down

0 comments on commit 8b5a09f

Please sign in to comment.