Merge pull request #90 from slava77/CMSSW_14_1_0_pre5/LSTb6-SoA-TCs

migrate TrackCandidate to use SoATemplate and PortableCollection
cms-sw · Oct 8, 2024 · 8338254 · 8338254
2 parents 3858cf3 + 69804f0
commit 8338254
Show file tree

Hide file tree

Showing 16 changed files with 304 additions and 348 deletions.
diff --git a/RecoTracker/LSTCore/interface/Constants.h b/RecoTracker/LSTCore/interface/Constants.h
@@ -2,6 +2,15 @@
 #define RecoTracker_LSTCore_interface_Constants_h
 
 #include "HeterogeneousCore/AlpakaInterface/interface/config.h"
+#include "DataFormats/Common/interface/StdArray.h"
+
+#if defined(FP16_Base)
+#if defined ALPAKA_ACC_GPU_CUDA_ENABLED
+#include <cuda_fp16.h>
+#elif defined ALPAKA_ACC_GPU_HIP_ENABLED
+#include <hip/hip_fp16.h>
+#endif
+#endif
 
 #ifdef CACHE_ALLOC
 #include "HeterogeneousCore/AlpakaInterface/interface/CachedBufAlloc.h"
@@ -55,6 +64,17 @@ namespace lst {
 
   constexpr unsigned int size_superbins = 45000;
 
+// Half precision wrapper functions.
+#if defined(FP16_Base)
+#define __F2H __float2half
+#define __H2F __half2float
+  typedef __half float FPX;
+#else
+#define __F2H
+#define __H2F
+  typedef float FPX;
+#endif
+
   // Defining the constant host device variables right up here
   // Currently pixel tracks treated as LSs with 2 double layers (IT layers 1+2 and 3+4) and 4 hits. To be potentially handled better in the future.
   struct Params_pLS {
@@ -74,8 +94,13 @@ namespace lst {
   };
   struct Params_pT5 {
     static constexpr int kLayers = 7, kHits = 14;
+    using ArrayU8xLayers = edm::StdArray<uint8_t, kLayers>;
+    using ArrayU16xLayers = edm::StdArray<uint16_t, kLayers>;
+    using ArrayUxHits = edm::StdArray<unsigned int, kHits>;
   };
 
+  using ArrayUx2 = edm::StdArray<unsigned int, 2>;
+
 }  //namespace lst
 
 #endif
diff --git a/RecoTracker/LSTCore/interface/TrackCandidatesHostCollection.h b/RecoTracker/LSTCore/interface/TrackCandidatesHostCollection.h
@@ -0,0 +1,10 @@
+#ifndef RecoTracker_LSTCore_interface_TrackCandidatesHostCollection_h
+#define RecoTracker_LSTCore_interface_TrackCandidatesHostCollection_h
+
+#include "RecoTracker/LSTCore/interface/TrackCandidatesSoA.h"
+#include "DataFormats/Portable/interface/PortableHostCollection.h"
+
+namespace lst {
+  using TrackCandidatesHostCollection = PortableHostCollection<TrackCandidatesSoA>;
+}  // namespace lst
+#endif
diff --git a/RecoTracker/LSTCore/interface/TrackCandidatesSoA.h b/RecoTracker/LSTCore/interface/TrackCandidatesSoA.h
@@ -0,0 +1,32 @@
+#ifndef RecoTracker_LSTCore_interface_TrackCandidatesSoA_h
+#define RecoTracker_LSTCore_interface_TrackCandidatesSoA_h
+
+#include <alpaka/alpaka.hpp>
+#include "DataFormats/Common/interface/StdArray.h"
+#include "DataFormats/SoATemplate/interface/SoALayout.h"
+
+#include "RecoTracker/LSTCore/interface/Constants.h"
+
+namespace lst {
+  GENERATE_SOA_LAYOUT(TrackCandidatesSoALayout,
+                      SOA_COLUMN(short, trackCandidateType),                  // 4-T5 5-pT3 7-pT5 8-pLS
+                      SOA_COLUMN(unsigned int, directObjectIndices),          // direct indices to each type containers
+                      SOA_COLUMN(ArrayUx2, objectIndices),                    // tracklet and  triplet indices
+                      SOA_COLUMN(Params_pT5::ArrayU8xLayers, logicalLayers),  //
+                      SOA_COLUMN(Params_pT5::ArrayUxHits, hitIndices),        //
+                      SOA_COLUMN(int, pixelSeedIndex),                        //
+                      SOA_COLUMN(Params_pT5::ArrayU16xLayers, lowerModuleIndices),  //
+                      SOA_COLUMN(FPX, centerX),                                     //
+                      SOA_COLUMN(FPX, centerY),                                     //
+                      SOA_COLUMN(FPX, radius),                                      //
+                      SOA_SCALAR(unsigned int, nTrackCandidates),                   //
+                      SOA_SCALAR(unsigned int, nTrackCandidatespT3),                //
+                      SOA_SCALAR(unsigned int, nTrackCandidatespT5),                //
+                      SOA_SCALAR(unsigned int, nTrackCandidatespLS),                //
+                      SOA_SCALAR(unsigned int, nTrackCandidatesT5))                 //
+
+  using TrackCandidatesSoA = TrackCandidatesSoALayout<>;
+  using TrackCandidates = TrackCandidatesSoA::View;
+  using TrackCandidatesConst = TrackCandidatesSoA::ConstView;
+}  // namespace lst
+#endif
diff --git a/RecoTracker/LSTCore/interface/alpaka/Constants.h b/RecoTracker/LSTCore/interface/alpaka/Constants.h
@@ -3,27 +3,10 @@
 
 #include "RecoTracker/LSTCore/interface/Constants.h"
 
-#if defined ALPAKA_ACC_GPU_CUDA_ENABLED
-#include <cuda_fp16.h>
-#elif defined ALPAKA_ACC_GPU_HIP_ENABLED
-#include <hip/hip_fp16.h>
-#endif
-
 namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
 
   using namespace ::lst;
 
-// Half precision wrapper functions.
-#if defined(FP16_Base)
-#define __F2H __float2half
-#define __H2F __half2float
-  typedef __half float FPX;
-#else
-#define __F2H
-#define __H2F
-  typedef float FPX;
-#endif
-
   Vec3D constexpr elementsPerThread(Vec3D::all(static_cast<Idx>(1)));
 
 // Needed for files that are compiled by g++ to not throw an error.

diff --git a/RecoTracker/LSTCore/interface/alpaka/LST.h b/RecoTracker/LSTCore/interface/alpaka/LST.h
@@ -66,10 +66,6 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
                       std::vector<float> const& ph2_z);
 
     void getOutput(Event& event);
-    std::vector<unsigned int> getHitIdxs(short trackCandidateType,
-                                         unsigned int TCIdx,
-                                         unsigned int const* TCHitIndices,
-                                         unsigned int const* hitIndices);
 
     // Input and output vectors
     std::vector<float> in_trkX_;

diff --git a/RecoTracker/LSTCore/interface/alpaka/TrackCandidatesDeviceCollection.h b/RecoTracker/LSTCore/interface/alpaka/TrackCandidatesDeviceCollection.h
@@ -0,0 +1,10 @@
+#ifndef RecoTracker_LSTCore_interface_TrackCandidatesDeviceCollection_h
+#define RecoTracker_LSTCore_interface_TrackCandidatesDeviceCollection_h
+
+#include "RecoTracker/LSTCore/interface/TrackCandidatesSoA.h"
+#include "DataFormats/Portable/interface/alpaka/PortableCollection.h"
+
+namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
+  using TrackCandidatesDeviceCollection = PortableCollection<TrackCandidatesSoA>;
+}  // namespace ALPAKA_ACCELERATOR_NAMESPACE::lst
+#endif