CUDATest/plugins/TestPortableTranscriber.cc

0001 #include "CUDADataFormats/Common/interface/Product.h"
0002 #include "CUDADataFormats/PortableTestObjects/interface/TestDeviceCollection.h"
0003 #include "CUDADataFormats/PortableTestObjects/interface/TestHostCollection.h"
0004 #include "FWCore/Framework/interface/Event.h"
0005 #include "FWCore/Framework/interface/EventSetup.h"
0006 #include "FWCore/Framework/interface/Frameworkfwd.h"
0007 #include "FWCore/Framework/interface/stream/EDProducer.h"
0008 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0009 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0010 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0011 #include "FWCore/Utilities/interface/EDGetToken.h"
0012 #include "FWCore/Utilities/interface/InputTag.h"
0013 #include "FWCore/Utilities/interface/StreamID.h"
0014 #include "HeterogeneousCore/CUDACore/interface/ScopedContext.h"
0015 #include "HeterogeneousCore/CUDAUtilities/interface/copyAsync.h"
0016
0017 class TestPortableTranscriber : public edm::stream::EDProducer<edm::ExternalWork> {
0018 public:
0019   TestPortableTranscriber(edm::ParameterSet const& config)
0020       : deviceToken_{consumes(config.getParameter<edm::InputTag>("source"))}, hostToken_{produces()} {}
0021
0022   void acquire(edm::Event const& event, edm::EventSetup const& setup, edm::WaitingTaskWithArenaHolder task) override {
0023     // create a context reusing the same device and queue as the producer of the input collection
0024     auto const& input = event.get(deviceToken_);
0025     cms::cuda::ScopedContextAcquire ctx{input, std::move(task)};
0026
0027     cudatest::TestDeviceCollection const& deviceProduct = ctx.get(input);
0028
0029     // allocate a host product based on the metadata of the device product
0030     hostProduct_ = cudatest::TestHostCollection{deviceProduct->metadata().size(), ctx.stream()};
0031
0032     // copy the content of the device product to the host product
0033     cms::cuda::copyAsync(hostProduct_.buffer(), deviceProduct.const_buffer(), deviceProduct.bufferSize(), ctx.stream());
0034
0035     // do not wait for the asynchronous operation to complete
0036   }
0037
0038   void produce(edm::Event& event, edm::EventSetup const&) override {
0039     // produce() is called once the asynchronous operation has completed, so there is no need for an explicit wait
0040     event.emplace(hostToken_, std::move(hostProduct_));
0041   }
0042
0043   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0044     edm::ParameterSetDescription desc;
0045     desc.add<edm::InputTag>("source");
0046     descriptions.addWithDefaultLabel(desc);
0047   }
0048
0049 private:
0050   const edm::EDGetTokenT<cms::cuda::Product<cudatest::TestDeviceCollection>> deviceToken_;
0051   const edm::EDPutTokenT<cudatest::TestHostCollection> hostToken_;
0052
0053   // hold the output product between acquire() and produce()
0054   cudatest::TestHostCollection hostProduct_;
0055 };
0056
0057 #include "FWCore/Framework/interface/MakerMacros.h"
0058 DEFINE_FWK_MODULE(TestPortableTranscriber);