add ScalarCache to generate different Dense scalar with different value type

Author: yhmtsai

core/base/dense_cache.cpp

@@ -4,6 +4,12 @@
 
 #include "ginkgo/core/base/dense_cache.hpp"
 
+#include <memory>
+#include <string>
+
+#include <ginkgo/core/base/array.hpp>
+#include <ginkgo/core/base/dim.hpp>
+#include <ginkgo/core/base/executor.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
 
 
@@ -70,6 +76,55 @@ std::shared_ptr<matrix::Dense<ValueType>> GenericDenseCache::get(
 }
 
 
+ScalarCache::ScalarCache(std::shared_ptr<const Executor> executor,
+                         double scalar_value)
+    : exec(std::move(executor)), value(scalar_value){};
+
+ScalarCache::ScalarCache(const ScalarCache& other) { *this = other; }
+
+
+ScalarCache::ScalarCache(ScalarCache&& other) noexcept
+{
+    *this = std::move(other);
+}
+
+
+ScalarCache& ScalarCache::operator=(const ScalarCache& other)
+{
+    exec = other.exec;
+    value = other.value;
+    return *this;
+}
+
+
+ScalarCache& ScalarCache::operator=(ScalarCache&& other) noexcept
+{
+    exec = std::exchange(other.exec, nullptr);
+    value = std::exchange(other.value, 0.0);
+    other.scalars.clear();
+    return *this;
+}
+
+
+template <typename ValueType>
+std::shared_ptr<const matrix::Dense<ValueType>> ScalarCache::get() const
+{
+    // using typeid name as key
+    std::string value_string = typeid(ValueType).name();
+    auto search = scalars.find(value_string);
+    if (search != scalars.end()) {
+        return std::dynamic_pointer_cast<const matrix::Dense<ValueType>>(
+            search->second);
+    } else {
+        auto new_scalar =
+            share(matrix::Dense<ValueType>::create(exec, dim<2>{1, 1}));
+        new_scalar->fill(static_cast<ValueType>(value));
+        scalars[value_string] = new_scalar;
+        return new_scalar;
+    }
+}
+
+
 #define GKO_DECLARE_DENSE_CACHE(_type) struct DenseCache<_type>
 GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_CACHE);
 
@@ -78,6 +133,10 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_DENSE_CACHE);
         std::shared_ptr<const Executor>, dim<2>) const
 GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_GENERIC_DENSE_CACHE_GET);
 
+#define GKO_DECLARE_SCALAR_CACHE_GET(_type) \
+    std::shared_ptr<const matrix::Dense<_type>> ScalarCache::get<_type>() const
+GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_SCALAR_CACHE_GET);
+
 
 }  // namespace detail
 }  // namespace gko

core/distributed/matrix.cpp

