Improve performance of eq distribution gather

Author: Mikolaj Komar

include/dr/mp/containers/matrix_formats/csr_eq_distribution.hpp

@@ -190,11 +190,22 @@ class csr_eq_distribution {
           auto first_row = row_offsets_[i];
           auto last_row = row_offsets_[i] + row_sizes_[i];
           auto row_size = row_sizes_[i];
-          for (auto j = first_row; j < last_row; j++) {
-            res[j + k * shape_[0]] +=
-                gathered_res_host[vals_width * current_offset + k * row_size +
-                                  j - first_row];
+          if (first_row < last_row) {
+            res[first_row + k * shape_[0]] +=
+                gathered_res_host[vals_width * current_offset + k * row_size];
           }
+          if (first_row < last_row - 1) {
+            auto piece_start = gathered_res_host + vals_width * current_offset +
+                               k * row_size + 1;
+            std::copy(piece_start, piece_start + last_row - first_row - 1,
+                      res.begin() + first_row + k * shape_[0] + 1);
+          }
+          // for (auto j = first_row; j < last_row; j++) {
+          //   res[j + k * shape_[0]] +=
+          //       gathered_res_host[vals_width * current_offset + k * row_size
+          //       +
+          //                         j - first_row];
+          // }
           current_offset += row_sizes_[i];
         }
       }