exp: simplify measurement access (#1168)

* Simplify measurement lookup host

* Simplify measurement access device

* Apply to other backedns and remove kernel

Author: niermann999

core/include/traccc/edm/measurement.hpp

@@ -96,6 +96,20 @@ struct measurement_sort_comp {
     }
 };
 
+struct measurement_sf_comp {
+    template <typename sf_descriptor_t>
+    TRACCC_HOST_DEVICE bool operator()(const sf_descriptor_t sf_desc,
+                                       const measurement& rhs) {
+        return sf_desc.barcode() < rhs.surface_link;
+    }
+
+    template <typename sf_descriptor_t>
+    TRACCC_HOST_DEVICE bool operator()(const measurement& lhs,
+                                       const sf_descriptor_t sf_desc) {
+        return lhs.surface_link < sf_desc.barcode();
+    }
+};
+
 struct measurement_equal_comp {
     TRACCC_HOST_DEVICE
     bool operator()(const measurement& lhs, const measurement& rhs) const {

core/include/traccc/finding/details/combinatorial_kalman_filter.hpp

@@ -87,41 +87,36 @@ combinatorial_kalman_filter(
     measurement_collection_types::const_device measurements{measurements_view};
 
     // Check contiguity of the measurements
-    assert(is_contiguous_on(measurement_module_projection(), measurements));
-
-    // Get copy of barcode uniques
-    std::vector<measurement> uniques;
-    uniques.resize(measurements.size());
-
-    std::vector<measurement>::iterator uniques_end =
-        std::unique_copy(measurements.begin(), measurements.end(),
-                         uniques.begin(), measurement_equal_comp());
-    const auto n_modules =
-        static_cast<unsigned int>(uniques_end - uniques.begin());
-
-    // Get upper bounds of unique elements
-    std::vector<unsigned int> upper_bounds;
-    upper_bounds.reserve(n_modules);
-    for (unsigned int i = 0; i < n_modules; i++) {
-        measurement_collection_types::const_device::iterator up =
-            std::upper_bound(measurements.begin(), measurements.end(),
-                             uniques[i], measurement_sort_comp());
-        upper_bounds.push_back(
+    assert(
+        host::is_contiguous_on(measurement_module_projection(), measurements));
+
+    // Get index ranges in the measurement container per detector surface
+    std::vector<unsigned int> meas_ranges;
+    meas_ranges.reserve(det.surfaces().size());
+
+    for (const auto& sf_desc : det.surfaces()) {
+        // Measurements can only be found on sensitive surfaces
+        if (!sf_desc.is_sensitive()) {
+            // Lower range index is the upper index of the previous range
+            // This is guaranteed by the measurement sorting step
+            const auto sf_idx{sf_desc.index()};
+            const unsigned int lo{sf_idx == 0u ? 0u : meas_ranges[sf_idx - 1u]};
+
+            // Hand the upper index of the previous range through to assign
+            // the lower index of the next valid range correctly
+            meas_ranges.push_back(lo);
+            continue;
+        }
+
+        auto up = std::upper_bound(measurements.begin(), measurements.end(),
+                                   sf_desc, measurement_sf_comp());
+        meas_ranges.push_back(
             static_cast<unsigned int>(std::distance(measurements.begin(), up)));
     }
+
     const measurement_collection_types::const_device::size_type n_meas =
         measurements.size();
 
-    // Get the number of measurements of each module
-    std::vector<unsigned int> sizes(n_modules);
-    std::adjacent_difference(upper_bounds.begin(), upper_bounds.end(),
-                             sizes.begin());
-
-    // Create barcode sequence
-    std::vector<detray::geometry::barcode> barcodes(n_modules);
-    std::transform(uniques.begin(), uniques_end, barcodes.begin(),
-                   [](const measurement& m) { return m.surface_link; });
-
     std::vector<std::vector<candidate_link>> links;
     links.resize(config.max_track_candidates_per_track);
 
@@ -225,48 +220,28 @@ combinatorial_kalman_filter(
                     sf);
             }
 
-            // Get barcode and measurements range on surface
-            const auto bcd = in_param.surface_link();
-            assert(!bcd.is_invalid());
-            std::pair<unsigned int, unsigned int> range;
-
-            // Find the corresponding index of bcd in barcode vector
-
-            const auto lo2 =
-                std::lower_bound(barcodes.begin(), barcodes.end(), bcd);
-
-            const auto bcd_id = std::distance(barcodes.begin(), lo2);
-
-            if (lo2 == barcodes.begin()) {
-                range.first = 0u;
-                range.second = upper_bounds[static_cast<std::size_t>(bcd_id)];
-            } else if (lo2 == barcodes.end()) {
-                range.first = 0u;
-                range.second = 0u;
-            } else {
-                range.first =
-                    upper_bounds[static_cast<std::size_t>(bcd_id - 1)];
-                range.second = upper_bounds[static_cast<std::size_t>(bcd_id)];
-            }
-
             /*****************************************************************
              * Find tracks (CKF)
              *****************************************************************/
 
+            // Iterate over the measurements for this surface
+            const auto sf_idx{sf.index()};
+            const unsigned int lo{sf_idx == 0u ? 0u : meas_ranges[sf_idx - 1]};
+            const unsigned int up{meas_ranges[sf_idx]};
+
             std::vector<std::tuple<candidate_link,
                                    bound_track_parameters<algebra_type>>>
                 best_links;
 
             // Iterate over the measurements
-            for (unsigned int item_id = range.first; item_id < range.second;
-                 item_id++) {
+            for (unsigned int meas_id = lo; meas_id < up; meas_id++) {
 
                 // The measurement on surface to handle.
-                const measurement& meas = measurements.at(item_id);
+                const measurement& meas = measurements.at(meas_id);
 
                 // Create a standalone track state object.
                 auto trk_state =
-                    edm::make_track_state<algebra_type>(measurements, item_id);
+                    edm::make_track_state<algebra_type>(measurements, meas_id);
 
                 const bool is_line = sf.template visit_mask<is_line_visitor>();
 
@@ -284,7 +259,7 @@ combinatorial_kalman_filter(
                     best_links.push_back(
                         {{.step = step,
                           .previous_candidate_idx = in_param_id,
-                          .meas_idx = item_id,
+                          .meas_idx = meas_id,
                           .seed_idx = orig_param_id,
                           .n_skipped = skip_counter,
                           .chi2 = chi2,

device/alpaka/src/finding/combinatorial_kalman_filter.hpp

@@ -25,7 +25,6 @@
 #include "traccc/finding/device/fill_finding_duplicate_removal_sort_keys.hpp"
 #include "traccc/finding/device/fill_finding_propagation_sort_keys.hpp"
 #include "traccc/finding/device/find_tracks.hpp"
-#include "traccc/finding/device/make_barcode_sequence.hpp"
 #include "traccc/finding/device/propagate_to_next_surface.hpp"
 #include "traccc/finding/device/remove_duplicates.hpp"
 #include "traccc/finding/finding_config.hpp"
@@ -40,19 +39,6 @@
 namespace traccc::alpaka::details {
 namespace kernels {
 
-/// Alpaka kernel functor for @c traccc::device::make_barcode_sequence
-struct make_barcode_sequence {
-    template <typename TAcc>
-    ALPAKA_FN_ACC void operator()(
-        TAcc const& acc,
-        const device::make_barcode_sequence_payload payload) const {
-
-        const device::global_index_t globalThreadIdx =
-            ::alpaka::getIdx<::alpaka::Grid, ::alpaka::Threads>(acc)[0];
-        device::make_barcode_sequence(globalThreadIdx, payload);
-    }
-};
-
 /// Alpaka kernel functor for @c traccc::device::apply_interaction
 template <typename detector_t>
 struct apply_interaction {
@@ -222,48 +208,21 @@ combinatorial_kalman_filter(
     const measurement_collection_types::const_view::size_type n_measurements =
         copy.get_size(measurements);
 
-    // Get copy of barcode uniques
-    measurement_collection_types::buffer uniques_buffer{n_measurements,
-                                                        mr.main};
-    copy.setup(uniques_buffer)->wait();
-    measurement_collection_types::device uniques(uniques_buffer);
-
-    measurement_collection_types::device::iterator uniques_end =
-        details::unique_copy(queue, mr, measurements.ptr(),
-                             measurements.ptr() + n_measurements,
-                             uniques.begin(), measurement_equal_comp());
-    const unsigned int n_modules =
-        static_cast<unsigned int>(uniques_end - uniques.begin());
-
-    // Get upper bounds of unique elements
-    vecmem::data::vector_buffer<unsigned int> upper_bounds_buffer{n_modules,
-                                                                  mr.main};
-    copy.setup(upper_bounds_buffer)->wait();
-    vecmem::device_vector<unsigned int> upper_bounds(upper_bounds_buffer);
-
-    details::upper_bound(queue, mr, measurements.ptr(),
-                         measurements.ptr() + n_measurements, uniques.begin(),
-                         uniques.begin() + n_modules, upper_bounds.begin(),
-                         measurement_sort_comp());
-
-    /*****************************************************************
-     * Kernel1: Create barcode sequence
-     *****************************************************************/
+    // Access the detector view as a detector object
+    detector_t device_det(det);
+    const unsigned int n_surfaces{device_det.surfaces().size()};
 
-    vecmem::data::vector_buffer<detray::geometry::barcode> barcodes_buffer{
-        n_modules, mr.main};
-    copy.setup(barcodes_buffer)->wait();
+    // Get upper bounds of measurement ranges per surface
+    vecmem::data::vector_buffer<unsigned int> meas_ranges_buffer{n_surfaces,
+                                                                 mr.main};
+    copy.setup(meas_ranges_buffer)->ignore();
+    vecmem::device_vector<unsigned int> measurement_ranges(meas_ranges_buffer);
 
-    {
-        Idx blocksPerGrid =
-            (barcodes_buffer.size() + threadsPerBlock - 1) / threadsPerBlock;
-        auto workDiv = makeWorkDiv<Acc>(blocksPerGrid, threadsPerBlock);
-
-        ::alpaka::exec<Acc>(queue, workDiv, kernels::make_barcode_sequence{},
-                            device::make_barcode_sequence_payload{
-                                uniques_buffer, barcodes_buffer});
-        ::alpaka::wait(queue);
-    }
+    // Get upper bounds of measurement ranges
+    details::upper_bound(
+        queue, mr, measurements.ptr(), measurements.ptr() + n_measurements,
+        device_det.surfaces().begin(), device_det.surfaces().end(),
+        measurement_ranges.begin(), measurement_sf_comp());
 
     const unsigned int n_seeds = copy.get_size(seeds);
 
@@ -384,8 +343,7 @@ combinatorial_kalman_filter(
                 .in_params_view = in_params_buffer,
                 .in_params_liveness_view = param_liveness_buffer,
                 .n_in_params = n_in_params,
-                .barcodes_view = barcodes_buffer,
-                .upper_bounds_view = upper_bounds_buffer,
+                .measurement_ranges_view = meas_ranges_buffer,
                 .links_view = links_buffer,
                 .prev_links_idx =
                     (step == 0 ? 0 : step_to_link_idx_map[step - 1]),

device/common/CMakeLists.txt

@@ -62,13 +62,11 @@ traccc_add_library( traccc_device_common device_common TYPE INTERFACE
    "include/traccc/finding/device/build_tracks.hpp"
    "include/traccc/finding/device/find_tracks.hpp"
    "include/traccc/finding/device/fill_finding_propagation_sort_keys.hpp"
-   "include/traccc/finding/device/make_barcode_sequence.hpp"
    "include/traccc/finding/device/propagate_to_next_surface.hpp"
    "include/traccc/finding/device/impl/apply_interaction.ipp"
    "include/traccc/finding/device/impl/build_tracks.ipp"
    "include/traccc/finding/device/impl/find_tracks.ipp"
    "include/traccc/finding/device/impl/fill_finding_propagation_sort_keys.ipp"
-   "include/traccc/finding/device/impl/make_barcode_sequence.ipp"
    "include/traccc/finding/device/impl/propagate_to_next_surface.ipp"
    # Track fitting funtions(s).
    "include/traccc/fitting/device/fit.hpp"

device/common/include/traccc/finding/device/find_tracks.hpp

@@ -60,15 +60,9 @@ struct find_tracks_payload {
     unsigned int n_in_params;
 
     /**
-     * @brief View object to the vector of barcodes for each measurement
+     * @brief View object to the vector of measurement index ranges per surface
      */
-    vecmem::data::vector_view<const detray::geometry::barcode> barcodes_view;
-
-    /**
-     * @brief View object to the vector of upper bounds of measurement indices
-     * per surface
-     */
-    vecmem::data::vector_view<const unsigned int> upper_bounds_view;
+    vecmem::data::vector_view<const unsigned int> measurement_ranges_view;
 
     /**
      * @brief View object to the link vector

device/common/include/traccc/finding/device/impl/find_tracks.ipp

@@ -61,10 +61,8 @@ TRACCC_HOST_DEVICE inline void find_tracks(
     vecmem::device_vector<candidate_link> tmp_links(payload.tmp_links_view);
     bound_track_parameters_collection_types::device tmp_params(
         payload.tmp_params_view);
-    vecmem::device_vector<const detray::geometry::barcode> barcodes(
-        payload.barcodes_view);
-    vecmem::device_vector<const unsigned int> upper_bounds(
-        payload.upper_bounds_view);
+    vecmem::device_vector<const unsigned int> meas_ranges(
+        payload.measurement_ranges_view);
     vecmem::device_vector<unsigned int> tips(payload.tips_view);
     vecmem::device_vector<unsigned int> tip_lengths(payload.tip_lengths_view);
     vecmem::device_vector<unsigned int> n_tracks_per_seed(
@@ -102,31 +100,16 @@ TRACCC_HOST_DEVICE inline void find_tracks(
          * Get the barcode of this thread's parameters, then find the first
          * measurement that matches it.
          */
-        const auto bcd = in_params.at(in_param_id).surface_link();
-        const auto lo = thrust::lower_bound(thrust::seq, barcodes.begin(),
-                                            barcodes.end(), bcd);
+        const unsigned int sf_idx{
+            in_params.at(in_param_id).surface_link().index()};
+        init_meas = sf_idx == 0u ? 0u : meas_ranges[sf_idx - 1];
+        num_meas = meas_ranges[sf_idx] - init_meas;
 
         /*
          * If we cannot find any corresponding measurements, set the number of
          * measurements to zero.
          */
-        if (lo == barcodes.end() || *lo != bcd) {
-            init_meas = 0;
-        }
-        /*
-         * If measurements are found, use the previously (outside this kernel)
-         * computed upper bound array to compute the range of measurements for
-         * this thread.
-         */
-        else {
-            const vecmem::device_vector<const unsigned int>::size_type bcd_id =
-                static_cast<
-                    vecmem::device_vector<const unsigned int>::size_type>(
-                    std::distance(barcodes.begin(), lo));
-
-            init_meas = lo == barcodes.begin() ? 0u : upper_bounds[bcd_id - 1];
-            num_meas = upper_bounds[bcd_id] - init_meas;
-        }
+        init_meas = (num_meas == 0u) ? 0u : init_meas;
     }
 
     /*