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+
17+ #ifndef KOKKOSKERNELS_UNIFIEDSCALARVIEW_HPP
18+ #define KOKKOSKERNELS_UNIFIEDSCALARVIEW_HPP
19+
20+ #include < type_traits>
21+
22+ #include < Kokkos_Core.hpp>
23+
24+ #include < KokkosKernels_AlwaysFalse.hpp>
25+ #include < KokkosKernels_IsKokkosComplex.hpp>
26+
27+
28+ namespace KokkosKernels {
29+ namespace Impl {
30+
31+
32+ template <typename ScalarLike, typename = void >
33+ struct is_scalar : std::false_type {};
34+
35+ // kokkos complex
36+ template <typename T>
37+ struct is_scalar <T, std::enable_if_t <is_kokkos_complex_v<T>>> : std::true_type {};
38+
39+ // other scalars
40+ template <typename ScalarLike>
41+ struct is_scalar <ScalarLike, std::enable_if_t <std::is_integral_v<ScalarLike> || std::is_floating_point_v<ScalarLike>>> : std::true_type {};
42+
43+ template <typename ScalarLike>
44+ inline constexpr bool is_scalar_v = is_scalar<ScalarLike>::value;
45+
46+
47+
48+
49+ template <typename ScalarLike, typename = void >
50+ struct is_scalar_view : std::false_type {};
51+
52+ // rank 0
53+ template <typename ScalarLike>
54+ struct is_scalar_view <ScalarLike, std::enable_if_t <0 == ScalarLike::rank>> : std::true_type {};
55+
56+ // rank 1 and static extent is 1
57+ template <typename ScalarLike>
58+ struct is_scalar_view <ScalarLike,
59+ std::enable_if_t <
60+ 1 == ScalarLike::rank && 1 == ScalarLike::static_extent(0 )
61+ >
62+ > : std::true_type {};
63+
64+ template <typename ScalarLike>
65+ inline constexpr bool is_scalar_view_v = is_scalar_view<ScalarLike>::value;
66+
67+ /* ! \brief true iff ScalarLike is a scalar or a 0D or 1D view of a single thing
68+ */
69+ template <typename ScalarLike>
70+ inline constexpr bool is_scalar_or_scalar_view = is_scalar_v<ScalarLike> || is_scalar_view_v<ScalarLike>;
71+
72+
73+
74+
75+
76+ template <typename Value, typename = void >
77+ struct unified_scalar ;
78+
79+ template <typename Value>
80+ struct unified_scalar <Value, std::enable_if_t <is_scalar_v<Value>>> {
81+
82+ using type = Value;
83+ using non_const_type = std::remove_const_t <type>;
84+ };
85+
86+ template <typename Value>
87+ struct unified_scalar <Value, std::enable_if_t <is_scalar_view_v<Value>>> {
88+
89+ using type = typename Value::value_type;
90+ using non_const_type = std::remove_const_t <type>;
91+ };
92+
93+ template <typename Value>
94+ using unified_scalar_t = typename unified_scalar<Value>::type;
95+
96+ template <typename Value>
97+ using non_const_unified_scalar_t = typename unified_scalar<Value>::non_const_type;
98+
99+
100+ template <typename Value>
101+ constexpr unified_scalar_t <Value> get_scalar (const Value &v) {
102+
103+ static_assert (is_scalar_or_scalar_view<Value>, " "
104+
105+ unified_scalar_t <Value> ref;
106+ if constexpr (is_scalar_view_v<Value>) {
107+ if (0 == Value::rank) {
108+ ref = *v;
109+ } else if (1 == Value::rank) {
110+ ref = v[0 ];
111+ } else {
112+ static_assert (KokkosKernels::Impl::always_false_v<Value>, " "
113+ }
114+ } else {
115+ ref = v;
116+ }
117+ return ref;
118+ }
119+
120+ } // namespace Impl
121+ } // namespace KokkosKernels
122+
123+ #endif // KOKKOSKERNELS_UNIFIEDSCALARVIEW_HPP
0 commit comments