HcalCPURecHitsProducer.cc

Go to the documentation of this file.

#include <iostream>
#include <string>
 
#include "CUDADataFormats/HcalRecHitSoA/interface/RecHitCollection.h"
#include "DataFormats/HcalRecHit/interface/HcalRecHitCollections.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/EventSetup.h"
#include "FWCore/Framework/interface/MakerMacros.h"
#include "FWCore/Framework/interface/stream/EDProducer.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "HeterogeneousCore/CUDACore/interface/ScopedContext.h"
#include "HeterogeneousCore/CUDAUtilities/interface/HostAllocator.h"
#include "HeterogeneousCore/CUDAUtilities/interface/cudaCheck.h"
 
class HcalCPURecHitsProducer : public edm::stream::EDProducer<edm::ExternalWork> {
public:
  explicit HcalCPURecHitsProducer(edm::ParameterSet const& ps);
  ~HcalCPURecHitsProducer() override = default;
 
  static void fillDescriptions(edm::ConfigurationDescriptions&);
 
private:
  void acquire(edm::Event const&, edm::EventSetup const&, edm::WaitingTaskWithArenaHolder) override;
  void produce(edm::Event&, edm::EventSetup const&) override;
 
private:
  const bool produceSoA_;
  const bool produceLegacy_;
 
  using IProductType = cms::cuda::Product<hcal::RecHitCollection<calo::common::DevStoragePolicy>>;
  const edm::EDGetTokenT<IProductType> recHitsM0TokenIn_;
 
  using OProductType = hcal::RecHitCollection<calo::common::VecStoragePolicy<calo::common::CUDAHostAllocatorAlias>>;
  const edm::EDPutTokenT<OProductType> recHitsM0TokenOut_;
  const edm::EDPutTokenT<HBHERecHitCollection> recHitsLegacyTokenOut_;
 
  // to pass from acquire to produce
  OProductType tmpRecHits_;
};
 
void HcalCPURecHitsProducer::fillDescriptions(edm::ConfigurationDescriptions& confDesc) {
  edm::ParameterSetDescription desc;
 
  desc.add<edm::InputTag>("recHitsM0LabelIn", edm::InputTag{"hbheRecHitProducerGPU"});
  desc.add<std::string>("recHitsM0LabelOut", "");
  desc.add<std::string>("recHitsLegacyLabelOut", "");
  desc.add<bool>("produceSoA", true);
  desc.add<bool>("produceLegacy", true);
 
  confDesc.addWithDefaultLabel(desc);
}
 
HcalCPURecHitsProducer::HcalCPURecHitsProducer(const edm::ParameterSet& ps)
    : produceSoA_{ps.getParameter<bool>("produceSoA")},
      produceLegacy_{ps.getParameter<bool>("produceLegacy")},
      recHitsM0TokenIn_{consumes<IProductType>(ps.getParameter<edm::InputTag>("recHitsM0LabelIn"))},
      recHitsM0TokenOut_{produceSoA_ ? produces<OProductType>(ps.getParameter<std::string>("recHitsM0LabelOut"))
                                     : edm::EDPutTokenT<OProductType>{}},  // empty token if disabled
      recHitsLegacyTokenOut_{produceLegacy_
                                 ? produces<HBHERecHitCollection>(ps.getParameter<std::string>("recHitsLegacyLabelOut"))
                                 : edm::EDPutTokenT<HBHERecHitCollection>{}}  // empty token if disabled
{}
 
void HcalCPURecHitsProducer::acquire(edm::Event const& event,
                                     edm::EventSetup const& setup,
                                     edm::WaitingTaskWithArenaHolder taskHolder) {
  // retrieve data/ctx
  auto const& recHitsProduct = event.get(recHitsM0TokenIn_);
  cms::cuda::ScopedContextAcquire ctx{recHitsProduct, std::move(taskHolder)};
  auto const& recHits = ctx.get(recHitsProduct);
 
  // resize tmp buffers
  tmpRecHits_.resize(recHits.size);
 
#ifdef HCAL_MAHI_CPUDEBUG
  std::cout << "num rec Hits = " << recHits.size << std::endl;
#endif
 
  auto lambdaToTransfer = [&ctx](auto& dest, auto* src) {
    using vector_type = typename std::remove_reference<decltype(dest)>::type;
    using src_data_type = typename std::remove_pointer<decltype(src)>::type;
    using type = typename vector_type::value_type;
    static_assert(std::is_same<src_data_type, type>::value && "Dest and Src data types do not match");
    cudaCheck(cudaMemcpyAsync(dest.data(), src, dest.size() * sizeof(type), cudaMemcpyDeviceToHost, ctx.stream()));
  };
 
  lambdaToTransfer(tmpRecHits_.energy, recHits.energy.get());
  lambdaToTransfer(tmpRecHits_.chi2, recHits.chi2.get());
  lambdaToTransfer(tmpRecHits_.energyM0, recHits.energyM0.get());
  lambdaToTransfer(tmpRecHits_.timeM0, recHits.timeM0.get());
  lambdaToTransfer(tmpRecHits_.did, recHits.did.get());
}
 
void HcalCPURecHitsProducer::produce(edm::Event& event, edm::EventSetup const& setup) {
  if (produceLegacy_) {
    // populate the legacy collection
    auto recHitsLegacy = std::make_unique<HBHERecHitCollection>();
    // did not set size with ctor as there is no setter for did
    recHitsLegacy->reserve(tmpRecHits_.did.size());
    for (uint32_t i = 0; i < tmpRecHits_.did.size(); i++) {
      recHitsLegacy->emplace_back(HcalDetId{tmpRecHits_.did[i]},
                                  tmpRecHits_.energy[i],
                                  0  // timeRising
      );
 
      // update newly pushed guy
      (*recHitsLegacy)[i].setChiSquared(tmpRecHits_.chi2[i]);
      (*recHitsLegacy)[i].setRawEnergy(tmpRecHits_.energyM0[i]);
    }
 
    // put the legacy collection
    event.put(recHitsLegacyTokenOut_, std::move(recHitsLegacy));
  }
 
  if (produceSoA_) {
    // put the SoA collection
    event.emplace(recHitsM0TokenOut_, std::move(tmpRecHits_));
  }
  // clear the temporary collection for the next event
  tmpRecHits_.resize(0);
}
 
DEFINE_FWK_MODULE(HcalCPURecHitsProducer);