Merge pull request #9 from neoblizz/dev

Cleaning up the container interface.

Author: neoblizz

include/loops/container/coo.hxx

@@ -3,6 +3,11 @@
 #include <loops/container/vector.hxx>
 #include <loops/memory.hxx>
 
+#include <thrust/execution_policy.h>
+#include <thrust/iterator/zip_iterator.h>
+#include <thrust/sort.h>
+#include <thrust/tuple.h>
+
 namespace loops {
 
 using namespace memory;
@@ -21,9 +26,9 @@ struct coo_t {
   std::size_t cols;
   std::size_t nnzs;
 
-  vector_t<index_t, space> row_indices;  // I
-  vector_t<index_t, space> col_indices;  // J
-  vector_t<value_t, space> values;       // V
+  vector_t<index_t, space> row_indices;  /// I
+  vector_t<index_t, space> col_indices;  /// J
+  vector_t<value_t, space> values;       /// V
 
   coo_t() : rows(0), cols(0), nnzs(0), row_indices(), col_indices(), values() {}
 
@@ -35,21 +40,36 @@ struct coo_t {
         col_indices(nnz),
         values(nnz) {}
 
-  coo_t(const coo_t<index_t, value_t, memory_space_t::device>& rhs)
+  template <auto rhs_space>
+  coo_t(const coo_t<index_t, value_t, rhs_space>& rhs)
       : rows(rhs.rows),
         cols(rhs.cols),
         nnzs(rhs.nnzs),
         row_indices(rhs.row_indices),
         col_indices(rhs.col_indices),
         values(rhs.values) {}
 
-  coo_t(const coo_t<index_t, value_t, memory_space_t::host>& rhs)
-      : rows(rhs.rows),
-        cols(rhs.cols),
-        nnzs(rhs.nnzs),
-        row_indices(rhs.row_indices),
-        col_indices(rhs.col_indices),
-        values(rhs.values) {}
+  void sort_by_row() {
+    auto begin = thrust::make_zip_iterator(
+        thrust::make_tuple(row_indices.begin(), col_indices.begin()));
+    auto end = thrust::make_zip_iterator(
+        thrust::make_tuple(row_indices.end(), col_indices.end()));
+    sort(begin, end);
+  }
+
+  void sort_by_column() {
+    auto begin = thrust::make_zip_iterator(
+        thrust::make_tuple(col_indices.begin(), row_indices.begin()));
+    auto end = thrust::make_zip_iterator(
+        thrust::make_tuple(col_indices.end(), row_indices.end()));
+    sort(begin, end);
+  }
+
+ private:
+  template <typename begin_it_t, typename end_it_t>
+  void sort(begin_it_t& begin, end_it_t& end) {
+    thrust::sort_by_key(begin, end, values.begin());
+  }
 };  // struct coo_t
 
 }  // namespace loops

include/loops/container/csr.hxx

@@ -52,25 +52,13 @@ struct csr_t {
         indices(nnz),
         values(nnz) {}
 
-  /**
-   * @brief Construct a new csr from another csr object on device.
-   *
-   * @param rhs csr_t<index_t, offset_t, value_t, memory_space_t::device>
-   */
-  csr_t(const csr_t<index_t, offset_t, value_t, memory_space_t::device>& rhs)
-      : rows(rhs.rows),
-        cols(rhs.cols),
-        nnzs(rhs.nnzs),
-        offsets(rhs.offsets),
-        indices(rhs.indices),
-        values(rhs.values) {}
-
   /**
    * @brief Construct a new csr from another csr object on host.
    *
-   * @param rhs csr_t<index_t, offset_t, value_t, memory_space_t::host>
+   * @param rhs csr_t<index_t, offset_t, value_t, rhs_space>
    */
-  csr_t(const csr_t<index_t, offset_t, value_t, memory_space_t::host>& rhs)
+  template <auto rhs_space>
+  csr_t(const csr_t<index_t, offset_t, value_t, rhs_space>& rhs)
       : rows(rhs.rows),
         cols(rhs.cols),
         nnzs(rhs.nnzs),
@@ -81,39 +69,19 @@ struct csr_t {
   /**
    * @brief Construct a new csr object from coordinate format (COO).
    *
-   * @param coo coo_t<index_t, value_t, memory_space_t::host>
-   */
-  csr_t(const coo_t<index_t, value_t, memory_space_t::host>& coo)
-      : rows(coo.rows),
-        cols(coo.cols),
-        nnzs(coo.nnzs),
-        indices(coo.col_indices),
-        values(coo.values) {
-    offsets.resize(coo.rows + 1);
-    /// TODO: Do not need this copy for all cases.
-    vector_t<index_t, space> _row_indices = coo.row_indices;
-    detail::indices_to_offsets<space>(_row_indices.data().get(),
-                                      _row_indices.size(), offsets.data().get(),
-                                      offsets.size());
-  }
-
-  /**
-   * @brief Construct a new csr object from coordinate format (COO).
-   *
-   * @param coo coo_t<index_t, value_t, memory_space_t::device>
+   * @param coo coo_t<index_t, value_t, auto>
    */
-  csr_t(const coo_t<index_t, value_t, memory_space_t::device>& coo)
+  template <auto rhs_space>
+  csr_t(const coo_t<index_t, value_t, rhs_space>& coo)
       : rows(coo.rows),
         cols(coo.cols),
         nnzs(coo.nnzs),
+        offsets(coo.rows + 1),
         indices(coo.col_indices),
         values(coo.values) {
-    offsets.resize(coo.rows + 1);
     /// TODO: Do not need this copy for all cases.
     vector_t<index_t, space> _row_indices = coo.row_indices;
-    detail::indices_to_offsets<space>(_row_indices.data().get(),
-                                      _row_indices.size(), offsets.data().get(),
-                                      offsets.size());
+    detail::indices_to_offsets(_row_indices, offsets);
   }
 };  // struct csr_t
 

include/loops/container/detail/convert.hxx

@@ -16,6 +16,9 @@
 #include <thrust/fill.h>
 #include <thrust/scatter.h>
 #include <thrust/execution_policy.h>
+#include <thrust/iterator/counting_iterator.h>
+#include <thrust/iterator/zip_iterator.h>
+#include <thrust/tuple.h>
 #include <loops/memory.hxx>
 
 namespace loops {
@@ -25,81 +28,53 @@ using namespace memory;
 /**
  * @brief Convert offsets to indices.
  *
- * @tparam space The memory space of offsets.
- * @tparam index_t The type of indices.
- * @tparam offset_t The type of offsets.
+ * @tparam index_v_t The type of vector indices.
+ * @tparam offset_v_t The type of vector offsets.
  * @param offsets The offsets.
- * @param size_of_offsets The size of offsets.
  * @param indices The indices.
- * @param size_of_indices The size of indices.
- * @param stream The stream.
  */
-template <memory_space_t space, typename index_t, typename offset_t>
-void offsets_to_indices(const offset_t* offsets,
-                        const std::size_t size_of_offsets,
-                        index_t* indices,
-                        const std::size_t size_of_indices,
-                        cudaStream_t stream = 0) {
-  // Execution policy (determines where to run the kernel).
-  using execution_policy_t =
-      std::conditional_t<(space == memory_space_t::device),
-                         decltype(thrust::cuda::par.on(stream)),
-                         decltype(thrust::host)>;
-
-  execution_policy_t exec;
+template <typename index_v_t, typename offset_v_t>
+void offsets_to_indices(const offset_v_t& offsets, index_v_t& indices) {
+  using offset_t = typename offset_v_t::value_type;
+  using index_t = typename index_v_t::value_type;
 
   // Convert compressed offsets into uncompressed indices.
-  thrust::fill(exec, indices + 0, indices + size_of_indices, offset_t(0));
+  thrust::fill(indices.begin(), indices.end(), offset_t(0));
 
   thrust::scatter_if(
-      exec,                                    // execution policy
       thrust::counting_iterator<offset_t>(0),  // begin iterator
-      thrust::counting_iterator<offset_t>(size_of_offsets - 1),  // end iterator
-      offsets + 0,  // where to scatter
+      thrust::counting_iterator<offset_t>(offsets.size() - 1),  // end iterator
+      offsets.begin(),  // where to scatter
       thrust::make_transform_iterator(
           thrust::make_zip_iterator(
-              thrust::make_tuple(offsets + 0, offsets + 1)),
+              thrust::make_tuple(offsets.begin(), offsets.begin() + 1)),
           [=] __host__ __device__(const thrust::tuple<offset_t, offset_t>& t) {
             thrust::not_equal_to<offset_t> comp;
             return comp(thrust::get<0>(t), thrust::get<1>(t));
           }),
-      indices + 0);
+      indices.begin());
 
-  thrust::inclusive_scan(exec, indices + 0, indices + size_of_indices,
-                         indices + 0, thrust::maximum<offset_t>());
+  thrust::inclusive_scan(indices.begin(), indices.end(), indices.begin(),
+                         thrust::maximum<offset_t>());
 }
 
 /**
  * @brief Converts "indices"-based array to "offsets"-based array.
  *
- * @tparam space The memory space of indices.
- * @tparam index_t The type of indices.
- * @tparam offset_t The type of offsets.
+ * @tparam index_v_t The type of vector indices.
+ * @tparam offset_v_t The type of vector offsets.
  * @param indices The indices.
- * @param size_of_indices The size of indices.
  * @param offsets The offsets.
- * @param size_of_offsets The size of offsets.
- * @param stream CUDA stream.
  */
-template <memory_space_t space, typename index_t, typename offset_t>
-void indices_to_offsets(index_t* indices,
-                        std::size_t size_of_indices,
-                        offset_t* offsets,
-                        std::size_t size_of_offsets,
-                        cudaStream_t stream = 0) {
-  // Execution policy (determines where to run the kernel).
-  using execution_policy_t =
-      std::conditional_t<(space == memory_space_t::device),
-                         decltype(thrust::cuda::par.on(stream)),
-                         decltype(thrust::host)>;
-
-  execution_policy_t exec;
+template <typename index_v_t, typename offset_v_t>
+void indices_to_offsets(const index_v_t& indices, offset_v_t& offsets) {
+  using offset_t = typename offset_v_t::value_type;
+  using index_t = typename index_v_t::value_type;
 
   // Convert uncompressed indices into compressed offsets.
-  thrust::lower_bound(exec, indices, indices + size_of_indices,
-                      thrust::counting_iterator<offset_t>(0),
-                      thrust::counting_iterator<offset_t>(size_of_offsets),
-                      offsets + 0);
+  thrust::lower_bound(
+      indices.begin(), indices.end(), thrust::counting_iterator<offset_t>(0),
+      thrust::counting_iterator<offset_t>(offsets.size()), offsets.begin());
 }
 
 }  // namespace detail

include/loops/util/generate.hxx

@@ -110,15 +110,15 @@ void csr(std::size_t rows,
   uniform_distribution(coo.values.begin(), coo.values.end(), value_t(0.0),
                        value_t(1.0));
 
-  // Construct row/col iterators to traverse.
+  // Sort COO matrix.
+  coo.sort_by_row();
+
+  // Zip Iterators.
   auto begin = thrust::make_zip_iterator(
       thrust::make_tuple(coo.row_indices.begin(), coo.col_indices.begin()));
   auto end = thrust::make_zip_iterator(
       thrust::make_tuple(coo.row_indices.end(), coo.col_indices.end()));
 
-  // Sort the COO matrix.
-  thrust::sort_by_key(begin, end, coo.values.begin());
-
   // Remove duplicates.
   auto new_it = thrust::unique_by_key(begin, end, coo.values.begin());
   auto first_it = thrust::get<1>(new_it);