GPU Metal: optionally limit the GPU work in flight

At 20 kHz occupancy the macOS GPU watchdog killed command buffers, and
reset the desktop once, while all 8 streams ran multi-kernel buffers of
up to ~1 s at the same time. Host-serialised execution (serializeGPU 3)
was the only mode that survived. Add two runtime limits, off by default:

- GPUCA_METAL_MAX_KERNELS_PER_BUFFER=N commits a stream's command buffer
  after N kernels, so that long chains (sector tracking, SectorRefit,
  TrackFit) become short buffers. Buffers of a stream stay ordered by its
  order event, the events of the other streams are untouched.
- GPUCA_METAL_MAX_STREAMS_IN_FLIGHT=M makes the host wait, before it
  commits a buffer on an idle stream, until fewer than M streams have
  committed buffers that are not complete. This cannot deadlock: a buffer
  only waits for events signalled by buffers committed before it, since
  RecordMarker commits the signal at once, so the work in flight completes
  without the host.

pbpb30, M1 Max, 6 rounds of 10 events (timeline: longest buffer, streams
running at once): default unchanged (931 ms/ev, -3 ms over two run orders);
N=1 M=1 1416 ms/ev, 87 ms, 1 stream (serializeGPU 3: 1413 ms/ev);
N=1 M=2 1100 ms/ev, 87 ms, 2; N=2 M=4 941 ms/ev, 123 ms, 4. The longest
buffers left hold the TrackFit kernel (~85 ms alone). Metal dumps byte-identical on
pbpb30 with and without the limits, sim5pp identical to the CPU.

Author: ktf

GPU/GPUTracking/Base/metal/GPUReconstructionMetal.mm

@@ -172,6 +172,8 @@ uint32_t metalGPUCoreCount()
     if (const char* timeline = getenv("GPUCA_METAL_TIMELINE")) {
       mInternals->timeline = fopen(timeline, "w");
     }
+    mInternals->maxKernelsPerBuffer = getenv("GPUCA_METAL_MAX_KERNELS_PER_BUFFER") ? atoi(getenv("GPUCA_METAL_MAX_KERNELS_PER_BUFFER")) : 0;
+    mInternals->maxStreamsInFlight = getenv("GPUCA_METAL_MAX_STREAMS_IN_FLIGHT") ? atoi(getenv("GPUCA_METAL_MAX_STREAMS_IN_FLIGHT")) : 0;
 
     mInternals->mem_host = metalNewSharedBuffer(mInternals->device, mHostMemorySize);
     if (mInternals->mem_host == nil) {
@@ -251,6 +253,9 @@ uint32_t metalGPUCoreCount()
       }
       [mInternals->mem_host release];
     }
+    // The completion handlers of the limit still use the internals
+    std::unique_lock<std::mutex> lock(mInternals->inFlightMutex);
+    mInternals->inFlightDone.wait(lock, [this] { return mInternals->streamsInFlight == 0; });
     if (mInternals->timeline) {
       fflush(mInternals->timeline);  // not closed: a completion handler may still write to it
     }
@@ -406,6 +411,9 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
   const uint64_t value = ++mInternals->streamOrderValue[stream];
   [cb encodeSignalEvent:mInternals->streamOrder[stream] value:value];
   metalSignalOnError(cb, mInternals->streamOrder[stream], value);
