Ifpack2 Templated Preconditioning Package Version 1.0
Loading...
Searching...
No Matches
Ifpack2_BlockComputeResidualVector_def.hpp
1// @HEADER
2// *****************************************************************************
3// Ifpack2: Templated Object-Oriented Algebraic Preconditioner Package
4//
5// Copyright 2009 NTESS and the Ifpack2 contributors.
6// SPDX-License-Identifier: BSD-3-Clause
7// *****************************************************************************
8// @HEADER
9
10#ifndef IFPACK2_BLOCKCOMPUTERES_VECTOR_DEF_HPP
11#define IFPACK2_BLOCKCOMPUTERES_VECTOR_DEF_HPP
12
13#include "Ifpack2_BlockComputeResidualVector_decl.hpp"
14
15namespace Ifpack2::BlockHelperDetails {
16
17template <typename impl_scalar_type, int B>
18struct ArrayHelper {
19 ArrayHelper(int /* blocksize_requested */) {}
20 impl_scalar_type data[B];
21 impl_scalar_type *get() { return data; }
22};
23
24template <typename impl_scalar_type>
25struct ArrayHelper<impl_scalar_type, 0> {
26 ArrayHelper(int blocksize_requested) {
27 data = (impl_scalar_type *)malloc(sizeof(impl_scalar_type) * blocksize_requested);
28 }
29 ~ArrayHelper() { free(data); }
30 impl_scalar_type *data;
31 impl_scalar_type *get() { return data; }
32};
33
34// Precompute offsets of each A and x entry to speed up residual.
35// (Applies for hasBlockCrsMatrix == true and OverlapTag/AsyncTag)
36// Reading A, x take up to 4, 6 levels of indirection respectively,
37// but precomputing the offsets reduces it to 2 for both.
38//
39// This function allocates and populates these members of AmD:
40// A_x_offsets, A_x_offsets_remote
41template <typename MatrixType>
42void ComputeResidualVector<MatrixType, BlockTriDiContainerDetails::ImplSimdTag>::precompute_A_x_offsets(
43 AmD<MatrixType> &amd,
44 const PartInterface<MatrixType> &interf,
45 const Teuchos::RCP<const typename impl_type::tpetra_crs_graph_type> &g,
46 const local_ordinal_type_1d_view &dm2cm,
47 int blocksize,
48 bool ownedRemoteSeparate) {
49 using impl_type = ImplType<MatrixType>;
50 using i64_3d_view = typename impl_type::i64_3d_view;
51 using size_type = typename impl_type::size_type;
52 using local_ordinal_type = typename impl_type::local_ordinal_type;
53 using execution_space = typename impl_type::execution_space;
54 auto local_graph = g->getLocalGraphDevice();
55 const auto A_block_rowptr = local_graph.row_map;
56 const auto A_colind = local_graph.entries;
57 local_ordinal_type numLocalRows = interf.rowidx2part.extent(0);
58 int blocksize_square = blocksize * blocksize;
59 // shallow-copying views to avoid capturing the amd, interf objects in lambdas
60 auto lclrow = interf.lclrow;
61 auto A_colindsub = amd.A_colindsub;
62 auto A_colindsub_remote = amd.A_colindsub_remote;
63 auto rowptr = amd.rowptr;
64 auto rowptr_remote = amd.rowptr_remote;
65 bool is_dm2cm_active = dm2cm.extent(0);
66 if (ownedRemoteSeparate) {
67 // amd.rowptr points to owned entries only, and amd.rowptr_remote points to nonowned.
68 local_ordinal_type maxOwnedEntriesPerRow = 0;
69 local_ordinal_type maxNonownedEntriesPerRow = 0;
70 Kokkos::parallel_reduce(
71 Kokkos::RangePolicy<execution_space>(0, numLocalRows),
72 KOKKOS_LAMBDA(local_ordinal_type i, local_ordinal_type & lmaxOwned, local_ordinal_type & lmaxNonowned) {
73 const local_ordinal_type lr = lclrow(i);
74 local_ordinal_type rowNumOwned = rowptr(lr + 1) - rowptr(lr);
75 if (rowNumOwned > lmaxOwned)
76 lmaxOwned = rowNumOwned;
77 // rowptr_remote won't be allocated for single-rank problems
78 if (rowptr_remote.extent(0)) {
79 local_ordinal_type rowNumNonowned = rowptr_remote(lr + 1) - rowptr_remote(lr);
80 if (rowNumNonowned > lmaxNonowned)
81 lmaxNonowned = rowNumNonowned;
82 } else {
83 lmaxNonowned = 0;
84 }
85 },
86 Kokkos::Max<local_ordinal_type>(maxOwnedEntriesPerRow), Kokkos::Max<local_ordinal_type>(maxNonownedEntriesPerRow));
87 // Allocate the two offsets views now that we know the dimensions
88 // For each one, the middle dimension is 0 for A offsets and 1 for x offsets.
89 // Packing them together in one view improves cache line utilization
90 amd.A_x_offsets = i64_3d_view("amd.A_x_offsets", numLocalRows, 2, maxOwnedEntriesPerRow);
91 amd.A_x_offsets_remote = i64_3d_view("amd.A_x_offsets_remote", numLocalRows, 2, maxNonownedEntriesPerRow);
92 auto A_x_offsets = amd.A_x_offsets;
93 auto A_x_offsets_remote = amd.A_x_offsets_remote;
94 // Now, populate all the offsets. Use ArithTraits<int64_t>::min to mark absent entries.
95 Kokkos::parallel_for(
96 Kokkos::RangePolicy<execution_space>(0, numLocalRows),
97 KOKKOS_LAMBDA(local_ordinal_type i) {
98 const local_ordinal_type lr = lclrow(i);
99 const size_type A_k0 = A_block_rowptr(lr);
100 // Owned entries
101 size_type rowBegin = rowptr(lr);
102 local_ordinal_type rowNumOwned = rowptr(lr + 1) - rowBegin;
103 for (local_ordinal_type entry = 0; entry < maxOwnedEntriesPerRow; entry++) {
104 if (entry < rowNumOwned) {
105 const size_type j = A_k0 + A_colindsub(rowBegin + entry);
106 const local_ordinal_type A_colind_at_j = A_colind(j);
107 const local_ordinal_type loc = is_dm2cm_active ? dm2cm(A_colind_at_j) : A_colind_at_j;
108 A_x_offsets(i, 0, entry) = int64_t(j) * blocksize_square;
109 A_x_offsets(i, 1, entry) = int64_t(loc) * blocksize;
110 } else {
111 A_x_offsets(i, 0, entry) = KokkosKernels::ArithTraits<int64_t>::min();
112 A_x_offsets(i, 1, entry) = KokkosKernels::ArithTraits<int64_t>::min();
113 }
114 }
115 // Nonowned entries
116 if (rowptr_remote.extent(0)) {
117 rowBegin = rowptr_remote(lr);
118 local_ordinal_type rowNumNonowned = rowptr_remote(lr + 1) - rowBegin;
119 for (local_ordinal_type entry = 0; entry < maxNonownedEntriesPerRow; entry++) {
120 if (entry < rowNumNonowned) {
121 const size_type j = A_k0 + A_colindsub_remote(rowBegin + entry);
122 const local_ordinal_type A_colind_at_j = A_colind(j);
123 const local_ordinal_type loc = A_colind_at_j - numLocalRows;
124 A_x_offsets_remote(i, 0, entry) = int64_t(j) * blocksize_square;
125 A_x_offsets_remote(i, 1, entry) = int64_t(loc) * blocksize;
126 } else {
127 A_x_offsets_remote(i, 0, entry) = KokkosKernels::ArithTraits<int64_t>::min();
128 A_x_offsets_remote(i, 1, entry) = KokkosKernels::ArithTraits<int64_t>::min();
129 }
130 }
131 }
132 });
133 } else {
134 // amd.rowptr points to both owned and nonowned entries, so it tells us how many columns (last dim) A_x_offsets should have.
135 local_ordinal_type maxEntriesPerRow = 0;
136 Kokkos::parallel_reduce(
137 Kokkos::RangePolicy<execution_space>(0, numLocalRows),
138 KOKKOS_LAMBDA(local_ordinal_type i, local_ordinal_type & lmax) {
139 const local_ordinal_type lr = lclrow(i);
140 local_ordinal_type rowNum = rowptr(lr + 1) - rowptr(lr);
141 if (rowNum > lmax)
142 lmax = rowNum;
143 },
144 Kokkos::Max<local_ordinal_type>(maxEntriesPerRow));
145 amd.A_x_offsets = i64_3d_view("amd.A_x_offsets", numLocalRows, 2, maxEntriesPerRow);
146 auto A_x_offsets = amd.A_x_offsets;
147 // Populate A,x offsets. Use ArithTraits<int64_t>::min to mark absent entries.
148 // For x offsets, add a shift blocksize*numLocalRows to represent that it indexes into x_remote instead of x.
149 Kokkos::parallel_for(
150 Kokkos::RangePolicy<execution_space>(0, numLocalRows),
151 KOKKOS_LAMBDA(local_ordinal_type i) {
152 const local_ordinal_type lr = lclrow(i);
153 const size_type A_k0 = A_block_rowptr(lr);
154 // Owned entries
155 size_type rowBegin = rowptr(lr);
156 local_ordinal_type rowOwned = rowptr(lr + 1) - rowBegin;
157 for (local_ordinal_type entry = 0; entry < maxEntriesPerRow; entry++) {
158 if (entry < rowOwned) {
159 const size_type j = A_k0 + A_colindsub(rowBegin + entry);
160 A_x_offsets(i, 0, entry) = j * blocksize_square;
161 const local_ordinal_type A_colind_at_j = A_colind(j);
162 if (A_colind_at_j < numLocalRows) {
163 const local_ordinal_type loc = is_dm2cm_active ? dm2cm[A_colind_at_j] : A_colind_at_j;
164 A_x_offsets(i, 1, entry) = int64_t(loc) * blocksize;
165 } else {
166 A_x_offsets(i, 1, entry) = int64_t(A_colind_at_j) * blocksize;
167 }
168 } else {
169 A_x_offsets(i, 0, entry) = KokkosKernels::ArithTraits<int64_t>::min();
170 A_x_offsets(i, 1, entry) = KokkosKernels::ArithTraits<int64_t>::min();
171 }
172 }
173 });
174 }
175}
176
180static inline int ComputeResidualVectorRecommendedCudaVectorSize(const int blksize,
181 const int team_size) {
182 int total_team_size(0);
183 if (blksize <= 5)
184 total_team_size = 32;
185 else if (blksize <= 9)
186 total_team_size = 32; // 64
187 else if (blksize <= 12)
188 total_team_size = 96;
189 else if (blksize <= 16)
190 total_team_size = 128;
191 else if (blksize <= 20)
192 total_team_size = 160;
193 else
194 total_team_size = 160;
195 return total_team_size / team_size;
196}
197
198static inline int ComputeResidualVectorRecommendedHIPVectorSize(const int blksize,
199 const int team_size) {
200 int total_team_size(0);
201 if (blksize <= 5)
202 total_team_size = 32;
203 else if (blksize <= 9)
204 total_team_size = 32; // 64
205 else if (blksize <= 12)
206 total_team_size = 96;
207 else if (blksize <= 16)
208 total_team_size = 128;
209 else if (blksize <= 20)
210 total_team_size = 160;
211 else
212 total_team_size = 160;
213 return total_team_size / team_size;
214}
215
216static inline int ComputeResidualVectorRecommendedSYCLVectorSize(const int blksize,
217 const int team_size) {
218 int total_team_size(0);
219 if (blksize <= 5)
220 total_team_size = 32;
221 else if (blksize <= 9)
222 total_team_size = 32; // 64
223 else if (blksize <= 12)
224 total_team_size = 96;
225 else if (blksize <= 16)
226 total_team_size = 128;
227 else if (blksize <= 20)
228 total_team_size = 160;
229 else
230 total_team_size = 160;
231 return total_team_size / team_size;
232}
233
234template <typename T>
235static inline int ComputeResidualVectorRecommendedVectorSize(const int blksize,
236 const int team_size) {
237 if (is_cuda<T>::value)
238 return ComputeResidualVectorRecommendedCudaVectorSize(blksize, team_size);
239 if (is_hip<T>::value)
240 return ComputeResidualVectorRecommendedHIPVectorSize(blksize, team_size);
241 if (is_sycl<T>::value)
242 return ComputeResidualVectorRecommendedSYCLVectorSize(blksize, team_size);
243 return -1;
244}
245
246template <typename MatrixType>
247struct ComputeResidualFunctor {
248 using impl_type = BlockHelperDetails::ImplType<MatrixType>;
249 using node_device_type = typename impl_type::node_device_type;
250 using execution_space = typename impl_type::execution_space;
251 using memory_space = typename impl_type::memory_space;
252
253 using local_ordinal_type = typename impl_type::local_ordinal_type;
254 using size_type = typename impl_type::size_type;
255 using impl_scalar_type = typename impl_type::impl_scalar_type;
256 using magnitude_type = typename impl_type::magnitude_type;
257 using btdm_scalar_type = typename impl_type::btdm_scalar_type;
258 using btdm_magnitude_type = typename impl_type::btdm_magnitude_type;
260 using local_ordinal_type_1d_view = typename impl_type::local_ordinal_type_1d_view;
261 using size_type_1d_view = typename impl_type::size_type_1d_view;
262 using tpetra_block_access_view_type = typename impl_type::tpetra_block_access_view_type; // block crs (layout right)
263 using impl_scalar_type_1d_view = typename impl_type::impl_scalar_type_1d_view;
264 using impl_scalar_type_2d_view_tpetra = typename impl_type::impl_scalar_type_2d_view_tpetra; // block multivector (layout left)
265 using vector_type_3d_view = typename impl_type::vector_type_3d_view;
266 using btdm_scalar_type_4d_view = typename impl_type::btdm_scalar_type_4d_view;
267 using i64_3d_view = typename impl_type::i64_3d_view;
268 static constexpr int vector_length = impl_type::vector_length;
269
271 using member_type = typename Kokkos::TeamPolicy<execution_space>::member_type;
272
273 ConstUnmanaged<impl_scalar_type_2d_view_tpetra> b;
274 ConstUnmanaged<impl_scalar_type_2d_view_tpetra> x; // x_owned
275 ConstUnmanaged<impl_scalar_type_2d_view_tpetra> x_remote;
276 Unmanaged<impl_scalar_type_2d_view_tpetra> y;
277 Unmanaged<vector_type_3d_view> y_packed;
278 Unmanaged<btdm_scalar_type_4d_view> y_packed_scalar;
279
280 // AmD information
281 const ConstUnmanaged<size_type_1d_view> rowptr, rowptr_remote;
282 const ConstUnmanaged<local_ordinal_type_1d_view> colindsub, colindsub_remote;
283 const ConstUnmanaged<impl_scalar_type_1d_view> tpetra_values;
284
285 // block crs graph information
286 // for cuda (kokkos crs graph uses a different size_type from size_t)
287 const ConstUnmanaged<Kokkos::View<size_t *, node_device_type>> A_block_rowptr;
288 const ConstUnmanaged<Kokkos::View<size_t *, node_device_type>> A_point_rowptr;
289 const ConstUnmanaged<Kokkos::View<local_ordinal_type *, node_device_type>> A_colind;
290
291 // blocksize
292 const local_ordinal_type blocksize_requested;
293
294 // part interface
295 const ConstUnmanaged<local_ordinal_type_1d_view> part2packrowidx0;
296 const ConstUnmanaged<local_ordinal_type_1d_view> part2rowidx0;
297 const ConstUnmanaged<local_ordinal_type_1d_view> rowidx2part;
298 const ConstUnmanaged<local_ordinal_type_1d_view> partptr;
299 const ConstUnmanaged<local_ordinal_type_1d_view> lclrow;
300 const ConstUnmanaged<local_ordinal_type_1d_view> dm2cm;
301
302 // block offsets
303 const ConstUnmanaged<i64_3d_view> A_x_offsets;
304 const ConstUnmanaged<i64_3d_view> A_x_offsets_remote;
305
306 const bool is_dm2cm_active;
307 const bool hasBlockCrsMatrix;
308
309 ComputeResidualFunctor(const ComputeResidualVector<MatrixType> &crv)
310 : rowptr(crv.rowptr)
311 , rowptr_remote(crv.rowptr_remote)
312 , colindsub(crv.colindsub)
313 , colindsub_remote(crv.colindsub_remote)
314 , tpetra_values(crv.tpetra_values)
315 , A_block_rowptr(crv.A_block_rowptr)
316 , A_point_rowptr(crv.A_point_rowptr)
317 , A_colind(crv.A_colind)
318 , blocksize_requested(crv.blocksize_requested)
319 , part2packrowidx0(crv.part2packrowidx0)
320 , part2rowidx0(crv.part2rowidx0)
321 , rowidx2part(crv.rowidx2part)
322 , partptr(crv.partptr)
323 , lclrow(crv.lclrow)
324 , dm2cm(crv.dm2cm)
325 , A_x_offsets(crv.A_x_offsets)
326 , A_x_offsets_remote(crv.A_x_offsets_remote)
327 , is_dm2cm_active(crv.is_dm2cm_active)
328 , hasBlockCrsMatrix(crv.hasBlockCrsMatrix) {}
329
330 inline void
331 SerialDot(const local_ordinal_type &blocksize,
332 const local_ordinal_type &lclRowID,
333 const local_ordinal_type &lclColID,
334 const local_ordinal_type &ii,
335 const ConstUnmanaged<local_ordinal_type_1d_view> colindsub_,
336 const impl_scalar_type *const KOKKOS_RESTRICT xx,
337 /* */ impl_scalar_type *KOKKOS_RESTRICT yy) const {
338 const size_type Aj_c = colindsub_(lclColID);
339 auto point_row_offset = A_point_rowptr(lclRowID * blocksize + ii) + Aj_c * blocksize;
340 impl_scalar_type val = 0;
341#if defined(KOKKOS_ENABLE_PRAGMA_IVDEP)
342#pragma ivdep
343#endif
344#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
345#pragma unroll
346#endif
347 for (local_ordinal_type k1 = 0; k1 < blocksize; ++k1)
348 val += tpetra_values(point_row_offset + k1) * xx[k1];
349 yy[ii] -= val;
350 }
351
352 inline void
353 SerialGemv(const local_ordinal_type &blocksize,
354 const impl_scalar_type *const KOKKOS_RESTRICT AA,
355 const impl_scalar_type *const KOKKOS_RESTRICT xx,
356 /* */ impl_scalar_type *KOKKOS_RESTRICT yy) const {
357 using tlb = BlockHelperDetails::TpetraLittleBlock<Tpetra::Impl::BlockCrsMatrixLittleBlockArrayLayout>;
358 for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0) {
359 impl_scalar_type val = 0;
360#if defined(KOKKOS_ENABLE_PRAGMA_IVDEP)
361#pragma ivdep
362#endif
363#if defined(KOKKOS_ENABLE_PRAGMA_UNROLL)
364#pragma unroll
365#endif
366 for (local_ordinal_type k1 = 0; k1 < blocksize; ++k1)
367 val += AA[tlb::getFlatIndex(k0, k1, blocksize)] * xx[k1];
368 yy[k0] -= val;
369 }
370 }
371
372 template <typename bbViewType, typename yyViewType>
373 KOKKOS_INLINE_FUNCTION void
374 VectorCopy(const member_type &member,
375 const local_ordinal_type &blocksize,
376 const bbViewType &bb,
377 const yyViewType &yy) const {
378 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member, blocksize), [&](const local_ordinal_type &k0) {
379 yy(k0) = static_cast<typename yyViewType::const_value_type>(bb(k0));
380 });
381 }
382
383 template <typename xxViewType, typename yyViewType>
384 KOKKOS_INLINE_FUNCTION void
385 VectorDot(const member_type &member,
386 const local_ordinal_type &blocksize,
387 const local_ordinal_type &lclRowID,
388 const local_ordinal_type &lclColID,
389 const local_ordinal_type &ii,
390 const ConstUnmanaged<local_ordinal_type_1d_view> colindsub_,
391 const xxViewType &xx,
392 const yyViewType &yy) const {
393 const size_type Aj_c = colindsub_(lclColID);
394 auto point_row_offset = A_point_rowptr(lclRowID * blocksize + ii) + Aj_c * blocksize;
395 impl_scalar_type val = 0;
396 Kokkos::parallel_reduce(
397 Kokkos::ThreadVectorRange(member, blocksize),
398 [&](const local_ordinal_type &k1, impl_scalar_type &update) {
399 update += tpetra_values(point_row_offset + k1) * xx(k1);
400 },
401 val);
402 Kokkos::single(Kokkos::PerThread(member),
403 [=]() {
404 Kokkos::atomic_add(&yy(ii), typename yyViewType::const_value_type(-val));
405 });
406 }
407
408 // BMK: This version coalesces accesses to AA for LayoutRight blocks.
409 template <typename AAViewType, typename xxViewType, typename yyViewType>
410 KOKKOS_INLINE_FUNCTION void
411 VectorGemv(const member_type &member,
412 const local_ordinal_type &blocksize,
413 const AAViewType &AA,
414 const xxViewType &xx,
415 const yyViewType &yy) const {
416 for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0) {
417 impl_scalar_type val = 0;
418 Kokkos::parallel_reduce(
419 Kokkos::ThreadVectorRange(member, blocksize),
420 [&](const local_ordinal_type &k1, impl_scalar_type &update) {
421 update += AA(k0, k1) * xx(k1);
422 },
423 val);
424 Kokkos::single(Kokkos::PerThread(member),
425 [=]() {
426 Kokkos::atomic_add(&yy(k0), -val);
427 });
428 }
429 }
430
431 struct SeqTag {};
432
433 void
434 operator()(const SeqTag &, const local_ordinal_type &i) const {
435 const local_ordinal_type blocksize = blocksize_requested;
436 const local_ordinal_type blocksize_square = blocksize * blocksize;
437
438 // constants
439 const Kokkos::pair<local_ordinal_type, local_ordinal_type> block_range(0, blocksize);
440 const local_ordinal_type num_vectors = y.extent(1);
441 const local_ordinal_type row = i * blocksize;
442 for (local_ordinal_type col = 0; col < num_vectors; ++col) {
443 // y := b
444 impl_scalar_type *yy = &y(row, col);
445 const impl_scalar_type *const bb = &b(row, col);
446 memcpy(yy, bb, sizeof(impl_scalar_type) * blocksize);
447
448 // y -= Rx
449 const size_type A_k0 = A_block_rowptr[i];
450 for (size_type k = rowptr[i]; k < rowptr[i + 1]; ++k) {
451 const size_type j = A_k0 + colindsub[k];
452 const impl_scalar_type *const xx = &x(A_colind[j] * blocksize, col);
453 if (hasBlockCrsMatrix) {
454 const impl_scalar_type *const AA = &tpetra_values(j * blocksize_square);
455 SerialGemv(blocksize, AA, xx, yy);
456 } else {
457 for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0)
458 SerialDot(blocksize, i, k, k0, colindsub, xx, yy);
459 }
460 }
461 }
462 }
463
464 KOKKOS_INLINE_FUNCTION
465 void
466 operator()(const SeqTag &, const member_type &member) const {
467 // constants
468 const local_ordinal_type blocksize = blocksize_requested;
469 const local_ordinal_type blocksize_square = blocksize * blocksize;
470
471 const local_ordinal_type lr = member.league_rank();
472 const Kokkos::pair<local_ordinal_type, local_ordinal_type> block_range(0, blocksize);
473 const local_ordinal_type num_vectors = y.extent(1);
474
475 // subview pattern
476 auto bb = Kokkos::subview(b, block_range, 0);
477 auto xx = bb;
478 auto A_block_cst = ConstUnmanaged<tpetra_block_access_view_type>(tpetra_values.data(), blocksize, blocksize);
479
480 const local_ordinal_type row = lr * blocksize;
481 for (local_ordinal_type col = 0; col < num_vectors; ++col) {
482 // y := b
483 auto yy = Kokkos::subview(y, Kokkos::make_pair(row, row + blocksize), col);
484 bb.assign_data(&b(row, col));
485 if (member.team_rank() == 0)
486 VectorCopy(member, blocksize, bb, yy);
487 member.team_barrier();
488
489 // y -= Rx
490 const size_type A_k0 = A_block_rowptr[lr];
491
492 if (hasBlockCrsMatrix) {
493 Kokkos::parallel_for(Kokkos::TeamThreadRange(member, rowptr[lr], rowptr[lr + 1]),
494 [&](const local_ordinal_type &k) {
495 const size_type j = A_k0 + colindsub[k];
496 xx.assign_data(&x(A_colind[j] * blocksize, col));
497 A_block_cst.assign_data(&tpetra_values(j * blocksize_square));
498 VectorGemv(member, blocksize, A_block_cst, xx, yy);
499 });
500 } else {
501 Kokkos::parallel_for(Kokkos::TeamThreadRange(member, rowptr[lr], rowptr[lr + 1]),
502 [&](const local_ordinal_type &k) {
503 const size_type j = A_k0 + colindsub[k];
504 xx.assign_data(&x(A_colind[j] * blocksize, col));
505
506 for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0)
507 VectorDot(member, blocksize, lr, k, k0, colindsub, xx, yy);
508 });
509 }
510 }
511 }
512
513 // * B: block size for compile-time specialization, or 0 for general case
514 // * async: true if using async importer. overlap is not used in this case.
515 // Whether a column is owned or nonowned is decided at runtime.
516 // * overlap: true if processing the columns that are not locally owned,
517 // false if processing locally owned columns.
518 // * haveBlockMatrix: true if A is a BlockCrsMatrix, false if it's CrsMatrix.
519 template <int B, bool async, bool overlap, bool haveBlockMatrix>
520 struct GeneralTag {
521 static_assert(!(async && overlap),
522 "ComputeResidualVector: async && overlap is not a valid configuration for GeneralTag");
523 };
524
525 // Define AsyncTag and OverlapTag in terms of GeneralTag:
526 // P == 0 means only compute on owned columns
527 // P == 1 means only compute on nonowned columns
528 template <int P, int B, bool haveBlockMatrix>
529 using OverlapTag = GeneralTag<B, false, P != 0, haveBlockMatrix>;
530
531 template <int B, bool haveBlockMatrix>
532 using AsyncTag = GeneralTag<B, true, false, haveBlockMatrix>;
533
534 // CPU implementation for all cases
535 template <int B, bool async, bool overlap, bool haveBlockMatrix>
536 void
537 operator()(const GeneralTag<B, async, overlap, haveBlockMatrix> &, const local_ordinal_type &rowidx) const {
538 // B == 0 means blocksize is a runtime value, provided via blocksize_requested.
539 // B != 0 is a compile-time block size.
540 const local_ordinal_type blocksize = (B == 0 ? blocksize_requested : B);
541
542 // constants
543 const local_ordinal_type partidx = rowidx2part(rowidx);
544 const local_ordinal_type pri = part2packrowidx0(partidx) + (rowidx - partptr(partidx));
545 const local_ordinal_type v = partidx % vector_length;
546
547 const local_ordinal_type num_vectors = y_packed.extent(2);
548 const local_ordinal_type num_local_rows = lclrow.extent(0);
549
550 // temporary buffer for y flat
551 ArrayHelper<impl_scalar_type, B> yy_alloc(blocksize_requested);
552 impl_scalar_type *yy = yy_alloc.get();
553
554 const local_ordinal_type lr = lclrow(rowidx);
555
556 auto colindsub_used = overlap ? colindsub_remote : colindsub;
557 auto rowptr_used = overlap ? rowptr_remote : rowptr;
558
559 for (local_ordinal_type col = 0; col < num_vectors; ++col) {
560 if constexpr (overlap) {
561 // y (temporary) := 0
562 memset((void *)yy, 0, sizeof(impl_scalar_type) * blocksize);
563 } else {
564 // y := b
565 const local_ordinal_type row = lr * blocksize;
566 memcpy(yy, &b(row, col), sizeof(impl_scalar_type) * blocksize);
567 }
568
569 // y -= Rx
570 const size_type A_k0 = A_block_rowptr[lr];
571 for (size_type k = rowptr_used[lr]; k < rowptr_used[lr + 1]; ++k) {
572 const size_type j = A_k0 + colindsub_used[k];
573 const local_ordinal_type A_colind_at_j = A_colind[j];
574 if constexpr (haveBlockMatrix) {
575 const local_ordinal_type blocksize_square = blocksize * blocksize;
576 const impl_scalar_type *const AA = &tpetra_values(j * blocksize_square);
577 if ((!async && !overlap) || (async && A_colind_at_j < num_local_rows)) {
578 const auto loc = is_dm2cm_active ? dm2cm[A_colind_at_j] : A_colind_at_j;
579 const impl_scalar_type *const xx = &x(loc * blocksize, col);
580 SerialGemv(blocksize, AA, xx, yy);
581 } else {
582 const auto loc = A_colind_at_j - num_local_rows;
583 const impl_scalar_type *const xx_remote = &x_remote(loc * blocksize, col);
584 SerialGemv(blocksize, AA, xx_remote, yy);
585 }
586 } else {
587 if ((!async && !overlap) || (async && A_colind_at_j < num_local_rows)) {
588 const auto loc = is_dm2cm_active ? dm2cm[A_colind_at_j] : A_colind_at_j;
589 const impl_scalar_type *const xx = &x(loc * blocksize, col);
590 for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0)
591 SerialDot(blocksize, lr, k, k0, colindsub_used, xx, yy);
592 } else {
593 const auto loc = A_colind_at_j - num_local_rows;
594 const impl_scalar_type *const xx_remote = &x_remote(loc * blocksize, col);
595 for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0)
596 SerialDot(blocksize, lr, k, k0, colindsub_used, xx_remote, yy);
597 }
598 }
599 }
600 // move yy to y_packed
601 if constexpr (overlap) {
602 for (local_ordinal_type k = 0; k < blocksize; ++k)
603 y_packed(pri, k, col)[v] += yy[k];
604 } else {
605 for (local_ordinal_type k = 0; k < blocksize; ++k)
606 y_packed(pri, k, col)[v] = yy[k];
607 }
608 }
609 }
610
611 // GPU implementation for hasBlockCrsMatrix == true
612 template <int B, bool async, bool overlap>
613 KOKKOS_INLINE_FUNCTION void
614 operator()(const GeneralTag<B, async, overlap, true> &, const member_type &member) const {
615 const local_ordinal_type blocksize = (B == 0 ? blocksize_requested : B);
616
617 // constants
618 const local_ordinal_type rowidx = member.league_rank();
619 const local_ordinal_type partidx = rowidx2part(rowidx);
620 const local_ordinal_type pri = part2packrowidx0(partidx) + (rowidx - partptr(partidx));
621 const local_ordinal_type v = partidx % vector_length;
622
623 const Kokkos::pair<local_ordinal_type, local_ordinal_type> block_range(0, blocksize);
624 const local_ordinal_type num_vectors = y_packed_scalar.extent(2);
625 const local_ordinal_type num_local_rows = lclrow.extent(0);
626
627 // subview pattern
628 auto bb = Kokkos::subview(b, block_range, 0);
629 auto xx = bb;
630 auto yy = Kokkos::subview(y_packed_scalar, 0, block_range, 0, 0);
631 auto A_block_cst = ConstUnmanaged<tpetra_block_access_view_type>(tpetra_values.data(), blocksize, blocksize);
632
633 // Get shared allocation for a local copy of x, Ax, and A
634 impl_scalar_type *local_Ax = reinterpret_cast<impl_scalar_type *>(member.team_scratch(0).get_shmem(blocksize * sizeof(impl_scalar_type)));
635 impl_scalar_type *local_x = reinterpret_cast<impl_scalar_type *>(member.thread_scratch(0).get_shmem(blocksize * sizeof(impl_scalar_type)));
636
637 const local_ordinal_type lr = lclrow(rowidx);
638 const local_ordinal_type row = lr * blocksize;
639 for (local_ordinal_type col = 0; col < num_vectors; ++col) {
640 if (col)
641 member.team_barrier();
642 // y -= Rx
643 // Initialize accumulation array
644 Kokkos::parallel_for(Kokkos::TeamVectorRange(member, blocksize), [&](const local_ordinal_type &i) {
645 local_Ax[i] = 0;
646 });
647 member.team_barrier();
648
649 int numEntries;
650 if constexpr (!overlap) {
651 numEntries = A_x_offsets.extent(2);
652 } else {
653 numEntries = A_x_offsets_remote.extent(2);
654 }
655
656 Kokkos::parallel_for(Kokkos::TeamThreadRange(member, 0, numEntries),
657 [&](const int k) {
658 int64_t A_offset = overlap ? A_x_offsets_remote(rowidx, 0, k) : A_x_offsets(rowidx, 0, k);
659 int64_t x_offset = overlap ? A_x_offsets_remote(rowidx, 1, k) : A_x_offsets(rowidx, 1, k);
660 if (A_offset != KokkosKernels::ArithTraits<int64_t>::min()) {
661 A_block_cst.assign_data(tpetra_values.data() + A_offset);
662 // Pull x into local memory
663 if constexpr (async) {
664 size_type remote_cutoff = blocksize * num_local_rows;
665 if (x_offset >= remote_cutoff)
666 xx.assign_data(&x_remote(x_offset - remote_cutoff, col));
667 else
668 xx.assign_data(&x(x_offset, col));
669 } else {
670 if constexpr (!overlap) {
671 xx.assign_data(&x(x_offset, col));
672 } else {
673 xx.assign_data(&x_remote(x_offset, col));
674 }
675 }
676
677 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member, blocksize), [&](const local_ordinal_type &i) {
678 local_x[i] = xx(i);
679 });
680
681 // MatVec op Ax += A*x
682 Kokkos::parallel_for(
683 Kokkos::ThreadVectorRange(member, blocksize),
684 [&](const local_ordinal_type &k0) {
685 impl_scalar_type val = 0;
686 for (int k1 = 0; k1 < blocksize; k1++)
687 val += A_block_cst(k0, k1) * local_x[k1];
688 Kokkos::atomic_add(local_Ax + k0, val);
689 });
690 }
691 });
692 member.team_barrier();
693 // Update y = b - local_Ax
694 yy.assign_data(&y_packed_scalar(pri, 0, col, v));
695 bb.assign_data(&b(row, col));
696 Kokkos::parallel_for(Kokkos::TeamVectorRange(member, blocksize), [&](const local_ordinal_type &i) {
697 if (!overlap)
698 yy(i) = bb(i) - local_Ax[i];
699 else
700 yy(i) -= local_Ax[i];
701 });
702 }
703 }
704
705 // GPU implementation for hasBlockCrsMatrix == false
706 template <int B, bool async, bool overlap>
707 KOKKOS_INLINE_FUNCTION void
708 operator()(const GeneralTag<B, async, overlap, false> &, const member_type &member) const {
709 const local_ordinal_type blocksize = (B == 0 ? blocksize_requested : B);
710
711 // constants
712 const local_ordinal_type rowidx = member.league_rank();
713 const local_ordinal_type partidx = rowidx2part(rowidx);
714 const local_ordinal_type pri = part2packrowidx0(partidx) + (rowidx - partptr(partidx));
715 const local_ordinal_type v = partidx % vector_length;
716
717 const Kokkos::pair<local_ordinal_type, local_ordinal_type> block_range(0, blocksize);
718 const local_ordinal_type num_vectors = y_packed_scalar.extent(2);
719 const local_ordinal_type num_local_rows = lclrow.extent(0);
720
721 // subview pattern
722 auto bb = Kokkos::subview(b, block_range, 0);
723 auto xx = bb;
724 auto xx_remote = bb;
725 auto yy = Kokkos::subview(y_packed_scalar, 0, block_range, 0, 0);
726 auto A_block_cst = ConstUnmanaged<tpetra_block_access_view_type>(tpetra_values.data(), blocksize, blocksize);
727 auto colindsub_used = overlap ? colindsub_remote : colindsub;
728 auto rowptr_used = overlap ? rowptr_remote : rowptr;
729
730 const local_ordinal_type lr = lclrow(rowidx);
731 const local_ordinal_type row = lr * blocksize;
732 for (local_ordinal_type col = 0; col < num_vectors; ++col) {
733 yy.assign_data(&y_packed_scalar(pri, 0, col, v));
734 if (!overlap) {
735 // y := b
736 bb.assign_data(&b(row, col));
737 if (member.team_rank() == 0)
738 VectorCopy(member, blocksize, bb, yy);
739 member.team_barrier();
740 }
741
742 // y -= Rx
743 const size_type A_k0 = A_block_rowptr[lr];
744 Kokkos::parallel_for(Kokkos::TeamThreadRange(member, rowptr_used[lr], rowptr_used[lr + 1]),
745 [&](const local_ordinal_type &k) {
746 const size_type j = A_k0 + colindsub_used[k];
747 const local_ordinal_type A_colind_at_j = A_colind[j];
748 if ((async && A_colind_at_j < num_local_rows) || (!async && !overlap)) {
749 const auto loc = is_dm2cm_active ? dm2cm[A_colind_at_j] : A_colind_at_j;
750 xx.assign_data(&x(loc * blocksize, col));
751 for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0)
752 VectorDot(member, blocksize, lr, k, k0, colindsub_used, xx, yy);
753 } else {
754 const auto loc = A_colind_at_j - num_local_rows;
755 xx_remote.assign_data(&x_remote(loc * blocksize, col));
756 for (local_ordinal_type k0 = 0; k0 < blocksize; ++k0)
757 VectorDot(member, blocksize, lr, k, k0, colindsub_used, xx_remote, yy);
758 }
759 });
760 }
761 }
762};
763
764// y = b - Rx; seq method
765template <typename MatrixType>
766void ComputeResidualVector<MatrixType, BlockTriDiContainerDetails::ImplSimdTag>::run(const impl_scalar_type_2d_view_tpetra &y_,
767 const Const<impl_scalar_type_2d_view_tpetra> &b_,
768 const impl_scalar_type_2d_view_tpetra &x_) {
769 IFPACK2_BLOCKHELPER_PROFILER_REGION_BEGIN;
770 IFPACK2_BLOCKHELPER_TIMER_WITH_FENCE("BlockTriDi::ComputeResidual::<SeqTag>", ComputeResidual0, execution_space);
771
772 using Functor = ComputeResidualFunctor<typename impl_type::matrix_type>;
773 Functor functor(*this);
774
775 functor.y = y_;
776 functor.b = b_;
777 functor.x = x_;
778 if constexpr (is_device<execution_space>::value) {
779 const local_ordinal_type blocksize = blocksize_requested;
780 const local_ordinal_type team_size = 8;
781 const local_ordinal_type vector_size = ComputeResidualVectorRecommendedVectorSize<execution_space>(blocksize, team_size);
782 const Kokkos::TeamPolicy<execution_space, typename Functor::SeqTag> policy(rowptr.extent(0) - 1, team_size, vector_size);
783 Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<SeqTag>", policy, functor);
784 } else {
785 const Kokkos::RangePolicy<execution_space, typename Functor::SeqTag> policy(0, rowptr.extent(0) - 1);
786 Kokkos::parallel_for("ComputeResidual::RangePolicy::run<SeqTag>", policy, functor);
787 }
788 IFPACK2_BLOCKHELPER_PROFILER_REGION_END;
789 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
790}
791
792// y = b - R (x , x_remote)
793template <typename MatrixType>
794void ComputeResidualVector<MatrixType, BlockTriDiContainerDetails::ImplSimdTag>::run(const vector_type_3d_view &y_packed_,
795 const Const<impl_scalar_type_2d_view_tpetra> &b_,
796 const impl_scalar_type_2d_view_tpetra &x_,
797 const impl_scalar_type_2d_view_tpetra &x_remote_) {
798 IFPACK2_BLOCKHELPER_PROFILER_REGION_BEGIN;
799 IFPACK2_BLOCKHELPER_TIMER_WITH_FENCE("BlockTriDi::ComputeResidual::<AsyncTag>", ComputeResidual0, execution_space);
800
801 using Functor = ComputeResidualFunctor<typename impl_type::matrix_type>;
802 Functor functor(*this);
803
804 functor.b = b_;
805 functor.x = x_;
806 functor.x_remote = x_remote_;
807 if constexpr (is_device<execution_space>::value) {
808 functor.y_packed_scalar = btdm_scalar_type_4d_view((btdm_scalar_type *)y_packed_.data(),
809 y_packed_.extent(0),
810 y_packed_.extent(1),
811 y_packed_.extent(2),
812 vector_length);
813 } else {
814 functor.y_packed = y_packed_;
815 }
816
817 if constexpr (is_device<execution_space>::value) {
818 const local_ordinal_type blocksize = blocksize_requested;
819 // local_ordinal_type vl_power_of_two = 1;
820 // for (;vl_power_of_two<=blocksize_requested;vl_power_of_two*=2);
821 // vl_power_of_two *= (vl_power_of_two < blocksize_requested ? 2 : 1);
822 // const local_ordinal_type vl = vl_power_of_two > vector_length ? vector_length : vl_power_of_two;
823#define BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(B) \
824 { \
825 if (this->hasBlockCrsMatrix) { \
826 const local_ordinal_type team_size = 8; \
827 const local_ordinal_type vector_size = 8; \
828 const size_t shmem_team_size = blocksize * sizeof(btdm_scalar_type); \
829 const size_t shmem_thread_size = blocksize * sizeof(btdm_scalar_type); \
830 Kokkos::TeamPolicy<execution_space, typename Functor::template AsyncTag<B, true>> \
831 policy(rowidx2part.extent(0), team_size, vector_size); \
832 policy.set_scratch_size(0, Kokkos::PerTeam(shmem_team_size), Kokkos::PerThread(shmem_thread_size)); \
833 Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<AsyncTag>", \
834 policy, functor); \
835 } else { \
836 const local_ordinal_type team_size = 8; \
837 const local_ordinal_type vector_size = ComputeResidualVectorRecommendedVectorSize<execution_space>(blocksize, team_size); \
838 const Kokkos::TeamPolicy<execution_space, typename Functor::template AsyncTag<B, false>> \
839 policy(rowidx2part.extent(0), team_size, vector_size); \
840 Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<AsyncTag>", \
841 policy, functor); \
842 } \
843 } \
844 break
845 switch (blocksize_requested) {
846 case 3: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(3);
847 case 5: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(5);
848 case 7: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(7);
849 case 9: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(9);
850 case 10: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(10);
851 case 11: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(11);
852 case 16: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(16);
853 case 17: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(17);
854 case 18: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(18);
855 default: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(0);
856 }
857#undef BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL
858 } else {
859#define BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(B) \
860 { \
861 if (this->hasBlockCrsMatrix) { \
862 const Kokkos::RangePolicy<execution_space, typename Functor::template AsyncTag<B, true>> policy(0, rowidx2part.extent(0)); \
863 Kokkos::parallel_for("ComputeResidual::RangePolicy::run<AsyncTag>", \
864 policy, functor); \
865 } else { \
866 const Kokkos::RangePolicy<execution_space, typename Functor::template AsyncTag<B, false>> policy(0, rowidx2part.extent(0)); \
867 Kokkos::parallel_for("ComputeResidual::RangePolicy::run<AsyncTag>", \
868 policy, functor); \
869 } \
870 } \
871 break
872
873 switch (blocksize_requested) {
874 case 3: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(3);
875 case 5: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(5);
876 case 7: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(7);
877 case 9: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(9);
878 case 10: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(10);
879 case 11: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(11);
880 case 16: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(16);
881 case 17: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(17);
882 case 18: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(18);
883 default: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(0);
884 }
885#undef BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL
886 }
887 IFPACK2_BLOCKHELPER_PROFILER_REGION_END;
888 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
889}
890
891// y = b - R (y , y_remote)
892template <typename MatrixType>
893void ComputeResidualVector<MatrixType, BlockTriDiContainerDetails::ImplSimdTag>::run(const vector_type_3d_view &y_packed_,
894 const Const<impl_scalar_type_2d_view_tpetra> &b_,
895 const impl_scalar_type_2d_view_tpetra &x_,
896 const impl_scalar_type_2d_view_tpetra &x_remote_,
897 const bool compute_owned) {
898 IFPACK2_BLOCKHELPER_PROFILER_REGION_BEGIN;
899 IFPACK2_BLOCKHELPER_TIMER_WITH_FENCE("BlockTriDi::ComputeResidual::<OverlapTag>", ComputeResidual0, execution_space);
900
901 using Functor = ComputeResidualFunctor<typename impl_type::matrix_type>;
902 Functor functor(*this);
903
904 functor.b = b_;
905 functor.x = x_;
906 functor.x_remote = x_remote_;
907 if constexpr (is_device<execution_space>::value) {
908 functor.y_packed_scalar = btdm_scalar_type_4d_view((btdm_scalar_type *)y_packed_.data(),
909 y_packed_.extent(0),
910 y_packed_.extent(1),
911 y_packed_.extent(2),
912 vector_length);
913 } else {
914 functor.y_packed = y_packed_;
915 }
916
917 if constexpr (is_device<execution_space>::value) {
918 const local_ordinal_type blocksize = blocksize_requested;
919 // local_ordinal_type vl_power_of_two = 1;
920 // for (;vl_power_of_two<=blocksize_requested;vl_power_of_two*=2);
921 // vl_power_of_two *= (vl_power_of_two < blocksize_requested ? 2 : 1);
922 // const local_ordinal_type vl = vl_power_of_two > vector_length ? vector_length : vl_power_of_two;
923#define BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(B) \
924 if (this->hasBlockCrsMatrix) { \
925 const local_ordinal_type team_size = 8; \
926 const local_ordinal_type vector_size = 8; \
927 const size_t shmem_team_size = blocksize * sizeof(btdm_scalar_type); \
928 const size_t shmem_thread_size = blocksize * sizeof(btdm_scalar_type); \
929 if (compute_owned) { \
930 Kokkos::TeamPolicy<execution_space, typename Functor::template OverlapTag<0, B, true>> \
931 policy(rowidx2part.extent(0), team_size, vector_size); \
932 policy.set_scratch_size(0, Kokkos::PerTeam(shmem_team_size), Kokkos::PerThread(shmem_thread_size)); \
933 Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<OverlapTag<0> >", policy, functor); \
934 } else { \
935 Kokkos::TeamPolicy<execution_space, typename Functor::template OverlapTag<1, B, true>> \
936 policy(rowidx2part.extent(0), team_size, vector_size); \
937 policy.set_scratch_size(0, Kokkos::PerTeam(shmem_team_size), Kokkos::PerThread(shmem_thread_size)); \
938 Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<OverlapTag<1> >", policy, functor); \
939 } \
940 } else { \
941 const local_ordinal_type team_size = 8; \
942 const local_ordinal_type vector_size = ComputeResidualVectorRecommendedVectorSize<execution_space>(blocksize, team_size); \
943 if (compute_owned) { \
944 const Kokkos::TeamPolicy<execution_space, typename Functor::template OverlapTag<0, B, false>> \
945 policy(rowidx2part.extent(0), team_size, vector_size); \
946 Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<OverlapTag<0> >", policy, functor); \
947 } else { \
948 const Kokkos::TeamPolicy<execution_space, typename Functor::template OverlapTag<1, B, false>> \
949 policy(rowidx2part.extent(0), team_size, vector_size); \
950 Kokkos::parallel_for("ComputeResidual::TeamPolicy::run<OverlapTag<1> >", policy, functor); \
951 } \
952 } \
953 break
954 switch (blocksize_requested) {
955 case 3: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(3);
956 case 5: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(5);
957 case 7: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(7);
958 case 9: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(9);
959 case 10: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(10);
960 case 11: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(11);
961 case 16: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(16);
962 case 17: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(17);
963 case 18: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(18);
964 default: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(0);
965 }
966#undef BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL
967 } else {
968#define BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(B) \
969 if (this->hasBlockCrsMatrix) { \
970 if (compute_owned) { \
971 const Kokkos::RangePolicy<execution_space, typename Functor::template OverlapTag<0, B, true>> \
972 policy(0, rowidx2part.extent(0)); \
973 Kokkos::parallel_for("ComputeResidual::RangePolicy::run<OverlapTag<0> >", policy, functor); \
974 } else { \
975 const Kokkos::RangePolicy<execution_space, typename Functor::template OverlapTag<1, B, true>> \
976 policy(0, rowidx2part.extent(0)); \
977 Kokkos::parallel_for("ComputeResidual::RangePolicy::run<OverlapTag<1> >", policy, functor); \
978 } \
979 } else { \
980 if (compute_owned) { \
981 const Kokkos::RangePolicy<execution_space, typename Functor::template OverlapTag<0, B, false>> \
982 policy(0, rowidx2part.extent(0)); \
983 Kokkos::parallel_for("ComputeResidual::RangePolicy::run<OverlapTag<0> >", policy, functor); \
984 } else { \
985 const Kokkos::RangePolicy<execution_space, typename Functor::template OverlapTag<1, B, false>> \
986 policy(0, rowidx2part.extent(0)); \
987 Kokkos::parallel_for("ComputeResidual::RangePolicy::run<OverlapTag<1> >", policy, functor); \
988 } \
989 } \
990 break
991
992 switch (blocksize_requested) {
993 case 3: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(3);
994 case 5: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(5);
995 case 7: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(7);
996 case 9: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(9);
997 case 10: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(10);
998 case 11: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(11);
999 case 16: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(16);
1000 case 17: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(17);
1001 case 18: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(18);
1002 default: BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL(0);
1003 }
1004#undef BLOCKTRIDICONTAINER_DETAILS_COMPUTERESIDUAL
1005 }
1006 IFPACK2_BLOCKHELPER_PROFILER_REGION_END;
1007 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
1008}
1009
1010} // namespace Ifpack2::BlockHelperDetails
1011
1012#define IFPACK2_BLOCKCOMPUTERESIDUALVECTOR_INSTANT(S, LO, GO, N) \
1013 template class Ifpack2::BlockHelperDetails::ComputeResidualVector<Tpetra::RowMatrix<S, LO, GO, N>>;
1014
1015#endif
Ifpack2::BlockHelperDetails::ImplType::size_type
size_t size_type
Definition Ifpack2_BlockHelper.hpp:278
Ifpack2::BlockHelperDetails::ImplType::node_device_type
node_type::device_type node_device_type
Definition Ifpack2_BlockHelper.hpp:302
Ifpack2::BlockHelperDetails::ImplType::impl_scalar_type
KokkosKernels::ArithTraits< scalar_type >::val_type impl_scalar_type
Definition Ifpack2_BlockHelper.hpp:288
Ifpack2::BlockHelperDetails::ImplType::size_type_1d_view
Kokkos::View< size_type *, device_type > size_type_1d_view
Definition Ifpack2_BlockHelper.hpp:351
Generated by doxygen 1.9.8