Reduce python transfer overhead in morgan fingerprints

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 ./)

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<const RDKit::ROMol*> 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<std::size_t>(numMols));
+  for (Py_ssize_t i = 0; i < numMols; ++i) {
+    const RDKit::ROMol* mol = extract<const RDKit::ROMol*>(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 <int nBits>
 nvMolKit::PyArray* makePyArrayFromFlatBitVects(nvMolKit::AsyncDeviceVector<nvMolKit::FlatBitVect<nBits>>& deviceVect) {
   using dtype                = typename nvMolKit::FlatBitVect<nBits>::StorageType;
@@ -38,6 +68,20 @@ nvMolKit::PyArray* makePyArrayFromFlatBitVects(nvMolKit::AsyncDeviceVector<nvMol
   return nvMolKit::makePyArray(deviceVect, dTypeStr, boost::python::make_tuple(nRows, nCols));
 }
 
+template <int nBits>
+nvMolKit::PyArray* getFingerprintsDevice(nvMolKit::MorganFingerprintGenerator&      generator,
+                                         const std::vector<const RDKit::ROMol*>&    mols,
+                                         cudaStream_t                               stream,
+                                         const nvMolKit::FingerprintComputeOptions& computeOptions) {
+  auto fingerprints = [&]() {
+    nvMolKit::ScopedNvtxRange range("MorganFPBindingNativeCompute", nvMolKit::NvtxColor::kOrange);
+    return generator.GetFingerprintsGpuBuffer<nBits>(mols, stream, computeOptions);
+  }();
+
+  nvMolKit::ScopedNvtxRange range("MorganFPBindingOutputWrapping", nvMolKit::NvtxColor::kCyan);
+  return makePyArrayFromFlatBitVects<nBits>(fingerprints);
+}
+
 }  // namespace
 
 BOOST_PYTHON_MODULE(_Fingerprints) {
@@ -53,46 +97,41 @@ 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<const RDKit::ROMol*> molsVec;
-        molsVec.reserve(len(mols));
-        for (int i = 0; i < len(mols); i++) {
-          molsVec.push_back(boost::python::extract<const RDKit::ROMol*>(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");
         }
         auto        stream  = *streamOpt;
         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) +

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.
 

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))