Review updates

Co-authored-by: Marcel Koch <marcel.koch@kit.edu> Co-authored-by: Yu-Hsiang Tsai <yhmtsai@gmail.com>
ginkgo-project · Oct 11, 2023 · 976d144 · 976d144
1 parent 5441457
commit 976d144
Show file tree

Hide file tree

Showing 12 changed files with 229 additions and 247 deletions.
diff --git a/common/cuda_hip/matrix/batch_ell_kernels.hpp.inc b/common/cuda_hip/matrix/batch_ell_kernels.hpp.inc
@@ -31,9 +31,9 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ******************************<GINKGO LICENSE>*******************************/
 
 
-template <typename ValueType>
+template <typename ValueType, typename IndexType>
 __device__ __forceinline__ void simple_apply(
-    const gko::batch::matrix::ell::batch_item<const ValueType>& mat,
+    const gko::batch::matrix::ell::batch_item<const ValueType, IndexType>& mat,
     const ValueType* const __restrict__ b, ValueType* const __restrict__ x)
 {
     const auto num_rows = mat.num_rows;
@@ -46,7 +46,7 @@ __device__ __forceinline__ void simple_apply(
         for (size_type idx = 0; idx < num_stored_elements_per_row; idx++) {
             const auto ind = tidx + idx * stride;
             const auto col_idx = col[ind];
-            if (col_idx == invalid_index<int>()) {
+            if (col_idx == invalid_index<IndexType>()) {
                 break;
             } else {
                 temp += val[ind] * b[col_idx];
@@ -56,12 +56,13 @@ __device__ __forceinline__ void simple_apply(
     }
 }
 
-template <typename ValueType>
+template <typename ValueType, typename IndexType>
 __global__ __launch_bounds__(
     default_block_size,
     sm_oversubscription) void simple_apply_kernel(const gko::batch::matrix::
                                                       ell::uniform_batch<
-                                                          const ValueType>
+                                                          const ValueType,
+                                                          IndexType>
                                                           mat,
                                                   const gko::batch::
                                                       multi_vector::
@@ -85,10 +86,10 @@ __global__ __launch_bounds__(
 }
 
 
-template <typename ValueType>
+template <typename ValueType, typename IndexType>
 __device__ __forceinline__ void advanced_apply(
     const ValueType alpha,
-    const gko::batch::matrix::ell::batch_item<const ValueType>& mat,
+    const gko::batch::matrix::ell::batch_item<const ValueType, IndexType>& mat,
     const ValueType* const __restrict__ b, const ValueType beta,
     ValueType* const __restrict__ x)
 {
@@ -102,7 +103,7 @@ __device__ __forceinline__ void advanced_apply(
         for (size_type idx = 0; idx < num_stored_elements_per_row; idx++) {
             const auto ind = tidx + idx * stride;
             const auto col_idx = col[ind];
-            if (col_idx == invalid_index<int>()) {
+            if (col_idx == invalid_index<IndexType>()) {
                 break;
             } else {
                 temp += alpha * val[ind] * b[col_idx];
@@ -112,7 +113,7 @@ __device__ __forceinline__ void advanced_apply(
     }
 }
 
-template <typename ValueType>
+template <typename ValueType, typename IndexType>
 __global__ __launch_bounds__(
     default_block_size,
     sm_oversubscription) void advanced_apply_kernel(const gko::batch::
@@ -122,7 +123,8 @@ __global__ __launch_bounds__(
                                                                 alpha,
                                                     const gko::batch::matrix::
                                                         ell::uniform_batch<
-                                                            const ValueType>
+                                                            const ValueType,
+                                                            IndexType>
                                                             mat,
                                                     const gko::batch::
                                                         multi_vector::

diff --git a/core/base/batch_utilities.hpp b/core/base/batch_utilities.hpp
@@ -165,12 +165,8 @@ std::vector<gko::matrix_data<ValueType, IndexType>> write(
 /**
  * Creates and initializes a batch of single column-vectors.
  *
- * This function first creates a temporary MultiVector, fills it with
- * passed in values, and then converts the vector to the requested type.
- *
- * @tparam Matrix  matrix type to initialize
- *                 (MultiVector has to implement the ConvertibleTo<Matrix>
- *                 interface)
+ * @tparam Matrix  matrix type to initialize (It has to implement the
+ *                 read<Matrix> function)
  * @tparam TArgs  argument types for Matrix::create method
  *                (not including the implied Executor as the first argument)
  *
@@ -180,7 +176,6 @@ std::vector<gko::matrix_data<ValueType, IndexType>> write(
  *                     including the Executor, which is passed as the first
  *                     argument
  *
- * @ingroup MultiVector
  * @ingroup mat_formats
  */
 template <typename Matrix, typename... TArgs>
@@ -220,23 +215,19 @@ std::unique_ptr<Matrix> initialize(
 
 
 /**
- * Creates and initializes a batch of multi-vectors.
- *
- * This function first creates a temporary MultiVector, fills it with
- * passed in values, and then converts the vector to the requested type.
+ * Creates and initializes a batch of matrices.
  *
- * @tparam Matrix  matrix type to initialize
- *                 (Dense has to implement the ConvertibleTo<Matrix> interface)
+ * @tparam Matrix  matrix type to initialize (It has to implement the
+ *                 read<Matrix> function)
  * @tparam TArgs  argument types for Matrix::create method
  *                (not including the implied Executor as the first argument)
  *
- * @param vals  values used to initialize the vector
- * @param exec  Executor associated to the vector
+ * @param vals  values used to initialize the matrix
+ * @param exec  Executor associated with the matrix
  * @param create_args  additional arguments passed to Matrix::create, not
  *                     including the Executor, which is passed as the first
  *                     argument
  *
- * @ingroup MultiVector
  * @ingroup mat_formats
  */
 template <typename Matrix, typename... TArgs>
@@ -290,23 +281,18 @@ std::unique_ptr<Matrix> initialize(
  * Creates and initializes a batch single column-vector by making copies of the
  * single input column vector.
  *
- * This function first creates a temporary batch multi-vector, fills it with
- * passed in values, and then converts the vector to the requested type.
- *
- * @tparam Matrix  matrix type to initialize
- *                 (MultiVector has to implement the ConvertibleTo<Matrix>
- *                  interface)
+ * @tparam Matrix  matrix type to initialize (It has to implement the
+ *                 read<Matrix> function)
  * @tparam TArgs  argument types for Matrix::create method
  *                (not including the implied Executor as the first argument)
  *
  * @param num_vectors  The number of times the input vector is to be duplicated
  * @param vals  values used to initialize each vector in the temp. batch
- * @param exec  Executor associated to the vector
+ * @param exec  Executor associated with the matrix
  * @param create_args  additional arguments passed to Matrix::create, not
  *                     including the Executor, which is passed as the first
  *                     argument
  *
- * @ingroup MultiVector
  * @ingroup mat_formats
  */
 template <typename Matrix, typename... TArgs>
@@ -343,23 +329,18 @@ std::unique_ptr<Matrix> initialize(
 /**
  * Creates and initializes a matrix from copies of a given matrix.
  *
- * This function first creates a temporary batch multi-vector, fills it with
- * passed in values, and then converts the vector to the requested type.
- *
- * @tparam Matrix  matrix type to initialize
- *                 (MultiVector has to implement the ConvertibleTo<Matrix>
- *                  interface)
+ * @tparam Matrix  matrix type to initialize (It has to implement the
+ *                 read<Matrix> function)
  * @tparam TArgs  argument types for Matrix::create method
  *                (not including the implied Executor as the first argument)
  *
  * @param num_batch_items The number of times the input matrix is duplicated
- * @param vals  values used to initialize each vector in the temp. batch
- * @param exec  Executor associated to the vector
+ * @param vals  values used to initialize each matrix in the temp. batch
+ * @param exec  Executor associated to the matrix
  * @param create_args  additional arguments passed to Matrix::create, not
  *                     including the Executor, which is passed as the first
  *                     argument
  *
- * @ingroup LinOp
  * @ingroup mat_formats
  */
 template <typename Matrix, typename... TArgs>

diff --git a/core/matrix/batch_struct.hpp b/core/matrix/batch_struct.hpp
@@ -89,10 +89,10 @@ namespace ell {
 /**
  * Encapsulates one matrix from a batch of ell matrices.
  */
-template <typename ValueType>
+template <typename ValueType, typename IndexType>
 struct batch_item {
     using value_type = ValueType;
-    using index_type = int32;
+    using index_type = IndexType;
 
     ValueType* values;
     const index_type* col_idxs;
@@ -106,11 +106,11 @@ struct batch_item {
 /**
  * A 'simple' structure to store a global uniform batch of ell matrices.
  */
-template <typename ValueType>
+template <typename ValueType, typename IndexType>
 struct uniform_batch {
     using value_type = ValueType;
-    using index_type = int32;
-    using entry_type = batch_item<value_type>;
+    using index_type = IndexType;
+    using entry_type = batch_item<value_type, index_type>;
 
     ValueType* values;
     const index_type* col_idxs;
@@ -164,27 +164,28 @@ GKO_ATTRIBUTES GKO_INLINE dense::batch_item<ValueType> extract_batch_item(
 }
 
 
-template <typename ValueType>
-GKO_ATTRIBUTES GKO_INLINE ell::batch_item<const ValueType> to_const(
-    const ell::batch_item<ValueType>& b)
+template <typename ValueType, typename IndexType>
+GKO_ATTRIBUTES GKO_INLINE ell::batch_item<const ValueType, IndexType> to_const(
+    const ell::batch_item<ValueType, IndexType>& b)
 {
     return {b.values,   b.col_idxs, b.stride,
             b.num_rows, b.num_cols, b.num_stored_elems_per_row};
 }
 
 
-template <typename ValueType>
-GKO_ATTRIBUTES GKO_INLINE ell::uniform_batch<const ValueType> to_const(
-    const ell::uniform_batch<ValueType>& ub)
+template <typename ValueType, typename IndexType>
+GKO_ATTRIBUTES GKO_INLINE ell::uniform_batch<const ValueType, IndexType>
+to_const(const ell::uniform_batch<ValueType, IndexType>& ub)
 {
     return {ub.values,   ub.col_idxs, ub.num_batch_items,         ub.stride,
             ub.num_rows, ub.num_cols, ub.num_stored_elems_per_row};
 }
 
 
-template <typename ValueType>
-GKO_ATTRIBUTES GKO_INLINE ell::batch_item<ValueType> extract_batch_item(
-    const ell::uniform_batch<ValueType>& batch, const size_type batch_idx)
+template <typename ValueType, typename IndexType>
+GKO_ATTRIBUTES GKO_INLINE ell::batch_item<ValueType, IndexType>
+extract_batch_item(const ell::uniform_batch<ValueType, IndexType>& batch,
+                   const size_type batch_idx)
 {
     return {batch.values +
                 batch_idx * batch.num_stored_elems_per_row * batch.num_rows,
@@ -195,11 +196,12 @@ GKO_ATTRIBUTES GKO_INLINE ell::batch_item<ValueType> extract_batch_item(
             batch.num_stored_elems_per_row};
 }
 
-template <typename ValueType>
-GKO_ATTRIBUTES GKO_INLINE ell::batch_item<ValueType> extract_batch_item(
-    ValueType* const batch_values, int* const batch_col_idxs, const int stride,
-    const int num_rows, const int num_cols, int num_elems_per_row,
-    const size_type batch_idx)
+template <typename ValueType, typename IndexType>
+GKO_ATTRIBUTES GKO_INLINE ell::batch_item<ValueType, IndexType>
+extract_batch_item(ValueType* const batch_values,
+                   IndexType* const batch_col_idxs, const int stride,
+                   const int num_rows, const int num_cols,
+                   int num_elems_per_row, const size_type batch_idx)
 {
     return {batch_values + batch_idx * num_elems_per_row * num_rows,
             batch_col_idxs,