GPU Metal: optionally limit the GPU work in flight

GPU/GPUTracking/Base/metal/GPUReconstructionMetal.mm

@@ -167,7 +167,13 @@ uint32_t metalGPUCoreCount()
       mInternals->streamOrderValue[i] = 0;
       mInternals->streamFence[i] = [mInternals->device newFence];
       mInternals->streamHasWork[i] = false;
+      mInternals->timelineKernels[i] = 0;
     }
+    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) {
@@ -212,7 +218,8 @@ uint32_t metalGPUCoreCount()
 
 int32_t GPUReconstructionMetal::ExitDevice_Runtime()
 {
-  // Tear down Metal
+  // Tear down Metal, also after a failure, without failing again
+  mInternals->exiting = true;
   SynchronizeGPU();
 
   for (uint32_t i = 0; i < mEvents.size(); i++) {
@@ -246,6 +253,12 @@ 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
+    }
 
     [mInternals->library release];
     [mInternals->device release];
@@ -312,6 +325,7 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
     id<MTLCommandBuffer> cb = mInternals->commandBuffers[realStream];
     id<MTLBlitCommandEncoder> blit = [cb blitCommandEncoder];
     metalBeginEncoder(mInternals, blit, realStream);
+    metalTimelineAdd(mInternals, realStream, "b", nullptr, sizeBytes);
     [blit copyFromBuffer:s.buffer
             sourceOffset:s.offset
                 toBuffer:d.buffer
@@ -345,6 +359,7 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
     id<MTLCommandBuffer> cb = mInternals->commandBuffers[realStream];
     id<MTLBlitCommandEncoder> blit = [cb blitCommandEncoder];
     metalBeginEncoder(mInternals, blit, realStream);
+    metalTimelineAdd(mInternals, realStream, "b", nullptr, size);
     [blit copyFromBuffer:s.buffer
             sourceOffset:s.offset
                 toBuffer:mInternals->mem_constant
@@ -380,7 +395,10 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
   // Does not change the retain count, so it's important we manage
   // the lifetime of the events outside here.
   auto mtlEvent = (__bridge id<MTLSharedEvent>)(ev->get<void*>());
-  [cb encodeSignalEvent:mtlEvent value:EventValue(*ev, true)];
+  const uint64_t value = EventValue(*ev, true);
+  [cb encodeSignalEvent:mtlEvent value:value];
+  metalSignalOnError(cb, mtlEvent, value);
+  metalTimelineAdd(mInternals, stream, "s", mtlEvent, value);
   CommitStream(stream);
 }
 
@@ -392,13 +410,32 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
   // in constant memory. So each buffer of a stream waits for the previous one.
   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;
+    const std::string ops = std::move(mInternals->timelineOps[stream]);
+    const uint32_t nKernels = mInternals->timelineKernels[stream];
+    const uint64_t seq = mInternals->timelineBuffers++;
+    const double commitTime = metalHostTime();
+    [cb addCompletedHandler:^(id<MTLCommandBuffer> c) {
+      fprintf(f, "CB %llu s%d st%d commit %.6f gpu %.6f %.6f nk %u:%s\n", (unsigned long long)seq, stream, (int)c.status, commitTime, c.GPUStartTime, c.GPUEndTime, nKernels, ops.c_str());
+    }];
+    mInternals->timelineOps[stream].clear();
+    mInternals->timelineKernels[stream] = 0;
+  }
   [cb commit];
   // Hand the committed buffer over to the stream's submission list, which owns
   // 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);
 }
 
 void GPUReconstructionMetal::DrainStream(int32_t stream)
@@ -417,19 +454,29 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
   // waitUntilCompleted, which sleeps until the completion has been processed.
   // Only short waits are worth a busy core, so give up after a while.
   bool done = false;
+  const double drainStart = mInternals->timeline && !mInternals->submittedBuffers[stream].empty() ? metalHostTime() : 0.;
   const auto t0 = std::chrono::steady_clock::now();
   while (!mInternals->submittedBuffers[stream].empty() && !(done = mInternals->streamOrder[stream].signaledValue >= mInternals->streamOrderValue[stream]) && std::chrono::steady_clock::now() - t0 < std::chrono::milliseconds(20)) {
   }
+  int64_t error = 0;
   for (id<MTLCommandBuffer> cb : mInternals->submittedBuffers[stream]) {
     if (!done) {
       [cb waitUntilCompleted];
     }
     if (cb.status == MTLCommandBufferStatusError) {
       GPUError("Metal command buffer on stream %d failed: %s", stream, [[cb.error description] UTF8String]);
+      error = error ? error : cb.error.code;
     }
     [cb release];
   }
   mInternals->submittedBuffers[stream].clear();
+  if (drainStart) {
+    fprintf(mInternals->timeline, "DR s%d %.6f %.6f\n", stream, drainStart, metalHostTime());
+  }
+  // The work of a failed buffer is lost, so the results are wrong: fail like the other backends.
+  if (error && !mInternals->exiting) {
+    GPUChkErr(error);
+  }
 }
 
 void GPUReconstructionMetal::SynchronizeGPU()
