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/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/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/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/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/Combination.H b/include/OGL/MatrixWrapper/Combination.H
index 82d67cac9..ded7c765f 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
@@ -18,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>;
@@ -66,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.
      *
@@ -176,7 +197,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 +208,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 +216,36 @@ 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;
+}
+
+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
new file mode 100644
index 000000000..76de22937
--- /dev/null
+++ b/include/OGL/MatrixWrapper/Distributed.H
@@ -0,0 +1,192 @@
+// 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 "OGL/CommunicationPattern.H"
+#include "OGL/MatrixWrapper/Combination.H"
+#include "OGL/MatrixWrapper/HostMatrix.H"
+#include "OGL/Repartitioner.H"
+
+
+/* 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 LocalMatrixType>
+class RepartDistMatrix
+    : public gko::experimental::EnableDistributedLinOp<
+          RepartDistMatrix<LocalMatrixType>>,
+      public gko::EnableCreateMethod<RepartDistMatrix<LocalMatrixType>>,
+      public gko::experimental::distributed::DistributedBase {
+    friend class gko::EnableCreateMethod<RepartDistMatrix<LocalMatrixType>>;
+    friend class gko::EnablePolymorphicObject<RepartDistMatrix<LocalMatrixType>,
+                                              gko::LinOp>;
+
+    using dist_mtx =
+        gko::experimental::distributed::Matrix<scalar, label, label>;
+
+    using part_type = gko::experimental::distributed::Partition<label, label>;
+
+    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<LocalMatrixType>>::convert_to;
+    using gko::experimental::EnableDistributedLinOp<
+        RepartDistMatrix<LocalMatrixType>>::move_to;
+
+    std::shared_ptr<const gko::LinOp> get_local_matrix()
+    {
+        return this->dist_mtx_->get_local_matrix();
+    }
+
+    std::shared_ptr<const gko::LinOp> get_non_local_matrix()
+    {
+        return 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<gko::LinOp> create(
+        const ExecutorHandler &exec_handler, const Repartitioner &repartitioner,
+        std::shared_ptr<const HostMatrixWrapper> host_A);
+
+    void update(const ExecutorHandler &exec_handler,
+                const Repartitioner &repartitioner,
+                std::shared_ptr<const HostMatrixWrapper> host_A);
+
+    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;
+
+    // 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_;
+};
+
+
+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/MatrixWrapper/HostMatrix.H b/include/OGL/MatrixWrapper/HostMatrix.H
index ef28fe389..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
@@ -216,12 +222,40 @@ 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_; }
 
     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_; }
+
+    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
+    {
+        return interface_ptr_[i];
+    }
+
+    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/MatrixWrapper/SparsityPattern.H b/include/OGL/MatrixWrapper/SparsityPattern.H
index a40ace6f3..c041b7450 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.H
+++ b/include/OGL/MatrixWrapper/SparsityPattern.H
@@ -15,37 +15,61 @@ 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.
  *
  * */
 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);
 
+void compress_cols(gko::array<label> &in);
+
+/* @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 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.
  *
  * Additionally it keeps track which parts of the sparsity pattern belongs to
  * which interface.
- *
  * */
 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())
+    SparsityPattern(std::shared_ptr<const gko::Executor> exec, gko::dim<2> dim_,
+                    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(spans_),
+          rank(rank_)
     {
-        // NOT fully implementented
-        FatalErrorInFunction << "Not implemented" << abort(FatalError);
+        // 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, 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)}
+    SparsityPattern(std::shared_ptr<const gko::Executor> exec)
+        : num_nnz(0),
+          row_idxs{exec},
+          col_idxs{exec},
+          ldu_mapping{exec},
+          dim(gko::dim<2>{})
     {}
 
     SparsityPattern(std::shared_ptr<const gko::Executor> exec, const label size,
@@ -54,10 +78,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());
@@ -96,10 +120,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;
 };
 
diff --git a/include/OGL/Repartitioner.H b/include/OGL/Repartitioner.H
index a43692300..d5fc4a4ef 100644
--- a/include/OGL/Repartitioner.H
+++ b/include/OGL/Repartitioner.H
@@ -11,26 +11,56 @@
 #include "CommunicationPattern.H"
 #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);
+
+/* 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);
+
+/* @brief exchange spans and rank between owner and non-owner ranks
+**
+** returns a pair vector spans
+*/
+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,
-                                 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);
 
