Ginkgo: ginkgo/extensions/kokkos/types.hpp Source File

| Ginkgo Generated from branch based on main. Ginkgo version 1.11.0 A numerical linear algebra library targeting many-core architectures

Loading...
Searching...
No Matches
types.hpp
1// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
2//
3// SPDX-License-Identifier: BSD-3-Clause
4
5#ifndef GINKGO_EXTENSIONS_KOKKOS_TYPES_HPP
6#define GINKGO_EXTENSIONS_KOKKOS_TYPES_HPP
7
8#include <Kokkos_Complex.hpp>
9#include <Kokkos_Core.hpp>
10
11#include <ginkgo/config.hpp>
12#include <ginkgo/core/base/array.hpp>
13#include <ginkgo/core/matrix/dense.hpp>
14#include <ginkgo/extensions/kokkos/config.hpp>
15
16
17namespace gko {
18namespace ext {
19namespace kokkos {
20namespace detail {
21
22
28template <typename T>
29struct value_type_impl {
30 using type = T;
31};
32
33template <typename T>
34struct value_type_impl<const T> {
35 using type = const typename value_type_impl<T>::type;
36};
37
38template <typename T>
39struct value_type_impl<std::complex<T>> {
40 using type = Kokkos::complex<T>;
41};
42
43
44template <typename T>
45struct value_type {
46 using type = typename value_type_impl<T>::type;
47
48 static_assert(sizeof(std::decay_t<T>) == sizeof(std::decay_t<type>),
49 "Can't handle C++ data type and corresponding Kokkos "
50 "type with mismatching type sizes.");
51#if GINKGO_EXTENSION_KOKKOS_CHECK_TYPE_ALIGNMENT
52 static_assert(
53 alignof(std::decay_t<T>) == alignof(std::decay_t<type>),
54 "Can't handle C++ data type and corresponding Kokkos type with "
55 "mismatching alignments. If std::complex is used, please make sure "
56 "to configure Kokkos with `KOKKOS_ENABLE_COMPLEX_ALIGN=ON`.\n"
57 "Alternatively, disable this check by setting the CMake option "
58 "-DGINKGO_EXTENSION_KOKKOS_CHECK_TYPE_ALIGNMENT=OFF.");
59#endif
60};
61
62template <typename T>
63using value_type_t = typename value_type<T>::type;
64
65
77template <typename T, typename MemorySpace>
78struct mapper {
79 static auto map(T&);
80
81 static auto map(const T&);
82};
83
84
93template <typename ValueType, typename MemorySpace>
94struct mapper<array<ValueType>, MemorySpace> {
95 template <typename ValueType_c>
96 using type =
97 Kokkos::View<typename value_type<ValueType_c>::type*, MemorySpace,
98 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
99
109 template <typename ValueType_c>
110 static type<ValueType_c> map(ValueType_c* data, size_type size);
111
115 static type<ValueType> map(array<ValueType>& arr);
116
120 static type<const ValueType> map(const array<ValueType>& arr);
121
125 static type<const ValueType> map(
126 const ::gko::detail::const_array_view<ValueType>& arr);
127};
128
129template <typename ValueType, typename MemorySpace>
130template <typename ValueType_c>
131typename mapper<array<ValueType>, MemorySpace>::template type<ValueType_c>
132mapper<array<ValueType>, MemorySpace>::map(ValueType_c* data, size_type size)
133{
134 return type<ValueType_c>{reinterpret_cast<value_type_t<ValueType_c>*>(data),
135 size};
136}
137
138
139template <typename ValueType, typename MemorySpace>
140typename mapper<array<ValueType>, MemorySpace>::template type<ValueType>
141mapper<array<ValueType>, MemorySpace>::map(array<ValueType>& arr)
142{
143 assert_compatibility<MemorySpace>(arr);
144
145 return map(arr.get_data(), arr.get_size());
146}
147
148
149template <typename ValueType, typename MemorySpace>
150typename mapper<array<ValueType>, MemorySpace>::template type<const ValueType>
151mapper<array<ValueType>, MemorySpace>::map(const array<ValueType>& arr)
152{
153 assert_compatibility<MemorySpace>(arr);
154
155 return map(arr.get_const_data(), arr.get_size());
156}
157
158
159template <typename ValueType, typename MemorySpace>
160typename mapper<array<ValueType>, MemorySpace>::template type<const ValueType>
161mapper<array<ValueType>, MemorySpace>::map(
162 const ::gko::detail::const_array_view<ValueType>& arr)
163{
164 assert_compatibility<MemorySpace>(arr);
165
166 return map(arr.get_const_data(), arr.get_size());
167}
168
169
178template <typename ValueType, typename MemorySpace>
179struct mapper<matrix::Dense<ValueType>, MemorySpace> {
180 template <typename ValueType_c>
181 using type = Kokkos::View<typename value_type<ValueType_c>::type**,
182 Kokkos::LayoutStride, MemorySpace,
183 Kokkos::MemoryTraits<Kokkos::Unmanaged>>;
184
188 static type<ValueType> map(matrix::Dense<ValueType>& m);
189
193 static type<const ValueType> map(const matrix::Dense<ValueType>& m);
194};
195
196template <typename ValueType, typename MemorySpace>
197typename mapper<matrix::Dense<ValueType>, MemorySpace>::template type<ValueType>
198mapper<matrix::Dense<ValueType>, MemorySpace>::map(matrix::Dense<ValueType>& m)
199{
200 assert_compatibility<MemorySpace>(m);
201
202 auto size = m.get_size();
203
204 return type<ValueType>{
205 reinterpret_cast<value_type_t<ValueType>*>(m.get_values()),
206 Kokkos::LayoutStride{size[0], m.get_stride(), size[1], 1}};
207}
208
209
210template <typename ValueType, typename MemorySpace>
211typename mapper<matrix::Dense<ValueType>,
212 MemorySpace>::template type<const ValueType>
213mapper<matrix::Dense<ValueType>, MemorySpace>::map(
214 const matrix::Dense<ValueType>& m)
215{
216 assert_compatibility<MemorySpace>(m);
217
218 auto size = m.get_size();
219
220 return type<const ValueType>{
221 reinterpret_cast<const value_type_t<ValueType>*>(m.get_const_values()),
222 Kokkos::LayoutStride{size[0], m.get_stride(), size[1], 1}};
223}
224
225
226} // namespace detail
227
228
239template <typename MemorySpace = Kokkos::DefaultExecutionSpace::memory_space,
240 typename T>
241inline auto map_data(T&& data)
242{
243 return detail::mapper<std::decay_t<T>, MemorySpace>::map(
244 std::forward<T>(data));
245}
246
250template <typename MemorySpace = Kokkos::DefaultExecutionSpace::memory_space,
251 typename T>
252inline auto map_data(T* data)
253{
254 return map_data<MemorySpace>(*data);
255}
256
260template <typename MemorySpace = Kokkos::DefaultExecutionSpace::memory_space,
261 typename T>
262inline auto map_data(std::unique_ptr<T>& data)
263{
264 return map_data<MemorySpace>(*data);
265}
266
270template <typename MemorySpace = Kokkos::DefaultExecutionSpace::memory_space,
271 typename T>
272inline auto map_data(std::shared_ptr<T>& data)
273{
274 return map_data<MemorySpace>(*data);
275}
276
277
278} // namespace kokkos
279} // namespace ext
280} // namespace gko
281
282
283#endif // GINKGO_EXTENSIONS_KOKKOS_TYPES_HPP
gko
The Ginkgo namespace.
Definition abstract_factory.hpp:20
gko::size_type
std::size_t size_type
Integral type used for allocation quantities.
Definition types.hpp:90
gko::layout_type::array
@ array
The matrix should be written as dense matrix in column-major order.
Definition mtx_io.hpp:97