HcalRecProducers/src/HcalCPURecHitsProducer.cc

0001 #include <iostream>
0002 #include <string>
0003
0004 #include "CUDADataFormats/HcalRecHitSoA/interface/RecHitCollection.h"
0005 #include "DataFormats/HcalRecHit/interface/HcalRecHitCollections.h"
0006 #include "FWCore/Framework/interface/Event.h"
0007 #include "FWCore/Framework/interface/EventSetup.h"
0008 #include "FWCore/Framework/interface/MakerMacros.h"
0009 #include "FWCore/Framework/interface/stream/EDProducer.h"
0010 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0011 #include "HeterogeneousCore/CUDACore/interface/ScopedContext.h"
0012 #include "HeterogeneousCore/CUDAUtilities/interface/HostAllocator.h"
0013 #include "HeterogeneousCore/CUDAUtilities/interface/cudaCheck.h"
0014
0015 class HcalCPURecHitsProducer : public edm::stream::EDProducer<edm::ExternalWork> {
0016 public:
0017   explicit HcalCPURecHitsProducer(edm::ParameterSet const& ps);
0018   ~HcalCPURecHitsProducer() override = default;
0019
0020   static void fillDescriptions(edm::ConfigurationDescriptions&);
0021
0022 private:
0023   void acquire(edm::Event const&, edm::EventSetup const&, edm::WaitingTaskWithArenaHolder) override;
0024   void produce(edm::Event&, edm::EventSetup const&) override;
0025
0026 private:
0027   const bool produceSoA_;
0028   const bool produceLegacy_;
0029
0030   using IProductType = cms::cuda::Product<hcal::RecHitCollection<calo::common::DevStoragePolicy>>;
0031   const edm::EDGetTokenT<IProductType> recHitsM0TokenIn_;
0032
0033   using OProductType = hcal::RecHitCollection<calo::common::VecStoragePolicy<calo::common::CUDAHostAllocatorAlias>>;
0034   const edm::EDPutTokenT<OProductType> recHitsM0TokenOut_;
0035   const edm::EDPutTokenT<HBHERecHitCollection> recHitsLegacyTokenOut_;
0036
0037   // to pass from acquire to produce
0038   OProductType tmpRecHits_;
0039 };
0040
0041 void HcalCPURecHitsProducer::fillDescriptions(edm::ConfigurationDescriptions& confDesc) {
0042   edm::ParameterSetDescription desc;
0043
0044   desc.add<edm::InputTag>("recHitsM0LabelIn", edm::InputTag{"hbheRecHitProducerGPU"});
0045   desc.add<std::string>("recHitsM0LabelOut", "");
0046   desc.add<std::string>("recHitsLegacyLabelOut", "");
0047   desc.add<bool>("produceSoA", true);
0048   desc.add<bool>("produceLegacy", true);
0049
0050   confDesc.addWithDefaultLabel(desc);
0051 }
0052
0053 HcalCPURecHitsProducer::HcalCPURecHitsProducer(const edm::ParameterSet& ps)
0054     : produceSoA_{ps.getParameter<bool>("produceSoA")},
0055       produceLegacy_{ps.getParameter<bool>("produceLegacy")},
0056       recHitsM0TokenIn_{consumes<IProductType>(ps.getParameter<edm::InputTag>("recHitsM0LabelIn"))},
0057       recHitsM0TokenOut_{produceSoA_ ? produces<OProductType>(ps.getParameter<std::string>("recHitsM0LabelOut"))
0058                                      : edm::EDPutTokenT<OProductType>{}},  // empty token if disabled
0059       recHitsLegacyTokenOut_{produceLegacy_
0060                                  ? produces<HBHERecHitCollection>(ps.getParameter<std::string>("recHitsLegacyLabelOut"))
0061                                  : edm::EDPutTokenT<HBHERecHitCollection>{}}  // empty token if disabled
0062 {}
0063
0064 void HcalCPURecHitsProducer::acquire(edm::Event const& event,
0065                                      edm::EventSetup const& setup,
0066                                      edm::WaitingTaskWithArenaHolder taskHolder) {
0067   // retrieve data/ctx
0068   auto const& recHitsProduct = event.get(recHitsM0TokenIn_);
0069   cms::cuda::ScopedContextAcquire ctx{recHitsProduct, std::move(taskHolder)};
0070   auto const& recHits = ctx.get(recHitsProduct);
0071
0072   // resize tmp buffers
0073   tmpRecHits_.resize(recHits.size);
0074
0075 #ifdef HCAL_MAHI_CPUDEBUG
0076   std::cout << "num rec Hits = " << recHits.size << std::endl;
0077 #endif
0078
0079   // do not try to copy the rechits if they are empty
0080   if (recHits.size == 0) {
0081     return;
0082   }
0083
0084   auto lambdaToTransfer = [&ctx](auto& dest, auto* src) {
0085     using vector_type = typename std::remove_reference<decltype(dest)>::type;
0086     using src_data_type = typename std::remove_pointer<decltype(src)>::type;
0087     using type = typename vector_type::value_type;
0088     static_assert(std::is_same<src_data_type, type>::value && "Dest and Src data types do not match");
0089     cudaCheck(cudaMemcpyAsync(dest.data(), src, dest.size() * sizeof(type), cudaMemcpyDeviceToHost, ctx.stream()));
0090   };
0091
0092   lambdaToTransfer(tmpRecHits_.energy, recHits.energy.get());
0093   lambdaToTransfer(tmpRecHits_.chi2, recHits.chi2.get());
0094   lambdaToTransfer(tmpRecHits_.energyM0, recHits.energyM0.get());
0095   lambdaToTransfer(tmpRecHits_.timeM0, recHits.timeM0.get());
0096   lambdaToTransfer(tmpRecHits_.did, recHits.did.get());
0097 }
0098
0099 void HcalCPURecHitsProducer::produce(edm::Event& event, edm::EventSetup const& setup) {
0100   if (produceLegacy_) {
0101     // populate the legacy collection
0102     auto recHitsLegacy = std::make_unique<HBHERecHitCollection>();
0103     // did not set size with ctor as there is no setter for did
0104     recHitsLegacy->reserve(tmpRecHits_.did.size());
0105     for (uint32_t i = 0; i < tmpRecHits_.did.size(); i++) {
0106       // skip bad channels
0107       if (tmpRecHits_.chi2[i] < 0)
0108         continue;
0109
0110       // build a legacy rechit with the computed detid and MAHI energy
0111       recHitsLegacy->emplace_back(HcalDetId{tmpRecHits_.did[i]},
0112                                   tmpRecHits_.energy[i],
0113                                   0  // timeRising
0114       );
0115       // update the legacy rechit with the Chi2 and M0 values
0116       recHitsLegacy->back().setChiSquared(tmpRecHits_.chi2[i]);
0117       recHitsLegacy->back().setRawEnergy(tmpRecHits_.energyM0[i]);
0118     }
0119
0120     // put the legacy collection
0121     event.put(recHitsLegacyTokenOut_, std::move(recHitsLegacy));
0122   }
0123
0124   if (produceSoA_) {
0125     // put the SoA collection
0126     event.emplace(recHitsM0TokenOut_, std::move(tmpRecHits_));
0127   }
0128   // clear the temporary collection for the next event
0129   tmpRecHits_.resize(0);
0130 }
0131
0132 DEFINE_FWK_MODULE(HcalCPURecHitsProducer);