@@ -50,7 +50,7 @@ Matrix<ValueType, LocalIndexType, GlobalIndexType>::Matrix(
       DistributedBase{comm},
       row_gatherer_{RowGatherer<LocalIndexType>::create(exec, comm)},
       imap_{exec},
-      one_scalar_{},
+      one_scalar_{exec, 1.0},
       local_mtx_{local_matrix_template->clone(exec)},
       non_local_mtx_{non_local_matrix_template->clone(exec)}
 {
@@ -60,8 +60,6 @@ Matrix<ValueType, LocalIndexType, GlobalIndexType>::Matrix(
     GKO_ASSERT(
         (dynamic_cast<ReadableFromMatrixData<ValueType, LocalIndexType>*>(
             non_local_mtx_.get())));
-    one_scalar_.init(exec, dim<2>{1, 1});
-    one_scalar_->fill(one<value_type>());
 }
 
 template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
@@ -72,13 +70,11 @@ Matrix<ValueType, LocalIndexType, GlobalIndexType>::Matrix(
       DistributedBase{comm},
       row_gatherer_{RowGatherer<LocalIndexType>::create(exec, comm)},
       imap_{exec},
-      one_scalar_{},
+      one_scalar_{exec, 1.0},
       non_local_mtx_(::gko::matrix::Coo<ValueType, LocalIndexType>::create(
           exec, dim<2>{local_linop->get_size()[0], 0}))
 {
     this->set_size(size);
-    one_scalar_.init(exec, dim<2>{1, 1});
-    one_scalar_->fill(one<value_type>());
     local_mtx_ = std::move(local_linop);
 }
 
@@ -91,13 +87,11 @@ Matrix<ValueType, LocalIndexType, GlobalIndexType>::Matrix(
       DistributedBase{comm},
       row_gatherer_(RowGatherer<LocalIndexType>::create(exec, comm)),
       imap_(std::move(imap)),
-      one_scalar_{}
+      one_scalar_{exec, 1.0}
 {
     this->set_size({imap_.get_global_size(), imap_.get_global_size()});
     local_mtx_ = std::move(local_linop);
     non_local_mtx_ = std::move(non_local_linop);
-    one_scalar_.init(exec, dim<2>{1, 1});
-    one_scalar_->fill(one<value_type>());
 
     row_gatherer_ = RowGatherer<LocalIndexType>::create(
         row_gatherer_->get_executor(),
@@ -427,8 +421,9 @@ void Matrix<ValueType, LocalIndexType, GlobalIndexType>::apply_impl(
             if (recv_ptr != recv_buffer.get()) {
                 recv_buffer->copy_from(host_recv_buffer.get());
             }
-            non_local_mtx_->apply(one_scalar_.get(), recv_buffer.get(),
-                                  one_scalar_.get(), local_x);
+            non_local_mtx_->apply(
+                one_scalar_.template get<ValueType>().get(), recv_buffer.get(),
+                one_scalar_.template get<x_value_type>().get(), local_x);
         },
         b, x);
 }
@@ -479,8 +474,9 @@ void Matrix<ValueType, LocalIndexType, GlobalIndexType>::apply_impl(
             if (recv_ptr != recv_buffer.get()) {
                 recv_buffer->copy_from(host_recv_buffer.get());
             }
-            non_local_mtx_->apply(local_alpha.get(), recv_buffer.get(),
-                                  one_scalar_.get(), local_x);
+            non_local_mtx_->apply(
+                local_alpha.get(), recv_buffer.get(),
+                one_scalar_.template get<x_value_type>().get(), local_x);
         },
         b, x);
 }
@@ -573,7 +569,8 @@ Matrix<ValueType, LocalIndexType, GlobalIndexType>::Matrix(const Matrix& other)
       DistributedBase{other.get_communicator()},
       row_gatherer_{RowGatherer<LocalIndexType>::create(
           other.get_executor(), other.get_communicator())},
-      imap_(other.get_executor())
+      imap_(other.get_executor()),
+      one_scalar_(other.get_executor(), 1.0)
 {
     *this = other;
 }
@@ -587,7 +584,8 @@ Matrix<ValueType, LocalIndexType, GlobalIndexType>::Matrix(
       DistributedBase{other.get_communicator()},
       row_gatherer_{RowGatherer<LocalIndexType>::create(
           other.get_executor(), other.get_communicator())},
-      imap_(other.get_executor())
+      imap_(other.get_executor()),
+      one_scalar_(other.get_executor(), 1.0)
 {
     *this = std::move(other);
 }
@@ -606,8 +604,6 @@ Matrix<ValueType, LocalIndexType, GlobalIndexType>::operator=(
         non_local_mtx_->copy_from(other.non_local_mtx_);
         row_gatherer_->copy_from(other.row_gatherer_);
         imap_ = other.imap_;
-        one_scalar_.init(this->get_executor(), dim<2>{1, 1});
-        one_scalar_->fill(one<value_type>());
     }
     return *this;
 }
@@ -626,8 +622,6 @@ Matrix<ValueType, LocalIndexType, GlobalIndexType>::operator=(Matrix&& other)
         non_local_mtx_->move_from(other.non_local_mtx_);
         row_gatherer_->move_from(other.row_gatherer_);
         imap_ = std::move(other.imap_);
-        one_scalar_.init(this->get_executor(), dim<2>{1, 1});
-        one_scalar_->fill(one<value_type>());
     }
     return *this;
 }

