diff --git a/benchmarks/substruct_bench.py b/benchmarks/substruct_bench.py index 4a88602e..c2840683 100644 --- a/benchmarks/substruct_bench.py +++ b/benchmarks/substruct_bench.py @@ -35,6 +35,9 @@ # Get all matches instead of just boolean: python substruct_bench.py --smiles --smarts --mode getSubstructMatches + # Select the adjacency-anchored DFS backend (GSI is the default): + python substruct_bench.py --smiles --smarts --algorithm dfs + # Limit to first 10 matches per target/query pair: python substruct_bench.py --smiles --smarts --mode getSubstructMatches --max_matches 10 @@ -51,7 +54,8 @@ python substruct_bench.py --smiles --smarts \ --rdkit_match_mode raw substructlib --rdkit_threads 1 4 16 - # Run multiple configurations from a dataframe (smarts, batch_size, workers, prep_threads, mode, num_gpus): + # Run multiple configurations from a dataframe + # (smarts, batch_size, workers, prep_threads, mode, num_gpus, optional algorithm): python substruct_bench.py --smiles --config """ @@ -67,6 +71,9 @@ import pandas as pd from bench_utils import add_rdkit_max_seconds_arg, load_pickle, load_smarts, load_smiles, time_it_bounded from benchmark_timing import time_it as _time_it +from rdkit import Chem +from rdkit.Chem import rdSubstructLibrary + from nvmolkit import autotune as nv_autotune from nvmolkit.substructure import ( SubstructSearchConfig, @@ -74,8 +81,6 @@ getSubstructMatches, hasSubstructMatch, ) -from rdkit import Chem -from rdkit.Chem import rdSubstructLibrary OPTUNA_AVAILABLE = nv_autotune.is_available() @@ -325,7 +330,7 @@ def main(): "--config", help=( "Path to config dataframe (.csv/.pkl/.pickle/.parquet) with columns: " - "smarts, batch_size, workers, prep_threads, mode, num_gpus" + "smarts, batch_size, workers, prep_threads, mode, num_gpus, optional algorithm" ), ) parser.add_argument("--num_mols", "-n", type=int, default=0, help="Max number of molecules (default: 0 = all)") @@ -346,6 +351,12 @@ def main(): default="hasSubstructMatch", help="Search mode (default: hasSubstructMatch)", ) + parser.add_argument( + "--algorithm", + choices=["gsi", "dfs"], + default=None, + help="nvmolkit matching backend (default: gsi)", + ) parser.add_argument( "--max_matches", type=int, default=0, help="Maximum matches per target/query pair, 0 = all (default: 0)" ) @@ -406,7 +417,8 @@ def main(): default=None, help=( "Path to a previously-saved SubstructSearchConfig JSON. " - "Overrides --batch_size/--workers/--prep_threads (and --num_gpus if gpuIds present in the file)." + "Overrides --batch_size/--workers/--prep_threads (and --num_gpus if gpuIds present in the file); " + "its algorithm is retained unless --algorithm is specified." ), ) parser.add_argument( @@ -509,7 +521,11 @@ def main(): print(f" SMARTS file: {args.smarts}") print(f" Mode: {args.mode}") if not args.no_nvmolkit: - print(f" nvmolkit config:") + print(" nvmolkit config:") + if args.autotune_load and args.algorithm is None: + print(" algorithm: from loaded config") + else: + print(f" algorithm: {args.algorithm or 'gsi'}") print(f" batch_size: {args.batch_size}") print(f" num_gpus: {args.num_gpus}") print(f" workers: {args.workers if args.workers >= 0 else 'auto'}") @@ -536,6 +552,7 @@ def main(): "prep_threads": args.prep_threads, "mode": args.mode, "num_gpus": args.num_gpus, + "algorithm": args.algorithm or "gsi", } ] @@ -545,12 +562,21 @@ def main(): for config_row in config_rows: smarts_path = config_row["smarts"] mode = config_row["mode"] + row_algorithm = config_row.get("algorithm", args.algorithm or "gsi") + algorithm = "gsi" if pd.isna(row_algorithm) else str(row_algorithm).lower() + if algorithm not in {"gsi", "dfs"}: + print(f"Error: algorithm must be 'gsi' or 'dfs', got {row_algorithm!r}") + sys.exit(1) print("\nRun configuration:") print(f" SMARTS file: {smarts_path}") print(f" Mode: {mode}") if not args.no_nvmolkit: - print(f" nvmolkit config:") + print(" nvmolkit config:") + if args.autotune_load and args.algorithm is None: + print(" algorithm: from loaded config") + else: + print(f" algorithm: {algorithm}") print(f" batch_size: {config_row['batch_size']}") print(f" num_gpus: {config_row['num_gpus']}") print(f" workers: {config_row['workers'] if config_row['workers'] >= 0 else 'auto'}") @@ -600,7 +626,8 @@ def main(): if not config.gpuIds: config.gpuIds = gpu_ids print( - f" Loaded: batchSize={config.batchSize}, workerThreads={config.workerThreads}, " + f" Loaded: algorithm={config.algorithm}, batchSize={config.batchSize}, " + f"workerThreads={config.workerThreads}, " f"preprocessingThreads={config.preprocessingThreads}, gpuIds={list(config.gpuIds)}" ) elif args.autotune: @@ -625,6 +652,7 @@ def main(): api=api_for_mode, maxMatches=args.max_matches, gpuIds=gpu_ids, + algorithm=algorithm, n_trials=args.autotune_trials, target_seconds_per_trial=args.autotune_time_budget, calibration_set=explicit_calibration, @@ -650,6 +678,9 @@ def main(): config.gpuIds = gpu_ids if args.max_matches > 0: config.maxMatches = args.max_matches + if not args.autotune_load or args.algorithm is not None: + config.algorithm = algorithm + algorithm = config.algorithm ran_nvmolkit = True torch_module = torch @@ -774,6 +805,7 @@ def main(): applied_workers = config_row["workers"] applied_prep_threads = config_row["prep_threads"] applied_num_gpus = config_row["num_gpus"] + applied_algorithm = algorithm if ran_nvmolkit else "N/A" if args.autotune: config_source = "autotuned" @@ -805,6 +837,7 @@ def main(): ( name, mode, + applied_algorithm if name == "nvmolkit" else "N/A", smarts_path, input_file, input_type, @@ -836,7 +869,7 @@ def main(): print("\n\nCSV Results:") print( - "method,mode,smarts,input_file,input_type,sanitize,num_mols,num_patterns," + "method,mode,algorithm,smarts,input_file,input_type,sanitize,num_mols,num_patterns," "max_matches,batch_size,num_gpus,workers,prep_threads,nvmolkit_config_source," "rdkit_threads,rdkit_match_mode,time_ms,std_ms," "pairs_processed,rdkit_max_seconds,pairs_per_second,vs_rdkit_throughput_ratio" @@ -845,6 +878,7 @@ def main(): ( name, mode, + algorithm, smarts_path, input_file, input_type, @@ -871,7 +905,7 @@ def main(): f"{rdkit_max_seconds:g}" if isinstance(rdkit_max_seconds, float) else str(rdkit_max_seconds) ) print( - f"{name},{mode},{smarts_path},{input_file},{input_type},{sanitize}," + f"{name},{mode},{algorithm},{smarts_path},{input_file},{input_type},{sanitize}," f"{num_mols},{num_patterns},{max_matches},{batch_size},{num_gpus},{workers},{prep_threads}," f"{nvmolkit_config_source},{rdkit_threads},{rdkit_match_mode},{avg_ms:.2f},{std_ms:.2f}," f"{pairs_done},{rdkit_max_seconds_str},{throughput:.2f},{vs_rdkit_str}" diff --git a/nvmolkit/autotune/tune_substructure.py b/nvmolkit/autotune/tune_substructure.py index dfbeffd7..786722e9 100644 --- a/nvmolkit/autotune/tune_substructure.py +++ b/nvmolkit/autotune/tune_substructure.py @@ -119,6 +119,7 @@ def tune_substructure( maxMatches: int = 0, uniquify: bool = False, gpuIds: Optional[Iterable[int]] = None, + algorithm: str = "gsi", calibration_set: Optional[Iterable[int]] = None, calibration_fraction: float = 0.1, calibration_max_size: int = 50_000, @@ -143,6 +144,7 @@ def tune_substructure( Held constant across trials. uniquify: ``uniquify`` flag forwarded to the resulting config. gpuIds: GPU device IDs to use. Fixed across the study. + algorithm: Matching backend to tune, either ``"gsi"`` or ``"dfs"``. calibration_set: Optional explicit indices into ``targets``. calibration_fraction: Fraction of the workload to auto-sample. calibration_max_size: Cap on the auto-sampled calibration size. @@ -194,6 +196,7 @@ def _make_config(values: dict[str, Any]) -> SubstructSearchConfig: maxMatches=int(maxMatches), uniquify=bool(uniquify), gpuIds=fixed_gpu_ids if fixed_gpu_ids else None, + algorithm=algorithm, ) queries_list = list(queries) @@ -209,6 +212,7 @@ def default_runner(state: CalibrationState) -> int: maxMatches=int(maxMatches), uniquify=bool(uniquify), gpuIds=fixed_gpu_ids if fixed_gpu_ids else None, + algorithm=algorithm, ), state, ) diff --git a/nvmolkit/substructure.cpp b/nvmolkit/substructure.cpp index a0bdd19c..98a2b5ca 100644 --- a/nvmolkit/substructure.cpp +++ b/nvmolkit/substructure.cpp @@ -20,6 +20,8 @@ #include #include #include +#include +#include #include #include "nvmolkit/boost_python_utils.h" @@ -92,6 +94,28 @@ void setGpuIdsPy(nvMolKit::SubstructSearchConfig& config, const object& iterable config.gpuIds = listFromIterable(iterable); } +std::string getAlgorithmPy(const nvMolKit::SubstructSearchConfig& config) { + switch (config.algorithm) { + case nvMolKit::SubstructAlgorithm::GSI: + return "gsi"; + case nvMolKit::SubstructAlgorithm::DFS: + return "dfs"; + case nvMolKit::SubstructAlgorithm::VF2: + throw std::invalid_argument("VF2 is not supported by the Python substructure bindings"); + } + throw std::invalid_argument("Unknown substructure algorithm"); +} + +void setAlgorithmPy(nvMolKit::SubstructSearchConfig& config, const std::string& algorithm) { + if (algorithm == "gsi") { + config.algorithm = nvMolKit::SubstructAlgorithm::GSI; + } else if (algorithm == "dfs") { + config.algorithm = nvMolKit::SubstructAlgorithm::DFS; + } else { + throw std::invalid_argument("algorithm must be 'gsi' or 'dfs'"); + } +} + } // namespace BOOST_PYTHON_MODULE(_substructure) { @@ -104,7 +128,8 @@ BOOST_PYTHON_MODULE(_substructure) { .def_readwrite("preprocessingThreads", &nvMolKit::SubstructSearchConfig::preprocessingThreads) .def_readwrite("maxMatches", &nvMolKit::SubstructSearchConfig::maxMatches) .def_readwrite("uniquify", &nvMolKit::SubstructSearchConfig::uniquify) - .add_property("gpuIds", &getGpuIdsPy, &setGpuIdsPy); + .add_property("gpuIds", &getGpuIdsPy, &setGpuIdsPy) + .add_property("algorithm", &getAlgorithmPy, &setAlgorithmPy); def( "getSubstructMatches", @@ -133,12 +158,7 @@ BOOST_PYTHON_MODULE(_substructure) { extractRange.pop(); nvMolKit::SubstructSearchResults results; - nvMolKit::getSubstructMatches(targetsVec, - queriesVec, - results, - nvMolKit::SubstructAlgorithm::GSI, - nullptr, - config); + nvMolKit::getSubstructMatches(targetsVec, queriesVec, results, config.algorithm, nullptr, config); auto csrPtr = std::make_unique(); csrPtr->numTargets = results.numTargets; @@ -244,12 +264,7 @@ BOOST_PYTHON_MODULE(_substructure) { const int numQueries = static_cast(queriesVec.size()); auto countsPtr = std::make_unique>(); - nvMolKit::countSubstructMatches(targetsVec, - queriesVec, - *countsPtr, - nvMolKit::SubstructAlgorithm::GSI, - nullptr, - config); + nvMolKit::countSubstructMatches(targetsVec, queriesVec, *countsPtr, config.algorithm, nullptr, config); nvMolKit::ScopedNvtxRange wrapRange("Python: wrap numpy array", nvMolKit::NvtxColor::kGreen); int* dataPtr = countsPtr->data(); @@ -304,12 +319,7 @@ BOOST_PYTHON_MODULE(_substructure) { extractRange.pop(); auto resultsPtr = std::make_unique(); - nvMolKit::hasSubstructMatch(targetsVec, - queriesVec, - *resultsPtr, - nvMolKit::SubstructAlgorithm::GSI, - nullptr, - config); + nvMolKit::hasSubstructMatch(targetsVec, queriesVec, *resultsPtr, config.algorithm, nullptr, config); nvMolKit::ScopedNvtxRange wrapRange("Python: wrap numpy array", nvMolKit::NvtxColor::kGreen); const int numTargets = resultsPtr->numTargets; diff --git a/nvmolkit/substructure.py b/nvmolkit/substructure.py index e4fa8b39..4c59f8e7 100644 --- a/nvmolkit/substructure.py +++ b/nvmolkit/substructure.py @@ -49,6 +49,13 @@ ] +def _normalize_algorithm(value: str) -> str: + algorithm = str(value).lower() + if algorithm not in {"gsi", "dfs"}: + raise ValueError("algorithm must be 'gsi' or 'dfs'") + return algorithm + + class SubstructSearchConfig: """Configuration for GPU substructure search execution. @@ -70,6 +77,7 @@ def __init__( maxMatches: int = 0, uniquify: bool = False, gpuIds: list[int] | None = None, + algorithm: str = "dfs", ) -> None: native = _NativeSubstructSearchConfig() native.batchSize = int(batchSize) @@ -78,6 +86,7 @@ def __init__( native.maxMatches = int(maxMatches) native.uniquify = bool(uniquify) native.gpuIds = list(gpuIds) if gpuIds is not None else [] + native.algorithm = _normalize_algorithm(algorithm) self._native = native @property @@ -134,6 +143,20 @@ def gpuIds(self) -> list[int]: def gpuIds(self, value: list[int]) -> None: self._native.gpuIds = list(value) + @property + def algorithm(self) -> str: + """Internal: matching backend, ``"dfs"`` (default) or ``"gsi"``. + + Not part of the public API -- exposed for benchmarking and testing + the two backends against each other, and may be removed once ``gsi`` + is retired. + """ + return self._native.algorithm + + @algorithm.setter + def algorithm(self, value: str) -> None: + self._native.algorithm = _normalize_algorithm(value) + def _as_native(self): """Internal: return the underlying native config object.""" return self._native @@ -147,12 +170,13 @@ def to_dict(self) -> dict[str, Any]: "maxMatches": self.maxMatches, "uniquify": self.uniquify, "gpuIds": list(self.gpuIds), + "algorithm": self.algorithm, } @classmethod def from_dict(cls, data: dict[str, Any]) -> "SubstructSearchConfig": """Create a :class:`SubstructSearchConfig` from a dictionary produced by :meth:`to_dict`.""" - known = {"batchSize", "workerThreads", "preprocessingThreads", "maxMatches", "uniquify", "gpuIds"} + known = {"batchSize", "workerThreads", "preprocessingThreads", "maxMatches", "uniquify", "gpuIds", "algorithm"} unknown = set(data) - known if unknown: raise ValueError(f"Unknown SubstructSearchConfig keys: {sorted(unknown)}") diff --git a/nvmolkit/tests/test_autotune.py b/nvmolkit/tests/test_autotune.py index e8c3833b..901752aa 100644 --- a/nvmolkit/tests/test_autotune.py +++ b/nvmolkit/tests/test_autotune.py @@ -163,6 +163,7 @@ def test_substruct_config_to_from_dict_roundtrip(): maxMatches=8, uniquify=True, gpuIds=[0], + algorithm="dfs", ) encoded = config.to_dict() assert encoded == { @@ -172,6 +173,7 @@ def test_substruct_config_to_from_dict_roundtrip(): "maxMatches": 8, "uniquify": True, "gpuIds": [0], + "algorithm": "dfs", } restored = SubstructSearchConfig.from_dict(encoded) assert restored.batchSize == 512 @@ -180,6 +182,7 @@ def test_substruct_config_to_from_dict_roundtrip(): assert restored.maxMatches == 8 assert restored.uniquify is True assert restored.gpuIds == [0] + assert restored.algorithm == "dfs" def test_save_load_hardware_options_roundtrip(tmp_path): diff --git a/nvmolkit/tests/test_substructure.py b/nvmolkit/tests/test_substructure.py index 2fd6161c..6b946c96 100644 --- a/nvmolkit/tests/test_substructure.py +++ b/nvmolkit/tests/test_substructure.py @@ -810,6 +810,7 @@ def test_default_config(self): assert config.preprocessingThreads == -1 # -1 = autoselect assert config.maxMatches == 0 # 0 = unlimited (like RDKit) assert config.uniquify is False + assert config.algorithm == "dfs" def test_config_modification(self): """Test configuration parameter modification.""" @@ -818,11 +819,19 @@ def test_config_modification(self): config.workerThreads = 4 config.preprocessingThreads = 8 config.maxMatches = 100 + config.algorithm = "gsi" assert config.batchSize == 512 assert config.workerThreads == 4 assert config.preprocessingThreads == 8 assert config.maxMatches == 100 + assert config.algorithm == "gsi" + + @pytest.mark.parametrize("algorithm", ["vf2", "unknown"]) + def test_algorithm_rejects_unsupported_value(self, algorithm): + """Test that only implemented high-level backends are accepted.""" + with pytest.raises(ValueError, match="algorithm must be"): + SubstructSearchConfig(algorithm=algorithm) def test_gpu_ids_property(self): """Test gpuIds property get/set.""" @@ -1040,6 +1049,11 @@ def validate_against_rdkit( return ValidationResult(total_pairs, count_mismatches, mapping_mismatches) +@pytest.fixture(params=["gsi", "dfs"]) +def substruct_algorithm(request: pytest.FixtureRequest) -> str: + return request.param + + @pytest.mark.long class TestIntegrationChemblSmarts: """Integration tests using ChEMBL molecules and SMARTS datasets.""" @@ -1060,7 +1074,7 @@ def chembl_mols(self) -> list[Chem.Mol]: "openbabel_functional_groups_supported.txt", ], ) - def test_chembl_vs_smarts_dataset(self, chembl_mols: list[Chem.Mol], smarts_file: str): + def test_chembl_vs_smarts_dataset(self, chembl_mols: list[Chem.Mol], smarts_file: str, substruct_algorithm: str): """Test ChEMBL molecules against SMARTS dataset files.""" smarts_path = TEST_DATA_DIR / "SMARTS" / smarts_file if not smarts_path.exists(): @@ -1070,7 +1084,7 @@ def test_chembl_vs_smarts_dataset(self, chembl_mols: list[Chem.Mol], smarts_file if not queries: pytest.skip(f"No valid queries in {smarts_file}") - config = SubstructSearchConfig() + config = SubstructSearchConfig(algorithm=substruct_algorithm) config.batchSize = 128 results = getSubstructMatches(chembl_mols, queries, config) @@ -1110,7 +1124,7 @@ def test_chembl_vs_smarts_dataset(self, chembl_mols: list[Chem.Mol], smarts_file f"Mapping mismatches: {len(validation.mapping_mismatches)}/{validation.total_pairs}\n{details}{more}" ) - def test_chembl_basic_queries(self, chembl_mols: list[Chem.Mol]): + def test_chembl_basic_queries(self, chembl_mols: list[Chem.Mol], substruct_algorithm: str): """Test ChEMBL molecules against basic SMARTS queries.""" smarts_strings = [ "[CX3]=[OX1]", # Carbonyl @@ -1121,7 +1135,7 @@ def test_chembl_basic_queries(self, chembl_mols: list[Chem.Mol]): ] queries = [Chem.MolFromSmarts(s) for s in smarts_strings] - config = SubstructSearchConfig() + config = SubstructSearchConfig(algorithm=substruct_algorithm) config.batchSize = 128 results = getSubstructMatches(chembl_mols, queries, config) @@ -1136,7 +1150,7 @@ def test_chembl_basic_queries(self, chembl_mols: list[Chem.Mol]): ) pytest.fail(f"Count mismatches: {len(validation.count_mismatches)}/{validation.total_pairs}\n{details}") - def test_chembl_has_substruct_match(self, chembl_mols: list[Chem.Mol]): + def test_chembl_has_substruct_match(self, chembl_mols: list[Chem.Mol], substruct_algorithm: str): """Test hasSubstructMatch on ChEMBL molecules.""" queries = [ Chem.MolFromSmarts("[CX3]=[OX1]"), @@ -1144,7 +1158,7 @@ def test_chembl_has_substruct_match(self, chembl_mols: list[Chem.Mol]): Chem.MolFromSmarts("[OX2H]"), ] - config = SubstructSearchConfig() + config = SubstructSearchConfig(algorithm=substruct_algorithm) config.batchSize = 128 results = hasSubstructMatch(chembl_mols, queries, config) @@ -1155,7 +1169,7 @@ def test_chembl_has_substruct_match(self, chembl_mols: list[Chem.Mol]): rdkit_has_match = target.HasSubstructMatch(query) assert bool(results[t_idx, q_idx]) == rdkit_has_match, f"Mismatch at target {t_idx}, query {q_idx}" - def test_chembl_count_substruct_matches(self, chembl_mols: list[Chem.Mol]): + def test_chembl_count_substruct_matches(self, chembl_mols: list[Chem.Mol], substruct_algorithm: str): """Test countSubstructMatches on ChEMBL molecules.""" queries = [ Chem.MolFromSmarts("[CX3]=[OX1]"), @@ -1163,7 +1177,7 @@ def test_chembl_count_substruct_matches(self, chembl_mols: list[Chem.Mol]): Chem.MolFromSmarts("[OX2H]"), ] - config = SubstructSearchConfig() + config = SubstructSearchConfig(algorithm=substruct_algorithm) config.batchSize = 128 results = countSubstructMatches(chembl_mols, queries, config) @@ -1196,11 +1210,11 @@ def test_mols(self) -> tuple[list[Chem.Mol], list[Chem.Mol]]: ] return targets, queries - def test_multithreaded_config(self, test_mols): + def test_multithreaded_config(self, test_mols, substruct_algorithm: str): """Test with multiple worker threads.""" targets, queries = test_mols - config = SubstructSearchConfig() + config = SubstructSearchConfig(algorithm=substruct_algorithm) config.workerThreads = 2 results = getSubstructMatches(targets, queries, config) @@ -1208,11 +1222,11 @@ def test_multithreaded_config(self, test_mols): validation = validate_against_rdkit(targets, queries, results) assert len(validation.count_mismatches) == 0 - def test_preprocessing_threads(self, test_mols): + def test_preprocessing_threads(self, test_mols, substruct_algorithm: str): """Test with preprocessing threads.""" targets, queries = test_mols - config = SubstructSearchConfig() + config = SubstructSearchConfig(algorithm=substruct_algorithm) config.preprocessingThreads = 4 results = getSubstructMatches(targets, queries, config) @@ -1220,11 +1234,11 @@ def test_preprocessing_threads(self, test_mols): validation = validate_against_rdkit(targets, queries, results) assert len(validation.count_mismatches) == 0 - def test_small_batch_size(self, test_mols): + def test_small_batch_size(self, test_mols, substruct_algorithm: str): """Test with smaller batch size.""" targets, queries = test_mols - config = SubstructSearchConfig() + config = SubstructSearchConfig(algorithm=substruct_algorithm) config.batchSize = 64 results = getSubstructMatches(targets, queries, config) @@ -1232,11 +1246,12 @@ def test_small_batch_size(self, test_mols): validation = validate_against_rdkit(targets, queries, results) assert len(validation.count_mismatches) == 0 - def test_results_exist(self, test_mols): + def test_results_exist(self, test_mols, substruct_algorithm: str): """Test basic match results exist.""" targets, queries = test_mols - results = getSubstructMatches(targets, queries) + config = SubstructSearchConfig(algorithm=substruct_algorithm) + results = getSubstructMatches(targets, queries, config) validation = validate_against_rdkit(targets, queries, results) assert len(validation.count_mismatches) == 0 diff --git a/src/subgraph/occupancy.cuh b/src/subgraph/occupancy.cuh new file mode 100644 index 00000000..f3ed7733 --- /dev/null +++ b/src/subgraph/occupancy.cuh @@ -0,0 +1,90 @@ +// 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. + +// Compile-time occupancy modelling for kernels whose residency is decided by +// one dominant __shared__ allocation, used to pick the second +// __launch_bounds__ argument. + +#ifndef NVMOLKIT_SUBGRAPH_OCCUPANCY_CUH +#define NVMOLKIT_SUBGRAPH_OCCUPANCY_CUH + +#include + +namespace nvMolKit { + +/** + * @brief Per-SM shared memory budget assumed when choosing minBlocksPerSM. + * + * Hopper and datacenter Blackwell (sm_90/100/103) have 228 KB per SM; everything + * else is modelled as A100's 164 KB, deliberately including the architectures + * with less (Turing 64 KB, Volta 96 KB, consumer Ampere/Ada/Blackwell 100 KB). + * ptxas does not validate .minnctapersm against shared memory, so on those an + * unreachable ask costs nothing at runtime -- occupancy just lands where shared + * memory puts it -- but it keeps the register cap as tight as the tuned + * configuration allows. Modelling the 228 KB parts does matter: the smaller + * assumed budget under-asks for the largest shapes there, and the relaxed + * register cap can drop real occupancy below what shared memory allows. + */ +constexpr std::size_t sharedBudgetPerSMBytes() { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1200) + return 228 * 1024; +#else + return 164 * 1024; +#endif +} + +/** + * @brief Max resident threads per SM, from the CUDA occupancy tables. + * + * Unlike shared memory, ptxas validates .minnctapersm against this limit and + * fails the build when minBlocks * blockSize exceeds it, so minBlocksPerSM + * must clamp by it. + */ +constexpr int maxThreadsPerSM() { +#if !defined(__CUDA_ARCH__) + return 2048; // Host pass; the value is never used. +#elif __CUDA_ARCH__ == 750 + return 1024; // Turing +#elif (__CUDA_ARCH__ >= 860 && __CUDA_ARCH__ < 900) || __CUDA_ARCH__ >= 1200 + return 1536; // Consumer Ampere/Ada, Orin, consumer Blackwell +#else + return 2048; // Volta, A100, Hopper, datacenter Blackwell +#endif +} + +/** + * @brief Second __launch_bounds__ argument: CTAs the kernel should fit per SM. + * + * @tparam SharedStateBytes Size of the kernel's dominant (assumed only) + * __shared__ allocation per CTA. + * @tparam BlockSize Threads per CTA. + * @tparam TargetResidentCTAs The residency the kernel is tuned at; asking for + * it also caps the compiler's register budget + * accordingly, so the two go together. The result is + * this value when shared memory and the SM thread + * limit allow it, otherwise whatever does fit + * (at least 1). + */ +template constexpr int minBlocksPerSM() { + constexpr std::size_t kBySharedMem = sharedBudgetPerSMBytes() / SharedStateBytes; + constexpr std::size_t kByThreads = static_cast(maxThreadsPerSM() / BlockSize); + constexpr std::size_t kResident = kBySharedMem < kByThreads ? kBySharedMem : kByThreads; + return kResident >= static_cast(TargetResidentCTAs) ? TargetResidentCTAs : + (kResident < 1 ? 1 : static_cast(kResident)); +} + +} // namespace nvMolKit + +#endif // NVMOLKIT_SUBGRAPH_OCCUPANCY_CUH diff --git a/src/subgraph/target_mask.cuh b/src/subgraph/target_mask.cuh new file mode 100644 index 00000000..1b794267 --- /dev/null +++ b/src/subgraph/target_mask.cuh @@ -0,0 +1,138 @@ +// 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 NVMOLKIT_SUBGRAPH_TARGET_MASK_CUH +#define NVMOLKIT_SUBGRAPH_TARGET_MASK_CUH + +#include +#include + +namespace nvMolKit { + +/** + * @brief Set of target atoms, bit t meaning target atom t. + * + * Specialised on the atom capacity rather than looped over an array so that each + * form is exactly as wide as it needs to be -- a bare uint32_t at 32 atoms, a + * uint64_t at 64 -- and the 128-atom form never dynamically indexes a + * register-resident array, which would spill it to local memory. The 32-atom + * form matters because these masks are the DFS inner loop: on 32-bit ALUs every + * 64-bit operation is two instructions, and popcount/ffs are two-instruction + * emulations, so the narrow form roughly halves the work per mask operation and + * halves the local-memory traffic of the per-depth candidate stack. + * + * setWord32(k, v) writes the k-th 32-bit lane group, which is how a label + * matrix transpose assembles a mask out of per-root ballots. + */ +template struct TargetMask; + +template <> struct TargetMask<32> { + uint32_t bits; + + __device__ __forceinline__ void clear() { bits = 0; } + __device__ __forceinline__ bool empty() const { return bits == 0; } + __device__ __forceinline__ void set(int bit) { bits |= 1u << bit; } + __device__ __forceinline__ void reset(int bit) { bits &= ~(1u << bit); } + /// Branch-free conditional set: @p value must be 0 or 1. + __device__ __forceinline__ void setIf(int bit, uint64_t value) { bits |= static_cast(value) << bit; } + __device__ __forceinline__ void setWord32(int, uint32_t value) { bits |= value; } + __device__ __forceinline__ int popcount() const { return __popc(bits); } + /// Lowest set atom index; undefined if empty. + __device__ __forceinline__ int lowest() const { return __ffs(static_cast(bits)) - 1; } + __device__ __forceinline__ void clearLowest() { bits &= bits - 1; } + __device__ __forceinline__ void andEq(const TargetMask& other) { bits &= other.bits; } + __device__ __forceinline__ void andNotEq(const TargetMask& other) { bits &= ~other.bits; } +}; + +template <> struct TargetMask<64> { + uint64_t lo; + + __device__ __forceinline__ void clear() { lo = 0; } + __device__ __forceinline__ bool empty() const { return lo == 0; } + __device__ __forceinline__ void set(int bit) { lo |= 1ULL << bit; } + __device__ __forceinline__ void reset(int bit) { lo &= ~(1ULL << bit); } + /// Branch-free conditional set: @p value must be 0 or 1. + __device__ __forceinline__ void setIf(int bit, uint64_t value) { lo |= value << bit; } + __device__ __forceinline__ void setWord32(int word, uint32_t value) { + lo |= static_cast(value) << (32 * word); + } + __device__ __forceinline__ int popcount() const { return __popcll(lo); } + /// Lowest set atom index; undefined if empty. + __device__ __forceinline__ int lowest() const { return __ffsll(static_cast(lo)) - 1; } + __device__ __forceinline__ void clearLowest() { lo &= lo - 1; } + __device__ __forceinline__ void andEq(const TargetMask& other) { lo &= other.lo; } + __device__ __forceinline__ void andNotEq(const TargetMask& other) { lo &= ~other.lo; } +}; + +template <> struct TargetMask<128> { + uint64_t lo; + uint64_t hi; + + __device__ __forceinline__ void clear() { lo = hi = 0; } + __device__ __forceinline__ bool empty() const { return (lo | hi) == 0; } + __device__ __forceinline__ void set(int bit) { + if (bit < 64) { + lo |= 1ULL << bit; + } else { + hi |= 1ULL << (bit - 64); + } + } + __device__ __forceinline__ void reset(int bit) { + if (bit < 64) { + lo &= ~(1ULL << bit); + } else { + hi &= ~(1ULL << (bit - 64)); + } + } + /// Branch-free conditional set: @p value must be 0 or 1. + __device__ __forceinline__ void setIf(int bit, uint64_t value) { + if (bit < 64) { + lo |= value << bit; + } else { + hi |= value << (bit - 64); + } + } + __device__ __forceinline__ void setWord32(int word, uint32_t value) { + if (word < 2) { + lo |= static_cast(value) << (32 * word); + } else { + hi |= static_cast(value) << (32 * (word - 2)); + } + } + __device__ __forceinline__ int popcount() const { return __popcll(lo) + __popcll(hi); } + /// Lowest set atom index; undefined if empty. + __device__ __forceinline__ int lowest() const { + return lo != 0 ? __ffsll(static_cast(lo)) - 1 : __ffsll(static_cast(hi)) + 63; + } + __device__ __forceinline__ void clearLowest() { + if (lo != 0) { + lo &= lo - 1; + } else { + hi &= hi - 1; + } + } + __device__ __forceinline__ void andEq(const TargetMask& other) { + lo &= other.lo; + hi &= other.hi; + } + __device__ __forceinline__ void andNotEq(const TargetMask& other) { + lo &= ~other.lo; + hi &= ~other.hi; + } +}; + +} // namespace nvMolKit + +#endif // NVMOLKIT_SUBGRAPH_TARGET_MASK_CUH diff --git a/src/subgraph/warp_dfs.cuh b/src/subgraph/warp_dfs.cuh new file mode 100644 index 00000000..82c9695d --- /dev/null +++ b/src/subgraph/warp_dfs.cuh @@ -0,0 +1,170 @@ +// 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. + +// Lane-local depth-first subgraph embedding search. +// +// Finds injective maps from query atoms to target atoms. Query atoms are +// matched in depth order: depth d means "choose the target atom for query atom +// d", and every query atom below d is already mapped. The caller owns what +// "query atom d" means -- the substructure backend matches query atoms in index +// order, while a seed-matching frontend (MCS) can search in any +// most-constrained-first permutation and apply the inverse permutation to the +// resulting mapping. +// +// All constraint knowledge lives behind the candidates oracle: at each descent +// the core asks it for the bitset of target atoms query atom `depth` may still +// map to given the mapping so far. The canonical oracle shape is one bitset +// intersection over target atoms: +// +// candidates(d) = labelCompatible(d) & ~used +// & AND over each back edge d--e (e < d) of +// { neighbours of mapping[e] reachable over a bond that +// edge's bond predicate accepts } +// +// Forward edges need no check; they are back edges of the deeper atom. How the +// per-edge neighbour sets are produced is the oracle's business: the +// substructure backend precomputes them per bond-mask class (see +// substruct_dfs.cuh), an MCS backend can walk a CSR row testing a +// bond-compatibility bit matrix. The core only ever sees the resulting masks. +// +// Each calling thread searches the subtrees rooted at the target atoms in its +// @p roots mask; those subtrees are disjoint, so all search state is +// lane-local registers and no synchronisation happens inside the search. The +// intended deployment is warp-per-pair with lane L owning roots L, L+32, ..., +// but the core is agnostic to how roots are distributed. +// +// An empty candidate bitset pops the stack; a non-empty one at the last depth +// means each of its bits completes an embedding, and the terminal handler +// decides what happens next (count, store, paint, stop -- see +// DfsTerminalVerdict). The abort hook is polled between roots so a frontend +// that only needs existence can stop every lane once one lane has found an +// embedding (e.g. by polling a warp-shared flag; a finer-grained frontend can +// additionally poll inside its oracle, which runs on every descent). + +#ifndef NVMOLKIT_SUBGRAPH_WARP_DFS_CUH +#define NVMOLKIT_SUBGRAPH_WARP_DFS_CUH + +#include "src/subgraph/target_mask.cuh" + +namespace nvMolKit { + +/// What the terminal handler tells the search to do after a terminal-depth +/// visit. +struct DfsTerminalVerdict { + bool rootDone; ///< Abandon the rest of the current root's subtree. + bool laneDone; ///< Stop this thread's search entirely; no further roots. +}; + +/** + * @brief Depth-first search over injective query->target atom maps, run + * independently by each calling thread over its @p roots. + * + * @tparam MaxDepth Capacity of the per-thread stack: the maximum number of + * query atoms (depths). Target atom indices are stored as + * bytes, so target capacity is bounded by the Mask width + * (<= 128 everywhere in nvMolKit). + * @tparam Mask Target-atom bitset, TargetMask<32/64/128> or anything + * matching its interface. + * @tparam CandidatesFn Mask candidatesAt(int depth, const unsigned char* mapping, + * const Mask& used, int prevTargetAtom). Must already + * exclude atoms in @p used. @p mapping[0..depth-1] are the + * committed target atoms; @p prevTargetAtom == mapping[depth-1], + * passed separately so chain-shaped oracles skip the load. + * @tparam TerminalFn DfsTerminalVerdict onTerminal(Mask terminals, + * const unsigned char* mapping). Called at @p lastDepth + * with the candidate set for the final query atom; every + * bit of @p terminals completes a distinct embedding with + * @p mapping[0..lastDepth-1]. mapping[0] is the root. + * @tparam AbortFn bool abortRequested(), polled before each root. + * + * @param roots Target atoms to anchor query atom 0 on, already filtered to + * label-compatible ones. + * @param lastDepth Index of the final query atom; must be >= 1 and < MaxDepth. + * Single-atom queries (lastDepth == 0) have no bond + * constraints and should be handled by the caller directly + * from its root mask. + */ +template +__device__ __forceinline__ void dfsFromRoots(Mask roots, + const int lastDepth, + CandidatesFn&& candidatesAt, + TerminalFn&& onTerminal, + AbortFn&& abortRequested) { + // The lane-local stack: mapping[d] is the target atom assigned to query atom + // d, remaining[d] the candidates at depth d not yet tried, and used the + // target atoms taken by depths 0..depth-1, which keeps the mapping injective. + Mask remaining[MaxDepth]; + unsigned char mapping[MaxDepth]; + + while (!roots.empty()) { + if (abortRequested()) { + return; + } + // Restart with query atom 0 anchored on the thread's next root. + const int rootAtom = roots.lowest(); + roots.clearLowest(); + + Mask used; + used.clear(); + used.set(rootAtom); + mapping[0] = static_cast(rootAtom); + + int depth = 1; + remaining[1] = candidatesAt(1, mapping, used, rootAtom); + bool laneDone = false; + + while (depth >= 1) { + if (depth == lastDepth) { + // Nothing deeper constrains the last query atom, so every remaining + // candidate completes a distinct embedding. The handler consumes the + // whole bitset; fall through to the pop below with this depth emptied. + const DfsTerminalVerdict verdict = onTerminal(remaining[depth], mapping); + remaining[depth].clear(); + if (verdict.rootDone || verdict.laneDone) { + laneDone = verdict.laneDone; + break; + } + } + + if (remaining[depth].empty()) { + // Exhausted: pop, releasing the atom the depth above had taken. Depth + // 0 is the root, which the outer loop advances, so the stack bottoms + // out at depth 1. + --depth; + if (depth >= 1) { + used.reset(mapping[depth]); + } + continue; + } + + // Otherwise commit the lowest untried candidate and descend, building + // the next depth's candidate set under the extended mapping. + const int candidate = remaining[depth].lowest(); + remaining[depth].clearLowest(); + mapping[depth] = static_cast(candidate); + used.set(candidate); + ++depth; + remaining[depth] = candidatesAt(depth, mapping, used, candidate); + } + + if (laneDone) { + return; + } + } +} + +} // namespace nvMolKit + +#endif // NVMOLKIT_SUBGRAPH_WARP_DFS_CUH diff --git a/src/subgraph/warp_reduce.cuh b/src/subgraph/warp_reduce.cuh new file mode 100644 index 00000000..a4f42fb2 --- /dev/null +++ b/src/subgraph/warp_reduce.cuh @@ -0,0 +1,88 @@ +// 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. + +// Full-warp reductions returning the aggregate to every lane. +// +// __reduce_*_sync requires sm_80. nvMolKit also builds for sm_70/sm_75, which +// fall back to a butterfly shuffle producing the same values. +// +// cub::WarpReduce is deliberately not used here. Its sm_80+ fast path lowers to +// these same __reduce_*_sync intrinsics, so there is nothing to gain, and it is +// worse in two ways: its contract returns the aggregate only to lane 0, while +// the min/max/or call sites need the all-lane broadcast that redux.sync +// provides for free (CUB would need an extra __shfl_sync); and CUB has no +// redux-backed OR reduction at all, so warpReduceOr would regress to a +// five-step shuffle loop even on sm_80+. + +#ifndef NVMOLKIT_SUBGRAPH_WARP_REDUCE_CUH +#define NVMOLKIT_SUBGRAPH_WARP_REDUCE_CUH + +#include + +namespace nvMolKit { + +constexpr uint32_t kFullWarpMask = 0xFFFFFFFFu; + +__device__ __forceinline__ uint32_t warpReduceAdd(uint32_t value) { +#if __CUDA_ARCH__ >= 800 + return __reduce_add_sync(kFullWarpMask, value); +#else +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_xor_sync(kFullWarpMask, value, offset); + } + return value; +#endif +} + +__device__ __forceinline__ uint32_t warpReduceOr(uint32_t value) { +#if __CUDA_ARCH__ >= 800 + return __reduce_or_sync(kFullWarpMask, value); +#else +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value |= __shfl_xor_sync(kFullWarpMask, value, offset); + } + return value; +#endif +} + +__device__ __forceinline__ uint32_t warpReduceMin(uint32_t value) { +#if __CUDA_ARCH__ >= 800 + return __reduce_min_sync(kFullWarpMask, value); +#else +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value = min(value, __shfl_xor_sync(kFullWarpMask, value, offset)); + } + return value; +#endif +} + +__device__ __forceinline__ uint32_t warpReduceMax(uint32_t value) { +#if __CUDA_ARCH__ >= 800 + return __reduce_max_sync(kFullWarpMask, value); +#else +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value = max(value, __shfl_xor_sync(kFullWarpMask, value, offset)); + } + return value; +#endif +} + +} // namespace nvMolKit + +#endif // NVMOLKIT_SUBGRAPH_WARP_REDUCE_CUH diff --git a/src/substruct/substruct_dfs.cuh b/src/substruct/substruct_dfs.cuh new file mode 100644 index 00000000..0760d686 --- /dev/null +++ b/src/substruct/substruct_dfs.cuh @@ -0,0 +1,755 @@ +// 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. + +// Warp-per-pair depth-first substructure search. +// +// Given a pair's label matrix -- bit (t, q) set iff target atom t satisfies +// query atom q's atom-level predicate, computed by populateLabelMatrix -- find +// the injective maps from query atoms to target atoms that also satisfy the +// query's bond constraints, which the label matrix does not cover. +// +// The search itself is nvMolKit::dfsFromRoots (src/subgraph/warp_dfs.cuh); +// this header is the substructure frontend: it builds the per-pair lookup +// tables the candidates oracle reads, supplies that oracle, and implements the +// output modes as terminal handlers. Query atoms are matched in index order, +// so depth d means "choose the target atom for query atom d". +// +// One warp per pair, kWarpsPerBlock pairs per CTA. Lane L searches the subtrees +// rooted at target atoms L, L+32, L+64, L+96; those subtrees are disjoint, so +// all search state is lane-local registers and lanes only meet at the final +// reduction. Shared memory holds the per-warp lookup tables and result counters. +// +// Three tables are built per pair before the search, each depending on the last: +// tryBuildQueryAtomCandidates (the label-compatibility term of the oracle's +// intersection), analyzeQueryBonds (the back edges), buildTargetAdjacency (the +// per-edge neighbour sets). +// +// Future optimization candidates: +// - Store target adjacency as the two packed degree-8 words at construction, +// deleting packAdjacencyRow. +// - Produce the label matrix query-major for this backend, deleting the ballot +// transpose in tryBuildQueryAtomCandidates. +// - Have the label kernel flag pairs with an empty column and compact the DFS +// pair list on device. +// - Precompute analyzeQueryBonds per query at host pack time (see its note). +// - Reorder query atoms most-constrained-first at construction, with an inverse +// permutation on mapping download. +// - Route size-1 queries to a popcount-over-label-column kernel. +// - Drop query>target pairs in the minibatch planner. + +#ifndef NVMOLKIT_SUBSTRUCT_DFS_CUH +#define NVMOLKIT_SUBSTRUCT_DFS_CUH + +#include +#include + +#include "src/subgraph/occupancy.cuh" +#include "src/subgraph/target_mask.cuh" +#include "src/subgraph/warp_dfs.cuh" +#include "src/subgraph/warp_reduce.cuh" +#include "src/substruct/molecules_device.cuh" +#include "src/substruct/packed_bonds_device.cuh" +#include "src/substruct/substruct_algos.cuh" + +namespace nvMolKit { +namespace dfs { + +/// Warps, and therefore independent pairs, per CTA. +constexpr int kWarpsPerBlock = 8; +constexpr int kBlockSize = kWarpsPerBlock * 32; + +/// How a warp records the embeddings it finds. +enum class DfsOutputMode { + Count, ///< Per-pair embedding count only; no mapping storage, no per-match atomics. + Store, ///< Full mappings written to the match index buffer. + Paint ///< Recursive SMARTS bit painted for mapping[0]; stops each root at its first embedding. +}; + +// ============================================================================= +// Query bond classification +// ============================================================================= + +// Entry layout in WarpSharedState::backEdges. A back edge of query atom d is an +// edge to a query atom with a smaller index, i.e. one already mapped at depth d. +constexpr unsigned char kBackEdgeAtomMask = 63u; ///< Low 6 bits: the earlier query atom's index. +constexpr unsigned char kBackEdgeAltMaskFlag = 64u; ///< Bit 6 (Dual only): use the Alt adjacency table. + +/** + * @brief How many distinct bond masks (QueryAtomBondsT::matchMask) the query's + * back edges carry. + * + * With at most two, the target neighbour sets those masks induce can be + * precomputed once per pair, turning the per-edge work in the inner loop into a + * table lookup. + */ +enum class BondMaskCase { + Uniform, ///< One mask. One adjacency table suffices. + Dual, ///< Two masks. Two adjacency tables; each back edge flags which it uses. + General ///< Three or more. Neighbour sets recomputed per edge during the search. +}; + +/// Classification of the query's bond constraints; identical in every lane. +struct QueryBondPlan { + BondMaskCase maskCase = BondMaskCase::General; + uint32_t minBondMask = 0; ///< Smallest bond mask on any back edge. + uint32_t maxBondMask = 0; ///< Largest bond mask on any back edge; equals minBondMask iff Uniform. + /// Bitset over query atoms: bit d set iff depth d's only back edge goes to + /// depth d-1 and resolves through targetAdjacency. Such a depth needs no + /// back-edge lookup, since the caller just chose depth d-1's target atom. + uint64_t chainDepths = 0; + /// As chainDepths, for depths resolving through targetAdjacencyAlt. + uint64_t chainDepthsAlt = 0; +}; + +// ============================================================================= +// Per-warp shared state +// ============================================================================= + +/** + * @brief Per-warp lookup tables and counters, sized for a whole CTA. + * + * Warp w touches only [w][...], so no synchronisation beyond __syncwarp() is + * needed. This is the kernel's only __shared__ allocation, so its size alone + * determines how many CTAs fit per SM (minBlocksPerSM). + * + * The adjacency tables hold different things per BondMaskCase; buildTargetAdjacency + * is the only place that fills them: + * Uniform: targetAdjacency[a] = atoms bonded to a over a bond the query's + * single mask accepts. Alt unused. + * Dual: targetAdjacency[a] for minBondMask, targetAdjacencyAlt[a] for + * maxBondMask; each back edge picks one via kBackEdgeAltMaskFlag. + * General: nothing is precomputable, so the tables cache the packed adjacency + * row instead -- targetAdjacency[a].lo the neighbour-index bytes, + * targetAdjacencyAlt[a].lo the bond-info bytes. + */ +template struct WarpSharedState { + static_assert(MaxQueryAtoms <= 64, "Back-edge indices pack into 6 bits and chain depths into a 64-bit mask"); + + /// Roots each lane owns: lane, lane + 32, ... below MaxTargetAtoms. + static constexpr int kRoots = static_cast(MaxTargetAtoms / 32); + using Mask = TargetMask; + + /** + * @brief One adjacency table slot, read as whichever member the case stores. + * + * A union because the General case caches packed adjacency rows, which are + * always 64 bits wide, whereas a mask is only as wide as MaxTargetAtoms. This + * keeps the table exactly as large as the wider of the two, so narrowing Mask + * costs no shared memory here while still narrowing it in registers. + */ + union AdjacencyEntry { + Mask mask; + uint64_t packedRow; + }; + + AdjacencyEntry targetAdjacency[kWarpsPerBlock][MaxTargetAtoms]; + AdjacencyEntry targetAdjacencyAlt[kWarpsPerBlock][MaxTargetAtoms]; + /// Transposed label matrix: the target atoms compatible with each query atom. + Mask queryAtomCandidates[kWarpsPerBlock][MaxQueryAtoms]; + /// Per query atom, its back edges, kBackEdge*-encoded. + unsigned char backEdges[kWarpsPerBlock][MaxQueryAtoms][kMaxBondsPerAtom]; + unsigned char backEdgeCounts[kWarpsPerBlock][MaxQueryAtoms]; ///< Valid entries in backEdges[w][q]. + int matchCount[kWarpsPerBlock]; ///< Embeddings found for the warp's pair. + int reportedCount[kWarpsPerBlock]; ///< Of those, how many fit in the output buffer. +}; + +/// Second __launch_bounds__ argument for the DFS kernels: eight resident CTAs +/// is the occupancy the search is tuned at, clamped to what WarpSharedState -- +/// the kernels' only shared allocation -- and the SM thread limit allow (see +/// src/subgraph/occupancy.cuh). The 128-atom target specialisations cannot +/// reach eight at any register count -- their adjacency tables alone are too +/// large -- so they ask for what fits. +template constexpr int minBlocksPerSM() { + return nvMolKit::minBlocksPerSM), kBlockSize, 8>(); +} + +// ============================================================================= +// Packed adjacency helpers +// ============================================================================= + +/** + * @brief Pack one target atom's adjacency into two words, one byte per bond slot. + * + * Byte j of @p neighbors is the j-th bonded atom's index, byte j of @p bondInfo + * its packed (bondType, isInRing) descriptor. Consumers then walk the bonds by + * shifting two registers instead of re-reading the 20-byte TargetAtomBonds. + * Slots past the atom's degree hold kNoNeighbor, which ends the walk. + */ +__device__ __forceinline__ void packAdjacencyRow(const TargetAtomBonds& bonds, + uint64_t& neighbors, + uint64_t& bondInfo) { + neighbors = ~0ULL; + bondInfo = 0; + const int deg = bonds.degree; +#pragma unroll + for (int j = 0; j < kMaxBondsPerAtom; ++j) { + if (j >= deg) { + break; + } + const int shift = 8 * j; + neighbors &= ~(0xFFULL << shift); + neighbors |= static_cast(bonds.neighborIdx[j]) << shift; + bondInfo |= static_cast(bonds.bondInfo[j]) << shift; + } +} + +/** + * @brief Neighbours of one target atom reachable over a bond @p queryBondMask accepts. + * + * @p neighbors and @p bondInfo are that atom's packed row from packAdjacencyRow. + * A query bond mask holds one bit per (isInRing, bondType) at position + * `isInRing * 16 + bondType`; see QueryAtomBondsT. + */ +template +__device__ __forceinline__ TargetMask neighborsMatchingBondMask(uint64_t neighbors, + uint64_t bondInfo, + uint32_t queryBondMask) { + TargetMask mask; + mask.clear(); +#pragma unroll 1 + for (int k = 0; k < kMaxBondsPerAtom; ++k) { + const uint32_t neighbor = static_cast(neighbors & 0xFFu); + if (neighbor == kNoNeighbor) { + break; + } + const uint32_t info = static_cast(bondInfo & 0xFFu); + const uint32_t code = ((info >> 4) & 1u) * 16u + (info & 15u); + mask.setIf(static_cast(neighbor), static_cast((queryBondMask >> code) & 1u)); + neighbors >>= 8; + bondInfo >>= 8; + } + return mask; +} + +// ============================================================================= +// Candidate construction +// ============================================================================= + +/** + * @brief The target atoms query atom @p depth may still map to: the candidates + * oracle handed to nvMolKit::dfsFromRoots. + * + * Evaluates the intersection described in src/subgraph/warp_dfs.cuh, in four + * shapes, cheapest first: + * - chain depth: the only back edge goes to depth-1, whose target atom the + * caller passes as @p prevTargetAtom, so the back-edge table is not read; + * - Uniform: every back edge intersects targetAdjacency; + * - Dual: each back edge's kBackEdgeAltMaskFlag picks its table; + * - General: recompute the neighbour set per edge from the cached packed row. + * + * Only the General case needs @p seen: it walks the query atom's raw bond slots, + * which can name the same neighbour twice, whereas analyzeQueryBonds already + * deduplicated the table. + */ +template +__device__ __forceinline__ TargetMask buildCandidates( + const WarpSharedState& shared, + int warp, + int depth, + const unsigned char* mapping, + const TargetMask& used, + const QueryMoleculeView& query, + const QueryBondPlan& plan, + int prevTargetAtom) { + using Mask = TargetMask; + + Mask candidates = shared.queryAtomCandidates[warp][depth]; + candidates.andNotEq(used); + + const uint64_t depthBit = 1ULL << depth; + if (plan.chainDepths & depthBit) { + candidates.andEq(shared.targetAdjacency[warp][prevTargetAtom].mask); + return candidates; + } + if (plan.chainDepthsAlt & depthBit) { + candidates.andEq(shared.targetAdjacencyAlt[warp][prevTargetAtom].mask); + return candidates; + } + + const int numBackEdges = shared.backEdgeCounts[warp][depth]; + + if (plan.maskCase == BondMaskCase::Uniform) { + for (int k = 0; k < numBackEdges && !candidates.empty(); ++k) { + candidates.andEq( + shared.targetAdjacency[warp][mapping[shared.backEdges[warp][depth][k] & kBackEdgeAtomMask]].mask); + } + return candidates; + } + + if (plan.maskCase == BondMaskCase::Dual) { + for (int k = 0; k < numBackEdges && !candidates.empty(); ++k) { + const uint32_t edge = shared.backEdges[warp][depth][k]; + const int mapped = mapping[edge & kBackEdgeAtomMask]; + candidates.andEq((edge & kBackEdgeAltMaskFlag) ? shared.targetAdjacencyAlt[warp][mapped].mask : + shared.targetAdjacency[warp][mapped].mask); + } + return candidates; + } + + const QueryAtomBonds& queryBonds = query.getQueryBonds(depth); + const int degree = queryBonds.degree; + uint64_t seen = 0; + for (int slot = 0; slot < kMaxBondsPerAtom && !candidates.empty(); ++slot) { + if (slot >= degree) { + break; + } + const uint32_t neighborQueryAtom = queryBonds.neighborIdx[slot]; + if (neighborQueryAtom < static_cast(depth) && !((seen >> neighborQueryAtom) & 1ULL)) { + seen |= 1ULL << neighborQueryAtom; + const int mapped = mapping[neighborQueryAtom]; + candidates.andEq(neighborsMatchingBondMask(shared.targetAdjacency[warp][mapped].packedRow, + shared.targetAdjacencyAlt[warp][mapped].packedRow, + queryBonds.matchMask[slot])); + } + } + return candidates; +} + +// ============================================================================= +// Per-pair setup +// ============================================================================= + +/** + * @brief Target atom @p atom 's label matrix row, as a bitset over query atoms. + * + * The label matrix is a BitMatrix2DView: + * row-major with the target atom as the row, packed least significant bit first, + * so a row is a contiguous MaxQueryAtoms-bit field. + */ +template +__device__ __forceinline__ uint64_t readLabelRow(const uint32_t* labelWords, int atom) { + if constexpr (MaxQueryAtoms == 16) { + return static_cast((labelWords[atom >> 1] >> ((atom & 1) * 16)) & 0xFFFFu); + } else if constexpr (MaxQueryAtoms == 32) { + return static_cast(labelWords[atom]); + } else { + return reinterpret_cast(labelWords)[atom]; + } +} + +/** + * @brief Transpose the pair's label matrix into shared.queryAtomCandidates. + * + * The label matrix is target-major (row t, bit q) but the search needs it + * query-major: at depth d, the target atoms compatible with query atom d. Each + * lane holds the rows of the target atoms it owns, so the transpose is one + * __ballot_sync per (query atom, root group) -- balloting bit d yields 32 bits + * of column d in target atom order. + * + * @return false if some query atom has no compatible target atom, in which case + * the pair cannot match and the caller must stop. The table is fully + * written either way. + */ +template +__device__ __forceinline__ bool tryBuildQueryAtomCandidates(WarpSharedState& shared, + int warp, + int lane, + const uint32_t* labelWords, + int numTargetAtoms, + int numQueryAtoms) { + constexpr int kRoots = WarpSharedState::kRoots; + + uint64_t rows[kRoots]; +#pragma unroll + for (int k = 0; k < kRoots; ++k) { + const int atom = lane + 32 * k; + rows[k] = atom < numTargetAtoms ? readLabelRow(labelWords, atom) : 0; + } + + uint32_t anyEmptyColumn = 0; + for (int depth = 0; depth < numQueryAtoms; ++depth) { + uint32_t ballots[kRoots]; + uint32_t occupied = 0; +#pragma unroll + for (int k = 0; k < kRoots; ++k) { + ballots[k] = __ballot_sync(kFullWarpMask, static_cast((rows[k] >> depth) & 1ULL)); + occupied |= ballots[k]; + } + anyEmptyColumn |= static_cast(occupied == 0u); + + // Ballots are uniform across the warp, so one lane assembles and stores. + if (lane == 0) { + TargetMask mask; + mask.clear(); +#pragma unroll + for (int k = 0; k < kRoots; ++k) { + mask.setWord32(k, ballots[k]); + } + shared.queryAtomCandidates[warp][depth] = mask; + } + } + + return anyEmptyColumn == 0u; +} + +/** + * @brief Build the back-edge table and classify the query's bond masks. + * + * Three passes over the query atoms, striped across the warp: + * 1. Record each query atom's deduplicated back edges into + * shared.backEdges/backEdgeCounts, reducing the smallest and largest bond + * mask across the warp. Min and max are only a cheap fingerprint for the + * number of distinct masks: all masks equal implies min == max, and if + * exactly two values occur they are the min and the max. + * 2. If min != max, flag the max-mask edges and watch for a third value, which + * demotes the pair to General. + * 3. Uniform and Dual only: find the chain depths (see QueryBondPlan). + * + * This is a pure function of the query but runs once per pair; it could be + * computed once per query at host pack time and loaded instead. + */ +template +__device__ __forceinline__ QueryBondPlan analyzeQueryBonds(WarpSharedState& shared, + int warp, + int lane, + const QueryMoleculeView& query, + int numQueryAtoms) { + QueryBondPlan plan; + + uint32_t laneMin = 0xFFFFFFFFu; + uint32_t laneMax = 0u; + + for (int depth = lane; depth < numQueryAtoms; depth += 32) { + const QueryAtomBonds& bonds = query.getQueryBonds(depth); + const int degree = bonds.degree; + int count = 0; + uint64_t seen = 0; + for (int slot = 0; slot < kMaxBondsPerAtom; ++slot) { + if (slot >= degree) { + break; + } + const uint32_t neighbor = bonds.neighborIdx[slot]; + if (neighbor < static_cast(depth) && !((seen >> neighbor) & 1ULL)) { + seen |= 1ULL << neighbor; + const uint32_t mask = bonds.matchMask[slot]; + shared.backEdges[warp][depth][count] = static_cast(neighbor); + laneMin = min(laneMin, mask); + laneMax = max(laneMax, mask); + ++count; + } + } + shared.backEdgeCounts[warp][depth] = static_cast(count); + } + + plan.minBondMask = warpReduceMin(laneMin); + plan.maxBondMask = warpReduceMax(laneMax); + const bool uniform = (plan.minBondMask == plan.maxBondMask); + __syncwarp(); + + // The walk here must mirror the first pass exactly: it re-derives the same + // slot ordering in order to flag the entries it wrote. + uint32_t sawThirdMask = 0; + if (!uniform) { + for (int depth = lane; depth < numQueryAtoms; depth += 32) { + const QueryAtomBonds& bonds = query.getQueryBonds(depth); + const int degree = bonds.degree; + int count = 0; + uint64_t seen = 0; + for (int slot = 0; slot < kMaxBondsPerAtom; ++slot) { + if (slot >= degree) { + break; + } + const uint32_t neighbor = bonds.neighborIdx[slot]; + if (neighbor < static_cast(depth) && !((seen >> neighbor) & 1ULL)) { + seen |= 1ULL << neighbor; + const uint32_t mask = bonds.matchMask[slot]; + if (mask != plan.minBondMask && mask != plan.maxBondMask) { + sawThirdMask = 1; + } + if (mask == plan.maxBondMask) { + shared.backEdges[warp][depth][count] |= kBackEdgeAltMaskFlag; + } + ++count; + } + } + } + } + const bool dual = !uniform && (__ballot_sync(kFullWarpMask, sawThirdMask) == 0u); + plan.maskCase = uniform ? BondMaskCase::Uniform : (dual ? BondMaskCase::Dual : BondMaskCase::General); + __syncwarp(); + + if (plan.maskCase != BondMaskCase::General) { + uint64_t lanePrimary = 0; + uint64_t laneAlt = 0; + for (int depth = lane; depth < numQueryAtoms; depth += 32) { + if (depth >= 1 && shared.backEdgeCounts[warp][depth] == 1) { + const uint32_t edge = shared.backEdges[warp][depth][0]; + if ((edge & kBackEdgeAtomMask) == static_cast(depth - 1)) { + if (edge & kBackEdgeAltMaskFlag) { + laneAlt |= 1ULL << depth; + } else { + lanePrimary |= 1ULL << depth; + } + } + } + } + // Two 32-bit OR reductions per 64-bit lane mask; the reductions are 32-bit. + plan.chainDepths = static_cast(warpReduceOr(static_cast(lanePrimary))) | + (static_cast(warpReduceOr(static_cast(lanePrimary >> 32))) << 32); + plan.chainDepthsAlt = static_cast(warpReduceOr(static_cast(laneAlt))) | + (static_cast(warpReduceOr(static_cast(laneAlt >> 32))) << 32); + } + + return plan; +} + +/** + * @brief Fill the per-warp adjacency tables for the classified plan, target + * atoms striped across the warp. See WarpSharedState for the contents. + */ +template +__device__ __forceinline__ void buildTargetAdjacency(WarpSharedState& shared, + int warp, + int lane, + const TargetMoleculeView& target, + int numTargetAtoms, + const QueryBondPlan& plan) { + for (int atom = lane; atom < numTargetAtoms; atom += 32) { + uint64_t neighbors; + uint64_t bondInfo; + packAdjacencyRow(target.targetAtomBonds[atom], neighbors, bondInfo); + + if (plan.maskCase == BondMaskCase::Uniform) { + shared.targetAdjacency[warp][atom].mask = + neighborsMatchingBondMask(neighbors, bondInfo, plan.maxBondMask); + } else if (plan.maskCase == BondMaskCase::Dual) { + shared.targetAdjacency[warp][atom].mask = + neighborsMatchingBondMask(neighbors, bondInfo, plan.minBondMask); + shared.targetAdjacencyAlt[warp][atom].mask = + neighborsMatchingBondMask(neighbors, bondInfo, plan.maxBondMask); + } else { + shared.targetAdjacency[warp][atom].packedRow = neighbors; + shared.targetAdjacencyAlt[warp][atom].packedRow = bondInfo; + } + } +} + +// ============================================================================= +// Warp-per-pair DFS +// ============================================================================= + +/** + * @brief Output plumbing for a single pair's DFS. + */ +struct DfsPairOutput { + int* matchCounts = nullptr; ///< Per-pair total embedding count + int* reportedCounts = nullptr; ///< Per-pair embeddings actually written + int16_t* matchIndices = nullptr; ///< Mapping storage + int matchOffset = 0; ///< Start of this pair's mapping storage + int storageCapacity = 0; ///< Mappings storable for this pair + int maxMatchesToFind = -1; ///< Stop once this many are found (-1 = unlimited) + bool countOnly = false; ///< Suppress mapping writes + PaintModeParams paint = {}; ///< Paint destination (Paint mode only) +}; + +/** + * @brief One warp searches one target/query pair. + * + * Builds the pair's three lookup tables, then runs nvMolKit::dfsFromRoots over + * each lane's roots, with the output mode expressed as the terminal handler. + * + * The caller must have decoded the pair and pointed @p labelWords at that pair's + * label matrix. All 32 lanes of the warp must call this. + */ +template +__device__ void dfsSearchPair(const TargetMoleculeView& target, + const QueryMoleculeView& query, + const uint32_t* labelWords, + WarpSharedState& shared, + int warp, + int lane, + int resultIdx, + const DfsPairOutput& out) { + constexpr int kRoots = WarpSharedState::kRoots; + using Mask = TargetMask; + + const int numTargetAtoms = target.numAtoms; + const int numQueryAtoms = query.numAtoms; + + // Count and Store leave the per-pair counters unwritten on every early exit + // path unless we zero them here; the result buffers are not pre-zeroed. + auto writeEmptyResult = [&]() { + if constexpr (Mode != DfsOutputMode::Paint) { + if (lane == 0) { + out.matchCounts[resultIdx] = 0; + if (!out.countOnly && out.reportedCounts != nullptr) { + out.reportedCounts[resultIdx] = 0; + } + } + } + }; + + if (numQueryAtoms > numTargetAtoms) { + writeEmptyResult(); + return; + } + + if (lane == 0) { + shared.matchCount[warp] = 0; + shared.reportedCount[warp] = 0; + } + + if (!tryBuildQueryAtomCandidates(shared, warp, lane, labelWords, numTargetAtoms, numQueryAtoms)) { + writeEmptyResult(); + return; + } + + const int limit = out.maxMatchesToFind; + const bool hasLimit = limit >= 0; + const int lastDepth = numQueryAtoms - 1; + + // Writes one complete mapping. Returns true once the pair's limit is reached. + auto emitMapping = [&](const unsigned char* mapping, int terminalAtom) -> bool { + const int slot = atomicAdd(&shared.matchCount[warp], 1); + if (!out.countOnly && slot < out.storageCapacity) { + const int writeOffset = out.matchOffset + slot * numQueryAtoms; + for (int q = 0; q < numQueryAtoms - 1; ++q) { + out.matchIndices[writeOffset + q] = static_cast(mapping[q]); + } + out.matchIndices[writeOffset + numQueryAtoms - 1] = static_cast(terminalAtom); + atomicAdd(&shared.reportedCount[warp], 1); + } + return hasLimit && (slot + 1) >= limit; + }; + + // Mode-dependent result write, shared by the single-atom and general paths. + auto finishPair = [&](uint32_t laneTotal) { + if constexpr (Mode == DfsOutputMode::Count) { + uint32_t pairTotal = warpReduceAdd(laneTotal); + if (hasLimit && pairTotal > static_cast(limit)) { + pairTotal = static_cast(limit); + } + if (lane == 0) { + out.matchCounts[resultIdx] = static_cast(pairTotal); + if (!out.countOnly && out.reportedCounts != nullptr) { + out.reportedCounts[resultIdx] = 0; + } + } + } else if constexpr (Mode == DfsOutputMode::Store) { + // Store accumulated into the shared counters as it went, so this is just a + // barrier before reading them. + __syncwarp(); + if (lane == 0) { + out.matchCounts[resultIdx] = shared.matchCount[warp]; + if (!out.countOnly && out.reportedCounts != nullptr) { + out.reportedCounts[resultIdx] = shared.reportedCount[warp]; + } + } + } + }; + + // Builds the roots this lane owns -- target atoms lane, lane+32, lane+64, + // lane+96 -- restricted to the ones label-compatible with query atom 0. + auto buildLaneRoots = [&]() -> Mask { + Mask rootMask; + rootMask.clear(); +#pragma unroll + for (int k = 0; k < kRoots; ++k) { + rootMask.set(lane + 32 * k); + } + Mask roots = shared.queryAtomCandidates[warp][0]; + roots.andEq(rootMask); + return roots; + }; + + if (lastDepth == 0) { + // Single-atom query: no bonds to satisfy, so every candidate root is already + // a complete embedding. Exits before the bond analysis and adjacency build, + // none of whose output it would read. (dfsFromRoots requires lastDepth >= 1, + // so single-atom queries are the frontend's job by contract.) + __syncwarp(); // The transpose's lane-0 shared writes. + Mask roots = buildLaneRoots(); + uint32_t laneTotal = 0; + if constexpr (Mode == DfsOutputMode::Count) { + laneTotal = static_cast(roots.popcount()); + } else if constexpr (Mode == DfsOutputMode::Paint) { + while (!roots.empty()) { + const int root = roots.lowest(); + roots.clearLowest(); + atomicOr(&out.paint.recursiveBits[out.paint.outputPairIdx * out.paint.maxTargetAtoms + root], + 1u << out.paint.patternId); + } + } else { + unsigned char mapping[1]; + while (!roots.empty()) { + const int root = roots.lowest(); + roots.clearLowest(); + if (emitMapping(mapping, root)) { + break; + } + } + } + finishPair(laneTotal); + return; + } + + const QueryBondPlan plan = analyzeQueryBonds(shared, warp, lane, query, numQueryAtoms); + + buildTargetAdjacency(shared, warp, lane, target, numTargetAtoms, plan); + __syncwarp(); + + uint32_t laneTotal = 0; + + auto candidatesAt = [&](int depth, const unsigned char* mapping, const Mask& used, int prevTargetAtom) -> Mask { + return buildCandidates(shared, warp, depth, mapping, used, query, plan, prevTargetAtom); + }; + + // The output mode as a terminal handler; every bit of @p terminals completes + // a distinct embedding with mapping[0..lastDepth-1]. + auto onTerminal = [&](Mask terminals, const unsigned char* mapping) -> DfsTerminalVerdict { + DfsTerminalVerdict verdict{false, false}; + if constexpr (Mode == DfsOutputMode::Count) { + laneTotal += static_cast(terminals.popcount()); + // A lane that has already reached the pair limit on its own roots + // cannot change the clamped answer, so it stops. Lanes cannot see + // each other's partial totals mid-search without a warp collective + // in divergent code, so the other lanes run to exhaustion. + if (hasLimit && laneTotal >= static_cast(limit)) { + verdict.laneDone = true; + } + } else if constexpr (Mode == DfsOutputMode::Paint) { + // Paint only asks whether the root can start an embedding, so the + // first hit settles the root and skips the rest of its subtree. + if (!terminals.empty()) { + atomicOr(&out.paint.recursiveBits[out.paint.outputPairIdx * out.paint.maxTargetAtoms + mapping[0]], + 1u << out.paint.patternId); + verdict.rootDone = true; + } + } else { + while (!terminals.empty()) { + const int terminalAtom = terminals.lowest(); + terminals.clearLowest(); + if (emitMapping(mapping, terminalAtom)) { + verdict.laneDone = true; + break; + } + } + } + return verdict; + }; + + dfsFromRoots(MaxQueryAtoms)>(buildLaneRoots(), lastDepth, candidatesAt, onTerminal, [] { + return false; + }); + + finishPair(laneTotal); +} + +} // namespace dfs +} // namespace nvMolKit + +#endif // NVMOLKIT_SUBSTRUCT_DFS_CUH diff --git a/src/substruct/substruct_kernels.cu b/src/substruct/substruct_kernels.cu index 61ffbe3c..c334784d 100644 --- a/src/substruct/substruct_kernels.cu +++ b/src/substruct/substruct_kernels.cu @@ -19,6 +19,7 @@ #include "src/substruct/sm_shared_mem_config.cuh" #include "src/substruct/substruct_algos.cuh" #include "src/substruct/substruct_debug.h" +#include "src/substruct/substruct_dfs.cuh" #include "src/substruct/substruct_kernels.h" #include "src/substruct/substruct_launch_config.h" #include "src/substruct/substruct_search_internal.h" @@ -57,6 +58,11 @@ template struct SubstructMatchResul return overflowBuffer + (blockIdx.x * overflowBuffersPerBlock + bufferIdx) * overflowEntriesPerBuffer; } + __device__ __forceinline__ PartialMatchT* getPairOverflowBuffer(int miniBatchIdx, + int bufferIdx = 0) const { + return overflowBuffer + (miniBatchIdx * overflowBuffersPerBlock + bufferIdx) * overflowEntriesPerBuffer; + } + __device__ __forceinline__ int getOverflowCapacity() const { return overflowEntriesPerBuffer; } __device__ __forceinline__ uint32_t getRecursiveMatchBits(int miniBatchIdx, int atomIdx) const { @@ -313,7 +319,7 @@ __global__ void labelMatrixPaintKernelT(TargetMoleculesDeviceView targets, * @tparam MaxTargetAtoms Maximum target atoms for label matrix sizing * @tparam MaxQueryAtoms Maximum query atoms for label matrix and partial match sizing * @tparam MaxBondsPerAtom Maximum bonds per atom for edge consistency loop unrolling - * @tparam Algo Algorithm to use (VF2 or GSI) + * @tparam Algo Algorithm to use (VF2, GSI, or DFS) */ template __global__ void substructMatchKernelT(TargetMoleculesDeviceView targets, @@ -495,7 +501,7 @@ __global__ void substructMatchKernelT(TargetMoleculesDeviceView * @tparam MaxTargetAtoms Maximum target atoms for label matrix sizing * @tparam MaxQueryAtoms Maximum query atoms for label matrix and partial match sizing * @tparam MaxBondsPerAtom Maximum bonds per atom for edge consistency loop unrolling - * @tparam Algo Algorithm to use (VF2 or GSI) + * @tparam Algo Algorithm to use (GSI or DFS) */ template __global__ void substructPaintKernelT(TargetMoleculesDeviceView targets, @@ -624,11 +630,153 @@ __global__ void substructPaintKernelT(TargetMoleculesDeviceView targets, } } +// ============================================================================= +// Dedicated DFS Kernels +// ============================================================================= +// +// DFS does not share the one-block-per-pair topology of the kernels above: it +// runs one warp per (target, query) pair, kWarpsPerBlock pairs per CTA, with all +// search state lane-local. See substruct_dfs.cuh for the algorithm. Everything a +// pair needs is derived inside its owning warp, so no per-pair global scratch +// exists and concurrent recursive depth groups on separate streams cannot +// collide. + +/** + * @brief Warp-per-pair DFS matching kernel. + * + * @tparam MaxTargetAtoms Maximum target atoms for label matrix and mask sizing + * @tparam MaxQueryAtoms Maximum query atoms for label matrix and mask sizing + * @tparam Mode Count (count-only) or Store (full mappings) + */ +template +__global__ +__launch_bounds__(dfs::kBlockSize, dfs::minBlocksPerSM()) void substructDfsMatchKernelT( + TargetMoleculesDeviceView targets, + QueryMoleculesDeviceView queries, + SubstructMatchResultsDeviceViewT results, + const int* pairIndices, + int numPairs, + int numQueries, + const int* batchLocalIndices) { + __shared__ dfs::WarpSharedState shared; + + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & (kWarpSize - 1); + const int launchIdx = blockIdx.x * dfs::kWarpsPerBlock + warp; + if (launchIdx >= numPairs) { + return; + } + + const int miniBatchIdx = batchLocalIndices ? batchLocalIndices[launchIdx] : launchIdx; + const int globalPairIdx = pairIndices[launchIdx]; + const int targetIdx = globalPairIdx / numQueries; + const int queryIdx = globalPairIdx % numQueries; + + if (targetIdx >= targets.numMolecules || queryIdx >= queries.numMolecules) { + return; + } + + const TargetMoleculeView target = getMolecule(targets, targetIdx); + const QueryMoleculeView query = getMolecule(queries, queryIdx); + + const bool countOnly = results.countOnly; + const int matchOffset = countOnly ? 0 : results.pairMatchStarts[miniBatchIdx]; + + dfs::DfsPairOutput out; + out.matchCounts = results.matchCounts; + out.reportedCounts = results.reportedCounts; + out.matchIndices = results.matchIndices; + out.matchOffset = matchOffset; + out.storageCapacity = countOnly ? 0 : (results.pairMatchStarts[miniBatchIdx + 1] - matchOffset) / query.numAtoms; + out.maxMatchesToFind = results.maxMatchesToFind; + out.countOnly = countOnly; + + dfs::dfsSearchPair(target, + query, + results.getLabelMatrixPtr(miniBatchIdx), + shared, + warp, + lane, + miniBatchIdx, + out); +} + +/** + * @brief Warp-per-pair DFS paint kernel for recursive SMARTS preprocessing. + * + * Pair indexing matches the one-block-per-pair label kernel that produced the + * label matrices: pair index = localTargetIdx * numPatterns + localPatternIdx. + */ +template +__global__ +__launch_bounds__(dfs::kBlockSize, dfs::minBlocksPerSM()) void substructDfsPaintKernelT( + TargetMoleculesDeviceView targets, + QueryMoleculesDeviceView patterns, + const BatchedPatternEntry* patternEntries, + int numPatterns, + int numPairs, + uint32_t* outputRecursiveBits, + int maxTargetAtoms, + int outputNumQueries, + int defaultPatternId, + int defaultMainQueryIdx, + int miniBatchPairOffset, + int miniBatchSize, + const uint32_t* labelMatrixBuffer, + int firstTargetIdx) { + constexpr std::size_t kLabelMatrixWordsT = MaxTargetAtoms * MaxQueryAtoms / 32; + + __shared__ dfs::WarpSharedState shared; + + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & (kWarpSize - 1); + const int pairIdx = blockIdx.x * dfs::kWarpsPerBlock + warp; + if (pairIdx >= numPairs) { + return; + } + + const int localTargetIdx = pairIdx / numPatterns; + const int targetIdx = firstTargetIdx + localTargetIdx; + const int localPatternIdx = pairIdx % numPatterns; + + if (targetIdx >= targets.numMolecules) { + return; + } + + const int mainQueryIdx = patternEntries ? patternEntries[localPatternIdx].mainQueryIdx : defaultMainQueryIdx; + const int patternId = patternEntries ? patternEntries[localPatternIdx].patternId : defaultPatternId; + const int patternMolIdx = patternEntries ? patternEntries[localPatternIdx].patternMolIdx : localPatternIdx; + + const int globalPairIdx = targetIdx * outputNumQueries + mainQueryIdx; + if (globalPairIdx < miniBatchPairOffset || globalPairIdx >= miniBatchPairOffset + miniBatchSize) { + return; + } + + const TargetMoleculeView target = getMolecule(targets, targetIdx); + const QueryMoleculeView pattern = getMolecule(patterns, patternMolIdx); + + dfs::DfsPairOutput out; + out.paint.recursiveBits = outputRecursiveBits; + out.paint.patternId = patternId; + out.paint.maxTargetAtoms = maxTargetAtoms; + out.paint.outputPairIdx = globalPairIdx - miniBatchPairOffset; + + dfs::dfsSearchPair( + target, + pattern, + labelMatrixBuffer + static_cast(pairIdx) * kLabelMatrixWordsT, + shared, + warp, + lane, + 0, + out); +} + // ============================================================================= // Explicit Template Instantiations for All 24 Valid Configurations // ============================================================================= -// Helper macro to instantiate both VF2 and GSI for a given configuration +// Helper macro to instantiate all supported algorithms for a given configuration #define INSTANTIATE_SUBSTRUCT_KERNELS(MaxT, MaxQ, MaxB) \ template __global__ void substructMatchKernelT( \ TargetMoleculesDeviceView, \ @@ -706,6 +854,51 @@ INSTANTIATE_SUBSTRUCT_KERNELS(128, 64, 8) #undef INSTANTIATE_SUBSTRUCT_KERNELS +// DFS kernels do not depend on MaxBondsPerAtom: the degree-8 packed rows are +// consumed whole, so there is no per-degree loop to unroll. +#define INSTANTIATE_SUBSTRUCT_DFS_KERNELS(MaxT, MaxQ) \ + template __global__ void substructDfsMatchKernelT( \ + TargetMoleculesDeviceView, \ + QueryMoleculesDeviceView, \ + SubstructMatchResultsDeviceViewT, \ + const int*, \ + int, \ + int, \ + const int*); \ + template __global__ void substructDfsMatchKernelT( \ + TargetMoleculesDeviceView, \ + QueryMoleculesDeviceView, \ + SubstructMatchResultsDeviceViewT, \ + const int*, \ + int, \ + int, \ + const int*); \ + template __global__ void substructDfsPaintKernelT(TargetMoleculesDeviceView, \ + QueryMoleculesDeviceView, \ + const BatchedPatternEntry*, \ + int, \ + int, \ + uint32_t*, \ + int, \ + int, \ + int, \ + int, \ + int, \ + int, \ + const uint32_t*, \ + int); + +INSTANTIATE_SUBSTRUCT_DFS_KERNELS(32, 16) +INSTANTIATE_SUBSTRUCT_DFS_KERNELS(32, 32) +INSTANTIATE_SUBSTRUCT_DFS_KERNELS(64, 16) +INSTANTIATE_SUBSTRUCT_DFS_KERNELS(64, 32) +INSTANTIATE_SUBSTRUCT_DFS_KERNELS(64, 64) +INSTANTIATE_SUBSTRUCT_DFS_KERNELS(128, 16) +INSTANTIATE_SUBSTRUCT_DFS_KERNELS(128, 32) +INSTANTIATE_SUBSTRUCT_DFS_KERNELS(128, 64) + +#undef INSTANTIATE_SUBSTRUCT_DFS_KERNELS + // Label matrix kernel instantiations (one per target/query combo, no MaxBonds needed) #define INSTANTIATE_LABEL_MATRIX_KERNEL(MaxT, MaxQ) \ template __global__ void labelMatrixKernelT(TargetMoleculesDeviceView, \ @@ -971,6 +1164,27 @@ void launchSubstructPaintKernelForConfig(SubstructAlgorithm algorithm, switch (algorithm) { case SubstructAlgorithm::VF2: break; + case SubstructAlgorithm::DFS: { + // numBlocks is the pair count from the caller's one-block-per-pair view; + // DFS packs eight pairs into each CTA instead. + const int dfsBlocks = (numBlocks + dfs::kWarpsPerBlock - 1) / dfs::kWarpsPerBlock; + substructDfsPaintKernelT + <<>>(targets, + patterns, + patternEntries, + numPatterns, + numBlocks, + outputRecursiveBits, + maxTargetAtoms, + outputNumQueries, + defaultPatternId, + defaultMainQueryIdx, + miniBatchPairOffset, + miniBatchSize, + labelMatrixBuffer, + firstTargetIdx); + break; + } case SubstructAlgorithm::GSI: substructPaintKernelT <<>>(targets, @@ -1046,6 +1260,9 @@ void configureSubstructKernelsSharedMem() { substructMatchKernelT); configureSharedMemCarveout( substructPaintKernelT); + configureSharedMemCarveout(substructDfsMatchKernelT); + configureSharedMemCarveout(substructDfsMatchKernelT); + configureSharedMemCarveout(substructDfsPaintKernelT); sharedMemCarveoutConfigured() = true; } @@ -1095,6 +1312,29 @@ void launchMatchKernelForConfig(SubstructAlgorithm algorithm, batchLocalIndices, timings); break; + case SubstructAlgorithm::DFS: { + const int dfsBlocks = (numPairs + dfs::kWarpsPerBlock - 1) / dfs::kWarpsPerBlock; + if (miniBatchResults.countOnly()) { + substructDfsMatchKernelT + <<>>(targets, + queries, + results, + pairIndices, + numPairs, + numQueries, + batchLocalIndices); + } else { + substructDfsMatchKernelT + <<>>(targets, + queries, + results, + pairIndices, + numPairs, + numQueries, + batchLocalIndices); + } + break; + } case SubstructAlgorithm::GSI: substructMatchKernelT <<>>(targets, diff --git a/src/substruct/substruct_results.h b/src/substruct/substruct_results.h index c4f85fe2..6793e402 100644 --- a/src/substruct/substruct_results.h +++ b/src/substruct/substruct_results.h @@ -27,7 +27,8 @@ namespace nvMolKit { */ enum class SubstructAlgorithm { VF2, ///< VF2 iterative stack-based DFS - GSI ///< GSI-inspired BFS level-by-level join + GSI, ///< GSI-inspired BFS level-by-level join + DFS ///< Adjacency-anchored DFS with injectivity bitsets }; /** @@ -52,9 +53,10 @@ struct SubstructSearchConfig { int workerThreads = -1; ///< GPU runner threads per GPU (-1 = autoselect) int preprocessingThreads = -1; ///< CPU threads for preprocessing and opportunistic RDKit fallback (-1 = autoselect) int executorsPerRunner = -1; ///< GPU executors per runner thread (-1 = auto: 3 for single runner, 2 otherwise) - std::vector gpuIds; ///< GPU device IDs to use (empty = current device only) - int maxMatches = 0; ///< Max matches per pair (0 = unlimited, like RDKit) - bool uniquify = false; ///< Remove duplicate matches differing only in atom enumeration order + std::vector gpuIds; ///< GPU device IDs to use (empty = current device only) + int maxMatches = 0; ///< Max matches per pair (0 = unlimited, like RDKit) + bool uniquify = false; ///< Remove duplicate matches differing only in atom enumeration order + SubstructAlgorithm algorithm = SubstructAlgorithm::GSI; ///< Matching backend used by high-level APIs }; /** diff --git a/src/substruct/substruct_search.cu b/src/substruct/substruct_search.cu index 99a88f23..10b4b912 100644 --- a/src/substruct/substruct_search.cu +++ b/src/substruct/substruct_search.cu @@ -145,6 +145,8 @@ void uploadAndLaunchMiniBatch(GpuExecutor& executor, SubstructAlgorithm algorithm) { ScopedNvtxRange uploadRange("uploadAndLaunchMiniBatch"); + // DFS keeps all per-pair state lane-local, so only GSI needs the ping-pong + // partial-match overflow slab. cudaStream_t executorStream = executor.stream(); const int numBuffersPerBlock = (algorithm == SubstructAlgorithm::GSI) ? 2 : 1; diff --git a/src/testutils/substruct_validation.cu b/src/testutils/substruct_validation.cu index 769fb4f9..9bc2239a 100644 --- a/src/testutils/substruct_validation.cu +++ b/src/testutils/substruct_validation.cu @@ -69,6 +69,8 @@ std::string algorithmName(SubstructAlgorithm algo) { return "VF2"; case SubstructAlgorithm::GSI: return "GSI"; + case SubstructAlgorithm::DFS: + return "DFS"; } return "Unknown"; } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 08676161..240efa13 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -350,6 +350,9 @@ target_link_libraries( test_substruct_algos PRIVATE substruct_definitions molecules device flatBitVect device_vector device_timings) +add_executable(test_warp_dfs test_warp_dfs.cu) +target_link_libraries(test_warp_dfs PRIVATE cuda_error_check CUDA::cudart) + add_executable(test_substructure_integration test_substruct_integration.cu) target_link_libraries( test_substructure_integration @@ -428,6 +431,7 @@ set(TEST_LIST test_rdkit_comp test_rdkit_bounds_matrix test_substruct_algos + test_warp_dfs test_substructure_search test_substructure_integration test_substruct_label_integration diff --git a/tests/test_substruct_integration.cu b/tests/test_substruct_integration.cu index 4ea37e2b..d41ad33e 100644 --- a/tests/test_substruct_integration.cu +++ b/tests/test_substruct_integration.cu @@ -285,7 +285,7 @@ const ThreadingConfig kMainTestThreadingConfigs[] = { INSTANTIATE_TEST_SUITE_P(AllCombinations, SubstructureIntegrationTest, - ::testing::Combine(::testing::Values(SubstructAlgorithm::GSI), + ::testing::Combine(::testing::Values(SubstructAlgorithm::GSI, SubstructAlgorithm::DFS), ::testing::ValuesIn(kDatasets), ::testing::ValuesIn(kMainTestThreadingConfigs), ::testing::Values(SubstructMode::Matches)), @@ -294,29 +294,32 @@ INSTANTIATE_TEST_SUITE_P(AllCombinations, std::get<1>(info.param).name + "_" + std::get<2>(info.param).name; }); -INSTANTIATE_TEST_SUITE_P(ConfigOptionTests, - SubstructureIntegrationTest, - ::testing::Values(SubstructParams{SubstructAlgorithm::GSI, - kDatasets[0], - kThreadingConfigs[2], - SubstructMode::Matches}, // Autoselect - SubstructParams{SubstructAlgorithm::GSI, - kDatasets[0], - kThreadingConfigs[0], - SubstructMode::HasMatch}, // SingleThreaded, bool - SubstructParams{SubstructAlgorithm::GSI, - kDatasets[3], - kThreadingConfigs[2], - SubstructMode::CountMatches}), // Autoselect, counts - [](const ::testing::TestParamInfo& info) { - const SubstructMode mode = std::get<3>(info.param); - const char* modeSuffix = - (mode == SubstructMode::HasMatch) ? +INSTANTIATE_TEST_SUITE_P( + ConfigOptionTests, + SubstructureIntegrationTest, + ::testing::Values( + SubstructParams{SubstructAlgorithm::GSI, kDatasets[0], kThreadingConfigs[2], SubstructMode::Matches}, // Autoselect + SubstructParams{SubstructAlgorithm::GSI, + kDatasets[0], + kThreadingConfigs[0], + SubstructMode::HasMatch}, // SingleThreaded, bool + SubstructParams{SubstructAlgorithm::GSI, + kDatasets[3], + kThreadingConfigs[2], + SubstructMode::CountMatches}, // Autoselect, counts + SubstructParams{SubstructAlgorithm::DFS, kDatasets[0], kThreadingConfigs[0], SubstructMode::HasMatch}, // DFS bool + SubstructParams{SubstructAlgorithm::DFS, + kDatasets[3], + kThreadingConfigs[2], + SubstructMode::CountMatches}), // DFS counts + [](const ::testing::TestParamInfo& info) { + const SubstructMode mode = std::get<3>(info.param); + const char* modeSuffix = (mode == SubstructMode::HasMatch) ? "HasSubstructMatch" : (mode == SubstructMode::CountMatches ? "CountSubstructMatches" : "Autoselect"); - return std::string(algorithmName(std::get<0>(info.param))) + "_" + - std::get<1>(info.param).name + "_" + std::get<2>(info.param).name + "_" + modeSuffix; - }); + return std::string(algorithmName(std::get<0>(info.param))) + "_" + std::get<1>(info.param).name + "_" + + std::get<2>(info.param).name + "_" + modeSuffix; + }); TEST_P(SubstructureIntegrationTest, ChemblVsSmarts) { const std::string smilesPath = testDataPath_ + "/chembl_1k.smi"; diff --git a/tests/test_substruct_search.cu b/tests/test_substruct_search.cu index 19b3c5f4..8ecab82b 100644 --- a/tests/test_substruct_search.cu +++ b/tests/test_substruct_search.cu @@ -196,7 +196,7 @@ class SubstructureSearchTest : public ::testing::TestWithParam& info) { return algorithmName(info.param); }); @@ -204,9 +204,9 @@ INSTANTIATE_TEST_SUITE_P(AllAlgorithms, // Fixture for recursive SMARTS tests - VF2 doesn't support recursion class RecursiveSubstructureSearchTest : public SubstructureSearchTest {}; -INSTANTIATE_TEST_SUITE_P(GSIOnly, +INSTANTIATE_TEST_SUITE_P(RecursiveAlgorithms, RecursiveSubstructureSearchTest, - ::testing::Values(SubstructAlgorithm::GSI), + ::testing::Values(SubstructAlgorithm::GSI, SubstructAlgorithm::DFS), [](const ::testing::TestParamInfo& info) { return algorithmName(info.param); }); @@ -354,6 +354,32 @@ TEST_P(SubstructureSearchTest, DifferentMoleculeSizes) { compareWithRDKit(results, targetMols, queryMols, true); } +TEST_P(SubstructureSearchTest, SizeBracketsRingWithLargeHydrocarbon) { + // Cyclohexane with a long alkane tail, sized to land in each target bracket + // (<=32, <=64, <=128 atoms), against queries sized for each query bracket + // (<=16, <=32, <=64 atoms). The 40-carbon chain query cannot fit in the + // 28-atom target at all, matches the 50-atom target only by threading + // through the ring, and matches the 100-atom target's tail directly -- so + // every (target bracket, query bracket) pair exercises a distinct outcome. + const std::string ring = "C1CCCCC1"; + + std::vector> targetMols; + std::vector> queryMols; + parseMolecules({ring + std::string(22, 'C'), // 28 atoms: 32-target bracket + ring + std::string(44, 'C'), // 50 atoms: 64-target bracket + ring + std::string(94, 'C')}, // 100 atoms: 128-target bracket + {ring, // 6 atoms: 16-query bracket + std::string(20, 'C'), // 20 atoms: 32-query bracket + std::string(40, 'C')}, // 40 atoms: 64-query bracket + targetMols, + queryMols); + + SubstructSearchResults results; + getSubstructMatches(getRawPtrs(targetMols), getRawPtrs(queryMols), results, algorithm(), stream_.stream()); + + compareWithRDKit(results, targetMols, queryMols, true); +} + TEST_P(SubstructureSearchTest, MultiAtomQuery) { std::vector> targetMols; std::vector> queryMols; diff --git a/tests/test_warp_dfs.cu b/tests/test_warp_dfs.cu new file mode 100644 index 00000000..08fd04f4 --- /dev/null +++ b/tests/test_warp_dfs.cu @@ -0,0 +1,553 @@ +// 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. + +// Tests nvMolKit::dfsFromRoots against a host brute-force enumerator on +// synthetic graphs, using a minimal oracle: candidates(d) = compat[d] & ~used +// & AND over back edges of adj[mapping[e]]. This exercises the core's stack +// discipline, terminal verdicts, abort hook, and every TargetMask width -- +// including the 128-atom form's lo/hi word boundary -- without any molecule +// plumbing; the substructure integration tests cover the production oracle. + +#include +#include + +#include +#include + +#include "src/subgraph/target_mask.cuh" +#include "src/subgraph/warp_dfs.cuh" +#include "src/utils/cuda_error_check.h" + +using nvMolKit::checkReturnCode; + +namespace { + +constexpr int kHostMaxTargetAtoms = 128; + +/// A (target graph, query constraints) pair small enough to pass as a kernel +/// argument. Target-atom sets are stored as (lo, hi) 64-bit word pairs; hi is +/// unused below 65 atoms. adj[t] is target atom t's neighbour set; compat[d] +/// the target atoms query atom d may map to; backEdges[d] the earlier query +/// atoms d is bonded to. Every edge accepts every bond (bond predicates are +/// the oracle's business, not the core's). +template struct TestProblem { + uint64_t adjLo[MaxTargetAtoms] = {}; + uint64_t adjHi[MaxTargetAtoms] = {}; + uint64_t compatLo[MaxDepth] = {}; + uint64_t compatHi[MaxDepth] = {}; + unsigned char backEdges[MaxDepth][4] = {}; + unsigned char backEdgeCounts[MaxDepth] = {}; + int numQueryAtoms = 0; + int numTargetAtoms = 0; +}; + +enum class TestMode { + CountAll, ///< Every lane counts all embeddings from its roots. + FirstOnly, ///< All roots on lane 0; stop the lane at the first embedding and record it. + RootExists, ///< Flag each root that starts at least one embedding, then abandon it. + AbortAll, ///< Abort hook returns true from the start; nothing may be found. + ExistsRace ///< All lanes race for one embedding via a found flag: the MCS existence pattern. +}; + +template +__device__ nvMolKit::TargetMask maskFromWords(uint64_t lo, uint64_t hi) { + nvMolKit::TargetMask mask; + mask.clear(); + mask.setWord32(0, static_cast(lo)); + if constexpr (MaxTargetAtoms >= 64) { + mask.setWord32(1, static_cast(lo >> 32)); + } + if constexpr (MaxTargetAtoms >= 128) { + mask.setWord32(2, static_cast(hi)); + mask.setWord32(3, static_cast(hi >> 32)); + } + return mask; +} + +template +__global__ void warpDfsTestKernel(TestProblem problem, + TestMode mode, + int* outCount, + int* outMapping, + int* outRootFlags, + int* foundFlag) { + const int lane = static_cast(threadIdx.x); + using Mask = nvMolKit::TargetMask; + + const Mask compatRoots = maskFromWords(problem.compatLo[0], problem.compatHi[0]); + + Mask roots; + roots.clear(); + if (mode == TestMode::FirstOnly) { + if (lane == 0) { + roots = compatRoots; + } + } else { + // Production distribution: lane L owns target atoms L, L+32, L+64, L+96. + Mask laneAtoms; + laneAtoms.clear(); +#pragma unroll + for (int k = 0; k < static_cast(MaxTargetAtoms / 32); ++k) { + laneAtoms.set(lane + 32 * k); + } + roots = compatRoots; + roots.andEq(laneAtoms); + } + + auto candidatesAt = [&](int depth, const unsigned char* mapping, const Mask& used, int /*prevTargetAtom*/) -> Mask { + Mask candidates = maskFromWords(problem.compatLo[depth], problem.compatHi[depth]); + candidates.andNotEq(used); + for (int k = 0; k < problem.backEdgeCounts[depth]; ++k) { + const int mapped = mapping[problem.backEdges[depth][k]]; + candidates.andEq(maskFromWords(problem.adjLo[mapped], problem.adjHi[mapped])); + } + return candidates; + }; + + auto onTerminal = [&](Mask terminals, const unsigned char* mapping) -> nvMolKit::DfsTerminalVerdict { + nvMolKit::DfsTerminalVerdict verdict{false, false}; + if (mode == TestMode::CountAll || mode == TestMode::AbortAll) { + atomicAdd(outCount, terminals.popcount()); + } else if (mode == TestMode::FirstOnly) { + if (!terminals.empty()) { + for (int d = 0; d < problem.numQueryAtoms - 1; ++d) { + outMapping[d] = mapping[d]; + } + outMapping[problem.numQueryAtoms - 1] = terminals.lowest(); + atomicAdd(outCount, 1); + verdict.laneDone = true; + } + } else if (mode == TestMode::RootExists) { + if (!terminals.empty()) { + outRootFlags[mapping[0]] = 1; + verdict.rootDone = true; + } + } else { // ExistsRace + if (!terminals.empty()) { + if (atomicCAS(foundFlag, 0, 1) == 0) { + // This lane won: it alone records the embedding. + for (int d = 0; d < problem.numQueryAtoms - 1; ++d) { + outMapping[d] = mapping[d]; + } + outMapping[problem.numQueryAtoms - 1] = terminals.lowest(); + atomicAdd(outCount, 1); + } + verdict.laneDone = true; + } + } + return verdict; + }; + + auto abortRequested = [&]() -> bool { + if (mode == TestMode::AbortAll) { + return true; + } + if (mode == TestMode::ExistsRace) { + return atomicAdd(foundFlag, 0) != 0; + } + return false; + }; + + nvMolKit::dfsFromRoots(roots, problem.numQueryAtoms - 1, candidatesAt, onTerminal, abortRequested); +} + +// ============================================================================= +// Host-side graph construction and brute force +// ============================================================================= + +struct WordPair { + uint64_t lo = 0; + uint64_t hi = 0; +}; + +void setBit(WordPair& words, int bit) { + if (bit < 64) { + words.lo |= 1ULL << bit; + } else { + words.hi |= 1ULL << (bit - 64); + } +} + +bool testBit(const WordPair& words, int bit) { + return bit < 64 ? ((words.lo >> bit) & 1ULL) != 0 : ((words.hi >> (bit - 64)) & 1ULL) != 0; +} + +/// Size-agnostic host mirror of TestProblem, converted per instantiation. +struct HostProblem { + WordPair adj[kHostMaxTargetAtoms] = {}; + std::vector compat; + std::vector> backEdges; + int numQueryAtoms = 0; + int numTargetAtoms = 0; + + void init(int numTarget, int numQuery) { + numTargetAtoms = numTarget; + numQueryAtoms = numQuery; + compat.assign(numQuery, {}); + backEdges.assign(numQuery, {}); + } + + void addUndirectedEdge(int a, int b) { + setBit(adj[a], b); + setBit(adj[b], a); + } + + void allowAllLabels() { + for (int d = 0; d < numQueryAtoms; ++d) { + for (int t = 0; t < numTargetAtoms; ++t) { + setBit(compat[d], t); + } + } + } + + /// Chain query: each atom bonded to the previous one. + void makePathQuery() { + for (int d = 1; d < numQueryAtoms; ++d) { + backEdges[d] = {d - 1}; + } + } +}; + +template +TestProblem toDeviceProblem(const HostProblem& host) { + TestProblem device; + device.numTargetAtoms = host.numTargetAtoms; + device.numQueryAtoms = host.numQueryAtoms; + for (int t = 0; t < host.numTargetAtoms; ++t) { + device.adjLo[t] = host.adj[t].lo; + device.adjHi[t] = host.adj[t].hi; + } + for (int d = 0; d < host.numQueryAtoms; ++d) { + device.compatLo[d] = host.compat[d].lo; + device.compatHi[d] = host.compat[d].hi; + for (std::size_t k = 0; k < host.backEdges[d].size(); ++k) { + device.backEdges[d][k] = static_cast(host.backEdges[d][k]); + } + device.backEdgeCounts[d] = static_cast(host.backEdges[d].size()); + } + return device; +} + +// Host-side mirror of the kernel's rules, enumerating all injective embeddings. +void bruteForceRecurse(const HostProblem& problem, std::vector& mapping, WordPair used, int depth, long& count) { + if (depth == problem.numQueryAtoms) { + ++count; + return; + } + for (int t = 0; t < problem.numTargetAtoms; ++t) { + if (!testBit(problem.compat[depth], t) || testBit(used, t)) { + continue; + } + bool edgesOk = true; + for (const int e : problem.backEdges[depth]) { + if (!testBit(problem.adj[mapping[e]], t)) { + edgesOk = false; + break; + } + } + if (!edgesOk) { + continue; + } + mapping[depth] = t; + WordPair nextUsed = used; + setBit(nextUsed, t); + bruteForceRecurse(problem, mapping, nextUsed, depth + 1, count); + } +} + +long bruteForceCount(const HostProblem& problem) { + long count = 0; + std::vector mapping(problem.numQueryAtoms, -1); + bruteForceRecurse(problem, mapping, {}, 0, count); + return count; +} + +bool rootCanEmbed(const HostProblem& problem, int root) { + if (!testBit(problem.compat[0], root)) { + return false; + } + HostProblem restricted = problem; + restricted.compat[0] = {}; + setBit(restricted.compat[0], root); + return bruteForceCount(restricted) > 0; +} + +/// Cyclohexane-like ring 0-1-2, closing 2-0, plus tail edge 2-3. +HostProblem triangleWithTailTarget(int numQueryAtoms) { + HostProblem problem; + problem.init(4, numQueryAtoms); + problem.addUndirectedEdge(0, 1); + problem.addUndirectedEdge(0, 2); + problem.addUndirectedEdge(1, 2); + problem.addUndirectedEdge(2, 3); + problem.allowAllLabels(); + return problem; +} + +/// Cycle over target atoms firstAtom..firstAtom+cycleLength-1 inside a +/// numTargetAtoms-atom graph whose other atoms are isolated. +HostProblem cycleTarget(int numTargetAtoms, int firstAtom, int cycleLength, int numQueryAtoms) { + HostProblem problem; + problem.init(numTargetAtoms, numQueryAtoms); + for (int i = 0; i < cycleLength; ++i) { + problem.addUndirectedEdge(firstAtom + i, firstAtom + (i + 1) % cycleLength); + } + problem.allowAllLabels(); + return problem; +} + +// ============================================================================= +// Test fixture +// ============================================================================= + +constexpr int kMaxRecordedDepth = 48; + +class WarpDfsTest : public ::testing::Test { + protected: + template void run(const HostProblem& host, TestMode mode) { + ASSERT_LE(host.numTargetAtoms, static_cast(MaxTargetAtoms)); + ASSERT_LE(host.numQueryAtoms, MaxDepth); + ASSERT_LE(host.numQueryAtoms, kMaxRecordedDepth); + cudaCheckError(cudaMemset(devCount_, 0, sizeof(int))); + cudaCheckError(cudaMemset(devMapping_, 0xFF, kMaxRecordedDepth * sizeof(int))); + cudaCheckError(cudaMemset(devRootFlags_, 0, kHostMaxTargetAtoms * sizeof(int))); + cudaCheckError(cudaMemset(devFoundFlag_, 0, sizeof(int))); + const auto problem = toDeviceProblem(host); + warpDfsTestKernel + <<<1, 32>>>(problem, mode, devCount_, devMapping_, devRootFlags_, devFoundFlag_); + cudaCheckError(cudaGetLastError()); + cudaCheckError(cudaDeviceSynchronize()); + cudaCheckError(cudaMemcpy(&count_, devCount_, sizeof(int), cudaMemcpyDeviceToHost)); + cudaCheckError(cudaMemcpy(mapping_, devMapping_, kMaxRecordedDepth * sizeof(int), cudaMemcpyDeviceToHost)); + cudaCheckError(cudaMemcpy(rootFlags_, devRootFlags_, kHostMaxTargetAtoms * sizeof(int), cudaMemcpyDeviceToHost)); + } + + void expectValidEmbedding(const HostProblem& problem) { + for (int d = 0; d < problem.numQueryAtoms; ++d) { + ASSERT_GE(mapping_[d], 0); + ASSERT_LT(mapping_[d], problem.numTargetAtoms); + EXPECT_TRUE(testBit(problem.compat[d], mapping_[d])); + for (int e = 0; e < d; ++e) { + EXPECT_NE(mapping_[d], mapping_[e]); + } + for (const int k : problem.backEdges[d]) { + EXPECT_TRUE(testBit(problem.adj[mapping_[k]], mapping_[d])); + } + } + } + + void SetUp() override { + cudaCheckError(cudaMalloc(&devCount_, sizeof(int))); + cudaCheckError(cudaMalloc(&devMapping_, kMaxRecordedDepth * sizeof(int))); + cudaCheckError(cudaMalloc(&devRootFlags_, kHostMaxTargetAtoms * sizeof(int))); + cudaCheckError(cudaMalloc(&devFoundFlag_, sizeof(int))); + } + + void TearDown() override { + cudaFree(devCount_); + cudaFree(devMapping_); + cudaFree(devRootFlags_); + cudaFree(devFoundFlag_); + } + + int count_ = 0; + int mapping_[kMaxRecordedDepth] = {}; + int rootFlags_[kHostMaxTargetAtoms] = {}; + int* devCount_ = nullptr; + int* devMapping_ = nullptr; + int* devRootFlags_ = nullptr; + int* devFoundFlag_ = nullptr; +}; + +// ============================================================================= +// 32-atom form +// ============================================================================= + +TEST_F(WarpDfsTest, CountsPathEmbeddings) { + // Query: path 0-1-2. In triangle+tail: sum over middle atoms of + // degree*(degree-1) = 2 + 2 + 6 + 0 = 10 ordered embeddings. + HostProblem problem = triangleWithTailTarget(3); + problem.makePathQuery(); + + run<32, 8>(problem, TestMode::CountAll); + EXPECT_EQ(count_, 10); + EXPECT_EQ(bruteForceCount(problem), 10); +} + +TEST_F(WarpDfsTest, CountsTriangleEmbeddings) { + // Query: triangle. Only the target triangle matches, in all 3! orders. + HostProblem problem = triangleWithTailTarget(3); + problem.makePathQuery(); + problem.backEdges[2] = {0, 1}; + + run<32, 8>(problem, TestMode::CountAll); + EXPECT_EQ(count_, 6); + EXPECT_EQ(bruteForceCount(problem), 6); +} + +TEST_F(WarpDfsTest, RespectsLabelCompatibility) { + // Path query whose first atom may only map to target atoms 2 or 3. + HostProblem problem = triangleWithTailTarget(3); + problem.makePathQuery(); + problem.compat[0] = {}; + setBit(problem.compat[0], 2); + setBit(problem.compat[0], 3); + + run<32, 8>(problem, TestMode::CountAll); + EXPECT_EQ(static_cast(count_), bruteForceCount(problem)); +} + +TEST_F(WarpDfsTest, RingClosureMatchesBruteForce) { + // 4-atom query with a ring closure (cycle 0-1-2-3-0 plus chord 0-2) forces + // backtracking through several depths; validate against brute force. + HostProblem problem; + problem.init(6, 4); + problem.addUndirectedEdge(0, 1); + problem.addUndirectedEdge(1, 2); + problem.addUndirectedEdge(2, 3); + problem.addUndirectedEdge(3, 0); + problem.addUndirectedEdge(0, 2); + problem.addUndirectedEdge(3, 4); + problem.addUndirectedEdge(4, 5); + problem.allowAllLabels(); + problem.makePathQuery(); + problem.backEdges[3] = {2, 0}; + + run<32, 8>(problem, TestMode::CountAll); + const long expected = bruteForceCount(problem); + EXPECT_GT(expected, 0); + EXPECT_EQ(static_cast(count_), expected); +} + +TEST_F(WarpDfsTest, LaneDoneStopsAfterFirstEmbedding) { + HostProblem problem = triangleWithTailTarget(3); + problem.makePathQuery(); + + run<32, 8>(problem, TestMode::FirstOnly); + EXPECT_EQ(count_, 1); + expectValidEmbedding(problem); +} + +TEST_F(WarpDfsTest, RootDoneFlagsExactlyTheRootsThatEmbed) { + // Triangle query: target atoms 0,1,2 can each start an embedding; the tail + // atom 3 cannot. + HostProblem problem = triangleWithTailTarget(3); + problem.makePathQuery(); + problem.backEdges[2] = {0, 1}; + + run<32, 8>(problem, TestMode::RootExists); + EXPECT_EQ(rootFlags_[0], 1); + EXPECT_EQ(rootFlags_[1], 1); + EXPECT_EQ(rootFlags_[2], 1); + EXPECT_EQ(rootFlags_[3], 0); +} + +TEST_F(WarpDfsTest, AbortHookStopsBeforeAnyRoot) { + HostProblem problem = triangleWithTailTarget(3); + problem.makePathQuery(); + + run<32, 8>(problem, TestMode::AbortAll); + EXPECT_EQ(count_, 0); +} + +// ============================================================================= +// 64- and 128-atom forms, multi-root lanes, deep stacks +// ============================================================================= + +TEST_F(WarpDfsTest, Cycle64CountsAcrossBothMaskWords) { + // 64-cycle, 5-atom path query: every atom starts one path per direction, so + // 128 ordered embeddings, half of them rooted in the mask's upper word. Two + // roots per lane exercises the root loop past its first root. + HostProblem problem = cycleTarget(64, 0, 64, 5); + problem.makePathQuery(); + + run<64, 8>(problem, TestMode::CountAll); + EXPECT_EQ(count_, 128); + EXPECT_EQ(bruteForceCount(problem), 128); +} + +TEST_F(WarpDfsTest, Cycle64RootDoneAdvancesToSecondRootOfLane) { + // Every atom of the 64-cycle can root an embedding. RootExists abandons a + // root at its first hit, so flagging all 64 proves each lane continued to + // its second root after rootDone on the first. + HostProblem problem = cycleTarget(64, 0, 64, 5); + problem.makePathQuery(); + + run<64, 8>(problem, TestMode::RootExists); + for (int t = 0; t < 64; ++t) { + EXPECT_EQ(rootFlags_[t], 1) << "root " << t; + } +} + +TEST_F(WarpDfsTest, Ring128StraddlesWordBoundary) { + // 12-cycle on atoms 58..69 of a 128-atom target: the ring's edges, roots, + // and used-set updates all cross the TargetMask<128> lo/hi boundary. The + // remaining 116 atoms are isolated, so a 4-path query embeds only in the + // ring: 12 starts x 2 directions = 24. + HostProblem problem = cycleTarget(128, 58, 12, 4); + problem.makePathQuery(); + + run<128, 8>(problem, TestMode::CountAll); + EXPECT_EQ(count_, 24); + EXPECT_EQ(bruteForceCount(problem), 24); + + run<128, 8>(problem, TestMode::RootExists); + for (int t = 0; t < 128; ++t) { + EXPECT_EQ(rootFlags_[t], rootCanEmbed(problem, t) ? 1 : 0) << "root " << t; + } +} + +TEST_F(WarpDfsTest, DeepPathQueryFillsTheStack) { + // 40-atom path query on a 64-cycle, with MaxDepth == numQueryAtoms so the + // terminal fires at exactly MaxDepth - 1: 64 starts x 2 directions = 128, + // found only by walking the stack through all 40 depths. + HostProblem problem = cycleTarget(64, 0, 64, 40); + problem.makePathQuery(); + + run<64, 40>(problem, TestMode::CountAll); + EXPECT_EQ(count_, 128); + EXPECT_EQ(bruteForceCount(problem), 128); +} + +TEST_F(WarpDfsTest, ExistsRaceFindsExactlyOneEmbedding) { + // The MCS existence pattern: lanes race via a found flag polled in the abort + // hook, and the atomicCAS winner alone records its embedding. + HostProblem problem = cycleTarget(64, 0, 64, 5); + problem.makePathQuery(); + + run<64, 8>(problem, TestMode::ExistsRace); + EXPECT_EQ(count_, 1); + expectValidEmbedding(problem); +} + +TEST_F(WarpDfsTest, ExistsRaceFindsNothingWhenNoEmbeddingExists) { + // Same shape, but the query path is one atom longer than the ring region + // can hold injectively... use a 6-path against a 5-cycle in a 64 target. + HostProblem problem = cycleTarget(64, 10, 5, 6); + problem.makePathQuery(); + // Restrict all depths to the cycle's atoms so isolated atoms cannot help. + for (int d = 0; d < problem.numQueryAtoms; ++d) { + problem.compat[d] = {}; + for (int i = 0; i < 5; ++i) { + setBit(problem.compat[d], 10 + i); + } + } + + ASSERT_EQ(bruteForceCount(problem), 0); + run<64, 8>(problem, TestMode::ExistsRace); + EXPECT_EQ(count_, 0); +} + +} // namespace