Merge pull request #31130 from cms-patatrack/patatrack_integration_4_…

…N_beamspot

Patatrack integration - GPU beamspot data format and transfer (4/N)

CUDADataFormats/BeamSpot/BuildFile.xml

@@ -0,0 +1,8 @@
+<use name="rootcore"/>
+<use name="CUDADataFormats/Common"/>
+<use name="DataFormats/Common"/>
+<use name="HeterogeneousCore/CUDAUtilities"/>
+
+<export>
+    <lib name="1"/>
+</export>

CUDADataFormats/BeamSpot/interface/BeamSpotCUDA.h

@@ -0,0 +1,33 @@
+#ifndef CUDADataFormats_BeamSpot_interface_BeamSpotCUDA_h
+#define CUDADataFormats_BeamSpot_interface_BeamSpotCUDA_h
+
+#include <cuda_runtime.h>
+
+#include "DataFormats/BeamSpot/interface/BeamSpotPOD.h"
+#include "HeterogeneousCore/CUDAUtilities/interface/device_unique_ptr.h"
+
+class BeamSpotCUDA {
+public:
+  // default constructor, required by cms::cuda::Product<BeamSpotCUDA>
+  BeamSpotCUDA() = default;
+
+  // constructor that allocates cached device memory on the given CUDA stream
+  BeamSpotCUDA(cudaStream_t stream) { data_d_ = cms::cuda::make_device_unique<BeamSpotPOD>(stream); }
+
+  // movable, non-copiable
+  BeamSpotCUDA(BeamSpotCUDA const&) = delete;
+  BeamSpotCUDA(BeamSpotCUDA&&) = default;
+  BeamSpotCUDA& operator=(BeamSpotCUDA const&) = delete;
+  BeamSpotCUDA& operator=(BeamSpotCUDA&&) = default;
+
+  BeamSpotPOD* data() { return data_d_.get(); }
+  BeamSpotPOD const* data() const { return data_d_.get(); }
+
+  cms::cuda::device::unique_ptr<BeamSpotPOD>& ptr() { return data_d_; }
+  cms::cuda::device::unique_ptr<BeamSpotPOD> const& ptr() const { return data_d_; }
+
+private:
+  cms::cuda::device::unique_ptr<BeamSpotPOD> data_d_;
+};
+
+#endif  // CUDADataFormats_BeamSpot_interface_BeamSpotCUDA_h

CUDADataFormats/BeamSpot/src/classes.h

@@ -0,0 +1,8 @@
+#ifndef CUDADataFormats_BeamSpot_classes_h
+#define CUDADataFormats_BeamSpot_classes_h
+
+#include "CUDADataFormats/Common/interface/Product.h"
+#include "CUDADataFormats/BeamSpot/interface/BeamSpotCUDA.h"
+#include "DataFormats/Common/interface/Wrapper.h"
+
+#endif  // CUDADataFormats_BeamSpot_classes_h

CUDADataFormats/BeamSpot/src/classes_def.xml

@@ -0,0 +1,4 @@
+<lcgdict>
+  <class name="cms::cuda::Product<BeamSpotCUDA>" persistent="false"/>
+  <class name="edm::Wrapper<cms::cuda::Product<BeamSpotCUDA>>" persistent="false"/>
+</lcgdict>

Configuration/StandardSequences/python/Reconstruction_cff.py

@@ -103,7 +103,7 @@
 from RecoLocalCalo.Castor.Castor_cff import *
 from RecoLocalCalo.Configuration.hcalGlobalReco_cff import *
 
