Tpetra parallel linear algebra Version of the Day
Loading...
Searching...
No Matches
Tpetra_Filter.hpp
1#ifndef TPETRA_FILTER_HPP
2#define TPETRA_FILTER_HPP
3
4#include "Kokkos_Core.hpp"
5#include "Teuchos_RCP.hpp"
6
7namespace Tpetra {
8
16template <class matrix_type, class filter_type>
17Teuchos::RCP<matrix_type> applyFilter_vals(const matrix_type& A, const filter_type& filter) {
18 using local_ordinal_type = typename matrix_type::local_ordinal_type;
19 using node_type = typename matrix_type::node_type;
20 using local_matrix_type = typename matrix_type::local_matrix_device_type;
21 using rowptr_type = typename local_matrix_type::row_map_type::non_const_type;
22 using colidx_type = typename local_matrix_type::index_type::non_const_type;
23 using values_type = typename local_matrix_type::values_type::non_const_type;
24 using exec_space = typename node_type::execution_space;
25
26 local_matrix_type newLclA;
27 {
29 auto lclA = A.getLocalMatrixDevice();
30
31 rowptr = rowptr_type("filtered_rowptr", lclA.numRows() + 1);
32 local_ordinal_type NNZ;
33
34 Kokkos::parallel_scan(
35 "filterMatrix::rowptr_construction", Kokkos::RangePolicy<exec_space>(0, lclA.numRows()), KOKKOS_LAMBDA(const local_ordinal_type rlid, local_ordinal_type& nnz, const bool update) {
36 auto row = lclA.rowConst(rlid);
37 for (local_ordinal_type k = 0; k < row.length; ++k) {
38 auto val = row.value(k);
39 if (filter(val)) {
40 ++nnz;
41 }
42 }
43 if (update && (rlid + 2 < lclA.numRows() + 1))
44 rowptr(rlid + 2) = nnz;
45 },
46 NNZ);
47
48 auto colidx = colidx_type(Kokkos::ViewAllocateWithoutInitializing("filtered_colidx"), NNZ);
49 auto values = values_type(Kokkos::ViewAllocateWithoutInitializing("filtered_values"), NNZ);
50
51 Kokkos::parallel_for(
52 "filterMatrix::matrix_fill", Kokkos::RangePolicy<exec_space>(0, lclA.numRows()), KOKKOS_LAMBDA(const local_ordinal_type rlid) {
53 auto row = lclA.rowConst(rlid);
54 for (local_ordinal_type k = 0; k < row.length; ++k) {
55 auto val = row.value(k);
56 if (filter(val)) {
57 auto clid = row.colidx(k);
58 colidx(rowptr(rlid + 1)) = clid;
59 values(rowptr(rlid + 1)) = val;
60 ++rowptr(rlid + 1);
61 }
62 }
63 });
64
65 newLclA = local_matrix_type("filtered_matrix", lclA.numRows(), lclA.numCols(), NNZ, values, rowptr, colidx);
66 }
67
68 return Teuchos::rcp(new matrix_type(newLclA, A.getRowMap(), A.getColMap(), A.getDomainMap(), A.getRangeMap()));
69}
70
78template <class matrix_type, class filter_type>
79Teuchos::RCP<matrix_type> applyFilter_LID(const matrix_type& A, const filter_type& filter) {
80 using local_ordinal_type = typename matrix_type::local_ordinal_type;
81 using node_type = typename matrix_type::node_type;
82 using local_matrix_type = typename matrix_type::local_matrix_device_type;
83 using rowptr_type = typename local_matrix_type::row_map_type::non_const_type;
84 using colidx_type = typename local_matrix_type::index_type::non_const_type;
85 using values_type = typename local_matrix_type::values_type::non_const_type;
86 using exec_space = typename node_type::execution_space;
87
88 local_matrix_type newLclA;
89 {
91 auto lclA = A.getLocalMatrixDevice();
92
93 rowptr = rowptr_type("filtered_rowptr", lclA.numRows() + 1);
94 local_ordinal_type NNZ;
95
96 Kokkos::parallel_scan(
97 "filterMatrix::rowptr_construction", Kokkos::RangePolicy<exec_space>(0, lclA.numRows()), KOKKOS_LAMBDA(const local_ordinal_type rlid, local_ordinal_type& nnz, const bool update) {
98 auto row = lclA.rowConst(rlid);
99 for (local_ordinal_type k = 0; k < row.length; ++k) {
100 auto clid = row.colidx(k);
101 auto val = row.value(k);
102 if (filter(rlid, clid, val)) {
103 ++nnz;
104 }
105 }
106 if (update && (rlid + 2 < lclA.numRows() + 1))
107 rowptr(rlid + 2) = nnz;
108 },
109 NNZ);
110
111 auto colidx = colidx_type(Kokkos::ViewAllocateWithoutInitializing("filtered_colidx"), NNZ);
112 auto values = values_type(Kokkos::ViewAllocateWithoutInitializing("filtered_values"), NNZ);
113
114 Kokkos::parallel_for(
115 "filterMatrix::matrix_fill", Kokkos::RangePolicy<exec_space>(0, lclA.numRows()), KOKKOS_LAMBDA(const local_ordinal_type rlid) {
116 auto row = lclA.rowConst(rlid);
117 for (local_ordinal_type k = 0; k < row.length; ++k) {
118 auto clid = row.colidx(k);
119 auto val = row.value(k);
120 if (filter(rlid, clid, val)) {
121 colidx(rowptr(rlid + 1)) = clid;
122 values(rowptr(rlid + 1)) = val;
123 ++rowptr(rlid + 1);
124 }
125 }
126 });
127
128 newLclA = local_matrix_type("filtered_matrix", lclA.numRows(), lclA.numCols(), NNZ, values, rowptr, colidx);
129 }
130
131 return Teuchos::rcp(new matrix_type(newLclA, A.getRowMap(), A.getColMap(), A.getDomainMap(), A.getRangeMap()));
132}
133
141template <class matrix_type, class filter_type>
142Teuchos::RCP<matrix_type> applyFilter_GID(const matrix_type& A, const filter_type& filter) {
143 using local_ordinal_type = typename matrix_type::local_ordinal_type;
144 using node_type = typename matrix_type::node_type;
145 using local_matrix_type = typename matrix_type::local_matrix_device_type;
146 using rowptr_type = typename local_matrix_type::row_map_type::non_const_type;
147 using colidx_type = typename local_matrix_type::index_type::non_const_type;
148 using values_type = typename local_matrix_type::values_type::non_const_type;
149 using exec_space = typename node_type::execution_space;
150
151 local_matrix_type newLclA;
152 {
154 auto lclA = A.getLocalMatrixDevice();
155
156 rowptr = rowptr_type("filtered_rowptr", lclA.numRows() + 1);
157 local_ordinal_type NNZ;
158
159 auto lclRowMap = A.getRowMap()->getLocalMap();
160 auto lclColMap = A.getColMap()->getLocalMap();
161
162 Kokkos::parallel_scan(
163 "filterMatrix::rowptr_construction", Kokkos::RangePolicy<exec_space>(0, lclA.numRows()), KOKKOS_LAMBDA(const local_ordinal_type rlid, local_ordinal_type& nnz, const bool update) {
164 auto row = lclA.rowConst(rlid);
165 auto rgid = lclRowMap.getGlobalElement(rlid);
166 for (local_ordinal_type k = 0; k < row.length; ++k) {
167 auto clid = row.colidx(k);
168 auto cgid = lclRowMap.getGlobalElement(clid);
169 auto val = row.value(k);
170 if (filter(rgid, cgid, val)) {
171 ++nnz;
172 }
173 }
174 if (update && (rlid + 2 < lclA.numRows() + 1))
175 rowptr(rlid + 2) = nnz;
176 },
177 NNZ);
178
179 auto colidx = colidx_type(Kokkos::ViewAllocateWithoutInitializing("filtered_colidx"), NNZ);
180 auto values = values_type(Kokkos::ViewAllocateWithoutInitializing("filtered_values"), NNZ);
181
182 Kokkos::parallel_for(
183 "filterMatrix::matrix_fill", Kokkos::RangePolicy<exec_space>(0, lclA.numRows()), KOKKOS_LAMBDA(const local_ordinal_type rlid) {
184 auto row = lclA.rowConst(rlid);
185 auto rgid = lclRowMap.getGlobalElement(rlid);
186 for (local_ordinal_type k = 0; k < row.length; ++k) {
187 auto clid = row.colidx(k);
188 auto cgid = lclRowMap.getGlobalElement(clid);
189 auto val = row.value(k);
190 if (filter(rgid, cgid, val)) {
191 colidx(rowptr(rlid + 1)) = clid;
192 values(rowptr(rlid + 1)) = val;
193 ++rowptr(rlid + 1);
194 }
195 }
196 });
197
198 newLclA = local_matrix_type("filtered_matrix", lclA.numRows(), lclA.numCols(), NNZ, values, rowptr, colidx);
199 }
200
201 return Teuchos::rcp(new matrix_type(newLclA, A.getRowMap(), A.getColMap(), A.getDomainMap(), A.getRangeMap()));
202}
203
204template <class scalar_type, class local_ordinal_type>
205struct AbsoluteMagnitudeFilter {
206#if KOKKOS_VERSION >= 40799
207 using ATS = KokkosKernels::ArithTraits<scalar_type>;
208 using impl_scalar_type = typename ATS::val_type;
209 using implATS = KokkosKernels::ArithTraits<impl_scalar_type>;
210#else
211 using ATS = Kokkos::ArithTraits<scalar_type>;
212 using impl_scalar_type = typename ATS::val_type;
213 using implATS = Kokkos::ArithTraits<impl_scalar_type>;
214#endif
215
216 typename implATS::magnitudeType tol_;
217
218 AbsoluteMagnitudeFilter(typename implATS::magnitudeType tol)
219 : tol_(tol) {}
220
221 KOKKOS_INLINE_FUNCTION
222 bool operator()(impl_scalar_type val) const {
223 return (implATS::magnitude(val) > tol_);
224 }
225};
226
227} // namespace Tpetra
228#endif
Struct that holds views of the contents of a CrsMatrix.
Namespace Tpetra contains the class and methods constituting the Tpetra library.
Teuchos::RCP< matrix_type > applyFilter_GID(const matrix_type &A, const filter_type &filter)
Constructs a matrix consisting only of the entries for which the filter function evaluates to true.
Teuchos::RCP< matrix_type > applyFilter_LID(const matrix_type &A, const filter_type &filter)
Constructs a matrix consisting only of the entries for which the filter function evaluates to true.
Teuchos::RCP< matrix_type > applyFilter_vals(const matrix_type &A, const filter_type &filter)
Constructs a matrix consisting only of the entries for which the filter function evaluates to true.