Cache linearised circles in seed finding

This commit improves the performance of our seed finding by computing
the linear circles for the doublets only once rather than computing them
multiple times.

Author: stephenswat

device/common/include/traccc/edm/device/triplet_counter.hpp

@@ -47,6 +47,9 @@ struct triplet_counter {
     /// The position in which these triplets will be added
     unsigned int posTriplets = 0;
 
+    /// Index of the bottom-middle doublet
+    unsigned int m_botMidIdx = 0;
+
 };  // struct triplet_counter
 
 /// Declare all triplet counter collection types

device/common/include/traccc/seeding/device/count_triplets.hpp

@@ -48,7 +48,9 @@ inline void count_triplets(
     const device_doublet_collection_types::const_view& mid_bot_doublet_view,
     const device_doublet_collection_types::const_view& mid_top_doublet_view,
     triplet_counter_spM_collection_types::view spM_tc,
-    triplet_counter_collection_types::view mb_tc);
+    triplet_counter_collection_types::view mb_tc,
+    vecmem::data::vector_view<const lin_circle> mid_bot_circles,
+    vecmem::data::vector_view<const lin_circle> mid_top_circles);
 
 }  // namespace traccc::device
 

device/common/include/traccc/seeding/device/find_triplets.hpp

@@ -49,6 +49,8 @@ inline void find_triplets(
     const device_doublet_collection_types::const_view& mid_top_doublet_view,
     const triplet_counter_spM_collection_types::const_view& spM_tc_view,
     const triplet_counter_collection_types::const_view& tc_view,
+    vecmem::data::vector_view<const lin_circle> mid_bot_circle_view,
+    vecmem::data::vector_view<const lin_circle> mid_top_circle_view,
     device_triplet_collection_types::view triplet_view);
 
 }  // namespace traccc::device

device/common/include/traccc/seeding/device/impl/count_triplets.ipp

