review updates:

- documentation
- use templated function
- fix test

Co-authored-by: Yu-Hsiang M. Tsai <[email protected]>

Author: MarcelKoch

core/matrix/dense.cpp

@@ -1447,27 +1447,8 @@ void Dense<ValueType>::row_gather(ptr_param<const LinOp> alpha,
 
 
 template <typename ValueType>
-void Dense<ValueType>::row_scatter(const array<int32>* row_idxs,
-                                   ptr_param<LinOp> row_collection) const
-{
-    gather_mixed_real_complex<ValueType>(
-        [&](auto dense) { this->row_scatter_impl(row_idxs, dense); },
-        row_collection.get());
-}
-
-
-template <typename ValueType>
-void Dense<ValueType>::row_scatter(const array<int64>* row_idxs,
-                                   ptr_param<LinOp> row_collection) const
-{
-    gather_mixed_real_complex<ValueType>(
-        [&](auto dense) { this->row_scatter_impl(row_idxs, dense); },
-        row_collection.get());
-}
-
-
-template <typename ValueType>
-void Dense<ValueType>::row_scatter(const index_set<int32>* row_idxs,
+template <typename IndexType>
+void Dense<ValueType>::row_scatter(const array<IndexType>* row_idxs,
                                    ptr_param<LinOp> row_collection) const
 {
     gather_mixed_real_complex<ValueType>(
@@ -1477,7 +1458,8 @@ void Dense<ValueType>::row_scatter(const index_set<int32>* row_idxs,
 
 
 template <typename ValueType>
-void Dense<ValueType>::row_scatter(const index_set<int64>* row_idxs,
+template <typename IndexType>
+void Dense<ValueType>::row_scatter(const index_set<IndexType>* row_idxs,
                                    ptr_param<LinOp> row_collection) const
 {
     gather_mixed_real_complex<ValueType>(
@@ -1786,6 +1768,18 @@ std::unique_ptr<Dense<ValueType>> Dense<ValueType>::create_submatrix_impl(
 #define GKO_DECLARE_DENSE_MATRIX(_type) class Dense<_type>
 GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_MATRIX);
 
+#define GKO_DECLARE_DENSE_ROW_SCATTER_ARRAY(_vtype, _itype)        \
+    void Dense<_vtype>::row_scatter(const array<_itype>* row_idxs, \
+                                    ptr_param<LinOp> row_collection) const
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_DENSE_ROW_SCATTER_ARRAY);
+
+#define GKO_DECLARE_DENSE_ROW_SCATTER_INDEX_SET(_vtype, _itype)        \
+    void Dense<_vtype>::row_scatter(const index_set<_itype>* row_idxs, \
+                                    ptr_param<LinOp> row_collection) const
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_DENSE_ROW_SCATTER_INDEX_SET);
+
 
 }  // namespace matrix
 }  // namespace gko

include/ginkgo/core/base/index_set.hpp

@@ -471,6 +471,8 @@ struct temporary_clone_helper<const index_set<T>> {
         return std::make_unique<const index_set<T>>(std::move(exec), *ptr);
     }
 };
+
+
 }  // namespace detail
 
 

include/ginkgo/core/matrix/dense.hpp

@@ -546,6 +546,8 @@ class Dense
     /**
      * Copies this matrix into the given rows of the target matrix.
      *
+     * @tparam IndexType  the index type, either int32 or int 64
+     *
      * @param scatter_indices  row indices of the target matrix. It must
      *                         have the same number of indices as rows in
      *                         this matrix.
@@ -554,33 +556,27 @@ class Dense
      *
      * @warning scatter_indices may not contain duplicates, unless if
      *          for indices `i, j` with `scatter_indices[i] ==
-     * scatter_indices[j]` the rows `i, j` of this matrix are identical.
+     *          scatter_indices[j]` the rows `i, j` of this matrix are
+     *          identical.
      */
-    void row_scatter(const array<int64>* scatter_indices,
-                     ptr_param<LinOp> target) const;
-
-    /**
-     * @copydoc row_scatter(const array<int64>*, ptr_param<LinOp>) const
-     */
-    void row_scatter(const array<int32>* scatter_indices,
+    template <typename IndexType>
+    void row_scatter(const array<IndexType>* scatter_indices,
                      ptr_param<LinOp> target) const;
 
     /**
      * Copies this matrix into the given rows of the target matrix.
      *
+     * @tparam IndexType  the index type, either int32 or int 64
+     *
      * @param scatter_indices  row indices of the target matrix. It must
      *                         have the same number of indices as rows in
      *                         this matrix.
      * @param target  matrix where the scattered rows are stored, i.e.
-     *                `target(scatter_indices[i], j) = this(i, j)`
-     */
-    void row_scatter(const index_set<int64>* scatter_indices,
-                     ptr_param<LinOp> target) const;
-
-    /**
-     * @copydoc row_scatter(const index_set<int64>*, ptr_param<LinOp>) const
+     *                `target(scatter_indices.get_global_index(i), j)
+     *                   = this(i, j)`
      */
-    void row_scatter(const index_set<int32>* scatter_indices,
+    template <typename IndexType>
+    void row_scatter(const index_set<IndexType>* scatter_indices,
                      ptr_param<LinOp> target) const;
 
     std::unique_ptr<LinOp> column_permute(

test/matrix/dense_kernels.cpp

@@ -163,7 +163,7 @@ class Dense : public CommonTestFixture {
         rscatter_idxs = std::unique_ptr<Arr>(
             new Arr{ref, tmp2.begin(), tmp2.begin() + u->get_size()[0]});
         rscatter_idxs_sub = std::unique_ptr<Arr>(
-            new Arr{ref, tmp2.begin(), tmp2.begin() + u->get_size()[0]});
+            new Arr{ref, tmp4.begin(), tmp4.begin() + u->get_size()[0]});
     }
 
     template <typename ConvertedType, typename InputType>
@@ -1324,8 +1324,8 @@ TEST_F(Dense, CanScatterRowsIntoDenseCrossExecutor)
 {
     set_up_apply_data();
 
-    u->row_scatter(rgather_idxs.get(), x);
-    u->row_scatter(rgather_idxs.get(), dx);
+    u->row_scatter(rscatter_idxs.get(), x);
+    u->row_scatter(rscatter_idxs.get(), dx);
 
     GKO_ASSERT_MTX_NEAR(x, dx, 0);
 }