Threads: add missing broadcast to TeamThreadRange parallel_scan (kokk…

…os#6601) * try * use reference
masterleinad · Nov 17, 2023 · 8fd8c94 · 8fd8c94
1 parent 1a14531
commit 8fd8c94
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/core/src/Threads/Kokkos_Threads_Team.hpp b/core/src/Threads/Kokkos_Threads_Team.hpp
@@ -1001,8 +1001,10 @@ KOKKOS_INLINE_FUNCTION void parallel_scan(
     lambda(i, scan_val, false);
   }
 
+  auto & team_member = loop_bounds.thread;
+
   // 'scan_val' output is the exclusive prefix sum
-  scan_val = loop_bounds.thread.team_scan(scan_val);
+  scan_val = team_member.team_scan(scan_val);
 
 #ifdef KOKKOS_ENABLE_PRAGMA_IVDEP
 #pragma ivdep
@@ -1012,6 +1014,8 @@ KOKKOS_INLINE_FUNCTION void parallel_scan(
     lambda(i, scan_val, true);
   }
 
+  team_member.team_broadcast(scan_val, team_member.team_size() - 1);
+
   return_val = scan_val;
 }