Metal: hold the GPU batch size fixed to stop MPSGraph recompile stalls

cpp/configs/analysis_example.cfg

@@ -260,6 +260,15 @@ nnRandomize = true
 #
 # Default (no config): 1 server thread, GPU-only mode (gpuIdx = 0).
 
+# Run every GPU (MPSGraph) evaluation at the max batch size, padding short batches.
+# MPSGraph compiles a new executable, which takes seconds and stalls the search, the
+# first time it sees each batch size. With padding there is one shape, and the default
+# nnMaxBatchSize becomes about half of numAnalysisThreads * numSearchThreadsPerAnalysisThread
+# so the padding stays cheap. The default (auto) pads only when nnMaxBatchSize is not set,
+# since padding every batch up to a size chosen for dynamic batching wastes GPU work; this
+# file sets nnMaxBatchSize above, so remove it to pad. ANE (CoreML) threads never pad.
+# metalPadBatch = auto
+
 # FP16 precision (default true). Currently this only affects the ANE/CoreML path:
 # the GPU/MPSGraph path always runs in FP32 regardless of this flag. Setting this
 # to false also disables ANE acceleration (the ANE is FP16-only hardware, so

cpp/configs/gtp_example.cfg

@@ -585,6 +585,14 @@ searchFactorWhenWinningThreshold = 0.95
 #
 # Default (no config): 1 server thread, GPU-only mode (gpuIdx = 0).
 
+# Run every GPU (MPSGraph) evaluation at the max batch size, padding short batches.
+# MPSGraph compiles a new executable, which takes seconds and stalls the search, the
+# first time it sees each batch size. With padding there is one shape, and the default
+# nnMaxBatchSize becomes about half of numSearchThreads so the padding stays cheap.
+# The default (auto) pads only when nnMaxBatchSize is not set, since padding every batch
+# up to a size chosen for dynamic batching wastes GPU work. ANE (CoreML) threads never pad.
+# metalPadBatch = auto
+
 # FP16 precision (default true). Currently this only affects the ANE/CoreML path:
 # the GPU/MPSGraph path always runs in FP32 regardless of this flag. Setting this
 # to false also disables ANE acceleration (the ANE is FP16-only hardware, so

cpp/neuralnet/metalbackend.cpp

