Add vector mean computation

common/unified/matrix/dense_kernels.instantiate.cpp

@@ -99,6 +99,8 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(
 // split
 GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(
     GKO_DECLARE_DENSE_COMPUTE_SQUARED_NORM2_KERNEL);
+// split
+GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_MEAN_KERNEL);
 // end
 
 

common/unified/matrix/dense_kernels.template.cpp

@@ -278,6 +278,22 @@ void compute_norm1(std::shared_ptr<const DefaultExecutor> exec,
 }
 
 
+template <typename ValueType>
+void compute_mean(std::shared_ptr<const DefaultExecutor> exec,
+                  const matrix::Dense<ValueType>* x,
+                  matrix::Dense<ValueType>* result, array<char>& tmp)
+{
+    using ValueType_nc = gko::remove_complex<ValueType>;
+    run_kernel_col_reduction_cached(
+        exec,
+        [] GKO_KERNEL(auto i, auto j, auto x, auto inv_total_size) {
+            return x(i, j) * inv_total_size;
+        },
+        GKO_KERNEL_REDUCE_SUM(ValueType), result->get_values(), x->get_size(),
+        tmp, x, ValueType_nc{1.} / x->get_size()[0]);
+}
+
+
 template <typename ValueType>
 void compute_max_nnz_per_row(std::shared_ptr<const DefaultExecutor> exec,
                              const matrix::Dense<ValueType>* source,

core/device_hooks/common_kernels.inc.cpp

@@ -331,6 +331,7 @@
 GKO_STUB_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_NORM2_KERNEL);
 GKO_STUB_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_NORM2_DISPATCH_KERNEL);
 GKO_STUB_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_NORM1_KERNEL);
+GKO_STUB_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_MEAN_KERNEL);
 GKO_STUB_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_SQUARED_NORM2_KERNEL);
 GKO_STUB_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_SQRT_KERNEL);
 GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_DENSE_FILL_IN_MATRIX_DATA_KERNEL);

core/distributed/vector.cpp

@@ -573,19 +573,62 @@
 }
 
 
+template <typename ValueType>
+void Vector<ValueType>::compute_mean(ptr_param<LinOp> result) const
+{
+    array<char> tmp{this->get_executor()};
+    this->compute_mean(result, tmp);
+}
+
+
+template <typename ValueType>
+void Vector<ValueType>::compute_mean(ptr_param<LinOp> result,
+                                     array<char>& tmp) const
+{
+    using MeanVector = local_vector_type;
+    const auto global_size = this->get_size()[0];
+    const auto local_size = this->get_local_vector()->get_size()[0];
+    const auto num_vecs = static_cast<int>(this->get_size()[1]);
+    GKO_ASSERT_EQUAL_COLS(result, this);
+    auto exec = this->get_executor();
+    const auto comm = this->get_communicator();
+    auto dense_res = make_temporary_clone(exec, as<MeanVector>(result));
+    this->get_local_vector()->compute_mean(dense_res.get());
+
+    // scale by its weight ie ratio of local to global size
+    auto weight = initialize<matrix::Dense<remove_complex<ValueType>>>(
+        {static_cast<remove_complex<ValueType>>(local_size) / global_size},
+        this->get_executor());
+    dense_res->scale(weight.get());
+
+    exec->synchronize();
+    if (mpi::requires_host_buffer(exec, comm)) {
+        host_reduction_buffer_.init(exec->get_master(), dense_res->get_size());
+        host_reduction_buffer_->copy_from(dense_res.get());
+        comm.all_reduce(exec->get_master(),
+                        host_reduction_buffer_->get_values(), num_vecs,
+                        MPI_SUM);
+        dense_res->copy_from(host_reduction_buffer_.get());
+    } else {
+        comm.all_reduce(exec, dense_res->get_values(), num_vecs, MPI_SUM);
+    }
+}
+
 template <typename ValueType>
 ValueType& Vector<ValueType>::at_local(size_type row, size_type col) noexcept
 {
     return local_.at(row, col);
 }
 
+
 template <typename ValueType>
 ValueType Vector<ValueType>::at_local(size_type row,
                                       size_type col) const noexcept
 {
     return local_.at(row, col);
 }
 
+
 template <typename ValueType>
 ValueType& Vector<ValueType>::at_local(size_type idx) noexcept
 {

core/matrix/dense.cpp

@@ -80,6 +80,7 @@
 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);
+GKO_REGISTER_OPERATION(compute_mean, dense::compute_mean);
 GKO_REGISTER_OPERATION(compute_squared_norm2, dense::compute_squared_norm2);
 GKO_REGISTER_OPERATION(compute_sqrt, dense::compute_sqrt);
 GKO_REGISTER_OPERATION(compute_max_nnz_per_row, dense::compute_max_nnz_per_row);
@@ -496,6 +497,29 @@
 }
 
 
