GPU Metal: run the merger's sorts on the host

The sorts of the merged tracks by quality and by q/pT, of the looper
candidates and of the O2 output tracks, which CUDA and HIP run with
thrust, ran on a single GPU thread otherwise: about 730 ms per event on
a 30-collision Pb-Pb time frame, at up to 140 times the CPU's time. The
GPU memory is shared with the host, so sort there once the stream is
drained, with the std::sort of the CPU backend, which also orders the
elements that compare equal as on the CPU. The number of elements is
read from the device copy of the counters, as the host copy of the O2
output count is still being transferred when the sort is launched.

Author: ktf

GPU/GPUTracking/Base/metal/GPUReconstructionMetal.h

@@ -60,6 +60,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,54 @@ 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());
+}