plugins/alpaka/TestAlpakaGlobalProducer.cc

0001 #include "DataFormats/PortableTestObjects/interface/alpaka/TestDeviceCollection.h"
0002 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0003 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0004 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0005 #include "FWCore/Utilities/interface/InputTag.h"
0006 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/global/EDProducer.h"
0007 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/EDPutToken.h"
0008 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/ESGetToken.h"
0009 #include "HeterogeneousCore/AlpakaInterface/interface/config.h"
0010 #include "HeterogeneousCore/AlpakaTest/interface/AlpakaESTestRecords.h"
0011 #include "HeterogeneousCore/AlpakaTest/interface/alpaka/AlpakaESTestData.h"
0012
0013 #include "TestAlgo.h"
0014
0015 namespace ALPAKA_ACCELERATOR_NAMESPACE {
0016   /**
0017    * This class demonstrates a global EDProducer that
0018    * - consumes a device ESProduct
0019    * - produces a device EDProduct (that can get transferred to host automatically)
0020    */
0021   class TestAlpakaGlobalProducer : public global::EDProducer<> {
0022   public:
0023     TestAlpakaGlobalProducer(edm::ParameterSet const& config)
0024         : EDProducer<>(config),
0025           esToken_(esConsumes(config.getParameter<edm::ESInputTag>("eventSetupSource"))),
0026           esMultiToken_(esConsumes(config.getParameter<edm::ESInputTag>("eventSetupSourceMulti"))),
0027           deviceToken_{produces()},
0028           deviceTokenMulti2_{produces()},
0029           deviceTokenMulti3_{produces()},
0030           size_{config.getParameter<edm::ParameterSet>("size").getParameter<int32_t>(
0031               EDM_STRINGIZE(ALPAKA_ACCELERATOR_NAMESPACE))},
0032           size2_{config.getParameter<edm::ParameterSet>("size").getParameter<int32_t>(
0033               EDM_STRINGIZE(ALPAKA_ACCELERATOR_NAMESPACE))},
0034           size3_{config.getParameter<edm::ParameterSet>("size").getParameter<int32_t>(
0035               EDM_STRINGIZE(ALPAKA_ACCELERATOR_NAMESPACE))} {}
0036
0037     void produce(edm::StreamID, device::Event& iEvent, device::EventSetup const& iSetup) const override {
0038       [[maybe_unused]] auto const& esData = iSetup.getData(esToken_);
0039       [[maybe_unused]] auto const& esMultiData = iSetup.getData(esMultiToken_);
0040
0041       portabletest::TestDeviceCollection deviceProduct{size_, iEvent.queue()};
0042       portabletest::TestDeviceMultiCollection2 deviceProductMulti2{{{size_, size2_}}, iEvent.queue()};
0043       portabletest::TestDeviceMultiCollection3 deviceProductMulti3{{{size_, size2_, size3_}}, iEvent.queue()};
0044
0045       // run the algorithm, potentially asynchronously
0046       algo_.fill(iEvent.queue(), deviceProduct);
0047       algo_.fillMulti2(iEvent.queue(), deviceProductMulti2);
0048       algo_.fillMulti3(iEvent.queue(), deviceProductMulti3);
0049
0050       iEvent.emplace(deviceToken_, std::move(deviceProduct));
0051       iEvent.emplace(deviceTokenMulti2_, std::move(deviceProductMulti2));
0052       iEvent.emplace(deviceTokenMulti3_, std::move(deviceProductMulti3));
0053     }
0054
0055     static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0056       edm::ParameterSetDescription desc;
0057       desc.add("eventSetupSource", edm::ESInputTag{});
0058       desc.add("eventSetupSourceMulti", edm::ESInputTag{});
0059
0060       edm::ParameterSetDescription psetSize;
0061       psetSize.add<int32_t>("alpaka_serial_sync");
0062       psetSize.add<int32_t>("alpaka_cuda_async");
0063       psetSize.add<int32_t>("alpaka_rocm_async");
0064       desc.add("size", psetSize);
0065
0066       descriptions.addWithDefaultLabel(desc);
0067     }
0068
0069   private:
0070     const device::ESGetToken<AlpakaESTestDataADevice, AlpakaESTestRecordA> esToken_;
0071     const device::ESGetToken<AlpakaESTestDataACMultiDevice, AlpakaESTestRecordA> esMultiToken_;
0072     const device::EDPutToken<portabletest::TestDeviceCollection> deviceToken_;
0073     const device::EDPutToken<portabletest::TestDeviceMultiCollection2> deviceTokenMulti2_;
0074     const device::EDPutToken<portabletest::TestDeviceMultiCollection3> deviceTokenMulti3_;
0075     const int32_t size_;
0076     const int32_t size2_;
0077     const int32_t size3_;
0078
0079     // implementation of the algorithm
0080     TestAlgo algo_;
0081   };
0082
0083 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE
0084
0085 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/MakerMacros.h"
0086 DEFINE_FWK_ALPAKA_MODULE(TestAlpakaGlobalProducer);