Skip to content

Track Embeddings for Improved Duplicate Removal in LST #48249

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cmsbuild merged 1 commit into from
Jun 5, 2025
Merged

Track Embeddings for Improved Duplicate Removal in LST #48249

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions RecoTracker/LSTCore/interface/Common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,9 @@ namespace lst {
};
struct Params_pLS {
static constexpr int kLayers = 2, kHits = 4;
static constexpr int kEmbed = 6;
using ArrayUxHits = edm::StdArray<unsigned int, kHits>;
using ArrayFxEmbed = edm::StdArray<float, kEmbed>;
};
struct Params_LS {
static constexpr int kLayers = 2, kHits = 4;
Expand All @@ -84,9 +86,11 @@ namespace lst {
};
struct Params_T5 {
static constexpr int kLayers = 5, kHits = 10;
static constexpr int kEmbed = 6;
using ArrayU8xLayers = edm::StdArray<uint8_t, kLayers>;
using ArrayU16xLayers = edm::StdArray<uint16_t, kLayers>;
using ArrayUxHits = edm::StdArray<unsigned int, kHits>;
using ArrayFxEmbed = edm::StdArray<float, kEmbed>;
};
struct Params_pT5 {
static constexpr int kLayers = 7, kHits = 14;
Expand Down
1 change: 1 addition & 0 deletions RecoTracker/LSTCore/interface/PixelSegmentsSoA.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ namespace lst {

GENERATE_SOA_LAYOUT(PixelSegmentsSoALayout,
SOA_COLUMN(uint4, pLSHitsIdxs),
SOA_COLUMN(Params_pLS::ArrayFxEmbed, plsEmbed),
SOA_COLUMN(char, isDup),
SOA_COLUMN(bool, partOfPT5),
SOA_COLUMN(float, score),
Expand Down
1 change: 1 addition & 0 deletions RecoTracker/LSTCore/interface/QuintupletsSoA.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ namespace lst {
SOA_COLUMN(Params_T5::ArrayU16xLayers, lowerModuleIndices), // lower module index in each layer
SOA_COLUMN(Params_T5::ArrayU8xLayers, logicalLayers), // layer ID
SOA_COLUMN(Params_T5::ArrayUxHits, hitIndices), // hit indices
SOA_COLUMN(Params_T5::ArrayFxEmbed, t5Embed), // t5 embedding vector
SOA_COLUMN(FPX, innerRadius), // inner triplet circle radius
SOA_COLUMN(FPX, bridgeRadius), // "middle"/bridge triplet radius
SOA_COLUMN(FPX, outerRadius), // outer triplet radius
Expand Down
11 changes: 7 additions & 4 deletions RecoTracker/LSTCore/interface/TripletsSoA.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,13 @@ namespace lst {
SOA_COLUMN(Params_T3::ArrayU16xLayers, lowerModuleIndices), // lower module index in each layer
SOA_COLUMN(Params_T3::ArrayU8xLayers, logicalLayers), // layer ID
SOA_COLUMN(Params_T3::ArrayUxHits, hitIndices), // hit indices
SOA_COLUMN(FPX, betaIn), // beta/chord angle of the inner segment
SOA_COLUMN(float, centerX), // lower/anchor-hit based circle center x
SOA_COLUMN(float, centerY), // lower/anchor-hit based circle center y
SOA_COLUMN(float, radius), // lower/anchor-hit based circle radius
SOA_COLUMN(FPX, betaIn), // beta/chord angle of the inner segment
SOA_COLUMN(float, centerX), // lower/anchor-hit based circle center x
SOA_COLUMN(float, centerY), // lower/anchor-hit based circle center y
SOA_COLUMN(float, radius), // lower/anchor-hit based circle radius
SOA_COLUMN(float, fakeScore), // DNN confidence score for fake t3
SOA_COLUMN(float, promptScore), // DNN confidence score for real (prompt) t3
SOA_COLUMN(float, displacedScore), // DNN confidence score for real (displaced) t3
#ifdef CUT_VALUE_DEBUG
SOA_COLUMN(float, zOut),
SOA_COLUMN(float, rtOut),
Expand Down
26 changes: 18 additions & 8 deletions RecoTracker/LSTCore/interface/alpaka/Common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,34 +58,44 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
namespace dnn {

// Common constants for both DNNs
HOST_DEVICE_CONSTANT float kEta_norm = 2.5f;
HOST_DEVICE_CONSTANT float kPhi_norm = kPi;
HOST_DEVICE_CONSTANT float kEtaSize = 0.25f; // Bin size in eta.
constexpr unsigned int kPtBins = 2;
constexpr unsigned int kEtaBins = 10;

namespace plsembdnn {
HOST_DEVICE_CONSTANT float kEta_norm = 4.0f;
HOST_DEVICE_CONSTANT float kEtaErr_norm = 0.00139f;
HOST_DEVICE_CONSTANT float kWP[kEtaBins] = {
0.9235f, 0.8974f, 0.9061f, 0.9431f, 0.8262f, 0.7998f, 0.7714f, 0.7017f, 0.6749f, 0.6624f};
} // namespace plsembdnn

namespace t3dnn {
HOST_DEVICE_CONSTANT float kEta_norm = 2.5f;
HOST_DEVICE_CONSTANT float kZ_max = 224.149505f;
HOST_DEVICE_CONSTANT float kR_max = 98.932365f;
HOST_DEVICE_CONSTANT unsigned int kOutputFeatures = 3;
HOST_DEVICE_CONSTANT float kWp_prompt[kPtBins][kEtaBins] = {
{0.4957, 0.5052, 0.5201, 0.5340, 0.4275, 0.4708, 0.4890, 0.4932, 0.5400, 0.5449},
{0.0302, 0.0415, 0.0994, 0.1791, 0.1960, 0.2467, 0.3227, 0.3242, 0.2367, 0.2187}};
{0.4957f, 0.5052f, 0.5201f, 0.5340f, 0.4275f, 0.4708f, 0.4890f, 0.4932f, 0.5400f, 0.5449f},
{0.0302f, 0.0415f, 0.0994f, 0.1791f, 0.1960f, 0.2467f, 0.3227f, 0.3242f, 0.2367f, 0.2187f}};
HOST_DEVICE_CONSTANT float kWp_displaced[kPtBins][kEtaBins] = {
{0.0334, 0.0504, 0.0748, 0.0994, 0.1128, 0.1123, 0.1118, 0.1525, 0.1867, 0.1847},
{0.0091, 0.0075, 0.0350, 0.0213, 0.0435, 0.0676, 0.1957, 0.1649, 0.1080, 0.1046}};
{0.0334f, 0.0504f, 0.0748f, 0.0994f, 0.1128f, 0.1123f, 0.1118f, 0.1525f, 0.1867f, 0.1847f},
{0.0091f, 0.0075f, 0.0350f, 0.0213f, 0.0435f, 0.0676f, 0.1957f, 0.1649f, 0.1080f, 0.1046f}};
} // namespace t3dnn

namespace t5dnn {
HOST_DEVICE_CONSTANT float kEta_norm = 2.5f;
HOST_DEVICE_CONSTANT float kZ_max = 267.2349854f;
HOST_DEVICE_CONSTANT float kR_max = 110.1099396f;
HOST_DEVICE_CONSTANT float kWp[kPtBins][kEtaBins] = {
{0.4493, 0.4939, 0.5715, 0.6488, 0.5709, 0.5938, 0.7164, 0.7565, 0.8103, 0.8593},
{0.4488, 0.4448, 0.5067, 0.5929, 0.4836, 0.4112, 0.4968, 0.4403, 0.5597, 0.5067}};
{0.4493f, 0.4939f, 0.5715f, 0.6488f, 0.5709f, 0.5938f, 0.7164f, 0.7565f, 0.8103f, 0.8593f},
{0.4488f, 0.4448f, 0.5067f, 0.5929f, 0.4836f, 0.4112f, 0.4968f, 0.4403f, 0.5597f, 0.5067f}};
} // namespace t5dnn

namespace pt3dnn {
HOST_DEVICE_CONSTANT float kEta_norm = 2.5f;
HOST_DEVICE_CONSTANT float kWp[kEtaBins] = {
0.189, 0.1805, 0.2267, 0.3104, 0.4719, 0.3159, 0.1372, 0.1571, 0.3198, 0.186};
0.189f, 0.1805f, 0.2267f, 0.3104f, 0.4719f, 0.3159f, 0.1372f, 0.1571f, 0.3198f, 0.186f};
HOST_DEVICE_CONSTANT float kWpHigh = 0.0473f;
} // namespace pt3dnn

Expand Down
11 changes: 10 additions & 1 deletion RecoTracker/LSTCore/src/alpaka/Kernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
#define RecoTracker_LSTCore_src_alpaka_Kernels_h

#include "HeterogeneousCore/AlpakaInterface/interface/workdivision.h"
#include "FWCore/Utilities/interface/CMSUnrollLoop.h"

#include "RecoTracker/LSTCore/interface/alpaka/Common.h"
#include "RecoTracker/LSTCore/interface/ModulesSoA.h"
Expand Down Expand Up @@ -249,7 +250,15 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
float dR2 = dEta * dEta + dPhi * dPhi;
int nMatched = checkHitsT5(ix, jx, quintuplets);
const int minNHitsForDup_T5 = 5;
if (dR2 < 0.001f || nMatched >= minNHitsForDup_T5) {

float d2 = 0.f;
CMS_UNROLL_LOOP
for (unsigned int k = 0; k < Params_T5::kEmbed; ++k) {
float diff = quintuplets.t5Embed()[ix][k] - quintuplets.t5Embed()[jx][k];
d2 += diff * diff;
}

if (((dR2 < 0.001f || nMatched >= minNHitsForDup_T5) && d2 < 1.0f) || (dR2 < 0.02f && d2 < 0.1f)) {
if (isPT5_jx || score_rphisum1 > score_rphisum2) {
rmQuintupletFromMemory(quintuplets, ix, true);
} else if (isPT5_ix || score_rphisum1 < score_rphisum2) {
Expand Down
162 changes: 151 additions & 11 deletions RecoTracker/LSTCore/src/alpaka/NeuralNetwork.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@
#include "T5NeuralNetworkWeights.h"
#include "T3NeuralNetworkWeights.h"
#include "pT3NeuralNetworkWeights.h"
#include "T5EmbedNetworkWeights.h"
#include "pLSEmbedNetworkWeights.h"

namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {

Expand Down Expand Up @@ -66,11 +68,11 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
const unsigned int mdIndex2,
const unsigned int mdIndex3,
const float radius,
const float betaIn) {
const float betaIn,
float (&output)[dnn::t3dnn::kOutputFeatures]) {
Copy link
Contributor

@fwyzard fwyzard Jun 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

isn't

Suggested change
float (&output)[dnn::t3dnn::kOutputFeatures]) {
float output[dnn::t3dnn::kOutputFeatures]) {

equivalent ?

Copy link
Contributor Author

@GNiendorf GNiendorf Jun 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see what you mean, but you do get a compile-time check from the reference form that the array passed to output is exactly kOutputFeatures elements.

// Constants for T3 DNN
constexpr unsigned int kInputFeatures = 14;
constexpr unsigned int kHiddenFeatures = 32;
constexpr unsigned int kOutputFeatures = 3;

// Extract hit information
float eta1 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex1]); // inner T3 anchor hit 1 eta
Expand All @@ -91,7 +93,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {

// Build input feature vector matching training order
float x[kInputFeatures] = {
eta1 / dnn::kEta_norm, // First hit eta normalized
eta1 / dnn::t3dnn::kEta_norm, // First hit eta normalized
alpaka::math::abs(acc, phi1) / dnn::kPhi_norm, // First hit phi normalized
z1 / dnn::t3dnn::kZ_max, // First hit z normalized
r1 / dnn::t3dnn::kR_max, // First hit r normalized
Expand All @@ -112,7 +114,6 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {

float x_1[kHiddenFeatures]; // Layer 1 output
float x_2[kHiddenFeatures]; // Layer 2 output
float x_3[kOutputFeatures]; // Layer 3 output (3 classes)

// Layer 1: Linear + Relu
linear_layer<kInputFeatures, kHiddenFeatures>(x, x_1, dnn::t3dnn::wgtT_layer1, dnn::t3dnn::bias_layer1);
Expand All @@ -123,17 +124,17 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
relu_activation<kHiddenFeatures>(x_2);

// Layer 3: Linear + Softmax
linear_layer<kHiddenFeatures, kOutputFeatures>(
x_2, x_3, dnn::t3dnn::wgtT_output_layer, dnn::t3dnn::bias_output_layer);
softmax_activation<kOutputFeatures>(acc, x_3);
linear_layer<kHiddenFeatures, dnn::t3dnn::kOutputFeatures>(
x_2, output, dnn::t3dnn::wgtT_output_layer, dnn::t3dnn::bias_output_layer);
softmax_activation<dnn::t3dnn::kOutputFeatures>(acc, output);

// Get pt and eta bin indices
float t3_pt = radius * lst::k2Rinv1GeVf * 2;
uint8_t pt_index = (t3_pt > 5);
uint8_t bin_index = (eta1 > 2.5f) ? (dnn::kEtaBins - 1) : static_cast<unsigned int>(eta1 / dnn::kEtaSize);

return x_3[1] > dnn::t3dnn::kWp_prompt[pt_index][bin_index] ||
x_3[2] > dnn::t3dnn::kWp_displaced[pt_index][bin_index];
return output[1] > dnn::t3dnn::kWp_prompt[pt_index][bin_index] ||
output[2] > dnn::t3dnn::kWp_displaced[pt_index][bin_index];
}
} // namespace t3dnn

Expand All @@ -157,7 +158,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
alpaka::math::log10(acc, pixelRadius),
alpaka::math::log10(acc, pixRadiusError),
alpaka::math::log10(acc, rzChiSquared),
alpaka::math::abs(acc, pixelEta) / dnn::kEta_norm};
alpaka::math::abs(acc, pixelEta) / dnn::pt3dnn::kEta_norm};

float x1[kHiddenFeatures];
float x2[kHiddenFeatures];
Expand Down Expand Up @@ -227,7 +228,7 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {

// Build the input feature vector using pairwise differences after the first hit
float x[kInputFeatures] = {
eta1 / dnn::kEta_norm, // inner T3: First hit eta normalized
eta1 / dnn::t5dnn::kEta_norm, // inner T3: First hit eta normalized
alpaka::math::abs(acc, phi1) / dnn::kPhi_norm, // inner T3: First hit phi normalized
z1 / dnn::t5dnn::kZ_max, // inner T3: First hit z normalized
r1 / dnn::t5dnn::kR_max, // inner T3: First hit r normalized
Expand Down Expand Up @@ -288,6 +289,145 @@ namespace ALPAKA_ACCELERATOR_NAMESPACE::lst {
}
} // namespace t5dnn

namespace t5embdnn {
template <typename TAcc>
ALPAKA_FN_ACC ALPAKA_FN_INLINE void runEmbed(TAcc const& acc,
MiniDoubletsConst mds,
unsigned int mdIndex1,
unsigned int mdIndex2,
unsigned int mdIndex3,
unsigned int mdIndex4,
unsigned int mdIndex5,
float innerRadius,
float outerRadius,
float bridgeRadius,
float fakeScore1,
float promptScore1,
float dispScore1,
float fakeScore2,
float promptScore2,
float dispScore2,
float (&embedding)[Params_T5::kEmbed]) {
constexpr unsigned int kInputFeatures = 30;
constexpr unsigned int kHiddenFeatures = 32;

float eta1 = mds.anchorEta()[mdIndex1];
float eta2 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex2]);
float eta3 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex3]);
float eta4 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex4]);
float eta5 = alpaka::math::abs(acc, mds.anchorEta()[mdIndex5]);

float phi1 = mds.anchorPhi()[mdIndex1];
float phi2 = mds.anchorPhi()[mdIndex2];
float phi3 = mds.anchorPhi()[mdIndex3];
float phi4 = mds.anchorPhi()[mdIndex4];
float phi5 = mds.anchorPhi()[mdIndex5];

float z1 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex1]);
float z2 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex2]);
float z3 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex3]);
float z4 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex4]);
float z5 = alpaka::math::abs(acc, mds.anchorZ()[mdIndex5]);

