@@ -57,6 +57,9 @@ class GPUReconstructionMetal : public GPUReconstructionProcessing::KernelInterfa
 
   int32_t GetMetalPrograms();
 
+  void CommitStream(int32_t stream);
+  void DrainStream(int32_t stream);
+
  private:
   int32_t AddKernels();
 };
 
@@ -9,10 +9,11 @@
 // granted to it by virtue of its status as an Intergovernmental Organization
 // or submit itself to any jurisdiction.
 
-#include "GPUReconstructionMetal.h"
+#include "GPUReconstructionMetalIncludesHost.h"  // must come first
+
 #include "GPUConstantMem.h"
 #include "GPUDefParametersLoad.inc"
-#include "GPUReconstructionMetalIncludesHost.h"
+#include "GPUReconstructionMetal.h"
 
 #include <cstddef>
 #include <map>
@@ -100,7 +101,7 @@
       if (mInternals->commandQueues[i] == nil) {
         GPUErrorReturn("Error creating Metal command queue");
       }
-      mInternals->commandBuffers[i] = [mInternals->commandQueues[i] commandBuffer];
+      mInternals->commandBuffers[i] = [[mInternals->commandQueues[i] commandBuffer] retain];
       if (mInternals->commandBuffers[i] == nil) {
         GPUErrorReturn("Error creating Metal command buffer");
       }
@@ -113,9 +114,11 @@
 
     mHostMemoryBase = mInternals->mem_host.contents;
     mHostMemorySize = mInternals->mem_host.allocatedSize;
-    mDeviceMemoryBase = mInternals->mem_gpu.contents;
+    // 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;
-    mDeviceConstantMem = (GPUConstantMem*)mInternals->mem_constant.contents;
+    mDeviceConstantMem = (GPUConstantMem*)(uintptr_t)mInternals->mem_constant.gpuAddress;
 
     if (GetProcessingSettings().debugLevel >= 1) {
       GPUInfo("Memory ptrs: GPU (%ld bytes): %p - Host (%ld bytes): %p", (int64_t)mDeviceMemorySize, mDeviceMemoryBase, (int64_t)mHostMemorySize, mHostMemoryBase);
@@ -153,11 +156,17 @@
         [mInternals->functions[i] release];
       }
       mInternals->functions.clear();
+      for (uint32_t i = 0; i < mInternals->pipelines.size(); i++) {
+        [mInternals->pipelines[i] release];
+      }
+      mInternals->pipelines.clear();
+      mInternals->maxThreadsPerGroup.clear();
     }
     if (mHostMemoryBase) {
       for (int32_t i = 0; i < mNStreams; i++) {
-        [mInternals->commandQueues[i] release];
+        DrainStream(i);
         [mInternals->commandBuffers[i] release];
+        [mInternals->commandQueues[i] release];
       }
       [mInternals->mem_host release];
     }
@@ -172,6 +181,31 @@
   return (0);
 }
 
