Skip to content

Commit

Permalink
Merge Update defaults for distributed multigrid
Browse files Browse the repository at this point in the history 
This PR updates the defaults for distributed multigrid for the smoother and the coarse solver. It also adds a test for the distributed pgm preconditioned solver.

The multigrid for distributed matrices uses a global scalar Jacobi smoother and a preconditioned GMRES solver as coarse solver.

Related PR: #1612
  • Loading branch information
@MarcelKoch
MarcelKoch authored May 28, 2024
2 parents 591ace6 + 1fc3d7d commit 6abd3f3
Show file tree
Hide file tree
Showing 5 changed files with 146 additions and 7 deletions.
2 changes: 1 addition & 1 deletion common/cuda_hip/matrix/csr_kernels.hpp.inc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1298,7 +1298,7 @@ void spgeam(std::shared_ptr<const DefaultExecutor> exec,
{
auto total_nnz =
a->get_num_stored_elements() + b->get_num_stored_elements();
auto nnz_per_row = total_nnz / a->get_size()[0];
auto nnz_per_row = a->get_size()[0] ? total_nnz / a->get_size()[0] : 0;
select_spgeam(
spgeam_kernels(),
[&](int compiled_subwarp_size) {
Expand Down
6 changes: 6 additions & 0 deletions core/multigrid/pgm.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,12 @@ std::shared_ptr<matrix::Csr<ValueType, IndexType>> generate_coarse(
gko::array<IndexType> col_idxs(exec, nnz);
gko::array<ValueType> vals(exec, nnz);
exec->copy_from(exec, nnz, fine_csr->get_const_values(), vals.get_data());

if (nnz == 0) {
return matrix::Csr<ValueType, IndexType>::create(
exec, dim<2>(num_agg, non_local_num_agg));
}

// map row_ptrs to coarse row index
exec->run(pgm::make_map_row(num, fine_csr->get_const_row_ptrs(),
agg.get_const_data(), row_idxs.get_data()));
Expand Down
66 changes: 62 additions & 4 deletions core/solver/multigrid.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#include <ginkgo/core/base/utils.hpp>
#include <ginkgo/core/base/utils_helper.hpp>
#include <ginkgo/core/distributed/matrix.hpp>
#include <ginkgo/core/distributed/preconditioner/schwarz.hpp>
#include <ginkgo/core/distributed/vector.hpp>
#include <ginkgo/core/factorization/lu.hpp>
#include <ginkgo/core/matrix/dense.hpp>
Expand Down Expand Up @@ -99,6 +100,28 @@ void handle_list(
auto list_size = smoother_list.size();
auto gen_default_smoother = [&] {
auto exec = matrix->get_executor();
#if GINKGO_BUILD_MPI
if (gko::detail::is_distributed(matrix.get())) {
using experimental::distributed::Matrix;
return run<Matrix<ValueType, int32, int32>,
Matrix<ValueType, int32, int64>,
Matrix<ValueType, int64, int64>>(
matrix, [exec, iteration, relaxation_factor](auto matrix) {
using Mtx = typename decltype(matrix)::element_type;
return share(
build_smoother(
experimental::distributed::preconditioner::Schwarz<
ValueType, typename Mtx::local_index_type,
typename Mtx::global_index_type>::build()
.with_local_solver(
preconditioner::Jacobi<ValueType>::build()
.with_max_block_size(1u))
.on(exec),
iteration, casting<ValueType>(relaxation_factor))
->generate(matrix));
});
}
#endif
return share(build_smoother(preconditioner::Jacobi<ValueType>::build()
.with_max_block_size(1u)
.on(exec),
Expand Down Expand Up @@ -180,7 +203,7 @@ namespace detail {
* @note it should only be used internally
*/
struct MultigridState {
MultigridState() : nrhs{0} {}
MultigridState() : nrhs{static_cast<size_type>(-1)} {}

/**
* Generate the cache for later usage.
Expand Down Expand Up @@ -638,7 +661,42 @@ void Multigrid::generate()
// default coarse grid solver, direct LU
// TODO: maybe remove fixed index type
auto gen_default_solver = [&]() -> std::unique_ptr<LinOp> {
// TODO: unify when dpcpp supports direct solver
// TODO: unify when dpcpp supports direct solver
#if GINKGO_BUILD_MPI
if (gko::detail::is_distributed(matrix.get())) {
using absolute_value_type = remove_complex<value_type>;
using experimental::distributed::Matrix;
return run<Matrix<value_type, int32, int32>,
Matrix<value_type, int32, int64>,
Matrix<value_type, int64,
int64>>(matrix, [exec](auto matrix) {
using Mtx = typename decltype(matrix)::element_type;
return solver::Gmres<value_type>::build()
.with_criteria(
stop::Iteration::build().with_max_iters(
matrix->get_size()[0]),
stop::ResidualNorm<value_type>::build()
.with_reduction_factor(
std::numeric_limits<
absolute_value_type>::epsilon() *
absolute_value_type{10}))
.with_krylov_dim(
std::min(size_type(100), matrix->get_size()[0]))
.with_preconditioner(
experimental::distributed::preconditioner::
Schwarz<value_type,
typename Mtx::local_index_type,
typename Mtx::global_index_type>::
build()
.with_local_solver(
preconditioner::Jacobi<
value_type>::build()
.with_max_block_size(1u)))
.on(exec)
->generate(matrix);
});
}
#endif
if (dynamic_cast<const DpcppExecutor*>(exec.get())) {
using absolute_value_type = remove_complex<value_type>;
return solver::Gmres<value_type>::build()
Expand Down Expand Up @@ -695,7 +753,7 @@ void Multigrid::apply_impl(const LinOp* b, LinOp* x) const
void Multigrid::apply_with_initial_guess_impl(const LinOp* b, LinOp* x,
initial_guess_mode guess) const
{
if (!this->get_system_matrix()) {
if (!this->get_system_matrix() || !this->get_system_matrix()->get_size()) {
return;
}

Expand Down Expand Up @@ -729,7 +787,7 @@ void Multigrid::apply_with_initial_guess_impl(const LinOp* alpha,
LinOp* x,
initial_guess_mode guess) const
{
if (!this->get_system_matrix()) {
if (!this->get_system_matrix() || !this->get_system_matrix()->get_size()) {
return;
}

Expand Down
17 changes: 17 additions & 0 deletions test/matrix/csr_kernels2.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -607,6 +607,23 @@ TEST_F(Csr, AdvancedApplyToIdentityMatrixIsEquivalentToRef)
}


TEST_F(Csr, AdvancedApplyZeroToIdentityMatrixIsEquivalentToRef)
{
set_up_apply_data<Mtx::classical>();
auto a = Mtx::create(ref);
auto b = Mtx::create(ref);
auto da = gko::clone(exec, a);
auto db = gko::clone(exec, b);
auto id = gko::matrix::Identity<Mtx::value_type>::create(ref);
auto did = gko::matrix::Identity<Mtx::value_type>::create(exec);

a->apply(alpha, id, beta, b);
da->apply(dalpha, did, dbeta, db);

GKO_ASSERT_MTX_NEAR(b, db, 0);
}


TEST_F(Csr, ApplyToComplexIsEquivalentToRef)
{
set_up_apply_data<Mtx::classical>();
Expand Down
62 changes: 60 additions & 2 deletions test/mpi/solver/solver.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,15 @@
#include <ginkgo/core/distributed/vector.hpp>
#include <ginkgo/core/matrix/csr.hpp>
#include <ginkgo/core/matrix/dense.hpp>
#include <ginkgo/core/multigrid/pgm.hpp>
#include <ginkgo/core/solver/bicgstab.hpp>
#include <ginkgo/core/solver/cg.hpp>
#include <ginkgo/core/solver/cgs.hpp>
#include <ginkgo/core/solver/fcg.hpp>
#include <ginkgo/core/solver/gcr.hpp>
#include <ginkgo/core/solver/gmres.hpp>
#include <ginkgo/core/solver/ir.hpp>
#include <ginkgo/core/solver/multigrid.hpp>
#include <ginkgo/core/stop/residual_norm.hpp>


Expand Down Expand Up @@ -93,6 +95,8 @@ struct SimpleSolverTest {
ASSERT_EQ(mtx->get_preconditioner(), nullptr);
ASSERT_EQ(mtx->get_stopping_criterion_factory(), nullptr);
}

static bool blacklisted(const std::string& test) { return false; }
};


Expand All @@ -106,6 +110,56 @@ struct Cg : SimpleSolverTest<gko::solver::Cg<solver_value_type>> {
};


struct CgWithMg : SimpleSolverTest<gko::solver::Cg<solver_value_type>> {
static void preprocess(
gko::matrix_data<value_type, global_index_type>& data)
{
// the MG preconditioner doesn's seem to be able to handle
// random HPD matrices well, so replace it with a 3-pt stencil matrix
data.nonzeros.clear();
for (int i = 0; i < data.size[0]; ++i) {
if (i > 0) {
data.nonzeros.emplace_back(i, i - 1, value_type(-1));
}
data.nonzeros.emplace_back(i, i, value_type(2));
if (i < data.size[0] - 1) {
data.nonzeros.emplace_back(i, i + 1, value_type(-1));
}
}
}

static typename solver_type::parameters_type build(
std::shared_ptr<const gko::Executor> exec)
{
return SimpleSolverTest<gko::solver::Cg<solver_value_type>>::build(exec)
.with_preconditioner(
gko::solver::Multigrid::build()
.with_mg_level(
gko::multigrid::Pgm<solver_value_type>::build()
.with_deterministic(false))
.with_min_coarse_rows(
16u) // necessary since the test matrices have less
// rows than the default value
.with_criteria(
gko::stop::Iteration::build().with_max_iters(
iteration_count()),
gko::stop::ResidualNorm<value_type>::build()
.with_baseline(gko::stop::mode::absolute)
.with_reduction_factor(2 * reduction_factor())));
}

static bool blacklisted(const std::string& test)
{
#ifdef GKO_COMPILING_HIP
// SPGEAM is broken for empty matrices on rocm <= 4.5
return test == "some-empty-partition";
#else
return false;
#endif
}
};


struct Cgs : SimpleSolverTest<gko::solver::Cgs<solver_value_type>> {};


Expand Down Expand Up @@ -307,6 +361,9 @@ class Solver : public CommonMpiTestFixture {
}
{
SCOPED_TRACE("Some empty partition");
if (Config::blacklisted("some-empty-partition")) {
return;
}
guarded_fn(gen_part(50, std::max(1, comm.size() - 1)));
}
}
Expand Down Expand Up @@ -475,8 +532,9 @@ class Solver : public CommonMpiTestFixture {
std::default_random_engine rand_engine;
};

using SolverTypes = ::testing::Types<Cg, Cgs, Fcg, Bicgstab, Ir, Gcr<10u>,
Gcr<100u>, Gmres<10u>, Gmres<100u>>;
using SolverTypes =
::testing::Types<Cg, CgWithMg, Cgs, Fcg, Bicgstab, Ir, Gcr<10u>, Gcr<100u>,
Gmres<10u>, Gmres<100u>>;

TYPED_TEST_SUITE(Solver, SolverTypes, TypenameNameGenerator);

Expand Down

0 comments on commit 6abd3f3

Please sign in to comment.