diff --git a/Configuration/Eras/python/Era_Phase2_cff.py b/Configuration/Eras/python/Era_Phase2_cff.py index 46432f200234e..7f0c3a6b6fcf4 100644 --- a/Configuration/Eras/python/Era_Phase2_cff.py +++ b/Configuration/Eras/python/Era_Phase2_cff.py @@ -20,6 +20,7 @@ from Configuration.Eras.Modifier_ctpps_2022_cff import ctpps_2022 from Configuration.ProcessModifiers.dd4hep_cff import dd4hep from Configuration.Eras.ModifierChain_trackingMkFitProd_cff import trackingMkFitProdPhase2 +from Configuration.ProcessModifiers.trackTorchClassifier_cff import trackTorchClassifier Phase2 = cms.ModifierChain(Run3_noMkFit.copyAndExclude([phase1Pixel,trackingPhase1,seedingDeepCore,displacedRegionalTracking,ctpps_2022,dd4hep]), - phase2_common, phase2_tracker, trackingPhase2PU140, phase2_ecal, phase2_hcal, phase2_hgcal, phase2_muon, phase2_GEM, hcalHardcodeConditions, phase2_timing, phase2_timing_layer, phase2_trigger, trackingMkFitProdPhase2) + phase2_common, phase2_tracker, trackingPhase2PU140, phase2_ecal, phase2_hcal, phase2_hgcal, phase2_muon, phase2_GEM, hcalHardcodeConditions, phase2_timing, phase2_timing_layer, phase2_trigger, trackingMkFitProdPhase2, trackTorchClassifier) diff --git a/Configuration/ProcessModifiers/python/trackTorchClassifier_cff.py b/Configuration/ProcessModifiers/python/trackTorchClassifier_cff.py new file mode 100644 index 0000000000000..d099fb16ad5dd --- /dev/null +++ b/Configuration/ProcessModifiers/python/trackTorchClassifier_cff.py @@ -0,0 +1,4 @@ +import FWCore.ParameterSet.Config as cms + +# This modifier sets the use of a deep neural network for high purity track selection +trackTorchClassifier = cms.Modifier() diff --git a/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackCutClassifier_cfi.py b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackCutClassifier_cfi.py index cafd4e7324cc5..1bac841e87f33 100644 --- a/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackCutClassifier_cfi.py +++ b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackCutClassifier_cfi.py @@ -47,3 +47,9 @@ hltPhase2LegacyTracking.toModify(hltInitialStepTrackCutClassifier, mva = dict(passThroughForAll=False, passThroughForDisplaced=False) ) + + +from Configuration.ProcessModifiers.trackTorchClassifier_cff import trackTorchClassifier +trackTorchClassifier.toModify(hltInitialStepTrackCutClassifier, + src = "hltInitialStepTrackTorchClassifierOutput" +) diff --git a/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackFeatureExtractor_cfi.py b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackFeatureExtractor_cfi.py new file mode 100644 index 0000000000000..85d667d4bb9f8 --- /dev/null +++ b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackFeatureExtractor_cfi.py @@ -0,0 +1,9 @@ +import FWCore.ParameterSet.Config as cms + +hltInitialStepTrackFeatureExtractor = cms.EDProducer("TrackFeatureExtractor@alpaka", + src = cms.InputTag("hltInitialStepTracks"), + beamSpot = cms.InputTag("hltOnlineBeamSpot"), + alpaka = cms.untracked.PSet( + backend = cms.untracked.string('') + ) +) diff --git a/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackSelectionHighPurity_cfi.py b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackSelectionHighPurity_cfi.py index d299885ff7454..2e18e05cda4cd 100644 --- a/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackSelectionHighPurity_cfi.py +++ b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackSelectionHighPurity_cfi.py @@ -8,3 +8,6 @@ originalQualVals = cms.InputTag("hltInitialStepTrackCutClassifier","QualityMasks"), originalSource = cms.InputTag("hltInitialStepTracks") ) + +from Configuration.ProcessModifiers.trackTorchClassifier_cff import trackTorchClassifier +trackTorchClassifier.toModify(hltInitialStepTrackSelectionHighPurity, originalSource = "hltInitialStepTrackTorchClassifierOutput") diff --git a/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackTorchClassifierOutput_cfi.py b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackTorchClassifierOutput_cfi.py new file mode 100644 index 0000000000000..bcde0318b1fb4 --- /dev/null +++ b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackTorchClassifierOutput_cfi.py @@ -0,0 +1,10 @@ +import FWCore.ParameterSet.Config as cms + +hltInitialStepTrackTorchClassifierOutput = cms.EDProducer("TrackTorchClassifierFromSoA", + src = cms.InputTag("hltInitialStepTracks"), + scores = cms.InputTag("hltInitialStepTrackTorchClassifier"), + features = cms.InputTag("hltInitialStepTrackFeatureExtractor"), + minScore = cms.double(0.377), + dxyThreshold = cms.double(0.5), + highDxyMinScore = cms.double(0.267) +) diff --git a/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackTorchClassifier_cfi.py b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackTorchClassifier_cfi.py new file mode 100644 index 0000000000000..81d41dd6d04ed --- /dev/null +++ b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackTorchClassifier_cfi.py @@ -0,0 +1,9 @@ +import FWCore.ParameterSet.Config as cms + +hltInitialStepTrackTorchClassifier = cms.EDProducer("TrackTorchClassifierAlpaka@alpaka", + modelPath = cms.FileInPath('RecoTracker/FinalTrackSelectors/data/TrackTorchClassifier/model.pt'), + features = cms.InputTag("hltInitialStepTrackFeatureExtractor"), + alpaka = cms.untracked.PSet( + backend = cms.untracked.string('') + ) +) diff --git a/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepHPSelectionSequence_cfi.py b/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepHPSelectionSequence_cfi.py new file mode 100644 index 0000000000000..167c898184620 --- /dev/null +++ b/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepHPSelectionSequence_cfi.py @@ -0,0 +1,9 @@ +import FWCore.ParameterSet.Config as cms + +from ..modules.hltInitialStepTrackCutClassifier_cfi import * +from ..modules.hltInitialStepTrackSelectionHighPurity_cfi import * + +HLTInitialStepHPSelectionSequence = cms.Sequence( + hltInitialStepTrackCutClassifier + +hltInitialStepTrackSelectionHighPurity +) diff --git a/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepSequence_cfi.py b/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepSequence_cfi.py index 3796ff867fb39..c5bb149bb4f45 100644 --- a/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepSequence_cfi.py +++ b/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepSequence_cfi.py @@ -7,14 +7,14 @@ from ..modules.hltInitialStepSeeds_cfi import * from ..modules.hltInitialStepTrackCandidates_cfi import * from ..modules.hltInitialStepTrackCandidatesMkFit_cfi import * -from ..modules.hltInitialStepTrackCutClassifier_cfi import * -from ..modules.hltInitialStepTrackSelectionHighPurity_cfi import * from ..modules.hltInitialStepTracks_cfi import * from ..modules.hltInitialStepTrajectorySeedsLST_cfi import * from ..modules.hltInitialStepTrajectorySeedsLSTTracks_cfi import * from ..modules.hltLST_cfi import * from ..modules.hltSiPhase2RecHits_cfi import * from ..sequences.HLTMkFitInputSequence_cfi import * +from ..sequences.HLTInitialStepHPSelectionSequence_cfi import * +from ..sequences.HLTInitialStepTorchClassifierHPSelectionSequence_cfi import * HLTInitialStepSequence = cms.Sequence( hltInitialStepSeeds @@ -28,8 +28,7 @@ +hltInitialStepTrackCandidatesMkFit +hltInitialStepTrackCandidates +hltInitialStepTracks - +hltInitialStepTrackCutClassifier - +hltInitialStepTrackSelectionHighPurity + +HLTInitialStepHPSelectionSequence ) @@ -97,8 +96,7 @@ +hltLST +hltInitialStepTrackCandidates +hltInitialStepTracks - +hltInitialStepTrackCutClassifier - +hltInitialStepTrackSelectionHighPurity + +HLTInitialStepHPSelectionSequence ) from Configuration.ProcessModifiers.trackingLST_cff import trackingLST @@ -133,9 +131,18 @@ +hltInitialStepTrackCandidatesMkFit +hltInitialStepTrackCandidatesMkFitFit +hltInitialStepTracks - +hltInitialStepTrackCutClassifier - +hltInitialStepTrackSelectionHighPurity + +HLTInitialStepHPSelectionSequence ) from Configuration.ProcessModifiers.trackingMkFitFit_cff import trackingMkFitFit trackingMkFitFit.toReplaceWith(HLTInitialStepSequence, _HLTInitialStepSequenceMkFitFit) + +_HLTInitialStepSequenceTrackTorchClassifier = HLTInitialStepSequence.copyAndExclude([HLTInitialStepHPSelectionSequence]) +_HLTInitialStepSequenceTrackTorchClassifier += HLTInitialStepTorchClassifierHPSelectionSequence + +_HLTInitialStepSequenceTrackTorchClassifierMkFitFit = _HLTInitialStepSequenceMkFitFit.copyAndExclude([HLTInitialStepHPSelectionSequence]) +_HLTInitialStepSequenceTrackTorchClassifierMkFitFit += HLTInitialStepTorchClassifierHPSelectionSequence + +from Configuration.ProcessModifiers.trackTorchClassifier_cff import trackTorchClassifier +(trackTorchClassifier & ~ngtScouting).toReplaceWith(HLTInitialStepSequence, _HLTInitialStepSequenceTrackTorchClassifier) +(trackTorchClassifier & trackingMkFitFit & ~ngtScouting).toReplaceWith(HLTInitialStepSequence, _HLTInitialStepSequenceTrackTorchClassifierMkFitFit) diff --git a/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepTorchClassifierHPSelectionSequence_cfi.py b/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepTorchClassifierHPSelectionSequence_cfi.py new file mode 100644 index 0000000000000..4a82808deddf1 --- /dev/null +++ b/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepTorchClassifierHPSelectionSequence_cfi.py @@ -0,0 +1,15 @@ +import FWCore.ParameterSet.Config as cms + +from ..modules.hltInitialStepTrackCutClassifier_cfi import * +from ..modules.hltInitialStepTrackSelectionHighPurity_cfi import * +from ..modules.hltInitialStepTrackFeatureExtractor_cfi import * +from ..modules.hltInitialStepTrackTorchClassifier_cfi import * +from ..modules.hltInitialStepTrackTorchClassifierOutput_cfi import * + +HLTInitialStepTorchClassifierHPSelectionSequence = cms.Sequence( + hltInitialStepTrackFeatureExtractor + +hltInitialStepTrackTorchClassifier + +hltInitialStepTrackTorchClassifierOutput + +hltInitialStepTrackCutClassifier + +hltInitialStepTrackSelectionHighPurity +) diff --git a/HLTrigger/Configuration/python/HLT_75e33/services/PyTorchService_cfi.py b/HLTrigger/Configuration/python/HLT_75e33/services/PyTorchService_cfi.py new file mode 100644 index 0000000000000..a353db979c92d --- /dev/null +++ b/HLTrigger/Configuration/python/HLT_75e33/services/PyTorchService_cfi.py @@ -0,0 +1,3 @@ +import FWCore.ParameterSet.Config as cms + +PyTorchService = cms.Service("PyTorchService") diff --git a/HLTrigger/Configuration/python/HLT_75e33_cff.py b/HLTrigger/Configuration/python/HLT_75e33_cff.py index 0abb71a072bc8..9bc9bb197b1bc 100644 --- a/HLTrigger/Configuration/python/HLT_75e33_cff.py +++ b/HLTrigger/Configuration/python/HLT_75e33_cff.py @@ -171,6 +171,7 @@ fragment.load("HLTrigger/Configuration/HLT_75e33/psets/seedFromProtoTracks_cfi") fragment.load("HLTrigger/Configuration/HLT_75e33/psets/SiStripClusterChargeCutLoose_cfi") fragment.load("HLTrigger/Configuration/HLT_75e33/psets/SiStripClusterChargeCutNone_cfi") +fragment.load("HLTrigger/Configuration/HLT_75e33/services/PyTorchService_cfi") fragment.load("HLTrigger/Configuration/HLT_75e33/sequences/HLTCalolocalrecoSequence_cfi") fragment.load("HLTrigger/Configuration/HLT_75e33/sequences/HLTCaloTowersRecSequence_cfi") diff --git a/HLTrigger/Configuration/python/HLT_75e33_timing_cff.py b/HLTrigger/Configuration/python/HLT_75e33_timing_cff.py index 3bc1b1f1ffff5..81da19f3e0896 100644 --- a/HLTrigger/Configuration/python/HLT_75e33_timing_cff.py +++ b/HLTrigger/Configuration/python/HLT_75e33_timing_cff.py @@ -261,6 +261,7 @@ fragment.load("HLTrigger/Configuration/HLT_75e33/sequences/HLTVertexRecoSequence_cfi") fragment.load("HLTrigger/Configuration/HLT_75e33/services/FastTimerService_cfi") fragment.load("HLTrigger/Configuration/HLT_75e33/services/ThroughputService_cfi") +fragment.load("HLTrigger/Configuration/HLT_75e33/services/PyTorchService_cfi") fragment.schedule = cms.Schedule(*[ diff --git a/RecoTracker/FinalTrackSelectors/BuildFile.xml b/RecoTracker/FinalTrackSelectors/BuildFile.xml index 8c55683a9a67d..67c012607d7f6 100644 --- a/RecoTracker/FinalTrackSelectors/BuildFile.xml +++ b/RecoTracker/FinalTrackSelectors/BuildFile.xml @@ -1,4 +1,8 @@ + + + + @@ -6,8 +10,11 @@ + + + diff --git a/RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h b/RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h new file mode 100644 index 0000000000000..3c806578461a3 --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h @@ -0,0 +1,30 @@ +#ifndef RecoTracker_FinalTrackSelectors_TrackTorchClassifierFeaturesSoA_h +#define RecoTracker_FinalTrackSelectors_TrackTorchClassifierFeaturesSoA_h + +#include "DataFormats/SoATemplate/interface/SoALayout.h" + +GENERATE_SOA_LAYOUT(TrackTorchClassifierFeaturesSoALayout, + SOA_COLUMN(float, dxyBeamSpot), + SOA_COLUMN(float, dzBeamSpot), + SOA_COLUMN(float, dxyError), + SOA_COLUMN(float, dzError), + SOA_COLUMN(float, normalizedChi2), + SOA_COLUMN(float, eta), + SOA_COLUMN(float, phi), + SOA_COLUMN(float, etaError), + SOA_COLUMN(float, phiError), + SOA_COLUMN(float, ndof), + SOA_COLUMN(float, lostInnerHits), + SOA_COLUMN(float, lostOuterHits), + SOA_COLUMN(float, layersWithoutMeas), + SOA_COLUMN(float, validPixelHits), + SOA_COLUMN(float, validStripHits)) + +using TrackTorchClassifierFeaturesSoA = TrackTorchClassifierFeaturesSoALayout<>; + +// Define the SoA layout for track scores (output) +GENERATE_SOA_LAYOUT(TrackTorchClassifierScoresSoALayout, SOA_COLUMN(float, score)) + +using TrackTorchClassifierScoresSoA = TrackTorchClassifierScoresSoALayout<>; + +#endif // RecoTracker_FinalTrackSelectors_TrackTorchClassifierFeaturesSoA_h diff --git a/RecoTracker/FinalTrackSelectors/interface/alpaka/TrackFeaturesDeviceCollection.h b/RecoTracker/FinalTrackSelectors/interface/alpaka/TrackFeaturesDeviceCollection.h new file mode 100644 index 0000000000000..4c11155fad1c5 --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/interface/alpaka/TrackFeaturesDeviceCollection.h @@ -0,0 +1,12 @@ +#ifndef RecoTracker_FinalTrackSelectors_alpaka_TrackFeaturesDeviceCollection_h +#define RecoTracker_FinalTrackSelectors_alpaka_TrackFeaturesDeviceCollection_h + +#include "DataFormats/Portable/interface/alpaka/PortableCollection.h" +#include "RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h" +#include "HeterogeneousCore/AlpakaInterface/interface/config.h" + +namespace ALPAKA_ACCELERATOR_NAMESPACE { + using TrackFeaturesDeviceCollection = PortableCollection; +} // namespace ALPAKA_ACCELERATOR_NAMESPACE + +#endif // RecoTracker_FinalTrackSelectors_alpaka_TrackFeaturesDeviceCollection_h diff --git a/RecoTracker/FinalTrackSelectors/interface/alpaka/TrackScoresDeviceCollection.h b/RecoTracker/FinalTrackSelectors/interface/alpaka/TrackScoresDeviceCollection.h new file mode 100644 index 0000000000000..cace4d74f155c --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/interface/alpaka/TrackScoresDeviceCollection.h @@ -0,0 +1,12 @@ +#ifndef RecoTracker_FinalTrackSelectors_alpaka_TrackScoresDeviceCollection_h +#define RecoTracker_FinalTrackSelectors_alpaka_TrackScoresDeviceCollection_h + +#include "DataFormats/Portable/interface/alpaka/PortableCollection.h" +#include "RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h" +#include "HeterogeneousCore/AlpakaInterface/interface/config.h" + +namespace ALPAKA_ACCELERATOR_NAMESPACE { + using TrackScoresDeviceCollection = PortableCollection; +} // namespace ALPAKA_ACCELERATOR_NAMESPACE + +#endif // RecoTracker_FinalTrackSelectors_alpaka_TrackScoresDeviceCollection_h diff --git a/RecoTracker/FinalTrackSelectors/plugins/BuildFile.xml b/RecoTracker/FinalTrackSelectors/plugins/BuildFile.xml index ed0db0add7b33..ef33efa1359a8 100644 --- a/RecoTracker/FinalTrackSelectors/plugins/BuildFile.xml +++ b/RecoTracker/FinalTrackSelectors/plugins/BuildFile.xml @@ -38,3 +38,17 @@ + + + + + + + + + + + + + + diff --git a/RecoTracker/FinalTrackSelectors/plugins/TrackTorchClassifierFromSoA.cc b/RecoTracker/FinalTrackSelectors/plugins/TrackTorchClassifierFromSoA.cc new file mode 100644 index 0000000000000..d73ceb4ca8f9f --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/plugins/TrackTorchClassifierFromSoA.cc @@ -0,0 +1,88 @@ +#include "FWCore/Framework/interface/stream/EDProducer.h" +#include "FWCore/Framework/interface/Event.h" +#include "FWCore/Framework/interface/MakerMacros.h" +#include "FWCore/ParameterSet/interface/ParameterSet.h" +#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h" +#include "FWCore/ParameterSet/interface/ParameterSetDescription.h" + +#include "DataFormats/TrackReco/interface/Track.h" +#include "DataFormats/TrackReco/interface/TrackFwd.h" + +#include "DataFormats/Portable/interface/PortableHostCollection.h" +#include "RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h" + +// This module consumes the HOST copy of the Alpaka device scores +// The framework automatically creates host copies of device PortableCollections +class TrackTorchClassifierFromSoA : public edm::stream::EDProducer<> { +public: + explicit TrackTorchClassifierFromSoA(const edm::ParameterSet& iConfig); + ~TrackTorchClassifierFromSoA() override = default; + + static void fillDescriptions(edm::ConfigurationDescriptions& descriptions); + +private: + void produce(edm::Event& iEvent, const edm::EventSetup& iSetup) override; + + const edm::EDGetTokenT tracks_token_; + const edm::EDGetTokenT> scores_token_; + const edm::EDGetTokenT> features_token_; + const float min_score_; + const float dxy_threshold_; + const float high_dxy_min_score_; + + const edm::EDPutTokenT filtered_tracks_token_; + const edm::EDPutTokenT> scores_output_token_; +}; + +TrackTorchClassifierFromSoA::TrackTorchClassifierFromSoA(const edm::ParameterSet& iConfig) + : tracks_token_(consumes(iConfig.getParameter("src"))), + scores_token_(consumes(iConfig.getParameter("scores"))), + features_token_(consumes(iConfig.getParameter("features"))), + min_score_(iConfig.getParameter("minScore")), + dxy_threshold_(iConfig.getParameter("dxyThreshold")), + high_dxy_min_score_(iConfig.getParameter("highDxyMinScore")), + filtered_tracks_token_(produces()), + scores_output_token_(produces>("MVAScores")) {} + +void TrackTorchClassifierFromSoA::fillDescriptions(edm::ConfigurationDescriptions& descriptions) { + edm::ParameterSetDescription desc; + desc.add("src", edm::InputTag("hltInitialStepTracks")); + desc.add("scores", edm::InputTag("hltInitialStepTrackTorchClassifier")); + desc.add("features", edm::InputTag("hltInitialStepTrackTorchClassifier")); + desc.add("minScore", 0.5)->setComment("Minimum DNN score to keep track (working point)"); + desc.add("dxyThreshold", 0.5)->setComment("Tracks with |dxy| > this value bypass the score cut"); + desc.add("highDxyMinScore", 0.5)->setComment("Minimum DNN score to keep high dxy track (working point)"); + descriptions.addWithDefaultLabel(desc); +} + +void TrackTorchClassifierFromSoA::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) { + const auto& tracks = iEvent.get(tracks_token_); + const auto& scores_host = iEvent.get(scores_token_); + const auto& features_host = iEvent.get(features_token_); + + const auto nTracks = tracks.size(); + + // Create filtered track collection + auto filtered_tracks = std::make_unique(); + auto all_scores = std::make_unique>(); + all_scores->reserve(nTracks); + + // Access scores and features from the host collection + auto scores_view = scores_host.const_view(); + auto features_view = features_host.const_view(); + + for (size_t i = 0; i < nTracks; ++i) { + float score = scores_view[i].score(); + float dxy = features_view[i].dxyBeamSpot(); + all_scores->push_back(score); + + if (score >= min_score_ || ((std::abs(dxy) > dxy_threshold_) && (score >= high_dxy_min_score_))) { + filtered_tracks->push_back(tracks[i]); + } + } + + iEvent.put(filtered_tracks_token_, std::move(filtered_tracks)); + iEvent.put(scores_output_token_, std::move(all_scores)); +} + +DEFINE_FWK_MODULE(TrackTorchClassifierFromSoA); diff --git a/RecoTracker/FinalTrackSelectors/plugins/alpaka/TrackFeatureExtractor.cc b/RecoTracker/FinalTrackSelectors/plugins/alpaka/TrackFeatureExtractor.cc new file mode 100644 index 0000000000000..a6322b9e5f690 --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/plugins/alpaka/TrackFeatureExtractor.cc @@ -0,0 +1,89 @@ +#include "DataFormats/TrackReco/interface/Track.h" +#include "DataFormats/TrackReco/interface/TrackFwd.h" +#include "DataFormats/BeamSpot/interface/BeamSpot.h" +#include "DataFormats/VertexReco/interface/Vertex.h" +#include "DataFormats/VertexReco/interface/VertexFwd.h" + +#include "FWCore/Framework/interface/Event.h" +#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h" +#include "FWCore/ParameterSet/interface/ParameterSet.h" +#include "FWCore/ParameterSet/interface/ParameterSetDescription.h" + +#include "HeterogeneousCore/AlpakaCore/interface/alpaka/EDPutToken.h" +#include "HeterogeneousCore/AlpakaCore/interface/alpaka/Event.h" +#include "HeterogeneousCore/AlpakaCore/interface/alpaka/EventSetup.h" +#include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/FixedQueueEDProducer.h" +#include "HeterogeneousCore/AlpakaInterface/interface/config.h" + +#include "DataFormats/Portable/interface/PortableHostCollection.h" +#include "RecoTracker/FinalTrackSelectors/interface/alpaka/TrackFeaturesDeviceCollection.h" +#include "RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h" + +namespace ALPAKA_ACCELERATOR_NAMESPACE { + + class TrackFeatureExtractor : public stream::FixedQueueEDProducer<> { + public: + TrackFeatureExtractor(const edm::ParameterSet& iConfig) + : FixedQueueEDProducer<>(iConfig), + tracksInput_token_(consumes(iConfig.getParameter("src"))), + beamspot_token_(consumes(iConfig.getParameter("beamSpot"))), + featuresPut_token_{produces()} {} + + static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) { + edm::ParameterSetDescription desc; + desc.add("src", edm::InputTag("hltInitialStepTracks")); + desc.add("beamSpot", edm::InputTag("hltOnlineBeamSpot")); + descriptions.addWithDefaultLabel(desc); + } + + void produce(device::Event& iEvent, const device::EventSetup& iSetup) override { + auto const& tracks = iEvent.get(tracksInput_token_); + auto const& beamspot = iEvent.get(beamspot_token_); + + const auto nTracks = tracks.size(); + + // Create HOST collection first, fill it, then copy to device + PortableHostCollection features_host(nTracks); + + auto features_view = features_host.view(); + + for (size_t i = 0; i < nTracks; ++i) { + const auto& track = tracks[i]; + + features_view[i].dxyBeamSpot() = track.dxy(beamspot.position()); + features_view[i].dzBeamSpot() = track.dz(beamspot.position()); + features_view[i].dxyError() = track.dxyError(); + features_view[i].dzError() = track.dzError(); + + features_view[i].normalizedChi2() = track.normalizedChi2(); + features_view[i].eta() = track.eta(); + features_view[i].phi() = track.phi(); + features_view[i].etaError() = track.etaError(); + features_view[i].phiError() = track.phiError(); + features_view[i].ndof() = track.ndof(); + + const auto& hitPattern = track.hitPattern(); + features_view[i].lostInnerHits() = hitPattern.numberOfLostTrackerHits(reco::HitPattern::MISSING_INNER_HITS); + features_view[i].lostOuterHits() = hitPattern.numberOfLostTrackerHits(reco::HitPattern::MISSING_OUTER_HITS); + features_view[i].layersWithoutMeas() = hitPattern.trackerLayersWithoutMeasurement(reco::HitPattern::TRACK_HITS); + features_view[i].validPixelHits() = hitPattern.numberOfValidPixelHits(); + features_view[i].validStripHits() = hitPattern.numberOfValidStripHits(); + } + + // Create device collection and copy from host + TrackFeaturesDeviceCollection features_device(iEvent.queue(), nTracks); + alpaka::memcpy(iEvent.queue(), features_device.buffer(), features_host.const_buffer()); + + iEvent.emplace(featuresPut_token_, std::move(features_device)); + } + + private: + const edm::EDGetTokenT tracksInput_token_; + const edm::EDGetTokenT beamspot_token_; + const device::EDPutToken featuresPut_token_; + }; + +} // namespace ALPAKA_ACCELERATOR_NAMESPACE + +#include "HeterogeneousCore/AlpakaCore/interface/alpaka/MakerMacros.h" +DEFINE_FWK_ALPAKA_MODULE(TrackFeatureExtractor); diff --git a/RecoTracker/FinalTrackSelectors/plugins/alpaka/TrackTorchClassifierAlpaka.cc b/RecoTracker/FinalTrackSelectors/plugins/alpaka/TrackTorchClassifierAlpaka.cc new file mode 100644 index 0000000000000..65da13d3050fc --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/plugins/alpaka/TrackTorchClassifierAlpaka.cc @@ -0,0 +1,80 @@ +#include "RecoTracker/FinalTrackSelectors/interface/alpaka/TrackFeaturesDeviceCollection.h" +#include "RecoTracker/FinalTrackSelectors/interface/alpaka/TrackScoresDeviceCollection.h" +#include "RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h" + +#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h" +#include "FWCore/ParameterSet/interface/ParameterSet.h" +#include "FWCore/ParameterSet/interface/ParameterSetDescription.h" + +#include "HeterogeneousCore/AlpakaCore/interface/alpaka/EDPutToken.h" +#include "HeterogeneousCore/AlpakaCore/interface/alpaka/Event.h" +#include "HeterogeneousCore/AlpakaCore/interface/alpaka/EventSetup.h" +#include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/FixedQueueEDProducer.h" +#include "HeterogeneousCore/AlpakaInterface/interface/config.h" + +#include "PhysicsTools/PyTorchAlpaka/interface/TensorCollection.h" +#include "PhysicsTools/PyTorchAlpaka/interface/alpaka/AlpakaModel.h" + +namespace ALPAKA_ACCELERATOR_NAMESPACE { + + class TrackTorchClassifierAlpaka : public stream::FixedQueueEDProducer<> { + public: + TrackTorchClassifierAlpaka(const edm::ParameterSet& iConfig) + : FixedQueueEDProducer<>(iConfig), + featuresInput_token_(consumes(iConfig.getParameter("features"))), + scoresPut_token_{produces()}, + model_(iConfig.getParameter("modelPath").fullPath()) {} + + static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) { + edm::ParameterSetDescription desc; + desc.add("modelPath", + edm::FileInPath("RecoTracker/FinalTrackSelectors/data/TrackTorchClassifier/model.pt")); + desc.add("features", edm::InputTag("hltInitialStepTrackFeatureExtractor")); + descriptions.addWithDefaultLabel(desc); + } + + void produce(device::Event& iEvent, const device::EventSetup& iSetup) override { + const auto& features = iEvent.get(featuresInput_token_); + const auto batch_size = features.const_view().metadata().size(); + + auto scores_device = TrackScoresDeviceCollection(iEvent.queue(), batch_size); + + auto input_records = features.const_view().records(); + auto output_records = scores_device.view().records(); + + cms::torch::alpakatools::TensorCollection inputs(batch_size); + inputs.add("features", + input_records.dxyBeamSpot(), + input_records.dzBeamSpot(), + input_records.dxyError(), + input_records.dzError(), + input_records.normalizedChi2(), + input_records.eta(), + input_records.phi(), + input_records.etaError(), + input_records.phiError(), + input_records.ndof(), + input_records.lostInnerHits(), + input_records.lostOuterHits(), + input_records.layersWithoutMeas(), + input_records.validPixelHits(), + input_records.validStripHits()); + + cms::torch::alpakatools::TensorCollection outputs(batch_size); + outputs.add("scores", output_records.score()); + + model_.forward(iEvent.queue(), inputs, outputs); + + iEvent.emplace(scoresPut_token_, std::move(scores_device)); + } + + private: + const device::EDGetToken featuresInput_token_; + const device::EDPutToken scoresPut_token_; + torch::AlpakaModel model_; + }; + +} // namespace ALPAKA_ACCELERATOR_NAMESPACE + +#include "HeterogeneousCore/AlpakaCore/interface/alpaka/MakerMacros.h" +DEFINE_FWK_ALPAKA_MODULE(TrackTorchClassifierAlpaka); diff --git a/RecoTracker/FinalTrackSelectors/src/alpaka/classes_cuda.h b/RecoTracker/FinalTrackSelectors/src/alpaka/classes_cuda.h new file mode 100644 index 0000000000000..596384201e2fb --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/src/alpaka/classes_cuda.h @@ -0,0 +1,9 @@ +#ifndef RecoTracker_FinalTrackSelectors_src_alpaka_classes_cuda_h +#define RecoTracker_FinalTrackSelectors_src_alpaka_classes_cuda_h + +#include "DataFormats/Common/interface/DeviceProduct.h" +#include "DataFormats/Common/interface/Wrapper.h" +#include "RecoTracker/FinalTrackSelectors/interface/alpaka/TrackFeaturesDeviceCollection.h" +#include "RecoTracker/FinalTrackSelectors/interface/alpaka/TrackScoresDeviceCollection.h" + +#endif diff --git a/RecoTracker/FinalTrackSelectors/src/alpaka/classes_cuda_def.xml b/RecoTracker/FinalTrackSelectors/src/alpaka/classes_cuda_def.xml new file mode 100644 index 0000000000000..af34dde7a2e0e --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/src/alpaka/classes_cuda_def.xml @@ -0,0 +1,9 @@ + + + + + + + + + diff --git a/RecoTracker/FinalTrackSelectors/src/alpaka/classes_rocm.h b/RecoTracker/FinalTrackSelectors/src/alpaka/classes_rocm.h new file mode 100644 index 0000000000000..6e163e2358b17 --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/src/alpaka/classes_rocm.h @@ -0,0 +1,9 @@ +#ifndef RecoTracker_FinalTrackSelectors_src_alpaka_classes_rocm_h +#define RecoTracker_FinalTrackSelectors_src_alpaka_classes_rocm_h + +#include "DataFormats/Common/interface/DeviceProduct.h" +#include "DataFormats/Common/interface/Wrapper.h" +#include "RecoTracker/FinalTrackSelectors/interface/alpaka/TrackFeaturesDeviceCollection.h" +#include "RecoTracker/FinalTrackSelectors/interface/alpaka/TrackScoresDeviceCollection.h" + +#endif diff --git a/RecoTracker/FinalTrackSelectors/src/alpaka/classes_rocm_def.xml b/RecoTracker/FinalTrackSelectors/src/alpaka/classes_rocm_def.xml new file mode 100644 index 0000000000000..2d8f93aacdc6e --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/src/alpaka/classes_rocm_def.xml @@ -0,0 +1,9 @@ + + + + + + + + + diff --git a/RecoTracker/FinalTrackSelectors/src/classes.h b/RecoTracker/FinalTrackSelectors/src/classes.h new file mode 100644 index 0000000000000..763ef7ba5b78a --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/src/classes.h @@ -0,0 +1,8 @@ +#ifndef RecoTracker_FinalTrackSelectors_src_classes_h +#define RecoTracker_FinalTrackSelectors_src_classes_h + +#include "DataFormats/Common/interface/Wrapper.h" +#include "DataFormats/Portable/interface/PortableHostCollection.h" +#include "RecoTracker/FinalTrackSelectors/interface/TrackTorchClassifierFeaturesSoA.h" + +#endif diff --git a/RecoTracker/FinalTrackSelectors/src/classes_def.xml b/RecoTracker/FinalTrackSelectors/src/classes_def.xml new file mode 100644 index 0000000000000..12286f55003d3 --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/src/classes_def.xml @@ -0,0 +1,9 @@ + + + + + + + + +