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

Jump to bottom

Reuse pgm and multigrid #1681

Draft
wants to merge 3 commits into
base: develop
Choose a base branch
from
Draft

Reuse pgm and multigrid #1681

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
48b282d
allow regenerate the pgm on pregenerated information
yhmtsai Aug 7, 2024
00f1c9f
update matrix value in multigrid
yhmtsai Aug 13, 2024
5b91a3a
add get_local_solver in schwarz
yhmtsai Aug 21, 2024
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
151 changes: 131 additions & 20 deletions core/multigrid/pgm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,10 +221,10 @@ Pgm<ValueType, IndexType>::generate_local(
exec->run(pgm::make_assign_to_exist_agg(weight_mtx.get(), diag.get(),
agg_, intermediate_agg));
}
IndexType num_agg = 0;
num_agg_ = 0;
// Renumber the index
exec->run(pgm::make_renumber(agg_, &num_agg));
gko::dim<2>::dimension_type coarse_dim = num_agg;
exec->run(pgm::make_renumber(agg_, &num_agg_));
gko::dim<2>::dimension_type coarse_dim = num_agg_;
auto fine_dim = local_matrix->get_size()[0];
// prolong_row_gather is the lightway implementation for prolongation
auto prolong_row_gather = share(matrix::RowGatherer<IndexType>::create(
Expand All @@ -234,13 +234,13 @@ Pgm<ValueType, IndexType>::generate_local(
auto restrict_sparsity =
share(matrix::SparsityCsr<ValueType, IndexType>::create(
exec, gko::dim<2>{coarse_dim, fine_dim}, fine_dim));
agg_to_restrict(exec, num_agg, agg_, restrict_sparsity->get_row_ptrs(),
agg_to_restrict(exec, num_agg_, agg_, restrict_sparsity->get_row_ptrs(),
restrict_sparsity->get_col_idxs());

// Construct the coarse matrix
// TODO: improve it
auto coarse_matrix =
generate_coarse(exec, local_matrix.get(), num_agg, agg_);
generate_coarse(exec, local_matrix.get(), num_agg_, agg_);

return std::tie(prolong_row_gather, coarse_matrix, restrict_sparsity);
}
Expand Down Expand Up @@ -457,40 +457,38 @@ void Pgm<ValueType, IndexType>::generate()
communicate(matrix, agg_, non_local_agg);
// generate non_local_col_map
non_local_agg.set_executor(exec->get_master());
array<IndexType> non_local_col_map(exec->get_master(),
non_local_size);
non_local_col_map_.set_executor(exec->get_master());
non_local_col_map_.resize_and_reset(non_local_size);
// add additional entry in tail such that the offset easily handle
// it.
array<IndexType> renumber(exec->get_master(), non_local_size + 1);
auto recv_offsets = matrix->recv_offsets_;
generate_non_local_map(recv_offsets, non_local_agg,
non_local_col_map, renumber);
non_local_col_map_, renumber);

// get new recv_size and recv_offsets
std::vector<experimental::distributed::comm_index_type>
new_recv_size(num_rank);
std::vector<experimental::distributed::comm_index_type>
new_recv_offsets(num_rank + 1);
array<IndexType> new_recv_gather_idxs(exec->get_master());
new_recv_size_.resize(num_rank);
new_recv_offsets_.resize(num_rank + 1);
new_recv_gather_idxs_.set_executor(exec->get_master());
compute_communication(recv_offsets, non_local_agg, renumber,
new_recv_size, new_recv_offsets,
new_recv_gather_idxs);
new_recv_size_, new_recv_offsets_,
new_recv_gather_idxs_);

non_local_col_map.set_executor(exec);
IndexType non_local_num_agg = new_recv_gather_idxs.get_size();
non_local_col_map_.set_executor(exec);
non_local_num_agg_ = new_recv_gather_idxs_.get_size();
// build csr from row and col map
// unlike non-distributed version, generate_coarse uses different
// row and col maps.
auto result_non_local_csr = generate_coarse(
exec, non_local_csr.get(),
static_cast<IndexType>(std::get<1>(result)->get_size()[0]),
agg_, non_local_num_agg, non_local_col_map);
agg_, non_local_num_agg_, non_local_col_map_);
// use local and non-local to build coarse matrix
// also restriction and prolongation (Local-only-global matrix)
auto coarse_size =
static_cast<int64>(std::get<1>(result)->get_size()[0]);
comm.all_reduce(exec->get_master(), &coarse_size, 1, MPI_SUM);
new_recv_gather_idxs.set_executor(exec);
new_recv_gather_idxs_.set_executor(exec);

