PixelSeeding/plugins/RiemannFitOnGPU.h

0001 //
0002 // Author: Felice Pantaleo, CERN
0003 //
0004
0005 #include <cstdint>
0006
0007 #include <cuda_runtime.h>
0008
0009 #include "CUDADataFormats/TrackingRecHit/interface/TrackingRecHitsUtilities.h"
0010 #include "CUDADataFormats/Track/interface/PixelTrackUtilities.h"
0011 #include "HeterogeneousCore/CUDAUtilities/interface/cudaCheck.h"
0012 #include "HeterogeneousCore/CUDAUtilities/interface/cuda_assert.h"
0013 #include "RecoLocalTracker/SiPixelRecHits/interface/pixelCPEforGPU.h"
0014 #include "RecoTracker/PixelTrackFitting/interface/RiemannFit.h"
0015
0016 #include "HelixFitOnGPU.h"
0017
0018 template <typename TrackerTraits>
0019 using Tuples = typename TrackSoA<TrackerTraits>::HitContainer;
0020 template <typename TrackerTraits>
0021 using OutputSoAView = TrackSoAView<TrackerTraits>;
0022 template <typename TrackerTraits>
0023 using TupleMultiplicity = caStructures::TupleMultiplicityT<TrackerTraits>;
0024
0025 template <int N, typename TrackerTraits>
0026 __global__ void kernel_FastFit(Tuples<TrackerTraits> const *__restrict__ foundNtuplets,
0027                                TupleMultiplicity<TrackerTraits> const *__restrict__ tupleMultiplicity,
0028                                uint32_t nHits,
0029                                TrackingRecHitSoAConstView<TrackerTraits> hh,
0030                                double *__restrict__ phits,
0031                                float *__restrict__ phits_ge,
0032                                double *__restrict__ pfast_fit,
0033                                uint32_t offset) {
0034   constexpr uint32_t hitsInFit = N;
0035
0036   assert(hitsInFit <= nHits);
0037
0038   assert(pfast_fit);
0039   assert(foundNtuplets);
0040   assert(tupleMultiplicity);
0041
0042   // look in bin for this hit multiplicity
0043   auto local_start = blockIdx.x * blockDim.x + threadIdx.x;
0044
0045 #ifdef RIEMANN_DEBUG
0046   if (0 == local_start)
0047     printf("%d Ntuple of size %d for %d hits to fit\n", tupleMultiplicity->size(nHits), nHits, hitsInFit);
0048 #endif
0049
0050   for (int local_idx = local_start, nt = riemannFit::maxNumberOfConcurrentFits; local_idx < nt;
0051        local_idx += gridDim.x * blockDim.x) {
0052     auto tuple_idx = local_idx + offset;
0053     if (tuple_idx >= tupleMultiplicity->size(nHits))
0054       break;
0055
0056     // get it from the ntuple container (one to one to helix)
0057     auto tkid = *(tupleMultiplicity->begin(nHits) + tuple_idx);
0058     assert(int(tkid) < foundNtuplets->nOnes());
0059
0060     assert(foundNtuplets->size(tkid) == nHits);
0061
0062     riemannFit::Map3xNd<N> hits(phits + local_idx);
0063     riemannFit::Map4d fast_fit(pfast_fit + local_idx);
0064     riemannFit::Map6xNf<N> hits_ge(phits_ge + local_idx);
0065
0066     // Prepare data structure
0067     auto const *hitId = foundNtuplets->begin(tkid);
0068     for (unsigned int i = 0; i < hitsInFit; ++i) {
0069       auto hit = hitId[i];
0070       float ge[6];
0071       hh.cpeParams().detParams(hh[hit].detectorIndex()).frame.toGlobal(hh[hit].xerrLocal(), 0, hh[hit].yerrLocal(), ge);
0072
0073       hits.col(i) << hh[hit].xGlobal(), hh[hit].yGlobal(), hh[hit].zGlobal();
0074       hits_ge.col(i) << ge[0], ge[1], ge[2], ge[3], ge[4], ge[5];
0075     }
0076     riemannFit::fastFit(hits, fast_fit);
0077
0078     // no NaN here....
0079     assert(fast_fit(0) == fast_fit(0));
0080     assert(fast_fit(1) == fast_fit(1));
0081     assert(fast_fit(2) == fast_fit(2));
0082     assert(fast_fit(3) == fast_fit(3));
0083   }
0084 }
0085
0086 template <int N, typename TrackerTraits>
0087 __global__ void kernel_CircleFit(TupleMultiplicity<TrackerTraits> const *__restrict__ tupleMultiplicity,
0088                                  uint32_t nHits,
0089                                  double bField,
0090                                  double *__restrict__ phits,
0091                                  float *__restrict__ phits_ge,
0092                                  double *__restrict__ pfast_fit_input,
0093                                  riemannFit::CircleFit *circle_fit,
0094                                  uint32_t offset) {
0095   assert(circle_fit);
0096   assert(N <= nHits);
0097
0098   // same as above...
0099
0100   // look in bin for this hit multiplicity
0101   auto local_start = blockIdx.x * blockDim.x + threadIdx.x;
0102   for (int local_idx = local_start, nt = riemannFit::maxNumberOfConcurrentFits; local_idx < nt;
0103        local_idx += gridDim.x * blockDim.x) {
0104     auto tuple_idx = local_idx + offset;
0105     if (tuple_idx >= tupleMultiplicity->size(nHits))
0106       break;
0107
0108     riemannFit::Map3xNd<N> hits(phits + local_idx);
0109     riemannFit::Map4d fast_fit(pfast_fit_input + local_idx);
0110     riemannFit::Map6xNf<N> hits_ge(phits_ge + local_idx);
0111
0112     riemannFit::VectorNd<N> rad = (hits.block(0, 0, 2, N).colwise().norm());
0113
0114     riemannFit::Matrix2Nd<N> hits_cov = riemannFit::Matrix2Nd<N>::Zero();
0115     riemannFit::loadCovariance2D(hits_ge, hits_cov);
0116
0117     circle_fit[local_idx] = riemannFit::circleFit(hits.block(0, 0, 2, N), hits_cov, fast_fit, rad, bField, true);
0118
0119 #ifdef RIEMANN_DEBUG
0120 //    auto tkid = *(tupleMultiplicity->begin(nHits) + tuple_idx);
0121 //  printf("kernelCircleFit circle.par(0,1,2): %d %f,%f,%f\n", tkid,
0122 //         circle_fit[local_idx].par(0), circle_fit[local_idx].par(1), circle_fit[local_idx].par(2));
0123 #endif
0124   }
0125 }
0126
0127 template <int N, typename TrackerTraits>
0128 __global__ void kernel_LineFit(TupleMultiplicity<TrackerTraits> const *__restrict__ tupleMultiplicity,
0129                                uint32_t nHits,
0130                                double bField,
0131                                OutputSoAView<TrackerTraits> results_view,
0132                                double *__restrict__ phits,
0133                                float *__restrict__ phits_ge,
0134                                double *__restrict__ pfast_fit_input,
0135                                riemannFit::CircleFit *__restrict__ circle_fit,
0136                                uint32_t offset) {
0137   assert(circle_fit);
0138   assert(N <= nHits);
0139
0140   // same as above...
0141
0142   // look in bin for this hit multiplicity
0143   auto local_start = (blockIdx.x * blockDim.x + threadIdx.x);
0144   for (int local_idx = local_start, nt = riemannFit::maxNumberOfConcurrentFits; local_idx < nt;
0145        local_idx += gridDim.x * blockDim.x) {
0146     auto tuple_idx = local_idx + offset;
0147     if (tuple_idx >= tupleMultiplicity->size(nHits))
0148       break;
0149
0150     // get it for the ntuple container (one to one to helix)
0151     int32_t tkid = *(tupleMultiplicity->begin(nHits) + tuple_idx);
0152
0153     riemannFit::Map3xNd<N> hits(phits + local_idx);
0154     riemannFit::Map4d fast_fit(pfast_fit_input + local_idx);
0155     riemannFit::Map6xNf<N> hits_ge(phits_ge + local_idx);
0156
0157     auto const &line_fit = riemannFit::lineFit(hits, hits_ge, circle_fit[local_idx], fast_fit, bField, true);
0158
0159     riemannFit::fromCircleToPerigee(circle_fit[local_idx]);
0160
0161     TracksUtilities<TrackerTraits>::copyFromCircle(results_view,
0162                                                    circle_fit[local_idx].par,
0163                                                    circle_fit[local_idx].cov,
0164                                                    line_fit.par,
0165                                                    line_fit.cov,
0166                                                    1.f / float(bField),
0167                                                    tkid);
0168     results_view[tkid].pt() = bField / std::abs(circle_fit[local_idx].par(2));
0169     results_view[tkid].eta() = asinhf(line_fit.par(0));
0170     results_view[tkid].chi2() = (circle_fit[local_idx].chi2 + line_fit.chi2) / (2 * N - 5);
0171
0172 #ifdef RIEMANN_DEBUG
0173     printf("kernelLineFit size %d for %d hits circle.par(0,1,2): %d %f,%f,%f\n",
0174            N,
0175            nHits,
0176            tkid,
0177            circle_fit[local_idx].par(0),
0178            circle_fit[local_idx].par(1),
0179            circle_fit[local_idx].par(2));
0180     printf("kernelLineFit line.par(0,1): %d %f,%f\n", tkid, line_fit.par(0), line_fit.par(1));
0181     printf("kernelLineFit chi2 cov %f/%f %e,%e,%e,%e,%e\n",
0182            circle_fit[local_idx].chi2,
0183            line_fit.chi2,
0184            circle_fit[local_idx].cov(0, 0),
0185            circle_fit[local_idx].cov(1, 1),
0186            circle_fit[local_idx].cov(2, 2),
0187            line_fit.cov(0, 0),
0188            line_fit.cov(1, 1));
0189 #endif
0190   }
0191 }