Cabana 0.8.0-dev
 
Loading...
Searching...
No Matches
Cabana_Parallel.hpp
Go to the documentation of this file.
1/****************************************************************************
2 * Copyright (c) 2018-2023 by the Cabana authors *
3 * All rights reserved. *
4 * *
5 * This file is part of the Cabana library. Cabana is distributed under a *
6 * BSD 3-clause license. For the licensing terms see the LICENSE file in *
7 * the top-level directory. *
8 * *
9 * SPDX-License-Identifier: BSD-3-Clause *
10 ****************************************************************************/
11
16#ifndef CABANA_PARALLEL_HPP
17#define CABANA_PARALLEL_HPP
18
19#include <Cabana_ExecutionPolicy.hpp>
20#include <Cabana_LinkedCellList.hpp>
21#include <Cabana_NeighborList.hpp>
22#include <Cabana_Types.hpp> // is_accessible_from
23#include <impl/Cabana_CartesianGrid.hpp>
24
25#include <Kokkos_Core.hpp>
26#include <Kokkos_Profiling_ScopedRegion.hpp>
27
28#include <cstdlib>
29#include <type_traits>
30
31namespace Cabana
32{
33//---------------------------------------------------------------------------//
34namespace Impl
35{
37
38// No work tag was provided so call without a tag argument.
39template <class WorkTag, class FunctorType, class... IndexTypes>
40KOKKOS_FORCEINLINE_FUNCTION
41 typename std::enable_if<std::is_same<WorkTag, void>::value>::type
42 functorTagDispatch( const FunctorType& functor, IndexTypes&&... indices )
43{
44 functor( std::forward<IndexTypes>( indices )... );
45}
46
47// The user gave us a tag so call the version using that.
48template <class WorkTag, class FunctorType, class... IndexTypes>
49KOKKOS_FORCEINLINE_FUNCTION
50 typename std::enable_if<!std::is_same<WorkTag, void>::value>::type
51 functorTagDispatch( const FunctorType& functor, IndexTypes&&... indices )
52{
53 const WorkTag t{};
54 functor( t, std::forward<IndexTypes>( indices )... );
55}
56
57// No work tag was provided so call reduce without a tag argument.
58template <class WorkTag, class FunctorType, class... IndexTypes,
59 class ReduceType>
60KOKKOS_FORCEINLINE_FUNCTION
61 typename std::enable_if<std::is_same<WorkTag, void>::value>::type
62 functorTagDispatch( const FunctorType& functor, IndexTypes&&... indices,
63 ReduceType& reduce_val )
64{
65 functor( std::forward<IndexTypes>( indices )..., reduce_val );
66}
67
68// The user gave us a tag so call the reduce version using that.
69template <class WorkTag, class FunctorType, class... IndexTypes,
70 class ReduceType>
71KOKKOS_FORCEINLINE_FUNCTION
72 typename std::enable_if<!std::is_same<WorkTag, void>::value>::type
73 functorTagDispatch( const FunctorType& functor, IndexTypes&&... indices,
74 ReduceType& reduce_val )
75{
76 const WorkTag t{};
77 functor( t, std::forward<IndexTypes>( indices )..., reduce_val );
78}
79
80template <class ExecutionPolicy, class Functor>
81struct ParallelFor;
82
83template <class Functor, int VectorLength, class... Properties>
84struct ParallelFor<SimdPolicy<VectorLength, Properties...>, Functor>
85{
86 using simd_policy = SimdPolicy<VectorLength, Properties...>;
87 using team_policy = typename simd_policy::base_type;
88 using work_tag = typename team_policy::work_tag;
89 using index_type = typename team_policy::index_type;
90 using member_type = typename team_policy::member_type;
91
92 simd_policy exec_policy_;
93 Functor functor_;
94
95 ParallelFor( std::string label, simd_policy exec_policy, Functor functor )
96 : exec_policy_( std::move( exec_policy ) )
97 , functor_( std::move( functor ) )
98 {
99 if ( label.empty() )
100 Kokkos::parallel_for(
101 dynamic_cast<const team_policy&>( exec_policy_ ), *this );
102 else
103 Kokkos::parallel_for(
104 label, dynamic_cast<const team_policy&>( exec_policy_ ),
105 *this );
106 }
107
108 template <class WorkTag>
109 KOKKOS_FUNCTION std::enable_if_t<!std::is_void<WorkTag>::value &&
110 std::is_same<WorkTag, work_tag>::value>
111 operator()( WorkTag, member_type const& team ) const
112 {
113 this->operator()( team );
114 }
115
116 KOKKOS_FUNCTION void operator()( member_type const& team ) const
117 {
118 index_type s = team.league_rank() + exec_policy_.structBegin();
119 Kokkos::parallel_for(
120 Kokkos::ThreadVectorRange( team, exec_policy_.arrayBegin( s ),
121 exec_policy_.arrayEnd( s ) ),
122 [&]( index_type a )
123 { Impl::functorTagDispatch<work_tag>( functor_, s, a ); } );
124 }
125};
126
128} // end namespace Impl
129
130//---------------------------------------------------------------------------//
131// SIMD Parallel For
132//---------------------------------------------------------------------------//
173template <class FunctorType, int VectorLength, class... ExecParameters>
174inline void simd_parallel_for(
175 const SimdPolicy<VectorLength, ExecParameters...>& exec_policy,
176 const FunctorType& functor, const std::string& str = "" )
177{
178 Kokkos::Profiling::ScopedRegion region( "Cabana::simd_parallel_for" );
179
180 Impl::ParallelFor<SimdPolicy<VectorLength, ExecParameters...>, FunctorType>(
181 str, exec_policy, functor );
182}
183
184//---------------------------------------------------------------------------//
185// Neighbor Parallel For
186//---------------------------------------------------------------------------//
187// Algorithm tags.
188
190class FirstNeighborsTag
191{
192};
193
195class SecondNeighborsTag
196{
197};
198
200class SerialOpTag
201{
202};
203
205class TeamOpTag
206{
207};
208
210class TeamVectorOpTag
211{
212};
213
214//---------------------------------------------------------------------------//
251template <class FunctorType, class NeighborListType, class... ExecParameters>
252inline void neighbor_parallel_for(
253 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
254 const FunctorType& functor, const NeighborListType& list,
255 const FirstNeighborsTag, const SerialOpTag, const std::string& str = "",
256 typename std::enable_if<( !is_linked_cell_list<NeighborListType>::value ),
257 int>::type* = 0 )
258{
259 Kokkos::Profiling::ScopedRegion region( "Cabana::neighbor_parallel_for" );
260
261 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
262
263 using execution_space =
264 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
265
266 using index_type =
267 typename Kokkos::RangePolicy<ExecParameters...>::index_type;
268
269 using neighbor_list_traits = NeighborList<NeighborListType>;
270
271 using memory_space = typename neighbor_list_traits::memory_space;
272
273 auto begin = exec_policy.begin();
274 auto end = exec_policy.end();
275 using linear_policy_type = Kokkos::RangePolicy<execution_space, void, void>;
276 linear_policy_type linear_exec_policy( begin, end );
277
278 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
279
280 auto neigh_func = KOKKOS_LAMBDA( const index_type i )
281 {
282 for ( index_type n = 0;
283 n < neighbor_list_traits::numNeighbor( list, i ); ++n )
284 Impl::functorTagDispatch<work_tag>(
285 functor, i,
286 static_cast<index_type>(
287 neighbor_list_traits::getNeighbor( list, i, n ) ) );
288 };
289 if ( str.empty() )
290 Kokkos::parallel_for( linear_exec_policy, neigh_func );
291 else
292 Kokkos::parallel_for( str, linear_exec_policy, neigh_func );
293}
294
295//---------------------------------------------------------------------------//
315template <class FunctorType, class NeighborListType, class... ExecParameters>
316inline void neighbor_parallel_for(
317 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
318 const FunctorType& functor, const NeighborListType& list,
319 const SecondNeighborsTag, const SerialOpTag, const std::string& str = "",
320 typename std::enable_if<( !is_linked_cell_list<NeighborListType>::value ),
321 int>::type* = 0 )
322{
323 Kokkos::Profiling::ScopedRegion region( "Cabana::neighbor_parallel_for" );
324
325 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
326
327 using execution_space =
328 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
329
330 using index_type =
331 typename Kokkos::RangePolicy<ExecParameters...>::index_type;
332
333 using neighbor_list_traits = NeighborList<NeighborListType>;
334
335 using memory_space = typename neighbor_list_traits::memory_space;
336
337 auto begin = exec_policy.begin();
338 auto end = exec_policy.end();
339 using linear_policy_type = Kokkos::RangePolicy<execution_space, void, void>;
340 linear_policy_type linear_exec_policy( begin, end );
341
342 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
343
344 auto neigh_func = KOKKOS_LAMBDA( const index_type i )
345 {
346 const index_type nn = neighbor_list_traits::numNeighbor( list, i );
347
348 for ( index_type n = 0; n < nn; ++n )
349 {
350 const index_type j =
351 neighbor_list_traits::getNeighbor( list, i, n );
352
353 for ( index_type a = n + 1; a < nn; ++a )
354 {
355 const index_type k =
356 neighbor_list_traits::getNeighbor( list, i, a );
357 Impl::functorTagDispatch<work_tag>( functor, i, j, k );
358 }
359 }
360 };
361 if ( str.empty() )
362 Kokkos::parallel_for( linear_exec_policy, neigh_func );
363 else
364 Kokkos::parallel_for( str, linear_exec_policy, neigh_func );
365}
366
367//---------------------------------------------------------------------------//
386template <class FunctorType, class NeighborListType, class... ExecParameters>
387inline void neighbor_parallel_for(
388 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
389 const FunctorType& functor, const NeighborListType& list,
390 const FirstNeighborsTag, const TeamOpTag, const std::string& str = "",
391 typename std::enable_if<( !is_linked_cell_list<NeighborListType>::value ),
392 int>::type* = 0 )
393{
394 Kokkos::Profiling::ScopedRegion region( "Cabana::neighbor_parallel_for" );
395
396 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
397
398 using execution_space =
399 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
400
401 using kokkos_policy =
402 Kokkos::TeamPolicy<execution_space, Kokkos::Schedule<Kokkos::Dynamic>>;
403 kokkos_policy team_policy( exec_policy.end() - exec_policy.begin(),
404 Kokkos::AUTO );
405
406 using index_type = typename kokkos_policy::index_type;
407
408 using neighbor_list_traits = NeighborList<NeighborListType>;
409
410 using memory_space = typename neighbor_list_traits::memory_space;
411
412 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
413
414 const auto range_begin = exec_policy.begin();
415
416 auto neigh_func =
417 KOKKOS_LAMBDA( const typename kokkos_policy::member_type& team )
418 {
419 index_type i = team.league_rank() + range_begin;
420 Kokkos::parallel_for(
421 Kokkos::TeamThreadRange(
422 team, neighbor_list_traits::numNeighbor( list, i ) ),
423 [&]( const index_type n )
424 {
425 Impl::functorTagDispatch<work_tag>(
426 functor, i,
427 static_cast<index_type>(
428 neighbor_list_traits::getNeighbor( list, i, n ) ) );
429 } );
430 };
431 if ( str.empty() )
432 Kokkos::parallel_for( team_policy, neigh_func );
433 else
434 Kokkos::parallel_for( str, team_policy, neigh_func );
435}
436
437//---------------------------------------------------------------------------//
458template <class FunctorType, class NeighborListType, class... ExecParameters>
459inline void neighbor_parallel_for(
460 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
461 const FunctorType& functor, const NeighborListType& list,
462 const SecondNeighborsTag, const TeamOpTag, const std::string& str = "",
463 typename std::enable_if<( !is_linked_cell_list<NeighborListType>::value ),
464 int>::type* = 0 )
465{
466 Kokkos::Profiling::ScopedRegion region( "Cabana::neighbor_parallel_for" );
467
468 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
469
470 using execution_space =
471 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
472
473 using kokkos_policy =
474 Kokkos::TeamPolicy<execution_space, Kokkos::Schedule<Kokkos::Dynamic>>;
475 kokkos_policy team_policy( exec_policy.end() - exec_policy.begin(),
476 Kokkos::AUTO );
477
478 using index_type = typename kokkos_policy::index_type;
479
480 using neighbor_list_traits = NeighborList<NeighborListType>;
481
482 using memory_space = typename neighbor_list_traits::memory_space;
483
484 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
485
486 const auto range_begin = exec_policy.begin();
487
488 auto neigh_func =
489 KOKKOS_LAMBDA( const typename kokkos_policy::member_type& team )
490 {
491 index_type i = team.league_rank() + range_begin;
492
493 const index_type nn = neighbor_list_traits::numNeighbor( list, i );
494 Kokkos::parallel_for(
495 Kokkos::TeamThreadRange( team, nn ),
496 [&]( const index_type n )
497 {
498 const index_type j =
499 neighbor_list_traits::getNeighbor( list, i, n );
500
501 for ( index_type a = n + 1; a < nn; ++a )
502 {
503 const index_type k =
504 neighbor_list_traits::getNeighbor( list, i, a );
505 Impl::functorTagDispatch<work_tag>( functor, i, j, k );
506 }
507 } );
508 };
509 if ( str.empty() )
510 Kokkos::parallel_for( team_policy, neigh_func );
511 else
512 Kokkos::parallel_for( str, team_policy, neigh_func );
513}
514
515//---------------------------------------------------------------------------//
536template <class FunctorType, class NeighborListType, class... ExecParameters>
537inline void neighbor_parallel_for(
538 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
539 const FunctorType& functor, const NeighborListType& list,
540 const SecondNeighborsTag, const TeamVectorOpTag,
541 const std::string& str = "",
542 typename std::enable_if<( !is_linked_cell_list<NeighborListType>::value ),
543 int>::type* = 0 )
544{
545 Kokkos::Profiling::ScopedRegion region( "Cabana::neighbor_parallel_for" );
546
547 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
548
549 using execution_space =
550 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
551
552 using kokkos_policy =
553 Kokkos::TeamPolicy<execution_space, Kokkos::Schedule<Kokkos::Dynamic>>;
554 kokkos_policy team_policy( exec_policy.end() - exec_policy.begin(),
555 Kokkos::AUTO );
556
557 using index_type = typename kokkos_policy::index_type;
558
559 using neighbor_list_traits = NeighborList<NeighborListType>;
560
561 using memory_space = typename neighbor_list_traits::memory_space;
562
563 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
564
565 const auto range_begin = exec_policy.begin();
566
567 auto neigh_func =
568 KOKKOS_LAMBDA( const typename kokkos_policy::member_type& team )
569 {
570 index_type i = team.league_rank() + range_begin;
571
572 const index_type nn = neighbor_list_traits::numNeighbor( list, i );
573 Kokkos::parallel_for(
574 Kokkos::TeamThreadRange( team, nn ),
575 [&]( const index_type n )
576 {
577 const index_type j =
578 neighbor_list_traits::getNeighbor( list, i, n );
579
580 Kokkos::parallel_for(
581 Kokkos::ThreadVectorRange( team, n + 1, nn ),
582 [&]( const index_type a )
583 {
584 const index_type k =
585 neighbor_list_traits::getNeighbor( list, i, a );
586 Impl::functorTagDispatch<work_tag>( functor, i, j, k );
587 } );
588 } );
589 };
590 if ( str.empty() )
591 Kokkos::parallel_for( team_policy, neigh_func );
592 else
593 Kokkos::parallel_for( str, team_policy, neigh_func );
594}
595
596//---------------------------------------------------------------------------//
597// Neighbor Parallel Reduce
598//---------------------------------------------------------------------------//
638template <class FunctorType, class NeighborListType, class ReduceType,
639 class... ExecParameters>
640inline void neighbor_parallel_reduce(
641 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
642 const FunctorType& functor, const NeighborListType& list,
643 const FirstNeighborsTag, const SerialOpTag, ReduceType& reduce_val,
644 const std::string& str = "",
645 typename std::enable_if<( !is_linked_cell_list<NeighborListType>::value ),
646 int>::type* = 0 )
647{
648 Kokkos::Profiling::ScopedRegion region(
649 "Cabana::neighbor_parallel_reduce" );
650
651 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
652
653 using execution_space =
654 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
655
656 using index_type =
657 typename Kokkos::RangePolicy<ExecParameters...>::index_type;
658
659 using neighbor_list_traits = NeighborList<NeighborListType>;
660
661 using memory_space = typename neighbor_list_traits::memory_space;
662
663 auto begin = exec_policy.begin();
664 auto end = exec_policy.end();
665 using linear_policy_type = Kokkos::RangePolicy<execution_space, void, void>;
666 linear_policy_type linear_exec_policy( begin, end );
667
668 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
669
670 auto neigh_reduce = KOKKOS_LAMBDA( const index_type i, ReduceType& ival )
671 {
672 for ( index_type n = 0;
673 n < neighbor_list_traits::numNeighbor( list, i ); ++n )
674 Impl::functorTagDispatch<work_tag>(
675 functor, i,
676 static_cast<index_type>(
677 neighbor_list_traits::getNeighbor( list, i, n ) ),
678 ival );
679 };
680 if ( str.empty() )
681 Kokkos::parallel_reduce( linear_exec_policy, neigh_reduce, reduce_val );
682 else
683 Kokkos::parallel_reduce( str, linear_exec_policy, neigh_reduce,
684 reduce_val );
685}
686
687//---------------------------------------------------------------------------//
709template <class FunctorType, class NeighborListType, class ReduceType,
710 class... ExecParameters>
711inline void neighbor_parallel_reduce(
712 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
713 const FunctorType& functor, const NeighborListType& list,
714 const SecondNeighborsTag, const SerialOpTag, ReduceType& reduce_val,
715 const std::string& str = "",
716 typename std::enable_if<( !is_linked_cell_list<NeighborListType>::value ),
717 int>::type* = 0 )
718{
719 Kokkos::Profiling::ScopedRegion region(
720 "Cabana::neighbor_parallel_reduce" );
721
722 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
723
724 using execution_space =
725 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
726
727 using index_type =
728 typename Kokkos::RangePolicy<ExecParameters...>::index_type;
729
730 using neighbor_list_traits = NeighborList<NeighborListType>;
731
732 using memory_space = typename neighbor_list_traits::memory_space;
733
734 auto begin = exec_policy.begin();
735 auto end = exec_policy.end();
736 using linear_policy_type = Kokkos::RangePolicy<execution_space, void, void>;
737 linear_policy_type linear_exec_policy( begin, end );
738
739 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
740
741 auto neigh_reduce = KOKKOS_LAMBDA( const index_type i, ReduceType& ival )
742 {
743 const index_type nn = neighbor_list_traits::numNeighbor( list, i );
744
745 for ( index_type n = 0; n < nn; ++n )
746 {
747 const index_type j =
748 neighbor_list_traits::getNeighbor( list, i, n );
749
750 for ( index_type a = n + 1; a < nn; ++a )
751 {
752 const index_type k =
753 neighbor_list_traits::getNeighbor( list, i, a );
754 Impl::functorTagDispatch<work_tag>( functor, i, j, k, ival );
755 }
756 }
757 };
758 if ( str.empty() )
759 Kokkos::parallel_reduce( linear_exec_policy, neigh_reduce, reduce_val );
760 else
761 Kokkos::parallel_reduce( str, linear_exec_policy, neigh_reduce,
762 reduce_val );
763}
764
765//---------------------------------------------------------------------------//
787template <class FunctorType, class NeighborListType, class ReduceType,
788 class... ExecParameters>
789inline void neighbor_parallel_reduce(
790 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
791 const FunctorType& functor, const NeighborListType& list,
792 const FirstNeighborsTag, const TeamOpTag, ReduceType& reduce_val,
793 const std::string& str = "",
794 typename std::enable_if<( !is_linked_cell_list<NeighborListType>::value ),
795 int>::type* = 0 )
796{
797 Kokkos::Profiling::ScopedRegion region(
798 "Cabana::neighbor_parallel_reduce" );
799
800 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
801
802 using execution_space =
803 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
804
805 using kokkos_policy =
806 Kokkos::TeamPolicy<execution_space, Kokkos::Schedule<Kokkos::Dynamic>>;
807 kokkos_policy team_policy( exec_policy.end() - exec_policy.begin(),
808 Kokkos::AUTO );
809
810 using index_type = typename kokkos_policy::index_type;
811
812 using neighbor_list_traits = NeighborList<NeighborListType>;
813
814 using memory_space = typename neighbor_list_traits::memory_space;
815
816 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
817
818 const auto range_begin = exec_policy.begin();
819
820 auto neigh_reduce = KOKKOS_LAMBDA(
821 const typename kokkos_policy::member_type& team, ReduceType& ival )
822 {
823 index_type i = team.league_rank() + range_begin;
824 ReduceType reduce_n = 0;
825
826 Kokkos::parallel_reduce(
827 Kokkos::TeamThreadRange(
828 team, neighbor_list_traits::numNeighbor( list, i ) ),
829 [&]( const index_type n, ReduceType& nval )
830 {
831 Impl::functorTagDispatch<work_tag>(
832 functor, i,
833 static_cast<index_type>(
834 neighbor_list_traits::getNeighbor( list, i, n ) ),
835 nval );
836 },
837 reduce_n );
838 Kokkos::single( Kokkos::PerTeam( team ), [&]() { ival += reduce_n; } );
839 };
840 if ( str.empty() )
841 Kokkos::parallel_reduce( team_policy, neigh_reduce, reduce_val );
842 else
843 Kokkos::parallel_reduce( str, team_policy, neigh_reduce, reduce_val );
844}
845
846//---------------------------------------------------------------------------//
869template <class FunctorType, class NeighborListType, class ReduceType,
870 class... ExecParameters>
871inline void neighbor_parallel_reduce(
872 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
873 const FunctorType& functor, const NeighborListType& list,
874 const SecondNeighborsTag, const TeamOpTag, ReduceType& reduce_val,
875 const std::string& str = "",
876 typename std::enable_if<( !is_linked_cell_list<NeighborListType>::value ),
877 int>::type* = 0 )
878{
879 Kokkos::Profiling::ScopedRegion region(
880 "Cabana::neighbor_parallel_reduce" );
881
882 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
883
884 using execution_space =
885 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
886
887 using kokkos_policy =
888 Kokkos::TeamPolicy<execution_space, Kokkos::Schedule<Kokkos::Dynamic>>;
889 kokkos_policy team_policy( exec_policy.end() - exec_policy.begin(),
890 Kokkos::AUTO );
891
892 using index_type = typename kokkos_policy::index_type;
893
894 using neighbor_list_traits = NeighborList<NeighborListType>;
895
896 using memory_space = typename neighbor_list_traits::memory_space;
897
898 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
899
900 const auto range_begin = exec_policy.begin();
901
902 auto neigh_reduce = KOKKOS_LAMBDA(
903 const typename kokkos_policy::member_type& team, ReduceType& ival )
904 {
905 index_type i = team.league_rank() + range_begin;
906 ReduceType reduce_n = 0;
907
908 const index_type nn = neighbor_list_traits::numNeighbor( list, i );
909 Kokkos::parallel_reduce(
910 Kokkos::TeamThreadRange( team, nn ),
911 [&]( const index_type n, ReduceType& nval )
912 {
913 const index_type j =
914 neighbor_list_traits::getNeighbor( list, i, n );
915
916 for ( index_type a = n + 1; a < nn; ++a )
917 {
918 const index_type k =
919 neighbor_list_traits::getNeighbor( list, i, a );
920 Impl::functorTagDispatch<work_tag>( functor, i, j, k,
921 nval );
922 }
923 },
924 reduce_n );
925 Kokkos::single( Kokkos::PerTeam( team ), [&]() { ival += reduce_n; } );
926 };
927 if ( str.empty() )
928 Kokkos::parallel_reduce( team_policy, neigh_reduce, reduce_val );
929 else
930 Kokkos::parallel_reduce( str, team_policy, neigh_reduce, reduce_val );
931}
932
933//---------------------------------------------------------------------------//
956template <class FunctorType, class NeighborListType, class ReduceType,
957 class... ExecParameters>
958inline void neighbor_parallel_reduce(
959 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
960 const FunctorType& functor, const NeighborListType& list,
961 const SecondNeighborsTag, const TeamVectorOpTag, ReduceType& reduce_val,
962 const std::string& str = "",
963 typename std::enable_if<( !is_linked_cell_list<NeighborListType>::value ),
964 int>::type* = 0 )
965{
966 Kokkos::Profiling::ScopedRegion region(
967 "Cabana::neighbor_parallel_reduce" );
968
969 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
970
971 using execution_space =
972 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
973
974 using kokkos_policy =
975 Kokkos::TeamPolicy<execution_space, Kokkos::Schedule<Kokkos::Dynamic>>;
976 kokkos_policy team_policy( exec_policy.end() - exec_policy.begin(),
977 Kokkos::AUTO );
978
979 using index_type = typename kokkos_policy::index_type;
980
981 using neighbor_list_traits = NeighborList<NeighborListType>;
982
983 using memory_space = typename neighbor_list_traits::memory_space;
984
985 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
986
987 const auto range_begin = exec_policy.begin();
988
989 auto neigh_reduce = KOKKOS_LAMBDA(
990 const typename kokkos_policy::member_type& team, ReduceType& ival )
991 {
992 index_type i = team.league_rank() + range_begin;
993 ReduceType reduce_n = 0;
994
995 const index_type nn = neighbor_list_traits::numNeighbor( list, i );
996 Kokkos::parallel_reduce(
997 Kokkos::TeamThreadRange( team, nn ),
998 [&]( const index_type n, ReduceType& nval )
999 {
1000 const index_type j =
1001 neighbor_list_traits::getNeighbor( list, i, n );
1002 ReduceType reduce_a = 0;
1003
1004 Kokkos::parallel_reduce(
1005 Kokkos::ThreadVectorRange( team, n + 1, nn ),
1006 [&]( const index_type a, ReduceType& aval )
1007 {
1008 const index_type k =
1009 neighbor_list_traits::getNeighbor( list, i, a );
1010 Impl::functorTagDispatch<work_tag>( functor, i, j, k,
1011 aval );
1012 },
1013 reduce_a );
1014 nval += reduce_a;
1015 },
1016 reduce_n );
1017 Kokkos::single( Kokkos::PerTeam( team ), [&]() { ival += reduce_n; } );
1018 };
1019 if ( str.empty() )
1020 Kokkos::parallel_reduce( team_policy, neigh_reduce, reduce_val );
1021 else
1022 Kokkos::parallel_reduce( str, team_policy, neigh_reduce, reduce_val );
1023}
1024
1025//---------------------------------------------------------------------------//
1058template <class IndexType, class FunctorType, class NeighborListType>
1059KOKKOS_INLINE_FUNCTION void
1060for_each_neighbor( const IndexType i, const FunctorType& neighbor_functor,
1061 const NeighborListType& list, const FirstNeighborsTag )
1062{
1063 using neighbor_list_traits = NeighborList<NeighborListType>;
1064
1065 for ( IndexType n = 0;
1066 n < static_cast<IndexType>(
1067 neighbor_list_traits::numNeighbor( list, i ) );
1068 ++n )
1069 neighbor_functor(
1070 i, static_cast<IndexType>(
1071 neighbor_list_traits::getNeighbor( list, i, n ) ) );
1072}
1073
1074//---------------------------------------------------------------------------//
1091template <class IndexType, class FunctorType, class NeighborListType,
1092 class TeamMemberType>
1093KOKKOS_INLINE_FUNCTION void
1094for_each_neighbor( const IndexType i, const TeamMemberType team,
1095 const FunctorType& neighbor_functor,
1096 const NeighborListType& list, const FirstNeighborsTag )
1097{
1098 using neighbor_list_traits = NeighborList<NeighborListType>;
1099
1100 Kokkos::parallel_for(
1101 Kokkos::TeamThreadRange( team,
1102 neighbor_list_traits::numNeighbor( list, i ) ),
1103 [&]( const IndexType n )
1104 {
1105 Impl::functorTagDispatch<void>(
1106 neighbor_functor, i,
1107 static_cast<IndexType>(
1108 neighbor_list_traits::getNeighbor( list, i, n ) ) );
1109 } );
1110}
1111
1112//---------------------------------------------------------------------------//
1113// Linked Cell Parallel For
1114//---------------------------------------------------------------------------//
1115namespace Impl
1116{
1118
1122template <class WorkTag, class Functor, class Policy, class LinkedCellType,
1123 class ViewType>
1124struct LinkedCellParallelFor
1125{
1127 using index_type = typename Policy::index_type;
1128
1130 Policy _exec_policy;
1132 Functor _functor;
1134 LinkedCellType _list;
1135
1137 index_type _begin;
1138
1140 NeighborDiscriminator<FullNeighborTag> _discriminator;
1141
1143 LinkedCellParallelFor( std::string label, Policy exec_policy,
1144 Functor functor, const LinkedCellType& list,
1145 const index_type begin = 0 )
1146 : _exec_policy( exec_policy )
1147 , _functor( functor )
1148 , _list( list )
1149 , _begin( begin )
1150 {
1151 if ( label.empty() )
1152 Kokkos::parallel_for( dynamic_cast<const Policy&>( exec_policy ),
1153 *this );
1154 else
1155 Kokkos::parallel_for(
1156 label, dynamic_cast<const Policy&>( exec_policy ), *this );
1157 }
1158
1160 KOKKOS_FUNCTION void operator()( SerialOpTag, const index_type i ) const
1161 {
1162 int imin, imax, jmin, jmax, kmin, kmax;
1163 _list.getStencilCells( _list.getParticleBin( i ), imin, imax, jmin,
1164 jmax, kmin, kmax );
1165
1166 // Loop over the cell stencil.
1167 for ( int gi = imin; gi < imax; ++gi )
1168 for ( int gj = jmin; gj < jmax; ++gj )
1169 for ( int gk = kmin; gk < kmax; ++gk )
1170 {
1171 // Check the particles in this bin to see if they are
1172 // neighbors.
1173 auto offset = _list.binOffset( gi, gj, gk );
1174 auto size = _list.binSize( gi, gj, gk );
1175 for ( std::size_t n = offset; n < offset + size; ++n )
1176 {
1177 // Get the true id of the candidate neighbor.
1178 auto j = _list.getParticle( n );
1179
1180 // Avoid self interactions (dummy position args).
1181 if ( _discriminator.isValid( i, 0, 0, 0, j, 0, 0, 0 ) )
1182 {
1183 Impl::functorTagDispatch<WorkTag>( _functor, i, j );
1184 }
1185 }
1186 }
1187 }
1188
1190 KOKKOS_FUNCTION void
1191 operator()( TeamOpTag, const typename Policy::member_type& team ) const
1192 {
1193 index_type i = team.league_rank() + _begin;
1194 int imin, imax, jmin, jmax, kmin, kmax;
1195 _list.getStencilCells( _list.getParticleBin( i ), imin, imax, jmin,
1196 jmax, kmin, kmax );
1197
1198 // Loop over the cell stencil.
1199 for ( int gi = imin; gi < imax; ++gi )
1200 for ( int gj = jmin; gj < jmax; ++gj )
1201 for ( int gk = kmin; gk < kmax; ++gk )
1202 {
1203 // Check the particles in this bin to see if they
1204 // are neighbors.
1205 auto offset = _list.binOffset( gi, gj, gk );
1206 auto size = _list.binSize( gi, gj, gk );
1207 Kokkos::parallel_for(
1208 Kokkos::TeamThreadRange( team, offset, offset + size ),
1209 [&]( const index_type n )
1210 {
1211 // Get the true id of the candidate neighbor.
1212 auto j = _list.getParticle( n );
1213
1214 // Avoid self interactions (dummy position args).
1215 if ( _discriminator.isValid( i, 0, 0, 0, j, 0, 0,
1216 0 ) )
1217 {
1218 Impl::functorTagDispatch<WorkTag>( _functor, i,
1219 j );
1220 }
1221 } );
1222 }
1223 };
1224};
1225
1230template <class WorkTag, class Functor, class Policy, class LinkedCellType,
1231 class ViewType, class ReduceType>
1232struct LinkedCellParallelReduce
1233{
1235 using index_type = typename Policy::index_type;
1236
1238 Policy _exec_policy;
1240 Functor _functor;
1242 LinkedCellType _list;
1243
1245 index_type _begin;
1246
1248 NeighborDiscriminator<FullNeighborTag> _discriminator;
1249
1251 LinkedCellParallelReduce( std::string label, Policy exec_policy,
1252 Functor functor, const LinkedCellType& list,
1253 ReduceType& reduce_val,
1254 const index_type begin = 0 )
1255 : _exec_policy( exec_policy )
1256 , _functor( functor )
1257 , _list( list )
1258 , _begin( begin )
1259 {
1260 if ( label.empty() )
1261 Kokkos::parallel_reduce( dynamic_cast<const Policy&>( exec_policy ),
1262 *this, reduce_val );
1263 else
1264 Kokkos::parallel_reduce( label,
1265 dynamic_cast<const Policy&>( exec_policy ),
1266 *this, reduce_val );
1267 }
1268
1270 KOKKOS_FUNCTION void operator()( SerialOpTag, const index_type i,
1271 ReduceType& ival ) const
1272 {
1273 int imin, imax, jmin, jmax, kmin, kmax;
1274 _list.getStencilCells( _list.getParticleBin( i ), imin, imax, jmin,
1275 jmax, kmin, kmax );
1276
1277 // Loop over the cell stencil.
1278 for ( int gi = imin; gi < imax; ++gi )
1279 for ( int gj = jmin; gj < jmax; ++gj )
1280 for ( int gk = kmin; gk < kmax; ++gk )
1281 {
1282 // Check the particles in this bin to see if they are
1283 // neighbors.
1284 auto offset = _list.binOffset( gi, gj, gk );
1285 auto size = _list.binSize( gi, gj, gk );
1286 for ( std::size_t n = offset; n < offset + size; ++n )
1287 {
1288 // Get the true id of the candidate neighbor.
1289 auto j = _list.getParticle( n );
1290
1291 // Avoid self interactions (dummy position args).
1292 if ( _discriminator.isValid( i, 0, 0, 0, j, 0, 0, 0 ) )
1293 {
1294 Impl::functorTagDispatch<WorkTag>( _functor, i, j,
1295 ival );
1296 }
1297 }
1298 }
1299 }
1300
1302 KOKKOS_FUNCTION void operator()( TeamOpTag,
1303 const typename Policy::member_type& team,
1304 ReduceType& ival ) const
1305 {
1306 index_type i = team.league_rank() + _begin;
1307 int imin, imax, jmin, jmax, kmin, kmax;
1308 _list.getStencilCells( _list.getParticleBin( i ), imin, imax, jmin,
1309 jmax, kmin, kmax );
1310
1311 // Loop over the cell stencil.
1312 for ( int gi = imin; gi < imax; ++gi )
1313 for ( int gj = jmin; gj < jmax; ++gj )
1314 for ( int gk = kmin; gk < kmax; ++gk )
1315 {
1316 // Check the particles in this bin to see if they
1317 // are neighbors.
1318 auto offset = _list.binOffset( gi, gj, gk );
1319 auto size = _list.binSize( gi, gj, gk );
1320 Kokkos::parallel_for(
1321 Kokkos::TeamThreadRange( team, offset, offset + size ),
1322 [&]( const index_type n )
1323 {
1324 // Get the true id of the candidate neighbor.
1325 auto j = _list.getParticle( n );
1326
1327 // Avoid self interactions (dummy position args).
1328 if ( _discriminator.isValid( i, 0, 0, 0, j, 0, 0,
1329 0 ) )
1330 {
1331 Impl::functorTagDispatch<WorkTag>( _functor, i,
1332 j, ival );
1333 }
1334 } );
1335 }
1336 };
1337};
1339} // namespace Impl
1340
1341//---------------------------------------------------------------------------//
1379
1380template <class FunctorType, class LinkedCellType, class... ExecParameters>
1381inline void neighbor_parallel_for(
1382 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
1383 const FunctorType& functor, const LinkedCellType& list,
1384 const FirstNeighborsTag, const SerialOpTag, const std::string& str = "",
1385 typename std::enable_if<( is_linked_cell_list<LinkedCellType>::value ),
1386 int>::type* = 0 )
1387{
1388 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
1389 using execution_space =
1390 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
1391
1392 using memory_space = typename LinkedCellType::memory_space;
1393
1394 auto begin = exec_policy.begin();
1395 auto end = exec_policy.end();
1396 // Cannot iterate over range that was not binned.
1397 assert( begin == list.getParticleBegin() );
1398 assert( end == list.getParticleEnd() );
1399
1400 using linear_policy_type =
1401 Kokkos::RangePolicy<SerialOpTag, execution_space>;
1402 linear_policy_type linear_exec_policy( begin, end );
1403
1404 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
1405
1406 Impl::LinkedCellParallelFor<work_tag, FunctorType, linear_policy_type,
1407 LinkedCellType,
1408 typename LinkedCellType::CountView>
1409 lcl_par( str, linear_exec_policy, functor, list, exec_policy.begin() );
1410}
1411
1412//---------------------------------------------------------------------------//
1433
1434template <class FunctorType, class LinkedCellType, class... ExecParameters>
1435inline void neighbor_parallel_for(
1436 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
1437 const FunctorType& functor, const LinkedCellType& list,
1438 const FirstNeighborsTag, const TeamOpTag, const std::string& str = "",
1439 typename std::enable_if<( is_linked_cell_list<LinkedCellType>::value ),
1440 int>::type* = 0 )
1441{
1442 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
1443 using execution_space =
1444 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
1445
1446 using team_policy_type =
1447 Kokkos::TeamPolicy<TeamOpTag, execution_space,
1448 Kokkos::Schedule<Kokkos::Dynamic>>;
1449 team_policy_type team_policy( exec_policy.end() - exec_policy.begin(),
1450 Kokkos::AUTO );
1451
1452 using memory_space = typename LinkedCellType::memory_space;
1453
1454 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
1455
1456 // Cannot iterate over range that was not binned.
1457 assert( exec_policy.begin() == list.getParticleBegin() );
1458 assert( exec_policy.end() == list.getParticleEnd() );
1459
1460 Impl::LinkedCellParallelFor<work_tag, FunctorType, team_policy_type,
1461 LinkedCellType,
1462 typename LinkedCellType::CountView>
1463 lcl_par( str, team_policy, functor, list, exec_policy.begin() );
1464}
1465
1466//---------------------------------------------------------------------------//
1505template <class FunctorType, class LinkedCellType, class ReduceType,
1506 class... ExecParameters>
1507inline void neighbor_parallel_reduce(
1508 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
1509 const FunctorType& functor, const LinkedCellType& list,
1510 const FirstNeighborsTag, const SerialOpTag, ReduceType& reduce_val,
1511 const std::string& str = "",
1512 typename std::enable_if<( is_linked_cell_list<LinkedCellType>::value ),
1513 int>::type* = 0 )
1514{
1515 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
1516 using execution_space =
1517 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
1518
1519 using memory_space = typename LinkedCellType::memory_space;
1520
1521 auto begin = exec_policy.begin();
1522 auto end = exec_policy.end();
1523 // Cannot iterate over range that was not binned.
1524 assert( begin == list.getParticleBegin() );
1525 assert( end == list.getParticleEnd() );
1526
1527 using linear_policy_type =
1528 Kokkos::RangePolicy<SerialOpTag, execution_space>;
1529 linear_policy_type linear_exec_policy( begin, end );
1530
1531 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
1532
1533 Impl::LinkedCellParallelReduce<
1534 work_tag, FunctorType, linear_policy_type, LinkedCellType,
1535 typename LinkedCellType::CountView, ReduceType>
1536 lcl_par( str, linear_exec_policy, functor, list, reduce_val,
1537 exec_policy.begin() );
1538}
1539
1540//---------------------------------------------------------------------------//
1562template <class FunctorType, class LinkedCellType, class ReduceType,
1563 class... ExecParameters>
1564inline void neighbor_parallel_reduce(
1565 const Kokkos::RangePolicy<ExecParameters...>& exec_policy,
1566 const FunctorType& functor, const LinkedCellType& list,
1567 const FirstNeighborsTag, const TeamOpTag, ReduceType& reduce_val,
1568 const std::string& str = "",
1569 typename std::enable_if<( is_linked_cell_list<LinkedCellType>::value ),
1570 int>::type* = 0 )
1571{
1572 using work_tag = typename Kokkos::RangePolicy<ExecParameters...>::work_tag;
1573 using execution_space =
1574 typename Kokkos::RangePolicy<ExecParameters...>::execution_space;
1575
1576 using team_policy_type =
1577 Kokkos::TeamPolicy<TeamOpTag, execution_space,
1578 Kokkos::Schedule<Kokkos::Dynamic>>;
1579 team_policy_type team_policy( exec_policy.end() - exec_policy.begin(),
1580 Kokkos::AUTO );
1581
1582 using memory_space = typename LinkedCellType::memory_space;
1583
1584 static_assert( is_accessible_from<memory_space, execution_space>{}, "" );
1585
1586 // Cannot iterate over range that was not binned.
1587 assert( exec_policy.begin() == list.getParticleBegin() );
1588 assert( exec_policy.end() == list.getParticleEnd() );
1589
1590 Impl::LinkedCellParallelReduce<
1591 work_tag, FunctorType, team_policy_type, LinkedCellType,
1592 typename LinkedCellType::CountView, ReduceType>
1593 lcl_par( str, team_policy, functor, list, reduce_val,
1594 exec_policy.begin() );
1595}
1596
1597} // end namespace Cabana
1598
1599#endif // end CABANA_PARALLEL_HPP
Cabana_ExecutionPolicy.hpp
SIMD execution policy.
Cabana_LinkedCellList.hpp
Linked cell list binning (spatial sorting) and neighbor iteration.
Cabana_NeighborList.hpp
Neighbor list interface.
Cabana_Types.hpp
Memory access type checking.
Cabana::FirstNeighborsTag
Loop over particle neighbors.
Definition Cabana_Parallel.hpp:191
Cabana::NeighborList
Neighbor list interface. Provides an interface callable at the functor level that gives access to nei...
Definition Cabana_NeighborList.hpp:114
Cabana::SecondNeighborsTag
Loop over particle neighbors (first) and neighbor's neighbors (second)
Definition Cabana_Parallel.hpp:196
Cabana::SerialOpTag
Neighbor operations are executed in serial on each particle thread.
Definition Cabana_Parallel.hpp:201
Cabana::SimdPolicy
Execution policy over a range of 2d indices.
Definition Cabana_ExecutionPolicy.hpp:82
Cabana::TeamOpTag
Neighbor operations are executed with team parallelism.
Definition Cabana_Parallel.hpp:206
Cabana::TeamVectorOpTag
Neighbor operations are executed with team vector parallelism.
Definition Cabana_Parallel.hpp:211
Cabana
Core: particle data structures and algorithms.
Definition Cabana_AoSoA.hpp:36
Cabana::for_each_neighbor
KOKKOS_INLINE_FUNCTION void for_each_neighbor(const IndexType i, const FunctorType &neighbor_functor, const NeighborListType &list, const FirstNeighborsTag)
Execute functor in serial within existing parallel kernel over particle first neighbors.
Definition Cabana_Parallel.hpp:1060
Cabana::neighbor_parallel_for
void neighbor_parallel_for(const Kokkos::RangePolicy< ExecParameters... > &exec_policy, const FunctorType &functor, const NeighborListType &list, const FirstNeighborsTag, const SerialOpTag, const std::string &str="", typename std::enable_if<(!is_linked_cell_list< NeighborListType >::value), int >::type *=0)
Execute functor in parallel according to the execution policy over particles with a thread-local seri...
Definition Cabana_Parallel.hpp:252
Cabana::size
auto size(SliceType slice, typename std::enable_if< is_slice< SliceType >::value, int >::type *=0)
Check slice size (differs from Kokkos View).
Definition Cabana_Slice.hpp:1012
Cabana::neighbor_parallel_reduce
void neighbor_parallel_reduce(const Kokkos::RangePolicy< ExecParameters... > &exec_policy, const FunctorType &functor, const NeighborListType &list, const FirstNeighborsTag, const SerialOpTag, ReduceType &reduce_val, const std::string &str="", typename std::enable_if<(!is_linked_cell_list< NeighborListType >::value), int >::type *=0)
Execute functor reduction in parallel according to the execution policy over particles with a thread-...
Definition Cabana_Parallel.hpp:640
Cabana::simd_parallel_for
void simd_parallel_for(const SimdPolicy< VectorLength, ExecParameters... > &exec_policy, const FunctorType &functor, const std::string &str="")
Execute a vectorized functor in parallel with a 2d execution policy.
Definition Cabana_Parallel.hpp:174
Cabana::is_accessible_from
Definition Cabana_Types.hpp:88
Cabana::is_linked_cell_list
LinkedCellList static type checker.
Definition Cabana_LinkedCellList.hpp:721