@@ -24,7 +24,9 @@ inline void count_triplets(
     const device_doublet_collection_types::const_view& mid_bot_doublet_view,
     const device_doublet_collection_types::const_view& mid_top_doublet_view,
     triplet_counter_spM_collection_types::view spM_tc_view,
-    triplet_counter_collection_types::view mb_tc_view) {
+    triplet_counter_collection_types::view mb_tc_view,
+    vecmem::data::vector_view<const lin_circle> mid_bot_circle_view,
+    vecmem::data::vector_view<const lin_circle> mid_top_circle_view) {
 
     // Create device copy of input parameters
     const device_doublet_collection_types::const_device mid_bot_doublet_device(
@@ -43,6 +45,10 @@ inline void count_triplets(
     const device_doublet_collection_types::const_device mid_top_doublet_device(
         mid_top_doublet_view);
     const doublet_counter_collection_types::const_device dc_device(dc_view);
+    const vecmem::device_vector<const lin_circle> mid_bot_circles(
+        mid_bot_circle_view);
+    const vecmem::device_vector<const lin_circle> mid_top_circles(
+        mid_top_circle_view);
 
     // Create device copy of output parameterss
     triplet_counter_collection_types::device mb_triplet_counter(mb_tc_view);
@@ -61,24 +67,16 @@ inline void count_triplets(
     const edm::spacepoint_collection::const_device::const_proxy_type spM =
         spacepoints.at(sp_device.bin(spM_loc.bin_idx)[spM_loc.sp_idx]);
     const sp_location spB_loc = mid_bot.sp2;
-    // bottom spacepoint
-    const edm::spacepoint_collection::const_device::const_proxy_type spB =
-        spacepoints.at(sp_device.bin(spB_loc.bin_idx)[spB_loc.sp_idx]);
 
     // Apply the conformal transformation to middle-bot doublet
-    traccc::lin_circle lb = doublet_finding_helper::transform_coordinates<
-        details::spacepoint_type::bottom>(spM, spB);
+    const traccc::lin_circle lb = mid_bot_circles.at(globalIndex);
 
     // Calculate some physical quantities required for triplet compatibility
     // check
     scalar iSinTheta2 = static_cast<scalar>(1.) + lb.cotTheta() * lb.cotTheta();
     scalar scatteringInRegion2 = config.maxScatteringAngle2 * iSinTheta2;
     scatteringInRegion2 *= config.sigmaScattering * config.sigmaScattering;
 
-    // These two quantities are used as output parameters in
-    // triplet_finding_helper::isCompatible but their values are irrelevant
-    scalar curvature, impact_parameter;
-
     // find the reference (start) index of the mid-top doublet container
     // item vector, where the doublets are recorded
     const unsigned int mt_start_idx = doublet_counts.m_posMidTop;
@@ -87,18 +85,16 @@ inline void count_triplets(
     // number of triplets per middle-bot doublet
     unsigned int num_triplets_per_mb = 0;
 
-    // iterate over mid-top doublets
-    for (unsigned int i = mt_start_idx; i < mt_end_idx; ++i) {
-        const traccc::sp_location spT_loc = mid_top_doublet_device[i].sp2;
-
-        const edm::spacepoint_collection::const_device::const_proxy_type spT =
-            spacepoints.at(sp_device.bin(spT_loc.bin_idx)[spT_loc.sp_idx]);
+    const unsigned int num_mt = mt_end_idx - mt_start_idx;
 
+    // iterate over mid-top doublets
+    for (unsigned int ri = 0; ri < num_mt; ++ri) {
         // Apply the conformal transformation to middle-top doublet
-        traccc::lin_circle lt = doublet_finding_helper::transform_coordinates<
-            details::spacepoint_type::top>(spM, spT);
+        const lin_circle& lt = mid_top_circles.at(ri + mt_start_idx);
 
-        // Check if mid-bot and mid-top doublets can form a triplet
+        // These two quantities are used as output parameters in
+        // triplet_finding_helper::isCompatible but their values are irrelevant
+        scalar curvature, impact_parameter;
         if (triplet_finding_helper::isCompatible(
                 spM, lb, lt, config, iSinTheta2, scatteringInRegion2, curvature,
                 impact_parameter)) {
@@ -114,8 +110,9 @@ inline void count_triplets(
         const unsigned int posTriplets =
             nTriplets.fetch_add(num_triplets_per_mb);
 
-        mb_triplet_counter.push_back(
-            {spB_loc, counter_link, num_triplets_per_mb, posTriplets});
+        mb_triplet_counter.push_back({spB_loc, counter_link,
+                                      num_triplets_per_mb, posTriplets,
+                                      globalIndex});
     }
 }
 

device/common/include/traccc/seeding/device/impl/find_triplets.ipp

@@ -25,6 +25,8 @@ inline void find_triplets(
     const device_doublet_collection_types::const_view& mid_top_doublet_view,
     const triplet_counter_spM_collection_types::const_view& spM_tc_view,
     const triplet_counter_collection_types::const_view& tc_view,
+    vecmem::data::vector_view<const lin_circle> mid_bot_circle_view,
+    vecmem::data::vector_view<const lin_circle> mid_top_circle_view,
     device_triplet_collection_types::view triplet_view) {
 
     // Check if anything needs to be done.
@@ -44,6 +46,10 @@ inline void find_triplets(
     const traccc::details::spacepoint_grid_types::const_device sp_grid(sp_view);
     const triplet_counter_spM_collection_types::const_device triplet_counts_spM(
         spM_tc_view);
+    const vecmem::device_vector<const lin_circle> mid_bot_circles(
+        mid_bot_circle_view);
+    const vecmem::device_vector<const lin_circle> mid_top_circles(
+        mid_top_circle_view);
 
     // Get the current work item information
     const triplet_counter mid_bot_counter = triplet_counts.at(globalIndex);
@@ -53,22 +59,17 @@ inline void find_triplets(
         doublet_counts.at(mid_bot_counter.spM_counter_link);
 
     const sp_location spM_loc = spM_counter.spM;
-    const sp_location spB_loc = mid_bot_counter.spB;
 
     // middle spacepoint
     const edm::spacepoint_collection::const_device::const_proxy_type spM =
         spacepoints.at(sp_grid.bin(spM_loc.bin_idx)[spM_loc.sp_idx]);
 
-    // bottom spacepoint
-    const edm::spacepoint_collection::const_device::const_proxy_type spB =
-        spacepoints.at(sp_grid.bin(spB_loc.bin_idx)[spB_loc.sp_idx]);
-
     // Set up the device result collection
     device_triplet_collection_types::device triplets(triplet_view);
 
     // Apply the conformal transformation to middle-bot doublet
-    const traccc::lin_circle lb = doublet_finding_helper::transform_coordinates<
-        details::spacepoint_type::bottom>(spM, spB);
+    const traccc::lin_circle lb =
+        mid_bot_circles.at(mid_bot_counter.m_botMidIdx);
 
     // Calculate some physical quantities required for triplet compatibility
     // check
@@ -77,10 +78,6 @@ inline void find_triplets(
                                        config.sigmaScattering *
                                        config.sigmaScattering;
 
-    // These two quantities are used as output parameters in
-    // triplet_finding_helper::isCompatible but their values are irrelevant
-    scalar curvature, impact_parameter;
-
     // find the reference (start) index of the mid-top doublet collection
     // item vector, where the doublets are recorded
     const unsigned int mt_start_idx = doublet_count.m_posMidTop;
@@ -93,24 +90,18 @@ inline void find_triplets(
 
     // iterate over mid-top doublets
     for (unsigned int i = mt_start_idx; i < mt_end_idx; ++i) {
-        const sp_location spT_loc = mid_top_doublet_device[i].sp2;
-
-        const edm::spacepoint_collection::const_device::const_proxy_type spT =
-            spacepoints.at(sp_grid.bin(spT_loc.bin_idx)[spT_loc.sp_idx]);
-
         // Apply the conformal transformation to middle-top doublet
-        const traccc::lin_circle lt =
-            doublet_finding_helper::transform_coordinates<
-                details::spacepoint_type::top>(spM, spT);
+        const traccc::lin_circle& lt = mid_top_circles.at(i);
 
         // Check if mid-bot and mid-top doublets can form a triplet
+        scalar curvature, impact_parameter;
         if (triplet_finding_helper::isCompatible(
                 spM, lb, lt, config, iSinTheta2, scatteringInRegion2, curvature,
                 impact_parameter)) {
 
             // Add triplet to jagged vector
             triplets.at(posTriplets++) = device_triplet(
-                {spT_loc, globalIndex, curvature,
+                {mid_top_doublet_device[i].sp2, globalIndex, curvature,
                  -impact_parameter * filter_config.impactWeightFactor,
                  lb.Zo()});
         }

device/cuda/src/seeding/seed_finding.cu

@@ -20,6 +20,7 @@
 #include "traccc/edm/device/doublet_counter.hpp"
 #include "traccc/edm/device/seeding_global_counter.hpp"
 #include "traccc/edm/device/triplet_counter.hpp"
+#include "traccc/seeding/detail/spacepoint_type.hpp"
 #include "traccc/seeding/device/count_doublets.hpp"
 #include "traccc/seeding/device/count_triplets.hpp"
 #include "traccc/seeding/device/find_doublets.hpp"
@@ -65,20 +66,93 @@ __global__ void find_doublets(
                           sp_grid, doublet_counter, mb_doublets, mt_doublets);
 }
 
+__global__ void make_mid_bot_lincircles(
+    device::device_doublet_collection_types::const_view mb_doublet_view,
+    device::doublet_counter_collection_types::const_view doublet_count_view,
+    edm::spacepoint_collection::const_view spacepoint_view,
+    traccc::details::spacepoint_grid_types::const_view sp_grid_view,
+    vecmem::data::vector_view<lin_circle> out_view) {
+
+    const device::device_doublet_collection_types::const_device doublets(
+        mb_doublet_view);
+    const device::doublet_counter_collection_types::const_device doublet_counts(
+        doublet_count_view);
+    const edm::spacepoint_collection::const_device spacepoints(spacepoint_view);
+    traccc::details::spacepoint_grid_types::const_device sp_grid(sp_grid_view);
+    vecmem::device_vector<lin_circle> out(out_view);
+
+    unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid >= doublets.size()) {
+        return;
+    }
+
+    const device::device_doublet dub = doublets.at(tid);
+    const unsigned int counter_link = dub.counter_link;
+    const device::doublet_counter count = doublet_counts.at(counter_link);
+    const sp_location spM_loc = count.m_spM;
+    const edm::spacepoint_collection::const_device::const_proxy_type spM =
+        spacepoints.at(sp_grid.bin(spM_loc.bin_idx)[spM_loc.sp_idx]);
+    const sp_location spB_loc = dub.sp2;
+    const edm::spacepoint_collection::const_device::const_proxy_type spB =
+        spacepoints.at(sp_grid.bin(spB_loc.bin_idx)[spB_loc.sp_idx]);
+
+    out.at(tid) = doublet_finding_helper::transform_coordinates<
+        traccc::details::spacepoint_type::bottom>(spM, spB);
+}
+
+__global__ void make_mid_top_lincircles(
+    device::device_doublet_collection_types::const_view mt_doublet_view,
+    device::doublet_counter_collection_types::const_view doublet_count_view,
+    edm::spacepoint_collection::const_view spacepoint_view,
+    traccc::details::spacepoint_grid_types::const_view sp_grid_view,
+    vecmem::data::vector_view<lin_circle> out_view) {
+
+    const device::device_doublet_collection_types::const_device doublets(
+        mt_doublet_view);
+    const device::doublet_counter_collection_types::const_device doublet_counts(
+        doublet_count_view);
+    const edm::spacepoint_collection::const_device spacepoints(spacepoint_view);
+    traccc::details::spacepoint_grid_types::const_device sp_grid(sp_grid_view);
+    vecmem::device_vector<lin_circle> out(out_view);
+
+    unsigned int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid >= doublets.size()) {
+        return;
+    }
+
+    const device::device_doublet dub = doublets.at(tid);
+    const unsigned int counter_link = dub.counter_link;
+    const device::doublet_counter count = doublet_counts.at(counter_link);
+    const sp_location spM_loc = count.m_spM;
+    const edm::spacepoint_collection::const_device::const_proxy_type spM =
+        spacepoints.at(sp_grid.bin(spM_loc.bin_idx)[spM_loc.sp_idx]);
+    const sp_location spT_loc = dub.sp2;
+    const edm::spacepoint_collection::const_device::const_proxy_type spT =
+        spacepoints.at(sp_grid.bin(spT_loc.bin_idx)[spT_loc.sp_idx]);
+
+    out.at(tid) = doublet_finding_helper::transform_coordinates<
+        traccc::details::spacepoint_type::top>(spM, spT);
+}
+
 /// CUDA kernel for running @c traccc::device::count_triplets
-__global__ void count_triplets(
+__global__ __launch_bounds__(128) void count_triplets(
     seedfinder_config config,
     edm::spacepoint_collection::const_view spacepoints,
     traccc::details::spacepoint_grid_types::const_view sp_grid,
     device::doublet_counter_collection_types::const_view doublet_counter,
     device::device_doublet_collection_types::const_view mb_doublets,
     device::device_doublet_collection_types::const_view mt_doublets,
     device::triplet_counter_spM_collection_types::view spM_counter,
-    device::triplet_counter_collection_types::view midBot_counter) {
+    device::triplet_counter_collection_types::view midBot_counter,
+    vecmem::data::vector_view<const lin_circle> midBot_circles,
+    vecmem::data::vector_view<const lin_circle> midTop_circles) {
 
     device::count_triplets(details::global_index1(), config, spacepoints,
                            sp_grid, doublet_counter, mb_doublets, mt_doublets,
-                           spM_counter, midBot_counter);
+                           spM_counter, midBot_counter, midBot_circles,
+                           midTop_circles);
 }
 
 /// CUDA kernel for running @c traccc::device::reduce_triplet_counts
@@ -92,19 +166,22 @@ __global__ void reduce_triplet_counts(
 }
 
 /// CUDA kernel for running @c traccc::device::find_triplets
-__global__ void find_triplets(
+__global__ __launch_bounds__(128) void find_triplets(
     seedfinder_config config, seedfilter_config filter_config,
     edm::spacepoint_collection::const_view spacepoints,
     traccc::details::spacepoint_grid_types::const_view sp_grid,
     device::doublet_counter_collection_types::const_view doublet_counter,
     device::device_doublet_collection_types::const_view mt_doublets,
     device::triplet_counter_spM_collection_types::const_view spM_tc,
     device::triplet_counter_collection_types::const_view midBot_tc,
+    vecmem::data::vector_view<const lin_circle> midBot_circles,
+    vecmem::data::vector_view<const lin_circle> midTop_circles,
     device::device_triplet_collection_types::view triplet_view) {
 
     device::find_triplets(details::global_index1(), config, filter_config,
                           spacepoints, sp_grid, doublet_counter, mt_doublets,
-                          spM_tc, midBot_tc, triplet_view);
+                          spM_tc, midBot_tc, midBot_circles, midTop_circles,
+                          triplet_view);
 }
 
 /// CUDA kernel for running @c traccc::device::update_triplet_weights
@@ -252,6 +329,33 @@ edm::seed_collection::buffer seed_finding::operator()(
             doublet_counter_buffer, doublet_buffer_mb, doublet_buffer_mt);
     TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
 
+    vecmem::data::vector_buffer<lin_circle> midBotLinCircles{
+        globalCounter_host->m_nMidBot, m_mr.main};
+    m_copy.setup(midBotLinCircles)->wait();
+    vecmem::data::vector_buffer<lin_circle> midTopLinCircles{
+        globalCounter_host->m_nMidTop, m_mr.main};
+    m_copy.setup(midBotLinCircles)->wait();
+
+    {
+        const unsigned int nThreads = 128;
+        const unsigned int nMidBotBlocks =
+            (globalCounter_host->m_nMidBot + nThreads - 1) / nThreads;
+        const unsigned int nMidTopBlocks =
+            (globalCounter_host->m_nMidTop + nThreads - 1) / nThreads;
+
+        kernels::
+            make_mid_bot_lincircles<<<nMidBotBlocks, nThreads, 0, stream>>>(
+                doublet_buffer_mb, doublet_counter_buffer, spacepoints_view,
+                g2_view, midBotLinCircles);
+        TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
+
+        kernels::
+            make_mid_top_lincircles<<<nMidTopBlocks, nThreads, 0, stream>>>(
+                doublet_buffer_mt, doublet_counter_buffer, spacepoints_view,
+                g2_view, midTopLinCircles);
+        TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
+    }
+
     // Set up the triplet counter buffers
     device::triplet_counter_spM_collection_types::buffer
         triplet_counter_spM_buffer = {doublet_counter_buffer_size, m_mr.main};
@@ -265,7 +369,7 @@ edm::seed_collection::buffer seed_finding::operator()(
 
     // Calculate the number of threads and thread blocks to run the doublet
     // counting kernel for.
-    const unsigned int nTripletCountThreads = m_warp_size * 2;
+    const unsigned int nTripletCountThreads = 128;
     const unsigned int nTripletCountBlocks =
         (globalCounter_host->m_nMidBot + nTripletCountThreads - 1) /
         nTripletCountThreads;
@@ -275,7 +379,7 @@ edm::seed_collection::buffer seed_finding::operator()(
                               stream>>>(
         m_seedfinder_config, spacepoints_view, g2_view, doublet_counter_buffer,
         doublet_buffer_mb, doublet_buffer_mt, triplet_counter_spM_buffer,
-        triplet_counter_midBot_buffer);
+        triplet_counter_midBot_buffer, midBotLinCircles, midTopLinCircles);
     TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
 
     // Calculate the number of threads and thread blocks to run the triplet
@@ -309,7 +413,7 @@ edm::seed_collection::buffer seed_finding::operator()(
 
     // Calculate the number of threads and thread blocks to run the triplet
     // finding kernel for.
-    const unsigned int nTripletFindThreads = m_warp_size * 2;
+    const unsigned int nTripletFindThreads = 128;
     const unsigned int nTripletFindBlocks =
         (m_copy.get_size(triplet_counter_midBot_buffer) + nTripletFindThreads -
          1) /
@@ -321,7 +425,7 @@ edm::seed_collection::buffer seed_finding::operator()(
             m_seedfinder_config, m_seedfilter_config, spacepoints_view, g2_view,
             doublet_counter_buffer, doublet_buffer_mt,
             triplet_counter_spM_buffer, triplet_counter_midBot_buffer,
-            triplet_buffer);
+            midBotLinCircles, midTopLinCircles, triplet_buffer);
     TRACCC_CUDA_ERROR_CHECK(cudaGetLastError());
 
     // Calculate the number of threads and thread blocks to run the weight