Skip to content

Commit 3b264ff

Browse files
committed
ITS: GPU: adapt seed refit
Signed-off-by: Felix Schlepper <felix.schlepper@cern.ch>
1 parent 687291f commit 3b264ff

File tree

7 files changed

+180
-92
lines changed

7 files changed

+180
-92
lines changed

Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TrackingKernels.h

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -209,13 +209,17 @@ void processNeighboursHandler(const int startLayer,
209209
template <int nLayers = 7>
210210
void trackSeedHandler(CellSeed<nLayers>* trackSeeds,
211211
const TrackingFrameInfo** foundTrackingFrameInfo,
212+
const Cluster** unsortedClusters,
212213
o2::its::TrackITSExt* tracks,
213-
std::vector<float>& minPtsHost,
214+
const std::vector<float>& layerRadiiHost,
215+
const std::vector<float>& minPtsHost,
214216
const unsigned int nSeeds,
215217
const float Bz,
216218
const int startLevel,
217-
float maxChi2ClusterAttachment,
218-
float maxChi2NDF,
219+
const float maxChi2ClusterAttachment,
220+
const float maxChi2NDF,
221+
const int reseedIfShorter,
222+
const bool shiftRefToCluster,
219223
const o2::base::Propagator* propagator,
220224
const o2::base::PropagatorF::MatCorrType matCorrType,
221225
const int nBlocks,

Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackerTraitsGPU.cxx

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -325,17 +325,21 @@ void TrackerTraitsGPU<nLayers>::findRoads(const int iteration)
325325
mTimeFrameGPU->createTrackITSExtDevice(trackSeeds);
326326
mTimeFrameGPU->loadTrackSeedsDevice(trackSeeds);
327327

328-
trackSeedHandler(mTimeFrameGPU->getDeviceTrackSeeds(), // CellSeed* trackSeeds
329-
mTimeFrameGPU->getDeviceArrayTrackingFrameInfo(), // TrackingFrameInfo** foundTrackingFrameInfo
330-
mTimeFrameGPU->getDeviceTrackITSExt(), // o2::its::TrackITSExt* tracks
331-
this->mTrkParams[iteration].MinPt, // std::vector<float>& minPtsHost,
328+
trackSeedHandler(mTimeFrameGPU->getDeviceTrackSeeds(), // CellSeed*
329+
mTimeFrameGPU->getDeviceArrayTrackingFrameInfo(), // TrackingFrameInfo**
330+
mTimeFrameGPU->getDeviceArrayUnsortedClusters(), // Cluster**
331+
mTimeFrameGPU->getDeviceTrackITSExt(), // o2::its::TrackITSExt*
332+
this->mTrkParams[iteration].LayerRadii, // const std::vector<float>&
333+
this->mTrkParams[iteration].MinPt, // const std::vector<float>&
332334
trackSeeds.size(), // const size_t nSeeds
333335
this->mBz, // const float Bz
334336
startLevel, // const int startLevel,
335337
this->mTrkParams[0].MaxChi2ClusterAttachment, // float maxChi2ClusterAttachment
336338
this->mTrkParams[0].MaxChi2NDF, // float maxChi2NDF
337-
mTimeFrameGPU->getDevicePropagator(), // const o2::base::Propagator* propagator
338-
this->mTrkParams[0].CorrType, // o2::base::PropagatorImpl<float>::MatCorrType
339+
this->mTrkParams[0].ReseedIfShorter,
340+
this->mTrkParams[0].ShiftRefToCluster,
341+
mTimeFrameGPU->getDevicePropagator(), // const o2::base::Propagator* propagator
342+
this->mTrkParams[0].CorrType, // o2::base::PropagatorImpl<float>::MatCorrType
339343
conf.nBlocksTracksSeeds[iteration],
340344
conf.nThreadsTracksSeeds[iteration]);
341345

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

Lines changed: 141 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -57,42 +57,67 @@ GPUdii() bool fitTrack(TrackITSExt& track,
5757
float bz,
5858
const TrackingFrameInfo** tfInfos,
5959
const o2::base::Propagator* prop,
60-
o2::base::PropagatorF::MatCorrType matCorrType)
60+
o2::base::PropagatorF::MatCorrType matCorrType,
61+
o2::track::TrackPar* linRef,
62+
const bool shiftRefToCluster)
6163
{
6264
for (int iLayer{start}; iLayer != end; iLayer += step) {
6365
if (track.getClusterIndex(iLayer) == constants::UnusedIndex) {
6466
continue;
6567
}
6668
const TrackingFrameInfo& trackingHit = tfInfos[iLayer][track.getClusterIndex(iLayer)];
67-
if (!track.o2::track::TrackParCovF::rotate(trackingHit.alphaTrackingFrame)) {
68-
return false;
69-
}
70-
71-
if (!prop->propagateToX(track,
72-
trackingHit.xTrackingFrame,
73-
bz,
74-
o2::base::PropagatorImpl<float>::MAX_SIN_PHI,
75-
o2::base::PropagatorImpl<float>::MAX_STEP,
76-
matCorrType)) {
77-
return false;
78-
}
69+
if (linRef) {
70+
if (!track.o2::track::TrackParCovF::rotate(trackingHit.alphaTrackingFrame, *linRef, bz)) {
71+
return false;
72+
}
73+
if (!prop->propagateToX(track,
74+
*linRef,
75+
trackingHit.xTrackingFrame,
76+
bz,
77+
o2::base::PropagatorImpl<float>::MAX_SIN_PHI,
78+
o2::base::PropagatorImpl<float>::MAX_STEP,
79+
matCorrType)) {
7980

80-
if (matCorrType == o2::base::PropagatorF::MatCorrType::USEMatCorrNONE) {
81-
const float xx0 = (iLayer > 2) ? 1.e-2f : 5.e-3f; // Rough layer thickness
82-
if (!track.correctForMaterial(xx0, xx0 * constants::Radl * constants::Rho, true)) {
8381
return false;
8482
}
83+
if (matCorrType == o2::base::PropagatorF::MatCorrType::USEMatCorrNONE) {
84+
const float xx0 = (iLayer > 2) ? 1.e-2f : 5.e-3f; // Rough layer thickness
85+
if (!track.correctForMaterial(*linRef, xx0, xx0 * constants::Radl * constants::Rho, true)) {
86+
return false;
87+
}
88+
}
89+
} else {
90+
if (!track.o2::track::TrackParCovF::rotate(trackingHit.alphaTrackingFrame)) {
91+
return false;
92+
}
93+
if (!prop->propagateToX(track,
94+
trackingHit.xTrackingFrame,
95+
bz,
96+
o2::base::PropagatorImpl<float>::MAX_SIN_PHI,
97+
o2::base::PropagatorImpl<float>::MAX_STEP,
98+
matCorrType)) {
99+
return false;
100+
}
101+
if (matCorrType == o2::base::PropagatorF::MatCorrType::USEMatCorrNONE) {
102+
const float xx0 = (iLayer > 2) ? 1.e-2f : 5.e-3f; // Rough layer thickness
103+
if (!track.correctForMaterial(xx0, xx0 * constants::Radl * constants::Rho, true)) {
104+
return false;
105+
}
106+
}
85107
}
86108

87109
auto predChi2{track.getPredictedChi2(trackingHit.positionTrackingFrame, trackingHit.covarianceTrackingFrame)};
88-
89110
if ((nCl >= 3 && predChi2 > chi2clcut) || predChi2 < 0.f) {
90111
return false;
91112
}
92113
track.setChi2(track.getChi2() + predChi2);
93114
if (!track.o2::track::TrackParCov::update(trackingHit.positionTrackingFrame, trackingHit.covarianceTrackingFrame)) {
94115
return false;
95116
}
117+
if (linRef && shiftRefToCluster) { // displace the reference to the last updated cluster
118+
linRef->setY(trackingHit.positionTrackingFrame[0]);
119+
linRef->setZ(trackingHit.positionTrackingFrame[1]);
120+
}
96121
nCl++;
97122
}
98123
return o2::gpu::CAMath::Abs(track.getQ2Pt()) < maxQoverPt && track.getChi2() < chi2ndfcut * (nCl * 2 - 5);
@@ -101,7 +126,8 @@ GPUdii() bool fitTrack(TrackITSExt& track,
101126
GPUdii() o2::track::TrackParCov buildTrackSeed(const Cluster& cluster1,
102127
const Cluster& cluster2,
103128
const TrackingFrameInfo& tf3,
104-
const float bz)
129+
const float bz,
130+
const bool reverse = false)
105131
{
106132
const float ca = o2::gpu::CAMath::Cos(tf3.alphaTrackingFrame), sa = o2::gpu::CAMath::Sin(tf3.alphaTrackingFrame);
107133
const float x1 = cluster1.xCoordinate * ca + cluster1.yCoordinate * sa;
@@ -115,12 +141,13 @@ GPUdii() o2::track::TrackParCov buildTrackSeed(const Cluster& cluster1,
115141
const float z3 = tf3.positionTrackingFrame[1];
116142

117143
const bool zeroField{o2::gpu::CAMath::Abs(bz) < o2::constants::math::Almost0};
118-
const float tgp = zeroField ? o2::gpu::CAMath::ATan2(y3 - y1, x3 - x1) : 1.f;
119-
const float crv = zeroField ? 1.f : math_utils::computeCurvature(x3, y3, x2, y2, x1, y1);
120-
const float snp = zeroField ? tgp / o2::gpu::CAMath::Sqrt(1.f + tgp * tgp) : crv * (x3 - math_utils::computeCurvatureCentreX(x3, y3, x2, y2, x1, y1));
144+
const float sign = (reverse) ? -1.f : 1.f;
145+
const float tgp = zeroField ? sign * o2::gpu::CAMath::ATan2(y3 - y1, x3 - x1) : 1.f;
146+
const float crv = sign * (zeroField ? 1.f : math_utils::computeCurvature(x3, y3, x2, y2, x1, y1));
147+
const float snp = (zeroField ? tgp / o2::gpu::CAMath::Sqrt(1.f + tgp * tgp) : crv * (x3 - math_utils::computeCurvatureCentreX(x3, y3, x2, y2, x1, y1)));
121148
const float tgl12 = math_utils::computeTanDipAngle(x1, y1, x2, y2, z1, z2);
122149
const float tgl23 = math_utils::computeTanDipAngle(x2, y2, x3, y3, z2, z3);
123-
const float q2pt = zeroField ? 1.f / o2::track::kMostProbablePt : crv / (bz * o2::constants::math::B2C);
150+
const float q2pt = zeroField ? sign / o2::track::kMostProbablePt : crv / (bz * o2::constants::math::B2C);
124151
const float q2pt2 = crv * crv;
125152
const float sg2q2pt = o2::track::kC1Pt2max * (q2pt2 > 0.0005 ? (q2pt2 < 1 ? q2pt2 : 1) : 0.0005);
126153
return track::TrackParCov(tf3.xTrackingFrame, tf3.alphaTrackingFrame,
@@ -132,6 +159,52 @@ GPUdii() o2::track::TrackParCov buildTrackSeed(const Cluster& cluster1,
132159
0.f, 0.f, 0.f, 0.f, sg2q2pt});
133160
}
134161

162+
template <int nLayers>
163+
GPUdii() TrackITSExt seedTrackForRefit(const CellSeed<nLayers>& seed,
164+
const TrackingFrameInfo** foundTrackingFrameInfo,
165+
const Cluster** unsortedClusters,
166+
const float* layerRadii,
167+
const float bz,
168+
const int reseedIfShorter)
169+
{
170+
TrackITSExt temporaryTrack(seed);
171+
int lrMin = nLayers, lrMax = 0, lrMid = 0;
172+
for (int iL{0}; iL < nLayers; ++iL) {
173+
const int idx = seed.getCluster(iL);
174+
temporaryTrack.setExternalClusterIndex(iL, idx, idx != constants::UnusedIndex);
175+
if (idx != constants::UnusedIndex) {
176+
// TODO only works if does not have holes
177+
lrMin = o2::gpu::CAMath::Min(lrMin, iL);
178+
lrMax = o2::gpu::CAMath::Max(lrMax, iL);
179+
}
180+
}
181+
const int ncl = temporaryTrack.getNClusters();
182+
if (ncl < reseedIfShorter && ncl > 0) { // need to check if there are any clusters since we keep invalidate seeeds around
183+
if (ncl == nLayers) {
184+
lrMin = 0;
185+
lrMax = nLayers - 1;
186+
lrMid = (lrMin + lrMax) / 2;
187+
} else {
188+
lrMid = lrMin + 1;
189+
float midR = 0.5f * (layerRadii[lrMax] + layerRadii[lrMin]), dstMidR = o2::gpu::CAMath::Abs(midR - layerRadii[lrMid]);
190+
for (int iL = lrMid + 1; iL < lrMax; ++iL) { // find the midpoint as closest to the midR
191+
auto dst = o2::gpu::GPUCommonMath::Abs(midR - layerRadii[iL]);
192+
if (dst < dstMidR) {
193+
lrMid = iL;
194+
dstMidR = dst;
195+
}
196+
}
197+
}
198+
const auto& cluster0_tf = foundTrackingFrameInfo[lrMin][seed.getCluster(lrMin)];
199+
const auto& cluster1_gl = unsortedClusters[lrMid][seed.getCluster(lrMid)];
200+
const auto& cluster2_gl = unsortedClusters[lrMax][seed.getCluster(lrMax)];
201+
temporaryTrack.getParamIn() = buildTrackSeed(cluster2_gl, cluster1_gl, cluster0_tf, bz, true);
202+
}
203+
temporaryTrack.resetCovariance();
204+
temporaryTrack.setCov(temporaryTrack.getQ2Pt() * temporaryTrack.getQ2Pt() * temporaryTrack.getCov()[o2::track::CovLabels::kSigQ2Pt2], o2::track::CovLabels::kSigQ2Pt2);
205+
return temporaryTrack;
206+
}
207+
135208
struct sort_tracklets {
136209
GPUhd() bool operator()(const Tracklet& a, const Tracklet& b)
137210
{
@@ -206,27 +279,23 @@ template <int nLayers>
206279
GPUg() void __launch_bounds__(256, 1) fitTrackSeedsKernel(
207280
CellSeed<nLayers>* trackSeeds,
208281
const TrackingFrameInfo** foundTrackingFrameInfo,
282+
const Cluster** unsortedClusters,
209283
o2::its::TrackITSExt* tracks,
284+
const float* layerRadii,
210285
const float* minPts,
211286
const unsigned int nSeeds,
212287
const float bz,
213288
const int startLevel,
214-
float maxChi2ClusterAttachment,
215-
float maxChi2NDF,
289+
const float maxChi2ClusterAttachment,
290+
const float maxChi2NDF,
291+
const int reseedIfShorter,
292+
const bool shifRefToCluster,
216293
const o2::base::Propagator* propagator,
217294
const o2::base::PropagatorF::MatCorrType matCorrType)
218295
{
219296
for (int iCurrentTrackSeedIndex = blockIdx.x * blockDim.x + threadIdx.x; iCurrentTrackSeedIndex < nSeeds; iCurrentTrackSeedIndex += blockDim.x * gridDim.x) {
220-
auto& seed = trackSeeds[iCurrentTrackSeedIndex];
221-
222-
TrackITSExt temporaryTrack{seed};
223-
224-
temporaryTrack.resetCovariance();
225-
temporaryTrack.setChi2(0);
226-
auto& clusters = seed.getClusters();
227-
for (int iL{0}; iL < nLayers; ++iL) {
228-
temporaryTrack.setExternalClusterIndex(iL, clusters[iL], clusters[iL] != constants::UnusedIndex);
229-
}
297+
TrackITSExt temporaryTrack = seedTrackForRefit<nLayers>(trackSeeds[iCurrentTrackSeedIndex], foundTrackingFrameInfo, unsortedClusters, layerRadii, bz, reseedIfShorter);
298+
o2::track::TrackPar linRef{temporaryTrack};
230299
bool fitSuccess = fitTrack(temporaryTrack, // TrackITSExt& track,
231300
0, // int lastLayer,
232301
nLayers, // int firstLayer,
@@ -238,14 +307,17 @@ GPUg() void __launch_bounds__(256, 1) fitTrackSeedsKernel(
238307
bz, // float bz,
239308
foundTrackingFrameInfo, // TrackingFrameInfo** trackingFrameInfo,
240309
propagator, // const o2::base::Propagator* propagator,
241-
matCorrType); // o2::base::PropagatorF::MatCorrType matCorrType
310+
matCorrType, // o2::base::PropagatorF::MatCorrType matCorrType
311+
&linRef,
312+
shifRefToCluster);
242313
if (!fitSuccess) {
243314
continue;
244315
}
245316
temporaryTrack.getParamOut() = temporaryTrack.getParamIn();
317+
linRef = temporaryTrack.getParamOut(); // use refitted track as lin.reference
246318
temporaryTrack.resetCovariance();
319+
temporaryTrack.setCov(temporaryTrack.getQ2Pt() * temporaryTrack.getQ2Pt() * temporaryTrack.getCov()[o2::track::CovLabels::kSigQ2Pt2], o2::track::CovLabels::kSigQ2Pt2);
247320
temporaryTrack.setChi2(0);
248-
249321
fitSuccess = fitTrack(temporaryTrack, // TrackITSExt& track,
250322
nLayers - 1, // int lastLayer,
251323
-1, // int firstLayer,
@@ -257,7 +329,9 @@ GPUg() void __launch_bounds__(256, 1) fitTrackSeedsKernel(
257329
bz, // float bz,
258330
foundTrackingFrameInfo, // TrackingFrameInfo** trackingFrameInfo,
259331
propagator, // const o2::base::Propagator* propagator,
260-
matCorrType); // o2::base::PropagatorF::MatCorrType matCorrType
332+
matCorrType, // o2::base::PropagatorF::MatCorrType matCorrType
333+
&linRef,
334+
shifRefToCluster);
261335
if (!fitSuccess || temporaryTrack.getPt() < minPts[nLayers - temporaryTrack.getNClusters()]) {
262336
continue;
263337
}
@@ -1088,34 +1162,42 @@ void processNeighboursHandler(const int startLayer,
10881162
template <int nLayers>
10891163
void trackSeedHandler(CellSeed<nLayers>* trackSeeds,
10901164
const TrackingFrameInfo** foundTrackingFrameInfo,
1165+
const Cluster** unsortedClusters,
10911166
o2::its::TrackITSExt* tracks,
1092-
std::vector<float>& minPtsHost,
1167+
const std::vector<float>& layerRadiiHost,
1168+
const std::vector<float>& minPtsHost,
10931169
const unsigned int nSeeds,
10941170
const float bz,
10951171
const int startLevel,
1096-
float maxChi2ClusterAttachment,
1097-
float maxChi2NDF,
1172+
const float maxChi2ClusterAttachment,
1173+
const float maxChi2NDF,
1174+
const int reseedIfShorter,
1175+
const bool shiftRefToCluster,
10981176
const o2::base::Propagator* propagator,
10991177
const o2::base::PropagatorF::MatCorrType matCorrType,
11001178
const int nBlocks,
11011179
const int nThreads)
11021180
{
11031181
thrust::device_vector<float> minPts(minPtsHost);
1182+
thrust::device_vector<float> layerRadii(layerRadiiHost);
11041183
gpu::fitTrackSeedsKernel<<<nBlocks, nThreads>>>(
1105-
trackSeeds, // CellSeed*
1106-
foundTrackingFrameInfo, // TrackingFrameInfo**
1107-
tracks, // TrackITSExt*
1108-
thrust::raw_pointer_cast(&minPts[0]), // const float* minPts,
1109-
nSeeds, // const unsigned int
1110-
bz, // const float
1111-
startLevel, // const int
1112-
maxChi2ClusterAttachment, // float
1113-
maxChi2NDF, // float
1114-
propagator, // const o2::base::Propagator*
1115-
matCorrType); // o2::base::PropagatorF::MatCorrType
1184+
trackSeeds, // CellSeed*
1185+
foundTrackingFrameInfo, // TrackingFrameInfo**
1186+
unsortedClusters, // Cluster**
1187+
tracks, // TrackITSExt*
1188+
thrust::raw_pointer_cast(&layerRadii[0]), // const float*
1189+
thrust::raw_pointer_cast(&minPts[0]), // const float*
1190+
nSeeds, // const unsigned int
1191+
bz, // const float
1192+
startLevel, // const int
1193+
maxChi2ClusterAttachment, // float
1194+
maxChi2NDF, // float
1195+
reseedIfShorter, // int
1196+
shiftRefToCluster, // bool
1197+
propagator, // const o2::base::Propagator*
1198+
matCorrType); // o2::base::PropagatorF::MatCorrType
11161199
thrust::device_ptr<o2::its::TrackITSExt> tr_ptr(tracks);
11171200
thrust::sort(tr_ptr, tr_ptr + nSeeds, gpu::compare_track_chi2());
1118-
GPUChkErrS(cudaStreamSynchronize(gpu::Stream::DefaultStream));
11191201
}
11201202

11211203
/// Explicit instantiation of ITS2 handlers
@@ -1281,13 +1363,17 @@ template void processNeighboursHandler<7>(const int startLayer,
12811363

12821364
template void trackSeedHandler(CellSeed<7>* trackSeeds,
12831365
const TrackingFrameInfo** foundTrackingFrameInfo,
1366+
const Cluster** unsortedClusters,
12841367
o2::its::TrackITSExt* tracks,
1285-
std::vector<float>& minPtsHost,
1368+
const std::vector<float>& layerRadiiHost,
1369+
const std::vector<float>& minPtsHost,
12861370
const unsigned int nSeeds,
12871371
const float bz,
12881372
const int startLevel,
1289-
float maxChi2ClusterAttachment,
1290-
float maxChi2NDF,
1373+
const float maxChi2ClusterAttachment,
1374+
const float maxChi2NDF,
1375+
const int reseedIfShorter,
1376+
const bool shiftRefToCluster,
12911377
const o2::base::Propagator* propagator,
12921378
const o2::base::PropagatorF::MatCorrType matCorrType,
12931379
const int nBlocks,

0 commit comments

Comments
 (0)