Adding recHit masking to extended pixel tracks in heterogeneous CA

Configuration/ProcessModifiers/python/allPixelTrackMask_cff.py

@@ -0,0 +1,8 @@
+import FWCore.ParameterSet.Config as cms
+
+from Configuration.ProcessModifiers.alpaka_cff import alpaka
+from Configuration.ProcessModifiers.pixelTrackMask_cff import pixelTrackMask
+from Configuration.ProcessModifiers.phase2CAExtension_cff import phase2CAExtension
+
+# collect all PixelTrackMask-related process modifiers here
+allPixelTrackMask = cms.ModifierChain(alpaka,pixelTrackMask,phase2CAExtension)

Configuration/ProcessModifiers/python/pixelTrackMask_cff.py

@@ -0,0 +1,3 @@
+import FWCore.ParameterSet.Config as cms
+
+pixelTrackMask =  cms.Modifier()

Configuration/PyReleaseValidation/python/upgradeWorkflowComponents.py

@@ -1513,6 +1513,27 @@ def setup_(self, step, stepName, stepDict, k, properties):
     offset = 0.4021,
 )
 
+# Pixel-only quadruplets with CA Extension and recHit masking workflow running on GPU (optional)
+#  - Pixel-only with CA Extension and recHit masking reconstruction with Alpaka, with DQM and validation
+#  - harvesting
+
+upgradeWFs['PatatrackPixelOnlyAlpakaCAExtensionMask'] = PatatrackWorkflow(
+    digi = {
+        '--procModifiers': 'allPixelTrackMask',
+        '--customise' : 'HeterogeneousCore/AlpakaServices/customiseAlpakaServiceMemoryFilling.customiseAlpakaServiceMemoryFilling',
+    },
+    reco = {
+        '-s': 'RAW2DIGI:RawToDigi_pixelOnly,RECO:reconstruction_pixelTrackingOnly,VALIDATION:@pixelTrackingOnlyValidation,DQM:@pixelTrackingOnlyDQM',
+        '--procModifiers': 'allPixelTrackMask',
+        '--customise' : 'HeterogeneousCore/AlpakaServices/customiseAlpakaServiceMemoryFilling.customiseAlpakaServiceMemoryFilling',
+    },
+    harvest = {
+        '-s': 'HARVESTING:@trackingOnlyValidation+@pixelTrackingOnlyDQM'
+    },
+    suffix = 'Patatrack_PixelOnlyAlpakaCAExtensionMask',
+    offset = 0.4022,
+)
+
 # Pixel-only quadruplets workflow running on GPU (optional)
 #  - Pixel-only reconstruction with Alpaka, with standard and CPUvsGPU DQM and validation
 #  - harvesting for CPUvsGPU validation

DataFormats/TrackSoA/interface/TrackDefinitions.h

@@ -20,6 +20,23 @@ namespace pixelTrack {
     return ret;
   }
 
+  enum class Iteration : uint8_t {
+    promptHighPt,
+    promptLowPt,
+    notIteration
+  };  // Not sure if a notIteration will be needed
+  constexpr uint32_t iterationSize{uint8_t(Iteration::notIteration)};
+  constexpr std::string_view iterationName[iterationSize]{"promptHighPt", "promptLowPt"};
+  inline Iteration iterationByName(std::string_view name) {
+    auto qp = std::find(iterationName, iterationName + iterationSize, name) - iterationName;
+    auto ret = static_cast<Iteration>(qp);
+
+    if (ret == pixelTrack::Iteration::notIteration)
+      throw std::invalid_argument(std::string(name) + " is not a pixelTrack::Iteration!");
+
+    return ret;
+  }
+
 #ifdef GPU_SMALL_EVENTS
   // kept for testing and debugging
   constexpr uint32_t maxNumber() { return 2 * 1024; }

DataFormats/TrackSoA/interface/TracksSoA.h

