Kernel updates and batch_random_matrix gen

common/cuda_hip/matrix/batch_ell_kernels.hpp.inc

@@ -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 < idx) {
+            if (col_idx == invalid_index<int>()) {
                 break;
             } else {
                 temp += val[ind] * b[col_idx];
@@ -102,7 +102,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 < idx) {
+            if (col_idx == invalid_index<int>()) {
                 break;
             } else {
                 temp += alpha * val[ind] * b[col_idx];

core/matrix/batch_ell.cpp

@@ -128,13 +128,6 @@ Ell<ValueType, IndexType>::create_const(
 }
 
 
-inline const batch_dim<2> get_col_sizes(const batch_dim<2>& sizes)
-{
-    return batch_dim<2>(sizes.get_num_batch_items(),
-                        dim<2>(1, sizes.get_common_size()[1]));
-}
-
-
 template <typename ValueType, typename IndexType>
 Ell<ValueType, IndexType>::Ell(std::shared_ptr<const Executor> exec,
                                const batch_dim<2>& size,

core/test/utils/batch_helpers.hpp

@@ -82,11 +82,22 @@ std::unique_ptr<MatrixType> generate_random_batch_matrix(
     auto result = MatrixType::create(
         exec, batch_dim<2>(num_batch_items, dim<2>(num_rows, num_cols)),
         std::forward<MatrixArgs>(args)...);
+    auto sp_mat = generate_random_device_matrix_data<value_type, index_type>(
+        num_rows, num_cols, nonzero_dist, value_dist, engine,
+        exec->get_master());
+    auto row_idxs = gko::array<index_type>::const_view(
+                        exec->get_master(), sp_mat.get_num_elems(),
+                        sp_mat.get_const_row_idxs())
+                        .copy_to_array();
+    auto col_idxs = gko::array<index_type>::const_view(
+                        exec->get_master(), sp_mat.get_num_elems(),
+                        sp_mat.get_const_col_idxs())
+                        .copy_to_array();
 
     for (size_type b = 0; b < num_batch_items; b++) {
-        auto rand_mat =
-            generate_random_matrix<typename MatrixType::unbatch_type>(
-                num_rows, num_cols, nonzero_dist, value_dist, engine, exec);
+        auto rand_mat = fill_random_matrix_with_sparsity_pattern<
+            typename MatrixType::unbatch_type, index_type>(
+            num_rows, num_cols, row_idxs, col_idxs, value_dist, engine, exec);
         result->create_view_for_item(b)->copy_from(rand_mat.get());
     }
 

core/test/utils/matrix_generator.hpp

@@ -42,6 +42,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <vector>
 
 
+#include <ginkgo/core/base/array.hpp>
 #include <ginkgo/core/base/math.hpp>
 #include <ginkgo/core/base/utils.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
@@ -54,6 +55,49 @@ namespace gko {
 namespace test {
 
 
+/**
+ * Fills matrix data for a random matrix given a sparsity pattern
+ *
+ * @tparam ValueType  the type for matrix values
+ * @tparam IndexType  the type for row and column indices
+ * @tparam ValueDistribution  type of value distribution
+ * @tparam Engine  type of random engine
+ *
+ * @param num_rows  number of rows
+ * @param num_cols  number of columns
+ * @param row_idxs  the row indices of the matrix
+ * @param col_idxs  the column indices of the matrix
+ * @param value_dist  distribution of matrix values
+ * @param engine  a random engine
+ *
+ * @return the generated matrix_data with entries according to the given
+ *         dimensions and nonzero count and value distributions.
+ */
+template <typename ValueType, typename IndexType, typename ValueDistribution,
+          typename Engine>
+matrix_data<ValueType, IndexType> fill_random_matrix_data(
+    size_type num_rows, size_type num_cols,
+    const gko::array<IndexType>& row_indices,
+    const gko::array<IndexType>& col_indices, ValueDistribution&& value_dist,
+    Engine&& engine)
+{
+    matrix_data<ValueType, IndexType> data{gko::dim<2>{num_rows, num_cols}, {}};
+    auto host_exec = row_indices.get_executor()->get_master();
+    auto host_row_indices = make_temporary_clone(host_exec, &row_indices);
+    auto host_col_indices = make_temporary_clone(host_exec, &col_indices);
+
+    for (int nnz = 0; nnz < row_indices.get_num_elems(); ++nnz) {
+        data.nonzeros.emplace_back(
+            host_row_indices->get_const_data()[nnz],
+            host_col_indices->get_const_data()[nnz],
+            detail::get_rand_value<ValueType>(value_dist, engine));
+    }
+
+    data.ensure_row_major_order();
+    return data;
+}
+
+
 /**
  * Generates matrix data for a random matrix.
  *
@@ -156,13 +200,59 @@ generate_random_device_matrix_data(gko::size_type num_rows,
 }
 
 
+/**
+ * Fills a random matrix with given sparsity pattern.
+ *
+ * @tparam MatrixType  type of matrix to generate (must implement
+ *                     the interface `ReadableFromMatrixData<>` and provide
+ *                     matching `value_type` and `index_type` type aliases)
+ *
+ * @param num_rows  number of rows
+ * @param num_cols  number of columns
+ * @param value_dist  distribution of matrix values
+ * @param row_idxs  the row indices of the matrix
+ * @param col_idxs  the column indices of the matrix
+ * @param exec  executor where the matrix should be allocated
+ * @param args  additional arguments for the matrix constructor
+ *
+ * The other (template) parameters match generate_random_matrix_data.
+ *
+ * @return the unique pointer of MatrixType
+ */
+template <typename MatrixType = matrix::Dense<>,
+          typename IndexType = typename MatrixType::index_type,
+          typename ValueDistribution, typename Engine, typename... MatrixArgs>
+std::unique_ptr<MatrixType> fill_random_matrix_with_sparsity_pattern(
+    size_type num_rows, size_type num_cols,
+    const gko::array<IndexType>& row_idxs,
+    const gko::array<IndexType>& col_idxs, ValueDistribution&& value_dist,
+    Engine&& engine, std::shared_ptr<const Executor> exec, MatrixArgs&&... args)
+{
+    using value_type = typename MatrixType::value_type;
+    using index_type = IndexType;
+
+    GKO_ASSERT(row_idxs.get_num_elems() == col_idxs.get_num_elems());
+    GKO_ASSERT(row_idxs.get_num_elems() < (num_rows * num_cols));
+    auto result = MatrixType::create(exec, std::forward<MatrixArgs>(args)...);
+    result->read(fill_random_matrix_data<value_type, index_type>(
+        num_rows, num_cols, row_idxs, col_idxs,
+        std::forward<ValueDistribution>(value_dist),
+        std::forward<Engine>(engine)));
+    return result;
+}
+
+
 /**
  * Generates a random matrix.
  *
  * @tparam MatrixType  type of matrix to generate (must implement
  *                     the interface `ReadableFromMatrixData<>` and provide
  *                     matching `value_type` and `index_type` type aliases)
  *
+ * @param num_rows  number of rows
+ * @param num_cols  number of columns
+ * @param nonzero_dist  distribution of nonzeros per row
+ * @param value_dist  distribution of matrix values
  * @param exec  executor where the matrix should be allocated
  * @param args  additional arguments for the matrix constructor
  *

cuda/matrix/batch_ell_kernels.cu

@@ -37,6 +37,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 
 #include <ginkgo/core/base/batch_multi_vector.hpp>
+#include <ginkgo/core/base/types.hpp>
 #include <ginkgo/core/matrix/batch_ell.hpp>
 
 

dpcpp/matrix/batch_ell_kernels.hpp.inc

@@ -42,38 +42,37 @@ __dpct_inline__ void simple_apply_kernel(
         auto temp = zero<ValueType>();
         for (size_type idx = 0; idx < mat.num_stored_elems_per_row; idx++) {
             const auto col_idx = mat.col_idxs[tidx + idx * mat.stride];
-            if (col_idx < idx)
+            if (col_idx == invalid_index<int>()) {
                 break;
-            else
-                temp += mat.values[tidx + idx * mat.stride] *
-                        b.values[col_idx * b.stride];
+                else temp += mat.values[tidx + idx * mat.stride] *
+                             b.values[col_idx * b.stride];
+            }
+            x.values[tidx * x.stride] = temp;
         }
-        x.values[tidx * x.stride] = temp;
     }
-}
 
 
-template <typename ValueType>
-__dpct_inline__ void advanced_apply_kernel(
-    const gko::batch::multi_vector::batch_item<const ValueType>& alpha,
-    const gko::batch::matrix::ell::batch_item<const ValueType>& mat,
-    const gko::batch::multi_vector::batch_item<const ValueType>& b,
-    const gko::batch::multi_vector::batch_item<const ValueType>& beta,
-    const gko::batch::multi_vector::batch_item<ValueType>& x,
-    sycl::nd_item<3>& item_ct1)
-{
-    for (int tidx = item_ct1.get_local_linear_id(); tidx < mat.num_rows;
-         tidx += item_ct1.get_local_range().size()) {
-        auto temp = zero<ValueType>();
-        for (size_type idx = 0; idx < mat.num_stored_elems_per_row; idx++) {
-            const auto col_idx = mat.col_idxs[tidx + idx * mat.stride];
-            if (col_idx < idx)
-                break;
-            else
-                temp += alpha.values[0] * mat.values[tidx + idx * mat.stride] *
-                        b.values[col_idx * b.stride];
+    template <typename ValueType>
+    __dpct_inline__ void advanced_apply_kernel(
+        const gko::batch::multi_vector::batch_item<const ValueType>& alpha,
+        const gko::batch::matrix::ell::batch_item<const ValueType>& mat,
+        const gko::batch::multi_vector::batch_item<const ValueType>& b,
+        const gko::batch::multi_vector::batch_item<const ValueType>& beta,
+        const gko::batch::multi_vector::batch_item<ValueType>& x,
+        sycl::nd_item<3>& item_ct1)
+    {
+        for (int tidx = item_ct1.get_local_linear_id(); tidx < mat.num_rows;
+             tidx += item_ct1.get_local_range().size()) {
+            auto temp = zero<ValueType>();
+            for (size_type idx = 0; idx < mat.num_stored_elems_per_row; idx++) {
+                const auto col_idx = mat.col_idxs[tidx + idx * mat.stride];
+                if (col_idx == invalid_index<int>()) {
+                    break;
+                    else temp += alpha.values[0] *
+                                 mat.values[tidx + idx * mat.stride] *
+                                 b.values[col_idx * b.stride];
+                }
+                x.values[tidx * x.stride] =
+                    temp + beta.values[0] * x.values[tidx * x.stride];
+            }
         }
-        x.values[tidx * x.stride] =
-            temp + beta.values[0] * x.values[tidx * x.stride];
-    }
-}

hip/matrix/batch_ell_kernels.hip.cpp

@@ -38,6 +38,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 
 #include <ginkgo/core/base/batch_multi_vector.hpp>
+#include <ginkgo/core/base/types.hpp>
 #include <ginkgo/core/matrix/batch_ell.hpp>
 
 

include/ginkgo/core/matrix/batch_ell.hpp

@@ -58,9 +58,14 @@ namespace matrix {
  * Ell is a sparse matrix format that stores the same number of nonzeros in each
  * row, enabling coalesced accesses. It is suitable for sparsity patterns that
  * have a similar number of nonzeros in every row. The values are stored in a
- * column-major fashion similar to the monolithic gko::matrix::Ell class. It is
- * also assumed that the sparsity pattern of all the items in the batch is the
- * same and therefore only a single copy of the sparsity pattern is stored.
+ * column-major fashion similar to the monolithic gko::matrix::Ell class.
+ *
+ * Similar to the monolithic gko::matrix::Ell class, invalid_index<IndexType> is
+ * used as the column index for padded zero entries.
+ *
+ * @note It is also assumed that the sparsity pattern of all the items in the
+ * batch is the same and therefore only a single copy of the sparsity pattern is
+ * stored.
  *
  * @tparam ValueType  value precision of matrix elements
  * @tparam IndexType  index precision of matrix elements
@@ -253,13 +258,13 @@ class Ell final
 
     /**
      * Creates a constant (immutable) batch ell matrix from a constant
-     * array.
+     * array. The column indices array needs to be the same for all batch items.
      *
      * @param exec  the executor to create the matrix on
      * @param size  the dimensions of the matrix
      * @param num_elems_per_row  the number of elements to be stored in each row
      * @param values  the value array of the matrix
-     * @param col_idxs the col_idxs array of the matrix
+     * @param col_idxs the col_idxs array of a single batch item of the matrix.
      *
      * @return A smart pointer to the constant matrix wrapping the input
      * array (if it resides on the same executor as the matrix) or a copy of the
@@ -325,15 +330,15 @@ class Ell final
 
     /**
      * Creates a Ell matrix from an already allocated (and initialized)
-     * array.
+     * array. The column indices array needs to be the same for all batch items.
      *
      * @tparam ValuesArray  type of array of values
      *
      * @param exec  Executor associated to the matrix
      * @param size  size of the matrix
      * @param num_elems_per_row  the number of elements to be stored in each row
      * @param values  array of matrix values
-     * @param col_idxs the col_idxs array of the matrix
+     * @param col_idxs the col_idxs array of a single batch item of the matrix.
      *
      * @note If `values` is not an rvalue, not an array of ValueType, or is on
      *       the wrong executor, an internal copy will be created, and the

test/matrix/batch_ell_kernels.cpp

@@ -55,7 +55,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 class Ell : public CommonTestFixture {
 protected:
-    using BMtx = gko::batch::matrix::Ell<value_type>;
+    using BMtx = gko::batch::matrix::Ell<value_type, gko::int32>;
     using BMVec = gko::batch::MultiVector<value_type>;
 
     Ell() : rand_engine(15) {}