plugins/alpaka/HelixFit.h

0001 #ifndef RecoTracker_PixelSeeding_plugins_alpaka_HelixFit_h
0002 #define RecoTracker_PixelSeeding_plugins_alpaka_HelixFit_h
0003
0004 #include <alpaka/alpaka.hpp>
0005
0006 #include <Eigen/Core>
0007
0008 #include "DataFormats/TrackSoA/interface/alpaka/TrackUtilities.h"
0009 #include "DataFormats/TrackingRecHitSoA/interface/TrackingRecHitsSoA.h"
0010 #include "RecoTracker/PixelTrackFitting/interface/alpaka/FitResult.h"
0011 #include "Geometry/CommonTopologies/interface/SimplePixelTopology.h"
0012 #include "HeterogeneousCore/AlpakaInterface/interface/config.h"
0013 #include "RecoTracker/PixelSeeding/interface/CAGeometrySoA.h"
0014
0015 #include "CAStructures.h"
0016
0017 namespace riemannFit {
0018
0019   // TODO: Can this be taken from TrackerTraits or somewhere else?
0020   // in case of memory issue can be made smaller
0021   constexpr uint32_t maxNumberOfConcurrentFits = 32 * 1024;
0022   constexpr uint32_t stride = maxNumberOfConcurrentFits;
0023   using Matrix3x4d = Eigen::Matrix<double, 3, 4>;
0024   using Map3x4d = Eigen::Map<Matrix3x4d, 0, Eigen::Stride<3 * stride, stride> >;
0025   using Matrix6x4f = Eigen::Matrix<float, 6, 4>;
0026   using Map6x4f = Eigen::Map<Matrix6x4f, 0, Eigen::Stride<6 * stride, stride> >;
0027
0028   // hits
0029   template <int N>
0030   using Matrix3xNd = Eigen::Matrix<double, 3, N>;
0031   template <int N>
0032   using Map3xNd = Eigen::Map<Matrix3xNd<N>, 0, Eigen::Stride<3 * stride, stride> >;
0033   // errors
0034   template <int N>
0035   using Matrix6xNf = Eigen::Matrix<float, 6, N>;
0036   template <int N>
0037   using Map6xNf = Eigen::Map<Matrix6xNf<N>, 0, Eigen::Stride<6 * stride, stride> >;
0038   // fast fit
0039   using Map4d = Eigen::Map<Vector4d, 0, Eigen::InnerStride<stride> >;
0040
0041   template <auto Start, auto End, auto Inc, class F>  //a compile-time bounded for loop
0042   constexpr void rolling_fits(F &&f) {
0043     if constexpr (Start < End) {
0044       f(std::integral_constant<decltype(Start), Start>());
0045       rolling_fits<Start + Inc, End, Inc>(f);
0046     }
0047   }
0048
0049 }  // namespace riemannFit
0050
0051 namespace ALPAKA_ACCELERATOR_NAMESPACE {
0052
0053   template <typename TrackerTraits>
0054   class HelixFit {
0055   public:
0056     using HitView = ::reco::TrackingRecHitView;
0057     using HitConstView = ::reco::TrackingRecHitConstView;
0058     using OutputSoAView = ::reco::TrackSoAView;
0059     using OutputHitSoAView = ::reco::TrackHitSoAView;
0060
0061     using Tuples = caStructures::SequentialContainer;
0062     using TupleMultiplicity = caStructures::GenericContainer;
0063
0064     explicit HelixFit(float bf, bool fitNas4) : bField_(bf), fitNas4_(fitNas4) {}
0065     ~HelixFit() { deallocate(); }
0066
0067     void setBField(double bField) { bField_ = bField; }
0068     void launchRiemannKernels(const HitConstView &hv,
0069                               const ::reco::CAModulesConstView &fr,
0070                               uint32_t nhits,
0071                               uint32_t maxNumberOfTuples,
0072                               Queue &queue);
0073     void launchBrokenLineKernels(const HitConstView &hv,
0074                                  const ::reco::CAModulesConstView &fr,
0075                                  uint32_t nhits,
0076                                  uint32_t maxNumberOfTuples,
0077                                  Queue &queue);
0078
0079     void allocate(TupleMultiplicity const *tupleMultiplicity,
0080                   OutputSoAView &helix_fit_results,
0081                   Tuples const *__restrict__ foundNtuplets);
0082     void deallocate();
0083
0084   private:
0085     static constexpr uint32_t maxNumberOfConcurrentFits_ = riemannFit::maxNumberOfConcurrentFits;
0086
0087     // fowarded
0088     Tuples const *tuples_ = nullptr;
0089     TupleMultiplicity const *tupleMultiplicity_ = nullptr;
0090     OutputSoAView outputSoa_;
0091     float bField_;
0092
0093     const bool fitNas4_;
0094   };
0095
0096 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE
0097
0098 #endif  // RecoTracker_PixelSeeding_plugins_alpaka_HelixFit_h