add work estimate framework

core/CMakeLists.txt

@@ -42,6 +42,7 @@ target_sources(${ginkgo_core}
     base/segmented_array.cpp
     base/timer.cpp
     base/version.cpp
+    base/work_estimate.cpp
     components/range_minimum_query.cpp
     config/config.cpp
     config/config_helper.cpp

core/base/work_estimate.cpp

@@ -0,0 +1,68 @@
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <ginkgo/core/base/work_estimate.hpp>
+
+
+namespace gko {
+
+
+compute_bound_work_estimate operator+(compute_bound_work_estimate a,
+                                      compute_bound_work_estimate b)
+{
+    return {a.flops + b.flops};
+}
+
+
+compute_bound_work_estimate& compute_bound_work_estimate::operator+=(
+    compute_bound_work_estimate other)
+{
+    *this = *this + other;
+    return *this;
+}
+
+
+memory_bound_work_estimate operator+(memory_bound_work_estimate a,
+                                     memory_bound_work_estimate b)
+{
+    return {a.bytes_read + b.bytes_read, a.bytes_written + b.bytes_written};
+}
+
+
+memory_bound_work_estimate& memory_bound_work_estimate::operator+=(
+    memory_bound_work_estimate other)
+{
+    *this = *this + other;
+    return *this;
+}
+
+
+custom_work_estimate operator+(custom_work_estimate a, custom_work_estimate b)
+{
+    GKO_ASSERT(a.operation_count_name == b.operation_count_name);
+    return {a.operation_count_name, a.operations + b.operations};
+}
+
+
+custom_work_estimate& custom_work_estimate::operator+=(
+    custom_work_estimate other)
+{
+    *this = *this + other;
+    return *this;
+}
+
+
+kernel_work_estimate operator+(kernel_work_estimate a, kernel_work_estimate b)
+{
+    // this fails with std::bad_variant_access if the two estimates are of
+    // different types
+    return std::visit(
+        [b](auto a) -> kernel_work_estimate {
+            return a + std::get<decltype(a)>(b);
+        },
+        a);
+}
+
+
+}  // namespace gko

core/components/prefix_sum_kernels.hpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -10,6 +10,7 @@
 
 #include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/base/types.hpp>
+#include <ginkgo/core/base/work_estimate.hpp>
 
 #include "core/base/kernel_declaration.hpp"
 
@@ -53,6 +54,19 @@ GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(components,
 #undef GKO_DECLARE_ALL_AS_TEMPLATES
 
 
+namespace work_estimate::components {
+
+
+template <typename IndexType>
+kernel_work_estimate prefix_sum_nonnegative(IndexType* counts,
+                                            size_type num_entries)
+{
+    return memory_bound_work_estimate{num_entries * sizeof(IndexType),
+                                      num_entries * sizeof(IndexType)};
+}
+
+
+}  // namespace work_estimate::components
 }  // namespace kernels
 }  // namespace gko
 

core/matrix/csr.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -83,8 +83,8 @@ GKO_REGISTER_OPERATION(is_sorted_by_column_index,
                        csr::is_sorted_by_column_index);
 GKO_REGISTER_OPERATION(extract_diagonal, csr::extract_diagonal);
 GKO_REGISTER_OPERATION(fill_array, components::fill_array);
