diff --git a/benchmarks/substruct_bench.py b/benchmarks/substruct_bench.py index a3424318..4a88602e 100644 --- a/benchmarks/substruct_bench.py +++ b/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", diff --git a/nvmolkit/autotune/tune_substructure.py b/nvmolkit/autotune/tune_substructure.py index 49418688..dfbeffd7 100644 --- a/nvmolkit/autotune/tune_substructure.py +++ b/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, diff --git a/nvmolkit/tests/test_autotune.py b/nvmolkit/tests/test_autotune.py index d4a52803..e8c3833b 100644 --- a/nvmolkit/tests/test_autotune.py +++ b/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: diff --git a/src/substruct/recursive_preprocessor.cu b/src/substruct/recursive_preprocessor.cu index a1bc2f52..5fcff86c 100644 --- a/src/substruct/recursive_preprocessor.cu +++ b/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(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(labelMatrixNeeded * 1.5)); } - if (scratch.patternEntries.size() < numPatternsInSubBatch) { - scratch.patternEntries.resize(static_cast(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 zeroBuffers; @@ -272,7 +272,7 @@ void RecursivePatternPreprocessor::preprocessMiniBatch( launchLabelMatrixPaintKernel(paintConfig, targetsDevice.view(), leafSubpatterns_.view(), - scratch.patternEntries.data(), + patternEntriesDevice, static_cast(numPatternsInSubBatch), numBlocksInSubBatch, numQueries, @@ -289,7 +289,7 @@ void RecursivePatternPreprocessor::preprocessMiniBatch( algorithm, targetsDevice.view(), leafSubpatterns_.view(), - scratch.patternEntries.data(), + patternEntriesDevice, static_cast(numPatternsInSubBatch), numBlocksInSubBatch, miniBatchResults.recursiveMatchBits(), diff --git a/src/substruct/recursive_preprocessor.h b/src/substruct/recursive_preprocessor.h index 2201c00f..5a5e37c1 100644 --- a/src/substruct/recursive_preprocessor.h +++ b/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, kMaxSmartsNestingDepth + 1> allQueriesPatternsAtDepthDevice; + int maxPatternAtoms_ = 0; LeafSubpatterns() = default; diff --git a/src/substruct/substruct_search.cu b/src/substruct/substruct_search.cu index 4dee07da..99a88f23 100644 --- a/src/substruct/substruct_search.cu +++ b/src/substruct/substruct_search.cu @@ -439,6 +439,11 @@ void runGpuCoordinator(int deviceId, localPreprocessor = std::make_unique(); 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(); }