single_store: added support for handling multiple contexts on the host (#939)

Added support for handling multiple contexts on the host

Author: tsulaiav

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/include/detray/test/common/detail/register_checks.hpp

@@ -19,8 +19,8 @@ template <template <typename> class check_t, typename detector_t,
           typename config_t = typename check_t<detector_t>::config>
 void register_checks(const detector_t &det,
                      const typename detector_t::name_map &vol_names,
-                     const config_t &cfg = {},
-                     const typename detector_t::geometry_context &gctx = {}) {
+                     const config_t &cfg,
+                     const typename detector_t::geometry_context &gctx) {
 
     const char *test_name = cfg.name().c_str();
     if (!test_name) {

tests/integration_tests/cpu/detectors/telescope_detector_navigation.cpp

@@ -51,11 +51,13 @@ int main(int argc, char **argv) {
         build_telescope_detector<test_algebra>(host_mr, tel_cfg);
 
     auto white_board = std::make_shared<test::whiteboard>();
+    tel_detector_t::geometry_context ctx{};
 
     // General data consistency of the detector
     test::consistency_check<tel_detector_t>::config cfg_cons{};
     detail::register_checks<test::consistency_check>(
-        tel_det, tel_names, cfg_cons.name("telescope_detector_consistency"));
+        tel_det, tel_names, cfg_cons.name("telescope_detector_consistency"),
+        ctx);
 
     // Navigation link consistency, discovered by ray intersection
     test::ray_scan<tel_detector_t>::config cfg_ray_scan{};
@@ -67,7 +69,8 @@ int main(int argc, char **argv) {
     cfg_ray_scan.track_generator().theta_range(0.f,
                                                0.25f * constant<scalar>::pi_4);
 
-    detail::register_checks<test::ray_scan>(tel_det, tel_names, cfg_ray_scan);
+    detail::register_checks<test::ray_scan>(tel_det, tel_names, cfg_ray_scan,
+                                            ctx);
 
     // Comparison of straight line navigation with ray scan
     test::straight_line_navigation<tel_detector_t>::config cfg_str_nav{};
@@ -80,7 +83,7 @@ int main(int argc, char **argv) {
         static_cast<float>(mask_tolerance[1]);
 
     detail::register_checks<test::straight_line_navigation>(tel_det, tel_names,
-                                                            cfg_str_nav);
+                                                            cfg_str_nav, ctx);
 
     // Navigation link consistency, discovered by helix intersection
     test::helix_scan<tel_detector_t>::config cfg_hel_scan{};
@@ -95,7 +98,8 @@ int main(int argc, char **argv) {
     cfg_hel_scan.track_generator().theta_range(0.f,
                                                0.25f * constant<scalar>::pi_4);
 
-    detail::register_checks<test::helix_scan>(tel_det, tel_names, cfg_hel_scan);
+    detail::register_checks<test::helix_scan>(tel_det, tel_names, cfg_hel_scan,
+                                              ctx);
 
     // Comparison of navigation in a constant B-field with helix
     test::helix_navigation<tel_detector_t>::config cfg_hel_nav{};
@@ -105,7 +109,7 @@ int main(int argc, char **argv) {
         -100.f * unit<float>::um;
 
     detail::register_checks<test::helix_navigation>(tel_det, tel_names,
-                                                    cfg_hel_nav);
+                                                    cfg_hel_nav, ctx);
 
     // Run the material validation
     test::material_scan<tel_detector_t>::config mat_scan_cfg{};
@@ -117,7 +121,7 @@ int main(int argc, char **argv) {
 
     // Record the material using a ray scan
     detail::register_checks<test::material_scan>(tel_det, tel_names,
-                                                 mat_scan_cfg);
+                                                 mat_scan_cfg, ctx);
 
     // Now trace the material during navigation and compare
     test::material_validation<tel_detector_t>::config mat_val_cfg{};
@@ -126,7 +130,7 @@ int main(int argc, char **argv) {
     mat_val_cfg.propagation() = cfg_str_nav.propagation();
 
     detail::register_checks<test::material_validation>(tel_det, tel_names,
-                                                       mat_val_cfg);
+                                                       mat_val_cfg, ctx);
 
     // Run the checks
     return RUN_ALL_TESTS();