Kokkos_Parallel.hpp

//@HEADER
// ************************************************************************
//
//                        Kokkos v. 4.0
//       Copyright (2022) National Technology & Engineering
//               Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// the U.S. Government retains certain rights in this software.
//
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
// See https://kokkos.org/LICENSE for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//@HEADER


#ifndef KOKKOS_IMPL_PUBLIC_INCLUDE
#include <Kokkos_Macros.hpp>
static_assert(false,
              "Including non-public Kokkos header files is not allowed.");
#endif
#ifndef KOKKOS_PARALLEL_HPP
#define KOKKOS_PARALLEL_HPP

#include <Kokkos_Core_fwd.hpp>
#include <Kokkos_DetectionIdiom.hpp>
#include <Kokkos_ExecPolicy.hpp>
#include <Kokkos_View.hpp>

#include <impl/Kokkos_Tools.hpp>
#include <impl/Kokkos_Tools_Generic.hpp>

#include <impl/Kokkos_Traits.hpp>
#include <impl/Kokkos_FunctorAnalysis.hpp>

#include <cstddef>
#include <type_traits>
#include <typeinfo>

//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

namespace Kokkos {
namespace Impl {

template <class T>
using execution_space_t = typename T::execution_space;

template <class T>
using device_type_t = typename T::device_type;

//----------------------------------------------------------------------------
template <class Functor, class Policy>
struct FunctorPolicyExecutionSpace {
  using policy_execution_space  = detected_t<execution_space_t, Policy>;
  using functor_execution_space = detected_t<execution_space_t, Functor>;
  using functor_device_type     = detected_t<device_type_t, Functor>;
  using functor_device_type_execution_space =
      detected_t<execution_space_t, functor_device_type>;

  static_assert(
      !is_detected<execution_space_t, Policy>::value ||
          !is_detected<execution_space_t, Functor>::value ||
          std::is_same<policy_execution_space, functor_execution_space>::value,
      "A policy with an execution space and a functor with an execution space "
      "are given but the execution space types do not match!");
  static_assert(!is_detected<execution_space_t, Policy>::value ||
                    !is_detected<device_type_t, Functor>::value ||
                    std::is_same<policy_execution_space,
                                 functor_device_type_execution_space>::value,
                "A policy with an execution space and a functor with a device "
                "type are given but the execution space types do not match!");
  static_assert(!is_detected<device_type_t, Functor>::value ||
                    !is_detected<execution_space_t, Functor>::value ||
                    std::is_same<functor_device_type_execution_space,
                                 functor_execution_space>::value,
                "A functor with both an execution space and device type is "
                "given but their execution space types do not match!");

  using execution_space = detected_or_t<
      detected_or_t<
          std::conditional_t<
              is_detected<device_type_t, Functor>::value,
              detected_t<execution_space_t, detected_t<device_type_t, Functor>>,
              Kokkos::DefaultExecutionSpace>,
          execution_space_t, Functor>,
      execution_space_t, Policy>;
};

}  // namespace Impl
}  // namespace Kokkos

//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

namespace Kokkos {

template <
    class ExecPolicy, class FunctorType,
    class Enable = std::enable_if_t<is_execution_policy<ExecPolicy>::value>>
inline void parallel_for(const std::string& str, const ExecPolicy& policy,
                         const FunctorType& functor) {
  uint64_t kpID = 0;

  ExecPolicy inner_policy = policy;
  Kokkos::Tools::Impl::begin_parallel_for(inner_policy, functor, str, kpID);

  Kokkos::Impl::shared_allocation_tracking_disable();
  Impl::ParallelFor<FunctorType, ExecPolicy> closure(functor, inner_policy);
  Kokkos::Impl::shared_allocation_tracking_enable();

  closure.execute();

  Kokkos::Tools::Impl::end_parallel_for(inner_policy, functor, str, kpID);
}

template <class ExecPolicy, class FunctorType>
inline void parallel_for(
    const ExecPolicy& policy, const FunctorType& functor,
    std::enable_if_t<is_execution_policy<ExecPolicy>::value>* = nullptr) {
  Kokkos::parallel_for("", policy, functor);
}

template <class FunctorType>
inline void parallel_for(const std::string& str, const size_t work_count,
                         const FunctorType& functor) {
  using execution_space =
      typename Impl::FunctorPolicyExecutionSpace<FunctorType,
                                                 void>::execution_space;
  using policy = RangePolicy<execution_space>;

  policy execution_policy = policy(0, work_count);
  ::Kokkos::parallel_for(str, execution_policy, functor);
}

template <class FunctorType>
inline void parallel_for(const size_t work_count, const FunctorType& functor) {
  ::Kokkos::parallel_for("", work_count, functor);
}

}  // namespace Kokkos

#include <Kokkos_Parallel_Reduce.hpp>
//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

