Fixes convergence/papi logger for distributed vectors

core/distributed/helpers.hpp

@@ -30,6 +30,10 @@ THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 ******************************<GINKGO LICENSE>*******************************/
 
+#ifndef GKO_CORE_DISTRIBUTED_HELPERS_HPP_
+#define GKO_CORE_DISTRIBUTED_HELPERS_HPP_
+
+
 #include <memory>
 
 
@@ -124,5 +128,48 @@ bool is_distributed(Arg* linop, Rest*... rest)
 }
 
 
+/**
+ * Cast an input linop to the correct underlying vector type (dense/distributed)
+ * and passes it to the given function.
+ *
+ * @tparam ValueType  The value type of the underlying dense or distributed
+ * vector.
+ * @tparam T  The linop type, either LinOp, or const LinOp.
+ * @tparam F  The function type.
+ * @tparam Args  The types for the additional arguments of f.
+ *
+ * @param linop  The linop to be casted into either a dense or distributed
+ *               vector.
+ * @param f  The function that is to be called with the correctly casted linop.
+ * @param args  The additional arguments of f.
+ */
+template <typename ValueType, typename T, typename F, typename... Args>
+void vector_dispatch(T* linop, F&& f, Args&&... args)
+{
+#if GINKGO_BUILD_MPI
+    if (is_distributed(linop)) {
+        using type = std::conditional_t<
+            std::is_const<T>::value,
+            const experimental::distributed::Vector<ValueType>,
+            experimental::distributed::Vector<ValueType>>;
+        f(dynamic_cast<type*>(linop), std::forward<Args>(args)...);
+    } else
+#endif
+    {
+        using type = std::conditional_t<std::is_const<T>::value,
+                                        const matrix::Dense<ValueType>,
+                                        matrix::Dense<ValueType>>;
+        if (auto concrete_linop = dynamic_cast<type*>(linop)) {
+            f(concrete_linop, std::forward<Args>(args)...);
+        } else {
+            GKO_NOT_SUPPORTED(linop);
+        }
+    }
+}
+
+
 }  // namespace detail
 }  // namespace gko
+
+
+#endif  // GKO_CORE_DISTRIBUTED_HELPERS_HPP_

core/log/convergence.cpp

@@ -35,10 +35,15 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 #include <ginkgo/core/base/array.hpp>
 #include <ginkgo/core/base/math.hpp>
+#include <ginkgo/core/distributed/vector.hpp>
 #include <ginkgo/core/stop/criterion.hpp>
 #include <ginkgo/core/stop/stopping_status.hpp>
 
 
