PixelVertexFinding/plugins/PixelVertexSoAFromCUDA.cc

0001 #include <cuda_runtime.h>
0002
0003 #include "CUDADataFormats/Common/interface/Product.h"
0004 #include "CUDADataFormats/Common/interface/HostProduct.h"
0005 #include "CUDADataFormats/Vertex/interface/ZVertexSoAHeterogeneousHost.h"
0006 #include "CUDADataFormats/Vertex/interface/ZVertexSoAHeterogeneousDevice.h"
0007 #include "DataFormats/Common/interface/Handle.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/stream/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 "HeterogeneousCore/CUDACore/interface/ScopedContext.h"
0020
0021 class PixelVertexSoAFromCUDA : public edm::stream::EDProducer<edm::ExternalWork> {
0022 public:
0023   explicit PixelVertexSoAFromCUDA(const edm::ParameterSet& iConfig);
0024   ~PixelVertexSoAFromCUDA() override = default;
0025
0026   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0027
0028 private:
0029   void acquire(edm::Event const& iEvent,
0030                edm::EventSetup const& iSetup,
0031                edm::WaitingTaskWithArenaHolder waitingTaskHolder) override;
0032   void produce(edm::Event& iEvent, edm::EventSetup const& iSetup) override;
0033
0034   edm::EDGetTokenT<cms::cuda::Product<ZVertexSoADevice>> tokenCUDA_;
0035   edm::EDPutTokenT<ZVertexSoAHost> tokenSOA_;
0036
0037   ZVertexSoAHost zvertex_h;
0038 };
0039
0040 PixelVertexSoAFromCUDA::PixelVertexSoAFromCUDA(const edm::ParameterSet& iConfig)
0041     : tokenCUDA_(consumes<cms::cuda::Product<ZVertexSoADevice>>(iConfig.getParameter<edm::InputTag>("src"))),
0042       tokenSOA_(produces<ZVertexSoAHost>()) {}
0043
0044 void PixelVertexSoAFromCUDA::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0045   edm::ParameterSetDescription desc;
0046
0047   desc.add<edm::InputTag>("src", edm::InputTag("pixelVerticesCUDA"));
0048   descriptions.add("pixelVerticesSoA", desc);
0049 }
0050
0051 void PixelVertexSoAFromCUDA::acquire(edm::Event const& iEvent,
0052                                      edm::EventSetup const& iSetup,
0053                                      edm::WaitingTaskWithArenaHolder waitingTaskHolder) {
0054   cms::cuda::Product<ZVertexSoADevice> const& inputDataWrapped = iEvent.get(tokenCUDA_);
0055   cms::cuda::ScopedContextAcquire ctx{inputDataWrapped, std::move(waitingTaskHolder)};
0056   auto const& zvertex_d = ctx.get(inputDataWrapped);  // Tracks on device
0057   zvertex_h = ZVertexSoAHost(ctx.stream());           // Create an instance of Tracks on Host, using the stream
0058   cudaCheck(cudaMemcpyAsync(zvertex_h.buffer().get(),
0059                             zvertex_d.const_buffer().get(),
0060                             zvertex_d.bufferSize(),
0061                             cudaMemcpyDeviceToHost,
0062                             ctx.stream()));  // Copy data from Device to Host
0063 }
0064
0065 void PixelVertexSoAFromCUDA::produce(edm::Event& iEvent, edm::EventSetup const& iSetup) {
0066   // No copies....
0067   iEvent.emplace(tokenSOA_, std::move(zvertex_h));
0068 }
0069
0070 DEFINE_FWK_MODULE(PixelVertexSoAFromCUDA);