Reduce memcpys in substruct search and improve benchmark autotuning

benchmarks/substruct_bench.py

@@ -335,11 +335,9 @@ def main():
         default=42,
         help="Random seed for sampling SMILES when --num_mols > 0 (default: 42)",
     )
-    parser.add_argument("--sanitize", action="store_true", dest="sanitize", help="Sanitize SMILES during parsing")
-    parser.add_argument(
-        "--no_sanitize", action="store_false", dest="sanitize", help="Skip sanitization (preprocessed SMILES)"
-    )
-    parser.set_defaults(sanitize=False)
+    sanitize_group = parser.add_mutually_exclusive_group()
+    sanitize_group.add_argument("--sanitize", action="store_true", dest="sanitize", default=True)
+    sanitize_group.add_argument("--no_sanitize", action="store_false", dest="sanitize")
     parser.add_argument("--runs", "-r", type=int, default=1, help="Number of timing runs (default: 1)")
     parser.add_argument(
         "--mode",

nvmolkit/autotune/tune_substructure.py

@@ -64,8 +64,8 @@ def _default_substruct_search_space(num_gpus: int, cpus: int) -> dict:
     at trial sampling time by clamping ``preprocessingThreads`` to whatever
     cores are left after ``workerThreads`` is chosen.
 
-    * ``batchSize``: stepped int range in multiples of 128 (kernels are
-      tile-tuned for these sizes); stepping preserves numeric ordering for TPE.
+    * ``batchSize``: target-query pairs per mini-batch, log-uniform since the
+      useful range spans two orders of magnitude.
     * ``workerThreads`` (per-GPU): max = ``min(8, cpus // num_gpus)``.
       8 is the empirical point of diminishing returns; the physical-core
       floor prevents oversubscribing across GPUs.
@@ -74,7 +74,7 @@ def _default_substruct_search_space(num_gpus: int, cpus: int) -> dict:
     """
     per_gpu_worker_max = max(1, min(8, cpus // max(1, num_gpus)))
     return {
-        "batchSize": (128, 1024, 128),
+        "batchSize": (128, 8192, "log"),
         "workerThreads": (1, per_gpu_worker_max),
         "preprocessingThreads": (1, cpus),
     }
@@ -121,7 +121,7 @@ def tune_substructure(
     gpuIds: Optional[Iterable[int]] = None,
     calibration_set: Optional[Iterable[int]] = None,
     calibration_fraction: float = 0.1,
-    calibration_max_size: int = 2000,
+    calibration_max_size: int = 50_000,
     target_seconds_per_trial: float = 10.0,
     n_trials: int = 30,
     search_space_overrides: Optional[dict[str, Any]] = None,

nvmolkit/tests/test_autotune.py

@@ -289,10 +289,11 @@ def test_default_substruct_search_space_caps_per_pool():
     assert space_4gpu["preprocessingThreads"] == (1, 16)
     assert space_64gpu["preprocessingThreads"] == (1, 16)
 
-    low, high, step = space_1gpu["batchSize"]
-    assert step % 64 == 0
-    assert low % step == 0 and high % step == 0
-    assert low <= high
+    # Unlike the FF and embed spaces, substruct batchSize is log-uniform rather
+    # than a stepped range: the useful span is two orders of magnitude.
+    low, high, distribution = space_1gpu["batchSize"]
+    assert distribution == "log"
+    assert 0 < low <= high
 
 
 class _RecordingTrial:

src/substruct/recursive_preprocessor.cu

@@ -167,6 +167,18 @@ void LeafSubpatterns::syncToDevice(cudaStream_t stream) {
     return;
   }
   patternsDevice.copyFromHost(patternsHost, stream);
+
+  // Upload the immutable pattern table once.
+  for (int depth = 0; depth <= kMaxSmartsNestingDepth; ++depth) {
+    const auto& entries      = allQueriesPatternsAtDepth[depth];
+    auto&       entriesOnGpu = allQueriesPatternsAtDepthDevice[depth];
+    entriesOnGpu.setStream(stream);
+    if (entries.empty()) {
+      continue;
+    }
+    entriesOnGpu.resize(entries.size());
+    entriesOnGpu.copyFromHost(entries.data(), entries.size());
+  }
 }
 
 void RecursivePatternPreprocessor::buildPatterns(const MoleculesHost& queriesHost) {
@@ -230,14 +242,9 @@ void RecursivePatternPreprocessor::preprocessMiniBatch(
       const size_t numPatternsInSubBatch = patternEnd - patternStart;
       const size_t numBlocksInSubBatch   = numTargetsInMiniBatch * numPatternsInSubBatch;
 
-      ScopedNvtxRange prepareRange("GPU: Upload pattern entries");
-      const int       bufferIdx = scratch.acquireBufferIndex();
-      scratch.waitForBuffer(bufferIdx);
-      scratch.ensureCapacity(bufferIdx, static_cast<int>(numPatternsInSubBatch));
-      for (size_t i = 0; i < numPatternsInSubBatch; ++i) {
-        scratch.patternsAtDepthHost[bufferIdx][i] = patternsForDepth[patternStart + i];
-      }
-      prepareRange.pop();
+      // Use the device-resident pattern table directly.
+      const BatchedPatternEntry* patternEntriesDevice =
+        leafSubpatterns_.allQueriesPatternsAtDepthDevice[currentDepth].data() + patternStart;
 
       const int    buffersPerBlock = gsiBuffersPerBlock;
       const size_t overflowNeeded  = numBlocksInSubBatch * buffersPerBlock * kOverflowEntriesPerBuffer;
@@ -251,13 +258,6 @@ void RecursivePatternPreprocessor::preprocessMiniBatch(
         scratch.labelMatrixBuffer.resize(static_cast<size_t>(labelMatrixNeeded * 1.5));
       }
 
-      if (scratch.patternEntries.size() < numPatternsInSubBatch) {
-        scratch.patternEntries.resize(static_cast<size_t>(numPatternsInSubBatch * 1.5));
-      }
-
-      scratch.patternEntries.copyFromHost(scratch.patternsAtDepthHost[bufferIdx].data(), numPatternsInSubBatch);
-      scratch.recordCopy(bufferIdx, scratch.patternEntries.stream());
-
       const uint32_t* recursiveBitsForLabel = (currentDepth > 0) ? miniBatchResults.recursiveMatchBits() : nullptr;
 
       std::optional<ZeroBuffersSpec> zeroBuffers;
@@ -272,7 +272,7 @@ void RecursivePatternPreprocessor::preprocessMiniBatch(
       launchLabelMatrixPaintKernel(paintConfig,
                                    targetsDevice.view<MoleculeType::Target>(),
                                    leafSubpatterns_.view(),
-                                   scratch.patternEntries.data(),
+                                   patternEntriesDevice,
                                    static_cast<int>(numPatternsInSubBatch),
                                    numBlocksInSubBatch,
                                    numQueries,
@@ -289,7 +289,7 @@ void RecursivePatternPreprocessor::preprocessMiniBatch(
                                  algorithm,
                                  targetsDevice.view<MoleculeType::Target>(),
                                  leafSubpatterns_.view(),
-                                 scratch.patternEntries.data(),
+                                 patternEntriesDevice,
                                  static_cast<int>(numPatternsInSubBatch),
                                  numBlocksInSubBatch,
                                  miniBatchResults.recursiveMatchBits(),

src/substruct/recursive_preprocessor.h

@@ -87,6 +87,9 @@ struct LeafSubpatterns {
   /// Max recursion depth across all queries
   int allQueriesMaxDepth = 0;
 
+  /// Device-resident pattern table shared across mini-batches.
+  std::array<AsyncDeviceVector<BatchedPatternEntry>, kMaxSmartsNestingDepth + 1> allQueriesPatternsAtDepthDevice;
+
   int maxPatternAtoms_ = 0;
 
   LeafSubpatterns() = default;

src/substruct/substruct_search.cu

@@ -439,6 +439,11 @@ void runGpuCoordinator(int                                 deviceId,
       localPreprocessor = std::make_unique<RecursivePatternPreprocessor>();
       localPreprocessor->buildPatterns(queriesHost);
       localPreprocessor->syncToDevice(nullptr);
+
+      // The workers use independent non-blocking streams, so fully publish the
+      // secondary-device queries and patterns before any worker can consume them.
+      cudaCheckError(cudaStreamSynchronize(nullptr));
+
       queriesPtr      = localQueries.get();
       preprocessorPtr = localPreprocessor.get();
     }