diff --git a/DataFormats/Detectors/ITSMFT/ITS/include/DataFormatsITS/TimeEstBC.h b/DataFormats/Detectors/ITSMFT/ITS/include/DataFormatsITS/TimeEstBC.h index 695d9aff42858..a37ffa63988b7 100644 --- a/DataFormats/Detectors/ITSMFT/ITS/include/DataFormatsITS/TimeEstBC.h +++ b/DataFormats/Detectors/ITSMFT/ITS/include/DataFormatsITS/TimeEstBC.h @@ -22,9 +22,9 @@ namespace o2::its { // Time estimates are given in BC -// error needs to cover maximum 1 orbit +// error needs to cover maximum 1 orbit (uint16_t), but increased due to 2 byte padding using TimeStampType = uint32_t; -using TimeStampErrorType = uint16_t; +using TimeStampErrorType = uint32_t; // this is an symmetric time error [t0-tE, t0+tE] using TimeStamp = o2::dataformats::TimeStampWithError; // this is an asymmetric time interval [t0, t0+tE] used for internal calculations @@ -95,7 +95,7 @@ class TimeEstBC : public o2::dataformats::TimeStampWithErrorsetTimeStampError(static_cast(hi - lo)); } - ClassDefNV(TimeEstBC, 1); + ClassDefNV(TimeEstBC, 2); }; } // namespace o2::its diff --git a/DataFormats/common/src/CommonDataFormatLinkDef.h b/DataFormats/common/src/CommonDataFormatLinkDef.h index d66e89af637cc..2f07b2f9c14b4 100644 --- a/DataFormats/common/src/CommonDataFormatLinkDef.h +++ b/DataFormats/common/src/CommonDataFormatLinkDef.h @@ -31,7 +31,7 @@ #pragma link C++ class o2::dataformats::TimeStampWithError < float, float> + ; #pragma link C++ class o2::dataformats::TimeStampWithError < double, double> + ; #pragma link C++ class o2::dataformats::TimeStampWithError < int, int> + ; -#pragma link C++ class o2::dataformats::TimeStampWithError < uint32_t, uint16_t> + ; +#pragma link C++ class o2::dataformats::TimeStampWithError < uint32_t, uint32_t> + ; #pragma link C++ class o2::dataformats::EvIndex < int, int> + ; #pragma link C++ class o2::dataformats::RangeReference < int, int> + ; diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TrackingKernels.h b/Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TrackingKernels.h index b31d9400657c1..94950c04877b9 100644 --- a/Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TrackingKernels.h +++ b/Detectors/ITSMFT/ITS/tracking/GPU/ITStrackingGPU/TrackingKernels.h @@ -101,6 +101,7 @@ struct TrackingKernels { const float maxChi2ClusterAttachment, const float bz, const unsigned int nCells, + o2::its::ExternalAllocator* alloc, gpu::Stream& stream); static void processNeighboursHandler(const int startLevel, diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/CMakeLists.txt b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/CMakeLists.txt index 435f1e9fab89c..85940efbf2310 100644 --- a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/CMakeLists.txt +++ b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/CMakeLists.txt @@ -32,6 +32,7 @@ if(CUDA_ENABLED) set_property(TARGET ${targetName} PROPERTY CUDA_SEPARABLE_COMPILATION ON) target_compile_options(${targetName} PRIVATE $<$:-diag-error=20014> + $<$:-lineinfo> # $<$:-G;-O0;-Xptxas=-O0> # $<$:-O0;-g> ) diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TimeFrameGPU.cu b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TimeFrameGPU.cu index c6ba5c27c3ef6..593f5529cf8ea 100644 --- a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TimeFrameGPU.cu +++ b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TimeFrameGPU.cu @@ -263,12 +263,12 @@ void TimeFrameGPU::createUsedClustersDevice(const int layer) template void TimeFrameGPU::loadUsedClustersDevice() { + GPUTimer timer("loading used clusters flags"); for (auto iLayer{0}; iLayer < NLayers; ++iLayer) { - GPUTimer timer(mGpuStreams[iLayer], "loading used clusters flags", iLayer); const auto& used = this->mUsedClusters[iLayer]; GPULog("gpu-transfer: loading {} used clusters flags on layer {}, for {:.2f} MB.", used.size(), iLayer, used.size() * sizeof(unsigned char) / constants::MB); if (!used.empty()) { - GPUChkErrS(cudaMemcpyAsync(mUsedClustersDevice[iLayer], used.data(), used.size() * sizeof(unsigned char), cudaMemcpyHostToDevice, mGpuStreams[iLayer].get())); + GPUChkErrS(cudaMemcpyAsync(mUsedClustersDevice[iLayer], used.data(), used.size() * sizeof(unsigned char), cudaMemcpyHostToDevice, Stream::DefaultStream)); } } } diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackerTraitsGPU.cxx b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackerTraitsGPU.cxx index 0eb289ad97169..b80fb6452f8a9 100644 --- a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackerTraitsGPU.cxx +++ b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackerTraitsGPU.cxx @@ -242,6 +242,7 @@ void TrackerTraitsGPU::findCellsNeighbours(const int iteration) this->mTrkParams[iteration].MaxChi2ClusterAttachment, this->mBz, mTimeFrameGPU->getNCells()[sourceCellTopologyId], + mTimeFrameGPU->getFrameworkAllocator(), stream); } return finalizeCellNeighboursHandler(mTimeFrameGPU->getDeviceNeighbours(targetCellTopologyId), diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackingKernels.cu b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackingKernels.cu index 4728572e46c77..a179fd0962dd3 100644 --- a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackingKernels.cu +++ b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackingKernels.cu @@ -176,17 +176,20 @@ GPUg() void __launch_bounds__(GPUThreads, (ExtendTracks ? MinBlocks.fitTrackSeed } } +/// A (source cell, target cell) pair that passed the index and time-stamp cuts and is worth fitting. +struct CellNeighbourCandidate { + int currentCell; + int nextCell; +}; template -GPUg() void __launch_bounds__(GPUThreads, MinBlocks.computeLayerCellNeighbours) computeLayerCellNeighboursKernel( +GPUg() void __launch_bounds__(GPUThreads, MinBlocks.computeLayerCellNeighbours) computeLayerCellNeighbourCandidatesKernel( CellSeed** cellSeedArray, int** cellsLUTs, - CellNeighbour* cellNeighbours, - int* outputCounter, - const int outputCapacity, const int sourceCellTopologyId, const int targetCellTopologyId, - const float maxChi2ClusterAttachment, - const float bz, + CellNeighbourCandidate* candidates, // nullptr on the counting pass + int* candidateCounter, + const int candidateCapacity, // 0 on the counting pass, so nothing is written const unsigned int nCells) { for (int iCurrentCellIndex = blockIdx.x * blockDim.x + threadIdx.x; iCurrentCellIndex < nCells; iCurrentCellIndex += blockDim.x * gridDim.x) { @@ -194,30 +197,55 @@ GPUg() void __launch_bounds__(GPUThreads, MinBlocks.computeLayerCellNeighbours) const int nextLayerTrackletIndex{currentCellSeed.getSecondTrackletIndex()}; const int nextLayerFirstCellIndex{cellsLUTs[targetCellTopologyId][nextLayerTrackletIndex]}; const int nextLayerLastCellIndex{cellsLUTs[targetCellTopologyId][nextLayerTrackletIndex + 1]}; + const auto currentTimeStamp{currentCellSeed.getTimeStamp()}; for (int iNextCell{nextLayerFirstCellIndex}; iNextCell < nextLayerLastCellIndex; ++iNextCell) { - auto nextCellSeed{cellSeedArray[targetCellTopologyId][iNextCell]}; // Copy - if (nextCellSeed.getFirstTrackletIndex() != nextLayerTrackletIndex || !currentCellSeed.getTimeStamp().isCompatible(nextCellSeed.getTimeStamp())) { + const auto& nextCellSeed{cellSeedArray[targetCellTopologyId][iNextCell]}; // No copy: only two accessors are read. + if (nextCellSeed.getFirstTrackletIndex() != nextLayerTrackletIndex || !currentTimeStamp.isCompatible(nextCellSeed.getTimeStamp())) { break; } - - if (!nextCellSeed.rotate(currentCellSeed.getAlpha()) || - !nextCellSeed.propagateTo(currentCellSeed.getX(), bz)) { - continue; + const int outputIndex = atomicAdd(candidateCounter, 1); + if (outputIndex < candidateCapacity) { + candidates[outputIndex] = {iCurrentCellIndex, iNextCell}; } + } + } +} - float chi2 = currentCellSeed.getPredictedChi2Fast(nextCellSeed); - if (chi2 > maxChi2ClusterAttachment) { - continue; - } +template +GPUg() void __launch_bounds__(GPUThreads, MinBlocks.computeLayerCellNeighbours) fitCellNeighboursKernel( + CellSeed** cellSeedArray, + const CellNeighbourCandidate* candidates, + const int nCandidates, + CellNeighbour* cellNeighbours, + int* outputCounter, + const int outputCapacity, + const int sourceCellTopologyId, + const int targetCellTopologyId, + const float maxChi2ClusterAttachment, + const float bz) +{ + for (int iCandidate = blockIdx.x * blockDim.x + threadIdx.x; iCandidate < nCandidates; iCandidate += blockDim.x * gridDim.x) { + const CellNeighbourCandidate candidate = candidates[iCandidate]; + const auto& currentCellSeed{cellSeedArray[sourceCellTopologyId][candidate.currentCell]}; + auto nextCellSeed{cellSeedArray[targetCellTopologyId][candidate.nextCell]}; // Copy - const int currentCellLevel{currentCellSeed.getLevel()}; - const int outputIndex = atomicAdd(outputCounter, 1); - if (outputIndex < outputCapacity) { - cellNeighbours[outputIndex] = {sourceCellTopologyId, iCurrentCellIndex, targetCellTopologyId, iNextCell, currentCellLevel + 1}; - } - if (currentCellLevel >= nextCellSeed.getLevel()) { - atomicMax(cellSeedArray[targetCellTopologyId][iNextCell].getLevelPtr(), currentCellLevel + 1); - } + if (!nextCellSeed.rotate(currentCellSeed.getAlpha()) || + !nextCellSeed.propagateTo(currentCellSeed.getX(), bz)) { + continue; + } + + float chi2 = currentCellSeed.getPredictedChi2Fast(nextCellSeed); + if (chi2 > maxChi2ClusterAttachment) { + continue; + } + + const int currentCellLevel{currentCellSeed.getLevel()}; + const int outputIndex = atomicAdd(outputCounter, 1); + if (outputIndex < outputCapacity) { + cellNeighbours[outputIndex] = {sourceCellTopologyId, candidate.currentCell, targetCellTopologyId, candidate.nextCell, currentCellLevel + 1}; + } + if (currentCellLevel >= nextCellSeed.getLevel()) { + atomicMax(cellSeedArray[targetCellTopologyId][candidate.nextCell].getLevelPtr(), currentCellLevel + 1); } } } @@ -227,15 +255,20 @@ struct CellCandidate { int firstTrackletIndex; int secondTrackletIndex; }; - -template +template GPUg() void __launch_bounds__(GPUThreads, MinBlocks.computeLayerCells) computeLayerCellCandidatesKernel( Tracklet** tracklets, int** trackletsLUT, const int nTrackletsCurrent, const int cellTopologyId, const typename TrackingTopology::View topology, + const Cluster** sortedClusters, + const Cluster** unsortedClusters, + const TrackingFrameInfo** tfInfo, + const float* layerxX0, + const float bz, CellCandidate* candidates, + unsigned int* candidateKeys, int* outputCounter, const int outputCapacity, const float cellDeltaTanLambdaSigma, @@ -260,9 +293,23 @@ GPUg() void __launch_bounds__(GPUThreads, MinBlocks.computeLayerCells) computeLa } const float deltaTanLambda{o2::gpu::CAMath::Abs(currentTracklet.tanLambda - nextTracklet.tanLambda)}; if (deltaTanLambda / cellDeltaTanLambdaSigma < nSigmaCut) { - const int outputIndex = atomicAdd(outputCounter, 1); - if (outputIndex < outputCapacity) { - candidates[outputIndex] = CellCandidate{iCurrentTrackletIndex, iNextTrackletIndex}; + if constexpr (Emit) { + const int outputIndex = atomicAdd(outputCounter, 1); + if (outputIndex < outputCapacity) { + candidates[outputIndex] = CellCandidate{iCurrentTrackletIndex, iNextTrackletIndex}; + const auto firstLink = topology.getLink(cellTopology.firstLink); + const auto secondLink = topology.getLink(cellTopology.secondLink); + const int layers[3] = {firstLink.fromLayer, firstLink.toLayer, secondLink.toLayer}; + const int clusId[3]{ + sortedClusters[layers[0]][currentTracklet.firstClusterIndex].clusterId, + sortedClusters[layers[1]][nextTracklet.firstClusterIndex].clusterId, + sortedClusters[layers[2]][nextTracklet.secondClusterIndex].clusterId}; + const auto seed{o2::its::track::buildTrackSeed(unsortedClusters[layers[0]][clusId[0]], unsortedClusters[layers[1]][clusId[1]], tfInfo[layers[2]][clusId[2]], bz)}; + candidateKeys[outputIndex] = static_cast( + seed.getELossSteps(layerxX0[layers[1]] * constants::Radl * constants::Rho, true)); + } + } else { + atomicAdd(outputCounter, 1); } } } @@ -562,8 +609,13 @@ struct cellNeighbourLess { } }; +/// A (current cell, neighbour-list entry) pair that passed every integer cut and is worth fitting. +struct NeighbourCandidate { + int currentCell; + int neighbourEntry; +}; template -GPUg() void __launch_bounds__(GPUThreads, (std::is_same_v ? MinBlocks.processNeighboursCellSeed : MinBlocks.processNeighboursTrackSeed)) processNeighboursKernel( +GPUg() void __launch_bounds__(GPUThreads, (std::is_same_v ? MinBlocks.processNeighboursCellSeed : MinBlocks.processNeighboursTrackSeed)) processNeighbourCandidatesKernel( const int defaultCellTopologyId, const int level, CellSeed** allCellSeeds, @@ -571,21 +623,12 @@ GPUg() void __launch_bounds__(GPUThreads, (std::is_same_v const int* currentCellIds, const int* currentCellTopologyIds, const unsigned int nCurrentCells, - TrackSeed* updatedCellSeeds, - int* updatedCellsIds, - int* updatedCellTopologyIds, - int* updatedSourceSeeds, - int* outputCounter, - const int outputCapacity, const unsigned char** usedClusters, CellNeighbour** neighbours, int** neighboursLUT, - const TrackingFrameInfo** foundTrackingFrameInfo, - const float* layerxX0, - const float bz, - const float maxChi2ClusterAttachment, - const o2::base::Propagator* propagator, - const o2::base::PropagatorF::MatCorrType matCorrType) + NeighbourCandidate* candidates, // nullptr on the counting pass + int* candidateCounter, + const int candidateCapacity) // 0 on the counting pass, so nothing is written { for (unsigned int iCurrentCell = blockIdx.x * blockDim.x + threadIdx.x; iCurrentCell < nCurrentCells; iCurrentCell += blockDim.x * gridDim.x) { const auto& currentCell{currentCellSeeds[iCurrentCell]}; @@ -613,9 +656,7 @@ GPUg() void __launch_bounds__(GPUThreads, (std::is_same_v for (int iNeighbourCell{startNeighbourId}; iNeighbourCell < endNeighbourId; ++iNeighbourCell) { const auto& neighbourRef = neighbours[cellTopologyId][iNeighbourCell]; - const int neighbourCellTopologyId = neighbourRef.cellTopology; - const int neighbourCellId = neighbourRef.cell; - const auto& neighbourCell = allCellSeeds[neighbourCellTopologyId][neighbourCellId]; + const auto& neighbourCell = allCellSeeds[neighbourRef.cellTopology][neighbourRef.cell]; if (neighbourCell.getSecondTrackletIndex() != currentCell.getFirstTrackletIndex()) { continue; @@ -626,11 +667,54 @@ GPUg() void __launch_bounds__(GPUThreads, (std::is_same_v if (currentCell.getLevel() - 1 != neighbourCell.getLevel()) { continue; } - const int neighbourLayer = neighbourCell.getInnerLayer(); - const int neighbourCluster = neighbourCell.getFirstClusterIndex(); - if (usedClusters[neighbourLayer][neighbourCluster]) { + if (usedClusters[neighbourCell.getInnerLayer()][neighbourCell.getFirstClusterIndex()]) { continue; } + const int outputIndex = atomicAdd(candidateCounter, 1); + if (outputIndex < candidateCapacity) { + candidates[outputIndex] = {static_cast(iCurrentCell), iNeighbourCell}; + } + } + } +} + +template +GPUg() void __launch_bounds__(GPUThreads, (std::is_same_v ? MinBlocks.processNeighboursCellSeed : MinBlocks.processNeighboursTrackSeed)) fitNeighbourCandidatesKernel( + const int defaultCellTopologyId, + CellSeed** allCellSeeds, + CurrentSeed* currentCellSeeds, + const int* currentCellTopologyIds, + const NeighbourCandidate* candidates, + const int* candidateCounter, + const int candidateCapacity, + CellNeighbour** neighbours, + TrackSeed* updatedCellSeeds, + int* updatedCellsIds, + int* updatedCellTopologyIds, + int* updatedSourceSeeds, + int* outputCounter, + const int outputCapacity, + const TrackingFrameInfo** foundTrackingFrameInfo, + const float* layerxX0, + const float bz, + const float maxChi2ClusterAttachment, + const o2::base::Propagator* propagator, + const o2::base::PropagatorF::MatCorrType matCorrType) +{ + const int filled = *candidateCounter < candidateCapacity ? *candidateCounter : candidateCapacity; + for (int iCandidate = blockIdx.x * blockDim.x + threadIdx.x; iCandidate < filled; iCandidate += blockDim.x * gridDim.x) { + const NeighbourCandidate candidate = candidates[iCandidate]; + const unsigned int iCurrentCell = static_cast(candidate.currentCell); + const auto& currentCell{currentCellSeeds[iCurrentCell]}; + const int cellTopologyId = currentCellTopologyIds == nullptr ? defaultCellTopologyId : currentCellTopologyIds[iCurrentCell]; + const auto& neighbourRef = neighbours[cellTopologyId][candidate.neighbourEntry]; + const int neighbourCellTopologyId = neighbourRef.cellTopology; + const int neighbourCellId = neighbourRef.cell; + const auto& neighbourCell = allCellSeeds[neighbourCellTopologyId][neighbourCellId]; + const int neighbourLayer = neighbourCell.getInnerLayer(); + const int neighbourCluster = neighbourCell.getFirstClusterIndex(); + + { TrackSeed seed{currentCell}; auto& trHit = foundTrackingFrameInfo[neighbourLayer][neighbourCluster]; @@ -803,10 +887,10 @@ int TrackingKernels::computeCellsHandler( const int candidateBlocks = gpu::gridBlocks(gpu::ResidentBlocks.computeLayerCells); GPUChkErrS(cudaMemsetAsync(outputCounter, 0, sizeof(int), stream.get())); - gpu::computeLayerCellCandidatesKernel<<>>( + gpu::computeLayerCellCandidatesKernel<<>>( tracklets, trackletsLUT, nTracklets, cellTopologyId, topology, - nullptr, // counting pass: capacity 0, so nothing is written - outputCounter, 0, cellDeltaTanLambdaSigma, nSigmaCut); + sortedClusters, unsortedClusters, tfInfo, layerxX0, bz, + nullptr, nullptr, outputCounter, 0, cellDeltaTanLambdaSigma, nSigmaCut); int nCandidates = 0; GPUChkErrS(cudaMemcpyAsync(&nCandidates, outputCounter, sizeof(int), cudaMemcpyDeviceToHost, stream.get())); stream.sync(); @@ -819,11 +903,20 @@ int TrackingKernels::computeCellsHandler( } auto candidates = candidateAllocator.allocate(nCandidates); + gpu::TypedAllocator candidateKeyAllocator(alloc); + auto candidateKeys = candidateKeyAllocator.allocate(nCandidates); GPUChkErrS(cudaMemsetAsync(outputCounter, 0, sizeof(int), stream.get())); - gpu::computeLayerCellCandidatesKernel<<>>( + gpu::computeLayerCellCandidatesKernel<<>>( tracklets, trackletsLUT, nTracklets, cellTopologyId, topology, - thrust::raw_pointer_cast(candidates), outputCounter, nCandidates, - cellDeltaTanLambdaSigma, nSigmaCut); + sortedClusters, unsortedClusters, tfInfo, layerxX0, bz, + thrust::raw_pointer_cast(candidates), thrust::raw_pointer_cast(candidateKeys), + outputCounter, nCandidates, cellDeltaTanLambdaSigma, nSigmaCut); + + // order the candidates by momentum before fitting them, so that the ELoss iteration count inside is uniform + { + auto candidatePolicy = THRUST_NAMESPACE::par_nosync(gpu::TypedAllocator(alloc)).on(stream.get()); + thrust::sort_by_key(candidatePolicy, candidateKeys, candidateKeys + nCandidates, candidates); + } GPUChkErrS(cudaMemsetAsync(outputCounter, 0, sizeof(int), stream.get())); gpu::fitLayerCellsKernel<<>>( @@ -885,19 +978,53 @@ void TrackingKernels::computeCellNeighboursHandler(CellSeed** cellsLaye const float maxChi2ClusterAttachment, const float bz, const unsigned int nCells, + o2::its::ExternalAllocator* alloc, gpu::Stream& stream) { - gpu::computeLayerCellNeighboursKernel<<>>( + const int neighbourBlocks = gpu::gridBlocks(gpu::ResidentBlocks.computeLayerCellNeighbours); + + constexpr uint64_t CandidateTag = qStr2Tag("ITSNGHCA"); + alloc->pushTagOnStack(CandidateTag); + gpu::TypedAllocator candidateAllocator(alloc); + gpu::TypedAllocator counterAllocator(alloc); + + auto candidateCounter = counterAllocator.allocate(1); + int* candidateCounterPtr = thrust::raw_pointer_cast(candidateCounter); + + GPUChkErrS(cudaMemsetAsync(candidateCounterPtr, 0, sizeof(int), stream.get())); + gpu::computeLayerCellNeighbourCandidatesKernel<<>>( + cellsLayersDevice, cellsLUTs, sourceCellTopologyId, targetCellTopologyId, + nullptr, // counting pass: capacity 0, so nothing is written + candidateCounterPtr, 0, nCells); + int nCandidates = 0; + GPUChkErrS(cudaMemcpyAsync(&nCandidates, candidateCounterPtr, sizeof(int), cudaMemcpyDeviceToHost, stream.get())); + stream.sync(); + + if (nCandidates == 0) { + alloc->popTagOffStack(CandidateTag); + return; + } + + auto candidates = candidateAllocator.allocate(nCandidates); + GPUChkErrS(cudaMemsetAsync(candidateCounterPtr, 0, sizeof(int), stream.get())); + gpu::computeLayerCellNeighbourCandidatesKernel<<>>( + cellsLayersDevice, cellsLUTs, sourceCellTopologyId, targetCellTopologyId, + thrust::raw_pointer_cast(candidates), candidateCounterPtr, nCandidates, nCells); + + gpu::fitCellNeighboursKernel<<>>( cellsLayersDevice, - cellsLUTs, + thrust::raw_pointer_cast(candidates), + nCandidates, cellNeighbours, outputCounter, capacity, sourceCellTopologyId, targetCellTopologyId, maxChi2ClusterAttachment, - bz, - nCells); + bz); + + stream.sync(); // the candidate slab must outlive the kernels reading it + alloc->popTagOffStack(CandidateTag); } int finalizeCellNeighboursHandler(CellNeighbour* cellNeighbours, @@ -980,6 +1107,9 @@ void TrackingKernels::processNeighboursHandler(const int startLevel, auto roadKey = [&](const int level) { return CapacityEstimator::makeKey(SlabSite::Roads, iteration, CapacityEstimator::makeVariant(startLevel, level), startCellTopologyId); }; + auto candidateKey = [&](const int level) { + return CapacityEstimator::makeKey(SlabSite::RoadCandidates, iteration, CapacityEstimator::makeVariant(startLevel, level), startCellTopologyId); + }; struct Slab { thrust::device_ptr> seeds{}; @@ -1035,35 +1165,53 @@ void TrackingKernels::processNeighboursHandler(const int startLevel, stagedCellIds = out.cellIds; stagedCellTopologyIds = out.cellTopologyIds; #endif + using LevelSeed = std::remove_pointer_t; + const int neighbourGrid = gpu::gridBlocks(std::is_same_v + ? gpu::ResidentBlocks.processNeighboursCellSeed + : gpu::ResidentBlocks.processNeighboursTrackSeed); GPUChkErrS(cudaMemsetAsync(thrust::raw_pointer_cast(outputCounter), 0, sizeof(int), gpu::Stream::DefaultStream)); - gpu::processNeighboursKernel><<, CellSeed> - ? gpu::ResidentBlocks.processNeighboursCellSeed - : gpu::ResidentBlocks.processNeighboursTrackSeed), - gpu::GPUThreads>>>(topologyId, - level, - allCellSeeds, - levelSeeds, - levelCellIds, - levelCellTopologyIds, - nLevelSeeds, - thrust::raw_pointer_cast(staged), - thrust::raw_pointer_cast(stagedCellIds), - thrust::raw_pointer_cast(stagedCellTopologyIds), - thrust::raw_pointer_cast(sourceSeeds), - thrust::raw_pointer_cast(outputCounter), - out.capacity, - usedClusters, - neighbours, - neighboursDeviceLUTs, - foundTrackingFrameInfo, - layerxX0, - bz, - maxChi2ClusterAttachment, - propagator, - matCorrType); + + constexpr uint64_t CandidateTag = qStr2Tag("ITS_PNCA"); + alloc->pushTagOnStack(CandidateTag); + auto allocCandidate = gpu::TypedAllocator(alloc); + auto candidateCounter = allocInt.allocate(1); + int* candidateCounterPtr = thrust::raw_pointer_cast(candidateCounter); + + gpu::NeighbourCandidate* candidatePtr = nullptr; + int candidateCapacity = 0; + const int nCandidates = runOnSlab( + estimator, candidateKey(level), static_cast(nLevelSeeds), [&](const int attemptCapacity) { + if (attemptCapacity > candidateCapacity) { + candidatePtr = thrust::raw_pointer_cast(allocCandidate.allocate(attemptCapacity)); + candidateCapacity = attemptCapacity; + } + GPUChkErrS(cudaMemsetAsync(candidateCounterPtr, 0, sizeof(int), gpu::Stream::DefaultStream)); + gpu::processNeighbourCandidatesKernel<<>>( + topologyId, level, allCellSeeds, levelSeeds, levelCellIds, levelCellTopologyIds, nLevelSeeds, + usedClusters, neighbours, neighboursDeviceLUTs, + candidatePtr, candidateCounterPtr, attemptCapacity); + int produced{0}; + GPUChkErrS(cudaMemcpyAsync(&produced, candidateCounterPtr, sizeof(int), cudaMemcpyDeviceToHost, gpu::Stream::DefaultStream)); + GPUChkErrS(cudaStreamSynchronize(gpu::Stream::DefaultStream)); + return produced; + }); + + if (nCandidates > 0) { + gpu::fitNeighbourCandidatesKernel<<>>( + topologyId, allCellSeeds, levelSeeds, levelCellTopologyIds, + candidatePtr, candidateCounterPtr, candidateCapacity, neighbours, + thrust::raw_pointer_cast(staged), + thrust::raw_pointer_cast(stagedCellIds), + thrust::raw_pointer_cast(stagedCellTopologyIds), + thrust::raw_pointer_cast(sourceSeeds), + thrust::raw_pointer_cast(outputCounter), + out.capacity, + foundTrackingFrameInfo, layerxX0, bz, maxChi2ClusterAttachment, propagator, matCorrType); + } int wanted{0}; GPUChkErrS(cudaMemcpyAsync(&wanted, thrust::raw_pointer_cast(outputCounter), sizeof(int), cudaMemcpyDeviceToHost, gpu::Stream::DefaultStream)); GPUChkErrS(cudaStreamSynchronize(gpu::Stream::DefaultStream)); + alloc->popTagOffStack(CandidateTag); return wanted; }, static_cast(out.capacity)); @@ -1151,6 +1299,7 @@ int TrackingKernels::computeTrackSeedHandler(TrackSeed* trackS o2::its::ExternalAllocator* alloc) { GPUChkErrS(cudaMemsetAsync(outputCounter, 0, sizeof(int), gpu::Stream::DefaultStream)); + // track follower is compiled out of the kernel when no iteration asks for it const auto launchFit = [&](auto extendTracks) { gpu::fitTrackSeedsKernel<<