+namespace
+{
+// Where a pointer lives as far as the backend is concerned. mDeviceMemoryBase is
+// a GPU address, so the CPU-visible address of owned memory has to be taken from
+// the buffer rather than from the pointer itself.
+struct MetalMemRef {
+  id<MTLBuffer> buffer;  // nil for memory the backend does not own
+  size_t offset;
+  char* cpu;
+};
+
+MetalMemRef metalResolve(const GPUReconstructionMetalInternals* internals, const void* p, const void* hostBase, size_t hostSize, const void* devBase, size_t devSize)
+{
+  if (p >= hostBase && p < ((const char*)hostBase + hostSize)) {
+    const size_t o = (const char*)p - (const char*)hostBase;
+    return {internals->mem_host, o, (char*)internals->mem_host.contents + o};
+  }
+  if (p >= devBase && p < ((const char*)devBase + devSize)) {
+    const size_t o = (const char*)p - (const char*)devBase;
+    return {internals->mem_gpu, o, (char*)internals->mem_gpu.contents + o};
+  }
+  return {nil, 0, (char*)p};
+}
+}  // namespace
+
 size_t GPUReconstructionMetal::GPUMemCpy(void* dst, const void* src, size_t sizeBytes, int32_t stream, int32_t toGPU, deviceEvent* ev, deviceEvent* evList, int32_t nEvents)
 {
   if (evList == nullptr) {
@@ -186,42 +220,27 @@
   }
 
   auto realStream = stream == -1 ? 0 : stream;
-  id<MTLCommandBuffer> cb = mInternals->commandBuffers[realStream];
-  id<MTLBlitCommandEncoder> blit = [cb blitCommandEncoder];
-  id<MTLBuffer> sourceBuffer = nil;
-  id<MTLBuffer> destBuffer = nil;
-  ptrdiff_t sourceOffset = 0;
-  ptrdiff_t destOffset = 0;
-
-  if (src >= mHostMemoryBase && src < ((char*)mHostMemoryBase + mHostMemorySize)) {
-    sourceBuffer = mInternals->mem_host;
-    sourceOffset = (char*)src - (char*)mHostMemoryBase;
-  } else if (src >= mDeviceMemoryBase && src < ((char*)mDeviceMemoryBase + mDeviceMemorySize)) {
-    sourceBuffer = mInternals->mem_gpu;
-    sourceOffset = (char*)src - (char*)mDeviceMemoryBase;
+  const MetalMemRef s = metalResolve(mInternals, src, mHostMemoryBase, mHostMemorySize, mDeviceMemoryBase, mDeviceMemorySize);
+  const MetalMemRef d = metalResolve(mInternals, dst, mHostMemoryBase, mHostMemorySize, mDeviceMemoryBase, mDeviceMemorySize);
+
+  if (s.buffer == nil || d.buffer == nil) {
+    // A blit moves data between MTLBuffers only, and the calibration objects
+    // live in ordinary memory. Both pools are shared storage on a unified
+    // architecture, so the copy can be done on the host once the stream is idle.
+    DrainStream(realStream);
+    memcpy(d.cpu, s.cpu, sizeBytes);
   } else {
-    GPUErrorReturn("Unknown buffer at %p", src);
+    id<MTLCommandBuffer> cb = mInternals->commandBuffers[realStream];
+    id<MTLBlitCommandEncoder> blit = [cb blitCommandEncoder];
+    [blit copyFromBuffer:s.buffer
+            sourceOffset:s.offset
+                toBuffer:d.buffer
+       destinationOffset:d.offset
+                    size:sizeBytes];
+    [blit endEncoding];
+    CommitStream(realStream);
   }
 
-  if (dst >= mHostMemoryBase && dst < ((char*)mHostMemoryBase + mHostMemorySize)) {
-    destBuffer = mInternals->mem_host;
-    destOffset = (char*)dst - (char*)mHostMemoryBase;
-  } else if (dst >= mDeviceMemoryBase && dst < ((char*)mDeviceMemoryBase + mDeviceMemorySize)) {
-    destBuffer = mInternals->mem_gpu;
-    destOffset = (char*)dst - (char*)mDeviceMemoryBase;
-  } else {
-    GPUErrorReturn("Unknown buffer at %p", dst);
-  }
-
-  [blit copyFromBuffer:sourceBuffer
-          sourceOffset:sourceOffset
-              toBuffer:destBuffer
-     destinationOffset:destOffset
-                  size:sizeBytes];
-
-  [blit endEncoding];
-  [cb commit];
-
   if (GetProcessingSettings().serializeGPU & 2) {
     GPUDebug(("GPUMemCpy " + std::to_string(toGPU)).c_str(), stream, true);
   }
@@ -235,27 +254,22 @@
   }
 
   auto realStream = stream == -1 ? 0 : stream;
-  id<MTLCommandBuffer> cb = mInternals->commandBuffers[realStream];
-  id<MTLBlitCommandEncoder> blit = [cb blitCommandEncoder];
-  id<MTLBuffer> sourceBuffer = nil;
-  ptrdiff_t sourceOffset = 0;
-  if (src >= mHostMemoryBase && src < ((char*)mHostMemoryBase + mHostMemorySize)) {
-    sourceBuffer = mInternals->mem_host;
-    sourceOffset = (char*)src - (char*)mHostMemoryBase;
-  } else if (src >= mDeviceMemoryBase && src < ((char*)mDeviceMemoryBase + mDeviceMemorySize)) {
-    sourceBuffer = mInternals->mem_gpu;
-    sourceOffset = (char*)src - (char*)mDeviceMemoryBase;
+  const MetalMemRef s = metalResolve(mInternals, src, mHostMemoryBase, mHostMemorySize, mDeviceMemoryBase, mDeviceMemorySize);
+
+  if (s.buffer == nil) {
+    DrainStream(realStream);
+    memcpy((char*)mInternals->mem_constant.contents + offset, s.cpu, size);
   } else {
-    GPUErrorReturn("Unknown buffer at %p", src);
+    id<MTLCommandBuffer> cb = mInternals->commandBuffers[realStream];
+    id<MTLBlitCommandEncoder> blit = [cb blitCommandEncoder];
+    [blit copyFromBuffer:s.buffer
+            sourceOffset:s.offset
+                toBuffer:mInternals->mem_constant
+       destinationOffset:offset
+                    size:size];
+    [blit endEncoding];
+    CommitStream(realStream);
   }
-  [blit copyFromBuffer:sourceBuffer
-          sourceOffset:sourceOffset
-              toBuffer:mInternals->mem_constant
-     destinationOffset:offset
-                  size:size];
-
-  [blit endEncoding];
-  [cb commit];
 
   if (GetProcessingSettings().serializeGPU & 2) {
     GPUDebug("WriteToConstantMemory", stream, true);
@@ -278,26 +292,46 @@
   // the lifetime of the events outside here.
   auto mtlEvent = (__bridge id<MTLSharedEvent>)(ev->get<void*>());
   [cb encodeSignalEvent:mtlEvent value:1];
+  CommitStream(stream);
+}
+
+void GPUReconstructionMetal::CommitStream(int32_t stream)
+{
+  id<MTLCommandBuffer> cb = mInternals->commandBuffers[stream];
   [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->commandBuffers[stream] = [[mInternals->commandQueues[stream] commandBuffer] retain];
+}
+
+void GPUReconstructionMetal::DrainStream(int32_t 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]) {
+    [cb waitUntilCompleted];
+    [cb release];
+  }
+  mInternals->submittedBuffers[stream].clear();
 }
 
 void GPUReconstructionMetal::SynchronizeGPU()
 {
   for (int32_t i = 0; i < mNStreams; i++) {
-    [mInternals->commandBuffers[i] waitUntilCompleted];
+    DrainStream(i);
   }
 }
 
 void GPUReconstructionMetal::SynchronizeStream(int32_t stream)
 {
-  [mInternals->commandBuffers[stream] waitUntilCompleted];
+  DrainStream(stream);
 }
 
 void GPUReconstructionMetal::SynchronizeEvents(deviceEvent* evList, int32_t nEvents)
 {
   for (int32_t si = 0; si < mNStreams; si++) {
-    id<MTLCommandBuffer> cb = mInternals->commandBuffers[si];
-    [cb waitUntilCompleted];
+    DrainStream(si);
   }
 }
 
@@ -309,8 +343,8 @@
     auto mtlEvent = (__bridge id<MTLSharedEvent>)(evList[ei].get<void*>());
     [cb encodeWaitForEvent:mtlEvent value:1];
   }