@@ -453,12 +500,18 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
 
 void GPUReconstructionMetal::StreamWaitForEvents(int32_t stream, deviceEvent* evList, int32_t nEvents)
 {
+  // Encoded after the stream's pending work, the waits would keep its command buffer
+  // from completing until the other streams are done: wait in a buffer of its own.
+  if (mInternals->streamHasWork[stream]) {
+    CommitStream(stream);
+  }
   // Encode commands to wait for all the events
   id<MTLCommandBuffer> cb = mInternals->commandBuffers[stream];
   for (int32_t ei = 0; ei < nEvents; ei++) {
     auto mtlEvent = (__bridge id<MTLSharedEvent>)(evList[ei].get<void*>());
     if (const uint64_t value = EventValue(evList[ei], false)) {
       [cb encodeWaitForEvent:mtlEvent value:value];
+      metalTimelineAdd(mInternals, stream, "w", mtlEvent, value);
     }
   }
   CommitStream(stream);

GPU/GPUTracking/Base/metal/GPUReconstructionMetalIncludesHost.h

@@ -33,7 +33,10 @@ using half = _Float16;
 #include <string>
 #include <memory>
 #include <mutex>
+#include <condition_variable>
 #include <unordered_map>
+#include <cstdio>
+#include <mach/mach_time.h>
 #include "GPULogging.h"
 
 #include "GPUReconstructionMetal.h"
@@ -66,6 +69,25 @@ struct GPUReconstructionMetalInternals {
   // collects the stream's work until something has to wait for it.
   std::array<id<MTLFence>, constants::GPU_MAX_STREAMS> streamFence;
   std::array<bool, constants::GPU_MAX_STREAMS> streamHasWork;
+  // Set while the device is torn down, when a failed command buffer is no longer fatal
+  bool exiting = false;
+  // Opt-in command-buffer timeline, GPUCA_METAL_TIMELINE=<file>: what each open
+  // buffer contains (kernels, blits, waits, signals), written out once it completes
+  FILE* timeline = nullptr;
+  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;
@@ -90,6 +112,50 @@ inline void metalBeginEncoder(GPUReconstructionMetalInternals* internals, E enco
 {
   [encoder waitForFence:internals->streamFence[stream]];
 }
+// A failed command buffer does not signal its events, and the work waiting for
+// them would wait forever: signal them from the host, so that everything
+// completes and the failure is reported when the host synchronises.
+inline void metalSignalOnError(id<MTLCommandBuffer> cb, id<MTLSharedEvent> event, uint64_t value)
+{
+  [cb addCompletedHandler:^(id<MTLCommandBuffer> c) {
+    if (c.status == MTLCommandBufferStatusError && event.signaledValue < value) {
+      event.signaledValue = value;
+    }
+  }];
+}
+// 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()
+{
+  static const double scale = [] { mach_timebase_info_data_t tb; mach_timebase_info(&tb); return 1e-9 * tb.numer / tb.denom; }();
+  return mach_absolute_time() * scale;
+}
+inline void metalTimelineAdd(GPUReconstructionMetalInternals* internals, int32_t stream, const char* op, const void* id, uint64_t value)
+{
+  if (internals->timeline) {
+    char buf[96];
+    snprintf(buf, sizeof(buf), " %s:%p=%llu", op, id, (unsigned long long)value);
+    internals->timelineOps[stream] += buf;
+  }
+}
 template <class E>
 inline void metalEndEncoder(GPUReconstructionMetalInternals* internals, E encoder, int32_t stream)
 {

GPU/GPUTracking/Base/metal/GPUReconstructionMetalKernels.mm

@@ -29,6 +29,11 @@
   if (kEvents.evList) {
     StreamWaitForEvents(kExec.stream, kEvents.evList, kEvents.nEvents);
   }
+  // Device timers: the kernel gets a command buffer of its own, whose GPU time is the kernel's
+  const bool timed = GetProcessingSettings().deviceTimers && GetProcessingSettings().debugLevel > 0;
+  if (timed && mInternals->streamHasWork[kExec.stream]) {
+    CommitStream(kExec.stream);
+  }
 
   id<MTLComputePipelineState> pso = mInternals->pipelines[num];
   id<MTLComputeCommandEncoder> computeEncoder = [mInternals->commandBuffers[kExec.stream] computeCommandEncoder];
@@ -50,12 +55,28 @@
   if (threadGroupSize > mInternals->maxThreadsPerGroup[num]) {
     threadGroupSize = mInternals->maxThreadsPerGroup[num];
   }
+  if (mInternals->timeline) {
+    mInternals->timelineOps[kExec.stream] += std::string(" k:") + GetKernelName<T, I>() + "/" + std::to_string(kExec.nBlocks);
+    mInternals->timelineKernels[kExec.stream]++;
+  }
   [computeEncoder dispatchThreadgroups:MTLSizeMake(kExec.nBlocks, 1, 1)
                  threadsPerThreadgroup:MTLSizeMake(threadGroupSize, 1, 1)];
   metalEndEncoder(mInternals, computeEncoder, kExec.stream);
+  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) {
+      CommitStream(kExec.stream);
+    }
+    DrainStream(kExec.stream);
+    [cb waitUntilCompleted];  // the times are only set once the completion is processed
+    _xyz.t = cb.GPUEndTime - cb.GPUStartTime;
   }
+  [cb release];
 }
 
 template <class T, int32_t I>