mirror of
https://github.com/slsdetectorgroup/aare.git
synced 2026-10-31 12:41:32 +01:00
kferjaoui
ea32b530e8
PixelHistogramCUDA_<input>_<storage> for uint16, float32 and float64 frames and the six storage types of PixelHistogram, behind an aare.PixelHistogramCUDA() factory. Frames are never converted: a wrong dtype raises TypeError, so a registered buffer is always the one copied. The GIL is released during GPU work. - C++: fill_histogram() and submit_batch() take const views, so read-only arrays (np.memmap mode="r") work; register_output_buffer() pins the destination of values(out). - Python: register_input_buffer()/register_output_buffer() keep their array alive; values(out=...) writes into a given array. With a pinned out, reading back a 640 MB histogram takes 28 ms instead of 230 ms. - Tests in test_PixelHistogramCUDA.py; PixelHistogramCUDA.ipynb compares CPU and GPU on 100k Moench frames (identical, ~13x faster on the GPU).
95 lines
4.3 KiB
Plaintext
95 lines
4.3 KiB
Plaintext
// SPDX-License-Identifier: MPL-2.0
|
|
//
|
|
// CUDA-only Python extension module. Registers ClusterFinderCUDA along with
|
|
// the ClusterVector and Cluster types it exposes in its return values, so
|
|
// the module is self-contained — users can call steal_clusters() and get
|
|
// back a usable ClusterVector without _aare needing to be imported first.
|
|
// Also registers PixelHistogramCUDA.
|
|
|
|
#include "bind_Cluster.hpp"
|
|
#include "bind_ClusterFinderCUDA.hpp"
|
|
#include "bind_ClusterFinderCUDAGraph.hpp"
|
|
#include "bind_ClusterVector.hpp"
|
|
#include "bind_PixelHistogramCUDA.hpp"
|
|
|
|
#include <pybind11/pybind11.h>
|
|
|
|
namespace py = pybind11;
|
|
|
|
// Register the Cluster + ClusterVector pair for one (T, N, M) combination.
|
|
// Subset of DEFINE_CLUSTER_BINDINGS from module.cpp: we register what
|
|
// ClusterFinderCUDA actually returns, nothing more. File I/O, eta and
|
|
// reduce_to_2x2 stay on the CPU side.
|
|
#define DEFINE_CUDA_CLUSTER_TYPES(T, N, M, U, TYPE_CODE) \
|
|
define_ClusterVector<T, N, M, U>(m, "Cluster" #N "x" #M #TYPE_CODE); \
|
|
define_Cluster<T, N, M, U>(m, #N "x" #M #TYPE_CODE);
|
|
|
|
#define DEFINE_BINDINGS_CLUSTERFINDER_CUDA(T, N, M, U, TYPE_CODE) \
|
|
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.
|
|
// SFINAE excludes 2x2 on ClusterFinderCUDA, so we skip it here too.
|
|
DEFINE_CUDA_CLUSTER_TYPES(int, 3, 3, uint16_t, i);
|
|
DEFINE_CUDA_CLUSTER_TYPES(double, 3, 3, uint16_t, d);
|
|
DEFINE_CUDA_CLUSTER_TYPES(float, 3, 3, uint16_t, f);
|
|
|
|
DEFINE_CUDA_CLUSTER_TYPES(int, 5, 5, uint16_t, i);
|
|
DEFINE_CUDA_CLUSTER_TYPES(double, 5, 5, uint16_t, d);
|
|
DEFINE_CUDA_CLUSTER_TYPES(float, 5, 5, uint16_t, f);
|
|
|
|
DEFINE_CUDA_CLUSTER_TYPES(int, 7, 7, uint16_t, i);
|
|
DEFINE_CUDA_CLUSTER_TYPES(double, 7, 7, uint16_t, d);
|
|
DEFINE_CUDA_CLUSTER_TYPES(float, 7, 7, uint16_t, f);
|
|
|
|
DEFINE_CUDA_CLUSTER_TYPES(int, 9, 9, uint16_t, i);
|
|
DEFINE_CUDA_CLUSTER_TYPES(double, 9, 9, uint16_t, d);
|
|
DEFINE_CUDA_CLUSTER_TYPES(float, 9, 9, uint16_t, f);
|
|
|
|
// Finders
|
|
DEFINE_BINDINGS_CLUSTERFINDER_CUDA(int, 3, 3, uint16_t, i);
|
|
DEFINE_BINDINGS_CLUSTERFINDER_CUDA(double, 3, 3, uint16_t, d);
|
|
DEFINE_BINDINGS_CLUSTERFINDER_CUDA(float, 3, 3, uint16_t, f);
|
|
|
|
DEFINE_BINDINGS_CLUSTERFINDER_CUDA(int, 5, 5, uint16_t, i);
|
|
DEFINE_BINDINGS_CLUSTERFINDER_CUDA(double, 5, 5, uint16_t, d);
|
|
DEFINE_BINDINGS_CLUSTERFINDER_CUDA(float, 5, 5, uint16_t, f);
|
|
|
|
DEFINE_BINDINGS_CLUSTERFINDER_CUDA(int, 7, 7, uint16_t, i);
|
|
DEFINE_BINDINGS_CLUSTERFINDER_CUDA(double, 7, 7, uint16_t, d);
|
|
DEFINE_BINDINGS_CLUSTERFINDER_CUDA(float, 7, 7, uint16_t, f);
|
|
|
|
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);
|
|
|
|
// Pixel histograms: PixelHistogramCUDA_<input>_<storage>
|
|
aare::define_PixelHistogramCUDA_bindings(m);
|
|
}
|
|
|
|
#undef DEFINE_CUDA_CLUSTER_TYPES
|
|
#undef DEFINE_BINDINGS_CLUSTERFINDER_CUDA
|
|
#undef DEFINE_BINDINGS_CLUSTERFINDER_CUDA_GRAPH |