mirror of
https://github.com/slsdetectorgroup/aare.git
synced 2026-11-01 13:41:28 +01:00
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.
This commit is contained in:
1 parent
1bf317f42a
commit
a42d71cf42
5 files changed
+732
-2
No files matched your search
@@ -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
|
||||
@@ -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
|
||||
Reference in new issue
Block a user