-GKO_REGISTER_OPERATION(prefix_sum_nonnegative,
-                       components::prefix_sum_nonnegative);
+GKO_REGISTER_OPERATION_WITH_WORK_ESTIMATE(prefix_sum_nonnegative,
+                                          components::prefix_sum_nonnegative);
 GKO_REGISTER_OPERATION(inplace_absolute_array,
                        components::inplace_absolute_array);
 GKO_REGISTER_OPERATION(outplace_absolute_array,

core/matrix/dense.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -40,17 +40,18 @@ namespace dense {
 namespace {
 
 
-GKO_REGISTER_OPERATION(simple_apply, dense::simple_apply);
+GKO_REGISTER_OPERATION_WITH_WORK_ESTIMATE(simple_apply, dense::simple_apply);
 GKO_REGISTER_OPERATION(apply, dense::apply);
-GKO_REGISTER_OPERATION(copy, dense::copy);
-GKO_REGISTER_OPERATION(fill, dense::fill);
+GKO_REGISTER_OPERATION_WITH_WORK_ESTIMATE(copy, dense::copy);
+GKO_REGISTER_OPERATION_WITH_WORK_ESTIMATE(fill, dense::fill);
 GKO_REGISTER_OPERATION(scale, dense::scale);
 GKO_REGISTER_OPERATION(inv_scale, dense::inv_scale);
 GKO_REGISTER_OPERATION(add_scaled, dense::add_scaled);
 GKO_REGISTER_OPERATION(sub_scaled, dense::sub_scaled);
 GKO_REGISTER_OPERATION(add_scaled_diag, dense::add_scaled_diag);
 GKO_REGISTER_OPERATION(sub_scaled_diag, dense::sub_scaled_diag);
-GKO_REGISTER_OPERATION(compute_dot, dense::compute_dot_dispatch);
+GKO_REGISTER_OPERATION_WITH_WORK_ESTIMATE(compute_dot,
+                                          dense::compute_dot_dispatch);
 GKO_REGISTER_OPERATION(compute_conj_dot, dense::compute_conj_dot_dispatch);
 GKO_REGISTER_OPERATION(compute_norm2, dense::compute_norm2_dispatch);
 GKO_REGISTER_OPERATION(compute_norm1, dense::compute_norm1);

core/matrix/dense_kernels.hpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -10,6 +10,7 @@
 
 #include <ginkgo/core/base/math.hpp>
 #include <ginkgo/core/base/types.hpp>
+#include <ginkgo/core/base/work_estimate.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
 #include <ginkgo/core/matrix/diagonal.hpp>
 
@@ -476,6 +477,53 @@ GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(dense, GKO_DECLARE_ALL_AS_TEMPLATES);
 #undef GKO_DECLARE_ALL_AS_TEMPLATES
 
 
+namespace work_estimate {
+namespace dense {
+
+
+template <typename ValueType>
+kernel_work_estimate simple_apply(const matrix::Dense<ValueType>* a,
+                                  const matrix::Dense<ValueType>* b,
+                                  matrix::Dense<ValueType>* c)
+{
+    const auto a_rows = a->get_size()[0];
+    const auto a_cols = a->get_size()[1];
+    const auto b_cols = b->get_size()[1];
+    return compute_bound_work_estimate{2 * a_rows * a_cols * b_cols};
+}
+
+
+template <typename InValueType, typename OutValueType>
+kernel_work_estimate copy(const matrix::Dense<InValueType>* input,
+                          matrix::Dense<OutValueType>* output)
+{
+    const auto memsize = input->get_size()[0] * input->get_size()[1];
+    return memory_bound_work_estimate{memsize * sizeof(InValueType),
+                                      memsize * sizeof(OutValueType)};
+}
+
+
+template <typename ValueType>
+kernel_work_estimate fill(matrix::Dense<ValueType>* mat, ValueType value)
+{
+    return memory_bound_work_estimate{
+        0, mat->get_size()[0] * mat->get_size()[1] * sizeof(ValueType)};
+}
+
+
+template <typename ValueType>
+kernel_work_estimate compute_dot_dispatch(const matrix::Dense<ValueType>* x,
+                                          const matrix::Dense<ValueType>* y,
+                                          matrix::Dense<ValueType>* result,
+                                          array<char>& tmp)
+{
+    const auto num_elements = x->get_size()[0] * x->get_size()[1];
+    return memory_bound_work_estimate{2 * num_elements * sizeof(ValueType), 0};
+}
+
+
+}  // namespace dense
+}  // namespace work_estimate
 }  // namespace kernels
 }  // namespace gko
 

core/matrix/ell.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -42,8 +42,8 @@ GKO_REGISTER_OPERATION(convert_to_csr, ell::convert_to_csr);
 GKO_REGISTER_OPERATION(count_nonzeros_per_row, ell::count_nonzeros_per_row);
 GKO_REGISTER_OPERATION(extract_diagonal, ell::extract_diagonal);
 GKO_REGISTER_OPERATION(fill_array, components::fill_array);
-GKO_REGISTER_OPERATION(prefix_sum_nonnegative,
-                       components::prefix_sum_nonnegative);
+GKO_REGISTER_OPERATION_WITH_WORK_ESTIMATE(prefix_sum_nonnegative,
+                                          components::prefix_sum_nonnegative);
 GKO_REGISTER_OPERATION(inplace_absolute_array,
                        components::inplace_absolute_array);
 GKO_REGISTER_OPERATION(outplace_absolute_array,

core/matrix/hybrid.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -42,8 +42,8 @@ GKO_REGISTER_OPERATION(compute_coo_row_ptrs, hybrid::compute_coo_row_ptrs);
 GKO_REGISTER_OPERATION(convert_idxs_to_ptrs, components::convert_idxs_to_ptrs);
 GKO_REGISTER_OPERATION(convert_to_csr, hybrid::convert_to_csr);
 GKO_REGISTER_OPERATION(fill_array, components::fill_array);
-GKO_REGISTER_OPERATION(prefix_sum_nonnegative,
-                       components::prefix_sum_nonnegative);
+GKO_REGISTER_OPERATION_WITH_WORK_ESTIMATE(prefix_sum_nonnegative,
+                                          components::prefix_sum_nonnegative);
 GKO_REGISTER_OPERATION(inplace_absolute_array,
                        components::inplace_absolute_array);
 GKO_REGISTER_OPERATION(outplace_absolute_array,

core/matrix/sellp.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -31,8 +31,8 @@ namespace {
 GKO_REGISTER_OPERATION(spmv, sellp::spmv);
 GKO_REGISTER_OPERATION(advanced_spmv, sellp::advanced_spmv);
 GKO_REGISTER_OPERATION(convert_idxs_to_ptrs, components::convert_idxs_to_ptrs);
-GKO_REGISTER_OPERATION(prefix_sum_nonnegative,
-                       components::prefix_sum_nonnegative);
+GKO_REGISTER_OPERATION_WITH_WORK_ESTIMATE(prefix_sum_nonnegative,
+                                          components::prefix_sum_nonnegative);
 GKO_REGISTER_OPERATION(compute_slice_sets, sellp::compute_slice_sets);
 GKO_REGISTER_OPERATION(fill_in_matrix_data, sellp::fill_in_matrix_data);
 GKO_REGISTER_OPERATION(fill_in_dense, sellp::fill_in_dense);