GPU Metal: collect the work of a stream in one command buffer

Every kernel, memory clear and copy was a command buffer of its own,
2,689 per event on a 5-event pp time frame, each committed at once and,
since they run in order, each waiting for the previous one's event. Keep
encoding into the stream's open buffer instead, ordered by a per-stream
MTLFence between encoders, and commit it only where something has to
wait for it: draining the stream, recording an event, waiting for
events. This makes 535 command buffers per event, and the time per
event about 3 to 6% less on an M1 Max.

Author: ktf

GPU/GPUTracking/Base/metal/GPUReconstructionMetal.mm

@@ -138,6 +138,8 @@ uint32_t metalGPUCoreCount()
       }
       mInternals->streamOrder[i] = [mInternals->device newSharedEvent];
       mInternals->streamOrderValue[i] = 0;
+      mInternals->streamFence[i] = [mInternals->device newFence];
+      mInternals->streamHasWork[i] = false;
     }
 
     mInternals->mem_host = [mInternals->device newBufferWithLength:mHostMemorySize options:MTLResourceStorageModeShared];
@@ -213,6 +215,7 @@ uint32_t metalGPUCoreCount()
         [mInternals->commandBuffers[i] release];
         [mInternals->commandQueues[i] release];
         [mInternals->streamOrder[i] release];
+        [mInternals->streamFence[i] release];
       }
       [mInternals->mem_host release];
     }
@@ -281,13 +284,13 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
   } else {
     id<MTLCommandBuffer> cb = mInternals->commandBuffers[realStream];
     id<MTLBlitCommandEncoder> blit = [cb blitCommandEncoder];
+    metalBeginEncoder(mInternals, blit, realStream);
     [blit copyFromBuffer:s.buffer
             sourceOffset:s.offset
                 toBuffer:d.buffer
        destinationOffset:d.offset
                     size:sizeBytes];
-    [blit endEncoding];
-    CommitStream(realStream);
+    metalEndEncoder(mInternals, blit, realStream);
   }
   if (ev) {
     RecordMarker(ev, realStream);
@@ -314,13 +317,13 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
   } else {
     id<MTLCommandBuffer> cb = mInternals->commandBuffers[realStream];
     id<MTLBlitCommandEncoder> blit = [cb blitCommandEncoder];
+    metalBeginEncoder(mInternals, blit, realStream);
     [blit copyFromBuffer:s.buffer
             sourceOffset:s.offset
                 toBuffer:mInternals->mem_constant
        destinationOffset:offset
                     size:size];
-    [blit endEncoding];
-    CommitStream(realStream);
+    metalEndEncoder(mInternals, blit, realStream);
   }
   if (ev) {
     RecordMarker(ev, realStream);
@@ -366,12 +369,16 @@ MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const
   // 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->commandBuffers[stream] = [[mInternals->commandQueues[stream] commandBuffer] retain];
   [mInternals->commandBuffers[stream] encodeWaitForEvent:mInternals->streamOrder[stream] value:value];
 }
 
 void GPUReconstructionMetal::DrainStream(int32_t stream)
 {
+  if (mInternals->streamHasWork[stream]) {
+    CommitStream(stream);
+  }
   // Waiting on the buffer that is currently open would block forever: it is not
   // committed, so it never completes. Only what was submitted can be waited for.
   for (id<MTLCommandBuffer> cb : mInternals->submittedBuffers[stream]) {

GPU/GPUTracking/Base/metal/GPUReconstructionMetalIncludesHost.h

@@ -62,6 +62,10 @@ struct GPUReconstructionMetalInternals {
   // Signalled by each committed buffer and waited for by the next one, see CommitStream
   std::array<id<MTLSharedEvent>, constants::GPU_MAX_STREAMS> streamOrder;
   std::array<uint64_t, constants::GPU_MAX_STREAMS> streamOrderValue;
+  // Orders the encoders within the open command buffer of a stream, which
+  // 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;
 
   std::vector<id<MTLFunction>> functions;
   std::vector<id<MTLComputePipelineState>> pipelines;
@@ -79,6 +83,20 @@ struct GPUReconstructionMetalInternals {
   std::unordered_map<void*, uint64_t> eventValues;
   std::mutex eventMutex;
 };
+// An encoder of a stream starts after the previous one of the same stream and
+// leaves its work in the open command buffer.
+template <class E>
+inline void metalBeginEncoder(GPUReconstructionMetalInternals* internals, E encoder, int32_t stream)
+{
+  [encoder waitForFence:internals->streamFence[stream]];
+}
+template <class E>
+inline void metalEndEncoder(GPUReconstructionMetalInternals* internals, E encoder, int32_t stream)
+{
+  [encoder updateFence:internals->streamFence[stream]];
+  [encoder endEncoding];
+  internals->streamHasWork[stream] = true;
+}
 } // namespace o2::gpu
 
 #endif // GPURECONSTRUCTIONOMETALINCLUDESHOST_H

GPU/GPUTracking/Base/metal/GPUReconstructionMetalKernels.mm

@@ -32,6 +32,7 @@
 
   id<MTLComputePipelineState> pso = mInternals->pipelines[num];
   id<MTLComputeCommandEncoder> computeEncoder = [mInternals->commandBuffers[kExec.stream] computeCommandEncoder];
+  metalBeginEncoder(mInternals, computeEncoder, kExec.stream);
 
   // 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.
@@ -51,10 +52,7 @@
   }
   [computeEncoder dispatchThreadgroups:MTLSizeMake(kExec.nBlocks, 1, 1)
                  threadsPerThreadgroup:MTLSizeMake(threadGroupSize, 1, 1)];
-  // Without ending the encoder the next one on this buffer is rejected, and
-  // without committing the kernel never reaches the device.
-  [computeEncoder endEncoding];
-  CommitStream(kExec.stream);
+  metalEndEncoder(mInternals, computeEncoder, kExec.stream);
   if (kEvents.ev) {
     RecordMarker(kEvents.ev, kExec.stream);
   }

GPU/GPUTracking/Base/metal/GPUReconstructionMetalKernelsSpecialize.inc

@@ -22,11 +22,10 @@ inline void GPUReconstructionMetal::runKernelBackend<GPUMemClean16, 0>(const krn
 
   id<MTLCommandBuffer> cb = mInternals->commandBuffers[_xyz.x.stream];
   id<MTLBlitCommandEncoder> blit = [cb blitCommandEncoder];
+  metalBeginEncoder(mInternals, blit, _xyz.x.stream);
 
   [blit fillBuffer:mInternals->mem_gpu range:NSMakeRange(offset, length) value:0];
-  [blit endEncoding];
-
-  CommitStream(_xyz.x.stream);
+  metalEndEncoder(mInternals, blit, _xyz.x.stream);
   if (_xyz.z.ev) {
     RecordMarker(_xyz.z.ev, _xyz.x.stream);
   }