diff --git a/core/multigrid/pgm.cpp b/core/multigrid/pgm.cpp index 9f1f5b50ba6..1bfdf718e97 100644 --- a/core/multigrid/pgm.cpp +++ b/core/multigrid/pgm.cpp @@ -221,10 +221,10 @@ Pgm::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::create( @@ -234,13 +234,13 @@ Pgm::generate_local( auto restrict_sparsity = share(matrix::SparsityCsr::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); } @@ -457,40 +457,38 @@ void Pgm::generate() communicate(matrix, agg_, non_local_agg); // generate non_local_col_map non_local_agg.set_executor(exec->get_master()); - array 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 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 - new_recv_size(num_rank); - std::vector - new_recv_offsets(num_rank + 1); - array 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(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(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 = @@ -500,7 +498,8 @@ void Pgm::generate() Matrix::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::create( @@ -540,6 +539,118 @@ void Pgm::generate() } +template +void Pgm::update_matrix_value( + std::shared_ptr new_matrix) +{ + using csr_type = matrix::Csr; + 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(matrix)->get_executor(); + auto comm = as(matrix) + ->get_communicator(); + auto fine = share( + experimental::distributed:: + Matrix::create( + exec, comm, + matrix::Csr::create(exec), + matrix::Csr::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( + matrix->get_local_matrix()); + auto non_local_csr = std::dynamic_pointer_cast( + 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::global_index_type; + convert_fine_op( + as>>(matrix)); + } + }; + + using fst_mtx_type = + experimental::distributed::Matrix; + using snd_mtx_type = + experimental::distributed::Matrix; + // 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(system_matrix_)) { + setup_fine_op(obj); + } else if (auto obj = std::dynamic_pointer_cast( + system_matrix_)) { + setup_fine_op(obj); + } else { + // handle other ValueTypes. + run(obj, + convert_fine_op); + } + + auto distributed_setup = [&](auto matrix) { + auto exec = gko::as(matrix)->get_executor(); + auto comm = + gko::as(matrix) + ->get_communicator(); + auto num_rank = comm.size(); + auto pgm_local_op = + gko::as(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(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::global_index_type; + auto coarse = share( + experimental::distributed:: + Matrix::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(this->get_fine_op(), distributed_setup); + } else +#endif + { + auto pgm_op = std::dynamic_pointer_cast(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( + 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); diff --git a/core/solver/multigrid.cpp b/core/solver/multigrid.cpp index 35ad7b5d1fe..835aa22eed5 100644 --- a/core/solver/multigrid.cpp +++ b/core/solver/multigrid.cpp @@ -833,6 +833,146 @@ void Multigrid::generate() } +void Multigrid::update_matrix_value(std::shared_ptr 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(mg_level)->update_matrix_value(matrix); + matrix = mg_level->get_coarse_op(); + run, std::complex>( + mg_level, + [this](auto mg_level, auto index, auto matrix) { + using value_type = + typename std::decay_t::value_type; + handle_list( + 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( + index, matrix, parameters_.mid_smoother, + mid_smoother_list_, parameters_.smoother_iters, + parameters_.smoother_relax); + } + if (!parameters_.post_uses_pre) { + handle_list( + 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, std::complex>( + last_mg_level, + [this](auto mg_level, auto level, auto matrix) { + using value_type = + typename std::decay_t::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 { + // TODO: unify when dpcpp supports direct solver +#if GINKGO_BUILD_MPI + if (gko::detail::is_distributed(matrix.get())) { + using absolute_value_type = remove_complex; + using experimental::distributed::Matrix; + return run, + Matrix, + Matrix>(matrix, [exec](auto matrix) { + using Mtx = typename decltype(matrix)::element_type; + return solver::Gmres::build() + .with_criteria( + stop::Iteration::build().with_max_iters( + matrix->get_size()[0]), + stop::ResidualNorm::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:: + build() + .with_local_solver( + preconditioner::Jacobi< + value_type>::build() + .with_max_block_size(1u))) + .on(exec) + ->generate(matrix); + }); + } +#endif + if (dynamic_cast(exec.get())) { + using absolute_value_type = remove_complex; + return solver::Gmres::build() + .with_criteria( + stop::Iteration::build().with_max_iters( + matrix->get_size()[0]), + stop::ResidualNorm::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::build() + .with_max_block_size(1u)) + .on(exec) + ->generate(matrix); + } else { + return experimental::solver::Direct::build() + .with_factorization( + experimental::factorization::Lu::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, diff --git a/include/ginkgo/core/distributed/preconditioner/schwarz.hpp b/include/ginkgo/core/distributed/preconditioner/schwarz.hpp index badd5ba7dd3..1de9971e7f2 100644 --- a/include/ginkgo/core/distributed/preconditioner/schwarz.hpp +++ b/include/ginkgo/core/distributed/preconditioner/schwarz.hpp @@ -78,6 +78,11 @@ class Schwarz GKO_ENABLE_LIN_OP_FACTORY(Schwarz, parameters, Factory); GKO_ENABLE_BUILD_METHOD(Factory); + std::shared_ptr get_local_solver() const + { + return local_solver_; + } + protected: /** * Creates an empty Schwarz preconditioner. diff --git a/include/ginkgo/core/multigrid/multigrid_level.hpp b/include/ginkgo/core/multigrid/multigrid_level.hpp index 7c5b7e09684..caa204295f7 100644 --- a/include/ginkgo/core/multigrid/multigrid_level.hpp +++ b/include/ginkgo/core/multigrid/multigrid_level.hpp @@ -17,6 +17,14 @@ namespace gko { + + +class UpdateMatrixValue { +public: + virtual void update_matrix_value(std::shared_ptr) = 0; +}; + + /** * @brief The multigrid components namespace. * diff --git a/include/ginkgo/core/multigrid/pgm.hpp b/include/ginkgo/core/multigrid/pgm.hpp index d07001be2f1..906926a261b 100644 --- a/include/ginkgo/core/multigrid/pgm.hpp +++ b/include/ginkgo/core/multigrid/pgm.hpp @@ -49,7 +49,8 @@ namespace multigrid { */ template class Pgm : public EnableLinOp>, - public EnableMultigridLevel { + public EnableMultigridLevel, + public UpdateMatrixValue { friend class EnableLinOp; friend class EnablePolymorphicObject; @@ -149,6 +150,8 @@ class Pgm : public EnableLinOp>, const config::type_descriptor& td_for_child = config::make_type_descriptor()); + void update_matrix_value(std::shared_ptr new_matrix) override; + protected: void apply_impl(const LinOp* b, LinOp* x) const override { @@ -172,6 +175,11 @@ class Pgm : public EnableLinOp>, parameters_{factory->get_parameters()}, system_matrix_{system_matrix}, agg_(factory->get_executor(), system_matrix_->get_size()[0]) +#if GINKGO_BUILD_MPI + , + non_local_col_map_(factory->get_executor()), + new_recv_gather_idxs_(factory->get_executor()) +#endif { GKO_ASSERT(parameters_.max_unassigned_ratio <= 1.0); GKO_ASSERT(parameters_.max_unassigned_ratio >= 0.0); @@ -205,6 +213,14 @@ class Pgm : public EnableLinOp>, private: std::shared_ptr system_matrix_{}; array agg_; + IndexType num_agg_; +#if GINKGO_BUILD_MPI + IndexType non_local_num_agg_; + array non_local_col_map_; + std::vector new_recv_size_; + std::vector new_recv_offsets_; + array new_recv_gather_idxs_; +#endif }; diff --git a/include/ginkgo/core/solver/multigrid.hpp b/include/ginkgo/core/solver/multigrid.hpp index 2d0278b538e..99e75b971ca 100644 --- a/include/ginkgo/core/solver/multigrid.hpp +++ b/include/ginkgo/core/solver/multigrid.hpp @@ -107,7 +107,8 @@ class MultigridState; class Multigrid : public EnableLinOp, public EnableSolverBase, public EnableIterativeBase, - public EnableApplyWithInitialGuess { + public EnableApplyWithInitialGuess, + public UpdateMatrixValue { friend class EnableLinOp; friend class EnablePolymorphicObject; friend class EnableApplyWithInitialGuess; @@ -133,7 +134,12 @@ class Multigrid : public EnableLinOp, std::vector> get_mg_level_list() const { - return mg_level_list_; + std::vector> + const_copy(mg_level_list_.size()); + for (int i = 0; i < mg_level_list_.size(); i++) { + const_copy.at(i) = mg_level_list_.at(i); + } + return const_copy; } /** @@ -400,6 +406,9 @@ class Multigrid : public EnableLinOp, const config::type_descriptor& td_for_child = config::make_type_descriptor<>()); + void update_matrix_value( + std::shared_ptr new_matrix) override; + protected: void apply_impl(const LinOp* b, LinOp* x) const override; @@ -447,7 +456,7 @@ class Multigrid : public EnableLinOp, void create_state() const; private: - std::vector> + std::vector> mg_level_list_{}; std::vector> pre_smoother_list_{}; std::vector> mid_smoother_list_{};