Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
Kokkos_Parallel.hpp
Go to the documentation of this file.
1//@HEADER
2// ************************************************************************
3//
4// Kokkos v. 4.0
5// Copyright (2022) National Technology & Engineering
6// Solutions of Sandia, LLC (NTESS).
7//
8// Under the terms of Contract DE-NA0003525 with NTESS,
9// the U.S. Government retains certain rights in this software.
10//
11// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12// See https://kokkos.org/LICENSE for license information.
13// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14//
15//@HEADER
16
19
20#ifndef KOKKOS_IMPL_PUBLIC_INCLUDE
21#include <Kokkos_Macros.hpp>
22static_assert(false,
23 "Including non-public Kokkos header files is not allowed.");
24#endif
25#ifndef KOKKOS_PARALLEL_HPP
26#define KOKKOS_PARALLEL_HPP
27
28#include <Kokkos_Core_fwd.hpp>
29#include <Kokkos_DetectionIdiom.hpp>
30#include <Kokkos_ExecPolicy.hpp>
31#include <Kokkos_View.hpp>
32
33#include <impl/Kokkos_Tools.hpp>
34#include <impl/Kokkos_Tools_Generic.hpp>
35
36#include <impl/Kokkos_Traits.hpp>
37#include <impl/Kokkos_FunctorAnalysis.hpp>
38
39#include <cstddef>
40#include <type_traits>
41#include <typeinfo>
42
43//----------------------------------------------------------------------------
44//----------------------------------------------------------------------------
45
46namespace Kokkos {
47namespace Impl {
48
49template <class T>
50using execution_space_t = typename T::execution_space;
51
52template <class T>
53using device_type_t = typename T::device_type;
54
55//----------------------------------------------------------------------------
64template <class Functor, class Policy>
71
72 static_assert(
73 !is_detected<execution_space_t, Policy>::value ||
74 !is_detected<execution_space_t, Functor>::value ||
75 std::is_same<policy_execution_space, functor_execution_space>::value,
76 "A policy with an execution space and a functor with an execution space "
77 "are given but the execution space types do not match!");
78 static_assert(!is_detected<execution_space_t, Policy>::value ||
79 !is_detected<device_type_t, Functor>::value ||
80 std::is_same<policy_execution_space,
82 "A policy with an execution space and a functor with a device "
83 "type are given but the execution space types do not match!");
84 static_assert(!is_detected<device_type_t, Functor>::value ||
85 !is_detected<execution_space_t, Functor>::value ||
88 "A functor with both an execution space and device type is "
89 "given but their execution space types do not match!");
90
91 using execution_space = detected_or_t<
92 detected_or_t<
93 std::conditional_t<
94 is_detected<device_type_t, Functor>::value,
96 Kokkos::DefaultExecutionSpace>,
97 execution_space_t, Functor>,
98 execution_space_t, Policy>;
99};
100
101} // namespace Impl
102} // namespace Kokkos
103
104//----------------------------------------------------------------------------
105//----------------------------------------------------------------------------
106
107namespace Kokkos {
108
130template <
131 class ExecPolicy, class FunctorType,
132 class Enable = std::enable_if_t<is_execution_policy<ExecPolicy>::value>>
133inline void parallel_for(const std::string& str, const ExecPolicy& policy,
134 const FunctorType& functor) {
135 uint64_t kpID = 0;
136
137 ExecPolicy inner_policy = policy;
138 Kokkos::Tools::Impl::begin_parallel_for(inner_policy, functor, str, kpID);
139
140 Kokkos::Impl::shared_allocation_tracking_disable();
142 Kokkos::Impl::shared_allocation_tracking_enable();
143
144 closure.execute();
145
146 Kokkos::Tools::Impl::end_parallel_for(inner_policy, functor, str, kpID);
147}
148
149template <class ExecPolicy, class FunctorType>
150inline void parallel_for(
151 const ExecPolicy& policy, const FunctorType& functor,
152 std::enable_if_t<is_execution_policy<ExecPolicy>::value>* = nullptr) {
153 Kokkos::parallel_for("", policy, functor);
154}
155
156template <class FunctorType>
157inline void parallel_for(const std::string& str, const size_t work_count,
158 const FunctorType& functor) {
159 using execution_space =
160 typename Impl::FunctorPolicyExecutionSpace<FunctorType,
161 void>::execution_space;
162 using policy = RangePolicy<execution_space>;
163
164 policy execution_policy = policy(0, work_count);
165 ::Kokkos::parallel_for(str, execution_policy, functor);
166}
167
168template <class FunctorType>
169inline void parallel_for(const size_t work_count, const FunctorType& functor) {
170 ::Kokkos::parallel_for("", work_count, functor);
171}
172
173} // namespace Kokkos
174
175#include <Kokkos_Parallel_Reduce.hpp>
176//----------------------------------------------------------------------------
177//----------------------------------------------------------------------------
178
179namespace Kokkos {
180
219// const value_type& input) const
329// i/ }
346template <class ExecutionPolicy, class FunctorType,
347 class Enable =
348 std::enable_if_t<is_execution_policy<ExecutionPolicy>::value>>
349inline void parallel_scan(const std::string& str, const ExecutionPolicy& policy,
350 const FunctorType& functor) {
351 uint64_t kpID = 0;
352 ExecutionPolicy inner_policy = policy;
353 Kokkos::Tools::Impl::begin_parallel_scan(inner_policy, functor, str, kpID);
354
355 Kokkos::Impl::shared_allocation_tracking_disable();
356 Impl::ParallelScan<FunctorType, ExecutionPolicy> closure(functor,
357 inner_policy);
358 Kokkos::Impl::shared_allocation_tracking_enable();
359
360 closure.execute();
361
362 Kokkos::Tools::Impl::end_parallel_scan(inner_policy, functor, str, kpID);
363}
364
365template <class ExecutionPolicy, class FunctorType>
366inline void parallel_scan(
367 const ExecutionPolicy& policy, const FunctorType& functor,
368 std::enable_if_t<is_execution_policy<ExecutionPolicy>::value>* = nullptr) {
369 ::Kokkos::parallel_scan("", policy, functor);
370}
371
372template <class FunctorType>
373inline void parallel_scan(const std::string& str, const size_t work_count,
374 const FunctorType& functor) {
375 using execution_space =
377 void>::execution_space;
378
380
381 policy execution_policy(0, work_count);
382 parallel_scan(str, execution_policy, functor);
383}
384
385template <class FunctorType>
386inline void parallel_scan(const size_t work_count, const FunctorType& functor) {
387 ::Kokkos::parallel_scan("", work_count, functor);
388}
389
390template <class ExecutionPolicy, class FunctorType, class ReturnType,
391 class Enable =
392 std::enable_if_t<is_execution_policy<ExecutionPolicy>::value>>
393inline void parallel_scan(const std::string& str, const ExecutionPolicy& policy,
394 const FunctorType& functor,
395 ReturnType& return_value) {
396 uint64_t kpID = 0;
397 ExecutionPolicy inner_policy = policy;
398 Kokkos::Tools::Impl::begin_parallel_scan(inner_policy, functor, str, kpID);
399
400 if constexpr (Kokkos::is_view<ReturnType>::value) {
401 Kokkos::Impl::shared_allocation_tracking_disable();
402 Impl::ParallelScanWithTotal<FunctorType, ExecutionPolicy, ReturnType>
403 closure(functor, inner_policy, return_value);
404 Kokkos::Impl::shared_allocation_tracking_enable();
405 closure.execute();
406 } else {
407 Kokkos::Impl::shared_allocation_tracking_disable();
409 Impl::ParallelScanWithTotal<FunctorType, ExecutionPolicy, ReturnType>
410 closure(functor, inner_policy, view);
411 Kokkos::Impl::shared_allocation_tracking_enable();
412 closure.execute();
413 }
414
415 Kokkos::Tools::Impl::end_parallel_scan(inner_policy, functor, str, kpID);
416
417 if (!Kokkos::is_view<ReturnType>::value)
418 policy.space().fence(
419 "Kokkos::parallel_scan: fence due to result being a value, not a view");
420}
421
422template <class ExecutionPolicy, class FunctorType, class ReturnType>
423inline void parallel_scan(
424 const ExecutionPolicy& policy, const FunctorType& functor,
425 ReturnType& return_value,
426 std::enable_if_t<is_execution_policy<ExecutionPolicy>::value>* = nullptr) {
427 ::Kokkos::parallel_scan("", policy, functor, return_value);
428}
429
430template <class FunctorType, class ReturnType>
431inline void parallel_scan(const std::string& str, const size_t work_count,
432 const FunctorType& functor,
433 ReturnType& return_value) {
434 using execution_space =
436 void>::execution_space;
437
439
440 policy execution_policy(0, work_count);
441 parallel_scan(str, execution_policy, functor, return_value);
442}
443
444template <class FunctorType, class ReturnType>
445inline void parallel_scan(const size_t work_count, const FunctorType& functor,
446 ReturnType& return_value) {
447 ::Kokkos::parallel_scan("", work_count, functor, return_value);
448}
449
450} // namespace Kokkos
451
452//----------------------------------------------------------------------------
453//----------------------------------------------------------------------------
454
455namespace Kokkos {
456namespace Impl {
457
458template <class FunctorType,
459 bool HasTeamShmemSize =
460 has_member_team_shmem_size<FunctorType>::value,
461 bool HasShmemSize = has_member_shmem_size<FunctorType>::value>
462struct FunctorTeamShmemSize {
463 KOKKOS_INLINE_FUNCTION static size_t value(const FunctorType&, int) {
464 return 0;
465 }
466};
467
468template <class FunctorType>
469struct FunctorTeamShmemSize<FunctorType, true, false> {
470 static inline size_t value(const FunctorType& f, int team_size) {
471 return f.team_shmem_size(team_size);
472 }
473};
474
475template <class FunctorType>
476struct FunctorTeamShmemSize<FunctorType, false, true> {
477 static inline size_t value(const FunctorType& f, int team_size) {
478 return f.shmem_size(team_size);
479 }
480};
481template <class FunctorType>
482struct FunctorTeamShmemSize<FunctorType, true, true> {
483 static inline size_t value(const FunctorType& /*f*/, int /*team_size*/) {
485 "Functor with both team_shmem_size and shmem_size defined is "
486 "not allowed");
487 return 0;
488 }
489};
490
491} // namespace Impl
492} // namespace Kokkos
493
494//----------------------------------------------------------------------------
495//----------------------------------------------------------------------------
496
497#endif /* KOKKOS_PARALLEL_HPP */
A thread safe view to a bitset.
Implementation of the ParallelFor operator that has a partial specialization for the device.
ReturnType
Given a Functor and Execution Policy query an execution space.