From b60a633df6273c6e4a4d72df53cf947b84aac7b9 Mon Sep 17 00:00:00 2001
From: Marcel Koch <marcel.koch@kit.edu>
Date: Wed, 10 Jan 2024 13:11:53 +0000
Subject: [PATCH 1/4] add exception for 2D access

---
 include/ginkgo/core/base/exception_helpers.hpp | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/include/ginkgo/core/base/exception_helpers.hpp b/include/ginkgo/core/base/exception_helpers.hpp
index c9d7e612de0..0482a50a334 100644
--- a/include/ginkgo/core/base/exception_helpers.hpp
+++ b/include/ginkgo/core/base/exception_helpers.hpp
@@ -774,6 +774,20 @@ inline T ensure_allocated_impl(T ptr, const std::string& file, int line,
                   "semi-colon warnings")
 
 
+/**
+ * Ensures that an access is within the specified 2D dimensions.
+ *
+ * @param _row  the row access
+ * @param _col  the column access
+ * @param _bound  the dimension bound
+ *
+ * @throw OutOfBoundsError  if `_row >= _bound[0] || _col >= _bound[1]`
+ */
+#define GKO_ENSURE_IN_DIMENSION_BOUNDS(_row, _col, _bound)          \
+    GKO_ENSURE_IN_BOUNDS(_row, ::gko::detail::get_size(_bound)[0]); \
+    GKO_ENSURE_IN_BOUNDS(_col, ::gko::detail::get_size(_bound)[1])
+
+
 /**
  * Creates a StreamError exception.
  * This macro sets the correct information about the location of the error

From 242542b9b0c6f9bfcce02dd66bc9163d71164ad4 Mon Sep 17 00:00:00 2001
From: Marcel Koch <marcel.koch@kit.edu>
Date: Wed, 18 Oct 2023 10:15:21 +0200
Subject: [PATCH 2/4] allow dispatch helper to return value

---
 core/base/dispatch_helper.hpp | 127 ++++++++++++++++++++++++++--------
 1 file changed, 98 insertions(+), 29 deletions(-)

diff --git a/core/base/dispatch_helper.hpp b/core/base/dispatch_helper.hpp
index 5c030dedd17..62eccfde0c6 100644
--- a/core/base/dispatch_helper.hpp
+++ b/core/base/dispatch_helper.hpp
@@ -13,20 +13,55 @@
 
 
 namespace gko {
+namespace detail {
+
+
+/**
+ *
+ * @copydoc run<typename ReturnType, typename K, typename... Types, typename T,
+ *              typename Func, typename... Args>
+ *
+ * @note this is the end case
+ */
+template <typename ReturnType, typename T, typename Func, typename... Args>
+ReturnType run_impl(T obj, Func&&, Args&&...)
+{
+    GKO_NOT_SUPPORTED(obj);
+}
+
+/**
+ * @copydoc run<typename K, typename... Types, typename T, typename Func,
+ *              typename... Args>
+ *
+ * @note This has additionally the return type encoded.
+ */
+template <typename ReturnType, typename K, typename... Types, typename T,
+          typename Func, typename... Args>
+ReturnType run_impl(T obj, Func&& f, Args&&... args)
+{
+    if (auto dobj = dynamic_cast<K>(obj)) {
+        return f(dobj, std::forward<Args>(args)...);
+    } else {
+        return run_impl<ReturnType, Types...>(obj, std::forward<Func>(f),
+                                              std::forward<Args>(args)...);
+    }
+}
 
 
 /**
  * run uses template to go through the list and select the valid
  * template and run it.
  *
- * @tparam T  the type of input object
- * @tparam Func  the function will run if the object can be converted to K
- * @tparam ...Args  the additional arguments for the Func
+ * @tparam Base  the Base class with one template
+ * @tparam T  the type of input object waiting converted
+ * @tparam Func  the validation
+ * @tparam ...Args  the variadic arguments.
  *
  * @note this is the end case
  */
-template <typename T, typename Func, typename... Args>
-void run(T obj, Func, Args...)
+template <typename ReturnType, template <typename> class Base, typename T,
+          typename Func, typename... Args>
+ReturnType run_impl(T obj, Func, Args...)
 {
     GKO_NOT_SUPPORTED(obj);
 }
@@ -35,45 +70,68 @@ void run(T obj, Func, Args...)
  * run uses template to go through the list and select the valid
  * template and run it.
  *
- * @tparam K  the current type tried in the conversion
+ * @tparam Base  the Base class with one template
+ * @tparam K  the current template type of B. pointer of const Base<K> is tried
+ *            in the conversion.
  * @tparam ...Types  the other types will be tried in the conversion if K fails
- * @tparam T  the type of input object
- * @tparam Func  the function will run if the object can be converted to K
+ * @tparam T  the type of input object waiting converted
+ * @tparam Func  the function will run if the object can be converted to pointer
+ *               of const Base<K>
  * @tparam ...Args  the additional arguments for the Func
  *
  * @param obj  the input object waiting converted
  * @param f  the function will run if obj can be converted successfully
  * @param args  the additional arguments for the function
  */
