Add tests and fixes to sp sparse matrix (#819)

Author: Xewar313

include/dr/detail/index.hpp

@@ -72,6 +72,17 @@ template <std::integral T = std::size_t> class index {
   }
 
   constexpr bool operator==(const index &) const noexcept = default;
+  constexpr bool operator<(const index &other) const noexcept
+    requires(std::totally_ordered<T>)
+  {
+    if (first < other.first) {
+      return true;
+    }
+    if (first == other.first && second < other.second) {
+      return true;
+    }
+    return false;
+  }
 
   template <std::size_t Index>
   constexpr T get() const noexcept

include/dr/sp/algorithms/matrix/gemv.hpp

@@ -19,7 +19,7 @@ namespace dr::sp {
 template <dr::distributed_range C, typename T, typename I,
           dr::distributed_range B>
 void flat_gemv(C &&c, dr::sp::sparse_matrix<T, I> &a, B &&b) {
-  assert(c.size() == b.size());
+  assert(a.shape()[0] == c.size());
   assert(a.shape()[1] == b.size());
   assert(a.grid_shape()[0] == c.segments().size());
   assert(a.grid_shape()[1] == 1);
@@ -79,7 +79,7 @@ template <dr::distributed_range C, typename T, typename I,
           dr::distributed_range B>
 void gemv(C &&c, dr::sp::sparse_matrix<T, I> &a, B &&b,
           sp::duplicated_vector<rng::range_value_t<B>> &scratch) {
-  assert(c.size() == b.size());
+  assert(a.shape()[0] == c.size());
   assert(a.shape()[1] == b.size());
   assert(a.grid_shape()[0] == c.segments().size());
   assert(a.grid_shape()[1] == 1);
@@ -104,6 +104,7 @@ void gemv(C &&c, dr::sp::sparse_matrix<T, I> &a, B &&b,
         dr::ranges::local(b_duplicated.local_vector(a_tile.rank()).begin());
     auto c_iter = dr::ranges::local(c.segments()[i].begin());
 
+    assert(c.segments()[i].size() == a_tile.shape()[0]);
     auto &&q = __detail::queue(a_tile.rank());
 
     auto event = __detail::local_gemv(q, a_tile, b_iter, c_iter,

include/dr/sp/containers/matrix_partition.hpp

@@ -86,7 +86,12 @@ class block_cyclic final : public matrix_partition {
 
   dr::index<> tile_shape_;
   dr::index<> grid_shape_;
-}; // namespace dr::sp
+};
+
+inline auto row_cyclic() {
+  return block_cyclic({dr::sp::tile::div, dr::sp::tile::div},
+                      {dr::sp::nprocs(), 1});
+}
 
 inline std::vector<block_cyclic> partition_matmul(std::size_t m, std::size_t n,
                                                   std::size_t k) {

include/dr/sp/containers/sparse_matrix.hpp

@@ -46,7 +46,13 @@ class distributed_range_accessor {
   constexpr distributed_range_accessor(Segments segments, size_type segment_id,
                                        size_type idx) noexcept
       : segments_(rng::views::all(std::forward<Segments>(segments))),
-        segment_id_(segment_id), idx_(idx) {}
+        segment_id_(segment_id), idx_(idx) {
+    while (idx_ >= rng::size((*(segments_.begin() + segment_id_))) &&
+           segment_id_ < rng::size(segments_)) {
+      segment_id_++;
+      idx_ = 0;
+    }
+  }
 
   constexpr distributed_range_accessor &
   operator+=(difference_type offset) noexcept {
@@ -59,7 +65,8 @@ class distributed_range_accessor {
       idx_ += current_offset;
       offset -= current_offset;
 
-      if (idx_ >= rng::size((*(segments_.begin() + segment_id_)))) {
+      while (idx_ >= rng::size((*(segments_.begin() + segment_id_))) &&
+             segment_id_ < rng::size(segments_)) {
         segment_id_++;
         idx_ = 0;
       }
@@ -71,14 +78,16 @@ class distributed_range_accessor {
 
       difference_type new_idx = difference_type(idx_) - current_offset;
 
-      if (new_idx < 0) {
+      while (new_idx < 0 && segment_id_ > 0) {
         segment_id_--;
         new_idx = rng::size(*(segments_.begin() + segment_id_)) - 1;
       }
 
       idx_ = new_idx;
+      offset += current_offset;
     }
 
+    assert(offset == 0);
     return *this;
   }
 
@@ -151,12 +160,12 @@ template <typename T, std::integral I = std::int64_t> class sparse_matrix {
       distributed_sparse_matrix_iterator<std::span<segment_type> &&>;
 
   sparse_matrix(key_type shape)
-      : shape_(shape), partition_(new dr::sp::block_cyclic()) {
+      : shape_(shape), partition_(default_partition_()) {
     init_();
   }
 
   sparse_matrix(key_type shape, double density)
-      : shape_(shape), partition_(new dr::sp::block_cyclic()) {
+      : shape_(shape), partition_(default_partition_()) {
     init_random_(density);
   }
 
@@ -393,6 +402,11 @@ template <typename T, std::integral I = std::int64_t> class sparse_matrix {
     segments_ = generate_segments_();
   }
 
+  std::unique_ptr<dr::sp::matrix_partition> default_partition_() {
+    auto ptr = dr::sp::row_cyclic();
+    return std::make_unique<dr::sp::block_cyclic>(ptr);
+  }
+
 private:
   key_type shape_;
   key_type grid_shape_;

include/dr/sp/util/matrix_io.hpp

@@ -279,11 +279,7 @@ auto mmread(std::string file_path, const matrix_partition &partition,
 
 template <typename T, typename I = std::size_t>
 auto mmread(std::string file_path, bool one_indexed = true) {
-  return mmread<T, I>(
-      file_path,
-      dr::sp::block_cyclic({dr::sp::tile::div, dr::sp::tile::div},
-                           {dr::sp::nprocs(), 1}),
-      one_indexed);
+  return mmread<T, I>(file_path, dr::sp::row_cyclic(), one_indexed);
 }
 
 } // namespace dr::sp

include/dr/sp/views/csr_matrix_view.hpp

@@ -156,11 +156,11 @@ class csr_matrix_view
   std::size_t rank() const { return rank_; }
 
   iterator begin() const {
-    return iterator(values_, rowptr_, colind_, 0, 0, shape()[1], idx_offset_);
+    return iterator(values_, rowptr_, colind_, 0, 0, shape()[0], idx_offset_);
   }
 
   iterator end() const {
-    return iterator(values_, rowptr_, colind_, nnz_, shape()[1], shape()[1],
+    return iterator(values_, rowptr_, colind_, nnz_, shape()[0], shape()[0],
                     idx_offset_);
   }
 

test/gtest/sp/CMakeLists.txt

@@ -14,12 +14,12 @@ add_executable(
   ../common/reduce.cpp ../common/sort.cpp ../common/subrange.cpp
   ../common/take.cpp ../common/transform.cpp ../common/transform_view.cpp
   ../common/zip.cpp ../common/zip_local.cpp containers.cpp algorithms.cpp
-  copy.cpp detail.cpp fill.cpp gemv.cpp transform.cpp)
+  copy.cpp detail.cpp fill.cpp gemv.cpp sparse.cpp transform.cpp)
 
 add_executable(sp-tests-3 sp-tests.cpp containers-3.cpp copy-3.cpp)
 
 # skeleton for rapid builds of individual tests, feel free to change this
-add_executable(sp-quick-test sp-tests.cpp ../common/equal.cpp)
+add_executable(sp-quick-test sp-tests.cpp sparse.cpp)
 target_compile_definitions(sp-quick-test PRIVATE QUICK_TEST)
 
 foreach(test-exec IN ITEMS sp-tests sp-tests-3 sp-quick-test)

test/gtest/sp/gemv.cpp

@@ -33,3 +33,152 @@ TEST(SparseMatrix, Gemv) {
   EXPECT_TRUE(fp_equal(c_ref, c_local))
       << fmt::format("Reference:\n  {}\nActual:\n  {}\n", c_ref, c_local);
 }
+
+TEST(SparseMatrix, EmptyGemv) {
+  std::size_t m = 100;
+  std::size_t k = 100;
+  using T = float;
+  using I = int;
+
+  dr::sp::__detail::coo_matrix<T, I> base;
+  auto csr = dr::sp::__detail::convert_to_csr(base, {m, k}, base.size(),
+                                              std::allocator<T>{});
+  dr::sp::sparse_matrix<T, I> a =
+      dr::sp::create_distributed(csr, dr::sp::row_cyclic());
+
+  dr::sp::distributed_vector<T> b(k, 1.f);
+  dr::sp::distributed_vector<T> c(m, 0.f);
+
+  dr::sp::gemv(c, a, b);
+
+  std::vector<T> c_local(m);
+
+  dr::sp::copy(c.begin(), c.end(), c_local.begin());
+
+  std::vector<T> c_ref(m, 0.f);
+
+  EXPECT_TRUE(fp_equal(c_ref, c_local))
+      << fmt::format("Reference:\n  {}\nActual:\n  {}\n", c_ref, c_local);
+}
+
+TEST(SparseMatrix, ZeroVector) {
+  std::size_t m = 100;
+  std::size_t k = 100;
+  using T = float;
+  using I = int;
+  std::vector<std::pair<std::pair<I, I>, T>> base;
+  for (int i = 0; i < m; i++) {
+    for (int j = 0; j < k; j++) {
+      base.push_back({{i, j}, i + j});
+    }
+  }
+
+  auto csr = dr::sp::__detail::convert_to_csr(base, {m, k}, base.size(),
+                                              std::allocator<T>{});
+  dr::sp::sparse_matrix<T, I> a =
+      dr::sp::create_distributed(csr, dr::sp::row_cyclic());
+
+  dr::sp::distributed_vector<T> b(k, 0.f);
+  dr::sp::distributed_vector<T> c(m, 0.f);
+
+  dr::sp::gemv(c, a, b);
+
+  std::vector<T> c_local(m);
+
+  dr::sp::copy(c.begin(), c.end(), c_local.begin());
+
+  std::vector<T> c_ref(m, 0.f);
+
+  EXPECT_TRUE(fp_equal(c_ref, c_local))
+      << fmt::format("Reference:\n  {}\nActual:\n  {}\n", c_ref, c_local);
+}
+
+TEST(SparseMatrix, NotSquareMatrix) {
+  std::size_t m = 10;
+  std::size_t k = 1000;
+
+  dr::sp::sparse_matrix<float> a(
+      {m, k}, 0.1f,
+      dr::sp::block_cyclic({dr::sp::tile::div, dr::sp::tile::div},
+                           {dr::sp::nprocs(), 1}));
+
+  dr::sp::distributed_vector<float> b(k, 1.f);
+  dr::sp::distributed_vector<float> c(m, 0.f);
+
+  dr::sp::gemv(c, a, b);
+
+  std::vector<float> c_local(m);
+
+  dr::sp::copy(c.begin(), c.end(), c_local.begin());
+
+  std::vector<float> c_ref(m, 0.f);
+
+  for (auto &&[index, v] : a) {
+    auto &&[i, k] = index;
+
+    c_ref[i] += v;
+  }
+
+  EXPECT_TRUE(fp_equal(c_ref, c_local))
+      << fmt::format("Reference:\n  {}\nActual:\n  {}\n", c_ref, c_local);
+}
+
+TEST(SparseMatrix, NotSquareMatrixOtherAxis) {
+  std::size_t m = 1000;
+  std::size_t k = 10;
+
+  dr::sp::sparse_matrix<float> a(
+      {m, k}, 0.1f,
+      dr::sp::block_cyclic({dr::sp::tile::div, dr::sp::tile::div},
+                           {dr::sp::nprocs(), 1}));
+
+  dr::sp::distributed_vector<float> b(k, 1.f);
+  dr::sp::distributed_vector<float> c(m, 0.f);
+
+  dr::sp::gemv(c, a, b);
+
+  std::vector<float> c_local(m);
+
+  dr::sp::copy(c.begin(), c.end(), c_local.begin());
+
+  std::vector<float> c_ref(m, 0.f);
+
+  for (auto &&[index, v] : a) {
+    auto &&[i, k] = index;
+
+    c_ref[i] += v;
+  }
+
+  EXPECT_TRUE(fp_equal(c_ref, c_local))
+      << fmt::format("Reference:\n  {}\nActual:\n  {}\n", c_ref, c_local);
+}
+
+TEST(SparseMatrix, VerySparseMatrix) {
+  std::size_t m = 100;
+  std::size_t k = 100;
+
+  dr::sp::sparse_matrix<float> a(
+      {m, k}, 0.001f,
+      dr::sp::block_cyclic({dr::sp::tile::div, dr::sp::tile::div},
+                           {dr::sp::nprocs(), 1}));
+
+  dr::sp::distributed_vector<float> b(k, 1.f);
+  dr::sp::distributed_vector<float> c(m, 0.f);
+
+  dr::sp::gemv(c, a, b);
+
+  std::vector<float> c_local(m);
+
+  dr::sp::copy(c.begin(), c.end(), c_local.begin());
+
+  std::vector<float> c_ref(m, 0.f);
+
+  for (auto &&[index, v] : a) {
+    auto &&[i, k] = index;
+
+    c_ref[i] += v;
+  }
+
+  EXPECT_TRUE(fp_equal(c_ref, c_local))
+      << fmt::format("Reference:\n  {}\nActual:\n  {}\n", c_ref, c_local);
+}

test/gtest/sp/sparse.cpp

@@ -0,0 +1,69 @@
+// SPDX-FileCopyrightText: Intel Corporation
+//
+// SPDX-License-Identifier: BSD-3-Clause
+#include "xp-tests.hpp"
+
+TEST(SparseMatrix, IterationForward) {
+  std::size_t m = 10;
+  std::size_t k = 10;
+  using T = float;
+  using I = int;
+  std::vector<std::pair<std::pair<I, I>, T>> base;
+  for (int i = 0; i < m; i++) {
+    for (int j = 0; j < k; j++) {
+      base.push_back({{i, j}, i + j});
+    }
+  }
+  std::vector<std::pair<std::pair<I, I>, T>> reference(base.size());
+  std::copy(base.begin(), base.end(), reference.begin());
+  auto csr = dr::sp::__detail::convert_to_csr(base, {m, k}, base.size(),
+                                              std::allocator<T>{});
+  dr::sp::sparse_matrix<T, I> a =
+      dr::sp::create_distributed(csr, dr::sp::row_cyclic());
+  int i = 0;
+  for (auto elem : a) {
+    auto [index, value] = elem;
+    auto [real_index, real_value] = reference[i];
+    auto [m, n] = index;
+    auto [r_m, r_n] = real_index;
+
+    EXPECT_TRUE(m == r_m && n == r_n) << fmt::format(
+        "Reference m, n:\n  {}, {}\nActual:\n  {}, {}\n", r_m, r_n, m, n);
+    EXPECT_TRUE(value == real_value) << fmt::format(
+        "Reference value:\n  {}\nActual:\n  {}\n", real_value, value);
+    i++;
+  }
+}
+
+TEST(SparseMatrix, IterationReverse) {
+  std::size_t m = 10;
+  std::size_t k = 10;
+  using T = float;
+  using I = int;
+  std::vector<std::pair<std::pair<I, I>, T>> base;
+  for (int i = 0; i < m; i++) {
+    for (int j = 0; j < k; j++) {
+      base.push_back({{i, j}, i + j});
+    }
+  }
+  std::vector<std::pair<std::pair<I, I>, T>> reference(base.size());
+  std::copy(base.begin(), base.end(), reference.begin());
+  auto csr = dr::sp::__detail::convert_to_csr(base, {m, k}, base.size(),
+                                              std::allocator<T>{});
+  dr::sp::sparse_matrix<T, I> a =
+      dr::sp::create_distributed(csr, dr::sp::row_cyclic());
+  int i = base.size();
+  auto iterator = a.end();
+  while (iterator > a.begin()) {
+    iterator--;
+    i--;
+    auto [index, value] = *iterator;
+    auto [real_index, real_value] = reference[i];
+    auto [m, n] = index;
+    auto [r_m, r_n] = real_index;
+    EXPECT_TRUE(m == r_m && n == r_n) << fmt::format(
+        "Reference m, n:\n  {}, {}\nActual:\n  {}, {}\n", r_m, r_n, m, n);
+    EXPECT_TRUE(value == real_value) << fmt::format(
+        "Reference value:\n  {}\nActual:\n  {}\n", real_value, value);
+  }
+}