PixelSeeding/plugins/HelixFitOnGPU.h

0001 #ifndef RecoTracker_PixelSeeding_plugins_HelixFitOnGPU_h
0002 #define RecoTracker_PixelSeeding_plugins_HelixFitOnGPU_h
0003
0004 #include "CUDADataFormats/Track/interface/PixelTrackUtilities.h"
0005 #include "CUDADataFormats/TrackingRecHit/interface/TrackingRecHitsUtilities.h"
0006 #include "RecoTracker/PixelTrackFitting/interface/FitResult.h"
0007 #include "Geometry/CommonTopologies/interface/SimplePixelTopology.h"
0008
0009 #include "CAStructures.h"
0010
0011 namespace riemannFit {
0012   // in case of memory issue can be made smaller
0013   constexpr uint32_t maxNumberOfConcurrentFits = 32 * 1024;
0014   constexpr uint32_t stride = maxNumberOfConcurrentFits;
0015   using Matrix3x4d = Eigen::Matrix<double, 3, 4>;
0016   using Map3x4d = Eigen::Map<Matrix3x4d, 0, Eigen::Stride<3 * stride, stride> >;
0017   using Matrix6x4f = Eigen::Matrix<float, 6, 4>;
0018   using Map6x4f = Eigen::Map<Matrix6x4f, 0, Eigen::Stride<6 * stride, stride> >;
0019
0020   // hits
0021   template <int N>
0022   using Matrix3xNd = Eigen::Matrix<double, 3, N>;
0023   template <int N>
0024   using Map3xNd = Eigen::Map<Matrix3xNd<N>, 0, Eigen::Stride<3 * stride, stride> >;
0025   // errors
0026   template <int N>
0027   using Matrix6xNf = Eigen::Matrix<float, 6, N>;
0028   template <int N>
0029   using Map6xNf = Eigen::Map<Matrix6xNf<N>, 0, Eigen::Stride<6 * stride, stride> >;
0030   // fast fit
0031   using Map4d = Eigen::Map<Vector4d, 0, Eigen::InnerStride<stride> >;
0032
0033   template <auto Start, auto End, auto Inc, class F>  //a compile-time bounded for loop
0034   constexpr void rolling_fits(F &&f) {
0035     if constexpr (Start < End) {
0036       f(std::integral_constant<decltype(Start), Start>());
0037       rolling_fits<Start + Inc, End, Inc>(f);
0038     }
0039   }
0040
0041 }  // namespace riemannFit
0042
0043 template <typename TrackerTraits>
0044 class HelixFitOnGPU {
0045 public:
0046   using TrackingRecHitSoAs = TrackingRecHitSoA<TrackerTraits>;
0047
0048   using HitView = TrackingRecHitSoAView<TrackerTraits>;
0049   using HitConstView = TrackingRecHitSoAConstView<TrackerTraits>;
0050
0051   using Tuples = typename TrackSoA<TrackerTraits>::HitContainer;
0052   using OutputSoAView = TrackSoAView<TrackerTraits>;
0053
0054   using TupleMultiplicity = caStructures::TupleMultiplicityT<TrackerTraits>;
0055
0056   explicit HelixFitOnGPU(float bf, bool fitNas4) : bField_(bf), fitNas4_(fitNas4) {}
0057   ~HelixFitOnGPU() { deallocateOnGPU(); }
0058
0059   void setBField(double bField) { bField_ = bField; }
0060   void launchRiemannKernels(const HitConstView &hv, uint32_t nhits, uint32_t maxNumberOfTuples, cudaStream_t cudaStream);
0061   void launchBrokenLineKernels(const HitConstView &hv,
0062                                uint32_t nhits,
0063                                uint32_t maxNumberOfTuples,
0064                                cudaStream_t cudaStream);
0065
0066   void launchRiemannKernelsOnCPU(const HitConstView &hv, uint32_t nhits, uint32_t maxNumberOfTuples);
0067   void launchBrokenLineKernelsOnCPU(const HitConstView &hv, uint32_t nhits, uint32_t maxNumberOfTuples);
0068
0069   void allocateOnGPU(TupleMultiplicity const *tupleMultiplicity, OutputSoAView &helix_fit_results);
0070   void deallocateOnGPU();
0071
0072 private:
0073   static constexpr uint32_t maxNumberOfConcurrentFits_ = riemannFit::maxNumberOfConcurrentFits;
0074
0075   // fowarded
0076   Tuples const *tuples_ = nullptr;
0077   TupleMultiplicity const *tupleMultiplicity_ = nullptr;
0078   OutputSoAView outputSoa_;
0079   float bField_;
0080
0081   const bool fitNas4_;
0082 };
0083
0084 #endif  // RecoTracker_PixelSeeding_plugins_HelixFitOnGPU_h