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); }