-  [cb commit];
-  [cb waitUntilCompleted];
+  CommitStream(stream);
+  DrainStream(stream);
 }
 
 bool GPUReconstructionMetal::IsEventDone(deviceEvent* evList, int32_t nEvents)
@@ -331,7 +365,7 @@
     return (0);
   }
   for (int32_t si = 0; si < mNStreams; si++) {
-    [mInternals->commandBuffers[si] waitUntilCompleted];
+    DrainStream(si);
   }
   if (GetProcessingSettings().debugLevel >= 3) {
     GPUInfo("GPU Sync Done");
 
@@ -12,12 +12,19 @@
 #ifndef GPURECONSTRUCTIONOMETALINCLUDESHOST_H
 #define GPURECONSTRUCTIONOMETALINCLUDESHOST_H
 
+#ifdef GPUCOMMONDEF_H
+#error "this header establishes GPUCA_GPUTYPE_METAL and __METAL_HOST__, so it must be included before any other GPU header"
+#endif
+
 #define GPUCA_GPUTYPE_METAL
 
 #import <Metal/Metal.h>
 #ifndef __METAL_VERSION__
 // __METAL_VERSION__ is only defined in device code.
 #define __METAL_HOST__
+// The parameter table picks MSL's `half` for the device layouts this translation
+// unit has to agree with; _Float16 is the same IEEE binary16 type.
+using half = _Float16;
 #endif
 
 #import <Foundation/Foundation.h>
@@ -45,10 +52,15 @@ struct GPUReconstructionMetalInternals {
   id<MTLDevice> device;
 
   std::array<id<MTLCommandQueue>, constants::GPU_MAX_STREAMS> commandQueues;
+  // The buffer currently open for encoding, plus the ones already committed and
+  // not yet waited for. A Metal command buffer takes one commit and cannot be
+  // encoded into afterwards, so the two cannot be the same object.
   std::array<id<MTLCommandBuffer>, constants::GPU_MAX_STREAMS> commandBuffers;
+  std::array<std::vector<id<MTLCommandBuffer>>, constants::GPU_MAX_STREAMS> submittedBuffers;
 
   std::vector<id<MTLFunction>> functions;
   std::vector<id<MTLComputePipelineState>> pipelines;
+  std::vector<uint32_t> maxThreadsPerGroup;
 
   id<MTLBuffer> mem_gpu;
   id<MTLBuffer> mem_constant;
 
@@ -9,9 +9,10 @@
 // granted to it by virtue of its status as an Intergovernmental Organization
 // or submit itself to any jurisdiction.
 
+#include "GPUReconstructionMetalIncludesHost.h"  // must come first
+
 #include <Foundation/Foundation.h>
 #include "GPUReconstructionKernelIncludes.h"
-#include "GPUReconstructionMetalIncludesHost.h"
 
 #include "GPUReconstructionMetalKernelsSpecialize.inc"
 #include "GPUReconstructionProcessingKernels.inc"
@@ -21,30 +22,35 @@
 template <class T, int32_t I, typename... Args>
 inline void GPUReconstructionMetal::runKernelBackend(const krnlSetupTime& _xyz, const Args&... args)
 {
-  id<MTLFunction> function = mInternals->functions[GetKernelNum<T, I>()];
+  const uint32_t num = GetKernelNum<T, I>();
   auto& kExec = _xyz.x;
   auto& runRange = _xyz.y;
 
-  NSError* error = nil;
-  auto pso = [mInternals->device newComputePipelineStateWithFunction:function error:&error];
+  id<MTLComputePipelineState> pso = mInternals->pipelines[num];
   id<MTLComputeCommandEncoder> computeEncoder = [mInternals->commandBuffers[kExec.stream] computeCommandEncoder];
 
-  // Map buffers and states
+  // 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];
   [computeEncoder setBuffer:mInternals->mem_gpu offset:0 atIndex:0];
   [computeEncoder setBuffer:mInternals->mem_constant offset:0 atIndex:1];
-  [computeEncoder setBuffer:mInternals->mem_host offset:0 atIndex:2];
-
-  MTLSize gridSize = MTLSizeMake(runRange.index, 1, 1);
+  const int32_t iSector = runRange.index;
+  [computeEncoder setBytes:&iSector length:sizeof(iSector) atIndex:2];
+  NSUInteger argIndex = 3;
+  ((void)[&](const auto& arg) { [computeEncoder setBytes:&arg length:sizeof(arg) atIndex:argIndex++]; }(args), ...);
 
-  NSUInteger threadGroupSize = pso.maxTotalThreadsPerThreadgroup;
-  if (threadGroupSize > runRange.index) {
-    threadGroupSize = runRange.index;
+  // runKernelInterface has already resolved the grid, so nBlocks and nThreads
+  // carry the tuned launch bounds and map onto Metal's threadgroups directly.
+  NSUInteger threadGroupSize = kExec.nThreads;
+  if (threadGroupSize > mInternals->maxThreadsPerGroup[num]) {
+    threadGroupSize = mInternals->maxThreadsPerGroup[num];
   }
-
-  MTLSize threadgroupSize = MTLSizeMake(threadGroupSize, 1, 1);
-  [computeEncoder dispatchThreads:gridSize
-            threadsPerThreadgroup:threadgroupSize];
+  [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);
 }
 
 template <class T, int32_t I>
@@ -58,7 +64,26 @@
     return 1;
   }
 
+  NSError* error = nil;
+  id<MTLComputePipelineState> pso = [mInternals->device newComputePipelineStateWithFunction:krnl error:&error];
+  if (pso == nil) {
+    NSLog(@"%@", error);
+    GPUError("Error creating Metal pipeline state: %s", [kname cStringUsingEncoding:NSUTF8StringEncoding]);
+    return 1;
+  }
+
+  // The kernel sizes its threadgroup memory from the tuning table at compile
+  // time, so a launch must not exceed what it was built for either way.
+  const uint32_t tunedThreads = this->template getKernelProperties<T, I>().nThreads;
+  if (tunedThreads > pso.maxTotalThreadsPerThreadgroup) {
+    // Fewer threads than the kernel was built for: safe for the shared arrays,
+    // but the scan and compaction kernels assume every slot has an owner.
+    GPUError("Metal kernel %s is tuned for %u threads but this device allows only %lu per threadgroup", GetKernelName<T, I>(), tunedThreads, (unsigned long)pso.maxTotalThreadsPerThreadgroup);
+  }
+
   mInternals->functions.emplace_back(krnl);
+  mInternals->pipelines.emplace_back(pso);
+  mInternals->maxThreadsPerGroup.emplace_back(pso.maxTotalThreadsPerThreadgroup);
   return 0;
 }