+template <typename ValueType>
+void Dense<ValueType>::compute_mean(ptr_param<LinOp> result) const
+{
+    auto exec = this->get_executor();
+    this->compute_mean_impl(make_temporary_output_clone(exec, result).get());
+}
+
+
+template <typename ValueType>
+void Dense<ValueType>::compute_mean(ptr_param<LinOp> result,
+                                    array<char>& tmp) const
+{
+    GKO_ASSERT_EQUAL_COLS(result, this);
+    auto exec = this->get_executor();
+    if (tmp.get_executor() != exec) {
+        tmp.clear();
+        tmp.set_executor(exec);
+    }
+    auto dense_res = make_temporary_conversion<ValueType>(result);
+    exec->run(dense::make_compute_mean(this, dense_res.get(), tmp));
+}
+
+
 template <typename ValueType>
 void Dense<ValueType>::compute_squared_norm2_impl(LinOp* result) const
 {
@@ -506,6 +530,15 @@
 }
 
 
+template <typename ValueType>
+void Dense<ValueType>::compute_mean_impl(LinOp* result) const
+{
+    auto exec = this->get_executor();
+    array<char> tmp{exec};
+    this->compute_mean(make_temporary_output_clone(exec, result).get(), tmp);
+}
+
+
 template <typename ValueType>
 Dense<ValueType>& Dense<ValueType>::operator=(const Dense& other)
 {

core/matrix/dense_kernels.hpp

@@ -146,6 +146,11 @@ namespace kernels {
                        matrix::Dense<remove_complex<_type>>* result, \
                        array<char>& tmp)
 
+#define GKO_DECLARE_DENSE_COMPUTE_MEAN_KERNEL(_type)               \
+    void compute_mean(std::shared_ptr<const DefaultExecutor> exec, \
+                      const matrix::Dense<_type>* x,               \
+                      matrix::Dense<_type>* result, array<char>& tmp)
+
 #define GKO_DECLARE_DENSE_FILL_IN_MATRIX_DATA_KERNEL(_type, _prec)         \
     void fill_in_matrix_data(std::shared_ptr<const DefaultExecutor> exec,  \
                              const device_matrix_data<_type, _prec>& data, \
@@ -349,6 +354,8 @@ namespace kernels {
     GKO_DECLARE_DENSE_COMPUTE_NORM2_DISPATCH_KERNEL(ValueType);             \
     template <typename ValueType>                                           \
     GKO_DECLARE_DENSE_COMPUTE_NORM1_KERNEL(ValueType);                      \
+    template <typename ValueType>                                           \
+    GKO_DECLARE_DENSE_COMPUTE_MEAN_KERNEL(ValueType);                       \
     template <typename ValueType, typename IndexType>                       \
     GKO_DECLARE_DENSE_FILL_IN_MATRIX_DATA_KERNEL(ValueType, IndexType);     \
     template <typename ValueType>                                           \

include/ginkgo/core/distributed/vector.hpp

@@ -404,6 +404,29 @@ class Vector
      */
     void compute_norm1(ptr_param<LinOp> result, array<char>& tmp) const;
 
+    /**
+     * Computes the column-wise mean of this (multi-)vector using a global
+     * reduction.
+     *
+     * @param result  a Dense row matrix, used to store the mean
+     *                (the number of columns in result must match the number
+     *                of columns of this)
+     */
+    void compute_mean(ptr_param<LinOp> result) const;
+
+    /**
+     * Computes the column-wise arithmetic mean of this (multi-)vector using a
+     * global reduction.
+     *
+     * @param result  a Dense row matrix, used to store the mean
+     *                (the number of columns in result must match the number
+     *                of columns of this)
+     * @param tmp  the temporary storage to use for partial sums during the
+     *             reduction computation. It may be resized and/or reset to the
+     *             correct executor.
+     */
+    void compute_mean(ptr_param<LinOp> result, array<char>& tmp) const;
+
     /**
      * Returns a single element of the multi-vector.
      *

include/ginkgo/core/matrix/dense.hpp

@@ -917,6 +917,27 @@ class Dense
      */
     void compute_squared_norm2(ptr_param<LinOp> result, array<char>& tmp) const;
 
+    /**
+     * Computes the column-wise arithmetic mean of this matrix.
+     *
+     * @param result  a Dense row vector, used to store the mean
+     *                (the number of columns in the vector must match the number
+     *                of columns of this)
+     */
+    void compute_mean(ptr_param<LinOp> result) const;
+
+    /**
+     * Computes the column-wise arithmetic mean of this matrix.
+     *
+     * @param result  a Dense row vector, used to store the mean
+     *                (the number of columns in the vector must match the
+     *                number of columns of this)
+     * @param tmp  the temporary storage to use for partial sums during the
+     *             reduction computation. It may be resized and/or reset to the
+     *             correct executor.
+     */
+    void compute_mean(ptr_param<LinOp> result, array<char>& tmp) const;
+
     /**
      * Create a submatrix from the original matrix.
      * Warning: defining stride for this create_submatrix method might cause
@@ -1215,6 +1236,11 @@ class Dense
      */
     virtual void compute_squared_norm2_impl(LinOp* result) const;
 
+    /**
+     * @copydoc compute_mean(LinOp*) const
+     */
+    virtual void compute_mean_impl(LinOp* result) const;
+
     /**
      * Resizes the matrix to the given size.
      *

reference/matrix/dense_kernels.cpp

@@ -397,6 +397,27 @@ void compute_norm1(std::shared_ptr<const ReferenceExecutor> exec,
 GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_NORM1_KERNEL);
 
 
+template <typename ValueType>
+void compute_mean(std::shared_ptr<const ReferenceExecutor> exec,
+                  const matrix::Dense<ValueType>* x,
+                  matrix::Dense<ValueType>* result, array<char>&)
+{
+    using ValueType_nc = gko::remove_complex<ValueType>;
+    for (size_type j = 0; j < x->get_size()[1]; ++j) {
+        result->at(0, j) = zero<ValueType>();
+    }
+
+    for (size_type i = 0; i < x->get_size()[1]; ++i) {
+        for (size_type j = 0; j < x->get_size()[0]; ++j) {
+            result->at(0, i) += x->at(j, i);
+        }
+        result->at(0, i) /= static_cast<ValueType_nc>(x->get_size()[0]);
+    }
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_MEAN_KERNEL);
+
+
 template <typename ValueType, typename IndexType>
 void fill_in_matrix_data(std::shared_ptr<const ReferenceExecutor> exec,
                          const device_matrix_data<ValueType, IndexType>& data,

reference/test/matrix/dense_kernels.cpp

@@ -35,6 +35,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 #include <complex>
 #include <memory>
+#include <numeric>
 #include <random>
 
 
@@ -700,6 +701,37 @@ TYPED_TEST(Dense, ComputesNorm1Mixed)
 }
 
 
+TYPED_TEST(Dense, ComputesMean)
+{
+    using Mtx = typename TestFixture::Mtx;
+    using T = typename TestFixture::value_type;
+
+    auto iota = Mtx::create(this->exec, gko::dim<2>{10, 1});
+    std::iota(iota->get_values(), iota->get_values() + 10, 1);
+    auto iota_result = Mtx::create(this->exec, gko::dim<2>{1, 1});
+    iota->compute_mean(iota_result.get());
+    GKO_EXPECT_NEAR(iota_result->at(0, 0), T{5.5}, r<T>::value * 10);
+
+    auto result = Mtx::create(this->exec, gko::dim<2>{1, 3});
+
+    this->mtx4->compute_mean(result.get());
+
+    GKO_EXPECT_NEAR(result->at(0, 0), T{0.5}, r<T>::value * 10);
+    GKO_EXPECT_NEAR(result->at(0, 1), T{4.0}, r<T>::value * 10);
+    GKO_EXPECT_NEAR(result->at(0, 2), T{1.0}, r<T>::value * 10);
+}
+
+
+TYPED_TEST(Dense, ComputesMeanFailsOnWrongResultSize)
+{
+    using Mtx = typename TestFixture::Mtx;
+    using T = typename TestFixture::value_type;
+    auto result = Mtx::create(this->exec, gko::dim<2>{1, 2});
+
+    ASSERT_THROW(this->mtx4->compute_mean(result), gko::DimensionMismatch);
+}
+
+
 TYPED_TEST(Dense, ComputeDotFailsOnWrongInputSize)
 {
     using Mtx = typename TestFixture::Mtx;

test/mpi/vector.cpp

@@ -676,6 +676,30 @@ TYPED_TEST(VectorReductions, ComputeSquaredNorm2WithTmpIsSameAsDense)
 }
 
 
+TYPED_TEST(VectorReductions, ComputesMeanIsSameAsDense)
+{
+    using value_type = typename TestFixture::value_type;
+    this->init_result();
+
+    this->x->compute_mean(this->res);
+    this->dense_x->compute_mean(this->dense_res);
+
+    GKO_ASSERT_MTX_NEAR(this->res, this->dense_res, r<value_type>::value);
+}
+
+
+TYPED_TEST(VectorReductions, ComputesMeanWithTmpIsSameAsDense)
+{
+    using value_type = typename TestFixture::value_type;
+    this->init_result();
+
+    this->x->compute_mean(this->res, this->tmp);
+    this->dense_x->compute_mean(this->dense_res, this->dense_tmp);
+
+    GKO_ASSERT_MTX_NEAR(this->res, this->dense_res, r<value_type>::value);
+}
+
+
 TYPED_TEST(VectorReductions, ComputeDotCopiesToHostOnlyIfNecessary)
 {
     this->init_result();