+/* 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)
-{
-    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
@@ -40,9 +70,9 @@ class Repartitioner {
 private:
     using partition = gko::experimental::distributed::Partition<label, label>;
 
-    const label size_;  //! Size (n matrix rows) before repartitioning
+    const label size_;  //! Local size (matrix rows) before repartitioning
 
-    const label repart_size_;  //! Size after repartitioning
+    const label repart_size_;  //! Local size (matrix rows) after repartitioning
 
     const label ranks_per_gpu_;
 
@@ -103,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
@@ -115,19 +147,30 @@ public:
     std::shared_ptr<const partition> get_orig_partition() const
     {
         return orig_partition_;
-    };
-
+    }
 
-    /* @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 gathered sparsity information copy local interfaces
+     *
+     * 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::pair<SparsityPattern, std::vector<bool>> build_non_local_interfaces(
-        const ExecutorHandler &exec_handler, SparsityPattern &loc,
-        const SparsityPattern &non_loc) const;
-
+    std::vector<std::pair<bool, label>> 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
      *
@@ -165,12 +208,9 @@ 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,
-        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/include/OGL/common.H b/include/OGL/common.H
index 08a31dbe1..ab3cd31d4 100644
--- a/include/OGL/common.H
+++ b/include/OGL/common.H
@@ -90,18 +90,18 @@ 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<<(
     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,
                    const gko::matrix::Dense<scalar> *x,
@@ -110,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);
 
@@ -131,4 +127,21 @@ 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);
+
+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 61598afbb..1ddca69c5 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
@@ -212,10 +211,27 @@ 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();
-    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..c0b18606d
--- /dev/null
+++ b/src/MatrixWrapper/Distributed.C
@@ -0,0 +1,346 @@
+// SPDX-FileCopyrightText: 2024 OGL authors
+//
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#include "OGL/MatrixWrapper/Distributed.H"
+
+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;
+    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);
+        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;
+}
+
+
+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>>
+        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) {
+        std::shared_ptr<const LocalMatrixType> local = gko::as<LocalMatrixType>(
+            gko::as<CombinationMatrix<LocalMatrixType>>(
+                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;
+    }
+
+    // copy interface values
+    auto comm = *exec_handler.get_communicator().get();
+    if (owner) {
+        std::shared_ptr<const LocalMatrixType> 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_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<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<LocalMatrixType>(
+                    gko::as<CombinationMatrix<LocalMatrixType>>(
+                        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 dim_type = gko::dim<2>::dimension_type;
+
+        std::shared_ptr<const LocalMatrixType> local = gko::as<LocalMatrixType>(
+            gko::as<CombinationMatrix<LocalMatrixType>>(
+                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 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 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();
+    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);
+
+    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.
+    auto device_exec = exec_handler.get_device_exec();
+    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();
+
+    // 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};
+
+    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<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));
+    }
+
+    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<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<gko::matrix::Coo<scalar, label>>;
+template class RepartDistMatrix<gko::matrix::Csr<scalar, label>>;
diff --git a/src/MatrixWrapper/HostMatrix.C b/src/MatrixWrapper/HostMatrix.C
index c97d9719d..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,
@@ -259,7 +268,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 +317,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 +460,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 +493,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 +502,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 +537,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 +545,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 +588,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 +598,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/MatrixWrapper/SparsityPattern.C b/src/MatrixWrapper/SparsityPattern.C
index 29209b812..cdd66eaf2 100644
--- a/src/MatrixWrapper/SparsityPattern.C
+++ b/src/MatrixWrapper/SparsityPattern.C
@@ -6,13 +6,25 @@
 
 namespace Foam {
 
+gko::dim<2> compute_dimensions(const std::vector<label> &rows)
+{
+    gko::size_type num_rows = rows.back() + 1;
+    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,
-                                  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 6d8279d06..38904381c 100644
--- a/src/Repartitioner.C
+++ b/src/Repartitioner.C
@@ -4,6 +4,103 @@
 
 #include "OGL/Repartitioner.H"
 
+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,
+                    const std::vector<label> &ranks)
+{
+    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());
+
+    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>
+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)
@@ -25,25 +122,227 @@ Repartitioner::repartition_sparsity(
     std::shared_ptr<const SparsityPattern> src_local_pattern,
     std::shared_ptr<const SparsityPattern> src_non_local_pattern) const
 {
-    FatalErrorInFunction << "Not implemented" << exit(FatalError);
+    LOG_1(verbose_, "start repartition sparsity pattern")
+
+    // 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_;
+
+    // early return if no repartitioning requested
+    if (ranks_per_gpu == 1) {
+        std::vector<std::pair<bool, label>> ret;
+        for ([[maybe_unused]] 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));
+    }
+
+    // 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/column index offset relative to its owner rank
+    auto offset = orig_partition_->get_range_bounds()[rank] -
+                  orig_partition_->get_range_bounds()[owner_rank];
+
+    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);
+
+    if (is_owner(exec_handler)) {
+        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> 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 [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);
+
+    // 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());
+
+    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);
+
+    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>(
+            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),
+            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::move(is_local));
+    }
 }
 
-std::pair<SparsityPattern, std::vector<bool>>
-Repartitioner::build_non_local_interfaces(const ExecutorHandler &exec_handler,
-                                          SparsityPattern &loc,
-                                          const SparsityPattern &non_loc) const
+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>>
+        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
 {
-    FatalErrorInFunction << "Not implemented" << exit(FatalError);
+    auto rank = exec_handler.get_rank();
+    std::vector<std::pair<bool, label>> 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]);
+
+        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);
+            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
+    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;
+        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,
+                             non_local_rows.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(repart_comm_rank);
+
+            // 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;
+        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;
 }
 
 
 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;
@@ -57,7 +356,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
@@ -107,8 +406,9 @@ 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(),
@@ -140,7 +440,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>(
@@ -149,7 +448,6 @@ Repartitioner::repartition_comm_pattern(
                                        target_ids.end()},
             gko::array<comm_size_type>{exec, target_sizes.begin(),
                                        target_sizes.end()},
-
             send_idxs);
     }
 
@@ -160,12 +458,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..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)
@@ -145,6 +124,49 @@ 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::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)
+{
+    return std::vector<label>(in.get_const_data(),
+                              in.get_const_data() + in.get_size());
+}
+
+
+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)
 {
@@ -152,7 +174,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) << ", ";
         }
@@ -163,13 +185,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();
@@ -177,7 +200,33 @@ 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] << ", ";
+        }
+        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;
+}
+
+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] << ", ";
         }
@@ -188,10 +237,11 @@ 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;
 };
+
 }  // namespace Foam
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
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..8f58d1674
--- /dev/null
+++ b/unitTests/MatrixWrapper/Distributed.C
@@ -0,0 +1,423 @@
+// SPDX-FileCopyrightText: 2024 OGL authors
+//
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+
+#include "OGL/MatrixWrapper/Distributed.H"
+#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;
+
+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()
+    {
+        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);
+
+        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(),
+            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;
+    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<Foam::fvMatrix<scalar>> fvMatrix;
+    std::shared_ptr<const HostMatrixWrapper> hostMatrix;
+};
+
+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();
+    label rank = exec.get_rank();
+    const gko::experimental::mpi::communicator comm =
+        *(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));
+
+TEST_P(DistributedMatrixFixture, canCreateDistributedMatrix)
+{
+    /* The test mesh is 6x6 grid decomposed into 4 3x3 subdomains */
+    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);
+
+    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});
+
+    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]);
+    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 */
+    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);
+
+    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});
+
+    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());
+
+    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 */
+    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);
+
+    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});
+
+    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());
+
+    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]);
+}
+
+TEST_P(DistributedMatrixFixture, distributedMatrixCanApplyCorrectly)
+{
+    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);
+
+    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(
+        exec.get_ref_exec(), comm, global_vec_dim, local_vec_dim, 1));
+    b->fill(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);
+
+    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}});
+
+    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);
+    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(x_local, exp_x_local[ranks_per_gpu][rank]);
+}
+
+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;
+}
diff --git a/unitTests/MatrixWrapper/HostMatrix.C b/unitTests/MatrixWrapper/HostMatrix.C
index e23fec25b..540d1343a 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,6 +103,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,28 +112,68 @@ 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);
 }
 
