Kokkos Core Kernels Package  Version of the Day
Kokkos_TaskScheduler.hpp
1 /*
2 //@HEADER
3 // ************************************************************************
4 //
5 // Kokkos v. 2.0
6 // Copyright (2014) Sandia Corporation
7 //
8 // Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
9 // the U.S. Government retains certain rights in this software.
10 //
11 // Redistribution and use in source and binary forms, with or without
12 // modification, are permitted provided that the following conditions are
13 // met:
14 //
15 // 1. Redistributions of source code must retain the above copyright
16 // notice, this list of conditions and the following disclaimer.
17 //
18 // 2. Redistributions in binary form must reproduce the above copyright
19 // notice, this list of conditions and the following disclaimer in the
20 // documentation and/or other materials provided with the distribution.
21 //
22 // 3. Neither the name of the Corporation nor the names of the
23 // contributors may be used to endorse or promote products derived from
24 // this software without specific prior written permission.
25 //
26 // THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
27 // EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
28 // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
29 // PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
30 // CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
31 // EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
32 // PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
33 // PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
34 // LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
35 // NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
36 // SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
37 //
38 // Questions? Contact H. Carter Edwards (hcedwar@sandia.gov)
39 //
40 // ************************************************************************
41 //@HEADER
42 */
43 
44 #ifndef KOKKOS_TASKPOLICY_HPP
45 #define KOKKOS_TASKPOLICY_HPP
46 
47 //----------------------------------------------------------------------------
48 
49 #include <Kokkos_Core_fwd.hpp>
50 
51 // If compiling with CUDA then must be using CUDA 8 or better
52 // and use relocateable device code to enable the task policy.
53 // nvcc relocatable device code option: --relocatable-device-code=true
54 
55 #if ( defined( KOKKOS_HAVE_CUDA ) )
56  #if ( 8000 <= CUDA_VERSION ) && \
57  defined( KOKKOS_CUDA_USE_RELOCATABLE_DEVICE_CODE )
58 
59  #define KOKKOS_ENABLE_TASKPOLICY
60 
61  #endif
62 #else
63  #define KOKKOS_ENABLE_TASKPOLICY
64 #endif
65 
66 
67 #if defined( KOKKOS_ENABLE_TASKPOLICY )
68 
69 //----------------------------------------------------------------------------
70 
71 #include <Kokkos_MemoryPool.hpp>
72 #include <impl/Kokkos_Tags.hpp>
73 #include <impl/Kokkos_TaskQueue.hpp>
74 
75 //----------------------------------------------------------------------------
76 
77 namespace Kokkos {
78 
79 enum TaskType { TaskTeam = Impl::TaskBase<void,void,void>::TaskTeam
80  , TaskSingle = Impl::TaskBase<void,void,void>::TaskSingle };
81 
82 enum TaskPriority { TaskHighPriority = 0
83  , TaskRegularPriority = 1
84  , TaskLowPriority = 2 };
85 
86 template< typename Space >
87 class TaskScheduler ;
88 
89 template< typename Space >
90 void wait( TaskScheduler< Space > const & );
91 
92 } // namespace Kokkos
93 
94 //----------------------------------------------------------------------------
95 
96 namespace Kokkos {
97 namespace Impl {
98 
99 /*\brief Implementation data for task data management, access, and execution.
100  *
101  * CRTP Inheritance structure to allow static_cast from the
102  * task root type and a task's FunctorType.
103  *
104  * TaskBase< Space , ResultType , FunctorType >
105  * : TaskBase< Space , ResultType , void >
106  * , FunctorType
107  * { ... };
108  *
109  * TaskBase< Space , ResultType , void >
110  * : TaskBase< Space , void , void >
111  * { ... };
112  */
113 template< typename Space , typename ResultType , typename FunctorType >
114 class TaskBase ;
115 
116 template< typename Space >
117 class TaskExec ;
118 
119 }} // namespace Kokkos::Impl
120 
121 //----------------------------------------------------------------------------
122 
123 namespace Kokkos {
124 
132 template< typename Arg1 /* = void */ , typename Arg2 /* = void */ >
133 class Future {
134 private:
135 
136  template< typename > friend class TaskScheduler ;
137  template< typename , typename > friend class Future ;
138  template< typename , typename , typename > friend class Impl::TaskBase ;
139 
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 };
146 
147  static_assert( ! ( Arg1_is_space && Arg2_is_space )
148  , "Future cannot be given two spaces" );
149 
150  static_assert( ! ( Arg1_is_value && Arg2_is_value )
151  , "Future cannot be given two value types" );
152 
153  using ValueType =
154  typename std::conditional< Arg1_is_value , Arg1 ,
155  typename std::conditional< Arg2_is_value , Arg2 , void
156  >::type >::type ;
157 
158  using Space =
159  typename std::conditional< Arg1_is_space , Arg1 ,
160  typename std::conditional< Arg2_is_space , Arg2 , void
161  >::type >::type ;
162 
163  using task_base = Impl::TaskBase< Space , ValueType , void > ;
164  using queue_type = Impl::TaskQueue< Space > ;
165 
166  task_base * m_task ;
167 
168  KOKKOS_INLINE_FUNCTION explicit
169  Future( task_base * task ) : m_task(0)
170  { if ( task ) queue_type::assign( & m_task , task ); }
171 
172  //----------------------------------------
173 
174 public:
175 
176  using execution_space = typename Space::execution_space ;
177  using value_type = ValueType ;
178 
179  //----------------------------------------
180 
181  KOKKOS_INLINE_FUNCTION
182  bool is_null() const { return 0 == m_task ; }
183 
184  KOKKOS_INLINE_FUNCTION
185  int reference_count() const
186  { return 0 != m_task ? m_task->reference_count() : 0 ; }
187 
188  //----------------------------------------
189 
190  KOKKOS_INLINE_FUNCTION
191  ~Future() { if ( m_task ) queue_type::assign( & m_task , (task_base*)0 ); }
192 
193  //----------------------------------------
194 
195  KOKKOS_INLINE_FUNCTION
196  constexpr Future() noexcept : m_task(0) {}
197 
198  KOKKOS_INLINE_FUNCTION
199  Future( Future && rhs )
200  : m_task( rhs.m_task ) { rhs.m_task = 0 ; }
201 
202  KOKKOS_INLINE_FUNCTION
203  Future( const Future & rhs )
204  : m_task(0)
205  { if ( rhs.m_task ) queue_type::assign( & m_task , rhs.m_task ); }
206 
207  KOKKOS_INLINE_FUNCTION
208  Future & operator = ( Future && rhs )
209  {
210  if ( m_task ) queue_type::assign( & m_task , (task_base*)0 );
211  m_task = rhs.m_task ;
212  rhs.m_task = 0 ;
213  return *this ;
214  }
215 
216  KOKKOS_INLINE_FUNCTION
217  Future & operator = ( const Future & rhs )
218  {
219  if ( m_task || rhs.m_task ) queue_type::assign( & m_task , rhs.m_task );
220  return *this ;
221  }
222 
223  //----------------------------------------
224 
225  template< class A1 , class A2 >
226  KOKKOS_INLINE_FUNCTION
227  Future( Future<A1,A2> && rhs )
228  : m_task( rhs.m_task )
229  {
230  static_assert
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" );
234 
235  static_assert
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" );
239 
240  rhs.m_task = 0 ;
241  }
242 
243  template< class A1 , class A2 >
244  KOKKOS_INLINE_FUNCTION
245  Future( const Future<A1,A2> & rhs )
246  : m_task(0)
247  {
248  static_assert
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" );
252 
253  static_assert
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" );
257 
258  if ( rhs.m_task ) queue_type::assign( & m_task , rhs.m_task );
259  }
260 
261  template< class A1 , class A2 >
262  KOKKOS_INLINE_FUNCTION
263  Future & operator = ( const Future<A1,A2> & rhs )
264  {
265  static_assert
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" );
269 
270  static_assert
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" );
274 
275  if ( m_task || rhs.m_task ) queue_type::assign( & m_task , rhs.m_task );
276  return *this ;
277  }
278 
279  template< class A1 , class A2 >
280  KOKKOS_INLINE_FUNCTION
281  Future & operator = ( Future<A1,A2> && rhs )
282  {
283  static_assert
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" );
287 
288  static_assert
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" );
292 
293  if ( m_task ) queue_type::assign( & m_task , (task_base*) 0 );
294  m_task = rhs.m_task ;
295  rhs.m_task = 0 ;
296  return *this ;
297  }
298 
299  //----------------------------------------
300 
301  KOKKOS_INLINE_FUNCTION
302  typename task_base::get_return_type
303  get() const
304  {
305  if ( 0 == m_task ) {
306  Kokkos::abort( "Kokkos:::Future::get ERROR: is_null()");
307  }
308  return m_task->get();
309  }
310 };
311 
312 } // namespace Kokkos
313 
314 //----------------------------------------------------------------------------
315 //----------------------------------------------------------------------------
316 
317 namespace Kokkos {
318 
319 template< typename ExecSpace >
320 class TaskScheduler
321 {
322 private:
323 
324  using track_type = Kokkos::Impl::SharedAllocationTracker ;
325  using queue_type = Kokkos::Impl::TaskQueue< ExecSpace > ;
326  using task_base = Impl::TaskBase< ExecSpace , void , void > ;
327 
328  track_type m_track ;
329  queue_type * m_queue ;
330 
331  //----------------------------------------
332  // Process optional arguments to spawn and respawn functions
333 
334  KOKKOS_INLINE_FUNCTION static
335  void assign( task_base * const ) {}
336 
337  // TaskTeam or TaskSingle
338  template< typename ... Options >
339  KOKKOS_INLINE_FUNCTION static
340  void assign( task_base * const task
341  , TaskType const & arg
342  , Options const & ... opts )
343  {
344  task->m_task_type = arg ;
345  assign( task , opts ... );
346  }
347 
348  // TaskHighPriority or TaskRegularPriority or TaskLowPriority
349  template< typename ... Options >
350  KOKKOS_INLINE_FUNCTION static
351  void assign( task_base * const task
352  , TaskPriority const & arg
353  , Options const & ... opts )
354  {
355  task->m_priority = arg ;
356  assign( task , opts ... );
357  }
358 
359  // Future for a dependence
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 )
365  {
366  // Assign dependence to task->m_next
367  // which will be processed within subsequent call to schedule.
368  // Error if the dependence is reset.
369 
370  if ( 0 != Kokkos::atomic_exchange(& task->m_next, arg.m_task) ) {
371  Kokkos::abort("TaskScheduler ERROR: resetting task dependence");
372  }
373 
374  if ( 0 != arg.m_task ) {
375  // The future may be destroyed upon returning from this call
376  // so increment reference count to track this assignment.
377  Kokkos::atomic_increment( &(arg.m_task->m_ref_count) );
378  }
379 
380  assign( task , opts ... );
381  }
382 
383  //----------------------------------------
384 
385 public:
386 
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 > ;
391 
392  KOKKOS_INLINE_FUNCTION
393  TaskScheduler() : m_track(), m_queue(0) {}
394 
395  KOKKOS_INLINE_FUNCTION
396  TaskScheduler( TaskScheduler && rhs ) = default ;
397 
398  KOKKOS_INLINE_FUNCTION
399  TaskScheduler( TaskScheduler const & rhs ) = default ;
400 
401  KOKKOS_INLINE_FUNCTION
402  TaskScheduler & operator = ( TaskScheduler && rhs ) = default ;
403 
404  KOKKOS_INLINE_FUNCTION
405  TaskScheduler & operator = ( TaskScheduler const & rhs ) = default ;
406 
407  TaskScheduler( memory_space const & arg_memory_space
408  , unsigned const arg_memory_pool_capacity
409  , unsigned const arg_memory_pool_log2_superblock = 12 )
410  : m_track()
411  , m_queue(0)
412  {
413  typedef Kokkos::Impl::SharedAllocationRecord
414  < memory_space , typename queue_type::Destroy >
415  record_type ;
416 
417  record_type * record =
418  record_type::allocate( arg_memory_space
419  , "TaskQueue"
420  , sizeof(queue_type)
421  );
422 
423  m_queue = new( record->data() )
424  queue_type( arg_memory_space
425  , arg_memory_pool_capacity
426  , arg_memory_pool_log2_superblock );
427 
428  record->m_destroy.m_queue = m_queue ;
429 
430  m_track.assign_allocated_record_to_uninitialized( record );
431  }
432 
433  //----------------------------------------
435  template< typename FunctorType >
436  KOKKOS_FUNCTION
437  size_t spawn_allocation_size() const
438  {
439  using task_type = Impl::TaskBase< execution_space
440  , typename FunctorType::value_type
441  , FunctorType > ;
442 
443  return m_queue->allocate_block_size( sizeof(task_type) );
444  }
445 
447  KOKKOS_FUNCTION
448  size_t when_all_allocation_size( int narg ) const
449  {
450  using task_base = Kokkos::Impl::TaskBase< ExecSpace , void , void > ;
451 
452  return m_queue->allocate_block_size( sizeof(task_base) + narg * sizeof(task_base*) );
453  }
454 
455  //----------------------------------------
456 
463  template< typename FunctorType , typename ... Options >
464  KOKKOS_FUNCTION
465  Future< typename FunctorType::value_type , ExecSpace >
466  task_spawn( FunctorType const & arg_functor
467  , Options const & ... arg_options
468  ) const
469  {
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
473  , value_type
474  , FunctorType > ;
475 
476  //----------------------------------------
477  // Give single-thread back-ends an opportunity to clear
478  // queue of ready tasks before allocating a new task
479 
480  m_queue->iff_single_thread_recursive_execute();
481 
482  //----------------------------------------
483 
484  future_type f ;
485 
486  // Allocate task from memory pool
487  f.m_task =
488  reinterpret_cast< task_type * >(m_queue->allocate(sizeof(task_type)));
489 
490  if ( f.m_task ) {
491 
492  // Placement new construction
493  new ( f.m_task ) task_type( arg_functor );
494 
495  // Reference count starts at two
496  // +1 for matching decrement when task is complete
497  // +1 for future
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);
501 
502  assign( f.m_task , arg_options... );
503 
504  // Spawning from within the execution space so the
505  // apply function pointer is guaranteed to be valid
506  f.m_task->m_apply = task_type::apply ;
507 
508  m_queue->schedule( f.m_task );
509  // this task may be updated or executed at any moment
510  }
511 
512  return f ;
513  }
514 
521  template< typename FunctorType , typename ... Options >
522  inline
523  Future< typename FunctorType::value_type , ExecSpace >
524  host_spawn( FunctorType const & arg_functor
525  , Options const & ... arg_options
526  ) const
527  {
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
531  , value_type
532  , FunctorType > ;
533 
534  if ( m_queue == 0 ) {
535  Kokkos::abort("Kokkos::TaskScheduler not initialized");
536  }
537 
538  future_type f ;
539 
540  // Allocate task from memory pool
541  f.m_task =
542  reinterpret_cast<task_type*>( m_queue->allocate(sizeof(task_type)) );
543 
544  if ( f.m_task ) {
545 
546  // Placement new construction
547  new( f.m_task ) task_type( arg_functor );
548 
549  // Reference count starts at two:
550  // +1 to match decrement when task completes
551  // +1 for the future
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);
555 
556  assign( f.m_task , arg_options... );
557 
558  // Potentially spawning outside execution space so the
559  // apply function pointer must be obtained from execution space.
560  // Required for Cuda execution space function pointer.
561  queue_type::specialization::template
562  proc_set_apply< FunctorType >( & f.m_task->m_apply );
563 
564  m_queue->schedule( f.m_task );
565  }
566  return f ;
567  }
568 
572  template< typename A1 , typename A2 >
573  KOKKOS_FUNCTION
574  Future< ExecSpace >
575  when_all( int narg , Future< A1 , A2 > const * const arg ) const
576  {
577  static_assert
578  ( std::is_same< execution_space
579  , typename Future< A1 , A2 >::execution_space
580  >::value
581  , "Future must have same execution space" );
582 
583  using future_type = Future< ExecSpace > ;
584  using task_base = Kokkos::Impl::TaskBase< ExecSpace , void , void > ;
585 
586  future_type f ;
587 
588  size_t const size = sizeof(task_base) + narg * sizeof(task_base*);
589 
590  f.m_task =
591  reinterpret_cast< task_base * >( m_queue->allocate( size ) );
592 
593  if ( f.m_task ) {
594 
595  new( f.m_task ) task_base();
596 
597  // Reference count starts at two:
598  // +1 to match decrement when task completes
599  // +1 for the future
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 ;
605 
606  task_base ** const dep = f.m_task->aggregate_dependences();
607 
608  // Assign dependences to increment their reference count
609  // The futures may be destroyed upon returning from this call
610  // so increment reference count to track this assignment.
611 
612  for ( int i = 0 ; i < narg ; ++i ) {
613  task_base * const t = dep[i] = arg[i].m_task ;
614  if ( 0 != t ) {
615  Kokkos::atomic_increment( &(t->m_ref_count) );
616  }
617  }
618 
619  m_queue->schedule( f.m_task );
620  // this when_all may be processed at any moment
621  }
622 
623  return f ;
624  }
625 
631  template< class FunctorType , typename ... Options >
632  KOKKOS_FUNCTION
633  void respawn( FunctorType * task_self
634  , Options const & ... arg_options ) const
635  {
636  using value_type = typename FunctorType::value_type ;
637  using task_type = Impl::TaskBase< execution_space
638  , value_type
639  , FunctorType > ;
640 
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 );
644 
645  // Precondition:
646  // task is in Executing state
647  // therefore m_next == LockTag
648  //
649  // Change to m_next == 0 for no dependence
650 
651  if ( lock != Kokkos::atomic_exchange( & task->m_next, zero ) ) {
652  Kokkos::abort("TaskScheduler::respawn ERROR: already respawned");
653  }
654 
655  assign( task , arg_options... );
656 
657  // Postcondition:
658  // task is in Executing-Respawn state
659  // therefore m_next == dependece or 0
660  }
661 
662  //----------------------------------------
663 
664  template< typename S >
665  friend
666  void Kokkos::wait( Kokkos::TaskScheduler< S > const & );
667 
668  //----------------------------------------
669 
670  inline
671  int allocation_capacity() const noexcept
672  { return m_queue->m_memory.get_mem_size(); }
673 
674  KOKKOS_INLINE_FUNCTION
675  int allocated_task_count() const noexcept
676  { return m_queue->m_count_alloc ; }
677 
678  KOKKOS_INLINE_FUNCTION
679  int allocated_task_count_max() const noexcept
680  { return m_queue->m_max_alloc ; }
681 
682  KOKKOS_INLINE_FUNCTION
683  long allocated_task_count_accum() const noexcept
684  { return m_queue->m_accum_alloc ; }
685 
686 };
687 
688 template< typename ExecSpace >
689 inline
690 void wait( TaskScheduler< ExecSpace > const & policy )
691 { policy.m_queue->execute(); }
692 
693 // For backward compatibility
694 template< typename ExecSpace >
695 using
696 TaskPolicy = TaskScheduler< ExecSpace > ;
697 
698 } // namespace Kokkos
699 
700 //----------------------------------------------------------------------------
701 //----------------------------------------------------------------------------
702 
703 #endif /* #if defined( KOKKOS_ENABLE_TASKPOLICY ) */
704 #endif /* #ifndef KOKKOS_TASKPOLICY_HPP */
705 
bool is_null(const boost::shared_ptr< T > &p)