10 #ifndef IFPACK2_BLOCKTRIDICONTAINER_IMPL_HPP
11 #define IFPACK2_BLOCKTRIDICONTAINER_IMPL_HPP
18 #include <Tpetra_Details_extractMpiCommFromTeuchos.hpp>
19 #include <Tpetra_Distributor.hpp>
20 #include <Tpetra_BlockMultiVector.hpp>
22 #include <Kokkos_ArithTraits.hpp>
23 #include <KokkosBatched_Util.hpp>
24 #include <KokkosBatched_Vector.hpp>
25 #include <KokkosBatched_Copy_Decl.hpp>
26 #include <KokkosBatched_Copy_Impl.hpp>
27 #include <KokkosBatched_AddRadial_Decl.hpp>
28 #include <KokkosBatched_AddRadial_Impl.hpp>
29 #include <KokkosBatched_SetIdentity_Decl.hpp>
30 #include <KokkosBatched_SetIdentity_Impl.hpp>
31 #include <KokkosBatched_Gemm_Decl.hpp>
32 #include <KokkosBatched_Gemm_Serial_Impl.hpp>
33 #include <KokkosBatched_Gemm_Team_Impl.hpp>
34 #include <KokkosBatched_Gemv_Decl.hpp>
35 #include <KokkosBatched_Gemv_Team_Impl.hpp>
36 #include <KokkosBatched_Trsm_Decl.hpp>
37 #include <KokkosBatched_Trsm_Serial_Impl.hpp>
38 #include <KokkosBatched_Trsm_Team_Impl.hpp>
39 #include <KokkosBatched_Trsv_Decl.hpp>
40 #include <KokkosBatched_Trsv_Serial_Impl.hpp>
41 #include <KokkosBatched_Trsv_Team_Impl.hpp>
42 #include <KokkosBatched_LU_Decl.hpp>
43 #include <KokkosBatched_LU_Serial_Impl.hpp>
44 #include <KokkosBatched_LU_Team_Impl.hpp>
46 #include <KokkosBlas1_nrm1.hpp>
47 #include <KokkosBlas1_nrm2.hpp>
51 #include "Ifpack2_BlockHelper.hpp"
52 #include "Ifpack2_BlockComputeResidualVector.hpp"
53 #include "Ifpack2_BlockComputeResidualAndSolve.hpp"
60 #if defined(KOKKOS_ENABLE_CUDA) && defined(IFPACK2_BLOCKTRIDICONTAINER_ENABLE_PROFILE)
61 #include "cuda_profiler_api.h"
66 #define IFPACK2_BLOCKTRIDICONTAINER_USE_MPI_3
74 #define IFPACK2_BLOCKTRIDICONTAINER_USE_PINNED_MEMORY_FOR_MPI
78 #define IFPACK2_BLOCKTRIDICONTAINER_USE_CUDA_SPACE
81 #if defined(HAVE_IFPACK2_BLOCKTRIDICONTAINER_SMALL_SCALAR)
82 #define IFPACK2_BLOCKTRIDICONTAINER_USE_SMALL_SCALAR_FOR_BLOCKTRIDIAG
86 #define IFPACK2_BLOCKTRIDICONTAINER_USE_EXEC_SPACE_INSTANCES
90 namespace BlockTriDiContainerDetails {
92 namespace KB = KokkosBatched;
99 template <
typename MemoryTraitsType, Kokkos::MemoryTraitsFlags flag>
100 using MemoryTraits = Kokkos::MemoryTraits<MemoryTraitsType::is_unmanaged |
101 MemoryTraitsType::is_random_access |
104 template <
typename ViewType>
105 using Unmanaged = Kokkos::View<
typename ViewType::data_type,
106 typename ViewType::array_layout,
107 typename ViewType::device_type,
108 MemoryTraits<typename ViewType::memory_traits,Kokkos::Unmanaged> >;
109 template <
typename ViewType>
110 using Atomic = Kokkos::View<
typename ViewType::data_type,
111 typename ViewType::array_layout,
112 typename ViewType::device_type,
113 MemoryTraits<typename ViewType::memory_traits,Kokkos::Atomic> >;
114 template <
typename ViewType>
115 using Const = Kokkos::View<
typename ViewType::const_data_type,
116 typename ViewType::array_layout,
117 typename ViewType::device_type,
118 typename ViewType::memory_traits>;
119 template <
typename ViewType>
120 using ConstUnmanaged = Const<Unmanaged<ViewType> >;
122 template <
typename ViewType>
123 using AtomicUnmanaged = Atomic<Unmanaged<ViewType> >;
125 template <
typename ViewType>
126 using Unmanaged = Kokkos::View<
typename ViewType::data_type,
127 typename ViewType::array_layout,
128 typename ViewType::device_type,
129 MemoryTraits<typename ViewType::memory_traits,Kokkos::Unmanaged> >;
132 template <
typename ViewType>
133 using Scratch = Kokkos::View<
typename ViewType::data_type,
134 typename ViewType::array_layout,
135 typename ViewType::execution_space::scratch_memory_space,
136 MemoryTraits<typename ViewType::memory_traits, Kokkos::Unmanaged> >;
142 #if defined(IFPACK2_BLOCKTRIDICONTAINER_USE_SMALL_SCALAR_FOR_BLOCKTRIDIAG)
147 #if defined(KOKKOS_ENABLE_CUDA) && defined(IFPACK2_BLOCKTRIDICONTAINER_ENABLE_PROFILE)
148 #define IFPACK2_BLOCKTRIDICONTAINER_PROFILER_REGION_BEGIN \
149 KOKKOS_IMPL_CUDA_SAFE_CALL(cudaProfilerStart());
151 #define IFPACK2_BLOCKTRIDICONTAINER_PROFILER_REGION_END \
152 { KOKKOS_IMPL_CUDA_SAFE_CALL( cudaProfilerStop() ); }
154 #define IFPACK2_BLOCKTRIDICONTAINER_PROFILER_REGION_BEGIN
156 #define IFPACK2_BLOCKTRIDICONTAINER_PROFILER_REGION_END
162 template<
typename MatrixType>
165 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::CreateBlockCrsTpetraImporter", CreateBlockCrsTpetraImporter);
167 using tpetra_map_type =
typename impl_type::tpetra_map_type;
168 using tpetra_mv_type =
typename impl_type::tpetra_block_multivector_type;
169 using tpetra_import_type =
typename impl_type::tpetra_import_type;
170 using crs_matrix_type =
typename impl_type::tpetra_crs_matrix_type;
171 using block_crs_matrix_type =
typename impl_type::tpetra_block_crs_matrix_type;
173 auto A_crs = Teuchos::rcp_dynamic_cast<
const crs_matrix_type>(A);
174 auto A_bcrs = Teuchos::rcp_dynamic_cast<
const block_crs_matrix_type>(A);
176 bool hasBlockCrsMatrix = ! A_bcrs.is_null ();
179 const auto g = hasBlockCrsMatrix ? A_bcrs->getCrsGraph() : *(A_crs->getCrsGraph());
181 const auto blocksize = hasBlockCrsMatrix ? A_bcrs->getBlockSize() : 1;
182 const auto src =
Teuchos::rcp(
new tpetra_map_type(tpetra_mv_type::makePointMap(*g.getDomainMap(), blocksize)));
183 const auto tgt =
Teuchos::rcp(
new tpetra_map_type(tpetra_mv_type::makePointMap(*g.getColMap() , blocksize)));
184 IFPACK2_BLOCKHELPER_TIMER_FENCE(
typename BlockHelperDetails::ImplType<MatrixType>::execution_space)
193 template<
typename MatrixType>
194 struct AsyncableImport {
202 #if !defined(HAVE_IFPACK2_MPI)
203 typedef int MPI_Request;
204 typedef int MPI_Comm;
206 using scalar_type =
typename impl_type::scalar_type;
210 static int isend(
const MPI_Comm comm,
const char* buf,
int count,
int dest,
int tag, MPI_Request* ireq) {
211 #ifdef HAVE_IFPACK2_MPI
213 int ret = MPI_Isend(const_cast<char*>(buf), count, MPI_CHAR, dest, tag, comm, ireq == NULL ? &ureq : ireq);
214 if (ireq == NULL) MPI_Request_free(&ureq);
221 static int irecv(
const MPI_Comm comm,
char* buf,
int count,
int src,
int tag, MPI_Request* ireq) {
222 #ifdef HAVE_IFPACK2_MPI
224 int ret = MPI_Irecv(buf, count, MPI_CHAR, src, tag, comm, ireq == NULL ? &ureq : ireq);
225 if (ireq == NULL) MPI_Request_free(&ureq);
232 static int waitany(
int count, MPI_Request* reqs,
int* index) {
233 #ifdef HAVE_IFPACK2_MPI
234 return MPI_Waitany(count, reqs, index, MPI_STATUS_IGNORE);
240 static int waitall(
int count, MPI_Request* reqs) {
241 #ifdef HAVE_IFPACK2_MPI
242 return MPI_Waitall(count, reqs, MPI_STATUS_IGNORE);
249 using tpetra_map_type =
typename impl_type::tpetra_map_type;
250 using tpetra_import_type =
typename impl_type::tpetra_import_type;
252 using local_ordinal_type =
typename impl_type::local_ordinal_type;
253 using global_ordinal_type =
typename impl_type::global_ordinal_type;
257 using int_1d_view_host = Kokkos::View<int*,Kokkos::HostSpace>;
258 using local_ordinal_type_1d_view_host = Kokkos::View<local_ordinal_type*,Kokkos::HostSpace>;
260 using execution_space =
typename impl_type::execution_space;
261 using memory_space =
typename impl_type::memory_space;
262 using local_ordinal_type_1d_view =
typename impl_type::local_ordinal_type_1d_view;
264 using size_type_1d_view_host = Kokkos::View<size_type*,Kokkos::HostSpace>;
266 #if defined(KOKKOS_ENABLE_CUDA)
267 using impl_scalar_type_1d_view =
268 typename std::conditional<std::is_same<execution_space,Kokkos::Cuda>::value,
269 # if defined(IFPACK2_BLOCKTRIDICONTAINER_USE_PINNED_MEMORY_FOR_MPI)
270 Kokkos::View<impl_scalar_type*,Kokkos::CudaHostPinnedSpace>,
271 # elif defined(IFPACK2_BLOCKTRIDICONTAINER_USE_CUDA_MEMORY_FOR_MPI)
272 Kokkos::View<impl_scalar_type*,Kokkos::CudaSpace>,
273 # else // no experimental macros are defined
274 typename impl_type::impl_scalar_type_1d_view,
276 typename impl_type::impl_scalar_type_1d_view>::type;
278 using impl_scalar_type_1d_view =
typename impl_type::impl_scalar_type_1d_view;
280 using impl_scalar_type_1d_view_host = Kokkos::View<impl_scalar_type*,Kokkos::HostSpace>;
281 using impl_scalar_type_2d_view =
typename impl_type::impl_scalar_type_2d_view;
282 using impl_scalar_type_2d_view_tpetra =
typename impl_type::impl_scalar_type_2d_view_tpetra;
284 #ifdef HAVE_IFPACK2_MPI
288 impl_scalar_type_2d_view_tpetra remote_multivector;
289 local_ordinal_type blocksize;
292 struct SendRecvPair {
297 SendRecvPair<int_1d_view_host> pids;
298 SendRecvPair<std::vector<MPI_Request> > reqs;
299 SendRecvPair<size_type_1d_view> offset;
300 SendRecvPair<size_type_1d_view_host> offset_host;
301 SendRecvPair<local_ordinal_type_1d_view> lids;
302 SendRecvPair<impl_scalar_type_1d_view> buffer;
303 SendRecvPair<impl_scalar_type_1d_view_host> buffer_host;
305 local_ordinal_type_1d_view dm2cm;
307 #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) || defined(KOKKOS_ENABLE_SYCL)
308 using exec_instance_1d_std_vector = std::vector<execution_space>;
309 exec_instance_1d_std_vector exec_instances;
315 const size_type_1d_view &offs) {
317 Kokkos::View<size_t*,Kokkos::HostSpace> lens_host(const_cast<size_t*>(lens.
getRawPtr()), lens.
size());
318 const auto lens_device = Kokkos::create_mirror_view_and_copy(memory_space(), lens_host);
321 const Kokkos::RangePolicy<execution_space> policy(0,offs.extent(0));
322 const local_ordinal_type lens_size = lens_device.extent(0);
323 Kokkos::parallel_scan
324 (
"AsyncableImport::RangePolicy::setOffsetValues",
325 policy, KOKKOS_LAMBDA(
const local_ordinal_type &i, size_type &update,
const bool &
final) {
328 update += (i < lens_size ? lens_device[i] : 0);
333 const size_type_1d_view_host &offs) {
335 Kokkos::View<size_t*,Kokkos::HostSpace> lens_host(const_cast<size_t*>(lens.
getRawPtr()), lens.
size());
336 const auto lens_device = Kokkos::create_mirror_view_and_copy(memory_space(), lens_host);
340 for (local_ordinal_type i=1,iend=offs.extent(0);i<iend;++i) {
341 offs(i) = offs(i-1) + lens[i-1];
346 void createMpiRequests(
const tpetra_import_type &
import) {
347 Tpetra::Distributor &distributor =
import.getDistributor();
350 const auto pids_from = distributor.getProcsFrom();
352 memcpy(pids.recv.data(), pids_from.getRawPtr(),
sizeof(int)*pids.recv.extent(0));
354 const auto pids_to = distributor.getProcsTo();
356 memcpy(pids.send.data(), pids_to.getRawPtr(),
sizeof(int)*pids.send.extent(0));
359 reqs.recv.resize(pids.recv.extent(0)); memset(reqs.recv.data(), 0, reqs.recv.size()*
sizeof(MPI_Request));
360 reqs.send.resize(pids.send.extent(0)); memset(reqs.send.data(), 0, reqs.send.size()*
sizeof(MPI_Request));
364 const auto lengths_to = distributor.getLengthsTo();
367 const auto lengths_from = distributor.getLengthsFrom();
370 setOffsetValues(lengths_to, offset.send);
371 offset_host.send = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), offset.send);
373 setOffsetValues(lengths_from, offset.recv);
374 offset_host.recv = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), offset.recv);
376 const auto lengths_to = distributor.getLengthsTo();
377 offset_host.send = size_type_1d_view_host(
do_not_initialize_tag(
"offset send"), lengths_to.size() + 1);
379 const auto lengths_from = distributor.getLengthsFrom();
380 offset_host.recv = size_type_1d_view_host(
do_not_initialize_tag(
"offset recv"), lengths_from.size() + 1);
382 setOffsetValuesHost(lengths_to, offset_host.send);
385 setOffsetValuesHost(lengths_from, offset_host.recv);
390 void createSendRecvIDs(
const tpetra_import_type &
import) {
392 const auto remote_lids =
import.getRemoteLIDs();
393 const local_ordinal_type_1d_view_host
394 remote_lids_view_host(const_cast<local_ordinal_type*>(remote_lids.getRawPtr()), remote_lids.size());
396 Kokkos::deep_copy(lids.recv, remote_lids_view_host);
399 auto epids =
import.getExportPIDs();
400 auto elids =
import.getExportLIDs();
403 auto lids_send_host = Kokkos::create_mirror_view(lids.send);
406 for (local_ordinal_type cnt=0,i=0,iend=pids.send.extent(0);i<iend;++i) {
407 const auto pid_send_value = pids.send[i];
408 for (local_ordinal_type j=0,jend=epids.size();j<jend;++j)
409 if (epids[j] == pid_send_value) lids_send_host[cnt++] = elids[j];
410 TEUCHOS_ASSERT(static_cast<size_t>(cnt) == offset_host.send[i+1]);
412 Kokkos::deep_copy(lids.send, lids_send_host);
415 void createExecutionSpaceInstances() {
416 #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) || defined(KOKKOS_ENABLE_SYCL)
419 Kokkos::Experimental::partition_space(execution_space(), 1, 1, 1, 1, 1, 1, 1, 1);
426 struct ToMultiVector {};
430 const local_ordinal_type blocksize_,
431 const local_ordinal_type_1d_view dm2cm_) {
432 blocksize = blocksize_;
435 #ifdef HAVE_IFPACK2_MPI
436 comm = Tpetra::Details::extractMpiCommFromTeuchos(*tgt_map->getComm());
438 const tpetra_import_type
import(src_map, tgt_map);
440 createMpiRequests(
import);
441 createSendRecvIDs(
import);
442 createExecutionSpaceInstances();
445 void createDataBuffer(
const local_ordinal_type &num_vectors) {
446 const size_type extent_0 = lids.recv.extent(0)*blocksize;
447 const size_type extent_1 = num_vectors;
448 if (remote_multivector.extent(0) == extent_0 &&
449 remote_multivector.extent(1) == extent_1) {
455 const auto send_buffer_size = offset_host.send[offset_host.send.extent(0)-1]*blocksize*num_vectors;
456 const auto recv_buffer_size = offset_host.recv[offset_host.recv.extent(0)-1]*blocksize*num_vectors;
461 if (!Tpetra::Details::Behavior::assumeMpiIsGPUAware()) {
462 buffer_host.send = impl_scalar_type_1d_view_host(
do_not_initialize_tag(
"buffer send"), send_buffer_size);
463 buffer_host.recv = impl_scalar_type_1d_view_host(
do_not_initialize_tag(
"buffer recv"), recv_buffer_size);
469 #ifdef HAVE_IFPACK2_MPI
470 waitall(reqs.recv.size(), reqs.recv.data());
471 waitall(reqs.send.size(), reqs.send.data());
479 #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) || defined(KOKKOS_ENABLE_SYCL)
480 template<
typename PackTag>
482 void copy(
const local_ordinal_type_1d_view &lids_,
483 const impl_scalar_type_1d_view &buffer_,
484 const local_ordinal_type ibeg_,
485 const local_ordinal_type iend_,
486 const impl_scalar_type_2d_view_tpetra &multivector_,
487 const local_ordinal_type blocksize_,
488 const execution_space &exec_instance_) {
489 const local_ordinal_type num_vectors = multivector_.extent(1);
490 const local_ordinal_type mv_blocksize = blocksize_*num_vectors;
491 const local_ordinal_type idiff = iend_ - ibeg_;
492 const auto abase = buffer_.data() + mv_blocksize*ibeg_;
494 using team_policy_type = Kokkos::TeamPolicy<execution_space>;
495 local_ordinal_type vector_size(0);
496 if (blocksize_ <= 4) vector_size = 4;
497 else if (blocksize_ <= 8) vector_size = 8;
498 else if (blocksize_ <= 16) vector_size = 16;
499 else vector_size = 32;
501 const auto work_item_property = Kokkos::Experimental::WorkItemProperty::HintLightWeight;
502 const team_policy_type policy(exec_instance_, idiff, 1, vector_size);
505 Kokkos::Experimental::require(policy, work_item_property),
506 KOKKOS_LAMBDA(
const typename team_policy_type::member_type &member) {
507 const local_ordinal_type i = member.league_rank();
509 (Kokkos::TeamThreadRange(member,num_vectors),[&](
const local_ordinal_type &j) {
510 auto aptr = abase + blocksize_*(i + idiff*j);
511 auto bptr = &multivector_(blocksize_*lids_(i + ibeg_), j);
512 if (std::is_same<PackTag,ToBuffer>::value)
514 (Kokkos::ThreadVectorRange(member,blocksize_),[&](
const local_ordinal_type &k) {
519 (Kokkos::ThreadVectorRange(member,blocksize_),[&](
const local_ordinal_type &k) {
526 void asyncSendRecvVar1(
const impl_scalar_type_2d_view_tpetra &mv) {
527 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::AsyncableImport::AsyncSendRecv", AsyncSendRecv);
529 #ifdef HAVE_IFPACK2_MPI
531 const local_ordinal_type num_vectors = mv.extent(1);
532 const local_ordinal_type mv_blocksize = blocksize*num_vectors;
535 for (local_ordinal_type i=0,iend=pids.recv.extent(0);i<iend;++i) {
536 if(Tpetra::Details::Behavior::assumeMpiIsGPUAware()) {
538 reinterpret_cast<char*>(buffer.recv.data() + offset_host.recv[i]*mv_blocksize),
539 (offset_host.recv[i+1] - offset_host.recv[i])*mv_blocksize*
sizeof(impl_scalar_type),
546 reinterpret_cast<char*>(buffer_host.recv.data() + offset_host.recv[i]*mv_blocksize),
547 (offset_host.recv[i+1] - offset_host.recv[i])*mv_blocksize*
sizeof(impl_scalar_type),
555 execution_space().fence();
558 for (local_ordinal_type i=0;i<static_cast<local_ordinal_type>(pids.send.extent(0));++i) {
560 if (i<8) exec_instances[i%8].fence();
561 copy<ToBuffer>(lids.send, buffer.send,
562 offset_host.send(i), offset_host.send(i+1),
565 exec_instances[i%8]);
566 if (!Tpetra::Details::Behavior::assumeMpiIsGPUAware()) {
568 const local_ordinal_type num_vectors = mv.extent(1);
569 const local_ordinal_type mv_blocksize = blocksize*num_vectors;
571 Kokkos::deep_copy(exec_instances[i%8],
572 Kokkos::subview(buffer_host.send,
573 Kokkos::pair<local_ordinal_type, local_ordinal_type>(
574 offset_host.send(i)*mv_blocksize,
575 offset_host.send(i+1)*mv_blocksize)),
576 Kokkos::subview(buffer.send,
577 Kokkos::pair<local_ordinal_type, local_ordinal_type>(
578 offset_host.send(i)*mv_blocksize,
579 offset_host.send(i+1)*mv_blocksize)));
584 for (local_ordinal_type i=0;i<static_cast<local_ordinal_type>(pids.send.extent(0));++i) {
586 if (i<8) exec_instances[i%8].fence();
587 if(Tpetra::Details::Behavior::assumeMpiIsGPUAware()) {
589 reinterpret_cast<const char*>(buffer.send.data() + offset_host.send[i]*mv_blocksize),
590 (offset_host.send[i+1] - offset_host.send[i])*mv_blocksize*
sizeof(impl_scalar_type),
597 reinterpret_cast<const char*>(buffer_host.send.data() + offset_host.send[i]*mv_blocksize),
598 (offset_host.send[i+1] - offset_host.send[i])*mv_blocksize*
sizeof(impl_scalar_type),
606 for (local_ordinal_type i=0,iend=pids.recv.extent(0);i<iend;++i) {
609 MPI_Iprobe(pids.recv[i], 42, comm, &flag, &stat);
611 #endif // HAVE_IFPACK2_MPI
612 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
615 void syncRecvVar1() {
616 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::AsyncableImport::SyncRecv", SyncRecv);
617 #ifdef HAVE_IFPACK2_MPI
619 for (local_ordinal_type i=0;i<static_cast<local_ordinal_type>(pids.recv.extent(0));++i) {
620 local_ordinal_type idx = i;
623 waitany(pids.recv.extent(0), reqs.recv.data(), &idx);
625 if (!Tpetra::Details::Behavior::assumeMpiIsGPUAware()) {
626 const local_ordinal_type num_vectors = remote_multivector.extent(1);
627 const local_ordinal_type mv_blocksize = blocksize*num_vectors;
630 Kokkos::subview(buffer.recv,
631 Kokkos::pair<local_ordinal_type, local_ordinal_type>(
632 offset_host.recv(idx)*mv_blocksize,
633 offset_host.recv(idx+1)*mv_blocksize)),
634 Kokkos::subview(buffer_host.recv,
635 Kokkos::pair<local_ordinal_type, local_ordinal_type>(
636 offset_host.recv(idx)*mv_blocksize,
637 offset_host.recv(idx+1)*mv_blocksize)));
641 copy<ToMultiVector>(lids.recv, buffer.recv,
642 offset_host.recv(idx), offset_host.recv(idx+1),
643 remote_multivector, blocksize,
644 exec_instances[idx%8]);
651 waitall(reqs.send.size(), reqs.send.data());
652 #endif // HAVE_IFPACK2_MPI
653 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
655 #endif //defined(KOKKOS_ENABLE_CUDA|HIP|SYCL)
662 template<
typename PackTag>
664 void copy(
const local_ordinal_type_1d_view &lids_,
665 const impl_scalar_type_1d_view &buffer_,
666 const local_ordinal_type &ibeg_,
667 const local_ordinal_type &iend_,
668 const impl_scalar_type_2d_view_tpetra &multivector_,
669 const local_ordinal_type blocksize_) {
670 const local_ordinal_type num_vectors = multivector_.extent(1);
671 const local_ordinal_type mv_blocksize = blocksize_*num_vectors;
672 const local_ordinal_type idiff = iend_ - ibeg_;
673 const auto abase = buffer_.data() + mv_blocksize*ibeg_;
674 if constexpr (BlockHelperDetails::is_device<execution_space>::value) {
675 using team_policy_type = Kokkos::TeamPolicy<execution_space>;
676 local_ordinal_type vector_size(0);
677 if (blocksize_ <= 4) vector_size = 4;
678 else if (blocksize_ <= 8) vector_size = 8;
679 else if (blocksize_ <= 16) vector_size = 16;
680 else vector_size = 32;
681 const team_policy_type policy(idiff, 1, vector_size);
683 (
"AsyncableImport::TeamPolicy::copy",
684 policy, KOKKOS_LAMBDA(
const typename team_policy_type::member_type &member) {
685 const local_ordinal_type i = member.league_rank();
687 (Kokkos::TeamThreadRange(member,num_vectors),[&](
const local_ordinal_type &j) {
688 auto aptr = abase + blocksize_*(i + idiff*j);
689 auto bptr = &multivector_(blocksize_*lids_(i + ibeg_), j);
690 if (std::is_same<PackTag,ToBuffer>::value)
692 (Kokkos::ThreadVectorRange(member,blocksize_),[&](
const local_ordinal_type &k) {
697 (Kokkos::ThreadVectorRange(member,blocksize_),[&](
const local_ordinal_type &k) {
703 const Kokkos::RangePolicy<execution_space> policy(0, idiff*num_vectors);
705 (
"AsyncableImport::RangePolicy::copy",
706 policy, KOKKOS_LAMBDA(
const local_ordinal_type &ij) {
707 const local_ordinal_type i = ij%idiff;
708 const local_ordinal_type j = ij/idiff;
709 auto aptr = abase + blocksize_*(i + idiff*j);
710 auto bptr = &multivector_(blocksize_*lids_(i + ibeg_), j);
711 auto from = std::is_same<PackTag,ToBuffer>::value ? bptr : aptr;
712 auto to = std::is_same<PackTag,ToBuffer>::value ? aptr : bptr;
713 memcpy(to, from,
sizeof(impl_scalar_type)*blocksize_);
722 void asyncSendRecvVar0(
const impl_scalar_type_2d_view_tpetra &mv) {
723 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::AsyncableImport::AsyncSendRecv", AsyncSendRecv);
725 #ifdef HAVE_IFPACK2_MPI
727 const local_ordinal_type num_vectors = mv.extent(1);
728 const local_ordinal_type mv_blocksize = blocksize*num_vectors;
731 for (local_ordinal_type i=0,iend=pids.recv.extent(0);i<iend;++i) {
732 if(Tpetra::Details::Behavior::assumeMpiIsGPUAware()) {
734 reinterpret_cast<char*>(buffer.recv.data() + offset_host.recv[i]*mv_blocksize),
735 (offset_host.recv[i+1] - offset_host.recv[i])*mv_blocksize*
sizeof(impl_scalar_type),
742 reinterpret_cast<char*>(buffer_host.recv.data() + offset_host.recv[i]*mv_blocksize),
743 (offset_host.recv[i+1] - offset_host.recv[i])*mv_blocksize*
sizeof(impl_scalar_type),
751 for (local_ordinal_type i=0,iend=pids.send.extent(0);i<iend;++i) {
752 copy<ToBuffer>(lids.send, buffer.send, offset_host.send(i), offset_host.send(i+1),
755 if(Tpetra::Details::Behavior::assumeMpiIsGPUAware()) {
757 reinterpret_cast<const char*>(buffer.send.data() + offset_host.send[i]*mv_blocksize),
758 (offset_host.send[i+1] - offset_host.send[i])*mv_blocksize*
sizeof(impl_scalar_type),
765 Kokkos::subview(buffer_host.send,
766 Kokkos::pair<local_ordinal_type, local_ordinal_type>(
767 offset_host.send(i)*mv_blocksize,
768 offset_host.send(i+1)*mv_blocksize)),
769 Kokkos::subview(buffer.send,
770 Kokkos::pair<local_ordinal_type, local_ordinal_type>(
771 offset_host.send(i)*mv_blocksize,
772 offset_host.send(i+1)*mv_blocksize)));
774 reinterpret_cast<const char*>(buffer_host.send.data() + offset_host.send[i]*mv_blocksize),
775 (offset_host.send[i+1] - offset_host.send[i])*mv_blocksize*
sizeof(impl_scalar_type),
784 for (local_ordinal_type i=0,iend=pids.recv.extent(0);i<iend;++i) {
787 MPI_Iprobe(pids.recv[i], 42, comm, &flag, &stat);
790 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
793 void syncRecvVar0() {
794 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::AsyncableImport::SyncRecv", SyncRecv);
795 #ifdef HAVE_IFPACK2_MPI
797 for (local_ordinal_type i=0,iend=pids.recv.extent(0);i<iend;++i) {
798 local_ordinal_type idx = i;
799 waitany(pids.recv.extent(0), reqs.recv.data(), &idx);
800 if (!Tpetra::Details::Behavior::assumeMpiIsGPUAware()) {
801 const local_ordinal_type num_vectors = remote_multivector.extent(1);
802 const local_ordinal_type mv_blocksize = blocksize*num_vectors;
804 Kokkos::subview(buffer.recv,
805 Kokkos::pair<local_ordinal_type, local_ordinal_type>(
806 offset_host.recv(idx)*mv_blocksize,
807 offset_host.recv(idx+1)*mv_blocksize)),
808 Kokkos::subview(buffer_host.recv,
809 Kokkos::pair<local_ordinal_type, local_ordinal_type>(
810 offset_host.recv(idx)*mv_blocksize,
811 offset_host.recv(idx+1)*mv_blocksize)));
813 copy<ToMultiVector>(lids.recv, buffer.recv, offset_host.recv(idx), offset_host.recv(idx+1),
814 remote_multivector, blocksize);
817 waitall(reqs.send.size(), reqs.send.data());
819 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
825 void asyncSendRecv(
const impl_scalar_type_2d_view_tpetra &mv) {
826 #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) || defined(KOKKOS_ENABLE_SYCL)
827 #if defined(IFPACK2_BLOCKTRIDICONTAINER_USE_EXEC_SPACE_INSTANCES)
828 asyncSendRecvVar1(mv);
830 asyncSendRecvVar0(mv);
833 asyncSendRecvVar0(mv);
837 #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) || defined(KOKKOS_ENABLE_SYCL)
838 #if defined(IFPACK2_BLOCKTRIDICONTAINER_USE_EXEC_SPACE_INSTANCES)
848 void syncExchange(
const impl_scalar_type_2d_view_tpetra &mv) {
849 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::AsyncableImport::SyncExchange", SyncExchange);
852 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
855 impl_scalar_type_2d_view_tpetra getRemoteMultiVectorLocalView()
const {
return remote_multivector; }
858 template <
typename ViewType1,
typename ViewType2>
859 struct are_same_struct {
863 are_same_struct(ViewType1 keys1_, ViewType2 keys2_) : keys1(keys1_), keys2(keys2_) {}
864 KOKKOS_INLINE_FUNCTION
865 void operator()(
int i,
unsigned int& count)
const {
866 if (keys1(i) != keys2(i)) count++;
870 template <
typename ViewType1,
typename ViewType2>
871 bool are_same (ViewType1 keys1, ViewType2 keys2) {
872 unsigned int are_same_ = 0;
874 Kokkos::parallel_reduce(Kokkos::RangePolicy<typename ViewType1::execution_space>(0, keys1.extent(0)),
875 are_same_struct(keys1, keys2),
883 template<
typename MatrixType>
888 using tpetra_map_type =
typename impl_type::tpetra_map_type;
889 using local_ordinal_type =
typename impl_type::local_ordinal_type;
890 using global_ordinal_type =
typename impl_type::global_ordinal_type;
891 using local_ordinal_type_1d_view =
typename impl_type::local_ordinal_type_1d_view;
892 using crs_matrix_type =
typename impl_type::tpetra_crs_matrix_type;
893 using block_crs_matrix_type =
typename impl_type::tpetra_block_crs_matrix_type;
894 using global_indices_array_device_type = Kokkos::View<const global_ordinal_type*, typename tpetra_map_type::device_type>;
896 auto A_crs = Teuchos::rcp_dynamic_cast<
const crs_matrix_type>(A);
897 auto A_bcrs = Teuchos::rcp_dynamic_cast<
const block_crs_matrix_type>(A);
899 bool hasBlockCrsMatrix = ! A_bcrs.is_null ();
902 const auto g = hasBlockCrsMatrix ? A_bcrs->getCrsGraph() : *(A_crs->getCrsGraph());
904 const auto blocksize = hasBlockCrsMatrix ? A_bcrs->getBlockSize() : 1;
905 const auto domain_map = g.getDomainMap();
906 const auto column_map = g.getColMap();
908 std::vector<global_ordinal_type> gids;
910 Kokkos::Subview<global_indices_array_device_type, std::pair<int,int>> column_map_global_iD_last;
912 bool separate_remotes =
true, found_first =
false, need_owned_permutation =
false;
914 IFPACK2_BLOCKHELPER_TIMER(
"createBlockCrsAsyncImporter::loop_over_local_elements", loop_over_local_elements);
916 global_indices_array_device_type column_map_global_iD = column_map->getMyGlobalIndicesDevice();
917 global_indices_array_device_type domain_map_global_iD = domain_map->getMyGlobalIndicesDevice();
919 if(are_same(domain_map_global_iD, column_map_global_iD)) {
921 separate_remotes =
true;
922 need_owned_permutation =
false;
924 column_map_global_iD_last = Kokkos::subview(column_map_global_iD,
925 std::pair<int,int>(domain_map_global_iD.extent(0), column_map_global_iD.extent(0)));
929 for (
size_t i=0;i<column_map->getLocalNumElements();++i) {
930 const global_ordinal_type gid = column_map->getGlobalElement(i);
931 if (!domain_map->isNodeGlobalElement(gid)) {
934 }
else if (found_first) {
935 separate_remotes =
false;
938 if (!found_first && !need_owned_permutation &&
939 domain_map->getLocalElement(gid) !=
static_cast<local_ordinal_type
>(i)) {
948 need_owned_permutation =
true;
952 IFPACK2_BLOCKHELPER_TIMER_FENCE(
typename BlockHelperDetails::ImplType<MatrixType>::execution_space)
955 if (separate_remotes) {
956 IFPACK2_BLOCKHELPER_TIMER(
"createBlockCrsAsyncImporter::separate_remotes", separate_remotes);
958 const auto parsimonious_col_map
959 = need_owned_permutation ?
960 Teuchos::rcp(
new tpetra_map_type(invalid, gids.data(), gids.size(), 0, domain_map->getComm())):
961 Teuchos::rcp(
new tpetra_map_type(invalid, column_map_global_iD_last, 0, domain_map->getComm()));
962 if (parsimonious_col_map->getGlobalNumElements() > 0) {
964 local_ordinal_type_1d_view dm2cm;
965 if (need_owned_permutation) {
966 dm2cm = local_ordinal_type_1d_view(
do_not_initialize_tag(
"dm2cm"), domain_map->getLocalNumElements());
967 const auto dm2cm_host = Kokkos::create_mirror_view(dm2cm);
968 for (
size_t i=0;i<domain_map->getLocalNumElements();++i)
969 dm2cm_host(i) = domain_map->getLocalElement(column_map->getGlobalElement(i));
970 Kokkos::deep_copy(dm2cm, dm2cm_host);
972 IFPACK2_BLOCKHELPER_TIMER_FENCE(
typename BlockHelperDetails::ImplType<MatrixType>::execution_space)
973 return Teuchos::rcp(
new AsyncableImport<MatrixType>(domain_map, parsimonious_col_map, blocksize, dm2cm));
976 IFPACK2_BLOCKHELPER_TIMER_FENCE(
typename BlockHelperDetails::ImplType<MatrixType>::execution_space)
977 return Teuchos::null;
980 template<
typename local_ordinal_type>
981 local_ordinal_type costTRSM(
const local_ordinal_type block_size) {
982 return block_size*block_size;
985 template<
typename local_ordinal_type>
986 local_ordinal_type costGEMV(
const local_ordinal_type block_size) {
987 return 2*block_size*block_size;
990 template<
typename local_ordinal_type>
991 local_ordinal_type costTriDiagSolve(
const local_ordinal_type subline_length,
const local_ordinal_type block_size) {
992 return 2 * subline_length * costTRSM(block_size) + 2 * (subline_length-1) * costGEMV(block_size);
995 template<
typename local_ordinal_type>
996 local_ordinal_type costSolveSchur(
const local_ordinal_type num_parts,
997 const local_ordinal_type num_teams,
998 const local_ordinal_type line_length,
999 const local_ordinal_type block_size,
1000 const local_ordinal_type n_subparts_per_part) {
1001 const local_ordinal_type subline_length = ceil(
double(line_length - (n_subparts_per_part-1) * 2) / n_subparts_per_part);
1002 if (subline_length < 1) {
1006 const local_ordinal_type p_n_lines = ceil(
double(num_parts)/num_teams);
1007 const local_ordinal_type p_n_sublines = ceil(
double(n_subparts_per_part)*num_parts/num_teams);
1008 const local_ordinal_type p_n_sublines_2 = ceil(
double(n_subparts_per_part-1)*num_parts/num_teams);
1010 const local_ordinal_type p_costApplyE = p_n_sublines_2 * subline_length * 2 * costGEMV(block_size);
1011 const local_ordinal_type p_costApplyS = p_n_lines * costTriDiagSolve((n_subparts_per_part-1)*2,block_size);
1012 const local_ordinal_type p_costApplyAinv = p_n_sublines * costTriDiagSolve(subline_length,block_size);
1013 const local_ordinal_type p_costApplyC = p_n_sublines_2 * 2 * costGEMV(block_size);
1015 if (n_subparts_per_part == 1) {
1016 return p_costApplyAinv;
1018 return p_costApplyE + p_costApplyS + p_costApplyAinv + p_costApplyC;
1021 template<
typename local_ordinal_type>
1022 local_ordinal_type getAutomaticNSubparts(
const local_ordinal_type num_parts,
1023 const local_ordinal_type num_teams,
1024 const local_ordinal_type line_length,
1025 const local_ordinal_type block_size) {
1026 local_ordinal_type n_subparts_per_part_0 = 1;
1027 local_ordinal_type flop_0 = costSolveSchur(num_parts, num_teams, line_length, block_size, n_subparts_per_part_0);
1028 local_ordinal_type flop_1 = costSolveSchur(num_parts, num_teams, line_length, block_size, n_subparts_per_part_0+1);
1029 while (flop_0 > flop_1) {
1031 flop_1 = costSolveSchur(num_parts, num_teams, line_length, block_size, (++n_subparts_per_part_0)+1);
1033 return n_subparts_per_part_0;
1036 template<
typename ArgActiveExecutionMemorySpace>
1037 struct SolveTridiagsDefaultModeAndAlgo;
1042 template<
typename MatrixType>
1043 BlockHelperDetails::PartInterface<MatrixType>
1045 const Teuchos::RCP<
const typename BlockHelperDetails::ImplType<MatrixType>::tpetra_crs_graph_type> &G,
1047 const typename BlockHelperDetails::ImplType<MatrixType>::local_ordinal_type n_subparts_per_part_in) {
1050 using local_ordinal_type =
typename impl_type::local_ordinal_type;
1051 using local_ordinal_type_1d_view =
typename impl_type::local_ordinal_type_1d_view;
1052 using local_ordinal_type_2d_view =
typename impl_type::local_ordinal_type_2d_view;
1053 using size_type =
typename impl_type::size_type;
1055 auto bA = Teuchos::rcp_dynamic_cast<
const typename BlockHelperDetails::ImplType<MatrixType>::tpetra_block_crs_matrix_type>(A);
1058 const local_ordinal_type blocksize = bA.is_null() ? A->getLocalNumRows() / G->getLocalNumRows() : A->getBlockSize();
1059 constexpr
int vector_length = impl_type::vector_length;
1060 constexpr
int internal_vector_length = impl_type::internal_vector_length;
1062 const auto comm = A->getRowMap()->getComm();
1064 BlockHelperDetails::PartInterface<MatrixType> interf;
1066 const bool jacobi = partitions.size() == 0;
1067 const local_ordinal_type A_n_lclrows = G->getLocalNumRows();
1068 const local_ordinal_type nparts = jacobi ? A_n_lclrows : partitions.size();
1070 typedef std::pair<local_ordinal_type,local_ordinal_type> size_idx_pair_type;
1071 std::vector<size_idx_pair_type> partsz(nparts);
1074 for (local_ordinal_type i=0;i<nparts;++i)
1075 partsz[i] = size_idx_pair_type(partitions[i].size(), i);
1076 std::sort(partsz.begin(), partsz.end(),
1077 [] (
const size_idx_pair_type& x,
const size_idx_pair_type& y) {
1078 return x.first > y.first;
1082 local_ordinal_type n_subparts_per_part;
1083 if (n_subparts_per_part_in == -1) {
1086 using execution_space =
typename impl_type::execution_space;
1088 const int line_length = partsz[0].first;
1090 const local_ordinal_type team_size =
1091 SolveTridiagsDefaultModeAndAlgo<typename execution_space::memory_space>::
1092 recommended_team_size(blocksize, vector_length, internal_vector_length);
1094 const local_ordinal_type num_teams = std::max(1, execution_space().concurrency() / (team_size * vector_length));
1096 n_subparts_per_part = getAutomaticNSubparts(nparts, num_teams, line_length, blocksize);
1098 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
1099 printf(
"Automatically chosen n_subparts_per_part = %d for nparts = %d, num_teams = %d, team_size = %d, line_length = %d, and blocksize = %d;\n", n_subparts_per_part, nparts, num_teams, team_size, line_length, blocksize);
1103 n_subparts_per_part = n_subparts_per_part_in;
1107 const local_ordinal_type n_sub_parts = nparts * n_subparts_per_part;
1110 const local_ordinal_type n_sub_parts_and_schur = n_sub_parts + nparts * (n_subparts_per_part-1);
1112 #if defined(BLOCKTRIDICONTAINER_DEBUG)
1113 local_ordinal_type nrows = 0;
1117 for (local_ordinal_type i=0;i<nparts;++i) nrows += partitions[i].size();
1120 (nrows != A_n_lclrows, BlockHelperDetails::get_msg_prefix(comm) <<
"The #rows implied by the local partition is not "
1121 <<
"the same as getLocalNumRows: " << nrows <<
" vs " << A_n_lclrows);
1125 std::vector<local_ordinal_type> p;
1127 interf.max_partsz = 1;
1128 interf.max_subpartsz = 0;
1129 interf.n_subparts_per_part = 1;
1130 interf.nparts = nparts;
1135 for (local_ordinal_type i=0;i<nparts;++i)
1136 p[i] = partsz[i].second;
1138 interf.max_partsz = partsz[0].first;
1140 constexpr local_ordinal_type connection_length = 2;
1141 const local_ordinal_type sub_line_length = (interf.max_partsz - (n_subparts_per_part - 1) * connection_length) / n_subparts_per_part;
1142 const local_ordinal_type last_sub_line_length = interf.max_partsz - (n_subparts_per_part - 1) * (connection_length + sub_line_length);
1144 interf.max_subpartsz = (sub_line_length > last_sub_line_length) ? sub_line_length : last_sub_line_length;
1145 interf.n_subparts_per_part = n_subparts_per_part;
1146 interf.nparts = nparts;
1152 interf.part2rowidx0 = local_ordinal_type_1d_view(
do_not_initialize_tag(
"part2rowidx0"), nparts + 1);
1153 interf.part2packrowidx0 = local_ordinal_type_1d_view(
do_not_initialize_tag(
"part2packrowidx0"), nparts + 1);
1156 interf.part2rowidx0_sub = local_ordinal_type_1d_view(
do_not_initialize_tag(
"part2rowidx0_sub"), n_sub_parts_and_schur + 1);
1157 interf.part2packrowidx0_sub = local_ordinal_type_2d_view(
do_not_initialize_tag(
"part2packrowidx0_sub"), nparts, 2 * n_subparts_per_part);
1158 interf.rowidx2part_sub = local_ordinal_type_1d_view(
do_not_initialize_tag(
"rowidx2part"), A_n_lclrows);
1160 interf.partptr_sub = local_ordinal_type_2d_view(
do_not_initialize_tag(
"partptr_sub"), n_sub_parts_and_schur, 2);
1163 const auto partptr = Kokkos::create_mirror_view(interf.partptr);
1164 const auto partptr_sub = Kokkos::create_mirror_view(interf.partptr_sub);
1166 const auto lclrow = Kokkos::create_mirror_view(interf.lclrow);
1167 const auto part2rowidx0 = Kokkos::create_mirror_view(interf.part2rowidx0);
1168 const auto part2packrowidx0 = Kokkos::create_mirror_view(interf.part2packrowidx0);
1169 const auto rowidx2part = Kokkos::create_mirror_view(interf.rowidx2part);
1171 const auto part2rowidx0_sub = Kokkos::create_mirror_view(interf.part2rowidx0_sub);
1172 const auto part2packrowidx0_sub = Kokkos::create_mirror_view(Kokkos::HostSpace(), interf.part2packrowidx0_sub);
1173 const auto rowidx2part_sub = Kokkos::create_mirror_view(interf.rowidx2part_sub);
1176 interf.row_contiguous =
true;
1178 part2rowidx0(0) = 0;
1179 part2packrowidx0(0) = 0;
1180 local_ordinal_type pack_nrows = 0;
1181 local_ordinal_type pack_nrows_sub = 0;
1183 IFPACK2_BLOCKHELPER_TIMER(
"compute part indices (Jacobi)", Jacobi);
1187 for (local_ordinal_type i=0; i <= nparts; ++i) {
1188 part2rowidx0(i) = i;
1191 for (local_ordinal_type i=0; i < nparts; ++i) {
1195 for (local_ordinal_type ip=0;ip<nparts;++ip) {
1197 if (ip % vector_length == 0) pack_nrows = 1;
1198 part2packrowidx0(ip+1) = part2packrowidx0(ip) + ((ip+1) % vector_length == 0 || ip+1 == nparts ? pack_nrows : 0);
1200 part2rowidx0_sub(0) = 0;
1201 partptr_sub(0, 0) = 0;
1203 for (local_ordinal_type ip=0;ip<nparts;++ip) {
1204 constexpr local_ordinal_type ipnrows = 1;
1205 const local_ordinal_type full_line_length = partptr(ip+1) - partptr(ip);
1208 (full_line_length != ipnrows, std::logic_error,
1209 "In the part " << ip );
1211 constexpr local_ordinal_type connection_length = 2;
1213 if (full_line_length < n_subparts_per_part + (n_subparts_per_part - 1) * connection_length )
1215 (
true, std::logic_error,
1216 "The part " << ip <<
" is too short to use " << n_subparts_per_part <<
" sub parts.");
1218 const local_ordinal_type sub_line_length = (full_line_length - (n_subparts_per_part - 1) * connection_length) / n_subparts_per_part;
1219 const local_ordinal_type last_sub_line_length = full_line_length - (n_subparts_per_part - 1) * (connection_length + sub_line_length);
1221 if (ip % vector_length == 0) pack_nrows_sub = ipnrows;
1223 for (local_ordinal_type local_sub_ip=0; local_sub_ip<n_subparts_per_part;++local_sub_ip) {
1224 const local_ordinal_type sub_ip = nparts*(2*local_sub_ip) + ip;
1225 const local_ordinal_type schur_ip = nparts*(2*local_sub_ip+1) + ip;
1226 if (local_sub_ip != n_subparts_per_part-1) {
1227 if (local_sub_ip != 0) {
1228 partptr_sub(sub_ip, 0) = partptr_sub(nparts*(2*local_sub_ip-1) + ip, 1);
1231 partptr_sub(sub_ip, 0) = partptr_sub(nparts*2*(n_subparts_per_part-1) + ip - 1, 1);
1233 partptr_sub(sub_ip, 1) = sub_line_length + partptr_sub(sub_ip, 0);
1234 partptr_sub(schur_ip, 0) = partptr_sub(sub_ip, 1);
1235 partptr_sub(schur_ip, 1) = connection_length + partptr_sub(schur_ip, 0);
1237 part2rowidx0_sub(sub_ip + 1) = part2rowidx0_sub(sub_ip) + sub_line_length;
1238 part2rowidx0_sub(sub_ip + 2) = part2rowidx0_sub(sub_ip + 1) + connection_length;
1240 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
1241 printf(
"Sub Part index = %d, first LID associated to the sub part = %d, sub part size = %d;\n", sub_ip, partptr_sub(ip, 2 * local_sub_ip), sub_line_length);
1242 printf(
"Sub Part index Schur = %d, first LID associated to the sub part = %d, sub part size = %d;\n", sub_ip + 1, partptr_sub(ip, 2 * local_sub_ip + 1), connection_length);
1246 if (local_sub_ip != 0) {
1247 partptr_sub(sub_ip, 0) = partptr_sub(nparts*(2*local_sub_ip-1) + ip, 1);
1250 partptr_sub(sub_ip, 0) = partptr_sub(nparts*2*(n_subparts_per_part-1) + ip - 1, 1);
1252 partptr_sub(sub_ip, 1) = last_sub_line_length + partptr_sub(sub_ip, 0);
1254 part2rowidx0_sub(sub_ip + 1) = part2rowidx0_sub(sub_ip) + last_sub_line_length;
1256 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
1257 printf(
"Sub Part index = %d, first LID associated to the sub part = %d, sub part size = %d;\n", sub_ip, partptr_sub(ip, 2 * local_sub_ip), last_sub_line_length);
1263 #ifdef IFPACK2_BLOCKTRIDICONTAINER_WRITE_MM
1264 std::cout <<
"partptr_sub = " << std::endl;
1265 for (size_type i = 0; i < partptr_sub.extent(0); ++i) {
1266 for (size_type j = 0; j < partptr_sub.extent(1); ++j) {
1267 std::cout << partptr_sub(i,j) <<
" ";
1269 std::cout << std::endl;
1271 std::cout <<
"partptr_sub end" << std::endl;
1275 local_ordinal_type npacks = ceil(
float(nparts)/vector_length);
1277 local_ordinal_type ip_max = nparts > vector_length ? vector_length : nparts;
1278 for (local_ordinal_type ip=0;ip<ip_max;++ip) {
1279 part2packrowidx0_sub(ip, 0) = 0;
1281 for (local_ordinal_type ipack=0;ipack<npacks;++ipack) {
1283 local_ordinal_type ip_min = ipack*vector_length;
1284 ip_max = nparts > (ipack+1)*vector_length ? (ipack+1)*vector_length : nparts;
1285 for (local_ordinal_type ip=ip_min;ip<ip_max;++ip) {
1286 part2packrowidx0_sub(ip, 0) = part2packrowidx0_sub(ip-vector_length, part2packrowidx0_sub.extent(1)-1);
1290 for (size_type local_sub_ip=0; local_sub_ip<part2packrowidx0_sub.extent(1)-1;++local_sub_ip) {
1291 local_ordinal_type ip_min = ipack*vector_length;
1292 ip_max = nparts > (ipack+1)*vector_length ? (ipack+1)*vector_length : nparts;
1294 const local_ordinal_type full_line_length = partptr(ip_min+1) - partptr(ip_min);
1296 constexpr local_ordinal_type connection_length = 2;
1298 const local_ordinal_type sub_line_length = (full_line_length - (n_subparts_per_part - 1) * connection_length) / n_subparts_per_part;
1299 const local_ordinal_type last_sub_line_length = full_line_length - (n_subparts_per_part - 1) * (connection_length + sub_line_length);
1301 if (local_sub_ip % 2 == 0) pack_nrows_sub = sub_line_length;
1302 if (local_sub_ip % 2 == 1) pack_nrows_sub = connection_length;
1303 if (local_sub_ip == part2packrowidx0_sub.extent(1)-2) pack_nrows_sub = last_sub_line_length;
1305 part2packrowidx0_sub(ip_min, local_sub_ip + 1) = part2packrowidx0_sub(ip_min, local_sub_ip) + pack_nrows_sub;
1307 for (local_ordinal_type ip=ip_min+1;ip<ip_max;++ip) {
1308 part2packrowidx0_sub(ip, local_sub_ip + 1) = part2packrowidx0_sub(ip_min, local_sub_ip + 1);
1313 Kokkos::deep_copy(interf.part2packrowidx0_sub, part2packrowidx0_sub);
1315 IFPACK2_BLOCKHELPER_TIMER_FENCE(
typename BlockHelperDetails::ImplType<MatrixType>::execution_space)
1317 IFPACK2_BLOCKHELPER_TIMER(
"compute part indices", indices);
1318 for (local_ordinal_type ip=0;ip<nparts;++ip) {
1319 const auto* part = &partitions[p[ip]];
1320 const local_ordinal_type ipnrows = part->size();
1321 TEUCHOS_ASSERT(ip == 0 || (ipnrows <= static_cast<local_ordinal_type>(partitions[p[ip-1]].size())));
1323 BlockHelperDetails::get_msg_prefix(comm)
1324 <<
"partition " << p[ip]
1325 <<
" is empty, which is not allowed.");
1327 part2rowidx0(ip+1) = part2rowidx0(ip) + ipnrows;
1330 if (ip % vector_length == 0) pack_nrows = ipnrows;
1331 part2packrowidx0(ip+1) = part2packrowidx0(ip) + ((ip+1) % vector_length == 0 || ip+1 == nparts ? pack_nrows : 0);
1332 const local_ordinal_type offset = partptr(ip);
1333 for (local_ordinal_type i=0;i<ipnrows;++i) {
1334 const auto lcl_row = (*part)[i];
1336 BlockHelperDetails::get_msg_prefix(comm)
1337 <<
"partitions[" << p[ip] <<
"]["
1338 << i <<
"] = " << lcl_row
1339 <<
" but input matrix implies limits of [0, " << A_n_lclrows-1
1341 lclrow(offset+i) = lcl_row;
1342 rowidx2part(offset+i) = ip;
1343 if (interf.row_contiguous && offset+i > 0 && lclrow((offset+i)-1) + 1 != lcl_row)
1344 interf.row_contiguous =
false;
1346 partptr(ip+1) = offset + ipnrows;
1348 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
1349 printf(
"Part index = ip = %d, first LID associated to the part = partptr(ip) = offset = %d, part->size() = ipnrows = %d;\n", ip, offset, ipnrows);
1350 printf(
"partptr(%d+1) = %d\n", ip, partptr(ip+1));
1354 part2rowidx0_sub(0) = 0;
1355 partptr_sub(0, 0) = 0;
1358 for (local_ordinal_type ip=0;ip<nparts;++ip) {
1359 const auto* part = &partitions[p[ip]];
1360 const local_ordinal_type ipnrows = part->size();
1361 const local_ordinal_type full_line_length = partptr(ip+1) - partptr(ip);
1364 (full_line_length != ipnrows, std::logic_error,
1365 "In the part " << ip );
1367 constexpr local_ordinal_type connection_length = 2;
1369 if (full_line_length < n_subparts_per_part + (n_subparts_per_part - 1) * connection_length )
1371 (
true, std::logic_error,
1372 "The part " << ip <<
" is too short to use " << n_subparts_per_part <<
" sub parts.");
1374 const local_ordinal_type sub_line_length = (full_line_length - (n_subparts_per_part - 1) * connection_length) / n_subparts_per_part;
1375 const local_ordinal_type last_sub_line_length = full_line_length - (n_subparts_per_part - 1) * (connection_length + sub_line_length);
1377 if (ip % vector_length == 0) pack_nrows_sub = ipnrows;
1379 for (local_ordinal_type local_sub_ip=0; local_sub_ip<n_subparts_per_part;++local_sub_ip) {
1380 const local_ordinal_type sub_ip = nparts*(2*local_sub_ip) + ip;
1381 const local_ordinal_type schur_ip = nparts*(2*local_sub_ip+1) + ip;
1382 if (local_sub_ip != n_subparts_per_part-1) {
1383 if (local_sub_ip != 0) {
1384 partptr_sub(sub_ip, 0) = partptr_sub(nparts*(2*local_sub_ip-1) + ip, 1);
1387 partptr_sub(sub_ip, 0) = partptr_sub(nparts*2*(n_subparts_per_part-1) + ip - 1, 1);
1389 partptr_sub(sub_ip, 1) = sub_line_length + partptr_sub(sub_ip, 0);
1390 partptr_sub(schur_ip, 0) = partptr_sub(sub_ip, 1);
1391 partptr_sub(schur_ip, 1) = connection_length + partptr_sub(schur_ip, 0);
1393 part2rowidx0_sub(sub_ip + 1) = part2rowidx0_sub(sub_ip) + sub_line_length;
1394 part2rowidx0_sub(sub_ip + 2) = part2rowidx0_sub(sub_ip + 1) + connection_length;
1396 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
1397 printf(
"Sub Part index = %d, first LID associated to the sub part = %d, sub part size = %d;\n", sub_ip, partptr_sub(sub_ip, 0), sub_line_length);
1398 printf(
"Sub Part index Schur = %d, first LID associated to the sub part = %d, sub part size = %d;\n", sub_ip + 1, partptr_sub(ip, 2 * local_sub_ip + 1), connection_length);
1402 if (local_sub_ip != 0) {
1403 partptr_sub(sub_ip, 0) = partptr_sub(nparts*(2*local_sub_ip-1) + ip, 1);
1406 partptr_sub(sub_ip, 0) = partptr_sub(nparts*2*(n_subparts_per_part-1) + ip - 1, 1);
1408 partptr_sub(sub_ip, 1) = last_sub_line_length + partptr_sub(sub_ip, 0);
1410 part2rowidx0_sub(sub_ip + 1) = part2rowidx0_sub(sub_ip) + last_sub_line_length;
1412 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
1413 printf(
"Sub Part index = %d, first LID associated to the sub part = %d, sub part size = %d;\n", sub_ip, partptr_sub(sub_ip, 0), last_sub_line_length);
1420 local_ordinal_type npacks = ceil(
float(nparts)/vector_length);
1422 local_ordinal_type ip_max = nparts > vector_length ? vector_length : nparts;
1423 for (local_ordinal_type ip=0;ip<ip_max;++ip) {
1424 part2packrowidx0_sub(ip, 0) = 0;
1426 for (local_ordinal_type ipack=0;ipack<npacks;++ipack) {
1428 local_ordinal_type ip_min = ipack*vector_length;
1429 ip_max = nparts > (ipack+1)*vector_length ? (ipack+1)*vector_length : nparts;
1430 for (local_ordinal_type ip=ip_min;ip<ip_max;++ip) {
1431 part2packrowidx0_sub(ip, 0) = part2packrowidx0_sub(ip-vector_length, part2packrowidx0_sub.extent(1)-1);
1435 for (size_type local_sub_ip=0; local_sub_ip<part2packrowidx0_sub.extent(1)-1;++local_sub_ip) {
1436 local_ordinal_type ip_min = ipack*vector_length;
1437 ip_max = nparts > (ipack+1)*vector_length ? (ipack+1)*vector_length : nparts;
1439 const local_ordinal_type full_line_length = partptr(ip_min+1) - partptr(ip_min);
1441 constexpr local_ordinal_type connection_length = 2;
1443 const local_ordinal_type sub_line_length = (full_line_length - (n_subparts_per_part - 1) * connection_length) / n_subparts_per_part;
1444 const local_ordinal_type last_sub_line_length = full_line_length - (n_subparts_per_part - 1) * (connection_length + sub_line_length);
1446 if (local_sub_ip % 2 == 0) pack_nrows_sub = sub_line_length;
1447 if (local_sub_ip % 2 == 1) pack_nrows_sub = connection_length;
1448 if (local_sub_ip == part2packrowidx0_sub.extent(1)-2) pack_nrows_sub = last_sub_line_length;
1450 part2packrowidx0_sub(ip_min, local_sub_ip + 1) = part2packrowidx0_sub(ip_min, local_sub_ip) + pack_nrows_sub;
1452 for (local_ordinal_type ip=ip_min+1;ip<ip_max;++ip) {
1453 part2packrowidx0_sub(ip, local_sub_ip + 1) = part2packrowidx0_sub(ip_min, local_sub_ip + 1);
1458 Kokkos::deep_copy(interf.part2packrowidx0_sub, part2packrowidx0_sub);
1460 IFPACK2_BLOCKHELPER_TIMER_FENCE(
typename BlockHelperDetails::ImplType<MatrixType>::execution_space)
1462 #if defined(BLOCKTRIDICONTAINER_DEBUG)
1465 if (lclrow(0) != 0) interf.row_contiguous =
false;
1467 Kokkos::deep_copy(interf.partptr, partptr);
1468 Kokkos::deep_copy(interf.lclrow, lclrow);
1470 Kokkos::deep_copy(interf.partptr_sub, partptr_sub);
1473 interf.part2rowidx0 = interf.partptr;
1474 Kokkos::deep_copy(interf.part2packrowidx0, part2packrowidx0);
1476 interf.part2packrowidx0_back = part2packrowidx0_sub(part2packrowidx0_sub.extent(0) - 1, part2packrowidx0_sub.extent(1) - 1);
1477 Kokkos::deep_copy(interf.rowidx2part, rowidx2part);
1480 IFPACK2_BLOCKHELPER_TIMER(
"Fill packptr", packptr0);
1481 local_ordinal_type npacks = ceil(
float(nparts)/vector_length) * (part2packrowidx0_sub.extent(1)-1);
1483 for (local_ordinal_type ip=1;ip<=nparts;++ip)
1484 if (part2packrowidx0(ip) != part2packrowidx0(ip-1))
1488 const auto packptr = Kokkos::create_mirror_view(interf.packptr);
1490 for (local_ordinal_type ip=1,k=1;ip<=nparts;++ip)
1491 if (part2packrowidx0(ip) != part2packrowidx0(ip-1))
1494 Kokkos::deep_copy(interf.packptr, packptr);
1496 local_ordinal_type npacks_per_subpart = ceil(
float(nparts)/vector_length);
1497 npacks = ceil(
float(nparts)/vector_length) * (part2packrowidx0_sub.extent(1)-1);
1499 interf.packindices_sub = local_ordinal_type_1d_view(
do_not_initialize_tag(
"packindices_sub"), npacks_per_subpart*n_subparts_per_part);
1500 interf.packindices_schur = local_ordinal_type_2d_view(
do_not_initialize_tag(
"packindices_schur"), npacks_per_subpart,n_subparts_per_part-1);
1502 const auto packindices_sub = Kokkos::create_mirror_view(interf.packindices_sub);
1503 const auto packindices_schur = Kokkos::create_mirror_view(interf.packindices_schur);
1507 for (local_ordinal_type local_sub_ip=0; local_sub_ip<n_subparts_per_part-1;++local_sub_ip) {
1508 for (local_ordinal_type local_pack_ip=0; local_pack_ip<npacks_per_subpart;++local_pack_ip) {
1509 packindices_sub(local_sub_ip * npacks_per_subpart + local_pack_ip) = 2 * local_sub_ip * npacks_per_subpart + local_pack_ip;
1510 packindices_schur(local_pack_ip,local_sub_ip) = 2 * local_sub_ip * npacks_per_subpart + local_pack_ip + npacks_per_subpart;
1514 for (local_ordinal_type local_pack_ip=0; local_pack_ip<npacks_per_subpart;++local_pack_ip) {
1515 packindices_sub((n_subparts_per_part-1) * npacks_per_subpart + local_pack_ip) = 2 * (n_subparts_per_part-1) * npacks_per_subpart + local_pack_ip;
1518 #ifdef IFPACK2_BLOCKTRIDICONTAINER_WRITE_MM
1519 std::cout <<
"packindices_sub = " << std::endl;
1520 for (size_type i = 0; i < packindices_sub.extent(0); ++i) {
1521 std::cout << packindices_sub(i) <<
" ";
1523 std::cout << std::endl;
1524 std::cout <<
"packindices_sub end" << std::endl;
1526 std::cout <<
"packindices_schur = " << std::endl;
1527 for (size_type i = 0; i < packindices_schur.extent(0); ++i) {
1528 for (size_type j = 0; j < packindices_schur.extent(1); ++j) {
1529 std::cout << packindices_schur(i,j) <<
" ";
1531 std::cout << std::endl;
1534 std::cout <<
"packindices_schur end" << std::endl;
1537 Kokkos::deep_copy(interf.packindices_sub, packindices_sub);
1538 Kokkos::deep_copy(interf.packindices_schur, packindices_schur);
1541 const auto packptr_sub = Kokkos::create_mirror_view(interf.packptr_sub);
1543 for (local_ordinal_type k=0;k<npacks + 1;++k)
1544 packptr_sub(k) = packptr(k%npacks_per_subpart) + (k / npacks_per_subpart) * packptr(npacks_per_subpart);
1546 Kokkos::deep_copy(interf.packptr_sub, packptr_sub);
1547 IFPACK2_BLOCKHELPER_TIMER_FENCE(
typename BlockHelperDetails::ImplType<MatrixType>::execution_space)
1549 IFPACK2_BLOCKHELPER_TIMER_FENCE(
typename BlockHelperDetails::ImplType<MatrixType>::execution_space)
1557 template <
typename MatrixType>
1560 using local_ordinal_type_1d_view =
typename impl_type::local_ordinal_type_1d_view;
1562 using size_type_2d_view =
typename impl_type::size_type_2d_view;
1563 using vector_type_3d_view =
typename impl_type::vector_type_3d_view;
1564 using vector_type_4d_view =
typename impl_type::vector_type_4d_view;
1565 using btdm_scalar_type_3d_view =
typename impl_type::btdm_scalar_type_3d_view;
1571 size_type_2d_view flat_td_ptr, pack_td_ptr, pack_td_ptr_schur;
1574 local_ordinal_type_1d_view A_colindsub;
1577 vector_type_3d_view values;
1580 vector_type_3d_view values_schur;
1582 vector_type_4d_view e_values;
1587 size_type_1d_view diag_offsets;
1591 btdm_scalar_type_3d_view d_inv;
1593 bool is_diagonal_only;
1599 template <
typename idx_type>
1600 static KOKKOS_FORCEINLINE_FUNCTION
1601 idx_type IndexToRow (
const idx_type& ind) {
return (ind + 1) / 3; }
1604 template <
typename idx_type>
1605 static KOKKOS_FORCEINLINE_FUNCTION
1606 idx_type RowToIndex (
const idx_type& row) {
return row > 0 ? 3*row - 1 : 0; }
1608 template <
typename idx_type>
1609 static KOKKOS_FORCEINLINE_FUNCTION
1610 idx_type NumBlocks (
const idx_type& nrows) {
return nrows > 0 ? 3*nrows - 2 : 0; }
1612 template <
typename idx_type>
1613 static KOKKOS_FORCEINLINE_FUNCTION
1614 idx_type NumBlocksSchur (
const idx_type& nrows) {
return nrows > 0 ? 3*nrows + 2 : 0; }
1621 template<
typename MatrixType>
1624 IFPACK2_BLOCKHELPER_TIMER(
"createBlockTridiags", createBlockTridiags0);
1626 using execution_space =
typename impl_type::execution_space;
1627 using local_ordinal_type =
typename impl_type::local_ordinal_type;
1628 using size_type =
typename impl_type::size_type;
1629 using size_type_2d_view =
typename impl_type::size_type_2d_view;
1631 constexpr
int vector_length = impl_type::vector_length;
1635 const local_ordinal_type ntridiags = interf.partptr_sub.extent(0);
1638 btdm.flat_td_ptr = size_type_2d_view(
do_not_initialize_tag(
"btdm.flat_td_ptr"), interf.nparts, 2*interf.n_subparts_per_part);
1639 const Kokkos::RangePolicy<execution_space> policy(0, 2 * interf.nparts * interf.n_subparts_per_part );
1640 Kokkos::parallel_scan
1641 (
"createBlockTridiags::RangePolicy::flat_td_ptr",
1642 policy, KOKKOS_LAMBDA(
const local_ordinal_type &i, size_type &update,
const bool &
final) {
1643 const local_ordinal_type partidx = i/(2 * interf.n_subparts_per_part);
1644 const local_ordinal_type local_subpartidx = i % (2 * interf.n_subparts_per_part);
1647 btdm.flat_td_ptr(partidx, local_subpartidx) = update;
1649 if (local_subpartidx != (2 * interf.n_subparts_per_part -1)) {
1650 const local_ordinal_type nrows = interf.partptr_sub(interf.nparts*local_subpartidx + partidx,1) - interf.partptr_sub(interf.nparts*local_subpartidx + partidx,0);
1651 if (local_subpartidx % 2 == 0)
1652 update += btdm.NumBlocks(nrows);
1654 update += btdm.NumBlocksSchur(nrows);
1658 const auto nblocks = Kokkos::create_mirror_view_and_copy
1659 (Kokkos::HostSpace(), Kokkos::subview(btdm.flat_td_ptr, interf.nparts-1, 2*interf.n_subparts_per_part-1));
1660 btdm.is_diagonal_only = (
static_cast<local_ordinal_type
>(nblocks()) == ntridiags);
1664 if (vector_length == 1) {
1665 btdm.pack_td_ptr = btdm.flat_td_ptr;
1669 local_ordinal_type npacks_per_subpart = 0;
1670 const auto part2packrowidx0 = Kokkos::create_mirror_view(interf.part2packrowidx0);
1671 Kokkos::deep_copy(part2packrowidx0, interf.part2packrowidx0);
1672 for (local_ordinal_type ip=1;ip<=interf.nparts;++ip)
1673 if (part2packrowidx0(ip) != part2packrowidx0(ip-1))
1674 ++npacks_per_subpart;
1676 btdm.pack_td_ptr = size_type_2d_view(
do_not_initialize_tag(
"btdm.pack_td_ptr"), interf.nparts, 2*interf.n_subparts_per_part);
1677 const Kokkos::RangePolicy<execution_space> policy(0,npacks_per_subpart);
1679 Kokkos::parallel_for
1680 (
"createBlockTridiags::RangePolicy::pack_td_ptr",
1681 policy, KOKKOS_LAMBDA(
const local_ordinal_type &i) {
1682 for (local_ordinal_type j = 0; j < 2*interf.n_subparts_per_part; ++j) {
1683 const local_ordinal_type pack_id = ( j == 2*interf.n_subparts_per_part-1 ) ? i+(j-1)*npacks_per_subpart : i+j*npacks_per_subpart;
1684 const local_ordinal_type nparts_in_pack = interf.packptr_sub(pack_id+1) - interf.packptr_sub(pack_id);
1686 const local_ordinal_type parti = interf.packptr_sub(pack_id);
1687 const local_ordinal_type partidx = parti%interf.nparts;
1689 for (local_ordinal_type pti=0;pti<nparts_in_pack;++pti) {
1690 btdm.pack_td_ptr(partidx+pti, j) = btdm.flat_td_ptr(i, j);
1696 btdm.pack_td_ptr_schur = size_type_2d_view(
do_not_initialize_tag(
"btdm.pack_td_ptr_schur"), interf.nparts, interf.n_subparts_per_part);
1698 const auto host_pack_td_ptr_schur = Kokkos::create_mirror_view(btdm.pack_td_ptr_schur);
1699 constexpr local_ordinal_type connection_length = 2;
1701 host_pack_td_ptr_schur(0,0) = 0;
1702 for (local_ordinal_type i = 0; i < interf.nparts; ++i) {
1703 if (i % vector_length == 0) {
1705 host_pack_td_ptr_schur(i,0) = host_pack_td_ptr_schur(i-1,host_pack_td_ptr_schur.extent(1)-1);
1706 for (local_ordinal_type j = 0; j < interf.n_subparts_per_part-1; ++j) {
1707 host_pack_td_ptr_schur(i,j+1) = host_pack_td_ptr_schur(i,j) + btdm.NumBlocks(connection_length) + (j != 0 ? 1 : 0) + (j != interf.n_subparts_per_part-2 ? 1 : 0);
1711 for (local_ordinal_type j = 0; j < interf.n_subparts_per_part; ++j) {
1712 host_pack_td_ptr_schur(i,j) = host_pack_td_ptr_schur(i-1,j);
1717 Kokkos::deep_copy(btdm.pack_td_ptr_schur, host_pack_td_ptr_schur);
1719 #ifdef IFPACK2_BLOCKTRIDICONTAINER_WRITE_MM
1720 const auto host_flat_td_ptr = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), btdm.flat_td_ptr);
1721 std::cout <<
"flat_td_ptr = " << std::endl;
1722 for (size_type i = 0; i < host_flat_td_ptr.extent(0); ++i) {
1723 for (size_type j = 0; j < host_flat_td_ptr.extent(1); ++j) {
1724 std::cout << host_flat_td_ptr(i,j) <<
" ";
1726 std::cout << std::endl;
1728 std::cout <<
"flat_td_ptr end" << std::endl;
1730 const auto host_pack_td_ptr = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), btdm.pack_td_ptr);
1732 std::cout <<
"pack_td_ptr = " << std::endl;
1733 for (size_type i = 0; i < host_pack_td_ptr.extent(0); ++i) {
1734 for (size_type j = 0; j < host_pack_td_ptr.extent(1); ++j) {
1735 std::cout << host_pack_td_ptr(i,j) <<
" ";
1737 std::cout << std::endl;
1739 std::cout <<
"pack_td_ptr end" << std::endl;
1742 std::cout <<
"pack_td_ptr_schur = " << std::endl;
1743 for (size_type i = 0; i < host_pack_td_ptr_schur.extent(0); ++i) {
1744 for (size_type j = 0; j < host_pack_td_ptr_schur.extent(1); ++j) {
1745 std::cout << host_pack_td_ptr_schur(i,j) <<
" ";
1747 std::cout << std::endl;
1749 std::cout <<
"pack_td_ptr_schur end" << std::endl;
1753 IFPACK2_BLOCKHELPER_TIMER_FENCE(
typename BlockHelperDetails::ImplType<MatrixType>::execution_space)
1768 template<
typename MatrixType>
1770 setTridiagsToIdentity
1771 (
const BlockTridiags<MatrixType>& btdm,
1772 const typename BlockHelperDetails::ImplType<MatrixType>::local_ordinal_type_1d_view& packptr)
1775 using execution_space =
typename impl_type::execution_space;
1776 using local_ordinal_type =
typename impl_type::local_ordinal_type;
1777 using size_type_2d_view =
typename impl_type::size_type_2d_view;
1779 const ConstUnmanaged<size_type_2d_view> pack_td_ptr(btdm.pack_td_ptr);
1780 const local_ordinal_type blocksize = btdm.values.extent(1);
1783 const int vector_length = impl_type::vector_length;
1784 const int internal_vector_length = impl_type::internal_vector_length;
1786 using btdm_scalar_type =
typename impl_type::btdm_scalar_type;
1787 using internal_vector_type =
typename impl_type::internal_vector_type;
1788 using internal_vector_type_4d_view =
1789 typename impl_type::internal_vector_type_4d_view;
1791 using team_policy_type = Kokkos::TeamPolicy<execution_space>;
1792 const internal_vector_type_4d_view values
1793 (reinterpret_cast<internal_vector_type*>(btdm.values.data()),
1794 btdm.values.extent(0),
1795 btdm.values.extent(1),
1796 btdm.values.extent(2),
1797 vector_length/internal_vector_length);
1798 const local_ordinal_type vector_loop_size = values.extent(3);
1799 #if defined(KOKKOS_ENABLE_CUDA) && defined(__CUDA_ARCH__)
1800 local_ordinal_type total_team_size(0);
1801 if (blocksize <= 5) total_team_size = 32;
1802 else if (blocksize <= 9) total_team_size = 64;
1803 else if (blocksize <= 12) total_team_size = 96;
1804 else if (blocksize <= 16) total_team_size = 128;
1805 else if (blocksize <= 20) total_team_size = 160;
1806 else total_team_size = 160;
1807 const local_ordinal_type team_size = total_team_size/vector_loop_size;
1808 const team_policy_type policy(packptr.extent(0)-1, team_size, vector_loop_size);
1809 #elif defined(KOKKOS_ENABLE_HIP)
1814 local_ordinal_type total_team_size(0);
1815 if (blocksize <= 5) total_team_size = 32;
1816 else if (blocksize <= 9) total_team_size = 64;
1817 else if (blocksize <= 12) total_team_size = 96;
1818 else if (blocksize <= 16) total_team_size = 128;
1819 else if (blocksize <= 20) total_team_size = 160;
1820 else total_team_size = 160;
1821 const local_ordinal_type team_size = total_team_size/vector_loop_size;
1822 const team_policy_type policy(packptr.extent(0)-1, team_size, vector_loop_size);
1823 #elif defined(KOKKOS_ENABLE_SYCL)
1825 local_ordinal_type total_team_size(0);
1826 if (blocksize <= 5) total_team_size = 32;
1827 else if (blocksize <= 9) total_team_size = 64;
1828 else if (blocksize <= 12) total_team_size = 96;
1829 else if (blocksize <= 16) total_team_size = 128;
1830 else if (blocksize <= 20) total_team_size = 160;
1831 else total_team_size = 160;
1832 const local_ordinal_type team_size = total_team_size/vector_loop_size;
1833 const team_policy_type policy(packptr.extent(0)-1, team_size, vector_loop_size);
1836 const team_policy_type policy(packptr.extent(0)-1, 1, 1);
1838 Kokkos::parallel_for
1839 (
"setTridiagsToIdentity::TeamPolicy",
1840 policy, KOKKOS_LAMBDA(
const typename team_policy_type::member_type &member) {
1841 const local_ordinal_type k = member.league_rank();
1842 const local_ordinal_type ibeg = pack_td_ptr(packptr(k),0);
1843 const local_ordinal_type iend = pack_td_ptr(packptr(k),pack_td_ptr.extent(1)-1);
1845 const local_ordinal_type diff = iend - ibeg;
1846 const local_ordinal_type icount = diff/3 + (diff%3 > 0);
1847 const btdm_scalar_type one(1);
1848 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
1849 Kokkos::parallel_for(Kokkos::TeamThreadRange(member,icount),[&](
const local_ordinal_type &ii) {
1850 const local_ordinal_type i = ibeg + ii*3;
1851 for (local_ordinal_type j=0;j<blocksize;++j) {
1852 values(i,j,j,v) = one;
1863 template<
typename MatrixType>
1866 const Teuchos::RCP<
const typename BlockHelperDetails::ImplType<MatrixType>::tpetra_crs_graph_type> &g,
1867 const BlockHelperDetails::PartInterface<MatrixType> &interf,
1870 const bool overlap_communication_and_computation,
1871 const Teuchos::RCP<AsyncableImport<MatrixType> > &async_importer,
1873 bool use_fused_jacobi) {
1874 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::SymbolicPhase", SymbolicPhase);
1878 using execution_space =
typename impl_type::execution_space;
1879 using host_execution_space =
typename impl_type::host_execution_space;
1881 using local_ordinal_type =
typename impl_type::local_ordinal_type;
1882 using global_ordinal_type =
typename impl_type::global_ordinal_type;
1883 using size_type =
typename impl_type::size_type;
1884 using local_ordinal_type_1d_view =
typename impl_type::local_ordinal_type_1d_view;
1885 using size_type_1d_view =
typename impl_type::size_type_1d_view;
1886 using vector_type_3d_view =
typename impl_type::vector_type_3d_view;
1887 using vector_type_4d_view =
typename impl_type::vector_type_4d_view;
1888 using crs_matrix_type =
typename impl_type::tpetra_crs_matrix_type;
1889 using block_crs_matrix_type =
typename impl_type::tpetra_block_crs_matrix_type;
1890 using btdm_scalar_type_3d_view =
typename impl_type::btdm_scalar_type_3d_view;
1892 constexpr
int vector_length = impl_type::vector_length;
1894 const auto comm = A->getRowMap()->getComm();
1896 auto A_crs = Teuchos::rcp_dynamic_cast<
const crs_matrix_type>(A);
1897 auto A_bcrs = Teuchos::rcp_dynamic_cast<
const block_crs_matrix_type>(A);
1899 bool hasBlockCrsMatrix = ! A_bcrs.is_null ();
1901 const local_ordinal_type blocksize = hasBlockCrsMatrix ? A->getBlockSize() : A->getLocalNumRows()/g->getLocalNumRows();
1904 const auto partptr = Kokkos::create_mirror_view_and_copy (Kokkos::HostSpace(), interf.partptr);
1905 const auto lclrow = Kokkos::create_mirror_view_and_copy (Kokkos::HostSpace(), interf.lclrow);
1906 const auto rowidx2part = Kokkos::create_mirror_view_and_copy (Kokkos::HostSpace(), interf.rowidx2part);
1907 const auto part2rowidx0 = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), interf.part2rowidx0);
1908 const auto packptr = Kokkos::create_mirror_view_and_copy (Kokkos::HostSpace(), interf.packptr);
1910 const local_ordinal_type nrows = partptr(partptr.extent(0) - 1);
1912 Kokkos::View<local_ordinal_type*,host_execution_space> col2row(
"col2row", A->getLocalNumCols());
1918 const auto rowmap = g->getRowMap();
1919 const auto colmap = g->getColMap();
1920 const auto dommap = g->getDomainMap();
1921 TEUCHOS_ASSERT( !(rowmap.is_null() || colmap.is_null() || dommap.is_null()));
1923 #if !defined(__CUDA_ARCH__) && !defined(__HIP_DEVICE_COMPILE__) && !defined(__SYCL_DEVICE_ONLY__)
1924 const Kokkos::RangePolicy<host_execution_space> policy(0,nrows);
1925 Kokkos::parallel_for
1926 (
"performSymbolicPhase::RangePolicy::col2row",
1927 policy, KOKKOS_LAMBDA(
const local_ordinal_type &lr) {
1928 const global_ordinal_type gid = rowmap->getGlobalElement(lr);
1930 if (dommap->isNodeGlobalElement(gid)) {
1931 const local_ordinal_type lc = colmap->getLocalElement(gid);
1932 # if defined(BLOCKTRIDICONTAINER_DEBUG)
1934 BlockHelperDetails::get_msg_prefix(comm) <<
"GID " << gid
1935 <<
" gives an invalid local column.");
1945 const auto local_graph = g->getLocalGraphHost();
1946 const auto local_graph_rowptr = local_graph.row_map;
1947 TEUCHOS_ASSERT(local_graph_rowptr.size() ==
static_cast<size_t>(nrows + 1));
1948 const auto local_graph_colidx = local_graph.entries;
1952 Kokkos::View<local_ordinal_type*,host_execution_space> lclrow2idx(
"lclrow2idx", nrows);
1954 const Kokkos::RangePolicy<host_execution_space> policy(0,nrows);
1955 Kokkos::parallel_for
1956 (
"performSymbolicPhase::RangePolicy::lclrow2idx",
1957 policy, KOKKOS_LAMBDA(
const local_ordinal_type &i) {
1958 lclrow2idx[lclrow(i)] = i;
1964 typename sum_reducer_type::value_type sum_reducer_value;
1966 const Kokkos::RangePolicy<host_execution_space> policy(0,nrows);
1967 Kokkos::parallel_reduce
1970 policy, KOKKOS_LAMBDA(
const local_ordinal_type &lr,
typename sum_reducer_type::value_type &update) {
1972 const local_ordinal_type ri0 = lclrow2idx[lr];
1973 const local_ordinal_type pi0 = rowidx2part(ri0);
1974 for (size_type j=local_graph_rowptr(lr);j<local_graph_rowptr(lr+1);++j) {
1975 const local_ordinal_type lc = local_graph_colidx(j);
1976 const local_ordinal_type lc2r = col2row[lc];
1977 bool incr_R =
false;
1979 if (lc2r == (local_ordinal_type) -1) {
1983 const local_ordinal_type ri = lclrow2idx[lc2r];
1984 const local_ordinal_type pi = rowidx2part(ri);
1992 if (ri0 + 1 >= ri && ri0 <= ri + 1)
1998 if (lc < nrows) ++update.v[1];
2002 }, sum_reducer_type(sum_reducer_value));
2004 size_type D_nnz = sum_reducer_value.v[0];
2005 size_type R_nnz_owned = sum_reducer_value.v[1];
2006 size_type R_nnz_remote = sum_reducer_value.v[2];
2008 if (!overlap_communication_and_computation) {
2009 R_nnz_owned += R_nnz_remote;
2015 const auto flat_td_ptr = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), btdm.flat_td_ptr);
2017 btdm.A_colindsub = local_ordinal_type_1d_view(
"btdm.A_colindsub", D_nnz);
2018 const auto D_A_colindsub = Kokkos::create_mirror_view(btdm.A_colindsub);
2020 #if defined(BLOCKTRIDICONTAINER_DEBUG)
2024 const local_ordinal_type nparts = partptr.extent(0) - 1;
2027 const Kokkos::RangePolicy<host_execution_space> policy(0, nparts);
2028 Kokkos::parallel_for
2029 (
"performSymbolicPhase::RangePolicy<host_execution_space>::D_graph",
2030 policy, KOKKOS_LAMBDA(
const local_ordinal_type &pi0) {
2031 const local_ordinal_type part_ri0 = part2rowidx0(pi0);
2032 local_ordinal_type offset = 0;
2033 for (local_ordinal_type ri0=partptr(pi0);ri0<partptr(pi0+1);++ri0) {
2034 const local_ordinal_type td_row_os = btdm.RowToIndex(ri0 - part_ri0) + offset;
2036 const local_ordinal_type lr0 = lclrow(ri0);
2037 const size_type j0 = local_graph_rowptr(lr0);
2038 for (size_type j=j0;j<local_graph_rowptr(lr0+1);++j) {
2039 const local_ordinal_type lc = local_graph_colidx(j);
2040 const local_ordinal_type lc2r = col2row[lc];
2041 if (lc2r == (local_ordinal_type) -1)
continue;
2042 const local_ordinal_type ri = lclrow2idx[lc2r];
2043 const local_ordinal_type pi = rowidx2part(ri);
2044 if (pi != pi0)
continue;
2045 if (ri + 1 < ri0 || ri > ri0 + 1)
continue;
2046 const local_ordinal_type row_entry = j - j0;
2047 D_A_colindsub(flat_td_ptr(pi0,0) + ((td_row_os + ri) - ri0)) = row_entry;
2052 #if defined(BLOCKTRIDICONTAINER_DEBUG)
2053 for (
size_t i=0;i<D_A_colindsub.extent(0);++i)
2056 Kokkos::deep_copy(btdm.A_colindsub, D_A_colindsub);
2060 const auto pack_td_ptr_last = Kokkos::subview(btdm.pack_td_ptr, btdm.pack_td_ptr.extent(0)-1, btdm.pack_td_ptr.extent(1)-1);
2061 const auto num_packed_blocks = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), pack_td_ptr_last);
2062 btdm.values = vector_type_3d_view(
"btdm.values", num_packed_blocks(), blocksize, blocksize);
2064 if (interf.n_subparts_per_part > 1) {
2065 const auto pack_td_ptr_schur_last = Kokkos::subview(btdm.pack_td_ptr_schur, btdm.pack_td_ptr_schur.extent(0)-1, btdm.pack_td_ptr_schur.extent(1)-1);
2066 const auto num_packed_blocks_schur = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), pack_td_ptr_schur_last);
2067 btdm.values_schur = vector_type_3d_view(
"btdm.values_schur", num_packed_blocks_schur(), blocksize, blocksize);
2070 if (vector_length > 1) setTridiagsToIdentity(btdm, interf.packptr);
2076 amd.rowptr = size_type_1d_view(
"amd.rowptr", nrows + 1);
2077 amd.A_colindsub = local_ordinal_type_1d_view(
do_not_initialize_tag(
"amd.A_colindsub"), R_nnz_owned);
2079 const auto R_rowptr = Kokkos::create_mirror_view(amd.rowptr);
2080 const auto R_A_colindsub = Kokkos::create_mirror_view(amd.A_colindsub);
2082 amd.rowptr_remote = size_type_1d_view(
"amd.rowptr_remote", overlap_communication_and_computation ? nrows + 1 : 0);
2083 amd.A_colindsub_remote = local_ordinal_type_1d_view(
do_not_initialize_tag(
"amd.A_colindsub_remote"), R_nnz_remote);
2085 const auto R_rowptr_remote = Kokkos::create_mirror_view(amd.rowptr_remote);
2086 const auto R_A_colindsub_remote = Kokkos::create_mirror_view(amd.A_colindsub_remote);
2089 const Kokkos::RangePolicy<host_execution_space> policy(0,nrows);
2090 Kokkos::parallel_for
2091 (
"performSymbolicPhase::RangePolicy<host_execution_space>::R_graph_count",
2092 policy, KOKKOS_LAMBDA(
const local_ordinal_type &lr) {
2093 const local_ordinal_type ri0 = lclrow2idx[lr];
2094 const local_ordinal_type pi0 = rowidx2part(ri0);
2095 const size_type j0 = local_graph_rowptr(lr);
2096 for (size_type j=j0;j<local_graph_rowptr(lr+1);++j) {
2097 const local_ordinal_type lc = local_graph_colidx(j);
2098 const local_ordinal_type lc2r = col2row[lc];
2099 if (lc2r != (local_ordinal_type) -1) {
2100 const local_ordinal_type ri = lclrow2idx[lc2r];
2101 const local_ordinal_type pi = rowidx2part(ri);
2102 if (pi == pi0 && ri + 1 >= ri0 && ri <= ri0 + 1) {
2107 if (!overlap_communication_and_computation || lc < nrows) {
2110 ++R_rowptr_remote(lr);
2119 Kokkos::RangePolicy<host_execution_space> policy(0,nrows+1);
2120 Kokkos::parallel_scan
2121 (
"performSymbolicPhase::RangePolicy<host_execution_space>::R_graph_fill",
2122 policy, KOKKOS_LAMBDA(
const local_ordinal_type &lr,
2123 update_type &update,
2124 const bool &
final) {
2126 val.v[0] = R_rowptr(lr);
2127 if (overlap_communication_and_computation)
2128 val.v[1] = R_rowptr_remote(lr);
2131 R_rowptr(lr) = update.v[0];
2132 if (overlap_communication_and_computation)
2133 R_rowptr_remote(lr) = update.v[1];
2136 const local_ordinal_type ri0 = lclrow2idx[lr];
2137 const local_ordinal_type pi0 = rowidx2part(ri0);
2139 size_type cnt_rowptr = R_rowptr(lr);
2140 size_type cnt_rowptr_remote = overlap_communication_and_computation ? R_rowptr_remote(lr) : 0;
2142 const size_type j0 = local_graph_rowptr(lr);
2143 for (size_type j=j0;j<local_graph_rowptr(lr+1);++j) {
2144 const local_ordinal_type lc = local_graph_colidx(j);
2145 const local_ordinal_type lc2r = col2row[lc];
2146 if (lc2r != (local_ordinal_type) -1) {
2147 const local_ordinal_type ri = lclrow2idx[lc2r];
2148 const local_ordinal_type pi = rowidx2part(ri);
2149 if (pi == pi0 && ri + 1 >= ri0 && ri <= ri0 + 1)
2152 const local_ordinal_type row_entry = j - j0;
2153 if (!overlap_communication_and_computation || lc < nrows)
2154 R_A_colindsub(cnt_rowptr++) = row_entry;
2156 R_A_colindsub_remote(cnt_rowptr_remote++) = row_entry;
2164 Kokkos::deep_copy(amd.rowptr, R_rowptr);
2165 Kokkos::deep_copy(amd.A_colindsub, R_A_colindsub);
2166 if (overlap_communication_and_computation) {
2168 Kokkos::deep_copy(amd.rowptr_remote, R_rowptr_remote);
2169 Kokkos::deep_copy(amd.A_colindsub_remote, R_A_colindsub_remote);
2173 if (hasBlockCrsMatrix)
2174 amd.tpetra_values = (
const_cast<block_crs_matrix_type*
>(A_bcrs.get())->getValuesDeviceNonConst());
2176 amd.tpetra_values = (
const_cast<crs_matrix_type*
>(A_crs.get()))->getLocalValuesDevice (Tpetra::Access::ReadWrite);
2182 if (interf.n_subparts_per_part > 1)
2183 btdm.e_values = vector_type_4d_view(
"btdm.e_values", 2, interf.part2packrowidx0_back, blocksize, blocksize);
2193 if(BlockHelperDetails::is_device<execution_space>::value && !useSeqMethod && hasBlockCrsMatrix)
2195 bool is_async_importer_active = !async_importer.is_null();
2196 local_ordinal_type_1d_view dm2cm = is_async_importer_active ? async_importer->dm2cm : local_ordinal_type_1d_view();
2197 bool ownedRemoteSeparate = overlap_communication_and_computation || !is_async_importer_active;
2198 BlockHelperDetails::precompute_A_x_offsets<MatrixType>(amd, interf, g, dm2cm, blocksize, ownedRemoteSeparate);
2202 if(use_fused_jacobi) {
2203 btdm.d_inv = btdm_scalar_type_3d_view(
do_not_initialize_tag(
"btdm.d_inv"), interf.nparts, blocksize, blocksize);
2204 auto rowptrs = A_bcrs->getCrsGraph().getLocalRowPtrsDevice();
2205 auto entries = A_bcrs->getCrsGraph().getLocalIndicesDevice();
2206 btdm.diag_offsets = BlockHelperDetails::findDiagOffsets<execution_space, size_type_1d_view>(rowptrs, entries, interf.nparts, blocksize);
2208 IFPACK2_BLOCKHELPER_TIMER_FENCE(
typename BlockHelperDetails::ImplType<MatrixType>::execution_space)
2215 template<
typename ArgActiveExecutionMemorySpace>
2220 typedef KB::Mode::Serial mode_type;
2221 #if defined(__KOKKOSBATCHED_INTEL_MKL_COMPACT_BATCHED__)
2222 typedef KB::Algo::Level3::CompactMKL algo_type;
2224 typedef KB::Algo::Level3::Blocked algo_type;
2226 static int recommended_team_size(
const int ,
2234 #if defined(KOKKOS_ENABLE_CUDA)
2235 static inline int ExtractAndFactorizeRecommendedCudaTeamSize(
const int blksize,
2236 const int vector_length,
2237 const int internal_vector_length) {
2238 const int vector_size = vector_length/internal_vector_length;
2239 int total_team_size(0);
2240 if (blksize <= 5) total_team_size = 32;
2241 else if (blksize <= 9) total_team_size = 32;
2242 else if (blksize <= 12) total_team_size = 96;
2243 else if (blksize <= 16) total_team_size = 128;
2244 else if (blksize <= 20) total_team_size = 160;
2245 else total_team_size = 160;
2246 return 2*total_team_size/vector_size;
2249 struct ExtractAndFactorizeTridiagsDefaultModeAndAlgo<Kokkos::CudaSpace> {
2250 typedef KB::Mode::Team mode_type;
2251 typedef KB::Algo::Level3::Unblocked algo_type;
2252 static int recommended_team_size(
const int blksize,
2253 const int vector_length,
2254 const int internal_vector_length) {
2255 return ExtractAndFactorizeRecommendedCudaTeamSize(blksize, vector_length, internal_vector_length);
2259 struct ExtractAndFactorizeTridiagsDefaultModeAndAlgo<Kokkos::CudaUVMSpace> {
2260 typedef KB::Mode::Team mode_type;
2261 typedef KB::Algo::Level3::Unblocked algo_type;
2262 static int recommended_team_size(
const int blksize,
2263 const int vector_length,
2264 const int internal_vector_length) {
2265 return ExtractAndFactorizeRecommendedCudaTeamSize(blksize, vector_length, internal_vector_length);
2270 #if defined(KOKKOS_ENABLE_HIP)
2271 static inline int ExtractAndFactorizeRecommendedHIPTeamSize(
const int blksize,
2272 const int vector_length,
2273 const int internal_vector_length) {
2274 const int vector_size = vector_length/internal_vector_length;
2275 int total_team_size(0);
2276 if (blksize <= 5) total_team_size = 32;
2277 else if (blksize <= 9) total_team_size = 32;
2278 else if (blksize <= 12) total_team_size = 96;
2279 else if (blksize <= 16) total_team_size = 128;
2280 else if (blksize <= 20) total_team_size = 160;
2281 else total_team_size = 160;
2282 return 2*total_team_size/vector_size;
2285 struct ExtractAndFactorizeTridiagsDefaultModeAndAlgo<Kokkos::HIPSpace> {
2286 typedef KB::Mode::Team mode_type;
2287 typedef KB::Algo::Level3::Unblocked algo_type;
2288 static int recommended_team_size(
const int blksize,
2289 const int vector_length,
2290 const int internal_vector_length) {
2291 return ExtractAndFactorizeRecommendedHIPTeamSize(blksize, vector_length, internal_vector_length);
2295 struct ExtractAndFactorizeTridiagsDefaultModeAndAlgo<Kokkos::HIPHostPinnedSpace> {
2296 typedef KB::Mode::Team mode_type;
2297 typedef KB::Algo::Level3::Unblocked algo_type;
2298 static int recommended_team_size(
const int blksize,
2299 const int vector_length,
2300 const int internal_vector_length) {
2301 return ExtractAndFactorizeRecommendedHIPTeamSize(blksize, vector_length, internal_vector_length);
2306 #if defined(KOKKOS_ENABLE_SYCL)
2307 static inline int ExtractAndFactorizeRecommendedSYCLTeamSize(
const int blksize,
2308 const int vector_length,
2309 const int internal_vector_length) {
2310 const int vector_size = vector_length/internal_vector_length;
2311 int total_team_size(0);
2312 if (blksize <= 5) total_team_size = 32;
2313 else if (blksize <= 9) total_team_size = 32;
2314 else if (blksize <= 12) total_team_size = 96;
2315 else if (blksize <= 16) total_team_size = 128;
2316 else if (blksize <= 20) total_team_size = 160;
2317 else total_team_size = 160;
2318 return 2*total_team_size/vector_size;
2321 struct ExtractAndFactorizeTridiagsDefaultModeAndAlgo<Kokkos::Experimental::SYCLDeviceUSMSpace> {
2322 typedef KB::Mode::Team mode_type;
2323 typedef KB::Algo::Level3::Unblocked algo_type;
2324 static int recommended_team_size(
const int blksize,
2325 const int vector_length,
2326 const int internal_vector_length) {
2327 return ExtractAndFactorizeRecommendedSYCLTeamSize(blksize, vector_length, internal_vector_length);
2331 struct ExtractAndFactorizeTridiagsDefaultModeAndAlgo<Kokkos::Experimental::SYCLSharedUSMSpace> {
2332 typedef KB::Mode::Team mode_type;
2333 typedef KB::Algo::Level3::Unblocked algo_type;
2334 static int recommended_team_size(
const int blksize,
2335 const int vector_length,
2336 const int internal_vector_length) {
2337 return ExtractAndFactorizeRecommendedSYCLTeamSize(blksize, vector_length, internal_vector_length);
2342 template<
typename impl_type,
typename WWViewType>
2343 KOKKOS_INLINE_FUNCTION
2345 solveMultiVector(
const typename Kokkos::TeamPolicy<typename impl_type::execution_space>::member_type &member,
2346 const typename impl_type::local_ordinal_type &,
2347 const typename impl_type::local_ordinal_type &i0,
2348 const typename impl_type::local_ordinal_type &r0,
2349 const typename impl_type::local_ordinal_type &nrows,
2350 const typename impl_type::local_ordinal_type &v,
2351 const ConstUnmanaged<typename impl_type::internal_vector_type_4d_view> D_internal_vector_values,
2352 const Unmanaged<typename impl_type::internal_vector_type_4d_view> X_internal_vector_values,
2353 const WWViewType &WW,
2354 const bool skip_first_pass=
false) {
2355 using execution_space =
typename impl_type::execution_space;
2356 using team_policy_type = Kokkos::TeamPolicy<execution_space>;
2357 using member_type =
typename team_policy_type::member_type;
2358 using local_ordinal_type =
typename impl_type::local_ordinal_type;
2360 typedef SolveTridiagsDefaultModeAndAlgo
2361 <
typename execution_space::memory_space> default_mode_and_algo_type;
2363 typedef typename default_mode_and_algo_type::mode_type default_mode_type;
2364 typedef typename default_mode_and_algo_type::multi_vector_algo_type default_algo_type;
2366 using btdm_magnitude_type =
typename impl_type::btdm_magnitude_type;
2369 const auto one = Kokkos::ArithTraits<btdm_magnitude_type>::one();
2370 const auto zero = Kokkos::ArithTraits<btdm_magnitude_type>::zero();
2373 auto A = Kokkos::subview(D_internal_vector_values, i0, Kokkos::ALL(), Kokkos::ALL(), v);
2374 auto X1 = Kokkos::subview(X_internal_vector_values, r0, Kokkos::ALL(), Kokkos::ALL(), v);
2377 local_ordinal_type i = i0, r = r0;
2382 if (skip_first_pass) {
2385 A.assign_data( &D_internal_vector_values(i+2,0,0,v) );
2386 X2.assign_data( &X_internal_vector_values(++r,0,0,v) );
2387 A.assign_data( &D_internal_vector_values(i+3,0,0,v) );
2388 KB::Trsm<member_type,
2389 KB::Side::Left,KB::Uplo::Lower,KB::Trans::NoTranspose,KB::Diag::Unit,
2390 default_mode_type,default_algo_type>
2391 ::invoke(member, one, A, X2);
2392 X1.assign_data( X2.data() );
2396 KB::Trsm<member_type,
2397 KB::Side::Left,KB::Uplo::Lower,KB::Trans::NoTranspose,KB::Diag::Unit,
2398 default_mode_type,default_algo_type>
2399 ::invoke(member, one, A, X1);
2400 for (local_ordinal_type tr=1;tr<nrows;++tr,i+=3) {
2401 A.assign_data( &D_internal_vector_values(i+2,0,0,v) );
2402 X2.assign_data( &X_internal_vector_values(++r,0,0,v) );
2403 member.team_barrier();
2404 KB::Gemm<member_type,
2405 KB::Trans::NoTranspose,KB::Trans::NoTranspose,
2406 default_mode_type,default_algo_type>
2407 ::invoke(member, -one, A, X1, one, X2);
2408 A.assign_data( &D_internal_vector_values(i+3,0,0,v) );
2409 KB::Trsm<member_type,
2410 KB::Side::Left,KB::Uplo::Lower,KB::Trans::NoTranspose,KB::Diag::Unit,
2411 default_mode_type,default_algo_type>
2412 ::invoke(member, one, A, X2);
2413 X1.assign_data( X2.data() );
2418 KB::Trsm<member_type,
2419 KB::Side::Left,KB::Uplo::Upper,KB::Trans::NoTranspose,KB::Diag::NonUnit,
2420 default_mode_type,default_algo_type>
2421 ::invoke(member, one, A, X1);
2422 for (local_ordinal_type tr=nrows;tr>1;--tr) {
2424 A.assign_data( &D_internal_vector_values(i+1,0,0,v) );
2425 X2.assign_data( &X_internal_vector_values(--r,0,0,v) );
2426 member.team_barrier();
2427 KB::Gemm<member_type,
2428 KB::Trans::NoTranspose,KB::Trans::NoTranspose,
2429 default_mode_type,default_algo_type>
2430 ::invoke(member, -one, A, X1, one, X2);
2432 A.assign_data( &D_internal_vector_values(i,0,0,v) );
2433 KB::Trsm<member_type,
2434 KB::Side::Left,KB::Uplo::Upper,KB::Trans::NoTranspose,KB::Diag::NonUnit,
2435 default_mode_type,default_algo_type>
2436 ::invoke(member, one, A, X2);
2437 X1.assign_data( X2.data() );
2441 auto W = Kokkos::subview(WW, Kokkos::ALL(), Kokkos::ALL(), v);
2442 KB::Copy<member_type,KB::Trans::NoTranspose,default_mode_type>
2443 ::invoke(member, X1, W);
2444 member.team_barrier();
2445 KB::Gemm<member_type,
2446 KB::Trans::NoTranspose,KB::Trans::NoTranspose,
2447 default_mode_type,default_algo_type>
2448 ::invoke(member, one, A, W, zero, X1);
2453 template<
typename impl_type,
typename WWViewType,
typename XViewType>
2454 KOKKOS_INLINE_FUNCTION
2456 solveSingleVectorNew(
const typename Kokkos::TeamPolicy<typename impl_type::execution_space>::member_type &member,
2457 const typename impl_type::local_ordinal_type &blocksize,
2458 const typename impl_type::local_ordinal_type &i0,
2459 const typename impl_type::local_ordinal_type &r0,
2460 const typename impl_type::local_ordinal_type &nrows,
2461 const typename impl_type::local_ordinal_type &v,
2462 const ConstUnmanaged<typename impl_type::internal_vector_type_4d_view> D_internal_vector_values,
2463 const XViewType &X_internal_vector_values,
2464 const WWViewType &WW) {
2465 using execution_space =
typename impl_type::execution_space;
2468 using local_ordinal_type =
typename impl_type::local_ordinal_type;
2470 typedef SolveTridiagsDefaultModeAndAlgo
2471 <
typename execution_space::memory_space> default_mode_and_algo_type;
2473 typedef typename default_mode_and_algo_type::mode_type default_mode_type;
2474 typedef typename default_mode_and_algo_type::single_vector_algo_type default_algo_type;
2476 using btdm_magnitude_type =
typename impl_type::btdm_magnitude_type;
2479 auto A = D_internal_vector_values.data();
2480 auto X = X_internal_vector_values.data();
2483 const auto one = Kokkos::ArithTraits<btdm_magnitude_type>::one();
2484 const auto zero = Kokkos::ArithTraits<btdm_magnitude_type>::zero();
2488 const local_ordinal_type astep = D_internal_vector_values.stride_0();
2489 const local_ordinal_type as0 = D_internal_vector_values.stride_1();
2490 const local_ordinal_type as1 = D_internal_vector_values.stride_2();
2491 const local_ordinal_type xstep = X_internal_vector_values.stride_0();
2492 const local_ordinal_type xs0 = X_internal_vector_values.stride_1();
2501 KOKKOSBATCHED_TRSV_LOWER_NO_TRANSPOSE_INTERNAL_INVOKE
2502 (default_mode_type,default_algo_type,
2505 blocksize,blocksize,
2510 for (local_ordinal_type tr=1;tr<nrows;++tr) {
2511 member.team_barrier();
2512 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
2513 (default_mode_type,default_algo_type,
2515 blocksize, blocksize,
2517 A+2*astep, as0, as1,
2521 KOKKOSBATCHED_TRSV_LOWER_NO_TRANSPOSE_INTERNAL_INVOKE
2522 (default_mode_type,default_algo_type,
2525 blocksize,blocksize,
2527 A+3*astep, as0, as1,
2535 KOKKOSBATCHED_TRSV_UPPER_NO_TRANSPOSE_INTERNAL_INVOKE
2536 (default_mode_type,default_algo_type,
2539 blocksize, blocksize,
2544 for (local_ordinal_type tr=nrows;tr>1;--tr) {
2546 member.team_barrier();
2547 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
2548 (default_mode_type,default_algo_type,
2550 blocksize, blocksize,
2552 A+1*astep, as0, as1,
2556 KOKKOSBATCHED_TRSV_UPPER_NO_TRANSPOSE_INTERNAL_INVOKE
2557 (default_mode_type,default_algo_type,
2560 blocksize, blocksize,
2569 const local_ordinal_type ws0 = WW.stride_0();
2570 auto W = WW.data() + v;
2571 KOKKOSBATCHED_COPY_VECTOR_NO_TRANSPOSE_INTERNAL_INVOKE
2573 member, blocksize, X, xs0, W, ws0);
2574 member.team_barrier();
2575 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
2576 (default_mode_type,default_algo_type,
2578 blocksize, blocksize,
2587 template<
typename local_ordinal_type,
typename ViewType>
2588 void writeBTDValuesToFile (
const local_ordinal_type &n_parts,
const ViewType &scalar_values_device, std::string fileName) {
2589 #ifdef IFPACK2_BLOCKTRIDICONTAINER_WRITE_MM
2590 auto scalar_values = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), scalar_values_device);
2591 std::ofstream myfile;
2592 myfile.open (fileName);
2594 const local_ordinal_type n_parts_per_pack = n_parts < (local_ordinal_type) scalar_values.extent(3) ? n_parts : scalar_values.extent(3);
2595 local_ordinal_type nnz = scalar_values.extent(0) * scalar_values.extent(1) * scalar_values.extent(2) * n_parts_per_pack;
2596 const local_ordinal_type n_blocks = scalar_values.extent(0)*n_parts_per_pack;
2597 const local_ordinal_type n_blocks_per_part = n_blocks/n_parts;
2599 const local_ordinal_type block_size = scalar_values.extent(1);
2601 const local_ordinal_type n_rows_per_part = (n_blocks_per_part+2)/3 * block_size;
2602 const local_ordinal_type n_rows = n_rows_per_part*n_parts;
2604 const local_ordinal_type n_packs = ceil(
float(n_parts)/n_parts_per_pack);
2606 myfile <<
"%%MatrixMarket matrix coordinate real general"<< std::endl;
2607 myfile <<
"%%nnz = " << nnz;
2608 myfile <<
" block size = " << block_size;
2609 myfile <<
" number of blocks = " << n_blocks;
2610 myfile <<
" number of parts = " << n_parts;
2611 myfile <<
" number of blocks per part = " << n_blocks_per_part;
2612 myfile <<
" number of rows = " << n_rows ;
2613 myfile <<
" number of cols = " << n_rows;
2614 myfile <<
" number of packs = " << n_packs << std::endl;
2616 myfile << n_rows <<
" " << n_rows <<
" " << nnz << std::setprecision(9) << std::endl;
2618 local_ordinal_type current_part_idx, current_block_idx, current_row_offset, current_col_offset, current_row, current_col;
2619 for (local_ordinal_type i_pack=0;i_pack<n_packs;++i_pack) {
2620 for (local_ordinal_type i_part_in_pack=0;i_part_in_pack<n_parts_per_pack;++i_part_in_pack) {
2621 current_part_idx = i_part_in_pack + i_pack * n_parts_per_pack;
2622 for (local_ordinal_type i_block_in_part=0;i_block_in_part<n_blocks_per_part;++i_block_in_part) {
2623 current_block_idx = i_block_in_part + i_pack * n_blocks_per_part;
2624 if (current_block_idx >= (local_ordinal_type) scalar_values.extent(0))
2626 if (i_block_in_part % 3 == 0) {
2627 current_row_offset = i_block_in_part/3 * block_size;
2628 current_col_offset = i_block_in_part/3 * block_size;
2630 else if (i_block_in_part % 3 == 1) {
2631 current_row_offset = (i_block_in_part-1)/3 * block_size;
2632 current_col_offset = ((i_block_in_part-1)/3+1) * block_size;
2634 else if (i_block_in_part % 3 == 2) {
2635 current_row_offset = ((i_block_in_part-2)/3+1) * block_size;
2636 current_col_offset = (i_block_in_part-2)/3 * block_size;
2638 current_row_offset += current_part_idx * n_rows_per_part;
2639 current_col_offset += current_part_idx * n_rows_per_part;
2640 for (local_ordinal_type i_in_block=0;i_in_block<block_size;++i_in_block) {
2641 for (local_ordinal_type j_in_block=0;j_in_block<block_size;++j_in_block) {
2642 current_row = current_row_offset + i_in_block + 1;
2643 current_col = current_col_offset + j_in_block + 1;
2644 myfile << current_row <<
" " << current_col <<
" " << scalar_values(current_block_idx,i_in_block,j_in_block,i_part_in_pack) << std::endl;
2655 template<
typename local_ordinal_type,
typename ViewType>
2656 void write4DMultiVectorValuesToFile (
const local_ordinal_type &n_parts,
const ViewType &scalar_values_device, std::string fileName) {
2657 #ifdef IFPACK2_BLOCKTRIDICONTAINER_WRITE_MM
2658 auto scalar_values = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), scalar_values_device);
2659 std::ofstream myfile;
2660 myfile.open (fileName);
2662 const local_ordinal_type n_parts_per_pack = n_parts < scalar_values.extent(3) ? n_parts : scalar_values.extent(3);
2663 const local_ordinal_type n_blocks = scalar_values.extent(0)*n_parts_per_pack;
2664 const local_ordinal_type n_blocks_per_part = n_blocks/n_parts;
2666 const local_ordinal_type block_size = scalar_values.extent(1);
2667 const local_ordinal_type n_cols = scalar_values.extent(2);
2669 const local_ordinal_type n_rows_per_part = n_blocks_per_part * block_size;
2670 const local_ordinal_type n_rows = n_rows_per_part*n_parts;
2672 const local_ordinal_type n_packs = ceil(
float(n_parts)/n_parts_per_pack);
2675 myfile <<
"%%MatrixMarket matrix array real general"<< std::endl;
2676 myfile <<
"%%block size = " << block_size;
2677 myfile <<
" number of blocks = " << n_blocks;
2678 myfile <<
" number of parts = " << n_parts;
2679 myfile <<
" number of blocks per part = " << n_blocks_per_part;
2680 myfile <<
" number of rows = " << n_rows ;
2681 myfile <<
" number of cols = " << n_cols;
2682 myfile <<
" number of packs = " << n_packs << std::endl;
2684 myfile << n_rows <<
" " << n_cols << std::setprecision(9) << std::endl;
2686 local_ordinal_type current_part_idx, current_block_idx, current_row_offset;
2687 (void) current_row_offset;
2688 (void) current_part_idx;
2689 for (local_ordinal_type j_in_block=0;j_in_block<n_cols;++j_in_block) {
2690 for (local_ordinal_type i_pack=0;i_pack<n_packs;++i_pack) {
2691 for (local_ordinal_type i_part_in_pack=0;i_part_in_pack<n_parts_per_pack;++i_part_in_pack) {
2692 current_part_idx = i_part_in_pack + i_pack * n_parts_per_pack;
2693 for (local_ordinal_type i_block_in_part=0;i_block_in_part<n_blocks_per_part;++i_block_in_part) {
2694 current_block_idx = i_block_in_part + i_pack * n_blocks_per_part;
2696 if (current_block_idx >= (local_ordinal_type) scalar_values.extent(0))
2698 for (local_ordinal_type i_in_block=0;i_in_block<block_size;++i_in_block) {
2699 myfile << scalar_values(current_block_idx,i_in_block,j_in_block,i_part_in_pack) << std::endl;
2709 template<
typename local_ordinal_type,
typename ViewType>
2710 void write5DMultiVectorValuesToFile (
const local_ordinal_type &n_parts,
const ViewType &scalar_values_device, std::string fileName) {
2711 #ifdef IFPACK2_BLOCKTRIDICONTAINER_WRITE_MM
2712 auto scalar_values = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), scalar_values_device);
2713 std::ofstream myfile;
2714 myfile.open (fileName);
2716 const local_ordinal_type n_parts_per_pack = n_parts < scalar_values.extent(4) ? n_parts : scalar_values.extent(4);
2717 const local_ordinal_type n_blocks = scalar_values.extent(1)*n_parts_per_pack;
2718 const local_ordinal_type n_blocks_per_part = n_blocks/n_parts;
2720 const local_ordinal_type block_size = scalar_values.extent(2);
2721 const local_ordinal_type n_blocks_cols = scalar_values.extent(0);
2722 const local_ordinal_type n_cols = n_blocks_cols * block_size;
2724 const local_ordinal_type n_rows_per_part = n_blocks_per_part * block_size;
2725 const local_ordinal_type n_rows = n_rows_per_part*n_parts;
2727 const local_ordinal_type n_packs = ceil(
float(n_parts)/n_parts_per_pack);
2729 myfile <<
"%%MatrixMarket matrix array real general"<< std::endl;
2730 myfile <<
"%%block size = " << block_size;
2731 myfile <<
" number of blocks = " << n_blocks;
2732 myfile <<
" number of parts = " << n_parts;
2733 myfile <<
" number of blocks per part = " << n_blocks_per_part;
2734 myfile <<
" number of rows = " << n_rows ;
2735 myfile <<
" number of cols = " << n_cols;
2736 myfile <<
" number of packs = " << n_packs << std::endl;
2738 myfile << n_rows <<
" " << n_cols << std::setprecision(9) << std::endl;
2740 local_ordinal_type current_part_idx, current_block_idx, current_row_offset;
2741 (void) current_row_offset;
2742 (void) current_part_idx;
2743 for (local_ordinal_type i_block_col=0;i_block_col<n_blocks_cols;++i_block_col) {
2744 for (local_ordinal_type j_in_block=0;j_in_block<block_size;++j_in_block) {
2745 for (local_ordinal_type i_pack=0;i_pack<n_packs;++i_pack) {
2746 for (local_ordinal_type i_part_in_pack=0;i_part_in_pack<n_parts_per_pack;++i_part_in_pack) {
2747 current_part_idx = i_part_in_pack + i_pack * n_parts_per_pack;
2748 for (local_ordinal_type i_block_in_part=0;i_block_in_part<n_blocks_per_part;++i_block_in_part) {
2749 current_block_idx = i_block_in_part + i_pack * n_blocks_per_part;
2751 if (current_block_idx >= (local_ordinal_type) scalar_values.extent(1))
2753 for (local_ordinal_type i_in_block=0;i_in_block<block_size;++i_in_block) {
2754 myfile << scalar_values(i_block_col,current_block_idx,i_in_block,j_in_block,i_part_in_pack) << std::endl;
2765 template<
typename local_ordinal_type,
typename member_type,
typename ViewType1,
typename ViewType2>
2766 KOKKOS_INLINE_FUNCTION
2768 copy3DView(
const member_type &member,
const ViewType1 &view1,
const ViewType2 &view2) {
2781 Kokkos::Experimental::local_deep_copy(member, view1, view2);
2783 template<
typename MatrixType,
int ScratchLevel>
2784 struct ExtractAndFactorizeTridiags {
2786 using impl_type = BlockHelperDetails::ImplType<MatrixType>;
2788 using execution_space =
typename impl_type::execution_space;
2789 using memory_space =
typename impl_type::memory_space;
2791 using local_ordinal_type =
typename impl_type::local_ordinal_type;
2794 using magnitude_type =
typename impl_type::magnitude_type;
2796 using row_matrix_type =
typename impl_type::tpetra_row_matrix_type;
2797 using crs_graph_type =
typename impl_type::tpetra_crs_graph_type;
2799 using local_ordinal_type_1d_view =
typename impl_type::local_ordinal_type_1d_view;
2800 using local_ordinal_type_2d_view =
typename impl_type::local_ordinal_type_2d_view;
2802 using size_type_2d_view =
typename impl_type::size_type_2d_view;
2803 using impl_scalar_type_1d_view_tpetra =
typename impl_type::impl_scalar_type_1d_view_tpetra;
2805 using btdm_scalar_type =
typename impl_type::btdm_scalar_type;
2806 using btdm_magnitude_type =
typename impl_type::btdm_magnitude_type;
2807 using vector_type_3d_view =
typename impl_type::vector_type_3d_view;
2808 using vector_type_4d_view =
typename impl_type::vector_type_4d_view;
2809 using internal_vector_type_4d_view =
typename impl_type::internal_vector_type_4d_view;
2810 using internal_vector_type_5d_view =
typename impl_type::internal_vector_type_5d_view;
2811 using btdm_scalar_type_2d_view =
typename impl_type::btdm_scalar_type_2d_view;
2812 using btdm_scalar_type_3d_view =
typename impl_type::btdm_scalar_type_3d_view;
2813 using btdm_scalar_type_4d_view =
typename impl_type::btdm_scalar_type_4d_view;
2814 using btdm_scalar_type_5d_view =
typename impl_type::btdm_scalar_type_5d_view;
2815 using internal_vector_scratch_type_3d_view = Scratch<typename impl_type::internal_vector_type_3d_view>;
2816 using btdm_scalar_scratch_type_3d_view = Scratch<typename impl_type::btdm_scalar_type_3d_view>;
2817 using tpetra_block_access_view_type =
typename impl_type::tpetra_block_access_view_type;
2818 using local_crs_graph_type =
typename impl_type::local_crs_graph_type;
2819 using colinds_view =
typename local_crs_graph_type::entries_type;
2821 using internal_vector_type =
typename impl_type::internal_vector_type;
2822 static constexpr
int vector_length = impl_type::vector_length;
2823 static constexpr
int internal_vector_length = impl_type::internal_vector_length;
2824 static_assert(vector_length >= internal_vector_length,
"Ifpack2 BlockTriDi Numeric: vector_length must be at least as large as internal_vector_length");
2825 static_assert(vector_length % internal_vector_length == 0,
"Ifpack2 BlockTriDi Numeric: vector_length must be divisible by internal_vector_length");
2830 static constexpr
int half_vector_length = impl_type::half_vector_length;
2833 using team_policy_type = Kokkos::TeamPolicy<execution_space>;
2834 using member_type =
typename team_policy_type::member_type;
2838 const ConstUnmanaged<local_ordinal_type_1d_view> partptr, lclrow, packptr, packindices_sub, packptr_sub;
2839 const ConstUnmanaged<local_ordinal_type_2d_view> partptr_sub, part2packrowidx0_sub, packindices_schur;
2840 const local_ordinal_type max_partsz;
2842 using size_type_1d_view_tpetra = Kokkos::View<size_t*,typename impl_type::node_device_type>;
2843 ConstUnmanaged<size_type_1d_view_tpetra> A_block_rowptr;
2844 ConstUnmanaged<size_type_1d_view_tpetra> A_point_rowptr;
2845 ConstUnmanaged<impl_scalar_type_1d_view_tpetra> A_values;
2847 const ConstUnmanaged<size_type_2d_view> pack_td_ptr, flat_td_ptr, pack_td_ptr_schur;
2848 const ConstUnmanaged<local_ordinal_type_1d_view> A_colindsub;
2849 const Unmanaged<internal_vector_type_4d_view> internal_vector_values, internal_vector_values_schur;
2850 const Unmanaged<internal_vector_type_5d_view> e_internal_vector_values;
2851 const Unmanaged<btdm_scalar_type_4d_view> scalar_values, scalar_values_schur;
2852 const Unmanaged<btdm_scalar_type_5d_view> e_scalar_values;
2853 const Unmanaged<btdm_scalar_type_3d_view> d_inv;
2854 const Unmanaged<size_type_1d_view> diag_offsets;
2856 const local_ordinal_type blocksize, blocksize_square;
2858 const magnitude_type tiny;
2859 const local_ordinal_type vector_loop_size;
2861 bool hasBlockCrsMatrix;
2864 ExtractAndFactorizeTridiags(
const BlockTridiags<MatrixType> &btdm_,
2865 const BlockHelperDetails::PartInterface<MatrixType> &interf_,
2868 const magnitude_type& tiny_) :
2870 partptr(interf_.partptr),
2871 lclrow(interf_.lclrow),
2872 packptr(interf_.packptr),
2873 packindices_sub(interf_.packindices_sub),
2874 packptr_sub(interf_.packptr_sub),
2875 partptr_sub(interf_.partptr_sub),
2876 part2packrowidx0_sub(interf_.part2packrowidx0_sub),
2877 packindices_schur(interf_.packindices_schur),
2878 max_partsz(interf_.max_partsz),
2880 pack_td_ptr(btdm_.pack_td_ptr),
2881 flat_td_ptr(btdm_.flat_td_ptr),
2882 pack_td_ptr_schur(btdm_.pack_td_ptr_schur),
2883 A_colindsub(btdm_.A_colindsub),
2884 internal_vector_values((internal_vector_type*)btdm_.values.data(),
2885 btdm_.values.extent(0),
2886 btdm_.values.extent(1),
2887 btdm_.values.extent(2),
2888 vector_length/internal_vector_length),
2889 internal_vector_values_schur((internal_vector_type*)btdm_.values_schur.data(),
2890 btdm_.values_schur.extent(0),
2891 btdm_.values_schur.extent(1),
2892 btdm_.values_schur.extent(2),
2893 vector_length/internal_vector_length),
2894 e_internal_vector_values((internal_vector_type*)btdm_.e_values.data(),
2895 btdm_.e_values.extent(0),
2896 btdm_.e_values.extent(1),
2897 btdm_.e_values.extent(2),
2898 btdm_.e_values.extent(3),
2899 vector_length/internal_vector_length),
2900 scalar_values((btdm_scalar_type*)btdm_.values.data(),
2901 btdm_.values.extent(0),
2902 btdm_.values.extent(1),
2903 btdm_.values.extent(2),
2905 scalar_values_schur((btdm_scalar_type*)btdm_.values_schur.data(),
2906 btdm_.values_schur.extent(0),
2907 btdm_.values_schur.extent(1),
2908 btdm_.values_schur.extent(2),
2910 e_scalar_values((btdm_scalar_type*)btdm_.e_values.data(),
2911 btdm_.e_values.extent(0),
2912 btdm_.e_values.extent(1),
2913 btdm_.e_values.extent(2),
2914 btdm_.e_values.extent(3),
2917 diag_offsets(btdm_.diag_offsets),
2918 blocksize(btdm_.values.extent(1)),
2919 blocksize_square(blocksize*blocksize),
2922 vector_loop_size(vector_length/internal_vector_length) {
2923 using crs_matrix_type =
typename impl_type::tpetra_crs_matrix_type;
2924 using block_crs_matrix_type =
typename impl_type::tpetra_block_crs_matrix_type;
2926 auto A_crs = Teuchos::rcp_dynamic_cast<
const crs_matrix_type>(A_);
2927 auto A_bcrs = Teuchos::rcp_dynamic_cast<
const block_crs_matrix_type>(A_);
2929 hasBlockCrsMatrix = ! A_bcrs.is_null ();
2931 A_block_rowptr = G_->getLocalGraphDevice().row_map;
2932 if (hasBlockCrsMatrix) {
2933 A_values =
const_cast<block_crs_matrix_type*
>(A_bcrs.get())->getValuesDeviceNonConst();
2936 A_point_rowptr = A_crs->getCrsGraph()->getLocalGraphDevice().row_map;
2937 A_values = A_crs->getLocalValuesDevice (Tpetra::Access::ReadOnly);
2943 KOKKOS_INLINE_FUNCTION
2945 extract(local_ordinal_type partidx,
2946 local_ordinal_type local_subpartidx,
2947 local_ordinal_type npacks)
const {
2948 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
2949 printf(
"extract partidx = %d, local_subpartidx = %d, npacks = %d;\n", partidx, local_subpartidx, npacks);
2951 using tlb = BlockHelperDetails::TpetraLittleBlock<Tpetra::Impl::BlockCrsMatrixLittleBlockArrayLayout>;
2952 const size_type kps = pack_td_ptr(partidx, local_subpartidx);
2953 local_ordinal_type kfs[vector_length] = {};
2954 local_ordinal_type ri0[vector_length] = {};
2955 local_ordinal_type nrows[vector_length] = {};
2957 for (local_ordinal_type vi=0;vi<npacks;++vi,++partidx) {
2958 kfs[vi] = flat_td_ptr(partidx,local_subpartidx);
2959 ri0[vi] = partptr_sub(pack_td_ptr.extent(0)*local_subpartidx + partidx,0);
2960 nrows[vi] = partptr_sub(pack_td_ptr.extent(0)*local_subpartidx + partidx,1) - ri0[vi];
2961 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
2962 printf(
"kfs[%d] = %d;\n", vi, kfs[vi]);
2963 printf(
"ri0[%d] = %d;\n", vi, ri0[vi]);
2964 printf(
"nrows[%d] = %d;\n", vi, nrows[vi]);
2967 local_ordinal_type tr_min = 0;
2968 local_ordinal_type tr_max = nrows[0];
2969 if (local_subpartidx % 2 == 1) {
2973 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
2974 printf(
"tr_min = %d and tr_max = %d;\n", tr_min, tr_max);
2976 for (local_ordinal_type tr=tr_min,j=0;tr<tr_max;++tr) {
2977 for (local_ordinal_type e=0;e<3;++e) {
2978 if (hasBlockCrsMatrix) {
2979 const impl_scalar_type* block[vector_length] = {};
2980 for (local_ordinal_type vi=0;vi<npacks;++vi) {
2981 const size_type Aj = A_block_rowptr(lclrow(ri0[vi] + tr)) + A_colindsub(kfs[vi] + j);
2983 block[vi] = &A_values(Aj*blocksize_square);
2985 const size_type pi = kps + j;
2986 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
2987 printf(
"Extract pi = %ld, ri0 + tr = %d, kfs + j = %d\n", pi, ri0[0] + tr, kfs[0] + j);
2990 for (local_ordinal_type ii=0;ii<blocksize;++ii) {
2991 for (local_ordinal_type jj=0;jj<blocksize;++jj) {
2992 const auto idx = tlb::getFlatIndex(ii, jj, blocksize);
2993 auto& v = internal_vector_values(pi, ii, jj, 0);
2994 for (local_ordinal_type vi=0;vi<npacks;++vi) {
2995 v[vi] =
static_cast<btdm_scalar_type
>(block[vi][idx]);
3001 const size_type pi = kps + j;
3003 for (local_ordinal_type vi=0;vi<npacks;++vi) {
3004 const size_type Aj_c = A_colindsub(kfs[vi] + j);
3006 for (local_ordinal_type ii=0;ii<blocksize;++ii) {
3007 auto point_row_offset = A_point_rowptr(lclrow(ri0[vi] + tr)*blocksize + ii);
3009 for (local_ordinal_type jj=0;jj<blocksize;++jj) {
3010 scalar_values(pi, ii, jj, vi) = A_values(point_row_offset + Aj_c*blocksize + jj);
3016 if (nrows[0] == 1)
break;
3017 if (local_subpartidx % 2 == 0) {
3018 if (e == 1 && (tr == 0 || tr+1 == nrows[0]))
break;
3019 for (local_ordinal_type vi=1;vi<npacks;++vi) {
3020 if ((e == 0 && nrows[vi] == 1) || (e == 1 && tr+1 == nrows[vi])) {
3027 if (e == 0 && (tr == -1 || tr == nrows[0]))
break;
3028 for (local_ordinal_type vi=1;vi<npacks;++vi) {
3029 if ((e == 0 && nrows[vi] == 1) || (e == 0 && tr == nrows[vi])) {
3039 KOKKOS_INLINE_FUNCTION
3041 extract(
const member_type &member,
3042 const local_ordinal_type &partidxbeg,
3043 local_ordinal_type local_subpartidx,
3044 const local_ordinal_type &npacks,
3045 const local_ordinal_type &vbeg)
const {
3046 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3047 printf(
"extract partidxbeg = %d, local_subpartidx = %d, npacks = %d, vbeg = %d;\n", partidxbeg, local_subpartidx, npacks, vbeg);
3049 using tlb = BlockHelperDetails::TpetraLittleBlock<Tpetra::Impl::BlockCrsMatrixLittleBlockArrayLayout>;
3050 local_ordinal_type kfs_vals[internal_vector_length] = {};
3051 local_ordinal_type ri0_vals[internal_vector_length] = {};
3052 local_ordinal_type nrows_vals[internal_vector_length] = {};
3054 const size_type kps = pack_td_ptr(partidxbeg,local_subpartidx);
3055 for (local_ordinal_type v=vbeg,vi=0;v<npacks && vi<internal_vector_length;++v,++vi) {
3056 kfs_vals[vi] = flat_td_ptr(partidxbeg+vi,local_subpartidx);
3057 ri0_vals[vi] = partptr_sub(pack_td_ptr.extent(0)*local_subpartidx + partidxbeg+vi,0);
3058 nrows_vals[vi] = partptr_sub(pack_td_ptr.extent(0)*local_subpartidx + partidxbeg+vi,1) - ri0_vals[vi];
3059 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3060 printf(
"kfs_vals[%d] = %d;\n", vi, kfs_vals[vi]);
3061 printf(
"ri0_vals[%d] = %d;\n", vi, ri0_vals[vi]);
3062 printf(
"nrows_vals[%d] = %d;\n", vi, nrows_vals[vi]);
3066 local_ordinal_type j_vals[internal_vector_length] = {};
3068 local_ordinal_type tr_min = 0;
3069 local_ordinal_type tr_max = nrows_vals[0];
3070 if (local_subpartidx % 2 == 1) {
3074 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3075 printf(
"tr_min = %d and tr_max = %d;\n", tr_min, tr_max);
3077 for (local_ordinal_type tr=tr_min;tr<tr_max;++tr) {
3078 for (local_ordinal_type v=vbeg,vi=0;v<npacks && vi<internal_vector_length;++v,++vi) {
3079 const local_ordinal_type nrows = (local_subpartidx % 2 == 0 ? nrows_vals[vi] : nrows_vals[vi]);
3080 if ((local_subpartidx % 2 == 0 && tr < nrows) || (local_subpartidx % 2 == 1 && tr < nrows+1)) {
3081 auto &j = j_vals[vi];
3082 const local_ordinal_type kfs = kfs_vals[vi];
3083 const local_ordinal_type ri0 = ri0_vals[vi];
3084 local_ordinal_type lbeg, lend;
3085 if (local_subpartidx % 2 == 0) {
3086 lbeg = (tr == tr_min ? 1 : 0);
3087 lend = (tr == nrows - 1 ? 2 : 3);
3096 else if (tr == nrows) {
3101 if (hasBlockCrsMatrix) {
3102 for (local_ordinal_type l=lbeg;l<lend;++l,++j) {
3103 const size_type Aj = A_block_rowptr(lclrow(ri0 + tr)) + A_colindsub(kfs + j);
3104 const impl_scalar_type* block = &A_values(Aj*blocksize_square);
3105 const size_type pi = kps + j;
3106 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3107 printf(
"Extract pi = %ld, ri0 + tr = %d, kfs + j = %d, tr = %d, lbeg = %d, lend = %d, l = %d\n", pi, ri0 + tr, kfs + j, tr, lbeg, lend, l);
3109 Kokkos::parallel_for
3110 (Kokkos::TeamThreadRange(member,blocksize),
3111 [&](
const local_ordinal_type &ii) {
3112 for (local_ordinal_type jj=0;jj<blocksize;++jj) {
3113 scalar_values(pi, ii, jj, v) =
static_cast<btdm_scalar_type
>(block[tlb::getFlatIndex(ii,jj,blocksize)]);
3119 for (local_ordinal_type l=lbeg;l<lend;++l,++j) {
3120 const size_type Aj_c = A_colindsub(kfs + j);
3121 const size_type pi = kps + j;
3122 Kokkos::parallel_for
3123 (Kokkos::TeamThreadRange(member,blocksize),
3124 [&](
const local_ordinal_type &ii) {
3125 auto point_row_offset = A_point_rowptr(lclrow(ri0 + tr)*blocksize + ii);
3126 for (local_ordinal_type jj=0;jj<blocksize;++jj) {
3127 scalar_values(pi, ii, jj, v) = A_values(point_row_offset + Aj_c*blocksize + jj);
3137 template<
typename AAViewType,
3138 typename WWViewType>
3139 KOKKOS_INLINE_FUNCTION
3141 factorize_subline(
const member_type &member,
3142 const local_ordinal_type &i0,
3143 const local_ordinal_type &nrows,
3144 const local_ordinal_type &v,
3145 const AAViewType &AA,
3146 const WWViewType &WW)
const {
3148 typedef ExtractAndFactorizeTridiagsDefaultModeAndAlgo
3149 <
typename execution_space::memory_space> default_mode_and_algo_type;
3151 typedef typename default_mode_and_algo_type::mode_type default_mode_type;
3152 typedef typename default_mode_and_algo_type::algo_type default_algo_type;
3155 const auto one = Kokkos::ArithTraits<btdm_magnitude_type>::one();
3157 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3158 printf(
"i0 = %d, nrows = %d, v = %d, AA.extent(0) = %ld;\n", i0, nrows, v, AA.extent(0));
3162 auto A = Kokkos::subview(AA, i0, Kokkos::ALL(), Kokkos::ALL(), v);
3164 default_mode_type,KB::Algo::LU::Unblocked>
3165 ::invoke(member, A , tiny);
3170 local_ordinal_type i = i0;
3171 for (local_ordinal_type tr=1;tr<nrows;++tr,i+=3) {
3172 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3173 printf(
"tr = %d, i = %d;\n", tr, i);
3175 B.assign_data( &AA(i+1,0,0,v) );
3176 KB::Trsm<member_type,
3177 KB::Side::Left,KB::Uplo::Lower,KB::Trans::NoTranspose,KB::Diag::Unit,
3178 default_mode_type,default_algo_type>
3179 ::invoke(member, one, A, B);
3180 C.assign_data( &AA(i+2,0,0,v) );
3181 KB::Trsm<member_type,
3182 KB::Side::Right,KB::Uplo::Upper,KB::Trans::NoTranspose,KB::Diag::NonUnit,
3183 default_mode_type,default_algo_type>
3184 ::invoke(member, one, A, C);
3185 A.assign_data( &AA(i+3,0,0,v) );
3187 member.team_barrier();
3188 KB::Gemm<member_type,
3189 KB::Trans::NoTranspose,KB::Trans::NoTranspose,
3190 default_mode_type,default_algo_type>
3191 ::invoke(member, -one, C, B, one, A);
3193 default_mode_type,KB::Algo::LU::Unblocked>
3194 ::invoke(member, A, tiny);
3198 auto W = Kokkos::subview(WW, Kokkos::ALL(), Kokkos::ALL(), v);
3199 KB::Copy<member_type,KB::Trans::NoTranspose,default_mode_type>
3200 ::invoke(member, A, W);
3201 KB::SetIdentity<member_type,default_mode_type>
3202 ::invoke(member, A);
3203 member.team_barrier();
3204 KB::Trsm<member_type,
3205 KB::Side::Left,KB::Uplo::Lower,KB::Trans::NoTranspose,KB::Diag::Unit,
3206 default_mode_type,default_algo_type>
3207 ::invoke(member, one, W, A);
3208 KB::Trsm<member_type,
3209 KB::Side::Left,KB::Uplo::Upper,KB::Trans::NoTranspose,KB::Diag::NonUnit,
3210 default_mode_type,default_algo_type>
3211 ::invoke(member, one, W, A);
3217 struct ExtractAndFactorizeSubLineTag {};
3218 struct ExtractAndFactorizeFusedJacobiTag {};
3219 struct ExtractBCDTag {};
3220 struct ComputeETag {};
3221 struct ComputeSchurTag {};
3222 struct FactorizeSchurTag {};
3224 KOKKOS_INLINE_FUNCTION
3226 operator() (
const ExtractAndFactorizeSubLineTag &,
const member_type &member)
const {
3228 const local_ordinal_type packidx = packindices_sub(member.league_rank());
3230 const local_ordinal_type subpartidx = packptr_sub(packidx);
3231 const local_ordinal_type n_parts = part2packrowidx0_sub.extent(0);
3232 const local_ordinal_type local_subpartidx = subpartidx/n_parts;
3233 const local_ordinal_type partidx = subpartidx%n_parts;
3235 const local_ordinal_type npacks = packptr_sub(packidx+1) - subpartidx;
3236 const local_ordinal_type i0 = pack_td_ptr(partidx,local_subpartidx);
3237 const local_ordinal_type nrows = partptr_sub(subpartidx,1) - partptr_sub(subpartidx,0);
3239 internal_vector_scratch_type_3d_view
3240 WW(member.team_scratch(ScratchLevel), blocksize, blocksize, vector_loop_size);
3242 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3243 printf(
"rank = %d, i0 = %d, npacks = %d, nrows = %d, packidx = %d, subpartidx = %d, partidx = %d, local_subpartidx = %d;\n", member.league_rank(), i0, npacks, nrows, packidx, subpartidx, partidx, local_subpartidx);
3244 printf(
"vector_loop_size = %d\n", vector_loop_size);
3247 if (vector_loop_size == 1) {
3248 extract(partidx, local_subpartidx, npacks);
3249 factorize_subline(member, i0, nrows, 0, internal_vector_values, WW);
3251 Kokkos::parallel_for
3252 (Kokkos::ThreadVectorRange(member, vector_loop_size),
3253 [&](
const local_ordinal_type &v) {
3254 const local_ordinal_type vbeg = v*internal_vector_length;
3255 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3256 printf(
"i0 = %d, npacks = %d, vbeg = %d;\n", i0, npacks, vbeg);
3259 extract(member, partidx+vbeg, local_subpartidx, npacks, vbeg);
3262 member.team_barrier();
3263 factorize_subline(member, i0, nrows, v, internal_vector_values, WW);
3268 KOKKOS_INLINE_FUNCTION
3270 operator() (
const ExtractAndFactorizeFusedJacobiTag&,
const member_type &member)
const {
3271 using default_mode_and_algo_type = ExtractAndFactorizeTridiagsDefaultModeAndAlgo<typename execution_space::memory_space>;
3272 using default_mode_type =
typename default_mode_and_algo_type::mode_type;
3273 using default_algo_type =
typename default_mode_and_algo_type::algo_type;
3276 btdm_scalar_scratch_type_3d_view WW1(member.team_scratch(ScratchLevel), half_vector_length, blocksize, blocksize);
3277 btdm_scalar_scratch_type_3d_view WW2(member.team_scratch(ScratchLevel), half_vector_length, blocksize, blocksize);
3278 const auto one = Kokkos::ArithTraits<btdm_magnitude_type>::one();
3279 const local_ordinal_type nrows = lclrow.extent(0);
3280 Kokkos::parallel_for
3281 (Kokkos::ThreadVectorRange(member, half_vector_length),
3282 [&](
const local_ordinal_type &v) {
3283 local_ordinal_type row = member.league_rank() * half_vector_length + v;
3285 auto W1 = Kokkos::subview(WW1, v, Kokkos::ALL(), Kokkos::ALL());
3286 auto W2 = Kokkos::subview(WW2, v, Kokkos::ALL(), Kokkos::ALL());
3289 const impl_scalar_type* A_diag = A_values.data() + diag_offsets(row);
3292 Kokkos::parallel_for(Kokkos::TeamThreadRange(member, blocksize * blocksize),
3295 W1.data()[i] = A_diag[i];
3298 KB::SetIdentity<member_type,default_mode_type>
3299 ::invoke(member, W2);
3305 KB::SetIdentity<member_type,default_mode_type>
3306 ::invoke(member, W1);
3308 member.team_barrier();
3310 KB::LU<member_type, default_mode_type,KB::Algo::LU::Unblocked>
3311 ::invoke(member, W1, tiny);
3312 member.team_barrier();
3313 KB::Trsm<member_type,
3314 KB::Side::Left,KB::Uplo::Lower,KB::Trans::NoTranspose,KB::Diag::Unit,
3315 default_mode_type,default_algo_type>
3316 ::invoke(member, one, W1, W2);
3317 KB::Trsm<member_type,
3318 KB::Side::Left,KB::Uplo::Upper,KB::Trans::NoTranspose,KB::Diag::NonUnit,
3319 default_mode_type,default_algo_type>
3320 ::invoke(member, one, W1, W2);
3321 member.team_barrier();
3323 Kokkos::parallel_for(Kokkos::TeamThreadRange(member, blocksize * blocksize),
3326 auto d_inv_block = &d_inv(row, 0, 0);
3327 d_inv_block[i] = W2.data()[i];
3333 KOKKOS_INLINE_FUNCTION
3335 operator() (
const ExtractBCDTag &,
const member_type &member)
const {
3337 const local_ordinal_type packindices_schur_i = member.league_rank() % packindices_schur.extent(0);
3338 const local_ordinal_type packindices_schur_j = member.league_rank() / packindices_schur.extent(0);
3339 const local_ordinal_type packidx = packindices_schur(packindices_schur_i, packindices_schur_j);
3341 const local_ordinal_type subpartidx = packptr_sub(packidx);
3342 const local_ordinal_type n_parts = part2packrowidx0_sub.extent(0);
3343 const local_ordinal_type local_subpartidx = subpartidx/n_parts;
3344 const local_ordinal_type partidx = subpartidx%n_parts;
3346 const local_ordinal_type npacks = packptr_sub(packidx+1) - subpartidx;
3350 if (vector_loop_size == 1) {
3351 extract(partidx, local_subpartidx, npacks);
3354 Kokkos::parallel_for
3355 (Kokkos::ThreadVectorRange(member, vector_loop_size),
3356 [&](
const local_ordinal_type &v) {
3357 const local_ordinal_type vbeg = v*internal_vector_length;
3358 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3359 const local_ordinal_type i0 = pack_td_ptr(partidx,local_subpartidx);
3360 printf(
"i0 = %d, npacks = %d, vbeg = %d;\n", i0, npacks, vbeg);
3363 extract(member, partidx+vbeg, local_subpartidx, npacks, vbeg);
3367 member.team_barrier();
3369 const size_type kps1 = pack_td_ptr(partidx, local_subpartidx);
3370 const size_type kps2 = pack_td_ptr(partidx, local_subpartidx+1)-1;
3372 const local_ordinal_type r1 = part2packrowidx0_sub(partidx,local_subpartidx)-1;
3373 const local_ordinal_type r2 = part2packrowidx0_sub(partidx,local_subpartidx)+2;
3375 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3376 printf(
"Copy for Schur complement part id = %d from kps1 = %ld to r1 = %d and from kps2 = %ld to r2 = %d partidx = %d local_subpartidx = %d;\n", packidx, kps1, r1, kps2, r2, partidx, local_subpartidx);
3380 copy3DView<local_ordinal_type>(member, Kokkos::subview(e_internal_vector_values, 0, r1, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()),
3381 Kokkos::subview(internal_vector_values, kps1, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()));
3383 copy3DView<local_ordinal_type>(member, Kokkos::subview(e_internal_vector_values, 1, r2, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()),
3384 Kokkos::subview(internal_vector_values, kps2, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()));
3388 KOKKOS_INLINE_FUNCTION
3390 operator() (
const ComputeETag &,
const member_type &member)
const {
3392 const local_ordinal_type packidx = packindices_sub(member.league_rank());
3394 const local_ordinal_type subpartidx = packptr_sub(packidx);
3395 const local_ordinal_type n_parts = part2packrowidx0_sub.extent(0);
3396 const local_ordinal_type local_subpartidx = subpartidx/n_parts;
3397 const local_ordinal_type partidx = subpartidx%n_parts;
3399 const local_ordinal_type npacks = packptr_sub(packidx+1) - subpartidx;
3400 const local_ordinal_type i0 = pack_td_ptr(partidx,local_subpartidx);
3401 const local_ordinal_type r0 = part2packrowidx0_sub(partidx,local_subpartidx);
3402 const local_ordinal_type nrows = partptr_sub(subpartidx,1) - partptr_sub(subpartidx,0);
3403 const local_ordinal_type num_vectors = blocksize;
3407 internal_vector_scratch_type_3d_view
3408 WW(member.team_scratch(ScratchLevel), blocksize, num_vectors, vector_loop_size);
3409 if (local_subpartidx == 0) {
3410 Kokkos::parallel_for
3411 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
3412 solveMultiVector<impl_type, internal_vector_scratch_type_3d_view> (member, blocksize, i0, r0, nrows, v, internal_vector_values, Kokkos::subview(e_internal_vector_values, 0, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()), WW,
true);
3415 else if (local_subpartidx == (local_ordinal_type) part2packrowidx0_sub.extent(1) - 2) {
3416 Kokkos::parallel_for
3417 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
3418 solveMultiVector<impl_type, internal_vector_scratch_type_3d_view> (member, blocksize, i0, r0, nrows, v, internal_vector_values, Kokkos::subview(e_internal_vector_values, 1, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()), WW);
3422 Kokkos::parallel_for
3423 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
3424 solveMultiVector<impl_type, internal_vector_scratch_type_3d_view> (member, blocksize, i0, r0, nrows, v, internal_vector_values, Kokkos::subview(e_internal_vector_values, 0, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()), WW,
true);
3425 solveMultiVector<impl_type, internal_vector_scratch_type_3d_view> (member, blocksize, i0, r0, nrows, v, internal_vector_values, Kokkos::subview(e_internal_vector_values, 1, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()), WW);
3430 KOKKOS_INLINE_FUNCTION
3432 operator() (
const ComputeSchurTag &,
const member_type &member)
const {
3434 const local_ordinal_type packindices_schur_i = member.league_rank() % packindices_schur.extent(0);
3435 const local_ordinal_type packindices_schur_j = member.league_rank() / packindices_schur.extent(0);
3436 const local_ordinal_type packidx = packindices_schur(packindices_schur_i, packindices_schur_j);
3438 const local_ordinal_type subpartidx = packptr_sub(packidx);
3439 const local_ordinal_type n_parts = part2packrowidx0_sub.extent(0);
3440 const local_ordinal_type local_subpartidx = subpartidx/n_parts;
3441 const local_ordinal_type partidx = subpartidx%n_parts;
3444 const local_ordinal_type i0 = pack_td_ptr(partidx,local_subpartidx);
3450 const local_ordinal_type local_subpartidx_schur = (local_subpartidx-1)/2;
3451 const local_ordinal_type i0_schur = local_subpartidx_schur == 0 ? pack_td_ptr_schur(partidx,local_subpartidx_schur) : pack_td_ptr_schur(partidx,local_subpartidx_schur) + 1;
3452 const local_ordinal_type i0_offset = local_subpartidx_schur == 0 ? i0+2 : i0+2;
3454 for (local_ordinal_type i = 0; i < 4; ++i) {
3455 copy3DView<local_ordinal_type>(member, Kokkos::subview(internal_vector_values_schur, i0_schur+i, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()),
3456 Kokkos::subview(internal_vector_values, i0_offset+i, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()));
3459 member.team_barrier();
3461 const auto one = Kokkos::ArithTraits<btdm_magnitude_type>::one();
3463 const size_type c_kps1 = pack_td_ptr(partidx, local_subpartidx)+1;
3464 const size_type c_kps2 = pack_td_ptr(partidx, local_subpartidx+1)-2;
3466 const local_ordinal_type e_r1 = part2packrowidx0_sub(partidx,local_subpartidx)-1;
3467 const local_ordinal_type e_r2 = part2packrowidx0_sub(partidx,local_subpartidx)+2;
3469 typedef ExtractAndFactorizeTridiagsDefaultModeAndAlgo
3470 <
typename execution_space::memory_space> default_mode_and_algo_type;
3472 typedef typename default_mode_and_algo_type::mode_type default_mode_type;
3473 typedef typename default_mode_and_algo_type::algo_type default_algo_type;
3475 Kokkos::parallel_for
3476 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
3477 for (size_type i = 0; i < pack_td_ptr_schur(partidx,local_subpartidx_schur+1)-pack_td_ptr_schur(partidx,local_subpartidx_schur); ++i) {
3478 local_ordinal_type e_r, e_c, c_kps;
3480 if ( local_subpartidx_schur == 0 ) {
3486 else if ( i == 3 ) {
3491 else if ( i == 4 ) {
3506 else if ( i == 1 ) {
3511 else if ( i == 4 ) {
3516 else if ( i == 5 ) {
3526 auto S = Kokkos::subview(internal_vector_values_schur, pack_td_ptr_schur(partidx,local_subpartidx_schur)+i, Kokkos::ALL(), Kokkos::ALL(), v);
3527 auto C = Kokkos::subview(internal_vector_values, c_kps, Kokkos::ALL(), Kokkos::ALL(), v);
3528 auto E = Kokkos::subview(e_internal_vector_values, e_c, e_r, Kokkos::ALL(), Kokkos::ALL(), v);
3529 KB::Gemm<member_type,
3530 KB::Trans::NoTranspose,KB::Trans::NoTranspose,
3531 default_mode_type,default_algo_type>
3532 ::invoke(member, -one, C, E, one, S);
3537 KOKKOS_INLINE_FUNCTION
3539 operator() (
const FactorizeSchurTag &,
const member_type &member)
const {
3540 const local_ordinal_type packidx = packindices_schur(member.league_rank(), 0);
3542 const local_ordinal_type subpartidx = packptr_sub(packidx);
3544 const local_ordinal_type n_parts = part2packrowidx0_sub.extent(0);
3545 const local_ordinal_type partidx = subpartidx%n_parts;
3547 const local_ordinal_type i0 = pack_td_ptr_schur(partidx,0);
3548 const local_ordinal_type nrows = 2*(pack_td_ptr_schur.extent(1)-1);
3550 internal_vector_scratch_type_3d_view
3551 WW(member.team_scratch(ScratchLevel), blocksize, blocksize, vector_loop_size);
3553 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3554 printf(
"FactorizeSchurTag rank = %d, i0 = %d, nrows = %d, vector_loop_size = %d;\n", member.league_rank(), i0, nrows, vector_loop_size);
3557 if (vector_loop_size == 1) {
3558 factorize_subline(member, i0, nrows, 0, internal_vector_values_schur, WW);
3560 Kokkos::parallel_for
3561 (Kokkos::ThreadVectorRange(member, vector_loop_size),
3562 [&](
const local_ordinal_type &v) {
3563 factorize_subline(member, i0, nrows, v, internal_vector_values_schur, WW);
3569 IFPACK2_BLOCKTRIDICONTAINER_PROFILER_REGION_BEGIN;
3570 const local_ordinal_type team_size =
3571 ExtractAndFactorizeTridiagsDefaultModeAndAlgo<typename execution_space::memory_space>::
3572 recommended_team_size(blocksize, vector_length, internal_vector_length);
3573 const local_ordinal_type per_team_scratch = internal_vector_scratch_type_3d_view::
3574 shmem_size(blocksize, blocksize, vector_loop_size);
3577 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3578 printf(
"Start ExtractAndFactorizeSubLineTag\n");
3580 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::NumericPhase::ExtractAndFactorizeSubLineTag", ExtractAndFactorizeSubLineTag0);
3581 Kokkos::TeamPolicy<execution_space,ExtractAndFactorizeSubLineTag>
3582 policy(packindices_sub.extent(0), team_size, vector_loop_size);
3585 const local_ordinal_type n_parts = part2packrowidx0_sub.extent(0);
3586 writeBTDValuesToFile(n_parts, scalar_values,
"before.mm");
3588 policy.set_scratch_size(ScratchLevel, Kokkos::PerTeam(per_team_scratch));
3589 Kokkos::parallel_for(
"ExtractAndFactorize::TeamPolicy::run<ExtractAndFactorizeSubLineTag>",
3591 execution_space().fence();
3593 writeBTDValuesToFile(n_parts, scalar_values,
"after.mm");
3594 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3595 printf(
"End ExtractAndFactorizeSubLineTag\n");
3599 if (packindices_schur.extent(1) > 0)
3602 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3603 printf(
"Start ExtractBCDTag\n");
3605 Kokkos::deep_copy(e_scalar_values, Kokkos::ArithTraits<btdm_magnitude_type>::zero());
3606 Kokkos::deep_copy(scalar_values_schur, Kokkos::ArithTraits<btdm_magnitude_type>::zero());
3608 write5DMultiVectorValuesToFile(part2packrowidx0_sub.extent(0), e_scalar_values,
"e_scalar_values_before_extract.mm");
3611 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::NumericPhase::ExtractBCDTag", ExtractBCDTag0);
3612 Kokkos::TeamPolicy<execution_space,ExtractBCDTag>
3613 policy(packindices_schur.extent(0)*packindices_schur.extent(1), team_size, vector_loop_size);
3615 policy.set_scratch_size(ScratchLevel, Kokkos::PerTeam(per_team_scratch));
3616 Kokkos::parallel_for(
"ExtractAndFactorize::TeamPolicy::run<ExtractBCDTag>",
3618 execution_space().fence();
3621 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3622 printf(
"End ExtractBCDTag\n");
3624 writeBTDValuesToFile(part2packrowidx0_sub.extent(0), scalar_values,
"after_extraction_of_BCD.mm");
3625 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3626 printf(
"Start ComputeETag\n");
3628 write5DMultiVectorValuesToFile(part2packrowidx0_sub.extent(0), e_scalar_values,
"e_scalar_values_after_extract.mm");
3630 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::NumericPhase::ComputeETag", ComputeETag0);
3631 Kokkos::TeamPolicy<execution_space,ComputeETag>
3632 policy(packindices_sub.extent(0), team_size, vector_loop_size);
3634 policy.set_scratch_size(ScratchLevel, Kokkos::PerTeam(per_team_scratch));
3635 Kokkos::parallel_for(
"ExtractAndFactorize::TeamPolicy::run<ComputeETag>",
3637 execution_space().fence();
3639 write5DMultiVectorValuesToFile(part2packrowidx0_sub.extent(0), e_scalar_values,
"e_scalar_values_after_compute.mm");
3641 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3642 printf(
"End ComputeETag\n");
3647 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3648 printf(
"Start ComputeSchurTag\n");
3650 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::NumericPhase::ComputeSchurTag", ComputeSchurTag0);
3651 writeBTDValuesToFile(part2packrowidx0_sub.extent(0), scalar_values_schur,
"before_schur.mm");
3652 Kokkos::TeamPolicy<execution_space,ComputeSchurTag>
3653 policy(packindices_schur.extent(0)*packindices_schur.extent(1), team_size, vector_loop_size);
3655 Kokkos::parallel_for(
"ExtractAndFactorize::TeamPolicy::run<ComputeSchurTag>",
3657 writeBTDValuesToFile(part2packrowidx0_sub.extent(0), scalar_values_schur,
"after_schur.mm");
3658 execution_space().fence();
3659 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3660 printf(
"End ComputeSchurTag\n");
3665 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3666 printf(
"Start FactorizeSchurTag\n");
3668 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::NumericPhase::FactorizeSchurTag", FactorizeSchurTag0);
3669 Kokkos::TeamPolicy<execution_space,FactorizeSchurTag>
3670 policy(packindices_schur.extent(0), team_size, vector_loop_size);
3671 policy.set_scratch_size(ScratchLevel, Kokkos::PerTeam(per_team_scratch));
3672 Kokkos::parallel_for(
"ExtractAndFactorize::TeamPolicy::run<FactorizeSchurTag>",
3674 execution_space().fence();
3675 writeBTDValuesToFile(part2packrowidx0_sub.extent(0), scalar_values_schur,
"after_factor_schur.mm");
3676 #ifdef IFPACK2_BLOCKTRIDICONTAINER_USE_PRINTF
3677 printf(
"End FactorizeSchurTag\n");
3682 IFPACK2_BLOCKTRIDICONTAINER_PROFILER_REGION_END;
3685 void run_fused_jacobi() {
3686 IFPACK2_BLOCKTRIDICONTAINER_PROFILER_REGION_BEGIN;
3687 const local_ordinal_type team_size =
3688 ExtractAndFactorizeTridiagsDefaultModeAndAlgo<typename execution_space::memory_space>::
3689 recommended_team_size(blocksize, half_vector_length, 1);
3690 const local_ordinal_type per_team_scratch =
3691 btdm_scalar_scratch_type_3d_view::shmem_size(blocksize, blocksize, 2 * half_vector_length);
3693 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::NumericPhase::ExtractAndFactorizeFusedJacobi", ExtractAndFactorizeFusedJacobiTag);
3694 Kokkos::TeamPolicy<execution_space, ExtractAndFactorizeFusedJacobiTag>
3695 policy((lclrow.extent(0) + half_vector_length - 1) / half_vector_length, team_size, half_vector_length);
3697 policy.set_scratch_size(ScratchLevel, Kokkos::PerTeam(per_team_scratch));
3698 Kokkos::parallel_for(
"ExtractAndFactorize::TeamPolicy::run<ExtractAndFactorizeFusedJacobiTag>",
3701 IFPACK2_BLOCKTRIDICONTAINER_PROFILER_REGION_END;
3708 template<
typename MatrixType>
3711 const Teuchos::RCP<
const typename BlockHelperDetails::ImplType<MatrixType>::tpetra_crs_graph_type> &G,
3712 const BlockHelperDetails::PartInterface<MatrixType> &interf,
3714 const typename BlockHelperDetails::ImplType<MatrixType>::magnitude_type tiny,
3715 bool use_fused_jacobi) {
3717 using execution_space =
typename impl_type::execution_space;
3718 using team_policy_type = Kokkos::TeamPolicy<execution_space>;
3719 using internal_vector_scratch_type_3d_view = Scratch<typename impl_type::internal_vector_type_3d_view>;
3720 using btdm_scalar_scratch_type_3d_view = Scratch<typename impl_type::btdm_scalar_type_3d_view>;
3722 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::NumericPhase", NumericPhase);
3724 int blocksize = btdm.values.extent(1);
3727 int scratch_required;
3728 if(!use_fused_jacobi) {
3730 scratch_required = internal_vector_scratch_type_3d_view::shmem_size(blocksize, blocksize, impl_type::vector_length / impl_type::internal_vector_length);
3734 scratch_required = btdm_scalar_scratch_type_3d_view::shmem_size(blocksize, blocksize, 2 * impl_type::half_vector_length);
3737 int max_scratch = team_policy_type::scratch_size_max(0);
3739 if(scratch_required < max_scratch) {
3741 ExtractAndFactorizeTridiags<MatrixType, 0>
function(btdm, interf, A, G, tiny);
3742 if(!use_fused_jacobi)
3745 function.run_fused_jacobi();
3749 ExtractAndFactorizeTridiags<MatrixType, 1>
function(btdm, interf, A, G, tiny);
3750 if(!use_fused_jacobi)
3753 function.run_fused_jacobi();
3755 IFPACK2_BLOCKHELPER_TIMER_FENCE(
typename BlockHelperDetails::ImplType<MatrixType>::execution_space)
3761 template<
typename MatrixType>
3765 using execution_space =
typename impl_type::execution_space;
3766 using memory_space =
typename impl_type::memory_space;
3768 using local_ordinal_type =
typename impl_type::local_ordinal_type;
3770 using btdm_scalar_type =
typename impl_type::btdm_scalar_type;
3771 using tpetra_multivector_type =
typename impl_type::tpetra_multivector_type;
3772 using local_ordinal_type_1d_view =
typename impl_type::local_ordinal_type_1d_view;
3773 using vector_type_3d_view =
typename impl_type::vector_type_3d_view;
3774 using impl_scalar_type_2d_view_tpetra =
typename impl_type::impl_scalar_type_2d_view_tpetra;
3775 using const_impl_scalar_type_2d_view_tpetra =
typename impl_scalar_type_2d_view_tpetra::const_type;
3776 static constexpr
int vector_length = impl_type::vector_length;
3778 using member_type =
typename Kokkos::TeamPolicy<execution_space>::member_type;
3782 const ConstUnmanaged<local_ordinal_type_1d_view> partptr;
3783 const ConstUnmanaged<local_ordinal_type_1d_view> packptr;
3784 const ConstUnmanaged<local_ordinal_type_1d_view> part2packrowidx0;
3785 const ConstUnmanaged<local_ordinal_type_1d_view> part2rowidx0;
3786 const ConstUnmanaged<local_ordinal_type_1d_view> lclrow;
3787 const local_ordinal_type blocksize;
3788 const local_ordinal_type num_vectors;
3791 vector_type_3d_view packed_multivector;
3792 const_impl_scalar_type_2d_view_tpetra scalar_multivector;
3794 template<
typename TagType>
3795 KOKKOS_INLINE_FUNCTION
3796 void copy_multivectors(
const local_ordinal_type &j,
3797 const local_ordinal_type &vi,
3798 const local_ordinal_type &pri,
3799 const local_ordinal_type &ri0)
const {
3800 for (local_ordinal_type col=0;col<num_vectors;++col)
3801 for (local_ordinal_type i=0;i<blocksize;++i)
3802 packed_multivector(pri, i, col)[vi] =
static_cast<btdm_scalar_type
>(scalar_multivector(blocksize*lclrow(ri0+j)+i,col));
3808 const vector_type_3d_view &pmv)
3809 : partptr(interf.partptr),
3810 packptr(interf.packptr),
3811 part2packrowidx0(interf.part2packrowidx0),
3812 part2rowidx0(interf.part2rowidx0),
3813 lclrow(interf.lclrow),
3814 blocksize(pmv.extent(1)),
3815 num_vectors(pmv.extent(2)),
3816 packed_multivector(pmv) {}
3819 KOKKOS_INLINE_FUNCTION
3821 operator() (
const local_ordinal_type &packidx)
const {
3822 local_ordinal_type partidx = packptr(packidx);
3823 local_ordinal_type npacks = packptr(packidx+1) - partidx;
3824 const local_ordinal_type pri0 = part2packrowidx0(partidx);
3826 local_ordinal_type ri0[vector_length] = {};
3827 local_ordinal_type nrows[vector_length] = {};
3828 for (local_ordinal_type v=0;v<npacks;++v,++partidx) {
3829 ri0[v] = part2rowidx0(partidx);
3830 nrows[v] = part2rowidx0(partidx+1) - ri0[v];
3832 for (local_ordinal_type j=0;j<nrows[0];++j) {
3833 local_ordinal_type cnt = 1;
3834 for (;cnt<npacks && j!= nrows[cnt];++cnt);
3836 const local_ordinal_type pri = pri0 + j;
3837 for (local_ordinal_type col=0;col<num_vectors;++col)
3838 for (local_ordinal_type i=0;i<blocksize;++i)
3839 for (local_ordinal_type v=0;v<npacks;++v)
3840 packed_multivector(pri, i, col)[v] =
static_cast<btdm_scalar_type
>(scalar_multivector(blocksize*lclrow(ri0[v]+j)+i,col));
3844 KOKKOS_INLINE_FUNCTION
3846 operator() (
const member_type &member)
const {
3847 const local_ordinal_type packidx = member.league_rank();
3848 const local_ordinal_type partidx_begin = packptr(packidx);
3849 const local_ordinal_type npacks = packptr(packidx+1) - partidx_begin;
3850 const local_ordinal_type pri0 = part2packrowidx0(partidx_begin);
3851 Kokkos::parallel_for(Kokkos::ThreadVectorRange(member, npacks), [&](
const local_ordinal_type &v) {
3852 const local_ordinal_type partidx = partidx_begin + v;
3853 const local_ordinal_type ri0 = part2rowidx0(partidx);
3854 const local_ordinal_type nrows = part2rowidx0(partidx+1) - ri0;
3857 const local_ordinal_type pri = pri0;
3858 for (local_ordinal_type col=0;col<num_vectors;++col) {
3859 Kokkos::parallel_for(Kokkos::TeamThreadRange(member, blocksize), [&](
const local_ordinal_type &i) {
3860 packed_multivector(pri, i, col)[v] =
static_cast<btdm_scalar_type
>(scalar_multivector(blocksize*lclrow(ri0)+i,col));
3864 Kokkos::parallel_for(Kokkos::TeamThreadRange(member, nrows), [&](
const local_ordinal_type &j) {
3865 const local_ordinal_type pri = pri0 + j;
3866 for (local_ordinal_type col=0;col<num_vectors;++col)
3867 for (local_ordinal_type i=0;i<blocksize;++i)
3868 packed_multivector(pri, i, col)[v] = static_cast<btdm_scalar_type>(scalar_multivector(blocksize*lclrow(ri0+j)+i,col));
3874 void run(
const const_impl_scalar_type_2d_view_tpetra &scalar_multivector_) {
3875 IFPACK2_BLOCKTRIDICONTAINER_PROFILER_REGION_BEGIN;
3876 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::MultiVectorConverter", MultiVectorConverter0);
3878 scalar_multivector = scalar_multivector_;
3879 if constexpr (BlockHelperDetails::is_device<execution_space>::value) {
3880 const local_ordinal_type vl = vector_length;
3881 const Kokkos::TeamPolicy<execution_space> policy(packptr.extent(0) - 1, Kokkos::AUTO(), vl);
3882 Kokkos::parallel_for
3883 (
"MultiVectorConverter::TeamPolicy", policy, *
this);
3885 const Kokkos::RangePolicy<execution_space> policy(0, packptr.extent(0) - 1);
3886 Kokkos::parallel_for
3887 (
"MultiVectorConverter::RangePolicy", policy, *
this);
3889 IFPACK2_BLOCKTRIDICONTAINER_PROFILER_REGION_END;
3890 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
3899 struct SolveTridiagsDefaultModeAndAlgo<Kokkos::HostSpace> {
3900 typedef KB::Mode::Serial mode_type;
3901 typedef KB::Algo::Level2::Unblocked single_vector_algo_type;
3902 #if defined(__KOKKOSBATCHED_INTEL_MKL_COMPACT_BATCHED__)
3903 typedef KB::Algo::Level3::CompactMKL multi_vector_algo_type;
3905 typedef KB::Algo::Level3::Blocked multi_vector_algo_type;
3907 static int recommended_team_size(
const int ,
3914 #if defined(KOKKOS_ENABLE_CUDA)
3915 static inline int SolveTridiagsRecommendedCudaTeamSize(
const int blksize,
3916 const int vector_length,
3917 const int internal_vector_length) {
3918 const int vector_size = vector_length/internal_vector_length;
3919 int total_team_size(0);
3920 if (blksize <= 5) total_team_size = 32;
3921 else if (blksize <= 9) total_team_size = 32;
3922 else if (blksize <= 12) total_team_size = 96;
3923 else if (blksize <= 16) total_team_size = 128;
3924 else if (blksize <= 20) total_team_size = 160;
3925 else total_team_size = 160;
3926 return total_team_size/vector_size;
3930 struct SolveTridiagsDefaultModeAndAlgo<Kokkos::CudaSpace> {
3931 typedef KB::Mode::Team mode_type;
3932 typedef KB::Algo::Level2::Unblocked single_vector_algo_type;
3933 typedef KB::Algo::Level3::Unblocked multi_vector_algo_type;
3934 static int recommended_team_size(
const int blksize,
3935 const int vector_length,
3936 const int internal_vector_length) {
3937 return SolveTridiagsRecommendedCudaTeamSize(blksize, vector_length, internal_vector_length);
3941 struct SolveTridiagsDefaultModeAndAlgo<Kokkos::CudaUVMSpace> {
3942 typedef KB::Mode::Team mode_type;
3943 typedef KB::Algo::Level2::Unblocked single_vector_algo_type;
3944 typedef KB::Algo::Level3::Unblocked multi_vector_algo_type;
3945 static int recommended_team_size(
const int blksize,
3946 const int vector_length,
3947 const int internal_vector_length) {
3948 return SolveTridiagsRecommendedCudaTeamSize(blksize, vector_length, internal_vector_length);
3953 #if defined(KOKKOS_ENABLE_HIP)
3954 static inline int SolveTridiagsRecommendedHIPTeamSize(
const int blksize,
3955 const int vector_length,
3956 const int internal_vector_length) {
3957 const int vector_size = vector_length/internal_vector_length;
3958 int total_team_size(0);
3959 if (blksize <= 5) total_team_size = 32;
3960 else if (blksize <= 9) total_team_size = 32;
3961 else if (blksize <= 12) total_team_size = 96;
3962 else if (blksize <= 16) total_team_size = 128;
3963 else if (blksize <= 20) total_team_size = 160;
3964 else total_team_size = 160;
3965 return total_team_size/vector_size;
3969 struct SolveTridiagsDefaultModeAndAlgo<Kokkos::HIPSpace> {
3970 typedef KB::Mode::Team mode_type;
3971 typedef KB::Algo::Level2::Unblocked single_vector_algo_type;
3972 typedef KB::Algo::Level3::Unblocked multi_vector_algo_type;
3973 static int recommended_team_size(
const int blksize,
3974 const int vector_length,
3975 const int internal_vector_length) {
3976 return SolveTridiagsRecommendedHIPTeamSize(blksize, vector_length, internal_vector_length);
3980 struct SolveTridiagsDefaultModeAndAlgo<Kokkos::HIPHostPinnedSpace> {
3981 typedef KB::Mode::Team mode_type;
3982 typedef KB::Algo::Level2::Unblocked single_vector_algo_type;
3983 typedef KB::Algo::Level3::Unblocked multi_vector_algo_type;
3984 static int recommended_team_size(
const int blksize,
3985 const int vector_length,
3986 const int internal_vector_length) {
3987 return SolveTridiagsRecommendedHIPTeamSize(blksize, vector_length, internal_vector_length);
3992 #if defined(KOKKOS_ENABLE_SYCL)
3993 static inline int SolveTridiagsRecommendedSYCLTeamSize(
const int blksize,
3994 const int vector_length,
3995 const int internal_vector_length) {
3996 const int vector_size = vector_length/internal_vector_length;
3997 int total_team_size(0);
3998 if (blksize <= 5) total_team_size = 32;
3999 else if (blksize <= 9) total_team_size = 32;
4000 else if (blksize <= 12) total_team_size = 96;
4001 else if (blksize <= 16) total_team_size = 128;
4002 else if (blksize <= 20) total_team_size = 160;
4003 else total_team_size = 160;
4004 return total_team_size/vector_size;
4008 struct SolveTridiagsDefaultModeAndAlgo<Kokkos::Experimental::SYCLSharedUSMSpace> {
4009 typedef KB::Mode::Team mode_type;
4010 typedef KB::Algo::Level2::Unblocked single_vector_algo_type;
4011 typedef KB::Algo::Level3::Unblocked multi_vector_algo_type;
4012 static int recommended_team_size(
const int blksize,
4013 const int vector_length,
4014 const int internal_vector_length) {
4015 return SolveTridiagsRecommendedSYCLTeamSize(blksize, vector_length, internal_vector_length);
4019 struct SolveTridiagsDefaultModeAndAlgo<Kokkos::Experimental::SYCLDeviceUSMSpace> {
4020 typedef KB::Mode::Team mode_type;
4021 typedef KB::Algo::Level2::Unblocked single_vector_algo_type;
4022 typedef KB::Algo::Level3::Unblocked multi_vector_algo_type;
4023 static int recommended_team_size(
const int blksize,
4024 const int vector_length,
4025 const int internal_vector_length) {
4026 return SolveTridiagsRecommendedSYCLTeamSize(blksize, vector_length, internal_vector_length);
4034 template<
typename MatrixType>
4035 struct SolveTridiags {
4037 using impl_type = BlockHelperDetails::ImplType<MatrixType>;
4038 using execution_space =
typename impl_type::execution_space;
4040 using local_ordinal_type =
typename impl_type::local_ordinal_type;
4043 using magnitude_type =
typename impl_type::magnitude_type;
4044 using btdm_scalar_type =
typename impl_type::btdm_scalar_type;
4045 using btdm_magnitude_type =
typename impl_type::btdm_magnitude_type;
4047 using local_ordinal_type_1d_view =
typename impl_type::local_ordinal_type_1d_view;
4048 using local_ordinal_type_2d_view =
typename impl_type::local_ordinal_type_2d_view;
4049 using size_type_2d_view =
typename impl_type::size_type_2d_view;
4051 using vector_type_3d_view =
typename impl_type::vector_type_3d_view;
4052 using internal_vector_type_4d_view =
typename impl_type::internal_vector_type_4d_view;
4053 using internal_vector_type_5d_view =
typename impl_type::internal_vector_type_5d_view;
4054 using btdm_scalar_type_4d_view =
typename impl_type::btdm_scalar_type_4d_view;
4056 using internal_vector_scratch_type_3d_view = Scratch<typename impl_type::internal_vector_type_3d_view>;
4058 using internal_vector_type =
typename impl_type::internal_vector_type;
4059 static constexpr
int vector_length = impl_type::vector_length;
4060 static constexpr
int internal_vector_length = impl_type::internal_vector_length;
4063 using impl_scalar_type_1d_view =
typename impl_type::impl_scalar_type_1d_view;
4064 using impl_scalar_type_2d_view_tpetra =
typename impl_type::impl_scalar_type_2d_view_tpetra;
4067 using team_policy_type = Kokkos::TeamPolicy<execution_space>;
4068 using member_type =
typename team_policy_type::member_type;
4072 local_ordinal_type n_subparts_per_part;
4073 const ConstUnmanaged<local_ordinal_type_1d_view> partptr;
4074 const ConstUnmanaged<local_ordinal_type_1d_view> packptr;
4075 const ConstUnmanaged<local_ordinal_type_1d_view> packindices_sub;
4076 const ConstUnmanaged<local_ordinal_type_2d_view> packindices_schur;
4077 const ConstUnmanaged<local_ordinal_type_1d_view> part2packrowidx0;
4078 const ConstUnmanaged<local_ordinal_type_2d_view> part2packrowidx0_sub;
4079 const ConstUnmanaged<local_ordinal_type_1d_view> lclrow;
4080 const ConstUnmanaged<local_ordinal_type_1d_view> packptr_sub;
4082 const ConstUnmanaged<local_ordinal_type_2d_view> partptr_sub;
4083 const ConstUnmanaged<size_type_2d_view> pack_td_ptr_schur;
4086 const ConstUnmanaged<size_type_2d_view> pack_td_ptr;
4089 const ConstUnmanaged<internal_vector_type_4d_view> D_internal_vector_values;
4090 const Unmanaged<internal_vector_type_4d_view> X_internal_vector_values;
4091 const Unmanaged<btdm_scalar_type_4d_view> X_internal_scalar_values;
4093 internal_vector_type_4d_view X_internal_vector_values_schur;
4095 const ConstUnmanaged<internal_vector_type_4d_view> D_internal_vector_values_schur;
4096 const ConstUnmanaged<internal_vector_type_5d_view> e_internal_vector_values;
4099 const local_ordinal_type vector_loop_size;
4102 Unmanaged<impl_scalar_type_2d_view_tpetra> Y_scalar_multivector;
4103 #if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) || defined(__SYCL_DEVICE_ONLY__)
4104 AtomicUnmanaged<impl_scalar_type_1d_view> Z_scalar_vector;
4106 Unmanaged<impl_scalar_type_1d_view> Z_scalar_vector;
4108 const impl_scalar_type df;
4109 const bool compute_diff;
4112 SolveTridiags(
const BlockHelperDetails::PartInterface<MatrixType> &interf,
4113 const BlockTridiags<MatrixType> &btdm,
4114 const vector_type_3d_view &pmv,
4115 const impl_scalar_type damping_factor,
4116 const bool is_norm_manager_active)
4119 n_subparts_per_part(interf.n_subparts_per_part),
4120 partptr(interf.partptr),
4121 packptr(interf.packptr),
4122 packindices_sub(interf.packindices_sub),
4123 packindices_schur(interf.packindices_schur),
4124 part2packrowidx0(interf.part2packrowidx0),
4125 part2packrowidx0_sub(interf.part2packrowidx0_sub),
4126 lclrow(interf.lclrow),
4127 packptr_sub(interf.packptr_sub),
4128 partptr_sub(interf.partptr_sub),
4129 pack_td_ptr_schur(btdm.pack_td_ptr_schur),
4131 pack_td_ptr(btdm.pack_td_ptr),
4132 D_internal_vector_values((internal_vector_type*)btdm.values.data(),
4133 btdm.values.extent(0),
4134 btdm.values.extent(1),
4135 btdm.values.extent(2),
4136 vector_length/internal_vector_length),
4137 X_internal_vector_values((internal_vector_type*)pmv.data(),
4141 vector_length/internal_vector_length),
4142 X_internal_scalar_values((btdm_scalar_type*)pmv.data(),
4148 2*(n_subparts_per_part-1) * part2packrowidx0_sub.extent(0),
4151 vector_length/internal_vector_length),
4152 D_internal_vector_values_schur((internal_vector_type*)btdm.values_schur.data(),
4153 btdm.values_schur.extent(0),
4154 btdm.values_schur.extent(1),
4155 btdm.values_schur.extent(2),
4156 vector_length/internal_vector_length),
4157 e_internal_vector_values((internal_vector_type*)btdm.e_values.data(),
4158 btdm.e_values.extent(0),
4159 btdm.e_values.extent(1),
4160 btdm.e_values.extent(2),
4161 btdm.e_values.extent(3),
4162 vector_length/internal_vector_length),
4163 vector_loop_size(vector_length/internal_vector_length),
4164 Y_scalar_multivector(),
4167 compute_diff(is_norm_manager_active)
4173 KOKKOS_INLINE_FUNCTION
4175 copyToFlatMultiVector(
const member_type &member,
4176 const local_ordinal_type partidxbeg,
4177 const local_ordinal_type npacks,
4178 const local_ordinal_type pri0,
4179 const local_ordinal_type v,
4180 const local_ordinal_type blocksize,
4181 const local_ordinal_type num_vectors)
const {
4182 const local_ordinal_type vbeg = v*internal_vector_length;
4183 if (vbeg < npacks) {
4184 local_ordinal_type ri0_vals[internal_vector_length] = {};
4185 local_ordinal_type nrows_vals[internal_vector_length] = {};
4186 for (local_ordinal_type vv=vbeg,vi=0;vv<npacks && vi<internal_vector_length;++vv,++vi) {
4187 const local_ordinal_type partidx = partidxbeg+vv;
4188 ri0_vals[vi] = partptr(partidx);
4189 nrows_vals[vi] = partptr(partidx+1) - ri0_vals[vi];
4192 impl_scalar_type z_partial_sum(0);
4193 if (nrows_vals[0] == 1) {
4194 const local_ordinal_type j=0, pri=pri0;
4196 for (local_ordinal_type vv=vbeg,vi=0;vv<npacks && vi<internal_vector_length;++vv,++vi) {
4197 const local_ordinal_type ri0 = ri0_vals[vi];
4198 const local_ordinal_type nrows = nrows_vals[vi];
4200 Kokkos::parallel_for
4201 (Kokkos::TeamThreadRange(member, blocksize),
4202 [&](
const local_ordinal_type &i) {
4203 const local_ordinal_type row = blocksize*lclrow(ri0+j)+i;
4204 for (local_ordinal_type col=0;col<num_vectors;++col) {
4205 impl_scalar_type &y = Y_scalar_multivector(row,col);
4206 const impl_scalar_type yd = X_internal_vector_values(pri, i, col, v)[vi] - y;
4210 const auto yd_abs = Kokkos::ArithTraits<impl_scalar_type>::abs(yd);
4211 z_partial_sum += yd_abs*yd_abs;
4219 Kokkos::parallel_for
4220 (Kokkos::TeamThreadRange(member, nrows_vals[0]),
4221 [&](
const local_ordinal_type &j) {
4222 const local_ordinal_type pri = pri0 + j;
4223 for (local_ordinal_type vv=vbeg,vi=0;vv<npacks && vi<internal_vector_length;++vv,++vi) {
4224 const local_ordinal_type ri0 = ri0_vals[vi];
4225 const local_ordinal_type nrows = nrows_vals[vi];
4227 for (local_ordinal_type col=0;col<num_vectors;++col) {
4228 for (local_ordinal_type i=0;i<blocksize;++i) {
4229 const local_ordinal_type row = blocksize*lclrow(ri0+j)+i;
4230 impl_scalar_type &y = Y_scalar_multivector(row,col);
4231 const impl_scalar_type yd = X_internal_vector_values(pri, i, col, v)[vi] - y;
4235 const auto yd_abs = Kokkos::ArithTraits<impl_scalar_type>::abs(yd);
4236 z_partial_sum += yd_abs*yd_abs;
4245 Z_scalar_vector(member.league_rank()) += z_partial_sum;
4252 template<
typename WWViewType>
4253 KOKKOS_INLINE_FUNCTION
4255 solveSingleVector(
const member_type &member,
4256 const local_ordinal_type &blocksize,
4257 const local_ordinal_type &i0,
4258 const local_ordinal_type &r0,
4259 const local_ordinal_type &nrows,
4260 const local_ordinal_type &v,
4261 const WWViewType &WW)
const {
4263 typedef SolveTridiagsDefaultModeAndAlgo
4264 <
typename execution_space::memory_space> default_mode_and_algo_type;
4266 typedef typename default_mode_and_algo_type::mode_type default_mode_type;
4267 typedef typename default_mode_and_algo_type::single_vector_algo_type default_algo_type;
4270 auto A = D_internal_vector_values.data();
4271 auto X = X_internal_vector_values.data();
4274 const auto one = Kokkos::ArithTraits<btdm_magnitude_type>::one();
4275 const auto zero = Kokkos::ArithTraits<btdm_magnitude_type>::zero();
4279 const local_ordinal_type astep = D_internal_vector_values.stride_0();
4280 const local_ordinal_type as0 = D_internal_vector_values.stride_1();
4281 const local_ordinal_type as1 = D_internal_vector_values.stride_2();
4282 const local_ordinal_type xstep = X_internal_vector_values.stride_0();
4283 const local_ordinal_type xs0 = X_internal_vector_values.stride_1();
4292 KOKKOSBATCHED_TRSV_LOWER_NO_TRANSPOSE_INTERNAL_INVOKE
4293 (default_mode_type,default_algo_type,
4296 blocksize,blocksize,
4301 for (local_ordinal_type tr=1;tr<nrows;++tr) {
4302 member.team_barrier();
4303 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
4304 (default_mode_type,default_algo_type,
4306 blocksize, blocksize,
4308 A+2*astep, as0, as1,
4312 KOKKOSBATCHED_TRSV_LOWER_NO_TRANSPOSE_INTERNAL_INVOKE
4313 (default_mode_type,default_algo_type,
4316 blocksize,blocksize,
4318 A+3*astep, as0, as1,
4326 KOKKOSBATCHED_TRSV_UPPER_NO_TRANSPOSE_INTERNAL_INVOKE
4327 (default_mode_type,default_algo_type,
4330 blocksize, blocksize,
4335 for (local_ordinal_type tr=nrows;tr>1;--tr) {
4337 member.team_barrier();
4338 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
4339 (default_mode_type,default_algo_type,
4341 blocksize, blocksize,
4343 A+1*astep, as0, as1,
4347 KOKKOSBATCHED_TRSV_UPPER_NO_TRANSPOSE_INTERNAL_INVOKE
4348 (default_mode_type,default_algo_type,
4351 blocksize, blocksize,
4360 const local_ordinal_type ws0 = WW.stride_0();
4361 auto W = WW.data() + v;
4362 KOKKOSBATCHED_COPY_VECTOR_NO_TRANSPOSE_INTERNAL_INVOKE
4364 member, blocksize, X, xs0, W, ws0);
4365 member.team_barrier();
4366 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
4367 (default_mode_type,default_algo_type,
4369 blocksize, blocksize,
4378 template<
typename WWViewType>
4379 KOKKOS_INLINE_FUNCTION
4381 solveMultiVector(
const member_type &member,
4382 const local_ordinal_type &,
4383 const local_ordinal_type &i0,
4384 const local_ordinal_type &r0,
4385 const local_ordinal_type &nrows,
4386 const local_ordinal_type &v,
4387 const WWViewType &WW)
const {
4389 typedef SolveTridiagsDefaultModeAndAlgo
4390 <
typename execution_space::memory_space> default_mode_and_algo_type;
4392 typedef typename default_mode_and_algo_type::mode_type default_mode_type;
4393 typedef typename default_mode_and_algo_type::multi_vector_algo_type default_algo_type;
4396 const auto one = Kokkos::ArithTraits<btdm_magnitude_type>::one();
4397 const auto zero = Kokkos::ArithTraits<btdm_magnitude_type>::zero();
4400 auto A = Kokkos::subview(D_internal_vector_values, i0, Kokkos::ALL(), Kokkos::ALL(), v);
4401 auto X1 = Kokkos::subview(X_internal_vector_values, r0, Kokkos::ALL(), Kokkos::ALL(), v);
4404 local_ordinal_type i = i0, r = r0;
4409 KB::Trsm<member_type,
4410 KB::Side::Left,KB::Uplo::Lower,KB::Trans::NoTranspose,KB::Diag::Unit,
4411 default_mode_type,default_algo_type>
4412 ::invoke(member, one, A, X1);
4413 for (local_ordinal_type tr=1;tr<nrows;++tr,i+=3) {
4414 A.assign_data( &D_internal_vector_values(i+2,0,0,v) );
4415 X2.assign_data( &X_internal_vector_values(++r,0,0,v) );
4416 member.team_barrier();
4417 KB::Gemm<member_type,
4418 KB::Trans::NoTranspose,KB::Trans::NoTranspose,
4419 default_mode_type,default_algo_type>
4420 ::invoke(member, -one, A, X1, one, X2);
4421 A.assign_data( &D_internal_vector_values(i+3,0,0,v) );
4422 KB::Trsm<member_type,
4423 KB::Side::Left,KB::Uplo::Lower,KB::Trans::NoTranspose,KB::Diag::Unit,
4424 default_mode_type,default_algo_type>
4425 ::invoke(member, one, A, X2);
4426 X1.assign_data( X2.data() );
4430 KB::Trsm<member_type,
4431 KB::Side::Left,KB::Uplo::Upper,KB::Trans::NoTranspose,KB::Diag::NonUnit,
4432 default_mode_type,default_algo_type>
4433 ::invoke(member, one, A, X1);
4434 for (local_ordinal_type tr=nrows;tr>1;--tr) {
4436 A.assign_data( &D_internal_vector_values(i+1,0,0,v) );
4437 X2.assign_data( &X_internal_vector_values(--r,0,0,v) );
4438 member.team_barrier();
4439 KB::Gemm<member_type,
4440 KB::Trans::NoTranspose,KB::Trans::NoTranspose,
4441 default_mode_type,default_algo_type>
4442 ::invoke(member, -one, A, X1, one, X2);
4444 A.assign_data( &D_internal_vector_values(i,0,0,v) );
4445 KB::Trsm<member_type,
4446 KB::Side::Left,KB::Uplo::Upper,KB::Trans::NoTranspose,KB::Diag::NonUnit,
4447 default_mode_type,default_algo_type>
4448 ::invoke(member, one, A, X2);
4449 X1.assign_data( X2.data() );
4453 auto W = Kokkos::subview(WW, Kokkos::ALL(), Kokkos::ALL(), v);
4454 KB::Copy<member_type,KB::Trans::NoTranspose,default_mode_type>
4455 ::invoke(member, X1, W);
4456 member.team_barrier();
4457 KB::Gemm<member_type,
4458 KB::Trans::NoTranspose,KB::Trans::NoTranspose,
4459 default_mode_type,default_algo_type>
4460 ::invoke(member, one, A, W, zero, X1);
4464 template<
int B>
struct SingleVectorTag {};
4465 template<
int B>
struct MultiVectorTag {};
4467 template<
int B>
struct SingleVectorSubLineTag {};
4468 template<
int B>
struct MultiVectorSubLineTag {};
4469 template<
int B>
struct SingleVectorApplyCTag {};
4470 template<
int B>
struct MultiVectorApplyCTag {};
4471 template<
int B>
struct SingleVectorSchurTag {};
4472 template<
int B>
struct MultiVectorSchurTag {};
4473 template<
int B>
struct SingleVectorApplyETag {};
4474 template<
int B>
struct MultiVectorApplyETag {};
4475 template<
int B>
struct SingleVectorCopyToFlatTag {};
4476 template<
int B>
struct SingleZeroingTag {};
4479 KOKKOS_INLINE_FUNCTION
4481 operator() (
const SingleVectorTag<B> &,
const member_type &member)
const {
4482 const local_ordinal_type packidx = member.league_rank();
4483 const local_ordinal_type partidx = packptr(packidx);
4484 const local_ordinal_type npacks = packptr(packidx+1) - partidx;
4485 const local_ordinal_type pri0 = part2packrowidx0(partidx);
4486 const local_ordinal_type i0 = pack_td_ptr(partidx,0);
4487 const local_ordinal_type r0 = part2packrowidx0(partidx);
4488 const local_ordinal_type nrows = partptr(partidx+1) - partptr(partidx);
4489 const local_ordinal_type blocksize = (B == 0 ? D_internal_vector_values.extent(1) : B);
4490 const local_ordinal_type num_vectors = 1;
4491 internal_vector_scratch_type_3d_view
4492 WW(member.team_scratch(0), blocksize, 1, vector_loop_size);
4493 Kokkos::single(Kokkos::PerTeam(member), [&]() {
4494 Z_scalar_vector(member.league_rank()) = impl_scalar_type(0);
4496 Kokkos::parallel_for
4497 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
4498 solveSingleVector(member, blocksize, i0, r0, nrows, v, WW);
4499 copyToFlatMultiVector(member, partidx, npacks, pri0, v, blocksize, num_vectors);
4504 KOKKOS_INLINE_FUNCTION
4506 operator() (
const MultiVectorTag<B> &,
const member_type &member)
const {
4507 const local_ordinal_type packidx = member.league_rank();
4508 const local_ordinal_type partidx = packptr(packidx);
4509 const local_ordinal_type npacks = packptr(packidx+1) - partidx;
4510 const local_ordinal_type pri0 = part2packrowidx0(partidx);
4511 const local_ordinal_type i0 = pack_td_ptr(partidx,0);
4512 const local_ordinal_type r0 = part2packrowidx0(partidx);
4513 const local_ordinal_type nrows = partptr(partidx+1) - partptr(partidx);
4514 const local_ordinal_type blocksize = (B == 0 ? D_internal_vector_values.extent(1) : B);
4515 const local_ordinal_type num_vectors = X_internal_vector_values.extent(2);
4517 internal_vector_scratch_type_3d_view
4518 WW(member.team_scratch(0), blocksize, num_vectors, vector_loop_size);
4519 Kokkos::single(Kokkos::PerTeam(member), [&]() {
4520 Z_scalar_vector(member.league_rank()) = impl_scalar_type(0);
4522 Kokkos::parallel_for
4523 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
4524 solveMultiVector(member, blocksize, i0, r0, nrows, v, WW);
4525 copyToFlatMultiVector(member, partidx, npacks, pri0, v, blocksize, num_vectors);
4530 KOKKOS_INLINE_FUNCTION
4532 operator() (
const SingleVectorSubLineTag<B> &,
const member_type &member)
const {
4534 const local_ordinal_type packidx = packindices_sub(member.league_rank());
4536 const local_ordinal_type subpartidx = packptr_sub(packidx);
4537 const local_ordinal_type n_parts = part2packrowidx0_sub.extent(0);
4538 const local_ordinal_type local_subpartidx = subpartidx/n_parts;
4539 const local_ordinal_type partidx = subpartidx%n_parts;
4541 const local_ordinal_type npacks = packptr_sub(packidx+1) - subpartidx;
4542 const local_ordinal_type i0 = pack_td_ptr(partidx,local_subpartidx);
4543 const local_ordinal_type r0 = part2packrowidx0_sub(partidx,local_subpartidx);
4544 const local_ordinal_type nrows = partptr_sub(subpartidx,1) - partptr_sub(subpartidx,0);
4545 const local_ordinal_type blocksize = e_internal_vector_values.extent(2);
4551 internal_vector_scratch_type_3d_view
4552 WW(member.team_scratch(0), blocksize, 1, vector_loop_size);
4554 Kokkos::parallel_for
4555 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
4556 solveSingleVectorNew<impl_type, internal_vector_scratch_type_3d_view> (member, blocksize, i0, r0, nrows, v, D_internal_vector_values, X_internal_vector_values, WW);
4561 KOKKOS_INLINE_FUNCTION
4563 operator() (
const SingleVectorApplyCTag<B> &,
const member_type &member)
const {
4566 const local_ordinal_type packidx = packindices_sub(member.league_rank());
4568 const local_ordinal_type subpartidx = packptr_sub(packidx);
4569 const local_ordinal_type n_parts = part2packrowidx0_sub.extent(0);
4570 const local_ordinal_type local_subpartidx = subpartidx/n_parts;
4571 const local_ordinal_type partidx = subpartidx%n_parts;
4572 const local_ordinal_type blocksize = e_internal_vector_values.extent(2);
4575 const local_ordinal_type i0 = pack_td_ptr(partidx,local_subpartidx);
4576 const local_ordinal_type r0 = part2packrowidx0_sub(partidx,local_subpartidx);
4577 const local_ordinal_type nrows = partptr_sub(subpartidx,1) - partptr_sub(subpartidx,0);
4579 internal_vector_scratch_type_3d_view
4580 WW(member.team_scratch(0), blocksize, blocksize, vector_loop_size);
4584 const local_ordinal_type local_subpartidx_schur = (local_subpartidx-1)/2;
4585 const local_ordinal_type i0_schur = local_subpartidx_schur == 0 ? pack_td_ptr_schur(partidx,local_subpartidx_schur) : pack_td_ptr_schur(partidx,local_subpartidx_schur) + 1;
4586 const local_ordinal_type i0_offset = local_subpartidx_schur == 0 ? i0+2 : i0+2;
4591 const auto one = Kokkos::ArithTraits<btdm_magnitude_type>::one();
4593 const size_type c_kps2 = local_subpartidx > 0 ? pack_td_ptr(partidx, local_subpartidx)-2 : 0;
4594 const size_type c_kps1 = pack_td_ptr(partidx, local_subpartidx+1)+1;
4596 typedef SolveTridiagsDefaultModeAndAlgo
4597 <
typename execution_space::memory_space> default_mode_and_algo_type;
4599 typedef typename default_mode_and_algo_type::mode_type default_mode_type;
4600 typedef typename default_mode_and_algo_type::single_vector_algo_type default_algo_type;
4602 if (local_subpartidx == 0) {
4603 Kokkos::parallel_for
4604 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
4605 auto v_1 = Kokkos::subview(X_internal_vector_values, r0+nrows-1, Kokkos::ALL(), 0, v);
4606 auto v_2 = Kokkos::subview(X_internal_vector_values, r0+nrows, Kokkos::ALL(), 0, v);
4607 auto C = Kokkos::subview(D_internal_vector_values, c_kps1, Kokkos::ALL(), Kokkos::ALL(), v);
4609 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
4610 (default_mode_type,default_algo_type,
4612 blocksize, blocksize,
4614 C.data(), C.stride_0(), C.stride_1(),
4615 v_1.data(), v_1.stride_0(),
4617 v_2.data(), v_2.stride_0());
4620 else if (local_subpartidx == (local_ordinal_type) part2packrowidx0_sub.extent(1) - 2) {
4621 Kokkos::parallel_for
4622 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
4623 auto v_1 = Kokkos::subview(X_internal_vector_values, r0, Kokkos::ALL(), 0, v);
4624 auto v_2 = Kokkos::subview(X_internal_vector_values, r0-1, Kokkos::ALL(), 0, v);
4625 auto C = Kokkos::subview(D_internal_vector_values, c_kps2, Kokkos::ALL(), Kokkos::ALL(), v);
4627 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
4628 (default_mode_type,default_algo_type,
4630 blocksize, blocksize,
4632 C.data(), C.stride_0(), C.stride_1(),
4633 v_1.data(), v_1.stride_0(),
4635 v_2.data(), v_2.stride_0());
4639 Kokkos::parallel_for
4640 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
4642 auto v_1 = Kokkos::subview(X_internal_vector_values, r0+nrows-1, Kokkos::ALL(), 0, v);
4643 auto v_2 = Kokkos::subview(X_internal_vector_values, r0+nrows, Kokkos::ALL(), 0, v);
4644 auto C = Kokkos::subview(D_internal_vector_values, c_kps1, Kokkos::ALL(), Kokkos::ALL(), v);
4646 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
4647 (default_mode_type,default_algo_type,
4649 blocksize, blocksize,
4651 C.data(), C.stride_0(), C.stride_1(),
4652 v_1.data(), v_1.stride_0(),
4654 v_2.data(), v_2.stride_0());
4657 auto v_1 = Kokkos::subview(X_internal_vector_values, r0, Kokkos::ALL(), 0, v);
4658 auto v_2 = Kokkos::subview(X_internal_vector_values, r0-1, Kokkos::ALL(), 0, v);
4659 auto C = Kokkos::subview(D_internal_vector_values, c_kps2, Kokkos::ALL(), Kokkos::ALL(), v);
4661 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
4662 (default_mode_type,default_algo_type,
4664 blocksize, blocksize,
4666 C.data(), C.stride_0(), C.stride_1(),
4667 v_1.data(), v_1.stride_0(),
4669 v_2.data(), v_2.stride_0());
4676 KOKKOS_INLINE_FUNCTION
4678 operator() (
const SingleVectorSchurTag<B> &,
const member_type &member)
const {
4679 const local_ordinal_type packidx = packindices_sub(member.league_rank());
4681 const local_ordinal_type partidx = packptr_sub(packidx);
4683 const local_ordinal_type blocksize = e_internal_vector_values.extent(2);
4685 const local_ordinal_type i0_schur = pack_td_ptr_schur(partidx,0);
4686 const local_ordinal_type nrows = 2*(n_subparts_per_part-1);
4688 const local_ordinal_type r0_schur = nrows * member.league_rank();
4690 internal_vector_scratch_type_3d_view
4691 WW(member.team_scratch(0), blocksize, blocksize, vector_loop_size);
4693 for (local_ordinal_type schur_sub_part = 0; schur_sub_part < n_subparts_per_part-1; ++schur_sub_part) {
4694 const local_ordinal_type r0 = part2packrowidx0_sub(partidx,2*schur_sub_part+1);
4695 for (local_ordinal_type i = 0; i < 2; ++i) {
4696 copy3DView<local_ordinal_type>(member,
4697 Kokkos::subview(X_internal_vector_values_schur, r0_schur+2*schur_sub_part+i, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()),
4698 Kokkos::subview(X_internal_vector_values, r0+i, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()));
4702 Kokkos::parallel_for
4703 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
4704 solveSingleVectorNew<impl_type, internal_vector_scratch_type_3d_view> (member, blocksize, i0_schur, r0_schur, nrows, v, D_internal_vector_values_schur, X_internal_vector_values_schur, WW);
4707 for (local_ordinal_type schur_sub_part = 0; schur_sub_part < n_subparts_per_part-1; ++schur_sub_part) {
4708 const local_ordinal_type r0 = part2packrowidx0_sub(partidx,2*schur_sub_part+1);
4709 for (local_ordinal_type i = 0; i < 2; ++i) {
4710 copy3DView<local_ordinal_type>(member,
4711 Kokkos::subview(X_internal_vector_values, r0+i, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()),
4712 Kokkos::subview(X_internal_vector_values_schur, r0_schur+2*schur_sub_part+i, Kokkos::ALL(), Kokkos::ALL(), Kokkos::ALL()));
4718 KOKKOS_INLINE_FUNCTION
4720 operator() (
const SingleVectorApplyETag<B> &,
const member_type &member)
const {
4721 const local_ordinal_type packidx = packindices_sub(member.league_rank());
4723 const local_ordinal_type subpartidx = packptr_sub(packidx);
4724 const local_ordinal_type n_parts = part2packrowidx0_sub.extent(0);
4725 const local_ordinal_type local_subpartidx = subpartidx/n_parts;
4726 const local_ordinal_type partidx = subpartidx%n_parts;
4727 const local_ordinal_type blocksize = e_internal_vector_values.extent(2);
4729 const local_ordinal_type r0 = part2packrowidx0_sub(partidx,local_subpartidx);
4730 const local_ordinal_type nrows = partptr_sub(subpartidx,1) - partptr_sub(subpartidx,0);
4732 internal_vector_scratch_type_3d_view
4733 WW(member.team_scratch(0), blocksize, blocksize, vector_loop_size);
4737 const auto one = Kokkos::ArithTraits<btdm_magnitude_type>::one();
4739 typedef SolveTridiagsDefaultModeAndAlgo
4740 <
typename execution_space::memory_space> default_mode_and_algo_type;
4742 typedef typename default_mode_and_algo_type::mode_type default_mode_type;
4743 typedef typename default_mode_and_algo_type::single_vector_algo_type default_algo_type;
4745 if (local_subpartidx == 0) {
4746 Kokkos::parallel_for
4747 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
4749 auto v_2 = Kokkos::subview(X_internal_vector_values, r0+nrows, Kokkos::ALL(), 0, v);
4751 for (local_ordinal_type row = 0; row < nrows; ++row) {
4752 auto v_1 = Kokkos::subview(X_internal_vector_values, r0+row, Kokkos::ALL(), 0, v);
4753 auto E = Kokkos::subview(e_internal_vector_values, 0, r0+row, Kokkos::ALL(), Kokkos::ALL(), v);
4755 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
4756 (default_mode_type,default_algo_type,
4758 blocksize, blocksize,
4760 E.data(), E.stride_0(), E.stride_1(),
4761 v_2.data(), v_2.stride_0(),
4763 v_1.data(), v_1.stride_0());
4767 else if (local_subpartidx == (local_ordinal_type) part2packrowidx0_sub.extent(1) - 2) {
4768 Kokkos::parallel_for
4769 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
4770 auto v_2 = Kokkos::subview(X_internal_vector_values, r0-1, Kokkos::ALL(), 0, v);
4772 for (local_ordinal_type row = 0; row < nrows; ++row) {
4773 auto v_1 = Kokkos::subview(X_internal_vector_values, r0+row, Kokkos::ALL(), 0, v);
4774 auto E = Kokkos::subview(e_internal_vector_values, 1, r0+row, Kokkos::ALL(), Kokkos::ALL(), v);
4776 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
4777 (default_mode_type,default_algo_type,
4779 blocksize, blocksize,
4781 E.data(), E.stride_0(), E.stride_1(),
4782 v_2.data(), v_2.stride_0(),
4784 v_1.data(), v_1.stride_0());
4789 Kokkos::parallel_for
4790 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
4792 auto v_2 = Kokkos::subview(X_internal_vector_values, r0+nrows, Kokkos::ALL(), 0, v);
4794 for (local_ordinal_type row = 0; row < nrows; ++row) {
4795 auto v_1 = Kokkos::subview(X_internal_vector_values, r0+row, Kokkos::ALL(), 0, v);
4796 auto E = Kokkos::subview(e_internal_vector_values, 0, r0+row, Kokkos::ALL(), Kokkos::ALL(), v);
4798 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
4799 (default_mode_type,default_algo_type,
4801 blocksize, blocksize,
4803 E.data(), E.stride_0(), E.stride_1(),
4804 v_2.data(), v_2.stride_0(),
4806 v_1.data(), v_1.stride_0());
4810 auto v_2 = Kokkos::subview(X_internal_vector_values, r0-1, Kokkos::ALL(), 0, v);
4812 for (local_ordinal_type row = 0; row < nrows; ++row) {
4813 auto v_1 = Kokkos::subview(X_internal_vector_values, r0+row, Kokkos::ALL(), 0, v);
4814 auto E = Kokkos::subview(e_internal_vector_values, 1, r0+row, Kokkos::ALL(), Kokkos::ALL(), v);
4816 KOKKOSBATCHED_GEMV_NO_TRANSPOSE_INTERNAL_INVOKE
4817 (default_mode_type,default_algo_type,
4819 blocksize, blocksize,
4821 E.data(), E.stride_0(), E.stride_1(),
4822 v_2.data(), v_2.stride_0(),
4824 v_1.data(), v_1.stride_0());
4832 KOKKOS_INLINE_FUNCTION
4834 operator() (
const SingleVectorCopyToFlatTag<B> &,
const member_type &member)
const {
4835 const local_ordinal_type packidx = member.league_rank();
4836 const local_ordinal_type partidx = packptr(packidx);
4837 const local_ordinal_type npacks = packptr(packidx+1) - partidx;
4838 const local_ordinal_type pri0 = part2packrowidx0(partidx);
4839 const local_ordinal_type blocksize = (B == 0 ? D_internal_vector_values.extent(1) : B);
4840 const local_ordinal_type num_vectors = 1;
4842 Kokkos::parallel_for
4843 (Kokkos::ThreadVectorRange(member, vector_loop_size),[&](
const int &v) {
4844 copyToFlatMultiVector(member, partidx, npacks, pri0, v, blocksize, num_vectors);
4849 KOKKOS_INLINE_FUNCTION
4851 operator() (
const SingleZeroingTag<B> &,
const member_type &member)
const {
4852 Kokkos::single(Kokkos::PerTeam(member), [&]() {
4853 Z_scalar_vector(member.league_rank()) = impl_scalar_type(0);
4857 void run(
const impl_scalar_type_2d_view_tpetra &Y,
4858 const impl_scalar_type_1d_view &Z) {
4859 IFPACK2_BLOCKTRIDICONTAINER_PROFILER_REGION_BEGIN;
4860 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::SolveTridiags", SolveTridiags);
4863 this->Y_scalar_multivector = Y;
4864 this->Z_scalar_vector = Z;
4866 const local_ordinal_type num_vectors = X_internal_vector_values.extent(2);
4867 const local_ordinal_type blocksize = D_internal_vector_values.extent(1);
4869 const local_ordinal_type team_size =
4870 SolveTridiagsDefaultModeAndAlgo<typename execution_space::memory_space>::
4871 recommended_team_size(blocksize, vector_length, internal_vector_length);
4872 const int per_team_scratch = internal_vector_scratch_type_3d_view
4873 ::shmem_size(blocksize, num_vectors, vector_loop_size);
4875 #if defined(KOKKOS_ENABLE_DEPRECATED_CODE)
4876 #define BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS(B) \
4877 if (num_vectors == 1) { \
4878 const Kokkos::TeamPolicy<execution_space,SingleVectorTag<B> > \
4879 policy(packptr.extent(0) - 1, team_size, vector_loop_size); \
4880 Kokkos::parallel_for \
4881 ("SolveTridiags::TeamPolicy::run<SingleVector>", \
4882 policy.set_scratch_size(0,Kokkos::PerTeam(per_team_scratch)), *this); \
4884 const Kokkos::TeamPolicy<execution_space,MultiVectorTag<B> > \
4885 policy(packptr.extent(0) - 1, team_size, vector_loop_size); \
4886 Kokkos::parallel_for \
4887 ("SolveTridiags::TeamPolicy::run<MultiVector>", \
4888 policy.set_scratch_size(0,Kokkos::PerTeam(per_team_scratch)), *this); \
4891 #define BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS(B) \
4892 if (num_vectors == 1) { \
4893 if (packindices_schur.extent(1) <= 0) { \
4894 Kokkos::TeamPolicy<execution_space,SingleVectorTag<B> > \
4895 policy(packptr.extent(0) - 1, team_size, vector_loop_size); \
4896 policy.set_scratch_size(0,Kokkos::PerTeam(per_team_scratch)); \
4897 Kokkos::parallel_for \
4898 ("SolveTridiags::TeamPolicy::run<SingleVector>", \
4904 Kokkos::TeamPolicy<execution_space,SingleZeroingTag<B> > \
4905 policy(packptr.extent(0) - 1, team_size, vector_loop_size); \
4906 Kokkos::parallel_for \
4907 ("SolveTridiags::TeamPolicy::run<SingleZeroingTag>", \
4911 IFPACK2_BLOCKHELPER_TIMER("BlockTriDi::ApplyInverseJacobi::SingleVectorSubLineTag", SingleVectorSubLineTag0); \
4912 write4DMultiVectorValuesToFile(part2packrowidx0_sub.extent(0), X_internal_scalar_values, "x_scalar_values_before_SingleVectorSubLineTag.mm"); \
4913 Kokkos::TeamPolicy<execution_space,SingleVectorSubLineTag<B> > \
4914 policy(packindices_sub.extent(0), team_size, vector_loop_size); \
4915 policy.set_scratch_size(0,Kokkos::PerTeam(per_team_scratch)); \
4916 Kokkos::parallel_for \
4917 ("SolveTridiags::TeamPolicy::run<SingleVector>", \
4919 write4DMultiVectorValuesToFile(part2packrowidx0_sub.extent(0), X_internal_scalar_values, "x_scalar_values_after_SingleVectorSubLineTag.mm"); \
4920 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space) \
4923 IFPACK2_BLOCKHELPER_TIMER("BlockTriDi::ApplyInverseJacobi::SingleVectorApplyCTag", SingleVectorApplyCTag0); \
4924 write4DMultiVectorValuesToFile(part2packrowidx0_sub.extent(0), X_internal_scalar_values, "x_scalar_values_before_SingleVectorApplyCTag.mm"); \
4925 Kokkos::TeamPolicy<execution_space,SingleVectorApplyCTag<B> > \
4926 policy(packindices_sub.extent(0), team_size, vector_loop_size); \
4927 policy.set_scratch_size(0,Kokkos::PerTeam(per_team_scratch)); \
4928 Kokkos::parallel_for \
4929 ("SolveTridiags::TeamPolicy::run<SingleVector>", \
4931 write4DMultiVectorValuesToFile(part2packrowidx0_sub.extent(0), X_internal_scalar_values, "x_scalar_values_after_SingleVectorApplyCTag.mm"); \
4932 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space) \
4935 IFPACK2_BLOCKHELPER_TIMER("BlockTriDi::ApplyInverseJacobi::SingleVectorSchurTag", SingleVectorSchurTag0); \
4936 write4DMultiVectorValuesToFile(part2packrowidx0_sub.extent(0), X_internal_scalar_values, "x_scalar_values_before_SingleVectorSchurTag.mm"); \
4937 Kokkos::TeamPolicy<execution_space,SingleVectorSchurTag<B> > \
4938 policy(packindices_schur.extent(0), team_size, vector_loop_size); \
4939 policy.set_scratch_size(0,Kokkos::PerTeam(per_team_scratch)); \
4940 Kokkos::parallel_for \
4941 ("SolveTridiags::TeamPolicy::run<SingleVector>", \
4943 write4DMultiVectorValuesToFile(part2packrowidx0_sub.extent(0), X_internal_scalar_values, "x_scalar_values_after_SingleVectorSchurTag.mm"); \
4944 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space) \
4947 IFPACK2_BLOCKHELPER_TIMER("BlockTriDi::ApplyInverseJacobi::SingleVectorApplyETag", SingleVectorApplyETag0); \
4948 write4DMultiVectorValuesToFile(part2packrowidx0_sub.extent(0), X_internal_scalar_values, "x_scalar_values_before_SingleVectorApplyETag.mm"); \
4949 Kokkos::TeamPolicy<execution_space,SingleVectorApplyETag<B> > \
4950 policy(packindices_sub.extent(0), team_size, vector_loop_size); \
4951 policy.set_scratch_size(0,Kokkos::PerTeam(per_team_scratch)); \
4952 Kokkos::parallel_for \
4953 ("SolveTridiags::TeamPolicy::run<SingleVector>", \
4955 write4DMultiVectorValuesToFile(part2packrowidx0_sub.extent(0), X_internal_scalar_values, "x_scalar_values_after_SingleVectorApplyETag.mm"); \
4956 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space) \
4960 Kokkos::TeamPolicy<execution_space,SingleVectorCopyToFlatTag<B> > \
4961 policy(packptr.extent(0) - 1, team_size, vector_loop_size); \
4962 Kokkos::parallel_for \
4963 ("SolveTridiags::TeamPolicy::run<SingleVectorCopyToFlatTag>", \
4968 Kokkos::TeamPolicy<execution_space,MultiVectorTag<B> > \
4969 policy(packptr.extent(0) - 1, team_size, vector_loop_size); \
4970 policy.set_scratch_size(0,Kokkos::PerTeam(per_team_scratch)); \
4971 Kokkos::parallel_for \
4972 ("SolveTridiags::TeamPolicy::run<MultiVector>", \
4976 switch (blocksize) {
4977 case 3: BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS( 3);
4978 case 5: BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS( 5);
4979 case 6: BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS( 6);
4980 case 7: BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS( 7);
4981 case 10: BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS(10);
4982 case 11: BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS(11);
4983 case 12: BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS(12);
4984 case 13: BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS(13);
4985 case 16: BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS(16);
4986 case 17: BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS(17);
4987 case 18: BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS(18);
4988 case 19: BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS(19);
4989 default : BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS( 0);
4991 #undef BLOCKTRIDICONTAINER_DETAILS_SOLVETRIDIAGS
4993 IFPACK2_BLOCKTRIDICONTAINER_PROFILER_REGION_END;
4994 IFPACK2_BLOCKHELPER_TIMER_FENCE(execution_space)
5001 template<
typename MatrixType>
5004 const Teuchos::RCP<
const typename BlockHelperDetails::ImplType<MatrixType>::tpetra_row_matrix_type> &A,
5005 const Teuchos::RCP<
const typename BlockHelperDetails::ImplType<MatrixType>::tpetra_crs_graph_type> &G,
5006 const Teuchos::RCP<
const typename BlockHelperDetails::ImplType<MatrixType>::tpetra_import_type> &tpetra_importer,
5007 const Teuchos::RCP<AsyncableImport<MatrixType> > &async_importer,
5008 const bool overlap_communication_and_computation,
5010 const typename BlockHelperDetails::ImplType<MatrixType>::tpetra_multivector_type &X,
5011 typename BlockHelperDetails::ImplType<MatrixType>::tpetra_multivector_type &Y,
5012 typename BlockHelperDetails::ImplType<MatrixType>::tpetra_multivector_type &Z,
5013 typename BlockHelperDetails::ImplType<MatrixType>::impl_scalar_type_1d_view &W,
5015 const BlockHelperDetails::PartInterface<MatrixType> &interf,
5018 typename BlockHelperDetails::ImplType<MatrixType>::vector_type_1d_view &work,
5023 const int max_num_sweeps,
5024 const typename BlockHelperDetails::ImplType<MatrixType>::magnitude_type tol,
5025 const int check_tol_every) {
5026 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::ApplyInverseJacobi", ApplyInverseJacobi);
5029 using node_memory_space =
typename impl_type::node_memory_space;
5030 using local_ordinal_type =
typename impl_type::local_ordinal_type;
5031 using size_type =
typename impl_type::size_type;
5032 using impl_scalar_type =
typename impl_type::impl_scalar_type;
5033 using magnitude_type =
typename impl_type::magnitude_type;
5034 using local_ordinal_type_1d_view =
typename impl_type::local_ordinal_type_1d_view;
5035 using vector_type_1d_view =
typename impl_type::vector_type_1d_view;
5036 using vector_type_3d_view =
typename impl_type::vector_type_3d_view;
5037 using tpetra_multivector_type =
typename impl_type::tpetra_multivector_type;
5039 using impl_scalar_type_1d_view =
typename impl_type::impl_scalar_type_1d_view;
5043 "Neither Tpetra importer nor Async importer is null.");
5046 "Maximum number of sweeps must be >= 1.");
5049 const bool is_seq_method_requested = !tpetra_importer.is_null();
5050 const bool is_async_importer_active = !async_importer.is_null();
5051 const bool is_norm_manager_active = tol > Kokkos::ArithTraits<magnitude_type>::zero();
5052 const magnitude_type tolerance = tol*tol;
5053 const local_ordinal_type blocksize = btdm.values.extent(1);
5054 const local_ordinal_type num_vectors = Y.getNumVectors();
5055 const local_ordinal_type num_blockrows = interf.part2packrowidx0_back;
5057 const impl_scalar_type zero(0.0);
5059 TEUCHOS_TEST_FOR_EXCEPT_MSG(is_norm_manager_active && is_seq_method_requested,
5060 "The seq method for applyInverseJacobi, " <<
5061 "which in any case is for developer use only, " <<
5062 "does not support norm-based termination.");
5063 const bool device_accessible_from_host = Kokkos::SpaceAccessibility<
5064 Kokkos::DefaultHostExecutionSpace, node_memory_space>::accessible;
5066 std::invalid_argument,
5067 "The seq method for applyInverseJacobi, " <<
5068 "which in any case is for developer use only, " <<
5069 "only supports memory spaces accessible from host.");
5072 const size_type work_span_required = num_blockrows*num_vectors*blocksize;
5073 if (work.span() < work_span_required)
5074 work = vector_type_1d_view(
"vector workspace 1d view", work_span_required);
5077 const local_ordinal_type W_size = interf.packptr.extent(0)-1;
5078 if (local_ordinal_type(W.extent(0)) < W_size)
5079 W = impl_scalar_type_1d_view(
"W", W_size);
5081 typename impl_type::impl_scalar_type_2d_view_tpetra remote_multivector;
5083 if (is_seq_method_requested) {
5084 if (Z.getNumVectors() != Y.getNumVectors())
5085 Z = tpetra_multivector_type(tpetra_importer->getTargetMap(), num_vectors,
false);
5087 if (is_async_importer_active) {
5089 async_importer->createDataBuffer(num_vectors);
5090 remote_multivector = async_importer->getRemoteMultiVectorLocalView();
5096 vector_type_3d_view pmv(work.data(), num_blockrows, blocksize, num_vectors);
5097 const auto XX = X.getLocalViewDevice(Tpetra::Access::ReadOnly);
5098 const auto YY = Y.getLocalViewDevice(Tpetra::Access::ReadWrite);
5099 const auto ZZ = Z.getLocalViewDevice(Tpetra::Access::ReadWrite);
5100 if (is_y_zero) Kokkos::deep_copy(YY, zero);
5103 SolveTridiags<MatrixType> solve_tridiags(interf, btdm, pmv,
5104 damping_factor, is_norm_manager_active);
5106 const local_ordinal_type_1d_view dummy_local_ordinal_type_1d_view;
5109 auto A_crs = Teuchos::rcp_dynamic_cast<
const typename impl_type::tpetra_crs_matrix_type>(A);
5110 auto A_bcrs = Teuchos::rcp_dynamic_cast<
const typename impl_type::tpetra_block_crs_matrix_type>(A);
5112 bool hasBlockCrsMatrix = ! A_bcrs.is_null ();
5115 const auto g = hasBlockCrsMatrix ? A_bcrs->getCrsGraph() : *(A_crs->getCrsGraph());
5117 BlockHelperDetails::ComputeResidualVector<MatrixType>
5118 compute_residual_vector(amd, G->getLocalGraphDevice(), g.getLocalGraphDevice(), blocksize, interf,
5119 is_async_importer_active ? async_importer->dm2cm : dummy_local_ordinal_type_1d_view,
5123 if (is_norm_manager_active)
5124 norm_manager.setCheckFrequency(check_tol_every);
5128 for (;sweep<max_num_sweeps;++sweep) {
5132 multivector_converter.run(XX);
5134 if (is_seq_method_requested) {
5138 Z.doImport(Y, *tpetra_importer, Tpetra::REPLACE);
5139 compute_residual_vector.run(YY, XX, ZZ);
5142 multivector_converter.run(YY);
5146 if (overlap_communication_and_computation || !is_async_importer_active) {
5147 if (is_async_importer_active) async_importer->asyncSendRecv(YY);
5149 compute_residual_vector.run(pmv, XX, YY, remote_multivector,
true);
5150 if (is_norm_manager_active && norm_manager.checkDone(sweep, tolerance)) {
5151 if (is_async_importer_active) async_importer->cancel();
5154 if (is_async_importer_active) {
5155 async_importer->syncRecv();
5157 compute_residual_vector.run(pmv, XX, YY, remote_multivector,
false);
5160 if (is_async_importer_active)
5161 async_importer->syncExchange(YY);
5162 if (is_norm_manager_active && norm_manager.checkDone(sweep, tolerance))
break;
5164 compute_residual_vector.run(pmv, XX, YY, remote_multivector);
5172 solve_tridiags.run(YY, W);
5175 if (is_norm_manager_active) {
5177 BlockHelperDetails::reduceVector<MatrixType>(W, norm_manager.getBuffer());
5178 if (sweep + 1 == max_num_sweeps) {
5179 norm_manager.ireduce(sweep,
true);
5180 norm_manager.checkDone(sweep + 1, tolerance,
true);
5182 norm_manager.ireduce(sweep);
5190 if (is_norm_manager_active) norm_manager.finalize();
5197 template<
typename MatrixType,
int B>
5199 applyFusedBlockJacobi_Impl(
5200 const Teuchos::RCP<
const typename BlockHelperDetails::ImplType<MatrixType>::tpetra_import_type> &tpetra_importer,
5201 const Teuchos::RCP<AsyncableImport<MatrixType> > &async_importer,
5202 const bool overlap_communication_and_computation,
5204 const typename BlockHelperDetails::ImplType<MatrixType>::tpetra_multivector_type &X,
5205 typename BlockHelperDetails::ImplType<MatrixType>::tpetra_multivector_type &Y,
5206 typename BlockHelperDetails::ImplType<MatrixType>::impl_scalar_type_1d_view &W,
5208 const BlockHelperDetails::PartInterface<MatrixType> &interf,
5209 const BlockTridiags<MatrixType> &btdm,
5211 typename BlockHelperDetails::ImplType<MatrixType>::impl_scalar_type_1d_view &work,
5216 const int max_num_sweeps,
5217 const typename BlockHelperDetails::ImplType<MatrixType>::magnitude_type tol,
5218 const int check_tol_every) {
5220 using node_memory_space =
typename impl_type::node_memory_space;
5221 using local_ordinal_type =
typename impl_type::local_ordinal_type;
5222 using size_type =
typename impl_type::size_type;
5223 using impl_scalar_type =
typename impl_type::impl_scalar_type;
5224 using magnitude_type =
typename impl_type::magnitude_type;
5225 using local_ordinal_type_1d_view =
typename impl_type::local_ordinal_type_1d_view;
5226 using tpetra_multivector_type =
typename impl_type::tpetra_multivector_type;
5227 using impl_scalar_type_1d_view =
typename impl_type::impl_scalar_type_1d_view;
5228 using impl_scalar_type_2d_view_tpetra =
typename impl_type::impl_scalar_type_2d_view_tpetra;
5232 "Neither Tpetra importer nor Async importer is null.");
5235 "Maximum number of sweeps must be >= 1.");
5238 const bool is_async_importer_active = !async_importer.is_null();
5239 const bool is_norm_manager_active = tol > Kokkos::ArithTraits<magnitude_type>::zero();
5240 const magnitude_type tolerance = tol*tol;
5241 const local_ordinal_type blocksize = btdm.d_inv.extent(1);
5242 const local_ordinal_type num_vectors = Y.getNumVectors();
5243 const local_ordinal_type num_blockrows = interf.nparts;
5245 typename impl_type::impl_scalar_type_2d_view_tpetra remote_multivector;
5247 if (is_async_importer_active) {
5249 async_importer->createDataBuffer(num_vectors);
5250 remote_multivector = async_importer->getRemoteMultiVectorLocalView();
5254 const auto XX = X.getLocalViewDevice(Tpetra::Access::ReadOnly);
5255 const auto YY = Y.getLocalViewDevice(Tpetra::Access::ReadWrite);
5257 const bool two_pass_residual =
5258 overlap_communication_and_computation && is_async_importer_active;
5263 size_t(num_blockrows) * blocksize * num_vectors != YY.extent(0) * YY.extent(1),
5264 "Local LHS vector (YY) has total size " << YY.extent(0) <<
"x" << YY.extent(1) <<
5265 " = " << YY.extent(0) * YY.extent(1) <<
",\n" <<
5266 "but expected " << num_blockrows <<
"x" << blocksize <<
"x" << num_vectors <<
5267 " = " << size_t(num_blockrows) * blocksize * num_vectors <<
'\n');
5268 size_type work_required = size_type(num_blockrows) * blocksize * num_vectors;
5269 if (work.extent(0) < work_required) {
5273 Unmanaged<impl_scalar_type_2d_view_tpetra> y_doublebuf(work.data(), num_blockrows * blocksize, num_vectors);
5276 if (W.extent(0) != size_t(num_blockrows))
5280 BlockHelperDetails::ComputeResidualAndSolve_SolveOnly<MatrixType, B>
5281 functor_solve_only(amd, btdm.d_inv, W, blocksize, damping_factor);
5282 BlockHelperDetails::ComputeResidualAndSolve_1Pass<MatrixType, B>
5283 functor_1pass(amd, btdm.d_inv, W, blocksize, damping_factor);
5284 BlockHelperDetails::ComputeResidualAndSolve_2Pass<MatrixType, B>
5285 functor_2pass(amd, btdm.d_inv, W, blocksize, damping_factor);
5288 if (is_norm_manager_active)
5289 norm_manager.setCheckFrequency(check_tol_every);
5294 Unmanaged<impl_scalar_type_2d_view_tpetra> y_buffers[2] = {YY, y_doublebuf};
5299 for (;sweep < max_num_sweeps; ++sweep) {
5302 functor_solve_only.run(XX, y_buffers[1-current_y]);
5305 if (overlap_communication_and_computation || !is_async_importer_active) {
5306 if (is_async_importer_active) async_importer->asyncSendRecv(y_buffers[current_y]);
5307 if(two_pass_residual) {
5310 functor_2pass.run_pass1(XX, y_buffers[current_y], y_buffers[1-current_y]);
5315 functor_1pass.run(XX, y_buffers[current_y], remote_multivector, y_buffers[1-current_y]);
5317 if (is_norm_manager_active && norm_manager.checkDone(sweep, tolerance)) {
5318 if (is_async_importer_active) async_importer->cancel();
5321 if (is_async_importer_active) {
5322 async_importer->syncRecv();
5324 functor_2pass.run_pass2(y_buffers[current_y], remote_multivector, y_buffers[1-current_y]);
5327 if (is_async_importer_active)
5328 async_importer->syncExchange(y_buffers[current_y]);
5329 if (is_norm_manager_active && norm_manager.checkDone(sweep, tolerance))
break;
5331 functor_1pass.run(XX, y_buffers[current_y], remote_multivector, y_buffers[1-current_y]);
5336 if (is_norm_manager_active) {
5337 BlockHelperDetails::reduceVector<MatrixType>(W, norm_manager.getBuffer());
5338 if (sweep + 1 == max_num_sweeps) {
5339 norm_manager.ireduce(sweep,
true);
5340 norm_manager.checkDone(sweep + 1, tolerance,
true);
5342 norm_manager.ireduce(sweep);
5347 current_y = 1 - current_y;
5349 if(current_y == 1) {
5351 Kokkos::deep_copy(YY, y_doublebuf);
5355 if (is_norm_manager_active) norm_manager.finalize();
5362 template<
typename MatrixType>
5365 const Teuchos::RCP<
const typename BlockHelperDetails::ImplType<MatrixType>::tpetra_import_type> &tpetra_importer,
5366 const Teuchos::RCP<AsyncableImport<MatrixType> > &async_importer,
5367 const bool overlap_communication_and_computation,
5369 const typename BlockHelperDetails::ImplType<MatrixType>::tpetra_multivector_type &X,
5370 typename BlockHelperDetails::ImplType<MatrixType>::tpetra_multivector_type &Y,
5371 typename BlockHelperDetails::ImplType<MatrixType>::impl_scalar_type_1d_view &W,
5373 const BlockHelperDetails::PartInterface<MatrixType> &interf,
5376 typename BlockHelperDetails::ImplType<MatrixType>::impl_scalar_type_1d_view &work,
5381 const int max_num_sweeps,
5382 const typename BlockHelperDetails::ImplType<MatrixType>::magnitude_type tol,
5383 const int check_tol_every) {
5384 IFPACK2_BLOCKHELPER_TIMER(
"BlockTriDi::ApplyFusedBlockJacobi", ApplyFusedBlockJacobi);
5385 int blocksize = btdm.d_inv.extent(1);
5387 #define BLOCKTRIDICONTAINER_APPLY_FUSED_JACOBI(B) { \
5388 sweep = applyFusedBlockJacobi_Impl<MatrixType, B>( \
5389 tpetra_importer, async_importer, overlap_communication_and_computation, \
5390 X, Y, W, interf, btdm, amd, work, \
5391 norm_manager, damping_factor, is_y_zero, \
5392 max_num_sweeps, tol, check_tol_every); \
5394 switch (blocksize) {
5395 case 3: BLOCKTRIDICONTAINER_APPLY_FUSED_JACOBI( 3);
5396 case 5: BLOCKTRIDICONTAINER_APPLY_FUSED_JACOBI( 5);
5397 case 7: BLOCKTRIDICONTAINER_APPLY_FUSED_JACOBI( 7);
5398 case 9: BLOCKTRIDICONTAINER_APPLY_FUSED_JACOBI( 9);
5399 case 10: BLOCKTRIDICONTAINER_APPLY_FUSED_JACOBI(10);
5400 case 11: BLOCKTRIDICONTAINER_APPLY_FUSED_JACOBI(11);
5401 case 16: BLOCKTRIDICONTAINER_APPLY_FUSED_JACOBI(16);
5402 case 17: BLOCKTRIDICONTAINER_APPLY_FUSED_JACOBI(17);
5403 case 18: BLOCKTRIDICONTAINER_APPLY_FUSED_JACOBI(18);
5404 default : BLOCKTRIDICONTAINER_APPLY_FUSED_JACOBI( 0);
5406 #undef BLOCKTRIDICONTAINER_APPLY_FUSED_JACOBI
5412 template<
typename MatrixType>
5415 using part_interface_type = BlockHelperDetails::PartInterface<MatrixType>;
5416 using block_tridiags_type = BlockTridiags<MatrixType>;
5419 using async_import_type = AsyncableImport<MatrixType>;
5426 bool overlap_communication_and_computation;
5429 mutable typename impl_type::tpetra_multivector_type Z;
5430 mutable typename impl_type::impl_scalar_type_1d_view W;
5433 part_interface_type part_interface;
5434 block_tridiags_type block_tridiags;
5438 bool use_fused_jacobi;
5441 mutable typename impl_type::vector_type_1d_view work;
5443 mutable typename impl_type::impl_scalar_type_1d_view work_flat;
5444 mutable norm_manager_type norm_manager;
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:141
void performNumericPhase(const Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_row_matrix_type > &A, const Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_crs_graph_type > &G, const BlockHelperDetails::PartInterface< MatrixType > &interf, BlockTridiags< MatrixType > &btdm, const typename BlockHelperDetails::ImplType< MatrixType >::magnitude_type tiny, bool use_fused_jacobi)
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:3710
#define TEUCHOS_TEST_FOR_EXCEPTION(throw_exception_test, Exception, msg)
size_t size_type
Definition: Ifpack2_BlockHelper.hpp:253
Teuchos::RCP< AsyncableImport< MatrixType > > createBlockCrsAsyncImporter(const Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_row_matrix_type > &A)
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:885
int applyInverseJacobi(const Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_row_matrix_type > &A, const Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_crs_graph_type > &G, const Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_import_type > &tpetra_importer, const Teuchos::RCP< AsyncableImport< MatrixType > > &async_importer, const bool overlap_communication_and_computation, const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_multivector_type &X, typename BlockHelperDetails::ImplType< MatrixType >::tpetra_multivector_type &Y, typename BlockHelperDetails::ImplType< MatrixType >::tpetra_multivector_type &Z, typename BlockHelperDetails::ImplType< MatrixType >::impl_scalar_type_1d_view &W, const BlockHelperDetails::PartInterface< MatrixType > &interf, const BlockTridiags< MatrixType > &btdm, const BlockHelperDetails::AmD< MatrixType > &amd, typename BlockHelperDetails::ImplType< MatrixType >::vector_type_1d_view &work, BlockHelperDetails::NormManager< MatrixType > &norm_manager, const typename BlockHelperDetails::ImplType< MatrixType >::impl_scalar_type &damping_factor, bool is_y_zero, const int max_num_sweeps, const typename BlockHelperDetails::ImplType< MatrixType >::magnitude_type tol, const int check_tol_every)
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:5003
TEUCHOS_DEPRECATED RCP< T > rcp(T *p, Dealloc_T dealloc, bool owns_mem)
#define TEUCHOS_TEST_FOR_EXCEPT_MSG(throw_exception_test, msg)
BlockTridiags< MatrixType > createBlockTridiags(const BlockHelperDetails::PartInterface< MatrixType > &interf)
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:1623
Definition: Ifpack2_BlockHelper.hpp:353
void performSymbolicPhase(const Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_row_matrix_type > &A, const Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_crs_graph_type > &g, const BlockHelperDetails::PartInterface< MatrixType > &interf, BlockTridiags< MatrixType > &btdm, BlockHelperDetails::AmD< MatrixType > &amd, const bool overlap_communication_and_computation, const Teuchos::RCP< AsyncableImport< MatrixType > > &async_importer, bool useSeqMethod, bool use_fused_jacobi)
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:1865
Kokkos::ViewAllocateWithoutInitializing do_not_initialize_tag
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:97
int applyFusedBlockJacobi(const Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_import_type > &tpetra_importer, const Teuchos::RCP< AsyncableImport< MatrixType > > &async_importer, const bool overlap_communication_and_computation, const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_multivector_type &X, typename BlockHelperDetails::ImplType< MatrixType >::tpetra_multivector_type &Y, typename BlockHelperDetails::ImplType< MatrixType >::impl_scalar_type_1d_view &W, const BlockHelperDetails::PartInterface< MatrixType > &interf, const BlockTridiags< MatrixType > &btdm, const BlockHelperDetails::AmD< MatrixType > &amd, typename BlockHelperDetails::ImplType< MatrixType >::impl_scalar_type_1d_view &work, BlockHelperDetails::NormManager< MatrixType > &norm_manager, const typename BlockHelperDetails::ImplType< MatrixType >::impl_scalar_type &damping_factor, bool is_y_zero, const int max_num_sweeps, const typename BlockHelperDetails::ImplType< MatrixType >::magnitude_type tol, const int check_tol_every)
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:5364
void send(const Packet sendBuffer[], const Ordinal count, const int destRank, const int tag, const Comm< Ordinal > &comm)
Kokkos::Details::ArithTraits< scalar_type >::val_type impl_scalar_type
Definition: Ifpack2_BlockHelper.hpp:262
Definition: Ifpack2_BlockHelper.hpp:188
BlockHelperDetails::PartInterface< MatrixType > createPartInterface(const Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_row_matrix_type > &A, const Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_crs_graph_type > &G, const Teuchos::Array< Teuchos::Array< typename BlockHelperDetails::ImplType< MatrixType >::local_ordinal_type > > &partitions, const typename BlockHelperDetails::ImplType< MatrixType >::local_ordinal_type n_subparts_per_part_in)
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:1044
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:2216
Kokkos::View< size_type *, device_type > size_type_1d_view
Definition: Ifpack2_BlockHelper.hpp:321
Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_import_type > createBlockCrsTpetraImporter(const Teuchos::RCP< const typename BlockHelperDetails::ImplType< MatrixType >::tpetra_row_matrix_type > &A)
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:164
RCP< CommRequest< Ordinal > > isend(const ArrayRCP< const Packet > &sendBuffer, const int destRank, const int tag, const Comm< Ordinal > &comm)
#define TEUCHOS_ASSERT(assertion_test)
Definition: Ifpack2_BlockHelper.hpp:215
Definition: Ifpack2_BlockHelper.hpp:249
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:1558
Definition: Ifpack2_BlockComputeResidualVector.hpp:23
Definition: Ifpack2_BlockTriDiContainer_impl.hpp:3762