Skip to content

Commit 6c94101

Browse files
f3schmconcas
f3sch
authored and
mconcas
committed
ITS: stream comp findCellsNeighbours
Signed-off-by: Felix Schlepper <felix.schlepper@cern.ch>
1 parent bfa2ec2 commit 6c94101

File tree

1 file changed

+107
-43
lines changed

1 file changed

+107
-43
lines changed

‎Detectors/ITSMFT/ITS/tracking/src/TrackerTraits.cxx‎

Lines changed: 107 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,11 @@
1212
/// \file TrackerTraits.cxx
1313
/// \brief
1414
///
15+
1516
#include <algorithm>
16-
#include <cassert>
1717
#include <iostream>
18+
#include <iterator>
19+
#include <ranges>
1820

1921
#ifdef OPTIMISATION_OUTPUT
2022
#include <format>
@@ -498,63 +500,125 @@ void TrackerTraits<nLayers>::findCellsNeighbours(const int iteration)
498500
#endif
499501
for (int iLayer{0}; iLayer < mTrkParams[iteration].CellsPerRoad() - 1; ++iLayer) {
500502
const int nextLayerCellsNum{static_cast<int>(mTimeFrame->getCells()[iLayer + 1].size())};
501-
mTimeFrame->getCellsNeighboursLUT()[iLayer].clear();
502-
mTimeFrame->getCellsNeighboursLUT()[iLayer].resize(nextLayerCellsNum, 0);
503+
deepVectorClear(mTimeFrame->getCellsNeighbours()[iLayer]);
504+
deepVectorClear(mTimeFrame->getCellsNeighboursLUT()[iLayer]);
503505
if (mTimeFrame->getCells()[iLayer + 1].empty() ||
504506
mTimeFrame->getCellsLookupTable()[iLayer].empty()) {
505-
mTimeFrame->getCellsNeighbours()[iLayer].clear();
506507
continue;
507508
}
508509

509-
int layerCellsNum{static_cast<int>(mTimeFrame->getCells()[iLayer].size())};
510-
bounded_vector<std::pair<int, int>> cellsNeighbours(mMemoryPool.get());
511-
cellsNeighbours.reserve(nextLayerCellsNum);
512-
513-
for (int iCell{0}; iCell < layerCellsNum; ++iCell) {
514-
const auto& currentCellSeed{mTimeFrame->getCells()[iLayer][iCell]};
515-
const int nextLayerTrackletIndex{currentCellSeed.getSecondTrackletIndex()};
516-
const int nextLayerFirstCellIndex{mTimeFrame->getCellsLookupTable()[iLayer][nextLayerTrackletIndex]};
517-
const int nextLayerLastCellIndex{mTimeFrame->getCellsLookupTable()[iLayer][nextLayerTrackletIndex + 1]};
518-
for (int iNextCell{nextLayerFirstCellIndex}; iNextCell < nextLayerLastCellIndex; ++iNextCell) {
510+
mTaskArena.execute([&] {
511+
int layerCellsNum{static_cast<int>(mTimeFrame->getCells()[iLayer].size())};
519512

520-
auto nextCellSeed{mTimeFrame->getCells()[iLayer + 1][iNextCell]}; /// copy
521-
if (nextCellSeed.getFirstTrackletIndex() != nextLayerTrackletIndex) {
522-
break;
523-
}
513+
bounded_vector<int> perCellCount(layerCellsNum + 1, 0, mMemoryPool.get());
514+
tbb::parallel_for(
515+
tbb::blocked_range<int>(0, layerCellsNum),
516+
[&](const tbb::blocked_range<int>& Cells) {
517+
for (int iCell = Cells.begin(); iCell < Cells.end(); ++iCell) {
518+
const auto& currentCellSeed{mTimeFrame->getCells()[iLayer][iCell]};
519+
const int nextLayerTrackletIndex{currentCellSeed.getSecondTrackletIndex()};
520+
const int nextLayerFirstCellIndex{mTimeFrame->getCellsLookupTable()[iLayer][nextLayerTrackletIndex]};
521+
const int nextLayerLastCellIndex{mTimeFrame->getCellsLookupTable()[iLayer][nextLayerTrackletIndex + 1]};
522+
523+
int foundNextCells{0};
524+
for (int iNextCell{nextLayerFirstCellIndex}; iNextCell < nextLayerLastCellIndex; ++iNextCell) {
525+
auto nextCellSeed{mTimeFrame->getCells()[iLayer + 1][iNextCell]}; /// copy
526+
if (nextCellSeed.getFirstTrackletIndex() != nextLayerTrackletIndex) {
527+
break;
528+
}
524529

525-
if (!nextCellSeed.rotate(currentCellSeed.getAlpha()) ||
526-
!nextCellSeed.propagateTo(currentCellSeed.getX(), getBz())) {
527-
continue;
528-
}
529-
float chi2 = currentCellSeed.getPredictedChi2(nextCellSeed); /// TODO: switch to the chi2 wrt cluster to avoid correlation
530+
if (!nextCellSeed.rotate(currentCellSeed.getAlpha()) ||
531+
!nextCellSeed.propagateTo(currentCellSeed.getX(), getBz())) {
532+
continue;
533+
}
534+
float chi2 = currentCellSeed.getPredictedChi2(nextCellSeed); /// TODO: switch to the chi2 wrt cluster to avoid correlation
530535

531536
#ifdef OPTIMISATION_OUTPUT
532-
bool good{mTimeFrame->getCellsLabel(iLayer)[iCell] == mTimeFrame->getCellsLabel(iLayer + 1)[iNextCell]};
533-
off << std::format("{}\t{:d}\t{}", iLayer, good, chi2) << std::endl;
537+
bool good{mTimeFrame->getCellsLabel(iLayer)[iCell] == mTimeFrame->getCellsLabel(iLayer + 1)[iNextCell]};
538+
off << std::format("{}\t{:d}\t{}", iLayer, good, chi2) << std::endl;
534539
#endif
535540

536-
if (chi2 > mTrkParams[0].MaxChi2ClusterAttachment) {
537-
continue;
538-
}
541+
if (chi2 > mTrkParams[0].MaxChi2ClusterAttachment) {
542+
continue;
543+
}
544+
++foundNextCells;
545+
}
546+
perCellCount[iCell] = foundNextCells;
547+
}
548+
});
539549

540-
mTimeFrame->getCellsNeighboursLUT()[iLayer][iNextCell]++;
541-
cellsNeighbours.push_back(std::make_pair(iCell, iNextCell));
542-
const int currentCellLevel{currentCellSeed.getLevel()};
550+
std::exclusive_scan(perCellCount.begin(), perCellCount.end(), perCellCount.begin(), 0);
551+
int totalCellNeighbours = perCellCount.back();
552+
if (totalCellNeighbours == 0) {
553+
deepVectorClear(mTimeFrame->getCellsNeighbours()[iLayer]);
554+
return;
555+
}
543556

544-
if (currentCellLevel >= nextCellSeed.getLevel()) {
545-
mTimeFrame->getCells()[iLayer + 1][iNextCell].setLevel(currentCellLevel + 1);
546-
}
557+
struct Neighbor {
558+
int cell{-1}, nextCell{-1}, level{-1};
559+
};
560+
bounded_vector<Neighbor> cellsNeighbours(mMemoryPool.get());
561+
cellsNeighbours.resize(totalCellNeighbours);
562+
563+
tbb::parallel_for(
564+
tbb::blocked_range<int>(0, layerCellsNum),
565+
[&](const tbb::blocked_range<int>& Cells) {
566+
for (int iCell = Cells.begin(); iCell < Cells.end(); ++iCell) {
567+
if (perCellCount[iCell] == perCellCount[iCell + 1]) {
568+
continue;
569+
}
570+
const auto& currentCellSeed{mTimeFrame->getCells()[iLayer][iCell]};
571+
const int nextLayerTrackletIndex{currentCellSeed.getSecondTrackletIndex()};
572+
const int nextLayerFirstCellIndex{mTimeFrame->getCellsLookupTable()[iLayer][nextLayerTrackletIndex]};
573+
const int nextLayerLastCellIndex{mTimeFrame->getCellsLookupTable()[iLayer][nextLayerTrackletIndex + 1]};
574+
575+
int position = perCellCount[iCell];
576+
for (int iNextCell{nextLayerFirstCellIndex}; iNextCell < nextLayerLastCellIndex; ++iNextCell) {
577+
auto nextCellSeed{mTimeFrame->getCells()[iLayer + 1][iNextCell]}; /// copy
578+
if (nextCellSeed.getFirstTrackletIndex() != nextLayerTrackletIndex) {
579+
break;
580+
}
581+
582+
if (!nextCellSeed.rotate(currentCellSeed.getAlpha()) ||
583+
!nextCellSeed.propagateTo(currentCellSeed.getX(), getBz())) {
584+
continue;
585+
}
586+
587+
float chi2 = currentCellSeed.getPredictedChi2(nextCellSeed); /// TODO: switch to the chi2 wrt cluster to avoid correlation
588+
if (chi2 > mTrkParams[0].MaxChi2ClusterAttachment) {
589+
continue;
590+
}
591+
592+
cellsNeighbours[position++] = {iCell, iNextCell, currentCellSeed.getLevel() + 1};
593+
}
594+
}
595+
});
596+
597+
tbb::parallel_sort(cellsNeighbours.begin(), cellsNeighbours.end(), [](const auto& a, const auto& b) {
598+
return a.nextCell < b.nextCell;
599+
});
600+
601+
auto& cellsNeighbourLUT = mTimeFrame->getCellsNeighboursLUT()[iLayer];
602+
cellsNeighbourLUT.assign(nextLayerCellsNum, 0);
603+
for (const auto& neigh : cellsNeighbours) {
604+
++cellsNeighbourLUT[neigh.nextCell];
605+
}
606+
std::inclusive_scan(cellsNeighbourLUT.begin(), cellsNeighbourLUT.end(), cellsNeighbourLUT.begin());
607+
608+
mTimeFrame->getCellsNeighbours()[iLayer].reserve(totalCellNeighbours);
609+
std::ranges::transform(cellsNeighbours, std::back_inserter(mTimeFrame->getCellsNeighbours()[iLayer]), [](const auto& neigh) { return neigh.cell; });
610+
611+
auto it = cellsNeighbours.begin();
612+
while (it != cellsNeighbours.end()) {
613+
const int current_nextCell = it->nextCell;
614+
auto group_end = std::find_if_not(it, cellsNeighbours.end(),
615+
[current_nextCell](const auto& nb) { return nb.nextCell == current_nextCell; });
616+
const auto max_level_it = std::max_element(it, group_end,
617+
[](const auto& a, const auto& b) { return a.level < b.level; });
618+
mTimeFrame->getCells()[iLayer + 1][current_nextCell].setLevel(max_level_it->level);
619+
it = group_end;
547620
}
548-
}
549-
std::sort(cellsNeighbours.begin(), cellsNeighbours.end(), [](const std::pair<int, int>& a, const std::pair<int, int>& b) {
550-
return a.second < b.second;
551621
});
552-
mTimeFrame->getCellsNeighbours()[iLayer].clear();
553-
mTimeFrame->getCellsNeighbours()[iLayer].reserve(cellsNeighbours.size());
554-
for (auto& cellNeighboursIndex : cellsNeighbours) {
555-
mTimeFrame->getCellsNeighbours()[iLayer].push_back(cellNeighboursIndex.first);
556-
}
557-
std::inclusive_scan(mTimeFrame->getCellsNeighboursLUT()[iLayer].begin(), mTimeFrame->getCellsNeighboursLUT()[iLayer].end(), mTimeFrame->getCellsNeighboursLUT()[iLayer].begin());
558622
}
559623
}
560624

0 commit comments

Comments
 (0)