Skip to content

Commit 94d9198

Browse files
mpucciompuccio
authored
ITS: add back the code for track following (#15456)
* ITS: re-enable the possibility of extending tracks * ITS: integrate track extension in road finding * ITS: implement code review * ITS: rename away from beam for the track follower * ITS: new round of comments * ITS: try to uniform GPU and CPU code for track following * ITS: introduce configuration context to reduce number of arguments
1 parent 6bc6de6 commit 94d9198

22 files changed

Lines changed: 1167 additions & 266 deletions

‎DataFormats/Detectors/ITSMFT/ITS/include/DataFormatsITS/TrackITS.h‎

Lines changed: 39 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,11 @@ namespace its
3535

3636
class TrackITS : public o2::track::TrackParCov
3737
{
38+
public:
39+
static constexpr unsigned int ExtendedPatternShift = 24;
40+
static constexpr int MaxLayersInTrackPattern = 8;
41+
42+
private:
3843
enum UserBits {
3944
kSharedClusters = 1 << 28
4045
};
@@ -106,16 +111,47 @@ class TrackITS : public o2::track::TrackParCov
106111
GPUhdi() uint32_t getPattern() const { return mPattern; }
107112
bool hasHitOnLayer(uint32_t i) const { return mPattern & (0x1 << i); }
108113
bool isFakeOnLayer(uint32_t i) const { return !(mPattern & (0x1 << (16 + i))); }
109-
bool isExtendedOnLayer(uint32_t i) const { return (mPattern & (0x1 << (24 + i))); } // only correct if getNClusters <= 8 on layers <= 8
110-
uint32_t getLastClusterLayer() const
114+
bool isExtendedOnLayer(uint32_t i) const { return (mPattern & (0x1 << (ExtendedPatternShift + i))); } // only correct if getNClusters <= 8 on layers <= 8
115+
template <int NLayers>
116+
GPUhdi() static constexpr uint32_t getLayerPatternMask()
117+
{
118+
return (NLayers >= 32) ? 0xffffffffu : ((1u << NLayers) - 1u);
119+
}
120+
template <int NLayers>
121+
GPUhdi() void setExtendedLayerPattern(uint32_t pattern)
122+
{
123+
pattern &= getLayerPatternMask<NLayers>();
124+
setUserField(static_cast<uint16_t>(pattern));
125+
if constexpr (NLayers <= MaxLayersInTrackPattern) {
126+
setPattern(getPattern() | (pattern << ExtendedPatternShift));
127+
}
128+
}
129+
template <int NLayers>
130+
GPUhdi() uint32_t getExtendedLayerPattern() const
131+
{
132+
const auto mask = getLayerPatternMask<NLayers>();
133+
if constexpr (NLayers <= MaxLayersInTrackPattern) {
134+
const auto pattern = (getPattern() >> ExtendedPatternShift) & mask;
135+
if (pattern) {
136+
return pattern;
137+
}
138+
}
139+
return getUserField() & mask;
140+
}
141+
GPUhdi() void clearExtendedLayerPattern()
142+
{
143+
setUserField(0);
144+
getParamOut().setUserField(0);
145+
}
146+
GPUhdi() uint32_t getLastClusterLayer() const
111147
{
112148
uint32_t r{0}, v{mPattern & ((1 << 16) - 1)};
113149
while (v >>= 1) {
114150
r++;
115151
}
116152
return r;
117153
}
118-
uint32_t getFirstClusterLayer() const
154+
GPUhdi() uint32_t getFirstClusterLayer() const
119155
{
120156
int s{0};
121157
while (!(mPattern & (1 << s))) {

‎Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TimeFrameGPU.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include "ITStracking/BoundedAllocator.h"
2020
#include "ITStracking/TimeFrame.h"
2121
#include "ITStracking/Configuration.h"
22+
#include "ITStracking/TrackExtensionHypothesis.h"
2223
#include "ITStrackingGPU/Utils.h"
2324

2425
namespace o2::its::gpu
@@ -90,6 +91,7 @@ class TimeFrameGPU : public TimeFrame<NLayers>
9091
void createNeighboursDevice(const unsigned int layer);
9192
void createNeighboursLUTDevice(const int, const unsigned int);
9293
void createTrackITSExtDevice(const size_t);
94+
void createTrackExtensionScratchDevice(const int nThreads, const int maxHypotheses);
9395
void downloadTrackITSExtDevice();
9496
void downloadCellsNeighboursDevice(std::vector<bounded_vector<CellNeighbour>>&, const int);
9597
void downloadNeighboursLUTDevice(bounded_vector<int>&, const int);
@@ -125,6 +127,8 @@ class TimeFrameGPU : public TimeFrame<NLayers>
125127

126128
// Hybrid
127129
TrackITSExt* getDeviceTrackITSExt() { return mTrackITSExtDevice; }
130+
TrackExtensionHypothesis<NLayers>* getDeviceActiveTrackExtensionHypotheses() { return mActiveTrackExtensionHypothesesDevice; }
131+
TrackExtensionHypothesis<NLayers>* getDeviceNextTrackExtensionHypotheses() { return mNextTrackExtensionHypothesesDevice; }
128132
int* getDeviceNeighboursLUT(const int layer) { return mNeighboursLUTDevice[layer]; }
129133
gsl::span<int*> getDeviceNeighboursLUTs() { return mNeighboursLUTDevice; }
130134
CellNeighbour** getDeviceArrayNeighbours() { return mNeighboursDeviceArray; }
@@ -222,6 +226,8 @@ class TimeFrameGPU : public TimeFrame<NLayers>
222226
float** mCellSeedsChi2DeviceArray;
223227

224228
TrackITSExt* mTrackITSExtDevice;
229+
TrackExtensionHypothesis<NLayers>* mActiveTrackExtensionHypothesesDevice{nullptr};
230+
TrackExtensionHypothesis<NLayers>* mNextTrackExtensionHypothesesDevice{nullptr};
225231
std::array<CellNeighbour*, MaxCells> mNeighboursDevice{};
226232
CellNeighbour** mNeighboursDeviceArray{nullptr};
227233
std::array<TrackingFrameInfo*, NLayers> mTrackingFrameInfoDevice;

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

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,13 @@
1313
#ifndef ITSTRACKINGGPU_TRACKINGKERNELS_H_
1414
#define ITSTRACKINGGPU_TRACKINGKERNELS_H_
1515

16+
#include <array>
1617
#include <gsl/gsl>
1718

1819
#include "ITStracking/BoundedAllocator.h"
1920
#include "ITStracking/ROFLookupTables.h"
2021
#include "ITStracking/TrackingTopology.h"
22+
#include "ITStracking/TrackExtensionHypothesis.h"
2123
#include "ITStrackingGPU/Utils.h"
2224
#include "DetectorsBase/Propagator.h"
2325

@@ -208,7 +210,6 @@ void countTrackSeedHandler(TrackSeed<NLayers>* trackSeeds,
208210
const std::vector<float>& layerxX0Host,
209211
const unsigned int nSeeds,
210212
const float Bz,
211-
const int startLevel,
212213
const float maxChi2ClusterAttachment,
213214
const float maxChi2NDF,
214215
const int reseedIfShorter,
@@ -222,20 +223,35 @@ template <int NLayers>
222223
void computeTrackSeedHandler(TrackSeed<NLayers>* trackSeeds,
223224
const TrackingFrameInfo** foundTrackingFrameInfo,
224225
const Cluster** unsortedClusters,
226+
const IndexTableUtils<NLayers>* utils,
227+
const typename ROFMaskTable<NLayers>::View& rofMask,
228+
const typename ROFOverlapTable<NLayers>::View& rofOverlaps,
229+
const Cluster** clusters,
230+
const unsigned char** usedClusters,
231+
const int** clustersIndexTables,
232+
const int** ROFClusters,
225233
o2::its::TrackITSExt* tracks,
226234
const int* seedLUT,
235+
TrackExtensionHypothesis<NLayers>* activeHypotheses,
236+
TrackExtensionHypothesis<NLayers>* nextHypotheses,
227237
const std::vector<float>& layerRadiiHost,
228238
const std::vector<float>& minPtsHost,
229239
const std::vector<float>& layerxX0Host,
230240
const unsigned int nSeeds,
231241
const unsigned int nTracks,
232242
const float Bz,
233-
const int startLevel,
234243
const float maxChi2ClusterAttachment,
235244
const float maxChi2NDF,
236245
const int reseedIfShorter,
237246
const bool repeatRefitOut,
238247
const bool shiftRefToCluster,
248+
const int nLayers,
249+
const int phiBins,
250+
const int maxHypotheses,
251+
const bool extendTop,
252+
const bool extendBot,
253+
const float nSigmaCutPhi,
254+
const float nSigmaCutZ,
239255
const o2::base::Propagator* propagator,
240256
const o2::base::PropagatorF::MatCorrType matCorrType,
241257
o2::its::ExternalAllocator* alloc);

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212

1313
#include <cuda_runtime.h>
1414

15+
#include <algorithm>
1516
#include <unistd.h>
1617
#include <vector>
1718

@@ -581,6 +582,21 @@ void TimeFrameGPU<NLayers>::createTrackITSExtDevice(const size_t nSeeds)
581582
GPUChkErrS(cudaMemset(mTrackITSExtDevice, 0, mNTracks * sizeof(o2::its::TrackITSExt)));
582583
}
583584

585+
template <int NLayers>
586+
void TimeFrameGPU<NLayers>::createTrackExtensionScratchDevice(const int nThreads, const int maxHypotheses)
587+
{
588+
GPUTimer timer("reserving track extension scratch");
589+
const size_t nHypotheses = static_cast<size_t>(std::max(1, nThreads)) * std::max(1, maxHypotheses);
590+
GPULog("gpu-allocation: reserving {} track extension hypotheses per scratch buffer, for {:.2f} MB each.", nHypotheses, nHypotheses * sizeof(o2::its::TrackExtensionHypothesis<NLayers>) / constants::MB);
591+
mActiveTrackExtensionHypothesesDevice = nullptr;
592+
mNextTrackExtensionHypothesesDevice = nullptr;
593+
if (nHypotheses == 0) {
594+
return;
595+
}
596+
allocMem(reinterpret_cast<void**>(&mActiveTrackExtensionHypothesesDevice), nHypotheses * sizeof(o2::its::TrackExtensionHypothesis<NLayers>), this->hasFrameworkAllocator(), (o2::gpu::GPUMemoryResource::MEMORY_GPU | o2::gpu::GPUMemoryResource::MEMORY_STACK));
597+
allocMem(reinterpret_cast<void**>(&mNextTrackExtensionHypothesesDevice), nHypotheses * sizeof(o2::its::TrackExtensionHypothesis<NLayers>), this->hasFrameworkAllocator(), (o2::gpu::GPUMemoryResource::MEMORY_GPU | o2::gpu::GPUMemoryResource::MEMORY_STACK));
598+
}
599+
584600
template <int NLayers>
585601
void TimeFrameGPU<NLayers>::downloadCellsDevice()
586602
{

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

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,15 @@
1212

1313
#include <unistd.h>
1414

15+
#include <algorithm>
16+
#include <array>
17+
1518
#include "ITStrackingGPU/TrackerTraitsGPU.h"
1619
#include "ITStrackingGPU/TrackingKernels.h"
1720
#include "ITStracking/Configuration.h"
1821

1922
namespace o2::its
2023
{
21-
2224
template <int NLayers>
2325
void TrackerTraitsGPU<NLayers>::initialiseTimeFrame(const int iteration)
2426
{
@@ -301,10 +303,11 @@ template <int NLayers>
301303
void TrackerTraitsGPU<NLayers>::findRoads(const int iteration)
302304
{
303305
bounded_vector<bounded_vector<int>> firstClusters(this->mTrkParams[iteration].NLayers, bounded_vector<int>(this->getMemoryPool().get()), this->getMemoryPool().get());
304-
bounded_vector<bounded_vector<int>> sharedFirstClusters(this->mTrkParams[iteration].NLayers, bounded_vector<int>(this->getMemoryPool().get()), this->getMemoryPool().get());
305306
firstClusters.resize(this->mTrkParams[iteration].NLayers);
306-
sharedFirstClusters.resize(this->mTrkParams[iteration].NLayers);
307307
const auto hostTopology = mTimeFrameGPU->getTrackingTopologyView();
308+
const bool extendTop = this->mTrkParams[iteration].PassFlags[IterationStep::TrackFollowerTop];
309+
const bool extendBot = this->mTrkParams[iteration].PassFlags[IterationStep::TrackFollowerBot];
310+
const bool extendTracks = extendTop || extendBot;
308311
for (int startLevel{this->mTrkParams[iteration].CellsPerRoad()}; startLevel >= this->mTrkParams[iteration].CellMinimumLevel(); --startLevel) {
309312
bounded_vector<TrackSeed<NLayers>> trackSeeds(this->getMemoryPool().get());
310313
for (int startCellTopologyId{0}; startCellTopologyId < hostTopology.nCells; ++startCellTopologyId) {
@@ -353,7 +356,6 @@ void TrackerTraitsGPU<NLayers>::findRoads(const int iteration)
353356
this->mTrkParams[iteration].LayerxX0,
354357
trackSeeds.size(),
355358
this->mBz,
356-
startLevel,
357359
this->mTrkParams[iteration].MaxChi2ClusterAttachment,
358360
this->mTrkParams[iteration].MaxChi2NDF,
359361
this->mTrkParams[iteration].ReseedIfShorter,
@@ -363,23 +365,41 @@ void TrackerTraitsGPU<NLayers>::findRoads(const int iteration)
363365
this->mTrkParams[iteration].CorrType,
364366
mTimeFrameGPU->getFrameworkAllocator());
365367
mTimeFrameGPU->createTrackITSExtDevice(trackSeeds.size());
368+
if (extendTracks) {
369+
mTimeFrameGPU->createTrackExtensionScratchDevice(constants::GPUThreadsTotal, this->mTrkParams[iteration].TrackFollowerMaxHypotheses);
370+
}
366371
computeTrackSeedHandler(mTimeFrameGPU->getDeviceTrackSeeds(),
367372
mTimeFrameGPU->getDeviceArrayTrackingFrameInfo(),
368373
mTimeFrameGPU->getDeviceArrayUnsortedClusters(),
374+
mTimeFrameGPU->getDeviceIndexTableUtils(),
375+
mTimeFrameGPU->getDeviceROFMaskTableView(),
376+
mTimeFrameGPU->getDeviceROFOverlapTableView(),
377+
mTimeFrameGPU->getDeviceArrayClusters(),
378+
(const unsigned char**)mTimeFrameGPU->getDeviceArrayUsedClusters(),
379+
mTimeFrameGPU->getDeviceArrayClustersIndexTables(),
380+
mTimeFrameGPU->getDeviceROFrameClusters(),
369381
mTimeFrameGPU->getDeviceTrackITSExt(),
370382
mTimeFrameGPU->getDeviceTrackSeedsLUT(),
383+
extendTracks ? mTimeFrameGPU->getDeviceActiveTrackExtensionHypotheses() : nullptr,
384+
extendTracks ? mTimeFrameGPU->getDeviceNextTrackExtensionHypotheses() : nullptr,
371385
this->mTrkParams[iteration].LayerRadii,
372386
this->mTrkParams[iteration].MinPt,
373387
this->mTrkParams[iteration].LayerxX0,
374388
trackSeeds.size(),
375389
mTimeFrameGPU->getNTrackSeeds(),
376390
this->mBz,
377-
startLevel,
378391
this->mTrkParams[iteration].MaxChi2ClusterAttachment,
379392
this->mTrkParams[iteration].MaxChi2NDF,
380393
this->mTrkParams[iteration].ReseedIfShorter,
381394
this->mTrkParams[iteration].RepeatRefitOut,
382395
this->mTrkParams[iteration].ShiftRefToCluster,
396+
this->mTrkParams[iteration].NLayers,
397+
this->mTrkParams[iteration].PhiBins,
398+
this->mTrkParams[iteration].TrackFollowerMaxHypotheses,
399+
extendTop,
400+
extendBot,
401+
this->mTrkParams[iteration].TrackFollowerNSigmaCutPhi,
402+
this->mTrkParams[iteration].TrackFollowerNSigmaCutZ,
383403
mTimeFrameGPU->getDevicePropagator(),
384404
this->mTrkParams[iteration].CorrType,
385405
mTimeFrameGPU->getFrameworkAllocator());

0 commit comments

Comments
 (0)