GPU Metal: optionally fail when a kernel ran on an unmapped memory pool

GPU/GPUTracking/Base/metal/GPUReconstructionMETAL.metal

@@ -82,5 +82,14 @@ using namespace metal;
 
 #include "GPUReconstructionKernelList.h"
 
+// Reads a word of the memory pool through its address, as the kernels do, and
+// counts a read that does not return it; the host fails the run on a count.
+kernel void krnl_metalCanary(device atomic_uint* bad [[buffer(0)]], constant uint64_t& word [[buffer(1)]])
+{
+  if (*(device uint32_t*)word != 0xA11CE0A5u) {
+    atomic_fetch_add_explicit(bad, 1u, memory_order_relaxed);
+  }
+}
+
 // clang-format on
 #pragma clang diagnostic pop

GPU/GPUTracking/Base/metal/GPUReconstructionMetal.mm

@@ -185,7 +185,12 @@ uint32_t metalGPUCoreCount()
     // The kernels reconstruct device pointers from these bases, so they must be
     // the addresses the GPU sees, not the CPU-visible ones of a shared buffer.
     mDeviceMemoryBase = (void*)(uintptr_t)mInternals->mem_gpu.gpuAddress;
-    mDeviceMemorySize = mInternals->mem_gpu.allocatedSize;
+    mDeviceMemorySize = mInternals->mem_gpu.allocatedSize - sizeof(uint32_t);
+    *(uint32_t*)((char*)mInternals->mem_gpu.contents + mDeviceMemorySize) = mInternals->canaryMagic;
+    if (getenv("GPUCA_METAL_CANARY")) {
+      mInternals->canaryBad = [mInternals->device newBufferWithLength:sizeof(uint32_t) options:MTLResourceStorageModeShared];
+      *(uint32_t*)mInternals->canaryBad.contents = 0;
+    }
     mDeviceConstantMem = (GPUConstantMem*)(uintptr_t)mInternals->mem_constant.gpuAddress;
 
     if (GetProcessingSettings().debugLevel >= 1) {
@@ -233,6 +238,7 @@ uint32_t metalGPUCoreCount()
     if (mDeviceMemoryBase) {
       [mInternals->mem_gpu release];
       [mInternals->mem_constant release];
+      [mInternals->canaryPipeline release];
       for (uint32_t i = 0; i < mInternals->functions.size(); i++) {
         [mInternals->functions[i] release];
       }
@@ -252,6 +258,7 @@ uint32_t metalGPUCoreCount()
         [mInternals->streamFence[i] release];
       }
       [mInternals->mem_host release];
+      [mInternals->canaryBad release];
     }
     // The completion handlers of the limit still use the internals
     std::unique_lock<std::mutex> lock(mInternals->inFlightMutex);
@@ -336,6 +343,9 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
   if (ev) {
     RecordMarker(ev, realStream);
   }
+  if (stream == -1) {  // Synchronous like cudaMemcpy: the caller reads the destination right away
+    DrainStream(realStream);
+  }
 
   if (GetProcessingSettings().serializeGPU & 2) {
     GPUDebug(("GPUMemCpy " + std::to_string(toGPU)).c_str(), stream, true);
@@ -370,6 +380,9 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
   if (ev) {
     RecordMarker(ev, realStream);
   }
+  if (stream == -1) {
+    DrainStream(realStream);
+  }
 
   if (GetProcessingSettings().serializeGPU & 2) {
     GPUDebug("WriteToConstantMemory", stream, true);
@@ -470,6 +483,10 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
     [cb release];
   }
   mInternals->submittedBuffers[stream].clear();
+  if (const uint32_t bad = mInternals->canaryBad ? *(volatile uint32_t*)mInternals->canaryBad.contents : 0) {
+    GPUError("Metal: %u kernels ran while the memory pool was not mapped for the GPU", bad);
+    error = error ? error : 1;
+  }
   if (drainStart) {
     fprintf(mInternals->timeline, "DR s%d %.6f %.6f\n", stream, drainStart, metalHostTime());
   }
@@ -606,5 +623,14 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
     return 1;
   }
 
