Benchmarking of CUDA cluster finder

python/aare/ClusterFinder.py

@@ -166,7 +166,34 @@ def ClusterFinderCUDAGraph(image_size, cluster_size=(3,3), n_sigma=5, dtype=np.i
                n_streams=n_streams)
 
 
-def ClusterCollector(clusterfindermt, dtype=np.int32): 
+def ClusterFinderCUDAOpt2(image_size, cluster_size=(3, 3), n_sigma=5, dtype=np.int32,
+                          max_clusters_per_frame=3000, n_streams=4):
+    """
+    Factory for the OPT2 snapshot finder — the pre-refactor pipeline (per-frame
+    pinned staging, round-robin streams with sync barriers, variable-length D2H),
+    kept only for benchmarking the optimization arc against the current finder.
+
+    It uses its own kernel snapshot (clusterfinder_kernel_opt2.cuh): f32 stencil
+    / f64 pedestal. The Test3 local-max gate has been backported so its cluster
+    counts match the CPU and current finders (correctness held constant across
+    the opt arc; only the pipeline differs).
+
+    Only the 3x3 cluster size is registered.
+    """
+    if not _cuda_available():
+        raise RuntimeError(
+            "ClusterFinderCUDAOpt2 is not available in this build of aare. "
+            "Rebuild with -DAARE_CUDA=ON (and -DAARE_PYTHON_BINDINGS=ON)."
+        )
+
+    cls = _get_class("ClusterFinderCUDAOpt2", 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 
     the templated ClusterCollector in C++.

python/aare/__init__.py

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

python/src/bind_ClusterFinderCUDAOpt2.hpp

@@ -0,0 +1,102 @@
+// SPDX-License-Identifier: MPL-2.0
+#pragma once
+#include "aare/ClusterFinderCUDAOpt2.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 {
+
+// Binding for the OPT2 snapshot finder (pre-refactor pipeline: per-frame pinned
+// staging, round-robin streams with sync barriers, variable-length D2H). Kept
+// only for benchmarking the optimization arc; not part of the shipped API.
+template <typename T, uint8_t ClusterSizeX, uint8_t ClusterSizeY,
+          typename CoordType = uint16_t>
+void define_ClusterFinderCUDAOpt2(py::module &m, const std::string &typestr) {
+    auto class_name = fmt::format("ClusterFinderCUDAOpt2_{}", typestr);
+
+    using ClusterType = Cluster<T, ClusterSizeX, ClusterSizeY, CoordType>;
+    using CF = ClusterFinderCUDAOpt2<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())
+        // ctor: (image_size, n_sigma, capacity, n_streams) — capacity is the
+        // per-stream device cluster buffer (upper bound on clusters/frame).
+        .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") = 3000, 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.)")
+
+        .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 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 (n_frames, nrows, ncols) round-robin across
+n_streams. Returns a list of ClusterVector, one per input frame.)")
+
+        .def("avg_kernel_time_ms", &CF::avg_kernel_time_ms)
+        .def("reset_timers", &CF::reset_timers);
+}
+
+} // namespace aare
+
+#pragma GCC diagnostic pop

python/src/cuda_bindings.cu

@@ -8,6 +8,7 @@
 #include "bind_Cluster.hpp"
 #include "bind_ClusterFinderCUDA.hpp"
 #include "bind_ClusterFinderCUDAGraph.hpp"
+#include "bind_ClusterFinderCUDAOpt2.hpp"
 #include "bind_ClusterVector.hpp"
 
 #include <pybind11/pybind11.h>
@@ -30,6 +31,10 @@ namespace py = pybind11;
     aare::define_ClusterFinderCUDAGraph<T, N, M, U>(m, "Cluster" #N            \
                                                        "x" #M #TYPE_CODE);
 
+#define DEFINE_BINDINGS_CLUSTERFINDER_CUDA_OPT2(T, N, M, U, TYPE_CODE)         \
+    aare::define_ClusterFinderCUDAOpt2<T, N, M, U>(m, "Cluster" #N             \
+                                                      "x" #M #TYPE_CODE);
+
 PYBIND11_MODULE(_aare_cuda, m) {
 
     // Types first — finders reference them in their signatures.
@@ -83,8 +88,14 @@ PYBIND11_MODULE(_aare_cuda, m) {
     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);
+
+    // OPT2 snapshot finder (benchmark only) — 3x3 is what the deck uses.
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_OPT2(int, 3, 3, uint16_t, i);
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_OPT2(double, 3, 3, uint16_t, d);
+    DEFINE_BINDINGS_CLUSTERFINDER_CUDA_OPT2(float, 3, 3, uint16_t, f);
 }
 
 #undef DEFINE_CUDA_CLUSTER_TYPES
 #undef DEFINE_BINDINGS_CLUSTERFINDER_CUDA
-#undef DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH
+#undef DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH
+#undef DEFINE_BINDINGS_CLUSTERFINDER_CUDA_OPT2

python/tests/nsys_kernel_probe.py

@@ -0,0 +1,32 @@
+# Minimal probe for nsys: train pedestal, run one batched pass, print summary.
+# Usage: nsys_kernel_probe.py [n_streams] [n_frames]
+import sys
+sys.path.append('/home/ferjao_k/aare/build')
+
+from pathlib import Path
+import time
+from aare import File, ClusterFinderCUDA
+
+n_streams = int(sys.argv[1]) if len(sys.argv) > 1 else 8
+N         = int(sys.argv[2]) if len(sys.argv) > 2 else 2000
+
+base = Path('/mnt/sls_det_storage/moench_data/2603_MaxIVBeamtime/2026032408/process/xrf/')
+f  = File(base / 'Cu_factor_10_data_master_0.json')
+pd = File(base / 'Cu_factor_10_pedestal_master_0.json')
+
+cf = ClusterFinderCUDA((f.rows, f.cols), (3, 3), n_sigma=5,
+                       max_clusters_per_frame=3000, n_streams=n_streams)
+for _ in range(1000):
+    cf.push_pedestal_frame(pd.read_frame().copy())
+
+data = f.read_n(N)
+cf.register_input_buffer(data)
+
+t0 = time.perf_counter()
+res = cf.find_clusters_batched(data, first_frame=0)
+t = time.perf_counter() - t0
+
+cf.unregister_input_buffer()
+n = sum(cv.size for cv in res)
+print(f'n_streams={n_streams}  N={N}  wall={t:.3f}s  ({N/t:.0f} FPS)  '
+      f'clusters/frame={n/N:.2f}  event kernel_ms={cf.avg_kernel_time_ms():.3f}')