CUDATest/plugins/TestCUDAProducerGPUtoCPU.cc

0001 #include "FWCore/Framework/interface/Event.h"
0002 #include "FWCore/Framework/interface/Frameworkfwd.h"
0003 #include "FWCore/Framework/interface/MakerMacros.h"
0004 #include "FWCore/Framework/interface/stream/EDProducer.h"
0005 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0006 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0007
0008 #include "CUDADataFormats/Common/interface/Product.h"
0009 #include "HeterogeneousCore/CUDAUtilities/interface/cudaCheck.h"
0010 #include "HeterogeneousCore/CUDACore/interface/ScopedContext.h"
0011 #include "HeterogeneousCore/CUDATest/interface/Thing.h"
0012 #include "HeterogeneousCore/CUDAUtilities/interface/host_unique_ptr.h"
0013
0014 #include "TestCUDAProducerGPUKernel.h"
0015
0016 class TestCUDAProducerGPUtoCPU : public edm::stream::EDProducer<edm::ExternalWork> {
0017 public:
0018   explicit TestCUDAProducerGPUtoCPU(edm::ParameterSet const& iConfig);
0019   ~TestCUDAProducerGPUtoCPU() override = default;
0020
0021   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0022
0023   void acquire(edm::Event const& iEvent,
0024                edm::EventSetup const& iSetup,
0025                edm::WaitingTaskWithArenaHolder waitingTaskHolder) override;
0026
0027   void produce(edm::Event& iEvent, const edm::EventSetup& iSetup) override;
0028
0029 private:
0030   std::string const label_;
0031   edm::EDGetTokenT<cms::cuda::Product<cms::cudatest::Thing>> const srcToken_;
0032   edm::EDPutTokenT<int> const dstToken_;
0033   cms::cuda::host::unique_ptr<float[]> buffer_;
0034 };
0035
0036 TestCUDAProducerGPUtoCPU::TestCUDAProducerGPUtoCPU(edm::ParameterSet const& iConfig)
0037     : label_{iConfig.getParameter<std::string>("@module_label")},
0038       srcToken_{consumes<cms::cuda::Product<cms::cudatest::Thing>>(iConfig.getParameter<edm::InputTag>("src"))},
0039       dstToken_{produces<int>()} {}
0040
0041 void TestCUDAProducerGPUtoCPU::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0042   edm::ParameterSetDescription desc;
0043   desc.add<edm::InputTag>("src", edm::InputTag())->setComment("Source for cms::cuda::Product<cms::cudatest::Thing>.");
0044   descriptions.addWithDefaultLabel(desc);
0045   descriptions.setComment(
0046       "This EDProducer is part of the TestCUDAProducer* family. It models the GPU->CPU data transfer and formatting of "
0047       "the data to legacy data format. Produces int, to be compatible with TestCUDAProducerCPU.");
0048 }
0049
0050 void TestCUDAProducerGPUtoCPU::acquire(edm::Event const& iEvent,
0051                                        edm::EventSetup const& iSetup,
0052                                        edm::WaitingTaskWithArenaHolder waitingTaskHolder) {
0053   edm::LogVerbatim("TestCUDAProducerGPUtoCPU") << label_ << " TestCUDAProducerGPUtoCPU::acquire begin event "
0054                                                << iEvent.id().event() << " stream " << iEvent.streamID();
0055
0056   auto const& in = iEvent.get(srcToken_);
0057   cms::cuda::ScopedContextAcquire ctx{in, std::move(waitingTaskHolder)};
0058   cms::cudatest::Thing const& device = ctx.get(in);
0059
0060   buffer_ = cms::cuda::make_host_unique<float[]>(TestCUDAProducerGPUKernel::NUM_VALUES, ctx.stream());
0061   // Enqueue async copy, continue in produce once finished
0062   cudaCheck(cudaMemcpyAsync(buffer_.get(),
0063                             device.get(),
0064                             TestCUDAProducerGPUKernel::NUM_VALUES * sizeof(float),
0065                             cudaMemcpyDeviceToHost,
0066                             ctx.stream()));
0067
0068   edm::LogVerbatim("TestCUDAProducerGPUtoCPU") << label_ << " TestCUDAProducerGPUtoCPU::acquire end event "
0069                                                << iEvent.id().event() << " stream " << iEvent.streamID();
0070 }
0071
0072 void TestCUDAProducerGPUtoCPU::produce(edm::Event& iEvent, edm::EventSetup const& iSetup) {
0073   edm::LogVerbatim("TestCUDAProducerGPUtoCPU") << label_ << " TestCUDAProducerGPUtoCPU::produce begin event "
0074                                                << iEvent.id().event() << " stream " << iEvent.streamID();
0075
0076   int counter = 0;
0077   for (int i = 0; i < TestCUDAProducerGPUKernel::NUM_VALUES; ++i) {
0078     counter += buffer_[i];
0079   }
0080   buffer_.reset();  // not so nice, but no way around?
0081
0082   iEvent.emplace(dstToken_, counter);
0083
0084   edm::LogVerbatim("TestCUDAProducerGPUtoCPU")
0085       << label_ << " TestCUDAProducerGPUtoCPU::produce end event " << iEvent.id().event() << " stream "
0086       << iEvent.streamID() << " result " << counter;
0087 }
0088
0089 DEFINE_FWK_MODULE(TestCUDAProducerGPUtoCPU);