From a63443d7f28e428c28f04d125c05ec2547ce7fe0 Mon Sep 17 00:00:00 2001
From: Pratik Nayak <pratikvn@protonmail.com>
Date: Sat, 21 Oct 2023 20:43:53 +0200
Subject: [PATCH] Add batch matrix gen utils and test helpers

---
 core/test/utils/batch_helpers.hpp           | 138 ++++++++++++++++++++
 reference/preconditioner/batch_identity.hpp |   1 -
 2 files changed, 138 insertions(+), 1 deletion(-)

diff --git a/core/test/utils/batch_helpers.hpp b/core/test/utils/batch_helpers.hpp
index 5b1fa60ed36..40eeae07a25 100644
--- a/core/test/utils/batch_helpers.hpp
+++ b/core/test/utils/batch_helpers.hpp
@@ -39,12 +39,14 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 
 #include <ginkgo/core/base/batch_multi_vector.hpp>
+#include <ginkgo/core/base/device_matrix_data.hpp>
 #include <ginkgo/core/base/matrix_data.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
 
 
 #include "core/test/utils/assertions.hpp"
 #include "core/test/utils/matrix_generator.hpp"
+#include "core/utils/matrix_utils.hpp"
 
 
 namespace gko {
@@ -106,6 +108,142 @@ std::unique_ptr<MatrixType> generate_random_batch_matrix(
 }
 
 
+/**
+ * Generate a batch of 1D Poisson (3pt stencil, {-1, 2, -1}) matrices in the
+ * given input matrix format.
+ *
+ * @tparam MatrixType  The concrete type of the output matrix.
+ *
+ * @param exec  The executor.
+ * @param num_rows  The size (number of rows) of the generated matrix
+ * @param num_batch_items  The number of Poisson matrices in the batch
+ * @param args The create args to be forwarded to the matrix
+ */
+template <typename MatrixType, typename... MatrixArgs>
+std::unique_ptr<MatrixType> generate_3pt_stencil_batch_matrix(
+    std::shared_ptr<const Executor> exec, const int num_rows,
+    const size_type num_batch_items, MatrixArgs&&... args)
+{
+    using value_type = typename MatrixType::value_type;
+    using index_type = typename MatrixType::index_type;
+    const int num_cols = num_rows;
+    gko::matrix_data<value_type, index_type> data{
+        gko::dim<2>{static_cast<size_type>(num_rows),
+                    static_cast<size_type>(num_cols)},
+        {}};
+    for (int row = 1; row < num_rows - 1; ++row) {
+        data.nonzeros.emplace_back(row, row + 1, value_type{-1.0});
+        data.nonzeros.emplace_back(row - 1, row, value_type{-1.0});
+        data.nonzeros.emplace_back(row, row, value_type{2.0});
+    }
+    data.nonzeros.emplace_back(0, 0, value_type{2.0});
+    data.nonzeros.emplace_back(num_rows - 1, num_rows - 1, value_type{2.0});
+    data.nonzeros.emplace_back(num_rows - 1, num_rows - 2, value_type{-1.0});
+    data.nonzeros.emplace_back(0, 1, value_type{-1.0});
+
+    std::vector<gko::matrix_data<value_type, index_type>> batch_data(
+        num_batch_items, data);
+    return gko::batch::read<value_type, index_type, MatrixType>(
+        exec, batch_data, std::forward<MatrixArgs>(args)...);
+}
+
+
+template <typename ValueType>
+struct BatchSystem {
+    using vec_type = batch::MultiVector<ValueType>;
+    std::unique_ptr<batch::BatchLinOp> A;
+    std::unique_ptr<vec_type> b;
+};
+
+
+template <typename MatrixType, typename... MatrixArgs>
+BatchSystem<typename MatrixType::value_type>
+generate_diag_dominant_batch_system(std::shared_ptr<const gko::Executor> exec,
+                                    const size_type num_batch_items,
+                                    const int num_rows, const int num_rhs,
+                                    const bool is_hermitian,
+                                    MatrixArgs&&... args)
+{
+    using value_type = typename MatrixType::value_type;
+    using index_type = typename MatrixType::index_type;
+    using unbatch_type = typename MatrixType::unbatch_type;
+    using real_type = remove_complex<value_type>;
+    const int num_cols = num_rows;
+    gko::matrix_data<value_type, index_type> data{
+        gko::dim<2>{static_cast<size_type>(num_rows),
+                    static_cast<size_type>(num_cols)},
+        {}};
+    auto engine = std::default_random_engine(42);
+    auto rand_diag_dist = std::normal_distribution<value_type>(4.0, 12.0);
+    for (int row = 1; row < num_rows - 1; ++row) {
+        auto rand_nnz_dist = std::normal_distribution<index_type>(1, row + 1);
+        auto k = detail::get_rand_value<index_type>(rand_nnz_dist, engine);
+        data.nonzeros.emplace_back(row, k, value_type{-1.0});
+        data.nonzeros.emplace_back(row, row + 1, value_type{-1.0});
+        data.nonzeros.emplace_back(row - 1, row, value_type{-1.0});
+        data.nonzeros.emplace_back(
+            row, row, detail::get_rand_value(rand_diag_dist, engine));
+    }
+    data.nonzeros.emplace_back(0, 0, value_type{2.0});
+    data.nonzeros.emplace_back(num_rows - 1, num_rows - 1, value_type{2.0});
+    data.nonzeros.emplace_back(num_rows - 1, num_rows - 2, value_type{-1.0});
+    data.nonzeros.emplace_back(0, 1, value_type{-1.0});
+
+    if (is_hermitian) {
+        gko::utils::make_hpd(data);
+    }
+    data.ensure_row_major_order();
+
+    auto soa_data =
+        gko::device_matrix_data<value_type, index_type>::create_from_host(
+            exec->get_master(), data);
+    auto row_idxs = gko::array<index_type>::const_view(
+                        exec->get_master(), soa_data.get_num_elems(),
+                        soa_data.get_const_row_idxs())
+                        .copy_to_array();
+    auto col_idxs = gko::array<index_type>::const_view(
+                        exec->get_master(), soa_data.get_num_elems(),
+                        soa_data.get_const_col_idxs())
+                        .copy_to_array();
+    auto result = MatrixType::create(
+        exec, batch_dim<2>(num_batch_items, dim<2>(num_rows, num_cols)),
+        std::forward<MatrixArgs>(args)...);
+    auto rand_val_dist = std::normal_distribution<value_type>(-0.5, 0.5);
+    std::vector<gko::matrix_data<value_type, index_type>> batch_data(
+        num_batch_items);
+    batch_data.reserve(num_batch_items);
+    BatchSystem<value_type> sys;
+
+    for (size_type b = 1; b < num_batch_items; b++) {
+        auto rand_data = fill_random_matrix_data<value_type, index_type>(
+            num_rows, num_cols, row_idxs, col_idxs, rand_val_dist, engine);
+        if (is_hermitian) {
+            gko::utils::make_hpd(rand_data);
+        } else {
+            gko::utils::make_diag_dominant(rand_data);
+        }
+        batch_data.emplace_back(rand_data);
+    }
+    sys.A = gko::give(gko::batch::read<value_type, index_type, MatrixType>(
+        exec, batch_data, std::forward<MatrixArgs>(args)...));
+
+    std::vector<gko::matrix_data<value_type, index_type>> batch_rhs_data(
+        num_batch_items);
+    batch_rhs_data.reserve(num_batch_items);
+    for (size_type b = 0; b < num_batch_items; b++) {
+        auto rand_data = generate_random_matrix_data<value_type, index_type>(
+            num_rows, num_cols,
+            std::normal_distribution<index_type>(num_rhs, num_rhs),
+            rand_val_dist, engine);
+        batch_data.emplace_back(rand_data);
+    }
+    sys.b = gko::give(gko::batch::read<value_type, index_type,
+                                       BatchSystem<value_type>::vec_type>(
+        exec, batch_rhs_data));
+    return sys;
+}
+
+
 }  // namespace test
 }  // namespace gko
 
diff --git a/reference/preconditioner/batch_identity.hpp b/reference/preconditioner/batch_identity.hpp
index c26beb155ac..e628997b0da 100644
--- a/reference/preconditioner/batch_identity.hpp
+++ b/reference/preconditioner/batch_identity.hpp
@@ -35,7 +35,6 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 
 #include "core/matrix/batch_struct.hpp"
-#include "reference/base/config.hpp"
 
 
 namespace gko {