44 #ifndef KOKKOS_TASKSCHEDULER_HPP 45 #define KOKKOS_TASKSCHEDULER_HPP 49 #include <Kokkos_Macros.hpp> 50 #if defined( KOKKOS_ENABLE_TASKDAG ) 52 #include <Kokkos_Core_fwd.hpp> 55 #include <Kokkos_MemoryPool.hpp> 56 #include <impl/Kokkos_Tags.hpp> 64 template<
typename Arg1 =
void ,
typename Arg2 =
void >
67 template<
typename Space >
70 template<
typename Space >
71 void wait( TaskScheduler< Space >
const & );
73 template<
typename Space >
74 struct is_scheduler :
public std::false_type {};
76 template<
typename Space >
77 struct is_scheduler< TaskScheduler< Space > > :
public std::true_type {};
81 #include <impl/Kokkos_TaskQueue.hpp> 103 template<
typename Space ,
typename ResultType ,
typename FunctorType >
120 template<
typename Arg1 ,
typename Arg2 >
124 template<
typename >
friend class TaskScheduler ;
125 template<
typename ,
typename >
friend class Future ;
126 template<
typename ,
typename ,
typename >
friend class Impl::TaskBase ;
128 enum { Arg1_is_space = Kokkos::is_space< Arg1 >::value };
129 enum { Arg2_is_space = Kokkos::is_space< Arg2 >::value };
130 enum { Arg1_is_value = ! Arg1_is_space &&
131 ! std::is_same< Arg1 , void >::value };
132 enum { Arg2_is_value = ! Arg2_is_space &&
133 ! std::is_same< Arg2 , void >::value };
135 static_assert( ! ( Arg1_is_space && Arg2_is_space )
136 ,
"Future cannot be given two spaces" );
138 static_assert( ! ( Arg1_is_value && Arg2_is_value )
139 ,
"Future cannot be given two value types" );
142 typename std::conditional< Arg1_is_value , Arg1 ,
143 typename std::conditional< Arg2_is_value , Arg2 ,
void 147 typename std::conditional< Arg1_is_space , Arg1 ,
148 typename std::conditional< Arg2_is_space , Arg2 ,
void 151 using task_base = Impl::TaskBase< Space , ValueType , void > ;
152 using queue_type = Impl::TaskQueue< Space > ;
156 KOKKOS_INLINE_FUNCTION
explicit 157 Future( task_base * task ) : m_task(0)
158 {
if ( task ) queue_type::assign( & m_task , task ); }
164 using execution_space =
typename Space::execution_space ;
165 using value_type = ValueType ;
169 KOKKOS_INLINE_FUNCTION
170 bool is_null()
const {
return 0 == m_task ; }
172 KOKKOS_INLINE_FUNCTION
173 int reference_count()
const 174 {
return 0 != m_task ? m_task->reference_count() : 0 ; }
178 KOKKOS_INLINE_FUNCTION
180 {
if ( m_task ) queue_type::assign( & m_task , (task_base*)0 ); }
184 KOKKOS_INLINE_FUNCTION
185 ~Future() { clear(); }
189 KOKKOS_INLINE_FUNCTION
190 constexpr Future() noexcept : m_task(0) {}
192 KOKKOS_INLINE_FUNCTION
193 Future( Future && rhs )
194 : m_task( rhs.m_task ) { rhs.m_task = 0 ; }
196 KOKKOS_INLINE_FUNCTION
197 Future(
const Future & rhs )
199 {
if ( rhs.m_task ) queue_type::assign( & m_task , rhs.m_task ); }
201 KOKKOS_INLINE_FUNCTION
202 Future & operator = ( Future && rhs )
205 m_task = rhs.m_task ;
210 KOKKOS_INLINE_FUNCTION
211 Future & operator = (
const Future & rhs )
213 if ( m_task || rhs.m_task ) queue_type::assign( & m_task , rhs.m_task );
219 template<
class A1 ,
class A2 >
220 KOKKOS_INLINE_FUNCTION
222 : m_task( rhs.m_task )
225 ( std::is_same< Space , void >::value ||
226 std::is_same< Space ,
typename Future<A1,A2>::Space >::value
227 ,
"Assigned Futures must have the same space" );
230 ( std::is_same< value_type , void >::value ||
231 std::is_same< value_type ,
typename Future<A1,A2>::value_type >::value
232 ,
"Assigned Futures must have the same value_type" );
237 template<
class A1 ,
class A2 >
238 KOKKOS_INLINE_FUNCTION
243 ( std::is_same< Space , void >::value ||
244 std::is_same< Space ,
typename Future<A1,A2>::Space >::value
245 ,
"Assigned Futures must have the same space" );
248 ( std::is_same< value_type , void >::value ||
249 std::is_same< value_type ,
typename Future<A1,A2>::value_type >::value
250 ,
"Assigned Futures must have the same value_type" );
252 if ( rhs.m_task ) queue_type::assign( & m_task , rhs.m_task );
255 template<
class A1 ,
class A2 >
256 KOKKOS_INLINE_FUNCTION
260 ( std::is_same< Space , void >::value ||
261 std::is_same< Space ,
typename Future<A1,A2>::Space >::value
262 ,
"Assigned Futures must have the same space" );
265 ( std::is_same< value_type , void >::value ||
266 std::is_same< value_type ,
typename Future<A1,A2>::value_type >::value
267 ,
"Assigned Futures must have the same value_type" );
269 if ( m_task || rhs.m_task ) queue_type::assign( & m_task , rhs.m_task );
273 template<
class A1 ,
class A2 >
274 KOKKOS_INLINE_FUNCTION
278 ( std::is_same< Space , void >::value ||
279 std::is_same< Space ,
typename Future<A1,A2>::Space >::value
280 ,
"Assigned Futures must have the same space" );
283 ( std::is_same< value_type , void >::value ||
284 std::is_same< value_type ,
typename Future<A1,A2>::value_type >::value
285 ,
"Assigned Futures must have the same value_type" );
288 m_task = rhs.m_task ;
295 KOKKOS_INLINE_FUNCTION
296 typename task_base::get_return_type
300 Kokkos::abort(
"Kokkos:::Future::get ERROR: is_null()");
302 return m_task->get();
307 template<
typename ,
typename ExecSpace =
void >
308 struct is_future :
public std::false_type {};
310 template<
typename Arg1 ,
typename Arg2 ,
typename ExecSpace >
311 struct is_future< Future<Arg1,Arg2> , ExecSpace >
312 :
public std::integral_constant
314 ( std::is_same< ExecSpace , void >::value ||
315 std::is_same< ExecSpace
316 , typename Future<Arg1,Arg2>::execution_space >::value )
326 enum class TaskPriority : int { High = 0
340 template<
int TaskEnum ,
typename DepFutureType >
341 struct TaskPolicyData
343 using execution_space =
typename DepFutureType::execution_space ;
344 using scheduler_type = TaskScheduler< execution_space > ;
346 enum :
int { m_task_type = TaskEnum };
348 scheduler_type
const * m_scheduler ;
349 DepFutureType
const m_dependence ;
352 TaskPolicyData() = delete ;
353 TaskPolicyData( TaskPolicyData && ) = default ;
354 TaskPolicyData( TaskPolicyData
const & ) = default ;
355 TaskPolicyData & operator = ( TaskPolicyData && ) = default ;
356 TaskPolicyData & operator = ( TaskPolicyData
const & ) = default ;
358 KOKKOS_INLINE_FUNCTION
359 TaskPolicyData( DepFutureType
const & arg_future
360 , Kokkos::TaskPriority
const & arg_priority )
362 , m_dependence( arg_future )
363 , m_priority( static_cast<int>( arg_priority ) )
366 KOKKOS_INLINE_FUNCTION
367 TaskPolicyData( scheduler_type
const & arg_scheduler
368 , Kokkos::TaskPriority
const & arg_priority )
369 : m_scheduler( & arg_scheduler )
371 , m_priority( static_cast<int>( arg_priority ) )
374 KOKKOS_INLINE_FUNCTION
375 TaskPolicyData( scheduler_type
const & arg_scheduler
376 , DepFutureType
const & arg_future
377 , Kokkos::TaskPriority
const & arg_priority )
378 : m_scheduler( & arg_scheduler )
379 , m_dependence( arg_future )
380 , m_priority( static_cast<int>( arg_priority ) )
392 template<
typename ExecSpace >
397 using track_type = Kokkos::Impl::SharedAllocationTracker ;
398 using queue_type = Kokkos::Impl::TaskQueue< ExecSpace > ;
399 using task_base = Impl::TaskBase< ExecSpace , void , void > ;
402 queue_type * m_queue ;
408 using execution_space = ExecSpace ;
409 using memory_space =
typename queue_type::memory_space ;
410 using memory_pool =
typename queue_type::memory_pool ;
412 typename Kokkos::Impl::TaskQueueSpecialization< ExecSpace >::member_type ;
414 KOKKOS_INLINE_FUNCTION
415 TaskScheduler() : m_track(), m_queue(0) {}
417 KOKKOS_INLINE_FUNCTION
418 TaskScheduler( TaskScheduler && rhs ) = default ;
420 KOKKOS_INLINE_FUNCTION
421 TaskScheduler( TaskScheduler
const & rhs ) = default ;
423 KOKKOS_INLINE_FUNCTION
424 TaskScheduler & operator = ( TaskScheduler && rhs ) = default ;
426 KOKKOS_INLINE_FUNCTION
427 TaskScheduler & operator = ( TaskScheduler
const & rhs ) = default ;
429 TaskScheduler( memory_pool
const & arg_memory_pool )
433 typedef Kokkos::Impl::SharedAllocationRecord
434 < memory_space ,
typename queue_type::Destroy >
437 record_type * record =
438 record_type::allocate( memory_space()
443 m_queue =
new( record->data() ) queue_type( arg_memory_pool );
445 record->m_destroy.m_queue = m_queue ;
447 m_track.assign_allocated_record_to_uninitialized( record );
450 TaskScheduler( memory_space
const & arg_memory_space
451 ,
size_t const mempool_capacity
452 ,
unsigned const mempool_min_block_size
453 ,
unsigned const mempool_max_block_size
454 ,
unsigned const mempool_superblock_size
456 : TaskScheduler( memory_pool( arg_memory_space
458 , mempool_min_block_size
459 , mempool_max_block_size
460 , mempool_superblock_size ) )
465 KOKKOS_INLINE_FUNCTION
466 memory_pool * memory() const noexcept
467 {
return m_queue ? m_queue->m_memory : (memory_pool*) 0 ; }
471 template<
typename FunctorType >
473 size_t spawn_allocation_size()
const 475 using task_type = Impl::TaskBase< execution_space
476 ,
typename FunctorType::value_type
479 return m_queue->allocate_block_size(
sizeof(task_type) );
484 size_t when_all_allocation_size(
int narg )
const 486 using task_base = Kokkos::Impl::TaskBase< ExecSpace , void , void > ;
488 return m_queue->allocate_block_size(
sizeof(task_base) + narg *
sizeof(task_base*) );
493 template<
int TaskEnum ,
typename DepFutureType ,
typename FunctorType >
494 KOKKOS_FUNCTION
static 496 spawn( Impl::TaskPolicyData<TaskEnum,DepFutureType>
const & arg_policy
497 ,
typename task_base::function_type arg_function
498 , FunctorType && arg_functor
501 using value_type =
typename FunctorType::value_type ;
503 using task_type = Impl::TaskBase< execution_space
507 queue_type *
const queue =
508 arg_policy.m_scheduler ? arg_policy.m_scheduler->m_queue : (
509 arg_policy.m_dependence.m_task
510 ? arg_policy.m_dependence.m_task->m_queue
514 Kokkos::abort(
"Kokkos spawn requires scheduler or non-null Future");
517 if ( arg_policy.m_dependence.m_task != 0 &&
518 arg_policy.m_dependence.m_task->m_queue != queue ) {
519 Kokkos::abort(
"Kokkos spawn given incompatible scheduler and Future");
526 queue->iff_single_thread_recursive_execute();
534 reinterpret_cast< task_type *
>(queue->allocate(
sizeof(task_type)));
543 task_type( arg_function
545 , arg_policy.m_dependence.m_task
547 ,
int(
sizeof(task_type))
548 ,
int(arg_policy.m_task_type)
549 , int(arg_policy.m_priority)
550 , std::move(arg_functor) );
557 queue->schedule_runnable( f.m_task );
565 template<
typename FunctorType ,
typename A1 ,
typename A2 >
566 KOKKOS_FUNCTION
static 568 respawn( FunctorType * arg_self
570 , TaskPriority
const & arg_priority
575 using value_type =
typename FunctorType::value_type ;
576 using task_type = Impl::TaskBase< execution_space
580 task_type *
const task =
static_cast< task_type *
>( arg_self );
582 task->m_priority =
static_cast<int>(arg_priority);
584 task->add_dependence( arg_dependence.m_task );
593 template<
typename A1 ,
typename A2 >
594 KOKKOS_FUNCTION
static 599 using task_base = Kokkos::Impl::TaskBase< execution_space , void , void > ;
605 queue_type * queue = 0 ;
607 for (
int i = 0 ; i < narg ; ++i ) {
608 task_base *
const t = arg[i].m_task ;
611 Kokkos::atomic_increment( &(t->m_ref_count) );
615 else if ( queue != t->m_queue ) {
616 Kokkos::abort(
"Kokkos when_all Futures must be in the same scheduler" );
623 size_t const size =
sizeof(task_base) + narg *
sizeof(task_base*);
626 reinterpret_cast< task_base *
>( queue->allocate( size ) );
633 new( f.m_task ) task_base( queue
641 task_base **
const dep = f.m_task->aggregate_dependences();
643 for (
int i = 0 ; i < narg ; ++i ) { dep[i] = arg[i].m_task ; }
645 queue->schedule_aggregate( f.m_task );
656 KOKKOS_INLINE_FUNCTION
657 int allocation_capacity() const noexcept
658 {
return m_queue->m_memory.capacity(); }
660 KOKKOS_INLINE_FUNCTION
661 int allocated_task_count() const noexcept
662 {
return m_queue->m_count_alloc ; }
664 KOKKOS_INLINE_FUNCTION
665 int allocated_task_count_max() const noexcept
666 {
return m_queue->m_max_alloc ; }
668 KOKKOS_INLINE_FUNCTION
669 long allocated_task_count_accum() const noexcept
670 {
return m_queue->m_accum_alloc ; }
674 template<
typename S >
676 void Kokkos::wait( Kokkos::TaskScheduler< S >
const & );
690 template<
typename T >
691 Kokkos::Impl::TaskPolicyData
692 < Kokkos::Impl::TaskBase<void,void,void>::TaskTeam
693 ,
typename std::conditional< Kokkos::is_future< T >::value , T ,
696 KOKKOS_INLINE_FUNCTION
697 TaskTeam( T
const & arg
698 , TaskPriority
const & arg_priority = TaskPriority::Regular
701 static_assert( Kokkos::is_future<T>::value ||
702 Kokkos::is_scheduler<T>::value
703 ,
"Kokkos TaskTeam argument must be Future or TaskScheduler" );
706 Kokkos::Impl::TaskPolicyData
707 < Kokkos::Impl::TaskBase<void,void,void>::TaskTeam
708 ,
typename std::conditional< Kokkos::is_future< T >::value , T ,
710 >( arg , arg_priority );
713 template<
typename E ,
typename F >
715 TaskPolicyData< Kokkos::Impl::TaskBase<void,void,void>::TaskTeam , F >
716 KOKKOS_INLINE_FUNCTION
717 TaskTeam( TaskScheduler<E>
const & arg_scheduler
718 , F
const & arg_future
719 ,
typename std::enable_if< Kokkos::is_future<F>::value ,
720 TaskPriority >::type
const & arg_priority = TaskPriority::Regular
724 Kokkos::Impl::TaskPolicyData
725 < Kokkos::Impl::TaskBase<void,void,void>::TaskTeam , F >
726 ( arg_scheduler , arg_future , arg_priority );
731 template<
typename T >
732 Kokkos::Impl::TaskPolicyData
733 < Kokkos::Impl::TaskBase<void,void,void>::TaskSingle
734 ,
typename std::conditional< Kokkos::is_future< T >::value , T ,
737 KOKKOS_INLINE_FUNCTION
738 TaskSingle( T
const & arg
739 , TaskPriority
const & arg_priority = TaskPriority::Regular
742 static_assert( Kokkos::is_future<T>::value ||
743 Kokkos::is_scheduler<T>::value
744 ,
"Kokkos TaskSingle argument must be Future or TaskScheduler" );
747 Kokkos::Impl::TaskPolicyData
748 < Kokkos::Impl::TaskBase<void,void,void>::TaskSingle
749 ,
typename std::conditional< Kokkos::is_future< T >::value , T ,
751 >( arg , arg_priority );
754 template<
typename E ,
typename F >
756 TaskPolicyData< Kokkos::Impl::TaskBase<void,void,void>::TaskSingle , F >
757 KOKKOS_INLINE_FUNCTION
758 TaskSingle( TaskScheduler<E>
const & arg_scheduler
759 , F
const & arg_future
760 ,
typename std::enable_if< Kokkos::is_future<F>::value ,
761 TaskPriority >::type
const & arg_priority = TaskPriority::Regular
765 Kokkos::Impl::TaskPolicyData
766 < Kokkos::Impl::TaskBase<void,void,void>::TaskSingle , F >
767 ( arg_scheduler , arg_future , arg_priority );
778 template<
int TaskEnum
779 ,
typename DepFutureType
780 ,
typename FunctorType >
781 Future<
typename FunctorType::value_type
782 ,
typename DepFutureType::execution_space >
783 host_spawn( Impl::TaskPolicyData<TaskEnum,DepFutureType>
const & arg_policy
784 , FunctorType && arg_functor
787 using exec_space =
typename DepFutureType::execution_space ;
788 using scheduler = TaskScheduler< exec_space > ;
790 typedef Impl::TaskBase< exec_space
791 ,
typename FunctorType::value_type
795 static_assert( TaskEnum == task_type::TaskTeam ||
796 TaskEnum == task_type::TaskSingle
797 ,
"Kokkos host_spawn requires TaskTeam or TaskSingle" );
801 typename task_type::function_type
const ptr =
802 Kokkos::Impl::TaskQueueSpecialization< exec_space >::
803 template get_function_pointer< task_type >();
805 return scheduler::spawn( arg_policy , ptr , std::move(arg_functor) );
814 template<
int TaskEnum
815 ,
typename DepFutureType
816 ,
typename FunctorType >
817 Future<
typename FunctorType::value_type
818 ,
typename DepFutureType::execution_space >
819 KOKKOS_INLINE_FUNCTION
820 task_spawn( Impl::TaskPolicyData<TaskEnum,DepFutureType>
const & arg_policy
821 , FunctorType && arg_functor
824 using exec_space =
typename DepFutureType::execution_space ;
825 using scheduler = TaskScheduler< exec_space > ;
827 typedef Impl::TaskBase< exec_space
828 ,
typename FunctorType::value_type
832 #if defined( KOKKOS_ACTIVE_EXECUTION_MEMORY_SPACE_HOST ) && \ 833 defined( KOKKOS_ENABLE_CUDA ) 835 static_assert( ! std::is_same< Kokkos::Cuda , exec_space >::value
836 ,
"Error calling Kokkos::task_spawn for Cuda space within Host code" );
840 static_assert( TaskEnum == task_type::TaskTeam ||
841 TaskEnum == task_type::TaskSingle
842 ,
"Kokkos host_spawn requires TaskTeam or TaskSingle" );
844 typename task_type::function_type
const ptr = task_type::apply ;
846 return scheduler::spawn( arg_policy , ptr , std::move(arg_functor) );
854 template<
typename FunctorType ,
typename T >
856 KOKKOS_INLINE_FUNCTION
859 , TaskPriority
const & arg_priority = TaskPriority::Regular
862 static_assert( Kokkos::is_future<T>::value ||
863 Kokkos::is_scheduler<T>::value
864 ,
"Kokkos respawn argument must be Future or TaskScheduler" );
866 TaskScheduler< typename T::execution_space >::
867 respawn( arg_self , arg , arg_priority );
872 template<
typename A1 ,
typename A2 >
873 KOKKOS_INLINE_FUNCTION
879 return TaskScheduler< typename Future<A1,A2>::execution_space >::
880 when_all( arg , narg );
886 template<
typename ExecSpace >
888 void wait( TaskScheduler< ExecSpace >
const & scheduler )
889 { scheduler.m_queue->execute(); }
Future< typename FunctorType::value_type, typename DepFutureType::execution_space > KOKKOS_INLINE_FUNCTION task_spawn(Impl::TaskPolicyData< TaskEnum, DepFutureType > const &arg_policy, FunctorType &&arg_functor)
A task spawns a task with options.
void KOKKOS_INLINE_FUNCTION respawn(FunctorType *arg_self, T const &arg, TaskPriority const &arg_priority=TaskPriority::Regular)
A task respawns itself with options.
Future< typename FunctorType::value_type, typename DepFutureType::execution_space > host_spawn(Impl::TaskPolicyData< TaskEnum, DepFutureType > const &arg_policy, FunctorType &&arg_functor)
A host control thread spawns a task with options.