GPU Metal: also run the border-range sorts on the host

GPU/GPUTracking/Base/metal/GPUReconstructionMetal.h

@@ -64,6 +64,10 @@ class GPUReconstructionMetal : public GPUReconstructionProcessing::KernelInterfa
   void CommitStream(int32_t stream);
   void DrainStream(int32_t stream);
   uint64_t EventValue(deviceEvent ev, bool next);
+  template <class T>
+  T* HostPointer(T* p);
+  template <class T, class N, class S>
+  void SortOnHost(const krnlSetupTime& _xyz, T* begin, const N& count, const S& comp);
 
  private:
   int32_t AddKernels();

GPU/GPUTracking/Base/metal/GPUReconstructionMetalKernelsSpecialize.inc

@@ -11,6 +11,8 @@
 
 /// \file GPUReconstructionMetalKernelsSpecialize.inc
 
+#include "GPUTPCGMMergerSortComp.h"
+
 template <>
 inline void GPUReconstructionMetal::runKernelBackend<GPUMemClean16, 0>(const krnlSetupTime& _xyz, void* const& ptr, uint64_t const& size)
 {
@@ -30,3 +32,64 @@ inline void GPUReconstructionMetal::runKernelBackend<GPUMemClean16, 0>(const krn
     RecordMarker(_xyz.z.ev, _xyz.x.stream);
   }
 }
+
+// The sorts that CUDA and HIP run with thrust (GPUCA_SPECIALIZE_THRUST_SORTS)
+// run on a single GPU thread otherwise. The memory is shared with the host, so
+// sort there, with the same std::sort as the CPU backend, which also orders
+// the elements that compare equal as on the CPU. The number of elements is
+// read once the stream is drained, from the device copy of the counters,
+// since the host copy may still be on its way.
+template <class T>
+inline T* GPUReconstructionMetal::HostPointer(T* p)
+{
+  const size_t offset = (const char*)p - (const char*)mDeviceMemoryBase;
+  return offset < mDeviceMemorySize ? (T*)((char*)mInternals->mem_gpu.contents + offset) : p;
+}
+
+template <class T, class N, class S>
+inline void GPUReconstructionMetal::SortOnHost(const krnlSetupTime& _xyz, T* begin, const N& count, const S& comp)
+{
+  if (_xyz.z.evList) {
+    StreamWaitForEvents(_xyz.x.stream, _xyz.z.evList, _xyz.z.nEvents);
+  }
+  DrainStream(_xyz.x.stream);
+  T* hostBegin = HostPointer(begin);
+  GPUCommonAlgorithm::sortDeviceDynamic(hostBegin, hostBegin + count(), comp);
+  if (_xyz.z.ev) {
+    RecordMarker(_xyz.z.ev, _xyz.x.stream);
+  }
+}
+
+template <>
+inline void GPUReconstructionMetal::runKernelBackend<GPUTPCGMMergerSortTracks, 0>(const krnlSetupTime& _xyz)
+{
+  SortOnHost(_xyz, mProcessorsShadow->tpcMerger.TrackOrderProcess(), [&] { return HostPointer(mProcessorsShadow->tpcMerger.Memory())->nMergedTracks; }, internal::GPUTPCGMMergerSortTracks_comp(HostPointer(mProcessorsShadow->tpcMerger.MergedTracks())));
+}
+
+template <>
+inline void GPUReconstructionMetal::runKernelBackend<GPUTPCGMMergerSortTracksQPt, 0>(const krnlSetupTime& _xyz)
+{
+  SortOnHost(_xyz, mProcessorsShadow->tpcMerger.TrackSort(), [&] { return HostPointer(mProcessorsShadow->tpcMerger.Memory())->nMergedTracks; }, internal::GPUTPCGMMergerSortTracksQPt_comp(HostPointer(mProcessorsShadow->tpcMerger.MergedTracks())));
+}
+
+template <>
+inline void GPUReconstructionMetal::runKernelBackend<GPUTPCGMMergerMergeLoopers, 1>(const krnlSetupTime& _xyz)
+{
+  SortOnHost(_xyz, mProcessorsShadow->tpcMerger.LooperCandidates(), [&] { return HostPointer(mProcessorsShadow->tpcMerger.Memory())->nLooperMatchCandidates; }, internal::GPUTPCGMMergerMergeLoopers_comp());
+}
+
+template <>
+inline void GPUReconstructionMetal::runKernelBackend<GPUTPCGMO2Output, GPUTPCGMO2Output::sort>(const krnlSetupTime& _xyz)
+{
+  SortOnHost(_xyz, mProcessorsShadow->tpcMerger.TrackSortO2(), [&] { return HostPointer(mProcessorsShadow->tpcMerger.Memory())->nO2Tracks; }, internal::GPUTPCGMO2OutputSort_comp());
+}
+
+template <>
+inline void GPUReconstructionMetal::runKernelBackend<GPUTPCGMMergerMergeBorders, 3>(const krnlSetupTime& _xyz, gputpcgmmergertypes::GPUTPCGMBorderRange* const& range, int32_t const& N, int32_t const& cmpMax)
+{
+  if (cmpMax) {
+    SortOnHost(_xyz, range, [&] { return N; }, internal::MergeBorderTracks_compMax());
+  } else {
+    SortOnHost(_xyz, range, [&] { return N; }, internal::MergeBorderTracks_compMin());
+  }
+}

