From 0ec5ae4861f040105677505ce45493f69d3f1748 Mon Sep 17 00:00:00 2001
From: Gregor Olenik <gregor.olenik@web.de>
Date: Wed, 11 Dec 2024 11:08:28 +0100
Subject: [PATCH] fix compress columns order based on orig ranks instead of
 interface id

---
 include/OGL/MatrixWrapper/SparsityPattern.hpp | 17 ++++++++++++-----
 src/Repartitioner.cpp                         | 10 ++++++----
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/include/OGL/MatrixWrapper/SparsityPattern.hpp b/include/OGL/MatrixWrapper/SparsityPattern.hpp
index 757ac4c8..8a20c224 100644
--- a/include/OGL/MatrixWrapper/SparsityPattern.hpp
+++ b/include/OGL/MatrixWrapper/SparsityPattern.hpp
@@ -36,6 +36,12 @@ std::vector<label> sort_permutation(const std::vector<T> &vec, Compare compare)
     return p;
 }
 
+/* @brief compute compress column indices
+**
+**@param in - vector of vectors of column indices per interface/sub-matrix
+**@param ids - vector of corresponding rank idxs on which the sub-matrix was
+*originally created
+*/
 std::pair<std::vector<std::vector<label>>, std::vector<label>> compress_cols(
     std::vector<std::vector<label>> in, std::vector<label> ids);
 }  // namespace detail
@@ -150,7 +156,7 @@ class SparsityPattern {
 
     std::vector<label> compute_to_global_map(bool fuse) const
     {
-        return std::get<1>(detail::compress_cols(cols_, id_));
+        return std::get<1>(detail::compress_cols(cols_, orig_rank_));
     }
 
     // make this a free function
@@ -163,7 +169,7 @@ class SparsityPattern {
         if (!repartioned) {
             return {rows_,
                     (compress_cols)
-                        ? std::get<0>(detail::compress_cols(cols_, id_))
+                        ? std::get<0>(detail::compress_cols(cols_, orig_rank_))
                         : cols_,
                     map_, id_};
         }
@@ -258,9 +264,10 @@ class SparsityPattern {
         std::vector<label> map;
         map.reserve(reserve_size);
 
-        auto or_cols = (compress_cols)
-                           ? std::get<0>(detail::compress_cols(cols_, id_))
-                           : cols_;
+        auto or_cols =
+            (compress_cols)
+                ? std::get<0>(detail::compress_cols(cols_, orig_rank_))
+                : cols_;
 
         label map_offset = 0;
 
diff --git a/src/Repartitioner.cpp b/src/Repartitioner.cpp
index c22adf6b..dfb0bd83 100644
--- a/src/Repartitioner.cpp
+++ b/src/Repartitioner.cpp
@@ -66,8 +66,11 @@ Repartitioner::repartition_sparsity(
     };
 
 
+    /* Helper function, create and return gathered sparsity pattern based on
+     * in_sparsity
+     *  */
     auto create_sparsity = [&](const auto &in_sparsity, auto &rows, auto &cols,
-                               auto &map, bool fuse) {
+                               auto &map) {
         auto lengths_tmp = in_sparsity->get_lengths();
         auto size_comm_pattern = compute_gather_to_owner_counts(
             exec_handler, ranks_per_gpu, lengths_tmp.size());
@@ -105,7 +108,7 @@ Repartitioner::repartition_sparsity(
     auto loc_row =
         gather_vector(loc_nnz, offset, src_local_pattern->get_rows());
     auto ret_local_sparsity =
-        create_sparsity(src_local_pattern, loc_row, loc_col, loc_map, true);
+        create_sparsity(src_local_pattern, loc_row, loc_col, loc_map);
 
     size_t non_loc_nnz = src_non_local_pattern->get_nnz();
     auto non_loc_map =
@@ -116,7 +119,7 @@ Repartitioner::repartition_sparsity(
     auto non_loc_col =
         gather_vector(non_loc_nnz, 0, src_non_local_pattern->get_cols());
     auto ret_non_local_sparsity = create_sparsity(
-        src_non_local_pattern, non_loc_row, non_loc_col, non_loc_map, false);
+        src_non_local_pattern, non_loc_row, non_loc_col, non_loc_map);
 
     for (auto &comm_rank : ret_non_local_sparsity->get_comm_rank()) {
         comm_rank = compute_owner_rank(comm_rank, ranks_per_gpu);
@@ -131,7 +134,6 @@ Repartitioner::repartition_sparsity(
         return out;
     };
 
-    // non owning ranks seems to have wrong non-local size
     if (ret_local_sparsity->get_nnz() != 0) {
         ret_local_sparsity->move_interface(ret_non_local_sparsity, rank,
                                            convert_to_local);