Graph-based CUDA ClusterFinder (ClusterFinderCUDAGraph)

- CUDA Graph variant of ClusterFinderCUDA: one pre-recorded graph per stream
  (memset + H2D + kernel + D2H), with per-frame src/dst pointers swapped via
  cudaGraphExecMemcpyNodeSetParams to cut per-frame launch overhead.
- Exposed through the _aare_cuda bindings and a ClusterFinderCUDAGraph factory.

python/aare/ClusterFinder.py

@@ -111,6 +111,47 @@ def ClusterFinderCUDA(image_size, cluster_size=(3,3), n_sigma=5, dtype=np.int32,
                max_clusters_per_frame=max_clusters_per_frame,
                n_streams=n_streams)
 
+def ClusterFinderCUDAGraph(image_size, cluster_size=(3,3), n_sigma=5, dtype=np.int32,
+                           max_clusters_per_frame=2048, n_streams=4):
+    """
+    Factory function to create a ClusterFinderCUDAGraph object. Uses pre-recorded
+    CUDA Graphs to reduce per-frame CPU API overhead (~23 µs vs ~31 µs for the
+    stream-based version), potentially improving throughput when processing is
+    CPU-overhead-bound.
+
+    Parameters
+    ----------
+    image_size : tuple of (int, int)
+        Detector shape as (nrows, ncols).
+    cluster_size : tuple of (int, int), optional
+        Cluster window size; default (3, 3).
+    n_sigma : float, optional
+        Threshold in units of per-pixel pedestal standard deviation.
+    dtype : numpy dtype, optional
+        Cluster value type (np.int32 or np.float32).
+    max_clusters_per_frame : int, optional
+        Hard upper bound on clusters per frame. Default 2048.
+    n_streams : int, optional
+        Number of CUDA streams (one graph per stream). Default 4.
+
+    Note
+    ----
+    avg_kernel_time_ms() always returns 0.0 for this variant — use
+    wall-clock timing around find_clusters_batched() instead.
+    """
+    if not _cuda_available():
+        raise RuntimeError(
+            "ClusterFinderCUDAGraph is not available in this build of aare. "
+            "Rebuild with -DAARE_CUDA=ON (and -DAARE_PYTHON_BINDINGS=ON)."
+        )
+
+    cls = _get_class("ClusterFinderCUDAGraph", cluster_size, dtype)
+    return cls(image_size,
+               n_sigma=n_sigma,
+               max_clusters_per_frame=max_clusters_per_frame,
+               n_streams=n_streams)
+
+
 def ClusterCollector(clusterfindermt, dtype=np.int32): 
     """ 
     Factory function to create a ClusterCollector object. Provides a cleaner syntax for 

python/aare/__init__.py

@@ -32,7 +32,7 @@ from ._aare import corner
 
 from ._version import __version__
 from .ClusterFinder import ClusterFinder, ClusterCollector, ClusterFinderMT, ClusterFileSink, ClusterFile
-from .ClusterFinder import ClusterFinderCUDA, _cuda_available
+from .ClusterFinder import ClusterFinderCUDA, ClusterFinderCUDAGraph, _cuda_available
 from .ClusterVector import ClusterVector
 from .Cluster import Cluster
 

python/src/bind_ClusterFinderCUDAGraph.hpp

@@ -0,0 +1,120 @@
+// SPDX-License-Identifier: MPL-2.0
+#pragma once
+#include "aare/ClusterFinderCUDA_graph.hpp"
+#include "aare/ClusterVector.hpp"
+#include "aare/NDView.hpp"
+#include "aare/Pedestal.hpp"
+#include "np_helper.hpp"
+
+#include <cstdint>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+using pd_type = double;
+
+using namespace aare;
+
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wunused-parameter"
+
+namespace aare {
+
+template <typename T, uint8_t ClusterSizeX, uint8_t ClusterSizeY,
+          typename CoordType = uint16_t>
+void define_ClusterFinderCUDAGraph(py::module &m, const std::string &typestr) {
+    auto class_name = fmt::format("ClusterFinderCUDAGraph_{}", typestr);
+
+    using ClusterType = Cluster<T, ClusterSizeX, ClusterSizeY, CoordType>;
+    using CF = ClusterFinderCUDAGraph<ClusterType, uint16_t, pd_type>;
+    using ContigArr =
+        py::array_t<uint16_t, py::array::c_style | py::array::forcecast>;
+
+    py::class_<CF>(m, class_name.c_str())
+        .def(py::init<Shape<2>, float, size_t, int>(), py::arg("image_size"),
+             py::arg("n_sigma") = 5.0f,
+             py::arg("max_clusters_per_frame") = 2048, py::arg("n_streams") = 4)
+
+        .def_property(
+            "nSigma", &CF::get_nSigma, &CF::set_nSigma,
+            R"(Number of sigma above the pedestal to consider a photon during cluster finding.)")
+
+        .def("push_pedestal_frame",
+             [](CF &self, ContigArr frame) {
+                 auto view = make_view_2d(frame);
+                 self.push_pedestal_frame(view);
+             })
+
+        .def("clear_pedestal", &CF::clear_pedestal)
+
+        .def_property_readonly("pedestal",
+                               [](CF &self) {
+                                   auto pd = new NDArray<pd_type, 2>{};
+                                   *pd = self.pedestal();
+                                   return return_image_data(pd);
+                               })
+
+        .def_property_readonly("noise",
+                               [](CF &self) {
+                                   auto arr = new NDArray<pd_type, 2>{};
+                                   *arr = self.noise();
+                                   return return_image_data(arr);
+                               })
+
+        .def(
+            "steal_clusters",
+            [](CF &self, bool realloc_same_capacity) {
+                return std::move(self.steal_clusters(realloc_same_capacity));
+            },
+            py::arg("realloc_same_capacity") = true)
+
+        .def(
+            "find_clusters",
+            [](CF &self, ContigArr frame, uint64_t frame_number) {
+                auto view = make_view_2d(frame);
+                self.find_clusters(view, frame_number);
+            },
+            py::arg("frame"), py::arg("frame_number") = 0,
+            py::call_guard<py::gil_scoped_release>())
+
+        .def(
+            "find_clusters_batched",
+            [](CF &self, ContigArr frames, uint64_t first_frame) {
+                auto view = make_view_3d(frames);
+                return self.find_clusters_batched(view, first_frame);
+            },
+            py::arg("frames"), py::arg("first_frame") = 0,
+            py::call_guard<py::gil_scoped_release>(),
+            R"(Process a 3D array of frames (n_frames, nrows, ncols) using
+            n_streams CUDA Graphs for H2D/kernel/D2H pipelining. Returns a
+            list of ClusterVector, one per input frame.)")
+
+        .def(
+            "avg_kernel_time_ms", &CF::avg_kernel_time_ms,
+            R"(Always returns 0.0 — graph version does not instrument individual kernel time.
+            Use wall-clock timing around find_clusters_batched instead.)")
+
+        .def("reset_timers", &CF::reset_timers)
+
+        .def(
+            "register_input_buffer",
+            [](CF &self, py::array arr) {
+                auto info = arr.request();
+                self.register_input_buffer(
+                    info.ptr, static_cast<size_t>(info.size) *
+                                  static_cast<size_t>(info.itemsize));
+            },
+            R"(Pin a numpy array as a locked host buffer so that
+            find_clusters_batched transfers it at full DMA bandwidth.
+            Call once before the processing loop with the full data array.
+            Slices passed to find_clusters_batched lie within the registered
+            region and benefit automatically. Call unregister_input_buffer()
+            when done.)")
+
+        .def("unregister_input_buffer", &CF::unregister_input_buffer,
+             "Release the previously pinned input buffer.");
+}
+
+} // namespace aare
+
+#pragma GCC diagnostic pop

python/src/cuda_bindings.cu

@@ -7,6 +7,7 @@
 
 #include "bind_Cluster.hpp"
 #include "bind_ClusterFinderCUDA.hpp"
+#include "bind_ClusterFinderCUDAGraph.hpp"
 #include "bind_ClusterVector.hpp"
 
 #include <pybind11/pybind11.h>
@@ -25,6 +26,10 @@ namespace py = pybind11;
     aare::define_ClusterFinderCUDA<T, N, M, U>(m, "Cluster" #N                 \
                                                   "x" #M #TYPE_CODE);
 
+#define DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(T, N, M, U, TYPE_CODE)        \
+    aare::define_ClusterFinderCUDAGraph<T, N, M, U>(m, "Cluster" #N            \
+                                                       "x" #M #TYPE_CODE);
+
 PYBIND11_MODULE(_aare_cuda, m) {
 
     // Types first — finders reference them in their signatures.
@@ -61,7 +66,25 @@ PYBIND11_MODULE(_aare_cuda, m) {
     DEFINE_BINDINGS_CLUSTERFINDER_CUDA(int, 9, 9, uint16_t, i);
     DEFINE_BINDINGS_CLUSTERFINDER_CUDA(double, 9, 9, uint16_t, d);
     DEFINE_BINDINGS_CLUSTERFINDER_CUDA(float, 9, 9, uint16_t, f);
+
+    // Graph-based finders
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(int, 3, 3, uint16_t, i);
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(double, 3, 3, uint16_t, d);
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(float, 3, 3, uint16_t, f);
+
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(int, 5, 5, uint16_t, i);
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(double, 5, 5, uint16_t, d);
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(float, 5, 5, uint16_t, f);
+
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(int, 7, 7, uint16_t, i);
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(double, 7, 7, uint16_t, d);
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(float, 7, 7, uint16_t, f);
+
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(int, 9, 9, uint16_t, i);
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(double, 9, 9, uint16_t, d);
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH(float, 9, 9, uint16_t, f);
 }
 
 #undef DEFINE_CUDA_CLUSTER_TYPES
-#undef DEFINE_BINDINGS_CLUSTERFINDER_CUDA
+#undef DEFINE_BINDINGS_CLUSTERFINDER_CUDA
+#undef DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH