Skip to content

Commit 7b34c90

Browse files
ariedel-cernariedel-cern
authored
[PWGCF] Speed up triplet mixing in femto framework (#16965)
1 parent 710ee27 commit 7b34c90

1 file changed

Lines changed: 45 additions & 18 deletions

File tree

‎PWGCF/Femto/Core/tripletProcessHelpers.h‎

Lines changed: 45 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
#include <Framework/Logger.h>
2424

2525
#include <cstdint>
26+
#include <optional>
2627

2728
namespace o2::analysis::femto::tripletprocesshelpers
2829
{
@@ -459,20 +460,34 @@ void processMixedEvent(T1 const& Collisions,
459460
T10& CtrManager,
460461
T11& TcManager)
461462
{
462-
int64_t lastCollisionIndex = -1;
463+
int64_t lastCollisionIndex1 = -1;
464+
int64_t lastCollisionIndex2 = -1;
463465
int windowSizeRaw = 0;
464466
int windowSizeEffective = 0;
465467

468+
// collision1 stays fixed across the outer mixing window, and collision2 stays fixed across each inner sub-window (it only advances once collision3 wraps
469+
// back to its start Both slices are therefore materialized once per window/sub-window and reused, instead of being re-sliced, i.e. a fresh arrow Slice + selection copy
470+
std::optional<decltype(Partition1->sliceByCached(o2::aod::femtobase::stored::fColId, 0, cache))> sliceParticle1;
471+
std::optional<decltype(Partition2->sliceByCached(o2::aod::femtobase::stored::fColId, 0, cache))> sliceParticle2;
472+
466473
for (auto const& [collision1, collision2, collision3] : o2::soa::selfCombinations(policy, depth, -1, Collisions, Collisions, Collisions)) {
467474

468-
// --- new window ---
469-
if (collision1.globalIndex() != lastCollisionIndex) {
470-
if (lastCollisionIndex != -1) {
475+
// outer window
476+
if (collision1.globalIndex() != lastCollisionIndex1) {
477+
if (lastCollisionIndex1 != -1) {
471478
TripletHistManager.fillMixingQaMePerMixingBin(windowSizeRaw, windowSizeEffective);
472479
}
473480
windowSizeRaw = 0;
474481
windowSizeEffective = 0;
475-
lastCollisionIndex = collision1.globalIndex();
482+
lastCollisionIndex1 = collision1.globalIndex();
483+
lastCollisionIndex2 = -1; // force sliceParticle2 to refresh below
484+
sliceParticle1.emplace(Partition1->sliceByCached(o2::aod::femtobase::stored::fColId, collision1.globalIndex(), cache));
485+
}
486+
487+
// inner sub-window
488+
if (collision2.globalIndex() != lastCollisionIndex2) {
489+
lastCollisionIndex2 = collision2.globalIndex();
490+
sliceParticle2.emplace(Partition2->sliceByCached(o2::aod::femtobase::stored::fColId, collision2.globalIndex(), cache));
476491
}
477492

478493
++windowSizeRaw;
@@ -486,21 +501,19 @@ void processMixedEvent(T1 const& Collisions,
486501

487502
CtrManager.setMagField(collision1.magField());
488503

489-
auto sliceParticle1 = Partition1->sliceByCached(o2::aod::femtobase::stored::fColId, collision1.globalIndex(), cache);
490-
auto sliceParticle2 = Partition2->sliceByCached(o2::aod::femtobase::stored::fColId, collision2.globalIndex(), cache);
491504
auto sliceParticle3 = Partition3->sliceByCached(o2::aod::femtobase::stored::fColId, collision3.globalIndex(), cache);
492505

493506
TripletHistManager.resetTrackedParticlesPerEvent();
494507

495-
if (sliceParticle1.size() == 0 || sliceParticle2.size() == 0 || sliceParticle3.size() == 0) {
508+
if (sliceParticle1->size() == 0 || sliceParticle2->size() == 0 || sliceParticle3.size() == 0) {
496509
TripletHistManager.fillMixingQaMePerEvent();
497510
continue;
498511
}
499512

500513
bool hasValidTriplet = false;
501514
TripletHistManager.fillMixingQaMe(collision1, collision2, collision3);
502515

503-
for (auto const& [p1, p2, p3] : o2::soa::combinations(o2::soa::CombinationsFullIndexPolicy(sliceParticle1, sliceParticle2, sliceParticle3))) {
516+
for (auto const& [p1, p2, p3] : o2::soa::combinations(o2::soa::CombinationsFullIndexPolicy(*sliceParticle1, *sliceParticle2, sliceParticle3))) {
504517
// pair cleaning
505518
if (!TcManager.isCleanTriplet(p1, p2, p3, TrackTable)) {
506519
continue;
@@ -528,7 +541,7 @@ void processMixedEvent(T1 const& Collisions,
528541
TripletHistManager.fillMixingQaMePerEvent();
529542
}
530543

531-
// --- final window ---
544+
// final window
532545
if (windowSizeRaw > 0) {
533546
TripletHistManager.fillMixingQaMePerMixingBin(windowSizeRaw, windowSizeEffective);
534547
}
@@ -563,19 +576,34 @@ void processMixedEvent(T1 const& Collisions,
563576
T12& CtrManager,
564577
T13& TcManager)
565578
{
566-
int64_t lastCollisionIndex = -1;
579+
int64_t lastCollisionIndex1 = -1;
580+
int64_t lastCollisionIndex2 = -1;
567581
int windowSizeRaw = 0;
568582
int windowSizeEffective = 0;
569583

584+
// collision1 stays fixed across the outer mixing window, and collision2 stays fixed across each inner sub-window (it only advances once collision3 wraps
585+
// back to its start Both slices are therefore materialized once per window/sub-window and reused, instead of being re-sliced, i.e. a fresh arrow Slice + selection copy
586+
std::optional<decltype(Partition1->sliceByCached(o2::aod::femtobase::stored::fColId, 0, cache))> sliceParticle1;
587+
std::optional<decltype(Partition2->sliceByCached(o2::aod::femtobase::stored::fColId, 0, cache))> sliceParticle2;
588+
570589
for (auto const& [collision1, collision2, collision3] : o2::soa::selfCombinations(policy, depth, -1, Collisions, Collisions, Collisions)) {
571590

572-
if (collision1.globalIndex() != lastCollisionIndex) {
573-
if (lastCollisionIndex != -1) {
591+
// outer window
592+
if (collision1.globalIndex() != lastCollisionIndex1) {
593+
if (lastCollisionIndex1 != -1) {
574594
TripletHistManager.fillMixingQaMePerMixingBin(windowSizeRaw, windowSizeEffective);
575595
}
576596
windowSizeRaw = 0;
577597
windowSizeEffective = 0;
578-
lastCollisionIndex = collision1.globalIndex();
598+
lastCollisionIndex1 = collision1.globalIndex();
599+
lastCollisionIndex2 = -1; // force sliceParticle2 to refresh below
600+
sliceParticle1.emplace(Partition1->sliceByCached(o2::aod::femtobase::stored::fColId, collision1.globalIndex(), cache));
601+
}
602+
603+
// inner sub-window
604+
if (collision2.globalIndex() != lastCollisionIndex2) {
605+
lastCollisionIndex2 = collision2.globalIndex();
606+
sliceParticle2.emplace(Partition2->sliceByCached(o2::aod::femtobase::stored::fColId, collision2.globalIndex(), cache));
579607
}
580608

581609
++windowSizeRaw;
@@ -589,21 +617,19 @@ void processMixedEvent(T1 const& Collisions,
589617

590618
CtrManager.setMagField(collision1.magField());
591619

592-
auto sliceParticle1 = Partition1->sliceByCached(o2::aod::femtobase::stored::fColId, collision1.globalIndex(), cache);
593-
auto sliceParticle2 = Partition2->sliceByCached(o2::aod::femtobase::stored::fColId, collision2.globalIndex(), cache);
594620
auto sliceParticle3 = Partition3->sliceByCached(o2::aod::femtobase::stored::fColId, collision3.globalIndex(), cache);
595621

596622
TripletHistManager.resetTrackedParticlesPerEvent();
597623

598-
if (sliceParticle1.size() == 0 || sliceParticle2.size() == 0 || sliceParticle3.size() == 0) {
624+
if (sliceParticle1->size() == 0 || sliceParticle2->size() == 0 || sliceParticle3.size() == 0) {
599625
TripletHistManager.fillMixingQaMePerEvent();
600626
continue;
601627
}
602628

603629
bool hasValidTriplet = false;
604630
TripletHistManager.fillMixingQaMe(collision1, collision2, collision3);
605631

606-
for (auto const& [p1, p2, p3] : o2::soa::combinations(o2::soa::CombinationsFullIndexPolicy(sliceParticle1, sliceParticle2, sliceParticle3))) {
632+
for (auto const& [p1, p2, p3] : o2::soa::combinations(o2::soa::CombinationsFullIndexPolicy(*sliceParticle1, *sliceParticle2, sliceParticle3))) {
607633
// pair cleaning
608634
if (!TcManager.isCleanTriplet(p1, p2, p3, TrackTable)) {
609635
continue;
@@ -631,6 +657,7 @@ void processMixedEvent(T1 const& Collisions,
631657
TripletHistManager.fillMixingQaMePerEvent();
632658
}
633659

660+
// --- final window ---
634661
if (windowSizeRaw > 0) {
635662
TripletHistManager.fillMixingQaMePerMixingBin(windowSizeRaw, windowSizeEffective);
636663
}

0 commit comments

Comments
 (0)