@@ -18,6 +18,7 @@ namespace reco {
   using Vector5f = Eigen::Matrix<float, 5, 1>;
   using Vector15f = Eigen::Matrix<float, 15, 1>;
   using Quality = pixelTrack::Quality;
+  using Iteration = pixelTrack::Iteration;
 
   GENERATE_SOA_LAYOUT(TrackLayout,
                       SOA_COLUMN(Quality, quality),
@@ -29,7 +30,8 @@ namespace reco {
                       SOA_EIGEN_COLUMN(Vector5f, state),
                       SOA_EIGEN_COLUMN(Vector15f, covariance),
                       SOA_SCALAR(int, nTracks),
-                      SOA_COLUMN(uint32_t, hitOffsets))
+                      SOA_COLUMN(uint32_t, hitOffsets),
+                      SOA_COLUMN(Iteration, iteration))
 
   GENERATE_SOA_LAYOUT(TrackHitsLayout, SOA_COLUMN(uint32_t, id), SOA_COLUMN(uint32_t, detId))
 

DataFormats/TrackingRecHitSoA/interface/TrackingRecHitsDevice.h

@@ -20,6 +20,9 @@ namespace reco {
   template <typename TDev>
   using HitPortableCollectionDevice = PortableDeviceCollection<TDev, reco::TrackingBlocksSoA>;
 
+  template <typename TDev>
+  using TrackingRecHitsMaskingDevice = PortableDeviceCollection<TDev, TrackingRecHitsMaskingSoA>;
+
   template <typename TDev>
   class TrackingRecHitDevice : public HitPortableCollectionDevice<TDev> {
   public:

DataFormats/TrackingRecHitSoA/interface/TrackingRecHitsHost.h

@@ -20,6 +20,8 @@ namespace reco {
 
   using HitPortableCollectionHost = PortableHostCollection<reco::TrackingBlocksSoA>;
 
+  using TrackingRecHitsMaskingHost = PortableHostCollection<TrackingRecHitsMaskingSoA>;
+
   class TrackingRecHitHost : public HitPortableCollectionHost {
   public:
     explicit TrackingRecHitHost(edm::Uninitialized)

DataFormats/TrackingRecHitSoA/interface/TrackingRecHitsSoA.h

@@ -49,6 +49,8 @@ namespace reco {
                       SOA_SCALAR(int32_t, endCapZPos),
                       SOA_SCALAR(int32_t, endCapZNeg))
 
+  GENERATE_SOA_LAYOUT(TrackingRecHitsMaskingLayout, SOA_COLUMN(uint32_t, recHitMask));
+
   using TrackingRecHitSoA = TrackingHitsLayout<>;
   using TrackingRecHitView = TrackingRecHitSoA::View;
   using TrackingRecHitConstView = TrackingRecHitSoA::ConstView;
@@ -65,6 +67,10 @@ namespace reco {
   using AverageGeometryView = AverageGeometrySoA::View;
   using AverageGeometryConstView = AverageGeometrySoA::ConstView;
 
+  using TrackingRecHitsMaskingSoA = TrackingRecHitsMaskingLayout<>;
+  using TrackingRecHitsMaskingView = TrackingRecHitsMaskingSoA::View;
+  using TrackingRecHitsMaskingConstView = TrackingRecHitsMaskingSoA::ConstView;
+
 };  // namespace reco
 
 #endif

DataFormats/TrackingRecHitSoA/interface/alpaka/TrackingRecHitsSoACollection.h

@@ -23,6 +23,10 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::reco {
   using TrackingRecHitsSoACollection = std::conditional_t<std::is_same_v<Device, alpaka::DevCpu>,
                                                           ::reco::TrackingRecHitHost,
                                                           ::reco::TrackingRecHitDevice<Device>>;
+
+  using TrackingRecHitsMaskingCollection = std::conditional_t<std::is_same_v<Device, alpaka::DevCpu>,
+                                                              ::reco::TrackingRecHitsMaskingHost,
+                                                              ::reco::TrackingRecHitsMaskingDevice<Device>>;
 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE::reco
 
 namespace cms::alpakatools {
@@ -87,5 +91,6 @@ namespace cms::alpakatools {
 }  // namespace cms::alpakatools
 
 ASSERT_DEVICE_MATCHES_HOST_COLLECTION(reco::TrackingRecHitsSoACollection, reco::TrackingRecHitHost);
+ASSERT_DEVICE_MATCHES_HOST_COLLECTION(reco::TrackingRecHitsMaskingCollection, reco::TrackingRecHitsMaskingHost);
 
 #endif  // DataFormats_TrackingRecHitSoA_interface_alpaka_TrackingRecHitsSoACollection_h

DataFormats/TrackingRecHitSoA/src/alpaka/classes_cuda_def.xml

@@ -2,5 +2,7 @@
   <class name="alpaka_cuda_async::reco::TrackingRecHitsSoACollection" persistent="false"/>
   <class name="edm::DeviceProduct<alpaka_cuda_async::reco::TrackingRecHitsSoACollection>" persistent="false"/>
   <class name="edm::Wrapper<edm::DeviceProduct<alpaka_cuda_async::reco::TrackingRecHitsSoACollection>>" persistent="false"/>
+  <class name="edm::DeviceProduct<alpaka_cuda_async::reco::TrackingRecHitsMaskingCollection>" persistent="false"/>
+  <class name="edm::Wrapper<edm::DeviceProduct<alpaka_cuda_async::reco::TrackingRecHitsMaskingCollection>>" persistent="false"/>
 </lcgdict>
 

DataFormats/TrackingRecHitSoA/src/alpaka/classes_rocm_def.xml

@@ -2,5 +2,7 @@
   <class name="alpaka_rocm_async::reco::TrackingRecHitsSoACollection" persistent="false"/>
   <class name="edm::DeviceProduct<alpaka_rocm_async::reco::TrackingRecHitsSoACollection>" persistent="false"/>
   <class name="edm::Wrapper<edm::DeviceProduct<alpaka_rocm_async::reco::TrackingRecHitsSoACollection>>" persistent="false"/>
+  <class name="edm::DeviceProduct<alpaka_rocm_async::reco::TrackingRecHitsMaskingCollection>" persistent="false"/>
+  <class name="edm::Wrapper<edm::DeviceProduct<alpaka_rocm_async::reco::TrackingRecHitsMaskingCollection>>" persistent="false"/>
 </lcgdict>
 

DataFormats/TrackingRecHitSoA/src/classes.cc

@@ -2,3 +2,4 @@
 #include "DataFormats/TrackingRecHitSoA/interface/TrackingRecHitsHost.h"
 
 SET_PORTABLEHOSTCOLLECTION_READ_RULES(reco::HitPortableCollectionHost);
+SET_PORTABLEHOSTCOLLECTION_READ_RULES(reco::TrackingRecHitsMaskingHost);

DataFormats/TrackingRecHitSoA/src/classes_def.xml

@@ -12,4 +12,8 @@
 
   <class name="SiPixelHitStatus"/>
   <class name="SiPixelHitStatusAndCharge"/>
+
+  <class name="reco::TrackingRecHitsMaskingLayout<128,false>"/>
+  <class name="PortableHostCollection<reco::TrackingRecHitsMaskingLayout<128,false> >"/>
+  <class name="edm::Wrapper<PortableHostCollection<reco::TrackingRecHitsMaskingLayout<128,false> > >"/>
 </lcgdict>

RecoLocalTracker/SiPixelRecHits/plugins/alpaka/PixelRecHitKernel.h

@@ -39,6 +39,18 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE {
                                                        ParamsOnDevice const* cpeParams,
                                                        Queue queue) const;
     };
+    class PixelRecHitMaskingKernel {
+    public:
+      PixelRecHitMaskingKernel() = default;
+      ~PixelRecHitMaskingKernel() = default;
+
+      PixelRecHitMaskingKernel(const PixelRecHitMaskingKernel&) = delete;
+      PixelRecHitMaskingKernel(PixelRecHitMaskingKernel&&) = delete;
+      PixelRecHitMaskingKernel& operator=(const PixelRecHitMaskingKernel&) = delete;
+      PixelRecHitMaskingKernel& operator=(PixelRecHitMaskingKernel&&) = delete;
+
+      reco::TrackingRecHitsMaskingCollection makeHitsMaskingAsync(uint32_t const nHits, Queue queue) const;
+    };
   }  // namespace pixelgpudetails
 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE
 

RecoLocalTracker/SiPixelRecHits/plugins/alpaka/PixelRecHitKernels.dev.cc

@@ -78,6 +78,30 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE {
       return hits_d;
     }
 
+    TrackingRecHitsMaskingCollection PixelRecHitMaskingKernel::makeHitsMaskingAsync(uint32_t const nHits,
+                                                                                    Queue queue) const {
+      using namespace pixelRecHits;
+
+      TrackingRecHitsMaskingCollection mask_d(queue, nHits);
+
+      int threadsPerBlock = 128;
+      int blocks = cms::alpakatools::divide_up_by(nHits, threadsPerBlock);
+      const auto workDiv1D = cms::alpakatools::make_workdiv<Acc1D>(blocks, threadsPerBlock);
+
+#ifdef GPU_DEBUG
+      std::cout << "launching LaunchZerosPixelMask kernel on " << alpaka::core::demangled<Acc1D> << " with " << blocks
+                << " blocks" << std::endl;
+#endif
+      alpaka::exec<Acc1D>(queue, workDiv1D, LaunchZerosPixelMask{}, mask_d.view());
+
+#ifdef GPU_DEBUG
+      alpaka::wait(queue);
+      std::cout << "makeHitsMaskingAsync -> DONE!" << std::endl;
+#endif
+
+      return mask_d;
+    }
+
     template class PixelRecHitKernel<pixelTopology::Phase1>;
     template class PixelRecHitKernel<pixelTopology::Phase2>;
     template class PixelRecHitKernel<pixelTopology::HIonPhase1>;

RecoLocalTracker/SiPixelRecHits/plugins/alpaka/PixelRecHits.h

@@ -220,6 +220,17 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE {
       }
     };
 
+    class LaunchZerosPixelMask {
+    public:
+      ALPAKA_FN_ACC void operator()(Acc1D const& acc, ::reco::TrackingRecHitsMaskingView mask) const {
+        for (uint32_t ic : cms::alpakatools::independent_group_elements(acc, mask.metadata().size())) {
+          assert(ic < (uint32_t)mask.metadata().size());
+          mask[ic].recHitMask() = 0;
+        }
+        alpaka::syncBlockThreads(acc);
+      }
+    };
+
   }  // namespace pixelRecHits
 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE
 

RecoLocalTracker/SiPixelRecHits/plugins/alpaka/SiPixelRecHitExtendedAlpaka.cc

@@ -23,6 +23,8 @@
 #include "RecoLocalTracker/SiPixelRecHits/interface/alpaka/PixelCPEFastParamsCollection.h"
 #include "RecoLocalTracker/SiPixelRecHits/interface/pixelCPEforDevice.h"
 
+#include "PixelRecHitKernel.h"
+
 //#define GPU_DEBUG
 
 namespace ALPAKA_ACCELERATOR_NAMESPACE {
@@ -41,13 +43,17 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE {
     const device::EDGetToken<reco::TrackingRecHitsSoACollection> trackerRecHitToken_;
 
     const device::EDPutToken<reco::TrackingRecHitsSoACollection> outputRecHitsSoAToken_;
+    const device::EDPutToken<reco::TrackingRecHitsMaskingCollection> outputRecHitsMaskToken_;
+
+    const pixelgpudetails::PixelRecHitMaskingKernel Algo_;
   };
 
   SiPixelRecHitExtendedAlpaka::SiPixelRecHitExtendedAlpaka(const edm::ParameterSet& iConfig)
       : EDProducer(iConfig),
         pixelRecHitToken_(consumes(iConfig.getParameter<edm::InputTag>("pixelRecHitsSoA"))),
         trackerRecHitToken_(consumes(iConfig.getParameter<edm::InputTag>("trackerRecHitsSoA"))),
-        outputRecHitsSoAToken_(produces()) {}
+        outputRecHitsSoAToken_(produces()),
+        outputRecHitsMaskToken_(produces()) {}
 
   void SiPixelRecHitExtendedAlpaka::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
     edm::ParameterSetDescription desc;
@@ -185,6 +191,10 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE {
 
     // emplace the merged SoA collection in the event
     iEvent.emplace(outputRecHitsSoAToken_, std::move(output));
+
+    // create masking vector with zeros and emplace in the event
+    iEvent.emplace(outputRecHitsMaskToken_,
+                   Algo_.makeHitsMaskingAsync(static_cast<uint32_t>(output.nHits()), iEvent.queue()));
   }
 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE
 

RecoTracker/PixelSeeding/plugins/alpaka/CAHitNtuplet.cc

@@ -102,6 +102,8 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE {
     using HitsOnDevice = reco::TrackingRecHitsSoACollection;
     using HitsOnHost = ::reco::TrackingRecHitHost;
 
+    using MapToHit = reco::TrackingRecHitsMaskingCollection;
+
     using TkSoAHost = ::reco::TracksHost;
     using TkSoADevice = reco::TracksSoACollection;
 
@@ -331,10 +333,13 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE {
     const edm::ESGetToken<MagneticField, IdealMagneticFieldRecord> tokenField_;
     const device::EDGetToken<HitsOnDevice> tokenHit_;
     const device::EDPutToken<TkSoADevice> tokenTrack_;
+    const device::EDGetToken<MapToHit> tokenHitMask_;
 
     const ::reco::FormulaEvaluator maxNumberOfDoublets_;
     const ::reco::FormulaEvaluator maxNumberOfTuples_;
 
+    const pixelTrack::Iteration iterationName_;
+
     Algo deviceAlgo_;
   };
 
@@ -345,8 +350,10 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE {
         tokenField_(esConsumes()),
         tokenHit_(consumes(iConfig.getParameter<edm::InputTag>("pixelRecHitSrc"))),
         tokenTrack_(produces()),
+        tokenHitMask_(consumes(iConfig.getParameter<edm::InputTag>("hitMask"))),
         maxNumberOfDoublets_(iConfig.getParameter<std::string>("maxNumberOfDoublets")),
         maxNumberOfTuples_(iConfig.getParameter<std::string>("maxNumberOfTuples")),
+        iterationName_(pixelTrack::iterationByName(iConfig.getParameter<std::string>("iterationName"))),
         deviceAlgo_(iConfig) {
     iCache->tokenGeometry_ = esConsumes<edm::Transition::BeginRun>();
     iCache->tokenTopology_ = esConsumes<edm::Transition::BeginRun>();
@@ -357,6 +364,14 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE {
     edm::ParameterSetDescription desc;
 
     desc.add<edm::InputTag>("pixelRecHitSrc", edm::InputTag("siPixelRecHitsPreSplittingAlpaka"));
+    desc.add<edm::InputTag>(
+        "hitMask",
+        edm::InputTag(
+            "hltPhase2PixelRecHitsExtendedSoA"));  // This is just an example, it has to be changed for each tracking iteration
+                                                   // Set as the HLT module to not modify the HLT menu
+    desc.add<std::string>(
+        "iterationName",
+        std::string("promptHighPt"));  // This is just an example, it has to be changed for each tracking iteration
 
     Algo::fillPSetDescription(desc);
     descriptions.addWithDefaultLabel(desc);
@@ -381,8 +396,11 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE {
       uint32_t const maxTuples = maxNumberOfTuples_.evaluate(nHitsV, emptyV);
       uint32_t const maxDoublets = maxNumberOfDoublets_.evaluate(nHitsV, emptyV);
 
+      auto const& mask = iEvent.get(tokenHitMask_);
+
       iEvent.emplace(tokenTrack_,
-                     deviceAlgo_.makeTuplesAsync(hits, geometry, bf, maxDoublets, maxTuples, iEvent.queue()));
+                     deviceAlgo_.makeTuplesAsync(
+                         hits, geometry, bf, maxDoublets, maxTuples, mask, iterationName_, iEvent.queue()));
 
     } else {
       edm::LogWarning("CAHitNtupletAlpaka") << "No hit on BPix1 (" << hits.offsetBPIX2()