// setup the generated linop.
using global_index_type =
Expand All @@ -500,7 +498,8 @@ void Pgm<ValueType, IndexType>::generate()
Matrix<ValueType, IndexType, global_index_type>::create(
exec, comm, gko::dim<2>(coarse_size, coarse_size),
std::get<1>(result), result_non_local_csr,
new_recv_size, new_recv_offsets, new_recv_gather_idxs));
new_recv_size_, new_recv_offsets_,
new_recv_gather_idxs_));
auto restrict_op = share(
experimental::distributed::
Matrix<ValueType, IndexType, global_index_type>::create(
Expand Down Expand Up @@ -540,6 +539,118 @@ void Pgm<ValueType, IndexType>::generate()
}


template <typename ValueType, typename IndexType>
void Pgm<ValueType, IndexType>::update_matrix_value(
std::shared_ptr<const LinOp> new_matrix)
{
using csr_type = matrix::Csr<ValueType, IndexType>;
system_matrix_ = new_matrix;
#if GINKGO_BUILD_MPI
if (std::dynamic_pointer_cast<
const experimental::distributed::DistributedBase>(system_matrix_)) {
auto convert_fine_op = [&](auto matrix) {
using global_index_type = typename std::decay_t<
decltype(*matrix)>::result_type::global_index_type;
auto exec = as<LinOp>(matrix)->get_executor();
auto comm = as<experimental::distributed::DistributedBase>(matrix)
->get_communicator();
auto fine = share(
experimental::distributed::
Matrix<ValueType, IndexType, global_index_type>::create(
exec, comm,
matrix::Csr<ValueType, IndexType>::create(exec),
matrix::Csr<ValueType, IndexType>::create(exec)));
matrix->convert_to(fine);
this->set_fine_op(fine);
};
auto setup_fine_op = [&](auto matrix) {
// Only support csr matrix currently.
auto local_csr = std::dynamic_pointer_cast<const csr_type>(
matrix->get_local_matrix());
auto non_local_csr = std::dynamic_pointer_cast<const csr_type>(
matrix->get_non_local_matrix());
// If system matrix is not csr or need sorting, generate the csr.
if (!parameters_.skip_sorting || !local_csr || !non_local_csr) {
using global_index_type =
typename std::decay_t<decltype(*matrix)>::global_index_type;
convert_fine_op(
as<ConvertibleTo<experimental::distributed::Matrix<
ValueType, IndexType, global_index_type>>>(matrix));
}
};

using fst_mtx_type =
experimental::distributed::Matrix<ValueType, IndexType, IndexType>;
using snd_mtx_type =
experimental::distributed::Matrix<ValueType, IndexType, int64>;
// setup the fine op using Csr with current ValueType
// we do not use dispatcher run in the first place because we have the
// fallback option for that.
if (auto obj =
std::dynamic_pointer_cast<const fst_mtx_type>(system_matrix_)) {
setup_fine_op(obj);
} else if (auto obj = std::dynamic_pointer_cast<const snd_mtx_type>(
system_matrix_)) {
setup_fine_op(obj);
} else {
// handle other ValueTypes.
run<ConvertibleTo, fst_mtx_type, snd_mtx_type>(obj,
convert_fine_op);
}

auto distributed_setup = [&](auto matrix) {
auto exec = gko::as<LinOp>(matrix)->get_executor();
auto comm =
gko::as<experimental::distributed::DistributedBase>(matrix)
->get_communicator();
auto num_rank = comm.size();
auto pgm_local_op =
gko::as<const csr_type>(matrix->get_local_matrix());

auto coarse_local_matrix = generate_coarse(
this->get_executor(), pgm_local_op.get(), num_agg_, agg_);

auto non_local_csr =
as<const csr_type>(matrix->get_non_local_matrix());

auto result_non_local_csr =
generate_coarse(exec, non_local_csr.get(), num_agg_, agg_,
non_local_num_agg_, non_local_col_map_);

// setup the generated linop.
using global_index_type =
typename std::decay_t<decltype(*matrix)>::global_index_type;
auto coarse = share(
experimental::distributed::
Matrix<ValueType, IndexType, global_index_type>::create(
exec, comm, this->get_coarse_op()->get_size(),
coarse_local_matrix, result_non_local_csr,
new_recv_size_, new_recv_offsets_,
new_recv_gather_idxs_));
this->set_multigrid_level(this->get_prolong_op(), coarse,
this->get_restrict_op());
};
// the fine op is using csr with the current ValueType
run<fst_mtx_type, snd_mtx_type>(this->get_fine_op(), distributed_setup);
} else
#endif
{
auto pgm_op = std::dynamic_pointer_cast<const csr_type>(system_matrix_);
// If system matrix is not csr or need sorting, generate the csr.
if (!parameters_.skip_sorting || !pgm_op) {
pgm_op = convert_to_with_sorting<csr_type>(
this->get_executor(), system_matrix_, parameters_.skip_sorting);
// keep the same precision data in fine_op
this->set_fine_op(pgm_op);
}
auto coarse_matrix =
generate_coarse(this->get_executor(), pgm_op.get(), num_agg_, agg_);
this->set_multigrid_level(this->get_prolong_op(), coarse_matrix,
this->get_restrict_op());
}
}


#define GKO_DECLARE_PGM(_vtype, _itype) class Pgm<_vtype, _itype>
GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_PGM);

