From 6779070dd4698f83ea5556c66fe674296a45b6f2 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Fri, 24 Jul 2026 10:13:06 -0400 Subject: [PATCH 1/3] Add independently tested fMCS seed growth --- src/mcs/fmcs_cuda/fmcs_grow.cuh | 208 +++++++++++++++ tests/CMakeLists.txt | 7 + tests/test_fmcs_grow.cu | 442 ++++++++++++++++++++++++++++++++ 3 files changed, 657 insertions(+) create mode 100644 src/mcs/fmcs_cuda/fmcs_grow.cuh create mode 100644 tests/test_fmcs_grow.cu diff --git a/src/mcs/fmcs_cuda/fmcs_grow.cuh b/src/mcs/fmcs_cuda/fmcs_grow.cuh new file mode 100644 index 00000000..73c8bfc9 --- /dev/null +++ b/src/mcs/fmcs_cuda/fmcs_grow.cuh @@ -0,0 +1,208 @@ +// 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_GROW_CUH +#define FMCS_CUDA_FMCS_GROW_CUH + +#include + +#include "fmcs_cuda/fmcs_seed.cuh" +#include "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); + if (bond.endAtomSeedIdx == NewBond::kNotInSeed) { + seedAddAtomWithinThread(seed, bond.newAtomIdx); + } +} + +// Seed growth follows RDKit Seed::grow()'s three stage shape: +// +// Stage 0 (fillNewBonds + biggest-child match): collect every outgoing +// query bond and build the all-bonds-added child. +// Stage 1: for each new bond build the singleton extension child; failed +// singleton matches remove that NewBond from Stage 2. +// Stage 2: enumerate the non-singleton subsets of the surviving NewBonds. +// +// The kernel owns the Stage 0/1/2 control flow because it has the queue, +// incumbent, and substructure scratch. This file supplies the shared NewBond +// enumeration and seed-patching helpers used by those stages. +// +// Every helper is cooperative-group-parallel across bonds. + +/// Cooperative: bond-centric scan that distributes query bonds across +/// @p group's lanes. A query bond becomes a NewBond entry iff it is +/// (a) not in @c seed.excludedBonds and (b) incident to at least one +/// bit in @c seed.lastAddedAtoms (i.e., one endpoint was newly added in +/// the most recent grow step). Each surviving bond is classified: +/// +/// - Both endpoints already in @c seed.atoms -> ring-closing. +/// @c endAtomSeedIdx is set to a non-sentinel value to flag this +/// to the kernel's Stage 0 patch loop (which then skips +/// @c seedAddAtom for this bond). +/// - Otherwise -> atom-adding. @c newAtomIdx holds the query atom +/// index that is NOT yet in the seed; @c endAtomSeedIdx is set to +/// @c NewBond::kNotInSeed. +/// +/// Lanes append to @p outBonds via @c atomicAdd on @p outCount. If +/// the count would exceed @p maxNewBonds the function returns false +/// and clamps @p outCount to @p maxNewBonds (slots beyond that bound +/// are not written). No I/O ordering on @p outBonds is guaranteed -- +/// the order depends on lane race outcomes. +template +__device__ __forceinline__ bool fillNewBondsCooperative(const GroupT& group, + const Seed& seed, + const QueryTopology& queryTopology, + NewBond* outBonds, + int* outCount, + int maxNewBonds) { + using SeedT = Seed; + using AtomWord = typename SeedT::atom_word_type; + using BondWord = typename SeedT::bond_word_type; + constexpr int kAtomBitsPerWord = SeedT::kAtomBitsPerWord; + constexpr int kBondBitsPerWord = SeedT::kBondBitsPerWord; + + const int laneRank = static_cast(group.thread_rank()); + const int laneCount = static_cast(group.num_threads()); + + if (laneRank == 0) + *outCount = 0; + group.sync(); + + for (int q = laneRank; q < queryTopology.numBonds; q += laneCount) { + // Excluded check first (cheapest). + const BondWord excludedBondsWord = seed.excludedBonds[q / kBondBitsPerWord]; + if ((excludedBondsWord >> (q % kBondBitsPerWord)) & 1) + continue; + + // 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); + + // Bond is a "new boundary" candidate iff at least one endpoint + // was added in the most recent grow step. + const AtomWord lastAddedWordU = seed.lastAddedAtoms[queryEndpointU / kAtomBitsPerWord]; + const AtomWord lastAddedWordV = seed.lastAddedAtoms[queryEndpointV / kAtomBitsPerWord]; + const bool uIsNewlyAdded = (lastAddedWordU >> (queryEndpointU % kAtomBitsPerWord)) & 1; + const bool vIsNewlyAdded = (lastAddedWordV >> (queryEndpointV % kAtomBitsPerWord)) & 1; + if (!uIsNewlyAdded && !vIsNewlyAdded) + continue; + + // Classify ring-closing vs atom-adding via seed.atoms membership. + const AtomWord seedAtomsWordU = seed.atoms[queryEndpointU / kAtomBitsPerWord]; + const AtomWord seedAtomsWordV = seed.atoms[queryEndpointV / kAtomBitsPerWord]; + const bool uInSeed = (seedAtomsWordU >> (queryEndpointU % kAtomBitsPerWord)) & 1; + const bool vInSeed = (seedAtomsWordV >> (queryEndpointV % kAtomBitsPerWord)) & 1; + + NewBond newBond; + newBond.bondIdx = static_cast(q); + newBond.alive = true; + if (uInSeed && vInSeed) { + // Ring-closing: both endpoints already in the seed. newAtomIdx + // is unused by the Stage 0 patch loop in this case; we just + // need endAtomSeedIdx to be non-sentinel to flag ring-closing. + newBond.newAtomIdx = static_cast(queryEndpointV); + newBond.endAtomSeedIdx = static_cast(queryEndpointU); + } else { + newBond.newAtomIdx = static_cast(uInSeed ? queryEndpointV : queryEndpointU); + newBond.endAtomSeedIdx = NewBond::kNotInSeed; + } + + const int slot = atomicAdd(outCount, 1); + if (slot < maxNewBonds) { + outBonds[slot] = newBond; + } + // Otherwise: no out-of-bounds write, but outCount has overrun + // maxNewBonds and we'll clamp + return false below. + } + group.sync(); + + // Read the pre-clamp count so all lanes can return a consistent + // success/overflow signal even after lane 0 clamps below. + const int totalAttempted = *outCount; + group.sync(); + if (laneRank == 0 && totalAttempted > maxNewBonds) { + *outCount = maxNewBonds; + } + group.sync(); + return totalAttempted <= maxNewBonds; +} + +/// Stage 1 driver. Despite the @c Cooperative suffix, the parallelism +/// here is shallow: the outer loop walks @p bonds serially (uniform +/// 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). +/// +/// In other words: this function is cooperative only for the +/// parent->child copy and for delegating to its callbacks; it does +/// not itself parallelize across @p bonds. Future refactors aiming +/// for cross-bond parallelism would need a different structure -- the +/// per-bond commits depend on shared visited-state from the previous +/// bond's match call, which is incompatible with running them +/// concurrently. +/// +/// @p childWorkspace must be a per-group shared slot the caller owns +/// -- the kernel reuses the same buffer it already allocated for the +/// Stage 0 biggest-child build. +template +__device__ __forceinline__ void pruneIndividualBondsCooperative( + const GroupT& group, + const QueuedSeed& parent, + QueuedSeed& childWorkspace, + const NewBond* bonds, + int nBonds, + MatchFn&& matchFn, + ChildSink&& childSink) { + using QueuedT = QueuedSeed; + for (int i = 0; i < nBonds; ++i) { + if (!bonds[i].alive) + continue; + + // Reset the workspace to a copy of the parent (group-cooperative). + warpCopy(group, &childWorkspace, &parent, sizeof(QueuedT)); + group.sync(); + + // Lane 0 patches in this single new bond / atom on top of the parent. + if (group.thread_rank() == 0) { + seedBeginGrowStepWithinThread(childWorkspace.seed); + seedAddNewBondWithinThread(childWorkspace.seed, bonds[i]); + } + group.sync(); + + // Try to extend the parent's recorded match by this single bond. + if (matchFn(childWorkspace)) { + childSink(childWorkspace); + } + group.sync(); + } +} + +} // namespace fmcs +} // namespace mcs + +#endif // FMCS_CUDA_FMCS_GROW_CUH diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2bb4fd26..692478e0 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -72,6 +72,12 @@ target_include_directories(test_fmcs_primitives PRIVATE ${CMAKE_SOURCE_DIR}/src/mcs) target_link_libraries(test_fmcs_primitives PRIVATE CUDA::cudart device_vector) +add_executable(test_fmcs_grow test_fmcs_grow.cu) +target_include_directories(test_fmcs_grow PRIVATE ${CMAKE_SOURCE_DIR}/src/mcs) +target_link_libraries(test_fmcs_grow PRIVATE CUDA::cudart device_vector) +target_compile_options(test_fmcs_grow + PRIVATE $<$:--extended-lambda>) + add_executable(test_openmp_helpers test_openmp_helpers.cpp) target_link_libraries(test_openmp_helpers PRIVATE openmp_helpers OpenMP::OpenMP_CXX) @@ -417,6 +423,7 @@ set(TEST_LIST test_similarity test_thread_safe_queue test_fmcs_primitives + test_fmcs_grow test_work_splitting test_fire_minimizer test_fire_minimizer_permol) diff --git a/tests/test_fmcs_grow.cu b/tests/test_fmcs_grow.cu new file mode 100644 index 00000000..4005e082 --- /dev/null +++ b/tests/test_fmcs_grow.cu @@ -0,0 +1,442 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include +#include +#include + +#include +#include +#include +#include + +#include "fmcs_cuda/fmcs_grow.cuh" +#include "src/utils/device_vector.h" + +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) { + host[i] = (static_cast(edges[i].first) << 16) | static_cast(edges[i].second); + } + AsyncDeviceVector device(edges.size()); + device.copyFromHost(host); + return device; +} + +} // namespace + +// --------------------------------------------------------------------------- +// fillNewBondsCooperative / pruneIndividualBondsCooperative +// --------------------------------------------------------------------------- + +namespace mcs_fmcs_grow_test { + +using mcs::fmcs::NewBond; +using SeedT = mcs::fmcs::Seed<16, 16>; +using QueuedT = mcs::fmcs::QueuedSeed<16, 16, 16, 16>; + +constexpr int kMaxNewBonds = 8; + +struct FillNewBondsOut { + NewBond bonds[kMaxNewBonds]; + int count; + bool ok; +}; + +// Driver: caller-provided seed-setup function fills the seed before +// calling fillNewBondsCooperative. Templated on the setup so each +// test can hand-construct its own scenario. +template +__device__ __forceinline__ void fillNewBondsRun(SeedSetup&& setup, + const std::uint32_t* qBondEndpoints, + int qNumAtoms, + int qNumBonds, + int maxNewBonds, + FillNewBondsOut* out) { + __shared__ SeedT seed; + __shared__ NewBond bonds[kMaxNewBonds]; + __shared__ int count; + + if (threadIdx.x == 0) { + mcs::fmcs::seedClearWithinThread(seed); + setup(seed); + } + __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); + __syncthreads(); + + if (threadIdx.x == 0) { + out->ok = ok; + out->count = count; + for (int i = 0; i < kMaxNewBonds; ++i) + out->bonds[i] = bonds[i]; + } +} + +// Test 1: 3-atom path query (atoms 0-1-2, bonds (0,1), (1,2)). Seed +// holds just atom 1, marked last-added. Both bonds touch atom 1; +// each should become an atom-adding NewBond. +__global__ void fillNewBondsAtomAddingDriver(const std::uint32_t* qBE, + int qNumAtoms, + int qNumBonds, + FillNewBondsOut* out) { + fillNewBondsRun([] __device__(SeedT & seed) { mcs::fmcs::seedAddAtomWithinThread(seed, 1); }, + qBE, + qNumAtoms, + qNumBonds, + kMaxNewBonds, + out); +} + +// Test 2: same query, seed has bond 0 in excludedBonds. Only bond 1 +// should appear in the output. +__global__ void fillNewBondsExcludedSkippedDriver(const std::uint32_t* qBE, + int qNumAtoms, + int qNumBonds, + FillNewBondsOut* out) { + fillNewBondsRun( + [] __device__(SeedT & seed) { + mcs::fmcs::seedAddAtomWithinThread(seed, 1); + // Mark bond 0 as excluded but DO NOT add it to seed.bonds. + // The exclusion alone is what fillNewBonds checks. + using BondWord = SeedT::bond_word_type; + constexpr int kBPW = SeedT::kBondBitsPerWord; + seed.excludedBonds[0 / kBPW] |= static_cast(1) << (0 % kBPW); + }, + qBE, + qNumAtoms, + qNumBonds, + kMaxNewBonds, + out); +} + +// Test 3: 3-atom triangle query (atoms 0,1,2; bonds (0,1), (1,2), (0,2)). +// Seed holds atoms 0 and 1 (both last-added) plus bond (0,1). Bond 2 +// = (0,2) is atom-adding (atom 2 not in seed). More importantly, this +// test verifies the seed.atoms vs seed.lastAddedAtoms split: if we +// only mark atom 1 as "last added" (atom 0 stays in seed.atoms but +// NOT in lastAddedAtoms), bond (0,2) should NOT be reported because +// neither endpoint is "newly added". Bond (1,2) SHOULD be reported. +__global__ void fillNewBondsLastAddedFilteringDriver(const std::uint32_t* qBE, + int qNumAtoms, + int qNumBonds, + FillNewBondsOut* out) { + fillNewBondsRun( + [] __device__(SeedT & seed) { + // Atom 0 is "old" (in seed but not last-added). + seed.atoms[0] |= 1u << 0; + seed.numAtoms = 1; + // Begin a grow step boundary. + mcs::fmcs::seedBeginGrowStepWithinThread(seed); + // Atom 1 is "newly added". + mcs::fmcs::seedAddAtomWithinThread(seed, 1); + // Bond 0 already mapped (in seed and excluded). + mcs::fmcs::seedAddBondWithinThread(seed, 0); + }, + qBE, + qNumAtoms, + qNumBonds, + kMaxNewBonds, + out); +} + +// Test 4: ring-closing. Seed has atoms 0, 1, 2 all marked last-added, +// bond 0 = (0,1) and bond 1 = (1,2) already in excludedBonds. Bond 2 +// = (0,2) closes the ring: both endpoints in seed. +__global__ void fillNewBondsRingClosingDriver(const std::uint32_t* qBE, + int qNumAtoms, + int qNumBonds, + FillNewBondsOut* out) { + fillNewBondsRun( + [] __device__(SeedT & seed) { + mcs::fmcs::seedAddAtomWithinThread(seed, 0); + mcs::fmcs::seedAddAtomWithinThread(seed, 1); + mcs::fmcs::seedAddAtomWithinThread(seed, 2); + mcs::fmcs::seedAddBondWithinThread(seed, 0); + mcs::fmcs::seedAddBondWithinThread(seed, 1); + // Bond 2 = (0,2) is the only candidate; it's ring-closing + // since both 0 and 2 are in seed.atoms. + }, + qBE, + qNumAtoms, + qNumBonds, + kMaxNewBonds, + out); +} + +// Test 5: overflow. 4-atom path with seed = {atom 0 newly-added}. +// Query bonds 0..3 all touch atom 0 (star graph: bonds (0,1) (0,2) +// (0,3) (0,4)). maxNewBonds = 2 -> first 2 win the race, function +// returns false, count clamped at 2. +__global__ void fillNewBondsOverflowDriver(const std::uint32_t* qBE, + int qNumAtoms, + int qNumBonds, + FillNewBondsOut* out) { + fillNewBondsRun([] __device__(SeedT & seed) { mcs::fmcs::seedAddAtomWithinThread(seed, 0); }, + qBE, + qNumAtoms, + qNumBonds, + /*maxNewBonds=*/2, + out); +} + +} // namespace mcs_fmcs_grow_test + +TEST(FMCSUnit, FillNewBondsAtomAddingFromBoundary) { + using namespace mcs_fmcs_grow_test; + // Path 0-1-2: bonds (0,1), (1,2). + auto qBE = makeBondEndpointsDevice({ + {0, 1}, + {1, 2} + }); + AsyncDevicePtr d_out; + + fillNewBondsAtomAddingDriver<<<1, 32>>>(qBE.data(), 3, 2, d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + FillNewBondsOut out{}; + ASSERT_EQ(cudaMemcpy(&out, d_out.data(), sizeof(out), cudaMemcpyDeviceToHost), cudaSuccess); + + EXPECT_TRUE(out.ok); + EXPECT_EQ(out.count, 2); + // Both bonds appeared; race-order is non-deterministic so collect + // by bondIdx. + std::set seenBonds; + for (int i = 0; i < out.count; ++i) { + seenBonds.insert(out.bonds[i].bondIdx); + EXPECT_EQ(out.bonds[i].endAtomSeedIdx, NewBond::kNotInSeed); + EXPECT_TRUE(out.bonds[i].alive); + } + EXPECT_TRUE(seenBonds.count(0)); + EXPECT_TRUE(seenBonds.count(1)); +} + +TEST(FMCSUnit, FillNewBondsExcludedSkipped) { + using namespace mcs_fmcs_grow_test; + auto qBE = makeBondEndpointsDevice({ + {0, 1}, + {1, 2} + }); + AsyncDevicePtr d_out; + + fillNewBondsExcludedSkippedDriver<<<1, 32>>>(qBE.data(), 3, 2, d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + FillNewBondsOut out{}; + ASSERT_EQ(cudaMemcpy(&out, d_out.data(), sizeof(out), cudaMemcpyDeviceToHost), cudaSuccess); + + EXPECT_TRUE(out.ok); + EXPECT_EQ(out.count, 1); + EXPECT_EQ(out.bonds[0].bondIdx, 1); // bond 0 excluded; only bond 1 left +} + +TEST(FMCSUnit, FillNewBondsLastAddedFiltering) { + using namespace mcs_fmcs_grow_test; + // Triangle: bonds 0=(0,1), 1=(1,2), 2=(0,2). + auto qBE = makeBondEndpointsDevice({ + {0, 1}, + {1, 2}, + {0, 2} + }); + AsyncDevicePtr d_out; + + fillNewBondsLastAddedFilteringDriver<<<1, 32>>>(qBE.data(), 3, 3, d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + FillNewBondsOut out{}; + ASSERT_EQ(cudaMemcpy(&out, d_out.data(), sizeof(out), cudaMemcpyDeviceToHost), cudaSuccess); + + EXPECT_TRUE(out.ok); + // Bond 0 is excluded. Bond 1 = (1,2) touches newly-added atom 1. + // Bond 2 = (0,2) touches atom 0 (NOT newly-added) and atom 2 (also + // not in seed) -> neither endpoint newly-added, must be skipped. + EXPECT_EQ(out.count, 1); + EXPECT_EQ(out.bonds[0].bondIdx, 1); + EXPECT_EQ(out.bonds[0].endAtomSeedIdx, NewBond::kNotInSeed); + EXPECT_EQ(out.bonds[0].newAtomIdx, 2); // atom 2 is the unmapped end +} + +TEST(FMCSUnit, FillNewBondsRingClosing) { + using namespace mcs_fmcs_grow_test; + // Triangle: bonds 0=(0,1), 1=(1,2), 2=(0,2). + auto qBE = makeBondEndpointsDevice({ + {0, 1}, + {1, 2}, + {0, 2} + }); + AsyncDevicePtr d_out; + + fillNewBondsRingClosingDriver<<<1, 32>>>(qBE.data(), 3, 3, d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + FillNewBondsOut out{}; + ASSERT_EQ(cudaMemcpy(&out, d_out.data(), sizeof(out), cudaMemcpyDeviceToHost), cudaSuccess); + + EXPECT_TRUE(out.ok); + // Only bond 2 = (0,2) is candidate (others are excluded). Both + // endpoints in seed -> ring-closing. + EXPECT_EQ(out.count, 1); + EXPECT_EQ(out.bonds[0].bondIdx, 2); + EXPECT_NE(out.bonds[0].endAtomSeedIdx, NewBond::kNotInSeed); +} + +TEST(FMCSUnit, FillNewBondsOverflowReturnsFalse) { + using namespace mcs_fmcs_grow_test; + // Star graph: 4 bonds from atom 0 -> 1, 2, 3, 4. numAtoms=5. + auto qBE = makeBondEndpointsDevice({ + {0, 1}, + {0, 2}, + {0, 3}, + {0, 4} + }); + AsyncDevicePtr d_out; + + fillNewBondsOverflowDriver<<<1, 32>>>(qBE.data(), 5, 4, d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + FillNewBondsOut out{}; + ASSERT_EQ(cudaMemcpy(&out, d_out.data(), sizeof(out), cudaMemcpyDeviceToHost), cudaSuccess); + + EXPECT_FALSE(out.ok); + // count is clamped to maxNewBonds (=2). + EXPECT_EQ(out.count, 2); +} + +// ---- pruneIndividualBondsCooperative ---- + +namespace mcs_fmcs_prune_test { + +using mcs::fmcs::NewBond; +using mcs_fmcs_grow_test::QueuedT; +using mcs_fmcs_grow_test::SeedT; + +constexpr int kMaxNewBondsPrune = 8; + +struct PruneOut { + // Identity of children that survived matchFn and reached childSink. + // We record the child seed's bondIdx-of-the-newly-added bond (the + // first set bit beyond the parent's bond bitset). That uniquely + // identifies which NewBond the child was built from. + int survivorBondIdx[kMaxNewBondsPrune]; + int numSurvivors; + int numMatchAttempts; +}; + +__global__ void pruneStage1Driver(const NewBond* hostBonds, int nBonds, PruneOut* out) { + __shared__ std::uint64_t parentStorage[(sizeof(QueuedT) + sizeof(std::uint64_t) - 1) / sizeof(std::uint64_t)]; + __shared__ std::uint64_t workspaceStorage[(sizeof(QueuedT) + sizeof(std::uint64_t) - 1) / sizeof(std::uint64_t)]; + __shared__ NewBond bonds[kMaxNewBondsPrune]; + __shared__ int matchAttempts; + __shared__ int survivorIdx; + QueuedT& parent = *reinterpret_cast(parentStorage); + QueuedT& workspace = *reinterpret_cast(workspaceStorage); + + if (threadIdx.x == 0) { + mcs::fmcs::seedClearWithinThread(parent.seed); + mcs::fmcs::matchResultClearWithinThread(parent.match); + // Parent has bond 0 already in its bitset (so pruneIndividualBonds + // appears as "extending past bond 0" for tracking). + mcs::fmcs::seedAddAtomWithinThread(parent.seed, 0); + mcs::fmcs::seedAddAtomWithinThread(parent.seed, 1); + mcs::fmcs::seedAddBondWithinThread(parent.seed, 0); + for (int i = 0; i < nBonds; ++i) + bonds[i] = hostBonds[i]; + matchAttempts = 0; + survivorIdx = 0; + for (int i = 0; i < kMaxNewBondsPrune; ++i) { + out->survivorBondIdx[i] = -1; + } + } + __syncthreads(); + + auto block = cooperative_groups::this_thread_block(); + auto warp = cooperative_groups::tiled_partition<32>(block); + + // Stub matchFn: returns true iff the new bond's bondIdx is even. + // Counts attempts so the test can verify every bond was tried. + auto matchFn = [&] __device__(QueuedT & child) -> bool { + if (warp.thread_rank() == 0) + ++matchAttempts; + warp.sync(); + // Find which bondIdx the workspace's seed has beyond the parent's. + // Parent has only bond 0 set; child has one more. + int childBond = -1; + for (int b = 0; b < 16; ++b) { + using BondWord = SeedT::bond_word_type; + const BondWord childWord = child.seed.bonds[0]; + const BondWord parentWord = parent.seed.bonds[0]; + const BondWord newBits = childWord & ~parentWord; + if ((newBits >> b) & 1) { + childBond = b; + break; + } + } + return (childBond % 2) == 0; + }; + auto childSink = [&] __device__(QueuedT & child) { + if (warp.thread_rank() == 0) { + // Record child's new-bond identity. + using BondWord = SeedT::bond_word_type; + const BondWord newBits = child.seed.bonds[0] & ~parent.seed.bonds[0]; + int childBond = -1; + for (int b = 0; b < 16; ++b) { + if ((newBits >> b) & 1) { + childBond = b; + break; + } + } + out->survivorBondIdx[survivorIdx++] = childBond; + } + warp.sync(); + }; + mcs::fmcs::pruneIndividualBondsCooperative(warp, parent, workspace, bonds, nBonds, matchFn, childSink); + + __syncthreads(); + if (threadIdx.x == 0) { + out->numMatchAttempts = matchAttempts; + out->numSurvivors = survivorIdx; + } +} + +} // namespace mcs_fmcs_prune_test + +TEST(FMCSUnit, PruneIndividualBondsStage1OnlySurvivorsSinked) { + using namespace mcs_fmcs_prune_test; + // 4 candidate bonds with bondIdx 1, 2, 3, 4. Stub matchFn passes + // even bondIdx (-> 2, 4 succeed; 1, 3 fail). All 4 must be tried + // (matchAttempts == 4); only 2 survivors should reach childSink. + const std::vector hostBonds = { + NewBond{1, 5, NewBond::kNotInSeed, true}, + NewBond{2, 6, NewBond::kNotInSeed, true}, + NewBond{3, 7, NewBond::kNotInSeed, true}, + NewBond{4, 8, NewBond::kNotInSeed, true}, + }; + AsyncDeviceVector d_bonds(hostBonds.size()); + d_bonds.copyFromHost(hostBonds); + + AsyncDevicePtr d_out; + pruneStage1Driver<<<1, 32>>>(d_bonds.data(), 4, d_out.data()); + ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); + PruneOut out{}; + ASSERT_EQ(cudaMemcpy(&out, d_out.data(), sizeof(out), cudaMemcpyDeviceToHost), cudaSuccess); + + EXPECT_EQ(out.numMatchAttempts, 4); // all four bonds were tried + EXPECT_EQ(out.numSurvivors, 2); + std::set survivors{out.survivorBondIdx[0], out.survivorBondIdx[1]}; + EXPECT_TRUE(survivors.count(2)); + EXPECT_TRUE(survivors.count(4)); + EXPECT_FALSE(survivors.count(1)); + EXPECT_FALSE(survivors.count(3)); +} + +// --------------------------------------------------------------------------- From 594be7c36ce0d38cbe5369589566bb816ca83bd5 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Fri, 24 Jul 2026 13:40:59 -0400 Subject: [PATCH 2/3] Use project-rooted includes in fMCS growth --- src/mcs/fmcs_cuda/fmcs_grow.cuh | 4 ++-- tests/test_fmcs_grow.cu | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/mcs/fmcs_cuda/fmcs_grow.cuh b/src/mcs/fmcs_cuda/fmcs_grow.cuh index 73c8bfc9..65eecf03 100644 --- a/src/mcs/fmcs_cuda/fmcs_grow.cuh +++ b/src/mcs/fmcs_cuda/fmcs_grow.cuh @@ -18,8 +18,8 @@ #include -#include "fmcs_cuda/fmcs_seed.cuh" -#include "mcs_common/mcs_cooperative_copy.cuh" +#include "src/mcs/fmcs_cuda/fmcs_seed.cuh" +#include "src/mcs/mcs_common/mcs_cooperative_copy.cuh" namespace mcs { namespace fmcs { diff --git a/tests/test_fmcs_grow.cu b/tests/test_fmcs_grow.cu index 4005e082..cf4a5c35 100644 --- a/tests/test_fmcs_grow.cu +++ b/tests/test_fmcs_grow.cu @@ -10,7 +10,7 @@ #include #include -#include "fmcs_cuda/fmcs_grow.cuh" +#include "src/mcs/fmcs_cuda/fmcs_grow.cuh" #include "src/utils/device_vector.h" namespace { From eb7b9340f1f2e44bf229ccdb3bdef035639e9e62 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Tue, 28 Jul 2026 12:12:26 -0400 Subject: [PATCH 3/3] Address review comments --- src/mcs/fmcs_cuda/fmcs_grow.cuh | 4 +++- tests/test_fmcs_grow.cu | 35 ++++++++++++++++++++++----------- 2 files changed, 26 insertions(+), 13 deletions(-) diff --git a/src/mcs/fmcs_cuda/fmcs_grow.cuh b/src/mcs/fmcs_cuda/fmcs_grow.cuh index 65eecf03..14b6b418 100644 --- a/src/mcs/fmcs_cuda/fmcs_grow.cuh +++ b/src/mcs/fmcs_cuda/fmcs_grow.cuh @@ -174,7 +174,7 @@ __device__ __forceinline__ void pruneIndividualBondsCooperative( const GroupT& group, const QueuedSeed& parent, QueuedSeed& childWorkspace, - const NewBond* bonds, + NewBond* bonds, int nBonds, MatchFn&& matchFn, ChildSink&& childSink) { @@ -197,6 +197,8 @@ __device__ __forceinline__ void pruneIndividualBondsCooperative( // Try to extend the parent's recorded match by this single bond. if (matchFn(childWorkspace)) { childSink(childWorkspace); + } else if (group.thread_rank() == 0) { + bonds[i].alive = false; } group.sync(); } diff --git a/tests/test_fmcs_grow.cu b/tests/test_fmcs_grow.cu index cf4a5c35..6f23af04 100644 --- a/tests/test_fmcs_grow.cu +++ b/tests/test_fmcs_grow.cu @@ -327,9 +327,10 @@ struct PruneOut { // We record the child seed's bondIdx-of-the-newly-added bond (the // first set bit beyond the parent's bond bitset). That uniquely // identifies which NewBond the child was built from. - int survivorBondIdx[kMaxNewBondsPrune]; - int numSurvivors; - int numMatchAttempts; + int survivorBondIdx[kMaxNewBondsPrune]; + bool bondAlive[kMaxNewBondsPrune]; + int numSurvivors; + int numMatchAttempts; }; __global__ void pruneStage1Driver(const NewBond* hostBonds, int nBonds, PruneOut* out) { @@ -405,6 +406,8 @@ __global__ void pruneStage1Driver(const NewBond* hostBonds, int nBonds, PruneOut if (threadIdx.x == 0) { out->numMatchAttempts = matchAttempts; out->numSurvivors = survivorIdx; + for (int i = 0; i < nBonds; ++i) + out->bondAlive[i] = bonds[i].alive; } } @@ -412,31 +415,39 @@ __global__ void pruneStage1Driver(const NewBond* hostBonds, int nBonds, PruneOut TEST(FMCSUnit, PruneIndividualBondsStage1OnlySurvivorsSinked) { using namespace mcs_fmcs_prune_test; - // 4 candidate bonds with bondIdx 1, 2, 3, 4. Stub matchFn passes - // even bondIdx (-> 2, 4 succeed; 1, 3 fail). All 4 must be tried - // (matchAttempts == 4); only 2 survivors should reach childSink. + // Five candidate bonds with bondIdx 1..5. Bond 3 is already dead and + // must be skipped without a match attempt. Stub matchFn passes even + // bondIdx (-> 2, 4 succeed; 1, 5 fail). The failed live bonds must + // also be marked dead for Stage 2. const std::vector hostBonds = { - NewBond{1, 5, NewBond::kNotInSeed, true}, - NewBond{2, 6, NewBond::kNotInSeed, true}, - NewBond{3, 7, NewBond::kNotInSeed, true}, - NewBond{4, 8, NewBond::kNotInSeed, true}, + NewBond{1, 5, NewBond::kNotInSeed, true}, + NewBond{2, 6, NewBond::kNotInSeed, true}, + NewBond{3, 7, NewBond::kNotInSeed, false}, + NewBond{4, 8, NewBond::kNotInSeed, true}, + NewBond{5, 9, NewBond::kNotInSeed, true}, }; AsyncDeviceVector d_bonds(hostBonds.size()); d_bonds.copyFromHost(hostBonds); AsyncDevicePtr d_out; - pruneStage1Driver<<<1, 32>>>(d_bonds.data(), 4, d_out.data()); + pruneStage1Driver<<<1, 32>>>(d_bonds.data(), 5, d_out.data()); ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); PruneOut out{}; ASSERT_EQ(cudaMemcpy(&out, d_out.data(), sizeof(out), cudaMemcpyDeviceToHost), cudaSuccess); - EXPECT_EQ(out.numMatchAttempts, 4); // all four bonds were tried + EXPECT_EQ(out.numMatchAttempts, 4); // pre-dead bond 3 was skipped EXPECT_EQ(out.numSurvivors, 2); std::set survivors{out.survivorBondIdx[0], out.survivorBondIdx[1]}; EXPECT_TRUE(survivors.count(2)); EXPECT_TRUE(survivors.count(4)); EXPECT_FALSE(survivors.count(1)); EXPECT_FALSE(survivors.count(3)); + EXPECT_FALSE(survivors.count(5)); + EXPECT_FALSE(out.bondAlive[0]); // singleton match failed + EXPECT_TRUE(out.bondAlive[1]); // singleton match survived + EXPECT_FALSE(out.bondAlive[2]); // already dead and remained skipped + EXPECT_TRUE(out.bondAlive[3]); // singleton match survived + EXPECT_FALSE(out.bondAlive[4]); // singleton match failed } // ---------------------------------------------------------------------------