Skip to content

Commit 2bd4652

Browse files
Felix Schlepperf3sch
Felix Schlepper
authored and
f3sch
committed
ITSGPU: count possible cells before
Signed-off-by: Felix Schlepper <felix.schlepper@cern.ch>
1 parent 66676ed commit 2bd4652

1 file changed

Lines changed: 137 additions & 78 deletions

File tree

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

Lines changed: 137 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -222,29 +222,26 @@ GPUg() void __launch_bounds__(GPUThreads, MinBlocks.computeLayerCellNeighbours)
222222
}
223223
}
224224

225+
/// A tracklet pair that passed the cheap cuts and is worth fitting.
226+
struct CellCandidate {
227+
int firstTrackletIndex;
228+
int secondTrackletIndex;
229+
};
230+
225231
template <int NLayers>
226-
GPUg() void __launch_bounds__(GPUThreads, MinBlocks.computeLayerCells) computeLayerCellsKernel(
227-
const Cluster** sortedClusters,
228-
const Cluster** unsortedClusters,
229-
const TrackingFrameInfo** tfInfo,
232+
GPUg() void __launch_bounds__(GPUThreads, MinBlocks.computeLayerCells) computeLayerCellCandidatesKernel(
230233
Tracklet** tracklets,
231234
int** trackletsLUT,
232235
const int nTrackletsCurrent,
233236
const int cellTopologyId,
234237
const typename TrackingTopology<NLayers>::View topology,
235-
CellSeed* cells,
238+
CellCandidate* candidates,
236239
int* outputCounter,
237240
const int outputCapacity,
238-
const float* layerxX0,
239-
const float bz,
240-
const float maxChi2ClusterAttachment,
241241
const float cellDeltaTanLambdaSigma,
242242
const float nSigmaCut)
243243
{
244244
const auto cellTopology = topology.getCell(cellTopologyId);
245-
const auto first = topology.getLink(cellTopology.firstLink);
246-
const auto second = topology.getLink(cellTopology.secondLink);
247-
const int layers[3] = {first.fromLayer, first.toLayer, second.toLayer};
248245
for (int iCurrentTrackletIndex = blockIdx.x * blockDim.x + threadIdx.x; iCurrentTrackletIndex < nTrackletsCurrent; iCurrentTrackletIndex += blockDim.x * gridDim.x) {
249246
const Tracklet& currentTracklet = tracklets[cellTopology.firstLink][iCurrentTrackletIndex];
250247
const int nextLayerClusterIndex{currentTracklet.secondClusterIndex};
@@ -262,56 +259,88 @@ GPUg() void __launch_bounds__(GPUThreads, MinBlocks.computeLayerCells) computeLa
262259
continue;
263260
}
264261
const float deltaTanLambda{o2::gpu::CAMath::Abs(currentTracklet.tanLambda - nextTracklet.tanLambda)};
265-
266262
if (deltaTanLambda / cellDeltaTanLambdaSigma < nSigmaCut) {
267-
const int clusId[3]{
268-
sortedClusters[layers[0]][currentTracklet.firstClusterIndex].clusterId,
269-
sortedClusters[layers[1]][nextTracklet.firstClusterIndex].clusterId,
270-
sortedClusters[layers[2]][nextTracklet.secondClusterIndex].clusterId};
271-
272-
const auto& cluster1Glo = unsortedClusters[layers[0]][clusId[0]];
273-
const auto& cluster2Glo = unsortedClusters[layers[1]][clusId[1]];
274-
const auto& cluster3Tf = tfInfo[layers[2]][clusId[2]];
275-
auto track{o2::its::track::buildTrackSeed(cluster1Glo, cluster2Glo, cluster3Tf, bz)};
276-
float chi2{0.f};
277-
bool good{false};
278-
for (int iC{2}; iC--;) {
279-
const TrackingFrameInfo& trackingHit = tfInfo[layers[iC]][clusId[iC]];
280-
if (!track.rotate(trackingHit.alphaTrackingFrame)) {
281-
break;
282-
}
283-
if (!track.propagateTo(trackingHit.xTrackingFrame, bz)) {
284-
break;
285-
}
286-
287-
if (!track.correctForMaterial(layerxX0[layers[iC]], layerxX0[layers[iC]] * constants::Radl * constants::Rho, true)) {
288-
break;
289-
}
290-
291-
const auto predChi2{track.getPredictedChi2Quiet(trackingHit.positionTrackingFrame, trackingHit.covarianceTrackingFrame)};
292-
if (!track.o2::track::TrackParCov::update(trackingHit.positionTrackingFrame, trackingHit.covarianceTrackingFrame)) {
293-
break;
294-
}
295-
if (!iC && predChi2 > maxChi2ClusterAttachment) {
296-
break;
297-
}
298-
good = !iC;
299-
chi2 += predChi2;
300-
}
301-
if (!good) {
302-
continue;
303-
}
304-
TimeEstBC ts = currentTracklet.getTimeStamp();
305-
ts += nextTracklet.getTimeStamp();
306263
const int outputIndex = atomicAdd(outputCounter, 1);
307264
if (outputIndex < outputCapacity) {
308-
new (cells + outputIndex) CellSeed{cellTopology.hitLayerMask, clusId[0], clusId[1], clusId[2], iCurrentTrackletIndex, iNextTrackletIndex, track, chi2, ts};
265+
candidates[outputIndex] = CellCandidate{iCurrentTrackletIndex, iNextTrackletIndex};
309266
}
310267
}
311268
}
312269
}
313270
}
314271

