49 #ifndef TPETRA_KOKKOS_REFACTOR_DETAILS_MULTI_VECTOR_DIST_OBJECT_KERNELS_HPP 50 #define TPETRA_KOKKOS_REFACTOR_DETAILS_MULTI_VECTOR_DIST_OBJECT_KERNELS_HPP 52 #include "Kokkos_Core.hpp" 53 #include "Kokkos_ArithTraits.hpp" 56 namespace KokkosRefactor {
62 template <
typename DstView,
typename SrcView,
typename IdxView>
63 struct PackArraySingleColumn {
64 typedef typename DstView::execution_space execution_space;
65 typedef typename execution_space::size_type size_type;
72 PackArraySingleColumn(
const DstView& dst_,
76 dst(dst_), src(src_), idx(idx_), col(col_) {}
78 KOKKOS_INLINE_FUNCTION
79 void operator()(
const size_type k )
const {
80 dst(k) = src(idx(k), col);
83 static void pack(
const DstView& dst,
87 Kokkos::parallel_for( idx.size(),
88 PackArraySingleColumn(dst,src,idx,col) );
94 template <
typename DstView,
typename SrcView,
typename IdxView>
95 void pack_array_single_column(
const DstView& dst,
99 PackArraySingleColumn<DstView,SrcView,IdxView>::pack(
103 template <
typename DstView,
typename SrcView,
typename IdxView>
104 struct PackArrayMultiColumn {
105 typedef typename DstView::execution_space execution_space;
106 typedef typename execution_space::size_type size_type;
113 PackArrayMultiColumn(
const DstView& dst_,
117 dst(dst_), src(src_), idx(idx_), numCols(numCols_) {}
119 KOKKOS_INLINE_FUNCTION
120 void operator()(
const size_type k )
const {
121 const typename IdxView::value_type localRow = idx(k);
122 const size_t offset = k*numCols;
123 for (
size_t j = 0; j < numCols; ++j)
124 dst(offset + j) = src(localRow, j);
127 static void pack(
const DstView& dst,
131 Kokkos::parallel_for( idx.size(),
132 PackArrayMultiColumn(dst,src,idx,numCols) );
138 template <
typename DstView,
typename SrcView,
typename IdxView>
139 void pack_array_multi_column(
const DstView& dst,
143 PackArrayMultiColumn<DstView,SrcView,IdxView>::pack(
144 dst, src, idx, numCols);
147 template <
typename DstView,
typename SrcView,
typename IdxView,
149 struct PackArrayMultiColumnVariableStride {
150 typedef typename DstView::execution_space execution_space;
151 typedef typename execution_space::size_type size_type;
159 PackArrayMultiColumnVariableStride(
const DstView& dst_,
164 dst(dst_), src(src_), idx(idx_), col(col_), numCols(numCols_) {}
166 KOKKOS_INLINE_FUNCTION
167 void operator()(
const size_type k )
const {
168 const typename IdxView::value_type localRow = idx(k);
169 const size_t offset = k*numCols;
170 for (
size_t j = 0; j < numCols; ++j)
171 dst(offset + j) = src(localRow, col(j));
174 static void pack(
const DstView& dst,
179 Kokkos::parallel_for( idx.size(),
180 PackArrayMultiColumnVariableStride(
181 dst,src,idx,col,numCols) );
187 template <
typename DstView,
typename SrcView,
typename IdxView,
189 void pack_array_multi_column_variable_stride(
const DstView& dst,
194 PackArrayMultiColumnVariableStride<DstView,SrcView,IdxView,ColView>::pack(
195 dst, src, idx, col, numCols);
199 template <
typename Scalar>
200 KOKKOS_INLINE_FUNCTION
201 void operator() (Scalar& dest,
const Scalar& src)
const {
202 Kokkos::atomic_assign(&dest, src);
206 template <
typename Scalar>
207 KOKKOS_INLINE_FUNCTION
208 void operator() (Scalar& dest,
const Scalar& src)
const {
209 Kokkos::atomic_add(&dest, src);
215 template <
typename T>
216 KOKKOS_INLINE_FUNCTION
217 T max(
const T& a,
const T& b)
const {
return a > b ? a : b; }
219 template <
typename Scalar>
220 KOKKOS_INLINE_FUNCTION
221 void operator() (Scalar& dest,
const Scalar& src)
const {
222 typedef Kokkos::Details::ArithTraits<Scalar> SCT;
223 Kokkos::atomic_assign(&dest, Scalar(max(SCT::abs(dest),SCT::abs(src))));
227 template <
typename DstView,
typename SrcView,
typename IdxView,
typename Op>
228 struct UnpackArrayMultiColumn {
229 typedef typename DstView::execution_space execution_space;
230 typedef typename execution_space::size_type size_type;
238 UnpackArrayMultiColumn(
const DstView& dst_,
243 dst(dst_), src(src_), idx(idx_), op(op_), numCols(numCols_) {}
245 KOKKOS_INLINE_FUNCTION
246 void operator()(
const size_type k )
const {
247 const typename IdxView::value_type localRow = idx(k);
248 const size_t offset = k*numCols;
249 for (
size_t j = 0; j < numCols; ++j)
250 op( dst(localRow,j), src(offset+j) );
253 static void unpack(
const DstView& dst,
258 Kokkos::parallel_for( idx.size(),
259 UnpackArrayMultiColumn(dst,src,idx,op,numCols) );
265 template <
typename DstView,
typename SrcView,
typename IdxView,
typename Op>
266 void unpack_array_multi_column(
const DstView& dst,
271 UnpackArrayMultiColumn<DstView,SrcView,IdxView,Op>::unpack(
272 dst, src, idx, op, numCols);
275 template <
typename DstView,
typename SrcView,
typename IdxView,
276 typename ColView,
typename Op>
277 struct UnpackArrayMultiColumnVariableStride {
278 typedef typename DstView::execution_space execution_space;
279 typedef typename execution_space::size_type size_type;
288 UnpackArrayMultiColumnVariableStride(
const DstView& dst_,
294 dst(dst_), src(src_), idx(idx_), col(col_), op(op_), numCols(numCols_) {}
296 KOKKOS_INLINE_FUNCTION
297 void operator()(
const size_type k )
const {
298 const typename IdxView::value_type localRow = idx(k);
299 const size_t offset = k*numCols;
300 for (
size_t j = 0; j < numCols; ++j)
301 op( dst(localRow,col(j)), src(offset+j) );
304 static void unpack(
const DstView& dst,
310 Kokkos::parallel_for( idx.size(),
311 UnpackArrayMultiColumnVariableStride(
312 dst,src,idx,col,op,numCols) );
318 template <
typename DstView,
typename SrcView,
typename IdxView,
319 typename ColView,
typename Op>
320 void unpack_array_multi_column_variable_stride(
const DstView& dst,
326 UnpackArrayMultiColumnVariableStride<DstView,SrcView,IdxView,ColView,Op>::unpack(
327 dst, src, idx, col, op, numCols);
330 template <
typename DstView,
typename SrcView,
331 typename DstIdxView,
typename SrcIdxView>
332 struct PermuteArrayMultiColumn {
333 typedef typename DstView::execution_space execution_space;
334 typedef typename execution_space::size_type size_type;
342 PermuteArrayMultiColumn(
const DstView& dst_,
344 const DstIdxView& dst_idx_,
345 const SrcIdxView& src_idx_,
347 dst(dst_), src(src_), dst_idx(dst_idx_), src_idx(src_idx_),
350 KOKKOS_INLINE_FUNCTION
351 void operator()(
const size_type k )
const {
352 const typename DstIdxView::value_type toRow = dst_idx(k);
353 const typename SrcIdxView::value_type fromRow = src_idx(k);
354 for (
size_t j = 0; j < numCols; ++j)
355 dst(toRow, j) = src(fromRow, j);
358 static void permute(
const DstView& dst,
360 const DstIdxView& dst_idx,
361 const SrcIdxView& src_idx,
363 const size_type n = std::min( dst_idx.size(), src_idx.size() );
364 Kokkos::parallel_for(
365 n, PermuteArrayMultiColumn(dst,src,dst_idx,src_idx,numCols) );
371 template <
typename DstView,
typename SrcView,
372 typename DstIdxView,
typename SrcIdxView>
373 void permute_array_multi_column(
const DstView& dst,
375 const DstIdxView& dst_idx,
376 const SrcIdxView& src_idx,
378 PermuteArrayMultiColumn<DstView,SrcView,DstIdxView,SrcIdxView>::permute(
379 dst, src, dst_idx, src_idx, numCols);
382 template <
typename DstView,
typename SrcView,
383 typename DstIdxView,
typename SrcIdxView,
384 typename DstColView,
typename SrcColView>
385 struct PermuteArrayMultiColumnVariableStride {
386 typedef typename DstView::execution_space execution_space;
387 typedef typename execution_space::size_type size_type;
397 PermuteArrayMultiColumnVariableStride(
const DstView& dst_,
399 const DstIdxView& dst_idx_,
400 const SrcIdxView& src_idx_,
401 const DstColView& dst_col_,
402 const SrcColView& src_col_,
404 dst(dst_), src(src_), dst_idx(dst_idx_), src_idx(src_idx_),
405 dst_col(dst_col_), src_col(src_col_),
408 KOKKOS_INLINE_FUNCTION
409 void operator()(
const size_type k )
const {
410 const typename DstIdxView::value_type toRow = dst_idx(k);
411 const typename SrcIdxView::value_type fromRow = src_idx(k);
412 for (
size_t j = 0; j < numCols; ++j)
413 dst(toRow, dst_col(j)) = src(fromRow, src_col(j));
416 static void permute(
const DstView& dst,
418 const DstIdxView& dst_idx,
419 const SrcIdxView& src_idx,
420 const DstColView& dst_col,
421 const SrcColView& src_col,
423 const size_type n = std::min( dst_idx.size(), src_idx.size() );
424 Kokkos::parallel_for(
425 n, PermuteArrayMultiColumnVariableStride(
426 dst,src,dst_idx,src_idx,dst_col,src_col,numCols) );
432 template <
typename DstView,
typename SrcView,
433 typename DstIdxView,
typename SrcIdxView,
434 typename DstColView,
typename SrcColView>
435 void permute_array_multi_column_variable_stride(
const DstView& dst,
437 const DstIdxView& dst_idx,
438 const SrcIdxView& src_idx,
439 const DstColView& dst_col,
440 const SrcColView& src_col,
442 PermuteArrayMultiColumnVariableStride<DstView,SrcView,
443 DstIdxView,SrcIdxView,DstColView,SrcColView>::permute(
444 dst, src, dst_idx, src_idx, dst_col, src_col, numCols);
451 #endif // TPETRA_KOKKOS_REFACTOR_DETAILS_MULTI_VECTOR_DIST_OBJECT_KERNELS_HPP Namespace Tpetra contains the class and methods constituting the Tpetra library.
Implementation details of Tpetra.