Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Muelu: use 4 map constructor for Kokkos-based TentativePFactory #2954

Merged
merged 2 commits into from
Jun 20, 2018
Merged

Muelu: use 4 map constructor for Kokkos-based TentativePFactory #2954

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a659d3b
Xpetra: add Kokkos-based 4 map constructor to CrsMatrix
aprokop Jun 15, 2018
27e8783
MueLu: kokkos: use 4 map constructor for TentativeP
aprokop Jun 15, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
287 changes: 145 additions & 142 deletions packages/muelu/src/Transfers/Smoothed-Aggregation/MueLu_TentativePFactory_kokkos_def.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -649,39 +649,32 @@ namespace MueLu {
// unnecessary).
const Kokkos::TeamPolicy<execution_space> policy(numAggregates, 1);

Kokkos::parallel_for("MueLu:TentativePF:BuildUncoupled:main_loop", policy,
KOKKOS_LAMBDA(const typename Kokkos::TeamPolicy<execution_space>::member_type &thread) {
auto agg = thread.league_rank();
if (doQRStep) {
Kokkos::parallel_for("MueLu:TentativePF:BuildUncoupled:main_loop", policy,
KOKKOS_LAMBDA(const typename Kokkos::TeamPolicy<execution_space>::member_type &thread) {
auto agg = thread.league_rank();

// size of the aggregate (number of DOFs in aggregate)
LO aggSize = aggRows(agg+1) - aggRows(agg);
// size of the aggregate (number of DOFs in aggregate)
LO aggSize = aggRows(agg+1) - aggRows(agg);

// Extract the piece of the nullspace corresponding to the aggregate, and
// put it in the flat array, "localQR" (in column major format) for the
// QR routine. Trivial in 1D.
if (goodMap) {
// Extract the piece of the nullspace corresponding to the aggregate, and
// put it in the flat array, "localQR" (in column major format) for the
// QR routine. Trivial in 1D.
auto norm = ATS::magnitude(zero);

if (doQRStep) {
// Calculate QR by hand
// FIXME: shouldn't there be stridedblock here?
// FIXME_KOKKOS: shouldn't there be stridedblock here?
for (decltype(aggSize) k = 0; k < aggSize; k++) {
auto dnorm = ATS::magnitude(fineNSRandom(agg2RowMapLO(aggRows(agg)+k),0));
norm += dnorm*dnorm;
}
norm = sqrt(norm);
// Calculate QR by hand
// FIXME: shouldn't there be stridedblock here?
// FIXME_KOKKOS: shouldn't there be stridedblock here?
for (decltype(aggSize) k = 0; k < aggSize; k++) {
auto dnorm = ATS::magnitude(fineNSRandom(agg2RowMapLO(aggRows(agg)+k),0));
norm += dnorm*dnorm;
}
norm = sqrt(norm);

if (norm == zero) {
// zero column; terminate the execution
statusAtomic(1) = true;
return;
}
} else {
// The easiest and less code churn way is to just declare the
// norm to one, that the column of tentative P is just scaled by
// identity
norm = ATS::magnitude(one);
if (norm == zero) {
// zero column; terminate the execution
statusAtomic(1) = true;
return;
}

// R = norm
Expand All @@ -697,16 +690,10 @@ namespace MueLu {
vals(localRow) = localVal;

}

} else {
// FIXME_KOKKOS: implement non-standard map QR
// Look at the original TentativeP for how to do that
statusAtomic(0) = true;
return;
}
});
});

typename status_type::HostMirror statusHost = Kokkos::create_mirror_view(status);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn't there a missing Kokkos::deep_copy (statusHost, status) after this? It might work only because of the UVM assumption, but I would like to purge that at some point.

Kokkos::deep_copy(statusHost, status);
for (decltype(statusHost.size()) i = 0; i < statusHost.size(); i++)
if (statusHost(i)) {
std::ostringstream oss;
Expand All @@ -718,102 +705,127 @@ namespace MueLu {
throw Exceptions::RuntimeError(oss.str());
}

} else { // NSdim > 1
// FIXME_KOKKOS: This code branch is completely unoptimized.
// Work to do:
// - Optimize QR decomposition
// - Remove INVALID usage similarly to CoalesceDropFactory_kokkos by
// packing new values in the beginning of each row
// We do use auxilary view in this case, so keep a second rows view for
// counting nonzeros in rows

// NOTE: the allocation (initialization) of these view takes noticeable time
size_t nnzEstimate = numRows * NSDim;
rows_type rowsAux("Ptent_aux_rows", numRows+1);
cols_type colsAux("Ptent_aux_cols", nnzEstimate);
vals_type valsAux("Ptent_aux_vals", nnzEstimate);
rows = rows_type("Ptent_rows", numRows+1);
{
// Stage 0: fill in views.
SubFactoryMonitor m2(*this, "Stage 0 (InitViews)", coarseLevel);

// The main thing to notice is initialization of vals with INVALID. These
// values will later be used to compress the arrays
Kokkos::parallel_for("MueLu:TentativePF:BuildPuncoupled:for1", range_type(0, numRows+1),
KOKKOS_LAMBDA(const LO row) {
rowsAux(row) = row*NSDim;
});
Kokkos::parallel_for("MueLu:TentativePF:BuildUncoupled:for2", range_type(0, nnzEstimate),
KOKKOS_LAMBDA(const LO j) {
colsAux(j) = INVALID;
valsAux(j) = zero;
});
}
} else {
Kokkos::parallel_for("MueLu:TentativePF:BuildUncoupled:main_loop_noqr", policy,
KOKKOS_LAMBDA(const typename Kokkos::TeamPolicy<execution_space>::member_type &thread) {
auto agg = thread.league_rank();

{
SubFactoryMonitor m2 = SubFactoryMonitor(*this, doQRStep ? "Stage 1 (LocalQR)" : "Stage 1 (Fill coarse nullspace and tentative P)", coarseLevel);
// Set up team policy with numAggregates teams and one thread per team.
// Each team handles a slice of the data associated with one aggregate
// and performs a local QR decomposition
const Kokkos::TeamPolicy<execution_space> policy(numAggregates,1); // numAggregates teams a 1 thread
LocalQRDecompFunctor<LocalOrdinal, GlobalOrdinal, Scalar, DeviceType, decltype(fineNSRandom),
decltype(aggDofSizes /*aggregate sizes in dofs*/), decltype(maxAggSize), decltype(agg2RowMapLO),
decltype(statusAtomic), decltype(rows), decltype(rowsAux), decltype(colsAux),
decltype(valsAux)>
localQRFunctor(fineNSRandom, coarseNS, aggDofSizes, maxAggSize, agg2RowMapLO, statusAtomic,
rows, rowsAux, colsAux, valsAux, doQRStep);
Kokkos::parallel_reduce("MueLu:TentativePF:BuildUncoupled:main_qr_loop", policy, localQRFunctor, nnz);
}
// size of the aggregate (number of DOFs in aggregate)
LO aggSize = aggRows(agg+1) - aggRows(agg);

// R = norm
coarseNS(agg, 0) = one;

// Q = localQR(:,0)/norm
for (decltype(aggSize) k = 0; k < aggSize; k++) {
LO localRow = agg2RowMapLO(aggRows(agg)+k);
SC localVal = fineNSRandom(agg2RowMapLO(aggRows(agg)+k),0);

rows(localRow+1) = localRow+1;
cols(localRow) = agg;
vals(localRow) = localVal;

typename status_type::HostMirror statusHost = Kokkos::create_mirror_view(status);
for (decltype(statusHost.size()) i = 0; i < statusHost.size(); i++)
if (statusHost(i)) {
std::ostringstream oss;
oss << "MueLu::TentativePFactory::MakeTentative: ";
switch(i) {
case 0: oss << "!goodMap is not implemented"; break;
case 1: oss << "fine level NS part has a zero column"; break;
}
throw Exceptions::RuntimeError(oss.str());
}
});
}

// Compress the cols and vals by ignoring INVALID column entries that correspond
// to 0 in QR.

// The real cols and vals are constructed using calculated (not estimated) nnz
cols = decltype(cols)("Ptent_cols", nnz);
vals = decltype(vals)("Ptent_vals", nnz);
{
// Stage 2: compress the arrays
SubFactoryMonitor m2(*this, "Stage 2 (CompressRows)", coarseLevel);

Kokkos::parallel_scan("MueLu:TentativePF:Build:compress_rows", range_type(0,numRows+1),
KOKKOS_LAMBDA(const LO i, LO& upd, const bool& final) {
upd += rows(i);
if (final)
rows(i) = upd;
});
}
} else { // NSdim > 1
// FIXME_KOKKOS: This code branch is completely unoptimized.
// Work to do:
// - Optimize QR decomposition
// - Remove INVALID usage similarly to CoalesceDropFactory_kokkos by
// packing new values in the beginning of each row
// We do use auxilary view in this case, so keep a second rows view for
// counting nonzeros in rows

// NOTE: the allocation (initialization) of these view takes noticeable time
size_t nnzEstimate = numRows * NSDim;
rows_type rowsAux("Ptent_aux_rows", numRows+1);
cols_type colsAux("Ptent_aux_cols", nnzEstimate);
vals_type valsAux("Ptent_aux_vals", nnzEstimate);
rows = rows_type("Ptent_rows", numRows+1);
{
// Stage 0: fill in views.
SubFactoryMonitor m2(*this, "Stage 0 (InitViews)", coarseLevel);

// The main thing to notice is initialization of vals with INVALID. These
// values will later be used to compress the arrays
Kokkos::parallel_for("MueLu:TentativePF:BuildPuncoupled:for1", range_type(0, numRows+1),
KOKKOS_LAMBDA(const LO row) {
rowsAux(row) = row*NSDim;
});
Kokkos::parallel_for("MueLu:TentativePF:BuildUncoupled:for2", range_type(0, nnzEstimate),
KOKKOS_LAMBDA(const LO j) {
colsAux(j) = INVALID;
valsAux(j) = zero;
});
}

{
SubFactoryMonitor m2(*this, "Stage 2 (CompressCols)", coarseLevel);

// FIXME_KOKKOS: this can be spedup by moving correct cols and vals values
// to the beginning of rows. See CoalesceDropFactory_kokkos for
// example.
Kokkos::parallel_for("MueLu:TentativePF:Build:compress_cols_vals", range_type(0,numRows),
KOKKOS_LAMBDA(const LO i) {
LO rowStart = rows(i);

size_t lnnz = 0;
for (auto j = rowsAux(i); j < rowsAux(i+1); j++)
if (colsAux(j) != INVALID) {
cols(rowStart+lnnz) = colsAux(j);
vals(rowStart+lnnz) = valsAux(j);
lnnz++;
}
});
{
SubFactoryMonitor m2 = SubFactoryMonitor(*this, doQRStep ? "Stage 1 (LocalQR)" : "Stage 1 (Fill coarse nullspace and tentative P)", coarseLevel);
// Set up team policy with numAggregates teams and one thread per team.
// Each team handles a slice of the data associated with one aggregate
// and performs a local QR decomposition
const Kokkos::TeamPolicy<execution_space> policy(numAggregates,1); // numAggregates teams a 1 thread
LocalQRDecompFunctor<LocalOrdinal, GlobalOrdinal, Scalar, DeviceType, decltype(fineNSRandom),
decltype(aggDofSizes /*aggregate sizes in dofs*/), decltype(maxAggSize), decltype(agg2RowMapLO),
decltype(statusAtomic), decltype(rows), decltype(rowsAux), decltype(colsAux),
decltype(valsAux)>
localQRFunctor(fineNSRandom, coarseNS, aggDofSizes, maxAggSize, agg2RowMapLO, statusAtomic,
rows, rowsAux, colsAux, valsAux, doQRStep);
Kokkos::parallel_reduce("MueLu:TentativePF:BuildUncoupled:main_qr_loop", policy, localQRFunctor, nnz);
}

typename status_type::HostMirror statusHost = Kokkos::create_mirror_view(status);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See above comment.

Kokkos::deep_copy(statusHost, status);
for (decltype(statusHost.size()) i = 0; i < statusHost.size(); i++)
if (statusHost(i)) {
std::ostringstream oss;
oss << "MueLu::TentativePFactory::MakeTentative: ";
switch(i) {
case 0: oss << "!goodMap is not implemented"; break;
case 1: oss << "fine level NS part has a zero column"; break;
}
throw Exceptions::RuntimeError(oss.str());
}

// Compress the cols and vals by ignoring INVALID column entries that correspond
// to 0 in QR.

// The real cols and vals are constructed using calculated (not estimated) nnz
cols = decltype(cols)("Ptent_cols", nnz);
vals = decltype(vals)("Ptent_vals", nnz);
{
// Stage 2: compress the arrays
SubFactoryMonitor m2(*this, "Stage 2 (CompressRows)", coarseLevel);

Kokkos::parallel_scan("MueLu:TentativePF:Build:compress_rows", range_type(0,numRows+1),
KOKKOS_LAMBDA(const LO i, LO& upd, const bool& final) {
upd += rows(i);
if (final)
rows(i) = upd;
});
}

{
SubFactoryMonitor m2(*this, "Stage 2 (CompressCols)", coarseLevel);

// FIXME_KOKKOS: this can be spedup by moving correct cols and vals values
// to the beginning of rows. See CoalesceDropFactory_kokkos for
// example.
Kokkos::parallel_for("MueLu:TentativePF:Build:compress_cols_vals", range_type(0,numRows),
KOKKOS_LAMBDA(const LO i) {
LO rowStart = rows(i);

size_t lnnz = 0;
for (auto j = rowsAux(i); j < rowsAux(i+1); j++)
if (colsAux(j) != INVALID) {
cols(rowStart+lnnz) = colsAux(j);
vals(rowStart+lnnz) = valsAux(j);
lnnz++;
}
});
}
}

GetOStream(Runtime1) << "TentativePFactory : aggregates do not cross process boundaries" << std::endl;
Expand All @@ -823,28 +835,19 @@ namespace MueLu {
SubFactoryMonitor m2(*this, "Stage 3 (LocalMatrix+FillComplete)", coarseLevel);

local_matrix_type lclMatrix = local_matrix_type("A", numRows, coarseMap->getNodeNumElements(), nnz, vals, rows, cols);
#if 1
// FIXME_KOKKOS: this should be gone once Xpetra propagate "local matrix + 4 maps" constructor
auto PtentCrs = CrsMatrixFactory::Build(rowMap, coarseMap, lclMatrix);
PtentCrs->resumeFill(); // we need that for rectangular matrices

// Managing labels & constants for ESFC
RCP<ParameterList> FCparams;
if(pL.isSublist("matrixmatrix: kernel params"))
FCparams=rcp(new ParameterList(pL.sublist("matrixmatrix: kernel params")));
if (pL.isSublist("matrixmatrix: kernel params"))
FCparams = rcp(new ParameterList(pL.sublist("matrixmatrix: kernel params")));
else
FCparams= rcp(new ParameterList);
FCparams = rcp(new ParameterList);

// By default, we don't need global constants for TentativeP
FCparams->set("compute global constants",FCparams->get("compute global constants",false));
std::string levelIDs = toString(levelID);
FCparams->set("Timer Label",std::string("MueLu::TentativeP-")+levelIDs);
RCP<const Export> dummy_e;
RCP<const Import> dummy_i;

PtentCrs->expertStaticFillComplete(coarseMap, A->getDomainMap(), dummy_i, dummy_e, FCparams);
#else
FCparams->set("compute global constants", FCparams->get("compute global constants", false));
FCparams->set("Timer Label", std::string("MueLu::TentativeP-") + toString(levelID));

auto PtentCrs = CrsMatrixFactory::Build(lclMatrix, rowMap, coarseMap, coarseMap, A->getDomainMap());
#endif
Ptentative = rcp(new CrsMatrixWrap(PtentCrs));
}
}
Expand Down
Loading