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..271189a4 100644 --- a/nvmolkit/fingerprints.cpp +++ b/nvmolkit/fingerprints.cpp @@ -21,11 +21,41 @@ #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) { + 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(); + } + + 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(items[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 +68,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 +97,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 +119,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..564a81d3 100644 --- a/nvmolkit/tests/test_fingerprints.py +++ b/nvmolkit/tests/test_fingerprints.py @@ -75,14 +75,47 @@ 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) + + +@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)) + 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))