+  if (mInternals->maxStreamsInFlight) {
+    metalLimitStreamsInFlight(mInternals, cb, stream);
+  }
   if (mInternals->timeline) {
     metalTimelineAdd(mInternals, stream, "s", mInternals->streamOrder[stream], value);
     FILE* f = mInternals->timeline;
@@ -424,6 +432,7 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
   // it until it is waited for, and open a fresh one to encode the next work in.
   mInternals->submittedBuffers[stream].emplace_back(cb);
   mInternals->streamHasWork[stream] = false;
+  mInternals->streamKernels[stream] = 0;
   mInternals->commandBuffers[stream] = [[mInternals->commandQueues[stream] commandBuffer] retain];
   [mInternals->commandBuffers[stream] encodeWaitForEvent:mInternals->streamOrder[stream] value:value];
   metalTimelineAdd(mInternals, stream, "w", mInternals->streamOrder[stream], value);

GPU/GPUTracking/Base/metal/GPUReconstructionMetalIncludesHost.h

@@ -33,6 +33,7 @@ using half = _Float16;
 #include <string>
 #include <memory>
 #include <mutex>
+#include <condition_variable>
 #include <unordered_map>
 #include <cstdio>
 #include <mach/mach_time.h>
@@ -76,6 +77,17 @@ struct GPUReconstructionMetalInternals {
   std::array<std::string, constants::GPU_MAX_STREAMS> timelineOps;
   std::array<uint32_t, constants::GPU_MAX_STREAMS> timelineKernels;
   uint64_t timelineBuffers = 0;
+  // Opt-in limits on the GPU work in flight, 0 = none, so that no stretch of GPU
+  // work runs long enough for the macOS GPU watchdog: GPUCA_METAL_MAX_KERNELS_PER_BUFFER
+  // commits a stream's buffer after that many kernels, GPUCA_METAL_MAX_STREAMS_IN_FLIGHT
+  // is how many streams may have committed buffers that are not complete yet.
+  uint32_t maxKernelsPerBuffer = 0;
+  uint32_t maxStreamsInFlight = 0;
+  std::array<uint32_t, constants::GPU_MAX_STREAMS> streamKernels{};
+  std::array<uint32_t, constants::GPU_MAX_STREAMS> streamBuffersInFlight{};
+  uint32_t streamsInFlight = 0;
+  std::mutex inFlightMutex;
+  std::condition_variable inFlightDone;
 
   std::vector<id<MTLFunction>> functions;
   std::vector<id<MTLComputePipelineState>> pipelines;
@@ -111,6 +123,25 @@ inline void metalSignalOnError(id<MTLCommandBuffer> cb, id<MTLSharedEvent> event
     }
   }];
 }
+// Before a buffer of an idle stream is committed, wait until fewer than the
+// limit of streams have work in flight. This cannot deadlock: a buffer only waits
+// for events that buffers committed before it signal (RecordMarker commits the
+// signal at once), so the work in flight completes without the host's help.
+inline void metalLimitStreamsInFlight(GPUReconstructionMetalInternals* internals, id<MTLCommandBuffer> cb, int32_t stream)
+{
+  std::unique_lock<std::mutex> lock(internals->inFlightMutex);
+  if (internals->streamBuffersInFlight[stream]++ == 0) {
+    internals->inFlightDone.wait(lock, [internals] { return internals->streamsInFlight < internals->maxStreamsInFlight; });
+    internals->streamsInFlight++;
+  }
+  [cb addCompletedHandler:^(id<MTLCommandBuffer>) {
+    std::lock_guard<std::mutex> l(internals->inFlightMutex);
+    if (--internals->streamBuffersInFlight[stream] == 0) {
+      internals->streamsInFlight--;
+      internals->inFlightDone.notify_all();
+    }
+  }];
+}
 // In the time base of the command buffers' GPUStartTime and GPUEndTime
 inline double metalHostTime()
 {

GPU/GPUTracking/Base/metal/GPUReconstructionMetalKernels.mm

@@ -65,6 +65,8 @@
   id<MTLCommandBuffer> cb = [mInternals->commandBuffers[kExec.stream] retain];
   if (kEvents.ev) {
     RecordMarker(kEvents.ev, kExec.stream);
+  } else if (mInternals->maxKernelsPerBuffer && ++mInternals->streamKernels[kExec.stream] >= mInternals->maxKernelsPerBuffer) {
+    CommitStream(kExec.stream);
   }
   if (timed) {
     if (mInternals->commandBuffers[kExec.stream] == cb) {