diff --git a/Configuration/ProcessModifiers/python/trackCutClassifier_cff.py b/Configuration/ProcessModifiers/python/trackCutClassifier_cff.py new file mode 100644 index 0000000000000..c2866cc1128c0 --- /dev/null +++ b/Configuration/ProcessModifiers/python/trackCutClassifier_cff.py @@ -0,0 +1,4 @@ +import FWCore.ParameterSet.Config as cms + +# This modifier sets the use of cut based high purity track selection +trackCutClassifier = 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..c77a993be430a 100644 --- a/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackCutClassifier_cfi.py +++ b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackCutClassifier_cfi.py @@ -32,7 +32,7 @@ minLayersForDisplaced = cms.int32(4) ), qualityCuts = cms.vdouble(-0.7, 0.1, 0.7), - src = cms.InputTag("hltInitialStepTracks"), + src = cms.InputTag("hltInitialStepTrackTorchClassifierOutput"), vertices = cms.InputTag("hltPhase2PixelVertices") ) @@ -47,3 +47,9 @@ hltPhase2LegacyTracking.toModify(hltInitialStepTrackCutClassifier, mva = dict(passThroughForAll=False, passThroughForDisplaced=False) ) + + +from Configuration.ProcessModifiers.trackCutClassifier_cff import trackCutClassifier +trackCutClassifier.toModify(hltInitialStepTrackCutClassifier, + src = "hltInitialStepTracks" +) 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..11e0e837a987d --- /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("alpaka_serial_sync::TrackFeatureExtractor", + 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 7e384417b31be..f1cf6d4709c7b 100644 --- a/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackSelectionHighPurity_cfi.py +++ b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackSelectionHighPurity_cfi.py @@ -6,8 +6,11 @@ minQuality = cms.string('highPurity'), originalMVAVals = cms.InputTag("hltInitialStepTrackCutClassifier","MVAValues"), originalQualVals = cms.InputTag("hltInitialStepTrackCutClassifier","QualityMasks"), - originalSource = cms.InputTag("hltInitialStepTracks") + originalSource = cms.InputTag("hltInitialStepTrackTorchClassifierOutput") ) from Configuration.ProcessModifiers.mtd_at_hlt_cff import mtd_at_hlt mtd_at_hlt.toModify(hltInitialStepTrackSelectionHighPurity, copyTrajectories = True) + +from Configuration.ProcessModifiers.trackCutClassifier_cff import trackCutClassifier +trackCutClassifier.toModify(hltInitialStepTrackSelectionHighPurity, originalSource = "hltInitialStepTracks") 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..c4b75ccfc69aa --- /dev/null +++ b/HLTrigger/Configuration/python/HLT_75e33/modules/hltInitialStepTrackTorchClassifierOutput_cfi.py @@ -0,0 +1,14 @@ +import FWCore.ParameterSet.Config as cms + +hltInitialStepTrackTorchClassifierOutput = cms.EDProducer("TrackTorchClassifierFromSoA", + src = cms.InputTag("hltInitialStepTracks"), + scores = cms.InputTag("hltInitialStepTrackTorchClassifier"), + features = cms.InputTag("hltInitialStepTrackFeatureExtractor"), + copyTrajectories = cms.bool(False), + minScore = cms.double(0.377), + dxyThreshold = cms.double(0.5), + highDxyMinScore = cms.double(0.267) +) + +from Configuration.ProcessModifiers.mtd_at_hlt_cff import mtd_at_hlt +mtd_at_hlt.toModify(hltInitialStepTrackTorchClassifierOutput, copyTrajectories = True) 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..98690ed4bc50c --- /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("alpaka_serial_sync::TrackTorchClassifierAlpaka", + 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/HLTInitialStepCutClassifierHPSelectionSequence_cfi.py b/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepCutClassifierHPSelectionSequence_cfi.py new file mode 100644 index 0000000000000..74ef4c3ca4b2c --- /dev/null +++ b/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepCutClassifierHPSelectionSequence_cfi.py @@ -0,0 +1,9 @@ +import FWCore.ParameterSet.Config as cms + +from ..modules.hltInitialStepTrackCutClassifier_cfi import * +from ..modules.hltInitialStepTrackSelectionHighPurity_cfi import * + +HLTInitialStepCutClassifierHPSelectionSequence = cms.Sequence( + hltInitialStepTrackCutClassifier + +hltInitialStepTrackSelectionHighPurity +) 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..3f4e985cde4a9 --- /dev/null +++ b/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepHPSelectionSequence_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 * + +HLTInitialStepHPSelectionSequence = cms.Sequence( + hltInitialStepTrackFeatureExtractor + +hltInitialStepTrackTorchClassifier + +hltInitialStepTrackTorchClassifierOutput + +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 574c199b9714c..cfbd5c700aec8 100644 --- a/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepSequence_cfi.py +++ b/HLTrigger/Configuration/python/HLT_75e33/sequences/HLTInitialStepSequence_cfi.py @@ -6,13 +6,13 @@ 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 ..sequences.HLTMkFitInputSequence_cfi import * +from .HLTInitialStepHPSelectionSequence_cfi import * +from .HLTInitialStepCutClassifierHPSelectionSequence_cfi import * HLTInitialStepSequence = cms.Sequence( hltInitialStepSeeds @@ -24,8 +24,7 @@ +hltInitialStepTrackCandidatesMkFit +hltInitialStepTrackCandidates +hltInitialStepTracks - +hltInitialStepTrackCutClassifier - +hltInitialStepTrackSelectionHighPurity + +HLTInitialStepHPSelectionSequence ) @@ -84,8 +83,7 @@ +hltLST +hltInitialStepTrackCandidates +hltInitialStepTracks - +hltInitialStepTrackCutClassifier - +hltInitialStepTrackSelectionHighPurity + +HLTInitialStepHPSelectionSequence ) from Configuration.ProcessModifiers.trackingLST_cff import trackingLST @@ -116,9 +114,18 @@ +hltInitialStepTrackCandidatesMkFit +hltInitialStepTrackCandidatesMkFitFit +hltInitialStepTracks - +hltInitialStepTrackCutClassifier - +hltInitialStepTrackSelectionHighPurity + +HLTInitialStepHPSelectionSequence ) from Configuration.ProcessModifiers.trackingMkFitFit_cff import trackingMkFitFit trackingMkFitFit.toReplaceWith(HLTInitialStepSequence, _HLTInitialStepSequenceMkFitFit) + +_HLTInitialStepSequenceTrackCutClassifier = HLTInitialStepSequence.copyAndExclude([HLTInitialStepHPSelectionSequence]) +_HLTInitialStepSequenceTrackCutClassifier += HLTInitialStepCutClassifierHPSelectionSequence + +_HLTInitialStepSequenceTrackCutClassifierMkFitFit = _HLTInitialStepSequenceMkFitFit.copyAndExclude([HLTInitialStepHPSelectionSequence]) +_HLTInitialStepSequenceTrackCutClassifierMkFitFit += HLTInitialStepCutClassifierHPSelectionSequence + +from Configuration.ProcessModifiers.trackCutClassifier_cff import trackCutClassifier +(trackCutClassifier & ~ngtScouting).toReplaceWith(HLTInitialStepSequence, _HLTInitialStepSequenceTrackCutClassifier) +(trackCutClassifier & trackingMkFitFit & ~ngtScouting).toReplaceWith(HLTInitialStepSequence, _HLTInitialStepSequenceTrackCutClassifierMkFitFit) 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 b1b646411f23e..e52e235792287 100644 --- a/HLTrigger/Configuration/python/HLT_75e33_cff.py +++ b/HLTrigger/Configuration/python/HLT_75e33_cff.py @@ -280,6 +280,7 @@ fragment.load("HLTrigger/Configuration/HLT_75e33/services/FastTimerService_cfi") fragment.load("HLTrigger/Configuration/HLT_75e33/services/MessageLogger_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/HLTrigger/Configuration/python/HLT_75e33_timing_cff.py b/HLTrigger/Configuration/python/HLT_75e33_timing_cff.py index 9f87e6aacae11..400ad0d8644cd 100644 --- a/HLTrigger/Configuration/python/HLT_75e33_timing_cff.py +++ b/HLTrigger/Configuration/python/HLT_75e33_timing_cff.py @@ -252,6 +252,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 c54e77d353eb0..7399dabef92dd 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..509ee53bf24e3 --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/plugins/TrackTorchClassifierFromSoA.cc @@ -0,0 +1,127 @@ +#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" +#include "TrackingTools/PatternTools/interface/Trajectory.h" +#include "TrackingTools/PatternTools/interface/TrajTrackAssociation.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_; + edm::EDGetTokenT> trajectories_token_; + const edm::EDGetTokenT> scores_token_; + const edm::EDGetTokenT> features_token_; + const bool copy_trajectories_; + const float min_score_; + const float dxy_threshold_; + const float high_dxy_min_score_; + + const edm::EDPutTokenT filtered_tracks_token_; + const edm::EDPutTokenT> filtered_trajectories_token_; + const edm::EDPutTokenT traj_track_associations_token_; + const edm::EDPutTokenT> scores_output_token_; +}; + +TrackTorchClassifierFromSoA::TrackTorchClassifierFromSoA(const edm::ParameterSet& iConfig) + : tracks_token_(consumes(iConfig.getParameter("src"))), + trajectories_token_(), + scores_token_(consumes(iConfig.getParameter("scores"))), + features_token_(consumes(iConfig.getParameter("features"))), + copy_trajectories_(iConfig.getParameter("copyTrajectories")), + min_score_(iConfig.getParameter("minScore")), + dxy_threshold_(iConfig.getParameter("dxyThreshold")), + high_dxy_min_score_(iConfig.getParameter("highDxyMinScore")), + filtered_tracks_token_(produces()), + filtered_trajectories_token_(produces>()), + traj_track_associations_token_(produces()), + scores_output_token_(produces>("MVAScores")) { + if (copy_trajectories_) { + trajectories_token_ = consumes(iConfig.getParameter("src")); + } +} + +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("copyTrajectories", false) + ->setComment("Whether to produce the filtered trajectory collection and associations"); + 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 std::vector* trajectories = nullptr; + if (copy_trajectories_) { + trajectories = &iEvent.get(trajectories_token_); + } + + const auto nTracks = tracks.size(); + + // Create filtered track collection and optionally copy the corresponding trajectories + auto filtered_tracks = std::make_unique(); + std::unique_ptr> filtered_trajectories; + std::unique_ptr traj_track_associations; + if (copy_trajectories_) { + filtered_trajectories = std::make_unique>(); + traj_track_associations = std::make_unique( + iEvent.getRefBeforePut>(), iEvent.getRefBeforePut()); + } + 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]); + if (copy_trajectories_) { + filtered_trajectories->push_back((*trajectories)[i]); + traj_track_associations->insert( + edm::Ref>(iEvent.getRefBeforePut>(), + filtered_trajectories->size() - 1), + reco::TrackRef(iEvent.getRefBeforePut(), filtered_tracks->size() - 1)); + } + } + } + + iEvent.put(filtered_tracks_token_, std::move(filtered_tracks)); + if (copy_trajectories_) { + iEvent.put(filtered_trajectories_token_, std::move(filtered_trajectories)); + iEvent.put(traj_track_associations_token_, std::move(traj_track_associations)); + } + 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..1fd09791156b8 --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/plugins/alpaka/TrackFeatureExtractor.cc @@ -0,0 +1,85 @@ +#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/EDProducer.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::EDProducer<> { + public: + TrackFeatureExtractor(const edm::ParameterSet& iConfig) + : EDProducer<>(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(); + } + + iEvent.emplace(featuresPut_token_, std::move(features_host)); + } + + private: + const edm::EDGetTokenT tracksInput_token_; + const edm::EDGetTokenT beamspot_token_; + const edm::EDPutTokenT> 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..cf254dc9fac00 --- /dev/null +++ b/RecoTracker/FinalTrackSelectors/plugins/alpaka/TrackTorchClassifierAlpaka.cc @@ -0,0 +1,118 @@ +#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)); + } + + void beginStream(edm::StreamID sid, Queue queue) override { + // Warmup the model with dummy data + const int warmupBatchSize = 4992; + // Allocate dummy input and output tensors on the device + auto features = TrackFeaturesDeviceCollection(queue, warmupBatchSize); + auto scores_device = TrackScoresDeviceCollection(queue, warmupBatchSize); + + auto input_records = features.view().records(); + auto output_records = scores_device.view().records(); + + for (auto it = 0; it < warmupIterations_; ++it) { + cms::torch::alpakatools::TensorCollection dummy_inputs(warmupBatchSize); + cms::torch::alpakatools::TensorCollection dummy_outputs(warmupBatchSize); + + dummy_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()); + + dummy_outputs.add("scores", output_records.score()); + + model_.forward(queue, dummy_inputs, dummy_outputs); + } + } + + private: + const device::EDGetToken featuresInput_token_; + const device::EDPutToken scoresPut_token_; + torch::AlpakaModel model_; + const int warmupIterations_ = 3; + }; + +} // 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 @@ + + + + + + + + +