Skip to content

Commit

Permalink
Improve ITS/TPC afterburner memory management
Browse files Browse the repository at this point in the history
  • Loading branch information
shahor02 committed Oct 24, 2023
1 parent 0f125fe commit 81a003d
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 29 deletions.
33 changes: 25 additions & 8 deletions Detectors/GlobalTracking/include/GlobalTracking/MatchTPCITS.h
Original file line number Diff line number Diff line change
Expand Up @@ -198,6 +198,8 @@ struct ABTrackLink : public o2::track::TrackParCov {
uint8_t ladderID = 0xff; ///< ladder ID in the layer (used for seeds with 2 hits in the layer)
float chi2 = 0.f; ///< chi2 after update

ABTrackLink() = default;
~ABTrackLink() = default;
ABTrackLink(const o2::track::TrackParCov& tr, int cl, int parID, int nextID, int lr, int nc, int ld, float _chi2)
: o2::track::TrackParCov(tr), clID(cl), parentID(parID), nextOnLr(nextID), layerID(int8_t(lr)), nContLayers(int8_t(nc)), ladderID(uint8_t(ld)), chi2(_chi2) {}

Expand All @@ -208,17 +210,25 @@ struct ABTrackLink : public o2::track::TrackParCov {
float chi2NormPredict(float chi2cl) const { return (chi2 + chi2cl) / (1 + o2::its::RecoGeomHelper::getNLayers() - layerID); }
};

struct LinksPoolMT {
std::vector<std::deque<ABTrackLink>> threadPool;
};

// AB primary seed: TPC track propagated to outermost ITS layer under specific InteractionCandidate hypothesis
struct TPCABSeed {
static constexpr int8_t NeedAlternative = -3;
int tpcWID = MinusOne; ///< TPC track ID
int ICCanID = MinusOne; ///< interaction candidate ID (they are sorted in increasing time)
int winLinkID = MinusOne; ///< ID of the validated link
uint32_t linksEntry = 0; ///< 1st entry of the link
int nLinks = 0; ///< number of links
int8_t lowestLayer = o2::its::RecoGeomHelper::getNLayers(); ///< lowest layer reached
int8_t status = MinusOne; ///< status (RS TODO)
uint8_t threadID = 0; ///< thread ID
o2::track::TrackParCov track{}; ///< Seed propagated to the outer layer under certain time constraint
std::array<int, o2::its::RecoGeomHelper::getNLayers()> firstInLr; ///< entry of 1st (best) hypothesis on each layer
std::vector<ABTrackLink> trackLinks{}; ///< links
static LinksPoolMT* gLinksPool; ///< pool of links per thread

TPCABSeed(int id, int ic, const o2::track::TrackParCov& trc) : tpcWID(id), ICCanID(ic), track(trc)
{
firstInLr.fill(MinusOne);
Expand All @@ -234,8 +244,14 @@ struct TPCABSeed {
int8_t getNLayers() const { return o2::its::RecoGeomHelper::getNLayers() - lowestLayer; }
bool needAlteranative() const { return status == NeedAlternative; }
void setNeedAlternative() { status = NeedAlternative; }
ABTrackLink& getLink(int i) { return trackLinks[i]; }
const ABTrackLink& getLink(int i) const { return trackLinks[i]; }
ABTrackLink& getLink(int i) { return gLinksPool->threadPool[threadID][i + linksEntry]; }
const ABTrackLink& getLink(int i) const { return gLinksPool->threadPool[threadID][i + linksEntry]; }
void addLink(const o2::track::TrackParCov& trc, int clID, int parentID, int nextID, int lr, int nc, int laddID, float chi2)
{
gLinksPool->threadPool[threadID].push_back({trc, clID, parentID, nextID, lr, nc, laddID, chi2});
nLinks++;
}
int getNLinks() const { return nLinks; }
auto getBestLinkID() const
{
return lowestLayer < o2::its::RecoGeomHelper::getNLayers() ? firstInLr[lowestLayer] : -1;
Expand All @@ -244,7 +260,7 @@ struct TPCABSeed {
{
// check if some clusters used by the link or its parents are forbidden (already used by validatet track)
while (linkID > MinusOne) {
const auto& link = trackLinks[linkID];
const auto& link = getLink(linkID);
if (link.clID > MinusOne && clStatus[link.clID] != MinusOne) {
return true;
}
Expand All @@ -256,15 +272,15 @@ struct TPCABSeed {
{
// check if some clusters used by the link or its parents are forbidden (already used by validated track)
while (linkID > MinusOne) {
const auto& link = trackLinks[linkID];
const auto& link = getLink(linkID);
if (link.clID > MinusOne) {
clStatus[link.clID] = MinusTen;
}
linkID = link.parentID;
}
}
size_t sizeInternal() const { return sizeof(ABTrackLink) * trackLinks.size(); }
size_t capInternal() const { return sizeof(ABTrackLink) * trackLinks.capacity(); }
size_t sizeInternal() const { return sizeof(ABTrackLink) * getNLinks(); }
size_t capInternal() const { return sizeof(ABTrackLink) * getNLinks(); }
};

struct InteractionCandidate : public o2::InteractionRecord {
Expand Down Expand Up @@ -528,7 +544,7 @@ class MatchTPCITS

// ========================= AFTERBURNER =========================
int prepareABSeeds();
void processABSeed(int sid, const ITSChipClustersRefs& itsChipClRefs);
void processABSeed(int sid, const ITSChipClustersRefs& itsChipClRefs, uint8_t tID);
int followABSeed(const o2::track::TrackParCov& seed, const ITSChipClustersRefs& itsChipClRefs, int seedID, int lrID, TPCABSeed& ABSeed);
int registerABTrackLink(TPCABSeed& ABSeed, const o2::track::TrackParCov& trc, int clID, int parentID, int lr, int laddID, float chi2Cl);
bool isBetter(float chi2A, float chi2B) { return chi2A < chi2B; } // RS FIMXE TODO
Expand Down Expand Up @@ -651,6 +667,7 @@ class MatchTPCITS
std::vector<int> mTPCABIndexCache;
std::vector<int> mABWinnersIDs;
std::vector<int> mABClusterLinkIndex; ///< index of 1st ABClusterLink for every cluster used by AfterBurner, -1: unused, -10: used by external ITS tracks
LinksPoolMT mABLinksPool;

///< per sector indices of TPC track entry in mTPCWork
std::array<std::vector<int>, o2::constants::math::NSectors> mTPCSectIndexCache;
Expand Down
51 changes: 30 additions & 21 deletions Detectors/GlobalTracking/src/MatchTPCITS.cxx
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,6 @@
#include "ITStracking/IOUtils.h"

#include "GPUO2Interface.h" // Needed for propper settings in GPUParam.h

#ifdef WITH_OPENMP
#include <omp.h>
#endif
Expand All @@ -58,6 +57,8 @@ using NAMES = o2::base::NameConf;
using GTrackID = o2::dataformats::GlobalTrackID;
constexpr float MatchTPCITS::Tan70, MatchTPCITS::Cos70I2, MatchTPCITS::MaxSnp, MatchTPCITS::MaxTgp;

LinksPoolMT* TPCABSeed::gLinksPool = nullptr;

//______________________________________________
MatchTPCITS::MatchTPCITS() = default;

Expand Down Expand Up @@ -173,6 +174,9 @@ void MatchTPCITS::clear()
mTPCLblWork.clear();
mITSLblWork.clear();
}
for (int i = 0; i < mNThreads; i++) {
mABLinksPool.threadPool[i].clear();
}
}

//______________________________________________
Expand Down Expand Up @@ -1759,13 +1763,13 @@ bool MatchTPCITS::runAfterBurner(pmr::vector<o2::dataformats::TrackTPCITS>& matc
continue;
}
#ifdef WITH_OPENMP
int tid = omp_get_thread_num();
uint8_t tid = (uint8_t)omp_get_thread_num();
#else
int tid = 0;
uint8_t tid = 0;
#endif
fillClustersForAfterBurner(intCand.rofITS, 1, itsChipClRefsBuff[tid]); // RS FIXME account for possibility of filling 2 ROFs
for (int is = intCand.seedsRef.getFirstEntry(); is < intCand.seedsRef.getEntriesBound(); is++) { // loop over all seeds of this interaction candidate
processABSeed(is, itsChipClRefsBuff[tid]);
processABSeed(is, itsChipClRefsBuff[tid], tid);
}
}
mTimer[SWABMatch].Stop();
Expand All @@ -1785,16 +1789,7 @@ bool MatchTPCITS::runAfterBurner(pmr::vector<o2::dataformats::TrackTPCITS>& matc
if (ABSeed.isDisabled()) {
continue;
}
if (ABSeed.lowestLayer > mParams->requireToReachLayerAB) {
ABSeed.disable();
continue;
}
auto candID = ABSeed.getBestLinkID();
if (candID < 0 || ABSeed.getLink(candID).nContLayers < mParams->minContributingLayersAB) {
ABSeed.disable();
continue;
}
candAB.emplace_back(SID{i, ABSeed.getLink(candID).chi2Norm()});
candAB.emplace_back(SID{i, ABSeed.getLink(ABSeed.getBestLinkID()).chi2Norm()});
}
std::sort(candAB.begin(), candAB.end(), [](SID a, SID b) { return a.chi2 < b.chi2; });
for (int i = 0; i < (int)candAB.size(); i++) {
Expand Down Expand Up @@ -1907,10 +1902,12 @@ void MatchTPCITS::refitABWinners(pmr::vector<o2::dataformats::TrackTPCITS>& matc
}

//______________________________________________
void MatchTPCITS::processABSeed(int sid, const ITSChipClustersRefs& itsChipClRefs)
void MatchTPCITS::processABSeed(int sid, const ITSChipClustersRefs& itsChipClRefs, uint8_t tID)
{
// prepare matching hypothesis tree for given seed
auto& ABSeed = mTPCABSeeds[sid];
ABSeed.threadID = tID;
ABSeed.linksEntry = mABLinksPool.threadPool[tID].size();
followABSeed(ABSeed.track, itsChipClRefs, MinusTen, NITSLayers - 1, ABSeed); // check matches on outermost layer
for (int ilr = NITSLayers - 1; ilr > mParams->lowestLayerAB; ilr--) {
int nextLinkID = ABSeed.firstInLr[ilr];
Expand All @@ -1928,6 +1925,16 @@ void MatchTPCITS::processABSeed(int sid, const ITSChipClustersRefs& itsChipClRef
nextLinkID = next2nextLinkID;
}
}
// is this seed has chance to be validated?
auto candID = ABSeed.getBestLinkID();
if (ABSeed.isDisabled() ||
ABSeed.lowestLayer > mParams->requireToReachLayerAB ||
candID < 0 ||
ABSeed.getLink(candID).nContLayers < mParams->minContributingLayersAB) { // free unused links
ABSeed.disable();
mABLinksPool.threadPool[tID].resize(size_t(ABSeed.linksEntry));
}

/* // RS FIXME remove on final clean-up
auto bestLinkID = ABSeed.getBestLinkID();
if (bestLinkID>MinusOne) {
Expand Down Expand Up @@ -2132,13 +2139,13 @@ int MatchTPCITS::registerABTrackLink(TPCABSeed& ABSeed, const o2::track::TrackPa
{
// registers new ABLink on the layer, assigning provided kinematics. The link will be registered in a
// way preserving the quality ordering of the links on the layer
int lnkID = ABSeed.trackLinks.size(), nextID = ABSeed.firstInLr[lr], nc = 1 + (parentID > MinusOne ? ABSeed.getLink(parentID).nContLayers : 0);
int lnkID = ABSeed.getNLinks(), nextID = ABSeed.firstInLr[lr], nc = 1 + (parentID > MinusOne ? ABSeed.getLink(parentID).nContLayers : 0);
float chi2 = chi2Cl + (parentID > MinusOne ? ABSeed.getLink(parentID).chi2 : 0.);
// LOG(info) << "Reg on lr " << lr << " nc = " << nc << " chi2cl=" << chi2Cl << " -> " << chi2; // RSTMP

if (ABSeed.firstInLr[lr] == MinusOne) { // no links on this layer yet
ABSeed.firstInLr[lr] = lnkID;
ABSeed.trackLinks.emplace_back(trc, clID, parentID, MinusOne, lr, nc, laddID, chi2);
ABSeed.addLink(trc, clID, parentID, MinusOne, lr, nc, laddID, chi2);
return lnkID;
}
// add new link sorting links of this layer in quality
Expand All @@ -2150,7 +2157,7 @@ int MatchTPCITS::registerABTrackLink(TPCABSeed& ABSeed, const o2::track::TrackPa
bool newIsBetter = parentID <= MinusOne ? isBetter(chi2, nextLink.chi2) : isBetter(ABSeed.getLink(parentID).chi2NormPredict(chi2Cl), nextLink.chi2Norm());
if (newIsBetter) { // need to insert new link before nextLink
if (count < mParams->maxABLinksOnLayer) { // will insert in front of nextID
ABSeed.trackLinks.emplace_back(trc, clID, parentID, nextID, lr, nc, laddID, chi2);
ABSeed.addLink(trc, clID, parentID, nextID, lr, nc, laddID, chi2);
if (topID == MinusOne) { // are we comparing new link with best link on the layer?
ABSeed.firstInLr[lr] = lnkID; // flag as best on the layer
} else {
Expand All @@ -2167,7 +2174,7 @@ int MatchTPCITS::registerABTrackLink(TPCABSeed& ABSeed, const o2::track::TrackPa
} while (nextID > MinusOne);
// new link is worse than all others, add it only if there is a room to expand
if (count < mParams->maxABLinksOnLayer) {
ABSeed.trackLinks.emplace_back(trc, clID, parentID, MinusOne, lr, nc, laddID, chi2);
ABSeed.addLink(trc, clID, parentID, MinusOne, lr, nc, laddID, chi2);
if (topID > MinusOne) {
ABSeed.getLink(topID).nextOnLr = lnkID; // point from previous one
}
Expand Down Expand Up @@ -2540,8 +2547,8 @@ void MatchTPCITS::reportSizes(pmr::vector<o2::dataformats::TrackTPCITS>& matched
for (const auto& a : mTPCABSeeds) {
siz += a.sizeInternal();
cap += a.capInternal();
cnt += a.trackLinks.size();
cntCap += a.trackLinks.capacity();
cnt += a.getNLinks();
cntCap += a.getNLinks();
}
sizTot += siz;
capTot += cap;
Expand Down Expand Up @@ -2615,6 +2622,8 @@ void MatchTPCITS::setNThreads(int n)
LOG(warning) << "Multithreading is not supported, imposing single thread";
mNThreads = 1;
#endif
mABLinksPool.threadPool.resize(mNThreads);
TPCABSeed::gLinksPool = &mABLinksPool;
}

//<<============================= AfterBurner for TPC-track / ITS cluster matching ===================<<
Expand Down

0 comments on commit 81a003d

Please sign in to comment.