14 #ifndef _ZOLTAN2_TPETRAROWMATRIXADAPTER_HPP_
15 #define _ZOLTAN2_TPETRAROWMATRIXADAPTER_HPP_
21 #include <Tpetra_RowMatrix.hpp>
41 template <
typename User,
typename UserCoord = User>
46 #ifndef DOXYGEN_SHOULD_SKIP_THIS
53 using device_t =
typename node_t::device_type;
54 using host_t =
typename Kokkos::HostSpace::memory_space;
56 using userCoord_t = UserCoord;
67 int nWeightsPerRow = 0);
103 void setWeightsHost(
typename Base::ConstWeightsHostView1D val,
int idx);
182 typename Base::ConstIdsHostView &rowIds)
const override;
185 typename Base::ConstIdsDeviceView &rowIds)
const override;
187 void getCRSView(ArrayRCP<const offset_t> &offsets,
188 ArrayRCP<const gno_t> &colIds)
const;
191 typename Base::ConstOffsetsHostView &offsets,
192 typename Base::ConstIdsHostView &colIds)
const override;
195 typename Base::ConstOffsetsDeviceView &offsets,
196 typename Base::ConstIdsDeviceView &colIds)
const override;
198 void getCRSView(ArrayRCP<const offset_t> &offsets,
199 ArrayRCP<const gno_t> &colIds,
200 ArrayRCP<const scalar_t> &values)
const;
203 typename Base::ConstOffsetsHostView &offsets,
204 typename Base::ConstIdsHostView &colIds,
205 typename Base::ConstScalarsHostView &values)
const override;
208 typename Base::ConstOffsetsDeviceView &offsets,
209 typename Base::ConstIdsDeviceView &colIds,
210 typename Base::ConstScalarsDeviceView &values)
const override;
221 typename Base::WeightsDeviceView &weights)
const override;
227 typename Base::WeightsHostView &weights)
const override;
231 template <
typename Adapter>
233 const User &in, User *&out,
236 template <
typename Adapter>
238 const User &in, RCP<User> &out,
244 const RCP<const User> &inmatrix)
268 virtual RCP<User>
doMigration(
const User &from,
size_t numLocalRows,
269 const gno_t *myNewRows)
const;
276 template <
typename User,
typename UserCoord>
278 const RCP<const User> &inmatrix,
int nWeightsPerRow):
279 matrix_(inmatrix), offset_(), columnIds_(),
280 nWeightsPerRow_(nWeightsPerRow), rowWeights_(),
281 mayHaveDiagonalEntries(true) {
283 using localInds_t =
typename User::nonconst_local_inds_host_view_type;
284 using localVals_t =
typename User::nonconst_values_host_view_type;
286 const auto nrows =
matrix_->getLocalNumRows();
287 const auto nnz =
matrix_->getLocalNumEntries();
288 auto maxNumEntries =
matrix_->getLocalMaxNumRowEntries();
293 colIdsHost_ =
typename Base::ConstIdsHostView(
"colIdsHost_", nnz);
294 offsHost_ =
typename Base::ConstOffsetsHostView(
"offsHost_", nrows + 1);
295 valuesHost_ =
typename Base::ScalarsHostView(
"valuesHost_", nnz);
297 localInds_t localColInds(
"localColInds", maxNumEntries);
298 localVals_t localVals(
"localVals", maxNumEntries);
300 for (
size_t r = 0; r < nrows; r++) {
301 size_t numEntries = 0;
302 matrix_->getLocalRowCopy(r, localColInds, localVals, numEntries);
308 for (
size_t j = 0; j < nnz; j++) {
324 "rowWeightsDevice_", nrows, nWeightsPerRow_);
336 template <
typename User,
typename UserCoord>
338 const scalar_t *weightVal,
int stride,
int idx) {
339 if (this->getPrimaryEntityType() ==
MATRIX_ROW)
340 setRowWeights(weightVal, stride, idx);
343 std::ostringstream emsg;
344 emsg << __FILE__ <<
"," << __LINE__
345 <<
" error: setWeights not yet supported for"
346 <<
" columns or nonzeros." << std::endl;
347 throw std::runtime_error(emsg.str());
352 template <
typename User,
typename UserCoord>
354 typename Base::ConstWeightsDeviceView1D val,
int idx) {
355 if (this->getPrimaryEntityType() ==
MATRIX_ROW)
356 setRowWeightsDevice(val, idx);
359 std::ostringstream emsg;
360 emsg << __FILE__ <<
"," << __LINE__
361 <<
" error: setWeights not yet supported for"
362 <<
" columns or nonzeros." << std::endl;
363 throw std::runtime_error(emsg.str());
368 template <
typename User,
typename UserCoord>
370 typename Base::ConstWeightsHostView1D val,
int idx) {
371 if (this->getPrimaryEntityType() ==
MATRIX_ROW)
372 setRowWeightsHost(val, idx);
375 std::ostringstream emsg;
376 emsg << __FILE__ <<
"," << __LINE__
377 <<
" error: setWeights not yet supported for"
378 <<
" columns or nonzeros." << std::endl;
379 throw std::runtime_error(emsg.str());
384 template <
typename User,
typename UserCoord>
386 const scalar_t *weightVal,
int stride,
int idx) {
389 "Invalid row weight index: " + std::to_string(idx));
391 size_t nrows = getLocalNumRows();
392 ArrayRCP<const scalar_t> weightV(weightVal, 0, nrows * stride,
false);
393 rowWeights_[idx] = input_t(weightV, stride);
397 template <
typename User,
typename UserCoord>
399 typename Base::ConstWeightsDeviceView1D
weights,
int idx) {
402 "Invalid row weight index: " + std::to_string(idx));
404 Kokkos::parallel_for(
405 rowWeightsDevice_.extent(0), KOKKOS_CLASS_LAMBDA(
const int rowID) {
406 rowWeightsDevice_(rowID, idx) =
weights(rowID);
413 template <
typename User,
typename UserCoord>
415 typename Base::ConstWeightsHostView1D weightsHost,
int idx) {
417 "Invalid row weight index: " + std::to_string(idx));
419 auto weightsDevice = Kokkos::create_mirror_view_and_copy(
422 setRowWeightsDevice(weightsDevice, idx);
426 template <
typename User,
typename UserCoord>
428 if (this->getPrimaryEntityType() ==
MATRIX_ROW)
429 setRowWeightIsNumberOfNonZeros(idx);
432 std::ostringstream emsg;
433 emsg << __FILE__ <<
"," << __LINE__
434 <<
" error: setWeightIsNumberOfNonZeros not yet supported for"
435 <<
" columns" << std::endl;
436 throw std::runtime_error(emsg.str());
441 template <
typename User,
typename UserCoord>
444 if (idx < 0 || idx >= nWeightsPerRow_) {
445 std::ostringstream emsg;
446 emsg << __FILE__ <<
":" << __LINE__ <<
" Invalid row weight index " << idx
448 throw std::runtime_error(emsg.str());
451 numNzWeight_(idx) =
true;
455 template <
typename User,
typename UserCoord>
457 return matrix_->getLocalNumRows();
461 template <
typename User,
typename UserCoord>
463 return matrix_->getLocalNumCols();
467 template <
typename User,
typename UserCoord>
469 return matrix_->getLocalNumEntries();
473 template <
typename User,
typename UserCoord>
477 template <
typename User,
typename UserCoord>
479 ArrayView<const gno_t> rowView = matrix_->getRowMap()->getLocalElementList();
480 rowIds = rowView.getRawPtr();
484 template <
typename User,
typename UserCoord>
486 typename Base::ConstIdsHostView &rowIds)
const {
487 auto idsDevice = matrix_->getRowMap()->getMyGlobalIndices();
488 auto tmpIds =
typename Base::IdsHostView(
"", idsDevice.extent(0));
490 Kokkos::deep_copy(tmpIds, idsDevice);
496 template <
typename User,
typename UserCoord>
498 typename Base::ConstIdsDeviceView &rowIds)
const {
500 auto idsDevice = matrix_->getRowMap()->getMyGlobalIndices();
501 auto tmpIds =
typename Base::IdsDeviceView(
"", idsDevice.extent(0));
503 Kokkos::deep_copy(tmpIds, idsDevice);
509 template <
typename User,
typename UserCoord>
511 ArrayRCP<const gno_t> &colIds)
const {
517 template <
typename User,
typename UserCoord>
519 typename Base::ConstOffsetsHostView &offsets,
520 typename Base::ConstIdsHostView &colIds)
const {
521 auto hostOffsets = Kokkos::create_mirror_view(offsDevice_);
522 Kokkos::deep_copy(hostOffsets, offsDevice_);
523 offsets = hostOffsets;
525 auto hostColIds = Kokkos::create_mirror_view(colIdsDevice_);
526 Kokkos::deep_copy(hostColIds, colIdsDevice_);
531 template <
typename User,
typename UserCoord>
533 typename Base::ConstOffsetsDeviceView &offsets,
534 typename Base::ConstIdsDeviceView &colIds)
const {
535 offsets = offsDevice_;
536 colIds = colIdsDevice_;
540 template <
typename User,
typename UserCoord>
542 ArrayRCP<const gno_t> &colIds,
543 ArrayRCP<const scalar_t> &values)
const {
550 template <
typename User,
typename UserCoord>
552 typename Base::ConstOffsetsHostView &offsets,
553 typename Base::ConstIdsHostView &colIds,
554 typename Base::ConstScalarsHostView &values)
const {
555 auto hostOffsets = Kokkos::create_mirror_view(offsDevice_);
556 Kokkos::deep_copy(hostOffsets, offsDevice_);
557 offsets = hostOffsets;
559 auto hostColIds = Kokkos::create_mirror_view(colIdsDevice_);
560 Kokkos::deep_copy(hostColIds, colIdsDevice_);
563 auto hostValues = Kokkos::create_mirror_view(valuesDevice_);
564 Kokkos::deep_copy(hostValues, valuesDevice_);
569 template <
typename User,
typename UserCoord>
571 typename Base::ConstOffsetsDeviceView &offsets,
572 typename Base::ConstIdsDeviceView &colIds,
573 typename Base::ConstScalarsDeviceView &values)
const {
574 offsets = offsDevice_;
575 colIds = colIdsDevice_;
576 values = valuesDevice_;
580 template <
typename User,
typename UserCoord>
584 template <
typename User,
typename UserCoord>
587 if (idx < 0 || idx >= nWeightsPerRow_) {
588 std::ostringstream emsg;
589 emsg << __FILE__ <<
":" << __LINE__ <<
" Invalid row weight index "
591 throw std::runtime_error(emsg.str());
595 rowWeights_[idx].getStridedList(length, weights, stride);
599 template <
typename User,
typename UserCoord>
601 typename Base::WeightsDeviceView1D &
weights,
int idx)
const {
603 "Invalid row weight index.");
605 const auto size = rowWeightsDevice_.extent(0);
606 weights =
typename Base::WeightsDeviceView1D(
"weights", size);
608 Kokkos::parallel_for(
609 size, KOKKOS_CLASS_LAMBDA(
const int id) {
610 weights(
id) = rowWeightsDevice_(
id, idx);
617 template <
typename User,
typename UserCoord>
619 typename Base::WeightsDeviceView &
weights)
const {
621 weights = rowWeightsDevice_;
625 template <
typename User,
typename UserCoord>
627 typename Base::WeightsHostView1D &
weights,
int idx)
const {
629 "Invalid row weight index.");
631 auto weightsDevice =
typename Base::WeightsDeviceView1D(
632 "weights", rowWeightsDevice_.extent(0));
633 getRowWeightsDeviceView(weightsDevice, idx);
635 weights = Kokkos::create_mirror_view(weightsDevice);
636 Kokkos::deep_copy(weights, weightsDevice);
640 template <
typename User,
typename UserCoord>
642 typename Base::WeightsHostView &
weights)
const {
644 weights = Kokkos::create_mirror_view(rowWeightsDevice_);
645 Kokkos::deep_copy(weights, rowWeightsDevice_);
649 template <
typename User,
typename UserCoord>
653 template <
typename User,
typename UserCoord>
654 template <
typename Adapter>
656 const User &in, User *&out,
660 ArrayRCP<gno_t> importList;
663 Zoltan2::getImportList<Adapter, TpetraRowMatrixAdapter<User, UserCoord>>(
664 solution,
this, importList);
669 RCP<User> outPtr = doMigration(in, numNewRows, importList.getRawPtr());
675 template <
typename User,
typename UserCoord>
676 template <
typename Adapter>
678 const User &in, RCP<User> &out,
682 ArrayRCP<gno_t> importList;
685 Zoltan2::getImportList<Adapter, TpetraRowMatrixAdapter<User, UserCoord>>(
686 solution,
this, importList);
691 out = doMigration(in, numNewRows, importList.getRawPtr());
695 template <
typename User,
typename UserCoord>
697 const User &from,
size_t numLocalRows,
const gno_t *myNewRows)
const {
698 typedef Tpetra::Map<lno_t, gno_t, node_t>
map_t;
699 typedef Tpetra::CrsMatrix<scalar_t, lno_t, gno_t, node_t> tcrsmatrix_t;
710 const tcrsmatrix_t *pCrsMatrix =
dynamic_cast<const tcrsmatrix_t *
>(&from);
713 throw std::logic_error(
"TpetraRowMatrixAdapter cannot migrate data for "
714 "your RowMatrix; it can migrate data only for "
715 "Tpetra::CrsMatrix. "
716 "You can inherit from TpetraRowMatrixAdapter and "
717 "implement migration for your RowMatrix.");
721 const RCP<const map_t> &smap = from.getRowMap();
722 gno_t numGlobalRows = smap->getGlobalNumElements();
723 gno_t base = smap->getMinAllGlobalIndex();
726 ArrayView<const gno_t> rowList(myNewRows, numLocalRows);
727 const RCP<const Teuchos::Comm<int>> &comm = from.getComm();
728 RCP<const map_t> tmap = rcp(
new map_t(numGlobalRows, rowList, base, comm));
731 Tpetra::Import<lno_t, gno_t, node_t> importer(smap, tmap);
733 int oldNumElts = smap->getLocalNumElements();
734 int newNumElts = numLocalRows;
737 typedef Tpetra::Vector<scalar_t, lno_t, gno_t, node_t> vector_t;
738 vector_t numOld(smap);
739 vector_t numNew(tmap);
740 for (
int lid = 0; lid < oldNumElts; lid++) {
741 numOld.replaceGlobalValue(smap->getGlobalElement(lid),
742 scalar_t(from.getNumEntriesInLocalRow(lid)));
744 numNew.doImport(numOld, importer, Tpetra::INSERT);
747 ArrayRCP<size_t> nnz(newNumElts);
748 if (newNumElts > 0) {
749 ArrayRCP<scalar_t> ptr = numNew.getDataNonConst(0);
750 for (
int lid = 0; lid < newNumElts; lid++) {
751 nnz[lid] =
static_cast<size_t>(ptr[lid]);
755 RCP<tcrsmatrix_t> M = rcp(
new tcrsmatrix_t(tmap, nnz()));
757 M->doImport(from, importer, Tpetra::INSERT);
760 return Teuchos::rcp_dynamic_cast<User>(M);
765 #endif // _ZOLTAN2_TPETRAROWMATRIXADAPTER_HPP_
bool CRSViewAvailable() const
Indicates whether the MatrixAdapter implements a view of the matrix in compressed sparse row (CRS) fo...
void setWeightsHost(typename Base::ConstWeightsHostView1D val, int idx)
Provide a host view of weights for the primary entity type.
ArrayRCP< StridedData< lno_t, scalar_t > > rowWeights_
bool mayHaveDiagonalEntries
void getRowWeightsHostView(typename Base::WeightsHostView1D &weights, int idx=0) const
virtual RCP< User > doMigration(const User &from, size_t numLocalRows, const gno_t *myNewRows) const
Base::ScalarsHostView valuesHost_
void setRowWeightsHost(typename Base::ConstWeightsHostView1D val, int idx)
Provide a host view to row weights.
Helper functions for Partitioning Problems.
Base::ConstIdsDeviceView colIdsDevice_
#define Z2_FORWARD_EXCEPTIONS
Forward an exception back through call stack.
typename InputTraits< User >::scalar_t scalar_t
MatrixAdapter defines the adapter interface for matrices.
static void AssertCondition(bool condition, const std::string &message, const char *file=__FILE__, int line=__LINE__)
typename Kokkos::HostSpace::memory_space host_t
map_t::global_ordinal_type gno_t
void setWeightsDevice(typename Base::ConstWeightsDeviceView1D val, int idx)
Provide a device view of weights for the primary entity type.
typename node_t::device_type device_t
ArrayRCP< scalar_t > values_
Base::ConstIdsHostView colIdsHost_
typename InputTraits< User >::part_t part_t
bool useNumNonzerosAsRowWeight(int idx) const
Indicate whether row weight with index idx should be the global number of nonzeros in the row...
Provides access for Zoltan2 to Tpetra::RowMatrix data.
typename InputTraits< User >::node_t node_t
void getRowWeightsView(const scalar_t *&weights, int &stride, int idx=0) const
Provide a pointer to the row weights, if any.
size_t getLocalNumEntries() const
Returns the number of nonzeros on this process.
A PartitioningSolution is a solution to a partitioning problem.
void getRowIDsView(const gno_t *&rowIds) const override
void setRowWeightsDevice(typename Base::ConstWeightsDeviceView1D val, int idx)
Provide a device view to row weights.
int getNumWeightsPerRow() const
Returns the number of weights per row (0 or greater). Row weights may be used when partitioning matri...
void getCRSDeviceView(typename Base::ConstOffsetsDeviceView &offsets, typename Base::ConstIdsDeviceView &colIds) const override
Base::ScalarsDeviceView valuesDevice_
size_t getLocalNumColumns() const
Returns the number of columns on this process.
Base::ConstOffsetsDeviceView offsDevice_
The StridedData class manages lists of weights or coordinates.
map_t::local_ordinal_type lno_t
void getRowIDsDeviceView(typename Base::ConstIdsDeviceView &rowIds) const override
ArrayRCP< gno_t > columnIds_
void applyPartitioningSolution(const User &in, User *&out, const PartitioningSolution< Adapter > &solution) const
void setRowWeights(const scalar_t *weightVal, int stride, int idx=0)
Specify a weight for each row.
RCP< const User > matrix_
typename InputTraits< User >::offset_t offset_t
void getCRSHostView(typename Base::ConstOffsetsHostView &offsets, typename Base::ConstIdsHostView &colIds) const override
void getRowWeightsDeviceView(typename Base::WeightsDeviceView1D &weights, int idx=0) const
void setWeightIsDegree(int idx)
Specify an index for which the weight should be the degree of the entity.
Base::ConstOffsetsHostView offsHost_
Defines the MatrixAdapter interface.
void setWeights(const scalar_t *weightVal, int stride, int idx=0)
Specify a weight for each entity of the primaryEntityType.
ArrayRCP< offset_t > offset_
Kokkos::View< bool *, host_t > numNzWeight_
void getRowIDsHostView(typename Base::ConstIdsHostView &rowIds) const override
size_t getLocalNumRows() const
Returns the number of rows on this process.
TpetraRowMatrixAdapter(int nWeightsPerRow, const RCP< const User > &inmatrix)
Base::WeightsDeviceView rowWeightsDevice_
void getCRSView(ArrayRCP< const offset_t > &offsets, ArrayRCP< const gno_t > &colIds) const
void setRowWeightIsNumberOfNonZeros(int idx)
Specify an index for which the row weight should be the global number of nonzeros in the row...
TpetraRowMatrixAdapter(const RCP< const User > &inmatrix, int nWeightsPerRow=0)
Constructor.
This file defines the StridedData class.