272+
/// Fit one tracklet pair per thread, emitting a cell for each pair that survives the fit.
273+
template <int NLayers>
274+
GPUg() void __launch_bounds__(GPUThreads, MinBlocks.computeLayerCells) fitLayerCellsKernel(
275+
const Cluster** sortedClusters,
276+
const Cluster** unsortedClusters,
277+
const TrackingFrameInfo** tfInfo,
278+
Tracklet** tracklets,
279+
const CellCandidate* candidates,
280+
const int nCandidates,
281+
const int cellTopologyId,
282+
const typename TrackingTopology<NLayers>::View topology,
283+
CellSeed* cells,
284+
int* outputCounter,
285+
const int outputCapacity,
286+
const float* layerxX0,
287+
const float bz,
288+
const float maxChi2ClusterAttachment)
289+
{
290+
const auto cellTopology = topology.getCell(cellTopologyId);
291+
const auto first = topology.getLink(cellTopology.firstLink);
292+
const auto second = topology.getLink(cellTopology.secondLink);
293+
const int layers[3] = {first.fromLayer, first.toLayer, second.toLayer};
294+
for (int iCandidate = blockIdx.x * blockDim.x + threadIdx.x; iCandidate < nCandidates; iCandidate += blockDim.x * gridDim.x) {
295+
const CellCandidate candidate = candidates[iCandidate];
296+
const Tracklet& currentTracklet = tracklets[cellTopology.firstLink][candidate.firstTrackletIndex];
297+
const Tracklet& nextTracklet = tracklets[cellTopology.secondLink][candidate.secondTrackletIndex];
298+
const int clusId[3]{
299+
sortedClusters[layers[0]][currentTracklet.firstClusterIndex].clusterId,
300+
sortedClusters[layers[1]][nextTracklet.firstClusterIndex].clusterId,
301+
sortedClusters[layers[2]][nextTracklet.secondClusterIndex].clusterId};
302+
303+
const auto& cluster1Glo = unsortedClusters[layers[0]][clusId[0]];
304+
const auto& cluster2Glo = unsortedClusters[layers[1]][clusId[1]];
305+
const auto& cluster3Tf = tfInfo[layers[2]][clusId[2]];
306+
auto track{o2::its::track::buildTrackSeed(cluster1Glo, cluster2Glo, cluster3Tf, bz)};
307+
float chi2{0.f};
308+
bool good{false};
309+
for (int iC{2}; iC--;) {
310+
const TrackingFrameInfo& trackingHit = tfInfo[layers[iC]][clusId[iC]];
311+
if (!track.rotate(trackingHit.alphaTrackingFrame)) {
312+
break;
313+
}
314+
if (!track.propagateTo(trackingHit.xTrackingFrame, bz)) {
315+
break;
316+
}
317+
318+
if (!track.correctForMaterial(layerxX0[layers[iC]], layerxX0[layers[iC]] * constants::Radl * constants::Rho, true)) {
319+
break;
320+
}
321+
322+
const auto predChi2{track.getPredictedChi2Quiet(trackingHit.positionTrackingFrame, trackingHit.covarianceTrackingFrame)};
323+
if (!track.o2::track::TrackParCov::update(trackingHit.positionTrackingFrame, trackingHit.covarianceTrackingFrame)) {
324+
break;
325+
}
326+
if (!iC && predChi2 > maxChi2ClusterAttachment) {
327+
break;
328+
}
329+
good = !iC;
330+
chi2 += predChi2;
331+
}
332+
if (!good) {
333+
continue;
334+
}
335+
TimeEstBC ts = currentTracklet.getTimeStamp();
336+
ts += nextTracklet.getTimeStamp();
337+
const int outputIndex = atomicAdd(outputCounter, 1);
338+
if (outputIndex < outputCapacity) {
339+
new (cells + outputIndex) CellSeed{cellTopology.hitLayerMask, clusId[0], clusId[1], clusId[2], candidate.firstTrackletIndex, candidate.secondTrackletIndex, track, chi2, ts};
340+
}
341+
}
342+
}
343+
315344
template <int NLayers>
316345
GPUg() void __launch_bounds__(GPUThreads, MinBlocks.computeLayerTracklets) computeLayerTrackletsMultiROFKernel(
317346
const IndexTableUtils<NLayers>* utils,
@@ -501,8 +530,16 @@ struct trackletClusterKey {
501530
}
502531
};
503532

