Simplify data access for host collections (celeritas-project#1029)

* add CollectionTraits for ldg-supported type and host memspace * fix collection test * cleanup
sethrj · Nov 23, 2023 · 2a6042a · 2a6042a
1 parent 5eaa199
commit 2a6042a
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 50 deletions.
diff --git a/src/corecel/data/Collection.hh b/src/corecel/data/Collection.hh
@@ -238,7 +238,7 @@ class Collection
     static_assert(std::is_trivially_destructible<T>::value || CELERITAS_USE_HIP,
                   "Collection element is not trivially destructible");
 
-    using CollectionTraitsT = detail::CollectionTraits<T, W>;
+    using CollectionTraitsT = detail::CollectionTraits<T, W, M>;
     using const_value_type = typename CollectionTraitsT::const_type;
 
   public:

diff --git a/src/corecel/data/detail/CollectionImpl.hh b/src/corecel/data/detail/CollectionImpl.hh
@@ -32,7 +32,7 @@ namespace celeritas
 namespace detail
 {
 //---------------------------------------------------------------------------//
-template<class T, Ownership W, typename = void>
+template<class T, Ownership W, MemSpace M, typename = void>
 struct CollectionTraits
 {
     using type = T;
@@ -44,8 +44,8 @@ struct CollectionTraits
 };
 
 //---------------------------------------------------------------------------//
-template<class T>
-struct CollectionTraits<T, Ownership::reference, void>
+template<class T, MemSpace M>
+struct CollectionTraits<T, Ownership::reference, M, void>
 {
     using type = T;
     using const_type = T;
@@ -56,9 +56,10 @@ struct CollectionTraits<T, Ownership::reference, void>
 };
 
 //---------------------------------------------------------------------------//
-template<class T>
+template<class T, MemSpace M>
 struct CollectionTraits<T,
                         Ownership::const_reference,
+                        M,
                         std::enable_if_t<!is_ldg_supported_v<std::add_const_t<T>>>>
 {
     using type = T const;
@@ -68,35 +69,46 @@ struct CollectionTraits<T,
     using SpanT = Span<type>;
     using SpanConstT = Span<const_type>;
 };
+
+//---------------------------------------------------------------------------//
+template<class T, MemSpace M>
+struct CollectionTraits<T,
+                        Ownership::const_reference,
+                        M,
+                        std::enable_if_t<is_ldg_supported_v<std::add_const_t<T>>>>
+{
+    using type = T const;
+    using const_type = T const;
+    using reference_type = type&;
+    using const_reference_type = const_type&;
+    using SpanT = Span<type>;
+    using SpanConstT = Span<const_type>;
+};
+
+//---------------------------------------------------------------------------//
 template<class T>
 struct CollectionTraits<T,
                         Ownership::const_reference,
+                        MemSpace::device,
                         std::enable_if_t<is_ldg_supported_v<std::add_const_t<T>>>>
 {
     using type = T const;
     using const_type = T const;
     using reference_type = type;
     using const_reference_type = const_type;
-    using SpanT = Span<LdgValue<const_type>>;
-    using SpanConstT = Span<LdgValue<const_type>>;
+    using SpanT = LdgSpan<const_type>;
+    using SpanConstT = LdgSpan<const_type>;
 };
 
 //---------------------------------------------------------------------------//
 //! Memspace-dependent storage for a collection
 template<class T, Ownership W, MemSpace M>
 struct CollectionStorage
 {
-    using type = typename CollectionTraits<T, W>::SpanT;
+    using type = typename CollectionTraits<T, W, M>::SpanT;
     type data;
 };
 
-template<class T>
-struct CollectionStorage<T, Ownership::value, MemSpace::host>;
-template<class T>
-struct CollectionStorage<T, Ownership::value, MemSpace::device>;
-template<class T>
-struct CollectionStorage<T, Ownership::value, MemSpace::mapped>;
-
 //---------------------------------------------------------------------------//
 //! Storage implementation for managed host data
 template<class T>
@@ -144,6 +156,30 @@ struct CollectionStorage<T, Ownership::value, MemSpace::mapped>
     type data;
 };
 
+//---------------------------------------------------------------------------//
+//! Check that sizes are acceptable when creating references from values
+template<Ownership W>
+struct CollectionStorageValidator
+{
+    template<class Size, class OtherSize>
+    void operator()(Size, OtherSize)
+    {
+        /* No validation needed */
+    }
+};
+
+template<>
+struct CollectionStorageValidator<Ownership::value>
+{
+    template<class Size, class OtherSize>
+    void operator()(Size dst, OtherSize src)
+    {
+        CELER_VALIDATE(dst == src,
+                       << "collection is too large (" << sizeof(Size)
+                       << "-byte int cannot hold " << src << " elements)");
+    }
+};
+
 //---------------------------------------------------------------------------//
 //! Assignment semantics for a collection
 template<Ownership W, MemSpace M>
@@ -172,37 +208,6 @@ struct CollectionAssigner
     }
 };
 
