Add an option to disable the Riemann fit in the CA (cms-sw#187)

rovere · Oct 6, 2018 · 627820c · 627820c
1 parent 0e3f658
commit 627820c
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 43 deletions.
diff --git a/RecoPixelVertexing/PixelTriplets/plugins/CAHitNtupletHeterogeneousEDProducer.cc b/RecoPixelVertexing/PixelTriplets/plugins/CAHitNtupletHeterogeneousEDProducer.cc
@@ -54,31 +54,32 @@ class CAHitNtupletHeterogeneousEDProducer
 private:
   edm::EDGetTokenT<edm::OwnVector<TrackingRegion>> regionToken_;
 
-  edm::EDGetTokenT<HeterogeneousProduct> tGpuHits;
-  edm::EDGetTokenT<SiPixelRecHitCollectionNew> tCpuHits;
+  edm::EDGetTokenT<HeterogeneousProduct> gpuHits_;
+  edm::EDGetTokenT<SiPixelRecHitCollectionNew> cpuHits_;
 
   edm::RunningAverage localRA_;
   CAHitQuadrupletGeneratorGPU GPUGenerator_;
 
   bool emptyRegions = false;
   std::unique_ptr<RegionsSeedingHitSets> seedingHitSets_;
 
-  bool enableTransfer_;
-  bool enableConversion_;
+  const bool doRiemannFit_;
+  const bool enableConversion_;
+  const bool enableTransfer_;
 };
 
 CAHitNtupletHeterogeneousEDProducer::CAHitNtupletHeterogeneousEDProducer(
     const edm::ParameterSet &iConfig)
     : HeterogeneousEDProducer(iConfig),
       regionToken_(consumes<edm::OwnVector<TrackingRegion>>(
           iConfig.getParameter<edm::InputTag>("trackingRegions"))),
