Use ‘const’ in Neighbor Table Value Types (#1796)

Additionally fixes missing `host_view` etc. in const data stores. Gives speedups when combined with #1795.
GridTools · Sep 30, 2024 · c4d2ced · c4d2ced
1 parent 8e05119
commit c4d2ced
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 20 deletions.
diff --git a/include/gridtools/storage/data_store.hpp b/include/gridtools/storage/data_store.hpp
@@ -173,38 +173,52 @@ namespace gridtools {
                 auto const_host_view() const { return const_target_view(); }
             };
 
-            template <class Traits, class T, class Info, class Kind, bool IsHostRefrenceable>
-            class data_store<Traits, T const, Info, Kind, true, IsHostRefrenceable>
-                : public base<Traits, T const, Info, Kind> {
-
-                template <class>
-                struct is_host_refrenceable : std::bool_constant<IsHostRefrenceable> {};
-
-                template <class Initializer, std::enable_if_t<!is_host_refrenceable<Initializer>::value, int> = 0>
-                void init(Initializer const &initializer) {
-                    auto host_ptr = std::make_unique<T[]>(this->info().length());
-                    initializer(host_ptr.get(), typename data_store::layout_t(), this->info());
-                    traits::update_target<Traits>(this->raw_target_ptr(), host_ptr.get(), this->info().length());
-                }
+            template <class Traits, class T, class Info, class Kind>
+            class data_store<Traits, T const, Info, Kind, true, false> : public base<Traits, T const, Info, Kind> {
+                std::unique_ptr<T[]> m_host_ptr;
 
-                template <class Initializer, std::enable_if_t<is_host_refrenceable<Initializer>::value, int> = 0>
-                void init(Initializer const &initializer) {
-                    initializer(this->raw_target_ptr(), typename data_store::layout_t(), this->info());
+              public:
+                template <class Halos>
+                data_store(std::string, Info, Halos const &, uninitialized const &) = delete;
+
+                template <class Initializer, class Halos>
+                data_store(std::string name, Info info, Halos const &halos, Initializer const &initializer)
+                    : data_store::base(std::move(name), std::move(info), halos),
+                      m_host_ptr(std::make_unique<T[]>(this->info().length())) {
+                    initializer(m_host_ptr.get(), typename data_store::layout_t(), this->info());
+                    traits::update_target<Traits>(this->raw_target_ptr(), m_host_ptr.get(), this->info().length());
                 }
 
+                T const *get_target_ptr() const { return this->raw_target_ptr(); }
+                T const *get_const_target_ptr() const { return this->raw_target_ptr(); }
+                auto target_view() const { return traits::make_target_view<Traits>(get_target_ptr(), this->info()); }
+                auto const_target_view() const { return target_view(); }
+
+                T const *get_host_ptr() const { return m_host_ptr.get(); }
+                T const *get_const_host_ptr() const { return get_host_ptr(); }
+                auto host_view() const { return make_host_view(get_host_ptr(), this->info()); }
+                auto const_host_view() const { return host_view(); }
+            };
+
+            template <class Traits, class T, class Info, class Kind>
+            class data_store<Traits, T const, Info, Kind, true, true> : public base<Traits, T const, Info, Kind> {
               public:
                 template <class Halos>
                 data_store(std::string, Info, Halos const &, uninitialized const &) = delete;
 
                 template <class Initializer, class Halos>
                 data_store(std::string name, Info info, Halos const &halos, Initializer const &initializer)
                     : base<Traits, T const, Info, Kind>(std::move(name), std::move(info), halos) {
-                    init(initializer);
+                    initializer(this->raw_target_ptr(), typename data_store::layout_t(), this->info());
                 }
                 T const *get_target_ptr() const { return this->raw_target_ptr(); }
+                T const *get_const_target_ptr() const { return get_target_ptr(); }
                 auto target_view() const { return traits::make_target_view<Traits>(get_target_ptr(), this->info()); }
-                auto get_const_target_ptr() const { return get_target_ptr(); }
                 auto const_target_view() const { return target_view(); }
+                T const *get_host_ptr() { return get_target_ptr(); }
+                T const *get_const_host_ptr() { return get_target_ptr(); }
+                auto host_view() const { return target_view(); }
+                auto const_host_view() const { return target_view(); }
             };
 
             template <class>

diff --git a/tests/include/fn_mesh.hpp b/tests/include/fn_mesh.hpp
@@ -120,11 +120,11 @@ namespace gridtools {
         }
 
         auto v2e_table() const {
-            return storage::builder<StorageTraits>.dimensions(nvertices(), max_v2e_neighbors_t()).template type<int>().initializer(v2e_initializer()).unknown_id().build();
+            return storage::builder<StorageTraits>.dimensions(nvertices(), max_v2e_neighbors_t()).template type<int const>().initializer(v2e_initializer()).unknown_id().build();
         }
 
         auto e2v_table() const {
-            return storage::builder<StorageTraits>.dimensions(nedges(), max_e2v_neighbors_t()).template type<int>().initializer(e2v_initializer()).unknown_id().build();
+            return storage::builder<StorageTraits>.dimensions(nedges(), max_e2v_neighbors_t()).template type<int const>().initializer(e2v_initializer()).unknown_id().build();
         }
     };