diff --git a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackingKernels.cu b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackingKernels.cu index 43ee8436892d4..bdb702d3b4427 100644 --- a/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackingKernels.cu +++ b/Detectors/ITSMFT/ITS/tracking/GPU/cuda/TrackingKernels.cu @@ -52,6 +52,7 @@ #include "ITStrackingGPU/TrackingKernels.h" #include "ITStrackingGPU/Utils.h" #include "MathUtils/Utils.h" +#include "GPUCommonAlgorithm.h" #include "ITStracking/ClusterLines.h" #include "ITStracking/LineProjection.h" #include "utils/strtag.h" @@ -1000,85 +1001,108 @@ GPUg() void fitPeaksKernel(const int* nPeaksDevice, VertexCand* cands) { const int nPeaks = *nPeaksDevice; - const o2::its::TimeEstBC* t = lineProjSorted.t; const int* idx = lineProjSorted.idx; - for (int p = blockIdx.x * blockDim.x + threadIdx.x; p < nPeaks; p += blockDim.x * gridDim.x) { - cands[p].ok = 0; - cands[p].nGood = 0; - const int k = peakLineIdx[p]; - cands[p].fine = isZPeakFine != nullptr ? isZPeakFine[k] : 0; + const o2::its::TimeEstBC* t = lineProjSorted.t; + const int lane = threadIdx.x % gpu::WarpSize; + const int nWarps = blockDim.x * gridDim.x / gpu::WarpSize; + const int iWarp = (blockIdx.x * blockDim.x + threadIdx.x) / gpu::WarpSize; + for (int iPeak = iWarp; iPeak < nPeaks; iPeak += nWarps) { + const int k = peakLineIdx[iPeak]; + if (lane == 0) { + cands[iPeak].ok = 0; + cands[iPeak].nGood = 0; + cands[iPeak].fine = isZPeakFine != nullptr ? isZPeakFine[k] : 0; + } const auto tk = t[k]; const LineWindow wk = win[k]; - - o2::its::ClusterLines seed; int nMembers = 0; - for (int j = wk.lo; j < wk.hi; ++j) { - const auto tj = t[j]; - if (tk.isCompatible(tj)) { - seed.add(lines[idx[j]]); - ++nMembers; + WarpFit seed; + for (int iLine = wk.lo + lane; iLine < wk.hi; iLine += gpu::WarpSize) { + const auto tj = t[iLine]; + if (!tk.isCompatible(tj)) { + continue; } + seed.add(lines[idx[iLine]]); + ++nMembers; } - float seedVertex[3]; - if (nMembers < 2 || !seed.solve(seedVertex)) { + seed.reduce(); + nMembers = warp_reduce_add(nMembers); + seed.computeClusterCentroid(); + if (nMembers < 2 || !seed.isValid()) { continue; } + const float* seedVertex = seed.getVertex(); - o2::its::ClusterLines fit; + WarpFit fit; int nKept = 0; int nGood = 0; - for (int j = wk.lo; j < wk.hi; ++j) { - const auto tj = t[j]; - if (tk.isCompatible(tj)) { - const o2::its::Line& line = lines[idx[j]]; - if (o2::its::Line::getDistance2FromPoint(line, seedVertex) < pairCut2) { - fit.add(line); - const float c = lineChi2[idx[j]]; - const float pt = linePt[idx[j]]; - const bool okChi2 = (goodLineChi2Cut <= 0.f || c <= goodLineChi2Cut); - const bool okPt = (goodLinePtCut <= 0.f || pt >= goodLinePtCut); - nGood += okChi2 && okPt; - ++nKept; - } + for (int iLine = wk.lo + lane; iLine < wk.hi; iLine += gpu::WarpSize) { + const auto tj = t[iLine]; + if (!tk.isCompatible(tj)) { + continue; + } + const o2::its::Line& line = lines[idx[iLine]]; + if (!(o2::its::Line::getDistance2FromPoint(line, seedVertex) < pairCut2)) { + continue; } + fit.add(line); + const float c = lineChi2[idx[iLine]]; + const float pt = linePt[idx[iLine]]; + const bool okChi2 = (goodLineChi2Cut <= 0.f || c <= goodLineChi2Cut); + const bool okPt = (goodLinePtCut <= 0.f || pt >= goodLinePtCut); + nGood += okChi2 && okPt; + ++nKept; } - float vertex[3]; - if (nKept < 2 || !fit.solve(vertex)) { + fit.reduce(); + nKept = warp_reduce_add(nKept); + nGood = warp_reduce_add(nGood); + fit.computeClusterCentroid(); + if (nKept < 2 || !fit.isValid()) { continue; } - cands[p].seed[0] = seedVertex[0]; - cands[p].seed[1] = seedVertex[1]; - cands[p].seed[2] = seedVertex[2]; + const float* vertex = fit.getVertex(); + + if (lane == 0) { + cands[iPeak].seed[0] = seedVertex[0]; + cands[iPeak].seed[1] = seedVertex[1]; + cands[iPeak].seed[2] = seedVertex[2]; + } + const float bd2 = (beamX - vertex[0]) * (beamX - vertex[0]) + (beamY - vertex[1]) * (beamY - vertex[1]); - if (nKept < minContributors || !(bd2 < nSigmaCut)) { + if (fit.getNContributors() < minContributors || !(bd2 < nSigmaCut)) { continue; } - if (fineMaxDrift > 0.f && cands[p].fine && + if (fineMaxDrift > 0.f && (isZPeakFine != nullptr ? isZPeakFine[k] : 0) && o2::gpu::GPUCommonMath::Abs(vertex[2] - seedVertex[2]) > fineMaxDrift) { continue; } - for (int j = wk.lo; j < wk.hi; ++j) { - const auto tj = t[j]; - if (tk.isCompatible(tj)) { - const o2::its::Line& line = lines[idx[j]]; - if (o2::its::Line::getDistance2FromPoint(line, seedVertex) < pairCut2) { - fit.addResidual(line, vertex); - } + for (int iLine = wk.lo + lane; iLine < wk.hi; iLine += gpu::WarpSize) { + const auto tj = t[iLine]; + if (!tk.isCompatible(tj)) { + continue; } + const o2::its::Line& line = lines[idx[iLine]]; + if (!(o2::its::Line::getDistance2FromPoint(line, seedVertex) < pairCut2)) { + continue; + } + fit.addResidual(line, vertex); } - - cands[p].x = vertex[0]; - cands[p].y = vertex[1]; - cands[p].z = vertex[2]; - for (int i = 0; i < 6; ++i) { - cands[p].rms2[i] = fit.getRMS2()[i]; + fit.reduceResiduals(nKept); + + if (lane == 0) { + cands[iPeak].x = vertex[0]; + cands[iPeak].y = vertex[1]; + cands[iPeak].z = vertex[2]; + for (int i = 0; i < 6; ++i) { + cands[iPeak].rms2[i] = fit.getRMS2()[i]; + } + cands[iPeak].avgDist2 = fit.getAvgDistance2(); + cands[iPeak].nGood = nGood; + cands[iPeak].time = fit.getTimeStamp(); + cands[iPeak].size = fit.getNContributors(); + cands[iPeak].ok = 1; } - cands[p].avgDist2 = fit.getAvgDistance2(); - cands[p].nGood = nGood; - cands[p].time = fit.getTimeStamp(); - cands[p].size = nKept; - cands[p].ok = 1; } } @@ -1721,6 +1745,13 @@ void TrackingKernels::fitPeaksHandler(const int* nPeaksDevice, gpu::VertexCand* cands, gpu::Stream& stream) { + // the kernel reduces with shuffles over gpu::WarpSize lanes: a device with another warp size would silently sum the wrong lanes + int dev = 0, warpSize = 0; + GPUChkErrS(cudaGetDevice(&dev)); + GPUChkErrS(cudaDeviceGetAttribute(&warpSize, cudaDevAttrWarpSize, dev)); + if (warpSize != gpu::WarpSize) { + LOGP(fatal, "fitPeaksKernel compiled for warp size {} but the device runs {}", gpu::WarpSize, warpSize); + } gpu::fitPeaksKernel<<>>( nPeaksDevice, peakLineIdx, win, sortedSoa, lines, lineChi2, linePt, goodLineChi2Cut, goodLinePtCut, pairCut2, nSigmaCut, minContributors, beamX, beamY, isPeakFine, fineMaxDrift, cands); } diff --git a/Detectors/ITSMFT/ITS/tracking/include/ITStracking/ClusterLines.h b/Detectors/ITSMFT/ITS/tracking/include/ITStracking/ClusterLines.h index f1255a4ee9d1f..5be9ab093cc91 100644 --- a/Detectors/ITSMFT/ITS/tracking/include/ITStracking/ClusterLines.h +++ b/Detectors/ITSMFT/ITS/tracking/include/ITStracking/ClusterLines.h @@ -17,6 +17,8 @@ #include #ifndef GPUCA_GPUCODE #include +#else +#include "GPUCommonAlgorithm.h" // warp reductions of WarpFit #endif #include "ITStracking/Cluster.h" #include "ITStracking/Cell.h" @@ -141,27 +143,19 @@ GPUdi() bool getCellLineAtBeam(const CellSeed& cell, const float beamX, const fl return true; } -/// Least-squares vertex fit over a set of lines (the normal equations AX = -B). -class ClusterLines final +/// Normal equations AX = -B of the point closest to a set of lines, and their solution. Base of the +/// sequential fit (ClusterLines) and of the device warp-parallel fit (WarpFit), which only differ in how they accumulate. +class LineFit { public: - GPUhdDefault() ClusterLines() = default; - /// Fit over lines[lineIndices[i]]; lineIndices must be sorted and unique. - ClusterLines(gsl::span lineIndices, gsl::span lines); - - /// Accumulate one line into the normal equations and the time estimate. - GPUhdi() void add(const Line& line) + /// Accumulate one line + GPUhdi() bool add(const Line& line) { const double d0 = line.cosinesDirector[0], d1 = line.cosinesDirector[1], d2 = line.cosinesDirector[2]; const double o0 = line.originPoint[0], o1 = line.originPoint[1], o2 = line.originPoint[2]; const double det = d0 * d0 + d1 * d1 + d2 * d2; // == 1 for a normalised director if (det <= 0.) { - return; - } - if (mNContributors <= 0) { - mTime = line.mTime; - } else { - mTime += line.mTime; + return false; } mA[0] += (det - d0 * d0) / det; mA[1] += (-d0 * d1) / det; @@ -173,20 +167,19 @@ class ClusterLines final mB[0] += (d0 * dDotO - det * o0) / det; mB[1] += (d1 * dDotO - det * o1) / det; mB[2] += (d2 * dDotO - det * o2) / det; - ++mNContributors; + return true; } - /// Solve the symmetric 3x3 system into an external vertex (= -A^-1 B) + /// Solve the symmetric 3x3 system into vertex (= -A^-1 B); false if it is singular GPUhdi() bool solve(float vertex[3]) const { - const double a = mA[0], b = mA[1], c = mA[2], d = mA[3], e = mA[4], f = mA[5]; - const double c00 = d * f - e * e; - const double c01 = c * e - b * f; - const double c02 = b * e - c * d; - const double c11 = a * f - c * c; - const double c12 = b * c - a * e; - const double c22 = a * d - b * b; - const double det = a * c00 + b * c01 + c * c02; + const double c00 = mA[3] * mA[5] - mA[4] * mA[4]; + const double c01 = mA[2] * mA[4] - mA[1] * mA[5]; + const double c02 = mA[1] * mA[4] - mA[2] * mA[3]; + const double c11 = mA[0] * mA[5] - mA[2] * mA[2]; + const double c12 = mA[1] * mA[2] - mA[0] * mA[4]; + const double c22 = mA[0] * mA[3] - mA[1] * mA[1]; + const double det = mA[0] * c00 + mA[1] * c01 + mA[2] * c02; if (o2::gpu::GPUCommonMath::Abs(det) < 1.e-12) { return false; } @@ -203,6 +196,49 @@ class ClusterLines final /// Solve into the stored vertex, setting the validity flag GPUhdi() void computeClusterCentroid() { mIsValid = solve(mVertex); } + GPUhdi() bool isValid() const noexcept { return mIsValid; } + GPUhdi() const float* getVertex() const noexcept { return mVertex; } + GPUhdi() int getSize() const noexcept { return mNContributors; } + GPUhdi() int getNContributors() const noexcept { return mNContributors; } + GPUhdi() const TimeEstBC& getTimeStamp() const noexcept { return mTime; } + GPUhdi() float getR2() const noexcept { return (mVertex[0] * mVertex[0]) + (mVertex[1] * mVertex[1]); } + GPUhdi() float getR() const noexcept { return o2::gpu::GPUCommonMath::Sqrt(getR2()); } + GPUhdi() const float* getRMS2() const noexcept { return mRMS2; } // {XX, XY, YY, XZ, YZ, ZZ} + GPUhdi() float getAvgDistance2() const noexcept { return mAvgDistance2; } + + protected: + double mA[6] = {0., 0., 0., 0., 0., 0.}; // AX=B, packed symmetric normal matrix + double mB[3] = {0., 0., 0.}; // AX=B, right-hand side + float mVertex[3] = {0.f, 0.f, 0.f}; // cluster centroid position + float mRMS2[6] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; // mean DCA components about the vertex + float mAvgDistance2 = 0.f; // mean squared distance about the vertex + int mNContributors = 0; // lines accumulated into the normal equations + bool mIsValid = false; // true if the linear system was solved successfully + TimeEstBC mTime; // intersection of the contributors' time intervals +}; + +/// Least-squares vertex fit over a set of lines (the normal equations AX = -B). +class ClusterLines final : public LineFit +{ + public: + GPUhdDefault() ClusterLines() = default; + /// Fit over lines[lineIndices[i]]; lineIndices must be sorted and unique. + ClusterLines(gsl::span lineIndices, gsl::span lines); + + /// Accumulate one line into the normal equations and the time estimate. + GPUhdi() void add(const Line& line) + { + if (!LineFit::add(line)) { + return; + } + if (mNContributors <= 0) { + mTime = line.mTime; + } else { + mTime += line.mTime; + } + ++mNContributors; + } + /// Running-mean update of the RMS2 and average distance about the stored vertex GPUhdi() void addResidual(const Line& line) { addResidual(line, mVertex); } @@ -220,28 +256,66 @@ class ClusterLines final mAvgDistance2 += (d2 - mAvgDistance2) / static_cast(mResidualCount); } - GPUhdi() bool isValid() const noexcept { return mIsValid; } - GPUhdi() const float* getVertex() const noexcept { return mVertex; } - GPUhdi() const float* getRMS2() const noexcept { return mRMS2; } // {XX, XY, YY, XZ, YZ, ZZ} - GPUhdi() float getAvgDistance2() const noexcept { return mAvgDistance2; } - GPUhdi() int getSize() const noexcept { return mNContributors; } - GPUhdi() int getNContributors() const noexcept { return mNContributors; } - GPUhdi() const TimeEstBC& getTimeStamp() const noexcept { return mTime; } - GPUhdi() float getR2() const noexcept { return (mVertex[0] * mVertex[0]) + (mVertex[1] * mVertex[1]); } - GPUhdi() float getR() const noexcept { return o2::gpu::GPUCommonMath::Sqrt(getR2()); } bool operator==(const ClusterLines& rhs) const noexcept; private: - double mA[6] = {0., 0., 0., 0., 0., 0.}; // AX=B, packed symmetric normal matrix - double mB[3] = {0., 0., 0.}; // AX=B, right-hand side - float mVertex[3] = {0.f, 0.f, 0.f}; // cluster centroid position - float mRMS2[6] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; - float mAvgDistance2 = 0.f; - int mNContributors = 0; - int mResidualCount = 0; - bool mIsValid = false; // true if the linear system was solved successfully - TimeEstBC mTime; + int mResidualCount = 0; // lines folded into the running means +}; + +#ifdef GPUCA_GPUCODE +/// Device counterpart of ClusterLines: each lane of a warp add()s its own subset of lines, then reduce() leaves +/// the full sums, contributor count and time intersection in every lane +struct WarpFit : public LineFit { + uint32_t tLo = 0, tHi = 0xffffffffu; + + GPUdi() void add(const Line& l) + { + if (LineFit::add(l)) { + tLo = o2::gpu::GPUCommonMath::Max(tLo, l.mTime.lower()); + tHi = o2::gpu::GPUCommonMath::Min(tHi, l.mTime.upper()); + ++mNContributors; + } + } + + template + GPUdi() void reduce() + { + for (int i = 0; i < 6; ++i) { + mA[i] = warp_reduce_add(mA[i]); + } + for (int i = 0; i < 3; ++i) { + mB[i] = warp_reduce_add(mB[i]); + } + tLo = warp_reduce_max(tLo); + tHi = warp_reduce_min(tHi); + mNContributors = warp_reduce_add(mNContributors); + mTime = TimeEstBC{tLo, tHi - tLo}; + } + + GPUdi() void addResidual(const Line& l, const float vertex[3]) + { + float dca[6]; + Line::getDCAComponents(l, vertex, dca); + for (int i = 0; i < 6; ++i) { + rms2Sum[i] += dca[i]; + } + dist2Sum += Line::getDistance2FromPoint(l, vertex); + } + + /// Means over the n lines of the whole warp, left in every lane + template + GPUdi() void reduceResiduals(const int n) + { + for (int i = 0; i < 6; ++i) { + mRMS2[i] = warp_reduce_add(rms2Sum[i]) / n; + } + mAvgDistance2 = warp_reduce_add(dist2Sum) / n; + } + + double rms2Sum[6] = {0., 0., 0., 0., 0., 0.}; + double dist2Sum = 0.; }; +#endif } // namespace o2::its #endif /* O2_ITS_CLUSTERLINES_H */ diff --git a/Detectors/ITSMFT/ITS/tracking/src/TrackingLinkDef.h b/Detectors/ITSMFT/ITS/tracking/src/TrackingLinkDef.h index 5ba4e36dda875..8adce18da865a 100644 --- a/Detectors/ITSMFT/ITS/tracking/src/TrackingLinkDef.h +++ b/Detectors/ITSMFT/ITS/tracking/src/TrackingLinkDef.h @@ -30,6 +30,7 @@ #pragma link C++ class o2::its::Line + ; #pragma link C++ class std::vector < o2::its::Line> + ; +#pragma link C++ class o2::its::LineFit + ; #pragma link C++ class o2::its::ClusterLines + ; #pragma link C++ class std::vector < o2::its::ClusterLines> + ; diff --git a/GPU/Common/GPUCommonAlgorithm.h b/GPU/Common/GPUCommonAlgorithm.h index f9af751d847cf..5145421087a4b 100644 --- a/GPU/Common/GPUCommonAlgorithm.h +++ b/GPU/Common/GPUCommonAlgorithm.h @@ -429,6 +429,47 @@ GPUdi() T warp_broadcast_FUNC(T v, int32_t i) #endif } +template +GPUdi() T warp_shuffle_xor(T v, int32_t laneMask) +{ +#ifdef __CUDACC__ + return __shfl_xor_sync(0xFFFFFFFF, v, laneMask); +#else // HIP + return __shfl_xor(v, laneMask); +#endif +} + +// Butterfly all-reduce over groups of Width lanes: every lane gets the result. +// Width is a power of two and must not exceed the hardware warp size. +template +GPUdi() T warp_reduce_add(T v) +{ + for (int32_t m = Width / 2; m > 0; m >>= 1) { + v += warp_shuffle_xor(v, m); + } + return v; +} + +template +GPUdi() T warp_reduce_min(T v) +{ + for (int32_t m = Width / 2; m > 0; m >>= 1) { + const T o = warp_shuffle_xor(v, m); + v = o < v ? o : v; + } + return v; +} + +template +GPUdi() T warp_reduce_max(T v) +{ + for (int32_t m = Width / 2; m > 0; m >>= 1) { + const T o = warp_shuffle_xor(v, m); + v = o > v ? o : v; + } + return v; +} + [[nodiscard]] GPUdi() bool work_group_any(bool pred) { return __syncthreads_or(pred); @@ -472,6 +513,24 @@ GPUdi() T warp_broadcast(T v, int32_t i) return v; } +template +GPUdi() T warp_reduce_add(T v) +{ + return v; +} + +template +GPUdi() T warp_reduce_min(T v) +{ + return v; +} + +template +GPUdi() T warp_reduce_max(T v) +{ + return v; +} + [[nodiscard]] GPUdi() bool work_group_any(bool pred) { return pred;