-globalreco_trackingTask = cms.Task(offlineBeamSpot,
+globalreco_trackingTask = cms.Task(offlineBeamSpotTask,
                           MeasurementTrackerEventPreSplitting, # unclear where to put this
                           siPixelClusterShapeCachePreSplitting, # unclear where to put this
                           standalonemuontrackingTask,
@@ -117,7 +117,7 @@
 ##########################################
 # offlineBeamSpot is reconstructed before mixing in fastSim
 ##########################################
-_fastSim_globalreco_trackingTask = globalreco_trackingTask.copyAndExclude([offlineBeamSpot,MeasurementTrackerEventPreSplitting,siPixelClusterShapeCachePreSplitting])
+_fastSim_globalreco_trackingTask = globalreco_trackingTask.copyAndExclude([offlineBeamSpotTask,MeasurementTrackerEventPreSplitting,siPixelClusterShapeCachePreSplitting])
 fastSim.toReplaceWith(globalreco_trackingTask,_fastSim_globalreco_trackingTask)
 
 _phase2_timing_layer_globalreco_trackingTask = globalreco_trackingTask.copy()
@@ -212,7 +212,7 @@
 reconstruction_trackingOnly = cms.Sequence(reconstruction_trackingOnlyTask)
 reconstruction_pixelTrackingOnlyTask = cms.Task(
     pixeltrackerlocalrecoTask,
-    offlineBeamSpot,
+    offlineBeamSpotTask,
     siPixelClusterShapeCachePreSplitting,
     recopixelvertexingTask
 )

DataFormats/BeamSpot/interface/BeamSpotPOD.h

@@ -0,0 +1,20 @@
+#ifndef DataFormats_BeamSpot_interface_BeamSpotPOD_h
+#define DataFormats_BeamSpot_interface_BeamSpotPOD_h
+
+// This struct is a transient-only, simplified representation of the beamspot
+// data used as the underlying type for data transfers and operations in
+// heterogeneous code (e.g. in CUDA code).
+
+// The covariance matrix is not used in that code, so is left out here.
+
+// align to the CUDA L1 cache line size
+struct alignas(128) BeamSpotPOD {
+  float x, y, z;  // position
+  float sigmaZ;
+  float beamWidthX, beamWidthY;
+  float dxdz, dydz;
+  float emittanceX, emittanceY;
+  float betaStar;
+};
+
+#endif  // DataFormats_BeamSpot_interface_BeamSpotPOD_h

DataFormats/BeamSpot/src/classes.h

@@ -1,8 +1,3 @@
-
-#include "DataFormats/Common/interface/Wrapper.h"
 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
-#include "Math/Cartesian3D.h"
-#include "DataFormats/Math/interface/Vector3D.h"
-#include "Math/Polar3D.h"
-#include "Math/CylindricalEta3D.h"
-#include "DataFormats/Math/interface/Vector.h"
+#include "DataFormats/BeamSpot/interface/BeamSpotPOD.h"
+#include "DataFormats/Common/interface/Wrapper.h"

DataFormats/BeamSpot/src/classes_def.xml

@@ -4,4 +4,7 @@
    <version ClassVersion="10" checksum="915260007"/>
   </class>
   <class name="edm::Wrapper<reco::BeamSpot>"/>
+
+  <class name="BeamSpotPOD" persistent="false"/>
+  <class name="edm::Wrapper<BeamSpotPOD>" persistent="false"/>
 </lcgdict>

RecoVertex/BeamSpotProducer/plugins/BeamSpotToCUDA.cc

@@ -0,0 +1,101 @@
+#include <cuda_runtime.h>
+
+#include "CUDADataFormats/BeamSpot/interface/BeamSpotCUDA.h"
+#include "CUDADataFormats/Common/interface/Product.h"
+#include "DataFormats/BeamSpot/interface/BeamSpot.h"
+#include "DataFormats/BeamSpot/interface/BeamSpotPOD.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+#include "FWCore/Framework/interface/global/EDProducer.h"
+#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
+#include "FWCore/ServiceRegistry/interface/Service.h"
+#include "HeterogeneousCore/CUDACore/interface/ScopedContext.h"
+#include "HeterogeneousCore/CUDAServices/interface/CUDAService.h"
+#include "HeterogeneousCore/CUDAUtilities/interface/copyAsync.h"
+#include "HeterogeneousCore/CUDAUtilities/interface/host_noncached_unique_ptr.h"
+
+namespace {
+
+  class BeamSpotHost {
+  public:
+    BeamSpotHost() : data_h_{cms::cuda::make_host_noncached_unique<BeamSpotPOD>(cudaHostAllocWriteCombined)} {}
+
+    BeamSpotHost(BeamSpotHost const&) = delete;
+    BeamSpotHost(BeamSpotHost&&) = default;
+
+    BeamSpotHost& operator=(BeamSpotHost const&) = delete;
+    BeamSpotHost& operator=(BeamSpotHost&&) = default;
+
+    BeamSpotPOD* data() { return data_h_.get(); }
+    BeamSpotPOD const* data() const { return data_h_.get(); }
+
+    cms::cuda::host::noncached::unique_ptr<BeamSpotPOD>& ptr() { return data_h_; }
+    cms::cuda::host::noncached::unique_ptr<BeamSpotPOD> const& ptr() const { return data_h_; }
+
+  private:
+    cms::cuda::host::noncached::unique_ptr<BeamSpotPOD> data_h_;
+  };
+
+}  // namespace
+
+class BeamSpotToCUDA : public edm::global::EDProducer<edm::StreamCache<BeamSpotHost>> {
+public:
+  explicit BeamSpotToCUDA(const edm::ParameterSet& iConfig);
+  ~BeamSpotToCUDA() override = default;
+
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
+
+  std::unique_ptr<BeamSpotHost> beginStream(edm::StreamID) const override {
+    edm::Service<CUDAService> cs;
+    if (cs->enabled()) {
+      return std::make_unique<BeamSpotHost>();
+    } else {
+      return nullptr;
+    }
+  }
+  void produce(edm::StreamID streamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const override;
+
+private:
+  const edm::EDGetTokenT<reco::BeamSpot> bsGetToken_;
+  const edm::EDPutTokenT<cms::cuda::Product<BeamSpotCUDA>> bsPutToken_;
+};
+
+BeamSpotToCUDA::BeamSpotToCUDA(const edm::ParameterSet& iConfig)
+    : bsGetToken_{consumes<reco::BeamSpot>(iConfig.getParameter<edm::InputTag>("src"))},
+      bsPutToken_{produces<cms::cuda::Product<BeamSpotCUDA>>()} {}
+
+void BeamSpotToCUDA::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  edm::ParameterSetDescription desc;
+  desc.add<edm::InputTag>("src", edm::InputTag("offlineBeamSpot"));
+  descriptions.add("offlineBeamSpotToCUDA", desc);
+}
+
+void BeamSpotToCUDA::produce(edm::StreamID streamID, edm::Event& iEvent, const edm::EventSetup& iSetup) const {
+  cms::cuda::ScopedContextProduce ctx{streamID};
+
+  const reco::BeamSpot& bs = iEvent.get(bsGetToken_);
+
+  auto& bsHost = streamCache(streamID)->ptr();
+
+  bsHost->x = bs.x0();
+  bsHost->y = bs.y0();
+  bsHost->z = bs.z0();
+
+  bsHost->sigmaZ = bs.sigmaZ();
+  bsHost->beamWidthX = bs.BeamWidthX();
+  bsHost->beamWidthY = bs.BeamWidthY();
+  bsHost->dxdz = bs.dxdz();
+  bsHost->dydz = bs.dydz();
+  bsHost->emittanceX = bs.emittanceX();
+  bsHost->emittanceY = bs.emittanceY();
+  bsHost->betaStar = bs.betaStar();
+
+  BeamSpotCUDA bsDevice(ctx.stream());
+  cms::cuda::copyAsync(bsDevice.ptr(), bsHost, ctx.stream());
+
+  ctx.emplace(iEvent, bsPutToken_, std::move(bsDevice));
+}
+
+DEFINE_FWK_MODULE(BeamSpotToCUDA);

RecoVertex/BeamSpotProducer/plugins/BuildFile.xml

@@ -1,13 +1,12 @@
+<use name="root"/>
+<use name="rootminuit"/>
+<use name="CondCore/DBOutputService"/>
+<use name="CondFormats/BeamSpotObjects"/>
+<use name="CondFormats/DataRecord"/>
 <use name="FWCore/Framework"/>
 <use name="FWCore/ParameterSet"/>
-<use name="FWCore/Utilities"/>
 <use name="FWCore/ServiceRegistry"/>
-<use name="CondFormats/BeamSpotObjects"/>
-<use name="CondFormats/DataRecord"/>
-<use name="CondCore/DBOutputService"/>
-
-<use name="root"/>
-<use name="rootminuit"/>
+<use name="FWCore/Utilities"/>
 <use name="RecoVertex/BeamSpotProducer"/>
 
 <library file="BeamSpotProducer.cc" name="BeamSpotProducer">
@@ -40,4 +39,10 @@
 <library file="OfflineToTransientBeamSpotESProducer.cc" name="OfflineToTransientBeamSpotESProducer">
   <flags EDM_PLUGIN="1"/>
 </library>
-
+<library file="BeamSpotToCUDA.cc" name="BeamSpotToCUDA">
+  <use name="cuda"/>
+  <use name="CUDADataFormats/BeamSpot"/>
+  <use name="HeterogeneousCore/CUDACore"/>
+  <use name="HeterogeneousCore/CUDAServices"/>
+  <flags EDM_PLUGIN="1"/>
+</library>

RecoVertex/BeamSpotProducer/python/BeamSpot_cff.py

@@ -1,4 +1,11 @@
 import FWCore.ParameterSet.Config as cms
 
 from RecoVertex.BeamSpotProducer.BeamSpot_cfi import *
+from RecoVertex.BeamSpotProducer.offlineBeamSpotToCUDA_cfi import offlineBeamSpotToCUDA
 
+offlineBeamSpotTask = cms.Task(offlineBeamSpot)
+
+from Configuration.ProcessModifiers.gpu_cff import gpu
+_offlineBeamSpotTask_gpu = offlineBeamSpotTask.copy()
+_offlineBeamSpotTask_gpu.add(offlineBeamSpotToCUDA)
+gpu.toReplaceWith(offlineBeamSpotTask, _offlineBeamSpotTask_gpu)