42 #ifndef STOKHOS_TPETRA_UTILITIES_UQ_PCE_HPP 43 #define STOKHOS_TPETRA_UTILITIES_UQ_PCE_HPP 47 #include "Tpetra_Map.hpp" 48 #include "Tpetra_MultiVector.hpp" 49 #include "Tpetra_CrsGraph.hpp" 50 #include "Tpetra_CrsMatrix.hpp" 54 #if defined(TPETRA_HAVE_KOKKOS_REFACTOR) 60 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
61 create_cijk_crs_graph(
const CijkType&
cijk,
62 const Teuchos::RCP<
const Teuchos::Comm<int> >& comm,
63 const size_t matrix_pce_size) {
65 using Teuchos::arrayView;
67 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
68 typedef Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> Map;
69 typedef Tpetra::CrsGraph<LocalOrdinal,GlobalOrdinal,Node> Graph;
71 const size_t pce_sz =
cijk.dimension();
73 Tpetra::createLocalMapWithNode<LocalOrdinal,GlobalOrdinal,Node>(pce_sz, comm);
75 if (matrix_pce_size == 1) {
76 graph = Tpetra::createCrsGraph(map, 1);
78 for (
size_t i=0; i<pce_sz; ++i) {
80 graph->insertGlobalIndices(row, arrayView(&row, 1));
87 size_t max_num_entry = 0;
88 for (
size_t i=0; i<pce_sz; ++i) {
89 const size_t num_entry =
cijk.num_entry(i);
90 max_num_entry = (num_entry > max_num_entry) ? num_entry : max_num_entry;
93 graph = Tpetra::createCrsGraph(map, max_num_entry);
95 for (
size_t i=0; i<pce_sz; ++i) {
97 const size_t num_entry =
cijk.num_entry(i);
98 const size_t entry_beg =
cijk.entry_begin(i);
99 const size_t entry_end = entry_beg + num_entry;
100 for (
size_t entry = entry_beg; entry < entry_end; ++entry) {
103 graph->insertGlobalIndices(row, arrayView(&
j, 1));
104 graph->insertGlobalIndices(row, arrayView(&k, 1));
108 graph->fillComplete();
119 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
120 create_flat_pce_graph(
122 const CijkType&
cijk,
123 Teuchos::RCP<
const Tpetra::Map<
LocalOrdinal,
GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_domain_map,
124 Teuchos::RCP<
const Tpetra::Map<
LocalOrdinal,
GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_range_map,
125 Teuchos::RCP<
const Tpetra::CrsGraph<
LocalOrdinal,
GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& cijk_graph,
126 const size_t matrix_pce_size) {
127 using Teuchos::ArrayView;
128 using Teuchos::ArrayRCP;
129 using Teuchos::Array;
133 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
134 typedef Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> Map;
135 typedef Tpetra::CrsGraph<LocalOrdinal,GlobalOrdinal,Node> Graph;
140 if (flat_domain_map == Teuchos::null)
141 flat_domain_map =
create_flat_map(*(graph.getDomainMap()), block_size);
144 if (flat_range_map == Teuchos::null)
148 RCP<const Map> flat_col_map =
154 RCP<const Map> flat_row_map;
155 if (graph.getRangeMap() == graph.getRowMap())
156 flat_row_map = flat_range_map;
161 if (cijk_graph == Teuchos::null)
162 cijk_graph = create_cijk_crs_graph<LocalOrdinal,GlobalOrdinal,Device>(
164 flat_domain_map->getComm(),
171 ArrayView<const LocalOrdinal> outer_cols;
172 ArrayView<const LocalOrdinal> inner_cols;
173 size_t max_num_row_entries = graph.getNodeMaxNumRowEntries()*block_size;
174 Array<LocalOrdinal> flat_col_indices;
175 flat_col_indices.reserve(max_num_row_entries);
176 RCP<Graph> flat_graph = rcp(
new Graph(flat_row_map, flat_col_map, max_num_row_entries));
177 const LocalOrdinal num_outer_rows = graph.getNodeNumRows();
178 for (
LocalOrdinal outer_row=0; outer_row < num_outer_rows; outer_row++) {
181 graph.getLocalRowView(outer_row, outer_cols);
185 for (
LocalOrdinal inner_row=0; inner_row < block_size; inner_row++) {
188 const LocalOrdinal flat_row = outer_row*block_size + inner_row;
191 cijk_graph->getLocalRowView(inner_row, inner_cols);
195 flat_col_indices.resize(0);
198 for (
LocalOrdinal outer_entry=0; outer_entry<num_outer_cols;
203 for (
LocalOrdinal inner_entry=0; inner_entry<num_inner_cols;
208 const LocalOrdinal flat_col = outer_col*block_size + inner_col;
209 flat_col_indices.push_back(flat_col);
215 flat_graph->insertLocalIndices(flat_row, flat_col_indices());
220 flat_graph->fillComplete(flat_domain_map, flat_range_map);
231 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
235 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
237 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
242 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
243 typedef Tpetra::MultiVector<BaseScalar,LocalOrdinal,GlobalOrdinal,Node> FlatVector;
244 typedef typename FlatVector::dual_view_type flat_view_type;
247 flat_view_type flat_vals (vec.getLocalViewDevice(), vec.getLocalViewHost());
248 if (vec.need_sync_device ()) {
249 flat_vals.modify_host ();
251 else if (vec.need_sync_host ()) {
252 flat_vals.modify_device ();
256 RCP<FlatVector> flat_vec = rcp(
new FlatVector(flat_map, flat_vals));
267 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
271 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
273 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
278 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
279 typedef Tpetra::MultiVector<BaseScalar,LocalOrdinal,GlobalOrdinal,Node> FlatVector;
280 typedef typename FlatVector::dual_view_type flat_view_type;
283 flat_view_type flat_vals (vec.getLocalViewDevice(), vec.getLocalViewHost());
284 if (vec.need_sync_device ()) {
285 flat_vals.modify_host ();
287 else if (vec.need_sync_host ()) {
288 flat_vals.modify_device ();
292 RCP<FlatVector> flat_vec = rcp(
new FlatVector(flat_map, flat_vals));
304 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
308 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
310 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
311 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
312 if (flat_map == Teuchos::null) {
317 const Teuchos::RCP< const Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> > const_flat_map = flat_map;
328 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
332 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
334 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
335 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
336 if (flat_map == Teuchos::null) {
341 const Teuchos::RCP< const Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> > const_flat_map = flat_map;
351 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
355 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec_const,
357 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
358 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
360 Teuchos::RCP< Tpetra::MultiVector<typename Storage::value_type,LocalOrdinal,GlobalOrdinal,Node> > flat_mv =
create_flat_vector_view(mv, flat_map);
361 return flat_mv->getVector(0);
371 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
375 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
377 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
378 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
379 if (flat_map == Teuchos::null) {
384 const Teuchos::RCP< const Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> > const_flat_map = flat_map;
394 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
398 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
400 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
401 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
403 Teuchos::RCP< Tpetra::MultiVector<typename Storage::value_type,LocalOrdinal,GlobalOrdinal,Node> > flat_mv =
create_flat_vector_view(mv, flat_map);
404 return flat_mv->getVectorNonConst(0);
414 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
418 Kokkos::Compat::KokkosDeviceWrapperNode<Device> >& vec,
420 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_map) {
421 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
422 if (flat_map == Teuchos::null) {
427 const Teuchos::RCP< const Tpetra::Map<LocalOrdinal,GlobalOrdinal,Node> > const_flat_map = flat_map;
434 typename Device,
typename CijkType>
437 Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >
441 const Teuchos::RCP<
const Tpetra::CrsGraph<
LocalOrdinal,
GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& flat_graph,
442 const Teuchos::RCP<
const Tpetra::CrsGraph<
LocalOrdinal,
GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<Device> > >& cijk_graph,
443 const CijkType&
cijk) {
444 using Teuchos::ArrayView;
445 using Teuchos::Array;
449 typedef Kokkos::Compat::KokkosDeviceWrapperNode<Device>
Node;
452 typedef Tpetra::CrsMatrix<BaseScalar,LocalOrdinal,GlobalOrdinal,Node> FlatMatrix;
459 RCP<FlatMatrix> flat_mat = rcp(
new FlatMatrix(flat_graph));
462 ArrayView<const Scalar> outer_values;
463 ArrayView<const LocalOrdinal> outer_cols;
464 ArrayView<const LocalOrdinal> inner_cols;
465 ArrayView<const LocalOrdinal> flat_cols;
466 Array<BaseScalar> flat_values;
467 flat_values.reserve(flat_graph->getNodeMaxNumRowEntries());
468 const LocalOrdinal num_outer_rows = mat.getNodeNumRows();
469 for (
LocalOrdinal outer_row=0; outer_row < num_outer_rows; outer_row++) {
472 mat.getLocalRowView(outer_row, outer_cols, outer_values);
476 for (
LocalOrdinal inner_row=0; inner_row < block_size; inner_row++) {
479 const LocalOrdinal flat_row = outer_row*block_size + inner_row;
482 cijk_graph->getLocalRowView(inner_row, inner_cols);
486 const LocalOrdinal num_flat_indices = num_outer_cols*num_inner_cols;
488 flat_values.assign(num_flat_indices, BaseScalar(0));
490 if (matrix_pce_size == 1) {
494 for (
LocalOrdinal outer_entry=0; outer_entry<num_outer_cols;
498 flat_values[outer_entry] = outer_values[outer_entry].coeff(0);
506 const size_t num_entry =
cijk.num_entry(inner_row);
507 const size_t entry_beg =
cijk.entry_begin(inner_row);
508 const size_t entry_end = entry_beg + num_entry;
509 for (
size_t entry = entry_beg; entry < entry_end; ++entry) {
512 const BaseScalar c =
cijk.value(entry);
515 typedef typename ArrayView<const LocalOrdinal>::iterator iterator;
517 std::find(inner_cols.begin(), inner_cols.end(),
j);
519 std::find(inner_cols.begin(), inner_cols.end(), k);
520 const LocalOrdinal j_offset = ptr_j - inner_cols.begin();
521 const LocalOrdinal k_offset = ptr_k - inner_cols.begin();
524 for (
LocalOrdinal outer_entry=0; outer_entry<num_outer_cols;
528 flat_values[outer_entry*num_inner_cols + j_offset] +=
529 c*outer_values[outer_entry].coeff(k);
530 flat_values[outer_entry*num_inner_cols + k_offset] +=
531 c*outer_values[outer_entry].coeff(
j);
540 flat_graph->getLocalRowView(flat_row, flat_cols);
541 flat_mat->replaceLocalValues(flat_row, flat_cols, flat_values());
546 flat_mat->fillComplete(flat_graph->getDomainMap(),
547 flat_graph->getRangeMap());
556 #endif // STOKHOS_TPETRA_UTILITIES_MP_VECTOR_HPP Stokhos::StandardStorage< int, double > Storage
Teuchos::RCP< const Tpetra::MultiVector< typename Storage::value_type, LocalOrdinal, GlobalOrdinal, Node > > create_flat_vector_view(const Tpetra::MultiVector< Sacado::MP::Vector< Storage >, LocalOrdinal, GlobalOrdinal, Node > &vec_const, const Teuchos::RCP< const Tpetra::Map< LocalOrdinal, GlobalOrdinal, Node > > &flat_map)
Top-level namespace for Stokhos classes and functions.
KokkosClassic::DefaultNode::DefaultNodeType Node
KOKKOS_INLINE_FUNCTION constexpr std::enable_if< is_view_uq_pce< View< T, P... > >::value, unsigned >::type dimension_scalar(const View< T, P... > &view)
Teuchos::RCP< Tpetra::CrsMatrix< typename Storage::value_type, LocalOrdinal, GlobalOrdinal, Node > > create_flat_matrix(const Tpetra::CrsMatrix< Sacado::MP::Vector< Storage >, LocalOrdinal, GlobalOrdinal, Node > &mat, const Teuchos::RCP< const Tpetra::CrsGraph< LocalOrdinal, GlobalOrdinal, Node > > &flat_graph, const LocalOrdinal block_size)
KOKKOS_INLINE_FUNCTION constexpr std::enable_if< is_view_uq_pce< view_type >::value, typename CijkType< view_type >::type >::type cijk(const view_type &view)
Teuchos::RCP< Tpetra::Map< LocalOrdinal, GlobalOrdinal, Node > > create_flat_map(const Tpetra::Map< LocalOrdinal, GlobalOrdinal, Node > &map, const LocalOrdinal block_size)