// SPDX-License-Identifier: MPL-2.0 #pragma once #include "aare/ClusterFinderCUDA.hpp" #include "aare/ClusterVector.hpp" #include "aare/NDView.hpp" #include "aare/Pedestal.hpp" #include "np_helper.hpp" #include #include #include // #include namespace py = pybind11; using pd_type = double; using namespace aare; #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wunused-parameter" namespace aare { template void define_ClusterFinderCUDA(py::module &m, const std::string &typestr) { auto class_name = fmt::format("ClusterFinderCUDA_{}", typestr); using ClusterType = Cluster; using CF = ClusterFinderCUDA; using ContigArr = py::array_t; // Opaque batch handle returned by submit_batch() and consumed by collect() py::class_(m, (class_name + "_BatchToken").c_str()); using VT = typename ClusterType::value_type; constexpr size_t NPIX = static_cast(ClusterSizeX) * static_cast(ClusterSizeY); // Zero-copy view into the finder's pinned D2H buffer. py::class_(m, (class_name + "_BatchView").c_str()) .def_property_readonly("n_frames", &CF::BatchView::n_frames) .def_property_readonly("first_frame", &CF::BatchView::first_frame) .def_property_readonly("valid", &CF::BatchView::valid) .def_property_readonly("total_clusters", &CF::BatchView::total_clusters) .def("count", &CF::BatchView::count, py::arg("frame_index")) .def("release", &CF::BatchView::release, R"(Give the slot back to the finder. The view is unusable afterwards and submit_batch() may reuse the buffer. Called automatically on destruction and on __exit__.)") .def("__enter__", [](py::object self) { return self; }) .def("__exit__", [](typename CF::BatchView &v, py::object, py::object, py::object) { v.release(); }) .def_property_readonly( "counts", [](const typename CF::BatchView &v) { py::array_t out(v.n_frames()); auto *o = out.mutable_data(); for (size_t i = 0; i < v.n_frames(); ++i) o[i] = v.count(i); return out; }, R"(Clusters found per frame, as a numpy array.)") .def( "sums", [](const typename CF::BatchView &v) { std::vector s; { // Reduce without the GIL, but take it back before building // the array — constructing a Python object without it is a // segfault, so this cannot be a call_guard. py::gil_scoped_release nogil; s = v.sums(); } return py::array_t(s.size(), s.data()); }, R"(Per-cluster sums for the whole batch, reduced in C++ straight out of the pinned buffer — the clusters are never materialised on the host. This is the fast path for spectra/histograms.)") .def( "frame_data", [](py::object self, size_t i) { auto &v = self.cast(); const uint32_t n = v.count(i); // Zero-copy (n, NPIX) view; stride skips each cluster's x/y. return py::array_t( {static_cast(n), NPIX}, {sizeof(ClusterType), sizeof(VT)}, n == 0 ? nullptr : v.clusters(i)->data.data(), self); }, py::arg("frame_index"), R"(Zero-copy (n_clusters, ClusterSizeX*ClusterSizeY) view of one frame's pixel data, straight out of the pinned buffer. Valid until this view is released — copy it if you need to keep it.)") .def( "frame_xy", [](py::object self, size_t i) { auto &v = self.cast(); const uint32_t n = v.count(i); return py::array_t( {static_cast(n), size_t{2}}, {sizeof(ClusterType), sizeof(CoordType)}, n == 0 ? nullptr : &v.clusters(i)->x, self); }, py::arg("frame_index"), R"(Zero-copy (n_clusters, 2) view of one frame's cluster centre coordinates as (x, y).)"); py::class_(m, class_name.c_str()) .def(py::init, float, size_t, int, bool>(), py::arg("image_size"), py::arg("n_sigma") = 5.0f, py::arg("max_clusters_per_frame") = 2048, py::arg("n_streams") = 4, py::arg("time_kernels") = false) .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 = self.pedestal(); return return_image_data(pd); }) .def_property_readonly("noise", [](CF &self) { auto arr = new NDArray{}; *arr = self.noise(); return return_image_data(arr); }) .def( "device_pedestal", [](CF &self, int stream) { auto pd = new NDArray{}; *pd = self.device_pedestal(stream); return return_image_data(pd); }, py::arg("stream") = 0, R"(Device pedestal MEAN for a stream — the pedestal the kernel actually decides with and updates each frame, unlike `pedestal` (the frozen host pedestal). Read it before find_clusters() for the decision-time state.)") .def( "device_noise", [](CF &self, int stream) { auto arr = new NDArray{}; *arr = self.device_noise(stream); return return_image_data(arr); }, py::arg("stream") = 0, R"(Device pedestal RMS for a stream, computed as the kernel does: sqrt(max(sum2/n - mean^2, 0)). Counterpart to `noise` for the device pedestal.)") .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()) .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(), R"(Process a 3D array of frames (n_frames, nrows, ncols) using n_streams CUDA streams for H2D/kernel/D2H pipelining. Returns a list of ClusterVector, one per input frame. The input array is converted to C-contiguous uint16 if needed.)") .def( "submit_batch", [](CF &self, ContigArr frames, uint64_t first_frame) { auto view = make_view_3d(frames); return self.submit_batch(view, first_frame); }, py::arg("frames"), py::arg("first_frame") = 0, py::call_guard(), R"(Enqueue one batch of frames onto the GPU without waiting for completion. Returns a BatchToken that must be passed to collect() to retrieve results and release the slot. At most 2 batches can be in flight simultaneously. The intended usage pattern to eliminate inter-batch GPU idle time is: tok = cf.submit_batch(buf_a, first_frame=0) for start in range(BATCH_SIZE, N, BATCH_SIZE): buf_b[:n] = data[start:start+n] # fill next buffer next_tok = cf.submit_batch(buf_b, first_frame=start) results += cf.collect(tok) # GPU runs buf_b tok = next_tok buf_a, buf_b = buf_b, buf_a # swap results += cf.collect(tok) # drain last batch Two separate input buffers must be used (one per in-flight batch) so that filling the next buffer does not corrupt the ongoing H2D transfer for the current batch.)") .def( "collect", [](CF &self, typename CF::BatchToken token) { return self.collect(token); }, py::arg("token"), py::call_guard(), R"(Wait for a previously submitted batch and return its results as a list of ClusterVector, one per input frame. Releases the batch slot so it can be reused by the next submit_batch() call. One allocation and one copy per frame; see collect_view() for neither.)") .def( "collect_view", [](CF &self, typename CF::BatchToken token) { return self.collect_view(token); }, py::arg("token"), py::call_guard(), R"(Like collect(), but copies nothing: returns a BatchView onto the finder's pinned D2H buffer. The slot stays reserved until the view is released, so with 2 slots you must finish with it before submitting two more batches. Use it as a context manager, or call release(): tok = cf.submit_batch(frames[a:b], first_frame=a) with cf.collect_view(tok) as v: hist.fill(v.sums()) # reduced in C++, nothing copied Anything you want to keep past release() must be copied out.)") .def("avg_kernel_time_ms", &CF::avg_kernel_time_ms, R"(Average per-frame kernel time in ms, or NaN if the finder was constructed with time_kernels=False (the default). WARNING: only meaningful with n_streams=1. The CUDA events bracket the kernel on its own stream, so under multi-stream contention the interval includes time spent queued behind other streams' kernels and over-reads by up to ~3.5x. Use Nsight Systems for exclusive kernel times.)") .def("kernel_timing_enabled", &CF::kernel_timing_enabled, R"(True if per-frame kernel timing was enabled at construction.)") .def("chunk_size_for", &CF::chunk_size_for, py::arg("n_frames"), R"(The chunk size find_clusters_batched() would use for n_frames. Use it to match its pipelining when driving submit_batch/ collect_view by hand.)") .def("reserve_output_slots", &CF::reserve_output_slots, py::arg("n_frames"), py::call_guard(), R"(Pre-allocate both pinned output slots for batches of n_frames. Processes nothing: no transfer, no kernel, and the pedestal is NOT advanced. Only the two cudaMallocHost calls that the first submit_batch() would make happen here. Page-locking runs about 1.0 us per 4 kB page (~66 ms for two 128 MiB slots), and that cost is charged to whoever triggers it. Call this before starting a timer so it does not land inside the measurement. Slots only grow, so pass the largest batch you intend to use: cf.reserve_output_slots(cf.chunk_size_for(len(data))) )") .def_property( "batch_chunk", &CF::get_batch_chunk, &CF::set_batch_chunk, R"(Frames per internally pipelined chunk in find_clusters_batched(). 0 (default) = auto: the batch is split into ~8 chunks so the host marshals one chunk while the GPU runs the next. Rounded up to a multiple of n_streams, which keeps the frame->stream assignment — and therefore the per-stream device pedestal each frame sees — identical to processing the batch in one go. Set equal to the batch size to disable chunking and recover the old submit-everything-then-copy-everything behaviour.)") .def("reset_timers", &CF::reset_timers, R"(Reset the internal kernel timing counters.)") .def( "register_input_buffer", [](CF &self, py::array arr) { auto info = arr.request(); self.register_input_buffer( info.ptr, static_cast(info.size) * static_cast(info.itemsize)); }, R"(Pin a numpy array as a locked host buffer so that find_clusters_batched transfers it at full DMA bandwidth (~22 GB/s) instead of going through the CUDA driver's internal staging (~15 GB/s for pageable memory). Call once before the processing loop with the full data array. Slices of that array 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.") .def( "pin_buffer", [](CF & /*self*/, py::array arr) { auto info = arr.request(); CUDA_CHECK( cudaHostRegister(info.ptr, static_cast(info.size) * static_cast(info.itemsize), cudaHostRegisterDefault)); }, R"(Pin an arbitrary numpy array as a locked host buffer for DMA-speed transfers. Unlike register_input_buffer(), does not unpin a previously registered buffer — use this to pin multiple buffers simultaneously (e.g. the two alternating buffers in an async pipeline). Call unpin_buffer() on each array when done.)") .def( "unpin_buffer", [](CF & /*self*/, py::array arr) { auto info = arr.request(); CUDA_CHECK(cudaHostUnregister(info.ptr)); }, "Release a buffer previously pinned with pin_buffer()."); } } // namespace aare #pragma GCC diagnostic pop