From 2ac2c53941fa90903b7aa7895c2f6b14d436459f Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Thu, 11 Jul 2024 10:20:36 +0200
Subject: [PATCH 1/9] wip add multigrid, gko.1.8.0

---
 CMakeLists.txt                                |   4 +
 DevicePersistent/Base/Base.C                  |   6 +
 .../CsrMatrixWrapper/CsrMatrixWrapper.H       |  66 ++++--
 .../ExecutorHandler/ExecutorHandler.H         |   2 +-
 DevicePersistent/Partition/Partition.H        |   6 +-
 HostMatrix/HostMatrix.C                       |   7 +-
 HostMatrix/HostMatrix.H                       |  43 ++++
 Solver/Multigrid/GKOMultigrid.H               | 208 +++++++++---------
 lduLduBase/lduLduBase.H                       |   4 +-
 9 files changed, 212 insertions(+), 134 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 95d9626fb..87d799b15 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -158,6 +158,8 @@ target_sources(
   PRIVATE common/common.C
           ${CMAKE_CURRENT_BINARY_DIR}/version.C
           lduLduBase/lduLduBase.C
+          HostMatrix/HostMatrix.C
+          HostMatrix/HostMatrixFreeFunctions.C
           StoppingCriterion/StoppingCriterion.C
           DevicePersistent/Base/Base.C
           DevicePersistent/Partition/Partition.C
@@ -170,6 +172,8 @@ target_sources(
           BaseWrapper/CoupledLduBase/GKOCoupledLduBase.C
           Solver/CG/GKOCG.C
           Solver/BiCGStab/GKOBiCGStab.C
+          # Solver / IR / GKOIR.C
+          Solver/Multigrid/GKOMultigrid.C
           Solver/GMRES/GKOGMRES.C
   PUBLIC common/common.H
          StoppingCriterion/StoppingCriterion.H
diff --git a/DevicePersistent/Base/Base.C b/DevicePersistent/Base/Base.C
index abf416eef..9801a6621 100644
--- a/DevicePersistent/Base/Base.C
+++ b/DevicePersistent/Base/Base.C
@@ -29,4 +29,10 @@ defineTemplateTypeNameWithName(DevicePersistentBase<gko::LinOp>,
 typedef gko::experimental::distributed::Matrix<scalar, label, label> GkoMatrix;
 defineTemplateTypeNameWithName(DevicePersistentBase<GkoMatrix>,
                                "PersistentMatrix");
+
+// typedef needed  to avoid confusion with the comma separated template
+// arguments as macro arguments
+typedef gko::experimental::distributed::Partition<label,label> Partition;
+defineTemplateTypeNameWithName(DevicePersistentBase<Partition>,
+                               "PersistentPartition");
 }  // namespace Foam
diff --git a/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H b/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H
index 8962e1b39..539a00810 100644
--- a/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H
+++ b/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H
@@ -38,6 +38,8 @@ struct MatrixInitFunctor {
 
     const PersistentArray<scalar> &non_local_coeffs_;
 
+    const CommunicationPattern &communication_pattern_;
+
     const word matrix_format_;
 
     const bool regenerate_;
@@ -54,6 +56,7 @@ struct MatrixInitFunctor {
                       const PersistentArray<label> &non_local_col_idxs,
                       const PersistentArray<label> &non_local_row_idxs,
                       const PersistentArray<scalar> &non_local_coeffs,
+                      const CommunicationPattern &communication_pattern,
                       const word matrix_format, const bool regenerate,
                       const label verbose, const word field_name)
         : db_(db),
@@ -65,6 +68,7 @@ struct MatrixInitFunctor {
           non_local_col_idxs_(non_local_col_idxs),
           non_local_row_idxs_(non_local_row_idxs),
           non_local_coeffs_(non_local_coeffs),
+          communication_pattern_(communication_pattern),
           matrix_format_(matrix_format),
           regenerate_(regenerate),
           verbose_(verbose),
@@ -137,27 +141,44 @@ struct MatrixInitFunctor {
 
     std::shared_ptr<dist_mtx> generate_dist_mtx_with_inner_type(
         std::shared_ptr<gko::Executor> exec,
-        std::shared_ptr<gko::experimental::mpi::communicator> comm) const
+        std::shared_ptr<gko::experimental::mpi::communicator> comm,
+        gko::dim<2> global_dim,
+        gko::device_matrix_data<scalar, label> &local_A_data,
+        gko::device_matrix_data<scalar, label> &non_local_A_data) const
     {
+        auto [recv_sizes, recv_offsets] =
+            communication_pattern_.recv_sizes_offsets(comm);
+        auto recv_gather_idxs =
+            communication_pattern_.total_rank_send_idx(exec->get_master());
+
         if (matrix_format_ == "Csr") {
-            return dist_mtx::create(exec, *comm.get(),
-                                    gko::with_matrix_type<gko::matrix::Csr>());
-        }
-        if (matrix_format_ == "Ell") {
-            return dist_mtx::create(exec, *comm.get(),
-                                    gko::with_matrix_type<gko::matrix::Ell>());
+            auto local_linop =
+                gko::share(gko::matrix::Csr<scalar, label>::create(exec));
+            local_linop->read(local_A_data);
+            auto non_local_linop =
+                gko::share(gko::matrix::Csr<scalar, label>::create(exec));
+            non_local_linop->read(non_local_A_data);
+            return dist_mtx::create(exec, *comm.get(), global_dim, local_linop,
+                                    non_local_linop, recv_sizes, recv_offsets,
+                                    recv_gather_idxs);
         }
-        // if (matrix_format_ == "Hybrid") {
-        //     return dist_mtx::create(
-        //         *comm.get(), gko::with_matrix_type<gko::matrix::Hybrid>());
-        // }
         if (matrix_format_ == "Coo") {
-            return dist_mtx::create(exec, *comm.get(),
-                                    gko::with_matrix_type<gko::matrix::Coo>());
+            auto local_linop =
+                gko::share(gko::matrix::Coo<scalar, label>::create(exec));
+            local_linop->read(local_A_data);
+            auto non_local_linop =
+                gko::share(gko::matrix::Coo<scalar, label>::create(exec));
+            non_local_linop->read(non_local_A_data);
+
+            auto dist = dist_mtx::create(
+                exec, *comm.get(), global_dim, local_linop, non_local_linop,
+                recv_sizes, recv_offsets, recv_gather_idxs);
+            return dist;
         }
 
-        FatalErrorInFunction << "Matrix format " << matrix_format_
-                             << " not supported " << abort(FatalError);
+
+        // FatalErrorInFunction << "Matrix format " << matrix_format_
+        //                      << " not supported " << abort(FatalError);
     }
 
     std::shared_ptr<dist_mtx> init() const
@@ -178,6 +199,7 @@ struct MatrixInitFunctor {
         auto exec = exec_.get_ref_exec();
 
         auto num_rows = partition_.get_local_size();
+        auto global_num_rows = partition_.get_total_size();
         gko::device_matrix_data<scalar, label> A_data(
             exec, gko::dim<2>(num_rows, num_rows), *rows.get(), *cols.get(),
             *coeffs.get());
@@ -187,15 +209,12 @@ struct MatrixInitFunctor {
             *non_local_rows.get(), *non_local_cols.get(),
             *non_local_coeffs.get());
 
-
         auto comm = exec_.get_gko_mpi_host_comm();
-        auto dist_A = generate_dist_mtx_with_inner_type(
-            exec_.get_ref_exec(), exec_.get_gko_mpi_host_comm());
 
-        auto device_mat = generate_dist_mtx_with_inner_type(
-            exec_.get_device_exec(), exec_.get_gko_mpi_device_comm());
-        dist_A->move_to(device_mat.get());
-        return device_mat;
+        return generate_dist_mtx_with_inner_type(
+            exec_.get_device_exec(), exec_.get_gko_mpi_host_comm(),
+            gko::dim<2>{global_num_rows, global_num_rows}, A_data,
+            non_local_A_data);
     }
 };
 
@@ -228,6 +247,7 @@ public:
                   const PersistentArray<label> &non_local_col_idxs,
                   const PersistentArray<label> &non_local_row_idxs,
                   const PersistentArray<scalar> &non_local_coeffs,
+                  const CommunicationPattern &communication_pattern,
                   const PersistentPartition &partition,
                   const dictionary &controlDict, const word sys_matrix_name,
                   const label verbose)
@@ -242,7 +262,7 @@ public:
               MatrixInitFunctor(
                   db, exec, partition, col_idxs, row_idxs, coeffs,
                   non_local_col_idxs, non_local_row_idxs, non_local_coeffs,
-                  matrix_format_,
+                  communication_pattern, matrix_format_,
                   controlDict.lookupOrDefault<Switch>("regenerate", false),
                   verbose_, sys_matrix_name),
               controlDict.lookupOrDefault<Switch>("updateSysMatrix", true),
diff --git a/DevicePersistent/ExecutorHandler/ExecutorHandler.H b/DevicePersistent/ExecutorHandler/ExecutorHandler.H
index 90df64a08..c1a6be41c 100644
--- a/DevicePersistent/ExecutorHandler/ExecutorHandler.H
+++ b/DevicePersistent/ExecutorHandler/ExecutorHandler.H
@@ -56,7 +56,7 @@ struct ExecutorInitFunctor {
             return gko::share(gko::CudaExecutor::create(
                 device_id_ % gko::CudaExecutor::get_num_devices(), host_exec));
         }
-        if (executor_name_ == "sycl" || "dpcpp") {
+        if (executor_name_ == "sycl" || executor_name_ == "dpcpp") {
             if (version.dpcpp_version.tag == not_compiled_tag) {
                 FatalErrorInFunction
                     << "SYCL Backend was not compiled. Recompile OGL/Ginkgo "
diff --git a/DevicePersistent/Partition/Partition.H b/DevicePersistent/Partition/Partition.H
index 9f6d08e77..fdf8c5876 100644
--- a/DevicePersistent/Partition/Partition.H
+++ b/DevicePersistent/Partition/Partition.H
@@ -62,7 +62,11 @@ struct PartitionInitFunctor {
                  std::to_string(local_size_)};
         LOG_1(verbose_, msg)
 
-        if (ranks_per_gpu_ == 1) return {};
+
+        auto host_partition = gko::share(gko::experimental::distributed::build_partition_from_local_size<label,label>(
+            exec, *comm.get(), local_size_));
+
+        return host_partition;
     }
 };
 
diff --git a/HostMatrix/HostMatrix.C b/HostMatrix/HostMatrix.C
index 5c003a797..74657a3df 100644
--- a/HostMatrix/HostMatrix.C
+++ b/HostMatrix/HostMatrix.C
@@ -283,10 +283,9 @@ HostMatrixWrapper<MatrixType>::create_communication_pattern(
     std::vector<std::pair<gko::array<label>, comm_size_type>> send_idxs;
     for (auto [proc, interface_cells] : interface_cell_map) {
         auto exec = exec_.get_ref_exec();
+        auto arr = gko::array<label>::const_view(exec, interface_cells.size(), interface_cells.data());
         send_idxs.push_back(std::pair<gko::array<label>, comm_size_type>(
-            gko::array<label>(exec, interface_cells.begin(),
-                              interface_cells.end()),
-            proc));
+            arr.copy_to_array(), proc));
     }
 
     // convert to gko::array
@@ -459,7 +458,7 @@ void HostMatrixWrapper<MatrixType>::init_non_local_sparsity_pattern(
     label element_ctr = 0;
     for (auto [interface_idx, row] : non_local_row_indices) {
         rows[element_ctr] = row;
-        cols[element_ctr] = interface_idx;
+        cols[element_ctr] =  interface_idx;
         permute[element_ctr] = interface_idx;
         element_ctr += 1;
     }
diff --git a/HostMatrix/HostMatrix.H b/HostMatrix/HostMatrix.H
index aaf803eab..aebf37b8d 100644
--- a/HostMatrix/HostMatrix.H
+++ b/HostMatrix/HostMatrix.H
@@ -76,6 +76,49 @@ struct CommunicationPattern {
     // send_idx stores the index_set ie which cells
     // are owned by the interface and the corresponding target rank
     std::vector<std::pair<gko::array<label>, comm_size_type>> send_idxs;
+
+    gko::array<label> total_rank_send_idx(std::shared_ptr<gko::Executor> exec) const
+    {
+        std::vector<label> tmp {};
+
+        for (auto [arr, id] : send_idxs) {
+            label arr_size = arr.get_size();
+            tmp.insert(tmp.end(), arr.get_const_data(),
+                       arr.get_const_data() + arr_size);
+        }
+
+        auto ret = gko::array<label>::const_view(exec, 
+                tmp.size(),
+                tmp.data()
+                );
+
+        return ret.copy_to_array();
+    }
+
+    std::pair<std::vector<label>, std::vector<label>> recv_sizes_offsets(
+        std::shared_ptr<gko::experimental::mpi::communicator> comm
+            ) const
+    {
+        std::vector<label> recv_counts(comm->size());
+        std::vector<label> recv_offsets(comm->size() + 1);
+
+        label comm_ranks = target_ids.get_size();
+        label tot_comm_size = 0;
+        for (label i = 0; i < comm_ranks; i++) {
+            auto comm_rank = target_ids.get_const_data()[i];
+            auto comm_size = target_sizes.get_const_data()[i];
+            tot_comm_size += comm_size;
+            recv_counts[comm_rank] = comm_size;
+        }
+
+        recv_offsets[comm->size()] = tot_comm_size;
+        std::partial_sum(recv_counts.begin(), recv_counts.end(),
+                         recv_offsets.begin() + 1);
+        recv_offsets[0] = 0;
+
+        return {recv_counts, recv_offsets};
+    }
+
 };
 
 
diff --git a/Solver/Multigrid/GKOMultigrid.H b/Solver/Multigrid/GKOMultigrid.H
index f01135a4d..90b93a858 100644
--- a/Solver/Multigrid/GKOMultigrid.H
+++ b/Solver/Multigrid/GKOMultigrid.H
@@ -15,12 +15,12 @@ class GKOMultigridFactory {
 private:
     using cg = gko::solver::Cg<scalar>;
     using mtx = gko::matrix::Csr<scalar>;
-    using vec = gko::matrix::Dense<scalar>;
     using ir = gko::solver::Ir<scalar>;
     using mg = gko::solver::Multigrid;
     using bj = gko::preconditioner::Jacobi<scalar, label>;
     using amgx_pgm = gko::multigrid::AmgxPgm<scalar, label>;
 
+    using vec = gko::matrix::Dense<scalar>;
     using dist_vec = gko::experimental::distributed::Vector<scalar>;
     using dist_mtx =
         gko::experimental::distributed::Matrix<scalar, label, label>;
@@ -87,114 +87,101 @@ public:
     {}
 
     std::unique_ptr<mg::Factory, std::default_delete<mg::Factory>>
-    create_dist_solver(std::shared_ptr<gko::Executor> exec,
-                       std::shared_ptr<gko::LinOp> gkomatrix,
-                       std::shared_ptr<dist_vec> x, std::shared_ptr<dist_vec> b,
+    create_dist_solver(std::shared_ptr<const gko::Executor> exec,
+                       std::shared_ptr<const gko::LinOp> gkomatrix,
+                       std::shared_ptr<const dist_vec> x, std::shared_ptr<const dist_vec> b,
                        const label verbose, const bool export_res,
                        std::shared_ptr<gko::LinOp> precond) const
     {
-        outerStoppingCriterionVec_.push_back(
-            outerStoppingCriterion_.build_dist_stopping_criterion(
-                exec, gkomatrix, x, b, verbose, export_res,
-                get_prev_number_of_iterations()));
-
-        auto inner_solver_gen =
-            gko::share(bj::build()
-                           .with_skip_sorting(true)
-                           .with_max_block_size(max_block_size_)
-                           .on(exec));
-
-        auto smoother_gen = gko::share(
-            ir::build()
-                .with_solver(inner_solver_gen)
-                .with_relaxation_factor(smoother_relaxation_factor_)
-                .with_criteria(gko::stop::Iteration::build()
-                                   .with_max_iters(smoother_max_iters_)
-                                   .on(exec))
-                .on(exec));
-        // Create MultigridLevel factory
-        auto mg_level_gen = amgx_pgm::build()
-                                .with_deterministic(true)
-                                .with_skip_sorting(true)
-                                .on(exec);
-        // Create CoarsestSolver factory
-        auto coarsest_gen =
-            gko::share(ir::build()
-                           .with_solver(inner_solver_gen)
-                           .with_relaxation_factor(inner_relaxation_factor_)
-                           .with_criteria(gko::stop::Iteration::build()
-                                              .with_max_iters(coarse_max_iters_)
-                                              .on(exec))
-                           .on(exec));
-
-        // Create multigrid factory
-        auto ret = mg::build()
-                       .with_max_levels(max_levels_)
-                       .with_min_coarse_rows(min_coarse_rows_)
-                       .with_pre_smoother(smoother_gen)
-                       .with_post_uses_pre(true)
-                       .with_mg_level(gko::share(mg_level_gen))
-                       .with_coarsest_solver(coarsest_gen)
-                       .with_zero_guess(false)
-                       .with_criteria(outerStoppingCriterionVec_)
-                       .on(exec);
-
+        std::cout << __FILE__ << ":" << __LINE__ << "create mg solver\n";
+        const scalar reduction_factor{1e-8};
+        label num_iters {100};
+        auto ret = gko::solver::Multigrid::build()
+                    .with_mg_level(    
+                        gko::multigrid::Pgm<scalar>::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(
+                            num_iters),
+                        gko::stop::ResidualNorm<scalar>::build()
+                            .with_baseline(gko::stop::mode::absolute)
+                            .with_reduction_factor(reduction_factor))
+                    .on(exec);
         return ret;
     }
 
     // Create solver
     std::unique_ptr<mg::Factory, std::default_delete<mg::Factory>>
-    create_solver(std::shared_ptr<gko::Executor> exec,
-                  std::shared_ptr<mtx> gkomatrix, std::shared_ptr<vec> x,
-                  std::shared_ptr<vec> b, const label verbose,
+    create_solver(std::shared_ptr<const gko::Executor> exec,
+                  std::shared_ptr<dist_mtx> gkomatrix, std::shared_ptr<dist_vec> x,
+                  std::shared_ptr<dist_vec> b, const label verbose,
                   const bool export_res,
                   std::shared_ptr<gko::LinOp> precond) const
     {
-        outerStoppingCriterionVec_.push_back(
-            outerStoppingCriterion_.build_stopping_criterion(
-                exec, gkomatrix, x, b, verbose, export_res, 0));
-
-        auto inner_solver_gen =
-            gko::share(bj::build()
-                           .with_skip_sorting(true)
-                           .with_max_block_size(max_block_size_)
-                           .on(exec));
-
-        auto smoother_gen = gko::share(
-            ir::build()
-                .with_solver(inner_solver_gen)
-                .with_relaxation_factor(smoother_relaxation_factor_)
-                .with_criteria(gko::stop::Iteration::build()
-                                   .with_max_iters(smoother_max_iters_)
-                                   .on(exec))
-                .on(exec));
-        // Create MultigridLevel factory
-        auto mg_level_gen = amgx_pgm::build()
-                                .with_deterministic(true)
-                                .with_skip_sorting(true)
-                                .on(exec);
-        // Create CoarsestSolver factory
-        auto coarsest_gen =
-            gko::share(ir::build()
-                           .with_solver(inner_solver_gen)
-                           .with_relaxation_factor(inner_relaxation_factor_)
-                           .with_criteria(gko::stop::Iteration::build()
-                                              .with_max_iters(coarse_max_iters_)
-                                              .on(exec))
-                           .on(exec));
-
-        // Create multigrid factory
-        auto ret = mg::build()
-                       .with_max_levels(max_levels_)
-                       .with_min_coarse_rows(min_coarse_rows_)
-                       .with_pre_smoother(smoother_gen)
-                       .with_post_uses_pre(true)
-                       .with_mg_level(gko::share(mg_level_gen))
-                       .with_coarsest_solver(coarsest_gen)
-                       .with_zero_guess(false)
-                       .with_criteria(outerStoppingCriterionVec_)
-                       .on(exec);
-
+        // outerStoppingCriterionVec_.push_back(
+        //     outerStoppingCriterion_.build_dist_stopping_criterion(
+        //         exec, gkomatrix, x, b, verbose, export_res, 0, 0));
+        //
+        // auto inner_solver_gen =
+        //     gko::share(bj::build()
+        //                    .with_skip_sorting(true)
+        //                    .with_max_block_size(max_block_size_)
+        //                    .on(exec));
+        //
+        // auto smoother_gen = gko::share(
+        //     ir::build()
+        //         .with_solver(inner_solver_gen)
+        //         .with_relaxation_factor(smoother_relaxation_factor_)
+        //         .with_criteria(gko::stop::Iteration::build()
+        //                            .with_max_iters(smoother_max_iters_)
+        //                            .on(exec))
+        //         .on(exec));
+        // // Create MultigridLevel factory
+        // auto mg_level_gen = amgx_pgm::build()
+        //                         .with_deterministic(true)
+        //                         .with_skip_sorting(true)
+        //                         .on(exec);
+        // // Create CoarsestSolver factory
+        // auto coarsest_gen =
+        //     gko::share(ir::build()
+        //                    .with_solver(inner_solver_gen)
+        //                    .with_relaxation_factor(inner_relaxation_factor_)
+        //                    .with_criteria(gko::stop::Iteration::build()
+        //                                       .with_max_iters(coarse_max_iters_)
+        //                                       .on(exec))
+        //                    .on(exec));
+        //
+        // // Create multigrid factory
+        // auto ret = mg::build()
+        //                .with_max_levels(max_levels_)
+        //                .with_min_coarse_rows(min_coarse_rows_)
+        //                .with_pre_smoother(smoother_gen)
+        //                .with_post_uses_pre(true)
+        //                .with_mg_level(gko::share(mg_level_gen))
+        //                .with_coarsest_solver(coarsest_gen)
+        //                .with_zero_guess(false)
+        //                .with_criteria(outerStoppingCriterionVec_)
+        //                .on(exec);
+        std::cout << __FILE__ << ":" << __LINE__ << "create mg solver\n";
+        const scalar reduction_factor{1e-8};
+        label num_iters {100};
+        auto ret = gko::solver::Multigrid::build()
+                    .with_mg_level(    
+                        gko::multigrid::Pgm<scalar>::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(
+                            num_iters),
+                        gko::stop::ResidualNorm<scalar>::build()
+                            .with_baseline(gko::stop::mode::absolute)
+                            .with_reduction_factor(reduction_factor))
+                    .on(exec);
         return ret;
     }
 
@@ -213,19 +200,34 @@ public:
         return outerStoppingCriterion_.get_res_norms();
     }
 
+    void set_prev_rel_res_cost(scalar prev_rel_res_cost) const
+    {
+        return ::Foam::set_solve_prev_rel_res_cost(sysMatrixName_, db_,
+                                                   prev_rel_res_cost);
+    }
+
     void store_number_of_iterations() const
     {
-        set_solve_prev_iters(sysMatrixName_, db_,
-                             stoppingCriterion_.get_num_iters(),
-                             stoppingCriterion_.get_is_final());
+        // set_solve_prev_iters(sysMatrixName_, db_,
+        //                      stoppingCriterion_.get_num_iters(),
+        //                      stoppingCriterion_.get_is_final());
     }
 
     label get_prev_number_of_iterations() const
     {
-        return get_solve_prev_iters(sysMatrixName_, db_,
-                                    stoppingCriterion_.get_is_final());
+        // return get_solve_prev_iters(sysMatrixName_, db_,
+        //                             stoppingCriterion_.get_is_final());
+        return 0;
     }
 
+    label get_res_norm_time() const
+    {
+        // return get_solve_prev_iters(sysMatrixName_, db_,
+        //                             stoppingCriterion_.get_is_final());
+        return 0;
+    }
+
+
     label get_number_of_iterations() const
     {
         return outerStoppingCriterion_.get_num_iters();
diff --git a/lduLduBase/lduLduBase.H b/lduLduBase/lduLduBase.H
index cf77709fc..225fe73c5 100644
--- a/lduLduBase/lduLduBase.H
+++ b/lduLduBase/lduLduBase.H
@@ -200,7 +200,6 @@ public:
             psi.size(),  // number of elements
             this->get_communication_pattern()};
 
-        // TODO this is now the local matrix part for dist_A
         MatrixWrapper dist_A{db_,
                              this->get_exec_handler(),  // get exec
                              this->get_local_col_idxs(),
@@ -209,6 +208,7 @@ public:
                              this->get_non_local_col_idxs(),
                              this->get_non_local_row_idxs(),
                              this->get_non_local_coeffs(),
+                             this->get_communication_pattern(),
                              partition,
                              solver_controls_,
                              this->fieldName(),
@@ -271,7 +271,7 @@ public:
 
         TIME_WITH_FIELDNAME(verbose_, generate_solver, this->fieldName(),
                             auto solver = solver_gen->generate(dist_A_v);)
-        LOG_1(verbose_, "solve solver done")
+        LOG_1(verbose_, "create solver done")
         TIME_WITH_FIELDNAME(verbose_, solve, this->fieldName(),
                             solver->apply(dist_b_v, dist_x_v);)
 

From 3b4c968fb608dc49cc297a2621ed61180d0a1207 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Tue, 23 Jul 2024 10:06:47 +0200
Subject: [PATCH 2/9] wip add multigrid, gko.1.8.0

---
 DevicePersistent/Array/Array.H                |  51 +++----
 .../CsrMatrixWrapper/CsrMatrixWrapper.H       |   3 +-
 DevicePersistent/Partition/Partition.H        |   5 +-
 DevicePersistent/Vector/Vector.H              |  17 +--
 Solver/Multigrid/GKOMultigrid.H               | 125 +++++-------------
 StoppingCriterion/StoppingCriterion.C         |  22 ++-
 StoppingCriterion/StoppingCriterion.H         |  18 +--
 lduLduBase/lduLduBase.H                       |  48 ++++---
 8 files changed, 120 insertions(+), 169 deletions(-)

diff --git a/DevicePersistent/Array/Array.H b/DevicePersistent/Array/Array.H
index 2981ad885..fd701f346 100644
--- a/DevicePersistent/Array/Array.H
+++ b/DevicePersistent/Array/Array.H
@@ -97,8 +97,6 @@ class PersistentArray
 
     mutable label size_;
 
-    const ExecutorHandler &exec_;
-
     // indicating if the underlying array needs to
     // updated even if was found in the object registry
     const bool update_;
@@ -133,7 +131,6 @@ public:
               ArrayInitFunctor<T>(exec, name, size, verbose, on_device), update,
               verbose),
           size_(size),
-          exec_(exec),
           update_(update)
     {}
 
@@ -160,7 +157,6 @@ public:
               update, verbose),
           memory_(memory),
           size_(size),
-          exec_(exec),
           update_(update)
     {}
 
@@ -183,16 +179,16 @@ public:
 
     /* Create a Dense<T> vector from underlying data
      * */
-    std::shared_ptr<gko::matrix::Dense<T>> get_dense_vec() const
-    {
-        auto array = get_array();
-        auto result = gko::share(vec::create(
-            exec_.get_device_exec(), gko::dim<2>(array->get_num_elems(), 1),
-            gko::array<T>::view(exec_.get_device_exec(), array->get_num_elems(),
-                                get_data()),
-            1));
-        return result;
-    }
+    // std::shared_ptr<gko::matrix::Dense<T>> get_dense_vec() const
+    // {
+    //     auto array = get_array();
+    //     auto result = gko::share(vec::create(
+    //         exec_.get_device_exec(), gko::dim<2>(array->get_num_elems(), 1),
+    //         gko::array<T>::view(exec_.get_device_exec(), array->get_num_elems(),
+    //                             get_data()),
+    //         1));
+    //     return result;
+    // }
 
     /* Copies from a distributed::Vector<T> back into raw host memory
      *
@@ -201,31 +197,28 @@ public:
     void copy_back(
         const std::shared_ptr<gko::experimental::distributed::Vector<T>> dist_x)
     {
-        auto host_view = gko::array<T>::view(exec_.get_ref_exec(), size_,
+        auto dist_device_exec = dist_x->get_executor();
+        auto host_view = gko::array<T>::view(dist_device_exec->get_master(), size_,
                                              const_cast<T *>(memory_));
 
-        auto dist_device_exec = dist_x->get_executor();
         auto x_view = gko::array<T>::view(dist_device_exec, size_,
                                           dist_x->get_local()->get_values());
 
         host_view = x_view;
     }
 
-    void copy_back(const std::shared_ptr<gko::matrix::Dense<T>> device_x)
-    {
-        auto host_view = gko::array<T>::view(exec_.get_ref_exec(), size_,
-                                             const_cast<T *>(memory_));
-
-        auto x_view = gko::array<T>::view(exec_.get_device_exec(), size_,
-                                          device_x->get_values());
-
-        host_view = x_view;
-    }
+    // void copy_back(const std::shared_ptr<gko::matrix::Dense<T>> device_x)
+    // {
+    //     auto host_view = gko::array<T>::view(exec_.get_ref_exec(), size_,
+    //                                          const_cast<T *>(memory_));
+    //
+    //     auto x_view = gko::array<T>::view(exec_.get_device_exec(), size_,
+    //                                       device_x->get_values());
+    //
+    //     host_view = x_view;
+    // }
 
     void set_size(const label size) const { size_ = size; }
-
-
-    const ExecutorHandler &get_exec_handler() const { return exec_; }
 };
 
 }  // namespace Foam
diff --git a/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H b/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H
index 539a00810..a2f46b706 100644
--- a/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H
+++ b/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H
@@ -80,6 +80,8 @@ struct MatrixInitFunctor {
         if (regenerate_) {
             TIME_WITH_FIELDNAME(verbose_, regenerate_distributed_matrix,
                                 field_name_, auto reinit = init();)
+            std::cout << __FILE__ << ":" << __LINE__ << " Reinit matrix \n";
+            // TODO NOTE use swap here
             persistent_device_matrix->copy_from(reinit.get());
             return;
         }
@@ -238,7 +240,6 @@ private:
 
     mutable label prev_solve_iters_ = 0;
 
-
 public:
     MatrixWrapper(const objectRegistry &db, const PersistentExecutor &exec,
                   const PersistentArray<label> &col_idxs,
diff --git a/DevicePersistent/Partition/Partition.H b/DevicePersistent/Partition/Partition.H
index fdf8c5876..6ac480416 100644
--- a/DevicePersistent/Partition/Partition.H
+++ b/DevicePersistent/Partition/Partition.H
@@ -88,8 +88,6 @@ class PersistentPartition
 
     mutable label global_elements_;
 
-    const ExecutorHandler &exec_;
-
 public:
     using part_type = gko::experimental::distributed::Partition<label, label>;
     /* PersistentPartition constructor using existing memory
@@ -111,8 +109,7 @@ public:
                                    communication_pattern),
               false, verbose),
           ranks_per_gpu_(ranks_per_gpu),
-          local_elements_(elements),
-          exec_(exec)
+          local_elements_(elements)
     {
         auto comm = exec.get_gko_mpi_host_comm();
         label local_elements = local_elements_;
diff --git a/DevicePersistent/Vector/Vector.H b/DevicePersistent/Vector/Vector.H
index 211a408eb..592370af7 100644
--- a/DevicePersistent/Vector/Vector.H
+++ b/DevicePersistent/Vector/Vector.H
@@ -99,8 +99,6 @@ class PersistentVector
 
     const PersistentPartition partition_;
 
-    const ExecutorHandler &exec_;
-
     // indicating if the underlying array needs to
     // updated even if was found in the object registry
     const bool update_;
@@ -135,33 +133,32 @@ public:
           name_(name),
           memory_(memory),
           partition_(partition),
-          exec_(exec),
           update_(update)
     {}
 
     /** Copies the content of the distributed vector back to the original source
      **/
-    void copy_back()
+    void copy_back(std::shared_ptr<gko::Executor> exec, gko::experimental::mpi::communicator& comm)
     {
+        auto master = exec->get_master();
         const auto local_size = partition_.get_local_size();
         auto global_size = partition_.get_total_size();
 
-        auto comm = exec_.get_gko_mpi_host_comm();
         auto host_device_vector =
-            dist_vec::create(exec_.get_ref_exec(), *comm.get());
+            dist_vec::create(master, comm);
         host_device_vector->operator=(*get_vector().get());
 
         // TODO why does it need its size
         // repartitioner size should be sufficient
-        auto target_vector = dist_vec::create(exec_.get_ref_exec(), *comm.get(),
+        auto target_vector = dist_vec::create(master, comm,
                                               gko::dim<2>{global_size, 1},
                                               gko::dim<2>{local_size, 1}, 1);
         scalar *host_buffer = host_device_vector->get_local_values();
 
         auto host_buffer_view =
-            gko::array<T>::view(exec_.get_ref_exec(), local_size, host_buffer);
+            gko::array<T>::view(master, local_size, host_buffer);
 
-        auto to_view = gko::array<T>::view(exec_.get_ref_exec(), local_size,
+        auto to_view = gko::array<T>::view(master, local_size,
                                            const_cast<T *>(memory_));
         to_view = host_buffer_view;
     }
@@ -188,8 +185,6 @@ public:
         return this->get_persistent_object()->get_const_data();
     }
 
-    const ExecutorHandler &get_exec_handler() const { return exec_; }
-
     std::shared_ptr<gko::experimental::distributed::Vector<T>> get_vector()
         const
     {
diff --git a/Solver/Multigrid/GKOMultigrid.H b/Solver/Multigrid/GKOMultigrid.H
index 90b93a858..0cb5ac9eb 100644
--- a/Solver/Multigrid/GKOMultigrid.H
+++ b/Solver/Multigrid/GKOMultigrid.H
@@ -54,10 +54,10 @@ private:
     const label min_coarse_rows_;
 
     mutable std::vector<std::shared_ptr<const gko::stop::CriterionFactory>>
-        outerStoppingCriterionVec_ = {};
+        outerStoppingCriterionVec_;
 
     mutable std::vector<std::shared_ptr<const gko::stop::CriterionFactory>>
-        innerStoppingCriterionVec_ = {};
+        innerStoppingCriterionVec_;
 
 public:
     GKOMultigridFactory(const dictionary &solverControls,
@@ -93,98 +93,43 @@ public:
                        const label verbose, const bool export_res,
                        std::shared_ptr<gko::LinOp> precond) const
     {
-        std::cout << __FILE__ << ":" << __LINE__ << "create mg solver\n";
-        const scalar reduction_factor{1e-8};
-        label num_iters {100};
-        auto ret = gko::solver::Multigrid::build()
-                    .with_mg_level(    
-                        gko::multigrid::Pgm<scalar>::build()
-                            .with_deterministic(false))
-                    .with_min_coarse_rows(                          
-                        16u)  // necessary since the test matrices have less
-                              // rows than the default value
+        outerStoppingCriterionVec_.push_back(
+            outerStoppingCriterion_.build_dist_stopping_criterion(
+                exec, gkomatrix, x, b, verbose, export_res,
+                get_prev_number_of_iterations(),
+                get_solve_prev_rel_res_cost()
+                ));
+
+        auto coarsest_solver =
+            gko::share(cg::build()
                     .with_criteria(
                         gko::stop::Iteration::build().with_max_iters(
-                            num_iters),
+                            coarse_max_iters_),
                         gko::stop::ResidualNorm<scalar>::build()
                             .with_baseline(gko::stop::mode::absolute)
-                            .with_reduction_factor(reduction_factor))
-                    .on(exec);
-        return ret;
-    }
+                            .with_reduction_factor(1e-18))
+                           .on(exec));
+
+        // std::cout << __FILE__ << __LINE__ 
+        //     << " coarse_max_iters " << coarse_max_iters_ << "\n";
 
-    // Create solver
-    std::unique_ptr<mg::Factory, std::default_delete<mg::Factory>>
-    create_solver(std::shared_ptr<const gko::Executor> exec,
-                  std::shared_ptr<dist_mtx> gkomatrix, std::shared_ptr<dist_vec> x,
-                  std::shared_ptr<dist_vec> b, const label verbose,
-                  const bool export_res,
-                  std::shared_ptr<gko::LinOp> precond) const
-    {
-        // outerStoppingCriterionVec_.push_back(
-        //     outerStoppingCriterion_.build_dist_stopping_criterion(
-        //         exec, gkomatrix, x, b, verbose, export_res, 0, 0));
-        //
-        // auto inner_solver_gen =
-        //     gko::share(bj::build()
-        //                    .with_skip_sorting(true)
-        //                    .with_max_block_size(max_block_size_)
-        //                    .on(exec));
-        //
-        // auto smoother_gen = gko::share(
-        //     ir::build()
-        //         .with_solver(inner_solver_gen)
-        //         .with_relaxation_factor(smoother_relaxation_factor_)
-        //         .with_criteria(gko::stop::Iteration::build()
-        //                            .with_max_iters(smoother_max_iters_)
-        //                            .on(exec))
-        //         .on(exec));
-        // // Create MultigridLevel factory
-        // auto mg_level_gen = amgx_pgm::build()
-        //                         .with_deterministic(true)
-        //                         .with_skip_sorting(true)
-        //                         .on(exec);
-        // // Create CoarsestSolver factory
-        // auto coarsest_gen =
-        //     gko::share(ir::build()
-        //                    .with_solver(inner_solver_gen)
-        //                    .with_relaxation_factor(inner_relaxation_factor_)
-        //                    .with_criteria(gko::stop::Iteration::build()
-        //                                       .with_max_iters(coarse_max_iters_)
-        //                                       .on(exec))
-        //                    .on(exec));
-        //
-        // // Create multigrid factory
-        // auto ret = mg::build()
-        //                .with_max_levels(max_levels_)
-        //                .with_min_coarse_rows(min_coarse_rows_)
-        //                .with_pre_smoother(smoother_gen)
-        //                .with_post_uses_pre(true)
-        //                .with_mg_level(gko::share(mg_level_gen))
-        //                .with_coarsest_solver(coarsest_gen)
-        //                .with_zero_guess(false)
-        //                .with_criteria(outerStoppingCriterionVec_)
-        //                .on(exec);
-        std::cout << __FILE__ << ":" << __LINE__ << "create mg solver\n";
-        const scalar reduction_factor{1e-8};
-        label num_iters {100};
         auto ret = gko::solver::Multigrid::build()
+                    .with_max_levels(max_levels_)
                     .with_mg_level(    
                         gko::multigrid::Pgm<scalar>::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(
-                            num_iters),
-                        gko::stop::ResidualNorm<scalar>::build()
-                            .with_baseline(gko::stop::mode::absolute)
-                            .with_reduction_factor(reduction_factor))
+                            .with_deterministic(true))
+                    .with_min_coarse_rows(min_coarse_rows_)
+                    .with_coarsest_solver(coarsest_solver)
+                    .with_criteria(outerStoppingCriterionVec_)
                     .on(exec);
         return ret;
     }
 
+    scalar get_solve_prev_rel_res_cost() const
+    {
+        return ::Foam::get_solve_prev_rel_res_cost(sysMatrixName_, db_);
+    }
+
     scalar get_init_res_norm() const
     {
         return outerStoppingCriterion_.get_init_res_norm();
@@ -208,23 +153,21 @@ public:
 
     void store_number_of_iterations() const
     {
-        // set_solve_prev_iters(sysMatrixName_, db_,
-        //                      stoppingCriterion_.get_num_iters(),
-        //                      stoppingCriterion_.get_is_final());
+        set_solve_prev_iters(sysMatrixName_, db_,
+                             outerStoppingCriterion_.get_num_iters(),
+                             outerStoppingCriterion_.get_is_final()
+                             );
     }
 
     label get_prev_number_of_iterations() const
     {
-        // return get_solve_prev_iters(sysMatrixName_, db_,
-        //                             stoppingCriterion_.get_is_final());
-        return 0;
+        return get_solve_prev_iters(sysMatrixName_, db_,
+                                    innerStoppingCriterion_.get_is_final());
     }
 
     label get_res_norm_time() const
     {
-        // return get_solve_prev_iters(sysMatrixName_, db_,
-        //                             stoppingCriterion_.get_is_final());
-        return 0;
+        return outerStoppingCriterion_.get_res_norm_time();
     }
 
 
diff --git a/StoppingCriterion/StoppingCriterion.C b/StoppingCriterion/StoppingCriterion.C
index 74406b5a6..ee6cd03fe 100644
--- a/StoppingCriterion/StoppingCriterion.C
+++ b/StoppingCriterion/StoppingCriterion.C
@@ -31,7 +31,7 @@ void StoppingCriterion::OpenFOAMDistStoppingCriterion::compute_Axref_dist(
 
 scalar
 StoppingCriterion::OpenFOAMDistStoppingCriterion::compute_normfactor_dist(
-    std::shared_ptr<const gko::Executor> device_exec, const dist_vec *r,
+    std::shared_ptr<const gko::Executor> device_exec, const dist_vec* r,
     std::shared_ptr<const gko::LinOp> gkomatrix,
     std::shared_ptr<const dist_vec> x, std::shared_ptr<const dist_vec> b) const
 {
@@ -65,7 +65,9 @@ StoppingCriterion::OpenFOAMDistStoppingCriterion::compute_normfactor_dist(
     auto res_host = vec::create(device_exec->get_master(), gko::dim<2>{1});
     res_host->copy_from(res.get());
 
-    return res_host->get_values()[0] + SMALL;
+    auto normfactor = res_host->get_values()[0] + SMALL;
+
+    return normfactor;
 }
 
 bool StoppingCriterion::OpenFOAMDistStoppingCriterion::check_impl(
@@ -74,6 +76,7 @@ bool StoppingCriterion::OpenFOAMDistStoppingCriterion::check_impl(
     const Criterion::Updater &updater)
 {
     // Dont check residual norm before minIter is reached
+
     if (*(parameters_.iter) > 0 &&
         *(parameters_.iter) < parameters_.openfoam_minIter) {
         *(parameters_.iter) += 1;
@@ -89,7 +92,20 @@ bool StoppingCriterion::OpenFOAMDistStoppingCriterion::check_impl(
     auto start_eval = std::chrono::steady_clock::now();
     const auto exec = this->get_executor();
 
-    auto *dense_r = gko::as<dist_vec>(updater.residual_);
+    std::shared_ptr<dist_vec> dense_r_vec;
+    // auto dense_r = gko::share(dist_vec::create_with_config_of(gko::as<dist_vec>(updater.solution_)));
+    // multigrid does not set residual for out iterations
+    if (updater.residual_ == nullptr) {
+
+        dense_r_vec = parameters_.b->clone();
+
+        auto one{gko::initialize<gko::matrix::Dense<scalar>>({1}, exec)};
+        auto neg_one{gko::initialize<gko::matrix::Dense<scalar>>({-1}, exec)};
+        parameters_.gkomatrix->apply(neg_one, updater.solution_, one, dense_r_vec);
+    }
+
+    const dist_vec* dense_r = (updater.residual_ == nullptr) ? dense_r_vec.get() : gko::as<dist_vec>(updater.residual_);
+
     auto norm1 = vec::create(exec, gko::dim<2>{1});
     dense_r->compute_norm1(norm1.get());
     auto norm1_host = vec::create(exec->get_master(), gko::dim<2>{1});
diff --git a/StoppingCriterion/StoppingCriterion.H b/StoppingCriterion/StoppingCriterion.H
index 2db8e8449..d5a3ac4d6 100644
--- a/StoppingCriterion/StoppingCriterion.H
+++ b/StoppingCriterion/StoppingCriterion.H
@@ -64,11 +64,11 @@ class StoppingCriterion {
 
             bool GKO_FACTORY_PARAMETER(export_res, false);
 
-            std::shared_ptr<gko::LinOp> GKO_FACTORY_PARAMETER(gkomatrix, {});
+            std::shared_ptr<const gko::LinOp> GKO_FACTORY_PARAMETER(gkomatrix, {});
 
-            std::shared_ptr<dist_vec> GKO_FACTORY_PARAMETER(x, {});
+            std::shared_ptr<const dist_vec> GKO_FACTORY_PARAMETER(x, {});
 
-            std::shared_ptr<dist_vec> GKO_FACTORY_PARAMETER(b, {});
+            std::shared_ptr<const dist_vec> GKO_FACTORY_PARAMETER(b, {});
         };
 
         GKO_ENABLE_CRITERION_FACTORY(OpenFOAMDistStoppingCriterion, parameters,
@@ -91,7 +91,8 @@ class StoppingCriterion {
          * or rewritten as || r - ( b - x* ) || + || (b - x*) ||
          *  */
         scalar compute_normfactor_dist(
-            std::shared_ptr<const gko::Executor> device_exec, const dist_vec *r,
+            std::shared_ptr<const gko::Executor> device_exec,
+            const dist_vec* r,
             std::shared_ptr<const gko::LinOp> gkomatrix,
             std::shared_ptr<const dist_vec> x,
             std::shared_ptr<const dist_vec> b) const;
@@ -130,6 +131,7 @@ class StoppingCriterion {
         mutable bool eval_norm_factor_ = true;
 
         mutable std::vector<scalar> res_norms_{};
+    
     };
 
     mutable label maxIter_;
@@ -189,10 +191,10 @@ public:
     }
 
     std::shared_ptr<const gko::stop::CriterionFactory>
-    build_dist_stopping_criterion(std::shared_ptr<gko::Executor> device_exec,
-                                  std::shared_ptr<gko::LinOp> gkomatrix,
-                                  std::shared_ptr<dist_vec> x,
-                                  std::shared_ptr<dist_vec> b, label verbose,
+    build_dist_stopping_criterion(std::shared_ptr<const gko::Executor> device_exec,
+                                  std::shared_ptr<const gko::LinOp> gkomatrix,
+                                  std::shared_ptr<const dist_vec> x,
+                                  std::shared_ptr<const dist_vec> b, label verbose,
                                   bool export_res, label prev_solve_iters,
                                   scalar prev_rel_cost) const
     {
diff --git a/lduLduBase/lduLduBase.H b/lduLduBase/lduLduBase.H
index 225fe73c5..c79b06b28 100644
--- a/lduLduBase/lduLduBase.H
+++ b/lduLduBase/lduLduBase.H
@@ -263,7 +263,6 @@ public:
             dist_b.write();
         }
 
-
         LOG_1(verbose_, "create solver")
         auto solver_gen = this->create_dist_solver(
             this->get_exec_handler().get_device_exec(), dist_A_v, dist_x_v,
@@ -275,34 +274,39 @@ public:
         TIME_WITH_FIELDNAME(verbose_, solve, this->fieldName(),
                             solver->apply(dist_b_v, dist_x_v);)
 
+        auto exec_handler = this->get_exec_handler();
+        auto exec = exec_handler.get_device_exec();
+        auto comm = *exec_handler.get_gko_mpi_host_comm().get();
         TIME_WITH_FIELDNAME(verbose_, copy_x_back, this->fieldName(),
-                            dist_x.copy_back();)
+                            dist_x.copy_back(
+                                exec, comm
+                                );)
         auto bandwidth_copy_back = sizeof(scalar) * partition.get_local_size() /
                                    delta_t_copy_x_back / 1000.0;
 
         solverPerf.initialResidual() = this->get_init_res_norm();
         solverPerf.finalResidual() = this->get_res_norm();
         solverPerf.nIterations() = this->get_number_of_iterations();
-        this->store_number_of_iterations();
-        auto time_for_res_norm_eval = this->get_res_norm_time();
-        auto time_per_iter =
-            delta_t_solve / max(this->get_number_of_iterations(), 1);
-        scalar prev_rel_res_cost = time_per_iter / time_for_res_norm_eval;
-        this->get_exec_handler().get_gko_mpi_host_comm()->broadcast(
-            this->get_exec_handler().get_ref_exec(), &prev_rel_res_cost, 1, 0);
-        this->set_prev_rel_res_cost(prev_rel_res_cost);
-        auto time_per_iter_and_dof =
-            time_per_iter * 1000.0 / partition.get_total_size();
-        word msg =
-            "\nStatistics:\n\tTime per iteration: " +
-            std::to_string(time_per_iter) +
-            std::string(" [mu s]\n\tTime per residual norm calculation: ") +
-            std::to_string(time_for_res_norm_eval) +
-            std::string(" [mu s]\n\tTime per iteration and DOF: ") +
-            std::to_string(time_per_iter_and_dof) + std::string(" [ns]") +
-            std::string("\n\tRetrieve results bandwidth ") +
-            std::to_string(bandwidth_copy_back) + std::string(" [GByte/s]");
-        MLOG_0(verbose_, msg)
+        // this->store_number_of_iterations();
+        // auto time_for_res_norm_eval = this->get_res_norm_time();
+        // auto time_per_iter =
+        //     delta_t_solve / max(this->get_number_of_iterations(), 1);
+        // scalar prev_rel_res_cost = time_per_iter / time_for_res_norm_eval;
+        // this->get_exec_handler().get_gko_mpi_host_comm()->broadcast(
+        //     this->get_exec_handler().get_ref_exec(), &prev_rel_res_cost, 1, 0);
+        // this->set_prev_rel_res_cost(prev_rel_res_cost);
+        // auto time_per_iter_and_dof =
+        //     time_per_iter * 1000.0 / partition.get_total_size();
+        // word msg =
+        //     "\nStatistics:\n\tTime per iteration: " +
+        //     std::to_string(time_per_iter) +
+        //     std::string(" [mu s]\n\tTime per residual norm calculation: ") +
+        //     std::to_string(time_for_res_norm_eval) +
+        //     std::string(" [mu s]\n\tTime per iteration and DOF: ") +
+        //     std::to_string(time_per_iter_and_dof) + std::string(" [ns]") +
+        //     std::string("\n\tRetrieve results bandwidth ") +
+        //     std::to_string(bandwidth_copy_back) + std::string(" [GByte/s]");
+        // MLOG_0(verbose_, msg)
 
         return solverPerf;
     }

From 0fcfacfd62817cbaaefa0520fcd5bff530928b30 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Thu, 1 Aug 2024 09:59:50 +0200
Subject: [PATCH 3/9] wip mg

---
 CMakeLists.txt                                |   2 +-
 .../CsrMatrixWrapper/CsrMatrixWrapper.H       |   8 +-
 HostMatrix/HostMatrix.C                       |  19 +++-
 Solver/Multigrid/GKOMultigrid.H               | 107 ++++++++++++------
 4 files changed, 92 insertions(+), 44 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 87d799b15..ec8414936 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -73,7 +73,7 @@ if((NOT OGL_ALLOW_REFERENCE_ONLY) AND (NOT OGL_USE_EXTERNAL_GINKGO))
 endif()
 
 set(GINKGO_CHECKOUT_VERSION
-    "ogl_rebase_1.8.0"
+    "develop"
     CACHE STRING "Use specific version of ginkgo")
 
 include(CheckIncludeFileCXX)
diff --git a/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H b/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H
index a2f46b706..aa752ae2d 100644
--- a/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H
+++ b/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H
@@ -80,8 +80,6 @@ struct MatrixInitFunctor {
         if (regenerate_) {
             TIME_WITH_FIELDNAME(verbose_, regenerate_distributed_matrix,
                                 field_name_, auto reinit = init();)
-            std::cout << __FILE__ << ":" << __LINE__ << " Reinit matrix \n";
-            // TODO NOTE use swap here
             persistent_device_matrix->copy_from(reinit.get());
             return;
         }
@@ -161,8 +159,10 @@ struct MatrixInitFunctor {
                 gko::share(gko::matrix::Csr<scalar, label>::create(exec));
             non_local_linop->read(non_local_A_data);
             return dist_mtx::create(exec, *comm.get(), global_dim, local_linop,
-                                    non_local_linop, recv_sizes, recv_offsets,
-                                    recv_gather_idxs);
+                                    non_local_linop,
+                                    recv_sizes, recv_offsets,
+                                    recv_gather_idxs
+                                    );
         }
         if (matrix_format_ == "Coo") {
             auto local_linop =
diff --git a/HostMatrix/HostMatrix.C b/HostMatrix/HostMatrix.C
index 74657a3df..476d5b359 100644
--- a/HostMatrix/HostMatrix.C
+++ b/HostMatrix/HostMatrix.C
@@ -255,6 +255,8 @@ HostMatrixWrapper<MatrixType>::create_communication_pattern(
 {
     // temp map, mapping from neighbour rank interface cells
     std::map<label, std::vector<label>> interface_cell_map{};
+    auto exec = exec_.get_device_exec();
+    auto comm = *exec_.get_gko_mpi_device_comm().get();
 
     // iterate all interfaces, count number of neighbour procs
     // and store rows to send to neighbour procs
@@ -265,26 +267,33 @@ HostMatrixWrapper<MatrixType>::create_communication_pattern(
             const auto &face_cells{iface->interface().faceCells()};
             const label neighbProcNo = patch.neighbProcNo();
 
+            auto recv_size = face_cells.size();
+            auto recv_cells = std::vector<label>(recv_size); 
+
+            // we need the cell idx of the other side
+            comm.send(exec, face_cells.begin(), recv_size, neighbProcNo, 0);
+            comm.recv(exec, recv_cells.data(), recv_size, neighbProcNo, 0);
+
             auto search = interface_cell_map.find(neighbProcNo);
             if (search == interface_cell_map.end()) {
                 n_procs++;
                 interface_cell_map.insert(std::pair{
                     neighbProcNo,
-                    std::vector<label>(face_cells.begin(), face_cells.end())});
+                    std::vector<label>(recv_cells.begin(), recv_cells.end())});
             } else {
                 auto &vec = search->second;
-                vec.insert(vec.end(), face_cells.begin(), face_cells.end());
+                vec.insert(vec.end(), recv_cells.begin(), recv_cells.end());
             }
         });
 
     using comm_size_type = CommunicationPattern::comm_size_type;
 
     // create index_sets
-    std::vector<std::pair<gko::array<label>, comm_size_type>> send_idxs;
+    std::vector<std::pair<gko::array<label>, comm_size_type>> recv_idxs;
     for (auto [proc, interface_cells] : interface_cell_map) {
         auto exec = exec_.get_ref_exec();
         auto arr = gko::array<label>::const_view(exec, interface_cells.size(), interface_cells.data());
-        send_idxs.push_back(std::pair<gko::array<label>, comm_size_type>(
+        recv_idxs.push_back(std::pair<gko::array<label>, comm_size_type>(
             arr.copy_to_array(), proc));
     }
 
@@ -301,7 +310,7 @@ HostMatrixWrapper<MatrixType>::create_communication_pattern(
         iter++;
     }
 
-    return CommunicationPattern{target_ids, target_sizes, send_idxs};
+    return CommunicationPattern{target_ids, target_sizes, recv_idxs};
 }
 
 
diff --git a/Solver/Multigrid/GKOMultigrid.H b/Solver/Multigrid/GKOMultigrid.H
index 0cb5ac9eb..48df9cac9 100644
--- a/Solver/Multigrid/GKOMultigrid.H
+++ b/Solver/Multigrid/GKOMultigrid.H
@@ -14,11 +14,15 @@ namespace Foam {
 class GKOMultigridFactory {
 private:
     using cg = gko::solver::Cg<scalar>;
+    using direct = gko::experimental::solver::Direct<scalar, label>;
     using mtx = gko::matrix::Csr<scalar>;
     using ir = gko::solver::Ir<scalar>;
     using mg = gko::solver::Multigrid;
     using bj = gko::preconditioner::Jacobi<scalar, label>;
     using amgx_pgm = gko::multigrid::AmgxPgm<scalar, label>;
+    using ras =
+        gko::experimental::distributed::preconditioner::Schwarz<scalar, label,
+                                                                label>;
 
     using vec = gko::matrix::Dense<scalar>;
     using dist_vec = gko::experimental::distributed::Vector<scalar>;
@@ -37,7 +41,7 @@ private:
 
     const StoppingCriterion innerStoppingCriterion_;
 
-    const word smoother_solver_;
+    const word coarsest_solver_;
 
     const label max_block_size_;
 
@@ -53,6 +57,8 @@ private:
 
     const label min_coarse_rows_;
 
+    const label cycle_;
+
     mutable std::vector<std::shared_ptr<const gko::stop::CriterionFactory>>
         outerStoppingCriterionVec_;
 
@@ -68,8 +74,8 @@ public:
           sysMatrixName_(sysMatrixName),
           outerStoppingCriterion_(solverControls),
           innerStoppingCriterion_(StoppingCriterion(innerSolverControls_)),
-          smoother_solver_(innerSolverControls_.lookupOrDefault(
-              "smootherSolver", word("CG"))),
+          coarsest_solver_(innerSolverControls_.lookupOrDefault(
+              "coarsestSolver", word("CG"))),
           max_block_size_(
               innerSolverControls_.lookupOrDefault("maxBlockSize", label(4))),
           inner_relaxation_factor_(innerSolverControls_.lookupOrDefault(
@@ -83,45 +89,79 @@ public:
           max_levels_(
               innerSolverControls_.lookupOrDefault("maxLevels", label(9))),
           min_coarse_rows_(
-              innerSolverControls_.lookupOrDefault("minCoarseRows", label(10)))
-    {}
+              innerSolverControls_.lookupOrDefault("minCoarseRows", label(10))),
+          cycle_(innerSolverControls_.lookupOrDefault("cycle", label(0)))
+    {
+        auto mtx_format =
+            solverControls.lookupOrDefault("matrixFormat", word("Coo"));
+        if (mtx_format != "Csr") {
+            FatalErrorInFunction
+                << "Ginkgos Multigrid solver currently only supports Csr "
+                   "matrices make sure to set: 'matrixFormat Csr;'"
+                << abort(FatalError);
+        }
+    }
 
     std::unique_ptr<mg::Factory, std::default_delete<mg::Factory>>
     create_dist_solver(std::shared_ptr<const gko::Executor> exec,
                        std::shared_ptr<const gko::LinOp> gkomatrix,
-                       std::shared_ptr<const dist_vec> x, std::shared_ptr<const dist_vec> b,
-                       const label verbose, const bool export_res,
+                       std::shared_ptr<const dist_vec> x,
+                       std::shared_ptr<const dist_vec> b, const label verbose,
+                       const bool export_res,
                        std::shared_ptr<gko::LinOp> precond) const
     {
         outerStoppingCriterionVec_.push_back(
             outerStoppingCriterion_.build_dist_stopping_criterion(
                 exec, gkomatrix, x, b, verbose, export_res,
                 get_prev_number_of_iterations(),
-                get_solve_prev_rel_res_cost()
-                ));
-
-        auto coarsest_solver =
-            gko::share(cg::build()
-                    .with_criteria(
-                        gko::stop::Iteration::build().with_max_iters(
-                            coarse_max_iters_),
-                        gko::stop::ResidualNorm<scalar>::build()
-                            .with_baseline(gko::stop::mode::absolute)
-                            .with_reduction_factor(1e-18))
-                           .on(exec));
-
-        // std::cout << __FILE__ << __LINE__ 
-        //     << " coarse_max_iters " << coarse_max_iters_ << "\n";
-
-        auto ret = gko::solver::Multigrid::build()
-                    .with_max_levels(max_levels_)
-                    .with_mg_level(    
-                        gko::multigrid::Pgm<scalar>::build()
-                            .with_deterministic(true))
-                    .with_min_coarse_rows(min_coarse_rows_)
-                    .with_coarsest_solver(coarsest_solver)
-                    .with_criteria(outerStoppingCriterionVec_)
-                    .on(exec);
+                get_solve_prev_rel_res_cost()));
+
+        std::shared_ptr<const gko::LinOpFactory> coarsest_solver{};
+
+        if (coarsest_solver_ == "CG") {
+            coarsest_solver = gko::share(
+                cg::build()
+                    // .with_preconditioner(ras::build().with_local_solver(
+                    //     bj::build().with_max_block_size(1u)))
+                    .with_criteria(gko::stop::Iteration::build().with_max_iters(
+                                       coarse_max_iters_),
+                                   gko::stop::ResidualNorm<scalar>::build()
+                                       .with_baseline(gko::stop::mode::absolute)
+                                       .with_reduction_factor(1e-18))
+                    .on(exec));
+        } else {
+            FatalErrorInFunction
+                << "Unsupported coarsestSolver: " << coarsest_solver_
+                << "! Only CG is supported" << abort(FatalError);
+        }
+        // if (coarsest_solver_ == "direct") {
+        //     coarsest_solver = gko::share(
+        //         ras::build()
+        //             .with_local_solver(direct::build().with_factorization(
+        //                 gko::experimental::factorization::Lu<scalar,
+        //                                                      label>::build()
+        //                     .with_symbolic_algorithm(
+        //                         gko::experimental::factorization::
+        //                             symbolic_type::symmetric)))
+        //             .on(exec));
+        // }
+
+        auto ret =
+            gko::solver::Multigrid::build()
+                .with_max_levels(max_levels_)
+                .with_mg_level(
+                    gko::multigrid::Pgm<scalar>::build().with_deterministic(
+                        false))
+                .with_min_coarse_rows(min_coarse_rows_)
+                .with_coarsest_solver(coarsest_solver)
+                .with_criteria(outerStoppingCriterionVec_)
+                .with_smoother_iters(smoother_max_iters_)
+                .with_cycle(cycle_)
+                // .with_pre_smoother(
+                //         smoother_gen
+                //         )
+                // .with_post_uses_pre(true)
+                .on(exec);
         return ret;
     }
 
@@ -155,8 +195,7 @@ public:
     {
         set_solve_prev_iters(sysMatrixName_, db_,
                              outerStoppingCriterion_.get_num_iters(),
-                             outerStoppingCriterion_.get_is_final()
-                             );
+                             outerStoppingCriterion_.get_is_final());
     }
 
     label get_prev_number_of_iterations() const

From 48fdf72825fce0aa54a59880b25a4b7550200d62 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Thu, 1 Aug 2024 10:31:18 +0200
Subject: [PATCH 4/9] fix Ell

---
 .../CsrMatrixWrapper/CsrMatrixWrapper.H       | 30 ++++++++++++-------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H b/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H
index aa752ae2d..a03c5b9db 100644
--- a/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H
+++ b/DevicePersistent/CsrMatrixWrapper/CsrMatrixWrapper.H
@@ -105,15 +105,6 @@ struct MatrixInitFunctor {
                 (gko::matrix::Csr<scalar, label> *)non_local_matrix.get();
             non_local_value_ptr = non_local_coo_ptr->get_values();
         }
-        if (matrix_format_ == "Ell") {
-            gko::matrix::Ell<scalar, label> *coo_ptr =
-                (gko::matrix::Ell<scalar, label> *)local_matrix.get();
-            value_ptr = coo_ptr->get_values();
-
-            gko::matrix::Ell<scalar, label> *non_local_coo_ptr =
-                (gko::matrix::Ell<scalar, label> *)non_local_matrix.get();
-            non_local_value_ptr = non_local_coo_ptr->get_values();
-        }
         if (matrix_format_ == "Coo") {
             gko::matrix::Coo<scalar, label> *coo_ptr =
                 (gko::matrix::Coo<scalar, label> *)local_matrix.get();
@@ -123,6 +114,10 @@ struct MatrixInitFunctor {
                 (gko::matrix::Coo<scalar, label> *)non_local_matrix.get();
             non_local_value_ptr = non_local_coo_ptr->get_values();
         }
+        if (matrix_format_ == "Ell") {
+        FatalErrorInFunction << "Updating Ell matrices not supported"
+                             << "set: 'regenerate true;'" << abort(FatalError);
+        }
 
         gko::array<scalar> device_values{device_exec, *coeffs.get()};
         gko::array<scalar> device_non_local_values{device_exec,
@@ -178,9 +173,22 @@ struct MatrixInitFunctor {
             return dist;
         }
 
+        if (matrix_format_ == "Ell") {
+            auto local_linop =
+                gko::share(gko::matrix::Ell<scalar, label>::create(exec));
+            local_linop->read(local_A_data);
+            auto non_local_linop =
+                gko::share(gko::matrix::Coo<scalar, label>::create(exec));
+            non_local_linop->read(non_local_A_data);
+
+            auto dist = dist_mtx::create(
+                exec, *comm.get(), global_dim, local_linop, non_local_linop,
+                recv_sizes, recv_offsets, recv_gather_idxs);
+            return dist;
+        }
 
-        // FatalErrorInFunction << "Matrix format " << matrix_format_
-        //                      << " not supported " << abort(FatalError);
+        FatalErrorInFunction << "Matrix format " << matrix_format_
+                             << " not supported " << abort(FatalError);
     }
 
     std::shared_ptr<dist_mtx> init() const

From 9a5e0e084d16ae50952e5bef05ca844357e2c2bf Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Thu, 1 Aug 2024 12:32:42 +0200
Subject: [PATCH 5/9] clean up

---
 Solver/Multigrid/GKOMultigrid.H | 29 ++++++++++++-----------------
 1 file changed, 12 insertions(+), 17 deletions(-)

diff --git a/Solver/Multigrid/GKOMultigrid.H b/Solver/Multigrid/GKOMultigrid.H
index 48df9cac9..e98410337 100644
--- a/Solver/Multigrid/GKOMultigrid.H
+++ b/Solver/Multigrid/GKOMultigrid.H
@@ -31,16 +31,12 @@ private:
 
     const dictionary &solverControls_;
 
-    const dictionary &innerSolverControls_;
-
     const objectRegistry &db_;
 
     const word sysMatrixName_;
 
     const StoppingCriterion outerStoppingCriterion_;
 
-    const StoppingCriterion innerStoppingCriterion_;
-
     const word coarsest_solver_;
 
     const label max_block_size_;
@@ -69,28 +65,26 @@ public:
     GKOMultigridFactory(const dictionary &solverControls,
                         const objectRegistry &db, word sysMatrixName)
         : solverControls_(solverControls),
-          innerSolverControls_(solverControls.subDict("inner")),
           db_(db),
           sysMatrixName_(sysMatrixName),
           outerStoppingCriterion_(solverControls),
-          innerStoppingCriterion_(StoppingCriterion(innerSolverControls_)),
-          coarsest_solver_(innerSolverControls_.lookupOrDefault(
+          coarsest_solver_(solverControls_.lookupOrDefault(
               "coarsestSolver", word("CG"))),
           max_block_size_(
-              innerSolverControls_.lookupOrDefault("maxBlockSize", label(4))),
-          inner_relaxation_factor_(innerSolverControls_.lookupOrDefault(
+              solverControls_.lookupOrDefault("maxBlockSize", label(4))),
+          inner_relaxation_factor_(solverControls_.lookupOrDefault(
               "innerRelaxationFactor", scalar(0.9))),
-          smoother_relaxation_factor_(innerSolverControls_.lookupOrDefault(
+          smoother_relaxation_factor_(solverControls_.lookupOrDefault(
               "smootherRelaxationFactor", scalar(0.9))),
-          smoother_max_iters_(innerSolverControls_.lookupOrDefault(
+          smoother_max_iters_(solverControls_.lookupOrDefault(
               "smootherMaxIters", label(2))),
           coarse_max_iters_(
-              innerSolverControls_.lookupOrDefault("coarseMaxIters", label(4))),
+              solverControls_.lookupOrDefault("coarseMaxIters", label(4))),
           max_levels_(
-              innerSolverControls_.lookupOrDefault("maxLevels", label(9))),
+              solverControls_.lookupOrDefault("maxLevels", label(9))),
           min_coarse_rows_(
-              innerSolverControls_.lookupOrDefault("minCoarseRows", label(10))),
-          cycle_(innerSolverControls_.lookupOrDefault("cycle", label(0)))
+              solverControls_.lookupOrDefault("minCoarseRows", label(10))),
+          cycle_(solverControls_.lookupOrDefault("cycle", label(0)))
     {
         auto mtx_format =
             solverControls.lookupOrDefault("matrixFormat", word("Coo"));
@@ -200,8 +194,9 @@ public:
 
     label get_prev_number_of_iterations() const
     {
-        return get_solve_prev_iters(sysMatrixName_, db_,
-                                    innerStoppingCriterion_.get_is_final());
+        return 0;
+        // get_solve_prev_iters(sysMatrixName_, db_,
+        //                             innerStoppingCriterion_.get_is_final());
     }
 
     label get_res_norm_time() const

From d6c1a774d41e9d0983a942113f385dace42a0f11 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Sat, 3 Aug 2024 08:29:35 +0200
Subject: [PATCH 6/9] print stats

---
 Preconditioner/Preconditioner.H | 136 ++++++++++++++++++--------------
 lduLduBase/lduLduBase.H         |  38 ++++-----
 2 files changed, 94 insertions(+), 80 deletions(-)

diff --git a/Preconditioner/Preconditioner.H b/Preconditioner/Preconditioner.H
index 7fa138d04..6a93b17d6 100644
--- a/Preconditioner/Preconditioner.H
+++ b/Preconditioner/Preconditioner.H
@@ -15,6 +15,7 @@ class Preconditioner {
     using bj = gko::preconditioner::Jacobi<>;
     using ic = gko::preconditioner::Ic<>;
     using ir = gko::solver::Ir<scalar>;
+    using it = gko::stop::Iteration;
     using cg = gko::solver::Cg<scalar>;
     using mg = gko::solver::Multigrid;
     using amgx_pgm = gko::multigrid::Pgm<scalar, label>;
@@ -258,26 +259,86 @@ public:
             return wrap_schwarz(gkomatrix, device_exec, std::move(pre_factory));
         }
         if (name == "Multigrid") {
-            auto inner_solver_gen =
-                gko::share(bj::build().with_max_block_size(1u).on(device_exec));
+            word type = controls.lookupOrDefault("type", word("Schwarz"));
 
-            auto smoother_gen = gko::share(
-                ir::build()
-                    .with_solver(inner_solver_gen)
-                    .with_relaxation_factor(0.9)
-                    .with_criteria(
-                        gko::stop::Iteration::build().with_max_iters(2u).on(
-                            device_exec))
+            const label coarse_solver_iters(
+                controls.lookupOrDefault("coarseSolverIters", label(4)));
+
+            auto cycle_name = controls.lookupOrDefault("cycle", word("v"));
+            gko::solver::multigrid::cycle cycle;
+            if (cycle_name == "v") cycle = gko::solver::multigrid::cycle::v;
+            if (cycle_name == "w") cycle = gko::solver::multigrid::cycle::w;
+            if (cycle_name == "f") cycle = gko::solver::multigrid::cycle::f;
+
+            const label max_levels(
+                controls.lookupOrDefault("maxLevels", label(5)));
+            const label min_coarse_rows(
+                controls.lookupOrDefault("minCoarseRows", label(10)));
+            bool zeroGuess(controls.lookupOrDefault<Switch>("zeroGuess", true));
+            const label smoother_max_iters(
+                controls.lookupOrDefault("smootherMaxIters", label(1)));
+
+            auto coarsest_gen = gko::share(
+                cg::build()
+                    .with_criteria(gko::stop::Iteration::build().with_max_iters(
+                        static_cast<gko::uint32>(coarse_solver_iters)))
                     .on(device_exec));
 
-            // Create MultigridLevel factory
-            auto mg_level_gen =
-                amgx_pgm::build().with_deterministic(true).on(device_exec);
+            word msg = "Generate preconditioner " + name + " MaxLevels " +
+                       std::to_string(max_levels) + " MinCoarseRows " +
+                       std::to_string(min_coarse_rows) + " ZeroGuess " +
+                       std::to_string(zeroGuess) + " Cycle " + cycle_name +
+                       " type: " + type;
+            MLOG_0(verbose_, msg)
+
+            if (type == "Schwarz") {
+                auto pre_factory =
+                    mg::build()
+                        .with_max_levels(static_cast<gko::uint32>(max_levels))
+                        .with_cycle(cycle)
+                        .with_min_coarse_rows(
+                            static_cast<gko::uint32>(min_coarse_rows))
+                        .with_pre_smoother(
+                            ir::build()
+                                .with_solver(
+                                    bj::build().with_max_block_size(1u))
+                                .with_relaxation_factor(0.9)
+                                .with_criteria(it::build().with_max_iters(
+                                    smoother_max_iters)))
+                        .with_post_uses_pre(true)
+                        .with_mg_level(
+                            amgx_pgm::build().with_deterministic(true).on(
+                                device_exec))
+                        .with_coarsest_solver(coarsest_gen)
+                        .with_criteria(it::build().with_max_iters(1u))
+                        .on(device_exec);
+
+                return wrap_schwarz(gkomatrix, device_exec,
+                                    std::move(pre_factory));
+            }
+
+            if (type == "Distributed") {
+                auto ret = gko::share(
+                    gko::solver::Multigrid::build()
+                        .with_max_levels(max_levels)
+                        .with_mg_level(gko::multigrid::Pgm<scalar>::build()
+                                           .with_deterministic(false))
+                        .with_min_coarse_rows(min_coarse_rows)
+                        .with_coarsest_solver(coarsest_gen)
+                        .with_criteria(it::build().with_max_iters(1u))
+                        .with_smoother_iters(smoother_max_iters)
+                        .with_cycle(cycle)
+                        .on(device_exec)
+                        ->generate(gkomatrix));
+                return ret;
+            }
+            FatalErrorInFunction << "Unknown Multigrid type: " << type
+                                 << "\nValid Choices: Schwarz, Distributed"
+                                 << abort(FatalError);
+
 
-            // Create CoarsestSolver factory
             // std::shared_ptr<gko::LinOp> coarsest_gen;
             // coarsest_gen =
-            //
             // (solverControls_.lookupOrDefault("PreconditionerMultigridUseIR",
             //                                      word("IR")) == "IR")
             //         ? gko::share(
@@ -294,53 +355,6 @@ public:
             //                                      .with_max_iters(4u)
             //                                      .on(device_exec))
             //                   .on(device_exec));
-
-            const label coarse_solver_iters(
-                controls.lookupOrDefault("coarseSolverIters", label(4)));
-            auto coarsest_gen = gko::share(
-                cg::build()
-                    .with_criteria(gko::stop::Iteration::build()
-                                       .with_max_iters(static_cast<gko::uint32>(
-                                           coarse_solver_iters))
-                                       .on(device_exec))
-                    .on(device_exec));
-
-            auto cycle_name = controls.lookupOrDefault("cycle", word("v"));
-            gko::solver::multigrid::cycle cycle;
-
-            if (cycle_name == "v") cycle = gko::solver::multigrid::cycle::v;
-            if (cycle_name == "w") cycle = gko::solver::multigrid::cycle::w;
-            if (cycle_name == "f") cycle = gko::solver::multigrid::cycle::f;
-
-            // Create multigrid factory
-            label max_levels(controls.lookupOrDefault("maxLevels", label(9)));
-            label min_coarse_rows(
-                controls.lookupOrDefault("minCoarseRows", label(10)));
-            bool zeroGuess(controls.lookupOrDefault<Switch>("zeroGuess", true));
-
-            word msg = "Generate preconditioner " + name + " MaxLevels " +
-                       std::to_string(max_levels) + " MinCoarseRows " +
-                       std::to_string(min_coarse_rows) + " ZeroGuess " +
-                       std::to_string(zeroGuess) + " Cycle " + cycle_name;
-            MLOG_0(verbose_, msg)
-
-            auto pre_factory =
-                mg::build()
-                    .with_max_levels(static_cast<gko::uint32>(max_levels))
-                    .with_cycle(cycle)
-                    .with_min_coarse_rows(
-                        static_cast<gko::uint32>(min_coarse_rows))
-                    .with_pre_smoother(smoother_gen)
-                    .with_post_uses_pre(true)
-                    .with_mg_level(
-                        amgx_pgm::build().with_deterministic(true).on(
-                            device_exec))
-                    .with_coarsest_solver(coarsest_gen)
-                    .with_criteria(
-                        gko::stop::Iteration::build().with_max_iters(1u).on(
-                            device_exec))
-                    .on(device_exec);
-            return wrap_schwarz(gkomatrix, device_exec, std::move(pre_factory));
         }
         if (name == "none") {
             return {};
diff --git a/lduLduBase/lduLduBase.H b/lduLduBase/lduLduBase.H
index c79b06b28..60e8583bd 100644
--- a/lduLduBase/lduLduBase.H
+++ b/lduLduBase/lduLduBase.H
@@ -288,25 +288,25 @@ public:
         solverPerf.finalResidual() = this->get_res_norm();
         solverPerf.nIterations() = this->get_number_of_iterations();
         // this->store_number_of_iterations();
-        // auto time_for_res_norm_eval = this->get_res_norm_time();
-        // auto time_per_iter =
-        //     delta_t_solve / max(this->get_number_of_iterations(), 1);
-        // scalar prev_rel_res_cost = time_per_iter / time_for_res_norm_eval;
-        // this->get_exec_handler().get_gko_mpi_host_comm()->broadcast(
-        //     this->get_exec_handler().get_ref_exec(), &prev_rel_res_cost, 1, 0);
-        // this->set_prev_rel_res_cost(prev_rel_res_cost);
-        // auto time_per_iter_and_dof =
-        //     time_per_iter * 1000.0 / partition.get_total_size();
-        // word msg =
-        //     "\nStatistics:\n\tTime per iteration: " +
-        //     std::to_string(time_per_iter) +
-        //     std::string(" [mu s]\n\tTime per residual norm calculation: ") +
-        //     std::to_string(time_for_res_norm_eval) +
-        //     std::string(" [mu s]\n\tTime per iteration and DOF: ") +
-        //     std::to_string(time_per_iter_and_dof) + std::string(" [ns]") +
-        //     std::string("\n\tRetrieve results bandwidth ") +
-        //     std::to_string(bandwidth_copy_back) + std::string(" [GByte/s]");
-        // MLOG_0(verbose_, msg)
+        auto time_for_res_norm_eval = this->get_res_norm_time();
+        auto time_per_iter =
+            delta_t_solve / max(this->get_number_of_iterations(), 1);
+        scalar prev_rel_res_cost = time_per_iter / time_for_res_norm_eval;
+        this->get_exec_handler().get_gko_mpi_host_comm()->broadcast(
+            this->get_exec_handler().get_ref_exec(), &prev_rel_res_cost, 1, 0);
+        this->set_prev_rel_res_cost(prev_rel_res_cost);
+        auto time_per_iter_and_dof =
+            time_per_iter * 1000.0 / partition.get_total_size();
+        word msg =
+            "\nStatistics:\n\tTime per iteration: " +
+            std::to_string(time_per_iter) +
+            std::string(" [mu s]\n\tTime per residual norm calculation: ") +
+            std::to_string(time_for_res_norm_eval) +
+            std::string(" [mu s]\n\tTime per iteration and DOF: ") +
+            std::to_string(time_per_iter_and_dof) + std::string(" [ns]") +
+            std::string("\n\tRetrieve results bandwidth ") +
+            std::to_string(bandwidth_copy_back) + std::string(" [GByte/s]");
+        MLOG_0(verbose_, msg)
 
         return solverPerf;
     }

From 86b89ad933aa9bff86947996b28b4eb6622d6d29 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Fri, 9 Aug 2024 09:35:07 +0200
Subject: [PATCH 7/9] print out mg info

---
 Solver/Multigrid/GKOMultigrid.H | 30 ++++++++++++++++++++----------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/Solver/Multigrid/GKOMultigrid.H b/Solver/Multigrid/GKOMultigrid.H
index e98410337..45c33e1d3 100644
--- a/Solver/Multigrid/GKOMultigrid.H
+++ b/Solver/Multigrid/GKOMultigrid.H
@@ -37,9 +37,9 @@ private:
 
     const StoppingCriterion outerStoppingCriterion_;
 
-    const word coarsest_solver_;
+    const label verbose_;
 
-    const label max_block_size_;
+    const word coarsest_solver_;
 
     const scalar inner_relaxation_factor_;
 
@@ -68,20 +68,18 @@ public:
           db_(db),
           sysMatrixName_(sysMatrixName),
           outerStoppingCriterion_(solverControls),
-          coarsest_solver_(solverControls_.lookupOrDefault(
-              "coarsestSolver", word("CG"))),
-          max_block_size_(
-              solverControls_.lookupOrDefault("maxBlockSize", label(4))),
+          verbose_(solverControls.lookupOrDefault<label>("verbose", 0)),
+          coarsest_solver_(
+              solverControls_.lookupOrDefault("coarsestSolver", word("CG"))),
           inner_relaxation_factor_(solverControls_.lookupOrDefault(
               "innerRelaxationFactor", scalar(0.9))),
           smoother_relaxation_factor_(solverControls_.lookupOrDefault(
               "smootherRelaxationFactor", scalar(0.9))),
-          smoother_max_iters_(solverControls_.lookupOrDefault(
-              "smootherMaxIters", label(2))),
+          smoother_max_iters_(
+              solverControls_.lookupOrDefault("smootherMaxIters", label(2))),
           coarse_max_iters_(
               solverControls_.lookupOrDefault("coarseMaxIters", label(4))),
-          max_levels_(
-              solverControls_.lookupOrDefault("maxLevels", label(9))),
+          max_levels_(solverControls_.lookupOrDefault("maxLevels", label(9))),
           min_coarse_rows_(
               solverControls_.lookupOrDefault("minCoarseRows", label(10))),
           cycle_(solverControls_.lookupOrDefault("cycle", label(0)))
@@ -94,6 +92,18 @@ public:
                    "matrices make sure to set: 'matrixFormat Csr;'"
                 << abort(FatalError);
         }
+
+        word msg = std::string("Multigrid parameters:") +
+                   std::string("\n\tcoarsestSolver: ") + coarsest_solver_ +
+                   std::string("\n\tcoarseMaxIters: ") +
+                   std::to_string(coarse_max_iters_) +
+                   std::string("\n\tsmootherMaxIters: ") +
+                   std::to_string(smoother_max_iters_) +
+                   std::string("\n\tmaxLevels: ") +
+                   std::to_string(max_levels_) +
+                   std::string("\n\tminCoasresRows: ") +
+                   std::to_string(min_coarse_rows_);
+        MLOG_0(verbose_, msg)
     }
 
     std::unique_ptr<mg::Factory, std::default_delete<mg::Factory>>

From b2fffe35fb8a58517d04a2dda25d22852ea8eff9 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 21 Aug 2024 16:12:13 +0200
Subject: [PATCH 8/9] use preconditioned coarse grid solver

---
 Preconditioner/Preconditioner.H | 53 ++++++++++++++++++++++++++++++---
 1 file changed, 49 insertions(+), 4 deletions(-)

diff --git a/Preconditioner/Preconditioner.H b/Preconditioner/Preconditioner.H
index 6a93b17d6..002d95cec 100644
--- a/Preconditioner/Preconditioner.H
+++ b/Preconditioner/Preconditioner.H
@@ -51,7 +51,6 @@ public:
         std::shared_ptr<gko::Executor> device_exec,
         std::unique_ptr<PrecondFactory> precond) const
     {
-        if (Pstream::parRun()) {
             return gko::share(
                 ras::build()
                     .with_generated_local_solver(gko::share(precond->generate(
@@ -60,10 +59,49 @@ public:
                             ->get_local_matrix())))
                     .on(device_exec)
                     ->generate(gkomatrix));
-        }
-        return gko::share(precond->generate(gkomatrix));
     }
 
+    template <typename PrecondFactory>
+    std::shared_ptr<gko::LinOp> wrap_multi_level_schwarz(
+        std::shared_ptr<gko::LinOp> gkomatrix,
+        std::shared_ptr<gko::Executor> device_exec,
+        std::unique_ptr<PrecondFactory> precond) const
+    {
+        using pgm = gko::multigrid::Pgm<scalar, label>;
+        using solver = gko::solver::Cg<scalar>;
+
+        auto pre_factory = gko::share(bj::build()
+                               .with_skip_sorting(true)
+                               .with_max_block_size(
+                                   static_cast<gko::uint32>(1))
+                               .on(device_exec));
+
+        auto coarse_solver = gko::share(
+            solver::build()
+                .with_preconditioner(
+                    ras::build().with_local_solver(pre_factory).on(device_exec))
+                .with_criteria(
+                    gko::stop::Iteration::build().with_max_iters(10).on(device_exec),
+                    gko::stop::ResidualNorm<scalar>::build()
+                        .with_reduction_factor(1e-3)
+                        .on(device_exec))
+                .on(device_exec));
+
+        auto pgm_fac = gko::share(pgm::build().on(device_exec));
+
+        return gko::share(
+            ras::build()
+                .with_generated_local_solver(gko::share(precond->generate(
+                    gko::as<gko::experimental::distributed::Matrix<
+                        scalar, label, label>>(gkomatrix)
+                        ->get_local_matrix())))
+                 .with_galerkin_ops(pgm_fac)
+                 .with_coarse_solver(coarse_solver)
+                .on(device_exec)
+                ->generate(gkomatrix));
+    }
+
+
     template <typename PrecondFactory, typename Factorization>
     std::shared_ptr<gko::LinOp> wrap_schwarz(
         std::shared_ptr<gko::LinOp> gkomatrix,
@@ -89,6 +127,9 @@ public:
         bool skip_sorting =
             controls.lookupOrDefault<Switch>("skipSorting", true);
 
+        bool multi_level_schwarz =
+            controls.lookupOrDefault<Switch>("multiLevelSchwarz", true);
+
         if (name == "BJ") {
             // TODO for non constant system matrix reuse block pointers
             label max_block_size(
@@ -103,7 +144,11 @@ public:
                                    .with_max_block_size(
                                        static_cast<gko::uint32>(max_block_size))
                                    .on(device_exec);
-            return wrap_schwarz(gkomatrix, device_exec, std::move(pre_factory));
+            if (multi_level_schwarz) {
+                return wrap_multi_level_schwarz(gkomatrix, device_exec, std::move(pre_factory));
+            } else {
+                return wrap_schwarz(gkomatrix, device_exec, std::move(pre_factory));
+            }
         }
         if (name == "ILU") {
             word msg = "Generate preconditioner " + name;

From 9e3f6f8fc92418c4d81e4442a3db6f4558d0cbe2 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 21 Aug 2024 16:12:35 +0200
Subject: [PATCH 9/9] use preconditioned coarse grid solver

---
 CMakeLists.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index ec8414936..4c1a3a268 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -73,7 +73,7 @@ if((NOT OGL_ALLOW_REFERENCE_ONLY) AND (NOT OGL_USE_EXTERNAL_GINKGO))
 endif()
 
 set(GINKGO_CHECKOUT_VERSION
-    "develop"
+    "multilevel-schwarz"
     CACHE STRING "Use specific version of ginkgo")
 
 include(CheckIncludeFileCXX)