core/test/base/dense_cache.cpp

@@ -362,3 +362,129 @@ TYPED_TEST(GenericDenseCache, WorkspaceIsNotMoveAssigned)
     ASSERT_EQ(cache.workspace.get_size(), 0);
     ASSERT_EQ(cache.workspace.get_executor(), nullptr);
 }
+
+
+template <typename ValueType>
+class ScalarCache : public ::testing::Test {
+protected:
+    using value_type = ValueType;
+
+    ScalarCache()
+        : ref(gko::ReferenceExecutor::create()), value(1.0), cache(ref, value)
+    {}
+
+    std::shared_ptr<gko::ReferenceExecutor> ref;
+    double value;
+    gko::detail::ScalarCache cache;
+};
+
+TYPED_TEST_SUITE(ScalarCache, gko::test::ValueTypes, TypenameNameGenerator);
+
+
+TYPED_TEST(ScalarCache, CanInitWithExecutorAndValue)
+{
+    using value_type = typename TestFixture::value_type;
+
+    gko::detail::ScalarCache cache(this->ref, 1.0);
+
+    ASSERT_EQ(this->cache.exec, this->ref);
+    ASSERT_EQ(this->cache.value, 1.0);
+    ASSERT_EQ(this->cache.scalars.size(), 0);
+}
+
+
+TYPED_TEST(ScalarCache, CanGetScalar)
+{
+    using value_type = typename TestFixture::value_type;
+
+    auto scalar = this->cache.template get<value_type>();
+
+    ASSERT_NE(scalar, nullptr);
+    GKO_ASSERT_EQUAL_DIMENSIONS(scalar->get_size(), gko::dim<2>(1, 1));
+    ASSERT_EQ(scalar->at(0, 0), static_cast<value_type>(this->value));
+    ASSERT_EQ(this->cache.scalars.size(), 1);
+}
+
+
+TYPED_TEST(ScalarCache, CanGetScalarWithDifferentType)
+{
+    using value_type = typename TestFixture::value_type;
+    using another_type = gko::next_precision<value_type>;
+
+    auto scalar = this->cache.template get<value_type>();
+    auto another_scalar = this->cache.template get<another_type>();
+
+    // The original one is still valid
+    ASSERT_NE(scalar, nullptr);
+    GKO_ASSERT_EQUAL_DIMENSIONS(scalar->get_size(), gko::dim<2>(1, 1));
+    ASSERT_EQ(scalar->at(0, 0), static_cast<value_type>(this->value));
+    ASSERT_NE(another_scalar, nullptr);
+    GKO_ASSERT_EQUAL_DIMENSIONS(another_scalar->get_size(), gko::dim<2>(1, 1));
+    ASSERT_EQ(another_scalar->at(0, 0), static_cast<another_type>(this->value));
+    // have two for different value type now
+    ASSERT_EQ(this->cache.scalars.size(), 2);
+}
+
+
+TYPED_TEST(ScalarCache, VectorIsNotCopied)
+{
+    using value_type = typename TestFixture::value_type;
+    auto scalar = this->cache.template get<value_type>();
+
+    gko::detail::ScalarCache cache(this->cache);
+
+    ASSERT_EQ(cache.scalars.size(), 0);
+    ASSERT_EQ(cache.value, this->cache.value);
+    ASSERT_EQ(cache.exec, this->cache.exec);
+    ASSERT_EQ(this->cache.scalars.size(), 1);
+}
+
+
+TYPED_TEST(ScalarCache, VectorIsNotMoved)
+{
+    using value_type = typename TestFixture::value_type;
+    auto scalar = this->cache.template get<value_type>();
+
+    gko::detail::ScalarCache cache(std::move(this->cache));
+
+    ASSERT_EQ(cache.scalars.size(), 0);
+    ASSERT_EQ(cache.value, this->value);
+    ASSERT_EQ(cache.exec, this->ref);
+    // The original one is cleared
+    ASSERT_EQ(this->cache.value, 0.0);
+    ASSERT_EQ(this->cache.exec, nullptr);
+    ASSERT_EQ(this->cache.scalars.size(), 0);
+}
+
+
+TYPED_TEST(ScalarCache, VectorIsNotCopyAssigned)
+{
+    using value_type = typename TestFixture::value_type;
+    auto scalar = this->cache.template get<value_type>();
+    gko::detail::ScalarCache cache(this->ref, 2.0);
+
+    cache = this->cache;
+
+    ASSERT_EQ(cache.scalars.size(), 0);
+    ASSERT_EQ(cache.value, this->cache.value);
+    ASSERT_EQ(cache.exec, this->cache.exec);
+    ASSERT_EQ(this->cache.scalars.size(), 1);
+}
+
+
+TYPED_TEST(ScalarCache, VectorIsNotMoveAssigned)
+{
+    using value_type = typename TestFixture::value_type;
+    auto scalar = this->cache.template get<value_type>();
+    gko::detail::ScalarCache cache(this->ref, 2.0);
+
+    cache = std::move(this->cache);
+
+    ASSERT_EQ(cache.scalars.size(), 0);
+    ASSERT_EQ(cache.value, this->value);
+    ASSERT_EQ(cache.exec, this->ref);
+    // The original one is cleared
+    ASSERT_EQ(this->cache.value, 0.0);
+    ASSERT_EQ(this->cache.exec, nullptr);
+    ASSERT_EQ(this->cache.scalars.size(), 0);
+}

