diff --git a/common/unified/multigrid/rs_kernels.cpp b/common/unified/multigrid/rs_kernels.cpp index 41a50afd713..2d05bb88248 100644 --- a/common/unified/multigrid/rs_kernels.cpp +++ b/common/unified/multigrid/rs_kernels.cpp @@ -72,7 +72,8 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( template void compute_soc_and_run_rs( std::shared_ptr exec, - matrix::view::csr A, double theta, + matrix::view::csr A, + matrix::view::csr off_diag, double theta, array& is_strong, array& lambda, array& cf_marker, IndexType& coarse_size) { @@ -80,6 +81,11 @@ void compute_soc_and_run_rs( const auto* a_row_ptrs = A.row_ptrs; const auto* a_col_idxs = A.col_idxs; const auto* a_vals = A.values; + // the off-diagonal block of a distributed matrix is stored on this rank + // and has the same rows as A, so it contributes to the strength threshold + // without any communication. It is empty for a non-distributed matrix. + const auto* od_row_ptrs = off_diag.row_ptrs; + const auto* od_vals = off_diag.values; bool* is_strong_vals = is_strong.get_data(); auto* lambda_vals = lambda.get_data(); auto* cf = cf_marker.get_data(); @@ -88,6 +94,7 @@ void compute_soc_and_run_rs( run_kernel( exec, [theta] GKO_KERNEL(auto i, auto row_ptrs, auto col_idxs, auto vals, + auto od_row_ptrs, auto od_vals, auto is_strong_vals) { auto max_offdiag = zero(); for (auto jj = row_ptrs[i]; jj < row_ptrs[i + 1]; ++jj) { @@ -95,6 +102,13 @@ void compute_soc_and_run_rs( max_offdiag = gko::max(max_offdiag, -real(vals[jj])); } } + // remote couplings of this row. None of them can be the diagonal + // entry, since they all live in other ranks' row ranges. + if (od_row_ptrs != nullptr) { + for (auto jj = od_row_ptrs[i]; jj < od_row_ptrs[i + 1]; ++jj) { + max_offdiag = gko::max(max_offdiag, -real(od_vals[jj])); + } + } const auto threshold = theta * static_cast(max_offdiag); for (auto jj = row_ptrs[i]; jj < row_ptrs[i + 1]; ++jj) { const auto j = col_idxs[jj]; @@ -103,7 +117,8 @@ void compute_soc_and_run_rs( static_cast(-real(vals[jj])) >= threshold); } }, - n, a_row_ptrs, a_col_idxs, a_vals, is_strong_vals); + n, a_row_ptrs, a_col_idxs, a_vals, od_row_ptrs, od_vals, + is_strong_vals); /// 2. COMPUTE lambda_i = number of strong nbrs run_kernel( @@ -207,6 +222,34 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( GKO_DECLARE_RS_COMPUTE_SOC_AND_RUN_RS_KERNEL); +template +void mark_forced_c_points(std::shared_ptr exec, + size_type num_forced, const IndexType* forced_rows, + array& cf_marker, IndexType& coarse_size) +{ + auto* cf = cf_marker.get_data(); + run_kernel( + exec, + [] GKO_KERNEL(auto i, auto forced_rows, auto cf) { + cf[forced_rows[i]] = 1; // C-point + }, + num_forced, forced_rows, cf); + + array d_coarse_size(exec, 1); + run_kernel_reduction( + exec, + [] GKO_KERNEL(auto i, auto cf) { + return cf[i] == 1 ? IndexType{1} : IndexType{0}; + }, + GKO_KERNEL_REDUCE_SUM(IndexType), d_coarse_size.get_data(), + cf_marker.get_size(), cf); + + coarse_size = get_element(d_coarse_size, 0); +} + +GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_RS_MARK_FORCED_C_POINTS_KERNEL); + + template void fill_coarse_and_compute_prolong_row_ptrs( std::shared_ptr exec, diff --git a/core/distributed/helpers.hpp b/core/distributed/helpers.hpp index 4895e97d429..56cbe728a04 100644 --- a/core/distributed/helpers.hpp +++ b/core/distributed/helpers.hpp @@ -9,6 +9,9 @@ #include #include +#include +#include +#include #include #include #include @@ -203,6 +206,68 @@ inline const LinOp* get_local(const LinOp* mtx) } +#if GINKGO_BUILD_MPI + + +/** + * Exchanges one value per halo index with the neighboring ranks. + * + * Distributed coarsenings need to know what the owner of a non-local index + * did with it: Pgm sends the aggregate an index was assigned to, Rs sends the + * coarse index a forced C-point was renumbered to. Both are a single value per + * halo index, so both reduce to this exchange. + * + * @param send_buffer one value per send index of the collective communicator, + * in its send index order + * + * @return one value per receive index, in the communicator's receive order + */ +template +array exchange_with_neighbors( + std::shared_ptr exec, + const experimental::mpi::communicator& comm, + const experimental::mpi::CollectiveCommunicator* coll_comm, + const array& send_buffer) +{ + const auto total_send_size = + static_cast(coll_comm->get_send_size()); + const auto total_recv_size = + static_cast(coll_comm->get_recv_size()); + GKO_ASSERT_EQ(send_buffer.get_size(), total_send_size); + array recv_buffer(exec, total_recv_size); + + // not every executor/MPI combination can send from device memory + auto use_host_buffer = experimental::mpi::requires_host_buffer(exec, comm); + array host_send_buffer(exec->get_master()); + array host_recv_buffer(exec->get_master()); + if (use_host_buffer) { + host_send_buffer.resize_and_reset(total_send_size); + host_recv_buffer.resize_and_reset(total_recv_size); + exec->get_master()->copy_from(exec, total_send_size, + send_buffer.get_const_data(), + host_send_buffer.get_data()); + } + + const auto send_ptr = use_host_buffer ? host_send_buffer.get_const_data() + : send_buffer.get_const_data(); + auto recv_ptr = + use_host_buffer ? host_recv_buffer.get_data() : recv_buffer.get_data(); + exec->synchronize(); + coll_comm + ->i_all_to_all_v(use_host_buffer ? exec->get_master() : exec, send_ptr, + recv_ptr) + .wait(); + if (use_host_buffer) { + exec->copy_from(exec->get_master(), total_recv_size, recv_ptr, + recv_buffer.get_data()); + } + return recv_buffer; +} + + +#endif + + } // namespace detail } // namespace gko diff --git a/core/multigrid/pgm.cpp b/core/multigrid/pgm.cpp index a895df87f0e..715400c63b3 100644 --- a/core/multigrid/pgm.cpp +++ b/core/multigrid/pgm.cpp @@ -30,6 +30,7 @@ #include "core/components/fill_array_kernels.hpp" #include "core/components/format_conversion_kernels.hpp" #include "core/config/config_helper.hpp" +#include "core/distributed/helpers.hpp" #include "core/distributed/index_map_kernels.hpp" #include "core/matrix/csr_builder.hpp" #include "core/multigrid/pgm_kernels.hpp" @@ -283,11 +284,9 @@ array Pgm::communicate_off_diag_agg( auto exec = matrix->get_executor(); const auto comm = matrix->get_communicator(); auto coll_comm = matrix->row_gatherer_->get_collective_communicator(); - auto total_send_size = coll_comm->get_send_size(); - auto total_recv_size = coll_comm->get_recv_size(); auto row_gatherer = matrix->row_gatherer_; - array send_agg(exec, total_send_size); + array send_agg(exec, coll_comm->get_send_size()); exec->run(pgm::make_gather_index( send_agg.get_size(), local_agg.get_const_data(), row_gatherer->get_const_send_idxs(), send_agg.get_data())); @@ -300,33 +299,8 @@ array Pgm::communicate_off_diag_agg( device_segmented_array{}, comm.rank(), send_agg, experimental::distributed::index_space::local, send_global_agg)); - array off_diag_agg(exec, total_recv_size); - - auto use_host_buffer = experimental::mpi::requires_host_buffer(exec, comm); - array host_recv_buffer(exec->get_master()); - array host_send_buffer(exec->get_master()); - if (use_host_buffer) { - host_recv_buffer.resize_and_reset(total_recv_size); - host_send_buffer.resize_and_reset(total_send_size); - exec->get_master()->copy_from(exec, total_send_size, - send_global_agg.get_data(), - host_send_buffer.get_data()); - } - - const auto send_ptr = use_host_buffer ? host_send_buffer.get_const_data() - : send_global_agg.get_const_data(); - auto recv_ptr = - use_host_buffer ? host_recv_buffer.get_data() : off_diag_agg.get_data(); - exec->synchronize(); - coll_comm - ->i_all_to_all_v(use_host_buffer ? exec->get_master() : exec, send_ptr, - recv_ptr) - .wait(); - if (use_host_buffer) { - exec->copy_from(exec->get_master(), total_recv_size, recv_ptr, - off_diag_agg.get_data()); - } - return off_diag_agg; + return gko::detail::exchange_with_neighbors(exec, comm, coll_comm.get(), + send_global_agg); } diff --git a/core/multigrid/rs.cpp b/core/multigrid/rs.cpp index 2c10366b7a8..2c502c8efe7 100644 --- a/core/multigrid/rs.cpp +++ b/core/multigrid/rs.cpp @@ -7,18 +7,29 @@ #include #include #include +#include #include #include #include +#include +#include +#include +#include +#include #include #include #include #include #include +#include "core/base/dispatch_helper.hpp" #include "core/base/utils.hpp" #include "core/components/fill_array_kernels.hpp" #include "core/config/config_helper.hpp" +#include "core/distributed/helpers.hpp" +#include "core/distributed/index_map_kernels.hpp" +#include "core/multigrid/pgm_kernels.hpp" +#include "core/multigrid/rs_helpers.hpp" #include "core/multigrid/rs_kernels.hpp" @@ -33,34 +44,46 @@ GKO_REGISTER_OPERATION(fill_seq_array, components::fill_seq_array); GKO_REGISTER_OPERATION(check_m_matrix, rs::check_m_matrix); GKO_REGISTER_OPERATION(compute_soc_and_run_rs, rs::compute_soc_and_run_rs); +GKO_REGISTER_OPERATION(mark_forced_c_points, rs::mark_forced_c_points); GKO_REGISTER_OPERATION(fill_coarse_and_compute_prolong_row_ptrs, rs::fill_coarse_and_compute_prolong_row_ptrs); GKO_REGISTER_OPERATION(compute_interpolation, rs::compute_interpolation); +// reuse Pgm's kernel +GKO_REGISTER_OPERATION(gather_index, pgm::gather_index); + } // anonymous namespace } // namespace rs +namespace index_map { +namespace { + + +GKO_REGISTER_OPERATION(map_to_global, index_map::map_to_global); + + +} +} // namespace index_map template -void Rs::generate() +std::tuple, std::shared_ptr, + std::shared_ptr> +Rs::generate_local( + std::shared_ptr> local_matrix, + const matrix::Csr* off_diag_matrix, + size_type num_forced_c_points, const IndexType* forced_c_points) { using csr_type = matrix::Csr; - using real_type = remove_complex; auto exec = this->get_executor(); - const auto fine_dim = this->system_matrix_->get_size()[0]; - - const csr_type* rs_op = dynamic_cast(system_matrix_.get()); - std::shared_ptr rs_op_shared_ptr{}; + const auto* rs_op = local_matrix.get(); + const auto fine_dim = rs_op->get_size()[0]; - if (!parameters_.skip_sorting || !rs_op) { - rs_op_shared_ptr = convert_to_with_sorting( - exec, system_matrix_, parameters_.skip_sorting); - rs_op = rs_op_shared_ptr.get(); - this->set_fine_op(rs_op_shared_ptr); - } array is_m_matrix_array(exec, 1); if (!parameters_.skip_m_matrix_check) { + // for a distributed matrix this only sees the local block: the + // off-diagonal block has no diagonal entries of its own, so it cannot + // be checked by the same kernel exec->run(rs::make_check_m_matrix(rs_op->get_const_device_view(), is_m_matrix_array)); if (!exec->copy_val_to_host(is_m_matrix_array.get_const_data())) { @@ -76,22 +99,36 @@ void Rs::generate() array lambda(exec, fine_dim); array cf_marker(exec, fine_dim); IndexType coarse_dim{}; + // only a distributed matrix has an off-diagonal block; the empty view + // makes the kernel skip the remote couplings + const auto off_diag_view = + off_diag_matrix ? off_diag_matrix->get_const_device_view() + : rs::no_off_diag_view(); // build Strength-of-Connection (SOC) mask, 1 byte per NNZ of the system // matrix, compute lambda, perform greedy RS C/F splitting: // 0 = undecided, 1 = C, -1 = F, // then extract coarse dims exec->run(rs::make_compute_soc_and_run_rs( - rs_op->get_const_device_view(), parameters_.strength_threshold, - is_strong, lambda, cf_marker, coarse_dim)); + rs_op->get_const_device_view(), off_diag_view, + parameters_.strength_threshold, is_strong, lambda, cf_marker, + coarse_dim)); + + if (num_forced_c_points > 0) { + // promote the rows other ranks couple to, so that their prolongation + // rows are unit vectors and P stays block-diagonal + exec->run(rs::make_mark_forced_c_points( + num_forced_c_points, forced_c_points, cf_marker, coarse_dim)); + } + const size_type coarse_dim_size = static_cast(coarse_dim); // fill in coarse_rows and fine_to_coarse, build prolongation using // interpolation array coarse_rows(exec, coarse_dim_size); - array fine_to_coarse(exec, fine_dim); + fine_to_coarse_ = array(exec, fine_dim); array prolong_row_ptrs(exec, fine_dim + 1); exec->run(rs::make_fill_coarse_and_compute_prolong_row_ptrs( - cf_marker, coarse_rows, fine_to_coarse, rs_op->get_const_device_view(), + cf_marker, coarse_rows, fine_to_coarse_, rs_op->get_const_device_view(), is_strong, prolong_row_ptrs)); IndexType prolong_nnz = @@ -106,7 +143,7 @@ void Rs::generate() exec->run(rs::make_compute_interpolation( rs_op->get_const_device_view(), is_strong.get_const_data(), cf_marker, - fine_to_coarse.get_const_data(), prolong_op->get_device_view())); + fine_to_coarse_.get_const_data(), prolong_op->get_device_view())); // build restriction as R = P^T auto restrict_op = share(as(prolong_op->transpose())); @@ -122,7 +159,246 @@ void Rs::generate() rs_op->apply(prolong_op, tmp); restrict_op->apply(tmp, coarse_matrix); - this->set_multigrid_level(prolong_op, coarse_matrix, restrict_op); + return std::make_tuple(std::shared_ptr(prolong_op), + std::shared_ptr(coarse_matrix), + std::shared_ptr(restrict_op)); +} + + +#if GINKGO_BUILD_MPI + + +template +template +array +Rs::communicate_off_diag_coarse_idxs( + std::shared_ptr> + matrix, + std::shared_ptr< + experimental::distributed::Partition> + coarse_partition, + const array& local_fine_to_coarse) +{ + auto exec = matrix->get_executor(); + const auto comm = matrix->get_communicator(); + auto coll_comm = matrix->row_gatherer_->get_collective_communicator(); + auto row_gatherer = matrix->row_gatherer_; + + // every send index is a forced C-point, so its fine_to_coarse entry is a + // valid coarse index rather than the -1 an F-point would carry + array send_coarse(exec, coll_comm->get_send_size()); + exec->run(rs::make_gather_index( + send_coarse.get_size(), local_fine_to_coarse.get_const_data(), + row_gatherer->get_const_send_idxs(), send_coarse.get_data())); + + // There is no index map on the coarse level yet, so map the local indices + // to global indices on the coarse level manually + array send_global_coarse(exec, send_coarse.get_size()); + exec->run(index_map::make_map_to_global( + to_device_const(coarse_partition.get()), + device_segmented_array{}, comm.rank(), + send_coarse, experimental::distributed::index_space::local, + send_global_coarse)); + + return gko::detail::exchange_with_neighbors(exec, comm, coll_comm.get(), + send_global_coarse); +} + + +#endif + + +template +void Rs::generate() +{ + using csr_type = matrix::Csr; +#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 diag_csr = std::dynamic_pointer_cast( + matrix->get_diag_matrix()); + auto off_diag_csr = std::dynamic_pointer_cast( + matrix->get_off_diag_matrix()); + // If system matrix is not csr or need sorting, generate the + // csr. + if (!parameters_.skip_sorting || !diag_csr || !off_diag_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(system_matrix_, + convert_fine_op); + } + + auto distributed_setup = [&](auto matrix) { + using global_index_type = + typename std::decay_t::global_index_type; + using dist_mtx_type = + experimental::distributed::Matrix; + + auto exec = gko::as(matrix)->get_executor(); + auto comm = + gko::as(matrix) + ->get_communicator(); + auto local_op = gko::as(matrix->get_diag_matrix()); + auto off_diag_op = + gko::as(matrix->get_off_diag_matrix()); + + // Coarsen the local block. The off-diagonal block only enters the + // strength threshold; the C/F splitting itself is process-local. + // Every local row a neighbor couples to is forced into the coarse + // set, which makes its prolongation row a unit vector and keeps P + // block-diagonal - so P needs no off-diagonal block, and the whole + // halo information is one coarse index per off-diag column. + auto row_gatherer = matrix->row_gatherer_; + auto result = this->generate_local( + local_op, off_diag_op.get(), row_gatherer->get_num_send_idxs(), + row_gatherer->get_const_send_idxs()); + + // create the coarse partition + // the coarse partition will have only one range per part + // and only one part per rank. + // The global indices are ordered block-wise by rank, i.e. rank + // 0 owns [0, ..., N_1), rank 1 [N_1, ..., N_2), ... + const auto coarse_local_size = std::get<1>(result)->get_size()[0]; + auto coarse_partition = gko::share( + experimental::distributed::build_partition_from_local_size< + IndexType, global_index_type>( + exec, comm, static_cast(coarse_local_size))); + + // get the coarse global index of every off-diag column + auto off_diag_coarse_idxs = communicate_off_diag_coarse_idxs( + matrix, coarse_partition, fine_to_coarse_); + + // create a coarse index map based on the connections given by the + // off-diag coarse indices + auto coarse_imap = + experimental::distributed::index_map( + exec, coarse_partition, comm.rank(), off_diag_coarse_idxs); + + // a mapping from the fine off-diag indices to the coarse off-diag + // indices. off_diag_coarse_idxs already maps the fine off-diag + // indices to coarse global indices, so mapping it with the coarse + // index map results in the coarse off-diag indices. + auto off_diag_map = coarse_imap.map_to_local( + off_diag_coarse_idxs, + experimental::distributed::index_space::non_local); + const auto coarse_non_local_size = coarse_imap.get_non_local_size(); + const auto fine_non_local_size = off_diag_op->get_size()[1]; + + // The coarse off-diagonal block is R * A_off_diag * P_non_local. + // Every off-diag column belongs to a forced C-point, so the + // neighbor's prolongation rows for them are unit vectors and + // P_non_local degenerates into the 0/1 matrix given by + // off_diag_map. Both products are process-local, which is what + // makes an explicit distributed triple product unnecessary. + std::shared_ptr coarse_off_diag_op; + if (fine_non_local_size == 0) { + coarse_off_diag_op = share(csr_type::create( + exec, + gko::dim<2>{coarse_local_size, coarse_non_local_size})); + } else { + auto non_local_prolong = csr_type::create( + exec, + gko::dim<2>{fine_non_local_size, coarse_non_local_size}, + fine_non_local_size); + exec->run( + rs::make_fill_seq_array(non_local_prolong->get_row_ptrs(), + fine_non_local_size + 1)); + exec->copy_from(exec, fine_non_local_size, + off_diag_map.get_const_data(), + non_local_prolong->get_col_idxs()); + exec->run(rs::make_fill_array(non_local_prolong->get_values(), + fine_non_local_size, + one())); + + auto tmp = + csr_type::create(exec, gko::dim<2>{local_op->get_size()[0], + coarse_non_local_size}); + off_diag_op->apply(non_local_prolong, tmp); + coarse_off_diag_op = share(csr_type::create( + exec, + gko::dim<2>{coarse_local_size, coarse_non_local_size})); + gko::as(std::get<2>(result)) + ->apply(tmp, coarse_off_diag_op); + } + + // setup the generated linop. The prolongation and restriction have + // no off-diagonal block, see above. + auto coarse = share( + dist_mtx_type::create(exec, comm, std::move(coarse_imap), + std::get<1>(result), coarse_off_diag_op)); + auto restrict_op = share(dist_mtx_type::create( + exec, comm, + dim<2>(coarse->get_size()[0], + gko::as(matrix)->get_size()[0]), + std::get<2>(result))); + auto prolong_op = share(dist_mtx_type::create( + exec, comm, + dim<2>(gko::as(matrix)->get_size()[0], + coarse->get_size()[0]), + std::get<0>(result))); + this->set_multigrid_level(prolong_op, coarse, restrict_op); + }; + + // the fine op is using csr with the current ValueType + run(this->get_fine_op(), distributed_setup); + } else +#endif // GINKGO_BUILD_MPI + { + auto exec = this->get_executor(); + // Only support csr matrix currently. + auto rs_op = std::dynamic_pointer_cast(system_matrix_); + // If system matrix is not csr or need sorting, generate the csr. + if (!parameters_.skip_sorting || !rs_op) { + rs_op = convert_to_with_sorting(exec, system_matrix_, + parameters_.skip_sorting); + // keep the same precision data in fine_op + this->set_fine_op(rs_op); + } + auto result = this->generate_local(rs_op); + this->set_multigrid_level(std::get<0>(result), std::get<1>(result), + std::get<2>(result)); + } } diff --git a/core/multigrid/rs_helpers.hpp b/core/multigrid/rs_helpers.hpp index 32855bfb818..000fdbab778 100644 --- a/core/multigrid/rs_helpers.hpp +++ b/core/multigrid/rs_helpers.hpp @@ -10,8 +10,10 @@ #include #include +#include #include #include +#include namespace gko { @@ -19,6 +21,21 @@ namespace multigrid { namespace rs { +/** + * Returns the empty CSR view that stands in for a missing off-diagonal block. + * + * The RS kernels take the off-diagonal block of a distributed matrix as a + * device view, which - unlike a pointer - cannot be null. A non-distributed + * matrix has no such block and passes this view instead: its null row pointers + * tell the kernels that the local rows have no remote couplings. + */ +template +constexpr matrix::view::csr no_off_diag_view() +{ + return {dim<2>{}, 0, nullptr, nullptr, nullptr}; +} + + /** * Runs the greedy Ruge-Stueben C/F splitting on host-resident data. * diff --git a/core/multigrid/rs_kernels.hpp b/core/multigrid/rs_kernels.hpp index fccfc4bd7f1..cdb2bab1238 100644 --- a/core/multigrid/rs_kernels.hpp +++ b/core/multigrid/rs_kernels.hpp @@ -27,13 +27,31 @@ namespace rs { matrix::view::csr matrix, \ array& is_m_matrix_array) -#define GKO_DECLARE_RS_COMPUTE_SOC_AND_RUN_RS_KERNEL(ValueType, IndexType) \ - void compute_soc_and_run_rs( \ - std::shared_ptr exec, \ - matrix::view::csr A, double theta, \ - array& is_strong, array& lambda, \ +// `off_diag` is the off-diagonal block of a distributed matrix, i.e. the +// couplings of the local rows to rows owned by other ranks. It only widens +// max_offdiag in the strength-of-connection test, so that the threshold - and +// with it the strength mask of the local block - is the same one the +// non-distributed kernel would compute for the full row. A non-distributed +// matrix passes the empty view returned by `rs::no_off_diag_view`. +#define GKO_DECLARE_RS_COMPUTE_SOC_AND_RUN_RS_KERNEL(ValueType, IndexType) \ + void compute_soc_and_run_rs( \ + std::shared_ptr exec, \ + matrix::view::csr A, \ + matrix::view::csr off_diag, \ + double theta, array& is_strong, array& lambda, \ array& cf_marker, IndexType& coarse_dim) +// Turns the given rows into C-points and recomputes the number of C-points. +// Used in the distributed case to force every local row another rank couples +// to into the coarse set. Promoting an F-point to a C-point is always a valid +// RS splitting: it only enlarges the coarse set, so every remaining F-point +// keeps its strong C-neighbours. +#define GKO_DECLARE_RS_MARK_FORCED_C_POINTS_KERNEL(IndexType) \ + void mark_forced_c_points( \ + std::shared_ptr exec, size_type num_forced, \ + const IndexType* forced_rows, array& cf_marker, \ + IndexType& coarse_dim) + #define GKO_DECLARE_RS_FILL_COARSE_AND_COMPUTE_PROLONG_ROW_PTRS_KERNEL( \ ValueType, IndexType) \ void fill_coarse_and_compute_prolong_row_ptrs( \ @@ -57,6 +75,8 @@ namespace rs { GKO_DECLARE_RS_CHECK_M_MATRIX_KERNEL(ValueType, IndexType); \ template \ GKO_DECLARE_RS_COMPUTE_SOC_AND_RUN_RS_KERNEL(ValueType, IndexType); \ + template \ + GKO_DECLARE_RS_MARK_FORCED_C_POINTS_KERNEL(IndexType); \ template \ GKO_DECLARE_RS_FILL_COARSE_AND_COMPUTE_PROLONG_ROW_PTRS_KERNEL(ValueType, \ IndexType); \ diff --git a/include/ginkgo/core/distributed/matrix.hpp b/include/ginkgo/core/distributed/matrix.hpp index 3a68650671a..b3636be0bb8 100644 --- a/include/ginkgo/core/distributed/matrix.hpp +++ b/include/ginkgo/core/distributed/matrix.hpp @@ -39,8 +39,11 @@ namespace multigrid { template class Pgm; +template +class Rs; -} + +} // namespace multigrid namespace detail { @@ -280,6 +283,7 @@ class Matrix GlobalIndexType>; friend class multigrid::Pgm; + friend class multigrid::Rs; GKO_ASSERT_SUPPORTED_VALUE_AND_DIST_INDEX_TYPE; public: diff --git a/include/ginkgo/core/multigrid/rs.hpp b/include/ginkgo/core/multigrid/rs.hpp index bf20ae6c63d..dee366de889 100644 --- a/include/ginkgo/core/multigrid/rs.hpp +++ b/include/ginkgo/core/multigrid/rs.hpp @@ -6,11 +6,14 @@ #define GKO_PUBLIC_CORE_MULTIGRID_RS_HPP_ +#include + #include #include #include #include #include +#include #include #include #include @@ -117,8 +120,63 @@ class Rs : public LinOp, public EnableMultigridLevel { void generate(); + /** + * Generates the coarsening operators for a single, process-local matrix. + * + * @param local_matrix the local (diagonal) block to coarsen + * @param off_diag_matrix the off-diagonal block of a distributed matrix, + * or nullptr. It only enters the + * strength-of-connection threshold, see + * GKO_DECLARE_RS_COMPUTE_SOC_AND_RUN_RS_KERNEL. + * @param num_forced_c_points number of rows in `forced_c_points` + * @param forced_c_points local rows that must end up in the coarse set, + * regardless of what the greedy pass decided + * + * @return a tuple with prolongation, coarse, and restriction linop + */ + std::tuple, std::shared_ptr, + std::shared_ptr> + generate_local( + std::shared_ptr> local_matrix, + const matrix::Csr* off_diag_matrix = nullptr, + size_type num_forced_c_points = 0, + const IndexType* forced_c_points = nullptr); + +#if GINKGO_BUILD_MPI + /** + * Communicates the coarse index of every local row a neighboring rank + * couples to, in the coarse matrix' global indexing. + * + * All of those rows are forced C-points, so each of them has exactly one + * coarse index - which is what keeps the prolongation block-diagonal and + * this exchange down to a single index per halo entry. + * + * @tparam GlobalIndexType Global index type + * + * @param matrix a distributed matrix + * @param coarse_partition the coarse partition, used to compute the new + * global indices + * @param local_fine_to_coarse the local fine-to-coarse map + * + * @return the coarse global index of every off-diag column + */ + template + array communicate_off_diag_coarse_idxs( + std::shared_ptr> + matrix, + std::shared_ptr< + experimental::distributed::Partition> + coarse_partition, + const array& local_fine_to_coarse); +#endif + private: std::shared_ptr system_matrix_{}; + // the fine-to-coarse map of the last generate_local call. The distributed + // path needs it after generate_local returned, to tell the neighbors which + // coarse index their halo rows became. + array fine_to_coarse_{}; }; diff --git a/reference/multigrid/rs_kernels.cpp b/reference/multigrid/rs_kernels.cpp index 460be515f66..a8babb38ce5 100644 --- a/reference/multigrid/rs_kernels.cpp +++ b/reference/multigrid/rs_kernels.cpp @@ -74,7 +74,8 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( template void compute_soc_and_run_rs( std::shared_ptr exec, - matrix::view::csr A, double theta, + matrix::view::csr A, + matrix::view::csr off_diag, double theta, array& is_strong, array& lambda, array& cf_marker, IndexType& coarse_size) { @@ -83,6 +84,11 @@ void compute_soc_and_run_rs( const auto* a_row_ptrs = A.row_ptrs; const auto* a_col_idxs = A.col_idxs; const auto* a_vals = A.values; + // the off-diagonal block of a distributed matrix is stored on this rank + // and has the same rows as A, so it contributes to the strength threshold + // without any communication. It is empty for a non-distributed matrix. + const auto* od_row_ptrs = off_diag.row_ptrs; + const auto* od_vals = off_diag.values; bool* is_strong_vals = is_strong.get_data(); auto* lambda_vals = lambda.get_data(); auto* cf = cf_marker.get_data(); @@ -92,12 +98,19 @@ void compute_soc_and_run_rs( for (IndexType i = 0; i < n; ++i) { real_type max_offdiag = zero(); - // pass 1: find max off-diagonal + // pass 1a: find max off-diagonal for (IndexType jj = a_row_ptrs[i]; jj < a_row_ptrs[i + 1]; ++jj) { if (A.col_idxs[jj] != i) { max_offdiag = std::max(max_offdiag, -real(a_vals[jj])); } } + // pass 1b: the remote couplings of this row. None of them can be the + // diagonal entry, since they all live in other ranks' row ranges. + if (od_row_ptrs) { + for (IndexType jj = od_row_ptrs[i]; jj < od_row_ptrs[i + 1]; ++jj) { + max_offdiag = std::max(max_offdiag, -real(od_vals[jj])); + } + } // pass 2: set mask for (IndexType jj = a_row_ptrs[i]; jj < a_row_ptrs[i + 1]; ++jj) { @@ -141,6 +154,28 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE( GKO_DECLARE_RS_COMPUTE_SOC_AND_RUN_RS_KERNEL); +template +void mark_forced_c_points(std::shared_ptr exec, + size_type num_forced, const IndexType* forced_rows, + array& cf_marker, IndexType& coarse_size) +{ + auto* cf = cf_marker.get_data(); + for (size_type i = 0; i < num_forced; ++i) { + cf[forced_rows[i]] = 1; // C-point + } + + IndexType count = 0; + for (size_type i = 0; i < cf_marker.get_size(); ++i) { + if (cf[i] == 1) { + count++; + } + } + coarse_size = count; +} + +GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_RS_MARK_FORCED_C_POINTS_KERNEL); + + template void fill_coarse_and_compute_prolong_row_ptrs( std::shared_ptr exec, diff --git a/reference/test/multigrid/rs_kernels.cpp b/reference/test/multigrid/rs_kernels.cpp index c40c88cd09a..b39850dd346 100644 --- a/reference/test/multigrid/rs_kernels.cpp +++ b/reference/test/multigrid/rs_kernels.cpp @@ -13,6 +13,8 @@ #include #include +#include "core/multigrid/rs_helpers.hpp" + namespace { @@ -66,10 +68,13 @@ TEST_F(Rs, ComputeSocAndRunRs) gko::array lambda(exec, 5); gko::array cf(exec, 5); index_type coarse{}; + // this matrix is not distributed, so it has no off-diagonal block + const auto no_off_diag = + gko::multigrid::rs::no_off_diag_view(); gko::kernels::reference::rs::compute_soc_and_run_rs( - exec, A->get_const_device_view(), 0.5, is_strong_empty, lambda, cf, - coarse); + exec, A->get_const_device_view(), no_off_diag, 0.5, is_strong_empty, + lambda, cf, coarse); // all off-diagonals are strong std::vector expected_soc{false, true, true, false, true, true, false, diff --git a/test/mpi/multigrid/CMakeLists.txt b/test/mpi/multigrid/CMakeLists.txt index 81920f8f481..d2dc3c7f3d3 100644 --- a/test/mpi/multigrid/CMakeLists.txt +++ b/test/mpi/multigrid/CMakeLists.txt @@ -7,3 +7,13 @@ ginkgo_create_common_and_reference_test( LABELS distributed ) + +ginkgo_create_common_and_reference_test( + rs + MPI_SIZE + 3 + DISABLE_EXECUTORS + dpcpp + LABELS + distributed +) diff --git a/test/mpi/multigrid/rs.cpp b/test/mpi/multigrid/rs.cpp new file mode 100644 index 00000000000..dd5ac5804b4 --- /dev/null +++ b/test/mpi/multigrid/rs.cpp @@ -0,0 +1,201 @@ +// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors +// +// SPDX-License-Identifier: BSD-3-Clause + +#include + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/test/utils.hpp" +#include "test/utils/mpi/common_fixture.hpp" + + +template +class Rs : public CommonMpiTestFixture { +protected: + using value_type = typename std::tuple_element< + 0, decltype(ValueLocalGlobalIndexType())>::type; + using local_index_type = typename std::tuple_element< + 1, decltype(ValueLocalGlobalIndexType())>::type; + using global_index_type = typename std::tuple_element< + 2, decltype(ValueLocalGlobalIndexType())>::type; + using dist_mtx_type = + gko::experimental::distributed::Matrix; + using dist_vec_type = gko::experimental::distributed::Vector; + using local_matrix_type = gko::matrix::Csr; + using local_vec_type = gko::matrix::Dense; + using Partition = + gko::experimental::distributed::Partition; + using rs = gko::multigrid::Rs; + + Rs() : size{num_rows, num_rows} + { + // A 1D Laplacian is an M-matrix, and with four rows per rank every + // rank keeps interior rows: forcing the interface rows into the coarse + // set still leaves something to coarsen. + mat_input = gko::matrix_data{size}; + const auto n = static_cast(size[0]); + for (global_index_type i = 0; i < n; ++i) { + if (i > 0) { + mat_input.nonzeros.push_back( + {i, i - 1, -gko::one()}); + } + mat_input.nonzeros.push_back({i, i, value_type{2}}); + if (i < n - 1) { + mat_input.nonzeros.push_back( + {i, i + 1, -gko::one()}); + } + } + + row_part = Partition::build_from_contiguous( + exec, gko::array( + exec, I{0, 4, 8, 12})); + + dist_mat = dist_mtx_type::create(exec, comm); + dist_mat->read_distributed(mat_input, row_part); + } + + void SetUp() override { ASSERT_EQ(comm.size(), 3); } + + gko::size_type local_rows_of(std::shared_ptr distributed) + { + return gko::as( + gko::as(distributed)->get_diag_matrix()) + ->get_size()[0]; + } + + static constexpr gko::size_type num_rows = 12; + + gko::dim<2> size; + std::shared_ptr row_part; + gko::matrix_data mat_input; + std::shared_ptr dist_mat; +}; + +TYPED_TEST_SUITE(Rs, gko::test::ValueLocalGlobalIndexTypes, + TupleTypenameNameGenerator); + + +TYPED_TEST(Rs, CoarseOperatorIsTheGalerkinProduct) +{ + using rs = typename TestFixture::rs; + using value_type = typename TestFixture::value_type; + using dist_vec_type = typename TestFixture::dist_vec_type; + using local_vec_type = typename TestFixture::local_vec_type; + auto level = rs::build().on(this->exec)->generate(this->dist_mat); + auto coarse = level->get_coarse_op(); + auto prolong = level->get_prolong_op(); + auto restrict_op = level->get_restrict_op(); + const auto fine_global = this->dist_mat->get_size()[0]; + const auto fine_local = this->local_rows_of(this->dist_mat); + const auto coarse_local = this->local_rows_of(coarse); + + // A deterministic, rank-dependent coarse vector + auto host_x = + local_vec_type::create(this->ref, gko::dim<2>{coarse_local, 1}); + for (gko::size_type i = 0; i < coarse_local; ++i) { + host_x->at(i, 0) = static_cast( + 1 + (i + 3 * static_cast(this->comm.rank())) % 5); + } + auto x = dist_vec_type::create(this->exec, this->comm, + gko::clone(this->exec, host_x)); + const auto coarse_global = x->get_size()[0]; + auto direct = dist_vec_type::create(this->exec, this->comm, + gko::dim<2>{coarse_global, 1}, + gko::dim<2>{coarse_local, 1}); + auto px = dist_vec_type::create(this->exec, this->comm, + gko::dim<2>{fine_global, 1}, + gko::dim<2>{fine_local, 1}); + auto apx = dist_vec_type::create(this->exec, this->comm, + gko::dim<2>{fine_global, 1}, + gko::dim<2>{fine_local, 1}); + auto galerkin = dist_vec_type::create(this->exec, this->comm, + gko::dim<2>{coarse_global, 1}, + gko::dim<2>{coarse_local, 1}); + + // Ac * x, against R * (A * (P * x)). The right hand side pulls the + // neighbors' prolongated values through A's off-diagonal block, so it is + // the true Galerkin product. It agrees with the assembled coarse operator + // only if replacing the neighbors' prolongation rows by unit vectors was + // legitimate, i.e. only if every halo row really is a C-point. + coarse->apply(x, direct); + prolong->apply(x, px); + this->dist_mat->apply(px, apx); + restrict_op->apply(apx, galerkin); + + // the two sides sum the same terms in a different order + GKO_ASSERT_MTX_NEAR(direct->get_local_vector(), + galerkin->get_local_vector(), + 10 * r::value); +} + + +TYPED_TEST(Rs, InterfaceRowsAreForcedToCPoints) +{ + using rs = typename TestFixture::rs; + using value_type = typename TestFixture::value_type; + using dist_mtx_type = typename TestFixture::dist_mtx_type; + using local_matrix_type = typename TestFixture::local_matrix_type; + auto level = rs::build().on(this->exec)->generate(this->dist_mat); + + auto prolong = gko::as(level->get_prolong_op()); + auto p_local = gko::clone( + this->ref, gko::as(prolong->get_diag_matrix())); + auto off_diag = gko::clone( + this->ref, + gko::as(this->dist_mat->get_off_diag_matrix())); + const auto* od_row_ptrs = off_diag->get_const_row_ptrs(); + const auto* p_row_ptrs = p_local->get_const_row_ptrs(); + const auto* p_vals = p_local->get_const_values(); + + // The prolongation must have no off-diagonal block at all, otherwise it + // would not be representable as a purely local operator. + ASSERT_EQ(gko::as(prolong->get_off_diag_matrix()) + ->get_num_stored_elements(), + 0); + for (gko::size_type i = 0; i < p_local->get_size()[0]; ++i) { + if (od_row_ptrs[i + 1] > od_row_ptrs[i]) { + // This row couples to a remote row, and the test matrix has a + // symmetric pattern, so the remote rank couples back to it: it is + // one of our send indices and must have become a C-point, i.e. its + // prolongation row is a unit vector. + ASSERT_EQ(p_row_ptrs[i + 1] - p_row_ptrs[i], 1); + ASSERT_EQ(p_vals[p_row_ptrs[i]], gko::one()); + } + } +} + + +TYPED_TEST(Rs, GeneratesConsistentlySizedOperators) +{ + using rs = typename TestFixture::rs; + auto level = rs::build().on(this->exec)->generate(this->dist_mat); + auto coarse = level->get_coarse_op(); + auto prolong = level->get_prolong_op(); + auto restrict_op = level->get_restrict_op(); + const auto fine_global = this->dist_mat->get_size()[0]; + + const auto coarse_global = coarse->get_size()[0]; + ASSERT_EQ(coarse->get_size()[1], coarse_global); + ASSERT_EQ(prolong->get_size()[0], fine_global); + ASSERT_EQ(prolong->get_size()[1], coarse_global); + ASSERT_EQ(restrict_op->get_size()[0], coarse_global); + ASSERT_EQ(restrict_op->get_size()[1], fine_global); + // the interior of every rank is coarsened, so the coarse grid is smaller + ASSERT_LT(coarse_global, fine_global); + ASSERT_GT(coarse_global, 0); +} diff --git a/test/multigrid/rs_kernels.cpp b/test/multigrid/rs_kernels.cpp index 08248a6d570..07ff02af47e 100644 --- a/test/multigrid/rs_kernels.cpp +++ b/test/multigrid/rs_kernels.cpp @@ -12,6 +12,7 @@ #include #include +#include "core/multigrid/rs_helpers.hpp" #include "core/test/utils.hpp" #include "core/test/utils/matrix_generator.hpp" #include "core/utils/matrix_utils.hpp" @@ -79,6 +80,11 @@ class Rs : public CommonTestFixture { std::shared_ptr d_m_matrix; std::shared_ptr non_m_matrix; std::shared_ptr d_non_m_matrix; + + // the matrices here are not distributed, so they have no off-diagonal + // block + const Csr::const_device_view no_off_diag = + gko::multigrid::rs::no_off_diag_view(); }; @@ -124,12 +130,12 @@ TEST_F(Rs, ComputeSocAndRunRsIsEquivalentToRef) index_type coarse_size_exec = 0; gko::kernels::reference::rs::compute_soc_and_run_rs( - ref, m_matrix->get_const_device_view(), theta, is_strong_ref, - lambda_ref, cf_marker_ref, coarse_size_ref); + ref, m_matrix->get_const_device_view(), no_off_diag, theta, + is_strong_ref, lambda_ref, cf_marker_ref, coarse_size_ref); gko::kernels::GKO_DEVICE_NAMESPACE::rs::compute_soc_and_run_rs( - exec, d_m_matrix->get_const_device_view(), theta, is_strong_exec, - lambda_exec, cf_marker_exec, coarse_size_exec); + exec, d_m_matrix->get_const_device_view(), no_off_diag, theta, + is_strong_exec, lambda_exec, cf_marker_exec, coarse_size_exec); GKO_ASSERT_ARRAY_EQ(is_strong_ref, is_strong_exec); GKO_ASSERT_ARRAY_EQ(lambda_ref, lambda_exec); @@ -150,8 +156,8 @@ TEST_F(Rs, FillCoarseAndComputeProlongRowPtrsIsEquivalentToRef) index_type coarse_size = 0; gko::kernels::reference::rs::compute_soc_and_run_rs( - ref, m_matrix->get_const_device_view(), theta, is_strong_ref, - lambda_ref, cf_marker_ref, coarse_size); + ref, m_matrix->get_const_device_view(), no_off_diag, theta, + is_strong_ref, lambda_ref, cf_marker_ref, coarse_size); gko::array is_strong_exec(exec, is_strong_ref); gko::array cf_marker_exec(exec, cf_marker_ref); @@ -191,8 +197,8 @@ TEST_F(Rs, ComputeInterpolationIsEquivalentToRef) index_type coarse_size = 0; gko::kernels::reference::rs::compute_soc_and_run_rs( - ref, m_matrix->get_const_device_view(), theta, is_strong_ref, - lambda_ref, cf_marker_ref, coarse_size); + ref, m_matrix->get_const_device_view(), no_off_diag, theta, + is_strong_ref, lambda_ref, cf_marker_ref, coarse_size); gko::array coarse_rows_ref(ref, coarse_size); gko::array fine_to_coarse_ref(ref, num_rows);