49 #ifndef TPETRA_DETAILS_MULTI_VECTOR_DIST_OBJECT_KERNELS_DECL_HPP 50 #define TPETRA_DETAILS_MULTI_VECTOR_DIST_OBJECT_KERNELS_DECL_HPP 52 #include "Tpetra_ConfigDefs.hpp" 53 #if TPETRA_USE_KOKKOS_DISTOBJECT 55 #include "Kokkos_Core.hpp" 66 template <
typename Scalar,
typename LocalOrdinal,
typename Device>
67 struct PackArraySingleColumnConstantStride {
68 typedef Device execution_space;
69 typedef typename execution_space::size_type size_type;
71 Kokkos::View<const LocalOrdinal*, execution_space> exportLIDs;
72 Kokkos::View<const Scalar*, execution_space, Kokkos::MemoryUnmanaged> src;
73 Kokkos::View<Scalar*, execution_space> exports;
75 KOKKOS_INLINE_FUNCTION
76 void operator()(
const size_type k )
const {
77 exports[k] = src[exportLIDs[k]];
83 template <
typename Scalar,
typename LocalOrdinal,
typename Device>
84 struct PackArraySingleColumnOffset {
85 typedef Device execution_space;
86 typedef typename execution_space::size_type size_type;
88 Kokkos::View<const LocalOrdinal*, execution_space> exportLIDs;
89 Kokkos::View<const Scalar*, execution_space, Kokkos::MemoryUnmanaged> src;
90 Kokkos::View<Scalar*, execution_space> exports;
93 KOKKOS_INLINE_FUNCTION
94 void operator()(
const size_type k )
const {
95 exports[k] = src[exportLIDs[k] + offset];
101 template <
typename Scalar,
typename LocalOrdinal,
typename Device>
102 struct PackArrayMultiColumnConstantStride {
103 typedef Device execution_space;
104 typedef typename execution_space::size_type size_type;
106 Kokkos::View<const LocalOrdinal*, execution_space> exportLIDs;
107 Kokkos::View<const Scalar*, execution_space, Kokkos::MemoryUnmanaged> src;
108 Kokkos::View<Scalar*, execution_space> exports;
109 size_t stride, numCols;
111 KOKKOS_INLINE_FUNCTION
112 void operator()(
const size_type k )
const {
113 const size_t localRow = exportLIDs[k];
114 const size_t offset = k*numCols;
115 for (
size_t j = 0; j < numCols; ++j)
116 exports[offset + j] = src[localRow + j*stride];
122 template <
typename Scalar,
typename LocalOrdinal,
typename Device>
123 struct PackArrayMultiColumnVariableStride {
124 typedef Device execution_space;
125 typedef typename execution_space::size_type size_type;
127 Kokkos::View<const LocalOrdinal*, execution_space> exportLIDs;
128 Kokkos::View<const size_t*, execution_space> srcWhichVectors;
129 Kokkos::View<const Scalar*, execution_space, Kokkos::MemoryUnmanaged> src;
130 Kokkos::View<Scalar*, execution_space> exports;
131 size_t stride, numCols;
133 KOKKOS_INLINE_FUNCTION
134 void operator()(
const size_type k )
const {
135 const size_t localRow = exportLIDs[k];
136 const size_t offset = k*numCols;
137 for (
size_t j = 0; j < numCols; ++j)
138 exports[offset + j] = src[localRow + srcWhichVectors[j]*stride];
145 template <
typename Scalar>
146 KOKKOS_INLINE_FUNCTION
147 void operator() (Scalar& dest,
const Scalar& src)
const {
149 Kokkos::atomic_exchange(&dest, src);
153 template <
typename Scalar>
154 KOKKOS_INLINE_FUNCTION
155 void operator() (Scalar& dest,
const Scalar& src)
const {
157 Kokkos::atomic_fetch_add(&dest, src);
161 template <
typename Scalar>
162 KOKKOS_INLINE_FUNCTION
163 Scalar max(
const Scalar& a,
const Scalar& b)
const {
164 return a > b ? a : b ;
167 template <
typename Scalar>
168 KOKKOS_INLINE_FUNCTION
169 void operator() (Scalar& dest,
const Scalar& src)
const {
170 typedef Teuchos::ScalarTraits<Scalar> SCT;
172 Kokkos::atomic_exchange(
173 &dest, max( SCT::magnitude(dest), SCT::magnitude(src) ) );
177 template <
typename Scalar,
typename LocalOrdinal,
typename Op,
typename Device>
178 struct UnpackArrayMultiColumnConstantStride {
179 typedef Device execution_space;
180 typedef typename execution_space::size_type size_type;
182 Kokkos::View<const LocalOrdinal*, execution_space> importLIDs;
183 Kokkos::View<const Scalar*, execution_space> imports;
184 Kokkos::View<Scalar*, execution_space, Kokkos::MemoryUnmanaged> dest;
185 size_t stride, numCols;
188 KOKKOS_INLINE_FUNCTION
189 void operator()(
const size_type k )
const {
190 const size_t localRow = importLIDs[k];
191 const size_t offset = k*numCols;
192 for (
size_t j = 0; j < numCols; ++j)
193 op(dest[localRow + j*stride], imports[offset + j]);
199 template <
typename Scalar,
typename LocalOrdinal,
typename Op,
typename Device>
200 struct UnpackArrayMultiColumnVariableStride {
201 typedef Device execution_space;
202 typedef typename execution_space::size_type size_type;
204 Kokkos::View<const LocalOrdinal*, execution_space> importLIDs;
205 Kokkos::View<const size_t*, execution_space> whichVectors;
206 Kokkos::View<const Scalar*, execution_space> imports;
207 Kokkos::View<Scalar*, execution_space, Kokkos::MemoryUnmanaged> dest;
208 size_t stride, numCols;
211 KOKKOS_INLINE_FUNCTION
212 void operator()(
const size_type k )
const {
213 const size_t localRow = importLIDs[k];
214 const size_t offset = k*numCols;
215 for (
size_t j = 0; j < numCols; ++j)
216 op(dest[localRow + whichVectors[j]*stride], imports[offset + j]);
222 template <
typename Scalar,
typename LocalOrdinal,
typename Device>
223 struct PermuteArrayMultiColumnConstantStride {
224 typedef Device execution_space;
225 typedef typename execution_space::size_type size_type;
227 Kokkos::View<const LocalOrdinal*, execution_space> permuteToLIDs;
228 Kokkos::View<const LocalOrdinal*, execution_space> permuteFromLIDs;
229 Kokkos::View<const Scalar*, execution_space, Kokkos::MemoryUnmanaged> src;
230 Kokkos::View<Scalar*, execution_space, Kokkos::MemoryUnmanaged> dest;
231 size_t src_stride, dest_stride, numCols;
233 KOKKOS_INLINE_FUNCTION
234 void operator()(
const size_type k )
const {
235 const size_t toRow = permuteToLIDs[k];
236 const size_t fromRow = permuteFromLIDs[k];
237 for (
size_t j = 0; j < numCols; ++j)
238 dest[toRow + j*dest_stride] = src[fromRow + j*src_stride];
244 template <
typename Scalar,
typename LocalOrdinal,
typename Device>
245 struct PermuteArrayMultiColumnVariableStride {
246 typedef Device execution_space;
247 typedef typename execution_space::size_type size_type;
249 Kokkos::View<const LocalOrdinal*, execution_space> permuteToLIDs;
250 Kokkos::View<const LocalOrdinal*, execution_space> permuteFromLIDs;
251 Kokkos::View<const size_t*, execution_space> src_whichVectors;
252 Kokkos::View<const size_t*, execution_space> dest_whichVectors;
253 Kokkos::View<const Scalar*, execution_space, Kokkos::MemoryUnmanaged> src;
254 Kokkos::View<Scalar*, execution_space, Kokkos::MemoryUnmanaged> dest;
255 size_t src_stride, dest_stride, numCols;
257 KOKKOS_INLINE_FUNCTION
258 void operator()(
const size_type k )
const {
259 const size_t toRow = permuteToLIDs[k];
260 const size_t fromRow = permuteFromLIDs[k];
261 for (
size_t j = 0; j < numCols; ++j)
262 dest[toRow + dest_whichVectors[j]*dest_stride] =
263 src[fromRow + src_whichVectors[j]*src_stride];
272 #endif // TPETRA_USE_KOKKOS_DISTOBJECT 274 #endif // TPETRA_DETAILS_MULTI_VECTOR_DIST_OBJECT_KERNELS_DECL_HPP Namespace Tpetra contains the class and methods constituting the Tpetra library.
Implementation details of Tpetra.