@@ -493,11 +493,20 @@ metalContext(createMetalComputeContext(nnX, nnY, useFP16Mode != enabled_t::False
   this->useFP16Mode = useFP16Mode;
   this->nnXLen = nnX;
   this->nnYLen = nnY;
+  this->padBatch = true;
 }
 
 ComputeContext::~ComputeContext() {
 }
 
+// metalPadBatch: whether GPU (MPSGraph) batches run padded up to nnMaxBatchSize, see getBatchPolicy.
+static bool metalPadsBatch(ConfigParser& cfg) {
+  enabled_t padBatchMode = cfg.contains("metalPadBatch") ? cfg.getEnabled("metalPadBatch") : enabled_t::Auto;
+  if(padBatchMode == enabled_t::Auto)
+    return !cfg.contains("nnMaxBatchSize");
+  return padBatchMode == enabled_t::True;
+}
+
 ComputeContext* NeuralNet::createComputeContext(
   const vector<int>& gpuIdxs,
   Logger* logger,
@@ -524,10 +533,10 @@ ComputeContext* NeuralNet::createComputeContext(
   (void)logger;
   (void)homeDataDirOverride;
   (void)loadedModel;
-  (void)cfg;
 
   ComputeContext* context = new ComputeContext(nnXLen, nnYLen, useFP16Mode);
   context->aneOnly = aneOnly;
+  context->padBatch = metalPadsBatch(cfg);
   return context;
 }
 
@@ -676,6 +685,7 @@ mpsGraphOnlyHandle(createMPSGraphHandleIfNeeded(context, loadedModel, requireExa
 coremlOnlyHandle(createCoreMLOnlyHandleIfNeeded(context, loadedModel, requireExactNNLen, maxBatchSize, gpuIdx, serverThreadIdx)) {
   bool hasMPSGraph = static_cast<bool>(mpsGraphOnlyHandle);
   bool hasCoreML = static_cast<bool>(coremlOnlyHandle);
+  padBatch = hasMPSGraph && context->padBatch;
   if(hasMPSGraph == hasCoreML) {
     throw runtime_error("Metal backend: Logic error - expected exactly one compute handle, got " + string(hasMPSGraph && hasCoreML ? "both" : "neither") + " (gpuIdx=" + to_string(gpuIdx) + ")");
   }
@@ -775,8 +785,14 @@ std::string NeuralNet::getRuntimeBackendDetail(ConfigParser& cfg) {
 }
 
 NeuralNet::BatchPolicy NeuralNet::getBatchPolicy(ConfigParser& cfg) {
-  (void)cfg;
-  return NeuralNet::BatchPolicy::Dynamic;
+  // MPSGraph compiles a new executable the first time it runs at each batch size, which takes
+  // seconds and stalls every search thread waiting on the GPU. So by default (metalPadBatch) the
+  // GPU path holds the batch size fixed and pads short batches, and the batch is sized the
+  // fixed-shape way (about half the concurrent evaluations) to keep the padding cheap. An
+  // nnMaxBatchSize set in the config is normally sized for dynamic batching (analysis_example.cfg
+  // asks for at least the total number of search threads), and padding every batch up to it would
+  // multiply the GPU work, so by default such a config keeps dynamic batches.
+  return metalPadsBatch(cfg) ? NeuralNet::BatchPolicy::FixedShape : NeuralNet::BatchPolicy::Dynamic;
 }
 
 int NeuralNet::getNumEffectiveDevices(ConfigParser& cfg, const std::vector<int>& gpuIdxByServerThread) {
@@ -1186,6 +1202,24 @@ void MetalProcess::getMetalOutput(
       batchSize);
   } else if(gpuHandle->mpsGraphOnlyHandle) {
     // GPU mode: Use MPSGraph (GPU)
+    // Run at a constant number of rows and throw the extra ones away, see getBatchPolicy. The
+    // padding rows copy row 0 rather than keeping whatever an earlier call left there (or zeros on
+    // the first call): an all-zero row has an empty mask, which the mask-weighted pooling divides by.
+    const int runBatchSize = gpuHandle->padBatch ? inputBuffers->maxBatchSize : batchSize;
+    for(int row = batchSize; row < runBatchSize; row++) {
+      copy(
+        inputBuffers->userInputBuffer, inputBuffers->userInputBuffer + inputBuffers->singleSpatialElts,
+        inputBuffers->userInputBuffer + inputBuffers->singleSpatialElts * row);
+      copy(
+        inputBuffers->userInputGlobalBuffer, inputBuffers->userInputGlobalBuffer + inputBuffers->singleInputGlobalElts,
+        inputBuffers->userInputGlobalBuffer + inputBuffers->singleInputGlobalElts * row);
+      copy(
+        inputBuffers->userInputMetaBuffer, inputBuffers->userInputMetaBuffer + inputBuffers->singleInputMetaElts,
+        inputBuffers->userInputMetaBuffer + inputBuffers->singleInputMetaElts * row);
+      copy(
+        inputBuffers->userInputMaskBuffer, inputBuffers->userInputMaskBuffer + inputBuffers->singleMaskElts,
+        inputBuffers->userInputMaskBuffer + inputBuffers->singleMaskElts * row);
+    }
     gpuHandle->mpsGraphOnlyHandle.get().apply(
       inputBuffers->userInputBuffer,
       inputBuffers->userInputGlobalBuffer,
@@ -1195,7 +1229,7 @@ void MetalProcess::getMetalOutput(
       inputBuffers->valueResults,
       inputBuffers->scoreValuesResults,
       inputBuffers->ownershipResults,
-      batchSize);
+      runBatchSize);
   } else {
     throw runtime_error("Metal backend: No valid compute handle available");
   }

cpp/neuralnet/metalbackend.h

@@ -113,6 +113,11 @@ struct ComputeContext {
    */
   MetalComputeContext metalContext;
 
+  /**
+   * @brief metalPadBatch: whether GPU (MPSGraph) threads run every batch at the max batch size.
+   */
+  bool padBatch;
+
   /**
    * @brief True only when every configured device is METAL_MUX_ANE, so no
    * MPSGraph (GPU) handle will ever read modelDesc weights. Gates the call to
@@ -155,6 +160,11 @@ struct ComputeHandle {
    */
   int gpuIndex;
 
+  /**
+   * @brief Whether this handle runs every MPSGraph batch at the max batch size (see metalPadBatch).
+   */
+  bool padBatch;
+
   /**
    * @brief The version of the loaded model.
    */