504-
struct cellFirstTrackletIndex {
505-
GPUhd() int operator()(const CellSeed& cell) const { return cell.getFirstTrackletIndex(); }
533+
struct cellTrackletKey {
534+
GPUhd() uint64_t operator()(const CellSeed& cell) const
535+
{
536+
return (static_cast<uint64_t>(cell.getFirstTrackletIndex()) << 32) | static_cast<uint32_t>(cell.getSecondTrackletIndex());
537+
}
538+
};
539+
540+
/// The first tracklet index recovered from a cellTrackletKey, for building the lookup table.
541+
struct cellKeyFirstTracklet {
542+
GPUhd() int operator()(const uint64_t key) const { return static_cast<int>(key >> 32); }
506543
};
507544

508545
struct cellNeighbourNextCell {
@@ -759,29 +796,49 @@ int TrackingKernels<NLayers>::computeCellsHandler(
759796
int emitted = 0;
760797
auto& stream = streams[cellTopologyId];
761798
int* outputCounter = cellsLUTsHost + nTracklets;
799+
800+
constexpr uint64_t CandidateTag = qStr2Tag("ITSCELCA");
801+
alloc->pushTagOnStack(CandidateTag);
802+
gpu::TypedAllocator<gpu::CellCandidate> candidateAllocator(alloc);
803+
804+
const int candidateBlocks = gpu::gridBlocks(gpu::ResidentBlocks.computeLayerCells);
762805
GPUChkErrS(cudaMemsetAsync(outputCounter, 0, sizeof(int), stream.get()));
763-
gpu::computeLayerCellsKernel<NLayers><<<gpu::gridBlocks(gpu::ResidentBlocks.computeLayerCells), gpu::GPUThreads, 0, stream.get()>>>(
764-
sortedClusters, // const Cluster**
765-
unsortedClusters, // const Cluster**
766-
tfInfo, // const TrackingFrameInfo**
767-
tracklets, // const Tracklets**
768-
trackletsLUT, // const int**
769-
nTracklets, // const int
770-
cellTopologyId, // const int
771-
topology,
772-
cells, // CellSeed*
773-
outputCounter, // int*
774-
capacity, // const int
775-
layerxX0,
776-
bz, // const float
777-
maxChi2ClusterAttachment, // const float
778-
cellDeltaTanLambdaSigma, // const float
779-
nSigmaCut); // const float
806+
gpu::computeLayerCellCandidatesKernel<NLayers><<<candidateBlocks, gpu::GPUThreads, 0, stream.get()>>>(
807+
tracklets, trackletsLUT, nTracklets, cellTopologyId, topology,
808+
nullptr, // counting pass: capacity 0, so nothing is written
809+
outputCounter, 0, cellDeltaTanLambdaSigma, nSigmaCut);
810+
int nCandidates = 0;
811+
GPUChkErrS(cudaMemcpyAsync(&nCandidates, outputCounter, sizeof(int), cudaMemcpyDeviceToHost, stream.get()));
812+
stream.sync();
813+
814+
if (nCandidates == 0) {
815+
GPUChkErrS(cudaMemsetAsync(cellsLUTsHost, 0, (nTracklets + 1) * sizeof(int), stream.get()));
816+
stream.sync();
817+
alloc->popTagOffStack(CandidateTag);
818+
return 0;
819+
}
820+
821+
auto candidates = candidateAllocator.allocate(nCandidates);
822+
GPUChkErrS(cudaMemsetAsync(outputCounter, 0, sizeof(int), stream.get()));
823+
gpu::computeLayerCellCandidatesKernel<NLayers><<<candidateBlocks, gpu::GPUThreads, 0, stream.get()>>>(
824+
tracklets, trackletsLUT, nTracklets, cellTopologyId, topology,
825+
thrust::raw_pointer_cast(candidates), outputCounter, nCandidates,
826+
cellDeltaTanLambdaSigma, nSigmaCut);
827+
828+
GPUChkErrS(cudaMemsetAsync(outputCounter, 0, sizeof(int), stream.get()));
829+
gpu::fitLayerCellsKernel<NLayers><<<candidateBlocks, gpu::GPUThreads, 0, stream.get()>>>(
830+
sortedClusters, unsortedClusters, tfInfo, tracklets,
831+
thrust::raw_pointer_cast(candidates), nCandidates,
832+
cellTopologyId, topology, cells, outputCounter, capacity,
833+
layerxX0, bz, maxChi2ClusterAttachment);
780834
GPUChkErrS(cudaMemcpyAsync(&emitted, outputCounter, sizeof(int), cudaMemcpyDeviceToHost, stream.get()));
781835
stream.sync();
836+
alloc->popTagOffStack(CandidateTag);
837+
782838
if (emitted > capacity) {
783839
return emitted;
784840
}
841+
785842
auto nosync_policy = THRUST_NAMESPACE::par_nosync(gpu::TypedAllocator<char>(alloc)).on(stream.get());
786843
GPUChkErrS(cudaMemsetAsync(cellsLUTsHost, 0, (nTracklets + 1) * sizeof(int), stream.get()));
787844
if (emitted == 0) {
@@ -790,20 +847,22 @@ int TrackingKernels<NLayers>::computeCellsHandler(
790847
constexpr uint64_t SortTag = qStr2Tag("ITSCELSR");
791848
alloc->pushTagOnStack(SortTag);
792849
gpu::TypedAllocator<int> keyAllocator(alloc);
850+
gpu::TypedAllocator<uint64_t> sortKeyAllocator(alloc);
793851
gpu::TypedAllocator<CellSeed> cellAllocator(alloc);
794-
auto keys = keyAllocator.allocate(emitted);
852+
auto keys = sortKeyAllocator.allocate(emitted);
795853
auto permutation = keyAllocator.allocate(emitted);
796854
thrust::device_ptr<CellSeed> cellsPtr(cells);
797-
thrust::transform(nosync_policy, cellsPtr, cellsPtr + emitted, keys, gpu::cellFirstTrackletIndex{});
855+
thrust::transform(nosync_policy, cellsPtr, cellsPtr + emitted, keys, gpu::cellTrackletKey{});
798856
thrust::sequence(nosync_policy, permutation, permutation + emitted);
799857
thrust::stable_sort_by_key(nosync_policy, keys, keys + emitted, permutation);
800-
gpu::compileLookupTableKernel<<<gpu::gridBlocks(gpu::ResidentBlocks.compileLookupTable), gpu::GPUThreads, 0, stream.get()>>>(
801-
thrust::raw_pointer_cast(keys),
802-
cellsLUTsHost,
803-
emitted);
804-
thrust::exclusive_scan(nosync_policy, cellsLUTsHost, cellsLUTsHost + nTracklets + 1, cellsLUTsHost);
805858
auto sortedCells = cellAllocator.allocate(emitted);
806859
thrust::gather(nosync_policy, permutation, permutation + emitted, cellsPtr, sortedCells);
860+
auto lutKeys = keyAllocator.allocate(emitted);
861+
thrust::transform(nosync_policy, keys, keys + emitted, lutKeys, gpu::cellKeyFirstTracklet{});
862+
gpu::compileLookupTableKernel<<<gpu::gridBlocks(gpu::ResidentBlocks.compileLookupTable), gpu::GPUThreads, 0, stream.get()>>>(thrust::raw_pointer_cast(lutKeys),
863+
cellsLUTsHost,
864+
emitted);
865+
thrust::exclusive_scan(nosync_policy, cellsLUTsHost, cellsLUTsHost + nTracklets + 1, cellsLUTsHost);
807866
GPUChkErrS(cudaMemcpyAsync(cells, thrust::raw_pointer_cast(sortedCells), emitted * sizeof(CellSeed), cudaMemcpyDeviceToDevice, stream.get()));
808867
stream.sync();
809868
alloc->popTagOffStack(SortTag);

0 commit comments

Comments
 (0)