float r1 = mds.anchorRt()[mdIndex1];
float r2 = mds.anchorRt()[mdIndex2];
float r3 = mds.anchorRt()[mdIndex3];
float r4 = mds.anchorRt()[mdIndex4];
float r5 = mds.anchorRt()[mdIndex5];

float x[kInputFeatures] = {eta1 / dnn::t5dnn::kEta_norm,
alpaka::math::cos(acc, phi1),
alpaka::math::sin(acc, phi1),
z1 / dnn::t5dnn::kZ_max,
r1 / dnn::t5dnn::kR_max,

eta2 - alpaka::math::abs(acc, eta1),
cms::alpakatools::deltaPhi(acc, phi2, phi1),
(z2 - z1) / dnn::t5dnn::kZ_max,
(r2 - r1) / dnn::t5dnn::kR_max,

eta3 - eta2,
cms::alpakatools::deltaPhi(acc, phi3, phi2),
(z3 - z2) / dnn::t5dnn::kZ_max,
(r3 - r2) / dnn::t5dnn::kR_max,

eta4 - eta3,
cms::alpakatools::deltaPhi(acc, phi4, phi3),
(z4 - z3) / dnn::t5dnn::kZ_max,
(r4 - r3) / dnn::t5dnn::kR_max,

eta5 - eta4,
cms::alpakatools::deltaPhi(acc, phi5, phi4),
(z5 - z4) / dnn::t5dnn::kZ_max,
(r5 - r4) / dnn::t5dnn::kR_max,

1.0f / innerRadius,
1.0f / bridgeRadius,
1.0f / outerRadius,

fakeScore1,
promptScore1,
dispScore1,
(fakeScore2 - fakeScore1),
(promptScore2 - promptScore1),
(dispScore2 - dispScore1)};

