Move rehash to the host

kokkos · Apr 4, 2023 · e5ba25f · e5ba25f
1 parent ddc1c4c
commit e5ba25f
Showing 1 changed file with 39 additions and 8 deletions.
diff --git a/sparse/src/KokkosSparse_coo2crs.hpp b/sparse/src/KokkosSparse_coo2crs.hpp
@@ -61,6 +61,7 @@ class Coo2Crs {
     UmapType *ptr;
   };
   using UmapViewType = Kokkos::View<UmapPtrType *, CrsET>;
+  using BmapViewType = Kokkos::View<bool *, CrsET>;
 
   using CrsRowMapView = Kokkos::View<CrsOT *, CrsET>;
   using CrsRowMapAtomicView =
@@ -73,6 +74,8 @@ class Coo2Crs {
   CrsValuesView m_crs_vals;
   CrsColIdsView m_crs_col_ids;
   UmapViewType m_umaps;
+  BmapViewType m_capacity_bmap;
+  BmapViewType m_tuple_bmap;
   UmapOpType m_insert_op;
   CrsOT m_nrows;
   CrsOT m_ncols;
@@ -86,14 +89,18 @@ class Coo2Crs {
  public:
   KOKKOS_INLINE_FUNCTION
   void operator()(const coo2crsRp1 &, const int &idx) const {
-    auto i = m_row(idx);
-    auto j = m_col(idx);
+    auto i           = m_row(idx);
+    auto j           = m_col(idx);
+    auto is_inserted = m_tuple_bmap(idx);
+
     if (i >= m_nrows || j >= m_ncols) {
       Kokkos::abort("tuple is out of bounds");
-    } else if (i >= 0 && j >= 0) {
-      auto umap = *m_umaps(i).ptr;
-      while (umap.insert(j, m_data(idx), m_insert_op).failed())
-        umap.rehash(umap.capacity() * 2);
+    } else if (!is_inserted && i >= 0 && j >= 0) {
+      auto map_ptr = m_umaps.data()[i].ptr;
+      if (map_ptr[0].insert(j, m_data(idx), m_insert_op).failed())
+        m_capacity_bmap(i) = true;  // hmap at index i reached capacity
+      else
+        m_tuple_bmap(idx) = true;  // check list of inserted tuples
     }
   }
 
@@ -155,6 +162,11 @@ class Coo2Crs {
     m_umaps = UmapViewType(
         Kokkos::view_alloc(Kokkos::WithoutInitializing, "m_umaps"), m_nrows);
 
+    m_capacity_bmap = BmapViewType("m_capacity_bmap", m_nrows);
+    typename BmapViewType::HostMirror m_capacity_bmap_mirror =
+        Kokkos::create_mirror_view(m_capacity_bmap);
+    m_tuple_bmap = BmapViewType("m_tuple_bmap", m_n_tuples);
+
     for (int i = 0; i < m_nrows; i++)
       m_umaps(i).ptr =
           new UmapType(arg_capacity_hint, arg_hasher, arg_equal_to);
@@ -164,8 +176,27 @@ class Coo2Crs {
         m_nrows + 1);
 
     using coo2crsRp1Pt = Kokkos::RangePolicy<coo2crsRp1, CrsET>;
-    Kokkos::parallel_for("coo2crsRp1", coo2crsRp1Pt(0, m_n_tuples), *this);
-    CrsET().fence();
+    bool rehashed      = true;
+    while (rehashed) {
+      Kokkos::parallel_for("coo2crsRp1", coo2crsRp1Pt(0, m_n_tuples), *this);
+
+      CrsET().fence();  // Wait for bitmap writes to land
+      Kokkos::deep_copy(m_capacity_bmap_mirror, m_capacity_bmap);
+      CrsET().fence();
+
+      rehashed = false;
+      // TODO: covert to host-level parallel for.
+      for (int i = 0; i < m_nrows; i++) {
+        if (m_capacity_bmap_mirror(i)) {
+          auto umap_ptr = m_umaps(i).ptr;
+          umap_ptr[0].rehash(umap_ptr[0].capacity() * 2);
+          rehashed                  = true;
+          m_capacity_bmap_mirror(i) = false;
+        }
+      }
+      Kokkos::deep_copy(m_capacity_bmap, m_capacity_bmap_mirror);
+      CrsET().fence();
+    }
 
     typename CrsRowMapView::HostMirror m_crs_row_map_h =
         Kokkos::create_mirror_view(m_crs_row_map);