From 41f807a9d4d6cf077ed01b2585a30c813ae244a1 Mon Sep 17 00:00:00 2001 From: Andrew Reisner Date: Tue, 3 Sep 2024 08:11:46 -0600 Subject: [PATCH] drop: adding rowmap --- raptor/multilevel/par_level.hpp | 1 + raptor/ruge_stuben/par_air_solver.hpp | 41 ++++++++++++- raptor/ruge_stuben/par_interpolation.cpp | 15 +++++ raptor/ruge_stuben/tests/test_air.cpp | 77 +++++++++++++++++++++++- raptor/ruge_stuben/tests/test_air.py | 28 +++++++++ 5 files changed, 160 insertions(+), 2 deletions(-) create mode 100644 raptor/ruge_stuben/tests/test_air.py diff --git a/raptor/multilevel/par_level.hpp b/raptor/multilevel/par_level.hpp index a116bb38..843bd429 100644 --- a/raptor/multilevel/par_level.hpp +++ b/raptor/multilevel/par_level.hpp @@ -34,6 +34,7 @@ namespace raptor ParCSRMatrix* A; ParCSRMatrix* P; + ParCSRMatrix* R; ParVector x; ParVector b; ParVector tmp; diff --git a/raptor/ruge_stuben/par_air_solver.hpp b/raptor/ruge_stuben/par_air_solver.hpp index b53d03f6..9f70c513 100644 --- a/raptor/ruge_stuben/par_air_solver.hpp +++ b/raptor/ruge_stuben/par_air_solver.hpp @@ -34,7 +34,7 @@ struct ParAIRSolver : ParMultilevel { ParCSRMatrix *A = levels[level_ctr]->A; ParBSRMatrix *A_bsr = dynamic_cast(A); if (A_bsr) extend_hier(*A_bsr); - // else extend_hier(*A); + else extend_hier(*A); } private: @@ -50,7 +50,46 @@ struct ParAIRSolver : ParMultilevel { auto P = one_point_interpolation(A, *S, split); auto R = local_air(A, *S, split, fpoint_distance::two); + levels.back()->P = P; + levels.back()->R = R; levels.emplace_back(new ParLevel()); + auto & level = *levels.back(); + auto AP = A.mult(P, tap_level); + auto coarse_A = R->mult(AP); + // auto coarse_A = AP->mult_T(P); + + level.A = coarse_A; + + finalize_coarse_op(tap_amg >= 0 && tap_amg <= level_ctr + 1); + init_coarse_vectors(); + + level.P = nullptr; + + delete AP; + delete S; + } + + void finalize_coarse_op(bool tap_init) { + auto & coarse_op = *levels.back()->A; + auto & fine_op = *(**(levels.end() - 2)).A; + + coarse_op.sort(); + coarse_op.on_proc->move_diag(); + coarse_op.comm = new ParComm(coarse_op.partition, coarse_op.off_proc_column_map, + coarse_op.on_proc_column_map, + fine_op.comm->key, fine_op.comm->mpi_comm); + + if (tap_init) { + coarse_op.init_tap_communicators(RAPtor_MPI_COMM_WORLD); + } + } + + void init_coarse_vectors() { + auto & clevel = *levels.back(); + [global_rows = clevel.A->global_num_rows, + local_rows = clevel.A->local_num_rows](auto & ... vecs) { + (vecs.resize(global_rows, local_rows),...); + }(clevel.x, clevel.b, clevel.tmp); } coarsen_t coarsen_type; diff --git a/raptor/ruge_stuben/par_interpolation.cpp b/raptor/ruge_stuben/par_interpolation.cpp index babadc23..0c0eea6e 100644 --- a/raptor/ruge_stuben/par_interpolation.cpp +++ b/raptor/ruge_stuben/par_interpolation.cpp @@ -2143,6 +2143,19 @@ T * create_R(const T & A, return global; }(local_rows); + + auto local_rowmap = [local_rows]() { + auto row_start = [](int local) { + MPI_Exscan(MPI_IN_PLACE, &local, 1, MPI_INT, MPI_SUM, RAPtor_MPI_COMM_WORLD); + int rank; MPI_Comm_rank(RAPtor_MPI_COMM_WORLD, &rank); + return (rank == 0) ? 0 : local; + }(local_rows); + + std::vector rowmap(local_rows); + std::iota(std::begin(rowmap), std::end(rowmap), row_start); + return rowmap; + }; + auto off_proc_num_cols = rowptr_and_colmap.offd_colmap.size(); auto move_data = [&](offd_map_rowptr && rac, ParMatrix & mat) { mat.on_proc->idx1 = std::move(rac.diag_rowptr); @@ -2163,11 +2176,13 @@ T * create_R(const T & A, local_rows, S.on_proc_num_cols, off_proc_num_cols, A.on_proc->b_rows, A.on_proc->b_cols); move_data(std::move(rowptr_and_colmap), *R); + R->local_row_map = local_rowmap(); return R; } else { auto * R = new ParCSRMatrix(S.partition, global_rows, S.global_num_cols, local_rows, S.on_proc_num_cols, off_proc_num_cols); move_data(std::move(rowptr_and_colmap), *R); + R->local_row_map = local_rowmap(); return R; } } diff --git a/raptor/ruge_stuben/tests/test_air.cpp b/raptor/ruge_stuben/tests/test_air.cpp index 498e1bf1..799cd6ab 100644 --- a/raptor/ruge_stuben/tests/test_air.cpp +++ b/raptor/ruge_stuben/tests/test_air.cpp @@ -1,6 +1,7 @@ #include "gtest/gtest.h" #include "raptor/raptor.hpp" +#include "raptor/ruge_stuben/par_air_solver.hpp" #include "raptor/tests/compare.hpp" using namespace raptor; @@ -14,6 +15,7 @@ int main(int argc, char** argv) return ret; } +#if 0 TEST(TestOnePointInterp, TestsInRuge_Stuben) { int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank); constexpr std::size_t n{16}; @@ -141,7 +143,6 @@ TEST(TestLocalAIR, TestsInRuge_Stuben) { } } -#if 0 TEST(TestLocalAIRBSR, TestsInRuge_Stuben) { int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank); constexpr std::size_t n{16}; @@ -198,3 +199,77 @@ TEST(TestLocalAIRBSR, TestsInRuge_Stuben) { auto R = local_air(*Absr, *S, split); } #endif + +void dump(const ParCSRMatrix & A) { + std::cout << A.global_num_cols << " " << A.global_num_rows << std::endl; + int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank); + std::ofstream ofile("coarse-" + std::to_string(rank)); + auto write_row = [&](int i, const Matrix & mat, const std::vector & colmap) { + const auto & rowmap = A.local_row_map; + for (int j = mat.idx1[i]; j < mat.idx1[i+1]; ++j) { + ofile << rowmap[i] << " " << colmap[mat.idx2[j]] << " " << std::scientific << mat.vals[j] << '\n'; + + } + }; + for (int i = 0; i < A.local_num_rows; ++i) { + write_row(i, *A.on_proc, A.on_proc_column_map); + write_row(i, *A.off_proc, A.off_proc_column_map); + } +} + +ParCSRMatrix * gen(std::size_t n) { + auto A = new ParCSRMatrix(n, n); + + int rank; MPI_Comm_rank(MPI_COMM_WORLD, &rank); + int size; MPI_Comm_size(MPI_COMM_WORLD, &size); + + auto num_local = A->partition->local_num_rows; + auto row_start = A->partition->first_local_row; + auto init = [&](Matrix & mat, int ncols) { + mat.n_rows = num_local; + mat.n_cols = ncols; + mat.nnz = 0; + mat.idx1.resize(num_local+1); + mat.idx2.reserve(num_local*2); + mat.vals.reserve(.3*num_local*2); + }; + init(*A->on_proc, num_local); + init(*A->off_proc, n); + + A->on_proc->idx1[0] = 0; + A->off_proc->idx1[0] = 0; + for (int i = 0; i < num_local; ++i) { + A->add_value(i, i + row_start, 1.); + if (!(rank == 0 && i == 0)) + A->add_value(i, i + row_start - 1, -1.); + + auto set_rowptr = [i](Matrix & mat) { + mat.idx1[i+1] = mat.idx2.size(); + }; + + set_rowptr(*A->on_proc); + set_rowptr(*A->off_proc); + } + + A->on_proc->nnz = A->on_proc->idx2.size(); + A->off_proc->nnz = A->off_proc->idx2.size(); + + A->finalize(); + + return A; +} +TEST(UpwindAdvection, TestsInRuge_Stuben) { + // constexpr std::size_t n{100}; + // std::vector grid; + // grid.resize(1, n); + // std::vector stencil{{-1, 1, 0}}; + // auto A = par_stencil_grid(stencil.data(), grid.data(), 1); + auto A = gen(100); + + // dump(*A); + ParAIRSolver ml; + // ParRugeStubenSolver ml; + ml.max_levels = 2; + ml.setup(A); + dump(*ml.levels[1]->A); +} diff --git a/raptor/ruge_stuben/tests/test_air.py b/raptor/ruge_stuben/tests/test_air.py new file mode 100644 index 00000000..91d44c81 --- /dev/null +++ b/raptor/ruge_stuben/tests/test_air.py @@ -0,0 +1,28 @@ +import numpy as np + +from numpy.testing import TestCase, assert_array_almost_equal +from scipy.sparse import bsr_matrix, diags, coo_matrix +from pyamg.classical.air import air_solver +from pyamg.gallery import poisson + +n = 100 +A = diags([np.ones((n,)), -1*np.ones((n-1,))], [0, -1]).tocsr() +f_relax = ('fc_jacobi', {'iterations': 1, 'f_iterations': 1, + 'c_iterations': 0}) +ml = air_solver(A, postsmoother=f_relax, max_levels=2) +# print(ml.levels[1].A) + +ra_coords = np.loadtxt('/home/andrew/src/raptor/build/raptor/ruge_stuben/tests/coarse') +I = np.array(ra_coords[:,0], dtype=int) +J = np.array(ra_coords[:,1], dtype=int) +# Jinds = np.unique(J) +# jmap = {} +# for i in range(Jinds.shape[0]): +# jmap[Jinds[i]] = i + +# for i in range(J.shape[0]): +# J[i] = jmap[J[i]] + +ra = coo_matrix((ra_coords[:,2], (I, J))) + +# print(np.max(ra - ml.levels[1].A))