single_store: added support for handling multiple contexts on the host

Author: Vakho Tsulaia

core/include/detray/core/detail/single_store.hpp

@@ -18,6 +18,7 @@
 #include <vecmem/memory/memory_resource.hpp>
 
 // System include(s)
+#include <iostream>
 #include <type_traits>
 
 namespace detray {
@@ -65,7 +66,9 @@ class single_store {
     template <typename allocator_t = vecmem::memory_resource>
     requires(std::derived_from<allocator_t, std::pmr::memory_resource>)
         DETRAY_HOST explicit single_store(allocator_t &resource)
-        : m_container(&resource) {}
+        : m_container(&resource) {
+        m_context_size = m_container.size();
+    }
 
     /// Copy Construct with a specific memory resource @param resource
     /// (host-side only)
@@ -74,12 +77,16 @@ class single_store {
     requires(std::is_same_v<C, std::vector<T>>
                  &&std::derived_from<allocator_t, std::pmr::memory_resource>)
         DETRAY_HOST single_store(allocator_t &resource, const T &arg)
-        : m_container(&resource, arg) {}
+        : m_container(&resource, arg) {
+        m_context_size = m_container.size();
+    }
 
     /// Construct from the container @param view . Mainly used device-side.
     template <concepts::device_view container_view_t>
     DETRAY_HOST_DEVICE explicit single_store(container_view_t &view)
-        : m_container(view) {}
+        : m_container(view) {
+        m_context_size = m_container.size();
+    }
 
     /// @returns a pointer to the underlying container - const
     DETRAY_HOST_DEVICE
@@ -95,26 +102,29 @@ class single_store {
     DETRAY_HOST_DEVICE
     constexpr auto size(const context_type & /*ctx*/ = {}) const noexcept
         -> dindex {
-        return static_cast<dindex>(m_container.size());
+        return static_cast<dindex>(m_context_size);
     }
 
     /// @returns true if the underlying container is empty
     DETRAY_HOST_DEVICE
-    constexpr auto empty(const context_type & /*ctx*/ = {}) const noexcept
-        -> bool {
-        return m_container.empty();
+    constexpr auto empty(const context_type &ctx = {}) const noexcept -> bool {
+        if (ctx.get() == 0u) {
+            return m_container.empty();
+        } else {
+            return ctx.get() > m_n_contexts;
+        }
     }
 
     /// @returns the collections iterator at the start position
     DETRAY_HOST_DEVICE
-    constexpr auto begin(const context_type & /*ctx*/ = {}) const {
-        return m_container.begin();
+    constexpr auto begin(const context_type &ctx = {}) const {
+        return m_container.begin() + ctx.get() * m_context_size;
     }
 
     /// @returns the collections iterator sentinel
     DETRAY_HOST_DEVICE
-    constexpr auto end(const context_type & /*ctx*/ = {}) const {
-        return m_container.end();
+    constexpr auto end(const context_type &ctx = {}) const {
+        return m_container.begin() + (ctx.get() + 1) * m_context_size;
     }
 
     /// @returns access to the underlying container - const
@@ -135,33 +145,34 @@ class single_store {
     constexpr auto at(const dindex i,
                       const context_type &ctx = {}) const noexcept
         -> const T & {
-        [[maybe_unused]] context_type tmp_ctx{
-            ctx};  // Temporary measure to avoid warnings
-        return m_container.at(i);
+        return m_container.at(ctx.get() * m_context_size + i);
     }
 
     /// @returns context based access to an element (also range checked)
     DETRAY_HOST_DEVICE
     constexpr auto at(const dindex i, const context_type &ctx = {}) noexcept
         -> T & {
-        [[maybe_unused]] context_type tmp_ctx{
-            ctx};  // Temporary measure to avoid warnings
-        return m_container.at(i);
+        return m_container.at(ctx.get() * m_context_size + i);
     }
 
     /// Removes and destructs all elements in the container.
     DETRAY_HOST void clear(const context_type & /*ctx*/) {
+        assert(m_n_contexts == 0u);
         m_container.clear();
+        m_context_size = 0;
     }
 
     /// Reserve memory of size @param n for a given geometry context
     DETRAY_HOST void reserve(std::size_t n, const context_type & /*ctx*/) {
+        assert(m_n_contexts == 0u);
         m_container.reserve(n);
     }
 
     /// Resize the underlying container to @param n for a given geometry context
     DETRAY_HOST void resize(std::size_t n, const context_type & /*ctx*/) {
+        assert(m_n_contexts == 0u);
         m_container.resize(n);
+        m_context_size = m_container.size();
     }
 
     /// Add a new element to the collection - copy
@@ -175,7 +186,9 @@ class single_store {
     DETRAY_HOST constexpr auto push_back(
         const U &arg, const context_type & /*ctx*/ = {}) noexcept(false)
         -> void {
+        assert(m_n_contexts == 0u);
         m_container.push_back(arg);
+        m_context_size = m_container.size();
     }
 
     /// Add a new element to the collection - move
@@ -188,7 +201,9 @@ class single_store {
     template <typename U>
     DETRAY_HOST constexpr auto push_back(
         U &&arg, const context_type & /*ctx*/ = {}) noexcept(false) -> void {
+        assert(m_n_contexts == 0u);
         m_container.push_back(std::forward<U>(arg));
+        m_context_size = m_container.size();
     }
 
     /// Add a new element to the collection in place
@@ -201,6 +216,8 @@ class single_store {
     template <typename... Args>
     DETRAY_HOST constexpr decltype(auto) emplace_back(
         const context_type & /*ctx*/ = {}, Args &&... args) noexcept(false) {
+        assert(m_n_contexts == 0u);
+        m_context_size++;
         return m_container.emplace_back(std::forward<Args>(args)...);
     }
 
@@ -215,8 +232,10 @@ class single_store {
     DETRAY_HOST auto insert(container_t<U> &new_data,
                             const context_type & /*ctx*/ = {}) noexcept(false)
         -> void {
+        assert(m_n_contexts == 0u);
         m_container.reserve(m_container.size() + new_data.size());
         m_container.insert(m_container.end(), new_data.begin(), new_data.end());
+        m_context_size = m_container.size();
     }
 
     /// Insert another collection - move
@@ -230,30 +249,61 @@ class single_store {
     DETRAY_HOST auto insert(container_t<U> &&new_data,
                             const context_type & /*ctx*/ = {}) noexcept(false)
         -> void {
+        assert(m_n_contexts == 0u);
         m_container.reserve(m_container.size() + new_data.size());
         m_container.insert(m_container.end(),
                            std::make_move_iterator(new_data.begin()),
                            std::make_move_iterator(new_data.end()));
+        m_context_size = m_container.size();
+    }
+
+    template <typename U>
+    DETRAY_HOST auto add_context(container_t<U> &context_data) noexcept(false)
+        -> void {
+        // Cannot add context data to an empty store
+        if (m_context_size == 0u) {
+            std::cout << "WARNING: Single Store. Cannot add a context to an "
+                         "empty store ";
+            return;
+        }
+        // Wrong size of the context_data vector
+        if (context_data.size() % m_context_size != 0u) {
+            std::cout << "WARNING: Single Store. Wrong size of the inserted "
+                         "vector. Must be multiple of the context size";
+            return;
+        }
+        // Drop previous contexts if any
+        if (m_container.size() > m_context_size)
+            m_container.resize(m_context_size);
+        // Add new contexts
+        m_n_contexts = context_data.size() / m_context_size;
+        m_container.reserve(m_container.size() + context_data.size());
+        m_container.insert(m_container.end(), context_data.begin(),
+                           context_data.end());
     }
 
     /// Append another store to the current one
     ///
     /// @param other The other container
     ///
     /// @note in general can throw an exception
+    /// ? Remove the ctx argument ?
     DETRAY_HOST void append(single_store &other,
                             const context_type &ctx = {}) noexcept(false) {
         insert(other.m_container, ctx);
+        m_context_size = m_container.size();
     }
 
     /// Append another store to the current one - move
     ///
     /// @param other The other container
     ///
     /// @note in general can throw an exception
+    /// ? Remove the ctx argument ?
     DETRAY_HOST void append(single_store &&other,
                             const context_type &ctx = {}) noexcept(false) {
         insert(std::move(other.m_container), ctx);
+        m_context_size = m_container.size();
     }
 
     /// @return the view on the underlying container - non-const
@@ -269,6 +319,8 @@ class single_store {
     private:
     /// The underlying container implementation
     base_type m_container;
+    size_type m_context_size{0u};
+    size_type m_n_contexts{0u};
 };
 
 }  // namespace detray

tests/unit_tests/cpu/core/transform_store.cpp

@@ -9,6 +9,7 @@
 #include "detray/core/detail/single_store.hpp"
 
 // Detray test include(s)
+#include "detray/test/common/assert.hpp"
 #include "detray/test/utils/types.hpp"
 
 // GTest include(s)
@@ -20,7 +21,8 @@ GTEST_TEST(detray_core, static_transform_store) {
     using transform3 = test::transform3;
     using point3 = test::point3;
 
-    using transform_store_t = single_store<transform3>;
+    using transform_store_t =
+        single_store<transform3, dvector, geometry_context>;
     transform_store_t static_store;
     typename transform_store_t::context_type ctx0{};
     typename transform_store_t::context_type ctx1{};
@@ -51,3 +53,63 @@ GTEST_TEST(detray_core, static_transform_store) {
     static_store.emplace_back(ctx0);
     ASSERT_EQ(static_store.size(ctx0), 5u);
 }
+
+// This tests the construction of a static transform store
+GTEST_TEST(detray_core, multicontext_transform_store) {
+    using namespace detray;
+    using transform3 = test::transform3;
+    using point3 = test::point3;
+
+    using transform_store_t =
+        single_store<transform3, dvector, geometry_context>;
+    transform_store_t xf_store;
+    typename transform_store_t::context_type ctx0{0};
+    typename transform_store_t::context_type ctx1{1};
+
+    // Populate the default context (context 0) of the transform store
+    // Add new elements one at a time
+    ASSERT_TRUE(xf_store.empty(ctx0));
+
+    point3 t0{0.f, 0.f, 0.f};
+    transform3 tf0{t0};
+    xf_store.push_back(tf0);
+
+    point3 t1{1.f, 10.f, 7.f};
+    transform3 tf1{t1};
+    xf_store.push_back(tf1);
+
+    point3 t2{2.f, 1.f, 4.f};
+    transform3 tf2{t2};
+    xf_store.push_back(tf2, ctx0);
+
+    point3 t3{3.f, 8.f, 6.f};
+    transform3 tf3{t3};
+    xf_store.push_back(tf3);
+
+    ASSERT_EQ(xf_store.size(ctx0), 4u);
+
+    // Create a vector of 'aligned' transforms
+    point3 shift{.11f, .12f, .13f};
+    typename transform_store_t::base_type aligned_transforms;
+    aligned_transforms.reserve(xf_store.size(ctx0));
+    aligned_transforms.push_back(transform3(t0 + shift));
+    aligned_transforms.push_back(transform3(t1 + shift));
+    aligned_transforms.push_back(transform3(t2 + shift));
+    aligned_transforms.push_back(transform3(t3 + shift));
+
+    // Add context 1 to the transform store
+    ASSERT_TRUE(xf_store.empty(ctx1));
+    xf_store.add_context(aligned_transforms);
+    ASSERT_EQ(xf_store.size(ctx1), 4u);
+
+    // Check that we can access context data as expected
+    auto tr0_0 = xf_store.at(0, ctx0).translation();
+    auto tr1_0 = xf_store.at(0, ctx1).translation();
+    auto tr_diff_0 = tr1_0 - tr0_0;
+    EXPECT_POINT3_NEAR(tr_diff_0, shift, 1e-6);
+
+    auto tr0_3 = xf_store.at(3, ctx0).translation();
+    auto tr1_3 = xf_store.at(3, ctx1).translation();
+    auto tr_diff_3 = tr1_3 - tr0_3;
+    EXPECT_POINT3_NEAR(tr_diff_3, shift, 1e-6);
+}

tests/unit_tests/cpu/navigation/intersection/intersection_kernel.cpp

@@ -73,7 +73,8 @@ using mask_container_t =
 using mask_link_t = typename mask_container_t::single_link;
 using material_link_t = dtyped_index<material_ids, dindex>;
 
-using transform_container_t = single_store<test::transform3>;
+using transform_container_t =
+    single_store<test::transform3, dvector, geometry_context>;
 
 /// The Surface definition:
 using surface_t =

tests/unit_tests/device/cuda/container_cuda.cpp

@@ -46,7 +46,7 @@ TEST(container_cuda, single_store) {
     EXPECT_EQ(mng_store.size(), 0u);
 
     // Test the managed memory allocation
-    empty_context ctx{};
+    geometry_context ctx{};
     mng_store.reserve(4, ctx);
     mng_store.emplace_back(ctx, 1.);
     mng_store.push_back(2., ctx);

tests/unit_tests/device/cuda/container_cuda_kernel.hpp

@@ -19,8 +19,9 @@ namespace detray {
 
 // Single store test
 /// @{
-using single_store_t = single_store<double, vecmem::vector>;
-using single_store_dev_t = single_store<double, vecmem::device_vector>;
+using single_store_t = single_store<double, vecmem::vector, geometry_context>;
+using single_store_dev_t =
+    single_store<double, vecmem::device_vector, geometry_context>;
 /// @}
 
 // Tuple container test

tests/unit_tests/device/cuda/transform_store_cuda_kernel.hpp

@@ -24,10 +24,11 @@ using scalar = dscalar<test_algebra>;
 using point3 = dpoint3D<test_algebra>;
 using transform3 = dtransform3D<test_algebra>;
 
-using host_transform_store_t = single_store<transform3, vecmem::vector>;
+using host_transform_store_t =
+    single_store<transform3, vecmem::vector, geometry_context>;
 
 using device_transform_store_t =
-    single_store<transform3, vecmem::device_vector>;
+    single_store<transform3, vecmem::device_vector, geometry_context>;
 
 void transform_test(vecmem::data::vector_view<point3> input_data,
                     typename host_transform_store_t::view_type store_data,