44 #ifndef KOKKOS_TASKPOLICY_HPP 45 #define KOKKOS_TASKPOLICY_HPP 49 #include <Kokkos_Core_fwd.hpp> 55 #if ( defined( KOKKOS_HAVE_CUDA ) ) 56 #if ( 8000 <= CUDA_VERSION ) && \ 57 defined( KOKKOS_CUDA_USE_RELOCATABLE_DEVICE_CODE ) 59 #define KOKKOS_ENABLE_TASKPOLICY 63 #define KOKKOS_ENABLE_TASKPOLICY 67 #if defined( KOKKOS_ENABLE_TASKPOLICY ) 71 #include <Kokkos_MemoryPool.hpp> 72 #include <impl/Kokkos_Tags.hpp> 73 #include <impl/Kokkos_TaskQueue.hpp> 79 enum TaskType { TaskTeam = Impl::TaskBase<void,void,void>::TaskTeam
80 , TaskSingle = Impl::TaskBase<void,void,void>::TaskSingle };
82 enum TaskPriority { TaskHighPriority = 0
83 , TaskRegularPriority = 1
84 , TaskLowPriority = 2 };
86 template<
typename Space >
89 template<
typename Space >
90 void wait( TaskScheduler< Space >
const & );
113 template<
typename Space ,
typename ResultType ,
typename FunctorType >
116 template<
typename Space >
132 template<
typename Arg1 ,
typename Arg2 >
136 template<
typename >
friend class TaskScheduler ;
137 template<
typename ,
typename >
friend class Future ;
138 template<
typename ,
typename ,
typename >
friend class Impl::TaskBase ;
140 enum { Arg1_is_space = Kokkos::Impl::is_space< Arg1 >::value };
141 enum { Arg2_is_space = Kokkos::Impl::is_space< Arg2 >::value };
142 enum { Arg1_is_value = ! Arg1_is_space &&
143 ! std::is_same< Arg1 , void >::value };
144 enum { Arg2_is_value = ! Arg2_is_space &&
145 ! std::is_same< Arg2 , void >::value };
147 static_assert( ! ( Arg1_is_space && Arg2_is_space )
148 ,
"Future cannot be given two spaces" );
150 static_assert( ! ( Arg1_is_value && Arg2_is_value )
151 ,
"Future cannot be given two value types" );
154 typename std::conditional< Arg1_is_value , Arg1 ,
155 typename std::conditional< Arg2_is_value , Arg2 ,
void 159 typename std::conditional< Arg1_is_space , Arg1 ,
160 typename std::conditional< Arg2_is_space , Arg2 ,
void 163 using task_base = Impl::TaskBase< Space , ValueType , void > ;
164 using queue_type = Impl::TaskQueue< Space > ;
168 KOKKOS_INLINE_FUNCTION
explicit 169 Future( task_base * task ) : m_task(0)
170 {
if ( task ) queue_type::assign( & m_task , task ); }
176 using execution_space =
typename Space::execution_space ;
177 using value_type = ValueType ;
181 KOKKOS_INLINE_FUNCTION
182 bool is_null()
const {
return 0 == m_task ; }
184 KOKKOS_INLINE_FUNCTION
185 int reference_count()
const 186 {
return 0 != m_task ? m_task->reference_count() : 0 ; }
190 KOKKOS_INLINE_FUNCTION
191 ~
Future() {
if ( m_task ) queue_type::assign( & m_task , (task_base*)0 ); }
195 KOKKOS_INLINE_FUNCTION
196 constexpr
Future() noexcept : m_task(0) {}
198 KOKKOS_INLINE_FUNCTION
200 : m_task( rhs.m_task ) { rhs.m_task = 0 ; }
202 KOKKOS_INLINE_FUNCTION
205 {
if ( rhs.m_task ) queue_type::assign( & m_task , rhs.m_task ); }
207 KOKKOS_INLINE_FUNCTION
210 if ( m_task ) queue_type::assign( & m_task , (task_base*)0 );
211 m_task = rhs.m_task ;
216 KOKKOS_INLINE_FUNCTION
219 if ( m_task || rhs.m_task ) queue_type::assign( & m_task , rhs.m_task );
225 template<
class A1 ,
class A2 >
226 KOKKOS_INLINE_FUNCTION
228 : m_task( rhs.m_task )
231 ( std::is_same< Space , void >::value ||
232 std::is_same< Space ,
typename Future<A1,A2>::Space >::value
233 ,
"Assigned Futures must have the same space" );
236 ( std::is_same< value_type , void >::value ||
237 std::is_same< value_type ,
typename Future<A1,A2>::value_type >::value
238 ,
"Assigned Futures must have the same value_type" );
243 template<
class A1 ,
class A2 >
244 KOKKOS_INLINE_FUNCTION
249 ( std::is_same< Space , void >::value ||
250 std::is_same< Space ,
typename Future<A1,A2>::Space >::value
251 ,
"Assigned Futures must have the same space" );
254 ( std::is_same< value_type , void >::value ||
255 std::is_same< value_type ,
typename Future<A1,A2>::value_type >::value
256 ,
"Assigned Futures must have the same value_type" );
258 if ( rhs.m_task ) queue_type::assign( & m_task , rhs.m_task );
261 template<
class A1 ,
class A2 >
262 KOKKOS_INLINE_FUNCTION
266 ( std::is_same< Space , void >::value ||
267 std::is_same< Space ,
typename Future<A1,A2>::Space >::value
268 ,
"Assigned Futures must have the same space" );
271 ( std::is_same< value_type , void >::value ||
272 std::is_same< value_type ,
typename Future<A1,A2>::value_type >::value
273 ,
"Assigned Futures must have the same value_type" );
275 if ( m_task || rhs.m_task ) queue_type::assign( & m_task , rhs.m_task );
279 template<
class A1 ,
class A2 >
280 KOKKOS_INLINE_FUNCTION
284 ( std::is_same< Space , void >::value ||
285 std::is_same< Space ,
typename Future<A1,A2>::Space >::value
286 ,
"Assigned Futures must have the same space" );
289 ( std::is_same< value_type , void >::value ||
290 std::is_same< value_type ,
typename Future<A1,A2>::value_type >::value
291 ,
"Assigned Futures must have the same value_type" );
293 if ( m_task ) queue_type::assign( & m_task , (task_base*) 0 );
294 m_task = rhs.m_task ;
301 KOKKOS_INLINE_FUNCTION
302 typename task_base::get_return_type
306 Kokkos::abort(
"Kokkos:::Future::get ERROR: is_null()");
308 return m_task->get();
319 template<
typename ExecSpace >
324 using track_type = Kokkos::Impl::SharedAllocationTracker ;
325 using queue_type = Kokkos::Impl::TaskQueue< ExecSpace > ;
326 using task_base = Impl::TaskBase< ExecSpace , void , void > ;
329 queue_type * m_queue ;
334 KOKKOS_INLINE_FUNCTION
static 335 void assign( task_base *
const ) {}
338 template<
typename ... Options >
339 KOKKOS_INLINE_FUNCTION
static 340 void assign( task_base *
const task
341 , TaskType
const & arg
342 , Options
const & ... opts )
344 task->m_task_type = arg ;
345 assign( task , opts ... );
349 template<
typename ... Options >
350 KOKKOS_INLINE_FUNCTION
static 351 void assign( task_base *
const task
352 , TaskPriority
const & arg
353 , Options
const & ... opts )
355 task->m_priority = arg ;
356 assign( task , opts ... );
360 template<
typename A1 ,
typename A2 ,
typename ... Options >
361 KOKKOS_INLINE_FUNCTION
static 362 void assign( task_base *
const task
363 , Future< A1 , A2 >
const & arg
364 , Options
const & ... opts )
370 if ( 0 != Kokkos::atomic_exchange(& task->m_next, arg.m_task) ) {
371 Kokkos::abort(
"TaskScheduler ERROR: resetting task dependence");
374 if ( 0 != arg.m_task ) {
377 Kokkos::atomic_increment( &(arg.m_task->m_ref_count) );
380 assign( task , opts ... );
387 using execution_policy = TaskScheduler ;
388 using execution_space = ExecSpace ;
389 using memory_space =
typename queue_type::memory_space ;
390 using member_type = Kokkos::Impl::TaskExec< ExecSpace > ;
392 KOKKOS_INLINE_FUNCTION
393 TaskScheduler() : m_track(), m_queue(0) {}
395 KOKKOS_INLINE_FUNCTION
396 TaskScheduler( TaskScheduler && rhs ) = default ;
398 KOKKOS_INLINE_FUNCTION
399 TaskScheduler( TaskScheduler
const & rhs ) = default ;
401 KOKKOS_INLINE_FUNCTION
402 TaskScheduler & operator = ( TaskScheduler && rhs ) = default ;
404 KOKKOS_INLINE_FUNCTION
405 TaskScheduler & operator = ( TaskScheduler
const & rhs ) = default ;
407 TaskScheduler( memory_space
const & arg_memory_space
408 ,
unsigned const arg_memory_pool_capacity
409 ,
unsigned const arg_memory_pool_log2_superblock = 12 )
413 typedef Kokkos::Impl::SharedAllocationRecord
414 < memory_space ,
typename queue_type::Destroy >
417 record_type * record =
418 record_type::allocate( arg_memory_space
423 m_queue =
new( record->data() )
424 queue_type( arg_memory_space
425 , arg_memory_pool_capacity
426 , arg_memory_pool_log2_superblock );
428 record->m_destroy.m_queue = m_queue ;
430 m_track.assign_allocated_record_to_uninitialized( record );
435 template<
typename FunctorType >
437 size_t spawn_allocation_size()
const 439 using task_type = Impl::TaskBase< execution_space
440 ,
typename FunctorType::value_type
443 return m_queue->allocate_block_size(
sizeof(task_type) );
448 size_t when_all_allocation_size(
int narg )
const 450 using task_base = Kokkos::Impl::TaskBase< ExecSpace , void , void > ;
452 return m_queue->allocate_block_size(
sizeof(task_base) + narg *
sizeof(task_base*) );
463 template<
typename FunctorType ,
typename ... Options >
465 Future< typename FunctorType::value_type , ExecSpace >
466 task_spawn( FunctorType
const & arg_functor
467 , Options
const & ... arg_options
470 using value_type =
typename FunctorType::value_type ;
471 using future_type = Future< value_type , execution_space > ;
472 using task_type = Impl::TaskBase< execution_space
480 m_queue->iff_single_thread_recursive_execute();
488 reinterpret_cast< task_type *
>(m_queue->allocate(
sizeof(task_type)));
493 new ( f.m_task ) task_type( arg_functor );
498 f.m_task->m_queue = m_queue ;
499 f.m_task->m_ref_count = 2 ;
500 f.m_task->m_alloc_size =
sizeof(task_type);
502 assign( f.m_task , arg_options... );
506 f.m_task->m_apply = task_type::apply ;
508 m_queue->schedule( f.m_task );
521 template<
typename FunctorType ,
typename ... Options >
523 Future< typename FunctorType::value_type , ExecSpace >
524 host_spawn( FunctorType
const & arg_functor
525 , Options
const & ... arg_options
528 using value_type =
typename FunctorType::value_type ;
529 using future_type = Future< value_type , execution_space > ;
530 using task_type = Impl::TaskBase< execution_space
534 if ( m_queue == 0 ) {
535 Kokkos::abort(
"Kokkos::TaskScheduler not initialized");
542 reinterpret_cast<task_type*
>( m_queue->allocate(
sizeof(task_type)) );
547 new( f.m_task ) task_type( arg_functor );
552 f.m_task->m_queue = m_queue ;
553 f.m_task->m_ref_count = 2 ;
554 f.m_task->m_alloc_size =
sizeof(task_type);
556 assign( f.m_task , arg_options... );
561 queue_type::specialization::template
562 proc_set_apply< FunctorType >( & f.m_task->m_apply );
564 m_queue->schedule( f.m_task );
572 template<
typename A1 ,
typename A2 >
575 when_all(
int narg , Future< A1 , A2 >
const *
const arg )
const 578 ( std::is_same< execution_space
579 ,
typename Future< A1 , A2 >::execution_space
581 ,
"Future must have same execution space" );
583 using future_type = Future< ExecSpace > ;
584 using task_base = Kokkos::Impl::TaskBase< ExecSpace , void , void > ;
588 size_t const size =
sizeof(task_base) + narg *
sizeof(task_base*);
591 reinterpret_cast< task_base *
>( m_queue->allocate( size ) );
595 new( f.m_task ) task_base();
600 f.m_task->m_queue = m_queue ;
601 f.m_task->m_ref_count = 2 ;
602 f.m_task->m_alloc_size = size ;
603 f.m_task->m_dep_count = narg ;
604 f.m_task->m_task_type = task_base::Aggregate ;
606 task_base **
const dep = f.m_task->aggregate_dependences();
612 for (
int i = 0 ; i < narg ; ++i ) {
613 task_base *
const t = dep[i] = arg[i].m_task ;
615 Kokkos::atomic_increment( &(t->m_ref_count) );
619 m_queue->schedule( f.m_task );
631 template<
class FunctorType ,
typename ... Options >
633 void respawn( FunctorType * task_self
634 , Options
const & ... arg_options )
const 636 using value_type =
typename FunctorType::value_type ;
637 using task_type = Impl::TaskBase< execution_space
641 task_base *
const zero = (task_base *) 0 ;
642 task_base *
const lock = (task_base *) task_base::LockTag ;
643 task_type *
const task =
static_cast< task_type *
>( task_self );
651 if ( lock != Kokkos::atomic_exchange( & task->m_next, zero ) ) {
652 Kokkos::abort(
"TaskScheduler::respawn ERROR: already respawned");
655 assign( task , arg_options... );
664 template<
typename S >
666 void Kokkos::wait( Kokkos::TaskScheduler< S >
const & );
671 int allocation_capacity() const noexcept
672 {
return m_queue->m_memory.get_mem_size(); }
674 KOKKOS_INLINE_FUNCTION
675 int allocated_task_count() const noexcept
676 {
return m_queue->m_count_alloc ; }
678 KOKKOS_INLINE_FUNCTION
679 int allocated_task_count_max() const noexcept
680 {
return m_queue->m_max_alloc ; }
682 KOKKOS_INLINE_FUNCTION
683 long allocated_task_count_accum() const noexcept
684 {
return m_queue->m_accum_alloc ; }
688 template<
typename ExecSpace >
690 void wait( TaskScheduler< ExecSpace >
const & policy )
691 { policy.m_queue->execute(); }
694 template<
typename ExecSpace >
696 TaskPolicy = TaskScheduler< ExecSpace > ;