+#include "core/base/dispatch_helper.hpp"
+#include "core/distributed/helpers.hpp"
+
+
 namespace gko {
 namespace log {
 
@@ -73,12 +78,14 @@ void Convergence<ValueType>::on_criterion_check_completed(
         if (residual_norm != nullptr) {
             this->residual_norm_.reset(residual_norm->clone().release());
         } else if (residual != nullptr) {
-            using Vector = matrix::Dense<ValueType>;
             using NormVector = matrix::Dense<remove_complex<ValueType>>;
-            this->residual_norm_ = NormVector::create(
-                residual->get_executor(), dim<2>{1, residual->get_size()[1]});
-            auto dense_r = as<Vector>(residual);
-            dense_r->compute_norm2(this->residual_norm_.get());
+            detail::vector_dispatch<ValueType>(
+                residual, [&](const auto* dense_r) {
+                    this->residual_norm_ =
+                        NormVector::create(residual->get_executor(),
+                                           dim<2>{1, residual->get_size()[1]});
+                    dense_r->compute_norm2(this->residual_norm_.get());
+                });
         }
     }
 }

core/log/papi.cpp

@@ -37,6 +37,9 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/matrix/dense.hpp>
 
 
+#include "core/distributed/helpers.hpp"
+
+
 namespace gko {
 namespace log {
 
@@ -243,12 +246,13 @@ void Papi<ValueType>::on_criterion_check_completed(
         residual_norm_d =
             static_cast<double>(std::real(dense_r_norm->at(0, 0)));
     } else if (residual != nullptr) {
-        auto tmp_res_norm = Vector::create(residual->get_executor(),
-                                           dim<2>{1, residual->get_size()[1]});
-        auto dense_r = as<Vector>(residual);
-        dense_r->compute_norm2(tmp_res_norm.get());
-        residual_norm_d =
-            static_cast<double>(std::real(tmp_res_norm->at(0, 0)));
+        detail::vector_dispatch<ValueType>(residual, [&](const auto* dense_r) {
+            auto tmp_res_norm = Vector::create(
+                residual->get_executor(), dim<2>{1, residual->get_size()[1]});
+            dense_r->compute_norm2(tmp_res_norm.get());
+            residual_norm_d =
+                static_cast<double>(std::real(tmp_res_norm->at(0, 0)));
+        });
     }
 
     const auto tmp = reinterpret_cast<uintptr>(criterion);

core/test/log/convergence.cpp

@@ -44,17 +44,81 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 namespace {
 
+
 template <typename T>
-class Convergence : public ::testing::Test {};
+class Convergence : public ::testing::Test {
+public:
+    using Dense = gko::matrix::Dense<T>;
+    using AbsoluteDense = gko::matrix::Dense<gko::remove_complex<T>>;
+
+    Convergence()
+    {
+        status.get_data()[0].reset();
+        status.get_data()[0].converge(0);
+    }
+
+    std::shared_ptr<gko::ReferenceExecutor> exec =
+        gko::ReferenceExecutor::create();
+
+    std::unique_ptr<Dense> residual = gko::initialize<Dense>({3, 4}, exec);
+    std::unique_ptr<AbsoluteDense> residual_norm =
+        gko::initialize<AbsoluteDense>({5}, exec);
+    std::unique_ptr<AbsoluteDense> implicit_sq_resnorm =
+        gko::initialize<AbsoluteDense>({6}, exec);
+    std::unique_ptr<Dense> solution = gko::initialize<Dense>({-2, 7}, exec);
+
+    gko::array<gko::stopping_status> status = {exec, 1};
+};
 
 TYPED_TEST_SUITE(Convergence, gko::test::ValueTypes, TypenameNameGenerator);
 
 
-TYPED_TEST(Convergence, CanGetData)
+TYPED_TEST(Convergence, CanGetEmptyData)
+{
+    auto logger = gko::log::Convergence<TypeParam>::create(
+        gko::log::Logger::criterion_events_mask);
+
+    ASSERT_EQ(logger->has_converged(), false);
+    ASSERT_EQ(logger->get_num_iterations(), 0);
+    ASSERT_EQ(logger->get_residual(), nullptr);
+    ASSERT_EQ(logger->get_residual_norm(), nullptr);
+    ASSERT_EQ(logger->get_implicit_sq_resnorm(), nullptr);
+}
+
+
+TYPED_TEST(Convergence, CanLogData)
+{
+    using Dense = gko::matrix::Dense<TypeParam>;
+    using AbsoluteDense = gko::matrix::Dense<gko::remove_complex<TypeParam>>;
+    auto logger = gko::log::Convergence<TypeParam>::create(
+        gko::log::Logger::criterion_events_mask);
+
+    logger->template on<gko::log::Logger::criterion_check_completed>(
+        nullptr, 100, this->residual.get(), this->residual_norm.get(),
+        this->implicit_sq_resnorm.get(), this->solution.get(), 0, false,
+        &this->status, false, true);
+
+    ASSERT_EQ(logger->has_converged(), true);
+    ASSERT_EQ(logger->get_num_iterations(), 100);
+    GKO_ASSERT_MTX_NEAR(gko::as<Dense>(logger->get_residual()),
+                        this->residual.get(), 0);
+    GKO_ASSERT_MTX_NEAR(gko::as<AbsoluteDense>(logger->get_residual_norm()),
+                        this->residual_norm.get(), 0);
+    GKO_ASSERT_MTX_NEAR(
+        gko::as<AbsoluteDense>(logger->get_implicit_sq_resnorm()),
+        this->implicit_sq_resnorm.get(), 0);
+}
+
+
+TYPED_TEST(Convergence, DoesNotLogIfNotStopped)
 {
-    auto exec = gko::ReferenceExecutor::create();
     auto logger = gko::log::Convergence<TypeParam>::create(
-        gko::log::Logger::iteration_complete_mask);
+        gko::log::Logger::criterion_events_mask);
+
+    logger->template on<gko::log::Logger::criterion_check_completed>(
+        nullptr, 100, this->residual.get(), this->residual_norm.get(),
+        this->implicit_sq_resnorm.get(), this->solution.get(), 0, false,
+        &this->status, false, false);
 
     ASSERT_EQ(logger->has_converged(), false);
     ASSERT_EQ(logger->get_num_iterations(), 0);
@@ -63,4 +127,19 @@ TYPED_TEST(Convergence, CanGetData)
 }
 
 
+TYPED_TEST(Convergence, CanComputeResidualNorm)
+{
+    using AbsoluteDense = gko::matrix::Dense<gko::remove_complex<TypeParam>>;
+    auto logger = gko::log::Convergence<TypeParam>::create(
+        gko::log::Logger::criterion_events_mask);
+
+    logger->template on<gko::log::Logger::criterion_check_completed>(
+        nullptr, 100, this->residual.get(), nullptr, nullptr, nullptr, 0, false,
+        &this->status, false, true);
+
+    GKO_ASSERT_MTX_NEAR(gko::as<AbsoluteDense>(logger->get_residual_norm()),
+                        this->residual_norm, r<TypeParam>::value);
+}
+
+
 }  // namespace

core/test/mpi/distributed/CMakeLists.txt

@@ -1 +1,2 @@
+ginkgo_create_test(helpers MPI_SIZE 1)
 ginkgo_create_test(matrix MPI_SIZE 1)

core/test/mpi/distributed/helpers.cpp

@@ -0,0 +1,134 @@
+/*******************************<GINKGO LICENSE>******************************
+Copyright (c) 2017-2022, the Ginkgo authors
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions
+are met:
+
+1. Redistributions of source code must retain the above copyright
+notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright
+notice, this list of conditions and the following disclaimer in the
+documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+contributors may be used to endorse or promote products derived from
+this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
+IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
+TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+******************************<GINKGO LICENSE>*******************************/
+
+#include <gtest/gtest.h>
+
+
+#include <ginkgo/core/matrix/csr.hpp>
+
+
+#include "core/distributed/helpers.hpp"
+#include "core/test/utils.hpp"
+
+
+int run_function(gko::experimental::distributed::Vector<>*) { return 1; }
+
+int run_function(const gko::experimental::distributed::Vector<>*) { return 2; }
+
+int run_function(gko::matrix::Dense<>*) { return 3; }
+
+int run_function(const gko::matrix::Dense<>*) { return 4; }
+
+
+class RunVector : public ::testing::Test {
+public:
+    std::shared_ptr<gko::ReferenceExecutor> exec =
+        gko::ReferenceExecutor::create();
+};
+
+
+TEST_F(RunVector, PicksDistributedVectorCorrectly)
+{
+    std::unique_ptr<gko::LinOp> dist_vector =
+        gko::experimental::distributed::Vector<>::create(exec, MPI_COMM_WORLD);
+    int result;
+
+    gko::detail::vector_dispatch<double>(
+        dist_vector.get(), [&](auto* dense) { result = run_function(dense); });
+
+    ASSERT_EQ(result,
+              run_function(gko::as<gko::experimental::distributed::Vector<>>(
+                  dist_vector.get())));
+}
+
+
+TEST_F(RunVector, PicksConstDistributedVectorCorrectly)
+{
+    std::unique_ptr<const gko::LinOp> const_dist_vector =
+        gko::experimental::distributed::Vector<>::create(exec, MPI_COMM_WORLD);
+    int result;
+
+    gko::detail::vector_dispatch<double>(
+        const_dist_vector.get(),
+        [&](auto* dense) { result = run_function(dense); });
+
+    ASSERT_EQ(
+        result,
+        run_function(gko::as<const gko::experimental::distributed::Vector<>>(
+            const_dist_vector.get())));
+}
+
+
+TEST_F(RunVector, PicksDenseVectorCorrectly)
+{
+    std::unique_ptr<gko::LinOp> dense_vector =
+        gko::matrix::Dense<>::create(exec);
+    int result;
+
+    gko::detail::vector_dispatch<double>(
+        dense_vector.get(), [&](auto* dense) { result = run_function(dense); });
+
+    ASSERT_EQ(result,
+              run_function(gko::as<gko::matrix::Dense<>>(dense_vector.get())));
+}
+
+
+TEST_F(RunVector, PicksConstDenseVectorCorrectly)
+{
+    std::unique_ptr<const gko::LinOp> const_dense_vector =
+        gko::matrix::Dense<>::create(exec);
+    int result;
+
+    gko::detail::vector_dispatch<double>(
+        const_dense_vector.get(),
+        [&](auto* dense) { result = run_function(dense); });
+
+    ASSERT_EQ(result, run_function(gko::as<const gko::matrix::Dense<>>(
+                          const_dense_vector.get())));
+}
+
+TEST_F(RunVector, ThrowsIfWrongType)
+{
+    std::unique_ptr<gko::LinOp> csr = gko::matrix::Csr<>::create(exec);
+
+    ASSERT_THROW(
+        gko::detail::vector_dispatch<double>(csr.get(), [&](auto* dense) {}),
+        gko::NotSupported);
+}
+
+
+TEST_F(RunVector, ThrowsIfNullptr)
+{
+    ASSERT_THROW(gko::detail::vector_dispatch<double>(
+                     static_cast<gko::LinOp*>(nullptr), [&](auto* dense) {}),
+                 gko::NotSupported);
+}