Review update.

Co-authored-by: Tobias Ribizel <[email protected]>
Co-authored-by: Marcel Koch <[email protected]>

core/base/index_set.cpp

@@ -105,23 +105,6 @@ IndexType IndexSet<IndexType>::get_local_index(const IndexType index) const
 }
 
 
-template <typename IndexType>
-IndexType IndexSet<IndexType>::get_subset_id(const IndexType index) const
-{
-    auto ss_end_host = make_temporary_clone<const Array<IndexType>>(
-        this->get_executor()->get_master(), &this->subsets_end_);
-    auto ss_begin_host = make_temporary_clone<const Array<IndexType>>(
-        this->get_executor()->get_master(), &this->subsets_begin_);
-    for (size_type id = 0; id < this->get_num_subsets(); ++id) {
-        if (index < ss_end_host->get_const_data()[id] &&
-            index >= ss_begin_host->get_const_data()[id]) {
-            return id;
-        }
-    }
-    return -1;
-}
-
-
 template <typename IndexType>
 Array<IndexType> IndexSet<IndexType>::to_global_indices() const
 {

core/matrix/csr.cpp

@@ -625,6 +625,9 @@ Csr<ValueType, IndexType>::create_submatrix(
 {
     using Mat = Csr<ValueType, IndexType>;
     auto exec = this->get_executor();
+    if (!row_index_set.get_num_elems() || !col_index_set.get_num_elems()) {
+        return Mat::create(exec);
+    }
     if (row_index_set.is_contiguous() && col_index_set.is_contiguous()) {
         auto row_st = row_index_set.get_executor()->copy_val_to_host(
             row_index_set.get_subsets_begin());
@@ -643,7 +646,7 @@ Csr<ValueType, IndexType>::create_submatrix(
         auto sub_mat_size = gko::dim<2>(submat_num_rows, submat_num_cols);
         Array<IndexType> row_ptrs(exec, submat_num_rows + 1);
         exec->run(csr::make_calculate_nonzeros_per_row_in_index_set(
-            this, row_index_set, col_index_set, &row_ptrs));
+            this, row_index_set, col_index_set, row_ptrs.get_data()));
         exec->run(
             csr::make_prefix_sum(row_ptrs.get_data(), submat_num_rows + 1));
         auto num_nnz =

core/matrix/csr_kernels.hpp

@@ -172,7 +172,7 @@ namespace kernels {
         std::shared_ptr<const DefaultExecutor> exec,                    \
         const matrix::Csr<ValueType, IndexType>* source,                \
         const IndexSet<IndexType>& row_index_set,                       \
-        const IndexSet<IndexType>& col_index_set, Array<IndexType>* row_nnz)
+        const IndexSet<IndexType>& col_index_set, IndexType* row_nnz)
 
 #define GKO_DECLARE_CSR_COMPUTE_SUB_MATRIX_KERNEL(ValueType, IndexType)     \
     void compute_submatrix(std::shared_ptr<const DefaultExecutor> exec,     \

cuda/matrix/csr_kernels.cu

@@ -1164,7 +1164,7 @@ void calculate_nonzeros_per_row_in_index_set(
     const matrix::Csr<ValueType, IndexType>* source,
     const IndexSet<IndexType>& row_index_set,
     const IndexSet<IndexType>& col_index_set,
-    Array<IndexType>* row_nnz) GKO_NOT_IMPLEMENTED;
+    IndexType* row_nnz) GKO_NOT_IMPLEMENTED;
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_CSR_CALC_NNZ_PER_ROW_IN_INDEX_SET_KERNEL);

dpcpp/matrix/csr_kernels.dp.cpp

@@ -1391,7 +1391,7 @@ void calculate_nonzeros_per_row_in_index_set(
     const matrix::Csr<ValueType, IndexType>* source,
     const IndexSet<IndexType>& row_index_set,
     const IndexSet<IndexType>& col_index_set,
-    Array<IndexType>* row_nnz) GKO_NOT_IMPLEMENTED;
+    IndexType* row_nnz) GKO_NOT_IMPLEMENTED;
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_CSR_CALC_NNZ_PER_ROW_IN_INDEX_SET_KERNEL);

