Kokkos_Concepts.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_CORE_CONCEPTS_HPP
#define KOKKOS_CORE_CONCEPTS_HPP

#include <type_traits>

// Needed for 'is_space<S>::host_mirror_space
#include <Kokkos_Core_fwd.hpp>

#include <Kokkos_DetectionIdiom.hpp>

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

namespace Kokkos {

// Schedules for Execution Policies
struct Static {};
struct Dynamic {};

// Schedule Wrapper Type
template <class T>
struct Schedule {
  static_assert(std::is_same_v<T, Static> || std::is_same_v<T, Dynamic>,
                "Kokkos: Invalid Schedule<> type.");
  using schedule_type = Schedule;
  using type          = T;
};

// Specify Iteration Index Type
template <typename T>
struct IndexType {
  static_assert(std::is_integral_v<T>, "Kokkos: Invalid IndexType<>.");
  using index_type = IndexType;
  using type       = T;
};

namespace Experimental {
struct WorkItemProperty {
  template <unsigned long Property>
  struct ImplWorkItemProperty {
    static const unsigned value = Property;
    using work_item_property    = ImplWorkItemProperty<Property>;
  };

  constexpr static const ImplWorkItemProperty<0> None =
      ImplWorkItemProperty<0>();
  constexpr static const ImplWorkItemProperty<1> HintLightWeight =
      ImplWorkItemProperty<1>();
  constexpr static const ImplWorkItemProperty<2> HintHeavyWeight =
      ImplWorkItemProperty<2>();
  constexpr static const ImplWorkItemProperty<4> HintRegular =
      ImplWorkItemProperty<4>();
  constexpr static const ImplWorkItemProperty<8> HintIrregular =
      ImplWorkItemProperty<8>();
  constexpr static const ImplWorkItemProperty<16> ImplForceGlobalLaunch =
      ImplWorkItemProperty<16>();
  using None_t                  = ImplWorkItemProperty<0>;
  using HintLightWeight_t       = ImplWorkItemProperty<1>;
  using HintHeavyWeight_t       = ImplWorkItemProperty<2>;
  using HintRegular_t           = ImplWorkItemProperty<4>;
  using HintIrregular_t         = ImplWorkItemProperty<8>;
  using ImplForceGlobalLaunch_t = ImplWorkItemProperty<16>;
};

template <unsigned long pv1, unsigned long pv2>
inline constexpr WorkItemProperty::ImplWorkItemProperty<pv1 | pv2> operator|(
    WorkItemProperty::ImplWorkItemProperty<pv1>,
    WorkItemProperty::ImplWorkItemProperty<pv2>) {
  return WorkItemProperty::ImplWorkItemProperty<pv1 | pv2>();
}

template <unsigned long pv1, unsigned long pv2>
inline constexpr WorkItemProperty::ImplWorkItemProperty<pv1 & pv2> operator&(
    WorkItemProperty::ImplWorkItemProperty<pv1>,
    WorkItemProperty::ImplWorkItemProperty<pv2>) {
  return WorkItemProperty::ImplWorkItemProperty<pv1 & pv2>();
}

template <unsigned long pv1, unsigned long pv2>
inline constexpr bool operator==(WorkItemProperty::ImplWorkItemProperty<pv1>,
                                 WorkItemProperty::ImplWorkItemProperty<pv2>) {
  return pv1 == pv2;
}

}  // namespace Experimental

template <unsigned int maxT = 0 /* Max threads per block */
          ,
          unsigned int minB = 0 /* Min blocks per SM */
          >
struct LaunchBounds {
  using launch_bounds = LaunchBounds;
  using type          = LaunchBounds<maxT, minB>;
  static constexpr unsigned int maxTperB{maxT};
  static constexpr unsigned int minBperSM{minB};
};

}  // namespace Kokkos

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

namespace Kokkos {

#define KOKKOS_IMPL_IS_CONCEPT(CONCEPT)                        \
  template <typename T>                                        \
  struct is_##CONCEPT {                                        \
   private:                                                    \
    template <typename U>                                      \
    using have_t = typename U::CONCEPT;                        \
    template <typename U>                                      \
    using have_type_t = typename U::CONCEPT##_type;            \
                                                               \
   public:                                                     \
    static constexpr bool value =                              \
        std::is_base_of_v<detected_t<have_t, T>, T> ||         \
        std::is_base_of_v<detected_t<have_type_t, T>, T>;      \
    constexpr operator bool() const noexcept { return value; } \
  };                                                           \
  template <typename T>                                        \
  inline constexpr bool is_##CONCEPT##_v = is_##CONCEPT<T>::value;

// Public concept:

KOKKOS_IMPL_IS_CONCEPT(memory_space)
KOKKOS_IMPL_IS_CONCEPT(memory_traits)
KOKKOS_IMPL_IS_CONCEPT(execution_space)
KOKKOS_IMPL_IS_CONCEPT(execution_policy)
KOKKOS_IMPL_IS_CONCEPT(array_layout)
KOKKOS_IMPL_IS_CONCEPT(reducer)
KOKKOS_IMPL_IS_CONCEPT(team_handle)
namespace Experimental {
KOKKOS_IMPL_IS_CONCEPT(work_item_property)
KOKKOS_IMPL_IS_CONCEPT(hooks_policy)
}  // namespace Experimental

namespace Impl {

// Implementation concept:

KOKKOS_IMPL_IS_CONCEPT(thread_team_member)
KOKKOS_IMPL_IS_CONCEPT(host_thread_team_member)
KOKKOS_IMPL_IS_CONCEPT(graph_kernel)

}  // namespace Impl