include/ginkgo/core/base/dense_cache.hpp

@@ -6,14 +6,22 @@
 #define GKO_PUBLIC_CORE_BASE_DENSE_CACHE_HPP_
 
 
+#include <map>
 #include <memory>
+#include <string>
 
 #include <ginkgo/core/base/array.hpp>
 #include <ginkgo/core/base/dim.hpp>
-#include <ginkgo/core/base/executor.hpp>
 
 
 namespace gko {
+
+
+class Executor;
+
+class LinOp;
+
+
 namespace matrix {
 
 
@@ -122,6 +130,38 @@ struct GenericDenseCache {
 };
 
 
+/**
+ * Manages a map to store Dense Scalar with different value_type by a
+ * user-specified value. The workspace is buffered and reused internally to
+ * avoid repeated allocations. Copying an instance will only yield an empty
+ * object since copying the cached vector would not make sense. The stored
+ * object is always mutable, so the cache can be used in a const-context.
+ *
+ * @internal  The struct is present to wrap cache-like buffer storage that will
+ *            not be copied when the outer object gets copied.
+ */
+struct ScalarCache {
+    ScalarCache(std::shared_ptr<const Executor> executor, double scalar_value);
+    ~ScalarCache() = default;
+    ScalarCache(const ScalarCache& other);
+    ScalarCache(ScalarCache&& other) noexcept;
+    ScalarCache& operator=(const ScalarCache& other);
+    ScalarCache& operator=(ScalarCache&& other) noexcept;
+    std::shared_ptr<const Executor> exec;
+    double value;
+    mutable std::map<std::string, std::shared_ptr<const gko::LinOp>> scalars;
+
+
+    /**
+     * Pointer access to the underlying vector with specific type.
+     *
+     * @return  Pointer to the vector view.
+     */
+    template <typename ValueType>
+    std::shared_ptr<const matrix::Dense<ValueType>> get() const;
+};
+
+
 }  // namespace detail
 }  // namespace gko
 

include/ginkgo/core/distributed/matrix.hpp

@@ -681,7 +681,7 @@ class Matrix
 private:
     std::shared_ptr<RowGatherer<LocalIndexType>> row_gatherer_;
     index_map<local_index_type, global_index_type> imap_;
-    gko::detail::DenseCache<value_type> one_scalar_;
+    gko::detail::ScalarCache one_scalar_;
     gko::detail::GenericDenseCache recv_buffer_;
     gko::detail::GenericDenseCache host_recv_buffer_;
     std::shared_ptr<LinOp> local_mtx_;