Moved local view to distribution

Author: Mikolaj Komar

examples/mp/local_issue.cpp

@@ -18,7 +18,7 @@ int main(int argc, char **argv) {
   dr::views::csr_matrix_view<double, long> local_data;
   auto root = 0;
   if (root == dr::mp::default_comm().rank()) {
-      local_data = dr::generate_band_csr<double, long>(100, 2, 2);
+      local_data = dr::generate_band_csr<double, long>(10, 0, 1);
   }
   {
     mp::distributed_sparse_matrix<
@@ -30,8 +30,13 @@ int main(int argc, char **argv) {
     auto [n, ma] = ind;
     fmt::print("some res 2 {} {} {}\n", val, n, ma);
 
+    auto mapper = [] (auto elem) { auto [a, b] = elem; auto [c, d] = a; return d;};
+    auto summer = [](auto x, auto y) { return x + y;};
+    auto z2 = dr::transform_view(m_row, mapper);
+    auto red2 = dr::mp::reduce(z2, 0, summer);
+    fmt::print("reduced row {} {}\n", red2, m_row.size());
   }
-
+  
   if (root == dr::mp::default_comm().rank()) {
     dr::__detail::destroy_csr_matrix_view(local_data, std::allocator<double>{});
   }

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

@@ -8,6 +8,42 @@
 #include <fmt/core.h>
 
 namespace dr::mp {
+namespace __detail {
+  template <typename T, typename V>
+  class transform_fn_1 {
+    public:
+    using value_type = V;
+    using index_type = T;
+    transform_fn_1(std::size_t offset, std::size_t row_size, T* row_ptr):
+      offset_(offset), row_size_(row_size), row_ptr_(row_ptr) {
+      assert(offset_ == 0);
+      }
+
+    ~transform_fn_1() {
+      destroyed = true;
+    }
+    template <typename P>
+    auto operator()(P entry) const {
+      assert(offset_ == 0);
+      assert(!destroyed);
+      auto [index, pair] = entry;
+      auto [val, column] = pair;
+      auto row = 0;
+      // auto row = rng::distance(
+      //           row_ptr_,
+      //           std::upper_bound(row_ptr_, row_ptr_ + row_size_, offset_ + index) -
+      //               1);
+      dr::index<index_type> index_obj(row, column);
+      value_type entry_obj(index_obj, val);
+      return entry_obj;
+    }
+    private:
+    bool destroyed = false;
+    std::size_t offset_;
+    std::size_t row_size_;
+    T* row_ptr_;
+  };
+}
 
 template <typename T, typename I, class BackendT = MpiBackend>
 class csr_row_distribution {
@@ -287,8 +323,41 @@ class csr_row_distribution {
           std::max(val_sizes_[i], static_cast<std::size_t>(1)));
     }
     fence();
+    local_view = get_elem_view(vals_size_, cols_data_, vals_data_, rows_data_, rank);
   }
 
+  static auto get_elem_view(std::size_t vals_size, 
+  index_type *local_cols, 
+  elem_type *local_vals,
+   std::shared_ptr<distributed_vector<I>> rows_data, 
+   std::size_t rank) {
+    auto row_size = rows_data->segment_size();
+    std::size_t offset = row_size * rank;
+    auto local_vals_range = rng::subrange(local_vals, local_vals + vals_size);
+    auto local_cols_range = rng::subrange(local_cols, local_cols + vals_size);
+    // auto local_rows = rows_data->segments()[rank].begin().local();
+    auto zipped_results = rng::views::zip(local_vals_range, local_cols_range);
+    auto enumerated_zipped = rng::views::enumerate(zipped_results);
+    auto transformer = [=](auto entry){
+      assert(offset == 0);
+      auto [index, pair] = entry;
+      auto [val, column] = pair;
+      auto row = 0;
+      // auto row = rng::distance(
+      //           local_rows,
+      //           std::upper_bound(local_rows, local_rows + row_size, offset_ + index) -
+      //               1);
+      dr::index<index_type> index_obj(row, column);
+      value_type entry_obj(index_obj, val);
+      return entry_obj;
+    };
+    //__detail::transform_fn_1<index_type, value_type>(offset, row_size, local_rows);
+    return rng::views::transform(enumerated_zipped, transformer);
+  }
+
+  using view_type = decltype(get_elem_view(0, nullptr, nullptr, std::shared_ptr<distributed_vector<I>>(nullptr),0));
+
+  view_type local_view;
   std::size_t segment_size_ = 0;
   std::size_t vals_size_ = 0;
   std::vector<std::size_t> val_offsets_;

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

@@ -5,48 +5,7 @@
 #pragma once
 
 
-int some_id_base =0;
 namespace dr::mp {
-namespace __detail {
-  template <typename T, typename V>
-  class transform_fn_1 {
-    public:
-    using value_type = V;
-    using index_type = T;
-    transform_fn_1(std::size_t offset, std::size_t row_size, T* row_ptr):
-      offset_(offset), row_size_(row_size), row_ptr_(row_ptr) {
-      assert(offset_ == 0);
-      myid = some_id_base++;
-      fmt::print("created {}\n", myid);
-      }
-
-    ~transform_fn_1() {
-      destroyed = true;
-      fmt::print("destroyed {}\n", myid);
-    }
-    template <typename P>
-    auto operator()(P entry) const {
-      fmt::print("called {}\n", myid);
-      assert(offset_ == 0);
-      assert(!destroyed);
-      auto [index, pair] = entry;
-      auto [val, column] = pair;
-      auto row = rng::distance(
-                row_ptr_,
-                std::upper_bound(row_ptr_, row_ptr_ + row_size_, offset_) -
-                    1);
-      dr::index<index_type> index_obj(row, column);
-      value_type entry_obj(index_obj, val);
-      return entry_obj;
-    }
-    private:
-    int myid = 0;
-    bool destroyed = false;
-    std::size_t offset_;
-    std::size_t row_size_;
-    T* row_ptr_;
-  };
-}
 template <typename DSM> class csr_row_segment_iterator;
 
 template <typename DSM> class csr_row_segment_reference {
@@ -96,10 +55,6 @@ template <typename DSM> class csr_row_segment_iterator {
     dsm_ = dsm;
     segment_index_ = segment_index;
     index_ = index;
-    if (dsm_->vals_backend_.getrank() == segment_index_) {
-      elem_view_ = get_elem_view(dsm_, segment_index);
-      base_iter = elem_view_.begin();
-    }
   }
 
   auto operator<=>(const csr_row_segment_iterator &other) const noexcept {
@@ -273,35 +228,13 @@ template <typename DSM> class csr_row_segment_iterator {
   auto local() const {
     const auto my_process_segment_index = dsm_->vals_backend_.getrank();
     assert(my_process_segment_index == segment_index_);
-    auto [a, b] = *base_iter;
-    auto [c, d] = a;
-    fmt::print("aqwsedrftgyhuji {} {} {}\n", b, c, d);
-    return base_iter;
+    return dsm_->local_view.begin();
   }
 
 private:
 
-  static auto get_elem_view(DSM *dsm, std::size_t segment_index) {
-    std::size_t offset = dsm->segment_size_ * segment_index;
-    auto row_size = dsm->segment_size_;
-    auto vals_size = dsm->vals_size_;
-    auto local_vals = dsm->vals_data_;
-    auto local_vals_range = rng::subrange(local_vals, local_vals + vals_size);
-    auto local_cols = dsm->cols_data_;
-    auto local_cols_range = rng::subrange(local_cols, local_cols + vals_size);
-    auto local_rows = dsm->rows_data_->segments()[segment_index].begin().local();
-    auto zipped_results = rng::views::zip(local_vals_range, local_cols_range);
-    auto enumerated_zipped = rng::views::enumerate(zipped_results);
-    auto transformer = __detail::transform_fn_1<index_type, value_type>(offset, row_size, local_rows);
-    return rng::views::transform(enumerated_zipped, transformer);
-  }
-
   // all fields need to be initialized by default ctor so every default
   // constructed iter is equal to any other default constructed iter
-  using view_type = decltype(get_elem_view(std::declval<DSM*>(), 0));
-  using iter_type = rng::iterator_t<view_type>;
-  view_type elem_view_;
-  iter_type base_iter;
   DSM *dsm_ = nullptr;
   std::size_t segment_index_ = 0;
   std::size_t index_ = 0;