Skip to content

Commit 001e143

Browse files
mpucciompuccio
committed
ITS: try to uniform GPU and CPU code for track following
1 parent 7984605 commit 001e143

3 files changed

Lines changed: 198 additions & 222 deletions

File tree

‎Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackingKernels.cu‎

Lines changed: 95 additions & 169 deletions
Original file line numberDiff line numberDiff line change
@@ -111,31 +111,65 @@ struct compare_track_chi2 {
111111
};
112112

113113
template <int NLayers>
114-
GPUdi() void finaliseTrackExtensionTrial(const uint32_t backupPattern,
115-
TrackITSInternal<NLayers>& trial,
116-
const TrackingFrameInfo* const* trackingFrameInfo,
117-
const float* layerxX0,
118-
const int nLayers,
119-
const float bz,
120-
const float maxChi2ClusterAttachment,
121-
const float maxChi2NDF,
122-
const o2::base::Propagator* propagator,
123-
const o2::base::PropagatorF::MatCorrType matCorrType,
124-
const bool shiftRefToCluster,
125-
const bool repeatRefitOut,
126-
TrackITSInternal<NLayers>& best,
127-
uint32_t& bestDiff)
128-
{
129-
const auto diff = (trial.getPattern() & ~backupPattern) & TrackITS::getLayerPatternMask<NLayers>();
130-
if (!diff ||
131-
!o2::its::track::refitTrack(trial, trackingFrameInfo, layerxX0, nLayers, bz, maxChi2ClusterAttachment, maxChi2NDF, propagator, matCorrType, shiftRefToCluster, repeatRefitOut)) {
132-
return;
133-
}
134-
if (o2::its::track::isBetter(trial, best)) {
135-
best = trial;
136-
bestDiff = diff;
114+
struct TrackExtensionDirectionFollowerDevice {
115+
GPUdi() bool operator()(TrackITSInternal<NLayers>& candidate, bool outward) const
116+
{
117+
const auto startHypothesis = TrackExtensionHypothesis<NLayers>{candidate, outward};
118+
TrackExtensionHypothesis<NLayers> bestHypothesis;
119+
if (!followTrackExtensionDirection<NLayers>(startHypothesis,
120+
*utils,
121+
rofMask,
122+
rofOverlaps,
123+
clusters,
124+
usedClusters,
125+
clustersIndexTables,
126+
ROFClusters,
127+
trackingFrameInfo,
128+
layerRadii,
129+
layerxX0,
130+
nLayers,
131+
phiBins,
132+
maxHypotheses,
133+
bz,
134+
maxChi2ClusterAttachment,
135+
maxChi2NDF,
136+
nSigmaCutPhi,
137+
nSigmaCutZ,
138+
outward,
139+
propagator,
140+
matCorrType,
141+
activeHypotheses,
142+
nextHypotheses,
143+
bestHypothesis)) {
144+
return false;
145+
}
146+
updateTrackFromExtensionHypothesis(bestHypothesis, outward, nLayers, candidate);
147+
return true;
137148
}
138-
}
149+
150+
const IndexTableUtils<NLayers>* utils{nullptr};
151+
typename ROFMaskTable<NLayers>::View rofMask;
152+
typename ROFOverlapTable<NLayers>::View rofOverlaps;
153+
const Cluster* const* clusters{nullptr};
154+
const unsigned char* const* usedClusters{nullptr};
155+
const int* const* clustersIndexTables{nullptr};
156+
const int* const* ROFClusters{nullptr};
157+
const TrackingFrameInfo* const* trackingFrameInfo{nullptr};
158+
const float* layerRadii{nullptr};
159+
const float* layerxX0{nullptr};
160+
int nLayers{0};
161+
int phiBins{0};
162+
int maxHypotheses{0};
163+
float bz{0.f};
164+
float maxChi2ClusterAttachment{0.f};
165+
float maxChi2NDF{0.f};
166+
float nSigmaCutPhi{0.f};
167+
float nSigmaCutZ{0.f};
168+
const o2::base::Propagator* propagator{nullptr};
169+
o2::base::PropagatorF::MatCorrType matCorrType{o2::base::PropagatorF::MatCorrType::USEMatCorrNONE};
170+
TrackExtensionHypothesis<NLayers>* activeHypotheses{nullptr};
171+
TrackExtensionHypothesis<NLayers>* nextHypotheses{nullptr};
172+
};
139173

140174
template <int NLayers>
141175
GPUg() void __launch_bounds__(constants::GPUThreads, 1) countTrackSeedsKernel(
@@ -240,153 +274,45 @@ GPUg() void __launch_bounds__(constants::GPUThreads, 1) fitTrackSeedsKernel(
240274
const int threadIndex = blockIdx.x * blockDim.x + threadIdx.x;
241275
auto* activeHypotheses = activeHypothesesScratch + threadIndex * maxHypotheses;
242276
auto* nextHypotheses = nextHypothesesScratch + threadIndex * maxHypotheses;
243-
const auto backupPattern = temporaryTrack.getPattern();
277+
const auto backup = temporaryTrack;
244278
auto best = temporaryTrack;
245279
uint32_t bestDiff{0};
246-
TrackITSInternal<NLayers> topResult;
247-
TrackITSInternal<NLayers> botResult;
248-
bool hasTopResult{false};
249-
bool hasBotResult{false};
250-
const uint32_t lastLayer = static_cast<uint32_t>(nLayers - 1);
251-
252-
if (extendTop && temporaryTrack.getLastClusterLayer() != lastLayer) {
253-
auto candidate = temporaryTrack;
254-
const auto startHypothesis = TrackExtensionHypothesis<NLayers>{temporaryTrack, true};
255-
TrackExtensionHypothesis<NLayers> bestHypothesis;
256-
if (followTrackExtensionDirection<NLayers>(startHypothesis,
257-
*utils,
258-
rofMask,
259-
rofOverlaps,
260-
clusters,
261-
usedClusters,
262-
clustersIndexTables,
263-
ROFClusters,
264-
foundTrackingFrameInfo,
265-
layerRadii,
266-
layerxX0,
267-
nLayers,
268-
phiBins,
269-
maxHypotheses,
270-
bz,
271-
maxChi2ClusterAttachment,
272-
maxChi2NDF,
273-
nSigmaCutPhi,
274-
nSigmaCutZ,
275-
true,
276-
propagator,
277-
matCorrType,
278-
activeHypotheses,
279-
nextHypotheses,
280-
bestHypothesis)) {
281-
updateTrackFromExtensionHypothesis(bestHypothesis, true, nLayers, candidate);
282-
topResult = candidate;
283-
hasTopResult = true;
284-
finaliseTrackExtensionTrial<NLayers>(backupPattern, candidate, foundTrackingFrameInfo, layerxX0, nLayers, bz, maxChi2ClusterAttachment, maxChi2NDF, propagator, matCorrType, shiftRefToCluster, repeatRefitOut, best, bestDiff);
285-
}
286-
}
287-
if (extendBot && temporaryTrack.getFirstClusterLayer() != 0) {
288-
auto candidate = temporaryTrack;
289-
const auto startHypothesis = TrackExtensionHypothesis<NLayers>{temporaryTrack, false};
290-
TrackExtensionHypothesis<NLayers> bestHypothesis;
291-
if (followTrackExtensionDirection<NLayers>(startHypothesis,
292-
*utils,
293-
rofMask,
294-
rofOverlaps,
295-
clusters,
296-
usedClusters,
297-
clustersIndexTables,
298-
ROFClusters,
299-
foundTrackingFrameInfo,
300-
layerRadii,
301-
layerxX0,
302-
nLayers,
303-
phiBins,
304-
maxHypotheses,
305-
bz,
306-
maxChi2ClusterAttachment,
307-
maxChi2NDF,
308-
nSigmaCutPhi,
309-
nSigmaCutZ,
310-
false,
311-
propagator,
312-
matCorrType,
313-
activeHypotheses,
314-
nextHypotheses,
315-
bestHypothesis)) {
316-
updateTrackFromExtensionHypothesis(bestHypothesis, false, nLayers, candidate);
317-
botResult = candidate;
318-
hasBotResult = true;
319-
finaliseTrackExtensionTrial<NLayers>(backupPattern, candidate, foundTrackingFrameInfo, layerxX0, nLayers, bz, maxChi2ClusterAttachment, maxChi2NDF, propagator, matCorrType, shiftRefToCluster, repeatRefitOut, best, bestDiff);
320-
}
321-
}
322-
if (extendTop && extendBot) {
323-
if (hasTopResult && topResult.getFirstClusterLayer() != 0) {
324-
auto candidate = topResult;
325-
const auto startHypothesis = TrackExtensionHypothesis<NLayers>{topResult, false};
326-
TrackExtensionHypothesis<NLayers> bestHypothesis;
327-
if (followTrackExtensionDirection<NLayers>(startHypothesis,
328-
*utils,
329-
rofMask,
330-
rofOverlaps,
331-
clusters,
332-
usedClusters,
333-
clustersIndexTables,
334-
ROFClusters,
335-
foundTrackingFrameInfo,
336-
layerRadii,
337-
layerxX0,
338-
nLayers,
339-
phiBins,
340-
maxHypotheses,
341-
bz,
342-
maxChi2ClusterAttachment,
343-
maxChi2NDF,
344-
nSigmaCutPhi,
345-
nSigmaCutZ,
346-
false,
347-
propagator,
348-
matCorrType,
349-
activeHypotheses,
350-
nextHypotheses,
351-
bestHypothesis)) {
352-
updateTrackFromExtensionHypothesis(bestHypothesis, false, nLayers, candidate);
353-
finaliseTrackExtensionTrial<NLayers>(backupPattern, candidate, foundTrackingFrameInfo, layerxX0, nLayers, bz, maxChi2ClusterAttachment, maxChi2NDF, propagator, matCorrType, shiftRefToCluster, repeatRefitOut, best, bestDiff);
354-
}
355-
}
356-
if (hasBotResult && botResult.getLastClusterLayer() != lastLayer) {
357-
auto candidate = botResult;
358-
const auto startHypothesis = TrackExtensionHypothesis<NLayers>{botResult, true};
359-
TrackExtensionHypothesis<NLayers> bestHypothesis;
360-
if (followTrackExtensionDirection<NLayers>(startHypothesis,
361-
*utils,
362-
rofMask,
363-
rofOverlaps,
364-
clusters,
365-
usedClusters,
366-
clustersIndexTables,
367-
ROFClusters,
368-
foundTrackingFrameInfo,
369-
layerRadii,
370-
layerxX0,
371-
nLayers,
372-
phiBins,
373-
maxHypotheses,
374-
bz,
375-
maxChi2ClusterAttachment,
376-
maxChi2NDF,
377-
nSigmaCutPhi,
378-
nSigmaCutZ,
379-
true,
380-
propagator,
381-
matCorrType,
382-
activeHypotheses,
383-
nextHypotheses,
384-
bestHypothesis)) {
385-
updateTrackFromExtensionHypothesis(bestHypothesis, true, nLayers, candidate);
386-
finaliseTrackExtensionTrial<NLayers>(backupPattern, candidate, foundTrackingFrameInfo, layerxX0, nLayers, bz, maxChi2ClusterAttachment, maxChi2NDF, propagator, matCorrType, shiftRefToCluster, repeatRefitOut, best, bestDiff);
387-
}
388-
}
389-
}
280+
TrackExtensionDirectionFollowerDevice<NLayers> followDirection{
281+
utils,
282+
rofMask,
283+
rofOverlaps,
284+
clusters,
285+
usedClusters,
286+
clustersIndexTables,
287+
ROFClusters,
288+
foundTrackingFrameInfo,
289+
layerRadii,
290+
layerxX0,
291+
nLayers,
292+
phiBins,
293+
maxHypotheses,
294+
bz,
295+
maxChi2ClusterAttachment,
296+
maxChi2NDF,
297+
nSigmaCutPhi,
298+
nSigmaCutZ,
299+
propagator,
300+
matCorrType,
301+
activeHypotheses,
302+
nextHypotheses};
303+
TrackExtensionBestTrial<NLayers> bestTrial{
304+
backup.getPattern(),
305+
foundTrackingFrameInfo,
306+
layerxX0,
307+
nLayers,
308+
bz,
309+
maxChi2ClusterAttachment,
310+
maxChi2NDF,
311+
propagator,
312+
matCorrType,
313+
shiftRefToCluster,
314+
repeatRefitOut};
315+
followTrackExtensionBranches(backup, extendTop, extendBot, nLayers, followDirection, bestTrial, best, bestDiff);
390316
temporaryTrack = best;
391317
tracks[seedLUT[iCurrentTrackSeedIndex]] = makeTrackITSExt(temporaryTrack);
392318
if (bestDiff) {

‎Detectors/ITSMFT/ITS/tracking/include/ITStracking/TrackFollower.h‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
#ifndef TRACKINGITSU_INCLUDE_TRACKFOLLOWER_H_
1616
#define TRACKINGITSU_INCLUDE_TRACKFOLLOWER_H_
1717

18+
#include <cstdint>
19+
1820
#include "GPUCommonDef.h"
1921
#include "GPUCommonMath.h"
2022
#include "DetectorsBase/Propagator.h"
@@ -72,6 +74,92 @@ GPUhdi() void updateTrackFromExtensionHypothesis(const TrackExtensionHypothesis<
7274
}
7375
}
7476

77+
template <int NLayers>
78+
struct TrackExtensionBestTrial {
79+
GPUdi() void update(TrackITSInternal<NLayers>& trial, TrackITSInternal<NLayers>& best, uint32_t& bestDiff) const
80+
{
81+
const auto diff = (trial.getPattern() & ~backupPattern) & TrackITS::getLayerPatternMask<NLayers>();
82+
if (!diff ||
83+
!track::refitTrack(trial,
84+
trackingFrameInfo,
85+
layerxX0,
86+
nLayers,
87+
bz,
88+
maxChi2ClusterAttachment,
89+
maxChi2NDF,
90+
propagator,
91+
matCorrType,
92+
shiftRefToCluster,
93+
repeatRefitOut)) {
94+
return;
95+
}
96+
if (track::isBetter(trial, best)) {
97+
best = trial;
98+
bestDiff = diff;
99+
}
100+
}
101+
102+
uint32_t backupPattern{0};
103+
const TrackingFrameInfo* const* trackingFrameInfo{nullptr};
104+
const float* layerxX0{nullptr};
105+
int nLayers{0};
106+
float bz{0.f};
107+
float maxChi2ClusterAttachment{0.f};
108+
float maxChi2NDF{0.f};
109+
const o2::base::Propagator* propagator{nullptr};
110+
o2::base::PropagatorF::MatCorrType matCorrType{o2::base::PropagatorF::MatCorrType::USEMatCorrNONE};
111+
bool shiftRefToCluster{false};
112+
bool repeatRefitOut{false};
113+
};
114+
115+
template <int NLayers, typename FollowDirection, typename BestTrial>
116+
GPUdi() void followTrackExtensionBranches(const TrackITSInternal<NLayers>& backup,
117+
const bool extendTop,
118+
const bool extendBot,
119+
const int nLayers,
120+
FollowDirection& followDirection,
121+
BestTrial& bestTrial,
122+
TrackITSInternal<NLayers>& best,
123+
uint32_t& bestDiff)
124+
{
125+
const uint32_t lastLayer = static_cast<uint32_t>(nLayers - 1);
126+
TrackITSInternal<NLayers> topResult;
127+
TrackITSInternal<NLayers> botResult;
128+
bool hasTopResult{false};
129+
bool hasBotResult{false};
130+
131+
if (extendTop && backup.getLastClusterLayer() != lastLayer) {
132+
auto candidate = backup;
133+
if (followDirection(candidate, true)) {
134+
topResult = candidate;
135+
hasTopResult = true;
136+
bestTrial.update(candidate, best, bestDiff);
137+
}
138+
}
139+
if (extendBot && backup.getFirstClusterLayer() != 0) {
140+
auto candidate = backup;
141+
if (followDirection(candidate, false)) {
142+
botResult = candidate;
143+
hasBotResult = true;
144+
bestTrial.update(candidate, best, bestDiff);
145+
}
146+
}
147+
if (extendTop && extendBot) {
148+
if (hasTopResult && topResult.getFirstClusterLayer() != 0) {
149+
auto candidate = topResult;
150+
if (followDirection(candidate, false)) {
151+
bestTrial.update(candidate, best, bestDiff);
152+
}
153+
}
154+
if (hasBotResult && botResult.getLastClusterLayer() != lastLayer) {
155+
auto candidate = botResult;
156+
if (followDirection(candidate, true)) {
157+
bestTrial.update(candidate, best, bestDiff);
158+
}
159+
}
160+
}
161+
}
162+
75163
template <int NLayers>
76164
GPUhdi() bool followTrackExtensionDirection(const TrackExtensionHypothesis<NLayers>& startHypothesis,
77165
const IndexTableUtils<NLayers>& utils,

0 commit comments

Comments
 (0)