plugins/alpaka/TestAlpakaStreamProducer.cc

0001 #include "DataFormats/PortableTestObjects/interface/alpaka/TestDeviceCollection.h"
0002 #include "DataFormats/TestObjects/interface/ToyProducts.h"
0003 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0004 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0005 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0006 #include "FWCore/Utilities/interface/InputTag.h"
0007 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/stream/EDProducer.h"
0008 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/EDPutToken.h"
0009 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/ESGetToken.h"
0010 #include "HeterogeneousCore/AlpakaInterface/interface/config.h"
0011 #include "HeterogeneousCore/AlpakaTest/interface/AlpakaESTestRecords.h"
0012 #include "HeterogeneousCore/AlpakaTest/interface/AlpakaESTestData.h"
0013
0014 #include "TestAlgo.h"
0015
0016 namespace ALPAKA_ACCELERATOR_NAMESPACE {
0017   /**
0018    * This class demonstrates a stream EDProducer that
0019    * - consumes a host EDProduct
0020    * - consumes a device ESProduct
0021    * - produces a device EDProduct (that gets transferred to host automatically if needed)
0022    * - optionally uses a product instance label
0023    */
0024   class TestAlpakaStreamProducer : public stream::EDProducer<> {
0025   public:
0026     TestAlpakaStreamProducer(edm::ParameterSet const& config)
0027         : EDProducer<>(config),
0028           size_{config.getParameter<edm::ParameterSet>("size").getParameter<int32_t>(
0029               EDM_STRINGIZE(ALPAKA_ACCELERATOR_NAMESPACE))},
0030           size2_{config.getParameter<edm::ParameterSet>("size").getParameter<int32_t>(
0031               EDM_STRINGIZE(ALPAKA_ACCELERATOR_NAMESPACE))},
0032           size3_{config.getParameter<edm::ParameterSet>("size").getParameter<int32_t>(
0033               EDM_STRINGIZE(ALPAKA_ACCELERATOR_NAMESPACE))} {
0034       getToken_ = consumes(config.getParameter<edm::InputTag>("source"));
0035       esToken_ = esConsumes(config.getParameter<edm::ESInputTag>("eventSetupSource"));
0036       devicePutToken_ = produces(config.getParameter<std::string>("productInstanceName"));
0037       devicePutTokenMulti2_ = produces(config.getParameter<std::string>("productInstanceName"));
0038       devicePutTokenMulti3_ = produces(config.getParameter<std::string>("productInstanceName"));
0039     }
0040
0041     void produce(device::Event& iEvent, device::EventSetup const& iSetup) override {
0042       [[maybe_unused]] auto inpData = iEvent.getHandle(getToken_);
0043       [[maybe_unused]] auto const& esData = iSetup.getData(esToken_);
0044
0045       auto deviceProduct = std::make_unique<portabletest::TestDeviceCollection>(size_, iEvent.queue());
0046       auto deviceProductMulti2 = std::make_unique<portabletest::TestDeviceMultiCollection2>(
0047           portabletest::TestDeviceMultiCollection2::SizesArray{{size_, size2_}}, iEvent.queue());
0048       auto deviceProductMulti3 = std::make_unique<portabletest::TestDeviceMultiCollection3>(
0049           portabletest::TestDeviceMultiCollection3::SizesArray{{size_, size2_, size3_}}, iEvent.queue());
0050
0051       // run the algorithm, potentially asynchronously
0052       algo_.fill(iEvent.queue(), *deviceProduct);
0053       algo_.fillMulti2(iEvent.queue(), *deviceProductMulti2);
0054       algo_.fillMulti3(iEvent.queue(), *deviceProductMulti3);
0055
0056       iEvent.put(devicePutToken_, std::move(deviceProduct));
0057       iEvent.put(devicePutTokenMulti2_, std::move(deviceProductMulti2));
0058       iEvent.put(devicePutTokenMulti3_, std::move(deviceProductMulti3));
0059     }
0060
0061     static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0062       edm::ParameterSetDescription desc;
0063       desc.add<edm::InputTag>("source");
0064       desc.add("eventSetupSource", edm::ESInputTag{});
0065       desc.add<std::string>("productInstanceName", "");
0066
0067       edm::ParameterSetDescription psetSize;
0068       psetSize.add<int32_t>("alpaka_serial_sync");
0069       psetSize.add<int32_t>("alpaka_cuda_async");
0070       psetSize.add<int32_t>("alpaka_rocm_async");
0071       desc.add("size", psetSize);
0072
0073       descriptions.addWithDefaultLabel(desc);
0074     }
0075
0076   private:
0077     edm::EDGetTokenT<edmtest::IntProduct> getToken_;
0078     device::ESGetToken<cms::alpakatest::AlpakaESTestDataB<Device>, AlpakaESTestRecordB> esToken_;
0079     device::EDPutToken<portabletest::TestDeviceCollection> devicePutToken_;
0080     device::EDPutToken<portabletest::TestDeviceMultiCollection2> devicePutTokenMulti2_;
0081     device::EDPutToken<portabletest::TestDeviceMultiCollection3> devicePutTokenMulti3_;
0082     const int32_t size_;
0083     const int32_t size2_;
0084     const int32_t size3_;
0085
0086     // implementation of the algorithm
0087     TestAlgo algo_;
0088   };
0089
0090 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE
0091
0092 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/MakerMacros.h"
0093 DEFINE_FWK_ALPAKA_MODULE(TestAlpakaStreamProducer);