namespace Kokkos {

//               const value_type& input) const
// i/     }
template <class ExecutionPolicy, class FunctorType,
          class Enable =
              std::enable_if_t<is_execution_policy<ExecutionPolicy>::value>>
inline void parallel_scan(const std::string& str, const ExecutionPolicy& policy,
                          const FunctorType& functor) {
  uint64_t kpID                = 0;
  ExecutionPolicy inner_policy = policy;
  Kokkos::Tools::Impl::begin_parallel_scan(inner_policy, functor, str, kpID);

  Kokkos::Impl::shared_allocation_tracking_disable();
  Impl::ParallelScan<FunctorType, ExecutionPolicy> closure(functor,
                                                           inner_policy);
  Kokkos::Impl::shared_allocation_tracking_enable();

  closure.execute();

  Kokkos::Tools::Impl::end_parallel_scan(inner_policy, functor, str, kpID);
}

template <class ExecutionPolicy, class FunctorType>
inline void parallel_scan(
    const ExecutionPolicy& policy, const FunctorType& functor,
    std::enable_if_t<is_execution_policy<ExecutionPolicy>::value>* = nullptr) {
  ::Kokkos::parallel_scan("", policy, functor);
}

template <class FunctorType>
inline void parallel_scan(const std::string& str, const size_t work_count,
                          const FunctorType& functor) {
  using execution_space =
      typename Kokkos::Impl::FunctorPolicyExecutionSpace<FunctorType,
                                                         void>::execution_space;

  using policy = Kokkos::RangePolicy<execution_space>;

  policy execution_policy(0, work_count);
  parallel_scan(str, execution_policy, functor);
}

template <class FunctorType>
inline void parallel_scan(const size_t work_count, const FunctorType& functor) {
  ::Kokkos::parallel_scan("", work_count, functor);
}

template <class ExecutionPolicy, class FunctorType, class ReturnType,
          class Enable =
              std::enable_if_t<is_execution_policy<ExecutionPolicy>::value>>
inline void parallel_scan(const std::string& str, const ExecutionPolicy& policy,
                          const FunctorType& functor,
                          ReturnType& return_value) {
  uint64_t kpID                = 0;
  ExecutionPolicy inner_policy = policy;
  Kokkos::Tools::Impl::begin_parallel_scan(inner_policy, functor, str, kpID);

  if constexpr (Kokkos::is_view<ReturnType>::value) {
    Kokkos::Impl::shared_allocation_tracking_disable();
    Impl::ParallelScanWithTotal<FunctorType, ExecutionPolicy,
                                typename ReturnType::value_type>
        closure(functor, inner_policy, return_value);
    Kokkos::Impl::shared_allocation_tracking_enable();
    closure.execute();
  } else {
    Kokkos::Impl::shared_allocation_tracking_disable();
    Kokkos::View<ReturnType, Kokkos::HostSpace> view(&return_value);
    Impl::ParallelScanWithTotal<FunctorType, ExecutionPolicy, ReturnType>
        closure(functor, inner_policy, view);
    Kokkos::Impl::shared_allocation_tracking_enable();
    closure.execute();
  }

  Kokkos::Tools::Impl::end_parallel_scan(inner_policy, functor, str, kpID);

  if (!Kokkos::is_view<ReturnType>::value)
    policy.space().fence(
        "Kokkos::parallel_scan: fence due to result being a value, not a view");
}

template <class ExecutionPolicy, class FunctorType, class ReturnType>
inline void parallel_scan(
    const ExecutionPolicy& policy, const FunctorType& functor,
    ReturnType& return_value,
    std::enable_if_t<is_execution_policy<ExecutionPolicy>::value>* = nullptr) {
  ::Kokkos::parallel_scan("", policy, functor, return_value);
}

template <class FunctorType, class ReturnType>
inline void parallel_scan(const std::string& str, const size_t work_count,
                          const FunctorType& functor,
                          ReturnType& return_value) {
  using execution_space =
      typename Kokkos::Impl::FunctorPolicyExecutionSpace<FunctorType,
                                                         void>::execution_space;

  using policy = Kokkos::RangePolicy<execution_space>;

  policy execution_policy(0, work_count);
  parallel_scan(str, execution_policy, functor, return_value);
}

template <class FunctorType, class ReturnType>
inline void parallel_scan(const size_t work_count, const FunctorType& functor,
                          ReturnType& return_value) {
  ::Kokkos::parallel_scan("", work_count, functor, return_value);
}

}  // namespace Kokkos

//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

namespace Kokkos {
namespace Impl {

template <class FunctorType,
          bool HasTeamShmemSize =
              has_member_team_shmem_size<FunctorType>::value,
          bool HasShmemSize = has_member_shmem_size<FunctorType>::value>
struct FunctorTeamShmemSize {
  KOKKOS_INLINE_FUNCTION static size_t value(const FunctorType&, int) {
    return 0;
  }
};

template <class FunctorType>
struct FunctorTeamShmemSize<FunctorType, true, false> {
  static inline size_t value(const FunctorType& f, int team_size) {
    return f.team_shmem_size(team_size);
  }
};

template <class FunctorType>
struct FunctorTeamShmemSize<FunctorType, false, true> {
  static inline size_t value(const FunctorType& f, int team_size) {
    return f.shmem_size(team_size);
  }
};
template <class FunctorType>
struct FunctorTeamShmemSize<FunctorType, true, true> {
  static inline size_t value(const FunctorType& /*f*/, int /*team_size*/) {
    Kokkos::abort(
        "Functor with both team_shmem_size and shmem_size defined is "
        "not allowed");
    return 0;
  }
};

}  // namespace Impl
}  // namespace Kokkos

//----------------------------------------------------------------------------
//----------------------------------------------------------------------------

#endif /* KOKKOS_PARALLEL_HPP */