From 1bb8b657cfb6ed39dce4bbf354ce24b1d6b8f409 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Fri, 2 Oct 2026 17:34:44 -0400 Subject: [PATCH] Add a persistent GPU SubstructLibrary SubstructLibrary is the GPU counterpart of RDKit's SubstructLibrary: molecules are copied in once and packed and uploaded at finalize(), then any number of queries run against them. Molecules get IDs in the order they are added and become searchable at the next finalize(); queries keep searching the previously finalized molecules until then. Each GPU holds its molecules as one packed batch, new molecules go to the GPU with the fewest, and results are merged in ID order. Molecules the packed format cannot represent are matched with RDKit on the CPU. Each GPU keeps one search workspace, and calls are serialized by the library lock. Also rename the resident-search API from #351 to the codebase's host and device vocabulary: ResidentTargetBatch / makeResidentTargetBatch become PersistentDeviceTargets / makePersistentDeviceTargets, hasSubstructMatchResident becomes a hasSubstructMatch overload, and ResidentSubstructSearchWorkspace becomes SubstructSearchWorkspace. Declare mergeBatch() in molecules.h so the library can append packed batches. --- src/substruct/CMakeLists.txt | 7 + src/substruct/molecules.h | 7 + src/substruct/substruct_library.cpp | 510 +++++++++++++++++++++++++++ src/substruct/substruct_library.h | 77 +++++ src/substruct/substruct_search.cu | 115 ++++--- src/substruct/substruct_search.h | 39 ++- tests/CMakeLists.txt | 5 + tests/test_substruct_library.cu | 515 ++++++++++++++++++++++++++++ tests/test_substruct_search.cu | 45 ++- 9 files changed, 1221 insertions(+), 99 deletions(-) create mode 100644 src/substruct/substruct_library.cpp create mode 100644 src/substruct/substruct_library.h create mode 100644 tests/test_substruct_library.cu diff --git a/src/substruct/CMakeLists.txt b/src/substruct/CMakeLists.txt index 2976e258..c2da06d9 100644 --- a/src/substruct/CMakeLists.txt +++ b/src/substruct/CMakeLists.txt @@ -51,3 +51,10 @@ target_link_libraries( PRIVATE ${RDKit_LIBS}) target_compile_options(substructure_search PRIVATE $<$:-Xcompiler=-fopenmp>) + +add_library(substruct_library substruct_library.cpp) +target_link_libraries( + substruct_library + PUBLIC substructure_search molecules + PRIVATE ${RDKit_LIBS} cuda_error_check device nvtx OpenMP::OpenMP_CXX + openmp_helpers) diff --git a/src/substruct/molecules.h b/src/substruct/molecules.h index 2db94d1b..ef4f84cd 100644 --- a/src/substruct/molecules.h +++ b/src/substruct/molecules.h @@ -292,6 +292,13 @@ class MoleculesDevice { */ void addToBatch(const RDKit::ROMol* mol, MoleculesHost& batch); +/** + * @brief Append every molecule of src to dest, adjusting all offsets. + * + * Works for both target and query batches. + */ +void mergeBatch(MoleculesHost& dest, const MoleculesHost& src); + /** * @brief Add a query molecule (from SMARTS) to an existing batch. * diff --git a/src/substruct/substruct_library.cpp b/src/substruct/substruct_library.cpp new file mode 100644 index 00000000..7ead077e --- /dev/null +++ b/src/substruct/substruct_library.cpp @@ -0,0 +1,510 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include "src/substruct/substruct_library.h" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "src/substruct/molecules.h" +#include "src/substruct/substruct_constants.h" +#include "src/substruct/substruct_search.h" +#include "src/utils/cuda_error_check.h" +#include "src/utils/device.h" +#include "src/utils/nvtx.h" +#include "src/utils/openmp_helpers.h" + +namespace nvMolKit { + +namespace { + +bool rdkitHasMatch(const RDKit::ROMol& target, const RDKit::ROMol& query) { + RDKit::SubstructMatchParameters params; + params.uniquify = false; + params.maxMatches = 1; + params.useChirality = false; + params.useQueryQueryMatches = false; + return !RDKit::SubstructMatch(target, query, params).empty(); +} + +//! A copy of molecule with the symmetrized SSSR ring info that GPU packing and RDKit ring queries need. RDKit's +//! SubstructLibrary computes it at search time; the library owns its copies, so it is computed once here. +std::unique_ptr copyWithRingInfo(const RDKit::ROMol& molecule) { + auto copy = std::make_unique(molecule); + if (!copy->getRingInfo()->isSymmSssr()) { + RDKit::MolOps::symmetrizeSSSR(*copy); + } + return copy; +} + +//! Whether the packed GPU target format can represent a molecule; the rest are matched with RDKit. +bool fitsGpu(const RDKit::ROMol& molecule) { + try { + return molecule.getNumAtoms() <= kMaxTargetAtoms && !requiresRDKitFallback(&molecule); + } catch (const std::runtime_error&) { + return false; + } +} + +//! The finalized molecules searched on one GPU, packed and uploaded as a single batch in ID order. +struct GpuTargets { + std::unique_ptr host; + std::unique_ptr device; + std::vector molecules; + std::vector ids; + std::shared_ptr searchTargets; +}; + +} // namespace + +class SubstructLibrary::Impl { + public: + explicit Impl(SubstructSearchConfig config) : config_(std::move(config)) { + if (config_.algorithm == SubstructAlgorithm::VF2) { + throw std::invalid_argument("Substructure libraries support the GSI and DFS production backends"); + } + std::unordered_set seenDevices; + for (const int deviceId : config_.gpuIds) { + if (deviceId < 0) { + throw std::invalid_argument("Substructure library GPU ID must be nonnegative"); + } + if (!seenDevices.insert(deviceId).second) { + throw std::invalid_argument("Substructure library GPU IDs must be unique"); + } + } + deviceIds_ = config_.gpuIds; + } + + ~Impl() noexcept { + releaseOnDevices( + gpus_.size(), + [&](std::size_t index) { return gpus_[index] != nullptr; }, + [&](std::size_t index) { gpus_[index].reset(); }); + releaseOnDevices( + workspaces_.size(), + [&](std::size_t index) { return workspaces_[index] != nullptr; }, + [&](std::size_t index) { workspaces_[index].reset(); }); + } + + unsigned int addMol(const RDKit::ROMol& molecule) { + const std::lock_guard lock(mutex_); + requireIdsAvailable(1); + molecules_.push_back(copyWithRingInfo(molecule)); + return static_cast(molecules_.size() - 1); + } + + std::vector addMols(const std::vector& molecules) { + const std::lock_guard lock(mutex_); + requireIdsAvailable(molecules.size()); + const std::size_t firstId = molecules_.size(); + molecules_.resize(firstId + molecules.size()); + detail::OpenMPExceptionRegistry exceptions; +#pragma omp parallel for num_threads(threads()) schedule(static) + for (std::int64_t index = 0; index < static_cast(molecules.size()); ++index) { + try { + const RDKit::ROMol* molecule = molecules[static_cast(index)]; + if (molecule == nullptr) { + throw std::invalid_argument("Substructure library molecule cannot be null"); + } + molecules_[firstId + static_cast(index)] = copyWithRingInfo(*molecule); + } catch (...) { + exceptions.store(std::current_exception()); + } + } + try { + exceptions.rethrow(); + } catch (...) { + molecules_.resize(firstId); + throw; + } + std::vector ids(molecules.size()); + for (std::size_t index = 0; index < ids.size(); ++index) { + ids[index] = static_cast(firstId + index); + } + return ids; + } + + void finalize(cudaStream_t stream) { + const std::lock_guard lock(mutex_); + useCurrentDeviceIfNoneConfigured(); + if (deviceIds_.size() > 1 && stream != nullptr) { + throw std::invalid_argument("A single external CUDA stream cannot be used with a multi-GPU substructure library"); + } + int deviceCount = 0; + cudaCheckError(cudaGetDeviceCount(&deviceCount)); + for (const int deviceId : deviceIds_) { + if (deviceId >= deviceCount) { + throw std::invalid_argument("Substructure library GPU ID is not available"); + } + } + gpus_.resize(deviceIds_.size()); + + // Molecules added since the last finalize(): those the GPU format cannot represent are matched with RDKit, + // and the rest go to the GPU searching the fewest molecules. + const std::size_t firstNew = numFinalized_; + const std::size_t numNew = molecules_.size() - firstNew; + std::vector onGpu(numNew); + detail::OpenMPExceptionRegistry screenExceptions; +#pragma omp parallel for num_threads(threads()) schedule(static) + for (std::int64_t index = 0; index < static_cast(numNew); ++index) { + try { + onGpu[static_cast(index)] = fitsGpu(*molecules_[firstNew + static_cast(index)]); + } catch (...) { + screenExceptions.store(std::current_exception()); + } + } + screenExceptions.rethrow(); + std::vector gpuLoad(deviceIds_.size(), 0); + for (std::size_t gpu = 0; gpu < deviceIds_.size(); ++gpu) { + gpuLoad[gpu] = gpus_[gpu] != nullptr ? gpus_[gpu]->ids.size() : 0; + } + std::vector> newGpuIds(deviceIds_.size()); + std::vector newRdkitIds; + for (std::size_t index = 0; index < numNew; ++index) { + const auto id = static_cast(firstNew + index); + if (onGpu[index] == 0) { + newRdkitIds.push_back(id); + continue; + } + const auto gpu = static_cast(std::min_element(gpuLoad.begin(), gpuLoad.end()) - gpuLoad.begin()); + newGpuIds[gpu].push_back(id); + ++gpuLoad[gpu]; + } + + // Pack one GPU at a time with every thread, rather than inside the per-GPU loop below, where OpenMP may not + // give the packer's nested parallel region more than one thread. + std::vector packed(deviceIds_.size()); + for (std::size_t gpu = 0; gpu < deviceIds_.size(); ++gpu) { + if (!newGpuIds[gpu].empty()) { + buildTargetBatchParallelInto(packed[gpu], threads(), moleculesWithIds(newGpuIds[gpu]), {}); + } + } + + // Build every GPU's replacement before touching the current ones, so a failure leaves the library as it was. + std::vector> replacements(deviceIds_.size()); + detail::OpenMPExceptionRegistry exceptions; +#pragma omp parallel for num_threads(static_cast(deviceIds_.size())) schedule(static) + for (std::int64_t gpu = 0; gpu < static_cast(deviceIds_.size()); ++gpu) { + try { + const auto index = static_cast(gpu); + if (!newGpuIds[index].empty()) { + const WithDevice device(deviceIds_[index]); + replacements[index] = + extendGpuTargets(gpus_[index].get(), newGpuIds[index], packed[index], validateStream(stream)); + } + } catch (...) { + // Clear a recoverable CUDA error such as out of memory, so this thread's next CUDA check does not report it. + static_cast(cudaGetLastError()); + exceptions.store(std::current_exception()); + } + } + try { + exceptions.rethrow(); + createWorkspaces(); + rdkitIds_.reserve(rdkitIds_.size() + newRdkitIds.size()); + } catch (...) { + static_cast(cudaGetLastError()); + releaseOnDevices( + replacements.size(), + [&](std::size_t index) { return replacements[index] != nullptr; }, + [&](std::size_t index) { replacements[index].reset(); }); + throw; + } + + releaseOnDevices( + gpus_.size(), + [&](std::size_t index) { return replacements[index] != nullptr; }, + [&](std::size_t index) { gpus_[index] = std::move(replacements[index]); }); + rdkitIds_.insert(rdkitIds_.end(), newRdkitIds.begin(), newRdkitIds.end()); + numFinalized_ = molecules_.size(); + finalized_ = true; + } + + [[nodiscard]] std::size_t size() const { + const std::lock_guard lock(mutex_); + return numFinalized_; + } + + [[nodiscard]] std::size_t pendingSize() const { + const std::lock_guard lock(mutex_); + return molecules_.size() - numFinalized_; + } + + [[nodiscard]] std::vector getMatches(const RDKit::ROMol& query, + int maxResults, + cudaStream_t stream) const { + const std::lock_guard lock(mutex_); + requireFinalized(); + if (maxResults == 0) { + return {}; + } + return matches(query, stream, maxResults); + } + + [[nodiscard]] std::size_t countMatches(const RDKit::ROMol& query, cudaStream_t stream) const { + const std::lock_guard lock(mutex_); + requireFinalized(); + return matches(query, stream, -1).size(); + } + + [[nodiscard]] bool hasMatch(const RDKit::ROMol& query, cudaStream_t stream) const { + const std::lock_guard lock(mutex_); + requireFinalized(); + return !matches(query, stream, 1).empty(); + } + + private: + void requireIdsAvailable(std::size_t count) const { + if (count > std::numeric_limits::max() - molecules_.size()) { + throw std::overflow_error("Substructure library molecule ID space exhausted"); + } + } + + void requireFinalized() const { + if (!finalized_) { + throw std::logic_error("Substructure library must be finalized before querying"); + } + } + + [[nodiscard]] int threads() const { + return config_.preprocessingThreads == -1 ? omp_get_max_threads() : std::max(1, config_.preprocessingThreads); + } + + void useCurrentDeviceIfNoneConfigured() { + if (deviceIds_.empty()) { + int currentDevice = 0; + cudaCheckError(cudaGetDevice(¤tDevice)); + deviceIds_.push_back(currentDevice); + config_.gpuIds = deviceIds_; + } + } + + [[nodiscard]] cudaStream_t validateStream(cudaStream_t stream) const { + const auto validated = acquireExternalStream(reinterpret_cast(stream)); + if (!validated.has_value()) { + throw std::invalid_argument("CUDA stream does not belong to the substructure library device"); + } + return *validated; + } + + [[nodiscard]] std::vector moleculesWithIds(const std::vector& ids) const { + std::vector result; + result.reserve(ids.size()); + for (const unsigned int id : ids) { + result.push_back(molecules_[id].get()); + } + return result; + } + + // The current device's targets plus newIds, whose packed form is packed, uploaded on stream and ready to search. + // The whole batch is copied and uploaded again, so the current targets stay intact if this throws. + [[nodiscard]] std::unique_ptr extendGpuTargets(const GpuTargets* current, + const std::vector& newIds, + const MoleculesHost& packed, + cudaStream_t stream) const { + const std::vector newMolecules = moleculesWithIds(newIds); + + auto next = std::make_unique(); + next->host = + current != nullptr ? std::make_unique(*current->host) : std::make_unique(); + mergeBatch(*next->host, packed); + if (current != nullptr) { + next->molecules = current->molecules; + next->ids = current->ids; + } + next->molecules.insert(next->molecules.end(), newMolecules.begin(), newMolecules.end()); + next->ids.insert(next->ids.end(), newIds.begin(), newIds.end()); + + next->device = std::make_unique(stream); + next->device->copyFromHost(*next->host, stream); + cudaCheckError(cudaStreamSynchronize(stream)); + // The device copy outlives the caller's stream, so release it on the default stream. + next->device->setStream(nullptr); + next->searchTargets = makePersistentDeviceTargets(next->molecules, *next->host, *next->device); + return next; + } + + // Search state for each GPU, created on the first finalize() and reused by every query. + void createWorkspaces() { + if (workspaces_.size() == deviceIds_.size()) { + return; + } + std::vector> created(deviceIds_.size()); + for (std::size_t gpu = 0; gpu < deviceIds_.size(); ++gpu) { + created[gpu] = makeSubstructSearchWorkspace(deviceIds_[gpu]); + } + workspaces_ = std::move(created); + matchFlags_.resize(deviceIds_.size()); + } + + [[nodiscard]] SubstructSearchConfig gpuConfig(std::size_t gpu) const { + SubstructSearchConfig result = config_; + result.gpuIds = {deviceIds_[gpu]}; + // Targets are already packed and each search runs one query, so a large preprocessing team does not help. + result.preprocessingThreads = config_.preprocessingThreads == -1 ? + 1 : + std::max(1, config_.preprocessingThreads / static_cast(deviceIds_.size())); + if (result.workerThreads == -1) { + result.workerThreads = 1; + } + return result; + } + + // IDs of molecules matching query, ascending and truncated to maxResults when positive. + [[nodiscard]] std::vector matches(const RDKit::ROMol& query, + cudaStream_t stream, + int maxResults) const { + if (deviceIds_.size() > 1 && stream != nullptr) { + throw std::invalid_argument("A single external CUDA stream cannot be used with a multi-GPU substructure library"); + } + std::vector> perGpu(deviceIds_.size()); + detail::OpenMPExceptionRegistry exceptions; +#pragma omp parallel for num_threads(static_cast(deviceIds_.size())) schedule(static) + for (std::int64_t gpu = 0; gpu < static_cast(deviceIds_.size()); ++gpu) { + try { + const auto index = static_cast(gpu); + if (gpus_[index] != nullptr) { + const WithDevice device(deviceIds_[index]); + perGpu[index] = gpuMatches(index, query, validateStream(stream)); + } + } catch (...) { + exceptions.store(std::current_exception()); + } + } + exceptions.rethrow(); + + std::vector result; + for (const auto& gpuResult : perGpu) { + result.insert(result.end(), gpuResult.begin(), gpuResult.end()); + } + std::sort(result.begin(), result.end()); + + // Molecules matched with RDKit. Once the result limit is filled, only IDs below the last kept match can + // change the answer. + const bool limitFilled = maxResults > 0 && result.size() >= static_cast(maxResults); + const std::uint64_t rdkitIdLimit = + limitFilled ? result[static_cast(maxResults) - 1] : std::numeric_limits::max(); + const std::size_t gpuMatchCount = result.size(); + ScopedNvtxRange rdkitRange("SubstructLibrary RDKit matching"); + for (const unsigned int id : rdkitIds_) { + if (id >= rdkitIdLimit) { + break; + } + if (rdkitHasMatch(*molecules_[id], query)) { + result.push_back(id); + } + } + std::inplace_merge(result.begin(), result.begin() + static_cast(gpuMatchCount), result.end()); + if (maxResults > 0 && result.size() > static_cast(maxResults)) { + result.resize(static_cast(maxResults)); + } + return result; + } + + // IDs of the molecules on one GPU that match query, ascending. + [[nodiscard]] std::vector gpuMatches(std::size_t gpu, + const RDKit::ROMol& query, + cudaStream_t stream) const { + ScopedNvtxRange searchRange("SubstructLibrary GPU search"); + const GpuTargets& targets = *gpus_[gpu]; + std::vector& flags = matchFlags_[gpu]; + const SubstructSearchConfig config = gpuConfig(gpu); + hasSubstructMatch(*targets.searchTargets, query, flags, config.algorithm, stream, config, workspaces_[gpu].get()); + std::vector result; + for (std::size_t target = 0; target < flags.size(); ++target) { + if (flags[target] != 0) { + result.push_back(targets.ids[target]); + } + } + return result; + } + + // For each of the first count configured devices i where holds(i), call release(i) with device i current, then + // restore the caller's device. Devices without resources are never selected, so a device ID rejected by + // finalize() leaves no CUDA error behind. Teardown cannot throw, so CUDA failures are logged and skipped. + template + void releaseOnDevices(std::size_t count, const Holds& holds, const Release& release) const noexcept { + int originalDevice = -1; + cudaCheckErrorNoThrow(cudaGetDevice(&originalDevice)); + for (std::size_t index = 0; index < std::min(count, deviceIds_.size()); ++index) { + if (!holds(index)) { + continue; + } + const cudaError_t selected = cudaSetDevice(deviceIds_[index]); + cudaCheckErrorNoThrow(selected); + if (selected == cudaSuccess) { + release(index); + } + } + if (originalDevice >= 0) { + cudaCheckErrorNoThrow(cudaSetDevice(originalDevice)); + } + } + + SubstructSearchConfig config_; + std::vector deviceIds_; + // Held by every call, so queries, additions, and finalize() run one at a time. + mutable std::mutex mutex_; + + //! Every added molecule, indexed by ID; the first numFinalized_ are searchable. + std::vector> molecules_; + std::size_t numFinalized_ = 0; + bool finalized_ = false; + + std::vector> gpus_; + //! Finalized molecules the GPU format cannot represent, ascending; matched with RDKit. + std::vector rdkitIds_; + std::vector> workspaces_; + //! Per-GPU match flags, reused across queries to avoid reallocating them. + mutable std::vector> matchFlags_; +}; + +SubstructLibrary::SubstructLibrary(SubstructSearchConfig config) : impl_(std::make_unique(std::move(config))) {} + +SubstructLibrary::~SubstructLibrary() = default; + +unsigned int SubstructLibrary::addMol(const RDKit::ROMol& molecule) { + return impl_->addMol(molecule); +} + +std::vector SubstructLibrary::addMols(const std::vector& molecules) { + return impl_->addMols(molecules); +} + +void SubstructLibrary::finalize(cudaStream_t stream) { + impl_->finalize(stream); +} + +std::size_t SubstructLibrary::size() const { + return impl_->size(); +} + +std::size_t SubstructLibrary::pendingSize() const { + return impl_->pendingSize(); +} + +std::vector SubstructLibrary::getMatches(const RDKit::ROMol& query, + int maxResults, + cudaStream_t stream) const { + return impl_->getMatches(query, maxResults, stream); +} + +std::size_t SubstructLibrary::countMatches(const RDKit::ROMol& query, cudaStream_t stream) const { + return impl_->countMatches(query, stream); +} + +bool SubstructLibrary::hasMatch(const RDKit::ROMol& query, cudaStream_t stream) const { + return impl_->hasMatch(query, stream); +} + +} // namespace nvMolKit diff --git a/src/substruct/substruct_library.h b/src/substruct/substruct_library.h new file mode 100644 index 00000000..58013596 --- /dev/null +++ b/src/substruct/substruct_library.h @@ -0,0 +1,77 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#ifndef NVMOLKIT_SUBSTRUCT_LIBRARY_H +#define NVMOLKIT_SUBSTRUCT_LIBRARY_H + +#include + +#include +#include +#include + +#include "src/substruct/substruct_results.h" + +namespace RDKit { +class ROMol; +} // namespace RDKit + +namespace nvMolKit { + +/** + * Target molecules kept on the GPU for repeated substructure queries, the GPU + * counterpart of RDKit's SubstructLibrary. Molecules get IDs in the order they + * are added and become searchable at the next finalize(). With several GPU IDs + * configured, molecules are split across those GPUs and results are merged in + * ID order. Molecules the GPU format cannot represent are matched with RDKit. + * + * Calls are serialized: queries, additions, and finalize() run one at a time. + */ +class SubstructLibrary { + public: + explicit SubstructLibrary(SubstructSearchConfig config = SubstructSearchConfig{}); + ~SubstructLibrary(); + + SubstructLibrary(const SubstructLibrary&) = delete; + SubstructLibrary& operator=(const SubstructLibrary&) = delete; + SubstructLibrary(SubstructLibrary&&) = delete; + SubstructLibrary& operator=(SubstructLibrary&&) = delete; + + /** Copy a molecule into the library and return its ID. It is searchable after the next finalize(). */ + unsigned int addMol(const RDKit::ROMol& molecule); + + /** + * Copy molecules into the library in parallel and return their IDs. They are searchable after the next + * finalize(). Uses preprocessingThreads from the search configuration (-1 selects all OpenMP threads). + */ + std::vector addMols(const std::vector& molecules); + + /** + * Pack and upload every molecule added since the last finalize() and make them searchable. If it throws, the + * library is unchanged and the molecules remain waiting for the next finalize(). A GPU that receives new + * molecules has all of its molecules copied and uploaded again, so add molecules in large groups between calls. + */ + void finalize(cudaStream_t stream = nullptr); + + /** Number of searchable molecules. */ + [[nodiscard]] std::size_t size() const; + + /** Number of molecules added since the last finalize(). */ + [[nodiscard]] std::size_t pendingSize() const; + + /** IDs of matching molecules, ascending. maxResults limits the count; -1 returns all, 0 returns none. */ + [[nodiscard]] std::vector getMatches(const RDKit::ROMol& query, + int maxResults = -1, + cudaStream_t stream = nullptr) const; + + [[nodiscard]] std::size_t countMatches(const RDKit::ROMol& query, cudaStream_t stream = nullptr) const; + [[nodiscard]] bool hasMatch(const RDKit::ROMol& query, cudaStream_t stream = nullptr) const; + + private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace nvMolKit + +#endif // NVMOLKIT_SUBSTRUCT_LIBRARY_H diff --git a/src/substruct/substruct_search.cu b/src/substruct/substruct_search.cu index 7b062c78..49b81d4b 100644 --- a/src/substruct/substruct_search.cu +++ b/src/substruct/substruct_search.cu @@ -64,10 +64,10 @@ void runPipelinedSubstructSearch(const std::vector& targets const SubstructSearchConfig& config, int effectivePreprocessingThreads, RDKitFallbackQueue* fallbackQueue, - HasSubstructMatchResults* boolResults = nullptr, - std::vector* countResults = nullptr, - const ResidentTargetBatch* resident = nullptr, - ResidentSubstructSearchWorkspace* workspace = nullptr); + HasSubstructMatchResults* boolResults = nullptr, + std::vector* countResults = nullptr, + const PersistentDeviceTargets* persistentTargets = nullptr, + SubstructSearchWorkspace* workspace = nullptr); } // anonymous namespace @@ -81,8 +81,8 @@ struct PreparedMiniBatch { std::shared_ptr> targetAtomCounts; ThreadWorkerContext ctx; MiniBatchPlan plan; - PinnedHostBuffer* pinnedBuffer = nullptr; - const MoleculesDevice* residentTargetsDevice = nullptr; + PinnedHostBuffer* pinnedBuffer = nullptr; + const MoleculesDevice* persistentTargetsDevice = nullptr; }; using PreparedBatchQueue = ThreadSafeQueue>; @@ -113,20 +113,20 @@ TargetBatchShape measureTargetBatch(const MoleculesHost& targetsHost, int numTar return shape; } -/// Resident targets are searched in place, so original index i is target i. -struct ResidentTargetMetadata { +/// Persistent targets are searched in place, so original index i is target i. +struct PersistentTargetMetadata { std::shared_ptr> originalIndices; TargetBatchShape shape; }; -ResidentTargetMetadata measureResidentTargets(const MoleculesHost& targetsHost, int numTargets) { - ResidentTargetMetadata metadata; +PersistentTargetMetadata measurePersistentTargets(const MoleculesHost& targetsHost, int numTargets) { + PersistentTargetMetadata metadata; metadata.originalIndices = std::make_shared>(static_cast(numTargets)); std::iota(metadata.originalIndices->begin(), metadata.originalIndices->end(), 0); metadata.shape = measureTargetBatch(targetsHost, numTargets); if (metadata.shape.maxTargetAtoms > kMaxTargetAtoms) { - // Resident targets skip the per-call RDKit fallback screen, so they must all fit the GPU path. - throw std::invalid_argument("Resident targets must have at most " + std::to_string(kMaxTargetAtoms) + " atoms"); + // Persistent targets skip the per-call RDKit fallback screen, so they must all fit the GPU path. + throw std::invalid_argument("Persistent targets must have at most " + std::to_string(kMaxTargetAtoms) + " atoms"); } return metadata; } @@ -165,8 +165,8 @@ std::unique_ptr makePreparedMiniBatch(PinnedHostBuffer* } // namespace -struct ResidentSubstructSearchWorkspace { - explicit ResidentSubstructSearchWorkspace(int gpuDeviceId) : deviceId(gpuDeviceId) {} +struct SubstructSearchWorkspace { + explicit SubstructSearchWorkspace(int gpuDeviceId) : deviceId(gpuDeviceId) {} /// Pinned query atom counts lent to each search, avoiding a pinned allocation per call. PinnedHostVector queryAtomCounts; @@ -199,12 +199,12 @@ struct ResidentSubstructSearchWorkspace { PinnedHostBufferPool bufferPool; }; -struct ResidentTargetBatch { +struct PersistentDeviceTargets { std::vector targets; const MoleculesHost* host; const MoleculesDevice* device; int deviceId; - ResidentTargetMetadata metadata; + PersistentTargetMetadata metadata; }; // ============================================================================= @@ -440,7 +440,7 @@ void runnerWorkerPipeline(int workerIdx, executor.consolidatedBuffer.allocate(batch->pinnedBuffer->consolidated.maxBatchSize, executor.stream()); executor.consolidatedBuffer.copyFromHost(*batch->pinnedBuffer, executor.stream()); - const MoleculesDevice* targetsDevice = batch->residentTargetsDevice; + const MoleculesDevice* targetsDevice = batch->persistentTargetsDevice; if (targetsDevice == nullptr) { executor.targetsDevice.copyFromHost(*batch->targetsHost, executor.stream()); cudaCheckError(cudaEventRecord(executor.targetsReadyEvent.event(), executor.stream())); @@ -538,7 +538,7 @@ void runGpuCoordinator(int deviceId, std::vector* countResults, std::vector& exceptions, std::atomic& pipelineAbort, - ResidentSubstructSearchWorkspace* workspace) { + SubstructSearchWorkspace* workspace) { try { ScopedNvtxRange coordRange("GPU" + std::to_string(deviceId) + " coordinator (pipeline)"); const WithDevice setDevice(deviceId); @@ -550,7 +550,7 @@ void runGpuCoordinator(int deviceId, executors.reserve(static_cast(numExecutorsThisGpu)); if (workspace != nullptr) { if (workspace->deviceId != deviceId || static_cast(workspace->executors.size()) < numExecutorsThisGpu) { - throw std::logic_error("Resident search workspace does not match the coordinator"); + throw std::logic_error("Search workspace does not match the coordinator"); } for (int i = 0; i < numExecutorsThisGpu; ++i) { executors.push_back(workspace->executors[static_cast(i)].get()); @@ -648,8 +648,8 @@ void runPipelinedSubstructSearch(const std::vector& targets RDKitFallbackQueue* fallbackQueue, HasSubstructMatchResults* boolResults, std::vector* countResults, - const ResidentTargetBatch* resident, - ResidentSubstructSearchWorkspace* workspace) { + const PersistentDeviceTargets* persistentTargets, + SubstructSearchWorkspace* workspace) { (void)stream; const bool countOnly = (boolResults != nullptr) || (countResults != nullptr); const char* rangeLabel = boolResults ? @@ -673,7 +673,7 @@ void runPipelinedSubstructSearch(const std::vector& targets } const int numGpus = static_cast(gpuIds.size()); if (workspace != nullptr && (numGpus != 1 || workspace->deviceId != gpuIds.front())) { - throw std::invalid_argument("A resident search workspace is bound to exactly one configured GPU"); + throw std::invalid_argument("A search workspace is bound to exactly one configured GPU"); } // Determine runner counts (per GPU, possibly limited by target count). @@ -830,11 +830,11 @@ void runPipelinedSubstructSearch(const std::vector& targets try { ScopedNvtxRange threadRange("Preprocess thread " + std::to_string(t)); - if (resident != nullptr) { + if (persistentTargets != nullptr) { if (t != 0) { return; } - const ResidentTargetMetadata* metadata = &resident->metadata; + const PersistentTargetMetadata* metadata = &persistentTargets->metadata; const int totalPairs = numTargets * numQueries; for (int pairOffset = 0; pairOffset < totalPairs; pairOffset += maxPairsPerBatch) { @@ -847,14 +847,14 @@ void runPipelinedSubstructSearch(const std::vector& targets } BufferReleaseGuard releaseGuard{&bufferPool, buffer}; - auto batch = makePreparedMiniBatch(buffer, + auto batch = makePreparedMiniBatch(buffer, metadata->originalIndices, metadata->shape, numTargets, queryContext, config, countOnly); - batch->residentTargetsDevice = resident->device; + batch->persistentTargetsDevice = persistentTargets->device; planner.prepareMiniBatch(batch->plan, *buffer, batch->ctx, leafSubpatterns, pairOffset, maxPairsPerBatch); releaseGuard.release(); batchQueue.push(std::move(batch)); @@ -1065,8 +1065,8 @@ void getSubstructMatchesImpl(const std::vector& targets, const SubstructSearchConfig& config, HasSubstructMatchResults* boolResults, std::vector* countResults, - const ResidentTargetBatch* resident = nullptr, - ResidentSubstructSearchWorkspace* workspace = nullptr) { + const PersistentDeviceTargets* persistentTargets = nullptr, + SubstructSearchWorkspace* workspace = nullptr) { // CUDA sources build with --default-stream=per-thread but C++ sources use the legacy default stream, so a null // stream means a different stream on each side. Resolve it to an explicit handle before handing it to host-compiled // code such as MoleculesDevice, or the synchronization below would not cover the query upload. @@ -1143,8 +1143,8 @@ void getSubstructMatchesImpl(const std::vector& targets, QueryPreprocessContext queryContext; // Borrow the workspace's pinned buffer for the duration of this search. struct LentQueryAtomCounts { - ResidentSubstructSearchWorkspace* workspace; - PinnedHostVector& counts; + SubstructSearchWorkspace* workspace; + PinnedHostVector& counts; ~LentQueryAtomCounts() { if (workspace != nullptr) { workspace->queryAtomCounts = std::move(counts); @@ -1198,10 +1198,10 @@ void getSubstructMatchesImpl(const std::vector& targets, queryContext.maxQueryAtoms = maxQueryAtoms; - // Resident workspaces hold recursive scratch only while a recursive query + // Search workspaces hold recursive scratch only while a recursive query // runs; query admission budgets it separately from the base workspace. struct RecursiveScratchRelease { - ResidentSubstructSearchWorkspace* workspace = nullptr; + SubstructSearchWorkspace* workspace = nullptr; ~RecursiveScratchRelease() { if (workspace != nullptr) { workspace->releaseRecursiveScratch(); @@ -1249,7 +1249,7 @@ void getSubstructMatchesImpl(const std::vector& targets, &fallbackQueue, boolResults, countResults, - resident, + persistentTargets, workspace); // Process any remaining fallback entries after GPU work completes. @@ -1337,40 +1337,41 @@ void hasSubstructMatch(const std::vector& targets, } } -std::shared_ptr makeResidentTargetBatch(const std::vector& targets, - const MoleculesHost& targetsHost, - const MoleculesDevice& targetsDevice) { +std::shared_ptr makePersistentDeviceTargets( + const std::vector& targets, + const MoleculesHost& targetsHost, + const MoleculesDevice& targetsDevice) { if (targetsHost.numMolecules() != targets.size()) { - throw std::invalid_argument("Resident target batch does not match its target molecules"); + throw std::invalid_argument("Persistent device targets do not match their molecules"); } - // Resident searches skip the per-call RDKit fallback screen, so every target must fit the GPU path. + // Searches of these targets skip the per-call RDKit fallback screen, so every target must fit the GPU path. for (const RDKit::ROMol* target : targets) { if (target == nullptr || target->getNumAtoms() > kMaxTargetAtoms || requiresRDKitFallback(target)) { - throw std::invalid_argument("Resident targets must not need the RDKit fallback"); + throw std::invalid_argument("Persistent targets must not need the RDKit fallback"); } } - auto batch = std::make_shared(); + auto batch = std::make_shared(); batch->targets = targets; batch->host = &targetsHost; batch->device = &targetsDevice; - batch->metadata = measureResidentTargets(targetsHost, static_cast(targets.size())); + batch->metadata = measurePersistentTargets(targetsHost, static_cast(targets.size())); cudaCheckError(cudaGetDevice(&batch->deviceId)); return batch; } -void hasSubstructMatchResident(const ResidentTargetBatch& batch, - const RDKit::ROMol& query, - std::vector& results, - SubstructAlgorithm algorithm, - cudaStream_t stream, - const SubstructSearchConfig& config, - ResidentSubstructSearchWorkspace* workspace) { +void hasSubstructMatch(const PersistentDeviceTargets& batch, + const RDKit::ROMol& query, + std::vector& results, + SubstructAlgorithm algorithm, + cudaStream_t stream, + const SubstructSearchConfig& config, + SubstructSearchWorkspace* workspace) { if (workspace != nullptr && workspace->deviceId != batch.deviceId) { - throw std::invalid_argument("A resident search workspace must be on the GPU holding the targets"); + throw std::invalid_argument("A search workspace must be on the GPU holding the targets"); } if (!config.gpuIds.empty() && config.gpuIds != std::vector{batch.deviceId}) { throw std::invalid_argument( - "Resident searches run on the GPU holding the targets; gpuIds must be empty or name it"); + "Searches of persistent device targets run on the GPU holding them; gpuIds must be empty or name that GPU"); } const WithDevice device(batch.deviceId); std::unique_lock workspaceLock; @@ -1379,12 +1380,12 @@ void hasSubstructMatchResident(const ResidentTargetBatch& batch, } std::vector queries{&query}; - HasSubstructMatchResults residentResults; + HasSubstructMatchResults flagResults; SubstructSearchResults unusedResults; SubstructSearchConfig hasMatchConfig = config; hasMatchConfig.maxMatches = 1; hasMatchConfig.gpuIds = {batch.deviceId}; - residentResults.resize(static_cast(batch.targets.size()), 1); + flagResults.resize(static_cast(batch.targets.size()), 1); try { getSubstructMatchesImpl(batch.targets, queries, @@ -1392,21 +1393,21 @@ void hasSubstructMatchResident(const ResidentTargetBatch& batch, algorithm, stream, hasMatchConfig, - &residentResults, + &flagResults, nullptr, &batch, workspace); } catch (...) { // A failed search can leave executor work in flight; drain it before the workspace is reused. - cudaDeviceSynchronize(); + cudaCheckErrorNoThrow(cudaDeviceSynchronize()); throw; } - results = std::move(residentResults.hasMatch); + results = std::move(flagResults.hasMatch); } -std::shared_ptr makeResidentSubstructSearchWorkspace(int deviceId) { +std::shared_ptr makeSubstructSearchWorkspace(int deviceId) { const WithDevice device(deviceId); - return std::make_shared(deviceId); + return std::make_shared(deviceId); } } // namespace nvMolKit diff --git a/src/substruct/substruct_search.h b/src/substruct/substruct_search.h index d7c9e49f..bc92a44a 100644 --- a/src/substruct/substruct_search.h +++ b/src/substruct/substruct_search.h @@ -31,7 +31,7 @@ namespace nvMolKit { struct MoleculesHost; class MoleculesDevice; -struct ResidentSubstructSearchWorkspace; +struct SubstructSearchWorkspace; /** * @brief Perform batch substructure matching on GPU. @@ -87,39 +87,40 @@ void hasSubstructMatch(const std::vector& targets, const SubstructSearchConfig& config = SubstructSearchConfig{}); /** - * Targets already packed and resident on one GPU, with the per-target shapes a resident search needs. - * Immutable: build one per packed batch with makeResidentTargetBatch(). + * Targets already packed and uploaded to one GPU, with the per-target shapes hasSubstructMatch() needs. + * Immutable: build one per packed batch with makePersistentDeviceTargets(). */ -struct ResidentTargetBatch; +struct PersistentDeviceTargets; /** - * @brief Describe packed targets resident on the current GPU for hasSubstructMatchResident(). + * @brief Describe packed targets uploaded to the current GPU for hasSubstructMatch(). * * targets, targetsHost, and targetsDevice must describe the same molecules in the same order, stay alive and * unmodified while the batch is used, and targetsDevice must be fully uploaded. Targets above kMaxTargetAtoms atoms * or needing the RDKit fallback (see requiresRDKitFallback()) are rejected with std::invalid_argument. */ -std::shared_ptr makeResidentTargetBatch(const std::vector& targets, - const MoleculesHost& targetsHost, - const MoleculesDevice& targetsDevice); +std::shared_ptr makePersistentDeviceTargets( + const std::vector& targets, + const MoleculesHost& targetsHost, + const MoleculesDevice& targetsDevice); /** - * @brief Check one query against a resident target batch. + * @brief Check one query against persistent device targets. * * Runs on the GPU holding the batch; config.gpuIds must be empty or name that GPU, and a workspace must belong to * it. A workspace keeps executors and pinned buffers alive across calls and serves one call at a time. results is * resized to the batch size and holds one flag per target. */ -void hasSubstructMatchResident(const ResidentTargetBatch& batch, - const RDKit::ROMol& query, - std::vector& results, - SubstructAlgorithm algorithm, - cudaStream_t stream, - const SubstructSearchConfig& config = SubstructSearchConfig{}, - ResidentSubstructSearchWorkspace* workspace = nullptr); - -/** Create reusable search state bound to one GPU for hasSubstructMatchResident(). */ -std::shared_ptr makeResidentSubstructSearchWorkspace(int deviceId); +void hasSubstructMatch(const PersistentDeviceTargets& batch, + const RDKit::ROMol& query, + std::vector& results, + SubstructAlgorithm algorithm, + cudaStream_t stream, + const SubstructSearchConfig& config = SubstructSearchConfig{}, + SubstructSearchWorkspace* workspace = nullptr); + +/** Create reusable search state bound to one GPU for hasSubstructMatch(). */ +std::shared_ptr makeSubstructSearchWorkspace(int deviceId); } // namespace nvMolKit diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index aff15ee3..730b0f71 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -373,6 +373,10 @@ target_link_libraries( test_substructure_search PRIVATE substructure_search molecules device substruct_validation ${RDKit_LIBS}) +add_executable(test_substruct_library test_substruct_library.cu) +target_link_libraries(test_substruct_library PRIVATE substruct_library + ${RDKit_LIBS}) + add_executable(test_substruct_algos test_substruct_algos.cu) target_link_libraries( test_substruct_algos PRIVATE substruct_definitions molecules device @@ -475,6 +479,7 @@ set(TEST_LIST test_substruct_algos test_warp_dfs test_substructure_search + test_substruct_library test_substructure_integration test_substruct_label_integration test_similarity diff --git a/tests/test_substruct_library.cu b/tests/test_substruct_library.cu new file mode 100644 index 00000000..212fdd22 --- /dev/null +++ b/tests/test_substruct_library.cu @@ -0,0 +1,515 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "src/substruct/substruct_library.h" +#include "src/substruct/substruct_search.h" +#include "src/utils/cuda_error_check.h" +#include "src/utils/device.h" + +namespace { + +using MoleculeId = unsigned int; + +std::unique_ptr molFromSmiles(const std::string& smiles) { + auto mol = std::unique_ptr(RDKit::SmilesToMol(smiles)); + EXPECT_NE(mol, nullptr) << "Failed to parse SMILES: " << smiles; + return mol; +} + +std::unique_ptr queryFromSmarts(const std::string& smarts) { + auto query = std::unique_ptr(RDKit::SmartsToMol(smarts)); + EXPECT_NE(query, nullptr) << "Failed to parse SMARTS: " << smarts; + return query; +} + +std::vector rdkitMatchingIds(const std::vector>& targets, + const RDKit::ROMol& query) { + std::vector expected; + for (std::size_t targetIdx = 0; targetIdx < targets.size(); ++targetIdx) { + RDKit::MatchVectType match; + if (RDKit::SubstructMatch(*targets[targetIdx], query, match)) { + expected.push_back(static_cast(targetIdx)); + } + } + return expected; +} + +std::string linearAlkane(std::size_t numAtoms) { + return std::string(numAtoms, 'C'); +} + +//! While alive, the current device's stream-ordered allocations come from a private pool that is capped at its +//! smallest size and already full, so they fail without taking memory from other processes sharing the GPU. +class FullMemoryPool { + public: + FullMemoryPool() { + EXPECT_EQ(cudaGetDevice(&device_), cudaSuccess); + EXPECT_EQ(cudaDeviceGetMemPool(&originalPool_, device_), cudaSuccess); + cudaMemPoolProps props = {}; + props.allocType = cudaMemAllocationTypePinned; + props.location.type = cudaMemLocationTypeDevice; + props.location.id = device_; + props.maxSize = 1; + EXPECT_EQ(cudaMemPoolCreate(&pool_, &props), cudaSuccess); + for (std::size_t chunk = std::size_t{1} << 20; chunk >= 256;) { + void* allocation = nullptr; + if (cudaMallocFromPoolAsync(&allocation, chunk, pool_, nullptr) == cudaSuccess) { + filler_.push_back(allocation); + } else { + chunk /= 2; + } + } + static_cast(cudaGetLastError()); + EXPECT_EQ(cudaDeviceSetMemPool(device_, pool_), cudaSuccess); + } + FullMemoryPool(const FullMemoryPool&) = delete; + FullMemoryPool& operator=(const FullMemoryPool&) = delete; + ~FullMemoryPool() { + EXPECT_EQ(cudaDeviceSetMemPool(device_, originalPool_), cudaSuccess); + for (void* allocation : filler_) { + EXPECT_EQ(cudaFreeAsync(allocation, nullptr), cudaSuccess); + } + EXPECT_EQ(cudaStreamSynchronize(nullptr), cudaSuccess); + EXPECT_EQ(cudaMemPoolDestroy(pool_), cudaSuccess); + } + + private: + int device_ = 0; + cudaMemPool_t originalPool_ = nullptr; + cudaMemPool_t pool_ = nullptr; + std::vector filler_; +}; + +std::unique_ptr highDegreeMolecule() { + RDKit::SmilesParserParams params; + params.sanitize = false; + auto mol = std::unique_ptr(RDKit::SmilesToMol("[Fe](C)(C)(C)(C)(C)(C)(C)(C)C", params)); + if (mol) { + RDKit::MolOps::symmetrizeSSSR(*mol); + } + return mol; +} + +TEST(SubstructLibraryState, RejectsQueriesBeforeFirstFinalize) { + nvMolKit::SubstructLibrary library; + auto target = molFromSmiles("CCO"); + auto query = queryFromSmarts("CO"); + + ASSERT_NE(target, nullptr); + ASSERT_NE(query, nullptr); + EXPECT_EQ(library.addMol(*target), 0U); + EXPECT_EQ(library.size(), 0U); + EXPECT_EQ(library.pendingSize(), 1U); + + EXPECT_THROW(static_cast(library.getMatches(*query)), std::logic_error); + EXPECT_THROW(static_cast(library.countMatches(*query)), std::logic_error); + EXPECT_THROW(static_cast(library.hasMatch(*query)), std::logic_error); +} + +TEST(SubstructLibraryState, FinalizedEmptyLibraryHasEmptyResults) { + nvMolKit::SubstructLibrary library; + auto query = queryFromSmarts("C"); + + ASSERT_NE(query, nullptr); + EXPECT_NO_THROW(library.finalize()); + EXPECT_EQ(library.size(), 0U); + EXPECT_EQ(library.pendingSize(), 0U); + EXPECT_TRUE(library.getMatches(*query).empty()); + EXPECT_EQ(library.countMatches(*query), 0U); + EXPECT_FALSE(library.hasMatch(*query)); + + EXPECT_NO_THROW(library.finalize()); + EXPECT_TRUE(library.getMatches(*query).empty()); +} + +TEST(SubstructLibraryState, AddedMoleculesBecomeSearchableAtFinalize) { + nvMolKit::SubstructLibrary library; + auto ethanol = molFromSmiles("CCO"); + auto benzene = molFromSmiles("c1ccccc1"); + auto phenol = molFromSmiles("Oc1ccccc1"); + auto aromaticCarbon = queryFromSmarts("c"); + + ASSERT_NE(ethanol, nullptr); + ASSERT_NE(benzene, nullptr); + ASSERT_NE(phenol, nullptr); + ASSERT_NE(aromaticCarbon, nullptr); + + EXPECT_EQ(library.addMol(*ethanol), 0U); + EXPECT_EQ(library.addMol(*benzene), 1U); + library.finalize(); + EXPECT_EQ(library.size(), 2U); + EXPECT_EQ(library.pendingSize(), 0U); + EXPECT_EQ(library.getMatches(*aromaticCarbon), std::vector({1U})); + + // Bulk adds continue the ID sequence after single adds and vice versa. + auto pyridine = molFromSmiles("c1ccncc1"); + auto methane = molFromSmiles("C"); + ASSERT_NE(pyridine, nullptr); + ASSERT_NE(methane, nullptr); + EXPECT_EQ(library.addMol(*phenol), 2U); + EXPECT_EQ(library.addMols({pyridine.get(), methane.get()}), std::vector({3U, 4U})); + EXPECT_EQ(library.addMol(*benzene), 5U); + EXPECT_EQ(library.size(), 2U); + EXPECT_EQ(library.pendingSize(), 4U); + + EXPECT_EQ(library.getMatches(*aromaticCarbon), std::vector({1U})); + EXPECT_EQ(library.countMatches(*aromaticCarbon), 1U); + + library.finalize(); + EXPECT_EQ(library.size(), 6U); + EXPECT_EQ(library.pendingSize(), 0U); + EXPECT_EQ(library.getMatches(*aromaticCarbon), std::vector({1U, 2U, 3U, 5U})); + + library.finalize(); + EXPECT_EQ(library.size(), 6U); + EXPECT_EQ(library.getMatches(*aromaticCarbon), std::vector({1U, 2U, 3U, 5U})); +} + +TEST(SubstructLibraryState, FailedIncrementalFinalizeKeepsPreviousMoleculesAndCanBeRetried) { + nvMolKit::SubstructLibrary library; + auto benzene = molFromSmiles("c1ccccc1"); + auto phenol = molFromSmiles("Oc1ccccc1"); + auto aromaticCarbon = queryFromSmarts("c"); + ASSERT_NE(benzene, nullptr); + ASSERT_NE(phenol, nullptr); + ASSERT_NE(aromaticCarbon, nullptr); + + library.addMol(*benzene); + library.finalize(); + library.addMol(*phenol); + + { + const FullMemoryPool fullPool; + EXPECT_THROW(library.finalize(), nvMolKit::CudaBadReturnCode); + } + + EXPECT_EQ(library.size(), 1U); + EXPECT_EQ(library.pendingSize(), 1U); + EXPECT_EQ(library.getMatches(*aromaticCarbon), std::vector({0U})); + + library.finalize(); + EXPECT_EQ(library.size(), 2U); + EXPECT_EQ(library.pendingSize(), 0U); + EXPECT_EQ(library.getMatches(*aromaticCarbon), std::vector({0U, 1U})); +} + +TEST(SubstructLibraryResults, PreservesInsertionOrderAndAppliesExactLimits) { + nvMolKit::SubstructLibrary library; + std::vector> targets; + for (const auto& smiles : {"CC", "O", "CCC", "N", "c1ccccc1", "CO"}) { + targets.push_back(molFromSmiles(smiles)); + ASSERT_NE(targets.back(), nullptr); + EXPECT_EQ(library.addMol(*targets.back()), targets.size() - 1); + } + auto query = queryFromSmarts("[#6]"); + ASSERT_NE(query, nullptr); + library.finalize(); + + const std::vector allExpected{0U, 2U, 4U, 5U}; + EXPECT_EQ(library.getMatches(*query), allExpected); + EXPECT_EQ(library.getMatches(*query, 0), std::vector{}); + EXPECT_EQ(library.getMatches(*query, 1), std::vector({0U})); + EXPECT_EQ(library.getMatches(*query, 3), std::vector({0U, 2U, 4U})); + EXPECT_EQ(library.getMatches(*query, 20), allExpected); + EXPECT_EQ(library.countMatches(*query), allExpected.size()); + EXPECT_TRUE(library.hasMatch(*query)); + + auto noMatch = queryFromSmarts("[Si]"); + ASSERT_NE(noMatch, nullptr); + EXPECT_TRUE(library.getMatches(*noMatch).empty()); + EXPECT_EQ(library.countMatches(*noMatch), 0U); + EXPECT_FALSE(library.hasMatch(*noMatch)); +} + +TEST(SubstructLibraryResults, MatchesRDKitAcrossRepresentativeQuerySemantics) { + nvMolKit::SubstructLibrary library; + std::vector> targets; + for (const auto& smiles : + {"CCO", "CC(=O)C", "c1ccccc1", "C1CCCCC1", "C[N+](C)(C)C", "CC(=O)[O-]", "[Na+].[Cl-]", "CCOC(=O)c1ccccc1O"}) { + targets.push_back(molFromSmiles(smiles)); + ASSERT_NE(targets.back(), nullptr); + EXPECT_EQ(library.addMol(*targets.back()), targets.size() - 1); + } + library.finalize(); + + for (const auto& smarts : + {"[#6]", "C=O", "c1ccccc1", "[R]", "[N+]", "[#8;H1]", "[$([CX3]=[OX1])]", "[Cl-]", "[Si]"}) { + auto query = queryFromSmarts(smarts); + ASSERT_NE(query, nullptr); + const auto expected = rdkitMatchingIds(targets, *query); + + EXPECT_EQ(library.getMatches(*query), expected) << "SMARTS: " << smarts; + EXPECT_EQ(library.countMatches(*query), expected.size()) << "SMARTS: " << smarts; + EXPECT_EQ(library.hasMatch(*query), !expected.empty()) << "SMARTS: " << smarts; + } +} + +TEST(SubstructLibraryResults, ReusedWorkspaceDoesNotCarryMatchesBetweenQueries) { + nvMolKit::SubstructLibrary library; + std::vector> targets; + for (const auto& smiles : {"CCO", "c1ccccc1", "[Na+].[Cl-]", "CCN", "ClCCl", "CC(=O)O"}) { + targets.push_back(molFromSmiles(smiles)); + ASSERT_NE(targets.back(), nullptr); + library.addMol(*targets.back()); + } + library.finalize(); + + // Alternate a recursive and a plain query over the same reused workspace. + auto broadRecursive = queryFromSmarts("[$([#6])]"); + auto narrowPlain = queryFromSmarts("[Cl-]"); + ASSERT_NE(broadRecursive, nullptr); + ASSERT_NE(narrowPlain, nullptr); + const auto expectedBroad = rdkitMatchingIds(targets, *broadRecursive); + const auto expectedNarrow = rdkitMatchingIds(targets, *narrowPlain); + ASSERT_EQ(expectedNarrow, std::vector({2U})); + for (std::size_t round = 0; round < 3; ++round) { + EXPECT_EQ(library.getMatches(*broadRecursive), expectedBroad) << "round " << round; + EXPECT_EQ(library.getMatches(*narrowPlain), expectedNarrow) << "round " << round; + } +} + +TEST(SubstructLibraryOwnership, DoesNotDependOnInputMoleculeLifetime) { + nvMolKit::SubstructLibrary library; + { + auto temporary = molFromSmiles("CC(=O)Oc1ccccc1C(=O)O"); + ASSERT_NE(temporary, nullptr); + EXPECT_EQ(library.addMol(*temporary), 0U); + } + + auto query = queryFromSmarts("C(=O)O"); + ASSERT_NE(query, nullptr); + library.finalize(); + EXPECT_EQ(library.getMatches(*query), std::vector({0U})); +} + +std::vector> unpackableAndPackableTargets() { + std::vector> targets; + targets.push_back(molFromSmiles(linearAlkane(128))); // Largest packable target. + targets.push_back(molFromSmiles(linearAlkane(129))); // Too many atoms. + targets.push_back(highDegreeMolecule()); // Atom degree above the packed limit. + targets.push_back(molFromSmiles("CCN")); + targets.push_back(molFromSmiles("[NH3]->[Cu]")); // Dative bond. + targets.push_back(molFromSmiles("[300C]CN")); // Isotope above 255. + targets.push_back(molFromSmiles("c1ccccc1N")); + return targets; +} + +TEST(SubstructLibraryResults, RingQueriesMatchRDKitForTargetsWithoutRingInfo) { + // Norbornane's bridgehead atoms are in two SSSR rings. Without ring info, it should match as RDKit's + // SubstructLibrary matches it, after symmetrizing its rings. + std::vector> targets; + targets.push_back(molFromSmiles("C1CC2CCC1C2")); + targets.push_back(molFromSmiles("c1ccccc1CC")); + for (const auto& target : targets) { + ASSERT_NE(target, nullptr); + } + targets[0]->getRingInfo()->reset(); + + nvMolKit::SubstructLibrary library; + for (const auto& target : targets) { + library.addMol(*target); + } + library.finalize(); + + std::vector> symmetrized; + for (const auto& target : targets) { + symmetrized.push_back(std::make_unique(*target)); + RDKit::MolOps::symmetrizeSSSR(*symmetrized.back()); + } + for (const char* smarts : {"[R2]", "[r5]", "[R1]", "[x3]", "C1CCCC1", "[!R]"}) { + auto query = queryFromSmarts(smarts); + ASSERT_NE(query, nullptr); + EXPECT_EQ(library.getMatches(*query), rdkitMatchingIds(symmetrized, *query)) << smarts; + } +} + +TEST(SubstructLibraryFallback, MatchesUnpackableTargetsFromEitherAddPathWithRDKit) { + const auto targets = unpackableAndPackableTargets(); + std::vector targetPtrs; + for (const auto& target : targets) { + ASSERT_NE(target, nullptr); + targetPtrs.push_back(target.get()); + } + ASSERT_GT(targets[2]->getAtomWithIdx(0)->getDegree(), 8U); + + nvMolKit::SubstructLibrary bulkLibrary; + EXPECT_EQ(bulkLibrary.addMols(targetPtrs), std::vector({0U, 1U, 2U, 3U, 4U, 5U, 6U})); + bulkLibrary.finalize(); + + nvMolKit::SubstructLibrary singleLibrary; + for (const auto* target : targetPtrs) { + singleLibrary.addMol(*target); + } + singleLibrary.finalize(); + + for (const char* smarts : {"CCCC", "[Fe]", "[#7]", "[Cu]", "[#6]-[#7]", "c"}) { + auto query = queryFromSmarts(smarts); + ASSERT_NE(query, nullptr); + const auto expected = rdkitMatchingIds(targets, *query); + for (auto* library : {&bulkLibrary, &singleLibrary}) { + EXPECT_EQ(library->getMatches(*query), expected) << smarts; + EXPECT_EQ(library->countMatches(*query), expected.size()) << smarts; + EXPECT_EQ(library->hasMatch(*query), !expected.empty()) << smarts; + } + } +} + +TEST(SubstructLibraryFallback, ResultLimitsMergeGpuAndFallbackMatchesInIdOrder) { + // GPU and RDKit-fallback matches alternate by ID. + std::vector> targets; + for (const char* smiles : {"[NH3]->[Cu]", "CCN", "[300C]CN", "CC", "NCCN", "[NH3]->[Zn]", "c1ccncc1"}) { + targets.push_back(molFromSmiles(smiles)); + ASSERT_NE(targets.back(), nullptr); + } + nvMolKit::SubstructLibrary library; + for (const auto& target : targets) { + library.addMol(*target); + } + library.finalize(); + + auto nitrogen = queryFromSmarts("[#7]"); + ASSERT_NE(nitrogen, nullptr); + const auto expected = rdkitMatchingIds(targets, *nitrogen); + ASSERT_EQ(expected, std::vector({0U, 1U, 2U, 4U, 5U, 6U})); + for (int limit = 1; limit <= static_cast(expected.size()) + 1; ++limit) { + const auto prefix = + std::vector(expected.begin(), expected.begin() + std::min(limit, expected.size())); + EXPECT_EQ(library.getMatches(*nitrogen, limit), prefix) << "limit " << limit; + } + EXPECT_EQ(library.countMatches(*nitrogen), expected.size()); +} + +TEST(SubstructLibraryFallback, LibraryOfOnlyFallbackTargetsAnswersQueries) { + nvMolKit::SubstructLibrary library; + std::vector> targets; + for (const char* smiles : {"[NH3]->[Cu]", "[300C]CN", "[NH3]->[Zn]"}) { + targets.push_back(molFromSmiles(smiles)); + ASSERT_NE(targets.back(), nullptr); + library.addMol(*targets.back()); + } + library.finalize(); + + for (const char* smarts : {"[#7]", "[Cu]", "[Si]"}) { + auto query = queryFromSmarts(smarts); + ASSERT_NE(query, nullptr); + const auto expected = rdkitMatchingIds(targets, *query); + EXPECT_EQ(library.getMatches(*query), expected) << smarts; + EXPECT_EQ(library.getMatches(*query, 1), + std::vector(expected.begin(), expected.begin() + std::min(1, expected.size()))) + << smarts; + EXPECT_EQ(library.countMatches(*query), expected.size()) << smarts; + EXPECT_EQ(library.hasMatch(*query), !expected.empty()) << smarts; + } +} + +TEST(SubstructLibraryConfiguration, SupportsProductionBackendsAndRejectsInvalidConstruction) { + nvMolKit::SubstructSearchConfig vf2Config; + vf2Config.algorithm = nvMolKit::SubstructAlgorithm::VF2; + EXPECT_THROW(nvMolKit::SubstructLibrary{vf2Config}, std::invalid_argument); + + nvMolKit::SubstructSearchConfig duplicateGpuConfig; + duplicateGpuConfig.gpuIds = {0, 0}; + EXPECT_THROW(nvMolKit::SubstructLibrary{duplicateGpuConfig}, std::invalid_argument); + + nvMolKit::SubstructSearchConfig invalidGpuConfig; + invalidGpuConfig.gpuIds = {std::numeric_limits::max()}; + nvMolKit::SubstructLibrary invalidGpuLibrary(invalidGpuConfig); + EXPECT_THROW(invalidGpuLibrary.finalize(), std::invalid_argument); + + auto target = molFromSmiles("CCOC(=O)C"); + auto query = queryFromSmarts("C=O"); + ASSERT_NE(target, nullptr); + ASSERT_NE(query, nullptr); + + for (const auto algorithm : {nvMolKit::SubstructAlgorithm::GSI, nvMolKit::SubstructAlgorithm::DFS}) { + nvMolKit::SubstructSearchConfig config; + config.algorithm = algorithm; + config.workerThreads = 1; + config.preprocessingThreads = 1; + nvMolKit::SubstructLibrary library(config); + EXPECT_EQ(library.addMol(*target), 0U); + library.finalize(); + EXPECT_EQ(library.getMatches(*query), std::vector({0U})); + } +} + +TEST(SubstructLibraryMultiGpu, ShardsTargetsAndMergesEveryOperationInInsertionOrder) { + int deviceCount = 0; + ASSERT_EQ(cudaGetDeviceCount(&deviceCount), cudaSuccess); + if (deviceCount < 2) { + GTEST_SKIP() << "Multi-GPU library test requires at least two CUDA devices"; + } + + nvMolKit::SubstructSearchConfig config; + config.gpuIds = {0, 1}; + config.workerThreads = 1; + config.preprocessingThreads = 4; + nvMolKit::SubstructLibrary library(config); + std::vector> targets; + std::vector pointers; + for (const auto& smiles : {"CC", "O", "CCC", "N", "c1ccccc1", "CO", "C=O", "[Si]"}) { + targets.push_back(molFromSmiles(smiles)); + ASSERT_NE(targets.back(), nullptr); + pointers.push_back(targets.back().get()); + } + EXPECT_EQ(library.addMols(pointers), std::vector({0U, 1U, 2U, 3U, 4U, 5U, 6U, 7U})); + library.finalize(); + + auto carbon = queryFromSmarts("[#6]"); + ASSERT_NE(carbon, nullptr); + const std::vector expected{0U, 2U, 4U, 5U, 6U}; + EXPECT_EQ(library.getMatches(*carbon), expected); + EXPECT_EQ(library.getMatches(*carbon, 3), std::vector({0U, 2U, 4U})); + EXPECT_EQ(library.countMatches(*carbon), expected.size()); + EXPECT_TRUE(library.hasMatch(*carbon)); + + auto phosphorus = queryFromSmarts("[P]"); + ASSERT_NE(phosphorus, nullptr); + EXPECT_TRUE(library.getMatches(*phosphorus).empty()); + EXPECT_EQ(library.countMatches(*phosphorus), 0U); + EXPECT_FALSE(library.hasMatch(*phosphorus)); +} + +TEST(SubstructLibraryStreams, SupportsFinalizeAndQueriesOnANondefaultStream) { + nvMolKit::ScopedStream stream("SubstructLibraryStreams"); + nvMolKit::SubstructLibrary library; + auto target = molFromSmiles("CCOC(=O)C"); + auto query = queryFromSmarts("C=O"); + ASSERT_NE(target, nullptr); + ASSERT_NE(query, nullptr); + + EXPECT_EQ(library.addMol(*target), 0U); + library.finalize(stream.stream()); + EXPECT_EQ(library.getMatches(*query, -1, stream.stream()), std::vector({0U})); + EXPECT_EQ(library.countMatches(*query, stream.stream()), 1U); + EXPECT_TRUE(library.hasMatch(*query, stream.stream())); +} + +} // namespace diff --git a/tests/test_substruct_search.cu b/tests/test_substruct_search.cu index e551cc1d..44509818 100644 --- a/tests/test_substruct_search.cu +++ b/tests/test_substruct_search.cu @@ -2123,16 +2123,16 @@ TEST_P(RecursiveSubstructureSearchTest, DeepRecursionMixedWithNormalQueries) { } } -class ResidentSubstructureSearchTest : public SubstructureSearchTest {}; +class PersistentTargetSubstructureSearchTest : public SubstructureSearchTest {}; -INSTANTIATE_TEST_SUITE_P(ResidentAlgorithms, - ResidentSubstructureSearchTest, +INSTANTIATE_TEST_SUITE_P(PersistentTargetAlgorithms, + PersistentTargetSubstructureSearchTest, ::testing::Values(SubstructAlgorithm::GSI, SubstructAlgorithm::DFS), [](const ::testing::TestParamInfo& info) { return algorithmName(info.param); }); -TEST_P(ResidentSubstructureSearchTest, MatchesBatchSearchAcrossCallsSharingAWorkspace) { +TEST_P(PersistentTargetSubstructureSearchTest, MatchesBatchSearchAcrossCallsSharingAWorkspace) { std::vector> targetMols; std::vector> queryMols; parseMolecules({"CCO", "c1ccccc1O", "CC(=O)O", "CC(=O)NC", "c1ccncc1", "CCCCCC", "OC(=O)c1ccccc1N", "C1CCNCC1"}, @@ -2149,11 +2149,11 @@ TEST_P(ResidentSubstructureSearchTest, MatchesBatchSearchAcrossCallsSharingAWork nvMolKit::MoleculesDevice targetsDevice(stream_.stream()); targetsDevice.copyFromHost(targetsHost); cudaCheckError(cudaStreamSynchronize(stream_.stream())); - const auto batch = nvMolKit::makeResidentTargetBatch(targets, targetsHost, targetsDevice); + const auto batch = nvMolKit::makePersistentDeviceTargets(targets, targetsHost, targetsDevice); int deviceId = 0; cudaCheckError(cudaGetDevice(&deviceId)); - const auto workspace = nvMolKit::makeResidentSubstructSearchWorkspace(deviceId); + const auto workspace = nvMolKit::makeSubstructSearchWorkspace(deviceId); // Both result buffers are reused across queries without clearing, so stale flags would show up as mismatches. std::vector withWorkspace; @@ -2162,14 +2162,14 @@ TEST_P(ResidentSubstructureSearchTest, MatchesBatchSearchAcrossCallsSharingAWork HasSubstructMatchResults expected; hasSubstructMatch(targets, {query.get()}, expected, algorithm(), stream_.stream()); - nvMolKit::hasSubstructMatchResident(*batch, - *query, - withWorkspace, - algorithm(), - stream_.stream(), - SubstructSearchConfig{}, - workspace.get()); - nvMolKit::hasSubstructMatchResident(*batch, *query, withoutWorkspace, algorithm(), stream_.stream()); + nvMolKit::hasSubstructMatch(*batch, + *query, + withWorkspace, + algorithm(), + stream_.stream(), + SubstructSearchConfig{}, + workspace.get()); + nvMolKit::hasSubstructMatch(*batch, *query, withoutWorkspace, algorithm(), stream_.stream()); ASSERT_EQ(withWorkspace.size(), targets.size()); ASSERT_EQ(withoutWorkspace.size(), targets.size()); @@ -2181,7 +2181,7 @@ TEST_P(ResidentSubstructureSearchTest, MatchesBatchSearchAcrossCallsSharingAWork } } -TEST_P(ResidentSubstructureSearchTest, RejectsMismatchedAndOversizedTargetBatches) { +TEST_P(PersistentTargetSubstructureSearchTest, RejectsMismatchedAndOversizedTargetBatches) { std::vector> targetMols; std::vector> queryMols; parseMolecules({std::string(65, 'C'), std::string(65, 'C')}, {"CC"}, targetMols, queryMols); @@ -2190,12 +2190,12 @@ TEST_P(ResidentSubstructureSearchTest, RejectsMismatchedAndOversizedTargetBatche nvMolKit::buildTargetBatchParallelInto(targetsHost, 1, targets, {}); nvMolKit::MoleculesDevice targetsDevice(stream_.stream()); - EXPECT_THROW(nvMolKit::makeResidentTargetBatch({targets[0]}, targetsHost, targetsDevice), std::invalid_argument); + EXPECT_THROW(nvMolKit::makePersistentDeviceTargets({targets[0]}, targetsHost, targetsDevice), std::invalid_argument); // The packer refuses targets above the atom limit, so join two packed targets into one 130-atom entry, as a // caller packing targets by other means could. targetsHost.batchAtomStarts.erase(targetsHost.batchAtomStarts.begin() + 1); - EXPECT_THROW(nvMolKit::makeResidentTargetBatch({targets[0]}, targetsHost, targetsDevice), std::invalid_argument); + EXPECT_THROW(nvMolKit::makePersistentDeviceTargets({targets[0]}, targetsHost, targetsDevice), std::invalid_argument); // A target that needs the RDKit fallback, here an atom above the packed bond limit, is rejected up front. RDKit::SmilesParserParams params; @@ -2205,11 +2205,11 @@ TEST_P(ResidentSubstructureSearchTest, RejectsMismatchedAndOversizedTargetBatche RDKit::MolOps::symmetrizeSSSR(*hypervalent); nvMolKit::MoleculesHost singleHost; nvMolKit::buildTargetBatchParallelInto(singleHost, 1, {targets[0]}, {}); - EXPECT_THROW(nvMolKit::makeResidentTargetBatch({hypervalent.get()}, singleHost, targetsDevice), + EXPECT_THROW(nvMolKit::makePersistentDeviceTargets({hypervalent.get()}, singleHost, targetsDevice), std::invalid_argument); } -TEST_P(ResidentSubstructureSearchTest, RejectsSearchesConfiguredForAnotherGpu) { +TEST_P(PersistentTargetSubstructureSearchTest, RejectsSearchesConfiguredForAnotherGpu) { std::vector> targetMols; std::vector> queryMols; parseMolecules({"CCO"}, {"CO"}, targetMols, queryMols); @@ -2219,14 +2219,13 @@ TEST_P(ResidentSubstructureSearchTest, RejectsSearchesConfiguredForAnotherGpu) { nvMolKit::MoleculesDevice targetsDevice(stream_.stream()); targetsDevice.copyFromHost(targetsHost); cudaCheckError(cudaStreamSynchronize(stream_.stream())); - const auto batch = nvMolKit::makeResidentTargetBatch(targets, targetsHost, targetsDevice); + const auto batch = nvMolKit::makePersistentDeviceTargets(targets, targetsHost, targetsDevice); int deviceId = 0; cudaCheckError(cudaGetDevice(&deviceId)); SubstructSearchConfig otherGpu; otherGpu.gpuIds = {deviceId + 1}; std::vector results; - EXPECT_THROW( - nvMolKit::hasSubstructMatchResident(*batch, *queryMols[0], results, algorithm(), stream_.stream(), otherGpu), - std::invalid_argument); + EXPECT_THROW(nvMolKit::hasSubstructMatch(*batch, *queryMols[0], results, algorithm(), stream_.stream(), otherGpu), + std::invalid_argument); }