Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2022-06-04 01:26:15

0001 #include <cuda_runtime.h>
0002 
0003 #include "CUDADataFormats/Common/interface/Product.h"
0004 #include "DataFormats/Common/interface/Handle.h"
0005 #include "FWCore/Framework/interface/ESHandle.h"
0006 #include "FWCore/Framework/interface/Event.h"
0007 #include "FWCore/Framework/interface/EventSetup.h"
0008 #include "FWCore/Framework/interface/MakerMacros.h"
0009 #include "FWCore/Framework/interface/global/EDProducer.h"
0010 #include "FWCore/Framework/interface/ConsumesCollector.h"
0011 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0012 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0013 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0014 #include "FWCore/Utilities/interface/InputTag.h"
0015 #include "FWCore/PluginManager/interface/ModuleDef.h"
0016 #include "FWCore/Utilities/interface/EDGetToken.h"
0017 #include "FWCore/Utilities/interface/RunningAverage.h"
0018 #include "HeterogeneousCore/CUDACore/interface/ScopedContext.h"
0019 #include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
0020 #include "RecoTracker/TkMSParametrization/interface/PixelRecoUtilities.h"
0021 
0022 #include "CAHitNtupletGeneratorOnGPU.h"
0023 #include "CUDADataFormats/Track/interface/PixelTrackHeterogeneous.h"
0024 #include "CUDADataFormats/TrackingRecHit/interface/TrackingRecHit2DHeterogeneous.h"
0025 
0026 class CAHitNtupletCUDA : public edm::global::EDProducer<> {
0027 public:
0028   explicit CAHitNtupletCUDA(const edm::ParameterSet& iConfig);
0029   ~CAHitNtupletCUDA() override = default;
0030 
0031   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0032 
0033 private:
0034   void beginJob() override;
0035   void endJob() override;
0036 
0037   void produce(edm::StreamID streamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const override;
0038 
0039   bool onGPU_;
0040 
0041   edm::ESGetToken<MagneticField, IdealMagneticFieldRecord> tokenField_;
0042   edm::EDGetTokenT<cms::cuda::Product<TrackingRecHit2DGPU>> tokenHitGPU_;
0043   edm::EDPutTokenT<cms::cuda::Product<PixelTrackHeterogeneous>> tokenTrackGPU_;
0044   edm::EDGetTokenT<TrackingRecHit2DCPU> tokenHitCPU_;
0045   edm::EDPutTokenT<PixelTrackHeterogeneous> tokenTrackCPU_;
0046 
0047   CAHitNtupletGeneratorOnGPU gpuAlgo_;
0048 };
0049 
0050 CAHitNtupletCUDA::CAHitNtupletCUDA(const edm::ParameterSet& iConfig)
0051     : onGPU_(iConfig.getParameter<bool>("onGPU")), tokenField_(esConsumes()), gpuAlgo_(iConfig, consumesCollector()) {
0052   if (onGPU_) {
0053     tokenHitGPU_ =
0054         consumes<cms::cuda::Product<TrackingRecHit2DGPU>>(iConfig.getParameter<edm::InputTag>("pixelRecHitSrc"));
0055     tokenTrackGPU_ = produces<cms::cuda::Product<PixelTrackHeterogeneous>>();
0056   } else {
0057     tokenHitCPU_ = consumes<TrackingRecHit2DCPU>(iConfig.getParameter<edm::InputTag>("pixelRecHitSrc"));
0058     tokenTrackCPU_ = produces<PixelTrackHeterogeneous>();
0059   }
0060 }
0061 
0062 void CAHitNtupletCUDA::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0063   edm::ParameterSetDescription desc;
0064 
0065   desc.add<bool>("onGPU", true);
0066   desc.add<edm::InputTag>("pixelRecHitSrc", edm::InputTag("siPixelRecHitsPreSplittingCUDA"));
0067 
0068   CAHitNtupletGeneratorOnGPU::fillDescriptions(desc);
0069   descriptions.add("pixelTracksCUDA", desc);
0070 }
0071 
0072 void CAHitNtupletCUDA::beginJob() { gpuAlgo_.beginJob(); }
0073 
0074 void CAHitNtupletCUDA::endJob() { gpuAlgo_.endJob(); }
0075 
0076 void CAHitNtupletCUDA::produce(edm::StreamID streamID, edm::Event& iEvent, const edm::EventSetup& es) const {
0077   auto bf = 1. / es.getData(tokenField_).inverseBzAtOriginInGeV();
0078 
0079   if (onGPU_) {
0080     auto hHits = iEvent.getHandle(tokenHitGPU_);
0081 
0082     cms::cuda::ScopedContextProduce ctx{*hHits};
0083     auto const& hits = ctx.get(*hHits);
0084 
0085     ctx.emplace(iEvent, tokenTrackGPU_, gpuAlgo_.makeTuplesAsync(hits, bf, ctx.stream()));
0086   } else {
0087     auto const& hits = iEvent.get(tokenHitCPU_);
0088     iEvent.emplace(tokenTrackCPU_, gpuAlgo_.makeTuples(hits, bf));
0089   }
0090 }
0091 
0092 DEFINE_FWK_MODULE(CAHitNtupletCUDA);