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 "RecoLocalTracker/SiPixelRecHits/interface/pixelCPEforDevice.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 TrackingRecHitSoAs = TrackingRecHitSoA<TrackerTraits>;
0057
0058     using HitView = TrackingRecHitSoAView<TrackerTraits>;
0059     using HitConstView = TrackingRecHitSoAConstView<TrackerTraits>;
0060
0061     using Tuples = typename reco::TrackSoA<TrackerTraits>::HitContainer;
0062     using OutputSoAView = reco::TrackSoAView<TrackerTraits>;
0063
0064     using TupleMultiplicity = caStructures::TupleMultiplicityT<TrackerTraits>;
0065
0066     using ParamsOnDevice = pixelCPEforDevice::ParamsOnDeviceT<TrackerTraits>;
0067
0068     explicit HelixFit(float bf, bool fitNas4) : bField_(bf), fitNas4_(fitNas4) {}
0069     ~HelixFit() { deallocate(); }
0070
0071     void setBField(double bField) { bField_ = bField; }
0072     void launchRiemannKernels(const HitConstView &hv,
0073                               ParamsOnDevice const *cpeParams,
0074                               uint32_t nhits,
0075                               uint32_t maxNumberOfTuples,
0076                               Queue &queue);
0077     void launchBrokenLineKernels(const HitConstView &hv,
0078                                  ParamsOnDevice const *cpeParams,
0079                                  uint32_t nhits,
0080                                  uint32_t maxNumberOfTuples,
0081                                  Queue &queue);
0082
0083     void allocate(TupleMultiplicity const *tupleMultiplicity, OutputSoAView &helix_fit_results);
0084     void deallocate();
0085
0086   private:
0087     static constexpr uint32_t maxNumberOfConcurrentFits_ = riemannFit::maxNumberOfConcurrentFits;
0088
0089     // fowarded
0090     Tuples const *tuples_ = nullptr;
0091     TupleMultiplicity const *tupleMultiplicity_ = nullptr;
0092     OutputSoAView outputSoa_;
0093     float bField_;
0094
0095     const bool fitNas4_;
0096   };
0097
0098 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE
0099
0100 #endif  // RecoTracker_PixelSeeding_plugins_alpaka_HelixFit_h