diff --git a/src/mcs/fmcs_cuda/fmcs_grow.cuh b/src/mcs/fmcs_cuda/fmcs_grow.cuh index 14b6b418..e72d44d7 100644 --- a/src/mcs/fmcs_cuda/fmcs_grow.cuh +++ b/src/mcs/fmcs_cuda/fmcs_grow.cuh @@ -19,14 +19,12 @@ #include #include "src/mcs/fmcs_cuda/fmcs_seed.cuh" +#include "src/mcs/fmcs_cuda/fmcs_topology.cuh" #include "src/mcs/mcs_common/mcs_cooperative_copy.cuh" namespace mcs { namespace fmcs { -constexpr int kGrowBondEndpointShift = 16; -constexpr std::uint32_t kGrowBondEndpointMask = 0xFFFFu; - template __device__ __forceinline__ void seedAddNewBondWithinThread(Seed& seed, const NewBond& bond) { seedAddBondWithinThread(seed, bond.bondIdx); @@ -68,10 +66,10 @@ __device__ __forceinline__ void seedAddNewBondWithinThread(Seed +template __device__ __forceinline__ bool fillNewBondsCooperative(const GroupT& group, const Seed& seed, - const QueryTopology& queryTopology, + const DeviceCsrView& queryTopology, NewBond* outBonds, int* outCount, int maxNewBonds) { @@ -96,8 +94,8 @@ __device__ __forceinline__ bool fillNewBondsCooperative(const GroupT& // Decode this query bond's endpoints. const std::uint32_t queryEndpoints = queryTopology.bondEndpoints[q]; - const int queryEndpointU = static_cast(queryEndpoints >> kGrowBondEndpointShift); - const int queryEndpointV = static_cast(queryEndpoints & kGrowBondEndpointMask); + const int queryEndpointU = static_cast(queryEndpoints >> kBondEndpointShift); + const int queryEndpointV = static_cast(queryEndpoints & kBondEndpointMask); // Bond is a "new boundary" candidate iff at least one endpoint // was added in the most recent grow step. @@ -153,10 +151,12 @@ __device__ __forceinline__ bool fillNewBondsCooperative(const GroupT& /// across all lanes of @p group), and within each iteration the only /// truly group-parallel step is the @c warpCopy of @c parent into /// @p childWorkspace. Lane 0 then patches in the singleton bond / -/// atom; the per-bond match call is whatever @p matchFn does (when -/// the kernel passes @ref matchIncrementalFastCooperative the inner -/// target-bond scan IS lane-parallel); @p childSink is typically lane -/// 0 only (cache.insert + queue.push). +/// atom; the per-bond match call is whatever @p matchFn does. The +/// kernel's callback may try @ref tryMatchIncrementalGreedyCooperative, +/// but it must own the exact fallback too: a greedy miss alone must not +/// prune the bond. The greedy helper's inner target-bond scan is +/// lane-parallel; @p childSink is typically lane 0 only (cache.insert + +/// queue.push). /// /// In other words: this function is cooperative only for the /// parent->child copy and for delegating to its callbacks; it does diff --git a/src/mcs/fmcs_cuda/fmcs_match.cuh b/src/mcs/fmcs_cuda/fmcs_match.cuh new file mode 100644 index 00000000..b800ca40 --- /dev/null +++ b/src/mcs/fmcs_cuda/fmcs_match.cuh @@ -0,0 +1,324 @@ +// 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. + +#ifndef FMCS_CUDA_FMCS_MATCH_CUH +#define FMCS_CUDA_FMCS_MATCH_CUH + +#include + +#include "src/mcs/fmcs_cuda/fmcs_match_tables.cuh" +#include "src/mcs/fmcs_cuda/fmcs_seed.cuh" +#include "src/mcs/fmcs_cuda/fmcs_topology.cuh" + +namespace mcs { +namespace fmcs { + +/// Resolved target endpoints for a successful single-(query bond, target +/// bond, orientation) compatibility check. Populated by +/// @ref matchSingleBondWithinThread on success only; contents are +/// unspecified on failure. @c targetAtomU is the target atom that the +/// query bond's u endpoint was mapped to (likewise V). +struct SingleBondMatch { + uint8_t targetAtomU; + uint8_t targetAtomV; +}; + +template struct FmcsSubstructureScratch { + // Scratch for the RDKit checkIfMatchAndAppend fallback. + std::uint8_t seedAtomList[maxAtoms]; + std::uint8_t seedAtoms[maxAtoms]; + std::uint8_t seedDegree[maxAtoms]; + std::uint8_t targetDegree[maxTargetAtoms]; + std::uint8_t orderedQueryAtom[maxAtoms]; + std::uint8_t queryOrderPos[maxAtoms]; + std::uint8_t targetAtomForQuery[maxAtoms]; + int currentCount; + int nextCount; + int found; + int overflowed; +}; + +template +__device__ __forceinline__ bool seedContainsBondWithinThread(const Seed& seed, + const int queryBondIdx) { + using SeedT = Seed; + using BondWord = typename SeedT::bond_word_type; + constexpr int kBondBitsPerWord = SeedT::kBondBitsPerWord; + if (queryBondIdx < 0 || queryBondIdx >= maxBonds) + return false; + const BondWord word = seed.bonds[queryBondIdx / kBondBitsPerWord]; + return ((word >> (queryBondIdx % kBondBitsPerWord)) & 1) != 0; +} + +/// Within-thread: per-lane single-(query bond, target bond, orientation) +/// compatibility check used by Phase 1 initial-seed enumeration. Writes +/// resolved target atom indices for the two endpoints of @p queryBondIdx +/// into @p outMatch and returns true on success; on false the caller +/// should not read @p outMatch. +/// +/// @p reversed selects the orientation: when false, the query bond's u +/// endpoint maps to the target bond's u endpoint; when true, to the +/// target bond's v endpoint. +__device__ __forceinline__ bool matchSingleBondWithinThread(const int queryBondIdx, + const int targetBondIdx, + const bool reversed, + const DeviceCsrView& queryTopology, + const DeviceCsrView& targetTopology, + const PairMatchTablesDevice& tables, + SingleBondMatch& outMatch) { + // Cheap bond-table check first; if the bond labels are incompatible + // we never need to touch the atom table or decode endpoints. + if (!tables.bonds.testBit(queryBondIdx, targetBondIdx)) + return false; + + // Decode the packed (u<<16 | v) endpoint encoding for both bonds. + const std::uint32_t queryEndpoints = queryTopology.bondEndpoints[queryBondIdx]; + const int queryEndpointU = static_cast(queryEndpoints >> kBondEndpointShift); + const int queryEndpointV = static_cast(queryEndpoints & kBondEndpointMask); + + const std::uint32_t targetEndpoints = targetTopology.bondEndpoints[targetBondIdx]; + const int targetEndpointU = static_cast(targetEndpoints >> kBondEndpointShift); + const int targetEndpointV = static_cast(targetEndpoints & kBondEndpointMask); + + // Pick which target endpoint the query's u maps to; v gets the other. + // reversed=false: queryU -> targetU, queryV -> targetV. + // reversed=true: queryU -> targetV, queryV -> targetU. + const int targetForQueryU = reversed ? targetEndpointV : targetEndpointU; + const int targetForQueryV = reversed ? targetEndpointU : targetEndpointV; + + // Both atom-pairings must be label-compatible; either failure rejects + // this orientation. + if (!tables.atoms.testBit(queryEndpointU, targetForQueryU)) + return false; + if (!tables.atoms.testBit(queryEndpointV, targetForQueryV)) + return false; + + outMatch.targetAtomU = static_cast(targetForQueryU); + outMatch.targetAtomV = static_cast(targetForQueryV); + return true; +} + +/// Cooperative greedy fast path: try to extend @p match by every query +/// bond in @p seed.bonds whose @c match.targetBondIdx[q] is still +/// @ref kUnmappedTargetIdx (i.e., unmapped by the parent's recorded +/// embedding). For each such bond: +/// - Both endpoints already mapped -> ring-closing case. The lanes +/// of @p group scan target bonds in parallel for one whose endpoint +/// pair exactly matches the mapped (queryU, queryV) target atoms +/// and is unvisited, with the bond-match-table bit set. First +/// compatible target bond commits. +/// - Exactly one endpoint mapped -> atom-adding case. Lane-parallel +/// scan for a target bond incident to the mapped target atom whose +/// other end is unvisited, atom-table-compatible with the unmapped +/// query atom, and bond-table-compatible. The first compatible +/// candidate commits both the new bond mapping and the new atom +/// mapping, and marks both visited. This path does not backtrack +/// if that greedy choice prevents a later bond from matching. +/// - Both endpoints unmapped -> defensive fail (shouldn't occur on +/// well-formed seeds, where Phase 1 maps both initial atoms before +/// pushing). +/// A true return proves that @p match was extended successfully. A false +/// return is inconclusive: the greedy choices may have missed another valid +/// embedding. The caller must run the exact substructure fallback before +/// rejecting the seed or marking a NewBond dead. On false, @p match is left +/// in an unspecified state and must not be reused as a valid embedding. +template +__device__ __forceinline__ bool tryMatchIncrementalGreedyCooperative( + const GroupT& group, + const Seed& seed, + const DeviceCsrView& queryTopology, + const DeviceCsrView& targetTopology, + const PairMatchTablesDevice& tables, + MatchResult& match) { + using SeedT = Seed; + using MatchT = MatchResult; + using BondWord = typename SeedT::bond_word_type; + + constexpr int kBondBitsPerWord = SeedT::kBondBitsPerWord; + constexpr int kBondWords = SeedT::kBondWords; + constexpr int kTargetAtomBitsPerWord = MatchT::kTargetAtomBitsPerWord; + constexpr int kTargetBondBitsPerWord = MatchT::kTargetBondBitsPerWord; + using TargetAtomWord = typename MatchT::target_atom_word; + using TargetBondWord = typename MatchT::target_bond_word; + + const int laneRank = static_cast(group.thread_rank()); + const int laneCount = static_cast(group.num_threads()); + + // Outer loop walks set bits of seed.bonds via __ffs/__ffsll. Each + // iteration handles one query bond q; if q is already mapped (parent + // saw it) we skip, otherwise we extend the match by one bond + + // possibly one atom. + for (int wordIdx = 0; wordIdx < kBondWords; ++wordIdx) { + BondWord remainingBondBits = seed.bonds[wordIdx]; + while (remainingBondBits != 0) { + int bitPosInWord; + if constexpr (sizeof(BondWord) == 4) { + bitPosInWord = __ffs(static_cast(remainingBondBits)) - 1; + } else { + bitPosInWord = __ffsll(static_cast(remainingBondBits)) - 1; + } + const int queryBondIdx = wordIdx * kBondBitsPerWord + bitPosInWord; + remainingBondBits &= remainingBondBits - 1; // clear lowest set bit + + // Keep this control decision uniform across the group; diverging before + // the later ballot/shuffle would make the winning lane undefined. + int mappedTargetBond = kUnmappedTargetIdx; + if (laneRank == 0) + mappedTargetBond = match.targetBondIdx[queryBondIdx]; + mappedTargetBond = group.shfl(mappedTargetBond, 0); + if (mappedTargetBond != kUnmappedTargetIdx) + continue; + + // Decode this bond's query endpoints from the packed (u<<16 | v). + const std::uint32_t queryEndpoints = queryTopology.bondEndpoints[queryBondIdx]; + const int queryEndpointU = static_cast(queryEndpoints >> kBondEndpointShift); + const int queryEndpointV = static_cast(queryEndpoints & kBondEndpointMask); + + // Look up the parent's atom mapping for both endpoints; either + // may already be mapped (from an earlier bond) or still unmapped + // (this bond is the one bringing it in). + int targetForQueryU = kUnmappedTargetIdx; + int targetForQueryV = kUnmappedTargetIdx; + if (laneRank == 0) { + targetForQueryU = match.targetAtomIdx[queryEndpointU]; + targetForQueryV = match.targetAtomIdx[queryEndpointV]; + } + targetForQueryU = group.shfl(targetForQueryU, 0); + targetForQueryV = group.shfl(targetForQueryV, 0); + const bool queryUIsMapped = targetForQueryU != kUnmappedTargetIdx; + const bool queryVIsMapped = targetForQueryV != kUnmappedTargetIdx; + + // Both endpoints unmapped means the seed is missing an earlier + // bond that should have brought one of them in. Phase 1 always + // anchors both initial atoms before pushing, so this never hits + // for well-formed seeds; defensive return false. + if (!queryUIsMapped && !queryVIsMapped) + return false; + + // Each lane records its first compatible target bond, and (for + // the atom-adding case) the resulting new target atom. After + // the per-lane scan, the warp ballots and the lowest-rank winner + // is broadcast as the committed extension. + int chosenTargetBond = -1; + int chosenTargetAtomForUnmapped = -1; // -1 means ring-closing. + + if (queryUIsMapped && queryVIsMapped) { + // Ring-closing: the new bond connects two atoms that are both + // already in the seed's mapping. Find a target bond whose + // endpoints are exactly the pair { targetForQueryU, targetForQueryV }. + const int srcTargetAtom = targetForQueryU; + const int dstTargetAtom = targetForQueryV; + if (srcTargetAtom >= 0 && srcTargetAtom < targetTopology.numAtoms) { + const int begin = static_cast(targetTopology.rowOffsets[srcTargetAtom]); + const int end = static_cast(targetTopology.rowOffsets[srcTargetAtom + 1]); + for (int adjIdx = begin + laneRank; adjIdx < end && chosenTargetBond < 0; adjIdx += laneCount) { + const int otherTargetAtom = static_cast(targetTopology.colIndices[adjIdx]); + if (otherTargetAtom != dstTargetAtom) + continue; + const int targetBondIdx = static_cast(targetTopology.bondIndices[adjIdx]); + if (targetBondIdx < 0 || targetBondIdx >= targetTopology.numBonds || targetBondIdx >= maxTB) { + continue; + } + const TargetBondWord visitedBondsWord = match.visitedTargetBonds[targetBondIdx / kTargetBondBitsPerWord]; + if ((visitedBondsWord >> (targetBondIdx % kTargetBondBitsPerWord)) & 1) { + continue; + } + if (!tables.bonds.testBit(queryBondIdx, targetBondIdx)) + continue; + chosenTargetBond = targetBondIdx; + } + } + } else { + // Atom-adding: one query endpoint (`src`) is already mapped to + // a target atom; the other (`dst`) is what we're trying to + // place. Scan target bonds incident to srcTargetAtom for one + // whose other endpoint is unvisited, atom-table-compatible + // with the unmapped query atom, and bond-table-compatible. + const int unmappedQueryAtom = queryUIsMapped ? queryEndpointV : queryEndpointU; + const int srcTargetAtom = queryUIsMapped ? targetForQueryU : targetForQueryV; + if (srcTargetAtom >= 0 && srcTargetAtom < targetTopology.numAtoms) { + const int begin = static_cast(targetTopology.rowOffsets[srcTargetAtom]); + const int end = static_cast(targetTopology.rowOffsets[srcTargetAtom + 1]); + for (int adjIdx = begin + laneRank; adjIdx < end && chosenTargetBond < 0; adjIdx += laneCount) { + const int targetBondIdx = static_cast(targetTopology.bondIndices[adjIdx]); + if (targetBondIdx < 0 || targetBondIdx >= targetTopology.numBonds || targetBondIdx >= maxTB) { + continue; + } + const TargetBondWord visitedBondsWord = match.visitedTargetBonds[targetBondIdx / kTargetBondBitsPerWord]; + if ((visitedBondsWord >> (targetBondIdx % kTargetBondBitsPerWord)) & 1) { + continue; + } + const int candidateTargetAtom = static_cast(targetTopology.colIndices[adjIdx]); + if (candidateTargetAtom < 0 || candidateTargetAtom >= targetTopology.numAtoms || + candidateTargetAtom >= maxTA) { + continue; + } + const TargetAtomWord visitedAtomsWord = + match.visitedTargetAtoms[candidateTargetAtom / kTargetAtomBitsPerWord]; + if ((visitedAtomsWord >> (candidateTargetAtom % kTargetAtomBitsPerWord)) & 1) { + continue; + } + if (!tables.bonds.testBit(queryBondIdx, targetBondIdx)) + continue; + if (!tables.atoms.testBit(unmappedQueryAtom, candidateTargetAtom)) + continue; + chosenTargetBond = targetBondIdx; + chosenTargetAtomForUnmapped = candidateTargetAtom; + } + } + } + + // Warp-wide pick: ballot for any lane with a candidate, broadcast + // the lowest-rank winner's choice to all lanes. No candidate + // anywhere -> this bond can't be extended -> abandon the seed. + const unsigned ballot = group.ballot(chosenTargetBond >= 0 ? 1u : 0u); + if (ballot == 0u) + return false; + const int firstWinningLane = __ffs(ballot) - 1; + const int committedTargetBond = group.shfl(chosenTargetBond, firstWinningLane); + const int committedTargetAtom = group.shfl(chosenTargetAtomForUnmapped, firstWinningLane); + if (committedTargetBond < 0 || committedTargetBond >= targetTopology.numBonds || committedTargetBond >= maxTB) + return false; + if (committedTargetAtom < -1 || committedTargetAtom >= targetTopology.numAtoms || committedTargetAtom >= maxTA) + return false; + + // Lane 0 commits the new mapping into the shared MatchResult. + // Subsequent bonds in this same call will read the updated + // visited bitsets after group.sync below. + if (laneRank == 0) { + match.targetBondIdx[queryBondIdx] = static_cast(committedTargetBond); + match.visitedTargetBonds[committedTargetBond / kTargetBondBitsPerWord] |= + static_cast(1) << (committedTargetBond % kTargetBondBitsPerWord); + match.matchedBondSize += 1; + match.empty = false; + if (committedTargetAtom >= 0) { + // Atom-adding case: also commit the new atom mapping. + const int unmappedQueryAtom = queryUIsMapped ? queryEndpointV : queryEndpointU; + match.targetAtomIdx[unmappedQueryAtom] = static_cast(committedTargetAtom); + match.visitedTargetAtoms[committedTargetAtom / kTargetAtomBitsPerWord] |= + static_cast(1) << (committedTargetAtom % kTargetAtomBitsPerWord); + match.matchedAtomSize += 1; + } + } + group.sync(); + } + } + return true; +} + +} // namespace fmcs +} // namespace mcs + +#endif // FMCS_CUDA_FMCS_MATCH_CUH diff --git a/src/mcs/fmcs_cuda/fmcs_topology.cuh b/src/mcs/fmcs_cuda/fmcs_topology.cuh new file mode 100644 index 00000000..dd26e14b --- /dev/null +++ b/src/mcs/fmcs_cuda/fmcs_topology.cuh @@ -0,0 +1,52 @@ +// 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. + +#ifndef FMCS_CUDA_FMCS_TOPOLOGY_CUH +#define FMCS_CUDA_FMCS_TOPOLOGY_CUH + +#include + +namespace mcs { +namespace fmcs { + +/// Bond endpoints are stored across the kernel as a single uint32 with +/// the u-endpoint atom index in the high 16 bits and the v-endpoint +/// atom index in the low 16 bits. Tiers cap maxAtoms at 128 so 16 bits +/// per index is plenty. +constexpr int kBondEndpointShift = 16; +constexpr std::uint32_t kBondEndpointMask = 0xFFFFu; + +/// Non-owning view over one side's CSR + bond-endpoint arrays. Passed to +/// the matcher and grow helpers as the query- or target-side topology. +/// +/// All four arrays are required; helpers index them unconditionally. +/// @c rowOffsets has @c numAtoms + 1 entries; @c colIndices and +/// @c bondIndices are parallel and hold one entry per directed CSR edge +/// (the neighbour atom and the undirected bond id respectively). +/// @c bondEndpoints has one packed (u << 16 | v) entry per undirected +/// bond, ordered to match the bond dimension of the match tables. +struct DeviceCsrView { + const std::uint32_t* rowOffsets = nullptr; + const std::uint32_t* colIndices = nullptr; + const std::uint32_t* bondIndices = nullptr; + const std::uint32_t* bondEndpoints = nullptr; + int numAtoms = 0; + int numBonds = 0; +}; + +} // namespace fmcs +} // namespace mcs + +#endif // FMCS_CUDA_FMCS_TOPOLOGY_CUH diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index f4480ba1..f4647a0c 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -81,6 +81,11 @@ target_compile_options(test_fmcs_grow add_executable(test_fmcs_foundations test_fmcs_foundations.cu) target_link_libraries(test_fmcs_foundations PRIVATE mcs_fmcs_foundations) +add_executable(test_fmcs_match test_fmcs_match.cu) +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_openmp_helpers test_openmp_helpers.cpp) target_link_libraries(test_openmp_helpers PRIVATE openmp_helpers OpenMP::OpenMP_CXX) @@ -428,6 +433,7 @@ set(TEST_LIST test_fmcs_primitives test_fmcs_foundations test_fmcs_grow + test_fmcs_match test_work_splitting test_fire_minimizer test_fire_minimizer_permol) diff --git a/tests/test_fmcs_grow.cu b/tests/test_fmcs_grow.cu index 6f23af04..0f9962fe 100644 --- a/tests/test_fmcs_grow.cu +++ b/tests/test_fmcs_grow.cu @@ -18,12 +18,6 @@ namespace { using nvMolKit::AsyncDevicePtr; using nvMolKit::AsyncDeviceVector; -struct TestCsrView { - const std::uint32_t* bondEndpoints = nullptr; - int numAtoms = 0; - int numBonds = 0; -}; - AsyncDeviceVector makeBondEndpointsDevice(const std::vector>& edges) { std::vector host(edges.size()); for (std::size_t i = 0; i < edges.size(); ++i) { @@ -74,10 +68,15 @@ __device__ __forceinline__ void fillNewBondsRun(SeedSetup&& setup, } __syncthreads(); - auto block = cooperative_groups::this_thread_block(); - auto warp = cooperative_groups::tiled_partition<32>(block); - TestCsrView qView{qBondEndpoints, qNumAtoms, qNumBonds}; - bool ok = mcs::fmcs::fillNewBondsCooperative(warp, seed, qView, bonds, &count, maxNewBonds); + auto block = cooperative_groups::this_thread_block(); + auto warp = cooperative_groups::tiled_partition<32>(block); + // fillNewBondsCooperative reads only the bond-endpoint side of the + // topology, so the CSR arrays stay null here. + mcs::fmcs::DeviceCsrView qView{}; + qView.bondEndpoints = qBondEndpoints; + qView.numAtoms = qNumAtoms; + qView.numBonds = qNumBonds; + bool ok = mcs::fmcs::fillNewBondsCooperative(warp, seed, qView, bonds, &count, maxNewBonds); __syncthreads(); if (threadIdx.x == 0) { diff --git a/tests/test_fmcs_match.cu b/tests/test_fmcs_match.cu new file mode 100644 index 00000000..004da3f0 --- /dev/null +++ b/tests/test_fmcs_match.cu @@ -0,0 +1,657 @@ +// 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 + +#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_{}; +}; + +// ---- matchSingleBondWithinThread ---- + +struct SingleBondTestOut { + bool ok; + SingleBondMatch match; +}; + +__global__ void matchSingleBondDriver(int qBondIdx, + int tBondIdx, + bool reversed, + DeviceCsrView qView, + DeviceCsrView tView, + PairMatchTablesDevice tables, + SingleBondTestOut* out) { + if (threadIdx.x != 0 || blockIdx.x != 0) + return; + SingleBondMatch sm{}; + out->ok = mcs::fmcs::matchSingleBondWithinThread(qBondIdx, tBondIdx, reversed, qView, tView, tables, sm); + out->match = sm; +} + +} // namespace + +TEST(FMCSUnit, MatchSingleBondForwardOrientation) { + // Query bond (0,1), target bond (0,1). All atoms / bonds compatible. + TestGraph query(2, + { + {0, 1} + }); + TestGraph target(2, + { + {0, 1} + }); + ManagedMatchTables tables; + tables.allocate(2, 2, 1, 1); + tables.setAllAtomBits(); + tables.setAllBondBits(); + + AsyncDevicePtr d_out; + matchSingleBondDriver<<<1, 1>>>(0, + 0, + /*reversed=*/false, + query.view(), + target.view(), + tables.device(), + d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + SingleBondTestOut out{}; + d_out.get(out); + + EXPECT_TRUE(out.ok); + EXPECT_EQ(out.match.targetAtomU, 0u); // qU=0 -> tU=0 + EXPECT_EQ(out.match.targetAtomV, 1u); // qV=1 -> tV=1 +} + +TEST(FMCSUnit, MatchSingleBondReverseOrientation) { + TestGraph query(2, + { + {0, 1} + }); + TestGraph target(2, + { + {0, 1} + }); + ManagedMatchTables tables; + tables.allocate(2, 2, 1, 1); + tables.setAllAtomBits(); + tables.setAllBondBits(); + + AsyncDevicePtr d_out; + matchSingleBondDriver<<<1, 1>>>(0, + 0, + /*reversed=*/true, + query.view(), + target.view(), + tables.device(), + d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + SingleBondTestOut out{}; + d_out.get(out); + + EXPECT_TRUE(out.ok); + EXPECT_EQ(out.match.targetAtomU, 1u); // qU=0 -> tV=1 (reversed) + EXPECT_EQ(out.match.targetAtomV, 0u); // qV=1 -> tU=0 +} + +TEST(FMCSUnit, MatchSingleBondBondTableRejection) { + TestGraph query(2, + { + {0, 1} + }); + TestGraph target(2, + { + {0, 1} + }); + ManagedMatchTables tables; + tables.allocate(2, 2, 1, 1); + tables.setAllAtomBits(); + // Deliberately leave bondData all zero -> bond-table reject. + + AsyncDevicePtr d_out; + matchSingleBondDriver<<<1, 1>>>(0, + 0, + /*reversed=*/false, + query.view(), + target.view(), + tables.device(), + d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + SingleBondTestOut out{}; + d_out.get(out); + + EXPECT_FALSE(out.ok); +} + +TEST(FMCSUnit, MatchSingleBondAtomTableRejection) { + TestGraph query(2, + { + {0, 1} + }); + TestGraph target(2, + { + {0, 1} + }); + ManagedMatchTables tables; + tables.allocate(2, 2, 1, 1); + tables.setAllBondBits(); + // Atom 0 compatible with target atom 0, but atom 1 is incompatible + // with target atom 1 -> forward orientation rejects on second atom. + tables.setAtomBit(0, 0); + tables.setAtomBit(0, 1); + tables.setAtomBit(1, 0); + // Note: (1, 1) deliberately left unset. + + AsyncDevicePtr d_out; + matchSingleBondDriver<<<1, 1>>>(0, + 0, + /*reversed=*/false, + query.view(), + target.view(), + tables.device(), + d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + SingleBondTestOut out{}; + d_out.get(out); + + EXPECT_FALSE(out.ok); +} + +// ---- tryMatchIncrementalGreedyCooperative ---- + +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_incremental_test { + +using QueuedT16 = mcs::fmcs::QueuedSeed<16, 16, 16, 16>; + +struct IncrementalTestOut { + bool ok; + QueuedT16 child; +}; + +// One-warp driver: builds parent match in shared mem, then constructs +// the child seed (parent + new bonds) and runs +// tryMatchIncrementalGreedyCooperative on it. +__global__ void matchIncrementalAtomAddingDriver(DeviceCsrView qView, + DeviceCsrView tView, + PairMatchTablesDevice tables, + IncrementalTestOut* out) { + __shared__ QueuedT16 child; + if (threadIdx.x == 0) { + mcs::fmcs::seedClearWithinThread(child.seed); + // Parent: bond (0,1) mapped 0->0, 1->1, bond 0->target bond 0. + int qA[] = {0, 1}; + int tA[] = {0, 1}; + int qB[] = {0}; + int tB[] = {0}; + buildParentMatch(child.match, qA, tA, 2, qB, tB, 1); + // Child seed has bonds {0, 1} and atoms {0, 1, 2}. Bond 1 is the + // new atom-adding bond (qU=1 already mapped, qV=2 unmapped). + mcs::fmcs::seedAddBondWithinThread(child.seed, 0); + mcs::fmcs::seedAddBondWithinThread(child.seed, 1); + mcs::fmcs::seedAddAtomWithinThread(child.seed, 0); + mcs::fmcs::seedAddAtomWithinThread(child.seed, 1); + mcs::fmcs::seedBeginGrowStepWithinThread(child.seed); + mcs::fmcs::seedAddAtomWithinThread(child.seed, 2); + } + __syncthreads(); + + auto block = cooperative_groups::this_thread_block(); + auto warp = cooperative_groups::tiled_partition<32>(block); + bool ok = mcs::fmcs::tryMatchIncrementalGreedyCooperative(warp, child.seed, qView, tView, tables, child.match); + __syncthreads(); + + if (threadIdx.x == 0) { + out->ok = ok; + out->child = child; + } +} + +__global__ void matchIncrementalRingClosingDriver(DeviceCsrView qView, + DeviceCsrView tView, + PairMatchTablesDevice tables, + IncrementalTestOut* out) { + __shared__ QueuedT16 child; + if (threadIdx.x == 0) { + mcs::fmcs::seedClearWithinThread(child.seed); + // 4-atom square query: bonds (0,1), (1,2), (2,3), (0,3). Parent + // has all 4 atoms mapped via 3 path bonds; child adds the + // ring-closing 4th bond (0,3). + int qA[] = {0, 1, 2, 3}; + int tA[] = {0, 1, 2, 3}; + int qB[] = {0, 1, 2}; + int tB[] = {0, 1, 2}; + buildParentMatch(child.match, qA, tA, 4, qB, tB, 3); + for (int b : {0, 1, 2, 3}) + mcs::fmcs::seedAddBondWithinThread(child.seed, b); + for (int a : {0, 1, 2, 3}) + mcs::fmcs::seedAddAtomWithinThread(child.seed, a); + } + __syncthreads(); + + auto block = cooperative_groups::this_thread_block(); + auto warp = cooperative_groups::tiled_partition<32>(block); + bool ok = mcs::fmcs::tryMatchIncrementalGreedyCooperative(warp, child.seed, qView, tView, tables, child.match); + __syncthreads(); + + if (threadIdx.x == 0) { + out->ok = ok; + out->child = child; + } +} + +__global__ void matchIncrementalVisitedConflictDriver(DeviceCsrView qView, + DeviceCsrView tView, + PairMatchTablesDevice tables, + IncrementalTestOut* out) { + __shared__ QueuedT16 child; + if (threadIdx.x == 0) { + mcs::fmcs::seedClearWithinThread(child.seed); + // Query: atoms 0,1,2; bonds (0,1), (0,2). Target: atoms 0,1; one + // bond (0,1). Parent has bond (0,1) mapped (qU=0->t=0, qV=1->t=1). + // Now child wants atom-adding bond (0,2), but the only target bond + // out of t=0 is (0,1) and t=1 is already visited -> must fail. + int qA[] = {0, 1}; + int tA[] = {0, 1}; + int qB[] = {0}; + int tB[] = {0}; + buildParentMatch(child.match, qA, tA, 2, qB, tB, 1); + mcs::fmcs::seedAddBondWithinThread(child.seed, 0); + mcs::fmcs::seedAddBondWithinThread(child.seed, 1); + mcs::fmcs::seedAddAtomWithinThread(child.seed, 0); + mcs::fmcs::seedAddAtomWithinThread(child.seed, 1); + mcs::fmcs::seedBeginGrowStepWithinThread(child.seed); + mcs::fmcs::seedAddAtomWithinThread(child.seed, 2); + } + __syncthreads(); + + auto block = cooperative_groups::this_thread_block(); + auto warp = cooperative_groups::tiled_partition<32>(block); + bool ok = mcs::fmcs::tryMatchIncrementalGreedyCooperative(warp, child.seed, qView, tView, tables, child.match); + __syncthreads(); + + if (threadIdx.x == 0) { + out->ok = ok; + out->child = child; + } +} + +__global__ void matchIncrementalTwoBondChainDriver(DeviceCsrView qView, + DeviceCsrView tView, + PairMatchTablesDevice tables, + IncrementalTestOut* out) { + __shared__ QueuedT16 child; + if (threadIdx.x == 0) { + mcs::fmcs::seedClearWithinThread(child.seed); + // 4-atom path 0-1-2-3. Parent has only bond (0,1) mapped. Child + // adds bonds (1,2) and (2,3) in one matchIncrementalFast call. + int qA[] = {0, 1}; + int tA[] = {0, 1}; + int qB[] = {0}; + int tB[] = {0}; + buildParentMatch(child.match, qA, tA, 2, qB, tB, 1); + mcs::fmcs::seedAddBondWithinThread(child.seed, 0); + mcs::fmcs::seedAddBondWithinThread(child.seed, 1); + mcs::fmcs::seedAddBondWithinThread(child.seed, 2); + mcs::fmcs::seedAddAtomWithinThread(child.seed, 0); + mcs::fmcs::seedAddAtomWithinThread(child.seed, 1); + mcs::fmcs::seedBeginGrowStepWithinThread(child.seed); + mcs::fmcs::seedAddAtomWithinThread(child.seed, 2); + mcs::fmcs::seedAddAtomWithinThread(child.seed, 3); + } + __syncthreads(); + + auto block = cooperative_groups::this_thread_block(); + auto warp = cooperative_groups::tiled_partition<32>(block); + bool ok = mcs::fmcs::tryMatchIncrementalGreedyCooperative(warp, child.seed, qView, tView, tables, child.match); + __syncthreads(); + + if (threadIdx.x == 0) { + out->ok = ok; + out->child = child; + } +} + +} // namespace mcs_fmcs_incremental_test + +TEST(FMCSUnit, TryMatchIncrementalGreedyAtomAdding) { + using mcs_fmcs_incremental_test::IncrementalTestOut; + using mcs_fmcs_incremental_test::matchIncrementalAtomAddingDriver; + + // Query/target are both a 3-atom path 0-1-2 with bonds (0,1), (1,2). + TestGraph query(3, + { + {0, 1}, + {1, 2} + }); + TestGraph target(3, + { + {0, 1}, + {1, 2} + }); + ManagedMatchTables tables; + tables.allocate(3, 3, 2, 2); + tables.setAllAtomBits(); + tables.setAllBondBits(); + + AsyncDevicePtr d_out; + matchIncrementalAtomAddingDriver<<<1, 32>>>(query.view(), target.view(), tables.device(), d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + IncrementalTestOut out{}; + d_out.get(out); + + EXPECT_TRUE(out.ok); + EXPECT_EQ(out.child.match.targetBondIdx[1], 1u); + EXPECT_EQ(out.child.match.targetAtomIdx[2], 2u); + EXPECT_EQ(out.child.match.matchedBondSize, 2); + EXPECT_EQ(out.child.match.matchedAtomSize, 3); +} + +TEST(FMCSUnit, TryMatchIncrementalGreedyRingClosing) { + using mcs_fmcs_incremental_test::IncrementalTestOut; + using mcs_fmcs_incremental_test::matchIncrementalRingClosingDriver; + + // 4-atom square with one diagonal-free closure. Bonds: (0,1) (1,2) + // (2,3) (0,3). + TestGraph query(4, + { + {0, 1}, + {1, 2}, + {2, 3}, + {0, 3} + }); + TestGraph target(4, + { + {0, 1}, + {1, 2}, + {2, 3}, + {0, 3} + }); + ManagedMatchTables tables; + tables.allocate(4, 4, 4, 4); + tables.setAllAtomBits(); + tables.setAllBondBits(); + + AsyncDevicePtr d_out; + matchIncrementalRingClosingDriver<<<1, 32>>>(query.view(), target.view(), tables.device(), d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + IncrementalTestOut out{}; + d_out.get(out); + + EXPECT_TRUE(out.ok); + EXPECT_EQ(out.child.match.targetBondIdx[3], 3u); + EXPECT_EQ(out.child.match.matchedBondSize, 4); + // No new atoms added by the ring-closing bond. + EXPECT_EQ(out.child.match.matchedAtomSize, 4); +} + +TEST(FMCSUnit, TryMatchIncrementalGreedyVisitedConflictNeedsFallback) { + using mcs_fmcs_incremental_test::IncrementalTestOut; + using mcs_fmcs_incremental_test::matchIncrementalVisitedConflictDriver; + + // Query: 3 atoms / 2 bonds. Target: 2 atoms / 1 bond. Parent has + // bond (0,1) mapped; trying to extend with bond (0,2) forces atom 2 + // onto target atom 1, which is already visited. + TestGraph query(3, + { + {0, 1}, + {0, 2} + }); + TestGraph target(2, + { + {0, 1} + }); + ManagedMatchTables tables; + tables.allocate(3, 2, 2, 1); + tables.setAllAtomBits(); + tables.setAllBondBits(); + + AsyncDevicePtr d_out; + matchIncrementalVisitedConflictDriver<<<1, 32>>>(query.view(), target.view(), tables.device(), d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + IncrementalTestOut out{}; + d_out.get(out); + + EXPECT_FALSE(out.ok); +} + +TEST(FMCSUnit, TryMatchIncrementalGreedyTwoBondChain) { + using mcs_fmcs_incremental_test::IncrementalTestOut; + using mcs_fmcs_incremental_test::matchIncrementalTwoBondChainDriver; + + // Both sides are the 4-atom path 0-1-2-3 with bonds 0=(0,1), 1=(1,2), + // 2=(2,3). + TestGraph query(4, + { + {0, 1}, + {1, 2}, + {2, 3} + }); + TestGraph target(4, + { + {0, 1}, + {1, 2}, + {2, 3} + }); + ManagedMatchTables tables; + tables.allocate(4, 4, 3, 3); + tables.setAllAtomBits(); + tables.setAllBondBits(); + + AsyncDevicePtr d_out; + matchIncrementalTwoBondChainDriver<<<1, 32>>>(query.view(), target.view(), tables.device(), d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + IncrementalTestOut out{}; + d_out.get(out); + + EXPECT_TRUE(out.ok); + EXPECT_EQ(out.child.match.targetBondIdx[1], 1u); + EXPECT_EQ(out.child.match.targetBondIdx[2], 2u); + EXPECT_EQ(out.child.match.targetAtomIdx[2], 2u); + EXPECT_EQ(out.child.match.targetAtomIdx[3], 3u); + EXPECT_EQ(out.child.match.matchedBondSize, 3); + EXPECT_EQ(out.child.match.matchedAtomSize, 4); +}