Tpetra parallel linear algebra Version of the Day
Loading...
Searching...
No Matches
TpetraExt_MatrixMatrix_Cuda.hpp
1// @HEADER
2// *****************************************************************************
3// Tpetra: Templated Linear Algebra Services Package
4//
5// Copyright 2008 NTESS and the Tpetra contributors.
6// SPDX-License-Identifier: BSD-3-Clause
7// *****************************************************************************
8// @HEADER
9
10#ifndef TPETRA_MATRIXMATRIX_CUDA_DEF_HPP
11#define TPETRA_MATRIXMATRIX_CUDA_DEF_HPP
12
13#include "Tpetra_Details_IntRowPtrHelper.hpp"
14
15#ifdef HAVE_TPETRA_INST_CUDA
16namespace Tpetra {
17namespace MMdetails {
18
19template <>
20struct KokkosKernelsSPGEMMBackend<Tpetra::KokkosCompat::KokkosCudaWrapperNode> {
21 static std::string parameter_prefix() { return "cuda"; }
22 static std::string algorithm_label() { return "Cuda"; }
23
24 template <class MatrixType>
25 static void pre_spgemm(MatrixType& Bmerged) {
26#if defined(KOKKOS_ENABLE_CUDA) && defined(KOKKOSKERNELS_ENABLE_TPL_CUSPARSE) && ((CUDA_VERSION < 11000) || (CUDA_VERSION >= 11040))
27 using device_t = typename MatrixType::device_type;
28 if constexpr (std::is_same_v<typename device_t::execution_space, Kokkos::Cuda>) {
29 if (!KokkosSparse::isCrsGraphSorted(Bmerged.graph.row_map, Bmerged.graph.entries)) {
30 Import_Util::sortCrsMatrix(Bmerged);
31 }
32 }
33#else
34 (void)Bmerged;
35#endif
36 }
37};
38
39/*********************************************************************************************************/
40// MMM KernelWrappers for Partial Specialization to CUDA
41template <class Scalar,
42 class LocalOrdinal,
43 class GlobalOrdinal,
44 class LocalOrdinalViewType>
45struct KernelWrappers<Scalar, LocalOrdinal, GlobalOrdinal, Tpetra::KokkosCompat::KokkosCudaWrapperNode, LocalOrdinalViewType> {
46 using Node = Tpetra::KokkosCompat::KokkosCudaWrapperNode;
47
48 static inline void mult_A_B_newmatrix_kernel_wrapper(CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
49 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
50 const LocalOrdinalViewType& Acol2Brow,
51 const LocalOrdinalViewType& Acol2Irow,
52 const LocalOrdinalViewType& Bcol2Ccol,
53 const LocalOrdinalViewType& Icol2Ccol,
54 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
55 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node> > Cimport,
56 const std::string& label = std::string(),
57 const Teuchos::RCP<Teuchos::ParameterList>& params = Teuchos::null) {
58 Tpetra::MMdetails::kokkos_kernels_mult_A_B_newmatrix(
59 Aview, Bview, Acol2Brow, Acol2Irow, Bcol2Ccol, Icol2Ccol, C, Cimport, label, params);
60 }
61
62 static inline void mult_A_B_reuse_kernel_wrapper(CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
63 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
64 const LocalOrdinalViewType& Acol2Brow,
65 const LocalOrdinalViewType& Acol2Irow,
66 const LocalOrdinalViewType& Bcol2Ccol,
67 const LocalOrdinalViewType& Icol2Ccol,
68 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
69 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node> > Cimport,
70 const std::string& label = std::string(),
71 const Teuchos::RCP<Teuchos::ParameterList>& params = Teuchos::null) {
72 Tpetra::MMdetails::host_mult_A_B_reuse(
73 Aview, Bview, Acol2Brow, Acol2Irow,
74 Bcol2Ccol, Icol2Ccol, C, Cimport, label, params);
75 }
76};
77
78// Jacobi KernelWrappers for Partial Specialization to Cuda
79template <class Scalar,
80 class LocalOrdinal,
81 class GlobalOrdinal, class LocalOrdinalViewType>
82struct KernelWrappers2<Scalar, LocalOrdinal, GlobalOrdinal, Tpetra::KokkosCompat::KokkosCudaWrapperNode, LocalOrdinalViewType> {
83 using Node = Tpetra::KokkosCompat::KokkosCudaWrapperNode;
84
85 static inline void jacobi_A_B_newmatrix_kernel_wrapper(typename Teuchos::ScalarTraits<Scalar>::magnitudeType omega,
86 const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Dinv,
87 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
88 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
89 const LocalOrdinalViewType& Acol2Brow,
90 const LocalOrdinalViewType& Acol2Irow,
91 const LocalOrdinalViewType& Bcol2Ccol,
92 const LocalOrdinalViewType& Icol2Ccol,
93 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
94 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node> > Cimport,
95 const std::string& label = std::string(),
96 const Teuchos::RCP<Teuchos::ParameterList>& params = Teuchos::null) {
97 // Node-specific code
98 using Teuchos::RCP;
99 using Teuchos::rcp;
100 RCP<Tpetra::Details::ProfilingRegion> MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM: Jacobi CudaWrapper"));
101
102 // Options
103 // int team_work_size = 16; // Defaults to 16 as per Deveci 12/7/16 - csiefer // unreferenced
104 std::string myalg("KK");
105 if (!params.is_null()) {
106 if (params->isParameter("cuda: jacobi algorithm"))
107 myalg = params->get("cuda: jacobi algorithm", myalg);
108 }
109
110 if (myalg == "MSAK") {
111 ::Tpetra::MatrixMatrix::ExtraKernels::jacobi_A_B_newmatrix_MultiplyScaleAddKernel(omega, Dinv, Aview, Bview, Acol2Brow, Acol2Irow, Bcol2Ccol, Icol2Ccol, C, Cimport, label, params);
112 } else if (myalg == "KK") {
113 kokkos_kernels_jacobi_A_B_newmatrix(omega, Dinv, Aview, Bview, Acol2Brow, Acol2Irow, Bcol2Ccol, Icol2Ccol, C, Cimport, label, params);
114 } else {
115 throw std::runtime_error("Tpetra::MatrixMatrix::Jacobi newmatrix unknown kernel");
116 }
117
118 MM = Teuchos::null;
119 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: Jacobi: Newmatrix CudaESFC"));
120
121 // Final Fillcomplete
122 RCP<Teuchos::ParameterList> labelList = rcp(new Teuchos::ParameterList);
123 labelList->set("Timer Label", label);
124 if (!params.is_null()) labelList->set("compute global constants", params->get("compute global constants", true));
125
126 // NOTE: MSAK already fillCompletes, so we have to check here
127 if (!C.isFillComplete()) {
128 RCP<const Export<LocalOrdinal, GlobalOrdinal, Node> > dummyExport;
129 C.expertStaticFillComplete(Bview.origMatrix->getDomainMap(), Aview.origMatrix->getRangeMap(), Cimport, dummyExport, labelList);
130 }
131 }
132
133 static inline void jacobi_A_B_reuse_kernel_wrapper(typename Teuchos::ScalarTraits<Scalar>::magnitudeType omega,
134 const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Dinv,
135 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
136 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
137 const LocalOrdinalViewType& Acol2Brow,
138 const LocalOrdinalViewType& Acol2Irow,
139 const LocalOrdinalViewType& Bcol2Ccol,
140 const LocalOrdinalViewType& Icol2Ccol,
141 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
142 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node> > Cimport,
143 const std::string& label = std::string(),
144 const Teuchos::RCP<Teuchos::ParameterList>& params = Teuchos::null) {
145 host_jacobi_A_B_reuse(omega, Dinv, Aview, Bview, Acol2Brow, Acol2Irow, Bcol2Ccol, Icol2Ccol, C, Cimport, label, params);
146 }
147};
148
149} // namespace MMdetails
150} // namespace Tpetra
151
152#endif // CUDA
153
154#endif
Namespace Tpetra contains the class and methods constituting the Tpetra library.