Expand Down
140 changes: 140 additions & 0 deletions core/solver/multigrid.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -833,6 +833,146 @@ void Multigrid::generate()
}


void Multigrid::update_matrix_value(std::shared_ptr<const LinOp> new_matrix)
{
this->set_system_matrix(new_matrix);
// generate coarse matrix until reaching max_level or min_coarse_rows
auto num_rows = this->get_system_matrix()->get_size()[0];
size_type level = 0;
auto matrix = this->get_system_matrix();
auto exec = this->get_executor();
// clean all smoother
pre_smoother_list_.clear();
mid_smoother_list_.clear();
post_smoother_list_.clear();
// Always generate smoother with size = level.
for (int i = 0; i < mg_level_list_.size(); i++) {
auto mg_level = mg_level_list_.at(i);
as<UpdateMatrixValue>(mg_level)->update_matrix_value(matrix);
matrix = mg_level->get_coarse_op();
run<gko::multigrid::EnableMultigridLevel, float, double,
std::complex<float>, std::complex<double>>(
mg_level,
[this](auto mg_level, auto index, auto matrix) {
using value_type =
typename std::decay_t<decltype(*mg_level)>::value_type;
handle_list<value_type>(
index, matrix, parameters_.pre_smoother, pre_smoother_list_,
parameters_.smoother_iters, parameters_.smoother_relax);
if (parameters_.mid_case ==
multigrid::mid_smooth_type::standalone) {
handle_list<value_type>(
index, matrix, parameters_.mid_smoother,
mid_smoother_list_, parameters_.smoother_iters,
parameters_.smoother_relax);
}
if (!parameters_.post_uses_pre) {
handle_list<value_type>(
index, matrix, parameters_.post_smoother,
post_smoother_list_, parameters_.smoother_iters,
parameters_.smoother_relax);
}
},
i, mg_level->get_fine_op());
}
if (parameters_.post_uses_pre) {
post_smoother_list_ = pre_smoother_list_;
}
auto last_mg_level = mg_level_list_.back();

// generate coarsest solver
run<gko::multigrid::EnableMultigridLevel, float, double,
std::complex<float>, std::complex<double>>(
last_mg_level,
[this](auto mg_level, auto level, auto matrix) {
using value_type =
typename std::decay_t<decltype(*mg_level)>::value_type;
auto exec = this->get_executor();
// 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
#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()
.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(
preconditioner::Jacobi<value_type>::build()
.with_max_block_size(1u))
.on(exec)
->generate(matrix);
} else {
return experimental::solver::Direct<value_type,
int32>::build()
.with_factorization(
experimental::factorization::Lu<value_type,
int32>::build())
.on(exec)
->generate(matrix);
}
};
if (parameters_.coarsest_solver.size() == 0) {
coarsest_solver_ = gen_default_solver();
} else {
auto temp_index = solver_selector_(level, matrix.get());
GKO_ENSURE_IN_BOUNDS(temp_index,
parameters_.coarsest_solver.size());
auto solver = parameters_.coarsest_solver.at(temp_index);
if (solver == nullptr) {
coarsest_solver_ = gen_default_solver();
} else {
coarsest_solver_ = solver->generate(matrix);
}
}
},
mg_level_list_.size(), matrix);
}


void Multigrid::apply_impl(const LinOp* b, LinOp* x) const
{
this->apply_with_initial_guess_impl(b, x,
Expand Down
5 changes: 5 additions & 0 deletions include/ginkgo/core/distributed/preconditioner/schwarz.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,11 @@ class Schwarz
GKO_ENABLE_LIN_OP_FACTORY(Schwarz, parameters, Factory);
GKO_ENABLE_BUILD_METHOD(Factory);

std::shared_ptr<const LinOp> get_local_solver() const
{
return local_solver_;
}

protected:
/**
* Creates an empty Schwarz preconditioner.
Expand Down
8 changes: 8 additions & 0 deletions include/ginkgo/core/multigrid/multigrid_level.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,14 @@


namespace gko {


class UpdateMatrixValue {
public:
virtual void update_matrix_value(std::shared_ptr<const gko::LinOp>) = 0;
};


/**
* @brief The multigrid components namespace.
*
Expand Down
Loading
Loading