From 4b009d55ecc74ca800718fa2bee19393767a1ddd Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Fri, 9 Aug 2024 09:43:08 +0200
Subject: [PATCH 01/67] add Repartitioner file

---
 include/OGL/Repartitioner.H | 606 ++++++++++++++++++++++++++++++++++++
 1 file changed, 606 insertions(+)
 create mode 100644 include/OGL/Repartitioner.H

diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
new file mode 100644
index 000000000..7bab1b863
--- /dev/null
+++ b/include/OGL/Repartitioner.H
@@ -0,0 +1,606 @@
+// SPDX-FileCopyrightText: 2024 OGL authors
+//
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#pragma once
+
+#include <ginkgo/ginkgo.hpp>
+
+#include "MatrixWrapper/CommunicationPattern/CommunicationPattern.H"
+#include "MatrixWrapper/LDUMatrix/HostMatrix.H"
+
+template <typename T>
+std::vector<T> apply_permutation(const std::vector<T> vec,
+                                 const std::vector<label> &p)
+{
+    std::vector<T> sorted_vec(vec.size());
+    std::transform(p.begin(), p.end(), sorted_vec.begin(),
+                   [&](label i) { return vec[i]; });
+    return sorted_vec;
+}
+
+template <typename T, typename Compare>
+std::vector<label> sort_permutation(const std::vector<T> &vec, Compare compare)
+{
+    std::vector<label> p(vec.size());
+    std::iota(p.begin(), p.end(), 0);
+    std::stable_sort(p.begin(), p.end(), [&](std::size_t i, std::size_t j) {
+        return compare(vec[i], vec[j]);
+    });
+    return p;
+}
+
+// TODO this should be persistent
+class Repartitioner {
+private:
+    using partition = gko::experimental::distributed::Partition<label, label>;
+
+    const label size_;  //! Size (n matrix rows ) before repartitioning
+
+    const label repart_size_;  //! Size after repartitioning
+
+    const label ranks_per_gpu_;
+
+    const label verbose_;
+
+    std::shared_ptr<const partition> orig_partition_;
+
+
+public:
+    Repartitioner(label size, label ranks_per_gpu, label verbose,
+                  const ExecutorHandler &exec_handler)
+        : size_(size),
+          repart_size_(Repartitioner::compute_repart_size(size, ranks_per_gpu,
+                                                          exec_handler)),
+          ranks_per_gpu_(ranks_per_gpu),
+          verbose_(verbose),
+          orig_partition_(gko::share(
+              gko::experimental::distributed::build_partition_from_local_size<
+                  label, label>(exec_handler.get_ref_exec(),
+                                *exec_handler.get_communicator().get(),
+                                size))){};
+
+    /* returns the owner rank for a given rank */
+    label get_owner_rank(label rank) const
+    {
+        return compute_owner_rank(rank, ranks_per_gpu_);
+    };
+
+    /* returns the owner rank for a given rank */
+    label get_owner_rank(const ExecutorHandler &exec_handler) const
+    {
+        return get_owner_rank(get_rank(exec_handler));
+    };
+
+    /* returns if current rank is an owner  */
+    bool is_owner(const ExecutorHandler &exec_handler) const
+    {
+        return get_rank(exec_handler) == get_owner_rank(get_rank(exec_handler));
+    };
+
+    /* @brief check if the given rank gets local after repartitioning
+     *
+     * */
+    bool reparts_to_local(const ExecutorHandler &exec_handler, label rank) const
+    {
+        return get_owner_rank(exec_handler) ==
+               compute_owner_rank(rank, ranks_per_gpu_);
+    };
+
+    /* shortcut to current rank */
+    label get_rank(const ExecutorHandler &exec_handler) const
+    {
+        return exec_handler.get_communicator().get()->rank();
+    };
+
+    label get_ranks_per_gpu() const { return ranks_per_gpu_; }
+
+    static label compute_repart_size(label local_size, label ranks_per_gpu,
+                                     const ExecutorHandler &exec_handler)
+    {
+        if (ranks_per_gpu == 1) {
+            return local_size;
+        }
+
+        CommCounts comm_pattern =
+            compute_send_recv_counts(exec_handler, ranks_per_gpu, 1);
+
+        std::vector<label> data{local_size};
+
+        auto local_sizes =
+            gather_to_owner(exec_handler, comm_pattern, 1, data.data());
+
+        return std::accumulate(local_sizes.begin(), local_sizes.end(), 0);
+    }
+
+    label get_repart_size() const { return repart_size_; }
+
+    gko::dim<2> get_repart_dim() const
+    {
+        return gko::dim<2>{repart_size_, repart_size_};
+    }
+
+    std::shared_ptr<const partition> get_orig_partition() const
+    {
+        return orig_partition_;
+    };
+
+
+    /* @brief given received interfaces this function sorts into local and non
+     *local
+     * returns: SparsityPatternVector storing new non_local sparsity pattern,
+     *  and a pair of locality information, where the first bool stores if the
+     *  interface performs a local communication and the second bool if the
+     *  the interface was originally from this rank
+     */
+    std::pair<SparsityPatternVector, std::vector<bool>>
+    build_non_local_interfaces(const ExecutorHandler &exec_handler,
+                               SparsityPatternVector &loc,
+                               const SparsityPatternVector &non_loc) const
+    {
+        LOG_1(verbose_, "start build non local interfaces")
+        std::vector<label> rows, cols, ldu_mapping, ranks, begins, ends;
+        std::vector<bool> is_local;
+        label merged_ranks_size = non_loc.ranks.size();
+
+        for (label i = 0; i < merged_ranks_size; i++) {
+            // these are the begin ends before merging they need to be
+            // offseted;
+            auto begin = non_loc.begin[i];
+            auto end = non_loc.end[i];
+            bool local = reparts_to_local(exec_handler, non_loc.ranks[i]);
+            is_local.push_back(local);
+            if (local) {
+                // TODO depending on the interface simple non
+                // transforming and interpolating interfaces like now
+                // local processor interfaces could be merged
+                loc.begin.push_back(loc.rows.size());
+                loc.rows.insert(loc.rows.end(), non_loc.rows.begin() + begin,
+                                non_loc.rows.begin() + end);
+                loc.cols.insert(loc.cols.end(), non_loc.cols.begin() + begin,
+                                non_loc.cols.begin() + end);
+                loc.mapping.insert(loc.mapping.end(),
+                                   non_loc.mapping.data() + begin,
+                                   non_loc.mapping.data() + end);
+                // TODO store from rank ie from which rank it came
+                // this is currently unused
+                loc.end.push_back(loc.rows.size());
+                loc.ranks.push_back(non_loc.ranks[i]);
+            } else {
+                begins.push_back(rows.size());
+                rows.insert(rows.end(), non_loc.rows.begin() + begin,
+                            non_loc.rows.begin() + end);
+                cols.insert(cols.end(), non_loc.cols.begin() + begin,
+                            non_loc.cols.begin() + end);
+                // NOTE we don't do anything with non local mapping
+                ldu_mapping.insert(ldu_mapping.end(),
+                                   non_loc.cols.begin() + begin,
+                                   non_loc.cols.begin() + end);
+                ends.push_back(rows.size());
+                ranks.push_back(get_owner_rank(non_loc.ranks[i]));
+            }
+        }
+        LOG_1(verbose_, "done build non local interfaces")
+        return std::make_pair(
+            SparsityPatternVector{rows, cols, ldu_mapping, begins, ends, ranks},
+            is_local);
+    }
+
+
+    /*
+    ** returns the the new sparsity patterns (local and non-local) and a vector
+    ** tracking the interfaces. This contains a pair where the first entry
+    *(bool)
+    ** signals whether this is a new local interface (no communication), and the
+    ** second entry (label) tracks the original rank of the interface
+    */
+    std::tuple<std::shared_ptr<SparsityPattern>,
+               std::shared_ptr<SparsityPattern>,
+               std::vector<std::pair<bool, label>>>
+    repartition_sparsity(
+        const ExecutorHandler &exec_handler,
+        std::shared_ptr<SparsityPattern> src_local_pattern,
+        std::shared_ptr<SparsityPattern> src_non_local_pattern) const
+    {
+        LOG_1(verbose_, "start repartition sparsity pattern")
+        // 1. obtain send recv sizes vector
+        // here we can reuse code from repartition_comm_pattern
+        //
+        // 2. initialize send and recv buffer
+        // if we keep interfaces separated row and col indices can be
+        //
+        // 3. ldu mapping needs to be update?
+        // 4. dim needs to be updated
+        // this is not implemented yet, but we can't fail here for
+        // debugging reasons
+        auto exec = exec_handler.get_ref_exec();
+        auto comm = *exec_handler.get_communicator().get();
+        label rank = get_rank(exec_handler);
+        label owner_rank = get_owner_rank(exec_handler);
+        label ranks_per_gpu = ranks_per_gpu_;
+        // TODO dont copy
+        // if (ranks_per_gpu == 1) {
+        //     std::vector<std::pair<bool, label>> ret;
+        //     for (auto comm_rank : src_non_local_pattern->rank) {
+        //         ret.emplace_back(false, rank);
+        //     }
+        //     return std::make_tuple<std::shared_ptr<SparsityPattern>,
+        //                            std::shared_ptr<SparsityPattern>,
+        //                            std::vector<std::pair<bool, label>>>(
+        //         std::make_shared<SparsityPattern>(src_local_pattern),
+        //         std::make_shared<SparsityPattern>(src_non_local_pattern),
+        //         std::move(ret));
+        // }
+
+        auto local_comm_pattern = compute_send_recv_counts(
+            exec_handler, ranks_per_gpu, src_local_pattern->size_);
+
+
+        auto offset = orig_partition_->get_range_bounds()[rank] -
+                      orig_partition_->get_range_bounds()[owner_rank];
+
+        auto gather_closure = [&](auto &comm_pattern, auto &data,
+                                  label offset) {
+            return gather_to_owner(exec_handler, comm_pattern, data.get_size(),
+                                   data.get_data(), offset);
+        };
+
+        SparsityPatternVector merged_local{
+            gather_closure(local_comm_pattern, src_local_pattern->row_idxs,
+                           offset),
+            gather_closure(local_comm_pattern, src_local_pattern->col_idxs,
+                           offset),
+            gather_closure(local_comm_pattern, src_local_pattern->ldu_mapping,
+                           0),
+            std::vector<label>{0},
+            std::vector<label>{0},
+            std::vector{rank}};
+        merged_local.end[0] = merged_local.rows.size();
+
+        if (is_owner(exec_handler)) {
+            make_ldu_mapping_consecutive(
+                local_comm_pattern, merged_local.mapping, rank, ranks_per_gpu);
+        }
+
+        label rows =
+            (is_owner(exec_handler)) ? merged_local.rows.back() + 1 : 0;
+        gko::dim<2> merged_local_dim{rows, rows};
+
+        auto non_local_comm_pattern = compute_send_recv_counts(
+            exec_handler, ranks_per_gpu, src_non_local_pattern->size_);
+
+        std::vector<label> spans_begin;
+        std::vector<label> spans_end;
+        auto span_comm_pattern = compute_send_recv_counts(
+            exec_handler, ranks_per_gpu,
+            src_non_local_pattern->interface_spans.size());
+        for (auto elem : src_non_local_pattern->interface_spans) {
+            spans_begin.push_back(elem.begin);
+            spans_end.push_back(elem.end);
+        }
+
+        // the non local cols are in local idx of other side
+        // thus we need the new offset of the other side
+        // NOTE TODO this modifies src_non_local_pattern better do the offsets
+        // on the owner side after gathering
+        for (label i = 0; i < src_non_local_pattern->rank.size(); i++) {
+            auto comm_rank = src_non_local_pattern->rank[i];
+            auto [begin, end] = src_non_local_pattern->interface_spans[i];
+            label local_offset =
+                orig_partition_->get_range_bounds()[comm_rank] -
+                orig_partition_->get_range_bounds()[owner_rank];
+            auto *data = src_non_local_pattern->col_idxs.get_data() + begin;
+            auto size = end - begin;
+            std::transform(data, data + size, data,
+                           [&](label idx) { return idx + local_offset; });
+        }
+
+        SparsityPatternVector merged_non_local{
+            gather_closure(non_local_comm_pattern,
+                           src_non_local_pattern->row_idxs, offset),
+            gather_closure(non_local_comm_pattern,
+                           src_non_local_pattern->col_idxs, 0),
+            gather_closure(non_local_comm_pattern,
+                           src_non_local_pattern->row_idxs, 0),
+            gather_to_owner(exec_handler, span_comm_pattern, spans_begin.size(),
+                            spans_begin.data()),
+            gather_to_owner(exec_handler, span_comm_pattern, spans_end.size(),
+                            spans_end.data()),
+            gather_to_owner(exec_handler, span_comm_pattern,
+                            src_non_local_pattern->rank.size(),
+                            src_non_local_pattern->rank.data())};
+
+        if (is_owner(exec_handler)) {
+            make_begin_and_ends_consecutive(
+                span_comm_pattern, merged_non_local.begin, merged_non_local.end,
+                rank, ranks_per_gpu);
+        }
+
+        auto [gathered_non_local, is_local] = build_non_local_interfaces(
+            exec_handler, merged_local, merged_non_local);
+
+        // build vector with locality information
+        std::vector<std::pair<bool, label>> locality;
+        label ctr{0};
+
+        if (is_owner(exec_handler)) {
+            auto recv_counts = std::get<1>(span_comm_pattern);
+            for (int i = 0; i < ranks_per_gpu; i++) {
+                label comm_rank = rank + i;
+                label num_elems = recv_counts[comm_rank];
+                for (int j = 0; j < num_elems; j++) {
+                    locality.emplace_back(is_local[ctr], comm_rank);
+                    ctr++;
+                }
+            }
+        }
+        // All interfaces that are in is_local need to be accounted for
+        ASSERT_EQ(ctr, is_local.size());
+
+        // TODO FIXME the correct way would be to check the communication
+        // pattern where a particular face is in the send idxs
+        // since we have already the row id of the other side
+        // it should be doable. Or alternatively we know that we
+        // keep an interface together thus we can just count the idx up to the
+        // size. But we have to make sure that the interfaces are in the same
+        // order on both communication sides.
+        for (int i = 0; i < gathered_non_local.cols.size(); i++) {
+            gathered_non_local.cols[i] = i;
+        }
+
+        LOG_1(verbose_, "done repartition sparsity pattern")
+        if (is_owner(exec_handler)) {
+            auto new_local_spars_pattern = std::make_shared<SparsityPattern>(
+                exec, merged_local_dim, merged_local);
+
+            auto new_non_local_spars_pattern =
+                std::make_shared<SparsityPattern>(
+                    exec,
+                    gko::dim<2>{merged_local_dim[0],
+                                gathered_non_local.rows.size()},
+                    gathered_non_local);
+
+            return std::make_tuple<std::shared_ptr<SparsityPattern>,
+                                   std::shared_ptr<SparsityPattern>,
+                                   std::vector<std::pair<bool, label>>>(
+                std::move(new_local_spars_pattern),
+                std::move(new_non_local_spars_pattern), std::move(locality));
+        } else {
+            auto new_local_spars_pattern =
+                std::make_shared<SparsityPattern>(exec);
+
+            auto new_non_local_spars_pattern =
+                std::make_shared<SparsityPattern>(exec);
+
+            // FIXME
+            // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
+            // comm.rank() <<
+            //     " merged_local dim " <<
+            new_local_spars_pattern->dim =
+                gko::dim<2>{0, 0};  // merged_local_dim;
+            new_non_local_spars_pattern->dim =
+                gko::dim<2>{0, 0};  // merged_local_dim;
+
+
+            return std::make_tuple<std::shared_ptr<SparsityPattern>,
+                                   std::shared_ptr<SparsityPattern>,
+                                   std::vector<std::pair<bool, label>>>(
+                std::move(new_local_spars_pattern),
+                std::move(new_non_local_spars_pattern),
+                std::vector<std::pair<bool, label>>{});
+        }
+        // TODO Some old documentation see what is still needed
+        // add interfaces
+        // to add interfaces we go through the non_local_sparsity
+        // pattern and check if interface is still non_local
+        //
+        // iterate interfaces and send to owner.
+        // on owner decide to move to local or non_local
+        // for sending to owner we can also use the all_to_all_v
+        // approach we than have
+        //
+        // row [ 4 , 8,  12 | 16 , 17, 18 ] <- local row
+        // col [ 1 , 2, 3 | 1, 2, 3 ] <- just the interface ctr
+        // from repartition_comm_pattern we could get
+        //
+        // or we split local / non-local first
+        // for this we could store ranges when computing
+        // repartition_comm_pattern
+        //
+    }
+
+    std::shared_ptr<const CommunicationPattern> repartition_comm_pattern(
+        const ExecutorHandler &exec_handler,
+        std::shared_ptr<const CommunicationPattern> src_comm_pattern,
+        std::shared_ptr<
+            const gko::experimental::distributed::Partition<label, label>>
+            partition) const
+    {
+        using comm_size_type = label;
+        auto exec = src_comm_pattern->target_ids.get_executor();
+        auto comm = src_comm_pattern->get_comm();
+
+        if (ranks_per_gpu_ == 1) {
+            std::cout << __FILE__ << ":" << __LINE__ << " rank " << comm.rank()
+                      << *src_comm_pattern.get() << "\n";
+            return src_comm_pattern;
+        }
+
+
+        label rank = comm.rank();
+        bool owner = is_owner(exec_handler);
+
+        // Step 1. Check if communication partner is non-local after
+        // repartitioning. If it is non-local we keep it. Otherwise local
+        // communication can be discarded. Here non-local means:
+        // the communication target proc id is != repartition proc id
+        std::vector<label> target_ids{};
+        std::vector<label> target_sizes{};
+        std::vector<std::pair<gko::array<label>, comm_size_type>> send_idxs;
+        label communication_partner = src_comm_pattern->target_ids.get_size();
+        for (int i = 0; i < communication_partner; i++) {
+            label target_id = src_comm_pattern->target_ids.get_const_data()[i];
+            if (!reparts_to_local(exec_handler, target_id)) {
+                // communication pattern is non local, hence we keep it
+                // after repartitioning we now have to communicate with a
+                // different rank store new owner rank to which repart_target_*
+                // needs to be send to
+                target_ids.push_back(get_owner_rank(target_id));
+                target_sizes.push_back(
+                    src_comm_pattern->target_sizes.get_const_data()[i]);
+                send_idxs.push_back(src_comm_pattern->send_idxs[i]);
+            }
+        }
+
+        // send all remaining non local ids and sizes to the new
+        // owner rank
+        auto comm_pattern = compute_send_recv_counts(
+            exec_handler, ranks_per_gpu_, target_ids.size());
+
+        auto gathered_target_ids = gather_to_owner(
+            exec_handler, comm_pattern, target_ids.size(), target_ids.data());
+
+        auto gathered_target_sizes =
+            gather_to_owner(exec_handler, comm_pattern, target_sizes.size(),
+                            target_sizes.data());
+
+        // next the send_ixs need to be updated we send them piecewise since
+        // the send_idxs are a vector of gko::arrays
+        if (owner) {
+            auto recv_counts = std::get<1>(comm_pattern);
+            label owner_recv_counts = recv_counts[rank];
+            // retrieved from i-th neighbor
+            for (int i = 1; i < ranks_per_gpu_; i++) {
+                // how many gko::arrays to with send_indexes to receive
+                // from i-th neighbor
+                label recv_count = recv_counts[rank + i];
+                for (int j = 0; j < recv_count; j++) {
+                    auto target_size =
+                        gathered_target_sizes[j + owner_recv_counts];
+                    std::vector<label> recv_buffer(target_size);
+
+                    // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
+                    // rank
+                    //           << " recv no  " << j
+                    //           << " recv from  " << rank + i
+                    //           << " recv_count " << recv_count
+                    //           << " target_size " << target_size
+                    //           << " \n";
+                    comm.recv(exec, recv_buffer.data(), target_size, rank + i,
+                              rank);
+
+                    // the new offset is
+                    auto offset = partition->get_range_bounds()[rank + i] -
+                                  partition->get_range_bounds()[rank];
+
+                    std::transform(recv_buffer.begin(), recv_buffer.end(),
+                                   recv_buffer.begin(),
+                                   [&](label idx) { return idx + offset; });
+
+                    auto target_id = gathered_target_ids[j + owner_recv_counts];
+                    send_idxs.push_back({gko::array<label>{
+                                             exec,
+                                             recv_buffer.begin(),
+                                             recv_buffer.end(),
+                                         },
+                                         target_id});
+                }
+            }
+        } else {
+            label owner = get_owner_rank(exec_handler);
+            auto send_count = std::get<0>(comm_pattern)[owner];
+            for (int i = 0; i < send_count; i++) {
+                auto send_buffer = send_idxs[i].first;
+                // std::cout << __FILE__ << ":" << __LINE__ << " rank " << rank
+                //           << " send to owner " << owner
+                //           << " count " << i << " of "  <<  send_count
+                //           << " send_buffer.size " << send_buffer.get_size()
+                //           << " \n";
+                comm.send(exec, send_buffer.get_const_data(),
+                          send_buffer.get_size(), owner, owner);
+            }
+        }
+
+        // clear communcation neighbors on non owning rank
+        if (!owner) {
+            // TODO NOTE should it be gathered_target_ids and is it needed?
+            target_ids.clear();
+            target_sizes.clear();
+            send_idxs.clear();
+        }
+
+
+        // early return if no communication partners are left
+        if (gathered_target_ids.size() == 0) {
+            return std::make_shared<CommunicationPattern>(
+                exec_handler,
+                gko::array<comm_size_type>{exec, target_ids.begin(),
+                                           target_ids.end()},
+                gko::array<comm_size_type>{exec, target_sizes.begin(),
+                                           target_sizes.end()},
+
+                send_idxs);
+        }
+
+        // merge communication
+        // we now might have communication partners multiple times, thus we can
+        // merge them sort repart_target_ids and use the sorting for
+        // target_sizes and repart_send_idxs
+        std::vector<label> merged_target_ids{};
+        std::vector<label> merged_target_sizes{};
+        std::vector<std::vector<label>> merged_send_idxs{};
+        auto p = sort_permutation(gathered_target_ids,
+                                  [](label a, label b) { return a < b; });
+
+        target_sizes = apply_permutation(gathered_target_sizes, p);
+        target_ids = apply_permutation(gathered_target_ids, p);
+        send_idxs = apply_permutation(send_idxs, p);
+
+        // Step 4.
+        // Merge communication pattern with corresponding neighbours
+        merged_target_ids.push_back(target_ids[0]);
+        merged_target_sizes.push_back(target_sizes[0]);
+        merged_send_idxs.emplace_back(std::vector<label>(
+            send_idxs[0].first.get_data(),
+            send_idxs[0].first.get_data() + send_idxs[0].first.get_size()));
+
+
+        for (int i = 1; i < target_ids.size(); i++) {
+            // communicates with same target rank
+            // thus we have only have to adapt the number
+            // of elemts and the send_ixs
+            auto *send_idx_begin = send_idxs[i].first.get_data();
+            auto *send_idx_end =
+                send_idxs[i].first.get_data() + send_idxs[i].first.get_size();
+            if (target_ids[i] == merged_target_ids.back()) {
+                merged_target_sizes.back() += target_sizes[i];
+                merged_send_idxs.back().insert(merged_send_idxs.back().end(),
+                                               send_idx_begin, send_idx_end);
+            } else {
+                merged_target_ids.push_back(target_ids[i]);
+                merged_target_sizes.push_back(target_sizes[i]);
+                merged_send_idxs.emplace_back(
+                    std::vector<label>(send_idx_begin, send_idx_end));
+            }
+        }
+
+        // recompute send_idxs
+        send_idxs.clear();
+
+        for (int i = 0; i < merged_target_ids.size(); i++) {
+            label target_id = merged_target_ids[i];
+            send_idxs.emplace_back(std::pair<gko::array<label>, comm_size_type>{
+                gko::array<label>{exec, merged_send_idxs[i].begin(),
+                                  merged_send_idxs[i].end()},
+                target_id});
+        }
+
+        return std::make_shared<CommunicationPattern>(
+            exec_handler,
+            gko::array<comm_size_type>{exec, merged_target_ids.begin(),
+                                       merged_target_ids.end()},
+            gko::array<comm_size_type>{exec, merged_target_sizes.begin(),
+                                       merged_target_sizes.end()},
+            send_idxs);
+    }
+};

From 5a447aa74704f68200be4f6df899269111389a7f Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Fri, 9 Aug 2024 09:47:10 +0200
Subject: [PATCH 02/67] move implementation to body

---
 CMakeLists.txt              |  1 +
 include/OGL/Repartitioner.H | 50 +-----------------------------
 src/Repartitioner.C         | 61 +++++++++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 49 deletions(-)
 create mode 100644 src/Repartitioner.C

diff --git a/CMakeLists.txt b/CMakeLists.txt
index fefa94217..3a58f1660 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -222,6 +222,7 @@ target_sources(
           src/DevicePersistent/DeviceIdGuard/DeviceIdGuard.C
           src/DevicePersistent/ExecutorHandler/ExecutorHandler.C
           src/MatrixWrapper/HostMatrix.C
+          src/Repartitioner.C
           src/Solver/CG/GKOCG.C
           src/Solver/BiCGStab/GKOBiCGStab.C
           src/Solver/GMRES/GKOGMRES.C
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index 7bab1b863..b75f77396 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -136,55 +136,7 @@ public:
     std::pair<SparsityPatternVector, std::vector<bool>>
     build_non_local_interfaces(const ExecutorHandler &exec_handler,
                                SparsityPatternVector &loc,
-                               const SparsityPatternVector &non_loc) const
-    {
-        LOG_1(verbose_, "start build non local interfaces")
-        std::vector<label> rows, cols, ldu_mapping, ranks, begins, ends;
-        std::vector<bool> is_local;
-        label merged_ranks_size = non_loc.ranks.size();
-
-        for (label i = 0; i < merged_ranks_size; i++) {
-            // these are the begin ends before merging they need to be
-            // offseted;
-            auto begin = non_loc.begin[i];
-            auto end = non_loc.end[i];
-            bool local = reparts_to_local(exec_handler, non_loc.ranks[i]);
-            is_local.push_back(local);
-            if (local) {
-                // TODO depending on the interface simple non
-                // transforming and interpolating interfaces like now
-                // local processor interfaces could be merged
-                loc.begin.push_back(loc.rows.size());
-                loc.rows.insert(loc.rows.end(), non_loc.rows.begin() + begin,
-                                non_loc.rows.begin() + end);
-                loc.cols.insert(loc.cols.end(), non_loc.cols.begin() + begin,
-                                non_loc.cols.begin() + end);
-                loc.mapping.insert(loc.mapping.end(),
-                                   non_loc.mapping.data() + begin,
-                                   non_loc.mapping.data() + end);
-                // TODO store from rank ie from which rank it came
-                // this is currently unused
-                loc.end.push_back(loc.rows.size());
-                loc.ranks.push_back(non_loc.ranks[i]);
-            } else {
-                begins.push_back(rows.size());
-                rows.insert(rows.end(), non_loc.rows.begin() + begin,
-                            non_loc.rows.begin() + end);
-                cols.insert(cols.end(), non_loc.cols.begin() + begin,
-                            non_loc.cols.begin() + end);
-                // NOTE we don't do anything with non local mapping
-                ldu_mapping.insert(ldu_mapping.end(),
-                                   non_loc.cols.begin() + begin,
-                                   non_loc.cols.begin() + end);
-                ends.push_back(rows.size());
-                ranks.push_back(get_owner_rank(non_loc.ranks[i]));
-            }
-        }
-        LOG_1(verbose_, "done build non local interfaces")
-        return std::make_pair(
-            SparsityPatternVector{rows, cols, ldu_mapping, begins, ends, ranks},
-            is_local);
-    }
+                               const SparsityPatternVector &non_loc) const;
 
 
     /*
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
new file mode 100644
index 000000000..8b997f6d4
--- /dev/null
+++ b/src/Repartitioner.C
@@ -0,0 +1,61 @@
+// SPDX-FileCopyrightText: 2024 OGL authors
+//
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+
+
+  std::pair<SparsityPatternVector, std::vector<bool>>
+ Repartitioner::   build_non_local_interfaces(const ExecutorHandler &exec_handler,
+                               SparsityPatternVector &loc,
+                               const SparsityPatternVector &non_loc) const
+{
+
+        LOG_1(verbose_, "start build non local interfaces")
+        std::vector<label> rows, cols, ldu_mapping, ranks, begins, ends;
+        std::vector<bool> is_local;
+        label merged_ranks_size = non_loc.ranks.size();
+
+        for (label i = 0; i < merged_ranks_size; i++) {
+            // these are the begin ends before merging they need to be
+            // offseted;
+            auto begin = non_loc.begin[i];
+            auto end = non_loc.end[i];
+            bool local = reparts_to_local(exec_handler, non_loc.ranks[i]);
+            is_local.push_back(local);
+            if (local) {
+                // TODO depending on the interface simple non
+                // transforming and interpolating interfaces like now
+                // local processor interfaces could be merged
+                loc.begin.push_back(loc.rows.size());
+                loc.rows.insert(loc.rows.end(), non_loc.rows.begin() + begin,
+                                non_loc.rows.begin() + end);
+                loc.cols.insert(loc.cols.end(), non_loc.cols.begin() + begin,
+                                non_loc.cols.begin() + end);
+                loc.mapping.insert(loc.mapping.end(),
+                                   non_loc.mapping.data() + begin,
+                                   non_loc.mapping.data() + end);
+                // TODO store from rank ie from which rank it came
+                // this is currently unused
+                loc.end.push_back(loc.rows.size());
+                loc.ranks.push_back(non_loc.ranks[i]);
+            } else {
+                begins.push_back(rows.size());
+                rows.insert(rows.end(), non_loc.rows.begin() + begin,
+                            non_loc.rows.begin() + end);
+                cols.insert(cols.end(), non_loc.cols.begin() + begin,
+                            non_loc.cols.begin() + end);
+                // NOTE we don't do anything with non local mapping
+                ldu_mapping.insert(ldu_mapping.end(),
+                                   non_loc.cols.begin() + begin,
+                                   non_loc.cols.begin() + end);
+                ends.push_back(rows.size());
+                ranks.push_back(get_owner_rank(non_loc.ranks[i]));
+            }
+        }
+        LOG_1(verbose_, "done build non local interfaces")
+        return std::make_pair(
+            SparsityPatternVector{rows, cols, ldu_mapping, begins, ends, ranks},
+            is_local);
+
+
+}

From ae2d38aeaca2920771fbcc458fa2f85f8c48d063 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Thu, 15 Aug 2024 20:28:28 +0200
Subject: [PATCH 03/67] start adding repartitioner unit tests

---
 include/OGL/Repartitioner.H |  6 ++-
 unitTests/CMakeLists.txt    |  1 +
 unitTests/Repartitioner.C   | 74 +++++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+), 2 deletions(-)
 create mode 100644 unitTests/Repartitioner.C

diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index b75f77396..84c36e119 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -30,7 +30,10 @@ std::vector<label> sort_permutation(const std::vector<T> &vec, Compare compare)
     return p;
 }
 
-// TODO this should be persistent
+/** @class Collects functions to repartion communication patterns 
+ * and matrices. Here repartitioning refers to changing a given partitioning on n 
+ * ranks to m ranks with n>=m
+ */
 class Repartitioner {
 private:
     using partition = gko::experimental::distributed::Partition<label, label>;
@@ -45,7 +48,6 @@ private:
 
     std::shared_ptr<const partition> orig_partition_;
 
-
 public:
     Repartitioner(label size, label ranks_per_gpu, label verbose,
                   const ExecutorHandler &exec_handler)
diff --git a/unitTests/CMakeLists.txt b/unitTests/CMakeLists.txt
index 6c40d60f5..1e3852659 100644
--- a/unitTests/CMakeLists.txt
+++ b/unitTests/CMakeLists.txt
@@ -31,5 +31,6 @@ macro(ogl_unit_test TEST FILE)
 endmacro()
 
 ogl_unit_test(CommunicationPattern "CommunicationPattern.C")
+ogl_unit_test(CommunicationPattern "Repartitioner.C")
 
 add_subdirectory(MatrixWrapper)
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
new file mode 100644
index 000000000..75c4625be
--- /dev/null
+++ b/unitTests/Repartitioner.C
@@ -0,0 +1,74 @@
+// SPDX-FileCopyrightText: 2024 OGL authors
+//
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "OGL/CommunicationPattern.H"
+
+#include "gtest/gtest.h"
+
+#include "mpi.h"
+
+#include <cstdlib>
+
+//---------------------------------------------
+// some_header.h
+extern int my_argc;
+extern char **my_argv;
+// eof
+//---------------------------------------------
+
+//---------------------------------------------
+// main.cpp
+int my_argc;
+char **my_argv;
+
+class RepartitionerEnvironment : public testing::Environment {
+public:
+    void SetUp()
+    {
+        args_ = std::make_shared<Foam::argList>(my_argc, my_argv);
+        time = std::make_shared<Foam::Time>("controlDict", *args_.get());
+        db = std::make_shared<Foam::objectRegistry>(*time.get());
+        Foam::dictionary dict;
+        dict.add("executor", "reference");
+        exec = std::make_shared<ExecutorHandler>(time->thisDb(), dict, "dummy",
+                                                 true);
+
+        auto comm = exec->get_gko_mpi_host_comm();
+        if (comm->size() < 2) {
+            std::cout << "At least 2 CPU processes should be used!"
+                      << std::endl;
+            std::abort();
+        }
+        // delete listener on ranks != 0
+        // to clean up output
+        ::testing::TestEventListeners &listeners =
+            ::testing::UnitTest::GetInstance()->listeners();
+        if (comm->rank() != 0) {
+            delete listeners.Release(listeners.default_result_printer());
+        }
+    }
+
+    std::shared_ptr<Foam::argList> args_;
+    std::shared_ptr<const Foam::Time> time;
+    std::shared_ptr<Foam::objectRegistry> db;
+    std::shared_ptr<ExecutorHandler> exec;
+};
+
+const testing::Environment *global_env =
+    AddGlobalTestEnvironment(new CommunicationPatternEnvironment);
+
+
+TEST(CommunicationPattern, can_create_repartitioner)
+{
+    // Arrange
+    auto exec = ((CommunicationPatternEnvironment *)global_env)->exec;
+    auto repartitioner = Repartitioner(0, 1, 0, exec);
+
+    // Act
+
+    // Assert
+    EXPECT_EQ(owner_rank, 0);
+}
+
+

From 5a7409b00e74795ecad987aef9a1209daf336046 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Thu, 15 Aug 2024 20:55:46 +0200
Subject: [PATCH 04/67] small fixes

---
 include/OGL/Repartitioner.H | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index 84c36e119..253ba75e0 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -7,7 +7,6 @@
 #include <ginkgo/ginkgo.hpp>
 
 #include "MatrixWrapper/CommunicationPattern/CommunicationPattern.H"
-#include "MatrixWrapper/LDUMatrix/HostMatrix.H"
 
 template <typename T>
 std::vector<T> apply_permutation(const std::vector<T> vec,
@@ -38,7 +37,7 @@ class Repartitioner {
 private:
     using partition = gko::experimental::distributed::Partition<label, label>;
 
-    const label size_;  //! Size (n matrix rows ) before repartitioning
+    const label size_;  //! Size (n matrix rows) before repartitioning
 
     const label repart_size_;  //! Size after repartitioning
 

From 004c4d42a45bda914133794384f8a6a8701ccf5b Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Fri, 16 Aug 2024 09:18:43 +0200
Subject: [PATCH 05/67] wip repartitioner implementations

- add make ldu consecutive
---
 include/OGL/MatrixWrapper/SparsityPattern.H | 45 ++++++++++
 include/OGL/Repartitioner.H                 | 74 ++++++++--------
 src/Repartitioner.C                         | 96 ++++++++++-----------
 unitTests/Repartitioner.C                   |  7 +-
 4 files changed, 133 insertions(+), 89 deletions(-)

diff --git a/include/OGL/MatrixWrapper/SparsityPattern.H b/include/OGL/MatrixWrapper/SparsityPattern.H
index e462ddeb7..edc2b2cd5 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.H
+++ b/include/OGL/MatrixWrapper/SparsityPattern.H
@@ -6,8 +6,39 @@
 
 #include <vector>
 
+#include "../common.H"
+
 namespace Foam {
 
+/* @brief based on a comm_pattern consecutive ldu_mappings are computed
+ *
+ * After merging sparsity patterns during repartitioning the ldu mapping are
+ * not consecutive ie they could look like [0 1 2 3 | 0 1 2 3 | 0 1], where
+ * | is the former rank boundary based on the comm pattern we compute a new
+ * mapping as  [ 0 1 2 3 | 4 5 6 7| 8 9 ] where the following offsets [ 0 |
+ * 4 | 8 ] based on the recv_counts are used.
+ *
+ * */
+void make_ldu_mapping_consecutive(
+    const AllToAllPattern & comm_pattern,
+    gko::array<label> &ldu_mapping, label rank, label ranks_per_gpu)
+    {
+        // TODO move to SparsityPattern.C
+    label ldu_offset = 0;
+    auto *data = ldu_mapping.get_data();
+    // TODO check if ldu_mapping is on host exec
+
+    for (label i = 0; i < ranks_per_gpu; i++) {
+        auto size = comm_pattern.recv_counts[i];
+        std::transform(data + ldu_offset, data + ldu_offset + size,
+                       data + ldu_offset,
+                       [&](label idx) { return idx + ldu_offset; });
+        ldu_offset += size;
+    }
+
+
+    }
+
 
 /* The SparsityPattern holds row and column index data using gko::arrays.
  * This struct is used for easy generation of (distributed) ginkgo matrices.
@@ -17,6 +48,19 @@ namespace Foam {
  *
  * */
 struct SparsityPattern {
+    /* constructor from vectors, assumes a single interface */
+    SparsityPattern(
+        std::vector<label> rows,
+        std::vector<label> cols,
+        std::vector<label> mapping,
+        std::vector<label> rank
+)
+        : num_nnz(rows.size())
+    {
+        // NOT fully implementented
+        FatalErrorInFunction << "Not implemented" << abort(FatalError);
+    }
+
     SparsityPattern(std::shared_ptr<const gko::Executor> exec, const label size)
         : num_nnz(size),
           row_idxs{exec, static_cast<gko::size_type>(size)},
@@ -56,6 +100,7 @@ struct SparsityPattern {
         rank = ranks;
     }
 
+    // TODO num_nnz makes no sense ->nnz;
     const label num_nnz;
 
     mutable gko::Array<label> row_idxs;
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index 253ba75e0..684432f6d 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -6,7 +6,10 @@
 
 #include <ginkgo/ginkgo.hpp>
 
-#include "MatrixWrapper/CommunicationPattern/CommunicationPattern.H"
+#include "fvCFD.H"
+
+#include "CommunicationPattern.H"
+#include "MatrixWrapper/SparsityPattern.H"
 
 template <typename T>
 std::vector<T> apply_permutation(const std::vector<T> vec,
@@ -29,9 +32,9 @@ std::vector<label> sort_permutation(const std::vector<T> &vec, Compare compare)
     return p;
 }
 
-/** @class Collects functions to repartion communication patterns 
- * and matrices. Here repartitioning refers to changing a given partitioning on n 
- * ranks to m ranks with n>=m
+/** @class Collects functions to repartion communication patterns
+ * and matrices. Here repartitioning refers to changing a given partitioning on
+ * n ranks to m ranks with n>=m
  */
 class Repartitioner {
 private:
@@ -59,7 +62,7 @@ public:
               gko::experimental::distributed::build_partition_from_local_size<
                   label, label>(exec_handler.get_ref_exec(),
                                 *exec_handler.get_communicator().get(),
-                                size))){};
+                                size))) {};
 
     /* returns the owner rank for a given rank */
     label get_owner_rank(label rank) const
@@ -96,6 +99,11 @@ public:
 
     label get_ranks_per_gpu() const { return ranks_per_gpu_; }
 
+    /* @brief computes the size after repartitioning
+     *
+     * Given a local size (can represent nrows or nnzs) the new size of
+     * this rank after repartitioning gets computed
+     */
     static label compute_repart_size(label local_size, label ranks_per_gpu,
                                      const ExecutorHandler &exec_handler)
     {
@@ -103,22 +111,18 @@ public:
             return local_size;
         }
 
-        CommCounts comm_pattern =
-            compute_send_recv_counts(exec_handler, ranks_per_gpu, 1);
+        auto all_to_all_pattern = compute_gather_to_owner_counts(
+            exec_handler, ranks_per_gpu, local_size);
 
-        std::vector<label> data{local_size};
-
-        auto local_sizes =
-            gather_to_owner(exec_handler, comm_pattern, 1, data.data());
-
-        return std::accumulate(local_sizes.begin(), local_sizes.end(), 0);
+        return all_to_all_pattern.recv_offsets.back();
     }
 
     label get_repart_size() const { return repart_size_; }
 
     gko::dim<2> get_repart_dim() const
     {
-        return gko::dim<2>{repart_size_, repart_size_};
+        return gko::dim<2>{static_cast<gko::size_type>(repart_size_),
+                           static_cast<gko::size_type>(repart_size_)};
     }
 
     std::shared_ptr<const partition> get_orig_partition() const
@@ -129,15 +133,14 @@ public:
 
     /* @brief given received interfaces this function sorts into local and non
      *local
-     * returns: SparsityPatternVector storing new non_local sparsity pattern,
+     * returns: SparsityPattern storing new non_local sparsity pattern,
      *  and a pair of locality information, where the first bool stores if the
      *  interface performs a local communication and the second bool if the
      *  the interface was originally from this rank
      */
-    std::pair<SparsityPatternVector, std::vector<bool>>
-    build_non_local_interfaces(const ExecutorHandler &exec_handler,
-                               SparsityPatternVector &loc,
-                               const SparsityPatternVector &non_loc) const;
+    std::pair<SparsityPattern, std::vector<bool>> build_non_local_interfaces(
+        const ExecutorHandler &exec_handler, SparsityPattern &loc,
+        const SparsityPattern &non_loc) const;
 
 
     /*
@@ -185,8 +188,8 @@ public:
         //         std::move(ret));
         // }
 
-        auto local_comm_pattern = compute_send_recv_counts(
-            exec_handler, ranks_per_gpu, src_local_pattern->size_);
+        auto local_comm_pattern = compute_gather_to_owner_counts(
+            exec_handler, ranks_per_gpu, src_local_pattern->num_nnz);
 
 
         auto offset = orig_partition_->get_range_bounds()[rank] -
@@ -194,40 +197,41 @@ public:
 
         auto gather_closure = [&](auto &comm_pattern, auto &data,
                                   label offset) {
-            return gather_to_owner(exec_handler, comm_pattern, data.get_size(),
-                                   data.get_data(), offset);
+            return gather_labels_to_owner(exec_handler, comm_pattern,
+                                          data.get_data(), data.get_size(),
+                                          offset);
         };
 
-        SparsityPatternVector merged_local{
+        SparsityPattern merged_local{
             gather_closure(local_comm_pattern, src_local_pattern->row_idxs,
                            offset),
             gather_closure(local_comm_pattern, src_local_pattern->col_idxs,
                            offset),
             gather_closure(local_comm_pattern, src_local_pattern->ldu_mapping,
                            0),
-            std::vector<label>{0},
-            std::vector<label>{0},
             std::vector{rank}};
-        merged_local.end[0] = merged_local.rows.size();
 
         if (is_owner(exec_handler)) {
             make_ldu_mapping_consecutive(
-                local_comm_pattern, merged_local.mapping, rank, ranks_per_gpu);
+                local_comm_pattern, merged_local.ldu_mapping, rank, ranks_per_gpu);
         }
 
+	auto get_back = [](const gko::array<label> & in){return *(in.get_const_data() + in.get_size());};
+
         label rows =
-            (is_owner(exec_handler)) ? merged_local.rows.back() + 1 : 0;
-        gko::dim<2> merged_local_dim{rows, rows};
+            (is_owner(exec_handler)) ? get_back(merged_local.row_idxs) + 1 : 0;
+        gko::dim<2> merged_local_dim{static_cast<gko::size_type>(rows),
+                           static_cast<gko::size_type>(rows)};
 
-        auto non_local_comm_pattern = compute_send_recv_counts(
+        auto non_local_comm_pattern = compute_gather_to_owner_counts(
             exec_handler, ranks_per_gpu, src_non_local_pattern->size_);
 
         std::vector<label> spans_begin;
         std::vector<label> spans_end;
-        auto span_comm_pattern = compute_send_recv_counts(
+        auto span_comm_pattern = compute_gather_to_owner_counts(
             exec_handler, ranks_per_gpu,
             src_non_local_pattern->interface_spans.size());
-        for (auto elem : src_non_local_pattern->interface_spans) {
+        for (auto elem : src_non_local_pattern->spans) {
             spans_begin.push_back(elem.begin);
             spans_end.push_back(elem.end);
         }
@@ -248,7 +252,7 @@ public:
                            [&](label idx) { return idx + local_offset; });
         }
 
-        SparsityPatternVector merged_non_local{
+        SparsityPattern merged_non_local{
             gather_closure(non_local_comm_pattern,
                            src_non_local_pattern->row_idxs, offset),
             gather_closure(non_local_comm_pattern,
@@ -407,7 +411,7 @@ public:
 
         // send all remaining non local ids and sizes to the new
         // owner rank
-        auto comm_pattern = compute_send_recv_counts(
+        auto comm_pattern = compute_gather_to_owner_counts(
             exec_handler, ranks_per_gpu_, target_ids.size());
 
         auto gathered_target_ids = gather_to_owner(
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index 8b997f6d4..719ac3c39 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -2,60 +2,56 @@
 //
 // SPDX-License-Identifier: GPL-3.0-or-later
 
+#include "OGL/Repartitioner.H"
 
-
-  std::pair<SparsityPatternVector, std::vector<bool>>
- Repartitioner::   build_non_local_interfaces(const ExecutorHandler &exec_handler,
-                               SparsityPatternVector &loc,
-                               const SparsityPatternVector &non_loc) const
+std::pair<SparsityPattern, std::vector<bool>>
+Repartitioner::build_non_local_interfaces(const ExecutorHandler &exec_handler,
+                                          SparsityPattern &loc,
+                                          const SparsityPattern &non_loc) const
 {
+    LOG_1(verbose_, "start build non local interfaces")
+    std::vector<label> rows, cols, ldu_mapping, ranks, begins, ends;
+    std::vector<bool> is_local;
+    label merged_ranks_size = non_loc.ranks.size();
 
-        LOG_1(verbose_, "start build non local interfaces")
-        std::vector<label> rows, cols, ldu_mapping, ranks, begins, ends;
-        std::vector<bool> is_local;
-        label merged_ranks_size = non_loc.ranks.size();
-
-        for (label i = 0; i < merged_ranks_size; i++) {
-            // these are the begin ends before merging they need to be
-            // offseted;
-            auto begin = non_loc.begin[i];
-            auto end = non_loc.end[i];
-            bool local = reparts_to_local(exec_handler, non_loc.ranks[i]);
-            is_local.push_back(local);
-            if (local) {
-                // TODO depending on the interface simple non
-                // transforming and interpolating interfaces like now
-                // local processor interfaces could be merged
-                loc.begin.push_back(loc.rows.size());
-                loc.rows.insert(loc.rows.end(), non_loc.rows.begin() + begin,
-                                non_loc.rows.begin() + end);
-                loc.cols.insert(loc.cols.end(), non_loc.cols.begin() + begin,
-                                non_loc.cols.begin() + end);
-                loc.mapping.insert(loc.mapping.end(),
-                                   non_loc.mapping.data() + begin,
-                                   non_loc.mapping.data() + end);
-                // TODO store from rank ie from which rank it came
-                // this is currently unused
-                loc.end.push_back(loc.rows.size());
-                loc.ranks.push_back(non_loc.ranks[i]);
-            } else {
-                begins.push_back(rows.size());
-                rows.insert(rows.end(), non_loc.rows.begin() + begin,
+    for (label i = 0; i < merged_ranks_size; i++) {
+        // these are the begin ends before merging they need to be
+        // offseted;
+        auto begin = non_loc.begin[i];
+        auto end = non_loc.end[i];
+        bool local = reparts_to_local(exec_handler, non_loc.ranks[i]);
+        is_local.push_back(local);
+        if (local) {
+            // TODO depending on the interface simple non
+            // transforming and interpolating interfaces like now
+            // local processor interfaces could be merged
+            loc.begin.push_back(loc.rows.size());
+            loc.rows.insert(loc.rows.end(), non_loc.rows.begin() + begin,
                             non_loc.rows.begin() + end);
-                cols.insert(cols.end(), non_loc.cols.begin() + begin,
+            loc.cols.insert(loc.cols.end(), non_loc.cols.begin() + begin,
                             non_loc.cols.begin() + end);
-                // NOTE we don't do anything with non local mapping
-                ldu_mapping.insert(ldu_mapping.end(),
-                                   non_loc.cols.begin() + begin,
-                                   non_loc.cols.begin() + end);
-                ends.push_back(rows.size());
-                ranks.push_back(get_owner_rank(non_loc.ranks[i]));
-            }
+            loc.mapping.insert(loc.mapping.end(),
+                               non_loc.mapping.data() + begin,
+                               non_loc.mapping.data() + end);
+            // TODO store from rank ie from which rank it came
+            // this is currently unused
+            loc.end.push_back(loc.rows.size());
+            loc.ranks.push_back(non_loc.ranks[i]);
+        } else {
+            begins.push_back(rows.size());
+            rows.insert(rows.end(), non_loc.rows.begin() + begin,
+                        non_loc.rows.begin() + end);
+            cols.insert(cols.end(), non_loc.cols.begin() + begin,
+                        non_loc.cols.begin() + end);
+            // NOTE we don't do anything with non local mapping
+            ldu_mapping.insert(ldu_mapping.end(), non_loc.cols.begin() + begin,
+                               non_loc.cols.begin() + end);
+            ends.push_back(rows.size());
+            ranks.push_back(get_owner_rank(non_loc.ranks[i]));
         }
-        LOG_1(verbose_, "done build non local interfaces")
-        return std::make_pair(
-            SparsityPatternVector{rows, cols, ldu_mapping, begins, ends, ranks},
-            is_local);
-
-
+    }
+    LOG_1(verbose_, "done build non local interfaces")
+    return std::make_pair(
+        SparsityPattern{rows, cols, ldu_mapping, begins, ends, ranks},
+        is_local);
 }
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 75c4625be..bc64a9255 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -62,13 +62,12 @@ const testing::Environment *global_env =
 TEST(CommunicationPattern, can_create_repartitioner)
 {
     // Arrange
+    label ranks_per_gpu = 1;
     auto exec = ((CommunicationPatternEnvironment *)global_env)->exec;
-    auto repartitioner = Repartitioner(0, 1, 0, exec);
-
-    // Act
+    auto repartitioner = Repartitioner(10, ranks_per_gpu, 0, *exec.get());
 
     // Assert
-    EXPECT_EQ(owner_rank, 0);
+    EXPECT_EQ(repartitioner.get_ranks_per_gpu(), 0);
 }
 
 

From 7f771e17b32b9cb1702c69c2b16a462bdaea23d6 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Fri, 16 Aug 2024 15:32:23 +0200
Subject: [PATCH 06/67] fix test name

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

diff --git a/unitTests/CMakeLists.txt b/unitTests/CMakeLists.txt
index 1e3852659..504c2c219 100644
--- a/unitTests/CMakeLists.txt
+++ b/unitTests/CMakeLists.txt
@@ -31,6 +31,6 @@ macro(ogl_unit_test TEST FILE)
 endmacro()
 
 ogl_unit_test(CommunicationPattern "CommunicationPattern.C")
-ogl_unit_test(CommunicationPattern "Repartitioner.C")
+ogl_unit_test(Repartitioner "Repartitioner.C")
 
 add_subdirectory(MatrixWrapper)

From fe2c9f820d82b89a05b8de74770c4b2acc0989d3 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Fri, 16 Aug 2024 20:28:31 +0200
Subject: [PATCH 07/67] make it compile

---
 CMakeLists.txt                              |   1 +
 include/OGL/DevicePersistent/Vector.H       |   8 +-
 include/OGL/MatrixWrapper/SparsityPattern.H |  31 +----
 include/OGL/Repartitioner.H                 | 144 +++++++++++---------
 include/OGL/common.H                        |   2 +-
 src/MatrixWrapper/HostMatrix.C              |  10 +-
 src/MatrixWrapper/SparsityPattern.C         |  22 +++
 src/Repartitioner.C                         |  86 ++++++------
 unitTests/Repartitioner.C                   |   6 +-
 9 files changed, 168 insertions(+), 142 deletions(-)
 create mode 100644 src/MatrixWrapper/SparsityPattern.C

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 3a58f1660..9621f37b9 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -222,6 +222,7 @@ target_sources(
           src/DevicePersistent/DeviceIdGuard/DeviceIdGuard.C
           src/DevicePersistent/ExecutorHandler/ExecutorHandler.C
           src/MatrixWrapper/HostMatrix.C
+	  src/MatrixWrapper/SparsityPattern.C
           src/Repartitioner.C
           src/Solver/CG/GKOCG.C
           src/Solver/BiCGStab/GKOBiCGStab.C
diff --git a/include/OGL/DevicePersistent/Vector.H b/include/OGL/DevicePersistent/Vector.H
index c6fc70d0f..7ac0ba4ab 100644
--- a/include/OGL/DevicePersistent/Vector.H
+++ b/include/OGL/DevicePersistent/Vector.H
@@ -154,8 +154,12 @@ public:
         // TODO why does it need its size
         // repartitioner size should be sufficient
         auto target_vector = dist_vec::create(exec_.get_ref_exec(), *comm.get(),
-                                              gko::dim<2>{global_size, 1},
-                                              gko::dim<2>{local_size, 1}, 1);
+                                              gko::dim<2>{
+  static_cast<gko::size_type>(
+					      global_size), 1},
+                                              gko::dim<2>{
+  static_cast<gko::size_type>(
+					      local_size), 1}, 1);
         scalar *host_buffer = host_device_vector->get_local_values();
 
         auto host_buffer_view =
diff --git a/include/OGL/MatrixWrapper/SparsityPattern.H b/include/OGL/MatrixWrapper/SparsityPattern.H
index edc2b2cd5..8d4793188 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.H
+++ b/include/OGL/MatrixWrapper/SparsityPattern.H
@@ -7,6 +7,7 @@
 #include <vector>
 
 #include "../common.H"
+#include "OGL/CommunicationPattern.H"
 
 namespace Foam {
 
@@ -19,25 +20,9 @@ namespace Foam {
  * 4 | 8 ] based on the recv_counts are used.
  *
  * */
-void make_ldu_mapping_consecutive(
-    const AllToAllPattern & comm_pattern,
-    gko::array<label> &ldu_mapping, label rank, label ranks_per_gpu)
-    {
-        // TODO move to SparsityPattern.C
-    label ldu_offset = 0;
-    auto *data = ldu_mapping.get_data();
-    // TODO check if ldu_mapping is on host exec
-
-    for (label i = 0; i < ranks_per_gpu; i++) {
-        auto size = comm_pattern.recv_counts[i];
-        std::transform(data + ldu_offset, data + ldu_offset + size,
-                       data + ldu_offset,
-                       [&](label idx) { return idx + ldu_offset; });
-        ldu_offset += size;
-    }
-
-
-    }
+void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
+                                  gko::array<label> &ldu_mapping, label rank,
+                                  label ranks_per_gpu);
 
 
 /* The SparsityPattern holds row and column index data using gko::arrays.
@@ -49,12 +34,8 @@ void make_ldu_mapping_consecutive(
  * */
 struct SparsityPattern {
     /* constructor from vectors, assumes a single interface */
-    SparsityPattern(
-        std::vector<label> rows,
-        std::vector<label> cols,
-        std::vector<label> mapping,
-        std::vector<label> rank
-)
+    SparsityPattern(std::vector<label> rows, std::vector<label> cols,
+                    std::vector<label> mapping, std::vector<label> rank)
         : num_nnz(rows.size())
     {
         // NOT fully implementented
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index 684432f6d..fe8561734 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -212,25 +212,27 @@ public:
             std::vector{rank}};
 
         if (is_owner(exec_handler)) {
-            make_ldu_mapping_consecutive(
-                local_comm_pattern, merged_local.ldu_mapping, rank, ranks_per_gpu);
+            make_ldu_mapping_consecutive(local_comm_pattern,
+                                         merged_local.ldu_mapping, rank,
+                                         ranks_per_gpu);
         }
 
-	auto get_back = [](const gko::array<label> & in){return *(in.get_const_data() + in.get_size());};
+        auto get_back = [](const gko::array<label> &in) {
+            return *(in.get_const_data() + in.get_size());
+        };
 
         label rows =
             (is_owner(exec_handler)) ? get_back(merged_local.row_idxs) + 1 : 0;
         gko::dim<2> merged_local_dim{static_cast<gko::size_type>(rows),
-                           static_cast<gko::size_type>(rows)};
+                                     static_cast<gko::size_type>(rows)};
 
         auto non_local_comm_pattern = compute_gather_to_owner_counts(
-            exec_handler, ranks_per_gpu, src_non_local_pattern->size_);
+            exec_handler, ranks_per_gpu, src_non_local_pattern->num_nnz);
 
         std::vector<label> spans_begin;
         std::vector<label> spans_end;
         auto span_comm_pattern = compute_gather_to_owner_counts(
-            exec_handler, ranks_per_gpu,
-            src_non_local_pattern->interface_spans.size());
+            exec_handler, ranks_per_gpu, src_non_local_pattern->spans.size());
         for (auto elem : src_non_local_pattern->spans) {
             spans_begin.push_back(elem.begin);
             spans_end.push_back(elem.end);
@@ -242,7 +244,7 @@ public:
         // on the owner side after gathering
         for (label i = 0; i < src_non_local_pattern->rank.size(); i++) {
             auto comm_rank = src_non_local_pattern->rank[i];
-            auto [begin, end] = src_non_local_pattern->interface_spans[i];
+            auto [begin, end] = src_non_local_pattern->spans[i];
             label local_offset =
                 orig_partition_->get_range_bounds()[comm_rank] -
                 orig_partition_->get_range_bounds()[owner_rank];
@@ -259,18 +261,23 @@ public:
                            src_non_local_pattern->col_idxs, 0),
             gather_closure(non_local_comm_pattern,
                            src_non_local_pattern->row_idxs, 0),
-            gather_to_owner(exec_handler, span_comm_pattern, spans_begin.size(),
-                            spans_begin.data()),
-            gather_to_owner(exec_handler, span_comm_pattern, spans_end.size(),
-                            spans_end.data()),
-            gather_to_owner(exec_handler, span_comm_pattern,
-                            src_non_local_pattern->rank.size(),
-                            src_non_local_pattern->rank.data())};
+            // {},
+            // this doesn't exist
+            // gather_labels_to_owner(exec_handler, span_comm_pattern,
+            // spans_begin.size(),
+            //                spans_begin.data()),
+            //  gather_labels_to_owner(exec_handler, span_comm_pattern,
+            //  spans_end.size(),
+            //                spans_end.data()),
+            gather_labels_to_owner(exec_handler, span_comm_pattern,
+                                   src_non_local_pattern->rank.data(),
+                                   src_non_local_pattern->rank.size())};
 
         if (is_owner(exec_handler)) {
-            make_begin_and_ends_consecutive(
-                span_comm_pattern, merged_non_local.begin, merged_non_local.end,
-                rank, ranks_per_gpu);
+            FatalErrorInFunction << "Not implemented" << abort(FatalError);
+            //  make_begin_and_ends_consecutive(
+            //      span_comm_pattern, merged_non_local.spans,
+            //     rank, ranks_per_gpu);
         }
 
         auto [gathered_non_local, is_local] = build_non_local_interfaces(
@@ -281,10 +288,9 @@ public:
         label ctr{0};
 
         if (is_owner(exec_handler)) {
-            auto recv_counts = std::get<1>(span_comm_pattern);
             for (int i = 0; i < ranks_per_gpu; i++) {
                 label comm_rank = rank + i;
-                label num_elems = recv_counts[comm_rank];
+                label num_elems = span_comm_pattern.recv_counts[comm_rank];
                 for (int j = 0; j < num_elems; j++) {
                     locality.emplace_back(is_local[ctr], comm_rank);
                     ctr++;
@@ -301,50 +307,54 @@ public:
         // keep an interface together thus we can just count the idx up to the
         // size. But we have to make sure that the interfaces are in the same
         // order on both communication sides.
-        for (int i = 0; i < gathered_non_local.cols.size(); i++) {
-            gathered_non_local.cols[i] = i;
+        for (int i = 0; i < gathered_non_local.col_idxs.get_size(); i++) {
+            gathered_non_local.col_idxs.get_data()[i] = i;
         }
 
         LOG_1(verbose_, "done repartition sparsity pattern")
         if (is_owner(exec_handler)) {
-            auto new_local_spars_pattern = std::make_shared<SparsityPattern>(
-                exec, merged_local_dim, merged_local);
-
-            auto new_non_local_spars_pattern =
-                std::make_shared<SparsityPattern>(
-                    exec,
-                    gko::dim<2>{merged_local_dim[0],
-                                gathered_non_local.rows.size()},
-                    gathered_non_local);
-
-            return std::make_tuple<std::shared_ptr<SparsityPattern>,
-                                   std::shared_ptr<SparsityPattern>,
-                                   std::vector<std::pair<bool, label>>>(
-                std::move(new_local_spars_pattern),
-                std::move(new_non_local_spars_pattern), std::move(locality));
+            FatalErrorInFunction << "Not implemented" << abort(FatalError);
+//auto new_local_spars_pattern = std::make_shared<SparsityPattern>(
+//    exec, merged_local_dim, merged_local);
+//
+//auto new_non_local_spars_pattern =
+//    std::make_shared<SparsityPattern>(
+//        exec,
+//        gko::dim<2>{merged_local_dim[0],
+//                    gathered_non_local.row_idxs.get_size()},
+//        gathered_non_local);
+//
+            //        return std::make_tuple<std::shared_ptr<SparsityPattern>,
+            //                               std::shared_ptr<SparsityPattern>,
+            //                               std::vector<std::pair<bool,
+            //                               label>>>(
+            //            std::move(new_local_spars_pattern),
+            //            std::move(new_non_local_spars_pattern),
+            //            std::move(locality));
         } else {
-            auto new_local_spars_pattern =
-                std::make_shared<SparsityPattern>(exec);
-
-            auto new_non_local_spars_pattern =
-                std::make_shared<SparsityPattern>(exec);
-
-            // FIXME
-            // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
-            // comm.rank() <<
-            //     " merged_local dim " <<
-            new_local_spars_pattern->dim =
-                gko::dim<2>{0, 0};  // merged_local_dim;
-            new_non_local_spars_pattern->dim =
-                gko::dim<2>{0, 0};  // merged_local_dim;
-
-
-            return std::make_tuple<std::shared_ptr<SparsityPattern>,
-                                   std::shared_ptr<SparsityPattern>,
-                                   std::vector<std::pair<bool, label>>>(
-                std::move(new_local_spars_pattern),
-                std::move(new_non_local_spars_pattern),
-                std::vector<std::pair<bool, label>>{});
+// auto new_local_spars_pattern =
+//     std::make_shared<SparsityPattern>(exec);
+// 
+// auto new_non_local_spars_pattern =
+//     std::make_shared<SparsityPattern>(exec);
+// 
+// // FIXME
+// // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
+// // comm.rank() <<
+// //     " merged_local dim " <<
+// new_local_spars_pattern->dim =
+//     gko::dim<2>{0, 0};  // merged_local_dim;
+// new_non_local_spars_pattern->dim =
+//     gko::dim<2>{0, 0};  // merged_local_dim;
+// 
+// 
+            FatalErrorInFunction << "Not implemented" << abort(FatalError);
+            // return std::make_tuple<std::shared_ptr<SparsityPattern>,
+            //                        std::shared_ptr<SparsityPattern>,
+            //                        std::vector<std::pair<bool, label>>>(
+            //     std::move(new_local_spars_pattern),
+            //     std::move(new_non_local_spars_pattern),
+            //     std::vector<std::pair<bool, label>>{});
         }
         // TODO Some old documentation see what is still needed
         // add interfaces
@@ -414,23 +424,22 @@ public:
         auto comm_pattern = compute_gather_to_owner_counts(
             exec_handler, ranks_per_gpu_, target_ids.size());
 
-        auto gathered_target_ids = gather_to_owner(
-            exec_handler, comm_pattern, target_ids.size(), target_ids.data());
+        auto gathered_target_ids = gather_labels_to_owner(
+            exec_handler, comm_pattern, target_ids.data(), target_ids.size());
 
         auto gathered_target_sizes =
-            gather_to_owner(exec_handler, comm_pattern, target_sizes.size(),
-                            target_sizes.data());
+            gather_labels_to_owner(exec_handler, comm_pattern,
+                                   target_sizes.data(), target_sizes.size());
 
         // next the send_ixs need to be updated we send them piecewise since
         // the send_idxs are a vector of gko::arrays
         if (owner) {
-            auto recv_counts = std::get<1>(comm_pattern);
-            label owner_recv_counts = recv_counts[rank];
+            label owner_recv_counts = comm_pattern.recv_counts[rank];
             // retrieved from i-th neighbor
             for (int i = 1; i < ranks_per_gpu_; i++) {
                 // how many gko::arrays to with send_indexes to receive
                 // from i-th neighbor
-                label recv_count = recv_counts[rank + i];
+                label recv_count = comm_pattern.recv_counts[rank + i];
                 for (int j = 0; j < recv_count; j++) {
                     auto target_size =
                         gathered_target_sizes[j + owner_recv_counts];
@@ -465,8 +474,7 @@ public:
             }
         } else {
             label owner = get_owner_rank(exec_handler);
-            auto send_count = std::get<0>(comm_pattern)[owner];
-            for (int i = 0; i < send_count; i++) {
+            for (int i = 0; i < comm_pattern.send_counts[owner]; i++) {
                 auto send_buffer = send_idxs[i].first;
                 // std::cout << __FILE__ << ":" << __LINE__ << " rank " << rank
                 //           << " send to owner " << owner
diff --git a/include/OGL/common.H b/include/OGL/common.H
index 8c0fb35e0..4172f97f3 100644
--- a/include/OGL/common.H
+++ b/include/OGL/common.H
@@ -3,6 +3,7 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 
 #pragma once
+#include <ginkgo/ginkgo.hpp>
 
 #include "fvCFD.H"
 #include "regIOobject.H"
@@ -12,7 +13,6 @@
 #endif
 
 #include <string.h>
-#include <ginkgo/ginkgo.hpp>
 
 namespace Foam {
 
diff --git a/src/MatrixWrapper/HostMatrix.C b/src/MatrixWrapper/HostMatrix.C
index d09a4ded4..b7f2d7888 100644
--- a/src/MatrixWrapper/HostMatrix.C
+++ b/src/MatrixWrapper/HostMatrix.C
@@ -452,7 +452,9 @@ std::shared_ptr<SparsityPattern> HostMatrixWrapper::compute_non_local_sparsity(
 {
     auto sparsity{std::make_shared<SparsityPattern>(exec->get_master(),
                                                     non_local_matrix_nnz_)};
-    sparsity->dim = gko::dim<2>{nrows_, non_local_matrix_nnz_};
+    sparsity->dim = gko::dim<2>{
+    static_cast<gko::size_type>(nrows_),
+     static_cast<gko::size_type>(non_local_matrix_nnz_)};
 
     auto non_local_indices = collect_cells_on_non_local_interface(interfaces_);
     auto rows = sparsity->row_idxs.get_data();
@@ -529,7 +531,11 @@ std::shared_ptr<SparsityPattern> HostMatrixWrapper::compute_local_sparsity(
     init_local_sparsity(nrows_, upper_nnz_, symmetric_, upper, lower, rows,
                         cols, permute);
 
-    sparsity->dim = gko::dim<2>{nrows_, nrows_};
+    sparsity->dim = gko::dim<2>{
+  static_cast<gko::size_type>(
+    nrows_),
+  static_cast<gko::size_type>(
+    nrows_)};
     sparsity->spans.emplace_back(0, local_matrix_nnz_);
 
     // if no local interfaces are present we are done here
diff --git a/src/MatrixWrapper/SparsityPattern.C b/src/MatrixWrapper/SparsityPattern.C
new file mode 100644
index 000000000..7f543e380
--- /dev/null
+++ b/src/MatrixWrapper/SparsityPattern.C
@@ -0,0 +1,22 @@
+// SPDX-FileCopyrightText: 2024 OGL authors
+//
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "OGL/MatrixWrapper/SparsityPattern.H"
+
+void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
+                                  gko::array<label> &ldu_mapping, label rank,
+                                  label ranks_per_gpu)
+{
+    label ldu_offset = 0;
+    auto *data = ldu_mapping.get_data();
+    // TODO check if ldu_mapping is on host exec
+
+    for (label i = 0; i < ranks_per_gpu; i++) {
+        auto size = comm_pattern.recv_counts[i];
+        std::transform(data + ldu_offset, data + ldu_offset + size,
+                       data + ldu_offset,
+                       [&](label idx) { return idx + ldu_offset; });
+        ldu_offset += size;
+    }
+}
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index 719ac3c39..843a269ab 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -12,46 +12,50 @@ Repartitioner::build_non_local_interfaces(const ExecutorHandler &exec_handler,
     LOG_1(verbose_, "start build non local interfaces")
     std::vector<label> rows, cols, ldu_mapping, ranks, begins, ends;
     std::vector<bool> is_local;
-    label merged_ranks_size = non_loc.ranks.size();
+    label merged_ranks_size = non_loc.rank.size();
 
-    for (label i = 0; i < merged_ranks_size; i++) {
-        // these are the begin ends before merging they need to be
-        // offseted;
-        auto begin = non_loc.begin[i];
-        auto end = non_loc.end[i];
-        bool local = reparts_to_local(exec_handler, non_loc.ranks[i]);
-        is_local.push_back(local);
-        if (local) {
-            // TODO depending on the interface simple non
-            // transforming and interpolating interfaces like now
-            // local processor interfaces could be merged
-            loc.begin.push_back(loc.rows.size());
-            loc.rows.insert(loc.rows.end(), non_loc.rows.begin() + begin,
-                            non_loc.rows.begin() + end);
-            loc.cols.insert(loc.cols.end(), non_loc.cols.begin() + begin,
-                            non_loc.cols.begin() + end);
-            loc.mapping.insert(loc.mapping.end(),
-                               non_loc.mapping.data() + begin,
-                               non_loc.mapping.data() + end);
-            // TODO store from rank ie from which rank it came
-            // this is currently unused
-            loc.end.push_back(loc.rows.size());
-            loc.ranks.push_back(non_loc.ranks[i]);
-        } else {
-            begins.push_back(rows.size());
-            rows.insert(rows.end(), non_loc.rows.begin() + begin,
-                        non_loc.rows.begin() + end);
-            cols.insert(cols.end(), non_loc.cols.begin() + begin,
-                        non_loc.cols.begin() + end);
-            // NOTE we don't do anything with non local mapping
-            ldu_mapping.insert(ldu_mapping.end(), non_loc.cols.begin() + begin,
-                               non_loc.cols.begin() + end);
-            ends.push_back(rows.size());
-            ranks.push_back(get_owner_rank(non_loc.ranks[i]));
-        }
-    }
-    LOG_1(verbose_, "done build non local interfaces")
-    return std::make_pair(
-        SparsityPattern{rows, cols, ldu_mapping, begins, ends, ranks},
-        is_local);
+    // TODO dont modify SparsityPattern but create vector
+    // and use the vector variant
+//   for (label i = 0; i < merged_ranks_size; i++) {
+//       // these are the begin ends before merging they need to be
+//       // offseted;
+//       auto begin = non_loc.spans[i].begin;
+//       auto end = non_loc.spans[i].end;
+//       bool local = reparts_to_local(exec_handler, non_loc.rank[i]);
+//       is_local.push_back(local);
+//       if (local) {
+//           // TODO depending on the interface simple non
+//           // transforming and interpolating interfaces like now
+//           // local processor interfaces could be merged
+//           loc.begin.push_back(loc.row_idxs.size());
+//           loc.row_idxs.insert(loc.row_idxs.end(), non_loc.row_idxs.get_data() + begin,
+//                           non_loc.row_idxs.get_data() + end);
+//           loc.col_idxs.insert(loc.cols.end(), non_loc.col_idxs.get_data() + begin,
+//                           non_loc.col_idxs.get_data() + end);
+//           loc.ldu_mapping.insert(loc.ldu_mapping.end(),
+//                              non_loc.ldu_mapping.get_data() + begin,
+//                              non_loc.ldu_mapping.get_data() + end);
+//           // TODO store from rank ie from which rank it came
+//           // this is currently unused
+//           // TODO reimplement end doesnt exist
+//           // loc.end.push_back(loc.row_idxs.size());
+//           // also cannot pushback
+//           // loc.rank.push_back(non_loc.rank[i]);
+//       } else {
+//           begins.push_back(row_idxs.size());
+//           rows.insert(rows.end(), non_loc.row_idxs.begin() + begin,
+//                       non_loc.row_idxs.begin() + end);
+//           cols.insert(cols.end(), non_loc.col_idxs.begin() + begin,
+//                       non_loc.col_idxs.begin() + end);
+//           // NOTE we don't do anything with non local mapping
+//           ldu_mapping.insert(ldu_mapping.end(), non_loc.col_idxs.begin() + begin,
+//                              non_loc.col_idxs.begin() + end);
+//           ends.push_back(rows.size());
+//           ranks.push_back(get_owner_rank(non_loc.rank[i]));
+//       }
+//   }
+//   LOG_1(verbose_, "done build non local interfaces")
+//   return std::make_pair(
+//       SparsityPattern{rows, cols, ldu_mapping, begins, ends, ranks},
+//       is_local);
 }
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index bc64a9255..1cf63aafe 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -2,7 +2,7 @@
 //
 // SPDX-License-Identifier: GPL-3.0-or-later
 
-#include "OGL/CommunicationPattern.H"
+#include "OGL/Repartitioner.H"
 
 #include "gtest/gtest.h"
 
@@ -56,14 +56,14 @@ public:
 };
 
 const testing::Environment *global_env =
-    AddGlobalTestEnvironment(new CommunicationPatternEnvironment);
+    AddGlobalTestEnvironment(new RepartitionerEnvironment);
 
 
 TEST(CommunicationPattern, can_create_repartitioner)
 {
     // Arrange
     label ranks_per_gpu = 1;
-    auto exec = ((CommunicationPatternEnvironment *)global_env)->exec;
+    auto exec = ((RepartitionerEnvironment *)global_env)->exec;
     auto repartitioner = Repartitioner(10, ranks_per_gpu, 0, *exec.get());
 
     // Assert

From 9883fe80944e9b04967805bd46631060cf630b8a Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Sat, 17 Aug 2024 16:17:53 +0200
Subject: [PATCH 08/67] format files

---
 include/OGL/DevicePersistent/Vector.H | 11 ++--
 include/OGL/Repartitioner.H           | 66 ++++++++++----------
 src/MatrixWrapper/HostMatrix.C        | 13 ++--
 src/Repartitioner.C                   | 87 ++++++++++++++-------------
 unitTests/MatrixWrapper/Combination.C |  3 +-
 unitTests/Repartitioner.C             | 17 +++++-
 6 files changed, 102 insertions(+), 95 deletions(-)

diff --git a/include/OGL/DevicePersistent/Vector.H b/include/OGL/DevicePersistent/Vector.H
index 7ac0ba4ab..ea632d915 100644
--- a/include/OGL/DevicePersistent/Vector.H
+++ b/include/OGL/DevicePersistent/Vector.H
@@ -153,13 +153,10 @@ public:
 
         // TODO why does it need its size
         // repartitioner size should be sufficient
-        auto target_vector = dist_vec::create(exec_.get_ref_exec(), *comm.get(),
-                                              gko::dim<2>{
-  static_cast<gko::size_type>(
-					      global_size), 1},
-                                              gko::dim<2>{
-  static_cast<gko::size_type>(
-					      local_size), 1}, 1);
+        auto target_vector = dist_vec::create(
+            exec_.get_ref_exec(), *comm.get(),
+            gko::dim<2>{static_cast<gko::size_type>(global_size), 1},
+            gko::dim<2>{static_cast<gko::size_type>(local_size), 1}, 1);
         scalar *host_buffer = host_device_vector->get_local_values();
 
         auto host_buffer_view =
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index fe8561734..ed86c0709 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -314,40 +314,40 @@ public:
         LOG_1(verbose_, "done repartition sparsity pattern")
         if (is_owner(exec_handler)) {
             FatalErrorInFunction << "Not implemented" << abort(FatalError);
-//auto new_local_spars_pattern = std::make_shared<SparsityPattern>(
-//    exec, merged_local_dim, merged_local);
-//
-//auto new_non_local_spars_pattern =
-//    std::make_shared<SparsityPattern>(
-//        exec,
-//        gko::dim<2>{merged_local_dim[0],
-//                    gathered_non_local.row_idxs.get_size()},
-//        gathered_non_local);
-//
-            //        return std::make_tuple<std::shared_ptr<SparsityPattern>,
-            //                               std::shared_ptr<SparsityPattern>,
-            //                               std::vector<std::pair<bool,
-            //                               label>>>(
-            //            std::move(new_local_spars_pattern),
-            //            std::move(new_non_local_spars_pattern),
-            //            std::move(locality));
+            // auto new_local_spars_pattern = std::make_shared<SparsityPattern>(
+            //     exec, merged_local_dim, merged_local);
+            //
+            // auto new_non_local_spars_pattern =
+            //     std::make_shared<SparsityPattern>(
+            //         exec,
+            //         gko::dim<2>{merged_local_dim[0],
+            //                     gathered_non_local.row_idxs.get_size()},
+            //         gathered_non_local);
+            //
+            //         return std::make_tuple<std::shared_ptr<SparsityPattern>,
+            //                                std::shared_ptr<SparsityPattern>,
+            //                                std::vector<std::pair<bool,
+            //                                label>>>(
+            //             std::move(new_local_spars_pattern),
+            //             std::move(new_non_local_spars_pattern),
+            //             std::move(locality));
         } else {
-// auto new_local_spars_pattern =
-//     std::make_shared<SparsityPattern>(exec);
-// 
-// auto new_non_local_spars_pattern =
-//     std::make_shared<SparsityPattern>(exec);
-// 
-// // FIXME
-// // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
-// // comm.rank() <<
-// //     " merged_local dim " <<
-// new_local_spars_pattern->dim =
-//     gko::dim<2>{0, 0};  // merged_local_dim;
-// new_non_local_spars_pattern->dim =
-//     gko::dim<2>{0, 0};  // merged_local_dim;
-// 
-// 
+            // auto new_local_spars_pattern =
+            //     std::make_shared<SparsityPattern>(exec);
+            //
+            // auto new_non_local_spars_pattern =
+            //     std::make_shared<SparsityPattern>(exec);
+            //
+            // // FIXME
+            // // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
+            // // comm.rank() <<
+            // //     " merged_local dim " <<
+            // new_local_spars_pattern->dim =
+            //     gko::dim<2>{0, 0};  // merged_local_dim;
+            // new_non_local_spars_pattern->dim =
+            //     gko::dim<2>{0, 0};  // merged_local_dim;
+            //
+            //
             FatalErrorInFunction << "Not implemented" << abort(FatalError);
             // return std::make_tuple<std::shared_ptr<SparsityPattern>,
             //                        std::shared_ptr<SparsityPattern>,
diff --git a/src/MatrixWrapper/HostMatrix.C b/src/MatrixWrapper/HostMatrix.C
index b7f2d7888..c97d9719d 100644
--- a/src/MatrixWrapper/HostMatrix.C
+++ b/src/MatrixWrapper/HostMatrix.C
@@ -452,9 +452,9 @@ std::shared_ptr<SparsityPattern> HostMatrixWrapper::compute_non_local_sparsity(
 {
     auto sparsity{std::make_shared<SparsityPattern>(exec->get_master(),
                                                     non_local_matrix_nnz_)};
-    sparsity->dim = gko::dim<2>{
-    static_cast<gko::size_type>(nrows_),
-     static_cast<gko::size_type>(non_local_matrix_nnz_)};
+    sparsity->dim =
+        gko::dim<2>{static_cast<gko::size_type>(nrows_),
+                    static_cast<gko::size_type>(non_local_matrix_nnz_)};
 
     auto non_local_indices = collect_cells_on_non_local_interface(interfaces_);
     auto rows = sparsity->row_idxs.get_data();
@@ -531,11 +531,8 @@ std::shared_ptr<SparsityPattern> HostMatrixWrapper::compute_local_sparsity(
     init_local_sparsity(nrows_, upper_nnz_, symmetric_, upper, lower, rows,
                         cols, permute);
 
-    sparsity->dim = gko::dim<2>{
-  static_cast<gko::size_type>(
-    nrows_),
-  static_cast<gko::size_type>(
-    nrows_)};
+    sparsity->dim = gko::dim<2>{static_cast<gko::size_type>(nrows_),
+                                static_cast<gko::size_type>(nrows_)};
     sparsity->spans.emplace_back(0, local_matrix_nnz_);
 
     // if no local interfaces are present we are done here
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index 843a269ab..79a0ea5d2 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -16,46 +16,49 @@ Repartitioner::build_non_local_interfaces(const ExecutorHandler &exec_handler,
 
     // TODO dont modify SparsityPattern but create vector
     // and use the vector variant
-//   for (label i = 0; i < merged_ranks_size; i++) {
-//       // these are the begin ends before merging they need to be
-//       // offseted;
-//       auto begin = non_loc.spans[i].begin;
-//       auto end = non_loc.spans[i].end;
-//       bool local = reparts_to_local(exec_handler, non_loc.rank[i]);
-//       is_local.push_back(local);
-//       if (local) {
-//           // TODO depending on the interface simple non
-//           // transforming and interpolating interfaces like now
-//           // local processor interfaces could be merged
-//           loc.begin.push_back(loc.row_idxs.size());
-//           loc.row_idxs.insert(loc.row_idxs.end(), non_loc.row_idxs.get_data() + begin,
-//                           non_loc.row_idxs.get_data() + end);
-//           loc.col_idxs.insert(loc.cols.end(), non_loc.col_idxs.get_data() + begin,
-//                           non_loc.col_idxs.get_data() + end);
-//           loc.ldu_mapping.insert(loc.ldu_mapping.end(),
-//                              non_loc.ldu_mapping.get_data() + begin,
-//                              non_loc.ldu_mapping.get_data() + end);
-//           // TODO store from rank ie from which rank it came
-//           // this is currently unused
-//           // TODO reimplement end doesnt exist
-//           // loc.end.push_back(loc.row_idxs.size());
-//           // also cannot pushback
-//           // loc.rank.push_back(non_loc.rank[i]);
-//       } else {
-//           begins.push_back(row_idxs.size());
-//           rows.insert(rows.end(), non_loc.row_idxs.begin() + begin,
-//                       non_loc.row_idxs.begin() + end);
-//           cols.insert(cols.end(), non_loc.col_idxs.begin() + begin,
-//                       non_loc.col_idxs.begin() + end);
-//           // NOTE we don't do anything with non local mapping
-//           ldu_mapping.insert(ldu_mapping.end(), non_loc.col_idxs.begin() + begin,
-//                              non_loc.col_idxs.begin() + end);
-//           ends.push_back(rows.size());
-//           ranks.push_back(get_owner_rank(non_loc.rank[i]));
-//       }
-//   }
-//   LOG_1(verbose_, "done build non local interfaces")
-//   return std::make_pair(
-//       SparsityPattern{rows, cols, ldu_mapping, begins, ends, ranks},
-//       is_local);
+    //   for (label i = 0; i < merged_ranks_size; i++) {
+    //       // these are the begin ends before merging they need to be
+    //       // offseted;
+    //       auto begin = non_loc.spans[i].begin;
+    //       auto end = non_loc.spans[i].end;
+    //       bool local = reparts_to_local(exec_handler, non_loc.rank[i]);
+    //       is_local.push_back(local);
+    //       if (local) {
+    //           // TODO depending on the interface simple non
+    //           // transforming and interpolating interfaces like now
+    //           // local processor interfaces could be merged
+    //           loc.begin.push_back(loc.row_idxs.size());
+    //           loc.row_idxs.insert(loc.row_idxs.end(),
+    //           non_loc.row_idxs.get_data() + begin,
+    //                           non_loc.row_idxs.get_data() + end);
+    //           loc.col_idxs.insert(loc.cols.end(), non_loc.col_idxs.get_data()
+    //           + begin,
+    //                           non_loc.col_idxs.get_data() + end);
+    //           loc.ldu_mapping.insert(loc.ldu_mapping.end(),
+    //                              non_loc.ldu_mapping.get_data() + begin,
+    //                              non_loc.ldu_mapping.get_data() + end);
+    //           // TODO store from rank ie from which rank it came
+    //           // this is currently unused
+    //           // TODO reimplement end doesnt exist
+    //           // loc.end.push_back(loc.row_idxs.size());
+    //           // also cannot pushback
+    //           // loc.rank.push_back(non_loc.rank[i]);
+    //       } else {
+    //           begins.push_back(row_idxs.size());
+    //           rows.insert(rows.end(), non_loc.row_idxs.begin() + begin,
+    //                       non_loc.row_idxs.begin() + end);
+    //           cols.insert(cols.end(), non_loc.col_idxs.begin() + begin,
+    //                       non_loc.col_idxs.begin() + end);
+    //           // NOTE we don't do anything with non local mapping
+    //           ldu_mapping.insert(ldu_mapping.end(), non_loc.col_idxs.begin()
+    //           + begin,
+    //                              non_loc.col_idxs.begin() + end);
+    //           ends.push_back(rows.size());
+    //           ranks.push_back(get_owner_rank(non_loc.rank[i]));
+    //       }
+    //   }
+    //   LOG_1(verbose_, "done build non local interfaces")
+    //   return std::make_pair(
+    //       SparsityPattern{rows, cols, ldu_mapping, begins, ends, ranks},
+    //       is_local);
 }
diff --git a/unitTests/MatrixWrapper/Combination.C b/unitTests/MatrixWrapper/Combination.C
index 2c2e71e45..6cba3c735 100644
--- a/unitTests/MatrixWrapper/Combination.C
+++ b/unitTests/MatrixWrapper/Combination.C
@@ -74,8 +74,7 @@ TEST(Combination, CanCreateConvertToCsr)
     auto cmb = CombinationMatrix<gko::matrix::Csr<scalar, label>>::create(
         exec, dim, linops);
 
-    auto out = CombinationMatrix<gko::matrix::Csr<scalar, label>>::create(
-        exec);
+    auto out = CombinationMatrix<gko::matrix::Csr<scalar, label>>::create(exec);
 
     cmb->convert_to(out);
 }
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 1cf63aafe..749e646f9 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -58,8 +58,7 @@ public:
 const testing::Environment *global_env =
     AddGlobalTestEnvironment(new RepartitionerEnvironment);
 
-
-TEST(CommunicationPattern, can_create_repartitioner)
+TEST(Repartitioner, can_create_repartitioner)
 {
     // Arrange
     label ranks_per_gpu = 1;
@@ -67,7 +66,19 @@ TEST(CommunicationPattern, can_create_repartitioner)
     auto repartitioner = Repartitioner(10, ranks_per_gpu, 0, *exec.get());
 
     // Assert
-    EXPECT_EQ(repartitioner.get_ranks_per_gpu(), 0);
+    EXPECT_EQ(repartitioner.get_ranks_per_gpu(), ranks_per_gpu);
 }
 
+int main(int argc, char *argv[])
+{
+    int result = 0;
+
+    ::testing::InitGoogleTest(&argc, argv);
 
+    my_argc = argc;
+    my_argv = argv;
+
+    result = RUN_ALL_TESTS();
+
+    return result;
+}

From 4f460a4d4c72bef98ace68bf095190ce84287d6a Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Sat, 17 Aug 2024 16:22:08 +0200
Subject: [PATCH 09/67] fix typos

---
 include/OGL/Repartitioner.H | 4 ++--
 src/Repartitioner.C         | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index ed86c0709..b141fb179 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -486,7 +486,7 @@ public:
             }
         }
 
-        // clear communcation neighbors on non owning rank
+        // clear communication neighbors on non owning rank
         if (!owner) {
             // TODO NOTE should it be gathered_target_ids and is it needed?
             target_ids.clear();
@@ -533,7 +533,7 @@ public:
         for (int i = 1; i < target_ids.size(); i++) {
             // communicates with same target rank
             // thus we have only have to adapt the number
-            // of elemts and the send_ixs
+            // of elements and the send_ixs
             auto *send_idx_begin = send_idxs[i].first.get_data();
             auto *send_idx_end =
                 send_idxs[i].first.get_data() + send_idxs[i].first.get_size();
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index 79a0ea5d2..a1d5c3258 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -18,7 +18,7 @@ Repartitioner::build_non_local_interfaces(const ExecutorHandler &exec_handler,
     // and use the vector variant
     //   for (label i = 0; i < merged_ranks_size; i++) {
     //       // these are the begin ends before merging they need to be
-    //       // offseted;
+    //       // offsetted;
     //       auto begin = non_loc.spans[i].begin;
     //       auto end = non_loc.spans[i].end;
     //       bool local = reparts_to_local(exec_handler, non_loc.rank[i]);

From 1d754e1d8be33ea4866dbdb690ab3cbc8d462aa6 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Sat, 17 Aug 2024 17:02:29 +0200
Subject: [PATCH 10/67] add more unitTests

---
 CMakeLists.txt                              |   5 +-
 include/OGL/MatrixWrapper/SparsityPattern.H |   3 +-
 include/OGL/Repartitioner.H                 | 232 +-------------------
 src/MatrixWrapper/SparsityPattern.C         |   6 +-
 src/Repartitioner.C                         | 232 ++++++++++++++++++++
 unitTests/Repartitioner.C                   |  37 ++++
 6 files changed, 279 insertions(+), 236 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 9621f37b9..7d055ad03 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -222,12 +222,11 @@ target_sources(
           src/DevicePersistent/DeviceIdGuard/DeviceIdGuard.C
           src/DevicePersistent/ExecutorHandler/ExecutorHandler.C
           src/MatrixWrapper/HostMatrix.C
-	  src/MatrixWrapper/SparsityPattern.C
+          src/MatrixWrapper/SparsityPattern.C
           src/Repartitioner.C
           src/Solver/CG/GKOCG.C
           src/Solver/BiCGStab/GKOBiCGStab.C
-          src/Solver/GMRES/GKOGMRES.C
-)
+          src/Solver/GMRES/GKOGMRES.C)
 
 enable_sanitizers(
   OGL ${OGL_ENABLE_SANITIZE_ADDRESS} ${OGL_ENABLE_SANITIZE_LEAK}
diff --git a/include/OGL/MatrixWrapper/SparsityPattern.H b/include/OGL/MatrixWrapper/SparsityPattern.H
index 8d4793188..46390b26d 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.H
+++ b/include/OGL/MatrixWrapper/SparsityPattern.H
@@ -6,7 +6,7 @@
 
 #include <vector>
 
-#include "../common.H"
+#include "OGL/common.H"
 #include "OGL/CommunicationPattern.H"
 
 namespace Foam {
@@ -24,7 +24,6 @@ void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
                                   gko::array<label> &ldu_mapping, label rank,
                                   label ranks_per_gpu);
 
-
 /* The SparsityPattern holds row and column index data using gko::arrays.
  * This struct is used for easy generation of (distributed) ginkgo matrices.
  *
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index b141fb179..d0d5310ee 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -105,17 +105,7 @@ public:
      * this rank after repartitioning gets computed
      */
     static label compute_repart_size(label local_size, label ranks_per_gpu,
-                                     const ExecutorHandler &exec_handler)
-    {
-        if (ranks_per_gpu == 1) {
-            return local_size;
-        }
-
-        auto all_to_all_pattern = compute_gather_to_owner_counts(
-            exec_handler, ranks_per_gpu, local_size);
-
-        return all_to_all_pattern.recv_offsets.back();
-    }
+                                     const ExecutorHandler &exec_handler);
 
     label get_repart_size() const { return repart_size_; }
 
@@ -156,225 +146,7 @@ public:
     repartition_sparsity(
         const ExecutorHandler &exec_handler,
         std::shared_ptr<SparsityPattern> src_local_pattern,
-        std::shared_ptr<SparsityPattern> src_non_local_pattern) const
-    {
-        LOG_1(verbose_, "start repartition sparsity pattern")
-        // 1. obtain send recv sizes vector
-        // here we can reuse code from repartition_comm_pattern
-        //
-        // 2. initialize send and recv buffer
-        // if we keep interfaces separated row and col indices can be
-        //
-        // 3. ldu mapping needs to be update?
-        // 4. dim needs to be updated
-        // this is not implemented yet, but we can't fail here for
-        // debugging reasons
-        auto exec = exec_handler.get_ref_exec();
-        auto comm = *exec_handler.get_communicator().get();
-        label rank = get_rank(exec_handler);
-        label owner_rank = get_owner_rank(exec_handler);
-        label ranks_per_gpu = ranks_per_gpu_;
-        // TODO dont copy
-        // if (ranks_per_gpu == 1) {
-        //     std::vector<std::pair<bool, label>> ret;
-        //     for (auto comm_rank : src_non_local_pattern->rank) {
-        //         ret.emplace_back(false, rank);
-        //     }
-        //     return std::make_tuple<std::shared_ptr<SparsityPattern>,
-        //                            std::shared_ptr<SparsityPattern>,
-        //                            std::vector<std::pair<bool, label>>>(
-        //         std::make_shared<SparsityPattern>(src_local_pattern),
-        //         std::make_shared<SparsityPattern>(src_non_local_pattern),
-        //         std::move(ret));
-        // }
-
-        auto local_comm_pattern = compute_gather_to_owner_counts(
-            exec_handler, ranks_per_gpu, src_local_pattern->num_nnz);
-
-
-        auto offset = orig_partition_->get_range_bounds()[rank] -
-                      orig_partition_->get_range_bounds()[owner_rank];
-
-        auto gather_closure = [&](auto &comm_pattern, auto &data,
-                                  label offset) {
-            return gather_labels_to_owner(exec_handler, comm_pattern,
-                                          data.get_data(), data.get_size(),
-                                          offset);
-        };
-
-        SparsityPattern merged_local{
-            gather_closure(local_comm_pattern, src_local_pattern->row_idxs,
-                           offset),
-            gather_closure(local_comm_pattern, src_local_pattern->col_idxs,
-                           offset),
-            gather_closure(local_comm_pattern, src_local_pattern->ldu_mapping,
-                           0),
-            std::vector{rank}};
-
-        if (is_owner(exec_handler)) {
-            make_ldu_mapping_consecutive(local_comm_pattern,
-                                         merged_local.ldu_mapping, rank,
-                                         ranks_per_gpu);
-        }
-
-        auto get_back = [](const gko::array<label> &in) {
-            return *(in.get_const_data() + in.get_size());
-        };
-
-        label rows =
-            (is_owner(exec_handler)) ? get_back(merged_local.row_idxs) + 1 : 0;
-        gko::dim<2> merged_local_dim{static_cast<gko::size_type>(rows),
-                                     static_cast<gko::size_type>(rows)};
-
-        auto non_local_comm_pattern = compute_gather_to_owner_counts(
-            exec_handler, ranks_per_gpu, src_non_local_pattern->num_nnz);
-
-        std::vector<label> spans_begin;
-        std::vector<label> spans_end;
-        auto span_comm_pattern = compute_gather_to_owner_counts(
-            exec_handler, ranks_per_gpu, src_non_local_pattern->spans.size());
-        for (auto elem : src_non_local_pattern->spans) {
-            spans_begin.push_back(elem.begin);
-            spans_end.push_back(elem.end);
-        }
-
-        // the non local cols are in local idx of other side
-        // thus we need the new offset of the other side
-        // NOTE TODO this modifies src_non_local_pattern better do the offsets
-        // on the owner side after gathering
-        for (label i = 0; i < src_non_local_pattern->rank.size(); i++) {
-            auto comm_rank = src_non_local_pattern->rank[i];
-            auto [begin, end] = src_non_local_pattern->spans[i];
-            label local_offset =
-                orig_partition_->get_range_bounds()[comm_rank] -
-                orig_partition_->get_range_bounds()[owner_rank];
-            auto *data = src_non_local_pattern->col_idxs.get_data() + begin;
-            auto size = end - begin;
-            std::transform(data, data + size, data,
-                           [&](label idx) { return idx + local_offset; });
-        }
-
-        SparsityPattern merged_non_local{
-            gather_closure(non_local_comm_pattern,
-                           src_non_local_pattern->row_idxs, offset),
-            gather_closure(non_local_comm_pattern,
-                           src_non_local_pattern->col_idxs, 0),
-            gather_closure(non_local_comm_pattern,
-                           src_non_local_pattern->row_idxs, 0),
-            // {},
-            // this doesn't exist
-            // gather_labels_to_owner(exec_handler, span_comm_pattern,
-            // spans_begin.size(),
-            //                spans_begin.data()),
-            //  gather_labels_to_owner(exec_handler, span_comm_pattern,
-            //  spans_end.size(),
-            //                spans_end.data()),
-            gather_labels_to_owner(exec_handler, span_comm_pattern,
-                                   src_non_local_pattern->rank.data(),
-                                   src_non_local_pattern->rank.size())};
-
-        if (is_owner(exec_handler)) {
-            FatalErrorInFunction << "Not implemented" << abort(FatalError);
-            //  make_begin_and_ends_consecutive(
-            //      span_comm_pattern, merged_non_local.spans,
-            //     rank, ranks_per_gpu);
-        }
-
-        auto [gathered_non_local, is_local] = build_non_local_interfaces(
-            exec_handler, merged_local, merged_non_local);
-
-        // build vector with locality information
-        std::vector<std::pair<bool, label>> locality;
-        label ctr{0};
-
-        if (is_owner(exec_handler)) {
-            for (int i = 0; i < ranks_per_gpu; i++) {
-                label comm_rank = rank + i;
-                label num_elems = span_comm_pattern.recv_counts[comm_rank];
-                for (int j = 0; j < num_elems; j++) {
-                    locality.emplace_back(is_local[ctr], comm_rank);
-                    ctr++;
-                }
-            }
-        }
-        // All interfaces that are in is_local need to be accounted for
-        ASSERT_EQ(ctr, is_local.size());
-
-        // TODO FIXME the correct way would be to check the communication
-        // pattern where a particular face is in the send idxs
-        // since we have already the row id of the other side
-        // it should be doable. Or alternatively we know that we
-        // keep an interface together thus we can just count the idx up to the
-        // size. But we have to make sure that the interfaces are in the same
-        // order on both communication sides.
-        for (int i = 0; i < gathered_non_local.col_idxs.get_size(); i++) {
-            gathered_non_local.col_idxs.get_data()[i] = i;
-        }
-
-        LOG_1(verbose_, "done repartition sparsity pattern")
-        if (is_owner(exec_handler)) {
-            FatalErrorInFunction << "Not implemented" << abort(FatalError);
-            // auto new_local_spars_pattern = std::make_shared<SparsityPattern>(
-            //     exec, merged_local_dim, merged_local);
-            //
-            // auto new_non_local_spars_pattern =
-            //     std::make_shared<SparsityPattern>(
-            //         exec,
-            //         gko::dim<2>{merged_local_dim[0],
-            //                     gathered_non_local.row_idxs.get_size()},
-            //         gathered_non_local);
-            //
-            //         return std::make_tuple<std::shared_ptr<SparsityPattern>,
-            //                                std::shared_ptr<SparsityPattern>,
-            //                                std::vector<std::pair<bool,
-            //                                label>>>(
-            //             std::move(new_local_spars_pattern),
-            //             std::move(new_non_local_spars_pattern),
-            //             std::move(locality));
-        } else {
-            // auto new_local_spars_pattern =
-            //     std::make_shared<SparsityPattern>(exec);
-            //
-            // auto new_non_local_spars_pattern =
-            //     std::make_shared<SparsityPattern>(exec);
-            //
-            // // FIXME
-            // // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
-            // // comm.rank() <<
-            // //     " merged_local dim " <<
-            // new_local_spars_pattern->dim =
-            //     gko::dim<2>{0, 0};  // merged_local_dim;
-            // new_non_local_spars_pattern->dim =
-            //     gko::dim<2>{0, 0};  // merged_local_dim;
-            //
-            //
-            FatalErrorInFunction << "Not implemented" << abort(FatalError);
-            // return std::make_tuple<std::shared_ptr<SparsityPattern>,
-            //                        std::shared_ptr<SparsityPattern>,
-            //                        std::vector<std::pair<bool, label>>>(
-            //     std::move(new_local_spars_pattern),
-            //     std::move(new_non_local_spars_pattern),
-            //     std::vector<std::pair<bool, label>>{});
-        }
-        // TODO Some old documentation see what is still needed
-        // add interfaces
-        // to add interfaces we go through the non_local_sparsity
-        // pattern and check if interface is still non_local
-        //
-        // iterate interfaces and send to owner.
-        // on owner decide to move to local or non_local
-        // for sending to owner we can also use the all_to_all_v
-        // approach we than have
-        //
-        // row [ 4 , 8,  12 | 16 , 17, 18 ] <- local row
-        // col [ 1 , 2, 3 | 1, 2, 3 ] <- just the interface ctr
-        // from repartition_comm_pattern we could get
-        //
-        // or we split local / non-local first
-        // for this we could store ranges when computing
-        // repartition_comm_pattern
-        //
-    }
+        std::shared_ptr<SparsityPattern> src_non_local_pattern) const;
 
     std::shared_ptr<const CommunicationPattern> repartition_comm_pattern(
         const ExecutorHandler &exec_handler,
diff --git a/src/MatrixWrapper/SparsityPattern.C b/src/MatrixWrapper/SparsityPattern.C
index 7f543e380..b56038e6f 100644
--- a/src/MatrixWrapper/SparsityPattern.C
+++ b/src/MatrixWrapper/SparsityPattern.C
@@ -4,13 +4,15 @@
 
 #include "OGL/MatrixWrapper/SparsityPattern.H"
 
+namespace Foam {
+
 void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
                                   gko::array<label> &ldu_mapping, label rank,
                                   label ranks_per_gpu)
 {
+    // TODO check if ldu_mapping is on host exec
     label ldu_offset = 0;
     auto *data = ldu_mapping.get_data();
-    // TODO check if ldu_mapping is on host exec
 
     for (label i = 0; i < ranks_per_gpu; i++) {
         auto size = comm_pattern.recv_counts[i];
@@ -20,3 +22,5 @@ void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
         ldu_offset += size;
     }
 }
+
+    }
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index a1d5c3258..e8e602efc 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -4,6 +4,238 @@
 
 #include "OGL/Repartitioner.H"
 
+
+label Repartitioner::compute_repart_size(label local_size, label ranks_per_gpu,
+                                         const ExecutorHandler &exec_handler)
+{
+    if (ranks_per_gpu == 1) {
+        return local_size;
+    }
+
+    auto all_to_all_pattern =
+        compute_gather_to_owner_counts(exec_handler, ranks_per_gpu, local_size);
+
+    return all_to_all_pattern.recv_offsets.back();
+}
+
+std::tuple<std::shared_ptr<SparsityPattern>, std::shared_ptr<SparsityPattern>,
+           std::vector<std::pair<bool, label>>>
+Repartitioner::repartition_sparsity(
+    const ExecutorHandler &exec_handler,
+    std::shared_ptr<SparsityPattern> src_local_pattern,
+    std::shared_ptr<SparsityPattern> src_non_local_pattern) const
+{
+    LOG_1(verbose_, "start repartition sparsity pattern")
+    // 1. obtain send recv sizes vector
+    // here we can reuse code from repartition_comm_pattern
+    //
+    // 2. initialize send and recv buffer
+    // if we keep interfaces separated row and col indices can be
+    //
+    // 3. ldu mapping needs to be update?
+    // 4. dim needs to be updated
+    // this is not implemented yet, but we can't fail here for
+    // debugging reasons
+    auto exec = exec_handler.get_ref_exec();
+    auto comm = *exec_handler.get_communicator().get();
+    label rank = get_rank(exec_handler);
+    label owner_rank = get_owner_rank(exec_handler);
+    label ranks_per_gpu = ranks_per_gpu_;
+    // TODO dont copy
+    // if (ranks_per_gpu == 1) {
+    //     std::vector<std::pair<bool, label>> ret;
+    //     for (auto comm_rank : src_non_local_pattern->rank) {
+    //         ret.emplace_back(false, rank);
+    //     }
+    //     return std::make_tuple<std::shared_ptr<SparsityPattern>,
+    //                            std::shared_ptr<SparsityPattern>,
+    //                            std::vector<std::pair<bool, label>>>(
+    //         std::make_shared<SparsityPattern>(src_local_pattern),
+    //         std::make_shared<SparsityPattern>(src_non_local_pattern),
+    //         std::move(ret));
+    // }
+
+    auto local_comm_pattern = compute_gather_to_owner_counts(
+        exec_handler, ranks_per_gpu, src_local_pattern->num_nnz);
+
+
+    auto offset = orig_partition_->get_range_bounds()[rank] -
+                  orig_partition_->get_range_bounds()[owner_rank];
+
+    auto gather_closure = [&](auto &comm_pattern, auto &data, label offset) {
+        return gather_labels_to_owner(exec_handler, comm_pattern,
+                                      data.get_data(), data.get_size(), offset);
+    };
+
+    SparsityPattern merged_local{
+        gather_closure(local_comm_pattern, src_local_pattern->row_idxs, offset),
+        gather_closure(local_comm_pattern, src_local_pattern->col_idxs, offset),
+        gather_closure(local_comm_pattern, src_local_pattern->ldu_mapping, 0),
+        std::vector{rank}};
+
+    if (is_owner(exec_handler)) {
+        make_ldu_mapping_consecutive(
+            local_comm_pattern, merged_local.ldu_mapping, rank, ranks_per_gpu);
+    }
+
+    auto get_back = [](const gko::array<label> &in) {
+        return *(in.get_const_data() + in.get_size());
+    };
+
+    label rows =
+        (is_owner(exec_handler)) ? get_back(merged_local.row_idxs) + 1 : 0;
+    gko::dim<2> merged_local_dim{static_cast<gko::size_type>(rows),
+                                 static_cast<gko::size_type>(rows)};
+
+    auto non_local_comm_pattern = compute_gather_to_owner_counts(
+        exec_handler, ranks_per_gpu, src_non_local_pattern->num_nnz);
+
+    std::vector<label> spans_begin;
+    std::vector<label> spans_end;
+    auto span_comm_pattern = compute_gather_to_owner_counts(
+        exec_handler, ranks_per_gpu, src_non_local_pattern->spans.size());
+    for (auto elem : src_non_local_pattern->spans) {
+        spans_begin.push_back(elem.begin);
+        spans_end.push_back(elem.end);
+    }
+
+    // the non local cols are in local idx of other side
+    // thus we need the new offset of the other side
+    // NOTE TODO this modifies src_non_local_pattern better do the offsets
+    // on the owner side after gathering
+    for (label i = 0; i < src_non_local_pattern->rank.size(); i++) {
+        auto comm_rank = src_non_local_pattern->rank[i];
+        auto [begin, end] = src_non_local_pattern->spans[i];
+        label local_offset = orig_partition_->get_range_bounds()[comm_rank] -
+                             orig_partition_->get_range_bounds()[owner_rank];
+        auto *data = src_non_local_pattern->col_idxs.get_data() + begin;
+        auto size = end - begin;
+        std::transform(data, data + size, data,
+                       [&](label idx) { return idx + local_offset; });
+    }
+
+    SparsityPattern merged_non_local{
+        gather_closure(non_local_comm_pattern, src_non_local_pattern->row_idxs,
+                       offset),
+        gather_closure(non_local_comm_pattern, src_non_local_pattern->col_idxs,
+                       0),
+        gather_closure(non_local_comm_pattern, src_non_local_pattern->row_idxs,
+                       0),
+        // {},
+        // this doesn't exist
+        // gather_labels_to_owner(exec_handler, span_comm_pattern,
+        // spans_begin.size(),
+        //                spans_begin.data()),
+        //  gather_labels_to_owner(exec_handler, span_comm_pattern,
+        //  spans_end.size(),
+        //                spans_end.data()),
+        gather_labels_to_owner(exec_handler, span_comm_pattern,
+                               src_non_local_pattern->rank.data(),
+                               src_non_local_pattern->rank.size())};
+
+    if (is_owner(exec_handler)) {
+        FatalErrorInFunction << "Not implemented" << abort(FatalError);
+        //  make_begin_and_ends_consecutive(
+        //      span_comm_pattern, merged_non_local.spans,
+        //     rank, ranks_per_gpu);
+    }
+
+    auto [gathered_non_local, is_local] = build_non_local_interfaces(
+        exec_handler, merged_local, merged_non_local);
+
+    // build vector with locality information
+    std::vector<std::pair<bool, label>> locality;
+    label ctr{0};
+
+    if (is_owner(exec_handler)) {
+        for (int i = 0; i < ranks_per_gpu; i++) {
+            label comm_rank = rank + i;
+            label num_elems = span_comm_pattern.recv_counts[comm_rank];
+            for (int j = 0; j < num_elems; j++) {
+                locality.emplace_back(is_local[ctr], comm_rank);
+                ctr++;
+            }
+        }
+    }
+    // All interfaces that are in is_local need to be accounted for
+    ASSERT_EQ(ctr, is_local.size());
+
+    // TODO FIXME the correct way would be to check the communication
+    // pattern where a particular face is in the send idxs
+    // since we have already the row id of the other side
+    // it should be doable. Or alternatively we know that we
+    // keep an interface together thus we can just count the idx up to the
+    // size. But we have to make sure that the interfaces are in the same
+    // order on both communication sides.
+    for (int i = 0; i < gathered_non_local.col_idxs.get_size(); i++) {
+        gathered_non_local.col_idxs.get_data()[i] = i;
+    }
+
+    LOG_1(verbose_, "done repartition sparsity pattern")
+    if (is_owner(exec_handler)) {
+        FatalErrorInFunction << "Not implemented" << abort(FatalError);
+        // auto new_local_spars_pattern = std::make_shared<SparsityPattern>(
+        //     exec, merged_local_dim, merged_local);
+        //
+        // auto new_non_local_spars_pattern =
+        //     std::make_shared<SparsityPattern>(
+        //         exec,
+        //         gko::dim<2>{merged_local_dim[0],
+        //                     gathered_non_local.row_idxs.get_size()},
+        //         gathered_non_local);
+        //
+        //         return std::make_tuple<std::shared_ptr<SparsityPattern>,
+        //                                std::shared_ptr<SparsityPattern>,
+        //                                std::vector<std::pair<bool,
+        //                                label>>>(
+        //             std::move(new_local_spars_pattern),
+        //             std::move(new_non_local_spars_pattern),
+        //             std::move(locality));
+    } else {
+        // auto new_local_spars_pattern =
+        //     std::make_shared<SparsityPattern>(exec);
+        //
+        // auto new_non_local_spars_pattern =
+        //     std::make_shared<SparsityPattern>(exec);
+        //
+        // // FIXME
+        // // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
+        // // comm.rank() <<
+        // //     " merged_local dim " <<
+        // new_local_spars_pattern->dim =
+        //     gko::dim<2>{0, 0};  // merged_local_dim;
+        // new_non_local_spars_pattern->dim =
+        //     gko::dim<2>{0, 0};  // merged_local_dim;
+        //
+        //
+        FatalErrorInFunction << "Not implemented" << abort(FatalError);
+        // return std::make_tuple<std::shared_ptr<SparsityPattern>,
+        //                        std::shared_ptr<SparsityPattern>,
+        //                        std::vector<std::pair<bool, label>>>(
+        //     std::move(new_local_spars_pattern),
+        //     std::move(new_non_local_spars_pattern),
+        //     std::vector<std::pair<bool, label>>{});
+    }
+    // TODO Some old documentation see what is still needed
+    // add interfaces
+    // to add interfaces we go through the non_local_sparsity
+    // pattern and check if interface is still non_local
+    //
+    // iterate interfaces and send to owner.
+    // on owner decide to move to local or non_local
+    // for sending to owner we can also use the all_to_all_v
+    // approach we than have
+    //
+    // row [ 4 , 8,  12 | 16 , 17, 18 ] <- local row
+    // col [ 1 , 2, 3 | 1, 2, 3 ] <- just the interface ctr
+    // from repartition_comm_pattern we could get
+    //
+    // or we split local / non-local first
+    // for this we could store ranges when computing
+    // repartition_comm_pattern
+    //
+}
+
 std::pair<SparsityPattern, std::vector<bool>>
 Repartitioner::build_non_local_interfaces(const ExecutorHandler &exec_handler,
                                           SparsityPattern &loc,
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 749e646f9..03724e450 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -69,6 +69,43 @@ TEST(Repartitioner, can_create_repartitioner)
     EXPECT_EQ(repartitioner.get_ranks_per_gpu(), ranks_per_gpu);
 }
 
+TEST(Repartitioner, has_correct_properties_for_1_rank)
+{
+    // Arrange
+    label ranks_per_gpu = 1;
+    label local_size = 10;
+    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
+
+    // Assert
+    EXPECT_EQ(
+        repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
+        local_size);
+
+    EXPECT_EQ(
+        repartitioner.is_owner(*exec),
+        true);
+}
+
+TEST(Repartitioner, has_correct_properties_for_4_rank)
+{
+    // Arrange
+    label ranks_per_gpu = 4;
+    label local_size = 10;
+    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
+    auto rank = repartitioner.get_rank(*exec);
+
+    // Assert
+    EXPECT_EQ(
+        repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
+        (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0 ));
+
+    EXPECT_EQ(
+        repartitioner.is_owner(*exec),
+        (rank == 0)? true: false);
+}
+
 int main(int argc, char *argv[])
 {
     int result = 0;

From ee329da2b09c2cfece6e8d17b972892f69157b28 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Sat, 17 Aug 2024 20:00:18 +0200
Subject: [PATCH 11/67] move impl to .C files

---
 include/OGL/Repartitioner.H         | 189 +--------------------------
 src/MatrixWrapper/SparsityPattern.C |   2 +-
 src/Repartitioner.C                 | 193 ++++++++++++++++++++++++++++
 3 files changed, 196 insertions(+), 188 deletions(-)

diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index d0d5310ee..f3affca90 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -109,6 +109,7 @@ public:
 
     label get_repart_size() const { return repart_size_; }
 
+    // TODO pass original dim as argument, if size can also be nnz
     gko::dim<2> get_repart_dim() const
     {
         return gko::dim<2>{static_cast<gko::size_type>(repart_size_),
@@ -153,191 +154,5 @@ public:
         std::shared_ptr<const CommunicationPattern> src_comm_pattern,
         std::shared_ptr<
             const gko::experimental::distributed::Partition<label, label>>
-            partition) const
-    {
-        using comm_size_type = label;
-        auto exec = src_comm_pattern->target_ids.get_executor();
-        auto comm = src_comm_pattern->get_comm();
-
-        if (ranks_per_gpu_ == 1) {
-            std::cout << __FILE__ << ":" << __LINE__ << " rank " << comm.rank()
-                      << *src_comm_pattern.get() << "\n";
-            return src_comm_pattern;
-        }
-
-
-        label rank = comm.rank();
-        bool owner = is_owner(exec_handler);
-
-        // Step 1. Check if communication partner is non-local after
-        // repartitioning. If it is non-local we keep it. Otherwise local
-        // communication can be discarded. Here non-local means:
-        // the communication target proc id is != repartition proc id
-        std::vector<label> target_ids{};
-        std::vector<label> target_sizes{};
-        std::vector<std::pair<gko::array<label>, comm_size_type>> send_idxs;
-        label communication_partner = src_comm_pattern->target_ids.get_size();
-        for (int i = 0; i < communication_partner; i++) {
-            label target_id = src_comm_pattern->target_ids.get_const_data()[i];
-            if (!reparts_to_local(exec_handler, target_id)) {
-                // communication pattern is non local, hence we keep it
-                // after repartitioning we now have to communicate with a
-                // different rank store new owner rank to which repart_target_*
-                // needs to be send to
-                target_ids.push_back(get_owner_rank(target_id));
-                target_sizes.push_back(
-                    src_comm_pattern->target_sizes.get_const_data()[i]);
-                send_idxs.push_back(src_comm_pattern->send_idxs[i]);
-            }
-        }
-
-        // send all remaining non local ids and sizes to the new
-        // owner rank
-        auto comm_pattern = compute_gather_to_owner_counts(
-            exec_handler, ranks_per_gpu_, target_ids.size());
-
-        auto gathered_target_ids = gather_labels_to_owner(
-            exec_handler, comm_pattern, target_ids.data(), target_ids.size());
-
-        auto gathered_target_sizes =
-            gather_labels_to_owner(exec_handler, comm_pattern,
-                                   target_sizes.data(), target_sizes.size());
-
-        // next the send_ixs need to be updated we send them piecewise since
-        // the send_idxs are a vector of gko::arrays
-        if (owner) {
-            label owner_recv_counts = comm_pattern.recv_counts[rank];
-            // retrieved from i-th neighbor
-            for (int i = 1; i < ranks_per_gpu_; i++) {
-                // how many gko::arrays to with send_indexes to receive
-                // from i-th neighbor
-                label recv_count = comm_pattern.recv_counts[rank + i];
-                for (int j = 0; j < recv_count; j++) {
-                    auto target_size =
-                        gathered_target_sizes[j + owner_recv_counts];
-                    std::vector<label> recv_buffer(target_size);
-
-                    // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
-                    // rank
-                    //           << " recv no  " << j
-                    //           << " recv from  " << rank + i
-                    //           << " recv_count " << recv_count
-                    //           << " target_size " << target_size
-                    //           << " \n";
-                    comm.recv(exec, recv_buffer.data(), target_size, rank + i,
-                              rank);
-
-                    // the new offset is
-                    auto offset = partition->get_range_bounds()[rank + i] -
-                                  partition->get_range_bounds()[rank];
-
-                    std::transform(recv_buffer.begin(), recv_buffer.end(),
-                                   recv_buffer.begin(),
-                                   [&](label idx) { return idx + offset; });
-
-                    auto target_id = gathered_target_ids[j + owner_recv_counts];
-                    send_idxs.push_back({gko::array<label>{
-                                             exec,
-                                             recv_buffer.begin(),
-                                             recv_buffer.end(),
-                                         },
-                                         target_id});
-                }
-            }
-        } else {
-            label owner = get_owner_rank(exec_handler);
-            for (int i = 0; i < comm_pattern.send_counts[owner]; i++) {
-                auto send_buffer = send_idxs[i].first;
-                // std::cout << __FILE__ << ":" << __LINE__ << " rank " << rank
-                //           << " send to owner " << owner
-                //           << " count " << i << " of "  <<  send_count
-                //           << " send_buffer.size " << send_buffer.get_size()
-                //           << " \n";
-                comm.send(exec, send_buffer.get_const_data(),
-                          send_buffer.get_size(), owner, owner);
-            }
-        }
-
-        // clear communication neighbors on non owning rank
-        if (!owner) {
-            // TODO NOTE should it be gathered_target_ids and is it needed?
-            target_ids.clear();
-            target_sizes.clear();
-            send_idxs.clear();
-        }
-
-
-        // early return if no communication partners are left
-        if (gathered_target_ids.size() == 0) {
-            return std::make_shared<CommunicationPattern>(
-                exec_handler,
-                gko::array<comm_size_type>{exec, target_ids.begin(),
-                                           target_ids.end()},
-                gko::array<comm_size_type>{exec, target_sizes.begin(),
-                                           target_sizes.end()},
-
-                send_idxs);
-        }
-
-        // merge communication
-        // we now might have communication partners multiple times, thus we can
-        // merge them sort repart_target_ids and use the sorting for
-        // target_sizes and repart_send_idxs
-        std::vector<label> merged_target_ids{};
-        std::vector<label> merged_target_sizes{};
-        std::vector<std::vector<label>> merged_send_idxs{};
-        auto p = sort_permutation(gathered_target_ids,
-                                  [](label a, label b) { return a < b; });
-
-        target_sizes = apply_permutation(gathered_target_sizes, p);
-        target_ids = apply_permutation(gathered_target_ids, p);
-        send_idxs = apply_permutation(send_idxs, p);
-
-        // Step 4.
-        // Merge communication pattern with corresponding neighbours
-        merged_target_ids.push_back(target_ids[0]);
-        merged_target_sizes.push_back(target_sizes[0]);
-        merged_send_idxs.emplace_back(std::vector<label>(
-            send_idxs[0].first.get_data(),
-            send_idxs[0].first.get_data() + send_idxs[0].first.get_size()));
-
-
-        for (int i = 1; i < target_ids.size(); i++) {
-            // communicates with same target rank
-            // thus we have only have to adapt the number
-            // of elements and the send_ixs
-            auto *send_idx_begin = send_idxs[i].first.get_data();
-            auto *send_idx_end =
-                send_idxs[i].first.get_data() + send_idxs[i].first.get_size();
-            if (target_ids[i] == merged_target_ids.back()) {
-                merged_target_sizes.back() += target_sizes[i];
-                merged_send_idxs.back().insert(merged_send_idxs.back().end(),
-                                               send_idx_begin, send_idx_end);
-            } else {
-                merged_target_ids.push_back(target_ids[i]);
-                merged_target_sizes.push_back(target_sizes[i]);
-                merged_send_idxs.emplace_back(
-                    std::vector<label>(send_idx_begin, send_idx_end));
-            }
-        }
-
-        // recompute send_idxs
-        send_idxs.clear();
-
-        for (int i = 0; i < merged_target_ids.size(); i++) {
-            label target_id = merged_target_ids[i];
-            send_idxs.emplace_back(std::pair<gko::array<label>, comm_size_type>{
-                gko::array<label>{exec, merged_send_idxs[i].begin(),
-                                  merged_send_idxs[i].end()},
-                target_id});
-        }
-
-        return std::make_shared<CommunicationPattern>(
-            exec_handler,
-            gko::array<comm_size_type>{exec, merged_target_ids.begin(),
-                                       merged_target_ids.end()},
-            gko::array<comm_size_type>{exec, merged_target_sizes.begin(),
-                                       merged_target_sizes.end()},
-            send_idxs);
-    }
+            partition) const;
 };
diff --git a/src/MatrixWrapper/SparsityPattern.C b/src/MatrixWrapper/SparsityPattern.C
index b56038e6f..29209b812 100644
--- a/src/MatrixWrapper/SparsityPattern.C
+++ b/src/MatrixWrapper/SparsityPattern.C
@@ -23,4 +23,4 @@ void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
     }
 }
 
-    }
+}  // namespace Foam
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index e8e602efc..4a11e0a52 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -294,3 +294,196 @@ Repartitioner::build_non_local_interfaces(const ExecutorHandler &exec_handler,
     //       SparsityPattern{rows, cols, ldu_mapping, begins, ends, ranks},
     //       is_local);
 }
+
+
+std::shared_ptr<const CommunicationPattern>
+Repartitioner::repartition_comm_pattern(
+    const ExecutorHandler &exec_handler,
+    std::shared_ptr<const CommunicationPattern> src_comm_pattern,
+    std::shared_ptr<
+        const gko::experimental::distributed::Partition<label, label>>
+        partition) const
+{
+    using comm_size_type = label;
+    auto exec = src_comm_pattern->target_ids.get_executor();
+    auto comm = src_comm_pattern->get_comm();
+
+    if (ranks_per_gpu_ == 1) {
+        std::cout << __FILE__ << ":" << __LINE__ << " rank " << comm.rank()
+                  << *src_comm_pattern.get() << "\n";
+        return src_comm_pattern;
+    }
+
+
+    label rank = comm.rank();
+    bool owner = is_owner(exec_handler);
+
+    // Step 1. Check if communication partner is non-local after
+    // repartitioning. If it is non-local we keep it. Otherwise local
+    // communication can be discarded. Here non-local means:
+    // the communication target proc id is != repartition proc id
+    std::vector<label> target_ids{};
+    std::vector<label> target_sizes{};
+    std::vector<std::pair<gko::array<label>, comm_size_type>> send_idxs;
+    label communication_partner = src_comm_pattern->target_ids.get_size();
+    for (int i = 0; i < communication_partner; i++) {
+        label target_id = src_comm_pattern->target_ids.get_const_data()[i];
+        if (!reparts_to_local(exec_handler, target_id)) {
+            // communication pattern is non local, hence we keep it
+            // after repartitioning we now have to communicate with a
+            // different rank store new owner rank to which repart_target_*
+            // needs to be send to
+            target_ids.push_back(get_owner_rank(target_id));
+            target_sizes.push_back(
+                src_comm_pattern->target_sizes.get_const_data()[i]);
+            send_idxs.push_back(src_comm_pattern->send_idxs[i]);
+        }
+    }
+
+    // send all remaining non local ids and sizes to the new
+    // owner rank
+    auto comm_pattern = compute_gather_to_owner_counts(
+        exec_handler, ranks_per_gpu_, target_ids.size());
+
+    auto gathered_target_ids = gather_labels_to_owner(
+        exec_handler, comm_pattern, target_ids.data(), target_ids.size());
+
+    auto gathered_target_sizes = gather_labels_to_owner(
+        exec_handler, comm_pattern, target_sizes.data(), target_sizes.size());
+
+    // next the send_ixs need to be updated we send them piecewise since
+    // the send_idxs are a vector of gko::arrays
+    if (owner) {
+        label owner_recv_counts = comm_pattern.recv_counts[rank];
+        // retrieved from i-th neighbor
+        for (int i = 1; i < ranks_per_gpu_; i++) {
+            // how many gko::arrays to with send_indexes to receive
+            // from i-th neighbor
+            label recv_count = comm_pattern.recv_counts[rank + i];
+            for (int j = 0; j < recv_count; j++) {
+                auto target_size = gathered_target_sizes[j + owner_recv_counts];
+                std::vector<label> recv_buffer(target_size);
+
+                // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
+                // rank
+                //           << " recv no  " << j
+                //           << " recv from  " << rank + i
+                //           << " recv_count " << recv_count
+                //           << " target_size " << target_size
+                //           << " \n";
+                comm.recv(exec, recv_buffer.data(), target_size, rank + i,
+                          rank);
+
+                // the new offset is
+                auto offset = partition->get_range_bounds()[rank + i] -
+                              partition->get_range_bounds()[rank];
+
+                std::transform(recv_buffer.begin(), recv_buffer.end(),
+                               recv_buffer.begin(),
+                               [&](label idx) { return idx + offset; });
+
+                auto target_id = gathered_target_ids[j + owner_recv_counts];
+                send_idxs.push_back({gko::array<label>{
+                                         exec,
+                                         recv_buffer.begin(),
+                                         recv_buffer.end(),
+                                     },
+                                     target_id});
+            }
+        }
+    } else {
+        label owner = get_owner_rank(exec_handler);
+        for (int i = 0; i < comm_pattern.send_counts[owner]; i++) {
+            auto send_buffer = send_idxs[i].first;
+            // std::cout << __FILE__ << ":" << __LINE__ << " rank " << rank
+            //           << " send to owner " << owner
+            //           << " count " << i << " of "  <<  send_count
+            //           << " send_buffer.size " << send_buffer.get_size()
+            //           << " \n";
+            comm.send(exec, send_buffer.get_const_data(),
+                      send_buffer.get_size(), owner, owner);
+        }
+    }
+
+    // clear communication neighbors on non owning rank
+    if (!owner) {
+        // TODO NOTE should it be gathered_target_ids and is it needed?
+        target_ids.clear();
+        target_sizes.clear();
+        send_idxs.clear();
+    }
+
+
+    // early return if no communication partners are left
+    if (gathered_target_ids.size() == 0) {
+        return std::make_shared<CommunicationPattern>(
+            exec_handler,
+            gko::array<comm_size_type>{exec, target_ids.begin(),
+                                       target_ids.end()},
+            gko::array<comm_size_type>{exec, target_sizes.begin(),
+                                       target_sizes.end()},
+
+            send_idxs);
+    }
+
+    // merge communication
+    // we now might have communication partners multiple times, thus we can
+    // merge them sort repart_target_ids and use the sorting for
+    // target_sizes and repart_send_idxs
+    std::vector<label> merged_target_ids{};
+    std::vector<label> merged_target_sizes{};
+    std::vector<std::vector<label>> merged_send_idxs{};
+    auto p = sort_permutation(gathered_target_ids,
+                              [](label a, label b) { return a < b; });
+
+    target_sizes = apply_permutation(gathered_target_sizes, p);
+    target_ids = apply_permutation(gathered_target_ids, p);
+    send_idxs = apply_permutation(send_idxs, p);
+
+    // Step 4.
+    // Merge communication pattern with corresponding neighbours
+    merged_target_ids.push_back(target_ids[0]);
+    merged_target_sizes.push_back(target_sizes[0]);
+    merged_send_idxs.emplace_back(std::vector<label>(
+        send_idxs[0].first.get_data(),
+        send_idxs[0].first.get_data() + send_idxs[0].first.get_size()));
+
+
+    for (int i = 1; i < target_ids.size(); i++) {
+        // communicates with same target rank
+        // thus we have only have to adapt the number
+        // of elements and the send_ixs
+        auto *send_idx_begin = send_idxs[i].first.get_data();
+        auto *send_idx_end =
+            send_idxs[i].first.get_data() + send_idxs[i].first.get_size();
+        if (target_ids[i] == merged_target_ids.back()) {
+            merged_target_sizes.back() += target_sizes[i];
+            merged_send_idxs.back().insert(merged_send_idxs.back().end(),
+                                           send_idx_begin, send_idx_end);
+        } else {
+            merged_target_ids.push_back(target_ids[i]);
+            merged_target_sizes.push_back(target_sizes[i]);
+            merged_send_idxs.emplace_back(
+                std::vector<label>(send_idx_begin, send_idx_end));
+        }
+    }
+
+    // recompute send_idxs
+    send_idxs.clear();
+
+    for (int i = 0; i < merged_target_ids.size(); i++) {
+        label target_id = merged_target_ids[i];
+        send_idxs.emplace_back(std::pair<gko::array<label>, comm_size_type>{
+            gko::array<label>{exec, merged_send_idxs[i].begin(),
+                              merged_send_idxs[i].end()},
+            target_id});
+    }
+
+    return std::make_shared<CommunicationPattern>(
+        exec_handler,
+        gko::array<comm_size_type>{exec, merged_target_ids.begin(),
+                                   merged_target_ids.end()},
+        gko::array<comm_size_type>{exec, merged_target_sizes.begin(),
+                                   merged_target_sizes.end()},
+        send_idxs);
+}

From 732479a543093eb9bdd54d0dbe833673cdf3bc31 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Sun, 18 Aug 2024 09:43:28 +0200
Subject: [PATCH 12/67] add repart comm pattern basic unit test

---
 include/OGL/CommunicationPattern.H            |  16 +++
 .../OGL/DevicePersistent/ExecutorHandler.H    |   2 +
 include/OGL/Repartitioner.H                   |  15 +--
 src/Repartitioner.C                           |  11 +-
 unitTests/Repartitioner.C                     | 122 ++++++++++++++++--
 5 files changed, 143 insertions(+), 23 deletions(-)

diff --git a/include/OGL/CommunicationPattern.H b/include/OGL/CommunicationPattern.H
index f33a61b83..70561d956 100644
--- a/include/OGL/CommunicationPattern.H
+++ b/include/OGL/CommunicationPattern.H
@@ -129,6 +129,22 @@ struct CommunicationPattern {
         // ASSERT_EQ(target_ids.get_size(), send_idxs.size());
     }
 
+    CommunicationPattern(
+        const ExecutorHandler &exec, const std::vector<label>& ranks,
+        const std::vector<std::vector<label>>& rows)
+        : exec_handler(exec),
+          target_ids(exec.get_ref_exec(), ranks.begin(), ranks.end()),
+          target_sizes(exec.get_ref_exec(), ranks.size())
+    {
+        ASSERT_EQ(ranks.size(), rows.size());
+
+        for (int i=0; i<rows.size();i++){
+            target_sizes.get_data()[i] = rows[i].size();
+            gko::array<label> row_array (exec.get_ref_exec(),rows[i].begin(), rows[i].end());
+            send_idxs.push_back({row_array, ranks[i]});
+        }
+    }
+
     const gko::experimental::mpi::communicator &get_comm() const
     {
         return *exec_handler.get_communicator().get();
diff --git a/include/OGL/DevicePersistent/ExecutorHandler.H b/include/OGL/DevicePersistent/ExecutorHandler.H
index a82ab57f6..390031a32 100644
--- a/include/OGL/DevicePersistent/ExecutorHandler.H
+++ b/include/OGL/DevicePersistent/ExecutorHandler.H
@@ -187,6 +187,8 @@ public:
     {
         return this->device_comm_;
     }
+
+    label get_rank() const { return get_communicator().get()->rank(); };
 };
 
 using PersistentExecutor = ExecutorHandler;
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index f3affca90..8cc0f1f97 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -73,13 +73,14 @@ public:
     /* returns the owner rank for a given rank */
     label get_owner_rank(const ExecutorHandler &exec_handler) const
     {
-        return get_owner_rank(get_rank(exec_handler));
+        return get_owner_rank(exec_handler.get_rank());
     };
 
     /* returns if current rank is an owner  */
     bool is_owner(const ExecutorHandler &exec_handler) const
     {
-        return get_rank(exec_handler) == get_owner_rank(get_rank(exec_handler));
+        return exec_handler.get_rank() ==
+               get_owner_rank(exec_handler.get_rank());
     };
 
     /* @brief check if the given rank gets local after repartitioning
@@ -91,12 +92,6 @@ public:
                compute_owner_rank(rank, ranks_per_gpu_);
     };
 
-    /* shortcut to current rank */
-    label get_rank(const ExecutorHandler &exec_handler) const
-    {
-        return exec_handler.get_communicator().get()->rank();
-    };
-
     label get_ranks_per_gpu() const { return ranks_per_gpu_; }
 
     /* @brief computes the size after repartitioning
@@ -149,6 +144,10 @@ public:
         std::shared_ptr<SparsityPattern> src_local_pattern,
         std::shared_ptr<SparsityPattern> src_non_local_pattern) const;
 
+    /* @brief repartion a communication pattern
+    **
+    ** returns the repartitioned communcation pattern
+    */
     std::shared_ptr<const CommunicationPattern> repartition_comm_pattern(
         const ExecutorHandler &exec_handler,
         std::shared_ptr<const CommunicationPattern> src_comm_pattern,
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index 4a11e0a52..c5c8aeaf0 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -38,7 +38,7 @@ Repartitioner::repartition_sparsity(
     // debugging reasons
     auto exec = exec_handler.get_ref_exec();
     auto comm = *exec_handler.get_communicator().get();
-    label rank = get_rank(exec_handler);
+    label rank = exec_handler.get_rank();
     label owner_rank = get_owner_rank(exec_handler);
     label ranks_per_gpu = ranks_per_gpu_;
     // TODO dont copy
@@ -304,16 +304,13 @@ Repartitioner::repartition_comm_pattern(
         const gko::experimental::distributed::Partition<label, label>>
         partition) const
 {
-    using comm_size_type = label;
-    auto exec = src_comm_pattern->target_ids.get_executor();
-    auto comm = src_comm_pattern->get_comm();
-
     if (ranks_per_gpu_ == 1) {
-        std::cout << __FILE__ << ":" << __LINE__ << " rank " << comm.rank()
-                  << *src_comm_pattern.get() << "\n";
         return src_comm_pattern;
     }
 
+    using comm_size_type = label;
+    auto exec = src_comm_pattern->target_ids.get_executor();
+    auto comm = src_comm_pattern->get_comm();
 
     label rank = comm.rank();
     bool owner = is_owner(exec_handler);
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 03724e450..45ec25a8f 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -82,28 +82,134 @@ TEST(Repartitioner, has_correct_properties_for_1_rank)
         repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
         local_size);
 
-    EXPECT_EQ(
-        repartitioner.is_owner(*exec),
-        true);
+    EXPECT_EQ(repartitioner.is_owner(*exec), true);
+}
+
+TEST(Repartitioner, can_repartition_comm_pattern_for_1_rank)
+{
+    // Arrange
+    label ranks_per_gpu = 1;
+    label local_size = 10;
+    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
+    auto rank = exec->get_rank();
+    auto ref_exec = exec->get_ref_exec();
+
+    // rank:     0     1     2     3
+    // cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
+    // cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- global row ids
+    std::vector<std::vector<label>> ids{{1}, {0, 2}, {1, 3}, {2}};
+    std::vector<std::vector<std::vector<label>>> rows{
+        {{1}}, {{0}, {1}}, {{0}, {1}}, {{0}}};
+
+    auto comm_pattern =
+        std::make_shared<CommunicationPattern>(*exec, ids[rank], rows[rank]);
+
+    auto comm = exec->get_communicator();
+    auto partition = gko::share(
+        gko::experimental::distributed::build_partition_from_local_size<label,
+                                                                        label>(
+            ref_exec, *comm.get(), 2));
+
+    // Act
+    auto repart_comm_pattern =
+        repartitioner.repartition_comm_pattern(*exec, comm_pattern, partition);
+
+    // Assert
+    // communication ranks are the same as before
+    std::vector<label> res_ids(
+        repart_comm_pattern->target_ids.get_const_data(),
+        repart_comm_pattern->target_ids.get_const_data() +
+            repart_comm_pattern->target_ids.get_size());
+    EXPECT_EQ(res_ids, ids[rank]);
+
+    std::vector<label> res_sizes(
+        repart_comm_pattern->target_sizes.get_const_data(),
+        repart_comm_pattern->target_sizes.get_const_data() +
+            repart_comm_pattern->target_sizes.get_size());
+    std::vector<std::vector<label>> exp_sizes{{1}, {1, 1}, {1, 1}, {1}};
+    EXPECT_EQ(res_sizes, exp_sizes[rank]);
 }
 
-TEST(Repartitioner, has_correct_properties_for_4_rank)
+TEST(Repartitioner, has_correct_properties_for_4_ranks)
 {
     // Arrange
     label ranks_per_gpu = 4;
     label local_size = 10;
     auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
     auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
-    auto rank = repartitioner.get_rank(*exec);
+    auto rank = exec->get_rank();
 
     // Assert
     EXPECT_EQ(
         repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
-        (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0 ));
+        (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0));
+
+    EXPECT_EQ(repartitioner.is_owner(*exec), (rank == 0) ? true : false);
+}
+
+TEST(Repartitioner, can_repartition_comm_pattern_for_4_rank)
+{
+    // Arrange
+    label ranks_per_gpu = 4;
+    label local_size = 10;
+    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
+    auto rank = exec->get_rank();
+    auto ref_exec = exec->get_ref_exec();
+
+    // rank:     0     1     2     3
+    // cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
+    // cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- global row ids
+    std::vector<std::vector<label>> ids{{1}, {0, 2}, {1, 3}, {2}};
+    std::vector<std::vector<std::vector<label>>> rows{
+        {{1}}, {{0}, {1}}, {{0}, {1}}, {{0}}};
+
+    auto comm_pattern =
+        std::make_shared<CommunicationPattern>(*exec, ids[rank], rows[rank]);
+
+    auto comm = exec->get_communicator();
+    auto partition = gko::share(
+        gko::experimental::distributed::build_partition_from_local_size<label,
+                                                                        label>(
+            ref_exec, *comm.get(), 2));
+
+    // Act
+    auto repart_comm_pattern =
+        repartitioner.repartition_comm_pattern(*exec, comm_pattern, partition);
+
+    // Assert
+    std::vector<label> res_ids(
+        repart_comm_pattern->target_ids.get_const_data(),
+        repart_comm_pattern->target_ids.get_const_data() +
+            repart_comm_pattern->target_ids.get_size());
+    // no communication ranks should be left
+    EXPECT_EQ(res_ids.size(), 0);
+
+    std::vector<label> res_sizes(
+        repart_comm_pattern->target_sizes.get_const_data(),
+        repart_comm_pattern->target_sizes.get_const_data() +
+            repart_comm_pattern->target_sizes.get_size());
+    // thus communication sizes is also empty
+    EXPECT_EQ(res_sizes.size(), 0);
+    EXPECT_EQ(repart_comm_pattern->send_idxs.size(), 0);
+}
+
+TEST(Repartitioner, has_correct_properties_for_2_ranks)
+{
+    // Arrange
+    label ranks_per_gpu = 2;
+    label local_size = 10;
+    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
+    auto rank = exec->get_rank();
 
+    // Assert
     EXPECT_EQ(
-        repartitioner.is_owner(*exec),
-        (rank == 0)? true: false);
+        repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
+        (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0));
+
+    EXPECT_EQ(repartitioner.is_owner(*exec), (rank % 2 == 0) ? true : false);
 }
 
 int main(int argc, char *argv[])

From bf46e9a5434a0cc1b3e41de69650a34a1274d34c Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Sun, 18 Aug 2024 21:33:48 +0200
Subject: [PATCH 13/67] add 2 of 4 repartition test

---
 src/CommunicationPattern.C | 10 ++++----
 src/Repartitioner.C        | 20 ++++------------
 unitTests/Repartitioner.C  | 49 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 19 deletions(-)

diff --git a/src/CommunicationPattern.C b/src/CommunicationPattern.C
index 5f92b38be..61598afbb 100644
--- a/src/CommunicationPattern.C
+++ b/src/CommunicationPattern.C
@@ -159,6 +159,10 @@ std::vector<label> gather_labels_to_owner(const ExecutorHandler &exec_handler,
                                           const label *send_buffer,
                                           label send_size, label offset)
 {
+    auto exec = exec_handler.get_ref_exec();
+    auto comm = *exec_handler.get_communicator().get();
+    auto rank = comm.rank();
+
     std::vector<label> send_buffer_copy;
     // create a copy if offset is needed
     if (offset > 0) {
@@ -168,18 +172,16 @@ std::vector<label> gather_labels_to_owner(const ExecutorHandler &exec_handler,
                        [&](label idx) { return idx + offset; });
     }
 
-    auto &[send_counts, recv_counts, send_offsets, recv_offsets] = comm_pattern;
+    auto &[send_counts, send_offsets, recv_counts, recv_offsets] = comm_pattern;
     auto total_recv_size =
         std::accumulate(recv_counts.begin(), recv_counts.end(), 0);
 
-    auto exec = exec_handler.get_ref_exec();
-    auto comm = *exec_handler.get_communicator().get();
     std::vector<label> recv_buffer(total_recv_size);
-    auto rank = comm.rank();
     comm.all_to_all_v(
         exec, (offset > 0) ? send_buffer_copy.data() : send_buffer,
         send_counts.data(), send_offsets.data(), recv_buffer.data(),
         recv_counts.data(), recv_offsets.data());
+
     return recv_buffer;
 }
 
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index c5c8aeaf0..c215013cb 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -316,9 +316,10 @@ Repartitioner::repartition_comm_pattern(
     bool owner = is_owner(exec_handler);
 
     // Step 1. Check if communication partner is non-local after
-    // repartitioning. If it is non-local we keep it. Otherwise local
-    // communication can be discarded. Here non-local means:
-    // the communication target proc id is != repartition proc id
+    // repartitioning. If it is non-local we keep it. Otherwise communcation
+    // partners that are local after repartitioning can be discarded. Here
+    // non-local means: the communication target rank id is != repartitioned
+    // rank id
     std::vector<label> target_ids{};
     std::vector<label> target_sizes{};
     std::vector<std::pair<gko::array<label>, comm_size_type>> send_idxs;
@@ -348,6 +349,7 @@ Repartitioner::repartition_comm_pattern(
     auto gathered_target_sizes = gather_labels_to_owner(
         exec_handler, comm_pattern, target_sizes.data(), target_sizes.size());
 
+
     // next the send_ixs need to be updated we send them piecewise since
     // the send_idxs are a vector of gko::arrays
     if (owner) {
@@ -361,13 +363,6 @@ Repartitioner::repartition_comm_pattern(
                 auto target_size = gathered_target_sizes[j + owner_recv_counts];
                 std::vector<label> recv_buffer(target_size);
 
-                // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
-                // rank
-                //           << " recv no  " << j
-                //           << " recv from  " << rank + i
-                //           << " recv_count " << recv_count
-                //           << " target_size " << target_size
-                //           << " \n";
                 comm.recv(exec, recv_buffer.data(), target_size, rank + i,
                           rank);
 
@@ -392,11 +387,6 @@ Repartitioner::repartition_comm_pattern(
         label owner = get_owner_rank(exec_handler);
         for (int i = 0; i < comm_pattern.send_counts[owner]; i++) {
             auto send_buffer = send_idxs[i].first;
-            // std::cout << __FILE__ << ":" << __LINE__ << " rank " << rank
-            //           << " send to owner " << owner
-            //           << " count " << i << " of "  <<  send_count
-            //           << " send_buffer.size " << send_buffer.get_size()
-            //           << " \n";
             comm.send(exec, send_buffer.get_const_data(),
                       send_buffer.get_size(), owner, owner);
         }
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 45ec25a8f..97b339dd7 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -131,6 +131,54 @@ TEST(Repartitioner, can_repartition_comm_pattern_for_1_rank)
     EXPECT_EQ(res_sizes, exp_sizes[rank]);
 }
 
+TEST(Repartitioner, can_repartition_comm_pattern_for_2_rank)
+{
+    // Arrange
+    label ranks_per_gpu = 2;
+    label local_size = 10;
+    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
+    auto rank = exec->get_rank();
+    auto ref_exec = exec->get_ref_exec();
+
+    // rank:     0     1     2     3
+    // cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
+    // cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- local row ids
+    std::vector<std::vector<label>> ids{{1}, {0, 2}, {1, 3}, {2}};
+    std::vector<std::vector<std::vector<label>>> rows{
+        {{1}}, {{0}, {1}}, {{0}, {1}}, {{0}}};
+
+    auto comm_pattern =
+        std::make_shared<CommunicationPattern>(*exec, ids[rank], rows[rank]);
+
+    auto comm = exec->get_communicator();
+    auto partition = gko::share(
+        gko::experimental::distributed::build_partition_from_local_size<label,
+                                                                        label>(
+            ref_exec, *comm.get(), 2));
+
+    // Act
+    auto repart_comm_pattern =
+        repartitioner.repartition_comm_pattern(*exec, comm_pattern, partition);
+
+    // Assert
+    std::vector<label> res_ids(
+        repart_comm_pattern->target_ids.get_const_data(),
+        repart_comm_pattern->target_ids.get_const_data() +
+            repart_comm_pattern->target_ids.get_size());
+    // on non owner no communication ranks should be left
+    std::vector<std::vector<label>> exp_ids {{2}, {}, {0}, {}};
+    EXPECT_EQ(res_ids, exp_ids[rank]);
+
+    std::vector<label> res_sizes(
+        repart_comm_pattern->target_sizes.get_const_data(),
+        repart_comm_pattern->target_sizes.get_const_data() +
+            repart_comm_pattern->target_sizes.get_size());
+    std::vector<std::vector<label>> exp_sizes {{1}, {}, {1}, {}};
+    // thus communication sizes is also empty
+    EXPECT_EQ(res_sizes, exp_sizes[rank]);
+}
+
 TEST(Repartitioner, has_correct_properties_for_4_ranks)
 {
     // Arrange
@@ -212,6 +260,7 @@ TEST(Repartitioner, has_correct_properties_for_2_ranks)
     EXPECT_EQ(repartitioner.is_owner(*exec), (rank % 2 == 0) ? true : false);
 }
 
+
 int main(int argc, char *argv[])
 {
     int result = 0;

From 05e62b639842fded6151ebf9f412450e853c9b6d Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Sun, 18 Aug 2024 21:38:19 +0200
Subject: [PATCH 14/67] add printer for vector of labels

---
 include/OGL/common.H |  5 +++++
 src/common.C         | 25 +++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/include/OGL/common.H b/include/OGL/common.H
index 4172f97f3..4fb8eafda 100644
--- a/include/OGL/common.H
+++ b/include/OGL/common.H
@@ -100,6 +100,11 @@ std::ostream &operator<<(
     std::ostream &os,
     const std::shared_ptr<gko::matrix::Dense<scalar>> array_in);
 
+std::ostream &operator<<(
+    std::ostream &os,
+    const std::vector<label>& in);
+
+
 void export_system(const word fieldName, const gko::matrix::Csr<scalar> *A,
                    const gko::matrix::Dense<scalar> *x,
                    const gko::matrix::Dense<scalar> *b, const word time);
diff --git a/src/common.C b/src/common.C
index bf29c463e..23ba8894a 100644
--- a/src/common.C
+++ b/src/common.C
@@ -169,4 +169,29 @@ std::ostream &operator<<(std::ostream &os,
     }
     return os;
 }
+
+std::ostream &operator<<(
+    std::ostream &os,
+    const std::vector<label>& in){
+    label size = in.size();
+    if (size == 0 ){ return os;}
+    os << size << " elements [";
+    if (size > 100) {
+        for (label i = 0; i < 9; i++) {
+            os << in[i] << ", ";
+        }
+        os << in[10] << " ... ";
+        for (label i = size - 9; i < size - 1; i++) {
+            os << in[i] << ", ";
+        }
+        os << in[size - 1] << "]\n";
+    } else {
+        for (label i = 0; i < size - 1; i++) {
+            os << "(" << i << ", " << in[i] << ") ";
+        }
+        os << "(" << size - 1 << ", " << in[size - 1] << ")]\n";
+    }
+    return os;
+
+};
 }  // namespace Foam

From 02ea8d4a9b2804626d84234d3c06481028adae8f Mon Sep 17 00:00:00 2001
From: chihta-wang <wangchihta@gmail.com>
Date: Mon, 12 Aug 2024 18:04:11 +0200
Subject: [PATCH 15/67] Add some description for get_coeffients function

---
 include/OGL/MatrixWrapper/Combination.H | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/include/OGL/MatrixWrapper/Combination.H b/include/OGL/MatrixWrapper/Combination.H
index 0d67288c7..f3835c142 100644
--- a/include/OGL/MatrixWrapper/Combination.H
+++ b/include/OGL/MatrixWrapper/Combination.H
@@ -66,6 +66,11 @@ public:
         return *this;
     }
 
+     /**
+     * Returns a list of coefficients of the combination.
+     *
+     * @return a list of coefficients
+     */
     const std::vector<std::shared_ptr<const gko::LinOp>> &get_coefficients()
         const noexcept
     {

From eff21e86ae073a5914c406688b7b6824c6481adb Mon Sep 17 00:00:00 2001
From: chihta-wang <wangchihta@gmail.com>
Date: Mon, 12 Aug 2024 18:17:19 +0200
Subject: [PATCH 16/67] - Add a unit test for Combination Matrix by using read
 function to read matrix data from a file

- Add matrix data for the unit test in Matrix Market Exchange Format (.mtx)

- Modify the CMake configuration file for HostMatrix in order to
copy matrix to the build directory of unit tests
---
 unitTests/MatrixWrapper/CMakeLists.txt |  6 ++++
 unitTests/MatrixWrapper/Combination.C  | 44 +++++++++++++++++++++++++-
 unitTests/MatrixWrapper/data/A.mtx     |  5 +++
 3 files changed, 54 insertions(+), 1 deletion(-)
 create mode 100644 unitTests/MatrixWrapper/data/A.mtx

diff --git a/unitTests/MatrixWrapper/CMakeLists.txt b/unitTests/MatrixWrapper/CMakeLists.txt
index ab74a9483..05eec064d 100644
--- a/unitTests/MatrixWrapper/CMakeLists.txt
+++ b/unitTests/MatrixWrapper/CMakeLists.txt
@@ -2,5 +2,11 @@
 #
 # SPDX-License-Identifier: GPL-3.0-or-later
 
+# Copy the data files to the execution directory
+configure_file(data/A.mtx data/A.mtx COPYONLY)
+
 ogl_unit_test(matrixConversion "HostMatrix.C")
 ogl_unit_test(combination "Combination.C")
+
+# Set working directory for the unit test suite of Combination Matrix
+set_tests_properties(combination PROPERTIES WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
diff --git a/unitTests/MatrixWrapper/Combination.C b/unitTests/MatrixWrapper/Combination.C
index 6cba3c735..eaa51c4d9 100644
--- a/unitTests/MatrixWrapper/Combination.C
+++ b/unitTests/MatrixWrapper/Combination.C
@@ -10,6 +10,9 @@
 
 #include <ginkgo/ginkgo.hpp>
 
+#include <fstream>
+
+#include <string>
 
 TEST(Combination, CanCreateEmptyCombination)
 {
@@ -54,7 +57,46 @@ TEST(Combination, CanCreateCombinationWithLinOpVector)
     ASSERT_EQ(x->at(4, 0), 0.0);
 }
 
-TEST(Combination, CanCreateConvertToCsr)
+TEST(Combination, CombinationMatrxWithRegularMatrix)
+{
+    auto dim = gko::dim<2>{5, 5};
+    // gko::matrix_data<double, int> m1
+    // (dim, {{0, 0, 2}, {1, 1, 0}, {2, 3, 5}});
+
+    auto exec = gko::share(gko::ReferenceExecutor::create());
+    // auto m1linop =
+    //     gko::share(gko::matrix::Csr<scalar, label>::create(exec, dim));
+    // m1linop->read(m1);
+    using ValueType = double;
+    using IndexType = int;
+    using mtx = gko::matrix::Csr<ValueType, IndexType>;
+
+    auto m1linop = gko::share(
+        gko::read<mtx>(std::ifstream("data/A.mtx"), exec));
+
+    std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m1linop};
+
+    auto cmb = CombinationMatrix<gko::matrix::Csr<scalar, label>>::create(
+        exec, dim, linops);
+
+    ASSERT_EQ(cmb->get_size(), gko::dim<2>(5, 5));
+    ASSERT_EQ(cmb->get_coefficients().size(), 2);
+    ASSERT_EQ(cmb->get_operators().size(), 2);
+
+    auto b = gko::matrix::Dense<scalar>::create(exec, gko::dim<2>{5, 1});
+    b->fill(1.0);
+    auto x = gko::matrix::Dense<scalar>::create(exec, gko::dim<2>{5, 1});
+    x->fill(0.0);
+    cmb->apply(b, x);
+
+    ASSERT_EQ(x->at(0, 0), 4.0);
+    ASSERT_EQ(x->at(1, 0), 0.0);
+    ASSERT_EQ(x->at(2, 0), 10.0);
+    ASSERT_EQ(x->at(3, 0), 0.0);
+    ASSERT_EQ(x->at(4, 0), 0.0);
+}
+
+TEST(Combination, CanConvertToCsr)
 {
     auto dim = gko::dim<2>{5, 5};
     gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, 0}, {2, 3, 5}});
diff --git a/unitTests/MatrixWrapper/data/A.mtx b/unitTests/MatrixWrapper/data/A.mtx
new file mode 100644
index 000000000..d813213b2
--- /dev/null
+++ b/unitTests/MatrixWrapper/data/A.mtx
@@ -0,0 +1,5 @@
+%%MatrixMarket matrix coordinate integer general
+5 5 3
+1 1 2
+2 2 0
+3 4 5

From d493a8dad7e344871f5245857d9d8b7113022746 Mon Sep 17 00:00:00 2001
From: chihta-wang <wangchihta@gmail.com>
Date: Wed, 14 Aug 2024 10:14:13 +0200
Subject: [PATCH 17/67] - Add a mtx file for storing the vector b - Enable
 reading the b vector data from the file - Modify the CMake file to copy
 vector data to the build directory

---
 unitTests/MatrixWrapper/CMakeLists.txt | 1 +
 unitTests/MatrixWrapper/Combination.C  | 9 ++++++---
 unitTests/MatrixWrapper/data/b.mtx     | 7 +++++++
 3 files changed, 14 insertions(+), 3 deletions(-)
 create mode 100644 unitTests/MatrixWrapper/data/b.mtx

diff --git a/unitTests/MatrixWrapper/CMakeLists.txt b/unitTests/MatrixWrapper/CMakeLists.txt
index 05eec064d..5cf05d2f0 100644
--- a/unitTests/MatrixWrapper/CMakeLists.txt
+++ b/unitTests/MatrixWrapper/CMakeLists.txt
@@ -4,6 +4,7 @@
 
 # Copy the data files to the execution directory
 configure_file(data/A.mtx data/A.mtx COPYONLY)
+configure_file(data/b.mtx data/b.mtx COPYONLY)
 
 ogl_unit_test(matrixConversion "HostMatrix.C")
 ogl_unit_test(combination "Combination.C")
diff --git a/unitTests/MatrixWrapper/Combination.C b/unitTests/MatrixWrapper/Combination.C
index eaa51c4d9..43fc769f1 100644
--- a/unitTests/MatrixWrapper/Combination.C
+++ b/unitTests/MatrixWrapper/Combination.C
@@ -57,7 +57,7 @@ TEST(Combination, CanCreateCombinationWithLinOpVector)
     ASSERT_EQ(x->at(4, 0), 0.0);
 }
 
-TEST(Combination, CombinationMatrxWithRegularMatrix)
+TEST(Combination, CombinationMatrixWithSparseMatrix)
 {
     auto dim = gko::dim<2>{5, 5};
     // gko::matrix_data<double, int> m1
@@ -70,9 +70,12 @@ TEST(Combination, CombinationMatrxWithRegularMatrix)
     using ValueType = double;
     using IndexType = int;
     using mtx = gko::matrix::Csr<ValueType, IndexType>;
+    using vec = gko::matrix::Dense<ValueType>;
 
     auto m1linop = gko::share(
         gko::read<mtx>(std::ifstream("data/A.mtx"), exec));
+    
+    auto b = gko::read<vec>(std::ifstream("data/b.mtx"), exec);
 
     std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m1linop};
 
@@ -83,8 +86,8 @@ TEST(Combination, CombinationMatrxWithRegularMatrix)
     ASSERT_EQ(cmb->get_coefficients().size(), 2);
     ASSERT_EQ(cmb->get_operators().size(), 2);
 
-    auto b = gko::matrix::Dense<scalar>::create(exec, gko::dim<2>{5, 1});
-    b->fill(1.0);
+    // auto b = gko::matrix::Dense<scalar>::create(exec, gko::dim<2>{5, 1});
+    // b->fill(1.0);
     auto x = gko::matrix::Dense<scalar>::create(exec, gko::dim<2>{5, 1});
     x->fill(0.0);
     cmb->apply(b, x);
diff --git a/unitTests/MatrixWrapper/data/b.mtx b/unitTests/MatrixWrapper/data/b.mtx
new file mode 100644
index 000000000..01420bf6f
--- /dev/null
+++ b/unitTests/MatrixWrapper/data/b.mtx
@@ -0,0 +1,7 @@
+%%MatrixMarket matrix array real general
+5 1
+1
+1
+1
+1
+1
\ No newline at end of file

From 4fdce852dcd446edd607fb3d7946e5b196a6af3b Mon Sep 17 00:00:00 2001
From: chihta-wang <wangchihta@gmail.com>
Date: Wed, 14 Aug 2024 11:21:25 +0200
Subject: [PATCH 18/67] - Use a sparse matrix with floating-point numbers to
 replace the original matrix which consists of integers - Create a mtx file
 for storing a vector - Remove commented lines of code and reorganize the unit
 test

---
 unitTests/MatrixWrapper/CMakeLists.txt    |  2 +-
 unitTests/MatrixWrapper/Combination.C     | 40 +++++++++++------------
 unitTests/MatrixWrapper/data/A.mtx        |  5 ---
 unitTests/MatrixWrapper/data/A_sparse.mtx | 10 ++++++
 unitTests/MatrixWrapper/data/b.mtx        |  8 ++---
 5 files changed, 35 insertions(+), 30 deletions(-)
 delete mode 100644 unitTests/MatrixWrapper/data/A.mtx
 create mode 100644 unitTests/MatrixWrapper/data/A_sparse.mtx

diff --git a/unitTests/MatrixWrapper/CMakeLists.txt b/unitTests/MatrixWrapper/CMakeLists.txt
index 5cf05d2f0..b4ac0325a 100644
--- a/unitTests/MatrixWrapper/CMakeLists.txt
+++ b/unitTests/MatrixWrapper/CMakeLists.txt
@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: GPL-3.0-or-later
 
 # Copy the data files to the execution directory
-configure_file(data/A.mtx data/A.mtx COPYONLY)
+configure_file(data/A_sparse.mtx data/A_sparse.mtx COPYONLY)
 configure_file(data/b.mtx data/b.mtx COPYONLY)
 
 ogl_unit_test(matrixConversion "HostMatrix.C")
diff --git a/unitTests/MatrixWrapper/Combination.C b/unitTests/MatrixWrapper/Combination.C
index 43fc769f1..ede535a88 100644
--- a/unitTests/MatrixWrapper/Combination.C
+++ b/unitTests/MatrixWrapper/Combination.C
@@ -57,46 +57,46 @@ TEST(Combination, CanCreateCombinationWithLinOpVector)
     ASSERT_EQ(x->at(4, 0), 0.0);
 }
 
-TEST(Combination, CombinationMatrixWithSparseMatrix)
+TEST(Combination, CanCreateCombinationMatrixWithSparseRealMatrix)
 {
-    auto dim = gko::dim<2>{5, 5};
-    // gko::matrix_data<double, int> m1
-    // (dim, {{0, 0, 2}, {1, 1, 0}, {2, 3, 5}});
-
-    auto exec = gko::share(gko::ReferenceExecutor::create());
-    // auto m1linop =
-    //     gko::share(gko::matrix::Csr<scalar, label>::create(exec, dim));
-    // m1linop->read(m1);
-    using ValueType = double;
-    using IndexType = int;
+    // Arrange
+    using ValueType = scalar;
+    using IndexType = label;
     using mtx = gko::matrix::Csr<ValueType, IndexType>;
     using vec = gko::matrix::Dense<ValueType>;
+    
+    int dim_size = 5;
+    auto dim = gko::dim<2>{dim_size, dim_size};
+
+    auto exec = gko::share(gko::ReferenceExecutor::create());
 
     auto m1linop = gko::share(
-        gko::read<mtx>(std::ifstream("data/A.mtx"), exec));
+        gko::read<mtx>(std::ifstream("data/A_sparse.mtx"), exec));
     
     auto b = gko::read<vec>(std::ifstream("data/b.mtx"), exec);
 
     std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m1linop};
 
+    // Act
     auto cmb = CombinationMatrix<gko::matrix::Csr<scalar, label>>::create(
         exec, dim, linops);
 
-    ASSERT_EQ(cmb->get_size(), gko::dim<2>(5, 5));
+    // Assert
+    ASSERT_EQ(cmb->get_size(), gko::dim<2>(dim_size, dim_size));
     ASSERT_EQ(cmb->get_coefficients().size(), 2);
     ASSERT_EQ(cmb->get_operators().size(), 2);
 
-    // auto b = gko::matrix::Dense<scalar>::create(exec, gko::dim<2>{5, 1});
-    // b->fill(1.0);
-    auto x = gko::matrix::Dense<scalar>::create(exec, gko::dim<2>{5, 1});
+    // Act
+    auto x = gko::matrix::Dense<scalar>::create(exec, gko::dim<2>{dim_size, 1});
     x->fill(0.0);
     cmb->apply(b, x);
 
-    ASSERT_EQ(x->at(0, 0), 4.0);
+    // Assert
+    ASSERT_EQ(x->at(0, 0), 38.0);
     ASSERT_EQ(x->at(1, 0), 0.0);
-    ASSERT_EQ(x->at(2, 0), 10.0);
-    ASSERT_EQ(x->at(3, 0), 0.0);
-    ASSERT_EQ(x->at(4, 0), 0.0);
+    ASSERT_EQ(x->at(2, 0), 2.0);
+    ASSERT_EQ(x->at(3, 0), -1414.4);
+    ASSERT_EQ(x->at(4, 0), 96.0);
 }
 
 TEST(Combination, CanConvertToCsr)
diff --git a/unitTests/MatrixWrapper/data/A.mtx b/unitTests/MatrixWrapper/data/A.mtx
deleted file mode 100644
index d813213b2..000000000
--- a/unitTests/MatrixWrapper/data/A.mtx
+++ /dev/null
@@ -1,5 +0,0 @@
-%%MatrixMarket matrix coordinate integer general
-5 5 3
-1 1 2
-2 2 0
-3 4 5
diff --git a/unitTests/MatrixWrapper/data/A_sparse.mtx b/unitTests/MatrixWrapper/data/A_sparse.mtx
new file mode 100644
index 000000000..5b029c6d1
--- /dev/null
+++ b/unitTests/MatrixWrapper/data/A_sparse.mtx
@@ -0,0 +1,10 @@
+%%MatrixMarket matrix coordinate real general
+5 5 8
+1 1 1
+1 4 6
+2 2 10.5
+3 3 0.5
+4 2 250.5
+4 4 -280
+4 5 33.2
+5 5 12
diff --git a/unitTests/MatrixWrapper/data/b.mtx b/unitTests/MatrixWrapper/data/b.mtx
index 01420bf6f..bed718eaf 100644
--- a/unitTests/MatrixWrapper/data/b.mtx
+++ b/unitTests/MatrixWrapper/data/b.mtx
@@ -1,7 +1,7 @@
 %%MatrixMarket matrix array real general
 5 1
 1
-1
-1
-1
-1
\ No newline at end of file
+0
+2
+3
+4
\ No newline at end of file

From 422ece2899cf3eb8f1d364448701750894a14032 Mon Sep 17 00:00:00 2001
From: chihta-wang <wangchihta@gmail.com>
Date: Wed, 14 Aug 2024 11:50:30 +0200
Subject: [PATCH 19/67] - Add a unit test for Combination Matrix which stores
 L, D, and U matrices - Add 3 mtx files for storing L, D, and U matrices,
 respectively - Modify the CMake file to copy L, D, and U matrices over to the
 build directory

---
 unitTests/MatrixWrapper/CMakeLists.txt |  3 ++
 unitTests/MatrixWrapper/Combination.C  | 49 ++++++++++++++++++++++++++
 unitTests/MatrixWrapper/data/D.mtx     |  7 ++++
 unitTests/MatrixWrapper/data/L.mtx     |  3 ++
 unitTests/MatrixWrapper/data/U.mtx     |  4 +++
 5 files changed, 66 insertions(+)
 create mode 100644 unitTests/MatrixWrapper/data/D.mtx
 create mode 100644 unitTests/MatrixWrapper/data/L.mtx
 create mode 100644 unitTests/MatrixWrapper/data/U.mtx

diff --git a/unitTests/MatrixWrapper/CMakeLists.txt b/unitTests/MatrixWrapper/CMakeLists.txt
index b4ac0325a..c288f036e 100644
--- a/unitTests/MatrixWrapper/CMakeLists.txt
+++ b/unitTests/MatrixWrapper/CMakeLists.txt
@@ -5,6 +5,9 @@
 # Copy the data files to the execution directory
 configure_file(data/A_sparse.mtx data/A_sparse.mtx COPYONLY)
 configure_file(data/b.mtx data/b.mtx COPYONLY)
+configure_file(data/L.mtx data/L.mtx COPYONLY)
+configure_file(data/D.mtx data/D.mtx COPYONLY)
+configure_file(data/U.mtx data/U.mtx COPYONLY)
 
 ogl_unit_test(matrixConversion "HostMatrix.C")
 ogl_unit_test(combination "Combination.C")
diff --git a/unitTests/MatrixWrapper/Combination.C b/unitTests/MatrixWrapper/Combination.C
index ede535a88..d8d1e82d0 100644
--- a/unitTests/MatrixWrapper/Combination.C
+++ b/unitTests/MatrixWrapper/Combination.C
@@ -99,6 +99,55 @@ TEST(Combination, CanCreateCombinationMatrixWithSparseRealMatrix)
     ASSERT_EQ(x->at(4, 0), 96.0);
 }
 
+
+TEST(Combination, CanCreateCombinationMatrixWithOpenFOAMLDUMatrix)
+{
+    // Arrange
+    using ValueType = scalar;
+    using IndexType = label;
+    using mtx = gko::matrix::Csr<ValueType, IndexType>;
+    using vec = gko::matrix::Dense<ValueType>;
+    
+    int dim_size = 5;
+    auto dim = gko::dim<2>{dim_size, dim_size};
+
+    auto exec = gko::share(gko::ReferenceExecutor::create());
+
+    auto Llinop = gko::share(
+        gko::read<mtx>(std::ifstream("data/L.mtx"), exec));
+
+    auto Dlinop = gko::share(
+        gko::read<mtx>(std::ifstream("data/D.mtx"), exec));
+    
+    auto Ulinop = gko::share(
+        gko::read<mtx>(std::ifstream("data/U.mtx"), exec));
+
+    auto b = gko::read<vec>(std::ifstream("data/b.mtx"), exec);
+
+    std::vector<std::shared_ptr<const gko::LinOp>> linops{Llinop, Dlinop, Ulinop};
+
+    // Act
+    auto cmb = CombinationMatrix<gko::matrix::Csr<scalar, label>>::create(
+        exec, dim, linops);
+
+    // Assert
+    ASSERT_EQ(cmb->get_size(), gko::dim<2>(dim_size, dim_size));
+    ASSERT_EQ(cmb->get_coefficients().size(), 3);
+    ASSERT_EQ(cmb->get_operators().size(), 3);
+
+    // Act
+    auto x = gko::matrix::Dense<scalar>::create(exec, gko::dim<2>{dim_size, 1});
+    x->fill(0.0);
+    cmb->apply(b, x);
+
+    // Assert
+    ASSERT_EQ(x->at(0, 0), 19.0);
+    ASSERT_EQ(x->at(1, 0), 0.0);
+    ASSERT_EQ(x->at(2, 0), 1.0);
+    ASSERT_EQ(x->at(3, 0), -707.2);
+    ASSERT_EQ(x->at(4, 0), 48.0);
+}
+
 TEST(Combination, CanConvertToCsr)
 {
     auto dim = gko::dim<2>{5, 5};
diff --git a/unitTests/MatrixWrapper/data/D.mtx b/unitTests/MatrixWrapper/data/D.mtx
new file mode 100644
index 000000000..784d16dec
--- /dev/null
+++ b/unitTests/MatrixWrapper/data/D.mtx
@@ -0,0 +1,7 @@
+%%MatrixMarket matrix coordinate real general
+5 5 5
+1 1 1
+2 2 10.5
+3 3 0.5
+4 4 -280
+5 5 12
\ No newline at end of file
diff --git a/unitTests/MatrixWrapper/data/L.mtx b/unitTests/MatrixWrapper/data/L.mtx
new file mode 100644
index 000000000..738df39b9
--- /dev/null
+++ b/unitTests/MatrixWrapper/data/L.mtx
@@ -0,0 +1,3 @@
+%%MatrixMarket matrix coordinate real general
+5 5 1
+4 2 250.5
\ No newline at end of file
diff --git a/unitTests/MatrixWrapper/data/U.mtx b/unitTests/MatrixWrapper/data/U.mtx
new file mode 100644
index 000000000..c0b5f6167
--- /dev/null
+++ b/unitTests/MatrixWrapper/data/U.mtx
@@ -0,0 +1,4 @@
+%%MatrixMarket matrix coordinate real general
+5 5 2
+1 4 6
+4 5 33.2
\ No newline at end of file

From 0b035b54b66e4b0b8142c6ce15df640fcea56b42 Mon Sep 17 00:00:00 2001
From: chihta-wang <wangchihta@gmail.com>
Date: Wed, 14 Aug 2024 15:34:37 +0200
Subject: [PATCH 20/67] Finish the verification of entry values in the unit
 test for the function convert_to COO format

---
 unitTests/MatrixWrapper/Combination.C | 31 +++++++++++++++++++++------
 1 file changed, 24 insertions(+), 7 deletions(-)

diff --git a/unitTests/MatrixWrapper/Combination.C b/unitTests/MatrixWrapper/Combination.C
index d8d1e82d0..2ab8d3405 100644
--- a/unitTests/MatrixWrapper/Combination.C
+++ b/unitTests/MatrixWrapper/Combination.C
@@ -148,27 +148,44 @@ TEST(Combination, CanCreateCombinationMatrixWithOpenFOAMLDUMatrix)
     ASSERT_EQ(x->at(4, 0), 48.0);
 }
 
-TEST(Combination, CanConvertToCsr)
+TEST(Combination, CanConvertToCoo)
 {
-    auto dim = gko::dim<2>{5, 5};
+    // Arrange
+    using ValueType = scalar;
+    using IndexType = label;
+    using InputFormat = gko::matrix::Csr<ValueType, IndexType>;
+    using OutputFormat = gko::matrix::Coo<ValueType, IndexType>;
+
+    int dim_size = 5;
+
+    auto dim = gko::dim<2>{dim_size, dim_size};
     gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, 0}, {2, 3, 5}});
     gko::matrix_data<double, int> m2(dim, {{3, 3, 1}, {4, 4, 2}});
 
     auto exec = gko::share(gko::ReferenceExecutor::create());
     auto m1linop =
-        gko::share(gko::matrix::Csr<scalar, label>::create(exec, dim));
+        gko::share(InputFormat::create(exec, dim));
     m1linop->read(m1);
 
     auto m2linop =
-        gko::share(gko::matrix::Csr<scalar, label>::create(exec, dim));
+        gko::share(InputFormat::create(exec, dim));
     m2linop->read(m2);
 
     std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m2linop};
 
-    auto cmb = CombinationMatrix<gko::matrix::Csr<scalar, label>>::create(
+    auto cmb = CombinationMatrix<InputFormat>::create(
         exec, dim, linops);
 
-    auto out = CombinationMatrix<gko::matrix::Csr<scalar, label>>::create(exec);
+    auto out = OutputFormat::create(
+        exec);
 
-    cmb->convert_to(out);
+    // Act
+    cmb->convert_to(out.get());
+
+    // Assert
+    ASSERT_EQ(out->get_values()[0], 2);
+    ASSERT_EQ(out->get_values()[1], 0);
+    ASSERT_EQ(out->get_values()[2], 5);
+    ASSERT_EQ(out->get_values()[3], 1);
+    ASSERT_EQ(out->get_values()[0], 2);
 }

From 040e7e8927faa7b1115023b803fce0bc63b27105 Mon Sep 17 00:00:00 2001
From: chihta-wang <wangchihta@gmail.com>
Date: Wed, 14 Aug 2024 16:00:01 +0200
Subject: [PATCH 21/67] Finish the unit test for the function convert_to COO
 format. - Add assertion tests for verifying row and column indices

---
 unitTests/MatrixWrapper/Combination.C | 71 +++++++++++++++++++++++++--
 1 file changed, 66 insertions(+), 5 deletions(-)

diff --git a/unitTests/MatrixWrapper/Combination.C b/unitTests/MatrixWrapper/Combination.C
index 2ab8d3405..9225be523 100644
--- a/unitTests/MatrixWrapper/Combination.C
+++ b/unitTests/MatrixWrapper/Combination.C
@@ -159,8 +159,8 @@ TEST(Combination, CanConvertToCoo)
     int dim_size = 5;
 
     auto dim = gko::dim<2>{dim_size, dim_size};
-    gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, 0}, {2, 3, 5}});
-    gko::matrix_data<double, int> m2(dim, {{3, 3, 1}, {4, 4, 2}});
+    gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, -3}, {2, 3, 5}});
+    gko::matrix_data<double, int> m2(dim, {{3, 2, 4}, {3, 3, 1}, {4, 4, 2}});
 
     auto exec = gko::share(gko::ReferenceExecutor::create());
     auto m1linop =
@@ -183,9 +183,70 @@ TEST(Combination, CanConvertToCoo)
     cmb->convert_to(out.get());
 
     // Assert
+    // Verify entry values
     ASSERT_EQ(out->get_values()[0], 2);
-    ASSERT_EQ(out->get_values()[1], 0);
+    ASSERT_EQ(out->get_values()[1], -3);
     ASSERT_EQ(out->get_values()[2], 5);
-    ASSERT_EQ(out->get_values()[3], 1);
-    ASSERT_EQ(out->get_values()[0], 2);
+    ASSERT_EQ(out->get_values()[3], 4);
+    ASSERT_EQ(out->get_values()[4], 1);
+    ASSERT_EQ(out->get_values()[5], 2);
+
+    // Verify row indices
+    ASSERT_EQ(out->get_row_idxs()[0], 0);
+    ASSERT_EQ(out->get_row_idxs()[1], 1);
+    ASSERT_EQ(out->get_row_idxs()[2], 2);
+    ASSERT_EQ(out->get_row_idxs()[3], 3);
+    ASSERT_EQ(out->get_row_idxs()[4], 3);
+    ASSERT_EQ(out->get_row_idxs()[5], 4);
+
+    // Verify column indices
+    ASSERT_EQ(out->get_col_idxs()[0], 0);
+    ASSERT_EQ(out->get_col_idxs()[1], 1);
+    ASSERT_EQ(out->get_col_idxs()[2], 3);
+    ASSERT_EQ(out->get_col_idxs()[3], 2);
+    ASSERT_EQ(out->get_col_idxs()[4], 3);
+    ASSERT_EQ(out->get_col_idxs()[5], 4);
 }
+
+// TEST(Combination, CanConvertToCsr)
+// {
+//     // Arrange
+//     using ValueType = scalar;
+//     using IndexType = label;
+//     using InputFormat = gko::matrix::Coo<ValueType, IndexType>;
+//     using OutputFormat = gko::matrix::Csr<ValueType, IndexType>;
+
+//     int dim_size = 5;
+
+//     auto dim = gko::dim<2>{dim_size, dim_size};
+//     gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, 0}, {2, 3, 5}});
+//     gko::matrix_data<double, int> m2(dim, {{3, 3, 1}, {4, 4, 2}});
+
+//     auto exec = gko::share(gko::ReferenceExecutor::create());
+//     auto m1linop =
+//         gko::share(InputFormat::create(exec, dim));
+//     m1linop->read(m1);
+
+//     auto m2linop =
+//         gko::share(InputFormat::create(exec, dim));
+//     m2linop->read(m2);
+
+//     std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m2linop};
+
+//     auto cmb = CombinationMatrix<InputFormat>::create(
+//         exec, dim, linops);
+
+//     auto out = OutputFormat::create(
+//         exec);
+
+//     // Act
+//     cmb->convert_to(out.get());
+
+//     // Assert
+//     ASSERT_EQ(out->get_values()[0], 2);
+//     ASSERT_EQ(out->get_values()[1], 0);
+//     ASSERT_EQ(out->get_values()[2], 5);
+//     ASSERT_EQ(out->get_values()[3], 1);
+//     ASSERT_EQ(out->get_values()[0], 2);
+    
+// }
\ No newline at end of file

From b416ff0b9413c514a2426f5e3da795cb2122cfe5 Mon Sep 17 00:00:00 2001
From: chihta-wang <wangchihta@gmail.com>
Date: Wed, 14 Aug 2024 16:17:50 +0200
Subject: [PATCH 22/67] - Implement the function convert_to CSR - Add a unit
 test for the function

---
 include/OGL/MatrixWrapper/Combination.H |  22 +++++-
 unitTests/MatrixWrapper/Combination.C   | 100 ++++++++++++++----------
 2 files changed, 79 insertions(+), 43 deletions(-)

diff --git a/include/OGL/MatrixWrapper/Combination.H b/include/OGL/MatrixWrapper/Combination.H
index f3835c142..0a1354c06 100644
--- a/include/OGL/MatrixWrapper/Combination.H
+++ b/include/OGL/MatrixWrapper/Combination.H
@@ -110,7 +110,27 @@ public:
         result->read(mat_data);
     }
 
-    void convert_to(gko::matrix::Csr<scalar, label> *result) const override {}
+    /*
+     * @warning assumes that all entries are unique
+     */
+    void convert_to(gko::matrix::Csr<scalar, label> *result) const override 
+    {
+        auto exec = this->get_executor();
+        auto operators = this->get_operators();
+
+        std::vector<gko::matrix_data_entry<scalar, label>> data;
+        for (auto op : operators) {
+            auto tmp_mat_data = gko::matrix_data<scalar>();
+            gko::as<InnerMatrixType>(op)->write(tmp_mat_data);
+            data.insert(data.end(), tmp_mat_data.nonzeros.begin(),
+                        tmp_mat_data.nonzeros.end());
+        }
+        auto mat_data = gko::matrix_data<scalar>();
+        mat_data.nonzeros = data;
+        mat_data.size = this->get_size();
+
+        result->read(mat_data);
+    }
 
     void move_to(gko::matrix::Coo<scalar, label> *result) override {}
 
diff --git a/unitTests/MatrixWrapper/Combination.C b/unitTests/MatrixWrapper/Combination.C
index 9225be523..b822e9644 100644
--- a/unitTests/MatrixWrapper/Combination.C
+++ b/unitTests/MatrixWrapper/Combination.C
@@ -208,45 +208,61 @@ TEST(Combination, CanConvertToCoo)
     ASSERT_EQ(out->get_col_idxs()[5], 4);
 }
 
-// TEST(Combination, CanConvertToCsr)
-// {
-//     // Arrange
-//     using ValueType = scalar;
-//     using IndexType = label;
-//     using InputFormat = gko::matrix::Coo<ValueType, IndexType>;
-//     using OutputFormat = gko::matrix::Csr<ValueType, IndexType>;
-
-//     int dim_size = 5;
-
-//     auto dim = gko::dim<2>{dim_size, dim_size};
-//     gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, 0}, {2, 3, 5}});
-//     gko::matrix_data<double, int> m2(dim, {{3, 3, 1}, {4, 4, 2}});
-
-//     auto exec = gko::share(gko::ReferenceExecutor::create());
-//     auto m1linop =
-//         gko::share(InputFormat::create(exec, dim));
-//     m1linop->read(m1);
-
-//     auto m2linop =
-//         gko::share(InputFormat::create(exec, dim));
-//     m2linop->read(m2);
-
-//     std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m2linop};
-
-//     auto cmb = CombinationMatrix<InputFormat>::create(
-//         exec, dim, linops);
-
-//     auto out = OutputFormat::create(
-//         exec);
-
-//     // Act
-//     cmb->convert_to(out.get());
-
-//     // Assert
-//     ASSERT_EQ(out->get_values()[0], 2);
-//     ASSERT_EQ(out->get_values()[1], 0);
-//     ASSERT_EQ(out->get_values()[2], 5);
-//     ASSERT_EQ(out->get_values()[3], 1);
-//     ASSERT_EQ(out->get_values()[0], 2);
-    
-// }
\ No newline at end of file
+TEST(Combination, CanConvertToCsr)
+{
+    // Arrange
+    using ValueType = scalar;
+    using IndexType = label;
+    using InputFormat = gko::matrix::Coo<ValueType, IndexType>;
+    using OutputFormat = gko::matrix::Csr<ValueType, IndexType>;
+
+    int dim_size = 5;
+
+    auto dim = gko::dim<2>{dim_size, dim_size};
+    gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, -3}, {2, 3, 5}});
+    gko::matrix_data<double, int> m2(dim, {{3, 2, 4}, {3, 3, 1}, {4, 4, 2}});
+
+    auto exec = gko::share(gko::ReferenceExecutor::create());
+    auto m1linop =
+        gko::share(InputFormat::create(exec, dim));
+    m1linop->read(m1);
+
+    auto m2linop =
+        gko::share(InputFormat::create(exec, dim));
+    m2linop->read(m2);
+
+    std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m2linop};
+
+    auto cmb = CombinationMatrix<InputFormat>::create(
+        exec, dim, linops);
+
+    auto out = OutputFormat::create(
+        exec);
+
+    // Act
+    cmb->convert_to(out.get());
+
+    // Assert
+    // Verify entry values
+    ASSERT_EQ(out->get_values()[0], 2);
+    ASSERT_EQ(out->get_values()[1], -3);
+    ASSERT_EQ(out->get_values()[2], 5);
+    ASSERT_EQ(out->get_values()[3], 4);
+    ASSERT_EQ(out->get_values()[4], 1);
+    ASSERT_EQ(out->get_values()[5], 2);
+
+    // Verify column indices
+    ASSERT_EQ(out->get_col_idxs()[0], 0);
+    ASSERT_EQ(out->get_col_idxs()[1], 1);
+    ASSERT_EQ(out->get_col_idxs()[2], 3);
+    ASSERT_EQ(out->get_col_idxs()[3], 2);
+    ASSERT_EQ(out->get_col_idxs()[4], 3);
+    ASSERT_EQ(out->get_col_idxs()[5], 4);
+
+    // Verify row indices
+    ASSERT_EQ(out->get_row_ptrs()[0], 0);
+    ASSERT_EQ(out->get_row_ptrs()[1], 1);
+    ASSERT_EQ(out->get_row_ptrs()[2], 2);
+    ASSERT_EQ(out->get_row_ptrs()[3], 3);
+    ASSERT_EQ(out->get_row_ptrs()[4], 5);
+}
\ No newline at end of file

From c84ec5eeb9c356ef55165bf945abd917838e81f2 Mon Sep 17 00:00:00 2001
From: chihta-wang <wangchihta@gmail.com>
Date: Wed, 14 Aug 2024 17:06:57 +0200
Subject: [PATCH 23/67] - Implement the function move_to COO format for
 CombinationMatrix class - Add a unit test for that function

---
 include/OGL/MatrixWrapper/Combination.H | 10 ++++-
 unitTests/MatrixWrapper/Combination.C   | 60 +++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/include/OGL/MatrixWrapper/Combination.H b/include/OGL/MatrixWrapper/Combination.H
index 0a1354c06..dc52bf067 100644
--- a/include/OGL/MatrixWrapper/Combination.H
+++ b/include/OGL/MatrixWrapper/Combination.H
@@ -132,7 +132,15 @@ public:
         result->read(mat_data);
     }
 
-    void move_to(gko::matrix::Coo<scalar, label> *result) override {}
+    /*
+     * @warning assumes that all entries are unique
+     */
+    void move_to(gko::matrix::Coo<scalar, label> *result) override 
+    {
+        this->convert_to(result);
+        this->comb_ = nullptr;
+        this->set_size(gko::dim<2>(0, 0));
+    }
 
     void move_to(gko::matrix::Csr<scalar, label> *result) override {}
 
diff --git a/unitTests/MatrixWrapper/Combination.C b/unitTests/MatrixWrapper/Combination.C
index b822e9644..7b1c3e37a 100644
--- a/unitTests/MatrixWrapper/Combination.C
+++ b/unitTests/MatrixWrapper/Combination.C
@@ -265,4 +265,64 @@ TEST(Combination, CanConvertToCsr)
     ASSERT_EQ(out->get_row_ptrs()[2], 2);
     ASSERT_EQ(out->get_row_ptrs()[3], 3);
     ASSERT_EQ(out->get_row_ptrs()[4], 5);
+}
+
+TEST(Combination, CanMoveToCoo)
+{
+    // Arrange
+    using ValueType = scalar;
+    using IndexType = label;
+    using InputFormat = gko::matrix::Csr<ValueType, IndexType>;
+    using OutputFormat = gko::matrix::Coo<ValueType, IndexType>;
+
+    int dim_size = 5;
+
+    auto dim = gko::dim<2>{dim_size, dim_size};
+    gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, -3}, {2, 3, 5}});
+    gko::matrix_data<double, int> m2(dim, {{3, 2, 4}, {3, 3, 1}, {4, 4, 2}});
+
+    auto exec = gko::share(gko::ReferenceExecutor::create());
+    auto m1linop =
+        gko::share(InputFormat::create(exec, dim));
+    m1linop->read(m1);
+
+    auto m2linop =
+        gko::share(InputFormat::create(exec, dim));
+    m2linop->read(m2);
+
+    std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m2linop};
+
+    auto cmb = CombinationMatrix<InputFormat>::create(
+        exec, dim, linops);
+
+    auto out = OutputFormat::create(
+        exec);
+
+    // Act
+    cmb->move_to(out.get());
+
+    // Assert
+    // Verify entry values
+    ASSERT_EQ(out->get_values()[0], 2);
+    ASSERT_EQ(out->get_values()[1], -3);
+    ASSERT_EQ(out->get_values()[2], 5);
+    ASSERT_EQ(out->get_values()[3], 4);
+    ASSERT_EQ(out->get_values()[4], 1);
+    ASSERT_EQ(out->get_values()[5], 2);
+
+    // Verify row indices
+    ASSERT_EQ(out->get_row_idxs()[0], 0);
+    ASSERT_EQ(out->get_row_idxs()[1], 1);
+    ASSERT_EQ(out->get_row_idxs()[2], 2);
+    ASSERT_EQ(out->get_row_idxs()[3], 3);
+    ASSERT_EQ(out->get_row_idxs()[4], 3);
+    ASSERT_EQ(out->get_row_idxs()[5], 4);
+
+    // Verify column indices
+    ASSERT_EQ(out->get_col_idxs()[0], 0);
+    ASSERT_EQ(out->get_col_idxs()[1], 1);
+    ASSERT_EQ(out->get_col_idxs()[2], 3);
+    ASSERT_EQ(out->get_col_idxs()[3], 2);
+    ASSERT_EQ(out->get_col_idxs()[4], 3);
+    ASSERT_EQ(out->get_col_idxs()[5], 4);
 }
\ No newline at end of file

From b3b7032963545d8a2a19184edade63bde1ab7143 Mon Sep 17 00:00:00 2001
From: chihta-wang <wangchihta@gmail.com>
Date: Wed, 14 Aug 2024 17:15:33 +0200
Subject: [PATCH 24/67] - Implement the function move_to CSR format for
 CombinationMatrix class - Add a unit test for that function

---
 include/OGL/MatrixWrapper/Combination.H | 10 ++++-
 unitTests/MatrixWrapper/Combination.C   | 60 +++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/include/OGL/MatrixWrapper/Combination.H b/include/OGL/MatrixWrapper/Combination.H
index dc52bf067..13b86a8cb 100644
--- a/include/OGL/MatrixWrapper/Combination.H
+++ b/include/OGL/MatrixWrapper/Combination.H
@@ -142,7 +142,15 @@ public:
         this->set_size(gko::dim<2>(0, 0));
     }
 
-    void move_to(gko::matrix::Csr<scalar, label> *result) override {}
+    /*
+     * @warning assumes that all entries are unique
+     */
+    void move_to(gko::matrix::Csr<scalar, label> *result) override 
+    {
+        this->convert_to(result);
+        this->comb_ = nullptr;
+        this->set_size(gko::dim<2>(0, 0));
+    }
 
 protected:
     CombinationMatrix(std::shared_ptr<const gko::Executor> exec)
diff --git a/unitTests/MatrixWrapper/Combination.C b/unitTests/MatrixWrapper/Combination.C
index 7b1c3e37a..e849bf945 100644
--- a/unitTests/MatrixWrapper/Combination.C
+++ b/unitTests/MatrixWrapper/Combination.C
@@ -325,4 +325,64 @@ TEST(Combination, CanMoveToCoo)
     ASSERT_EQ(out->get_col_idxs()[3], 2);
     ASSERT_EQ(out->get_col_idxs()[4], 3);
     ASSERT_EQ(out->get_col_idxs()[5], 4);
+}
+
+
+TEST(Combination, CanMoveToCsr)
+{
+    // Arrange
+    using ValueType = scalar;
+    using IndexType = label;
+    using InputFormat = gko::matrix::Coo<ValueType, IndexType>;
+    using OutputFormat = gko::matrix::Csr<ValueType, IndexType>;
+
+    int dim_size = 5;
+
+    auto dim = gko::dim<2>{dim_size, dim_size};
+    gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, -3}, {2, 3, 5}});
+    gko::matrix_data<double, int> m2(dim, {{3, 2, 4}, {3, 3, 1}, {4, 4, 2}});
+
+    auto exec = gko::share(gko::ReferenceExecutor::create());
+    auto m1linop =
+        gko::share(InputFormat::create(exec, dim));
+    m1linop->read(m1);
+
+    auto m2linop =
+        gko::share(InputFormat::create(exec, dim));
+    m2linop->read(m2);
+
+    std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m2linop};
+
+    auto cmb = CombinationMatrix<InputFormat>::create(
+        exec, dim, linops);
+
+    auto out = OutputFormat::create(
+        exec);
+
+    // Act
+    cmb->move_to(out.get());
+
+    // Assert
+    // Verify entry values
+    ASSERT_EQ(out->get_values()[0], 2);
+    ASSERT_EQ(out->get_values()[1], -3);
+    ASSERT_EQ(out->get_values()[2], 5);
+    ASSERT_EQ(out->get_values()[3], 4);
+    ASSERT_EQ(out->get_values()[4], 1);
+    ASSERT_EQ(out->get_values()[5], 2);
+
+    // Verify column indices
+    ASSERT_EQ(out->get_col_idxs()[0], 0);
+    ASSERT_EQ(out->get_col_idxs()[1], 1);
+    ASSERT_EQ(out->get_col_idxs()[2], 3);
+    ASSERT_EQ(out->get_col_idxs()[3], 2);
+    ASSERT_EQ(out->get_col_idxs()[4], 3);
+    ASSERT_EQ(out->get_col_idxs()[5], 4);
+
+    // Verify row indices
+    ASSERT_EQ(out->get_row_ptrs()[0], 0);
+    ASSERT_EQ(out->get_row_ptrs()[1], 1);
+    ASSERT_EQ(out->get_row_ptrs()[2], 2);
+    ASSERT_EQ(out->get_row_ptrs()[3], 3);
+    ASSERT_EQ(out->get_row_ptrs()[4], 5);
 }
\ No newline at end of file

From 0450d2a497077fa4d839f8cfdfb457c0b8efca45 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Sat, 17 Aug 2024 16:17:53 +0200
Subject: [PATCH 25/67] format files

---
 unitTests/MatrixWrapper/Combination.C | 5 ++---
 unitTests/Repartitioner.C             | 2 +-
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/unitTests/MatrixWrapper/Combination.C b/unitTests/MatrixWrapper/Combination.C
index e849bf945..f80bc9702 100644
--- a/unitTests/MatrixWrapper/Combination.C
+++ b/unitTests/MatrixWrapper/Combination.C
@@ -176,8 +176,7 @@ TEST(Combination, CanConvertToCoo)
     auto cmb = CombinationMatrix<InputFormat>::create(
         exec, dim, linops);
 
-    auto out = OutputFormat::create(
-        exec);
+    auto out = OutputFormat::create(exec);
 
     // Act
     cmb->convert_to(out.get());
@@ -385,4 +384,4 @@ TEST(Combination, CanMoveToCsr)
     ASSERT_EQ(out->get_row_ptrs()[2], 2);
     ASSERT_EQ(out->get_row_ptrs()[3], 3);
     ASSERT_EQ(out->get_row_ptrs()[4], 5);
-}
\ No newline at end of file
+}
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 97b339dd7..f1ec63f29 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -67,6 +67,7 @@ TEST(Repartitioner, can_create_repartitioner)
 
     // Assert
     EXPECT_EQ(repartitioner.get_ranks_per_gpu(), ranks_per_gpu);
+<<<<<<< HEAD
 }
 
 TEST(Repartitioner, has_correct_properties_for_1_rank)
@@ -260,7 +261,6 @@ TEST(Repartitioner, has_correct_properties_for_2_ranks)
     EXPECT_EQ(repartitioner.is_owner(*exec), (rank % 2 == 0) ? true : false);
 }
 
-
 int main(int argc, char *argv[])
 {
     int result = 0;

From 9e4fc0906118044a72b3318285cff6d99b50fec0 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <go@hpsim.de>
Date: Tue, 20 Aug 2024 12:07:57 +0200
Subject: [PATCH 26/67] Apply suggestions from code review

Co-authored-by: chihta <97195768+chihta-wang@users.noreply.github.com>
---
 include/OGL/Repartitioner.H | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index 8cc0f1f97..11336e78f 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -32,7 +32,7 @@ std::vector<label> sort_permutation(const std::vector<T> &vec, Compare compare)
     return p;
 }
 
-/** @class Collects functions to repartion communication patterns
+/** @class Collects functions to repartition communication patterns
  * and matrices. Here repartitioning refers to changing a given partitioning on
  * n ranks to m ranks with n>=m
  */
@@ -94,7 +94,7 @@ public:
 
     label get_ranks_per_gpu() const { return ranks_per_gpu_; }
 
-    /* @brief computes the size after repartitioning
+    /* @brief computes the size of the submatrix owned by a given rank after repartitioning
      *
      * Given a local size (can represent nrows or nnzs) the new size of
      * this rank after repartitioning gets computed
@@ -130,7 +130,7 @@ public:
 
 
     /*
-    ** returns the the new sparsity patterns (local and non-local) and a vector
+    ** returns the new sparsity patterns (local and non-local) and a vector
     ** tracking the interfaces. This contains a pair where the first entry
     *(bool)
     ** signals whether this is a new local interface (no communication), and the
@@ -144,7 +144,7 @@ public:
         std::shared_ptr<SparsityPattern> src_local_pattern,
         std::shared_ptr<SparsityPattern> src_non_local_pattern) const;
 
-    /* @brief repartion a communication pattern
+    /* @brief repartition a communication pattern
     **
     ** returns the repartitioned communcation pattern
     */

From 9d3640089e36a3aaaab75116effafb714886297a Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Tue, 20 Aug 2024 14:16:47 +0200
Subject: [PATCH 27/67] start value parameterization

---
 unitTests/Repartitioner.C | 63 ++++++++++++++++++++++++++-------------
 1 file changed, 43 insertions(+), 20 deletions(-)

diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index f1ec63f29..1552025fc 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -58,10 +58,15 @@ public:
 const testing::Environment *global_env =
     AddGlobalTestEnvironment(new RepartitionerEnvironment);
 
-TEST(Repartitioner, can_create_repartitioner)
+class RepartitionerFixture : public testing::TestWithParam<int> {};
+
+INSTANTIATE_TEST_SUITE_P(RepartitionerFixtureInstantiation,
+                         RepartitionerFixture, testing::Values(1, 2, 4));
+
+TEST_P(RepartitionerFixture, can_create_repartitioner)
 {
     // Arrange
-    label ranks_per_gpu = 1;
+    label ranks_per_gpu = GetParam();
     auto exec = ((RepartitionerEnvironment *)global_env)->exec;
     auto repartitioner = Repartitioner(10, ranks_per_gpu, 0, *exec.get());
 
@@ -86,6 +91,40 @@ TEST(Repartitioner, has_correct_properties_for_1_rank)
     EXPECT_EQ(repartitioner.is_owner(*exec), true);
 }
 
+TEST(Repartitioner, has_correct_properties_for_4_ranks)
+{
+    // Arrange
+    label ranks_per_gpu = 4;
+    label local_size = 10;
+    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
+    auto rank = exec->get_rank();
+
+    // Assert
+    EXPECT_EQ(
+        repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
+        (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0));
+
+    EXPECT_EQ(repartitioner.is_owner(*exec), (rank == 0) ? true : false);
+}
+
+TEST(Repartitioner, has_correct_properties_for_2_ranks)
+{
+    // Arrange
+    label ranks_per_gpu = 2;
+    label local_size = 10;
+    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
+    auto rank = exec->get_rank();
+
+    // Assert
+    EXPECT_EQ(
+        repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
+        (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0));
+
+    EXPECT_EQ(repartitioner.is_owner(*exec), (rank % 2 == 0) ? true : false);
+}
+
 TEST(Repartitioner, can_repartition_comm_pattern_for_1_rank)
 {
     // Arrange
@@ -168,34 +207,18 @@ TEST(Repartitioner, can_repartition_comm_pattern_for_2_rank)
         repart_comm_pattern->target_ids.get_const_data() +
             repart_comm_pattern->target_ids.get_size());
     // on non owner no communication ranks should be left
-    std::vector<std::vector<label>> exp_ids {{2}, {}, {0}, {}};
+    std::vector<std::vector<label>> exp_ids{{2}, {}, {0}, {}};
     EXPECT_EQ(res_ids, exp_ids[rank]);
 
     std::vector<label> res_sizes(
         repart_comm_pattern->target_sizes.get_const_data(),
         repart_comm_pattern->target_sizes.get_const_data() +
             repart_comm_pattern->target_sizes.get_size());
-    std::vector<std::vector<label>> exp_sizes {{1}, {}, {1}, {}};
+    std::vector<std::vector<label>> exp_sizes{{1}, {}, {1}, {}};
     // thus communication sizes is also empty
     EXPECT_EQ(res_sizes, exp_sizes[rank]);
 }
 
-TEST(Repartitioner, has_correct_properties_for_4_ranks)
-{
-    // Arrange
-    label ranks_per_gpu = 4;
-    label local_size = 10;
-    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
-    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
-    auto rank = exec->get_rank();
-
-    // Assert
-    EXPECT_EQ(
-        repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
-        (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0));
-
-    EXPECT_EQ(repartitioner.is_owner(*exec), (rank == 0) ? true : false);
-}
 
 TEST(Repartitioner, can_repartition_comm_pattern_for_4_rank)
 {

From b55c138dd70a4aa6cb130e8150324828189d2c4e Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Tue, 20 Aug 2024 14:17:27 +0200
Subject: [PATCH 28/67] fix issue with narrowing conversion

---
 unitTests/MatrixWrapper/Combination.C | 97 ++++++++++++---------------
 1 file changed, 44 insertions(+), 53 deletions(-)

diff --git a/unitTests/MatrixWrapper/Combination.C b/unitTests/MatrixWrapper/Combination.C
index f80bc9702..24c2e25c8 100644
--- a/unitTests/MatrixWrapper/Combination.C
+++ b/unitTests/MatrixWrapper/Combination.C
@@ -64,15 +64,16 @@ TEST(Combination, CanCreateCombinationMatrixWithSparseRealMatrix)
     using IndexType = label;
     using mtx = gko::matrix::Csr<ValueType, IndexType>;
     using vec = gko::matrix::Dense<ValueType>;
-    
+
     int dim_size = 5;
-    auto dim = gko::dim<2>{dim_size, dim_size};
+    auto dim = gko::dim<2>{static_cast<gko::size_type>(dim_size),
+                           static_cast<gko::size_type>(dim_size)};
 
     auto exec = gko::share(gko::ReferenceExecutor::create());
 
-    auto m1linop = gko::share(
-        gko::read<mtx>(std::ifstream("data/A_sparse.mtx"), exec));
-    
+    auto m1linop =
+        gko::share(gko::read<mtx>(std::ifstream("data/A_sparse.mtx"), exec));
+
     auto b = gko::read<vec>(std::ifstream("data/b.mtx"), exec);
 
     std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m1linop};
@@ -82,12 +83,13 @@ TEST(Combination, CanCreateCombinationMatrixWithSparseRealMatrix)
         exec, dim, linops);
 
     // Assert
-    ASSERT_EQ(cmb->get_size(), gko::dim<2>(dim_size, dim_size));
+    ASSERT_EQ(cmb->get_size(), dim);
     ASSERT_EQ(cmb->get_coefficients().size(), 2);
     ASSERT_EQ(cmb->get_operators().size(), 2);
 
     // Act
-    auto x = gko::matrix::Dense<scalar>::create(exec, gko::dim<2>{dim_size, 1});
+    auto x = gko::matrix::Dense<scalar>::create(
+        exec, gko::dim<2>{static_cast<gko::size_type>(dim_size), 1});
     x->fill(0.0);
     cmb->apply(b, x);
 
@@ -107,36 +109,36 @@ TEST(Combination, CanCreateCombinationMatrixWithOpenFOAMLDUMatrix)
     using IndexType = label;
     using mtx = gko::matrix::Csr<ValueType, IndexType>;
     using vec = gko::matrix::Dense<ValueType>;
-    
+
     int dim_size = 5;
-    auto dim = gko::dim<2>{dim_size, dim_size};
+    auto dim = gko::dim<2>{static_cast<gko::size_type>(dim_size),
+                           static_cast<gko::size_type>(dim_size)};
 
     auto exec = gko::share(gko::ReferenceExecutor::create());
 
-    auto Llinop = gko::share(
-        gko::read<mtx>(std::ifstream("data/L.mtx"), exec));
+    auto Llinop = gko::share(gko::read<mtx>(std::ifstream("data/L.mtx"), exec));
 
-    auto Dlinop = gko::share(
-        gko::read<mtx>(std::ifstream("data/D.mtx"), exec));
-    
-    auto Ulinop = gko::share(
-        gko::read<mtx>(std::ifstream("data/U.mtx"), exec));
+    auto Dlinop = gko::share(gko::read<mtx>(std::ifstream("data/D.mtx"), exec));
+
+    auto Ulinop = gko::share(gko::read<mtx>(std::ifstream("data/U.mtx"), exec));
 
     auto b = gko::read<vec>(std::ifstream("data/b.mtx"), exec);
 
-    std::vector<std::shared_ptr<const gko::LinOp>> linops{Llinop, Dlinop, Ulinop};
+    std::vector<std::shared_ptr<const gko::LinOp>> linops{Llinop, Dlinop,
+                                                          Ulinop};
 
     // Act
     auto cmb = CombinationMatrix<gko::matrix::Csr<scalar, label>>::create(
         exec, dim, linops);
 
     // Assert
-    ASSERT_EQ(cmb->get_size(), gko::dim<2>(dim_size, dim_size));
+    ASSERT_EQ(cmb->get_size(), dim);
     ASSERT_EQ(cmb->get_coefficients().size(), 3);
     ASSERT_EQ(cmb->get_operators().size(), 3);
 
     // Act
-    auto x = gko::matrix::Dense<scalar>::create(exec, gko::dim<2>{dim_size, 1});
+    auto x = gko::matrix::Dense<scalar>::create(
+        exec, gko::dim<2>{static_cast<gko::size_type>(dim_size), 1});
     x->fill(0.0);
     cmb->apply(b, x);
 
@@ -158,23 +160,21 @@ TEST(Combination, CanConvertToCoo)
 
     int dim_size = 5;
 
-    auto dim = gko::dim<2>{dim_size, dim_size};
+    auto dim = gko::dim<2>{static_cast<gko::size_type>(dim_size),
+                           static_cast<gko::size_type>(dim_size)};
     gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, -3}, {2, 3, 5}});
     gko::matrix_data<double, int> m2(dim, {{3, 2, 4}, {3, 3, 1}, {4, 4, 2}});
 
     auto exec = gko::share(gko::ReferenceExecutor::create());
-    auto m1linop =
-        gko::share(InputFormat::create(exec, dim));
+    auto m1linop = gko::share(InputFormat::create(exec, dim));
     m1linop->read(m1);
 
-    auto m2linop =
-        gko::share(InputFormat::create(exec, dim));
+    auto m2linop = gko::share(InputFormat::create(exec, dim));
     m2linop->read(m2);
 
     std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m2linop};
 
-    auto cmb = CombinationMatrix<InputFormat>::create(
-        exec, dim, linops);
+    auto cmb = CombinationMatrix<InputFormat>::create(exec, dim, linops);
 
     auto out = OutputFormat::create(exec);
 
@@ -217,26 +217,23 @@ TEST(Combination, CanConvertToCsr)
 
     int dim_size = 5;
 
-    auto dim = gko::dim<2>{dim_size, dim_size};
+    auto dim = gko::dim<2>{static_cast<gko::size_type>(dim_size),
+                           static_cast<gko::size_type>(dim_size)};
     gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, -3}, {2, 3, 5}});
     gko::matrix_data<double, int> m2(dim, {{3, 2, 4}, {3, 3, 1}, {4, 4, 2}});
 
     auto exec = gko::share(gko::ReferenceExecutor::create());
-    auto m1linop =
-        gko::share(InputFormat::create(exec, dim));
+    auto m1linop = gko::share(InputFormat::create(exec, dim));
     m1linop->read(m1);
 
-    auto m2linop =
-        gko::share(InputFormat::create(exec, dim));
+    auto m2linop = gko::share(InputFormat::create(exec, dim));
     m2linop->read(m2);
 
     std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m2linop};
 
-    auto cmb = CombinationMatrix<InputFormat>::create(
-        exec, dim, linops);
+    auto cmb = CombinationMatrix<InputFormat>::create(exec, dim, linops);
 
-    auto out = OutputFormat::create(
-        exec);
+    auto out = OutputFormat::create(exec);
 
     // Act
     cmb->convert_to(out.get());
@@ -276,26 +273,23 @@ TEST(Combination, CanMoveToCoo)
 
     int dim_size = 5;
 
-    auto dim = gko::dim<2>{dim_size, dim_size};
+    auto dim = gko::dim<2>{static_cast<gko::size_type>(dim_size),
+                           static_cast<gko::size_type>(dim_size)};
     gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, -3}, {2, 3, 5}});
     gko::matrix_data<double, int> m2(dim, {{3, 2, 4}, {3, 3, 1}, {4, 4, 2}});
 
     auto exec = gko::share(gko::ReferenceExecutor::create());
-    auto m1linop =
-        gko::share(InputFormat::create(exec, dim));
+    auto m1linop = gko::share(InputFormat::create(exec, dim));
     m1linop->read(m1);
 
-    auto m2linop =
-        gko::share(InputFormat::create(exec, dim));
+    auto m2linop = gko::share(InputFormat::create(exec, dim));
     m2linop->read(m2);
 
     std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m2linop};
 
-    auto cmb = CombinationMatrix<InputFormat>::create(
-        exec, dim, linops);
+    auto cmb = CombinationMatrix<InputFormat>::create(exec, dim, linops);
 
-    auto out = OutputFormat::create(
-        exec);
+    auto out = OutputFormat::create(exec);
 
     // Act
     cmb->move_to(out.get());
@@ -337,26 +331,23 @@ TEST(Combination, CanMoveToCsr)
 
     int dim_size = 5;
 
-    auto dim = gko::dim<2>{dim_size, dim_size};
+    auto dim = gko::dim<2>{static_cast<gko::size_type>(dim_size),
+                           static_cast<gko::size_type>(dim_size)};
     gko::matrix_data<double, int> m1(dim, {{0, 0, 2}, {1, 1, -3}, {2, 3, 5}});
     gko::matrix_data<double, int> m2(dim, {{3, 2, 4}, {3, 3, 1}, {4, 4, 2}});
 
     auto exec = gko::share(gko::ReferenceExecutor::create());
-    auto m1linop =
-        gko::share(InputFormat::create(exec, dim));
+    auto m1linop = gko::share(InputFormat::create(exec, dim));
     m1linop->read(m1);
 
-    auto m2linop =
-        gko::share(InputFormat::create(exec, dim));
+    auto m2linop = gko::share(InputFormat::create(exec, dim));
     m2linop->read(m2);
 
     std::vector<std::shared_ptr<const gko::LinOp>> linops{m1linop, m2linop};
 
-    auto cmb = CombinationMatrix<InputFormat>::create(
-        exec, dim, linops);
+    auto cmb = CombinationMatrix<InputFormat>::create(exec, dim, linops);
 
-    auto out = OutputFormat::create(
-        exec);
+    auto out = OutputFormat::create(exec);
 
     // Act
     cmb->move_to(out.get());

From 34d1df2fa6b3b90cd98fa3b61041ff606971d496 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Tue, 20 Aug 2024 14:18:43 +0200
Subject: [PATCH 29/67] improve doc strings, const correctnes

---
 include/OGL/Repartitioner.H | 48 ++++++++++++++++++++++++++-----------
 src/Repartitioner.C         |  4 ++--
 2 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index 11336e78f..437a7c880 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -94,7 +94,8 @@ public:
 
     label get_ranks_per_gpu() const { return ranks_per_gpu_; }
 
-    /* @brief computes the size of the submatrix owned by a given rank after repartitioning
+    /* @brief computes the size of the submatrix owned by a given rank after
+     * repartitioning
      *
      * Given a local size (can represent nrows or nnzs) the new size of
      * this rank after repartitioning gets computed
@@ -117,7 +118,7 @@ public:
     };
 
 
-    /* @brief given received interfaces this function sorts into local and non
+    /* @brief given received interfaces this function sorts interfaces into local and non
      *local
      * returns: SparsityPattern storing new non_local sparsity pattern,
      *  and a pair of locality information, where the first bool stores if the
@@ -129,25 +130,44 @@ public:
         const SparsityPattern &non_loc) const;
 
 
-    /*
-    ** returns the new sparsity patterns (local and non-local) and a vector
-    ** tracking the interfaces. This contains a pair where the first entry
-    *(bool)
-    ** signals whether this is a new local interface (no communication), and the
-    ** second entry (label) tracks the original rank of the interface
-    */
+    /* @brief function to repartition a sparsity pattern
+     *
+     * function to repartition a sparsity pattern based on ranks_per_gpu
+     * ie changing a sparsity pattern which has rows on each rank to a sparsity
+     * pattern which has rows only on every ranks_per_gpu-th rank (owner)
+     *
+     * @param exec_handler the executor handler
+     * @param src_local_pattern the original sparsity pattern of the local
+     *matrix
+     * @param src_non_local_pattern the original sparsity pattern of the non
+     *local matrix
+     * returns the new sparsity patterns (local and non-local) and a vector
+     * tracking the interfaces. This contains a pair where the first entry
+     *(bool)
+     * signals whether this is a new local interface (no communication), and
+     *the
+     * second entry (label) tracks the original rank of the interface
+     */
     std::tuple<std::shared_ptr<SparsityPattern>,
                std::shared_ptr<SparsityPattern>,
                std::vector<std::pair<bool, label>>>
     repartition_sparsity(
         const ExecutorHandler &exec_handler,
-        std::shared_ptr<SparsityPattern> src_local_pattern,
-        std::shared_ptr<SparsityPattern> src_non_local_pattern) const;
+        std::shared_ptr<const SparsityPattern> src_local_pattern,
+        std::shared_ptr<const SparsityPattern> src_non_local_pattern) const;
 
     /* @brief repartition a communication pattern
-    **
-    ** returns the repartitioned communcation pattern
-    */
+     * function to repartition a communication pattern based on ranks_per_gpu
+     * ie changing a communication pattern which has rows on each rank and where
+     * every rank communicates to a communication pattern which has rows only on
+     * every ranks_per_gpu-th rank and thus only communicates between owner
+     * ranks
+     *
+     * @param src_comm_pattern the original communication pattern of the current
+     * rank belonging to the distributed matrix
+     * @param partition the original partition belonging to the distributed
+     * matrix returns the repartitioned communcation pattern
+     */
     std::shared_ptr<const CommunicationPattern> repartition_comm_pattern(
         const ExecutorHandler &exec_handler,
         std::shared_ptr<const CommunicationPattern> src_comm_pattern,
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index c215013cb..b05e0f5ae 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -22,8 +22,8 @@ std::tuple<std::shared_ptr<SparsityPattern>, std::shared_ptr<SparsityPattern>,
            std::vector<std::pair<bool, label>>>
 Repartitioner::repartition_sparsity(
     const ExecutorHandler &exec_handler,
-    std::shared_ptr<SparsityPattern> src_local_pattern,
-    std::shared_ptr<SparsityPattern> src_non_local_pattern) const
+    std::shared_ptr<const SparsityPattern> src_local_pattern,
+    std::shared_ptr<const SparsityPattern> src_non_local_pattern) const
 {
     LOG_1(verbose_, "start repartition sparsity pattern")
     // 1. obtain send recv sizes vector

From 9c6961d0205acc79079f1bae4c90f363e7ad9560 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Tue, 20 Aug 2024 14:24:12 +0200
Subject: [PATCH 30/67] parameterize has_correct_properties test

---
 unitTests/Repartitioner.C | 43 ++++++++-------------------------------
 1 file changed, 8 insertions(+), 35 deletions(-)

diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 1552025fc..c172a6a1c 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -68,6 +68,8 @@ TEST_P(RepartitionerFixture, can_create_repartitioner)
     // Arrange
     label ranks_per_gpu = GetParam();
     auto exec = ((RepartitionerEnvironment *)global_env)->exec;
+
+    // Act
     auto repartitioner = Repartitioner(10, ranks_per_gpu, 0, *exec.get());
 
     // Assert
@@ -75,56 +77,27 @@ TEST_P(RepartitionerFixture, can_create_repartitioner)
 <<<<<<< HEAD
 }
 
-TEST(Repartitioner, has_correct_properties_for_1_rank)
-{
-    // Arrange
-    label ranks_per_gpu = 1;
-    label local_size = 10;
-    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
-    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
-
-    // Assert
-    EXPECT_EQ(
-        repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
-        local_size);
-
-    EXPECT_EQ(repartitioner.is_owner(*exec), true);
-}
-
-TEST(Repartitioner, has_correct_properties_for_4_ranks)
+TEST_P(RepartitionerFixture, has_correct_properties_for_n_rank)
 {
     // Arrange
-    label ranks_per_gpu = 4;
+    label ranks_per_gpu = GetParam();
     label local_size = 10;
     auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
-    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
     auto rank = exec->get_rank();
 
-    // Assert
-    EXPECT_EQ(
-        repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
-        (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0));
-
-    EXPECT_EQ(repartitioner.is_owner(*exec), (rank == 0) ? true : false);
-}
-
-TEST(Repartitioner, has_correct_properties_for_2_ranks)
-{
-    // Arrange
-    label ranks_per_gpu = 2;
-    label local_size = 10;
-    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
+    // Act
     auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
-    auto rank = exec->get_rank();
 
     // Assert
+    EXPECT_EQ(repartitioner.is_owner(*exec), (rank % ranks_per_gpu == 0) ? true : false);
+
     EXPECT_EQ(
         repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
         (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0));
 
-    EXPECT_EQ(repartitioner.is_owner(*exec), (rank % 2 == 0) ? true : false);
 }
 
+
 TEST(Repartitioner, can_repartition_comm_pattern_for_1_rank)
 {
     // Arrange

From b7cef8a8f7150fea8bf7bc62ec4f774393c3334f Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Tue, 20 Aug 2024 14:24:34 +0200
Subject: [PATCH 31/67] parameterize has_correct_properties test

---
 unitTests/Repartitioner.C | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index c172a6a1c..7e1af5e41 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -89,12 +89,12 @@ TEST_P(RepartitionerFixture, has_correct_properties_for_n_rank)
     auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
 
     // Assert
-    EXPECT_EQ(repartitioner.is_owner(*exec), (rank % ranks_per_gpu == 0) ? true : false);
+    EXPECT_EQ(repartitioner.is_owner(*exec),
+              (rank % ranks_per_gpu == 0) ? true : false);
 
     EXPECT_EQ(
         repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
         (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0));
-
 }
 
 

From 2ef0fcd2a7c5f54b5501572c37814b6895dc7d81 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Tue, 20 Aug 2024 15:10:41 +0200
Subject: [PATCH 32/67] removed solver compilation since it registers solver
 multiple times

---
 CMakeLists.txt | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 7d055ad03..b585d04ce 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -224,9 +224,10 @@ target_sources(
           src/MatrixWrapper/HostMatrix.C
           src/MatrixWrapper/SparsityPattern.C
           src/Repartitioner.C
-          src/Solver/CG/GKOCG.C
-          src/Solver/BiCGStab/GKOBiCGStab.C
-          src/Solver/GMRES/GKOGMRES.C)
+	  # src/Solver/CG/GKOCG.C
+	  # src/Solver/BiCGStab/GKOBiCGStab.C
+	  # src/Solver/GMRES/GKOGMRES.C
+	  )
 
 enable_sanitizers(
   OGL ${OGL_ENABLE_SANITIZE_ADDRESS} ${OGL_ENABLE_SANITIZE_LEAK}

From 42a8c54eed03bd19c1e8dcc36b5818b3b584bcfd Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Tue, 20 Aug 2024 20:31:26 +0200
Subject: [PATCH 33/67] implement can_repartition_1d_comm_pattern n ranks

---
 unitTests/Repartitioner.C | 133 ++++++++------------------------------
 1 file changed, 28 insertions(+), 105 deletions(-)

diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 7e1af5e41..f05ea5e55 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -58,7 +58,11 @@ public:
 const testing::Environment *global_env =
     AddGlobalTestEnvironment(new RepartitionerEnvironment);
 
-class RepartitionerFixture : public testing::TestWithParam<int> {};
+class RepartitionerFixture : public testing::TestWithParam<int> {
+        public:
+            label local_size = 10;
+
+};
 
 INSTANTIATE_TEST_SUITE_P(RepartitionerFixtureInstantiation,
                          RepartitionerFixture, testing::Values(1, 2, 4));
@@ -77,11 +81,11 @@ TEST_P(RepartitionerFixture, can_create_repartitioner)
 <<<<<<< HEAD
 }
 
+
 TEST_P(RepartitionerFixture, has_correct_properties_for_n_rank)
 {
     // Arrange
     label ranks_per_gpu = GetParam();
-    label local_size = 10;
     auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
     auto rank = exec->get_rank();
 
@@ -98,11 +102,10 @@ TEST_P(RepartitionerFixture, has_correct_properties_for_n_rank)
 }
 
 
-TEST(Repartitioner, can_repartition_comm_pattern_for_1_rank)
+TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
 {
     // Arrange
-    label ranks_per_gpu = 1;
-    label local_size = 10;
+    label ranks_per_gpu = GetParam();
     auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
     auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
     auto rank = exec->get_rank();
@@ -112,55 +115,26 @@ TEST(Repartitioner, can_repartition_comm_pattern_for_1_rank)
     // cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
     // cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- global row ids
     std::vector<std::vector<label>> ids{{1}, {0, 2}, {1, 3}, {2}};
-    std::vector<std::vector<std::vector<label>>> rows{
-        {{1}}, {{0}, {1}}, {{0}, {1}}, {{0}}};
-
-    auto comm_pattern =
-        std::make_shared<CommunicationPattern>(*exec, ids[rank], rows[rank]);
 
-    auto comm = exec->get_communicator();
-    auto partition = gko::share(
-        gko::experimental::distributed::build_partition_from_local_size<label,
-                                                                        label>(
-            ref_exec, *comm.get(), 2));
-
-    // Act
-    auto repart_comm_pattern =
-        repartitioner.repartition_comm_pattern(*exec, comm_pattern, partition);
+    // expected communcation ranks
+    using vec_vec =std::vector<std::vector<label>>;
+    std::map<label, vec_vec> exp_res_ids {};
 
-    // Assert
-    // communication ranks are the same as before
-    std::vector<label> res_ids(
-        repart_comm_pattern->target_ids.get_const_data(),
-        repart_comm_pattern->target_ids.get_const_data() +
-            repart_comm_pattern->target_ids.get_size());
-    EXPECT_EQ(res_ids, ids[rank]);
-
-    std::vector<label> res_sizes(
-        repart_comm_pattern->target_sizes.get_const_data(),
-        repart_comm_pattern->target_sizes.get_const_data() +
-            repart_comm_pattern->target_sizes.get_size());
-    std::vector<std::vector<label>> exp_sizes{{1}, {1, 1}, {1, 1}, {1}};
-    EXPECT_EQ(res_sizes, exp_sizes[rank]);
-}
+    exp_res_ids[1] = ids; // in the ranks_per_gpu==1 case nothing changes
+    // only communication partners are 0-2 and 2-0
+    exp_res_ids.emplace(2, vec_vec{{2}, {}, {0}, {}});
+    // no communication if all ranks are repartitioned to single owner
+    exp_res_ids.emplace(4, vec_vec{{}, {}, {}, {}});
 
-TEST(Repartitioner, can_repartition_comm_pattern_for_2_rank)
-{
-    // Arrange
-    label ranks_per_gpu = 2;
-    label local_size = 10;
-    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
-    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
-    auto rank = exec->get_rank();
-    auto ref_exec = exec->get_ref_exec();
+    // expected communication sizes
+    std::map<label, std::vector<std::vector<label>>> exp_res_sizes;
+    exp_res_sizes.emplace(1, vec_vec({{1}, {1, 1}, {1, 1}, {1}}));
+    exp_res_sizes.emplace(2, vec_vec({{1}, {}, {1}, {}}));
+    exp_res_sizes.emplace(4, vec_vec({{}, {}, {}, {}}));
 
-    // rank:     0     1     2     3
-    // cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
-    // cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- local row ids
-    std::vector<std::vector<label>> ids{{1}, {0, 2}, {1, 3}, {2}};
+    // expected rows
     std::vector<std::vector<std::vector<label>>> rows{
         {{1}}, {{0}, {1}}, {{0}, {1}}, {{0}}};
-
     auto comm_pattern =
         std::make_shared<CommunicationPattern>(*exec, ids[rank], rows[rank]);
 
@@ -179,65 +153,14 @@ TEST(Repartitioner, can_repartition_comm_pattern_for_2_rank)
         repart_comm_pattern->target_ids.get_const_data(),
         repart_comm_pattern->target_ids.get_const_data() +
             repart_comm_pattern->target_ids.get_size());
-    // on non owner no communication ranks should be left
-    std::vector<std::vector<label>> exp_ids{{2}, {}, {0}, {}};
-    EXPECT_EQ(res_ids, exp_ids[rank]);
-
-    std::vector<label> res_sizes(
-        repart_comm_pattern->target_sizes.get_const_data(),
-        repart_comm_pattern->target_sizes.get_const_data() +
-            repart_comm_pattern->target_sizes.get_size());
-    std::vector<std::vector<label>> exp_sizes{{1}, {}, {1}, {}};
-    // thus communication sizes is also empty
-    EXPECT_EQ(res_sizes, exp_sizes[rank]);
-}
-
-
-TEST(Repartitioner, can_repartition_comm_pattern_for_4_rank)
-{
-    // Arrange
-    label ranks_per_gpu = 4;
-    label local_size = 10;
-    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
-    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
-    auto rank = exec->get_rank();
-    auto ref_exec = exec->get_ref_exec();
-
-    // rank:     0     1     2     3
-    // cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
-    // cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- global row ids
-    std::vector<std::vector<label>> ids{{1}, {0, 2}, {1, 3}, {2}};
-    std::vector<std::vector<std::vector<label>>> rows{
-        {{1}}, {{0}, {1}}, {{0}, {1}}, {{0}}};
-
-    auto comm_pattern =
-        std::make_shared<CommunicationPattern>(*exec, ids[rank], rows[rank]);
-
-    auto comm = exec->get_communicator();
-    auto partition = gko::share(
-        gko::experimental::distributed::build_partition_from_local_size<label,
-                                                                        label>(
-            ref_exec, *comm.get(), 2));
 
-    // Act
-    auto repart_comm_pattern =
-        repartitioner.repartition_comm_pattern(*exec, comm_pattern, partition);
+    EXPECT_EQ(res_ids, exp_res_ids[ranks_per_gpu][rank]);
 
-    // Assert
-    std::vector<label> res_ids(
-        repart_comm_pattern->target_ids.get_const_data(),
-        repart_comm_pattern->target_ids.get_const_data() +
-            repart_comm_pattern->target_ids.get_size());
-    // no communication ranks should be left
-    EXPECT_EQ(res_ids.size(), 0);
-
-    std::vector<label> res_sizes(
-        repart_comm_pattern->target_sizes.get_const_data(),
-        repart_comm_pattern->target_sizes.get_const_data() +
-            repart_comm_pattern->target_sizes.get_size());
-    // thus communication sizes is also empty
-    EXPECT_EQ(res_sizes.size(), 0);
-    EXPECT_EQ(repart_comm_pattern->send_idxs.size(), 0);
+   std::vector<label> res_sizes(
+       repart_comm_pattern->target_sizes.get_const_data(),
+       repart_comm_pattern->target_sizes.get_const_data() +
+           repart_comm_pattern->target_sizes.get_size());
+   EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
 }
 
 TEST(Repartitioner, has_correct_properties_for_2_ranks)

From 6e5fd07466b4774e72532fa0d6d37211a2b87158 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 21 Aug 2024 08:47:59 +0200
Subject: [PATCH 34/67] add rows check for 1d cells array

---
 unitTests/Repartitioner.C | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index f05ea5e55..569b1aa3c 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -113,7 +113,7 @@ TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
 
     // rank:     0     1     2     3
     // cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
-    // cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- global row ids
+    // cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- local row ids
     std::vector<std::vector<label>> ids{{1}, {0, 2}, {1, 3}, {2}};
 
     // expected communcation ranks
@@ -127,14 +127,22 @@ TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
     exp_res_ids.emplace(4, vec_vec{{}, {}, {}, {}});
 
     // expected communication sizes
-    std::map<label, std::vector<std::vector<label>>> exp_res_sizes;
+    std::map<label, vec_vec> exp_res_sizes;
     exp_res_sizes.emplace(1, vec_vec({{1}, {1, 1}, {1, 1}, {1}}));
     exp_res_sizes.emplace(2, vec_vec({{1}, {}, {1}, {}}));
     exp_res_sizes.emplace(4, vec_vec({{}, {}, {}, {}}));
 
     // expected rows
-    std::vector<std::vector<std::vector<label>>> rows{
-        {{1}}, {{0}, {1}}, {{0}, {1}}, {{0}}};
+    using vec_vec_vec = std::vector<vec_vec>;
+    vec_vec_vec rows{{{1}}, {{0}, {1}}, {{0}, {1}}, {{0}}};
+    std::map<label, vec_vec> exp_res_rows;
+    exp_res_rows.emplace(1, vec_vec{{1}, {0, 1}, {0, 1}, {0}});
+    // after repartitioning with 2 ranks_per_gpu [ 0 1 2 3 | 0 1 2 3 ] <- local row ids
+    exp_res_rows.emplace(2, vec_vec{{3}, {}, {0}, {}});
+    // after repartitioning with 4 ranks_per_gpu [ 0 1 2 3 4 5 6 7 ] <- local row ids
+    exp_res_rows.emplace(4, vec_vec{{}, {}, {}, {}});
+
+    // the original comm_pattern
     auto comm_pattern =
         std::make_shared<CommunicationPattern>(*exec, ids[rank], rows[rank]);
 
@@ -161,6 +169,14 @@ TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
        repart_comm_pattern->target_sizes.get_const_data() +
            repart_comm_pattern->target_sizes.get_size());
    EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
+
+   auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
+    std::vector<label> res_rows(
+        total_rank_send_idx.get_const_data(),
+        total_rank_send_idx.get_const_data() +
+            total_rank_send_idx.get_size());
+
+    EXPECT_EQ(res_rows, exp_res_rows[ranks_per_gpu][rank]);
 }
 
 TEST(Repartitioner, has_correct_properties_for_2_ranks)

From 9435c5ecc36d3b9323481dadbf0d6e4a479c9741 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 21 Aug 2024 09:14:27 +0200
Subject: [PATCH 35/67] add unit test for 2d cells array

---
 src/Repartitioner.C       |  1 -
 unitTests/Repartitioner.C | 82 +++++++++++++++++++++++++++++++++++----
 2 files changed, 75 insertions(+), 8 deletions(-)

diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index b05e0f5ae..c002812c4 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -349,7 +349,6 @@ Repartitioner::repartition_comm_pattern(
     auto gathered_target_sizes = gather_labels_to_owner(
         exec_handler, comm_pattern, target_sizes.data(), target_sizes.size());
 
-
     // next the send_ixs need to be updated we send them piecewise since
     // the send_idxs are a vector of gko::arrays
     if (owner) {
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 569b1aa3c..d915cbf2b 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -179,21 +179,89 @@ TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
     EXPECT_EQ(res_rows, exp_res_rows[ranks_per_gpu][rank]);
 }
 
-TEST(Repartitioner, has_correct_properties_for_2_ranks)
+TEST_P(RepartitionerFixture, can_repartition_2D_comm_pattern_for_n_ranks)
 {
     // Arrange
-    label ranks_per_gpu = 2;
-    label local_size = 10;
+    label ranks_per_gpu = GetParam();
     auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
     auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
     auto rank = exec->get_rank();
+    auto ref_exec = exec->get_ref_exec();
+
+    // cells: [ 2 3 | 2 3 ]
+    //   2    [ 0 1 | 0 1 ]  3
+    //        -------------
+    //        [ 2 3 | 2 3 ]
+    //   0    [ 0 1 | 0 1 ]  1
+    //
+    std::vector<std::vector<label>> ids{{1, 2}, {0, 3}, {0, 3}, {1, 2}};
+
+    // expected communcation ranks
+    using vec_vec =std::vector<std::vector<label>>;
+    std::map<label, vec_vec> exp_res_ids {};
+
+    exp_res_ids[1] = ids; // in the ranks_per_gpu==1 case nothing changes
+    // only communication partners are 0-2 and 2-0
+    exp_res_ids.emplace(2, vec_vec{{2}, {}, {0}, {}});
+    // no communication if all ranks are repartitioned to single owner
+    exp_res_ids.emplace(4, vec_vec{{}, {}, {}, {}});
+
+    // expected communication sizes
+    std::map<label, vec_vec> exp_res_sizes;
+    exp_res_sizes.emplace(1, vec_vec({{2, 2}, {2, 2}, {2, 2}, {2, 2}}));
+    exp_res_sizes.emplace(2, vec_vec({{4}, {}, {4}, {}}));
+    exp_res_sizes.emplace(4, vec_vec({{}, {}, {}, {}}));
+
+    // expected rows
+    using vec_vec_vec = std::vector<vec_vec>;
+    vec_vec_vec rows{{{1, 3}, {2, 3}}, {{0, 2}, {2, 3}}, {{0, 1}, {1, 3}}, {{0, 1}, {0, 2}}};
+    std::map<label, vec_vec> exp_res_rows;
+    exp_res_rows.emplace(1, vec_vec{{1, 3, 2, 3}, {0, 2, 2, 3}, {0, 1, 1, 3}, {0, 1, 0 ,2 }});
+    // after repartitioning with 2 ranks_per_gpu
+    // cells: [ 2 3  6 7 ]
+    //   2    [ 0 1  4 5 ]  3
+    //        -------------
+    //        [ 2 3  6 7 ]
+    //   0    [ 0 1  4 5 ]  1
+    exp_res_rows.emplace(2, vec_vec{{2, 3, 6, 7}, {}, {0, 1, 4, 5}, {}});
+    // after repartitioning with 4 ranks_per_gpu [ 0 1 2 3 4 5 6 7 ] <- local row ids
+    exp_res_rows.emplace(4, vec_vec{{}, {}, {}, {}});
+
+    // the original comm_pattern
+    auto comm_pattern =
+        std::make_shared<CommunicationPattern>(*exec, ids[rank], rows[rank]);
+
+    auto comm = exec->get_communicator();
+    auto partition = gko::share(
+        gko::experimental::distributed::build_partition_from_local_size<label,
+                                                                        label>(
+            ref_exec, *comm.get(), 4));
+
+    // Act
+    auto repart_comm_pattern =
+        repartitioner.repartition_comm_pattern(*exec, comm_pattern, partition);
 
     // Assert
-    EXPECT_EQ(
-        repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
-        (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0));
+    std::vector<label> res_ids(
+        repart_comm_pattern->target_ids.get_const_data(),
+        repart_comm_pattern->target_ids.get_const_data() +
+            repart_comm_pattern->target_ids.get_size());
 
-    EXPECT_EQ(repartitioner.is_owner(*exec), (rank % 2 == 0) ? true : false);
+    EXPECT_EQ(res_ids, exp_res_ids[ranks_per_gpu][rank]);
+
+   std::vector<label> res_sizes(
+       repart_comm_pattern->target_sizes.get_const_data(),
+       repart_comm_pattern->target_sizes.get_const_data() +
+           repart_comm_pattern->target_sizes.get_size());
+   EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
+
+   auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
+    std::vector<label> res_rows(
+        total_rank_send_idx.get_const_data(),
+        total_rank_send_idx.get_const_data() +
+            total_rank_send_idx.get_size());
+
+    EXPECT_EQ(res_rows, exp_res_rows[ranks_per_gpu][rank]);
 }
 
 int main(int argc, char *argv[])

From 5fcd46aa717d661c58139ec0b8fb55d4b6226add Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 21 Aug 2024 09:15:09 +0200
Subject: [PATCH 36/67] add unit test for 2d cells array

---
 unitTests/Repartitioner.C | 63 ++++++++++++++++++++-------------------
 1 file changed, 32 insertions(+), 31 deletions(-)

diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index d915cbf2b..299e58f2d 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -59,9 +59,8 @@ const testing::Environment *global_env =
     AddGlobalTestEnvironment(new RepartitionerEnvironment);
 
 class RepartitionerFixture : public testing::TestWithParam<int> {
-        public:
-            label local_size = 10;
-
+public:
+    label local_size = 10;
 };
 
 INSTANTIATE_TEST_SUITE_P(RepartitionerFixtureInstantiation,
@@ -117,10 +116,10 @@ TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
     std::vector<std::vector<label>> ids{{1}, {0, 2}, {1, 3}, {2}};
 
     // expected communcation ranks
-    using vec_vec =std::vector<std::vector<label>>;
-    std::map<label, vec_vec> exp_res_ids {};
+    using vec_vec = std::vector<std::vector<label>>;
+    std::map<label, vec_vec> exp_res_ids{};
 
-    exp_res_ids[1] = ids; // in the ranks_per_gpu==1 case nothing changes
+    exp_res_ids[1] = ids;  // in the ranks_per_gpu==1 case nothing changes
     // only communication partners are 0-2 and 2-0
     exp_res_ids.emplace(2, vec_vec{{2}, {}, {0}, {}});
     // no communication if all ranks are repartitioned to single owner
@@ -137,9 +136,11 @@ TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
     vec_vec_vec rows{{{1}}, {{0}, {1}}, {{0}, {1}}, {{0}}};
     std::map<label, vec_vec> exp_res_rows;
     exp_res_rows.emplace(1, vec_vec{{1}, {0, 1}, {0, 1}, {0}});
-    // after repartitioning with 2 ranks_per_gpu [ 0 1 2 3 | 0 1 2 3 ] <- local row ids
+    // after repartitioning with 2 ranks_per_gpu [ 0 1 2 3 | 0 1 2 3 ] <- local
+    // row ids
     exp_res_rows.emplace(2, vec_vec{{3}, {}, {0}, {}});
-    // after repartitioning with 4 ranks_per_gpu [ 0 1 2 3 4 5 6 7 ] <- local row ids
+    // after repartitioning with 4 ranks_per_gpu [ 0 1 2 3 4 5 6 7 ] <- local
+    // row ids
     exp_res_rows.emplace(4, vec_vec{{}, {}, {}, {}});
 
     // the original comm_pattern
@@ -164,17 +165,16 @@ TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
 
     EXPECT_EQ(res_ids, exp_res_ids[ranks_per_gpu][rank]);
 
-   std::vector<label> res_sizes(
-       repart_comm_pattern->target_sizes.get_const_data(),
-       repart_comm_pattern->target_sizes.get_const_data() +
-           repart_comm_pattern->target_sizes.get_size());
-   EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
+    std::vector<label> res_sizes(
+        repart_comm_pattern->target_sizes.get_const_data(),
+        repart_comm_pattern->target_sizes.get_const_data() +
+            repart_comm_pattern->target_sizes.get_size());
+    EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
 
-   auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
+    auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
     std::vector<label> res_rows(
         total_rank_send_idx.get_const_data(),
-        total_rank_send_idx.get_const_data() +
-            total_rank_send_idx.get_size());
+        total_rank_send_idx.get_const_data() + total_rank_send_idx.get_size());
 
     EXPECT_EQ(res_rows, exp_res_rows[ranks_per_gpu][rank]);
 }
@@ -193,14 +193,13 @@ TEST_P(RepartitionerFixture, can_repartition_2D_comm_pattern_for_n_ranks)
     //        -------------
     //        [ 2 3 | 2 3 ]
     //   0    [ 0 1 | 0 1 ]  1
-    //
     std::vector<std::vector<label>> ids{{1, 2}, {0, 3}, {0, 3}, {1, 2}};
 
     // expected communcation ranks
-    using vec_vec =std::vector<std::vector<label>>;
-    std::map<label, vec_vec> exp_res_ids {};
+    using vec_vec = std::vector<std::vector<label>>;
+    std::map<label, vec_vec> exp_res_ids{};
 
-    exp_res_ids[1] = ids; // in the ranks_per_gpu==1 case nothing changes
+    exp_res_ids[1] = ids;  // in the ranks_per_gpu==1 case nothing changes
     // only communication partners are 0-2 and 2-0
     exp_res_ids.emplace(2, vec_vec{{2}, {}, {0}, {}});
     // no communication if all ranks are repartitioned to single owner
@@ -214,9 +213,11 @@ TEST_P(RepartitionerFixture, can_repartition_2D_comm_pattern_for_n_ranks)
 
     // expected rows
     using vec_vec_vec = std::vector<vec_vec>;
-    vec_vec_vec rows{{{1, 3}, {2, 3}}, {{0, 2}, {2, 3}}, {{0, 1}, {1, 3}}, {{0, 1}, {0, 2}}};
+    vec_vec_vec rows{
+        {{1, 3}, {2, 3}}, {{0, 2}, {2, 3}}, {{0, 1}, {1, 3}}, {{0, 1}, {0, 2}}};
     std::map<label, vec_vec> exp_res_rows;
-    exp_res_rows.emplace(1, vec_vec{{1, 3, 2, 3}, {0, 2, 2, 3}, {0, 1, 1, 3}, {0, 1, 0 ,2 }});
+    exp_res_rows.emplace(
+        1, vec_vec{{1, 3, 2, 3}, {0, 2, 2, 3}, {0, 1, 1, 3}, {0, 1, 0, 2}});
     // after repartitioning with 2 ranks_per_gpu
     // cells: [ 2 3  6 7 ]
     //   2    [ 0 1  4 5 ]  3
@@ -224,7 +225,8 @@ TEST_P(RepartitionerFixture, can_repartition_2D_comm_pattern_for_n_ranks)
     //        [ 2 3  6 7 ]
     //   0    [ 0 1  4 5 ]  1
     exp_res_rows.emplace(2, vec_vec{{2, 3, 6, 7}, {}, {0, 1, 4, 5}, {}});
-    // after repartitioning with 4 ranks_per_gpu [ 0 1 2 3 4 5 6 7 ] <- local row ids
+    // after repartitioning with 4 ranks_per_gpu [ 0 1 2 3 4 5 6 7 ] <- local
+    // row ids
     exp_res_rows.emplace(4, vec_vec{{}, {}, {}, {}});
 
     // the original comm_pattern
@@ -249,17 +251,16 @@ TEST_P(RepartitionerFixture, can_repartition_2D_comm_pattern_for_n_ranks)
 
     EXPECT_EQ(res_ids, exp_res_ids[ranks_per_gpu][rank]);
 
-   std::vector<label> res_sizes(
-       repart_comm_pattern->target_sizes.get_const_data(),
-       repart_comm_pattern->target_sizes.get_const_data() +
-           repart_comm_pattern->target_sizes.get_size());
-   EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
+    std::vector<label> res_sizes(
+        repart_comm_pattern->target_sizes.get_const_data(),
+        repart_comm_pattern->target_sizes.get_const_data() +
+            repart_comm_pattern->target_sizes.get_size());
+    EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
 
-   auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
+    auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
     std::vector<label> res_rows(
         total_rank_send_idx.get_const_data(),
-        total_rank_send_idx.get_const_data() +
-            total_rank_send_idx.get_size());
+        total_rank_send_idx.get_const_data() + total_rank_send_idx.get_size());
 
     EXPECT_EQ(res_rows, exp_res_rows[ranks_per_gpu][rank]);
 }

From f40969c6cd472b0e547753a31972dfe6c48d8028 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 21 Aug 2024 09:19:13 +0200
Subject: [PATCH 37/67] format files

---
 include/OGL/CommunicationPattern.H          | 11 ++++++-----
 include/OGL/MatrixWrapper/Combination.H     |  8 ++++----
 include/OGL/MatrixWrapper/SparsityPattern.H |  2 +-
 include/OGL/Repartitioner.H                 | 11 +++++------
 include/OGL/common.H                        |  4 +---
 5 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/include/OGL/CommunicationPattern.H b/include/OGL/CommunicationPattern.H
index 70561d956..7fa83c2a0 100644
--- a/include/OGL/CommunicationPattern.H
+++ b/include/OGL/CommunicationPattern.H
@@ -129,18 +129,19 @@ struct CommunicationPattern {
         // ASSERT_EQ(target_ids.get_size(), send_idxs.size());
     }
 
-    CommunicationPattern(
-        const ExecutorHandler &exec, const std::vector<label>& ranks,
-        const std::vector<std::vector<label>>& rows)
+    CommunicationPattern(const ExecutorHandler &exec,
+                         const std::vector<label> &ranks,
+                         const std::vector<std::vector<label>> &rows)
         : exec_handler(exec),
           target_ids(exec.get_ref_exec(), ranks.begin(), ranks.end()),
           target_sizes(exec.get_ref_exec(), ranks.size())
     {
         ASSERT_EQ(ranks.size(), rows.size());
 
-        for (int i=0; i<rows.size();i++){
+        for (int i = 0; i < rows.size(); i++) {
             target_sizes.get_data()[i] = rows[i].size();
-            gko::array<label> row_array (exec.get_ref_exec(),rows[i].begin(), rows[i].end());
+            gko::array<label> row_array(exec.get_ref_exec(), rows[i].begin(),
+                                        rows[i].end());
             send_idxs.push_back({row_array, ranks[i]});
         }
     }
diff --git a/include/OGL/MatrixWrapper/Combination.H b/include/OGL/MatrixWrapper/Combination.H
index 13b86a8cb..82d67cac9 100644
--- a/include/OGL/MatrixWrapper/Combination.H
+++ b/include/OGL/MatrixWrapper/Combination.H
@@ -66,7 +66,7 @@ public:
         return *this;
     }
 
-     /**
+    /**
      * Returns a list of coefficients of the combination.
      *
      * @return a list of coefficients
@@ -113,7 +113,7 @@ public:
     /*
      * @warning assumes that all entries are unique
      */
-    void convert_to(gko::matrix::Csr<scalar, label> *result) const override 
+    void convert_to(gko::matrix::Csr<scalar, label> *result) const override
     {
         auto exec = this->get_executor();
         auto operators = this->get_operators();
@@ -135,7 +135,7 @@ public:
     /*
      * @warning assumes that all entries are unique
      */
-    void move_to(gko::matrix::Coo<scalar, label> *result) override 
+    void move_to(gko::matrix::Coo<scalar, label> *result) override
     {
         this->convert_to(result);
         this->comb_ = nullptr;
@@ -145,7 +145,7 @@ public:
     /*
      * @warning assumes that all entries are unique
      */
-    void move_to(gko::matrix::Csr<scalar, label> *result) override 
+    void move_to(gko::matrix::Csr<scalar, label> *result) override
     {
         this->convert_to(result);
         this->comb_ = nullptr;
diff --git a/include/OGL/MatrixWrapper/SparsityPattern.H b/include/OGL/MatrixWrapper/SparsityPattern.H
index 46390b26d..a40ace6f3 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.H
+++ b/include/OGL/MatrixWrapper/SparsityPattern.H
@@ -6,8 +6,8 @@
 
 #include <vector>
 
-#include "OGL/common.H"
 #include "OGL/CommunicationPattern.H"
+#include "OGL/common.H"
 
 namespace Foam {
 
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index 437a7c880..0873825c9 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -118,12 +118,11 @@ public:
     };
 
 
-    /* @brief given received interfaces this function sorts interfaces into local and non
-     *local
-     * returns: SparsityPattern storing new non_local sparsity pattern,
-     *  and a pair of locality information, where the first bool stores if the
-     *  interface performs a local communication and the second bool if the
-     *  the interface was originally from this rank
+    /* @brief given received interfaces this function sorts interfaces into
+     *local and non local returns: SparsityPattern storing new non_local
+     *sparsity pattern, and a pair of locality information, where the first bool
+     *stores if the interface performs a local communication and the second bool
+     *if the the interface was originally from this rank
      */
     std::pair<SparsityPattern, std::vector<bool>> build_non_local_interfaces(
         const ExecutorHandler &exec_handler, SparsityPattern &loc,
diff --git a/include/OGL/common.H b/include/OGL/common.H
index 4fb8eafda..08a31dbe1 100644
--- a/include/OGL/common.H
+++ b/include/OGL/common.H
@@ -100,9 +100,7 @@ std::ostream &operator<<(
     std::ostream &os,
     const std::shared_ptr<gko::matrix::Dense<scalar>> array_in);
 
-std::ostream &operator<<(
-    std::ostream &os,
-    const std::vector<label>& in);
+std::ostream &operator<<(std::ostream &os, const std::vector<label> &in);
 
 
 void export_system(const word fieldName, const gko::matrix::Csr<scalar> *A,

From da511805ca8791e6abf287255326a56105f1e642 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 21 Aug 2024 09:23:34 +0200
Subject: [PATCH 38/67] format files, update reuse file

---
 .reuse/dep5  |  4 ++++
 src/common.C | 10 +++++-----
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/.reuse/dep5 b/.reuse/dep5
index 5d8501672..932aa9c2f 100644
--- a/.reuse/dep5
+++ b/.reuse/dep5
@@ -43,6 +43,10 @@ Files: unitTests/case/*
 Copyright: 2024 OGL authors
 License: GPL-3.0-or-later
 
+Files: unitTests/MatrixWrapper/data/*mtx
+Copyright: 2024 OGL authors
+License: GPL-3.0-or-later
+
 # Sample paragraph, commented out:
 #
 # Files: src/*
diff --git a/src/common.C b/src/common.C
index 23ba8894a..d8f3ef85f 100644
--- a/src/common.C
+++ b/src/common.C
@@ -170,11 +170,12 @@ std::ostream &operator<<(std::ostream &os,
     return os;
 }
 
-std::ostream &operator<<(
-    std::ostream &os,
-    const std::vector<label>& in){
+std::ostream &operator<<(std::ostream &os, const std::vector<label> &in)
+{
     label size = in.size();
-    if (size == 0 ){ return os;}
+    if (size == 0) {
+        return os;
+    }
     os << size << " elements [";
     if (size > 100) {
         for (label i = 0; i < 9; i++) {
@@ -192,6 +193,5 @@ std::ostream &operator<<(
         os << "(" << size - 1 << ", " << in[size - 1] << ")]\n";
     }
     return os;
-
 };
 }  // namespace Foam

From cc567e94ce3a7900c2a1c5ef2f961f537855991a Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 21 Aug 2024 09:23:59 +0200
Subject: [PATCH 39/67] format file

---
 include/OGL/Repartitioner.H | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index 0873825c9..a43692300 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -62,7 +62,7 @@ public:
               gko::experimental::distributed::build_partition_from_local_size<
                   label, label>(exec_handler.get_ref_exec(),
                                 *exec_handler.get_communicator().get(),
-                                size))) {};
+                                size))){};
 
     /* returns the owner rank for a given rank */
     label get_owner_rank(label rank) const

From 8b4d4060ab0e00702f53e9fdf16c9c05c991db17 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 21 Aug 2024 09:24:36 +0200
Subject: [PATCH 40/67] format cmake files

---
 CMakeLists.txt                         | 7 +++----
 cmake/ginkgo.cmake                     | 3 +--
 cmake/version.cmake                    | 6 ++----
 unitTests/MatrixWrapper/CMakeLists.txt | 3 ++-
 4 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index b585d04ce..0d7dfd81a 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -224,10 +224,9 @@ target_sources(
           src/MatrixWrapper/HostMatrix.C
           src/MatrixWrapper/SparsityPattern.C
           src/Repartitioner.C
-	  # src/Solver/CG/GKOCG.C
-	  # src/Solver/BiCGStab/GKOBiCGStab.C
-	  # src/Solver/GMRES/GKOGMRES.C
-	  )
+  # src/Solver/CG/GKOCG.C src/Solver/BiCGStab/GKOBiCGStab.C
+  # src/Solver/GMRES/GKOGMRES.C
+)
 
 enable_sanitizers(
   OGL ${OGL_ENABLE_SANITIZE_ADDRESS} ${OGL_ENABLE_SANITIZE_LEAK}
diff --git a/cmake/ginkgo.cmake b/cmake/ginkgo.cmake
index 2a52fc81a..47923c4f9 100644
--- a/cmake/ginkgo.cmake
+++ b/cmake/ginkgo.cmake
@@ -10,8 +10,7 @@ if(NOT ${OGL_USE_EXTERNAL_GINKGO})
     QUITE
     GIT_SHALLOW ON
     GIT_REPOSITORY "https://github.com/ginkgo-project/ginkgo.git"
-    GIT_TAG ${GINKGO_CHECKOUT_VERSION}
-    )
+    GIT_TAG ${GINKGO_CHECKOUT_VERSION})
 
   FetchContent_MakeAvailable(Ginkgo)
 endif()
diff --git a/cmake/version.cmake b/cmake/version.cmake
index 84e61a7f0..e8eb6ae50 100644
--- a/cmake/version.cmake
+++ b/cmake/version.cmake
@@ -19,10 +19,8 @@ if("${GIT_REV}" STREQUAL "")
 else()
   execute_process(COMMAND bash -c "git diff --quiet --exit-code || echo +"
                   OUTPUT_VARIABLE GIT_DIFF)
-  execute_process(
-    COMMAND git describe --exact-match --tags
-    OUTPUT_VARIABLE GIT_TAG
-    ERROR_QUIET)
+  execute_process(COMMAND git describe --exact-match --tags
+                  OUTPUT_VARIABLE GIT_TAG ERROR_QUIET)
   execute_process(COMMAND git rev-parse --abbrev-ref HEAD
                   OUTPUT_VARIABLE GIT_BRANCH)
 
diff --git a/unitTests/MatrixWrapper/CMakeLists.txt b/unitTests/MatrixWrapper/CMakeLists.txt
index c288f036e..90a548f78 100644
--- a/unitTests/MatrixWrapper/CMakeLists.txt
+++ b/unitTests/MatrixWrapper/CMakeLists.txt
@@ -13,4 +13,5 @@ ogl_unit_test(matrixConversion "HostMatrix.C")
 ogl_unit_test(combination "Combination.C")
 
 # Set working directory for the unit test suite of Combination Matrix
-set_tests_properties(combination PROPERTIES WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})
+set_tests_properties(combination PROPERTIES WORKING_DIRECTORY
+                                            ${CMAKE_CURRENT_BINARY_DIR})

From df9f4f6c40441452f70694987caf41cdb27849ab Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 21 Aug 2024 09:25:03 +0200
Subject: [PATCH 41/67] format files .mtx

---
 unitTests/MatrixWrapper/data/D.mtx | 2 +-
 unitTests/MatrixWrapper/data/L.mtx | 2 +-
 unitTests/MatrixWrapper/data/U.mtx | 2 +-
 unitTests/MatrixWrapper/data/b.mtx | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/unitTests/MatrixWrapper/data/D.mtx b/unitTests/MatrixWrapper/data/D.mtx
index 784d16dec..9c2289ca1 100644
--- a/unitTests/MatrixWrapper/data/D.mtx
+++ b/unitTests/MatrixWrapper/data/D.mtx
@@ -4,4 +4,4 @@
 2 2 10.5
 3 3 0.5
 4 4 -280
-5 5 12
\ No newline at end of file
+5 5 12
diff --git a/unitTests/MatrixWrapper/data/L.mtx b/unitTests/MatrixWrapper/data/L.mtx
index 738df39b9..ca71205ad 100644
--- a/unitTests/MatrixWrapper/data/L.mtx
+++ b/unitTests/MatrixWrapper/data/L.mtx
@@ -1,3 +1,3 @@
 %%MatrixMarket matrix coordinate real general
 5 5 1
-4 2 250.5
\ No newline at end of file
+4 2 250.5
diff --git a/unitTests/MatrixWrapper/data/U.mtx b/unitTests/MatrixWrapper/data/U.mtx
index c0b5f6167..33df26e15 100644
--- a/unitTests/MatrixWrapper/data/U.mtx
+++ b/unitTests/MatrixWrapper/data/U.mtx
@@ -1,4 +1,4 @@
 %%MatrixMarket matrix coordinate real general
 5 5 2
 1 4 6
-4 5 33.2
\ No newline at end of file
+4 5 33.2
diff --git a/unitTests/MatrixWrapper/data/b.mtx b/unitTests/MatrixWrapper/data/b.mtx
index bed718eaf..3ab5be355 100644
--- a/unitTests/MatrixWrapper/data/b.mtx
+++ b/unitTests/MatrixWrapper/data/b.mtx
@@ -4,4 +4,4 @@
 0
 2
 3
-4
\ No newline at end of file
+4

From afdf6ade70897ec97fbe0aa39cb8baf67ba9b405 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 21 Aug 2024 09:26:43 +0200
Subject: [PATCH 42/67] update docstring

---
 include/OGL/MatrixWrapper/SparsityPattern.H | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/include/OGL/MatrixWrapper/SparsityPattern.H b/include/OGL/MatrixWrapper/SparsityPattern.H
index a40ace6f3..752d33b19 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.H
+++ b/include/OGL/MatrixWrapper/SparsityPattern.H
@@ -15,7 +15,7 @@ namespace Foam {
  *
  * After merging sparsity patterns during repartitioning the ldu mapping are
  * not consecutive ie they could look like [0 1 2 3 | 0 1 2 3 | 0 1], where
- * | is the former rank boundary based on the comm pattern we compute a new
+ * | is the former rank boundary. Based on the comm pattern we compute a new
  * mapping as  [ 0 1 2 3 | 4 5 6 7| 8 9 ] where the following offsets [ 0 |
  * 4 | 8 ] based on the recv_counts are used.
  *
@@ -29,7 +29,6 @@ void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
  *
  * Additionally it keeps track which parts of the sparsity pattern belongs to
  * which interface.
- *
  * */
 struct SparsityPattern {
     /* constructor from vectors, assumes a single interface */

From 341a5dfee0db67e0e59b612d07131e2f1c0a3c4a Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Sat, 24 Aug 2024 14:25:31 +0200
Subject: [PATCH 43/67] add docstrings, constructor, and compute dimensions
 function

---
 include/OGL/MatrixWrapper/SparsityPattern.H | 34 +++++++++++++++------
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/include/OGL/MatrixWrapper/SparsityPattern.H b/include/OGL/MatrixWrapper/SparsityPattern.H
index 752d33b19..81197c476 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.H
+++ b/include/OGL/MatrixWrapper/SparsityPattern.H
@@ -24,6 +24,14 @@ void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
                                   gko::array<label> &ldu_mapping, label rank,
                                   label ranks_per_gpu);
 
+/* @brief computes the dimensions of square matrix based row index array
+ *
+ * @note assumes that rows are ordered
+ * @params ordered array of row indices
+ * returns dimensions
+ * */
+gko::dim<2> compute_dimensions(const gko::array<label> &rows);
+
 /* The SparsityPattern holds row and column index data using gko::arrays.
  * This struct is used for easy generation of (distributed) ginkgo matrices.
  *
@@ -32,13 +40,17 @@ void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
  * */
 struct SparsityPattern {
     /* constructor from vectors, assumes a single interface */
-    SparsityPattern(std::vector<label> rows, std::vector<label> cols,
-                    std::vector<label> mapping, std::vector<label> rank)
-        : num_nnz(rows.size())
-    {
-        // NOT fully implementented
-        FatalErrorInFunction << "Not implemented" << abort(FatalError);
-    }
+    SparsityPattern(std::shared_ptr<const gko::Executor> exec, gko::dim<2> dim_,
+                    std::vector<label> rows, std::vector<label> cols,
+                    std::vector<label> mapping, std::vector<label> rank_)
+        : num_nnz(rows.size()),
+          row_idxs(exec, rows.begin(), rows.end()),
+          col_idxs(exec, cols.begin(), cols.end()),
+          ldu_mapping(exec, mapping.begin(), mapping.end()),
+          dim(dim_),
+          spans(std::vector{gko::span(0, num_nnz)}),
+          rank(rank_)
+    {}
 
     SparsityPattern(std::shared_ptr<const gko::Executor> exec, const label size)
         : num_nnz(size),
@@ -95,10 +107,14 @@ struct SparsityPattern {
 
     // A vector of spans indicating begin and end of each interface
     // this is used to keep all col and row idx consecutive in memory
+    // TODO this can be refactored, since its original motivation
+    // was to support the distributed_read function, which is not
+    // needed anymore. However, removing would require row and col_idxs to
+    // be stored in a std::vector<gko::Array>
     mutable std::vector<gko::span> spans;
 
-    // keep track of original (communication) rank of each interface or matrix
-    // block
+    // keep track of original (local sparsity) or communication (non-local) rank
+    // of each interface or matrix block
     mutable std::vector<label> rank;
 };
 

From edc9debca64287b50bbfc165f1b142e2d7fb7bb4 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Sat, 24 Aug 2024 14:30:28 +0200
Subject: [PATCH 44/67] implement exchange spans method

---
 include/OGL/Repartitioner.H |  53 ++++++++----
 src/Repartitioner.C         | 163 ++++++++++++++++++++++--------------
 unitTests/Repartitioner.C   |  74 ++++++++++++++++
 3 files changed, 210 insertions(+), 80 deletions(-)

diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index a43692300..94860a4b5 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -11,26 +11,23 @@
 #include "CommunicationPattern.H"
 #include "MatrixWrapper/SparsityPattern.H"
 
+namespace detail {
+/* @brief exchange spans and rank between owner and non-owner ranks
+**
+** returns a pair vector spans
+*/
+std::pair<std::vector<gko::span>, std::vector<label>> exchange_span_ranks(
+    const ExecutorHandler &exec_handler, label ranks_per_gpu,
+    const std::vector<gko::span> &spans, const std::vector<label> &ranks) const;
+
 template <typename T>
 std::vector<T> apply_permutation(const std::vector<T> vec,
-                                 const std::vector<label> &p)
-{
-    std::vector<T> sorted_vec(vec.size());
-    std::transform(p.begin(), p.end(), sorted_vec.begin(),
-                   [&](label i) { return vec[i]; });
-    return sorted_vec;
-}
+                                 const std::vector<label> &p);
 
 template <typename T, typename Compare>
-std::vector<label> sort_permutation(const std::vector<T> &vec, Compare compare)
-{
-    std::vector<label> p(vec.size());
-    std::iota(p.begin(), p.end(), 0);
-    std::stable_sort(p.begin(), p.end(), [&](std::size_t i, std::size_t j) {
-        return compare(vec[i], vec[j]);
-    });
-    return p;
-}
+std::vector<label> sort_permutation(const std::vector<T> &vec, Compare compare);
+
+}  // namespace detail
 
 /** @class Collects functions to repartition communication patterns
  * and matrices. Here repartitioning refers to changing a given partitioning on
@@ -115,7 +112,27 @@ public:
     std::shared_ptr<const partition> get_orig_partition() const
     {
         return orig_partition_;
-    };
+    }
+
+
+    // std::vector<label> offset_column_indices() {
+    // // the non local cols are in local idx of other side
+    // // thus we need the new offset of the other side
+    // // NOTE TODO this modifies src_non_local_pattern better do the offsets
+    // // on the owner side after gathering
+    // for (label i = 0; i < src_non_local_pattern->rank.size(); i++) {
+    //     auto comm_rank = src_non_local_pattern->rank[i];
+    //     auto [begin, end] = src_non_local_pattern->spans[i];
+    //     label local_offset = orig_partition_->get_range_bounds()[comm_rank] -
+    //                          orig_partition_->get_range_bounds()[owner_rank];
+    //     auto *data = src_non_local_pattern->col_idxs.get_data() + begin;
+    //     auto size = end - begin;
+    //     std::transform(data, data + size, data,
+    //                    [&](label idx) { return idx + local_offset; });
+    // }
+
+
+    // }
 
 
     /* @brief given received interfaces this function sorts interfaces into
@@ -145,7 +162,7 @@ public:
      *(bool)
      * signals whether this is a new local interface (no communication), and
      *the
-     * second entry (label) tracks the original rank of the interface
+     * second entry (label) stores the original rank of the interface
      */
     std::tuple<std::shared_ptr<SparsityPattern>,
                std::shared_ptr<SparsityPattern>,
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index c002812c4..e0fe56113 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -4,6 +4,58 @@
 
 #include "OGL/Repartitioner.H"
 
+namespace detail {
+std::pair<std::vector<gko::span>, std::vector<label>> exchange_span_ranks(
+    const ExecutorHandler &exec_handler, label ranks_per_gpu,
+    const std::vector<gko::span> &spans, const std::vector<label> &ranks) const
+{
+    auto comm_pattern = compute_gather_to_owner_counts(
+        exec_handler, ranks_per_gpu, spans.size());
+
+    std::vector<label> size{};
+
+    for (auto &elem : spans) {
+        size.push_back(elem.length());
+    }
+
+    auto gathered_size = gather_labels_to_owner(exec_handler, comm_pattern,
+                                                size.data(), size.size());
+
+    std::vector<gko::span> out_spans{};
+    int count = 0;
+    for (auto length : gathered_size) {
+        out_spans.emplace_back(count, count + length);
+        count += length;
+    }
+
+    auto new_ranks = gather_labels_to_owner(exec_handler, comm_pattern,
+                                            ranks.data(), ranks.size());
+
+    return {out_spans, new_ranks};
+}
+
+template <typename T>
+std::vector<T> apply_permutation(const std::vector<T> vec,
+                                 const std::vector<label> &p)
+{
+    std::vector<T> sorted_vec(vec.size());
+    std::transform(p.begin(), p.end(), sorted_vec.begin(),
+                   [&](label i) { return vec[i]; });
+    return sorted_vec;
+}
+
+template <typename T, typename Compare>
+std::vector<label> sort_permutation(const std::vector<T> &vec, Compare compare)
+{
+    std::vector<label> p(vec.size());
+    std::iota(p.begin(), p.end(), 0);
+    std::stable_sort(p.begin(), p.end(), [&](std::size_t i, std::size_t j) {
+        return compare(vec[i], vec[j]);
+    });
+    return p;
+}
+
+}  // namespace detail
 
 label Repartitioner::compute_repart_size(label local_size, label ranks_per_gpu,
                                          const ExecutorHandler &exec_handler)
@@ -26,6 +78,7 @@ Repartitioner::repartition_sparsity(
     std::shared_ptr<const SparsityPattern> src_non_local_pattern) const
 {
     LOG_1(verbose_, "start repartition sparsity pattern")
+
     // 1. obtain send recv sizes vector
     // here we can reuse code from repartition_comm_pattern
     //
@@ -41,33 +94,37 @@ Repartitioner::repartition_sparsity(
     label rank = exec_handler.get_rank();
     label owner_rank = get_owner_rank(exec_handler);
     label ranks_per_gpu = ranks_per_gpu_;
+
     // TODO dont copy
-    // if (ranks_per_gpu == 1) {
-    //     std::vector<std::pair<bool, label>> ret;
-    //     for (auto comm_rank : src_non_local_pattern->rank) {
-    //         ret.emplace_back(false, rank);
-    //     }
-    //     return std::make_tuple<std::shared_ptr<SparsityPattern>,
-    //                            std::shared_ptr<SparsityPattern>,
-    //                            std::vector<std::pair<bool, label>>>(
-    //         std::make_shared<SparsityPattern>(src_local_pattern),
-    //         std::make_shared<SparsityPattern>(src_non_local_pattern),
-    //         std::move(ret));
-    // }
+    if (ranks_per_gpu == 1) {
+        std::vector<std::pair<bool, label>> ret;
+        for (auto comm_rank : src_non_local_pattern->rank) {
+            ret.emplace_back(false, rank);
+        }
+        return std::make_tuple<std::shared_ptr<SparsityPattern>,
+                               std::shared_ptr<SparsityPattern>,
+                               std::vector<std::pair<bool, label>>>(
+            std::make_shared<SparsityPattern>(*src_local_pattern.get()),
+            std::make_shared<SparsityPattern>(*src_non_local_pattern.get()),
+            std::move(ret));
+    }
 
     auto local_comm_pattern = compute_gather_to_owner_counts(
         exec_handler, ranks_per_gpu, src_local_pattern->num_nnz);
 
-
+    // row index offset relative to its owner rank
     auto offset = orig_partition_->get_range_bounds()[rank] -
                   orig_partition_->get_range_bounds()[owner_rank];
 
-    auto gather_closure = [&](auto &comm_pattern, auto &data, label offset) {
+    auto gather_closure = [exec_handler](auto &comm_pattern, auto &data,
+                                         label offset) {
         return gather_labels_to_owner(exec_handler, comm_pattern,
                                       data.get_data(), data.get_size(), offset);
     };
 
     SparsityPattern merged_local{
+        exec,
+        src_local_pattern->dim,
         gather_closure(local_comm_pattern, src_local_pattern->row_idxs, offset),
         gather_closure(local_comm_pattern, src_local_pattern->col_idxs, offset),
         gather_closure(local_comm_pattern, src_local_pattern->ldu_mapping, 0),
@@ -78,60 +135,42 @@ Repartitioner::repartition_sparsity(
             local_comm_pattern, merged_local.ldu_mapping, rank, ranks_per_gpu);
     }
 
-    auto get_back = [](const gko::array<label> &in) {
-        return *(in.get_const_data() + in.get_size());
-    };
-
-    label rows =
-        (is_owner(exec_handler)) ? get_back(merged_local.row_idxs) + 1 : 0;
-    gko::dim<2> merged_local_dim{static_cast<gko::size_type>(rows),
-                                 static_cast<gko::size_type>(rows)};
+    gko::dim<2> merged_local_dim =
+        (is_owner(exec_handler)) ? compute_dimensions(merged_local.row_idxs)
+                                 : gko::dim<2>{0, 0};
 
     auto non_local_comm_pattern = compute_gather_to_owner_counts(
         exec_handler, ranks_per_gpu, src_non_local_pattern->num_nnz);
 
-    std::vector<label> spans_begin;
-    std::vector<label> spans_end;
-    auto span_comm_pattern = compute_gather_to_owner_counts(
-        exec_handler, ranks_per_gpu, src_non_local_pattern->spans.size());
-    for (auto elem : src_non_local_pattern->spans) {
-        spans_begin.push_back(elem.begin);
-        spans_end.push_back(elem.end);
-    }
 
-    // the non local cols are in local idx of other side
-    // thus we need the new offset of the other side
-    // NOTE TODO this modifies src_non_local_pattern better do the offsets
-    // on the owner side after gathering
-    for (label i = 0; i < src_non_local_pattern->rank.size(); i++) {
-        auto comm_rank = src_non_local_pattern->rank[i];
-        auto [begin, end] = src_non_local_pattern->spans[i];
-        label local_offset = orig_partition_->get_range_bounds()[comm_rank] -
-                             orig_partition_->get_range_bounds()[owner_rank];
-        auto *data = src_non_local_pattern->col_idxs.get_data() + begin;
-        auto size = end - begin;
-        std::transform(data, data + size, data,
-                       [&](label idx) { return idx + local_offset; });
-    }
+    // auto offseted_column_indices = offset_column_indices();
+
+
+    SparsityPattern merged_non_local{*src_non_local_pattern.get()};
+    //   SparsityPattern merged_non_local{exec, src_non_local_pattern->dim,
+    //       gather_closure(non_local_comm_pattern,
+    //       src_non_local_pattern->row_idxs,
+    //                      offset),
+    //       gather_closure(non_local_comm_pattern,
+    //       src_non_local_pattern->col_idxs,
+    //                      0),
+    //       gather_closure(non_local_comm_pattern,
+    //       src_non_local_pattern->row_idxs,
+    //                      0),
+    //       // {},
+    //       // this doesn't exist
+    //       // gather_labels_to_owner(exec_handler, span_comm_pattern,
+    //       // spans_begin.size(),
+    //       //                spans_begin.data()),
+    //       //  gather_labels_to_owner(exec_handler, span_comm_pattern,
+    //       //  spans_end.size(),
+    //       //                spans_end.data()),
+    //       gather_labels_to_owner(exec_handler, span_comm_pattern,
+    //                              src_non_local_pattern->rank.data(),
+    //                              src_non_local_pattern->rank.size())};
 
-    SparsityPattern merged_non_local{
-        gather_closure(non_local_comm_pattern, src_non_local_pattern->row_idxs,
-                       offset),
-        gather_closure(non_local_comm_pattern, src_non_local_pattern->col_idxs,
-                       0),
-        gather_closure(non_local_comm_pattern, src_non_local_pattern->row_idxs,
-                       0),
-        // {},
-        // this doesn't exist
-        // gather_labels_to_owner(exec_handler, span_comm_pattern,
-        // spans_begin.size(),
-        //                spans_begin.data()),
-        //  gather_labels_to_owner(exec_handler, span_comm_pattern,
-        //  spans_end.size(),
-        //                spans_end.data()),
-        gather_labels_to_owner(exec_handler, span_comm_pattern,
-                               src_non_local_pattern->rank.data(),
-                               src_non_local_pattern->rank.size())};
+    auto span_comm_pattern = compute_gather_to_owner_counts(
+        exec_handler, ranks_per_gpu, src_non_local_pattern->spans.size());
 
     if (is_owner(exec_handler)) {
         FatalErrorInFunction << "Not implemented" << abort(FatalError);
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 299e58f2d..c10055fab 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -100,6 +100,80 @@ TEST_P(RepartitionerFixture, has_correct_properties_for_n_rank)
         (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0));
 }
 
+TEST_P(RepartitionerFixture, can_exchange_spans_and_ranks_for_n_ranks)
+{
+    using vec_vec = std::vector<std::vector<label>>;
+    // Arrange
+    label ranks_per_gpu = GetParam();
+    auto exec = *((RepartitionerEnvironment *)global_env)->exec.get();
+    auto rank = exec.get_rank();
+    auto ref_exec = exec.get_ref_exec();
+
+    std::vector<gko::span> spans{{0, 5}, {5, 10}};
+    std::vector<label> ranks{rank + 1, rank + 2};
+
+    std::map<label, vec_vec> exp_ranks;
+    exp_ranks.emplace(1, vec_vec{{1, 2}, {2, 3}, {3, 4}, {4, 5}});
+    exp_ranks.emplace(2, vec_vec{{1, 2, 2, 3}, {}, {3, 4, 4, 5}, {}});
+    exp_ranks.emplace(4, vec_vec{{1, 2, 2, 3, 3, 4, 4, 5}, {}, {}, {}});
+
+    // Act
+    auto [new_spans, new_ranks] =
+        detail::exchange_span_ranks(exec, ranks_per_gpu, spans, ranks);
+
+    // Assert
+    ASSERT_EQ(new_ranks, exp_ranks[ranks_per_gpu][rank]);
+}
+
+TEST_P(RepartitionerFixture, can_repartition_sparsity_pattern_for_n_ranks)
+{
+    // Arrange
+    label ranks_per_gpu = GetParam();
+    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
+    auto rank = exec->get_rank();
+    auto ref_exec = exec->get_ref_exec();
+
+    // mesh:
+    // rank:     0     1     2     3
+    // cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
+    // cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- local row ids
+
+    std::vector<label> rows{0, 0, 1, 1};
+    std::vector<label> cols{0, 1, 0, 1};
+    std::vector<label> mapping{0, 1, 2, 3};
+    std::vector<label> ranks{rank};
+    auto local_sparsity = std::make_shared<SparsityPattern>(
+        ref_exec, gko::dim<2>{2, 2}, rows, cols, mapping, ranks);
+
+    std::vector<label> non_local_rows{1};
+    std::vector<label> non_local_cols{0};
+    std::vector<label> non_local_mapping{0};
+    std::vector<label> non_local_ranks{rank};
+    auto non_local_sparsity = std::make_shared<SparsityPattern>(
+        ref_exec, gko::dim<2>{2, 1}, non_local_rows, non_local_cols,
+        non_local_mapping, non_local_ranks);
+
+    std::map<label, std::vector<label>> exp_local_nnz;
+    exp_local_nnz.emplace(1, std::vector<label>{4, 4, 4, 4});
+    exp_local_nnz.emplace(2, std::vector<label>{8, 0, 8, 0});
+    exp_local_nnz.emplace(4, std::vector<label>{16, 0, 0, 0});
+
+    std::map<label, std::vector<label>> exp_non_local_nnz;
+    exp_local_nnz.emplace(1, std::vector<label>{1, 2, 2, 1});
+    exp_local_nnz.emplace(2, std::vector<label>{1, 0, 1, 0});
+    exp_local_nnz.emplace(4, std::vector<label>{0, 0, 0, 0});
+
+    // Act
+    auto [repart_local, repart_non_local, tracking] =
+        repartitioner.repartition_sparsity(*exec, local_sparsity,
+                                           non_local_sparsity);
+
+    // Assert
+    ASSERT_EQ(repart_local->num_nnz, exp_local_nnz[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_non_local->num_nnz,
+              exp_non_local_nnz[ranks_per_gpu][rank]);
+}
 
 TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
 {

From 124e32423cd88db7e8c9a7adf2e6ee23c6cbaf2b Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Sat, 24 Aug 2024 14:30:52 +0200
Subject: [PATCH 45/67] add compute_dimensions implementation

---
 src/MatrixWrapper/SparsityPattern.C | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/src/MatrixWrapper/SparsityPattern.C b/src/MatrixWrapper/SparsityPattern.C
index 29209b812..3f88a480c 100644
--- a/src/MatrixWrapper/SparsityPattern.C
+++ b/src/MatrixWrapper/SparsityPattern.C
@@ -6,6 +6,19 @@
 
 namespace Foam {
 
+gko::dim<2> compute_dimensions(const gko::array<label> &rows)
+{
+    // helper function to get the last element of a ginkgo array
+    auto get_back = [](const gko::array<label> &in) {
+        return *(in.get_const_data() + in.get_size());
+    };
+
+    gko::size_type num_rows = get_back(rows) + 1;
+
+    return gko::dim<2>{num_rows, num_rows};
+}
+
+
 void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
                                   gko::array<label> &ldu_mapping, label rank,
                                   label ranks_per_gpu)

From 52a29fdb2a4db72a51e00fc9f79c7f927f7ab68d Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Mon, 26 Aug 2024 21:07:11 +0200
Subject: [PATCH 46/67] WIP repartition sparsisty

---
 include/OGL/MatrixWrapper/SparsityPattern.H |  21 +-
 include/OGL/Repartitioner.H                 |  92 +++--
 include/OGL/common.H                        |   3 +
 src/MatrixWrapper/SparsityPattern.C         |  15 +-
 src/Repartitioner.C                         | 435 ++++++++++----------
 src/common.C                                |   5 +
 unitTests/Repartitioner.C                   | 174 +++++++-
 7 files changed, 482 insertions(+), 263 deletions(-)

diff --git a/include/OGL/MatrixWrapper/SparsityPattern.H b/include/OGL/MatrixWrapper/SparsityPattern.H
index 81197c476..8f3645edd 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.H
+++ b/include/OGL/MatrixWrapper/SparsityPattern.H
@@ -21,7 +21,7 @@ namespace Foam {
  *
  * */
 void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
-                                  gko::array<label> &ldu_mapping, label rank,
+                                  std::vector<label> &ldu_mapping, label rank,
                                   label ranks_per_gpu);
 
 /* @brief computes the dimensions of square matrix based row index array
@@ -30,7 +30,7 @@ void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
  * @params ordered array of row indices
  * returns dimensions
  * */
-gko::dim<2> compute_dimensions(const gko::array<label> &rows);
+gko::dim<2> compute_dimensions(const std::vector<label> &rows);
 
 /* The SparsityPattern holds row and column index data using gko::arrays.
  * This struct is used for easy generation of (distributed) ginkgo matrices.
@@ -41,17 +41,28 @@ gko::dim<2> compute_dimensions(const gko::array<label> &rows);
 struct SparsityPattern {
     /* constructor from vectors, assumes a single interface */
     SparsityPattern(std::shared_ptr<const gko::Executor> exec, gko::dim<2> dim_,
-                    std::vector<label> rows, std::vector<label> cols,
-                    std::vector<label> mapping, std::vector<label> rank_)
+                    const std::vector<label> &rows,
+                    const std::vector<label> &cols,
+                    const std::vector<label> &mapping,
+                    const std::vector<gko::span> &spans_,
+                    const std::vector<label> &rank_)
         : num_nnz(rows.size()),
           row_idxs(exec, rows.begin(), rows.end()),
           col_idxs(exec, cols.begin(), cols.end()),
           ldu_mapping(exec, mapping.begin(), mapping.end()),
           dim(dim_),
-          spans(std::vector{gko::span(0, num_nnz)}),
+          spans(spans_),
           rank(rank_)
     {}
 
+    SparsityPattern(std::shared_ptr<const gko::Executor> exec)
+        : num_nnz(0),
+          dim(gko::dim<2>{}),
+          row_idxs{exec},
+          col_idxs{exec},
+          ldu_mapping{exec}
+    {}
+
     SparsityPattern(std::shared_ptr<const gko::Executor> exec, const label size)
         : num_nnz(size),
           row_idxs{exec, static_cast<gko::size_type>(size)},
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index 94860a4b5..40ef76ffc 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -12,18 +12,69 @@
 #include "MatrixWrapper/SparsityPattern.H"
 
 namespace detail {
+
+/* Convert to global
+** Given an array of column indices local to the communication rank
+** this function offsets these array
+** @param idx pointer to gko::array holding the indices which need to be
+*converted from local to global ids
+** @param spans start and ends of the interfaces
+** @param ranks the rank to which the interface index is a local row index
+*/
+std::vector<label> convert_to_global(
+    std::shared_ptr<
+        const gko::experimental::distributed::Partition<label, label>>
+        partition,
+    const std::vector<label> &idx, const std::vector<gko::span> &spans,
+    const std::vector<label> &ranks)
+{
+    std::vector<label> ret;
+    ret.reserve(idx.size());
+
+    for (label i = 0; i < ranks.size(); i++) {
+        auto rank = ranks[i];
+        auto [begin, end] = spans[i];
+        label offset = partition->get_range_bounds()[rank];
+        for (label j = begin; j < end; j++) {
+            ret.push_back(idx.data()[j] + offset);
+        }
+    }
+    return ret;
+}
+
+/* Convert to local
+** Given an array of column indices local to the communication rank
+** this function offsets these array
+** @param idx pointer to gko::array holding the indices which need to be
+*converted from local to global ids
+** @param spans start and ends of the interfaces
+** @param ranks the rank to which the interface index is a local row index
+*/
+void convert_to_local(
+    std::shared_ptr<
+        const gko::experimental::distributed::Partition<label, label>>
+        partition,
+    std::vector<label> &in, label rank)
+{
+    label offset = partition->get_range_bounds()[rank];
+    std::transform(in.begin(), in.end(), in.begin(),
+                   [&](label idx) { return idx - offset; });
+}
+
 /* @brief exchange spans and rank between owner and non-owner ranks
 **
 ** returns a pair vector spans
 */
 std::pair<std::vector<gko::span>, std::vector<label>> exchange_span_ranks(
     const ExecutorHandler &exec_handler, label ranks_per_gpu,
-    const std::vector<gko::span> &spans, const std::vector<label> &ranks) const;
+    const std::vector<gko::span> &spans, const std::vector<label> &ranks);
 
 template <typename T>
 std::vector<T> apply_permutation(const std::vector<T> vec,
                                  const std::vector<label> &p);
 
+/* Sorts a vector and returns the permutation for reuse
+ */
 template <typename T, typename Compare>
 std::vector<label> sort_permutation(const std::vector<T> &vec, Compare compare);
 
@@ -115,36 +166,27 @@ public:
     }
 
 
-    // std::vector<label> offset_column_indices() {
-    // // the non local cols are in local idx of other side
-    // // thus we need the new offset of the other side
-    // // NOTE TODO this modifies src_non_local_pattern better do the offsets
-    // // on the owner side after gathering
-    // for (label i = 0; i < src_non_local_pattern->rank.size(); i++) {
-    //     auto comm_rank = src_non_local_pattern->rank[i];
-    //     auto [begin, end] = src_non_local_pattern->spans[i];
-    //     label local_offset = orig_partition_->get_range_bounds()[comm_rank] -
-    //                          orig_partition_->get_range_bounds()[owner_rank];
-    //     auto *data = src_non_local_pattern->col_idxs.get_data() + begin;
-    //     auto size = end - begin;
-    //     std::transform(data, data + size, data,
-    //                    [&](label idx) { return idx + local_offset; });
-    // }
-
-
-    // }
-
-
     /* @brief given received interfaces this function sorts interfaces into
+     *
      *local and non local returns: SparsityPattern storing new non_local
      *sparsity pattern, and a pair of locality information, where the first bool
      *stores if the interface performs a local communication and the second bool
      *if the the interface was originally from this rank
+     *
      */
-    std::pair<SparsityPattern, std::vector<bool>> build_non_local_interfaces(
-        const ExecutorHandler &exec_handler, SparsityPattern &loc,
-        const SparsityPattern &non_loc) const;
-
+    std::vector<bool> build_non_local_interfaces(
+        const ExecutorHandler &exec_handler,
+        std::shared_ptr<
+            const gko::experimental::distributed::Partition<label, label>>
+            partition,
+        std::vector<label> &local_rows, std::vector<label> &local_cols,
+        std::vector<label> &local_mapping, std::vector<label> &local_ranks,
+        std::vector<gko::span> &local_spans, std::vector<label> &non_local_rows,
+        std::vector<label> &non_local_cols,
+        std::vector<label> &non_local_mapping,
+        std::vector<label> &non_local_ranks,
+        std::vector<label> &non_local_rank_origin,
+        std::vector<gko::span> &non_local_spans) const;
 
     /* @brief function to repartition a sparsity pattern
      *
diff --git a/include/OGL/common.H b/include/OGL/common.H
index 08a31dbe1..4175f794f 100644
--- a/include/OGL/common.H
+++ b/include/OGL/common.H
@@ -131,4 +131,7 @@ void set_next_caching(word sys_matrix_name, const objectRegistry &db,
                       label caching);
 
 label get_next_caching(word sys_matrix_name, const objectRegistry &db);
+
+std::shared_ptr<gko::array<label>> convert_to_array(
+    const std::vector<label> &in);
 }  // namespace Foam
diff --git a/src/MatrixWrapper/SparsityPattern.C b/src/MatrixWrapper/SparsityPattern.C
index 3f88a480c..6e6cb5408 100644
--- a/src/MatrixWrapper/SparsityPattern.C
+++ b/src/MatrixWrapper/SparsityPattern.C
@@ -6,26 +6,19 @@
 
 namespace Foam {
 
-gko::dim<2> compute_dimensions(const gko::array<label> &rows)
+gko::dim<2> compute_dimensions(const std::vector<label> &rows)
 {
-    // helper function to get the last element of a ginkgo array
-    auto get_back = [](const gko::array<label> &in) {
-        return *(in.get_const_data() + in.get_size());
-    };
-
-    gko::size_type num_rows = get_back(rows) + 1;
-
+    gko::size_type num_rows = rows.back() + 1;
     return gko::dim<2>{num_rows, num_rows};
 }
 
 
 void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
-                                  gko::array<label> &ldu_mapping, label rank,
+                                  std::vector<label> &ldu_mapping, label rank,
                                   label ranks_per_gpu)
 {
-    // TODO check if ldu_mapping is on host exec
     label ldu_offset = 0;
-    auto *data = ldu_mapping.get_data();
+    auto *data = ldu_mapping.data();
 
     for (label i = 0; i < ranks_per_gpu; i++) {
         auto size = comm_pattern.recv_counts[i];
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index e0fe56113..c7613a5f4 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -5,9 +5,10 @@
 #include "OGL/Repartitioner.H"
 
 namespace detail {
+
 std::pair<std::vector<gko::span>, std::vector<label>> exchange_span_ranks(
     const ExecutorHandler &exec_handler, label ranks_per_gpu,
-    const std::vector<gko::span> &spans, const std::vector<label> &ranks) const
+    const std::vector<gko::span> &spans, const std::vector<label> &ranks)
 {
     auto comm_pattern = compute_gather_to_owner_counts(
         exec_handler, ranks_per_gpu, spans.size());
@@ -79,23 +80,20 @@ Repartitioner::repartition_sparsity(
 {
     LOG_1(verbose_, "start repartition sparsity pattern")
 
-    // 1. obtain send recv sizes vector
-    // here we can reuse code from repartition_comm_pattern
-    //
-    // 2. initialize send and recv buffer
-    // if we keep interfaces separated row and col indices can be
-    //
-    // 3. ldu mapping needs to be update?
-    // 4. dim needs to be updated
-    // this is not implemented yet, but we can't fail here for
-    // debugging reasons
+    // helper function to simplify calling gather_labels_to_owner
+    auto gather_closure = [exec_handler](auto &comm_pattern, auto &data,
+                                         label offset) {
+        return gather_labels_to_owner(exec_handler, comm_pattern,
+                                      data.get_data(), data.get_size(), offset);
+    };
+
     auto exec = exec_handler.get_ref_exec();
     auto comm = *exec_handler.get_communicator().get();
     label rank = exec_handler.get_rank();
     label owner_rank = get_owner_rank(exec_handler);
     label ranks_per_gpu = ranks_per_gpu_;
 
-    // TODO dont copy
+    // early return if no repartitioning requested
     if (ranks_per_gpu == 1) {
         std::vector<std::pair<bool, label>> ret;
         for (auto comm_rank : src_non_local_pattern->rank) {
@@ -109,229 +107,242 @@ Repartitioner::repartition_sparsity(
             std::move(ret));
     }
 
+    // Step 1. gather all local sparsity pattern to owner rank
+    // this is a all in collected into std::vectors, thus we can
+    // append all the interfaces which became local after repartitioning
     auto local_comm_pattern = compute_gather_to_owner_counts(
         exec_handler, ranks_per_gpu, src_local_pattern->num_nnz);
 
-    // row index offset relative to its owner rank
+    // row/column index offset relative to its owner rank
     auto offset = orig_partition_->get_range_bounds()[rank] -
                   orig_partition_->get_range_bounds()[owner_rank];
 
-    auto gather_closure = [exec_handler](auto &comm_pattern, auto &data,
-                                         label offset) {
-        return gather_labels_to_owner(exec_handler, comm_pattern,
-                                      data.get_data(), data.get_size(), offset);
-    };
+    auto tmp_local_rows =
+        gather_closure(local_comm_pattern, src_local_pattern->row_idxs, offset);
+    auto tmp_local_cols =
+        gather_closure(local_comm_pattern, src_local_pattern->col_idxs, offset);
+    auto tmp_local_mapping =
+        gather_closure(local_comm_pattern, src_local_pattern->ldu_mapping, 0);
 
-    SparsityPattern merged_local{
-        exec,
-        src_local_pattern->dim,
-        gather_closure(local_comm_pattern, src_local_pattern->row_idxs, offset),
-        gather_closure(local_comm_pattern, src_local_pattern->col_idxs, offset),
-        gather_closure(local_comm_pattern, src_local_pattern->ldu_mapping, 0),
-        std::vector{rank}};
 
     if (is_owner(exec_handler)) {
-        make_ldu_mapping_consecutive(
-            local_comm_pattern, merged_local.ldu_mapping, rank, ranks_per_gpu);
+        make_ldu_mapping_consecutive(local_comm_pattern, tmp_local_mapping,
+                                     rank, ranks_per_gpu);
     }
+    std::vector<gko::span> tmp_local_span{gko::span{0, tmp_local_rows.size()}};
+    std::vector<label> tmp_local_rank{rank};
 
-    gko::dim<2> merged_local_dim =
-        (is_owner(exec_handler)) ? compute_dimensions(merged_local.row_idxs)
-                                 : gko::dim<2>{0, 0};
+    gko::dim<2> tmp_local_dim = (is_owner(exec_handler))
+                                    ? compute_dimensions(tmp_local_rows)
+                                    : gko::dim<2>{0, 0};
+
+    // Done with Step 1. Next we gather all non local sparsity. Again
+    // all of it is collected into temporary vectors first.
 
     auto non_local_comm_pattern = compute_gather_to_owner_counts(
         exec_handler, ranks_per_gpu, src_non_local_pattern->num_nnz);
 
+    auto tmp_non_local_rows = gather_closure(
+        non_local_comm_pattern, src_non_local_pattern->row_idxs, offset);
 
-    // auto offseted_column_indices = offset_column_indices();
-
-
-    SparsityPattern merged_non_local{*src_non_local_pattern.get()};
-    //   SparsityPattern merged_non_local{exec, src_non_local_pattern->dim,
-    //       gather_closure(non_local_comm_pattern,
-    //       src_non_local_pattern->row_idxs,
-    //                      offset),
-    //       gather_closure(non_local_comm_pattern,
-    //       src_non_local_pattern->col_idxs,
-    //                      0),
-    //       gather_closure(non_local_comm_pattern,
-    //       src_non_local_pattern->row_idxs,
-    //                      0),
-    //       // {},
-    //       // this doesn't exist
-    //       // gather_labels_to_owner(exec_handler, span_comm_pattern,
-    //       // spans_begin.size(),
-    //       //                spans_begin.data()),
-    //       //  gather_labels_to_owner(exec_handler, span_comm_pattern,
-    //       //  spans_end.size(),
-    //       //                spans_end.data()),
-    //       gather_labels_to_owner(exec_handler, span_comm_pattern,
-    //                              src_non_local_pattern->rank.data(),
-    //                              src_non_local_pattern->rank.size())};
-
-    auto span_comm_pattern = compute_gather_to_owner_counts(
-        exec_handler, ranks_per_gpu, src_non_local_pattern->spans.size());
+    auto [new_spans, new_ranks] = detail::exchange_span_ranks(
+        exec_handler, ranks_per_gpu_, src_non_local_pattern->spans,
+        src_non_local_pattern->rank);
 
-    if (is_owner(exec_handler)) {
-        FatalErrorInFunction << "Not implemented" << abort(FatalError);
-        //  make_begin_and_ends_consecutive(
-        //      span_comm_pattern, merged_non_local.spans,
-        //     rank, ranks_per_gpu);
-    }
+    auto tmp_non_local_cols = detail::convert_to_global(
+        orig_partition_,
+        gather_closure(non_local_comm_pattern, src_non_local_pattern->col_idxs,
+                       0),
+        new_spans, new_ranks);
 
-    auto [gathered_non_local, is_local] = build_non_local_interfaces(
-        exec_handler, merged_local, merged_non_local);
+    std::vector<label> tmp_non_local_mapping =
+        std::vector<label>(tmp_non_local_rows.size());
 
-    // build vector with locality information
-    std::vector<std::pair<bool, label>> locality;
-    label ctr{0};
+    std::vector<label> tmp_non_local_origin{};
 
-    if (is_owner(exec_handler)) {
-        for (int i = 0; i < ranks_per_gpu; i++) {
-            label comm_rank = rank + i;
-            label num_elems = span_comm_pattern.recv_counts[comm_rank];
-            for (int j = 0; j < num_elems; j++) {
-                locality.emplace_back(is_local[ctr], comm_rank);
-                ctr++;
-            }
-        }
-    }
-    // All interfaces that are in is_local need to be accounted for
-    ASSERT_EQ(ctr, is_local.size());
-
-    // TODO FIXME the correct way would be to check the communication
-    // pattern where a particular face is in the send idxs
-    // since we have already the row id of the other side
-    // it should be doable. Or alternatively we know that we
-    // keep an interface together thus we can just count the idx up to the
-    // size. But we have to make sure that the interfaces are in the same
-    // order on both communication sides.
-    for (int i = 0; i < gathered_non_local.col_idxs.get_size(); i++) {
-        gathered_non_local.col_idxs.get_data()[i] = i;
-    }
+    auto is_local = build_non_local_interfaces(
+        exec_handler, orig_partition_, tmp_local_rows, tmp_local_cols,
+        tmp_local_mapping, tmp_local_rank, tmp_local_span, tmp_non_local_rows,
+        tmp_non_local_cols, tmp_non_local_mapping, new_ranks,
+        tmp_non_local_origin, new_spans);
+
+    // TODO FIXME
+    gko::dim<2> tmp_non_local_dim {0, 0};
 
-    LOG_1(verbose_, "done repartition sparsity pattern")
     if (is_owner(exec_handler)) {
-        FatalErrorInFunction << "Not implemented" << abort(FatalError);
-        // auto new_local_spars_pattern = std::make_shared<SparsityPattern>(
-        //     exec, merged_local_dim, merged_local);
-        //
-        // auto new_non_local_spars_pattern =
-        //     std::make_shared<SparsityPattern>(
-        //         exec,
-        //         gko::dim<2>{merged_local_dim[0],
-        //                     gathered_non_local.row_idxs.get_size()},
-        //         gathered_non_local);
-        //
-        //         return std::make_tuple<std::shared_ptr<SparsityPattern>,
-        //                                std::shared_ptr<SparsityPattern>,
-        //                                std::vector<std::pair<bool,
-        //                                label>>>(
-        //             std::move(new_local_spars_pattern),
-        //             std::move(new_non_local_spars_pattern),
-        //             std::move(locality));
+        auto local_pattern = std::make_shared<SparsityPattern>(
+            exec, tmp_local_dim, tmp_local_rows, tmp_local_cols,
+            tmp_local_mapping, tmp_local_span, tmp_local_rank);
+        auto non_local_pattern = std::make_shared<SparsityPattern>(
+            exec, tmp_non_local_dim, tmp_non_local_rows, tmp_non_local_cols,
+            tmp_non_local_mapping, new_spans, new_ranks);
+
+        return std::make_tuple<std::shared_ptr<SparsityPattern>,
+                               std::shared_ptr<SparsityPattern>,
+                               std::vector<std::pair<bool, label>>>(
+            std::move(local_pattern), std::move(non_local_pattern), {});
     } else {
-        // auto new_local_spars_pattern =
-        //     std::make_shared<SparsityPattern>(exec);
-        //
-        // auto new_non_local_spars_pattern =
-        //     std::make_shared<SparsityPattern>(exec);
-        //
-        // // FIXME
-        // // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
-        // // comm.rank() <<
-        // //     " merged_local dim " <<
-        // new_local_spars_pattern->dim =
-        //     gko::dim<2>{0, 0};  // merged_local_dim;
-        // new_non_local_spars_pattern->dim =
-        //     gko::dim<2>{0, 0};  // merged_local_dim;
-        //
-        //
-        FatalErrorInFunction << "Not implemented" << abort(FatalError);
-        // return std::make_tuple<std::shared_ptr<SparsityPattern>,
-        //                        std::shared_ptr<SparsityPattern>,
-        //                        std::vector<std::pair<bool, label>>>(
-        //     std::move(new_local_spars_pattern),
-        //     std::move(new_non_local_spars_pattern),
-        //     std::vector<std::pair<bool, label>>{});
+        return std::make_tuple<std::shared_ptr<SparsityPattern>,
+                               std::shared_ptr<SparsityPattern>,
+                               std::vector<std::pair<bool, label>>>(
+            std::make_shared<SparsityPattern>(exec),
+            std::make_shared<SparsityPattern>(exec), {});
     }
-    // TODO Some old documentation see what is still needed
-    // add interfaces
-    // to add interfaces we go through the non_local_sparsity
-    // pattern and check if interface is still non_local
-    //
-    // iterate interfaces and send to owner.
-    // on owner decide to move to local or non_local
-    // for sending to owner we can also use the all_to_all_v
-    // approach we than have
-    //
-    // row [ 4 , 8,  12 | 16 , 17, 18 ] <- local row
-    // col [ 1 , 2, 3 | 1, 2, 3 ] <- just the interface ctr
-    // from repartition_comm_pattern we could get
-    //
-    // or we split local / non-local first
-    // for this we could store ranges when computing
-    // repartition_comm_pattern
-    //
+    // auto span_comm_pattern = compute_gather_to_owner_counts(
+    //     exec_handler, ranks_per_gpu, src_non_local_pattern->spans.size());
+    // // build vector with locality information
+    // std::vector<std::pair<bool, label>> locality;
+    // label ctr{0};
+
+    // if (is_owner(exec_handler)) {
+    //     for (int i = 0; i < ranks_per_gpu; i++) {
+    //         label comm_rank = rank + i;
+    //         label num_elems = span_comm_pattern.recv_counts[comm_rank];
+    //         for (int j = 0; j < num_elems; j++) {
+    //             locality.emplace_back(is_local[ctr], comm_rank);
+    //             ctr++;
+    //         }
+    //     }
+    // }
+    // // All interfaces that are in is_local need to be accounted for
+    // ASSERT_EQ(ctr, is_local.size());
+
+    // // TODO FIXME the correct way would be to check the communication
+    // // pattern where a particular face is in the send idxs
+    // // since we have already the row id of the other side
+    // // it should be doable. Or alternatively we know that we
+    // // keep an interface together thus we can just count the idx up to the
+    // // size. But we have to make sure that the interfaces are in the same
+    // // order on both communication sides.
+    // for (int i = 0; i < gathered_non_local.col_idxs.get_size(); i++) {
+    //     gathered_non_local.col_idxs.get_data()[i] = i;
+    // }
+
+    // LOG_1(verbose_, "done repartition sparsity pattern")
+    // if (is_owner(exec_handler)) {
+    //     FatalErrorInFunction << "Not implemented" << abort(FatalError);
+    //     // auto new_local_spars_pattern = std::make_shared<SparsityPattern>(
+    //     //     exec, merged_local_dim, merged_local);
+    //     //
+    //     // auto new_non_local_spars_pattern =
+    //     //     std::make_shared<SparsityPattern>(
+    //     //         exec,
+    //     //         gko::dim<2>{merged_local_dim[0],
+    //     //                     gathered_non_local.row_idxs.get_size()},
+    //     //         gathered_non_local);
+    //     //
+    //     //         return std::make_tuple<std::shared_ptr<SparsityPattern>,
+    //     //                                std::shared_ptr<SparsityPattern>,
+    //     //                                std::vector<std::pair<bool,
+    //     //                                label>>>(
+    //     //             std::move(new_local_spars_pattern),
+    //     //             std::move(new_non_local_spars_pattern),
+    //     //             std::move(locality));
+    // } else {
+    //     // auto new_local_spars_pattern =
+    //     //     std::make_shared<SparsityPattern>(exec);
+    //     //
+    //     // auto new_non_local_spars_pattern =
+    //     //     std::make_shared<SparsityPattern>(exec);
+    //     //
+    //     // // FIXME
+    //     // // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
+    //     // // comm.rank() <<
+    //     // //     " merged_local dim " <<
+    //     // new_local_spars_pattern->dim =
+    //     //     gko::dim<2>{0, 0};  // merged_local_dim;
+    //     // new_non_local_spars_pattern->dim =
+    //     //     gko::dim<2>{0, 0};  // merged_local_dim;
+    //     //
+    //     //
+    //     FatalErrorInFunction << "Not implemented" << abort(FatalError);
+    //     // return std::make_tuple<std::shared_ptr<SparsityPattern>,
+    //     //                        std::shared_ptr<SparsityPattern>,
+    //     //                        std::vector<std::pair<bool, label>>>(
+    //     //     std::move(new_local_spars_pattern),
+    //     //     std::move(new_non_local_spars_pattern),
+    //     //     std::vector<std::pair<bool, label>>{});
+    // }
+    // // TODO Some old documentation see what is still needed
+    // // add interfaces
+    // // to add interfaces we go through the non_local_sparsity
+    // // pattern and check if interface is still non_local
+    // //
+    // // iterate interfaces and send to owner.
+    // // on owner decide to move to local or non_local
+    // // for sending to owner we can also use the all_to_all_v
+    // // approach we than have
+    // //
+    // // row [ 4 , 8,  12 | 16 , 17, 18 ] <- local row
+    // // col [ 1 , 2, 3 | 1, 2, 3 ] <- just the interface ctr
+    // // from repartition_comm_pattern we could get
+    // //
+    // // or we split local / non-local first
+    // // for this we could store ranges when computing
+    // // repartition_comm_pattern
+    // //
 }
 
-std::pair<SparsityPattern, std::vector<bool>>
-Repartitioner::build_non_local_interfaces(const ExecutorHandler &exec_handler,
-                                          SparsityPattern &loc,
-                                          const SparsityPattern &non_loc) const
+std::vector<bool> Repartitioner::build_non_local_interfaces(
+    const ExecutorHandler &exec_handler,
+    std::shared_ptr<
+        const gko::experimental::distributed::Partition<label, label>>
+        partition,
+    std::vector<label> &local_rows, std::vector<label> &local_cols,
+    std::vector<label> &local_mapping, std::vector<label> &local_ranks,
+    std::vector<gko::span> &local_spans, std::vector<label> &non_local_rows,
+    std::vector<label> &non_local_cols, std::vector<label> &non_local_mapping,
+    std::vector<label> &non_local_ranks,
+    std::vector<label> &non_local_rank_origin,
+    std::vector<gko::span> &non_local_spans) const
 {
-    LOG_1(verbose_, "start build non local interfaces")
-    std::vector<label> rows, cols, ldu_mapping, ranks, begins, ends;
+    // LOG_1(verbose_, "start build non local interfaces")
+    auto rank = exec_handler.get_rank();
     std::vector<bool> is_local;
-    label merged_ranks_size = non_loc.rank.size();
-
-    // TODO dont modify SparsityPattern but create vector
-    // and use the vector variant
-    //   for (label i = 0; i < merged_ranks_size; i++) {
-    //       // these are the begin ends before merging they need to be
-    //       // offsetted;
-    //       auto begin = non_loc.spans[i].begin;
-    //       auto end = non_loc.spans[i].end;
-    //       bool local = reparts_to_local(exec_handler, non_loc.rank[i]);
-    //       is_local.push_back(local);
-    //       if (local) {
-    //           // TODO depending on the interface simple non
-    //           // transforming and interpolating interfaces like now
-    //           // local processor interfaces could be merged
-    //           loc.begin.push_back(loc.row_idxs.size());
-    //           loc.row_idxs.insert(loc.row_idxs.end(),
-    //           non_loc.row_idxs.get_data() + begin,
-    //                           non_loc.row_idxs.get_data() + end);
-    //           loc.col_idxs.insert(loc.cols.end(), non_loc.col_idxs.get_data()
-    //           + begin,
-    //                           non_loc.col_idxs.get_data() + end);
-    //           loc.ldu_mapping.insert(loc.ldu_mapping.end(),
-    //                              non_loc.ldu_mapping.get_data() + begin,
-    //                              non_loc.ldu_mapping.get_data() + end);
-    //           // TODO store from rank ie from which rank it came
-    //           // this is currently unused
-    //           // TODO reimplement end doesnt exist
-    //           // loc.end.push_back(loc.row_idxs.size());
-    //           // also cannot pushback
-    //           // loc.rank.push_back(non_loc.rank[i]);
-    //       } else {
-    //           begins.push_back(row_idxs.size());
-    //           rows.insert(rows.end(), non_loc.row_idxs.begin() + begin,
-    //                       non_loc.row_idxs.begin() + end);
-    //           cols.insert(cols.end(), non_loc.col_idxs.begin() + begin,
-    //                       non_loc.col_idxs.begin() + end);
-    //           // NOTE we don't do anything with non local mapping
-    //           ldu_mapping.insert(ldu_mapping.end(), non_loc.col_idxs.begin()
-    //           + begin,
-    //                              non_loc.col_idxs.begin() + end);
-    //           ends.push_back(rows.size());
-    //           ranks.push_back(get_owner_rank(non_loc.rank[i]));
-    //       }
-    //   }
-    //   LOG_1(verbose_, "done build non local interfaces")
-    //   return std::make_pair(
-    //       SparsityPattern{rows, cols, ldu_mapping, begins, ends, ranks},
-    //       is_local);
+    std::vector<gko::span> mark_delete;
+
+    for (int i = 0; i < non_local_spans.size(); i++) {
+        auto [begin, end] = non_local_spans[i];
+        bool local = reparts_to_local(exec_handler, non_local_ranks[i]);
+        is_local.push_back(local);
+        if (local) {
+            label rows_start = local_rows.size();
+            local_rows.insert(local_rows.end(), non_local_rows.data() + begin,
+                              non_local_rows.data() + end);
+
+            std::vector<label> tmp_rank_local_cols(
+                non_local_cols.data() + begin, non_local_cols.data() + end);
+
+            detail::convert_to_local(partition, tmp_rank_local_cols, rank);
+
+            local_cols.insert(local_cols.end(), tmp_rank_local_cols.begin(),
+                              tmp_rank_local_cols.end());
+
+            local_mapping.insert(local_mapping.end(),
+                                 non_local_mapping.data() + begin,
+                                 non_local_mapping.data() + end);
+
+            local_spans.push_back(
+                gko::span{static_cast<gko::size_type>(rows_start),
+                          static_cast<gko::size_type>(end - begin)});
+            // TODO
+            // local_ranks.push_back(non_local_rank_origin[i]);
+        }
+    }
+
+    // remove data from non_local vectors
+    for (int i = mark_delete.size(); i > 0; i--) {
+        auto [begin, end] = non_local_spans[i];
+        non_local_rows.erase(non_local_rows.begin() + begin,
+                             non_local_rows.begin() + end);
+        non_local_cols.erase(non_local_cols.begin() + begin,
+                             non_local_cols.begin() + end);
+        non_local_mapping.erase(non_local_cols.begin() + begin,
+                                non_local_cols.begin() + end);
+        non_local_ranks.erase(non_local_ranks.begin() + begin);
+        // non_local_spans.erase(non_local_spans.begin() + begin, );
+    }
+    return is_local;
 }
 
 
@@ -458,12 +469,12 @@ Repartitioner::repartition_comm_pattern(
     std::vector<label> merged_target_ids{};
     std::vector<label> merged_target_sizes{};
     std::vector<std::vector<label>> merged_send_idxs{};
-    auto p = sort_permutation(gathered_target_ids,
-                              [](label a, label b) { return a < b; });
+    auto p = detail::sort_permutation(gathered_target_ids,
+                                      [](label a, label b) { return a < b; });
 
-    target_sizes = apply_permutation(gathered_target_sizes, p);
-    target_ids = apply_permutation(gathered_target_ids, p);
-    send_idxs = apply_permutation(send_idxs, p);
+    target_sizes = detail::apply_permutation(gathered_target_sizes, p);
+    target_ids = detail::apply_permutation(gathered_target_ids, p);
+    send_idxs = detail::apply_permutation(send_idxs, p);
 
     // Step 4.
     // Merge communication pattern with corresponding neighbours
diff --git a/src/common.C b/src/common.C
index d8f3ef85f..72ad6c546 100644
--- a/src/common.C
+++ b/src/common.C
@@ -145,6 +145,11 @@ label get_solve_prev_iters(word sys_matrix_name, const objectRegistry &db,
     return get_gko_solver_property(sys_matrix_name, iters_name, db, label(1));
 }
 
+std::shared_ptr<gko::array<label>> convert_to_array(const std::vector<label> & in) {
+    return std::make_shared<gko::array<label>>(gko::ReferenceExecutor::create(), in.begin(), in.end());
+}
+
+
 std::ostream &operator<<(std::ostream &os,
                          const std::shared_ptr<gko::matrix::Dense<scalar>> in)
 {
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index c10055fab..cb94361c5 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -3,6 +3,7 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 
 #include "OGL/Repartitioner.H"
+#include "OGL/common.H"
 
 #include "gtest/gtest.h"
 
@@ -100,6 +101,137 @@ TEST_P(RepartitionerFixture, has_correct_properties_for_n_rank)
         (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0));
 }
 
+TEST(Repartitioner, can_convert_to_global)
+{
+    // Arrange
+    auto exec = *((RepartitionerEnvironment *)global_env)->exec;
+    auto rank = exec.get_rank();
+    auto ref_exec = exec.get_ref_exec();
+    auto comm = exec.get_communicator();
+    auto partition = gko::share(
+        gko::experimental::distributed::build_partition_from_local_size<label,
+                                                                        label>(
+            ref_exec, *comm.get(), 4));
+
+    //         local ids          |   global ids
+    // cells: [ 2 3 | 2 3 ]       |  [ 10 11 | 14 15 ]
+    //   2    [ 0 1 | 0 1 ]  3    |  [  8  9 | 12 13 ]
+    //        -------------       |  -----------------
+    //        [ 2 3 | 2 3 ]       |  [  2  3 |  6  7 ]
+    //   0    [ 0 1 | 0 1 ]  1    |  [  0  1 |  4  5 ]
+    //
+
+    // vector
+    std::vector<std::vector<label>> idxs{
+        {0, 2, 0, 1},  // rank 0
+        {1, 3, 0, 1},  // rank 1
+        {2, 3, 0, 3},  // rank 2
+        {2, 3, 1, 3}   // rank 3
+    };
+    // vector
+    std::vector<std::vector<label>> ranks{
+        {1, 2},  // rank 0
+        {0, 3},  // rank 1
+        {0, 3},  // rank 2
+        {2, 1}   // rank 3
+    };
+    std::vector<std::vector<gko::span>> spans{
+        {gko::span{0, 2}, gko::span{2, 4}},  // rank 0
+        {gko::span{0, 2}, gko::span{2, 4}},  // rank 1
+        {gko::span{0, 2}, gko::span{2, 4}},  // rank 2
+        {gko::span{0, 2}, gko::span{2, 4}},  // rank 3
+    };
+
+    std::vector<std::vector<label>> exp_global_idxs{
+        {4, 6, 8, 9},    // rank 0
+        {1, 3, 12, 13},  // rank 1
+        {2, 3, 12, 14},  // rank 2
+        {9, 11, 6, 7}    // rank 3
+    };
+
+    // Act
+    auto global_idxs = detail::convert_to_global(partition, idxs[rank],
+                                                 spans[rank], ranks[rank]);
+
+    // Assert
+    EXPECT_EQ(global_idxs, exp_global_idxs[rank]);
+}
+
+// TEST_P(RepartitionerFixture, can_repartition_sparsity)
+// {
+//     using vec_vec = std::vector<std::vector<label>>;
+//     // Arrange
+//     label ranks_per_gpu = GetParam();
+//     auto exec = *((RepartitionerEnvironment *)global_env)->exec.get();
+//     auto rank = exec.get_rank();
+//     auto comm = exec.get_communicator();
+//     auto ref_exec = exec.get_ref_exec();
+//     auto partition = gko::share(
+//         gko::experimental::distributed::build_partition_from_local_size<label,
+//                                                                         label>(
+//             ref_exec, *comm.get(), 4));
+
+//     //         local ids          |   global ids
+//     // cells: [ 2 3 | 2 3 ]       |  [ 10 11 | 14 15 ]
+//     //   2    [ 0 1 | 0 1 ]  3    |  [  8  9 | 12 13 ]
+//     //        -------------       |  -----------------
+//     //        [ 2 3 | 2 3 ]       |  [  2  3 |  6  7 ]
+//     //   0    [ 0 1 | 0 1 ]  1    |  [  0  1 |  4  5 ]
+
+//     std::vector<label> local_rows{0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3};
+//     std::vector<label> local_cols{0, 1, 2, 0, 1, 3, 0, 2, 3, 1, 2, 3};
+//     std::vector<label> local_mapping{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11};
+//     std::vector<label> local_ranks{rank};
+//     std::vector<gko::span> local_spans{gko::span{0, local_rows.size()}};
+
+//     vec_vec non_local_rows{
+//         {1, 3, 2, 3},
+//         {0, 2, 2, 3},
+//         {0, 1, 1, 3},
+//         {0, 1, 0, 2},
+//     };
+//     vec_vec non_local_cols{
+//         {4, 6, 8, 9},    // rank 0
+//         {1, 3, 12, 13},  // rank 1
+//         {2, 3, 12, 14},  // rank 2
+//         {6, 7, 9, 11}    // rank 3
+//     };
+//     std::vector<label> non_local_mapping{0, 1, 2, 3};
+//     vec_vec non_local_ranks{
+//         {1, 2},  // rank 0
+//         {0, 3},  // rank 1
+//         {0, 3},  // rank 2
+//         {1, 2}   // rank 3
+//     };
+//     vec_vec non_local_ranks{
+//         {1, 2},  // rank 0
+//         {0, 3},  // rank 1
+//         {0, 3},  // rank 2
+//         {1, 2}   // rank 3
+//     };
+//     std::vector<gko::span> non_local_spans{
+//         {1, 2},  // rank 0
+//         {0, 3},  // rank 1
+//         {0, 3},  // rank 2
+//         {1, 2}   // rank 3
+//     };
+//     std::vector<gko::span> non_local_spans{
+//         {1, 2},  // rank 0
+//         {0, 3},  // rank 1
+//         {0, 3},  // rank 2
+//         {1, 2}   // rank 3
+//     };
+
+//     auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, exec);
+
+//     auto is_local = repartitioner.build_non_local_interfaces(
+//         exec, partition, local_rows, local_cols, local_mapping, local_ranks,
+//         local_spans, non_local_rows, non_local_cols, non_local_mapping,
+//         non_local_ranks, non_local_origin, non_local_spans);
+
+
+// }
+
 TEST_P(RepartitionerFixture, can_exchange_spans_and_ranks_for_n_ranks)
 {
     using vec_vec = std::vector<std::vector<label>>;
@@ -117,15 +249,36 @@ TEST_P(RepartitionerFixture, can_exchange_spans_and_ranks_for_n_ranks)
     exp_ranks.emplace(2, vec_vec{{1, 2, 2, 3}, {}, {3, 4, 4, 5}, {}});
     exp_ranks.emplace(4, vec_vec{{1, 2, 2, 3, 3, 4, 4, 5}, {}, {}, {}});
 
+    std::map<label, vec_vec> exp_spans_begin;
+    exp_spans_begin.emplace(1, vec_vec{{0, 5}, {0, 5}, {0, 5}, {0, 5}});
+    exp_spans_begin.emplace(2, vec_vec{{0, 5, 10, 15}, {}, {0, 5, 10, 15}, {}});
+    exp_spans_begin.emplace(
+        4, vec_vec{{0, 5, 10, 15, 20, 25, 30, 35}, {}, {}, {}});
+    std::map<label, vec_vec> exp_spans_end;
+    exp_spans_end.emplace(1, vec_vec{{5, 10}, {5, 10}, {5, 10}, {5, 10}});
+    exp_spans_end.emplace(2, vec_vec{{5, 10, 15, 20}, {}, {5, 10, 15, 15}, {}});
+    exp_spans_end.emplace(4,
+                          vec_vec{{5, 10, 15, 20, 25, 30, 35, 40}, {}, {}, {}});
+
     // Act
     auto [new_spans, new_ranks] =
         detail::exchange_span_ranks(exec, ranks_per_gpu, spans, ranks);
 
+    std::vector<label> res_spans_begin{};
+    std::vector<label> res_spans_end{};
+
+    for (auto span : new_spans) {
+        res_spans_begin.push_back(span.begin);
+        res_spans_end.push_back(span.end);
+    }
+
     // Assert
     ASSERT_EQ(new_ranks, exp_ranks[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_spans_begin, exp_spans_begin[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_spans_end, exp_spans_end[ranks_per_gpu][rank]);
 }
 
-TEST_P(RepartitionerFixture, can_repartition_sparsity_pattern_for_n_ranks)
+TEST_P(RepartitionerFixture, can_repartition_sparsity_pattern_1D_for_n_ranks)
 {
     // Arrange
     label ranks_per_gpu = GetParam();
@@ -142,27 +295,29 @@ TEST_P(RepartitionerFixture, can_repartition_sparsity_pattern_for_n_ranks)
     std::vector<label> rows{0, 0, 1, 1};
     std::vector<label> cols{0, 1, 0, 1};
     std::vector<label> mapping{0, 1, 2, 3};
+    std::vector<gko::span> spans{gko::span {0, 4}};
     std::vector<label> ranks{rank};
     auto local_sparsity = std::make_shared<SparsityPattern>(
-        ref_exec, gko::dim<2>{2, 2}, rows, cols, mapping, ranks);
+        ref_exec, gko::dim<2>{2, 2}, rows, cols, mapping, spans, ranks);
 
     std::vector<label> non_local_rows{1};
     std::vector<label> non_local_cols{0};
     std::vector<label> non_local_mapping{0};
+    std::vector<gko::span> non_local_spans{gko::span {0, 1}};
     std::vector<label> non_local_ranks{rank};
     auto non_local_sparsity = std::make_shared<SparsityPattern>(
         ref_exec, gko::dim<2>{2, 1}, non_local_rows, non_local_cols,
-        non_local_mapping, non_local_ranks);
+        non_local_mapping, non_local_spans, non_local_ranks);
 
     std::map<label, std::vector<label>> exp_local_nnz;
     exp_local_nnz.emplace(1, std::vector<label>{4, 4, 4, 4});
-    exp_local_nnz.emplace(2, std::vector<label>{8, 0, 8, 0});
-    exp_local_nnz.emplace(4, std::vector<label>{16, 0, 0, 0});
+    exp_local_nnz.emplace(2, std::vector<label>{10, 0, 10, 0});
+    exp_local_nnz.emplace(4, std::vector<label>{24, 0, 0, 0});
 
     std::map<label, std::vector<label>> exp_non_local_nnz;
-    exp_local_nnz.emplace(1, std::vector<label>{1, 2, 2, 1});
-    exp_local_nnz.emplace(2, std::vector<label>{1, 0, 1, 0});
-    exp_local_nnz.emplace(4, std::vector<label>{0, 0, 0, 0});
+    exp_non_local_nnz.emplace(1, std::vector<label>{1, 2, 2, 1});
+    exp_non_local_nnz.emplace(2, std::vector<label>{1, 0, 1, 0});
+    exp_non_local_nnz.emplace(4, std::vector<label>{0, 0, 0, 0});
 
     // Act
     auto [repart_local, repart_non_local, tracking] =
@@ -171,8 +326,7 @@ TEST_P(RepartitionerFixture, can_repartition_sparsity_pattern_for_n_ranks)
 
     // Assert
     ASSERT_EQ(repart_local->num_nnz, exp_local_nnz[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_non_local->num_nnz,
-              exp_non_local_nnz[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_non_local->num_nnz, exp_non_local_nnz[ranks_per_gpu][rank]);
 }
 
 TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)

From eba1ce2d50e59fc7eae5942f67028cb8da1de4b4 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Tue, 27 Aug 2024 14:26:39 +0200
Subject: [PATCH 47/67] WIP repartition sparsisty

---
 .../OGL/DevicePersistent/ExecutorHandler.H    |   9 +-
 include/OGL/MatrixWrapper/SparsityPattern.H   |   8 +-
 include/OGL/Repartitioner.H                   |   7 +-
 include/OGL/common.H                          |   8 +-
 src/Repartitioner.C                           |  88 +++++++----
 unitTests/Repartitioner.C                     | 141 ++++++------------
 6 files changed, 125 insertions(+), 136 deletions(-)

diff --git a/include/OGL/DevicePersistent/ExecutorHandler.H b/include/OGL/DevicePersistent/ExecutorHandler.H
index 390031a32..7bfa141ca 100644
--- a/include/OGL/DevicePersistent/ExecutorHandler.H
+++ b/include/OGL/DevicePersistent/ExecutorHandler.H
@@ -176,22 +176,21 @@ public:
             MPI_COMM_WORLD, gko_force_host_buffer_);
     }
 
-    std::shared_ptr<gko::experimental::mpi::communicator>
+    std::shared_ptr<const gko::experimental::mpi::communicator>
     get_gko_mpi_device_comm() const
     {
         return this->device_comm_;
     }
 
-    std::shared_ptr<gko::experimental::mpi::communicator> get_communicator()
-        const
+    std::shared_ptr<const gko::experimental::mpi::communicator>
+    get_communicator() const
     {
         return this->device_comm_;
     }
 
-    label get_rank() const { return get_communicator().get()->rank(); };
+    label get_rank() const { return get_communicator()->rank(); };
 };
 
 using PersistentExecutor = ExecutorHandler;
 
 }  // namespace Foam
-// namespace Foam
diff --git a/include/OGL/MatrixWrapper/SparsityPattern.H b/include/OGL/MatrixWrapper/SparsityPattern.H
index 8f3645edd..dd7789cc8 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.H
+++ b/include/OGL/MatrixWrapper/SparsityPattern.H
@@ -53,7 +53,13 @@ struct SparsityPattern {
           dim(dim_),
           spans(spans_),
           rank(rank_)
-    {}
+    {
+        // For every rank there should be a span
+        ASSERT_EQ(spans.size(), rank.size());
+
+        ASSERT_EQ(rows.size(), cols.size());
+        ASSERT_EQ(rows.size(), mapping.size());
+    }
 
     SparsityPattern(std::shared_ptr<const gko::Executor> exec)
         : num_nnz(0),
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index 40ef76ffc..d14747e98 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -65,9 +65,10 @@ void convert_to_local(
 **
 ** returns a pair vector spans
 */
-std::pair<std::vector<gko::span>, std::vector<label>> exchange_span_ranks(
-    const ExecutorHandler &exec_handler, label ranks_per_gpu,
-    const std::vector<gko::span> &spans, const std::vector<label> &ranks);
+std::tuple<std::vector<gko::span>, std::vector<label>, std::vector<label>>
+exchange_span_ranks(const ExecutorHandler &exec_handler, label ranks_per_gpu,
+                    const std::vector<gko::span> &spans,
+                    const std::vector<label> &ranks);
 
 template <typename T>
 std::vector<T> apply_permutation(const std::vector<T> vec,
diff --git a/include/OGL/common.H b/include/OGL/common.H
index 4175f794f..cbe99e83c 100644
--- a/include/OGL/common.H
+++ b/include/OGL/common.H
@@ -90,10 +90,10 @@ namespace Foam {
 
 #define UNUSED(x) (void)(x)
 
-#define ASSERT_EQ(_val1, _val2)                                                \
-    if (_val1 != _val2) {                                                      \
-        throw ::gko::ValueMismatch(__FILE__, __LINE__, __func__, _val1, _val2, \
-                                   "expected equal values");                   \
+#define ASSERT_EQ(_val1, _val2)                                               \
+    if (_val1 != _val2) {                                                     \
+        FatalErrorInFunction << "Expected equal values: " << _val1 << " and " \
+                             << _val2 << abort(FatalError);                   \
     }
 
 std::ostream &operator<<(
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index c7613a5f4..a8d6dffcd 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -6,7 +6,7 @@
 
 namespace detail {
 
-std::pair<std::vector<gko::span>, std::vector<label>> exchange_span_ranks(
+std::tuple<std::vector<gko::span>, std::vector<label>, std::vector<label>> exchange_span_ranks(
     const ExecutorHandler &exec_handler, label ranks_per_gpu,
     const std::vector<gko::span> &spans, const std::vector<label> &ranks)
 {
@@ -32,7 +32,18 @@ std::pair<std::vector<gko::span>, std::vector<label>> exchange_span_ranks(
     auto new_ranks = gather_labels_to_owner(exec_handler, comm_pattern,
                                             ranks.data(), ranks.size());
 
-    return {out_spans, new_ranks};
+    std::vector<label> origins{};
+    if (gathered_size.size() > 0) {
+    auto rank = exec_handler.get_rank();
+    for (int i = 0; i<ranks_per_gpu; i++) {
+        auto j = comm_pattern.recv_counts[i+rank];
+        for (int k=0;k<j;k++) {
+            origins.push_back(i+rank);
+        }
+    }
+    }
+
+    return {out_spans, new_ranks, origins};
 }
 
 template <typename T>
@@ -145,7 +156,7 @@ Repartitioner::repartition_sparsity(
     auto tmp_non_local_rows = gather_closure(
         non_local_comm_pattern, src_non_local_pattern->row_idxs, offset);
 
-    auto [new_spans, new_ranks] = detail::exchange_span_ranks(
+    auto [new_spans, new_ranks, tmp_non_local_origin] = detail::exchange_span_ranks(
         exec_handler, ranks_per_gpu_, src_non_local_pattern->spans,
         src_non_local_pattern->rank);
 
@@ -158,16 +169,15 @@ Repartitioner::repartition_sparsity(
     std::vector<label> tmp_non_local_mapping =
         std::vector<label>(tmp_non_local_rows.size());
 
-    std::vector<label> tmp_non_local_origin{};
-
     auto is_local = build_non_local_interfaces(
         exec_handler, orig_partition_, tmp_local_rows, tmp_local_cols,
         tmp_local_mapping, tmp_local_rank, tmp_local_span, tmp_non_local_rows,
         tmp_non_local_cols, tmp_non_local_mapping, new_ranks,
         tmp_non_local_origin, new_spans);
 
+
     // TODO FIXME
-    gko::dim<2> tmp_non_local_dim {0, 0};
+    gko::dim<2> tmp_non_local_dim{0, 0};
 
     if (is_owner(exec_handler)) {
         auto local_pattern = std::make_shared<SparsityPattern>(
@@ -296,52 +306,80 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
     std::vector<label> &non_local_rank_origin,
     std::vector<gko::span> &non_local_spans) const
 {
-    // LOG_1(verbose_, "start build non local interfaces")
     auto rank = exec_handler.get_rank();
     std::vector<bool> is_local;
-    std::vector<gko::span> mark_delete;
+    std::vector<label> mark_keep;
 
     for (int i = 0; i < non_local_spans.size(); i++) {
         auto [begin, end] = non_local_spans[i];
         bool local = reparts_to_local(exec_handler, non_local_ranks[i]);
         is_local.push_back(local);
+
         if (local) {
-            label rows_start = local_rows.size();
             local_rows.insert(local_rows.end(), non_local_rows.data() + begin,
                               non_local_rows.data() + end);
 
             std::vector<label> tmp_rank_local_cols(
-                non_local_cols.data() + begin, non_local_cols.data() + end);
+                non_local_cols.data() + begin,
+                non_local_cols.data() + end);
 
             detail::convert_to_local(partition, tmp_rank_local_cols, rank);
 
             local_cols.insert(local_cols.end(), tmp_rank_local_cols.begin(),
                               tmp_rank_local_cols.end());
-
             local_mapping.insert(local_mapping.end(),
                                  non_local_mapping.data() + begin,
                                  non_local_mapping.data() + end);
 
-            local_spans.push_back(
-                gko::span{static_cast<gko::size_type>(rows_start),
-                          static_cast<gko::size_type>(end - begin)});
-            // TODO
-            // local_ranks.push_back(non_local_rank_origin[i]);
+            gko::size_type rows_start = local_rows.size();
+            local_spans.emplace_back(
+                rows_start,
+                          static_cast<gko::size_type>(end - begin));
+
+            local_ranks.push_back(non_local_rank_origin[i]);
+        } else {
+            mark_keep.push_back(i);
         }
     }
 
     // remove data from non_local vectors
-    for (int i = mark_delete.size(); i > 0; i--) {
+    if (mark_keep.size() == 0) {
+        std::cout << "rank " << rank << " \n";
+    non_local_rows.clear();
+    non_local_cols.clear();
+    non_local_mapping.clear();
+    non_local_ranks.clear();
+    non_local_spans.clear();
+        } else {
+    std::vector<label> copy_rows, copy_cols, copy_mapping, copy_ranks;
+    std::vector<gko::span> copy_spans;
+    for (label i : mark_keep) {
         auto [begin, end] = non_local_spans[i];
-        non_local_rows.erase(non_local_rows.begin() + begin,
-                             non_local_rows.begin() + end);
-        non_local_cols.erase(non_local_cols.begin() + begin,
-                             non_local_cols.begin() + end);
-        non_local_mapping.erase(non_local_cols.begin() + begin,
-                                non_local_cols.begin() + end);
-        non_local_ranks.erase(non_local_ranks.begin() + begin);
-        // non_local_spans.erase(non_local_spans.begin() + begin, );
+        std::cout << "rank " << rank << " begin " << begin << " end " << end << "\n";
+        copy_rows.insert(copy_rows.end(),
+                         non_local_rows.data() + begin,
+                         non_local_rows.data() + end);
+        copy_cols.insert(copy_cols.end(),
+                         non_local_cols.data() + begin,
+                         non_local_cols.data() + end);
+        copy_mapping.insert(copy_mapping.end(),
+                            non_local_mapping.data() + begin,
+                            non_local_mapping.data() + end);
+        copy_ranks.push_back(non_local_ranks[i]);
+        copy_spans.emplace_back(non_local_spans[i].begin,
+                                non_local_spans[i].end);
     }
+
+    non_local_rows = copy_rows;
+    non_local_cols = copy_cols;
+    non_local_mapping = copy_mapping;
+    non_local_ranks = copy_ranks;
+    non_local_spans.clear();
+    for (auto &span : copy_spans) {
+        non_local_spans.emplace_back(span.begin, span.end);
+    }
+    }
+
     return is_local;
 }
 
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index cb94361c5..03ec1a8ac 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -23,6 +23,9 @@ extern char **my_argv;
 int my_argc;
 char **my_argv;
 
+using vec_vec = std::vector<std::vector<label>>;
+using vec_vec_vec = std::vector<vec_vec>;
+
 class RepartitionerEnvironment : public testing::Environment {
 public:
     void SetUp()
@@ -78,7 +81,6 @@ TEST_P(RepartitionerFixture, can_create_repartitioner)
 
     // Assert
     EXPECT_EQ(repartitioner.get_ranks_per_gpu(), ranks_per_gpu);
-<<<<<<< HEAD
 }
 
 
@@ -104,7 +106,7 @@ TEST_P(RepartitionerFixture, has_correct_properties_for_n_rank)
 TEST(Repartitioner, can_convert_to_global)
 {
     // Arrange
-    auto exec = *((RepartitionerEnvironment *)global_env)->exec;
+    auto& exec = *((RepartitionerEnvironment *)global_env)->exec;
     auto rank = exec.get_rank();
     auto ref_exec = exec.get_ref_exec();
     auto comm = exec.get_communicator();
@@ -125,7 +127,7 @@ TEST(Repartitioner, can_convert_to_global)
     std::vector<std::vector<label>> idxs{
         {0, 2, 0, 1},  // rank 0
         {1, 3, 0, 1},  // rank 1
-        {2, 3, 0, 3},  // rank 2
+        {2, 3, 0, 2},  // rank 2
         {2, 3, 1, 3}   // rank 3
     };
     // vector
@@ -133,7 +135,7 @@ TEST(Repartitioner, can_convert_to_global)
         {1, 2},  // rank 0
         {0, 3},  // rank 1
         {0, 3},  // rank 2
-        {2, 1}   // rank 3
+        {1, 2}   // rank 3
     };
     std::vector<std::vector<gko::span>> spans{
         {gko::span{0, 2}, gko::span{2, 4}},  // rank 0
@@ -146,7 +148,7 @@ TEST(Repartitioner, can_convert_to_global)
         {4, 6, 8, 9},    // rank 0
         {1, 3, 12, 13},  // rank 1
         {2, 3, 12, 14},  // rank 2
-        {9, 11, 6, 7}    // rank 3
+        {6, 7, 9, 11}    // rank 3
     };
 
     // Act
@@ -157,84 +159,9 @@ TEST(Repartitioner, can_convert_to_global)
     EXPECT_EQ(global_idxs, exp_global_idxs[rank]);
 }
 
-// TEST_P(RepartitionerFixture, can_repartition_sparsity)
-// {
-//     using vec_vec = std::vector<std::vector<label>>;
-//     // Arrange
-//     label ranks_per_gpu = GetParam();
-//     auto exec = *((RepartitionerEnvironment *)global_env)->exec.get();
-//     auto rank = exec.get_rank();
-//     auto comm = exec.get_communicator();
-//     auto ref_exec = exec.get_ref_exec();
-//     auto partition = gko::share(
-//         gko::experimental::distributed::build_partition_from_local_size<label,
-//                                                                         label>(
-//             ref_exec, *comm.get(), 4));
-
-//     //         local ids          |   global ids
-//     // cells: [ 2 3 | 2 3 ]       |  [ 10 11 | 14 15 ]
-//     //   2    [ 0 1 | 0 1 ]  3    |  [  8  9 | 12 13 ]
-//     //        -------------       |  -----------------
-//     //        [ 2 3 | 2 3 ]       |  [  2  3 |  6  7 ]
-//     //   0    [ 0 1 | 0 1 ]  1    |  [  0  1 |  4  5 ]
-
-//     std::vector<label> local_rows{0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3};
-//     std::vector<label> local_cols{0, 1, 2, 0, 1, 3, 0, 2, 3, 1, 2, 3};
-//     std::vector<label> local_mapping{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11};
-//     std::vector<label> local_ranks{rank};
-//     std::vector<gko::span> local_spans{gko::span{0, local_rows.size()}};
-
-//     vec_vec non_local_rows{
-//         {1, 3, 2, 3},
-//         {0, 2, 2, 3},
-//         {0, 1, 1, 3},
-//         {0, 1, 0, 2},
-//     };
-//     vec_vec non_local_cols{
-//         {4, 6, 8, 9},    // rank 0
-//         {1, 3, 12, 13},  // rank 1
-//         {2, 3, 12, 14},  // rank 2
-//         {6, 7, 9, 11}    // rank 3
-//     };
-//     std::vector<label> non_local_mapping{0, 1, 2, 3};
-//     vec_vec non_local_ranks{
-//         {1, 2},  // rank 0
-//         {0, 3},  // rank 1
-//         {0, 3},  // rank 2
-//         {1, 2}   // rank 3
-//     };
-//     vec_vec non_local_ranks{
-//         {1, 2},  // rank 0
-//         {0, 3},  // rank 1
-//         {0, 3},  // rank 2
-//         {1, 2}   // rank 3
-//     };
-//     std::vector<gko::span> non_local_spans{
-//         {1, 2},  // rank 0
-//         {0, 3},  // rank 1
-//         {0, 3},  // rank 2
-//         {1, 2}   // rank 3
-//     };
-//     std::vector<gko::span> non_local_spans{
-//         {1, 2},  // rank 0
-//         {0, 3},  // rank 1
-//         {0, 3},  // rank 2
-//         {1, 2}   // rank 3
-//     };
-
-//     auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, exec);
-
-//     auto is_local = repartitioner.build_non_local_interfaces(
-//         exec, partition, local_rows, local_cols, local_mapping, local_ranks,
-//         local_spans, non_local_rows, non_local_cols, non_local_mapping,
-//         non_local_ranks, non_local_origin, non_local_spans);
-
-
-// }
 
 TEST_P(RepartitionerFixture, can_exchange_spans_and_ranks_for_n_ranks)
 {
-    using vec_vec = std::vector<std::vector<label>>;
     // Arrange
     label ranks_per_gpu = GetParam();
     auto exec = *((RepartitionerEnvironment *)global_env)->exec.get();
@@ -249,6 +176,8 @@ TEST_P(RepartitionerFixture, can_exchange_spans_and_ranks_for_n_ranks)
     exp_ranks.emplace(2, vec_vec{{1, 2, 2, 3}, {}, {3, 4, 4, 5}, {}});
     exp_ranks.emplace(4, vec_vec{{1, 2, 2, 3, 3, 4, 4, 5}, {}, {}, {}});
 
+    // TODO test origins
+
     std::map<label, vec_vec> exp_spans_begin;
     exp_spans_begin.emplace(1, vec_vec{{0, 5}, {0, 5}, {0, 5}, {0, 5}});
     exp_spans_begin.emplace(2, vec_vec{{0, 5, 10, 15}, {}, {0, 5, 10, 15}, {}});
@@ -256,12 +185,11 @@ TEST_P(RepartitionerFixture, can_exchange_spans_and_ranks_for_n_ranks)
         4, vec_vec{{0, 5, 10, 15, 20, 25, 30, 35}, {}, {}, {}});
     std::map<label, vec_vec> exp_spans_end;
     exp_spans_end.emplace(1, vec_vec{{5, 10}, {5, 10}, {5, 10}, {5, 10}});
-    exp_spans_end.emplace(2, vec_vec{{5, 10, 15, 20}, {}, {5, 10, 15, 15}, {}});
-    exp_spans_end.emplace(4,
-                          vec_vec{{5, 10, 15, 20, 25, 30, 35, 40}, {}, {}, {}});
+    exp_spans_end.emplace(2, vec_vec{{5, 10, 15, 20}, {}, {5, 10, 15, 20}, {}});
+    exp_spans_end.emplace(4, vec_vec{{5, 10, 15, 20, 25, 30, 35, 40}, {}, {}, {}});
 
     // Act
-    auto [new_spans, new_ranks] =
+    auto [new_spans, new_ranks, origins] =
         detail::exchange_span_ranks(exec, ranks_per_gpu, spans, ranks);
 
     std::vector<label> res_spans_begin{};
@@ -272,6 +200,8 @@ TEST_P(RepartitionerFixture, can_exchange_spans_and_ranks_for_n_ranks)
         res_spans_end.push_back(span.end);
     }
 
+
+
     // Assert
     ASSERT_EQ(new_ranks, exp_ranks[ranks_per_gpu][rank]);
     ASSERT_EQ(res_spans_begin, exp_spans_begin[ranks_per_gpu][rank]);
@@ -295,24 +225,42 @@ TEST_P(RepartitionerFixture, can_repartition_sparsity_pattern_1D_for_n_ranks)
     std::vector<label> rows{0, 0, 1, 1};
     std::vector<label> cols{0, 1, 0, 1};
     std::vector<label> mapping{0, 1, 2, 3};
-    std::vector<gko::span> spans{gko::span {0, 4}};
+    std::vector<gko::span> spans{gko::span{0, 2}};
     std::vector<label> ranks{rank};
     auto local_sparsity = std::make_shared<SparsityPattern>(
         ref_exec, gko::dim<2>{2, 2}, rows, cols, mapping, spans, ranks);
 
-    std::vector<label> non_local_rows{1};
-    std::vector<label> non_local_cols{0};
-    std::vector<label> non_local_mapping{0};
-    std::vector<gko::span> non_local_spans{gko::span {0, 1}};
-    std::vector<label> non_local_ranks{rank};
+    // rank local rows
+    vec_vec non_local_rows{{1}, {0, 1}, {0, 1}, {0}};
+
+    // other side local rows
+    vec_vec non_local_cols{{0}, {1, 0}, {1, 0}, {1}};
+
+    vec_vec non_local_mapping{{0}, {0, 1}, {0, 1}, {0}};
+
+    std::vector<std::vector<gko::span>> non_local_spans{
+        {gko::span{0, 1}},                   // rank 0
+        {gko::span{0, 1}, gko::span{1, 2}},  // rank 1
+        {gko::span{0, 1}, gko::span{1, 2}},  // rank 2
+        {gko::span{0, 1}},                   // rank 3
+    };
+
+    vec_vec non_local_ranks{
+        {1},     // rank 0
+        {0, 2},  // rank 1
+        {1, 3},  // rank 2
+        {2},     // rank 3
+    };
+
     auto non_local_sparsity = std::make_shared<SparsityPattern>(
-        ref_exec, gko::dim<2>{2, 1}, non_local_rows, non_local_cols,
-        non_local_mapping, non_local_spans, non_local_ranks);
+        ref_exec, gko::dim<2>{2, non_local_ranks[rank].size()}, non_local_rows[rank],
+        non_local_cols[rank], non_local_mapping[rank], non_local_spans[rank],
+        non_local_ranks[rank]);
 
     std::map<label, std::vector<label>> exp_local_nnz;
     exp_local_nnz.emplace(1, std::vector<label>{4, 4, 4, 4});
     exp_local_nnz.emplace(2, std::vector<label>{10, 0, 10, 0});
-    exp_local_nnz.emplace(4, std::vector<label>{24, 0, 0, 0});
+    exp_local_nnz.emplace(4, std::vector<label>{22, 0, 0, 0});
 
     std::map<label, std::vector<label>> exp_non_local_nnz;
     exp_non_local_nnz.emplace(1, std::vector<label>{1, 2, 2, 1});
@@ -326,7 +274,8 @@ TEST_P(RepartitionerFixture, can_repartition_sparsity_pattern_1D_for_n_ranks)
 
     // Assert
     ASSERT_EQ(repart_local->num_nnz, exp_local_nnz[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_non_local->num_nnz, exp_non_local_nnz[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_non_local->num_nnz,
+              exp_non_local_nnz[ranks_per_gpu][rank]);
 }
 
 TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
@@ -344,7 +293,6 @@ TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
     std::vector<std::vector<label>> ids{{1}, {0, 2}, {1, 3}, {2}};
 
     // expected communcation ranks
-    using vec_vec = std::vector<std::vector<label>>;
     std::map<label, vec_vec> exp_res_ids{};
 
     exp_res_ids[1] = ids;  // in the ranks_per_gpu==1 case nothing changes
@@ -360,7 +308,6 @@ TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
     exp_res_sizes.emplace(4, vec_vec({{}, {}, {}, {}}));
 
     // expected rows
-    using vec_vec_vec = std::vector<vec_vec>;
     vec_vec_vec rows{{{1}}, {{0}, {1}}, {{0}, {1}}, {{0}}};
     std::map<label, vec_vec> exp_res_rows;
     exp_res_rows.emplace(1, vec_vec{{1}, {0, 1}, {0, 1}, {0}});
@@ -424,7 +371,6 @@ TEST_P(RepartitionerFixture, can_repartition_2D_comm_pattern_for_n_ranks)
     std::vector<std::vector<label>> ids{{1, 2}, {0, 3}, {0, 3}, {1, 2}};
 
     // expected communcation ranks
-    using vec_vec = std::vector<std::vector<label>>;
     std::map<label, vec_vec> exp_res_ids{};
 
     exp_res_ids[1] = ids;  // in the ranks_per_gpu==1 case nothing changes
@@ -440,7 +386,6 @@ TEST_P(RepartitionerFixture, can_repartition_2D_comm_pattern_for_n_ranks)
     exp_res_sizes.emplace(4, vec_vec({{}, {}, {}, {}}));
 
     // expected rows
-    using vec_vec_vec = std::vector<vec_vec>;
     vec_vec_vec rows{
         {{1, 3}, {2, 3}}, {{0, 2}, {2, 3}}, {{0, 1}, {1, 3}}, {{0, 1}, {0, 2}}};
     std::map<label, vec_vec> exp_res_rows;

From 1235ae54aacf8d91bf6d7de29c6cd8d85179b72a Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 28 Aug 2024 10:53:58 +0200
Subject: [PATCH 48/67] refactor unit test, add WIP 1d sparsity pattern test

---
 include/OGL/Repartitioner.H |   4 +-
 include/OGL/common.H        |   2 +
 src/Repartitioner.C         |   4 -
 src/common.C                |   9 ++
 unitTests/Repartitioner.C   | 311 +++++++++++++++++-------------------
 5 files changed, 158 insertions(+), 172 deletions(-)

diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index d14747e98..e6249f404 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -89,9 +89,9 @@ class Repartitioner {
 private:
     using partition = gko::experimental::distributed::Partition<label, label>;
 
-    const label size_;  //! Size (n matrix rows) before repartitioning
+    const label size_;  //! Size (matrix rows) before repartitioning
 
-    const label repart_size_;  //! Size after repartitioning
+    const label repart_size_;  //! Size (matrix rows) after repartitioning
 
     const label ranks_per_gpu_;
 
diff --git a/include/OGL/common.H b/include/OGL/common.H
index cbe99e83c..237807777 100644
--- a/include/OGL/common.H
+++ b/include/OGL/common.H
@@ -134,4 +134,6 @@ label get_next_caching(word sys_matrix_name, const objectRegistry &db);
 
 std::shared_ptr<gko::array<label>> convert_to_array(
     const std::vector<label> &in);
+
+std::vector<label> convert_to_vector(const gko::array<label> &in);
 }  // namespace Foam
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index a8d6dffcd..84f04d6dd 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -135,7 +135,6 @@ Repartitioner::repartition_sparsity(
     auto tmp_local_mapping =
         gather_closure(local_comm_pattern, src_local_pattern->ldu_mapping, 0);
 
-
     if (is_owner(exec_handler)) {
         make_ldu_mapping_consecutive(local_comm_pattern, tmp_local_mapping,
                                      rank, ranks_per_gpu);
@@ -175,7 +174,6 @@ Repartitioner::repartition_sparsity(
         tmp_non_local_cols, tmp_non_local_mapping, new_ranks,
         tmp_non_local_origin, new_spans);
 
-
     // TODO FIXME
     gko::dim<2> tmp_non_local_dim{0, 0};
 
@@ -344,7 +342,6 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
 
     // remove data from non_local vectors
     if (mark_keep.size() == 0) {
-        std::cout << "rank " << rank << " \n";
     non_local_rows.clear();
     non_local_cols.clear();
     non_local_mapping.clear();
@@ -355,7 +352,6 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
     std::vector<gko::span> copy_spans;
     for (label i : mark_keep) {
         auto [begin, end] = non_local_spans[i];
-        std::cout << "rank " << rank << " begin " << begin << " end " << end << "\n";
         copy_rows.insert(copy_rows.end(),
                          non_local_rows.data() + begin,
                          non_local_rows.data() + end);
diff --git a/src/common.C b/src/common.C
index 72ad6c546..25de64bb2 100644
--- a/src/common.C
+++ b/src/common.C
@@ -150,6 +150,15 @@ std::shared_ptr<gko::array<label>> convert_to_array(const std::vector<label> & i
 }
 
 
+std::vector<label> convert_to_vector(
+ const gko::array<label>& in
+) {
+    return std::vector<label> (
+        in.get_const_data(),
+        in.get_const_data() + in.get_size());
+}
+
+
 std::ostream &operator<<(std::ostream &os,
                          const std::shared_ptr<gko::matrix::Dense<scalar>> in)
 {
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 03ec1a8ac..960ffe128 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -38,8 +38,7 @@ public:
         exec = std::make_shared<ExecutorHandler>(time->thisDb(), dict, "dummy",
                                                  true);
 
-        auto comm = exec->get_gko_mpi_host_comm();
-        if (comm->size() < 2) {
+        if (exec->get_communicator()->size() < 2) {
             std::cout << "At least 2 CPU processes should be used!"
                       << std::endl;
             std::abort();
@@ -48,7 +47,7 @@ public:
         // to clean up output
         ::testing::TestEventListeners &listeners =
             ::testing::UnitTest::GetInstance()->listeners();
-        if (comm->rank() != 0) {
+        if (exec->get_communicator()->rank() != 0) {
             delete listeners.Release(listeners.default_result_printer());
         }
     }
@@ -62,89 +61,135 @@ public:
 const testing::Environment *global_env =
     AddGlobalTestEnvironment(new RepartitionerEnvironment);
 
-class RepartitionerFixture : public testing::TestWithParam<int> {
+class RepartitionerFixture : public testing::Test {
 public:
-    label local_size = 10;
+    ExecutorHandler exec =
+        *((RepartitionerEnvironment *)global_env)->exec.get();
+    label rank = exec.get_rank();
+    const gko::experimental::mpi::communicator comm =
+        *(exec.get_communicator().get());
 };
 
-INSTANTIATE_TEST_SUITE_P(RepartitionerFixtureInstantiation,
-                         RepartitionerFixture, testing::Values(1, 2, 4));
+/* @brief Test fixture class for 1D mesh
+ *
+ * The mesh has the following structure
+ * ranks:    0     1     2     3
+ * cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
+ * cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- local row ids
+ */
+class RepartitionerFixture1D : public RepartitionerFixture,
+                               public testing::WithParamInterface<int> {
+public:
+    label local_size = 2;  // matrix rows, same for all ranks
+
+    // local data
+    std::vector<label> rows{0, 0, 1, 1};
+    std::vector<label> cols{0, 1, 0, 1};
+    std::vector<label> mapping{0, 1, 2, 3};
+    std::vector<gko::span> spans{gko::span{0, 2}};
+
+    // non local data
+    vec_vec non_local_rows{{1}, {0, 1}, {0, 1}, {0}};
+    vec_vec non_local_cols{{0}, {1, 0}, {1, 0}, {1}};
+    vec_vec non_local_mapping{{0}, {0, 1}, {0, 1}, {0}};
+    // communication partners (ranks)
+    vec_vec ids{{1}, {0, 2}, {1, 3}, {2}};
+
+    std::vector<std::vector<gko::span>> non_local_spans{
+        {gko::span{0, 1}},                   // rank 0
+        {gko::span{0, 1}, gko::span{1, 2}},  // rank 1
+        {gko::span{0, 1}, gko::span{1, 2}},  // rank 2
+        {gko::span{0, 1}},                   // rank 3
+    };
+
+    vec_vec non_local_ranks{
+        {1},     // rank 0
+        {0, 2},  // rank 1
+        {1, 3},  // rank 2
+        {2},     // rank 3
+    };
+};
+
+/* @brief Test fixture class for 1D mesh
+ *
+ * The mesh has the following structure
+ *         local ids          |   global ids
+ * cells: [ 2 3 | 2 3 ]       |  [ 10 11 | 14 15 ]
+ *   2    [ 0 1 | 0 1 ]  3    |  [  8  9 | 12 13 ]
+ *        -------------       |  -----------------
+ *        [ 2 3 | 2 3 ]       |  [  2  3 |  6  7 ]
+ *   0    [ 0 1 | 0 1 ]  1    |  [  0  1 |  4  5 ]
+ *
+ */
+class RepartitionerFixture2D : public RepartitionerFixture,
+                               public testing::WithParamInterface<int> {
+public:
+    label local_size = 4;
+
+    vec_vec ids{{1, 2}, {0, 3}, {0, 3}, {1, 2}};
 
-TEST_P(RepartitionerFixture, can_create_repartitioner)
+    vec_vec idxs{
+        {0, 2, 0, 1},  // rank 0
+        {1, 3, 0, 1},  // rank 1
+        {2, 3, 0, 2},  // rank 2
+        {2, 3, 1, 3}   // rank 3
+    };
+    // vector
+    vec_vec ranks{
+        {1, 2},  // rank 0
+        {0, 3},  // rank 1
+        {0, 3},  // rank 2
+        {1, 2}   // rank 3
+    };
+    std::vector<std::vector<gko::span>> spans{
+        {gko::span{0, 2}, gko::span{2, 4}},  // rank 0
+        {gko::span{0, 2}, gko::span{2, 4}},  // rank 1
+        {gko::span{0, 2}, gko::span{2, 4}},  // rank 2
+        {gko::span{0, 2}, gko::span{2, 4}},  // rank 3
+    };
+};
+
+INSTANTIATE_TEST_SUITE_P(RepartitionerFixture1DInstantiation,
+                         RepartitionerFixture1D, testing::Values(1, 2, 4));
+INSTANTIATE_TEST_SUITE_P(RepartitionerFixture2DInstantiation,
+                         RepartitionerFixture2D, testing::Values(1, 2, 4));
+
+TEST_P(RepartitionerFixture1D, can_create_repartitioner)
 {
     // Arrange
     label ranks_per_gpu = GetParam();
-    auto exec = ((RepartitionerEnvironment *)global_env)->exec;
-
     // Act
-    auto repartitioner = Repartitioner(10, ranks_per_gpu, 0, *exec.get());
-
+    auto repartitioner = Repartitioner(10, ranks_per_gpu, 0, exec);
     // Assert
     EXPECT_EQ(repartitioner.get_ranks_per_gpu(), ranks_per_gpu);
 }
 
 
-TEST_P(RepartitionerFixture, has_correct_properties_for_n_rank)
+TEST_P(RepartitionerFixture1D, has_correct_properties_for_n_rank)
 {
     // Arrange
     label ranks_per_gpu = GetParam();
-    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
-    auto rank = exec->get_rank();
-
     // Act
-    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, exec);
 
     // Assert
-    EXPECT_EQ(repartitioner.is_owner(*exec),
+    EXPECT_EQ(repartitioner.is_owner(exec),
               (rank % ranks_per_gpu == 0) ? true : false);
 
     EXPECT_EQ(
-        repartitioner.compute_repart_size(local_size, ranks_per_gpu, *exec),
-        (repartitioner.is_owner(*exec) ? ranks_per_gpu * local_size : 0));
+        repartitioner.compute_repart_size(local_size, ranks_per_gpu, exec),
+        (repartitioner.is_owner(exec) ? ranks_per_gpu * local_size : 0));
 }
 
-TEST(Repartitioner, can_convert_to_global)
+TEST_P(RepartitionerFixture2D, can_convert_to_global)
 {
     // Arrange
-    auto& exec = *((RepartitionerEnvironment *)global_env)->exec;
-    auto rank = exec.get_rank();
+    using namespace gko::experimental::distributed;
     auto ref_exec = exec.get_ref_exec();
-    auto comm = exec.get_communicator();
     auto partition = gko::share(
-        gko::experimental::distributed::build_partition_from_local_size<label,
-                                                                        label>(
-            ref_exec, *comm.get(), 4));
-
-    //         local ids          |   global ids
-    // cells: [ 2 3 | 2 3 ]       |  [ 10 11 | 14 15 ]
-    //   2    [ 0 1 | 0 1 ]  3    |  [  8  9 | 12 13 ]
-    //        -------------       |  -----------------
-    //        [ 2 3 | 2 3 ]       |  [  2  3 |  6  7 ]
-    //   0    [ 0 1 | 0 1 ]  1    |  [  0  1 |  4  5 ]
-    //
+        build_partition_from_local_size<label, label>(ref_exec, comm, 4));
 
-    // vector
-    std::vector<std::vector<label>> idxs{
-        {0, 2, 0, 1},  // rank 0
-        {1, 3, 0, 1},  // rank 1
-        {2, 3, 0, 2},  // rank 2
-        {2, 3, 1, 3}   // rank 3
-    };
-    // vector
-    std::vector<std::vector<label>> ranks{
-        {1, 2},  // rank 0
-        {0, 3},  // rank 1
-        {0, 3},  // rank 2
-        {1, 2}   // rank 3
-    };
-    std::vector<std::vector<gko::span>> spans{
-        {gko::span{0, 2}, gko::span{2, 4}},  // rank 0
-        {gko::span{0, 2}, gko::span{2, 4}},  // rank 1
-        {gko::span{0, 2}, gko::span{2, 4}},  // rank 2
-        {gko::span{0, 2}, gko::span{2, 4}},  // rank 3
-    };
-
-    std::vector<std::vector<label>> exp_global_idxs{
+    vec_vec exp_global_idxs{
         {4, 6, 8, 9},    // rank 0
         {1, 3, 12, 13},  // rank 1
         {2, 3, 12, 14},  // rank 2
@@ -160,13 +205,10 @@ TEST(Repartitioner, can_convert_to_global)
 }
 
 
-TEST_P(RepartitionerFixture, can_exchange_spans_and_ranks_for_n_ranks)
+TEST_P(RepartitionerFixture1D, can_exchange_spans_and_ranks_for_n_ranks)
 {
     // Arrange
     label ranks_per_gpu = GetParam();
-    auto exec = *((RepartitionerEnvironment *)global_env)->exec.get();
-    auto rank = exec.get_rank();
-    auto ref_exec = exec.get_ref_exec();
 
     std::vector<gko::span> spans{{0, 5}, {5, 10}};
     std::vector<label> ranks{rank + 1, rank + 2};
@@ -186,7 +228,8 @@ TEST_P(RepartitionerFixture, can_exchange_spans_and_ranks_for_n_ranks)
     std::map<label, vec_vec> exp_spans_end;
     exp_spans_end.emplace(1, vec_vec{{5, 10}, {5, 10}, {5, 10}, {5, 10}});
     exp_spans_end.emplace(2, vec_vec{{5, 10, 15, 20}, {}, {5, 10, 15, 20}, {}});
-    exp_spans_end.emplace(4, vec_vec{{5, 10, 15, 20, 25, 30, 35, 40}, {}, {}, {}});
+    exp_spans_end.emplace(4,
+                          vec_vec{{5, 10, 15, 20, 25, 30, 35, 40}, {}, {}, {}});
 
     // Act
     auto [new_spans, new_ranks, origins] =
@@ -200,62 +243,27 @@ TEST_P(RepartitionerFixture, can_exchange_spans_and_ranks_for_n_ranks)
         res_spans_end.push_back(span.end);
     }
 
-
-
     // Assert
     ASSERT_EQ(new_ranks, exp_ranks[ranks_per_gpu][rank]);
     ASSERT_EQ(res_spans_begin, exp_spans_begin[ranks_per_gpu][rank]);
     ASSERT_EQ(res_spans_end, exp_spans_end[ranks_per_gpu][rank]);
 }
 
-TEST_P(RepartitionerFixture, can_repartition_sparsity_pattern_1D_for_n_ranks)
+TEST_P(RepartitionerFixture1D, can_repartition_sparsity_pattern_1D_for_n_ranks)
 {
     // Arrange
     label ranks_per_gpu = GetParam();
-    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
-    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
-    auto rank = exec->get_rank();
-    auto ref_exec = exec->get_ref_exec();
-
-    // mesh:
-    // rank:     0     1     2     3
-    // cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
-    // cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- local row ids
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, exec);
+    auto ref_exec = exec.get_ref_exec();
 
-    std::vector<label> rows{0, 0, 1, 1};
-    std::vector<label> cols{0, 1, 0, 1};
-    std::vector<label> mapping{0, 1, 2, 3};
-    std::vector<gko::span> spans{gko::span{0, 2}};
     std::vector<label> ranks{rank};
     auto local_sparsity = std::make_shared<SparsityPattern>(
         ref_exec, gko::dim<2>{2, 2}, rows, cols, mapping, spans, ranks);
 
-    // rank local rows
-    vec_vec non_local_rows{{1}, {0, 1}, {0, 1}, {0}};
-
-    // other side local rows
-    vec_vec non_local_cols{{0}, {1, 0}, {1, 0}, {1}};
-
-    vec_vec non_local_mapping{{0}, {0, 1}, {0, 1}, {0}};
-
-    std::vector<std::vector<gko::span>> non_local_spans{
-        {gko::span{0, 1}},                   // rank 0
-        {gko::span{0, 1}, gko::span{1, 2}},  // rank 1
-        {gko::span{0, 1}, gko::span{1, 2}},  // rank 2
-        {gko::span{0, 1}},                   // rank 3
-    };
-
-    vec_vec non_local_ranks{
-        {1},     // rank 0
-        {0, 2},  // rank 1
-        {1, 3},  // rank 2
-        {2},     // rank 3
-    };
-
     auto non_local_sparsity = std::make_shared<SparsityPattern>(
-        ref_exec, gko::dim<2>{2, non_local_ranks[rank].size()}, non_local_rows[rank],
-        non_local_cols[rank], non_local_mapping[rank], non_local_spans[rank],
-        non_local_ranks[rank]);
+        ref_exec, gko::dim<2>{2, non_local_ranks[rank].size()},
+        non_local_rows[rank], non_local_cols[rank], non_local_mapping[rank],
+        non_local_spans[rank], non_local_ranks[rank]);
 
     std::map<label, std::vector<label>> exp_local_nnz;
     exp_local_nnz.emplace(1, std::vector<label>{4, 4, 4, 4});
@@ -267,34 +275,41 @@ TEST_P(RepartitionerFixture, can_repartition_sparsity_pattern_1D_for_n_ranks)
     exp_non_local_nnz.emplace(2, std::vector<label>{1, 0, 1, 0});
     exp_non_local_nnz.emplace(4, std::vector<label>{0, 0, 0, 0});
 
+    std::map<label, vec_vec> exp_local_rows;
+    exp_local_rows.emplace(1, vec_vec{rows, rows, rows, rows});
+    exp_local_rows.emplace(2, vec_vec{{0, 0, 1, 1, 2, 2, 3, 3, 1, 2},
+                                      {},
+                                      {0, 0, 1, 1, 2, 2, 3, 3, 1, 2},
+                                      {}});
+    exp_local_rows.emplace(4, vec_vec{{0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5,
+                                       5, 6, 6, 7, 7, 1, 2, 3, 4, 5, 6},
+                                      {},
+                                      {},
+                                      {}});
+
     // Act
     auto [repart_local, repart_non_local, tracking] =
-        repartitioner.repartition_sparsity(*exec, local_sparsity,
+        repartitioner.repartition_sparsity(exec, local_sparsity,
                                            non_local_sparsity);
+    auto res_local_rows = convert_to_vector(repart_local->row_idxs);
+    auto res_local_cols = convert_to_vector(repart_local->col_idxs);
 
     // Assert
     ASSERT_EQ(repart_local->num_nnz, exp_local_nnz[ranks_per_gpu][rank]);
     ASSERT_EQ(repart_non_local->num_nnz,
               exp_non_local_nnz[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_local_rows, exp_local_rows[ranks_per_gpu][rank]);
 }
 
-TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
+TEST_P(RepartitionerFixture1D, can_repartition_1D_comm_pattern_for_n_ranks)
 {
     // Arrange
     label ranks_per_gpu = GetParam();
-    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
-    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
-    auto rank = exec->get_rank();
-    auto ref_exec = exec->get_ref_exec();
-
-    // rank:     0     1     2     3
-    // cells: [ 0 1 | 2 3 | 4 5 | 6 7 ] <- global row ids
-    // cells: [ 0 1 | 0 1 | 0 1 | 0 1 ] <- local row ids
-    std::vector<std::vector<label>> ids{{1}, {0, 2}, {1, 3}, {2}};
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, exec);
+    auto ref_exec = exec.get_ref_exec();
 
     // expected communcation ranks
     std::map<label, vec_vec> exp_res_ids{};
-
     exp_res_ids[1] = ids;  // in the ranks_per_gpu==1 case nothing changes
     // only communication partners are 0-2 and 2-0
     exp_res_ids.emplace(2, vec_vec{{2}, {}, {0}, {}});
@@ -320,55 +335,32 @@ TEST_P(RepartitionerFixture, can_repartition_1D_comm_pattern_for_n_ranks)
 
     // the original comm_pattern
     auto comm_pattern =
-        std::make_shared<CommunicationPattern>(*exec, ids[rank], rows[rank]);
+        std::make_shared<CommunicationPattern>(exec, ids[rank], rows[rank]);
 
-    auto comm = exec->get_communicator();
-    auto partition = gko::share(
-        gko::experimental::distributed::build_partition_from_local_size<label,
-                                                                        label>(
-            ref_exec, *comm.get(), 2));
+    auto partition = repartitioner.get_orig_partition();
 
     // Act
     auto repart_comm_pattern =
-        repartitioner.repartition_comm_pattern(*exec, comm_pattern, partition);
+        repartitioner.repartition_comm_pattern(exec, comm_pattern, partition);
 
     // Assert
-    std::vector<label> res_ids(
-        repart_comm_pattern->target_ids.get_const_data(),
-        repart_comm_pattern->target_ids.get_const_data() +
-            repart_comm_pattern->target_ids.get_size());
-
+    auto res_ids = convert_to_vector(repart_comm_pattern->target_ids);
     EXPECT_EQ(res_ids, exp_res_ids[ranks_per_gpu][rank]);
 
-    std::vector<label> res_sizes(
-        repart_comm_pattern->target_sizes.get_const_data(),
-        repart_comm_pattern->target_sizes.get_const_data() +
-            repart_comm_pattern->target_sizes.get_size());
+    auto res_sizes = convert_to_vector(repart_comm_pattern->target_sizes);
     EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
 
     auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
-    std::vector<label> res_rows(
-        total_rank_send_idx.get_const_data(),
-        total_rank_send_idx.get_const_data() + total_rank_send_idx.get_size());
-
+    auto res_rows = convert_to_vector(total_rank_send_idx);
     EXPECT_EQ(res_rows, exp_res_rows[ranks_per_gpu][rank]);
 }
 
-TEST_P(RepartitionerFixture, can_repartition_2D_comm_pattern_for_n_ranks)
+TEST_P(RepartitionerFixture2D, can_repartition_2D_comm_pattern_for_n_ranks)
 {
     // Arrange
     label ranks_per_gpu = GetParam();
-    auto exec = ((RepartitionerEnvironment *)global_env)->exec.get();
-    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, *exec);
-    auto rank = exec->get_rank();
-    auto ref_exec = exec->get_ref_exec();
-
-    // cells: [ 2 3 | 2 3 ]
-    //   2    [ 0 1 | 0 1 ]  3
-    //        -------------
-    //        [ 2 3 | 2 3 ]
-    //   0    [ 0 1 | 0 1 ]  1
-    std::vector<std::vector<label>> ids{{1, 2}, {0, 3}, {0, 3}, {1, 2}};
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, exec);
+    auto ref_exec = exec.get_ref_exec();
 
     // expected communcation ranks
     std::map<label, vec_vec> exp_res_ids{};
@@ -404,37 +396,24 @@ TEST_P(RepartitionerFixture, can_repartition_2D_comm_pattern_for_n_ranks)
 
     // the original comm_pattern
     auto comm_pattern =
-        std::make_shared<CommunicationPattern>(*exec, ids[rank], rows[rank]);
+        std::make_shared<CommunicationPattern>(exec, ids[rank], rows[rank]);
 
-    auto comm = exec->get_communicator();
-    auto partition = gko::share(
-        gko::experimental::distributed::build_partition_from_local_size<label,
-                                                                        label>(
-            ref_exec, *comm.get(), 4));
+    auto partition = repartitioner.get_orig_partition();
 
     // Act
     auto repart_comm_pattern =
-        repartitioner.repartition_comm_pattern(*exec, comm_pattern, partition);
+        repartitioner.repartition_comm_pattern(exec, comm_pattern, partition);
 
     // Assert
-    std::vector<label> res_ids(
-        repart_comm_pattern->target_ids.get_const_data(),
-        repart_comm_pattern->target_ids.get_const_data() +
-            repart_comm_pattern->target_ids.get_size());
+    auto res_ids = convert_to_vector(repart_comm_pattern->target_ids);
 
     EXPECT_EQ(res_ids, exp_res_ids[ranks_per_gpu][rank]);
 
-    std::vector<label> res_sizes(
-        repart_comm_pattern->target_sizes.get_const_data(),
-        repart_comm_pattern->target_sizes.get_const_data() +
-            repart_comm_pattern->target_sizes.get_size());
+    auto res_sizes = convert_to_vector(repart_comm_pattern->target_sizes);
     EXPECT_EQ(res_sizes, exp_res_sizes[ranks_per_gpu][rank]);
 
     auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
-    std::vector<label> res_rows(
-        total_rank_send_idx.get_const_data(),
-        total_rank_send_idx.get_const_data() + total_rank_send_idx.get_size());
-
+    auto res_rows = convert_to_vector(total_rank_send_idx);
     EXPECT_EQ(res_rows, exp_res_rows[ranks_per_gpu][rank]);
 }
 

From 7953f05aea213b39755d9f79cd721fda4a4b27c4 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 28 Aug 2024 11:41:03 +0200
Subject: [PATCH 49/67] add non local column test

---
 src/Repartitioner.C       | 192 ++++++++++----------------------------
 unitTests/Repartitioner.C |  43 ++++++++-
 2 files changed, 87 insertions(+), 148 deletions(-)

diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index 84f04d6dd..087dbf3b0 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -6,9 +6,10 @@
 
 namespace detail {
 
-std::tuple<std::vector<gko::span>, std::vector<label>, std::vector<label>> exchange_span_ranks(
-    const ExecutorHandler &exec_handler, label ranks_per_gpu,
-    const std::vector<gko::span> &spans, const std::vector<label> &ranks)
+std::tuple<std::vector<gko::span>, std::vector<label>, std::vector<label>>
+exchange_span_ranks(const ExecutorHandler &exec_handler, label ranks_per_gpu,
+                    const std::vector<gko::span> &spans,
+                    const std::vector<label> &ranks)
 {
     auto comm_pattern = compute_gather_to_owner_counts(
         exec_handler, ranks_per_gpu, spans.size());
@@ -34,14 +35,14 @@ std::tuple<std::vector<gko::span>, std::vector<label>, std::vector<label>> excha
 
     std::vector<label> origins{};
     if (gathered_size.size() > 0) {
-    auto rank = exec_handler.get_rank();
-    for (int i = 0; i<ranks_per_gpu; i++) {
-        auto j = comm_pattern.recv_counts[i+rank];
-        for (int k=0;k<j;k++) {
-            origins.push_back(i+rank);
+        auto rank = exec_handler.get_rank();
+        for (int i = 0; i < ranks_per_gpu; i++) {
+            auto j = comm_pattern.recv_counts[i + rank];
+            for (int k = 0; k < j; k++) {
+                origins.push_back(i + rank);
+            }
         }
     }
-    }
 
     return {out_spans, new_ranks, origins};
 }
@@ -148,16 +149,16 @@ Repartitioner::repartition_sparsity(
 
     // Done with Step 1. Next we gather all non local sparsity. Again
     // all of it is collected into temporary vectors first.
-
     auto non_local_comm_pattern = compute_gather_to_owner_counts(
         exec_handler, ranks_per_gpu, src_non_local_pattern->num_nnz);
 
     auto tmp_non_local_rows = gather_closure(
         non_local_comm_pattern, src_non_local_pattern->row_idxs, offset);
 
-    auto [new_spans, new_ranks, tmp_non_local_origin] = detail::exchange_span_ranks(
-        exec_handler, ranks_per_gpu_, src_non_local_pattern->spans,
-        src_non_local_pattern->rank);
+    auto [new_spans, new_ranks, tmp_non_local_origin] =
+        detail::exchange_span_ranks(exec_handler, ranks_per_gpu_,
+                                    src_non_local_pattern->spans,
+                                    src_non_local_pattern->rank);
 
     auto tmp_non_local_cols = detail::convert_to_global(
         orig_partition_,
@@ -174,8 +175,7 @@ Repartitioner::repartition_sparsity(
         tmp_non_local_cols, tmp_non_local_mapping, new_ranks,
         tmp_non_local_origin, new_spans);
 
-    // TODO FIXME
-    gko::dim<2> tmp_non_local_dim{0, 0};
+    gko::dim<2> tmp_non_local_dim{tmp_local_dim[0], tmp_non_local_rows.size()};
 
     if (is_owner(exec_handler)) {
         auto local_pattern = std::make_shared<SparsityPattern>(
@@ -196,99 +196,6 @@ Repartitioner::repartition_sparsity(
             std::make_shared<SparsityPattern>(exec),
             std::make_shared<SparsityPattern>(exec), {});
     }
-    // auto span_comm_pattern = compute_gather_to_owner_counts(
-    //     exec_handler, ranks_per_gpu, src_non_local_pattern->spans.size());
-    // // build vector with locality information
-    // std::vector<std::pair<bool, label>> locality;
-    // label ctr{0};
-
-    // if (is_owner(exec_handler)) {
-    //     for (int i = 0; i < ranks_per_gpu; i++) {
-    //         label comm_rank = rank + i;
-    //         label num_elems = span_comm_pattern.recv_counts[comm_rank];
-    //         for (int j = 0; j < num_elems; j++) {
-    //             locality.emplace_back(is_local[ctr], comm_rank);
-    //             ctr++;
-    //         }
-    //     }
-    // }
-    // // All interfaces that are in is_local need to be accounted for
-    // ASSERT_EQ(ctr, is_local.size());
-
-    // // TODO FIXME the correct way would be to check the communication
-    // // pattern where a particular face is in the send idxs
-    // // since we have already the row id of the other side
-    // // it should be doable. Or alternatively we know that we
-    // // keep an interface together thus we can just count the idx up to the
-    // // size. But we have to make sure that the interfaces are in the same
-    // // order on both communication sides.
-    // for (int i = 0; i < gathered_non_local.col_idxs.get_size(); i++) {
-    //     gathered_non_local.col_idxs.get_data()[i] = i;
-    // }
-
-    // LOG_1(verbose_, "done repartition sparsity pattern")
-    // if (is_owner(exec_handler)) {
-    //     FatalErrorInFunction << "Not implemented" << abort(FatalError);
-    //     // auto new_local_spars_pattern = std::make_shared<SparsityPattern>(
-    //     //     exec, merged_local_dim, merged_local);
-    //     //
-    //     // auto new_non_local_spars_pattern =
-    //     //     std::make_shared<SparsityPattern>(
-    //     //         exec,
-    //     //         gko::dim<2>{merged_local_dim[0],
-    //     //                     gathered_non_local.row_idxs.get_size()},
-    //     //         gathered_non_local);
-    //     //
-    //     //         return std::make_tuple<std::shared_ptr<SparsityPattern>,
-    //     //                                std::shared_ptr<SparsityPattern>,
-    //     //                                std::vector<std::pair<bool,
-    //     //                                label>>>(
-    //     //             std::move(new_local_spars_pattern),
-    //     //             std::move(new_non_local_spars_pattern),
-    //     //             std::move(locality));
-    // } else {
-    //     // auto new_local_spars_pattern =
-    //     //     std::make_shared<SparsityPattern>(exec);
-    //     //
-    //     // auto new_non_local_spars_pattern =
-    //     //     std::make_shared<SparsityPattern>(exec);
-    //     //
-    //     // // FIXME
-    //     // // std::cout << __FILE__ << ":" << __LINE__ << " rank " <<
-    //     // // comm.rank() <<
-    //     // //     " merged_local dim " <<
-    //     // new_local_spars_pattern->dim =
-    //     //     gko::dim<2>{0, 0};  // merged_local_dim;
-    //     // new_non_local_spars_pattern->dim =
-    //     //     gko::dim<2>{0, 0};  // merged_local_dim;
-    //     //
-    //     //
-    //     FatalErrorInFunction << "Not implemented" << abort(FatalError);
-    //     // return std::make_tuple<std::shared_ptr<SparsityPattern>,
-    //     //                        std::shared_ptr<SparsityPattern>,
-    //     //                        std::vector<std::pair<bool, label>>>(
-    //     //     std::move(new_local_spars_pattern),
-    //     //     std::move(new_non_local_spars_pattern),
-    //     //     std::vector<std::pair<bool, label>>{});
-    // }
-    // // TODO Some old documentation see what is still needed
-    // // add interfaces
-    // // to add interfaces we go through the non_local_sparsity
-    // // pattern and check if interface is still non_local
-    // //
-    // // iterate interfaces and send to owner.
-    // // on owner decide to move to local or non_local
-    // // for sending to owner we can also use the all_to_all_v
-    // // approach we than have
-    // //
-    // // row [ 4 , 8,  12 | 16 , 17, 18 ] <- local row
-    // // col [ 1 , 2, 3 | 1, 2, 3 ] <- just the interface ctr
-    // // from repartition_comm_pattern we could get
-    // //
-    // // or we split local / non-local first
-    // // for this we could store ranges when computing
-    // // repartition_comm_pattern
-    // //
 }
 
 std::vector<bool> Repartitioner::build_non_local_interfaces(
@@ -318,8 +225,7 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
                               non_local_rows.data() + end);
 
             std::vector<label> tmp_rank_local_cols(
-                non_local_cols.data() + begin,
-                non_local_cols.data() + end);
+                non_local_cols.data() + begin, non_local_cols.data() + end);
 
             detail::convert_to_local(partition, tmp_rank_local_cols, rank);
 
@@ -330,9 +236,8 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
                                  non_local_mapping.data() + end);
 
             gko::size_type rows_start = local_rows.size();
-            local_spans.emplace_back(
-                rows_start,
-                          static_cast<gko::size_type>(end - begin));
+            local_spans.emplace_back(rows_start,
+                                     static_cast<gko::size_type>(end - begin));
 
             local_ranks.push_back(non_local_rank_origin[i]);
         } else {
@@ -342,38 +247,35 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
 
     // remove data from non_local vectors
     if (mark_keep.size() == 0) {
-    non_local_rows.clear();
-    non_local_cols.clear();
-    non_local_mapping.clear();
-    non_local_ranks.clear();
-    non_local_spans.clear();
-        } else {
-    std::vector<label> copy_rows, copy_cols, copy_mapping, copy_ranks;
-    std::vector<gko::span> copy_spans;
-    for (label i : mark_keep) {
-        auto [begin, end] = non_local_spans[i];
-        copy_rows.insert(copy_rows.end(),
-                         non_local_rows.data() + begin,
-                         non_local_rows.data() + end);
-        copy_cols.insert(copy_cols.end(),
-                         non_local_cols.data() + begin,
-                         non_local_cols.data() + end);
-        copy_mapping.insert(copy_mapping.end(),
-                            non_local_mapping.data() + begin,
-                            non_local_mapping.data() + end);
-        copy_ranks.push_back(non_local_ranks[i]);
-        copy_spans.emplace_back(non_local_spans[i].begin,
-                                non_local_spans[i].end);
-    }
-
-    non_local_rows = copy_rows;
-    non_local_cols = copy_cols;
-    non_local_mapping = copy_mapping;
-    non_local_ranks = copy_ranks;
-    non_local_spans.clear();
-    for (auto &span : copy_spans) {
-        non_local_spans.emplace_back(span.begin, span.end);
-    }
+        non_local_rows.clear();
+        non_local_cols.clear();
+        non_local_mapping.clear();
+        non_local_ranks.clear();
+        non_local_spans.clear();
+    } else {
+        std::vector<label> copy_rows, copy_cols, copy_mapping, copy_ranks;
+        std::vector<gko::span> copy_spans;
+        for (label i : mark_keep) {
+            auto [begin, end] = non_local_spans[i];
+            copy_rows.insert(copy_rows.end(), non_local_rows.data() + begin,
+                             non_local_rows.data() + end);
+            copy_cols.insert(copy_cols.end(), non_local_cols.data() + begin,
+                             non_local_cols.data() + end);
+            copy_mapping.insert(copy_mapping.end(),
+                                non_local_mapping.data() + begin,
+                                non_local_mapping.data() + end);
+            copy_ranks.push_back(non_local_ranks[i]);
+            copy_spans.emplace_back(non_local_spans[i].begin,
+                                    non_local_spans[i].end);
+        }
+        non_local_rows = copy_rows;
+        non_local_cols = copy_cols;
+        non_local_mapping = copy_mapping;
+        non_local_ranks = copy_ranks;
+        non_local_spans.clear();
+        for (auto &span : copy_spans) {
+            non_local_spans.emplace_back(span.begin, span.end);
+        }
     }
 
     return is_local;
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 960ffe128..b04c33bc7 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -23,6 +23,7 @@ extern char **my_argv;
 int my_argc;
 char **my_argv;
 
+using vec = std::vector<label>;
 using vec_vec = std::vector<std::vector<label>>;
 using vec_vec_vec = std::vector<vec_vec>;
 
@@ -287,18 +288,54 @@ TEST_P(RepartitionerFixture1D, can_repartition_sparsity_pattern_1D_for_n_ranks)
                                       {},
                                       {}});
 
+    std::map<label, vec_vec> exp_local_cols;
+    exp_local_cols.emplace(1, vec_vec{cols, cols, cols, cols});
+    exp_local_cols.emplace(2, vec_vec{{0, 0, 1, 1, 2, 2, 3, 3, 1, 2},
+                                      {},
+                                      {0, 0, 1, 1, 2, 2, 3, 3, 1, 2},
+                                      {}});
+    exp_local_cols.emplace(4, vec_vec{{0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5,
+                                       5, 6, 6, 7, 7, 1, 2, 3, 4, 5, 6},
+                                      {},
+                                      {},
+                                      {}});
+
+    std::map<label, vec> exp_local_dim_rows;
+    exp_local_dim_rows.emplace(1, vec{local_size, local_size, local_size, local_size});
+    exp_local_dim_rows.emplace(2, vec{2*local_size, 0, 2*local_size, 0});
+    exp_local_dim_rows.emplace(4, vec{4*local_size, 0, 0, 0});
+
+    std::map<label, vec_vec> exp_non_local_rows;
+    exp_non_local_rows.emplace(1, non_local_rows);
+    exp_non_local_rows.emplace(2, vec_vec{{3}, {}, {0}, {}});
+    exp_non_local_rows.emplace(4, vec_vec{{}, {}, {}, {}});
+    std::map<label, vec_vec> exp_non_local_cols;
+    exp_non_local_cols.emplace(1, non_local_cols);
+    exp_non_local_cols.emplace(2, vec_vec{{0}, {}, {3}, {}});
+    exp_non_local_cols.emplace(4, vec_vec{{}, {}, {}, {}});
+
     // Act
     auto [repart_local, repart_non_local, tracking] =
         repartitioner.repartition_sparsity(exec, local_sparsity,
                                            non_local_sparsity);
+    // Assert
+    // local properties
+    ASSERT_EQ(repart_local->num_nnz, exp_local_nnz[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_local->dim[0], exp_local_dim_rows[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_local->dim[1], exp_local_dim_rows[ranks_per_gpu][rank]);
     auto res_local_rows = convert_to_vector(repart_local->row_idxs);
     auto res_local_cols = convert_to_vector(repart_local->col_idxs);
+    ASSERT_EQ(res_local_rows, exp_local_rows[ranks_per_gpu][rank]);
 
-    // Assert
-    ASSERT_EQ(repart_local->num_nnz, exp_local_nnz[ranks_per_gpu][rank]);
+    // non local properties
     ASSERT_EQ(repart_non_local->num_nnz,
               exp_non_local_nnz[ranks_per_gpu][rank]);
-    ASSERT_EQ(res_local_rows, exp_local_rows[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_non_local->dim[0], exp_local_dim_rows[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_non_local->dim[1], exp_non_local_nnz[ranks_per_gpu][rank]);
+    auto res_non_local_rows = convert_to_vector(repart_non_local->row_idxs);
+    auto res_non_local_cols = convert_to_vector(repart_non_local->col_idxs);
+    ASSERT_EQ(res_non_local_rows, exp_non_local_rows[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_non_local_cols, exp_non_local_cols[ranks_per_gpu][rank]);
 }
 
 TEST_P(RepartitionerFixture1D, can_repartition_1D_comm_pattern_for_n_ranks)

From e81098d4ffd8a10be6c332010fc2eee8cd2168a8 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 28 Aug 2024 12:15:26 +0200
Subject: [PATCH 50/67] fix typos

---
 include/OGL/Repartitioner.H | 2 +-
 unitTests/Repartitioner.C   | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index e93b8746f..3ec2d0adc 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -223,7 +223,7 @@ public:
      * @param src_comm_pattern the original communication pattern of the current
      * rank belonging to the distributed matrix
      * @param partition the original partition belonging to the distributed
-     * matrix returns the repartitioned communcation pattern
+     * matrix returns the repartitioned communication pattern
      */
     std::shared_ptr<const CommunicationPattern> repartition_comm_pattern(
         const ExecutorHandler &exec_handler,
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index a4423d287..70fb0e53b 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -345,7 +345,7 @@ TEST_P(RepartitionerFixture1D, can_repartition_1D_comm_pattern_for_n_ranks)
     auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, exec);
     auto ref_exec = exec.get_ref_exec();
 
-    // expected communcation ranks
+    // expected communication ranks
     std::map<label, vec_vec> exp_res_ids{};
 
     exp_res_ids[1] = ids;  // in the ranks_per_gpu==1 case nothing changes
@@ -400,7 +400,7 @@ TEST_P(RepartitionerFixture2D, can_repartition_2D_comm_pattern_for_n_ranks)
     auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, exec);
     auto ref_exec = exec.get_ref_exec();
 
-    // expected communcation ranks
+    // expected communication ranks
     std::map<label, vec_vec> exp_res_ids{};
     exp_res_ids[1] = ids;  // in the ranks_per_gpu==1 case nothing changes
     // only communication partners are 0-2 and 2-0

From 3c6074b6119af56ce6268b348db2090a5360a125 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 28 Aug 2024 12:16:20 +0200
Subject: [PATCH 51/67] format file

---
 src/common.C | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/src/common.C b/src/common.C
index 25de64bb2..1fec3f7a2 100644
--- a/src/common.C
+++ b/src/common.C
@@ -145,17 +145,18 @@ label get_solve_prev_iters(word sys_matrix_name, const objectRegistry &db,
     return get_gko_solver_property(sys_matrix_name, iters_name, db, label(1));
 }
 
-std::shared_ptr<gko::array<label>> convert_to_array(const std::vector<label> & in) {
-    return std::make_shared<gko::array<label>>(gko::ReferenceExecutor::create(), in.begin(), in.end());
+std::shared_ptr<gko::array<label>> convert_to_array(
+    const std::vector<label> &in)
+{
+    return std::make_shared<gko::array<label>>(gko::ReferenceExecutor::create(),
+                                               in.begin(), in.end());
 }
 
 
-std::vector<label> convert_to_vector(
- const gko::array<label>& in
-) {
-    return std::vector<label> (
-        in.get_const_data(),
-        in.get_const_data() + in.get_size());
+std::vector<label> convert_to_vector(const gko::array<label> &in)
+{
+    return std::vector<label>(in.get_const_data(),
+                              in.get_const_data() + in.get_size());
 }
 
 

From 492761f0c491000347bb90cbac56c75b0cd2162d Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 28 Aug 2024 16:49:01 +0200
Subject: [PATCH 52/67] fix non local issue

---
 include/OGL/Repartitioner.H | 12 ++++++-----
 src/Repartitioner.C         | 41 +++++++++++++++++++++++++++----------
 unitTests/Repartitioner.C   | 16 ++++++++-------
 3 files changed, 46 insertions(+), 23 deletions(-)

diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index 3ec2d0adc..f1881d8ae 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -166,12 +166,14 @@ public:
         return orig_partition_;
     }
 
-    /* @brief given received interfaces this function sorts interfaces into
+    /* @brief given gathered sparsity information copy local interfaces
      *
-     *local and non local returns: SparsityPattern storing new non_local
-     *sparsity pattern, and a pair of locality information, where the first bool
-     *stores if the interface performs a local communication and the second bool
-     *if the the interface was originally from this rank
+     * After gathering local and non-local rows, cols, and mapping
+     * we have interfaces that are now local in the non-local rows, cols and
+     * mapping. Thus this function copies the rows, cols and mapping, spans etc
+     * to the corresponding local vectors.
+     * returns a vector of locality information.
+     * @param non_local_cols vector in comm rank local indices
      */
     std::vector<bool> build_non_local_interfaces(
         const ExecutorHandler &exec_handler,
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index 8d16e47f0..3fb755ff8 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -70,7 +70,6 @@ std::vector<label> sort_permutation(const std::vector<T> &vec, Compare compare)
 
 }  // namespace detail
 
-
 label Repartitioner::compute_repart_size(label local_size, label ranks_per_gpu,
                                          const ExecutorHandler &exec_handler)
 {
@@ -161,11 +160,10 @@ Repartitioner::repartition_sparsity(
                                     src_non_local_pattern->spans,
                                     src_non_local_pattern->rank);
 
-    auto tmp_non_local_cols = detail::convert_to_global(
-        orig_partition_,
-        gather_closure(non_local_comm_pattern, src_non_local_pattern->col_idxs,
-                       0),
-        new_spans, new_ranks);
+    // comm ranks are based on non repartitioned ranks
+
+    auto tmp_non_local_cols = gather_closure(
+        non_local_comm_pattern, src_non_local_pattern->col_idxs, 0);
 
     std::vector<label> tmp_non_local_mapping =
         std::vector<label>(tmp_non_local_rows.size());
@@ -216,6 +214,9 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
     std::vector<bool> is_local;
     std::vector<label> mark_keep;
 
+    auto tmp_non_local_cols = detail::convert_to_global(
+        orig_partition_, non_local_cols, non_local_spans, non_local_ranks);
+
     for (int i = 0; i < non_local_spans.size(); i++) {
         auto [begin, end] = non_local_spans[i];
         bool local = reparts_to_local(exec_handler, non_local_ranks[i]);
@@ -226,7 +227,8 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
                               non_local_rows.data() + end);
 
             std::vector<label> tmp_rank_local_cols(
-                non_local_cols.data() + begin, non_local_cols.data() + end);
+                tmp_non_local_cols.data() + begin,
+                tmp_non_local_cols.data() + end);
 
             detail::convert_to_local(partition, tmp_rank_local_cols, rank);
 
@@ -260,12 +262,30 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
             auto [begin, end] = non_local_spans[i];
             copy_rows.insert(copy_rows.end(), non_local_rows.data() + begin,
                              non_local_rows.data() + end);
-            copy_cols.insert(copy_cols.end(), non_local_cols.data() + begin,
-                             non_local_cols.data() + end);
+
+            auto orig_comm_rank = non_local_ranks[i];
+            auto repart_comm_rank = get_owner_rank(orig_comm_rank);
+            std::vector<label> tmp_rank_local_cols(
+                tmp_non_local_cols.data() + begin,
+                tmp_non_local_cols.data() + end);
+
+            // this is per interface
+            // need to add offset of original comm rank
+            label rank_offset = partition->get_range_bounds()[orig_comm_rank] -
+                                partition->get_range_bounds()[repart_comm_rank];
+            detail::convert_to_local(partition, tmp_rank_local_cols,
+                                     non_local_ranks[i]);
+            std::transform(
+                tmp_rank_local_cols.begin(), tmp_rank_local_cols.end(),
+                tmp_rank_local_cols.begin(),
+                [rank_offset](label idx) { return idx + rank_offset; });
+
+            copy_cols.insert(copy_cols.end(), tmp_rank_local_cols.begin(),
+                             tmp_rank_local_cols.end());
             copy_mapping.insert(copy_mapping.end(),
                                 non_local_mapping.data() + begin,
                                 non_local_mapping.data() + end);
-            copy_ranks.push_back(non_local_ranks[i]);
+            copy_ranks.push_back(repart_comm_rank);
             copy_spans.emplace_back(non_local_spans[i].begin,
                                     non_local_spans[i].end);
         }
@@ -413,7 +433,6 @@ Repartitioner::repartition_comm_pattern(
     target_ids = detail::apply_permutation(gathered_target_ids, p);
     send_idxs = detail::apply_permutation(send_idxs, p);
 
-
     // Step 4.
     // Merge communication pattern with corresponding neighbours
     merged_target_ids.push_back(target_ids[0]);
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 70fb0e53b..26cf8db60 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -49,9 +49,9 @@ public:
         // to clean up output
         ::testing::TestEventListeners &listeners =
             ::testing::UnitTest::GetInstance()->listeners();
-        if (exec->get_communicator()->rank() != 0) {
-            delete listeners.Release(listeners.default_result_printer());
-        }
+        // if (exec->get_communicator()->rank() != 0) {
+        //     delete listeners.Release(listeners.default_result_printer());
+        // }
     }
 
     std::shared_ptr<Foam::argList> args_;
@@ -301,9 +301,10 @@ TEST_P(RepartitionerFixture1D, can_repartition_sparsity_pattern_1D_for_n_ranks)
                                       {}});
 
     std::map<label, vec> exp_local_dim_rows;
-    exp_local_dim_rows.emplace(1, vec{local_size, local_size, local_size, local_size});
-    exp_local_dim_rows.emplace(2, vec{2*local_size, 0, 2*local_size, 0});
-    exp_local_dim_rows.emplace(4, vec{4*local_size, 0, 0, 0});
+    exp_local_dim_rows.emplace(
+        1, vec{local_size, local_size, local_size, local_size});
+    exp_local_dim_rows.emplace(2, vec{2 * local_size, 0, 2 * local_size, 0});
+    exp_local_dim_rows.emplace(4, vec{4 * local_size, 0, 0, 0});
 
     std::map<label, vec_vec> exp_non_local_rows;
     exp_non_local_rows.emplace(1, non_local_rows);
@@ -330,7 +331,8 @@ TEST_P(RepartitionerFixture1D, can_repartition_sparsity_pattern_1D_for_n_ranks)
     // non local properties
     ASSERT_EQ(repart_non_local->num_nnz,
               exp_non_local_nnz[ranks_per_gpu][rank]);
-    ASSERT_EQ(repart_non_local->dim[0], exp_local_dim_rows[ranks_per_gpu][rank]);
+    ASSERT_EQ(repart_non_local->dim[0],
+              exp_local_dim_rows[ranks_per_gpu][rank]);
     ASSERT_EQ(repart_non_local->dim[1], exp_non_local_nnz[ranks_per_gpu][rank]);
     auto res_non_local_rows = convert_to_vector(repart_non_local->row_idxs);
     auto res_non_local_cols = convert_to_vector(repart_non_local->col_idxs);

From 1c1f5055ea2bbc3f220daa3484e4fb95889ae4eb Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 28 Aug 2024 20:07:17 +0200
Subject: [PATCH 53/67] fix typos

---
 src/Repartitioner.C | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index 3fb755ff8..8c53616e7 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -323,7 +323,7 @@ Repartitioner::repartition_comm_pattern(
     bool owner = is_owner(exec_handler);
 
     // Step 1. Check if communication partner is non-local after
-    // repartitioning. If it is non-local we keep it. Otherwise communcation
+    // repartitioning. If it is non-local we keep it. Otherwise communication
     // partners that are local after repartitioning can be discarded. Here
     // non-local means: the communication target rank id is != repartitioned
     // rank id

From ef48e6eb5b49c5df8752fafcc68c3f5ceeb93978 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 28 Aug 2024 20:09:34 +0200
Subject: [PATCH 54/67] format

---
 src/Repartitioner.C | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index 8c53616e7..2192a12b0 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -406,7 +406,6 @@ Repartitioner::repartition_comm_pattern(
         send_idxs.clear();
     }
 
-
     // early return if no communication partners are left
     if (gathered_target_ids.size() == 0) {
         return std::make_shared<CommunicationPattern>(
@@ -415,7 +414,6 @@ Repartitioner::repartition_comm_pattern(
                                        target_ids.end()},
             gko::array<comm_size_type>{exec, target_sizes.begin(),
                                        target_sizes.end()},
-
             send_idxs);
     }
 

From 45d70e9ae427291af30adf3b7de45904ead382cf Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Fri, 9 Aug 2024 09:37:59 +0200
Subject: [PATCH 55/67] add Distributed file

---
 include/OGL/MatrixWrapper/Distributed.H | 616 ++++++++++++++++++++++++
 1 file changed, 616 insertions(+)
 create mode 100644 include/OGL/MatrixWrapper/Distributed.H

diff --git a/include/OGL/MatrixWrapper/Distributed.H b/include/OGL/MatrixWrapper/Distributed.H
new file mode 100644
index 000000000..ac760c597
--- /dev/null
+++ b/include/OGL/MatrixWrapper/Distributed.H
@@ -0,0 +1,616 @@
+// SPDX-FileCopyrightText: 2024 OGL authors
+//
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#pragma once
+
+#include <map>
+
+#include <ginkgo/core/distributed/lin_op.hpp>
+#include <ginkgo/ginkgo.hpp>
+
+
+#include "MatrixWrapper/CommunicationPattern/CommunicationPattern.H"
+#include "MatrixWrapper/LDUMatrix/HostMatrix.H"
+#include "Repartitioner/Repartitioner.H"
+
+
+template <typename DistMtxType>
+std::unique_ptr<DistMtxType> generate_dist_mtx_with_inner_type(
+    word matrix_format, std::shared_ptr<const gko::Executor> exec,
+    gko::experimental::mpi::communicator comm,
+    std::shared_ptr<const SparsityPattern> local_sparsity,
+    std::shared_ptr<const SparsityPattern> non_local_sparsity)
+{
+    auto local_dim = local_sparsity->dim;
+    auto non_local_dim = non_local_sparsity->dim;
+    auto local_interfaces = local_sparsity->interface_spans;
+    auto non_local_interfaces = non_local_sparsity->interface_spans;
+
+    if (matrix_format == "Ell") {
+        using mtx_type = gko::matrix::Ell<scalar, label>;
+
+        return DistMtxType::create(
+            exec, comm,
+            CombinationMatrix<scalar, label, mtx_type>::create(
+                exec, local_dim, local_interfaces),
+            CombinationMatrix<scalar, label, mtx_type>::create(
+                exec, non_local_dim, non_local_interfaces));
+    }
+    if (matrix_format == "Csr") {
+        using mtx_type = gko::matrix::Csr<scalar, label>;
+
+        return DistMtxType::create(
+            exec, comm,
+            CombinationMatrix<scalar, label, mtx_type>::create(
+                exec, local_dim, local_interfaces),
+            CombinationMatrix<scalar, label, mtx_type>::create(
+                exec, non_local_dim, non_local_interfaces));
+    }
+    if (matrix_format == "Coo") {
+        using mtx_type = gko::matrix::Coo<scalar, label>;
+
+        return DistMtxType::create(
+            exec, comm,
+            CombinationMatrix<scalar, label, mtx_type>::create(
+                exec, local_dim, local_interfaces),
+            CombinationMatrix<scalar, label, mtx_type>::create(
+                exec, non_local_dim, non_local_interfaces));
+    }
+
+    FatalErrorInFunction << "Matrix format " << matrix_format
+                         << " not supported! Supported formats: Csr, Ell, Coo"
+                         << abort(FatalError);
+    return {};
+}
+
+
+// The RepartDistMatrix class is a wrapper around Ginkgos distributed Matrix
+// class
+//
+// It adds functionality for repeated read and repartitioning operatitions. As a
+// constraint it is required that the inner matrix types of the distributed
+// matrix are of RepartDistMatrix type.
+template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
+class RepartDistMatrix
+    : public gko::experimental::EnableDistributedLinOp<
+          RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>>,
+      public gko::EnableCreateMethod<
+          RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>>,
+      public gko::experimental::distributed::DistributedBase {
+    friend class gko::EnableCreateMethod<
+        RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>>;
+    friend class gko::EnablePolymorphicObject<
+        RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>,
+        gko::LinOp>;
+
+    using dist_mtx =
+        gko::experimental::distributed::Matrix<ValueType, LocalIndexType,
+                                               GlobalIndexType>;
+
+    using part_type = gko::experimental::distributed::Partition<label, label>;
+
+    using vec = gko::matrix::Dense<ValueType>;
+    using device_matrix_data = gko::device_matrix_data<scalar, label>;
+    using communicator = gko::experimental::mpi::communicator;
+
+public:
+    using gko::experimental::EnableDistributedLinOp<RepartDistMatrix<
+        ValueType, LocalIndexType, GlobalIndexType>>::convert_to;
+    using gko::experimental::EnableDistributedLinOp<
+        RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>>::move_to;
+
+    std::shared_ptr<const gko::LinOp> get_local_matrix()
+    {
+        this->dist_mtx_->get_local_matrix();
+    }
+
+    std::shared_ptr<const gko::LinOp> get_non_local_matrix()
+    {
+        this->dist_mtx_->get_non_local_matrix();
+    }
+
+    /**
+     * Copy-assigns a CombinationMatrix matrix. Preserves executor, copies
+     * everything else.
+     */
+    RepartDistMatrix &operator=(const RepartDistMatrix &other)
+    {
+        if (&other != this) {
+            // FatalErrorInFunction << "Copying the RepartDistMatrix is
+            // disallowed "
+            //                         "for performance reasons"
+            //                      << abort(FatalError);
+            gko::experimental::EnableDistributedLinOp<
+                RepartDistMatrix>::operator=(other);
+            this->dist_mtx_ = other.dist_mtx_;
+            this->local_sparsity_ = other.local_sparsity_;
+            this->non_local_sparsity_ = other.non_local_sparsity_;
+            this->src_comm_pattern_ = other.src_comm_pattern_;
+            this->local_interfaces_ = other.local_interfaces_;
+        }
+        return *this;
+    }
+
+    /**
+     * Move-assigns a CombinationMatrix matrix. Preserves executor, moves the
+     * data and leaves the moved-from object in an empty state (0x0 LinOp with
+     * unchanged executor and strategy, no nonzeros and valid row pointers).
+     */
+    RepartDistMatrix &operator=(RepartDistMatrix &&other)
+    {
+        if (&other != this) {
+            FatalErrorInFunction << "Not implemented" << abort(FatalError);
+            gko::experimental::EnableDistributedLinOp<
+                RepartDistMatrix>::operator=(std::move(other));
+            this->dist_mtx_ = std::move(other.dist_mtx_);
+            this->local_sparsity_ = std::move(other.local_sparsity_);
+            this->non_local_sparsity_ = std::move(other.non_local_sparsity_);
+            this->src_comm_pattern_ = std::move(other.src_comm_pattern_);
+            this->local_interfaces_ = std::move(other.local_interfaces_);
+        }
+        return *this;
+    }
+
+    static std::shared_ptr<RepartDistMatrix> create(
+        const ExecutorHandler &exec_handler, word matrix_format,
+        const Repartitioner &repartitioner,
+        std::shared_ptr<const HostMatrixWrapper> host_A)
+    {
+        auto exec = exec_handler.get_ref_exec();
+        auto comm = *exec_handler.get_communicator().get();
+
+        // repartition things here first by creating device_matrix_data on
+        // device with correct size on each rank
+        //
+        // NOTE Start by copying first and repartition later
+        auto local_sparsity_ = host_A->compute_local_sparsity(exec);
+        auto non_local_sparsity_ = host_A->compute_non_local_sparsity(exec);
+
+        auto orig_partition = repartitioner.get_orig_partition();
+
+        // create original communicator pattern
+        auto ranks_per_gpu = repartitioner.get_ranks_per_gpu();
+        label rank{comm.rank()};
+
+        auto src_comm_pattern =
+            host_A.get()->create_communication_pattern(exec_handler);
+        // create partiton here and pass to constructor
+        //
+        auto dst_comm_pattern = repartitioner.repartition_comm_pattern(
+            exec_handler, src_comm_pattern, orig_partition);
+
+        label owner_rank = repartitioner.get_owner_rank(exec_handler);
+        bool owner = repartitioner.is_owner(exec_handler);
+
+        auto [repart_loc_sparsity, repart_non_loc_sparsity, local_interfaces] =
+            repartitioner.repartition_sparsity(exec_handler, local_sparsity_,
+                                               non_local_sparsity_);
+
+        local_sparsity_ = repart_loc_sparsity;
+        non_local_sparsity_ = repart_non_loc_sparsity;
+
+
+        auto device_exec = exec_handler.get_device_exec();
+        auto dist_A = gko::share(generate_dist_mtx_with_inner_type<dist_mtx>(
+            matrix_format, device_exec, comm, local_sparsity_,
+            non_local_sparsity_));
+
+        auto local_coeffs = gko::array<scalar>(exec, local_sparsity_->size_);
+        auto non_local_coeffs =
+            gko::array<scalar>(exec, non_local_sparsity_->size_);
+
+
+        if (owner) {
+            local_coeffs.fill(0.0);
+            non_local_coeffs.fill(0.0);
+        }
+
+        std::cout << __FILE__ << " rank " << rank << " local sparsity size "
+                  << local_sparsity_->size_ << " local sparsity dim ["
+                  << local_sparsity_->dim[0] << "x" << local_sparsity_->dim[1]
+                  << "] non_local sparsity size " << non_local_sparsity_->size_
+                  << " non local sparsity dim [" << non_local_sparsity_->dim[0]
+                  << "x" << non_local_sparsity_->dim[1] << "] target_ids "
+                  << dst_comm_pattern->target_ids << " target_sizes "
+                  << dst_comm_pattern->target_sizes << " target_send_idxs.size "
+                  << dst_comm_pattern->send_idxs.size()
+                  << " non_local_sparsity.size " << non_local_sparsity_->size_
+                  << " \n";
+
+        // FIXME make sure that we work on the device executor
+        auto comm_pattern = compute_send_recv_counts(
+            exec_handler, 1, non_local_sparsity_->size_);
+
+        auto send_idxs = dst_comm_pattern->total_rank_send_idx();
+        auto comm_counts = dst_comm_pattern->send_recv_pattern();
+
+
+        dist_A->read_distributed(
+            device_matrix_data(exec, local_sparsity_->dim,
+                               local_sparsity_->row_idxs,
+                               local_sparsity_->col_idxs, local_coeffs),
+            device_matrix_data(exec, non_local_sparsity_->dim,
+                               non_local_sparsity_->row_idxs,
+                               non_local_sparsity_->col_idxs, non_local_coeffs),
+            std::get<0>(comm_counts), std::get<1>(comm_counts),
+            std::get<2>(comm_counts), std::get<3>(comm_counts), send_idxs);
+
+
+        // auto dist_A = DistMtxType::create(
+        //        exec, comm,
+        //        CombinationMatrix<scalar, label, mtx_type>::create(
+        //            exec, local_dim, local_interfaces),
+        //        CombinationMatrix<scalar, label, mtx_type>::create(
+        //            exec, non_local_dim, non_local_interfaces));
+
+        update_impl(exec_handler, matrix_format, repartitioner, host_A, dist_A,
+                    local_sparsity_, non_local_sparsity_, src_comm_pattern,
+                    local_interfaces);
+
+        auto ret = std::make_shared<RepartDistMatrix>(
+            exec, comm, repartitioner.get_repart_dim(), dist_A->get_size(),
+            std::move(dist_A), local_sparsity_, non_local_sparsity_,
+            src_comm_pattern, local_interfaces);
+
+        return ret;
+    };
+
+    void update(const ExecutorHandler &exec_handler, word matrix_format,
+                const Repartitioner &repartitioner,
+                std::shared_ptr<const HostMatrixWrapper> host_A)
+    {
+        update_impl(exec_handler, matrix_format, repartitioner, host_A,
+                    dist_mtx_, local_sparsity_, non_local_sparsity_,
+                    src_comm_pattern_, local_interfaces_);
+    };
+
+    // TODO matrix_format can be data member
+    static void update_impl(
+        const ExecutorHandler &exec_handler, word matrix_format,
+        const Repartitioner &repartitioner,
+        std::shared_ptr<const HostMatrixWrapper> host_A,
+        std::shared_ptr<
+            gko::experimental::distributed::Matrix<scalar, label, label>>
+            dist_A,
+        std::shared_ptr<const SparsityPattern> local_sparsity,
+        std::shared_ptr<const SparsityPattern> non_local_sparsity,
+        std::shared_ptr<const CommunicationPattern> src_comm_pattern,
+        std::vector<std::pair<bool, label>> local_interfaces)
+    {
+        auto exec = exec_handler.get_ref_exec();
+        auto device_exec = exec_handler.get_device_exec();
+        auto ranks_per_gpu = repartitioner.get_ranks_per_gpu();
+        bool requires_host_buffer = exec_handler.get_gko_force_host_buffer();
+
+        label rank{repartitioner.get_rank(exec_handler)};
+        label owner_rank = repartitioner.get_owner_rank(exec_handler);
+        bool owner = repartitioner.is_owner(exec_handler);
+        label nrows = host_A->get_local_nrows();
+        label local_matrix_nnz = host_A->get_local_matrix_nnz();
+        // size + padding has to be local_matrix_nnz
+        auto diag_comm_pattern = compute_send_recv_counts(
+            exec_handler, ranks_per_gpu, nrows, local_matrix_nnz,
+            local_matrix_nnz - nrows, 0);
+        label upper_nnz = host_A->get_upper_nnz();
+        auto upper_comm_pattern = compute_send_recv_counts(
+            exec_handler, ranks_per_gpu, upper_nnz, local_matrix_nnz, 0,
+            local_matrix_nnz - upper_nnz);
+        auto lower_comm_pattern =
+            compute_send_recv_counts(exec_handler, ranks_per_gpu, upper_nnz,
+                                     local_matrix_nnz, upper_nnz, nrows);
+
+        scalar *local_ptr;
+        scalar *local_ptr_2;
+        label nnz = 0;
+
+        // update main values
+        std::vector<scalar> loc_buffer;
+        if (owner) {
+            using Coo = gko::matrix::Coo<scalar, label>;
+            auto local_mtx = dist_A->get_local_matrix();
+
+
+            std::shared_ptr<const Coo> local =
+                gko::as<Coo>(gko::as<CombinationMatrix<scalar, label, Coo>>(
+                                 dist_A->get_local_matrix())
+                                 ->get_combination()
+                                 ->get_operators()[0]);
+            nnz = local->get_num_stored_elements();
+            if (requires_host_buffer) {
+                loc_buffer.resize(nnz);
+                local_ptr = loc_buffer.data();
+                local_ptr_2 = const_cast<scalar *>(local->get_const_values());
+            } else {
+                local_ptr = const_cast<scalar *>(local->get_const_values());
+            }
+        }
+        communicate_values(exec_handler, diag_comm_pattern, host_A->get_diag(),
+                           local_ptr);
+        communicate_values(exec_handler, upper_comm_pattern,
+                           host_A->get_upper(), local_ptr);
+
+        if (host_A->symmetric()) {
+            // TODO FIXME
+            // if symmetric we can reuse already copied data
+            communicate_values(exec_handler, lower_comm_pattern,
+                               host_A->get_lower(), local_ptr);
+        } else {
+            communicate_values(exec_handler, lower_comm_pattern,
+                               host_A->get_lower(), local_ptr);
+        }
+
+        if (requires_host_buffer) {
+            auto host_buffer_view =
+                gko::array<scalar>::view(exec, nnz, local_ptr);
+            auto target_buffer_view =
+                gko::array<scalar>::view(device_exec, nnz, local_ptr_2);
+            target_buffer_view = host_buffer_view;
+        }
+
+        if (requires_host_buffer) {
+            auto host_buffer_view =
+                gko::array<scalar>::view(exec, nnz, local_ptr);
+            auto target_buffer_view =
+                gko::array<scalar>::view(device_exec, nnz, local_ptr_2);
+            target_buffer_view = host_buffer_view;
+        }
+        // copy interface values
+        auto comm = *exec_handler.get_communicator().get();
+        if (owner) {
+            using Coo = gko::matrix::Coo<scalar, label>;
+            std::shared_ptr<const Coo> mtx;
+            label loc_ctr{1};
+            label nloc_ctr{0};
+            label host_interface_ctr{0};
+            label tag = 0;
+            label comm_rank, comm_size;
+            scalar *recv_buffer_ptr;
+            scalar *recv_buffer_ptr_2;
+            std::vector<scalar> host_recv_buffer;
+            label remain_host_interfaces = host_A->get_interface_size();
+            for (auto [is_local, comm_rank] : local_interfaces) {
+                label &ctr = (is_local) ? loc_ctr : nloc_ctr;
+                if (is_local) {
+                    mtx = gko::as<Coo>(
+                        gko::as<CombinationMatrix<scalar, label, Coo>>(
+                            dist_A->get_local_matrix())
+                            ->get_combination()
+                            ->get_operators()[ctr]);
+                    comm_size = local_sparsity->interface_spans[ctr].length();
+                } else {
+                    mtx = gko::as<Coo>(
+                        gko::as<CombinationMatrix<scalar, label, Coo>>(
+                            dist_A->get_non_local_matrix())
+                            ->get_combination()
+                            ->get_operators()[ctr]);
+                    comm_size =
+                        non_local_sparsity->interface_spans[ctr].length();
+                }
+
+                if (requires_host_buffer) {
+                    host_recv_buffer.resize(comm_size);
+                    recv_buffer_ptr = host_recv_buffer.data();
+                    recv_buffer_ptr_2 =
+                        const_cast<scalar *>(mtx->get_const_values());
+                } else {
+                    recv_buffer_ptr =
+                        const_cast<scalar *>(mtx->get_const_values());
+                }
+
+                if (comm_rank != rank) {
+                    comm.recv(device_exec, recv_buffer_ptr, comm_size,
+                              comm_rank, tag);
+                    if (requires_host_buffer) {
+                        auto host_buffer_view = gko::array<scalar>::view(
+                            exec, comm_size, recv_buffer_ptr);
+                        auto target_buffer_view = gko::array<scalar>::view(
+                            device_exec, comm_size, recv_buffer_ptr_2);
+                        target_buffer_view = host_buffer_view;
+                    }
+
+                } else {
+                    // if data is already on this rank
+                    auto data_view = gko::array<scalar>::const_view(
+                        exec, comm_size,
+                        host_A->get_interface_data(host_interface_ctr));
+
+                    // TODO FIXME this needs target executor
+                    recv_buffer_ptr =
+                        const_cast<scalar *>(mtx->get_const_values());
+                    auto target_view = gko::array<scalar>::view(
+                        device_exec, comm_size, recv_buffer_ptr);
+
+                    target_view = data_view;
+
+                    host_interface_ctr++;
+                    remain_host_interfaces--;
+                }
+
+                ctr++;
+                // interface values need to be multiplied by -1
+                using vec = gko::matrix::Dense<ValueType>;
+                recv_buffer_ptr = const_cast<scalar *>(mtx->get_const_values());
+                auto neg_one = gko::initialize<vec>({-1.0}, device_exec);
+                auto interface_dense =
+                    vec::create(device_exec, gko::dim<2>{comm_size, 1},
+                                gko::array<scalar>::view(device_exec, comm_size,
+                                                         recv_buffer_ptr),
+                                1);
+
+                interface_dense->scale(neg_one);
+            }
+        } else {
+            // the non owner has send all its interfaces to owner
+            // thus all values need to be communicated to the owner as well
+            label num_interfaces = src_comm_pattern->target_ids.get_size();
+            label tag = 0;
+            for (int i = 0; i < num_interfaces; i++) {
+                label comm_size =
+                    src_comm_pattern->target_sizes.get_const_data()[i];
+                const scalar *send_buffer_ptr = host_A->get_interface_data(i);
+                comm.send(device_exec, send_buffer_ptr, comm_size, owner_rank,
+                          tag);
+            }
+        }
+        // reorder updated values
+        if (owner) {
+            // NOTE local sparsity size includes the interfaces
+            using Coo = gko::matrix::Coo<scalar, label>;
+            using dim_type = gko::dim<2>::dimension_type;
+            std::shared_ptr<const Coo> local =
+                gko::as<Coo>(gko::as<CombinationMatrix<scalar, label, Coo>>(
+                                 dist_A->get_local_matrix())
+                                 ->get_combination()
+                                 ->get_operators()[0]);
+            auto local_elements = local->get_num_stored_elements();
+            local_ptr = const_cast<scalar *>(local->get_const_values());
+            // TODO make sure this doesn't copy
+            // create a non owning dense matrix of local_values
+
+            auto row_collection = gko::share(gko::matrix::Dense<scalar>::create(
+                device_exec,
+                gko::dim<2>{static_cast<dim_type>(local_elements), 1},
+                gko::array<scalar>::view(device_exec, local_elements,
+                                         local_ptr),
+                1));
+            auto mapping_view = gko::array<label>::view(
+                exec, local_elements, local_sparsity->ldu_mapping.get_data());
+
+
+            // TODO this needs to copy ldu_mapping to the device
+            auto dense_vec = row_collection->clone();
+            // auto dense_vec =
+            // gko::share(gko::matrix::Dense<scalar>::create(exec,
+            // row_collection->get_size()));
+
+            dense_vec->row_gather(&mapping_view, row_collection.get());
+        }
+    };
+
+    RepartDistMatrix(
+        std::shared_ptr<const gko::Executor> exec, communicator comm,
+        gko::dim<2> local_size, gko::dim<2> global_size,
+        std::shared_ptr<dist_mtx> dist_mtx,
+        std::shared_ptr<const SparsityPattern> local_sparsity,
+        std::shared_ptr<const SparsityPattern> non_local_sparsity,
+        std::shared_ptr<const CommunicationPattern> src_comm_pattern,
+        std::vector<std::pair<bool, label>> &local_interfaces)
+        : gko::experimental::EnableDistributedLinOp<RepartDistMatrix>(exec),
+          gko::experimental::distributed::DistributedBase(comm),
+          dist_mtx_(std::move(dist_mtx)),
+          local_sparsity_(local_sparsity),
+          non_local_sparsity_(non_local_sparsity),
+          src_comm_pattern_(src_comm_pattern),
+          local_interfaces_(local_interfaces)
+    {
+        this->set_size(global_size);
+    }
+
+    void write(const word field_name_, const objectRegistry &db_) const
+    {
+        auto dist_matrix = gko::as<
+            gko::experimental::distributed::Matrix<scalar, label, label>>(
+            this->dist_mtx_.get());
+        // TODO FIXME dont leave hardcoded value here
+        word matrix_format_{"Coo"};
+
+        if (matrix_format_ == "Coo") {
+            using Coo = gko::matrix::Coo<scalar, label>;
+            std::vector<std::shared_ptr<const gko::LinOp>> local_interfaces =
+                gko::as<CombinationMatrix<scalar, label, Coo>>(
+                    dist_matrix->get_local_matrix())
+                    ->get_combination()
+                    ->get_operators();
+            export_mtx<Coo>(word(field_name_ + "_local"), local_interfaces,
+                            db_);
+            auto non_local_interfaces =
+                gko::as<CombinationMatrix<scalar, label, Coo>>(
+                    dist_matrix->get_non_local_matrix())
+                    ->get_combination()
+                    ->get_operators();
+            export_mtx<Coo>(field_name_ + "_non_local", non_local_interfaces,
+                            db_);
+        }
+
+        if (matrix_format_ == "Csr") {
+            using Csr = gko::matrix::Csr<scalar, label>;
+            auto local_interfaces =
+                gko::as<CombinationMatrix<scalar, label, Csr>>(
+                    dist_matrix->get_local_matrix().get())
+                    ->get_combination()
+                    ->get_operators();
+            auto non_local_interfaces =
+                gko::as<CombinationMatrix<scalar, label, Csr>>(
+                    dist_matrix->get_non_local_matrix().get())
+                    ->get_combination()
+                    ->get_operators();
+            export_mtx<Csr>(word(field_name_ + "_local"), local_interfaces,
+                            db_);
+            export_mtx<Csr>(field_name_ + "_non_local", non_local_interfaces,
+                            db_);
+        }
+
+        if (matrix_format_ == "Ell") {
+            using Ell = gko::matrix::Ell<scalar, label>;
+            auto local_interfaces =
+                gko::as<CombinationMatrix<scalar, label, Ell>>(
+                    dist_matrix->get_local_matrix().get())
+                    ->get_combination()
+                    ->get_operators();
+            auto non_local_interfaces =
+                gko::as<CombinationMatrix<scalar, label, Ell>>(
+                    dist_matrix->get_non_local_matrix().get())
+                    ->get_combination()
+                    ->get_operators();
+            export_mtx<Ell>(field_name_ + "_local", local_interfaces, db_);
+            export_mtx<Ell>(field_name_ + "_non_local", non_local_interfaces,
+                            db_);
+        }
+    }
+
+    // Needed for distributed/polymorphic_object.hpp
+    RepartDistMatrix(std::shared_ptr<const gko::Executor> exec,
+                     communicator comm)
+        : gko::experimental::EnableDistributedLinOp<RepartDistMatrix>(exec),
+          gko::experimental::distributed::DistributedBase{comm}
+    {}
+
+protected:
+    // Here we implement the application of the linear operator, x = A * b.
+    // apply_impl will be called by the apply method, after the arguments
+    // have been moved to the correct executor and the operators checked for
+    // conforming sizes.
+    //
+    // For simplicity, we assume that there is always only one right hand
+    // side and the stride of consecutive elements in the vectors is 1 (both
+    // of these are always true in this example).
+    void apply_impl(const gko::LinOp *b, gko::LinOp *x) const override
+    {
+        this->dist_mtx_->apply(b, x);
+    }
+
+
+    // There is also a version of the apply function which does the
+    // operation x = alpha * A * b + beta * x. This function is commonly
+    // used and can often be better optimized than implementing it using x =
+    // A * b. However, for simplicity, we will implement it exactly like
+    // that in this example.
+    void apply_impl(const gko::LinOp *alpha, const gko::LinOp *b,
+                    const gko::LinOp *beta, gko::LinOp *x) const override
+    {
+        this->dist_mtx_->apply(alpha, b, beta, x);
+    }
+
+
+private:
+    std::shared_ptr<dist_mtx> dist_mtx_;
+
+    std::shared_ptr<const SparsityPattern> local_sparsity_;
+
+    std::shared_ptr<const SparsityPattern> non_local_sparsity_;
+
+    std::shared_ptr<const CommunicationPattern> src_comm_pattern_;
+
+    std::vector<std::pair<bool, label>> local_interfaces_;
+};

From fa04ac864b4ce0f8a8519f7c8cf683d196fbebe0 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Thu, 29 Aug 2024 09:43:23 +0200
Subject: [PATCH 56/67] make it compile

---
 CMakeLists.txt                               |   1 +
 cmake/ginkgo.cmake                           |   1 +
 include/OGL/DevicePersistent/DeviceIdGuard.H |   5 +-
 include/OGL/DevicePersistent/Partition.H     |  21 +-
 include/OGL/MatrixWrapper/Distributed.H      | 743 +++++++++----------
 include/OGL/MatrixWrapper/HostMatrix.H       |   3 +-
 include/OGL/MatrixWrapper/SparsityPattern.H  |  15 +-
 include/OGL/Repartitioner.H                  |  32 +-
 src/CommunicationPattern.C                   |   2 -
 src/MatrixWrapper/Distributed.C              |  73 ++
 src/MatrixWrapper/HostMatrix.C               |  66 +-
 src/Repartitioner.C                          |  44 +-
 unitTests/MatrixWrapper/HostMatrix.C         |  12 +-
 unitTests/Repartitioner.C                    |   8 +-
 14 files changed, 530 insertions(+), 496 deletions(-)
 create mode 100644 src/MatrixWrapper/Distributed.C

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 0d7dfd81a..c90fbdf6a 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -223,6 +223,7 @@ target_sources(
           src/DevicePersistent/ExecutorHandler/ExecutorHandler.C
           src/MatrixWrapper/HostMatrix.C
           src/MatrixWrapper/SparsityPattern.C
+          src/MatrixWrapper/Distributed.C
           src/Repartitioner.C
   # src/Solver/CG/GKOCG.C src/Solver/BiCGStab/GKOBiCGStab.C
   # src/Solver/GMRES/GKOGMRES.C
diff --git a/cmake/ginkgo.cmake b/cmake/ginkgo.cmake
index 47923c4f9..f0bd096e4 100644
--- a/cmake/ginkgo.cmake
+++ b/cmake/ginkgo.cmake
@@ -7,6 +7,7 @@ include(FetchContent)
 if(NOT ${OGL_USE_EXTERNAL_GINKGO})
   FetchContent_Declare(
     Ginkgo
+    SYSTEM
     QUITE
     GIT_SHALLOW ON
     GIT_REPOSITORY "https://github.com/ginkgo-project/ginkgo.git"
diff --git a/include/OGL/DevicePersistent/DeviceIdGuard.H b/include/OGL/DevicePersistent/DeviceIdGuard.H
index 4ab42c188..8b094fa5c 100644
--- a/include/OGL/DevicePersistent/DeviceIdGuard.H
+++ b/include/OGL/DevicePersistent/DeviceIdGuard.H
@@ -26,7 +26,8 @@ struct DeviceIdGuardInitFunctor {
     std::shared_ptr<gko::scoped_device_id_guard> init() const
     {
         return std::make_shared<gko::scoped_device_id_guard>(
-            std::move(executor_->get_scoped_device_id_guard()));
+            executor_->get_scoped_device_id_guard());
+        // std::move(executor_->get_scoped_device_id_guard()));
     }
 };
 
@@ -41,6 +42,4 @@ public:
     {}
 };
 
-
 }  // namespace Foam
-   // namespace Foam
diff --git a/include/OGL/DevicePersistent/Partition.H b/include/OGL/DevicePersistent/Partition.H
index f2c502c7c..afa8615d5 100644
--- a/include/OGL/DevicePersistent/Partition.H
+++ b/include/OGL/DevicePersistent/Partition.H
@@ -55,14 +55,18 @@ struct PartitionInitFunctor {
 
     std::shared_ptr<part_type> init() const
     {
-        auto exec = exec_.get_ref_exec();
-        auto comm = exec_.get_gko_mpi_host_comm();
+        FatalErrorInFunction << "Not implemented" << abort(FatalError);
+        // auto exec = exec_.get_ref_exec();
+        // auto comm = exec_.get_gko_mpi_host_comm();
 
-        word msg{"initialising partition of size " +
-                 std::to_string(local_size_)};
-        LOG_1(verbose_, msg)
+        // word msg{"initialising partition of size " +
+        //          std::to_string(local_size_)};
+        // LOG_1(verbose_, msg)
 
-        if (ranks_per_gpu_ == 1) return {};
+        // if (ranks_per_gpu_ == 1) return {};
+
+
+        return {};
     }
 };
 
@@ -84,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
@@ -107,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/include/OGL/MatrixWrapper/Distributed.H b/include/OGL/MatrixWrapper/Distributed.H
index ac760c597..3fdf7aab0 100644
--- a/include/OGL/MatrixWrapper/Distributed.H
+++ b/include/OGL/MatrixWrapper/Distributed.H
@@ -9,61 +9,18 @@
 #include <ginkgo/core/distributed/lin_op.hpp>
 #include <ginkgo/ginkgo.hpp>
 
-
-#include "MatrixWrapper/CommunicationPattern/CommunicationPattern.H"
-#include "MatrixWrapper/LDUMatrix/HostMatrix.H"
-#include "Repartitioner/Repartitioner.H"
-
-
-template <typename DistMtxType>
-std::unique_ptr<DistMtxType> generate_dist_mtx_with_inner_type(
-    word matrix_format, std::shared_ptr<const gko::Executor> exec,
-    gko::experimental::mpi::communicator comm,
-    std::shared_ptr<const SparsityPattern> local_sparsity,
-    std::shared_ptr<const SparsityPattern> non_local_sparsity)
-{
-    auto local_dim = local_sparsity->dim;
-    auto non_local_dim = non_local_sparsity->dim;
-    auto local_interfaces = local_sparsity->interface_spans;
-    auto non_local_interfaces = non_local_sparsity->interface_spans;
-
-    if (matrix_format == "Ell") {
-        using mtx_type = gko::matrix::Ell<scalar, label>;
-
-        return DistMtxType::create(
-            exec, comm,
-            CombinationMatrix<scalar, label, mtx_type>::create(
-                exec, local_dim, local_interfaces),
-            CombinationMatrix<scalar, label, mtx_type>::create(
-                exec, non_local_dim, non_local_interfaces));
-    }
-    if (matrix_format == "Csr") {
-        using mtx_type = gko::matrix::Csr<scalar, label>;
-
-        return DistMtxType::create(
-            exec, comm,
-            CombinationMatrix<scalar, label, mtx_type>::create(
-                exec, local_dim, local_interfaces),
-            CombinationMatrix<scalar, label, mtx_type>::create(
-                exec, non_local_dim, non_local_interfaces));
-    }
-    if (matrix_format == "Coo") {
-        using mtx_type = gko::matrix::Coo<scalar, label>;
-
-        return DistMtxType::create(
-            exec, comm,
-            CombinationMatrix<scalar, label, mtx_type>::create(
-                exec, local_dim, local_interfaces),
-            CombinationMatrix<scalar, label, mtx_type>::create(
-                exec, non_local_dim, non_local_interfaces));
-    }
-
-    FatalErrorInFunction << "Matrix format " << matrix_format
-                         << " not supported! Supported formats: Csr, Ell, Coo"
-                         << abort(FatalError);
-    return {};
-}
-
+#include "OGL/CommunicationPattern.H"
+#include "OGL/MatrixWrapper/HostMatrix.H"
+#include "OGL/Repartitioner.H"
+
+namespace detail {
+/* @brief generates a vector of ginkgo linops
+**
+*/
+std::vector<std::unique_ptr<gko::LinOp>> generate_inner_linops(
+    word matrix_format, bool fuse, std::shared_ptr<const gko::Executor> exec,
+    std::shared_ptr<const SparsityPattern> sparsity);
+}  // namespace detail
 
 // The RepartDistMatrix class is a wrapper around Ginkgos distributed Matrix
 // class
@@ -164,77 +121,67 @@ public:
         // device with correct size on each rank
         //
         // NOTE Start by copying first and repartition later
-        auto local_sparsity_ = host_A->compute_local_sparsity(exec);
-        auto non_local_sparsity_ = host_A->compute_non_local_sparsity(exec);
-
-        auto orig_partition = repartitioner.get_orig_partition();
-
-        // create original communicator pattern
-        auto ranks_per_gpu = repartitioner.get_ranks_per_gpu();
-        label rank{comm.rank()};
+        auto local_sparsity = host_A->compute_local_sparsity(exec);
+        auto non_local_sparsity = host_A->compute_non_local_sparsity(exec);
 
-        auto src_comm_pattern =
-            host_A.get()->create_communication_pattern(exec_handler);
-        // create partiton here and pass to constructor
-        //
+        auto src_comm_pattern = host_A->create_communication_pattern();
         auto dst_comm_pattern = repartitioner.repartition_comm_pattern(
-            exec_handler, src_comm_pattern, orig_partition);
+            exec_handler, src_comm_pattern);
 
-        label owner_rank = repartitioner.get_owner_rank(exec_handler);
-        bool owner = repartitioner.is_owner(exec_handler);
+        // bool owner = repartitioner.is_owner(exec_handler);
 
         auto [repart_loc_sparsity, repart_non_loc_sparsity, local_interfaces] =
-            repartitioner.repartition_sparsity(exec_handler, local_sparsity_,
-                                               non_local_sparsity_);
-
-        local_sparsity_ = repart_loc_sparsity;
-        non_local_sparsity_ = repart_non_loc_sparsity;
+            repartitioner.repartition_sparsity(exec_handler, local_sparsity,
+                                               non_local_sparsity);
 
+        // create vector of inner type linops
+        // if fuse the vector contains only a single element
+        // thus we can unwrap it.
 
         auto device_exec = exec_handler.get_device_exec();
-        auto dist_A = gko::share(generate_dist_mtx_with_inner_type<dist_mtx>(
-            matrix_format, device_exec, comm, local_sparsity_,
-            non_local_sparsity_));
-
-        auto local_coeffs = gko::array<scalar>(exec, local_sparsity_->size_);
-        auto non_local_coeffs =
-            gko::array<scalar>(exec, non_local_sparsity_->size_);
-
-
-        if (owner) {
-            local_coeffs.fill(0.0);
-            non_local_coeffs.fill(0.0);
-        }
-
-        std::cout << __FILE__ << " rank " << rank << " local sparsity size "
-                  << local_sparsity_->size_ << " local sparsity dim ["
-                  << local_sparsity_->dim[0] << "x" << local_sparsity_->dim[1]
-                  << "] non_local sparsity size " << non_local_sparsity_->size_
-                  << " non local sparsity dim [" << non_local_sparsity_->dim[0]
-                  << "x" << non_local_sparsity_->dim[1] << "] target_ids "
-                  << dst_comm_pattern->target_ids << " target_sizes "
-                  << dst_comm_pattern->target_sizes << " target_send_idxs.size "
-                  << dst_comm_pattern->send_idxs.size()
-                  << " non_local_sparsity.size " << non_local_sparsity_->size_
-                  << " \n";
-
-        // FIXME make sure that we work on the device executor
-        auto comm_pattern = compute_send_recv_counts(
-            exec_handler, 1, non_local_sparsity_->size_);
-
-        auto send_idxs = dst_comm_pattern->total_rank_send_idx();
-        auto comm_counts = dst_comm_pattern->send_recv_pattern();
-
-
-        dist_A->read_distributed(
-            device_matrix_data(exec, local_sparsity_->dim,
-                               local_sparsity_->row_idxs,
-                               local_sparsity_->col_idxs, local_coeffs),
-            device_matrix_data(exec, non_local_sparsity_->dim,
-                               non_local_sparsity_->row_idxs,
-                               non_local_sparsity_->col_idxs, non_local_coeffs),
-            std::get<0>(comm_counts), std::get<1>(comm_counts),
-            std::get<2>(comm_counts), std::get<3>(comm_counts), send_idxs);
+        auto local_linops = detail::generate_inner_linops(
+            matrix_format, false, device_exec, repart_loc_sparsity);
+
+        // auto local_coeffs =
+        //     gko::array<scalar>(exec, repart_local_sparsity->num_nnz);
+        // auto non_local_coeffs =
+        //     gko::array<scalar>(exec, repart_non_local_sparsity->num_nnz);
+
+
+        // std::cout << __FILE__ << " rank " << rank << " local sparsity size "
+        //           << local_sparsity_->size_ << " local sparsity dim ["
+        //           << local_sparsity_->dim[0] << "x" <<
+        //           local_sparsity_->dim[1]
+        //           << "] non_local sparsity size " <<
+        //           non_local_sparsity_->size_
+        //           << " non local sparsity dim [" <<
+        //           non_local_sparsity_->dim[0]
+        //           << "x" << non_local_sparsity_->dim[1] << "] target_ids "
+        //           << dst_comm_pattern->target_ids << " target_sizes "
+        //           << dst_comm_pattern->target_sizes << "
+        //           target_send_idxs.size "
+        //           << dst_comm_pattern->send_idxs.size()
+        //           << " non_local_sparsity.size " <<
+        //           non_local_sparsity_->size_
+        //           << " \n";
+
+        // // FIXME make sure that we work on the device executor
+        // auto comm_pattern = compute_send_recv_counts(
+        //     exec_handler, 1, non_local_sparsity_->size_);
+
+        // auto send_idxs = dst_comm_pattern->total_rank_send_idx();
+        // auto comm_counts = dst_comm_pattern->send_recv_pattern();
+
+        // dist_A->read_distributed(
+        //     device_matrix_data(exec, local_sparsity_->dim,
+        //                        local_sparsity_->row_idxs,
+        //                        local_sparsity_->col_idxs, local_coeffs),
+        //     device_matrix_data(exec, non_local_sparsity_->dim,
+        //                        non_local_sparsity_->row_idxs,
+        //                        non_local_sparsity_->col_idxs,
+        //                        non_local_coeffs),
+        //     std::get<0>(comm_counts), std::get<1>(comm_counts),
+        //     std::get<2>(comm_counts), std::get<3>(comm_counts), send_idxs);
 
 
         // auto dist_A = DistMtxType::create(
@@ -244,16 +191,17 @@ public:
         //        CombinationMatrix<scalar, label, mtx_type>::create(
         //            exec, non_local_dim, non_local_interfaces));
 
-        update_impl(exec_handler, matrix_format, repartitioner, host_A, dist_A,
-                    local_sparsity_, non_local_sparsity_, src_comm_pattern,
-                    local_interfaces);
+        // update_impl(exec_handler, matrix_format, repartitioner, host_A,
+        // dist_A,
+        //             repart_local_sparsity, repart_non_local_sparsity,
+        //             src_comm_pattern, local_interfaces);
 
-        auto ret = std::make_shared<RepartDistMatrix>(
-            exec, comm, repartitioner.get_repart_dim(), dist_A->get_size(),
-            std::move(dist_A), local_sparsity_, non_local_sparsity_,
-            src_comm_pattern, local_interfaces);
+        // auto ret = std::make_shared<RepartDistMatrix>(
+        //     exec, comm, repartitioner.get_repart_dim(), dist_A->get_size(),
+        //     std::move(dist_A), local_sparsity_, non_local_sparsity_,
+        //     src_comm_pattern, local_interfaces);
 
-        return ret;
+        return {};
     };
 
     void update(const ExecutorHandler &exec_handler, word matrix_format,
@@ -276,216 +224,223 @@ public:
         std::shared_ptr<const SparsityPattern> local_sparsity,
         std::shared_ptr<const SparsityPattern> non_local_sparsity,
         std::shared_ptr<const CommunicationPattern> src_comm_pattern,
-        std::vector<std::pair<bool, label>> local_interfaces)
-    {
-        auto exec = exec_handler.get_ref_exec();
-        auto device_exec = exec_handler.get_device_exec();
-        auto ranks_per_gpu = repartitioner.get_ranks_per_gpu();
-        bool requires_host_buffer = exec_handler.get_gko_force_host_buffer();
-
-        label rank{repartitioner.get_rank(exec_handler)};
-        label owner_rank = repartitioner.get_owner_rank(exec_handler);
-        bool owner = repartitioner.is_owner(exec_handler);
-        label nrows = host_A->get_local_nrows();
-        label local_matrix_nnz = host_A->get_local_matrix_nnz();
-        // size + padding has to be local_matrix_nnz
-        auto diag_comm_pattern = compute_send_recv_counts(
-            exec_handler, ranks_per_gpu, nrows, local_matrix_nnz,
-            local_matrix_nnz - nrows, 0);
-        label upper_nnz = host_A->get_upper_nnz();
-        auto upper_comm_pattern = compute_send_recv_counts(
-            exec_handler, ranks_per_gpu, upper_nnz, local_matrix_nnz, 0,
-            local_matrix_nnz - upper_nnz);
-        auto lower_comm_pattern =
-            compute_send_recv_counts(exec_handler, ranks_per_gpu, upper_nnz,
-                                     local_matrix_nnz, upper_nnz, nrows);
-
-        scalar *local_ptr;
-        scalar *local_ptr_2;
-        label nnz = 0;
-
-        // update main values
-        std::vector<scalar> loc_buffer;
-        if (owner) {
-            using Coo = gko::matrix::Coo<scalar, label>;
-            auto local_mtx = dist_A->get_local_matrix();
-
-
-            std::shared_ptr<const Coo> local =
-                gko::as<Coo>(gko::as<CombinationMatrix<scalar, label, Coo>>(
-                                 dist_A->get_local_matrix())
-                                 ->get_combination()
-                                 ->get_operators()[0]);
-            nnz = local->get_num_stored_elements();
-            if (requires_host_buffer) {
-                loc_buffer.resize(nnz);
-                local_ptr = loc_buffer.data();
-                local_ptr_2 = const_cast<scalar *>(local->get_const_values());
-            } else {
-                local_ptr = const_cast<scalar *>(local->get_const_values());
-            }
-        }
-        communicate_values(exec_handler, diag_comm_pattern, host_A->get_diag(),
-                           local_ptr);
-        communicate_values(exec_handler, upper_comm_pattern,
-                           host_A->get_upper(), local_ptr);
-
-        if (host_A->symmetric()) {
-            // TODO FIXME
-            // if symmetric we can reuse already copied data
-            communicate_values(exec_handler, lower_comm_pattern,
-                               host_A->get_lower(), local_ptr);
-        } else {
-            communicate_values(exec_handler, lower_comm_pattern,
-                               host_A->get_lower(), local_ptr);
-        }
-
-        if (requires_host_buffer) {
-            auto host_buffer_view =
-                gko::array<scalar>::view(exec, nnz, local_ptr);
-            auto target_buffer_view =
-                gko::array<scalar>::view(device_exec, nnz, local_ptr_2);
-            target_buffer_view = host_buffer_view;
-        }
-
-        if (requires_host_buffer) {
-            auto host_buffer_view =
-                gko::array<scalar>::view(exec, nnz, local_ptr);
-            auto target_buffer_view =
-                gko::array<scalar>::view(device_exec, nnz, local_ptr_2);
-            target_buffer_view = host_buffer_view;
-        }
-        // copy interface values
-        auto comm = *exec_handler.get_communicator().get();
-        if (owner) {
-            using Coo = gko::matrix::Coo<scalar, label>;
-            std::shared_ptr<const Coo> mtx;
-            label loc_ctr{1};
-            label nloc_ctr{0};
-            label host_interface_ctr{0};
-            label tag = 0;
-            label comm_rank, comm_size;
-            scalar *recv_buffer_ptr;
-            scalar *recv_buffer_ptr_2;
-            std::vector<scalar> host_recv_buffer;
-            label remain_host_interfaces = host_A->get_interface_size();
-            for (auto [is_local, comm_rank] : local_interfaces) {
-                label &ctr = (is_local) ? loc_ctr : nloc_ctr;
-                if (is_local) {
-                    mtx = gko::as<Coo>(
-                        gko::as<CombinationMatrix<scalar, label, Coo>>(
-                            dist_A->get_local_matrix())
-                            ->get_combination()
-                            ->get_operators()[ctr]);
-                    comm_size = local_sparsity->interface_spans[ctr].length();
-                } else {
-                    mtx = gko::as<Coo>(
-                        gko::as<CombinationMatrix<scalar, label, Coo>>(
-                            dist_A->get_non_local_matrix())
-                            ->get_combination()
-                            ->get_operators()[ctr]);
-                    comm_size =
-                        non_local_sparsity->interface_spans[ctr].length();
-                }
-
-                if (requires_host_buffer) {
-                    host_recv_buffer.resize(comm_size);
-                    recv_buffer_ptr = host_recv_buffer.data();
-                    recv_buffer_ptr_2 =
-                        const_cast<scalar *>(mtx->get_const_values());
-                } else {
-                    recv_buffer_ptr =
-                        const_cast<scalar *>(mtx->get_const_values());
-                }
-
-                if (comm_rank != rank) {
-                    comm.recv(device_exec, recv_buffer_ptr, comm_size,
-                              comm_rank, tag);
-                    if (requires_host_buffer) {
-                        auto host_buffer_view = gko::array<scalar>::view(
-                            exec, comm_size, recv_buffer_ptr);
-                        auto target_buffer_view = gko::array<scalar>::view(
-                            device_exec, comm_size, recv_buffer_ptr_2);
-                        target_buffer_view = host_buffer_view;
-                    }
-
-                } else {
-                    // if data is already on this rank
-                    auto data_view = gko::array<scalar>::const_view(
-                        exec, comm_size,
-                        host_A->get_interface_data(host_interface_ctr));
-
-                    // TODO FIXME this needs target executor
-                    recv_buffer_ptr =
-                        const_cast<scalar *>(mtx->get_const_values());
-                    auto target_view = gko::array<scalar>::view(
-                        device_exec, comm_size, recv_buffer_ptr);
-
-                    target_view = data_view;
-
-                    host_interface_ctr++;
-                    remain_host_interfaces--;
-                }
-
-                ctr++;
-                // interface values need to be multiplied by -1
-                using vec = gko::matrix::Dense<ValueType>;
-                recv_buffer_ptr = const_cast<scalar *>(mtx->get_const_values());
-                auto neg_one = gko::initialize<vec>({-1.0}, device_exec);
-                auto interface_dense =
-                    vec::create(device_exec, gko::dim<2>{comm_size, 1},
-                                gko::array<scalar>::view(device_exec, comm_size,
-                                                         recv_buffer_ptr),
-                                1);
-
-                interface_dense->scale(neg_one);
-            }
-        } else {
-            // the non owner has send all its interfaces to owner
-            // thus all values need to be communicated to the owner as well
-            label num_interfaces = src_comm_pattern->target_ids.get_size();
-            label tag = 0;
-            for (int i = 0; i < num_interfaces; i++) {
-                label comm_size =
-                    src_comm_pattern->target_sizes.get_const_data()[i];
-                const scalar *send_buffer_ptr = host_A->get_interface_data(i);
-                comm.send(device_exec, send_buffer_ptr, comm_size, owner_rank,
-                          tag);
-            }
-        }
-        // reorder updated values
-        if (owner) {
-            // NOTE local sparsity size includes the interfaces
-            using Coo = gko::matrix::Coo<scalar, label>;
-            using dim_type = gko::dim<2>::dimension_type;
-            std::shared_ptr<const Coo> local =
-                gko::as<Coo>(gko::as<CombinationMatrix<scalar, label, Coo>>(
-                                 dist_A->get_local_matrix())
-                                 ->get_combination()
-                                 ->get_operators()[0]);
-            auto local_elements = local->get_num_stored_elements();
-            local_ptr = const_cast<scalar *>(local->get_const_values());
-            // TODO make sure this doesn't copy
-            // create a non owning dense matrix of local_values
-
-            auto row_collection = gko::share(gko::matrix::Dense<scalar>::create(
-                device_exec,
-                gko::dim<2>{static_cast<dim_type>(local_elements), 1},
-                gko::array<scalar>::view(device_exec, local_elements,
-                                         local_ptr),
-                1));
-            auto mapping_view = gko::array<label>::view(
-                exec, local_elements, local_sparsity->ldu_mapping.get_data());
-
-
-            // TODO this needs to copy ldu_mapping to the device
-            auto dense_vec = row_collection->clone();
-            // auto dense_vec =
-            // gko::share(gko::matrix::Dense<scalar>::create(exec,
-            // row_collection->get_size()));
-
-            dense_vec->row_gather(&mapping_view, row_collection.get());
-        }
+        std::vector<std::pair<bool, label>> local_interfaces){
+        // auto exec = exec_handler.get_ref_exec();
+        // auto device_exec = exec_handler.get_device_exec();
+        // auto ranks_per_gpu = repartitioner.get_ranks_per_gpu();
+        // bool requires_host_buffer = exec_handler.get_gko_force_host_buffer();
+
+        // label rank{exec_handler.get_rank()};
+        // label owner_rank = repartitioner.get_owner_rank(exec_handler);
+        // bool owner = repartitioner.is_owner(exec_handler);
+        // label nrows = host_A->get_local_nrows();
+        // label local_matrix_nnz = host_A->get_local_matrix_nnz();
+        // // size + padding has to be local_matrix_nnz
+        // auto diag_comm_pattern = compute_send_recv_counts(
+        //     exec_handler, ranks_per_gpu, nrows, local_matrix_nnz,
+        //     local_matrix_nnz - nrows, 0);
+        // label upper_nnz = host_A->get_upper_nnz();
+        // auto upper_comm_pattern = compute_send_recv_counts(
+        //     exec_handler, ranks_per_gpu, upper_nnz, local_matrix_nnz, 0,
+        //     local_matrix_nnz - upper_nnz);
+        // auto lower_comm_pattern =
+        //     compute_send_recv_counts(exec_handler, ranks_per_gpu, upper_nnz,
+        //                              local_matrix_nnz, upper_nnz, nrows);
+
+        // scalar *local_ptr;
+        // scalar *local_ptr_2;
+        // label nnz = 0;
+
+        // // update main values
+        // std::vector<scalar> loc_buffer;
+        // if (owner) {
+        //     using Coo = gko::matrix::Coo<scalar, label>;
+        //     auto local_mtx = dist_A->get_local_matrix();
+
+
+        //     std::shared_ptr<const Coo> local =
+        //         gko::as<Coo>(gko::as<CombinationMatrix<scalar, label, Coo>>(
+        //                          dist_A->get_local_matrix())
+        //                          ->get_combination()
+        //                          ->get_operators()[0]);
+        //     nnz = local->get_num_stored_elements();
+        //     if (requires_host_buffer) {
+        //         loc_buffer.resize(nnz);
+        //         local_ptr = loc_buffer.data();
+        //         local_ptr_2 = const_cast<scalar
+        //         *>(local->get_const_values());
+        //     } else {
+        //         local_ptr = const_cast<scalar *>(local->get_const_values());
+        //     }
+        // }
+        // communicate_values(exec_handler, diag_comm_pattern,
+        // host_A->get_diag(),
+        //                    local_ptr);
+        // communicate_values(exec_handler, upper_comm_pattern,
+        //                    host_A->get_upper(), local_ptr);
+
+        // if (host_A->symmetric()) {
+        //     // TODO FIXME
+        //     // if symmetric we can reuse already copied data
+        //     communicate_values(exec_handler, lower_comm_pattern,
+        //                        host_A->get_lower(), local_ptr);
+        // } else {
+        //     communicate_values(exec_handler, lower_comm_pattern,
+        //                        host_A->get_lower(), local_ptr);
+        // }
+
+        // if (requires_host_buffer) {
+        //     auto host_buffer_view =
+        //         gko::array<scalar>::view(exec, nnz, local_ptr);
+        //     auto target_buffer_view =
+        //         gko::array<scalar>::view(device_exec, nnz, local_ptr_2);
+        //     target_buffer_view = host_buffer_view;
+        // }
+
+        // if (requires_host_buffer) {
+        //     auto host_buffer_view =
+        //         gko::array<scalar>::view(exec, nnz, local_ptr);
+        //     auto target_buffer_view =
+        //         gko::array<scalar>::view(device_exec, nnz, local_ptr_2);
+        //     target_buffer_view = host_buffer_view;
+        // }
+        // // copy interface values
+        // auto comm = *exec_handler.get_communicator().get();
+        // if (owner) {
+        //     using Coo = gko::matrix::Coo<scalar, label>;
+        //     std::shared_ptr<const Coo> mtx;
+        //     label loc_ctr{1};
+        //     label nloc_ctr{0};
+        //     label host_interface_ctr{0};
+        //     label tag = 0;
+        //     label comm_rank, comm_size;
+        //     scalar *recv_buffer_ptr;
+        //     scalar *recv_buffer_ptr_2;
+        //     std::vector<scalar> host_recv_buffer;
+        //     label remain_host_interfaces = host_A->get_interface_size();
+        //     for (auto [is_local, comm_rank] : local_interfaces) {
+        //         label &ctr = (is_local) ? loc_ctr : nloc_ctr;
+        //         if (is_local) {
+        //             mtx = gko::as<Coo>(
+        //                 gko::as<CombinationMatrix<scalar, label, Coo>>(
+        //                     dist_A->get_local_matrix())
+        //                     ->get_combination()
+        //                     ->get_operators()[ctr]);
+        //             comm_size =
+        //             local_sparsity->interface_spans[ctr].length();
+        //         } else {
+        //             mtx = gko::as<Coo>(
+        //                 gko::as<CombinationMatrix<scalar, label, Coo>>(
+        //                     dist_A->get_non_local_matrix())
+        //                     ->get_combination()
+        //                     ->get_operators()[ctr]);
+        //             comm_size =
+        //                 non_local_sparsity->interface_spans[ctr].length();
+        //         }
+
+        //         if (requires_host_buffer) {
+        //             host_recv_buffer.resize(comm_size);
+        //             recv_buffer_ptr = host_recv_buffer.data();
+        //             recv_buffer_ptr_2 =
+        //                 const_cast<scalar *>(mtx->get_const_values());
+        //         } else {
+        //             recv_buffer_ptr =
+        //                 const_cast<scalar *>(mtx->get_const_values());
+        //         }
+
+        //         if (comm_rank != rank) {
+        //             comm.recv(device_exec, recv_buffer_ptr, comm_size,
+        //                       comm_rank, tag);
+        //             if (requires_host_buffer) {
+        //                 auto host_buffer_view = gko::array<scalar>::view(
+        //                     exec, comm_size, recv_buffer_ptr);
+        //                 auto target_buffer_view = gko::array<scalar>::view(
+        //                     device_exec, comm_size, recv_buffer_ptr_2);
+        //                 target_buffer_view = host_buffer_view;
+        //             }
+
+        //         } else {
+        //             // if data is already on this rank
+        //             auto data_view = gko::array<scalar>::const_view(
+        //                 exec, comm_size,
+        //                 host_A->get_interface_data(host_interface_ctr));
+
+        //             // TODO FIXME this needs target executor
+        //             recv_buffer_ptr =
+        //                 const_cast<scalar *>(mtx->get_const_values());
+        //             auto target_view = gko::array<scalar>::view(
+        //                 device_exec, comm_size, recv_buffer_ptr);
+
+        //             target_view = data_view;
+
+        //             host_interface_ctr++;
+        //             remain_host_interfaces--;
+        //         }
+
+        //         ctr++;
+        //         // interface values need to be multiplied by -1
+        //         using vec = gko::matrix::Dense<ValueType>;
+        //         recv_buffer_ptr = const_cast<scalar
+        //         *>(mtx->get_const_values()); auto neg_one =
+        //         gko::initialize<vec>({-1.0}, device_exec); auto
+        //         interface_dense =
+        //             vec::create(device_exec, gko::dim<2>{comm_size, 1},
+        //                         gko::array<scalar>::view(device_exec,
+        //                         comm_size,
+        //                                                  recv_buffer_ptr),
+        //                         1);
+
+        //         interface_dense->scale(neg_one);
+        //     }
+        // } else {
+        //     // the non owner has send all its interfaces to owner
+        //     // thus all values need to be communicated to the owner as well
+        //     label num_interfaces = src_comm_pattern->target_ids.get_size();
+        //     label tag = 0;
+        //     for (int i = 0; i < num_interfaces; i++) {
+        //         label comm_size =
+        //             src_comm_pattern->target_sizes.get_const_data()[i];
+        //         const scalar *send_buffer_ptr =
+        //         host_A->get_interface_data(i); comm.send(device_exec,
+        //         send_buffer_ptr, comm_size, owner_rank,
+        //                   tag);
+        //     }
+        // }
+        // // reorder updated values
+        // if (owner) {
+        //     // NOTE local sparsity size includes the interfaces
+        //     using Coo = gko::matrix::Coo<scalar, label>;
+        //     using dim_type = gko::dim<2>::dimension_type;
+        //     std::shared_ptr<const Coo> local =
+        //         gko::as<Coo>(gko::as<CombinationMatrix<scalar, label, Coo>>(
+        //                          dist_A->get_local_matrix())
+        //                          ->get_combination()
+        //                          ->get_operators()[0]);
+        //     auto local_elements = local->get_num_stored_elements();
+        //     local_ptr = const_cast<scalar *>(local->get_const_values());
+        //     // TODO make sure this doesn't copy
+        //     // create a non owning dense matrix of local_values
+
+        //     auto row_collection =
+        //     gko::share(gko::matrix::Dense<scalar>::create(
+        //         device_exec,
+        //         gko::dim<2>{static_cast<dim_type>(local_elements), 1},
+        //         gko::array<scalar>::view(device_exec, local_elements,
+        //                                  local_ptr),
+        //         1));
+        //     auto mapping_view = gko::array<label>::view(
+        //         exec, local_elements,
+        //         local_sparsity->ldu_mapping.get_data());
+
+
+        //     // TODO this needs to copy ldu_mapping to the device
+        //     auto dense_vec = row_collection->clone();
+        //     // auto dense_vec =
+        //     // gko::share(gko::matrix::Dense<scalar>::create(exec,
+        //     // row_collection->get_size()));
+
+        //     dense_vec->row_gather(&mapping_view, row_collection.get());
+        // }
     };
 
     RepartDistMatrix(
@@ -509,64 +464,64 @@ public:
 
     void write(const word field_name_, const objectRegistry &db_) const
     {
-        auto dist_matrix = gko::as<
-            gko::experimental::distributed::Matrix<scalar, label, label>>(
-            this->dist_mtx_.get());
-        // TODO FIXME dont leave hardcoded value here
-        word matrix_format_{"Coo"};
-
-        if (matrix_format_ == "Coo") {
-            using Coo = gko::matrix::Coo<scalar, label>;
-            std::vector<std::shared_ptr<const gko::LinOp>> local_interfaces =
-                gko::as<CombinationMatrix<scalar, label, Coo>>(
-                    dist_matrix->get_local_matrix())
-                    ->get_combination()
-                    ->get_operators();
-            export_mtx<Coo>(word(field_name_ + "_local"), local_interfaces,
-                            db_);
-            auto non_local_interfaces =
-                gko::as<CombinationMatrix<scalar, label, Coo>>(
-                    dist_matrix->get_non_local_matrix())
-                    ->get_combination()
-                    ->get_operators();
-            export_mtx<Coo>(field_name_ + "_non_local", non_local_interfaces,
-                            db_);
-        }
-
-        if (matrix_format_ == "Csr") {
-            using Csr = gko::matrix::Csr<scalar, label>;
-            auto local_interfaces =
-                gko::as<CombinationMatrix<scalar, label, Csr>>(
-                    dist_matrix->get_local_matrix().get())
-                    ->get_combination()
-                    ->get_operators();
-            auto non_local_interfaces =
-                gko::as<CombinationMatrix<scalar, label, Csr>>(
-                    dist_matrix->get_non_local_matrix().get())
-                    ->get_combination()
-                    ->get_operators();
-            export_mtx<Csr>(word(field_name_ + "_local"), local_interfaces,
-                            db_);
-            export_mtx<Csr>(field_name_ + "_non_local", non_local_interfaces,
-                            db_);
-        }
-
-        if (matrix_format_ == "Ell") {
-            using Ell = gko::matrix::Ell<scalar, label>;
-            auto local_interfaces =
-                gko::as<CombinationMatrix<scalar, label, Ell>>(
-                    dist_matrix->get_local_matrix().get())
-                    ->get_combination()
-                    ->get_operators();
-            auto non_local_interfaces =
-                gko::as<CombinationMatrix<scalar, label, Ell>>(
-                    dist_matrix->get_non_local_matrix().get())
-                    ->get_combination()
-                    ->get_operators();
-            export_mtx<Ell>(field_name_ + "_local", local_interfaces, db_);
-            export_mtx<Ell>(field_name_ + "_non_local", non_local_interfaces,
-                            db_);
-        }
+        // auto dist_matrix = gko::as<
+        //     gko::experimental::distributed::Matrix<scalar, label, label>>(
+        //     this->dist_mtx_.get());
+        // // TODO FIXME dont leave hardcoded value here
+        // word matrix_format_{"Coo"};
+
+        // if (matrix_format_ == "Coo") {
+        //     using Coo = gko::matrix::Coo<scalar, label>;
+        //     std::vector<std::shared_ptr<const gko::LinOp>> local_interfaces =
+        //         gko::as<CombinationMatrix<scalar, label, Coo>>(
+        //             dist_matrix->get_local_matrix())
+        //             ->get_combination()
+        //             ->get_operators();
+        //     export_mtx<Coo>(word(field_name_ + "_local"), local_interfaces,
+        //                     db_);
+        //     auto non_local_interfaces =
+        //         gko::as<CombinationMatrix<scalar, label, Coo>>(
+        //             dist_matrix->get_non_local_matrix())
+        //             ->get_combination()
+        //             ->get_operators();
+        //     export_mtx<Coo>(field_name_ + "_non_local", non_local_interfaces,
+        //                     db_);
+        // }
+
+        // if (matrix_format_ == "Csr") {
+        //     using Csr = gko::matrix::Csr<scalar, label>;
+        //     auto local_interfaces =
+        //         gko::as<CombinationMatrix<scalar, label, Csr>>(
+        //             dist_matrix->get_local_matrix().get())
+        //             ->get_combination()
+        //             ->get_operators();
+        //     auto non_local_interfaces =
+        //         gko::as<CombinationMatrix<scalar, label, Csr>>(
+        //             dist_matrix->get_non_local_matrix().get())
+        //             ->get_combination()
+        //             ->get_operators();
+        //     export_mtx<Csr>(word(field_name_ + "_local"), local_interfaces,
+        //                     db_);
+        //     export_mtx<Csr>(field_name_ + "_non_local", non_local_interfaces,
+        //                     db_);
+        // }
+
+        // if (matrix_format_ == "Ell") {
+        //     using Ell = gko::matrix::Ell<scalar, label>;
+        //     auto local_interfaces =
+        //         gko::as<CombinationMatrix<scalar, label, Ell>>(
+        //             dist_matrix->get_local_matrix().get())
+        //             ->get_combination()
+        //             ->get_operators();
+        //     auto non_local_interfaces =
+        //         gko::as<CombinationMatrix<scalar, label, Ell>>(
+        //             dist_matrix->get_non_local_matrix().get())
+        //             ->get_combination()
+        //             ->get_operators();
+        //     export_mtx<Ell>(field_name_ + "_local", local_interfaces, db_);
+        //     export_mtx<Ell>(field_name_ + "_non_local", non_local_interfaces,
+        //                     db_);
+        // }
     }
 
     // Needed for distributed/polymorphic_object.hpp
diff --git a/include/OGL/MatrixWrapper/HostMatrix.H b/include/OGL/MatrixWrapper/HostMatrix.H
index ef28fe389..cbcb0939a 100644
--- a/include/OGL/MatrixWrapper/HostMatrix.H
+++ b/include/OGL/MatrixWrapper/HostMatrix.H
@@ -216,7 +216,8 @@ public:
      ** @param interfaces The list of interfaces for the search
      ** @return the CommunicationPattern
      */
-    [[nodiscard]] CommunicationPattern create_communication_pattern() const;
+    [[nodiscard]] std::shared_ptr<CommunicationPattern>
+    create_communication_pattern() const;
 
     bool get_verbose() const { return verbose_; }
 
diff --git a/include/OGL/MatrixWrapper/SparsityPattern.H b/include/OGL/MatrixWrapper/SparsityPattern.H
index ba7295b77..da68be505 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.H
+++ b/include/OGL/MatrixWrapper/SparsityPattern.H
@@ -64,17 +64,10 @@ struct SparsityPattern {
 
     SparsityPattern(std::shared_ptr<const gko::Executor> exec)
         : num_nnz(0),
-          dim(gko::dim<2>{}),
           row_idxs{exec},
           col_idxs{exec},
-          ldu_mapping{exec}
-    {}
-
-    SparsityPattern(std::shared_ptr<const gko::Executor> exec, const label size)
-        : num_nnz(size),
-          row_idxs{exec, static_cast<gko::size_type>(size)},
-          col_idxs{exec, static_cast<gko::size_type>(size)},
-          ldu_mapping{exec, static_cast<gko::size_type>(size)}
+          ldu_mapping{exec},
+          dim(gko::dim<2>{})
     {}
 
     SparsityPattern(std::shared_ptr<const gko::Executor> exec, const label size,
@@ -83,10 +76,10 @@ struct SparsityPattern {
                     std::vector<label> &spans_begin,
                     std::vector<label> &spans_end, std::vector<label> &ranks)
         : num_nnz(size),
-          dim(dim),
           row_idxs{exec, static_cast<gko::size_type>(size)},
           col_idxs{exec, static_cast<gko::size_type>(size)},
-          ldu_mapping{exec, static_cast<gko::size_type>(size)}
+          ldu_mapping{exec, static_cast<gko::size_type>(size)},
+          dim(dim)
     {
         ASSERT_EQ(size, rows.size());
         std::copy(rows.begin(), rows.end(), row_idxs.get_data());
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index f1881d8ae..b7d36470c 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -26,21 +26,7 @@ std::vector<label> convert_to_global(
         const gko::experimental::distributed::Partition<label, label>>
         partition,
     const std::vector<label> &idx, const std::vector<gko::span> &spans,
-    const std::vector<label> &ranks)
-{
-    std::vector<label> ret;
-    ret.reserve(idx.size());
-
-    for (label i = 0; i < ranks.size(); i++) {
-        auto rank = ranks[i];
-        auto [begin, end] = spans[i];
-        label offset = partition->get_range_bounds()[rank];
-        for (label j = begin; j < end; j++) {
-            ret.push_back(idx.data()[j] + offset);
-        }
-    }
-    return ret;
-}
+    const std::vector<label> &ranks);
 
 /* Convert to local
 ** Given an array of column indices local to the communication rank
@@ -54,12 +40,7 @@ void convert_to_local(
     std::shared_ptr<
         const gko::experimental::distributed::Partition<label, label>>
         partition,
-    std::vector<label> &in, label rank)
-{
-    label offset = partition->get_range_bounds()[rank];
-    std::transform(in.begin(), in.end(), in.begin(),
-                   [&](label idx) { return idx - offset; });
-}
+    std::vector<label> &in, label rank);
 
 /* @brief exchange spans and rank between owner and non-owner ranks
 **
@@ -89,7 +70,7 @@ class Repartitioner {
 private:
     using partition = gko::experimental::distributed::Partition<label, label>;
 
-    const label size_;  //! Size (matrix rows) before repartitioning
+    // const label size_;  //! Size (matrix rows) before repartitioning
 
     const label repart_size_;  //! Size (matrix rows) after repartitioning
 
@@ -102,7 +83,7 @@ private:
 public:
     Repartitioner(label size, label ranks_per_gpu, label verbose,
                   const ExecutorHandler &exec_handler)
-        : size_(size),
+        :  // size_(size),
           repart_size_(Repartitioner::compute_repart_size(size, ranks_per_gpu,
                                                           exec_handler)),
           ranks_per_gpu_(ranks_per_gpu),
@@ -229,8 +210,5 @@ public:
      */
     std::shared_ptr<const CommunicationPattern> repartition_comm_pattern(
         const ExecutorHandler &exec_handler,
-        std::shared_ptr<const CommunicationPattern> src_comm_pattern,
-        std::shared_ptr<
-            const gko::experimental::distributed::Partition<label, label>>
-            partition) const;
+        std::shared_ptr<const CommunicationPattern> src_comm_pattern) const;
 };
diff --git a/src/CommunicationPattern.C b/src/CommunicationPattern.C
index 61598afbb..e0aad54ea 100644
--- a/src/CommunicationPattern.C
+++ b/src/CommunicationPattern.C
@@ -161,7 +161,6 @@ std::vector<label> gather_labels_to_owner(const ExecutorHandler &exec_handler,
 {
     auto exec = exec_handler.get_ref_exec();
     auto comm = *exec_handler.get_communicator().get();
-    auto rank = comm.rank();
 
     std::vector<label> send_buffer_copy;
     // create a copy if offset is needed
@@ -215,7 +214,6 @@ gko::array<label> CommunicationPattern::total_rank_send_idx() const
 AllToAllPattern CommunicationPattern::send_recv_pattern() const
 {
     auto comm = *exec_handler.get_communicator().get();
-    label total_ranks{comm.size()};
 
     std::vector<label> send_counts(comm.size());
     std::vector<label> send_offsets(comm.size() + 1);
diff --git a/src/MatrixWrapper/Distributed.C b/src/MatrixWrapper/Distributed.C
new file mode 100644
index 000000000..c7f0b0978
--- /dev/null
+++ b/src/MatrixWrapper/Distributed.C
@@ -0,0 +1,73 @@
+// SPDX-FileCopyrightText: 2024 OGL authors
+//
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "OGL/MatrixWrapper/Distributed.H"
+
+std::vector<std::unique_ptr<gko::LinOp>> detail::generate_inner_linops(
+    word matrix_format, bool fuse, std::shared_ptr<const gko::Executor> exec,
+    std::shared_ptr<const SparsityPattern> sparsity)
+{
+    std::vector<std::unique_ptr<gko::LinOp>> lin_ops;
+
+    for (int i = 0; i < sparsity->spans.size(); i++)
+    {
+        auto [begin, end] = sparsity->spans[i];
+        gko::array<scalar> coeffs(exec, end - begin);
+        coeffs.fill(0.0);
+
+        gko::array<label> row_idxs_view(exec->get_master(), end - begin);
+        gko::array<label> col_idxs_view(exec->get_master(), end - begin);
+
+        auto create_function = [&](word format) {
+            return gko::matrix::Coo<scalar, label>::create(
+                exec, sparsity->dim, coeffs, col_idxs_view, row_idxs_view);
+        };
+        lin_ops.push_back(create_function(matrix_format));
+    }
+    return lin_ops;
+}
+
+
+// auto local_dim = local_sparsity->dim;
+// auto non_local_dim = non_local_sparsity->dim;
+// auto local_interfaces = local_sparsity->interface_spans;
+// auto non_local_interfaces = non_local_sparsity->interface_spans;
+
+// if (matrix_format == "Ell") {
+//     using mtx_type = gko::matrix::Ell<scalar, label>;
+
+//     return DistMtxType::create(
+//         exec, comm,
+//         CombinationMatrix<scalar, label, mtx_type>::create(
+//             exec, local_dim, local_interfaces),
+//         CombinationMatrix<scalar, label, mtx_type>::create(
+//             exec, non_local_dim, non_local_interfaces));
+// }
+// if (matrix_format == "Csr") {
+//     using mtx_type = gko::matrix::Csr<scalar, label>;
+
+//     return DistMtxType::create(
+//         exec, comm,
+//         CombinationMatrix<scalar, label, mtx_type>::create(
+//             exec, local_dim, local_interfaces),
+//         CombinationMatrix<scalar, label, mtx_type>::create(
+//             exec, non_local_dim, non_local_interfaces));
+// }
+// if (matrix_format == "Coo") {
+//     using mtx_type = gko::matrix::Coo<scalar, label>;
+
+//     return DistMtxType::create(
+//         exec, comm,
+//         CombinationMatrix<scalar, label, mtx_type>::create(
+//             exec, local_dim, local_interfaces),
+//         CombinationMatrix<scalar, label, mtx_type>::create(
+//             exec, non_local_dim, non_local_interfaces));
+// }
+
+//     FatalErrorInFunction << "Matrix format " << matrix_format
+//                          << " not supported! Supported formats: Csr, Ell,
+//                          Coo"
+//                          << abort(FatalError);
+//     return {};
+// }
diff --git a/src/MatrixWrapper/HostMatrix.C b/src/MatrixWrapper/HostMatrix.C
index c97d9719d..d5a368372 100644
--- a/src/MatrixWrapper/HostMatrix.C
+++ b/src/MatrixWrapper/HostMatrix.C
@@ -259,7 +259,8 @@ void neg_interface_iterator(const lduInterfaceFieldPtrsList &interfaces,
 }
 
 
-CommunicationPattern HostMatrixWrapper::create_communication_pattern() const
+std::shared_ptr<CommunicationPattern>
+HostMatrixWrapper::create_communication_pattern() const
 {
     using comm_size_type = CommunicationPattern::comm_size_type;
     // temp map, mapping from neighbour rank interface cells
@@ -307,8 +308,8 @@ CommunicationPattern HostMatrixWrapper::create_communication_pattern() const
         iter++;
     }
 
-    return CommunicationPattern{get_exec_handler(), target_ids, target_sizes,
-                                send_idxs};
+    return std::make_shared<CommunicationPattern>(
+        get_exec_handler(), target_ids, target_sizes, send_idxs);
 }
 
 
@@ -450,22 +451,22 @@ HostMatrixWrapper::collect_cells_on_non_local_interface(
 std::shared_ptr<SparsityPattern> HostMatrixWrapper::compute_non_local_sparsity(
     std::shared_ptr<const gko::Executor> exec) const
 {
-    auto sparsity{std::make_shared<SparsityPattern>(exec->get_master(),
-                                                    non_local_matrix_nnz_)};
-    sparsity->dim =
-        gko::dim<2>{static_cast<gko::size_type>(nrows_),
-                    static_cast<gko::size_type>(non_local_matrix_nnz_)};
-
     auto non_local_indices = collect_cells_on_non_local_interface(interfaces_);
-    auto rows = sparsity->row_idxs.get_data();
-    auto cols = sparsity->col_idxs.get_data();
-    auto permute = sparsity->ldu_mapping.get_data();
+    std::vector<label> rows_vec(non_local_matrix_nnz_);
+    std::vector<label> cols_vec(non_local_matrix_nnz_);
+    std::vector<label> mapping_vec(non_local_matrix_nnz_);
+    std::vector<label> ranks{};
+    std::vector<gko::span> spans{};
+
+    auto rows = rows_vec.data();
+    auto cols = cols_vec.data();
+    auto permute = mapping_vec.data();
     label prev_interface_ctr{0};
     label end{0};
     label start{0};
 
     label element_ctr = 0;
-    label interface_ctr{0};
+    // label interface_ctr{0};
     label prev_rank{0};
 
     // TODO currently we set permute eventhough this is not required
@@ -483,8 +484,8 @@ std::shared_ptr<SparsityPattern> HostMatrixWrapper::compute_non_local_sparsity(
             // the end of the non_local_indices thus we need to increment
             // the element_ctr once more
             end = (last_element) ? element_ctr + 1 : element_ctr;
-            sparsity->spans.emplace_back(start, end);
-            sparsity->rank.emplace_back(prev_rank);
+            spans.emplace_back(start, end);
+            ranks.emplace_back(prev_rank);
             start = end;
             prev_interface_ctr = interface_idx;
         }
@@ -492,16 +493,28 @@ std::shared_ptr<SparsityPattern> HostMatrixWrapper::compute_non_local_sparsity(
         element_ctr++;
     }
 
-    return sparsity;
+    gko::dim<2> dim{static_cast<gko::size_type>(nrows_),
+                    static_cast<gko::size_type>(non_local_matrix_nnz_)};
+    return std::make_shared<SparsityPattern>(
+        exec->get_master(), dim, rows_vec, cols_vec, mapping_vec, spans, ranks);
 }
 
+
 std::shared_ptr<SparsityPattern> HostMatrixWrapper::compute_local_sparsity(
     std::shared_ptr<const gko::Executor> exec) const
 {
     LOG_1(verbose_, "start init host sparsity pattern")
 
-    auto sparsity{std::make_shared<SparsityPattern>(
-        exec->get_master(), local_matrix_w_interfaces_nnz_)};
+    std::vector<label> rows_vec(local_matrix_w_interfaces_nnz_);
+    std::vector<label> cols_vec(local_matrix_w_interfaces_nnz_);
+    std::vector<label> mapping_vec(local_matrix_w_interfaces_nnz_);
+    std::vector<label> ranks{exec_.get_rank()};
+    std::vector<gko::span> spans{gko::span{
+        0, static_cast<gko::size_type>(local_matrix_w_interfaces_nnz_)}};
+
+    auto rows = rows_vec.data();
+    auto cols = cols_vec.data();
+    auto permute = mapping_vec.data();
 
     auto lower_local = idx_array::view(
         exec, upper_nnz_, const_cast<label *>(addr_.lowerAddr().begin()));
@@ -515,10 +528,6 @@ std::shared_ptr<SparsityPattern> HostMatrixWrapper::compute_local_sparsity(
     // col of upper, row of lower
     const auto upper = upper_local.get_const_data();
 
-    auto rows = sparsity->row_idxs.get_data();
-    auto cols = sparsity->col_idxs.get_data();
-    const auto permute = sparsity->ldu_mapping.get_data();
-
     // Scan through given rows and insert row and column indices into array
     //
     // position after all local offdiagonal elements, needed for
@@ -527,14 +536,10 @@ std::shared_ptr<SparsityPattern> HostMatrixWrapper::compute_local_sparsity(
     // TODO in order to simplify when local interfaces exists set
     // local_sparsity to size of nrows_w_interfaces, if interfaces exist
     // local_sparsity is only valid till nrows_
-    label after_neighbours = (symmetric_) ? upper_nnz_ : 2 * upper_nnz_;
+    // label after_neighbours = (symmetric_) ? upper_nnz_ : 2 * upper_nnz_;
     init_local_sparsity(nrows_, upper_nnz_, symmetric_, upper, lower, rows,
                         cols, permute);
 
-    sparsity->dim = gko::dim<2>{static_cast<gko::size_type>(nrows_),
-                                static_cast<gko::size_type>(nrows_)};
-    sparsity->spans.emplace_back(0, local_matrix_nnz_);
-
     // if no local interfaces are present we are done here
     // otherwise we need to add local interfaces to local_sparsity in order
     // of the interfaces to end of the col and row idx arrays. This will produce
@@ -574,7 +579,8 @@ std::shared_ptr<SparsityPattern> HostMatrixWrapper::compute_local_sparsity(
                 local_interface_ctr + 1 == local_interfaces.size()) {
                 end = start + local_interface_ctr;
                 local_interface_ctr = 0;
-                sparsity->spans.emplace_back(start, end);
+                spans.emplace_back(start, end);
+                ranks.emplace_back(exec_.get_rank());
                 start = end;
                 prev_interface_idx = interface_idx;
             }
@@ -583,7 +589,9 @@ std::shared_ptr<SparsityPattern> HostMatrixWrapper::compute_local_sparsity(
     }
 
     LOG_1(verbose_, "done init host sparsity pattern")
-    return sparsity;
+    return std::make_shared<SparsityPattern>(exec->get_master(), get_size(),
+                                             rows_vec, cols_vec, mapping_vec,
+                                             spans, ranks);
 }
 
 }  // namespace Foam
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index 2192a12b0..d788ee185 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -6,6 +6,38 @@
 
 namespace detail {
 
+std::vector<label> convert_to_global(
+    std::shared_ptr<
+        const gko::experimental::distributed::Partition<label, label>>
+        partition,
+    const std::vector<label> &idx, const std::vector<gko::span> &spans,
+    const std::vector<label> &ranks)
+{
+    std::vector<label> ret;
+    ret.reserve(idx.size());
+
+    for (label i = 0; i < ranks.size(); i++) {
+        auto rank = ranks[i];
+        auto [begin, end] = spans[i];
+        label offset = partition->get_range_bounds()[rank];
+        for (label j = begin; j < end; j++) {
+            ret.push_back(idx.data()[j] + offset);
+        }
+    }
+    return ret;
+}
+
+void convert_to_local(
+    std::shared_ptr<
+        const gko::experimental::distributed::Partition<label, label>>
+        partition,
+    std::vector<label> &in, label rank)
+{
+    label offset = partition->get_range_bounds()[rank];
+    std::transform(in.begin(), in.end(), in.begin(),
+                   [&](label idx) { return idx - offset; });
+}
+
 std::tuple<std::vector<gko::span>, std::vector<label>, std::vector<label>>
 exchange_span_ranks(const ExecutorHandler &exec_handler, label ranks_per_gpu,
                     const std::vector<gko::span> &spans,
@@ -108,7 +140,7 @@ Repartitioner::repartition_sparsity(
     // early return if no repartitioning requested
     if (ranks_per_gpu == 1) {
         std::vector<std::pair<bool, label>> ret;
-        for (auto comm_rank : src_non_local_pattern->rank) {
+        for ([[maybe_unused]] auto comm_rank : src_non_local_pattern->rank) {
             ret.emplace_back(false, rank);
         }
         return std::make_tuple<std::shared_ptr<SparsityPattern>,
@@ -306,10 +338,8 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
 std::shared_ptr<const CommunicationPattern>
 Repartitioner::repartition_comm_pattern(
     const ExecutorHandler &exec_handler,
-    std::shared_ptr<const CommunicationPattern> src_comm_pattern,
-    std::shared_ptr<
-        const gko::experimental::distributed::Partition<label, label>>
-        partition) const
+    std::shared_ptr<const CommunicationPattern> src_comm_pattern
+) const
 {
     if (ranks_per_gpu_ == 1) {
         return src_comm_pattern;
@@ -373,8 +403,8 @@ Repartitioner::repartition_comm_pattern(
                           rank);
 
                 // the new offset is
-                auto offset = partition->get_range_bounds()[rank + i] -
-                              partition->get_range_bounds()[rank];
+                auto offset = get_orig_partition()->get_range_bounds()[rank + i] -
+                              get_orig_partition()->get_range_bounds()[rank];
 
                 std::transform(recv_buffer.begin(), recv_buffer.end(),
                                recv_buffer.begin(),
diff --git a/unitTests/MatrixWrapper/HostMatrix.C b/unitTests/MatrixWrapper/HostMatrix.C
index e23fec25b..c5bd9ef43 100644
--- a/unitTests/MatrixWrapper/HostMatrix.C
+++ b/unitTests/MatrixWrapper/HostMatrix.C
@@ -116,20 +116,20 @@ TEST(HostMatrix, canCreateCommunicationPattern)
     std::shared_ptr<const HostMatrixWrapper> hostMatrix =
         ((HostMatrixEnvironment *)global_env)->hostMatrix;
     auto commPattern = hostMatrix->create_communication_pattern();
-    auto comm = commPattern.get_comm();
+    auto comm = commPattern->get_comm();
 
-    EXPECT_EQ(commPattern.send_idxs.size(), 2);
+    EXPECT_EQ(commPattern->send_idxs.size(), 2);
 
     std::vector<std::vector<label>> target_ids_exp{
         {1, 2}, {0, 3}, {0, 3}, {1, 2}};
-    std::vector<label> target_ids_res(commPattern.target_ids.get_data(),
-                                      commPattern.target_ids.get_data() + 2);
+    std::vector<label> target_ids_res(commPattern->target_ids.get_data(),
+                                      commPattern->target_ids.get_data() + 2);
     EXPECT_EQ(target_ids_exp[comm.rank()], target_ids_res);
 
     std::vector<std::vector<label>> target_sizes_exp{
         {3, 3}, {3, 3}, {3, 3}, {3, 3}};
-    std::vector<label> target_size_res(commPattern.target_sizes.get_data(),
-                                       commPattern.target_sizes.get_data() + 2);
+    std::vector<label> target_size_res(commPattern->target_sizes.get_data(),
+                                       commPattern->target_sizes.get_data() + 2);
     EXPECT_EQ(target_sizes_exp[comm.rank()], target_size_res);
 }
 
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 26cf8db60..ce6e96929 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -377,11 +377,9 @@ TEST_P(RepartitionerFixture1D, can_repartition_1D_comm_pattern_for_n_ranks)
     auto comm_pattern =
         std::make_shared<CommunicationPattern>(exec, ids[rank], rows[rank]);
 
-    auto partition = repartitioner.get_orig_partition();
-
     // Act
     auto repart_comm_pattern =
-        repartitioner.repartition_comm_pattern(exec, comm_pattern, partition);
+        repartitioner.repartition_comm_pattern(exec, comm_pattern);
 
     // Assert
     auto res_ids = convert_to_vector(repart_comm_pattern->target_ids);
@@ -437,11 +435,9 @@ TEST_P(RepartitionerFixture2D, can_repartition_2D_comm_pattern_for_n_ranks)
     auto comm_pattern =
         std::make_shared<CommunicationPattern>(exec, ids[rank], rows[rank]);
 
-    auto partition = repartitioner.get_orig_partition();
-
     // Act
     auto repart_comm_pattern =
-        repartitioner.repartition_comm_pattern(exec, comm_pattern, partition);
+        repartitioner.repartition_comm_pattern(exec, comm_pattern);
 
     // Assert
     auto res_ids = convert_to_vector(repart_comm_pattern->target_ids);

From 93a9b5b56fc884e0350a6e3fbcf91af03b203408 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Thu, 29 Aug 2024 11:02:22 +0200
Subject: [PATCH 57/67] add unit test

---
 include/OGL/MatrixWrapper/Distributed.H |  78 ++++++---------
 src/MatrixWrapper/Distributed.C         |  12 ++-
 unitTests/MatrixWrapper/CMakeLists.txt  |   1 +
 unitTests/MatrixWrapper/Distributed.C   | 128 ++++++++++++++++++++++++
 4 files changed, 169 insertions(+), 50 deletions(-)
 create mode 100644 unitTests/MatrixWrapper/Distributed.C

diff --git a/include/OGL/MatrixWrapper/Distributed.H b/include/OGL/MatrixWrapper/Distributed.H
index 3fdf7aab0..751db6ef9 100644
--- a/include/OGL/MatrixWrapper/Distributed.H
+++ b/include/OGL/MatrixWrapper/Distributed.H
@@ -10,6 +10,7 @@
 #include <ginkgo/ginkgo.hpp>
 
 #include "OGL/CommunicationPattern.H"
+#include "OGL/MatrixWrapper/Combination.H"
 #include "OGL/MatrixWrapper/HostMatrix.H"
 #include "OGL/Repartitioner.H"
 
@@ -125,7 +126,7 @@ public:
         auto non_local_sparsity = host_A->compute_non_local_sparsity(exec);
 
         auto src_comm_pattern = host_A->create_communication_pattern();
-        auto dst_comm_pattern = repartitioner.repartition_comm_pattern(
+        auto repart_comm_pattern = repartitioner.repartition_comm_pattern(
             exec_handler, src_comm_pattern);
 
         // bool owner = repartitioner.is_owner(exec_handler);
@@ -137,17 +138,42 @@ public:
         // create vector of inner type linops
         // if fuse the vector contains only a single element
         // thus we can unwrap it.
-
         auto device_exec = exec_handler.get_device_exec();
         auto local_linops = detail::generate_inner_linops(
             matrix_format, false, device_exec, repart_loc_sparsity);
+        auto non_local_linops = detail::generate_inner_linops(
+            matrix_format, false, device_exec, repart_non_loc_sparsity);
+
+        auto [send_counts, send_offsets, recv_sizes, recv_offsets] =
+            repart_comm_pattern->send_recv_pattern();
+        auto recv_gather_idxs = repart_comm_pattern->total_rank_send_idx();
+
+        auto global_rows = repartitioner.get_orig_partition()->get_size();
+        gko::dim<2> global_dim{global_rows, global_rows};
+
+        std::unique_ptr<dist_mtx> dist_A;
+        bool fuse = false;
+        if (fuse) {
+        } else {
+            dist_A = dist_mtx::create(
+                exec, comm, global_dim,
+                gko::share(
+                    CombinationMatrix<gko::matrix::Coo<scalar, label>>::create(
+                        exec, repart_loc_sparsity->dim, local_linops)),
+                gko::share(
+                    CombinationMatrix<gko::matrix::Coo<scalar, label>>::create(
+                        exec, repart_non_loc_sparsity->dim, non_local_linops)),
+                recv_sizes, recv_offsets, recv_gather_idxs);
+        }
 
-        // auto local_coeffs =
-        //     gko::array<scalar>(exec, repart_local_sparsity->num_nnz);
-        // auto non_local_coeffs =
-        //     gko::array<scalar>(exec, repart_non_local_sparsity->num_nnz);
-
+        update_impl(exec_handler, matrix_format, repartitioner, host_A, dist_A,
+                    repart_loc_sparsity, repart_non_loc_sparsity,
+                    src_comm_pattern, local_interfaces);
 
+        return std::make_shared<RepartDistMatrix>(
+            exec, comm, repartitioner.get_repart_dim(), dist_A->get_size(),
+            std::move(dist_A), repart_loc_sparsity, repart_non_loc_sparsity,
+            src_comm_pattern, local_interfaces);
         // std::cout << __FILE__ << " rank " << rank << " local sparsity size "
         //           << local_sparsity_->size_ << " local sparsity dim ["
         //           << local_sparsity_->dim[0] << "x" <<
@@ -164,44 +190,6 @@ public:
         //           << " non_local_sparsity.size " <<
         //           non_local_sparsity_->size_
         //           << " \n";
-
-        // // FIXME make sure that we work on the device executor
-        // auto comm_pattern = compute_send_recv_counts(
-        //     exec_handler, 1, non_local_sparsity_->size_);
-
-        // auto send_idxs = dst_comm_pattern->total_rank_send_idx();
-        // auto comm_counts = dst_comm_pattern->send_recv_pattern();
-
-        // dist_A->read_distributed(
-        //     device_matrix_data(exec, local_sparsity_->dim,
-        //                        local_sparsity_->row_idxs,
-        //                        local_sparsity_->col_idxs, local_coeffs),
-        //     device_matrix_data(exec, non_local_sparsity_->dim,
-        //                        non_local_sparsity_->row_idxs,
-        //                        non_local_sparsity_->col_idxs,
-        //                        non_local_coeffs),
-        //     std::get<0>(comm_counts), std::get<1>(comm_counts),
-        //     std::get<2>(comm_counts), std::get<3>(comm_counts), send_idxs);
-
-
-        // auto dist_A = DistMtxType::create(
-        //        exec, comm,
-        //        CombinationMatrix<scalar, label, mtx_type>::create(
-        //            exec, local_dim, local_interfaces),
-        //        CombinationMatrix<scalar, label, mtx_type>::create(
-        //            exec, non_local_dim, non_local_interfaces));
-
-        // update_impl(exec_handler, matrix_format, repartitioner, host_A,
-        // dist_A,
-        //             repart_local_sparsity, repart_non_local_sparsity,
-        //             src_comm_pattern, local_interfaces);
-
-        // auto ret = std::make_shared<RepartDistMatrix>(
-        //     exec, comm, repartitioner.get_repart_dim(), dist_A->get_size(),
-        //     std::move(dist_A), local_sparsity_, non_local_sparsity_,
-        //     src_comm_pattern, local_interfaces);
-
-        return {};
     };
 
     void update(const ExecutorHandler &exec_handler, word matrix_format,
diff --git a/src/MatrixWrapper/Distributed.C b/src/MatrixWrapper/Distributed.C
index c7f0b0978..274191317 100644
--- a/src/MatrixWrapper/Distributed.C
+++ b/src/MatrixWrapper/Distributed.C
@@ -9,15 +9,17 @@ std::vector<std::unique_ptr<gko::LinOp>> detail::generate_inner_linops(
     std::shared_ptr<const SparsityPattern> sparsity)
 {
     std::vector<std::unique_ptr<gko::LinOp>> lin_ops;
-
-    for (int i = 0; i < sparsity->spans.size(); i++)
-    {
+    for (int i = 0; i < sparsity->spans.size(); i++) {
         auto [begin, end] = sparsity->spans[i];
         gko::array<scalar> coeffs(exec, end - begin);
         coeffs.fill(0.0);
 
-        gko::array<label> row_idxs_view(exec->get_master(), end - begin);
-        gko::array<label> col_idxs_view(exec->get_master(), end - begin);
+        auto row_idxs_view =
+            gko::array<label>::view(exec->get_master(), end - begin,
+                                    sparsity->row_idxs.get_data() + begin);
+        auto col_idxs_view =
+            gko::array<label>::view(exec->get_master(), end - begin,
+                                    sparsity->col_idxs.get_data() + begin);
 
         auto create_function = [&](word format) {
             return gko::matrix::Coo<scalar, label>::create(
diff --git a/unitTests/MatrixWrapper/CMakeLists.txt b/unitTests/MatrixWrapper/CMakeLists.txt
index 90a548f78..bd0e0ba75 100644
--- a/unitTests/MatrixWrapper/CMakeLists.txt
+++ b/unitTests/MatrixWrapper/CMakeLists.txt
@@ -11,6 +11,7 @@ configure_file(data/U.mtx data/U.mtx COPYONLY)
 
 ogl_unit_test(matrixConversion "HostMatrix.C")
 ogl_unit_test(combination "Combination.C")
+ogl_unit_test(distributed "Distributed.C")
 
 # Set working directory for the unit test suite of Combination Matrix
 set_tests_properties(combination PROPERTIES WORKING_DIRECTORY
diff --git a/unitTests/MatrixWrapper/Distributed.C b/unitTests/MatrixWrapper/Distributed.C
new file mode 100644
index 000000000..778efc38b
--- /dev/null
+++ b/unitTests/MatrixWrapper/Distributed.C
@@ -0,0 +1,128 @@
+// SPDX-FileCopyrightText: 2024 OGL authors
+//
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+
+#include "OGL/MatrixWrapper/HostMatrix.H"
+#include "OGL/Repartitioner.H"
+
+#include "gtest/gtest.h"
+
+//---------------------------------------------
+// some_header.h
+extern int my_argc;
+extern char **my_argv;
+// eof
+//---------------------------------------------
+
+//---------------------------------------------
+// main.cpp
+int my_argc;
+char **my_argv;
+
+
+class Environment : public testing::Environment {
+public:
+    void SetUp()
+    {
+        dict.add("executor", "reference");
+
+        args_ = std::make_shared<Foam::argList>(my_argc, my_argv);
+        if (args_->size() != 1) {
+            std::cout << "Wrong number of arguments detected: " << args_->size()
+                      << ", make sure to run "
+                         "with -parallel"
+                      << std::endl;
+            std::abort();
+        }
+
+        runTime_ = std::make_shared<Foam::Time>("controlDict", *args_.get());
+
+        mesh = std::make_shared<Foam::fvMesh>(
+            Foam::IOobject(word{""}, runTime_->timeName(), *runTime_.get(),
+                           Foam::IOobject::MUST_READ),
+            false);
+
+        exec = std::make_shared<ExecutorHandler>(runTime_->thisDb(), dict,
+                                                 "dummy", true);
+
+        auto comm = exec->get_gko_mpi_host_comm();
+        if (comm->size() != 4 || Pstream::nProcs() != 4) {
+            std::cout << "This unit test expects to be run on 4 ranks"
+                      << std::endl;
+            std::abort();
+        }
+
+        // delete listener on ranks != 0
+        // to clean up output
+        ::testing::TestEventListeners &listeners =
+            ::testing::UnitTest::GetInstance()->listeners();
+        if (Foam::Pstream::myProcNo() != 0) {
+            delete listeners.Release(listeners.default_result_printer());
+        }
+
+        word fieldName{"p"};
+        dimensionSet ds{0, 0, 0, 0, 0};
+        field = std::make_shared<
+            GeometricField<scalar, Foam::fvPatchField, Foam::volMesh>>(
+            Foam::IOobject(fieldName, runTime_->timeName(), runTime_->thisDb(),
+                           Foam::IOobject::MUST_READ),
+            *mesh.get(), ds);
+
+        Foam::fvMatrix<scalar> fvMatrix{*field.get(), ds};
+
+        interfaces = field->boundaryField().scalarInterfaces();
+
+        hostMatrix = std::make_shared<HostMatrixWrapper>(
+            *exec.get(), runTime_->thisDb(), mesh->lduAddr(), true,
+            fvMatrix.diag().data(), fvMatrix.upper().data(),
+            fvMatrix.lower().data(), fvMatrix.boundaryCoeffs(),
+            fvMatrix.internalCoeffs(), interfaces, dict, "fieldName", 0);
+    }
+
+    Foam::lduInterfaceFieldPtrsList interfaces;
+    Foam::PtrList<Foam::lduInterfaceField> newInterfaces;
+    std::shared_ptr<Foam::argList> args_;
+    std::shared_ptr<Foam::Time> runTime_;
+    std::shared_ptr<fvMesh> mesh;
+    Foam::dictionary dict;
+    std::shared_ptr<GeometricField<scalar, Foam::fvPatchField, Foam::volMesh>>
+        field;
+    std::shared_ptr<const ExecutorHandler> exec;
+    std::shared_ptr<const HostMatrixWrapper> hostMatrix;
+};
+
+const testing::Environment *global_env =
+    AddGlobalTestEnvironment(new Environment);
+
+class DistributedMatrixFixture : public testing::Test {
+public:
+    ExecutorHandler exec =
+        *((Environment *)global_env)->exec.get();
+    label rank = exec.get_rank();
+    const gko::experimental::mpi::communicator comm =
+        *(exec.get_communicator().get());
+};
+
+TEST_F(DistributedMatrixFixture, returnsCorrectSize)
+{
+    /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
+    auto mesh = ((Environment *)global_env)->mesh;
+    auto hostMatrix = ((Environment *)global_env)->hostMatrix;
+    auto repartitioner = Repartitioner(
+        hostMatrix->get_local_nrows(), 1, 0, exec);
+}
+
+int main(int argc, char *argv[])
+{
+    int result = 0;
+
+    ::testing::InitGoogleTest(&argc, argv);
+
+    my_argc = argc;
+    my_argv = argv;
+
+    result = RUN_ALL_TESTS();
+
+    return result;
+}

From 8318b16ccb70dd1e8954ca77b13deb46859f04a6 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Thu, 29 Aug 2024 11:54:36 +0200
Subject: [PATCH 58/67] extend unit test

---
 include/OGL/MatrixWrapper/Distributed.H | 12 ++++++------
 src/MatrixWrapper/Distributed.C         |  8 ++++----
 unitTests/MatrixWrapper/Distributed.C   |  7 ++++++-
 3 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/include/OGL/MatrixWrapper/Distributed.H b/include/OGL/MatrixWrapper/Distributed.H
index 751db6ef9..017a663e8 100644
--- a/include/OGL/MatrixWrapper/Distributed.H
+++ b/include/OGL/MatrixWrapper/Distributed.H
@@ -18,7 +18,7 @@ namespace detail {
 /* @brief generates a vector of ginkgo linops
 **
 */
-std::vector<std::unique_ptr<gko::LinOp>> generate_inner_linops(
+std::vector<std::shared_ptr<const gko::LinOp>> generate_inner_linops(
     word matrix_format, bool fuse, std::shared_ptr<const gko::Executor> exec,
     std::shared_ptr<const SparsityPattern> sparsity);
 }  // namespace detail
@@ -60,12 +60,12 @@ public:
 
     std::shared_ptr<const gko::LinOp> get_local_matrix()
     {
-        this->dist_mtx_->get_local_matrix();
+        return this->dist_mtx_->get_local_matrix();
     }
 
     std::shared_ptr<const gko::LinOp> get_non_local_matrix()
     {
-        this->dist_mtx_->get_non_local_matrix();
+        return this->dist_mtx_->get_non_local_matrix();
     }
 
     /**
@@ -151,11 +151,11 @@ public:
         auto global_rows = repartitioner.get_orig_partition()->get_size();
         gko::dim<2> global_dim{global_rows, global_rows};
 
-        std::unique_ptr<dist_mtx> dist_A;
+        std::shared_ptr<dist_mtx> dist_A;
         bool fuse = false;
         if (fuse) {
         } else {
-            dist_A = dist_mtx::create(
+            dist_A = gko::share(dist_mtx::create(
                 exec, comm, global_dim,
                 gko::share(
                     CombinationMatrix<gko::matrix::Coo<scalar, label>>::create(
@@ -163,7 +163,7 @@ public:
                 gko::share(
                     CombinationMatrix<gko::matrix::Coo<scalar, label>>::create(
                         exec, repart_non_loc_sparsity->dim, non_local_linops)),
-                recv_sizes, recv_offsets, recv_gather_idxs);
+                recv_sizes, recv_offsets, recv_gather_idxs));
         }
 
         update_impl(exec_handler, matrix_format, repartitioner, host_A, dist_A,
diff --git a/src/MatrixWrapper/Distributed.C b/src/MatrixWrapper/Distributed.C
index 274191317..1b310f387 100644
--- a/src/MatrixWrapper/Distributed.C
+++ b/src/MatrixWrapper/Distributed.C
@@ -4,11 +4,11 @@
 
 #include "OGL/MatrixWrapper/Distributed.H"
 
-std::vector<std::unique_ptr<gko::LinOp>> detail::generate_inner_linops(
+std::vector<std::shared_ptr<const gko::LinOp>> detail::generate_inner_linops(
     word matrix_format, bool fuse, std::shared_ptr<const gko::Executor> exec,
     std::shared_ptr<const SparsityPattern> sparsity)
 {
-    std::vector<std::unique_ptr<gko::LinOp>> lin_ops;
+    std::vector<std::shared_ptr<const gko::LinOp>> lin_ops;
     for (int i = 0; i < sparsity->spans.size(); i++) {
         auto [begin, end] = sparsity->spans[i];
         gko::array<scalar> coeffs(exec, end - begin);
@@ -22,8 +22,8 @@ std::vector<std::unique_ptr<gko::LinOp>> detail::generate_inner_linops(
                                     sparsity->col_idxs.get_data() + begin);
 
         auto create_function = [&](word format) {
-            return gko::matrix::Coo<scalar, label>::create(
-                exec, sparsity->dim, coeffs, col_idxs_view, row_idxs_view);
+            return gko::share(gko::matrix::Coo<scalar, label>::create(
+                exec, sparsity->dim, coeffs, col_idxs_view, row_idxs_view));
         };
         lin_ops.push_back(create_function(matrix_format));
     }
diff --git a/unitTests/MatrixWrapper/Distributed.C b/unitTests/MatrixWrapper/Distributed.C
index 778efc38b..b42971733 100644
--- a/unitTests/MatrixWrapper/Distributed.C
+++ b/unitTests/MatrixWrapper/Distributed.C
@@ -4,6 +4,7 @@
 
 
 #include "OGL/MatrixWrapper/HostMatrix.H"
+#include "OGL/MatrixWrapper/Distributed.H"
 #include "OGL/Repartitioner.H"
 
 #include "gtest/gtest.h"
@@ -104,13 +105,17 @@ public:
         *(exec.get_communicator().get());
 };
 
-TEST_F(DistributedMatrixFixture, returnsCorrectSize)
+TEST_F(DistributedMatrixFixture, canCreateDistributeMatrix)
 {
     /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
     auto mesh = ((Environment *)global_env)->mesh;
     auto hostMatrix = ((Environment *)global_env)->hostMatrix;
     auto repartitioner = Repartitioner(
         hostMatrix->get_local_nrows(), 1, 0, exec);
+
+    auto distributed = RepartDistMatrix<scalar, label, label>::create(exec, "Coo", repartitioner, hostMatrix);
+
+    ASSERT_EQ(distributed->get_local_matrix()->get_size()[0], 9);
 }
 
 int main(int argc, char *argv[])

From 731b75d6e6b819f733597620c3736e3a33889174 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Fri, 30 Aug 2024 09:28:07 +0200
Subject: [PATCH 59/67] Fix local span computation

---
 src/Repartitioner.C       |  9 +++------
 unitTests/Repartitioner.C | 29 ++++++++++++++++++++++++++---
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index d788ee185..e0a2b45c9 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -255,25 +255,22 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
         is_local.push_back(local);
 
         if (local) {
+            gko::size_type rows_start = local_rows.size();
             local_rows.insert(local_rows.end(), non_local_rows.data() + begin,
                               non_local_rows.data() + end);
 
             std::vector<label> tmp_rank_local_cols(
                 tmp_non_local_cols.data() + begin,
                 tmp_non_local_cols.data() + end);
-
             detail::convert_to_local(partition, tmp_rank_local_cols, rank);
-
             local_cols.insert(local_cols.end(), tmp_rank_local_cols.begin(),
                               tmp_rank_local_cols.end());
+
             local_mapping.insert(local_mapping.end(),
                                  non_local_mapping.data() + begin,
                                  non_local_mapping.data() + end);
-
-            gko::size_type rows_start = local_rows.size();
             local_spans.emplace_back(rows_start,
-                                     static_cast<gko::size_type>(end - begin));
-
+                                     rows_start + static_cast<gko::size_type>(end - begin));
             local_ranks.push_back(non_local_rank_origin[i]);
         } else {
             mark_keep.push_back(i);
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index ce6e96929..677393c6d 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -88,7 +88,7 @@ public:
     std::vector<label> rows{0, 0, 1, 1};
     std::vector<label> cols{0, 1, 0, 1};
     std::vector<label> mapping{0, 1, 2, 3};
-    std::vector<gko::span> spans{gko::span{0, 2}};
+    std::vector<gko::span> spans{gko::span{0, 5}};
 
     // non local data
     vec_vec non_local_rows{{1}, {0, 1}, {0, 1}, {0}};
@@ -219,8 +219,6 @@ TEST_P(RepartitionerFixture1D, can_exchange_spans_and_ranks_for_n_ranks)
     exp_ranks.emplace(2, vec_vec{{1, 2, 2, 3}, {}, {3, 4, 4, 5}, {}});
     exp_ranks.emplace(4, vec_vec{{1, 2, 2, 3, 3, 4, 4, 5}, {}, {}, {}});
 
-    // TODO test origins
-
     std::map<label, vec_vec> exp_spans_begin;
     exp_spans_begin.emplace(1, vec_vec{{0, 5}, {0, 5}, {0, 5}, {0, 5}});
     exp_spans_begin.emplace(2, vec_vec{{0, 5, 10, 15}, {}, {0, 5, 10, 15}, {}});
@@ -306,6 +304,22 @@ TEST_P(RepartitionerFixture1D, can_repartition_sparsity_pattern_1D_for_n_ranks)
     exp_local_dim_rows.emplace(2, vec{2 * local_size, 0, 2 * local_size, 0});
     exp_local_dim_rows.emplace(4, vec{4 * local_size, 0, 0, 0});
 
+    std::map<label, vec_vec> exp_local_spans_begin;
+    exp_local_spans_begin.emplace(
+        1, vec_vec{{0}, {0}, {0}, {0}});
+    exp_local_spans_begin.emplace(
+        2, vec_vec{{0, 8, 9}, {}, {0, 8, 9}, {}});
+    exp_local_spans_begin.emplace(
+        4, vec_vec{{0, 16, 17, 18, 19, 20, 21}, {}, {}, {}});
+
+    std::map<label, vec_vec> exp_local_spans_end;
+    exp_local_spans_end.emplace(
+        1, vec_vec{{5}, {5}, {5}, {5}});
+    exp_local_spans_end.emplace(
+        2, vec_vec{{8, 9, 10}, {}, {8, 9 ,10}, {}});
+    exp_local_spans_end.emplace(
+        4, vec_vec{{16, 17, 18, 19, 20, 21, 22}, {}, {}, {}});
+
     std::map<label, vec_vec> exp_non_local_rows;
     exp_non_local_rows.emplace(1, non_local_rows);
     exp_non_local_rows.emplace(2, vec_vec{{3}, {}, {0}, {}});
@@ -328,6 +342,15 @@ TEST_P(RepartitionerFixture1D, can_repartition_sparsity_pattern_1D_for_n_ranks)
     auto res_local_cols = convert_to_vector(repart_local->col_idxs);
     ASSERT_EQ(res_local_rows, exp_local_rows[ranks_per_gpu][rank]);
 
+    std::vector<label> res_local_spans_begin {};
+    std::vector<label> res_local_spans_end {};
+    for (auto [begin, end]: repart_local->spans) {
+        res_local_spans_begin.push_back(begin);
+        res_local_spans_end.push_back(end);
+    }
+    ASSERT_EQ(res_local_spans_begin, exp_local_spans_begin[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_local_spans_end, exp_local_spans_end[ranks_per_gpu][rank]);
+
     // non local properties
     ASSERT_EQ(repart_non_local->num_nnz,
               exp_non_local_nnz[ranks_per_gpu][rank]);

From 7836e5f14179ca25c212467a87b3e51c14a39c34 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Fri, 30 Aug 2024 09:33:31 +0200
Subject: [PATCH 60/67] parameterize unit test, clean up

---
 src/MatrixWrapper/Distributed.C       | 44 ---------------------------
 unitTests/MatrixWrapper/Distributed.C | 30 +++++++++++++++---
 2 files changed, 26 insertions(+), 48 deletions(-)

diff --git a/src/MatrixWrapper/Distributed.C b/src/MatrixWrapper/Distributed.C
index 1b310f387..dd0cf48af 100644
--- a/src/MatrixWrapper/Distributed.C
+++ b/src/MatrixWrapper/Distributed.C
@@ -29,47 +29,3 @@ std::vector<std::shared_ptr<const gko::LinOp>> detail::generate_inner_linops(
     }
     return lin_ops;
 }
-
-
-// auto local_dim = local_sparsity->dim;
-// auto non_local_dim = non_local_sparsity->dim;
-// auto local_interfaces = local_sparsity->interface_spans;
-// auto non_local_interfaces = non_local_sparsity->interface_spans;
-
-// if (matrix_format == "Ell") {
-//     using mtx_type = gko::matrix::Ell<scalar, label>;
-
-//     return DistMtxType::create(
-//         exec, comm,
-//         CombinationMatrix<scalar, label, mtx_type>::create(
-//             exec, local_dim, local_interfaces),
-//         CombinationMatrix<scalar, label, mtx_type>::create(
-//             exec, non_local_dim, non_local_interfaces));
-// }
-// if (matrix_format == "Csr") {
-//     using mtx_type = gko::matrix::Csr<scalar, label>;
-
-//     return DistMtxType::create(
-//         exec, comm,
-//         CombinationMatrix<scalar, label, mtx_type>::create(
-//             exec, local_dim, local_interfaces),
-//         CombinationMatrix<scalar, label, mtx_type>::create(
-//             exec, non_local_dim, non_local_interfaces));
-// }
-// if (matrix_format == "Coo") {
-//     using mtx_type = gko::matrix::Coo<scalar, label>;
-
-//     return DistMtxType::create(
-//         exec, comm,
-//         CombinationMatrix<scalar, label, mtx_type>::create(
-//             exec, local_dim, local_interfaces),
-//         CombinationMatrix<scalar, label, mtx_type>::create(
-//             exec, non_local_dim, non_local_interfaces));
-// }
-
-//     FatalErrorInFunction << "Matrix format " << matrix_format
-//                          << " not supported! Supported formats: Csr, Ell,
-//                          Coo"
-//                          << abort(FatalError);
-//     return {};
-// }
diff --git a/unitTests/MatrixWrapper/Distributed.C b/unitTests/MatrixWrapper/Distributed.C
index b42971733..61da19ea1 100644
--- a/unitTests/MatrixWrapper/Distributed.C
+++ b/unitTests/MatrixWrapper/Distributed.C
@@ -21,6 +21,9 @@ extern char **my_argv;
 int my_argc;
 char **my_argv;
 
+using vec = std::vector<label>;
+using vec_vec = std::vector<std::vector<label>>;
+using vec_vec_vec = std::vector<vec_vec>;
 
 class Environment : public testing::Environment {
 public:
@@ -96,7 +99,8 @@ public:
 const testing::Environment *global_env =
     AddGlobalTestEnvironment(new Environment);
 
-class DistributedMatrixFixture : public testing::Test {
+class DistributedMatrixFixture :
+     public testing::TestWithParam<int> {
 public:
     ExecutorHandler exec =
         *((Environment *)global_env)->exec.get();
@@ -105,17 +109,35 @@ public:
         *(exec.get_communicator().get());
 };
 
-TEST_F(DistributedMatrixFixture, canCreateDistributeMatrix)
+INSTANTIATE_TEST_SUITE_P(DistributedMatrixFixtureInstantiation,
+                         DistributedMatrixFixture, testing::Values(1, 2, 4));
+
+TEST_P(DistributedMatrixFixture, canCreateDistributeMatrix)
 {
     /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
+    label ranks_per_gpu = GetParam();
     auto mesh = ((Environment *)global_env)->mesh;
     auto hostMatrix = ((Environment *)global_env)->hostMatrix;
     auto repartitioner = Repartitioner(
-        hostMatrix->get_local_nrows(), 1, 0, exec);
+        hostMatrix->get_local_nrows(), ranks_per_gpu, 0, exec);
+
+    std::map<label, vec> exp_local_size;
+    exp_local_size.emplace(1, vec{9, 9, 9, 9});
+    exp_local_size.emplace(2, vec{18, 0, 18, 0});
+    exp_local_size.emplace(4, vec{36, 0 ,0 ,0});
+
+    std::map<label, vec> exp_non_local_size;
+    exp_non_local_size.emplace(1, vec{6, 6, 6, 6});
+    exp_non_local_size.emplace(2, vec{6, 0, 6, 0});
+    exp_non_local_size.emplace(4, vec{0, 0 ,0 ,0});
 
     auto distributed = RepartDistMatrix<scalar, label, label>::create(exec, "Coo", repartitioner, hostMatrix);
 
-    ASSERT_EQ(distributed->get_local_matrix()->get_size()[0], 9);
+    ASSERT_EQ(distributed->get_local_matrix()->get_size()[0], exp_local_size[ranks_per_gpu][rank]);
+    ASSERT_EQ(distributed->get_local_matrix()->get_size()[1], exp_local_size[ranks_per_gpu][rank]);
+    ASSERT_EQ(distributed->get_non_local_matrix()->get_size()[0], exp_local_size[ranks_per_gpu][rank]);
+    ASSERT_EQ(distributed->get_non_local_matrix()->get_size()[1], exp_non_local_size[ranks_per_gpu][rank]);
+
 }
 
 int main(int argc, char *argv[])

From aa0dd6e83bdc20d33c3d90d14586b62b3f363341 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Fri, 30 Aug 2024 13:47:46 +0200
Subject: [PATCH 61/67] WIP make update_impl compile, start apply unit test

---
 include/OGL/MatrixWrapper/Distributed.H | 112 ++-------
 include/OGL/MatrixWrapper/HostMatrix.H  |  20 ++
 include/OGL/Repartitioner.H             |   8 +-
 src/MatrixWrapper/Distributed.C         | 305 ++++++++++++++++++++++++
 unitTests/MatrixWrapper/Distributed.C   |  20 ++
 5 files changed, 374 insertions(+), 91 deletions(-)

diff --git a/include/OGL/MatrixWrapper/Distributed.H b/include/OGL/MatrixWrapper/Distributed.H
index 017a663e8..0119958f5 100644
--- a/include/OGL/MatrixWrapper/Distributed.H
+++ b/include/OGL/MatrixWrapper/Distributed.H
@@ -21,14 +21,27 @@ namespace detail {
 std::vector<std::shared_ptr<const gko::LinOp>> generate_inner_linops(
     word matrix_format, bool fuse, std::shared_ptr<const gko::Executor> exec,
     std::shared_ptr<const SparsityPattern> sparsity);
+
+void update_impl(
+    const ExecutorHandler &exec_handler, word matrix_format,
+    const Repartitioner &repartitioner,
+    std::shared_ptr<const HostMatrixWrapper> host_A,
+    std::shared_ptr<
+        gko::experimental::distributed::Matrix<scalar, label, label>>
+        dist_A,
+    std::shared_ptr<const SparsityPattern> local_sparsity,
+    std::shared_ptr<const SparsityPattern> non_local_sparsity,
+    std::shared_ptr<const CommunicationPattern> src_comm_pattern,
+    std::vector<std::pair<bool, label>> local_interfaces);
 }  // namespace detail
 
-// The RepartDistMatrix class is a wrapper around Ginkgos distributed Matrix
-// class
-//
-// It adds functionality for repeated read and repartitioning operatitions. As a
-// constraint it is required that the inner matrix types of the distributed
-// matrix are of RepartDistMatrix type.
+/* The RepartDistMatrix class is a wrapper around Ginkgos distributed Matrix
+ * class
+ *
+ * It adds functionality for repeated read and repartitioning operatitions. As a
+ * constraint it is required that the inner matrix types of the distributed
+ * matrix are of RepartDistMatrix type.
+ * */
 template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
 class RepartDistMatrix
     : public gko::experimental::EnableDistributedLinOp<
@@ -113,96 +126,19 @@ public:
     static std::shared_ptr<RepartDistMatrix> create(
         const ExecutorHandler &exec_handler, word matrix_format,
         const Repartitioner &repartitioner,
-        std::shared_ptr<const HostMatrixWrapper> host_A)
-    {
-        auto exec = exec_handler.get_ref_exec();
-        auto comm = *exec_handler.get_communicator().get();
-
-        // repartition things here first by creating device_matrix_data on
-        // device with correct size on each rank
-        //
-        // NOTE Start by copying first and repartition later
-        auto local_sparsity = host_A->compute_local_sparsity(exec);
-        auto non_local_sparsity = host_A->compute_non_local_sparsity(exec);
-
-        auto src_comm_pattern = host_A->create_communication_pattern();
-        auto repart_comm_pattern = repartitioner.repartition_comm_pattern(
-            exec_handler, src_comm_pattern);
-
-        // bool owner = repartitioner.is_owner(exec_handler);
-
-        auto [repart_loc_sparsity, repart_non_loc_sparsity, local_interfaces] =
-            repartitioner.repartition_sparsity(exec_handler, local_sparsity,
-                                               non_local_sparsity);
-
-        // create vector of inner type linops
-        // if fuse the vector contains only a single element
-        // thus we can unwrap it.
-        auto device_exec = exec_handler.get_device_exec();
-        auto local_linops = detail::generate_inner_linops(
-            matrix_format, false, device_exec, repart_loc_sparsity);
-        auto non_local_linops = detail::generate_inner_linops(
-            matrix_format, false, device_exec, repart_non_loc_sparsity);
-
-        auto [send_counts, send_offsets, recv_sizes, recv_offsets] =
-            repart_comm_pattern->send_recv_pattern();
-        auto recv_gather_idxs = repart_comm_pattern->total_rank_send_idx();
-
-        auto global_rows = repartitioner.get_orig_partition()->get_size();
-        gko::dim<2> global_dim{global_rows, global_rows};
-
-        std::shared_ptr<dist_mtx> dist_A;
-        bool fuse = false;
-        if (fuse) {
-        } else {
-            dist_A = gko::share(dist_mtx::create(
-                exec, comm, global_dim,
-                gko::share(
-                    CombinationMatrix<gko::matrix::Coo<scalar, label>>::create(
-                        exec, repart_loc_sparsity->dim, local_linops)),
-                gko::share(
-                    CombinationMatrix<gko::matrix::Coo<scalar, label>>::create(
-                        exec, repart_non_loc_sparsity->dim, non_local_linops)),
-                recv_sizes, recv_offsets, recv_gather_idxs));
-        }
-
-        update_impl(exec_handler, matrix_format, repartitioner, host_A, dist_A,
-                    repart_loc_sparsity, repart_non_loc_sparsity,
-                    src_comm_pattern, local_interfaces);
-
-        return std::make_shared<RepartDistMatrix>(
-            exec, comm, repartitioner.get_repart_dim(), dist_A->get_size(),
-            std::move(dist_A), repart_loc_sparsity, repart_non_loc_sparsity,
-            src_comm_pattern, local_interfaces);
-        // std::cout << __FILE__ << " rank " << rank << " local sparsity size "
-        //           << local_sparsity_->size_ << " local sparsity dim ["
-        //           << local_sparsity_->dim[0] << "x" <<
-        //           local_sparsity_->dim[1]
-        //           << "] non_local sparsity size " <<
-        //           non_local_sparsity_->size_
-        //           << " non local sparsity dim [" <<
-        //           non_local_sparsity_->dim[0]
-        //           << "x" << non_local_sparsity_->dim[1] << "] target_ids "
-        //           << dst_comm_pattern->target_ids << " target_sizes "
-        //           << dst_comm_pattern->target_sizes << "
-        //           target_send_idxs.size "
-        //           << dst_comm_pattern->send_idxs.size()
-        //           << " non_local_sparsity.size " <<
-        //           non_local_sparsity_->size_
-        //           << " \n";
-    };
+        std::shared_ptr<const HostMatrixWrapper> host_A);
 
     void update(const ExecutorHandler &exec_handler, word matrix_format,
                 const Repartitioner &repartitioner,
                 std::shared_ptr<const HostMatrixWrapper> host_A)
     {
-        update_impl(exec_handler, matrix_format, repartitioner, host_A,
-                    dist_mtx_, local_sparsity_, non_local_sparsity_,
-                    src_comm_pattern_, local_interfaces_);
+        detail::update_impl(exec_handler, matrix_format, repartitioner, host_A,
+                            dist_mtx_, local_sparsity_, non_local_sparsity_,
+                            src_comm_pattern_, local_interfaces_);
     };
 
     // TODO matrix_format can be data member
-    static void update_impl(
+    void update_impl(
         const ExecutorHandler &exec_handler, word matrix_format,
         const Repartitioner &repartitioner,
         std::shared_ptr<const HostMatrixWrapper> host_A,
diff --git a/include/OGL/MatrixWrapper/HostMatrix.H b/include/OGL/MatrixWrapper/HostMatrix.H
index cbcb0939a..2fc9af02d 100644
--- a/include/OGL/MatrixWrapper/HostMatrix.H
+++ b/include/OGL/MatrixWrapper/HostMatrix.H
@@ -223,6 +223,26 @@ public:
 
     label get_local_nrows() const { return nrows_; }
 
+    label get_local_matrix_nnz() const
+    {
+        return local_matrix_w_interfaces_nnz_;
+    }
+
+    label get_upper_nnz() const { return upper_nnz_; }
+
+    const scalar *get_lower() const { return lower_; }
+
+    const scalar *get_upper() const { return upper_; }
+
+    const scalar *get_diag() const { return diag_; }
+
+    const scalar *get_interface_data(label i) const
+    {
+        FatalErrorInFunction << " not implemented " << abort(FatalError);
+    }
+
+    bool get_symmetric() const { return symmetric_; }
+
     gko::dim<2> get_size() const { return gko::dim<2>(nrows_, nrows_); }
 
     const ExecutorHandler &get_exec_handler() const { return exec_; }
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index b7d36470c..621cdecc8 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -70,9 +70,9 @@ class Repartitioner {
 private:
     using partition = gko::experimental::distributed::Partition<label, label>;
 
-    // const label size_;  //! Size (matrix rows) before repartitioning
+    const label size_;  //! Local size (matrix rows) before repartitioning
 
-    const label repart_size_;  //! Size (matrix rows) after repartitioning
+    const label repart_size_;  //! Local size (matrix rows) after repartitioning
 
     const label ranks_per_gpu_;
 
@@ -83,7 +83,7 @@ private:
 public:
     Repartitioner(label size, label ranks_per_gpu, label verbose,
                   const ExecutorHandler &exec_handler)
-        :  // size_(size),
+        : size_(size),
           repart_size_(Repartitioner::compute_repart_size(size, ranks_per_gpu,
                                                           exec_handler)),
           ranks_per_gpu_(ranks_per_gpu),
@@ -133,6 +133,8 @@ public:
     static label compute_repart_size(label local_size, label ranks_per_gpu,
                                      const ExecutorHandler &exec_handler);
 
+    label get_orig_size() const { return size_; }
+
     label get_repart_size() const { return repart_size_; }
 
     // TODO pass original dim as argument, if size can also be nnz
diff --git a/src/MatrixWrapper/Distributed.C b/src/MatrixWrapper/Distributed.C
index dd0cf48af..5338feec3 100644
--- a/src/MatrixWrapper/Distributed.C
+++ b/src/MatrixWrapper/Distributed.C
@@ -29,3 +29,308 @@ std::vector<std::shared_ptr<const gko::LinOp>> detail::generate_inner_linops(
     }
     return lin_ops;
 }
+
+void detail::update_impl(
+    const ExecutorHandler &exec_handler, word matrix_format,
+    const Repartitioner &repartitioner,
+    std::shared_ptr<const HostMatrixWrapper> host_A,
+    std::shared_ptr<
+        gko::experimental::distributed::Matrix<scalar, label, label>>
+        dist_A,
+    std::shared_ptr<const SparsityPattern> local_sparsity,
+    std::shared_ptr<const SparsityPattern> non_local_sparsity,
+    std::shared_ptr<const CommunicationPattern> src_comm_pattern,
+    std::vector<std::pair<bool, label>> local_interfaces) {
+    auto exec = exec_handler.get_ref_exec();
+    auto device_exec = exec_handler.get_device_exec();
+    auto ranks_per_gpu = repartitioner.get_ranks_per_gpu();
+    bool requires_host_buffer = exec_handler.get_gko_force_host_buffer();
+
+    label rank{exec_handler.get_rank()};
+    label owner_rank = repartitioner.get_owner_rank(exec_handler);
+    bool owner = repartitioner.is_owner(exec_handler);
+    label nrows = host_A->get_local_nrows();
+    label local_matrix_nnz = host_A->get_local_matrix_nnz();
+    // size + padding has to be local_matrix_nnz
+    auto diag_comm_pattern = compute_gather_to_owner_counts(
+        exec_handler, ranks_per_gpu, nrows, local_matrix_nnz,
+        local_matrix_nnz - nrows, 0);
+    label upper_nnz = host_A->get_upper_nnz();
+    auto upper_comm_pattern = compute_gather_to_owner_counts(
+        exec_handler, ranks_per_gpu, upper_nnz, local_matrix_nnz, 0,
+        local_matrix_nnz - upper_nnz);
+    auto lower_comm_pattern =
+        compute_gather_to_owner_counts(exec_handler, ranks_per_gpu, upper_nnz,
+                                 local_matrix_nnz, upper_nnz, nrows);
+
+    scalar *local_ptr;
+    scalar *local_ptr_2;
+    label nnz = 0;
+
+    // update main values
+    std::vector<scalar> loc_buffer;
+    if (owner) {
+        using Coo = gko::matrix::Coo<scalar, label>;
+        auto local_mtx = dist_A->get_local_matrix();
+
+        // TODO make it work with other matrix types
+        std::shared_ptr<const Coo> local =
+            gko::as<Coo>(gko::as<CombinationMatrix<Coo>>(
+                             dist_A->get_local_matrix())
+                             ->get_combination()
+                             ->get_operators()[0]);
+        nnz = local->get_num_stored_elements();
+        if (requires_host_buffer) {
+            loc_buffer.resize(nnz);
+            local_ptr = loc_buffer.data();
+            local_ptr_2 = const_cast<scalar
+            *>(local->get_const_values());
+        } else {
+            local_ptr = const_cast<scalar *>(local->get_const_values());
+        }
+    }
+    communicate_values(exec_handler, diag_comm_pattern,
+    host_A->get_diag(),
+                       local_ptr);
+    communicate_values(exec_handler, upper_comm_pattern,
+                       host_A->get_upper(), local_ptr);
+
+    if (host_A->get_symmetric()) {
+        // TODO FIXME
+        // if symmetric we can reuse already copied data
+        communicate_values(exec_handler, lower_comm_pattern,
+                           host_A->get_lower(), local_ptr);
+    } else {
+        communicate_values(exec_handler, lower_comm_pattern,
+                           host_A->get_lower(), local_ptr);
+    }
+
+    if (requires_host_buffer) {
+        auto host_buffer_view =
+            gko::array<scalar>::view(exec, nnz, local_ptr);
+        auto target_buffer_view =
+            gko::array<scalar>::view(device_exec, nnz, local_ptr_2);
+        target_buffer_view = host_buffer_view;
+    }
+
+    if (requires_host_buffer) {
+        auto host_buffer_view =
+            gko::array<scalar>::view(exec, nnz, local_ptr);
+        auto target_buffer_view =
+            gko::array<scalar>::view(device_exec, nnz, local_ptr_2);
+        target_buffer_view = host_buffer_view;
+    }
+    // copy interface values
+    auto comm = *exec_handler.get_communicator().get();
+    if (owner) {
+        using Coo = gko::matrix::Coo<scalar, label>;
+        std::shared_ptr<const Coo> mtx;
+        label loc_ctr{1};
+        label nloc_ctr{0};
+        label host_interface_ctr{0};
+        label tag = 0;
+        label comm_rank, comm_size;
+        scalar *recv_buffer_ptr;
+        scalar *recv_buffer_ptr_2;
+        std::vector<scalar> host_recv_buffer;
+        // TODO FIXME can this be taken from sparsity/comm_pattern
+        // label remain_host_interfaces = host_A->get_interface_size();
+        label remain_host_interfaces = 0; // host_A->get_interface_size();
+        for (auto [is_local, comm_rank] : local_interfaces) {
+            label &ctr = (is_local) ? loc_ctr : nloc_ctr;
+            if (is_local) {
+                // TODO make it work for other types
+                mtx = gko::as<Coo>(
+                    gko::as<CombinationMatrix<Coo>>(
+                        dist_A->get_local_matrix())
+                        ->get_combination()
+                        ->get_operators()[ctr]);
+                comm_size =
+                local_sparsity->spans[ctr].length();
+            } else {
+                mtx = gko::as<Coo>(
+                    gko::as<CombinationMatrix<Coo>>(
+                        dist_A->get_non_local_matrix())
+                        ->get_combination()
+                        ->get_operators()[ctr]);
+                comm_size =
+                    non_local_sparsity->spans[ctr].length();
+            }
+
+            if (requires_host_buffer) {
+                host_recv_buffer.resize(comm_size);
+                recv_buffer_ptr = host_recv_buffer.data();
+                recv_buffer_ptr_2 =
+                    const_cast<scalar *>(mtx->get_const_values());
+            } else {
+                recv_buffer_ptr =
+                    const_cast<scalar *>(mtx->get_const_values());
+            }
+
+            if (comm_rank != rank) {
+                comm.recv(device_exec, recv_buffer_ptr, comm_size,
+                          comm_rank, tag);
+                if (requires_host_buffer) {
+                    auto host_buffer_view = gko::array<scalar>::view(
+                        exec, comm_size, recv_buffer_ptr);
+                    auto target_buffer_view = gko::array<scalar>::view(
+                        device_exec, comm_size, recv_buffer_ptr_2);
+                    target_buffer_view = host_buffer_view;
+                }
+
+            } else {
+                // if data is already on this rank
+                auto data_view = gko::array<scalar>::const_view(
+                    exec, comm_size,
+                    host_A->get_interface_data(host_interface_ctr));
+
+                // TODO FIXME this needs target executor
+                recv_buffer_ptr =
+                    const_cast<scalar *>(mtx->get_const_values());
+                auto target_view = gko::array<scalar>::view(
+                    device_exec, comm_size, recv_buffer_ptr);
+
+                target_view = data_view;
+
+                host_interface_ctr++;
+                remain_host_interfaces--;
+            }
+
+            ctr++;
+            // interface values need to be multiplied by -1
+            using vec = gko::matrix::Dense<scalar>;
+            recv_buffer_ptr = const_cast<scalar
+            *>(mtx->get_const_values()); auto neg_one =
+            gko::initialize<vec>({-1.0}, device_exec); auto
+            interface_dense =
+                vec::create(device_exec, gko::dim<2>{static_cast<gko::size_type>(comm_size), 1},
+                            gko::array<scalar>::view(device_exec,
+                            comm_size,
+                                                     recv_buffer_ptr),
+                            1);
+
+            interface_dense->scale(neg_one);
+        }
+    } else {
+        // the non owner has send all its interfaces to owner
+        // thus all values need to be communicated to the owner as well
+        label num_interfaces = src_comm_pattern->target_ids.get_size();
+        label tag = 0;
+        for (int i = 0; i < num_interfaces; i++) {
+            label comm_size =
+                src_comm_pattern->target_sizes.get_const_data()[i];
+            const scalar *send_buffer_ptr =
+            host_A->get_interface_data(i); comm.send(device_exec,
+            send_buffer_ptr, comm_size, owner_rank,
+                      tag);
+        }
+    }
+    // reorder updated values
+    if (owner) {
+        // NOTE local sparsity size includes the interfaces
+        using Coo = gko::matrix::Coo<scalar, label>;
+        using dim_type = gko::dim<2>::dimension_type;
+        std::shared_ptr<const Coo> local =
+            gko::as<Coo>(gko::as<CombinationMatrix<Coo>>(
+                             dist_A->get_local_matrix())
+                             ->get_combination()
+                             ->get_operators()[0]);
+        auto local_elements = local->get_num_stored_elements();
+        local_ptr = const_cast<scalar *>(local->get_const_values());
+        // TODO make sure this doesn't copy
+        // create a non owning dense matrix of local_values
+
+        auto row_collection =
+        gko::share(gko::matrix::Dense<scalar>::create(
+            device_exec,
+            gko::dim<2>{static_cast<dim_type>(local_elements), 1},
+            gko::array<scalar>::view(device_exec, local_elements,
+                                     local_ptr),
+            1));
+        auto mapping_view = gko::array<label>::view(
+            exec, local_elements,
+            local_sparsity->ldu_mapping.get_data());
+
+
+        // TODO this needs to copy ldu_mapping to the device
+        auto dense_vec = row_collection->clone();
+        // auto dense_vec =
+        // gko::share(gko::matrix::Dense<scalar>::create(exec,
+        // row_collection->get_size()));
+
+        dense_vec->row_gather(&mapping_view, row_collection.get());
+    }
+};
+
+
+template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
+static std::shared_ptr<
+    RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>>
+RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>::create(
+    const ExecutorHandler &exec_handler, word matrix_format,
+    const Repartitioner &repartitioner,
+    std::shared_ptr<const HostMatrixWrapper> host_A)
+{
+    label rank = exec_handler.get_rank();
+    auto exec = exec_handler.get_ref_exec();
+    auto comm = *exec_handler.get_communicator().get();
+
+    // repartition things here first by creating device_matrix_data on
+    // device with correct size on each rank
+    //
+    // NOTE Start by copying first and repartition later
+    auto local_sparsity = host_A->compute_local_sparsity(exec);
+    auto non_local_sparsity = host_A->compute_non_local_sparsity(exec);
+
+    auto src_comm_pattern = host_A->create_communication_pattern();
+    auto repart_comm_pattern =
+        repartitioner.repartition_comm_pattern(exec_handler, src_comm_pattern);
+
+    // bool owner = repartitioner.is_owner(exec_handler);
+
+    auto [repart_loc_sparsity, repart_non_loc_sparsity, local_interfaces] =
+        repartitioner.repartition_sparsity(exec_handler, local_sparsity,
+                                           non_local_sparsity);
+
+    // create vector of inner type linops
+    // if fuse the vector contains only a single element
+    // thus we can unwrap it.
+    auto device_exec = exec_handler.get_device_exec();
+    auto local_linops = detail::generate_inner_linops(
+        matrix_format, false, device_exec, repart_loc_sparsity);
+    auto non_local_linops = detail::generate_inner_linops(
+        matrix_format, false, device_exec, repart_non_loc_sparsity);
+
+    auto [send_counts, send_offsets, recv_sizes, recv_offsets] =
+        repart_comm_pattern->send_recv_pattern();
+    auto recv_gather_idxs = repart_comm_pattern->total_rank_send_idx();
+
+    auto global_rows = repartitioner.get_orig_partition()->get_size();
+    gko::dim<2> global_dim{global_rows, global_rows};
+
+    std::shared_ptr<dist_mtx> dist_A;
+    bool fuse = false;
+    if (fuse) {
+    } else {
+        dist_A = gko::share(dist_mtx::create(
+            exec, comm, global_dim,
+            gko::share(
+                CombinationMatrix<gko::matrix::Coo<scalar, label>>::create(
+                    exec, repart_loc_sparsity->dim, local_linops)),
+            gko::share(
+                CombinationMatrix<gko::matrix::Coo<scalar, label>>::create(
+                    exec, repart_non_loc_sparsity->dim, non_local_linops)),
+            recv_sizes, recv_offsets, recv_gather_idxs));
+    }
+
+    detail::update_impl(exec_handler, matrix_format, repartitioner, host_A,
+                        dist_A, repart_loc_sparsity, repart_non_loc_sparsity,
+                        src_comm_pattern, local_interfaces);
+
+    return std::make_shared<RepartDistMatrix>(
+        exec, comm, repartitioner.get_repart_dim(), dist_A->get_size(),
+        std::move(dist_A), repart_loc_sparsity, repart_non_loc_sparsity,
+        src_comm_pattern, local_interfaces);
+}
+
+template class RepartDistMatrix<scalar, label, label>;
diff --git a/unitTests/MatrixWrapper/Distributed.C b/unitTests/MatrixWrapper/Distributed.C
index 61da19ea1..e2138667c 100644
--- a/unitTests/MatrixWrapper/Distributed.C
+++ b/unitTests/MatrixWrapper/Distributed.C
@@ -121,6 +121,9 @@ TEST_P(DistributedMatrixFixture, canCreateDistributeMatrix)
     auto repartitioner = Repartitioner(
         hostMatrix->get_local_nrows(), ranks_per_gpu, 0, exec);
 
+    gko::dim<2> global_vec_dim {repartitioner.get_orig_partition()->get_size(), 1};
+    gko::dim<2> local_vec_dim {repartitioner.get_repart_dim()[0], 1};
+
     std::map<label, vec> exp_local_size;
     exp_local_size.emplace(1, vec{9, 9, 9, 9});
     exp_local_size.emplace(2, vec{18, 0, 18, 0});
@@ -131,13 +134,30 @@ TEST_P(DistributedMatrixFixture, canCreateDistributeMatrix)
     exp_non_local_size.emplace(2, vec{6, 0, 6, 0});
     exp_non_local_size.emplace(4, vec{0, 0 ,0 ,0});
 
+    auto b = gko::experimental::distributed::Vector<scalar>::create(
+        exec.get_ref_exec(),
+        comm, global_vec_dim, local_vec_dim, 1
+        );
+    b->fill(1);
+
+    auto x = gko::experimental::distributed::Vector<scalar>::create(
+        exec.get_ref_exec(),
+        comm, global_vec_dim, local_vec_dim, 1
+        );
+    x->fill(0);
+
     auto distributed = RepartDistMatrix<scalar, label, label>::create(exec, "Coo", repartitioner, hostMatrix);
+    distributed->apply(b, x);
 
     ASSERT_EQ(distributed->get_local_matrix()->get_size()[0], exp_local_size[ranks_per_gpu][rank]);
     ASSERT_EQ(distributed->get_local_matrix()->get_size()[1], exp_local_size[ranks_per_gpu][rank]);
     ASSERT_EQ(distributed->get_non_local_matrix()->get_size()[0], exp_local_size[ranks_per_gpu][rank]);
     ASSERT_EQ(distributed->get_non_local_matrix()->get_size()[1], exp_non_local_size[ranks_per_gpu][rank]);
 
+    auto x_local = std::vector<scalar>(x->get_local_vector()->get_const_values(), x->get_local_vector()->get_const_values() + local_vec_dim[0]);
+    std::vector<scalar> exp_x_local {1};
+
+    ASSERT_EQ(x_local, exp_x_local);
 }
 
 int main(int argc, char *argv[])

From 7e77323e8ac8f2a93abfcc1f7cda4d3eff820d50 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Fri, 30 Aug 2024 17:03:23 +0200
Subject: [PATCH 62/67] add some more hostmatrix tests and getters

---
 include/OGL/MatrixWrapper/HostMatrix.H | 15 ++++++++++++++-
 src/MatrixWrapper/HostMatrix.C         | 17 +++++++++++++----
 unitTests/MatrixWrapper/HostMatrix.C   |  8 ++++++++
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/include/OGL/MatrixWrapper/HostMatrix.H b/include/OGL/MatrixWrapper/HostMatrix.H
index 2fc9af02d..30376adfb 100644
--- a/include/OGL/MatrixWrapper/HostMatrix.H
+++ b/include/OGL/MatrixWrapper/HostMatrix.H
@@ -110,6 +110,12 @@ private:
 
     const FieldField<Field, scalar> &interfaceBouCoeffs_;
 
+    /* number of elements on each interfach */
+    mutable std::vector<label> interface_length_;
+
+    /* Pointer to interface data*/
+    mutable std::vector<const scalar *> interface_ptr_;
+
     /* Iterates all interfaces and collects the coefficients into a vector
     **
     ** @param local whether local or non local coefficients should be collected
@@ -236,9 +242,16 @@ public:
 
     const scalar *get_diag() const { return diag_; }
 
+    std::vector<label> get_interface_length() const
+    {
+        return interface_length_;
+    };
+
+    label get_num_interfaces() const { return interface_ptr_.size(); };
+
     const scalar *get_interface_data(label i) const
     {
-        FatalErrorInFunction << " not implemented " << abort(FatalError);
+        return interface_ptr_[i];
     }
 
     bool get_symmetric() const { return symmetric_; }
diff --git a/src/MatrixWrapper/HostMatrix.C b/src/MatrixWrapper/HostMatrix.C
index d5a368372..9bef881e4 100644
--- a/src/MatrixWrapper/HostMatrix.C
+++ b/src/MatrixWrapper/HostMatrix.C
@@ -140,16 +140,25 @@ HostMatrixWrapper::HostMatrixWrapper(
       lower_(lower),
       scaling_(solverControls.lookupOrDefault<scalar>("scaling", 1)),
       nrows_(nrows),
-      local_interface_nnz_{count_interface_nnz(interfaces, false)},
       upper_nnz_(upper_nnz),
       symmetric_(symmetric),
+      local_interface_nnz_{count_interface_nnz(interfaces, false)},
       non_diag_nnz_(2 * upper_nnz_),
       local_matrix_nnz_(nrows_ + 2 * upper_nnz_),
       local_matrix_w_interfaces_nnz_(local_matrix_nnz_ + local_interface_nnz_),
+      non_local_matrix_nnz_{count_interface_nnz(interfaces, true)},
       interfaces_(interfaces),
-      interfaceBouCoeffs_(interfaceBouCoeffs),
-      non_local_matrix_nnz_{count_interface_nnz(interfaces, true)}
-{}
+      interfaceBouCoeffs_(interfaceBouCoeffs)
+{
+    for (label i = 0; i < interfaces.size(); i++) {
+        if (interface_getter(interfaces, i) == nullptr) {
+            continue;
+        }
+        const auto iface{interface_getter(interfaces, i)};
+        interface_length_.push_back(iface->interface().faceCells().size());
+        interface_ptr_.push_back(interfaceBouCoeffs[i].begin());
+    }
+}
 
 HostMatrixWrapper::HostMatrixWrapper(
     const ExecutorHandler &exec, const objectRegistry &db,
diff --git a/unitTests/MatrixWrapper/HostMatrix.C b/unitTests/MatrixWrapper/HostMatrix.C
index c5bd9ef43..29aa3c7e3 100644
--- a/unitTests/MatrixWrapper/HostMatrix.C
+++ b/unitTests/MatrixWrapper/HostMatrix.C
@@ -101,6 +101,8 @@ TEST(HostMatrix, returnsCorrectSize)
     auto mesh = ((HostMatrixEnvironment *)global_env)->mesh;
     auto hostMatrix = ((HostMatrixEnvironment *)global_env)->hostMatrix;
 
+    std::vector<label> exp_interface_length {3, 3};
+
     // the local size is 9
     EXPECT_EQ(mesh->C().size(), 9);
     // which results in a 9x9 matrix
@@ -108,6 +110,12 @@ TEST(HostMatrix, returnsCorrectSize)
     EXPECT_EQ(hostMatrix->get_size()[1], 9);
 
     EXPECT_EQ(hostMatrix->get_local_nrows(), 9);
+
+    // each rank has 2 interfaces to neighbouring rank
+    EXPECT_EQ(hostMatrix->get_num_interfaces(), 2);
+    EXPECT_EQ(hostMatrix->get_local_matrix_nnz(), 33);
+
+    EXPECT_EQ(hostMatrix->get_interface_length(), exp_interface_length);
 }
 
 

From edd60a95db75c5715179e7a750d77e20b53c71b7 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Tue, 3 Sep 2024 15:23:59 +0200
Subject: [PATCH 63/67] add working distributed unit test

---
 include/OGL/CommunicationPattern.H          |   6 +
 include/OGL/MatrixWrapper/Combination.H     |  19 +-
 include/OGL/MatrixWrapper/SparsityPattern.H |   2 +
 include/OGL/Repartitioner.H                 |   2 +-
 include/OGL/common.H                        |  15 +-
 src/CommunicationPattern.C                  |  18 ++
 src/MatrixWrapper/Distributed.C             | 143 +++++----
 src/MatrixWrapper/SparsityPattern.C         |   6 +
 src/Repartitioner.C                         |  33 ++-
 src/common.C                                |  68 ++++-
 unitTests/MatrixWrapper/Distributed.C       | 304 +++++++++++++++++---
 unitTests/MatrixWrapper/HostMatrix.C        |  85 ++++--
 unitTests/Repartitioner.C                   |  37 ++-
 13 files changed, 568 insertions(+), 170 deletions(-)

diff --git a/include/OGL/CommunicationPattern.H b/include/OGL/CommunicationPattern.H
index 7fa83c2a0..6a88bf289 100644
--- a/include/OGL/CommunicationPattern.H
+++ b/include/OGL/CommunicationPattern.H
@@ -157,6 +157,12 @@ struct CommunicationPattern {
      */
     gko::array<label> total_rank_send_idx() const;
 
+    /* @brief concatenate all separate send idxs arrays into one contiguous
+     * array
+     */
+    gko::array<label> compute_recv_gather_idxs(
+        const ExecutorHandler &exec_handler) const;
+
 
     /** Computes the send and recv pattern for ginkgos distributed matrix SpMV
      */
diff --git a/include/OGL/MatrixWrapper/Combination.H b/include/OGL/MatrixWrapper/Combination.H
index 82d67cac9..a992ddeb2 100644
--- a/include/OGL/MatrixWrapper/Combination.H
+++ b/include/OGL/MatrixWrapper/Combination.H
@@ -8,6 +8,7 @@
 
 #include "fvCFD.H"
 
+
 /* @brief The CombinationMatrix class is a wrapper around Ginkgos combination
  * linear operator. It provides an interface for creating linear combinations of
  * a given inner matrix type and a read implementation
@@ -176,7 +177,7 @@ protected:
      * */
     void apply_impl(const gko::LinOp *b, gko::LinOp *x) const override
     {
-        if (this->get_size()[0] > 0) {
+        if (this->get_size()[0] > 0 && this->get_size()[1] > 0) {
             this->comb_->apply(b, x);
         }
     }
@@ -187,7 +188,7 @@ protected:
     void apply_impl(const gko::LinOp *alpha, const gko::LinOp *b,
                     const gko::LinOp *beta, gko::LinOp *x) const override
     {
-        if (this->get_size()[0] > 0) {
+        if (this->get_size()[0] > 0 && this->get_size()[1] > 0) {
             this->comb_->apply(alpha, b, beta, x);
         }
     }
@@ -195,3 +196,17 @@ protected:
 private:
     std::shared_ptr<gko::Combination<scalar>> comb_;
 };
+
+
+namespace detail {
+std::shared_ptr<gko::matrix::Coo<scalar, label>> convert_combination_to_coo(
+    std::shared_ptr<const gko::Executor> exec,
+    std::shared_ptr<const gko::LinOp> in)
+{
+    auto out = gko::share(gko::matrix::Coo<scalar, label>::create(exec));
+    gko::as<CombinationMatrix<gko::matrix::Coo<scalar, label>>>(in)->convert_to(
+        out.get());
+
+    return out;
+}
+}  // namespace detail
diff --git a/include/OGL/MatrixWrapper/SparsityPattern.H b/include/OGL/MatrixWrapper/SparsityPattern.H
index da68be505..c041b7450 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.H
+++ b/include/OGL/MatrixWrapper/SparsityPattern.H
@@ -24,6 +24,8 @@ void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
                                   std::vector<label> &ldu_mapping, label rank,
                                   label ranks_per_gpu);
 
+void compress_cols(gko::array<label> &in);
+
 /* @brief computes the dimensions of square matrix based row index array
  *
  * @note assumes that rows are ordered
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index 621cdecc8..d5fc4a4ef 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -158,7 +158,7 @@ public:
      * returns a vector of locality information.
      * @param non_local_cols vector in comm rank local indices
      */
-    std::vector<bool> build_non_local_interfaces(
+    std::vector<std::pair<bool, label>> build_non_local_interfaces(
         const ExecutorHandler &exec_handler,
         std::shared_ptr<
             const gko::experimental::distributed::Partition<label, label>>
diff --git a/include/OGL/common.H b/include/OGL/common.H
index 237807777..3a23e0903 100644
--- a/include/OGL/common.H
+++ b/include/OGL/common.H
@@ -100,7 +100,8 @@ std::ostream &operator<<(
     std::ostream &os,
     const std::shared_ptr<gko::matrix::Dense<scalar>> array_in);
 
-std::ostream &operator<<(std::ostream &os, const std::vector<label> &in);
+template <typename ValueType>
+std::ostream &operator<<(std::ostream &os, const std::vector<ValueType> &in);
 
 
 void export_system(const word fieldName, const gko::matrix::Csr<scalar> *A,
@@ -135,5 +136,17 @@ label get_next_caching(word sys_matrix_name, const objectRegistry &db);
 std::shared_ptr<gko::array<label>> convert_to_array(
     const std::vector<label> &in);
 
+std::pair<label, const scalar *> get_val(
+    std::shared_ptr<const gko::matrix::Coo<scalar, label>> in);
+
+std::pair<label, const label *> get_col(
+    std::shared_ptr<const gko::matrix::Coo<scalar, label>> in);
+
+std::pair<label, const label *> get_row(
+    std::shared_ptr<const gko::matrix::Coo<scalar, label>> in);
+
 std::vector<label> convert_to_vector(const gko::array<label> &in);
+
+std::vector<scalar> convert_to_vector(
+    std::shared_ptr<const gko::matrix::Diagonal<scalar>> in);
 }  // namespace Foam
diff --git a/src/CommunicationPattern.C b/src/CommunicationPattern.C
index e0aad54ea..1ddca69c5 100644
--- a/src/CommunicationPattern.C
+++ b/src/CommunicationPattern.C
@@ -211,6 +211,24 @@ gko::array<label> CommunicationPattern::total_rank_send_idx() const
 }
 
 
+gko::array<label> CommunicationPattern::compute_recv_gather_idxs(
+    const ExecutorHandler &exec_handler) const
+{
+    auto exec = exec_handler.get_ref_exec();
+    auto comm = *exec_handler.get_communicator().get();
+    auto rs_idx = total_rank_send_idx();
+    auto pattern = send_recv_pattern();
+    auto recv_buffer =
+        gko::array<label>(exec_handler.get_ref_exec(), rs_idx.get_size());
+
+    comm.all_to_all_v(exec, rs_idx.get_const_data(), pattern.send_counts.data(),
+                      pattern.send_offsets.data(), recv_buffer.get_data(),
+                      pattern.recv_counts.data(), pattern.recv_offsets.data());
+
+    return recv_buffer;
+}
+
+
 AllToAllPattern CommunicationPattern::send_recv_pattern() const
 {
     auto comm = *exec_handler.get_communicator().get();
diff --git a/src/MatrixWrapper/Distributed.C b/src/MatrixWrapper/Distributed.C
index 5338feec3..de22879c2 100644
--- a/src/MatrixWrapper/Distributed.C
+++ b/src/MatrixWrapper/Distributed.C
@@ -14,16 +14,16 @@ std::vector<std::shared_ptr<const gko::LinOp>> detail::generate_inner_linops(
         gko::array<scalar> coeffs(exec, end - begin);
         coeffs.fill(0.0);
 
-        auto row_idxs_view =
-            gko::array<label>::view(exec->get_master(), end - begin,
-                                    sparsity->row_idxs.get_data() + begin);
-        auto col_idxs_view =
-            gko::array<label>::view(exec->get_master(), end - begin,
-                                    sparsity->col_idxs.get_data() + begin);
+        auto row_idxs_arr = gko::array<label>(
+            exec->get_master(), sparsity->row_idxs.get_const_data() + begin,
+            sparsity->row_idxs.get_const_data() + end);
+        auto col_idxs_arr = gko::array<label>(
+            exec->get_master(), sparsity->col_idxs.get_const_data() + begin,
+            sparsity->col_idxs.get_const_data() + end);
 
         auto create_function = [&](word format) {
             return gko::share(gko::matrix::Coo<scalar, label>::create(
-                exec, sparsity->dim, coeffs, col_idxs_view, row_idxs_view));
+                exec, sparsity->dim, coeffs, col_idxs_arr, row_idxs_arr));
         };
         lin_ops.push_back(create_function(matrix_format));
     }
@@ -40,7 +40,8 @@ void detail::update_impl(
     std::shared_ptr<const SparsityPattern> local_sparsity,
     std::shared_ptr<const SparsityPattern> non_local_sparsity,
     std::shared_ptr<const CommunicationPattern> src_comm_pattern,
-    std::vector<std::pair<bool, label>> local_interfaces) {
+    std::vector<std::pair<bool, label>> local_interfaces)
+{
     auto exec = exec_handler.get_ref_exec();
     auto device_exec = exec_handler.get_device_exec();
     auto ranks_per_gpu = repartitioner.get_ranks_per_gpu();
@@ -51,6 +52,7 @@ void detail::update_impl(
     bool owner = repartitioner.is_owner(exec_handler);
     label nrows = host_A->get_local_nrows();
     label local_matrix_nnz = host_A->get_local_matrix_nnz();
+
     // size + padding has to be local_matrix_nnz
     auto diag_comm_pattern = compute_gather_to_owner_counts(
         exec_handler, ranks_per_gpu, nrows, local_matrix_nnz,
@@ -61,7 +63,7 @@ void detail::update_impl(
         local_matrix_nnz - upper_nnz);
     auto lower_comm_pattern =
         compute_gather_to_owner_counts(exec_handler, ranks_per_gpu, upper_nnz,
-                                 local_matrix_nnz, upper_nnz, nrows);
+                                       local_matrix_nnz, upper_nnz, nrows);
 
     scalar *local_ptr;
     scalar *local_ptr_2;
@@ -71,29 +73,27 @@ void detail::update_impl(
     std::vector<scalar> loc_buffer;
     if (owner) {
         using Coo = gko::matrix::Coo<scalar, label>;
-        auto local_mtx = dist_A->get_local_matrix();
 
         // TODO make it work with other matrix types
-        std::shared_ptr<const Coo> local =
-            gko::as<Coo>(gko::as<CombinationMatrix<Coo>>(
-                             dist_A->get_local_matrix())
-                             ->get_combination()
-                             ->get_operators()[0]);
+        std::shared_ptr<const Coo> local = gko::as<Coo>(
+            gko::as<CombinationMatrix<Coo>>(dist_A->get_local_matrix())
+                ->get_combination()
+                ->get_operators()[0]);
+
         nnz = local->get_num_stored_elements();
         if (requires_host_buffer) {
             loc_buffer.resize(nnz);
             local_ptr = loc_buffer.data();
-            local_ptr_2 = const_cast<scalar
-            *>(local->get_const_values());
+            local_ptr_2 = const_cast<scalar *>(local->get_const_values());
         } else {
             local_ptr = const_cast<scalar *>(local->get_const_values());
         }
     }
-    communicate_values(exec_handler, diag_comm_pattern,
-    host_A->get_diag(),
+
+    communicate_values(exec_handler, diag_comm_pattern, host_A->get_diag(),
+                       local_ptr);
+    communicate_values(exec_handler, upper_comm_pattern, host_A->get_upper(),
                        local_ptr);
-    communicate_values(exec_handler, upper_comm_pattern,
-                       host_A->get_upper(), local_ptr);
 
     if (host_A->get_symmetric()) {
         // TODO FIXME
@@ -106,20 +106,12 @@ void detail::update_impl(
     }
 
     if (requires_host_buffer) {
-        auto host_buffer_view =
-            gko::array<scalar>::view(exec, nnz, local_ptr);
+        auto host_buffer_view = gko::array<scalar>::view(exec, nnz, local_ptr);
         auto target_buffer_view =
             gko::array<scalar>::view(device_exec, nnz, local_ptr_2);
         target_buffer_view = host_buffer_view;
     }
 
-    if (requires_host_buffer) {
-        auto host_buffer_view =
-            gko::array<scalar>::view(exec, nnz, local_ptr);
-        auto target_buffer_view =
-            gko::array<scalar>::view(device_exec, nnz, local_ptr_2);
-        target_buffer_view = host_buffer_view;
-    }
     // copy interface values
     auto comm = *exec_handler.get_communicator().get();
     if (owner) {
@@ -133,28 +125,22 @@ void detail::update_impl(
         scalar *recv_buffer_ptr;
         scalar *recv_buffer_ptr_2;
         std::vector<scalar> host_recv_buffer;
-        // TODO FIXME can this be taken from sparsity/comm_pattern
-        // label remain_host_interfaces = host_A->get_interface_size();
-        label remain_host_interfaces = 0; // host_A->get_interface_size();
+        label remain_host_interfaces = host_A->get_num_interfaces();
         for (auto [is_local, comm_rank] : local_interfaces) {
             label &ctr = (is_local) ? loc_ctr : nloc_ctr;
             if (is_local) {
                 // TODO make it work for other types
                 mtx = gko::as<Coo>(
-                    gko::as<CombinationMatrix<Coo>>(
-                        dist_A->get_local_matrix())
+                    gko::as<CombinationMatrix<Coo>>(dist_A->get_local_matrix())
                         ->get_combination()
                         ->get_operators()[ctr]);
-                comm_size =
-                local_sparsity->spans[ctr].length();
+                comm_size = local_sparsity->spans[ctr].length();
             } else {
-                mtx = gko::as<Coo>(
-                    gko::as<CombinationMatrix<Coo>>(
-                        dist_A->get_non_local_matrix())
-                        ->get_combination()
-                        ->get_operators()[ctr]);
-                comm_size =
-                    non_local_sparsity->spans[ctr].length();
+                mtx = gko::as<Coo>(gko::as<CombinationMatrix<Coo>>(
+                                       dist_A->get_non_local_matrix())
+                                       ->get_combination()
+                                       ->get_operators()[ctr]);
+                comm_size = non_local_sparsity->spans[ctr].length();
             }
 
             if (requires_host_buffer) {
@@ -163,13 +149,12 @@ void detail::update_impl(
                 recv_buffer_ptr_2 =
                     const_cast<scalar *>(mtx->get_const_values());
             } else {
-                recv_buffer_ptr =
-                    const_cast<scalar *>(mtx->get_const_values());
+                recv_buffer_ptr = const_cast<scalar *>(mtx->get_const_values());
             }
 
             if (comm_rank != rank) {
-                comm.recv(device_exec, recv_buffer_ptr, comm_size,
-                          comm_rank, tag);
+                comm.recv(device_exec, recv_buffer_ptr, comm_size, comm_rank,
+                          tag);
                 if (requires_host_buffer) {
                     auto host_buffer_view = gko::array<scalar>::view(
                         exec, comm_size, recv_buffer_ptr);
@@ -185,8 +170,7 @@ void detail::update_impl(
                     host_A->get_interface_data(host_interface_ctr));
 
                 // TODO FIXME this needs target executor
-                recv_buffer_ptr =
-                    const_cast<scalar *>(mtx->get_const_values());
+                recv_buffer_ptr = const_cast<scalar *>(mtx->get_const_values());
                 auto target_view = gko::array<scalar>::view(
                     device_exec, comm_size, recv_buffer_ptr);
 
@@ -199,15 +183,14 @@ void detail::update_impl(
             ctr++;
             // interface values need to be multiplied by -1
             using vec = gko::matrix::Dense<scalar>;
-            recv_buffer_ptr = const_cast<scalar
-            *>(mtx->get_const_values()); auto neg_one =
-            gko::initialize<vec>({-1.0}, device_exec); auto
-            interface_dense =
-                vec::create(device_exec, gko::dim<2>{static_cast<gko::size_type>(comm_size), 1},
-                            gko::array<scalar>::view(device_exec,
-                            comm_size,
-                                                     recv_buffer_ptr),
-                            1);
+            recv_buffer_ptr = const_cast<scalar *>(mtx->get_const_values());
+            auto neg_one = gko::initialize<vec>({-1.0}, device_exec);
+            auto interface_dense = vec::create(
+                device_exec,
+                gko::dim<2>{static_cast<gko::size_type>(comm_size), 1},
+                gko::array<scalar>::view(device_exec, comm_size,
+                                         recv_buffer_ptr),
+                1);
 
             interface_dense->scale(neg_one);
         }
@@ -219,37 +202,33 @@ void detail::update_impl(
         for (int i = 0; i < num_interfaces; i++) {
             label comm_size =
                 src_comm_pattern->target_sizes.get_const_data()[i];
-            const scalar *send_buffer_ptr =
-            host_A->get_interface_data(i); comm.send(device_exec,
-            send_buffer_ptr, comm_size, owner_rank,
-                      tag);
+            const scalar *send_buffer_ptr = host_A->get_interface_data(i);
+            comm.send(device_exec, send_buffer_ptr, comm_size, owner_rank, tag);
         }
     }
+
     // reorder updated values
     if (owner) {
         // NOTE local sparsity size includes the interfaces
         using Coo = gko::matrix::Coo<scalar, label>;
         using dim_type = gko::dim<2>::dimension_type;
-        std::shared_ptr<const Coo> local =
-            gko::as<Coo>(gko::as<CombinationMatrix<Coo>>(
-                             dist_A->get_local_matrix())
-                             ->get_combination()
-                             ->get_operators()[0]);
+
+        std::shared_ptr<const Coo> local = gko::as<Coo>(
+            gko::as<CombinationMatrix<Coo>>(dist_A->get_local_matrix())
+                ->get_combination()
+                ->get_operators()[0]);
+
         auto local_elements = local->get_num_stored_elements();
         local_ptr = const_cast<scalar *>(local->get_const_values());
+
         // TODO make sure this doesn't copy
         // create a non owning dense matrix of local_values
-
-        auto row_collection =
-        gko::share(gko::matrix::Dense<scalar>::create(
-            device_exec,
-            gko::dim<2>{static_cast<dim_type>(local_elements), 1},
-            gko::array<scalar>::view(device_exec, local_elements,
-                                     local_ptr),
+        auto row_collection = gko::share(gko::matrix::Dense<scalar>::create(
+            device_exec, gko::dim<2>{static_cast<dim_type>(local_elements), 1},
+            gko::array<scalar>::view(device_exec, local_elements, local_ptr),
             1));
         auto mapping_view = gko::array<label>::view(
-            exec, local_elements,
-            local_sparsity->ldu_mapping.get_data());
+            exec, local_elements, local_sparsity->ldu_mapping.get_data());
 
 
         // TODO this needs to copy ldu_mapping to the device
@@ -292,6 +271,8 @@ RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>::create(
         repartitioner.repartition_sparsity(exec_handler, local_sparsity,
                                            non_local_sparsity);
 
+    compress_cols(repart_non_loc_sparsity->col_idxs);
+
     // create vector of inner type linops
     // if fuse the vector contains only a single element
     // thus we can unwrap it.
@@ -303,7 +284,13 @@ RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>::create(
 
     auto [send_counts, send_offsets, recv_sizes, recv_offsets] =
         repart_comm_pattern->send_recv_pattern();
-    auto recv_gather_idxs = repart_comm_pattern->total_rank_send_idx();
+
+    // recv_gather_idxs are send upon creation to ginkgo distributed matrix
+    // to partner ranks. This sets the send_sizes_ on the partner ranks.
+    // Thus recv_gather_idxs are local indices of comm partner rank of
+    // interfaces.
+    auto recv_gather_idxs =
+        repart_comm_pattern->compute_recv_gather_idxs(exec_handler);
 
     auto global_rows = repartitioner.get_orig_partition()->get_size();
     gko::dim<2> global_dim{global_rows, global_rows};
diff --git a/src/MatrixWrapper/SparsityPattern.C b/src/MatrixWrapper/SparsityPattern.C
index 6e6cb5408..cdd66eaf2 100644
--- a/src/MatrixWrapper/SparsityPattern.C
+++ b/src/MatrixWrapper/SparsityPattern.C
@@ -12,6 +12,12 @@ gko::dim<2> compute_dimensions(const std::vector<label> &rows)
     return gko::dim<2>{num_rows, num_rows};
 }
 
+void compress_cols(gko::array<label> &in)
+{
+    for (label i = 0; i < in.get_size(); i++) {
+        in.get_data()[i] = i;
+    }
+};
 
 void make_ldu_mapping_consecutive(const AllToAllPattern &comm_pattern,
                                   std::vector<label> &ldu_mapping, label rank,
diff --git a/src/Repartitioner.C b/src/Repartitioner.C
index e0a2b45c9..38904381c 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -219,17 +219,18 @@ Repartitioner::repartition_sparsity(
         return std::make_tuple<std::shared_ptr<SparsityPattern>,
                                std::shared_ptr<SparsityPattern>,
                                std::vector<std::pair<bool, label>>>(
-            std::move(local_pattern), std::move(non_local_pattern), {});
+            std::move(local_pattern), std::move(non_local_pattern),
+            std::move(is_local));
     } else {
         return std::make_tuple<std::shared_ptr<SparsityPattern>,
                                std::shared_ptr<SparsityPattern>,
                                std::vector<std::pair<bool, label>>>(
             std::make_shared<SparsityPattern>(exec),
-            std::make_shared<SparsityPattern>(exec), {});
+            std::make_shared<SparsityPattern>(exec), std::move(is_local));
     }
 }
 
-std::vector<bool> Repartitioner::build_non_local_interfaces(
+std::vector<std::pair<bool, label>> Repartitioner::build_non_local_interfaces(
     const ExecutorHandler &exec_handler,
     std::shared_ptr<
         const gko::experimental::distributed::Partition<label, label>>
@@ -243,7 +244,7 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
     std::vector<gko::span> &non_local_spans) const
 {
     auto rank = exec_handler.get_rank();
-    std::vector<bool> is_local;
+    std::vector<std::pair<bool, label>> is_local;
     std::vector<label> mark_keep;
 
     auto tmp_non_local_cols = detail::convert_to_global(
@@ -252,7 +253,6 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
     for (int i = 0; i < non_local_spans.size(); i++) {
         auto [begin, end] = non_local_spans[i];
         bool local = reparts_to_local(exec_handler, non_local_ranks[i]);
-        is_local.push_back(local);
 
         if (local) {
             gko::size_type rows_start = local_rows.size();
@@ -269,12 +269,14 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
             local_mapping.insert(local_mapping.end(),
                                  non_local_mapping.data() + begin,
                                  non_local_mapping.data() + end);
-            local_spans.emplace_back(rows_start,
-                                     rows_start + static_cast<gko::size_type>(end - begin));
+            local_spans.emplace_back(
+                rows_start,
+                rows_start + static_cast<gko::size_type>(end - begin));
             local_ranks.push_back(non_local_rank_origin[i]);
         } else {
             mark_keep.push_back(i);
         }
+        is_local.emplace_back(local, non_local_rank_origin[i]);
     }
 
     // remove data from non_local vectors
@@ -287,6 +289,7 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
     } else {
         std::vector<label> copy_rows, copy_cols, copy_mapping, copy_ranks;
         std::vector<gko::span> copy_spans;
+        label span_ctr{0};
         for (label i : mark_keep) {
             auto [begin, end] = non_local_spans[i];
             copy_rows.insert(copy_rows.end(), non_local_rows.data() + begin,
@@ -315,8 +318,12 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
                                 non_local_mapping.data() + begin,
                                 non_local_mapping.data() + end);
             copy_ranks.push_back(repart_comm_rank);
-            copy_spans.emplace_back(non_local_spans[i].begin,
-                                    non_local_spans[i].end);
+
+            // the spans are now consecutive based on all gathered spans,
+            // thus we need to make them consecutive based on kept interfaces
+            copy_spans.emplace_back(span_ctr,
+                                    span_ctr + non_local_spans[i].length());
+            span_ctr += non_local_spans[i].length();
         }
         non_local_rows = copy_rows;
         non_local_cols = copy_cols;
@@ -335,8 +342,7 @@ std::vector<bool> Repartitioner::build_non_local_interfaces(
 std::shared_ptr<const CommunicationPattern>
 Repartitioner::repartition_comm_pattern(
     const ExecutorHandler &exec_handler,
-    std::shared_ptr<const CommunicationPattern> src_comm_pattern
-) const
+    std::shared_ptr<const CommunicationPattern> src_comm_pattern) const
 {
     if (ranks_per_gpu_ == 1) {
         return src_comm_pattern;
@@ -400,8 +406,9 @@ Repartitioner::repartition_comm_pattern(
                           rank);
 
                 // the new offset is
-                auto offset = get_orig_partition()->get_range_bounds()[rank + i] -
-                              get_orig_partition()->get_range_bounds()[rank];
+                auto offset =
+                    get_orig_partition()->get_range_bounds()[rank + i] -
+                    get_orig_partition()->get_range_bounds()[rank];
 
                 std::transform(recv_buffer.begin(), recv_buffer.end(),
                                recv_buffer.begin(),
diff --git a/src/common.C b/src/common.C
index 1fec3f7a2..a242d1a57 100644
--- a/src/common.C
+++ b/src/common.C
@@ -152,6 +152,26 @@ std::shared_ptr<gko::array<label>> convert_to_array(
                                                in.begin(), in.end());
 }
 
+std::pair<label, const scalar *> get_val(
+    std::shared_ptr<const gko::matrix::Coo<scalar, label>> in)
+{
+    return std::make_pair(in->get_num_stored_elements(),
+                          in->get_const_values());
+}
+
+std::pair<label, const label *> get_col(
+    std::shared_ptr<const gko::matrix::Coo<scalar, label>> in)
+{
+    return std::make_pair(in->get_num_stored_elements(),
+                          in->get_const_col_idxs());
+}
+
+std::pair<label, const label *> get_row(
+    std::shared_ptr<const gko::matrix::Coo<scalar, label>> in)
+{
+    return std::make_pair(in->get_num_stored_elements(),
+                          in->get_const_row_idxs());
+}
 
 std::vector<label> convert_to_vector(const gko::array<label> &in)
 {
@@ -160,6 +180,14 @@ std::vector<label> convert_to_vector(const gko::array<label> &in)
 }
 
 
+std::vector<scalar> convert_to_vector(
+    std::shared_ptr<const gko::matrix::Diagonal<scalar>> in)
+{
+    return std::vector<scalar>(in->get_const_values(),
+                               in->get_const_values() + in->get_size()[0]);
+}
+
+
 std::ostream &operator<<(std::ostream &os,
                          const std::shared_ptr<gko::matrix::Dense<scalar>> in)
 {
@@ -167,7 +195,7 @@ std::ostream &operator<<(std::ostream &os,
     auto array = in->clone(ref_exec);
     label size = array->get_size()[0];
     os << size << " elements [";
-    if (size > 100) {
+    if (size > 200) {
         for (label i = 0; i < 9; i++) {
             os << array->at(i) << ", ";
         }
@@ -178,13 +206,14 @@ std::ostream &operator<<(std::ostream &os,
         os << array->at(size - 1) << "]\n";
     } else {
         for (label i = 0; i < size - 1; i++) {
-            os << "(" << i << ", " << array->at(i) << ") ";
+            os << array->at(i) << ", ";
         }
-        os << "(" << size - 1 << ", " << array->at(size - 1) << ")]\n";
+        os << ", " << array->at(size - 1) << "]\n";
     }
     return os;
 }
 
+template <>
 std::ostream &operator<<(std::ostream &os, const std::vector<label> &in)
 {
     label size = in.size();
@@ -192,7 +221,7 @@ std::ostream &operator<<(std::ostream &os, const std::vector<label> &in)
         return os;
     }
     os << size << " elements [";
-    if (size > 100) {
+    if (size > 200) {
         for (label i = 0; i < 9; i++) {
             os << in[i] << ", ";
         }
@@ -203,10 +232,37 @@ std::ostream &operator<<(std::ostream &os, const std::vector<label> &in)
         os << in[size - 1] << "]\n";
     } else {
         for (label i = 0; i < size - 1; i++) {
-            os << "(" << i << ", " << in[i] << ") ";
+            os << in[i] << ", ";
         }
-        os << "(" << size - 1 << ", " << in[size - 1] << ")]\n";
+        os << ", " << in[size - 1] << "]\n";
+    }
+    return os;
+}
+
+template <>
+std::ostream &operator<<(std::ostream &os, const std::vector<scalar> &in)
+{
+    label size = in.size();
+    if (size == 0) {
+        return os;
+    }
+    os << size << " elements [";
+    if (size > 200) {
+        for (label i = 0; i < 9; i++) {
+            os << in[i] << ", ";
+        }
+        os << in[10] << " ... ";
+        for (label i = size - 9; i < size - 1; i++) {
+            os << in[i] << ", ";
+        }
+        os << in[size - 1] << "]\n";
+    } else {
+        for (label i = 0; i < size - 1; i++) {
+            os << in[i] << ", ";
+        }
+        os << ", " << in[size - 1] << "]\n";
     }
     return os;
 };
+
 }  // namespace Foam
diff --git a/unitTests/MatrixWrapper/Distributed.C b/unitTests/MatrixWrapper/Distributed.C
index e2138667c..0729f8934 100644
--- a/unitTests/MatrixWrapper/Distributed.C
+++ b/unitTests/MatrixWrapper/Distributed.C
@@ -3,8 +3,8 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 
 
-#include "OGL/MatrixWrapper/HostMatrix.H"
 #include "OGL/MatrixWrapper/Distributed.H"
+#include "OGL/MatrixWrapper/HostMatrix.H"
 #include "OGL/Repartitioner.H"
 
 #include "gtest/gtest.h"
@@ -25,6 +25,18 @@ using vec = std::vector<label>;
 using vec_vec = std::vector<std::vector<label>>;
 using vec_vec_vec = std::vector<vec_vec>;
 
+using vec_vec_s = std::vector<std::vector<scalar>>;
+
+
+template <typename ValueType, typename IndexType>
+std::vector<ValueType> convert_to_vector(
+    std::pair<IndexType, const ValueType *> in)
+{
+    auto [size, ptr] = in;
+    return std::vector<ValueType>(ptr, ptr + size);
+}
+
+
 class Environment : public testing::Environment {
 public:
     void SetUp()
@@ -61,9 +73,9 @@ public:
         // to clean up output
         ::testing::TestEventListeners &listeners =
             ::testing::UnitTest::GetInstance()->listeners();
-        if (Foam::Pstream::myProcNo() != 0) {
-            delete listeners.Release(listeners.default_result_printer());
-        }
+        // if (Foam::Pstream::myProcNo() != 0) {
+        //     delete listeners.Release(listeners.default_result_printer());
+        // }
 
         word fieldName{"p"};
         dimensionSet ds{0, 0, 0, 0, 0};
@@ -73,15 +85,40 @@ public:
                            Foam::IOobject::MUST_READ),
             *mesh.get(), ds);
 
-        Foam::fvMatrix<scalar> fvMatrix{*field.get(), ds};
+        fvMatrix = std::make_shared<Foam::fvMatrix<scalar>>(*field.get(), ds);
 
         interfaces = field->boundaryField().scalarInterfaces();
 
         hostMatrix = std::make_shared<HostMatrixWrapper>(
-            *exec.get(), runTime_->thisDb(), mesh->lduAddr(), true,
-            fvMatrix.diag().data(), fvMatrix.upper().data(),
-            fvMatrix.lower().data(), fvMatrix.boundaryCoeffs(),
-            fvMatrix.internalCoeffs(), interfaces, dict, "fieldName", 0);
+            *exec.get(), runTime_->thisDb(), mesh->lduAddr(),
+            fvMatrix->symmetric(), fvMatrix->diag().data(),
+            fvMatrix->upper().data(), fvMatrix->lower().data(),
+            fvMatrix->boundaryCoeffs(), fvMatrix->internalCoeffs(), interfaces,
+            dict, "fieldName", 0);
+
+        // Set all matrix coefficients to some value
+        // in order to make the apply test give actual results
+        // this is required since creating the host matrix
+        // without actual DSL sets all values to zero
+        for (int i = 0; i < 9; i++) {
+            fvMatrix->diag().data()[i] = 2.0;
+        }
+        for (int i = 0; i < 12; i++) {
+            fvMatrix->upper().data()[i] = 1.0;
+        }
+        for (int i = 0; i < 12; i++) {
+            fvMatrix->lower().data()[i] = 1.0;
+        }
+        // set the interface value, we use get_interface_data here
+        // because that is more comfortable.
+        for (int i = 0; i < 2; i++) {
+            scalar *data =
+                const_cast<scalar *>(hostMatrix->get_interface_data(i));
+            for (int j = 0; j < 3; j++) {
+                // values are *-1.0 on interface on OF side
+                data[j] = -1.0;
+            }
+        }
     }
 
     Foam::lduInterfaceFieldPtrsList interfaces;
@@ -93,17 +130,16 @@ public:
     std::shared_ptr<GeometricField<scalar, Foam::fvPatchField, Foam::volMesh>>
         field;
     std::shared_ptr<const ExecutorHandler> exec;
+    std::shared_ptr<Foam::fvMatrix<scalar>> fvMatrix;
     std::shared_ptr<const HostMatrixWrapper> hostMatrix;
 };
 
 const testing::Environment *global_env =
     AddGlobalTestEnvironment(new Environment);
 
-class DistributedMatrixFixture :
-     public testing::TestWithParam<int> {
+class DistributedMatrixFixture : public testing::TestWithParam<int> {
 public:
-    ExecutorHandler exec =
-        *((Environment *)global_env)->exec.get();
+    ExecutorHandler exec = *((Environment *)global_env)->exec.get();
     label rank = exec.get_rank();
     const gko::experimental::mpi::communicator comm =
         *(exec.get_communicator().get());
@@ -112,52 +148,244 @@ public:
 INSTANTIATE_TEST_SUITE_P(DistributedMatrixFixtureInstantiation,
                          DistributedMatrixFixture, testing::Values(1, 2, 4));
 
-TEST_P(DistributedMatrixFixture, canCreateDistributeMatrix)
+TEST_P(DistributedMatrixFixture, canCreateDistributedMatrix)
 {
     /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
     label ranks_per_gpu = GetParam();
     auto mesh = ((Environment *)global_env)->mesh;
     auto hostMatrix = ((Environment *)global_env)->hostMatrix;
-    auto repartitioner = Repartitioner(
-        hostMatrix->get_local_nrows(), ranks_per_gpu, 0, exec);
+    auto repartitioner =
+        Repartitioner(hostMatrix->get_local_nrows(), ranks_per_gpu, 0, exec);
 
-    gko::dim<2> global_vec_dim {repartitioner.get_orig_partition()->get_size(), 1};
-    gko::dim<2> local_vec_dim {repartitioner.get_repart_dim()[0], 1};
+    gko::dim<2> global_vec_dim{repartitioner.get_orig_partition()->get_size(),
+                               1};
+    gko::dim<2> local_vec_dim{repartitioner.get_repart_dim()[0], 1};
 
     std::map<label, vec> exp_local_size;
     exp_local_size.emplace(1, vec{9, 9, 9, 9});
     exp_local_size.emplace(2, vec{18, 0, 18, 0});
-    exp_local_size.emplace(4, vec{36, 0 ,0 ,0});
+    exp_local_size.emplace(4, vec{36, 0, 0, 0});
+
+    auto distributed = RepartDistMatrix<scalar, label, label>::create(
+        exec, "Coo", repartitioner, hostMatrix);
+
+    ASSERT_EQ(distributed->get_local_matrix()->get_size()[0],
+              exp_local_size[ranks_per_gpu][rank]);
+    ASSERT_EQ(distributed->get_local_matrix()->get_size()[1],
+              exp_local_size[ranks_per_gpu][rank]);
+    ASSERT_EQ(distributed->get_non_local_matrix()->get_size()[0],
+              exp_local_size[ranks_per_gpu][rank]);
+    ASSERT_EQ(distributed->get_local_matrix()->get_size()[0],
+              exp_local_size[ranks_per_gpu][rank]);
+
+}
+
+TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectLocalMatrix)
+{
+    /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
+    label ranks_per_gpu = GetParam();
+    auto mesh = ((Environment *)global_env)->mesh;
+    auto hostMatrix = ((Environment *)global_env)->hostMatrix;
+    auto repartitioner =
+        Repartitioner(hostMatrix->get_local_nrows(), ranks_per_gpu, 0, exec);
+
+    gko::dim<2> global_vec_dim{repartitioner.get_orig_partition()->get_size(),
+                               1};
+    gko::dim<2> local_vec_dim{repartitioner.get_repart_dim()[0], 1};
+
+    std::map<label, vec> exp_local_size;
+    exp_local_size.emplace(1, vec{9, 9, 9, 9});
+    exp_local_size.emplace(2, vec{18, 0, 18, 0});
+    exp_local_size.emplace(4, vec{36, 0, 0, 0});
+
+    auto distributed = RepartDistMatrix<scalar, label, label>::create(
+        exec, "Coo", repartitioner, hostMatrix);
+
+    auto local = detail::convert_combination_to_coo(
+        exec.get_ref_exec(), distributed->get_local_matrix());
+
+    auto res_local_values = convert_to_vector(get_val(local));
+    std::map<label, vec_vec_s> exp_local_values;
+    std::vector<scalar> exp_val_1{2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1,
+                                  2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2,
+                                  1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2};
+    exp_local_values.emplace(
+        1, vec_vec_s{exp_val_1, exp_val_1, exp_val_1, exp_val_1});
+    std::vector<scalar> exp_val_2{
+        2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1,
+        2, 1, 1, 1, 2, 1, 1, 1, 2, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1,
+        1, 2, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 1, 1, 1};
+    exp_local_values.emplace(2, vec_vec_s{exp_val_2, {}, exp_val_2, {}});
+    std::vector<scalar> exp_val_4{
+        2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1,
+        1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1,
+        1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 2, 1, 1,
+        1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 1, 1, 2, 1, 1, 2, 1,
+        1, 1, 2, 1, 1, 1, 2, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1,
+        2, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 2, 1, 1, 1, 1, 1, 1,
+        1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1};
+    exp_local_values.emplace(4, vec_vec_s{exp_val_4, {}, {}, {}});
+
+    auto res_local_cols = convert_to_vector(get_col(local));
+    std::map<label, vec_vec> exp_local_cols;
+    vec exp_cols_1 = {0, 1, 3, 0, 1, 2, 4, 1, 2, 5, 0, 3, 4, 6, 1, 3, 4,
+                      5, 7, 2, 4, 5, 8, 3, 6, 7, 4, 6, 7, 8, 5, 7, 8};
+    exp_local_cols.emplace(
+        1, vec_vec{exp_cols_1, exp_cols_1, exp_cols_1, exp_cols_1});
+    vec exp_cols_2 = {
+        0,  1,  3,  0,  1,  2,  4,  1,  2,  5,  0,  3,  4,  6,  1,  3,  4,  5,
+        7,  2,  4,  5,  8,  3,  6,  7,  4,  6,  7,  8,  5,  7,  8,  9,  10, 12,
+        9,  10, 11, 13, 10, 11, 14, 9,  12, 13, 15, 10, 12, 13, 14, 16, 11, 13,
+        14, 17, 12, 15, 16, 13, 15, 16, 17, 14, 16, 17, 9,  12, 15, 2,  5,  8};
+    exp_local_cols.emplace(2, vec_vec{exp_cols_2, {}, exp_cols_2, {}});
+    vec exp_col_4 = {
+        0,  1,  3,  0,  1,  2,  4,  1,  2,  5,  0,  3,  4,  6,  1,  3,  4,  5,
+        7,  2,  4,  5,  8,  3,  6,  7,  4,  6,  7,  8,  5,  7,  8,  9,  10, 12,
+        9,  10, 11, 13, 10, 11, 14, 9,  12, 13, 15, 10, 12, 13, 14, 16, 11, 13,
+        14, 17, 12, 15, 16, 13, 15, 16, 17, 14, 16, 17, 18, 19, 21, 18, 19, 20,
+        22, 19, 20, 23, 18, 21, 22, 24, 19, 21, 22, 23, 25, 20, 22, 23, 26, 21,
+        24, 25, 22, 24, 25, 26, 23, 25, 26, 27, 28, 30, 27, 28, 29, 31, 28, 29,
+        32, 27, 30, 31, 33, 28, 30, 31, 32, 34, 29, 31, 32, 35, 30, 33, 34, 31,
+        33, 34, 35, 32, 34, 35, 9,  12, 15, 18, 19, 20, 2,  5,  8,  27, 28, 29,
+        6,  7,  8,  27, 30, 33, 15, 16, 17, 20, 23, 26};
+    exp_local_cols.emplace(4, vec_vec{exp_col_4, {}, {}, {}});
+
+    auto res_local_rows = convert_to_vector(get_row(local));
+    std::map<label, vec_vec> exp_local_rows;
+    vec exp_row_1 = {0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4,
+                     4, 4, 5, 5, 5, 5, 6, 6, 6, 7, 7, 7, 7, 8, 8, 8};
+    exp_local_rows.emplace(1,
+                           vec_vec{exp_row_1, exp_row_1, exp_row_1, exp_row_1});
+    vec exp_row_2 = {0,  0,  0,  1,  1,  1,  1,  2,  2,  2,  3,  3,  3,  3,  4,
+                     4,  4,  4,  4,  5,  5,  5,  5,  6,  6,  6,  7,  7,  7,  7,
+                     8,  8,  8,  9,  9,  9,  10, 10, 10, 10, 11, 11, 11, 12, 12,
+                     12, 12, 13, 13, 13, 13, 13, 14, 14, 14, 14, 15, 15, 15, 16,
+                     16, 16, 16, 17, 17, 17, 2,  5,  8,  9,  12, 15};
+    exp_local_rows.emplace(2, vec_vec{exp_row_2, {}, exp_row_2, {}});
+    vec exp_row_4 = {
+        0,  0,  0,  1,  1,  1,  1,  2,  2,  2,  3,  3,  3,  3,  4,  4,  4,  4,
+        4,  5,  5,  5,  5,  6,  6,  6,  7,  7,  7,  7,  8,  8,  8,  9,  9,  9,
+        10, 10, 10, 10, 11, 11, 11, 12, 12, 12, 12, 13, 13, 13, 13, 13, 14, 14,
+        14, 14, 15, 15, 15, 16, 16, 16, 16, 17, 17, 17, 18, 18, 18, 19, 19, 19,
+        19, 20, 20, 20, 21, 21, 21, 21, 22, 22, 22, 22, 22, 23, 23, 23, 23, 24,
+        24, 24, 25, 25, 25, 25, 26, 26, 26, 27, 27, 27, 28, 28, 28, 28, 29, 29,
+        29, 30, 30, 30, 30, 31, 31, 31, 31, 31, 32, 32, 32, 32, 33, 33, 33, 34,
+        34, 34, 34, 35, 35, 35, 2,  5,  8,  6,  7,  8,  9,  12, 15, 15, 16, 17,
+        18, 19, 20, 20, 23, 26, 27, 28, 29, 27, 30, 33};
+    exp_local_rows.emplace(4, vec_vec{exp_row_4, {}, {}, {}});
+
+    ASSERT_EQ(distributed->get_local_matrix()->get_size()[1],
+              exp_local_size[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_local_values, exp_local_values[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_local_rows, exp_local_rows[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_local_cols, exp_local_cols[ranks_per_gpu][rank]);
+}
+
+
+TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectNonLocalMatrix)
+{
+    /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
+    label ranks_per_gpu = GetParam();
+    auto mesh = ((Environment *)global_env)->mesh;
+    auto hostMatrix = ((Environment *)global_env)->hostMatrix;
+    auto repartitioner =
+        Repartitioner(hostMatrix->get_local_nrows(), ranks_per_gpu, 0, exec);
 
     std::map<label, vec> exp_non_local_size;
     exp_non_local_size.emplace(1, vec{6, 6, 6, 6});
     exp_non_local_size.emplace(2, vec{6, 0, 6, 0});
-    exp_non_local_size.emplace(4, vec{0, 0 ,0 ,0});
+    exp_non_local_size.emplace(4, vec{0, 0, 0, 0});
+
+    auto distributed = RepartDistMatrix<scalar, label, label>::create(
+        exec, "Coo", repartitioner, hostMatrix);
+
+    auto non_local = detail::convert_combination_to_coo(
+        exec.get_ref_exec(), distributed->get_non_local_matrix());
+
+    auto res_non_local_values = convert_to_vector(get_val(non_local));
+
+    std::map<label, vec_vec_s> exp_non_local_values;
+    std::vector<scalar> exp_non_local(6, 1.0);
+    exp_non_local_values.emplace(1, vec_vec_s(4, exp_non_local));
+    exp_non_local_values.emplace(
+        2, vec_vec_s{exp_non_local, {}, exp_non_local, {}});
+    exp_non_local_values.emplace(4, vec_vec_s{{}, {}, {}, {}});
+
+    auto res_non_local_cols = convert_to_vector(get_col(non_local));
+    std::map<label, vec_vec> exp_non_local_cols;
+    exp_non_local_cols.emplace(1, vec_vec{{0, 1, 2, 3, 4, 5},
+                                          {0, 1, 2, 3, 4, 5},
+                                          {0, 1, 2, 3, 4, 5},
+                                          {0, 1, 2, 3, 4, 5}});
+    exp_non_local_cols.emplace(
+        2, vec_vec{{0, 1, 2, 3, 4, 5}, {}, {0, 1, 2, 3, 4, 5}, {}});
+    exp_non_local_cols.emplace(4, vec_vec{{}, {}, {}, {}});
+
+    auto res_non_local_rows = convert_to_vector(get_row(non_local));
+    std::map<label, vec_vec> exp_non_local_rows;
+    exp_non_local_rows.emplace(1, vec_vec{{2, 5, 8, 6, 7, 8},
+                                          {0, 3, 6, 6, 7, 8},
+                                          {0, 1, 2, 2, 5, 8},
+                                          {0, 1, 2, 0, 3, 6}});
+    exp_non_local_rows.emplace(
+        2, vec_vec{{6, 7, 8, 15, 16, 17}, {}, {0, 1, 2, 9, 10, 11}, {}});
+    exp_non_local_rows.emplace(4, vec_vec{{}, {}, {}, {}});
+
+    ASSERT_EQ(distributed->get_non_local_matrix()->get_size()[1],
+              exp_non_local_size[ranks_per_gpu][rank]);
+
+    ASSERT_EQ(res_non_local_values, exp_non_local_values[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_non_local_rows, exp_non_local_rows[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_non_local_cols, exp_non_local_cols[ranks_per_gpu][rank]);
+}
 
-    auto b = gko::experimental::distributed::Vector<scalar>::create(
-        exec.get_ref_exec(),
-        comm, global_vec_dim, local_vec_dim, 1
-        );
+TEST_P(DistributedMatrixFixture, distributedMatrixCanApplyCorrectly)
+{
+    label ranks_per_gpu = GetParam();
+    auto mesh = ((Environment *)global_env)->mesh;
+    auto hostMatrix = ((Environment *)global_env)->hostMatrix;
+    auto repartitioner =
+        Repartitioner(hostMatrix->get_local_nrows(), ranks_per_gpu, 0, exec);
+    auto distributed = RepartDistMatrix<scalar, label, label>::create(
+        exec, "Coo", repartitioner, hostMatrix);
+
+    gko::dim<2> global_vec_dim{repartitioner.get_orig_partition()->get_size(),
+                               1};
+    gko::dim<2> local_vec_dim{repartitioner.get_repart_dim()[0], 1};
+
+    auto b =gko::share(gko::experimental::distributed::Vector<scalar>::create(
+        exec.get_ref_exec(), comm, global_vec_dim, local_vec_dim, 1));
     b->fill(1);
 
-    auto x = gko::experimental::distributed::Vector<scalar>::create(
-        exec.get_ref_exec(),
-        comm, global_vec_dim, local_vec_dim, 1
-        );
+    auto x = gko::share(gko::experimental::distributed::Vector<scalar>::create(
+        exec.get_ref_exec(), comm, global_vec_dim, local_vec_dim, 1));
     x->fill(0);
 
-    auto distributed = RepartDistMatrix<scalar, label, label>::create(exec, "Coo", repartitioner, hostMatrix);
+    std::map<int, vec_vec_s> exp_x_local;
+
+    exp_x_local.emplace(1, vec_vec_s{{4, 5, 5, 5, 6, 6, 5, 6, 6},
+                                     {5, 5, 4, 6, 6, 5, 6, 6, 5},
+                                     {5, 6, 6, 5, 6, 6, 4, 5, 5},
+                                     {6, 6, 5, 6, 6, 5, 5, 5, 4}});
+    exp_x_local.emplace(
+        2, vec_vec_s{{4, 5, 5, 5, 6, 6, 5, 6, 6, 5, 5, 4, 6, 6, 5, 6, 6, 5},
+                     {},
+                     {5, 6, 6, 5, 6, 6, 4, 5, 5, 6, 6, 5, 6, 6, 5, 5, 5, 4},
+                     {}});
+    exp_x_local.emplace(
+        4, vec_vec_s{{4, 5, 5, 5, 6, 6, 5, 6, 6, 5, 5, 4, 6, 6, 5, 6, 6, 5,
+                      5, 6, 6, 5, 6, 6, 4, 5, 5, 6, 6, 5, 6, 6, 5, 5, 5, 4},
+                     {},
+                     {},
+                     {}});
+
+    // Act
     distributed->apply(b, x);
+    auto x_local = std::vector<scalar>(
+        x->get_local_vector()->get_const_values(),
+        x->get_local_vector()->get_const_values() + local_vec_dim[0]);
 
-    ASSERT_EQ(distributed->get_local_matrix()->get_size()[0], exp_local_size[ranks_per_gpu][rank]);
-    ASSERT_EQ(distributed->get_local_matrix()->get_size()[1], exp_local_size[ranks_per_gpu][rank]);
-    ASSERT_EQ(distributed->get_non_local_matrix()->get_size()[0], exp_local_size[ranks_per_gpu][rank]);
-    ASSERT_EQ(distributed->get_non_local_matrix()->get_size()[1], exp_non_local_size[ranks_per_gpu][rank]);
-
-    auto x_local = std::vector<scalar>(x->get_local_vector()->get_const_values(), x->get_local_vector()->get_const_values() + local_vec_dim[0]);
-    std::vector<scalar> exp_x_local {1};
-
-    ASSERT_EQ(x_local, exp_x_local);
+    ASSERT_EQ(x_local, exp_x_local[ranks_per_gpu][rank]);
 }
 
 int main(int argc, char *argv[])
diff --git a/unitTests/MatrixWrapper/HostMatrix.C b/unitTests/MatrixWrapper/HostMatrix.C
index 29aa3c7e3..2fd43e2d9 100644
--- a/unitTests/MatrixWrapper/HostMatrix.C
+++ b/unitTests/MatrixWrapper/HostMatrix.C
@@ -2,10 +2,11 @@
 //
 // SPDX-License-Identifier: GPL-3.0-or-later
 
+#include <gmock/gmock.h>
+#include "gtest/gtest.h"
 
 #include "OGL/MatrixWrapper/HostMatrix.H"
 
-#include "gtest/gtest.h"
 
 //---------------------------------------------
 // some_header.h
@@ -68,15 +69,15 @@ public:
                            Foam::IOobject::MUST_READ),
             *mesh.get(), ds);
 
-        Foam::fvMatrix<scalar> fvMatrix{*field.get(), ds};
+        fvMatrix = std::make_shared<Foam::fvMatrix<scalar>>(*field.get(), ds);
 
         interfaces = field->boundaryField().scalarInterfaces();
 
         hostMatrix = std::make_shared<HostMatrixWrapper>(
             *exec.get(), runTime_->thisDb(), mesh->lduAddr(), true,
-            fvMatrix.diag().data(), fvMatrix.upper().data(),
-            fvMatrix.lower().data(), fvMatrix.boundaryCoeffs(),
-            fvMatrix.internalCoeffs(), interfaces, dict, "fieldName", 0);
+            fvMatrix->diag().data(), fvMatrix->upper().data(),
+            fvMatrix->lower().data(), fvMatrix->boundaryCoeffs(),
+            fvMatrix->internalCoeffs(), interfaces, dict, "fieldName", 0);
     }
 
     Foam::lduInterfaceFieldPtrsList interfaces;
@@ -88,6 +89,7 @@ public:
     std::shared_ptr<GeometricField<scalar, Foam::fvPatchField, Foam::volMesh>>
         field;
     std::shared_ptr<const ExecutorHandler> exec;
+    std::shared_ptr<Foam::fvMatrix<scalar>> fvMatrix;
     std::shared_ptr<const HostMatrixWrapper> hostMatrix;
 };
 
@@ -101,7 +103,7 @@ TEST(HostMatrix, returnsCorrectSize)
     auto mesh = ((HostMatrixEnvironment *)global_env)->mesh;
     auto hostMatrix = ((HostMatrixEnvironment *)global_env)->hostMatrix;
 
-    std::vector<label> exp_interface_length {3, 3};
+    std::vector<label> exp_interface_length{3, 3};
 
     // the local size is 9
     EXPECT_EQ(mesh->C().size(), 9);
@@ -118,6 +120,39 @@ TEST(HostMatrix, returnsCorrectSize)
     EXPECT_EQ(hostMatrix->get_interface_length(), exp_interface_length);
 }
 
+TEST(HostMatrix, givesAccesToData)
+{
+    auto const max_abs_error = 1e-12;
+    auto hostMatrix = ((HostMatrixEnvironment *)global_env)->hostMatrix;
+
+    label nrows = hostMatrix->get_local_nrows();
+    label upper_nnz = hostMatrix->get_upper_nnz();
+
+    auto exp_diag = std::vector<scalar>(nrows, 0);
+    std::vector<scalar> diag(hostMatrix->get_diag(),
+                             hostMatrix->get_diag() + nrows);
+
+    auto exp_upper = std::vector<scalar>(upper_nnz, 0);
+    std::vector<scalar> upper(hostMatrix->get_upper(),
+                              hostMatrix->get_upper() + upper_nnz);
+
+    std::vector<std::vector<scalar>> exp_interfaces{{0, 0, 0}, {0, 0, 0}};
+
+    ASSERT_THAT(
+        diag, testing::Pointwise(testing::FloatNear(max_abs_error), exp_diag));
+    ASSERT_THAT(upper, testing::Pointwise(testing::FloatNear(max_abs_error),
+                                          exp_upper));
+
+    auto i_length = hostMatrix->get_interface_length();
+    for (int i = 0; i < hostMatrix->get_num_interfaces(); i++) {
+        std::vector<scalar> interface(
+            hostMatrix->get_interface_data(i),
+            hostMatrix->get_interface_data(i) + i_length[i]);
+        ASSERT_THAT(interface,
+                    testing::Pointwise(testing::FloatNear(max_abs_error),
+                                       exp_interfaces[i]));
+    }
+}
 
 TEST(HostMatrix, canCreateCommunicationPattern)
 {
@@ -136,8 +171,9 @@ TEST(HostMatrix, canCreateCommunicationPattern)
 
     std::vector<std::vector<label>> target_sizes_exp{
         {3, 3}, {3, 3}, {3, 3}, {3, 3}};
-    std::vector<label> target_size_res(commPattern->target_sizes.get_data(),
-                                       commPattern->target_sizes.get_data() + 2);
+    std::vector<label> target_size_res(
+        commPattern->target_sizes.get_data(),
+        commPattern->target_sizes.get_data() + 2);
     EXPECT_EQ(target_sizes_exp[comm.rank()], target_size_res);
 }
 
@@ -155,6 +191,18 @@ TEST(HostMatrix, canGenerateLocalSparsityPattern)
                                       3, 4, 6, 1, 3, 4, 5, 7, 2, 4, 5,
                                       8, 3, 6, 7, 4, 6, 7, 8, 5, 7, 8});
 
+    std::vector<label> mapping_expected({
+        12, 0,  1,          // cell 0
+        0,  13, 2,  3,      // cell 1
+        2,  14, 4,          // cell 2
+        1,  15, 5,  6,      // cell 3
+        3,  5,  16, 7,  8,  // cell 4
+        4,  7,  17, 9,      // cell 5
+        6,  18, 10,         // cell 6
+        8,  10, 19, 11,     // cell 7
+        9,  11, 20          // cell 8
+    });
+
     // we have 9x9 matrix with 33 nnz entries
     EXPECT_EQ(localSparsity->num_nnz, 33);
     EXPECT_EQ(localSparsity->dim[0], 9);
@@ -167,13 +215,14 @@ TEST(HostMatrix, canGenerateLocalSparsityPattern)
     EXPECT_EQ(localSparsity->spans[0].end, 33);
 
     auto res_size{localSparsity->col_idxs.get_size()};
-    std::vector<label> rows_res(localSparsity->row_idxs.get_data(),
-                                localSparsity->row_idxs.get_data() + res_size);
-    std::vector<label> cols_res(localSparsity->col_idxs.get_data(),
-                                localSparsity->col_idxs.get_data() + res_size);
+
+    auto rows_res = convert_to_vector(localSparsity->row_idxs);
+    auto cols_res = convert_to_vector(localSparsity->col_idxs);
+    auto mapping_res = convert_to_vector(localSparsity->ldu_mapping);
 
     EXPECT_EQ(rows_expected, rows_res);
     EXPECT_EQ(cols_expected, cols_res);
+    EXPECT_EQ(mapping_expected, mapping_res);
 }
 
 TEST(HostMatrix, canGenerateNonLocalSparsityPattern)
@@ -196,6 +245,8 @@ TEST(HostMatrix, canGenerateNonLocalSparsityPattern)
                                                    {6, 7, 8, 0, 3, 6},
                                                    {6, 7, 8, 2, 5, 8}});
 
+    std::vector<label> mapping_expected({0, 1, 2, 3, 4, 5});
+
     // we dont test the cols expected for now,
     // as they are in compressed format
     EXPECT_EQ(nonLocalSparsity->num_nnz, 6);
@@ -207,14 +258,12 @@ TEST(HostMatrix, canGenerateNonLocalSparsityPattern)
     EXPECT_EQ(nonLocalSparsity->spans[1].end, 6);
 
     auto res_size{nonLocalSparsity->row_idxs.get_size()};
-    std::vector<label> rows_res(
-        nonLocalSparsity->row_idxs.get_data(),
-        nonLocalSparsity->row_idxs.get_data() + res_size);
-    std::vector<label> cols_res(
-        nonLocalSparsity->col_idxs.get_data(),
-        nonLocalSparsity->col_idxs.get_data() + res_size);
+    auto rows_res = convert_to_vector(nonLocalSparsity->row_idxs);
+    auto cols_res = convert_to_vector(nonLocalSparsity->col_idxs);
+    auto mapping_res = convert_to_vector(nonLocalSparsity->ldu_mapping);
     EXPECT_EQ(rows_expected[comm->rank()], rows_res);
     EXPECT_EQ(cols_expected[comm->rank()], cols_res);
+    EXPECT_EQ(mapping_expected, mapping_res);
 }
 
 
diff --git a/unitTests/Repartitioner.C b/unitTests/Repartitioner.C
index 677393c6d..dd0ab18be 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -118,7 +118,7 @@ public:
  *         local ids          |   global ids
  * cells: [ 2 3 | 2 3 ]       |  [ 10 11 | 14 15 ]
  *   2    [ 0 1 | 0 1 ]  3    |  [  8  9 | 12 13 ]
- *        -------------       |  -----------------
+ *        ------+------       |  --------+--------
  *        [ 2 3 | 2 3 ]       |  [  2  3 |  6  7 ]
  *   0    [ 0 1 | 0 1 ]  1    |  [  0  1 |  4  5 ]
  *
@@ -305,18 +305,14 @@ TEST_P(RepartitionerFixture1D, can_repartition_sparsity_pattern_1D_for_n_ranks)
     exp_local_dim_rows.emplace(4, vec{4 * local_size, 0, 0, 0});
 
     std::map<label, vec_vec> exp_local_spans_begin;
-    exp_local_spans_begin.emplace(
-        1, vec_vec{{0}, {0}, {0}, {0}});
-    exp_local_spans_begin.emplace(
-        2, vec_vec{{0, 8, 9}, {}, {0, 8, 9}, {}});
+    exp_local_spans_begin.emplace(1, vec_vec{{0}, {0}, {0}, {0}});
+    exp_local_spans_begin.emplace(2, vec_vec{{0, 8, 9}, {}, {0, 8, 9}, {}});
     exp_local_spans_begin.emplace(
         4, vec_vec{{0, 16, 17, 18, 19, 20, 21}, {}, {}, {}});
 
     std::map<label, vec_vec> exp_local_spans_end;
-    exp_local_spans_end.emplace(
-        1, vec_vec{{5}, {5}, {5}, {5}});
-    exp_local_spans_end.emplace(
-        2, vec_vec{{8, 9, 10}, {}, {8, 9 ,10}, {}});
+    exp_local_spans_end.emplace(1, vec_vec{{5}, {5}, {5}, {5}});
+    exp_local_spans_end.emplace(2, vec_vec{{8, 9, 10}, {}, {8, 9, 10}, {}});
     exp_local_spans_end.emplace(
         4, vec_vec{{16, 17, 18, 19, 20, 21, 22}, {}, {}, {}});
 
@@ -324,6 +320,8 @@ TEST_P(RepartitionerFixture1D, can_repartition_sparsity_pattern_1D_for_n_ranks)
     exp_non_local_rows.emplace(1, non_local_rows);
     exp_non_local_rows.emplace(2, vec_vec{{3}, {}, {0}, {}});
     exp_non_local_rows.emplace(4, vec_vec{{}, {}, {}, {}});
+
+    // expected cols are in compressed format
     std::map<label, vec_vec> exp_non_local_cols;
     exp_non_local_cols.emplace(1, non_local_cols);
     exp_non_local_cols.emplace(2, vec_vec{{0}, {}, {3}, {}});
@@ -342,13 +340,14 @@ TEST_P(RepartitionerFixture1D, can_repartition_sparsity_pattern_1D_for_n_ranks)
     auto res_local_cols = convert_to_vector(repart_local->col_idxs);
     ASSERT_EQ(res_local_rows, exp_local_rows[ranks_per_gpu][rank]);
 
-    std::vector<label> res_local_spans_begin {};
-    std::vector<label> res_local_spans_end {};
-    for (auto [begin, end]: repart_local->spans) {
+    std::vector<label> res_local_spans_begin{};
+    std::vector<label> res_local_spans_end{};
+    for (auto [begin, end] : repart_local->spans) {
         res_local_spans_begin.push_back(begin);
         res_local_spans_end.push_back(end);
     }
-    ASSERT_EQ(res_local_spans_begin, exp_local_spans_begin[ranks_per_gpu][rank]);
+    ASSERT_EQ(res_local_spans_begin,
+              exp_local_spans_begin[ranks_per_gpu][rank]);
     ASSERT_EQ(res_local_spans_end, exp_local_spans_end[ranks_per_gpu][rank]);
 
     // non local properties
@@ -440,6 +439,7 @@ TEST_P(RepartitionerFixture2D, can_repartition_2D_comm_pattern_for_n_ranks)
     // expected rows
     vec_vec_vec rows{
         {{1, 3}, {2, 3}}, {{0, 2}, {2, 3}}, {{0, 1}, {1, 3}}, {{0, 1}, {0, 2}}};
+
     std::map<label, vec_vec> exp_res_rows;
     exp_res_rows.emplace(
         1, vec_vec{{1, 3, 2, 3}, {0, 2, 2, 3}, {0, 1, 1, 3}, {0, 1, 0, 2}});
@@ -454,6 +454,12 @@ TEST_P(RepartitionerFixture2D, can_repartition_2D_comm_pattern_for_n_ranks)
     // row ids
     exp_res_rows.emplace(4, vec_vec{{}, {}, {}, {}});
 
+    std::map<label, vec_vec> exp_gather_idx;
+    exp_gather_idx.emplace(
+        1, vec_vec{{0, 2, 0, 1}, {1, 3, 0, 1}, {2, 3, 0, 2}, {2, 3, 1, 3}});
+    exp_gather_idx.emplace(2, vec_vec{{0, 1, 4, 5}, {}, {2, 3, 6, 7}, {}});
+    exp_gather_idx.emplace(4, vec_vec{{}, {}, {}, {}});
+
     // the original comm_pattern
     auto comm_pattern =
         std::make_shared<CommunicationPattern>(exec, ids[rank], rows[rank]);
@@ -473,6 +479,11 @@ TEST_P(RepartitionerFixture2D, can_repartition_2D_comm_pattern_for_n_ranks)
     auto total_rank_send_idx = repart_comm_pattern->total_rank_send_idx();
     auto res_rows = convert_to_vector(total_rank_send_idx);
     EXPECT_EQ(res_rows, exp_res_rows[ranks_per_gpu][rank]);
+
+    auto total_recv_gather_idx =
+        repart_comm_pattern->compute_recv_gather_idxs(exec);
+    auto res_gather_idx = convert_to_vector(total_recv_gather_idx);
+    EXPECT_EQ(res_gather_idx, exp_gather_idx[ranks_per_gpu][rank]);
 }
 
 

From 315f59119666cce298a09a89bdd01f7256cb663f Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Tue, 3 Sep 2024 15:24:17 +0200
Subject: [PATCH 64/67] add working distributed unit test

---
 unitTests/CMakeLists.txt | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/unitTests/CMakeLists.txt b/unitTests/CMakeLists.txt
index 504c2c219..8805278e4 100644
--- a/unitTests/CMakeLists.txt
+++ b/unitTests/CMakeLists.txt
@@ -19,6 +19,8 @@ macro(ogl_unit_test TEST FILE)
     ${TEST}
     OpenFOAM
     GTest::gtest_main
+    GTest::gtest
+    GTest::gmock
     OGL
     Ginkgo::ginkgo
     ginkgo_hip

From 0d5db932f18a640173ea8f361ad44c29f55a6789 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 4 Sep 2024 10:20:53 +0200
Subject: [PATCH 65/67] add export function

---
 include/OGL/MatrixWrapper/Combination.H |  41 ++-
 include/OGL/MatrixWrapper/Distributed.H | 369 +++---------------------
 include/OGL/common.H                    |   5 -
 src/MatrixWrapper/Distributed.C         | 135 +++++----
 src/common.C                            |  21 --
 unitTests/MatrixWrapper/Distributed.C   |  69 +++--
 6 files changed, 195 insertions(+), 445 deletions(-)

diff --git a/include/OGL/MatrixWrapper/Combination.H b/include/OGL/MatrixWrapper/Combination.H
index a992ddeb2..ded7c765f 100644
--- a/include/OGL/MatrixWrapper/Combination.H
+++ b/include/OGL/MatrixWrapper/Combination.H
@@ -19,7 +19,8 @@ class CombinationMatrix
       public gko::EnableCreateMethod<CombinationMatrix<InnerMatrixType>>,
       public gko::ReadableFromMatrixData<scalar, label>,
       public gko::ConvertibleTo<gko::matrix::Coo<scalar, label>>,
-      public gko::ConvertibleTo<gko::matrix::Csr<scalar, label>> {
+      public gko::ConvertibleTo<gko::matrix::Csr<scalar, label>>,
+      public gko::WritableToMatrixData<scalar, label> {
     friend class gko::EnableCreateMethod<CombinationMatrix<InnerMatrixType>>;
     friend class gko::EnablePolymorphicObject<
         CombinationMatrix<InnerMatrixType>, gko::LinOp>;
@@ -67,6 +68,25 @@ public:
         return *this;
     }
 
+
+    void write(gko::matrix_data<scalar> &res) const override
+    {
+        auto exec = this->get_executor();
+        auto operators = this->get_operators();
+
+        std::vector<gko::matrix_data_entry<scalar, label>> data;
+        for (auto op : operators) {
+            auto tmp_mat_data = gko::matrix_data<scalar>();
+            gko::as<InnerMatrixType>(op)->write(tmp_mat_data);
+            data.insert(data.end(), tmp_mat_data.nonzeros.begin(),
+                        tmp_mat_data.nonzeros.end());
+        }
+        auto mat_data = gko::matrix_data<scalar>();
+        res.nonzeros = data;
+        res.size = this->get_size();
+    }
+
+
     /**
      * Returns a list of coefficients of the combination.
      *
@@ -209,4 +229,23 @@ std::shared_ptr<gko::matrix::Coo<scalar, label>> convert_combination_to_coo(
 
     return out;
 }
+
+template <typename InnerMatrixType>
+void export_mtx(const word fieldName,
+                std::shared_ptr<CombinationMatrix<InnerMatrixType>> out,
+                const word local, const objectRegistry &db,
+                const word matrixFormat)
+{
+    std::string folder{db.time().timePath()};
+    std::filesystem::create_directories(folder);
+    std::string fn{folder + "/" + fieldName + "_A_" + local + ".mtx"};
+    std::cout << "exporting " << fn << std::endl;
+    std::ofstream stream{fn};
+    stream << std::setprecision(15);
+
+    auto coo = gko::share(
+        gko::matrix::Coo<scalar, label>::create(out->get_executor()));
+    out->convert_to(coo);
+}
+
 }  // namespace detail
diff --git a/include/OGL/MatrixWrapper/Distributed.H b/include/OGL/MatrixWrapper/Distributed.H
index 0119958f5..76de22937 100644
--- a/include/OGL/MatrixWrapper/Distributed.H
+++ b/include/OGL/MatrixWrapper/Distributed.H
@@ -14,26 +14,6 @@
 #include "OGL/MatrixWrapper/HostMatrix.H"
 #include "OGL/Repartitioner.H"
 
-namespace detail {
-/* @brief generates a vector of ginkgo linops
-**
-*/
-std::vector<std::shared_ptr<const gko::LinOp>> generate_inner_linops(
-    word matrix_format, bool fuse, std::shared_ptr<const gko::Executor> exec,
-    std::shared_ptr<const SparsityPattern> sparsity);
-
-void update_impl(
-    const ExecutorHandler &exec_handler, word matrix_format,
-    const Repartitioner &repartitioner,
-    std::shared_ptr<const HostMatrixWrapper> host_A,
-    std::shared_ptr<
-        gko::experimental::distributed::Matrix<scalar, label, label>>
-        dist_A,
-    std::shared_ptr<const SparsityPattern> local_sparsity,
-    std::shared_ptr<const SparsityPattern> non_local_sparsity,
-    std::shared_ptr<const CommunicationPattern> src_comm_pattern,
-    std::vector<std::pair<bool, label>> local_interfaces);
-}  // namespace detail
 
 /* The RepartDistMatrix class is a wrapper around Ginkgos distributed Matrix
  * class
@@ -42,34 +22,30 @@ void update_impl(
  * constraint it is required that the inner matrix types of the distributed
  * matrix are of RepartDistMatrix type.
  * */
-template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
+template <typename LocalMatrixType>
 class RepartDistMatrix
     : public gko::experimental::EnableDistributedLinOp<
-          RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>>,
-      public gko::EnableCreateMethod<
-          RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>>,
+          RepartDistMatrix<LocalMatrixType>>,
+      public gko::EnableCreateMethod<RepartDistMatrix<LocalMatrixType>>,
       public gko::experimental::distributed::DistributedBase {
-    friend class gko::EnableCreateMethod<
-        RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>>;
-    friend class gko::EnablePolymorphicObject<
-        RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>,
-        gko::LinOp>;
+    friend class gko::EnableCreateMethod<RepartDistMatrix<LocalMatrixType>>;
+    friend class gko::EnablePolymorphicObject<RepartDistMatrix<LocalMatrixType>,
+                                              gko::LinOp>;
 
     using dist_mtx =
-        gko::experimental::distributed::Matrix<ValueType, LocalIndexType,
-                                               GlobalIndexType>;
+        gko::experimental::distributed::Matrix<scalar, label, label>;
 
     using part_type = gko::experimental::distributed::Partition<label, label>;
 
-    using vec = gko::matrix::Dense<ValueType>;
+    using vec = gko::matrix::Dense<scalar>;
     using device_matrix_data = gko::device_matrix_data<scalar, label>;
     using communicator = gko::experimental::mpi::communicator;
 
 public:
-    using gko::experimental::EnableDistributedLinOp<RepartDistMatrix<
-        ValueType, LocalIndexType, GlobalIndexType>>::convert_to;
     using gko::experimental::EnableDistributedLinOp<
-        RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>>::move_to;
+        RepartDistMatrix<LocalMatrixType>>::convert_to;
+    using gko::experimental::EnableDistributedLinOp<
+        RepartDistMatrix<LocalMatrixType>>::move_to;
 
     std::shared_ptr<const gko::LinOp> get_local_matrix()
     {
@@ -123,249 +99,13 @@ public:
         return *this;
     }
 
-    static std::shared_ptr<RepartDistMatrix> create(
-        const ExecutorHandler &exec_handler, word matrix_format,
-        const Repartitioner &repartitioner,
+    static std::shared_ptr<gko::LinOp> create(
+        const ExecutorHandler &exec_handler, const Repartitioner &repartitioner,
         std::shared_ptr<const HostMatrixWrapper> host_A);
 
-    void update(const ExecutorHandler &exec_handler, word matrix_format,
+    void update(const ExecutorHandler &exec_handler,
                 const Repartitioner &repartitioner,
-                std::shared_ptr<const HostMatrixWrapper> host_A)
-    {
-        detail::update_impl(exec_handler, matrix_format, repartitioner, host_A,
-                            dist_mtx_, local_sparsity_, non_local_sparsity_,
-                            src_comm_pattern_, local_interfaces_);
-    };
-
-    // TODO matrix_format can be data member
-    void update_impl(
-        const ExecutorHandler &exec_handler, word matrix_format,
-        const Repartitioner &repartitioner,
-        std::shared_ptr<const HostMatrixWrapper> host_A,
-        std::shared_ptr<
-            gko::experimental::distributed::Matrix<scalar, label, label>>
-            dist_A,
-        std::shared_ptr<const SparsityPattern> local_sparsity,
-        std::shared_ptr<const SparsityPattern> non_local_sparsity,
-        std::shared_ptr<const CommunicationPattern> src_comm_pattern,
-        std::vector<std::pair<bool, label>> local_interfaces){
-        // auto exec = exec_handler.get_ref_exec();
-        // auto device_exec = exec_handler.get_device_exec();
-        // auto ranks_per_gpu = repartitioner.get_ranks_per_gpu();
-        // bool requires_host_buffer = exec_handler.get_gko_force_host_buffer();
-
-        // label rank{exec_handler.get_rank()};
-        // label owner_rank = repartitioner.get_owner_rank(exec_handler);
-        // bool owner = repartitioner.is_owner(exec_handler);
-        // label nrows = host_A->get_local_nrows();
-        // label local_matrix_nnz = host_A->get_local_matrix_nnz();
-        // // size + padding has to be local_matrix_nnz
-        // auto diag_comm_pattern = compute_send_recv_counts(
-        //     exec_handler, ranks_per_gpu, nrows, local_matrix_nnz,
-        //     local_matrix_nnz - nrows, 0);
-        // label upper_nnz = host_A->get_upper_nnz();
-        // auto upper_comm_pattern = compute_send_recv_counts(
-        //     exec_handler, ranks_per_gpu, upper_nnz, local_matrix_nnz, 0,
-        //     local_matrix_nnz - upper_nnz);
-        // auto lower_comm_pattern =
-        //     compute_send_recv_counts(exec_handler, ranks_per_gpu, upper_nnz,
-        //                              local_matrix_nnz, upper_nnz, nrows);
-
-        // scalar *local_ptr;
-        // scalar *local_ptr_2;
-        // label nnz = 0;
-
-        // // update main values
-        // std::vector<scalar> loc_buffer;
-        // if (owner) {
-        //     using Coo = gko::matrix::Coo<scalar, label>;
-        //     auto local_mtx = dist_A->get_local_matrix();
-
-
-        //     std::shared_ptr<const Coo> local =
-        //         gko::as<Coo>(gko::as<CombinationMatrix<scalar, label, Coo>>(
-        //                          dist_A->get_local_matrix())
-        //                          ->get_combination()
-        //                          ->get_operators()[0]);
-        //     nnz = local->get_num_stored_elements();
-        //     if (requires_host_buffer) {
-        //         loc_buffer.resize(nnz);
-        //         local_ptr = loc_buffer.data();
-        //         local_ptr_2 = const_cast<scalar
-        //         *>(local->get_const_values());
-        //     } else {
-        //         local_ptr = const_cast<scalar *>(local->get_const_values());
-        //     }
-        // }
-        // communicate_values(exec_handler, diag_comm_pattern,
-        // host_A->get_diag(),
-        //                    local_ptr);
-        // communicate_values(exec_handler, upper_comm_pattern,
-        //                    host_A->get_upper(), local_ptr);
-
-        // if (host_A->symmetric()) {
-        //     // TODO FIXME
-        //     // if symmetric we can reuse already copied data
-        //     communicate_values(exec_handler, lower_comm_pattern,
-        //                        host_A->get_lower(), local_ptr);
-        // } else {
-        //     communicate_values(exec_handler, lower_comm_pattern,
-        //                        host_A->get_lower(), local_ptr);
-        // }
-
-        // if (requires_host_buffer) {
-        //     auto host_buffer_view =
-        //         gko::array<scalar>::view(exec, nnz, local_ptr);
-        //     auto target_buffer_view =
-        //         gko::array<scalar>::view(device_exec, nnz, local_ptr_2);
-        //     target_buffer_view = host_buffer_view;
-        // }
-
-        // if (requires_host_buffer) {
-        //     auto host_buffer_view =
-        //         gko::array<scalar>::view(exec, nnz, local_ptr);
-        //     auto target_buffer_view =
-        //         gko::array<scalar>::view(device_exec, nnz, local_ptr_2);
-        //     target_buffer_view = host_buffer_view;
-        // }
-        // // copy interface values
-        // auto comm = *exec_handler.get_communicator().get();
-        // if (owner) {
-        //     using Coo = gko::matrix::Coo<scalar, label>;
-        //     std::shared_ptr<const Coo> mtx;
-        //     label loc_ctr{1};
-        //     label nloc_ctr{0};
-        //     label host_interface_ctr{0};
-        //     label tag = 0;
-        //     label comm_rank, comm_size;
-        //     scalar *recv_buffer_ptr;
-        //     scalar *recv_buffer_ptr_2;
-        //     std::vector<scalar> host_recv_buffer;
-        //     label remain_host_interfaces = host_A->get_interface_size();
-        //     for (auto [is_local, comm_rank] : local_interfaces) {
-        //         label &ctr = (is_local) ? loc_ctr : nloc_ctr;
-        //         if (is_local) {
-        //             mtx = gko::as<Coo>(
-        //                 gko::as<CombinationMatrix<scalar, label, Coo>>(
-        //                     dist_A->get_local_matrix())
-        //                     ->get_combination()
-        //                     ->get_operators()[ctr]);
-        //             comm_size =
-        //             local_sparsity->interface_spans[ctr].length();
-        //         } else {
-        //             mtx = gko::as<Coo>(
-        //                 gko::as<CombinationMatrix<scalar, label, Coo>>(
-        //                     dist_A->get_non_local_matrix())
-        //                     ->get_combination()
-        //                     ->get_operators()[ctr]);
-        //             comm_size =
-        //                 non_local_sparsity->interface_spans[ctr].length();
-        //         }
-
-        //         if (requires_host_buffer) {
-        //             host_recv_buffer.resize(comm_size);
-        //             recv_buffer_ptr = host_recv_buffer.data();
-        //             recv_buffer_ptr_2 =
-        //                 const_cast<scalar *>(mtx->get_const_values());
-        //         } else {
-        //             recv_buffer_ptr =
-        //                 const_cast<scalar *>(mtx->get_const_values());
-        //         }
-
-        //         if (comm_rank != rank) {
-        //             comm.recv(device_exec, recv_buffer_ptr, comm_size,
-        //                       comm_rank, tag);
-        //             if (requires_host_buffer) {
-        //                 auto host_buffer_view = gko::array<scalar>::view(
-        //                     exec, comm_size, recv_buffer_ptr);
-        //                 auto target_buffer_view = gko::array<scalar>::view(
-        //                     device_exec, comm_size, recv_buffer_ptr_2);
-        //                 target_buffer_view = host_buffer_view;
-        //             }
-
-        //         } else {
-        //             // if data is already on this rank
-        //             auto data_view = gko::array<scalar>::const_view(
-        //                 exec, comm_size,
-        //                 host_A->get_interface_data(host_interface_ctr));
-
-        //             // TODO FIXME this needs target executor
-        //             recv_buffer_ptr =
-        //                 const_cast<scalar *>(mtx->get_const_values());
-        //             auto target_view = gko::array<scalar>::view(
-        //                 device_exec, comm_size, recv_buffer_ptr);
-
-        //             target_view = data_view;
-
-        //             host_interface_ctr++;
-        //             remain_host_interfaces--;
-        //         }
-
-        //         ctr++;
-        //         // interface values need to be multiplied by -1
-        //         using vec = gko::matrix::Dense<ValueType>;
-        //         recv_buffer_ptr = const_cast<scalar
-        //         *>(mtx->get_const_values()); auto neg_one =
-        //         gko::initialize<vec>({-1.0}, device_exec); auto
-        //         interface_dense =
-        //             vec::create(device_exec, gko::dim<2>{comm_size, 1},
-        //                         gko::array<scalar>::view(device_exec,
-        //                         comm_size,
-        //                                                  recv_buffer_ptr),
-        //                         1);
-
-        //         interface_dense->scale(neg_one);
-        //     }
-        // } else {
-        //     // the non owner has send all its interfaces to owner
-        //     // thus all values need to be communicated to the owner as well
-        //     label num_interfaces = src_comm_pattern->target_ids.get_size();
-        //     label tag = 0;
-        //     for (int i = 0; i < num_interfaces; i++) {
-        //         label comm_size =
-        //             src_comm_pattern->target_sizes.get_const_data()[i];
-        //         const scalar *send_buffer_ptr =
-        //         host_A->get_interface_data(i); comm.send(device_exec,
-        //         send_buffer_ptr, comm_size, owner_rank,
-        //                   tag);
-        //     }
-        // }
-        // // reorder updated values
-        // if (owner) {
-        //     // NOTE local sparsity size includes the interfaces
-        //     using Coo = gko::matrix::Coo<scalar, label>;
-        //     using dim_type = gko::dim<2>::dimension_type;
-        //     std::shared_ptr<const Coo> local =
-        //         gko::as<Coo>(gko::as<CombinationMatrix<scalar, label, Coo>>(
-        //                          dist_A->get_local_matrix())
-        //                          ->get_combination()
-        //                          ->get_operators()[0]);
-        //     auto local_elements = local->get_num_stored_elements();
-        //     local_ptr = const_cast<scalar *>(local->get_const_values());
-        //     // TODO make sure this doesn't copy
-        //     // create a non owning dense matrix of local_values
-
-        //     auto row_collection =
-        //     gko::share(gko::matrix::Dense<scalar>::create(
-        //         device_exec,
-        //         gko::dim<2>{static_cast<dim_type>(local_elements), 1},
-        //         gko::array<scalar>::view(device_exec, local_elements,
-        //                                  local_ptr),
-        //         1));
-        //     auto mapping_view = gko::array<label>::view(
-        //         exec, local_elements,
-        //         local_sparsity->ldu_mapping.get_data());
-
-
-        //     // TODO this needs to copy ldu_mapping to the device
-        //     auto dense_vec = row_collection->clone();
-        //     // auto dense_vec =
-        //     // gko::share(gko::matrix::Dense<scalar>::create(exec,
-        //     // row_collection->get_size()));
-
-        //     dense_vec->row_gather(&mapping_view, row_collection.get());
-        // }
-    };
+                std::shared_ptr<const HostMatrixWrapper> host_A);
 
     RepartDistMatrix(
         std::shared_ptr<const gko::Executor> exec, communicator comm,
@@ -386,67 +126,7 @@ public:
         this->set_size(global_size);
     }
 
-    void write(const word field_name_, const objectRegistry &db_) const
-    {
-        // auto dist_matrix = gko::as<
-        //     gko::experimental::distributed::Matrix<scalar, label, label>>(
-        //     this->dist_mtx_.get());
-        // // TODO FIXME dont leave hardcoded value here
-        // word matrix_format_{"Coo"};
-
-        // if (matrix_format_ == "Coo") {
-        //     using Coo = gko::matrix::Coo<scalar, label>;
-        //     std::vector<std::shared_ptr<const gko::LinOp>> local_interfaces =
-        //         gko::as<CombinationMatrix<scalar, label, Coo>>(
-        //             dist_matrix->get_local_matrix())
-        //             ->get_combination()
-        //             ->get_operators();
-        //     export_mtx<Coo>(word(field_name_ + "_local"), local_interfaces,
-        //                     db_);
-        //     auto non_local_interfaces =
-        //         gko::as<CombinationMatrix<scalar, label, Coo>>(
-        //             dist_matrix->get_non_local_matrix())
-        //             ->get_combination()
-        //             ->get_operators();
-        //     export_mtx<Coo>(field_name_ + "_non_local", non_local_interfaces,
-        //                     db_);
-        // }
-
-        // if (matrix_format_ == "Csr") {
-        //     using Csr = gko::matrix::Csr<scalar, label>;
-        //     auto local_interfaces =
-        //         gko::as<CombinationMatrix<scalar, label, Csr>>(
-        //             dist_matrix->get_local_matrix().get())
-        //             ->get_combination()
-        //             ->get_operators();
-        //     auto non_local_interfaces =
-        //         gko::as<CombinationMatrix<scalar, label, Csr>>(
-        //             dist_matrix->get_non_local_matrix().get())
-        //             ->get_combination()
-        //             ->get_operators();
-        //     export_mtx<Csr>(word(field_name_ + "_local"), local_interfaces,
-        //                     db_);
-        //     export_mtx<Csr>(field_name_ + "_non_local", non_local_interfaces,
-        //                     db_);
-        // }
-
-        // if (matrix_format_ == "Ell") {
-        //     using Ell = gko::matrix::Ell<scalar, label>;
-        //     auto local_interfaces =
-        //         gko::as<CombinationMatrix<scalar, label, Ell>>(
-        //             dist_matrix->get_local_matrix().get())
-        //             ->get_combination()
-        //             ->get_operators();
-        //     auto non_local_interfaces =
-        //         gko::as<CombinationMatrix<scalar, label, Ell>>(
-        //             dist_matrix->get_non_local_matrix().get())
-        //             ->get_combination()
-        //             ->get_operators();
-        //     export_mtx<Ell>(field_name_ + "_local", local_interfaces, db_);
-        //     export_mtx<Ell>(field_name_ + "_non_local", non_local_interfaces,
-        //                     db_);
-        // }
-    }
+    void write(const word field_name_, const objectRegistry &db_) const;
 
     // Needed for distributed/polymorphic_object.hpp
     RepartDistMatrix(std::shared_ptr<const gko::Executor> exec,
@@ -493,3 +173,20 @@ private:
 
     std::vector<std::pair<bool, label>> local_interfaces_;
 };
+
+
+std::shared_ptr<gko::LinOp> create_distributed(
+    const ExecutorHandler &exec_handler, const Repartitioner &repartitioner,
+    std::shared_ptr<const HostMatrixWrapper> hostMatrix, word matrix_format)
+{
+    if (matrix_format == "Coo") {
+        return RepartDistMatrix<gko::matrix::Coo<scalar, label>>::create(
+            exec_handler, repartitioner, hostMatrix);
+    }
+    if (matrix_format == "Csr") {
+        return RepartDistMatrix<gko::matrix::Csr<scalar, label>>::create(
+            exec_handler, repartitioner, hostMatrix);
+    }
+
+    return {};
+}
diff --git a/include/OGL/common.H b/include/OGL/common.H
index 3a23e0903..ab3cd31d4 100644
--- a/include/OGL/common.H
+++ b/include/OGL/common.H
@@ -103,7 +103,6 @@ std::ostream &operator<<(
 template <typename ValueType>
 std::ostream &operator<<(std::ostream &os, const std::vector<ValueType> &in);
 
-
 void export_system(const word fieldName, const gko::matrix::Csr<scalar> *A,
                    const gko::matrix::Dense<scalar> *x,
                    const gko::matrix::Dense<scalar> *b, const word time);
@@ -111,10 +110,6 @@ void export_system(const word fieldName, const gko::matrix::Csr<scalar> *A,
 void export_vec(const word fieldName, const gko::matrix::Dense<scalar> *x,
                 const objectRegistry &db);
 
-void export_mtx(const word fieldName, std::shared_ptr<const gko::LinOp> A,
-                const word local, const objectRegistry &db,
-                const word matrixFormat);
-
 void set_solve_prev_iters(const word sys_matrix_name, const objectRegistry &db,
                           label prev_solve_iters, const bool is_final);
 
diff --git a/src/MatrixWrapper/Distributed.C b/src/MatrixWrapper/Distributed.C
index de22879c2..c0b18606d 100644
--- a/src/MatrixWrapper/Distributed.C
+++ b/src/MatrixWrapper/Distributed.C
@@ -4,8 +4,9 @@
 
 #include "OGL/MatrixWrapper/Distributed.H"
 
-std::vector<std::shared_ptr<const gko::LinOp>> detail::generate_inner_linops(
-    word matrix_format, bool fuse, std::shared_ptr<const gko::Executor> exec,
+template <typename MatrixType>
+std::vector<std::shared_ptr<const gko::LinOp>> generate_inner_linops(
+    bool fuse, std::shared_ptr<const gko::Executor> exec,
     std::shared_ptr<const SparsityPattern> sparsity)
 {
     std::vector<std::shared_ptr<const gko::LinOp>> lin_ops;
@@ -13,26 +14,42 @@ std::vector<std::shared_ptr<const gko::LinOp>> detail::generate_inner_linops(
         auto [begin, end] = sparsity->spans[i];
         gko::array<scalar> coeffs(exec, end - begin);
         coeffs.fill(0.0);
-
-        auto row_idxs_arr = gko::array<label>(
-            exec->get_master(), sparsity->row_idxs.get_const_data() + begin,
-            sparsity->row_idxs.get_const_data() + end);
-        auto col_idxs_arr = gko::array<label>(
-            exec->get_master(), sparsity->col_idxs.get_const_data() + begin,
-            sparsity->col_idxs.get_const_data() + end);
-
-        auto create_function = [&](word format) {
-            return gko::share(gko::matrix::Coo<scalar, label>::create(
-                exec, sparsity->dim, coeffs, col_idxs_arr, row_idxs_arr));
-        };
-        lin_ops.push_back(create_function(matrix_format));
+        lin_ops.push_back(gko::share(MatrixType::create(
+            exec, sparsity->dim, coeffs,
+            gko::array<label>(exec->get_master(),
+                              sparsity->col_idxs.get_const_data() + begin,
+                              sparsity->col_idxs.get_const_data() + end),
+            gko::array<label>(exec->get_master(),
+                              sparsity->row_idxs.get_const_data() + begin,
+                              sparsity->row_idxs.get_const_data() + end))));
     }
     return lin_ops;
 }
 
-void detail::update_impl(
-    const ExecutorHandler &exec_handler, word matrix_format,
-    const Repartitioner &repartitioner,
+
+template <typename LocalMatrixType>
+void RepartDistMatrix<LocalMatrixType>::write(const word field_name_,
+                                              const objectRegistry &db_) const
+{
+    // std::vector<std::shared_ptr<const gko::LinOp>> local_interfaces =
+    // auto local = dist_mtx_->get_local_matrix()->get_operators();
+    // if (fuse){
+
+    // }
+    // gko::as<CombinationMatrix<LocalMatrixType>>(dist_mtx_->get_local_matrix())
+    //     ->get_operators();
+    // export_mtx<Coo>(word(field_name_ + "_local"), local_interfaces, db_);
+    // auto non_local_interfaces = gko::as<CombinationMatrix<scalar, label,
+    // Coo>>(
+    //                                 dist_matrix->get_non_local_matrix())
+    //                                 ->get_combination()
+    //                                 ->get_operators();
+    // export_mtx<Coo>(field_name_ + "_non_local", non_local_interfaces, db_);
+}
+
+template <typename LocalMatrixType>
+void update_impl(
+    const ExecutorHandler &exec_handler, const Repartitioner &repartitioner,
     std::shared_ptr<const HostMatrixWrapper> host_A,
     std::shared_ptr<
         gko::experimental::distributed::Matrix<scalar, label, label>>
@@ -72,11 +89,9 @@ void detail::update_impl(
     // update main values
     std::vector<scalar> loc_buffer;
     if (owner) {
-        using Coo = gko::matrix::Coo<scalar, label>;
-
-        // TODO make it work with other matrix types
-        std::shared_ptr<const Coo> local = gko::as<Coo>(
-            gko::as<CombinationMatrix<Coo>>(dist_A->get_local_matrix())
+        std::shared_ptr<const LocalMatrixType> local = gko::as<LocalMatrixType>(
+            gko::as<CombinationMatrix<LocalMatrixType>>(
+                dist_A->get_local_matrix())
                 ->get_combination()
                 ->get_operators()[0]);
 
@@ -115,8 +130,7 @@ void detail::update_impl(
     // copy interface values
     auto comm = *exec_handler.get_communicator().get();
     if (owner) {
-        using Coo = gko::matrix::Coo<scalar, label>;
-        std::shared_ptr<const Coo> mtx;
+        std::shared_ptr<const LocalMatrixType> mtx;
         label loc_ctr{1};
         label nloc_ctr{0};
         label host_interface_ctr{0};
@@ -130,16 +144,18 @@ void detail::update_impl(
             label &ctr = (is_local) ? loc_ctr : nloc_ctr;
             if (is_local) {
                 // TODO make it work for other types
-                mtx = gko::as<Coo>(
-                    gko::as<CombinationMatrix<Coo>>(dist_A->get_local_matrix())
+                mtx = gko::as<LocalMatrixType>(
+                    gko::as<CombinationMatrix<LocalMatrixType>>(
+                        dist_A->get_local_matrix())
                         ->get_combination()
                         ->get_operators()[ctr]);
                 comm_size = local_sparsity->spans[ctr].length();
             } else {
-                mtx = gko::as<Coo>(gko::as<CombinationMatrix<Coo>>(
-                                       dist_A->get_non_local_matrix())
-                                       ->get_combination()
-                                       ->get_operators()[ctr]);
+                mtx = gko::as<LocalMatrixType>(
+                    gko::as<CombinationMatrix<LocalMatrixType>>(
+                        dist_A->get_non_local_matrix())
+                        ->get_combination()
+                        ->get_operators()[ctr]);
                 comm_size = non_local_sparsity->spans[ctr].length();
             }
 
@@ -210,11 +226,11 @@ void detail::update_impl(
     // reorder updated values
     if (owner) {
         // NOTE local sparsity size includes the interfaces
-        using Coo = gko::matrix::Coo<scalar, label>;
         using dim_type = gko::dim<2>::dimension_type;
 
-        std::shared_ptr<const Coo> local = gko::as<Coo>(
-            gko::as<CombinationMatrix<Coo>>(dist_A->get_local_matrix())
+        std::shared_ptr<const LocalMatrixType> local = gko::as<LocalMatrixType>(
+            gko::as<CombinationMatrix<LocalMatrixType>>(
+                dist_A->get_local_matrix())
                 ->get_combination()
                 ->get_operators()[0]);
 
@@ -239,15 +255,23 @@ void detail::update_impl(
 
         dense_vec->row_gather(&mapping_view, row_collection.get());
     }
-};
+}
+
+
+template <typename LocalMatrixType>
+void RepartDistMatrix<LocalMatrixType>::update(
+    const ExecutorHandler &exec_handler, const Repartitioner &repartitioner,
+    std::shared_ptr<const HostMatrixWrapper> host_A)
+{
+    update_impl<LocalMatrixType>(exec_handler, repartitioner, host_A, dist_mtx_,
+                                 local_sparsity_, non_local_sparsity_,
+                                 src_comm_pattern_, local_interfaces_);
+}
 
 
-template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
-static std::shared_ptr<
-    RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>>
-RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>::create(
-    const ExecutorHandler &exec_handler, word matrix_format,
-    const Repartitioner &repartitioner,
+template <typename LocalMatrixType>
+static std::shared_ptr<gko::LinOp> RepartDistMatrix<LocalMatrixType>::create(
+    const ExecutorHandler &exec_handler, const Repartitioner &repartitioner,
     std::shared_ptr<const HostMatrixWrapper> host_A)
 {
     label rank = exec_handler.get_rank();
@@ -277,10 +301,10 @@ RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>::create(
     // if fuse the vector contains only a single element
     // thus we can unwrap it.
     auto device_exec = exec_handler.get_device_exec();
-    auto local_linops = detail::generate_inner_linops(
-        matrix_format, false, device_exec, repart_loc_sparsity);
-    auto non_local_linops = detail::generate_inner_linops(
-        matrix_format, false, device_exec, repart_non_loc_sparsity);
+    auto local_linops = generate_inner_linops<LocalMatrixType>(
+        false, device_exec, repart_loc_sparsity);
+    auto non_local_linops = generate_inner_linops<LocalMatrixType>(
+        false, device_exec, repart_non_loc_sparsity);
 
     auto [send_counts, send_offsets, recv_sizes, recv_offsets] =
         repart_comm_pattern->send_recv_pattern();
@@ -301,23 +325,22 @@ RepartDistMatrix<ValueType, LocalIndexType, GlobalIndexType>::create(
     } else {
         dist_A = gko::share(dist_mtx::create(
             exec, comm, global_dim,
-            gko::share(
-                CombinationMatrix<gko::matrix::Coo<scalar, label>>::create(
-                    exec, repart_loc_sparsity->dim, local_linops)),
-            gko::share(
-                CombinationMatrix<gko::matrix::Coo<scalar, label>>::create(
-                    exec, repart_non_loc_sparsity->dim, non_local_linops)),
+            gko::share(CombinationMatrix<LocalMatrixType>::create(
+                exec, repart_loc_sparsity->dim, local_linops)),
+            gko::share(CombinationMatrix<LocalMatrixType>::create(
+                exec, repart_non_loc_sparsity->dim, non_local_linops)),
             recv_sizes, recv_offsets, recv_gather_idxs));
     }
 
-    detail::update_impl(exec_handler, matrix_format, repartitioner, host_A,
-                        dist_A, repart_loc_sparsity, repart_non_loc_sparsity,
-                        src_comm_pattern, local_interfaces);
+    update_impl<LocalMatrixType>(exec_handler, repartitioner, host_A, dist_A,
+                                 repart_loc_sparsity, repart_non_loc_sparsity,
+                                 src_comm_pattern, local_interfaces);
 
-    return std::make_shared<RepartDistMatrix>(
+    return std::make_shared<RepartDistMatrix<LocalMatrixType>>(
         exec, comm, repartitioner.get_repart_dim(), dist_A->get_size(),
         std::move(dist_A), repart_loc_sparsity, repart_non_loc_sparsity,
         src_comm_pattern, local_interfaces);
 }
 
-template class RepartDistMatrix<scalar, label, label>;
+template class RepartDistMatrix<gko::matrix::Coo<scalar, label>>;
+template class RepartDistMatrix<gko::matrix::Csr<scalar, label>>;
diff --git a/src/common.C b/src/common.C
index a242d1a57..ab171113e 100644
--- a/src/common.C
+++ b/src/common.C
@@ -36,27 +36,6 @@ void export_vec(const word fieldName, const gko::matrix::Dense<scalar> *x,
     export_x(fn, x);
 }
 
-void export_mtx(const word fieldName, std::shared_ptr<const gko::LinOp> A,
-                const word local, const objectRegistry &db,
-                const word matrixFormat)
-{
-    std::string folder{db.time().timePath()};
-    std::filesystem::create_directories(folder);
-    std::string fn{folder + "/" + fieldName + "_A_" + local + ".mtx"};
-    std::cout << "exporting " << fn << std::endl;
-    std::ofstream stream{fn};
-    stream << std::setprecision(15);
-    if (matrixFormat == "Coo") {
-        gko::write(stream, (const gko::matrix::Coo<scalar> *)A.get());
-    }
-    if (matrixFormat == "Csr") {
-        gko::write(stream, (const gko::matrix::Csr<scalar> *)A.get());
-    }
-    if (matrixFormat == "Ell") {
-        gko::write(stream, (const gko::matrix::Ell<scalar> *)A.get());
-    }
-}
-
 void export_system(const word fieldName, const gko::matrix::Csr<scalar> *A,
                    const gko::matrix::Dense<scalar> *x,
                    const gko::matrix::Dense<scalar> *b, const word time)
diff --git a/unitTests/MatrixWrapper/Distributed.C b/unitTests/MatrixWrapper/Distributed.C
index 0729f8934..0e059ef2a 100644
--- a/unitTests/MatrixWrapper/Distributed.C
+++ b/unitTests/MatrixWrapper/Distributed.C
@@ -137,6 +137,14 @@ public:
 const testing::Environment *global_env =
     AddGlobalTestEnvironment(new Environment);
 
+class DistributedMatrixFixtureMatrixFormat : public testing::TestWithParam<std::tuple<int, string>> {
+public:
+    ExecutorHandler exec = *((Environment *)global_env)->exec.get();
+    label rank = exec.get_rank();
+    const gko::experimental::mpi::communicator comm =
+        *(exec.get_communicator().get());
+};
+
 class DistributedMatrixFixture : public testing::TestWithParam<int> {
 public:
     ExecutorHandler exec = *((Environment *)global_env)->exec.get();
@@ -145,13 +153,22 @@ public:
         *(exec.get_communicator().get());
 };
 
+
+// INSTANTIATE_TEST_SUITE_P(DistributedMatrixFixtureInstantiationMatrixFormat,
+//                          DistributedMatrixFixtureMatrixFormat,
+//                          testing::Combine(
+//                          testing::Values(1, 2, 4),
+//                          testing::Values("Coo", "Csr")));
+
 INSTANTIATE_TEST_SUITE_P(DistributedMatrixFixtureInstantiation,
-                         DistributedMatrixFixture, testing::Values(1, 2, 4));
+                         DistributedMatrixFixture,
+                         testing::Values(1, 2, 4)
+);
 
 TEST_P(DistributedMatrixFixture, canCreateDistributedMatrix)
 {
     /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
-    label ranks_per_gpu = GetParam();
+    auto ranks_per_gpu = GetParam();
     auto mesh = ((Environment *)global_env)->mesh;
     auto hostMatrix = ((Environment *)global_env)->hostMatrix;
     auto repartitioner =
@@ -166,8 +183,8 @@ TEST_P(DistributedMatrixFixture, canCreateDistributedMatrix)
     exp_local_size.emplace(2, vec{18, 0, 18, 0});
     exp_local_size.emplace(4, vec{36, 0, 0, 0});
 
-    auto distributed = RepartDistMatrix<scalar, label, label>::create(
-        exec, "Coo", repartitioner, hostMatrix);
+    using repart_dist = RepartDistMatrix<gko::matrix::Coo<scalar, label>>;
+    auto distributed = gko::as<repart_dist>(repart_dist::create(exec, repartitioner, hostMatrix));
 
     ASSERT_EQ(distributed->get_local_matrix()->get_size()[0],
               exp_local_size[ranks_per_gpu][rank]);
@@ -177,13 +194,12 @@ TEST_P(DistributedMatrixFixture, canCreateDistributedMatrix)
               exp_local_size[ranks_per_gpu][rank]);
     ASSERT_EQ(distributed->get_local_matrix()->get_size()[0],
               exp_local_size[ranks_per_gpu][rank]);
-
 }
 
 TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectLocalMatrix)
 {
     /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
-    label ranks_per_gpu = GetParam();
+    auto ranks_per_gpu = GetParam();
     auto mesh = ((Environment *)global_env)->mesh;
     auto hostMatrix = ((Environment *)global_env)->hostMatrix;
     auto repartitioner =
@@ -198,8 +214,8 @@ TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectLocalMatrix)
     exp_local_size.emplace(2, vec{18, 0, 18, 0});
     exp_local_size.emplace(4, vec{36, 0, 0, 0});
 
-    auto distributed = RepartDistMatrix<scalar, label, label>::create(
-        exec, "Coo", repartitioner, hostMatrix);
+    using repart_dist = RepartDistMatrix<gko::matrix::Coo<scalar, label>>;
+    auto distributed = gko::as<repart_dist>(repart_dist::create(exec, repartitioner, hostMatrix));
 
     auto local = detail::convert_combination_to_coo(
         exec.get_ref_exec(), distributed->get_local_matrix());
@@ -285,7 +301,7 @@ TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectLocalMatrix)
 TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectNonLocalMatrix)
 {
     /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
-    label ranks_per_gpu = GetParam();
+    auto ranks_per_gpu = GetParam();
     auto mesh = ((Environment *)global_env)->mesh;
     auto hostMatrix = ((Environment *)global_env)->hostMatrix;
     auto repartitioner =
@@ -296,8 +312,8 @@ TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectNonLocalMatrix)
     exp_non_local_size.emplace(2, vec{6, 0, 6, 0});
     exp_non_local_size.emplace(4, vec{0, 0, 0, 0});
 
-    auto distributed = RepartDistMatrix<scalar, label, label>::create(
-        exec, "Coo", repartitioner, hostMatrix);
+    using repart_dist = RepartDistMatrix<gko::matrix::Coo<scalar, label>>;
+    auto distributed = gko::as<repart_dist>(repart_dist::create(exec, repartitioner, hostMatrix));
 
     auto non_local = detail::convert_combination_to_coo(
         exec.get_ref_exec(), distributed->get_non_local_matrix());
@@ -341,19 +357,20 @@ TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectNonLocalMatrix)
 
 TEST_P(DistributedMatrixFixture, distributedMatrixCanApplyCorrectly)
 {
-    label ranks_per_gpu = GetParam();
+    auto ranks_per_gpu = GetParam();
     auto mesh = ((Environment *)global_env)->mesh;
     auto hostMatrix = ((Environment *)global_env)->hostMatrix;
     auto repartitioner =
         Repartitioner(hostMatrix->get_local_nrows(), ranks_per_gpu, 0, exec);
-    auto distributed = RepartDistMatrix<scalar, label, label>::create(
-        exec, "Coo", repartitioner, hostMatrix);
+
+    using repart_dist = RepartDistMatrix<gko::matrix::Coo<scalar, label>>;
+    auto distributed = gko::as<repart_dist>(repart_dist::create(exec, repartitioner, hostMatrix));
 
     gko::dim<2> global_vec_dim{repartitioner.get_orig_partition()->get_size(),
                                1};
     gko::dim<2> local_vec_dim{repartitioner.get_repart_dim()[0], 1};
 
-    auto b =gko::share(gko::experimental::distributed::Vector<scalar>::create(
+    auto b = gko::share(gko::experimental::distributed::Vector<scalar>::create(
         exec.get_ref_exec(), comm, global_vec_dim, local_vec_dim, 1));
     b->fill(1);
 
@@ -367,17 +384,17 @@ TEST_P(DistributedMatrixFixture, distributedMatrixCanApplyCorrectly)
                                      {5, 5, 4, 6, 6, 5, 6, 6, 5},
                                      {5, 6, 6, 5, 6, 6, 4, 5, 5},
                                      {6, 6, 5, 6, 6, 5, 5, 5, 4}});
-    exp_x_local.emplace(
-        2, vec_vec_s{{4, 5, 5, 5, 6, 6, 5, 6, 6, 5, 5, 4, 6, 6, 5, 6, 6, 5},
-                     {},
-                     {5, 6, 6, 5, 6, 6, 4, 5, 5, 6, 6, 5, 6, 6, 5, 5, 5, 4},
-                     {}});
-    exp_x_local.emplace(
-        4, vec_vec_s{{4, 5, 5, 5, 6, 6, 5, 6, 6, 5, 5, 4, 6, 6, 5, 6, 6, 5,
-                      5, 6, 6, 5, 6, 6, 4, 5, 5, 6, 6, 5, 6, 6, 5, 5, 5, 4},
-                     {},
-                     {},
-                     {}});
+
+    std::vector<scalar> exp_x_local_2_1 = {4, 5, 5, 5, 6, 6, 5, 6, 6,
+                                           5, 5, 4, 6, 6, 5, 6, 6, 5};
+    std::vector<scalar> exp_x_local_2_2 = {5, 6, 6, 5, 6, 6, 4, 5, 5,
+                                           6, 6, 5, 6, 6, 5, 5, 5, 4};
+    exp_x_local.emplace(2, vec_vec_s{exp_x_local_2_1, {}, exp_x_local_2_2, {}});
+
+    std::vector<scalar> exp_x_local_4 = {4, 5, 5, 5, 6, 6, 5, 6, 6, 5, 5, 4,
+                                         6, 6, 5, 6, 6, 5, 5, 6, 6, 5, 6, 6,
+                                         4, 5, 5, 6, 6, 5, 6, 6, 5, 5, 5, 4};
+    exp_x_local.emplace(4, vec_vec_s{exp_x_local_4, {}, {}, {}});
 
     // Act
     distributed->apply(b, x);

From fa722560eb2bcc99056e73e799c6ccf46a3dece4 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 4 Sep 2024 10:22:40 +0200
Subject: [PATCH 66/67] format

---
 unitTests/MatrixWrapper/Distributed.C | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/unitTests/MatrixWrapper/Distributed.C b/unitTests/MatrixWrapper/Distributed.C
index 0e059ef2a..8f58d1674 100644
--- a/unitTests/MatrixWrapper/Distributed.C
+++ b/unitTests/MatrixWrapper/Distributed.C
@@ -137,7 +137,8 @@ public:
 const testing::Environment *global_env =
     AddGlobalTestEnvironment(new Environment);
 
-class DistributedMatrixFixtureMatrixFormat : public testing::TestWithParam<std::tuple<int, string>> {
+class DistributedMatrixFixtureMatrixFormat
+    : public testing::TestWithParam<std::tuple<int, string>> {
 public:
     ExecutorHandler exec = *((Environment *)global_env)->exec.get();
     label rank = exec.get_rank();
@@ -161,9 +162,7 @@ public:
 //                          testing::Values("Coo", "Csr")));
 
 INSTANTIATE_TEST_SUITE_P(DistributedMatrixFixtureInstantiation,
-                         DistributedMatrixFixture,
-                         testing::Values(1, 2, 4)
-);
+                         DistributedMatrixFixture, testing::Values(1, 2, 4));
 
 TEST_P(DistributedMatrixFixture, canCreateDistributedMatrix)
 {
@@ -184,7 +183,8 @@ TEST_P(DistributedMatrixFixture, canCreateDistributedMatrix)
     exp_local_size.emplace(4, vec{36, 0, 0, 0});
 
     using repart_dist = RepartDistMatrix<gko::matrix::Coo<scalar, label>>;
-    auto distributed = gko::as<repart_dist>(repart_dist::create(exec, repartitioner, hostMatrix));
+    auto distributed = gko::as<repart_dist>(
+        repart_dist::create(exec, repartitioner, hostMatrix));
 
     ASSERT_EQ(distributed->get_local_matrix()->get_size()[0],
               exp_local_size[ranks_per_gpu][rank]);
@@ -215,7 +215,8 @@ TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectLocalMatrix)
     exp_local_size.emplace(4, vec{36, 0, 0, 0});
 
     using repart_dist = RepartDistMatrix<gko::matrix::Coo<scalar, label>>;
-    auto distributed = gko::as<repart_dist>(repart_dist::create(exec, repartitioner, hostMatrix));
+    auto distributed = gko::as<repart_dist>(
+        repart_dist::create(exec, repartitioner, hostMatrix));
 
     auto local = detail::convert_combination_to_coo(
         exec.get_ref_exec(), distributed->get_local_matrix());
@@ -313,7 +314,8 @@ TEST_P(DistributedMatrixFixture, distributedMatrixHasCorrectNonLocalMatrix)
     exp_non_local_size.emplace(4, vec{0, 0, 0, 0});
 
     using repart_dist = RepartDistMatrix<gko::matrix::Coo<scalar, label>>;
-    auto distributed = gko::as<repart_dist>(repart_dist::create(exec, repartitioner, hostMatrix));
+    auto distributed = gko::as<repart_dist>(
+        repart_dist::create(exec, repartitioner, hostMatrix));
 
     auto non_local = detail::convert_combination_to_coo(
         exec.get_ref_exec(), distributed->get_non_local_matrix());
@@ -364,7 +366,8 @@ TEST_P(DistributedMatrixFixture, distributedMatrixCanApplyCorrectly)
         Repartitioner(hostMatrix->get_local_nrows(), ranks_per_gpu, 0, exec);
 
     using repart_dist = RepartDistMatrix<gko::matrix::Coo<scalar, label>>;
-    auto distributed = gko::as<repart_dist>(repart_dist::create(exec, repartitioner, hostMatrix));
+    auto distributed = gko::as<repart_dist>(
+        repart_dist::create(exec, repartitioner, hostMatrix));
 
     gko::dim<2> global_vec_dim{repartitioner.get_orig_partition()->get_size(),
                                1};

From e29975ee3e68753008fb343aaf84aa8e541ee646 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 4 Sep 2024 10:23:39 +0200
Subject: [PATCH 67/67] Fix typo

---
 unitTests/MatrixWrapper/HostMatrix.C | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/unitTests/MatrixWrapper/HostMatrix.C b/unitTests/MatrixWrapper/HostMatrix.C
index 2fd43e2d9..540d1343a 100644
--- a/unitTests/MatrixWrapper/HostMatrix.C
+++ b/unitTests/MatrixWrapper/HostMatrix.C
@@ -120,7 +120,7 @@ TEST(HostMatrix, returnsCorrectSize)
     EXPECT_EQ(hostMatrix->get_interface_length(), exp_interface_length);
 }
 
-TEST(HostMatrix, givesAccesToData)
+TEST(HostMatrix, givesAccessToData)
 {
     auto const max_abs_error = 1e-12;
     auto hostMatrix = ((HostMatrixEnvironment *)global_env)->hostMatrix;