float h1[kHiddenFeatures];
float h2[kHiddenFeatures];

linear_layer<kInputFeatures, kHiddenFeatures>(x, h1, dnn::t5embdnn::wgtT_fc1, dnn::t5embdnn::bias_fc1);
relu_activation<kHiddenFeatures>(h1);

linear_layer<kHiddenFeatures, kHiddenFeatures>(h1, h2, dnn::t5embdnn::wgtT_fc2, dnn::t5embdnn::bias_fc2);
relu_activation<kHiddenFeatures>(h2);

linear_layer<kHiddenFeatures, Params_T5::kEmbed>(h2, embedding, dnn::t5embdnn::wgtT_fc3, dnn::t5embdnn::bias_fc3);
}

} // namespace t5embdnn

namespace plsembdnn {
template <typename TAcc>
ALPAKA_FN_ACC ALPAKA_FN_INLINE void runEmbed(TAcc const& acc,
const float eta,
const float etaErr,
const float phi,
const float circleCenterX,
const float circleCenterY,
const float circleRadius,
const float ptIn,
const float ptErr,
const bool isQuad,
float (&embedding)[Params_pLS::kEmbed]) {
constexpr unsigned int kInputFeatures = 10;
constexpr unsigned int kHiddenFeatures = 32;

float x[kInputFeatures] = {eta / dnn::plsembdnn::kEta_norm,
etaErr / dnn::plsembdnn::kEtaErr_norm,
alpaka::math::cos(acc, phi),
alpaka::math::sin(acc, phi),
1.0f / ptIn,
alpaka::math::log10(acc, ptErr),
isQuad ? 1.0f : 0.0f,
alpaka::math::log10(acc, alpaka::math::abs(acc, circleCenterX)),
alpaka::math::log10(acc, alpaka::math::abs(acc, circleCenterY)),
alpaka::math::log10(acc, circleRadius)};

float h1[kHiddenFeatures];
float h2[kHiddenFeatures];

linear_layer<kInputFeatures, kHiddenFeatures>(x, h1, dnn::plsembdnn::wgtT_fc1, dnn::plsembdnn::bias_fc1);
relu_activation<kHiddenFeatures>(h1);

linear_layer<kHiddenFeatures, kHiddenFeatures>(h1, h2, dnn::plsembdnn::wgtT_fc2, dnn::plsembdnn::bias_fc2);
relu_activation<kHiddenFeatures>(h2);

linear_layer<kHiddenFeatures, Params_pLS::kEmbed>(
h2, embedding, dnn::plsembdnn::wgtT_fc3, dnn::plsembdnn::bias_fc3);
}

} // namespace plsembdnn

} // namespace ALPAKA_ACCELERATOR_NAMESPACE::lst

#endif
Loading