+  id<MTLFunction> canary = [mInternals->library newFunctionWithName:@"krnl_metalCanary"];
+  mInternals->canaryPipeline = canary ? [mInternals->device newComputePipelineStateWithFunction:canary error:&error] : nil;
+  [canary release];
+  if (mInternals->canaryPipeline == nil) {
+    NSLog(@"%@", error);
+    GPUError("Error creating the Metal canary pipeline state");
+    return 1;
+  }
+
   return AddKernels();
 }

GPU/GPUTracking/Base/metal/GPUReconstructionMetalIncludesHost.h

@@ -96,6 +96,12 @@ struct GPUReconstructionMetalInternals {
   id<MTLBuffer> mem_gpu;
   id<MTLBuffer> mem_constant;
   id<MTLBuffer> mem_host;
+  // Opt-in, GPUCA_METAL_CANARY=1: a kernel that reads a word past the end of the
+  // device pool, where the host put canaryMagic, and counts the reads that return
+  // something else. canaryBad is nil when it is off.
+  static constexpr uint32_t canaryMagic = 0xA11CE0A5;
+  id<MTLComputePipelineState> canaryPipeline;
+  id<MTLBuffer> canaryBad = nil;
 
   id<MTLLibrary> library;
 
@@ -156,6 +162,16 @@ inline void metalTimelineAdd(GPUReconstructionMetalInternals* internals, int32_t
     internals->timelineOps[stream] += buf;
   }
 }
+// Runs after each kernel in its encoder, so a kernel that ran while the pool
+// was not mapped for the GPU is counted and the run fails instead of going on
+// with the results it did not produce.
+inline void metalEncodeCanary(GPUReconstructionMetalInternals* internals, id<MTLComputeCommandEncoder> encoder, uint64_t word)
+{
+  [encoder setComputePipelineState:internals->canaryPipeline];
+  [encoder setBuffer:internals->canaryBad offset:0 atIndex:0];
+  [encoder setBytes:&word length:sizeof(word) atIndex:1];
+  [encoder dispatchThreadgroups:MTLSizeMake(1, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
+}
 template <class E>
 inline void metalEndEncoder(GPUReconstructionMetalInternals* internals, E encoder, int32_t stream)
 {

GPU/GPUTracking/Base/metal/GPUReconstructionMetalKernels.mm

@@ -42,6 +42,10 @@
   // Slots 0 to 2 are what GPUCA_KRNL_BUFFERS and GPUCA_KRNL_SECTOR_ARG declare
   // in the generated signature; the kernel's own arguments follow from slot 3.
   [computeEncoder setComputePipelineState:pso];
+  // The kernels reach the pools through the addresses kept in constant memory,
+  // which Metal cannot see, so they have to be made resident explicitly.
+  [computeEncoder useResource:mInternals->mem_gpu usage:MTLResourceUsageRead | MTLResourceUsageWrite];
+  [computeEncoder useResource:mInternals->mem_host usage:MTLResourceUsageRead | MTLResourceUsageWrite];
   [computeEncoder setBuffer:mInternals->mem_gpu offset:0 atIndex:0];
   [computeEncoder setBuffer:mInternals->mem_constant offset:0 atIndex:1];
   const int32_t iSector = runRange.index;
@@ -61,6 +65,9 @@
   }
   [computeEncoder dispatchThreadgroups:MTLSizeMake(kExec.nBlocks, 1, 1)
                  threadsPerThreadgroup:MTLSizeMake(threadGroupSize, 1, 1)];
+  if (mInternals->canaryBad) {
+    metalEncodeCanary(mInternals, computeEncoder, (uint64_t)mDeviceMemoryBase + mDeviceMemorySize);
+  }
   metalEndEncoder(mInternals, computeEncoder, kExec.stream);
   id<MTLCommandBuffer> cb = [mInternals->commandBuffers[kExec.stream] retain];
   if (kEvents.ev) {