#undef KOKKOS_IMPL_IS_CONCEPT

}  // namespace Kokkos

namespace Kokkos {
namespace Impl {

template <class Object>
class has_member_team_shmem_size {
  template <typename T>
  static int32_t test_for_member(decltype(&T::team_shmem_size)) {
    return int32_t(0);
  }
  template <typename T>
  static int64_t test_for_member(...) {
    return int64_t(0);
  }

 public:
  constexpr static bool value =
      sizeof(test_for_member<Object>(nullptr)) == sizeof(int32_t);
};

template <class Object>
class has_member_shmem_size {
  template <typename T>
  static int32_t test_for_member(decltype(&T::shmem_size_me)) {
    return int32_t(0);
  }
  template <typename T>
  static int64_t test_for_member(...) {
    return int64_t(0);
  }

 public:
  constexpr static bool value =
      sizeof(test_for_member<Object>(0)) == sizeof(int32_t);
};

}  // namespace Impl
}  // namespace Kokkos
//----------------------------------------------------------------------------

namespace Kokkos {

template <class ExecutionSpace, class MemorySpace>
struct Device {
  static_assert(Kokkos::is_execution_space<ExecutionSpace>::value,
                "Execution space is not valid");
  static_assert(Kokkos::is_memory_space<MemorySpace>::value,
                "Memory space is not valid");
  using execution_space = ExecutionSpace;
  using memory_space    = MemorySpace;
  using device_type     = Device<execution_space, memory_space>;
};

namespace Impl {

template <typename T>
struct is_device_helper : std::false_type {};

template <typename ExecutionSpace, typename MemorySpace>
struct is_device_helper<Device<ExecutionSpace, MemorySpace>> : std::true_type {
};

}  // namespace Impl

template <typename T>
using is_device = typename Impl::is_device_helper<std::remove_cv_t<T>>::type;

template <typename T>
inline constexpr bool is_device_v = is_device<T>::value;

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

template <typename T>
struct is_space {
 private:
  template <typename, typename = void>
  struct exe : std::false_type {
    using space = void;
  };

  template <typename, typename = void>
  struct mem : std::false_type {
    using space = void;
  };

  template <typename, typename = void>
  struct dev : std::false_type {
    using space = void;
  };

  template <typename U>
  struct exe<U, std::conditional_t<true, void, typename U::execution_space>>
      : std::is_same<U, typename U::execution_space>::type {
    using space = typename U::execution_space;
  };

  template <typename U>
  struct mem<U, std::conditional_t<true, void, typename U::memory_space>>
      : std::is_same<U, typename U::memory_space>::type {
    using space = typename U::memory_space;
  };

  template <typename U>
  struct dev<U, std::conditional_t<true, void, typename U::device_type>>
      : std::is_same<U, typename U::device_type>::type {
    using space = typename U::device_type;
  };

  using is_exe = typename is_space<T>::template exe<std::remove_cv_t<T>>;
  using is_mem = typename is_space<T>::template mem<std::remove_cv_t<T>>;
  using is_dev = typename is_space<T>::template dev<std::remove_cv_t<T>>;

 public:
  static constexpr bool value = is_exe::value || is_mem::value || is_dev::value;

  constexpr operator bool() const noexcept { return value; }

  using execution_space = typename is_exe::space;
  using memory_space    = typename is_mem::space;
};

}  // namespace Kokkos

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

namespace Kokkos {
namespace Impl {

template <typename DstMemorySpace, typename SrcMemorySpace>
struct MemorySpaceAccess {
  static_assert(Kokkos::is_memory_space<DstMemorySpace>::value &&
                    Kokkos::is_memory_space<SrcMemorySpace>::value,
                "template arguments must be memory spaces");

  enum { assignable = std::is_same_v<DstMemorySpace, SrcMemorySpace> };

  enum { accessible = assignable };

  enum { deepcopy = assignable };
};

}  // namespace Impl
}  // namespace Kokkos

namespace Kokkos {

template <typename AccessSpace, typename MemorySpace>
struct SpaceAccessibility {
 private:
  static_assert(Kokkos::is_space<AccessSpace>::value,
                "template argument #1 must be a Kokkos space");

  static_assert(Kokkos::is_memory_space<MemorySpace>::value,
                "template argument #2 must be a Kokkos memory space");

  // The input AccessSpace may be a Device<ExecSpace,MemSpace>
  // verify that it is a valid combination of spaces.
  static_assert(Kokkos::Impl::MemorySpaceAccess<
                    typename AccessSpace::execution_space::memory_space,
                    typename AccessSpace::memory_space>::accessible,
                "template argument #1 is an invalid space");

  using exe_access = Kokkos::Impl::MemorySpaceAccess<
      typename AccessSpace::execution_space::memory_space, MemorySpace>;

  using mem_access =
      Kokkos::Impl::MemorySpaceAccess<typename AccessSpace::memory_space,
                                      MemorySpace>;

 public:
  enum { accessible = exe_access::accessible };

  enum {
    assignable = is_memory_space<AccessSpace>::value && mem_access::assignable
  };

  enum { deepcopy = mem_access::deepcopy };

  // What intercessory space for AccessSpace::execution_space
  // to be able to access MemorySpace?
  // If same memory space or not accessible use the AccessSpace
  // else construct a device with execution space and memory space.
  using space = std::conditional_t<
      std::is_same_v<typename AccessSpace::memory_space, MemorySpace> ||
          !exe_access::accessible,
      AccessSpace,
      Kokkos::Device<typename AccessSpace::execution_space, MemorySpace>>;
};

}  // namespace Kokkos

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

#endif  // KOKKOS_CORE_CONCEPTS_HPP