Skip to content

Commit

Permalink
Amesos2 : filter out "extra columns" during reindex & gather
Browse files Browse the repository at this point in the history 
Signed-off-by: iyamazaki <[email protected]>
  • Loading branch information
@iyamazaki
iyamazaki committed Feb 6, 2025
1 parent e595265 commit 0f7ef94
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 28 deletions.
8 changes: 2 additions & 6 deletions packages/amesos2/src/Amesos2_KLU2_def.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -463,16 +463,12 @@ KLU2<Matrix,Vector>::loadA_impl(EPhase current_phase)
bool column_major = true;
if (!is_contiguous_) {
auto contig_mat = this->matrixA_->reindex(this->contig_rowmap_, this->contig_colmap_, current_phase);
if (!contig_mat.is_null()) {
// reindex did not fail (e.g., not rectangular matrix)
nnz_ret = contig_mat->gather(host_nzvals_view_, host_rows_view_, host_col_ptr_view_, this->recvCounts, this->recvDispls, this->transpose_map, this->nzvals_t,
column_major, current_phase);
}
nnz_ret = contig_mat->gather(host_nzvals_view_, host_rows_view_, host_col_ptr_view_, this->recvCounts, this->recvDispls, this->transpose_map, this->nzvals_t,
column_major, current_phase);
} else {
nnz_ret = this->matrixA_->gather(host_nzvals_view_, host_rows_view_, host_col_ptr_view_, this->recvCounts, this->recvDispls, this->transpose_map, this->nzvals_t,
column_major, current_phase);
}
// reindex failed (e.g., rectangular matrix) or
// gather failed (e.g., not implemened for KokkosCrsMatrix)
// in case of the failure, it falls back to the original "do_get"
if (nnz_ret < 0) gather_supported = false;
Expand Down
8 changes: 2 additions & 6 deletions packages/amesos2/src/Amesos2_ShyLUBasker_def.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -597,16 +597,12 @@ ShyLUBasker<Matrix,Vector>::loadA_impl(EPhase current_phase)
bool column_major = true;
if (!is_contiguous_) {
auto contig_mat = this->matrixA_->reindex(this->contig_rowmap_, this->contig_colmap_, current_phase);
if (!contig_mat.is_null()) {
// reindex did not fail (e.g., not rectangular matrix)
nnz_ret = contig_mat->gather(nzvals_view_, rowind_view_, colptr_view_, this->recvCounts, this->recvDispls, this->transpose_map, this->nzvals_t,
column_major, current_phase);
}
nnz_ret = contig_mat->gather(nzvals_view_, rowind_view_, colptr_view_, this->recvCounts, this->recvDispls, this->transpose_map, this->nzvals_t,
column_major, current_phase);
} else {
nnz_ret = this->matrixA_->gather(nzvals_view_, rowind_view_, colptr_view_, this->recvCounts, this->recvDispls, this->transpose_map, this->nzvals_t,
column_major, current_phase);
}
// reindex failed (e.g., rectangular matrix) or
// gather failed (e.g., not implemened for KokkosCrsMatrix)
// in case of the failure, it falls back to the original "do_get"
if (nnz_ret < 0) gather_supported = false;
Expand Down
42 changes: 26 additions & 16 deletions packages/amesos2/src/Amesos2_TpetraCrsMatrix_MatrixAdapter_def.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -128,8 +128,10 @@ namespace Amesos2 {
}
// Create new GID list for ColMap
Kokkos::View<global_ordinal_t*, HostExecSpaceType> colIndexList ("colIndexList", nCols);
// initialize to catch col GIDs that are not in row GIDs
// they will be all assigned to (n+1)th columns
for (local_ordinal_t k = 0; k < nCols; k++) {
colIndexList(k) = indexBase-1; // initialize to catch col GIDs that are not in row GIDs
colIndexList(k) = numDoFs+indexBase;
}
typedef Tpetra::MultiVector<global_ordinal_t,
local_ordinal_t,
Expand All @@ -146,12 +148,6 @@ namespace Amesos2 {
// col_mv is imported from rowIndexList, which is based on index-base of rowMap
auto col_view = col_mv.getLocalViewHost(Tpetra::Access::ReadOnly);
for(int i=0; i<nCols; i++) {
if (col_view(i,0) < indexBase) {
// This indicates colMap has GIDs that are not in rowMap (rectangular matrix)
contigRowMap = RCP<map_t>();
contigColMap = RCP<map_t>();
return RCP<ConcreteMatrixAdapter<matrix_t>>();
}
colIndexList(i) = col_view(i,0);
}
}
Expand Down Expand Up @@ -324,9 +320,11 @@ namespace Amesos2 {
recvDispls(p+1) = recvDispls(p) + recvCounts(p);
}
}
// -- convert to global colids & convert to 0-base
// -- convert to global colids & ** convert to base-zero **
KV_GO lclColind_ ("localColind_", lclColind.extent(0));
for (int i = 0; i < int(lclColind.extent(0)); i++) lclColind_(i) = (colMap->getGlobalElement((lclColind(i))) - colIndexBase);
for (int i = 0; i < int(lclColind.extent(0)); i++) {
lclColind_(i) = (colMap->getGlobalElement((lclColind(i))) - colIndexBase);
}
if (column_major || need_to_perm) {
Kokkos::resize(indices_t, indices.extent(0));
Teuchos::gatherv<int, LocalOrdinal> (lclColind_.data(), lclColind_.extent(0), indices_t.data(),
Expand All @@ -343,6 +341,7 @@ namespace Amesos2 {
Teuchos::RCP< Teuchos::Time > gatherTime = Teuchos::TimeMonitor::getNewCounter ("Amesos2::gather(transpose index)");
Teuchos::TimeMonitor GatherTimer(*gatherTime);
#endif
// (note: column idexes are now in base-0)
if (column_major) {
// Map to transpose
Kokkos::resize(transpose_map, ret);
Expand All @@ -364,9 +363,14 @@ namespace Amesos2 {
int i = perm_l2g(row);
for (int k=pointers_t(i); k<pointers_t(i+1); k++) {
int col = indices_t(k);
transpose_map(k) = pointers(1+col);
indices(pointers(1+col)) = row;
pointers(1+col) ++;
if (col < nRows) {
transpose_map(k) = pointers(1+col);
indices(pointers(1+col)) = row;
pointers(1+col) ++;
} else {
// extra columns
transpose_map(k) = -1;
}
}
}
} else if (need_to_perm) {
Expand All @@ -383,9 +387,13 @@ namespace Amesos2 {
int row = perm_g2l(i);
for (int k=pointers_t(i); k<pointers_t(i+1); k++) {
int col = indices_t(k);
transpose_map(k) = pointers(1+row);
indices(pointers(1+row)) = col;
pointers(1+row) ++;
if (col < nRows) {
transpose_map(k) = pointers(1+row);
indices(pointers(1+row)) = col;
pointers(1+row) ++;
} else {
transpose_map(k) = -1;
}
}
}
} else {
Expand Down Expand Up @@ -426,7 +434,9 @@ namespace Amesos2 {
#endif
if (transpose_map.extent(0) > 0) {
for (int k=0; k<ret; k++) {
nzvals(transpose_map(k)) = nzvals_t(k);
if (transpose_map(k) >= 0) {
nzvals(transpose_map(k)) = nzvals_t(k);
}
}
}
}
Expand Down

0 comments on commit 0f7ef94

Please sign in to comment.