From ea4d277b55fc4788b612eee90c43a814e6ae46bd Mon Sep 17 00:00:00 2001 From: Christian Sonnabend Date: Mon, 6 Jul 2026 19:36:13 +0200 Subject: [PATCH 1/2] Momentum vector usage created by neural network --- .../include/DataFormatsTPC/ClusterNative.h | 31 +++++++++++++++++++ GPU/GPUTracking/Base/GPUConstantMem.h | 5 +++ GPU/GPUTracking/Base/GPUProcessor.h | 6 ++++ GPU/GPUTracking/Definitions/GPUSettingsList.h | 2 +- .../SectorTracker/GPUTPCNeighboursFinder.cxx | 10 +++--- GPU/GPUTracking/SectorTracker/GPUTPCTracker.h | 17 ++++++++++ .../GPUTPCTrackletConstructor.cxx | 11 +++++-- .../TPCClusterFinder/GPUTPCNNClusterizer.h | 1 + .../GPUTPCNNClusterizerHost.cxx | 3 ++ .../GPUTPCNNClusterizerKernels.cxx | 14 +++++++++ GPU/Workflow/src/GPUWorkflowSpec.cxx | 3 +- 11 files changed, 95 insertions(+), 8 deletions(-) diff --git a/DataFormats/Detectors/TPC/include/DataFormatsTPC/ClusterNative.h b/DataFormats/Detectors/TPC/include/DataFormatsTPC/ClusterNative.h index fb81afdf67587..298358d26fcb0 100644 --- a/DataFormats/Detectors/TPC/include/DataFormatsTPC/ClusterNative.h +++ b/DataFormats/Detectors/TPC/include/DataFormatsTPC/ClusterNative.h @@ -66,6 +66,9 @@ struct ClusterNative { static constexpr int scaleSaturatedQtot = 8; static constexpr int maxRegularQtot = 25 * 1024; static constexpr int maxSaturatedQtot = (USHRT_MAX - maxRegularQtot) * scaleSaturatedQtot; + static constexpr float scaleNNDirectionPacked = 2048.f; + static constexpr int16_t invalidNNDirectionPacked = -32768; + static constexpr float maxNNDirection = 32767.f / scaleNNDirectionPacked; uint32_t timeFlagsPacked; //< Contains the time in the lower 24 bits in a packed format, contains the flags in the // upper 8 bits @@ -74,11 +77,19 @@ struct ClusterNative { uint8_t sigmaPadPacked; //< Sigma of the pad in packed format uint16_t qMax; //< QMax of the cluster uint16_t qTotPacked; //< Total charge of the cluster + int16_t nnDydxPacked = invalidNNDirectionPacked; + int16_t nnDzdxPacked = invalidNNDirectionPacked; GPUd() static uint16_t packPad(float pad) { return (uint16_t)(pad * scalePadPacked + 0.5); } GPUd() static uint32_t packTime(float time) { return (uint32_t)(time * scaleTimePacked + 0.5); } GPUd() static float unpackPad(uint16_t pad) { return float(pad) * (1.f / scalePadPacked); } GPUd() static float unpackTime(uint32_t time) { return float(time) * (1.f / scaleTimePacked); } + GPUd() static int16_t packNNDirection(float v) + { + v = v < -maxNNDirection ? -maxNNDirection : (v > maxNNDirection ? maxNNDirection : v); + return static_cast(v * scaleNNDirectionPacked + (v >= 0.f ? 0.5f : -0.5f)); + } + GPUd() static float unpackNNDirection(int16_t v) { return float(v) * (1.f / scaleNNDirectionPacked); } GPUdDefault() ClusterNative() = default; GPUd() ClusterNative(uint32_t time, uint8_t flags, uint16_t pad, uint8_t sigmaTime, uint8_t sigmaPad, uint16_t qmax, uint16_t qtotPacked) : padPacked(pad), sigmaTimePacked(sigmaTime), sigmaPadPacked(sigmaPad), qMax(qmax), qTotPacked(qtotPacked) @@ -86,6 +97,20 @@ struct ClusterNative { setTimePackedFlags(time, flags); } + GPUd() bool hasNNDirection() const { return nnDydxPacked != invalidNNDirectionPacked && nnDzdxPacked != invalidNNDirectionPacked; } + GPUd() float getNNDydx() const { return unpackNNDirection(nnDydxPacked); } + GPUd() float getNNDzdx() const { return unpackNNDirection(nnDzdxPacked); } + GPUd() void setNNDirection(float dydx, float dzdx) + { + nnDydxPacked = packNNDirection(dydx); + nnDzdxPacked = packNNDirection(dzdx); + } + GPUd() void clearNNDirection() + { + nnDydxPacked = invalidNNDirectionPacked; + nnDzdxPacked = invalidNNDirectionPacked; + } + GPUd() uint16_t getQmax() const { return qMax; } GPUd() uint32_t getQtot() const { return isSaturated() ? getSaturatedQtot() : (uint32_t)qTotPacked; } GPUd() uint8_t getFlags() const { return timeFlagsPacked >> 24; } @@ -187,6 +212,10 @@ struct ClusterNative { return (this->qMax < rhs.qMax); } else if (this->qTotPacked != rhs.qTotPacked) { return (this->getQtot() < rhs.getQtot()); + } else if (this->nnDydxPacked != rhs.nnDydxPacked) { + return (this->nnDydxPacked < rhs.nnDydxPacked); + } else if (this->nnDzdxPacked != rhs.nnDzdxPacked) { + return (this->nnDzdxPacked < rhs.nnDzdxPacked); } else { return (this->getFlags() < rhs.getFlags()); } @@ -200,6 +229,8 @@ struct ClusterNative { this->sigmaPadPacked == rhs.sigmaPadPacked && this->qMax == rhs.qMax && this->qTotPacked == rhs.qTotPacked && + this->nnDydxPacked == rhs.nnDydxPacked && + this->nnDzdxPacked == rhs.nnDzdxPacked && this->getFlags() == rhs.getFlags(); } diff --git a/GPU/GPUTracking/Base/GPUConstantMem.h b/GPU/GPUTracking/Base/GPUConstantMem.h index 14c388e450d73..8e8c9d9be46c8 100644 --- a/GPU/GPUTracking/Base/GPUConstantMem.h +++ b/GPU/GPUTracking/Base/GPUConstantMem.h @@ -111,6 +111,11 @@ GPUdi() GPUconstantref() const GPUParam& GPUProcessor::Param() const return GetConstantMem()->param; } +GPUdi() const o2::tpc::ClusterNativeAccess* GPUProcessor::GetClustersNative() const +{ + return GetConstantMem()->ioPtrs.clustersNative; +} + GPUdi() void GPUProcessor::raiseError(uint32_t code, uint32_t param1, uint32_t param2, uint32_t param3) const { GetConstantMem()->errorCodes.raiseError(code, param1, param2, param3); diff --git a/GPU/GPUTracking/Base/GPUProcessor.h b/GPU/GPUTracking/Base/GPUProcessor.h index 337ecfc61f79d..512d3e1d03807 100644 --- a/GPU/GPUTracking/Base/GPUProcessor.h +++ b/GPU/GPUTracking/Base/GPUProcessor.h @@ -23,6 +23,11 @@ #include #endif +namespace o2::tpc +{ +struct ClusterNativeAccess; +} + namespace o2::gpu { struct GPUTrackingInOutPointers; @@ -49,6 +54,7 @@ class GPUProcessor #endif GPUd() GPUconstantref() const GPUConstantMem* GetConstantMem() const; // Body in GPUConstantMem.h to avoid circular headers + GPUd() const o2::tpc::ClusterNativeAccess* GetClustersNative() const; // ... GPUd() GPUconstantref() const GPUParam& Param() const; // ... GPUd() void raiseError(uint32_t code, uint32_t param1 = 0, uint32_t param2 = 0, uint32_t param3 = 0) const; const GPUReconstruction& GetRec() const { return *mRec; } diff --git a/GPU/GPUTracking/Definitions/GPUSettingsList.h b/GPU/GPUTracking/Definitions/GPUSettingsList.h index 40f7a34e699c9..7a2a77bd1400c 100644 --- a/GPU/GPUTracking/Definitions/GPUSettingsList.h +++ b/GPU/GPUTracking/Definitions/GPUSettingsList.h @@ -295,7 +295,7 @@ AddOption(nnLoadFromCCDB, int, 0, "", 0, "If 1 networks are fetched from ccdb, e AddOption(nnCCDBDumpToFile, int, 0, "", 0, "If 1, additionally dump fetched CCDB networks to nnLocalFolder") AddOption(nnLocalFolder, std::string, ".", "", 0, "Local folder in which the networks will be fetched") AddOption(nnCCDBPath, std::string, "Users/c/csonnabe/TPC/Clusterization", "", 0, "Folder path containing the networks") -AddOption(nnCCDBWithMomentum, std::string, "", "", 0, "Distinguishes between the network with and without momentum output for the regression") +AddOption(nnCCDBWithMomentum, std::string, "0", "", 0, "Distinguishes between the network with and without momentum output for the regression") AddOption(nnCCDBClassificationLayerType, std::string, "FC", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN") AddOption(nnCCDBRegressionLayerType, std::string, "FC", "", 0, "Distinguishes between network with different layer types. Options: FC, CNN") AddOption(nnCCDBBeamType, std::string, "pp", "", 0, "Distinguishes between networks trained for different beam types. Options: pp, pPb, PbPb") diff --git a/GPU/GPUTracking/SectorTracker/GPUTPCNeighboursFinder.cxx b/GPU/GPUTracking/SectorTracker/GPUTPCNeighboursFinder.cxx index 5151427377b05..96f0caf9c3be6 100644 --- a/GPU/GPUTracking/SectorTracker/GPUTPCNeighboursFinder.cxx +++ b/GPU/GPUTracking/SectorTracker/GPUTPCNeighboursFinder.cxx @@ -120,6 +120,8 @@ GPUdii() void GPUTPCNeighboursFinder::Thread<0>(int32_t /*nBlocks*/, int32_t nTh const GPUglobalref() cahit2& hitData = pHitData[lHitNumberOffset + ih]; const float y = y0 + hitData.x * stepY; const float z = z0 + hitData.y * stepZ; + float nnDydx = 0.f, nnDzdx = 0.f; + const bool useNNDir = tracker.HitNNDirection(row, ih, nnDydx, nnDzdx) && CAMath::Abs(nnDydx) < 10.f && CAMath::Abs(nnDzdx) < 10.f; uint32_t nNeighUp = 0; float minZ, maxZ, minY, maxY; @@ -127,8 +129,8 @@ GPUdii() void GPUTPCNeighboursFinder::Thread<0>(int32_t /*nBlocks*/, int32_t nTh int32_t nY; { // area in the upper row - const float yy = y * s.mUpTx; - const float zz = z * kAreaSlopeZUp; + const float yy = useNNDir ? y + nnDydx * s.mUpDx : y * s.mUpTx; + const float zz = useNNDir ? z + nnDzdx * s.mUpDx : z * kAreaSlopeZUp; minZ = zz - kAreaSizeZUp; maxZ = zz + kAreaSizeZUp; minY = yy - kAreaSizeY; @@ -196,8 +198,8 @@ GPUdii() void GPUTPCNeighboursFinder::Thread<0>(int32_t /*nBlocks*/, int32_t nTh } { // area in the lower row - const float yy = y * s.mDnTx; - const float zz = z * kAreaSlopeZDn; + const float yy = useNNDir ? y + nnDydx * s.mDnDx : y * s.mDnTx; + const float zz = useNNDir ? z + nnDzdx * s.mDnDx : z * kAreaSlopeZDn; minZ = zz - kAreaSizeZDn; maxZ = zz + kAreaSizeZDn; minY = yy - kAreaSizeY; diff --git a/GPU/GPUTracking/SectorTracker/GPUTPCTracker.h b/GPU/GPUTracking/SectorTracker/GPUTPCTracker.h index cb0c8d4a76fdb..d18131fe3ed47 100644 --- a/GPU/GPUTracking/SectorTracker/GPUTPCTracker.h +++ b/GPU/GPUTracking/SectorTracker/GPUTPCTracker.h @@ -26,6 +26,7 @@ #include "GPUTPCTrackParam.h" #include "GPUTPCTracklet.h" #include "GPUProcessor.h" +#include "DataFormatsTPC/ClusterNative.h" namespace o2::gpu { @@ -151,6 +152,22 @@ class GPUTPCTracker : public GPUProcessor GPUhd() int32_t HitInputID(const GPUTPCRow& row, int32_t hitIndex) const { return mData.ClusterDataIndex(row, hitIndex); } + GPUd() bool HitNNDirection(const GPUTPCRow& row, int32_t hitIndex, float& dydx, float& dzdx) const + { + const o2::tpc::ClusterNativeAccess* native = GetClustersNative(); + if (native == nullptr || native->clustersLinear == nullptr) { + return false; + } + const uint32_t sectorOffset = native->clusterOffset[mISector][0]; + const o2::tpc::ClusterNative& cluster = native->clustersLinear[sectorOffset + mData.ClusterDataIndex(row, hitIndex)]; + if (!cluster.hasNNDirection()) { + return false; + } + dydx = cluster.getNNDydx(); + dzdx = cluster.getNNDzdx(); + return true; + } + /** * The hit weight is used to determine whether a hit belongs to a certain tracklet or another one * competing for the same hit. The tracklet that has a higher weight wins. Comparison is done diff --git a/GPU/GPUTracking/SectorTracker/GPUTPCTrackletConstructor.cxx b/GPU/GPUTracking/SectorTracker/GPUTPCTrackletConstructor.cxx index 33a3264a87ab3..8fd823155e129 100644 --- a/GPU/GPUTracking/SectorTracker/GPUTPCTrackletConstructor.cxx +++ b/GPU/GPUTracking/SectorTracker/GPUTPCTrackletConstructor.cxx @@ -170,9 +170,16 @@ GPUdic(2, 1) void GPUTPCTrackletConstructor::UpdateTracklet(int32_t /*nBlocks*/, float ri = 1.f / CAMath::Sqrt(dx * dx + dy * dy); if (iRow == (int32_t)r.mStartRow + 2) { - tParam.SetSinPhi(dy * ri); + float nnDydx = 0.f, nnDzdx = 0.f; + if (tracker.HitNNDirection(row, seedIH, nnDydx, nnDzdx) && CAMath::Abs(nnDydx) < 10.f && CAMath::Abs(nnDzdx) < 10.f) { + const float nnNormI = CAMath::InvSqrt(1.f + nnDydx * nnDydx); + tParam.SetSinPhi(0.5f * (dy * ri + nnDydx * nnNormI)); + tParam.SetDzDs(0.5f * (dz * ri + nnDzdx * nnNormI)); + } else { + tParam.SetSinPhi(dy * ri); + tParam.SetDzDs(dz * ri); + } tParam.SetSignCosPhi(dx); - tParam.SetDzDs(dz * ri); float err2Y, err2Z; tracker.GetErrors2Seeding(iRow, tParam, -1.f, err2Y, err2Z); // Use correct time tParam.SetCov(0, err2Y); diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h index 7aa23489eb2f4..92c10e06600dc 100644 --- a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h +++ b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizer.h @@ -54,6 +54,7 @@ class GPUTPCNNClusterizer : public GPUProcessor int32_t mNnClusterizerBoundaryFillValue = -1; int32_t mNnClusterizerModelClassNumOutputNodes = -1; int32_t mNnClusterizerModelReg1NumOutputNodes = -1; + int32_t mNnClusterizerUseMomentumVector = 0; int32_t mNnClusterizerModelReg2NumOutputNodes = -1; int32_t mNnInferenceInputDType = 0; // 0: float16, 1: float32 int32_t mNnInferenceOutputDType = 0; // 0: float16, 1: float32 diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx index 96b8a1d7ed2fd..6552ede836ec4 100644 --- a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx +++ b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerHost.cxx @@ -138,6 +138,9 @@ void GPUTPCNNClusterizerHost::initClusterizer(const GPUSettingsProcessingNNclust if (!settings.nnClusterizerUseCfRegression) { if (mModelClass.getNumOutputNodes()[0][1] == 1 || !mModelReg2.isInitialized()) { clustererNN.mNnClusterizerModelReg1NumOutputNodes = mModelReg1.getNumOutputNodes()[0][1]; + if (clustererNN.mNnClusterizerModelReg1NumOutputNodes > 5) { + clustererNN.mNnClusterizerUseMomentumVector = true; + } } else { clustererNN.mNnClusterizerModelReg1NumOutputNodes = mModelReg1.getNumOutputNodes()[0][1]; clustererNN.mNnClusterizerModelReg2NumOutputNodes = mModelReg2.getNumOutputNodes()[0][1]; diff --git a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.cxx b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.cxx index 693ee4dd78e8d..669cf5f4c254b 100644 --- a/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.cxx +++ b/GPU/GPUTracking/TPCClusterFinder/GPUTPCNNClusterizerKernels.cxx @@ -38,6 +38,15 @@ using namespace o2::gpu::tpccf; static_assert(GPUTPCNNClusterizerKernels::SCRATCH_PAD_WORK_GROUP_SIZE == GPUTPCCFClusterizer::SCRATCH_PAD_WORK_GROUP_SIZE, "Work group sizes do not match"); +GPUd() static void setClass1NNDirection(const GPUTPCNNClusterizer& nn, int8_t dtype, int32_t base, o2::tpc::ClusterNative& cluster) +{ + if (dtype == 0) { + cluster.setNNDirection(nn.mOutputDataReg1_32[base + 5], nn.mOutputDataReg1_32[base + 5]); + } else { + cluster.setNNDirection(nn.mOutputDataReg1_16[base + 6].ToFloat(), nn.mOutputDataReg1_16[base + 6].ToFloat()); + } +} + // Defining individual thread functions for data filling, determining the class label and running the CF clusterizer template <> GPUdii() void GPUTPCNNClusterizerKernels::Thread(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUSharedMemory& smem, processorType& processors, uint8_t sector, int8_t dtype, int8_t withMC, uint32_t batchStart) @@ -476,6 +485,9 @@ GPUdii() void GPUTPCNNClusterizerKernels::Thread metadata; metadata["inputDType"] = nnClusterizerSettings.nnInferenceInputDType; // FP16 or FP32 metadata["outputDType"] = nnClusterizerSettings.nnInferenceOutputDType; // FP16 or FP32 - metadata["nnCCDBWithMomentum"] = nnClusterizerSettings.nnCCDBWithMomentum; // 0, 1 -> Only for regression model metadata["nnCCDBLayerType"] = nnClusterizerSettings.nnCCDBClassificationLayerType; // FC, CNN metadata["nnCCDBInteractionRate"] = nnClusterizerSettings.nnCCDBInteractionRate; // in kHz metadata["nnCCDBBeamType"] = nnClusterizerSettings.nnCCDBBeamType; // pp, pPb, PbPb @@ -1286,6 +1285,7 @@ Inputs GPURecoWorkflowSpec::inputs() printSettings(metadata); } + metadata["nnCCDBWithMomentum"] = "0"; if (mSpecConfig.nnEvalMode[0] == "c1") { metadata["nnCCDBEvalType"] = "classification_c1"; convert_map_to_metadata(metadata, ccdb_metadata); @@ -1297,6 +1297,7 @@ Inputs GPURecoWorkflowSpec::inputs() inputs.emplace_back("nn_classification_c2", gDataOriginTPC, "NNCLUSTERIZER_C2", 0, Lifetime::Condition, ccdbParamSpec(nnClusterizerSettings.nnCCDBPath + "/" + metadata["nnCCDBEvalType"], ccdb_metadata, 0)); } + metadata["nnCCDBWithMomentum"] = nnClusterizerSettings.nnCCDBWithMomentum; // 0, 1 -> Only for regression model metadata["nnCCDBEvalType"] = "regression_c1"; metadata["nnCCDBLayerType"] = nnClusterizerSettings.nnCCDBRegressionLayerType; convert_map_to_metadata(metadata, ccdb_metadata); From f97e0946f47e7acb9346f46737e84a2d8aa21345 Mon Sep 17 00:00:00 2001 From: ALICE Action Bot Date: Mon, 6 Jul 2026 17:38:11 +0000 Subject: [PATCH 2/2] Please consider the following formatting changes --- GPU/Workflow/src/GPUWorkflowSpec.cxx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/GPU/Workflow/src/GPUWorkflowSpec.cxx b/GPU/Workflow/src/GPUWorkflowSpec.cxx index 8514b03b52855..6b88d45207951 100644 --- a/GPU/Workflow/src/GPUWorkflowSpec.cxx +++ b/GPU/Workflow/src/GPUWorkflowSpec.cxx @@ -1297,7 +1297,7 @@ Inputs GPURecoWorkflowSpec::inputs() inputs.emplace_back("nn_classification_c2", gDataOriginTPC, "NNCLUSTERIZER_C2", 0, Lifetime::Condition, ccdbParamSpec(nnClusterizerSettings.nnCCDBPath + "/" + metadata["nnCCDBEvalType"], ccdb_metadata, 0)); } - metadata["nnCCDBWithMomentum"] = nnClusterizerSettings.nnCCDBWithMomentum; // 0, 1 -> Only for regression model + metadata["nnCCDBWithMomentum"] = nnClusterizerSettings.nnCCDBWithMomentum; // 0, 1 -> Only for regression model metadata["nnCCDBEvalType"] = "regression_c1"; metadata["nnCCDBLayerType"] = nnClusterizerSettings.nnCCDBRegressionLayerType; convert_map_to_metadata(metadata, ccdb_metadata);