-template<>
-struct CollectionAssigner<Ownership::value, MemSpace::host>;
-template<>
-struct CollectionAssigner<Ownership::value, MemSpace::device>;
-template<>
-struct CollectionAssigner<Ownership::value, MemSpace::mapped>;
-
-//---------------------------------------------------------------------------//
-//! Check that sizes are acceptable when creating references from values
-template<Ownership W>
-struct CollectionStorageValidator
-{
-    template<class Size, class OtherSize>
-    void operator()(Size, OtherSize)
-    {
-        /* No validation needed */
-    }
-};
-
-template<>
-struct CollectionStorageValidator<Ownership::value>
-{
-    template<class Size, class OtherSize>
-    void operator()(Size dst, OtherSize src)
-    {
-        CELER_VALIDATE(dst == src,
-                       << "collection is too large (" << sizeof(Size)
-                       << "-byte int cannot hold " << src << " elements)");
-    }
-};
-
 //---------------------------------------------------------------------------//
 //! Assignment semantics for copying to host memory
 template<>

diff --git a/test/corecel/data/Collection.test.cc b/test/corecel/data/Collection.test.cc
@@ -309,11 +309,10 @@ TEST_F(SimpleCollectionTest, accessors)
     EXPECT_EQ(321, host_ref_cref[AllInts<host>{}].back());
 
     CRef<host> host_cref{host_val};
-    EXPECT_TRUE((std::is_same_v<decltype(host_cref[IntId{0}]), int>));
+    EXPECT_TRUE((std::is_same_v<decltype(host_cref[IntId{0}]), int const&>));
+    EXPECT_TRUE((std::is_same_v<decltype(host_cref[irange]), Span<int const>>));
     EXPECT_TRUE((
-        std::is_same_v<decltype(host_cref[irange]), Span<LdgValue<int const>>>));
-    EXPECT_TRUE((std::is_same_v<decltype(host_cref[AllInts<host>{}]),
-                                Span<LdgValue<int const>>>));
+        std::is_same_v<decltype(host_cref[AllInts<host>{}]), Span<int const>>));
     EXPECT_EQ(4, host_ref.size());
     EXPECT_EQ(123, host_cref[IntId{0}]);
     EXPECT_EQ(123, host_cref[irange].front());
@@ -342,6 +341,14 @@ TEST_F(SimpleCollectionTest, TEST_IF_CELER_DEVICE(algo_device))
     resize(&src, 2);
     fill(123, &src);
 
+    CRef<device> device_cref{src};
+    EXPECT_TRUE((std::is_same_v<decltype(device_cref[IntId{0}]), int>));
+    EXPECT_TRUE(
+        (std::is_same_v<decltype(device_cref[IntRange{IntId{0}, IntId{2}}]),
+                        LdgSpan<int const>>));
+    EXPECT_TRUE((std::is_same_v<decltype(device_cref[AllInts<device>{}]),
+                                LdgSpan<int const>>));
+
     // Test 'copy_to_host'
     std::vector<int> dst(src.size());
     copy_to_host(src, make_span(dst));