-template <typename K, typename... Types, typename T, typename Func,
-          typename... Args>
-void run(T obj, Func f, Args... args)
+template <typename ReturnType, template <typename> class Base, typename K,
+          typename... Types, typename T, typename Func, typename... Args>
+ReturnType run_impl(T obj, Func&& f, Args&&... args)
 {
-    if (auto dobj = dynamic_cast<K>(obj)) {
-        f(dobj, args...);
+    if (auto dobj = std::dynamic_pointer_cast<const Base<K>>(obj)) {
+        return f(dobj, args...);
     } else {
-        run<Types...>(obj, f, args...);
+        return run_impl<ReturnType, Base, Types...>(
+            obj, std::forward<Func>(f), std::forward<Args>(args)...);
     }
 }
 
+
+}  // namespace detail
+
+
 /**
  * run uses template to go through the list and select the valid
  * template and run it.
  *
- * @tparam Base  the Base class with one template
- * @tparam T  the type of input object waiting converted
- * @tparam Func  the validation
- * @tparam ...Args  the variadic arguments.
+ * @tparam K  the current type tried in the conversion
+ * @tparam ...Types  the other types will be tried in the conversion if K fails
+ * @tparam T  the type of input object
+ * @tparam Func  the function will run if the object can be converted to K
+ * @tparam ...Args  the additional arguments for the Func
  *
- * @note this is the end case
+ * @param obj  the input object waiting converted
+ * @param f  the function will run if obj can be converted successfully
+ * @param args  the additional arguments for the function
+ *
+ * @note  This assumes that each invocation of f with types (K, Types...)
+ *        returns the same type
+ *
+ * @return  the result of f invoked with obj cast to the first matching type
  */
-template <template <typename> class Base, typename T, typename Func,
+template <typename K, typename... Types, typename T, typename Func,
           typename... Args>
-void run(T obj, Func, Args...)
+auto run(T obj, Func&& f, Args&&... args)
 {
-    GKO_NOT_SUPPORTED(obj);
+#if __cplusplus < 201703L
+    using ReturnType = std::result_of_t<Func(K, Args...)>;
+#else
+    using ReturnType = std::invoke_result_t<Func, K, Args...>;
+#endif
+    return detail::run_impl<ReturnType, K, Types...>(
+        obj, std::forward<Func>(f), std::forward<Args>(args)...);
 }
 
+
 /**
  * run uses template to go through the list and select the valid
  * template and run it.
@@ -90,16 +148,27 @@ void run(T obj, Func, Args...)
  * @param obj  the input object waiting converted
  * @param f  the function will run if obj can be converted successfully
  * @param args  the additional arguments for the function
+ *
+ * @note  This assumes that each invocation of f with types (smart_ptr<Base<K>>,
+ *        smart_ptr<Base<Types>>...) returns the same type
+ *
+ * @return  the result of f invoked with obj cast to the first matching type
  */
 template <template <typename> class Base, typename K, typename... Types,
-          typename T, typename func, typename... Args>
-void run(T obj, func f, Args... args)
+          typename T, typename Func, typename... Args>
+auto run(T obj, Func&& f, Args&&... args)
 {
-    if (auto dobj = std::dynamic_pointer_cast<const Base<K>>(obj)) {
-        f(dobj, args...);
-    } else {
-        run<Base, Types...>(obj, f, args...);
-    }
+    // Since T is a smart pointer, the type used to invoke f also has to be a
+    // smart pointer. unique_ptr is used because it can be converted into a
+    // shared_ptr, but not the other way around.
+#if __cplusplus < 201703L
+    using ReturnType = std::result_of<Func(std::unique_ptr<Base<K>>, Args...)>;
+#else
+    using ReturnType =
+        std::invoke_result_t<Func, std::unique_ptr<Base<K>>, Args...>;
+#endif
+    return detail::run_impl<ReturnType, Base, K, Types...>(
+        obj, std::forward<Func>(f), std::forward<Args>(args)...);
 }
 
 

From 6bc95881f3b3a4547abbba99fadebc94f2c99d10 Mon Sep 17 00:00:00 2001
From: Marcel Koch <marcel.koch@kit.edu>
Date: Mon, 16 Oct 2023 09:47:06 +0200
Subject: [PATCH 3/4] adds block operator

Co-authored-by: Yu-Hsiang M. Tsai <yhmtsai@gmail.com>
---
 core/CMakeLists.txt                         |   1 +
 core/base/block_operator.cpp                | 302 ++++++++++++++++++++
 include/ginkgo/core/base/block_operator.hpp | 181 ++++++++++++
 include/ginkgo/ginkgo.hpp                   |   1 +
 4 files changed, 485 insertions(+)
 create mode 100644 core/base/block_operator.cpp
 create mode 100644 include/ginkgo/core/base/block_operator.hpp