GPU/GPUTracking/CMakeLists.txt

@@ -145,6 +145,7 @@ set(HDRS_INSTALL
     Global/GPUChainTrackingGetters.inc
     Global/GPUErrorCodes.h
     Merger/GPUTPCGMMergerTypes.h
+    Merger/GPUTPCGMMergerSortComp.h
     qa/GPUQAHelper.h
     qconfigoptions.h
     Refit/GPUTrackParamConvert.h

GPU/GPUTracking/Merger/GPUTPCGMMerger.cxx

@@ -63,84 +63,7 @@ using namespace o2::gpu;
 using namespace o2::tpc;
 using namespace gputpcgmmergertypes;
 
-namespace o2::gpu::internal
-{
-struct MergeLooperParam {
-  float refz;
-  float x;
-  float y;
-  uint32_t id;
-};
-
-struct MergeBorderTracks_compMax {
-  GPUd() bool operator()(const GPUTPCGMBorderRange& a, const GPUTPCGMBorderRange& b)
-  {
-    return GPUCA_DETERMINISTIC_CODE((a.fMax != b.fMax) ? (a.fMax < b.fMax) : (a.fId < b.fId), a.fMax < b.fMax);
-  }
-};
-struct MergeBorderTracks_compMin {
-  GPUd() bool operator()(const GPUTPCGMBorderRange& a, const GPUTPCGMBorderRange& b)
-  {
-    return GPUCA_DETERMINISTIC_CODE((a.fMin != b.fMin) ? (a.fMin < b.fMin) : (a.fId < b.fId), a.fMin < b.fMin);
-  }
-};
-
-struct GPUTPCGMMergerSortTracks_comp {
-  const GPUTPCGMMergedTrack* const mCmp;
-  GPUhd() GPUTPCGMMergerSortTracks_comp(GPUTPCGMMergedTrack* cmp) : mCmp(cmp) {}
-  GPUd() bool operator()(const int32_t aa, const int32_t bb)
-  {
-    const GPUTPCGMMergedTrack& GPUrestrict() a = mCmp[aa];
-    const GPUTPCGMMergedTrack& GPUrestrict() b = mCmp[bb];
-    if (a.OK() != b.OK()) {
-      return a.OK();
-    }
-    if (a.CCE() != b.CCE()) {
-      return a.CCE() > b.CCE();
-    }
-    GPUCA_DETERMINISTIC_CODE( // clang-format off
-      if (a.NClusters() != b.NClusters()) {
-        return a.NClusters() > b.NClusters();
-      } if (CAMath::Abs(a.GetParam().GetQPt()) != CAMath::Abs(b.GetParam().GetQPt())) {
-        return CAMath::Abs(a.GetParam().GetQPt()) > CAMath::Abs(b.GetParam().GetQPt());
-      } if (a.GetParam().GetY() != b.GetParam().GetY()) {
-        return a.GetParam().GetY() > b.GetParam().GetY();
-      }
-      return aa > bb;
-    , // !GPUCA_DETERMINISTIC_CODE
-      return a.NClusters() > b.NClusters();
-    ) // clang-format on
-  }
-};
-
-struct GPUTPCGMMergerSortTracksQPt_comp {
-  const GPUTPCGMMergedTrack* const mCmp;
-  GPUhd() GPUTPCGMMergerSortTracksQPt_comp(GPUTPCGMMergedTrack* cmp) : mCmp(cmp) {}
-  GPUd() bool operator()(const int32_t aa, const int32_t bb)
-  {
-    const GPUTPCGMMergedTrack& GPUrestrict() a = mCmp[aa];
-    const GPUTPCGMMergedTrack& GPUrestrict() b = mCmp[bb];
-    GPUCA_DETERMINISTIC_CODE( // clang-format off
-      if (CAMath::Abs(a.GetParam().GetQPt()) != CAMath::Abs(b.GetParam().GetQPt())) {
-        return CAMath::Abs(a.GetParam().GetQPt()) > CAMath::Abs(b.GetParam().GetQPt());
-      } if (a.GetParam().GetY() != b.GetParam().GetY()) {
-        return a.GetParam().GetY() > b.GetParam().GetY();
-      }
-      return a.GetParam().GetZ() > b.GetParam().GetZ();
-    , // !GPUCA_DETERMINISTIC_CODE
-      return CAMath::Abs(a.GetParam().GetQPt()) > CAMath::Abs(b.GetParam().GetQPt());
-    ) // clang-format on
-  }
-};
-
-struct GPUTPCGMMergerMergeLoopers_comp {
-  GPUd() bool operator()(const MergeLooperParam& a, const MergeLooperParam& b)
-  {
-    return GPUCA_DETERMINISTIC_CODE(CAMath::Abs(a.refz) != CAMath::Abs(b.refz) ? CAMath::Abs(a.refz) < CAMath::Abs(b.refz) : a.id < b.id, CAMath::Abs(a.refz) < CAMath::Abs(b.refz));
-  }
-};
-
-} // namespace o2::gpu::internal
+#include "GPUTPCGMMergerSortComp.h"
 
 using namespace o2::gpu::internal;
 

GPU/GPUTracking/Merger/GPUTPCGMMergerSortComp.h

@@ -0,0 +1,110 @@
+// Copyright 2019-2025 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// In applying this license CERN does not waive the privileges and immunities
+// granted to it by virtue of its status as an Intergovernmental Organization
+// or submit itself to any jurisdiction.
+
+/// \file GPUTPCGMMergerSortComp.h
+/// \brief Orderings of the merger's sorts, for the kernels and for the backends that run these sorts themselves
+
+#ifndef GPUTPCGMMERGERSORTCOMP_H
+#define GPUTPCGMMERGERSORTCOMP_H
+
+#include "GPUCommonDef.h"
+#include "GPUCommonMath.h"
+#include "GPUTPCGMMergedTrack.h"
+#include "GPUTPCGMMergerTypes.h"
+#include "GPUTPCGMMerger.h"
+
+namespace o2::gpu::internal
+{
+struct MergeLooperParam {
+  float refz;
+  float x;
+  float y;
+  uint32_t id;
+};
+
+struct MergeBorderTracks_compMax {
+  GPUd() bool operator()(const gputpcgmmergertypes::GPUTPCGMBorderRange& a, const gputpcgmmergertypes::GPUTPCGMBorderRange& b)
+  {
+    return GPUCA_DETERMINISTIC_CODE((a.fMax != b.fMax) ? (a.fMax < b.fMax) : (a.fId < b.fId), a.fMax < b.fMax);
+  }
+};
+struct MergeBorderTracks_compMin {
+  GPUd() bool operator()(const gputpcgmmergertypes::GPUTPCGMBorderRange& a, const gputpcgmmergertypes::GPUTPCGMBorderRange& b)
+  {
+    return GPUCA_DETERMINISTIC_CODE((a.fMin != b.fMin) ? (a.fMin < b.fMin) : (a.fId < b.fId), a.fMin < b.fMin);
+  }
+};
+
+struct GPUTPCGMMergerSortTracks_comp {
+  const GPUTPCGMMergedTrack* const mCmp;
+  GPUhd() GPUTPCGMMergerSortTracks_comp(GPUTPCGMMergedTrack* cmp) : mCmp(cmp) {}
+  GPUd() bool operator()(const int32_t aa, const int32_t bb)
+  {
+    const GPUTPCGMMergedTrack& GPUrestrict() a = mCmp[aa];
+    const GPUTPCGMMergedTrack& GPUrestrict() b = mCmp[bb];
+    if (a.OK() != b.OK()) {
+      return a.OK();
+    }
+    if (a.CCE() != b.CCE()) {
+      return a.CCE() > b.CCE();
+    }
+    GPUCA_DETERMINISTIC_CODE( // clang-format off
+      if (a.NClusters() != b.NClusters()) {
+        return a.NClusters() > b.NClusters();
+      } if (CAMath::Abs(a.GetParam().GetQPt()) != CAMath::Abs(b.GetParam().GetQPt())) {
+        return CAMath::Abs(a.GetParam().GetQPt()) > CAMath::Abs(b.GetParam().GetQPt());
+      } if (a.GetParam().GetY() != b.GetParam().GetY()) {
+        return a.GetParam().GetY() > b.GetParam().GetY();
+      }
+      return aa > bb;
+    , // !GPUCA_DETERMINISTIC_CODE
+      return a.NClusters() > b.NClusters();
+    ) // clang-format on
+  }
+};
+
+struct GPUTPCGMMergerSortTracksQPt_comp {
+  const GPUTPCGMMergedTrack* const mCmp;
+  GPUhd() GPUTPCGMMergerSortTracksQPt_comp(GPUTPCGMMergedTrack* cmp) : mCmp(cmp) {}
+  GPUd() bool operator()(const int32_t aa, const int32_t bb)
+  {
+    const GPUTPCGMMergedTrack& GPUrestrict() a = mCmp[aa];
+    const GPUTPCGMMergedTrack& GPUrestrict() b = mCmp[bb];
+    GPUCA_DETERMINISTIC_CODE( // clang-format off
+      if (CAMath::Abs(a.GetParam().GetQPt()) != CAMath::Abs(b.GetParam().GetQPt())) {
+        return CAMath::Abs(a.GetParam().GetQPt()) > CAMath::Abs(b.GetParam().GetQPt());
+      } if (a.GetParam().GetY() != b.GetParam().GetY()) {
+        return a.GetParam().GetY() > b.GetParam().GetY();
+      }
+      return a.GetParam().GetZ() > b.GetParam().GetZ();
+    , // !GPUCA_DETERMINISTIC_CODE
+      return CAMath::Abs(a.GetParam().GetQPt()) > CAMath::Abs(b.GetParam().GetQPt());
+    ) // clang-format on
+  }
+};
+
+struct GPUTPCGMMergerMergeLoopers_comp {
+  GPUd() bool operator()(const MergeLooperParam& a, const MergeLooperParam& b)
+  {
+    return GPUCA_DETERMINISTIC_CODE(CAMath::Abs(a.refz) != CAMath::Abs(b.refz) ? CAMath::Abs(a.refz) < CAMath::Abs(b.refz) : a.id < b.id, CAMath::Abs(a.refz) < CAMath::Abs(b.refz));
+  }
+};
+
+struct GPUTPCGMO2OutputSort_comp {
+  GPUd() bool operator()(const GPUTPCGMMerger::tmpSort& a, const GPUTPCGMMerger::tmpSort& b)
+  {
+    return GPUCA_DETERMINISTIC_CODE(a.y != b.y ? a.y > b.y : a.x > b.x, a.y > b.y);
+  }
+};
+
+} // namespace o2::gpu::internal
+
+#endif

GPU/GPUTracking/Merger/GPUTPCGMO2Output.cxx

@@ -33,17 +33,7 @@ using namespace o2::tpc::constants;
 
 GPUdi() static constexpr uint8_t getFlagsReject() { return GPUTPCGMMergedTrackHit::flagReject | GPUTPCGMMergedTrackHit::flagHighIncl; }
 
-namespace o2::gpu::internal
-{
-
-struct GPUTPCGMO2OutputSort_comp {
-  GPUd() bool operator()(const GPUTPCGMMerger::tmpSort& a, const GPUTPCGMMerger::tmpSort& b)
-  {
-    return GPUCA_DETERMINISTIC_CODE(a.y != b.y ? a.y > b.y : a.x > b.x, a.y > b.y);
-  }
-};
-
-} // namespace o2::gpu::internal
+#include "GPUTPCGMMergerSortComp.h"
 
 template <>
 GPUdii() void GPUTPCGMO2Output::Thread<GPUTPCGMO2Output::prepare>(int32_t nBlocks, int32_t nThreads, int32_t iBlock, int32_t iThread, GPUsharedref() GPUSharedMemory& smem, processorType& GPUrestrict() merger)

GPU/GPUTracking/TPCClusterFinder/ClusterAccumulator.cxx

@@ -62,13 +62,26 @@ GPUd() void ClusterAccumulator::finalize(const CfChargePos& pos, const Charge q,
 {
   mQtot += q;
 
+#ifdef __METAL__
+  // Metal's fast-math division and sqrt are not correctly rounded, which moves
+  // widths that sit within an ulp of a packing step. Do the operations the
+  // arm64 -ffast-math CPU build does here: reciprocal, products, fused variance.
+  const float invQ = metal::precise::divide(1.f, mQtot);
+  mPadMean *= invQ;
+  mTimeMean *= invQ;
+  const float padMean2 = mPadMean * mPadMean;
+  const float timeMean2 = mTimeMean * mTimeMean;
+  mPadSigma = metal::precise::sqrt(metal::fma(mPadSigma, invQ, -padMean2));
+  mTimeSigma = metal::precise::sqrt(metal::fma(mTimeSigma, invQ, -timeMean2));
+#else
   mPadMean /= mQtot;
   mTimeMean /= mQtot;
   mPadSigma /= mQtot;
   mTimeSigma /= mQtot;
 
   mPadSigma = CAMath::Sqrt(mPadSigma - mPadMean * mPadMean);
   mTimeSigma = CAMath::Sqrt(mTimeSigma - mTimeMean * mTimeMean);
+#endif
 
   Pad pad = pos.pad();
   mPadMean += pad;