+TEST(HostMatrix, givesAccessToData)
+{
+    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)
 {
     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);
 }
 
@@ -147,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);
@@ -159,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)
@@ -188,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);
@@ -199,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 37a1ddaa7..dd0ab18be 100644
--- a/unitTests/Repartitioner.C
+++ b/unitTests/Repartitioner.C
@@ -3,6 +3,8 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 
 #include "OGL/Repartitioner.H"
+#include "OGL/common.H"
+
 
 #include "gtest/gtest.h"
 
@@ -22,6 +24,10 @@ 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 RepartitionerEnvironment : public testing::Environment {
 public:
     void SetUp()
@@ -34,8 +40,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();
@@ -44,9 +49,9 @@ public:
         // to clean up output
         ::testing::TestEventListeners &listeners =
             ::testing::UnitTest::GetInstance()->listeners();
-        if (comm->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_;
@@ -58,64 +63,313 @@ public:
 const testing::Environment *global_env =
     AddGlobalTestEnvironment(new RepartitionerEnvironment);
 
-class RepartitionerFixture : public testing::TestWithParam<int> {
+class RepartitionerFixture : public testing::Test {
+public:
+    ExecutorHandler exec =
+        *((RepartitionerEnvironment *)global_env)->exec.get();
+    label rank = exec.get_rank();
+    const gko::experimental::mpi::communicator comm =
+        *(exec.get_communicator().get());
+};
+
+/* @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 = 10;
+    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, 5}};
+
+    // 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
+    };
 };
 
-INSTANTIATE_TEST_SUITE_P(RepartitionerFixtureInstantiation,
-                         RepartitionerFixture, testing::Values(1, 2, 4));
+/* @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}};
+
+    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
+    };
+};
 
-TEST_P(RepartitionerFixture, can_create_repartitioner)
+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_P(RepartitionerFixture2D, can_convert_to_global)
+{
+    // Arrange
+    using namespace gko::experimental::distributed;
+    auto ref_exec = exec.get_ref_exec();
+    auto partition = gko::share(
+        build_partition_from_local_size<label, label>(ref_exec, comm, 4));
+
+    vec_vec exp_global_idxs{
+        {4, 6, 8, 9},    // rank 0
+        {1, 3, 12, 13},  // rank 1
+        {2, 3, 12, 14},  // rank 2
+        {6, 7, 9, 11}    // 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_1D_comm_pattern_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 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}};
-
-    // expected communcation ranks
-    using vec_vec = std::vector<std::vector<label>>;
+
+    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}, {}, {}, {}});
+
+    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, 20}, {}});
+    exp_spans_end.emplace(4,
+                          vec_vec{{5, 10, 15, 20, 25, 30, 35, 40}, {}, {}, {}});
+
+    // Act
+    auto [new_spans, new_ranks, origins] =
+        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(RepartitionerFixture1D, can_repartition_sparsity_pattern_1D_for_n_ranks)
+{
+    // Arrange
+    label ranks_per_gpu = GetParam();
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, exec);
+    auto ref_exec = exec.get_ref_exec();
+
+    std::vector<label> ranks{rank};
+    auto local_sparsity = std::make_shared<SparsityPattern>(
+        ref_exec, gko::dim<2>{2, 2}, rows, cols, mapping, spans, ranks);
+
+    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]);
+
+    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>{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});
+    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},
+                                      {},
+                                      {},
+                                      {}});
+
+    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_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}, {}});
+    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}, {}});
+    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]);
+
+    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]);
+    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)
+{
+    // Arrange
+    label ranks_per_gpu = GetParam();
+    auto repartitioner = Repartitioner(local_size, ranks_per_gpu, 0, exec);
+    auto ref_exec = exec.get_ref_exec();
+
+    // 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
@@ -131,7 +385,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}});
@@ -144,60 +397,33 @@ 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]);
-
-    auto comm = exec->get_communicator();
-    auto partition = gko::share(
-        gko::experimental::distributed::build_partition_from_local_size<label,
-                                                                        label>(
-            ref_exec, *comm.get(), 2));
+        std::make_shared<CommunicationPattern>(exec, ids[rank], rows[rank]);
 
     // Act
     auto repart_comm_pattern =
-        repartitioner.repartition_comm_pattern(*exec, comm_pattern, partition);
+        repartitioner.repartition_comm_pattern(exec, comm_pattern);
 
     // 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
-    using vec_vec = std::vector<std::vector<label>>;
+    // 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
     exp_res_ids.emplace(2, vec_vec{{2}, {}, {0}, {}});
@@ -211,9 +437,9 @@ 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;
     exp_res_rows.emplace(
         1, vec_vec{{1, 3, 2, 3}, {0, 2, 2, 3}, {0, 1, 1, 3}, {0, 1, 0, 2}});
@@ -228,40 +454,36 @@ TEST_P(RepartitionerFixture, 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]);
-
-    auto comm = exec->get_communicator();
-    auto partition = gko::share(
-        gko::experimental::distributed::build_partition_from_local_size<label,
-                                                                        label>(
-            ref_exec, *comm.get(), 4));
+        std::make_shared<CommunicationPattern>(exec, ids[rank], rows[rank]);
 
     // Act
     auto repart_comm_pattern =
-        repartitioner.repartition_comm_pattern(*exec, comm_pattern, partition);
+        repartitioner.repartition_comm_pattern(exec, comm_pattern);
 
     // 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]);
+
+    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]);
 }