interface/alpaka/ProducerBase.h

0001 #ifndef HeterogeneousCore_AlpakaCore_interface_ProducerBase_h
0002 #define HeterogeneousCore_AlpakaCore_interface_ProducerBase_h
0003
0004 #include "DataFormats/Common/interface/DeviceProduct.h"
0005 #include "FWCore/Framework/interface/FrameworkfwdMostUsed.h"
0006 #include "FWCore/Framework/interface/moduleAbilities.h"
0007 #include "FWCore/Framework/interface/Event.h"
0008 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0009 #include "FWCore/Utilities/interface/EDPutToken.h"
0010 #include "FWCore/Utilities/interface/Transition.h"
0011 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/DeviceProductType.h"
0012 #include "HeterogeneousCore/AlpakaCore/interface/alpaka/EDMetadataAcquireSentry.h"
0013 #include "HeterogeneousCore/AlpakaCore/interface/modulePrevalidate.h"
0014 #include "HeterogeneousCore/AlpakaInterface/interface/Backend.h"
0015 #include "HeterogeneousCore/AlpakaInterface/interface/CopyToDevice.h"
0016 #include "HeterogeneousCore/AlpakaInterface/interface/CopyToHost.h"
0017
0018 #include <memory>
0019 #include <tuple>
0020
0021 namespace ALPAKA_ACCELERATOR_NAMESPACE {
0022   template <typename Producer, edm::Transition Tr>
0023   class ProducerBaseAdaptor;
0024
0025   /**
0026    * The ProducerBase acts as a common base class for all Alpaka
0027    * EDProducers. The main benefit is to have a single place for the
0028    * definition of produces() functions.
0029    *
0030    * The produces() functions return a custom ProducerBaseAdaptor in
0031    * order to call the deviceProduces(). For device or asynchronous
0032    * backends the deviceProduces() registers the automatic copy to
0033    * host and a transformation from edm::DeviceProduct<T> to U, where
0034    * U is the host-equivalent of T. The transformation from T to U is
0035    * done by a specialization of cms::alpakatools::CopyToHost<T> class
0036    * template, that should be provided in the same file where T is
0037    * defined
0038    *
0039    * TODO: add "override" for labelsForToken()
0040    */
0041   template <template <typename...> class BaseT, typename... Args>
0042   class ProducerBase : public BaseT<Args..., edm::Transformer> {
0043     static_assert(not edm::CheckAbility<edm::module::Abilities::kTransformer>::kHasIt,
0044                   "ALPAKA_ACCELERATOR_NAMESPACE::ProducerBase can not be used with Transformer ability (as it is "
0045                   "used internally)");
0046     using Base = BaseT<Args..., edm::Transformer>;
0047
0048   public:
0049     ProducerBase(edm::ParameterSet const& iConfig)
0050         : backendToken_(Base::produces("backend")),
0051           // The 'synchronize' parameter can be unset in Alpaka
0052           // modules specified with the namespace prefix instead if
0053           // '@alpaka' suffix
0054           synchronize_(iConfig.getUntrackedParameter<edm::ParameterSet>("alpaka").getUntrackedParameter<bool>(
0055               "synchronize", false)) {}
0056
0057     template <edm::Transition Tr = edm::Transition::Event>
0058     [[nodiscard]] auto produces() noexcept {
0059       return ProducerBaseAdaptor<ProducerBase, Tr>(*this);
0060     }
0061
0062     template <edm::Transition Tr = edm::Transition::Event>
0063     [[nodiscard]] auto produces(std::string instanceName) noexcept {
0064       return ProducerBaseAdaptor<ProducerBase, Tr>(*this, std::move(instanceName));
0065     }
0066
0067     static void prevalidate(edm::ConfigurationDescriptions& descriptions) {
0068       Base::prevalidate(descriptions);
0069       cms::alpakatools::modulePrevalidate(descriptions);
0070     }
0071
0072   protected:
0073     void putBackend(edm::Event& iEvent) const {
0074       iEvent.emplace(this->backendToken_, static_cast<unsigned short>(kBackend));
0075     }
0076
0077     bool synchronize() const { return synchronize_; }
0078
0079   private:
0080     edm::EDPutTokenT<unsigned short> const backendToken_;
0081     bool const synchronize_ = false;
0082
0083     template <typename TProducer, edm::Transition Tr>
0084     friend class ProducerBaseAdaptor;
0085
0086     // Host products
0087     //
0088     // intentionally not returning BranchAliasSetter
0089     // can think of it later if really needed
0090     template <typename TProduct, edm::Transition Tr>
0091     edm::EDPutTokenT<TProduct> produces(std::string instanceName) {
0092       constexpr bool hasCopy = requires(Queue& queue, TProduct const& prod) {
0093         cms::alpakatools::CopyToDevice<TProduct>::copyAsync(queue, prod);
0094       };
0095
0096       if constexpr (detail::useProductDirectly or not hasCopy) {
0097         return Base::template produces<TProduct, Tr>(std::move(instanceName));
0098       } else {
0099         edm::EDPutTokenT<TProduct> hostToken = Base::template produces<TProduct, Tr>(instanceName);
0100         this->registerTransformAsync(
0101             hostToken,
0102             [synchronize = this->synchronize()](
0103                 edm::StreamID streamID, TProduct const& hostProduct, edm::WaitingTaskWithArenaHolder holder) {
0104               detail::EDMetadataAcquireSentry sentry(streamID, std::move(holder), synchronize);
0105               using CopyT = cms::alpakatools::CopyToDevice<TProduct>;
0106               auto productOnDevice = CopyT::copyAsync(sentry.metadata()->queue(), hostProduct);
0107               // Need to keep the EDMetadata object from sentry.finish()
0108               // alive until the synchronization
0109               using TplType = std::tuple<std::shared_ptr<EDMetadata>, decltype(productOnDevice)>;
0110               // Wrap possibly move-only type into a copyable type
0111               return std::make_shared<TplType>(sentry.finish(), std::move(productOnDevice));
0112             },
0113             [](edm::StreamID, auto tplPtr) {
0114               using DeviceObject = std::tuple_element_t<1, std::remove_cvref_t<decltype(*tplPtr)>>;
0115               using DeviceProductType = detail::DeviceProductType<DeviceObject>;
0116               return DeviceProductType(std::move(std::get<0>(*tplPtr)), std::move(std::get<1>(*tplPtr)));
0117             },
0118             std::move(instanceName));
0119         return hostToken;
0120       }
0121     }
0122
0123     // Device products
0124     //
0125     // intentionally not returning BranchAliasSetter
0126     // can think of it later if really needed
0127     template <typename TProduct, typename TToken, edm::Transition Tr>
0128     edm::EDPutTokenT<TToken> deviceProduces(std::string instanceName) {
0129       if constexpr (detail::useProductDirectly) {
0130         return Base::template produces<TToken, Tr>(std::move(instanceName));
0131       } else {
0132         edm::EDPutTokenT<TToken> token = Base::template produces<TToken, Tr>(instanceName);
0133         using CopyT = cms::alpakatools::CopyToHost<TProduct>;
0134         this->registerTransformAsync(
0135             token,
0136             [](edm::StreamID, TToken const& deviceProduct, edm::WaitingTaskWithArenaHolder holder) {
0137               auto const& device = alpaka::getDev(deviceProduct.template metadata<EDMetadata>().queue());
0138               detail::EDMetadataAcquireSentry sentry(device, std::move(holder));
0139               auto metadataPtr = sentry.metadata();
0140               constexpr bool tryReuseQueue = true;
0141               TProduct const& productOnDevice =
0142                   deviceProduct.template getSynchronized<EDMetadata>(*metadataPtr, tryReuseQueue);
0143
0144               auto productOnHost = CopyT::copyAsync(metadataPtr->queue(), productOnDevice);
0145
0146               // Need to keep the EDMetadata object from sentry.finish()
0147               // alive until the synchronization
0148               using TplType = std::tuple<decltype(productOnHost), std::shared_ptr<EDMetadata>>;
0149               // Wrap possibly move-only type into a copyable type
0150               return std::make_shared<TplType>(std::move(productOnHost), sentry.finish());
0151             },
0152             [](edm::StreamID, auto tplPtr) {
0153               auto& productOnHost = std::get<0>(*tplPtr);
0154               if constexpr (requires { CopyT::postCopy(productOnHost); }) {
0155                 CopyT::postCopy(productOnHost);
0156               }
0157               return std::move(productOnHost);
0158             },
0159             std::move(instanceName));
0160         return token;
0161       }
0162     }
0163   };
0164
0165   // Adaptor class to make the type-deducing produces() calls to work
0166   template <typename TProducer, edm::Transition Tr>
0167   class ProducerBaseAdaptor {
0168   public:
0169     // for host-only products
0170     template <typename Type>
0171     edm::EDPutTokenT<Type> produces() {
0172       return producer_.template produces<Type, Tr>(label_);
0173     }
0174
0175     // for device products
0176     template <typename TProduct, typename TToken>
0177     edm::EDPutTokenT<TToken> deviceProduces() {
0178       return producer_.template deviceProduces<TProduct, TToken, Tr>(label_);
0179     }
0180
0181   private:
0182     // only ProducerBase is allowed to make an instance of this class
0183     friend TProducer;
0184
0185     ProducerBaseAdaptor(TProducer& producer, std::string label) : producer_(producer), label_(std::move(label)) {}
0186     explicit ProducerBaseAdaptor(TProducer& producer) : producer_(producer) {}
0187
0188     TProducer& producer_;
0189     std::string const label_;
0190   };
0191 }  // namespace ALPAKA_ACCELERATOR_NAMESPACE
0192
0193 #endif