PixelTrackFitting/plugins/PixelTrackSoAFromCUDA.cc

0001 #include <cuda_runtime.h>
0002 #include <Eigen/Core>  // needed here by soa layout
0003
0004 #include "CUDADataFormats/Common/interface/Product.h"
0005 #include "CUDADataFormats/Common/interface/HostProduct.h"
0006 #include "CUDADataFormats/Track/interface/TrackSoAHeterogeneousHost.h"
0007 #include "CUDADataFormats/Track/interface/TrackSoAHeterogeneousDevice.h"
0008 #include "CUDADataFormats/Track/interface/PixelTrackUtilities.h"
0009 #include "DataFormats/Common/interface/Handle.h"
0010 #include "FWCore/Framework/interface/Event.h"
0011 #include "FWCore/Framework/interface/EventSetup.h"
0012 #include "FWCore/Framework/interface/MakerMacros.h"
0013 #include "FWCore/Framework/interface/stream/EDProducer.h"
0014 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0015 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0016 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0017 #include "FWCore/PluginManager/interface/ModuleDef.h"
0018 #include "FWCore/Utilities/interface/EDGetToken.h"
0019 #include "FWCore/Utilities/interface/InputTag.h"
0020 #include "HeterogeneousCore/CUDACore/interface/ScopedContext.h"
0021
0022 // Switch on to enable checks and printout for found tracks
0023 // #define PIXEL_DEBUG_PRODUCE
0024
0025 template <typename TrackerTraits>
0026 class PixelTrackSoAFromCUDAT : public edm::stream::EDProducer<edm::ExternalWork> {
0027   using TrackSoAHost = TrackSoAHeterogeneousHost<TrackerTraits>;
0028   using TrackSoADevice = TrackSoAHeterogeneousDevice<TrackerTraits>;
0029
0030 public:
0031   explicit PixelTrackSoAFromCUDAT(const edm::ParameterSet& iConfig);
0032   ~PixelTrackSoAFromCUDAT() override = default;
0033
0034   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0035
0036 private:
0037   void acquire(edm::Event const& iEvent,
0038                edm::EventSetup const& iSetup,
0039                edm::WaitingTaskWithArenaHolder waitingTaskHolder) override;
0040   void produce(edm::Event& iEvent, edm::EventSetup const& iSetup) override;
0041
0042   edm::EDGetTokenT<cms::cuda::Product<TrackSoADevice>> tokenCUDA_;
0043   edm::EDPutTokenT<TrackSoAHost> tokenSOA_;
0044
0045   TrackSoAHost tracks_h_;
0046 };
0047
0048 template <typename TrackerTraits>
0049 PixelTrackSoAFromCUDAT<TrackerTraits>::PixelTrackSoAFromCUDAT(const edm::ParameterSet& iConfig)
0050     : tokenCUDA_(consumes(iConfig.getParameter<edm::InputTag>("src"))), tokenSOA_(produces<TrackSoAHost>()) {}
0051
0052 template <typename TrackerTraits>
0053 void PixelTrackSoAFromCUDAT<TrackerTraits>::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0054   edm::ParameterSetDescription desc;
0055
0056   desc.add<edm::InputTag>("src", edm::InputTag("pixelTracksCUDA"));
0057   descriptions.addWithDefaultLabel(desc);
0058 }
0059
0060 template <typename TrackerTraits>
0061 void PixelTrackSoAFromCUDAT<TrackerTraits>::acquire(edm::Event const& iEvent,
0062                                                     edm::EventSetup const& iSetup,
0063                                                     edm::WaitingTaskWithArenaHolder waitingTaskHolder) {
0064   cms::cuda::Product<TrackSoADevice> const& inputDataWrapped = iEvent.get(tokenCUDA_);
0065   cms::cuda::ScopedContextAcquire ctx{inputDataWrapped, std::move(waitingTaskHolder)};
0066   auto const& tracks_d = ctx.get(inputDataWrapped);  // Tracks on device
0067   tracks_h_ = TrackSoAHost(ctx.stream());            // Create an instance of Tracks on Host, using the stream
0068   cudaCheck(cudaMemcpyAsync(tracks_h_.buffer().get(),
0069                             tracks_d.const_buffer().get(),
0070                             tracks_d.bufferSize(),
0071                             cudaMemcpyDeviceToHost,
0072                             ctx.stream()));  // Copy data from Device to Host
0073 }
0074
0075 template <typename TrackerTraits>
0076 void PixelTrackSoAFromCUDAT<TrackerTraits>::produce(edm::Event& iEvent, edm::EventSetup const& iSetup) {
0077   auto maxTracks = tracks_h_.view().metadata().size();
0078   auto nTracks = tracks_h_.view().nTracks();
0079
0080   assert(nTracks < maxTracks);
0081   if (nTracks == maxTracks - 1) {
0082     edm::LogWarning("PixelTracks") << "Unsorted reconstructed pixel tracks truncated to " << maxTracks - 1
0083                                    << " candidates";
0084   }
0085
0086 #ifdef PIXEL_DEBUG_PRODUCE
0087   std::cout << "size of SoA " << sizeof(tsoa) << " stride " << maxTracks << std::endl;
0088   std::cout << "found " << nTracks << " tracks in cpu SoA at " << &tsoa << std::endl;
0089
0090   int32_t nt = 0;
0091   for (int32_t it = 0; it < maxTracks; ++it) {
0092     auto nHits = TracksUtilities<TrackerTraits>::nHits(tracks_h_.view(), it);
0093     assert(nHits == int(tracks_h_.view().hitIndices().size(it)));
0094     if (nHits == 0)
0095       break;  // this is a guard: maybe we need to move to nTracks...
0096     nt++;
0097   }
0098   assert(nTracks == nt);
0099 #endif
0100
0101   // DO NOT  make a copy  (actually TWO....)
0102   iEvent.emplace(tokenSOA_, std::move(tracks_h_));
0103   assert(!tracks_h_.buffer());
0104 }
0105
0106 using PixelTrackSoAFromCUDAPhase1 = PixelTrackSoAFromCUDAT<pixelTopology::Phase1>;
0107 DEFINE_FWK_MODULE(PixelTrackSoAFromCUDAPhase1);
0108
0109 using PixelTrackSoAFromCUDAPhase2 = PixelTrackSoAFromCUDAT<pixelTopology::Phase2>;
0110 DEFINE_FWK_MODULE(PixelTrackSoAFromCUDAPhase2);
0111
0112 using PixelTrackSoAFromCUDAHIonPhase1 = PixelTrackSoAFromCUDAT<pixelTopology::HIonPhase1>;
0113 DEFINE_FWK_MODULE(PixelTrackSoAFromCUDAHIonPhase1);