-
Notifications
You must be signed in to change notification settings - Fork 35
/
Copy pathPixelTrackSoAFromCUDA.cc
65 lines (52 loc) · 2.36 KB
/
PixelTrackSoAFromCUDA.cc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
#include <cuda_runtime.h>
#include "CUDACore/Product.h"
#include "CUDACore/HostProduct.h"
#include "CUDADataFormats/PixelTrackHeterogeneous.h"
#include "Framework/EventSetup.h"
#include "Framework/Event.h"
#include "Framework/PluginFactory.h"
#include "Framework/EDProducer.h"
#include "CUDACore/ScopedContext.h"
class PixelTrackSoAFromCUDA : public edm::EDProducerExternalWork {
public:
explicit PixelTrackSoAFromCUDA(edm::ProductRegistry& reg);
~PixelTrackSoAFromCUDA() override = default;
private:
void acquire(edm::Event const& iEvent,
edm::EventSetup const& iSetup,
edm::WaitingTaskWithArenaHolder waitingTaskHolder) override;
void produce(edm::Event& iEvent, edm::EventSetup const& iSetup) override;
edm::EDGetTokenT<cms::cuda::Product<PixelTrackHeterogeneous>> tokenCUDA_;
edm::EDPutTokenT<PixelTrackHeterogeneous> tokenSOA_;
cms::cuda::host::unique_ptr<pixelTrack::TrackSoA> m_soa;
};
PixelTrackSoAFromCUDA::PixelTrackSoAFromCUDA(edm::ProductRegistry& reg)
: tokenCUDA_(reg.consumes<cms::cuda::Product<PixelTrackHeterogeneous>>()),
tokenSOA_(reg.produces<PixelTrackHeterogeneous>()) {}
void PixelTrackSoAFromCUDA::acquire(edm::Event const& iEvent,
edm::EventSetup const& iSetup,
edm::WaitingTaskWithArenaHolder waitingTaskHolder) {
cms::cuda::Product<PixelTrackHeterogeneous> const& inputDataWrapped = iEvent.get(tokenCUDA_);
cms::cuda::ScopedContextAcquire ctx{inputDataWrapped, std::move(waitingTaskHolder)};
auto const& inputData = ctx.get(inputDataWrapped);
m_soa = inputData.toHostAsync(ctx.stream());
}
void PixelTrackSoAFromCUDA::produce(edm::Event& iEvent, edm::EventSetup const& iSetup) {
/*
auto const & tsoa = *m_soa;
auto maxTracks = tsoa.stride();
std::cout << "size of SoA" << sizeof(tsoa) << " stride " << maxTracks << std::endl;
int32_t nt = 0;
for (int32_t it = 0; it < maxTracks; ++it) {
auto nHits = tsoa.nHits(it);
assert(nHits==int(tsoa.hitIndices.size(it)));
if (nHits == 0) break; // this is a guard: maybe we need to move to nTracks...
nt++;
}
std::cout << "found " << nt << " tracks in cpu SoA at " << &tsoa << std::endl;
*/
// DO NOT make a copy (actually TWO....)
iEvent.emplace(tokenSOA_, PixelTrackHeterogeneous(std::move(m_soa)));
assert(!m_soa);
}
DEFINE_FWK_MODULE(PixelTrackSoAFromCUDA);