Fix issue when distributed vector is too small

Author: Mikolaj Komar

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

@@ -301,7 +301,6 @@ class csr_eq_distribution {
       segments_.emplace_back(this, segment_index++,
                              std::min(segment_size_, nnz_ - i), segment_size_);
     }
-    fence();
     auto local_rows = rows_data_;
     auto real_val_size = std::min(vals_data_->segment_size(),
                                   nnz_ - vals_data_->segment_size() * rank);
@@ -317,10 +316,15 @@ class csr_eq_distribution {
       view_helper_const[0] = my_tuple;
     }
 
-    auto local_cols = cols_data_->segments()[rank].begin().local();
-    auto local_vals = vals_data_->segments()[rank].begin().local();
-    local_view = std::make_shared<view_type>(get_elem_view(
-        real_val_size, view_helper_const, local_cols, local_vals, rank));
+    auto local_cols = static_cast<I *>(nullptr);
+    auto local_vals = static_cast<T *>(nullptr);
+    if (cols_data_->segments().size() > rank) {
+      local_cols = cols_data_->segments()[rank].begin().local();
+      local_vals = vals_data_->segments()[rank].begin().local();
+      local_view = std::make_shared<view_type>(get_elem_view(
+          real_val_size, view_helper_const, local_cols, local_vals, rank));
+    }
+    fence();
   }
 
   static auto get_elem_view(std::size_t vals_size, view_tuple *helper_tuple,
@@ -360,7 +364,7 @@ class csr_eq_distribution {
 
   dr::mp::__detail::allocator<view_tuple> tuple_alloc;
   view_tuple *view_helper_const;
-  std::shared_ptr<view_type> local_view;
+  std::shared_ptr<view_type> local_view = nullptr;
 
   std::size_t segment_size_ = 0;
   std::size_t row_size_ = 0;

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

@@ -239,6 +239,9 @@ template <typename DSM> class csr_eq_segment_iterator {
   auto local() const {
     const auto my_process_segment_index = dsm_->rows_backend_.getrank();
     assert(my_process_segment_index == segment_index_);
+    if (dsm_->local_view == nullptr) {
+      return nullptr;
+    }
     return dsm_->local_view->begin();
   }
 

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

@@ -289,8 +289,11 @@ class csr_row_distribution {
           this, segment_index++, val_sizes_[i],
           std::max(val_sizes_[i], static_cast<std::size_t>(1)));
     }
-    fence();
-    auto local_rows = rows_data_->segments()[rank].begin().local();
+
+    auto local_rows = static_cast<I *>(nullptr);
+    if (rows_data_->segments().size() > rank) {
+      local_rows = rows_data_->segments()[rank].begin().local();
+    }
     auto offset = val_offsets_[rank];
     auto real_row_size =
         std::min(rows_data_->segment_size(),
@@ -307,8 +310,11 @@ class csr_row_distribution {
       view_helper_const[0] = my_tuple;
     }
 
-    local_view = std::make_shared<view_type>(get_elem_view(
-        vals_size_, view_helper_const, cols_data_, vals_data_, rank));
+    if (rows_data_->segments().size() > rank) {
+      local_view = std::make_shared<view_type>(get_elem_view(
+          vals_size_, view_helper_const, cols_data_, vals_data_, rank));
+    }
+    fence();
   }
 
   static auto get_elem_view(std::size_t vals_size, view_tuple *helper_tuple,
@@ -365,6 +371,6 @@ class csr_row_distribution {
   dr::index<size_t> shape_;
   std::size_t nnz_;
   std::vector<segment_type> segments_;
-  std::shared_ptr<distributed_vector<I>> rows_data_;
+  std::shared_ptr<distributed_vector<I>> rows_data_ = nullptr;
 };
 } // namespace dr::mp

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

@@ -227,6 +227,9 @@ 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_);
+    if (dsm_->local_view == nullptr) {
+      return nullptr;
+    }
     return dsm_->local_view->begin();
   }