File indexing completed on 2024-04-06 12:28:33
0001 #include <cuda_runtime.h>
0002
0003 #include "CUDADataFormats/Common/interface/Product.h"
0004 #include "CUDADataFormats/Track/interface/TrackSoAHeterogeneousDevice.h"
0005 #include "CUDADataFormats/Track/interface/TrackSoAHeterogeneousHost.h"
0006 #include "DataFormats/Common/interface/Handle.h"
0007 #include "FWCore/Framework/interface/ConsumesCollector.h"
0008 #include "FWCore/Framework/interface/ESHandle.h"
0009 #include "FWCore/Framework/interface/Event.h"
0010 #include "FWCore/Framework/interface/EventSetup.h"
0011 #include "FWCore/Framework/interface/MakerMacros.h"
0012 #include "FWCore/Framework/interface/global/EDProducer.h"
0013 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0014 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0015 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0016 #include "FWCore/PluginManager/interface/ModuleDef.h"
0017 #include "FWCore/Utilities/interface/EDGetToken.h"
0018 #include "FWCore/Utilities/interface/InputTag.h"
0019 #include "FWCore/Utilities/interface/RunningAverage.h"
0020 #include "HeterogeneousCore/CUDACore/interface/ScopedContext.h"
0021 #include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
0022 #include "RecoTracker/TkMSParametrization/interface/PixelRecoUtilities.h"
0023
0024 #include "CAHitNtupletGeneratorOnGPU.h"
0025
0026 template <typename TrackerTraits>
0027 class CAHitNtupletCUDAT : public edm::global::EDProducer<> {
0028 using HitsConstView = TrackingRecHitSoAConstView<TrackerTraits>;
0029 using HitsOnDevice = TrackingRecHitSoADevice<TrackerTraits>;
0030 using HitsOnHost = TrackingRecHitSoAHost<TrackerTraits>;
0031
0032 using TrackSoAHost = TrackSoAHeterogeneousHost<TrackerTraits>;
0033 using TrackSoADevice = TrackSoAHeterogeneousDevice<TrackerTraits>;
0034
0035 using GPUAlgo = CAHitNtupletGeneratorOnGPU<TrackerTraits>;
0036
0037 public:
0038 explicit CAHitNtupletCUDAT(const edm::ParameterSet& iConfig);
0039 ~CAHitNtupletCUDAT() override = default;
0040
0041 static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0042
0043 private:
0044 void beginJob() override;
0045 void endJob() override;
0046
0047 void produce(edm::StreamID streamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const override;
0048
0049 bool onGPU_;
0050
0051 edm::ESGetToken<MagneticField, IdealMagneticFieldRecord> tokenField_;
0052 edm::EDGetTokenT<cms::cuda::Product<HitsOnDevice>> tokenHitGPU_;
0053 edm::EDPutTokenT<cms::cuda::Product<TrackSoADevice>> tokenTrackGPU_;
0054 edm::EDGetTokenT<HitsOnHost> tokenHitCPU_;
0055 edm::EDPutTokenT<TrackSoAHost> tokenTrackCPU_;
0056
0057 GPUAlgo gpuAlgo_;
0058 };
0059
0060 template <typename TrackerTraits>
0061 CAHitNtupletCUDAT<TrackerTraits>::CAHitNtupletCUDAT(const edm::ParameterSet& iConfig)
0062 : onGPU_(iConfig.getParameter<bool>("onGPU")), tokenField_(esConsumes()), gpuAlgo_(iConfig, consumesCollector()) {
0063 if (onGPU_) {
0064 tokenHitGPU_ = consumes(iConfig.getParameter<edm::InputTag>("pixelRecHitSrc"));
0065 tokenTrackGPU_ = produces<cms::cuda::Product<TrackSoADevice>>();
0066 } else {
0067 tokenHitCPU_ = consumes(iConfig.getParameter<edm::InputTag>("pixelRecHitSrc"));
0068 tokenTrackCPU_ = produces<TrackSoAHost>();
0069 }
0070 }
0071
0072 template <typename TrackerTraits>
0073 void CAHitNtupletCUDAT<TrackerTraits>::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0074 edm::ParameterSetDescription desc;
0075
0076 desc.add<bool>("onGPU", true);
0077 desc.add<edm::InputTag>("pixelRecHitSrc", edm::InputTag("siPixelRecHitsPreSplittingCUDA"));
0078
0079 GPUAlgo::fillDescriptions(desc);
0080 descriptions.addWithDefaultLabel(desc);
0081 }
0082
0083 template <typename TrackerTraits>
0084 void CAHitNtupletCUDAT<TrackerTraits>::beginJob() {
0085 gpuAlgo_.beginJob();
0086 }
0087
0088 template <typename TrackerTraits>
0089 void CAHitNtupletCUDAT<TrackerTraits>::endJob() {
0090 gpuAlgo_.endJob();
0091 }
0092
0093 template <typename TrackerTraits>
0094 void CAHitNtupletCUDAT<TrackerTraits>::produce(edm::StreamID streamID,
0095 edm::Event& iEvent,
0096 const edm::EventSetup& es) const {
0097 auto bf = 1. / es.getData(tokenField_).inverseBzAtOriginInGeV();
0098
0099 if (onGPU_) {
0100 auto const& hits = iEvent.get(tokenHitGPU_);
0101
0102 cms::cuda::ScopedContextProduce ctx{hits};
0103 auto& hits_d = ctx.get(hits);
0104 ctx.emplace(iEvent, tokenTrackGPU_, gpuAlgo_.makeTuplesAsync(hits_d, bf, ctx.stream()));
0105 } else {
0106 auto& hits_h = iEvent.get(tokenHitCPU_);
0107 iEvent.emplace(tokenTrackCPU_, gpuAlgo_.makeTuples(hits_h, bf));
0108 }
0109 }
0110
0111 using CAHitNtupletCUDAPhase1 = CAHitNtupletCUDAT<pixelTopology::Phase1>;
0112 DEFINE_FWK_MODULE(CAHitNtupletCUDAPhase1);
0113
0114 using CAHitNtupletCUDAPhase2 = CAHitNtupletCUDAT<pixelTopology::Phase2>;
0115 DEFINE_FWK_MODULE(CAHitNtupletCUDAPhase2);
0116
0117 using CAHitNtupletCUDAHIonPhase1 = CAHitNtupletCUDAT<pixelTopology::HIonPhase1>;
0118 DEFINE_FWK_MODULE(CAHitNtupletCUDAHIonPhase1);