diff --git a/core/CMakeLists.txt b/core/CMakeLists.txt
index 21f188ccc5f..047d362e8dd 100644
--- a/core/CMakeLists.txt
+++ b/core/CMakeLists.txt
@@ -5,6 +5,7 @@ target_sources(ginkgo
     PRIVATE
     base/array.cpp
     base/batch_multi_vector.cpp
+    base/block_operator.cpp
     base/combination.cpp
     base/composition.cpp
     base/dense_cache.cpp
diff --git a/core/base/block_operator.cpp b/core/base/block_operator.cpp
new file mode 100644
index 00000000000..1efc5e3b152
--- /dev/null
+++ b/core/base/block_operator.cpp
@@ -0,0 +1,302 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <ginkgo/core/base/block_operator.hpp>
+
+
+#include <utility>
+
+
+#include <ginkgo/core/base/precision_dispatch.hpp>
+#include <ginkgo/core/matrix/dense.hpp>
+
+
+#include "core/base/dispatch_helper.hpp"
+
+
+namespace gko {
+namespace {
+
+
+template <typename Fn>
+auto dispatch_dense(Fn&& fn, LinOp* v)
+{
+    return run<matrix::Dense<float>*, matrix::Dense<double>*,
+               matrix::Dense<std::complex<float>>*,
+               matrix::Dense<std::complex<double>>*>(v, std::forward<Fn>(fn));
+}
+
+
+template <typename LinOpType>
+auto create_vector_blocks(LinOpType* vector,
+                          const std::vector<detail::value_span>& spans)
+{
+    return [=](size_type i) {
+        return dispatch_dense(
+            [&](auto* dense) -> std::unique_ptr<LinOpType> {
+                GKO_ENSURE_IN_BOUNDS(i, spans.size());
+                return dense->create_submatrix(spans[i],
+                                               {0, dense->get_size()[1]});
+            },
+            const_cast<LinOp*>(vector));
+    };
+}
+
+
+const LinOp* find_non_zero_in_row(
+    const std::vector<std::vector<std::shared_ptr<const LinOp>>>& blocks,
+    size_type row)
+{
+    auto it = std::find_if(blocks[row].begin(), blocks[row].end(),
+                           [](const auto& b) { return b.get() != nullptr; });
+    GKO_THROW_IF_INVALID(it != blocks[row].end(),
+                         "Encountered row with only nullptrs.");
+    return it->get();
+}
+
+
+const LinOp* find_non_zero_in_col(
+    const std::vector<std::vector<std::shared_ptr<const LinOp>>>& blocks,
+    size_type col)
+{
+    auto it = std::find_if(blocks.begin(), blocks.end(), [col](const auto& b) {
+        return b[col].get() != nullptr;
+    });
+    GKO_THROW_IF_INVALID(it != blocks.end(),
+                         "Encountered columns with only nullptrs.");
+    return it->at(col).get();
+}
+
+
+void validate_blocks(
+    const std::vector<std::vector<std::shared_ptr<const LinOp>>>& blocks)
+{
+    GKO_THROW_IF_INVALID(blocks.empty() || !blocks.front().empty(),
+                         "Blocks must either be empty, or a 2D std::vector.");
+    // all rows have same number of columns
+    for (size_type row = 1; row < blocks.size(); ++row) {
+        GKO_ASSERT_EQ(blocks[row].size(), blocks.front().size());
+    }
+    // within each row and each column the blocks have the same number of rows
+    // and columns respectively
+    for (size_type row = 0; row < blocks.size(); ++row) {
+        auto non_zero_row = find_non_zero_in_row(blocks, row);
+        for (size_type col = 0; col < blocks.front().size(); ++col) {
+            auto non_zero_col = find_non_zero_in_col(blocks, col);
+            if (blocks[row][col]) {
+                GKO_ASSERT_EQUAL_COLS(blocks[row][col], non_zero_col);
+                GKO_ASSERT_EQUAL_ROWS(blocks[row][col], non_zero_row);
+            }
+        }
+    }
+}
+
+
+template <typename Fn>
+std::vector<detail::value_span> compute_local_spans(
+    size_type num_blocks,
+    const std::vector<std::vector<std::shared_ptr<const LinOp>>>& blocks,
+    Fn&& get_size)
+{
+    validate_blocks(blocks);
+    std::vector<detail::value_span> local_spans;
+    size_type offset = 0;
+    for (size_type i = 0; i < num_blocks; ++i) {
+        auto local_size = get_size(i);
+        local_spans.emplace_back(offset, offset + local_size);
+        offset += local_size;
+    }
+    return local_spans;
+}
+
+
+dim<2> compute_global_size(
+    const std::vector<std::vector<std::shared_ptr<const LinOp>>>& blocks)
+{
+    validate_blocks(blocks);
+    if (blocks.empty()) {
+        return {};
+    }
+    size_type num_rows = 0;
+    for (size_type row = 0; row < blocks.size(); ++row) {
+        num_rows += find_non_zero_in_row(blocks, row)->get_size()[0];
+    }
+    size_type num_cols = 0;
+    for (size_type col = 0; col < blocks.front().size(); ++col) {
+        num_cols += find_non_zero_in_col(blocks, col)->get_size()[1];
+    }
+    return {num_rows, num_cols};
+}
+
+
+}  // namespace
+
+
+std::unique_ptr<BlockOperator> BlockOperator::create(
+    std::shared_ptr<const Executor> exec)
+{
+    return std::unique_ptr<BlockOperator>(new BlockOperator(std::move(exec)));
+}
+
+
+std::unique_ptr<BlockOperator> BlockOperator::create(
+    std::shared_ptr<const Executor> exec,
+    std::vector<std::vector<std::shared_ptr<const LinOp>>> blocks)
+{
+    return std::unique_ptr<BlockOperator>(
+        new BlockOperator(std::move(exec), std::move(blocks)));
+}
+
+
+BlockOperator::BlockOperator(std::shared_ptr<const Executor> exec)
+    : EnableLinOp<BlockOperator>(std::move(exec))
+{}
+
+
+BlockOperator::BlockOperator(
+    std::shared_ptr<const Executor> exec,
+    std::vector<std::vector<std::shared_ptr<const LinOp>>> blocks)
+    : EnableLinOp<BlockOperator>(exec, compute_global_size(blocks)),
+      block_size_(blocks.empty()
+                      ? dim<2>{}
+                      : dim<2>(blocks.size(), blocks.front().size())),
+      row_spans_(compute_local_spans(
+          block_size_[0], blocks,
+          [&](auto i) {
+              return find_non_zero_in_row(blocks, i)->get_size()[0];
+          })),
+      col_spans_(compute_local_spans(block_size_[1], blocks, [&](auto i) {
+          return find_non_zero_in_col(blocks, i)->get_size()[1];
+      }))
+{
+    for (auto& row : blocks) {
+        for (auto& block : row) {
+            if (block && block->get_executor() != exec) {
+                blocks_.push_back(gko::clone(exec, block));
+            } else {
+                blocks_.push_back(std::move(block));
+            }
+        }
+    }
+}
+
+
+void init_one_cache(std::shared_ptr<const Executor> exec,
+                    const detail::DenseCache<default_precision>& one_cache)
+{
+    if (one_cache.get() == nullptr) {
+        one_cache.init(std::move(exec), {1, 1});
+        one_cache->fill(one<default_precision>());
+    }
+}
+
+
+void BlockOperator::apply_impl(const LinOp* b, LinOp* x) const
+{
+    auto block_b = create_vector_blocks(b, col_spans_);
+    auto block_x = create_vector_blocks(x, row_spans_);
+
+    init_one_cache(this->get_executor(), one_);
+    for (size_type row = 0; row < block_size_[0]; ++row) {
+        bool first_in_row = true;
+        for (size_type col = 0; col < block_size_[1]; ++col) {
+            if (!block_at(row, col)) {
+                continue;
+            }
+            if (first_in_row) {
+                block_at(row, col)->apply(block_b(col), block_x(row));
+                first_in_row = false;
+            } else {
+                block_at(row, col)->apply(one_.get(), block_b(col), one_.get(),
+                                          block_x(row));
+            }
+        }
+    }
+}
+
+
+void BlockOperator::apply_impl(const LinOp* alpha, const LinOp* b,
+                               const LinOp* beta, LinOp* x) const
+{
+    auto block_b = create_vector_blocks(b, col_spans_);
+    auto block_x = create_vector_blocks(x, row_spans_);
+
+    init_one_cache(this->get_executor(), one_);
+    for (size_type row = 0; row < block_size_[0]; ++row) {
+        bool first_in_row = true;
+        for (size_type col = 0; col < block_size_[1]; ++col) {
+            if (!block_at(row, col)) {
+                continue;
+            }
+            if (first_in_row) {
+                block_at(row, col)->apply(alpha, block_b(col), beta,
+                                          block_x(row));
+                first_in_row = false;
+            } else {
+                block_at(row, col)->apply(alpha, block_b(col), one_.get(),
+                                          block_x(row));
+            }
+        }
+    }
+}
+
+
+BlockOperator::BlockOperator(const BlockOperator& other)
+    : EnableLinOp<BlockOperator>(other.get_executor())
+{
+    *this = other;
+}
+
+
+BlockOperator::BlockOperator(BlockOperator&& other) noexcept
+    : EnableLinOp<BlockOperator>(other.get_executor())
+{
+    *this = std::move(other);
+}
+
+
+BlockOperator& BlockOperator::operator=(const BlockOperator& other)
+{
+    if (this != &other) {
+        auto exec = this->get_executor();
+
+        set_size(other.get_size());
+        block_size_ = other.get_block_size();
+        col_spans_ = other.col_spans_;
+        row_spans_ = other.row_spans_;
+        blocks_.clear();
+        for (const auto& block : other.blocks_) {
+            blocks_.emplace_back(block == nullptr ? nullptr
+                                                  : gko::clone(exec, block));
+        }
+    }
+    return *this;
+}
+
+
+BlockOperator& BlockOperator::operator=(BlockOperator&& other)
+{
+    if (this != &other) {
+        auto exec = this->get_executor();
+
+        set_size(other.get_size());
+        other.set_size({});
+
+        block_size_ = std::exchange(other.block_size_, dim<2>{});
+        col_spans_ = std::move(other.col_spans_);
+        row_spans_ = std::move(other.row_spans_);
+        blocks_ = std::move(other.blocks_);
+        if (exec != other.get_executor()) {
+            for (auto& block : blocks_) {
+                if (block != nullptr) {
+                    block = gko::clone(exec, block);
+                }
+            }
+        }
+    }
+    return *this;
+}
+
+
+}  // namespace gko
diff --git a/include/ginkgo/core/base/block_operator.hpp b/include/ginkgo/core/base/block_operator.hpp
new file mode 100644
index 00000000000..c6932d45aaf
--- /dev/null
+++ b/include/ginkgo/core/base/block_operator.hpp
@@ -0,0 +1,181 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#ifndef GINKGO_BLOCK_MATRIX_HPP
+#define GINKGO_BLOCK_MATRIX_HPP
+
+#include <ginkgo/config.hpp>
+#include <ginkgo/core/base/dense_cache.hpp>
+#include <ginkgo/core/base/lin_op.hpp>
+#include <ginkgo/core/base/polymorphic_object.hpp>
+
+
+namespace gko {
+namespace detail {
+
+
+/**
+ * Carbon copy og gko::span, except that the members are not const.
+ * Necessary since the normal gko::span is not copy/movable.
+ */
+struct value_span {
+    constexpr value_span(size_type point) noexcept
+        : value_span{point, point + 1}
+    {}
+
+    constexpr value_span(size_type begin, size_type end) noexcept
+        : begin{begin}, end{end}
+    {}
+
+    constexpr operator span() const { return {begin, end}; }
+
+    constexpr value_span(const span& s) noexcept : value_span(s.begin, s.end) {}
+
+    constexpr bool is_valid() const { return begin <= end; }
+
+    constexpr size_type length() const { return end - begin; }
+
+    size_type begin;
+    size_type end;
+};
+
+
+}  // namespace detail
+
+
+/**
+ * A BlockOperator represents a linear operator that is partitioned into
+ * multiple blocks.
+ *
+ * For example, a BlockOperator can be used to define the operator:
+ * ```
+ * | A B |
+ * | C D |
+ * ```
+ * where `A, B, C, D` itself are matrices of compatible size. This can be
+ * created with:
+ * ```c++
+ * std::shared_ptr<LinOp> A = ...;
+ * std::shared_ptr<LinOp> B = ...;
+ * std::shared_ptr<LinOp> C = ...;
+ * std::shared_ptr<LinOp> D = ...;
+ * auto bop = BlockOperator::create(exec, {{A, B}, {C, D}});
+ * ```
+ * The requirements on the individual blocks passed to the create method are:
+ * - In each block-row, all blocks have the same number of rows
+ * - In each block-column, all blocks have the same number of columns
+ * - Each block-row must have the same number of blocks
+ * It is possible to set blocks to zero, by passing in a nullptr. But every
+ * block-row and block-column must contain at least one non-nullptr block.
+ *
+ * The constructor will store all passed in blocks on the same executor as the
+ * BlockOperator, which will requires copying any block that is associated with
+ * a different executor.
+ */
+class BlockOperator final : public EnableLinOp<BlockOperator> {
+    friend class EnablePolymorphicObject<BlockOperator, LinOp>;
+
+public:
+    /**
+     * Get the block dimension of this, i.e. the number of blocks per row and
+     * column.
+     *
+     * @return  The block dimension of this.
+     */
+    dim<2> get_block_size() const { return block_size_; }
+
+    /**
+     * Const access to a specific block.
+     *
+     * @param i  block row.
+     * @param j  block column.
+     *
+     * @return  the block stored at (i, j).
+     */
+    const LinOp* block_at(size_type i, size_type j) const
+    {
+        GKO_ENSURE_IN_DIMENSION_BOUNDS(i, j, block_size_);
+        return blocks_[i * block_size_[1] + j].get();
+    }
+
+    /**
+     * Create empty BlockOperator.
+     *
+     * @param exec  the executor of this.
+     *
+     * @return  empty BlockOperator.
+     */
+    static std::unique_ptr<BlockOperator> create(
+        std::shared_ptr<const Executor> exec);
+
+    /**
+     * Create BlockOperator from the given blocks.
+     *
+     * @param exec  the executor of this.
+     * @param blocks  the blocks of this operator. The blocks will be used in a
+     *                row-major form.
+     *
+     * @return  BlockOperator with the given blocks.
+     */
+    static std::unique_ptr<BlockOperator> create(
+        std::shared_ptr<const Executor> exec,
+        std::vector<std::vector<std::shared_ptr<const LinOp>>> blocks);
+
+    /**
+     * Copy constructs a BlockOperator. The executor of other is used for this.
+     * The blocks of other are deep-copied into this, using clone.
+     */
+    BlockOperator(const BlockOperator& other);
+
+    /**
+     * Move constructs a BlockOperator. The executor of other is used for this.
+     * All remaining data of other is moved into this. After this operation,
+     * other will be empty.
+     */
+    BlockOperator(BlockOperator&& other) noexcept;
+
+    /**
+     * Copy assigns a BlockOperator. The executor of this is not modified.
+     * The blocks of other are deep-copied into this, using clone.
+     */
+    BlockOperator& operator=(const BlockOperator& other);
+
+    /**
+     * Move assigns a BlockOperator. The executor of this is not modified.
+     * All data of other (except its executor) is moved into this. If the
+     * executor of this and other differ, the blocks will be copied to the
+     * executor of this. After this operation, other will be empty.
+     */
+    BlockOperator& operator=(BlockOperator&& other);
+
+private:
+    explicit BlockOperator(std::shared_ptr<const Executor> exec);
+
+    BlockOperator(
+        std::shared_ptr<const Executor> exec,
+        std::vector<std::vector<std::shared_ptr<const LinOp>>> blocks);
+
+    void apply_impl(const LinOp* b, LinOp* x) const override;
+
+    void apply_impl(const LinOp* alpha, const LinOp* b, const LinOp* beta,
+                    LinOp* x) const override;
+
+    dim<2> block_size_;
+    std::vector<detail::value_span> row_spans_;
+    std::vector<detail::value_span> col_spans_;
+    std::vector<std::shared_ptr<const LinOp>> blocks_;
+    /**
+     * @internal Using a fixed precision here may lead to temporary
+     *           conversions, since there is no value_type information
+     *           of the stored block available.
+     *
+     * @todo fix when better value_type information is available
+     */
+    detail::DenseCache<default_precision> one_;
+};
+
+
+}  // namespace gko
+
+#endif  // GINKGO_BLOCK_MATRIX_HPP
diff --git a/include/ginkgo/ginkgo.hpp b/include/ginkgo/ginkgo.hpp
index 048b188a6d4..57add89ac08 100644
--- a/include/ginkgo/ginkgo.hpp
+++ b/include/ginkgo/ginkgo.hpp
@@ -14,6 +14,7 @@
 #include <ginkgo/core/base/batch_dim.hpp>
 #include <ginkgo/core/base/batch_lin_op.hpp>
 #include <ginkgo/core/base/batch_multi_vector.hpp>
+#include <ginkgo/core/base/block_operator.hpp>
 #include <ginkgo/core/base/combination.hpp>
 #include <ginkgo/core/base/composition.hpp>
 #include <ginkgo/core/base/dense_cache.hpp>

From d54f5131ac1c05ee76fa083c279bb7c98d705e1c Mon Sep 17 00:00:00 2001
From: Marcel Koch <marcel.koch@kit.edu>
Date: Mon, 16 Oct 2023 10:29:59 +0200
Subject: [PATCH 4/4] adds block operator tests

Co-authored-by: Yu-Hsiang M. Tsai <yhmtsai@gmail.com>
---
 core/test/base/CMakeLists.txt     |   1 +
 core/test/base/block_operator.cpp | 336 ++++++++++++++++++++++++++++++
 2 files changed, 337 insertions(+)
 create mode 100644 core/test/base/block_operator.cpp

diff --git a/core/test/base/CMakeLists.txt b/core/test/base/CMakeLists.txt
index 4e611852be5..d2630ac65bc 100644
--- a/core/test/base/CMakeLists.txt
+++ b/core/test/base/CMakeLists.txt
@@ -4,6 +4,7 @@ ginkgo_create_test(array)
 ginkgo_create_test(batch_dim)
 ginkgo_create_test(batch_lin_op)
 ginkgo_create_test(batch_multi_vector)
+ginkgo_create_test(block_operator)
 ginkgo_create_test(combination)
 ginkgo_create_test(composition)
 ginkgo_create_test(deferred_factory)
diff --git a/core/test/base/block_operator.cpp b/core/test/base/block_operator.cpp
new file mode 100644
index 00000000000..54ac3649d9e
--- /dev/null
+++ b/core/test/base/block_operator.cpp
@@ -0,0 +1,336 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <ginkgo/core/base/block_operator.hpp>
+
+
+#include <vector>
+
+
+#include <gtest/gtest.h>
+
+
+#include <ginkgo/core/matrix/dense.hpp>
+
+
+#include "core/test/utils.hpp"
+
+
+struct DummyOperator : public gko::EnableLinOp<DummyOperator> {
+    explicit DummyOperator(std::shared_ptr<const gko::Executor> exec)
+        : DummyOperator(std::move(exec), {3, 3})
+    {}
+
+    explicit DummyOperator(std::shared_ptr<const gko::Executor> exec,
+                           gko::dim<2> size)
+        : gko::EnableLinOp<DummyOperator>(std::move(exec), size)
+    {}
+
+    void apply_impl(const LinOp* b, LinOp* x) const override {}
+
+    void apply_impl(const LinOp* alpha, const LinOp* b, const LinOp* beta,
+                    LinOp* x) const override
+    {}
+};
+
+
+class BlockOperator : public ::testing::Test {
+protected:
+    std::shared_ptr<const gko::Executor> exec =
+        gko::ReferenceExecutor::create();
+    std::vector<std::vector<std::shared_ptr<const gko::LinOp>>> operators = {
+        {std::make_shared<DummyOperator>(exec),
+         std::make_shared<DummyOperator>(exec)},
+        {std::make_shared<DummyOperator>(exec),
+         std::make_shared<DummyOperator>(exec)}};
+    std::default_random_engine engine;
+};
+
+
+template <typename Fn>
+void enumerate_for_all_blocks(gko::ptr_param<const gko::BlockOperator> bop,
+                              Fn&& fn)
+{
+    auto size = bop->get_block_size();
+    for (gko::size_type row = 0; row < size[0]; ++row) {
+        for (gko::size_type col = 0; col < size[1]; ++col) {
+            fn(row, col);
+        }
+    }
+}
+
+
+template <typename Fn>
+void for_all_blocks(gko::ptr_param<const gko::BlockOperator> bop, Fn&& fn)
+{
+    enumerate_for_all_blocks(
+        bop, [&](auto row, auto col) { fn(bop->block_at(row, col)); });
+}
+
+
+TEST_F(BlockOperator, CanBeEmpty)
+{
+    auto bop = gko::BlockOperator::create(exec);
+
+    ASSERT_EQ(bop->get_executor(), exec);
+    ASSERT_FALSE(bop->get_size());
+    ASSERT_FALSE(bop->get_block_size());
+}
+
+
+TEST_F(BlockOperator, CanBeConstructedFromVectors)
+{
+    auto bop = gko::BlockOperator::create(exec, operators);
+
+    gko::dim<2> block_size{2u, 2u};
+    gko::dim<2> global_size{6u, 6u};
+    GKO_ASSERT_EQUAL_DIMENSIONS(bop->get_block_size(), block_size);
+    GKO_ASSERT_EQUAL_DIMENSIONS(bop->get_size(), global_size);
+    for_all_blocks(bop, [](auto block) { ASSERT_TRUE(block); });
+}
+
+
+TEST_F(BlockOperator, CanBeConstructedFromInitializerList)
+{
+    auto bop = gko::BlockOperator::create(
+        exec, {{std::make_shared<DummyOperator>(exec),
+                std::make_shared<DummyOperator>(exec)},
+               {std::make_shared<DummyOperator>(exec),
+                std::make_shared<DummyOperator>(exec)}});
+
+    gko::dim<2> block_size{2u, 2u};
+    gko::dim<2> global_size{6u, 6u};
+    GKO_ASSERT_EQUAL_DIMENSIONS(bop->get_block_size(), block_size);
+    GKO_ASSERT_EQUAL_DIMENSIONS(bop->get_size(), global_size);
+    for_all_blocks(bop, [](auto block) { ASSERT_TRUE(block); });
+}
+
+
+TEST_F(BlockOperator, CanBeConstructedWithNull)
+{
+    auto bop = gko::BlockOperator::create(
+        exec, {{std::make_shared<DummyOperator>(exec), nullptr},
+               {nullptr, std::make_shared<DummyOperator>(exec)}});
+
+    gko::dim<2> block_size{2u, 2u};
+    gko::dim<2> global_size{6u, 6u};
+    GKO_ASSERT_EQUAL_DIMENSIONS(bop->get_block_size(), block_size);
+    GKO_ASSERT_EQUAL_DIMENSIONS(bop->get_size(), global_size);
+    ASSERT_TRUE(bop->block_at(0, 0));
+    ASSERT_TRUE(bop->block_at(1, 1));
+    ASSERT_FALSE(bop->block_at(0, 1));
+    ASSERT_FALSE(bop->block_at(1, 0));
+}
+
+
+TEST_F(BlockOperator, ThrowsOnUnequalNumberOfColumns)
+{
+    // I need to pull out the temporary because NVHPC can't handle that.
+    std::vector<std::vector<std::shared_ptr<const gko::LinOp>>> ops{
+        {std::make_shared<DummyOperator>(exec),
+         std::make_shared<DummyOperator>(exec)},
+        {std::make_shared<DummyOperator>(exec)}};
+    ASSERT_THROW(gko::BlockOperator::create(exec, ops), gko::ValueMismatch);
+}
+
+
+TEST_F(BlockOperator, ThrowsOnEmptyRow)
+{
+    std::vector<std::vector<std::shared_ptr<const gko::LinOp>>> ops{
+        {std::make_shared<DummyOperator>(exec),
+         std::make_shared<DummyOperator>(exec)},
+        {nullptr, nullptr}};
+    ASSERT_THROW(gko::BlockOperator::create(exec, ops), gko::InvalidStateError);
+}
+
+
+TEST_F(BlockOperator, ThrowsOnEmptyColumn)
+{
+    std::vector<std::vector<std::shared_ptr<const gko::LinOp>>> ops{
+        {std::make_shared<DummyOperator>(exec), nullptr},
+        {std::make_shared<DummyOperator>(exec), nullptr}};
+    ASSERT_THROW(gko::BlockOperator::create(exec, ops), gko::InvalidStateError);
+}
+
+
+TEST_F(BlockOperator, ThrowsOnNonMatchingRowsWithinBlocks)
+{
+    std::vector<std::vector<std::shared_ptr<const gko::LinOp>>> ops{
+        {nullptr, std::make_shared<DummyOperator>(exec, gko::dim<2>{2, 2}),
+         std::make_shared<DummyOperator>(exec, gko::dim<2>{1, 2})},
+        {std::make_shared<DummyOperator>(exec, gko::dim<2>{2, 2}), nullptr,
+         std::make_shared<DummyOperator>(exec, gko::dim<2>{2, 2})}};
+    ASSERT_THROW(gko::BlockOperator::create(exec, ops), gko::DimensionMismatch);
+}
+
+
+TEST_F(BlockOperator, ThrowsOnNonMatchingColsWithinBlocks)
+{
+    std::vector<std::vector<std::shared_ptr<const gko::LinOp>>> ops{
+        {nullptr, std::make_shared<DummyOperator>(exec, gko::dim<2>{2, 2}),
+         std::make_shared<DummyOperator>(exec, gko::dim<2>{2, 1})},
+        {std::make_shared<DummyOperator>(exec, gko::dim<2>{2, 2}), nullptr,
+         std::make_shared<DummyOperator>(exec, gko::dim<2>{2, 2})}};
+    ASSERT_THROW(gko::BlockOperator::create(exec, ops), gko::DimensionMismatch);
+}
+
+
+TEST_F(BlockOperator, ThrowsOnOutOfBoundsBlockAccess)
+{
+    auto bop = gko::BlockOperator::create(
+        exec, {{std::make_shared<DummyOperator>(exec, gko::dim<2>{8, 8}),
+                std::make_shared<DummyOperator>(exec, gko::dim<2>{8, 8})},
+               {std::make_shared<DummyOperator>(exec, gko::dim<2>{8, 8}),
+                std::make_shared<DummyOperator>(exec, gko::dim<2>{8, 8})}});
+
+    ASSERT_THROW(bop->block_at(5, 5), gko::OutOfBoundsError);
+}
+
+
+TEST_F(BlockOperator, CanBeCopied)
+{
+    using Mtx = gko::matrix::Dense<>;
+    auto bop = gko::BlockOperator::create(
+        exec, {{gko::initialize<Mtx>({{1, 2}, {2, 1}}, exec), nullptr},
+               {nullptr, gko::initialize<Mtx>({{3, 4}, {4, 3}}, exec)}});
+    auto bop_copy = gko::BlockOperator::create(exec);
+
+    bop_copy->copy_from(bop);
+
+    GKO_ASSERT_EQUAL_DIMENSIONS(bop, bop_copy);
+    GKO_ASSERT_EQUAL_DIMENSIONS(bop->get_block_size(),
+                                bop_copy->get_block_size());
+    enumerate_for_all_blocks(bop, [&](auto row, auto col) {
+        if (bop->block_at(row, col) == nullptr) {
+            ASSERT_TRUE(bop_copy->block_at(row, col) == nullptr);
+        } else {
+            ASSERT_NE(bop->block_at(row, col), bop_copy->block_at(row, col));
+            GKO_ASSERT_MTX_NEAR(gko::as<Mtx>(bop->block_at(row, col)),
+                                gko::as<Mtx>(bop_copy->block_at(row, col)),
+                                0.0);
+        }
+    });
+}
+
+
+TEST_F(BlockOperator, CanBeMoved)
+{
+    using Mtx = gko::matrix::Dense<>;
+    auto bop = gko::BlockOperator::create(
+        exec, {{gko::initialize<Mtx>({{1, 2}, {2, 1}}, exec), nullptr},
+               {nullptr, gko::initialize<Mtx>({{3, 4}, {4, 3}}, exec)}});
+    auto bop_copy = gko::clone(bop);
+    auto bop_move = gko::BlockOperator::create(exec);
+
+    bop_move->move_from(bop);
+
+    GKO_ASSERT_EQUAL_DIMENSIONS(bop_copy, bop_move);
+    GKO_ASSERT_EQUAL_DIMENSIONS(bop_copy->get_block_size(),
+                                bop_move->get_block_size());
+    GKO_ASSERT_EQUAL_DIMENSIONS(bop, gko::dim<2>{});
+    GKO_ASSERT_EQUAL_DIMENSIONS(bop->get_block_size(), gko::dim<2>{});
+    enumerate_for_all_blocks(bop_copy, [&](auto row, auto col) {
+        if (bop_copy->block_at(row, col) == nullptr) {
+            ASSERT_TRUE(bop_move->block_at(row, col) == nullptr);
+        } else {
+            GKO_ASSERT_MTX_NEAR(gko::as<Mtx>(bop_copy->block_at(row, col)),
+                                gko::as<Mtx>(bop_move->block_at(row, col)),
+                                0.0);
+        }
+    });
+}
+
+
+TEST_F(BlockOperator, CanApply)
+{
+    using vtype = double;
+    using Mtx = gko::matrix::Dense<vtype>;
+    auto bop = gko::BlockOperator::create(
+        exec, {{gko::initialize<Mtx>({{1, 2}, {2, 1}}, exec),
+                gko::initialize<Mtx>({{5, 6}, {6, 5}}, exec)},
+               {nullptr, gko::initialize<Mtx>({{3, 4}, {4, 3}}, exec)}});
+    auto x = gko::initialize<Mtx>({{1, 10}, {2, 20}, {3, 30}, {4, 40}}, exec);
+    auto y = Mtx::create_with_config_of(x);
+
+    bop->apply(x, y);
+
+    GKO_ASSERT_MTX_NEAR(
+        y, l(I<I<vtype>>{{44, 440}, {42, 420}, {25, 250}, {24, 240}}),
+        r<vtype>::value);
+}
+
+
+TEST_F(BlockOperator, CanAdvancedApply)
+{
+    using vtype = double;
+    using Mtx = gko::matrix::Dense<vtype>;
+    auto bop = gko::BlockOperator::create(
+        exec, {{gko::initialize<Mtx>({{1, 2}, {2, 1}}, exec),
+                gko::initialize<Mtx>({{5, 6}, {6, 5}}, exec)},
+               {nullptr, gko::initialize<Mtx>({{3, 4}, {4, 3}}, exec)}});
+    auto x = gko::initialize<Mtx>({1, 2, 3, 4}, exec);
+    auto y = gko::initialize<Mtx>({-4, -3, -2, -1}, exec);
+    auto alpha = gko::initialize<Mtx>({0.5}, exec);
+    auto beta = gko::initialize<Mtx>({-1}, exec);
+
+    bop->apply(alpha, x, beta, y);
+
+    GKO_ASSERT_MTX_NEAR(y, l(I<vtype>{26, 24, 14.5, 13}), r<vtype>::value);
+}
+
+
+TEST_F(BlockOperator, CanApplyAndAdvancedApplyLarge)
+{
+    using vtype = double;
+    using Mtx = gko::matrix::Dense<vtype>;
+    gko::size_type local_num_rows = 4;
+    gko::size_type local_num_cols = 6;
+    gko::size_type block_num_rows = 8;
+    gko::size_type block_num_cols = 5;
+    auto dense = gko::test::generate_random_dense_matrix<vtype>(
+        block_num_rows * local_num_rows, block_num_cols * local_num_cols,
+        std::uniform_real_distribution<vtype>(-1, 1), engine, exec);
+    auto get_submatrix = [&](auto i, auto j) {
+        return dense->create_submatrix(
+            {i * local_num_rows, (i + 1) * local_num_rows},
+            {j * local_num_cols, (j + 1) * local_num_cols});
+    };
+    std::vector<std::vector<std::shared_ptr<const gko::LinOp>>> blocks(
+        block_num_rows,
+        std::vector<std::shared_ptr<const gko::LinOp>>(block_num_cols));
+    for (gko::size_type i = 0; i < block_num_rows; ++i) {
+        for (gko::size_type j = 0; j < block_num_cols; ++j) {
+            blocks[i][j] = get_submatrix(i, j);
+        }
+    }
+    auto bop = gko::BlockOperator::create(exec, blocks);
+    auto x = gko::test::generate_random_dense_matrix<vtype>(
+        block_num_cols * local_num_cols, 3,
+        std::uniform_real_distribution<vtype>(-1, 1), engine, exec);
+    auto y = Mtx::create(exec, gko::dim<2>{block_num_rows * local_num_rows, 3});
+    auto result_y = gko::clone(y);
+
+    {
+        SCOPED_TRACE("Apply");
+        y->fill(1.0);
+        result_y->fill(1.0);
+
+        dense->apply(x, result_y);
+        bop->apply(x, y);
+
+        GKO_ASSERT_MTX_NEAR(y, result_y, r<vtype>::value);
+    }
+    {
+        SCOPED_TRACE("Advanced Apply");
+        auto alpha = gko::initialize<Mtx>({0.5}, exec);
+        auto beta = gko::initialize<Mtx>({-1}, exec);
+        y->fill(1.0);
+        result_y->fill(1.0);
+
+        dense->apply(alpha, x, beta, result_y);
+        bop->apply(alpha, x, beta, y);
+
+        GKO_ASSERT_MTX_NEAR(y, result_y, r<vtype>::value);
+    }
+}