diff --git a/src/mcs/fmcs_cuda/fmcs_match.cuh b/src/mcs/fmcs_cuda/fmcs_match.cuh index b800ca40..be6f679a 100644 --- a/src/mcs/fmcs_cuda/fmcs_match.cuh +++ b/src/mcs/fmcs_cuda/fmcs_match.cuh @@ -318,6 +318,596 @@ __device__ __forceinline__ bool tryMatchIncrementalGreedyCooperative( return true; } +__device__ __forceinline__ bool findTargetBondBetweenAtomsWithinThread(const int targetAtomA, + const int targetAtomB, + const int queryBondIdx, + const DeviceCsrView& targetTopology, + const PairMatchTablesDevice& tables, + int& outTargetBondIdx) { + if (targetAtomA < 0 || targetAtomA >= targetTopology.numAtoms) { + return false; + } + const int begin = static_cast(targetTopology.rowOffsets[targetAtomA]); + const int end = static_cast(targetTopology.rowOffsets[targetAtomA + 1]); + for (int adjIdx = begin; adjIdx < end; ++adjIdx) { + const int otherTargetAtom = static_cast(targetTopology.colIndices[adjIdx]); + if (otherTargetAtom != targetAtomB) + continue; + const int targetBondIdx = static_cast(targetTopology.bondIndices[adjIdx]); + if (targetBondIdx < 0 || targetBondIdx >= targetTopology.numBonds) { + continue; + } + if (!tables.bonds.testBit(queryBondIdx, targetBondIdx)) + continue; + outTargetBondIdx = targetBondIdx; + return true; + } + return false; +} + +template +__device__ __forceinline__ bool rebuildMatchFromSubstructureMappingWithinThread( + const Seed& seed, + const DeviceCsrView& queryTopology, + const DeviceCsrView& targetTopology, + const PairMatchTablesDevice& tables, + MatchResult& match, + FmcsSubstructureScratch& scratch) { + using SeedT = Seed; + using MatchT = MatchResult; + using BondWord = typename SeedT::bond_word_type; + using TargetBondWord = typename MatchT::target_bond_word; + using TargetAtomWord = typename MatchT::target_atom_word; + + constexpr int kBondBitsPerWord = SeedT::kBondBitsPerWord; + constexpr int kBondWords = SeedT::kBondWords; + constexpr int kTargetAtomBitsPerWord = MatchT::kTargetAtomBitsPerWord; + constexpr int kTargetBondBitsPerWord = MatchT::kTargetBondBitsPerWord; + + matchResultClearWithinThread(match); + for (int i = 0; i < seed.numAtoms; ++i) { + const int queryAtomIdx = scratch.seedAtomList[i]; + const int targetAtomIdx = scratch.targetAtomForQuery[queryAtomIdx]; + if (targetAtomIdx == kUnmappedTargetIdx) + return false; + match.targetAtomIdx[queryAtomIdx] = static_cast(targetAtomIdx); + match.visitedTargetAtoms[targetAtomIdx / kTargetAtomBitsPerWord] |= static_cast(1) + << (targetAtomIdx % kTargetAtomBitsPerWord); + } + match.matchedAtomSize = seed.numAtoms; + + int matchedBondCount = 0; + for (int wordIdx = 0; wordIdx < kBondWords; ++wordIdx) { + BondWord remaining = seed.bonds[wordIdx]; + while (remaining != 0) { + int bitPosInWord; + if constexpr (sizeof(BondWord) == 4) { + bitPosInWord = __ffs(static_cast(remaining)) - 1; + } else { + bitPosInWord = __ffsll(static_cast(remaining)) - 1; + } + const int queryBondIdx = wordIdx * kBondBitsPerWord + bitPosInWord; + remaining &= remaining - 1; + + const std::uint32_t queryEndpoints = queryTopology.bondEndpoints[queryBondIdx]; + const int queryEndpointU = static_cast(queryEndpoints >> kBondEndpointShift); + const int queryEndpointV = static_cast(queryEndpoints & kBondEndpointMask); + const int targetEndpointU = match.targetAtomIdx[queryEndpointU]; + const int targetEndpointV = match.targetAtomIdx[queryEndpointV]; + if (targetEndpointU == kUnmappedTargetIdx || targetEndpointV == kUnmappedTargetIdx) { + matchResultClearWithinThread(match); + return false; + } + + int targetBondIdx = -1; + if (!findTargetBondBetweenAtomsWithinThread(targetEndpointU, + targetEndpointV, + queryBondIdx, + targetTopology, + tables, + targetBondIdx)) { + matchResultClearWithinThread(match); + return false; + } + const TargetBondWord visitedWord = match.visitedTargetBonds[targetBondIdx / kTargetBondBitsPerWord]; + if ((visitedWord >> (targetBondIdx % kTargetBondBitsPerWord)) & 1) { + matchResultClearWithinThread(match); + return false; + } + match.targetBondIdx[queryBondIdx] = static_cast(targetBondIdx); + match.visitedTargetBonds[targetBondIdx / kTargetBondBitsPerWord] |= static_cast(1) + << (targetBondIdx % kTargetBondBitsPerWord); + ++matchedBondCount; + } + } + match.matchedBondSize = static_cast(matchedBondCount); + if (matchedBondCount != seed.numBonds) { + matchResultClearWithinThread(match); + return false; + } + match.empty = false; + return true; +} + +template +__device__ __forceinline__ void initializeSeedSubstructureScratchCooperative( + const GroupT& group, + const DeviceCsrView& targetTopology, + FmcsSubstructureScratch& scratch) { + const int laneRank = static_cast(group.thread_rank()); + const int laneCount = static_cast(group.num_threads()); + + for (int i = laneRank; i < maxAtoms; i += laneCount) { + scratch.seedDegree[i] = 0; + scratch.orderedQueryAtom[i] = 0; + scratch.queryOrderPos[i] = kUnmappedTargetIdx; + scratch.targetAtomForQuery[i] = kUnmappedTargetIdx; + } + + for (int targetAtomIdx = laneRank; targetAtomIdx < targetTopology.numAtoms; targetAtomIdx += laneCount) { + scratch.targetDegree[targetAtomIdx] = static_cast(targetTopology.rowOffsets[targetAtomIdx + 1] - + targetTopology.rowOffsets[targetAtomIdx]); + } + + if (laneRank == 0) { + scratch.currentCount = 0; + scratch.nextCount = 0; + scratch.found = 0; + scratch.overflowed = 0; + } + group.sync(); +} + +template +__device__ __forceinline__ bool prepareSeedSubstructureSearchWithinThread( + const Seed& seed, + const DeviceCsrView& queryTopology, + const DeviceCsrView& targetTopology, + const PairMatchTablesDevice& tables, + FmcsSubstructureScratch& scratch, + int& numSeedAtoms) { + using SeedT = Seed; + using AtomWord = typename SeedT::atom_word_type; + using BondWord = typename SeedT::bond_word_type; + + constexpr int kAtomBitsPerWord = SeedT::kAtomBitsPerWord; + constexpr int kAtomWords = SeedT::kAtomWords; + constexpr int kBondBitsPerWord = SeedT::kBondBitsPerWord; + constexpr int kBondWords = SeedT::kBondWords; + + if (seed.numAtoms > targetTopology.numAtoms || seed.numBonds > targetTopology.numBonds) { + return false; + } + + numSeedAtoms = 0; + for (int wordIdx = 0; wordIdx < kAtomWords; ++wordIdx) { + AtomWord remaining = seed.atoms[wordIdx]; + while (remaining != 0) { + int bitPosInWord; + if constexpr (sizeof(AtomWord) == 4) { + bitPosInWord = __ffs(static_cast(remaining)) - 1; + } else { + bitPosInWord = __ffsll(static_cast(remaining)) - 1; + } + const int queryAtomIdx = wordIdx * kAtomBitsPerWord + bitPosInWord; + remaining &= remaining - 1; + if (queryAtomIdx < queryTopology.numAtoms) { + scratch.seedAtomList[numSeedAtoms++] = static_cast(queryAtomIdx); + } + } + } + if (numSeedAtoms != seed.numAtoms) + return false; + + for (int wordIdx = 0; wordIdx < kBondWords; ++wordIdx) { + BondWord remaining = seed.bonds[wordIdx]; + while (remaining != 0) { + int bitPosInWord; + if constexpr (sizeof(BondWord) == 4) { + bitPosInWord = __ffs(static_cast(remaining)) - 1; + } else { + bitPosInWord = __ffsll(static_cast(remaining)) - 1; + } + const int queryBondIdx = wordIdx * kBondBitsPerWord + bitPosInWord; + remaining &= remaining - 1; + + const std::uint32_t queryEndpoints = queryTopology.bondEndpoints[queryBondIdx]; + const int queryEndpointU = static_cast(queryEndpoints >> kBondEndpointShift); + const int queryEndpointV = static_cast(queryEndpoints & kBondEndpointMask); + ++scratch.seedDegree[queryEndpointU]; + ++scratch.seedDegree[queryEndpointV]; + } + } + + for (int orderPos = 0; orderPos < numSeedAtoms; ++orderPos) { + int bestAtom = -1; + int bestMappedNeighborCount = -1; + int bestDegree = -1; + int bestCandidateCount = maxTA + 1; + + for (int atomListIdx = 0; atomListIdx < numSeedAtoms; ++atomListIdx) { + const int queryAtomIdx = scratch.seedAtomList[atomListIdx]; + if (scratch.orderedQueryAtom[queryAtomIdx]) + continue; + + int mappedNeighborCount = 0; + if (orderPos > 0) { + const int begin = static_cast(queryTopology.rowOffsets[queryAtomIdx]); + const int end = static_cast(queryTopology.rowOffsets[queryAtomIdx + 1]); + for (int adjIdx = begin; adjIdx < end; ++adjIdx) { + const int queryBondIdx = static_cast(queryTopology.bondIndices[adjIdx]); + if (queryBondIdx >= queryTopology.numBonds || + !seedContainsBondWithinThread(seed, queryBondIdx)) { + continue; + } + const int otherQueryAtom = static_cast(queryTopology.colIndices[adjIdx]); + if (otherQueryAtom >= 0 && otherQueryAtom < queryTopology.numAtoms && + scratch.orderedQueryAtom[otherQueryAtom]) { + ++mappedNeighborCount; + } + } + } + if (orderPos > 0 && mappedNeighborCount == 0) + continue; + + int candidateCount = 0; + for (int targetAtomIdx = 0; targetAtomIdx < targetTopology.numAtoms; ++targetAtomIdx) { + if (scratch.targetDegree[targetAtomIdx] < scratch.seedDegree[queryAtomIdx]) { + continue; + } + if (!tables.atoms.testBit(queryAtomIdx, targetAtomIdx)) + continue; + ++candidateCount; + } + if (candidateCount == 0) + return false; + + const int degree = scratch.seedDegree[queryAtomIdx]; + const bool better = bestAtom < 0 || mappedNeighborCount > bestMappedNeighborCount || + (mappedNeighborCount == bestMappedNeighborCount && degree > bestDegree) || + (mappedNeighborCount == bestMappedNeighborCount && degree == bestDegree && + candidateCount < bestCandidateCount) || + (mappedNeighborCount == bestMappedNeighborCount && degree == bestDegree && + candidateCount == bestCandidateCount && queryAtomIdx < bestAtom); + if (better) { + bestAtom = queryAtomIdx; + bestMappedNeighborCount = mappedNeighborCount; + bestDegree = degree; + bestCandidateCount = candidateCount; + } + } + + if (bestAtom < 0) { + for (int atomListIdx = 0; atomListIdx < numSeedAtoms; ++atomListIdx) { + const int queryAtomIdx = scratch.seedAtomList[atomListIdx]; + if (!scratch.orderedQueryAtom[queryAtomIdx]) { + bestAtom = queryAtomIdx; + break; + } + } + } + if (bestAtom < 0) + return false; + scratch.seedAtoms[orderPos] = static_cast(bestAtom); + scratch.orderedQueryAtom[bestAtom] = 1; + scratch.queryOrderPos[bestAtom] = static_cast(orderPos); + } + return true; +} + +__device__ __forceinline__ bool partialUsesTargetAtomWithinThread(const std::uint8_t* partial, + const int depth, + const int targetAtomIdx) { + for (int i = 0; i < depth; ++i) { + if (partial[i] == targetAtomIdx) + return true; + } + return false; +} + +template +__device__ __forceinline__ bool findMappedQueryNeighborWithinThread( + const Seed& seed, + const DeviceCsrView& queryTopology, + const FmcsSubstructureScratch& scratch, + const int depth, + const int queryAtomIdx, + int& outNeighborOrderPos) { + outNeighborOrderPos = -1; + const int begin = static_cast(queryTopology.rowOffsets[queryAtomIdx]); + const int end = static_cast(queryTopology.rowOffsets[queryAtomIdx + 1]); + for (int adjIdx = begin; adjIdx < end; ++adjIdx) { + const int queryBondIdx = static_cast(queryTopology.bondIndices[adjIdx]); + if (queryBondIdx >= queryTopology.numBonds || + !seedContainsBondWithinThread(seed, queryBondIdx)) { + continue; + } + const int otherQueryAtom = static_cast(queryTopology.colIndices[adjIdx]); + if (otherQueryAtom < 0 || otherQueryAtom >= queryTopology.numAtoms) { + continue; + } + const int otherOrderPos = scratch.queryOrderPos[otherQueryAtom]; + if (otherOrderPos != kUnmappedTargetIdx && otherOrderPos < depth) { + outNeighborOrderPos = otherOrderPos; + return true; + } + } + return false; +} + +template +__device__ __forceinline__ bool substructurePartialEdgeConsistentWithinThread( + const Seed& seed, + const DeviceCsrView& queryTopology, + const DeviceCsrView& targetTopology, + const PairMatchTablesDevice& tables, + const FmcsSubstructureScratch& scratch, + const std::uint8_t* partial, + const int depth, + const int queryAtomIdx, + const int targetAtomIdx) { + const int begin = static_cast(queryTopology.rowOffsets[queryAtomIdx]); + const int end = static_cast(queryTopology.rowOffsets[queryAtomIdx + 1]); + for (int adjIdx = begin; adjIdx < end; ++adjIdx) { + const int queryBondIdx = static_cast(queryTopology.bondIndices[adjIdx]); + if (queryBondIdx >= queryTopology.numBonds || + !seedContainsBondWithinThread(seed, queryBondIdx)) { + continue; + } + const int otherQueryAtom = static_cast(queryTopology.colIndices[adjIdx]); + if (otherQueryAtom < 0 || otherQueryAtom >= queryTopology.numAtoms) { + continue; + } + + const int otherOrderPos = scratch.queryOrderPos[otherQueryAtom]; + if (otherOrderPos == kUnmappedTargetIdx || otherOrderPos >= depth) { + continue; + } + const int otherTargetAtom = partial[otherOrderPos]; + + int targetBondIdx = -1; + if (!findTargetBondBetweenAtomsWithinThread(targetAtomIdx, + otherTargetAtom, + queryBondIdx, + targetTopology, + tables, + targetBondIdx)) { + return false; + } + } + return true; +} + +template +__device__ __forceinline__ bool matchSeedSubstructureCooperative(const GroupT& group, + const Seed& seed, + const DeviceCsrView& queryTopology, + const DeviceCsrView& targetTopology, + const PairMatchTablesDevice& tables, + MatchResult& match, + FmcsSubstructureScratch& scratch, + std::uint8_t* partialStorage, + int partialCapacity, + bool* overflowedFlag) { + const int laneRank = static_cast(group.thread_rank()); + const int laneCount = static_cast(group.num_threads()); + + int numSeedAtoms = 0; + int prepared = 0; + if (laneRank == 0) { + matchResultClearWithinThread(match); + } + if (seed.numAtoms != 0) { + initializeSeedSubstructureScratchCooperative(group, targetTopology, scratch); + } + if (laneRank == 0) { + if (seed.numAtoms == 0) { + prepared = (seed.numBonds == 0) ? 2 : -1; + } else { + prepared = + prepareSeedSubstructureSearchWithinThread(seed, queryTopology, targetTopology, tables, scratch, numSeedAtoms) ? + 1 : + -1; + } + if ((partialStorage == nullptr || partialCapacity <= 0) && prepared == 1) { + scratch.overflowed = 1; + prepared = -1; + } + } + group.sync(); + prepared = group.shfl(prepared, 0); + numSeedAtoms = group.shfl(numSeedAtoms, 0); + if (prepared == 2) + return true; + if (prepared < 0) { + if (laneRank == 0 && scratch.overflowed && overflowedFlag != nullptr) { + *overflowedFlag = true; + } + return false; + } + + const int stride = numSeedAtoms; + const int halfBytes = partialCapacity * maxAtoms; + const int effectiveCapacity = halfBytes / stride; + std::uint8_t* currentPartials = partialStorage; + std::uint8_t* nextPartials = partialStorage + halfBytes; + + const int firstQueryAtom = scratch.seedAtoms[0]; + for (int targetAtomIdx = laneRank; targetAtomIdx < targetTopology.numAtoms; targetAtomIdx += laneCount) { + if (scratch.targetDegree[targetAtomIdx] < scratch.seedDegree[firstQueryAtom]) { + continue; + } + if (!tables.atoms.testBit(firstQueryAtom, targetAtomIdx)) + continue; + + const int slot = atomicAdd(&scratch.currentCount, 1); + if (slot < effectiveCapacity) { + currentPartials[slot * stride] = static_cast(targetAtomIdx); + } else { + atomicExch(&scratch.overflowed, 1); + } + } + group.sync(); + + if (numSeedAtoms == 1) { + if (scratch.currentCount > 0) { + if (laneRank == 0) { + scratch.targetAtomForQuery[firstQueryAtom] = currentPartials[0]; + prepared = + rebuildMatchFromSubstructureMappingWithinThread(seed, queryTopology, targetTopology, tables, match, scratch) ? + 1 : + -1; + } + group.sync(); + prepared = group.shfl(prepared, 0); + return prepared == 1; + } + return false; + } + + for (int depth = 1; depth < numSeedAtoms; ++depth) { + if (scratch.currentCount == 0 || scratch.found != 0) + break; + if (laneRank == 0) + scratch.nextCount = 0; + group.sync(); + + const int queryAtomIdx = scratch.seedAtoms[depth]; + const int numPartials = scratch.currentCount < effectiveCapacity ? scratch.currentCount : effectiveCapacity; + + for (int partialIdx = 0; partialIdx < numPartials && scratch.found == 0; ++partialIdx) { + const std::uint8_t* partial = currentPartials + partialIdx * stride; + int neighborOrderPos = -1; + const bool hasMappedNeighbor = + findMappedQueryNeighborWithinThread(seed, queryTopology, scratch, depth, queryAtomIdx, neighborOrderPos); + const bool scanAdjacency = + hasMappedNeighbor && targetTopology.rowOffsets != nullptr && targetTopology.colIndices != nullptr; + const int targetScanBegin = + scanAdjacency ? static_cast(targetTopology.rowOffsets[partial[neighborOrderPos]]) : 0; + const int targetScanEnd = scanAdjacency ? + static_cast(targetTopology.rowOffsets[partial[neighborOrderPos] + 1]) : + targetTopology.numAtoms; + + for (int targetScanIdx = targetScanBegin + laneRank; targetScanIdx < targetScanEnd && scratch.found == 0; + targetScanIdx += laneCount) { + const int targetAtomIdx = + scanAdjacency ? static_cast(targetTopology.colIndices[targetScanIdx]) : targetScanIdx; + if (scratch.targetDegree[targetAtomIdx] < scratch.seedDegree[queryAtomIdx]) { + continue; + } + if (!tables.atoms.testBit(queryAtomIdx, targetAtomIdx)) + continue; + if (partialUsesTargetAtomWithinThread(partial, depth, targetAtomIdx)) { + continue; + } + if (!substructurePartialEdgeConsistentWithinThread(seed, + queryTopology, + targetTopology, + tables, + scratch, + partial, + depth, + queryAtomIdx, + targetAtomIdx)) { + continue; + } + + if (depth == numSeedAtoms - 1) { + if (atomicCAS(&scratch.found, 0, 1) == 0) { + for (int orderPos = 0; orderPos < depth; ++orderPos) { + const int mappedQueryAtom = scratch.seedAtoms[orderPos]; + scratch.targetAtomForQuery[mappedQueryAtom] = partial[orderPos]; + } + scratch.targetAtomForQuery[queryAtomIdx] = static_cast(targetAtomIdx); + } + } else { + const int slot = atomicAdd(&scratch.nextCount, 1); + if (slot < effectiveCapacity) { + std::uint8_t* next = nextPartials + slot * stride; + for (int orderPos = 0; orderPos < depth; ++orderPos) { + next[orderPos] = partial[orderPos]; + } + next[depth] = static_cast(targetAtomIdx); + } else { + atomicExch(&scratch.overflowed, 1); + } + } + } + } + group.sync(); + + if (scratch.found != 0) + break; + if (laneRank == 0) + scratch.currentCount = scratch.nextCount; + std::uint8_t* tmp = currentPartials; + currentPartials = nextPartials; + nextPartials = tmp; + group.sync(); + } + + int found = scratch.found; + if (found != 0) { + if (laneRank == 0) { + found = + rebuildMatchFromSubstructureMappingWithinThread(seed, queryTopology, targetTopology, tables, match, scratch) ? + 1 : + 0; + } + group.sync(); + found = group.shfl(found, 0); + return found != 0; + } + + if (laneRank == 0) { + matchResultClearWithinThread(match); + if (scratch.overflowed && overflowedFlag != nullptr) { + *overflowedFlag = true; + } + } + group.sync(); + return false; +} + +template +__device__ __forceinline__ bool matchSeedWithSubstructureFallbackCooperative( + const GroupT& group, + const Seed& seed, + const DeviceCsrView& queryTopology, + const DeviceCsrView& targetTopology, + const PairMatchTablesDevice& tables, + MatchResult& match, + FmcsSubstructureScratch& scratch, + int* scratchLock, + std::uint8_t* partialStorage, + int partialCapacity, + bool* overflowedFlag) { + if (tryMatchIncrementalGreedyCooperative(group, seed, queryTopology, targetTopology, tables, match)) { + return true; + } + if (group.thread_rank() == 0) { + while (atomicCAS(scratchLock, 0, 1) != 0) { + } + } + group.sync(); + const bool ok = matchSeedSubstructureCooperative(group, + seed, + queryTopology, + targetTopology, + tables, + match, + scratch, + partialStorage, + partialCapacity, + overflowedFlag); + group.sync(); + if (group.thread_rank() == 0) { + atomicExch(scratchLock, 0); + } + group.sync(); + return ok; +} + } // namespace fmcs } // namespace mcs diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 08676161..0445863a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -89,6 +89,12 @@ target_include_directories(test_fmcs_match PRIVATE ${CMAKE_SOURCE_DIR}/src/mcs) target_link_libraries(test_fmcs_match PRIVATE mcs_fmcs_foundations device_vector) +add_executable(test_fmcs_substructure test_fmcs_substructure.cu) +target_include_directories(test_fmcs_substructure + PRIVATE ${CMAKE_SOURCE_DIR}/src/mcs) +target_link_libraries(test_fmcs_substructure PRIVATE mcs_fmcs_foundations + device_vector) + add_executable(test_openmp_helpers test_openmp_helpers.cpp) target_link_libraries(test_openmp_helpers PRIVATE openmp_helpers OpenMP::OpenMP_CXX) @@ -438,6 +444,7 @@ set(TEST_LIST test_fmcs_grow test_fmcs_match test_fmcs_policy + test_fmcs_substructure test_work_splitting test_fire_minimizer test_fire_minimizer_permol) diff --git a/tests/test_fmcs_substructure.cu b/tests/test_fmcs_substructure.cu new file mode 100644 index 00000000..d50c768d --- /dev/null +++ b/tests/test_fmcs_substructure.cu @@ -0,0 +1,599 @@ +// 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. +// +// Unit tests for the per-helper device functions in fmcs_cuda/. Each test +// launches a tiny __global__ driver that constructs inputs, calls one +// helper, and copies results back to host memory for assertion. This +// file is populated incrementally as Steps 1-5 land real implementations. + +#include +#include + +#include +#include +#include +#include + +#include "src/mcs/fmcs_cuda/fmcs_match.cuh" +#include "src/mcs/fmcs_cuda/fmcs_match_tables.cuh" +#include "src/mcs/fmcs_cuda/fmcs_seed.cuh" +#include "src/utils/device_vector.h" + +namespace { + +using nvMolKit::AsyncDevicePtr; +using nvMolKit::AsyncDeviceVector; + +using mcs::fmcs::MatchResult; +using mcs::fmcs::Seed; + +} // namespace + +// --------------------------------------------------------------------------- +// matchSingleBondWithinThread / tryMatchIncrementalGreedyCooperative +// --------------------------------------------------------------------------- + +namespace { + +using mcs::fmcs::MatchTableDevice; +using mcs::fmcs::PairMatchTablesDevice; +using mcs::fmcs::SingleBondMatch; + +using mcs::fmcs::DeviceCsrView; + +// Match tables staged on the host and uploaded to device memory on +// demand. Both atom and bond tables are 32-bit row-packed bitmasks +// (one bit per (q, t) pair). Tests mutate the host staging vectors +// via the set*Bit helpers; device() uploads (if dirty) and returns the +// device view to pass to kernels. +struct ManagedMatchTables { + std::vector atomHost; + std::vector bondHost; + AsyncDeviceVector atomData; + AsyncDeviceVector bondData; + int qNumAtoms = 0; + int qNumBonds = 0; + + void allocate(int qAtoms, int tAtoms, int qBonds, int tBonds) { + qNumAtoms = qAtoms; + qNumBonds = qBonds; + const int atomWordsPerRow = (tAtoms + 31) / 32; + const int bondWordsPerRow = (tBonds + 31) / 32; + atomHost.assign(qAtoms * atomWordsPerRow, 0); + bondHost.assign(qBonds * bondWordsPerRow, 0); + dev_.atoms = MatchTableDevice{nullptr, qAtoms, tAtoms, atomWordsPerRow}; + dev_.bonds = MatchTableDevice{nullptr, qBonds, tBonds, bondWordsPerRow}; + dirty_ = true; + } + + void setAtomBit(int qAtom, int tAtom) { + atomHost[qAtom * dev_.atoms.wordsPerRow + tAtom / 32] |= (1u << (tAtom % 32)); + dirty_ = true; + } + void setBondBit(int qBond, int tBond) { + bondHost[qBond * dev_.bonds.wordsPerRow + tBond / 32] |= (1u << (tBond % 32)); + dirty_ = true; + } + void setAllAtomBits() { + for (int q = 0; q < dev_.atoms.nRows; ++q) + for (int t = 0; t < dev_.atoms.nCols; ++t) + setAtomBit(q, t); + } + void setAllBondBits() { + for (int q = 0; q < dev_.bonds.nRows; ++q) + for (int t = 0; t < dev_.bonds.nCols; ++t) + setBondBit(q, t); + } + + PairMatchTablesDevice device() { + if (dirty_) { + atomData.setFromVector(atomHost); + bondData.setFromVector(bondHost); + dev_.atoms.data = atomData.data(); + dev_.bonds.data = bondData.data(); + dirty_ = false; + } + return dev_; + } + + private: + PairMatchTablesDevice dev_{}; + bool dirty_ = true; +}; + +// Owns the device buffers behind a DeviceCsrView. Built from an +// undirected edge list: bond i is edges[i], and the CSR is the symmetric +// expansion of that list (each edge contributes one entry to each +// endpoint's row), matching what the host-side graph builder uploads for +// real molecules. +class TestGraph { + public: + TestGraph(int numAtoms, const std::vector>& edges) { + const int numBonds = static_cast(edges.size()); + + std::vector bondEndpointsHost(numBonds); + for (int i = 0; i < numBonds; ++i) { + bondEndpointsHost[i] = + (static_cast(edges[i].first) << 16) | static_cast(edges[i].second); + } + + // Counting sort into CSR: degree pass, prefix sum, then scatter. + std::vector rowOffsetsHost(numAtoms + 1, 0); + for (const auto& edge : edges) { + ++rowOffsetsHost[edge.first + 1]; + ++rowOffsetsHost[edge.second + 1]; + } + for (int atom = 0; atom < numAtoms; ++atom) { + rowOffsetsHost[atom + 1] += rowOffsetsHost[atom]; + } + + std::vector colIndicesHost(2 * numBonds); + std::vector bondIndicesHost(2 * numBonds); + std::vector cursor(rowOffsetsHost.begin(), rowOffsetsHost.end() - 1); + for (int bond = 0; bond < numBonds; ++bond) { + const int u = edges[bond].first; + const int v = edges[bond].second; + colIndicesHost[cursor[u]] = static_cast(v); + bondIndicesHost[cursor[u]] = static_cast(bond); + ++cursor[u]; + colIndicesHost[cursor[v]] = static_cast(u); + bondIndicesHost[cursor[v]] = static_cast(bond); + ++cursor[v]; + } + + bondEndpoints_.setFromVector(bondEndpointsHost); + rowOffsets_.setFromVector(rowOffsetsHost); + colIndices_.setFromVector(colIndicesHost); + bondIndices_.setFromVector(bondIndicesHost); + + view_.bondEndpoints = bondEndpoints_.data(); + view_.rowOffsets = rowOffsets_.data(); + view_.colIndices = colIndices_.data(); + view_.bondIndices = bondIndices_.data(); + view_.numAtoms = numAtoms; + view_.numBonds = numBonds; + } + + DeviceCsrView view() const { return view_; } + + private: + AsyncDeviceVector bondEndpoints_; + AsyncDeviceVector rowOffsets_; + AsyncDeviceVector colIndices_; + AsyncDeviceVector bondIndices_; + DeviceCsrView view_{}; +}; + +} // namespace + +namespace { + +// Helper: build a parent MatchResult that records the mapping +// {qAtom[i] -> tAtom[i]} and {qBond[i] -> tBond[i]} from caller-supplied +// parallel arrays. Used by the incremental tests to set up "what the +// parent already had matched" before adding new bonds. +template +__device__ __forceinline__ void buildParentMatch(mcs::fmcs::MatchResult& match, + const int* qAtoms, + const int* tAtoms, + int nAtomMaps, + const int* qBonds, + const int* tBonds, + int nBondMaps) { + using MatchT = mcs::fmcs::MatchResult; + mcs::fmcs::matchResultClearWithinThread(match); + for (int i = 0; i < nAtomMaps; ++i) { + match.targetAtomIdx[qAtoms[i]] = static_cast(tAtoms[i]); + const int t = tAtoms[i]; + match.visitedTargetAtoms[t / MatchT::kTargetAtomBitsPerWord] |= (typename MatchT::target_atom_word{1}) + << (t % MatchT::kTargetAtomBitsPerWord); + } + for (int i = 0; i < nBondMaps; ++i) { + match.targetBondIdx[qBonds[i]] = static_cast(tBonds[i]); + const int t = tBonds[i]; + match.visitedTargetBonds[t / MatchT::kTargetBondBitsPerWord] |= (typename MatchT::target_bond_word{1}) + << (t % MatchT::kTargetBondBitsPerWord); + } + match.matchedAtomSize = static_cast(nAtomMaps); + match.matchedBondSize = static_cast(nBondMaps); + match.empty = (nAtomMaps == 0 && nBondMaps == 0); +} + +} // namespace +namespace mcs_fmcs_substructure_test { + +using QueuedT16 = mcs::fmcs::QueuedSeed<16, 16, 16, 16>; + +struct SubstructureTestOut { + bool ok; + bool overflowed; + QueuedT16 child; +}; + +constexpr int kTestSubstructurePartialCapacity = 64; + +__device__ __forceinline__ void addMaskSeed(QueuedT16& child, std::uint32_t atomMask, std::uint32_t bondMask) { + mcs::fmcs::seedClearWithinThread(child.seed); + mcs::fmcs::matchResultClearWithinThread(child.match); + for (int a = 0; a < 16; ++a) { + if ((atomMask >> a) & 1u) { + mcs::fmcs::seedAddAtomWithinThread(child.seed, a); + } + } + for (int b = 0; b < 16; ++b) { + if ((bondMask >> b) & 1u) { + mcs::fmcs::seedAddBondWithinThread(child.seed, b); + } + } +} + +__global__ void matchSubstructureMaskDriver(DeviceCsrView qView, + DeviceCsrView tView, + PairMatchTablesDevice tables, + std::uint32_t atomMask, + std::uint32_t bondMask, + std::uint8_t* partialStorage, + int partialCapacity, + SubstructureTestOut* out) { + __shared__ QueuedT16 child; + __shared__ mcs::fmcs::FmcsSubstructureScratch<16, 16> scratch; + if (threadIdx.x == 0) { + addMaskSeed(child, atomMask, bondMask); + } + __syncthreads(); + + auto block = cooperative_groups::this_thread_block(); + auto warp = cooperative_groups::tiled_partition<32>(block); + bool overflowed = false; + bool ok = mcs::fmcs::matchSeedSubstructureCooperative(warp, + child.seed, + qView, + tView, + tables, + child.match, + scratch, + partialStorage, + partialCapacity, + &overflowed); + __syncthreads(); + + if (threadIdx.x == 0) { + out->ok = ok; + out->overflowed = overflowed; + out->child = child; + } +} + +__global__ void matchFallbackBadParentDriver(DeviceCsrView qView, + DeviceCsrView tView, + PairMatchTablesDevice tables, + std::uint8_t* partialStorage, + int partialCapacity, + SubstructureTestOut* out) { + __shared__ QueuedT16 child; + __shared__ mcs::fmcs::FmcsSubstructureScratch<16, 16> scratch; + __shared__ int scratchLock; + if (threadIdx.x == 0) { + // Query seed is the 4-edge path inside the triangle-with-leaves + // repro: query bonds 1,2,3,4 and all five atoms. The stored parent + // match maps q bond 1 = (0,2) onto target bond 0 = (0,3), which is + // locally valid but blocks q bond 2 = (0,4) in the fast extender. + addMaskSeed(child, /*atomMask=*/0x1Fu, /*bondMask=*/0x1Eu); + mcs::fmcs::matchResultClearWithinThread(child.match); + using MatchT = decltype(child.match); + child.match.targetAtomIdx[0] = 0; + child.match.targetAtomIdx[2] = 3; + child.match.visitedTargetAtoms[0 / MatchT::kTargetAtomBitsPerWord] |= typename MatchT::target_atom_word{1} + << (0 % MatchT::kTargetAtomBitsPerWord); + child.match.visitedTargetAtoms[3 / MatchT::kTargetAtomBitsPerWord] |= typename MatchT::target_atom_word{1} + << (3 % MatchT::kTargetAtomBitsPerWord); + child.match.targetBondIdx[1] = 0; + child.match.visitedTargetBonds[0 / MatchT::kTargetBondBitsPerWord] |= typename MatchT::target_bond_word{1} + << (0 % MatchT::kTargetBondBitsPerWord); + child.match.matchedAtomSize = 2; + child.match.matchedBondSize = 1; + child.match.empty = false; + scratchLock = 0; + } + __syncthreads(); + + auto block = cooperative_groups::this_thread_block(); + auto warp = cooperative_groups::tiled_partition<32>(block); + bool overflowed = false; + bool ok = mcs::fmcs::matchSeedWithSubstructureFallbackCooperative(warp, + child.seed, + qView, + tView, + tables, + child.match, + scratch, + &scratchLock, + partialStorage, + partialCapacity, + &overflowed); + __syncthreads(); + + if (threadIdx.x == 0) { + out->ok = ok; + out->overflowed = overflowed; + out->child = child; + } +} + +} // namespace mcs_fmcs_substructure_test + +TEST(FMCSUnit, MatchSeedSubstructurePath) { + using namespace mcs_fmcs_substructure_test; + + TestGraph query(4, + { + {0, 1}, + {1, 2}, + {2, 3} + }); + TestGraph target(5, + { + {0, 1}, + {1, 2}, + {2, 3}, + {3, 4} + }); + ManagedMatchTables tables; + tables.allocate(4, 5, 3, 4); + tables.setAllAtomBits(); + tables.setAllBondBits(); + + AsyncDevicePtr d_out; + AsyncDeviceVector partials(2 * kTestSubstructurePartialCapacity * 16); + matchSubstructureMaskDriver<<<1, 32>>>(query.view(), + target.view(), + tables.device(), + /*atomMask=*/0xFu, + /*bondMask=*/0x7u, + partials.data(), + kTestSubstructurePartialCapacity, + d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + SubstructureTestOut out{}; + d_out.get(out); + + EXPECT_TRUE(out.ok); + EXPECT_FALSE(out.overflowed); + EXPECT_EQ(out.child.match.matchedAtomSize, 4); + EXPECT_EQ(out.child.match.matchedBondSize, 3); + for (int q = 0; q < 4; ++q) { + EXPECT_NE(out.child.match.targetAtomIdx[q], mcs::fmcs::kUnmappedTargetIdx); + } + for (int q = 0; q < 3; ++q) { + EXPECT_NE(out.child.match.targetBondIdx[q], mcs::fmcs::kUnmappedTargetIdx); + } +} + +TEST(FMCSUnit, MatchSeedSubstructureRejectsNoMatch) { + using namespace mcs_fmcs_substructure_test; + + TestGraph query(3, + { + {0, 1}, + {1, 2}, + {0, 2} + }); + TestGraph target(3, + { + {0, 1}, + {1, 2} + }); + ManagedMatchTables tables; + tables.allocate(3, 3, 3, 2); + tables.setAllAtomBits(); + tables.setAllBondBits(); + + AsyncDevicePtr d_out; + AsyncDeviceVector partials(2 * kTestSubstructurePartialCapacity * 16); + matchSubstructureMaskDriver<<<1, 32>>>(query.view(), + target.view(), + tables.device(), + /*atomMask=*/0x7u, + /*bondMask=*/0x7u, + partials.data(), + kTestSubstructurePartialCapacity, + d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + SubstructureTestOut out{}; + d_out.get(out); + + EXPECT_FALSE(out.ok); + EXPECT_FALSE(out.overflowed); + EXPECT_TRUE(out.child.match.empty); +} + +TEST(FMCSUnit, MatchSeedSubstructureRespectsAtomTable) { + using namespace mcs_fmcs_substructure_test; + + TestGraph query(3, + { + {0, 1}, + {1, 2} + }); + TestGraph target(3, + { + {0, 1}, + {1, 2} + }); + ManagedMatchTables tables; + tables.allocate(3, 3, 2, 2); + tables.setAtomBit(0, 0); + tables.setAtomBit(1, 1); + tables.setAtomBit(2, 2); + tables.setAllBondBits(); + + AsyncDevicePtr d_out; + AsyncDeviceVector partials(2 * kTestSubstructurePartialCapacity * 16); + matchSubstructureMaskDriver<<<1, 32>>>(query.view(), + target.view(), + tables.device(), + /*atomMask=*/0x7u, + /*bondMask=*/0x3u, + partials.data(), + kTestSubstructurePartialCapacity, + d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + SubstructureTestOut out{}; + d_out.get(out); + + EXPECT_TRUE(out.ok); + EXPECT_FALSE(out.overflowed); + EXPECT_EQ(out.child.match.targetAtomIdx[0], 0u); + EXPECT_EQ(out.child.match.targetAtomIdx[1], 1u); + EXPECT_EQ(out.child.match.targetAtomIdx[2], 2u); + EXPECT_EQ(out.child.match.matchedAtomSize, 3); + EXPECT_EQ(out.child.match.matchedBondSize, 2); +} + +TEST(FMCSUnit, MatchSeedSubstructureRespectsBondTable) { + using namespace mcs_fmcs_substructure_test; + + TestGraph query(3, + { + {0, 1}, + {1, 2} + }); + TestGraph target(3, + { + {0, 1}, + {1, 2}, + {0, 2} + }); + ManagedMatchTables tables; + tables.allocate(3, 3, 2, 3); + tables.setAllAtomBits(); + tables.setBondBit(0, 0); + tables.setBondBit(1, 1); + + AsyncDevicePtr d_out; + AsyncDeviceVector partials(2 * kTestSubstructurePartialCapacity * 16); + matchSubstructureMaskDriver<<<1, 32>>>(query.view(), + target.view(), + tables.device(), + /*atomMask=*/0x7u, + /*bondMask=*/0x3u, + partials.data(), + kTestSubstructurePartialCapacity, + d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + SubstructureTestOut out{}; + d_out.get(out); + + EXPECT_TRUE(out.ok); + EXPECT_FALSE(out.overflowed); + EXPECT_EQ(out.child.match.targetBondIdx[0], 0u); + EXPECT_EQ(out.child.match.targetBondIdx[1], 1u); + EXPECT_EQ(out.child.match.matchedAtomSize, 3); + EXPECT_EQ(out.child.match.matchedBondSize, 2); +} + +TEST(FMCSUnit, MatchSeedSubstructureFindsPathInsideTriangleWithLeaves) { + using namespace mcs_fmcs_substructure_test; + + TestGraph query(5, + { + {0, 1}, + {0, 2}, + {0, 4}, + {1, 2}, + {1, 3} + }); + TestGraph target(5, + { + {0, 3}, + {1, 2}, + {1, 4}, + {2, 3} + }); + ManagedMatchTables tables; + tables.allocate(5, 5, 5, 4); + tables.setAllAtomBits(); + tables.setAllBondBits(); + + AsyncDevicePtr d_out; + AsyncDeviceVector partials(2 * kTestSubstructurePartialCapacity * 16); + matchSubstructureMaskDriver<<<1, 32>>>(query.view(), + target.view(), + tables.device(), + /*atomMask=*/0x1Fu, + /*bondMask=*/0x1Eu, + partials.data(), + kTestSubstructurePartialCapacity, + d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + SubstructureTestOut out{}; + d_out.get(out); + + EXPECT_TRUE(out.ok); + EXPECT_FALSE(out.overflowed); + EXPECT_EQ(out.child.match.matchedAtomSize, 5); + EXPECT_EQ(out.child.match.matchedBondSize, 4); + EXPECT_EQ(out.child.match.targetBondIdx[0], mcs::fmcs::kUnmappedTargetIdx); + for (int q : {1, 2, 3, 4}) { + EXPECT_NE(out.child.match.targetBondIdx[q], mcs::fmcs::kUnmappedTargetIdx); + } +} + +TEST(FMCSUnit, MatchSeedFallbackRebuildsAfterGreedyFailure) { + using namespace mcs_fmcs_substructure_test; + + TestGraph query(5, + { + {0, 1}, + {0, 2}, + {0, 4}, + {1, 2}, + {1, 3} + }); + TestGraph target(5, + { + {0, 3}, + {1, 2}, + {1, 4}, + {2, 3} + }); + ManagedMatchTables tables; + tables.allocate(5, 5, 5, 4); + tables.setAllAtomBits(); + tables.setAllBondBits(); + + AsyncDevicePtr d_out; + AsyncDeviceVector partials(2 * kTestSubstructurePartialCapacity * 16); + matchFallbackBadParentDriver<<<1, 32>>>(query.view(), + target.view(), + tables.device(), + partials.data(), + kTestSubstructurePartialCapacity, + d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + SubstructureTestOut out{}; + d_out.get(out); + + EXPECT_TRUE(out.ok); + EXPECT_FALSE(out.overflowed); + EXPECT_EQ(out.child.match.matchedAtomSize, 5); + EXPECT_EQ(out.child.match.matchedBondSize, 4); + for (int q = 0; q < 5; ++q) { + EXPECT_NE(out.child.match.targetAtomIdx[q], mcs::fmcs::kUnmappedTargetIdx); + } + for (int q : {1, 2, 3, 4}) { + EXPECT_NE(out.child.match.targetBondIdx[q], mcs::fmcs::kUnmappedTargetIdx); + } +}