Fixup test for complex values

Add vec mean implementation

Fixup

Add a preliminary version of compute_mean

Add reference implementation of compute_mean

Format files

Add vec mean tests

common/unified/matrix/dense_kernels.template.cpp

@@ -278,6 +278,24 @@ 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 total_size) {
+            return x(i, j) / static_cast<ValueType_nc>(total_size);
+        },
+        GKO_KERNEL_REDUCE_SUM(ValueType), result->get_values(), x->get_size(),
+        tmp, x, x->get_size()[0]);
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_COMPUTE_MEAN_KERNEL);
+
+
 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_CONJ_DOT_DISPATCH_KERNEL);
 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

@@ -30,10 +30,8 @@ THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ******************************<GINKGO LICENSE>*******************************/
 
-#include <ginkgo/core/distributed/vector.hpp>
-
-
 #include <ginkgo/core/distributed/partition.hpp>
+#include <ginkgo/core/distributed/vector.hpp>
 
 
 #include "core/distributed/vector_kernels.hpp"
@@ -573,19 +571,60 @@ void Vector<ValueType>::compute_squared_norm2(ptr_param<LinOp> result,
 }
 
 
+template <typename ValueType>
+void Vector<ValueType>::compute_mean(LinOp* result) const
+{
+    array<char> tmp{this->get_executor()};
+    this->compute_mean(result, tmp);
+}
+
+
+void Vector<ValueType>::compute_mean(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_DIMENSIONS(result, dim<2>(1, num_vecs));
+    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>>>(
+        1, {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_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);
+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);
@@ -235,6 +236,14 @@ void Dense<ValueType>::compute_squared_norm2(ptr_param<LinOp> result) const
 }
 
 
+template <typename ValueType>
+void Dense<ValueType>::compute_mean(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>::inv_scale_impl(const LinOp* alpha)
 {
@@ -480,29 +489,48 @@ void Dense<ValueType>::compute_norm1_impl(LinOp* result) const
 
 
 template <typename ValueType>
+<<<<<<< HEAD
 void Dense<ValueType>::compute_squared_norm2(ptr_param<LinOp> result,
                                              array<char>& tmp) const
+=======
+void Dense<ValueType>::compute_mean(LinOp* result, array<char>& tmp) const
+>>>>>>> 52be8e43b (Fixup test for complex values)
 {
     GKO_ASSERT_EQUAL_DIMENSIONS(result, dim<2>(1, this->get_size()[1]));
     auto exec = this->get_executor();
     if (tmp.get_executor() != exec) {
         tmp.clear();
         tmp.set_executor(exec);
     }
+<<<<<<< HEAD
     auto local_result = make_temporary_clone(exec, result);
     auto dense_res = make_temporary_conversion<remove_complex<ValueType>>(
         local_result.get());
     exec->run(dense::make_compute_squared_norm2(this, dense_res.get(), tmp));
+=======
+    auto dense_res = make_temporary_conversion<ValueType>(result);
+    exec->run(dense::make_compute_mean(this, dense_res.get(), tmp));
+>>>>>>> 52be8e43b (Fixup test for complex values)
 }
 
 
 template <typename ValueType>
+<<<<<<< HEAD
 void Dense<ValueType>::compute_squared_norm2_impl(LinOp* result) const
 {
     auto exec = this->get_executor();
     array<char> tmp{exec};
     this->compute_squared_norm2(make_temporary_output_clone(exec, result).get(),
                                 tmp);
+=======
+void Dense<ValueType>::compute_mean_impl(LinOp* result) const
+{
+    GKO_ASSERT_EQUAL_DIMENSIONS(result, dim<2>(1, this->get_size()[1]));
+    auto exec = this->get_executor();
+    auto dense_res = make_temporary_conversion<ValueType>(result);
+    array<char> tmp{exec};
+    exec->run(dense::make_compute_mean(this, dense_res.get(), tmp));
+>>>>>>> 52be8e43b (Fixup test for complex values)
 }
 
 

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,32 @@ 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)
+     * @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(LinOp* result) 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)
+     * @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(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 mean of this matrix.
+     *
+     * @param result  a Dense row vector, used to store the norm
+     *                (the number of columns in the vector must match the number
+     *                of columns of this)
+     */
+    void compute_mean(LinOp* result) const;
+
+    /**
+     * Computes the column-wise mean of this matrix.
+     *
+     * @param result  a Dense row vector, used to store the norm
+     *                (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(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,14 @@ class Dense
      */
     virtual void compute_squared_norm2_impl(LinOp* result) const;
 
+    /**
+     * @copydoc compute_mean(LinOp*) const
+     *
+     * @deprecated  This function will be removed in the future,
+     *              we will instead always use Ginkgo's implementation.
+     */
+    virtual void compute_mean_impl(LinOp* result) const;
+
     /**
      * Resizes the matrix to the given size.
      *

reference/matrix/dense_kernels.cpp

@@ -397,6 +397,28 @@ 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()[0]; ++i) {
+        const ValueType_nc alpha = static_cast<ValueType_nc>(i) / (i + 1);
+        const ValueType_nc beta = static_cast<ValueType_nc>(1) / (i + 1);
+        for (size_type j = 0; j < x->get_size()[1]; ++j) {
+            result->at(0, j) = alpha * result->at(0, j) + beta * x->at(i, j);
+        }
+    }
+}
+
+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

@@ -700,6 +700,22 @@ TYPED_TEST(Dense, ComputesNorm1Mixed)
 }
 
 
+TYPED_TEST(Dense, ComputesMean)
+{
+    using Mtx = typename TestFixture::Mtx;
+    using T = typename TestFixture::value_type;
+    using MeanVector = gko::matrix::Dense<T>;
+    auto mtx(gko::initialize<Mtx>(
+        {I<T>{1.0, 0.0}, I<T>{2.0, 3.0}, I<T>{2.0, 4.0}, I<T>{-1.0, -1.0}},
+        this->exec));
+    auto result = MeanVector::create(this->exec, gko::dim<2>{1, 2});
+
+    mtx->compute_mean(result.get());
+
+    GKO_ASSERT_MTX_NEAR(result, l<T>({{1.0, 1.5}}), 1e-2);
+}
+
+
 TYPED_TEST(Dense, ComputeDotFailsOnWrongInputSize)
 {
     using Mtx = typename TestFixture::Mtx;

test/mpi/vector.cpp

@@ -675,6 +675,29 @@ TYPED_TEST(VectorReductions, ComputeSquaredNorm2WithTmpIsSameAsDense)
                         r<value_type>::value);
 }
 
+TYPED_TEST(VectorReductions, ComputesMeanIsSameAsDense)
+{
+    using value_type = typename TestFixture::value_type;
+    this->init_result();
+
+    this->x->compute_mean(this->res.get());
+    this->dense_x->compute_mean(this->dense_res.get());
+
+    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.get(), this->tmp);
+    this->dense_x->compute_mean(this->dense_res.get(), this->dense_tmp);
+
+    GKO_ASSERT_MTX_NEAR(this->res, this->dense_res,
+                        r<value_type>::value);
+}
 
 TYPED_TEST(VectorReductions, ComputeDotCopiesToHostOnlyIfNecessary)
 {