hip/matrix/csr_kernels.hip.cpp

@@ -951,7 +951,7 @@ void calculate_nonzeros_per_row_in_index_set(
     const matrix::Csr<ValueType, IndexType>* source,
     const IndexSet<IndexType>& row_index_set,
     const IndexSet<IndexType>& col_index_set,
-    Array<IndexType>* row_nnz) GKO_NOT_IMPLEMENTED;
+    IndexType* row_nnz) GKO_NOT_IMPLEMENTED;
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_CSR_CALC_NNZ_PER_ROW_IN_INDEX_SET_KERNEL);

include/ginkgo/core/base/index_set.hpp

@@ -82,10 +82,17 @@ namespace gko {
  * @ingroup IndexSet
  */
 template <typename IndexType = int32>
-class IndexSet : public EnablePolymorphicObject<IndexSet<IndexType>> {
+class IndexSet : public EnablePolymorphicObject<IndexSet<IndexType>>,
+                 public EnablePolymorphicAssignment<IndexSet<IndexType>>,
+                 public EnableCreateMethod<IndexSet<IndexType>> {
     friend class EnablePolymorphicObject<IndexSet>;
+    friend class EnableCreateMethod<IndexSet>;
 
 public:
+    using EnableCreateMethod<IndexSet>::create;
+    using EnablePolymorphicAssignment<IndexSet>::convert_to;
+    using EnablePolymorphicAssignment<IndexSet>::move_to;
+
     /**
      * The type of elements stored in the index set.
      */
@@ -97,7 +104,9 @@ class IndexSet : public EnablePolymorphicObject<IndexSet<IndexType>> {
      * @param exec  the Executor where the IndexSet data is allocated
      */
     IndexSet(std::shared_ptr<const Executor> exec)
-        : EnablePolymorphicObject<IndexSet>(std::move(exec))
+        : EnablePolymorphicObject<IndexSet>(std::move(exec)),
+          index_space_size_{0},
+          num_stored_indices_{0}
     {}
 
     /**
@@ -109,14 +118,18 @@ class IndexSet : public EnablePolymorphicObject<IndexSet<IndexType>> {
      * @param is_sorted  a parameter that specifies if the indices array is
      *                   sorted or not. `true` if sorted.
      */
-    IndexSet(std::shared_ptr<const gko::Executor> executor,
-             std::initializer_list<IndexType> init_list,
-             const bool is_sorted = false)
+    explicit IndexSet(std::shared_ptr<const gko::Executor> executor,
+                      std::initializer_list<IndexType> init_list,
+                      const bool is_sorted = false)
         : EnablePolymorphicObject<IndexSet>(std::move(executor)),
-          index_space_size_(
-              *(std::max_element(std::begin(init_list), std::end(init_list))) +
-              1)
+          index_space_size_(init_list.size() > 0
+                                ? *(std::max_element(std::begin(init_list),
+                                                     std::end(init_list))) +
+                                      1
+                                : 0),
+          num_stored_indices_{static_cast<IndexType>(init_list.size())}
     {
+        GKO_ASSERT(index_space_size_ > 0);
         this->populate_subsets(
             Array<IndexType>(this->get_executor(), init_list), is_sorted);
     }
@@ -131,9 +144,10 @@ class IndexSet : public EnablePolymorphicObject<IndexSet<IndexType>> {
      * @param is_sorted  a parameter that specifies if the indices array is
      *                   sorted or not. `true` if sorted.
      */
-    IndexSet(std::shared_ptr<const gko::Executor> executor,
-             const index_type size, const gko::Array<index_type>& indices,
-             const bool is_sorted = false)
+    explicit IndexSet(std::shared_ptr<const gko::Executor> executor,
+                      const index_type size,
+                      const gko::Array<index_type>& indices,
+                      const bool is_sorted = false)
         : EnablePolymorphicObject<IndexSet>(std::move(executor)),
           index_space_size_(size)
     {
@@ -217,28 +231,6 @@ class IndexSet : public EnablePolymorphicObject<IndexSet<IndexType>> {
      */
     index_type get_local_index(index_type global_index) const;
 
-    /**
-     * Return which set the global index belongs to.
-     *
-     * Consider the set idx_set = (0, 1, 2, 4, 6, 7, 8, 9). This function
-     * returns the subset id in the index set of the input global index. For
-     * example, `idx_set.get_subset_id(0) == 0` `idx_set.get_subset_id(4)
-     * == 1` and `idx_set.get_subset_id(6) == 2`.
-     *
-     * @note This function returns a scalar value and needs a scalar value.
-     *       For repeated queries, it is more efficient to use the Array
-     *       functions that take and return arrays which allow for more
-     *       throughput.
-     *
-     * @param global_index  the global index.
-     *
-     * @return  the local index of the element in the index set.
-     *
-     * @warning This single entry query can have significant kernel launch
-     *          overheads and should be avoided if possible.
-     */
-    index_type get_subset_id(index_type global_index) const;
-
     /**
      * This is an array version of the scalar function above.
      *

include/ginkgo/core/matrix/csr.hpp

@@ -793,13 +793,6 @@ class Csr : public EnableLinOp<Csr<ValueType, IndexType>>,
 
     std::unique_ptr<Diagonal<ValueType>> extract_diagonal() const override;
 
-    std::unique_ptr<Csr<ValueType, IndexType>> create_submatrix(
-        const gko::IndexSet<IndexType>& row_index_set,
-        const gko::IndexSet<IndexType>& column_index_set) const;
-
-    std::unique_ptr<Csr<ValueType, IndexType>> create_submatrix(
-        const gko::span& row_span, const gko::span& column_span) const;
-
     std::unique_ptr<absolute_type> compute_absolute() const override;
 
     void compute_absolute_inplace() override;
@@ -959,7 +952,7 @@ class Csr : public EnableLinOp<Csr<ValueType, IndexType>>,
         this->inv_scale_impl(make_temporary_clone(exec, alpha).get());
     }
 
-    /*
+    /**
      * Creates a constant (immutable) Csr matrix from a set of constant arrays.
      *
      * @param exec  the executor to create the matrix on
@@ -987,7 +980,7 @@ class Csr : public EnableLinOp<Csr<ValueType, IndexType>>,
             gko::detail::array_const_cast(std::move(row_ptrs)), strategy});
     }
 
-    /*
+    /**
      * This is version of create_const with a default strategy.
      */
     static std::unique_ptr<const Csr> create_const(
@@ -1001,6 +994,36 @@ class Csr : public EnableLinOp<Csr<ValueType, IndexType>>,
                                  Csr::make_default_strategy(exec));
     }
 
+    /**
+     * Creates a submatrix from this Csr matrix given row and column IndexSet
+     * objects.
+     *
+     * @param row_index_set  the row index set containing the set of rows to be
+     *                       in the submatrix.
+     * @param column_index_set  the col index set containing the set of columns
+     *                          to be in the submatrix.
+     * @return A new CSR matrix with the elements that belong to the row and
+     *          columns of this matrix as specified by the index sets.
+     * @note This is not a view but creates a new, separate CSR matrix.
+     */
+    std::unique_ptr<Csr<ValueType, IndexType>> create_submatrix(
+        const gko::IndexSet<IndexType>& row_index_set,
+        const gko::IndexSet<IndexType>& column_index_set) const;
+
+    /**
+     * Creates a submatrix from this Csr matrix given row and column spans
+     *
+     * @param row_span  the row span containing the contiguous set of rows to be
+     *                  in the submatrix.
+     * @param column_span  the column span containing the contiguous set of
+     *                     columns to be in the submatrix.
+     * @return A new CSR matrix with the elements that belong to the row and
+     *          columns of this matrix as specified by the index sets.
+     * @note This is not a view but creates a new, separate CSR matrix.
+     */
+    std::unique_ptr<Csr<ValueType, IndexType>> create_submatrix(
+        const gko::span& row_span, const gko::span& column_span) const;
+
 protected:
     /**
      * Creates an uninitialized CSR matrix of the specified size.

omp/base/index_set_kernels.cpp

@@ -73,10 +73,11 @@ void to_global_indices(std::shared_ptr<const DefaultExecutor> exec,
 {
 #pragma omp parallel for
     for (size_type subset = 0; subset < num_subsets; ++subset) {
-        for (size_type i = 0;
-             i < superset_indices[subset + 1] - superset_indices[subset]; ++i) {
-            decomp_indices[superset_indices[subset] + i] =
-                subset_begin[subset] + i;
+        IndexType local_i{};
+        for (auto i = superset_indices[subset];
+             i < superset_indices[subset + 1]; ++i) {
+            decomp_indices[i] = local_i + subset_begin[subset];
+            local_i++;
         }
     }
 }
@@ -157,20 +158,19 @@ void global_to_local(std::shared_ptr<const DefaultExecutor> exec,
 #pragma omp parallel for
     for (size_type i = 0; i < num_indices; ++i) {
         auto index = global_indices[i];
-        if (index >= index_space_size) {
+        if (index < 0 || index >= index_space_size) {
             local_indices[i] = invalid_index<IndexType>();
             continue;
         }
         const auto bucket = std::distance(
-            subset_begin,
-            std::upper_bound(subset_begin, subset_begin + num_subsets, index));
-        auto shifted_bucket = bucket == 0 ? 0 : (bucket - 1);
-        if (subset_end[shifted_bucket] <= index ||
-            index < subset_begin[shifted_bucket]) {
+            subset_begin + 1,
+            std::upper_bound(subset_begin + 1, subset_begin + num_subsets + 1,
+                             index));
+        if (index >= subset_end[bucket] || index < subset_begin[bucket]) {
             local_indices[i] = invalid_index<IndexType>();
         } else {
-            local_indices[i] = index - subset_begin[shifted_bucket] +
-                               superset_indices[shifted_bucket];
+            local_indices[i] =
+                index - subset_begin[bucket] + superset_indices[bucket];
         }
     }
 }
@@ -190,17 +190,16 @@ void local_to_global(std::shared_ptr<const DefaultExecutor> exec,
 #pragma omp parallel for
     for (size_type i = 0; i < num_indices; ++i) {
         auto index = local_indices[i];
-        if (index >= superset_indices[num_subsets]) {
+        if (index < 0 || index >= superset_indices[num_subsets]) {
             global_indices[i] = invalid_index<IndexType>();
             continue;
         }
         const auto bucket = std::distance(
-            superset_indices,
-            std::upper_bound(superset_indices,
+            superset_indices + 1,
+            std::upper_bound(superset_indices + 1,
                              superset_indices + num_subsets + 1, index));
-        auto shifted_bucket = bucket == 0 ? 0 : (bucket - 1);
-        global_indices[i] = subset_begin[shifted_bucket] + index -
-                            superset_indices[shifted_bucket];
+        global_indices[i] =
+            subset_begin[bucket] + index - superset_indices[bucket];
     }
 }
 

omp/matrix/csr_kernels.cpp

@@ -749,7 +749,7 @@ void calculate_nonzeros_per_row_in_index_set(
     std::shared_ptr<const DefaultExecutor> exec,
     const matrix::Csr<ValueType, IndexType>* source,
     const IndexSet<IndexType>& row_index_set,
-    const IndexSet<IndexType>& col_index_set, Array<IndexType>* row_nnz)
+    const IndexSet<IndexType>& col_index_set, IndexType* row_nnz)
 {
     auto num_row_subsets = row_index_set.get_num_subsets();
     auto num_col_subsets = col_index_set.get_num_subsets();
@@ -765,7 +765,7 @@ void calculate_nonzeros_per_row_in_index_set(
         size_type res_row = row_superset_indices[set];
         for (auto row = row_subset_begin[set]; row < row_subset_end[set];
              ++row) {
-            row_nnz->get_data()[res_row] = zero<IndexType>();
+            row_nnz[res_row] = zero<IndexType>();
             for (size_type i = src_ptrs[row]; i < src_ptrs[row + 1]; ++i) {
                 auto index = source->get_const_col_idxs()[i];
                 if (index >= col_index_set.get_size()) {
@@ -777,11 +777,9 @@ void calculate_nonzeros_per_row_in_index_set(
                                      col_subset_begin + num_col_subsets,
                                      index));
                 auto shifted_bucket = bucket == 0 ? 0 : (bucket - 1);
-                if (col_subset_end[shifted_bucket] <= index ||
-                    (index < col_subset_begin[shifted_bucket])) {
-                    continue;
-                } else {
-                    row_nnz->get_data()[res_row]++;
+                if (index < col_subset_end[shifted_bucket] &&
+                    index >= col_subset_begin[shifted_bucket]) {
+                    row_nnz[res_row]++;
                 }
             }
             res_row++;
@@ -869,10 +867,8 @@ void compute_submatrix_from_index_set(
                                      col_subset_begin + num_col_subsets,
                                      index));
                 auto shifted_bucket = bucket == 0 ? 0 : (bucket - 1);
-                if (col_subset_end[shifted_bucket] <= index ||
-                    (index < col_subset_begin[shifted_bucket])) {
-                    continue;
-                } else {
+                if (index < col_subset_end[shifted_bucket] &&
+                    (index >= col_subset_begin[shifted_bucket])) {
                     res_col_idxs[res_nnz] =
                         index - col_subset_begin[shifted_bucket] +
                         col_superset_indices[shifted_bucket];

omp/test/matrix/csr_kernels.cpp

@@ -773,9 +773,9 @@ TEST_F(Csr, CalculateNnzPerRowInIndexSetIsEquivalentToRef)
     auto drow_nnz = gko::Array<int>(this->omp, row_nnz);
 
     gko::kernels::reference::csr::calculate_nonzeros_per_row_in_index_set(
-        this->ref, this->mtx2.get(), rset, cset, &row_nnz);
+        this->ref, this->mtx2.get(), rset, cset, row_nnz.get_data());
     gko::kernels::omp::csr::calculate_nonzeros_per_row_in_index_set(
-        this->omp, this->dmtx2.get(), drset, dcset, &drow_nnz);
+        this->omp, this->dmtx2.get(), drset, dcset, drow_nnz.get_data());
 
     GKO_ASSERT_ARRAY_EQ(row_nnz, drow_nnz);
 }
@@ -797,7 +797,7 @@ TEST_F(Csr, ComputeSubmatrixFromIndexSetIsEquivalentToRef)
     auto row_nnz = gko::Array<int>(this->ref, rset.get_num_elems() + 1);
     row_nnz.fill(gko::zero<int>());
     gko::kernels::reference::csr::calculate_nonzeros_per_row_in_index_set(
-        this->ref, this->mtx2.get(), rset, cset, &row_nnz);
+        this->ref, this->mtx2.get(), rset, cset, row_nnz.get_data());
     gko::kernels::reference::components::prefix_sum(
         this->ref, row_nnz.get_data(), row_nnz.get_num_elems());
     auto num_nnz = row_nnz.get_data()[rset.get_num_elems()];

reference/base/index_set_kernels.cpp

@@ -181,7 +181,7 @@ void global_to_local(std::shared_ptr<const DefaultExecutor> exec,
             shifted_bucket = 0;
         }
         auto index = global_indices[i];
-        if (index >= index_space_size) {
+        if (index < 0 || index >= index_space_size) {
             local_indices[i] = invalid_index<IndexType>();
             continue;
         }
@@ -221,7 +221,7 @@ void local_to_global(std::shared_ptr<const DefaultExecutor> exec,
             shifted_bucket = 0;
         }
         auto index = local_indices[i];
-        if (index >= superset_indices[num_subsets]) {
+        if (index < 0 || index >= superset_indices[num_subsets]) {
             global_indices[i] = invalid_index<IndexType>();
             continue;
         }