-      tGpuHits(consumesHeterogeneous(iConfig.getParameter<edm::InputTag>("heterogeneousPixelRecHitSrc"))),
-      tCpuHits(consumes<SiPixelRecHitCollectionNew>(iConfig.getParameter<edm::InputTag>("heterogeneousPixelRecHitSrc"))),
-      GPUGenerator_(iConfig, consumesCollector()) {
-
-  enableConversion_ = iConfig.getParameter<bool>("gpuEnableConversion");
-  enableTransfer_ = enableConversion_ || iConfig.getParameter<bool>("gpuEnableTransfer");
-
+      gpuHits_(consumesHeterogeneous(iConfig.getParameter<edm::InputTag>("heterogeneousPixelRecHitSrc"))),
+      cpuHits_(consumes<SiPixelRecHitCollectionNew>(iConfig.getParameter<edm::InputTag>("heterogeneousPixelRecHitSrc"))),
+      GPUGenerator_(iConfig, consumesCollector()),
+      doRiemannFit_(iConfig.getParameter<bool>("doRiemannFit")),
+      enableConversion_(iConfig.getParameter<bool>("gpuEnableConversion")),
+      enableTransfer_(enableConversion_ || iConfig.getParameter<bool>("gpuEnableTransfer"))
+{
   produces<RegionsSeedingHitSets>();
 }
 
@@ -88,9 +89,8 @@ void CAHitNtupletHeterogeneousEDProducer::fillDescriptions(
 
   desc.add<edm::InputTag>("doublets", edm::InputTag(""))->setComment("Not really used, kept to keep the python parameters");
   desc.add<edm::InputTag>("trackingRegions", edm::InputTag("globalTrackingRegionFromBeamSpot"));
-
   desc.add<edm::InputTag>("heterogeneousPixelRecHitSrc", edm::InputTag("siPixelRecHitsPreSplitting"));
-
+  desc.add<bool>("doRiemannFit", true);
   desc.add<bool>("gpuEnableTransfer", true);
   desc.add<bool>("gpuEnableConversion", true);
 
@@ -123,13 +123,10 @@ void CAHitNtupletHeterogeneousEDProducer::acquireGPUCuda(
 
   const TrackingRegion &region = regions[0];
 
-
   edm::Handle<siPixelRecHitsHeterogeneousProduct::GPUProduct> gh;
-  iEvent.getByToken<siPixelRecHitsHeterogeneousProduct::HeterogeneousPixelRecHit>(tGpuHits, gh);
+  iEvent.getByToken<siPixelRecHitsHeterogeneousProduct::HeterogeneousPixelRecHit>(gpuHits_, gh);
   auto const & gHits = *gh;
-//  auto nhits = gHits.nHits;
 
-  // move inside hitNtuplets???
   GPUGenerator_.buildDoublets(gHits,cudaStream.id());
 
   seedingHitSets_->reserve(regions.size(), localRA_.upper());
@@ -139,8 +136,7 @@ void CAHitNtupletHeterogeneousEDProducer::acquireGPUCuda(
         << "Creating ntuplets for " << regions.size()
         << " regions";
 
-  GPUGenerator_.hitNtuplets(region, gHits, iSetup, enableTransfer_, cudaStream.id());
-
+  GPUGenerator_.hitNtuplets(region, gHits, iSetup, doRiemannFit_, enableTransfer_, cudaStream.id());
 }
 
 void CAHitNtupletHeterogeneousEDProducer::produceGPUCuda(
@@ -155,7 +151,7 @@ void CAHitNtupletHeterogeneousEDProducer::produceGPUCuda(
     const auto &regions = *hregions;
 
     edm::Handle<SiPixelRecHitCollectionNew> gh;
-    iEvent.getByToken(tCpuHits, gh);
+    iEvent.getByToken(cpuHits_, gh);
     auto const & rechits = *gh;
 
     std::vector<OrderedHitSeeds> ntuplets(regions.size());

diff --git a/RecoPixelVertexing/PixelTriplets/plugins/CAHitQuadrupletGeneratorGPU.cc b/RecoPixelVertexing/PixelTriplets/plugins/CAHitQuadrupletGeneratorGPU.cc
@@ -85,12 +85,13 @@ void CAHitQuadrupletGeneratorGPU::hitNtuplets(
     TrackingRegion const& region,
     HitsOnCPU const& hh,
     edm::EventSetup const& es,
+    bool doRiemannFit,
     bool transferToCPU,
     cudaStream_t cudaStream)
 {
   hitsOnCPU = &hh;
   int index = 0;
-  launchKernels(region, index, hh, transferToCPU, cudaStream);
+  launchKernels(region, index, hh, doRiemannFit, transferToCPU, cudaStream);
 }
 
 void CAHitQuadrupletGeneratorGPU::fillResults(

diff --git a/RecoPixelVertexing/PixelTriplets/plugins/CAHitQuadrupletGeneratorGPU.cu b/RecoPixelVertexing/PixelTriplets/plugins/CAHitQuadrupletGeneratorGPU.cu
@@ -321,6 +321,7 @@ void CAHitQuadrupletGeneratorGPU::allocateOnGPU()
 
 void CAHitQuadrupletGeneratorGPU::launchKernels(const TrackingRegion &region,
                                                 int regionIndex, HitsOnCPU const & hh,
+                                                bool doRiemannFit,
                                                 bool transferToCPU,
                                                 cudaStream_t cudaStream)
 {
@@ -361,26 +362,28 @@ void CAHitQuadrupletGeneratorGPU::launchKernels(const TrackingRegion &region,
 
   // kernel_print_found_ntuplets<<<1, 1, 0, cudaStream>>>(d_foundNtupletsVec_[regionIndex], 10);
 
-  kernelFastFitAllHits<<<numberOfBlocks, 512, 0, cudaStream>>>(
-      d_foundNtupletsVec_[regionIndex], hh.gpu_d, 4, bField_, helix_fit_resultsGPU_,
-      hitsGPU_, hits_covGPU_, circle_fit_resultsGPU_, fast_fit_resultsGPU_,
-      line_fit_resultsGPU_);
-  cudaCheck(cudaGetLastError());
-
-  blockSize = 256;
-  numberOfBlocks = (maxNumberOfQuadruplets_ + blockSize - 1) / blockSize;
-
-  kernelCircleFitAllHits<<<numberOfBlocks, blockSize, 0, cudaStream>>>(
-      d_foundNtupletsVec_[regionIndex], 4, bField_, helix_fit_resultsGPU_,
-      hitsGPU_, hits_covGPU_, circle_fit_resultsGPU_, fast_fit_resultsGPU_,
-      line_fit_resultsGPU_);
-  cudaCheck(cudaGetLastError());
-
-  kernelLineFitAllHits<<<numberOfBlocks, blockSize, 0, cudaStream>>>(
-      d_foundNtupletsVec_[regionIndex], bField_, helix_fit_resultsGPU_,
-      hitsGPU_, hits_covGPU_, circle_fit_resultsGPU_, fast_fit_resultsGPU_,
-      line_fit_resultsGPU_);
-  cudaCheck(cudaGetLastError());
+  if (doRiemannFit) {
+    kernelFastFitAllHits<<<numberOfBlocks, 512, 0, cudaStream>>>(
+        d_foundNtupletsVec_[regionIndex], hh.gpu_d, 4, bField_, helix_fit_resultsGPU_,
+        hitsGPU_, hits_covGPU_, circle_fit_resultsGPU_, fast_fit_resultsGPU_,
+        line_fit_resultsGPU_);
+    cudaCheck(cudaGetLastError());
+
+    blockSize = 256;
+    numberOfBlocks = (maxNumberOfQuadruplets_ + blockSize - 1) / blockSize;
+
+    kernelCircleFitAllHits<<<numberOfBlocks, blockSize, 0, cudaStream>>>(
+        d_foundNtupletsVec_[regionIndex], 4, bField_, helix_fit_resultsGPU_,
+        hitsGPU_, hits_covGPU_, circle_fit_resultsGPU_, fast_fit_resultsGPU_,
+        line_fit_resultsGPU_);
+    cudaCheck(cudaGetLastError());
+
+    kernelLineFitAllHits<<<numberOfBlocks, blockSize, 0, cudaStream>>>(
+        d_foundNtupletsVec_[regionIndex], bField_, helix_fit_resultsGPU_,
+        hitsGPU_, hits_covGPU_, circle_fit_resultsGPU_, fast_fit_resultsGPU_,
+        line_fit_resultsGPU_);
+    cudaCheck(cudaGetLastError());
+  }
 
   if (transferToCPU) {
     cudaCheck(cudaMemcpyAsync(h_foundNtupletsVec_[regionIndex], d_foundNtupletsVec_[regionIndex],

diff --git a/RecoPixelVertexing/PixelTriplets/plugins/CAHitQuadrupletGeneratorGPU.h b/RecoPixelVertexing/PixelTriplets/plugins/CAHitQuadrupletGeneratorGPU.h
@@ -61,6 +61,7 @@ class CAHitQuadrupletGeneratorGPU {
 
     void hitNtuplets(const TrackingRegion &region, HitsOnCPU const & hh,
                      const edm::EventSetup& es,
+                     bool doRiemannFit,
                      bool transferToCPU,
                      cudaStream_t stream);
     void cleanup(cudaStream_t stream);
@@ -139,10 +140,9 @@ class CAHitQuadrupletGeneratorGPU {
         const bool enabled_;
     };
 
-    void  launchKernels(const TrackingRegion &, int, HitsOnCPU const & hh, bool transferToCPU, cudaStream_t);
+    void launchKernels(const TrackingRegion &, int, HitsOnCPU const & hh, bool doRiemannFit, bool transferToCPU, cudaStream_t);
     std::vector<std::array<int,4>> fetchKernelResult(int);
 
-
     float bField_;
 
     const float extraHitRPhitolerance;