From b0adc27e1065ff487a95f3e9390346e976549e8f Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Wed, 2 Sep 2026 13:59:27 -0400 Subject: [PATCH 1/3] Optimize Morgan Python input conversion --- nvmolkit/CMakeLists.txt | 1 + nvmolkit/fingerprints.cpp | 74 +++++++++++++++++++++-------- nvmolkit/fingerprints.py | 6 ++- nvmolkit/tests/test_fingerprints.py | 37 ++++++++++++--- 4 files changed, 90 insertions(+), 28 deletions(-) diff --git a/nvmolkit/CMakeLists.txt b/nvmolkit/CMakeLists.txt index 27733915..265ad620 100644 --- a/nvmolkit/CMakeLists.txt +++ b/nvmolkit/CMakeLists.txt @@ -30,6 +30,7 @@ target_link_libraries( ${RDKit_LIBS} morganFingerprint device + nvtx CUDA::cudart) target_include_directories(_Fingerprints PRIVATE ${Python_INCLUDE_DIRS}) installpythontarget(_Fingerprints ./) diff --git a/nvmolkit/fingerprints.cpp b/nvmolkit/fingerprints.cpp index 4482d3c4..a95f1a7a 100644 --- a/nvmolkit/fingerprints.cpp +++ b/nvmolkit/fingerprints.cpp @@ -21,11 +21,36 @@ #include "nvmolkit/array_helpers.h" #include "src/morgan_fingerprint.h" #include "src/utils/device.h" +#include "src/utils/nvtx.h" namespace { using namespace boost::python; +struct PythonMoleculeSequence { + object owner; + std::vector molecules; +}; + +PythonMoleculeSequence convertMolecules(const object& mols) { + PyObject* sequence = PySequence_Fast(mols.ptr(), "mols must be a sequence of RDKit molecules"); + if (sequence == nullptr) { + throw_error_already_set(); + } + + PythonMoleculeSequence result{object(handle<>(sequence)), {}}; + const Py_ssize_t numMols = PySequence_Fast_GET_SIZE(sequence); + result.molecules.reserve(static_cast(numMols)); + for (Py_ssize_t i = 0; i < numMols; ++i) { + const RDKit::ROMol* mol = extract(PySequence_Fast_GET_ITEM(sequence, i)); + if (mol == nullptr) { + throw std::invalid_argument("Invalid molecule at index " + std::to_string(i)); + } + result.molecules.push_back(mol); + } + return result; +} + template nvMolKit::PyArray* makePyArrayFromFlatBitVects(nvMolKit::AsyncDeviceVector>& deviceVect) { using dtype = typename nvMolKit::FlatBitVect::StorageType; @@ -38,6 +63,20 @@ nvMolKit::PyArray* makePyArrayFromFlatBitVects(nvMolKit::AsyncDeviceVector +nvMolKit::PyArray* getFingerprintsDevice(nvMolKit::MorganFingerprintGenerator& generator, + const std::vector& mols, + cudaStream_t stream, + const nvMolKit::FingerprintComputeOptions& computeOptions) { + auto fingerprints = [&]() { + nvMolKit::ScopedNvtxRange range("MorganFPBindingNativeCompute", nvMolKit::NvtxColor::kOrange); + return generator.GetFingerprintsGpuBuffer(mols, stream, computeOptions); + }(); + + nvMolKit::ScopedNvtxRange range("MorganFPBindingOutputWrapping", nvMolKit::NvtxColor::kCyan); + return makePyArrayFromFlatBitVects(fingerprints); +} + } // namespace BOOST_PYTHON_MODULE(_Fingerprints) { @@ -53,21 +92,21 @@ BOOST_PYTHON_MODULE(_Fingerprints) { .def( "GetFingerprintsDevice", +[](nvMolKit::MorganFingerprintGenerator& selfref, - boost::python::list& mols, + const boost::python::object& mols, int numThreads, std::uintptr_t streamPtr) { - std::vector molsVec; - molsVec.reserve(len(mols)); - for (int i = 0; i < len(mols); i++) { - molsVec.push_back(boost::python::extract(boost::python::object(mols[i]))); - if (molsVec.back() == nullptr) { - throw std::invalid_argument("Invalid molecule at index " + std::to_string(i)); - } - } + auto convertedMols = [&]() { + nvMolKit::ScopedNvtxRange range("MorganFPBindingInputConversion", nvMolKit::NvtxColor::kYellow); + return convertMolecules(mols); + }(); + nvMolKit::FingerprintComputeOptions computeOptions; computeOptions.backend = nvMolKit::FingerprintComputeBackend::GPU; computeOptions.numCpuThreads = numThreads; - auto streamOpt = nvMolKit::acquireExternalStream(streamPtr); + auto streamOpt = [&]() { + nvMolKit::ScopedNvtxRange range("MorganFPBindingStreamAcquisition", nvMolKit::NvtxColor::kBlue); + return nvMolKit::acquireExternalStream(streamPtr); + }(); if (!streamOpt) { throw std::invalid_argument("Invalid CUDA stream"); } @@ -75,24 +114,19 @@ BOOST_PYTHON_MODULE(_Fingerprints) { const auto& options = selfref.GetOptions(); switch (options.fpSize) { case 128: { - auto array = selfref.GetFingerprintsGpuBuffer<128>(molsVec, stream, computeOptions); - return makePyArrayFromFlatBitVects<128>(array); + return getFingerprintsDevice<128>(selfref, convertedMols.molecules, stream, computeOptions); } case 256: { - auto array = selfref.GetFingerprintsGpuBuffer<256>(molsVec, stream, computeOptions); - return makePyArrayFromFlatBitVects<256>(array); + return getFingerprintsDevice<256>(selfref, convertedMols.molecules, stream, computeOptions); } case 512: { - auto array = selfref.GetFingerprintsGpuBuffer<512>(molsVec, stream, computeOptions); - return makePyArrayFromFlatBitVects<512>(array); + return getFingerprintsDevice<512>(selfref, convertedMols.molecules, stream, computeOptions); } case 1024: { - auto array = selfref.GetFingerprintsGpuBuffer<1024>(molsVec, stream, computeOptions); - return makePyArrayFromFlatBitVects<1024>(array); + return getFingerprintsDevice<1024>(selfref, convertedMols.molecules, stream, computeOptions); } case 2048: { - auto array = selfref.GetFingerprintsGpuBuffer<2048>(molsVec, stream, computeOptions); - return makePyArrayFromFlatBitVects<2048>(array); + return getFingerprintsDevice<2048>(selfref, convertedMols.molecules, stream, computeOptions); } default: throw std::invalid_argument("Invalid fpSize: " + std::to_string(options.fpSize) + diff --git a/nvmolkit/fingerprints.py b/nvmolkit/fingerprints.py index c03ab64c..35243432 100644 --- a/nvmolkit/fingerprints.py +++ b/nvmolkit/fingerprints.py @@ -15,6 +15,8 @@ """GPU-accelerated fingerprint generation.""" +from collections.abc import Sequence + import torch from nvmolkit._Fingerprints import MorganFingerprintGenerator as InternalFPGen @@ -92,7 +94,7 @@ def __init__(self, radius: int, fpSize: int): self._internal = InternalFPGen(radius, fpSize) def GetFingerprints( - self, mols: list, num_threads: int = 0, stream: torch.cuda.Stream | None = None + self, mols: Sequence, num_threads: int = 0, stream: torch.cuda.Stream | None = None ) -> AsyncGpuResult: """Compute Morgan fingerprints for a list of molecules. @@ -104,7 +106,7 @@ def GetFingerprints( via `unpack_fingerprint`. Args: - mols: List of RDKit molecules to generate fingerprints for + mols: Sequence of RDKit molecules to generate fingerprints for num_threads: Number of CPU threads to use for fingerprint generation. If 0, uses all available threads. stream: CUDA stream to use. If None, uses the current stream. diff --git a/nvmolkit/tests/test_fingerprints.py b/nvmolkit/tests/test_fingerprints.py index 46defe5a..1258fd3d 100644 --- a/nvmolkit/tests/test_fingerprints.py +++ b/nvmolkit/tests/test_fingerprints.py @@ -20,7 +20,6 @@ from nvmolkit.fingerprints import MorganFingerprintGenerator, pack_fingerprint, unpack_fingerprint - def test_roundtrip_pack_unpack(): n_fps = 10 fp_size = 128 @@ -75,14 +74,40 @@ def test_nvmolkit_fingerprint_throws_on_invalid_fpsize(fpSize, size_limited_mols def test_empty_input(): fpgen = MorganFingerprintGenerator(radius=3, fpSize=2048) - fps = fpgen.GetFingerprints([]).torch() - assert fps.shape == (0, 2048 // 32) + for mols in ([], ()): + fps = fpgen.GetFingerprints(mols).torch() + assert fps.shape == (0, 2048 // 32) -def test_invalid_input(): +@pytest.mark.parametrize("mols", ([None], (Chem.MolFromSmiles("CC"), None))) +def test_invalid_input(mols): fpgen = MorganFingerprintGenerator(radius=3, fpSize=2048) - with pytest.raises(ValueError, match="Invalid molecule at index 0"): - fpgen.GetFingerprints([None]) + invalid_index = mols.index(None) + with pytest.raises(ValueError, match=rf"Invalid molecule at index {invalid_index}"): + fpgen.GetFingerprints(mols) + + +def test_fingerprints_accept_tuple_input(size_limited_mols): + mols = size_limited_mols[:8] + reordered = tuple(reversed(mols)) + gen = MorganFingerprintGenerator(radius=2, fpSize=256) + + actual = gen.GetFingerprints(reordered, num_threads=2).torch() + expected = gen.GetFingerprints(list(reversed(mols)), num_threads=2).torch() + + torch.testing.assert_close(actual, expected) + + +def test_fingerprints_large_sequence_preserves_every_row(size_limited_mols): + source = size_limited_mols[:16] + mols = tuple(source * 32) + gen = MorganFingerprintGenerator(radius=2, fpSize=256) + + actual = gen.GetFingerprints(mols, num_threads=2).torch() + expected = gen.GetFingerprints(source, num_threads=2).torch().repeat((32, 1)) + + assert actual.shape == (512, 256 // 32) + torch.testing.assert_close(actual, expected) @pytest.mark.parametrize("fpSize", (128, 1024, 2048)) From 9108a5daa25b31977cab4715ad1163910d1bd028 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Wed, 2 Sep 2026 14:05:45 -0400 Subject: [PATCH 2/3] Test Morgan non-sequence input rejection --- nvmolkit/fingerprints.cpp | 4 ++++ nvmolkit/tests/test_fingerprints.py | 8 ++++++++ 2 files changed, 12 insertions(+) diff --git a/nvmolkit/fingerprints.cpp b/nvmolkit/fingerprints.cpp index a95f1a7a..5e9f29da 100644 --- a/nvmolkit/fingerprints.cpp +++ b/nvmolkit/fingerprints.cpp @@ -33,6 +33,10 @@ struct PythonMoleculeSequence { }; PythonMoleculeSequence convertMolecules(const object& mols) { + if (!PySequence_Check(mols.ptr())) { + PyErr_SetString(PyExc_TypeError, "mols must be a sequence of RDKit molecules"); + throw_error_already_set(); + } PyObject* sequence = PySequence_Fast(mols.ptr(), "mols must be a sequence of RDKit molecules"); if (sequence == nullptr) { throw_error_already_set(); diff --git a/nvmolkit/tests/test_fingerprints.py b/nvmolkit/tests/test_fingerprints.py index 1258fd3d..564a81d3 100644 --- a/nvmolkit/tests/test_fingerprints.py +++ b/nvmolkit/tests/test_fingerprints.py @@ -20,6 +20,7 @@ from nvmolkit.fingerprints import MorganFingerprintGenerator, pack_fingerprint, unpack_fingerprint + def test_roundtrip_pack_unpack(): n_fps = 10 fp_size = 128 @@ -87,6 +88,13 @@ def test_invalid_input(mols): fpgen.GetFingerprints(mols) +@pytest.mark.parametrize("mols", (42, iter(()))) +def test_non_sequence_input(mols): + fpgen = MorganFingerprintGenerator(radius=3, fpSize=2048) + with pytest.raises(TypeError, match="mols must be a sequence of RDKit molecules"): + fpgen.GetFingerprints(mols) + + def test_fingerprints_accept_tuple_input(size_limited_mols): mols = size_limited_mols[:8] reordered = tuple(reversed(mols)) From e4a129ce34aa2b5c1117d491dac7e74d3526a683 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Wed, 2 Sep 2026 14:38:50 -0400 Subject: [PATCH 3/3] Hoist Morgan Python sequence item lookup --- nvmolkit/fingerprints.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/nvmolkit/fingerprints.cpp b/nvmolkit/fingerprints.cpp index 5e9f29da..271189a4 100644 --- a/nvmolkit/fingerprints.cpp +++ b/nvmolkit/fingerprints.cpp @@ -44,9 +44,10 @@ PythonMoleculeSequence convertMolecules(const object& mols) { PythonMoleculeSequence result{object(handle<>(sequence)), {}}; const Py_ssize_t numMols = PySequence_Fast_GET_SIZE(sequence); + PyObject* const* items = PySequence_Fast_ITEMS(sequence); result.molecules.reserve(static_cast(numMols)); for (Py_ssize_t i = 0; i < numMols; ++i) { - const RDKit::ROMol* mol = extract(PySequence_Fast_GET_ITEM(sequence, i)); + const RDKit::ROMol* mol = extract(items[i]); if (mol == nullptr) { throw std::invalid_argument("Invalid molecule at index " + std::to_string(i)); }